diff --git a/tests/integration_tests/models.py b/tests/integration_tests/models.py index ff4b28c3a3..ef041693f9 100755 --- a/tests/integration_tests/models.py +++ b/tests/integration_tests/models.py @@ -13,9 +13,7 @@ def _enable_spmd_backend(t: OverrideDefinitions, backend: str) -> OverrideDefinitions: """Use ``backend`` for every variant, or return an unsupported test unchanged.""" if backend == "spmd_types" and any( - "--module kimi_k2_7" in arg or "--module muse_glimmer" in arg - for variant in t.override_args - for arg in variant + "--module muse_glimmer" in arg for variant in t.override_args for arg in variant ): return t diff --git a/torchtitan/models/common/attention.py b/torchtitan/models/common/attention.py index 839b779b3b..b1f42793ed 100644 --- a/torchtitan/models/common/attention.py +++ b/torchtitan/models/common/attention.py @@ -156,13 +156,15 @@ def forward( max_q = attention_masks.max_q max_k = attention_masks.max_k - B, L, _, H = q_BLNH.shape + B, L, _, qk_head_dim = q_BLNH.shape + value_head_dim = v_BLNH.shape[-1] T = B * L - # varlen attention expects (T, N, H) - q_TNH = q_BLNH.reshape(T, -1, H) - k_TNH = k_BLNH.reshape(T, -1, H) - v_TNH = v_BLNH.reshape(T, -1, H) + # varlen attention expects (T, N, H). The value head dimension may + # differ from the query/key head dimension, as in MLA. + q_TNH = q_BLNH.reshape(T, -1, qk_head_dim) + k_TNH = k_BLNH.reshape(T, -1, qk_head_dim) + v_TNH = v_BLNH.reshape(T, -1, value_head_dim) # Some operators can upcast under AMP, but varlen attention currently only # supports bf16/fp16 inputs. If this changes, or fp16 training support @@ -219,7 +221,7 @@ def forward( # rejected during q_TNH reshape propagation. q_ps = get_partition_spec(q_TNH) spmd.assert_type(result, q_local, q_ps) - out_BLNH = result.view(B, L, -1, H).to(q_BLNH.dtype) + out_BLNH = result.view(B, L, -1, value_head_dim).to(q_BLNH.dtype) return out_BLNH out_TNH, lse_NT = result @@ -231,7 +233,7 @@ def forward( lse_ps = None if q_ps is None else spmd.PartitionSpec(q_ps[1], q_ps[0]) spmd.assert_type(lse_NT, q_local, lse_ps) - out_BLNH = out_TNH.view(B, L, -1, H).to(q_BLNH.dtype) + out_BLNH = out_TNH.view(B, L, -1, value_head_dim).to(q_BLNH.dtype) # FA varlen returns the LSE as (N, T); reorder to (B, L, N) so # out_transform can broadcast per (token, head). lse_BLN = lse_NT.transpose(0, 1).reshape(B, L, -1) diff --git a/torchtitan/models/kimi_k2_7/__init__.py b/torchtitan/models/kimi_k2_7/__init__.py index 2373a5a0ef..d612f6b7ca 100644 --- a/torchtitan/models/kimi_k2_7/__init__.py +++ b/torchtitan/models/kimi_k2_7/__init__.py @@ -17,6 +17,7 @@ Embedding, Linear, RMSNorm, + ScaledBiasRowwiseLinear, TransformerBlock, ) from torchtitan.models.common.nn_modules import LayerNorm @@ -104,6 +105,17 @@ def _vl_linear(in_features: int, out_features: int) -> Linear.Config: ) +def _scaled_bias_rowwise_linear( + in_features: int, out_features: int +) -> ScaledBiasRowwiseLinear.Config: + return ScaledBiasRowwiseLinear.Config( + in_features=in_features, + out_features=out_features, + bias=True, + param_init=_LINEAR_INIT, + ) + + def _vl_layernorm(dim: int, eps: float = 1e-5) -> LayerNorm.Config: return LayerNorm.Config(normalized_shape=dim, eps=eps) @@ -138,11 +150,11 @@ def _vision_encoder_config( wq=_vl_linear(dim, dim), wk=_vl_linear(dim, dim), wv=_vl_linear(dim, dim), - proj=_vl_linear(dim, dim), + proj=_scaled_bias_rowwise_linear(dim, dim), ), mlp=VisionMLP.Config( fc1=_vl_linear(dim, ffn_dim), - fc2=_vl_linear(ffn_dim, dim), + fc2=_scaled_bias_rowwise_linear(ffn_dim, dim), ), ) @@ -168,7 +180,7 @@ def _vision_encoder_config( merged_dim=merged_dim, pre_norm=_vl_layernorm(dim), linear_1=_vl_linear(merged_dim, merged_dim), - linear_2=_vl_linear(merged_dim, text_hidden_size), + linear_2=_scaled_bias_rowwise_linear(merged_dim, text_hidden_size), ), ) diff --git a/torchtitan/models/kimi_k2_7/model.py b/torchtitan/models/kimi_k2_7/model.py index bfa6ffed65..9643960637 100644 --- a/torchtitan/models/kimi_k2_7/model.py +++ b/torchtitan/models/kimi_k2_7/model.py @@ -8,10 +8,14 @@ https://github.com/sgl-project/sglang/blob/e0c0c0a45cb1bda90392bfa2bba4184f5b0638a0/python/sglang/srt/models/kimi_k25.py """ +import contextlib from dataclasses import dataclass +import spmd_types as spmd import torch +from torchtitan.distributed.spmd_types import spmd_mesh_size +from torchtitan.distributed.utils import get_spmd_backend from torchtitan.models.common.attention import AttentionMasksType from torchtitan.models.common.decoder import Decoder from torchtitan.models.common.multimodal import ( @@ -20,7 +24,10 @@ ) from torchtitan.models.deepseek_v3.model import DeepSeekV3Model -from .sharding import set_kimi_k2_5_sharding_config +from .sharding import ( + annotate_multimodal_input_spmd_types, + set_kimi_k2_5_sharding_config, +) from .vision_encoder import KimiK25VisionEncoder @@ -76,6 +83,12 @@ def __init__(self, config: Config): config.vision_encoder.build() if config.vision_encoder is not None else None ) + def multimodal_context(self) -> contextlib.AbstractContextManager[None]: + """Use local DP typechecking while preparing multimodal inputs.""" + if get_spmd_backend() == "spmd_types" and spmd_mesh_size("dp") > 1: + return spmd.set_current_mesh(local_axes=("dp",)) + return contextlib.nullcontext() + def _prepare_multimodal_embeds( self, tokens: torch.Tensor, @@ -166,17 +179,31 @@ def forward( # pyrefly: ignore [bad-override] Returns: (batch, seq_len, vocab_size) logits. """ - if self.tok_embeddings is not None: - x = self._prepare_multimodal_embeds( - tokens, - pixel_values=pixel_values, - grid_thw=grid_thw, - pixel_values_videos=pixel_values_videos, - grid_thw_videos=grid_thw_videos, - special_tokens=special_tokens, # pyrefly: ignore [bad-argument-type] - ) - else: - x = tokens + with self.multimodal_context(): + if get_spmd_backend() == "spmd_types": + annotate_multimodal_input_spmd_types( + pixel_values=pixel_values, + grid_thw=grid_thw, + pixel_values_videos=pixel_values_videos, + grid_thw_videos=grid_thw_videos, + ) + + if self.tok_embeddings is not None: + x = self._prepare_multimodal_embeds( + tokens, + pixel_values=pixel_values, + grid_thw=grid_thw, + pixel_values_videos=pixel_values_videos, + grid_thw_videos=grid_thw_videos, + special_tokens=special_tokens, # pyrefly: ignore [bad-argument-type] + ) + else: + x = tokens + + if get_spmd_backend() == "spmd_types": + # The scatter restores a token-aligned tensor, so text-model DP + # resumes as global batch sharding after the multimodal region. + spmd.assert_type(x, {"dp": spmd.S(0), "tp": spmd.R}) for layer in self.layers.values(): x = layer(x, attention_masks, positions) diff --git a/torchtitan/models/kimi_k2_7/parallelize.py b/torchtitan/models/kimi_k2_7/parallelize.py index a671db3867..ab82749e2f 100644 --- a/torchtitan/models/kimi_k2_7/parallelize.py +++ b/torchtitan/models/kimi_k2_7/parallelize.py @@ -32,6 +32,7 @@ from torchtitan.distributed.full_dtensor import ( resolve_fsdp_mesh, resolve_sparse_fsdp_mesh, + validate_config, ) @@ -67,11 +68,10 @@ def parallelize_kimi_k2_5( compile_config.enable and "model" in compile_config.components ) - if ( - parallelism.spmd_backend == "spmd_types" - or parallel_dims.tp_enabled - or parallel_dims.ep_enabled - ): + if parallelism.spmd_backend == "spmd_types": + validate_config(parallel_dims, model) + model.parallelize(parallel_dims) # pyrefly: ignore [not-callable] + elif parallel_dims.tp_enabled or parallel_dims.ep_enabled: model.parallelize(parallel_dims) # pyrefly: ignore [not-callable] if ac_config is not None: diff --git a/torchtitan/models/kimi_k2_7/sharding.py b/torchtitan/models/kimi_k2_7/sharding.py index cb4092adf3..109ae6ddc9 100644 --- a/torchtitan/models/kimi_k2_7/sharding.py +++ b/torchtitan/models/kimi_k2_7/sharding.py @@ -12,39 +12,62 @@ - Decoder (MLA + MoE): reuses ``set_deepseek_v3_sharding_config``. Multimodal configs keep the token embedding ``Replicate`` for the vision scatter and resume SP at layer 0 (see ``_shard_decoder_after_embedding_scatter``). -- Vision encoder: activations flow ``Replicate`` (no SP -- the patch sequence is +- Vision encoder: activations flow ``Invariant`` (no SP -- the patch sequence is short, so sequence-sharding would add gather/scatter around the block-diagonal attention for little memory gain). Only the linear layers are Colwise/Rowwise - sharded for memory; norms and position embeddings stay ``Replicate``. + sharded for memory; norms and position embeddings stay ``Invariant``. """ from typing import TYPE_CHECKING import spmd_types as spmd +import torch +from torchtitan.distributed.parallel_dims import MeshAxisName from torchtitan.models.common.decoder_sharding import ( - colwise_config, dense_activation_placement, dense_param_placement, - rowwise_config, - set_gqa_inner_attention_local_map, + dense_sequence_parallel_placement, +) +from torchtitan.models.common.vision_encoder_sharding import ( + invariant_norm_config, + set_vision_transformer_block_sharding_config, + vision_colwise_config, + vision_invariant_linear_config, + vision_scaled_bias_rowwise_config, ) from torchtitan.models.deepseek_v3.sharding import set_deepseek_v3_sharding_config -from torchtitan.protocols.sharding import LocalMapConfig, ShardingConfig +from torchtitan.protocols.sharding import LocalMapConfig, ShardingConfig, SpmdLayout + +DP = MeshAxisName.DP +TP = MeshAxisName.TP if TYPE_CHECKING: from torchtitan.models.kimi_k2_7.model import KimiK25Model -_REPLICATE_PARAM = dense_param_placement(tp=spmd.R) _REPLICATE_ACT = dense_activation_placement(tp=spmd.R) -_REPLICATE_NORM = ShardingConfig( - state_shardings={"weight": _REPLICATE_PARAM, "bias": _REPLICATE_PARAM}, - in_src_shardings={"input": _REPLICATE_ACT}, - in_dst_shardings={"input": _REPLICATE_ACT}, - out_src_shardings=_REPLICATE_ACT, - out_dst_shardings=_REPLICATE_ACT, -) + +def annotate_multimodal_input_spmd_types( + *, + pixel_values: torch.Tensor | None, + grid_thw: torch.Tensor | None, + pixel_values_videos: torch.Tensor | None, + grid_thw_videos: torch.Tensor | None, +) -> None: + """Annotate Kimi K2.5 multimodal inputs with their local SPMD types.""" + multimodal_type = { + MeshAxisName.DP: spmd.V, + MeshAxisName.TP: spmd.I, + } + for tensor in ( + pixel_values, + grid_thw, + pixel_values_videos, + grid_thw_videos, + ): + if tensor is not None: + spmd.assert_type(tensor, multimodal_type) def set_kimi_k2_5_sharding_config( @@ -65,15 +88,7 @@ def set_kimi_k2_5_sharding_config( def _shard_decoder_after_embedding_scatter(config: "KimiK25Model.Config") -> None: - """Keep ``tok_embeddings`` ``Replicate`` and resume SP at layer 0's output. - - The vision scatter writes features at arbitrary sequence positions, so it - needs the full (``Replicate``) embedding -- a ``Shard(1)`` one cannot be - indexed by sequence position locally. Layer 0 then takes a ``Replicate`` - input and its rowwise ``wo`` reduce-scatters back to ``Shard(1)``, so the - residual is sequence-parallel from layer 0's output and layers ``1..N-1`` - are unchanged full SP. - """ + """Keep the full embedding through vision scatter, then resume SP at layer 0.""" config.tok_embeddings.sharding_config = ShardingConfig( state_shardings={"weight": dense_param_placement(tp=spmd.S(0))}, in_src_shardings={"input": _REPLICATE_ACT}, @@ -84,62 +99,46 @@ def _shard_decoder_after_embedding_scatter(config: "KimiK25Model.Config") -> Non ) layer0 = config.layers[0] - layer0.attention_norm.sharding_config = ShardingConfig( - state_shardings={"weight": _REPLICATE_PARAM}, - in_src_shardings={"input": _REPLICATE_ACT}, - out_src_shardings=_REPLICATE_ACT, - ) - layer0.attention.sharding_config = ShardingConfig( + layer0.sharding_config = ShardingConfig( in_src_shardings={"x": _REPLICATE_ACT}, - in_dst_shardings={"x": _REPLICATE_ACT}, + in_dst_shardings={"x": dense_sequence_parallel_placement()}, + out_src_shardings=dense_sequence_parallel_placement(), ) def _set_vision_encoder_sharding(ve_cfg) -> None: - """Replicate-activation TP plan for the MoonViT3d vision encoder. + """Invariant-activation TP plan for the MoonViT3d vision encoder. Linear layers are Colwise/Rowwise sharded for memory; norms and the - learnable position table are Replicate. ``patch_embed`` wraps the plain - ``pixel_values`` input as ``DTensor(Replicate)`` so the rest of the encoder - runs in DTensor space. + learnable position table stay Invariant. ``patch_embed`` wraps the plain + ``pixel_values`` input as a TP-invariant tensor so the rest of the encoder + runs in distributed tensor space. """ - # The encoder's own ``pos_embed`` table is Replicate (F.interpolate runs on it). + # The encoder's own ``pos_embed`` table is invariant across TP ranks. ve_cfg.sharding_config = ShardingConfig( - state_shardings={"pos_embed": _REPLICATE_PARAM}, + state_shardings={ + "pos_embed": SpmdLayout({DP: spmd.R, TP: spmd.I}), + }, + out_src_shardings=SpmdLayout({DP: spmd.V, TP: spmd.I}), + out_dst_shardings=SpmdLayout({DP: spmd.V, TP: spmd.R}), ) - - # patch_embed (Linear): receives plain pixel_values -> wrap as Replicate. - ve_cfg.patch_embed_proj.sharding_config = ShardingConfig( - state_shardings={"weight": _REPLICATE_PARAM, "bias": _REPLICATE_PARAM}, - in_src_shardings={"input": _REPLICATE_ACT}, - in_dst_shardings={"input": _REPLICATE_ACT}, - out_src_shardings=_REPLICATE_ACT, - out_dst_shardings=_REPLICATE_ACT, + ve_cfg.rotary_pos_emb.sharding_config = ShardingConfig( + state_shardings={ + "inv_freq": SpmdLayout({DP: spmd.R, TP: spmd.I}), + }, + out_src_shardings=SpmdLayout({DP: spmd.R, TP: spmd.I}), ) - # Transformer block sub-modules (shared VisionTransformerBlock: norm1/norm2). - block = ve_cfg.block - block.norm1.sharding_config = _REPLICATE_NORM - block.norm2.sharding_config = _REPLICATE_NORM + ve_cfg.patch_embed_proj.sharding_config = vision_invariant_linear_config() - # The stacked 2D rope_cache enters the attention as a plain (Replicate) - # tensor input so it is DTensor-wrapped before meeting head-sharded q/k. - block.attn.sharding_config = ShardingConfig( - in_src_shardings={"rope_cache": _REPLICATE_ACT}, - in_dst_shardings={"rope_cache": _REPLICATE_ACT}, + set_vision_transformer_block_sharding_config( + ve_cfg.block, + rope_cache_dp=spmd.V, ) - block.attn.wq.sharding_config = colwise_config() - block.attn.wk.sharding_config = colwise_config() - block.attn.wv.sharding_config = colwise_config() - block.attn.proj.sharding_config = rowwise_config(output_sp=False) - set_gqa_inner_attention_local_map(block.attn.inner_attention) - - block.mlp.fc1.sharding_config = colwise_config() - block.mlp.fc2.sharding_config = rowwise_config(output_sp=False) # Final norm + projector. - ve_cfg.final_norm.sharding_config = _REPLICATE_NORM + ve_cfg.final_norm.sharding_config = invariant_norm_config() proj = ve_cfg.projector - proj.pre_norm.sharding_config = _REPLICATE_NORM - proj.linear_1.sharding_config = colwise_config() - proj.linear_2.sharding_config = rowwise_config(output_sp=False) + proj.pre_norm.sharding_config = invariant_norm_config() + proj.linear_1.sharding_config = vision_colwise_config() + proj.linear_2.sharding_config = vision_scaled_bias_rowwise_config() diff --git a/torchtitan/models/kimi_k2_7/vision_encoder.py b/torchtitan/models/kimi_k2_7/vision_encoder.py index 5accd91f85..682754238a 100644 --- a/torchtitan/models/kimi_k2_7/vision_encoder.py +++ b/torchtitan/models/kimi_k2_7/vision_encoder.py @@ -18,14 +18,17 @@ """ from dataclasses import dataclass, field +from typing import cast +import spmd_types as spmd import torch import torch.nn as nn import torch.nn.functional as F +from torchtitan.distributed.utils import get_spmd_backend from torchtitan.models.common import Linear from torchtitan.models.common.nn_modules import GELU, LayerNorm -from torchtitan.models.common.rope import ComplexRoPE +from torchtitan.models.common.rope import _maybe_wrap_positions, ComplexRoPE from torchtitan.models.common.vision_encoder import ( compiled_create_block_mask, get_vision_block_mask_mod, @@ -86,6 +89,8 @@ def _compute_learned_pos_embeds( """ height, width, dim = pos_embed.shape pos = pos_embed.new_zeros(len(grids), max_num_patch, dim) + if get_spmd_backend() == "spmd_types" and spmd.is_type_checking(): + pos = spmd.mutate_type(pos, src=spmd.R, dst={"dp": spmd.V, "tp": spmd.I}) # (dim, height, width) for F.interpolate; .float() for bicubic. grid_table = pos_embed.permute(2, 0, 1).unsqueeze(0).float() @@ -155,9 +160,9 @@ def _compute_2d_rope_cache( """ device = freq_table.device - angles = torch.zeros( - len(grids), max_num_patch, head_dim // 2, device=device, dtype=freq_table.dtype - ) + angles = freq_table.new_zeros(len(grids), max_num_patch, head_dim // 2) + if get_spmd_backend() == "spmd_types" and spmd.is_type_checking(): + angles = spmd.mutate_type(angles, src=spmd.R, dst={"dp": spmd.V, "tp": spmd.I}) # Group by (h, w) so the per-resolution angle grid is built once. hw_to_indices: dict[tuple[int, int], list[int]] = {} @@ -168,6 +173,9 @@ def _compute_2d_rope_cache( # Raster order: position p -> (row = p // w, col = p % w). Gather each # axis's angles from the precomputed table (freq_table[pos] = pos*inv_freq). flat = torch.arange(h * w, device=device) + flat = cast(torch.Tensor, _maybe_wrap_positions(flat, freq_table)) + if get_spmd_backend() == "spmd_types" and spmd.is_type_checking(): + flat = spmd.mutate_type(flat, "tp", src=spmd.R, dst=spmd.I) x_ang = freq_table[flat % w] # (h*w, head_dim/4) column y_ang = freq_table[flat // w] # (h*w, head_dim/4) row # Interleave x/y so pair 2k uses x-position, pair 2k+1 uses y-position. @@ -211,6 +219,8 @@ def _tpool_patch_merger( max_merged = max((h // kh) * (w // kw) for _, h, w in grids) merged = hidden_NPD.new_zeros(num_vision, max_merged, merged_dim) + if get_spmd_backend() == "spmd_types" and spmd.is_type_checking(): + merged = spmd.mutate_type(merged, src=spmd.R, dst={"dp": spmd.V, "tp": spmd.I}) for i, (t, h, w) in enumerate(grids): seq = hidden_NPD[i, : t * h * w] @@ -271,6 +281,9 @@ def forward(self, seqlen: int) -> torch.Tensor: seq = torch.arange( seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype ) + seq = cast(torch.Tensor, _maybe_wrap_positions(seq, self.inv_freq)) + if get_spmd_backend() == "spmd_types" and spmd.is_type_checking(): + seq = spmd.mutate_type(seq, "tp", src=spmd.R, dst=spmd.I) return torch.outer(seq, self.inv_freq) @@ -442,14 +455,15 @@ def forward( x = self.patch_embed(pixel_values) + learned_pos mask_mod = get_vision_block_mask_mod(num_patch) - attention_mask = compiled_create_block_mask( - mask_mod, - num_vision, - None, - max_num_patch, - max_num_patch, - device=x.device, - ) + with spmd.no_typecheck(): + attention_mask = compiled_create_block_mask( + mask_mod, + num_vision, + None, + max_num_patch, + max_num_patch, + device=x.device, + ) for block in self.layers.values(): x = block(