[Feature] Add support for Kimi K2.6
- Dominant language
- Python
- Stars
- 2k
- Forks
- 561
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 145
Description
# [Feature] Add support for Kimi K2.6
## Summary
Add NeMo RL support for `moonshotai/Kimi-K2.6` on the Megatron Core / Megatron backend, including checkpoint conversion or loading, model configuration mapping, example training recipes, and backend consistency validation for RL workflows.
## Motivation
Kimi K2.6 is a recent open-source MoE model from Moonshot AI with strong long-horizon coding and agentic-task performance. The Hugging Face model card describes it as a native multimodal agentic model with:
- 1T total parameters, 32B activated parameters
- 384 experts, 8 selected experts per token, 1 shared expert
- 61 layers, 1 dense layer
- 256K context length
- MLA attention, SwiGLU activation, 160K vocabulary
- MoonViT vision encoder
NeMo RL already exposes a Megatron Core path for large models, long context, MoE, sequence packing, and RL training. Adding Kimi K2.6 support would make it possible to post-train and evaluate a frontier-scale MoE coding/agent model in NeMo RL without users hand-rolling the Megatron mapping and validation.
## Proposed Scope
### MCore model integration
- Add or reuse the Megatron Core model mapping required for Kimi K2.6's architecture.
- Map the HF config fields to the appropriate Megatron Bridge / MCore configuration:
- MoE expert count, selected experts, shared experts, expert hidden size
- dense layer placement
- MLA attention settings
- vocab size and tokenizer settings
- long-context settings
- Confirm whether Kimi K2.6 can reuse the Kimi K2.5 architecture path, since the Kimi model card says K2.6 has the same architecture as K2.5 and can reuse its deployment method.
### Checkpoint conversion / loading
- Support conversion from `moonshotai/Kimi-K2.6` Hugging Face checkpoints into the MCore format used by NeMo RL, or document the supported loading path through Megatron Bridge.
- Preserve tied/untied embedding semantics, MoE tensor layout, tokenizer assets, and chat template behavior.
- Document any unsupported weight formats. Kimi K2.6 includes native INT4 quantization; if INT4 is out of scope for training, please document the expected BF16/FP8 path.
### Training and rollout recipes
- Add a minimal SFT recipe for Kimi K2.6 on the Megatron backend.
- Add a minimal GRPO recipe for Kimi K2.6 on the Megatron backend.
- Include recommended parallelism settings for a smoke test and for a realistic multi-node run where possible:
- TP / PP / CP / EP / FSDP
- sequence packing on/off guidance for MoE
- long-context guidance
- Clarify supported rollout backends:
- vLLM and/or SGLang for generation
- Megatron inference if supported to avoid weight conversion
### Validation
Follow the NeMo RL "Add New Models" validation workflow:
- Verify HF vs rollout-backend log probability consistency.
- Verify Megatron vs rollout-backend log probability consistency.
- Run across real and synthetic prompts, greedy and sampling modes, multiple batch sizes, and at least short/medium/long sequence lengths.
- Use the documented `1.05` error threshold for equal-precision log-probability consistency unless Kimi-specific precision caveats require a different threshold.
- Run the existing model diagnostics where applicable:
- `max_model_len_respected.py`
- `long_generation_decode_vs_prefill.py`
- `check_hf_model_embeddings_untrained.py`
- `vllm_precision_compilation_test.py`
## Acceptance Criteria
- `moonshotai/Kimi-K2.6` can be loaded or converted for Megatron Core training in NeMo RL.
- A Kimi K2.6 Megatron config can complete a small SFT smoke test.
- A Kimi K2.6 Megatron config can complete a small GRPO smoke test.
- Training-backend and rollout-backend log probabilities are validated and documented.
- Example configs are added under `examples/configs/`.
- Documentation is updated to list Kimi K2.6 support, known limitations, and recommended precision/parallelism settings.
- Tests or diagnostics are added so regressions in the model mapping or checkpoint conversion path are caught.
## References
- Kimi K2.6 model card: https://huggingface.co/moonshotai/Kimi-K2.6
- NeMo RL repository: https://github.com/NVIDIA-NeMo/RL
- NeMo RL Add New Models guide: https://docs.nvidia.com/nemo/rl/latest/adding-new-models.html
- NeMo RL Megatron/MCore backend docs and README: https://docs.nvidia.com/nemo/rl/latest/index.html
Contributor guide
Assessment
This issue has not been assessed yet.