开源项目

Megatron提供FlashInfer量化MoE训练与生成log-prob一致路径

feat(moe): Train-gen parity for MXFP8 and per-token NVFP4 MoE on FlashInfer (#7798)

模型训练模型优化AI 基础设施强化学习量化分布式训练基础设施

概述

Megatron-LM用于分布式大模型训练。 此次变化前,量化MoE训练前向与推理生成采用不同内核,训练计算的token概率可能与产生轨迹时的概率不一致。 本次交付:新增默认关闭的moe_training_matches_inference,使MXFP8和逐token NVFP4训练前向使用同一FlashInfer路由内核;反向仍用TE重算。 使用条件与边界:验证和速度为4卡GB300、Nano-v3形状的单层测试;MXFP8另有权重重建与内存代价,不能外推整模型训练倍数。