diff --git a/tests/integration_tests/features.py b/tests/integration_tests/features.py index 481e520f9b..e89c716116 100755 --- a/tests/integration_tests/features.py +++ b/tests/integration_tests/features.py @@ -303,4 +303,12 @@ def build_features_test_list() -> list[OverrideDefinitions]: timeout=30, use_real_pg=True, ), + OverrideDefinitions( + configs=[recipes.kimi_k3_debugmodel_cp2], + test_descr="Kimi K3, context parallel cp2", + test_name="kimi_k3_cp2", + ngpu=2, + use_real_pg=True, + skip_rocm_test=True, + ), ] diff --git a/tests/unit_tests/cpu/test_kimi_k3_cp_contracts.py b/tests/unit_tests/cpu/test_kimi_k3_cp_contracts.py new file mode 100644 index 0000000000..cae2c13d9a --- /dev/null +++ b/tests/unit_tests/cpu/test_kimi_k3_cp_contracts.py @@ -0,0 +1,63 @@ +# 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. + +"""The CP contracts, on CPU. + +These pin the folded token layout. This model carries no batch axis, so the +Ulysses pair moves the shard between tensor dims 0 and 1; the batched spelling +of the same contract used dims 1 and 2, and nothing in a shape check would +catch that swap -- both dims exist and the all-to-all would produce a +plausible tensor with the heads and the sequence exchanged. +""" + +from __future__ import annotations + +import unittest + +import torch + +from torchtitan.models.kimi_k3.sharding import ( + contract_for_mode, + cp_all_to_all_headseq, + HEAD_DIM, + KCP, + SEQ_DIM, + ULYSSES, +) + + +class TestCPContracts(unittest.TestCase): + def test_dims_are_the_folded_ones(self): + self.assertEqual((SEQ_DIM, HEAD_DIM), (0, 1)) + + def test_ulysses_swaps_the_sharded_axis(self): + self.assertEqual(ULYSSES.in_dims(), (SEQ_DIM, HEAD_DIM)) + self.assertEqual(ULYSSES.out_dims(), (HEAD_DIM, SEQ_DIM)) + self.assertTrue(ULYSSES.redistributes()) + self.assertTrue(ULYSSES.head_sharded) + + def test_kcp_is_an_identity_pair(self): + """The recurrence passes state rank to rank, which is a sequential + dependency rather than a redistribution, so no placement pair + describes it and the contract is declared as an identity.""" + self.assertEqual(KCP.in_dims(), (SEQ_DIM, SEQ_DIM)) + self.assertFalse(KCP.redistributes()) + self.assertFalse(KCP.head_sharded) + + def test_unknown_mode_is_rejected(self): + with self.assertRaises(ValueError): + contract_for_mode("ring") + + def test_unimplemented_dim_pair_raises(self): + """A contract naming a pair with no implementation must raise here + rather than being quietly ignored.""" + x = torch.zeros(4, 2, 3) + with self.assertRaises(ValueError): + cp_all_to_all_headseq(x, None, src_dim=SEQ_DIM, dst_dim=2) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit_tests/cpu/test_kimi_k3_cp_document_mask.py b/tests/unit_tests/cpu/test_kimi_k3_cp_document_mask.py new file mode 100644 index 0000000000..5c36e5d31a --- /dev/null +++ b/tests/unit_tests/cpu/test_kimi_k3_cp_document_mask.py @@ -0,0 +1,89 @@ +# 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. + +"""The Ulysses full-sequence mask preserves packed-document boundaries. + +Two gloo ranks each hold a contiguous positions shard; the gathered mask must +equal the mask built directly from the global positions. The synthetic stream +packs three documents so that one boundary falls ON the shard cut and one +falls INSIDE a shard -- the two cases a causal-only rebuild gets wrong. +""" + +import os +import tempfile +import unittest + +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +# Three documents over 8 tokens: [0,1,2,3] [0,1] [0,1]. The cut at token 4 +# lands exactly on the second document's start; the third document starts +# inside rank 1's shard. +_POSITIONS = [0, 1, 2, 3, 0, 1, 0, 1] + + +def _worker(rank: int, world_size: int, init_file: str, out_dir: str) -> None: + dist.init_process_group( + "gloo", + init_method=f"file://{init_file}", + rank=rank, + world_size=world_size, + ) + try: + from torchtitan.models.common.attention import ( + create_attention_mask, + get_efficient_causal_mask_mod_for_packed_document, + ) + from torchtitan.models.kimi_k3.sharding import full_sequence_document_mask + + positions_full = torch.tensor(_POSITIONS, dtype=torch.int64) + shard = positions_full.chunk(world_size)[rank] + + gathered_mask = full_sequence_document_mask(None, shard, dist.group.WORLD) + reference_mask = create_attention_mask( + get_efficient_causal_mask_mod_for_packed_document(positions_full), + None, + None, + len(_POSITIONS), + len(_POSITIONS), + device=positions_full.device, + ) + from torch.nn.attention.flex_attention import create_mask + + n = len(_POSITIONS) + dense = create_mask(gathered_mask.mask_mod, 1, 1, n, n, device="cpu") + expected = create_mask(reference_mask.mask_mod, 1, 1, n, n, device="cpu") + torch.save( + {"equal": bool(torch.equal(dense, expected)), "dense": dense}, + os.path.join(out_dir, f"rank{rank}.pt"), + ) + finally: + dist.destroy_process_group() + + +class TestUlyssesDocumentMask(unittest.TestCase): + def test_gathered_mask_matches_global_positions(self): + with tempfile.TemporaryDirectory() as tmp: + init_file = os.path.join(tmp, "rdzv") + mp.spawn(_worker, args=(2, init_file, tmp), nprocs=2, join=True) + results = [torch.load(os.path.join(tmp, f"rank{r}.pt")) for r in (0, 1)] + for r, res in enumerate(results): + self.assertTrue(res["equal"], f"rank {r} mask differs from reference") + # Both ranks reassemble the same full sequence, so the masks agree. + self.assertTrue(torch.equal(results[0]["dense"], results[1]["dense"])) + # The boundary cases themselves: token 4 (doc 2 start, ON the cut) must + # not attend to token 3; token 6 (doc 3 start, inside rank 1) must not + # attend to token 5. A causal-only mask allows both. + dense = results[0]["dense"].reshape(len(_POSITIONS), len(_POSITIONS)) + self.assertFalse(bool(dense[4, 3])) + self.assertFalse(bool(dense[6, 5])) + self.assertTrue(bool(dense[3, 0])) + self.assertTrue(bool(dense[7, 6])) + + +if __name__ == "__main__": + unittest.main() diff --git a/torchtitan/distributed/fsdp.py b/torchtitan/distributed/fsdp.py index 1b00fa41d6..752570b6a1 100644 --- a/torchtitan/distributed/fsdp.py +++ b/torchtitan/distributed/fsdp.py @@ -165,6 +165,30 @@ def apply_fsdp_to_vision_encoder( ) +def add_zero_valued_dependency( + output: torch.Tensor, + unused_output: torch.Tensor, +) -> torch.Tensor: + """Keep a partly consumed FSDP module in the autograd graph. + + FSDP2 issues a module's all-gather from its pre-forward hook and its + reduce-scatter from the autograd hooks on that module's output. A rank that + consumes none of that output -- or only part of it -- would otherwise leave + the unconsumed rows outside the graph, so the collectives are issued by a + subset of the process group and the step deadlocks. + + Scaling by zero leaves ``output`` numerically unchanged while preserving the + graph edge, so every rank issues the same collectives and the module + receives zero gradients for the rows nobody used -- which is also their + correct contribution to the data-parallel average. + + Args: + output: the tensor the caller actually wants to return. + unused_output: a tensor produced by the module being kept alive. + """ + return output + unused_output.sum().to(output.dtype) * 0.0 + + def apply_fsdp_to_decoder( model: "Decoder", dp_mesh: DeviceMesh, diff --git a/torchtitan/models/common/decoder.py b/torchtitan/models/common/decoder.py index a208df1244..db2502e48b 100644 --- a/torchtitan/models/common/decoder.py +++ b/torchtitan/models/common/decoder.py @@ -136,6 +136,12 @@ def max_context_length(self) -> int: ) return rope_cfg.max_context_length + def _validate_cp_backend(self, parallelism) -> None: + """Overridable CP backend check. TODO: remove once linear-attention CP kernels are torch-native and models migrate to spmd_types.""" + from torchtitan.distributed.context_parallel import validate_cp_backend + + validate_cp_backend(parallelism) + def update_from_config( self, *, @@ -152,7 +158,6 @@ def update_from_config( that case the training/debug setup is skipped. """ from torchtitan.config import ParallelismConfig - from torchtitan.distributed.context_parallel import validate_cp_backend from torchtitan.trainer import Trainer assert hasattr(config, "parallelism"), ( @@ -171,8 +176,7 @@ def update_from_config( ) if parallelism.context_parallel_degree > 1: - # ShardingConfig-based CP requires the spmd_types backend. - validate_cp_backend(parallelism) + self._validate_cp_backend(parallelism) if any(self.traverse(ScaledDotProductAttention.Config)) or any( self.traverse(VarlenAttention.Config) ): diff --git a/torchtitan/models/common/vision_encoder.py b/torchtitan/models/common/vision_encoder.py index 37f446a2ac..9e197d5c5a 100644 --- a/torchtitan/models/common/vision_encoder.py +++ b/torchtitan/models/common/vision_encoder.py @@ -127,7 +127,13 @@ def forward( rope_cache: torch.Tensor, rope_apply: RopeApply, attention_mask: BlockMask, + cp_plan: object | None = None, ) -> torch.Tensor: + # cp_plan is ignored here and consumed by subclasses that partition an + # image across ranks. It travels as an argument, not as module state, + # because activation checkpointing recomputes this forward from the + # arguments it saved -- state set around the call is gone by then. + del cp_plan num_tokens = x.shape[0] # -1 infers the head count locally (= num_heads / TP under tensor @@ -170,12 +176,14 @@ def forward( rope_cache: torch.Tensor, rope_apply: RopeApply, attention_mask: BlockMask, + cp_plan: object | None = None, ) -> torch.Tensor: x = x + self.attn( self.norm1(x), rope_cache=rope_cache, rope_apply=rope_apply, attention_mask=attention_mask, + cp_plan=cp_plan, ) x = x + self.mlp(self.norm2(x)) return x diff --git a/torchtitan/models/kimi_k3/__init__.py b/torchtitan/models/kimi_k3/__init__.py index e559f1469c..364792fae4 100644 --- a/torchtitan/models/kimi_k3/__init__.py +++ b/torchtitan/models/kimi_k3/__init__.py @@ -16,11 +16,7 @@ from torchtitan.models.common.moe import RoutedExperts, TokenChoiceTopKRouter from torchtitan.models.common.nn_modules import GELU, RMSNorm from torchtitan.models.common.token_dispatcher import LocalTokenDispatcher -from torchtitan.models.common.vision_encoder import ( - VisionAttention, - VisionMLP, - VisionTransformerBlock, -) +from torchtitan.models.common.vision_encoder import VisionMLP, VisionTransformerBlock from torchtitan.models.kimi_k2_7.vision_encoder import VisionRotaryEmbedding2D from torchtitan.models.utils import validate_converter_order from torchtitan.protocols.model import ModelConfigConverter @@ -31,7 +27,11 @@ from .moe import KimiFeedForward, KimiGroupedExperts, KimiLatentMoE from .parallelize import parallelize_kimi_k3 from .state_dict_adapter import KimiK3StateDictAdapter -from .vision_encoder import KimiK3VisionEncoder, KimiK3VisionProjector +from .vision_encoder import ( + KimiK3VisionCPAttention, + KimiK3VisionEncoder, + KimiK3VisionProjector, +) __all__ = [ "KIMI_K3_SPECIAL_TOKENS", @@ -286,7 +286,7 @@ def _vision_encoder_config( block = VisionTransformerBlock.Config( norm1=vision_norm, norm2=vision_norm, - attn=VisionAttention.Config( + attn=KimiK3VisionCPAttention.Config( dim=qkv_dim, num_heads=num_heads, wq=_linear(dim, qkv_dim), diff --git a/torchtitan/models/kimi_k3/kda.py b/torchtitan/models/kimi_k3/kda.py index f31a8eb9b1..1804d5fdb3 100644 --- a/torchtitan/models/kimi_k3/kda.py +++ b/torchtitan/models/kimi_k3/kda.py @@ -9,6 +9,7 @@ from dataclasses import dataclass import torch +import torch.distributed as dist import torch.nn.functional as F from fla.ops.kda import chunk_kda from torch import nn @@ -66,7 +67,11 @@ def forward( beta_BLH: torch.Tensor, A_log_H: torch.Tensor, dt_bias_HK: torch.Tensor, + cp_context=None, ) -> torch.Tensor: + # cp_context turns the scan into fla's prefix-scan over rank-local + # fragments. output_final_state is unsupported there, and unneeded in + # training: the final state only matters for decoding. out_BLHV, _ = chunk_kda( q_BLHK, k_BLHK, @@ -80,10 +85,55 @@ def forward( use_beta_sigmoid_in_kernel=True, safe_gate=self.lower_bound is not None, lower_bound=self.lower_bound, + **( + {"cp_context": cp_context, "cu_seqlens": cp_context.cu_seqlens} + if cp_context is not None + else {} + ), ) return out_BLHV +def conv_with_halo(conv, x_local, cp_context, activation: str | None = None): + """Depthwise causal conv on a sequence-sharded input, exactly: the CP op + exchanges the previous rank's tail as a fixed-size halo. ``activation`` + defaults to ``conv.activation`` (fla's ``ShortConvolution`` carries one; + a plain ``nn.Conv1d`` does not).""" + from attn_gym.linear.kda.fla_cp import causal_conv1d_cp + + return causal_conv1d_cp( + x_local, + conv.weight, + conv.bias, + cp_context, + activation=getattr(conv, "activation", None) + if activation is None + else activation, + ) + + +def build_kcp_context( + seq_len_local: int, + group, + device, + conv1d_kernel_size: int | None = None, + cu_seqlens=None, +): + """CP context for one evenly split sequence. ``cu_seqlens`` must be + GLOBAL boundaries of the packed sequence; the default is one document + spanning the whole sequence, matching the non-CP call sites, which also + pass no boundaries.""" + from attn_gym.linear.kda.fla_cp import build_fla_cp_context + + return build_fla_cp_context( + seq_len_local, + group, + device, + conv1d_kernel_size=conv1d_kernel_size, + cu_seqlens=cu_seqlens, + ) + + class KimiDeltaAttention(Module): @dataclass(kw_only=True, slots=True) class Config(Module.Config): @@ -105,6 +155,9 @@ class Config(Module.Config): output_norm: KimiRMSNormGated.Config output_proj: Linear.Config + # Set by apply_cp_kimi_k3; None means the layer runs without CP. + _cp_group = None + def __init__(self, config: Config): super().__init__() self.num_heads = config.num_heads @@ -144,6 +197,10 @@ def forward( "Kimi K3 reference KDA does not support packed-document masks." ) + cp_group = self._cp_group + if cp_group is not None and dist.get_world_size(cp_group) > 1: + return self._forward_kcp(x_TD, cp_group) + num_tokens = x_TD.shape[0] q_THK = self._causal_conv(self.q_proj(x_TD), self.q_conv).view( num_tokens, self.num_heads, self.head_dim @@ -173,3 +230,65 @@ def forward( ) out_THV = self.output_norm(out_THV, output_gate_THV) return self.output_proj(out_THV.reshape(num_tokens, -1)) + + def _forward_kcp(self, x_TD: torch.Tensor, cp_group) -> torch.Tensor: + """KCP forward: the sequence stays sharded (report sec 5.1.2). + + No rank holds the full sequence. The two cross-rank dependencies have + different structure and are handled separately: the causal convolutions + need only the previous rank's tail, one fixed-size halo; the delta-rule + recurrence needs the true incoming state, which does not decompose by + summation, so fla's cp_context prefix-scans over (cumulative transition, + zero-started state) fragments. + + The folded token stream is already one packed sequence, which is exactly + what fla's CP ops assume, so this path has no batch loop. + """ + t_loc = x_TD.shape[0] + # One context serves both the conv halo and the recurrence; the conv + # needs the kernel width, the recurrence ignores it. + ctx = build_kcp_context( + t_loc, + cp_group, + x_TD.device, + conv1d_kernel_size=self.conv_kernel_size, + ) + + def conv(proj, conv_module) -> torch.Tensor: + # fla's CP conv wants [1, T, C] and applies the activation itself; + # the reference model applies SiLU outside its Conv1d, so the name + # is passed explicitly. + y_1TC = conv_with_halo( + conv_module, proj(x_TD).unsqueeze(0), ctx, activation="silu" + ) + return y_1TC.squeeze(0) + + q_THK = conv(self.q_proj, self.q_conv).view( + t_loc, self.num_heads, self.head_dim + ) + k_THK = conv(self.k_proj, self.k_conv).view( + t_loc, self.num_heads, self.head_dim + ) + v_THV = conv(self.v_proj, self.v_conv).view( + t_loc, self.num_heads, self.head_dim + ) + forget_THK = self.forget_b(self.forget_a(x_TD)).view( + t_loc, self.num_heads, self.head_dim + ) + beta_TH = self.beta(x_TD).float() + + out_THV = self.kernel( + q_THK.unsqueeze(0), + k_THK.unsqueeze(0), + v_THV.unsqueeze(0), + forget_THK.unsqueeze(0), + beta_TH.unsqueeze(0), + self.A_log, + self.dt_bias, + cp_context=ctx, + ).squeeze(0) + output_gate_THV = self.output_gate(x_TD).view( + t_loc, self.num_heads, self.head_dim + ) + out_THV = self.output_norm(out_THV, output_gate_THV) + return self.output_proj(out_THV.reshape(t_loc, -1)) diff --git a/torchtitan/models/kimi_k3/model.py b/torchtitan/models/kimi_k3/model.py index 25f473ff21..d6613d1bcb 100644 --- a/torchtitan/models/kimi_k3/model.py +++ b/torchtitan/models/kimi_k3/model.py @@ -7,10 +7,13 @@ from dataclasses import dataclass, field import torch +import torch.distributed as dist from torch import nn +from torch.distributed.tensor import DTensor -from torchtitan.hf_datasets.multimodal.mm_datasets import MMSamplePackingConfig +from torchtitan.distributed.fsdp import add_zero_valued_dependency +from torchtitan.hf_datasets.multimodal.mm_datasets import MMSamplePackingConfig from torchtitan.models.common import Linear from torchtitan.models.common.attention import ( AttentionMasksType, @@ -25,9 +28,11 @@ from torchtitan.models.common.nn_modules import RMSNorm from torchtitan.models.utils import get_moe_model_nparams_and_flops from torchtitan.protocols.module import Module +from torchtitan.tools.logging import logger from .kda import KimiDeltaAttention from .moe import KimiFeedForward, KimiLatentMoE +from .sharding import mla_ulysses_attention from .vision_encoder import KimiK3VisionEncoder # Shape suffixes: @@ -62,6 +67,11 @@ class Config(BaseAttention.Config): wo: Linear.Config inner_attention: Module.Config = field(default_factory=FlexAttention.Config) + # Set by apply_cp_kimi_k3; None means the layer runs without CP. MLA is + # Ulysses under either KDA CP mode -- KCP describes a recurrence that MLA + # does not have. + _cp_group = None + def __init__(self, config: Config): super().__init__() self.n_heads = config.n_heads @@ -88,12 +98,15 @@ def forward( attention_masks: AttentionMasksType | None = None, positions: torch.Tensor | None = None, ) -> torch.Tensor: - del positions num_tokens = x_TD.shape[0] - q_THK = self.wq_b(self.q_norm(self.wq_a(x_TD))).view( - num_tokens, self.n_heads, self.q_head_dim - ) + # The head count is DERIVED from the projection width, not read off + # self.n_heads. Ulysses splits whatever head count this rank actually + # holds, and the two differ once another parallelism has already split + # the head axis; deriving works either way and needs no branch. + q_proj_TE = self.wq_b(self.q_norm(self.wq_a(x_TD))) + h_local = q_proj_TE.shape[-1] // self.q_head_dim + q_THK = q_proj_TE.view(num_tokens, h_local, self.q_head_dim) compressed_kv_TC = self.wkv_a(x_TD) kv_latent_TC, k_rope_TK = torch.split( @@ -103,7 +116,7 @@ def forward( ) kv_THC = self.wkv_b(self.kv_norm(kv_latent_TC)).view( num_tokens, - self.n_heads, + h_local, self.qk_nope_head_dim + self.v_head_dim, ) k_nope_THK, v_THV = torch.split( @@ -112,18 +125,24 @@ def forward( dim=-1, ) k_rope_THK = k_rope_TK.view(num_tokens, 1, self.qk_rope_head_dim).expand( - -1, self.n_heads, -1 + -1, h_local, -1 ) k_THK = torch.cat((k_nope_THK, k_rope_THK), dim=-1) - out_THV = self.inner_attention( - q_THK, - k_THK, - v_THV, - attention_masks=attention_masks, - scale=self.scale, - ) - out_TD = out_THV.reshape(num_tokens, self.n_heads * self.v_head_dim) + cp_group = self._cp_group + if cp_group is not None and dist.get_world_size(cp_group) > 1: + out_THV = mla_ulysses_attention( + self, q_THK, kv_THC, k_rope_TK, cp_group, positions + ) + else: + out_THV = self.inner_attention( + q_THK, + k_THK, + v_THV, + attention_masks=attention_masks, + scale=self.scale, + ) + out_TD = out_THV.reshape(num_tokens, h_local * self.v_head_dim) out_TD = out_TD * torch.sigmoid(self.gate(x_TD)) return self.wo(out_TD) @@ -259,6 +278,29 @@ def forward( return prefix_sum_TD + h_TD, block_residual_TND +class _PlainGradBoundary(torch.autograd.Function): + """Identity forward; forces the incoming gradient to be a plain tensor. + + The vision tower must stay plain in BOTH directions. Its dynamic CP is a + separate mechanism from the decoder's, and the CP path runs hand-written + collectives whose transpose is a reduce_scatter -- + _c10d_functional.reduce_scatter_tensor has no DTensor sharding strategy. + + to_local() alone is not enough and grad_placements is the wrong knob: the + first re-wraps the gradient with the forward placements, the second states + which placements to re-wrap WITH. Neither can say "do not re-wrap". That is + what this states, and only an autograd.Function can. + """ + + @staticmethod + def forward(ctx, x): # type: ignore[override] + return x + + @staticmethod + def backward(ctx, grad): # type: ignore[override] + return grad.to_local() if isinstance(grad, DTensor) else grad + + class KimiK3Model(Decoder): @dataclass(kw_only=True, slots=True) class Config(Decoder.Config): @@ -266,6 +308,15 @@ class Config(Decoder.Config): output_res_norm: RMSNorm.Config output_res_proj: Linear.Config vision_encoder: KimiK3VisionEncoder.Config | None = None + # Smallest image worth partitioning across CP ranks. Below it the + # replicated encode is cheaper: splitting buys one gather per layer. + dynamic_cp_min_patches: int = 256 + + def _validate_cp_backend(self, parallelism) -> None: + """This model's CP is not ShardingConfig-driven -- the KDA kernels + are fla triton and never see a DTensor -- so the spmd_types + requirement does not apply; apply_cp_kimi_k3 checks its own + preconditions at wiring time.""" def update_from_config(self, *, config, **kwargs) -> None: dataset = config.dataloader.dataset @@ -273,6 +324,21 @@ def update_from_config(self, *, config, **kwargs) -> None: # and KDA recurrent states at document boundaries. if isinstance(dataset, MMSamplePackingConfig): raise ValueError("Kimi K3 does not yet support sample packing.") + parallelism = config.parallelism + if ( + parallelism.context_parallel_degree > 1 + and parallelism.context_parallel_load_balancer is not None + ): + # Both CP algorithms here read the sequence as rank-ordered + # contiguous chunks: the Ulysses all-to-all reassembles it in + # rank order, and KDA's recurrence passes state from rank r to + # rank r+1. A load balancer permutes tokens across ranks, which + # silently breaks both -- the shapes still line up. + raise ValueError( + "Kimi K3 context parallel requires " + "parallelism.context_parallel_load_balancer=None; " + f"got {parallelism.context_parallel_load_balancer!r}." + ) Decoder.Config.update_from_config(self, config=config, **kwargs) def get_nparams_and_flops( @@ -296,6 +362,9 @@ def get_nparams_and_flops( seq_len, ) + # Set by apply_cp_kimi_k3 to this model's context-parallel process group. + _cp_group = None + def __init__(self, config: Config): super().__init__(config) self.output_res_norm = config.output_res_norm.build() @@ -303,6 +372,42 @@ def __init__(self, config: Config): self.vision_encoder = ( config.vision_encoder.build() if config.vision_encoder is not None else None ) + self.dynamic_cp_min_patches = config.dynamic_cp_min_patches + self._dyncp_logged = False + + def _tower_needs_collectives(self) -> bool: + """Is the tower wrapped in something that issues per-forward collectives? + + True once FSDP has sharded it, which is when skipping it desynchronizes + the process group. A replicated DTensor -- what a tp-invariant module + holds -- issues no all-gather to match, so the test is on the placement + and not merely on the type. + """ + return any( + isinstance(p, DTensor) and any(pl.is_shard() for pl in p.placements) + for p in self.vision_encoder.parameters() + ) + + def _tower_placeholder(self) -> tuple[torch.Tensor, torch.Tensor]: + """The smallest input the tower accepts, for a rank with no images.""" + kernel_h, kernel_w = self.vision_encoder.merge_kernel_size + grid = torch.tensor( + [[1, kernel_h, kernel_w]], dtype=torch.long, device=self._device() + ) + weight = self.vision_encoder.patch_embed.weight + # A plain tensor, not weight.new_zeros: once FSDP has sharded the + # tower the weight is a DTensor, and a placeholder inheriting that + # meets the tower's own plain tensors as "aten.mm got mixed". + patches = torch.zeros( + kernel_h * kernel_w, + weight.shape[-1], + dtype=weight.dtype, + device=self._device(), + ) + return patches, grid + + def _device(self) -> torch.device: + return next(self.parameters()).device def _prepare_multimodal_embeds( self, @@ -319,6 +424,13 @@ def _prepare_multimodal_embeds( "both be omitted." ) if pixel_values is None: + # An image-free batch is normal, but FSDP2 issues the tower's + # all-gather from its pre-forward hook, so every rank must run it. + # A zero-valued placeholder keeps collectives and the DP average right. + if self.vision_encoder is not None and self._tower_needs_collectives(): + placeholder, placeholder_grid = self._tower_placeholder() + unused = self.vision_encoder(placeholder, grid_thw=placeholder_grid) + return add_zero_valued_dependency(embeddings_TD, unused) return embeddings_TD assert grid_thw is not None if self.vision_encoder is None: @@ -327,13 +439,29 @@ def _prepare_multimodal_embeds( raise ValueError("special_tokens are required for multimodal inputs.") pixel_values = pixel_values.to(self.vision_encoder.patch_embed.weight.dtype) - vision_embeds = self.vision_encoder(pixel_values, grid_thw=grid_thw) + vision_embeds = self._encode_images(pixel_values, grid_thw) # MoonViT collapses time and merges spatially, so the text-side token # count per item is (h/kh)*(w/kw), independent of t. kernel_h, kernel_w = self.vision_encoder.merge_kernel_size num_tokens_per_item = (grid_thw[:, 1] // kernel_h) * ( grid_thw[:, 2] // kernel_w ) + if self._cp_group is not None and dist.get_world_size(self._cp_group) > 1: + # This rank holds a sequence shard but encoded every image: take the + # feature slice its placeholders correspond to and scatter that. + # get_vision_positions needs whole visual items, which a shard does + # not have -- it raises "found N contiguous run(s) ... but received M + # visual item(s)" as soon as a shard splits or omits an item. + local_mask = tokens == special_tokens["image_id"] + counts = self._exchange_sentinel_counts(int(local_mask.sum().item())) + mine = self._select_cp_shard(vision_embeds, counts) + embeddings_TD = embeddings_TD.masked_scatter( + local_mask.unsqueeze(-1), mine.to(embeddings_TD.dtype) + ) + # Rows this rank did not consume still have to reach the graph, or + # the tower's reduce-scatter is issued by a subset of the group. + return add_zero_valued_dependency(embeddings_TD, vision_embeds) + vision_positions = get_vision_positions( tokens, num_tokens_per_item, @@ -345,6 +473,230 @@ def _prepare_multimodal_embeds( vision_positions=vision_positions, ) + def _exchange_sentinel_counts(self, local: int) -> torch.Tensor: + """Per-rank vision-placeholder counts across the CP group. + + Called whenever CP is on, including on ranks holding no placeholders: + the collective's participants are decided by the mesh, never by the data. + """ + group = self._cp_group + counts = torch.zeros( + dist.get_world_size(group), + dtype=torch.long, + device=torch.cuda.current_device(), + ) + counts[dist.get_rank(group)] = local + dist.all_reduce(counts, group=group) + return counts + + def _select_cp_shard( + self, vision_embeds: torch.Tensor, counts: torch.Tensor + ) -> torch.Tensor: + """Keep only the visual features belonging to this CP rank's shard. + + ``prepare_context_parallel_input`` shards inputs, labels and positions + along the sequence but leaves ``pixel_values`` whole, so every rank + encodes every image while holding only a slice of the placeholders. The + features are ordered by sequence position and the shards are contiguous + and equal -- the config rejects a load balancer under CP precisely + because a permuting one would break that -- so this rank's slice starts + after however many placeholders the lower ranks hold. + + This is correctness, not an optimization: the encoder still runs + redundantly on every CP rank. + """ + num_rows = vision_embeds.shape[0] + if int(counts.sum().item()) != num_rows: + raise ValueError( + f"CP ranks hold {int(counts.sum().item())} vision " + f"placeholder(s) in total but {num_rows} visual token(s) were " + "encoded; the sequence shard and the image batch disagree" + ) + rank = dist.get_rank(self._cp_group) + start = int(counts[:rank].sum().item()) + local = int(counts[rank].item()) + return vision_embeds[start : start + local] + + def _encode_images( + self, pixel_values: torch.Tensor, grid_thw: torch.Tensor + ) -> torch.Tensor: + """Encode every image, partitioning the large ones (report sec 5.2.3). + + Report 5.2.3 has two halves and both are needed: a single large image + is split along the patch dimension with attention gathering key-value + pairs across ranks, AND each CP group is divided into sub-CP groups + with the large images distributed across them, which is what keeps the + communication fraction from growing with scale. + + Every large image is encoded by one sub-CP group, with its patches split + across that sub-group's ranks. Images below the threshold, or whose grid + height does not divide the merge kernel, stay whole and are encoded + replicated -- splitting one buys a gather per layer and saves nothing. + """ + import torch.distributed._functional_collectives as funcol + + from torchtitan.models.kimi_k3.vision_encoder import CPPatchPlan + from torchtitan.models.kimi_k3.vit_cp_plan import ( + balance_images, + classify, + merged_tokens, + row_partition, + subgroup_layout, + ) + + grids = grid_thw.tolist() + counts = [t * h * w for t, h, w in grids] + kh, kw = self.vision_encoder.merge_kernel_size + offsets = [0] + for c in counts: + offsets.append(offsets[-1] + c) + + def _replicated(which: list[int]) -> dict[int, torch.Tensor]: + """Encode a subset redundantly on every rank.""" + out = {} + for i in which: + item = pixel_values[offsets[i] : offsets[i + 1]] + item_grid = torch.tensor( + [grids[i]], dtype=grid_thw.dtype, device=grid_thw.device + ) + out[i] = self.vision_encoder(item, grid_thw=item_grid) + return out + + subgroups = getattr(self, "_cp_subgroups", None) + group_all = self._cp_group + cp_size = dist.get_world_size(group_all) if group_all is not None else 1 + if not subgroups or cp_size <= 1: + return torch.cat( + [_replicated(list(range(len(counts))))[i] for i in range(len(counts))], + dim=0, + ) + + large = classify(counts, cp_size, min_patches=self.dynamic_cp_min_patches) + # Grid heights must divide the merge kernel for a partition to be legal. + # An image that fails it is left replicated instead of being cut unsafely. + large = [i for i in large if grids[i][1] % kh == 0] + if not large: + return torch.cat( + [_replicated(list(range(len(counts))))[i] for i in range(len(counts))], + dim=0, + ) + + n_sub, g = subgroup_layout(len(large), cp_size) + group = subgroups.get(n_sub) + if group is None or g <= 1: + # No usable sub-group of size > 1 means there is nothing to partition + # across. + return torch.cat( + [_replicated(list(range(len(counts))))[i] for i in range(len(counts))], + dim=0, + ) + + cp_rank = dist.get_rank(group_all) + my_sub = cp_rank // g + rank_in_sub = cp_rank % g + group_of = balance_images([counts[i] for i in large], n_sub) + my_large = [img for img, sub in zip(large, group_of) if sub == my_sub] + + if not self._dyncp_logged: + self._dyncp_logged = True + logger.info( + "Dynamic CP: %d large image(s) of %d over %d sub-CP group(s) of " + "%d rank(s); min_patches=%d.", + len(large), + len(counts), + n_sub, + g, + self.dynamic_cp_min_patches, + ) + + out: dict[int, torch.Tensor] = {} + # Every sub-group must run the same NUMBER of passes or the collectives + # inside them desynchronise. The count is the max over sub-groups, and a + # sub-group with fewer images pads with an empty pass. + per_sub = [sum(1 for s in group_of if s == k) for k in range(n_sub)] + n_passes = max(per_sub) if per_sub else 0 + + for p in range(n_passes): + img = my_large[p] if p < len(my_large) else None + if img is None: + # An empty pass still joins this sub-group's collectives. One + # merge block keeps every shape valid; the output is discarded. + local = pixel_values.new_zeros(kh * kw, *pixel_values.shape[1:]) + local_grid = torch.tensor( + [[1, kh, kw]], dtype=grid_thw.dtype, device=grid_thw.device + ) + plan = CPPatchPlan( + group=group, + valid_total=kh * kw * g, + full_grid=(1, kh * g, kw), + row_start=0, + band=kh, + real_rows=kh, + ) + else: + t, h, w = grids[img] + shards = row_partition(t, h, w, kh=kh, group_size=g) + sh = shards[rank_in_sub] + bands = [s.row_end - s.row_start for s in shards] + band = max(bands) + # The ceiling split keeps any deficit on the TRAILING ranks, so + # every rank's padding lands at the end of the gathered stream + # rather than inside it. Taking a prefix below depends on that. + if bands != sorted(bands, reverse=True): + raise AssertionError( + f"bands {bands} are not non-increasing; padding would land " + "inside the gathered token stream and corrupt the order" + ) + flat = pixel_values[offsets[img] : offsets[img + 1]] + # This rank's rows of EVERY frame: the projector's temporal mean + # spans all frames, so splitting by frame would give each rank the + # mean of its own frames instead. + pad_rows = band - (sh.row_end - sh.row_start) + pieces = [] + for a, b in sh.ranges: + pieces.append(flat[a:b]) + if pad_rows: + pieces.append(flat.new_zeros(pad_rows * w, *flat.shape[1:])) + local = torch.cat(pieces, dim=0) + local_grid = torch.tensor( + [[t, band, w]], dtype=grid_thw.dtype, device=grid_thw.device + ) + plan = CPPatchPlan( + group=group, + valid_total=counts[img], + full_grid=(t, h, w), + row_start=sh.row_start, + band=band, + real_rows=sh.row_end - sh.row_start, + ) + + local = local.to(self.vision_encoder.patch_embed.weight.dtype) + feats = self.vision_encoder(local, grid_thw=local_grid, cp_plan=plan) + # to_local unwraps the value but its backward re-wraps the gradient, + # and the all_gather below has a reduce_scatter transpose with no + # DTensor rule. + if isinstance(feats, DTensor): + feats = feats.to_local() + local_feat = _PlainGradBoundary.apply(feats) + # The boundary belongs on the OUTPUT too: the gradient arrives from + # downstream, so sealing only the input leaves the transpose + # receiving a DTensor. + gathered = _PlainGradBoundary.apply( + funcol.all_gather_tensor( + local_feat.contiguous(), gather_dim=0, group=group + ) + ) + if img is not None: + t, h, w = grids[img] + # NOT counts // merge: the projector collapses time, so a video's + # token count carries no t. + out[img] = gathered[: merged_tokens(h, w, kh, kw)] + + rest = [i for i in range(len(counts)) if i not in out] + if rest: + out.update(_replicated(rest)) + return torch.cat([out[i] for i in range(len(counts))], dim=0) + def forward( # pyrefly: ignore [bad-override] self, tokens: torch.Tensor, diff --git a/torchtitan/models/kimi_k3/parallelize.py b/torchtitan/models/kimi_k3/parallelize.py index 8a7d604a02..532eb189bc 100644 --- a/torchtitan/models/kimi_k3/parallelize.py +++ b/torchtitan/models/kimi_k3/parallelize.py @@ -4,6 +4,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import torch.distributed as dist import torch.nn as nn from torchtitan.config import ( @@ -18,7 +19,10 @@ apply_fsdp_to_decoder, apply_fsdp_to_vision_encoder, ) -from .model import KimiK3Model +from torchtitan.tools.logging import logger + +from .kda import KimiDeltaAttention +from .model import KimiK3Model, KimiMLAAttention def parallelize_kimi_k3( @@ -31,22 +35,21 @@ def parallelize_kimi_k3( ac_config: ActivationCheckpointingConfig, dump_folder: str, ) -> nn.Module: - """Apply FSDP2 to the Kimi K3 decoder and vision encoder.""" + """Apply FSDP2 and context parallelism to the Kimi K3 decoder and vision encoder.""" unsupported_parallelisms = [ name for name, enabled in ( ("tensor parallel", parallel_dims.tp_enabled), ("pipeline parallel", parallel_dims.pp_enabled), - ("context parallel", parallel_dims.cp_enabled), ("expert parallel", parallel_dims.ep_enabled), ) if enabled ] if unsupported_parallelisms: raise NotImplementedError( - "Kimi K3 currently supports FSDP2 data parallelism " - f"only; disable {', '.join(unsupported_parallelisms)}." + "Kimi K3 currently supports FSDP2 data parallelism and context " + f"parallelism only; disable {', '.join(unsupported_parallelisms)}." ) if parallelism.spmd_backend != "partial_dtensor": raise NotImplementedError( @@ -62,6 +65,9 @@ def parallelize_kimi_k3( dp_mesh = parallel_dims.get_mesh(dp_mesh_names) assert isinstance(model, KimiK3Model) + if parallel_dims.cp_enabled: + apply_cp_kimi_k3(model, parallel_dims) + if ac_config is not None: ac_policy = ac_config.build(dump_folder=dump_folder) ac_policy.apply(model) @@ -95,3 +101,119 @@ def parallelize_kimi_k3( ) return model + + +def apply_cp_kimi_k3( + model: nn.Module, + parallel_dims: ParallelDims, +) -> None: + """Wire context parallelism: KCP on the KDA layers, Ulysses on the MLA layers. + + Both at once, on disjoint layer kinds. Imperative rather than declared: + KDA's kernels are fla triton and never see a DTensor, so no ShardingConfig + can drive them (the model config overrides ``_validate_cp_backend`` for the + same reason). + """ + cp_group = parallel_dims.get_mesh("cp").get_group() + cp_degree = parallel_dims.cp + model._cp_group = cp_group + model._cp_subgroups = _build_cp_subgroups(cp_group) + + num_mla = 0 + kda_modules = [] + for module in model.modules(): + if isinstance(module, KimiMLAAttention): + if module.n_heads % cp_degree != 0: + raise ValueError( + f"MLA n_heads={module.n_heads} must be divisible by " + f"cp={cp_degree} for Ulysses head sharding" + ) + module._cp_group = cp_group + num_mla += 1 + elif isinstance(module, KimiDeltaAttention): + kda_modules.append(module) + + if kda_modules: + # Checked at wiring time so the message is actionable, rather than an + # ImportError from inside a layer's first forward. + try: + from attn_gym.linear.kda.fla_cp import ( # noqa: F401 + build_fla_cp_context, + causal_conv1d_cp, + ) + except ImportError as err: + raise ValueError( + "KDA context parallelism needs attention-gym's fla CP wrappers " + "(attn_gym.linear.kda.fla_cp), which wrap fla-core >= 0.5.1; " + f"import failed with: {err}." + ) from err + + for module in kda_modules: + module._cp_group = cp_group + if num_mla + len(kda_modules) == 0: + raise ValueError( + "context parallel is enabled but no attention layer was found to " + "wire it onto." + ) + logger.info( + "Applied context parallel to %d MLA and %d KDA layer(s).", + num_mla, + len(kda_modules), + ) + + +def _build_cp_subgroups(cp_group) -> dict[int, object]: + """Pre-create every sub-CP group layout this CP group could use. + + Report 5.2.3 divides each CP group into sub-CP groups so gather-KV runs inside + a sub-group instead of across the whole group. Which layout a step wants + depends on how many large images the BATCH holds, and building a process group + per batch is not an option: ``new_group`` must be called by every process in + the default group, with the same rank list, in the same order. A per-batch call + would have each rank passing its own CP group's ranks, which is exactly the + mismatch that hangs. + + So every layout is built once here and looked up per batch. The layouts are the + divisors of ``cp_size`` -- for cp=8 that is 1, 2, 4, 8 sub-groups -- so the set + is small, and an unused group costs nothing because NCCL creates its + communicator lazily on first use. + + Uniformity across ranks is achieved by all-gathering the CP rank lists first, + so every rank iterates the same global list of sub-groups in the same order and + keeps the one it belongs to. Returns ``{num_subgroups: this rank's group}``. + """ + if cp_group is None: + return {} + cp_ranks = dist.get_process_group_ranks(cp_group) + cp_size = len(cp_ranks) + if cp_size <= 1: + return {} + + # Every rank needs every CP group's membership, or the new_group calls below + # would differ between ranks. + world = dist.get_world_size() + all_cp: list[list[int] | None] = [None] * world + dist.all_gather_object(all_cp, cp_ranks) + # Deduplicate while keeping a deterministic order: identical CP groups appear + # once per member rank. + seen: list[list[int]] = [] + for entry in all_cp: + if entry and list(entry) not in seen: + seen.append(list(entry)) + seen.sort() + + my_rank = dist.get_rank() + out: dict[int, object] = {} + for n_sub in [d for d in range(1, cp_size + 1) if cp_size % d == 0]: + g = cp_size // n_sub + mine = None + for ranks in seen: + for s in range(n_sub): + members = ranks[s * g : (s + 1) * g] + # Called on every rank, same order, same lists. + pg = dist.new_group(ranks=members) + if my_rank in members: + mine = pg + if mine is not None: + out[n_sub] = mine + return out diff --git a/torchtitan/models/kimi_k3/sharding.py b/torchtitan/models/kimi_k3/sharding.py new file mode 100644 index 0000000000..fafb6396ec --- /dev/null +++ b/torchtitan/models/kimi_k3/sharding.py @@ -0,0 +1,265 @@ +# 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. + +"""Declarative CP contracts for the K3 attention layers. + +Two CP algorithms run at once on disjoint layer kinds: Ulysses on the MLA +layers, KCP on the KDA layers. Each is stated here as a placement pair on the +CP mesh axis plus the preconditions that pair implies, so ``apply_cp_kimi_k3`` +resolves a contract per module instead of branching per algorithm. + +Only the CP axis is declared. The CP collectives run on plain local tensors +after the TP-wrapped projections, at the same gap the TP plan already strips +DTensor, so TP's own head sharding is not this contract's to describe -- and +declaring both here would be two mesh axes on tensor dim 2, which SpmdLayout +rejects without an explicit partition_spec. +""" + +from dataclasses import dataclass + +import spmd_types as spmd + +import torch +import torch.distributed as dist +import torch.distributed.nn.functional as dist_nn + +from torchtitan.distributed.parallel_dims import MeshAxisName, SpmdLayout +from torchtitan.models.common.attention import ( + create_attention_mask, + get_efficient_causal_mask_mod_for_packed_document, +) + + +__all__ = [ + "CPContract", + "KCP", + "ULYSSES", + "contract_for_mode", + "cp_all_to_all_headseq", +] + +CP = MeshAxisName.CP + +# Tensor dims of the [T, H, K] activations the contracts talk about. This model +# carries a folded token stream with no batch axis, so the sequence is dim 0. +SEQ_DIM = 0 +HEAD_DIM = 1 + + +def _cp(axis_type: spmd.PerMeshAxisSpmdType) -> SpmdLayout: + return SpmdLayout(axis_types={CP: axis_type}) + + +@dataclass(frozen=True, slots=True) +class CPContract: + """What one CP algorithm does to the [B, T, H, K] activations. + + Attributes: + name: the mode spelling, and what the wiring log reports. + in_src: Placement entering the attention body. + in_dst: Placement the body computes at. + out_src: Placement leaving the body. + out_dst: Placement at the module boundary. + head_sharded: Whether the body splits heads across CP, i.e. whether + the head-divisibility precondition applies. + """ + + name: str + in_src: SpmdLayout + in_dst: SpmdLayout + out_src: SpmdLayout + out_dst: SpmdLayout + head_sharded: bool + + def redistributes(self) -> bool: + """False when in_dst == in_src, i.e. the boundary moves no data.""" + return self.in_src.axis_types != self.in_dst.axis_types + + def in_dims(self) -> tuple[int, int]: + """(src, dst) tensor dims the CP axis shards on the way in.""" + return _shard_dim(self.in_src), _shard_dim(self.in_dst) + + def out_dims(self) -> tuple[int, int]: + """(src, dst) tensor dims the CP axis shards on the way out.""" + return _shard_dim(self.out_src), _shard_dim(self.out_dst) + + +def _shard_dim(layout: SpmdLayout) -> int: + axis_type = layout.axis_types[CP] + if not isinstance(axis_type, spmd.Shard): + raise ValueError( + f"CP contract expects a Shard on the CP axis, got {axis_type!r}" + ) + return axis_type.dim + + +# Ulysses: projections run seq-local, then one all-to-all trades the sharded +# axis -- sequence for heads -- so the body sees the full sequence for its head +# subset. The output pair is the same swap reversed. +ULYSSES = CPContract( + name="ulysses", + in_src=_cp(spmd.S(SEQ_DIM)), + in_dst=_cp(spmd.S(HEAD_DIM)), + out_src=_cp(spmd.S(HEAD_DIM)), + out_dst=_cp(spmd.S(SEQ_DIM)), + head_sharded=True, +) + +# KCP: the sequence stays sharded end to end (report sec 5.1.2). The delta-rule +# recurrence carries state rank to rank -- a sequential dependency no placement +# pair describes -- so the contract is an identity, declared to keep one shape. +KCP = CPContract( + name="kcp", + in_src=_cp(spmd.S(SEQ_DIM)), + in_dst=_cp(spmd.S(SEQ_DIM)), + out_src=_cp(spmd.S(SEQ_DIM)), + out_dst=_cp(spmd.S(SEQ_DIM)), + head_sharded=False, +) + +_BY_MODE = {c.name: c for c in (ULYSSES, KCP)} + + +def contract_for_mode(mode: str) -> CPContract: + if mode not in _BY_MODE: + raise ValueError(f"cp mode must be one of {sorted(_BY_MODE)}, got {mode!r}") + return _BY_MODE[mode] + + +def cp_all_to_all_headseq( + x: torch.Tensor, cp_group, *, src_dim: int, dst_dim: int +) -> torch.Tensor: + """Differentiable Ulysses all-to-all moving the CP shard between tensor dims. + + ``(0, 1)``: ``[T/cp, H, K]`` (seq-sharded) -> ``[T, H/cp, K]``. + ``(1, 0)``: ``[T, H/cp, K]`` -> ``[T/cp, H, K]``. + + The dims come from the CP contract's placement pair rather than a flag, so a + contract that names a pair with no implementation raises here instead of being + quietly ignored. + + Numerics (round-trip and per-head chunk_kda parity) validated + bit-exact against a single-rank reference; backward is the + transposed all-to-all via torch.distributed.nn.functional. + """ + + if (src_dim, dst_dim) not in ((SEQ_DIM, HEAD_DIM), (HEAD_DIM, SEQ_DIM)): + raise ValueError( + f"no Ulysses all-to-all for CP shard dims {src_dim} -> {dst_dim}; " + f"implemented pairs are {SEQ_DIM} <-> {HEAD_DIM}" + ) + cp = dist.get_world_size(cp_group) + d0, d1, K = x.shape + if (src_dim, dst_dim) == (SEQ_DIM, HEAD_DIM): + t_loc, num_heads = d0, d1 + # [T/cp, H, K] -> [cp, T/cp, H/cp, K] (split heads by destination rank) + x_split = x.reshape(t_loc, cp, num_heads // cp, K).permute(1, 0, 2, 3) + out = dist_nn.all_to_all_single( + torch.empty_like(x_split.contiguous()), x_split.contiguous(), group=cp_group + ) + # recv[s] holds source s's T/cp for THIS rank's head subset, and s is + # already the sequence-chunk order, so the reshape stacks the sequence. + return out.reshape(cp * t_loc, num_heads // cp, K).contiguous() + t_full, h_loc = d0, d1 + t_loc = t_full // cp + # dim 0 is the destination rank: which sequence chunk each rank receives. + x_split = x.reshape(cp, t_loc, h_loc, K).contiguous() + out = dist_nn.all_to_all_single(torch.empty_like(x_split), x_split, group=cp_group) + # out[s] = source s's head subset for THIS rank's sequence chunk; put T/cp + # first so the reshape stacks heads in ascending source order. + return out.permute(1, 0, 2, 3).reshape(t_loc, cp * h_loc, K).contiguous() + + +def full_sequence_document_mask(attn, positions_L, cp_group): + """Document-aware mask for the sequence Ulysses reassembles. + + Each rank holds the full sequence after the all-to-all, so the global + packed-document mask applies as-is; only the positions must be gathered + (contiguous shards, no load balancer under CP). Rebuilt per call: the + mask follows the data, not the shape. + """ + if positions_L is None: + raise ValueError( + "context parallel needs positions to rebuild the packed-document " + "mask for the reassembled sequence, but the attention layer " + "received None." + ) + gathered = [torch.empty_like(positions_L) for _ in range(dist.get_world_size(cp_group))] + dist.all_gather(gathered, positions_L.contiguous(), group=cp_group) + positions_full = torch.cat(gathered, dim=0) + num_tokens = positions_full.shape[0] + return create_attention_mask( + get_efficient_causal_mask_mod_for_packed_document(positions_full), + None, + None, + num_tokens, + num_tokens, + device=positions_full.device, + ) + + +def mla_ulysses_attention( + attn, + q_LHQ: torch.Tensor, + kv_LHC: torch.Tensor, + k_rope_LR: torch.Tensor, + cp_group, + positions_L: torch.Tensor | None, +) -> torch.Tensor: + """MLA attention over the full sequence for this rank's head subset. + + * One fused all-to-all trades the sharded axis, sequence for heads; the + attention backend runs unchanged; a second all-to-all trades back. + * The rotary slice stays OUT of the exchange: it is headless (one vector + per token), so it is all-gathered along the sequence and expanded onto + local heads. Packing the expanded key instead reassembles it against the + wrong head subset. + * Shape suffixes beyond the legend: L local sequence (T/cp), G this rank's + head count, W packed channel width, R rotary width. + """ + cp_size = dist.get_world_size(cp_group) + t_loc = q_LHQ.shape[0] + t_full = t_loc * cp_size + # q_LHQ already carries this rank's TP-local heads, so cp splits those. + h_cp = q_LHQ.shape[1] // cp_size + + packed_LHW = torch.cat([q_LHQ, kv_LHC], dim=-1) + src_dim, dst_dim = ULYSSES.in_dims() + packed_TGW = cp_all_to_all_headseq( + packed_LHW, cp_group, src_dim=src_dim, dst_dim=dst_dim + ) + q_TGQ, k_nope_TGN, v_TGV = torch.split( + packed_TGW, + [attn.q_head_dim, attn.qk_nope_head_dim, attn.v_head_dim], + dim=-1, + ) + + # Differentiable all-gather: backward is the reduce-scatter a value every + # rank consumed needs. + k_rope_TR = torch.cat( + dist_nn.all_gather(k_rope_LR.contiguous(), group=cp_group), dim=0 + ) + k_TGQ = torch.cat( + [ + k_nope_TGN, + k_rope_TR.view(t_full, 1, attn.qk_rope_head_dim).expand( + t_full, h_cp, attn.qk_rope_head_dim + ), + ], + dim=-1, + ) + + out_TGV = attn.inner_attention( + q_TGQ, + k_TGQ, + v_TGV, + attention_masks=full_sequence_document_mask(attn, positions_L, cp_group), + scale=attn.scale, + ) + out_src_dim, out_dst_dim = ULYSSES.out_dims() + return cp_all_to_all_headseq( + out_TGV.contiguous(), cp_group, src_dim=out_src_dim, dst_dim=out_dst_dim + ) diff --git a/torchtitan/models/kimi_k3/vision_encoder.py b/torchtitan/models/kimi_k3/vision_encoder.py index 6c480c8815..f428f92111 100644 --- a/torchtitan/models/kimi_k3/vision_encoder.py +++ b/torchtitan/models/kimi_k3/vision_encoder.py @@ -15,13 +15,208 @@ from dataclasses import dataclass, field import torch +import torch.distributed as dist from torchtitan.models.common import Linear +from torchtitan.models.common.attention import create_attention_mask from torchtitan.models.common.nn_modules import GELU, RMSNorm -from torchtitan.models.kimi_k2_7.vision_encoder import MoonViTEncoder +from torchtitan.models.common.rope import ComplexRoPE +from torchtitan.models.common.vision_encoder import local_head_split, VisionAttention +from torchtitan.models.kimi_k2_7.vision_encoder import ( + _tpool_patch_merger, + MoonViTEncoder, +) from torchtitan.protocols.module import Module +@dataclass +class CPPatchPlan: + """Dynamic CP: one large image split along the PATCH dimension (report 5.2.3). + + Report 5.2.3, verbatim in substance: "A single large image is partitioned + along the patch dimension across multiple devices, and attention is computed + by gathering key-value pairs (gather-KV) across CP ranks." + + This is the half that the earlier image-level round-robin did NOT provide, and + it is the load-bearing one -- the report's stated purpose for it is to reduce + "the encoder latency of large visual samples and the cross-device load + imbalance, allowing the remaining encoder computation to be hidden in pipeline + bubbles", so DEP depends on it rather than the other way round. + + Each rank holds ``shard_len`` consecutive patches of the image and computes q + for those alone; k and v are all-gathered across ``group`` so every rank + attends over the whole image. The gather is differentiable, so its transpose + is the reduce-scatter that returns each rank the gradient for the patches it + owns. + + ``valid_total`` is the image's true patch count. A partition needs an equal + shard on every rank for a fixed-shape collective, so the tail is padded and + the padded KEY positions are masked out of attention. Without the mask the + padding would contribute to every softmax -- silently, since the shapes are + all correct. + + ``full_grid`` and ``row_start`` exist because KimiK3VisionEncoder carries position + information TWICE -- the divided_fixed absolute embedding added at the patch + embed, and 2-D RoPE applied to q/k in every block -- and both are built from + the grid starting at row 0. Describing a shard as a standalone image therefore + gives every rank the same positions, so rank 1's patches would be encoded as + if they were rank 0's. Measured before this was carried: the partitioned path + differed from the replicated one by 2.3e-03 in step-1 loss, which is far too + large for a reduction-order effect. The tables are built for the whole image + and sliced. + """ + + group: dist.ProcessGroup + valid_total: int + """The image's true patch count, for the padded-key mask.""" + + full_grid: tuple[int, int, int] = (0, 0, 0) + """(t, h, w) of the WHOLE image, not of this shard.""" + + row_start: int = 0 + """First patch-grid ROW this rank owns, in the whole image's coordinates.""" + + band: int = 0 + """Rows in this rank's tensor, including padding: the shard is (t, band, w).""" + + real_rows: int = 0 + """How many of ``band`` are real; the rest are padding.""" + + +def _slice_for_shard(table: torch.Tensor, plan: "CPPatchPlan"): + """Take this rank's ROW BAND out of a table built for the WHOLE image. + + The band is strided once the image is a video: the rank owns rows + ``[row_start, row_start + real_rows)`` of EVERY frame, because the projector's + temporal mean spans all frames and splitting by frame would break it. So the + table is gathered frame by frame and padded to ``band`` rows per frame, exactly + mirroring how the caller lays out the pixels. + + Padding rows repeat the last real row rather than being zeroed: a zeroed RoPE + factor is not a rotation. Neither choice changes the result -- padded queries + are discarded and padded keys are masked -- but staying in range keeps a NaN + out of the softmax, where it would reach real rows. + """ + t, h, w = plan.full_grid + per_frame = [] + for f in range(t): + base = f * h * w + lo = base + plan.row_start * w + hi = lo + plan.real_rows * w + rows = table[lo:hi] + pad_rows = plan.band - plan.real_rows + if pad_rows > 0: + src = rows[-1:] if rows.size(0) else table[base : base + 1] + rows = torch.cat([rows, src.expand(pad_rows * w, *table.shape[1:])], dim=0) + per_frame.append(rows) + return torch.cat(per_frame, dim=0) + + +def _padded_key_keep(plan: "CPPatchPlan", total: int, device) -> torch.Tensor: + """Which of the gathered key positions are real patches. + + NOT a prefix. ``_slice_for_shard`` pads PER FRAME -- it takes each frame's + band rows, tops that frame up to ``band``, and only then concatenates the + frames -- so a deficit rank's stream is [frame0 real, frame0 pad, frame1 + real, frame1 pad, ...] and the padding is INTERLEAVED. A prefix mask admits + frame 0's padding into the softmax and masks frame 1's real keys instead: + silently wrong encoder output whenever t > 1 and some rank is short. (t == 1 + has one frame, so a prefix happens to be right, which is why a test with a + single frame would pass either way.) + + Each rank's real row count follows from the same ceiling split + ``row_partition`` performs, so this needs no extra field and no collective: + rank r holds ``min(band, max(0, h - r * band))`` real rows. + """ + keep = torch.zeros(total, dtype=torch.bool, device=device) + t, h, _w = plan.full_grid + group_size = dist.get_world_size(plan.group) + if t > 0 and plan.band > 0 and total % (group_size * t * plan.band) == 0: + row_len = total // (group_size * t * plan.band) + pos = 0 + for r in range(group_size): + real_rows = min(plan.band, max(0, h - r * plan.band)) + for _frame in range(t): + keep[pos : pos + real_rows * row_len] = True + pos += plan.band * row_len + else: + # A plan with no grid describes a flat patch split where the padding IS + # a trailing run. Falling back matters: the branch above would compute + # from zeros, mark nothing valid, and mask every key. + keep[: plan.valid_total] = True + return keep + + +class KimiK3VisionCPAttention(VisionAttention): + """Vision attention over an image whose patches are split across ranks. + + q is this rank's patch shard; k and v are gathered so the shard attends over + the whole image (report sec 5.2.3's gather-KV). The gather is + differentiable, so its transpose is the reduce-scatter that returns each + rank the gradient for the patches it owns -- ``dist.all_gather`` would + detach and the tower would train on gradients missing every other rank's + contribution. + + Without a plan this is exactly ``VisionAttention``. + """ + + @dataclass(kw_only=True, slots=True) + class Config(VisionAttention.Config): + """Same fields; a distinct Config so build() returns this class.""" + + def forward( + self, + x: torch.Tensor, + *, + rope_cache: torch.Tensor, + rope_apply, + attention_mask, + cp_plan: "CPPatchPlan | None" = None, + ) -> torch.Tensor: + # An argument rather than module state: activation checkpointing + # recomputes this forward from its saved arguments, and state set + # around the call has been cleared by recompute time. + plan = cp_plan + if plan is None: + return super().forward( + x, + rope_cache=rope_cache, + rope_apply=rope_apply, + attention_mask=attention_mask, + ) + + import torch.distributed.nn.functional as dist_nn + + num_tokens = x.shape[0] + q_THDh = local_head_split(self.wq(x), self.head_dim) + k_THDh = local_head_split(self.wk(x), self.head_dim) + v_THDh = local_head_split(self.wv(x), self.head_dim) + q_THDh, k_THDh = rope_apply(q_THDh, k_THDh, rope_cache) + + k_full = torch.cat( + dist_nn.all_gather(k_THDh.contiguous(), group=plan.group), dim=0 + ) + v_full = torch.cat( + dist_nn.all_gather(v_THDh.contiguous(), group=plan.group), dim=0 + ) + total = k_full.size(0) + + # Built unconditionally: flex requires a BlockMask, and with no padding + # the keep vector is all true, which is the dense mask the replicated + # path would use for a single image. + keep = _padded_key_keep(plan, total, x.device) + + def _mask_mod(b, h, q_idx, kv_idx): + del b, h, q_idx + return keep[kv_idx] + + mask = create_attention_mask( + _mask_mod, None, None, q_THDh.size(0), total, device=x.device + ) + out_THDh = self.flex_attention(q_THDh, k_full, v_full, attention_masks=mask) + return self.proj(out_THDh.reshape(num_tokens, -1)) + + class KimiK3VisionProjector(Module): """PatchMergerMLPV2 projector from merged vision features to text width.""" @@ -54,3 +249,66 @@ class Config(MoonViTEncoder.Config): final_norm: RMSNorm.Config # pyrefly: ignore [bad-override] projector: KimiK3VisionProjector.Config # pyrefly: ignore [bad-override] + + def forward( + self, + pixel_values: torch.Tensor, + *, + grid_thw: torch.Tensor, + cp_plan: "CPPatchPlan | None" = None, + ) -> torch.Tensor: + """The shared tower's forward, plus report sec 5.2.3's patch partition. + + Without ``cp_plan`` this defers to the tower unchanged. With one, + ``pixel_values`` is this rank's row band of ONE image and ``grid_thw`` + still describes the whole image: the position tables are built for the + whole image and sliced, because both the learned absolute embedding and + the 2-D RoPE index from row 0, and describing a shard as a standalone + image would give every rank rank 0's positions. + + Kept here rather than in the shared tower because the partition is a + Kimi K3 feature -- it is what its report describes -- and k2.5 has no + use for it. + """ + if cp_plan is None: + return super().forward(pixel_values, grid_thw=grid_thw) + + grids = grid_thw.tolist() + if len(grids) != 1: + raise ValueError( + "a CP patch plan describes one image, but grid_thw carries " + f"{len(grids)}; a mixed stream needs a per-segment plan." + ) + # Position tables are built for the WHOLE image, then sliced to this + # rank's band: ``grid_thw`` describes only the local shard, and building + # from it gives every rank rank 0's positions. The full grid rides the plan. + x, rope_cache = self._embed_patches(pixel_values, cp_plan) + x = self._run_blocks(x, rope_cache=rope_cache, cp_plan=cp_plan) + return self._merge_and_project(self.final_norm(x), cp_plan) + + def _embed_patches(self, pixel_values, cp_plan): + """Patch embed plus position tables. Returns (x, rope_cache).""" + full_grid = [list(cp_plan.full_grid)] + learned_pos, rope_cache = self.compute_position_embeddings(full_grid) + learned_pos = _slice_for_shard(learned_pos, cp_plan) + rope_cache = _slice_for_shard(rope_cache, cp_plan) + return self.patch_embed(pixel_values) + learned_pos, rope_cache + + def _run_blocks(self, x, *, rope_cache, cp_plan): + for block in self.layers.values(): + x = block( + x, + rope_cache=rope_cache, + rope_apply=ComplexRoPE.apply_rotary_emb, + attention_mask=None, + cp_plan=cp_plan, + ) + return x + + def _merge_and_project(self, x, cp_plan): + # The merge sees the SHARD's grid: this rank holds a band of rows for + # every frame, so the (kh, kw) blocking and the temporal mean are over + # its own (t, band, w). The positions above needed the whole image. + t, _, w = cp_plan.full_grid + merged = _tpool_patch_merger(x, [[t, cp_plan.band, w]], self.merge_kernel_size) + return self.projector(merged) diff --git a/torchtitan/models/kimi_k3/vit_cp_plan.py b/torchtitan/models/kimi_k3/vit_cp_plan.py new file mode 100644 index 0000000000..e7c740754e --- /dev/null +++ b/torchtitan/models/kimi_k3/vit_cp_plan.py @@ -0,0 +1,170 @@ +# 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. + +"""Planning for dynamic CP in the vision encoder (report 5.2.3). + +Pure functions, no collectives and no torch tensors in the signatures, so the +scheduling decisions can be tested without spawning ranks. The distributed half +lives in ``model`` and ``vision_encoder``. + +Report 5.2.3 asks for two things: + +1. "A single large image is partitioned along the patch dimension across multiple + devices, and attention is computed by gathering key-value pairs (gather-KV) + across CP ranks." +2. "we divide each CP group into several sub-CP groups and distribute multiple + large images across them in a load-balanced manner, preventing the + communication fraction from growing with scale." + +The reason (2) exists is in its own clause: gather-KV over the WHOLE CP group +makes every rank exchange every large image's keys, so the communication fraction +grows with the group. Partitioning one image over a sub-group of 2 while another +image occupies a different sub-group keeps the exchange local and the ranks busy. + +**The merge kernel constrains where a partition may cut.** The projector merges +each ``(kh, kw)`` block of patches into one output token, so a cut inside a block +would ask two ranks to merge halves of the same block. The safe unit is a +MERGE-ROW BLOCK -- ``kh`` consecutive grid rows, ``kh * w`` patches. Since patches +are laid out row-major over ``(t, h, w)``, such a block is contiguous in the +packed stream and consecutive blocks abut, including across a video's frame +boundary. Cutting on arbitrary patch counts is merge-unsafe; cutting "rows r0..r1 +of every frame" is merge-safe but NOT contiguous once ``t > 1``. +""" + +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class ImageShard: + """One rank's slice of one partitioned image: a row band, all frames.""" + + row_start: int + row_end: int + """Half-open band of patch-grid ROWS, a multiple of ``kh``. Empty when the + image has fewer row blocks than the sub-group has ranks.""" + + grid: tuple[int, int, int] + """The shard's own (t, h, w) -- all frames, this rank's rows.""" + + ranges: tuple[tuple[int, int], ...] + """One flat ``[start, end)`` range per frame. A still gives one range; a video + gives ``t`` of them, because the band is strided in the packed stream.""" + + +def row_partition( + t: int, h: int, w: int, *, kh: int, group_size: int +) -> list[ImageShard]: + """Split one image across ``group_size`` ranks along the SPATIAL rows. + + Every rank keeps every frame and takes a band of rows. Two constraints force + this shape, and both were learned by measuring: + + * **The merge kernel.** The projector merges each ``(kh, kw)`` block, so a cut + inside ``kh`` rows would ask two ranks to merge halves of one block. Bands are + therefore multiples of ``kh``. + * **The temporal pool.** ``_temporal_pool_and_merge`` is ``sd2_tpool``: it takes + ``mean(dim=0)`` over ALL frames, collapsing time completely. Splitting a video + by FRAMES therefore gives each rank the mean of its own frames and the + concatenation is ``t`` times too many tokens, not the mean -- measured as a + 100% mismatch on t=2. Keeping all frames on every rank makes each rank's + temporal mean the true one for its rows. + + So the output token count of a partitioned image is ``(h/kh) * (w/kw)``, + independent of ``t``, and each rank contributes ``(band/kh) * (w/kw)`` of it. + + A band is strided in the packed stream once ``t > 1``, hence ``ranges`` rather + than one offset pair. An image with fewer row blocks than ranks leaves the tail + ranks empty; the caller pads for the fixed-shape collective. The ceiling split + keeps any deficit on the TRAILING ranks, so padding lands at the end of the + gathered stream rather than inside it. + """ + if h % kh: + raise ValueError( + f"patch grid height {h} must divide the merge kernel height {kh}; " + "the projector merges (kh, kw) blocks and a partition cannot cut " + "inside one" + ) + blocks = h // kh + per = -(-blocks // group_size) + frame = h * w + shards: list[ImageShard] = [] + for r in range(group_size): + b0 = min(r * per, blocks) + b1 = min((r + 1) * per, blocks) + r0, r1 = b0 * kh, b1 * kh + ranges = tuple((f * frame + r0 * w, f * frame + r1 * w) for f in range(t)) + shards.append( + ImageShard( + row_start=r0, + row_end=r1, + grid=(t, r1 - r0, w), + ranges=ranges, + ) + ) + return shards + + +def merged_tokens(h: int, w: int, kh: int, kw: int) -> int: + """Tokens the projector emits for one image -- time is collapsed, so ``t`` + does not appear. ``patch_count // (kh*kw)`` is only right when ``t == 1``.""" + return (h // kh) * (w // kw) + + +def subgroup_layout(num_large: int, cp_size: int) -> tuple[int, int]: + """Choose (number of sub-CP groups, ranks per sub-group). + + One large image and a CP group of 8 gives (1, 8) -- the report's "a single + large image is partitioned across multiple devices". Four large images and 8 + ranks gives (4, 2), so each image is exchanged inside a pair instead of across + all eight, which is the communication-fraction argument. + + Sub-groups are equal in size because a process group is formed from a rank + list and an uneven split would leave a sub-group whose gather is a different + shape on different ranks. So the count is the largest divisor of ``cp_size`` + that does not exceed ``num_large``. + """ + if num_large <= 0 or cp_size <= 1: + return (1, cp_size) + best = 1 + for n in range(1, cp_size + 1): + if cp_size % n == 0 and n <= num_large: + best = n + return (best, cp_size // best) + + +def balance_images(sizes: list[int], num_groups: int) -> list[int]: + """Assign each image to a sub-group, longest-processing-time-first. + + Returns ``group_of[i]`` for every entry of ``sizes``. LPT rather than + round-robin: round-robin on sizes [100, 10, 10, 10] with two groups gives 110 + against 20, while LPT gives 100 against 30. The report asks for "a + load-balanced manner" and the imbalance it is trying to remove is exactly this + one. + """ + if num_groups <= 1: + return [0] * len(sizes) + load = [0] * num_groups + group_of = [0] * len(sizes) + for i in sorted(range(len(sizes)), key=lambda j: -sizes[j]): + g = min(range(num_groups), key=lambda x: load[x]) + group_of[i] = g + load[g] += sizes[i] + return group_of + + +def classify(counts: list[int], cp_size: int, *, min_patches: int) -> list[int]: + """Indices of the images worth partitioning within a sub-group. + + An image is only worth splitting if the split leaves each rank real work: the + threshold is on the image's own patch count, not on the batch. Below it the + image-level round-robin already balances better, because splitting a small + image buys one gather per layer for nothing. + """ + if cp_size <= 1: + return [] + return [i for i, c in enumerate(counts) if c >= min_patches] diff --git a/torchtitan_recipes/tests/features.py b/torchtitan_recipes/tests/features.py index 8ecfc5140f..294ae942e3 100644 --- a/torchtitan_recipes/tests/features.py +++ b/torchtitan_recipes/tests/features.py @@ -421,3 +421,16 @@ def llama3_debugmodel_seed_checkpoint() -> Trainer.Config: config.checkpoint.create_seed_checkpoint = True config.training.disable_cuda_graphs = True return config + +def kimi_k3_debugmodel_cp2() -> Trainer.Config: + """Kimi K3 text decoder with context parallel over two ranks. + + KCP on the KDA layers and Ulysses on the MLA layers; the load balancer is + rejected under CP because both paths assume contiguous, equal shards. + """ + from torchtitan.models.kimi_k3.config_registry import kimi_k3_debugmodel + + config = kimi_k3_debugmodel() + config.parallelism.context_parallel_degree = 2 + config.parallelism.context_parallel_load_balancer = None + return config