fix: avoid MCore CUDA assumptions in CPU-only model conversion

[cuichenx]

What changed

This PR removes CUDA construction/save assumptions from the MCore portions of a CPU-only model-conversion workflow while preserving the required normal Transformer Engine (TE) reload on GPU:

• construct YARN caches on the requested initialization device, make frequency rebuilds follow a registered cache buffer after module migration, and invalidate memoized forwards so an identical post-migration call cannot return a tensor from the old device
• create non-vanilla attention softmax offsets on CPU during explicit CPU initialization; keep the non-learnable offset as a nonpersistent buffer so it follows module migration
• avoid MoE CUDA metadata/stream construction only during explicit CPU initialization on a host where CUDA is unavailable; preserve the normal CUDA resources when GPUs are visible
• avoid unconditional CUDA synchronization for all-CPU DCP preload and use the active process group's collective device for DCP finalization
• preserve the local-to-TE checkpoint schema needed by the tested dense Qwen2.5 and GPT-OSS paths: standalone MoE pre-MLP norm keys and optional empty TE norm/attention _extra_state
• preserve TENorm tensor-parallel replica identity when using its dynamically bound sharded-state method
• add focused CPU, CUDA migration/runtime, checkpoint-key, and TE extra-state regression tests, then validate with real Gloo and TP=2 TE DCP runs

