feat: add Kimi K3 model support

[HuiyingLi]

Summary

• add a standalone native Kimi K3 text-model implementation without requiring trust_remote_code
• support the K3 text decoder, KDA and MLA attention, 896-expert MoE, tokenizer/encoding, MXFP4 state loading, and context-parallel helpers
• add K3-aware pipeline partitioning that keeps each attention-residual block within one stage and carries the residual accumulator across stages
• register K3 model, config, and tokenizer classes with AutoModel
• add the text-only examples/llm_finetune/kimi/k3_hellaswag.yaml EP32/PP8 recipe
• add CP support

Runtime Changes

• allow grouped experts to apply router weights after expert down projection, matching K3 MoE semantics
• isolate strict-FP32 KDA parameter holders into homogeneous FSDP units
• select a CUDA/PP-mesh device for gradient clipping on pipeline stages without local gradients

Validation

• ruff check: passed
• ruff format --check: passed
• focused unit tests: 91 passed
• pipeline tests cover exact 12-layer attention-residual boundaries, residual handoff, partial checkpoint loading, and forward/backward parity
• EP32/PP8 HellaSwag training, Slurm job 5639240: completed 100 steps, loss 2.0011 -> 1.4307
  • W&B run olnfgmh1

• CP1/CP2 parity
  cp1: https://wandb.ai/Nemo-automodel/huiyingl_workspace/runs/fcguk1ys
  cp2: https://wandb.ai/Nemo-automodel/huiyingl_workspace/runs/n7n3dtc0

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[HuiyingLi]

/claude review

[claude]

Review summary

Light review of the Kimi K3 model-onboarding PR (17 files). Registry wiring (MODEL_ARCH_MAPPING, _CUSTOM_CONFIG_REGISTRATIONS for kimi_k3/kimi_linear), tokenizer registration, config classes, PP block-boundary partitioning, CP sharding, and the state-dict adapter (KDA fp32-holder routing, MXFP4 dequant) all look coherent and are covered by focused tests. RoPE fp32 handling is correct: KimiRMSNorm and KDA gates compute in fp32, _keep_in_fp32_modules lists _fp32_params, and initialize_weights casts via cast_model_to_dtype rather than a raw self.to(dtype). ModelCapabilities correctly declares cp/pp/ep support matching the wired parallelism.

A few points worth addressing (all posted inline):

1. Missing MoE forward coverage (test_pipeline_parallel.py) — every test config sets first_k_dense_replace=num_hidden_layers + 1, so no MoE layer is ever built. The routed-expert path (KimiK3MoE, KimiK3Gate, _forward_reference_order) and the new apply_router_weight_after_down index math added to the shared moe/experts.py have no numerical test.

2. Import-time monkeypatch of transformers.activations (vision.py:42) — a process-global side effect on module import.

3. Dead fallback import (tokenization.py:19) — from encoding_k3 import ... references a nonexistent module.

The moe/experts.py change defaults apply_router_weight_after_down=False, so existing MoE models are behavior-preserved. The two encoding.py/vision.py/tokenization.py files carry Moonshot/HF upstream markers and ruff: noqa for the vendored portions, which is acceptable for backported reference code.

[claude]

Import-time monkeypatching of a third-party (transformers) global. activations.PytorchGELUTanh = GELUTanh mutates transformers.activations for the whole process the moment this module is imported (it is imported transitively by multimodal.py, the registered VLM). Even though it only fires on older transformers where the symbol is missing, a process-global side effect at import time is fragile — the alias leaks into any other code that inspects transformers.activations. Prefer binding a module-local name only (PytorchGELUTanh = GELUTanh) without writing back onto the activations module.

[claude]

The except ImportError fallback imports from encoding_k3, but no such module exists anywhere in the package (the sibling module is encoding.py). If the relative import ever fails, this raises ModuleNotFoundError rather than a working fallback, so the branch is dead/misleading. Either drop the fallback (the module is only ever imported as part of the package) or point it at the correct name.

