[DO NOT review, pending K3 text CP side merging] [Kimi K3] Add K3 MoonViT Dynamic CP support

tests/integration_tests/features.py

@@ -303,4 +303,12 @@ def build_features_test_list() -> list[OverrideDefinitions]:
             timeout=30,
             use_real_pg=True,
         ),
+        OverrideDefinitions(
+            configs=[recipes.kimi_k3_debugmodel_cp2],
+            test_descr="Kimi K3, context parallel cp2",
+            test_name="kimi_k3_cp2",
+            ngpu=2,
+            use_real_pg=True,
+            skip_rocm_test=True,
+        ),
     ]

tests/unit_tests/cpu/test_kimi_k3_cp_contracts.py

@@ -0,0 +1,63 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""The CP contracts, on CPU.
+
+These pin the folded token layout. This model carries no batch axis, so the
+Ulysses pair moves the shard between tensor dims 0 and 1; the batched spelling
+of the same contract used dims 1 and 2, and nothing in a shape check would
+catch that swap -- both dims exist and the all-to-all would produce a
+plausible tensor with the heads and the sequence exchanged.
+"""
+
+from __future__ import annotations
+
+import unittest
+
+import torch
+
+from torchtitan.models.kimi_k3.sharding import (
+    contract_for_mode,
+    cp_all_to_all_headseq,
+    HEAD_DIM,
+    KCP,
+    SEQ_DIM,
+    ULYSSES,
+)
+
+
+class TestCPContracts(unittest.TestCase):
+    def test_dims_are_the_folded_ones(self):
+        self.assertEqual((SEQ_DIM, HEAD_DIM), (0, 1))
+
+    def test_ulysses_swaps_the_sharded_axis(self):
+        self.assertEqual(ULYSSES.in_dims(), (SEQ_DIM, HEAD_DIM))
+        self.assertEqual(ULYSSES.out_dims(), (HEAD_DIM, SEQ_DIM))
+        self.assertTrue(ULYSSES.redistributes())
+        self.assertTrue(ULYSSES.head_sharded)
+
+    def test_kcp_is_an_identity_pair(self):
+        """The recurrence passes state rank to rank, which is a sequential
+        dependency rather than a redistribution, so no placement pair
+        describes it and the contract is declared as an identity."""
+        self.assertEqual(KCP.in_dims(), (SEQ_DIM, SEQ_DIM))
+        self.assertFalse(KCP.redistributes())
+        self.assertFalse(KCP.head_sharded)
+
+    def test_unknown_mode_is_rejected(self):
+        with self.assertRaises(ValueError):
+            contract_for_mode("ring")
+
+    def test_unimplemented_dim_pair_raises(self):
+        """A contract naming a pair with no implementation must raise here
+        rather than being quietly ignored."""
+        x = torch.zeros(4, 2, 3)
+        with self.assertRaises(ValueError):
+            cp_all_to_all_headseq(x, None, src_dim=SEQ_DIM, dst_dim=2)
+
+
+if __name__ == "__main__":
+    unittest.main()

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()

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,

torchtitan/models/common/decoder.py

@@ -136,6 +136,12 @@ def max_context_length(self) -> int:
                 )
             return rope_cfg.max_context_length
 
+        def _validate_cp_backend(self, parallelism) -> None:
+            """Overridable CP backend check. TODO: remove once linear-attention CP kernels are torch-native and models migrate to spmd_types."""
+            from torchtitan.distributed.context_parallel import validate_cp_backend
+
+            validate_cp_backend(parallelism)
+
         def update_from_config(
             self,
             *,
@@ -152,7 +158,6 @@ def update_from_config(
             that case the training/debug setup is skipped.
             """
             from torchtitan.config import ParallelismConfig
-            from torchtitan.distributed.context_parallel import validate_cp_backend
             from torchtitan.trainer import Trainer
 
             assert hasattr(config, "parallelism"), (
@@ -171,8 +176,7 @@ def update_from_config(
                 )
 
             if parallelism.context_parallel_degree > 1:
-                # ShardingConfig-based CP requires the spmd_types backend.
-                validate_cp_backend(parallelism)
+                self._validate_cp_backend(parallelism)
                 if any(self.traverse(ScaledDotProductAttention.Config)) or any(
                     self.traverse(VarlenAttention.Config)
                 ):

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

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