SuperAPI AI 工具集

nemo-mbridge-mlm-bridge-training

智能体与模型

查看项目

更新于 10-16 07:24

技能介绍

用途:在 NVIDIA NeMo Megatron Bridge 仓库中运行 Megatron-LM(MLM)与 Bridge 两种训练流程,并做损失相关性对比测试。

主要能力:提供 MLM 与 Bridge 的入口点、启动方式与参数映射说明;给出相关性测试方法,使用 vanilla_gpt_pretrain_config 配方,使 Bridge 与 MLM 在匹配参数下 LM loss 在 BF16 舍入误差内一致;提供单 GPU 与多 GPU(TP=2)示例;列出常用 recipe,如 vanilla_gpt_pretrain_config、llama32_1b_pretrain_config、llama3_8b_pretrain_config、qwen3_8b_pretrain_config、deepseek_v2_lite_pretrain_config,SFT/PEFT 变体使用 _sft_config / _peft_config 后缀;说明 Megatron-Core 子模块的版本切换。

使用方式:Bridge 入口为 scripts/training/run_recipe.py,MLM 入口为 3rdparty/Megatron-LM/pretrain_gpt.py,两者均通过 uv run python -m torch.distributed.run 启动。MLM 需设置 PYTHONPATH=3rdparty/Megatron-LM:$PYTHONPATH。Bridge 运行前需 rm -rf nemo_experiments 清理旧检查点。子模块版本可用 ./scripts/switch_mcore.sh status|dev|main 查看或切换,切到 dev 后用 uv sync(不加 --locked)。拉取最新 main 后执行 git submodule update --init 3rdparty/Megatron-LM。

适用场景:MLM 与 Bridge 的损失相关性验证、自定义配置测试、多 GPU 张量并行训练示例参考。

限制与注意:必须使用 uv run,不能用裸 torchrun 或 python;覆盖 train.train_iters 为小值时需同时设置 scheduler.lr_warmup_iters 与 scheduler.lr_decay_iters,否则报断言错误;预训练与微调数据集 CLI 覆盖均使用 dataset.seq_length;大型 MoE 模型需完整激活重计算且通常需多节点 EP,TP 不减少每 GPU 专家内存;3rdparty/Megatron-LM/ 下文件不应从本仓库修改。