From 7eea0b18ec6c7bbd25556b267ab7d45c90292b60 Mon Sep 17 00:00:00 2001 From: QIU023 Date: Mon, 24 Aug 2026 19:43:48 +0000 Subject: [PATCH 01/14] [kimi k3] text-decoder context parallelism (PR-4025) KCP on the KDA layers, Ulysses on the MLA layers, on the plain CP group. Text-only slice: no vision dynamic CP, no vision-tower CP attention. --- torchtitan/models/common/decoder.py | 18 +- torchtitan/models/kimi_k3/__init__.py | 13 ++ torchtitan/models/kimi_k3/config_registry.py | 46 ++++ torchtitan/models/kimi_k3/dtensor_ops.py | 77 +++++++ torchtitan/models/kimi_k3/kcp.py | 104 ++++++++++ torchtitan/models/kimi_k3/kda.py | 196 +++++++++++++++++- torchtitan/models/kimi_k3/model.py | 192 +++++++++++++++-- torchtitan/models/kimi_k3/parallelize.py | 120 ++++++++++- torchtitan/models/kimi_k3/sharding.py | 172 +++++++++++++++ torchtitan/models/kimi_k3/tests/__init__.py | 5 + .../models/kimi_k3/tests/test_cp_contracts.py | 63 ++++++ 11 files changed, 985 insertions(+), 21 deletions(-) create mode 100644 torchtitan/models/kimi_k3/dtensor_ops.py create mode 100644 torchtitan/models/kimi_k3/kcp.py create mode 100644 torchtitan/models/kimi_k3/sharding.py create mode 100644 torchtitan/models/kimi_k3/tests/__init__.py create mode 100644 torchtitan/models/kimi_k3/tests/test_cp_contracts.py diff --git a/torchtitan/models/common/decoder.py b/torchtitan/models/common/decoder.py index a208df1244..fd28efb34c 100644 --- a/torchtitan/models/common/decoder.py +++ b/torchtitan/models/common/decoder.py @@ -81,6 +81,21 @@ class Config(BaseModel.Config): # that support it set this True in their config factories; the tying # itself is handled by ``Decoder.__init__`` / ``Decoder.init_states``. enable_weight_tying: bool = False + # Whether this model's context parallel is driven by ShardingConfig. + # ``validate_cp_backend`` is documented as being for "the models that + # declare CP in ShardingConfig", but the check below runs for every + # decoder, so a model whose CP cannot be declarative has no way to say + # so. Kernels that do not dispatch through DTensor force the case: + # nothing in a ShardingConfig can reach them, and the model implements + # CP itself. Such a model sets this False and takes on its own CP + # preconditions. + # + # Deriving it instead -- asking whether any placement in the config + # names the CP axis -- was tried and is wrong: every model here sets + # its sharding AFTER delegating to this method, llama3 included, so + # at this point no model declares anything and the check would be + # silently skipped for all of them. + cp_via_sharding_config: bool = True @property def first_attention(self) -> BaseAttention.Config | None: @@ -172,7 +187,8 @@ def update_from_config( if parallelism.context_parallel_degree > 1: # ShardingConfig-based CP requires the spmd_types backend. - validate_cp_backend(parallelism) + if self.cp_via_sharding_config: + validate_cp_backend(parallelism) if any(self.traverse(ScaledDotProductAttention.Config)) or any( self.traverse(VarlenAttention.Config) ): diff --git a/torchtitan/models/kimi_k3/__init__.py b/torchtitan/models/kimi_k3/__init__.py index e559f1469c..932fc728e5 100644 --- a/torchtitan/models/kimi_k3/__init__.py +++ b/torchtitan/models/kimi_k3/__init__.py @@ -488,6 +488,18 @@ def _debugmodel(attn_backend: str) -> KimiK3Model.Config: ) +def _debugmodel_text(attn_backend: str) -> KimiK3Model.Config: + """The debug decoder with no vision tower. + + The text arm of the context-parallel matrix needs a flavor with no vision + path, so that a failure there is attributable to the decoder's CP rather + than to the tower or to the image/text token interleaving. + """ + config = _debugmodel(attn_backend) + config.vision_encoder = None + return config + + def _kimi_k3(attn_backend: str) -> KimiK3Model.Config: dim = 7168 return _kimi_k3_config( @@ -526,6 +538,7 @@ def _kimi_k3(attn_backend: str) -> KimiK3Model.Config: kimi_k3_configs = { "debugmodel": _debugmodel, + "debugmodel_text": _debugmodel_text, "Kimi-K3": _kimi_k3, } diff --git a/torchtitan/models/kimi_k3/config_registry.py b/torchtitan/models/kimi_k3/config_registry.py index f1eac6b4fb..4ca621071b 100644 --- a/torchtitan/models/kimi_k3/config_registry.py +++ b/torchtitan/models/kimi_k3/config_registry.py @@ -8,6 +8,7 @@ from torchtitan.components.checkpointer import CheckpointManager from torchtitan.components.data import GrainDataLoader, SingleDatasetConfig +from torchtitan.components.data.packing import ConcatThenSplitPackingConfig from torchtitan.components.loss import ChunkedLossWrapper, CrossEntropyLoss from torchtitan.components.metrics import MetricsProcessor from torchtitan.components.optimizer import default_adamw, LRSchedulersContainer @@ -20,6 +21,7 @@ MultiModalProcessor, ) from torchtitan.hf_datasets.multimodal.utils.image import resize_to_patch_budget +from torchtitan.hf_datasets.text_datasets import DATASETS as TEXT_DATASETS from torchtitan.models.common.config_utils import decoder_vocab_size from torchtitan.trainer import Trainer @@ -94,3 +96,47 @@ def kimi_k3_debugmodel() -> Trainer.Config: ), activation_checkpoint=SelectiveAC.Config(), ) + + +def kimi_k3_debugmodel_text() -> Trainer.Config: + """The debug model with no vision tower, on packed text. + + The text arm of the context-parallel matrix: a CP failure here is the + decoder's, not the vision tower's. + """ + model_spec = model_registry("debugmodel_text") + return Trainer.Config( + loss=ChunkedLossWrapper.Config( + loss_fn=CrossEntropyLoss.Config( + global_vocab_size=decoder_vocab_size(model_spec), + ), + ), + hf_assets_path="./tests/assets/tokenizer", + tokenizer=MultiModalTokenizer.Config(**KIMI_K3_SPECIAL_TOKENS), + metrics=MetricsProcessor.Config(log_freq=1), + model_spec=model_spec, + dataloader=GrainDataLoader.Config( + dataset=ConcatThenSplitPackingConfig(dataset=TEXT_DATASETS["c4_test"]), + ), + optimizer=default_adamw(lr=8e-4), + lr_scheduler=LRSchedulersContainer.Config( + warmup_steps=2, + decay_ratio=0.8, + decay_type="linear", + min_lr_factor=0.0, + ), + # TODO: Kimi K3 has no spmd_types annotations yet. + parallelism=ParallelismConfig(spmd_backend="partial_dtensor"), + training=TrainingConfig( + num_tokens_per_microbatch_per_dp_rank=256, + max_context_length=256, + steps=10, + dtype="bfloat16", + disable_cuda_graphs=True, + ), + checkpoint=CheckpointManager.Config( + interval=10, + last_save_model_only=False, + ), + activation_checkpoint=SelectiveAC.Config(), + ) diff --git a/torchtitan/models/kimi_k3/dtensor_ops.py b/torchtitan/models/kimi_k3/dtensor_ops.py new file mode 100644 index 0000000000..f2ca5e01bd --- /dev/null +++ b/torchtitan/models/kimi_k3/dtensor_ops.py @@ -0,0 +1,77 @@ +# 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. + +"""Leaving DTensor land for the fla-core triton kernels. + +Both K3 attention kinds call kernels that do not dispatch through DTensor, so +both have to unwrap at the kernel call site. The two spellings differ only in +the gradient placement they hand back, and that difference is the whole point of +keeping them next to each other -- see the docstrings. +""" + +from torch.distributed.tensor import DTensor +from torch.distributed.tensor.placement_types import Partial, Replicate + + +__all__ = ["to_local_if_dtensor", "to_local_partial_grad"] + + +def to_local_if_dtensor(t): + """Strip DTensor wrapping for fla-core triton kernels. + + fla-core's chunk_kda / fused_kda_gate / ShortConvolution are Triton + kernels that don't dispatch through DTensor. Under TP, KDA's + self_attn is NoParallel-wrapped (params become DTensor(Replicate) + on tp_mesh) and incoming x is also DTensor at the parent's + boundary. KDA forward stashes the DTensor mesh+placements, strips + DTensor from x and from each weight at the kernel call site, runs + the kernels on plain tensors (each rank computes redundantly under + Replicate), and re-DTensors at the end so the parent NoParallel + output hook composes correctly. + + isinstance(t, DTensor) is the safe check that dynamo's fake-tensor + mode honors (``hasattr(t, "to_local")`` is unreliable: dynamo's + type tracking can elide attribute lookups on DTensor parameters). + + ``grad_placements`` is passed rather than left to default, which the + distributed rules ask for on every ``to_local``. It is the forward + placement, which is also what the default would pick -- and that is the + right answer only because every rank does the SAME work with the unwrapped + value. When ranks diverge the gradient of a replicated value is their sum, + and :func:`to_local_partial_grad` is the spelling for that; stating this + one explicitly is what makes the pair readable as a choice. + """ + if isinstance(t, DTensor): + return t.to_local(grad_placements=list(t.placements)) + return t + + +def to_local_partial_grad(t): + """``to_local`` for a value each rank then consumes DIFFERENTLY. + + ``to_local()`` defaults the incoming gradient's placement to the forward + placement. For a Replicate value that is correct only when every rank does the + SAME work with it -- which is exactly KDA's redundant kernels, and why + ``to_local_if_dtensor`` keeps the default. + + It is wrong when the ranks diverge. MLA's CP path expands the replicated + ``k_rot`` onto this rank's head subset, so each rank's gradient is one partial + contribution and the gradient of the replicated value is their sum: Partial, + not Replicate. Keeping the default drops that all-reduce silently, because the + placement still reads Replicate afterwards. + + Measured on ``kimi_k3_debugmodel_report_arch`` at tp2 x cp2: all four MLA + layers' ``kv_a_proj_with_mqa`` gradients differed across the tp pair by 1-6% + relative on every step, while tp2 alone was bit-identical -- the non-CP path + never leaves DTensor, so DTensor reduces it there. + """ + if not isinstance(t, DTensor): + return t + return t.to_local( + grad_placements=[ + Partial() if isinstance(p, Replicate) else p for p in t.placements + ] + ) diff --git a/torchtitan/models/kimi_k3/kcp.py b/torchtitan/models/kimi_k3/kcp.py new file mode 100644 index 0000000000..87a1dce039 --- /dev/null +++ b/torchtitan/models/kimi_k3/kcp.py @@ -0,0 +1,104 @@ +# 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. + +"""KCP: KDA Context Parallelism (report sec 5.1.2). + + Two cross-rank dependencies with different shapes. The recurrence needs each rank's + true incoming state, which does NOT decompose by summation -- the delta rule applies a + token-dependent transition, so a prefix scan over (cumulative transition, zero-started + state) fragments recovers it. The short convolutions need only the previous rank's + tail, one fixed-size exchange. + + See ``phase13_k3like_48b_posttrain/KCP_DESIGN.md``. + """ + +from __future__ import annotations + +import torch +import torch.distributed as dist +from torch.distributed.tensor import DTensor + + +def conv_with_halo( + conv, x_local: torch.Tensor, cp_context, activation: str | None = None +) -> torch.Tensor: + """Run a depthwise causal conv on a sequence-sharded input, exactly. + + Thin adapter over fla's ``causal_conv1d_cp``: unpack the depthwise weight + the way ``ShortConvolution.forward`` does and hand over the CP context, + which must have been built with ``conv1d_kernel_size`` set. + + ``activation`` defaults to reading ``conv.activation``, which fla's + ``ShortConvolution`` carries. A plain ``nn.Conv1d`` does not -- the upstream + K3 model applies its SiLU outside the conv -- so those call sites pass the + name explicitly rather than getting a second copy of this function. + + The weight and bias are unwrapped to local first. Under TP the KDA layers are + NoParallel, so these are DTensor(Replicate), and handing a DTensor to fla's + triton kernel does not raise anything legible -- it surfaces as + ``CUBLAS_STATUS_INTERNAL_ERROR`` or an illegal memory access from inside the + kernel. The Ulysses path unwraps them in its own ``conv_subset``; this one did + not, which is why KCP worked in every cell that had no TP and broke every cell + that had both. + """ + from einops import rearrange + from fla.modules.conv.cp.ops import causal_conv1d_cp + + weight = conv.weight + if isinstance(weight, DTensor): + weight = weight.to_local() + bias = conv.bias + if bias is not None and isinstance(bias, DTensor): + bias = bias.to_local() + + return causal_conv1d_cp( + x=x_local, + weight=rearrange(weight, "d 1 w -> d w"), + bias=bias, + activation=getattr(conv, "activation", None) + if activation is None + else activation, + cp_context=cp_context, + ) + + +def build_kcp_context( + seq_len_local: int, + group, + device, + conv1d_kernel_size: int | None = None, + cu_seqlens: "torch.Tensor | None" = None, +) -> object: + """fla CP context for one evenly-split sequence. + + ``chunk_kda`` needs the GLOBAL cu_seqlens of the packed sequence plus the + process group; ``build_cp_context`` derives each rank's slice from them. + ``conv1d_kernel_size`` is required by ``causal_conv1d_cp`` and otherwise + unused, so it is optional here. + + ``cu_seqlens`` defaults to ``[0, seq_len_local * world]``, i.e. ONE document + spanning the whole sequence. Pass real boundaries to describe a packed + (multi-document) sequence -- they must be GLOBAL, since that is what fla + slices per rank. + + Whether the default is right is a property of the caller, not of this + helper, and worth stating plainly: nothing in this repo hands KDA document + boundaries in ANY mode. Both non-CP call sites pass ``cu_seqlens=None`` to + ``chunk_kda``, so a packed SFT batch already carries the delta-rule state + across document boundaries with or without CP. The default here matches that + behaviour rather than introducing a hole of its own; fixing it means + threading the dataloader's boundaries through every KDA call site, not + changing this default. + """ + from fla.ops.cp.context import build_cp_context + + if cu_seqlens is None: + world = dist.get_world_size(group) + total = seq_len_local * world + cu_seqlens = torch.tensor([0, total], dtype=torch.int32, device=device) + return build_cp_context( + cu_seqlens, group=group, conv1d_kernel_size=conv1d_kernel_size + ) diff --git a/torchtitan/models/kimi_k3/kda.py b/torchtitan/models/kimi_k3/kda.py index f31a8eb9b1..b8b0bd3ef8 100644 --- a/torchtitan/models/kimi_k3/kda.py +++ b/torchtitan/models/kimi_k3/kda.py @@ -9,12 +9,20 @@ from dataclasses import dataclass import torch +import torch.distributed as dist import torch.nn.functional as F +from torch.distributed.tensor import DTensor from fla.ops.kda import chunk_kda from torch import nn from torchtitan.models.common import Conv1d, Linear from torchtitan.models.common.attention import AttentionMasksType +from torchtitan.models.kimi_k3.dtensor_ops import to_local_if_dtensor +from torchtitan.models.kimi_k3.sharding import ( + contract_for_mode, + cp_all_to_all_headseq, + ULYSSES, +) from torchtitan.protocols.module import Module # Shape suffixes: @@ -66,7 +74,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,6 +92,11 @@ 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 @@ -104,9 +121,17 @@ class Config(Module.Config): kernel: Module.Config output_norm: KimiRMSNormGated.Config output_proj: Linear.Config + cp_mode: str = "kcp" + + # Set by apply_cp_kimi_k3; None means the layer runs without CP. + _cp_group = None def __init__(self, config: Config): super().__init__() + self.cp_mode = config.cp_mode + # Validate against the declared contracts rather than restating the + # accepted spellings here. + contract_for_mode(self.cp_mode) self.num_heads = config.num_heads self.head_dim = config.head_dim self.conv_kernel_size = config.conv_kernel_size @@ -144,6 +169,14 @@ 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) + if self.cp_mode == "kcp" + else self._forward_ulysses(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 @@ -159,6 +192,78 @@ def forward( ) beta_TH = self.beta(x_TD).float() + # The kernel is fla triton and does not dispatch through DTensor. + # Under TP these arrive wrapped, and handing a DTensor to it produces + # an illegal memory access rather than anything legible, so the unwrap + # happens at the call site and the result is re-wrapped for the + # module's declared output layout. + out_THV = self.kernel( + to_local_if_dtensor(q_THK).unsqueeze(0), + to_local_if_dtensor(k_THK).unsqueeze(0), + to_local_if_dtensor(v_THV).unsqueeze(0), + to_local_if_dtensor(forget_THK).unsqueeze(0), + to_local_if_dtensor(beta_TH).unsqueeze(0), + to_local_if_dtensor(self.A_log), + to_local_if_dtensor(self.dt_bias), + ).squeeze(0) + if isinstance(q_THK, DTensor): + out_THV = DTensor.from_local( + out_THV, q_THK.device_mesh, q_THK.placements, run_check=False + ) + output_gate_THV = self.output_gate(x_TD).view( + num_tokens, self.num_heads, self.head_dim + ) + 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. + """ + from torchtitan.models.kimi_k3.kcp import build_kcp_context, conv_with_halo + + 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), @@ -167,9 +272,96 @@ def forward( beta_TH.unsqueeze(0), self.A_log, self.dt_bias, + cp_context=ctx, ).squeeze(0) output_gate_THV = self.output_gate(x_TD).view( - num_tokens, self.num_heads, self.head_dim + 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(num_tokens, -1)) + return self.output_proj(out_THV.reshape(t_loc, -1)) + + def _forward_ulysses(self, x_TD: torch.Tensor, cp_group) -> torch.Tensor: + """Ulysses CP forward: trade the sharded axis, sequence for heads. + + Projections run sequence-local, one fused all-to-all moves everything to + full-sequence head-subset layout, and the convolutions then run on the + full sequence for this rank's heads -- so no halo is needed here, unlike + KCP. The convolutions are depthwise, so restricting them to a contiguous + head subset is a contiguous channel slice of the weight and is exact. + + Shape suffixes beyond the file legend: L local sequence (T/cp), G this + rank's head count (H/cp), W the packed per-head channel width. + """ + # Head divisibility is checked at wiring time, against tp*cp rather + # than cp -- under TP the head axis is already split once. Repeating a + # cp-only version of that test here would reject configurations that + # run. + cp_size = dist.get_world_size(cp_group) + cp_rank = dist.get_rank(cp_group) + t_loc = x_TD.shape[0] + num_heads, head_dim = self.num_heads, self.head_dim + h_cp = num_heads // cp_size + h0 = cp_rank * h_cp + + def heads(t_LC: torch.Tensor) -> torch.Tensor: + return t_LC.view(t_loc, num_heads, head_dim) + + # 1) Sequence-local projections, pre-convolution. + q_LHK = heads(self.q_proj(x_TD)) + k_LHK = heads(self.k_proj(x_TD)) + v_LHV = heads(self.v_proj(x_TD)) + forget_LHK = heads(self.forget_b(self.forget_a(x_TD))) + gate_LHV = heads(self.output_gate(x_TD)) + beta_LH1 = self.beta(x_TD).unsqueeze(-1) + + # 2) One fused all-to-all instead of six. + packed_LHW = torch.cat( + [q_LHK, k_LHK, v_LHV, forget_LHK, gate_LHV, beta_LH1], 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_TGK, k_TGK, v_TGV, forget_TGK, gate_TGV, beta_TG1 = torch.split( + packed_TGW, [head_dim] * 5 + [1], dim=-1 + ) + t_full = t_loc * cp_size + + # 3) Causal convolution on the full sequence, channels sliced to this + # rank's heads. + def conv_subset(conv: Conv1d, x_TGK: torch.Tensor) -> torch.Tensor: + lo, hi = h0 * head_dim, (h0 + h_cp) * head_dim + w_C1W = to_local_if_dtensor(conv.weight)[lo:hi] + b_C = ( + to_local_if_dtensor(conv.bias)[lo:hi] if conv.bias is not None else None + ) + x_1CT = F.pad( + x_TGK.reshape(t_full, h_cp * head_dim).T.unsqueeze(0), + (self.conv_kernel_size - 1, 0), + ) + y_1CT = F.conv1d(x_1CT, w_C1W, b_C, groups=h_cp * head_dim) + return F.silu(y_1CT).squeeze(0).T.view(t_full, h_cp, head_dim) + + q_TGK = conv_subset(self.q_conv, q_TGK) + k_TGK = conv_subset(self.k_conv, k_TGK) + v_TGV = conv_subset(self.v_conv, v_TGV) + + # 4) The scan runs on this rank's heads over the full sequence, so + # A_log and dt_bias are sliced to the same subset. + out_TGV = self.kernel( + q_TGK.unsqueeze(0), + k_TGK.unsqueeze(0), + v_TGV.unsqueeze(0), + forget_TGK.unsqueeze(0), + beta_TG1.squeeze(-1).float().unsqueeze(0), + to_local_if_dtensor(self.A_log)[h0 : h0 + h_cp], + to_local_if_dtensor(self.dt_bias)[h0 : h0 + h_cp], + ).squeeze(0) + out_TGV = self.output_norm(out_TGV, gate_TGV) + + # 5) Back to sequence-sharded full-head layout. + out_src_dim, out_dst_dim = ULYSSES.out_dims() + out_LHV = cp_all_to_all_headseq( + out_TGV, cp_group, src_dim=out_src_dim, dst_dim=out_dst_dim + ) + return self.output_proj(out_LHV.reshape(t_loc, num_heads * head_dim)) diff --git a/torchtitan/models/kimi_k3/model.py b/torchtitan/models/kimi_k3/model.py index 25f473ff21..bc778b69ac 100644 --- a/torchtitan/models/kimi_k3/model.py +++ b/torchtitan/models/kimi_k3/model.py @@ -7,6 +7,7 @@ from dataclasses import dataclass, field import torch +import torch.distributed as dist from torch import nn from torchtitan.hf_datasets.multimodal.mm_datasets import MMSamplePackingConfig @@ -15,7 +16,9 @@ from torchtitan.models.common.attention import ( AttentionMasksType, BaseAttention, + create_attention_mask, FlexAttention, + get_causal_mask_mod, ) from torchtitan.models.common.decoder import Decoder from torchtitan.models.common.multimodal import ( @@ -28,6 +31,7 @@ from .kda import KimiDeltaAttention from .moe import KimiFeedForward, KimiLatentMoE +from .sharding import cp_all_to_all_headseq, ULYSSES from .vision_encoder import KimiK3VisionEncoder # Shape suffixes: @@ -62,6 +66,12 @@ 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 + _cp_mask = None + def __init__(self, config: Config): super().__init__() self.n_heads = config.n_heads @@ -91,9 +101,13 @@ def forward( 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 +117,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,21 +126,150 @@ 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 = self._ulysses_attention(q_THK, kv_THC, k_rope_TK, cp_group) + 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) + def _full_sequence_causal_mask(self, num_tokens: int, device): + """Causal mask for the sequence Ulysses reassembles. + + The mask the layer is handed has been sharded for context parallel by + ``cp_shard``, which cuts it the way ring attention wants: local queries + against global keys. Ulysses reassembles the whole sequence on every + rank instead, so it needs the whole causal mask. Rebuilding it is + correct here only because this model rejects sample packing, so the + sequence is one document and the mask carries no boundaries; a packed + sequence would need the global boundaries threaded down instead. + + Cached per (length, device) because the shape is constant across layers + and steps, and create_block_mask is compiled. + """ + # The mask the decoder builds at dp1 is causal AND packed-document + # (common/decoder._create_flex_attention_mask_for_document). This + # rebuild is causal only, which is equivalent exactly when the folded + # stream holds ONE document. Sample packing is already rejected in + # update_from_config, but a microbatch wider than the context window + # folds several documents into one stream as well, and then CP would + # let a sample attend to the previous one while dp1 would not -- + # silently, since every shape stays valid. Caught here rather than + # documented. + limit = getattr(self, "_cp_max_context_length", None) + if limit is not None and num_tokens > limit: + raise NotImplementedError( + f"context parallel folds {num_tokens} tokens into one stream " + f"but the context window is {limit}, so the " + "stream holds more than one document. The CP path rebuilds a " + "causal-only mask and cannot see document boundaries; use a " + "microbatch no wider than the context window." + ) + key = (num_tokens, device) + if self._cp_mask is None or self._cp_mask[0] != key: + mask = create_attention_mask( + get_causal_mask_mod(), + None, + None, + num_tokens, + num_tokens, + device=device, + ) + self._cp_mask = (key, mask) + return self._cp_mask[1] + + def _ulysses_attention( + self, + q_LHQ: torch.Tensor, + kv_LHC: torch.Tensor, + k_rope_LR: torch.Tensor, + cp_group, + ) -> torch.Tensor: + """Attention over the full sequence for this rank's head subset. + + One fused all-to-all trades the sharded axis, sequence for heads, then + the backend runs unchanged, then a second trades back. The gate and the + output projection stay sequence-local, so they are outside this. + + The rotary slice is deliberately not in the all-to-all. It is headless + -- one vector per token, shared by every head -- so it is all-gathered + along the sequence and expanded onto this rank's heads afterwards. + Packing the already-expanded key instead sends the same values once per + head and reassembles them against the wrong head subset, which shows up + as a forward that diverges from the same layer run without CP. + + Shape suffixes beyond the file legend: L local sequence (T/cp), G this + rank's head count (H/cp), W the packed per-head channel width, R the + rotary width. + """ + import torch.distributed.nn.functional as dist_nn + + from torchtitan.models.kimi_k3.dtensor_ops import to_local_partial_grad + + # Head divisibility is checked at wiring time; see apply_cp_kimi_k3. + cp_size = dist.get_world_size(cp_group) + t_loc = q_LHQ.shape[0] + t_full = t_loc * cp_size + # Local head count: q_LHQ already carries this rank's local heads, so + # the CP split is over that, not over the global n_heads. + 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, + [self.q_head_dim, self.qk_nope_head_dim, self.v_head_dim], + dim=-1, + ) + + # k_rope is produced by a module every rank of the head-splitting axis + # ran on the same input, so its gradient is the SUM across those ranks, + # i.e. Partial. A no-op when the input is a plain tensor, which is the + # CP-only case, so the reachable path here is unchanged. + k_rope_LR = to_local_partial_grad(k_rope_LR) + + # Differentiable all-gather: the backward is a reduce-scatter, which is + # what 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, self.qk_rope_head_dim).expand( + t_full, h_cp, self.qk_rope_head_dim + ), + ], + dim=-1, + ) + + out_TGV = self.inner_attention( + q_TGQ, + k_TGQ, + v_TGV, + attention_masks=self._full_sequence_causal_mask(t_full, q_TGQ.device), + scale=self.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 + ) + def _apply_attention_residual( prefix_sum_TD: torch.Tensor, @@ -266,6 +409,11 @@ class Config(Decoder.Config): output_res_norm: RMSNorm.Config output_res_proj: Linear.Config vision_encoder: KimiK3VisionEncoder.Config | None = None + # KDA runs on fla triton kernels, which do not dispatch through + # DTensor, so no ShardingConfig can drive its context parallel -- the + # layer implements both CP modes itself, and the preconditions that + # replaces the backend check with are enforced below. + cp_via_sharding_config: bool = False def update_from_config(self, *, config, **kwargs) -> None: dataset = config.dataloader.dataset @@ -273,6 +421,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 +459,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() diff --git a/torchtitan/models/kimi_k3/parallelize.py b/torchtitan/models/kimi_k3/parallelize.py index 8a7d604a02..7281e23bc0 100644 --- a/torchtitan/models/kimi_k3/parallelize.py +++ b/torchtitan/models/kimi_k3/parallelize.py @@ -18,7 +18,11 @@ 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 +from .sharding import contract_for_mode, ULYSSES 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, training.max_context_length) + if ac_config is not None: ac_policy = ac_config.build(dump_folder=dump_folder) ac_policy.apply(model) @@ -95,3 +101,107 @@ def parallelize_kimi_k3( ) return model + + +def _check_head_divisibility( + contract, num_heads: int, divisor: int, divisor_expr: str, kind: str, field: str +) -> None: + """Enforce the head split a contract asks for, if it asks for one.""" + if not contract.head_sharded: + return + if num_heads % divisor != 0: + raise ValueError( + f"{kind} {field}={num_heads} must be divisible by " + f"{divisor_expr}={divisor} for {contract.name} CP head sharding" + ) + + +def apply_cp_kimi_k3( + model: nn.Module, + parallel_dims: ParallelDims, + max_context_length: int | None = None, +) -> None: + """Wire context parallelism: KCP on the KDA layers, Ulysses on the MLA layers. + + Both at once, on disjoint layer kinds. KCP decomposes the delta-rule + recurrence and says nothing about softmax attention, so it does not replace + Ulysses; ``cp_mode="ulysses"`` runs the KDA layers the second way and is + kept as an A/B. + + Imperative rather than declared because KDA's kernels are fla triton and + never see a DTensor; see ``cp_via_sharding_config`` on the model config for + why the declarative path cannot serve them. + """ + cp_group = parallel_dims.get_mesh("cp").get_group() + cp_degree, tp_degree = parallel_dims.cp, parallel_dims.tp + model._cp_group = cp_group + # The CP mask rebuild is causal-only; hand the layers the context window + # so they can reject a folded stream that holds several documents. + # The window comes from the training config: K3's MLA is nope, so the + # decoder's RoPE-derived max_context_length raises rather than returning + # one. A folded stream longer than this holds more than one document, which + # the causal-only CP mask cannot represent. + max_ctx = max_context_length + + num_mla = 0 + kda_modules = [] + for module in model.modules(): + if isinstance(module, KimiMLAAttention): + module._cp_max_context_length = max_ctx + # The head axis may already be split by another parallelism, so + # Ulysses splits what is left: heads must divide by tp*cp. + _check_head_divisibility( + ULYSSES, + module.n_heads, + tp_degree * cp_degree, + "tp*cp", + "MLA", + "n_heads", + ) + module._cp_group = cp_group + num_mla += 1 + elif isinstance(module, KimiDeltaAttention): + kda_modules.append(module) + + modes = {m.cp_mode for m in kda_modules} + for mode in modes: + contract = contract_for_mode(mode) + for module in kda_modules: + if module.cp_mode == mode: + _check_head_divisibility( + contract, + module.num_heads, + tp_degree * cp_degree, + "tp*cp", + "KDA", + "num_heads", + ) + if "kcp" in modes: + # Checked here rather than at the first forward: the message is + # actionable at wiring time and the failure is otherwise an ImportError + # from inside a layer. + try: + from fla.modules.conv.cp.ops import causal_conv1d_cp # noqa: F401 + from fla.ops.cp.context import build_cp_context # noqa: F401 + except ImportError as err: + raise ValueError( + "cp_mode='kcp' needs fla-core's CP ops " + "(fla.ops.cp.context.build_cp_context and " + "fla.modules.conv.cp.ops.causal_conv1d_cp), which ship in " + f"fla-core >= 0.5.1; import failed with: {err}. Install a " + "newer fla-core or use cp_mode='ulysses'." + ) 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), modes=%s.", + num_mla, + len(kda_modules), + sorted(modes) or ["-"], + ) diff --git a/torchtitan/models/kimi_k3/sharding.py b/torchtitan/models/kimi_k3/sharding.py new file mode 100644 index 0000000000..46c50cf34d --- /dev/null +++ b/torchtitan/models/kimi_k3/sharding.py @@ -0,0 +1,172 @@ +# 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. + +See CP_DECLARATIVE.md in the logbook for why KCP is an identity pair. +""" + +from dataclasses import dataclass + +import spmd_types as spmd + +import torch +import torch.distributed as dist + +from torchtitan.distributed.parallel_dims import MeshAxisName, SpmdLayout + + +__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: ``kda_cp_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, which is a sequential dependency, not a +# redistribution -- no placement pair describes it, so it stays inside the op and +# the contract is an identity. Declared anyway to keep one shape for both modes. +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"kda_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. + """ + import torch.distributed.nn.functional as dist_nn + + 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() diff --git a/torchtitan/models/kimi_k3/tests/__init__.py b/torchtitan/models/kimi_k3/tests/__init__.py new file mode 100644 index 0000000000..2e41cd717f --- /dev/null +++ b/torchtitan/models/kimi_k3/tests/__init__.py @@ -0,0 +1,5 @@ +# 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. diff --git a/torchtitan/models/kimi_k3/tests/test_cp_contracts.py b/torchtitan/models/kimi_k3/tests/test_cp_contracts.py new file mode 100644 index 0000000000..cae2c13d9a --- /dev/null +++ b/torchtitan/models/kimi_k3/tests/test_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() From 383770b9bbb240b7383e7588b31ade50a68161c2 Mon Sep 17 00:00:00 2001 From: QIU023 Date: Tue, 25 Aug 2026 00:08:19 +0000 Subject: [PATCH 02/14] kimi_k3: move the CP unit test to tests/unit_tests and add a cp2 integration cell The contract test sat in torchtitan/models/kimi_k3/tests/, the only in-model tests directory in the tree; the convention is tests/unit_tests/. Moved and renamed to carry the model name. Adds a cp2 recipe and the integration cell that runs it, so context parallel is covered by CI rather than only by a local matrix. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01WBy1d9YVu44nYCVqykRqL1 --- tests/integration_tests/features.py | 6 ++++++ .../unit_tests/test_kimi_k3_cp_contracts.py | 0 torchtitan/models/kimi_k3/tests/__init__.py | 5 ----- torchtitan_recipes/tests/features.py | 13 +++++++++++++ 4 files changed, 19 insertions(+), 5 deletions(-) rename torchtitan/models/kimi_k3/tests/test_cp_contracts.py => tests/unit_tests/test_kimi_k3_cp_contracts.py (100%) delete mode 100644 torchtitan/models/kimi_k3/tests/__init__.py diff --git a/tests/integration_tests/features.py b/tests/integration_tests/features.py index 481e520f9b..a5efe6a97b 100755 --- a/tests/integration_tests/features.py +++ b/tests/integration_tests/features.py @@ -303,4 +303,10 @@ def build_features_test_list() -> list[OverrideDefinitions]: timeout=30, use_real_pg=True, ), + OverrideDefinitions( + configs=[recipes.kimi_k3_debugmodel_text_cp2], + test_descr="Kimi K3 text decoder, context parallel cp2", + test_name="kimi_k3_text_cp2", + ngpu=2, + ), ] diff --git a/torchtitan/models/kimi_k3/tests/test_cp_contracts.py b/tests/unit_tests/test_kimi_k3_cp_contracts.py similarity index 100% rename from torchtitan/models/kimi_k3/tests/test_cp_contracts.py rename to tests/unit_tests/test_kimi_k3_cp_contracts.py diff --git a/torchtitan/models/kimi_k3/tests/__init__.py b/torchtitan/models/kimi_k3/tests/__init__.py deleted file mode 100644 index 2e41cd717f..0000000000 --- a/torchtitan/models/kimi_k3/tests/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -# 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. diff --git a/torchtitan_recipes/tests/features.py b/torchtitan_recipes/tests/features.py index 8ecfc5140f..c7e8f0e890 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_text_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_text + + config = kimi_k3_debugmodel_text() + config.parallelism.context_parallel_degree = 2 + config.parallelism.context_parallel_load_balancer = None + return config From 6ed6c606f674bf273308fb652464aa7e3361e82f Mon Sep 17 00:00:00 2001 From: QIU023 Date: Tue, 25 Aug 2026 01:34:37 +0000 Subject: [PATCH 03/14] kimi_k3: put the CP tests where CI runs them, after the test reorg The CPU unit-test workflow collects tests/unit_tests/cpu, not the top level, so the contract test would have sat there un-run; move it. The cp2 cell declares use_real_pg because the thing it covers is the all-to-all and the KCP state pass; under Fake PG the collectives are no-ops and the cell would pass without exercising either. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01WBy1d9YVu44nYCVqykRqL1 --- tests/integration_tests/features.py | 1 + tests/unit_tests/{ => cpu}/test_kimi_k3_cp_contracts.py | 0 2 files changed, 1 insertion(+) rename tests/unit_tests/{ => cpu}/test_kimi_k3_cp_contracts.py (100%) diff --git a/tests/integration_tests/features.py b/tests/integration_tests/features.py index a5efe6a97b..6c86ba0a6b 100755 --- a/tests/integration_tests/features.py +++ b/tests/integration_tests/features.py @@ -308,5 +308,6 @@ def build_features_test_list() -> list[OverrideDefinitions]: test_descr="Kimi K3 text decoder, context parallel cp2", test_name="kimi_k3_text_cp2", ngpu=2, + use_real_pg=True, ), ] diff --git a/tests/unit_tests/test_kimi_k3_cp_contracts.py b/tests/unit_tests/cpu/test_kimi_k3_cp_contracts.py similarity index 100% rename from tests/unit_tests/test_kimi_k3_cp_contracts.py rename to tests/unit_tests/cpu/test_kimi_k3_cp_contracts.py From 55f82b30aacd15e0f8ff42aa83f536adf03b7ada Mon Sep 17 00:00:00 2001 From: QIU023 Date: Tue, 25 Aug 2026 01:39:46 -0400 Subject: [PATCH 04/14] kimi_k3: CP without tensor-parallel content, and no new core config field Review round on the CP change, four structural points. The core accommodation becomes a protected method instead of a config field. Decoder.Config._validate_cp_backend holds the spmd_types check and a model whose CP is not ShardingConfig-driven overrides it -- no field appears on every other model's config surface, and when this model's kernels move to a torch-native KDA and its CP goes declarative, the override is deleted rather than a core field deprecated. Tensor parallel is not on this PR, so everything that existed only for the TP interaction goes: dtensor_ops.py whole (both helpers guard DTensor inputs that cannot occur without TP), the kernel-call unwraps and re-wrap in kda, the k_rope Partial-gradient boundary in the MLA CP path, and the DTensor unwraps in conv_with_halo. KDA's alternate Ulysses mode goes with its cp_mode enum: under CP the KDA layers run KCP, full stop, which is also all the PR body ever described. MLA's Ulysses stays, as the method on the attention class that orchestrates its projections; the reusable transport already lives in sharding.py. kcp.py's two helpers move into kda.py next to their only caller, with the debugging narratives cut to what they established. The head-divisibility check inlines into apply_cp_kimi_k3, which shrinks to one loop. --- torchtitan/models/common/decoder.py | 27 ++-- torchtitan/models/kimi_k3/dtensor_ops.py | 77 ---------- torchtitan/models/kimi_k3/kcp.py | 104 -------------- torchtitan/models/kimi_k3/kda.py | 170 +++++++---------------- torchtitan/models/kimi_k3/model.py | 19 +-- torchtitan/models/kimi_k3/parallelize.py | 82 +++-------- 6 files changed, 82 insertions(+), 397 deletions(-) delete mode 100644 torchtitan/models/kimi_k3/dtensor_ops.py delete mode 100644 torchtitan/models/kimi_k3/kcp.py diff --git a/torchtitan/models/common/decoder.py b/torchtitan/models/common/decoder.py index fd28efb34c..d4889c6640 100644 --- a/torchtitan/models/common/decoder.py +++ b/torchtitan/models/common/decoder.py @@ -81,22 +81,6 @@ class Config(BaseModel.Config): # that support it set this True in their config factories; the tying # itself is handled by ``Decoder.__init__`` / ``Decoder.init_states``. enable_weight_tying: bool = False - # Whether this model's context parallel is driven by ShardingConfig. - # ``validate_cp_backend`` is documented as being for "the models that - # declare CP in ShardingConfig", but the check below runs for every - # decoder, so a model whose CP cannot be declarative has no way to say - # so. Kernels that do not dispatch through DTensor force the case: - # nothing in a ShardingConfig can reach them, and the model implements - # CP itself. Such a model sets this False and takes on its own CP - # preconditions. - # - # Deriving it instead -- asking whether any placement in the config - # names the CP axis -- was tried and is wrong: every model here sets - # its sharding AFTER delegating to this method, llama3 included, so - # at this point no model declares anything and the check would be - # silently skipped for all of them. - cp_via_sharding_config: bool = True - @property def first_attention(self) -> BaseAttention.Config | None: """Attention config of the first layer that has one, else None. @@ -151,6 +135,13 @@ def max_context_length(self) -> int: ) return rope_cfg.max_context_length + + def _validate_cp_backend(self, parallelism) -> None: + """ShardingConfig-driven CP requires the spmd_types backend. A model + whose CP is not ShardingConfig-driven overrides this and takes on + its own preconditions.""" + validate_cp_backend(parallelism) + def update_from_config( self, *, @@ -186,9 +177,7 @@ def update_from_config( ) if parallelism.context_parallel_degree > 1: - # ShardingConfig-based CP requires the spmd_types backend. - if self.cp_via_sharding_config: - 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/kimi_k3/dtensor_ops.py b/torchtitan/models/kimi_k3/dtensor_ops.py deleted file mode 100644 index f2ca5e01bd..0000000000 --- a/torchtitan/models/kimi_k3/dtensor_ops.py +++ /dev/null @@ -1,77 +0,0 @@ -# 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. - -"""Leaving DTensor land for the fla-core triton kernels. - -Both K3 attention kinds call kernels that do not dispatch through DTensor, so -both have to unwrap at the kernel call site. The two spellings differ only in -the gradient placement they hand back, and that difference is the whole point of -keeping them next to each other -- see the docstrings. -""" - -from torch.distributed.tensor import DTensor -from torch.distributed.tensor.placement_types import Partial, Replicate - - -__all__ = ["to_local_if_dtensor", "to_local_partial_grad"] - - -def to_local_if_dtensor(t): - """Strip DTensor wrapping for fla-core triton kernels. - - fla-core's chunk_kda / fused_kda_gate / ShortConvolution are Triton - kernels that don't dispatch through DTensor. Under TP, KDA's - self_attn is NoParallel-wrapped (params become DTensor(Replicate) - on tp_mesh) and incoming x is also DTensor at the parent's - boundary. KDA forward stashes the DTensor mesh+placements, strips - DTensor from x and from each weight at the kernel call site, runs - the kernels on plain tensors (each rank computes redundantly under - Replicate), and re-DTensors at the end so the parent NoParallel - output hook composes correctly. - - isinstance(t, DTensor) is the safe check that dynamo's fake-tensor - mode honors (``hasattr(t, "to_local")`` is unreliable: dynamo's - type tracking can elide attribute lookups on DTensor parameters). - - ``grad_placements`` is passed rather than left to default, which the - distributed rules ask for on every ``to_local``. It is the forward - placement, which is also what the default would pick -- and that is the - right answer only because every rank does the SAME work with the unwrapped - value. When ranks diverge the gradient of a replicated value is their sum, - and :func:`to_local_partial_grad` is the spelling for that; stating this - one explicitly is what makes the pair readable as a choice. - """ - if isinstance(t, DTensor): - return t.to_local(grad_placements=list(t.placements)) - return t - - -def to_local_partial_grad(t): - """``to_local`` for a value each rank then consumes DIFFERENTLY. - - ``to_local()`` defaults the incoming gradient's placement to the forward - placement. For a Replicate value that is correct only when every rank does the - SAME work with it -- which is exactly KDA's redundant kernels, and why - ``to_local_if_dtensor`` keeps the default. - - It is wrong when the ranks diverge. MLA's CP path expands the replicated - ``k_rot`` onto this rank's head subset, so each rank's gradient is one partial - contribution and the gradient of the replicated value is their sum: Partial, - not Replicate. Keeping the default drops that all-reduce silently, because the - placement still reads Replicate afterwards. - - Measured on ``kimi_k3_debugmodel_report_arch`` at tp2 x cp2: all four MLA - layers' ``kv_a_proj_with_mqa`` gradients differed across the tp pair by 1-6% - relative on every step, while tp2 alone was bit-identical -- the non-CP path - never leaves DTensor, so DTensor reduces it there. - """ - if not isinstance(t, DTensor): - return t - return t.to_local( - grad_placements=[ - Partial() if isinstance(p, Replicate) else p for p in t.placements - ] - ) diff --git a/torchtitan/models/kimi_k3/kcp.py b/torchtitan/models/kimi_k3/kcp.py deleted file mode 100644 index 87a1dce039..0000000000 --- a/torchtitan/models/kimi_k3/kcp.py +++ /dev/null @@ -1,104 +0,0 @@ -# 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. - -"""KCP: KDA Context Parallelism (report sec 5.1.2). - - Two cross-rank dependencies with different shapes. The recurrence needs each rank's - true incoming state, which does NOT decompose by summation -- the delta rule applies a - token-dependent transition, so a prefix scan over (cumulative transition, zero-started - state) fragments recovers it. The short convolutions need only the previous rank's - tail, one fixed-size exchange. - - See ``phase13_k3like_48b_posttrain/KCP_DESIGN.md``. - """ - -from __future__ import annotations - -import torch -import torch.distributed as dist -from torch.distributed.tensor import DTensor - - -def conv_with_halo( - conv, x_local: torch.Tensor, cp_context, activation: str | None = None -) -> torch.Tensor: - """Run a depthwise causal conv on a sequence-sharded input, exactly. - - Thin adapter over fla's ``causal_conv1d_cp``: unpack the depthwise weight - the way ``ShortConvolution.forward`` does and hand over the CP context, - which must have been built with ``conv1d_kernel_size`` set. - - ``activation`` defaults to reading ``conv.activation``, which fla's - ``ShortConvolution`` carries. A plain ``nn.Conv1d`` does not -- the upstream - K3 model applies its SiLU outside the conv -- so those call sites pass the - name explicitly rather than getting a second copy of this function. - - The weight and bias are unwrapped to local first. Under TP the KDA layers are - NoParallel, so these are DTensor(Replicate), and handing a DTensor to fla's - triton kernel does not raise anything legible -- it surfaces as - ``CUBLAS_STATUS_INTERNAL_ERROR`` or an illegal memory access from inside the - kernel. The Ulysses path unwraps them in its own ``conv_subset``; this one did - not, which is why KCP worked in every cell that had no TP and broke every cell - that had both. - """ - from einops import rearrange - from fla.modules.conv.cp.ops import causal_conv1d_cp - - weight = conv.weight - if isinstance(weight, DTensor): - weight = weight.to_local() - bias = conv.bias - if bias is not None and isinstance(bias, DTensor): - bias = bias.to_local() - - return causal_conv1d_cp( - x=x_local, - weight=rearrange(weight, "d 1 w -> d w"), - bias=bias, - activation=getattr(conv, "activation", None) - if activation is None - else activation, - cp_context=cp_context, - ) - - -def build_kcp_context( - seq_len_local: int, - group, - device, - conv1d_kernel_size: int | None = None, - cu_seqlens: "torch.Tensor | None" = None, -) -> object: - """fla CP context for one evenly-split sequence. - - ``chunk_kda`` needs the GLOBAL cu_seqlens of the packed sequence plus the - process group; ``build_cp_context`` derives each rank's slice from them. - ``conv1d_kernel_size`` is required by ``causal_conv1d_cp`` and otherwise - unused, so it is optional here. - - ``cu_seqlens`` defaults to ``[0, seq_len_local * world]``, i.e. ONE document - spanning the whole sequence. Pass real boundaries to describe a packed - (multi-document) sequence -- they must be GLOBAL, since that is what fla - slices per rank. - - Whether the default is right is a property of the caller, not of this - helper, and worth stating plainly: nothing in this repo hands KDA document - boundaries in ANY mode. Both non-CP call sites pass ``cu_seqlens=None`` to - ``chunk_kda``, so a packed SFT batch already carries the delta-rule state - across document boundaries with or without CP. The default here matches that - behaviour rather than introducing a hole of its own; fixing it means - threading the dataloader's boundaries through every KDA call site, not - changing this default. - """ - from fla.ops.cp.context import build_cp_context - - if cu_seqlens is None: - world = dist.get_world_size(group) - total = seq_len_local * world - cu_seqlens = torch.tensor([0, total], dtype=torch.int32, device=device) - return build_cp_context( - cu_seqlens, group=group, conv1d_kernel_size=conv1d_kernel_size - ) diff --git a/torchtitan/models/kimi_k3/kda.py b/torchtitan/models/kimi_k3/kda.py index b8b0bd3ef8..ab4789a31a 100644 --- a/torchtitan/models/kimi_k3/kda.py +++ b/torchtitan/models/kimi_k3/kda.py @@ -11,18 +11,11 @@ import torch import torch.distributed as dist import torch.nn.functional as F -from torch.distributed.tensor import DTensor from fla.ops.kda import chunk_kda from torch import nn from torchtitan.models.common import Conv1d, Linear from torchtitan.models.common.attention import AttentionMasksType -from torchtitan.models.kimi_k3.dtensor_ops import to_local_if_dtensor -from torchtitan.models.kimi_k3.sharding import ( - contract_for_mode, - cp_all_to_all_headseq, - ULYSSES, -) from torchtitan.protocols.module import Module # Shape suffixes: @@ -101,6 +94,47 @@ def forward( 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: fla's + ``causal_conv1d_cp`` 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 einops import rearrange + from fla.modules.conv.cp.ops import causal_conv1d_cp + + return causal_conv1d_cp( + x=x_local, + weight=rearrange(conv.weight, "d 1 w -> d w"), + bias=conv.bias, + activation=getattr(conv, "activation", None) + if activation is None + else activation, + cp_context=cp_context, + ) + + +def build_kcp_context( + seq_len_local: int, + group, + device, + conv1d_kernel_size: int | None = None, + cu_seqlens=None, +): + """fla 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 fla.ops.cp.context import build_cp_context + + if cu_seqlens is None: + world = dist.get_world_size(group) + total = seq_len_local * world + cu_seqlens = torch.tensor([0, total], dtype=torch.int32, device=device) + return build_cp_context( + cu_seqlens, group=group, conv1d_kernel_size=conv1d_kernel_size + ) + + class KimiDeltaAttention(Module): @dataclass(kw_only=True, slots=True) class Config(Module.Config): @@ -121,17 +155,12 @@ class Config(Module.Config): kernel: Module.Config output_norm: KimiRMSNormGated.Config output_proj: Linear.Config - cp_mode: str = "kcp" # Set by apply_cp_kimi_k3; None means the layer runs without CP. _cp_group = None def __init__(self, config: Config): super().__init__() - self.cp_mode = config.cp_mode - # Validate against the declared contracts rather than restating the - # accepted spellings here. - contract_for_mode(self.cp_mode) self.num_heads = config.num_heads self.head_dim = config.head_dim self.conv_kernel_size = config.conv_kernel_size @@ -171,11 +200,7 @@ def forward( 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) - if self.cp_mode == "kcp" - else self._forward_ulysses(x_TD, cp_group) - ) + 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( @@ -192,24 +217,15 @@ def forward( ) beta_TH = self.beta(x_TD).float() - # The kernel is fla triton and does not dispatch through DTensor. - # Under TP these arrive wrapped, and handing a DTensor to it produces - # an illegal memory access rather than anything legible, so the unwrap - # happens at the call site and the result is re-wrapped for the - # module's declared output layout. out_THV = self.kernel( - to_local_if_dtensor(q_THK).unsqueeze(0), - to_local_if_dtensor(k_THK).unsqueeze(0), - to_local_if_dtensor(v_THV).unsqueeze(0), - to_local_if_dtensor(forget_THK).unsqueeze(0), - to_local_if_dtensor(beta_TH).unsqueeze(0), - to_local_if_dtensor(self.A_log), - to_local_if_dtensor(self.dt_bias), + 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, ).squeeze(0) - if isinstance(q_THK, DTensor): - out_THV = DTensor.from_local( - out_THV, q_THK.device_mesh, q_THK.placements, run_check=False - ) output_gate_THV = self.output_gate(x_TD).view( num_tokens, self.num_heads, self.head_dim ) @@ -229,8 +245,6 @@ def _forward_kcp(self, x_TD: torch.Tensor, cp_group) -> torch.Tensor: 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. """ - from torchtitan.models.kimi_k3.kcp import build_kcp_context, conv_with_halo - 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. @@ -279,89 +293,3 @@ def conv(proj, conv_module) -> torch.Tensor: ) out_THV = self.output_norm(out_THV, output_gate_THV) return self.output_proj(out_THV.reshape(t_loc, -1)) - - def _forward_ulysses(self, x_TD: torch.Tensor, cp_group) -> torch.Tensor: - """Ulysses CP forward: trade the sharded axis, sequence for heads. - - Projections run sequence-local, one fused all-to-all moves everything to - full-sequence head-subset layout, and the convolutions then run on the - full sequence for this rank's heads -- so no halo is needed here, unlike - KCP. The convolutions are depthwise, so restricting them to a contiguous - head subset is a contiguous channel slice of the weight and is exact. - - Shape suffixes beyond the file legend: L local sequence (T/cp), G this - rank's head count (H/cp), W the packed per-head channel width. - """ - # Head divisibility is checked at wiring time, against tp*cp rather - # than cp -- under TP the head axis is already split once. Repeating a - # cp-only version of that test here would reject configurations that - # run. - cp_size = dist.get_world_size(cp_group) - cp_rank = dist.get_rank(cp_group) - t_loc = x_TD.shape[0] - num_heads, head_dim = self.num_heads, self.head_dim - h_cp = num_heads // cp_size - h0 = cp_rank * h_cp - - def heads(t_LC: torch.Tensor) -> torch.Tensor: - return t_LC.view(t_loc, num_heads, head_dim) - - # 1) Sequence-local projections, pre-convolution. - q_LHK = heads(self.q_proj(x_TD)) - k_LHK = heads(self.k_proj(x_TD)) - v_LHV = heads(self.v_proj(x_TD)) - forget_LHK = heads(self.forget_b(self.forget_a(x_TD))) - gate_LHV = heads(self.output_gate(x_TD)) - beta_LH1 = self.beta(x_TD).unsqueeze(-1) - - # 2) One fused all-to-all instead of six. - packed_LHW = torch.cat( - [q_LHK, k_LHK, v_LHV, forget_LHK, gate_LHV, beta_LH1], 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_TGK, k_TGK, v_TGV, forget_TGK, gate_TGV, beta_TG1 = torch.split( - packed_TGW, [head_dim] * 5 + [1], dim=-1 - ) - t_full = t_loc * cp_size - - # 3) Causal convolution on the full sequence, channels sliced to this - # rank's heads. - def conv_subset(conv: Conv1d, x_TGK: torch.Tensor) -> torch.Tensor: - lo, hi = h0 * head_dim, (h0 + h_cp) * head_dim - w_C1W = to_local_if_dtensor(conv.weight)[lo:hi] - b_C = ( - to_local_if_dtensor(conv.bias)[lo:hi] if conv.bias is not None else None - ) - x_1CT = F.pad( - x_TGK.reshape(t_full, h_cp * head_dim).T.unsqueeze(0), - (self.conv_kernel_size - 1, 0), - ) - y_1CT = F.conv1d(x_1CT, w_C1W, b_C, groups=h_cp * head_dim) - return F.silu(y_1CT).squeeze(0).T.view(t_full, h_cp, head_dim) - - q_TGK = conv_subset(self.q_conv, q_TGK) - k_TGK = conv_subset(self.k_conv, k_TGK) - v_TGV = conv_subset(self.v_conv, v_TGV) - - # 4) The scan runs on this rank's heads over the full sequence, so - # A_log and dt_bias are sliced to the same subset. - out_TGV = self.kernel( - q_TGK.unsqueeze(0), - k_TGK.unsqueeze(0), - v_TGV.unsqueeze(0), - forget_TGK.unsqueeze(0), - beta_TG1.squeeze(-1).float().unsqueeze(0), - to_local_if_dtensor(self.A_log)[h0 : h0 + h_cp], - to_local_if_dtensor(self.dt_bias)[h0 : h0 + h_cp], - ).squeeze(0) - out_TGV = self.output_norm(out_TGV, gate_TGV) - - # 5) Back to sequence-sharded full-head layout. - out_src_dim, out_dst_dim = ULYSSES.out_dims() - out_LHV = cp_all_to_all_headseq( - out_TGV, cp_group, src_dim=out_src_dim, dst_dim=out_dst_dim - ) - return self.output_proj(out_LHV.reshape(t_loc, num_heads * head_dim)) diff --git a/torchtitan/models/kimi_k3/model.py b/torchtitan/models/kimi_k3/model.py index bc778b69ac..2678357e5d 100644 --- a/torchtitan/models/kimi_k3/model.py +++ b/torchtitan/models/kimi_k3/model.py @@ -216,8 +216,6 @@ def _ulysses_attention( """ import torch.distributed.nn.functional as dist_nn - from torchtitan.models.kimi_k3.dtensor_ops import to_local_partial_grad - # Head divisibility is checked at wiring time; see apply_cp_kimi_k3. cp_size = dist.get_world_size(cp_group) t_loc = q_LHQ.shape[0] @@ -237,12 +235,6 @@ def _ulysses_attention( dim=-1, ) - # k_rope is produced by a module every rank of the head-splitting axis - # ran on the same input, so its gradient is the SUM across those ranks, - # i.e. Partial. A no-op when the input is a plain tensor, which is the - # CP-only case, so the reachable path here is unchanged. - k_rope_LR = to_local_partial_grad(k_rope_LR) - # Differentiable all-gather: the backward is a reduce-scatter, which is # what a value every rank consumed needs. k_rope_TR = torch.cat( @@ -409,11 +401,12 @@ class Config(Decoder.Config): output_res_norm: RMSNorm.Config output_res_proj: Linear.Config vision_encoder: KimiK3VisionEncoder.Config | None = None - # KDA runs on fla triton kernels, which do not dispatch through - # DTensor, so no ShardingConfig can drive its context parallel -- the - # layer implements both CP modes itself, and the preconditions that - # replaces the backend check with are enforced below. - cp_via_sharding_config: bool = False + + 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 diff --git a/torchtitan/models/kimi_k3/parallelize.py b/torchtitan/models/kimi_k3/parallelize.py index 7281e23bc0..613823c390 100644 --- a/torchtitan/models/kimi_k3/parallelize.py +++ b/torchtitan/models/kimi_k3/parallelize.py @@ -22,7 +22,6 @@ from .kda import KimiDeltaAttention from .model import KimiK3Model, KimiMLAAttention -from .sharding import contract_for_mode, ULYSSES def parallelize_kimi_k3( @@ -103,19 +102,6 @@ def parallelize_kimi_k3( return model -def _check_head_divisibility( - contract, num_heads: int, divisor: int, divisor_expr: str, kind: str, field: str -) -> None: - """Enforce the head split a contract asks for, if it asks for one.""" - if not contract.head_sharded: - return - if num_heads % divisor != 0: - raise ValueError( - f"{kind} {field}={num_heads} must be divisible by " - f"{divisor_expr}={divisor} for {contract.name} CP head sharding" - ) - - def apply_cp_kimi_k3( model: nn.Module, parallel_dims: ParallelDims, @@ -123,73 +109,44 @@ def apply_cp_kimi_k3( ) -> None: """Wire context parallelism: KCP on the KDA layers, Ulysses on the MLA layers. - Both at once, on disjoint layer kinds. KCP decomposes the delta-rule - recurrence and says nothing about softmax attention, so it does not replace - Ulysses; ``cp_mode="ulysses"`` runs the KDA layers the second way and is - kept as an A/B. - - Imperative rather than declared because KDA's kernels are fla triton and - never see a DTensor; see ``cp_via_sharding_config`` on the model config for - why the declarative path cannot serve them. + 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, tp_degree = parallel_dims.cp, parallel_dims.tp + cp_degree = parallel_dims.cp model._cp_group = cp_group - # The CP mask rebuild is causal-only; hand the layers the context window - # so they can reject a folded stream that holds several documents. - # The window comes from the training config: K3's MLA is nope, so the - # decoder's RoPE-derived max_context_length raises rather than returning - # one. A folded stream longer than this holds more than one document, which - # the causal-only CP mask cannot represent. - max_ctx = max_context_length num_mla = 0 kda_modules = [] for module in model.modules(): if isinstance(module, KimiMLAAttention): - module._cp_max_context_length = max_ctx - # The head axis may already be split by another parallelism, so - # Ulysses splits what is left: heads must divide by tp*cp. - _check_head_divisibility( - ULYSSES, - module.n_heads, - tp_degree * cp_degree, - "tp*cp", - "MLA", - "n_heads", - ) + # The CP mask rebuild is causal-only; the layer uses the context + # window to reject a folded stream holding several documents. + module._cp_max_context_length = max_context_length + 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) - modes = {m.cp_mode for m in kda_modules} - for mode in modes: - contract = contract_for_mode(mode) - for module in kda_modules: - if module.cp_mode == mode: - _check_head_divisibility( - contract, - module.num_heads, - tp_degree * cp_degree, - "tp*cp", - "KDA", - "num_heads", - ) - if "kcp" in modes: - # Checked here rather than at the first forward: the message is - # actionable at wiring time and the failure is otherwise an ImportError - # from inside a layer. + 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 fla.modules.conv.cp.ops import causal_conv1d_cp # noqa: F401 from fla.ops.cp.context import build_cp_context # noqa: F401 except ImportError as err: raise ValueError( - "cp_mode='kcp' needs fla-core's CP ops " + "KDA context parallelism needs fla-core's CP ops " "(fla.ops.cp.context.build_cp_context and " "fla.modules.conv.cp.ops.causal_conv1d_cp), which ship in " - f"fla-core >= 0.5.1; import failed with: {err}. Install a " - "newer fla-core or use cp_mode='ulysses'." + f"fla-core >= 0.5.1; import failed with: {err}." ) from err for module in kda_modules: @@ -200,8 +157,7 @@ def apply_cp_kimi_k3( "wire it onto." ) logger.info( - "Applied context parallel to %d MLA and %d KDA layer(s), modes=%s.", + "Applied context parallel to %d MLA and %d KDA layer(s).", num_mla, len(kda_modules), - sorted(modes) or ["-"], ) From 48afbd1d820c8e092b534dc60da08ed2762a59ec Mon Sep 17 00:00:00 2001 From: QIU023 Date: Tue, 25 Aug 2026 07:38:37 +0000 Subject: [PATCH 05/14] kimi_k3: MLA's Ulysses body moves to sharding.py, and the CP-backend override carries its import The attention keeps the branch and the call; the exchange itself -- two all-to-alls with the attention backend between them, and the rotary slice gathered outside them -- lives next to the contract that describes it. The causal mask rebuild moves with it. _validate_cp_backend called validate_cp_backend as a bare name, but that name was only bound by a function-local import in update_from_config, so a method body resolving to module globals would not find it: any run with context_parallel_degree > 1 raised NameError before reaching the check. The method now carries the import, and the dead local one is gone. Also: the module docstring no longer claims a contract is resolved per module, the error text stops naming a config field that no longer exists, and the text flavor derives from the debug flavor. Configuration-identical, 4810 and 4678 fields over the union of both key sets. --- torchtitan/models/common/decoder.py | 5 +- torchtitan/models/kimi_k3/config_registry.py | 47 ++----- torchtitan/models/kimi_k3/model.py | 126 +------------------ torchtitan/models/kimi_k3/sharding.py | 103 +++++++++++++-- 4 files changed, 109 insertions(+), 172 deletions(-) diff --git a/torchtitan/models/common/decoder.py b/torchtitan/models/common/decoder.py index d4889c6640..d1aeed67f4 100644 --- a/torchtitan/models/common/decoder.py +++ b/torchtitan/models/common/decoder.py @@ -81,6 +81,7 @@ class Config(BaseModel.Config): # that support it set this True in their config factories; the tying # itself is handled by ``Decoder.__init__`` / ``Decoder.init_states``. enable_weight_tying: bool = False + @property def first_attention(self) -> BaseAttention.Config | None: """Attention config of the first layer that has one, else None. @@ -135,11 +136,12 @@ def max_context_length(self) -> int: ) return rope_cfg.max_context_length - def _validate_cp_backend(self, parallelism) -> None: """ShardingConfig-driven CP requires the spmd_types backend. A model whose CP is not ShardingConfig-driven overrides this and takes on its own preconditions.""" + from torchtitan.distributed.context_parallel import validate_cp_backend + validate_cp_backend(parallelism) def update_from_config( @@ -158,7 +160,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"), ( diff --git a/torchtitan/models/kimi_k3/config_registry.py b/torchtitan/models/kimi_k3/config_registry.py index 4ca621071b..d014ceacf3 100644 --- a/torchtitan/models/kimi_k3/config_registry.py +++ b/torchtitan/models/kimi_k3/config_registry.py @@ -99,44 +99,11 @@ def kimi_k3_debugmodel() -> Trainer.Config: def kimi_k3_debugmodel_text() -> Trainer.Config: - """The debug model with no vision tower, on packed text. - - The text arm of the context-parallel matrix: a CP failure here is the - decoder's, not the vision tower's. - """ - model_spec = model_registry("debugmodel_text") - return Trainer.Config( - loss=ChunkedLossWrapper.Config( - loss_fn=CrossEntropyLoss.Config( - global_vocab_size=decoder_vocab_size(model_spec), - ), - ), - hf_assets_path="./tests/assets/tokenizer", - tokenizer=MultiModalTokenizer.Config(**KIMI_K3_SPECIAL_TOKENS), - metrics=MetricsProcessor.Config(log_freq=1), - model_spec=model_spec, - dataloader=GrainDataLoader.Config( - dataset=ConcatThenSplitPackingConfig(dataset=TEXT_DATASETS["c4_test"]), - ), - optimizer=default_adamw(lr=8e-4), - lr_scheduler=LRSchedulersContainer.Config( - warmup_steps=2, - decay_ratio=0.8, - decay_type="linear", - min_lr_factor=0.0, - ), - # TODO: Kimi K3 has no spmd_types annotations yet. - parallelism=ParallelismConfig(spmd_backend="partial_dtensor"), - training=TrainingConfig( - num_tokens_per_microbatch_per_dp_rank=256, - max_context_length=256, - steps=10, - dtype="bfloat16", - disable_cuda_graphs=True, - ), - checkpoint=CheckpointManager.Config( - interval=10, - last_save_model_only=False, - ), - activation_checkpoint=SelectiveAC.Config(), + """The debug model with no vision tower, trained on the packed text dataset.""" + config = kimi_k3_debugmodel() + config.model_spec = model_registry("debugmodel_text") + config.loss.loss_fn.global_vocab_size = decoder_vocab_size(config.model_spec) + config.dataloader = GrainDataLoader.Config( + dataset=ConcatThenSplitPackingConfig(dataset=TEXT_DATASETS["c4_test"]), ) + return config diff --git a/torchtitan/models/kimi_k3/model.py b/torchtitan/models/kimi_k3/model.py index 2678357e5d..2cbeea0750 100644 --- a/torchtitan/models/kimi_k3/model.py +++ b/torchtitan/models/kimi_k3/model.py @@ -16,9 +16,7 @@ from torchtitan.models.common.attention import ( AttentionMasksType, BaseAttention, - create_attention_mask, FlexAttention, - get_causal_mask_mod, ) from torchtitan.models.common.decoder import Decoder from torchtitan.models.common.multimodal import ( @@ -31,7 +29,7 @@ from .kda import KimiDeltaAttention from .moe import KimiFeedForward, KimiLatentMoE -from .sharding import cp_all_to_all_headseq, ULYSSES +from .sharding import mla_ulysses_attention from .vision_encoder import KimiK3VisionEncoder # Shape suffixes: @@ -132,7 +130,9 @@ def forward( cp_group = self._cp_group if cp_group is not None and dist.get_world_size(cp_group) > 1: - out_THV = self._ulysses_attention(q_THK, kv_THC, k_rope_TK, cp_group) + out_THV = mla_ulysses_attention( + self, q_THK, kv_THC, k_rope_TK, cp_group + ) else: out_THV = self.inner_attention( q_THK, @@ -145,124 +145,6 @@ def forward( out_TD = out_TD * torch.sigmoid(self.gate(x_TD)) return self.wo(out_TD) - def _full_sequence_causal_mask(self, num_tokens: int, device): - """Causal mask for the sequence Ulysses reassembles. - - The mask the layer is handed has been sharded for context parallel by - ``cp_shard``, which cuts it the way ring attention wants: local queries - against global keys. Ulysses reassembles the whole sequence on every - rank instead, so it needs the whole causal mask. Rebuilding it is - correct here only because this model rejects sample packing, so the - sequence is one document and the mask carries no boundaries; a packed - sequence would need the global boundaries threaded down instead. - - Cached per (length, device) because the shape is constant across layers - and steps, and create_block_mask is compiled. - """ - # The mask the decoder builds at dp1 is causal AND packed-document - # (common/decoder._create_flex_attention_mask_for_document). This - # rebuild is causal only, which is equivalent exactly when the folded - # stream holds ONE document. Sample packing is already rejected in - # update_from_config, but a microbatch wider than the context window - # folds several documents into one stream as well, and then CP would - # let a sample attend to the previous one while dp1 would not -- - # silently, since every shape stays valid. Caught here rather than - # documented. - limit = getattr(self, "_cp_max_context_length", None) - if limit is not None and num_tokens > limit: - raise NotImplementedError( - f"context parallel folds {num_tokens} tokens into one stream " - f"but the context window is {limit}, so the " - "stream holds more than one document. The CP path rebuilds a " - "causal-only mask and cannot see document boundaries; use a " - "microbatch no wider than the context window." - ) - key = (num_tokens, device) - if self._cp_mask is None or self._cp_mask[0] != key: - mask = create_attention_mask( - get_causal_mask_mod(), - None, - None, - num_tokens, - num_tokens, - device=device, - ) - self._cp_mask = (key, mask) - return self._cp_mask[1] - - def _ulysses_attention( - self, - q_LHQ: torch.Tensor, - kv_LHC: torch.Tensor, - k_rope_LR: torch.Tensor, - cp_group, - ) -> torch.Tensor: - """Attention over the full sequence for this rank's head subset. - - One fused all-to-all trades the sharded axis, sequence for heads, then - the backend runs unchanged, then a second trades back. The gate and the - output projection stay sequence-local, so they are outside this. - - The rotary slice is deliberately not in the all-to-all. It is headless - -- one vector per token, shared by every head -- so it is all-gathered - along the sequence and expanded onto this rank's heads afterwards. - Packing the already-expanded key instead sends the same values once per - head and reassembles them against the wrong head subset, which shows up - as a forward that diverges from the same layer run without CP. - - Shape suffixes beyond the file legend: L local sequence (T/cp), G this - rank's head count (H/cp), W the packed per-head channel width, R the - rotary width. - """ - import torch.distributed.nn.functional as dist_nn - - # Head divisibility is checked at wiring time; see apply_cp_kimi_k3. - cp_size = dist.get_world_size(cp_group) - t_loc = q_LHQ.shape[0] - t_full = t_loc * cp_size - # Local head count: q_LHQ already carries this rank's local heads, so - # the CP split is over that, not over the global n_heads. - 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, - [self.q_head_dim, self.qk_nope_head_dim, self.v_head_dim], - dim=-1, - ) - - # Differentiable all-gather: the backward is a reduce-scatter, which is - # what 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, self.qk_rope_head_dim).expand( - t_full, h_cp, self.qk_rope_head_dim - ), - ], - dim=-1, - ) - - out_TGV = self.inner_attention( - q_TGQ, - k_TGQ, - v_TGV, - attention_masks=self._full_sequence_causal_mask(t_full, q_TGQ.device), - scale=self.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 - ) - - def _apply_attention_residual( prefix_sum_TD: torch.Tensor, block_residual_TND: torch.Tensor, diff --git a/torchtitan/models/kimi_k3/sharding.py b/torchtitan/models/kimi_k3/sharding.py index 46c50cf34d..5d1acd4fa4 100644 --- a/torchtitan/models/kimi_k3/sharding.py +++ b/torchtitan/models/kimi_k3/sharding.py @@ -16,8 +16,6 @@ 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. - -See CP_DECLARATIVE.md in the logbook for why KCP is an identity pair. """ from dataclasses import dataclass @@ -26,8 +24,13 @@ 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_causal_mask_mod, +) __all__ = [ @@ -55,7 +58,7 @@ class CPContract: """What one CP algorithm does to the [B, T, H, K] activations. Attributes: - name: ``kda_cp_mode`` spelling, and what the wiring log reports. + 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. @@ -106,9 +109,8 @@ def _shard_dim(layout: SpmdLayout) -> int: ) # KCP: the sequence stays sharded end to end (report sec 5.1.2). The delta-rule -# recurrence carries state rank to rank, which is a sequential dependency, not a -# redistribution -- no placement pair describes it, so it stays inside the op and -# the contract is an identity. Declared anyway to keep one shape for both modes. +# 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)), @@ -123,7 +125,7 @@ def _shard_dim(layout: SpmdLayout) -> int: def contract_for_mode(mode: str) -> CPContract: if mode not in _BY_MODE: - raise ValueError(f"kda_cp_mode must be one of {sorted(_BY_MODE)}, got {mode!r}") + raise ValueError(f"cp mode must be one of {sorted(_BY_MODE)}, got {mode!r}") return _BY_MODE[mode] @@ -143,7 +145,6 @@ def cp_all_to_all_headseq( bit-exact against a single-rank reference; backward is the transposed all-to-all via torch.distributed.nn.functional. """ - import torch.distributed.nn.functional as dist_nn if (src_dim, dst_dim) not in ((SEQ_DIM, HEAD_DIM), (HEAD_DIM, SEQ_DIM)): raise ValueError( @@ -170,3 +171,89 @@ def cp_all_to_all_headseq( # 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_causal_mask(attn, num_tokens: int, device): + """Causal-only mask for the sequence Ulysses reassembles, cached on + ``attn`` per (length, device). Correct only while the folded stream holds + ONE document, so a stream wider than the context window is rejected -- a + causal-only rebuild cannot see document boundaries.""" + limit = getattr(attn, "_cp_max_context_length", None) + if limit is not None and num_tokens > limit: + raise NotImplementedError( + f"context parallel folds {num_tokens} tokens into one stream " + f"but the context window is {limit}, so the " + "stream holds more than one document. The CP path rebuilds a " + "causal-only mask and cannot see document boundaries; use a " + "microbatch no wider than the context window." + ) + key = (num_tokens, device) + if attn._cp_mask is None or attn._cp_mask[0] != key: + mask = create_attention_mask( + get_causal_mask_mod(), None, None, num_tokens, num_tokens, device=device + ) + attn._cp_mask = (key, mask) + return attn._cp_mask[1] + + +def mla_ulysses_attention( + attn, + q_LHQ: torch.Tensor, + kv_LHC: torch.Tensor, + k_rope_LR: torch.Tensor, + cp_group, +) -> 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_causal_mask(attn, t_full, q_TGQ.device), + 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 + ) From 3557bbe7dc80440d54aab3a3de82dd51c115e9d4 Mon Sep 17 00:00:00 2001 From: QIU023 Date: Tue, 25 Aug 2026 10:24:03 -0400 Subject: [PATCH 06/14] common: the CP-backend seam's docstring is one line, and says when it dies --- torchtitan/models/common/decoder.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/torchtitan/models/common/decoder.py b/torchtitan/models/common/decoder.py index d1aeed67f4..f2201a2bfc 100644 --- a/torchtitan/models/common/decoder.py +++ b/torchtitan/models/common/decoder.py @@ -137,9 +137,7 @@ def max_context_length(self) -> int: return rope_cfg.max_context_length def _validate_cp_backend(self, parallelism) -> None: - """ShardingConfig-driven CP requires the spmd_types backend. A model - whose CP is not ShardingConfig-driven overrides this and takes on - its own preconditions.""" + """Overridable CP backend check. TODO: remove after KDA becomes torch-native and Kimi K3 migrates to spmd_types.""" from torchtitan.distributed.context_parallel import validate_cp_backend validate_cp_backend(parallelism) From 29cfb66e32cfee32e6359663fa08f4e3d782558d Mon Sep 17 00:00:00 2001 From: QIU023 Date: Wed, 26 Aug 2026 18:25:44 +0000 Subject: [PATCH 07/14] kimi_k3: align the vision splice with the sequence shard under context parallel Upstream ships one kimi_k3 flavor and it is multimodal. Turning context parallelism on with it raises upstream's own alignment check: torchtitan/models/common/multimodal.py:72, in get_vision_positions ValueError: Multimodal misalignment: found 0 contiguous run(s) of placeholder id 2016 in the token sequence but received 1 visual item(s). 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 -- and a slice may split a visual item or contain none at all. get_vision_positions needs whole items, so it refuses. Each rank now scatters the feature slice its own placeholders correspond to. The features are ordered by sequence position and CP shards are contiguous and equal -- the config already rejects a load balancer under CP, because a permuting one would break exactly that -- so a rank's slice starts after however many placeholders the lower ranks hold, which one all-reduce establishes. The rows a rank does not consume still reach the graph through add_zero_valued_dependency: FSDP2 issues the tower's reduce-scatter from the autograd hooks on its output, so leaving them out would have a subset of the process group issue the collective. This is correctness, not the report's sec 5.2.3 vision parallelism: the encoder still runs redundantly on every CP rank. Splitting the tower itself belongs to the DEP and dynamic-CP work, and nothing here anticipates it. Verified on the multimodal debug flavor at cp2, two steps, exit 0. Before this the same command dies in the check quoted above. --- torchtitan/distributed/fsdp.py | 24 ++++++++++++ torchtitan/models/kimi_k3/model.py | 61 ++++++++++++++++++++++++++++++ 2 files changed, 85 insertions(+) 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/kimi_k3/model.py b/torchtitan/models/kimi_k3/model.py index 2cbeea0750..db7a61356c 100644 --- a/torchtitan/models/kimi_k3/model.py +++ b/torchtitan/models/kimi_k3/model.py @@ -12,6 +12,7 @@ from torchtitan.hf_datasets.multimodal.mm_datasets import MMSamplePackingConfig +from torchtitan.distributed.fsdp import add_zero_valued_dependency from torchtitan.models.common import Linear from torchtitan.models.common.attention import ( AttentionMasksType, @@ -375,6 +376,22 @@ def _prepare_multimodal_embeds( 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, @@ -386,6 +403,50 @@ 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 forward( # pyrefly: ignore [bad-override] self, tokens: torch.Tensor, From 6c65c8b22e6f47b89f198672ff756621dca0ec78 Mon Sep 17 00:00:00 2001 From: QIU023 Date: Thu, 27 Aug 2026 10:20:53 +0000 Subject: [PATCH 08/14] kimi_k3: CP runs on the upstream multimodal debug flavor; the text flavor goes Same review direction as the EP PR: no text-only flavor. With the vision splice aligned to the sequence shard (previous commit), everything here runs on the flavor upstream ships: the cp2 CI cell and its recipe move to kimi_k3_debugmodel, and the flavor plus its model-registry entry leave the branch. Measured on this tip before the switch: dp1/cp2/cp4/cp8/dp2/ fsdp2_cp2/fsdp2_cp4 all train from one seed (mx3_cp_mm_0826_182604), with cp2/cp4/cp8 at 1.30e-2/1.24e-2/8.78e-3 from dp1 against 2.47e-2 for dp2 measured the same way. --- tests/integration_tests/features.py | 6 +++--- torchtitan/models/kimi_k3/__init__.py | 13 ------------- torchtitan/models/kimi_k3/config_registry.py | 12 ------------ torchtitan_recipes/tests/features.py | 6 +++--- 4 files changed, 6 insertions(+), 31 deletions(-) diff --git a/tests/integration_tests/features.py b/tests/integration_tests/features.py index 6c86ba0a6b..3084089a66 100755 --- a/tests/integration_tests/features.py +++ b/tests/integration_tests/features.py @@ -304,9 +304,9 @@ def build_features_test_list() -> list[OverrideDefinitions]: use_real_pg=True, ), OverrideDefinitions( - configs=[recipes.kimi_k3_debugmodel_text_cp2], - test_descr="Kimi K3 text decoder, context parallel cp2", - test_name="kimi_k3_text_cp2", + configs=[recipes.kimi_k3_debugmodel_cp2], + test_descr="Kimi K3, context parallel cp2", + test_name="kimi_k3_cp2", ngpu=2, use_real_pg=True, ), diff --git a/torchtitan/models/kimi_k3/__init__.py b/torchtitan/models/kimi_k3/__init__.py index 932fc728e5..e559f1469c 100644 --- a/torchtitan/models/kimi_k3/__init__.py +++ b/torchtitan/models/kimi_k3/__init__.py @@ -488,18 +488,6 @@ def _debugmodel(attn_backend: str) -> KimiK3Model.Config: ) -def _debugmodel_text(attn_backend: str) -> KimiK3Model.Config: - """The debug decoder with no vision tower. - - The text arm of the context-parallel matrix needs a flavor with no vision - path, so that a failure there is attributable to the decoder's CP rather - than to the tower or to the image/text token interleaving. - """ - config = _debugmodel(attn_backend) - config.vision_encoder = None - return config - - def _kimi_k3(attn_backend: str) -> KimiK3Model.Config: dim = 7168 return _kimi_k3_config( @@ -538,7 +526,6 @@ def _kimi_k3(attn_backend: str) -> KimiK3Model.Config: kimi_k3_configs = { "debugmodel": _debugmodel, - "debugmodel_text": _debugmodel_text, "Kimi-K3": _kimi_k3, } diff --git a/torchtitan/models/kimi_k3/config_registry.py b/torchtitan/models/kimi_k3/config_registry.py index d014ceacf3..dd6f4e1936 100644 --- a/torchtitan/models/kimi_k3/config_registry.py +++ b/torchtitan/models/kimi_k3/config_registry.py @@ -8,7 +8,6 @@ from torchtitan.components.checkpointer import CheckpointManager from torchtitan.components.data import GrainDataLoader, SingleDatasetConfig -from torchtitan.components.data.packing import ConcatThenSplitPackingConfig from torchtitan.components.loss import ChunkedLossWrapper, CrossEntropyLoss from torchtitan.components.metrics import MetricsProcessor from torchtitan.components.optimizer import default_adamw, LRSchedulersContainer @@ -21,7 +20,6 @@ MultiModalProcessor, ) from torchtitan.hf_datasets.multimodal.utils.image import resize_to_patch_budget -from torchtitan.hf_datasets.text_datasets import DATASETS as TEXT_DATASETS from torchtitan.models.common.config_utils import decoder_vocab_size from torchtitan.trainer import Trainer @@ -97,13 +95,3 @@ def kimi_k3_debugmodel() -> Trainer.Config: activation_checkpoint=SelectiveAC.Config(), ) - -def kimi_k3_debugmodel_text() -> Trainer.Config: - """The debug model with no vision tower, trained on the packed text dataset.""" - config = kimi_k3_debugmodel() - config.model_spec = model_registry("debugmodel_text") - config.loss.loss_fn.global_vocab_size = decoder_vocab_size(config.model_spec) - config.dataloader = GrainDataLoader.Config( - dataset=ConcatThenSplitPackingConfig(dataset=TEXT_DATASETS["c4_test"]), - ) - return config diff --git a/torchtitan_recipes/tests/features.py b/torchtitan_recipes/tests/features.py index c7e8f0e890..294ae942e3 100644 --- a/torchtitan_recipes/tests/features.py +++ b/torchtitan_recipes/tests/features.py @@ -422,15 +422,15 @@ def llama3_debugmodel_seed_checkpoint() -> Trainer.Config: config.training.disable_cuda_graphs = True return config -def kimi_k3_debugmodel_text_cp2() -> Trainer.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_text + from torchtitan.models.kimi_k3.config_registry import kimi_k3_debugmodel - config = kimi_k3_debugmodel_text() + config = kimi_k3_debugmodel() config.parallelism.context_parallel_degree = 2 config.parallelism.context_parallel_load_balancer = None return config From 3f9c29442d4cf51f7acce4bb9f58e1b568268a87 Mon Sep 17 00:00:00 2001 From: QIU023 Date: Fri, 28 Aug 2026 01:28:07 +0000 Subject: [PATCH 09/14] kimi_k3: config_registry returns to the upstream byte stream The branch's only remaining delta here was a trailing blank line, which pre-commit would re-normalize on the next touch anyway -- as it did on the EP branch. Restoring the upstream bytes takes the file out of the diff. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01WBy1d9YVu44nYCVqykRqL1 --- torchtitan/models/kimi_k3/config_registry.py | 1 - 1 file changed, 1 deletion(-) diff --git a/torchtitan/models/kimi_k3/config_registry.py b/torchtitan/models/kimi_k3/config_registry.py index dd6f4e1936..f1eac6b4fb 100644 --- a/torchtitan/models/kimi_k3/config_registry.py +++ b/torchtitan/models/kimi_k3/config_registry.py @@ -94,4 +94,3 @@ def kimi_k3_debugmodel() -> Trainer.Config: ), activation_checkpoint=SelectiveAC.Config(), ) - From cb7ecfd16ee72f23b1c46de8240ed2e035ded28c Mon Sep 17 00:00:00 2001 From: QIU023 Date: Sat, 29 Aug 2026 06:08:55 +0000 Subject: [PATCH 10/14] common: model-neutral wording for the CP backend check TODO --- torchtitan/models/common/decoder.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torchtitan/models/common/decoder.py b/torchtitan/models/common/decoder.py index f2201a2bfc..db2502e48b 100644 --- a/torchtitan/models/common/decoder.py +++ b/torchtitan/models/common/decoder.py @@ -137,7 +137,7 @@ 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 after KDA becomes torch-native and Kimi K3 migrates to spmd_types.""" + """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) From e26879345e35070e5bb62d3c66a06bf828afa36a Mon Sep 17 00:00:00 2001 From: QIU023 Date: Sat, 29 Aug 2026 17:27:14 +0000 Subject: [PATCH 11/14] kimi_k3: the cp2 integration cell skips on ROCm --- tests/integration_tests/features.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/integration_tests/features.py b/tests/integration_tests/features.py index 3084089a66..e89c716116 100755 --- a/tests/integration_tests/features.py +++ b/tests/integration_tests/features.py @@ -309,5 +309,6 @@ def build_features_test_list() -> list[OverrideDefinitions]: test_name="kimi_k3_cp2", ngpu=2, use_real_pg=True, + skip_rocm_test=True, ), ] From a1a8140738444eb703b336f0e39b64cccb4e4e0f Mon Sep 17 00:00:00 2001 From: QIU023 Date: Sat, 29 Aug 2026 21:55:37 +0000 Subject: [PATCH 12/14] kimi_k3: the Ulysses full-sequence mask preserves packed-document boundaries The CP path rebuilt a causal-only mask for the reassembled sequence and used the context window to reject streams that might hold several documents. After the all-to-all every rank holds the full sequence, so the global packed-document mask applies as-is: gather the contiguous positions shards (no load balancer under CP) and build the same causal x document mask the non-CP path uses. The window guard, its config plumbing and the shape-keyed mask cache go away -- the mask follows the data now. A two-rank gloo test packs three documents with one boundary on the shard cut and one inside a shard; the gathered mask equals the mask built from the global positions, and both boundary attentions are refused where a causal-only mask lets them through. --- .../cpu/test_kimi_k3_cp_document_mask.py | 89 +++++++++++++++++++ torchtitan/models/kimi_k3/model.py | 4 +- torchtitan/models/kimi_k3/parallelize.py | 6 +- torchtitan/models/kimi_k3/sharding.py | 50 ++++++----- 4 files changed, 119 insertions(+), 30 deletions(-) create mode 100644 tests/unit_tests/cpu/test_kimi_k3_cp_document_mask.py 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/models/kimi_k3/model.py b/torchtitan/models/kimi_k3/model.py index db7a61356c..3a102d5fa0 100644 --- a/torchtitan/models/kimi_k3/model.py +++ b/torchtitan/models/kimi_k3/model.py @@ -69,7 +69,6 @@ class Config(BaseAttention.Config): # Ulysses under either KDA CP mode -- KCP describes a recurrence that MLA # does not have. _cp_group = None - _cp_mask = None def __init__(self, config: Config): super().__init__() @@ -97,7 +96,6 @@ def forward( attention_masks: AttentionMasksType | None = None, positions: torch.Tensor | None = None, ) -> torch.Tensor: - del positions num_tokens = x_TD.shape[0] # The head count is DERIVED from the projection width, not read off @@ -132,7 +130,7 @@ def forward( 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 + self, q_THK, kv_THC, k_rope_TK, cp_group, positions ) else: out_THV = self.inner_attention( diff --git a/torchtitan/models/kimi_k3/parallelize.py b/torchtitan/models/kimi_k3/parallelize.py index 613823c390..d937b8f6d9 100644 --- a/torchtitan/models/kimi_k3/parallelize.py +++ b/torchtitan/models/kimi_k3/parallelize.py @@ -65,7 +65,7 @@ def parallelize_kimi_k3( assert isinstance(model, KimiK3Model) if parallel_dims.cp_enabled: - apply_cp_kimi_k3(model, parallel_dims, training.max_context_length) + apply_cp_kimi_k3(model, parallel_dims) if ac_config is not None: ac_policy = ac_config.build(dump_folder=dump_folder) @@ -105,7 +105,6 @@ def parallelize_kimi_k3( def apply_cp_kimi_k3( model: nn.Module, parallel_dims: ParallelDims, - max_context_length: int | None = None, ) -> None: """Wire context parallelism: KCP on the KDA layers, Ulysses on the MLA layers. @@ -122,9 +121,6 @@ def apply_cp_kimi_k3( kda_modules = [] for module in model.modules(): if isinstance(module, KimiMLAAttention): - # The CP mask rebuild is causal-only; the layer uses the context - # window to reject a folded stream holding several documents. - module._cp_max_context_length = max_context_length if module.n_heads % cp_degree != 0: raise ValueError( f"MLA n_heads={module.n_heads} must be divisible by " diff --git a/torchtitan/models/kimi_k3/sharding.py b/torchtitan/models/kimi_k3/sharding.py index 5d1acd4fa4..fafb6396ec 100644 --- a/torchtitan/models/kimi_k3/sharding.py +++ b/torchtitan/models/kimi_k3/sharding.py @@ -29,7 +29,7 @@ from torchtitan.distributed.parallel_dims import MeshAxisName, SpmdLayout from torchtitan.models.common.attention import ( create_attention_mask, - get_causal_mask_mod, + get_efficient_causal_mask_mod_for_packed_document, ) @@ -173,27 +173,32 @@ def cp_all_to_all_headseq( return out.permute(1, 0, 2, 3).reshape(t_loc, cp * h_loc, K).contiguous() -def full_sequence_causal_mask(attn, num_tokens: int, device): - """Causal-only mask for the sequence Ulysses reassembles, cached on - ``attn`` per (length, device). Correct only while the folded stream holds - ONE document, so a stream wider than the context window is rejected -- a - causal-only rebuild cannot see document boundaries.""" - limit = getattr(attn, "_cp_max_context_length", None) - if limit is not None and num_tokens > limit: - raise NotImplementedError( - f"context parallel folds {num_tokens} tokens into one stream " - f"but the context window is {limit}, so the " - "stream holds more than one document. The CP path rebuilds a " - "causal-only mask and cannot see document boundaries; use a " - "microbatch no wider than the context window." - ) - key = (num_tokens, device) - if attn._cp_mask is None or attn._cp_mask[0] != key: - mask = create_attention_mask( - get_causal_mask_mod(), None, None, num_tokens, num_tokens, device=device +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." ) - attn._cp_mask = (key, mask) - return attn._cp_mask[1] + 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( @@ -202,6 +207,7 @@ def mla_ulysses_attention( 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. @@ -250,7 +256,7 @@ def mla_ulysses_attention( q_TGQ, k_TGQ, v_TGV, - attention_masks=full_sequence_causal_mask(attn, t_full, q_TGQ.device), + attention_masks=full_sequence_document_mask(attn, positions_L, cp_group), scale=attn.scale, ) out_src_dim, out_dst_dim = ULYSSES.out_dims() From 837a662351ac9d3b80a830bd6dae377c2801ff38 Mon Sep 17 00:00:00 2001 From: QIU023 Date: Sun, 30 Aug 2026 06:01:01 +0000 Subject: [PATCH 13/14] kimi_k3: KCP's fla CP machinery routes through attention-gym's fla_cp wrappers --- torchtitan/models/kimi_k3/kda.py | 37 ++++++++++++------------ torchtitan/models/kimi_k3/parallelize.py | 13 +++++---- 2 files changed, 25 insertions(+), 25 deletions(-) diff --git a/torchtitan/models/kimi_k3/kda.py b/torchtitan/models/kimi_k3/kda.py index ab4789a31a..1804d5fdb3 100644 --- a/torchtitan/models/kimi_k3/kda.py +++ b/torchtitan/models/kimi_k3/kda.py @@ -95,21 +95,20 @@ def forward( def conv_with_halo(conv, x_local, cp_context, activation: str | None = None): - """Depthwise causal conv on a sequence-sharded input, exactly: fla's - ``causal_conv1d_cp`` 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 einops import rearrange - from fla.modules.conv.cp.ops import causal_conv1d_cp + """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=x_local, - weight=rearrange(conv.weight, "d 1 w -> d w"), - bias=conv.bias, + x_local, + conv.weight, + conv.bias, + cp_context, activation=getattr(conv, "activation", None) if activation is None else activation, - cp_context=cp_context, ) @@ -120,18 +119,18 @@ def build_kcp_context( conv1d_kernel_size: int | None = None, cu_seqlens=None, ): - """fla CP context for one evenly split sequence. ``cu_seqlens`` must be + """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 fla.ops.cp.context import build_cp_context - - if cu_seqlens is None: - world = dist.get_world_size(group) - total = seq_len_local * world - cu_seqlens = torch.tensor([0, total], dtype=torch.int32, device=device) - return build_cp_context( - cu_seqlens, group=group, conv1d_kernel_size=conv1d_kernel_size + 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, ) diff --git a/torchtitan/models/kimi_k3/parallelize.py b/torchtitan/models/kimi_k3/parallelize.py index d937b8f6d9..066c307224 100644 --- a/torchtitan/models/kimi_k3/parallelize.py +++ b/torchtitan/models/kimi_k3/parallelize.py @@ -135,14 +135,15 @@ def apply_cp_kimi_k3( # Checked at wiring time so the message is actionable, rather than an # ImportError from inside a layer's first forward. try: - from fla.modules.conv.cp.ops import causal_conv1d_cp # noqa: F401 - from fla.ops.cp.context import build_cp_context # noqa: F401 + 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 fla-core's CP ops " - "(fla.ops.cp.context.build_cp_context and " - "fla.modules.conv.cp.ops.causal_conv1d_cp), which ship in " - f"fla-core >= 0.5.1; import failed with: {err}." + "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: From a5339256e8eb6e18d6ca9a9a4aea26f3c2fefa8f Mon Sep 17 00:00:00 2001 From: QIU023 Date: Sun, 30 Aug 2026 06:16:23 +0000 Subject: [PATCH 14/14] kimi_k3: dynamic CP for the vision encoder -- patch partition, gather-KV, sub-CP groups --- torchtitan/models/common/vision_encoder.py | 8 + torchtitan/models/kimi_k3/__init__.py | 14 +- torchtitan/models/kimi_k3/model.py | 258 ++++++++++++++++++- torchtitan/models/kimi_k3/parallelize.py | 59 +++++ torchtitan/models/kimi_k3/vision_encoder.py | 260 +++++++++++++++++++- torchtitan/models/kimi_k3/vit_cp_plan.py | 170 +++++++++++++ 6 files changed, 758 insertions(+), 11 deletions(-) create mode 100644 torchtitan/models/kimi_k3/vit_cp_plan.py 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/model.py b/torchtitan/models/kimi_k3/model.py index 3a102d5fa0..d6613d1bcb 100644 --- a/torchtitan/models/kimi_k3/model.py +++ b/torchtitan/models/kimi_k3/model.py @@ -9,10 +9,11 @@ import torch import torch.distributed as dist from torch import nn - -from torchtitan.hf_datasets.multimodal.mm_datasets import MMSamplePackingConfig +from torch.distributed.tensor import DTensor 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, @@ -27,6 +28,7 @@ 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 @@ -144,6 +146,7 @@ def forward( out_TD = out_TD * torch.sigmoid(self.gate(x_TD)) return self.wo(out_TD) + def _apply_attention_residual( prefix_sum_TD: torch.Tensor, block_residual_TND: torch.Tensor, @@ -275,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): @@ -282,6 +308,9 @@ 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 @@ -343,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, @@ -359,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: @@ -367,7 +439,7 @@ 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 @@ -445,6 +517,186 @@ def _select_cp_shard( 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 066c307224..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 ( @@ -116,6 +117,7 @@ def apply_cp_kimi_k3( 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 = [] @@ -158,3 +160,60 @@ def apply_cp_kimi_k3( 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/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]