[qwen3.5] fix DTensor TP vision position caches (#3899)

Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.12.0)
(oldest at bottom):
* #3895
* __->__ #3899

Main TP was failing with `FLA_TILELANG=0 MODULE=qwen3_5
CONFIG=qwen35_debugmodel NGPU=2 ./run_train.sh
--parallelism.tensor_parallel_degree 2
--parallelism.data_parallel_shard_degree 1 --debug.seed 42
--debug.deterministic --training.steps 1`
```
[rank0]:[rank0]:   File "/data/users/pianpwk/full_spmd_types_torchtitan/tt4_spmd_types/torchtitan/torchtitan/models/qwen3_5/vision_encoder.py", line 91, in _compute_learned_pos_embeds
[rank0]:[rank0]:     pos_hw = F.interpolate(
[rank0]:[rank0]:              ^^^^^^^^^^^^^^
[rank0]:[rank0]:   File "/data/users/pianpwk/full_spmd_types_torchtitan/pytorch/torch/nn/functional.py", line 5273, in interpolate
[rank0]:[rank0]:     )._upsample_linear_vec(input, output_size, align_corners, scale_factors)
[rank0]:[rank0]:       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:[rank0]:   File "/data/users/pianpwk/full_spmd_types_torchtitan/pytorch/torch/_decomp/decompositions.py", line 4507, in _upsample_linear_vec
[rank0]:[rank0]:     return _upsample_linear(input, osize, align_corners, scales)
[rank0]:[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:[rank0]:   File "/data/users/pianpwk/full_spmd_types_torchtitan/pytorch/torch/_decomp/decompositions.py", line 91, in inner
[rank0]:[rank0]:     r = f(*tree_map(increase_prec, args), **tree_map(increase_prec, kwargs))
[rank0]:[rank0]:         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:[rank0]:   File "/data/users/pianpwk/full_spmd_types_torchtitan/pytorch/torch/_decomp/decompositions.py", line 4627, in _upsample_linear
[rank0]:[rank0]:     v = aten._unsafe_index(input, idx)
[rank0]:[rank0]:         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:[rank0]:   File "/data/users/pianpwk/full_spmd_types_torchtitan/pytorch/torch/_ops.py", line 1350, in __call__
[rank0]:[rank0]:     return self._op(*args, **kwargs)
[rank0]:[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:[rank0]: RuntimeError: aten._unsafe_index.Tensor got mixed torch.Tensor and DTensor, need to convert all torch.Tensor to DTensor before calling distributed operators!
```

- DTensor decomposition goes to unsafe_index which creates local plain
tensors - fixes by local-mapping interpolate
- upcoming spmd_types PR would parallelize inv_freq, get ahead of this
by also wrapping positions as DTensors - previously that portion was
done in plain tensor

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,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