视觉大模型实践案例

本工程用于探索如何在消费级显卡上成功训练相对大的视觉模型。

虽然视觉模型并没有像 LLM 那样有极其夸张的参数量，但是即使常用的以 Swin Large 为 backbone 的模型，都需要在 A100 上才能成功训练，这无疑阻碍了用户在视觉大模型上的探索和实验。因此本工程将探索在 3090 等 24G 甚至更小显存的消级显卡上如何训练视觉大模型。

本工程主要涉及到的训练技术有 FSDP、DeepSpeed 和 ColossalAI 等常用大模型训练技术。

本工程将不断更新完善，如果你有比较好的探索和意见，也非常欢迎提 PR

依赖

mmengine >=0.9.0 # 案例 1
deepspeed # 案例 2
fairscale # 案例 2

案例 1： 采用 8 张 24G 3090 显卡结合 FSDP 训练 dino-5scale_swin-l_fsdp_8xb2-12e_coco.py

cd mmdetection
./tools/dist_train.sh projects/example_largemodel/dino-5scale_swin-l_fsdp_8xb2-12e_coco.py 8
./tools/dist_train.sh projects/example_largemodel/dino-5scale_swin-l_fsdp_8xb2-12e_coco.py 8 --amp

• AMP: 混合精度训练
• GC: 梯度/激活值检查点
• FSDP: ZeRO-3 结合梯度检查点
• Iter Time: 一次迭代训练总时间

从上表可以看出：

1. 采用 FSDP 结合 AMP 和 GC 技术，可以将最初的 49G 显存降低为 8.5G，但是会增加 1.7 倍训练时间
2. 在目标检测视觉模型中，占据最大显存的是激活值，而不是优化器状态，这和 LLM 不同，因此用户应该首选梯度检查点，而不是 FSDP
3. 如果不开启梯度检查点，仅开启 FSDP 的话依然会 OOM，即使尝试了更加细致的参数切分策略
4. 虽然 AMP 可以减少不少显存，但是有些算法使用 AMP 会导致精度下降而 FSDP 不会

案例 2： 采用 8 张 24G 3090 显卡结合 DeepSpeed 训练 dino-5scale_swin-l_deepspeed_8xb2-12e_coco.py

cd mmdetection
./tools/dist_train.sh projects/example_largemodel/dino-5scale_swin-l_deepspeed_8xb2-12e_coco.py 8

很遗憾，到目前为止这依然是一个失败的案例，因为梯度始终会溢出导致精度很低。

从上表可以看出：

1. DeepSpeed 易用性上相比于 FSDP 有很大提升，因为梯度检查点可以用 torch 原生的而不需要修改特殊定制，同时也没有 auto_wrap_policy 这个需要用户自行设置的参数
2. DeepSpeed ZeRO 系列必须要采用 FP16 模式，其底层是采用了 NVIDIA’s Apex package, 其使用 Apex 的 AMP O2 模式，这导致需要修改代码，并且 O2 模式采用大量 FP16 计算导致 DINO 算法无法正常训练，但是它的这种模式可以显著节省显存，相比于 torch 官方的 AMP，类型转换更加彻底

从上述分析可知，如果 DeepSpeed 能够在不降低性能情况下成功训练 DINO 模型，那么其将比 FSDP 具备比较大的优势。如果您对 DeepSpeed 和 Apex 有比较深入的了解同时有兴趣排查精度问题，欢迎反馈或者提 PR

前面说过由于 Apex AMP O2 的特殊性，目前的 MMDetection 无法训练 DINO 模型，考虑到这是一个失败的案例，因此将修改的代码放在了 https://github.com/hhaAndroid/mmdetection/tree/dino_deepspeed 分支，其对应修改见 commit。如果您有兴趣尝试，可以拉取该分支进行试验。