变更摘要

本 PR 引入面向 Hugging Face Qwen3-MoE 模型的注意力激活交换（attention activation swap）能力：新增 activation_swap 配置项（取值 none / attention，默认 none），并新增 hyper_models/components/distributed/attention_swap.py 模块，利用 swap_wrapper、SwapManager 与 CheckpointPolicy 对 Qwen3MoeAttention 进行包装和逐层 prefetch 调度，将注意力前向中的大张量激活换出以降低显存占用。同时将该能力接入模型构建链路（auto_model.py、infrastructure.py、base.py），并在应用 sharding、torch.compile、FSDP 之前完成包装，且校验其与编译、激活重计算、流水线并行等特性的互斥关系。

主要改动

• 新增配置项 activation_swap：在 TrainerConfig 中新增 activation_swap: Literal["none", "attention"] = "none" 字段，并在 examples/training_demo/train.yaml 中增加 activation_swap: none # attention, none 配置项。
• 新增 attention_swap.py 模块：实现 qwen3_attention_swap_policy（按张量是否可导、维度、存储是否独立、大小是否超过 _MIN_SWAP_TENSOR_BYTES 决定 MUST_SWAP / MUST_SAVE）、validate_attention_swap（校验非法模式及与 torch.compile、activation checkpointing、pp_size != 1 的互斥）以及 apply_qwen3_moe_attention_swap（校验模型层级结构、用 swap_wrapper 替换各层 self_attn 并通过 SwapManager.set_forward_prefetch_layer 建立层间 prefetch）。
• 模型构建链路接入：auto_model.py 的 from_pretrained、from_config 及核心构建函数均新增 activation_swap 参数并透传；apply_model_infrastructure 新增参数并在 Step 9-2 于 sharding、compile、FSDP 之前执行校验与 attention swap 包装。
• 训练器透传配置：base.py 在构建模型时将 self.config.activation_swap 传入，测试辅助函数 _model_target 及对应断言同步增加 activation_swap 字段。
• 新增测试覆盖：新增 tests/hyper_models/trainer/test_attention_swap.py，覆盖 policy 过滤、互斥特性校验、仅包装 attention 且调度 47 层 prefetch 等场景；test_resolver.py 增加 activation_swap 默认值、合法值及非法值解析测试。

代码审查

我已逐一审查全部 9 个变更文件。下面是审查结论汇总。

审查结论（按文件）

1. examples/training_demo/train.yaml — 无问题。activation_swap: none 为顶层键，缩进正确，与 TrainerConfig.activation_swap 字段对应；none 会被 PyYAML 解析为字符串而非 null，能通过 resolver 的 Literal 校验。
2. hyper_models/_transformers/auto_model.py — 发现问题：_build_model 删除了 distributed_setup=distributed_setup 透传（P1）。
3. hyper_models/_transformers/infrastructure.py — 发现问题：attention swap 被放在 sharding 之后，违反其自身注释的"sharding 前安装"前置条件（P2）。
4. hyper_models/components/distributed/attention_swap.py — 发现问题：未使用的 get_platform/PlatformType 导入与模块级 get_platform() 副作用（P3）。
5. hyper_models/trainer/base.py — 无问题。activation_swap=self.config.activation_swap 透传正确，且 Target.build 会按签名过滤未知参数，不会破坏自定义 model target。
6. hyper_models/trainer/config.py — 无问题。activation_swap: Literal["none", "attention"] = "none" 与 resolver 的 Literal 处理及错误信息格式一致。
7. tests/hyper_models/trainer/test_attention_swap.py — 无问题。policy 边界用例（1MiB 阈值、view/共享存储、requires_grad）与 48 层→47 次 prefetch 断言均正确。
8. tests/hyper_models/trainer/test_data_build.py — 无问题。_model_target 新增 activation_swap 形参，_build_model 恒会传入，断言一致。
9. tests/ut/config/test_resolver.py — 无问题。默认值/合法值/非法值三测与 resolver 行为匹配。

发现数量统计

• P0：0
• P1：1
• P2：1
• P3：1

整体风险判断

中高风险。核心风险是 P1 回归：distributed_setup 不再透传给 apply_model_infrastructure，导致 plan_overrides 中的 replace_module 模块替换被静默跳过，属于无报错的生产功能回归，建议合入前必须修复。P2 的 attention swap 安装顺序问题在启用 TP/CP/EP 分片时会触发精确类型校验失败，也应在合入前澄清设计（前移安装或扩展 validate 拒绝矩阵）。P3 为死代码/import 期副作用，属于可选清理。

⛔ 需要修改