[models] fix: support CPU-only checkpoint conversion by yaoyu-33 · Pull Request #3853 · NVIDIA-NeMo/Megatron-Bridge

yaoyu-33 · 2026-05-15T19:55:40Z

Summary

Fixes #3295.

Use Gloo and CPU initialization when CUDA is not visible for provider-based model construction.
Disable MCore CPU-offload hooks/flags and TE-only persist_layer_norm behavior during CPU-only initialization.
Fall back to the local GPT layer spec for CPU-init/no-CUDA builds so TransformerEngine SelfAttention is not instantiated.
Add focused unit coverage for no-visible-CUDA provider initialization and GPT layer-spec selection.

Validation

cw job 11784762 completed 0:0
- log: /lustre/fsw/portfolios/coreai/projects/coreai_dlalgo_llm/users/yuya/logs/issue3295/issue3295-cpu7_11784762.log
- uv sync --frozen --no-install-project --no-install-package megatron-core
- focused pytest: 6 passed, 30 deselected
- no-CUDA tiny Llama AutoBridge.to_megatron_model(wrap_with_ddp=False, use_cpu_initialization=True) completed with REPRO_DONE
cw job 11784769 completed 0:0
- log: /lustre/fsw/portfolios/coreai/projects/coreai_dlalgo_llm/users/yuya/logs/issue3295/issue3295-pc2_11784769.log
- uv run --no-sync pre-commit run --all-files

copy-pr-bot · 2026-05-15T19:55:44Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

yaoyu-33 · 2026-05-15T19:55:51Z

Progress update for issue #3295 CPU-only conversion:

Fixed provider initialization so no-visible-CUDA paths use CPU initialization and Gloo instead of CUDA/NCCL assumptions.
Disabled MCore CPU-offload context/flags and TE-only persist_layer_norm behavior for CPU-only initialization.
Added a no-CUDA GPT local layer-spec fallback so CPU conversion does not instantiate TransformerEngine SelfAttention.
Added focused regression coverage in test_model_provider_mixin.py and test_gpt_provider.py.

Validation on rebased HEAD 761a2f1ba96a959c2f233f104c6a4f09a963bc3a:

cw job 11784762 completed 0:0: focused pytest 6 passed, 30 deselected; no-CUDA tiny Llama AutoBridge.to_megatron_model(... use_cpu_initialization=True) reached REPRO_DONE.
Log: /lustre/fsw/portfolios/coreai/projects/coreai_dlalgo_llm/users/yuya/logs/issue3295/issue3295-cpu7_11784762.log
cw job 11784769 completed 0:0: uv run --no-sync pre-commit run --all-files passed.
Log: /lustre/fsw/portfolios/coreai/projects/coreai_dlalgo_llm/users/yuya/logs/issue3295/issue3295-pc2_11784769.log

Residual risk: this validates CPU-only model construction/conversion setup for the scoped wrap_with_ddp=False + use_cpu_initialization=True path with a tiny Llama repro. Broader CPU-only distributed conversion and non-GPT architectures may still need follow-up coverage.

yaoyu-33 · 2026-05-15T19:56:08Z

/ok to test 761a2f1

claude · 2026-05-15T19:58:23Z

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


nit: The existing mtp_block_spec tests (lines 461, 494, 497, 526 in test_gpt_provider.py) all still assert use_transformer_engine=True — they pass because the test providers default to use_cpu_initialization=False. Consider adding one test that sets use_cpu_initialization=True with CUDA mocked away, so the use_transformer_engine=False path through get_gpt_mtp_block_spec and get_gpt_decoder_layer_specs is covered.

Added this coverage in test_mtp_block_spec_uses_local_decoder_spec_for_cpu_only_initialization.

The test forces use_cpu_initialization=True with no visible CUDA, drives the empty layer_specs fallback, and asserts both get_gpt_decoder_layer_specs() and get_gpt_mtp_block_spec() receive use_transformer_engine=False.

claude · 2026-05-15T19:59:32Z

Review: [models] fix: support CPU-only checkpoint conversion

Clean PR - the backend selection, TE feature disabling, and layer-spec fallback logic all look correct. The context manager in _disable_te_cpu_offload_context_for_cpu_only_initialization properly restores the monkey-patched attribute in the finally block, and the guard chain (import, getattr, patch) is defensive in the right places.

One gap: the mtp_block_spec path was changed from hardcoded use_transformer_engine=True to a computed value, but no test exercises the False branch. The existing MTP tests all pass because their providers default to use_cpu_initialization=False. See inline comment.

