本文描述 hyper_parallel/compile（Graph Mode）的终版设计，包括设计边界、目录结构、
配置模型、图捕获与 pass 流水线、训练接入链路、运行边界和后续演进。

---

0. 背景

Graph Mode 是与 eager 并行路径（fully_shard / tensor_parallel / pipeline_parallel）
互补的第二条接入方式：不采用 module wrapper，而是把整个训练步（forward + backward）
捕获成一张 joint FX graph，在图上完成并行切分与通信优化，再以单图执行。

用户只写模型代码 + 并行配置，框架负责：图捕获 → 并行分区 → 通信计算重叠 → 执行。

当前覆盖 FSDP（含与 automodel TP 的混合接入）；TP/PP/EP 的图内 pass 为后续演进项。

---

1. 目标

• 声明式：PassPlan 以 FQN / 通配 pattern 标记分片模块，无侵入式 wrap，
  用户不改模型代码。
• 单图执行：forward + backward 捕获为一张图，运行期在 torch.no_grad() 下执行，
  消除 autograd 图维护开销与中间值保活。
• FSDP-agnostic optimizer：FSDP pass 对 live model 参数原地物理分片，
  model.parameters() 即本地分片，optimizer / trainer 无 FSDP 感知。
• 通信计算重叠：通信发出后延迟 wait_tensor，与独立计算重叠。
• 可组合：pass 流水线由 PassConfig 驱动，pass 可插拔、可扩展。

---

2. 设计边界

Graph Mode 负责：

• joint fwd+bwd FX graph 捕获，参数/buffers 静态输入化；
• GraphPass 抽象与 PassPipeline 编排；
• FSDP 分区 pass（参数 all_gather、梯度 reduce_scatter、live model 原地分片）；
• 通信-计算重叠 pass（wait_tensor 下沉调度）；
• PassConfig / PassPlan 配置模型与 YAML 解析；
• GraphTrainer：懒编译、训练循环、梯度累积、optimizer step；
• 与上层 MeshContext 的 TP+FSDP 混合接入（TP 组复用、fsdp 子组注册）。

Graph Mode 不负责：

• TP collectives 本身（TP 通信在 boundary forward 内由上层烘焙，不进图）；
• PP / EP 图内 pass（规划项，本阶段不落地）；
• activation checkpoint / activation swap；
• checkpoint 权重读写与转换；
• 通用 Trainer / 数据 / 优化器生态（内置 Adam 仅为最小闭环）；
• DeviceMesh 生命周期管理（复用上层 MeshContext 或 init_device_mesh）；
• MindSpore 后端（torch-only 例外，见 §12）。

---

3. 目录和文件清单

3.1 compile 包文件树

hyper_parallel/compile/
├── __init__.py            # 对外 API：GraphTrainer / PassConfig / PassPlan / ...
├── parallel_config.py     # PassConfig：并行配置（torch-free 纯数据类）
├── sharding_config.py     # PassPlan / FSDPModuleConfig + YAML 解析
├── trainer.py             # GraphTrainer：编译与训练循环
├── tracer/
│   └── graph_tracer.py    # joint fwd+bwd 图捕获（make_fx + autograd.grad）
├── passes/
│   ├── base.py            # GraphPass 抽象基类
│   ├── pipeline.py        # PassPipeline 编排 + 基础优化 pass
│   ├── parallel/
│   │   └── fsdp_pass.py   # FSDP 分区 pass
│   └── overlap/
│       └── schedule.py    # AutoOverlapPass（wait_tensor 调度，当前为占位实现）
├── README.md              # 架构说明（当前 Graph Mode 文档的 SoT）
└── examples/
    ├── simple_model/      # 纯 FSDP 图模式（YAML + 8 卡脚本）
    └── automodel_tp/      # TP+FSDP 混合（automodel MeshContext，tp=2 dp=2）

3.2 非 compile 接入文件

hyper_parallel/distributed/
└── mesh.py  # MeshContext：GraphTrainer 复用其 TP 组，将 fsdp_non_moe_mesh 的
             # fsdp_shard 维（或 device_mesh 的 dp 维）子组注册为 "fsdp"，
             # 并回填 pass_config.fsdp_degree

tests/ut/compile/         # Graph Mode 单元测试（见 §13）

3.3 文档文件

docs/guide/
└── graph_mode.md          # TODO：用户指南；落地前以 compile/README.md 为准

---

4. 配置模型

