[shengliangxu]

Note

Split on 2026-10-30 at the reviewers' request: this PR is now [2/2], the Kimi-K3 expert-parallel adapter itself. [1/2] is #3132: the fixes the checkpoint's own modeling code needs on current transformers and fla-core, moved unchanged out of this PR into modelopt/torch/models/kimi_k3/_remote_fixes.py, with their tests. #3132 merged on 2026-10-30.

Base: #2748's branch, shengliangx/hfep-11b-kimi-k25-vl, since 2026-10-30: #2748 runs a model spec's prepare only once the plan is decided expert-parallel, which this PR's Kimi-K3 spec now relies on. Below #2748 sit #3138 and shengliangx/hfep-base-11b-kimi, a branch (not a PR) that merges #2746 and #2747 with main, which holds #3130 and #3132 since they merged. The diff below is this PR's alone, and it merges after #3138 and #2748.

Earlier re-cuts: on 2026-10-29 this PR moved off #2748 onto a branch merging #2746 and #2747, carrying the one test row and the import-list line it had used from #2748. The series first became a tree on 2026-10-27, from main 1521053 (this PR was [20/21]); it was rebased onto main 7c3c868 later that day, after #2739 merged, and onto main f03651e that night.

Series layout

• main ([OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745, [OMNIML-6064] Load each rank's checkpoint shard straight into a meta-built model #2739, [OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736, [OMNIML-6054] Decide how a model is distributed, from the model built on the meta device #2737, [OMNIML-6082] [1/3] Let a reverse conversion record a merge rule no key matched #3100, [OMNIML-6082] [2/3] Move each tensor of a sharded model whole to one owner rank #3101, [OMNIML-6081] [1/2] Pack sharded weights in FSDP2 weight-update windows, and NVFP4 DTensor shards #3099, [OMNIML-6081] [2/2] Fold fused experts' quantizers into their weights shard-locally #2750, [OMNIML-6063] [1/3] Stage and keep patched a checkpoint's own modeling code #3130 and [OMNIML-6079] [1/2] Fix Kimi-K3's own modeling code for current transformers and fla-core #3132 have merged)
  • [OMNIML-6060] [1/3] Add require_transformers, and build a VLM without its LM head #3131 [1/3] require_transformers, and a VLM built without its LM head
    • [OMNIML-6076] [1/3] Describe expert parallelism in the model specs #3134 [1/3] expert parallelism in the model specs
      • [OMNIML-6076] [2/3] Prepare a model's config for expert parallelism from its spec #3135 [2/3] preparing a model's config from its spec
  • [OMNIML-6073] [1/2] Calibrate and fold expert-parallel fused experts on each rank's own experts #3122 [1/2] calibration on each rank's own experts
    • [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 [2/2] calibration on grouped_mm
      • [OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067 Nemotron-H's experts under grouped_mm
      • [OMNIML-6071] Calibrate a distributed model on per-rank data #2740 calibration on per-rank data
  • [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779 load routing
  • [OMNIML-6082] [3/3] Write a sharded model's HF checkpoint shards from each tensor's owner rank #2751 distributed-writer pieces
    • [OMNIML-6083] Write a sharded model's HF checkpoint without gathering it #2752 no-gather export (on [OMNIML-6082] [3/3] Write a sharded model's HF checkpoint shards from each tensor's owner rank #2751 and [OMNIML-6071] Calibrate a distributed model on per-rank data #2740)
      • [OMNIML-6084] Export an expert-parallel model through the no-gather writer #2753 expert-parallel export
  • [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 [2/3] expert-parallel loader (on [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779, [OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067, [OMNIML-6071] Calibrate a distributed model on per-rank data #2740 and [OMNIML-6060] [1/3] Add require_transformers, and build a VLM without its LM head #3131)
    • [OMNIML-6060] [3/3] Refuse an expert-parallel load that would calibrate the experts wrong #2741 [3/3] post-load refusals
    • [OMNIML-6061] Load the dense part sharded, and size the load warmup to each rank #2742 dense sharding
      • [OMNIML-6062] Dequantize block-FP8 checkpoints as they load expert-parallel #2743 block-FP8 dequantization
    • [OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780 expert-parallel load routing (merges after [OMNIML-6060] [3/3] Refuse an expert-parallel load that would calibrate the experts wrong #2741)
      • [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746 [3/3] the spec hook in the plan decision and the load (on [OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780 and [OMNIML-6076] [2/3] Prepare a model's config for expert parallelism from its spec #3135)
    • [OMNIML-6063] [2/3] Decompress compressed-tensors MoE experts as they load #3137 [2/3] on-the-fly expert decompression
      • [OMNIML-6063] [3/3] Load a checkpoint its model spec decompresses through the expert-parallel loader #2747 [3/3] decompressed checkpoints through the expert-parallel loader
  • The Kimi adapters (on [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746 and [OMNIML-6063] [3/3] Load a checkpoint its model spec decompresses through the expert-parallel loader #2747)
    • [OMNIML-6078] [1/2] Add the Kimi-VL expert-parallel adapter #3138 [1/2] Kimi-VL adapter
      • [OMNIML-6078] [2/2] Add the Kimi-K2.5 expert-parallel adapter #2748 [2/2] Kimi-K2.5 adapter, and prepare on the expert-parallel route only
        • [OMNIML-6079] [2/2] Add the Kimi-K3 expert-parallel adapter #2749 [2/2] Kimi-K3 adapter (this PR)
  • [OMNIML-6085] [1/3] Add forward= to distributed_forward_loop, and fix the GPU memory monitor's NVML probe #3133 [1/3] → [2/3] → [OMNIML-6085] [3/3] Run expert-parallel PTQ from hf_ptq.py #2754 [3/3] expert-parallel PTQ from hf_ptq.py (on [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746, [OMNIML-6061] Load the dense part sharded, and size the load warmup to each rank #2742 and [OMNIML-6084] Export an expert-parallel model through the no-gather writer #2753)

What does this PR do?

Type of change: New feature

Kimi-K3's text backbone, kimi_linear, has no native transformers modeling to swap in, so the adapter patches the checkpoint's remote code in place:

• It replaces the MoE block's per-expert ModuleList with fused experts that dispatch through the experts interface, keeping K3's latent expert space and its custom SwiGLU.
• It registers converters that fuse the checkpoint's per-expert w1/w3 into gate_up_proj and stack w2 into down_proj. The MXFP4 experts are decompressed as they load ([OMNIML-6063] [3/3] Load a checkpoint its model spec decompresses through the expert-parallel loader #2747).
• It registers torch_a2a, the experts implementation it names on the text config, with transformers. transformers refuses to build a model whose experts implementation it has not registered; prepare runs only on the expert-parallel route, so that is the only route that registers it.
• Before it builds the model, it applies the fixes the remote model needs for the expert-parallel load and calibration: KDA's A_log shape and sigmoid(beta), the causal-mask and cache signatures, the kwargs newer transformers pass to tie_weights, and the OutputRecorder import. They are [1/2] of this split, [OMNIML-6079] [1/2] Fix Kimi-K3's own modeling code for current transformers and fla-core #3132, which describes each.
• prepare runs only once the plan is decided expert-parallel ([OMNIML-6078] [2/2] Add the Kimi-K2.5 expert-parallel adapter #2748). Before deciding, the plan decision runs only the spec's make_buildable: the OutputRecorder shim, the staging of the code's imports, and its VLM class's tie_weights adapted to the arguments post_init passes it, which change nothing the model computes or loads. It reads the plan the spec states, on text_config, and finds the experts in the model the checkpoint's own code builds. So a whole-model or FSDP2 load never meets Kimi-K3's patches unless an expert-parallel load of it ran first in the same process; since the patches are process-wide and stay after that load, the spec says so (prepare_patches_every_load), and load_pretrained then refuses a whole-model or FSDP2-sharded load of Kimi-K3 in that process before anything is built. A block built before prepare, as a model loaded earlier in the process holds, keeps its per-expert experts and its own forward. The other fixes are what the remote code needs on current transformers and fla-core anyway. torch_a2a goes only on the prepared copy, and the decompressor rewrites only configs it marked.
• It keeps K3's latent-expert projections (routed_expert_down_proj / routed_expert_up_proj) and its scalar residual gates (self_attention_res_proj, mlp_res_proj, output_attn_res_proj) unquantized in every built-in format.
  • vLLM builds them unquantized and refuses a checkpoint that quantizes them.
  • The checkpoint-mirror recipe behind nvidia/Kimi-K3-NVFP4 already keeps them in BF16.
• It adds Kimi's mm_projector to the default unit's vision block.
  • That block is meant to keep every multimodal projector in BF16, but it named only gemma's and Llama-4's spellings, so --recipe runs, which keep the whole VLM, quantized Kimi's.
  • The published Kimi-K2.5 and Kimi-K3 NVFP4 checkpoints keep it in BF16.

It needs transformers 5.13, the release it was developed and validated against; its spec declares it, so an older release is refused before prepare runs. Its spec states the plan prepare attaches, as #2748's do, so a run at an EP degree of 1 still finds its experts (#2746).

Usage

Picked automatically by apply_distributed_support for kimi_k3 configs under torchrun; see #2754.

Testing

• New tests/unit/torch/models/test_kimi_k3_ep.py (+592 lines), against the stand-in remote module that follows the checkpoint's own code (tests/_test_utils/torch/kimi_k3_remote.py, from [OMNIML-6079] [1/2] Fix Kimi-K3's own modeling code for current transformers and fla-core #3132), covering:
  • the fused experts through the experts interface;
  • the patched MoE block against the remote one, with latent and plain experts and with the situ activation. The block derives latent MoE from routed_expert_hidden_size, as the checkpoint's modeling does; the config has no use_latent_moe field;
  • the w1/w3/w2 conversions;
  • the plan naming exactly the fused params;
  • prepare patching the defining module once, failing loudly when none defines the MoE block, and refusing a call without the checkpoint's path before it changes anything;
  • the plan prepare attaches equal to the one the spec states;
  • prepare leaving the loader alone for bf16 experts, which now runs without compressed-tensors installed;
  • the conversions' registration raising when transformers' conversion API is missing. Before, it only warned, and the experts loaded as missing;
  • a model built on the meta device from the prepared config, in a process that has not registered torch_a2a, passing transformers' experts-implementation check. Without the registration it fails with the experts_implementation="torch_a2a" is not supported error that a tiny Kimi-K3 on 2 GPUs (transformers 5.18) hit while deciding its plan.
• tests/unit/torch/utils/test_model_load_utils.py (in test_resolve_distribution_plan.py before [OMNIML-6054] Decide how a model is distributed, from the model built on the meta device #2737 split it up): a remote-code Kimi-K3 left whole on every rank at an EP degree of 1 is refused, as [OMNIML-6078] [2/2] Add the Kimi-K2.5 expert-parallel adapter #2748's Kimi models are. The test is Kimi-K3's own, with its own stand-in model (block_sparse_moe), and sits apart from [OMNIML-6078] [2/2] Add the Kimi-K2.5 expert-parallel adapter #2748's, so the two PRs merge in either order.
• tests/unit/recipe/test_presets.py: the fp8 and nvfp4 presets leave those five projections and mm_projector unquantized.
• Mutation-checked: each test fails on at least one deliberate break of the code.
• Each checkpoint revision loaded in one process patches its own remote module, found in the package of the class its auto_map loads.
• Duplicates fold into the tests that cover them: the per-expert reference into the patched-block test, the two beta tests into one, the mm_projector preset into the Kimi projections preset. Each drop is checked by mutation.
• transformers 5.12 is refused for Kimi-K3 by its spec.
• Locally on this revision, the stack rebased on main c6fa613: pre-commit over the whole PR against main, at every PR of the stack; and at the stack's tip the full tests/unit on transformers 5.18 and 5.5 and on torch 2.8, and the hf_ptq example tests, each with the same failures as main in that environment. CI runs the full unit suite.
• The GPU results here were measured on earlier revisions of this branch, except one run on this revision: a tiny Kimi-K3 (bf16, 4 layers, 32 experts, KDA and full-attention layers) on 2 GPUs at EP 2 with the dense part replicated, transformers 5.18 and fla-core 0.5.0, --qformat nvfp4 --kv_cache_qformat fp8, max calibration on 64 samples of the default text blend (sequence length 128, batch 2). It decided its plan, previewed before and after PTQ on both ranks, calibrated in 51 s and exported, in 4 min 28 s at 13.9 GB per GPU. Before this revision it failed while deciding its plan (torch_a2a not registered), and with that fixed, in generate (the cache's query offset). GPU CI does not run on this PR while its base is a stacked branch, so the rest of this revision's GPU tests have not run at this head.
• On GPUs (16 GB200 nodes, EP64, dense FSDP2-sharded), Kimi-K3 NVFP4 loads and calibrates with no NaN, and its export checks out against the source checkpoint:

  • all 92 router correction biases match the checkpoint, to BF16 rounding of its F32 values;

  • routed experts export as experts.N.w1/w2/w3 and shared experts as gate_proj/up_proj/down_proj, the naming vLLM 0.28's Kimi-K3 loader reads;

  • the latent-expert projections, the residual gates and mm_projector (373 modules) come out BF16, matching nvidia/Kimi-K3-NVFP4;

    • every other weight and weight scale is byte-identical to the export made before this change, in 120 hashed tensors;
    • only calibration-derived input scales vary, within run-to-run noise;

  • vLLM 0.28 loads the export and starts serving at TP=32. It cannot serve one with an FP8 KV cache, however:

    • its Kimi-K3 FP8-KV path needs an FP8 prefill query, which its NVFP4 linears reject;
    • nvidia/Kimi-K3-NVFP4 ships without a quantized KV cache.

    TBD: generation from the checkpoint-mirror recipe models/moonshotai/Kimi-K3/ptq/nvfp4_experts-fp8_pb_attention, with [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074's expert weight-scale fix.

• Re-cut on 2026-10-27 (see the note at the top): pre-commit, mypy included, passes at this head over the whole change against main, and the series' PRs merged together give a tree byte-identical to the former stack's tip. This PR's own unit tests pass at this head on CPU (torch 2.14, transformers 5.18); CI runs the full unit suite.
• Rebased on 2026-10-27 onto main 7c3c868, after [OMNIML-6064] Load each rank's checkpoint shard straight into a meta-built model #2739 merged: pre-commit, mypy included, and this PR's own unit tests pass at this head, and the series' PRs merged together add no unit-test failure over main (tests/unit/torch utils, export, quantization and models, on transformers 5.18 and 5.5.4 and on torch 2.8).
• Review suggestion, 2026-10-27 night: both kernel wrappers carry functools.wraps, so signatures and tracebacks show the kernel's own; the wrapper test checks the name and signature. Both now ship in [OMNIML-6079] [1/2] Fix Kimi-K3's own modeling code for current transformers and fla-core #3132.
• Re-cut on 2026-10-29 onto [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746 and [OMNIML-6063] [3/3] Load a checkpoint its model spec decompresses through the expert-parallel loader #2747 (see the note at the top): the Kimi-K3 case of [OMNIML-6078] [2/2] Add the Kimi-K2.5 expert-parallel adapter #2748's refusal test is now its own copy of that test, in test_model_load_utils.py, with the stand-in model's tie_weights. Pre-commit, mypy included, passes at this head over the whole change against main, and tests/unit/torch/utils, tests/unit/torch/models and tests/unit/recipe/test_presets.py pass on CPU at this head (torch 2.14, transformers 5.18 and 5.5.4).
• Split on 2026-10-30 (see the note at the top): the remote-code fixes and their tests moved to [OMNIML-6079] [1/2] Fix Kimi-K3's own modeling code for current transformers and fla-core #3132 unchanged, and the stand-in remote module to tests/_test_utils/torch/kimi_k3_remote.py, which both PRs' tests import: [OMNIML-6079] [1/2] Fix Kimi-K3's own modeling code for current transformers and fla-core #3132 keeps what its fixes patch, and this PR adds back the per-expert MoE block, its gate and the situ activation its tests build. Every test of the former PR is in exactly one of the two, unchanged apart from names and imports. The Kimi-K3 refusal test is now its own (_remote_kimi_k3(), with block_sparse_moe), apart from [OMNIML-6078] [2/2] Add the Kimi-K2.5 expert-parallel adapter #2748's. Pre-commit, mypy included, passes at this head over the whole change against main. On CPU (torch 2.14), tests/unit/torch/models, tests/unit/torch/utils/test_model_load_utils.py and tests/unit/recipe/test_presets.py pass on transformers 5.18 and 5.5.4; the adapter's own tests skip below transformers 5.8, as before.
• 2026-10-30: rebased onto main 99c3834, through the new [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746, [OMNIML-6063] [3/3] Load a checkpoint its model spec decompresses through the expert-parallel loader #2747 (which now refuses compressed tensors the model leaves unplaced), [OMNIML-6063] [1/3] Stage and keep patched a checkpoint's own modeling code #3130 and [OMNIML-6079] [1/2] Fix Kimi-K3's own modeling code for current transformers and fla-core #3132, with no conflicts and the own diff unchanged. Pre-commit against main passes, mypy included; tests/unit/torch/models, test_model_load_utils.py, test_ep_loader.py and tests/unit/recipe/test_presets.py pass on CPU, torch 2.14, on transformers 5.18 and 5.5.4.
• 2026-10-30, review round: the Kimi-K3 refusal test's config names its expert count num_experts, as Kimi-K3's does. Pre-commit against main passes, mypy included; tests/unit/torch/models, test_model_load_utils.py and test_ep_loader.py and tests/unit/recipe/test_presets.py pass on CPU, torch 2.14, on transformers 5.18 and 5.5.4.
• 2026-10-30: rebased onto the rebuilt base: main 7168107, where [OMNIML-6063] [1/3] Stage and keep patched a checkpoint's own modeling code #3130 and [OMNIML-6079] [1/2] Fix Kimi-K3's own modeling code for current transformers and fla-core #3132 have merged, joined with [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746 5d51cd6 and [OMNIML-6063] [3/3] Load a checkpoint its model spec decompresses through the expert-parallel loader #2747 4494480. There were no conflicts in this PR and its own diff is unchanged. Pre-commit against main passes, mypy included. On CPU, torch 2.14, with transformers 5.18 and 5.5.4, these pass: tests/unit/torch/models, test_model_load_utils.py, test_ep_loader.py, test_unplaced_source_keys.py and tests/unit/recipe/test_presets.py.
• 2026-10-30, cjluo-nv's review round, at 7a50ed7 on the rebuilt base (main 62a5c48 with the new [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746 and [OMNIML-6063] [3/3] Load a checkpoint its model spec decompresses through the expert-parallel loader #2747, whose decompressor now reverses its converters as the ones it rewrote):
  • a Kimi-K3 MoE block built before prepare keeps its own forward: the patched-block test calls its pre-patch block again after patching, which fails without the fix;
  • the export names packed experts back to the checkpoint's w1/w3/w2: a unit test maps ModelOpt's per-expert names through the exporter's reverse name mapper, over Kimi-K3's registered converters rewritten as the decompressor rewrites them, and fails with the arity mismatch without [OMNIML-6063] [3/3] Load a checkpoint its model spec decompresses through the expert-parallel loader #2747's decompressor fix;
  • new 2-GPU case test_kimi_k3_packed_experts_load_calibrate_and_name_back_expert_parallel in tests/gpu/torch/utils/test_expert_parallel_load.py: a tiny Kimi-K3-shaped remote checkpoint (tests/_test_utils/torch/kimi_k3_tiny_modeling.py) with 8 packed MXFP4 experts, loaded expert-parallel through the spec; each rank checks its decompressed experts, its logits against the whole model on one GPU, FP8 calibration against the whole model's, and the exported expert names. Not yet run on GPUs: a cluster run at this head is pending. Its non-distributed parts run on CPU (transformers 5.18): the experts decompress exactly, the forward runs, FP8 calibration picks up the fused experts, and the names map back. One GPU skips it;
  • print_rank_0 is imported at module scope; KimiK3Experts is attributed to transformers' MiniMaxM3VLExperts, which it is adapted from (file header, LICENSE, license-hook exclusion);
  • pre-commit against main 584a80c passes, mypy included; on CPU, torch 2.14, tests/unit/torch/models, test_model_load_utils.py, test_ep_loader.py, test_unplaced_source_keys.py and tests/unit/recipe/test_presets.py pass: 605 on transformers 5.18, and 470 passed, 56 skipped on 5.5.4.
• 2026-10-30, cjluo-nv's second review round, at 6b8fc1e on the rebuilt base (main 62a5c48 joined with [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746 d3a058e and [OMNIML-6063] [3/3] Load a checkpoint its model spec decompresses through the expert-parallel loader #2747 dc0a618):
  • a packed Kimi-K3 planned on four ranks through AutoModel gets no prepare: the plan is the whole model on every rank, and the config returned is the one read, unmarked and still holding its compressed-tensors scheme. Without [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746's model-class check it fails, on prepare;
  • the tiny Kimi-K3 the GPU case writes copies its scheme to the top level, as Kimi-K3's config class does;
  • the left-whole refusal test expects the error this transformers raises: a ValueError at or above the spec's 5.13 floor, the RuntimeError naming the release below it;
  • the GPU file imports the tiny modeling below its transformers guard, the unit tests import build_reverse_name_mapper at module scope, and both restore only the loader hooks the decompressor patches;
  • pre-commit against main 059588b passes, mypy included; on CPU, torch 2.14, tests/unit/torch/models, tests/unit/torch/utils and tests/unit/recipe/test_presets.py pass on transformers 5.18 and 5.5.4;
  • the compressed-expert GPU cases need compressed-tensors, which the CI GPU container does not have, so they skip in CI and run on our cluster. There, at 6b8fc1e with transformers 5.18.0 and compressed-tensors 0.19.0, the whole tests/gpu/torch/utils/test_expert_parallel_load.py passed: on 2 GPUs, 4 passed and [ep-half] skipped, as it needs 4 GPUs; on 4 GPUs, all 5 passed, [ep-half] included. test_kimi_k3_packed_experts_load_calibrate_and_name_back_expert_parallel passed on both. It failed at 7a50ed7, with the packed weight cast to bf16 on the expert-parallel shard path, before [OMNIML-6063] [2/3] Decompress compressed-tensors MoE experts as they load #3137 kept it uncast there;
  • 44084ef adds only that docstring note on the patches staying for the process.
• 2026-10-30, Claude's and cjluo-nv's review round, at 342c4fb on the rebuilt base (main 059588b, with [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746 0eadf4c and [OMNIML-6063] [3/3] Load a checkpoint its model spec decompresses through the expert-parallel loader #2747 23eb10d):
  • prepare installs the decompressor before it patches anything, so a refusal leaves the process unchanged (a78d656). The refusal test has the real decompressor refuse, for an unparsable scheme and for a dtype other than bf16, and checks that the remote block, the config, the converter registry, the experts implementations and the loader hooks are unchanged. Both cases fail with the old order;
  • prepare also reads whether the checkpoint is compressed from its own config.json (ea287d8). A config that has lost its scheme and the mark, as a load's returned model.config has, is refused instead of loading the experts newly initialized. The unit test builds the config as the load returns it and fails without the change; the 2-GPU Kimi-K3 job passes the returned model.config to another load and checks the refusal (342c4fb);
  • pre-commit against main passes, mypy included; on CPU, torch 2.14, tests/unit/torch/models, tests/unit/torch/utils and tests/unit/recipe/test_presets.py pass on transformers 5.18 and 5.5.4;
  • at this head, with transformers 5.18.0 and compressed-tensors 0.19.0, the whole tests/gpu/torch/utils/test_expert_parallel_load.py passed: on 2 GPUs, 4 passed and [ep-half] skipped, as it needs 4 GPUs; on 4 GPUs, all 5 passed. The Kimi-K3 case passed on both, including the reload with the returned model.config, which every rank refused.
• 2026-10-31: restacked onto [OMNIML-6078] [2/2] Add the Kimi-K2.5 expert-parallel adapter #2748, which now runs prepare only on the expert-parallel route; this PR's base moved to [OMNIML-6078] [2/2] Add the Kimi-K2.5 expert-parallel adapter #2748's branch. At 63c8267:
  • Kimi-K3's spec names its import step (make_importable, now make_buildable: the OutputRecorder shim and the staging of its imports), puts its stated plan on text_config, and names compressed-tensors among the formats its prepare decompresses (c459add, 63c8267). torch_a2a is registered only on the expert-parallel route, and the patches-persist note narrows to a later load after an expert-parallel one;
  • test_only_an_expert_parallel_plan_patches_kimi_k3s_code[fsdp2|whole|expert-parallel|expert-parallel-packed] runs the real prepare through the decision. FSDP2-sharded, or left whole and refused, the remote MoE block, the converter registry, torch_a2a, the loader hooks and the config as read are untouched. An expert-parallel plan fuses the experts, and a packed checkpoint is planned expert-parallel and gets the decompressor. The fsdp2 and whole cases fail on the old decision, the packed one without the spec's formats;
  • pre-commit against main passes, mypy included; on CPU, torch 2.14, tests/unit/torch/models, tests/unit/torch/utils and tests/unit/recipe/test_presets.py pass on transformers 5.18 and 5.5.4;
  • at this head, with transformers 5.18.0 and compressed-tensors 0.19.0, the whole tests/gpu/torch/utils/test_expert_parallel_load.py passed: on 2 GPUs, 4 passed and [ep-half] skipped, as it needs 4; on 4 GPUs, all 5 passed. The Kimi-K3 case passed on both, its reload refusal included.
• 2026-10-31, cjluo-nv's and Claude's reviews, at 183bf95, rebased onto [OMNIML-6078] [2/2] Add the Kimi-K2.5 expert-parallel adapter #2748 27961b6:
  • a whole-model or FSDP2-sharded load of Kimi-K3 after an expert-parallel one in the same process is refused before anything is built (EPSpec.prepare_patches_every_load, recorded by apply_distributed_support, checked by load_pretrained). test_a_load_after_a_prepare_that_patched_every_load_is_refused[whole|fsdp2] fails without it, and the 2-GPU Kimi-K3 job runs the review's exact continuation, which every rank must refuse;
  • Kimi-K3's build step also adapts its VLM class's tie_weights, which the decision's build needs on transformers 5.18, and the checkpoint-config read covers llm_config too;
  • the adapter's tests move to tests/unit/torch/models/kimi_k3/test_ep.py;
  • pre-commit against main passes, mypy included; on CPU, torch 2.14, tests/unit/torch/models, tests/unit/torch/utils and tests/unit/recipe/test_presets.py pass on transformers 5.18 and 5.5.4. At this head, with transformers 5.18.0 and compressed-tensors 0.19.0, the whole GPU file passed: on 2 GPUs, 5 passed and [ep-half] skipped, as it needs 4; on 4 GPUs, all 6 passed. The Kimi-K3 case passed on both, the same-process whole-model reload refused on every rank. It loads through the spec and the loader directly, not through the plan decision.
• 2026-10-31, Claude's review, at 4389363: Kimi-K3's prepare records its lasting patch where it makes it, so a step after it that raises still leaves the refusal armed (test_a_prepare_that_raises_after_patching_still_records_it fails without it), and _checkpoint_quantized's docstring says which revision it reads. Pre-commit against main passes, mypy included; tests/unit/torch/models, tests/unit/torch/utils and tests/unit/recipe/test_presets.py pass on CPU, torch 2.14, on transformers 5.18 and 5.5.4. At this head, the whole GPU file passed, with transformers 5.18.0 and compressed-tensors 0.19.0: on 2 GPUs, 5 passed and [ep-half] skipped, as it needs 4; on 4 GPUs, all 6 passed; both Kimi cases passed on both.
• 2026-10-31, Claude's review, at f186098: the early record of prepare's lasting patch keys on kimi_k3.ep.MODEL_TYPE, the one definition the spec registers, and the docs say the refusal comes before the load builds or loads the model. The plan decision's throwaway meta build runs earlier. The Kimi-K3 case's behaviour is unchanged since the GPU run at 4389363: the key has the same value. Pre-commit against main passes, mypy included; tests/unit/torch/models, tests/unit/torch/utils and tests/unit/recipe/test_presets.py pass on CPU, torch 2.14, on transformers 5.18 and 5.5.4.
• 2026-10-31, at a60c01c: restacked onto [OMNIML-6078] [2/2] Add the Kimi-K2.5 expert-parallel adapter #2748 b88d349 on the rebuilt base (main e352cac, [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746 d8059b3, [OMNIML-6063] [3/3] Load a checkpoint its model spec decompresses through the expert-parallel loader #2747 6b87b69). test_model_load_utils.py now holds the moved expert-parallel loader tests; its one import conflict kept both sides' imports, and the own diff is otherwise unchanged. Pre-commit against main passes, mypy included; tests/unit/torch/models, tests/unit/torch/utils and tests/unit/recipe/test_presets.py pass on CPU, torch 2.14, on transformers 5.18 and 5.5.4. The base changed under the Kimi GPU cases ([OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746's load now checks the spec first), so the whole GPU file was rerun at this head, and passed, with transformers 5.18.0 and compressed-tensors 0.19.0: on 2 GPUs, 5 passed and [ep-half] skipped, as it needs 4; on 4 GPUs, all 6 passed. Both Kimi cases passed on both; they load through the spec and the loader directly, not through the plan decision.
• Copied code: KimiK3Experts' fused parameters and eager loop are adapted from transformers v5.18.0's MiniMaxM3VLExperts (Apache-2.0). The file carries the source link and notice, and LICENSE lists the copyright holder. The human licensing sign-off (NVBug 6046893) is pending.

Before your PR is "Ready for review"

• Is this change backward compatible?: ✅ Additive. Two effects reach beyond the expert-parallel path: the shared default disabled-quantizer unit now also keeps Kimi's latent-expert projections, residual gates and mm_projector unquantized in every preset that imports it; and prepare, which only an expert-parallel plan runs, patches the checkpoint's own classes for the process (see above).
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: ✅ transformers' MiniMaxM3VLExperts (Apache-2.0), attributed in the file header and LICENSE; human licensing sign-off pending (NVBug 6046893). flash-linear-attention is a runtime requirement of the checkpoint's own remote code, not of ModelOpt.
• Did you write any new necessary tests?: ✅
• Did you update Changelog?: N/A. The feature's entry lands with [OMNIML-6085] [3/3] Run expert-parallel PTQ from hf_ptq.py #2754.
• Did you get Claude approval on this PR?: ❌ Not yet.

Additional Information

Part of expert-parallel PTQ for Hugging Face MoE models (was part 20 of 21). This PR is [2/2]; it merges after #3138 and #2748, which merge after #2746 and #2747 (#3132 has merged).

Summary by CodeRabbit

• New Features
  • Added support for running Kimi-K3 models with expert parallelism, including supported remote model implementations.
  • Updated FP8 and NVFP4 quantization presets for Kimi-K3. Selected projections and multimodal components remain unquantized by default, while other layers can still be quantized.
• Compatibility
  • Kimi-K3 checkpoint support requires Transformers 5.13 or later. Loading checkpoint-specific remote code requires explicitly enabling trusted remote code.
• Bug Fixes
  • Improved compatibility with compressed model weights and remote-model loading.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

• Last reviewed commit: 78dad39
• Latest commit, not reviewed: a60c01c

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

Walkthrough

This change registers Kimi-K3 remote-code modeling and adds an expert-parallel adapter with fused experts, checkpoint weight converters, and remote API compatibility patches. It also updates default-disabled quantizer rules for Kimi-K3 projections and adds adapter and preset tests.

Changes: Kimi-K3 expert parallelism

Layer / File(s)	Summary
Fused expert execution and weight conversion modelopt/torch/models/kimi_k3/ep.py, tests/unit/torch/models/test_kimi_k3_ep.py	The adapter defines the expert-parallel plan, replaces per-expert modules with fused experts, and registers checkpoint converters. Tests validate expert execution and converted weight targets.
Remote model compatibility modelopt/torch/models/kimi_k3/ep.py, tests/unit/torch/models/test_kimi_k3_ep.py	The adapter adds wrappers for KDA functions, causal-mask arguments, cache mask sizing, and tie_weights. Tests cover the wrapper behavior.
Preparation and model registration modelopt/torch/models/kimi_k3/*, modelopt/torch/models/__init__.py, tests/unit/torch/models/test_kimi_k3_ep.py, tests/unit/torch/utils/test_resolve_distribution_plan.py	The preparation function loads and patches remote model classes, configures expert parallelism and conversions, and handles quantization configuration. The model package imports the spec that selects the preparation function. Tests cover preparation, loading, failure cases, and distribution-plan handling.

Kimi-K3 quantizer defaults

Layer / File(s)	Summary
Quantizer rules and preset checks modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml, tests/unit/recipe/test_presets.py	Default-disabled rules cover specified routed-expert, residual, and mm_projector layers. FP8 and NVFP4 tests check these placements and verify that self_attn.o_proj remains quantized.

Estimated code review effort: 4 (Complex) | ~60 minutes

Change: Feature

Sequence Diagram(s)

sequenceDiagram
  participant ModelSpec
  participant Prepare
  participant RemoteModel
  ModelSpec->>Prepare: select Kimi-K3 preparation
  Prepare->>RemoteModel: load and patch remote classes
  Prepare->>Prepare: configure expert-parallel plan and conversions

Loading

Possibly related issues

• NVIDIA/Megatron-LM#6757: The issue describes planned Kimi-K3 MoE support and expert parallelism; this change adds Kimi-K3 expert-parallel support.
• NVIDIA/Megatron-LM#8427: The issue describes routed-only MegaMoE extensions, latent paths, and expert-parallel integration; this change adds expert-parallel support with latent-expert handling.

Merge Risk: Low · up to 78dad

A malformed compressed checkpoint configuration could pass preparation without the required decompression setup, risking incorrect expert loading. Narrow the exception handling before merging.

Pre-merge checks: 5

Check name	Status	Explanation
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns		No listed security anti-pattern is introduced in the changed ModelOpt Python code. prepare exposes trust_remote_code as a caller-configurable parameter defaulting to False and passes that variab…
Title check		The title clearly and concisely describes the main change: adding the Kimi-K3 expert-parallel adapter.
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.

Full details: Security Anti-Patterns

Explanation

No listed security anti-pattern is introduced in the changed ModelOpt Python code. prepare exposes trust_remote_code as a caller-configurable parameter defaulting to False and passes that variable to get_class_from_dynamic_module; it does not hardcode trust_remote_code=True. The diff adds no torch.load(..., weights_only=False), numpy.load(..., allow_pickle=True), external-input eval/exec, or # nosec comments. No dependency manifest changes were made.

Finishing Touches

Generate docstrings

• Commit to this branch
• Create a new PR

Fix failing CI checks

• Commit to this branch
• Create a new PR

🧪 Generate unit tests (beta)

• Commit to this branch
• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

[github-actions]

PR Preview Action v1.8.1
🚀 View preview at https://NVIDIA.github.io/Model-Optimizer/pr-preview/pr-2749/
Built to branch gh-pages at 2026-10-27 00:06 UTC. Preview will be ready when the GitHub Pages deployment is complete.

[shengliangxu]

/ok to test 784e815

[codecov]

Codecov Report

❌ Patch coverage is 93.41085% with 17 lines in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (shengliangx/hfep-11b-kimi-k25-vl@e83762c). Learn more about missing BASE report.

Files with missing lines	Patch %	Lines
modelopt/torch/models/kimi_k3/ep.py	93.28%	17 Missing ⚠️

Additional details and impacted files

@@                         Coverage Diff                         @@
##             shengliangx/hfep-11b-kimi-k25-vl    #2749   +/-   ##
===================================================================
  Coverage                                    ?   62.63%           
===================================================================
  Files                                       ?      668           
  Lines                                       ?    74485           
  Branches                                    ?        0           
===================================================================
  Hits                                        ?    46654           
  Misses                                      ?    27831           
  Partials                                    ?        0           

Flag	Coverage Δ
unit	62.63% <93.41%> (?)

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[shengliangxu]

/ok to test 57494b6

[shengliangxu]

/ok to test 53b85aa

[shengliangxu]

/ok to test ba6939d

[shengliangxu]

/ok to test 4389363

[shengliangxu]

/claude review

[claude]

[SUGGESTION] The refusal's key is a literal here that has to equal ModelSpec.model_type in specs.py, and nothing checks the two agree.

check_not_patched_by_prepare is called with the model_type load_pretrained resolved from the checkpoint's config, so if the spec's model_type were ever renamed (or a second Kimi-K3 spelling registered) this early mark would land under a key no load looks up. The loss is silent and exactly the case this call exists for: apply_distributed_support still marks with the real model_type for a prepare that returns, so only the raised-mid-prepare path — the one test_a_prepare_that_raises_after_patching_still_records_it guards — would quietly stop refusing. The test asserts the same literal, so it would not catch the divergence either.

Cheapest fix is to derive it rather than spell it:

from modelopt.torch.models import get_spec  # or thread the model_type through prepare
...
mark_patched_every_load(get_spec("kimi_k3").model_type)

or, if the literal is preferred, have the spec test assert it: assert get_spec("kimi_k3").model_type == "kimi_k3" alongside the existing ep_spec.prepare is kimi_k3.prepare check in test_the_kimi_k3_spec_routes_expert_parallelism_through_this_adapter. Non-blocking — the placement itself is right, and I verified the three assignments in _patch_kimi_moe_block cannot raise between patching and this call.

[shengliangxu]

Taken in f186098. kimi_k3.ep.MODEL_TYPE is now the one definition: the spec registers it and the early record keys on it. The test checks the record against get_spec("kimi_k3").model_type.

[claude]

[SUGGESTION] "refused before anything is built" is true of the load but not of the plan decision that precedes it, and the ordering is the whole point of this guard — so the claim is worth narrowing.

In load_pretrained, this check sits inside the _refused_together("load checks", ...) block at model_load_utils.py:1055, which runs after resolve_distribution_plan (:990). That decision calls build_meta_causal_lm, so under torchrun a meta model has already been constructed from the patched KimiSparseMoeBlock by the time the refusal fires. Harmless in practice — it is a throwaway on meta, and _spec_ep_plan names the fused params so the decision comes out the same — but a reader taking the docstring literally would conclude nothing touches the patched classes first.

Related and worth a sentence while you are here: with MODELOPT_EP_DEGREE=1 under torchrun the decision raises its own generic ValueError ("can be expert-parallel, but the plan resolved …") before reaching this check, so the clear "load it in a new process" message is only what the FSDP2 and single-process routes get. The two unit cases cover exactly those, so coverage is fine; it is the docstring that over-promises.

Suggested change

	load of the model is refused before anything is built.
	expert-parallel load builds from: once it ran in this process, a whole-model or FSDP2-sharded
	load of the model is refused before it builds or loads one. The plan decision's throwaway meta
	build runs earlier, so an EP degree of 1 under torchrun meets its own refusal first.

[shengliangxu]

Taken in f186098, with your wording. check_not_patched_by_prepare, the spec field and prepare's docstring now say the refusal comes before the load builds or loads the model. check_not_patched_by_prepare adds that the plan decision's throwaway meta build runs earlier, so under torchrun an EP degree of 1 meets the decision's own refusal first.

[shengliangxu]

/ok to test f186098

[shengliangxu]

/ok to test a60c01c