tuner
所有候选配置 (exps)
│
▼
Tuner 选出下一批配置
┌──────────────────────────────┐
│ RandomTuner → 随机选 │
│ GridSearchTuner → 顺序选 │
│ ModelBasedTuner → XGBoost │
│ 预测性能最好的配置优先选 │
└──────────────────────────────┘
│
▼
ResourceManager 实际跑训练实验
│
▼
解析结果 (吞吐量/延迟)
│
▼
更新 cost model，继续迭代
│
▼
找到 best_exp（最优 DeepSpeed 配置）passes/ 是 DeepSpeed 编译优化的核心，通过对 FX 计算图的图变换（graph transformation），将 ZeRO 显存优化、参数/激活 offload、通信预取、长上下文 checkpointing 等功能"编译进"模型执行图，从而在 torch.compile 路径下仍能享受 DeepSpeed 的所有显存优化特性。
torch.compile 捕获 FX 计算图
│
▼
passes/ 中的各个 Pass 依次处理
┌─────────────────────────────────────┐
│ zero3_compile → 插入 all-gather │
│ selective_gather→ 智能参数持久化 │
│ prefetch → 通信计算重叠 │
│ offload_* → 显存 ↔ CPU卸载 │
│ checkpointing → 重计算策略 │
└─────────────────────────────────────┘
│
▼
优化后的计算图交给 inductor 编译执行整体逻辑
所有候选配置 (exps)
│
▼
Tuner 选出下一批配置
┌──────────────────────────────┐
│ RandomTuner → 随机选 │
│ GridSearchTuner → 顺序选 │
│ ModelBasedTuner → XGBoost │
│ 预测性能最好的配置优先选 │
└──────────────────────────────┘
│
▼
ResourceManager 实际跑训练实验
│
▼
解析结果 (吞吐量/延迟)
│
▼
更新 cost model，继续迭代
│
▼
找到 best_exp（最优 DeepSpeed 配置）本质上是在 tune DeepSpeed 的训练配置 ，目标是用尽可能少的实际实验次数，找到让模型训练速度最快的那组 ds_config。
这个 passes/ 文件夹是 DeepSpeed 编译后端（compile）的图优化 Pass 集合 ，核心作用是：在 torch.compile 生成的计算图（FX Graph）上，插入/改写节点，以实现显存优化、通信优化等功能。
文件功能总结
| 文件 | 功能 |
|---|---|
| init.py | Pass 调度器 run_opt_passes：串行执行多个 pass，每次执行后做内存 profiling 验证 |
| zero3_compile.py | ZeRO-3 图编译 pass ：在 FX 图中插入 allgather_param / wait_allgather / reduce_scatter 等通信节点，实现参数的按需 All-Gather（分片参数在使用前聚合，用完后释放） |
| selective_gather.py | 选择性参数持久化 pass：根据显存预算，决定哪些参数在整个前向/反向过程中常驻 GPU（不需要每次 all-gather），以减少通信开销 |
| prefetch.py | All-Gather 预取 pass ：分析计算图中通信与计算的时序，提前调度 allgather_param，让通信与计算重叠，隐藏通信延迟 |
| offload_parameters.py | 参数 offload pass ：在图中插入 offload_parameter / reload_parameter 节点，将暂时不用的参数卸载到 CPU 内存 |
| offload_activation.py | Activation offload pass：将前向中间激活值卸载到 CPU，在反向需要时再取回，节省 GPU 显存 |
offload_adam_states.py |
Adam 优化器状态 offload pass：将一阶/二阶矩等优化器状态卸载到 CPU |
| long_context_checkpointing.py | 长上下文 Activation Checkpointing pass ：针对长序列场景，定制 torch.compile 的重计算策略（修改 solve_min_cut 中的 should_ban_recomputation 逻辑），强制保存 attention 等关键节点，避免长序列下的 OOM |
sp_compile.py |
Sequence Parallel（序列并行）相关编译 pass |
zero1_compile.py |
ZeRO-1（优化器状态分片）相关编译 pass |
整体定位
torch.compile 捕获 FX 计算图
│
▼
passes/ 中的各个 Pass 依次处理
┌─────────────────────────────────────┐
│ zero3_compile → 插入 all-gather │
│ selective_gather→ 智能参数持久化 │
│ prefetch → 通信计算重叠 │
│ offload_* → 显存 ↔ CPU卸载 │
│ checkpointing → 重计算策略 │
└─────────────────────────────────────┘
│
▼
优化后的计算图交给 inductor 编译执行支持的压缩技术
| 技术 | 对应常量 | 作用 |
|---|---|---|
| 层裁剪（Layer Reduction） | LAYER_REDUCTION |
减少 Transformer 层数，用教师模型的权重初始化学生模型（你选中的代码就在这里） |
| 权重量化（Weight Quantization） | WEIGHT_QUANTIZATION |
将权重从 float32 量化为 INT8/INT4 等低精度，减少内存占用 |
| 激活量化（Activation Quantization） | ACTIVATION_QUANTIZATION |
对激活值做量化 |
| 稀疏剪枝（Sparse Pruning） | SPARSE_PRUNING |
将权重中不重要的值置零，产生稀疏矩阵 |
| 行剪枝（Row Pruning） | ROW_PRUNING |
裁掉整行权重，真正减少矩阵维度 |
| Head 剪枝（Head Pruning） | HEAD_PRUNING |
裁剪 Attention Head |
| 通道剪枝（Channel Pruning） | CHANNEL_PRUNING |
裁剪卷积/线性层的输出通道 |
一句话总结 ：这里的压缩是量化 + 剪枝 + 层裁剪的综合框架，蒸馏只是层裁剪中用于权重初始化的辅助手段，不是唯一的压缩方式。
1. 权重量化（Weight/Activation Quantization）
原理：精度冗余
神经网络训练用 float32（32位），但推理时模型对数值精度不敏感。
float32: 0.3141592653... (32 bits)
INT8: 0.31 (8 bits) → 精度略降，但体积缩小 4x实验证明，大多数权重的数值范围很窄，用低精度表示损失极小，而内存和计算量大幅减少。
2. 稀疏剪枝（Sparse Pruning）
原理：数值冗余（接近零的权重无贡献）
网络中大量权重的绝对值非常小，对输出几乎没有贡献：
W = [0.001, 0.85, -0.002, 0.76, 0.0003]
↓ 剪枝
W = [0, 0.85, 0, 0.76, 0 ]将这些接近零的权重强制置零，用稀疏矩阵存储，可节省内存和计算。
3. 结构化剪枝（Row / Head / Channel Pruning）
原理：结构冗余（整个神经元/通道/头部不重要）
稀疏剪枝产生的零散零值在硬件上难以加速。结构化剪枝则直接删除整行/整列/整个 Attention Head：
原矩阵 [768 × 768] → 删掉不重要的行 → [512 × 768]评估重要性的标准通常是梯度大小或激活值的 L1/L2 范数。删除后矩阵物理变小，硬件可以直接加速。
4. 层裁剪（Layer Reduction）
原理：层间冗余（相邻层学到相似的表示）
研究发现 Transformer 的中间层存在大量表示相似性，某些层几乎在做重复的事：
12层 BERT → 删掉冗余的中间层 → 6层 Student BERT直接删层后精度会大幅下降，所以需要用教师模型对应层的权重做初始化（即 student_initialization），再继续蒸馏训练来恢复精度。
对比总结
| 技术 | 利用的冗余类型 | 压缩维度 | 硬件加速友好度 |
|---|---|---|---|
| 权重量化 | 精度冗余 | bit 宽度 | ⭐⭐⭐⭐⭐ |
| 稀疏剪枝 | 数值冗余（零散） | 非零元素数量 | ⭐⭐（需稀疏硬件支持） |
| 结构化剪枝 | 结构冗余（整块） | 矩阵物理尺寸 | ⭐⭐⭐⭐ |
| 层裁剪 | 层间冗余 | 网络深度 | ⭐⭐⭐⭐⭐ |
根本思想 ：神经网络是过参数化（over-parameterized）的，训练时需要大量参数保证收敛，但推理时大量参数是冗余的，压缩就是在精度损失可接受的前提下，系统性地去除这些冗余。
autotp
一、模型感知与自动策略推导
原来的 TP： 需要用户为每个模型提供明确的 replace_policy（容器策略），手动写出哪些层用行并行、哪些用列并行。
AutoTP 的改变：
tp_parser(model)自动遍历模型的ModuleList，通过分析子模块名称（如o_proj、down_proj、out_proj等）来推断哪些层需要 AllReduce（行并行），其余默认为列并行。supported(model)自动检测模型是否在不支持列表内，无需用户手动判断。
二、基于正则表达式的可配置分区 API（AutoTPConfig + TPLayerSpec）
这是此版本最核心的新增机制，通过 autotp_config.py 实现：
| 概念 | 说明 |
|---|---|
TPLayerSpec |
用正则 patterns 匹配参数名，声明 COLUMN/ROW/SKIP 三种分区方式 |
shape 参数 |
支持 (3, -1) 等分组描述，用于 fused QKV / packed MLP 等融合层 |
AutoTPConfig |
汇总所有 TPLayerSpec，支持 strict_mode（未匹配层报错）、use_default_specs 等开关 |
AutoTPPresets |
内置 llama / bloom / chatglm / mixtral / deepseek_v2 等常见模型的预设配置 |
相比原来仅凭层名称用 if/elif 硬编码路由，现在完全由配置驱动，用户只需传入 AutoTPConfig 即可覆盖或扩展行为。
三、_replace_module 增加双路由路径
python
# 新路径：partition_config 不为 None 时，按全名 + 正则匹配
if self.partition_config is not None:
new_child = self._replace_with_config(child, full_name)
# 传统路径：按 linear_policies 类型字典匹配
elif child.__class__ in self.linear_policies:
...- 新路径 ：追踪完整参数路径（
full_name），交给_replace_with_config用正则做细粒度匹配，支持同类型层在不同位置有不同策略（如 MoE 的 gate 层 SKIP，expert 层 ROW/COLUMN）。 - 传统路径 ：只按
nn.Linear/nn.Embedding类型分发，无法区分同类型但语义不同的层。
四、支持 SubParamLinearLayer / SubParamLinearAllreduce 处理融合层
原来 fused QKV 要靠 require_tp_fused_qkvw 启发式判断。新 API 中，通过 TPLayerSpec.shape 显式描述分组方式（如 (3, -1) 代表 Q/K/V 三等分，((q_size, kv_size, kv_size), -1) 代表 GQA 非等分），由 SubParamLinearLayer / SubParamLinearAllreduce 负责实际的 reshape + split，逻辑清晰且可扩展。
五、训练模式支持
python
if is_autotp_training_mode():
self.mp_group = groups.get_tensor_model_parallel_group()
self.mp_size = groups.get_tensor_model_parallel_world_size()原始 TP 主要面向推理，AutoTP 在 set_tensor_parallel_config 中增加了训练模式的进程组感知，可与 DeepSpeed 的训练并行组无缝集成。
一句话总结： AutoTP 的核心是把原来散落在各处的硬编码 if/elif 模型特判，重构为一套「正则匹配 + 声明式分区配置 + 内置预设」的统一框架，从而让 TP 的适配不再需要修改内部代码，只需传入一份 AutoTPConfig 即可。