[DO NOT MERGE, pending rebase and sequential smaller PRs after eager-model merge] Kimi K3: full tree with 5D parallelism multimodal, LoRA, MoonViT Dynamic CP&DEP, MXFP4/MXFP8 QAT, quantile balancing

tests/unit_tests/test_moe_routing_map_placement.py

@@ -0,0 +1,83 @@
+# 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 routing-map scatter must preserve the router's shard placement.
+
+Under TP/SP with EP the router outputs are DTensors sharded on the token
+dim. A plain ``zeros_like(...).scatter_(...)`` cannot keep ``Shard(1)``:
+the in-place form has no DTensor strategy, and the out-of-place form
+redistributes to ``Replicate``, which turns the per-shard token counts
+that ``RoutedExperts`` consumes as ``Partial(sum)`` into full-sequence
+counts on every rank. See ``MoE.forward``.
+"""
+
+import torch
+from torch.distributed.tensor import distribute_tensor, DTensor, Replicate, Shard
+from torch.testing._internal.distributed._tensor.common_dtensor import (
+    DTensorTestBase,
+    with_comms,
+)
+
+
+def _routing_map(scores_BLE, topk_expert_ids_BLK):
+    """The construction under test, lifted from ``MoE.forward``."""
+    if isinstance(scores_BLE, DTensor):
+        local_map = torch.zeros_like(
+            scores_BLE.to_local(), dtype=torch.bool
+        ).scatter_(-1, topk_expert_ids_BLK.to_local(), True)
+        return DTensor.from_local(
+            local_map, scores_BLE.device_mesh, scores_BLE.placements
+        )
+    return torch.zeros_like(scores_BLE, dtype=torch.bool).scatter_(
+        -1, topk_expert_ids_BLK, True
+    )
+
+
+class TestRoutingMapPlacement(DTensorTestBase):
+    @property
+    def world_size(self) -> int:
+        return 2
+
+    @with_comms
+    def test_token_sharded_scores_keep_placement(self):
+        mesh = self.build_device_mesh()
+        B, L, E, K = 2, 8, 4, 2
+        torch.manual_seed(0)
+        scores = torch.rand(B, L, E, device=self.device_type)
+        topk = scores.topk(K, dim=-1).indices
+
+        expected = _routing_map(scores, topk)
+
+        scores_dt = distribute_tensor(scores, mesh, [Shard(1)])
+        topk_dt = distribute_tensor(topk, mesh, [Shard(1)])
+        got = _routing_map(scores_dt, topk_dt)
+
+        self.assertEqual(got.placements, scores_dt.placements)
+        self.assertEqual(got.full_tensor(), expected)
+        # RoutedExperts reduces the counts as Partial(sum), so each rank must
+        # hold LOCAL counts, not full-sequence counts.
+        self.assertEqual(
+            got.to_local().sum(dim=(0, 1)).sum().item(),
+            B * (L // self.world_size) * K,
+        )
+
+    @with_comms
+    def test_replicated_scores_unchanged(self):
+        mesh = self.build_device_mesh()
+        torch.manual_seed(0)
+        scores = torch.rand(2, 8, 4, device=self.device_type)
+        topk = scores.topk(2, dim=-1).indices
+        scores_dt = distribute_tensor(scores, mesh, [Replicate()])
+        topk_dt = distribute_tensor(topk, mesh, [Replicate()])
+        got = _routing_map(scores_dt, topk_dt)
+        self.assertEqual(got.placements, scores_dt.placements)
+        self.assertEqual(got.full_tensor(), _routing_map(scores, topk))
+
+
+if __name__ == "__main__":
+    import unittest
+
+    unittest.main()

torchtitan/components/optimizer/lr_scheduler.py

@@ -196,9 +196,11 @@ def linear_warmup_stable_decay(
     schedulers: list[LRScheduler]
 
     def __init__(self, optimizers: OptimizersContainer, lr_lambda: Callable) -> None:
-        assert (
-            len(optimizers) > 0
-        ), "Must have at least one optimizer to create LRScheduler"
+        # No assert on len(optimizers) > 0. A pipeline stage that owns only frozen
+        # weights gets no optimizer (see OptimizersContainer._build_param_groups), and it
+        # has no learning rate to schedule either. The comprehension below then yields
+        # zero schedulers, and step() / get_last_lr() iterate, so an empty container is
+        # well defined. LoRA plus PP is what produces such a stage.
 
         self.schedulers = [LambdaLR(optimizer, lr_lambda) for optimizer in optimizers]
 
@@ -224,13 +226,24 @@ def step(self) -> None:
             scheduler.step()
 
     def state_dict(self) -> dict[str, Any]:
+        # A container with no schedulers has no last_epoch to report, and returning an
+        # empty dict is the honest answer rather than a zero that a later load would
+        # apply as real progress. DCP is fine with a rank contributing no keys.
+        if not self.schedulers:
+            return {}
+
         # Only last_epoch is needed — each scheduler recomputes its lr from
         # its own optimizer's base_lrs on load. Per-scheduler state (base_lrs,
         # _last_lr) is not saved because it's reconstructed from the optimizer
         # config at construction time.
         return {"last_epoch": self.schedulers[0].last_epoch}
 
     def load_state_dict(self, state_dict: dict[str, Any]) -> None:
+        # Nothing to restore on a stage that schedules nothing. Returning before reading
+        # the key matters: a frozen-only stage's own state_dict() never wrote one.
+        if not self.schedulers:
+            return
+
         # Only restore last_epoch. Each scheduler recomputes _last_lr from its
         # own optimizer's base_lrs and the shared lambda. This is correct for
         # mixed optimizers (different base_lrs) and resharding (different number

torchtitan/components/optimizer/optimizer.py

@@ -198,10 +198,25 @@ def _build_param_groups(
                     claimed.add(name)
 
             if not params:
-                raise ValueError(
-                    f"Optimizer param_groups pattern '{pg.pattern}' "
-                    f"matched no parameters"
+                # An empty match means one of two different things, and only one is an
+                # error. If the model has trainable parameters and none matched, the
+                # pattern is wrong -- keep raising, that is what this check is for. If
+                # the model has NO trainable parameters at all, the pattern is fine and
+                # this model part simply has nothing to optimize: a pipeline stage that
+                # owns only frozen weights. LoRA plus PP produces exactly that, e.g. a
+                # DEP vision stage holding an unadapted tower and frozen embeddings.
+                if any(p.requires_grad for p in model.parameters()):
+                    raise ValueError(
+                        f"Optimizer param_groups pattern '{pg.pattern}' "
+                        f"matched no parameters"
+                    )
+                logger.warning(
+                    "Optimizer param_groups pattern '%s' matched no parameters and "
+                    "this model part has none that require grad; it gets no optimizer. "
+                    "Expected for a pipeline stage that owns only frozen weights.",
+                    pg.pattern,
                 )
+                continue
 
             groups[pg.optimizer_name].append(
                 {
@@ -327,7 +342,12 @@ def load_state_dict(self, state_dict: dict[str, Any]) -> None:
     def _post_init(self, all_params: list[nn.Parameter]) -> None:
         # We need to call Optimizer.__init__() to initialize some necessary optimizer
         # functionality such as hooks (e.g. register_step_pre_hook for MoE load balancing).
-        Optimizer.__init__(self, all_params, {})
+        #
+        # torch rejects an empty params LIST but accepts an empty param GROUP, and the
+        # difference matters here: a pipeline stage owning only frozen weights has nothing
+        # to optimize, and it still has to end up a properly initialized Optimizer so the
+        # hooks and param_groups exist. LoRA plus PP produces such a stage.
+        Optimizer.__init__(self, all_params or [{"params": []}], {})
 
     def _register_bf16_optimizer_state_hook(self) -> None:
         """Register a step pre-hook to create Adam optimizer states in bfloat16.