Graph Mode 不引入第二份独立配置体系：配置只来自 PassConfig + PassPlan。
两者均为纯数据类（torch-free），可独立于运行时导入与序列化。

4.1 PassConfig

字段	默认	语义
enable_overlap	True	驱动 AutoOverlapPass 移动 wait_tensor
fsdp_enabled	True	驱动 FSDPPass；置 False 用于纯 TP/PP 图模式。分布式未初始化或 world_size == 1 时 FSDPPass 自身早退
fsdp_degree	None	FSDP 组大小。None 表示运行期回填：外部 Mesh 场景由 GraphTrainer 从子 mesh 回填，纯 FSDP 场景回退 world_size
tp_size	1	信息字段（TP collectives 在 boundary forward 内），供未来 TP-aware pass 读取，避免 API churn
sequence_parallel	False	TP 轴 sequence parallel
loss_parallel	False	TP 轴 loss parallel

validate() 校验 tp_size >= 1、fsdp_degree 为 None 或正整数。
fsdp_enabled 不探测 torch.distributed（分布式初始化判断在 FSDPPass.run 内），
保持配置类 torch-free。

4.2 PassPlan 与 YAML

• 双注册表：fsdp_modules（exact FQN）+ fsdp_patterns（fnmatch 通配）。
  查询时 exact 优先，pattern 按插入序取首个匹配。
• merge()：合并两个 plan，key 冲突时 other 胜，不原地修改。
• YAML 契约（create_sharding_plan_from_yaml）：

fsdp:
  enabled: true        # 缺省时：modules / patterns 非空即视为 enabled
  modules:
    - name: tok_embeddings
  patterns:
    - pattern: "layers.*"

仅 name / pattern 被解析消费。历史示例中的 reshard_after_forward /
use_cpu_offload 等字段已随死配置面移除，parser 不读取；待真实 reshard /
offload pass 落地时再带消费方加回。

---

5. 编译流程

编译由 GraphTrainer 驱动，首个 batch 懒触发（也可显式调用 compile()）：

GraphTrainer(model, train_fn, pass_config, pass_plan)
  └─ compile(sample_input, sample_label)
       1. _pytree_pre_hook()            # 可选：模型特定 pytree 注册（如 BlockMask）
       2. fsdp_enabled 且 dist 已初始化 -> _init_device_mesh()
       3. trace_model_graph(...)        # 捕获 joint FX graph
       4. PassPipeline.from_config(pass_config, pass_plan)
       5. pipeline.run(graph_module, model=..., fsdp_group_name="fsdp")
       6. _init_optimizer()             # 建立在已分片的 live model 参数上
  └─ train(): train_step -> _run_graph -> _accumulate_grads -> optimizer_step

pass 流水线固定顺序：

DeadCodeElimination -> CanonicalizeGraph -> [FSDPPass] -> [AutoOverlapPass]

FSDPPass 仅在 fsdp_enabled 时加入，AutoOverlapPass 仅在 enable_overlap
时加入。后端编译槽位（inductor 类后端）刻意预留未接线。

5.1 Device mesh 初始化（双模式）

• 外部 MeshContext（automodel 路径）：TP 组已由上层创建，直接复用；解析
  fsdp_non_moe_mesh 的 fsdp_shard 维（或 device_mesh 的 dp 维）子组，
  以 _register_process_group("fsdp", pg) 注册，并回填 fsdp_degree。
  TP+FSDP 混合时 FSDP 组是全世界的真子组，不能用 world_size，否则会沿 TP
  轴过度分片。
• 兜底：init_device_mesh 构建一维 ("fsdp",) 全 world mesh（纯 FSDP 路径）。

5.2 梯度映射与累积

• 图按 state_fqns 序输出梯度（共享参数按 FQN 重复出现）；
  _accumulate_grads 按 FQN 对齐到 live model 参数，数量不一致时直接报错拒绝
  赋值，防止 zip 静默截断。
• 累积而非覆盖：多个 micro-batch 的 train_step 可先跑后 optimizer_step
  （以 zero_grad 收尾）。
• 单进程或未初始化分布式时 FSDPPass 早退，退化为纯图模式
  （图捕获与执行照常，无分片）。

---

6. 图捕获（tracer）

设计参考 torchtitan 的 experiments/graph_trainer/minimal_fx_tracer：

1. 参数静态输入化：参数/buffers 经
   torch.nn.utils.stateless._reparametrize_module 从 live model 抽出为 flat
   tensors，作为图的 leading placeholders（而非 get_attr 节点）。图内因此
   无参数 get_attr 节点，pass 可通过 reshape placeholder 切分图。
