显卡|1块显卡＋几行代码：大模型训练提速40%！( 二 )

tensor_placement_policy="auto"),
optimizer_config=dict(gpu_margin_mem_ratio=0.8))
第二步，是在配置文件准备好后，插入几行代码来启动新功能。
首先，通过一行代码，使用配置文件来启动Colossal-AI 。
Colossal-AI会自动初始化分布式环境，读取相关配置，然后将配置里的功能自动注入到模型及优化器等组件中。
colossalai.launch_from_torch(config='./configs/colossalai_zero.py')
然后，还是像往常一样定义数据集、模型、优化器、损失函数等。
比如直接使用原生PyTorch代码，在定义模型时，只需将模型放置于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)
还是得靠GPU+CPU异构
而能够让用户实现如上“傻瓜式”操作的关键，还是AI系统本身要足够聪明。
发挥核心作用的是Colossal-AI系统的高效异构内存管理子系统Gemini 。
它就像是系统内的一个总管，在收集好计算所需的信息后，动态分配CPU、GPU的内存使用。
具体工作原理，就是在前面几个step进行预热，收集PyTorch动态计算图中的内存消耗信息。
在预热结束后，计算一个算子前，利用收集的内存使用记录，Gemini将预留出这个算子在计算设备上所需的峰值内存，并同时从GPU显存移动一些模型张量到CPU内存。