Files
laoyao_2b_moe/scripts/README.md

5.8 KiB
Raw Blame History

Scripts

  • sync_pretrain_data_into_repo.sh: 200B 数据构建完成后,把数据目录同步到 dataset/pretrain/data/,默认优先 hardlink。
  • wait_and_sync_pretrain_data.sh: 后台等待当前 200B 构建进程结束,然后自动同步数据。
  • preprocess_megatron_bridge_pretrain.sh: 旧的 Megatron indexed dataset 预处理入口,保留用于对照。
  • preprocess_megatron_bridge_pretrain_direct.sh: 直接从 parquet 生成 Megatron indexed dataset不落中间 JSONL。
  • train_megatron_bridge_2b_moe.sh: 当前主训练入口,使用 NeMo 26.06 镜像中的 Megatron-Bridge。
  • train_nemo_megatron_2b_moe.sh: NeMo/Megatron 训练入口占位,包含 image、mount、路径检查。
  • g0050_download_and_setup_from_modelscope.sh: 在 g0050 上一键准备训练环境并从 ModelScope 下载未 tokenize parquet 数据。
  • g0050_wait_and_tokenize_glm52_8192.sh: 等待 g0050 ModelScope parquet 下载完成后,按 GLM-5.2 tokenizer 和 max_seq_len=8192 转成 Megatron indexed dataset。
  • resume_pretrain_8192_8gpu_mbs14.sh: g0050 当前主恢复训练入口,默认 8 卡 DP、mbs=14gbs=112、full recompute、distributed optimizer、2500 iter 保存、15000 iter validation。
  • serve_laoyao_megatron.sh / serve_laoyao_megatron_inner.sh: 启动 Megatron 原生 text generation server用于 checkpoint 推理 smoke。
  • laoyao_megatron_inference_setup.sh: 容器内热补丁 Megatron inference server 的参数兼容和返回格式问题。
  • check_laoyao_megatron_inference_ready.sh: 检查 checkpoint、端口、GPU 和相关容器状态。
  • query_laoyao_megatron_server.sh: 对 Megatron inference server 发送简单 prompt。
  • stop_laoyao_megatron_server.sh: 停止 Megatron inference server 容器。
  • graceful_stop_megatron_training.sh: 向 torchrun 的训练 worker 发送信号,在当前 iteration 完成后临时保存 checkpoint 并安全退出。
  • record_active_training.sh: launcher 启动容器后原子记录当前 container、run 和 checkpoint 路径,供状态检查与无参数 graceful stop 使用。

g0050 下载与部署

在 Mac 侧通过 B300 跳转到 g0050

ssh B300 'ssh ubuntu@g0050 "cd /ssd/workspace/yi/laoyao_2b_moe && MODELSCOPE_API_TOKEN=ms-... bash scripts/g0050_download_and_setup_from_modelscope.sh"'

默认行为:

  • repo 路径:/ssd/workspace/yi/laoyao_2b_moe
  • 数据路径:/ssd/workspace/yi/laoyao_2b_moe_pretraining_dataset
  • ModelScope dataseteigentom/laoyao_2b_moe_pretrain_parquet_20260702
  • 训练镜像:nvcr.io/nvidia/nemo:26.06
  • 下载镜像ModelScope CUDA 13.0 / Swift 4.3.1 官方镜像
  • 代理:默认使用 B300/g0050 侧的 http://100.72.0.101:8888

私有 Gitea clone 时不要把 token 写进脚本,运行时通过环境变量传入:

GIT_REPO_URL=https://yi_lu:<token>@git.deeepseek.net/yi_lu/laoyao_2b_moe.git \
MODELSCOPE_API_TOKEN=ms-... \
bash scripts/g0050_download_and_setup_from_modelscope.sh

g0050 当前训练恢复

在 g0050 上:

cd /ssd/workspace/yi/laoyao_2b_moe
bash scripts/resume_pretrain_8192_8gpu_mbs14.sh

默认配置与当前长跑实验一致:

  • data manifest: /ssd/workspace/yi/laoyao_2b_moe_pretraining_dataset/megatron_bridge/pretrain_8192_glm52_direct_v1/manifest.json
  • validation prefix: dataset/val/megatron_8192_glm52/heldout_2p8k_text_document
  • train iters: 184209
  • seq_length=8192
  • micro_batch_size=14, global_batch_size=112
  • save_interval=2500, eval_interval=15000, eval_iters=10
  • --use-distributed-optimizer --overlap-grad-reduce --overlap-param-gather
  • --recompute-granularity full --recompute-method uniform --recompute-num-layers 1

Stage1-2 多 shard 恢复训练

resume_stage1_2_pretrain_8192_8gpu_mbs14.sh 用于从 stage1-2 Megatron checkpoint 恢复。stage1-2 weighted manifest 包含数百个 indexed dataset prefix每个 rank 会同时 mmap 对应的 .bin.idx 和 dataset index 文件。Docker 默认 nofile=1024 会在数据集初始化阶段触发:

OSError: [Errno 24] Too many open files

因此该入口固定设置:

--ulimit nofile=1048576:1048576

不要通过减少 shard 或丢弃数据绕过这个限制。恢复前应确认 LOAD_DIR 下的 latest_checkpointed_iteration.txt 指向目标 checkpoint并保留 load_optim=Trueload_rng=True 以恢复优化器与 RNG 状态。

信号触发的临时 checkpoint

laoyao_2b_moe_pretrain.py 默认启用 Megatron-Bridge 原生 distributed signal handler并监听 SIGINT。收到信号后,各 rank 会:

  1. 完成当前 iteration。
  2. 使用标准 checkpoint 路径保存模型、优化器、LR scheduler、RNG 和数据迭代状态。
  3. 完成分布式 barrier 后退出。

不要直接执行 docker kill --signal=SIGINT <container>。容器 PID 1 是 torchrun;它收到 SIGINT 后会启动 elastic worker 回收流程,可能在 checkpoint 写完前终止 rank。应运行

bash scripts/graceful_stop_megatron_training.sh

脚本通过 docker top 只选择 torchrun 的直接 worker 子进程,不会误发给 dataloader worker也不会先终止 torchrun。默认等待 900 秒且不会在超时后自动 强杀容器。状态默认读取:

.runtime/active_training.txt

所有长期训练 launcher 的默认容器名固定为:

laoyao-2b-pretraining

训练 launcher 确认容器启动后还会原子更新状态文件,记录 checkpoint 目录、run 目录和日志路径。即使状态文件缺失,停止脚本也会操作上述固定容器名。多实验并存时 仍可显式覆盖:

CHECKPOINT_DIR=/path/to/checkpoints \
bash scripts/graceful_stop_megatron_training.sh <container-name>

训练入口也支持显式切换信号:

--graceful-exit-signal SIGINT   # 默认
--graceful-exit-signal SIGTERM
--disable-graceful-exit-on-signal