mindspore-ai / mindspore-ai/hyper-parallel

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

Open
#190 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
53
Forks
63
Avg merge
23h 45m
Merged PRs (30d)
63

Description

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 保存恢复;本期不提前引入这些状态、配置或
生命周期接口。

schema_version: 1
source: gitcode
gitcode_repo: mindspore/hyper-parallel
gitcode_issue: 329
source_url: https://gitcode.com/mindspore/hyper-parallel/issues/329

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

The issue names hyper_models/trainer/config.py, _transformers/infrastructure.py, and the components tree; begin by reading those entry points and tracing the existing model-plan, FSDP2, and low-precision extension protocols. Then verify the proposed DeepSeek-V3 replacement, MXFP8/HiF8 paths, and checkpoint lifecycle against current interfaces. Done means the design is implemented and validated against the stated supported and unsupported cases on Ascend A5.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
distributed-systems, machine-learning, performance
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.