HyperParallel FSDP 测试方案

版本：v1.1 | 范围：fully_shard API + HSDPModule 接口

---

Level 1：HSDPModule 接口单点测试

原则：覆盖 fully_shard() 函数入参 + HSDPModule 方法的覆盖 gap；语义相关接口合并到一个 case；2–4 卡；简单 MLP；目标是接口行为正确性，不做完整精度对标。

fully_shard() 入参 gap 分析

参数	当前覆盖状态	Gap
mesh=None	已有（test_fully_shard_none_mesh）	—
mesh=DeviceMesh	已有（多个用例）	—
reshard_after_forward=True/False	静态构造时已覆盖；动态 setter 未独立测试	L1-01 补充
shard_placement_fn	_test_hsdp_param 中有单元，但 fully_shard API 层未独立测试	L1-07 补充
mp_policy 全字段组合	严重不足：仅有 fp16/fp32/fp32 一种组合	L1-10~13 重点补充
offload_policy	有 pin_memory=True 一种	—
ignored_params	无独立单点测试	L1-06 补充
replicate_params	有精度测试（1D/2D），行为语义已基本覆盖	—
comm_fusion	有 ON/OFF 精度测试；无 ON vs OFF 数值一致性对比	L1-08 补充
module=List[Module]	有精度测试；HSDPModule 接口覆盖不完整	L1-09 补充

HSDPModule 方法 gap 分析

方法	当前覆盖状态	Gap
set_requires_gradient_sync	已有（test_fully_shard_module）	—
zero_grad	未独立测试	L1-04 补充
set_modules_to_forward_prefetch	已有	—
set_modules_to_backward_prefetch	已有	—
reshard() / unshard()	HSDPParam 层有；API 层幂等性未测	L1-05 补充
set_is_last_backward	已有	—
set_requires_all_reduce	未独立测试	L1-02 补充
set_reshard_after_forward（动态）	未独立测试	L1-01 补充
set_reshard_after_backward（动态）	未独立测试	L1-01 补充
set_reduce_op_type	未独立测试	L1-03 补充

---

L1 用例列表

#	用例名	目标文件	卡数	核心断言
L1-01	test_dynamic_reshard_control	_test_fully_shard_module.py	2	set_reshard_after_forward(False) → forward 后 param.data.numel()==full_numel；toggle 回 True → shard_numel；set_reshard_after_backward 同理；动态切换不抛异常
L1-02	test_set_requires_all_reduce	_test_fully_shard_module.py	4（2×2 HSDP）	关闭时两 replica-group grad 不等；开启后 grad atol=1e-6 一致；recurse=False 只影响顶层 module
L1-03	test_set_reduce_op_type	_test_fully_shard_module.py	2	"sum" → grad == ref_grad × world_size；"avg" → grad == ref_grad；mid-train 切换不污染 optimizer state
L1-04	test_zero_grad	_test_fully_shard_module.py	2	N 步累积后调用 zero_grad()：param.grad is None；param.data 与 snapshot 一致；后续 backward 产生新鲜 grad
L1-05	test_manual_reshard_unshard_lifecycle	_test_fully_shard_module.py	2	幂等性（重复调用不报错）；numel 正确（shard/full）；async_op=True 时 handle wait 后方可读取
L1-06	test_ignored_params	_test_fully_shard_precision.py	2	ignored 参数为普通 Tensor 非 DTensor；numel==full_numel；grad 不跨 rank 同步；其余参数正常 reduce-scatter
L1-07	test_custom_shard_placement_fn	_test_hsdp_param.py	2	dim-1 分片：param.shape[1]==full_shape[1]//world_size；loss + grad atol=1e-5 vs 单卡；Replicate() 返回时 shape == full_shape
L1-08	test_comm_fusion_numerical_correctness	_test_fully_shard_precision.py	2	abs(loss_fusion - loss_nofusion) < 1e-6；per-param grad atol=1e-6（comm_fusion 开/关数值等价性）
L1-09	test_fully_shard_list_input	_test_fully_shard_module.py	2	return 是 list，两元素均有 HSDPModule 接口；独立控制互不干扰；两者 grad 均有限