2. joint graph：trace 函数体内执行 torch.autograd.grad(loss, params)。
   为使其可被追踪，需要：
   • patched autograd engine：torch.compiler._patch_engine_backward（torchtitan
     patched torch 提供；stock torch 缺失时 RuntimeWarning 降级）；
   • backward 单线程：torch.autograd.set_multithreading_enabled(False)，
     否则 C++ autograd 引擎将 backward 派发到 worker 线程并以全新
     contextvars.Context 运行，tracing context 丢失。
3. 追踪上下文：FakeTensorMode(allow_non_fake_inputs=True) + 非 strict
   tracing + preserve_node_meta() + _skip_nested_compile()（dynamo 内联
   torch.compile 包装的函数，如 compiled attention kernel）。
4. 叶子类型白名单：Tensor / int / float / bool / str / None。其余
   （callable、自定义对象）必须在 train_fn 闭包内捕获，或注册为 pytree 节点。
5. 运行期：图以 torch.no_grad() 执行——图内已含显式 backward ops，
   再叠加 autograd 图会通过 grad_fn 引用保活全部前向中间值。
6. torch 版本兼容：make_fx 的 record_stack_traces /
   record_module_stack kwargs 按版本探测转发；fwd→bwd 元数据拷贝容忍缺失
   stack_trace。
7. 元数据：GraphModule 附带 state_fqns / state_is_param，供梯度映射；
   state_is_param 缺失（旧 trace）时降级为仅参数、跳过 buffer，而非抛错。

---

7. FSDP 分区 pass

全部并行逻辑集中在 pass 内，trainer 只喂 model.parameters()：

1. 依据 PassPlan 识别 FSDP 参数 placeholder（exact 优先，pattern 插入序
   首个匹配；plan 为 None 时全量分片）。
2. 在每个该类 placeholder 之后插入 all_gather（Shard → Replicate）：计算
   主体使用全量参数，图输入保持分片。
3. 在该类参数的梯度输出上插入 reduce_scatter（Replicate → Shard）；
   非 FSDP 参数的梯度保持 full。
4. live model 参数原地物理分片（dim 0，按 FSDP rank）：
   model.parameters() 即本地分片，optimizer / trainer 无 FSDP 感知。

组解析：按注册名 "fsdp" 解析 process group；fsdp_degree 从
pass_config.fsdp_degree 读取，缺省回退 world_size。

约束：各参数 dim 0 必须能被 fsdp_degree 整除。

---

8. 通信-计算重叠 pass

• FSDPPass 插入通信时 wait_tensor 紧跟通信 op（保证正确性）。
• AutoOverlapPass 的目标：把每个 wait_tensor 下沉到其首个使用者之前，
  使通信与独立计算重叠。
• 现状：_find_wait_nodes 已完成 wait_tensor 节点识别；
  _move_wait_later 为占位 no-op，调度器落地时替换该函数即可，
  pipeline 结构不变。
• 后端编译槽位（inductor / NPU 图编译）在 PassPipeline 中预留，待后端
  集成落地时加入。

---

9. TP / SP 混合接入

• TP collectives 不进图：TP 通信在 boundary forward 内由上层
  （automodel 路径）烘焙完成，boundary forward 直接持有组对象。
  graph-mode 只负责 FSDP 维。
• GraphTrainer 接收 mesh_context 后：
  1. 取 mesh_context.fsdp_non_moe_mesh（缺省回退 device_mesh）；
  2. 选 fsdp_shard 维（automodel）或 dp 维（cp=1 场景）；
  3. 注册 "fsdp" 组并回填 fsdp_degree。
• 上层保证 device_mesh["tp"] 与 fsdp_non_moe_mesh["tp"] 的 rank 组与
  本地 rank 一致性校验。
• PassConfig.tp_size / sequence_parallel / loss_parallel 为信息字段：
  本阶段不驱动图内 pass，保留给未来 TP-aware pass。

---

10. 训练接入点

• 入口：

GraphTrainer(
    model,                    # nn.Module
    train_fn,                 # (model, input, label) -> loss
    pass_config,              # PassConfig
    pass_plan=None,           # PassPlan（声明式分片）
    optimizer_config=None,    # {"lr": ..., "grad_clip": ...}
    device=None,              # 缺省探测 npu -> cpu
    mesh_context=None,        # automodel MeshContext（TP+FSDP 混合）
)

