# 参数卸载
## 📖 概述
Lightx2v 实现了先进的参数卸载机制，专为在有限硬件资源下处理大型模型推理而设计。该系统通过智能管理不同内存层次中的模型权重，提供了优秀的速度-内存平衡。
**核心特性：**
- **分block/phase卸载**：高效地以block/phase为单位管理模型权重，实现最优内存使用
- **Block**：Transformer模型的基本计算单元，包含完整的Transformer层（自注意力、交叉注意力、前馈网络等），是较大的内存管理单位
- **Phase**：Block内部的更细粒度计算阶段，包含单个计算组件（如自注意力、交叉注意力、前馈网络等），提供更精细的内存控制
- **多级存储支持**：GPU → CPU → 磁盘层次结构，配合智能缓存
- **异步操作**：使用 CUDA 流实现计算和数据传输的重叠
- **磁盘/NVMe 序列化**：当内存不足时支持二级存储
## 🎯 卸载策略
### 策略一：GPU-CPU 分block/phase卸载
**适用场景**：GPU 显存不足但系统内存充足
**工作原理**：在 GPU 和 CPU 内存之间以block或phase为单位管理模型权重，利用 CUDA 流实现计算和数据传输的重叠。Block包含完整的Transformer层，而Phase则是Block内部的单个计算组件。
**Block vs Phase 说明**：
- **Block粒度**：较大的内存管理单位，包含完整的Transformer层（自注意力、交叉注意力、前馈网络等），适合内存充足的情况，减少管理开销
- **Phase粒度**：更细粒度的内存管理，包含单个计算组件（如自注意力、交叉注意力、前馈网络等），适合内存受限的情况，提供更灵活的内存控制
**关键特性：**
- **异步传输**：使用三个不同优先级的CUDA流实现计算和传输的并行
- 计算流(priority=-1)：高优先级，负责当前计算
- GPU加载流(priority=0)：中优先级，负责从CPU到GPU的预取
- CPU加载流(priority=0)：中优先级，负责从GPU到CPU的卸载
- **预取机制**：提前将下一个block/phase加载到 GPU
- **智能缓存**：在 CPU 内存中维护权重缓存
- **流同步**：确保数据传输和计算的正确性
- **Swap操作**：计算完成后轮换block/phase位置，实现连续计算
### 策略二：磁盘-CPU-GPU 分block/phase卸载（延迟加载）
**适用场景**：GPU 显存和系统内存都不足
**工作原理**：在策略一的基础上引入磁盘存储，实现三级存储层次（磁盘 → CPU → GPU）。CPU继续作为缓存池，但大小可配置，适用于CPU内存受限的设备。
**关键特性：**
- **延迟加载**：模型权重按需从磁盘加载，避免一次性加载全部模型
- **智能缓存**：CPU内存缓冲区使用FIFO策略管理，可配置大小
- **多线程预取**：使用多个磁盘工作线程并行加载
- **异步传输**：使用CUDA流实现计算和数据传输的重叠
- **Swap轮换**：通过位置轮换实现连续计算，避免重复加载/卸载
**工作步骤**：
- **磁盘存储**：模型权重按block存储在SSD/NVMe上，每个block一个.safetensors文件
- **任务调度**：当需要某个block/phase时，优先级任务队列分配磁盘工作线程
- **异步加载**：多个磁盘线程并行从磁盘读取权重文件到CPU内存缓冲区
- **智能缓存**：CPU内存缓冲区使用FIFO策略管理缓存，可配置大小
- **缓存命中**：如果权重已在缓存中，直接传输到GPU，无需磁盘读取
- **预取传输**：缓存中的权重异步传输到GPU内存（使用GPU加载流）
- **计算执行**：GPU上的权重进行计算（使用计算流），同时后台继续预取下一个block/phase
- **Swap轮换**：计算完成后轮换block/phase位置，实现连续计算
- **内存管理**：当CPU缓存满时，自动淘汰最早使用的权重block/phase
## ⚙️ 配置参数
### GPU-CPU 卸载配置
```python
config = {
"cpu_offload": True,
"offload_ratio": 1.0, # 卸载比例（0.0-1.0）
"offload_granularity": "block", # 卸载粒度："block"或"phase"
"lazy_load": False, # 禁用延迟加载
}
```
### 磁盘-CPU-GPU 卸载配置
```python
config = {
"cpu_offload": True,
"lazy_load": True, # 启用延迟加载
"offload_ratio": 1.0, # 卸载比例
"offload_granularity": "phase", # 推荐使用phase粒度
"num_disk_workers": 2, # 磁盘工作线程数
"offload_to_disk": True, # 启用磁盘卸载
}
```
**智能缓存关键参数：**
- `max_memory`：控制CPU缓存大小，影响缓存命中率和内存使用
- `num_disk_workers`：控制磁盘加载线程数，影响预取速度
- `offload_granularity`：控制缓存粒度（block或phase），影响缓存效率
- `"block"`：以完整的Transformer层为单位进行缓存管理
- `"phase"`：以单个计算组件为单位进行缓存管理
**非 DIT 模型组件（T5、CLIP、VAE）的卸载配置：**
这些组件的卸载行为遵循以下规则：
- **默认行为**：如果没有单独指定，T5、CLIP、VAE 会跟随 `cpu_offload` 的设置
- **独立配置**：可以为每个组件单独设置卸载策略，实现精细控制
**配置示例**：
```json
{
"cpu_offload": true, // DIT 模型卸载开关
"t5_cpu_offload": false, // T5 编码器独立设置
"clip_cpu_offload": false, // CLIP 编码器独立设置
"vae_cpu_offload": false // VAE 编码器独立设置
}
```
在显存受限的设备上，建议采用渐进式卸载策略：
1. **第一步**：仅开启 `cpu_offload`，关闭 `t5_cpu_offload`、`clip_cpu_offload`、`vae_cpu_offload`
2. **第二步**：如果显存仍不足，逐步开启 T5、CLIP、VAE 的 CPU 卸载
3. **第三步**：如果显存仍然不够，考虑使用量化 + CPU 卸载或启用 `lazy_load`
**实践经验**：
- **RTX 4090 24GB + 14B 模型**：通常只需开启 `cpu_offload`，其他组件卸载需要手动设为 `false`，同时使用 FP8 量化版本
- **更小显存的 GPU**：需要组合使用量化、CPU 卸载和延迟加载
- **量化方案**：建议参考[量化技术文档](../method_tutorials/quantization.md)选择合适的量化策略
**配置文件参考**：
- **Wan2.1 系列模型**：参考 [offload 配置文件](https://github.com/ModelTC/lightx2v/tree/main/configs/offload)
- **Wan2.2 系列模型**：参考 [wan22 配置文件](https://github.com/ModelTC/lightx2v/tree/main/configs/wan22) 中以 `4090` 结尾的配置文件
## 🎯 使用建议
- 🔄 GPU-CPU分block/phase卸载：适合GPU显存不足(RTX 3090/4090 24G)但系统内存(>64/128G)充足
- 💾 磁盘-CPU-GPU分block/phase卸载：适合GPU显存(RTX 3060/4090 8G)和系统内存(16/32G)都不足
- 🚫 无Offload：适合高端硬件配置，追求最佳性能
## 🔍 故障排除
### 常见问题及解决方案
1. **磁盘I/O瓶颈**
- 解决方案：使用NVMe SSD，增加num_disk_workers
2. **内存缓冲区溢出**
- 解决方案：增加max_memory或减少num_disk_workers
3. **加载超时**
- 解决方案：检查磁盘性能，优化文件系统
**注意**：本卸载机制专为Lightx2v设计，充分利用了现代硬件的异步计算能力，能够显著降低大模型推理的硬件门槛。
