feat: support pipeline-stage-local HF export (#6291)

Signed-off-by: yangbosong <2112536022@qq.com>
Signed-off-by: Chen Cui <chcui@nvidia.com>
Co-authored-by: Chen Cui <chcui@nvidia.com>

Author: Bsgg1

docs/bridge-guide.md

@@ -346,7 +346,7 @@ AutoBridge.to_megatron_model(load_weights: bool = True, hf_path: str | Path | No
 AutoBridge.load_hf_weights(model: list[MegatronModule], hf_path: str | Path | None = None) -> None
 
 # Megatron → HF conversion
-AutoBridge.export_hf_weights(model: list[MegatronModule], cpu: bool = False, show_progress: bool = True, conversion_tasks: Optional[list[WeightConversionTask]] = None) -> Iterable[HFWeightTuple]
+AutoBridge.export_hf_weights(model: list[MegatronModule], cpu: bool = False, show_progress: bool = True, conversion_tasks: Optional[list[WeightConversionTask]] = None, current_pp_stage_only: bool = False) -> Iterable[HFWeightTuple]
 AutoBridge.save_hf_pretrained(model: list[MegatronModule], path: str | Path, show_progress: bool = True) -> None
 AutoBridge.save_hf_weights(model: list[MegatronModule], path: str | Path, show_progress: bool = True) -> None
 

docs/bridge-rl-integration.md

@@ -352,7 +352,28 @@ for name, weight in bridge.export_hf_weights(megatron_model, cpu=True, show_prog
 
 Point your inference engine (e.g., vLLM) to `"/path/to/hf_export"`.
 
-### B) Zero-copy streaming via ZMQ (fast refit, colocated)
+### B) Pipeline-stage-local streaming for RL weight synchronization
+
+When inference workers collectively receive shards from every training pipeline stage, avoid replicating the full
+HF stream across PP ranks:
+
+```python
+for name, weight in bridge.export_hf_weights(
+    megatron_model,
+    show_progress=False,
+    current_pp_stage_only=True,
+):
+    send_to_inference_workers(name, weight)
+```
+
+- `current_pp_stage_only` changes ownership only across the PP axis.
+- TP/EP ranks in the owning stage still gather shards and perform normal HF layout conversion.
+- Combining one representative stream from every PP stage reconstructs the normal complete export.
+- Source-only HF passthrough tensors are assigned to PP stage 0; duplicate Megatron parameters use the lowest PP owner.
+
+This mode is opt-in. Omitting the flag preserves the complete stream on every rank.
+
+### C) Zero-copy streaming via ZMQ (fast refit, colocated)
 
 Stream tensors from the training side to your inference runtime without writing to disk. The transport is ZMQ peer-to-peer with async send/recv and ping‑pong buffers for overlap; Ray is used only for lightweight coordination. This replaces the earlier ad‑hoc per‑tensor IPC handle passing and aligns with the refactor in [NVIDIA-NeMo/RL#1267](https://github.com/NVIDIA-NeMo/RL/pull/1267).
 

docs/fern/versions/nightly/pages/bridge-guide.mdx

@@ -309,7 +309,7 @@ AutoBridge.to_megatron_model(load_weights: bool = True, hf_path: str | Path | No
 AutoBridge.load_hf_weights(model: list[MegatronModule], hf_path: str | Path | None = None) -> None
 
 # Megatron → HF conversion
-AutoBridge.export_hf_weights(model: list[MegatronModule], cpu: bool = False, show_progress: bool = True, conversion_tasks: Optional[list[WeightConversionTask]] = None) -> Iterable[HFWeightTuple]
+AutoBridge.export_hf_weights(model: list[MegatronModule], cpu: bool = False, show_progress: bool = True, conversion_tasks: Optional[list[WeightConversionTask]] = None, current_pp_stage_only: bool = False) -> Iterable[HFWeightTuple]
 AutoBridge.save_hf_pretrained(model: list[MegatronModule], path: str | Path, show_progress: bool = True) -> None
 AutoBridge.save_hf_weights(model: list[MegatronModule], path: str | Path, show_progress: bool = True) -> None
 

src/megatron/bridge/models/conversion/auto_bridge.py

@@ -294,29 +294,34 @@ def _drop_readonly_config_properties(
     return {key: value for key, value in config_dict.items() if key not in readonly_properties}
 
 
-def _check_with_megatron_names_support(stream_fn: Callable[..., Any], owner: object) -> None:
-    """Fail up front when a bridge's streaming override cannot accept ``with_megatron_names``.
-
-    Bridges that override the streaming export with an explicit signature (rather than
-    ``*args, **kwargs``) have to forward the keyword themselves; without this check the
-    request would only surface as an opaque ``TypeError`` once the generator is consumed.
-    """
+def _check_streaming_argument_support(stream_fn: Callable[..., Any], owner: object, argument: str) -> None:
+    """Fail early when a streaming override cannot accept an export argument."""
     try:
         parameters = inspect.signature(stream_fn).parameters
     except (TypeError, ValueError):
         # Not introspectable (C callables); let the call itself decide.
         return
-    if "with_megatron_names" in parameters or any(
+    if argument in parameters or any(
         parameter.kind is inspect.Parameter.VAR_KEYWORD for parameter in parameters.values()
     ):
         return
     raise TypeError(
         f"{type(owner).__name__}.{getattr(stream_fn, '__name__', 'stream_weights')} does not accept "
-        "'with_megatron_names'; a bridge that overrides the streaming export must forward this flag "
-        "before its exported weights can carry source Megatron parameter names."
+        f"'{argument}'; a bridge that overrides the streaming export must forward this flag "
+        "before the requested export mode can be used."
     )
 
 
+def _check_with_megatron_names_support(stream_fn: Callable[..., Any], owner: object) -> None:
+    """Fail up front when a streamer cannot accept ``with_megatron_names``."""
+    _check_streaming_argument_support(stream_fn, owner, "with_megatron_names")
+
+
+def _check_current_pp_stage_only_support(stream_fn: Callable[..., Any], owner: object) -> None:
+    """Fail early when a streamer cannot accept PP-local export."""
+    _check_streaming_argument_support(stream_fn, owner, "current_pp_stage_only")
+
+
 class AutoBridge(Generic[MegatronModelT]):
     """
     Automatically select and instantiate the appropriate bridge for a model.
@@ -801,13 +806,15 @@ def export_hf_weights(
         merge_adapter_weights: bool = True,
         weight_dtype: Optional[torch.dtype] = None,
         with_megatron_names: bool = False,
+        current_pp_stage_only: bool = False,
     ) -> Iterable["HFWeightTuple"]:
         """
         Export Megatron model weights to HuggingFace format.
 
         This method yields weight tensors in HuggingFace format, handling the
         gathering of distributed tensors and format conversion. It's useful for
-        streaming weight export or custom processing. All ranks get full tensors.
+        streaming weight export or custom processing. By default all ranks get full tensors;
+        set ``current_pp_stage_only=True`` to yield only the current pipeline stage's tensors.
 
         If the model contains LoRA adapters, they will be automatically merged
         into the base weights before export. This ensures the exported model
@@ -832,6 +839,10 @@ def export_hf_weights(
                 a grouped-expert export, and none for HF-only passthrough tensors. Bridges that
                 override ``stream_weights_megatron_to_hf`` must accept the flag; otherwise a
                 ``TypeError`` is raised before any weight is streamed.
+            current_pp_stage_only: Export only parameters owned by the current pipeline stage and
+                suppress cross-stage PP broadcasts. TP/EP gathering and HF layout conversion remain
+                enabled. Defaults to False. Bridges that override the streamer must accept this
+                flag when it is enabled.
 
         Yields:
             HFWeightTuple: Named tuples of (param_name, weight_tensor), or HFSourcedWeightTuple
@@ -859,6 +870,8 @@ def export_hf_weights(
         bridge = self._model_bridge
         if with_megatron_names:
             _check_with_megatron_names_support(bridge.stream_weights_megatron_to_hf, bridge)
+        if current_pp_stage_only:
+            _check_current_pp_stage_only_support(bridge.stream_weights_megatron_to_hf, bridge)
         return bridge.stream_weights_megatron_to_hf(
             model,
             self.hf_pretrained,
@@ -869,6 +882,7 @@ def export_hf_weights(
             weight_dtype=weight_dtype,
             # Only forward the flag when set so bridges with a custom streamer keep working by default.
             **({"with_megatron_names": True} if with_megatron_names else {}),
+            **({"current_pp_stage_only": True} if current_pp_stage_only else {}),
         )
 
     def export_hf_weights_modelopt(