---

L1-10~13：混合精度策略充分校验（重点补强）

MixedPrecisionPolicy 有 5 个字段：param_dtype、reduce_dtype、output_dtype、cast_forward_inputs、apply_grad_on_fp32_main_grad。
当前仅覆盖了 fp16/fp32/fp32 一种组合，以下补强：

#	用例名	mp_policy 配置	卡数	核心断言
L1-10	test_mp_bf16_param_fp32_reduce	param_dtype=bf16, reduce_dtype=fp32, output_dtype=fp32	2	参数存储为 bf16；reduce 后 grad 为 fp32；loss + grad norm atol=1e-3 vs 单卡 bf16 参考
L1-11	test_mp_bf16_all_bf16	param_dtype=bf16, reduce_dtype=bf16, output_dtype=bf16	2	全程 bf16；output tensor dtype 为 bf16；grad dtype 为 bf16；loss atol=1e-3
L1-12	test_mp_cast_forward_inputs_false	param_dtype=fp16, reduce_dtype=fp32, cast_forward_inputs=False	2	输入不被 cast（forward hook 确认 input dtype 不变）；loss finite；无 dtype 不匹配报错
L1-13	test_mp_apply_grad_on_fp32_main_grad	param_dtype=fp16, reduce_dtype=fp32, apply_grad_on_fp32_main_grad=True	2	grad 累积到 fp32 main grad；param.grad dtype 为 fp32；精度 atol=1e-5 vs fp32 单卡参考

混合精度覆盖矩阵（补强后）：

param_dtype	reduce_dtype	cast_forward_inputs	apply_grad_on_fp32_main_grad	覆盖用例
None（fp32）	None	True	False	已有
fp16	fp32	True	False	已有
bf16	fp32	True	False	L1-10
bf16	bf16	True	False	L1-11
fp16	fp32	False	False	L1-12
fp16	fp32	True	True	L1-13

目标文件：L1-10~13 均归入 tests/torch/fully_shard/_test_fully_shard_precision.py

---

Level 2：特性耦合测试

原则：简单网络（≤6 层）；精度统一对标单卡/DDP；按特性命名归档文件。

L2 用例列表

#	用例名	目标文件	卡数	核心断言
L2-01	test_gradient_accumulation	_test_fully_shard_precision.py	2	N micro-step 累积后 merged grad == large-batch DDP grad atol=1e-5；中间步 rank 间 grad 不等（确认未同步）
L2-02	test_frozen_params_mixed_training	_test_fully_shard_precision.py	2	frozen 参数 grad is None；data 不变；非 frozen 层 grad atol=1e-5 vs 单卡；FSDP 不为 frozen 参数发起 reduce-scatter
L2-03	test_reentrant_ac_fsdp	_test_fully_shard_with_ac.py	2	loss atol=1e-5；per-layer grad norm atol=1e-5；首/末层 grad tensor atol=1e-5 vs 单卡
L2-04	test_nonreentrant_ac_fsdp	_test_fully_shard_with_ac.py	2	同 L2-03；额外：峰值内存低于无 AC baseline
L2-05	test_selective_ac_fsdp	_test_fully_shard_with_ac.py	2	同 L2-03；额外：checkpoint 层 forward-hook 触发 2次/backward，非 checkpoint 层 1次
L2-06	test_ac_activation_swap_fsdp	_test_fully_shard_with_ac.py	2	loss atol=1e-5；grad norm atol=1e-5；swap 回 GPU 后无 device mismatch；recompute 时激活在 GPU 上
L2-07	test_grad_clip_with_dtensor_grads	_test_fully_shard_precision.py	2	clip 后 global norm ≤ max_norm+1e-6；per-param grad atol=1e-5 vs 单卡；全局 norm 而非 per-shard norm
L2-08	test_dynamic_sequence_length	_test_fully_shard_precision.py	2	连续 5 步不同 seq_len 无 hang（60s timeout）；每步 loss finite；grad norm finite；param shape 不变
L2-09	test_pp_fsdp_coupling	_test_fully_shard_with_pp.py	4（2PP×2FSDP）	loss atol=1e-5 vs 4卡 DDP；stage boundary grad atol=1e-5；无死锁（60s timeout）
L2-10	test_ep_fsdp_coupling	_test_fully_shard_with_ep.py	4（2EP×2FSDP）	loss atol=1e-4；per-expert grad norm atol=1e-4；无 NCCL 报错

