文本到图像生成样例“一个被猫绊倒的希腊人雕像”(左侧两列为Imagen,右侧两列为DALL·E 2)
tensor_placement_policy
决定了我们异构训练的策略,这个参数可以为
cuda
、
cpu
以及
auto
。各个策略有不同的优点:
cuda
: 将全部模型参数都放置于GPU上,适合不offload时仍然能进行训练的传统场景;
cpu
则会将模型参数都放置在CPU内存中,仅在GPU显存中保留当前参与计算的权重,适合超大模型的训练;
auto
则会根据实时的内存信息,自动决定保留在GPU显存中的参数量,这样能最大化利用GPU显存,同时减少CPU-GPU之间的数据传输。
auto
策略,由Colossal-AI自动化地实时动态选择最佳异构策略,最大化计算效率。
from colossalai.zero.shard_utils import TensorShardStrategy
zero = dict(model_config=dict(shard_strategy=TensorShardStrategy(),
tensor_placement_policy="auto"),
optimizer_config=dict(gpu_margin_mem_ratio=0.8))
colossalai.launch_from_torch(config='./configs/colossalai_zero.py')
ZeroInitContext
下初始化即可。在例子里,我们使用Hugging Face提供的OPTForCausalLM模型以及预训练权重,在Wikitext数据集上进行微调。
with ZeroInitContext(target_device=torch.cuda.current_device(),
shard_strategy=shard_strategy,
shard_param=True):
model = OPTForCausalLM.from_pretrained(
'facebook/opt-1.3b'
config=config
)
colossalai.initialize
,便可将配置文件里定义的异构内存功能统一注入到训练引擎中,即可启动相应功能。
engine, train_dataloader, eval_dataloader, lr_scheduler = colossalai.initialize(model=model,
optimizer=optimizer,
criterion=criterion,
train_dataloader=train_dataloader,
test_dataloader=eval_dataloader,
lr_scheduler=lr_scheduler)
-nprocs 8
即可实现!
parallel = dict(
pipeline=2,
tensor=dict(mode='2.5d', depth = 1, size=4)
)
|