Suggested test cases: No perf tests impacted.

yaoyu-33 · 2026-05-15T21:05:12Z

 ModelT = TypeVar("ModelT", bound=MegatronModule)


+def _cuda_is_available() -> bool:


can you help to move utils functions to utils file?

Moved the CPU-only helper functions out of model_provider.py into src/megatron/bridge/models/utils.py, with model_provider.py importing them from there. I also updated the affected mixin tests to patch the relocated helper module.

yaoyu-33 · 2026-05-15T21:29:34Z

Review follow-up pushed in 73080e44f841b486b5f3de4f3600b40269009a2d.

Changes:

Moved CPU-only provider helper functions from model_provider.py to src/megatron/bridge/models/utils.py.
Added focused MTP coverage for the no-CUDA + use_cpu_initialization=True path, including the empty-layer-spec decoder fallback and use_transformer_engine=False propagation.
Updated affected provider mixin tests for the relocated utility calls.

Validation:

cw 11785831 completed 0:0: focused pytest selected 11 provider tests (11 passed, 26 deselected) and the no-CUDA tiny Llama AutoBridge.to_megatron_model(... use_cpu_initialization=True) repro completed with REPRO_DONE.
Log: /lustre/fsw/portfolios/coreai/projects/coreai_dlalgo_llm/users/yuya/logs/issue3295/issue3295-review-tests-container_11785831.log
cw 11785878 completed 0:0: uv run --no-sync pre-commit run --all-files passed.
Log: /lustre/fsw/portfolios/coreai/projects/coreai_dlalgo_llm/users/yuya/logs/issue3295/issue3295-review-precommit-container_11785878.log

Residual risk: validation is focused on the CPU-only conversion/provider path and review-requested MTP fallback; full CI is still the broader safety net.

yaoyu-33 · 2026-05-15T21:29:58Z

/ok to test 73080e4

yaoyu-33 · 2026-05-15T22:35:15Z

@@ -374,6 +374,58 @@ def test_default_layer_spec_default_case(self, mock_te_full_spec, mock_te_spec):
        mock_te_spec.assert_called_once_with(provider)


add a cpu only functional conversion tests please. Can be L1.

Added L1 CPU-only functional conversion coverage in 2e33ee6250c0bfb3b7154aecf1859043a20f4b2a.

What changed:

Added tests/functional_tests/test_groups/converter/test_cpu_only_conversion.py, which creates a tiny local HF Llama model and runs AutoBridge.to_megatron_model(wrap_with_ddp=False, use_cpu_initialization=True) with CUDA hidden.

Added tests/functional_tests/launch_scripts/h100/active/L1_Launch_converter_cpu_only.sh so this is picked up as an H100 L1 functional launch script.

cw validation:

11787115 completed 0:0: new L1 launch script passed (1 passed) and completed the CPU-only tiny Llama conversion.
Log: /lustre/fsw/portfolios/coreai/projects/coreai_dlalgo_llm/users/yuya/logs/issue3295/issue3295-functional-cpu-container_11787115.log

11787138 completed 0:0: uv run --no-sync pre-commit run --all-files passed.
Log: /lustre/fsw/portfolios/coreai/projects/coreai_dlalgo_llm/users/yuya/logs/issue3295/issue3295-functional-precommit-container_11787138.log

Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>

yaoyu-33 · 2026-05-15T22:58:12Z

/ok to test 2e33ee6

yaoyu-33 · 2026-05-18T21:51:16Z

/ok to test f0788ae

cuichenx · 2026-05-18T21:52:53Z

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


can a model converted without TE work with TE training?

yaoyu-33 · 2026-05-18T21:58:10Z

Closing this PR without merge.

The scoped CPU-only path explored here gets model construction farther, but the compatibility tradeoff is too large to merge right now. If no GPU is available, we cannot use Transformer Engine, so this path has to fall back to the local spec. That means the imported model can have issues or differ from the model a user later trains on GPU with TE enabled.

We should defer this approach until there is a stronger request for CPU-only checkpoint conversion and a better compatibility design that preserves the intended GPU/TE training behavior.

yaoyu-33 mentioned this pull request May 15, 2026

[support] CPU-only Checkpoint Conversion #3295

Closed

copy-pr-bot Bot temporarily deployed to public May 15, 2026 19:56 Inactive

copy-pr-bot Bot temporarily deployed to test May 15, 2026 19:57 Inactive

