diff --git a/tests/unit_tests/test_state_dict_adapter.py b/tests/unit_tests/test_state_dict_adapter.py new file mode 100644 index 0000000000..457f27fb17 --- /dev/null +++ b/tests/unit_tests/test_state_dict_adapter.py @@ -0,0 +1,72 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import tempfile +import unittest + +import torch +import torch.distributed as dist +from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.tensor import DTensor, Replicate, Shard + +from torchtitan.models.deepseek_v3 import deepseekv3_configs +from torchtitan.models.deepseek_v3.state_dict_adapter import DeepSeekV3StateDictAdapter + + +class DeepSeekV3StateDictAdapterTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + cls._temporary_directory = tempfile.TemporaryDirectory() + cls._owns_process_group = not dist.is_initialized() + if cls._owns_process_group: + dist.init_process_group( + backend="gloo", + init_method=f"file://{cls._temporary_directory.name}/rendezvous", + rank=0, + world_size=1, + ) + + @classmethod + def tearDownClass(cls) -> None: + if cls._owns_process_group: + dist.destroy_process_group() + cls._temporary_directory.cleanup() + + def test_to_hf_handles_replicated_grouped_experts(self) -> None: + config = deepseekv3_configs["debugmodel"]( + attn_backend="flex", + moe_comm_backend="standard", + ) + adapter = DeepSeekV3StateDictAdapter(config, hf_assets_path=None) + mesh = init_device_mesh( + "cpu", + (1, 1), + mesh_dim_names=("replicate", "shard"), + ) + local_weight = torch.arange(8 * 2 * 3, dtype=torch.float32).reshape(8, 2, 3) + grouped_expert_weight = DTensor.from_local( + local_weight, + mesh, + (Replicate(), Shard(0)), + run_check=False, + ) + + hf_state_dict = adapter.to_hf( + {"layers.1.moe.routed_experts.inner_experts.w1_EFD": grouped_expert_weight} + ) + + expected_keys = { + f"model.layers.1.mlp.experts.{expert}.gate_proj.weight" + for expert in range(8) + } + self.assertEqual(set(hf_state_dict), expected_keys) + for expert in range(8): + key = f"model.layers.1.mlp.experts.{expert}.gate_proj.weight" + self.assertIsInstance(hf_state_dict[key], DTensor) + torch.testing.assert_close( + hf_state_dict[key].to_local(), + local_weight[expert], + ) diff --git a/torchtitan/models/utils.py b/torchtitan/models/utils.py index ba6da084bc..ba1a0e489e 100644 --- a/torchtitan/models/utils.py +++ b/torchtitan/models/utils.py @@ -127,7 +127,7 @@ def _caculate_indices_from_placements( # pyrefly: ignore [bad-argument-type] for i, name in enumerate(device_mesh.mesh_dim_names): placement = dtensor_placements[i] - if placement.dim == dim: + if isinstance(placement, (Shard, _StridedShard)) and placement.dim == dim: mesh_names.append(name) dim_i_placements.append(placement)