bobo0810 / LearnColossalAI

ColossalAI教程 & 示例注释 & 学习笔记 （大模型高效训练）

大模型食用手册

• 该仓库收录于PytorchNetHub
• ColossalAI官网 文档

大模型方案

总结

• GPU总数量= 数据并行大小 × 张量并行大小 × 流水并行大小

  指定 张量并行、流水并行的大小，则自动推断数据并行大小

• Engine： 对模型、优化器、损失函数的封装类。

• fp16与ZeRO配置不兼容

注意

• 配置文件的batch 指单个进程的batch。 故总batch 是Nx batch
• 配置文件的epoch 指单个GPU上的epoch。 故总epoch 是Nx epoch

举例：每个gpu上的dataloader都是完整的数据集，未做拆分。 epoch=3 数据并行=2时模型实际上过了6遍数据集

实验

模型名称	参数量	输入尺寸	batch（数据并行->流水并行）	提升	备注
swinv2_large	196M	256×256	90->120	33%↑	SwinV2系列最大规格
convnext_xlarge	350M	224×224	140->200	43%↑	ConvNext系列最大规格
tf_efficientnet_l2	480M	800×800	2->4	100%↑	EfficientNet系列最大规格
vit_huge	632M	336×336	20->46	130%↑	ViT系列最大规格

更多详情

混合并行示例

• 数据并行+流水并行+Tensor并行 examples/hybrid

• 解决问题1：自定义拆分教程

• 解决问题2： 模型库

• 注：colossalai的官方API无此内容，属于隐藏的高级API。"补足算子"、"自定义多级拆分" 特性是基于读源码总结得出。

其他示例

Tensor并行

• 1D并行 examples/tp.py

流水并行

examples/pp.py

零冗余优化器ZeRO

examples/zero

异构内存空间管理器Gemini

examples/zero

自动并行AutoParallel(实验性)

api暂未稳定