task: hyper_parallel.fully_shard兼容接口泛化能力提升

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-1013, L2-0102, 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

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

Start by reading hyper_parallel/core/fully_shard/api.py and utils.py, then run the existing tests under tests/torch/fully_shard/. Extend the named module, precision, and parameter test files, and add the specified AC, PP, EP, LLM, and ViT test files. Done means the listed L1–L3 cases pass with their stated accuracy, gradient, determinism, and timeout checks.