执行摘要
本 PR 修复了 ModelOpt FP8 量化扩散模型中 layerwise offload 的兼容性问题，通过保留权重张量的非连续步幅并添加 32 字节对齐，确保 FP8 GEMM 内核能正确运行。变更允许用户在量化模型上启用 offload 以降低内存占用，同时更新了相关文档和测试。
功能与动机
之前，当用户在 ModelOpt FP8 量化检查点上启用 --dit-layerwise-offload true 时，会出现 "Misaligned Tensor data ... expected data alignment=32 bytes" 错误。这是因为 offload 管理器会扁平化权重，破坏了 CUTLASS 内核所需的列主序布局。修复后，用户可以在 FLUX.1-dev 和 Wan2.2 等模型上安全使用 layerwise offload，减少 GPU 内存峰值使用量。
实现拆解
- 核心 offload 管理：在
layerwise_offload.py 中，LayerwiseOffloadManager 新增了 _strided_cpu_weights 字典来存储保留原始步幅的权重，并引入了 _align_numel_offset 方法确保连续张量切片在缓冲区中 32 字节对齐。
- 适配器逻辑调整：
transformer_load_utils.py 中的 _maybe_disable_incompatible_dit_offload_modes 函数现在只禁用 dit_cpu_offload，而保持 dit_layerwise_offload 启用，更新警告信息以反映新行为。
- 文档与测试：更新
quantization.md 和技能文件 SKILL.md，新增单元测试 test_layerwise_offload.py 覆盖步幅保留、对齐和适配器逻辑。
评论区精华
- 性能优化：gemini-code-assist[bot] 指出，在 GPU 到 CPU 拷贝时，使用
.cpu() 会创建额外中间张量，建议直接使用 copy_ 提高效率。代码中已采纳此建议。
- 文档清理：mickqian 在文档中提出删除重复行的建议，BBuf 回应并进行了修改。
风险与影响
- 技术风险：对齐逻辑依赖于硬编码的 32 字节要求，若未来内核对齐需求变化可能引发错误；步幅保留增加了代码复杂性，可能在其他量化类型中引入未测试的边缘情况。
- 影响范围：用户现在可以在 ModelOpt FP8 量化模型上启用 layerwise offload，验证显示内存使用降低（如 FLUX.1-dev 的 peak allocated MB 从 18873.95 MB 减少）。系统性能因优化拷贝而略有提升，团队需更新相关技能和文档以确保一致性。
关联脉络
本 PR 是扩散模型量化支持演进的一部分。关联 PR 20082 "Enable modelopt quantized FLUX deployment" 引入了 ModelOpt 量化初始支持，而本 PR 解决了其 offload 兼容性问题。同时，PR 22182 "[diffusion] model: support LTX2.3 two stage" 展示了扩散模型功能的持续扩展，表明仓库正加强对多阶段生成和量化集成的投入。
参与讨论