[kimi k3] add eager reference model with FSDP2 by JavaZeroo · Pull Request #4025 · pytorch/torchtitan

JavaZeroo · 2026-07-30T07:49:26Z

Summary

Adds a PyTorch-native numerical reference for the released Kimi K3 architecture,
together with pure FSDP2 data parallelism.

Implements KDA, gated MLA, LatentMoE, block attention residuals, MoonViT3d,
the multimodal projector, and image-feature scatter.
Adds a topology-complete 13-layer debugmodel that preserves the dense/MoE
transition, KDA/MLA cadence, two attention-residual blocks, routed and shared
experts, and vision path while reducing decoder widths, depth, and expert
count. It retains the released vocabulary and has about 100M parameters.
Adds an unquantized Hugging Face state-dict adapter.
Applies FSDP2 to the vision encoder and the MoE-aware decoder using
TorchTitan's shared FSDP helpers, without changing model math or state-dict
names.
Adds numerical parity against the released Hugging Face implementation, exact
BF16 eager/FSDP forward and per-parameter gradient coverage, and a two-GPU
multimodal FSDP integration test.

KDA runs on FLA's chunked kernel (fla.ops.kda.chunk_kda), the same split
Qwen3.5 uses: the kernel is the training path and the pure-PyTorch recurrence
lives in the unit tests as the numerical reference. Every other operator --
MLA, the KDA short convolutions, LatentMoE, routing, and the whole vision path
-- is ordinary eager PyTorch.

Why

Related to #3029, which tracks the broader Kimi K3 pre-training,
post-training, and multi-dimensional parallelism effort.

This PR intentionally establishes a reduced, inspectable numerical baseline
before expert parallelism lands. FSDP2 provides a useful distributed training
path while preserving the same forward contract used for comparison with the
released Hugging Face implementation.

Validation

Environment: PyTorch 2.14.0.dev20260729+cu130, fla-core
0.5.2.

Numerical parity with the released implementation

tests/unit_tests/test_kimi_k3_hf_parity.py freezes float32 outputs from a
deterministic reduced model evaluated with the released Hugging Face code at
commit c5d1dd4c428bd1ce8b88c5044f3b6ccde9e3b721, covering text logits, routed
expert IDs, projected vision features, and end-to-end image+text logits. The
model is loaded strictly from the state dict produced by
KimiK3StateDictAdapter, so no checkpoint or network access is required. Every
refactor in this PR was accepted only with those values unchanged.

test_kimi_k3.py compares the FLA kernel against the explicit recurrence on
CUDA, forward and backward, for both gate activations.

Training

Single-GPU kimi_k3_debugmodel, 10 steps, --debug.seed 42 --debug.deterministic: loss tracks the pure-PyTorch KDA baseline to within
3e-3 with matching grad norms. Not bitwise equal, since a chunked kernel
replaces a sequential FP32 recurrence.
The scaled flavor builds as 100,051,368 parameters.

FSDP2

Single-rank FSDP2 forward and per-parameter gradients are bitwise equal to an
unsharded BF16 reference, including the attention-residual pass-through path.
Two-rank FSDP2 forward and backward with one rank holding an image and the
other text-only.
tests/unit_tests is 562 passed; pre-commit run --all-files and Pyrefly
over the changed files are clean.

Initial scope

Eager single-device execution and pure FSDP2 data parallelism are supported.
HSDP, TP, PP, CP, EP, activation checkpointing, torch.compile, and CPU
offload are rejected explicitly.
Image inputs are supported; video inputs are not.
Routed experts keep one module per expert, matching the released
KimiSparseMoeBlock layout, which is what lets the state-dict adapter map
experts directly. Expert parallelism will need the grouped-GEMM layout and is
left to a follow-up.
The released routed-expert checkpoint is MXFP4; direct loading of those
compressed weights and a full 2.8T flavor are follow-up work.

pytorch-bot · 2026-07-30T07:49:34Z

~~Workflows were awaiting approval.~~ CI has now been triggered for the ciflow labels on this PR.

shuhuayu · 2026-07-31T00:27:19Z

@JavaZeroo Thanks for the pr. It looks workable to me in general, have some general comments, do you have bandwidth to work on them timely?

Can you reuse more common modules in torchtitan today, like MoonViT in common/vision_encoder, common/multi_modal, and reuse MoE router.
this pr should include a reproducible numerical parity check test, likely using some reduced hf reference, since the original hf ckpt is too big.
The current code only supports fsdp2, is there bottleneck to support other parallelisms?
Currently fsdps wrap vision encoder independently, but vit is used conditionally, this will be problematic of one dp rank gets text only but another dp rank gets images.

shuhuayu · 2026-07-31T00:29:03Z

@claude please review.

claude · 2026-07-31T00:29:29Z

Claude finished @shuhuayu's task in 4m 3s —— View job

Code Review: Kimi K3 eager reference + FSDP2

I read through model.py, vision_encoder.py, parallelize.py, state_dict_adapter.py, config_registry.py, both unit tests, and the relevant shared infra (common/moe.py, common/vision_encoder.py, common/multimodal.py, distributed/fsdp.py, components/optimizer.py). Overall this is a clean, well-documented eager reference that follows the repo's shape-suffix and config conventions and already reuses Decoder, VisionMLP, get_vision_positions, and the shared FSDP helpers. Below I fold in @shuhuayu's four questions plus concrete findings.

1. Reuse of common modules (maintainer Q1) — largely agree, several concrete duplications

