mindspore-ai / mindspore-ai/hyper-parallel
DeepSeek-V3 16B FSDP2 MXFP8/HiF8 训练设计
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 + currentscaling;- 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
对应TELinear、TEColumnParallelLinear和TERowParallelLinear。MLA 实现还明确处理
linear_proj保存量化输入,
说明 projection Linear 位于 FP8 路径。 - Megatron 将核心 attention 的低精单独控制,
fp8_dot_product_attention和fp8_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 转换成 MXFP8Tensor 或 HiF8Tensor。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_FWD、WEIGHT_FWD |
| dgrad | row | col(Dense)/ row(Grouped) | GRAD_OUTPUT、WEIGHT_BWD |
| wgrad | row/col,按算子契约选择 | 不读取旧 weight | GRAD_OUTPUT、INPUT_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 正确,必须分层验收:
- 单个 Dense 和 Grouped 算子分别与 BF16 golden 对比 output、dinput、dweight;报告 max
abs、MSE、NRMSE、cosine 和 nonfinite count。 - 固定 1 层小模型,比较 BF16/MXFP8/HiF8 的单 step loss、grad norm 和关键参数梯度。
- 16B 固定数据短程训练至少 20 optimizer steps,报告三条 loss/grad-norm 曲线,不允许 NaN、
Inf、持续爆炸或从第一个 step 起完全不更新。 - 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
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- 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