Training Kimi-K2.6 on H200

Author

ATTUNE (cite)

GPUs

256 x H200 141 GB (AWS p5en.48xlarge, 8 per node)

Layout

TP 8, PP 16 (first stage 3 layers, last stage 2, the 14 others 4), CP 1, EP 16, ETP 1, DP 2, sequence parallel; one 8,192-token row per microbatch (max 8,192 tokens/GPU); no recompute; precision-aware distributed optimizer, no offload; alltoall dispatcher; expandable segments off

Tokens/s per GPU

353.40

$ / 1M tokens

$6.22

MFU

6.8% (N = 31,686,066,176 active parameters per token)

Step-time CV

1.58%

Convergence check

correlation 0.9999977, final loss difference 0.00013, 0 pointwise failures over 20 steps

Measured

2026-10-22

GPU-hours

63.72

Rows

kimi-k26-tulu3-fixed8192-v1 (rows_sha256 a24c09c268a1256482a379cee91397a7ff3b2e845d8d1a76eff023da753f15f2)

export CUDA_DEVICE_MAX_CONNECTIONS=1
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:False
export NCCL_NVLS_ENABLE=1
export NCCL_TIMEOUT=3600
export FI_PROVIDER=efa
export FI_EFA_USE_DEVICE_RDMA=1

MODEL_ARGS_LINE="$(python3 miles/utils/external_utils/model_args_utils.py kimi-k25)"
read -ra MODEL_ARGS <<< "$MODEL_ARGS_LINE"
python3 train.py \
  --actor-num-nodes 32 \
  --actor-num-gpus-per-node 8 \
  --num-gpus-per-node 8 \
  "${MODEL_ARGS[@]}" \
  --hf-checkpoint "$HF_CKPT" \
  --ref-load "$HF_CKPT" \
  --megatron-to-hf-mode bridge \
  --model-name kimi_k25 \
  --loss-type sft_loss \
  --calculate-per-token-loss \
  --disable-compute-advantages-and-returns \
  --debug-train-only \
  --rollout-function-path pretokenized_sft_rollout.generate_rollout \
  --prompt-data "$SFT_DATA" \
  --input-key messages \
  --metadata-key metadata \
  --num-rollout 20 \
  --rollout-batch-size 128 \
  --global-batch-size 128 \
  --tensor-model-parallel-size 8 \
  --sequence-parallel \
  --pipeline-model-parallel-size 16 \
  --decoder-first-pipeline-num-layers 3 \
  --decoder-last-pipeline-num-layers 2 \
  --context-parallel-size 1 \
  --expert-model-parallel-size 16 \
  --expert-tensor-parallel-size 1 \
  --use-dynamic-batch-size \
  --max-tokens-per-gpu 8192 \
  --optimizer adam --lr 1e-5 --lr-decay-style cosine --min-lr 1e-6 --lr-warmup-fraction 0.1 \
  --weight-decay 0.1 --adam-beta1 0.9 --adam-beta2 0.98 \
  --use-precision-aware-optimizer \
  --use-distributed-optimizer \
  --attention-dropout 0.0 --hidden-dropout 0.0 \
  --accumulate-allreduce-grads-in-fp32 \
  --attention-softmax-in-fp32 \
  --attention-backend flash \
  --no-check-for-nan-in-loss-and-grad \
  --no-gradient-accumulation-fusion

# recipes/kimi-k2.6/h200/attune-tp8-pp16-ep16.yaml
# An OpenRecipes recipe; the fields are documented in schema/recipe.schema.yaml.
schema_version: 1
model: Kimi-K2.6
gpu:
  type: H200
  memory_gb: 141
  count: 256
  per_node: 8
  instance: AWS p5en.48xlarge
author:
  name: ATTUNE
  url: https://arxiv.org/abs/2610.17528
