Update

[ghstack-poisoned]

Author: pianpwk

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(

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,24 @@ 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],
+                mode="bilinear",
+                align_corners=True,
+            )
+        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,6 +265,7 @@ def forward(self, seqlen: int) -> torch.Tensor:
         seq = torch.arange(
             seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype
         )
+        seq = _maybe_wrap_positions(seq, self.inv_freq)
         return torch.outer(seq, self.inv_freq)
 
 
@@ -485,13 +500,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,
+                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