[claude]

first_k_dense_replace=num_hidden_layers + 1 forces every layer to be dense, so no KimiK3MoE layer is ever instantiated in the test suite. Combined with test_state_dict_adapter.py (which only exercises key mapping / MXFP4 dequant), the routed-expert forward path has no numerical coverage: KimiK3MoE.forward, KimiK3Gate (fp32 sigmoid + group top-k), _forward_reference_order, and — most importantly — the new apply_router_weight_after_down grouped-experts logic added to the shared moe/experts.py are all untested.

The experts.py change reshapes the scatter output to [tokens, topk, hidden], applies the fp32 router weight after the down projection, then sums over topk, with new slot_id scatter indexing in both _forward_loop and _forward_grouped_mm. This is subtle index math on a shared component consumed by every MoE model. Please add a small CPU unit test that builds a tiny K3 model with first_k_dense_replace low enough to create at least one MoE layer and asserts a finite forward (and ideally parity of apply_router_weight_after_down=True vs. the equivalent before-down weighting on the reference loop path).

[HuiyingLi]

/claude review

[claude]

Light review — Kimi K3 (KDA/MLA hybrid MoE) onboarding. Read the full diff for all 17 changed files (excluding uv.lock); revision confirmed current at head 38031cd.

The model package is well-structured: registry wiring (MODEL_ARCH_MAPPING + _CUSTOM_CONFIG_REGISTRATIONS) is present, ModelCapabilities correctly declares cp/pp/ep for the new parallelism paths, RoPE/decay tables are kept in fp32 via _keep_in_fp32_modules and cast_model_to_dtype, tensor-layout docstrings are thorough, and the PP/state-dict/MXFP4 paths have focused CPU tests.

Two findings posted inline:

1. thinking_effort="medium" inconsistency (encoding.py): the emitted system message advertises medium as supported but _VALID_THINKING_EFFORTS omits it, so that value trips an assert; the assert is also user-input validation that -O strips.
2. Missing config-resolution test (registry.py): the new kimi_k3/kimi_linear model_type registrations lack the focused AutoConfig/resolve_custom_config_cls test that proves the local config resolves from a checkpoint config.json, including the stale-Transformers case the registry exists to handle.

[claude]

"medium" is advertised as a supported thinking_effort in the emitted system message ("supported values include low, medium, high, and max" at line 511), but _VALID_THINKING_EFFORTS = {"low", "high", "max"} omits it. Passing thinking_effort="medium" therefore trips this assert and raises AssertionError, and even if it passed the assert the elif thinking and thinking_effort in _VALID_THINKING_EFFORTS guard at line 505 would silently drop the effort message. Either add "medium" to _VALID_THINKING_EFFORTS or remove it from the advertised list so the two agree.

Separately, this is user-supplied input validated with assert, which is stripped under python -O; prefer raising a ValueError that names the invalid value.

[claude]

New _CUSTOM_CONFIG_REGISTRATIONS entries kimi_k3 and kimi_linear are added, but there is no focused test proving that AutoConfig/resolve_custom_config_cls resolves KimiK3Config/KimiK3TextConfig from a checkpoint-style config.json (model_type: "kimi_linear" / "kimi_k3"). Since Automodel owns these config classes and the K3 checkpoint advertises kimi_linear (a model_type not guaranteed present in the installed Transformers CONFIG_MAPPING), the exact failure mode the registry guards against — a stale Transformers release where AutoConfig.from_pretrained cannot find the type — is untested. Please add a test that writes a config.json with these model_types and asserts the local config class is resolved (covering the builtin-absent case).

[HuiyingLi]

/claude review

[claude]

Reviewed all 18 changed files (excluding uv.lock). This is a large, well-structured new-model onboarding (Kimi K3): registry + custom-config registrations are complete and covered by a stale-Transformers AutoConfig test, the apply_router_weight_after_down addition to shared GroupedExperts defaults to False so existing MoE models are unaffected, and the PP block-boundary handoff has forward+backward parity tests on CPU. The vision/encoding/tokenization files are clearly-marked vendored Moonshot reference code and were reviewed as such.