status: verified
stack:
  framework: miles 0.1.0 (radixark/miles commit 78527b91, a slime fork), no local changes
  megatron_lm: radixark/Megatron-LM commit 37f91617 (megatron-core 0.16.0rc0), no local changes
  megatron_bridge: 0.5.0 (commit 7f0fb345; the Kimi K2.5 vision-language bridge, with the transformers 5 shim for the
    checkpoint's model code)
  transformer_engine: 2.17.0
  pytorch: 2.11.0 (CUDA 13.0 build)
  cuda: '13.0'
  nccl: 2.28.9 (CUDA 13.0 build; printed by the measured run)
  deepep: 1.2.1 (installed; the alltoall dispatcher does not use it)
  python: 3.12.3
  transformers: 5.12.1
  network: 8 GPUs per node, EFA (aws-ofi-nccl 1.20), 3.2 Tb/s RDMA per node
  source: 'version check of the measured run''s training image (same digest), 2026-10-22: miles and Megatron-LM git commits
    (clean trees), installed TE, PyTorch, megatron-core, Megatron-Bridge, DeepEP and transformers, Python, and miles''
    Kimi K2.5 launcher and model arguments; NCCL from the measured run''s log. The image is radixark/miles v0.1.0 (amd64
    digest sha256:45ae7833f97a01369dbd8f336784830a87d94587bc010c7b8065291cc15cd7ca) with EFA 1.49 added'
workload:
  row_tokens: 8192
  rows: SFT conversations
  global_batch: 128
  precision: bf16
  dataset:
    name: Tulu 3 SFT mixture
    source: allenai/tulu-3-sft-mixture
    revision: b14afda60f1bbebe55d5d2fa1e4df5042f97f8be
    license: ODC-By 1.0
    rows_sha256: a24c09c268a1256482a379cee91397a7ff3b2e845d8d1a76eff023da753f15f2
    build:
      name: kimi-k26-tulu3-fixed8192-v1
      tokenizer: Kimi-K2.6 (revision 7eb5002f6aadc958aed6a9177b7ed26bb94011bb; tiktoken.model sha256 b6c497a7469b33ced9c38afb1ad6e47f03f5e5dc05f15930799210ec050c5103;
        tokenizer_config.json sha256 12fcab43d2b6068f46769f5ff373960bf7c17a94d7abbc50e2491306b2f6cf58; tokenization_kimi.py
        sha256 2ab1ffb6f5c4380758bd8d9752ff1041c09024182676a4311528fbdf92fb9599; chat_template.jinja sha256 8bf859698fd4781c0e1e1c63ce74422aab27e53ccc5f47116317c64cda06132f)
      loss_mask: kimi_k25
      rows: 6400
      row_tokens: 8192
      padding_tokens: 0
      seed: 42
      packing: source conversations in a seeded global permutation, concatenated with an EOS token ([EOS]) between them
        and cut into consecutive 8,192-token rows; rows without supervised targets are skipped
      supervised_token_ratio: 0.666
      parquet_sha256: 7ee719299ac3ca0a386f79f0b18c11094b65416571436874b84189135f14cc69
      parquet_bytes: 91955447
      builder: data/generate_fixed_sft.py with --loss-mask-type kimi_k25 --trust-remote-code (data/README.md)
layout:
  label: TP8 PP16 EP16 DP2
  summary: TP 8, PP 16 (first stage 3 layers, last stage 2, the 14 others 4), CP 1, EP 16, ETP 1, DP 2, sequence parallel;
    one 8,192-token row per microbatch (max 8,192 tokens/GPU); no recompute; precision-aware distributed optimizer, no
    offload; alltoall dispatcher; expandable segments off
  highlights: PP 16, no recompute, expandable segments off
  tp: 8
  pp: 16
  vpp: 1
  cp: 1
  ep: 16
  etp: 1
  dp: 2
  sequence_parallel: true
  layer_split:
    first: 3
    last: 2
  max_tokens_per_gpu: 8192
  rows_per_microbatch: 1
  recompute: null
  offload: null
  dispatcher: alltoall
result:
  tokens_per_s_per_gpu: 353.4
  usd_per_million_tokens: 6.22
  step_time_cv_percent: 1.58
  mfu: 0.0679
  mfu_note: null
  active_params:
    count: 31686066176
    source: https://huggingface.co/moonshotai/Kimi-K2.6
    revision: 7eb5002f6aadc958aed6a9177b7ed26bb94011bb
  measured: 2026-10-22
  steps_per_run: 20
  window:
  - 5
  - 19
  window_tokens: 15728640
  runs: 1
  profiler: false
  gpu_hours: 63.72
  gpu_hours_note: GPUs x (scheduler admission to finish), including startup; the same run is the candidate arm of the
    convergence check
correctness:
  result: pass
  summary: passed
  reference: 'miles'' layout for Kimi K2.5 and K2.6 (`scripts/run_kimi_k25.py`, the layout it validated on 32 nodes of
    8 H200), here in BF16 without INT4 fake quantization and with Adam on the GPU: TP 8 with sequence parallel, PP 8 (5
    layers on the last stage), CP 4, EP 32, ETP 1, DP 1, full recompute (uniform, 1 layer) and max 4,096 tokens/GPU (two
    rows per microbatch), with `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True`. Model arguments, rows and schedule
    are the same in both recipes.'
  configuration: pass
  component_tests: []
  convergence:
    steps: 20
    result: pass
    correlation: 0.9999977
    final_loss_difference: 0.00013
    pointwise_failures: 0
    note: 'candidate arm: the measured run; reference arm: a 20-step run of the reference recipe on the same image and
      rows; the parity probe off in both'
  rules: argument classes of ATTUNE's skill/scripts/correctness/rules/numerics_list.yaml (amazon-science/attune aff5b81);
    20-step check by Megatron-Bridge validate_convergence with its pipeline thresholds
