可复现性

可复现性是科研进展的基础。vime 通过结合 vLLM 的 batch-invariant 确定性推理 与 Megatron-LM 的 deterministic 模式，支持 bitwise 级的实验复现。

为了开启确定性训练，你需要通过 pip uninstall flash_attn_3 -y 卸载 flash attention 3，并设置：

  # vLLM config
  --vllm-enable-deterministic-inference
  --vllm-attention-backend flashinfer

  # megatron config
  --deterministic-mode

以及设置如下环境变量：

     "env_vars": {
        ...,
        "NCCL_ALGO": "Ring",
        "NVTE_ALLOW_NONDETERMINISTIC_ALGO": "0",
        "CUBLAS_WORKSPACE_CONFIG": ":4096:8"
     }

我们提供了一个完全确定性的，用 Qwen2.5 0.5B 训练 GSM8K 的脚本。

可以用如下脚本初始化训练数据和 ckpt：

# download
hf download --repo-type dataset zhuzilin/gsm8k --local-dir /root/gsm8k
hf download Qwen/Qwen2.5-0.5B-Instruct --local-dir /root/Qwen2.5-0.5B-Instruct

# convert ckpt
cd vime/
source scripts/models/qwen2.5-0.5B.sh
PYTHONPATH=/root/Megatron-LM/ python \
   tools/convert_hf_to_torch_dist.py \
   ${MODEL_ARGS[@]} \
   --hf-checkpoint /root/Qwen2.5-0.5B-Instruct \
   --save /root/Qwen2.5-0.5B-Instruct_torch_dist/

可以使用如下脚本进行训练：

bash scripts/run-qwen2.5-0.5B-reproducibility.sh

这个 PR 中记录了 wandb 的截图 pull#370。