用途:本技能是 Megatron-Bridge 中为新的 LLM 或 VLM 模型添加支持的指南,覆盖 bridge、provider、recipe、测试、文档与示例。
主要能力与流程:
- 发现阶段:获取 HuggingFace 模型链接或 config.json,提取 model_type、architectures、tie_word_embeddings、层数、隐藏维度、注意力头数、MoE 与 MLA 字段等;判断 VLM(含 text_config 与 vision_config)还是 LLM;检查 FP8/FP4 量化权重。
- 桥接支持:LLM 通常只需 bridge 文件,实现 provider_bridge() 与 mapping_registry();VLM 可能需要 provider 子类与组合视觉+语言的 model.py。模型专属代码应放在 modeling_<model> 命名空间下,避免修改共享文件。可通过本地映射子类或桥接钩子(如 maybe_modify_loaded_hf_weight()、megatron_to_hf_config())扩展。
- Recipe:在 recipes/<family>/<model>.py 中为各尺寸与训练模式定义 sft、peft、pretrain 配置函数。
- 测试:单元测试验证 provider 映射,功能测试做 HF↔Megatron 往返转换。
- 文档与示例:提供 README、inference.sh、slurm_sft.sh、slurm_peft.sh,README 需含真实模型的导入导出命令。
限制与注意:量化权重不会自动反量化,需在导入时处理,否则模型静默损坏;新增计算块需更新 flop_utils.py 中的 FLOPs 计算器;共享文件仅在可复用于两个以上模型家族时才修改。具体命令与代码示例以原文为准。