env:
  CUDA_DEVICE_MAX_CONNECTIONS: '1'
  PYTORCH_CUDA_ALLOC_CONF: expandable_segments:False
  NCCL_NVLS_ENABLE: '1'
  NCCL_TIMEOUT: '3600'
  FI_PROVIDER: efa
  FI_EFA_USE_DEVICE_RDMA: '1'
patches: []
run: |
  export CUDA_DEVICE_MAX_CONNECTIONS=1
  export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:False
  export NCCL_NVLS_ENABLE=1
  export NCCL_TIMEOUT=3600
  export FI_PROVIDER=efa
  export FI_EFA_USE_DEVICE_RDMA=1

  MODEL_ARGS_LINE="$(python3 miles/utils/external_utils/model_args_utils.py kimi-k25)"
  read -ra MODEL_ARGS <<< "$MODEL_ARGS_LINE"
  python3 train.py \
    --actor-num-nodes 32 \
    --actor-num-gpus-per-node 8 \
    --num-gpus-per-node 8 \
    "${MODEL_ARGS[@]}" \
    --hf-checkpoint "$HF_CKPT" \
    --ref-load "$HF_CKPT" \
    --megatron-to-hf-mode bridge \
    --model-name kimi_k25 \
    --loss-type sft_loss \
    --calculate-per-token-loss \
    --disable-compute-advantages-and-returns \
    --debug-train-only \
    --rollout-function-path pretokenized_sft_rollout.generate_rollout \
    --prompt-data "$SFT_DATA" \
    --input-key messages \
    --metadata-key metadata \
    --num-rollout 20 \
    --rollout-batch-size 128 \
    --global-batch-size 128 \
    --tensor-model-parallel-size 8 \
    --sequence-parallel \
    --pipeline-model-parallel-size 16 \
    --decoder-first-pipeline-num-layers 3 \
    --decoder-last-pipeline-num-layers 2 \
    --context-parallel-size 1 \
    --expert-model-parallel-size 16 \
    --expert-tensor-parallel-size 1 \
    --use-dynamic-batch-size \
    --max-tokens-per-gpu 8192 \
    --optimizer adam --lr 1e-5 --lr-decay-style cosine --min-lr 1e-6 --lr-warmup-fraction 0.1 \
    --weight-decay 0.1 --adam-beta1 0.9 --adam-beta2 0.98 \
    --use-precision-aware-optimizer \
    --use-distributed-optimizer \
    --attention-dropout 0.0 --hidden-dropout 0.0 \
    --accumulate-allreduce-grads-in-fp32 \
    --attention-softmax-in-fp32 \
    --attention-backend flash \
    --no-check-for-nan-in-loss-and-grad \
    --no-gradient-accumulation-fusion
notes: |
  - Starting point: miles' Kimi launcher, `scripts/run_kimi_k25.py` with `scripts/models/kimi-k25.py` at miles 0.1.0
    (model page https://miles.radixark.com/docs/models/kimi/kimi-k2.5, "Kimi K2.5 / K2.6"). miles publishes it as a
    GRPO recipe validated on 32 nodes of 8 H200 that trains the INT4 checkpoint with fake quantization and Adam on the
    CPU, with no tokens/s figure. Ported to this workload in BF16 with its layout (TP 8 with sequence parallel, PP 8,
    CP 4, EP 32, full recompute) and Adam on the GPU on 256 H200, it measured 149.6 tokens/s/GPU (149.8 in an earlier
    run of the same configuration); ATTUNE tuned it to 353.4 (+136%). A 23-step run of this recipe, profiled on its
    last step, measured 354.6 over steps 5-19.
  - Each pipeline stage spans two nodes (TP 8 x DP 2), and EP 16 spreads each MoE layer's experts over those 16 GPUs.
    With 16 stages each GPU holds 2 to 4 of the 61 layers, which leaves the memory to drop both recompute and context
    parallelism; the ranks the measured run's log shows use up to 117 GB of 140 before the end-of-step cache release.
  - With `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True`, as in the reference run, a tuning run of this layout
    measured 216.3 tokens/s/GPU. miles empties the CUDA cache after every step (`clear_memory` in
    `miles/utils/memory_utils.py`), so the allocator takes its memory from the driver again each step.