千卡级训练配置：Kimi-K2

本文介绍如何在 MTT S5000 集群上配置 Kimi-K2 千卡级 MoE 训练，并以 62 台机器、496 张卡的训练脚本为例。内容基于现有多机启动脚本的并行方式，并结合 Kimi-K2 官方模型配置说明模型结构。文中的主机名、目录和数据路径均为占位符。

本文不提供吞吐量、MFU、训练耗时等性能指标。

本文中的命令分为两类：标注“宿主机执行”的命令在 Docker 宿主机运行，其他命令默认在训练容器内运行；带有 <...> 的内容必须替换为当前环境的实际值后再执行。正式训练前请记录镜像 digest、代码 commit、DeepEP 版本、数据版本和启动参数，保证任务可复现。

1. 环境准备

1.1 千卡级拓扑

run_kimi_496GPU.sh 采用以下并行组合：

参数	示例值	说明
WORLD_SIZE	496	应与 hostfile 总 slot 一致
TP	1	张量并行
PP	31	流水线并行，调整时必须重新检查层分配
EP	8	专家并行
DP	2	496 ÷ (TP × PP × EP)
Micro batch	1	单次前向/反向的每卡样本数
Microbatch 数	32	用于计算 global batch

必须满足：

WORLD_SIZE = TP × PP × EP × DP
GLOBAL_BATCH_SIZE = MICRO_BATCH_SIZE × NUM_MICROBATCHES × DP

PP=31 与现有模型分层和启动脚本相关，不能在不检查流水线分层的情况下直接修改。

WORLD_SIZE 必须等于 hostfile 中所有 slots 的总和。调整 TP、PP 或 EP 后，应重新计算 DP，并先执行小规模 dry-run。EP、DP 的 rank 分组方式由当前 Megatron MUSA Patch 和启动脚本决定，不能只依据公式手工推断。

1.2 创建容器

以下命令在每台 Docker 宿主机执行。示例使用 v2.1.7 镜像；实际训练时应替换为近期已验证、且预先安装了与训练代码匹配 DeepEP 的 Training Suite 镜像：

docker pull registry.mthreads.com/mcctest/training-suite:v2.1.7

# 可选：记录实际使用的镜像摘要，确保所有节点使用同一镜像
docker image inspect registry.mthreads.com/mcctest/training-suite:v2.1.7 \
  --format '{{index .RepoDigests 0}}'

docker run -d \
  --name training_kimi_k2 \
  --runtime=mthreads \
  --privileged \
  --network=host \
  --shm-size=256g \
  -e MTHREADS_VISIBLE_DEVICES=all \
  -v <SHARED_ROOT>:/data/shared \
  -v <CODE_ROOT>:/data/code \
  registry.mthreads.com/mcctest/training-suite:v2.1.7 \
  sleep infinity

进入容器：

docker exec -it training_kimi_k2 /bin/bash

1.3 检查 DeepEP 和运行环境

Kimi-K2 MoE 训练需要确认镜像中包含与训练代码匹配的 DeepEP：

mthreads-gmi
cat /usr/local/musa/version.json

python3 - <<'PY'
import importlib.util
import torch
import torch_musa

print("PyTorch:", torch.__version__)
print("torch_musa:", torch_musa.__version__)
print("MUSA available:", torch.musa.is_available())
print("DeepEP module:", importlib.util.find_spec("deep_ep"))
PY

如果 deep_ep 未找到，停止后续操作，改用包含已验证 DeepEP 的镜像或按项目发布说明安装匹配版本。

还应记录实际包版本：

python3 -m pip show deep-ep 2>/dev/null || true
git -C /data/code/megatron-lm-musa-patch rev-parse HEAD

2. 准备目录和 hostfile

2.1 目录约定

/data/shared/kimi-k2/
├── data/
├── tokenizer/
├── checkpoints/
└── logs/

/data/code/megatron-lm-musa-patch/
└── examples/kimi-k2/

2.2 创建 hostfile

本示例使用 62 个节点、每节点 8 张 GPU，总计 496 个 slot：

train-node-000 slots=8
train-node-001 slots=8
...
train-node-061 slots=8

检查 hostfile：

HOSTFILE=/data/shared/kimi-k2/hostfile

grep -Ev '^#|^$' "${HOSTFILE}" | wc -l
awk -F'slots=' '/slots=/{sum += $2} END{print sum}' "${HOSTFILE}"

不要在公开文档中写入生产 IP、用户名或 SSH 凭据。

建议在提交任务前校验规模：

NODE_COUNT=$(grep -Ev '^#|^$' "${HOSTFILE}" | wc -l)
SLOT_COUNT=$(awk -F'slots=' '/slots=/{sum += $2} END{print sum}' "${HOSTFILE}")
test "${NODE_COUNT}" -eq 62
test "${SLOT_COUNT}" -eq 496

