[models] Handle replicated MoE placements in HF export

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],
+            )

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)