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

megatron/core/dist_checkpointing/strategies/filesystem_async.py

@@ -234,17 +234,19 @@ def preload_tensors(write_buckets: List[WriteBucket], non_blocking=True) -> List
             non_blocking (bool, optional): knob to enable pinned D2H memcpy. Default is True.
         """
         result = []
+        synchronize_cuda = False
 
         for bucket in write_buckets:
             file_name, storage_key, (bytes_data, tensor_data) = bucket
             tensor_list = []
             for item, tensor in tensor_data:
                 # we belive these tensors are detached from the model trainers
+                synchronize_cuda |= tensor.is_cuda
                 tensor_list.append((item, tensor.to("cpu", non_blocking=non_blocking)))
                 # This is required for `PersistentAsyncCaller` to remove reference
                 del tensor
             result.append((file_name, storage_key, (bytes_data, tensor_list)))
-        if non_blocking:
+        if non_blocking and synchronize_cuda:
             torch.cuda.synchronize()
         return result
 

megatron/core/dist_checkpointing/strategies/state_dict_saver.py

@@ -1,6 +1,6 @@
 # Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.
 
-""" State dict saver for PyT Distributed format allowing asynchronous save. """
+"""State dict saver for PyT Distributed format allowing asynchronous save."""
 
 from logging import getLogger
 from time import time
@@ -13,6 +13,7 @@
 from torch.distributed.checkpoint.metadata import STATE_DICT_TYPE, Metadata
 from torch.distributed.checkpoint.planner import SavePlan, SavePlanner
 from torch.distributed.checkpoint.utils import _DistWrapper, _get_failure_dict
+from torch.distributed.distributed_c10d import _get_object_coll_device
 
 if TYPE_CHECKING:
     from .filesystem_async import FileSystemWriterAsync
@@ -249,7 +250,9 @@ def save_state_dict_async_finalize(
     # Broadcast failure status to all ranks to raise exceptions everywhere if needed.
     # The failure details are only raised on the coordinator.
     failures_occurred = torch.tensor(
-        [int(len(node_failures) > 0)], dtype=torch.int, device=torch.cuda.current_device()
+        [int(len(node_failures) > 0)],
+        dtype=torch.int,
+        device=_get_object_coll_device(dist_wrapper.group),
     )
     torch.distributed.broadcast(
         failures_occurred, src=dist_wrapper.coordinator_rank, group=dist_wrapper.group

megatron/core/extensions/transformer_engine.py

@@ -11,6 +11,7 @@
 import re
 import warnings
 from contextlib import contextmanager, nullcontext
+from types import MethodType
 from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Sequence, Set, Tuple, cast
 
 import torch
@@ -20,7 +21,11 @@
 from torch.nn.parameter import Parameter
 from typing_extensions import override
 
-from megatron.core.dist_checkpointing.mapping import ShardedObject, ShardedStateDict
+from megatron.core.dist_checkpointing.mapping import (
+    LocalNonpersistentObject,
+    ShardedObject,
+    ShardedStateDict,
+)
 from megatron.core.dist_checkpointing.utils import replace_prefix_for_sharding
 from megatron.core.enums import Fp4Recipe, Fp8Recipe
 from megatron.core.model_parallel_config import ModelParallelConfig
@@ -1015,6 +1020,45 @@ def forward(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor, torch.Tens
     TEFusedResidualRMSNorm = None  # type: ignore[assignment, misc]
 
 
+def _set_empty_te_extra_state_nonpersistent(
+    state_dict: dict, sharded_state_dict: ShardedStateDict, prefix: str
+) -> None:
+    """Keep empty TE quantization metadata local when loading non-TE checkpoints."""
+    if '_extra_state' not in state_dict:
+        return
+
+    extra_state = state_dict['_extra_state']
+    if extra_state is None or (isinstance(extra_state, torch.Tensor) and extra_state.numel() == 0):
+        sharded_state_dict[f'{prefix}_extra_state'] = LocalNonpersistentObject(extra_state)
+
+
+def _tenorm_sharded_state_dict(
+    module: torch.nn.Module,
+    prefix: str = '',
+    sharded_offsets: tuple = (),
+    metadata: Optional[dict] = None,
+    tp_group: Optional[torch.distributed.ProcessGroup] = None,
+) -> ShardedStateDict:
+    """Build a TE norm state dict without requiring empty quantization state on load."""
+    state_dict = module.state_dict(prefix='', keep_vars=True)
+    sharded_state_dict = make_sharded_tensors_for_checkpoint(
+        state_dict,
+        prefix,
+        sharded_offsets=sharded_offsets,
+        tp_group=get_tensor_model_parallel_group_if_none(tp_group),
+        dp_cp_group=(metadata or {}).get('dp_cp_group'),
+    )
+
+    _set_empty_te_extra_state_nonpersistent(state_dict, sharded_state_dict, prefix)
+    return sharded_state_dict
+
+
+def _bind_tenorm_sharded_state_dict(module: torch.nn.Module) -> None:
+    """Attach MCore distributed-checkpoint handling to a TE norm instance."""
+    module.sharded_state_dict = MethodType(_tenorm_sharded_state_dict, module)
+    module._mcore_sharded_state_dict_accepts_tp_group = True
+
+
 class TENorm:
     """A conditional wrapper to initialize an instance of
     Transformer-Engine's `LayerNorm` or `RMSNorm` based on input.