hostfile 的顺序通常会影响 rank 分配。第一个节点是否作为 master、节点名是否可被所有节点解析，以及 launcher 使用的 SSH 用户和端口，必须以 run_kimi_496GPU.sh 的实现为准。

3. 准备模型和数据

3.1 模型结构

Kimi-K2 官方配置包含以下核心参数：

MODEL_ARGS=(
  --num-layers 61
  --hidden-size 7168
  --num-attention-heads 64
  --num-experts 384
  --moe-router-topk 8
  --seq-length <SEQ_LENGTH>
  --max-position-embeddings 131072
  --bf16
)

Kimi-K2 使用 MLA 和 MoE。实际 Megatron 参数名称可能因 MUSA Patch 版本不同而变化，应以当前代码的帮助信息和已验证脚本为准，不要只根据 Hugging Face 字段机械转换。

3.2 数据和 Tokenizer

export DATA_PATH=<MEGATRON_DATA_PREFIX>
export TOKENIZER_MODEL=<KIMI_K2_TOKENIZER_DIR>

ls -lh "${DATA_PATH}.bin" "${DATA_PATH}.idx"
test -d "${TOKENIZER_MODEL}"

多数据集场景需要确认数据权重、路径顺序和所有节点的可见性一致。

DATA_PATH 通常是不带扩展名的 Megatron 数据前缀。例如，若数据文件为 train_text_document.bin 和 train_text_document.idx，应设置为：

export DATA_PATH=/data/shared/kimi-k2/data/train_text_document

Kimi-K2 的模型权重、Tokenizer 和数据预处理方式必须与启动脚本保持一致。本文未提供内部模型和数据下载地址时，不应直接使用任意公开权重替代。

4. 配置训练任务

4.1 并行和 batch 参数

TP_SIZE=1
PP_SIZE=31
EP_SIZE=8
WORLD_SIZE=496
MICRO_BATCH_SIZE=1
NUM_MICROBATCHES=32

DP_SIZE=$((WORLD_SIZE / (TP_SIZE * PP_SIZE * EP_SIZE)))
GLOBAL_BATCH_SIZE=$((MICRO_BATCH_SIZE * NUM_MICROBATCHES * DP_SIZE))

test $((TP_SIZE * PP_SIZE * EP_SIZE * DP_SIZE)) -eq "${WORLD_SIZE}"

4.2 MoE 相关能力

现有训练方案包含以下能力，使用前需要逐项确认当前脚本确实启用：

• Full recompute；
• SwiGLU fusion；
• DeepEP 专家通信；
• 分布式优化器；
• 专家负载统计与异常检测。

不要在没有精度验证的情况下随意切换 MoE 路由、专家数量、top-k 或重计算粒度。

4.3 检查启动脚本

cd /data/code/megatron-lm-musa-patch/examples/kimi-k2
find . -maxdepth 3 -type f -name '*.sh' -print

若未找到，496 卡脚本压缩包可从这里下载：

wget -O run_kimi_496GPU.tar.gz \
  https://mt-ai-infra.tos-cn-beijing.volces.com/run_kimi_496GPU.tar.gz
tar -xzvf run_kimi_496GPU.tar.gz

检查启动脚本并设置任务参数：

export HOSTFILE=/data/shared/kimi-k2/hostfile
export DATA_PATH=<MEGATRON_DATA_PREFIX>
export TOKENIZER_MODEL=<KIMI_K2_TOKENIZER_DIR>
export CHECKPOINT_DIR=/data/shared/kimi-k2/checkpoints
export LOG_DIR=/data/shared/kimi-k2/logs

编辑 run_kimi_496GPU.sh，将其中的 DATA_PATH 修改为当前环境可访问的数据集路径。所有训练节点都必须能够访问该路径，并使用一致的数据配置。脚本文件名不能代表实际规模，必须检查其中的 WORLD_SIZE、hostfile、TP、PP、EP 和 batch 计算。

同时确认脚本实际使用的 launcher 以及 MASTER_ADDR、MASTER_PORT、NNODES、NODE_RANK、NPROC_PER_NODE 等变量。不要在多个节点上重复手工执行该脚本，除非脚本明确要求每个节点分别启动。

5. 启动训练

5.1 缩小规模 dry-run

正式扩展前先使用可整除 TP × PP × EP 的较小规模验证：

• 模型能够初始化；
• MoE 专家路由和 DeepEP 通信正常；
• loss、梯度和专家负载为有限值；
• 检查点能够保存和恢复；
• 数据读取、共享存储和日志路径正常。

5.2 启动千卡级任务

cd /data/code/megatron-lm-musa-patch/examples/kimi-k2
bash run_kimi_496GPU.sh

该脚本适用于 62 台机器、496 张卡的任务，并开启 Full recompute、SwiGLU fusion 和 DeepEP。如果脚本通过 SSH 在所有节点启动进程，应先在少量节点验证远端工作目录、环境变量和日志重定向行为一致。