torchtitan/distributed/fsdp.py

@@ -165,6 +165,35 @@ def apply_fsdp_to_vision_encoder(
     )
 
 
+def add_zero_valued_dependency(
+    output: torch.Tensor,
+    unused_output: torch.Tensor,
+) -> torch.Tensor:
+    """Keep a conditionally executed 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 module
+    that only some data-parallel ranks execute -- a VLM vision encoder on a
+    batch that happens to carry no images, for example -- would therefore
+    issue collectives on a subset of the process group and deadlock the step.
+
+    A rank with no real work for such a module runs it on a placeholder input
+    and routes the result through this helper. 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 --
+    which is also its 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.
+
+    VENDORED verbatim from pytorch/torchtitan#4025 (the Kimi K3 eager
+    reference PR), which adds it to this file. Delete this copy and import the
+    upstream one when that PR lands -- kimi_k3 is the only caller.
+    """
+    return output + unused_output.sum().to(output.dtype) * 0.0
+
+
 def apply_fsdp_to_decoder(
     model: "Decoder",
     dp_mesh: DeviceMesh,

torchtitan/distributed/utils.py

@@ -591,6 +591,58 @@ def set_pg_timeouts(
         torch.distributed.set_timeout(timeout, group)
 
 
+@torch.no_grad()
+def _get_total_norm_fp32(
+    tensors: Iterable[torch.Tensor],
+    norm_type: float,
+    error_if_nonfinite: bool,
+    foreach: bool | None,
+) -> torch.Tensor:
+    """``torch.nn.utils.get_total_norm`` with the reduction carried in float32.
+
+    That function returns the norm in the TENSORS' dtype. With bf16 gradients both the
+    per-tensor norms and the norm-of-norms are bf16 -- three to four significant digits --
+    so the total is wrong by a few tenths of a percent, and by an amount that depends on
+    how the tensors are GROUPED. Under PP or EP each rank norms its own share, so the
+    value a rank reports depends on where the pipeline was cut rather than only on the
+    gradients. Measured on 394 bf16 tensors: 0.184% error overall, and 0.057% / 0.142% /
+    0.174% for three different splits of the same tensors.
+
+    With clipping active the error is not cosmetic -- it scales the update. A run with
+    ``max_norm=1.0`` against a true norm near 10 clips by ``max_norm / total_norm``, so a
+    0.2% error in the norm is a 0.2% error in that step's effective learning rate.
+
+    Structure follows the upstream function so the only difference is precision: same
+    empty-list result, same error_if_nonfinite behaviour, and the foreach fast path where
+    it applies. DTensors go through ``vector_norm`` one at a time because upstream's
+    foreach support check excludes them too; the dtype argument preserves their
+    ``_NormPartial`` placement, so the caller's ``full_tensor()`` still reduces correctly.
+    """
+    tensors = list(tensors)
+    if not tensors:
+        # Upstream returns a plain CPU scalar here, and callers depend on that: it is how
+        # a PP rank owning no gradients of one group signals "nothing to contribute".
+        return torch.tensor(0.0)
+    if any(isinstance(t, DTensor) for t in tensors):
+        norms: list[torch.Tensor] = [
+            torch.linalg.vector_norm(t, norm_type, dtype=torch.float32) for t in tensors
+        ]
+    else:
+        norms = list(torch._foreach_norm(tensors, norm_type, dtype=torch.float32))
+    first_device = tensors[0].device
+    total_norm = torch.linalg.vector_norm(
+        torch.stack([norm.to(first_device) for norm in norms]), norm_type
+    )
+    if error_if_nonfinite and torch.logical_or(total_norm.isnan(), total_norm.isinf()):
+        raise RuntimeError(
+            f"The total norm of order {norm_type} for gradients from "
+            "`parameters` is non-finite, so it cannot be clipped. To disable "
+            "this error and scale the gradients by the non-finite norm anyway, "
+            "set `error_if_nonfinite=False`"
+        )
+    return total_norm
+
+
 @torch.no_grad()
 def clip_grad_norm_(
     parameters: torch.Tensor | Iterable[torch.Tensor],
@@ -645,9 +697,7 @@ def clip_grad_norm_(
         # prevent generators from being exhausted
         parameters = list(parameters)
     grads = [p.grad for p in parameters if p.grad is not None]
-    total_norm = torch.nn.utils.get_total_norm(
-        grads, norm_type, error_if_nonfinite, foreach
-    )
+    total_norm = _get_total_norm_fp32(grads, norm_type, error_if_nonfinite, foreach)
 
     # If total_norm is a DTensor, the placements must be `torch.distributed._tensor.ops.math_ops._NormPartial`.
     # We can simply reduce the DTensor to get the total norm in this tensor's process group
@@ -661,6 +711,15 @@ def clip_grad_norm_(
         total_norm = total_norm.full_tensor()
 
     if pp_mesh is not None:
+        # Normalise dtype and device before the collective. get_total_norm returns a CPU
+        # float32 tensor(0.) for an empty gradient list, so a PP rank whose share of the
+        # model contributes no gradients to one of these groups reaches this all_reduce
+        # with float32 while a peer holding bf16 gradients carries bfloat16. NCCL returns
+        # GARBAGE for a dtype mismatch instead of raising: the observed symptom is one
+        # side of the pipeline reporting a plausible norm -- its own PP shard's sum only --
+        # and the other reporting NaN, while every individual gradient is finite. float32
+        # is also the right width for a sum of squares.
+        total_norm = total_norm.to(device=pp_mesh.device_type, dtype=torch.float32)
         if math.isinf(norm_type):
             dist.all_reduce(total_norm, op=dist.ReduceOp.MAX, group=pp_mesh.get_group())
         else:
@@ -704,14 +763,14 @@ def _clip_grad_norm_with_ep(
     # - In autoparallel, all params may live on a single sparse mesh with "ep" dimension,
     #   so non_ep_grads would be empty
     # - In PP + EP setups, certain PP ranks may only own EP or non-EP layers
-    ep_grads_total_norm = torch.nn.utils.get_total_norm(
+    ep_grads_total_norm = _get_total_norm_fp32(
         ep_grads, norm_type, error_if_nonfinite, foreach
     )
     # get_total_norm returns tensor(0.) for empty list, which is a non-DTensor
     if isinstance(ep_grads_total_norm, DTensor):
         ep_grads_total_norm = ep_grads_total_norm.full_tensor()
 
-    non_ep_grads_total_norm = torch.nn.utils.get_total_norm(
+    non_ep_grads_total_norm = _get_total_norm_fp32(
         non_ep_grads, norm_type, error_if_nonfinite, foreach
     )
     # get_total_norm returns tensor(0.) for empty list, which is a non-DTensor
@@ -727,6 +786,15 @@ def _clip_grad_norm_with_ep(
         total_norm **= 1.0 / norm_type
 
     if pp_mesh is not None:
+        # Normalise dtype and device before the collective. get_total_norm returns a CPU
+        # float32 tensor(0.) for an empty gradient list, so a PP rank whose share of the
+        # model contributes no gradients to one of these groups reaches this all_reduce
+        # with float32 while a peer holding bf16 gradients carries bfloat16. NCCL returns
+        # GARBAGE for a dtype mismatch instead of raising: the observed symptom is one
+        # side of the pipeline reporting a plausible norm -- its own PP shard's sum only --
+        # and the other reporting NaN, while every individual gradient is finite. float32
+        # is also the right width for a sum of squares.
+        total_norm = total_norm.to(device=pp_mesh.device_type, dtype=torch.float32)
         if math.isinf(norm_type):
             dist.all_reduce(total_norm, op=dist.ReduceOp.MAX, group=pp_mesh.get_group())
         else:

torchtitan/models/__init__.py

@@ -10,6 +10,7 @@
         "flux",
         "gpt_oss",
         "kimi_k2_7",
+        "kimi_k3",
         "llama3",
         "muse_glimmer",
         "qwen3",