MoE router. KimiMoERouter (model.py:426-460) re-implements what common/moe.py::TokenChoiceTopKRouter already does: score_func="sigmoid", route_norm, route_scale, and expert_bias_E-based routing with topk(sorted=False) + gather. The only real difference is that Kimi computes the gate in fp32 via F.linear(x.float(), self.gate.weight.float()) instead of an autocast(float32) around self.gate(x). That is a very thin reason to fork the router — recommend reusing TokenChoiceTopKRouter (it already returns scores_BLE, which you need for the routing map) and, if the fp32 gate matters for HF parity, adding that as an option there rather than a parallel class.
Vision block. KimiK3VisionBlock (vision_encoder.py:320-349) is structurally identical to common/vision_encoder.py::VisionTransformerBlock (norm -> attn -> residual -> norm -> mlp). The reason for the fork is the eager, block-diagonal attention vs. the shared FlexAttention path. The shared VisionAttention already injects RoPE through a rope_apply callable and takes an inner_attention config — an eager block-diagonal attention module implementing that same interface would let you reuse VisionTransformerBlock/VisionAttention verbatim. Worth attempting; if the padded per-item Python loop genuinely can't fit the attention_mask: BlockMask contract, document why in the module docstring.
Vision embed scatter. _replace_vision_embeds (model.py:38-65) duplicates common/multimodal.py::scatter_vision_embeds. The only difference is out-of-place index_copy vs. in-place assignment (added in commit cb8d92c). This belongs in common/multimodal.py as the canonical out-of-place variant (or make the existing one autograd-safe), not a private copy in the model — other VLMs will want the same. Fix this →
KimiRMSNorm subclassing RMSNorm just to expose kimi_eps and force the fp32 reduction is reasonable, but check whether common/nn_modules.py::RMSNorm already reduces in fp32 — if so the subclass may be unnecessary.

2. Reproducible numerical parity test (maintainer Q2) — missing, should be added

The PR body reports strong parity numbers (text logits max abs err 1.7e-4, exact routed-expert IDs), but there is no committed test that reproduces them. test_kimi_k3.py covers topology, the KDA kernel vs. a local recurrence, GELU, and a to_hf/from_hf round-trip — all good — but nothing pins numerics against an HF reference. As @shuhuayu notes, the full checkpoint is too big, so the right move is a tiny randomly-initialized config loaded into both the TorchTitan module and a minimal HF-equivalent reference (or a saved golden-logits tensor) with an explicit atol/rtol. Without this, silent math drift in future refactors won't be caught, which is exactly the guarantee CLAUDE.md asks for.

3. Bottlenecks for other parallelisms (maintainer Q3)

The eager reference has two structural blockers to TP/CP/PP/EP beyond "not wired up yet":

