diff --git a/3rdparty/Megatron-LM b/3rdparty/Megatron-LM
index 6513e3e23d..d57d882fd0 160000
--- a/3rdparty/Megatron-LM
+++ b/3rdparty/Megatron-LM
@@ -1 +1 @@
-Subproject commit 6513e3e23d6b5eda6a1c934990b15e804237732b
+Subproject commit d57d882fd08ee1762eab6144260cf9693f58ac9d
diff --git a/skills/create-model-verification-card/SKILL.md b/skills/create-model-verification-card/SKILL.md
index 6d782f2372..aafb5193a8 100644
--- a/skills/create-model-verification-card/SKILL.md
+++ b/skills/create-model-verification-card/SKILL.md
@@ -1,6 +1,6 @@
 ---
 name: create-model-verification-card
-description: Create or update concise, agent-readable Megatron Bridge model verification cards. Use when adding a model support card, auditing cross-model convergence comparability or verification coverage, recording conversion, deterministic inference, training, checkpoint resume, post-SFT export, or performance results, or preparing a model-support PR. Enforce the required core inventory, convergence-versus-performance contracts, optional canonical performance item, public Slurm launcher commands, training metrics, important-feature allowlist, and a strict privacy boundary that excludes private runtime wiring, internal paths, credentials, and job metadata.
+description: Create or update concise, agent-readable Megatron Bridge model verification cards. Use when adding a model support card, auditing cross-model convergence comparability or verification coverage, recording conversion, deterministic inference, training, checkpoint resume, post-SFT export, or performance results, or preparing a model-support PR. Enforce publicly reachable clean source provenance, upstream ownership of discovered fixes, the required core inventory, convergence-versus-performance contracts, public Slurm launcher commands, training metrics, important-feature allowlist, and a strict privacy boundary.
 ---
 
 # Create Model Verification Card
@@ -45,6 +45,42 @@ the leaf's optional `bridge_commit` field. Omit the field when it would repeat
 the top-level value, and never use a commit field to disguise uncommitted
 runtime changes. Items that are not verified must not carry a commit override.
 
+#### Enforce source integrity and upstream ownership
+
+Treat source provenance as a verification gate, not a reporting detail. Every
+verification run must use either an upstream commit or the exact pushed head of
+an open upstream PR. Before and after the run, require a clean tracked source
+tree and record the exact Bridge commit plus every relevant submodule and
+dependency revision in the private durable run record. Confirm that each
+recorded commit is reachable from its stated public remote or PR; a copied
+source directory without public Git provenance is insufficient.
+
+Do not count a run that depends on an uncommitted edit, unpublished commit,
+source overlay, monkeypatch, bind-mounted replacement file, or locally rebuilt
+dependency as model verification. Do not reclassify such a run as a
+feasibility experiment. It is an invalid verification attempt and leaves the
+item `unverified`.
+
+When verification exposes a product defect:
+
+1. Stop the affected verification item and record the unchanged failing
+   command, exact public source revisions, failure, and owning upstream
+   repository in the private durable record.
+2. Implement the fix through the owning repository's normal branch, test,
+   review, and PR workflow. Never keep a required fix only in the verification
+   checkout.
+3. Label runs of a pushed fix as candidate-PR validation and link the upstream
+   PR and exact head commit. A local patched run cannot satisfy any card
+   checkbox, metric, or expected result.
+4. Rerun the original failing workload from the exact clean, pushed PR head.
+   Only that rerun may become verification evidence; after further edits, rerun
+   from the new pushed head.
+
+Do not hide a newly discovered blocker by changing the command, container,
+model, dataset, scale, or expected result. If another layer is responsible,
+file or link its upstream issue or PR and keep the card item `unverified` until
+the original workload passes on publicly reviewable source.
+
 ### 2. Create the core inventory and add performance when available
 
 Include these twelve required items, even when their status is `unsupported` or
@@ -760,6 +796,13 @@ an item verified merely to make validation pass.
 - Pin a public immutable HF revision, minimum Transformers version, public base
   container, and exact Bridge verification commit; use an item override only
   for a verified workload run from a different clean commit.
+- Require every verification commit and dependency revision to be publicly
+  reachable from upstream or an open upstream PR. Reject local patches,
+  unpublished commits, source overlays, monkeypatches, replacement mounts, and
+  copied trees without public Git provenance as verification evidence.
+- For every defect found during verification, preserve the unchanged failure,
+  link the owning upstream issue or PR, and rerun the original workload from
+  the exact clean pushed fix commit before changing the item to `verified`.
 - Use the public model name in commands.
 - Include commands and concrete expected results for verified items.
 - For manual forward pass, require a next-token match and cosine similarity of
diff --git a/src/megatron/bridge/models/conversion/auto_bridge.py b/src/megatron/bridge/models/conversion/auto_bridge.py
index 3685b34b6a..8521fc6d1b 100644
--- a/src/megatron/bridge/models/conversion/auto_bridge.py
+++ b/src/megatron/bridge/models/conversion/auto_bridge.py
@@ -1215,6 +1215,7 @@ def _filter_quant(gen):
             ignored_source_key_prefixes = (
                 _mtp_source_key_prefixes(source, hf_config, model_config) if mtp_disabled else ()
             ) or None
+            source_key_replacements = bridge.hf_export_source_key_replacements(self.hf_pretrained) or None
             source.save_generator(
                 generator,
                 path,
@@ -1222,6 +1223,7 @@ def _filter_quant(gen):
                 distributed_save=distributed_save,
                 save_every_n_ranks=save_every_n_ranks,
                 ignored_source_key_prefixes=ignored_source_key_prefixes,
+                source_key_replacements=source_key_replacements,
             )
         else:
             # Config-only path: shard and write safetensors directly