L2 目标文件命名说明：

文件	归档用例
_test_fully_shard_precision.py（已有）	L2-01、L2-02、L2-07、L2-08
_test_fully_shard_with_ac.py（新建）	L2-03、L2-04、L2-05、L2-06
_test_fully_shard_with_pp.py（新建）	L2-09
_test_fully_shard_with_ep.py（新建）	L2-10

---

Level 3：真实网络结构测试

原则：

• 双校验（loss 标量 + global grad norm），两者须同时通过
• 至少跑 2 步（step0 + optimizer 更新后的 step1），验证 DTensor→optimizer 交接正确性
• 固定 seed，同进程同卡数重跑须 bitwise 一致，FSDP 下的不确定性是 blocking bug
• 每个网络结构须覆盖 1D-mesh（FSDP）和 2D-mesh（HSDP）vs DDP 的精度自洽性

L3 用例列表

#	用例名	目标文件	卡数	架构摘要	并行策略
L3-01	test_llm_decoder_hsdp_vs_ddp	_test_fully_shard_with_llm.py	4	4层 Transformer Decoder（RMSNorm+MHA+SwiGLU）+ Embedding + 权重共享输出层；vocab=1024, hidden=256, seq=128	HSDP (replicate=2, shard=2) vs DDP
L3-02	test_llm_decoder_fsdp_1d_vs_ddp	_test_fully_shard_with_llm.py	2	同 L3-01 架构	FSDP 1D vs DDP
L3-03	test_moe_llm_ep_fsdp_vs_ddp	_test_fully_shard_with_llm.py	4	4层 Decoder（2层标准FFN + 2层MoE：4 expert, top-2 routing, expert_dim=2×hidden）；hidden=256, seq=64	3D mesh (replicate=1, fsdp=2, ep=2) vs 单卡
L3-04	test_vit_fsdp_vs_ddp	_test_fully_shard_with_vit.py	2	patch embedding + 6 encoder block（LN+MHA+LN+MLP-GELU）+ CLS token + 分类头；hidden=256, img=224	FSDP 1D vs DDP；patch embedding 和分类头不分片
L3-05	test_chunked_ce_fsdp_vs_ddp	_test_fully_shard_with_llm.py	2	2层 Decoder + 大 vocab 输出层（vocab=32768）；CE 分 chunk_size=512 计算	FSDP 1D vs DDP；per-chunk loss 一致性额外验证
L3-06	test_multimodal_deepstack_fsdp	_test_fully_shard_with_vit.py	2	4-block ViT（hidden=128）→ 2层 projection MLP → 2层 Decoder（hidden=256）；两子模型独立 fully_shard	FSDP 1D；projection MLP 不分片

各 L3 用例断言规格

L3-01 / L3-02（LLM 1D/2D-mesh vs DDP）：

• step0 loss atol=1e-5；step1 loss（optimizer 更新后）atol=1e-5
• global grad norm atol=1e-5
• 首层 attention weight grad tensor atol=1e-5
• 末层 FFN weight grad tensor atol=1e-5

L3-03（MoE EP+FSDP vs 单卡）：

• total loss atol=1e-4；dense 层 grad norm atol=1e-5；per-expert grad norm atol=1e-4
• 无 all-to-all / all-gather 通信冲突（无 NCCL 报错）