Data-dependent, per-token Python loops. The KDA recurrence (KimiKDAKernel, model.py:317-330) loops over sequence length, and KimiLatentMoE.forward (model.py:548-558) loops over experts with torch.nonzero/index_select. These are inherently incompatible with torch.compile and CP, and the list-of-modules expert layout can't use the grouped-GEMM EP path in common/moe.py. EP specifically will require moving to GroupedExperts (or an equivalent grouped layout) — the list-backed KimiRoutedExperts is fine for a reference but is a dead end for EP.
KimiK3Model.forward overrides Decoder.forward with an incompatible signature (model.py:802, # pyrefly: ignore [bad-override]) and threads a block_residual_TND tuple through every layer. PP's pipelining stage-splitting assumes the standard decoder contract, so PP will need this reworked. Reasonable to defer, but worth stating explicitly in the README's "Initial scope" as the reason, not just "rejected."

4. Vision encoder FSDP + conditional execution (maintainer Q4) — real hang risk, agree

This is the most important correctness concern. parallelize.py:76-85 wraps vision_encoder as its own FSDP unit, but KimiK3Model.forward only calls it when pixel_values is not None (model.py:774-783). Under FSDP2 the unshard all-gather (forward) and reduce-scatter (backward) for the vision params only fire on ranks that actually run the encoder. If DP rank A gets an image batch and rank B gets a text-only batch, the collectives are mismatched and training deadlocks. The kimi_k3_mm_fsdp integration test never exercises this because the cc12m-test dataloader always yields images on both ranks, so CI won't catch it. Options: require every rank to run the encoder every step (e.g. a zero-token dummy forward when no images) and document the invariant, or gate on a per-rank flag with an explicit collective barrier. At minimum, add a test/assert for the mixed text-only/image case rather than leaving it latent.

Other findings

tokens_per_expert_E is reassigned out-of-place in KimiLatentMoE.forward (model.py:539-541: self.tokens_per_expert_E = self.tokens_per_expert_E + ...) whereas common/moe.py uses in-place .add_(). The load-balancing hook (optimizer.py:520-521) does moe.expert_bias_E.add_() and moe.tokens_per_expert_E.zero_(), and register_moe_load_balancing_hook is wired via post_optimizer_build_fn. Reassigning replaces the buffer object each step; confirm the hook still sees the live buffer (the unit test asserts identity after one forward, but the hook reads it after multiple forwards / grad-accum). In-place add_ as in common MoE would be safer and matches the documented AC double-count behavior.
KimiMoERouter reads self.gate.weight directly instead of calling self.gate(x). This works under pure FSDP (params are unsharded plain tensors inside forward) but silently breaks under TP, where weight is a sharded DTensor and F.linear(plain, dtensor) won't do what you want. TP is rejected today, so this is latent, but it's the kind of direct-.weight access that the distributed rules warn against.
MLA allocates a full L x L causal mask every forward (model.py:244-250) and vision attention runs a Python per-item loop. Both are fine as a reference; just flagging the O(L^2)/O(num_items) cost so it isn't mistaken for a training-ready path.
Param budget is dominated by the vocab. vocab_size=163840 with dim=256 puts ~84M of the ~100M params in tok_embeddings + lm_head, so the "100M debug model" exercises a very small transformer. That's a defensible choice (retains released vocab for exact tokenizer/state-dict parity), but consider noting it so reviewers don't read "100M" as transformer capacity.
Unused config field: KimiMLAAttention.Config.q_lora_rank is stored in the config but never used in __init__ (dims are derived from the projection configs). Either use it for an assertion or drop it.

Style / minor

ASCII-only in new comments/docstrings looks respected. Good.
parallelize.py:62 if ac_config is not None: rejects AC, but config_registry.py:76 sets activation_checkpoint=None, so the debug config is internally consistent — good.
README documents scope clearly; please add the parity-test reproduction steps once test Add FSDP grad scaler to the train loop #2 lands.

None of the above are blocking for an explicitly-scoped eager reference, but I'd prioritize (4) the conditional-vision-encoder FSDP hang and (2) a committed parity test before merge, and treat the router/vision-block/scatter reuse as the main follow-ups for the "reduce duplication" goal.
· agent/add-kimi-k3-reference-model

JavaZeroo · 2026-07-31T03:34:22Z

@JavaZeroo Thanks for the pr. It looks workable to me in general, have some general comments, do you have bandwidth to work on them timely?

Can you reuse more common modules in torchtitan today, like MoonViT in common/vision_encoder, common/multi_modal, and reuse MoE router.

this pr should include a reproducible numerical parity check test, likely using some reduced hf reference, since the original hf ckpt is too big.

The current code only supports fsdp2, is there bottleneck to support other parallelisms?

Currently fsdps wrap vision encoder independently, but vit is used conditionally, this will be problematic of one dp rank gets text only but another dp rank gets images.

Thanks @shuhuayu. Yes, I have bandwidth and am working on these items now.

I am reusing VisionTransformerBlock/VisionMLP, the common multimodal scatter helper, and TokenChoiceTopKRouter. I am currently running numerical validation for these changes.
Agreed. I am adding a reproducible reduced-model numerical parity test against the pinned Hugging Face Kimi K3 reference.
I would like to implement a version of fsdp with correct precision first. other parallel implementations on k3 may require more adaptation work, especially cp and pp. i think it is possible to quickly support a version of fsdp for k3 first, enough to support our experiments with small models. In the meantime I'm happy to be able to add ep tp support to this pr, do you see any need to add it in this pr?
I reproduced the mixed-modality FSDP issue and am working on the fix.

shuhuayu · 2026-07-31T06:05:35Z

I would like to implement a version of fsdp with correct precision first. other parallel implementations on k3 may require more adaptation work, especially cp and pp. i think it is possible to quickly support a version of fsdp for k3 first, enough to support our experiments with small models. In the meantime I'm happy to be able to add ep tp support to this pr, do you see any need to add it in this pr?

Support fsdp first sounds good to me. I think we should target for a training ready version using kernels for kda, one option is to use from fla.ops.kda import chunk_kda (similarly we used fla kernels for qwen 3.5, and put the pytorch native reference implementation into the numerical tests). We may have plan to use our own kernel for kda in the future.

cc: @tianyu-l

Addresses the review feedback on pytorch#4025. - KDA now dispatches to fla.ops.kda.chunk_kda with the gate activation, beta sigmoid, and q/k L2 norm fused into the kernel, following how Qwen3.5 uses FLA. The pure-PyTorch recurrence becomes ReferenceKimiKDAKernel in the unit tests, which the CPU suite builds the model with, and a CUDA-only test checks the kernel against it forward and backward for both gate activations. FLA cannot compile head dimensions below 16, so the config now rejects those with a clear error instead of a Triton compilation failure. - The vision encoder runs on every batch rather than only when images are present. It is its own FSDP unit, and the shared multimodal collator can hand one data-parallel rank a text-only batch, so conditional execution issued collectives on a subset of the process group and could deadlock the step. Batches without images use the smallest mergeable grid and contribute through add_zero_valued_dependency, which leaves the text embeddings numerically unchanged. This replaces the flag parallelize() used to set, so single-GPU and multi-GPU take the same forward path. - KimiMoERouter is replaced by the common TokenChoiceTopKRouter, which also removes a direct self.gate.weight read that would break under TP. - The private out-of-place vision scatter is dropped for the shared scatter_vision_embeds. FSDP2 only loses its pre-backward hook when a wrapped module returns a view, and Embedding returns a fresh tensor from F.embedding, so the fork was unnecessary. Its test now covers the shared helper instead. - tokens_per_expert_E is updated in place so the load-balancing hook keeps referring to the live buffer, and the unused q_lora_rank field is removed. Validated on 1x RTX 5080 with PyTorch 2.14.0.dev20260729+cu130 and fla-core 0.5.2: the frozen HuggingFace parity values are unchanged, the kimi_k3 tests pass (13, including the CUDA kernel comparison), and a 10-step debugmodel run tracks the previous losses to within 3e-3 with matching grad norms. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

Follow-up review pass over the Kimi K3 change. - kimi_k3/model.py imports FLA at module scope, which is a per-model dependency rather than a core one. Without it the three Kimi test modules failed collection, and pytest treats a collection error as fatal, so `pytest tests/unit_tests` aborted entirely on a machine that had not installed it. Guard the imports and raise unittest.SkipTest instead, as test_qwen3_5_deltanet.py already does for the same reason. - Move the scatter_vision_embeds test to tests/unit_tests/test_multimodal.py. It covers a helper shared by three VLMs, so it does not belong behind the Kimi FLA guard, and models/common/multimodal.py had no tests of its own. - Cover the attention-residual pass-through path in the FSDP comparison. _small_model_config used attn_res_block_size=1, so every layer extended the residual and no layer passed it through. Layers that pass it through return it back out across the FSDP module boundary, and that routing is what keeps FSDP gradients bitwise equal to eager: returning None instead reassociates the residual's gradient accumulation and moves tok_embeddings.weight.grad by ~2e-3 relative, which measurably breaks the eager comparison. Run the comparison at attn_res_block_size=2 and record the reasoning where the tempting change would be made, since the arrangement exists only to satisfy a constraint the previous config could not observe. - Reject a spatial_merge_size that disagrees with the vision encoder's merge_kernel_size. The two are independent config fields that must match; a mismatch previously surfaced as a placeholder-run misalignment error that blamed the prompt rather than the configuration. - Correct the README: the token embedding and the output projection are separate parameters, not tied. Drop a stale "device-neutral" backend message and a note about qwen3_5 that will silently rot once that model is fixed. Verified with PyTorch 2.14.0.dev20260729+cu130 and fla-core 0.5.2: the frozen HuggingFace parity values are unchanged, the debugmodel reproduces 12.55150 / 12.33794 / 11.61755 over three deterministic steps, tests/unit_tests is 562 passed with the 18 failures all pre-existing missing-dependency ones, and pyrefly reports no errors over the changed files. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

FLA's chunked KDA kernel fails below head_dim 16 because chunk_intra.py sets BK = next_power_of_2(K) without the floor of 16 that wy_fast.py applies, and triton's tl.dot needs a contraction dimension of at least 16. That is an unhandled small-head case in FLA rather than a constraint worth encoding here: released KDA head dimensions are far above it, the reduced flavor uses 32, and the guard would have to be revisited once FLA clamps the block size. Report it upstream instead of carrying a local check. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

JavaZeroo · 2026-08-01T09:18:59Z

I would like to implement a version of fsdp with correct precision first. other parallel implementations on k3 may require more adaptation work, especially cp and pp. i think it is possible to quickly support a version of fsdp for k3 first, enough to support our experiments with small models. In the meantime I'm happy to be able to add ep tp support to this pr, do you see any need to add it in this pr?

Support fsdp first sounds good to me. I think we should target for a training ready version using kernels for kda, one option is to use from fla.ops.kda import chunk_kda (similarly we used fla kernels for qwen 3.5, and put the pytorch native reference implementation into the numerical tests). We may have plan to use our own kernel for kda in the future.

cc: @tianyu-l

Thanks @shuhuayu, I have already made the changes you requested.

Router is TokenChoiceTopKRouter now, and the vision scatter uses the shared
scatter_vision_embeds. The vision block is not folded into VisionTransformerBlock yet,
because that one hardcodes LayerNorm for the norms and BlockMask for the mask, and in eager
the padding rows would softmax to NaN. Reusing it would be a fairly big change. The MoE is
not reused either, I kept the for-loop form. Should that switch to grouped_mm, or stay eager?
Added test_kimi_k3_hf_parity.py, which freezes fp32 outputs from the released hf code, covering text logits, routed expert ids, vision features and multimodal logits.
Fixed.
kda uses chunk_kda now, with the gate, beta sigmoid and qk l2norm fused into the kernel,
same split as qwen3.5. The torch recurrence moved into the tests as the reference, and I
added a cuda test comparing fwd/bwd.

…ation pytorch/torchtitan#4025 adds Kimi K3 upstream, and checking it against the reasons this branch existed inverts the argument. It constructs nn.Linear positionally rather than through a config tree -- it has config dataclasses but does not declare child Linear.Config fields or build them -- so "return to the titan standard" was never true; upstream's own K3 does what ours does. And it supports FSDP2 only, explicitly rejecting HSDP, TP, PP, CP, EP, activation checkpointing, torch.compile and CPU offload, with the author noting TP/PP/CP would need significant adaptation because of data-dependent Python loops and incompatible forward signatures. So the parallelism work upstream declines to do is exactly what this fork has: 14/14 matrix legs producing loss, PP verified per-parameter at 0.00000 over 548 parameters, and two TP defects found and fixed, one of which also fixes upstream deepseek_v3. Refactoring toward a style upstream does not use, at the cost of breaking that, is the wrong trade. The cost was measured rather than estimated: converting three MLA linears to Linear.Config(...).build() failed 12 of 14 legs with silent exit=0 hangs. The branch keeps its two gated commits in case the LoRAConverter question returns. Work moves back to finishing veRL. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01WBy1d9YVu44nYCVqykRqL1

QIU023 · 2026-08-03T21:31:36Z

Thanks folks for connecting this PR with the broader RFC I have opened for the entire K3 support, I will review this asap in recent days and get it landed and aligned with my broad parallelism for all the reasonable interfaces needed, post-training and QAT support on top of this model backbone

torchtitan/distributed/fsdp.py already has apply_fsdp_to_vision_encoder. This folder carried its own apply_fsdp_vision, a 48-line duplicate of it that no caller ever reached, so the tower rode along inside the root wrap fully replicated on every DP rank. Invisible at the debug tower's 4 layers / hidden 256; not an option at MoonViT-V2's real 447.4M against k3mini's 80.9M text side, where the encoder is 5.5x the model it serves. Deleted the duplicate and called the core helper before the decoder, as its docstring asks. This also matches how pytorch#4025 wires the same thing, so the rebase is a deletion rather than a merge. Vendored add_zero_valued_dependency from that PR verbatim, with a note to drop it when the PR lands. It covers a hazard our own CP fix does not: FSDP2 issues the tower's all-gather from its pre-forward hook and its reduce-scatter from the output's autograd hooks, so once the tower is actually sharded, a rank that skips it desynchronizes the process group. Our fix only aligned our own all_reduce. One trap on the way: with the tower sharded its params are DTensors too, so encode_images' "is the weight a DTensor" test no longer distinguished TP's replication from FSDP's sharding. It lifted the input onto the FSDP mesh, where it met the plain all-gathered weight inside the conv. parallelize now records the tp mesh explicitly instead. 12/12 multimodal legs, 10 steps, seed 42 deterministic: bit-identical to the unsharded run on every leg (mm_fsdp2 7.73923 -> 5.32836 ... mm_ep2_fsdp2_pp2_cp2 7.71223 -> 5.26428). Vision confirmed live. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01WBy1d9YVu44nYCVqykRqL1

…multimodal Evidence index for the update to pytorch/torchtitan#3029, covering what is implemented and reproduced today rather than proposed: 13 text parallelism combinations and 12 multimodal ones, 10 steps each, seed 42 deterministic, all monotone; PP8xVP4 at |Dloss| 0.0018 against the no-PP reference; CP built on fla's merged KCP (fla-org/flash-linear-attention#691) rather than a private recurrence. Records the defects alongside, because each one passes every check that reads a loss curve: the Block AttnRes 1/tp over-reduction, the moe_sharding in_grad_placements drop that also reproduces on unmodified deepseek_v3, the non-autograd-aware conv halo that left ~60% gradient error on W-1 boundary tokens while the forward stayed bit-exact, and ten multimodal defects of which six silently reverted forward to its text-only branch. Also states what the matrices are NOT: bf16 with fp32 reduction, no QAT. K3's MXFP4 is post-training only -- the report puts QAT across SFT and RL, not pretraining -- so a pretraining-shaped matrix should not carry it. The kimi_k3_mini_qat_mxfp4 flavor implements the released scheme separately. Open gaps stated rather than omitted: LoRA's TP gradient defect (ratio up to 2.26 at tp4 on the rowwise lora_b, invisible to cold-seed checks because B is zero at init), and the report's sec 5.2.3 encoder optimizations. Refs: pytorch/torchtitan#3029, pytorch/torchtitan#4025 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01WBy1d9YVu44nYCVqykRqL1

Compiled 13-leg matrix: 11 pass with a worst eager-vs-compiled delta of 0.013, against the 0.10-0.40 spread the parallelism configurations show among themselves -- compile is numerically fine where it runs. Two fail, both EP with pipeline parallel, on _grouped_mm receiving a [224, 0] operand. First reading was that this is a core limitation: the call site, models/common/moe.py:95/101/106, is byte-identical in this fork and in #4025's tree, and has no empty-group guard. That reading is wrong. The rest of the file is not identical -- this fork rewrote the routing-map scatter under TP+EP (129e29de0), and that map determines the group boundaries _grouped_mm is handed. Control on #4025's tree, which carries upstream's unmodified moe.py: deepseek_v3_debugmodel at dp2 x ep2 x tp2 x pp2 with --compile.enable passes (loss 8.13452). Same call site, same parallelism, same compile flag. So the defect is in this fork, and the routing-map change is the prime suspect. Also records that #4025 declares torch.compile out of scope and defaults CompileConfig to enable=False, so the published comparison stays compile-off on both sides and needs no adjustment. Not fixed. Next step is to instrument num_tokens_per_expert_E under the failing configuration and find which expert goes empty, rather than adding a guard that hides the cause. Refs: pytorch/torchtitan#3029, pytorch/torchtitan#4025 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01WBy1d9YVu44nYCVqykRqL1

kimi_k3_debugmodel_pr_4025 mirrors pytorch#4025's debugmodel architecture exactly -- 13 layers at dim 256, 4 heads, q_lora 128 / kv_lora 64, qk_nope 32 / qk_rope 16 / v 32, full attention on {4, 8, 12} with KDA elsewhere, AttnRes block 12, LatentMoE latent 128 / 8 experts top-2 / 2 shared, vocab 163840, and a 4-layer 3-head MoonViT at dim 256 / qkv 384 / hidden 1024. Same model on both sides, so the comparison is our parallelism against theirs rather than two different debug models. The first version inherited k3mini's kda_layers, a 15-entry list, into a 13-layer model -- two descriptions of the same stack contradicting each other. Deriving it from full_attn_layers fixes ep2_fsdp2, which now runs all 5 steps. Verified: FSDP2 runs with vision live (20 encode_images calls, 30/30 tower parameters with gradients), starting loss 12.06 against pytorch#4025's own 12.48 on the same vocab. Refs: pytorch#3029, pytorch#4025 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01WBy1d9YVu44nYCVqykRqL1

Two desynchronizations on the PR-4025 twin flavor under CP, both surfacing as a 100-second NCCL watchdog timeout rather than an error. The sentinel-count all_reduce sat after forward's `pixel_values is None` early return, so a rank whose batch happened to carry no images returned without entering it while its CP peers waited there forever (NumelIn=2 on mesh_cp). Hoisted to the top of forward, gated on cp_world_size > 1 -- a property every rank agrees on before looking at any data. Second, now that the tower is FSDP-sharded, skipping it also skips the all-gather FSDP2 issues from its pre-forward hook (_ALLGATHER_BASE, NumelIn=10486144 on mesh_fsdp). An image-free batch now runs the tower on a minimal placeholder and keeps the graph edge through add_zero_valued_dependency, so every rank issues the same collectives and the tower's contribution to the data-parallel average is a correct zero. That is the hazard pytorch#4025 added that helper for, reached here by a second route. Both are real and both are fixed. They are NOT sufficient: fsdp2_tp2_cp2 and ep2_fsdp2_tp2_cp2 still hang at step 2 on the same NumelIn=2 all_reduce, so a third path leaves a rank out of it. Ruled out: it is not KCP (that is fla's KDA recurrence, not this collective) and not the sentinel-count assertion (which never fires in the logs). Next step is per-rank instrumentation of the entry to _exchange_sentinel_counts rather than more hypotheses. No regression: fsdp2 on the twin flavor is bit-identical (12.05716 12.04941 12.04791 11.98434 11.78795), and ep2_fsdp2 -- which the kda_layers fix repaired -- still runs all 5 steps. Refs: pytorch#3029, pytorch#4025 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01WBy1d9YVu44nYCVqykRqL1

Per-rank instrumentation of the sentinel-count exchange on the PR-4025 twin. The sharding premise holds: rank 0 and rank 2 are a CP pair reporting local counts 255 and 34, summing to 289 -- exactly 17x17, one 34x34-patch image after 2x2 merge. Each rank does hold a complementary slice. What does not hold is the number of times the exchange runs. forward executes several times per step over different microbatches (pixel_values of 1120, 1140, 1156 and 1092 patches were observed), and the entry counts differ between ranks within a step. A collective whose count differs across participants hangs the same way as one whose participants differ, which is why fixing the two data-dependent entry conditions was necessary but not sufficient. So the remaining defect is in how many times a per-forward collective runs relative to the microbatch loop, not in which slice a rank takes. Refs: pytorch/torchtitan#3029, pytorch/torchtitan#4025 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01WBy1d9YVu44nYCVqykRqL1

…aths State file so the diagnosis can resume without re-deriving it: the failing command, the two collectives that time out, the four hypotheses killed by measurement (KCP, the sentinel assertion, the shard arithmetic, the call counts), the two defects fixed on the way, and the exact next probe -- flush per collective rather than per step, so a partial step-2 trace survives. Also answers the question the twin's failure raised about the published 12/12 multimodal matrix: if the same code paths hang there, was that result luck? kimi_k3_mini_vl at dp2 x tp2 x cp2 runs 30 steps clean (7.73550 -> 2.78104), three times the published horizon, on the leg most likely to be fragile. So the difference between the two flavors is configuration, not chance, and the 12/12 holds. Which configuration difference triggers it is still open. max_patches and seq_len are identical in both, so the obvious candidate is out; what remains is vocab 2020 vs 163840, dim 512 vs 256, 21 vs 13 layers, 15 vs 10 KDA layers, and local_batch_size. Bisecting the twin one field at a time toward k3mini is ~2 minutes per run. Refs: pytorch/torchtitan#3029, pytorch/torchtitan#4025 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01WBy1d9YVu44nYCVqykRqL1

… k3mini too Bisected the twin flavor toward kimi_k3_mini_vl one field at a time. The difference was not in the model at all: the published multimodal matrix passed --training.local-batch-size 4, the twin matrix did not, and both flavors default to 1. At global batch 8 over dp2 that is one forward per step versus four gradient-accumulation microbatches -- exactly the forward=4 the per-rank probe recorded. kimi_k3_debugmodel_pr_4025 local_batch 4: 5 steps pass local_batch 1: hangs kimi_k3_mini_vl local_batch 4: 30 steps pass local_batch 1: hangs So the defect reproduces on kimi_k3_mini_vl as well. The published 12-leg multimodal matrix did not exercise it because that run's local batch was large enough to avoid accumulation entirely. That qualifies the published number: it is "passes without gradient accumulation", not "passes", and the qualification has to travel with it -- accumulation is standard at any real scale. Why accumulation breaks it is still open. The shard arithmetic is correct and step-1 call counts match across all eight ranks, so the suspect is state carried across microbatches within a step. Refs: pytorch/torchtitan#3029, pytorch/torchtitan#4025 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01WBy1d9YVu44nYCVqykRqL1

…tuple Refs: pytorch#3029 A non-last PP stage returns (hidden_state, block_residuals) -- the AttnRes adapter ships the block payload alongside the activation -- so handing that straight to add_zero_valued_dependency raised AttributeError: 'tuple' object has no attribute 'dtype'. Both tower-alive call sites did it: the image-free path (latent, never reached by a passing leg) and the zero-sentinel CP path added in the previous commit, which is what surfaced it. Route both through a local helper that puts the graph edge on the hidden state and rebuilds the tuple, the same thing the adapter's own _keepalive_touch does. Kept out of add_zero_valued_dependency so that helper stays byte-identical to pytorch#4025's and the rebase stays a clean delete. Twin-flavor multimodal matrix, 3 steps, seed 42, deterministic: the three PP+CP legs go from FAIL to passing, taking the matrix to 10/13. tp2_pp2_cp2 12.07205 12.02185 11.98569 fsdp2_pp2_cp2 12.05744 12.03855 11.97321 ep2_fsdp2_pp2_cp2 12.06617 12.02984 11.96501

…its architecture Refs: pytorch#3029, pytorch#4025 The flavor reproduced pytorch#4025's model exactly but kept k3mini_vl's training and data settings, which left three differences that change what actually runs: training.dtype float32 (inherited) -> bfloat16 (pytorch#4025 sets this) max_patches 1024 / 64 per side -> 256 / 16, max_pixels 224x224 seq_len 8192 chain default -> 256, local_batch_size 1 dtype is the one with teeth. training.dtype applies to the model itself, while mixed_precision_param only reaches parameters through FSDP -- so on a layout with no FSDP (dp_shard 1 and no CP, since torchtitan's FSDP mesh is dp_shard x cp) the twin ran KDA on fp32 operands. That kernel asks for 108160 bytes of dynamic shared memory against the 101376 this GPU permits, so dp1, pp2 and tp2 died where fsdp2 and cp2 passed. That was diagnosed as a hardware ceiling. It is a hardware ceiling only for a configuration pytorch#4025 does not use: predicted that forcing --training.mixed_precision_param float32 onto the passing fsdp2 leg would reproduce it, and it does, with the identical 108160. Batch size is not involved -- local_batch_size 1, 2 and 4 all request the same. The image budget mattered for a different reason: at 1024 patches one image nearly fills the sequence, which is a different data distribution from the one that PR trains on. local_batch_size 1 against global batch 8 also makes gradient accumulation the default here, which is the configuration that exposed the zero-sentinel CP defect. The matrix should run it rather than avoid it.

Refs: pytorch#3029, pytorch#4025 Both were declared with an explicit dtype=torch.float32 while every other parameter follows the default. That reads like a precision decision but is an artifact: the fp32 was on the empty() that the uniform_ init draws into, and the parameter inherited it. The comment above it only discusses shape. It becomes load-bearing under training.dtype=bfloat16, where the other 380 parameters are bf16 and these 20 are not, and FSDP2 refuses outright: "FSDP expects uniform original parameter dtype but got {bfloat16, float32}". That blocked nine of the thirteen twin legs. The init math still happens in fp32 and is then stored at the default dtype, so under the fp32 default this is a no-op. Verified as one rather than argued: kimi_k3_mini_diag_4l_moe_depth at dp2, 3 steps, seed 42, deterministic, run with and without this change on the same tree -- with: 7.67804 7.27652 6.28913 without: 7.67804 7.27652 6.28913 Open question against pytorch#4025, recorded in KDA_GATE_DTYPE_2026-08-04.md rather than silently resolved in their favour: their debugmodel sets training.dtype=bfloat16, which puts A_log and dt_bias in bf16, and fla's gate kernel computes the forward in fp32 but casts the gradient back to the parameter's own dtype (ops/kda/gate.py: dA.sum(0).view_as(A_log).type_as(A_log)). We measure those gradients two to three orders below the model median, so bf16 storage rounds them where fp32 would not. Worth raising on that PR.

…inst #4025 Refs: pytorch/torchtitan#3029, pytorch/torchtitan#4025 Four documents. CP_MULTIMODAL_HANG_RESOLVED: the zero-sentinel CP shard dropped the vision tower out of the loss graph, so that rank issued one fewer FSDP reduce-scatter than its peers. Found by diffing per-rank collective traces rather than by reading stacks, which is why two earlier probes could not localize it and one published a wrong conclusion. Corrects the previous attribution to gradient accumulation: accumulation is the trigger only because it shrinks the microbatch to a single sequence. TWIN_FIDELITY: the #4025 twin flavor copied that PR's architecture but inherited k3mini_vl's training and data config, which is not the same thing. training.dtype in particular was left at float32, so three legs ran KDA on fp32 operands and hit this GPU's shared-memory ceiling. Recorded earlier as an unavoidable hardware limit; it was ours. KDA_GATE_DTYPE and REVIEW_4025_FINDINGS: things in #4025's tree we think are wrong or worth questioning, kept apart from the interface asks so they can be sent separately. The load-bearing one is get_vision_positions in models/common/multimodal.py, which raises on both of the states a CP rank normally sees -- fewer placeholder runs than items, and a run truncated at a shard boundary. It sits in models/common, so it becomes the contract for every multimodal model upstream, which is why it is worth raising before our CP PR rather than after. Standing rule recorded there: upstream being upstream does not make it right. Two changes here were made to match their tree and one of them we now think is the wrong direction on the merits.

… and not #4025's Refs: pytorch/torchtitan#3029, pytorch/torchtitan#4025 Written up as "likely blocking for CP, worth raising against #4025". Both halves were wrong. Provenance: models/common/multimodal.py came in with #3532 ("[kimi k2_7] add kimi k2_7", 2026-07-29) and is already on main. #4025 is its third consumer after qwen3_5 and kimi_k2_7, not its author. Classification: all three consumers refuse context parallel with NotImplementedError, and two name this exact reason in the error text -- qwen3_5 "multimodal CP needs vision scatter before CP sharding", kimi_k2_7 "vision scatter needs the full sequence before CP would shard it". The whole-sequence assumption is deliberate and the guard exists so nobody reaches the helper with a shard. Nothing is broken. So the comment becomes "here is the gap you documented, filled" rather than "here is a fault", which is both accurate and a better position to argue from. Also records that their intended design is not ours: they describe scattering before CP shards, we shard first and select each rank's slice by a prefix sum over per-rank sentinel counts. Theirs needs the vision encode to happen before prepare_context_parallel_input, i.e. in the trainer rather than the model. The PR should say that outright instead of quietly shipping the other split. Standing note added at the top: check git log --diff-filter=A before attributing a file to this PR. Two of the three findings here are in code it only consumes.

…/13 compiled Refs: pytorch/torchtitan#3029, pytorch/torchtitan#4025 Run at the flavor's own settings rather than CLI overrides, so this is #4025's configuration -- including local_batch_size 1 against global batch 8, which makes gradient accumulation the default and exercises the zero-sentinel CP path that used to deadlock. Eager 13/13, and two independent full matrix runs agree bit-for-bit on all thirteen legs, which is a stronger claim than a single pass. Step-1 spread 0.041 across every parallelism combination, all just above ln(163840)=12.006. tp2_pp2_cp2 reported 0/3 in one earlier run and passed three times since with identical losses. Recorded as a transient launch failure with the retries shown, rather than dropped or called flaky without evidence. Compiled 10/13. Worst eager-vs-compiled delta on the ten that run is 0.0086, inside the 0.041 spread across layouts. The three failures are all EP and all one cause: torch._grouped_mm rejects a zero-length contraction dim. Not worked around -- a guard in models/common/moe.py would turn them green and hide it. States plainly what this is not: no numerical comparison against #4025's own tree, since that tree refuses TP/CP/PP and has one runnable cell.

QIU023 · 2026-08-04T09:06:09Z

Hi @JavaZeroo , thanks for raising initial changes for K3 reference models as a good start point! I have reviewed this change and only have one major comment about model architecture repro for K3:

Review comment 1 — anchor at `init.py` (`full_attention_layers`)

Should full_attention_layers include the last layer? The released config.json lists full_attn_layers explicitly as [4, 8, 12, ..., 84, 88, 92, 93] -- 24 entries, with 92 and 93 both global -- which matches sec 2.1's "an additional Gated MLA layer is placed at the end of the backbone, ensuring that the final layer always performs global attention". As written here, {4, 8, 12} over 13 layers ends the stack on KDA.

Raising it at config level rather than as a debug-model nit: the released list is not expressible as "every (ratio+1)-th layer", so a 2.8T config needs either the trailing layer appended or the list carried verbatim.

Review comment 2 — anchor at `distributed/fsdp.py` (`add_zero_valued_dependency`)

Strong agree with this helper. One field note from running this architecture under CP: the trigger is not only "a batch that happens to carry no images" -- under context parallelism a rank's sequence shard can hold zero vision sentinels even when every rank received images, which reaches the same subset-collective deadlock by a route that only appears at cp>1. Might be worth a sentence in the docstring, since the CP route is easy to miss when reading from the DP example. Our CP follow-up includes a regression test for exactly this route.

QIU023 · 2026-08-04T10:10:11Z

In my forked repo, the same K3 debug model architecture is already running under broad parallelism combinations (except full 5D parallelism yet as I need to get a 16 GPUs machine or a multi-node setup), including FSDP2, PP, EP, TP and CP on our fork -- same matrix, same seeds, results in the RFC update linked above. Although stacking those onto this folder may need minor changes on top of the current interfaces for PP and CP support, and I will raise PRs for further discussion with maintainers and @JavaZeroo during rebasing for parallelism supports, once this #4025 got merged.

Please check the details in RFC latest content on top and let me know all of your thoughts, thanks a lot!

The released config lists full_attn_layers as [4, 8, ..., 88, 92, 93]: 92 and 93 are both global, so the backbone always ends on a full-attention layer. The debug model stopped at {4, 8, 12} over 13 layers and closed on KDA instead. That list is not expressible as "every n-th layer", so the topology is now assembled by _kimi_k3_config, which takes the 1-based indices verbatim. A future full-scale flavor can pass the released 24-entry list unchanged. 13 is also one past the attention-residual block size of 12, so the short trailing block the released 93-layer stack has is now covered too. Also records the float32 parity run at the released head dimensions, which needs two independent TF32 switches closed: the cuBLAS one that NGC containers enable through TORCH_ALLOW_TF32_CUBLAS_OVERRIDE, and TRITON_F32_DEFAULT, which the torch flag does not reach and which FLA's chunk_kda relies on. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

The docstring only described the data-parallel trigger, where a rank's batch carries no images. Context parallelism reaches the same subset-collective deadlock by a second route: a rank's sequence shard can hold zero vision placeholders even when every rank received images. Reported by QIU023 from running this architecture at cp>1. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

JavaZeroo · 2026-08-04T12:23:32Z

Hi @JavaZeroo , thanks for raising initial changes for K3 reference models as a good start point! I have reviewed this change and only have one major comment about model architecture repro for K3:

Review comment 1 — anchor at __init__.py (full_attention_layers)

Should full_attention_layers include the last layer? The released config.json lists full_attn_layers explicitly as [4, 8, 12, ..., 84, 88, 92, 93] -- 24 entries, with 92 and 93 both global -- which matches sec 2.1's "an additional Gated MLA layer is placed at the end of the backbone, ensuring that the final layer always performs global attention". As written here, {4, 8, 12} over 13 layers ends the stack on KDA.

Raising it at config level rather than as a debug-model nit: the released list is not expressible as "every (ratio+1)-th layer", so a 2.8T config needs either the trailing layer appended or the list carried verbatim.

Review comment 2 — anchor at distributed/fsdp.py (add_zero_valued_dependency)

Strong agree with this helper. One field note from running this architecture under CP: the trigger is not only "a batch that happens to carry no images" -- under context parallelism a rank's sequence shard can hold zero vision sentinels even when every rank received images, which reaches the same subset-collective deadlock by a route that only appears at cp>1. Might be worth a sentence in the docstring, since the CP route is easy to miss when reading from the DP example. Our CP follow-up includes a regression test for exactly this route.

all fixed

JavaZeroo · 2026-08-05T07:42:03Z

@QIU023 @shuhuayu @tianyu-l Thanks for reviewing. Is there any further work here? Do I need to run ci?

QIU023 · 2026-08-05T08:53:23Z

Looks fine to me, looking forward to rebase all my parallelism supports and other further supports on it! Please wait for maintainers to approve and keep me posted once this is merged so I will raise subsequent PRs very soon

JavaZeroo · 2026-08-07T03:37:27Z

@pytorchbot rerun

pytorch-bot · 2026-08-07T03:37:30Z

❌ 🤖 pytorchbot command failed:

@pytorchbot: error: argument command: invalid choice: 'rerun' (choose from 'merge', 'revert', 'rebase', 'label', 'drci', 'lint', 'fix-lint', 'apply-lint', 'cherry-pick')

usage: @pytorchbot [-h]
                   
                   {merge,revert,rebase,label,drci,lint,fix-lint,apply-lint,cherry-pick}
                   ...

Try @pytorchbot --help for more info.

Add Kimi K3 reference model

a2c160f

pytorch-bot Bot added the ciflow/8gpu label Jul 30, 2026

meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label Jul 30, 2026

JavaZeroo added 2 commits July 30, 2026 16:07

Remove development-only Kimi K3 diagnostics

7d8667d

Add FSDP2 support for Kimi K3

85eba0b

JavaZeroo changed the title ~~[kimi k3] add Kimi K3 reference model~~ [kimi k3] add eager reference model with FSDP2 Jul 30, 2026

JavaZeroo mentioned this pull request Jul 30, 2026

[RFC] Kimi K3 Architecture support for torchtitan based pre-training and post-training (with veRL) #3029

Open

JavaZeroo added 3 commits July 30, 2026 20:33

Scale Kimi K3 debug model

956249f

Use out-of-place operations in Kimi K3 eager paths

cb8d92c

Keep unused Kimi experts in the FSDP autograd graph

eae7ea9

JavaZeroo force-pushed the agent/add-kimi-k3-reference-model branch from 3c6de86 to eae7ea9 Compare July 30, 2026 14:45

tianyu-l requested a review from shuhuayu July 30, 2026 20:44

JavaZeroo and others added 4 commits July 31, 2026 16:20

Add Kimi K3 numerical and mixed-modality FSDP coverage

61de88a

JavaZeroo marked this pull request as ready for review August 1, 2026 09:19

JavaZeroo requested review from fegin, tianyu-l, wconstab and wwwjn as code owners August 1, 2026 09:19

JavaZeroo and others added 2 commits August 4, 2026 11:55

QIU023 approved these changes Aug 5, 2026

View reviewed changes

Uh oh!

Conversation

JavaZeroo commented Jul 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Why

Validation

Numerical parity with the released implementation

Training

FSDP2

Initial scope

Uh oh!

pytorch-bot Bot commented Jul 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

shuhuayu commented Jul 31, 2026

Uh oh!

shuhuayu commented Jul 31, 2026

Uh oh!

claude Bot commented Jul 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Code Review: Kimi K3 eager reference + FSDP2

1. Reuse of common modules (maintainer Q1) — largely agree, several concrete duplications

2. Reproducible numerical parity test (maintainer Q2) — missing, should be added

3. Bottlenecks for other parallelisms (maintainer Q3)

4. Vision encoder FSDP + conditional execution (maintainer Q4) — real hang risk, agree

Other findings

Style / minor

Uh oh!

JavaZeroo commented Jul 31, 2026

Uh oh!

shuhuayu commented Jul 31, 2026

Uh oh!

JavaZeroo commented Aug 1, 2026

Uh oh!

QIU023 commented Aug 3, 2026

Uh oh!

QIU023 commented Aug 4, 2026

Review comment 1 — anchor at __init__.py (full_attention_layers)

Review comment 2 — anchor at distributed/fsdp.py (add_zero_valued_dependency)

Uh oh!

QIU023 commented Aug 4, 2026

Uh oh!

JavaZeroo commented Aug 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Review comment 1 — anchor at __init__.py (full_attention_layers)

Review comment 2 — anchor at distributed/fsdp.py (add_zero_valued_dependency)

Uh oh!

JavaZeroo commented Aug 5, 2026

Uh oh!

QIU023 commented Aug 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

JavaZeroo commented Aug 7, 2026

Uh oh!

pytorch-bot Bot commented Aug 7, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

JavaZeroo commented Jul 30, 2026 •

edited

Loading

pytorch-bot Bot commented Jul 30, 2026 •

edited

Loading

claude Bot commented Jul 31, 2026 •

edited

Loading

Review comment 1 — anchor at `init.py` (`full_attention_layers`)

Review comment 2 — anchor at `distributed/fsdp.py` (`add_zero_valued_dependency`)

JavaZeroo commented Aug 4, 2026 •

edited

Loading

Review comment 1 — anchor at `init.py` (`full_attention_layers`)

Review comment 2 — anchor at `distributed/fsdp.py` (`add_zero_valued_dependency`)

QIU023 commented Aug 5, 2026 •

edited

Loading