@@ -1450,27 +1452,31 @@ def import_ckpt(
         # Load the HuggingFace model
         bridge = cls.from_hf_pretrained(hf_model_id, **kwargs)
 
-        # Convert to Megatron model
-        megatron_model = bridge.to_megatron_model(wrap_with_ddp=False, use_cpu_initialization=True)
+        from megatron.bridge.training.model_load_save import temporary_distributed_context
+
+        model_context = nullcontext() if dist.is_initialized() else temporary_distributed_context(backend="gloo")
+        with model_context:
+            # Convert to Megatron model
+            megatron_model = bridge.to_megatron_model(wrap_with_ddp=False, use_cpu_initialization=True)
 
-        # Save as Megatron checkpoint
-        hf_tokenizer_kwargs = {}
-        if hasattr(bridge._model_bridge, "get_hf_tokenizer_kwargs"):
-            hf_tokenizer_kwargs = bridge._model_bridge.get_hf_tokenizer_kwargs()
-        if hf_tokenizer_kwargs is None:
+            # Save as Megatron checkpoint
             hf_tokenizer_kwargs = {}
-        if kwargs.get("revision") is not None:
-            hf_tokenizer_kwargs.setdefault("revision", kwargs["revision"])
-        # Forward trust_remote_code to the tokenizer (needed for repos with custom code)
-        if kwargs.get("trust_remote_code"):
-            hf_tokenizer_kwargs.setdefault("trust_remote_code", True)
-        bridge.save_megatron_model(
-            megatron_model,
-            megatron_path,
-            hf_tokenizer_path=hf_model_id,
-            hf_tokenizer_kwargs=hf_tokenizer_kwargs,
-            low_memory_save=low_memory_save,
-        )
+            if hasattr(bridge._model_bridge, "get_hf_tokenizer_kwargs"):
+                hf_tokenizer_kwargs = bridge._model_bridge.get_hf_tokenizer_kwargs()
+            if hf_tokenizer_kwargs is None:
+                hf_tokenizer_kwargs = {}
+            if kwargs.get("revision") is not None:
+                hf_tokenizer_kwargs.setdefault("revision", kwargs["revision"])
+            # Forward trust_remote_code to the tokenizer (needed for repos with custom code)
+            if kwargs.get("trust_remote_code"):
+                hf_tokenizer_kwargs.setdefault("trust_remote_code", True)
+            bridge.save_megatron_model(
+                megatron_model,
+                megatron_path,
+                hf_tokenizer_path=hf_model_id,
+                hf_tokenizer_kwargs=hf_tokenizer_kwargs,
+                low_memory_save=low_memory_save,
+            )
 
     def export_ckpt(
         self,
diff --git a/src/megatron/bridge/models/conversion/model_bridge.py b/src/megatron/bridge/models/conversion/model_bridge.py
index 5cea6145f9..97ae4e6546 100644
--- a/src/megatron/bridge/models/conversion/model_bridge.py
+++ b/src/megatron/bridge/models/conversion/model_bridge.py
@@ -1011,6 +1011,22 @@ def maybe_modify_loaded_hf_weight(
             hf_weights = {k: hf_state_dict[v] for k, v in hf_param.items()}
         return hf_weights
 
+    def hf_export_source_key_replacements(self, hf_pretrained: HFPreTrained) -> Mapping[str, Iterable[str]]:
+        """Describe exported keys that intentionally replace source-checkpoint keys.
+
+        Streaming HF export normally preserves the source checkpoint's strict tensor-key
+        inventory and shard placement. Model families that export a different physical
+        representation can override this hook to replace multiple source keys with one
+        exported key while retaining the source shard assignment.
+
+        Args:
+            hf_pretrained: Hugging Face model wrapper containing the source state.
+
+        Returns:
+            Mapping from each exported key to the source keys it replaces.
+        """
+        return {}
+
     def maybe_modify_converted_hf_weight(
         self,
         task: WeightConversionTask,
diff --git a/src/megatron/bridge/models/conversion/utils.py b/src/megatron/bridge/models/conversion/utils.py
index f3ae70f4db..ca90cde663 100644
--- a/src/megatron/bridge/models/conversion/utils.py
+++ b/src/megatron/bridge/models/conversion/utils.py
@@ -46,13 +46,20 @@ def unwrap_model(model, module_instances=None):
     if module_instances is None:
         from megatron.core.distributed import DistributedDataParallel as DDP
         from megatron.core.distributed import TorchFullyShardedDataParallel as torch_FSDP
-        from megatron.core.distributed.fsdp.mcore_fsdp_adapter import (
-            FullyShardedDataParallel as megatron_FSDP,
-        )
+        from megatron.core.distributed.fsdp import mcore_fsdp_adapter
         from megatron.core.distributed.fsdp.src.megatron_fsdp.megatron_fsdp import MegatronFSDP
         from megatron.core.transformer.module import Float16Module
 
-        module_instances = (DDP, torch_FSDP, megatron_FSDP, Float16Module, MegatronFSDP)
+        megatron_fsdp = mcore_fsdp_adapter.FullyShardedDataParallel
+        if isinstance(megatron_fsdp, type):
+            megatron_fsdp_classes = (megatron_fsdp,)
+        else:
+            megatron_fsdp_classes = (
+                mcore_fsdp_adapter.FullyShardedDataParallelV1,
+                mcore_fsdp_adapter.FullyShardedDataParallelV2,
+            )
+
+        module_instances = (DDP, torch_FSDP, *megatron_fsdp_classes, Float16Module, MegatronFSDP)
 
     return_list = True
     if not isinstance(model, list):
diff --git a/src/megatron/bridge/models/gpt_oss/gpt_oss_bridge.py b/src/megatron/bridge/models/gpt_oss/gpt_oss_bridge.py
index 48d941a177..73ad14cd81 100644
--- a/src/megatron/bridge/models/gpt_oss/gpt_oss_bridge.py
+++ b/src/megatron/bridge/models/gpt_oss/gpt_oss_bridge.py
@@ -152,6 +152,28 @@ def maybe_modify_loaded_hf_weight(
             )
         return {k: hf_state_dict[v] for k, v in hf_param.items()}
 
+    def hf_export_source_key_replacements(self, hf_pretrained: PreTrainedCausalLM) -> Mapping[str, Tuple[str, str]]:
+        """Replace source MXFP4 block/scale pairs with dequantized export tensors."""
+        source = getattr(getattr(hf_pretrained, "state", None), "source", None)
+        if source is None or not hasattr(source, "get_all_keys"):
+            return {}
+
+        source_keys = set(source.get_all_keys())
+        replacements = {}
+        quantized_expert_suffixes = (
+            ".mlp.experts.gate_up_proj_blocks",
+            ".mlp.experts.down_proj_blocks",
+        )
+        for blocks_key in source_keys:
+            if not blocks_key.endswith(quantized_expert_suffixes):
+                continue
+            output_key = blocks_key.removesuffix("_blocks")
+            scales_key = f"{output_key}_scales"
+            if scales_key in source_keys:
+                replacements[output_key] = (blocks_key, scales_key)
+
+        return replacements
+
     def mapping_registry(self) -> MegatronMappingRegistry:
         """
         Return MegatronMappingRegistry containing parameter mappings from HF to Megatron format.
diff --git a/src/megatron/bridge/models/gpt_provider.py b/src/megatron/bridge/models/gpt_provider.py
index 6bb4a9145c..e5ea9f21dd 100644
--- a/src/megatron/bridge/models/gpt_provider.py
+++ b/src/megatron/bridge/models/gpt_provider.py
@@ -97,6 +97,11 @@ def local_layer_spec(config: "GPTModelProvider") -> ModuleSpec:
     )
 
 
+def _use_local_spec_for_cpu_only_initialization(config: "GPTModelProvider") -> bool:
+    """Return whether CPU initialization must avoid Transformer Engine modules."""
+    return bool(config.use_cpu_initialization) and not torch.cuda.is_available()
+
+
 def modelopt_transformer_layer_spec(config: "GPTModelProvider") -> ModuleSpec:
     """Layer specification for quantization with ModelOpt."""
     # arbitrary attention mask is used for speculative decoding training
@@ -119,6 +124,8 @@ def modelopt_transformer_layer_spec(config: "GPTModelProvider") -> ModuleSpec:
 
 def default_layer_spec(config: "GPTModelProvider") -> ModuleSpec:
     """Determine the most appropriate layer specification based on availability."""
+    if _use_local_spec_for_cpu_only_initialization(config):
+        return local_layer_spec(config)
     if config.use_transformer_engine_full_layer_spec:
         return transformer_engine_full_layer_spec(config)
     else:
@@ -352,6 +359,7 @@ def mtp_block_spec(config: "GPTModelProvider", vp_stage: Optional[int] = None) -
     if getattr(config, "mtp_num_layers", None):
         from megatron.core.models.gpt.gpt_layer_specs import get_gpt_mtp_block_spec
 
+        use_transformer_engine = not _use_local_spec_for_cpu_only_initialization(config)
         if isinstance(config.transformer_layer_spec, Callable):
             if "vp_stage" in inspect.signature(config.transformer_layer_spec).parameters:
                 spec = config.transformer_layer_spec(config, vp_stage=vp_stage)
@@ -368,12 +376,12 @@ def mtp_block_spec(config: "GPTModelProvider", vp_stage: Optional[int] = None) -
 
             decoder_layer_specs = get_gpt_decoder_layer_specs(
                 config,
-                use_transformer_engine=True,
+                use_transformer_engine=use_transformer_engine,
                 normalization=config.normalization,
                 qk_l2_norm=config.qk_l2_norm,
             )
             spec = decoder_layer_specs[-1]
-        return get_gpt_mtp_block_spec(config, spec, use_transformer_engine=True, vp_stage=vp_stage)
+        return get_gpt_mtp_block_spec(config, spec, use_transformer_engine=use_transformer_engine, vp_stage=vp_stage)
     else:
         return None
 
diff --git a/src/megatron/bridge/models/hf_pretrained/state.py b/src/megatron/bridge/models/hf_pretrained/state.py
index 1145f7fad3..ecd0b9c4e7 100644
--- a/src/megatron/bridge/models/hf_pretrained/state.py
+++ b/src/megatron/bridge/models/hf_pretrained/state.py
@@ -459,6 +459,45 @@ def _ignore_source_key_prefixes(
         prefixes = tuple(ignored_source_key_prefixes)
         return {key: filename for key, filename in key_to_filename_map.items() if not key.startswith(prefixes)}
 
+    @classmethod
+    def _prepare_output_key_map(
+        cls,
+        key_to_filename_map: Mapping[str, str] | None,
+        ignored_source_key_prefixes: Iterable[str] | None,
+        source_key_replacements: Mapping[str, Iterable[str]] | None,
+    ) -> Dict[str, str]:
+        """Build the strict output key map after intentional source-format changes."""
+        output_key_map = cls._ignore_source_key_prefixes(key_to_filename_map, ignored_source_key_prefixes)
+        if not source_key_replacements:
+            return output_key_map
+
+        for output_key, replaced_keys_iterable in source_key_replacements.items():
+            replaced_keys = tuple(replaced_keys_iterable)
+            if not replaced_keys:
+                raise ValueError(f"Source-key replacement for '{output_key}' must contain at least one source key.")
+
+            missing_keys = [key for key in replaced_keys if key not in output_key_map]
+            if missing_keys:
+                raise KeyError(
+                    f"Cannot replace source keys for '{output_key}'; missing from the source sharding map: "
+                    f"{missing_keys}."
+                )
+            if output_key in output_key_map and output_key not in replaced_keys:
+                raise ValueError(f"Cannot replace source keys with existing output key '{output_key}'.")
+
+            source_filenames = {output_key_map[key] for key in replaced_keys}
+            if len(source_filenames) != 1:
+                raise ValueError(
+                    f"Source keys replaced by '{output_key}' span multiple shards: {sorted(source_filenames)}."
+                )
+
+            output_filename = source_filenames.pop()
+            for replaced_key in replaced_keys:
+                del output_key_map[replaced_key]
+            output_key_map[output_key] = output_filename
+
+        return output_key_map
+
     @property
     def path(self) -> Path:
         """
@@ -690,6 +729,7 @@ def save_generator(
         distributed_save: bool = False,
         save_every_n_ranks: int = 1,
         ignored_source_key_prefixes: Iterable[str] | None = None,
+        source_key_replacements: Mapping[str, Iterable[str]] | None = None,
     ):
         """
         Saves tensors from a generator to `.safetensors` files, preserving the
@@ -718,6 +758,8 @@ def save_generator(
                 For example, if set to 2, only ranks 0, 2, 4, ... will save weights.
             ignored_source_key_prefixes: Source tensor key prefixes to omit from the expected
                 source sharding map when saving.
+            source_key_replacements: Mapping from an exported tensor key to source tensor keys
+                that it intentionally replaces. Replaced keys must share one source shard.
 
         """
         if distributed_save:
@@ -727,6 +769,7 @@ def save_generator(
                 strict,
                 save_every_n_ranks=save_every_n_ranks,
                 ignored_source_key_prefixes=ignored_source_key_prefixes,
+                source_key_replacements=source_key_replacements,
             )
 
         # In a distributed environment, only rank 0 should write to disk.
@@ -746,7 +789,11 @@ def save_generator(
         output_path = Path(output_path)
         output_path.mkdir(parents=True, exist_ok=True)
 
-        key_to_filename_map = self._ignore_source_key_prefixes(self.key_to_filename_map, ignored_source_key_prefixes)
+        key_to_filename_map = self._prepare_output_key_map(
+            self.key_to_filename_map,
+            ignored_source_key_prefixes,
+            source_key_replacements,
+        )
         all_expected_keys = set(key_to_filename_map.keys())
 
         if not key_to_filename_map:
@@ -926,6 +973,7 @@ def _save_generator_distributed(
         strict: bool = True,
         save_every_n_ranks: int = 1,
         ignored_source_key_prefixes: Iterable[str] | None = None,
+        source_key_replacements: Mapping[str, Iterable[str]] | None = None,
     ):
         is_distributed = torch.distributed.is_available() and torch.distributed.is_initialized()
         if is_distributed:
@@ -952,7 +1000,11 @@ def _save_generator_distributed(
         if is_distributed:
             torch.distributed.barrier()
 
-        key_to_filename_map = self._ignore_source_key_prefixes(self.key_to_filename_map, ignored_source_key_prefixes)
+        key_to_filename_map = self._prepare_output_key_map(
+            self.key_to_filename_map,
+            ignored_source_key_prefixes,
+            source_key_replacements,
+        )
 
         # Fallback: no sharding map, single-file save
         if not key_to_filename_map:
diff --git a/src/megatron/bridge/training/checkpointing.py b/src/megatron/bridge/training/checkpointing.py
index ac4508884f..72a594b17e 100644
--- a/src/megatron/bridge/training/checkpointing.py
+++ b/src/megatron/bridge/training/checkpointing.py
@@ -1152,12 +1152,14 @@ def save_checkpoint(
     # Collect rng state across data parallel ranks.
     if pg_collection is None:
         pg_collection = get_pg_collection(model)
-    rng_state = get_rng_state(
-        data_parallel_random_init=cfg.rng.data_parallel_random_init,
-        ckpt_format=ckpt_cfg.ckpt_format,
-        pg_collection=pg_collection,
-        module_name=module_name,
-    )
+    rng_state = None
+    if ckpt_cfg.save_rng:
+        rng_state = get_rng_state(
+            data_parallel_random_init=cfg.rng.data_parallel_random_init,
+            ckpt_format=ckpt_cfg.ckpt_format,
+            pg_collection=pg_collection,
+            module_name=module_name,
+        )
 
     # Collect rerun state across all ranks
     rerun_state_machine = get_rerun_state_machine()
diff --git a/src/megatron/bridge/training/model_load_save.py b/src/megatron/bridge/training/model_load_save.py
index a077b8a520..d9c638b935 100644
--- a/src/megatron/bridge/training/model_load_save.py
+++ b/src/megatron/bridge/training/model_load_save.py
@@ -622,17 +622,13 @@ def save_megatron_model(
         from megatron.bridge.training.checkpointing import (
             _build_sharded_state_dict_metadata,
             generate_state_dict,
-            get_rng_state,
         )
         from megatron.bridge.training.utils.pg_utils import get_pg_collection
 
         logger.info("[LOW_MEMORY_SAVE] Generating state dict...")
 
-        # Get RNG state (minimal, since save_rng=False)
+        # Conversion checkpoints intentionally omit RNG state.
         pg_collection = get_pg_collection(model)
-        rng_state = get_rng_state(
-            data_parallel_random_init=False, ckpt_format=ckpt_format, pg_collection=pg_collection
-        )
 
         # Build sharded state dict metadata
         sharded_sd_metadata = _build_sharded_state_dict_metadata(False, state.cfg.checkpoint)
@@ -643,7 +639,7 @@ def save_megatron_model(
             model,
             optimizer=None,
             opt_param_scheduler=None,
-            rng_state=rng_state,
+            rng_state=None,
             iteration=0,
             optim_sd_kwargs=dict(metadata=sharded_sd_metadata),
             model_sd_kwargs=dict(metadata=sharded_sd_metadata),
diff --git a/tests/unit_tests/models/gpt_oss/test_gpt_oss_bridges.py b/tests/unit_tests/models/gpt_oss/test_gpt_oss_bridges.py
index ff7c35a0af..f10bceb758 100644
--- a/tests/unit_tests/models/gpt_oss/test_gpt_oss_bridges.py
+++ b/tests/unit_tests/models/gpt_oss/test_gpt_oss_bridges.py
@@ -67,3 +67,27 @@ def test_provider_bridge_maps_config(self, mock_pretrained):
         # dtype mapping
         assert provider.bf16 is True
         assert provider.params_dtype == torch.bfloat16
+
+    def test_export_replaces_mxfp4_expert_keys(self, mock_pretrained):
+        source = Mock()
+        source.get_all_keys.return_value = [
+            "model.layers.0.mlp.experts.gate_up_proj_blocks",
+            "model.layers.0.mlp.experts.gate_up_proj_scales",
+            "model.layers.0.mlp.experts.down_proj_blocks",
+            "model.layers.0.mlp.experts.down_proj_scales",
+            "model.layers.0.mlp.experts.gate_up_proj_bias",
+        ]
+        mock_pretrained.state = Mock(source=source)
+
+        replacements = GPTOSSBridge().hf_export_source_key_replacements(mock_pretrained)
+
+        assert replacements == {
+            "model.layers.0.mlp.experts.gate_up_proj": (
+                "model.layers.0.mlp.experts.gate_up_proj_blocks",
+                "model.layers.0.mlp.experts.gate_up_proj_scales",
+            ),
+            "model.layers.0.mlp.experts.down_proj": (
+                "model.layers.0.mlp.experts.down_proj_blocks",
+                "model.layers.0.mlp.experts.down_proj_scales",
+            ),
+        }
diff --git a/tests/unit_tests/models/test_auto_bridge.py b/tests/unit_tests/models/test_auto_bridge.py
index f8a8ee225f..4fbe62b24d 100644
--- a/tests/unit_tests/models/test_auto_bridge.py
+++ b/tests/unit_tests/models/test_auto_bridge.py
@@ -348,7 +348,20 @@ def test_save_hf_weights_keeps_all_keys_when_mtp_enabled(self, tmp_path):
 
         assert source.save_generator_kwargs["ignored_source_key_prefixes"] is None
 
-    def _run_save_hf_weights(self, source, tmp_path, *, mtp_num_layers):
+    def test_save_hf_weights_passes_source_key_replacements(self, tmp_path):
+        source = _make_fake_source(present=set())
+        replacements = {"model.weight": ("model.weight_blocks", "model.weight_scales")}
+
+        self._run_save_hf_weights(
+            source,
+            tmp_path,
+            mtp_num_layers=1,
+            source_key_replacements=replacements,
+        )
+
+        assert source.save_generator_kwargs["source_key_replacements"] == replacements
+
+    def _run_save_hf_weights(self, source, tmp_path, *, mtp_num_layers, source_key_replacements=None):
         """Drive ``save_hf_weights`` with a stubbed bridge/model so the only
         behavior under test is the MTP prefix-resolution wiring.
 
@@ -366,6 +379,7 @@ def _run_save_hf_weights(self, source, tmp_path, *, mtp_num_layers):
 
         fake_model_bridge = Mock()
         fake_model_bridge.stream_weights_megatron_to_hf.return_value = iter([])
+        fake_model_bridge.hf_export_source_key_replacements.return_value = source_key_replacements or {}
 
         with (
             # ``state`` is a read-only property on PreTrainedBase, so patch it
@@ -1793,9 +1807,18 @@ def test_import_ckpt_basic(self, mock_from_hf_pretrained, mock_to_megatron_model
         mock_bridge.save_megatron_model = Mock()
 
         # Test import_ckpt
-        AutoBridge.import_ckpt("meta-llama/Meta-Llama-3-8B", "./megatron_checkpoint")
+        with (
+            patch("megatron.bridge.models.conversion.auto_bridge.dist.is_initialized", return_value=False),
+            patch(
+                "megatron.bridge.training.model_load_save.temporary_distributed_context"
+            ) as mock_distributed_context,
+        ):
+            AutoBridge.import_ckpt("meta-llama/Meta-Llama-3-8B", "./megatron_checkpoint")
 
         # Assertions
+        mock_distributed_context.assert_called_once_with(backend="gloo")
+        mock_distributed_context.return_value.__enter__.assert_called_once_with()
+        mock_distributed_context.return_value.__exit__.assert_called_once()
         mock_from_hf_pretrained.assert_called_once_with("meta-llama/Meta-Llama-3-8B")
         mock_bridge.to_megatron_model.assert_called_once_with(wrap_with_ddp=False, use_cpu_initialization=True)
         mock_bridge.save_megatron_model.assert_called_once_with(
@@ -1821,13 +1844,17 @@ def test_import_ckpt_with_kwargs(self, mock_from_hf_pretrained, mock_to_megatron
         mock_bridge._model_bridge.get_hf_tokenizer_kwargs.return_value = {}
 
         # Test import_ckpt with kwargs
-        AutoBridge.import_ckpt(
-            "./local_model",
-            "./megatron_checkpoint",
-            torch_dtype=torch.float16,
-            device_map="auto",
-            revision="0123456789abcdef",  # pragma: allowlist secret
-        )
+        with (
+            patch("megatron.bridge.models.conversion.auto_bridge.dist.is_initialized", return_value=False),
+            patch("megatron.bridge.training.model_load_save.temporary_distributed_context"),
+        ):
+            AutoBridge.import_ckpt(
+                "./local_model",
+                "./megatron_checkpoint",
+                torch_dtype=torch.float16,
+                device_map="auto",
+                revision="0123456789abcdef",  # pragma: allowlist secret
+            )
 
         # Assertions
         mock_from_hf_pretrained.assert_called_once_with(
@@ -1859,12 +1886,16 @@ def test_import_ckpt_with_low_memory_save(
         mock_bridge.save_megatron_model = Mock()
         mock_bridge._model_bridge.get_hf_tokenizer_kwargs.return_value = {}
 
-        AutoBridge.import_ckpt(
-            "meta-llama/Meta-Llama-3-8B",
-            "./megatron_checkpoint",
-            low_memory_save=True,
-            torch_dtype=torch.bfloat16,
-        )
+        with (
+            patch("megatron.bridge.models.conversion.auto_bridge.dist.is_initialized", return_value=False),
+            patch("megatron.bridge.training.model_load_save.temporary_distributed_context"),
+        ):
+            AutoBridge.import_ckpt(
+                "meta-llama/Meta-Llama-3-8B",
+                "./megatron_checkpoint",
+                low_memory_save=True,
+                torch_dtype=torch.bfloat16,
+            )
 
         mock_from_hf_pretrained.assert_called_once_with(
             "meta-llama/Meta-Llama-3-8B",
@@ -1878,6 +1909,31 @@ def test_import_ckpt_with_low_memory_save(
             low_memory_save=True,
         )
 
+    @patch.object(AutoBridge, "save_megatron_model")
+    @patch.object(AutoBridge, "to_megatron_model")
+    @patch.object(AutoBridge, "from_hf_pretrained")
+    def test_import_ckpt_reuses_existing_distributed_context(
+        self, mock_from_hf_pretrained, mock_to_megatron_model, mock_save_megatron_model
+    ):
+        """Test import_ckpt does not replace a caller-managed process group."""
+        mock_bridge = Mock(spec=AutoBridge)
+        mock_from_hf_pretrained.return_value = mock_bridge
+        mock_bridge.to_megatron_model.return_value = [Mock()]
+        mock_bridge.save_megatron_model = Mock()
+        mock_bridge._model_bridge.get_hf_tokenizer_kwargs.return_value = {}
+
+        with (
+            patch("megatron.bridge.models.conversion.auto_bridge.dist.is_initialized", return_value=True),
+            patch(
+                "megatron.bridge.training.model_load_save.temporary_distributed_context"
+            ) as mock_distributed_context,
+        ):
+            AutoBridge.import_ckpt("meta-llama/Meta-Llama-3-8B", "./megatron_checkpoint")
+
+        mock_distributed_context.assert_not_called()
+        mock_bridge.to_megatron_model.assert_called_once_with(wrap_with_ddp=False, use_cpu_initialization=True)
+        mock_bridge.save_megatron_model.assert_called_once()
+
     def test_export_ckpt_basic(self):
         """Test basic export_ckpt functionality."""
         # Setup mocks
diff --git a/tests/unit_tests/models/test_conversion_utils.py b/tests/unit_tests/models/test_conversion_utils.py
index e232eb067d..c755d3bca0 100644
--- a/tests/unit_tests/models/test_conversion_utils.py
+++ b/tests/unit_tests/models/test_conversion_utils.py
@@ -13,7 +13,9 @@
 # limitations under the License.
 
 import pytest
+import torch
 
+from megatron.bridge.models.conversion import utils as conversion_utils
 from megatron.bridge.models.conversion.utils import mcore_to_hf_window_size
 
 
@@ -33,3 +35,37 @@ def test_mcore_to_hf_window_size(window_size, expected):
 def test_mcore_to_hf_window_size_rejects_malformed_pair():
     with pytest.raises(ValueError, match="two-element MCore window"):
         mcore_to_hf_window_size([2047])
+
+
+def test_unwrap_model_supports_mcore_fsdp_factory(monkeypatch):
+    from megatron.core.distributed.fsdp import mcore_fsdp_adapter
+
+    class FakeFullyShardedDataParallelV1(torch.nn.Module):
+        def __init__(self, module):
+            super().__init__()
+            self.module = module
+
+    class FakeFullyShardedDataParallelV2(FakeFullyShardedDataParallelV1):
+        pass
+
+    def fsdp_factory(*args, **kwargs):
+        return FakeFullyShardedDataParallelV2(*args, **kwargs)
+
+    monkeypatch.setattr(mcore_fsdp_adapter, "FullyShardedDataParallel", fsdp_factory)
+    monkeypatch.setattr(
+        mcore_fsdp_adapter,
+        "FullyShardedDataParallelV1",
+        FakeFullyShardedDataParallelV1,
+        raising=False,
+    )
+    monkeypatch.setattr(
+        mcore_fsdp_adapter,
+        "FullyShardedDataParallelV2",
+        FakeFullyShardedDataParallelV2,
+        raising=False,
+    )
+
+    model = torch.nn.Linear(2, 2)
+    wrapped_model = FakeFullyShardedDataParallelV2(model)
+
+    assert conversion_utils.unwrap_model(wrapped_model) is model
diff --git a/tests/unit_tests/models/test_gpt_provider.py b/tests/unit_tests/models/test_gpt_provider.py
index 48ecff5b8a..b26416f370 100644
--- a/tests/unit_tests/models/test_gpt_provider.py
+++ b/tests/unit_tests/models/test_gpt_provider.py
@@ -374,6 +374,57 @@ def test_default_layer_spec_default_case(self, mock_te_full_spec, mock_te_spec):
         mock_te_spec.assert_called_once_with(provider)
         assert result == "te_spec"
 
+    @patch("megatron.bridge.models.gpt_provider.local_layer_spec")
+    @patch("megatron.bridge.models.gpt_provider.transformer_engine_layer_spec")
+    @patch("megatron.bridge.models.gpt_provider.transformer_engine_full_layer_spec")
+    @patch("megatron.bridge.models.gpt_provider.torch.cuda.is_available", return_value=False)
+    def test_default_layer_spec_uses_local_spec_for_cpu_only_initialization(
+        self, mock_cuda_available, mock_te_full_spec, mock_te_spec, mock_local_spec
+    ):
+        """CPU-only model construction must not instantiate Transformer Engine modules."""
+        from megatron.bridge.models.gpt_provider import default_layer_spec
+
+        provider = GPTModelProvider(
+            num_layers=12,
+            hidden_size=768,
+            num_attention_heads=12,
+            use_cpu_initialization=True,
+            use_transformer_engine_full_layer_spec=True,
+        )
+        mock_local_spec.return_value = "local_spec"
+
+        result = default_layer_spec(provider)
+
+        mock_cuda_available.assert_called_once_with()
+        mock_local_spec.assert_called_once_with(provider)
+        mock_te_full_spec.assert_not_called()
+        mock_te_spec.assert_not_called()
+        assert result == "local_spec"
+
+    @patch("megatron.bridge.models.gpt_provider.local_layer_spec")
+    @patch("megatron.bridge.models.gpt_provider.transformer_engine_layer_spec")
+    @patch("megatron.bridge.models.gpt_provider.torch.cuda.is_available", return_value=True)
+    def test_default_layer_spec_keeps_te_spec_for_cpu_initialization_with_cuda(
+        self, mock_cuda_available, mock_te_spec, mock_local_spec
+    ):
+        """CPU initialization on a CUDA host preserves the configured TE model layout."""
+        from megatron.bridge.models.gpt_provider import default_layer_spec
+
+        provider = GPTModelProvider(
+            num_layers=12,
+            hidden_size=768,
+            num_attention_heads=12,
+            use_cpu_initialization=True,
+        )
+        mock_te_spec.return_value = "te_spec"
+
+        result = default_layer_spec(provider)
+
+        mock_cuda_available.assert_called_once_with()
+        mock_local_spec.assert_not_called()
+        mock_te_spec.assert_called_once_with(provider)
+        assert result == "te_spec"
+
     def test_mtp_block_spec_returns_none_when_mtp_disabled(self):
         """mtp_block_spec returns None when mtp_num_layers is unset."""
         from megatron.bridge.models.gpt_provider import mtp_block_spec
@@ -456,6 +507,42 @@ def test_mtp_block_spec_re_derives_last_decoder_spec_when_layer_specs_empty(
         mock_get_mtp.assert_called_once_with(provider, moe_layer_spec, use_transformer_engine=True, vp_stage=2)
         assert result == "mtp_spec"
 
+    @patch("megatron.bridge.models.gpt_provider.torch.cuda.is_available", return_value=False)
+    @patch("megatron.core.models.gpt.gpt_layer_specs.get_gpt_decoder_layer_specs")
+    @patch("megatron.core.models.gpt.gpt_layer_specs.get_gpt_mtp_block_spec")
+    def test_mtp_block_spec_uses_local_spec_for_cpu_only_initialization(
+        self, mock_get_mtp, mock_get_decoder_specs, mock_cuda_available
+    ):
+        """CPU-only MTP construction must propagate the non-TE layer selection."""
+        from megatron.bridge.models.gpt_provider import mtp_block_spec
+
+        provider = GPTModelProvider(
+            num_layers=2,
+            hidden_size=128,
+            num_attention_heads=4,
+            mtp_num_layers=1,
+            use_cpu_initialization=True,
+        )
+        empty_block_spec = Mock()
+        empty_block_spec.layer_specs = []
+        provider.transformer_layer_spec = lambda config: empty_block_spec
+
+        local_moe_spec = Mock(name="local_moe_spec")
+        mock_get_decoder_specs.return_value = [local_moe_spec]
+        mock_get_mtp.return_value = "mtp_spec"
+
+        result = mtp_block_spec(provider, vp_stage=2)
+
+        mock_cuda_available.assert_called_once_with()
+        mock_get_decoder_specs.assert_called_once_with(
+            provider,
+            use_transformer_engine=False,
+            normalization=provider.normalization,
+            qk_l2_norm=provider.qk_l2_norm,
+        )
+        mock_get_mtp.assert_called_once_with(provider, local_moe_spec, use_transformer_engine=False, vp_stage=2)
+        assert result == "mtp_spec"
+
     @patch("megatron.core.models.gpt.gpt_layer_specs.get_gpt_mtp_block_spec")
     def test_mtp_block_spec_passes_vp_stage_to_callable_spec(self, mock_get_mtp):
         """When the transformer_layer_spec callable accepts vp_stage, it is forwarded."""
diff --git a/tests/unit_tests/models/test_hf_pretrained_state.py b/tests/unit_tests/models/test_hf_pretrained_state.py
index 7356f5f6e8..05bda24724 100644
--- a/tests/unit_tests/models/test_hf_pretrained_state.py
+++ b/tests/unit_tests/models/test_hf_pretrained_state.py
@@ -162,6 +162,69 @@ def gather_rank_zero(output: list[object | None], value: object) -> None:
     assert index_data["metadata"] == {"format": "pt", "total_size": expected_total_size}
 
 
+@pytest.mark.parametrize("distributed_save", [False, True])
+def test_save_generator_replaces_quantized_source_keys(tmp_path, monkeypatch, distributed_save: bool) -> None:
+    shard_filename = "model-00001-of-00001.safetensors"
+    blocks_key = "model.layers.0.mlp.experts.gate_up_proj_blocks"
+    scales_key = "model.layers.0.mlp.experts.gate_up_proj_scales"
+    output_key = "model.layers.0.mlp.experts.gate_up_proj"
+    bias_key = "model.layers.0.mlp.experts.gate_up_proj_bias"
+    _write_safetensors_index(
+        tmp_path,
+        {
+            blocks_key: shard_filename,
+            scales_key: shard_filename,
+            bias_key: shard_filename,
+        },
+    )
+    source = SafeTensorsStateSource(tmp_path)
+    output_path = tmp_path / "output"
+    if distributed_save:
+        _mock_single_rank_distributed(monkeypatch)
+
+    tensors = {
+        output_key: torch.ones((2, 2)),
+        bias_key: torch.zeros(2),
+    }
+    source.save_generator(
+        iter(tensors.items()),
+        output_path,
+        distributed_save=distributed_save,
+        source_key_replacements={output_key: (blocks_key, scales_key)},
+    )
+
+    with safe_open(output_path / shard_filename, framework="pt", device="cpu") as shard:
+        assert set(shard.keys()) == set(tensors)
+        for key, tensor in tensors.items():
+            torch.testing.assert_close(shard.get_tensor(key), tensor)
+
+    index_data = json.loads((output_path / "model.safetensors.index.json").read_text(encoding="utf-8"))
+    assert index_data["weight_map"] == {
+        output_key: shard_filename,
+        bias_key: shard_filename,
+    }
+
+
+def test_save_generator_rejects_source_key_replacement_across_shards(tmp_path) -> None:
+    blocks_key = "model.weight_blocks"
+    scales_key = "model.weight_scales"
+    _write_safetensors_index(
+        tmp_path,
+        {
+            blocks_key: "model-00001-of-00002.safetensors",
+            scales_key: "model-00002-of-00002.safetensors",
+        },
+    )
+    source = SafeTensorsStateSource(tmp_path)
+
+    with pytest.raises(ValueError, match="span multiple shards"):
+        source.save_generator(
+            iter([("model.weight", torch.ones(1))]),
+            tmp_path / "output",
+            source_key_replacements={"model.weight": (blocks_key, scales_key)},
+        )
+
+
 def test_save_generator_writes_shard_as_soon_as_its_remaining_keys_arrive(tmp_path, monkeypatch) -> None:
     class _SubsetForbiddenSet(set[str]):
         def issubset(self, _other: Iterable[object]) -> bool:
diff --git a/tests/unit_tests/training/test_checkpointing.py b/tests/unit_tests/training/test_checkpointing.py
index a47083692e..ac18f0eb98 100644
--- a/tests/unit_tests/training/test_checkpointing.py
+++ b/tests/unit_tests/training/test_checkpointing.py
@@ -600,6 +600,7 @@ def save_checkpoint_fixtures():
 class TestSaveCheckpoint:
     """Test checkpoint saving functionality."""
 
+    @pytest.mark.parametrize("save_rng", [True, False])
     @patch("megatron.bridge.training.checkpointing.wandb_utils")
     @patch("megatron.bridge.training.checkpointing.is_last_rank")
     @patch("builtins.open", new_callable=mock_open)
@@ -647,6 +648,7 @@ def test_save_checkpoint_global(
         mock_is_last_rank,
         mock_wandb,
         save_checkpoint_fixtures,
+        save_rng,
     ):
         """Test saving a global checkpoint."""
         # Setup mocks
@@ -678,6 +680,7 @@ def test_save_checkpoint_global(
         # Add wandb logger to state
         save_checkpoint_fixtures["mock_state"].wandb_logger = Mock()
         save_checkpoint_fixtures["mock_state"].cfg.checkpoint.most_recent_k = -1
+        save_checkpoint_fixtures["mock_state"].cfg.checkpoint.save_rng = save_rng
 
         # Call save_checkpoint
         save_checkpoint(
@@ -694,6 +697,10 @@ def test_save_checkpoint_global(
         mock_ft.on_checkpointing_start.assert_called_once()
         mock_gen_state.assert_called_once()
         mock_dist_ckpt.save.assert_called_once()
+        if save_rng:
+            mock_get_rng.assert_called_once()
+        else:
+            mock_get_rng.assert_not_called()
 
         # Verify that the tracker file was written with the correct iteration
         tracker_calls = [
diff --git a/tests/unit_tests/training/test_model_load_save.py b/tests/unit_tests/training/test_model_load_save.py
index ee93303cb5..0ddb79fe56 100644
--- a/tests/unit_tests/training/test_model_load_save.py
+++ b/tests/unit_tests/training/test_model_load_save.py
@@ -846,16 +846,53 @@ def test_load_megatron_model_applies_overrides(self, mock_load_model_config, moc
 class TestSaveMegatronModel:
     """Test save_megatron_model function.
 
-    Note: These tests use low_memory_save=False because the low_memory_save=True path
-    requires parallel state to be initialized (get_rng_state calls mpu.get_pipeline_model_parallel_rank()).
-    Testing the low_memory_save=True path would require either:
-    1. Full distributed initialization, or
-    2. Extensive mocking of checkpointing internals (get_rng_state, generate_state_dict, etc.)
-
-    The low_memory_save=False path tests the core save_checkpoint integration without
-    those dependencies, which is sufficient for unit testing the function's API and behavior.
+    Most tests use low_memory_save=False to exercise save_checkpoint integration
+    without mocking the incremental state-dict processing machinery.
     """
 
+    def test_low_memory_save_omits_rng_collection(self):
+        """Low-memory conversion saves must not initialize CUDA for disabled RNG state."""
+
+        class MockModelConfig(ModelProviderMixin, Mock):
+            def provide(self, pre_process=None, post_process=None, vp_stage=None):
+                return Mock()
+
+            def finalize(self) -> None:
+                pass
+
+        mock_model = Mock()
+        mock_model.named_parameters.return_value = []
+        mock_model.parameters.return_value = []
+        mock_pg_collection = Mock()
+
+        with (
+            tempfile.TemporaryDirectory() as temp_dir,
+            patch(
+                "megatron.bridge.training.model_load_save.get_model_config",
+                return_value=MockModelConfig(),
+            ),
+            patch(
+                "megatron.bridge.training.utils.pg_utils.get_pg_collection",
+                return_value=mock_pg_collection,
+            ),
+            patch(
+                "megatron.bridge.training.checkpointing.get_rng_state",
+            ) as mock_get_rng_state,
+            patch(
+                "megatron.bridge.training.checkpointing._build_sharded_state_dict_metadata",
+                return_value={},
+            ),
+            patch(
+                "megatron.bridge.training.checkpointing.generate_state_dict",
+                return_value={},
+            ) as mock_generate_state_dict,
+            patch("megatron.bridge.training.model_load_save.save_checkpoint"),
+        ):
+            save_megatron_model([mock_model], temp_dir, ckpt_format="torch_dist", low_memory_save=True)
+
+        mock_get_rng_state.assert_not_called()
+        assert mock_generate_state_dict.call_args.kwargs["rng_state"] is None
+
     @patch("megatron.bridge.training.model_load_save.save_checkpoint")
     @patch("megatron.bridge.training.model_load_save.get_model_config")
     @patch("megatron.bridge.training.model_load_save.GlobalState")
diff --git a/uv.lock b/uv.lock
index 0f1626ea0b..54026cfabe 100644
--- a/uv.lock
+++ b/uv.lock
@@ -1,5 +1,5 @@
 version = 1
-revision = 2
+revision = 3
 requires-python = "==3.12.*"
 
 [options]
@@ -2226,6 +2226,7 @@ requires-dist = [
 provides-extras = ["training", "mlm", "dev", "lts", "te", "ssm"]
 
 [package.metadata.requires-dev]
+batch-invariant = [{ name = "deep-gemm", git = "https://github.com/deepseek-ai/DeepGEMM.git?rev=714dd1a4a980f7937a74343d19a8eba4fe321480" }]
 build = [
     { name = "cython", specifier = ">=3.0.0" },
     { name = "hatchling" },
@@ -2250,13 +2251,14 @@ docs = [
     { name = "sphinx-copybutton" },
 ]
 linting = [
-    { name = "black", specifier = "==24.4.2" },
+    { name = "black", specifier = "==26.3.0" },
     { name = "flake8", specifier = "==7.1.0" },
     { name = "isort", specifier = "==5.13.2" },
     { name = "pylint", specifier = "==3.2.6" },
     { name = "ruff", specifier = "~=0.9.0" },
 ]
 no-pypi-wheels = [
+    { name = "deep-gemm", git = "https://github.com/deepseek-ai/DeepGEMM.git?rev=714dd1a4a980f7937a74343d19a8eba4fe321480" },
     { name = "emerging-optimizers", git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git?rev=v0.2.0" },
     { name = "flash-mla", git = "https://github.com/deepseek-ai/FlashMLA?rev=nv_dev" },
 ]
@@ -2267,7 +2269,7 @@ test = [
     { name = "nltk" },
     { name = "pydantic" },
     { name = "pygithub" },
-    { name = "pytest", specifier = "==8.3.5" },
+    { name = "pytest", specifier = "==9.1.1" },
     { name = "pytest-asyncio" },
     { name = "pytest-cov" },
     { name = "pytest-mock" },