ColossalAI是一个具有高效并行化技术的综合大规模模型训练系统。旨在无缝整合不同的并行化技术范式,包括数据并行、管道并行、多张量并行和序列并行。
Colossal-AI的目标是支持人工智能社区以与他们正常编写模型相同的方式编写分布式模型。这使得他们可以专注于开发模型架构,并将分布式训练的问题从开发过程中分离出来。
ColossalAI提供了一组并行训练组件。旨在支持用户编写分布式深度学习模型,就像编写单GPU模型一样。提供友好的工具,只需几行即可启动分布式培训。数据并行管道并行1D、2D、2.5D、3D和序列并行友好的trainer和engine可扩展以实现新的并行性混合精度训练零冗余优化器(ZeRO)importcolossalaifromcolossalai.engineimportEnginefromcolossalai.trainerimportTrainerfromcolossalai.coreimportglobal_contextasgpcmodel,train_dataloader,test_dataloader,criterion,optimizer,schedule,lr_scheduler=colossalai.initialize()engine=Engine(model=model,criterion=criterion,optimizer=optimizer,lr_scheduler=lr_scheduler,schedule=schedule)trainer=Trainer(engine=engine,hooks_cfg=gpc.config.hooks,verbose=True)trainer.fit(train_dataloader=train_dataloader,test_dataloader=test_dataloader,max_epochs=gpc.config.num_epochs,display_progress=True,test_interval=5)展示样例ViT14倍批大小和5倍训练速度(张量并行=64)GPT-3释放50%GPU资源占用,或10.7%加速GPT-2降低11倍GPU显存占用,或超线性扩展BERT2倍训练速度,或1.5倍序列长度
评论