One inline finding (medium severity): KimiK3TextModel.init_weights unconditionally accesses output_attn_res_norm/output_attn_res_proj, which __init__ only defines when attn_res_block_size is not None. Every other consumer guards that mode, so the checkpoint-free init path would AttributeError for an attn_res_block_size=None model. All current tests set a non-None block size, so the path is uncovered.

Nothing else rose to a high-confidence blocking finding.

[claude]

init_weights reads self.output_attn_res_norm / self.output_attn_res_proj unconditionally, but __init__ only assigns those attributes when self.use_attn_residuals is true (attn_res_block_size is not None, line 1402-1409). Every other consumer of these attributes guards the None-block-size mode — forward branches on use_attn_residuals, and customize_pipeline_stage_modules / get_pipeline_stage_metas use getattr(..., None) — so a model configured with attn_res_block_size=None is a supported configuration, yet the checkpoint-free init path here raises AttributeError before the first forward. Guard the access, e.g.:

            if self.use_attn_residuals:
                self.output_attn_res_norm.reset_parameters()
                nn.init.normal_(self.output_attn_res_proj.weight, mean=0.0, std=init_std)

(matching the getattr pattern used elsewhere). The existing tests always set a non-None block size, so this path is uncovered.

[HuiyingLi]

/claude review

[claude]

Reviewed the full diff (all 19 non-uv.lock files). This is a large but carefully-structured and well-tested model onboarding — registry wiring, config resolution (incl. the stale-Transformers path), cast_model_to_dtype/_keep_in_fp32_modules handling of the KDA fp32 params, checkpoint-free init, PP block partitioning, and the new apply_router_weight_after_down MoE math all check out and carry focused tests. No correctness, security, or repository-invariant issues that I can confidently verify.

One non-blocking coverage observation:

The KDA linear-attention path — the defining "Kimi Linear" component — is not exercised by any test. Every model test in test_pipeline_parallel.py builds _tiny_config(...) with linear_attn_config["kda_layers"] = [], so KimiDeltaAttention, _kda_core, _forward_with_cp, the _torch_kda_gate/_fused_kda_gate helpers, and document_causal_flex_attention are never instantiated or run. The same applies to the context-parallel batch sharding (shard_batch_for_kimi_cp) and the multimodal/vision tower + _merge_input_ids_with_image_features path.

This is understandable — KimiDeltaAttention.__init__ calls _require_fla() and the kernels (chunk_kda, fused_recurrent_kda, ShortConvolution, FusedRMSNormGated) are FLA/GPU-only, so a pure-CPU unit test can't cover the forward. Given that, consider adding a scheduled L1/L2 GPU (or parity) test that exercises at least one KDA layer forward+backward and the CP-sharded MLA path, so the model's novel machinery has some numerical guard rather than relying solely on the MLA-only CPU tests.

[HuiyingLi]

/ok to test 181e86a

[HuiyingLi]

/ok to test fd2ea50

[HuiyingLi]

/ok to test c6f144e

[HuiyingLi]

/ok to test 7c08b9b

[HuiyingLi]

/ok to test 7439737

[HuiyingLi]

/ok to test a4d357e

[HuiyingLi]

/ok to test b4b19d1

[HuiyingLi]

/ok to test b4b19d1

[HuiyingLi]

/ok to test c9fac9d

[HuiyingLi]

/ok to test 29de6d4997022544647c13356f9c8e0466628214

[copy-pr-bot]

/ok to test 29de6d4997022544647c13356f9c8e0466628214

@HuiyingLi, there was an error processing your request: E2

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/

[HuiyingLi]

/ok to test 29de6d4

[github-actions]

🌿 Preview your docs: https://nvidia-preview-huiyingl-feat-kimi-k3.docs.buildwithfern.com/nemo/automodel