江苏龙网

  1. 主页 > 生活百科 > >

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


在这里，使用的是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内存 。
显卡|1块显卡＋几行代码：大模型训练提速40%！
文章图片

Gemini内置的内存管理器给每个张量都标记一个状态信息，包括HOLD、COMPUTE、FREE等 。
然后，根据动态查询到的内存使用情况，不断动态转换张量状态、调整张量位置 。
带来的直接好处，就是能在硬件非常有限的情况下，最大化模型容量和平衡训练速度 。
要知道，业界主流方法ZeRO （Zero Reduency Optimizer），尽管也利用CPU+GPU异构内存的方法，但是由于是静态划分，还是会引起系统崩溃、不必要通信量等问题 。
而且，使用动态异构CPU+GPU内存的办法，还能用加内存条的办法来扩充内存 。
怎么也比买高端显卡划算多了 。
显卡|1块显卡＋几行代码：大模型训练提速40%！
文章图片

目前，使用Colossal-AI的方法，RTX 2060 6GB普通游戏本能训练15亿参数模型；RTX 3090 24GB主机直接单挑180亿参数大模型；Tesla V100 32GB连240亿参数都能拿下 。
除了最大化利用内存外，Colossal-AI还使用分布式并行的方法，让训练速度不断提升 。
它提出同时使用数据并行、流水并行、2.5维张量并行等复杂并行策略 。
方法虽复杂，但上手却还是非常“傻瓜操作”，只需简单声明，就能自动实现 。
无需像其他系统和框架侵入代码，手动处理复杂的底层逻辑 。
parallel = dict(
pipeline=2,
tensor=dict(mode='2.5d', depth = 1, size=4)


推荐阅读


上一篇：显卡|GPU-Z 2.47.0升级发布：假冒N卡无所遁形、优化国产显卡

下一篇：气质精神小伙经典语录有哪些？