首次执行建议先检查脚本参数而不提交全量任务：

bash run_kimi_496GPU.sh --help

如果脚本不支持 --help，应先阅读脚本中实际执行的 launcher 和参数，再进行 dry-run。

6. 监控训练

tail -F /data/shared/kimi-k2/logs/<RANK0_LOG>

重点检查：

• 所有 rank 是否完成初始化；
• EP、PP 和 DP 通信是否出现 hang 或超时；
• loss、梯度范数及专家负载是否正常；
• 是否出现 NaN、Inf、OOM 或数据读取中断；
• 检查点是否完整写入共享存储。

本文不记录或展示性能指标。

7. 断点恢复和停止任务

恢复训练时必须使用同一套模型、优化器、RNG、数据位置和拓扑信息。改变 PP 或 EP 后不能直接假设旧检查点可加载，应先完成重分片验证。

恢复参数名以脚本实现为准，常见形式如下：

export CHECKPOINT_DIR=/data/shared/kimi-k2/checkpoints/<STEP_DIR>
bash run_kimi_496GPU.sh --load "${CHECKPOINT_DIR}"

如果脚本没有 --load 参数，应在脚本中设置对应的 checkpoint 环境变量。恢复前先复制或备份 checkpoint 目录，避免试运行覆盖原始检查点。

停止任务：

pkill -f 'run_kimi_496GPU.sh'

pkill 可能误杀同一节点上的其他任务。生产环境应优先使用 launcher 返回的任务 ID 或 PID 文件停止任务，并在所有节点确认没有残留训练进程。

FAQ

Q1: 如何确认所有节点的软件版本完全一致？

除操作系统、驱动、MUSA SDK、PyTorch、torch_musa、镜像和代码提交外，Kimi-K2 还必须检查 DeepEP 与 Megatron MUSA Patch 版本。

HOSTFILE=/data/shared/kimi-k2/hostfile

while read -r host _; do
  echo "===== ${host} ====="
  ssh "${host}" '
    uname -r
    cat /usr/local/musa/version.json
    python3 -c "import torch, torch_musa; print(torch.__version__, torch_musa.__version__)"
    python3 -m pip show deep-ep
    git -C /data/code/megatron-lm-musa-patch rev-parse HEAD
  '
done < <(grep -Ev "^#|^$" "${HOSTFILE}")

任一节点缺少 DeepEP、版本不同或代码提交不同，都应先修复或从 hostfile 中移除。

Q2: 如何确认全部 GPU 都可用？

在每个节点检查：

mthreads-gmi

python3 - <<'PY'
import torch
import torch_musa

assert torch.musa.is_available()
count = torch.musa.device_count()
print("MUSA device count:", count)
assert count == 8, f"expected 8 devices, got {count}"
PY

同时检查是否存在掉卡、硬件错误、异常温度或残留进程。hostfile 的 slot 数必须与容器实际可见 GPU 数一致。

Q3: 如何确认 MCCL 和 DeepEP 通信正常？

先执行单节点 MCCL 测试，再扩展到两节点、一个 EP 组和完整训练集群。至少覆盖 AllReduce、AllGather、ReduceScatter 和点对点通信：

<MCCL_TEST_COMMAND>

MCCL 正常后，再用缩小模型执行 DeepEP 专家通信 dry-run。若只有 EP 训练失败，应重点检查 DeepEP 安装、专家 rank 分组、网络环境和进程组初始化。

Q4: 启动时部分专家 rank 报错怎么办？

检查：

WORLD_SIZE = TP × PP × EP × DP

并确认：

• WORLD_SIZE 能被 TP × PP × EP 整除；
• hostfile 总 slot 等于 WORLD_SIZE；
• DeepEP、MCCL、rank map 使用同一组节点；
• PP 分层与模型层数匹配；
• 所有节点读取相同的数据和 Tokenizer。

Q5: 专家负载异常或 loss 出现 NaN 怎么办？

先暂停扩容并检查：

1. MoE top-k、专家数量和路由配置是否与模型一致；
2. DeepEP 和 SwiGLU fusion 是否使用匹配版本；
3. BF16、FP8 和 loss scale 设置；
4. 数据样本和 Tokenizer 是否异常；
5. 梯度裁剪和学习率；
6. 各节点是否加载了同一模型配置和检查点。

不要同时修改路由、batch 和精度配置，否则难以定位根因。

Q6: 为什么不直接给出固定的数据和检查点路径？

数据和检查点路径属于部署配置，而且必须在所有节点上保持一致可见。本文使用占位符，避免公开内部存储地址。实际部署时应把路径写入统一配置或环境变量，不要逐节点手工修改。

相关资料

• Kimi-K2-Base 官方模型配置
• TrainingSuite v2.1.7 发布说明