3.0 KiB
Booster 插件
作者: Hongxin Liu
前置教程:
引言
正如 Booster API 中提到的，我们可以使用 booster 插件来自定义并行训练。在本教程中，我们将介绍如何使用 booster 插件。
我们现在提供以下插件:
- Low Level Zero 插件: 它包装了
colossalai.zero.low_level.LowLevelZeroOptimizer，可用于使用 Zero-dp 训练模型。它仅支持 Zero 阶段1和阶段2。 - Gemini 插件: 它包装了 Gemini，Gemini 实现了基于Chunk内存管理和异构内存管理的 Zero-3。
- Torch DDP 插件: 它包装了
torch.nn.parallel.DistributedDataParallel并且可用于使用数据并行训练模型。 - Torch FSDP 插件: 它包装了
torch.distributed.fsdp.FullyShardedDataParallel并且可用于使用 Zero-dp 训练模型。
更多插件即将推出。
插件
Low Level Zero 插件
该插件实现了 Zero-1 和 Zero-2（使用/不使用 CPU 卸载），使用reduce和gather来同步梯度和权重。
Zero-1 可以看作是 Torch DDP 更好的替代品，内存效率更高，速度更快。它可以很容易地用于混合并行。
Zero-2 不支持局部梯度累积。如果您坚持使用，虽然可以积累梯度，但不能降低通信成本。也就是说，同时使用流水线并行和 Zero-2 并不是一个好主意。
{{ autodoc:colossalai.booster.plugin.LowLevelZeroPlugin }}
我们已经测试了一些主流模型的兼容性，可能不支持以下模型：
timm.models.convit_base- dlrm and deepfm models in
torchrec diffusers.VQModeltransformers.AlbertModeltransformers.AlbertForPreTrainingtransformers.BertModeltransformers.BertForPreTrainingtransformers.GPT2DoubleHeadsModel
兼容性问题将在未来修复。
⚠ 该插件现在只能加载自己保存的且具有相同进程数的优化器 Checkpoint。这将在未来得到解决。
Gemini 插件
这个插件实现了基于Chunk内存管理和异构内存管理的 Zero-3。它可以训练大型模型而不会损失太多速度。它也不支持局部梯度累积。更多详细信息，请参阅 Gemini 文档.
{{ autodoc:colossalai.booster.plugin.GeminiPlugin }}
Torch DDP 插件
更多详细信息，请参阅 Pytorch 文档.
{{ autodoc:colossalai.booster.plugin.TorchDDPPlugin }}
Torch FSDP 插件
⚠ 如果 torch 版本低于 1.12.0，此插件将不可用。
⚠ 该插件现在还不支持保存/加载分片的模型 checkpoint。
⚠ 该插件现在还不支持使用了multi params group的optimizer。
更多详细信息，请参阅 Pytorch 文档.
{{ autodoc:colossalai.booster.plugin.TorchFSDPPlugin }}