From cb9318f01e4acbea4b0a616913328f9af1f4812f Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Tue, 4 Aug 2026 15:01:38 -0700 Subject: [PATCH 01/14] fix(conversion): initialize CPU imports with Gloo Signed-off-by: Chen Cui --- .../bridge/models/conversion/auto_bridge.py | 44 ++++++------ tests/unit_tests/models/test_auto_bridge.py | 70 +++++++++++++++---- 2 files changed, 80 insertions(+), 34 deletions(-) diff --git a/src/megatron/bridge/models/conversion/auto_bridge.py b/src/megatron/bridge/models/conversion/auto_bridge.py index 3685b34b6a..7fbb42e310 100644 --- a/src/megatron/bridge/models/conversion/auto_bridge.py +++ b/src/megatron/bridge/models/conversion/auto_bridge.py @@ -1450,27 +1450,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) - - # 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: + 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 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/tests/unit_tests/models/test_auto_bridge.py b/tests/unit_tests/models/test_auto_bridge.py index f8a8ee225f..a9ef81279f 100644 --- a/tests/unit_tests/models/test_auto_bridge.py +++ b/tests/unit_tests/models/test_auto_bridge.py @@ -1793,9 +1793,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 +1830,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 +1872,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 +1895,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 From b73e19e25ff82e15db31a1d24ed89b68d08c163c Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Tue, 4 Aug 2026 15:01:44 -0700 Subject: [PATCH 02/14] docs(skill): require public source for model verification Signed-off-by: Chen Cui --- .../create-model-verification-card/SKILL.md | 45 ++++++++++++++++++- 1 file changed, 44 insertions(+), 1 deletion(-) 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 From 781a6162b14edea3d74e5aef83b659523e9964ba Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Tue, 4 Aug 2026 15:28:52 -0700 Subject: [PATCH 03/14] [build] chore: pin MCore CPU YARN fix for validation Signed-off-by: Chen Cui --- 3rdparty/Megatron-LM | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/Megatron-LM b/3rdparty/Megatron-LM index 6513e3e23d..579c4624c0 160000 --- a/3rdparty/Megatron-LM +++ b/3rdparty/Megatron-LM @@ -1 +1 @@ -Subproject commit 6513e3e23d6b5eda6a1c934990b15e804237732b +Subproject commit 579c4624c07266ea6076fdc47b09ee04a9c83edb From ac12ef8b7d22c36edad7e1e6ee5e616a2000b345 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Tue, 4 Aug 2026 15:48:26 -0700 Subject: [PATCH 04/14] [model] fix: use local GPT spec for CPU-only initialization Avoid constructing CUDA-only Transformer Engine modules when GPT models are initialized on a host without visible CUDA. Keep the TE path unchanged when CUDA is available, and propagate the same backend choice to MTP specs. This revives the narrow spec-selection portion of #3853 now that full-model CPU import and GPU/TE checkpoint-load compatibility are explicit validation gates. Co-authored-by: Yu Yao <54727607+yaoyu-33@users.noreply.github.com> Signed-off-by: Chen Cui --- src/megatron/bridge/models/gpt_provider.py | 12 ++- tests/unit_tests/models/test_gpt_provider.py | 87 ++++++++++++++++++++ 2 files changed, 97 insertions(+), 2 deletions(-) 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/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.""" From 49700fa8a00d9a306a2dfd2637ee5e695511423f Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Tue, 4 Aug 2026 16:21:00 -0700 Subject: [PATCH 05/14] [build] chore: refresh MCore CPU initialization pin Signed-off-by: Chen Cui --- 3rdparty/Megatron-LM | 2 +- uv.lock | 8 +++++--- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/3rdparty/Megatron-LM b/3rdparty/Megatron-LM index 579c4624c0..14450c8af9 160000 --- a/3rdparty/Megatron-LM +++ b/3rdparty/Megatron-LM @@ -1 +1 @@ -Subproject commit 579c4624c07266ea6076fdc47b09ee04a9c83edb +Subproject commit 14450c8af9c8d4e5214b08f6bf00f2fa1b780521 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" }, From 5c8bc9c257c6af5648b14a0d3d032e7067ae1868 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Tue, 4 Aug 2026 16:42:32 -0700 Subject: [PATCH 06/14] [build] chore: refresh MCore CPU-only MoE pin Signed-off-by: Chen Cui --- 3rdparty/Megatron-LM | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/Megatron-LM b/3rdparty/Megatron-LM index 14450c8af9..593bc0aa8a 160000 --- a/3rdparty/Megatron-LM +++ b/3rdparty/Megatron-LM @@ -1 +1 @@ -Subproject commit 14450c8af9c8d4e5214b08f6bf00f2fa1b780521 +Subproject commit 593bc0aa8ab28cd4d4e2bb44426280f5f1e4d787 From e3205ac9893a8a60ee9b521240420345a8e21d1b Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Tue, 4 Aug 2026 16:57:51 -0700 Subject: [PATCH 07/14] fix(conversion): support MCore FSDP factory wrappers Signed-off-by: Chen Cui (cherry picked from commit 62097ac80ae1d4de370412951e5deee328e8039b) --- .../bridge/models/conversion/utils.py | 15 +++++--- .../models/test_conversion_utils.py | 36 +++++++++++++++++++ 2 files changed, 47 insertions(+), 4 deletions(-) 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/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 From b6957155ec9a322b0f4efd8ece68cca403b11e89 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Tue, 4 Aug 2026 17:24:37 -0700 Subject: [PATCH 08/14] fix(checkpointing): honor disabled RNG saves Signed-off-by: Chen Cui (cherry picked from commit 3c5b4bb48979128e8031b48bb595919466e0ca44) --- src/megatron/bridge/training/checkpointing.py | 14 ++--- .../bridge/training/model_load_save.py | 8 +-- .../unit_tests/training/test_checkpointing.py | 7 +++ .../training/test_model_load_save.py | 53 ++++++++++++++++--- 4 files changed, 62 insertions(+), 20 deletions(-) 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/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") From d77315e70bf5e6b3ed4f0165784c52ec200d16f3 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Tue, 4 Aug 2026 18:08:15 -0700 Subject: [PATCH 09/14] [build] chore: refresh MCore CPU checkpoint pin Signed-off-by: Chen Cui --- 3rdparty/Megatron-LM | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/Megatron-LM b/3rdparty/Megatron-LM index 593bc0aa8a..72b9e3c5df 160000 --- a/3rdparty/Megatron-LM +++ b/3rdparty/Megatron-LM @@ -1 +1 @@ -Subproject commit 593bc0aa8ab28cd4d4e2bb44426280f5f1e4d787 +Subproject commit 72b9e3c5df752739e6f399e291675c207c6d3e0d From 4df2bbcb31a4458ea080dd79806fa4e7b119d4b3 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Tue, 4 Aug 2026 18:46:34 -0700 Subject: [PATCH 10/14] [build] chore: refresh MCore CPU finalizer pin Signed-off-by: Chen Cui --- 3rdparty/Megatron-LM | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/Megatron-LM b/3rdparty/Megatron-LM index 72b9e3c5df..c081d1a5c1 160000 --- a/3rdparty/Megatron-LM +++ b/3rdparty/Megatron-LM @@ -1 +1 @@ -Subproject commit 72b9e3c5df752739e6f399e291675c207c6d3e0d +Subproject commit c081d1a5c1b7aad72614e701ad39b67c8762e50b From 027e284beff1e6b5d7ddb1b2ebebc462a9c2bc5b Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Tue, 4 Aug 2026 19:35:02 -0700 Subject: [PATCH 11/14] [build] chore: refresh MCore MoE checkpoint pin Signed-off-by: Chen Cui --- 3rdparty/Megatron-LM | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/Megatron-LM b/3rdparty/Megatron-LM index c081d1a5c1..4220c840f5 160000 --- a/3rdparty/Megatron-LM +++ b/3rdparty/Megatron-LM @@ -1 +1 @@ -Subproject commit c081d1a5c1b7aad72614e701ad39b67c8762e50b +Subproject commit 4220c840f520f621ca73415c8e73212998856b3c From dbe0a5b7a4c692f8dc8695fcb53601c7edbfc02d Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Tue, 4 Aug 2026 20:25:27 -0700 Subject: [PATCH 12/14] [build] chore: refresh MCore TE norm state pin Signed-off-by: Chen Cui --- 3rdparty/Megatron-LM | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/Megatron-LM b/3rdparty/Megatron-LM index 4220c840f5..9ba6247b74 160000 --- a/3rdparty/Megatron-LM +++ b/3rdparty/Megatron-LM @@ -1 +1 @@ -Subproject commit 4220c840f520f621ca73415c8e73212998856b3c +Subproject commit 9ba6247b7495c5a2cc1692ca8d14fd21947f305e From 99433894a7c6ec98a25079e123038b60cdff53c3 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Tue, 4 Aug 2026 21:17:32 -0700 Subject: [PATCH 13/14] chore: pin MCore TE attention checkpoint fix Signed-off-by: Chen Cui --- 3rdparty/Megatron-LM | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/Megatron-LM b/3rdparty/Megatron-LM index 9ba6247b74..d57d882fd0 160000 --- a/3rdparty/Megatron-LM +++ b/3rdparty/Megatron-LM @@ -1 +1 @@ -Subproject commit 9ba6247b7495c5a2cc1692ca8d14fd21947f305e +Subproject commit d57d882fd08ee1762eab6144260cf9693f58ac9d From d3346db7c4a8eedd43311d85d888b39687469f22 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Tue, 4 Aug 2026 22:10:39 -0700 Subject: [PATCH 14/14] fix(conversion): preserve dequantized GPT-OSS export keys Signed-off-by: Chen Cui --- .../bridge/models/conversion/auto_bridge.py | 2 + .../bridge/models/conversion/model_bridge.py | 16 +++++ .../bridge/models/gpt_oss/gpt_oss_bridge.py | 22 +++++++ .../bridge/models/hf_pretrained/state.py | 56 ++++++++++++++++- .../models/gpt_oss/test_gpt_oss_bridges.py | 24 +++++++ tests/unit_tests/models/test_auto_bridge.py | 16 ++++- .../models/test_hf_pretrained_state.py | 63 +++++++++++++++++++ 7 files changed, 196 insertions(+), 3 deletions(-) diff --git a/src/megatron/bridge/models/conversion/auto_bridge.py b/src/megatron/bridge/models/conversion/auto_bridge.py index 7fbb42e310..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 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/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/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/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 a9ef81279f..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 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: