已开启
DeepSeek-V3 16B FSDP2 MXFP8/HiF8 训练设计 #329
zhouyaqiang0创建于  8月12日
zhouyaqiang0
8月12日 创建

DeepSeek-V3 16B FSDP2 MXFP8/HiF8 训练设计

1. 目标与范围

本设计在 HyperParallel Trainer 中打通 DeepSeek-V3 16B 的 MXFP8/HiF8 在线低精预训练,
目标运行环境为 Ascend A5 + PyTorch/torch_npu。两种低精格式复用同一套模型加工、
forward/dgrad/wgrad、FSDP2、Optimizer 和 DCP 生命周期,只在量化策略、物理存储和 NPU
算子参数上分开实现。

本期支持:

  • 根据 Hugging Face DeepseekV3Config 构建冻结的 16B 结构并随机初始化;
  • TP=1、CP=1、EP=1、PP=1,使用 FSDP2 做数据并行分片;
  • Dense FFN、shared experts、routed experts 和 MLA projection Linear 的
    forward/dgrad/wgrad 低精计算;
  • mxfp8_e4m3 + mx_block
  • hif8 + current scaling;
  • FP32 master 权重和 FP32 AdamW 状态;
  • 梯度累积、activation recompute;
  • DCP 模型与优化器 checkpoint 保存、恢复。

本期不支持:

  • TP、CP、EP 和 PP;
  • MXFP4 或其他低精格式;
  • HiF8 delayed scaling;
  • FSDP 通信融合和低精梯度通信;
  • 低精 All-to-All;
  • PEFT wrapper、Linear 子类、自定义 fused module 和 forward 中直接调用的 F.linear/mm
  • 非 DCP checkpoint;
  • Precision Observer;
  • 低精 RMSNorm、SwiGLU 或 attention 融合算子。

不提供 BF16/FP32 计算 fallback。fallback_to_unsharded 只允许不满足低精分片约束的参数
保持复制,参数对应的 MM/GMM 仍使用所选低精格式。

1.1 模型低精边界

低精组件范围同时参考论文边界和开源工程实现:

  • DeepSeek-V3 Technical Report 3.3.1
    将主要 GEMM 的 Fprop、Dgrad 和 Wgrad 放到 FP8,同时将 embedding、output head、MoE
    gate、normalization 和 attention operators 保持高精。论文没有逐个枚举 MLA projection,
    因而不能仅凭该表判断 projection Linear 是否属于这里的 attention operators。
  • Megatron-LM 的
    MLA layer spec
    将 Q、KV 的 down/up projection 和 output projection 都交给 backend Linear 构建;其
    Transformer Engine backend
    对应 TELinearTEColumnParallelLinearTERowParallelLinear。MLA 实现还明确处理
    linear_proj 保存量化输入
    说明 projection Linear 位于 FP8 路径。
  • Megatron 将核心 attention 的低精单独控制,
    fp8_dot_product_attentionfp8_multi_head_attention
    均默认关闭。因此其默认边界是“projection Linear 可低精,QK、Softmax、PV 等核心
    attention 计算保持高精”。这是工程实现证据,不反向证明 DeepSeek-V3 原始训练对每个
    MLA projection 都使用了 FP8。

本设计对 MXFP8 和 HiF8 使用相同的模块边界:MLA projection 作为普通 GEMM 进入低精;
核心 attention、RoPE 和 Q/K normalization 保持 BF16/FP32。目标必须按完整 FQN 显式列出,
不使用 self_attn.*_proj 或全模型 *_proj 通配。

DeepSeek-V3 组件 本期精度 处理方式
Dense FFN gate_proj/up_proj/down_proj MXFP8 或 HiF8 exact nn.Linear → LowPrecisionLinear
Shared experts gate_proj/up_proj/down_proj MXFP8 或 HiF8 exact nn.Linear → LowPrecisionLinear
Routed experts gate_up_proj/down_proj MXFP8 或 HiF8 整体替换 Experts,使用对应格式 GMM
MLA q/kv down/up projection、o_proj MXFP8 或 HiF8 按完整 FQN 替换 exact nn.Linear
Embedding、LM head BF16/FP32 不进入 replacement
MoE gate/router BF16/FP32 不进入 replacement
RMSNorm 与其他 normalization BF16/FP32 不进入 replacement
MLA 核心 QK、Softmax、PV、RoPE BF16/FP32 不进入 replacement

2. 总体架构

flowchart LR
    YAML[Trainer YAML] --> PLAN[Model Plan / Replacement Plan]
    PLAN --> DENSE[LowPrecisionLinear]
    PLAN --> DSV3[DeepseekV3LowPrecisionExperts]
    DENSE --> FUNC[Low-precision Functional]
    DSV3 --> FUNC
    FUNC --> RECIPE[Recipe / Role Quantizers]
    FUNC --> OPS[NPU MM/GMM Ops]
    RECIPE --> QT[MXFP8Tensor / HiF8Tensor]
    QT --> LIFE[FSDP / Optimizer / DCP Adapter]
层次 负责 不负责
模型加工框架 FQN 匹配、严格类型检查、模块替换、应用顺序 低精算法和 DeepSeek 语义
DeepSeek 模型适配 Experts 参数布局、forward 签名、top-k 输出聚合 量化数学和 NPU kernel
低精通用模块 Tensor、Recipe、Quantizer、MM/GMM、autograd 模型遍历、Router、集合通信
格式实现 MXFP8/HiF8 量化格式、scale 语义和算子参数 Trainer 和模型语义
生命周期适配 将低精 Tensor 接入 FSDP、Optimizer、DCP 复制这些组件的内部实现

依赖只能沿上图向下。DeepSeek adapter 可以依赖低精通用模块,低精通用模块不得反向导入
DeepSeek;HyperParallel core 只识别通用扩展协议,不得依赖 MXFP8 或 HiF8 类型。

3. 两种格式的统一与差异

项目 MXFP8 HiF8 current scaling
data dtype E4M3 torch_npu.hifloat8
scale dtype E8M0 FP32
scale 粒度 1×32/32×1 block Dense 标量;Grouped 按 leading group
scale 来源 当前输入的 MX block 当前输入 amax
量化方向 row/col 不同物理表示 data 可复用,operand 描述方向
持久化权重存储 row_data/row_scale + col_data/col_scale 一份 data + scale
forward/backward 使用方式 按 GEMM 方向选择不同物理表示 复用 data,通过 layout/transpose 选择方向
FSDP all-gather 按阶段通信 row 或 row+col 通信 data+scale,不区分前反向方向
分片主要约束 MX scale tile 必须与本地 shard 边界一致 Dense scale 跨 FSDP rank 一致;Grouped 不切开单个 expert
运行时状态
checkpoint 额外状态

3.1 核心差别

MXFP8 对每个 32 元素 block 独立缩放。更细的 scale 粒度能够适应同一 Tensor 内不同区域的
数值范围,但 row/col 量化方向对应不同的物理 data/scale,权重需要同时保存两个方向:

MXFP8 weight
├── row_data + row_scale(E8M0)
└── col_data + col_scale(E8M0)

FSDP forward 可以根据计算类型只收集一个方向;backward 同时覆盖 dgrad/wgrad,需要完整的
row 和 col 表示。分片边界还必须与 MX scale tile 对齐。

HiF8 current scaling 的 Dense Tensor 使用当前 Tensor 的一个 FP32 scale,Grouped Tensor
使用按 expert/group 划分的 scale。scale 不依赖 row/col block,因此同一份量化 data 可以用于
不同 GEMM 方向:

HiF8 weight
├── data(hifloat8)
└── scale(FP32 scalar or [E])

FSDP forward/backward 均收集同一份 data+scale。Dense 权重在各 rank 从 FP32 master shard
刷新前,需要对 amax 做 FSDP-group MAX reduce,保证所有 shard 使用同一 scale;Grouped
权重的 scale 与完整 expert 一起分片和重建。

从格式特征推测,MXFP8 的细粒度 block scale 对局部数值分布差异适应更强,但物理存储、
量化和通信组织更复杂;HiF8 的权重表示和通信更简单,但 Dense per-tensor scale 更容易受
异常大值影响。该描述只是设计层面的预期,最终精度、吞吐和显存结论必须通过 A5 相同模型、
相同 seed 的对比测试确定。

3.2 共用训练流程

两种格式不建立两套模型或 Trainer 流程,共用:

  • plan_overrides 和 Module Replacement Plan;
  • LowPrecisionLinear 与 DeepSeek Experts adapter;
  • forward/dgrad/wgrad autograd 框架;
  • FP32 master、FP32 AdamW 状态和 optimizer-step 后权重刷新;
  • FSDP2 通用多存储 Tensor 扩展协议;
  • DCP 模型与优化器 checkpoint;
  • 不支持场景快速失败且禁止 BF16/FP32 GEMM fallback。

实现分叉只收敛在格式层:

LowPrecisionRecipe
├── MXFP8Quantizer → MXFP8Tensor → npu_mxfp8
└── HiF8Quantizer  → HiF8Tensor  → npu_hif8

HiF8 初始数值策略冻结为:

  • input/weight 的格式最大值为 15.0
  • gradient 的格式最大值为 224.0
  • scale = valid(amax) ? amax / format_max : 1.0
  • scale 使用 FP32;
  • Dense 使用 per-tensor scale;3D Grouped 权重使用 per-leading-group scale。

这些数值不作为普通模型配置开放;它们属于经过算子与精度验证的 HiF8 Recipe。后续若硬件
契约变化,通过新增 recipe 版本演进,避免同一个 recipe 名称在不同环境中产生不同语义。

4. 配置入口

三类配置各自只有一个职责:

配置 回答的问题 产物
HF model config 构建什么模型、是否加载预训练权重 原始 DeepSeek-V3 模块树和高精 Parameter
LowPrecisionConfig 使用什么格式和缩放算法 只读 LowPrecisionPolicy
plan_overrides 哪些模块采用低精代码实现 ModuleReplacementPlan

4.1 HF 模型配置

新增 from_hf_config() 入口,只加载 HF 配置并随机初始化模型,不加载 DeepSeek-V3 671B
checkpoint:

HyperAutoModelForCausalLM.from_hf_config(
    config_name_or_path,
    *,
    config_overrides=None,
    torch_dtype=torch.bfloat16,
)

冻结的 16B 规格示例:

model:
  _target_: hyper_models._transformers.HyperAutoModelForCausalLM.from_hf_config
  config_name_or_path: <hf-deepseek-v3-config-path>
  config_overrides:
    vocab_size: 102400
    hidden_size: 2048
    intermediate_size: 10944
    moe_intermediate_size: 1408
    num_hidden_layers: 28
    first_k_dense_replace: 1
    num_attention_heads: 16
    num_key_value_heads: 16
    n_routed_experts: 64
    n_shared_experts: 2
    q_lora_rank: null
    kv_lora_rank: 3072
    qk_nope_head_dim: 64
    qk_rope_head_dim: 64
    v_head_dim: 128
    num_experts_per_tok: 6
    n_group: 1
    topk_group: 1
    norm_topk_prob: true
    routed_scaling_factor: 2.5
    max_position_embeddings: 4096
    use_cache: false
  torch_dtype: bfloat16
  attn_implementation: sdpa
  force_hf: true

模型构建后必须校验层数、Dense/MoE 排布、expert 数、投影 shape 和总参数量,不能让来源
HF 配置隐式改变 16B 规格。

4.2 低精策略配置

LowPrecisionConfig 只回答“怎么算”,不包含模块 FQN:

@dataclass(frozen=True)
class LowPrecisionConfig:
    enabled: bool = False
    format: Literal["mxfp8_e4m3", "hif8"] = "mxfp8_e4m3"
    scaling: Literal["mx_block", "current"] = "mx_block"

合法组合:

format scaling
mxfp8_e4m3 mx_block
hif8 current

格式与 scaling 组合不合法时,配置解析阶段直接报错。HiF8 current scaling 固定使用当前
amax,不开放安全余量等调节参数。解析后构建只读 LowPrecisionPolicy,再由模型加工上下文
传给 replacement factory。低精模块不直接读取 Trainer YAML。

YAML 示例:

# 二选一:MXFP8
low_precision:
  enabled: true
  format: mxfp8_e4m3
  scaling: mx_block
# 二选一:HiF8 current scaling
low_precision:
  enabled: true
  format: hif8
  scaling: current

4.3 目标模块与 FSDP 配置

模块目标不随格式变化,因此 replacement factory 使用格式无关名称:

fsdp_config:
  dp_shard_size: 8
  reshard_after_forward: true
  comm_fusion: false
  fallback_to_unsharded: false

plan_overrides:
  - match:
      - "model.layers.*.self_attn.q_proj"
      - "model.layers.*.self_attn.kv_a_proj_with_mqa"
      - "model.layers.*.self_attn.kv_b_proj"
      - "model.layers.*.self_attn.o_proj"
    when: low_precision
    module_type: torch.nn.Linear
    exact_type: true
    replace_module:
      _target_: hyper_models.components.training.low_precision.modules.linear.replace_low_precision_linear

  - match:
      - "model.layers.0.mlp.*_proj"
      - "model.layers.*.mlp.shared_experts.*_proj"
    when: low_precision
    module_type: torch.nn.Linear
    exact_type: true
    replace_module:
      _target_: hyper_models.components.training.low_precision.modules.linear.replace_low_precision_linear

  - match: "model.layers.*.mlp.experts"
    when: low_precision
    module_type: transformers.models.deepseek_v3.modeling_deepseek_v3.DeepseekV3Experts
    exact_type: true
    replace_module:
      _target_: hyper_models.components.models.adapters.deepseek_v3.low_precision.replace_deepseek_v3_low_precision_experts

配置约束:

  • plan_overrides 不重复配置 format/scaling;
  • MLA projection 使用模型版本对应的完整 FQN 列表;
  • q_lora_rank=null 只配置 q_proj,非空时改配 q_a_proj/q_b_proj
  • 已启用 replacement 规则零命中时直接报错;
  • 一个模块命中多个 replacement factory 时直接报错;
  • exact_type=true 使用 type(module) is configured_type
  • enabled=false 时不调用 factory 和 NPU capability 检查;
  • Target 必须 import-safe,模块导入阶段不得强制导入 torch_npu

构建阶段的交接关系:

HF config + overrides
    → meta device 构建原始 DeepSeek-V3 模块树
              │
plan_overrides ─┴→ 编译 ModuleReplacementPlan(只读取结构)
              │
              └→ 按 HF 语义物化并初始化高精 weight/bias

LowPrecisionConfig
    → LowPrecisionPolicy / Recipe ─────┐
ModuleReplacementPlan ─────────────────┤
                                       ↓
    → LowPrecisionLinear.from_linear(..., policy)
    → DeepseekV3LowPrecisionExperts(..., policy)
    → finalize 高精 weight 为对应格式 QuantizedWeightTensor
    → FSDP2 fully_shard

5. 模块替换

低精替换复用 Model Plan/Apply,不保留独立 Converter:

完整 model.named_modules() 扫描
    → match、module_type、exact_type 校验
    → 按模块 identity 合并 alias FQN
    → 生成不可变 Replacement Plan
    → 原模型完成高精参数初始化
    → factory 构造并校验全部 replacement
    → 原子安装 replacement
    → 在稳定的新模块树上执行 Sharding Plan 和 FSDP2

replacement 必须保持模块 FQN、Parameter/Buffer 名称、共享关系、requires_grad、checkpoint key
和 forward 签名。它允许改变具体模块类型并增加 policy、recipe 和 quantizer 等普通 Python
属性。

5.1 Dense 与 shared experts

只将 YAML 显式选中的 exact nn.Linear 替换为 LowPrecisionLinear

class LowPrecisionLinear(nn.Linear):
    policy: LowPrecisionPolicy
    recipe: LowPrecisionRecipe
    weight: QuantizedWeightTensor

    def forward(self, inputs): ...
    def refresh_weight_storage(self, master_shard): ...

from_linear() 复用原始 weight/bias,不重新初始化或复制 Parameter。高精权重物化后,
finalize 根据 policy 转换成 MXFP8TensorHiF8Tensor。bias 始终保持 BF16/FP32。

5.2 Routed experts

HF DeepseekV3Experts 已持有堆叠的 3D 参数:

gate_up_proj: [E, 2F, H]
down_proj:    [E, H, F]

因此整体替换为 DeepseekV3LowPrecisionExperts,不递归查找内部 Linear,也不再次堆叠
expert。该模块只适配 DeepSeek 参数名、forward 签名和 top-k 聚合;MXFP8/HiF8 Grouped
SwiGLU 的计算与 autograd 放在通用 functional/grouped_linear.py

shared experts 不是 routed expert container,其内部 exact Linear 继续走 Dense replacement。

6. 内部目录

hyper_models/
├── trainer/config.py
├── _transformers/infrastructure.py
└── components/
    ├── distributed/
    │   ├── module_replacement.py
    │   └── injection.py
    ├── models/adapters/deepseek_v3/
    │   └── low_precision.py
    └── training/low_precision/
        ├── config.py
        ├── policy.py
        ├── recipe.py
        ├── tensor/
        │   ├── base.py
        │   ├── operand.py
        │   ├── mxfp8_tensor.py
        │   └── hif8_tensor.py
        ├── quantizers/
        │   ├── base.py
        │   ├── mxfp8.py
        │   └── hif8.py
        ├── ops/
        │   ├── npu_mxfp8.py
        │   └── npu_hif8.py
        ├── functional/
        │   ├── linear.py
        │   └── grouped_linear.py
        ├── modules/
        │   └── linear.py
        └── integration/
            ├── fsdp_adapter.py
            ├── optimizer_adapter.py
            └── checkpoint_adapter.py
文件 内容
distributed/module_replacement.py 通用匹配、校验和原子安装,不导入低精类型
modules/linear.py 格式无关 LowPrecisionLinear 与 Dense factory
models/adapters/deepseek_v3/low_precision.py DeepseekV3LowPrecisionExperts 与模型 factory
recipe.py 建立计算角色到 Quantizer 的不可变映射
functional/linear.py Dense forward/dgrad/wgrad autograd
functional/grouped_linear.py Grouped SwiGLU forward/dgrad/wgrad autograd
quantizers/mxfp8.py MX block quantization 和 dual-axis 权重量化
quantizers/hif8.py HiF8 current-scaling quantization
ops/npu_mxfp8.py MXFP8 quant/MM/GMM 薄封装
ops/npu_hif8.py HiF8 quant/MM/GMM 薄封装
integration/*_adapter.py 连接低精 Tensor 与外部生命周期组件

首期不增加通用 modules/grouped_linear.py:PyTorch/HF 没有统一的 Grouped Experts Module
接口。通用 GMM 放在 functional,DeepSeek 的有状态模块放在模型 adapter。

7. Recipe、角色与数据结构

7.1 计算角色

LowPrecisionRecipe 按角色创建 Quantizer,而不是由 Linear 在运行时判断格式:

class QuantizationRole(str, Enum):
    INPUT_FWD = "input_fwd"
    WEIGHT_FWD = "weight_fwd"
    GRAD_OUTPUT = "grad_output"
    INPUT_BWD = "input_bwd"
    WEIGHT_BWD = "weight_bwd"

MXFP8 可让多个角色共享同一个无状态 Quantizer。HiF8 current scaling 中 input/weight 的
format max 相同,可以共享不可变配置;gradient 使用不同 format max,使用独立 Quantizer。

“角色”和“方向”是两个概念:角色决定 dtype、format max 和 scaling 状态;row/col 决定当前
MM/GMM 的物理布局。

7.2 MXFP8Tensor

MXFP8Tensor 是跨 micro-batch 持续存在的权重 Owner:

逻辑 shape / dtype / requires_grad
├── row_data
├── row_scale   E8M0
├── col_data
└── col_scale   E8M0

四份物理存储由 FP32 master 派生,不进入 checkpoint。

7.3 HiF8Tensor

HiF8 per-tensor/per-group scale 与 row/col block 无关,因此不复制两份相同 data:

逻辑 shape / dtype / requires_grad
├── data        hifloat8
└── scale       FP32 scalar 或 [E]

Dense 权重使用标量 scale;3D routed-expert 权重使用 [E] scale。select(layout) 通过视图和
transpose 元数据构造 operand,不重新量化或复制 data。

7.4 QuantizedOperand

QuantizedOperand 表示一次 MM/GMM 使用的短生命周期输入:

@dataclass(frozen=True)
class QuantizedOperand:
    data: torch.Tensor
    scale: torch.Tensor
    logical_dtype: torch.dtype
    format: LowPrecisionFormat
    orientation: Literal["row", "col"]

activation 量化结果和权重 select() 均返回 operand。FSDP 只处理持久化
QuantizedWeightTensor,不处理 activation operand。

8. 计算数据流

module.forward
    → functional autograd
        ├── 根据 QuantizationRole 取得 Quantizer
        ├── activation/gradient → quantize → QuantizedOperand
        ├── QuantizedWeightTensor.select(row/col) → QuantizedOperand
        └── format ops → BF16 output
计算 activation/grad weight 使用角色
forward row row(Dense)/ col(Grouped) INPUT_FWDWEIGHT_FWD
dgrad row col(Dense)/ row(Grouped) GRAD_OUTPUTWEIGHT_BWD
wgrad row/col,按算子契约选择 不读取旧 weight GRAD_OUTPUTINPUT_BWD

方向表是 GEMM/GMM 契约,对 MXFP8 决定选择哪份 physical data;对 HiF8 决定 data 的逻辑视图
和算子 transpose 参数。

8.1 MXFP8

activation/gradient 每次动态执行 MX block 量化。权重在 optimizer step 后从 FP32 master
一次性生成 row/col data/scale,在后续多个 micro-batch 和 activation recompute 中复用。

8.2 HiF8 current scaling

每次 activation/gradient 量化都从当前 Tensor 计算 amax 和 FP32 scale。持久化权重只在
optimizer step 后重新计算 scale 和 data,不在每个 forward 重复量化。

9. FSDP2 多存储扩展

HyperParallel core 只提供格式无关协议:

@dataclass(frozen=True)
class FSDPGatherContext:
    phase: Literal["forward", "backward"]
    reshard_after_forward: bool
    param_fqn: str


class FSDPLocalTensorExtension(Protocol):
    def fsdp_pre_all_gather(self, context):
        """返回本阶段需要通信的 physical tensors 和 metadata。"""

    def fsdp_post_all_gather(self, outputs, metadata, out=None):
        """由通信输出重建完整逻辑 Tensor。"""

FSDP 不理解 row/col、format 或 scale,只分别 all-gather pre 返回的物理 Tensor,并把
metadata 原样传回 post

9.1 MXFP8 gather

阶段 reshard_after_forward Dense Grouped
forward true row col
forward false row + col row + col
backward 任意 row + col row + col

reshard_after_forward=false 时,首次 all-gather 直接准备两个方向。backward 同时覆盖 dgrad
和 wgrad 所需数据,也准备两个方向。

MXFP8 分片合法性:

  • 2D [M, N]:本地 M_local 满足 scale tile 边界,当前要求 M_local % 64 == 0
  • 3D [E_local, D1, D2]:前两维为展平行维,要求 (E_local * D1) % 64 == 0
  • 不使用 padding 改变逻辑或物理 shape;
  • 不满足约束且 fallback_to_unsharded=false 时直接报错,为 true 时保持复制。

9.2 HiF8 gather

HiF8 权重只有一份 data/scale,不按 forward/backward 选择不同方向:

权重 本地状态 all-gather 输出
2D Dense data_shard + replicated scalar scale 完整 data + 经一致性校验的 scalar scale
3D Grouped data[E_local,...] + scale[E_local] 完整 data + scale[E]

2D Dense 的每个 FSDP rank 必须使用相同 scale。刷新权重前先对本地 master-shard amax 在
FSDP group 做 MAX reduce,再以同一 scale 分别量化各 data shard;这样 all-gather 后与使用
同一 scale 量化完整权重数学等价。post 校验各 rank gather 到的 scalar scale 一致后折叠为
一个标量。

3D Grouped 参数只允许沿 expert 维切分,并保持单个 expert 完整;scale 随 expert shard 一起
all-gather。若 FSDP 布局会切开单个 expert,则直接报错或由 fallback_to_unsharded 保持复制。

9.3 常驻与临时状态

本地低精 shard 在整个梯度累积周期内常驻;all-gather 得到的完整物理 Tensor 只在当前
unshard 窗口存在,由 FSDP 按 reshard 策略释放。二者不能互相释放。

10. Optimizer 生命周期

训练状态:

FP32 master shard        optimizer 更新的权威权重
FP32 exp_avg/exp_avg_sq  optimizer 状态
低精 weight shard        由 master 派生的计算存储

首次构建时不能从低精权重反量化得到 FP32 master。高精参数 finalize 时临时保留初始化值;
FSDP 完成本地分片后,用对应高精 shard 创建 FP32 master,随后释放初始化临时值。

统一 step 数据流:

梯度累积完成
    → FSDP reduce-scatter 得到本地 BF16/FP32 gradient
    → optimizer 使用 FP32 master 和 FP32 状态更新
    → 从新 master 刷新格式专属低精权重
         MXFP8: dual-axis data + E8M0 scales
         HiF8 current: current amax + FP32 scale + data
    → 下一 optimizer step 内复用

刷新只发生在 optimizer.step() 成功后,不公开“已失效但尚未重建”状态。zero_grad()
activation recompute 不刷新权重。若 optimizer step 被 overflow/异常跳过,则不刷新权重。

11. Checkpoint

仅支持 DCP:

  • 模型 checkpoint 保存 BF16 逻辑权重,不保存 wrapper 或派生低精 data/scale;
  • optimizer checkpoint 保存 FP32 master 和 FP32 exp_avg/exp_avg_sq
  • MXFP8 和 HiF8 current scaling 的运行时权重均可由 master 重建,不保存额外状态;
  • 恢复时先恢复模型与 optimizer 高精状态,再生成本地低精权重;
  • checkpoint key 与原 HF 模型保持一致;
  • fallback_to_unsharded 参数按普通复制参数处理。

12. 构建与运行顺序

12.1 构建

1. 读取 HF config,覆盖并校验冻结的 16B 规格
2. 在 meta device 构建原始 HF 模型
3. 编译 ModuleReplacementPlan
4. 按原 HF 语义物化并初始化高精参数
5. 应用 Dense/Experts replacement
6. 根据 policy finalize MXFP8Tensor 或 HiF8Tensor
7. 执行 Sharding Plan 和 FSDP2 fully_shard
8. 创建 FP32 master 与 FP32 optimizer state

12.2 单个训练 step

forward/recompute
    → 按角色量化 activation
    → unshard 低精 weight
    → MXFP8/HiF8 MM/GMM
backward
    → 按角色量化 grad_output/input
    → MXFP8/HiF8 dgrad/wgrad
    → FSDP reduce-scatter gradient
optimizer.step
    → 更新 FP32 master/state
    → 刷新低精 weight shard
checkpoint(如到保存步)
    → 保存逻辑模型和 optimizer

任何读取真实数据、调用 NPU kernel 或创建通信组的动作,都不得发生在 meta 构建和
replacement 编译阶段。

13. 对外接口设计

本节冻结 Q3 转测依赖的接口。以下是语义契约,最终 Python 包路径可随目录实现调整,但调用方
不应依赖具体格式类的私有字段。

13.1 Trainer 配置接口

@dataclass(frozen=True)
class LowPrecisionConfig:
    enabled: bool = False
    format: Literal["mxfp8_e4m3", "hif8"] = "mxfp8_e4m3"
    scaling: Literal["mx_block", "current"] = "mx_block"

输入来自 Trainer YAML;输出为不可变 LowPrecisionPolicy。配置层只允许两个组合:

mxfp8_e4m3 + mx_block
hif8       + current

其他组合在模型构建前报错。enabled=false 时不得导入 torch_npu、替换模块或创建低精状态。

13.2 模块替换接口

def replace_low_precision_linear(
    source: nn.Linear,
    *,
    fqn: str,
    policy: LowPrecisionPolicy,
) -> LowPrecisionLinear: ...


def replace_deepseek_v3_low_precision_experts(
    source: nn.Module,
    *,
    fqn: str,
    policy: LowPrecisionPolicy,
) -> DeepseekV3LowPrecisionExperts: ...

调用方是 ModuleReplacementPlan。factory 必须复用原 Parameter/Buffer,不得改变 FQN、共享
关系、requires_grad、state_dict key 和 forward 输入输出契约;失败时 Plan 不得部分修改模型。

13.3 Recipe 与 Quantizer 接口

class LowPrecisionRecipe(Protocol):
    def quantizer(self, role: QuantizationRole) -> Quantizer: ...


class Quantizer(Protocol):
    def quantize(
        self,
        tensor: torch.Tensor,
        *,
        orientation: Literal["row", "col"],
        group_list: torch.Tensor | None = None,
    ) -> QuantizedOperand: ...

QuantizationRole 决定 input/weight/gradient 的格式参数,orientation 决定 GEMM/GMM layout。
Quantizer 不遍历模型、不执行 collective,也不负责 optimizer 或 checkpoint。

13.4 低精 Tensor 与算子接口

class QuantizedWeightTensor(torch.Tensor):
    def select(self, orientation: Literal["row", "col"]) -> QuantizedOperand: ...
    def refresh_from_master(self, master_shard: torch.Tensor) -> None: ...


def low_precision_matmul(
    left: QuantizedOperand,
    right: QuantizedOperand,
    *,
    layout: Literal["NN", "NT", "TN"],
    output_dtype: torch.dtype,
) -> torch.Tensor: ...


def low_precision_grouped_matmul(
    left: QuantizedOperand,
    right: QuantizedOperand,
    group_list: torch.Tensor,
    *,
    layout: Literal["NN", "NT", "TN"],
    output_dtype: torch.dtype,
) -> torch.Tensor: ...

公开计算接口使用 QuantizedOperand,不让上层直接拼装 data/scale。格式专属 NPU adapter
负责 dtype、scale dtype、transpose 和算子参数校验,任何失败都不得退回高精 GEMM。

13.5 FSDP2 扩展接口

FSDP2 只调用第 9 章的 fsdp_pre_all_gather(context)
fsdp_post_all_gather(outputs, metadata, out)。扩展返回物理 Tensor 的有序 tuple 和可序列化
metadata;post 必须校验输出数量、shape、dtype 和顺序。普通 Tensor 不进入该接口。

13.6 Optimizer 与 DCP 接口

Optimizer adapter 只在成功更新 FP32 master 后调用:

def refresh_low_precision_weights(
    params: Iterable[QuantizedWeightTensor],
    master_params: Iterable[torch.Tensor],
) -> None: ...

DCP adapter 对外提供高精 checkpoint 视图:模型为 BF16 逻辑权重,optimizer 为 FP32 master、
FP32 moments 和 step。MXFP8/HiF8 data/scale 不属于持久化接口。

14. 与其他模块的相关性

相关模块 本功能使用的能力 本次需要的适配 Owner 与边界
HF/AutoModel DeepSeek-V3 config、meta 构建、参数初始化 from_hf_config、16B override 和结构校验 模型 Owner;低精不实现模型算法
Model Plan/Apply YAML 匹配、严格类型筛选、有序加工 module replacement action、冲突和零命中检查 加工框架 Owner;低精不再自建 Converter
DeepSeek adapter routed experts 参数布局、router/top-k 语义 DeepseekV3LowPrecisionExperts 模型适配 Owner;通用低精不导入 DeepSeek
FSDP2 Parameter 分片、all-gather、reduce-scatter、reshard 格式无关多存储 Tensor pre/post 协议 FSDP Owner;格式层决定物理 Tensor
Optimizer AdamW 更新、梯度累积、overflow/skip 语义 FP32 master/state 和成功 step 后刷新回调 Optimizer Owner;低精不复制 AdamW
Activation recompute 重放 forward 和 RNG 状态 确认不额外刷新权重,通信 phase 正确 Recompute Owner;低精无专用开关
DCP 分片模型与 optimizer 保存恢复 高精视图、FP32 master 恢复后重新量化 DCP Owner;不持久化派生 data/scale
torch_npu/CANN quant、MM、GMM MXFP8/HiF8 薄封装和 capability gate 算子 Owner;HP 不实现 kernel
Precision Observer 量化误差观测 本期不接入 不属于 Q3 转测范围

必须保持的加工顺序:

HF/meta 构建
    → 高精参数物化
    → Module Replacement Plan
    → 低精权重 finalize
    → Sharding Plan / FSDP2
    → activation checkpoint wrapper
    → optimizer 与 DCP 注册

顺序改变可能导致 meta 阶段调用 NPU、FSDP 看不到低精 Tensor、hook 重复注册或 optimizer
拿不到高精 master,属于启动时必须检查的集成错误。

15. 约束限制与失败规则

启动时按 policy 校验,不通过直接失败:

  • A5 设备和 torch_npu 版本;
  • MXFP8:E4M3、E8M0、dynamic MX quant、dual-axis quant、quant MM/GMM;
  • HiF8:torch_npu.hifloat8、current quant、指定 scale quant、FP32 scale MM/GMM;
  • 格式与 scaling 组合以及附加参数;
  • replacement 命中数、严格类型、规则冲突和 Parameter schema;
  • TP/CP/EP/PP > 1
  • FSDP 通信融合或非 DCP checkpoint;
  • 分片不满足格式约束且未开启 fallback_to_unsharded

运行阶段不得因 capability、shape 或算子不支持退回 BF16/FP32 GEMM。

完整限制矩阵:

维度 Q3 支持值 不支持时行为
硬件/后端 Ascend A5,PyTorch + torch_npu 启动失败
模型 冻结的 DeepSeek-V3 16B 结构 规格校验失败
低精格式 MXFP8 E4M3、HiF8 current 配置解析失败
低精目标 显式 Dense/shared/MLA Linear 和 routed experts 零命中或类型不符失败
并行 FSDP2,TP/CP/EP/PP 均为 1 启动失败
FSDP comm fusion false 启动失败
分片不合法 fallback_to_unsharded 控制复制或失败 不允许改用高精 GEMM
Optimizer AdamW,FP32 master/moments/update 其他状态 dtype 失败
Checkpoint HP DCP 其他路径失败
Recompute off、full 其他策略不作为 Q3 承诺
PEFT/自定义模块 不支持 命中后失败,不盲目替换
编译/融合 torch.compile、低精融合算子不承诺 不进入 Q3 转测

16. UT 与组件测试设计

UT 默认在 CPU/mock NPU Ops 上执行;验证 Python 契约、数据流和失败语义。真实 kernel 数值只在
A5 ST/转测验证,UT 不伪装成硬件验收。

16.1 配置与 Replacement Plan UT

用例 输入/操作 检查点
合法格式组合 两种合法 format/scaling 生成对应不可变 policy
非法格式组合 MXFP8+current、HiF8+mx_block、HiF8+delayed 构建模型前报错
关闭低精 enabled=false 不导入 torch_npu、不替换、不调用 capability gate
FQN 精确选择 Dense/shared/MLA 显式 glob 目标全部命中,非目标不命中
严格类型 exact Linear、Linear 子类、PEFT/custom wrapper 只接受 exact Linear;其余明确失败
零命中/重复命中 错误 FQN、两个 factory 命中同一模块 原子失败,模型未部分修改
alias Parameter 两个 FQN 指向同一模块/权重 replacement identity 和共享关系保留
replacement schema 替换前后比较 named parameters/buffers/modules FQN、key、shape、requires_grad、training 状态一致

16.2 Tensor、Quantizer 与 Ops UT

用例 MXFP8 检查点 HiF8 current 检查点
量化构造 row/col data-scale 成对且 E8M0 data 为 hifloat8、scale 为 FP32
role Fprop/dgrad/wgrad 方向正确 input/weight max=15,gradient max=224
Dense scale block scale shape 正确 amax/format_max,零/NaN/Inf 安全处理
Grouped scale group_list 与 GMM 参数正确 scale [E] 与 expert 顺序一致
operand select(row/col) 返回正确 data/scale row/col 复用 storage,layout 正确
MM/GMM adapter dtype、scale dtype、transpose、layout 原样下发 HiF8 dtype、FP32 scale 和 layout 原样下发
非法输入 shape、pair、layout、dtype 错误时报错 shape、scale、group_list、dtype 错误时报错
无 fallback mock kernel 抛错 异常原样归因,不调用高精 linear/mm

Quantizer 数值 UT 使用确定性小 Tensor 手算 scale 和反量化误差;MM/GMM Python UT 使用 fake
operator 记录调用参数,不把 fake 输出作为真实低精精度结论。

16.3 Dense/Grouped Autograd UT

  • 2D/3D activation 的 output、dinput、dweight shape/dtype;
  • Dense NN/NT/TN 三条 Fprop/Dgrad/Wgrad 路径和量化角色;
  • Grouped gate/up、SwiGLU、down 的 Fprop/Dgrad/Wgrad 调用顺序;
  • 空 expert、单 expert、token 数不均和多个 top-k 分支;
  • bias 保持高精且梯度正确归属;
  • 多 micro-batch 只累计梯度,不刷新权重;成功 optimizer step 后每个权重只刷新一次;
  • activation recompute 不创建第二份持久权重或触发 optimizer 刷新。

16.4 FSDP2 Extension UT

  • 普通 Tensor 不进入扩展路径,行为保持不变;
  • fake 多存储 Tensor 的协议识别、缺失方法校验;
  • pre 返回不同 dtype/shape 的多个物理 Tensor,post 按 metadata 重建;
  • physical tensor 数量、顺序、shape、dtype 不一致时报错;
  • forward/backward、reshard_after_forward=true/false 和 prefetch phase;
  • MXFP8 forward 单/双方向和 backward 双方向;
  • HiF8 Dense data+复制 scale、Grouped data+[E] scale;
  • fallback_to_unsharded 的复制/报错分支;
  • extension 与 comm_fusion=true 组合快速失败。

16.5 Optimizer、Recompute 与 DCP UT/ST

  • optimizer 以 FP32 master 创建,exp_avg/exp_avg_sq 均为 FP32;
  • master 更新成功后刷新低精 shard;step 被跳过或失败时不刷新;
  • HiF8 Dense 刷新前的 amax MAX reduce 和跨 rank scale 一致性;
  • checkpoint 模型视图不包含 wrapper/data/scale;
  • DCP optimizer 状态包含 FP32 master/moments/step;
  • 保存、恢复后由 master 重建低精 shard,参数 FQN 和逻辑 shape 不变;
  • activation_checkpoint=off/full 下 loss、梯度和 kernel 路径一致,full 不重复刷新权重。

16.6 A5 端到端 ST 伪代码

ST 必须经过生产 Trainer 入口。STProbe 只观测模块类型、kernel、FSDP physical tensors、scale
和 refresh 次数,不在测试中手工执行 replacement、量化、all-gather 或权重刷新。

# 伪代码:接口名称表达测试职责,不要求与最终实现逐字一致。

@dataclass(frozen=True)
class STCase:
    name: str
    low_precision: dict | None       # None 表示 BF16 baseline
    reshard_after_forward: bool
    activation_checkpoint: str      # "off" / "full"
    grad_accum_steps: int


CASES = [
    STCase("bf16", None, True, "off", 1),
    STCase(
        "mxfp8",
        {"enabled": True, "format": "mxfp8_e4m3", "scaling": "mx_block"},
        True,
        "off",
        1,
    ),
    STCase(
        "hif8_current",
        {"enabled": True, "format": "hif8", "scaling": "current"},
        True,
        "off",
        1,
    ),
]


def run_st(case: STCase, checkpoint_dir: Path) -> STResult:
    init_process_group(backend="hccl")
    set_deterministic_seed(2026)  # 所有 case 使用相同 seed 和 batch 顺序

    cfg = load_yaml("dsv3_16b_fsdp2_st.yaml")
    cfg.low_precision = case.low_precision
    cfg.fsdp_config.reshard_after_forward = case.reshard_after_forward
    cfg.gradient_checkpointing.activation_checkpoint = case.activation_checkpoint
    cfg.gradient_accumulation_steps = case.grad_accum_steps

    # setup() 必须走正式顺序:HF 构建 → replacement → finalize → FSDP2 → optimizer。
    trainer = Trainer()
    probe = STProbe(read_only=True)
    trainer.register_test_probe(probe)
    trainer.setup(cfg)

    assert_model_scope(trainer.model, case, probe)
    assert_optimizer_fp32_state(trainer.optimizer)
    if case.low_precision is not None:
        assert_no_high_precision_gemm_fallback(probe)

    step_records = []
    for step in range(WARMUP_STEPS + MEASURE_STEPS):
        trainer.optimizer.zero_grad(set_to_none=True)
        probe.begin_optimizer_step(step)

        for micro_step in range(case.grad_accum_steps):
            batch = deterministic_batch(step, micro_step)
            loss = trainer.forward_backward(
                batch,
                loss_scale=1.0 / case.grad_accum_steps,
            )
            assert torch.isfinite(loss)

        assert_all_gradients_finite(trainer.model)

        # 由生产 optimizer adapter 在成功 step 后刷新;ST 只记录事件。
        updated = trainer.optimizer_step()
        probe.end_optimizer_step(updated=updated)

        if updated and case.low_precision is not None:
            probe.assert_one_weight_refresh_per_parameter(step)
        if not updated:
            probe.assert_no_weight_refresh(step)

        torch.npu.synchronize()
        step_records.append(
            collect_step_record(
                step=step,
                loss=loss,
                grad_norm=trainer.grad_norm,
                elapsed=probe.step_elapsed,
                peak_memory=torch.npu.max_memory_allocated(),
            )
        )

        if step == SAVE_STEP:
            trainer.save_dcp(checkpoint_dir)

    if case.low_precision is not None:
        assert_low_precision_kernel_contract(case, probe)
        assert_fsdp_physical_tensor_contract(case, probe)
        assert_scale_contract(case, probe)

    return STResult(
        case=case,
        records=step_records,
        checkpoint_dir=checkpoint_dir,
        trace=probe.export_trace(),
    )

关键断言伪代码:

def assert_low_precision_kernel_contract(case, probe):
    # Dense 与 routed experts 的三种反向角色必须全部真实执行。
    expected = {
        ("dense", "fprop"), ("dense", "dgrad"), ("dense", "wgrad"),
        ("grouped", "fprop"), ("grouped", "dgrad"), ("grouped", "wgrad"),
    }
    assert expected <= probe.executed_kernel_roles()
    assert probe.high_precision_fallback_count == 0

    if case.name == "mxfp8":
        probe.assert_all_kernel_format("mxfp8_e4m3")
        probe.assert_mxfp8_scale_dtype("e8m0")
    elif case.name == "hif8_current":
        probe.assert_all_kernel_format("hif8")
        probe.assert_hif8_scale_dtype(torch.float32)
        probe.assert_hif8_format_max(input_weight=15.0, gradient=224.0)


def assert_fsdp_physical_tensor_contract(case, probe):
    for event in probe.fsdp_gather_events:
        event.assert_post_matches_pre_schema()

    if case.name == "mxfp8":
        probe.assert_mxfp8_forward_direction_by_compute_kind()
        probe.assert_mxfp8_backward_has_row_and_col()
    elif case.name == "hif8_current":
        probe.assert_hif8_forward_backward_use_data_and_scale()
        probe.assert_hif8_dense_scale_equal_across_ranks()
        probe.assert_hif8_grouped_scale_matches_expert_order()


def assert_scale_contract(case, probe):
    if case.name != "hif8_current":
        return
    for event in probe.hif8_quant_events:
        expected_max = 224.0 if event.role == "grad_output" else 15.0
        expected_scale = safe_scale(event.current_amax, expected_max)
        torch.testing.assert_close(event.scale, expected_scale)
        assert event.scale.dtype == torch.float32
        assert torch.isfinite(event.scale).all() and (event.scale > 0).all()

DCP 续训使用新的 Trainer 实例,防止旧进程内对象掩盖恢复问题:

def run_resume_st(case, continuous_result):
    resumed = Trainer()
    resumed.setup(load_same_case_config(case))
    resumed.load_dcp(continuous_result.checkpoint_dir)

    assert resumed.global_step == SAVE_STEP
    assert_fp32_master_and_moments(resumed.optimizer)
    assert_checkpoint_has_no_derived_low_precision_storage(
        continuous_result.checkpoint_dir
    )

    # load_dcp 后由生产 adapter 从 FP32 master 重建低精权重。
    assert_rebuilt_weight_storage(resumed.model, case)

    for step in range(SAVE_STEP, SAVE_STEP + RESUME_STEPS):
        batch = deterministic_batch(step, 0)
        resumed_record = resumed.train_step(batch)
        continuous_record = continuous_result.record(step)
        assert_resume_metrics_close(resumed_record, continuous_record)

性能汇总伪代码:

bf16 = run_st(BF16_CASE, tmp_path / "bf16")
mxfp8 = run_st(MXFP8_CASE, tmp_path / "mxfp8")
hif8 = run_st(HIF8_CURRENT_CASE, tmp_path / "hif8")

for low_precision in (mxfp8, hif8):
    # 丢弃 warmup,只统计相同 token 数的连续 optimizer steps。
    assert low_precision.steady_tokens_per_second >= bf16.steady_tokens_per_second
    assert low_precision.average_step_time <= bf16.average_step_time

8 卡启动形式:

torchrun --nproc_per_node=8 tests/st/low_precision/test_dsv3_16b_fsdp2.py \
  --config configs/st/dsv3_16b_fsdp2.yaml \
  --case hif8_current

17. Q3 转测范围与验收标准

17.1 转测环境与固定输入

Q3 转测使用单机 8 卡 Ascend A5(每卡 96 GiB)、PyTorch + torch_npu,固定 DeepSeek-V3
16B 规格和同一份确定性训练数据。每次测试记录代码 commit、CANN/torch_npu/PyTorch 版本、
完整 YAML、随机种子和算子 trace。若 16B 因环境资源问题无法启动,只能用缩小模型定位,不能
以缩小模型代替最终验收。

固定配置维度:

format: BF16 baseline / MXFP8 / HiF8 current
reshard_after_forward: true / false
activation_checkpoint: off / full
fallback_to_unsharded: false(主路径)/ true(专项)
gradient_accumulation_steps: 1 / >1

17.2 阻塞验收矩阵

ID 场景 配置/步骤 必须采集的证据 通过标准
Q3-01 BF16 基线 8 卡 FSDP2,至少 20 optimizer steps loss、grad norm、NPU error、checkpoint 全程无异常;loss/grad norm 均 finite;可保存恢复并继续 5 steps
Q3-02 MXFP8 端到端 MXFP8,主配置,至少 20 steps 模块命中清单、MM/GMM trace、loss、grad norm 目标模块全部替换;Fprop/Dgrad/Wgrad 命中 MXFP8 MM/GMM;无高精 GEMM fallback;指标 finite
Q3-03 HiF8 端到端 HiF8 current,主配置,至少 20 steps 模块命中清单、MM/GMM trace、scale 统计、loss、grad norm 目标模块全部替换;三阶段命中 HiF8 MM/GMM;scale 为 FP32 且 finite/positive;无 fallback
Q3-04 模型边界 对两种低精格式检查模块树与 trace replacement 结果和高精算子 trace Dense/shared/routed/显式 MLA 命中;embedding/head/router/norm/核心 attention 未命中
Q3-05 MXFP8 FSDP reshard_after_forward=true/false 各运行 5 steps pre/post gather trace、物理 tensor schema true 时 forward 按表选择、backward 双方向;false 时首次 gather 双方向;重建 shape/dtype 正确
Q3-06 HiF8 FSDP reshard_after_forward=true/false 各运行 5 steps data/scale gather trace、rank scale Dense 各 rank scale 一致;Grouped scale 顺序与 expert 一致;前反向使用同一 data+scale
Q3-07 梯度累积 accumulation >1,各格式运行 5 optimizer steps micro-step/optimizer-step/refresh 计数 micro-step 不刷新;每个成功 optimizer step 每个低精权重只刷新一次
Q3-08 全量重计算 activation_checkpoint=full,各格式 5 steps kernel、gather、refresh 计数与 loss 可完成 F/B;无额外持久权重;不重复刷新;loss/grad finite
Q3-09 DCP 续训 连续训练 N+5;训练 N 后保存再恢复 5 model/optimizer key、dtype、step、恢复后 loss checkpoint 无派生 data/scale;master/moments FP32;恢复 step=N;首个恢复 loss 与连续路径在约定容差内
Q3-10 非法配置 逐项覆盖第 15 章 错误消息和退出阶段 均在约定阶段失败;错误包含配置项/FQN/shape/API;无 silent skip/fallback
Q3-11 性能不劣化 BF16/MXFP8/HiF8 使用完全相同的有效训练配置 稳态 tokens/s、平均 step time、原始逐 step 数据 MXFP8 和 HiF8 的稳态 tokens/s 均不低于 BF16,且平均 step time 均不高于 BF16

这里的“约定容差”必须在首轮 A5 BF16、MXFP8、HiF8 固定 seed 数据采集后,由算法负责人给出
具体数值并写回转测用例;在阈值冻结前,Q3-09 不得仅凭“曲线看起来接近”判定通过。

17.3 数值与精度专项

17.3.1 HiF8 专项验收范围

Q3 HiF8 专项验收仅覆盖 hif8 + current,不覆盖 delayed scaling。转测不需要构造、检查或
恢复 amax history、pending amax、scale update interval 等 delayed 状态。

检查项 测试方法 通过标准
Dense 计算 分别触发 Fprop、Dgrad、Wgrad 并采集 kernel trace 三条路径均命中 HiF8 quant MM,无 BF16/FP32 GEMM fallback
Routed experts 覆盖 gate/up、SwiGLU、down 的 Fprop、Dgrad、Wgrad 均命中 HiF8 GMM;空 expert 和 token 不均场景可执行
Current scale 对确定性 input、weight、grad_output 记录 amax 和 scale input/weight 使用 format max 15,gradient 使用 224;FP32 scale 与当前 amax 一致且 finite/positive
FSDP Dense 检查各 rank 本地 amax 归并和 all-gather 输出 同一 Dense 权重的各 rank 使用相同 scale,重建 data/shape 正确
FSDP Grouped 检查 [E_local,...] data、[E_local] scale 的 gather scale 顺序与全局 expert 顺序一致,不切开单个 expert
权重生命周期 梯度累积和 recompute 下记录 refresh 次数 micro-step 和 recompute 不刷新权重;每个成功 optimizer step 每个权重只刷新一次
DCP 续训 保存后恢复模型和 optimizer,再运行相同后续 batch checkpoint 不含 HiF8 data/scale;从 FP32 master 重建后可继续训练并满足恢复容差
性能 与相同有效训练配置的 BF16 比较 稳态吞吐不低于 BF16,平均 step time 不高于 BF16

HiF8 delayed scaling 配置必须在启动阶段报“不支持”,而不是被忽略、映射为 current 或回退到
高精计算。delayed scaling 的实现与验收统一放到第 18 章后续演进。

17.3.2 分层数值验收

端到端 loss 不能单独证明 Fprop/Dgrad/Wgrad 正确,必须分层验收:

  1. 单个 Dense 和 Grouped 算子分别与 BF16 golden 对比 output、dinput、dweight;报告 max
    abs、MSE、NRMSE、cosine 和 nonfinite count。
  2. 固定 1 层小模型,比较 BF16/MXFP8/HiF8 的单 step loss、grad norm 和关键参数梯度。
  3. 16B 固定数据短程训练至少 20 optimizer steps,报告三条 loss/grad-norm 曲线,不允许 NaN、
    Inf、持续爆炸或从第一个 step 起完全不更新。
  4. checkpoint 连续/恢复路径使用相同后续 batch,比较恢复后的第一个 loss、关键参数和
    optimizer step。

Q3 首轮允许“采集阈值”与功能转测并行,但发布前必须冻结上述指标的数值门槛。文档不预设
未经 A5 实测的统一误差阈值,避免用任意阈值掩盖格式和算子差异。

17.4 性能与显存专项

性能是 Q3 阻塞验收项。比较时必须固定模型、训练数据、sequence length、micro/global batch、
梯度累积、FSDP、reshard 和 recompute 配置;不得通过缩小 batch、减少有效 token 或改变并行
策略得到表面加速。关闭 profiling 和一次性初始化影响,预热后至少采集 20 个连续 optimizer
step,以相同统计方法计算:

throughput_ratio = low_precision_tokens_per_second / bf16_tokens_per_second
step_time_ratio  = low_precision_average_step_time / bf16_average_step_time

通过条件:
    throughput_ratio >= 1.0
    step_time_ratio  <= 1.0

MXFP8 和 HiF8 current 必须分别满足上述条件,不能用一种格式的收益抵消另一种格式的劣化。
同时必须产出:

  • BF16、MXFP8、HiF8 在相同 global batch/sequence length 下的稳态 tokens/s、step time;
  • 单卡和全局峰值显存;
  • quant、MM/GMM、all-gather、reduce-scatter 的耗时占比;
  • reshard_after_forward 和 recompute 开关的差异;
  • 低精关闭时相对 BF16 的额外依赖、常驻对象和性能回退。

任一低精格式吞吐低于 BF16 或平均 step time 高于 BF16,均判定 Q3 性能验收失败。峰值显存
本期要求记录并解释,暂不设置相对 BF16 的阻塞阈值;若发现错误 fallback、重复量化、重复
all-gather 或缓存未释放,则同时按功能问题阻塞。

17.5 转测交付物

  • 两套可直接运行的 MXFP8/HiF8 YAML 和一套 BF16 baseline YAML;
  • A5 启动命令、环境版本和依赖安装说明;
  • Q3-01~Q3-11 自动化脚本或明确的用例步骤;
  • kernel trace、模块 replacement 结果和 FSDP physical-tensor trace;
  • 数值、性能、显存和 checkpoint 恢复报告;
  • 已知限制、失败规则及未关闭问题清单。

18. 后续演进

HiF8 delayed scaling 后续单独设计和交付。届时需要新增角色级 amax history、跨 rank amax
归并、optimizer-step 更新时序和 DCP scaling-state 保存恢复;本期不提前引入这些状态、配置或
生命周期接口。

likedislike
missyuanzi
8月12日 评论:

需要补充说明Q3转测的范围、与其他模块相关性、UT设计,最好有个使用样例

likedislike
Zzhouyaqiang0
8月12日 修改了issue 的描述
Zzhouyaqiang0
8月12日 修改了issue 的描述