feat(distributed): Kimi K3 pipeline stage executor and full-model multi-node validation
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
Part of #1331. Follow-up to #1334.
## Problem / Background
#1334 adds the Kimi K3 text backbone (`model_type: "kimi_k3"`, `moonshotai/Kimi-K3`, 2.8T parameters, 93 layers: 69 Kimi Delta Attention layers and 24 gated NoPE-MLA layers, latent SiTU MoE with 896 experts in the `mxfp4-pack-quantized` layout, Attention Residuals every 12 layers). The checkpoint is about 1.56 TB on disk and about 1.4 TB resident at 4-bit, so the full model needs at least three 512 GB nodes through the pipeline in `docs/distributed.md`. No such nodes are reachable at the time of the #1334 PR, so that PR validates only a layer-truncated local copy (`num_hidden_layers` reduced to what fits on a 128 GB machine: load, shape, causality and finite-logits gate). Its acceptance criterion "fluent output on the full model through the distributed pipeline" is discharged by this issue.
The layer-truncated gate is not a substitute for the full run. Until this issue lands, Kimi K3 is single-process-only in mlxcel and cannot actually be served at full size on any hardware.
## Current Behavior
- The #1334 PR adds only a `build_per_layer_bytes` arm in `src/distributed/pipeline/partition_profile_heuristics.rs:32`, so the partitioner can size Kimi K3 stages. It adds no `StageFamily` variant (`src/distributed/pipeline/stage_executor/mod.rs:239`), no `resolve_stage_family` arm (`mod.rs:345`), and no executor under `src/distributed/pipeline/stage_executor/`, so `LoadedStageExecutor::load` bails on a `kimi_k3` model and `--pp-size`, `--pp-layers` and `--distributed-config` cannot run it.
- Stages exchange exactly one hidden-state tensor per micro-batch: `StageExecutionInput::HiddenStates(&MlxArray)` / `StageExecutionOutput::HiddenStates` (`mod.rs:77-88`), serialized by `serialize_mlx_array` (`src/distributed/pipeline/wire_tensor.rs:27`) into `ActivationMessage.tensor_data` (`src/distributed/pipeline/activation_transfer.rs:46`). Attention Residuals need more than that: layer `l` mixes its input against every frozen block stored by layers `0, 12, 24, ...` below it, so a stage that starts at layer 24 needs the two blocks stored by layers 0 and 12 on an earlier node.
- `language_model.model.output_attn_res_proj.weight` and `language_model.model.output_attn_res_norm.weight` classify as `WeightClass::Other` in `classify_weight_key` (`src/distributed/pipeline/partial_loading.rs:145`), which `should_load_key` loads on the first stage, but the final AttnRes mix runs after the last layer on the lm_head stage.
- `mlx_dtype_to_tensor_dtype` (`src/distributed/kv_cache_serde/types.rs:438`) has no arm for MLX dtype codes 1 (uint8) or 3 (uint32) and `TensorDtype` (`src/distributed/tensor_protocol.rs:152`) has no `UInt32` variant, so the mxfp4 expert planes (uint32 packed codes, uint8 E8M0 scales) cannot cross the wire today; they fail loudly rather than up-cast, but nothing tests that.
## Proposed Solution
1. `StageFamily::KimiK3` with `name() == "kimi_k3"`, in `supported_families()` sorted by name (it goes between `Jamba` and `Llama`), `ModelType::KimiK3 => StageFamily::KimiK3` in `resolve_stage_family`, an arm in `load_family_backend` guarded by `ensure_no_adapter` like `Qwen3Next` (`mod.rs:511`), and a `kimi_k3_family_is_registered` floor test in `family_registry_tests.rs` asserting membership and the name string.
2. `src/distributed/pipeline/stage_executor/kimi_k3.rs`, modelled on `qwen3_next.rs`: `KimiK3StageExecutor { model: KimiK3StageModel, cache_store: PointerOwnedCacheStore }` where `C` is the per-layer cache type `KimiK3Model::make_caches()` returns in #1334 (`KimiDeltaCache` for KDA layers, `KVCache` holding `(latent, k_pe)` for MLA layers, wrapped in one enum with an `offset(&C) -> i32` accessor). `make_caches` returns one external `KVCache::new()` per local layer; `execute` calls `caches_for_sequence`, dispatches `TokenIds` / `HiddenStates`, and `sync_external_offsets` back, exactly as `qwen3_next.rs:61-101`.
3. `KimiK3StageModel::load(model_dir, filter: &LayerFilter, stage_index) -> Result` in `src/models/kimi_k3.rs` next to the full model, following `Qwen3NextStageModel` (`src/models/qwen3_next.rs:1827-2055`): read the config, `identify_required_shards` from `model.safetensors.index.json` so a stage opens only the shards holding its layer range, `filter_weight_map`, then run the #1334 sanitize (fused QKV concat, `A_log[:96]`, expert stacking to `mlp.switch_mlp.*` uint32 / uint8) on the filtered map only. Layer `i` is KDA iff `i + 1` is in `linear_attn_config.kda_layers`, decided from the GLOBAL index so the cache variant and mask match the single-process path. MLA layers get `create_causal_mask(seq_len, offset)` with the offset from the stage's first local MLA cache; reuse the `stage_attention_offset` pattern (`qwen3_next.rs:1857`) rather than a second copy. Add `"model.output_attn_res_"` and `"language_model.model.output_attn_res_"` to `NORM_PREFIXES` in `partial_loading.rs` so the lm_head stage loads them, with a `partial_loading_tests.rs` case.
4. Attention Residual transport. The stage output is one tensor `[B, T, (K + 1) * D]` with `D = hidden_size = 7168`: the layer output first, then the `K` frozen blocks in storage order, concatenated on the last axis in the activation dtype. `K` after executing layers `0..end` is `ceil(end / 12)` (layer 0 stores the embeddings, layer 12 stores its input, and so on), so a stage starting at layer `s` expects `K = ceil(s / 12)` and rejects any other last-axis width with an error naming `s`, `K` and the received shape. The receiving stage splits the tensor, recomputes `inv_rms_k = rsqrt(mean(raw_k^2, -1) + eps)` in f32 from `raw_k` (the same formula the single-process path uses when it stores a block), and continues. The entry stage starts with an empty list; the lm_head stage applies the `output_attn_res_*` mix before `norm` and `lm_head`. Packing into the existing single tensor was chosen over adding an `aux_tensors` field to `ActivationMessage` because it keeps the wire struct and the in-process `StageExecutionOutput` unchanged and needs no capability-protocol bump; the largest payload is decode at `K = 7` for the last stage, `8 * 7168 * 2` bytes per token.
5. Wire dtype check: add `TensorDtype::UInt32 = 9` (`element_size` 4) and the `1 => UInt8`, `3 => UInt32` arms in `mlx_dtype_to_tensor_dtype`, and a round-trip test in `activation_transfer_tests.rs` (or a new `wire_tensor_tests.rs`) that a uint32 `[2, 3, 448]` array and a uint8 `[2, 3, 112]` array pass through `serialize_mlx_array` / `deserialize_wire_tensor` byte-identical with the same dtype code. Stages load their own weights from local safetensors, so the mxfp4 planes never travel today; this guards KV-cache serde and any future stage migration against a silent float promotion.
6. Full-model run on at least three 512 GB nodes with the multi-host path in `docs/distributed.md` (a three-node TOML next to `examples/distributed/pipeline_remote_2node_tcp.toml`), then `./target/release/mlxcel generate -m models/Kimi-K3 -p "Write a Python retry wrapper with exponential backoff." -n 64` with the XTML prompt rendered by #1338.
7. `docs/distributed.md` (Pipeline parallelism section) and `docs/supported-models.md` (family list at line 64 and the Distributed support summary table at line 607) record the measured per-node resident memory, the validated topology and layer split, and the single-process limitation removed by this issue.
## Scope
**In scope:** `src/distributed/pipeline/stage_executor/{mod.rs,kimi_k3.rs,family_registry_tests.rs}`, `src/models/kimi_k3.rs` (stage model only), `src/distributed/pipeline/partial_loading.rs` and its tests, `src/distributed/kv_cache_serde/types.rs`, `src/distributed/tensor_protocol.rs`, the wire round-trip test, a three-node example TOML, and the two docs.
**Out of scope:** vision under PP (#1342 is single-process first), PP+LoRA for this family (keep `ensure_no_adapter`), fused kernels for the AttnRes mix or the KDA short-conv step, tensor parallelism for Kimi K3.
## Implementation Notes
- **Reuse**: `PointerOwnedCacheStore` (`stage_executor/common.rs:31`), `filter_weight_map` / `identify_required_shards` (`partial_loading.rs`), `Qwen3NextStageModel` as the load and execute template, the #1334 layer forward and sanitize functions (call them, do not fork them; update their `// Used by:` comment).
- **Edge cases**: a stage whose range holds only KDA layers has no attention offset and no mask; a stage boundary exactly at a multiple of 12 stores the block on the sending side, never the receiving side; a stage of width 12 or less may hold zero storing layers and must forward the block list unchanged; decode (`T == 1`) and prefill share the packing code.
- **Error handling**: wrong last-axis width, cache-count mismatch, and `output_attn_res_*` missing on the lm_head stage are `Err(String)` from the stage model, surfaced through the existing `RemoteStageResponse::Error` path; nothing is silently zero-filled.
- **Blocked**: `status:blocked` until three 512 GB nodes are available; items 1 to 5 can land first behind the truncated-copy tests, but the issue closes only after item 6.
## Acceptance Criteria
- [ ] `mlxcel generate -m models/Kimi-K3-truncated --pp-size 2` (8-layer local copy from #1334, split `0-3,4-7` so a block crosses the boundary) produces logits within bf16 jitter of the single-process forward on the same copy.
- [ ] On the truncated copy, a per-layer mean-abs trace of the first 8 layers' outputs matches an f32 scalar reference of the #1334 formulas run on the real weights within bf16 jitter, with no Python in the path.
- [ ] The three-node run above returns fluent, correct Python (an exponential-backoff retry wrapper) for the coding prompt; the output and the per-node resident memory are recorded in the PR.
- [ ] `supported_families()` lists `kimi_k3` in sorted position; `family_registry_tests.rs` and the wire round-trip test pass.
- [ ] The two docs state the measured memory and validated topology, and no longer describe Kimi K3 as single-process-only.
## Verification
```bash
cargo build --release --features metal,accelerate
cargo test --workspace --profile test-fast --features metal,accelerate distributed::pipeline
cargo clippy --workspace --all-targets -- -D warnings && cargo fmt --check
./target/release/mlxcel generate -m models/Kimi-K3-truncated --pp-size 2 --pp-layers 0-3,4-7 -p "Hello" -n 8
# three 512 GB nodes, per docs/distributed.md multi-host path
./target/release/mlxcel generate -m models/Kimi-K3 -p "Write a Python retry wrapper with exponential backoff." -n 64
```
## Technical Considerations
Related: #1334 (backbone), #1338 (XTML prompt rendering), #1331 (epic). `qwen3_next.rs` is the closest sibling because it also carries a model-owned mixed linear/attention cache across the stage boundary; `nemotron_h.rs` and `jamba.rs` show the same store for SSM hybrids.
Contributor guide
Research direction
Start with src/distributed/pipeline/stage_executor/qwen3_next.rs and the Kimi K3 paths in src/models/kimi_k3.rs, then review the listed registry, partial-loading, tensor-protocol, and wire-transfer tests. Validate the truncated two-stage run before attempting the three-node full-model command. Done means the acceptance tests pass, the full run is recorded, and both distributed-support docs include the measured topology and memory.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- distributed-systems, documentation
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 35/100