claude Bot reviewed May 15, 2026

View reviewed changes

copy-pr-bot Bot temporarily deployed to public May 15, 2026 20:04 Inactive

copy-pr-bot Bot temporarily deployed to public May 15, 2026 20:18 Inactive

yaoyu-33 commented May 15, 2026

View reviewed changes

yaoyu-33 force-pushed the dispatch/issue3295-cpu-conversion branch from 761a2f1 to 73080e4 Compare May 15, 2026 21:29

copy-pr-bot Bot temporarily deployed to public May 15, 2026 21:30 Inactive

copy-pr-bot Bot temporarily deployed to test May 15, 2026 21:31 Inactive

copy-pr-bot Bot temporarily deployed to public May 15, 2026 21:38 Inactive

copy-pr-bot Bot temporarily deployed to public May 15, 2026 21:53 Inactive

yaoyu-33 added the needs-more-tests Requires additional L0 and L1 test coverage before merge label May 15, 2026

yaoyu-33 commented May 15, 2026

View reviewed changes

yaoyu-33 added 5 commits May 15, 2026 15:39

[models] fix: support CPU-only provider initialization

ebaded5

Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>

[models] fix: disable CPU offload flags without CUDA

c7c0bac

Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>

[models] fix: use local GPT spec for CPU conversion

8b74b59

Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>

[models] refactor: move CPU init helpers to utils

123e547

Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>

[test] add CPU-only conversion functional test

2e33ee6

Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>

yaoyu-33 force-pushed the dispatch/issue3295-cpu-conversion branch from 73080e4 to 2e33ee6 Compare May 15, 2026 22:57

copy-pr-bot Bot temporarily deployed to public May 15, 2026 22:58 Inactive

copy-pr-bot Bot temporarily deployed to test May 15, 2026 22:59 Inactive

yaoyu-33 added area:ckpt Checkpoint conversion, loading, export, and save paths bug Something isn't working needs-review PR is ready for code review and waiting on a reviewer labels May 15, 2026

copy-pr-bot Bot temporarily deployed to public May 15, 2026 23:09 Inactive

copy-pr-bot Bot temporarily deployed to public May 15, 2026 23:23 Inactive

Merge branch 'main' into dispatch/issue3295-cpu-conversion

f0788ae

copy-pr-bot Bot temporarily deployed to public May 18, 2026 21:51 Inactive

copy-pr-bot Bot temporarily deployed to test May 18, 2026 21:52 Inactive

cuichenx reviewed May 18, 2026

View reviewed changes

yaoyu-33 closed this May 18, 2026

copy-pr-bot Bot temporarily deployed to public May 18, 2026 21:59 Inactive

copy-pr-bot Bot temporarily deployed to public May 18, 2026 22:12 Inactive

This was referenced Aug 4, 2026

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

Closed

[model] fix: use local GPT spec for CPU-only initialization #5317

Closed

		ModelT = TypeVar("ModelT", bound=MegatronModule)


		def _cuda_is_available() -> bool:

		@@ -374,6 +374,58 @@ def test_default_layer_spec_default_case(self, mock_te_full_spec, mock_te_spec):
		mock_te_spec.assert_called_once_with(provider)

Uh oh!

Conversation

yaoyu-33 commented May 15, 2026

Summary

Validation

Uh oh!

copy-pr-bot Bot commented May 15, 2026

Uh oh!

yaoyu-33 commented May 15, 2026

Uh oh!

yaoyu-33 commented May 15, 2026

Uh oh!

claude Bot May 15, 2026

Choose a reason for hiding this comment

Uh oh!

yaoyu-33 May 15, 2026

Choose a reason for hiding this comment

Uh oh!

claude Bot commented May 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

yaoyu-33 May 15, 2026

Choose a reason for hiding this comment

Uh oh!

yaoyu-33 May 15, 2026

Choose a reason for hiding this comment

Uh oh!

yaoyu-33 commented May 15, 2026

Uh oh!

yaoyu-33 commented May 15, 2026

Uh oh!

yaoyu-33 May 15, 2026

Choose a reason for hiding this comment

Uh oh!

yaoyu-33 May 15, 2026

Choose a reason for hiding this comment

Uh oh!

yaoyu-33 commented May 15, 2026

Uh oh!

yaoyu-33 commented May 18, 2026

Uh oh!

cuichenx May 18, 2026

Choose a reason for hiding this comment

Uh oh!

yaoyu-33 commented May 18, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

claude Bot commented May 15, 2026 •

edited

Loading