The public branch history was preserved. MCore main at f30005143640bf5a00ea0a5ddb4f6ed2b0288158 (85 commits beyond the PR's previous merged base) was merged normally in signed commit cc0ddc7d7947a1d0e122af22a06165ace91fc605. The only textual conflict was in the TENorm constructor; the semantic resolution preserves both upstream's returns_residual assignment and this PR's checkpoint hook. A focused regression guards both adjacent behaviors. main has since added one unrelated activation-offload commit through 28797d06ae3ddc2ab5110b9a60b159552edc0c15; the PR remains mergeable. The exact candidate head is signed commit 68855c052a3f1b035c462268d255a918c889dd3c.

Scope and compatibility contract

The device and DCP fixes are generic hardening. The local-spec/TE adaptations are retained because a CPU-produced checkpoint must subsequently load into the normal GPU/TE model, but their evidence is intentionally model-specific:

• dense Qwen2.5-0.5B covers the dense local-to-TE route
• GPT-OSS 20B covers YARN, learnable non-vanilla softmax offsets, MoE dispatcher construction, the standalone MoE norm mapping, and optional TE norm/attention state
• an actual TE norm at TP=2 covers replicated checkpoint identity and integrity-validated DCP save/load

This does not claim arbitrary local/TE schema compatibility across model families, strict-loading modes, or MCore/TE versions. Focused fake-module tests classify empty and nonempty TE state; actual TE evidence covers the empty-state path. Nonempty state remains checkpointed by construction but has no real-FP8 model oracle here.

YARN's existing lru_cache on forward is now invalidated through nn.Module._apply after a successful device or dtype migration. A regression test performs CPU forward -> .cuda() -> identical direct forward and verifies that the result and both frequency tensors are on CUDA with numerical parity.

Bridge conversion requests a synchronous save, but MCore implements it by constructing its writer request and calling AsyncRequest.execute_sync(). The conversion saves therefore execute the changed FileSystemWriterAsync.preload_tensors and save_state_dict_async_finalize functions. This PR does not make the separate AsyncCallsQueue scheduler driverless-safe; that broader path still contains existing unconditional CUDA operations.

The finalizer relies on PyTorch's private _get_object_coll_device helper, verified in the tested PyTorch 2.13.0a0+8145d630e8.nv26.06 runtime. Its private-API status is an explicit maintenance risk.

Current-head validation

Unless explicitly labeled historical support, every result below uses exact candidate MCore head 68855c052..., merged main parent f30005143..., and the exact validation image noted below. Model-level runs additionally use the stated Bridge companion.

Focused MCore and DCP tests

• genuine driverless DFW cpu_short job 16203848, 8 CPUs/32 GiB, no GPU allocation or exclusivity: 13 passed, 4 intentional CUDA-only skips, covering YARN, non-vanilla attention, MoE construction, both DCP changes, checkpoint mappings, TE wrapper classification, and the TENorm merge regression; TE imported while CUDA stayed unavailable, device count zero, and CUDA uninitialized
• one-H100 job 16203852: 9 passed, including CPU forward -> .cuda() -> identical YARN forward, CUDA-visible MoE resource preservation, checkpoint-key/schema tests, and the TENorm merge regression
• two-H100 job 16203853 with an actual TE norm at TP=2: passed on both ranks, including TP replica-id checks and integrity-validated DCP save/load
• genuine driverless two-rank Gloo job 16203849: synchronous writer/finalizer save/load completed with access-integrity validation, CPU collective device, CUDA unavailable, device count zero, and CUDA uninitialized

Dense Qwen2.5-0.5B

• genuine driverless cpu_short job 16203850 completed 0:0; HF -> MCore -> HF constructed 494,032,768 parameters and matched all 291 tensors exactly
• original/exported CPU logits were finite, selected the same next token (12095), had cosine 0.999930322, and maximum absolute difference 0
• CUDA stayed unavailable/uninitialized and mapped driver-library checks stayed empty across import, save, reload, export, and comparison
• one-H100 job 16204194 completed 0:0; the CPU-produced checkpoint loaded into a normal model containing 217 TE-backed modules
• HF and TE/MCore selected the same next token (12095); logit cosine was 0.999936
• the fresh H100 re-export again matched all 291 HF tensors exactly, and its exported-HF logits matched the original exactly

GPT-OSS 20B

• genuine driverless cpu_short job 16203851 completed 0:0; import constructed all 20,914,757,184 parameters with the local spec and saved the complete torch-dist checkpoint (41,839,125,360-byte shard)
• CUDA stayed unavailable, device count stayed zero, and CUDA stayed uninitialized before and after
• one-H100 job 16204251 completed 0:0; the CPU-produced checkpoint loaded into the complete normal TE model and exported 411 BF16 HF tensors
• all 411 exported tensors matched the cached BF16 reference exactly and were finite; all 17 normalized structural/RoPE config fields matched
• one-H100 job 16204252 completed 0:0 after loading 241 TE-backed modules; original HF and loaded TE/MCore selected the same next token (12650), with maximum absolute logit difference 0.1875, mean absolute difference 0.027503010, and cosine 0.999975801

The validation companion predates a separate Bridge-side output-config correction: its generated GPT-OSS config writes attention_bias=false even though the exported files contain all 96 reference attention-bias tensors. The MCore/TE checkpoint load and exact tensor comparator remain valid oracles, but that companion-generated HF directory is not claimed directly reloadable. Current Bridge main at 2e77041c19... now writes attention_bias=true; this historical companion limitation must not be attributed to current Bridge.

Formatting and CI

• exact candidate, git diff --check: passed; every commit carries a DCO sign-off
• fresh public CI on 68855c052...: preflight, lint, copyright, DCO, wheel build/tests, pip/UV installation, and docs passed; the H100 functional workflow is waiting in the maintainer-controlled queue (the current /nvskills-ci requester requires maintain/admin permission)
• the previous c50b1c8bd... H100 transformer and MoE buckets passed, including this PR's new async-DCP tests; its broad distributed-checkpoint bucket failed/cancelled after three existing tests/unit_tests/dist_checkpointing/test_local.py cases, with no fix: avoid MCore CUDA assumptions in CPU-only model conversion #6270 test failure reported
• codeowner and multi-approval failures are review-governance gates, not code failures

Required external prerequisite and integration boundary

Stock NeMo 26.08.rc9 uses Transformer Engine 2.17.1+4329ff84 built with NCCL-EP. That build hard-links libcuda.so.1, so plain TE import fails on a genuine driverless CPU node before MCore code runs. This PR does not fix that first blocker; Transformer Engine #3381 remains open.

The driverless evidence therefore uses the same rc9 TE source rebuilt with NVTE_WITH_NCCL_EP=0:

nvcr.io/nvidian/nemo:26.08.rc9-te-no-nccl-ep-omar-20260813
sha256:99356ac6fa9260dab3fbd5dba24021a4786ad8baecada98121a53f0dfe933dc1

That workaround makes TE import driverless, but TE modules still cannot be constructed on CPU. Bridge must select a model-appropriate local PyTorch spec for the CPU stage.

Model validation uses cumulative Bridge integration head d3346db7c4a8eedd43311d85d888b39687469f22. This is a validation companion, not a claim about universal Bridge behavior: Bridge main contains the merged Gloo import scope from NVIDIA-NeMo/Megatron-Bridge#5007 and the merged save_rng=False fix from #5323; automatic local GPT/MTP selection #5317 was closed because it is not proven across families. A caller must still select a model-appropriate local spec for the CPU stage.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[cuichenx]

Reproduction context: full openai/gpt-oss-20b CPU import using NVIDIA-NeMo/Megatron-Bridge#5313 reached YarnRotaryEmbedding construction and failed because CPU initialization called torch.cuda.current_device() while creating the initial cache.

The focused regression test file in this PR has now passed with the repository-standard eight ranks in the nvcr.io/nvidian/nemo:26.08.rc4 container from exact pushed commit 579c4624c07266ea6076fdc47b09ee04a9c83edb (both tests passed on all eight ranks). End-to-end GPT-OSS CPU import remains pending and is not claimed as verified.

[cuichenx]

Integration update: a full openai/gpt-oss-20b CPU import using the exact public head of this PR progressed past YARN construction, confirming the original torch.cuda.current_device() failure is cleared. The run then stopped at a separate Megatron Bridge layer-spec issue: Bridge selected Transformer Engine attention on the no-CUDA host.

End-to-end import is still not verified, but the full workload no longer fails at the YARN code changed by this PR.

[cuichenx]

The exact-public-source GPT-OSS CPU import progressed past the Gloo, YARN, and local-spec blockers, then exposed a second generic MCore CPU-initialization defect: DotProductAttention called torch.cuda.current_device() while allocating off-by-one or learnable softmax offsets despite use_cpu_initialization=True.

Commit 14450c8af9c8d4e5214b08f6bf00f2fa1b780521 is now pushed on this PR. It allocates those tensors on CPU only for explicit CPU initialization and adds a regression test that makes any CUDA-device access fail. Both parameterized cases pass in a driverless container, and repository Black/isort/pylint/Ruff checks pass.

This is not an end-to-end verification claim. The repository-standard eight-rank focused run and a fresh full GPT-OSS CPU import from an exact clean public integration commit are next.

[cuichenx]

Focused validation of exact public head 14450c8af9c8d4e5214b08f6bf00f2fa1b780521 is complete. From a fresh clean clone in nvcr.io/nvidian/nemo:26.08.rc4, all eight ranks passed all four targeted YARN and DotProductAttention CPU-initialization tests. The checkout remained clean and publicly reachable after the run.

This validates the focused fixes only. Full GPT-OSS CPU import and GPU/TE loading of its produced checkpoint remain pending and are not claimed as verified.

[cuichenx]

Exact public head 9ba6247b7495c5a2cc1692ca8d14fd21947f305e passed the focused checkpoint-compatibility cluster validation.

• DFW job: 15110325
• Container: nvcr.io/nvidian/nemo:26.08.rc4
• Node: pool0-01804, 8 H100s / 8 ranks
• Result: COMPLETED 0:0 after 4m43s
• Tests: all eight ranks reported all four focused dense/MoE checkpoint-key and empty/non-empty TE norm extra-state cases passed
• Provenance: the job used a fresh clone of NVIDIA's public refs/pull/6270/head at the exact commit above; no source overlay or uncommitted patch was used

This is focused exact-public PR-head validation. End-to-end GPT-OSS conversion remains not verified. Next I will advance the Bridge integration gitlink to this exact public commit, push that change normally, and rerun the full CPU-import and GPU/TE-load workflow from fresh clean public source.

[cuichenx]

The fresh exact-public full GPU/TE gate exposed an eighth upstream checkpoint-compatibility defect.

• Bridge: clean public dbe0a5b7a4c692f8dc8695fcb53601c7edbfc02d
• MCore: clean public 9ba6247b7495c5a2cc1692ca8d14fd21947f305e
• CPU checkpoint: newly produced by zero-GPU job 15110835 from those exact heads
• GPU gate: DFW job 15111919, one H100 on pool0-01101, nvcr.io/nvidian/nemo:26.08.rc4, PyTorch 2.13.0a0+8145d630e8.nv26.06, TE 2.18.0+e7c550c5

The real 20,914,757,184-parameter TE model cleared the standalone MoE norm weight mismatch and the empty standalone TE norm extra-state mismatch. Checkpoint planning then failed on another optional empty TE entry:

RuntimeError: Missing key in checkpoint state_dict: decoder.layers.self_attention.core_attention._extra_state/shard_0_24.

Job 15111919 failed 1:0 after 4m14s. No HF output was created, and both source trees remained clean at the exact public commits. No local patch, overlay, or monkeypatch was used.

This is a failed exact-public candidate validation, not verification. I will fix the owning generic MCore TE-module checkpoint handling with focused empty/non-empty-state coverage, push it normally to this PR, then repeat the public-head validation, Bridge pin, fresh full CPU import, and H100/TE gate.

[cuichenx]

/ok to test d57d882

[cuichenx]

Exact-public focused H100 validation passed for current PR head d57d882fd08ee1762eab6144260cf9693f58ac9d.

• Fresh public-ref source: detached Git worktree at the exact SHA; fetched NVIDIA refs/pull/6270/head resolved to the same SHA
• Post-run source status: clean
• Container: nvcr.io/nvidian/nemo:26.08.rc4
• Hardware: eight H100 ranks on DFW node pool0-01637
• Command scope: tests/unit_tests/transformer/test_spec_customization.py -k 'GptLayerCheckpointKeys or TECheckpointCompatibility'
• Result: all eight ranks reported 5 passed, 5 deselected; Slurm job 15112718 completed 0:0 in 4m17s

This validates the focused public candidate only. It is not an end-to-end model-verification claim; the Bridge integration pin and full clean CPU-import plus H100/TE-load workflow still need to be rerun from the updated public heads.

[cuichenx]

/ok to test 350e619

[cuichenx]

/ok to test 8b36e73

[cuichenx]

/ok to test c50b1c8

[cuichenx]

/ok to test 68855c0

[cuichenx]

/nvskills-ci