From 5a168b73ff4ec2389916e75b283ffead3e971d7c Mon Sep 17 00:00:00 2001 From: Ivy Zhou Date: Wed, 19 Aug 2026 21:52:31 -0700 Subject: [PATCH] [models] Handle replicated MoE placements in HF export Summary: Ports D115876261 to OSS. Needed by the DCPv2 512-GPU panel: the consolidation (Hugging Face export) cases cannot run without it. `_get_dtensor_shard_mesh_names` walked every placement on the mesh and read `placement.dim`, which only exists on `Shard` and `_StridedShard`. A MoE mesh that carries a `Replicate()` placement therefore raised `AttributeError` during HF export. Guard the comparison with an `isinstance` check so replicated placements are skipped instead of dereferenced. Adds `tests/unit_tests/test_state_dict_adapter.py` covering a DeepSeek V3 adapter export over a mesh that mixes `Replicate()` with sharded placements, which reproduces the failure without the fix. Test Plan: `pytest -q tests/unit_tests/test_state_dict_adapter.py` -> 1 passed in 3.31s. `pre-commit run --files torchtitan/models/utils.py tests/unit_tests/test_state_dict_adapter.py` passes all 13 hooks, including pyrefly. Verified the port is faithful: the fbsource patch restricted to `fbcode/pytorch/torchtitan` is byte-identical between the local draft and the published D115876261, and applies to the OSS tree with no fuzz after stripping the mirror prefix. `Shard` and `_StridedShard` were already imported in `torchtitan/models/utils.py`, and both symbols the new test imports (`deepseekv3_configs`, `DeepSeekV3StateDictAdapter`) exist in this tree. --- tests/unit_tests/test_state_dict_adapter.py | 72 +++++++++++++++++++++ torchtitan/models/utils.py | 2 +- 2 files changed, 73 insertions(+), 1 deletion(-) create mode 100644 tests/unit_tests/test_state_dict_adapter.py 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)