From 719414fea04e7750508211230f1327044abf0999 Mon Sep 17 00:00:00 2001 From: Pian Pawakapan Date: Mon, 13 Jul 2026 13:08:00 -0700 Subject: [PATCH 1/2] Update (base update) [ghstack-poisoned] --- torchtitan/models/qwen3_5/sharding.py | 4 ++ torchtitan/models/qwen3_5/vision_encoder.py | 54 +++++++++++++++------ 2 files changed, 43 insertions(+), 15 deletions(-) diff --git a/torchtitan/models/qwen3_5/sharding.py b/torchtitan/models/qwen3_5/sharding.py index 366dc62ce3..c3b7f3b45a 100644 --- a/torchtitan/models/qwen3_5/sharding.py +++ b/torchtitan/models/qwen3_5/sharding.py @@ -205,6 +205,10 @@ def _set_vision_encoder_sharding(ve_cfg: "Qwen35VisionEncoder.Config") -> None: ve_cfg.sharding_config = ShardingConfig( state_shardings={"pos_embed": dense_param_placement(tp=spmd.R)}, ) + ve_cfg.rotary_pos_emb.sharding_config = ShardingConfig( + state_shardings={"inv_freq": dense_param_placement(tp=spmd.I)}, + out_src_shardings=dense_param_placement(tp=spmd.I), + ) # patch_embed receives plain pixel_values — wrap as DTensor(Replicate) ve_cfg.patch_embed_proj.sharding_config = ShardingConfig( diff --git a/torchtitan/models/qwen3_5/vision_encoder.py b/torchtitan/models/qwen3_5/vision_encoder.py index 83ceea35ee..d4624f280e 100644 --- a/torchtitan/models/qwen3_5/vision_encoder.py +++ b/torchtitan/models/qwen3_5/vision_encoder.py @@ -10,12 +10,14 @@ import torch import torch.nn as nn import torch.nn.functional as F +from torch.distributed.tensor import DTensor +from torch.distributed.tensor.experimental import local_map from torch.nn.attention.flex_attention import BlockMask, create_block_mask from torchtitan.models.common import Linear from torchtitan.models.common.attention import FlexAttention from torchtitan.models.common.nn_modules import GELU, LayerNorm -from torchtitan.models.common.rope import CosSinRoPE +from torchtitan.models.common.rope import _maybe_wrap_positions, CosSinRoPE from torchtitan.protocols.module import Module, ModuleDict _compiled_create_block_mask = torch.compile(create_block_mask) @@ -88,12 +90,21 @@ def _compute_learned_pos_embeds( ) for (h, w), indices in hw_to_indices.items(): - pos_hw = F.interpolate( - pos_grid, - size=[h, w], - mode="bilinear", - align_corners=True, - ) + if isinstance(pos_grid, DTensor): + pos_hw = local_map(F.interpolate, out_placements=(pos_grid.placements,),)( + pos_grid, + size=[h, w], # pyrefly: ignore [unexpected-keyword] + mode="bilinear", # pyrefly: ignore [unexpected-keyword] + align_corners=True, # pyrefly: ignore [unexpected-keyword] + ) + else: + pos_hw = F.interpolate( + pos_grid, + size=[h, w], + mode="bilinear", + align_corners=True, + ) + # (1, dim, h, w) → (h*w, dim) pos_hw = pos_hw.squeeze(0).permute(1, 2, 0).reshape(-1, dim).to(dtype) @@ -251,7 +262,8 @@ def forward(self, seqlen: int) -> torch.Tensor: seq = torch.arange( seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype ) - return torch.outer(seq, self.inv_freq) + seq = _maybe_wrap_positions(seq, self.inv_freq) + return torch.outer(seq, self.inv_freq) # pyrefly: ignore class PatchMerger(Module): @@ -485,13 +497,25 @@ def compute_position_embeddings( self.config.dim, ) - rope_cache = _compute_2d_rope_cache( - self._cached_freq_table, - grid_thw, - max_num_patch, - self.spatial_merge_size, - head_dim, - ) + if isinstance(self._cached_freq_table, DTensor): + rope_cache = local_map( + _compute_2d_rope_cache, + out_placements=(self._cached_freq_table.placements,), + )( + self._cached_freq_table, + grid_thw, # pyrefly: ignore [bad-argument-count] + max_num_patch, + self.spatial_merge_size, + head_dim, + ) + else: + rope_cache = _compute_2d_rope_cache( + self._cached_freq_table, + grid_thw, + max_num_patch, + self.spatial_merge_size, + head_dim, + ) return learned_pos, rope_cache From 4372470c53457906d7ae6ebe46b96e372e75edde Mon Sep 17 00:00:00 2001 From: Pian Pawakapan Date: Mon, 13 Jul 2026 13:08:00 -0700 Subject: [PATCH 2/2] Update [ghstack-poisoned] --- tests/unit_tests/test_parallel_dims.py | 33 +++++++++++++------------ torchtitan/distributed/parallel_dims.py | 29 +++++++++++++++------- torchtitan/distributed/spmd_types.py | 26 ------------------- torchtitan/protocols/sharding.py | 19 +++++++++----- 4 files changed, 50 insertions(+), 57 deletions(-) diff --git a/tests/unit_tests/test_parallel_dims.py b/tests/unit_tests/test_parallel_dims.py index ad68dcac94..5ad2bd8a0c 100644 --- a/tests/unit_tests/test_parallel_dims.py +++ b/tests/unit_tests/test_parallel_dims.py @@ -12,7 +12,7 @@ import torch import torch.distributed as dist from torch.distributed.device_mesh import init_device_mesh -from torch.distributed.tensor import Shard +from torch.distributed.tensor import Replicate from torch.distributed.tensor.debug import CommDebugMode from torch.testing._internal.distributed._tensor.common_dtensor import ( DTensorTestBase, @@ -28,7 +28,6 @@ ) from torchtitan.distributed.spmd_types import ( spmd_distribute_tensor, - spmd_layout_to_dtensor_placements, spmd_redistribute_per_axis, spmd_validate_redistributions, ) @@ -37,7 +36,7 @@ dense_sequence_parallel_placement, ) from torchtitan.models.llama3 import model_registry -from torchtitan.protocols.sharding import ShardingConfig +from torchtitan.protocols.sharding import resolve_placements, ShardingConfig class TestParallelDimsValidation(unittest.TestCase): @@ -236,19 +235,6 @@ class TestSpmdLayout(DTensorTestBase): def world_size(self): return 4 - def test_converts_partition_spec_to_dtensor_shard(self): - """PartitionSpec refines V into concrete DTensor Shard placement.""" - layout = SpmdLayout( - {MeshAxisName.TP: spmd.V}, - partition_spec=spmd.PartitionSpec(MeshAxisName.TP), - ) - - self.assertEqual(layout.per_axis_spmd_types(), {MeshAxisName.TP: spmd.S(0)}) - self.assertEqual( - spmd_layout_to_dtensor_placements(layout), - {MeshAxisName.TP: Shard(0)}, - ) - def test_seq_parallel_activation_per_axis_spmd_types(self): """PartitionSpec can map multiple mesh axes to one tensor dim.""" layout = SpmdLayout( @@ -280,6 +266,21 @@ def test_unfold_dp_axes(self): ["dp_replicate", "dp_shard", "cp", "tp"], ) + @with_comms + def test_resolve_placements_ignores_extra_untranslatable_axes(self): + """Extra layout axes are ignored before converting to DTensor placements.""" + mesh = init_device_mesh( + self.device_type, (self.world_size,), mesh_dim_names=("tp",) + ) + layout = SpmdLayout( + { + MeshAxisName.DP: spmd.V, + MeshAxisName.TP: spmd.I, + } + ) + + self.assertEqual(resolve_placements(layout, mesh), (Replicate(),)) + def test_rejects_partition_spec_reorder_redistribute(self): """((DP, CP), None) -> ((CP, DP), None) not supported by a single redistribute call.""" with self.assertRaises(ValueError) as cm: diff --git a/torchtitan/distributed/parallel_dims.py b/torchtitan/distributed/parallel_dims.py index 34d7b58009..7c5fedcfd6 100644 --- a/torchtitan/distributed/parallel_dims.py +++ b/torchtitan/distributed/parallel_dims.py @@ -19,7 +19,13 @@ from torchtitan.tools.utils import device_type -__all__ = ["MeshAxisName", "ParallelDims", "SpmdLayout", "unfold_dp_axes"] +__all__ = [ + "MeshAxisName", + "ParallelDims", + "SpmdLayout", + "unfold_dp_axis", + "unfold_dp_axes", +] class StrEnum(str, Enum): @@ -115,16 +121,21 @@ def per_axis_spmd_types(self) -> dict[MeshAxisName, spmd.PerMeshAxisSpmdType]: return result +def unfold_dp_axis(axis: MeshAxisName | str) -> tuple[MeshAxisName, ...]: + """Expand logical ``dp`` into concrete dense storage mesh axes.""" + axis_name = MeshAxisName(axis) + if axis_name == MeshAxisName.DP: + return (MeshAxisName.DP_REPLICATE, MeshAxisName.DP_SHARD) + return (axis_name,) + + def unfold_dp_axes(axes: Iterable[MeshAxisName | str]) -> list[str]: """Expand logical ``dp`` into concrete dense storage mesh axes.""" - result: list[str] = [] - for axis in axes: - axis_value = axis.value if isinstance(axis, MeshAxisName) else axis - if axis_value == "dp": - result.extend(("dp_replicate", "dp_shard")) - else: - result.append(axis_value) - return result + return [ + concrete_axis.value + for axis in axes + for concrete_axis in unfold_dp_axis(axis) + ] @dataclass diff --git a/torchtitan/distributed/spmd_types.py b/torchtitan/distributed/spmd_types.py index e0e2b50e62..2cf9e56032 100644 --- a/torchtitan/distributed/spmd_types.py +++ b/torchtitan/distributed/spmd_types.py @@ -16,7 +16,6 @@ import spmd_types as spmd import torch from torch.distributed.device_mesh import DeviceMesh -from torch.distributed.tensor import Partial, Placement, Replicate, Shard from torchtitan.distributed.utils import get_spmd_backend @@ -40,7 +39,6 @@ "set_current_spmd_mesh", "set_spmd_meshes", "maybe_set_sparse_mesh", - "spmd_layout_to_dtensor_placements", ] @@ -134,30 +132,6 @@ def maybe_set_sparse_mesh() -> Iterator[None]: yield -def spmd_layout_to_dtensor_placements( - layout: "SpmdLayout", -) -> dict["MeshAxisName", Placement]: - """Convert an SPMD layout to DTensor placements keyed by mesh axis name.""" - from torchtitan.distributed.parallel_dims import MeshAxisName - - result: dict[MeshAxisName, Placement] = {} - for axis_name, axis_type in layout.per_axis_spmd_types().items(): - if axis_type == spmd.R or axis_type == spmd.I: - dtensor_placement: Placement = Replicate() - elif axis_type == spmd.P: - dtensor_placement = Partial() - else: - assert isinstance(axis_type, spmd.Shard) - dtensor_placement = Shard(axis_type.dim) - - if axis_name == MeshAxisName.DP: - result[MeshAxisName.DP_REPLICATE] = dtensor_placement - result[MeshAxisName.DP_SHARD] = dtensor_placement - else: - result[axis_name] = dtensor_placement - return result - - def annotate_input_spmd_types( parallel_dims: "ParallelDims", inputs: torch.Tensor, diff --git a/torchtitan/protocols/sharding.py b/torchtitan/protocols/sharding.py index 118d43d964..4f3af43cfd 100644 --- a/torchtitan/protocols/sharding.py +++ b/torchtitan/protocols/sharding.py @@ -14,12 +14,15 @@ from dataclasses import dataclass, field +import spmd_types as spmd from torch.distributed.device_mesh import DeviceMesh from torch.distributed.tensor import Partial, Placement, Replicate, Shard -from torchtitan.distributed.parallel_dims import MeshAxisName, SpmdLayout - -from torchtitan.distributed.spmd_types import spmd_layout_to_dtensor_placements +from torchtitan.distributed.parallel_dims import ( + MeshAxisName, + SpmdLayout, + unfold_dp_axis, +) __all__ = [ @@ -143,18 +146,22 @@ def resolve_placements( # TODO(fegin): remove the size-1 ``Shard(d)``/``Partial`` to ``Replicate()`` # conversion once FlexShard replaces ``fully_shard``. assert mesh.mesh_dim_names is not None, "DeviceMesh must have named axes" - placements = spmd_layout_to_dtensor_placements(layout) + axis_types = {} + for axis_name, axis_type in layout.per_axis_spmd_types().items(): + for concrete_axis_name in unfold_dp_axis(axis_name): + axis_types[concrete_axis_name] = axis_type + result = [] for i, axis_name in enumerate(mesh.mesh_dim_names): key = MeshAxisName(axis_name) - if key not in placements: + if key not in axis_types: raise ValueError( f"ShardingConfig does not declare a placement for mesh axis " f"{axis_name!r}. Declared: " f"{sorted(k.value for k in layout.axes())}; " f"required: {list(mesh.mesh_dim_names)}." ) - p = placements[key] + p = spmd.spmd_type_to_dtensor_placement(axis_types[key]) if isinstance(p, (Shard, Partial)) and mesh.size(i) == 1: p = Replicate() result.append(p)