@@ -1072,6 +1116,8 @@ def __new__(
         )
 
         instance.returns_residual = use_fused_residual
+        _bind_tenorm_sharded_state_dict(instance)
+
         return cast(LayerNormInterface, instance)
 
 
@@ -2328,14 +2374,16 @@ def sharded_state_dict(
             state_dict = self.state_dict(prefix="", keep_vars=True)
         else:
             state_dict = {}
-        return make_sharded_tensors_for_checkpoint(
+        sharded_state_dict = make_sharded_tensors_for_checkpoint(
             state_dict,
             prefix,
             {'softmax_offset': 0},
             sharded_offsets,
             tp_group=self._tp_group,
             dp_cp_group=metadata["dp_cp_group"],
         )
+        _set_empty_te_extra_state_nonpersistent(state_dict, sharded_state_dict, prefix)
+        return sharded_state_dict
 
 
 if HAVE_TE and is_te_min_version("1.9.0.dev0"):

megatron/core/models/common/embeddings/yarn_rotary_pos_embedding.py

@@ -103,6 +103,14 @@ def __init__(
             # method causes a memory leak in NeMo-RL.
             self.forward.cache_clear()
 
+    def _apply(self, fn, recurse=True):
+        module = super()._apply(fn, recurse=recurse)
+        # ``forward`` is cached independently of registered buffers. Invalidate it after
+        # device or dtype migration so an identical call cannot return a tensor from the
+        # module's previous device.
+        self.forward.cache_clear()
+        return module
+
     def get_emb(self, max_seq_len: int, offset: int = 0) -> Tensor:
         """Forward pass of Yarn Rotary Embedding.
 
@@ -117,13 +125,16 @@ def get_emb(self, max_seq_len: int, offset: int = 0) -> Tensor:
             not self.rotary_interleaved
         ), "Yarn RoPE does not support interleaved rotary embeddings"
 
-        if self.inv_freq_extra.device.type == 'cpu':
-            # move `inv_freq_extra` to GPU once at the first micro-batch forward pass
-            self.inv_freq_extra = self.inv_freq_extra.to(device=torch.cuda.current_device())
-
-        if self.inv_freq_inter.device.type == 'cpu':
-            # move `inv_freq_inter` to GPU once at the first micro-batch forward pass
-            self.inv_freq_inter = self.inv_freq_inter.to(device=torch.cuda.current_device())
+        # The initial cache is built during construction, before ``cos_cached`` exists. Keep
+        # that work on the requested initialization device. On later cache rebuilds, follow
+        # the registered cache buffer so CPU-initialized models still migrate correctly.
+        target_device = (
+            self.cos_cached.device if hasattr(self, 'cos_cached') else self.inv_freq_extra.device
+        )
+        if self.inv_freq_extra.device != target_device:
+            self.inv_freq_extra = self.inv_freq_extra.to(device=target_device)
+        if self.inv_freq_inter.device != target_device:
+            self.inv_freq_inter = self.inv_freq_inter.to(device=target_device)
 
         low, high = _yarn_find_correction_range(
             self.beta_fast,

megatron/core/models/gpt/gpt_layer_specs.py

@@ -435,6 +435,10 @@ def get_gpt_layer_local_submodules(
             mlp_bda=get_bias_dropout_add,
         )
     else:
+        sharded_state_dict_keys_map = {"input_layernorm.": "self_attention.linear_qkv.layer_norm_"}
+        # TE MoE layers keep the pre-MLP norm standalone; dense TE layers fuse it into FC1.
+        if num_experts is None:
+            sharded_state_dict_keys_map["pre_mlp_layernorm."] = "mlp.linear_fc1.layer_norm_"
         return TransformerLayerSubmodules(
             input_layernorm=layer_norm,
             self_attention=ModuleSpec(
@@ -456,10 +460,7 @@ def get_gpt_layer_local_submodules(
             pre_mlp_layernorm=layer_norm,
             mlp=mlp,
             mlp_bda=get_bias_dropout_add,
-            sharded_state_dict_keys_map={
-                "input_layernorm.": "self_attention.linear_qkv.layer_norm_",
-                "pre_mlp_layernorm.": "mlp.linear_fc1.layer_norm_",
-            },
+            sharded_state_dict_keys_map=sharded_state_dict_keys_map,
         )
 
 

megatron/core/transformer/dot_product_attention.py

@@ -118,18 +118,28 @@ def __init__(
         if self.config.softmax_type == "vanilla":
             self.softmax_offset = None
         elif self.config.softmax_type == "off-by-one":
-            self.softmax_offset = torch.zeros(
-                self.num_attention_heads_per_partition,
-                device=torch.cuda.current_device(),
-                dtype=self.config.params_dtype,
+            self.register_buffer(
+                "softmax_offset",
+                torch.zeros(
+                    self.num_attention_heads_per_partition,
+                    device=(
+                        "cpu" if self.config.use_cpu_initialization else torch.cuda.current_device()
+                    ),
+                    dtype=self.config.params_dtype,
+                ),
+                persistent=False,
             )
         elif self.config.softmax_type == "learnable":
             self.register_parameter(
                 "softmax_offset",
                 torch.nn.Parameter(
                     torch.empty(
                         self.num_attention_heads_per_partition,
-                        device=torch.cuda.current_device(),
+                        device=(
+                            "cpu"
+                            if self.config.use_cpu_initialization
+                            else torch.cuda.current_device()
+                        ),
                         dtype=self.config.params_dtype,
                     )
                 ),

megatron/core/transformer/moe/token_dispatcher.py

@@ -427,7 +427,14 @@ def __init__(
         # [tp_size]. Represents the number of tokens received by the current rank from
         # other TP ranks.
         self.output_splits_tp = None
-        self.permute_idx_device = torch.device("cuda") if self.config.moe_permute_fusion else "cpu"
+        cpu_only_initialization = (
+            self.config.use_cpu_initialization and not torch.cuda.is_available()
+        )
+        self.permute_idx_device = (
+            torch.device("cuda")
+            if self.config.moe_permute_fusion and not cpu_only_initialization
+            else "cpu"
+        )
         input_chunk_idxs = torch.arange(
             self.num_experts * self.tp_size, device=self.permute_idx_device
         )
@@ -467,7 +474,7 @@ def __init__(
             or not self.config.cuda_graph_modules
         ):
             self.cuda_dtoh_point = "before_ep_alltoall"
-        if MoEAlltoAllTokenDispatcher.cuda_dtoh_stream is None:
+        if not cpu_only_initialization and MoEAlltoAllTokenDispatcher.cuda_dtoh_stream is None:
             MoEAlltoAllTokenDispatcher.cuda_dtoh_stream = torch.cuda.Stream()
 
         # Attributes that need to be captured in cudagraph. These attributes are returned

megatron/core/transformer/utils.py

@@ -281,9 +281,14 @@ def sharded_state_dict_default(
     metadata = ensure_metadata_has_dp_cp_group(metadata)
 
     if hasattr(module, 'sharded_state_dict'):
-        module_sharded_sd = module.sharded_state_dict(
-            prefix=prefix, sharded_offsets=sharded_offsets, metadata=metadata
-        )
+        sharded_state_dict_kwargs: Dict[str, Any] = {
+            'prefix': prefix,
+            'sharded_offsets': sharded_offsets,
+            'metadata': metadata,
+        }
+        if getattr(module, '_mcore_sharded_state_dict_accepts_tp_group', False):
+            sharded_state_dict_kwargs['tp_group'] = tp_group
+        module_sharded_sd = module.sharded_state_dict(**sharded_state_dict_kwargs)
     else:
         module_sd = module.state_dict(prefix='', keep_vars=True)
         module_sharded_sd = make_sharded_tensors_for_checkpoint(

tests/unit_tests/dist_checkpointing/test_async_save.py

@@ -1,5 +1,6 @@
 # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
 import sys
+from pathlib import Path
 from unittest import mock
 
 import pytest
@@ -11,6 +12,9 @@
 from megatron.core.dist_checkpointing.strategies.async_utils import AsyncCallsQueue
 from megatron.core.dist_checkpointing.strategies.filesystem_async import FileSystemWriterAsync
 from megatron.core.dist_checkpointing.strategies.nvrx import has_nvrx_async_support
+from megatron.core.dist_checkpointing.strategies.state_dict_saver import (
+    save_state_dict_async_finalize,
+)
 from megatron.core.dist_checkpointing.strategies.torch import (
     TorchDistSaveShardedStrategy,
     get_async_strategy,
@@ -163,3 +167,63 @@ def test_version_check_fails(self):
         ):
             with pytest.raises(AssertionError, match="Minimum required nvidia-resiliency-ext"):
                 has_nvrx_async_support()
+
+
+class TestFileSystemWriterAsync:
+    @staticmethod
+    def _write_buckets(tensor):
+        return [(Path("checkpoint"), "storage-key", ([], [(mock.sentinel.item, tensor)]))]
+
+    def test_preload_cpu_tensors_does_not_synchronize_cuda(self):
+        tensor = torch.ones(2)
+
+        with mock.patch.object(torch.cuda, "synchronize") as synchronize:
+            result = FileSystemWriterAsync.preload_tensors(self._write_buckets(tensor))
+
+        synchronize.assert_not_called()
+        assert result[0][2][1][0][1] is tensor
+
+    def test_preload_cuda_tensors_synchronizes_cuda(self):
+        tensor = mock.MagicMock()
+        tensor.is_cuda = True
+        tensor.to.return_value = torch.ones(2)
+
+        with mock.patch.object(torch.cuda, "synchronize") as synchronize:
+            FileSystemWriterAsync.preload_tensors(self._write_buckets(tensor))
+
+        tensor.to.assert_called_once_with("cpu", non_blocking=True)
+        synchronize.assert_called_once_with()
+
+
+class TestSaveStateDictAsyncFinalize:
+    @pytest.mark.parametrize("collective_device", ["cpu", "cuda"])
+    def test_failure_status_uses_process_group_device(self, collective_device):
+        storage_writer = mock.Mock()
+        storage_writer.retrieve_write_results.return_value = [mock.sentinel.local_result]
+        all_results = [mock.sentinel.all_results]
+        dist_wrapper = mock.Mock(is_coordinator=True, coordinator_rank=0, group=mock.sentinel.group)
+        dist_wrapper.gather_object.return_value = all_results
+        failures_occurred = mock.MagicMock()
+        failures_occurred.__bool__.return_value = False
+
+        with (
+            mock.patch(
+                "megatron.core.dist_checkpointing.strategies.state_dict_saver._get_failure_dict",
+                return_value={},
+            ),
+            mock.patch(
+                "megatron.core.dist_checkpointing.strategies.state_dict_saver._get_object_coll_device",
+                return_value=collective_device,
+            ) as get_collective_device,
+            mock.patch.object(torch, "tensor", return_value=failures_occurred) as tensor,
+            mock.patch.object(torch.distributed, "get_rank", return_value=0),
+            mock.patch.object(torch.distributed, "broadcast") as broadcast,
+        ):
+            save_state_dict_async_finalize(storage_writer, mock.sentinel.metadata, dist_wrapper)
+
+        storage_writer.finish.assert_called_once_with(mock.sentinel.metadata, all_results)
+        get_collective_device.assert_called_once_with(dist_wrapper.group)
+        tensor.assert_called_once_with([0], dtype=torch.int, device=collective_device)
+        broadcast.assert_called_once_with(
+            failures_occurred, src=dist_wrapper.coordinator_rank, group=dist_wrapper.group
+        )