L3-04（ViT FSDP vs DDP）：

• loss atol=1e-5；grad norm atol=1e-5
• patch embedding grad（不分片）：所有 rank bitwise 一致且等于 DDP 参考
• 末层 MLP weight grad atol=1e-5

L3-05（Chunked CE FSDP vs DDP）：

• total loss atol=1e-5；per-chunk loss atol=1e-5；sum(per_chunk_losses) == total_loss（atol=1e-6，一致性验证）
• grad norm atol=1e-5；输出层 weight grad norm atol=1e-5

L3-06（多模态 Deep Stack）：

• E2E loss atol=1e-5；ViT grad norm atol=1e-5；Decoder grad norm atol=1e-5
• projection MLP grad（不分片）atol=1e-5
• 跨模块梯度流通验证：projection MLP grad ≠ 0（梯度未在模块边界断流）

L3 目标文件命名说明：

文件	归档用例
_test_fully_shard_with_llm.py（新建）	L3-01、L3-02、L3-03、L3-05
_test_fully_shard_with_vit.py（新建）	L3-04、L3-06

---

精度自洽矩阵

并行配置	loss tol	grad norm tol	覆盖用例
FSDP 1D, fp32	1e-5	1e-5	L1-03, L2-01, L3-02, L3-04, L3-05
FSDP 1D, bf16	1e-3	1e-3	L1-10, L1-11
HSDP 2D, fp32	1e-5	1e-5	L1-02, L3-01
TP + FSDP, fp32	1e-5	1e-5	已有 TP+FSDP E2E 用例
EP + FSDP, fp32	1e-4	1e-4	L2-10, L3-03
AC + FSDP, fp32	1e-5	1e-5	L2-03, L2-04, L2-05
AC + swap + FSDP, fp32	1e-5	1e-5	L2-06
PP + FSDP, fp32	1e-5	1e-5	L2-09

精度自洽规则：

1. baseline 锚点固定：baseline 固定为单卡 fp32 或等价 DDP，禁止 FSDP-vs-FSDP 互比
2. 不允许静默放宽容忍度：若测试在 1e-5 失败但 1e-4 通过，必须 triage 根因，说明存在数值累积 bug，不允许静默放宽
3. 双指标独立通过：grad norm 通过 ≠ 精度正确——grad norm 通过但张量元素级比对失败，说明存在 norm-cancellation bug；两个指标须独立通过
4. 2 步验证：Level 3 至少跑 2 步，验证 DTensor→optimizer 交接不断
5. 确定性要求：Level 3 固定 seed，同进程同卡数重跑须 bitwise 一致；FSDP 下的不确定性是 blocking bug

---

关键文件映射

角色	路径
核心 API	hyper_parallel/core/fully_shard/api.py
MixedPrecisionPolicy	hyper_parallel/core/fully_shard/utils.py
Torch 调度器	hyper_parallel/platform/torch/fully_shard/
现有接口测试（扩充 L1-01~05, L1-09）	tests/torch/fully_shard/_test_fully_shard_module.py
现有精度测试（扩充 L1-06, L1-08, L1-10~13, L2-01~02, L2-07~08）	tests/torch/fully_shard/_test_fully_shard_precision.py
现有参数单元（扩充 L1-07）	tests/torch/fully_shard/_test_hsdp_param.py
AC 耦合（新建，L2-03~06）	tests/torch/fully_shard/_test_fully_shard_with_ac.py
PP 耦合（新建，L2-09）	tests/torch/fully_shard/_test_fully_shard_with_pp.py
EP 耦合（新建，L2-10）	tests/torch/fully_shard/_test_fully_shard_with_ep.py
LLM 场景（新建，L3-01~03, L3-05）	tests/torch/fully_shard/_test_fully_shard_with_llm.py
ViT/多模态场景（新建，L3-04, L3-06）	tests/torch/fully_shard/_test_fully_shard_with_vit.py