From f123130668d5de9db97ca1d7f7601f7c2eff8e68 Mon Sep 17 00:00:00 2001
From: Chen Cui <chcui@nvidia.com>
Date: Mon, 15 Jun 2026 13:34:24 -0700
Subject: [PATCH 1/5] feat(conversion): support distributed adapter export
 (#4221)

Signed-off-by: Chen Cui <chcui@nvidia.com>
(cherry picked from commit 9bfa760bc2af050ef57636d605196e74073ef887)
---
 examples/conversion/adapter/README.md         |  27 +-
 examples/conversion/adapter/export_adapter.py | 195 +++-
 examples/peft/merge_lora.py                   |  11 +-
 .../bridge/models/conversion/auto_bridge.py   |  53 +-
 .../bridge/models/conversion/model_bridge.py  |   3 +
 .../bridge/models/conversion/peft_bridge.py   |  13 +-
 src/megatron/bridge/peft/utils.py             | 968 ++++++++++++++++--
 .../unit_tests/models/test_adapter_export.py  | 519 +++++++++-
 tests/unit_tests/models/test_auto_bridge.py   |   1 +
 .../models/test_model_bridge_lora.py          | 208 +++-
 tests/unit_tests/peft/test_utils.py           | 775 +++++++++++---
 11 files changed, 2516 insertions(+), 257 deletions(-)

diff --git a/examples/conversion/adapter/README.md b/examples/conversion/adapter/README.md
index 1f0bf5a276..4f02874e1c 100644
--- a/examples/conversion/adapter/README.md
+++ b/examples/conversion/adapter/README.md
@@ -29,14 +29,30 @@ model = PeftModel.from_pretrained(base, "./my_adapter")
 
 ### 1. `export_adapter.py` — Checkpoint Export
 
-Converts a Megatron-Bridge PEFT checkpoint to HuggingFace PEFT format. Runs
-entirely on CPU — no GPU required.
+Converts a Megatron-Bridge PEFT checkpoint to HuggingFace PEFT format. The
+default path runs on CPU. For large hybrid models or sharded checkpoints that
+require CUDA-backed model materialization, launch the distributed GPU path with
+TP/PP/EP settings matching the checkpoint.
 
 ```bash
 uv run python examples/conversion/adapter/export_adapter.py \
     --hf-model-path meta-llama/Llama-3.2-1B \
     --lora-checkpoint /path/to/finetune_ckpt \
-    --output ./my_adapter
+    --output ./my_adapter \
+    --exclude-adapter-base-prefix mtp.layers
+```
+
+```bash
+# Multi-GPU export matching a TP=2, EP=16 checkpoint
+uv run python -m torch.distributed.run --nproc_per_node=16 \
+    examples/conversion/adapter/export_adapter.py \
+    --hf-model-path nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16 \
+    --lora-checkpoint /path/to/finetune_ckpt/iter_0000300 \
+    --output ./my_adapter \
+    --trust-remote-code \
+    --dtype bf16 \
+    --exclude-adapter-base-prefix mtp.layers \
+    --tp 2 --pp 1 --ep 16 --sequence-parallel
 ```
 
 | Argument | Description |
@@ -45,6 +61,10 @@ uv run python examples/conversion/adapter/export_adapter.py \
 | `--lora-checkpoint` | Path to the Megatron-Bridge distributed checkpoint containing LoRA adapter weights |
 | `--output` | Output directory (default: `./my_adapter`) |
 | `--trust-remote-code` | Allow custom code from the HuggingFace repository |
+| `--dtype` | Dtype used to materialize the model for distributed GPU export (default: `float32`) |
+| `--exclude-adapter-base-prefix` | Megatron adapter base prefix to skip during export; can be repeated |
+| `--tp`, `--pp`, `--ep`, `--etp` | Distributed GPU export parallelism; use values matching the checkpoint |
+| `--sequence-parallel` | Enable sequence parallelism for distributed GPU export |
 
 **Output structure:**
 
@@ -141,6 +161,7 @@ bridge = AutoBridge.from_hf_pretrained("meta-llama/Llama-3.2-1B")
 bridge.export_adapter_ckpt(
     peft_checkpoint="/path/to/finetune_ckpt",
     output_path="./my_adapter",
+    exclude_adapter_base_prefixes=("mtp.layers",),
 )
 
 # Or, if you already have a model in memory:
diff --git a/examples/conversion/adapter/export_adapter.py b/examples/conversion/adapter/export_adapter.py
index f4a152fc3e..894ddc90c5 100644
--- a/examples/conversion/adapter/export_adapter.py
+++ b/examples/conversion/adapter/export_adapter.py
@@ -17,7 +17,9 @@
 Export LoRA adapter weights from a Megatron-Bridge PEFT checkpoint to
 HuggingFace PEFT format (``adapter_config.json`` + ``adapter_model.safetensors``).
 
-No GPU required -- runs entirely on CPU.
+The default path runs on CPU and is suitable for models whose architecture can be
+materialized without CUDA. Large hybrid models can use the distributed GPU path
+by passing TP/PP/EP settings that match the adapter checkpoint.
 
 The output can be loaded directly with::
 
@@ -32,15 +34,50 @@
     uv run python examples/conversion/adapter/export_adapter.py \\
         --hf-model-path meta-llama/Llama-3.2-1B \\
         --lora-checkpoint /path/to/finetune_ckpt \\
-        --output ./my_adapter
+        --output ./my_adapter \\
+        --exclude-adapter-base-prefix mtp.layers
 """
 
 from __future__ import annotations
 
 import argparse
+import logging
+from collections.abc import Mapping
 from pathlib import Path
 
+import torch
+import torch.distributed as dist
+from megatron.core import dist_checkpointing, parallel_state
+from transformers import AutoConfig
+
 from megatron.bridge import AutoBridge
+from megatron.bridge.peft.lora import LoRA, VLMLoRA
+from megatron.bridge.peft.utils import enable_legacy_shared_expert_adapter_loading
+from megatron.bridge.training.checkpointing import (
+    _generate_model_state_dict,
+    apply_peft_adapter_filter_to_state_dict,
+)
+from megatron.bridge.training.utils.checkpoint_utils import read_run_config
+from megatron.bridge.utils.activation_map import str_to_dtype
+from megatron.bridge.utils.common_utils import get_local_rank_preinit
+
+
+logger = logging.getLogger(__name__)
+
+_SUPPORTED_EXPORT_DTYPES = {torch.float32, torch.float16, torch.bfloat16}
+
+
+def _parse_dtype(dtype: str) -> torch.dtype:
+    try:
+        parsed_dtype = str_to_dtype(dtype.lower())
+    except ValueError as err:
+        raise argparse.ArgumentTypeError(str(err)) from err
+    if parsed_dtype not in _SUPPORTED_EXPORT_DTYPES:
+        supported = ", ".join(sorted(str(dtype).replace("torch.", "") for dtype in _SUPPORTED_EXPORT_DTYPES))
+        raise argparse.ArgumentTypeError(
+            f"Unsupported adapter export dtype: {parsed_dtype}. Supported values: {supported}"
+        )
+    return parsed_dtype
 
 
 def parse_args() -> argparse.Namespace:
@@ -61,18 +98,162 @@ def parse_args() -> argparse.Namespace:
     )
     parser.add_argument("--output", type=Path, default=Path("./my_adapter"))
     parser.add_argument("--trust-remote-code", action="store_true")
+    parser.add_argument(
+        "--dtype",
+        type=_parse_dtype,
+        default=torch.float32,
+        help="Dtype used to materialize the model for distributed GPU export.",
+    )
+    parser.add_argument(
+        "--exclude-adapter-base-prefix",
+        action="append",
+        default=[],
+        help=(
+            "Megatron adapter base prefix to skip during export, before HF mapping lookup. "
+            "Can be specified multiple times; e.g. `mtp.layers` excludes MTP adapters."
+        ),
+    )
+    parser.add_argument("--tp", type=int, default=1, help="Tensor parallel size for distributed GPU export.")
+    parser.add_argument("--pp", type=int, default=1, help="Pipeline parallel size for distributed GPU export.")
+    parser.add_argument("--ep", type=int, default=1, help="Expert parallel size for distributed GPU export.")
+    parser.add_argument("--etp", type=int, default=1, help="Expert tensor parallel size for distributed GPU export.")
+    parser.add_argument("--sequence-parallel", action="store_true", help="Enable sequence parallelism.")
     return parser.parse_args()
 
 
+def _load_lora_config(ckpt_path: Path) -> LoRA | VLMLoRA:
+    peft_class: type[LoRA | VLMLoRA] = LoRA
+    peft_cfg: dict = {}
+    cfg_file = ckpt_path / "run_config.yaml"
+    if not cfg_file.exists() and ckpt_path.parent != ckpt_path:
+        cfg_file = ckpt_path.parent / "run_config.yaml"
+    if cfg_file.exists():
+        try:
+            run_cfg_dict = read_run_config(str(cfg_file))
+        except Exception as err:
+            logger.warning("Failed to read LoRA settings from %s: %s. Using defaults.", cfg_file, err)
+        else:
+            peft_cfg = run_cfg_dict.get("peft", {}) or {}
+            if "VLMLoRA" in peft_cfg.get("_target_", ""):
+                peft_class = VLMLoRA
+            vlm_only_keys = {"freeze_language_model", "freeze_vision_model", "freeze_vision_projection"}
+            allowed_keys = {
+                "target_modules",
+                "exclude_modules",
+                "dim",
+                "alpha",
+                "dropout",
+                "dropout_position",
+                "normalize_moe_lora",
+                "share_expert_adapters",
+            }
+            if peft_class is VLMLoRA:
+                allowed_keys |= vlm_only_keys
+            peft_cfg = {key: value for key, value in peft_cfg.items() if key in allowed_keys}
+    return peft_class(**peft_cfg)
+
+
+def _get_loaded_model_key(loaded_sd: Mapping[str, object], ckpt_path: Path) -> str:
+    if "model" in loaded_sd:
+        return "model"
+
+    model_key = next((key for key in loaded_sd if key.startswith("model")), None)
+    if model_key is None:
+        raise RuntimeError(f"Checkpoint at {ckpt_path} has no 'model' key. Available keys: {list(loaded_sd.keys())}")
+    return model_key
+
+
+def _uses_distributed_export(args: argparse.Namespace) -> bool:
+    return args.tp > 1 or args.pp > 1 or args.ep > 1 or args.etp > 1
+
+
+def _configure_cuda_device() -> torch.device:
+    if not torch.cuda.is_available():
+        raise RuntimeError("Distributed adapter export requires CUDA. Use the default CPU path for TP=PP=EP=ETP=1.")
+    local_rank = get_local_rank_preinit()
+    torch.cuda.set_device(local_rank)
+    return torch.device("cuda", local_rank)
+
+
+def _export_adapter_distributed(args: argparse.Namespace) -> None:
+    device = _configure_cuda_device()
+    ckpt_path = Path(args.lora_checkpoint).expanduser().resolve()
+    if not ckpt_path.exists():
+        raise FileNotFoundError(f"PEFT checkpoint not found: {ckpt_path}")
+    config = AutoConfig.from_pretrained(args.hf_model_path, trust_remote_code=args.trust_remote_code)
+    bridge = AutoBridge.from_hf_config(config)
+    lora = _load_lora_config(ckpt_path)
+
+    provider = bridge.to_megatron_provider(load_weights=False)
+    provider.tensor_model_parallel_size = args.tp
+    provider.pipeline_model_parallel_size = args.pp
+    provider.expert_model_parallel_size = args.ep
+    provider.expert_tensor_parallel_size = args.etp
+    provider.sequence_parallel = args.sequence_parallel
+    provider.pipeline_dtype = args.dtype
+    provider.params_dtype = args.dtype
+    provider.finalize()
+    provider.register_pre_wrap_hook(lambda chunks: lora(chunks, training=False))
+    try:
+        provider.initialize_model_parallel(seed=0)
+
+        model = provider.provide_distributed_model(
+            wrap_with_ddp=False,
+            use_cpu_initialization=False,
+            init_model_with_meta_device=False,
+        )
+        model = [chunk.to(device) for chunk in model]
+        if len(model) != 1:
+            raise RuntimeError(
+                "Distributed adapter export currently supports exactly one local model chunk; "
+                f"got {len(model)}. Use pipeline parallel size 1 without virtual pipeline parallelism."
+            )
+
+        sharded_state_dict = _generate_model_state_dict(model, {})
+        sharded_state_dict = apply_peft_adapter_filter_to_state_dict(sharded_state_dict, lora)
+        legacy_shared_expert_adapter = enable_legacy_shared_expert_adapter_loading(
+            model, sharded_state_dict, ckpt_path
+        )
+        if legacy_shared_expert_adapter:
+            sharded_state_dict = _generate_model_state_dict(model, {})
+            sharded_state_dict = apply_peft_adapter_filter_to_state_dict(sharded_state_dict, lora)
+        loaded_sd = dist_checkpointing.load(
+            sharded_state_dict,
+            str(ckpt_path),
+            validate_access_integrity=not legacy_shared_expert_adapter,
+        )
+        model_key = _get_loaded_model_key(loaded_sd, ckpt_path)
+        model[0].load_state_dict(loaded_sd[model_key], strict=False)
+
+        bridge.save_hf_adapter(
+            model,
+            path=args.output,
+            peft_config=lora,
+            base_model_name_or_path=args.hf_model_path,
+            exclude_adapter_base_prefixes=tuple(args.exclude_adapter_base_prefix),
+        )
+    finally:
+        if parallel_state.is_initialized():
+            parallel_state.destroy_model_parallel()
+        if dist.is_initialized():
+            dist.destroy_process_group()
+
+
 def main() -> None:
     """Export a Megatron-Bridge PEFT checkpoint to HuggingFace PEFT format."""
     args = parse_args()
 
-    bridge = AutoBridge.from_hf_pretrained(args.hf_model_path, trust_remote_code=args.trust_remote_code)
-    bridge.export_adapter_ckpt(
-        peft_checkpoint=args.lora_checkpoint,
-        output_path=args.output,
-    )
+    if _uses_distributed_export(args):
+        _export_adapter_distributed(args)
+    else:
+        if args.dtype != torch.float32:
+            raise ValueError("--dtype is only supported by distributed GPU export; CPU export uses float32.")
+        bridge = AutoBridge.from_hf_pretrained(args.hf_model_path, trust_remote_code=args.trust_remote_code)
+        bridge.export_adapter_ckpt(
+            peft_checkpoint=args.lora_checkpoint,
+            output_path=args.output,
+            exclude_adapter_base_prefixes=tuple(args.exclude_adapter_base_prefix),
+        )
 
 
 if __name__ == "__main__":
diff --git a/examples/peft/merge_lora.py b/examples/peft/merge_lora.py
index c63e5156c8..0f2f8653de 100644
--- a/examples/peft/merge_lora.py
+++ b/examples/peft/merge_lora.py
@@ -57,6 +57,7 @@
 
 from megatron.bridge.models.conversion.auto_bridge import AutoBridge
 from megatron.bridge.peft.lora import LoRA, VLMLoRA
+from megatron.bridge.peft.utils import enable_legacy_shared_expert_adapter_loading
 from megatron.bridge.training.checkpointing import (
     _generate_model_state_dict,
     apply_peft_adapter_filter_to_state_dict,
@@ -223,9 +224,17 @@ def merge_lora(
     sharded_state_dict = _generate_model_state_dict(model, {})
     # Keep only LoRA adapter tensors (and any other trainable parameters) so we don't read unnecessary dense weights.
     sharded_state_dict = apply_peft_adapter_filter_to_state_dict(sharded_state_dict, lora_peft)
+    legacy_shared_expert_adapter = enable_legacy_shared_expert_adapter_loading(model, sharded_state_dict, lora_dir)
+    if legacy_shared_expert_adapter:
+        sharded_state_dict = _generate_model_state_dict(model, {})
+        sharded_state_dict = apply_peft_adapter_filter_to_state_dict(sharded_state_dict, lora_peft)
 
     # Load those tensors from the checkpoint directory
-    loaded_sd = dist_checkpointing.load(sharded_state_dict, str(lora_dir))
+    loaded_sd = dist_checkpointing.load(
+        sharded_state_dict,
+        str(lora_dir),
+        validate_access_integrity=not legacy_shared_expert_adapter,
+    )
     # dist_checkpointing.load returns the same nested dict structure; we need the model section
     model_section_key = "model" if "model" in loaded_sd else next(k for k in loaded_sd if k.startswith("model"))
     adapter_sd = loaded_sd[model_section_key]
diff --git a/src/megatron/bridge/models/conversion/auto_bridge.py b/src/megatron/bridge/models/conversion/auto_bridge.py
index 81c2acbb8b..a5173745ed 100644
--- a/src/megatron/bridge/models/conversion/auto_bridge.py
+++ b/src/megatron/bridge/models/conversion/auto_bridge.py
@@ -597,6 +597,7 @@ def export_adapter_weights(
         model: list[MegatronModelT],
         cpu: bool = True,
         show_progress: bool = True,
+        exclude_adapter_base_prefixes: Iterable[str] | None = None,
     ) -> Iterable["HFWeightTuple"]:
         """
         Export only adapter weights from a Megatron model without merging them into base tensors.
@@ -608,12 +609,19 @@ def export_adapter_weights(
             model: Megatron model instance or list of instances
             cpu: Whether to move tensors to CPU before yielding
             show_progress: Display progress bar during export
+            exclude_adapter_base_prefixes: Megatron adapter base prefixes to
+                skip before resolving HuggingFace parameter mappings.
 
         Yields:
             HFWeightTuple: Named tuples of (param_name, weight_tensor) for adapter parameters
         """
         bridge = self._model_bridge
-        return bridge.stream_adapter_weights_megatron_to_hf(model, cpu=cpu, show_progress=show_progress)
+        return bridge.stream_adapter_weights_megatron_to_hf(
+            model,
+            cpu=cpu,
+            show_progress=show_progress,
+            exclude_adapter_base_prefixes=exclude_adapter_base_prefixes,
+        )
 
     def save_hf_adapter(
         self,
@@ -622,6 +630,7 @@ def save_hf_adapter(
         peft_config: "PEFT",
         base_model_name_or_path: Optional[str] = None,
         show_progress: bool = True,
+        exclude_adapter_base_prefixes: Iterable[str] | None = None,
     ) -> None:
         """Save LoRA adapter weights as a HuggingFace PEFT-compatible directory.
 
@@ -638,6 +647,8 @@ def save_hf_adapter(
                 of the base model this adapter was trained on.  If *None*, the
                 value is inferred from ``hf_pretrained.model_name_or_path``.
             show_progress: Display progress bar during export.
+            exclude_adapter_base_prefixes: Megatron adapter base prefixes to
+                skip before resolving HuggingFace parameter mappings.
 
         Example:
             >>> bridge.save_hf_adapter(
@@ -672,8 +683,13 @@ def save_hf_adapter(
             dist.barrier()
 
         raw_adapter_weights = [
-            HFWeightTuple(exported_weight.param_name, exported_weight.weight.clone().float())
-            for exported_weight in self.export_adapter_weights(model, cpu=True, show_progress=show_progress)
+            HFWeightTuple(exported_weight.param_name, exported_weight.weight.detach().clone())
+            for exported_weight in self.export_adapter_weights(
+                model,
+                cpu=True,
+                show_progress=show_progress,
+                exclude_adapter_base_prefixes=exclude_adapter_base_prefixes,
+            )
         ]
         if not raw_adapter_weights:
             raise RuntimeError(
@@ -1251,6 +1267,7 @@ def export_adapter_ckpt(
         peft_checkpoint: str | Path,
         output_path: str | Path,
         show_progress: bool = True,
+        exclude_adapter_base_prefixes: Iterable[str] | None = None,
     ) -> None:
         """Export LoRA adapter weights from a Megatron PEFT checkpoint to HuggingFace PEFT format.
 
@@ -1270,6 +1287,8 @@ def export_adapter_ckpt(
                 directory.
             output_path: Directory where the adapter files will be saved.
             show_progress: Display progress bar during export.
+            exclude_adapter_base_prefixes: Megatron adapter base prefixes to
+                skip before resolving HuggingFace parameter mappings.
 
         Example:
             >>> bridge = AutoBridge.from_hf_pretrained("meta-llama/Llama-3.2-1B")
@@ -1288,6 +1307,7 @@ def export_adapter_ckpt(
         from megatron.core import dist_checkpointing
 
         from megatron.bridge.peft.lora import LoRA, VLMLoRA
+        from megatron.bridge.peft.utils import enable_legacy_shared_expert_adapter_loading
         from megatron.bridge.training.checkpointing import (
             _generate_model_state_dict,
             apply_peft_adapter_filter_to_state_dict,
@@ -1313,6 +1333,7 @@ def export_adapter_ckpt(
                 peft_cfg = run_cfg_dict.get("peft", {}) or {}
                 if "VLMLoRA" in peft_cfg.get("_target_", ""):
                     peft_class = VLMLoRA
+                vlm_only_keys = {"freeze_language_model", "freeze_vision_model", "freeze_vision_projection"}
                 allowed_keys = {
                     "target_modules",
                     "exclude_modules",
@@ -1322,10 +1343,9 @@ def export_adapter_ckpt(
                     "dropout_position",
                     "normalize_moe_lora",
                     "share_expert_adapters",
-                    "freeze_language_model",
-                    "freeze_vision_model",
-                    "freeze_vision_projection",
                 }
+                if peft_class is VLMLoRA:
+                    allowed_keys |= vlm_only_keys
                 peft_cfg = {k: v for k, v in peft_cfg.items() if k in allowed_keys}
             except Exception as err:
                 _logger.warning(f"Failed to read LoRA settings from {cfg_file}: {err}. Using defaults.")
@@ -1334,10 +1354,10 @@ def export_adapter_ckpt(
 
         lora = peft_class(**peft_cfg)
 
-        # Materialise model with base weights + LoRA structure.
-        # Use float32 so adapter weights are exported at full precision;
-        # bfloat16 matmul in downstream PEFT merges causes ~1e-3 weight
-        # errors that compound into large logit diffs.
+        # Materialise model with base weights + LoRA structure. Use float32 so
+        # adapter weights are exported at full precision; bfloat16 downstream
+        # PEFT merges can introduce ~1e-3 weight errors that compound into
+        # large logit diffs.
         provider = self.to_megatron_provider(load_weights=True)
         provider.pipeline_dtype = torch.float32
         provider.params_dtype = torch.float32
@@ -1350,7 +1370,17 @@ def _load_and_export_adapter(model):
             # Load adapter weights from the PEFT checkpoint
             sharded_state_dict = _generate_model_state_dict(model, {})
             sharded_state_dict = apply_peft_adapter_filter_to_state_dict(sharded_state_dict, lora)
-            loaded_sd = dist_checkpointing.load(sharded_state_dict, str(ckpt_path))
+            legacy_shared_expert_adapter = enable_legacy_shared_expert_adapter_loading(
+                model, sharded_state_dict, ckpt_path
+            )
+            if legacy_shared_expert_adapter:
+                sharded_state_dict = _generate_model_state_dict(model, {})
+                sharded_state_dict = apply_peft_adapter_filter_to_state_dict(sharded_state_dict, lora)
+            loaded_sd = dist_checkpointing.load(
+                sharded_state_dict,
+                str(ckpt_path),
+                validate_access_integrity=not legacy_shared_expert_adapter,
+            )
             model_key = "model" if "model" in loaded_sd else next(k for k in loaded_sd if k.startswith("model"))
             model[0].load_state_dict(loaded_sd[model_key], strict=False)
 
@@ -1365,6 +1395,7 @@ def _load_and_export_adapter(model):
                 peft_config=lora,
                 base_model_name_or_path=base_model_name,
                 show_progress=show_progress,
+                exclude_adapter_base_prefixes=exclude_adapter_base_prefixes,
             )
 
         model_context = (
diff --git a/src/megatron/bridge/models/conversion/model_bridge.py b/src/megatron/bridge/models/conversion/model_bridge.py
index f3782eee6f..df241474b7 100644
--- a/src/megatron/bridge/models/conversion/model_bridge.py
+++ b/src/megatron/bridge/models/conversion/model_bridge.py
@@ -2018,6 +2018,7 @@ def stream_adapter_weights_megatron_to_hf(
     megatron_model: Union[MegatronModel, List[MegatronModel]],
     cpu: bool = True,
     show_progress: bool = True,
+    exclude_adapter_base_prefixes: Optional[Iterable[str]] = None,
 ) -> Iterable[HFWeightTuple]:
     """Bridge only adapter weights from Megatron to HuggingFace format."""
     ...
@@ -2109,12 +2110,14 @@ def _adapter_stream_registered_impl(
         megatron_model: Union[MegatronModel, List[MegatronModel]],
         cpu: bool = True,
         show_progress: bool = True,
+        exclude_adapter_base_prefixes: Optional[Iterable[str]] = None,
     ) -> Iterable[HFWeightTuple]:
         bridge = bridge_class()
         return bridge.stream_adapter_weights_megatron_to_hf(
             megatron_model,
             cpu=cpu,
             show_progress=show_progress,
+            exclude_adapter_base_prefixes=exclude_adapter_base_prefixes,
         )
 
     # Set meaningful names for debugging
diff --git a/src/megatron/bridge/models/conversion/peft_bridge.py b/src/megatron/bridge/models/conversion/peft_bridge.py
index 7a873972b1..d9bde27748 100644
--- a/src/megatron/bridge/models/conversion/peft_bridge.py
+++ b/src/megatron/bridge/models/conversion/peft_bridge.py
@@ -632,7 +632,9 @@ def _construct_adapters_names(self, prefix: str, adapter_key: Optional[str]) ->
         return linear_in_name, linear_out_name
 
     def build_adapter_conversion_tasks(
-        self, megatron_model: Union[MegatronModel, List[MegatronModel]]
+        self,
+        megatron_model: Union[MegatronModel, List[MegatronModel]],
+        exclude_adapter_base_prefixes: Iterable[str] | None = None,
     ) -> Dict[str, List[AdapterWeightConversionTask]]:
         """Construct adapter merge tasks keyed by their base parameter.
 
@@ -647,6 +649,7 @@ def build_adapter_conversion_tasks(
 
         adapters_info = self._megatron_global_adapters_info_all_pp_ranks(megatron_model)
         tasks_by_base: Dict[str, List[AdapterWeightConversionTask]] = defaultdict(list)  # type: ignore[name-defined]
+        excluded_prefixes = tuple(exclude_adapter_base_prefixes or ())
 
         from megatron.bridge.models.conversion.model_bridge import WeightConversionTask
 
@@ -667,6 +670,8 @@ def build_adapter_conversion_tasks(
         ) in adapters_info:
             # global_base_name example: decoder.layers.0.mlp.linear_fc1.adapter.adapter_q
             global_base_prefix, _, adapter_suffix = global_base_name.partition(".adapter")
+            if excluded_prefixes and global_base_prefix.startswith(excluded_prefixes):
+                continue
 
             adapter_key = None
             if adapter_suffix:
@@ -822,6 +827,7 @@ def stream_adapter_weights_megatron_to_hf(
         megatron_model: Union[MegatronModel, List[MegatronModel]],
         cpu: bool = True,
         show_progress: bool = True,
+        exclude_adapter_base_prefixes: Iterable[str] | None = None,
     ) -> Iterable["HFWeightTuple"]:
         """Stream only adapter weights without merging them into base tensors."""
 
@@ -832,7 +838,10 @@ def stream_adapter_weights_megatron_to_hf(
             megatron_model = [megatron_model]
 
         num_moe_experts = megatron_model[0].config.num_moe_experts
-        adapter_tasks_by_base = self.build_adapter_conversion_tasks(megatron_model)
+        adapter_tasks_by_base = self.build_adapter_conversion_tasks(
+            megatron_model,
+            exclude_adapter_base_prefixes=exclude_adapter_base_prefixes,
+        )
         adapter_tasks = list(itertools.chain.from_iterable(adapter_tasks_by_base.values()))
         if not adapter_tasks:
             return
diff --git a/src/megatron/bridge/peft/utils.py b/src/megatron/bridge/peft/utils.py
index b3b8329580..87227c7f8b 100644
--- a/src/megatron/bridge/peft/utils.py
+++ b/src/megatron/bridge/peft/utils.py
@@ -24,10 +24,10 @@
 import packaging
 import torch
 import torch.nn as nn
-from megatron.core import ModelParallelConfig, parallel_state
+from megatron.core import ModelParallelConfig, dist_checkpointing, parallel_state
 from megatron.core.dist_checkpointing.mapping import ShardedStateDict, ShardedTensor, ShardedTensorFactory
 from megatron.core.process_groups_config import ProcessGroupCollection
-from megatron.core.tensor_parallel import ColumnParallelLinear, RowParallelLinear
+from megatron.core.tensor_parallel import ColumnParallelLinear, RowParallelLinear, set_tensor_model_parallel_attributes
 from megatron.core.tensor_parallel.mappings import (
     gather_from_sequence_parallel_region,
     scatter_to_sequence_parallel_region,
@@ -35,6 +35,7 @@
 from megatron.core.transformer.mlp import apply_swiglu_sharded_factory
 from megatron.core.transformer.module import MegatronModule
 from megatron.core.transformer.moe.router import TopKRouter
+from megatron.core.utils import get_pg_rank, get_pg_size
 
 from megatron.bridge.utils.activation_map import str_to_dtype
 from megatron.bridge.utils.import_utils import safe_import_from
@@ -46,6 +47,8 @@
 ModelHook = Callable[[ModelList], ModelList | None]
 CheckpointPath = str | Path
 
+_LEGACY_SHARED_EXPERT_ADAPTER_CHECKPOINT_ATTR = "use_legacy_shared_expert_adapter_checkpoint"
+
 
 TEColumnParallelLinear, HAVE_TE_COL_LINEAR = safe_import_from(
     "megatron.core.extensions.transformer_engine", "TEColumnParallelLinear"
@@ -81,6 +84,234 @@
     )
 )
 
+
+def _get_pg_collection_from_module(module: object | None) -> ProcessGroupCollection | None:
+    """Return the process-group collection attached to a module or its config."""
+
+    for owner in (module, getattr(module, "config", None)):
+        if owner is None:
+            continue
+        for attr in ("pg_collection", "_pg_collection"):
+            pg_collection = getattr(owner, attr, None)
+            if pg_collection is not None:
+                return pg_collection
+    return None
+
+
+def _get_pg_collection(
+    pg_collection: ProcessGroupCollection | None = None,
+    source: object | None = None,
+    *,
+    required_pgs: List[str],
+) -> ProcessGroupCollection | None:
+    """Return the explicit PG collection or MCore's default collection fallback."""
+
+    pg_collection = pg_collection or _get_pg_collection_from_module(source)
+    if pg_collection is None:
+        # TODO: Once LoRA/DoRA transforms carry the model-level ProcessGroupCollection,
+        # pass it into adapter constructors explicitly and remove this default-MPU fallback.
+        pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=required_pgs)
+    return pg_collection
+
+
+def _iter_sharded_tensor_factories(state_dict: object) -> list[ShardedTensorFactory]:
+    """Return all sharded tensor factories in a nested state dict."""
+
+    if isinstance(state_dict, ShardedTensorFactory):
+        return [state_dict]
+    if isinstance(state_dict, Mapping):
+        factories = []
+        for value in state_dict.values():
+            factories.extend(_iter_sharded_tensor_factories(value))
+        return factories
+    if isinstance(state_dict, list | tuple):
+        factories = []
+        for value in state_dict:
+            factories.extend(_iter_sharded_tensor_factories(value))
+        return factories
+    return []
+
+
+def _checkpoint_tensor_shape(checkpoint_metadata: Mapping[str, ShardedTensor], key: str) -> tuple[int, ...] | None:
+    """Return checkpoint global tensor shape for a key, tolerating model-section prefixes."""
+
+    for candidate in (key, f"model.{key}"):
+        metadata = checkpoint_metadata.get(candidate)
+        if metadata is not None:
+            return tuple(metadata.global_shape)
+    return None
+
+
+def _legacy_shared_expert_adapter_key(factory: ShardedTensorFactory) -> str | None:
+    """Return the adapter module key if a factory represents a shared expert LoRA tensor."""
+
+    for suffix in (".linear_in.weight", ".linear_out.weight"):
+        if not factory.key.endswith(suffix):
+            continue
+        built = factory.build()
+        shards = built if isinstance(built, list) else [built]
+        if not shards or not isinstance(shards[0], ShardedTensor):
+            continue
+        expected_shape = tuple(shards[0].global_shape)
+        local_shape = tuple(factory.data.shape)
+        if len(expected_shape) == len(local_shape) + 1:
+            return factory.key[: -len(suffix)]
+    return None
+
+
+def _legacy_shared_expert_adapter_matches(
+    adapters_by_name: Mapping[str, "ParallelLinearAdapter"], adapter_key: str
+) -> list["ParallelLinearAdapter"]:
+    """Return adapter modules matching a legacy shared-expert checkpoint key."""
+
+    adapter = adapters_by_name.get(adapter_key)
+    if adapter is not None:
+        return [adapter]
+
+    adapter_base_key = adapter_key.removesuffix(".adapter")
+    matched_adapters = []
+    for module_name, module in adapters_by_name.items():
+        module_base_key = module_name.removesuffix(".adapter")
+        base_linear_name = module.base_linear_name
+        if (
+            adapter_key.endswith(module_name)
+            or module_name.endswith(adapter_key)
+            or adapter_base_key.endswith(module_base_key)
+            or module_base_key.endswith(adapter_base_key)
+            or adapter_base_key.endswith(base_linear_name)
+            or base_linear_name.endswith(adapter_base_key)
+        ):
+            matched_adapters.append(module)
+
+    if matched_adapters:
+        return matched_adapters
+
+    return list(adapters_by_name.values())
+
+
+def enable_legacy_shared_expert_adapter_loading(
+    megatron_model: list[nn.Module] | nn.Module,
+    sharded_state_dict: ShardedStateDict,
+    checkpoint_path: str | Path,
+) -> bool:
+    """Enable legacy 2D checkpoint loading for old shared grouped-expert adapters.
+
+    New shared grouped-expert LoRA checkpoints expose a leading global expert axis
+    so they can be resharded across EP changes. Older checkpoints saved the same
+    shared adapter as a plain 2D tensor. This helper detects that old metadata
+    shape and marks only the matching shared adapter modules to emit the legacy
+    2D sharded state dict for loading.
+
+    Args:
+        megatron_model: Model module or model chunks containing PEFT adapters.
+        sharded_state_dict: Current adapter-only sharded state dict.
+        checkpoint_path: Distributed checkpoint directory to inspect.
+
+    Returns:
+        True if at least one shared expert adapter was marked for legacy loading.
+    """
+
+    checkpoint_metadata = dist_checkpointing.load_tensors_metadata(str(checkpoint_path))
+    models = megatron_model if isinstance(megatron_model, list) else [megatron_model]
+    adapters_by_name: dict[str, ParallelLinearAdapter] = {}
+    for model in models:
+        for name, module in model.named_modules():
+            if isinstance(module, ParallelLinearAdapter) and module._uses_grouped_expert_sharding():
+                adapters_by_name[name.removeprefix("module.")] = module
+
+    enabled = False
+    for factory in _iter_sharded_tensor_factories(sharded_state_dict):
+        adapter_key = _legacy_shared_expert_adapter_key(factory)
+        if adapter_key is None:
+            continue
+        built = factory.build()
+        shards = built if isinstance(built, list) else [built]
+        expected_shape = tuple(shards[0].global_shape)
+        legacy_shape = expected_shape[1:]
+        if _checkpoint_tensor_shape(checkpoint_metadata, factory.key) == legacy_shape:
+            for adapter in _legacy_shared_expert_adapter_matches(adapters_by_name, adapter_key):
+                setattr(adapter, _LEGACY_SHARED_EXPERT_ADAPTER_CHECKPOINT_ATTR, True)
+                enabled = True
+
+    return enabled
+
+
+def _get_process_group(pg_collection: ProcessGroupCollection | None, *names: str) -> object | None:
+    """Return the first named process group available on a collection."""
+
+    if pg_collection is None:
+        return None
+    for name in names:
+        group = getattr(pg_collection, name, None)
+        if group is not None:
+            return group
+    return None
+
+
+def _process_group_size(group: object | None, fallback: int = 1) -> int:
+    """Return a process-group size without consulting global parallel state."""
+
+    if group is None:
+        return int(fallback or 1)
+    size = None
+    size_attr = getattr(group, "size", None)
+    try:
+        size = size_attr() if callable(size_attr) else size_attr
+    except (RuntimeError, ValueError, TypeError):
+        size = None
+    if size is None:
+        try:
+            size = get_pg_size(group)
+        except (RuntimeError, ValueError, TypeError):
+            size = None
+    return int(size if size is not None else fallback or 1)
+
+
+def _process_group_rank(group: object | None, fallback: int = 0) -> int:
+    """Return this rank within a process group without consulting global parallel state."""
+
+    if group is None:
+        return int(fallback or 0)
+    rank = None
+    rank_attr = getattr(group, "rank", None)
+    try:
+        rank = rank_attr() if callable(rank_attr) else rank_attr
+    except (RuntimeError, ValueError, TypeError):
+        rank = None
+    if rank is None:
+        try:
+            rank = get_pg_rank(group)
+        except (RuntimeError, ValueError, TypeError):
+            rank = None
+    return int(rank if rank is not None else fallback or 0)
+
+
+def _get_tensor_parallel_group(
+    pg_collection: ProcessGroupCollection | None, *, is_expert: bool = False
+) -> object | None:
+    """Return the tensor-parallel group for dense or expert linear layers."""
+
+    if is_expert:
+        return _get_process_group(pg_collection, "expt_tp", "etp")
+    return _get_process_group(pg_collection, "tp")
+
+
+def _get_tensor_parallel_group_from_module(
+    module: nn.Module, *, is_expert: bool = False, pg_collection: ProcessGroupCollection | None = None
+) -> object | None:
+    """Return the TP group passed to the wrapped module, falling back to its collection."""
+
+    pg_collection = _get_pg_collection(
+        pg_collection,
+        module,
+        required_pgs=["expt_tp"] if is_expert else ["tp"],
+    )
+    group = _get_tensor_parallel_group(pg_collection, is_expert=is_expert)
+    if group is not None:
+        return group
+    return getattr(module, "_tp_group", None) or getattr(module, "tp_group", None)
+
+
 MixedFusedLayerNorm, HAVE_APEX = safe_import_from("apex.normalization.fused_layer_norm", "MixedFusedLayerNorm")
 ModelOptLinear, HAVE_MODELOPT_LINEAR = safe_import_from("megatron.core.post_training.modelopt.layers", "Linear")
 
@@ -158,11 +389,14 @@ def load_peft_adapter_checkpoint(
     checkpoint_path = str(adapter_checkpoint_path)
     if load_strategy is None:
         load_strategy = get_default_load_sharded_strategy(checkpoint_path)
-        if fully_parallel_load and parallel_state.is_initialized():
-            load_strategy = FullyParallelLoadStrategyWrapper(
-                load_strategy,
-                parallel_state.get_data_parallel_group(with_context_parallel=True),
-            )
+        if pg_collection is None and fully_parallel_load:
+            try:
+                pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=["dp_cp"])
+            except AssertionError:
+                pg_collection = None
+        dp_cp_group = _get_process_group(pg_collection, "dp_cp")
+        if fully_parallel_load and dp_cp_group is not None:
+            load_strategy = FullyParallelLoadStrategyWrapper(load_strategy, dp_cp_group)
 
     loaded_state_dict = dist_checkpointing.load(sharded_state_dict, checkpoint_path, load_strategy)
     for vpp_rank, model_chunk in enumerate(model_chunks):
@@ -256,7 +490,11 @@ class AdapterAttributes:
     base_linear_is_parallel: bool
 
 
-def get_adapter_attributes_from_linear(m: nn.Module, is_expert: bool = False) -> AdapterAttributes:
+def get_adapter_attributes_from_linear(
+    m: nn.Module,
+    is_expert: bool = False,
+    pg_collection: ProcessGroupCollection | None = None,
+) -> AdapterAttributes:
     """Returns attributes from the base layer as an AdapterAttributes dataclass.
 
     input_is_parallel, in_features, out_features, disable_tensor_parallel_comm,
@@ -302,10 +540,15 @@ def get_adapter_attributes_from_linear(m: nn.Module, is_expert: bool = False) ->
             base_linear_is_parallel=False,
         )
 
-    if is_expert:
-        tp_size = parallel_state.get_expert_tensor_parallel_world_size()
-    else:
-        tp_size = parallel_state.get_tensor_model_parallel_world_size()
+    tp_group = _get_tensor_parallel_group_from_module(m, is_expert=is_expert, pg_collection=pg_collection)
+    tp_size = _process_group_size(
+        tp_group,
+        getattr(
+            m.config,
+            "expert_tensor_parallel_size" if is_expert else "tensor_model_parallel_size",
+            1,
+        ),
+    )
     if isinstance(m, TopKRouter):
         input_is_parallel = False
         in_features = m.weight.shape[1]
@@ -422,15 +665,15 @@ def align_expert_dim_for_tp(
     normalize_moe_lora: bool,
     is_expert: bool,
     input_is_parallel: bool,
+    pg_collection: ProcessGroupCollection | None = None,
 ) -> int:
     """Round normalized expert LoRA ranks up to the expert-TP granularity when needed."""
 
     if not normalize_moe_lora or not is_expert or input_is_parallel:
         return dim
 
-    expert_tp_size = (
-        parallel_state.get_expert_tensor_parallel_world_size() or module.config.expert_tensor_parallel_size or 1
-    )
+    expert_tp_group = _get_tensor_parallel_group_from_module(module, is_expert=True, pg_collection=pg_collection)
+    expert_tp_size = _process_group_size(expert_tp_group, module.config.expert_tensor_parallel_size or 1)
     if expert_tp_size <= 1 or dim % expert_tp_size == 0:
         return dim
 
@@ -561,7 +804,7 @@ class _All2AllHp2Sp(torch.autograd.Function):
     """
 
     @staticmethod
-    def forward(ctx, input_: torch.Tensor) -> torch.Tensor:
+    def forward(ctx, input_: torch.Tensor, group: object | None) -> torch.Tensor:
         """Forward pass: All-to-All from Hidden Parallel to Sequence Parallel.
 
         Args:
@@ -571,8 +814,8 @@ def forward(ctx, input_: torch.Tensor) -> torch.Tensor:
         Returns:
             Output tensor in sequence parallel layout.
         """
-        world_size = parallel_state.get_tensor_model_parallel_world_size()
-        group = parallel_state.get_tensor_model_parallel_group()
+        ctx.group = group
+        world_size = _process_group_size(group)
         send_list = list(input_.chunk(world_size, dim=0))
         send_list = [tensor.contiguous() for tensor in send_list]
         receive_list = [torch.empty_like(send_list[0]) for _ in range(world_size)]
@@ -592,18 +835,18 @@ def backward(ctx, grad_output: torch.Tensor) -> torch.Tensor:
         Returns:
             Gradient tensor in hidden parallel layout.
         """
-        world_size = parallel_state.get_tensor_model_parallel_world_size()
-        group = parallel_state.get_tensor_model_parallel_group()
+        group = ctx.group
+        world_size = _process_group_size(group)
         send_list = list(grad_output.chunk(world_size, dim=-1))
         send_list = [tensor.contiguous() for tensor in send_list]
         receive_list = [torch.empty_like(send_list[0]) for _ in range(world_size)]
         torch.distributed.all_to_all(receive_list, send_list, group=group)
         x = torch.cat(receive_list, dim=0)
 
-        return x
+        return x, None
 
 
-def all2all_hp2sp(input_: torch.Tensor) -> torch.Tensor:
+def all2all_hp2sp(input_: torch.Tensor, tensor_parallel_group: object | None = None) -> torch.Tensor:
     """Perform All-to-All communication from Hidden Parallel to Sequence Parallel.
 
     Args:
@@ -612,7 +855,7 @@ def all2all_hp2sp(input_: torch.Tensor) -> torch.Tensor:
     Returns:
         Output tensor in sequence parallel layout.
     """
-    return _All2AllHp2Sp.apply(input_)
+    return _All2AllHp2Sp.apply(input_, tensor_parallel_group)
 
 
 class ParallelLinearAdapter(nn.Module):
@@ -663,6 +906,7 @@ def __init__(
         disable_tensor_parallel_comm: bool = False,
         disable_sequence_parallel_comm: bool = True,
         base_linear_is_parallel: bool = True,
+        pg_collection: ProcessGroupCollection | None = None,
     ) -> None:
         """Initialize the ParallelLinearAdapter.
 
@@ -695,11 +939,22 @@ def __init__(
         self.use_a2a = a2a_experimental
         self.is_expert = is_expert
         self.base_linear_is_parallel = base_linear_is_parallel
+        self.use_legacy_shared_expert_adapter_checkpoint = False
 
         # megatron_gpt_peft_models will provide this arg, but deprecated ones do not.
         # in case this arg is not provided, use the dummy default config.
         if model_parallel_config is None:
             model_parallel_config = ModelParallelConfig()
+        # TODO: When the PEFT transform API has explicit PG plumbing, pass the
+        # model-level collection here instead of relying on config/default discovery.
+        self.pg_collection = _get_pg_collection(
+            pg_collection,
+            model_parallel_config,
+            required_pgs=["ep", "expt_tp", "expt_dp"] if is_expert else ["tp"],
+        )
+        self.tp_group = _get_tensor_parallel_group(self.pg_collection, is_expert=is_expert)
+        self.ep_group = _get_process_group(self.pg_collection, "ep")
+        self.expert_dp_group = _get_process_group(self.pg_collection, "expt_dp")
         _sequence_parallel = model_parallel_config.sequence_parallel
         model_parallel_config.sequence_parallel = False  # SP is irrelevant for the lora linear layer
         self.config = model_parallel_config
@@ -718,6 +973,7 @@ def __init__(
                 bias=False,
                 init_method=self._get_init_fn(column_init_method),
                 is_expert=is_expert,
+                tp_group=self.tp_group,
             )
         else:
             self.linear_in = ColumnParallelLinear(
@@ -729,6 +985,7 @@ def __init__(
                 init_method=self._get_init_fn(column_init_method),
                 disable_grad_reduce=_sequence_parallel,
                 is_expert=is_expert,
+                tp_group=self.tp_group,
             )
 
         # (@adithyare) we use this option to mirror the behavior
@@ -755,6 +1012,7 @@ def __init__(
             gather_output=lin_out_gather_output,
             init_method=self._get_init_fn(row_init_method),
             is_expert=is_expert,
+            tp_group=self.tp_group,
         )
 
         if dropout > 0.0:
@@ -768,6 +1026,9 @@ def __init__(
         elif model_parallel_config.fp16:
             self.half()
 
+        if self._uses_grouped_expert_sharding():
+            self._register_shared_expert_grad_sync_hooks()
+
         # revert config change in case it is read elsewhere
         model_parallel_config.sequence_parallel = _sequence_parallel
         self.disable_sequence_parallel_comm = disable_sequence_parallel_comm
@@ -850,7 +1111,7 @@ def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:
             # layernorm before lora is impacted by sequence parallel,
             # hence seq dim need to be gathered right before lora linear layers
             # this function also handles the backward pass correctly
-            x = gather_from_sequence_parallel_region(x)
+            x = gather_from_sequence_parallel_region(x, group=self.tp_group)
 
         if self.config.cpu_offloading and self.config.cpu_offloading_activations:
             x.activation_offloading = True
@@ -869,9 +1130,9 @@ def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:
             # this function also handles the backward pass correctly
             if self.use_a2a:
                 # all2all hidden_size / TP to seq_len / TP
-                x = all2all_hp2sp(x)
+                x = all2all_hp2sp(x, self.tp_group)
             else:
-                x = scatter_to_sequence_parallel_region(x)
+                x = scatter_to_sequence_parallel_region(x, group=self.tp_group)
 
         # Add dropout if available
         if self.dropout_position == "post":
@@ -885,6 +1146,207 @@ def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:
 
         return x
 
+    def local_experts_per_rank(self) -> int:
+        """Return the number of global expert slots owned by this EP rank."""
+
+        ep_size = _process_group_size(self.ep_group, self.config.expert_model_parallel_size or 1)
+        num_global_experts = getattr(self.config, "num_moe_experts", None)
+        if num_global_experts is None:
+            return 1
+        if int(num_global_experts) % int(ep_size) != 0:
+            raise ValueError(
+                f"num_moe_experts={num_global_experts} must be divisible by expert_model_parallel_size={ep_size}"
+            )
+        return int(num_global_experts) // int(ep_size)
+
+    def _uses_grouped_expert_sharding(self) -> bool:
+        """Return whether this shared adapter needs an explicit expert axis."""
+
+        return self.is_expert and is_grouped_expert_linear(self.base_linear_name)
+
+    def _allreduce_shared_expert_grad(self, grad: torch.Tensor) -> torch.Tensor:
+        """Sum shared expert adapter grads across EP before expert-DP reduction."""
+
+        if not torch.distributed.is_available() or not torch.distributed.is_initialized():
+            return grad
+        if self.ep_group is None or _process_group_size(self.ep_group) <= 1:
+            return grad
+        # Sum across EP first; MCore expert DDP then reduces across expert-DP
+        # and scales expert buffers by 1 / dp_cp_group.size(), i.e. the full
+        # EP x expert-DP data-parallel world, not just expert-DP.
+        torch.distributed.all_reduce(grad, group=self.ep_group)
+        return grad
+
+    def _register_shared_expert_grad_sync_hooks(self) -> None:
+        """Keep shared grouped-expert adapters synchronized across EP ranks."""
+
+        for module in (self.linear_in, self.linear_out):
+            weight = getattr(module, "weight", None)
+            if isinstance(weight, torch.Tensor) and weight.requires_grad:
+                weight.register_hook(self._allreduce_shared_expert_grad)
+
+    def _expert_axis_info(self, sharded_offsets: Tuple) -> Tuple[int, int, int]:
+        """Return the global expert-axis sharding metadata for this rank."""
+
+        ep_size = _process_group_size(self.ep_group, self.config.expert_model_parallel_size or 1)
+        ep_rank = _process_group_rank(self.ep_group)
+        local_experts = self.local_experts_per_rank()
+        if local_experts <= 0:
+            raise ValueError(f"local_experts_per_rank must be positive, got {local_experts}")
+
+        num_global_experts = getattr(self.config, "num_moe_experts", None)
+        if num_global_experts is None:
+            num_global_experts = ep_size * local_experts
+        num_global_experts = int(num_global_experts)
+        first_expert_slot = ep_rank * local_experts
+        if first_expert_slot >= num_global_experts:
+            raise ValueError(
+                f"Invalid expert adapter sharding for {self.base_linear_name}: "
+                f"ep_rank={ep_rank}, local_experts_per_rank={local_experts}, "
+                f"num_global_experts={num_global_experts}"
+            )
+
+        expert_axis = len(sharded_offsets)
+        return expert_axis, first_expert_slot, num_global_experts
+
+    def _keep_expert_extra_state(self) -> bool:
+        """Keep one unsharded adapter extra-state entry."""
+
+        tp_rank = _process_group_rank(self.tp_group)
+        ep_rank = _process_group_rank(self.ep_group)
+        return tp_rank == 0 and ep_rank == 0
+
+    def _set_expert_replica_ids(self, *state_dicts: ShardedStateDict) -> None:
+        """Mark expert adapter replicas across expert data-parallel ranks."""
+
+        edp_rank = _process_group_rank(self.expert_dp_group)
+        for state_dict in state_dicts:
+            for value in state_dict.values():
+                if not hasattr(value, "replica_id"):
+                    continue
+                replica_id = value.replica_id
+                if isinstance(replica_id, int):
+                    replica_id = (0, 0, replica_id)
+                if len(replica_id) != 3:
+                    raise ValueError(
+                        f"Expected replica_id for {self.base_linear_name} to be in "
+                        f"(PP, TP, DP) format, got: {replica_id}"
+                    )
+                dp_replica_id = 0 if getattr(value, "is_data_parallel_fully_shard", False) else edp_rank
+                value.replica_id = (*replica_id[:2], dp_replica_id)
+
+    def _apply_expert_axis_factory(
+        self,
+        sharded_tensor: ShardedTensor,
+        sharded_offsets: Tuple,
+        *,
+        split_swiglu: bool = False,
+    ) -> ShardedTensorFactory:
+        """Map one shared 2D adapter tensor to this rank's global expert slots."""
+
+        expert_axis, first_expert_slot, num_global_experts = self._expert_axis_info(sharded_offsets)
+        local_experts = self.local_experts_per_rank()
+        base_prepend_axis_num = len(sharded_offsets)
+        output_prepend_axis_num = base_prepend_axis_num + 1
+        swiglu_shard_axis = 0
+
+        preserved_rank_offsets = []
+        for axis, local_axis_shape in enumerate(sharded_tensor.local_shape):
+            base_global_axis_idx = axis + base_prepend_axis_num
+            output_global_axis_idx = base_global_axis_idx + 1
+            axis_fragments = sharded_tensor.axis_fragmentations[base_global_axis_idx]
+            if axis_fragments <= 1:
+                continue
+            global_offset = sharded_tensor.global_offset[base_global_axis_idx]
+            if global_offset % local_axis_shape != 0:
+                raise ValueError(
+                    f"Cannot preserve non-integral sharding for {sharded_tensor.key}: "
+                    f"offset={global_offset}, local_axis_shape={local_axis_shape}"
+                )
+            preserved_rank_offsets.append((output_global_axis_idx, global_offset // local_axis_shape, axis_fragments))
+
+        swiglu_axis_frag = None
+        swiglu_rank_offset = None
+        base_swiglu_global_axis = swiglu_shard_axis + base_prepend_axis_num
+        output_swiglu_global_axis = swiglu_shard_axis + output_prepend_axis_num
+        if split_swiglu:
+            local_axis_size = sharded_tensor.local_shape[swiglu_shard_axis]
+            if sharded_tensor.global_offset[base_swiglu_global_axis] % local_axis_size != 0:
+                raise ValueError(
+                    f"Cannot split SwiGLU tensor {sharded_tensor.key}: "
+                    f"offset={sharded_tensor.global_offset[base_swiglu_global_axis]}, local_axis_size={local_axis_size}"
+                )
+            swiglu_rank_offset = sharded_tensor.global_offset[base_swiglu_global_axis] // local_axis_size
+            swiglu_axis_frag = sharded_tensor.axis_fragmentations[base_swiglu_global_axis]
+            preserved_rank_offsets = [
+                rank_offset for rank_offset in preserved_rank_offsets if rank_offset[0] != output_swiglu_global_axis
+            ]
+
+        @torch.no_grad()
+        def sh_ten_build_fn(key: str, tensor: torch.Tensor, replica_id, flattened_range):
+            del flattened_range
+            sharded_tensors = []
+            for expert_index in range(local_experts):
+                expert_offset = (expert_axis, first_expert_slot + expert_index, num_global_experts)
+                if not split_swiglu:
+                    sharded_tensors.append(
+                        ShardedTensor.from_rank_offsets(
+                            key,
+                            tensor,
+                            *sharded_offsets,
+                            *preserved_rank_offsets,
+                            expert_offset,
+                            replica_id=replica_id,
+                            prepend_axis_num=output_prepend_axis_num,
+                        )
+                    )
+                    continue
+
+                tensor_w, tensor_v = torch.chunk(tensor, 2, dim=swiglu_shard_axis)
+                offset_w = (output_swiglu_global_axis, swiglu_rank_offset, swiglu_axis_frag * 2)
+                offset_v = (
+                    output_swiglu_global_axis,
+                    swiglu_rank_offset + swiglu_axis_frag,
+                    swiglu_axis_frag * 2,
+                )
+                for tensor_part, swiglu_offset in ((tensor_w, offset_w), (tensor_v, offset_v)):
+                    sharded_tensors.append(
+                        ShardedTensor.from_rank_offsets(
+                            key,
+                            tensor_part,
+                            *sharded_offsets,
+                            *preserved_rank_offsets,
+                            expert_offset,
+                            swiglu_offset,
+                            replica_id=replica_id,
+                            prepend_axis_num=output_prepend_axis_num,
+                        )
+                    )
+            return sharded_tensors
+
+        def sh_ten_merge_fn(sub_state_dict):
+            if not isinstance(sub_state_dict, list):
+                sub_state_dict = [sub_state_dict]
+            if split_swiglu:
+                if len(sub_state_dict) % 2 != 0:
+                    raise ValueError(f"Expected even number of SwiGLU shards for {sharded_tensor.key}")
+                sub_state_dict = [
+                    torch.cat(sub_state_dict[index : index + 2], dim=swiglu_shard_axis)
+                    for index in range(0, len(sub_state_dict), 2)
+                ]
+            if len(sub_state_dict) == 1:
+                return sub_state_dict[0]
+            return torch.stack(sub_state_dict, dim=0).mean(dim=0)
+
+        return ShardedTensorFactory(
+            sharded_tensor.key,
+            sharded_tensor.data,
+            sh_ten_build_fn,
+            sh_ten_merge_fn,
+            sharded_tensor.replica_id,
+            flattened_range=sharded_tensor.flattened_range,
+        )
+
     def sharded_state_dict(
         self,
         prefix: str = "",
@@ -906,40 +1368,38 @@ def sharded_state_dict(
             Sharded state dictionary for distributed checkpointing.
         """
         sharded_state_dict = {}
+        # Shared grouped-expert adapters have one 2D weight per EP rank, but the
+        # checkpoint must expose the global expert axis so EP changes can reshard it.
+        # Non-grouped expert adapters already sit under .local_experts.* and keep
+        # their existing expert-DP replica metadata instead.
+        use_expert_axis = self._uses_grouped_expert_sharding() and not self.use_legacy_shared_expert_adapter_checkpoint
+        split_swiglu = "linear_fc1" in self.base_linear_name and getattr(self.config, "gated_linear_unit", False)
         linear_in_sd = self.linear_in.sharded_state_dict(f"{prefix}linear_in.", sharded_offsets, metadata)
         linear_out_sd = self.linear_out.sharded_state_dict(f"{prefix}linear_out.", sharded_offsets, metadata)
 
-        # The experts.py code in Megatron-LM set replica_id = (PP, ETP, EDP),
-        # but it will cause errors as mentioned in https://github.com/volcengine/verl/issues/4303,
-        # since adapter weights are not EP sharded and it assumes that it will
-        # replicate along DP modulo EP (sharded by EP)
-        if self.is_expert:
-            from megatron.core import parallel_state
-
-            ep_rank = parallel_state.get_expert_model_parallel_rank()
-            edp_rank = parallel_state.get_expert_data_parallel_rank()
-            dp_size = parallel_state.get_data_parallel_world_size()
-            # TODO: This modification logic is in question and needs further verification.
-            rank = (ep_rank + 1) * (edp_rank + 1) - 1 if dp_size == 1 else ep_rank
-            for sd in [linear_in_sd, linear_out_sd]:
-                for v in sd.values():
-                    if hasattr(v, "replica_id"):
-                        old_rid = v.replica_id
-                        v.replica_id = (old_rid[0], rank, old_rid[2])
-
-            # TE FP8 _extra_state is a ShardedObject (no replica_id) that is
-            # identical across all TP ranks in the same EP group — experts are
-            # EP-sharded, not TP-sharded, so every TP rank sees the same experts
-            # and produces the same shard_X_N keys.  Keep it only on TP rank 0
-            # to avoid CheckpointingException: Duplicate ShardedObject keys.
-            tp_rank = parallel_state.get_tensor_model_parallel_rank()
-            if tp_rank > 0:
-                for sd in [linear_in_sd, linear_out_sd]:
-                    extra_state_keys = [k for k in sd if "_extra_state" in k]
-                    for k in extra_state_keys:
-                        del sd[k]
-
-        if "linear_fc1" in self.base_linear_name:
+        if use_expert_axis:
+            keep_extra_state = self._keep_expert_extra_state()
+            if not keep_extra_state:
+                for sd in (linear_in_sd, linear_out_sd):
+                    for key in [k for k in sd if "_extra_state" in k]:
+                        del sd[key]
+            for key, value in list(linear_in_sd.items()):
+                if isinstance(value, ShardedTensor):
+                    linear_in_sd[key] = self._apply_expert_axis_factory(value, sharded_offsets)
+            for key, value in list(linear_out_sd.items()):
+                if isinstance(value, ShardedTensor):
+                    linear_out_sd[key] = self._apply_expert_axis_factory(
+                        value,
+                        sharded_offsets,
+                        split_swiglu=split_swiglu,
+                    )
+        elif self.is_expert:
+            if _process_group_rank(self.tp_group) > 0:
+                for state_dict in (linear_in_sd, linear_out_sd):
+                    for key in [k for k in state_dict if "_extra_state" in k]:
+                        del state_dict[key]
+
+        if split_swiglu and not use_expert_axis:
             for k, v in linear_out_sd.items():
                 if k in (f"{prefix}linear_out.weight", f"{prefix}linear_out.bias"):
                     linear_out_sd[k] = apply_swiglu_sharded_factory(v, sharded_offsets)
@@ -973,6 +1433,9 @@ def sharded_state_dict(
                             0,  # split along dimension 0
                         )
 
+        if self.is_expert:
+            self._set_expert_replica_ids(linear_in_sd, linear_out_sd)
+
         sharded_state_dict.update(linear_in_sd)
         sharded_state_dict.update(linear_out_sd)
         return sharded_state_dict
@@ -1118,6 +1581,9 @@ def _make_grouped_expert_sharded_tensor(
     *,
     tp_axis: Optional[int],
     sharded_offsets: Tuple,
+    pg_collection: ProcessGroupCollection | None,
+    ep_size_fallback: int = 1,
+    etp_size_fallback: int = 1,
 ) -> ShardedTensor:
     """Build a sharded tensor for packed grouped-expert weights.
 
@@ -1128,32 +1594,52 @@ def _make_grouped_expert_sharded_tensor(
     prepend_axis_num = len(sharded_offsets)
     rank_offsets = list(sharded_offsets)
 
-    ep_size = parallel_state.get_expert_model_parallel_world_size() or 1
+    ep_group = _get_process_group(pg_collection, "ep")
+    ep_size = _process_group_size(ep_group, ep_size_fallback)
     _append_rank_offset(
         rank_offsets,
         prepend_axis_num,
-        parallel_state.get_expert_model_parallel_rank() or 0,
+        _process_group_rank(ep_group),
         ep_size,
     )
 
     if tp_axis is not None:
-        etp_size = parallel_state.get_expert_tensor_parallel_world_size() or 1
+        etp_group = _get_tensor_parallel_group(pg_collection, is_expert=True)
+        etp_size = _process_group_size(etp_group, etp_size_fallback)
         _append_rank_offset(
             rank_offsets,
             prepend_axis_num + tp_axis,
-            parallel_state.get_expert_tensor_parallel_rank() or 0,
+            _process_group_rank(etp_group),
             etp_size,
         )
 
+    expt_dp_group = _get_process_group(pg_collection, "expt_dp")
     return ShardedTensor.from_rank_offsets(
         key,
         tensor,
         *rank_offsets,
-        replica_id=(0, 0, parallel_state.get_expert_data_parallel_rank() or 0),
+        replica_id=(0, 0, _process_group_rank(expt_dp_group)),
         prepend_axis_num=prepend_axis_num,
     )
 
 
+class _GroupedExpertAdapterWeight(nn.Module):
+    """Callable parameter container so DDP forward pre-hooks see grouped LoRA weights."""
+
+    # Overlapped param gather is driven by module forward pre-hooks. Calling this
+    # container before reading the weight makes expert-DP LoRA params participate
+    # in the normal training-time gather instead of only forced eval/checkpoint sync.
+
+    def __init__(self, weight: torch.Tensor) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(weight)
+
+    def forward(self, indices: Optional[List[int]] = None) -> torch.Tensor:
+        if indices is None:
+            return self.weight
+        return self.weight[indices]
+
+
 class GroupedExpertLinearAdapter(nn.Module):
     """LoRA adapter with one low-rank pair per local grouped MoE expert."""
 
@@ -1176,6 +1662,7 @@ def __init__(
         base_linear_is_parallel: bool = True,
         params_device: Optional[torch.device] = None,
         params_dtype: Optional[torch.dtype] = None,
+        pg_collection: ProcessGroupCollection | None = None,
     ) -> None:
         """Initialize grouped-expert LoRA weights for one adapter per local expert."""
 
@@ -1197,11 +1684,22 @@ def __init__(
         if model_parallel_config is None:
             model_parallel_config = ModelParallelConfig()
         self.config = model_parallel_config
+        # TODO: When the PEFT transform API has explicit PG plumbing, pass the
+        # model-level collection here instead of relying on config/default discovery.
+        self.pg_collection = _get_pg_collection(
+            pg_collection,
+            model_parallel_config,
+            required_pgs=["ep", "expt_tp", "expt_dp"],
+        )
+        self.expert_tp_group = _get_tensor_parallel_group(self.pg_collection, is_expert=True)
+        self.ep_group = _get_process_group(self.pg_collection, "ep")
+        self.expert_dp_group = _get_process_group(self.pg_collection, "expt_dp")
 
         model_parallel_config.perform_initialization = True
 
-        expert_tp_size = (
-            parallel_state.get_expert_tensor_parallel_world_size() or model_parallel_config.expert_tensor_parallel_size
+        expert_tp_size = _process_group_size(
+            self.expert_tp_group,
+            model_parallel_config.expert_tensor_parallel_size or 1,
         )
         linear_in_tp_axis = 2 if input_is_parallel else 1
         linear_out_tp_axis = 1
@@ -1225,11 +1723,12 @@ def __init__(
         )
 
         if params_device is None:
+            distributed_initialized = torch.distributed.is_available() and torch.distributed.is_initialized()
             params_device = (
                 torch.device("cpu")
                 if model_parallel_config.use_cpu_initialization
                 or not torch.cuda.is_available()
-                or not parallel_state.is_initialized()
+                or not distributed_initialized
                 else torch.device("cuda", torch.cuda.current_device())
             )
         dtype = params_dtype or model_parallel_config.params_dtype
@@ -1240,22 +1739,23 @@ def __init__(
         ParallelLinearAdapter._get_init_fn(self, row_init_method)(linear_out_weight)
 
         expert_parallel = (
-            parallel_state.get_expert_model_parallel_world_size() or model_parallel_config.expert_model_parallel_size
-        ) > 1
+            _process_group_size(
+                self.ep_group,
+                model_parallel_config.expert_model_parallel_size or 1,
+            )
+            > 1
+        )
         self._linear_in_tp_axis = linear_in_tp_axis
         self._linear_out_tp_axis = linear_out_tp_axis
-        self.linear_in = nn.Module()
-        self.linear_in.weight = nn.Parameter(linear_in_weight)
-        self.linear_out = nn.Module()
-        self.linear_out.weight = nn.Parameter(linear_out_weight)
+        self.linear_in = _GroupedExpertAdapterWeight(linear_in_weight)
+        self.linear_out = _GroupedExpertAdapterWeight(linear_out_weight)
         for weight, tp_axis in (
             (self.linear_in.weight, linear_in_tp_axis),
             (self.linear_out.weight, linear_out_tp_axis),
         ):
             setattr(weight, "allreduce", not expert_parallel)
             if tp_axis is not None:
-                setattr(weight, "partition_dim", tp_axis)
-                setattr(weight, "partition_stride", 1)
+                set_tensor_model_parallel_attributes(weight, True, tp_axis, 1)
 
         if dropout > 0.0:
             self.dropout = nn.Dropout(dropout)
@@ -1287,12 +1787,10 @@ def _extract_expert_splits(self, args: Tuple, kwargs: Dict) -> List[int]:
     def _gather_along_last_dim(self, tensor: torch.Tensor) -> torch.Tensor:
         """Gather a tensor across expert TP ranks by concatenating its last dimension."""
 
-        expert_tp_size = (
-            parallel_state.get_expert_tensor_parallel_world_size() or self.config.expert_tensor_parallel_size
-        )
+        expert_tp_size = _process_group_size(self.expert_tp_group, self.config.expert_tensor_parallel_size or 1)
         if expert_tp_size == 1:
             return tensor
-        expert_tp_group = parallel_state.get_expert_tensor_parallel_group(check_initialized=False)
+        expert_tp_group = self.expert_tp_group
         if expert_tp_group is None:
             raise ValueError(
                 f"{self.base_linear_name} requires initialized expert tensor parallel state "
@@ -1477,6 +1975,8 @@ def _forward_per_expert(
     ) -> torch.Tensor:
         """Apply the adapter using the per-expert fallback path."""
 
+        linear_in_weight = self.linear_in()
+        linear_out_weight = self.linear_out()
         outputs = []
         start = 0
         for expert_idx, split_size in enumerate(expert_splits):
@@ -1489,14 +1989,14 @@ def _forward_per_expert(
                 if self.config.cpu_offloading and self.config.cpu_offloading_activations:
                     expert_input.activation_offloading = True
 
-            hidden = nn.functional.linear(expert_input, self.linear_in.weight[expert_idx])
+            hidden = nn.functional.linear(expert_input, linear_in_weight[expert_idx])
             if not self.input_is_parallel:
                 hidden = self._gather_along_last_dim(hidden)
             hidden = self.activation(hidden)
 
             if self.config.cpu_offloading and self.config.cpu_offloading_activations:
                 hidden.activation_offloading = True
-            expert_output = nn.functional.linear(hidden, self.linear_out.weight[expert_idx])
+            expert_output = nn.functional.linear(hidden, linear_out_weight[expert_idx])
             if self.input_is_parallel:
                 expert_output = self._gather_along_last_dim(expert_output)
 
@@ -1523,14 +2023,15 @@ def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:
         if self.dropout_position == "pre":
             x = self.dropout(x)
 
-        expert_tp_size = (
-            parallel_state.get_expert_tensor_parallel_world_size() or self.config.expert_tensor_parallel_size
-        )
+        expert_tp_size = _process_group_size(self.expert_tp_group, self.config.expert_tensor_parallel_size or 1)
         output_features = self.linear_out.weight.shape[1]
         if self.input_is_parallel:
             output_features *= expert_tp_size
         if x.shape[0] == 0:
-            return x.new_empty((0, output_features)) * (self.alpha / self.dim)
+            linear_in_weight = self.linear_in()
+            linear_out_weight = self.linear_out()
+            grad_anchor = linear_in_weight.reshape(-1)[0] + linear_out_weight.reshape(-1)[0]
+            return (x.new_empty((0, output_features)) + grad_anchor * 0.0) * (self.alpha / self.dim)
 
         if not use_te_grouped_linear and not self._can_use_grouped_mm(x):
             return self._forward_per_expert(x, expert_splits=expert_splits, expert_tp_size=expert_tp_size) * (
@@ -1561,7 +2062,7 @@ def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:
         if not use_te_grouped_linear:
             offs = self._build_grouped_mm_offsets(padded_splits, device=x.device)
 
-        active_linear_in = self.linear_in.weight[active_expert_indices]
+        active_linear_in = self.linear_in(active_expert_indices)
         hidden = self._forward_grouped_projection(
             grouped_input,
             weight=active_linear_in,
@@ -1575,7 +2076,7 @@ def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:
 
         if self.config.cpu_offloading and self.config.cpu_offloading_activations:
             hidden.activation_offloading = True
-        active_linear_out = self.linear_out.weight[active_expert_indices]
+        active_linear_out = self.linear_out(active_expert_indices)
         expert_output = self._forward_grouped_projection(
             hidden,
             weight=active_linear_out,
@@ -1616,6 +2117,9 @@ def sharded_state_dict(
                 f"{prefix}linear_in.weight",
                 tp_axis=self._linear_in_tp_axis,
                 sharded_offsets=sharded_offsets,
+                pg_collection=self.pg_collection,
+                ep_size_fallback=self.config.expert_model_parallel_size or 1,
+                etp_size_fallback=self.config.expert_tensor_parallel_size or 1,
             )
         }
         linear_out_sd = {
@@ -1624,10 +2128,13 @@ def sharded_state_dict(
                 f"{prefix}linear_out.weight",
                 tp_axis=self._linear_out_tp_axis,
                 sharded_offsets=sharded_offsets,
+                pg_collection=self.pg_collection,
+                ep_size_fallback=self.config.expert_model_parallel_size or 1,
+                etp_size_fallback=self.config.expert_tensor_parallel_size or 1,
             )
         }
 
-        if "linear_fc1" in self.base_linear_name:
+        if "linear_fc1" in self.base_linear_name and getattr(self.config, "gated_linear_unit", False):
             singleton_local_shards = (metadata or {}).get("singleton_local_shards", False)
             linear_out_key = f"{prefix}linear_out.weight"
             linear_out_sd[linear_out_key] = _apply_grouped_expert_swiglu_sharded_factory(
@@ -1639,3 +2146,292 @@ def sharded_state_dict(
         sharded_state_dict.update(linear_in_sd)
         sharded_state_dict.update(linear_out_sd)
         return sharded_state_dict
+
+
+def _make_cross_ep_replicated(weight: nn.Parameter) -> None:
+    """Mark a weight as logically replicated across the intra-PP-stage group.
+
+    Megatron's DDP routes ``is_expert=True`` parameters through the expert
+    data-parallel group only, which does not span the EP axis. A weight
+    that must stay bit-identical across all EP ranks (e.g., the shared
+    side of :class:`SharedOuterGroupedExpertAdapter`, which a serving
+    engine consumes as a single global LoRA tensor) is otherwise left
+    unsynced. This helper closes that gap with two primitives:
+
+      * a one-shot broadcast from group rank 0 so every rank starts with
+        bit-identical values despite per-rank RNG forks;
+      * a backward hook that SUM all-reduces the gradient across the group
+        so the optimizer step on every rank applies the same update.
+
+    SUM is the correct reduction: each rank's local gradient is the partial
+    loss gradient over its (token, expert) subset, and the total gradient
+    is the sum of those partials. AVG would train at 1/N the intended rate.
+
+    The intra-PP-stage group is ``tensor_and_data_parallel_group`` with
+    context parallel included, which by Megatron's construction equals
+    ETP × EP × EDP — all ranks within the current pipeline stage.
+
+    Args:
+        weight: The parameter to keep replicated across the group. Must
+            be a leaf parameter so the backward hook fires when its
+            gradient is computed.
+    """
+
+    if not (torch.distributed.is_available() and torch.distributed.is_initialized()):
+        return
+    try:
+        group = parallel_state.get_tensor_and_data_parallel_group(with_context_parallel=True)
+    except AssertionError:
+        return
+    if torch.distributed.get_world_size(group=group) <= 1:
+        return
+
+    if weight.is_cuda:
+        # NCCL requires CUDA tensors; pre-GPU construction relies on
+        # deterministic init matching across ranks.
+        src_rank = torch.distributed.get_global_rank(group, 0)
+        with torch.no_grad():
+            torch.distributed.broadcast(weight.data, src=src_rank, group=group)
+
+    def _all_reduce_grad(grad: torch.Tensor) -> torch.Tensor:
+        grad = grad.contiguous()
+        torch.distributed.all_reduce(grad, op=torch.distributed.ReduceOp.SUM, group=group)
+        return grad
+
+    weight.register_hook(_all_reduce_grad)
+
+
+class PackedPerExpertLinear(nn.Module):
+    """Per-expert linear with a packed 3D weight ``[N_local, out, in]``.
+
+    Used as the per-expert side of :class:`SharedOuterGroupedExpertAdapter`.
+    Stores one ``nn.Parameter`` (3D) so Bridge's adapter export sees a single
+    ``.weight`` per side, matching the ``linear_in.weight`` / ``linear_out.weight``
+    convention in :mod:`megatron.bridge.models.conversion.peft_bridge`. Forward
+    dispatches to :func:`torch._grouped_mm` (the same grouped GEMM kernel TE's
+    :class:`te.pytorch.GroupedLinear` calls) via a single fused op with native
+    autograd, which keeps rank kernel launch counts in lockstep so CP's ring
+    P2P does not deadlock.
+    """
+
+    def __init__(
+        self,
+        num_local_experts: int,
+        in_features: int,
+        out_features: int,
+        *,
+        init_method: Optional[Callable] = None,
+        dtype: Optional[torch.dtype] = None,
+        device: Optional[torch.device] = None,
+    ):
+        super().__init__()
+        if not hasattr(torch, "_grouped_mm"):
+            raise RuntimeError("PackedPerExpertLinear requires torch._grouped_mm (torch >= 2.9).")
+        self.num_local_experts = num_local_experts
+        self.in_features = in_features
+        self.out_features = out_features
+        weight = torch.empty(num_local_experts, out_features, in_features, dtype=dtype, device=device)
+        if init_method is not None:
+            for e in range(num_local_experts):
+                init_method(weight[e])
+        else:
+            nn.init.zeros_(weight)
+        self.weight = nn.Parameter(weight)
+        # DDP routes ``is_expert`` weights through the EDP group; the cross-EP
+        # axis is naturally distinct here (different experts on each EP rank).
+        setattr(self.weight, "allreduce", False)
+
+    def forward(self, x: torch.Tensor, m_splits) -> Tuple[torch.Tensor, None]:
+        # torch._grouped_mm expects mat2 as [num_groups, K, N]; our weight is
+        # [N_local, out, in] so transpose the last two dims.
+        if isinstance(m_splits, torch.Tensor):
+            m_splits_i32 = m_splits.to(device=x.device, dtype=torch.int32)
+        else:
+            m_splits_i32 = torch.tensor(m_splits, device=x.device, dtype=torch.int32)
+        offs = torch.cumsum(m_splits_i32, dim=0, dtype=torch.int32)
+        out = torch._grouped_mm(x, self.weight.transpose(1, 2), offs=offs)
+        return out, None
+
+    def sharded_state_dict(
+        self, prefix: str = "", sharded_offsets: Tuple = (), metadata: Optional[Dict] = None
+    ) -> ShardedStateDict:
+        """Shard the packed 3D weight along dim 0 (experts) across EP ranks."""
+        key = f"{prefix}weight"
+        return {
+            key: _make_grouped_expert_sharded_tensor(
+                self.weight.data, key, tp_axis=None, sharded_offsets=sharded_offsets
+            )
+        }
+
+
+class SharedOuterGroupedExpertAdapter(nn.Module):
+    """LoRA adapter for grouped expert MLP with shared-outer semantics.
+
+    Matches SGLang PR #21466's ``experts_shared_outer_loras=True`` contract:
+
+    * fc1 (gate_up):  linear_in  = SHARED     (hidden -> rank)
+                      linear_out = PER-EXPERT (rank -> 2*intermediate)
+    * fc2 (down):     linear_in  = PER-EXPERT (intermediate -> rank)
+                      linear_out = SHARED     (rank -> hidden)
+
+    The shared side is an ``is_expert=True`` ``ColumnParallelLinear`` (fc1)
+    or ``RowParallelLinear`` (fc2): the TP group is ETP (ETP=1 → local
+    forward), DDP routes the weight through the EDP group, and the
+    logically-replicated cross-EP axis is covered by
+    :func:`_make_cross_ep_replicated`.
+
+    The per-expert side is :class:`PackedPerExpertLinear` (packed 3D weight
+    + :func:`torch._grouped_mm`) — kept as a single ``.weight`` Parameter so
+    Bridge's adapter-export materializer (which reads ``linear_in.weight`` /
+    ``linear_out.weight``) sees a standard single-weight linear per side.
+
+    Differs from ``ParallelLinearAdapter`` in ``__init__`` and ``forward``;
+    ``sharded_state_dict`` is specialized for the packed 3D per-expert side.
+    """
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        dim: int,
+        *,
+        num_local_experts: int,
+        base_linear_name: str,
+        activation: str = "swish",
+        column_init_method: str = "xavier",
+        row_init_method: str = "zero",
+        input_is_parallel: bool = False,
+        dropout: float = 0.0,
+        model_parallel_config: Optional[ModelParallelConfig] = None,
+        alpha: Optional[float] = None,
+        dropout_position: str = "pre",
+        base_linear_is_parallel: bool = True,
+        params_device: Optional[torch.device] = None,
+        params_dtype: Optional[torch.dtype] = None,
+    ) -> None:
+        """Initialize shared-outer LoRA weights with one shared and one per-expert side."""
+
+        super().__init__()
+        self.base_linear_name = base_linear_name
+        self.activation = ParallelLinearAdapter._get_activation_fn(self, activation)
+        self.dim = dim
+        self.alpha = alpha if alpha is not None else self.dim
+        self.dropout_position = dropout_position
+        self.num_local_experts = num_local_experts
+        self.base_linear_is_parallel = base_linear_is_parallel
+        # ``is_expert=True`` is observed by param_mapping.py and by inherited
+        # checkpoint helpers; the per-expert side's grad routing is set on its
+        # 3D weight directly inside :class:`PackedPerExpertLinear`.
+        self.is_expert = True
+
+        if model_parallel_config is None:
+            model_parallel_config = ModelParallelConfig()
+        model_parallel_config.perform_initialization = True
+        self.config = model_parallel_config
+
+        # ``input_is_parallel`` selects fc1 (column-parallel base) vs fc2
+        # (row-parallel base). Mirrors :class:`ParallelLinearAdapter` and
+        # :class:`GroupedExpertLinearAdapter`.
+        self._is_fc1 = not input_is_parallel
+
+        column_init = ParallelLinearAdapter._get_init_fn(self, column_init_method)
+        row_init = ParallelLinearAdapter._get_init_fn(self, row_init_method)
+        if self._is_fc1:
+            # Shared A (hidden → rank); per-expert B (rank → 2*intermediate).
+            self.linear_in = ColumnParallelLinear(
+                in_features,
+                dim,
+                config=model_parallel_config,
+                bias=False,
+                gather_output=True,
+                init_method=column_init,
+                is_expert=True,
+            )
+            self.linear_out = PackedPerExpertLinear(
+                num_local_experts,
+                dim,
+                out_features,
+                init_method=row_init,
+                device=params_device,
+                dtype=params_dtype,
+            )
+        else:
+            # Per-expert A (intermediate → rank); shared B (rank → hidden).
+            self.linear_in = PackedPerExpertLinear(
+                num_local_experts,
+                in_features,
+                dim,
+                init_method=column_init,
+                device=params_device,
+                dtype=params_dtype,
+            )
+            self.linear_out = RowParallelLinear(
+                dim,
+                out_features,
+                config=model_parallel_config,
+                bias=False,
+                input_is_parallel=True,
+                skip_bias_add=True,
+                init_method=row_init,
+                is_expert=True,
+            )
+
+        self.dropout = nn.Dropout(dropout) if dropout > 0.0 else nn.Identity()
+
+        if model_parallel_config.bf16:
+            self.bfloat16()
+        elif model_parallel_config.fp16:
+            self.half()
+
+        # The shared weight is logically replicated across EP; close the gap
+        # that Megatron's expert-DDP routing leaves open.
+        shared_weight = self.linear_in.weight if self._is_fc1 else self.linear_out.weight
+        _make_cross_ep_replicated(shared_weight)
+
+    def forward(self, x: torch.Tensor, m_splits=None) -> torch.Tensor:
+        """Forward. ``m_splits`` is the tokens-per-expert split passed through
+        from the base TEGroupedLinear; required for the per-expert side.
+        """
+        if self.dropout_position == "pre":
+            x = self.dropout(x)
+
+        if self._is_fc1:
+            # Shared A → activation → per-expert B.
+            x, _ = self.linear_in(x)
+            x = self.activation(x)
+            x, _ = self.linear_out(x, m_splits)
+        else:
+            # Per-expert A → activation → shared B.
+            x, _ = self.linear_in(x, m_splits)
+            x = self.activation(x)
+            x, _ = self.linear_out(x)
+
+        if self.dropout_position == "post":
+            x = self.dropout(x)
+
+        return x * (self.alpha / self.dim)
+
+    def sharded_state_dict(
+        self,
+        prefix: str = "",
+        sharded_offsets: Tuple = (),
+        metadata: Optional[Dict] = None,
+    ) -> ShardedStateDict:
+        """Create sharded state dictionary for mixed shared/per-expert adapter weights."""
+
+        linear_in_sd = self.linear_in.sharded_state_dict(f"{prefix}linear_in.", sharded_offsets, metadata)
+        linear_out_sd = self.linear_out.sharded_state_dict(f"{prefix}linear_out.", sharded_offsets, metadata)
+
+        if self._is_fc1:
+            singleton_local_shards = (metadata or {}).get("singleton_local_shards", False)
+            linear_out_key = f"{prefix}linear_out.weight"
+            linear_out_sd[linear_out_key] = _apply_grouped_expert_swiglu_sharded_factory(
+                linear_out_sd[linear_out_key],
+                sharded_offsets,
+                singleton_local_shards,
+            )
+
+        sharded_state_dict = {}
+        sharded_state_dict.update(linear_in_sd)
+        sharded_state_dict.update(linear_out_sd)
+        return sharded_state_dict
diff --git a/tests/unit_tests/models/test_adapter_export.py b/tests/unit_tests/models/test_adapter_export.py
index a43ecd75d7..cb3021df54 100644
--- a/tests/unit_tests/models/test_adapter_export.py
+++ b/tests/unit_tests/models/test_adapter_export.py
@@ -16,7 +16,9 @@
 
 from __future__ import annotations
 
+import argparse
 import json
+import logging
 from contextlib import nullcontext
 from dataclasses import dataclass
 from pathlib import Path
@@ -444,6 +446,71 @@ def test_save_creates_files(self, tmp_path):
         assert "target_parameters" not in cfg
         assert cfg["base_model_name_or_path"] == "test/model"
 
+    def test_save_preserves_adapter_dtype_and_detaches(self, tmp_path):
+        """save_hf_adapter should write detached adapter tensors in their exported dtype."""
+        from safetensors.torch import load_file
+
+        from megatron.bridge.peft.lora import LoRA
+
+        output_dir = tmp_path / "adapter_out_dtype"
+        fake_weights = [
+            _adapter_export(
+                "model.layers.0.self_attn.q_proj.lora_A.weight",
+                torch.randn(8, 64, dtype=torch.bfloat16).requires_grad_(),
+            ),
+            _adapter_export(
+                "model.layers.0.self_attn.q_proj.lora_B.weight",
+                torch.randn(64, 8, dtype=torch.bfloat16).requires_grad_(),
+            ),
+        ]
+
+        mock_bridge = MagicMock()
+        mock_bridge.export_adapter_weights.return_value = iter(fake_weights)
+        mock_bridge.hf_pretrained = _ToyAdapterModel(model_name_or_path="test/model")
+
+        with patch("torch.distributed.is_initialized", return_value=False):
+            from megatron.bridge.models.conversion.auto_bridge import AutoBridge
+
+            AutoBridge.save_hf_adapter(
+                mock_bridge,
+                model=[MagicMock()],
+                path=output_dir,
+                peft_config=LoRA(),
+                base_model_name_or_path="test/model",
+            )
+
+        state = load_file(str(output_dir / "adapter_model.safetensors"))
+        assert state
+        assert {tensor.dtype for tensor in state.values()} == {torch.bfloat16}
+
+    def test_save_passes_exclude_adapter_base_prefixes(self, tmp_path):
+        """save_hf_adapter should pass adapter-base exclusions to adapter streaming."""
+        from megatron.bridge.peft.lora import LoRA
+
+        output_dir = tmp_path / "adapter_out_excluded"
+        fake_weights = [
+            _adapter_export("model.layers.0.self_attn.q_proj.lora_A.weight", torch.randn(8, 64)),
+            _adapter_export("model.layers.0.self_attn.q_proj.lora_B.weight", torch.randn(64, 8)),
+        ]
+
+        mock_bridge = MagicMock()
+        mock_bridge.export_adapter_weights.return_value = iter(fake_weights)
+        mock_bridge.hf_pretrained = _ToyAdapterModel(model_name_or_path="test/model")
+
+        with patch("torch.distributed.is_initialized", return_value=False):
+            from megatron.bridge.models.conversion.auto_bridge import AutoBridge
+
+            AutoBridge.save_hf_adapter(
+                mock_bridge,
+                model=[MagicMock()],
+                path=output_dir,
+                peft_config=LoRA(),
+                base_model_name_or_path="test/model",
+                exclude_adapter_base_prefixes=("mtp.layers",),
+            )
+
+        assert mock_bridge.export_adapter_weights.call_args.kwargs["exclude_adapter_base_prefixes"] == ("mtp.layers",)
+
     def test_save_with_nonzero_dropout_keeps_linear_target_modules(self, tmp_path):
         from megatron.bridge.peft.lora import LoRA
 
@@ -745,10 +812,11 @@ def _patch_heavy_deps(self):
         """Mock out dist_checkpointing, distributed context, and checkpoint helpers."""
         fake_sd = {"model": {}}
         with (
+            patch("megatron.core.dist_checkpointing.load", return_value=fake_sd) as mock_dist_load,
             patch(
-                "megatron.core.dist_checkpointing.load",
-                return_value=fake_sd,
-            ) as self.mock_dist_load,
+                "megatron.bridge.peft.utils.enable_legacy_shared_expert_adapter_loading",
+                return_value=False,
+            ),
             patch(
                 "megatron.bridge.training.checkpointing._generate_model_state_dict",
                 return_value={"model": {}},
@@ -762,6 +830,7 @@ def _patch_heavy_deps(self):
                 return_value=nullcontext(),
             ),
         ):
+            self.mock_dist_load = mock_dist_load
             yield
 
     def test_basic_export_calls_save_hf_adapter(self, bridge, ckpt_dir, tmp_path):
@@ -813,6 +882,32 @@ def test_vlmlora_selected_when_target_matches(self, bridge, tmp_path):
         )
         assert isinstance(peft_config, VLMLoRA)
 
+    def test_plain_lora_filters_vlm_only_keys(self, bridge, tmp_path):
+        """VLM-only keys should not be passed into plain LoRA configs."""
+        ckpt = tmp_path / "plain_lora_ckpt"
+        ckpt.mkdir()
+        run_cfg = {
+            "peft": {
+                "_target_": "megatron.bridge.peft.lora.LoRA",
+                "dim": 8,
+                "alpha": 16,
+                "freeze_language_model": False,
+                "freeze_vision_model": False,
+                "freeze_vision_projection": False,
+            }
+        }
+        (ckpt / "run_config.yaml").write_text(yaml.dump(run_cfg))
+
+        bridge.export_adapter_ckpt(str(ckpt), tmp_path / "out")
+
+        peft_config = bridge.save_hf_adapter.call_args.kwargs.get(
+            "peft_config",
+            bridge.save_hf_adapter.call_args.args[2] if len(bridge.save_hf_adapter.call_args.args) > 2 else None,
+        )
+        assert peft_config.dim == 8
+        assert peft_config.alpha == 16
+        assert not hasattr(peft_config, "freeze_language_model")
+
     def test_missing_checkpoint_raises(self, bridge, tmp_path):
         with pytest.raises(FileNotFoundError, match="PEFT checkpoint not found"):
             bridge.export_adapter_ckpt(str(tmp_path / "nonexistent"), tmp_path / "out")
@@ -874,8 +969,8 @@ def test_corrupt_run_config_falls_back_to_defaults(self, bridge, tmp_path):
         assert isinstance(peft_config, LoRA)
         assert peft_config.dim == LoRA().dim
 
-    def test_provider_set_to_float32(self, bridge, ckpt_dir, tmp_path):
-        """Provider dtypes must be forced to float32 for full-precision adapter export."""
+    def test_provider_defaults_to_float32(self, bridge, ckpt_dir, tmp_path):
+        """Provider dtypes default to float32 for full-precision adapter export."""
         bridge.export_adapter_ckpt(str(ckpt_dir), tmp_path / "out")
 
         provider = bridge.to_megatron_provider.return_value
@@ -883,6 +978,24 @@ def test_provider_set_to_float32(self, bridge, ckpt_dir, tmp_path):
         assert provider.params_dtype == torch.float32
         provider.finalize.assert_called_once()
 
+    def test_export_adapter_ckpt_does_not_pass_output_dtype(self, bridge, ckpt_dir, tmp_path):
+        """CPU checkpoint export should keep full precision and not override saved dtype."""
+        bridge.export_adapter_ckpt(str(ckpt_dir), tmp_path / "out")
+
+        save_kwargs = bridge.save_hf_adapter.call_args.kwargs
+        assert "output_dtype" not in save_kwargs
+
+    def test_exclude_adapter_base_prefixes_passed_to_save(self, bridge, ckpt_dir, tmp_path):
+        """Adapter-base exclusions should be threaded from checkpoint export to save_hf_adapter."""
+        bridge.export_adapter_ckpt(
+            str(ckpt_dir),
+            tmp_path / "out",
+            exclude_adapter_base_prefixes=("mtp.layers",),
+        )
+
+        save_kwargs = bridge.save_hf_adapter.call_args.kwargs
+        assert save_kwargs["exclude_adapter_base_prefixes"] == ("mtp.layers",)
+
     def test_dist_checkpointing_called_with_ckpt_path(self, bridge, ckpt_dir, tmp_path):
         bridge.export_adapter_ckpt(str(ckpt_dir), tmp_path / "out")
 
@@ -915,3 +1028,399 @@ def test_multi_rank_export_runs_in_current_process(self, bridge, ckpt_dir, tmp_p
             bridge.export_adapter_ckpt(str(ckpt_dir), output)
 
         bridge.save_hf_adapter.assert_called_once()
+
+
+# ---------------------------------------------------------------------------
+# examples/conversion/adapter/export_adapter.py
+# ---------------------------------------------------------------------------
+
+
+class TestExportAdapterScript:
+    def test_parse_dtype_uses_shared_dtype_map(self):
+        from examples.conversion.adapter import export_adapter
+
+        assert export_adapter._parse_dtype("bf16") == torch.bfloat16
+        assert export_adapter._parse_dtype("torch.float32") == torch.float32
+
+    def test_parse_dtype_rejects_invalid_dtype(self):
+        from examples.conversion.adapter import export_adapter
+
+        with pytest.raises(argparse.ArgumentTypeError, match="Unknown dtype"):
+            export_adapter._parse_dtype("float8-ish")
+
+    def test_parse_dtype_rejects_unsupported_export_dtype(self):
+        from examples.conversion.adapter import export_adapter
+
+        with pytest.raises(argparse.ArgumentTypeError, match="Unsupported adapter export dtype"):
+            export_adapter._parse_dtype("fp8")
+
+    def test_load_lora_config_falls_back_to_defaults_on_parse_error(self, tmp_path, caplog):
+        from examples.conversion.adapter import export_adapter
+
+        from megatron.bridge.peft.lora import LoRA
+
+        ckpt = tmp_path / "adapter_ckpt"
+        ckpt.mkdir()
+        (ckpt / "run_config.yaml").write_text("not: valid: yaml: [[[")
+
+        with (
+            patch("examples.conversion.adapter.export_adapter.read_run_config", side_effect=ValueError("bad yaml")),
+            caplog.at_level(logging.WARNING),
+        ):
+            lora = export_adapter._load_lora_config(ckpt)
+
+        assert isinstance(lora, LoRA)
+        assert lora.dim == LoRA().dim
+        assert "Using defaults" in caplog.text
+
+    def test_load_lora_config_reads_parent_run_config(self, tmp_path):
+        from examples.conversion.adapter import export_adapter
+
+        parent = tmp_path / "run"
+        iter_dir = parent / "iter_0000001"
+        iter_dir.mkdir(parents=True)
+        (parent / "run_config.yaml").write_text(yaml.dump({"peft": {"_target_": "LoRA", "dim": 4, "alpha": 8}}))
+
+        lora = export_adapter._load_lora_config(iter_dir)
+
+        assert lora.dim == 4
+        assert lora.alpha == 8
+
+    def test_load_lora_config_filters_vlm_keys_for_plain_lora(self, tmp_path):
+        from examples.conversion.adapter import export_adapter
+
+        ckpt = tmp_path / "adapter_ckpt"
+        ckpt.mkdir()
+        (ckpt / "run_config.yaml").write_text(
+            yaml.dump(
+                {
+                    "peft": {
+                        "_target_": "megatron.bridge.peft.lora.LoRA",
+                        "dim": 8,
+                        "alpha": 16,
+                        "freeze_language_model": False,
+                        "freeze_vision_model": False,
+                        "freeze_vision_projection": False,
+                    }
+                }
+            )
+        )
+
+        lora = export_adapter._load_lora_config(ckpt)
+
+        assert lora.dim == 8
+        assert lora.alpha == 16
+        assert not hasattr(lora, "freeze_language_model")
+
+    def test_get_loaded_model_key_prefers_exact_model_key(self, tmp_path):
+        from examples.conversion.adapter import export_adapter
+
+        assert export_adapter._get_loaded_model_key({"model": {}, "model0": {}}, tmp_path) == "model"
+
+    def test_get_loaded_model_key_accepts_prefixed_model_key(self, tmp_path):
+        from examples.conversion.adapter import export_adapter
+
+        assert export_adapter._get_loaded_model_key({"model0": {}}, tmp_path) == "model0"
+
+    def test_get_loaded_model_key_raises_clear_error_when_missing(self, tmp_path):
+        from examples.conversion.adapter import export_adapter
+
+        with pytest.raises(RuntimeError, match="has no 'model' key"):
+            export_adapter._get_loaded_model_key({"optimizer": {}}, tmp_path)
+
+    @pytest.mark.parametrize(
+        ("tp", "pp", "ep", "etp", "expected"),
+        [
+            (1, 1, 1, 1, False),
+            (2, 1, 1, 1, True),
+            (1, 2, 1, 1, True),
+            (1, 1, 2, 1, True),
+            (1, 1, 1, 2, True),
+        ],
+    )
+    def test_uses_distributed_export(self, tp, pp, ep, etp, expected):
+        from examples.conversion.adapter import export_adapter
+
+        args = SimpleNamespace(tp=tp, pp=pp, ep=ep, etp=etp)
+
+        assert export_adapter._uses_distributed_export(args) is expected
+
+    def test_main_rejects_non_fp32_dtype_for_cpu_export(self, tmp_path):
+        from examples.conversion.adapter import export_adapter
+
+        args = SimpleNamespace(
+            hf_model_path="test/model",
+            trust_remote_code=False,
+            lora_checkpoint=str(tmp_path),
+            output=tmp_path / "out",
+            tp=1,
+            pp=1,
+            ep=1,
+            etp=1,
+            sequence_parallel=False,
+            dtype=torch.bfloat16,
+            exclude_adapter_base_prefix=[],
+        )
+
+        with (
+            patch("examples.conversion.adapter.export_adapter.parse_args", return_value=args),
+            patch("examples.conversion.adapter.export_adapter.AutoBridge.from_hf_pretrained") as mock_from_hf,
+            pytest.raises(ValueError, match="only supported by distributed GPU export"),
+        ):
+            export_adapter.main()
+
+        mock_from_hf.assert_not_called()
+
+    def test_main_cpu_export_does_not_pass_dtype_to_autobridge(self, tmp_path):
+        from examples.conversion.adapter import export_adapter
+
+        args = SimpleNamespace(
+            hf_model_path="test/model",
+            trust_remote_code=False,
+            lora_checkpoint=str(tmp_path),
+            output=tmp_path / "out",
+            tp=1,
+            pp=1,
+            ep=1,
+            etp=1,
+            sequence_parallel=False,
+            dtype=torch.float32,
+            exclude_adapter_base_prefix=["mtp.layers"],
+        )
+        bridge = MagicMock()
+
+        with (
+            patch("examples.conversion.adapter.export_adapter.parse_args", return_value=args),
+            patch(
+                "examples.conversion.adapter.export_adapter.AutoBridge.from_hf_pretrained",
+                return_value=bridge,
+            ) as mock_from_hf,
+        ):
+            export_adapter.main()
+
+        mock_from_hf.assert_called_once_with("test/model", trust_remote_code=False)
+        bridge.export_adapter_ckpt.assert_called_once_with(
+            peft_checkpoint=str(tmp_path),
+            output_path=tmp_path / "out",
+            exclude_adapter_base_prefixes=("mtp.layers",),
+        )
+
+    def test_configure_cuda_device_requires_cuda(self):
+        from examples.conversion.adapter import export_adapter
+
+        with (
+            patch("examples.conversion.adapter.export_adapter.torch.cuda.is_available", return_value=False),
+            pytest.raises(RuntimeError, match="requires CUDA"),
+        ):
+            export_adapter._configure_cuda_device()
+
+    def test_configure_cuda_device_uses_local_rank(self):
+        from examples.conversion.adapter import export_adapter
+
+        with (
+            patch("examples.conversion.adapter.export_adapter.torch.cuda.is_available", return_value=True),
+            patch("examples.conversion.adapter.export_adapter.get_local_rank_preinit", return_value=2),
+            patch("examples.conversion.adapter.export_adapter.torch.cuda.set_device") as mock_set_device,
+        ):
+            device = export_adapter._configure_cuda_device()
+
+        mock_set_device.assert_called_once_with(2)
+        assert device == torch.device("cuda", 2)
+
+    def test_export_adapter_distributed_missing_checkpoint_raises_clear_error(self, tmp_path):
+        from examples.conversion.adapter import export_adapter
+
+        args = SimpleNamespace(
+            hf_model_path="test/model",
+            trust_remote_code=False,
+            lora_checkpoint=str(tmp_path / "missing"),
+            output=tmp_path / "out",
+            tp=2,
+            pp=1,
+            ep=1,
+            etp=1,
+            sequence_parallel=False,
+            dtype=torch.float32,
+            exclude_adapter_base_prefix=[],
+        )
+
+        with (
+            patch(
+                "examples.conversion.adapter.export_adapter._configure_cuda_device", return_value=torch.device("cpu")
+            ),
+            patch("examples.conversion.adapter.export_adapter.AutoConfig.from_pretrained") as mock_from_pretrained,
+            pytest.raises(FileNotFoundError, match="PEFT checkpoint not found"),
+        ):
+            export_adapter._export_adapter_distributed(args)
+
+        mock_from_pretrained.assert_not_called()
+
+    def test_export_adapter_distributed_rejects_multiple_model_chunks(self, tmp_path):
+        from examples.conversion.adapter import export_adapter
+
+        ckpt = tmp_path / "adapter_ckpt"
+        ckpt.mkdir()
+        args = SimpleNamespace(
+            hf_model_path="test/model",
+            trust_remote_code=False,
+            lora_checkpoint=str(ckpt),
+            output=tmp_path / "out",
+            tp=2,
+            pp=2,
+            ep=1,
+            etp=1,
+            sequence_parallel=False,
+            dtype=torch.float32,
+            exclude_adapter_base_prefix=[],
+        )
+        model_chunks = [MagicMock(), MagicMock()]
+        for chunk in model_chunks:
+            chunk.to.return_value = chunk
+        provider = MagicMock()
+        provider.provide_distributed_model.return_value = model_chunks
+        bridge = MagicMock()
+        bridge.to_megatron_provider.return_value = provider
+
+        with (
+            patch(
+                "examples.conversion.adapter.export_adapter._configure_cuda_device", return_value=torch.device("cpu")
+            ),
+            patch("examples.conversion.adapter.export_adapter.AutoConfig.from_pretrained", return_value=MagicMock()),
+            patch("examples.conversion.adapter.export_adapter.AutoBridge.from_hf_config", return_value=bridge),
+            patch("examples.conversion.adapter.export_adapter._load_lora_config", return_value=MagicMock()),
+            patch("examples.conversion.adapter.export_adapter._generate_model_state_dict") as mock_state_dict,
+            patch("examples.conversion.adapter.export_adapter.dist_checkpointing.load") as mock_load,
+            patch("examples.conversion.adapter.export_adapter.parallel_state.is_initialized", return_value=True),
+            patch(
+                "examples.conversion.adapter.export_adapter.parallel_state.destroy_model_parallel"
+            ) as mock_destroy_mp,
+            patch("examples.conversion.adapter.export_adapter.dist.is_initialized", return_value=True),
+            patch("examples.conversion.adapter.export_adapter.dist.destroy_process_group") as mock_destroy_pg,
+            pytest.raises(RuntimeError, match="exactly one local model chunk"),
+        ):
+            export_adapter._export_adapter_distributed(args)
+
+        mock_state_dict.assert_not_called()
+        mock_load.assert_not_called()
+        bridge.save_hf_adapter.assert_not_called()
+        for chunk in model_chunks:
+            chunk.load_state_dict.assert_not_called()
+        mock_destroy_mp.assert_called_once()
+        mock_destroy_pg.assert_called_once()
+
+    def test_export_adapter_distributed_raises_clear_error_for_missing_model_key(self, tmp_path):
+        from examples.conversion.adapter import export_adapter
+
+        args = SimpleNamespace(
+            hf_model_path="test/model",
+            trust_remote_code=False,
+            lora_checkpoint=str(tmp_path),
+            output=tmp_path / "out",
+            tp=2,
+            pp=1,
+            ep=1,
+            etp=1,
+            sequence_parallel=False,
+            dtype=torch.float32,
+            exclude_adapter_base_prefix=[],
+        )
+        model_chunk = MagicMock()
+        model_chunk.to.return_value = model_chunk
+        provider = MagicMock()
+        provider.provide_distributed_model.return_value = [model_chunk]
+        bridge = MagicMock()
+        bridge.to_megatron_provider.return_value = provider
+
+        with (
+            patch(
+                "examples.conversion.adapter.export_adapter._configure_cuda_device", return_value=torch.device("cpu")
+            ),
+            patch("examples.conversion.adapter.export_adapter.AutoConfig.from_pretrained", return_value=MagicMock()),
+            patch("examples.conversion.adapter.export_adapter.AutoBridge.from_hf_config", return_value=bridge),
+            patch("examples.conversion.adapter.export_adapter._load_lora_config", return_value=MagicMock()),
+            patch("examples.conversion.adapter.export_adapter._generate_model_state_dict", return_value={"model": {}}),
+            patch(
+                "examples.conversion.adapter.export_adapter.apply_peft_adapter_filter_to_state_dict",
+                side_effect=lambda state_dict, _lora: state_dict,
+            ),
+            patch(
+                "examples.conversion.adapter.export_adapter.enable_legacy_shared_expert_adapter_loading",
+                return_value=False,
+            ),
+            patch(
+                "examples.conversion.adapter.export_adapter.dist_checkpointing.load", return_value={"optimizer": {}}
+            ),
+            patch("examples.conversion.adapter.export_adapter.parallel_state.is_initialized", return_value=True),
+            patch(
+                "examples.conversion.adapter.export_adapter.parallel_state.destroy_model_parallel"
+            ) as mock_destroy_mp,
+            patch("examples.conversion.adapter.export_adapter.dist.is_initialized", return_value=True),
+            patch("examples.conversion.adapter.export_adapter.dist.destroy_process_group") as mock_destroy_pg,
+            pytest.raises(RuntimeError, match="has no 'model' key"),
+        ):
+            export_adapter._export_adapter_distributed(args)
+
+        model_chunk.load_state_dict.assert_not_called()
+        mock_destroy_mp.assert_called_once()
+        mock_destroy_pg.assert_called_once()
+
+    def test_export_adapter_distributed_enables_legacy_shared_expert_adapter_loading(self, tmp_path):
+        from examples.conversion.adapter import export_adapter
+
+        args = SimpleNamespace(
+            hf_model_path="test/model",
+            trust_remote_code=False,
+            lora_checkpoint=str(tmp_path),
+            output=tmp_path / "out",
+            tp=2,
+            pp=1,
+            ep=1,
+            etp=1,
+            sequence_parallel=False,
+            dtype=torch.float32,
+            exclude_adapter_base_prefix=[],
+        )
+        model_chunk = MagicMock()
+        model_chunk.to.return_value = model_chunk
+        provider = MagicMock()
+        provider.provide_distributed_model.return_value = [model_chunk]
+        bridge = MagicMock()
+        bridge.to_megatron_provider.return_value = provider
+        lora = MagicMock()
+        state_dicts = [{"model": {"new": object()}}, {"model": {"legacy": object()}}]
+
+        with (
+            patch(
+                "examples.conversion.adapter.export_adapter._configure_cuda_device", return_value=torch.device("cpu")
+            ),
+            patch("examples.conversion.adapter.export_adapter.AutoConfig.from_pretrained", return_value=MagicMock()),
+            patch("examples.conversion.adapter.export_adapter.AutoBridge.from_hf_config", return_value=bridge),
+            patch("examples.conversion.adapter.export_adapter._load_lora_config", return_value=lora),
+            patch(
+                "examples.conversion.adapter.export_adapter._generate_model_state_dict",
+                side_effect=state_dicts,
+            ) as mock_state_dict,
+            patch(
+                "examples.conversion.adapter.export_adapter.apply_peft_adapter_filter_to_state_dict",
+                side_effect=lambda state_dict, _lora: state_dict,
+            ) as mock_filter,
+            patch(
+                "examples.conversion.adapter.export_adapter.enable_legacy_shared_expert_adapter_loading",
+                return_value=True,
+            ) as mock_enable_legacy,
+            patch(
+                "examples.conversion.adapter.export_adapter.dist_checkpointing.load",
+                return_value={"model": {"adapter": "weights"}},
+            ) as mock_load,
+            patch("examples.conversion.adapter.export_adapter.parallel_state.is_initialized", return_value=True),
+            patch("examples.conversion.adapter.export_adapter.parallel_state.destroy_model_parallel"),
+            patch("examples.conversion.adapter.export_adapter.dist.is_initialized", return_value=True),
+            patch("examples.conversion.adapter.export_adapter.dist.destroy_process_group"),
+        ):
+            export_adapter._export_adapter_distributed(args)
+
+        assert mock_state_dict.call_count == 2
+        assert mock_filter.call_count == 2
+        mock_enable_legacy.assert_called_once_with([model_chunk], state_dicts[0], tmp_path)
+        mock_load.assert_called_once_with(state_dicts[1], str(tmp_path), validate_access_integrity=False)
+        model_chunk.load_state_dict.assert_called_once_with({"adapter": "weights"}, strict=False)
diff --git a/tests/unit_tests/models/test_auto_bridge.py b/tests/unit_tests/models/test_auto_bridge.py
index 4e922b5ad5..a55ca3db59 100644
--- a/tests/unit_tests/models/test_auto_bridge.py
+++ b/tests/unit_tests/models/test_auto_bridge.py
@@ -875,6 +875,7 @@ def test_export_adapter_weights(self):
                         mock_megatron_model,
                         cpu=False,
                         show_progress=False,
+                        exclude_adapter_base_prefixes=None,
                     )
 
     def test_get_causal_lm_architecture(self):
diff --git a/tests/unit_tests/models/test_model_bridge_lora.py b/tests/unit_tests/models/test_model_bridge_lora.py
index 2bbadc431d..c53764ca0b 100644
--- a/tests/unit_tests/models/test_model_bridge_lora.py
+++ b/tests/unit_tests/models/test_model_bridge_lora.py
@@ -766,6 +766,36 @@ def test_build_adapter_conversion_tasks(monkeypatch):
     assert task.linear_out_task.param_weight.shape == torch.Size([2, 2])
 
 
+def test_build_adapter_conversion_tasks_excludes_base_prefix_before_mapping(monkeypatch):
+    bridge = DummyBridge()
+    bridge.hf_pretrained = SimpleNamespace()
+    bridge.hf_config = bridge.hf_pretrained
+
+    adapters_info = [
+        (
+            "mtp.layers.0.mtp_model_layer.layers.0.self_attention.linear_proj.adapter",
+            "mtp.layers.0.mtp_model_layer.layers.0.self_attention.linear_proj",
+            False,
+            False,
+            False,
+            4,
+            8,
+            0,
+            0,
+        )
+    ]
+
+    monkeypatch.setattr(bridge, "_megatron_global_adapters_info_all_pp_ranks", lambda *_: adapters_info)
+    monkeypatch.setattr(bridge, "mapping_registry", lambda: MegatronMappingRegistry())
+
+    tasks_by_base = bridge.build_adapter_conversion_tasks(
+        [Mock()],
+        exclude_adapter_base_prefixes=("mtp.layers",),
+    )
+
+    assert tasks_by_base == {}
+
+
 def test_materialize_adapter_weights(monkeypatch):
     bridge = DummyBridge()
 
@@ -1022,7 +1052,7 @@ def test_stream_adapter_weights_megatron_to_hf(monkeypatch):
     monkeypatch.setattr(
         bridge,
         "build_adapter_conversion_tasks",
-        lambda *_: {"decoder.layers.0.mlp.linear_fc1": [adapter_task]},
+        lambda *_args, **_kwargs: {"decoder.layers.0.mlp.linear_fc1": [adapter_task]},
     )
     monkeypatch.setattr(
         bridge,
@@ -1087,7 +1117,7 @@ def test_stream_adapter_weights_megatron_to_hf_qkv(monkeypatch):
     monkeypatch.setattr(
         bridge,
         "build_adapter_conversion_tasks",
-        lambda *_: {"decoder.layers.0.self_attn.linear_qkv": [adapter_task]},
+        lambda *_args, **_kwargs: {"decoder.layers.0.self_attn.linear_qkv": [adapter_task]},
     )
     monkeypatch.setattr(bridge, "materialize_adapter_weights", lambda *_: [adapter_weight])
     monkeypatch.setattr(
@@ -1166,7 +1196,7 @@ def test_stream_adapter_weights_megatron_to_hf_fused_fc1(monkeypatch):
     monkeypatch.setattr(
         bridge,
         "build_adapter_conversion_tasks",
-        lambda *_: {"decoder.layers.0.mlp.linear_fc1": [adapter_task]},
+        lambda *_args, **_kwargs: {"decoder.layers.0.mlp.linear_fc1": [adapter_task]},
     )
     monkeypatch.setattr(bridge, "materialize_adapter_weights", lambda *_: [adapter_weight])
     monkeypatch.setattr(
@@ -1236,7 +1266,7 @@ def test_stream_adapter_weights_megatron_to_hf_fused_fc1_minimax_w13(monkeypatch
     monkeypatch.setattr(
         bridge,
         "build_adapter_conversion_tasks",
-        lambda *_: {"decoder.layers.0.mlp.experts.linear_fc1": [adapter_task]},
+        lambda *_args, **_kwargs: {"decoder.layers.0.mlp.experts.linear_fc1": [adapter_task]},
     )
     monkeypatch.setattr(bridge, "materialize_adapter_weights", lambda *_: [adapter_weight])
     monkeypatch.setattr(
@@ -1317,7 +1347,7 @@ def test_stream_adapter_weights_megatron_to_hf_packed_expert_stacks(monkeypatch)
     monkeypatch.setattr(
         bridge,
         "build_adapter_conversion_tasks",
-        lambda *_: {"decoder.layers.0.mlp.experts.linear_fc2": [adapter_task]},
+        lambda *_args, **_kwargs: {"decoder.layers.0.mlp.experts.linear_fc2": [adapter_task]},
     )
     monkeypatch.setattr(bridge, "materialize_adapter_weights", lambda *_: [adapter_weight])
     monkeypatch.setattr(
@@ -1377,7 +1407,7 @@ def test_stream_adapter_weights_megatron_to_hf_grouped_expert_exports_per_expert
     monkeypatch.setattr(
         bridge,
         "build_adapter_conversion_tasks",
-        lambda *_: {"decoder.layers.0.mlp.experts.linear_fc2": [adapter_task]},
+        lambda *_args, **_kwargs: {"decoder.layers.0.mlp.experts.linear_fc2": [adapter_task]},
     )
     monkeypatch.setattr(bridge, "materialize_adapter_weights", lambda *_: [adapter_weight])
     monkeypatch.setattr(
@@ -1405,6 +1435,172 @@ def test_stream_adapter_weights_megatron_to_hf_grouped_expert_exports_per_expert
     assert weights[1].param_name == "model.layers.0.mlp.experts.0.down_proj.lora_B.weight"
 
 
+def test_stream_adapter_weights_megatron_to_hf_shared_outer_fc1_gate_up(monkeypatch):
+    # Shared-outer FC1 (SGLang PR #21466): linear_in (lora_A) is a single 2D matrix
+    # shared across experts, while linear_out (lora_B) is a per-expert 3D pack of the
+    # fused gate/up projection. The shared side is emitted once under an
+    # expert-agnostic name; the per-expert side is split into gate/up per expert.
+    bridge = DummyBridge()
+
+    adapter_task = AdapterWeightConversionTask(
+        global_base_prefix="decoder.layers.0.mlp.experts.linear_fc1",
+        adapter_key=None,
+        alpha=2,
+        dim=4,
+        linear_in_task=WeightConversionTask(
+            param_name="local_in",
+            global_param_name="decoder.layers.0.mlp.experts.linear_fc1.adapter.linear_in.weight",
+            mapping=Mock(),
+        ),
+        linear_out_task=WeightConversionTask(
+            param_name="local_out",
+            global_param_name="decoder.layers.0.mlp.experts.linear_fc1.adapter.linear_out.weight",
+            mapping=Mock(),
+        ),
+    )
+
+    # Shared lora_A: [rank=2, hidden=3]. Per-expert lora_B: [num_experts=2, 2*inter=4, rank=2],
+    # gate = first 2 rows, up = last 2 rows, with distinct values per expert/projection.
+    shared_lora_a = torch.ones(2, 3)
+    expert0_lora_b = torch.cat([torch.full((2, 2), 10.0), torch.full((2, 2), 20.0)], dim=0)
+    expert1_lora_b = torch.cat([torch.full((2, 2), 30.0), torch.full((2, 2), 40.0)], dim=0)
+    per_expert_lora_b = torch.stack([expert0_lora_b, expert1_lora_b], dim=0)
+
+    adapter_weight = AdapterWeight(
+        global_base_prefix="decoder.layers.0.mlp.experts.linear_fc1",
+        adapter_key=None,
+        alpha=2,
+        dim=4,
+        linear_in_weight=MegatronWeightTuple("local_in", shared_lora_a, vp_stage=0),
+        linear_out_weight=MegatronWeightTuple("local_out", per_expert_lora_b, vp_stage=0),
+    )
+
+    def fake_base_names(_registry, _prefix, _adapter_key, base_suffix):
+        # base_suffix is ".weight0"/".weight1"/...; reflect the expert index into the
+        # HF names so the per-expert side keeps experts.<idx> and the shared side strips it.
+        idx = base_suffix[len(".weight") :]
+        return [
+            f"model.layers.0.mlp.experts.{idx}.gate_proj.weight",
+            f"model.layers.0.mlp.experts.{idx}.up_proj.weight",
+        ]
+
+    monkeypatch.setattr(
+        bridge,
+        "build_adapter_conversion_tasks",
+        lambda *_args, **_kwargs: {"decoder.layers.0.mlp.experts.linear_fc1": [adapter_task]},
+    )
+    monkeypatch.setattr(bridge, "materialize_adapter_weights", lambda *_: [adapter_weight])
+    monkeypatch.setattr(bridge, "_get_base_hf_param_names_for_adapter", fake_base_names)
+    monkeypatch.setattr(
+        "megatron.bridge.models.conversion.peft_bridge.parallel_state.get_expert_model_parallel_world_size",
+        lambda: 1,
+    )
+
+    weights = list(
+        bridge.stream_adapter_weights_megatron_to_hf(
+            [SimpleNamespace(config=SimpleNamespace(num_moe_experts=2))],
+            cpu=False,
+            show_progress=False,
+        )
+    )
+
+    # Shared lora_A emitted once per fused projection under the expert-agnostic name,
+    # then per-expert lora_B split into gate/up for each expert.
+    assert [w.param_name for w in weights] == [
+        "model.layers.0.mlp.experts.gate_proj.lora_A.weight",
+        "model.layers.0.mlp.experts.up_proj.lora_A.weight",
+        "model.layers.0.mlp.experts.0.gate_proj.lora_B.weight",
+        "model.layers.0.mlp.experts.0.up_proj.lora_B.weight",
+        "model.layers.0.mlp.experts.1.gate_proj.lora_B.weight",
+        "model.layers.0.mlp.experts.1.up_proj.lora_B.weight",
+    ]
+
+    # Shared side is unsqueezed to [1, ...] and identical for gate and up.
+    assert weights[0].weight.shape == (1, 2, 3)
+    torch.testing.assert_close(weights[0].weight, shared_lora_a.unsqueeze(0))
+    torch.testing.assert_close(weights[1].weight, shared_lora_a.unsqueeze(0))
+    # Per-expert side carries each expert's gate/up halves.
+    torch.testing.assert_close(weights[2].weight, torch.full((2, 2), 10.0))
+    torch.testing.assert_close(weights[3].weight, torch.full((2, 2), 20.0))
+    torch.testing.assert_close(weights[4].weight, torch.full((2, 2), 30.0))
+    torch.testing.assert_close(weights[5].weight, torch.full((2, 2), 40.0))
+
+
+def test_stream_adapter_weights_megatron_to_hf_shared_outer_fc2_down(monkeypatch):
+    # Shared-outer FC2 is the mirror of FC1: linear_in (lora_A) is the per-expert 3D
+    # pack and linear_out (lora_B) is the shared 2D matrix. down_proj is not fused, so
+    # the per-expert side emits one weight per expert and the shared side emits once.
+    bridge = DummyBridge()
+
+    adapter_task = AdapterWeightConversionTask(
+        global_base_prefix="decoder.layers.0.mlp.experts.linear_fc2",
+        adapter_key=None,
+        alpha=2,
+        dim=4,
+        linear_in_task=WeightConversionTask(
+            param_name="local_in",
+            global_param_name="decoder.layers.0.mlp.experts.linear_fc2.adapter.linear_in.weight",
+            mapping=Mock(),
+        ),
+        linear_out_task=WeightConversionTask(
+            param_name="local_out",
+            global_param_name="decoder.layers.0.mlp.experts.linear_fc2.adapter.linear_out.weight",
+            mapping=Mock(),
+        ),
+    )
+
+    # Per-expert lora_A: [num_experts=2, rank=2, inter=3]. Shared lora_B: [hidden=3, rank=2].
+    per_expert_lora_a = torch.stack([torch.full((2, 3), 5.0), torch.full((2, 3), 6.0)], dim=0)
+    shared_lora_b = torch.full((3, 2), 7.0)
+
+    adapter_weight = AdapterWeight(
+        global_base_prefix="decoder.layers.0.mlp.experts.linear_fc2",
+        adapter_key=None,
+        alpha=2,
+        dim=4,
+        linear_in_weight=MegatronWeightTuple("local_in", per_expert_lora_a, vp_stage=0),
+        linear_out_weight=MegatronWeightTuple("local_out", shared_lora_b, vp_stage=0),
+    )
+
+    def fake_base_names(_registry, _prefix, _adapter_key, base_suffix):
+        idx = base_suffix[len(".weight") :]
+        return [f"model.layers.0.mlp.experts.{idx}.down_proj.weight"]
+
+    monkeypatch.setattr(
+        bridge,
+        "build_adapter_conversion_tasks",
+        lambda *_args, **_kwargs: {"decoder.layers.0.mlp.experts.linear_fc2": [adapter_task]},
+    )
+    monkeypatch.setattr(bridge, "materialize_adapter_weights", lambda *_: [adapter_weight])
+    monkeypatch.setattr(bridge, "_get_base_hf_param_names_for_adapter", fake_base_names)
+    monkeypatch.setattr(
+        "megatron.bridge.models.conversion.peft_bridge.parallel_state.get_expert_model_parallel_world_size",
+        lambda: 1,
+    )
+
+    weights = list(
+        bridge.stream_adapter_weights_megatron_to_hf(
+            [SimpleNamespace(config=SimpleNamespace(num_moe_experts=2))],
+            cpu=False,
+            show_progress=False,
+        )
+    )
+
+    # Per-expert lora_A emitted once per expert (experts.<idx>), then the shared lora_B
+    # emitted once under the expert-agnostic name.
+    assert [w.param_name for w in weights] == [
+        "model.layers.0.mlp.experts.0.down_proj.lora_A.weight",
+        "model.layers.0.mlp.experts.1.down_proj.lora_A.weight",
+        "model.layers.0.mlp.experts.down_proj.lora_B.weight",
+    ]
+
+    torch.testing.assert_close(weights[0].weight, torch.full((2, 3), 5.0))
+    torch.testing.assert_close(weights[1].weight, torch.full((2, 3), 6.0))
+    # Shared side is unsqueezed to [1, hidden, rank].
+    assert weights[2].weight.shape == (1, 3, 2)
+    torch.testing.assert_close(weights[2].weight, shared_lora_b.unsqueeze(0))
+
+
 def test_split_gdn_in_proj_linear_out_weight_roundtrip(monkeypatch):
     bridge = DummyBridge()
     config = SimpleNamespace(
diff --git a/tests/unit_tests/peft/test_utils.py b/tests/unit_tests/peft/test_utils.py
index 250af3207c..4ebfb82cdc 100644
--- a/tests/unit_tests/peft/test_utils.py
+++ b/tests/unit_tests/peft/test_utils.py
@@ -20,12 +20,13 @@
 """
 
 import math
+from types import SimpleNamespace
 from unittest.mock import Mock, patch
 
 import pytest
 import torch
 import torch.nn as nn
-from megatron.core.dist_checkpointing.mapping import ShardedTensorFactory
+from megatron.core.dist_checkpointing.mapping import ShardedTensor, ShardedTensorFactory
 from megatron.core.tensor_parallel import ColumnParallelLinear, RowParallelLinear
 
 from megatron.bridge.peft import utils as peft_utils
@@ -33,6 +34,7 @@
     GroupedExpertLinearAdapter,
     ParallelLinearAdapter,
     all2all_hp2sp,
+    enable_legacy_shared_expert_adapter_loading,
     get_adapter_attributes_from_linear,
     init_method_const,
     init_method_kaiming_uniform,
@@ -68,6 +70,42 @@ def __init__(self):
         self.gradient_accumulation_fusion = False
 
 
+class MockProcessGroup:
+    """Small process-group stand-in with MCore-style size/rank methods."""
+
+    def __init__(self, size: int = 1, rank: int = 0):
+        self._size = size
+        self._rank = rank
+
+    def size(self) -> int:
+        return self._size
+
+    def rank(self) -> int:
+        return self._rank
+
+
+def make_mock_pg_collection(
+    *,
+    tp_size: int = 1,
+    tp_rank: int = 0,
+    ep_size: int = 1,
+    ep_rank: int = 0,
+    etp_size: int = 1,
+    etp_rank: int = 0,
+    edp_size: int = 1,
+    edp_rank: int = 0,
+) -> SimpleNamespace:
+    """Build the subset of ProcessGroupCollection used by PEFT tests."""
+
+    return SimpleNamespace(
+        tp=MockProcessGroup(tp_size, tp_rank),
+        ep=MockProcessGroup(ep_size, ep_rank),
+        expt_tp=MockProcessGroup(etp_size, etp_rank),
+        expt_dp=MockProcessGroup(edp_size, edp_rank),
+        dp_cp=MockProcessGroup(),
+    )
+
+
 class MockColumnParallelLinear(ColumnParallelLinear):
     """Mock ColumnParallelLinear for testing."""
 
@@ -249,15 +287,12 @@ def test_pad_unpad_roundtrip(self):
         assert torch.equal(unpadded, original)
 
 
-@patch("megatron.bridge.peft.utils.parallel_state")
 class TestAll2AllCommunication:
     """Test All2All communication functions."""
 
-    def test_all2all_hp2sp_mock(self, mock_parallel_state):
-        """Test all2all_hp2sp with mocked parallel state."""
-        # Mock parallel state
-        mock_parallel_state.get_tensor_model_parallel_world_size.return_value = 2
-        mock_parallel_state.get_tensor_model_parallel_group.return_value = None
+    def test_all2all_hp2sp_mock(self):
+        """Test all2all_hp2sp with an explicit tensor-parallel process group."""
+        tp_group = MockProcessGroup(size=2)
 
         # Mock torch.distributed.all_to_all
         with patch("torch.distributed.all_to_all") as mock_all_to_all:
@@ -270,7 +305,7 @@ def side_effect(receive_list, send_list, group):
             mock_all_to_all.side_effect = side_effect
 
             x = torch.randn(4, 8)  # Input tensor
-            result = all2all_hp2sp(x)
+            result = all2all_hp2sp(x, tp_group)
 
             assert result.shape == (2, 16)  # Should reshape appropriately
 
@@ -278,10 +313,8 @@ def side_effect(receive_list, send_list, group):
 class TestGetAdapterAttributes:
     """Test get_adapter_attributes_from_linear function."""
 
-    @patch("megatron.bridge.peft.utils.parallel_state")
-    def test_get_adapter_attributes_column_parallel(self, mock_parallel_state):
+    def test_get_adapter_attributes_column_parallel(self):
         """Test with ColumnParallelLinear."""
-        mock_parallel_state.get_tensor_model_parallel_world_size.return_value = 1
         linear = MockColumnParallelLinear(input_size=100, output_size=50)
 
         attrs = get_adapter_attributes_from_linear(linear)
@@ -293,10 +326,8 @@ def test_get_adapter_attributes_column_parallel(self, mock_parallel_state):
         assert attrs.disable_sequence_parallel_comm  # Should be True when sequence_parallel is False
         assert attrs.base_linear_is_parallel  # Should be True for parallel linear layers
 
-    @patch("megatron.bridge.peft.utils.parallel_state")
-    def test_get_adapter_attributes_row_parallel(self, mock_parallel_state):
+    def test_get_adapter_attributes_row_parallel(self):
         """Test with RowParallelLinear."""
-        mock_parallel_state.get_tensor_model_parallel_world_size.return_value = 1
         linear = MockRowParallelLinear(input_size=100, output_size=50)
 
         attrs = get_adapter_attributes_from_linear(linear)
@@ -308,10 +339,8 @@ def test_get_adapter_attributes_row_parallel(self, mock_parallel_state):
         assert attrs.disable_sequence_parallel_comm
         assert attrs.base_linear_is_parallel  # Should be True for parallel linear layers
 
-    @patch("megatron.bridge.peft.utils.parallel_state")
-    def test_get_adapter_attributes_sequence_parallel(self, mock_parallel_state):
+    def test_get_adapter_attributes_sequence_parallel(self):
         """Test with sequence parallel enabled."""
-        mock_parallel_state.get_tensor_model_parallel_world_size.return_value = 1
         linear = MockColumnParallelLinear(input_size=100, output_size=50)
         linear.config.sequence_parallel = True
 
@@ -321,20 +350,16 @@ def test_get_adapter_attributes_sequence_parallel(self, mock_parallel_state):
         assert not attrs.disable_sequence_parallel_comm  # Should be False when sequence_parallel is True
         assert attrs.base_linear_is_parallel  # Should be True for parallel linear layers
 
-    @patch("megatron.bridge.peft.utils.parallel_state")
-    def test_get_adapter_attributes_unsupported_module(self, mock_parallel_state):
+    def test_get_adapter_attributes_unsupported_module(self):
         """Test with unsupported module type."""
-        mock_parallel_state.get_tensor_model_parallel_world_size.return_value = 1
         linear = nn.Conv2d(3, 3, 3)
         linear.config = MockModelParallelConfig()
 
         with pytest.raises(NotImplementedError):
             get_adapter_attributes_from_linear(linear)
 
-    @patch("megatron.bridge.peft.utils.parallel_state")
-    def test_get_adapter_attributes_base_linear_is_parallel_flag(self, mock_parallel_state):
+    def test_get_adapter_attributes_base_linear_is_parallel_flag(self):
         """Test that base_linear_is_parallel flag is correctly returned."""
-        mock_parallel_state.get_tensor_model_parallel_world_size.return_value = 1
         # Test with ColumnParallelLinear - should return True for base_linear_is_parallel
         column_linear = MockColumnParallelLinear(input_size=100, output_size=50)
         assert get_adapter_attributes_from_linear(
@@ -573,16 +598,10 @@ def test_parallel_linear_adapter_forward_basic(self, mock_row_linear, mock_col_l
         expected_scale = adapter.alpha / adapter.dim
         assert expected_scale > 0
 
-    @patch("megatron.bridge.peft.utils.parallel_state")
     @patch("megatron.bridge.peft.utils.ColumnParallelLinear")
     @patch("megatron.bridge.peft.utils.RowParallelLinear")
-    def test_parallel_linear_adapter_expert_mode(
-        self, mock_row_linear, mock_col_linear, mock_parallel_state, mock_config
-    ):
+    def test_parallel_linear_adapter_expert_mode(self, mock_row_linear, mock_col_linear, mock_config):
         """Test adapter in expert mode (MoE)."""
-        # Mock parallel state for expert mode
-        mock_parallel_state.get_expert_tensor_parallel_world_size.return_value = 4
-
         # Set tensor_model_parallel_size to 4 so that sequence length 7 gets padded to 8
         mock_config.tensor_model_parallel_size = 4
         mock_config.expert_tensor_parallel_size = 4
@@ -656,6 +675,7 @@ def test_parallel_linear_adapter_sharded_state_dict_fc1_special_case(
 
         # Mock the swiglu factory
         mock_swiglu_factory.return_value = "swiglu_processed_tensor"
+        mock_config.gated_linear_unit = True
 
         adapter = ParallelLinearAdapter(
             in_features=20, out_features=10, dim=16, base_linear_name="linear_fc1", model_parallel_config=mock_config
@@ -667,6 +687,392 @@ def test_parallel_linear_adapter_sharded_state_dict_fc1_special_case(
         mock_swiglu_factory.assert_called()
         assert result["adapter.linear_out.weight"] == "swiglu_processed_tensor"
 
+    @patch("megatron.bridge.peft.utils.ColumnParallelLinear")
+    @patch("megatron.bridge.peft.utils.RowParallelLinear")
+    def test_parallel_linear_adapter_grouped_expert_sharded_state_dict_uses_expert_axis(
+        self, mock_row_linear, mock_col_linear, mock_config
+    ):
+        """Shared grouped-expert adapters should add a stable checkpoint expert axis."""
+        mock_linear_in = Mock()
+        mock_linear_out = Mock()
+        linear_in_weight = torch.arange(4, dtype=torch.float32).reshape(2, 2)
+        linear_out_weight = torch.arange(4, 8, dtype=torch.float32).reshape(2, 2)
+        mock_linear_in.sharded_state_dict.return_value = {
+            "adapter.linear_in.weight": ShardedTensor.from_rank_offsets(
+                "adapter.linear_in.weight", linear_in_weight, replica_id=(0, 0, 0)
+            ),
+            "adapter.linear_in._extra_state": torch.tensor([1.0]),
+        }
+        mock_linear_out.sharded_state_dict.return_value = {
+            "adapter.linear_out.weight": ShardedTensor.from_rank_offsets(
+                "adapter.linear_out.weight", linear_out_weight, replica_id=(0, 0, 0)
+            ),
+            "adapter.linear_out._extra_state": torch.tensor([2.0]),
+        }
+        mock_col_linear.side_effect = [mock_linear_in, mock_linear_out]
+        mock_config.num_moe_experts = 4
+        mock_config._pg_collection = make_mock_pg_collection(ep_size=2, ep_rank=1, etp_rank=0, edp_rank=3)
+
+        adapter = ParallelLinearAdapter(
+            in_features=2,
+            out_features=2,
+            dim=2,
+            base_linear_name="decoder.layers.0.mlp.experts.linear_fc2",
+            is_expert=True,
+            model_parallel_config=mock_config,
+        )
+        result = adapter.sharded_state_dict(prefix="adapter.")
+
+        assert "adapter.linear_in._extra_state" not in result
+        assert "adapter.linear_out._extra_state" not in result
+        factory = result["adapter.linear_in.weight"]
+        assert isinstance(factory, ShardedTensorFactory)
+        assert factory.replica_id == (0, 0, 3)
+
+        built = factory.build()
+        assert len(built) == 2
+        assert built[0].global_shape == (4, 2, 2)
+        assert built[0].global_offset == (2, 0, 0)
+        assert built[1].global_offset == (3, 0, 0)
+
+        merged = factory.merge_fn([torch.ones(2, 2), torch.full((2, 2), 3.0)])
+        torch.testing.assert_close(merged, torch.full((2, 2), 2.0))
+
+    @patch("megatron.bridge.peft.utils.ColumnParallelLinear")
+    @patch("megatron.bridge.peft.utils.RowParallelLinear")
+    def test_parallel_linear_adapter_grouped_expert_shared_adapter_syncs_grad_across_ep(
+        self, mock_row_linear, mock_col_linear, mock_config
+    ):
+        """Shared grouped-expert adapters must not drift across EP ranks.
+
+        A shared expert adapter is one logical weight used by every EP rank, but
+        MCore's expert DDP sync only covers expert-DP replicas. The EP grad hook
+        keeps EP>1 ranks from updating that shared adapter from different local
+        token subsets.
+        """
+        mock_linear_in = Mock()
+        mock_linear_out = Mock()
+        mock_linear_in.weight = nn.Parameter(torch.ones(2, 2))
+        mock_linear_out.weight = nn.Parameter(torch.ones(2, 2))
+        mock_col_linear.side_effect = [mock_linear_in, mock_linear_out]
+        mock_config._pg_collection = make_mock_pg_collection(ep_size=2)
+
+        ParallelLinearAdapter(
+            in_features=2,
+            out_features=2,
+            dim=2,
+            base_linear_name="decoder.layers.0.mlp.experts.linear_fc2",
+            is_expert=True,
+            model_parallel_config=mock_config,
+        )
+
+        ep_group = mock_config._pg_collection.ep
+
+        with (
+            patch("torch.distributed.is_available", return_value=True),
+            patch("torch.distributed.is_initialized", return_value=True),
+            patch("torch.distributed.all_reduce") as mock_all_reduce,
+        ):
+            (mock_linear_in.weight.sum() + mock_linear_out.weight.sum()).backward()
+
+        assert mock_all_reduce.call_count == 2
+        for call in mock_all_reduce.call_args_list:
+            grad = call.args[0]
+            torch.testing.assert_close(grad, torch.ones_like(grad))
+            assert call.kwargs["group"] is ep_group
+
+    @patch("megatron.bridge.peft.utils.ColumnParallelLinear")
+    @patch("megatron.bridge.peft.utils.RowParallelLinear")
+    def test_parallel_linear_adapter_grouped_expert_swiglu_sharded_state_dict_uses_expert_axis(
+        self, mock_row_linear, mock_col_linear, mock_config
+    ):
+        """Shared grouped-expert SwiGLU adapters should split gate/up shards inside the expert axis."""
+        mock_linear_in = Mock()
+        mock_linear_out = Mock()
+        linear_in_weight = torch.arange(4, dtype=torch.float32).reshape(2, 2)
+        linear_out_weight = torch.arange(8, dtype=torch.float32).reshape(4, 2)
+        mock_linear_in.sharded_state_dict.return_value = {
+            "adapter.linear_in.weight": ShardedTensor.from_rank_offsets(
+                "adapter.linear_in.weight", linear_in_weight, replica_id=(0, 0, 0)
+            ),
+            "adapter.linear_in._extra_state": torch.tensor([1.0]),
+        }
+        mock_linear_out.sharded_state_dict.return_value = {
+            "adapter.linear_out.weight": ShardedTensor.from_rank_offsets(
+                "adapter.linear_out.weight", linear_out_weight, replica_id=(0, 0, 0)
+            ),
+            "adapter.linear_out._extra_state": torch.tensor([2.0]),
+        }
+        mock_col_linear.side_effect = [mock_linear_in, mock_linear_out]
+        mock_config.num_moe_experts = 4
+        mock_config.gated_linear_unit = True
+        mock_config._pg_collection = make_mock_pg_collection(ep_size=2, ep_rank=1, etp_rank=0, edp_rank=2)
+
+        adapter = ParallelLinearAdapter(
+            in_features=2,
+            out_features=4,
+            dim=2,
+            base_linear_name="decoder.layers.0.mlp.experts.linear_fc1",
+            is_expert=True,
+            model_parallel_config=mock_config,
+        )
+        result = adapter.sharded_state_dict(prefix="adapter.")
+
+        assert "adapter.linear_in._extra_state" not in result
+        assert "adapter.linear_out._extra_state" not in result
+        factory = result["adapter.linear_out.weight"]
+        assert isinstance(factory, ShardedTensorFactory)
+        assert factory.replica_id == (0, 0, 2)
+
+        built = factory.build()
+        assert len(built) == 4
+        assert built[0].global_shape == (4, 4, 2)
+        assert [shard.global_offset for shard in built] == [
+            (2, 0, 0),
+            (2, 2, 0),
+            (3, 0, 0),
+            (3, 2, 0),
+        ]
+
+        merged = factory.merge_fn(
+            [
+                torch.full((2, 2), 1.0),
+                torch.full((2, 2), 2.0),
+                torch.full((2, 2), 3.0),
+                torch.full((2, 2), 5.0),
+            ]
+        )
+        expected = torch.cat([torch.full((2, 2), 2.0), torch.full((2, 2), 3.5)], dim=0)
+        torch.testing.assert_close(merged, expected)
+
+    @patch("megatron.bridge.peft.utils.ColumnParallelLinear")
+    @patch("megatron.bridge.peft.utils.RowParallelLinear")
+    def test_parallel_linear_adapter_grouped_expert_sharded_state_dict_keeps_extra_state_on_main_expert_rank(
+        self, mock_row_linear, mock_col_linear, mock_config
+    ):
+        """Shared grouped-expert adapter extra state should be kept only on EP0/ETP0."""
+        mock_linear_in = Mock()
+        mock_linear_out = Mock()
+        linear_in_weight = torch.arange(4, dtype=torch.float32).reshape(2, 2)
+        linear_out_weight = torch.arange(4, 8, dtype=torch.float32).reshape(2, 2)
+        linear_in_extra_state = torch.tensor([1.0])
+        linear_out_extra_state = torch.tensor([2.0])
+        mock_linear_in.sharded_state_dict.return_value = {
+            "adapter.linear_in.weight": ShardedTensor.from_rank_offsets(
+                "adapter.linear_in.weight", linear_in_weight, replica_id=(0, 0, 0)
+            ),
+            "adapter.linear_in._extra_state": linear_in_extra_state,
+        }
+        mock_linear_out.sharded_state_dict.return_value = {
+            "adapter.linear_out.weight": ShardedTensor.from_rank_offsets(
+                "adapter.linear_out.weight", linear_out_weight, replica_id=(0, 0, 0)
+            ),
+            "adapter.linear_out._extra_state": linear_out_extra_state,
+        }
+        mock_col_linear.side_effect = [mock_linear_in, mock_linear_out]
+        mock_config.num_moe_experts = 4
+        mock_config._pg_collection = make_mock_pg_collection(ep_size=2, ep_rank=0, etp_rank=0, edp_rank=0)
+
+        adapter = ParallelLinearAdapter(
+            in_features=2,
+            out_features=2,
+            dim=2,
+            base_linear_name="decoder.layers.0.mlp.experts.linear_fc2",
+            is_expert=True,
+            model_parallel_config=mock_config,
+        )
+        result = adapter.sharded_state_dict(prefix="adapter.")
+
+        assert result["adapter.linear_in._extra_state"] is linear_in_extra_state
+        assert result["adapter.linear_out._extra_state"] is linear_out_extra_state
+        built = result["adapter.linear_in.weight"].build()
+        assert [shard.global_offset for shard in built] == [(0, 0, 0), (1, 0, 0)]
+
+    @patch("megatron.bridge.peft.utils.ColumnParallelLinear")
+    @patch("megatron.bridge.peft.utils.RowParallelLinear")
+    def test_parallel_linear_adapter_legacy_shared_expert_state_dict_uses_2d_shape(
+        self, mock_row_linear, mock_col_linear, mock_config
+    ):
+        """Legacy shared grouped-expert adapter checkpoints should load as 2D tensors."""
+        mock_linear_in = Mock()
+        mock_linear_out = Mock()
+        linear_in_weight = torch.arange(4, dtype=torch.float32).reshape(2, 2)
+        linear_out_weight = torch.arange(4, 8, dtype=torch.float32).reshape(2, 2)
+        mock_linear_in.sharded_state_dict.return_value = {
+            "adapter.linear_in.weight": ShardedTensor.from_rank_offsets(
+                "adapter.linear_in.weight", linear_in_weight, replica_id=(0, 0, 0)
+            ),
+        }
+        mock_linear_out.sharded_state_dict.return_value = {
+            "adapter.linear_out.weight": ShardedTensor.from_rank_offsets(
+                "adapter.linear_out.weight", linear_out_weight, replica_id=(0, 0, 0)
+            ),
+        }
+        mock_col_linear.side_effect = [mock_linear_in, mock_linear_out]
+        mock_config.num_moe_experts = 4
+        mock_config._pg_collection = make_mock_pg_collection(ep_size=2, ep_rank=1, etp_rank=0, edp_rank=3)
+
+        adapter = ParallelLinearAdapter(
+            in_features=2,
+            out_features=2,
+            dim=2,
+            base_linear_name="decoder.layers.0.mlp.experts.linear_fc2",
+            is_expert=True,
+            model_parallel_config=mock_config,
+        )
+        adapter.use_legacy_shared_expert_adapter_checkpoint = True
+
+        result = adapter.sharded_state_dict(prefix="adapter.")
+
+        sharded_weight = result["adapter.linear_in.weight"]
+        assert isinstance(sharded_weight, ShardedTensor)
+        assert sharded_weight.global_shape == (2, 2)
+
+    @patch("megatron.bridge.peft.utils.ColumnParallelLinear")
+    @patch("megatron.bridge.peft.utils.RowParallelLinear")
+    def test_enable_legacy_shared_expert_adapter_loading_detects_2d_checkpoint_metadata(
+        self, mock_row_linear, mock_col_linear, mock_config, monkeypatch
+    ):
+        """A 2D checkpoint tensor should opt only its shared expert adapter into legacy loading."""
+        mock_linear_in = Mock()
+        mock_linear_out = Mock()
+        linear_in_weight = torch.arange(4, dtype=torch.float32).reshape(2, 2)
+        linear_out_weight = torch.arange(4, 8, dtype=torch.float32).reshape(2, 2)
+        mock_linear_in.sharded_state_dict.return_value = {
+            "decoder.layers.0.mlp.experts.linear_fc2.adapter.linear_in.weight": ShardedTensor.from_rank_offsets(
+                "decoder.layers.0.mlp.experts.linear_fc2.adapter.linear_in.weight",
+                linear_in_weight,
+                replica_id=(0, 0, 0),
+            ),
+        }
+        mock_linear_out.sharded_state_dict.return_value = {
+            "decoder.layers.0.mlp.experts.linear_fc2.adapter.linear_out.weight": ShardedTensor.from_rank_offsets(
+                "decoder.layers.0.mlp.experts.linear_fc2.adapter.linear_out.weight",
+                linear_out_weight,
+                replica_id=(0, 0, 0),
+            ),
+        }
+        mock_col_linear.side_effect = [mock_linear_in, mock_linear_out]
+        mock_config.num_moe_experts = 4
+        mock_config._pg_collection = make_mock_pg_collection(ep_size=2, ep_rank=0, etp_rank=0, edp_rank=0)
+
+        adapter = ParallelLinearAdapter(
+            in_features=2,
+            out_features=2,
+            dim=2,
+            base_linear_name="decoder.layers.0.mlp.experts.linear_fc2",
+            is_expert=True,
+            model_parallel_config=mock_config,
+        )
+        sharded_state_dict = {
+            "model": adapter.sharded_state_dict(prefix="decoder.layers.0.mlp.experts.linear_fc2.adapter.")
+        }
+        metadata = {
+            "decoder.layers.0.mlp.experts.linear_fc2.adapter.linear_in.weight": ShardedTensor.from_rank_offsets(
+                "decoder.layers.0.mlp.experts.linear_fc2.adapter.linear_in.weight",
+                torch.empty(2, 2),
+            ).without_data()
+        }
+        monkeypatch.setattr(peft_utils.dist_checkpointing, "load_tensors_metadata", lambda _: metadata)
+
+        enabled = enable_legacy_shared_expert_adapter_loading(
+            [SimpleNamespace(named_modules=lambda: [("decoder.layers.0.mlp.experts.linear_fc2.adapter", adapter)])],
+            sharded_state_dict,
+            "/checkpoint",
+        )
+
+        assert enabled is True
+        assert adapter.use_legacy_shared_expert_adapter_checkpoint is True
+
+    @patch("megatron.bridge.peft.utils.ColumnParallelLinear")
+    @patch("megatron.bridge.peft.utils.RowParallelLinear")
+    def test_enable_legacy_shared_expert_adapter_loading_tolerates_key_name_mismatch(
+        self, mock_row_linear, mock_col_linear, mock_config, monkeypatch
+    ):
+        """Legacy loading should still work when checkpoint keys and module names differ."""
+        mock_linear_in = Mock()
+        mock_linear_out = Mock()
+        linear_in_weight = torch.arange(4, dtype=torch.float32).reshape(2, 2)
+        linear_out_weight = torch.arange(4, 8, dtype=torch.float32).reshape(2, 2)
+        global_key = "decoder.layers.8.mlp.experts.linear_fc2.adapter.linear_in.weight"
+        mock_linear_in.sharded_state_dict.return_value = {
+            global_key: ShardedTensor.from_rank_offsets(global_key, linear_in_weight, replica_id=(0, 0, 0)),
+        }
+        mock_linear_out.sharded_state_dict.return_value = {
+            "decoder.layers.8.mlp.experts.linear_fc2.adapter.linear_out.weight": ShardedTensor.from_rank_offsets(
+                "decoder.layers.8.mlp.experts.linear_fc2.adapter.linear_out.weight",
+                linear_out_weight,
+                replica_id=(0, 0, 0),
+            ),
+        }
+        mock_col_linear.side_effect = [mock_linear_in, mock_linear_out]
+        mock_config.num_moe_experts = 4
+        mock_config._pg_collection = make_mock_pg_collection(ep_size=2, ep_rank=0, etp_rank=0, edp_rank=0)
+
+        adapter = ParallelLinearAdapter(
+            in_features=2,
+            out_features=2,
+            dim=2,
+            base_linear_name="decoder.layers.0.mlp.experts.linear_fc2",
+            is_expert=True,
+            model_parallel_config=mock_config,
+        )
+        sharded_state_dict = {
+            "model": adapter.sharded_state_dict(prefix="decoder.layers.8.mlp.experts.linear_fc2.adapter.")
+        }
+        metadata = {global_key: ShardedTensor.from_rank_offsets(global_key, torch.empty(2, 2)).without_data()}
+        monkeypatch.setattr(peft_utils.dist_checkpointing, "load_tensors_metadata", lambda _: metadata)
+
+        enabled = enable_legacy_shared_expert_adapter_loading(
+            [SimpleNamespace(named_modules=lambda: [("decoder.layers.0.mlp.experts.linear_fc2.adapter", adapter)])],
+            sharded_state_dict,
+            "/checkpoint",
+        )
+
+        assert enabled is True
+        assert adapter.use_legacy_shared_expert_adapter_checkpoint is True
+
+    @patch("megatron.bridge.peft.utils.ColumnParallelLinear")
+    @patch("megatron.bridge.peft.utils.RowParallelLinear")
+    def test_parallel_linear_adapter_non_grouped_expert_sharded_state_dict_uses_expert_dp_replica_id(
+        self, mock_row_linear, mock_col_linear, mock_config
+    ):
+        """Sequential local expert adapters should keep base sharding and use expert-DP replica ids."""
+        mock_linear_in = Mock()
+        mock_linear_out = Mock()
+        linear_in_weight = torch.arange(4, dtype=torch.float32).reshape(2, 2)
+        linear_out_weight = torch.arange(4, 8, dtype=torch.float32).reshape(2, 2)
+        mock_linear_in.sharded_state_dict.return_value = {
+            "adapter.linear_in.weight": ShardedTensor.from_rank_offsets(
+                "adapter.linear_in.weight", linear_in_weight, replica_id=(0, 1, 99)
+            ),
+            "adapter.linear_in._extra_state": torch.tensor([1.0]),
+        }
+        mock_linear_out.sharded_state_dict.return_value = {
+            "adapter.linear_out.weight": ShardedTensor.from_rank_offsets(
+                "adapter.linear_out.weight", linear_out_weight, replica_id=(0, 1, 99)
+            ),
+            "adapter.linear_out._extra_state": torch.tensor([2.0]),
+        }
+        mock_col_linear.side_effect = [mock_linear_in, mock_linear_out]
+        mock_config._pg_collection = make_mock_pg_collection(etp_rank=0, edp_rank=3)
+
+        adapter = ParallelLinearAdapter(
+            in_features=2,
+            out_features=2,
+            dim=2,
+            base_linear_name="decoder.layers.0.mlp.experts.local_experts.0.linear_fc2",
+            is_expert=True,
+            model_parallel_config=mock_config,
+        )
+        result = adapter.sharded_state_dict(prefix="adapter.")
+
+        sharded_weight = result["adapter.linear_in.weight"]
+        assert isinstance(sharded_weight, ShardedTensor)
+        assert sharded_weight.global_shape == (2, 2)
+        assert sharded_weight.replica_id == (0, 1, 3)
+        assert "adapter.linear_in._extra_state" in result
+
 
 class TestGroupedExpertLinearAdapter:
     """Tests for grouped-expert per-expert LoRA adapters."""
@@ -746,6 +1152,109 @@ def test_grouped_expert_linear_adapter_forward_uses_per_expert_weights(self):
         )
         torch.testing.assert_close(output, expected)
 
+    def test_grouped_expert_linear_adapter_keeps_checkpoint_keys_after_weight_module_wrap(self):
+        """Calling weight containers should not change existing adapter checkpoint keys."""
+        adapter = GroupedExpertLinearAdapter(
+            in_features=2,
+            out_features=2,
+            dim=2,
+            num_local_experts=2,
+            base_linear_name="decoder.layers.0.mlp.experts.linear_fc2",
+            activation="identity",
+            input_is_parallel=False,
+            model_parallel_config=MockModelParallelConfig(),
+        )
+
+        state_dict = adapter.state_dict()
+
+        assert sorted(state_dict) == ["linear_in.weight", "linear_out.weight"]
+        old_style_state_dict = {
+            "linear_in.weight": torch.ones_like(adapter.linear_in.weight),
+            "linear_out.weight": torch.full_like(adapter.linear_out.weight, 2.0),
+        }
+        missing, unexpected = adapter.load_state_dict(old_style_state_dict, strict=True)
+        assert missing == []
+        assert unexpected == []
+        torch.testing.assert_close(adapter.linear_in.weight, old_style_state_dict["linear_in.weight"])
+        torch.testing.assert_close(adapter.linear_out.weight, old_style_state_dict["linear_out.weight"])
+
+    def test_grouped_expert_linear_adapter_forward_calls_weight_modules_for_param_sync_hooks(self):
+        """Grouped per-expert adapters must participate in training-time param gather.
+
+        With EP plus expert-DP, distributed optimizer param gather is driven by
+        DDP forward pre-hooks. The weight containers need to be called so live
+        training weights refresh during normal forwards, not only at forced
+        eval/checkpoint sync boundaries.
+        """
+        adapter = GroupedExpertLinearAdapter(
+            in_features=2,
+            out_features=2,
+            dim=2,
+            num_local_experts=3,
+            base_linear_name="decoder.layers.0.mlp.experts.linear_fc2",
+            activation="identity",
+            input_is_parallel=False,
+            model_parallel_config=MockModelParallelConfig(),
+        )
+
+        calls = []
+        adapter.linear_in.register_forward_pre_hook(lambda module, inputs: calls.append("linear_in"))
+        adapter.linear_out.register_forward_pre_hook(lambda module, inputs: calls.append("linear_out"))
+
+        def fake_grouped_mm(inputs, weights, *, offs):
+            chunks = []
+            start = 0
+            for weight_idx, end in enumerate(offs.tolist()):
+                chunks.append(inputs[start:end] @ weights[weight_idx])
+                start = end
+            return torch.cat(chunks, dim=0)
+
+        x = torch.tensor(
+            [
+                [1.0, 2.0],
+                [3.0, 4.0],
+                [5.0, 6.0],
+            ]
+        )
+        with (
+            patch.object(GroupedExpertLinearAdapter, "_can_use_grouped_mm", return_value=True),
+            patch(
+                "megatron.bridge.peft.utils.nn.functional.grouped_mm",
+                side_effect=fake_grouped_mm,
+                create=True,
+            ),
+        ):
+            adapter(x, [1, 0, 2])
+
+        assert calls == ["linear_in", "linear_out"]
+
+    def test_grouped_expert_linear_adapter_zero_token_batch_keeps_weight_grad_dependency(self):
+        """Empty local expert batches should still produce zero grads for DDP hooks.
+
+        EP routing can leave a local grouped adapter with no tokens on a step.
+        The zero-sized output still needs a zero-valued dependency on the LoRA
+        weights so DDP sees ready gradients instead of leaving replicas stale.
+        """
+        adapter = GroupedExpertLinearAdapter(
+            in_features=2,
+            out_features=2,
+            dim=2,
+            num_local_experts=2,
+            base_linear_name="decoder.layers.0.mlp.experts.linear_fc2",
+            activation="identity",
+            input_is_parallel=False,
+            model_parallel_config=MockModelParallelConfig(),
+        )
+
+        output = adapter(torch.empty(0, 2), [0, 0])
+
+        assert output.shape == (0, 2)
+        output.sum().backward()
+        assert adapter.linear_in.weight.grad is not None
+        assert adapter.linear_out.weight.grad is not None
+        torch.testing.assert_close(adapter.linear_in.weight.grad, torch.zeros_like(adapter.linear_in.weight))
+        torch.testing.assert_close(adapter.linear_out.weight.grad, torch.zeros_like(adapter.linear_out.weight))
+
     def test_grouped_expert_linear_adapter_grouped_mm_falls_back_on_cpu(self):
         """CPU inputs should not enter the grouped_mm fast path."""
         adapter = GroupedExpertLinearAdapter(
@@ -935,17 +1444,7 @@ def test_grouped_expert_linear_adapter_requires_expert_tp_group_for_gather(self)
             model_parallel_config=config,
         )
 
-        with (
-            patch(
-                "megatron.bridge.peft.utils.parallel_state.get_expert_tensor_parallel_world_size",
-                return_value=None,
-            ),
-            patch(
-                "megatron.bridge.peft.utils.parallel_state.get_expert_tensor_parallel_group",
-                return_value=None,
-            ),
-            patch("megatron.bridge.peft.utils.torch.distributed.all_gather") as mock_all_gather,
-        ):
+        with patch("megatron.bridge.peft.utils.torch.distributed.all_gather") as mock_all_gather:
             with pytest.raises(
                 ValueError,
                 match="requires initialized expert tensor parallel state when expert_tensor_parallel_size=2",
@@ -956,36 +1455,21 @@ def test_grouped_expert_linear_adapter_requires_expert_tp_group_for_gather(self)
 
     def test_grouped_expert_linear_fc1_sharded_state_dict_preserves_expert_axis(self):
         """Grouped expert fc1 checkpoints should split SwiGLU on the hidden axis, not the expert axis."""
-        with (
-            patch(
-                "megatron.bridge.peft.utils.parallel_state.get_expert_model_parallel_world_size",
-                return_value=2,
-            ),
-            patch(
-                "megatron.bridge.peft.utils.parallel_state.get_expert_model_parallel_rank",
-                return_value=1,
-            ),
-            patch(
-                "megatron.bridge.peft.utils.parallel_state.get_expert_data_parallel_rank",
-                return_value=0,
-            ),
-            patch(
-                "megatron.bridge.peft.utils.parallel_state.get_expert_tensor_parallel_world_size",
-                return_value=1,
-            ),
-        ):
-            adapter = GroupedExpertLinearAdapter(
-                in_features=2,
-                out_features=4,
-                dim=2,
-                num_local_experts=2,
-                base_linear_name="decoder.layers.0.mlp.experts.linear_fc1",
-                activation="identity",
-                input_is_parallel=False,
-                model_parallel_config=MockModelParallelConfig(),
-            )
+        config = MockModelParallelConfig()
+        config.gated_linear_unit = True
+        config._pg_collection = make_mock_pg_collection(ep_size=2, ep_rank=1, edp_rank=0, etp_size=1)
+        adapter = GroupedExpertLinearAdapter(
+            in_features=2,
+            out_features=4,
+            dim=2,
+            num_local_experts=2,
+            base_linear_name="decoder.layers.0.mlp.experts.linear_fc1",
+            activation="identity",
+            input_is_parallel=False,
+            model_parallel_config=config,
+        )
 
-            result = adapter.sharded_state_dict("adapter.")
+        result = adapter.sharded_state_dict("adapter.")
 
         factory = result["adapter.linear_out.weight"]
         assert isinstance(factory, ShardedTensorFactory)
@@ -1001,40 +1485,21 @@ def test_grouped_expert_linear_fc1_sharded_state_dict_preserves_expert_axis(self
 
     def test_grouped_expert_linear_fc1_factory_merge_restores_gate_up_order(self):
         """Grouped expert fc1 checkpoint reload should de-interleave gate/up expert-TP shards."""
-        with (
-            patch(
-                "megatron.bridge.peft.utils.parallel_state.get_expert_model_parallel_world_size",
-                return_value=1,
-            ),
-            patch(
-                "megatron.bridge.peft.utils.parallel_state.get_expert_model_parallel_rank",
-                return_value=0,
-            ),
-            patch(
-                "megatron.bridge.peft.utils.parallel_state.get_expert_data_parallel_rank",
-                return_value=0,
-            ),
-            patch(
-                "megatron.bridge.peft.utils.parallel_state.get_expert_tensor_parallel_world_size",
-                return_value=2,
-            ),
-            patch(
-                "megatron.bridge.peft.utils.parallel_state.get_expert_tensor_parallel_rank",
-                return_value=0,
-            ),
-        ):
-            adapter = GroupedExpertLinearAdapter(
-                in_features=2,
-                out_features=8,
-                dim=2,
-                num_local_experts=1,
-                base_linear_name="decoder.layers.0.mlp.experts.linear_fc1",
-                activation="identity",
-                input_is_parallel=False,
-                model_parallel_config=MockModelParallelConfig(),
-            )
+        config = MockModelParallelConfig()
+        config.gated_linear_unit = True
+        config._pg_collection = make_mock_pg_collection(ep_size=1, ep_rank=0, edp_rank=0, etp_size=2, etp_rank=0)
+        adapter = GroupedExpertLinearAdapter(
+            in_features=2,
+            out_features=8,
+            dim=2,
+            num_local_experts=1,
+            base_linear_name="decoder.layers.0.mlp.experts.linear_fc1",
+            activation="identity",
+            input_is_parallel=False,
+            model_parallel_config=config,
+        )
 
-            factory = adapter.sharded_state_dict("adapter.")["adapter.linear_out.weight"]
+        factory = adapter.sharded_state_dict("adapter.")["adapter.linear_out.weight"]
 
         fused_tp0 = torch.tensor([[[1.0, 1.0], [1.0, 1.0], [2.0, 2.0], [2.0, 2.0]]])
         fused_tp1 = torch.tensor([[[3.0, 3.0], [3.0, 3.0], [4.0, 4.0], [4.0, 4.0]]])
@@ -1045,41 +1510,79 @@ def test_grouped_expert_linear_fc1_factory_merge_restores_gate_up_order(self):
         )
         torch.testing.assert_close(merged, expected)
 
+    @pytest.mark.parametrize(("ep_size", "expected_allreduce"), [(1, True), (2, False)])
+    def test_grouped_expert_linear_adapter_allreduce_flag_tracks_expert_parallelism(self, ep_size, expected_allreduce):
+        """Per-expert grouped adapters should use expert-DP grad sync only when EP is enabled."""
+        config = MockModelParallelConfig()
+        config._pg_collection = make_mock_pg_collection(ep_size=ep_size, etp_size=1)
+        adapter = GroupedExpertLinearAdapter(
+            in_features=2,
+            out_features=2,
+            dim=2,
+            num_local_experts=2,
+            base_linear_name="decoder.layers.0.mlp.experts.linear_fc2",
+            activation="identity",
+            input_is_parallel=False,
+            model_parallel_config=config,
+        )
+
+        assert adapter.linear_in.weight.allreduce is expected_allreduce
+        assert adapter.linear_out.weight.allreduce is expected_allreduce
+        assert adapter.linear_in.weight.tensor_model_parallel is True
+        assert adapter.linear_out.weight.tensor_model_parallel is True
+        assert adapter.linear_in.weight.partition_dim == 1
+        assert adapter.linear_out.weight.partition_dim == 1
+
+    def test_grouped_expert_linear_adapter_groups_as_expert_ddp_buffer_when_ep_enabled(self):
+        """Per-expert adapter params must sync on expert-DP, not dense DP.
+
+        EP plus DP replicates each local expert across expert-DP ranks. Marking
+        these params as expert-parallel keeps replicas for the same expert in
+        sync without mixing different EP-owned experts.
+        """
+        from megatron.core.distributed.param_and_grad_buffer import group_params_for_buffers
+
+        config = MockModelParallelConfig()
+        config._pg_collection = make_mock_pg_collection(ep_size=8, etp_size=1)
+        adapter = GroupedExpertLinearAdapter(
+            in_features=2,
+            out_features=2,
+            dim=2,
+            num_local_experts=2,
+            base_linear_name="decoder.layers.0.mlp.experts.linear_fc2",
+            activation="identity",
+            input_is_parallel=False,
+            model_parallel_config=config,
+        )
+
+        buffer_groups = group_params_for_buffers(
+            [adapter.linear_in.weight, adapter.linear_out.weight],
+            grad_reduce_in_fp32=False,
+        )
+
+        assert len(buffer_groups) == 1
+        buffer_key, (params, _param_indices) = next(iter(buffer_groups.items()))
+        assert buffer_key.is_expert_parallel
+        assert [id(param) for param in params] == [
+            id(adapter.linear_in.weight),
+            id(adapter.linear_out.weight),
+        ]
+
     def test_grouped_expert_linear_sharded_state_dict_uses_expert_parallel_offsets(self):
         """Grouped-expert weights should shard only across expert EP/ETP and use expert-DP replica ids."""
-        with (
-            patch(
-                "megatron.bridge.peft.utils.parallel_state.get_expert_model_parallel_world_size",
-                return_value=2,
-            ),
-            patch(
-                "megatron.bridge.peft.utils.parallel_state.get_expert_model_parallel_rank",
-                return_value=1,
-            ),
-            patch(
-                "megatron.bridge.peft.utils.parallel_state.get_expert_data_parallel_rank",
-                return_value=4,
-            ),
-            patch(
-                "megatron.bridge.peft.utils.parallel_state.get_expert_tensor_parallel_world_size",
-                return_value=1,
-            ),
-            patch(
-                "megatron.bridge.peft.utils.parallel_state.get_expert_tensor_parallel_rank",
-                return_value=0,
-            ),
-        ):
-            adapter = GroupedExpertLinearAdapter(
-                in_features=2,
-                out_features=2,
-                dim=2,
-                num_local_experts=2,
-                base_linear_name="decoder.layers.0.mlp.experts.linear_fc2",
-                activation="identity",
-                input_is_parallel=False,
-                model_parallel_config=MockModelParallelConfig(),
-            )
-            result = adapter.sharded_state_dict("adapter.")
+        config = MockModelParallelConfig()
+        config._pg_collection = make_mock_pg_collection(ep_size=2, ep_rank=1, edp_rank=4, etp_size=1, etp_rank=0)
+        adapter = GroupedExpertLinearAdapter(
+            in_features=2,
+            out_features=2,
+            dim=2,
+            num_local_experts=2,
+            base_linear_name="decoder.layers.0.mlp.experts.linear_fc2",
+            activation="identity",
+            input_is_parallel=False,
+            model_parallel_config=config,
+        )
+        result = adapter.sharded_state_dict("adapter.")
 
         sharded_weight = result["adapter.linear_in.weight"]
         assert sharded_weight.local_shape == (2, 2, 2)

From 19cbc75213c4dc9c7e889e02b79074ab8d51f444 Mon Sep 17 00:00:00 2001
From: Chen Cui <chcui@nvidia.com>
Date: Wed, 10 Jun 2026 13:26:45 -0700
Subject: [PATCH 2/5] fix(checkpointing): include optimizer scaffold while
 loading (#4222)

Signed-off-by: Chen Cui <chcui@nvidia.com>
(cherry picked from commit 0917258dec36df1acf61f3f3c3c34c18ec92c042)
---
 src/megatron/bridge/training/checkpointing.py |  6 ++--
 .../unit_tests/training/test_checkpointing.py | 32 +++++++++++++++++++
 .../training/utils/test_checkpoint_utils.py   |  9 +++---
 3 files changed, 41 insertions(+), 6 deletions(-)

diff --git a/src/megatron/bridge/training/checkpointing.py b/src/megatron/bridge/training/checkpointing.py
index 056cdb2c56..10f6765a9e 100644
--- a/src/megatron/bridge/training/checkpointing.py
+++ b/src/megatron/bridge/training/checkpointing.py
@@ -1580,8 +1580,10 @@ def generate_state_dict(
         _generate_model_state_dict(model, model_sd_kwargs, ckpt_cfg.ckpt_format, pg_collection=pg_collection)
     )
 
-    # Optimizer stuff.
-    if ckpt_cfg.save_optim:
+    # Optimizer stuff. During load, optimizer sharded-state scaffolding is
+    # required even if the next checkpoint should not save optimizer state.
+    include_optimizer_state = ckpt_cfg.save_optim or bool((optim_sd_kwargs or {}).get("is_loading"))
+    if include_optimizer_state:
         if optimizer is not None and not getattr(optimizer, "is_stub_optimizer", False):
             if ckpt_cfg.ckpt_format == "torch_dist":
                 state_dict["optimizer"] = optimizer.sharded_state_dict(state_dict, **(optim_sd_kwargs or {}))
diff --git a/tests/unit_tests/training/test_checkpointing.py b/tests/unit_tests/training/test_checkpointing.py
index 3e6397f8e4..9c60c35dca 100644
--- a/tests/unit_tests/training/test_checkpointing.py
+++ b/tests/unit_tests/training/test_checkpointing.py
@@ -2696,6 +2696,38 @@ def test_generate_state_dict_torch_dist_no_preprocessing(self):
             mock_model.sharded_state_dict.assert_called_once()
             assert "model" in result
 
+    def test_generate_state_dict_includes_optimizer_scaffold_when_loading(self):
+        """Load-time optimizer scaffolding should not depend on save_optim."""
+        from unittest.mock import Mock
+
+        from megatron.bridge.training.checkpointing import generate_state_dict
+        from megatron.bridge.training.config import CheckpointConfig
+
+        mock_model = Mock()
+        mock_model.sharded_state_dict.return_value = {"test_param": torch.tensor([1.0])}
+        mock_optimizer = Mock()
+        mock_optimizer.is_stub_optimizer = False
+        mock_optimizer.sharded_state_dict.return_value = {"optimizer": {"param_groups": []}}
+        mock_scheduler = Mock()
+        mock_scheduler.state_dict.return_value = {"scheduler": "state"}
+
+        ckpt_cfg = CheckpointConfig(ckpt_format="torch_dist", save_optim=False, save_rng=False)
+        result = generate_state_dict(
+            ckpt_cfg=ckpt_cfg,
+            model=[mock_model],
+            optimizer=mock_optimizer,
+            opt_param_scheduler=mock_scheduler,
+            rng_state=None,
+            optim_sd_kwargs={
+                "is_loading": True,
+                "metadata": {"distrib_optim_sharding_type": "dp_zero_gather_scatter"},
+            },
+        )
+
+        assert result["optimizer"] == {"optimizer": {"param_groups": []}}
+        assert result["opt_param_scheduler"] == {"scheduler": "state"}
+        mock_optimizer.sharded_state_dict.assert_called_once()
+
 
 class TestCheckpointPathOverride:
     """Test checkpoint_path_override parameter in loading functions."""
diff --git a/tests/unit_tests/training/utils/test_checkpoint_utils.py b/tests/unit_tests/training/utils/test_checkpoint_utils.py
index 8495f52524..07df2b9162 100644
--- a/tests/unit_tests/training/utils/test_checkpoint_utils.py
+++ b/tests/unit_tests/training/utils/test_checkpoint_utils.py
@@ -468,10 +468,12 @@ def test_checkpoint_exists_edge_cases(self, checkpoint_path, expected):
 
     # ===== ADVANCED TEST SCENARIOS =====
 
-    def test_concurrent_access_to_cached_functions(self):
+    def test_concurrent_access_to_cached_functions(self, tmp_path):
         """Test concurrent access to cached functions for thread safety."""
         config_data = {"model": {"type": "concurrent_test"}}
-        config_yaml = yaml.dump(config_data)
+        config_file = tmp_path / "concurrent_config.yaml"
+        with open(config_file, "w") as f:
+            yaml.dump(config_data, f)
 
         results = []
         errors = []
@@ -484,9 +486,8 @@ def read_config_worker():
                         "megatron.bridge.training.utils.checkpoint_utils.torch.distributed.is_initialized",
                         return_value=False,
                     ),
-                    patch("builtins.open", mock_open(read_data=config_yaml)),
                 ):
-                    result = read_run_config("concurrent_config.yaml")
+                    result = read_run_config(str(config_file))
                     results.append(result)
             except Exception as e:
                 errors.append(e)

From cb10e4d1c9061d7b285421d8fad1ad4d94d1922f Mon Sep 17 00:00:00 2001
From: Chen Cui <chcui@nvidia.com>
Date: Wed, 10 Jun 2026 13:27:25 -0700
Subject: [PATCH 3/5] feat(scripts): add Ultra script prerequisites (#4223)

Signed-off-by: Chen Cui <chcui@nvidia.com>
(cherry picked from commit 1d0dab636ee44d6f16f8f73b3cce51b9565de7d7)
---
 scripts/inference/text_generation.py          |  21 +-
 scripts/training/pack_sft_data.py             |  85 ++++++-
 .../unit_tests/scripts/test_pack_sft_data.py  | 220 ++++++++++++++++++
 .../scripts/test_text_generation.py           |  55 +++++
 4 files changed, 372 insertions(+), 9 deletions(-)
 create mode 100644 tests/unit_tests/scripts/test_pack_sft_data.py

diff --git a/scripts/inference/text_generation.py b/scripts/inference/text_generation.py
index f2e12e1745..d2f7f1a1fb 100644
--- a/scripts/inference/text_generation.py
+++ b/scripts/inference/text_generation.py
@@ -46,7 +46,15 @@
 from megatron.bridge import AutoBridge
 from megatron.bridge.models.hf_pretrained.utils import is_safe_repo
 from megatron.bridge.training.utils.checkpoint_utils import get_hf_model_id_from_checkpoint
-from megatron.bridge.utils.common_utils import disable_mtp_for_inference, get_local_rank_preinit, print_rank_0
+from megatron.bridge.utils.common_utils import (
+    disable_mtp_for_inference,
+    get_local_rank_preinit,
+    get_master_addr_safe,
+    get_master_port_safe,
+    get_rank_safe,
+    get_world_size_safe,
+    print_rank_0,
+)
 
 
 logger = logging.getLogger(__name__)
@@ -258,11 +266,12 @@ def _maybe_initialize_distributed(timeout_minutes: int) -> None:
     if not dist.is_available() or dist.is_initialized():
         return
 
-    os.environ["RANK"] = os.environ.get("RANK", "0")
-    os.environ["WORLD_SIZE"] = os.environ.get("WORLD_SIZE", "1")
-    os.environ["MASTER_ADDR"] = os.environ.get("MASTER_ADDR", "localhost")
-    os.environ["MASTER_PORT"] = os.environ.get("MASTER_PORT", "12355")
-    torch.cuda.set_device(get_local_rank_preinit())
+    os.environ["RANK"] = os.environ.get("RANK", str(get_rank_safe()))
+    os.environ["WORLD_SIZE"] = os.environ.get("WORLD_SIZE", str(get_world_size_safe()))
+    os.environ["LOCAL_RANK"] = os.environ.get("LOCAL_RANK", str(get_local_rank_preinit()))
+    os.environ["MASTER_ADDR"] = os.environ.get("MASTER_ADDR", get_master_addr_safe())
+    os.environ["MASTER_PORT"] = os.environ.get("MASTER_PORT", str(get_master_port_safe()))
+    torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
     dist.init_process_group("nccl", timeout=timedelta(minutes=timeout_minutes))
 
 
diff --git a/scripts/training/pack_sft_data.py b/scripts/training/pack_sft_data.py
index 525260dfd3..978e716e48 100644
--- a/scripts/training/pack_sft_data.py
+++ b/scripts/training/pack_sft_data.py
@@ -28,8 +28,10 @@
 """
 
 import argparse
+import inspect
 import sys
 from dataclasses import fields
+from pathlib import Path
 
 
 def main() -> None:
@@ -38,13 +40,29 @@ def main() -> None:
     parser.add_argument(
         "--recipe",
         required=True,
-        help="Recipe name, e.g. gpt_oss_20b_sft_openmathinstruct2_thinking_packed_config",
+        help="Recipe name for a packed-sequence SFT dataset config.",
     )
+    parser.add_argument("--seq-length", type=int, default=None, help="Optional sequence length override.")
+    parser.add_argument("--hf-path", default=None, help="Optional Hugging Face model ID or local snapshot path.")
+    parser.add_argument(
+        "--train-input-path", default=None, help="Optional processed JSONL path for the training split."
+    )
+    parser.add_argument(
+        "--val-input-path", default=None, help="Optional processed JSONL path for the validation split."
+    )
+    parser.add_argument(
+        "--packed-train-data-path", default=None, help="Optional output path for packed training data."
+    )
+    parser.add_argument(
+        "--packed-val-data-path", default=None, help="Optional output path for packed validation data."
+    )
+    parser.add_argument("--packed-metadata-path", default=None, help="Optional output path for packing metadata.")
     args = parser.parse_args()
 
     import megatron.bridge.recipes as all_recipes
     from megatron.bridge.data.builders.finetuning_dataset import FinetuningDatasetBuilder
     from megatron.bridge.data.builders.hf_dataset import HFDatasetBuilder, HFDatasetConfig
+    from megatron.bridge.data.datasets.packed_sequence import prepare_packed_sequence_data
     from megatron.bridge.training.config import DataloaderConfig
     from megatron.bridge.training.tokenizers.tokenizer import build_tokenizer
 
@@ -52,7 +70,22 @@ def main() -> None:
     if recipe_fn is None:
         sys.exit(f"Error: recipe '{args.recipe}' not found. Check available recipes in megatron.bridge.recipes.")
 
-    cfg = recipe_fn()
+    sig = inspect.signature(recipe_fn)
+    params = sig.parameters
+    has_var_keyword = any(param.kind == inspect.Parameter.VAR_KEYWORD for param in params.values())
+    recipe_kwargs = {}
+    if args.seq_length is not None:
+        if "seq_length" in params or has_var_keyword:
+            recipe_kwargs["seq_length"] = args.seq_length
+        else:
+            sys.exit(f"Error: recipe '{args.recipe}' does not accept a 'seq_length' parameter.")
+    if args.hf_path is not None:
+        if "hf_path" in params or has_var_keyword:
+            recipe_kwargs["hf_path"] = args.hf_path
+        else:
+            sys.exit(f"Error: recipe '{args.recipe}' does not accept an 'hf_path' parameter.")
+
+    cfg = recipe_fn(**recipe_kwargs)
 
     if cfg.dataset is None:
         sys.exit("Error: recipe has no dataset configuration.")
@@ -87,7 +120,53 @@ def main() -> None:
         },
     )
 
-    # For HF datasets, download + apply process_example_fn → training.jsonl. The
+    custom_pack_paths = [
+        args.train_input_path,
+        args.val_input_path,
+        args.packed_train_data_path,
+        args.packed_val_data_path,
+        args.packed_metadata_path,
+    ]
+    if any(custom_pack_paths):
+        if not args.train_input_path:
+            sys.exit("Error: --train-input-path is required when using explicit pack paths.")
+        if not args.packed_train_data_path:
+            sys.exit("Error: --packed-train-data-path is required when using explicit pack paths.")
+
+        packed_metadata_path = Path(args.packed_metadata_path) if args.packed_metadata_path else builder.pack_metadata
+        prepare_packed_sequence_data(
+            input_path=Path(args.train_input_path),
+            output_path=Path(args.packed_train_data_path),
+            output_metadata_path=packed_metadata_path,
+            packed_sequence_size=cfg.dataset.packed_sequence_specs.packed_sequence_size,
+            tokenizer=tokenizer,
+            max_seq_length=cfg.dataset.seq_length,
+            seed=cfg.dataset.seed,
+            dataset_kwargs=cfg.dataset.dataset_kwargs,
+            pad_seq_to_mult=cfg.dataset.packed_sequence_specs.pad_seq_to_mult,
+            num_tokenizer_workers=cfg.dataset.packed_sequence_specs.num_tokenizer_workers,
+        )
+
+        if args.val_input_path and args.packed_val_data_path:
+            prepare_packed_sequence_data(
+                input_path=Path(args.val_input_path),
+                output_path=Path(args.packed_val_data_path),
+                output_metadata_path=packed_metadata_path,
+                packed_sequence_size=cfg.dataset.packed_sequence_specs.packed_sequence_size,
+                tokenizer=tokenizer,
+                max_seq_length=cfg.dataset.seq_length,
+                seed=cfg.dataset.seed,
+                dataset_kwargs=cfg.dataset.dataset_kwargs,
+                pad_seq_to_mult=cfg.dataset.packed_sequence_specs.pad_seq_to_mult,
+                num_tokenizer_workers=cfg.dataset.packed_sequence_specs.num_tokenizer_workers,
+            )
+        elif args.val_input_path or args.packed_val_data_path:
+            sys.exit("Error: --val-input-path and --packed-val-data-path must be provided together.")
+
+        print("Done.")
+        return
+
+    # For HF datasets, download + apply process_example_fn -> training.jsonl. The
     # packer in prepare_packed_data() reads that JSONL; without this step the
     # cache directory is empty and packing fails with FileNotFoundError.
     if isinstance(builder, HFDatasetBuilder):
diff --git a/tests/unit_tests/scripts/test_pack_sft_data.py b/tests/unit_tests/scripts/test_pack_sft_data.py
new file mode 100644
index 0000000000..1683868324
--- /dev/null
+++ b/tests/unit_tests/scripts/test_pack_sft_data.py
@@ -0,0 +1,220 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unit tests for ``scripts/training/pack_sft_data.py``."""
+
+from __future__ import annotations
+
+import importlib.util
+import sys
+import types
+from dataclasses import dataclass, field
+from pathlib import Path
+from unittest.mock import Mock
+
+import pytest
+
+
+_REPO_ROOT = Path(__file__).resolve().parents[3]
+_SCRIPT_PATH = _REPO_ROOT / "scripts" / "training" / "pack_sft_data.py"
+
+
+@dataclass
+class _DataloaderConfig:
+    micro_batch_size: int = 1
+
+
+@dataclass
+class _PackedSequenceSpecs:
+    packed_sequence_size: int = 2048
+    pad_seq_to_mult: int = 8
+    num_tokenizer_workers: int = -1
+
+
+@dataclass
+class _DatasetConfig:
+    seq_length: int = 2048
+    seed: int = 123
+    dataset_kwargs: dict[str, object] = field(default_factory=lambda: {"chat": "template"})
+    packed_sequence_specs: _PackedSequenceSpecs | None = field(default_factory=_PackedSequenceSpecs)
+
+
+@dataclass
+class _RecipeConfig:
+    dataset: _DatasetConfig | None = field(default_factory=_DatasetConfig)
+    tokenizer: object = "tokenizer-config"
+
+
+class _FinetuningDatasetBuilder:
+    pack_metadata = Path("default-pack-metadata.json")
+
+    def __init__(self, *, tokenizer: object, **kwargs: object) -> None:
+        self.tokenizer = tokenizer
+        self.kwargs = kwargs
+
+    def prepare_packed_data(self) -> None:
+        raise AssertionError("explicit pack-path tests should not call prepare_packed_data")
+
+
+class _HFDatasetConfig:
+    pass
+
+
+class _HFDatasetBuilder(_FinetuningDatasetBuilder):
+    def prepare_data(self) -> None:
+        raise AssertionError("explicit pack-path tests should not call prepare_data")
+
+
+def _load_module():
+    spec = importlib.util.spec_from_file_location("pack_sft_data_under_test", _SCRIPT_PATH)
+    assert spec is not None and spec.loader is not None
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[spec.name] = module
+    try:
+        spec.loader.exec_module(module)
+        return module
+    finally:
+        sys.modules.pop(spec.name, None)
+
+
+def _install_pack_sft_stubs(monkeypatch: pytest.MonkeyPatch, recipe_fn) -> Mock:
+    megatron = sys.modules.get("megatron", types.ModuleType("megatron"))
+    bridge = types.ModuleType("megatron.bridge")
+    data = types.ModuleType("megatron.bridge.data")
+    builders = types.ModuleType("megatron.bridge.data.builders")
+    datasets = types.ModuleType("megatron.bridge.data.datasets")
+    training = types.ModuleType("megatron.bridge.training")
+    tokenizers = types.ModuleType("megatron.bridge.training.tokenizers")
+    recipes = types.ModuleType("megatron.bridge.recipes")
+    recipes.unit_recipe = recipe_fn
+
+    finetuning_dataset = types.ModuleType("megatron.bridge.data.builders.finetuning_dataset")
+    finetuning_dataset.FinetuningDatasetBuilder = _FinetuningDatasetBuilder
+
+    hf_dataset = types.ModuleType("megatron.bridge.data.builders.hf_dataset")
+    hf_dataset.HFDatasetBuilder = _HFDatasetBuilder
+    hf_dataset.HFDatasetConfig = _HFDatasetConfig
+
+    packed_sequence = types.ModuleType("megatron.bridge.data.datasets.packed_sequence")
+    prepare_packed_sequence_data = Mock()
+    packed_sequence.prepare_packed_sequence_data = prepare_packed_sequence_data
+
+    training_config = types.ModuleType("megatron.bridge.training.config")
+    training_config.DataloaderConfig = _DataloaderConfig
+
+    tokenizer_module = types.ModuleType("megatron.bridge.training.tokenizers.tokenizer")
+    tokenizer_module.build_tokenizer = Mock(return_value="tokenizer")
+
+    monkeypatch.setitem(sys.modules, "megatron", megatron)
+    monkeypatch.setitem(sys.modules, "megatron.bridge", bridge)
+    monkeypatch.setitem(sys.modules, "megatron.bridge.data", data)
+    monkeypatch.setitem(sys.modules, "megatron.bridge.data.builders", builders)
+    monkeypatch.setitem(sys.modules, "megatron.bridge.data.datasets", datasets)
+    monkeypatch.setitem(sys.modules, "megatron.bridge.training", training)
+    monkeypatch.setitem(sys.modules, "megatron.bridge.training.tokenizers", tokenizers)
+    monkeypatch.setitem(sys.modules, "megatron.bridge.recipes", recipes)
+    monkeypatch.setitem(sys.modules, "megatron.bridge.data.builders.finetuning_dataset", finetuning_dataset)
+    monkeypatch.setitem(sys.modules, "megatron.bridge.data.builders.hf_dataset", hf_dataset)
+    monkeypatch.setitem(sys.modules, "megatron.bridge.data.datasets.packed_sequence", packed_sequence)
+    monkeypatch.setitem(sys.modules, "megatron.bridge.training.config", training_config)
+    monkeypatch.setitem(sys.modules, "megatron.bridge.training.tokenizers.tokenizer", tokenizer_module)
+    monkeypatch.setattr(megatron, "bridge", bridge, raising=False)
+    monkeypatch.setattr(bridge, "recipes", recipes, raising=False)
+    monkeypatch.setattr(bridge, "data", data, raising=False)
+    monkeypatch.setattr(data, "builders", builders, raising=False)
+    monkeypatch.setattr(data, "datasets", datasets, raising=False)
+    monkeypatch.setattr(builders, "finetuning_dataset", finetuning_dataset, raising=False)
+    monkeypatch.setattr(builders, "hf_dataset", hf_dataset, raising=False)
+    monkeypatch.setattr(datasets, "packed_sequence", packed_sequence, raising=False)
+    monkeypatch.setattr(bridge, "training", training, raising=False)
+    monkeypatch.setattr(training, "config", training_config, raising=False)
+    monkeypatch.setattr(training, "tokenizers", tokenizers, raising=False)
+    monkeypatch.setattr(tokenizers, "tokenizer", tokenizer_module, raising=False)
+    return prepare_packed_sequence_data
+
+
+def test_pack_sft_data_rejects_unsupported_seq_length(monkeypatch):
+    module = _load_module()
+
+    def unit_recipe():
+        raise AssertionError("recipe should not run when explicit override cannot be forwarded")
+
+    _install_pack_sft_stubs(monkeypatch, unit_recipe)
+    monkeypatch.setattr(sys, "argv", ["pack_sft_data.py", "--recipe", "unit_recipe", "--seq-length", "4096"])
+
+    with pytest.raises(SystemExit) as exc_info:
+        module.main()
+
+    assert str(exc_info.value) == "Error: recipe 'unit_recipe' does not accept a 'seq_length' parameter."
+
+
+def test_pack_sft_data_rejects_unsupported_hf_path(monkeypatch):
+    module = _load_module()
+
+    def unit_recipe():
+        raise AssertionError("recipe should not run when explicit override cannot be forwarded")
+
+    _install_pack_sft_stubs(monkeypatch, unit_recipe)
+    monkeypatch.setattr(sys, "argv", ["pack_sft_data.py", "--recipe", "unit_recipe", "--hf-path", "nvidia/unit"])
+
+    with pytest.raises(SystemExit) as exc_info:
+        module.main()
+
+    assert str(exc_info.value) == "Error: recipe 'unit_recipe' does not accept an 'hf_path' parameter."
+
+
+def test_pack_sft_data_forwards_supported_overrides_and_explicit_paths(monkeypatch, tmp_path):
+    module = _load_module()
+    recipe_calls = []
+
+    def unit_recipe(seq_length: int, hf_path: str):
+        recipe_calls.append({"seq_length": seq_length, "hf_path": hf_path})
+        return _RecipeConfig(dataset=_DatasetConfig(seq_length=seq_length))
+
+    prepare_packed_sequence_data = _install_pack_sft_stubs(monkeypatch, unit_recipe)
+
+    train_input = tmp_path / "train.jsonl"
+    train_output = tmp_path / "packed-train.parquet"
+    metadata_output = tmp_path / "metadata.json"
+    monkeypatch.setattr(
+        sys,
+        "argv",
+        [
+            "pack_sft_data.py",
+            "--recipe",
+            "unit_recipe",
+            "--seq-length",
+            "4096",
+            "--hf-path",
+            "nvidia/unit",
+            "--train-input-path",
+            str(train_input),
+            "--packed-train-data-path",
+            str(train_output),
+            "--packed-metadata-path",
+            str(metadata_output),
+        ],
+    )
+
+    module.main()
+
+    assert recipe_calls == [{"seq_length": 4096, "hf_path": "nvidia/unit"}]
+    prepare_packed_sequence_data.assert_called_once()
+    _, kwargs = prepare_packed_sequence_data.call_args
+    assert kwargs["input_path"] == train_input
+    assert kwargs["output_path"] == train_output
+    assert kwargs["output_metadata_path"] == metadata_output
+    assert kwargs["packed_sequence_size"] == 2048
+    assert kwargs["max_seq_length"] == 4096
+    assert kwargs["num_tokenizer_workers"] == 1
diff --git a/tests/unit_tests/scripts/test_text_generation.py b/tests/unit_tests/scripts/test_text_generation.py
index bc24ce6bab..e04d9a01ea 100644
--- a/tests/unit_tests/scripts/test_text_generation.py
+++ b/tests/unit_tests/scripts/test_text_generation.py
@@ -19,6 +19,7 @@
 import importlib.util
 import sys
 import types
+from datetime import timedelta
 from enum import Enum
 from pathlib import Path
 
@@ -180,3 +181,57 @@ def test_megatron_checkpoint_overrides_preserve_attention_backend(text_generatio
     assert overrides["fp16"] is False
     assert overrides["cache_mla_latents"] is True
     assert overrides["inference_moe_token_dispatcher_type"] == "nvls"
+
+
+def test_maybe_initialize_distributed_populates_env_from_safe_helpers(monkeypatch, text_generation):
+    init_calls = []
+    set_device_calls = []
+
+    for key in ("RANK", "WORLD_SIZE", "LOCAL_RANK", "MASTER_ADDR", "MASTER_PORT"):
+        monkeypatch.delenv(key, raising=False)
+
+    monkeypatch.setattr(text_generation.dist, "is_available", lambda: True)
+    monkeypatch.setattr(text_generation.dist, "is_initialized", lambda: False)
+    monkeypatch.setattr(text_generation, "get_rank_safe", lambda: 7)
+    monkeypatch.setattr(text_generation, "get_world_size_safe", lambda: 16)
+    monkeypatch.setattr(text_generation, "get_local_rank_preinit", lambda: 3)
+    monkeypatch.setattr(text_generation, "get_master_addr_safe", lambda: "node-0")
+    monkeypatch.setattr(text_generation, "get_master_port_safe", lambda: 23456)
+    monkeypatch.setattr(text_generation.torch.cuda, "set_device", lambda device: set_device_calls.append(device))
+    monkeypatch.setattr(
+        text_generation.dist,
+        "init_process_group",
+        lambda backend, timeout: init_calls.append({"backend": backend, "timeout": timeout}),
+    )
+
+    text_generation._maybe_initialize_distributed(timeout_minutes=11)
+
+    assert (
+        dict(
+            RANK="7",
+            WORLD_SIZE="16",
+            LOCAL_RANK="3",
+            MASTER_ADDR="node-0",
+            MASTER_PORT="23456",
+        ).items()
+        <= dict(text_generation.os.environ).items()
+    )
+    assert set_device_calls == [3]
+    assert init_calls == [{"backend": "nccl", "timeout": timedelta(minutes=11)}]
+
+
+def test_maybe_initialize_distributed_is_noop_when_dist_unavailable(monkeypatch, text_generation):
+    monkeypatch.setattr(text_generation.dist, "is_available", lambda: False)
+    monkeypatch.setattr(text_generation.dist, "is_initialized", lambda: False)
+    monkeypatch.setattr(
+        text_generation.dist,
+        "init_process_group",
+        lambda *args, **kwargs: pytest.fail("init_process_group should not be called"),
+    )
+    monkeypatch.setattr(
+        text_generation.torch.cuda,
+        "set_device",
+        lambda device: pytest.fail("set_device should not be called"),
+    )
+
+    text_generation._maybe_initialize_distributed(timeout_minutes=1)

From 461fff528bd655900d802895c3dba2a2df3e8718 Mon Sep 17 00:00:00 2001
From: Chen Cui <chcui@nvidia.com>
Date: Tue, 23 Jun 2026 09:45:13 -0700
Subject: [PATCH 4/5] feat(nemotronh): add Nemotron 3 Ultra recipes and
 examples (#4224)

Signed-off-by: Chen Cui <chcui@nvidia.com>
(cherry picked from commit b35a97b2a2183ee655b201ee0997fbefb373b997)
---
 docs/models/nemotron/index.md                 |   2 +
 docs/models/nemotron/nemotron3-ultra.md       |  11 +
 examples/models/nemotron/nemotron_3/README.md |  15 +-
 .../nemotron/nemotron_3/ultra/README.md       | 194 +++++++++++
 .../nemotron/nemotron_3/ultra/conversion.sh   |  41 +++
 .../nemotron_3/ultra/pack_data_job.sh         |  62 ++++
 .../nemotron_3/ultra/slurm_conversion.sh      | 120 +++++++
 .../nemotron_3/ultra/slurm_inference.sh       | 137 ++++++++
 .../nemotron/nemotron_3/ultra/slurm_peft.sh   | 179 ++++++++++
 .../nemotron_3/ultra/slurm_pretrain.sh        | 199 +++++++++++
 .../nemotron/nemotron_3/ultra/slurm_sft.sh    | 202 +++++++++++
 .../bridge/recipes/nemotronh/__init__.py      |   9 +
 .../recipes/nemotronh/nemotron_3_ultra.py     | 316 ++++++++++++++++++
 .../nemotronh/test_nemotron_3_ultra.py        | 155 +++++++++
 14 files changed, 1641 insertions(+), 1 deletion(-)
 create mode 100644 docs/models/nemotron/nemotron3-ultra.md
 create mode 100644 examples/models/nemotron/nemotron_3/ultra/README.md
 create mode 100755 examples/models/nemotron/nemotron_3/ultra/conversion.sh
 create mode 100755 examples/models/nemotron/nemotron_3/ultra/pack_data_job.sh
 create mode 100755 examples/models/nemotron/nemotron_3/ultra/slurm_conversion.sh
 create mode 100755 examples/models/nemotron/nemotron_3/ultra/slurm_inference.sh
 create mode 100755 examples/models/nemotron/nemotron_3/ultra/slurm_peft.sh
 create mode 100755 examples/models/nemotron/nemotron_3/ultra/slurm_pretrain.sh
 create mode 100755 examples/models/nemotron/nemotron_3/ultra/slurm_sft.sh
 create mode 100644 src/megatron/bridge/recipes/nemotronh/nemotron_3_ultra.py
 create mode 100644 tests/unit_tests/recipes/nemotronh/test_nemotron_3_ultra.py

diff --git a/docs/models/nemotron/index.md b/docs/models/nemotron/index.md
index 66089cf1a2..28d978789e 100644
--- a/docs/models/nemotron/index.md
+++ b/docs/models/nemotron/index.md
@@ -9,6 +9,7 @@ llama-nemotron.md
 nemotronh.md
 nemotron3-nano.md
 nemotron3-super.md
+nemotron3-ultra.md
 nemotron-nano-v2-vl.md
 nemotron-3-omni.md
 ```
@@ -19,6 +20,7 @@ nemotron-3-omni.md
 | Nemotron H and Nemotron Nano v2 | [nemotronh.md](nemotronh.md) |
 | Nemotron-3 Nano | [nemotron3-nano.md](nemotron3-nano.md) |
 | Nemotron-3 Super | [nemotron3-super.md](nemotron3-super.md) |
+| Nemotron-3 Ultra | [nemotron3-ultra.md](nemotron3-ultra.md) |
 | Nemotron Nano V2 VL | [nemotron-nano-v2-vl.md](nemotron-nano-v2-vl.md) |
 | Nemotron-3 Nano Omni | [nemotron-3-omni.md](nemotron-3-omni.md) |
 
diff --git a/docs/models/nemotron/nemotron3-ultra.md b/docs/models/nemotron/nemotron3-ultra.md
new file mode 100644
index 0000000000..7545f97d03
--- /dev/null
+++ b/docs/models/nemotron/nemotron3-ultra.md
@@ -0,0 +1,11 @@
+# Nemotron 3 Ultra
+
+[Nemotron 3 Ultra](https://docs.nvidia.com/nemotron/nightly/usage-cookbook/Nemotron-3-Ultra-Base/README.html)
+is a 550B total / A55B active hybrid Mamba-Transformer MoE model.
+
+Megatron Bridge provides Nemotron 3 Ultra recipes and examples for
+Hugging Face to Megatron conversion, inference, DCLM pretraining,
+packed OpenMathInstruct-2 full SFT, and packed OpenMathInstruct-2 LoRA PEFT.
+
+Use the main example README for setup and scripts:
+[`examples/models/nemotron/nemotron_3/ultra/README.md`](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/models/nemotron/nemotron_3/ultra/README.md).
diff --git a/examples/models/nemotron/nemotron_3/README.md b/examples/models/nemotron/nemotron_3/README.md
index c2446a0369..cfefc5aba9 100644
--- a/examples/models/nemotron/nemotron_3/README.md
+++ b/examples/models/nemotron/nemotron_3/README.md
@@ -6,6 +6,7 @@ This directory contains example scripts for Nemotron 3 language models:
 |-------|-----------|-------------------|--------------|
 | Nemotron 3 Nano | 30B | A3B | [nano/](nano/) |
 | Nemotron 3 Super | 120B | A12B | [super/](super/) |
+| Nemotron 3 Ultra | 550B | A55B | [ultra/](ultra/) |
 
 ## Workspace Configuration
 
@@ -21,7 +22,7 @@ Directory structure:
 
 ## Checkpoint Conversion
 
-Each model has its own conversion script: [nano/conversion.sh](nano/conversion.sh), [super/conversion.sh](super/conversion.sh).
+Nano and Super have conversion scripts: [nano/conversion.sh](nano/conversion.sh), [super/conversion.sh](super/conversion.sh). Ultra has Slurm examples for multi-node conversion, inference, and OpenMath training; see [ultra/](ultra/) and [Ultra documentation](../../../../docs/models/nemotron/nemotron3-ultra.md).
 
 ## Training Recipes
 
@@ -37,6 +38,11 @@ Available recipes:
 - `nemotron_3_super_sft_config`: Supervised fine-tuning
 - `nemotron_3_super_peft_config`: PEFT with LoRA support
 
+**Ultra** ([source](../../../../src/megatron/bridge/recipes/nemotronh/nemotron_3_ultra.py)):
+- `nemotron_3_ultra_pretrain_config`: Pretraining
+- `nemotron_3_ultra_sft_openmathinstruct2_packed_config`: Packed OpenMathInstruct-2 SFT
+- `nemotron_3_ultra_peft_openmathinstruct2_packed_config`: Packed OpenMathInstruct-2 PEFT
+
 Before training, ensure the following are configured:
 1. **Container Image**: Set `CONTAINER_IMAGE` in the SLURM scripts to your container path
 2. **Container Mounts**: (optional) Set `CONTAINER_MOUNTS` for data and workspace directories
@@ -55,6 +61,13 @@ See the SLURM scripts in [nano/](nano/): [slurm_pretrain.sh](nano/slurm_pretrain
 
 See the SLURM scripts in [super/](super/): [slurm_pretrain.sh](super/slurm_pretrain.sh), [slurm_sft.sh](super/slurm_sft.sh), [slurm_peft.sh](super/slurm_peft.sh).
 
+### Ultra
+
+See [ultra/slurm_inference.sh](ultra/slurm_inference.sh) for the 4-node inference pattern.
+For OpenMath training, use [ultra/slurm_sft.sh](ultra/slurm_sft.sh) and
+[ultra/slurm_peft.sh](ultra/slurm_peft.sh), which default to the current
+OpenMath tuning starting points.
+
 ## Evaluation
 
 Coming soon.
diff --git a/examples/models/nemotron/nemotron_3/ultra/README.md b/examples/models/nemotron/nemotron_3/ultra/README.md
new file mode 100644
index 0000000000..6d11e9262f
--- /dev/null
+++ b/examples/models/nemotron/nemotron_3/ultra/README.md
@@ -0,0 +1,194 @@
+# Nemotron 3 Ultra Examples
+
+This directory contains examples for Nemotron 3 Ultra conversion, inference,
+DCLM pretraining, packed OpenMathInstruct-2 full SFT, and packed
+OpenMathInstruct-2 LoRA PEFT.
+
+Nemotron 3 Ultra is a 550B total / A55B active hybrid Mamba-Transformer MoE
+model. See the
+[Nemotron 3 Ultra Base model guide](https://docs.nvidia.com/nemotron/nightly/usage-cookbook/Nemotron-3-Ultra-Base/README.html)
+for model details.
+
+## Workspace Configuration
+
+The scripts use `WORKSPACE` as the base directory for checkpoints, packed data,
+and results. Defaults:
+
+```bash
+export WORKSPACE=/workspace
+export MODEL_HOME=${WORKSPACE}/models/nvidia
+export HF_MODEL_PATH=nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16
+export MEGATRON_MODEL_PATH=${MODEL_HOME}/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16-megatron
+export PRETRAINED_CHECKPOINT=${MEGATRON_MODEL_PATH}
+```
+
+Use shared filesystems for multi-node jobs:
+
+```bash
+export HF_HOME=${WORKSPACE}/cache/hf
+export NEMO_HOME=${WORKSPACE}/cache/nemo
+export UV_CACHE_DIR=${WORKSPACE}/cache/uv
+```
+
+The BF16 Hugging Face cache and imported Megatron checkpoint are each about
+1.1 TB. Reserve at least 2.5 TB for model storage before starting checkpoint
+conversion, plus additional space for training outputs and logs. Full-model
+training checkpoints can each require several TB, so set `WORKSPACE` to a
+filesystem with enough quota before running SFT or pretraining.
+
+## Hardware Starting Points
+
+The checked-in Slurm scripts default to 8xH100 nodes unless noted below.
+When running on 4xGB200 nodes, update the `#SBATCH --nodes`,
+`#SBATCH --ntasks-per-node`, `#SBATCH --gpus-per-node`, and parallelism
+environment variables to the GB200 values in this table.
+
+| Workflow | 8xH100 nodes | 4xGB200 nodes |
+| --- | --- | --- |
+| Checkpoint import | 1 node with [conversion.sh](conversion.sh), CPU import path | 6 nodes with [slurm_conversion.sh](slurm_conversion.sh), `TP=1 PP=6 EP=4` |
+| Base inference | 4 nodes, `TP=1 PP=4 EP=8`, `KV_CACHE_BUFFER_SIZE_GB=4` | 3 nodes, `TP=1 PP=3 EP=4` |
+| DCLM pretraining | 48 nodes, `TP=4 PP=12 EP=16`, full uniform recompute with `RECOMPUTE_GRANULARITY=full RECOMPUTE_METHOD=uniform RECOMPUTE_NUM_LAYERS=1 RECOMPUTE_MODULES=""` | 24 nodes, `TP=2 PP=3 EP=32`, selective recompute on `moe+layernorm+core_attn+moe_act+mlp+shared_experts` |
+| OpenMath SFT | 48 nodes, `TP=2 PP=12 EP=16`, full uniform recompute with `RECOMPUTE_GRANULARITY=full RECOMPUTE_METHOD=uniform RECOMPUTE_NUM_LAYERS=1 RECOMPUTE_MODULES=""` | 48 nodes, `TP=2 PP=3 EP=32`, selective recompute on `moe+layernorm+core_attn+moe_act` |
+| OpenMath PEFT | 4 nodes, `TP=2 PP=4 EP=8`, selective recompute on `moe+layernorm+core_attn+moe_act+mlp+shared_experts` | 4 nodes, `TP=2 PP=1 EP=16`, selective recompute on `moe+layernorm+core_attn+moe_act` |
+
+These are bring-up and convergence starting points, not universal optima.
+Keep `TP` within a node-local NVLink domain and scale with `PP`, `EP`, and
+data parallelism when moving between hardware. For MoE sizing, the minimum GPU
+count is `PP * max(TP * CP, EP * ETP)`, then additional GPUs increase dense
+DP and expert DP.
+
+## Checkpoint Conversion
+
+Use [conversion.sh](conversion.sh) for CPU checkpoint import when the node has
+enough host RAM to materialize Nemotron 3 Ultra, for example an 8xH100 node.
+This is the preferred path when available because it avoids distributed GPU
+memory pressure during import.
+
+```bash
+bash conversion.sh
+```
+
+Set these variables for your environment:
+
+- `WORKSPACE`
+- `HF_HOME`
+- `UV_CACHE_DIR`
+- `HF_MODEL_PATH`
+- `MEGATRON_MODEL_PATH`
+
+Use [slurm_conversion.sh](slurm_conversion.sh) for distributed GPU checkpoint
+import when host RAM is not large enough, for example a 4xGB200 setup with less
+than 1 TB of host RAM. The checked-in distributed example defaults to 6
+8-GPU nodes with `TP=1 PP=6 EP=8`; for 4xGB200, use 6 nodes and set
+`#SBATCH --ntasks-per-node=4`, `#SBATCH --gpus-per-node=4`, and
+`TP=1 PP=6 EP=4`.
+
+```bash
+sbatch slurm_conversion.sh
+```
+
+Set these variables for your environment:
+
+- `CONTAINER_IMAGE`
+- `CONTAINER_MOUNTS`
+- `WORKDIR`
+- `WORKSPACE`
+- `HF_MODEL_PATH`
+- `MEGATRON_MODEL_PATH`
+
+## Inference
+
+Use [slurm_inference.sh](slurm_inference.sh) for 4-node text generation with
+`TP=1 PP=4 EP=8`. The script defaults `KV_CACHE_BUFFER_SIZE_GB=4` to keep
+the inference KV/context buffer within H100 memory for the default prompt
+lengths. On 4xGB200 nodes, use 3 nodes with `TP=1 PP=3 EP=4`.
+
+```bash
+sbatch slurm_inference.sh
+```
+
+Set `MEGATRON_MODEL_PATH` to generate from an imported Megatron checkpoint.
+Leave it unset to load from the Hugging Face checkpoint path.
+
+## DCLM Pretraining
+
+Use [slurm_pretrain.sh](slurm_pretrain.sh) for DCLM pretraining with
+`TP=4 PP=12 EP=16` and full uniform recompute on 8xH100 nodes. On
+4xGB200 nodes, use 24 nodes with `TP=2 PP=3 EP=32`.
+
+```bash
+sbatch slurm_pretrain.sh
+```
+
+Set `DCLM_DATA_DIR` to a preprocessed DCLM directory containing
+`*_text_document.bin` / `*_text_document.idx` files. The script defaults to
+matching `dclm_01_*_text_document.bin`. Async checkpoint saving is enabled by
+the recipe; the script defaults `SAVE_INTERVAL=1000` to save one checkpoint for
+the default 1000-iteration starter run.
+
+## OpenMath Packed Data
+
+Pre-pack OpenMath data before training:
+
+```bash
+sbatch pack_data_job.sh
+```
+
+Use the same `SEQ_LENGTH`, `HF_MODEL_PATH`, and `NEMO_HOME` for packing and
+training. `NEMO_HOME` must point at a shared filesystem visible on all nodes.
+
+## Training
+
+PEFT:
+
+```bash
+sbatch slurm_peft.sh
+```
+
+Full SFT:
+
+```bash
+sbatch slurm_sft.sh
+```
+
+The scripts default to OpenMath convergence settings: `TRAIN_ITERS=1000`,
+`GLOBAL_BATCH_SIZE=128`, `SEQ_LENGTH=4096`, and `LR_WARMUP_ITERS=250`. W&B
+logging is disabled by default. SFT and PEFT save at the final training
+iteration by default; the SFT script removes older intermediate `iter_*`
+checkpoints after a successful run to avoid retaining multiple full-model
+checkpoints.
+
+Current OpenMath starting points are:
+
+- PEFT: 4 nodes, `TP=2 PP=4 EP=8`, selective recompute on
+  `moe+layernorm+core_attn+moe_act+mlp+shared_experts`.
+- Full SFT: 48 nodes, `TP=2 PP=12 EP=16`, full uniform recompute with
+  `RECOMPUTE_GRANULARITY=full RECOMPUTE_METHOD=uniform
+  RECOMPUTE_NUM_LAYERS=1 RECOMPUTE_MODULES=""`. This is the current H100
+  starting point for 4096-token packed OpenMath SFT.
+
+For 4xGB200 nodes:
+
+- PEFT: 4 nodes, `TP=2 PP=1 EP=16`, selective recompute on
+  `moe+layernorm+core_attn+moe_act`.
+- Full SFT: 48 nodes, `TP=2 PP=3 EP=32`, selective recompute on
+  `moe+layernorm+core_attn+moe_act`.
+
+Advanced VPP, pipeline-layout, and recompute sweeps are intentionally left out
+of these starter scripts; add those overrides only for targeted performance
+experiments.
+
+## W&B
+
+W&B logging is disabled by default:
+
+```bash
+WANDB_ENTITY=nvidia-nemo-fw-public
+WANDB_PROJECT=megatron-bridge-nemotron-ultra
+WANDB_MODE=disabled
+```
+
+To enable online W&B logging, set `WANDB_MODE=online` and make `WANDB_API_KEY`
+visible in the submit environment.
+
+Run names include model, OpenMath, mode, TP/PP/EP, recompute, and Slurm job ID.
diff --git a/examples/models/nemotron/nemotron_3/ultra/conversion.sh b/examples/models/nemotron/nemotron_3/ultra/conversion.sh
new file mode 100755
index 0000000000..be701694cb
--- /dev/null
+++ b/examples/models/nemotron/nemotron_3/ultra/conversion.sh
@@ -0,0 +1,41 @@
+#!/usr/bin/env bash
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+set -euo pipefail
+
+WORKSPACE=${WORKSPACE:-/workspace}
+MODEL_HOME=${MODEL_HOME:-${WORKSPACE}/models/nvidia}
+HF_MODEL_PATH=${HF_MODEL_PATH:-nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16}
+MEGATRON_MODEL_PATH=${MEGATRON_MODEL_PATH:-${MODEL_HOME}/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16-megatron}
+
+[ -n "${HF_HOME:-}" ] && export HF_HOME
+[ -n "${UV_CACHE_DIR:-}" ] && export UV_CACHE_DIR
+
+mkdir -p "$(dirname "$MEGATRON_MODEL_PATH")"
+
+if [ -e "${MEGATRON_MODEL_PATH}/latest_checkpointed_iteration.txt" ] || [ -e "${MEGATRON_MODEL_PATH}/latest_train_state.pt" ]; then
+    echo "ERROR: target already contains a Megatron checkpoint: ${MEGATRON_MODEL_PATH}"
+    exit 1
+fi
+
+echo "Nemotron 3 Ultra CPU import"
+echo "HF_MODEL_PATH=${HF_MODEL_PATH}"
+echo "MEGATRON_MODEL_PATH=${MEGATRON_MODEL_PATH}"
+
+uv run --no-sync python examples/conversion/convert_checkpoints.py import \
+    --hf-model "$HF_MODEL_PATH" \
+    --megatron-path "$MEGATRON_MODEL_PATH" \
+    --torch-dtype bfloat16 \
+    --device-map cpu
diff --git a/examples/models/nemotron/nemotron_3/ultra/pack_data_job.sh b/examples/models/nemotron/nemotron_3/ultra/pack_data_job.sh
new file mode 100755
index 0000000000..2243ca7ac7
--- /dev/null
+++ b/examples/models/nemotron/nemotron_3/ultra/pack_data_job.sh
@@ -0,0 +1,62 @@
+#!/bin/bash
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+#SBATCH --job-name=nemotron-ultra-pack
+#SBATCH --nodes=1
+#SBATCH --ntasks-per-node=1
+#SBATCH --cpus-per-task=16
+#SBATCH --time=04:00:00
+#SBATCH --account=<your-account>
+#SBATCH --partition=cpu
+#SBATCH --output=logs/nemotron_ultra_pack_%j.log
+
+set -euo pipefail
+
+CONTAINER_IMAGE=${CONTAINER_IMAGE:-}
+CONTAINER_MOUNTS=${CONTAINER_MOUNTS:-}
+WORKDIR=${WORKDIR:-/opt/Megatron-Bridge}
+HF_MODEL_PATH=${HF_MODEL_PATH:-nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16}
+RECIPE_NAME=${RECIPE_NAME:-nemotron_3_ultra_sft_openmathinstruct2_packed_config}
+SEQ_LENGTH=${SEQ_LENGTH:-4096}
+
+[ -n "${HF_HOME:-}" ] && export HF_HOME
+[ -n "${NEMO_HOME:-}" ] && export NEMO_HOME
+[ -n "${UV_CACHE_DIR:-}" ] && export UV_CACHE_DIR
+export WORKDIR HF_MODEL_PATH RECIPE_NAME SEQ_LENGTH
+
+if [ -z "$CONTAINER_IMAGE" ]; then
+    echo "ERROR: CONTAINER_IMAGE must be set."
+    exit 1
+fi
+
+mkdir -p logs
+
+SRUN_CMD=(srun --mpi=pmix --container-image="${CONTAINER_IMAGE}" --no-container-mount-home)
+if [ -n "$CONTAINER_MOUNTS" ]; then
+    SRUN_CMD+=(--container-mounts="${CONTAINER_MOUNTS}")
+fi
+
+"${SRUN_CMD[@]}" bash -c '
+set -euo pipefail
+cd "$WORKDIR"
+export PYTHONPATH="$WORKDIR/src:$WORKDIR/3rdparty/Megatron-LM:${PYTHONPATH:-}"
+
+uv run --no-sync python scripts/training/pack_sft_data.py \
+    --recipe "$RECIPE_NAME" \
+    --seq-length "$SEQ_LENGTH" \
+    --hf-path "$HF_MODEL_PATH"
+'
+
+echo PACK_DATA_DONE
diff --git a/examples/models/nemotron/nemotron_3/ultra/slurm_conversion.sh b/examples/models/nemotron/nemotron_3/ultra/slurm_conversion.sh
new file mode 100755
index 0000000000..83f7d4f2eb
--- /dev/null
+++ b/examples/models/nemotron/nemotron_3/ultra/slurm_conversion.sh
@@ -0,0 +1,120 @@
+#!/bin/bash
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# ==============================================================================
+# Nemotron 3 Ultra Distributed GPU Checkpoint Import
+#
+# Use this path when host RAM is not large enough for CPU checkpoint import
+# and the checkpoint must be materialized across GPUs.
+#
+# Usage:
+#   1. Modify the #SBATCH directives for your cluster.
+#   2. Set CONTAINER_IMAGE and optional CONTAINER_MOUNTS.
+#   3. Submit: sbatch slurm_conversion.sh
+# ==============================================================================
+
+#SBATCH --job-name=nemotron-ultra-gpu-import
+#SBATCH --nodes=6
+#SBATCH --ntasks-per-node=8
+#SBATCH --gpus-per-node=8
+#SBATCH --time=04:00:00
+#SBATCH --account=<your-account>
+#SBATCH --partition=batch
+#SBATCH --output=logs/nemotron_ultra_import_%j.log
+#SBATCH --exclusive
+
+set -euo pipefail
+
+CONTAINER_IMAGE=${CONTAINER_IMAGE:-}
+CONTAINER_MOUNTS=${CONTAINER_MOUNTS:-}
+WORKDIR=${WORKDIR:-/opt/Megatron-Bridge}
+MODEL_HOME=${MODEL_HOME:-${WORKSPACE:-/workspace}/models/nvidia}
+
+HF_MODEL_PATH=${HF_MODEL_PATH:-nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16}
+MEGATRON_MODEL_PATH=${MEGATRON_MODEL_PATH:-${MODEL_HOME}/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16-megatron}
+
+TP=${TP:-1}
+PP=${PP:-6}
+EP=${EP:-8}
+ETP=${ETP:-1}
+GPUS_PER_NODE=${GPUS_PER_NODE:-8}
+
+[ -n "${HF_HOME:-}" ] && export HF_HOME
+[ -n "${UV_CACHE_DIR:-}" ] && export UV_CACHE_DIR
+export TORCH_NCCL_AVOID_RECORD_STREAMS=1
+export NCCL_NVLS_ENABLE=0
+export NCCL_DEBUG=${NCCL_DEBUG:-WARN}
+export NCCL_TIMEOUT=${NCCL_TIMEOUT:-1800000}
+export PYTORCH_CUDA_ALLOC_CONF=${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}
+export CUDA_DEVICE_MAX_CONNECTIONS=${CUDA_DEVICE_MAX_CONNECTIONS:-1}
+
+if [ -z "$CONTAINER_IMAGE" ]; then
+    echo "ERROR: CONTAINER_IMAGE must be set."
+    exit 1
+fi
+
+if [ "$((TP * PP * EP))" -ne "$((SLURM_JOB_NUM_NODES * GPUS_PER_NODE))" ]; then
+    echo "ERROR: TP*PP*EP must equal nodes*GPUS_PER_NODE for this script."
+    echo "TP=$TP PP=$PP EP=$EP nodes=$SLURM_JOB_NUM_NODES GPUS_PER_NODE=$GPUS_PER_NODE"
+    exit 2
+fi
+
+if [ -e "${MEGATRON_MODEL_PATH}/latest_checkpointed_iteration.txt" ] || [ -e "${MEGATRON_MODEL_PATH}/latest_train_state.pt" ]; then
+    echo "ERROR: target already contains a Megatron checkpoint: ${MEGATRON_MODEL_PATH}"
+    exit 3
+fi
+
+mkdir -p logs "$(dirname "$MEGATRON_MODEL_PATH")"
+
+MASTER_ADDR=$(python3 - <<'PY'
+import os
+import re
+
+nodelist = os.environ.get("SLURM_NODELIST", "")
+match = re.match(r"([A-Za-z0-9_-]+)\[(\d+)", nodelist)
+print(match.group(1) + match.group(2) if match else nodelist.split(",")[0])
+PY
+)
+MASTER_PORT=$((18000 + SLURM_JOB_ID % 40000))
+export MASTER_ADDR MASTER_PORT HF_MODEL_PATH MEGATRON_MODEL_PATH TP PP EP ETP GPUS_PER_NODE WORKDIR
+
+echo "Nemotron 3 Ultra distributed GPU import"
+echo "Job ${SLURM_JOB_ID} nodes=${SLURM_JOB_NUM_NODES} GPUs/node=${GPUS_PER_NODE} TP=${TP} PP=${PP} EP=${EP} ETP=${ETP}"
+echo "HF_MODEL_PATH=${HF_MODEL_PATH}"
+echo "MEGATRON_MODEL_PATH=${MEGATRON_MODEL_PATH}"
+
+SRUN_CMD=(srun --mpi=pmix --no-kill --container-image="${CONTAINER_IMAGE}" --no-container-mount-home)
+if [ -n "$CONTAINER_MOUNTS" ]; then
+    SRUN_CMD+=(--container-mounts="${CONTAINER_MOUNTS}")
+fi
+
+"${SRUN_CMD[@]}" bash -c '
+set -euo pipefail
+cd "$WORKDIR"
+export PYTHONPATH="$WORKDIR/src:$WORKDIR/3rdparty/Megatron-LM:${PYTHONPATH:-}"
+export RANK="${SLURM_PROCID}"
+export WORLD_SIZE="${SLURM_NTASKS}"
+export LOCAL_RANK="${SLURM_LOCALID}"
+export LOCAL_WORLD_SIZE="$GPUS_PER_NODE"
+
+uv run --no-sync python examples/conversion/convert_checkpoints_multi_gpu.py \
+    import \
+    --hf-model "$HF_MODEL_PATH" \
+    --megatron-path "$MEGATRON_MODEL_PATH" \
+    --tp "$TP" --pp "$PP" --ep "$EP" --etp "$ETP" \
+    --torch-dtype bfloat16
+'
+
+echo IMPORT_DONE
diff --git a/examples/models/nemotron/nemotron_3/ultra/slurm_inference.sh b/examples/models/nemotron/nemotron_3/ultra/slurm_inference.sh
new file mode 100755
index 0000000000..be5ff72756
--- /dev/null
+++ b/examples/models/nemotron/nemotron_3/ultra/slurm_inference.sh
@@ -0,0 +1,137 @@
+#!/bin/bash
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# ==============================================================================
+# Nemotron 3 Ultra Inference (4 nodes / 32 GPUs via Slurm)
+#
+# Usage:
+#   1. Set CONTAINER_IMAGE, CONTAINER_MOUNTS, and cache/token environment variables.
+#   2. Optionally set HF_MODEL_PATH to a local Hugging Face snapshot.
+#   3. Submit with: sbatch examples/models/nemotron/nemotron_3/ultra/slurm_inference.sh
+# ==============================================================================
+
+#SBATCH --job-name=nemotron-ultra-inference
+#SBATCH --nodes=4
+#SBATCH --ntasks-per-node=8
+#SBATCH --gpus-per-node=8
+#SBATCH --time=04:00:00
+#SBATCH --account=<your-account>
+#SBATCH --partition=batch
+#SBATCH --output=logs/nemotron_ultra_inference_%j.log
+#SBATCH --exclusive
+
+set -euo pipefail
+
+CONTAINER_IMAGE=${CONTAINER_IMAGE:-}
+CONTAINER_MOUNTS=${CONTAINER_MOUNTS:-}
+WORKDIR=${WORKDIR:-/opt/Megatron-Bridge}
+
+HF_MODEL_PATH=${HF_MODEL_PATH:-nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16}
+MEGATRON_MODEL_PATH=${MEGATRON_MODEL_PATH:-}
+PROMPT=${PROMPT:-"Solve 2x + 3 = 11. Show the reasoning briefly."}
+MAX_NEW_TOKENS=${MAX_NEW_TOKENS:-64}
+KV_CACHE_BUFFER_SIZE_GB=${KV_CACHE_BUFFER_SIZE_GB:-4}
+INFERENCE_MOE_TOKEN_DISPATCHER_TYPE=${INFERENCE_MOE_TOKEN_DISPATCHER_TYPE:-nccl}
+
+TP=${TP:-1}
+PP=${PP:-4}
+EP=${EP:-8}
+ETP=${ETP:-1}
+GPUS_PER_NODE=${GPUS_PER_NODE:-8}
+
+[ -n "${HF_HOME:-}" ] && export HF_HOME
+[ -n "${UV_CACHE_DIR:-}" ] && export UV_CACHE_DIR
+[ -n "${NEMO_HOME:-}" ] && export NEMO_HOME
+export TORCH_NCCL_AVOID_RECORD_STREAMS=1
+export NCCL_NVLS_ENABLE=0
+export NCCL_DEBUG=${NCCL_DEBUG:-WARN}
+export NCCL_TIMEOUT=${NCCL_TIMEOUT:-1800000}
+export PYTORCH_CUDA_ALLOC_CONF=${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}
+export CUDA_DEVICE_MAX_CONNECTIONS=${CUDA_DEVICE_MAX_CONNECTIONS:-1}
+
+if [ -z "$CONTAINER_IMAGE" ]; then
+    echo "ERROR: CONTAINER_IMAGE must be set."
+    exit 1
+fi
+
+TOTAL_GPUS=$((SLURM_JOB_NUM_NODES * GPUS_PER_NODE))
+EXPERT_MESH=$((EP * ETP))
+if [ "$EXPERT_MESH" -gt "$TP" ]; then
+    STAGE_GPUS=$EXPERT_MESH
+else
+    STAGE_GPUS=$TP
+fi
+MIN_GPUS=$((PP * STAGE_GPUS))
+if [ "$((TOTAL_GPUS % MIN_GPUS))" -ne 0 ]; then
+    echo "ERROR: nodes*GPUS_PER_NODE must be a multiple of PP*max(TP,EP*ETP) for MoE inference."
+    echo "TP=$TP PP=$PP EP=$EP ETP=$ETP nodes=$SLURM_JOB_NUM_NODES GPUS_PER_NODE=$GPUS_PER_NODE min_gpus=$MIN_GPUS total_gpus=$TOTAL_GPUS"
+    exit 2
+fi
+
+mkdir -p logs
+
+export HF_MODEL_PATH MEGATRON_MODEL_PATH PROMPT MAX_NEW_TOKENS KV_CACHE_BUFFER_SIZE_GB INFERENCE_MOE_TOKEN_DISPATCHER_TYPE
+export TP PP EP ETP GPUS_PER_NODE WORKDIR
+[ -n "${COORDINATOR_HOST:-}" ] && export COORDINATOR_HOST
+
+echo "Nemotron 3 Ultra inference"
+echo "Job ${SLURM_JOB_ID} nodes=${SLURM_JOB_NUM_NODES} GPUs/node=${GPUS_PER_NODE} TP=${TP} PP=${PP} EP=${EP} ETP=${ETP}"
+echo "HF_MODEL_PATH=${HF_MODEL_PATH}"
+echo "MEGATRON_MODEL_PATH=${MEGATRON_MODEL_PATH:-<load HF weights directly>}"
+echo "KV_CACHE_BUFFER_SIZE_GB=${KV_CACHE_BUFFER_SIZE_GB}"
+echo "COORDINATOR_HOST=${COORDINATOR_HOST:-<auto local IP>}"
+
+SRUN_CMD=(srun --mpi=pmix --no-kill --container-image="${CONTAINER_IMAGE}" --no-container-mount-home)
+if [ -n "$CONTAINER_MOUNTS" ]; then
+    SRUN_CMD+=(--container-mounts="${CONTAINER_MOUNTS}")
+fi
+
+"${SRUN_CMD[@]}" bash -c '
+set -euo pipefail
+cd "$WORKDIR"
+
+rm -f /opt/venv/lib/python3.12/site-packages/__editable__*megatron*.pth \
+      /opt/venv/lib/python3.12/site-packages/__editable__*megatron*.py 2>/dev/null || true
+
+export PYTHONPATH="$WORKDIR/src:$WORKDIR/3rdparty/Megatron-LM:${PYTHONPATH:-}"
+
+MEGATRON_MODEL_ARGS=()
+if [ -n "${MEGATRON_MODEL_PATH:-}" ]; then
+    MEGATRON_MODEL_ARGS=(--megatron_model_path "$MEGATRON_MODEL_PATH")
+fi
+
+COORDINATOR_ARGS=()
+if [ -z "${COORDINATOR_HOST:-}" ]; then
+    COORDINATOR_HOST=$(python3 - <<'"'"'PY'"'"'
+import socket
+
+print(socket.gethostbyname(socket.gethostname()))
+PY
+)
+fi
+COORDINATOR_ARGS=(--coordinator-host "$COORDINATOR_HOST")
+
+uv run --no-sync python scripts/inference/text_generation.py \
+    --hf_model_path "$HF_MODEL_PATH" \
+    "${MEGATRON_MODEL_ARGS[@]}" \
+    --prompt "$PROMPT" \
+    --max_new_tokens "$MAX_NEW_TOKENS" \
+    --kv_cache_buffer_size_gb "$KV_CACHE_BUFFER_SIZE_GB" \
+    --tp "$TP" --pp "$PP" --ep "$EP" --etp "$ETP" \
+    --use-coordinator \
+    "${COORDINATOR_ARGS[@]}" \
+    --inference-moe-token-dispatcher-type "$INFERENCE_MOE_TOKEN_DISPATCHER_TYPE" \
+    --distributed-timeout-minutes 90
+'
diff --git a/examples/models/nemotron/nemotron_3/ultra/slurm_peft.sh b/examples/models/nemotron/nemotron_3/ultra/slurm_peft.sh
new file mode 100755
index 0000000000..f8fa1dd389
--- /dev/null
+++ b/examples/models/nemotron/nemotron_3/ultra/slurm_peft.sh
@@ -0,0 +1,179 @@
+#!/bin/bash
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# ==============================================================================
+# Nemotron 3 Ultra OpenMath LoRA PEFT
+#
+# Usage:
+#   1. Modify the #SBATCH directives for your cluster.
+#   2. Set CONTAINER_IMAGE and optional CONTAINER_MOUNTS.
+#   3. Submit: sbatch slurm_peft.sh
+# ==============================================================================
+
+#SBATCH --job-name=nemotron-ultra-openmath-peft
+#SBATCH --nodes=4
+#SBATCH --ntasks-per-node=8
+#SBATCH --gpus-per-node=8
+#SBATCH --time=24:00:00
+#SBATCH --account=<your-account>
+#SBATCH --partition=batch
+#SBATCH --output=logs/nemotron_ultra_openmath_peft_%j.log
+#SBATCH --exclusive
+
+set -euo pipefail
+
+# ==============================================================================
+# CONFIGURATION
+# ==============================================================================
+
+WORKSPACE=${WORKSPACE:-/workspace}
+WORKDIR=${WORKDIR:-/opt/Megatron-Bridge}
+MODEL_HOME=${MODEL_HOME:-${WORKSPACE}/models/nvidia}
+
+HF_MODEL_PATH=${HF_MODEL_PATH:-nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16}
+PRETRAINED_CHECKPOINT=${PRETRAINED_CHECKPOINT:-${MODEL_HOME}/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16-megatron}
+MODEL_NAME=nemotron_3_ultra
+DATASET_NAME=openmathinstruct2
+RECIPE_NAME=nemotron_3_ultra_peft_openmathinstruct2_packed_config
+PEFT_SCHEME=lora
+
+SEQ_LENGTH=4096
+TRAIN_ITERS=${TRAIN_ITERS:-1000}
+GLOBAL_BATCH_SIZE=128
+MICRO_BATCH_SIZE=1
+EVAL_ITERS=32
+EVAL_INTERVAL=50
+LR_WARMUP_ITERS=250
+LR_DECAY_ITERS=${LR_DECAY_ITERS:-$TRAIN_ITERS}
+LR=1e-4
+MIN_LR=1e-5
+SAVE_INTERVAL=${SAVE_INTERVAL:-$TRAIN_ITERS}
+ASYNC_SAVE=${ASYNC_SAVE:-True}
+ASYNC_STRATEGY=${ASYNC_STRATEGY:-nvrx}
+LOG_INTERVAL=1
+
+TP=${TP:-2}
+PP=${PP:-4}
+EP=${EP:-8}
+ETP=${ETP:-1}
+CP=${CP:-1}
+SP=${SP:-True}
+GPUS_PER_NODE=${GPUS_PER_NODE:-8}
+
+RECOMPUTE_GRANULARITY=${RECOMPUTE_GRANULARITY:-selective}
+RECOMPUTE_MODULES=${RECOMPUTE_MODULES:-"[moe,layernorm,core_attn,moe_act,mlp,shared_experts]"}
+RECOMPUTE_TAG=${RECOMPUTE_TAG:-recompute_selective_moe_layernorm_core_attn_moe_act_mlp_shared_experts}
+
+WANDB_ENTITY=${WANDB_ENTITY:-nvidia-nemo-fw-public}
+WANDB_PROJECT=${WANDB_PROJECT:-megatron-bridge-nemotron-ultra}
+WANDB_MODE=${WANDB_MODE:-disabled}
+
+CONTAINER_IMAGE=${CONTAINER_IMAGE:-}
+CONTAINER_MOUNTS=${CONTAINER_MOUNTS:-}
+
+# ==============================================================================
+# Environment Setup
+# ==============================================================================
+
+[ -n "${HF_HOME:-}" ] && export HF_HOME
+[ -n "${NEMO_HOME:-}" ] && export NEMO_HOME
+[ -n "${UV_CACHE_DIR:-}" ] && export UV_CACHE_DIR
+export WANDB_MODE
+export TORCH_NCCL_AVOID_RECORD_STREAMS=1
+export NCCL_NVLS_ENABLE=0
+export NCCL_DEBUG=${NCCL_DEBUG:-WARN}
+export NCCL_TIMEOUT=${NCCL_TIMEOUT:-1800000}
+export PYTORCH_CUDA_ALLOC_CONF=${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}
+export CUDA_DEVICE_MAX_CONNECTIONS=${CUDA_DEVICE_MAX_CONNECTIONS:-1}
+
+# ==============================================================================
+# Job Execution
+# ==============================================================================
+
+mkdir -p logs
+
+if [ -z "$CONTAINER_IMAGE" ]; then
+    echo "ERROR: CONTAINER_IMAGE must be set."
+    exit 1
+fi
+
+if [ "$WANDB_MODE" = "online" ] && [ -z "${WANDB_API_KEY:-}" ]; then
+    echo "ERROR: WANDB_API_KEY must be visible in the submit environment for online W&B logging."
+    exit 2
+fi
+
+SAVE_DIR="${WORKSPACE}/results/${MODEL_NAME}_${DATASET_NAME}_peft_tp${TP}_pp${PP}_ep${EP}_sp${SP}_cp${CP}_${RECOMPUTE_TAG}_${SLURM_JOB_ID}"
+WANDB_EXP_NAME="${MODEL_NAME}_${DATASET_NAME}_peft_tp${TP}_pp${PP}_ep${EP}_${RECOMPUTE_TAG}_${SLURM_JOB_ID}"
+
+CLI_OVERRIDES="\
+    checkpoint.pretrained_checkpoint=${PRETRAINED_CHECKPOINT} \
+    checkpoint.save=${SAVE_DIR} \
+    checkpoint.save_interval=${SAVE_INTERVAL} \
+    checkpoint.async_save=${ASYNC_SAVE} \
+    checkpoint.async_strategy=${ASYNC_STRATEGY} \
+    train.train_iters=${TRAIN_ITERS} \
+    train.global_batch_size=${GLOBAL_BATCH_SIZE} \
+    train.micro_batch_size=${MICRO_BATCH_SIZE} \
+    validation.eval_interval=${EVAL_INTERVAL} \
+    validation.eval_iters=${EVAL_ITERS} \
+    scheduler.lr_warmup_iters=${LR_WARMUP_ITERS} \
+    scheduler.lr_decay_iters=${LR_DECAY_ITERS} \
+    optimizer.lr=${LR} \
+    optimizer.min_lr=${MIN_LR} \
+    logger.log_interval=${LOG_INTERVAL} \
+    logger.tensorboard_dir=${SAVE_DIR}/tb_logs \
+    logger.wandb_entity=${WANDB_ENTITY} \
+    logger.wandb_project=${WANDB_PROJECT} \
+    logger.wandb_exp_name=${WANDB_EXP_NAME} \
+    logger.wandb_save_dir=${SAVE_DIR}/wandb \
+    model.tensor_model_parallel_size=${TP} \
+    model.pipeline_model_parallel_size=${PP} \
+    model.expert_model_parallel_size=${EP} \
+    model.expert_tensor_parallel_size=${ETP} \
+    model.sequence_parallel=${SP} \
+    model.context_parallel_size=${CP} \
+    model.seq_length=${SEQ_LENGTH} \
+    model.recompute_granularity=${RECOMPUTE_GRANULARITY} \
+    model.recompute_modules=${RECOMPUTE_MODULES} \
+    dist.distributed_timeout_minutes=90"
+
+CMD="cd ${WORKDIR} && mkdir -p ${WORKSPACE}/results ${SAVE_DIR}/wandb ${SAVE_DIR}/tb_logs && \
+export PYTHONPATH=${WORKDIR}/src:${WORKDIR}/3rdparty/Megatron-LM:\${PYTHONPATH:-} && \
+uv run --no-sync python scripts/training/run_recipe.py \
+--recipe ${RECIPE_NAME} --peft_scheme ${PEFT_SCHEME} --seq_length ${SEQ_LENGTH} --hf_path ${HF_MODEL_PATH} \
+${CLI_OVERRIDES}"
+
+SRUN_CMD="srun --mpi=pmix --no-kill --container-image=${CONTAINER_IMAGE} --no-container-mount-home"
+if [ -n "$CONTAINER_MOUNTS" ]; then
+    SRUN_CMD="${SRUN_CMD} --container-mounts=${CONTAINER_MOUNTS}"
+fi
+
+echo "======================================"
+echo "Nemotron 3 Ultra OpenMath LoRA PEFT"
+echo "======================================"
+echo "Job ID: ${SLURM_JOB_ID}"
+echo "Nodes: ${SLURM_JOB_NUM_NODES}"
+echo "GPUs/node: ${GPUS_PER_NODE}"
+echo "Recipe: ${RECIPE_NAME}"
+echo "Parallelism: TP=${TP} PP=${PP} EP=${EP} ETP=${ETP} CP=${CP} SP=${SP}"
+echo "Recompute: ${RECOMPUTE_GRANULARITY} ${RECOMPUTE_MODULES}"
+echo "Async save: ${ASYNC_SAVE} (${ASYNC_STRATEGY})"
+echo "Save dir: ${SAVE_DIR}"
+echo "W&B: ${WANDB_ENTITY}/${WANDB_PROJECT} (${WANDB_MODE})"
+echo "======================================"
+
+$SRUN_CMD bash -c "$CMD"
+
+echo OPENMATH_PEFT_DONE
diff --git a/examples/models/nemotron/nemotron_3/ultra/slurm_pretrain.sh b/examples/models/nemotron/nemotron_3/ultra/slurm_pretrain.sh
new file mode 100755
index 0000000000..d58f27daa6
--- /dev/null
+++ b/examples/models/nemotron/nemotron_3/ultra/slurm_pretrain.sh
@@ -0,0 +1,199 @@
+#!/bin/bash
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# ==============================================================================
+# Nemotron 3 Ultra DCLM Pretraining
+#
+# Usage:
+#   1. Modify the #SBATCH directives for your cluster.
+#   2. Set CONTAINER_IMAGE and DCLM_DATA_DIR.
+#   3. Optionally set CONTAINER_MOUNTS, WORKSPACE, WORKDIR, or HF_MODEL_PATH.
+#   4. Submit: sbatch slurm_pretrain.sh
+# ==============================================================================
+
+#SBATCH --job-name=nemotron-ultra-dclm-pretrain
+#SBATCH --nodes=48
+#SBATCH --ntasks-per-node=8
+#SBATCH --gpus-per-node=8
+#SBATCH --time=24:00:00
+#SBATCH --account=<your-account>
+#SBATCH --partition=batch
+#SBATCH --output=logs/nemotron_ultra_dclm_pretrain_%j.log
+#SBATCH --exclusive
+
+set -euo pipefail
+
+# Required for most clusters:
+CONTAINER_IMAGE=${CONTAINER_IMAGE:-}
+DCLM_DATA_DIR=${DCLM_DATA_DIR:-}
+
+# Optional environment-specific paths:
+WORKSPACE=${WORKSPACE:-/workspace}
+WORKDIR=${WORKDIR:-/opt/Megatron-Bridge}
+CONTAINER_MOUNTS=${CONTAINER_MOUNTS:-}
+HF_MODEL_PATH=${HF_MODEL_PATH:-nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16}
+EXTRA_OVERRIDES=${EXTRA_OVERRIDES:-}
+
+# Starter profile:
+MODEL_NAME=nemotron_3_ultra
+RECIPE_NAME=nemotron_3_ultra_pretrain_config
+SEQ_LENGTH=4096
+DCLM_PATTERN=${DCLM_PATTERN:-dclm_01_*_text_document.bin}
+DCLM_CACHE="${WORKSPACE}/data_cache/dclm"
+TRAIN_ITERS=${TRAIN_ITERS:-1000}
+SAVE_INTERVAL=${SAVE_INTERVAL:-1000}
+TP=${TP:-4}
+PP=${PP:-12}
+EP=${EP:-16}
+ETP=${ETP:-1}
+CP=${CP:-1}
+SP=${SP:-True}
+RECOMPUTE_GRANULARITY=${RECOMPUTE_GRANULARITY:-full}
+RECOMPUTE_METHOD=${RECOMPUTE_METHOD:-uniform}
+if [ -z "${RECOMPUTE_MODULES+x}" ]; then
+    RECOMPUTE_MODULES=""
+fi
+RECOMPUTE_NUM_LAYERS=${RECOMPUTE_NUM_LAYERS:-1}
+RECOMPUTE_TAG=${RECOMPUTE_TAG:-recompute_full_uniform1}
+GPUS_PER_NODE=${GPUS_PER_NODE:-8}
+SAVE_DIR="${WORKSPACE}/results/${MODEL_NAME}_dclm_pretrain_tp${TP}_pp${PP}_ep${EP}_${RECOMPUTE_TAG}_${SLURM_JOB_ID}"
+WANDB_ENTITY=${WANDB_ENTITY:-nvidia-nemo-fw-public}
+WANDB_PROJECT=${WANDB_PROJECT:-megatron-bridge-nemotron-ultra}
+WANDB_EXP_NAME="${MODEL_NAME}_dclm_pretrain_tp${TP}_pp${PP}_ep${EP}_${RECOMPUTE_TAG}_${SLURM_JOB_ID}"
+WANDB_MODE=${WANDB_MODE:-disabled}
+
+[ -n "${HF_HOME:-}" ] && export HF_HOME
+[ -n "${NEMO_HOME:-}" ] && export NEMO_HOME
+[ -n "${UV_CACHE_DIR:-}" ] && export UV_CACHE_DIR
+export WANDB_MODE
+export TORCH_NCCL_AVOID_RECORD_STREAMS=1
+export NCCL_NVLS_ENABLE=0
+export NCCL_DEBUG=${NCCL_DEBUG:-WARN}
+export NCCL_TIMEOUT=${NCCL_TIMEOUT:-1800000}
+export PYTORCH_CUDA_ALLOC_CONF=${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}
+export CUDA_DEVICE_MAX_CONNECTIONS=${CUDA_DEVICE_MAX_CONNECTIONS:-1}
+export PYTHONWARNINGS="${PYTHONWARNINGS:+${PYTHONWARNINGS},}ignore:The AccumulateGrad node:UserWarning,ignore:The pad token id in the tokenizer collides:UserWarning"
+
+mkdir -p logs
+
+if [ -z "$CONTAINER_IMAGE" ]; then
+    echo "ERROR: set CONTAINER_IMAGE to a valid sqsh image."
+    exit 1
+fi
+
+if [ -z "$DCLM_DATA_DIR" ]; then
+    echo "ERROR: set DCLM_DATA_DIR to a preprocessed DCLM directory."
+    exit 2
+fi
+
+if [ "$WANDB_MODE" = "online" ] && [ -z "${WANDB_API_KEY:-}" ]; then
+    echo "ERROR: WANDB_API_KEY must be visible in the submit environment for online W&B logging."
+    exit 3
+fi
+
+export DCLM_DATA_DIR DCLM_PATTERN DCLM_CACHE EXTRA_OVERRIDES HF_MODEL_PATH RECIPE_NAME SAVE_DIR SEQ_LENGTH
+export TRAIN_ITERS SAVE_INTERVAL TP PP EP ETP CP SP RECOMPUTE_GRANULARITY RECOMPUTE_METHOD RECOMPUTE_MODULES
+export RECOMPUTE_NUM_LAYERS GPUS_PER_NODE
+export WANDB_ENTITY WANDB_PROJECT WANDB_EXP_NAME WORKDIR WORKSPACE
+
+CMD='
+set -euo pipefail
+cd "$WORKDIR"
+mkdir -p "$WORKSPACE/results" "$SAVE_DIR/wandb" "$SAVE_DIR/tb_logs" "$DCLM_CACHE"
+export PYTHONPATH="$WORKDIR/src:$WORKDIR/3rdparty/Megatron-LM:${PYTHONPATH:-}"
+
+BLEND_PATHS=""
+shopt -s nullglob
+for BIN_PATH in "$DCLM_DATA_DIR"/$DCLM_PATTERN; do
+    PREFIX=${BIN_PATH%.bin}
+    BLEND_PATHS="${BLEND_PATHS}\"${PREFIX}\","
+done
+shopt -u nullglob
+BLEND_PATHS="${BLEND_PATHS%,}"
+
+if [ -z "$BLEND_PATHS" ]; then
+    echo "ERROR: no DCLM shards matching ${DCLM_DATA_DIR}/${DCLM_PATTERN}"
+    exit 4
+fi
+
+OPTIONAL_OVERRIDES=()
+if [ -n "$RECOMPUTE_METHOD" ]; then
+    OPTIONAL_OVERRIDES+=(model.recompute_method="$RECOMPUTE_METHOD")
+fi
+if [ -n "$RECOMPUTE_MODULES" ]; then
+    OPTIONAL_OVERRIDES+=(model.recompute_modules="$RECOMPUTE_MODULES")
+fi
+if [ -n "$RECOMPUTE_NUM_LAYERS" ]; then
+    OPTIONAL_OVERRIDES+=(model.recompute_num_layers="$RECOMPUTE_NUM_LAYERS")
+fi
+
+uv run --no-sync python scripts/training/run_recipe.py \
+    --recipe "$RECIPE_NAME" \
+    --dataset llm-pretrain \
+    --seq_length "$SEQ_LENGTH" \
+    --hf_path "$HF_MODEL_PATH" \
+    checkpoint.save="$SAVE_DIR" \
+    checkpoint.save_interval="$SAVE_INTERVAL" \
+    train.train_iters="$TRAIN_ITERS" \
+    train.global_batch_size=128 \
+    train.micro_batch_size=1 \
+    validation.eval_interval=100 \
+    validation.eval_iters=10 \
+    logger.log_interval=1 \
+    logger.tensorboard_dir="$SAVE_DIR/tb_logs" \
+    logger.wandb_entity="$WANDB_ENTITY" \
+    logger.wandb_project="$WANDB_PROJECT" \
+    logger.wandb_exp_name="$WANDB_EXP_NAME" \
+    logger.wandb_save_dir="$SAVE_DIR/wandb" \
+    model.tensor_model_parallel_size="$TP" \
+    model.pipeline_model_parallel_size="$PP" \
+    model.expert_model_parallel_size="$EP" \
+    model.expert_tensor_parallel_size="$ETP" \
+    model.sequence_parallel="$SP" \
+    model.context_parallel_size="$CP" \
+    model.seq_length="$SEQ_LENGTH" \
+    dataset.sequence_length="$SEQ_LENGTH" \
+    model.recompute_granularity="$RECOMPUTE_GRANULARITY" \
+    "${OPTIONAL_OVERRIDES[@]}" \
+    dist.distributed_timeout_minutes=90 \
+    "dataset.blend=[[${BLEND_PATHS}],null]" \
+    dataset.split=\"9999,8,2\" \
+    dataset.path_to_cache="$DCLM_CACHE" \
+    ${EXTRA_OVERRIDES}
+'
+
+SRUN_CMD="srun --mpi=pmix --no-kill --container-image=${CONTAINER_IMAGE} --no-container-mount-home"
+if [ -n "$CONTAINER_MOUNTS" ]; then
+    SRUN_CMD="${SRUN_CMD} --container-mounts=${CONTAINER_MOUNTS}"
+fi
+
+echo "======================================"
+echo "Nemotron 3 Ultra DCLM Pretraining"
+echo "======================================"
+echo "Job ID: ${SLURM_JOB_ID}"
+echo "Nodes: ${SLURM_JOB_NUM_NODES}"
+echo "GPUs/node: ${GPUS_PER_NODE}"
+echo "Recipe: ${RECIPE_NAME}"
+echo "HF model: ${HF_MODEL_PATH}"
+echo "DCLM data: ${DCLM_DATA_DIR}"
+echo "Parallelism: TP=${TP} PP=${PP} EP=${EP} ETP=${ETP} CP=${CP} SP=${SP}"
+echo "Recompute: ${RECOMPUTE_GRANULARITY} ${RECOMPUTE_METHOD:-} ${RECOMPUTE_MODULES}"
+echo "Save dir: ${SAVE_DIR}"
+echo "W&B: ${WANDB_ENTITY}/${WANDB_PROJECT} (${WANDB_MODE})"
+echo "======================================"
+
+$SRUN_CMD bash -c "$CMD"
+
+echo DCLM_PRETRAIN_DONE
diff --git a/examples/models/nemotron/nemotron_3/ultra/slurm_sft.sh b/examples/models/nemotron/nemotron_3/ultra/slurm_sft.sh
new file mode 100755
index 0000000000..02ca86e9f9
--- /dev/null
+++ b/examples/models/nemotron/nemotron_3/ultra/slurm_sft.sh
@@ -0,0 +1,202 @@
+#!/bin/bash
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# ==============================================================================
+# Nemotron 3 Ultra OpenMath Full SFT
+#
+# Usage:
+#   1. Modify the #SBATCH directives for your cluster.
+#   2. Set CONTAINER_IMAGE and optional CONTAINER_MOUNTS.
+#   3. Submit: sbatch slurm_sft.sh
+# ==============================================================================
+
+#SBATCH --job-name=nemotron-ultra-openmath-sft
+#SBATCH --nodes=48
+#SBATCH --ntasks-per-node=8
+#SBATCH --gpus-per-node=8
+#SBATCH --time=24:00:00
+#SBATCH --account=<your-account>
+#SBATCH --partition=batch
+#SBATCH --output=logs/nemotron_ultra_openmath_sft_%j.log
+#SBATCH --exclusive
+
+set -euo pipefail
+
+# ==============================================================================
+# CONFIGURATION
+# ==============================================================================
+
+WORKSPACE=${WORKSPACE:-/workspace}
+WORKDIR=${WORKDIR:-/opt/Megatron-Bridge}
+MODEL_HOME=${MODEL_HOME:-${WORKSPACE}/models/nvidia}
+
+HF_MODEL_PATH=${HF_MODEL_PATH:-nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16}
+PRETRAINED_CHECKPOINT=${PRETRAINED_CHECKPOINT:-${MODEL_HOME}/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16-megatron}
+MODEL_NAME=nemotron_3_ultra
+DATASET_NAME=openmathinstruct2
+RECIPE_NAME=nemotron_3_ultra_sft_openmathinstruct2_packed_config
+
+SEQ_LENGTH=4096
+TRAIN_ITERS=${TRAIN_ITERS:-1000}
+GLOBAL_BATCH_SIZE=128
+MICRO_BATCH_SIZE=1
+EVAL_ITERS=32
+EVAL_INTERVAL=50
+LR_WARMUP_ITERS=250
+LR_DECAY_ITERS=${LR_DECAY_ITERS:-$TRAIN_ITERS}
+LR=5e-6
+MIN_LR=5e-7
+SAVE_INTERVAL=${SAVE_INTERVAL:-$TRAIN_ITERS}
+LOG_INTERVAL=1
+
+TP=${TP:-2}
+PP=${PP:-12}
+EP=${EP:-16}
+ETP=${ETP:-1}
+CP=${CP:-1}
+SP=${SP:-True}
+GPUS_PER_NODE=${GPUS_PER_NODE:-8}
+
+RECOMPUTE_GRANULARITY=${RECOMPUTE_GRANULARITY:-full}
+RECOMPUTE_METHOD=${RECOMPUTE_METHOD:-uniform}
+if [ -z "${RECOMPUTE_MODULES+x}" ]; then
+    RECOMPUTE_MODULES=""
+fi
+RECOMPUTE_NUM_LAYERS=${RECOMPUTE_NUM_LAYERS:-1}
+RECOMPUTE_TAG=${RECOMPUTE_TAG:-recompute_full_uniform1}
+
+WANDB_ENTITY=${WANDB_ENTITY:-nvidia-nemo-fw-public}
+WANDB_PROJECT=${WANDB_PROJECT:-megatron-bridge-nemotron-ultra}
+WANDB_MODE=${WANDB_MODE:-disabled}
+
+CONTAINER_IMAGE=${CONTAINER_IMAGE:-}
+CONTAINER_MOUNTS=${CONTAINER_MOUNTS:-}
+EXTRA_OVERRIDES=${EXTRA_OVERRIDES:-}
+
+# ==============================================================================
+# Environment Setup
+# ==============================================================================
+
+[ -n "${HF_HOME:-}" ] && export HF_HOME
+[ -n "${NEMO_HOME:-}" ] && export NEMO_HOME
+[ -n "${UV_CACHE_DIR:-}" ] && export UV_CACHE_DIR
+export WANDB_MODE
+export TORCH_NCCL_AVOID_RECORD_STREAMS=1
+export NCCL_NVLS_ENABLE=0
+export NCCL_DEBUG=${NCCL_DEBUG:-WARN}
+export NCCL_TIMEOUT=${NCCL_TIMEOUT:-1800000}
+export PYTORCH_CUDA_ALLOC_CONF=${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}
+export CUDA_DEVICE_MAX_CONNECTIONS=${CUDA_DEVICE_MAX_CONNECTIONS:-1}
+export PYTHONWARNINGS="${PYTHONWARNINGS:+${PYTHONWARNINGS},}ignore:The AccumulateGrad node:UserWarning"
+
+# ==============================================================================
+# Job Execution
+# ==============================================================================
+
+mkdir -p logs
+
+if [ -z "$CONTAINER_IMAGE" ]; then
+    echo "ERROR: CONTAINER_IMAGE must be set."
+    exit 1
+fi
+
+if [ "$WANDB_MODE" = "online" ] && [ -z "${WANDB_API_KEY:-}" ]; then
+    echo "ERROR: WANDB_API_KEY must be visible in the submit environment for online W&B logging."
+    exit 2
+fi
+
+SAVE_DIR="${WORKSPACE}/results/${MODEL_NAME}_${DATASET_NAME}_sft_tp${TP}_pp${PP}_ep${EP}_sp${SP}_cp${CP}_${RECOMPUTE_TAG}_${SLURM_JOB_ID}"
+WANDB_EXP_NAME="${MODEL_NAME}_${DATASET_NAME}_sft_tp${TP}_pp${PP}_ep${EP}_${RECOMPUTE_TAG}_${SLURM_JOB_ID}"
+
+CLI_OVERRIDES="\
+    checkpoint.pretrained_checkpoint=${PRETRAINED_CHECKPOINT} \
+    checkpoint.save=${SAVE_DIR} \
+    checkpoint.save_interval=${SAVE_INTERVAL} \
+    train.train_iters=${TRAIN_ITERS} \
+    train.global_batch_size=${GLOBAL_BATCH_SIZE} \
+    train.micro_batch_size=${MICRO_BATCH_SIZE} \
+    validation.eval_interval=${EVAL_INTERVAL} \
+    validation.eval_iters=${EVAL_ITERS} \
+    scheduler.lr_warmup_iters=${LR_WARMUP_ITERS} \
+    scheduler.lr_decay_iters=${LR_DECAY_ITERS} \
+    optimizer.lr=${LR} \
+    optimizer.min_lr=${MIN_LR} \
+    logger.log_interval=${LOG_INTERVAL} \
+    logger.tensorboard_dir=${SAVE_DIR}/tb_logs \
+    logger.wandb_entity=${WANDB_ENTITY} \
+    logger.wandb_project=${WANDB_PROJECT} \
+    logger.wandb_exp_name=${WANDB_EXP_NAME} \
+    logger.wandb_save_dir=${SAVE_DIR}/wandb \
+    model.tensor_model_parallel_size=${TP} \
+    model.pipeline_model_parallel_size=${PP} \
+    model.expert_model_parallel_size=${EP} \
+    model.expert_tensor_parallel_size=${ETP} \
+    model.sequence_parallel=${SP} \
+    model.context_parallel_size=${CP} \
+    model.seq_length=${SEQ_LENGTH} \
+    model.recompute_granularity=${RECOMPUTE_GRANULARITY} \
+    dist.distributed_timeout_minutes=90"
+
+if [ -n "$RECOMPUTE_METHOD" ]; then
+    CLI_OVERRIDES="${CLI_OVERRIDES} model.recompute_method=${RECOMPUTE_METHOD}"
+fi
+if [ -n "$RECOMPUTE_MODULES" ]; then
+    CLI_OVERRIDES="${CLI_OVERRIDES} model.recompute_modules=${RECOMPUTE_MODULES}"
+fi
+if [ -n "$RECOMPUTE_NUM_LAYERS" ]; then
+    CLI_OVERRIDES="${CLI_OVERRIDES} model.recompute_num_layers=${RECOMPUTE_NUM_LAYERS}"
+fi
+CLI_OVERRIDES="${CLI_OVERRIDES} ${EXTRA_OVERRIDES}"
+
+CMD="cd ${WORKDIR} && mkdir -p ${WORKSPACE}/results ${SAVE_DIR}/wandb ${SAVE_DIR}/tb_logs && \
+export PYTHONPATH=${WORKDIR}/src:${WORKDIR}/3rdparty/Megatron-LM:\${PYTHONPATH:-} && \
+uv run --no-sync python scripts/training/run_recipe.py \
+--recipe ${RECIPE_NAME} --seq_length ${SEQ_LENGTH} --hf_path ${HF_MODEL_PATH} \
+${CLI_OVERRIDES}"
+
+SRUN_CMD="srun --mpi=pmix --no-kill --container-image=${CONTAINER_IMAGE} --no-container-mount-home"
+if [ -n "$CONTAINER_MOUNTS" ]; then
+    SRUN_CMD="${SRUN_CMD} --container-mounts=${CONTAINER_MOUNTS}"
+fi
+
+echo "======================================"
+echo "Nemotron 3 Ultra OpenMath Full SFT"
+echo "======================================"
+echo "Job ID: ${SLURM_JOB_ID}"
+echo "Nodes: ${SLURM_JOB_NUM_NODES}"
+echo "GPUs/node: ${GPUS_PER_NODE}"
+echo "Recipe: ${RECIPE_NAME}"
+echo "Parallelism: TP=${TP} PP=${PP} EP=${EP} ETP=${ETP} CP=${CP} SP=${SP}"
+echo "Recompute: ${RECOMPUTE_GRANULARITY} ${RECOMPUTE_METHOD:-} ${RECOMPUTE_MODULES}"
+echo "Save dir: ${SAVE_DIR}"
+echo "W&B: ${WANDB_ENTITY}/${WANDB_PROJECT} (${WANDB_MODE})"
+echo "======================================"
+
+$SRUN_CMD bash -c "$CMD"
+
+LATEST_ITER_FILE="${SAVE_DIR}/latest_checkpointed_iteration.txt"
+if [ -f "$LATEST_ITER_FILE" ]; then
+    LATEST_ITER=$(tr -d '[:space:]' < "$LATEST_ITER_FILE")
+    if [[ "$LATEST_ITER" =~ ^[0-9]+$ ]]; then
+        LATEST_DIR=$(printf "iter_%07d" "$LATEST_ITER")
+        find "$SAVE_DIR" -mindepth 1 -maxdepth 1 -type d -name "iter_*" ! -name "$LATEST_DIR" -exec rm -rf {} +
+    else
+        echo "Skipping SFT intermediate checkpoint cleanup: latest checkpoint marker is not numeric: ${LATEST_ITER}"
+    fi
+else
+    echo "Skipping SFT intermediate checkpoint cleanup: missing ${LATEST_ITER_FILE}"
+fi
+
+echo OPENMATH_SFT_DONE
diff --git a/src/megatron/bridge/recipes/nemotronh/__init__.py b/src/megatron/bridge/recipes/nemotronh/__init__.py
index 46ae25a152..91bfe4feeb 100644
--- a/src/megatron/bridge/recipes/nemotronh/__init__.py
+++ b/src/megatron/bridge/recipes/nemotronh/__init__.py
@@ -26,6 +26,11 @@
     nemotron_3_super_pretrain_config,
     nemotron_3_super_sft_config,
 )
+from megatron.bridge.recipes.nemotronh.nemotron_3_ultra import (
+    nemotron_3_ultra_peft_openmathinstruct2_packed_config,
+    nemotron_3_ultra_pretrain_config,
+    nemotron_3_ultra_sft_openmathinstruct2_packed_config,
+)
 from megatron.bridge.recipes.nemotronh.nemotron_nano_v2 import (
     nemotron_nano_9b_v2_peft_config,
     nemotron_nano_9b_v2_pretrain_config,
@@ -81,4 +86,8 @@
     "nemotron_3_super_pretrain_config",
     "nemotron_3_super_sft_config",
     "nemotron_3_super_peft_config",
+    # Nemotron 3 Ultra models
+    "nemotron_3_ultra_pretrain_config",
+    "nemotron_3_ultra_sft_openmathinstruct2_packed_config",
+    "nemotron_3_ultra_peft_openmathinstruct2_packed_config",
 ]
diff --git a/src/megatron/bridge/recipes/nemotronh/nemotron_3_ultra.py b/src/megatron/bridge/recipes/nemotronh/nemotron_3_ultra.py
new file mode 100644
index 0000000000..f99c8fa30f
--- /dev/null
+++ b/src/megatron/bridge/recipes/nemotronh/nemotron_3_ultra.py
@@ -0,0 +1,316 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+
+from megatron.bridge import AutoBridge
+from megatron.bridge.peft.base import PEFT
+from megatron.bridge.recipes.common import _peft_common, _pretrain_common, _sft_common
+from megatron.bridge.recipes.utils.finetune_utils import default_openmathinstruct2_config, default_peft_config
+from megatron.bridge.training.config import ConfigContainer
+
+
+NEMOTRON_3_ULTRA_HF_MODEL_ID = "nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16"
+NEMOTRON_3_ULTRA_TOKENIZER_NAME = "nvidia--NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16"
+
+
+def nemotron_3_ultra_pretrain_config(
+    *,
+    hf_path: str | None = None,
+    seq_length: int = 8192,
+) -> ConfigContainer:
+    """Return a pre-training config for Nemotron 3 Ultra.
+
+    Args:
+        hf_path: Optional Hugging Face model ID or local snapshot path.
+        seq_length: Sequence length for model and dataset settings.
+
+    Returns:
+        Pre-training configuration for Nemotron 3 Ultra.
+    """
+    cfg = _pretrain_common()
+    model_source = hf_path or NEMOTRON_3_ULTRA_HF_MODEL_ID
+
+    cfg.model = AutoBridge.from_hf_pretrained(model_source).to_megatron_provider(load_weights=False)
+    cfg.model.tensor_model_parallel_size = 1
+    cfg.model.pipeline_model_parallel_size = 3
+    cfg.model.pipeline_dtype = torch.bfloat16
+    cfg.model.virtual_pipeline_model_parallel_size = None
+    cfg.model.context_parallel_size = 1
+    cfg.model.sequence_parallel = True
+    cfg.model.expert_tensor_parallel_size = 1
+    cfg.model.expert_model_parallel_size = 8
+    cfg.model.pipeline_model_parallel_layout = None
+    cfg.model.seq_length = seq_length
+    cfg.model.apply_rope_fusion = False
+    cfg.model.attention_backend = "fused"
+    cfg.model.gradient_accumulation_fusion = True
+    cfg.model.init_method_std = 0.014
+    cfg.model.use_fused_weighted_squared_relu = True
+    cfg.model.calculate_per_token_loss = True
+    cfg.model.moe_token_dispatcher_type = "flex"
+    cfg.model.moe_shared_expert_overlap = False
+    cfg.model.moe_flex_dispatcher_backend = "hybridep"
+    cfg.model.cuda_graph_impl = "none"
+    cfg.model.cuda_graph_scope = []
+    cfg.model.mtp_num_layers = 2
+    cfg.model.keep_mtp_spec_in_bf16 = True
+    cfg.model.mtp_loss_scaling_factor = 0.3
+    cfg.model.mtp_use_repeated_layer = True
+    cfg.model.use_te_rng_tracker = True
+
+    cfg.tokenizer.tokenizer_model = model_source
+    cfg.dataset.seq_length = seq_length
+    cfg.dataset.blend = None
+    cfg.dataset.num_workers = 1
+    cfg.dataset.mmap_bin_files = False
+
+    cfg.train.train_iters = 39735
+    cfg.train.global_batch_size = 3072
+    cfg.train.micro_batch_size = 1
+    cfg.train.manual_gc = False
+    cfg.train.manual_gc_interval = 0
+    cfg.validation.eval_interval = 1000
+
+    cfg.model.transformer_impl = "transformer_engine"
+    cfg.model.cross_entropy_fusion_impl = "te"
+    cfg.mixed_precision = "bf16_mixed"
+
+    cfg.optimizer.lr = 2.5e-4
+    cfg.optimizer.min_lr = 2.5e-4
+    cfg.optimizer.weight_decay = 0.1
+    cfg.optimizer.adam_beta1 = 0.9
+    cfg.optimizer.adam_beta2 = 0.95
+    cfg.optimizer.adam_eps = 1e-8
+    cfg.scheduler.lr_warmup_iters = 0
+    cfg.scheduler.start_weight_decay = 0.1
+    cfg.scheduler.end_weight_decay = 0.1
+    cfg.scheduler.lr_decay_style = "constant"
+
+    cfg.checkpoint.save_interval = 200
+    cfg.checkpoint.ckpt_assume_constant_structure = True
+    cfg.checkpoint.dist_ckpt_strictness = "log_all"
+    cfg.checkpoint.async_save = True
+    cfg.checkpoint.async_strategy = "mcore"
+
+    cfg.ddp.overlap_grad_reduce = True
+    cfg.ddp.overlap_param_gather = True
+    cfg.ddp.check_for_nan_in_grad = True
+    cfg.ddp.use_distributed_optimizer = True
+    cfg.ddp.average_in_collective = False
+
+    return cfg
+
+
+def nemotron_3_ultra_sft_openmathinstruct2_packed_config(
+    *,
+    hf_path: str | None = None,
+    seq_length: int = 4096,
+) -> ConfigContainer:
+    """Return a packed OpenMathInstruct-2 full SFT config for Nemotron 3 Ultra.
+
+    Args:
+        hf_path: Optional Hugging Face model ID or local snapshot path.
+        seq_length: Packed sequence length.
+
+    Returns:
+        Full-parameter SFT configuration for OpenMathInstruct-2.
+    """
+    cfg = _sft_common()
+    model_source = hf_path or NEMOTRON_3_ULTRA_HF_MODEL_ID
+
+    cfg.model = AutoBridge.from_hf_pretrained(model_source).to_megatron_provider(load_weights=False)
+    cfg.model.tensor_model_parallel_size = 2
+    cfg.model.pipeline_model_parallel_size = 6
+    cfg.model.pipeline_dtype = torch.bfloat16
+    cfg.model.virtual_pipeline_model_parallel_size = None
+    cfg.model.context_parallel_size = 1
+    cfg.model.sequence_parallel = True
+    cfg.model.expert_tensor_parallel_size = 1
+    cfg.model.expert_model_parallel_size = 32
+    cfg.model.pipeline_model_parallel_layout = None
+    cfg.model.seq_length = seq_length
+    cfg.model.apply_rope_fusion = False
+    cfg.model.attention_backend = "fused"
+    cfg.model.gradient_accumulation_fusion = True
+    cfg.model.init_method_std = 0.014
+    cfg.model.use_fused_weighted_squared_relu = True
+    cfg.model.calculate_per_token_loss = True
+    cfg.model.moe_token_dispatcher_type = "flex"
+    cfg.model.moe_shared_expert_overlap = False
+    cfg.model.moe_flex_dispatcher_backend = "hybridep"
+    cfg.model.cuda_graph_impl = "none"
+    cfg.model.cuda_graph_scope = []
+    cfg.model.mtp_num_layers = 2
+    cfg.model.keep_mtp_spec_in_bf16 = True
+    cfg.model.mtp_loss_scaling_factor = 0.3
+    cfg.model.mtp_use_repeated_layer = True
+    cfg.model.use_te_rng_tracker = True
+    cfg.model.recompute_granularity = "selective"
+    cfg.model.recompute_method = None
+    cfg.model.recompute_num_layers = None
+    cfg.model.recompute_modules = ["moe", "layernorm", "core_attn", "moe_act"]
+
+    cfg.tokenizer.tokenizer_model = model_source
+    cfg.dataset = default_openmathinstruct2_config(seq_length=seq_length, packed_sequence=True)
+    if cfg.dataset.packed_sequence_specs is not None:
+        cfg.dataset.packed_sequence_specs.packed_sequence_size = seq_length
+        cfg.dataset.packed_sequence_specs.tokenizer_model_name = NEMOTRON_3_ULTRA_TOKENIZER_NAME
+
+    cfg.train.train_iters = 1000
+    cfg.train.global_batch_size = 128
+    cfg.train.micro_batch_size = 1
+    cfg.validation.eval_interval = 50
+    cfg.validation.eval_iters = 32
+
+    cfg.optimizer.lr = 5e-6
+    cfg.optimizer.min_lr = 5e-7
+    cfg.optimizer.adam_beta1 = 0.9
+    cfg.optimizer.adam_beta2 = 0.98
+    cfg.optimizer.adam_eps = 1e-8
+    cfg.optimizer.weight_decay = 0.1
+    cfg.scheduler.start_weight_decay = 0.1
+    cfg.scheduler.end_weight_decay = 0.1
+    cfg.scheduler.lr_decay_style = "cosine"
+    cfg.scheduler.lr_warmup_iters = 250
+    cfg.scheduler.lr_decay_iters = 1000
+
+    cfg.checkpoint.save_interval = 250
+    cfg.checkpoint.ckpt_format = "torch_dist"
+    cfg.checkpoint.dist_ckpt_strictness = "log_all"
+    cfg.checkpoint.ckpt_assume_constant_structure = True
+    cfg.checkpoint.async_save = True
+    cfg.checkpoint.async_strategy = "mcore"
+
+    cfg.logger.log_interval = 1
+    cfg.rng.seed = 5678
+
+    cfg.ddp.check_for_nan_in_grad = True
+    cfg.ddp.grad_reduce_in_fp32 = True
+    cfg.ddp.overlap_grad_reduce = True
+    cfg.ddp.overlap_param_gather = True
+    cfg.ddp.use_distributed_optimizer = True
+
+    return cfg
+
+
+def nemotron_3_ultra_peft_openmathinstruct2_packed_config(
+    *,
+    peft: str | PEFT | None = "lora",
+    hf_path: str | None = None,
+    seq_length: int = 4096,
+) -> ConfigContainer:
+    """Return a packed OpenMathInstruct-2 PEFT config for Nemotron 3 Ultra.
+
+    Args:
+        peft: PEFT scheme, PEFT instance, or "none".
+        hf_path: Optional Hugging Face model ID or local snapshot path.
+        seq_length: Packed sequence length.
+
+    Returns:
+        PEFT configuration for OpenMathInstruct-2.
+    """
+    cfg = _peft_common()
+    model_source = hf_path or NEMOTRON_3_ULTRA_HF_MODEL_ID
+
+    cfg.model = AutoBridge.from_hf_pretrained(model_source).to_megatron_provider(load_weights=False)
+    cfg.model.tensor_model_parallel_size = 2
+    cfg.model.pipeline_model_parallel_size = 4
+    cfg.model.pipeline_dtype = torch.bfloat16
+    cfg.model.virtual_pipeline_model_parallel_size = None
+    cfg.model.context_parallel_size = 1
+    cfg.model.sequence_parallel = True
+    cfg.model.expert_tensor_parallel_size = 1
+    cfg.model.expert_model_parallel_size = 8
+    cfg.model.pipeline_model_parallel_layout = None
+    cfg.model.seq_length = seq_length
+    cfg.model.apply_rope_fusion = False
+    cfg.model.attention_backend = "fused"
+    cfg.model.gradient_accumulation_fusion = True
+    cfg.model.init_method_std = 0.014
+    cfg.model.use_fused_weighted_squared_relu = True
+    cfg.model.calculate_per_token_loss = True
+    cfg.model.moe_token_dispatcher_type = "flex"
+    cfg.model.moe_shared_expert_overlap = False
+    cfg.model.moe_flex_dispatcher_backend = "hybridep"
+    cfg.model.cuda_graph_impl = "none"
+    cfg.model.cuda_graph_scope = []
+    cfg.model.mtp_num_layers = 2
+    cfg.model.keep_mtp_spec_in_bf16 = True
+    cfg.model.mtp_loss_scaling_factor = 0.3
+    cfg.model.mtp_use_repeated_layer = True
+    cfg.model.use_te_rng_tracker = True
+    cfg.model.recompute_granularity = "selective"
+    cfg.model.recompute_method = None
+    cfg.model.recompute_num_layers = None
+    cfg.model.recompute_modules = ["moe", "layernorm", "core_attn", "moe_act", "mlp", "shared_experts"]
+
+    target_modules = [
+        "linear_qkv",
+        "linear_proj",
+        "linear_fc1",
+        "linear_fc2",
+        "in_proj",
+        "out_proj",
+    ]
+    cfg.peft = default_peft_config(peft, target_modules=target_modules)
+
+    cfg.tokenizer.tokenizer_model = model_source
+    cfg.dataset = default_openmathinstruct2_config(seq_length=seq_length, packed_sequence=True)
+    if cfg.dataset.packed_sequence_specs is not None:
+        cfg.dataset.packed_sequence_specs.packed_sequence_size = seq_length
+        cfg.dataset.packed_sequence_specs.tokenizer_model_name = NEMOTRON_3_ULTRA_TOKENIZER_NAME
+
+    cfg.train.train_iters = 1000
+    cfg.train.global_batch_size = 128
+    cfg.train.micro_batch_size = 1
+    cfg.validation.eval_interval = 50
+    cfg.validation.eval_iters = 32
+
+    cfg.optimizer.lr = 1e-4
+    cfg.optimizer.min_lr = 1e-5
+    cfg.optimizer.adam_beta1 = 0.9
+    cfg.optimizer.adam_beta2 = 0.98
+    cfg.optimizer.adam_eps = 1e-8
+    cfg.optimizer.weight_decay = 0.1
+    cfg.scheduler.start_weight_decay = 0.1
+    cfg.scheduler.end_weight_decay = 0.1
+    cfg.scheduler.lr_decay_style = "cosine"
+    cfg.scheduler.lr_warmup_iters = 250
+    cfg.scheduler.lr_decay_iters = 1000
+
+    cfg.checkpoint.save_interval = 250
+    cfg.checkpoint.ckpt_format = "torch_dist"
+    cfg.checkpoint.dist_ckpt_strictness = "log_all"
+    cfg.checkpoint.ckpt_assume_constant_structure = True
+    cfg.checkpoint.async_save = True
+    cfg.checkpoint.async_strategy = "nvrx"
+
+    cfg.logger.log_interval = 1
+    cfg.rng.seed = 5678
+
+    cfg.ddp.check_for_nan_in_grad = True
+    cfg.ddp.grad_reduce_in_fp32 = True
+    cfg.ddp.overlap_grad_reduce = True
+    cfg.ddp.overlap_param_gather = True
+    cfg.ddp.use_distributed_optimizer = True
+
+    return cfg
+
+
+__all__ = [
+    "nemotron_3_ultra_pretrain_config",
+    "nemotron_3_ultra_sft_openmathinstruct2_packed_config",
+    "nemotron_3_ultra_peft_openmathinstruct2_packed_config",
+]
diff --git a/tests/unit_tests/recipes/nemotronh/test_nemotron_3_ultra.py b/tests/unit_tests/recipes/nemotronh/test_nemotron_3_ultra.py
new file mode 100644
index 0000000000..3546790849
--- /dev/null
+++ b/tests/unit_tests/recipes/nemotronh/test_nemotron_3_ultra.py
@@ -0,0 +1,155 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import importlib
+
+import pytest
+
+from megatron.bridge.recipes.nemotronh.nemotron_3_ultra import (
+    NEMOTRON_3_ULTRA_TOKENIZER_NAME,
+    nemotron_3_ultra_peft_openmathinstruct2_packed_config,
+    nemotron_3_ultra_pretrain_config,
+    nemotron_3_ultra_sft_openmathinstruct2_packed_config,
+)
+
+
+class _FakeUltraProvider:
+    """Fake model provider for testing recipe field overrides without HF Hub I/O."""
+
+    def __init__(self) -> None:
+        self.vocab_size = 256
+
+    def finalize(self) -> None:
+        return None
+
+
+class _FakeAutoBridge:
+    """Fake AutoBridge that returns an Ultra provider without loading a model."""
+
+    @classmethod
+    def from_hf_pretrained(cls, *args, **kwargs):
+        return cls()
+
+    def to_megatron_provider(self, *args, **kwargs):
+        return _FakeUltraProvider()
+
+
+@pytest.fixture(autouse=True)
+def _patch_autobridge(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Patch AutoBridge in the recipe module to avoid Hugging Face access."""
+    mod = importlib.import_module("megatron.bridge.recipes.nemotronh.nemotron_3_ultra")
+    monkeypatch.setattr(mod, "AutoBridge", _FakeAutoBridge)
+
+
+@pytest.mark.unit
+def test_pretrain_uses_initial_parallelism_values() -> None:
+    cfg = nemotron_3_ultra_pretrain_config()
+
+    assert cfg.model.tensor_model_parallel_size == 1
+    assert cfg.model.pipeline_model_parallel_size == 3
+    assert cfg.model.expert_model_parallel_size == 8
+    assert cfg.model.moe_token_dispatcher_type == "flex"
+    assert cfg.model.moe_flex_dispatcher_backend == "hybridep"
+    assert cfg.model.sequence_parallel is True
+    assert cfg.model.virtual_pipeline_model_parallel_size is None
+    assert cfg.model.mtp_num_layers == 2
+    assert cfg.model.mtp_loss_scaling_factor == 0.3
+    assert cfg.model.mtp_use_repeated_layer is True
+
+    assert cfg.optimizer.lr == 2.5e-4
+    assert cfg.optimizer.min_lr == 2.5e-4
+    assert cfg.optimizer.weight_decay == 0.1
+    assert cfg.scheduler.lr_decay_style == "constant"
+    assert cfg.scheduler.lr_warmup_iters == 0
+    assert cfg.checkpoint.async_save is True
+    assert cfg.checkpoint.async_strategy == "mcore"
+    assert cfg.checkpoint.save_interval == 200
+
+    assert cfg.train.global_batch_size == 3072
+    assert cfg.train.micro_batch_size == 1
+    assert cfg.dataset.seq_length == 8192
+    assert cfg.dataset.blend is None
+
+
+@pytest.mark.unit
+def test_openmath_sft_uses_initial_parallelism_values() -> None:
+    cfg = nemotron_3_ultra_sft_openmathinstruct2_packed_config()
+
+    assert cfg.model.tensor_model_parallel_size == 2
+    assert cfg.model.pipeline_model_parallel_size == 6
+    assert cfg.model.expert_model_parallel_size == 32
+    assert cfg.model.moe_token_dispatcher_type == "flex"
+    assert cfg.model.moe_flex_dispatcher_backend == "hybridep"
+    assert cfg.model.sequence_parallel is True
+    assert cfg.model.virtual_pipeline_model_parallel_size is None
+    assert cfg.model.recompute_granularity == "selective"
+    assert cfg.model.recompute_method is None
+    assert cfg.model.recompute_num_layers is None
+    assert cfg.model.recompute_modules == ["moe", "layernorm", "core_attn", "moe_act"]
+
+    assert cfg.train.train_iters == 1000
+    assert cfg.train.global_batch_size == 128
+    assert cfg.checkpoint.async_save is True
+    assert cfg.checkpoint.async_strategy == "mcore"
+    assert cfg.dataset.dataset_name == "nvidia/OpenMathInstruct-2"
+    assert cfg.dataset.packed_sequence_specs.packed_sequence_size == 4096
+    assert cfg.dataset.packed_sequence_specs.tokenizer_model_name == NEMOTRON_3_ULTRA_TOKENIZER_NAME
+
+
+@pytest.mark.unit
+def test_openmath_peft_uses_validated_parallelism_values() -> None:
+    cfg = nemotron_3_ultra_peft_openmathinstruct2_packed_config()
+
+    assert cfg.model.tensor_model_parallel_size == 2
+    assert cfg.model.pipeline_model_parallel_size == 4
+    assert cfg.model.expert_model_parallel_size == 8
+    assert cfg.model.moe_token_dispatcher_type == "flex"
+    assert cfg.model.moe_flex_dispatcher_backend == "hybridep"
+    assert cfg.model.sequence_parallel is True
+    assert cfg.model.virtual_pipeline_model_parallel_size is None
+    assert cfg.model.recompute_granularity == "selective"
+    assert cfg.model.recompute_method is None
+    assert cfg.model.recompute_num_layers is None
+    assert cfg.model.recompute_modules == ["moe", "layernorm", "core_attn", "moe_act", "mlp", "shared_experts"]
+
+    assert cfg.optimizer.lr == 1e-4
+    assert cfg.optimizer.min_lr == 1e-5
+    assert cfg.train.train_iters == 1000
+    assert cfg.train.global_batch_size == 128
+    assert cfg.checkpoint.async_save is True
+    assert cfg.checkpoint.async_strategy == "nvrx"
+
+
+@pytest.mark.unit
+def test_openmath_peft_none_disables_adapter() -> None:
+    cfg = nemotron_3_ultra_peft_openmathinstruct2_packed_config(peft="none")
+    assert cfg.peft is None
+
+
+@pytest.mark.unit
+def test_openmath_peft_recompute_modules_are_not_shared() -> None:
+    cfg = nemotron_3_ultra_peft_openmathinstruct2_packed_config()
+    cfg.model.recompute_modules.append("sentinel")
+
+    fresh_cfg = nemotron_3_ultra_peft_openmathinstruct2_packed_config()
+    assert fresh_cfg.model.recompute_modules == ["moe", "layernorm", "core_attn", "moe_act", "mlp", "shared_experts"]
+
+
+@pytest.mark.unit
+def test_openmath_sft_recompute_modules_are_not_shared() -> None:
+    cfg = nemotron_3_ultra_sft_openmathinstruct2_packed_config()
+    cfg.model.recompute_modules.append("sentinel")
+
+    fresh_cfg = nemotron_3_ultra_sft_openmathinstruct2_packed_config()
+    assert fresh_cfg.model.recompute_modules == ["moe", "layernorm", "core_attn", "moe_act"]

From fc2a92d70de6d1b562d76eac05afdca57fd4fd60 Mon Sep 17 00:00:00 2001
From: Chen Cui <chcui@nvidia.com>
Date: Wed, 24 Jun 2026 11:05:04 -0700
Subject: [PATCH 5/5] fix(peft): complete LoRA backport fixes

Signed-off-by: Chen Cui <chcui@nvidia.com>
---
 .../bridge/models/conversion/peft_bridge.py   | 94 +++++++++++++++++++
 src/megatron/bridge/peft/utils.py             |  2 +-
 tests/unit_tests/peft/test_canonical_lora.py  |  9 +-
 tests/unit_tests/peft/test_lora.py            | 12 +--
 4 files changed, 104 insertions(+), 13 deletions(-)

diff --git a/src/megatron/bridge/models/conversion/peft_bridge.py b/src/megatron/bridge/models/conversion/peft_bridge.py
index d9bde27748..5b1c758690 100644
--- a/src/megatron/bridge/models/conversion/peft_bridge.py
+++ b/src/megatron/bridge/models/conversion/peft_bridge.py
@@ -341,6 +341,24 @@ def _infer_hf_expert_idx(self, hf_name: str) -> Optional[int]:
         except ValueError:
             return None
 
+    def _strip_hf_expert_index(self, hf_name: str) -> str:
+        """Drop the ``experts.<idx>`` index from an HF MoE weight name.
+
+        A shared-outer adapter's shared side is replicated across experts, so it
+        is exported under the expert-agnostic name (``experts.gate_proj`` rather
+        than ``experts.0.gate_proj``) that the serving loader keys its 3D-shared
+        branch on. Mirrors :meth:`_infer_hf_expert_idx`'s name parsing.
+        """
+
+        parts = hf_name.split(".")
+        try:
+            experts_idx = parts.index("experts")
+        except ValueError:
+            return hf_name
+        if experts_idx + 1 < len(parts) and parts[experts_idx + 1].isdigit():
+            del parts[experts_idx + 1]
+        return ".".join(parts)
+
     def _split_qkv_linear_out_weight(
         self,
         megatron_model: Union[MegatronModel, List[MegatronModel]],
@@ -856,6 +874,20 @@ def stream_adapter_weights_megatron_to_hf(
             linear_out_tensor = adapter_weight.linear_out_weight.weight
             is_expert = is_expert_linear(adapter_task.global_base_prefix)
             is_grouped_expert = is_expert and ".local_experts." not in adapter_task.global_base_prefix
+            is_shared_outer_lora = is_grouped_expert and linear_in_tensor.ndim != linear_out_tensor.ndim
+
+            if is_shared_outer_lora:
+                yield from self._stream_shared_outer_adapter_weights(
+                    megatron_model,
+                    mapping_registry,
+                    adapter_task,
+                    linear_in_tensor,
+                    linear_out_tensor,
+                    num_moe_experts,
+                    cpu,
+                )
+                continue
+
             expert_linear_in_gathered = None
             expert_linear_out_gathered = None
             if is_grouped_expert:
@@ -983,6 +1015,68 @@ def stream_adapter_weights_megatron_to_hf(
                 yield HFWeightTuple(linear_in_hf_names[0], current_linear_in_tensor)
                 yield HFWeightTuple(linear_out_hf_names[0], current_linear_out_tensor)
 
+    def _stream_shared_outer_adapter_weights(
+        self,
+        megatron_model: List[MegatronModel],
+        mapping_registry: "MegatronMappingRegistry",
+        adapter_task: AdapterWeightConversionTask,
+        linear_in_tensor: torch.Tensor,
+        linear_out_tensor: torch.Tensor,
+        num_moe_experts: int,
+        cpu: bool,
+    ) -> Iterable["HFWeightTuple"]:
+        """Stream a shared-outer grouped-expert LoRA adapter (SGLang PR #21466).
+
+        One side is a 2D LoRA matrix replicated across local experts; the other
+        is a per-expert 3D pack. The shared side is emitted once as a ``[1, ...]``
+        tensor under the expert-agnostic HF name (so the serving loader takes its
+        3D-shared branch); the per-expert side is gathered across EP ranks and
+        emitted once per global expert.
+        """
+
+        from megatron.bridge.models.conversion.model_bridge import HFWeightTuple
+
+        is_expert = is_expert_linear(adapter_task.global_base_prefix)
+        for side_tensor, side_suffix in (
+            (linear_in_tensor, ".linear_in.weight"),
+            (linear_out_tensor, ".linear_out.weight"),
+        ):
+            if side_tensor.ndim == 2:
+                current = side_tensor.cpu() if cpu else side_tensor
+                current = current.unsqueeze(0)
+
+                base_hf_weight_names = self._get_base_hf_param_names_for_adapter(
+                    mapping_registry, adapter_task.global_base_prefix, adapter_task.adapter_key, ".weight0"
+                )
+                for base_name in base_hf_weight_names:
+                    hf_name = self._make_lora_param_name(self._strip_hf_expert_index(base_name), side_suffix)
+                    yield HFWeightTuple(hf_name, current)
+                continue
+
+            gathered = self._gather_expert_adapter_weight(side_tensor)
+            for expert_idx in range(num_moe_experts):
+                current = self._select_expert_adapter_weight(side_tensor, gathered, expert_idx, num_moe_experts)
+                if cpu:
+                    current = current.cpu()
+
+                base_hf_weight_names = self._get_base_hf_param_names_for_adapter(
+                    mapping_registry, adapter_task.global_base_prefix, adapter_task.adapter_key, f".weight{expert_idx}"
+                )
+                side_hf_names = [self._make_lora_param_name(name, side_suffix) for name in base_hf_weight_names]
+
+                per_base = None
+                if side_suffix == ".linear_out.weight" and adapter_task.adapter_key is None:
+                    per_base = self._get_fused_adapter_linear_out_slices(
+                        megatron_model, base_hf_weight_names, current, is_expert=is_expert
+                    )
+                if per_base is None:
+                    yield HFWeightTuple(side_hf_names[0], current)
+                    continue
+                for index, base_name in enumerate(base_hf_weight_names):
+                    chunk = per_base.get(base_name)
+                    assert chunk is not None, f"unknown projection name: {base_name!r}"
+                    yield HFWeightTuple(side_hf_names[index], chunk)
+
     def _get_fused_adapter_linear_out_slices(
         self,
         megatron_model: List[MegatronModel],
diff --git a/src/megatron/bridge/peft/utils.py b/src/megatron/bridge/peft/utils.py
index 87227c7f8b..2e7a5a9f3e 100644
--- a/src/megatron/bridge/peft/utils.py
+++ b/src/megatron/bridge/peft/utils.py
@@ -2259,7 +2259,7 @@ def sharded_state_dict(
         key = f"{prefix}weight"
         return {
             key: _make_grouped_expert_sharded_tensor(
-                self.weight.data, key, tp_axis=None, sharded_offsets=sharded_offsets
+                self.weight.data, key, tp_axis=None, sharded_offsets=sharded_offsets, pg_collection=None
             )
         }
 
diff --git a/tests/unit_tests/peft/test_canonical_lora.py b/tests/unit_tests/peft/test_canonical_lora.py
index c3aeb489a0..2e6f30e815 100644
--- a/tests/unit_tests/peft/test_canonical_lora.py
+++ b/tests/unit_tests/peft/test_canonical_lora.py
@@ -14,6 +14,7 @@
 
 import datetime
 import os
+from types import SimpleNamespace
 from unittest.mock import patch
 
 import megatron.core.parallel_state as parallel_state
@@ -461,6 +462,10 @@ def mock_get_attrs(module, is_expert=False):
     def test_canonical_lora_normalize_moe_lora_aligns_expert_dim_to_expert_tp(self):
         """Normalized canonical expert fc1 adapters should round up to the expert-TP granularity when needed."""
         model = MoEMegatronStyleModel(moe_router_topk=8)
+        for module in model.modules():
+            if hasattr(module, "config"):
+                module.config.expert_tensor_parallel_size = 2
+                module.config._pg_collection = SimpleNamespace(expt_tp=SimpleNamespace(size=lambda: 2))
         lora = CanonicalLoRA(
             target_modules=["linear_fc1_up", "linear_fc1_gate"],
             dim=8,
@@ -482,10 +487,6 @@ def mock_get_attrs(module, is_expert=False):
                 "megatron.bridge.peft.canonical_lora.get_adapter_attributes_from_linear",
                 side_effect=mock_get_attrs,
             ),
-            patch(
-                "megatron.bridge.peft.utils.parallel_state.get_expert_tensor_parallel_world_size",
-                return_value=2,
-            ),
             patch("megatron.bridge.peft.canonical_lora.ParallelLinearAdapter") as mock_adapter,
         ):
             mock_adapter.return_value = nn.Linear(1, 1)
diff --git a/tests/unit_tests/peft/test_lora.py b/tests/unit_tests/peft/test_lora.py
index beb76c9c78..26465e0750 100644
--- a/tests/unit_tests/peft/test_lora.py
+++ b/tests/unit_tests/peft/test_lora.py
@@ -709,6 +709,9 @@ def mock_get_attrs(module, is_expert=False):
     def test_normalize_moe_lora_aligns_shared_expert_dim_to_expert_tp(self):
         """Normalized expert fc1 adapters should round up to the expert-TP granularity when needed."""
         model = MoEModel(moe_router_topk=8)
+        for module in model.modules():
+            if hasattr(module, "config"):
+                module.config.expert_tensor_parallel_size = 2
         lora = LoRA(target_modules=["linear_fc1"], dim=8, normalize_moe_lora=True)
 
         def mock_get_attrs(module, is_expert=False):
@@ -723,10 +726,6 @@ def mock_get_attrs(module, is_expert=False):
 
         with (
             patch("megatron.bridge.peft.lora.get_adapter_attributes_from_linear", side_effect=mock_get_attrs),
-            patch(
-                "megatron.bridge.peft.lora.parallel_state.get_expert_tensor_parallel_world_size",
-                return_value=2,
-            ),
             patch("megatron.bridge.peft.lora.ParallelLinearAdapter") as mock_adapter,
         ):
             mock_adapter.return_value = nn.Linear(1, 1)
@@ -747,6 +746,7 @@ def test_lora_grouped_expert_normalized_dim_aligns_to_expert_tp(self):
         """Per-expert grouped adapters should round normalized dims up to expert-TP granularity."""
         model = GroupedExpertModel()
         model.decoder.layers[0].mlp.experts.linear_fc2.config.moe_router_topk = 8
+        model.decoder.layers[0].mlp.experts.linear_fc2.config.expert_tensor_parallel_size = 2
         lora = LoRA(target_modules=["linear_fc2"], dim=8, normalize_moe_lora=True, share_expert_adapters=False)
 
         def mock_get_attrs(module, is_expert=False):
@@ -761,10 +761,6 @@ def mock_get_attrs(module, is_expert=False):
 
         with (
             patch("megatron.bridge.peft.lora.get_adapter_attributes_from_linear", side_effect=mock_get_attrs),
-            patch(
-                "megatron.bridge.peft.lora.parallel_state.get_expert_tensor_parallel_world_size",
-                return_value=2,
-            ),
             patch("megatron.bridge.peft.lora.GroupedExpertLinearAdapter") as mock_adapter,
         ):
             mock_adapter.return_value = nn.Identity()