Kimi-K3 lora RL day-0 support

[yueming-yuan]

Blog

check details and experiment results in:
https://www.lmsys.org/blog/2026-07-27-kimi-k3-day0-support

Usage

sglang branch: https://github.com/sgl-project/sglang/tree/sglang-miles-k3
Docker image: docker.io/radixark/miles:kimi-k3 (multi-arch amd64/arm64; on Hopper set SGLANG_K3_ATTN_RES_MODE=jit)

About

Megatron training backend for Kimi K3 plus the colocated SGLang rollout path
it needs for RL. KDA and MLA attention per layer, the attention-residual
snapshot bank, and native LoRA adapters under TP/EP/PP/CP — with shared-A /
per-expert-B factors for the 896-expert MoE, exported to the rollout engines
as HF-named chunks over CUDA IPC.

Reproduce

The image ships miles (/root/miles, editable) and sglang
(/sgl-workspace/sglang, editable) at the validated pins, flashinfer
0.6.15.post1, and a small shm_unlink shim preloaded via
/etc/ld.so.preload (tolerates a benign PyTorch CUDA-IPC unlink race that
otherwise aborts colocated weight sync at scale).

Prepare (once). Only external asset: the Kimi K3 MXFP4 HF checkpoint.
Everything else derives in-repo (run from /root/miles, Megatron on
PYTHONPATH):

python scripts/run_kimi_k3_lora.py prepare-data --task dapo-math --data-dir <datasets>

python tools/convert_mxfp4_to_bf16.py --model-dir <native-mxfp4> --save-dir <bf16-hf>

source scripts/models/kimi-k3.sh   # defines MODEL_ARGS
torchrun --nnodes=8 --nproc-per-node=4 ... \
    tools/convert_hf_to_torch_dist.py "${MODEL_ARGS[@]}" \
    --hf-checkpoint <bf16-hf> --save <torch-dist-dcp> \
    --bf16 --tensor-model-parallel-size 32 --sequence-parallel \
    --pipeline-model-parallel-size 1 --context-parallel-size 1 \
    --expert-model-parallel-size 32 --expert-tensor-parallel-size 1 \
    --megatron-to-hf-mode raw   # 32 ranks; the torch_dist output re-shards at load

Run (validated on 16 nodes x 4 GPUs). One container per node, bring up a
ray cluster across the nodes, export MILES_SCRIPT_EXTERNAL_RAY=1, then:

python scripts/run_kimi_k3_lora.py train \
  --mode normal --model-variant full --task dapo-math --reward-model deepscaler \
  --num-nodes 16 --num-gpus-per-node 4 \
  --pipeline-parallel-size 8 --context-parallel-size 2 \
  --rollout-tp-size 16 --rollout-max-concurrency 8 \
  --lora-rank 32 --lora-alpha 64 \
  --num-rollout 1000 --rollout-batch-size 8 --n-samples-per-prompt 8 \
  --rollout-max-response-len 4096 --sglang-max-total-tokens 65536 \
  --global-batch-size 64 --lr 1e-05 --eval-interval 10 \
  --distributed-timeout-minutes 60 \
  --hf-checkpoint <native-mxfp4> --ref-load <torch-dist-dcp> \
  --sglang-path /sgl-workspace/sglang/python \
  --data-dir <datasets> --enable-wandb

(--rollout-max-concurrency 8 explicitly: the field default is 64, the
validated runs pin 8.)

Expected. Resolved config at startup shows expert_model_parallel_size 8,
max_tokens_per_gpu 8192, colocate_memory_peak_device gpu,
lora_base_cpu_backup True. On our GB300 validation runs: ~11-13 min per
rollout cycle; trainer allocated returns to ~91 GB after every weight sync;
rollout/raw_reward 0.5-0.75 from rollout 0; eval/aime 0.37-0.43 at eval@0
(temp-0 nondeterminism band), rising >= +0.06 by eval@9 (we measured
0.367 -> 0.467).

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.