已开启
[Feature]: FLOPs-per-token 由模型配置几何推导,支撑吞吐/TFLOPS/MFU 指标 #383
Robin创建于  10 天前
Robin
Robin
10 天前 创建

🚀 功能描述

训练指标的 TFLOPS/MFU 计算依赖每 token FLOPs(6N 约定)。当前该值需要在训练 YAML 中逐模型手工填写(model_flops_per_token),属于魔法数字:不同模型族、同一模型的 tiny/full 几何覆盖都需要重新计算并手工同步,容易填错或过期,导致 MFU 报表失真。

提议:FLOPs-per-token 不再作为训练配置传入,而是由模型配置几何结构自动推导——新增后端无关的估算模块 hyper_parallel/models/flops.py(支持 MLA 与 GQA/MHA 投影布局、DeepSeek 与 Qwen-MoE 两种 MoE 字段约定、可选 attention 二次项),模型族自带的 hp_flops_per_token 属性优先;序列长度取首个训练 batch 的实测值,回退 max_position_embeddings

现有替代方案

在训练配置中手工填写 model_flops_per_token(现状);或由每个模型实现各自硬编码常数。两者都需要逐模型、逐配置人工维护,且无法随几何覆盖(如调试用的缩层配置)自动适配。

与DFX相关性DF

  • 该功能属于训练可观测性(吞吐/TFLOPS/MFU 报表),上层应用(训练 recipe)无需再感知 FLOPs 计算细节;
  • 纯 Python 估算,无新增算子、无图优化,不影响训练主路径性能;
  • 功能逻辑封闭(配置几何 → 标量),不改动现有并行/训练行为,与其他功能无冲突。

提议的新API

  • hyper_parallel.models.flops.estimate_flops_per_token(config, seq_len=None)
  • hyper_parallel.models.flops.resolve_flops_per_token(model, model_config=None, seq_len=None)
  • hyper_parallel.models.flops.batch_seq_len(micro_batches)
  • BaseTrainer.add_callback(callback);环境变量 HP_THROUGHPUT_MFU=1 自动注册 ThroughputMFUCallback

是否影响现有API

移除 TrainingConfig.model_flops_per_token(手工覆盖入口),由配置推导取代;模型可选暴露 hp_flops_per_token 属性作为精确覆盖。其余训练配置与回调接口不变。

补充信息

实现 PR:#1392。估算器经 Kimi-K2.6 全量几何对拍,6N 部分与既有手工标定值 1.901164e11 完全一致。

likedislike
RobinRobin
10 天前 关联了pull request:feat: derive FLOPs-per-token from model config for throughput/MFU metrics
Mmindspore-ci-bot成员
3 天前 关联了pull request:[mirror] mindspore-ai/hyper-parallel#121: feat: derive FLOPs-per-token from model config for throughput/MFU metrics