• 训练循环：train(dataloader, max_steps=None, log_interval=None, log_fn=None)。
  每个 batch 移至 self.device，随后 train_step（懒编译）+ optimizer_step。
• 扩展点：
  • set_pytree_pre_hook(hook)：编译前注册模型特定 pytree
    （如 flex-attention BlockMask）；
  • to(device)：模型与后续 batch 落位。
• 与 eager 路径的关系：GraphTrainer 独立于 BaseTrainer / eager FSDP；
  TP 已分片的模型可经 mesh_context 注入，graph-mode 只补 FSDP 维。
• optimizer：内置 Adam（lr / grad_clip 来自 optimizer_config），
  建立于分片后参数之上。

---

11. 运行环境要求

• torch：joint graph 捕获依赖 torch.compiler._patch_engine_backward
  （torchtitan patched torch 或较新版本）。stock torch 会 RuntimeWarning，
  且 autograd.grad 的 joint capture 可能不正确。
• 分布式：FSDP 需要已初始化的 torch.distributed；world_size == 1
  或未初始化时 FSDPPass 早退（纯图模式）。FSDP 组大小需整除各参数 dim 0。
• 设备：device 缺省探测 npu（torch.npu.is_available()）→ cpu；
  mesh device_type 同规则。
• CI 例外：torch_npu 已安装但无 NPU 时，Adam 的 foreach/fused 内核选择
  会探测 torch_npu.npu.current_device() 并触发 _lazy_init 失败；通过
  foreach=False 显式绕过。NPU 运行保持默认 foreach/fused 路径。

已知限制：

• 仅支持 FSDP（TP/EP/PP 图内 pass 规划中）；
• 一次编译冻结 shape，无动态 shape / 重编译策略；
• mesh device_type 探测当前仅覆盖 npu/cpu，纯 GPU 环境会回落到 cpu
  （待补齐，见 §14）。

---

12. 平台边界

• compile/ 是平台分层规则的显式 Torch-only 例外（与 core/multicore
  同类处理）：依赖 patched autograd engine、torch.fx 与 functional
  collectives 的图内表示，这些均无平台无关抽象。模块内直接 import torch
  是有意为之，compile/README.md 已记录该例外。
• 待图 trace / FSDP-pass API 出现 get_platform() 级抽象后，再评估
  MindSpore 支持。
• 编排层保持 torch-free：PassConfig / PassPlan / PassPipeline 不 import
  torch（torch.fx 仅 TYPE_CHECKING 延迟导入），使配置与编排可被其他
  后端复用。

---

13. 测试覆盖

tests/ut/compile/（单卡 UT）：

• test_tracer.py：捕获契约（静态输入、joint 图、元数据）；
• test_fsdp_pass.py：collective 插入、live model 分片；
• test_pass_pipeline.py：流水线组装与 pass 选择；
• test_parallel_config.py / test_sharding_plan.py：配置与 YAML 契约；
• test_auto_overlap.py：wait_tensor 识别与调度；
• test_tpsp.py：TP/SP 适配契约；
• test_trainer.py：训练循环、梯度映射、FSDP 门控。

后续补齐：

• overlap 调度器真实下沉后的重叠性测试；
• TP / PP / EP pass 落地后的分布式 ST（torchrun_case）；
• MindSpore 后端（若引入）。

---

14. 后续演进（TODO）

按优先级：

1. overlap 调度器落地：_move_wait_later 真实现——依赖分析 + 将
   wait_tensor 下沉到首个使用者之前。
2. reshard / offload pass：FSDPModuleConfig 曾以 dead surface 移除
   reshard_after_forward / use_cpu_offload / wrap_separately；
   待真实 pass 落地时带真实消费方再加回。
3. TP-aware / PP / EP 图内 pass：README Limitations 已声明 TP/EP/PP
   规划中；PP 需要 schedule 与 micro-batch 在图上的表示，执行层扩展位在
   PassPipeline 的分区层。
4. 后端编译槽位：PassPipeline 接入 inductor / NPU 图编译后端。
5. GPU device_type 探测：补齐 cuda 分支（当前 npu/cpu 两分支）。
6. 动态 shape 与重编译策略：当前一次编译冻结 shape。
7. optimizer 抽象：脱离内置 Adam，接入上层 optimizer 生态。
8. 文档：docs/guide/graph_mode.md 用户指南；清理示例 YAML 中未被
   parser 消费的字段。