diff --git a/megatron/core/dist_checkpointing/strategies/filesystem_async.py b/megatron/core/dist_checkpointing/strategies/filesystem_async.py index 1b144822839..6a4bdc30e98 100644 --- a/megatron/core/dist_checkpointing/strategies/filesystem_async.py +++ b/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 diff --git a/megatron/core/dist_checkpointing/strategies/state_dict_saver.py b/megatron/core/dist_checkpointing/strategies/state_dict_saver.py index 31ed0bb6b17..0dbd074c9fe 100644 --- a/megatron/core/dist_checkpointing/strategies/state_dict_saver.py +++ b/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 diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index c63f0c75dc2..6655f148f47 100644 --- a/megatron/core/extensions/transformer_engine.py +++ b/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,7 +2374,7 @@ 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}, @@ -2336,6 +2382,8 @@ def sharded_state_dict( 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"): diff --git a/megatron/core/models/common/embeddings/yarn_rotary_pos_embedding.py b/megatron/core/models/common/embeddings/yarn_rotary_pos_embedding.py index cb8a03d0b2b..36c6b32d6cf 100644 --- a/megatron/core/models/common/embeddings/yarn_rotary_pos_embedding.py +++ b/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, diff --git a/megatron/core/models/gpt/gpt_layer_specs.py b/megatron/core/models/gpt/gpt_layer_specs.py index a61c36bff5b..9ccf84399e0 100755 --- a/megatron/core/models/gpt/gpt_layer_specs.py +++ b/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, ) diff --git a/megatron/core/transformer/dot_product_attention.py b/megatron/core/transformer/dot_product_attention.py index 69039e0bfd0..49d7393464d 100644 --- a/megatron/core/transformer/dot_product_attention.py +++ b/megatron/core/transformer/dot_product_attention.py @@ -118,10 +118,16 @@ 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( @@ -129,7 +135,11 @@ def __init__( 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, ) ), diff --git a/megatron/core/transformer/moe/token_dispatcher.py b/megatron/core/transformer/moe/token_dispatcher.py index 73fbafc9d71..ff2dc810de4 100644 --- a/megatron/core/transformer/moe/token_dispatcher.py +++ b/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 diff --git a/megatron/core/transformer/utils.py b/megatron/core/transformer/utils.py index 9983f2f6dc0..4141418fda4 100644 --- a/megatron/core/transformer/utils.py +++ b/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( diff --git a/tests/unit_tests/dist_checkpointing/test_async_save.py b/tests/unit_tests/dist_checkpointing/test_async_save.py index 3c5f37c0133..dcfdcaaf3ef 100644 --- a/tests/unit_tests/dist_checkpointing/test_async_save.py +++ b/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 + ) diff --git a/tests/unit_tests/transformer/moe/test_token_dispatcher_cpu_initialization.py b/tests/unit_tests/transformer/moe/test_token_dispatcher_cpu_initialization.py new file mode 100644 index 00000000000..e3fcd436f01 --- /dev/null +++ b/tests/unit_tests/transformer/moe/test_token_dispatcher_cpu_initialization.py @@ -0,0 +1,70 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +from types import SimpleNamespace + +import pytest +import torch + +from megatron.core.transformer.moe.token_dispatcher import MoEAlltoAllTokenDispatcher +from megatron.core.transformer.transformer_config import TransformerConfig + + +def test_alltoall_dispatcher_cpu_only_initialization_does_not_access_cuda(monkeypatch): + config = TransformerConfig( + num_layers=1, + hidden_size=16, + num_attention_heads=4, + num_moe_experts=4, + moe_router_topk=2, + moe_permute_fusion=True, + use_cpu_initialization=True, + ) + process_groups = SimpleNamespace(ep=None, expt_tp=None, tp_ep=None) + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + monkeypatch.setattr( + torch.cuda, + "Stream", + lambda: pytest.fail("CPU-only dispatcher construction must not create a CUDA stream"), + ) + monkeypatch.setattr(MoEAlltoAllTokenDispatcher, "cuda_dtoh_stream", None) + + dispatcher = MoEAlltoAllTokenDispatcher( + num_local_experts=4, + local_expert_indices=[0, 1, 2, 3], + config=config, + pg_collection=process_groups, + ) + + assert dispatcher.permute_idx_device == "cpu" + assert dispatcher.sort_input_by_local_experts.device.type == "cpu" + assert dispatcher.restore_output_by_local_experts.device.type == "cpu" + assert dispatcher.cuda_dtoh_stream is None + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_alltoall_dispatcher_cpu_weight_initialization_preserves_cuda_runtime_resources( + monkeypatch, +): + config = TransformerConfig( + num_layers=1, + hidden_size=16, + num_attention_heads=4, + num_moe_experts=4, + moe_router_topk=2, + moe_permute_fusion=True, + use_cpu_initialization=True, + ) + process_groups = SimpleNamespace(ep=None, expt_tp=None, tp_ep=None) + monkeypatch.setattr(MoEAlltoAllTokenDispatcher, "cuda_dtoh_stream", None) + + dispatcher = MoEAlltoAllTokenDispatcher( + num_local_experts=4, + local_expert_indices=[0, 1, 2, 3], + config=config, + pg_collection=process_groups, + ) + + assert dispatcher.permute_idx_device.type == "cuda" + assert dispatcher.sort_input_by_local_experts.device.type == "cuda" + assert dispatcher.restore_output_by_local_experts.device.type == "cuda" + assert dispatcher.cuda_dtoh_stream is not None diff --git a/tests/unit_tests/transformer/test_dot_product_attention.py b/tests/unit_tests/transformer/test_dot_product_attention.py new file mode 100644 index 00000000000..fa251393cc7 --- /dev/null +++ b/tests/unit_tests/transformer/test_dot_product_attention.py @@ -0,0 +1,70 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +from types import SimpleNamespace + +import pytest +import torch + +from megatron.core.transformer.dot_product_attention import DotProductAttention +from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.transformer_config import TransformerConfig + + +@pytest.mark.parametrize("softmax_type", ["off-by-one", "learnable"]) +def test_cpu_initialization_keeps_softmax_offset_on_cpu(monkeypatch, softmax_type): + config = TransformerConfig( + num_layers=1, + hidden_size=16, + num_attention_heads=4, + perform_initialization=False, + softmax_type=softmax_type, + use_cpu_initialization=True, + ) + tp_group = SimpleNamespace(size=lambda: 1) + process_groups = SimpleNamespace(tp=tp_group) + monkeypatch.setattr( + torch.cuda, + "current_device", + lambda: pytest.fail("CPU-initialized attention must not access CUDA"), + ) + + attention = DotProductAttention( + config=config, + layer_number=1, + attn_mask_type=AttnMaskType.causal, + attention_type="self", + pg_collection=process_groups, + ) + + assert attention.softmax_offset.shape == (4,) + assert attention.softmax_offset.device.type == "cpu" + assert attention.softmax_offset.dtype == config.params_dtype + assert isinstance(attention.softmax_offset, torch.nn.Parameter) is (softmax_type == "learnable") + if softmax_type == "off-by-one": + assert dict(attention.named_buffers())["softmax_offset"] is attention.softmax_offset + assert "softmax_offset" not in attention.state_dict() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_cpu_initialized_off_by_one_softmax_offset_follows_module_to_cuda(): + config = TransformerConfig( + num_layers=1, + hidden_size=16, + num_attention_heads=4, + perform_initialization=False, + softmax_type="off-by-one", + use_cpu_initialization=True, + ) + tp_group = SimpleNamespace(size=lambda: 1) + process_groups = SimpleNamespace(tp=tp_group) + + attention = DotProductAttention( + config=config, + layer_number=1, + attn_mask_type=AttnMaskType.causal, + attention_type="self", + pg_collection=process_groups, + ).cuda() + + assert attention.softmax_offset.device.type == "cuda" + assert "softmax_offset" not in attention.state_dict() diff --git a/tests/unit_tests/transformer/test_spec_customization.py b/tests/unit_tests/transformer/test_spec_customization.py index bc2e6edee74..80fbec330bb 100755 --- a/tests/unit_tests/transformer/test_spec_customization.py +++ b/tests/unit_tests/transformer/test_spec_customization.py @@ -2,19 +2,32 @@ import sys from dataclasses import fields +from types import SimpleNamespace +import pytest import torch import transformer_engine as te +from megatron.core.dist_checkpointing import load, save +from megatron.core.dist_checkpointing.mapping import ( + LocalNonpersistentObject, + ShardedObject, + ShardedTensor, +) from megatron.core.extensions.transformer_engine import ( TEDotProductAttention, TELayerNormColumnParallelLinear, TENorm, TERowParallelLinear, + _bind_tenorm_sharded_state_dict, ) from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_local_submodules -from megatron.core.parallel_state import get_context_parallel_group, get_tensor_model_parallel_group +from megatron.core.parallel_state import ( + get_context_parallel_group, + get_data_parallel_group, + get_tensor_model_parallel_group, +) from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from megatron.core.transformer.attention import SelfAttention, SelfAttentionSubmodules @@ -25,10 +38,128 @@ from megatron.core.transformer.transformer_block import TransformerBlock, TransformerBlockSubmodules from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.transformer.transformer_layer import TransformerLayer, TransformerLayerSubmodules +from megatron.core.transformer.utils import sharded_state_dict_default from megatron.core.utils import is_te_min_version +from tests.unit_tests.dist_checkpointing import TempNamedDir from tests.unit_tests.test_utilities import Utils +def _fake_process_group() -> SimpleNamespace: + return SimpleNamespace(rank=lambda: 0, size=lambda: 1) + + +class TestGptLayerCheckpointKeys: + def test_dense_local_spec_uses_te_fused_layernorm_keys(self): + submodules = get_gpt_layer_local_submodules() + + assert submodules.sharded_state_dict_keys_map == { + "input_layernorm.": "self_attention.linear_qkv.layer_norm_", + "pre_mlp_layernorm.": "mlp.linear_fc1.layer_norm_", + } + + def test_moe_local_spec_keeps_te_standalone_pre_mlp_layernorm_key(self): + submodules = get_gpt_layer_local_submodules(num_experts=8) + + assert submodules.sharded_state_dict_keys_map == { + "input_layernorm.": "self_attention.linear_qkv.layer_norm_" + } + + +class TestTECheckpointCompatibility: + class _FakeTEModule(torch.nn.Module): + def __init__(self, extra_state): + super().__init__() + self.weight = torch.nn.Parameter(torch.ones(4)) + self._test_extra_state = extra_state + + def get_extra_state(self): + return self._test_extra_state + + def set_extra_state(self, state): + self._test_extra_state = state + + def test_empty_extra_state_is_local_nonpersistent(self): + norm = self._FakeTEModule(torch.empty(0, dtype=torch.uint8)) + _bind_tenorm_sharded_state_dict(norm) + + sharded_state_dict = norm.sharded_state_dict( + prefix="norm.", metadata={"dp_cp_group": _fake_process_group()} + ) + + assert isinstance(sharded_state_dict["norm._extra_state"], LocalNonpersistentObject) + assert sharded_state_dict["norm._extra_state"].unwrap().numel() == 0 + + def test_nonempty_extra_state_remains_checkpointed(self): + norm = self._FakeTEModule(torch.ones(1, dtype=torch.uint8)) + _bind_tenorm_sharded_state_dict(norm) + + sharded_state_dict = norm.sharded_state_dict( + prefix="norm.", metadata={"dp_cp_group": _fake_process_group()} + ) + + assert isinstance(sharded_state_dict["norm._extra_state"], ShardedObject) + + def test_attention_empty_extra_state_is_local_nonpersistent(self): + class FakeConfig: + softmax_type = "learnable" + + attention = self._FakeTEModule(torch.empty(0, dtype=torch.uint8)) + attention.config = FakeConfig() + attention.softmax_offset = torch.nn.Parameter(torch.ones(4)) + attention._tp_group = None + + sharded_state_dict = TEDotProductAttention.sharded_state_dict( + attention, prefix="attention.", metadata={"dp_cp_group": _fake_process_group()} + ) + + assert "attention.softmax_offset" in sharded_state_dict + assert isinstance(sharded_state_dict["attention._extra_state"], LocalNonpersistentObject) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_tenorm_keeps_runtime_and_checkpoint_hooks(self): + config = TransformerConfig(num_layers=1, hidden_size=8, num_attention_heads=2) + + norm = TENorm(config=config, hidden_size=config.hidden_size) + + assert norm.returns_residual is False + assert norm._mcore_sharded_state_dict_accepts_tp_group is True + assert norm.sharded_state_dict.__self__ is norm + + +class TestTENormTensorParallelCheckpoint: + def setup_method(self, method): + if Utils.world_size < 2: + pytest.skip("requires at least two distributed ranks") + Utils.initialize_model_parallel(2, 1) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + def test_save_load_uses_tensor_parallel_replica_rank(self, tmp_path_dist_ckpt): + config = TransformerConfig(num_layers=1, hidden_size=8, num_attention_heads=2) + norm = TENorm(config=config, hidden_size=config.hidden_size) + tp_group = get_tensor_model_parallel_group() + sharded_state_dict = sharded_state_dict_default( + norm, + prefix="norm.", + metadata={"dp_cp_group": get_data_parallel_group(with_context_parallel=True)}, + tp_group=tp_group, + ) + + assert isinstance(sharded_state_dict["norm.weight"], ShardedTensor) + assert sharded_state_dict["norm.weight"].replica_id[1] == tp_group.rank() + assert isinstance(sharded_state_dict["norm._extra_state"], LocalNonpersistentObject) + + expected_weight = norm.weight.detach().clone() + with TempNamedDir(tmp_path_dist_ckpt / "test_tenorm_tp_save_load") as checkpoint_dir: + save(sharded_state_dict, checkpoint_dir, validate_access_integrity=True) + loaded_state_dict = load( + sharded_state_dict, checkpoint_dir, validate_access_integrity=True + ) + + torch.testing.assert_close(loaded_state_dict["norm.weight"], expected_weight) + + class TestSpecCustomization: def setup_method(self, method): Utils.initialize_model_parallel(1, 1) diff --git a/tests/unit_tests/transformer/test_yarn_rope.py b/tests/unit_tests/transformer/test_yarn_rope.py new file mode 100644 index 00000000000..bf815f0282b --- /dev/null +++ b/tests/unit_tests/transformer/test_yarn_rope.py @@ -0,0 +1,54 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +import pytest +import torch + +from megatron.core.models.common.embeddings.yarn_rotary_pos_embedding import YarnRotaryEmbedding + + +class TestYarnRotaryEmbedding: + def test_cpu_initialization_keeps_cache_on_cpu(self, monkeypatch): + monkeypatch.setattr( + torch.cuda, + "current_device", + lambda: pytest.fail("CPU-initialized YARN must not access CUDA"), + ) + rope = YarnRotaryEmbedding( + kv_channels=8, use_cpu_initialization=True, original_max_position_embeddings=64 + ) + + assert rope.inv_freq_extra.device.type == 'cpu' + assert rope.inv_freq_inter.device.type == 'cpu' + assert rope.cos_cached.device.type == 'cpu' + assert rope.sin_cached.device.type == 'cpu' + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_cpu_initialized_frequencies_follow_migrated_cache(self): + rope = YarnRotaryEmbedding( + kv_channels=8, use_cpu_initialization=True, original_max_position_embeddings=64 + ).cuda() + + rope.get_cached_cos_sin(128) + + assert rope.inv_freq_extra.device.type == 'cuda' + assert rope.inv_freq_inter.device.type == 'cuda' + assert rope.cos_cached.device.type == 'cuda' + assert rope.sin_cached.device.type == 'cuda' + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_forward_cache_is_invalidated_after_device_migration(self): + rope = YarnRotaryEmbedding( + kv_channels=8, use_cpu_initialization=True, original_max_position_embeddings=64 + ) + + cpu_embedding, cpu_mscale = rope(64) + assert cpu_embedding.device.type == 'cpu' + + rope.cuda() + cuda_embedding, cuda_mscale = rope(64) + + assert cuda_embedding.device.type == 'cuda' + assert rope.inv_freq_extra.device.type == 'cuda' + assert rope.inv_freq_inter.device.type == 'cuda' + assert cuda_mscale == cpu_mscale + torch.testing.assert_close(cuda_embedding.cpu(), cpu_embedding)