🚀 功能描述

在大模型分布式训练中，用户通常需要先在真实 NPU/GPU 集群上完成模型构建、并行切分和训练步骤，才能发现显存不足、策略组合不合理或流水线阶段内存不均衡等问题。该过程成本高，且难以在资源受限场景下提前定位问题。

希望提供 Dry-run 能力：仅基于模型 config.json 在 Meta/FakeTensor 上模拟一次训练步骤，不加载模型权重、tokenizer 或真实数据集，即可输出各 rank 的内存摘要和内存生命周期报告。

该能力应覆盖常见分布式策略及组合，包括 TP、CP、EP、FSDP/HSDP、PP，以及 PP 的 GPipe、1F1B、VPP 等调度场景；同时支持通过配置描述 attention、MoE routing、TP cross-entropy 等影响 shape、通信量和内存生命周期的值依赖。

输出结果包括：

• 内存生命周期，展示 Parameter、Gradient、Optimizer State、Activation 等对象的申请和释放过程；
• 未配置的 FakeTensor 值依赖诊断信息，辅助定位模型分支问题。

该能力用于训练前的内存风险评估和并行策略分析，不能替代真实设备上的性能 Profiling。

现有替代方案

当前可通过真实设备训练或 ProfilingCallback 采集内存数据，但需要准备模型权重、数据集和多卡 NPU/GPU 资源，且 OOM 往往只能在运行后发现。

也可以根据参数量和激活公式手工估算显存，但难以覆盖 FSDP/HSDP 重分片、PP micro-batch 生命周期、MoE 路由不均衡及混合并行组合，结果不够直观和可复现。

与DFX相关性DF

• 在不占用真实加速卡的情况下提前识别 OOM 风险；
• 通过 rank 级内存摘要和生命周期详细统计展示峰值来源与内存变化过程；
• 对比不同并行策略、PP 调度和 micro-batch 配置下的内存差异；
• 在 FakeTensor 遇到依赖真实数值的模型分支时，输出模块路径、ATen 算子和源码位置，提升问题定位效率；

提议的新API

新增 dry_run 配置项及 DryRunConfig，用于启用内存模拟并指定：

• output_dir：报告输出目录；
• pipeline_stage_builder：PP 场景的 stage 构建方法；
• value_dependencies.rules：影响 shape、通信或内存生命周期的值依赖规则。

是否影响现有API

不影响现有训练和并行 API 的兼容性。
Dry-run 仅在 dry_run.enabled=true 时生效；
未启用时，仍使用原有 Trainer 训练路径；
不修改 MeshContext.device_mesh、apply_sharding_plan() 返回值或 FSDP/HSDP 的现有语义；
PP Dry-run 的 stage 装配逻辑作为内部实现，不新增正式 PP 公共 API。

补充信息

Thanks for contributing 🎉!