SuperAPI AI 工具集

nemo-mbridge-mlm-bridge-training

智能体与模型

查看项目

更新于 10-16 03:49

技能介绍

用途:本技能用于运行 Megatron-LM(MLM)与 Megatron Bridge 的训练,支持 mock 数据或真实数据,并用于在两者之间做损失相关性(correlation)测试。

先答清单:涉及 MLM 与 Bridge 相关性问题时,需先指出 Bridge recipe 为 vanilla_gpt_pretrain_config;Bridge 入口为 scripts/training/run_recipe.py;MLM 入口为 3rdparty/Megatron-LM/pretrain_gpt.py;两者启动封装均为 uv run python -m torch.distributed.run;Bridge 运行前需 rm -rf nemo_experiments。MLM 需要设置 PYTHONPATH=3rdparty/Megatron-LM:$PYTHONPATH;参数匹配时两者损失应在 BF16 舍入误差内一致;不应修改 3rdparty/Megatron-LM/ 下文件。

主要能力与用法:用 vanilla_gpt_pretrain_config 做损失相关性测试,该 recipe 使用裸 GPTModelProvider 默认值(LayerNorm、GeLU、学习式绝对位置嵌入、词表继承自 tokenizer),与无参数 MLM pretrain_gpt.py 默认一致。文档给出 1 GPU 的 MLM 与 Bridge 对照命令,以及 TP=2 的 2 GPU 示例。验证方式是比对两侧日志中的 lm loss,参数匹配时每轮应几乎相同,一致程度受 BF16 舍入影响。

可用 recipe:vanilla_gpt_pretrain_config(最小 GPT,适合相关性测试与自定义配置)、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 子模块:可用 ./scripts/switch_mcore.sh status 查看版本,./scripts/switch_mcore.sh dev 切换到 dev 后执行 uv sync(不加 --locked,因锁文件对应 main),./scripts/switch_mcore.sh main 切回;拉取最新 main 后用 git submodule update --init 3rdparty/Megatron-LM 重新同步子模块。

限制与注意:新相关性运行前必须 rm -rf nemo_experiments,否则 Bridge 会静默从旧 checkpoint 自动恢复;必须用 uv run python -m torch.distributed.run,而非裸 torchrun 或 python;覆盖 train.train_iters 为小值时须同时设置 scheduler.lr_warmup_iters 与 scheduler.lr_decay_iters,否则触发断言错误;CLI 覆盖使用 dataset.sequence_length 而非 dataset.seq_length;大型 MoE 需要完整激活重计算且通常多节点 EP,TP 不降低每 GPU 专家内存;切到 dev 后 uv sync --locked 会失败。其余信息暂未提供。