[build] chore: pin CPU-only conversion fixes for validation

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

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

@@ -1215,13 +1215,15 @@ 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,
                 strict=strict,
                 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,

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,

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):

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.

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
 

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:

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()