恕我直言,你们的模型训练都还不够快

简介: 恕我直言,你们的模型训练都还不够快

周末在家没事干,也没人约了打游戏,于是打开了gayhub闲逛,哦不,是github。

然后发现了一个挺有意思的项目:

「也就是将你模型中的参数全部存储为一个连续的内存块,加速你的模型训练。」

于是我抱着试试看的心态,基于Fairseq和LightSeq分别实现了两个单层的Transformer编码层模型,简单写了一个例子试了一下。

安装

首先为了运行我这个例子,你需要安装上面提到的contiguous-params库。然后还需要安装fairseqlightseq库。

pip install contiguous-params fairseq lightseq

一个简单的例子

我这里创建了一个模型,就是单层的Transformer编码层,然后随机输入一个向量,损失函数就是输出向量的所有元素的平方均值。

然后测试了采用参数连续化前后,前向传播、反向传播、梯度更新三部分的时间消耗。

import time
from dataclasses import dataclass
import copy
import torch
from fairseq.modules.transformer_layer import TransformerEncoderLayer
from lightseq.training.ops.pytorch.transformer_encoder_layer import LSTransformerEncoderLayer
from contiguous_params import ContiguousParams
def get_time():
    '''CUDA同步并获取当前时间'''
    torch.cuda.synchronize(device="cuda:0")
    return time.time()
def ls_config_to_fs_args(config):
    '''将LightSeq的config转换为Fairseq的args'''
    @dataclass
    class Args:
        encoder_embed_dim: int
        encoder_ffn_embed_dim: int
        encoder_attention_heads: int
        dropout: float
        attention_dropout: float
        activation_dropout: float
        encoder_normalize_before: bool
    args = Args(
        config.hidden_size,
        config.intermediate_size,
        config.nhead,
        config.hidden_dropout_ratio,
        config.attn_prob_dropout_ratio,
        config.activation_dropout_ratio,
        config.pre_layer_norm
    )
    return args
def train(model, inputs, masks, contiguous=False):
    '''训练过程'''
    model.to(device="cuda:0")
    model.train()
    if contiguous:
        parameters = ContiguousParams(model.parameters())
        opt = torch.optim.Adam(parameters.contiguous(), lr=1e-3)
    else:
        opt = torch.optim.Adam(model.parameters(), lr=1e-3)
    fw_time, bw_time, step_time = 0, 0, 0
    for epoch in range(1000):
        opt.zero_grad()
        start_time = get_time()
        outputs = model(inputs, masks)
        loss = torch.square(outputs).mean()
        fw_time += get_time() - start_time
        start_time = get_time()
        loss.backward()
        bw_time += get_time() - start_time
        start_time = get_time()
        opt.step()
        step_time += get_time() - start_time
        if epoch % 200 == 0:
            print("epoch {:>3d}: loss = {:>5.3f}".format(epoch, loss))
    return fw_time, bw_time, step_time
if __name__ == "__main__":
    # 定义LightSeq的config
    config = LSTransformerEncoderLayer.get_config(
        max_batch_tokens=4096,
        max_seq_len=256,
        hidden_size=128,
        intermediate_size=512,
        nhead=16,
        attn_prob_dropout_ratio=0.1,
        activation_dropout_ratio=0.1,
        hidden_dropout_ratio=0.1,
        pre_layer_norm=True,
        fp16=False,
        local_rank=0
    )
    # 将LightSeq的config转换为Fairseq的args
    args = ls_config_to_fs_args(config)
    # 随机生成输入
    bsz, sl = 50, 80
    inputs = torch.randn(bsz, sl, config.hidden_size).to(device="cuda:0")
    masks = torch.zeros(bsz, sl).to(device="cuda:0")
    # 定义LightSeq模型并训练
    ls_model = LSTransformerEncoderLayer(config)
    ls_fw_time, ls_bw_time, ls_step_time = train(ls_model, inputs, masks)
    # 定义连续化参数的LightSeq模型并训练
    config_cont = copy.deepcopy(config)
    ls_model_cont = LSTransformerEncoderLayer(config_cont)
    ls_c_fw_time, ls_c_bw_time, ls_c_step_time = train(ls_model_cont, inputs, masks, contiguous=True)
    inputs = inputs.transpose(0, 1)
    masks = masks > 0.5
    # 定义Fairseq模型并训练
    fs_model = TransformerEncoderLayer(args)
    fs_fw_time, fs_bw_time, fs_step_time = train(fs_model, inputs, masks)
    # 定义连续化参数的Fairseq模型并训练
    fs_model_cont = TransformerEncoderLayer(args)
    fs_c_fw_time, fs_c_bw_time, fs_c_step_time = train(fs_model_cont, inputs, masks, contiguous=True)
    print("LightSeq time:         {:.3f}s, {:.3f}s, {:.3f}s".format(ls_fw_time, ls_bw_time, ls_step_time))
    print("LightSeq (cont) time:  {:.3f}s, {:.3f}s, {:.3f}s".format(ls_c_fw_time, ls_c_bw_time, ls_c_step_time))
    print("Fairseq time:          {:.3f}s, {:.3f}s, {:.3f}s".format(fs_fw_time, fs_bw_time, fs_step_time))
    print("Fairseq (cont) time:   {:.3f}s, {:.3f}s, {:.3f}s".format(fs_c_fw_time, fs_c_bw_time, fs_c_step_time))

详细讲解

这里最主要的地方就两行:

parameters = ContiguousParams(model.parameters())
opt = torch.optim.Adam(parameters.contiguous(), lr=1e-3)

首先用ContiguousParams类封装model.parameters(),然后将封装后的parameters.contiguous()送进优化器中,这里送进去的就已经是连续存储的一整块参数了。

我们详细阅读ContiguousParams的源码,可以发现实现很简单:https://github.com/PhilJd/contiguous_pytorch_params/blob/master/contiguous_params/params.py

核心代码就是下面这个函数,注释中我都详细解释了每一步在干嘛:

def make_params_contiguous(self):
    index = 0
    # 遍历所有的参数
    for p in self._parameters:
        # 计算参数p的大小
        size = p.numel()
        # 在连续参数块中的对应位置赋值参数p
        self._param_buffer[index:index + size] = p.data.view(-1)
        # 将参数p的数值和梯度都重新指向连续参数块和连续梯度块的对应位置
        p.data = self._param_buffer[index:index + size].view(p.data.shape)
        p.grad = self._grad_buffer[index:index + size].view(p.data.shape)
        # 连续内存块位置偏移到下一个参数
        index += size
    # 连续参数块的梯度设置为连续梯度块
    self._param_buffer.grad = self._grad_buffer

所以在封装了原始参数之后,之后模型计算就会从连续内存块中对应位置取出数值,然后进行计算。

运行结果

我在V100显卡上运行了一下上面的例子,结果如下:

可以看出,LightSeq在采用参数连续化前后,三部分运行时间几乎没有任何变化,这主要是由于LightSeq已经在模型内部做过参数连续化了,因此速度已经很快了。

而Fairseq前后的第三部分,也就是参数更新部分时间缩减非常多,从1.5秒缩短到了0.1秒,总的训练时间几乎缩短了将近一半。

最后对比LightSeq和Fairseq可以明显发现,LightSeq的训练时间比Fairseq快非常多。主要是因为LightSeq采用了算子融合等各种技术,加速了Transformer模型的训练。

总结

所以在你的「任意」PyTorch模型中,都可以用上面的参数连续化技术大大加快训练速度。

而如果你的模型是Transformer类模型,那还可以直接用字节跳动开源的LightSeq训练加速引擎,更加方便。

如果你是TensorFlow爱好者,还可以直接用字节跳动开源的NeurST序列生成库进行训练,里面直接集成了LightSeq,所以训练很快。

目录
打赏
0
0
0
0
14
分享
相关文章
论文介绍:训练计算最优的大型语言模型
【2月更文挑战第30天】研究人员发现,在有限计算资源下,优化大型语言模型的训练需使模型大小和训练数据量成比例增长,以达到计算最优。通过训练700亿参数的Chinchilla模型并对比GPT-3等,验证了该策略的有效性。论文强调数据集质量和伦理隐私问题,并提出预测模型扩展的方法。这一发现对AI领域的模型训练策略提供了新思路,但也面临数据质量和伦理挑战。
86 2
论文介绍:训练计算最优的大型语言模型
【科普向】模型蒸馏和模型量化到底是什么???
在数字化快速发展的时代,人工智能(AI)技术已广泛应用,但大型深度学习模型对计算资源的需求日益增长,增加了部署成本并限制了其在资源有限环境下的应用。为此,研究人员提出了模型蒸馏和模型量化两种关键技术。 模型蒸馏通过将大型教师模型的知识传递给小型学生模型,利用软标签指导训练,使学生模型在保持较高准确性的同时显著减少计算需求,特别适用于移动设备和嵌入式系统。 模型量化则是通过降低模型权重的精度(如从32位浮点数到8位整数),大幅减少模型大小和计算量,提高运行速度,并能更好地适应低配置设备。量化分为后训练量化和量化感知训练等多种方法,各有优劣。
|
5月前
|
两个小模型互相验证,直接比肩大模型?微软的rStar甚至没用CoT和微调
【9月更文挑战第10天】微软研究院亚洲院与哈佛大学研究团队提出了rStar,一种创新的方法,旨在提升小型语言模型(SLMs)的推理能力。rStar采用自我对弈的相互生成-判别过程,利用增强版蒙特卡洛树搜索(MCTS)算法生成高质量推理轨迹,并由另一个相似能力的SLM验证这些轨迹的一致性。这种方法在多个模型与任务中显著提升了推理准确率,如在GSM8K上,LLaMA2-7B的准确率从12.51%跃升至63.91%。rStar的独特之处在于无需微调或依赖更强大的模型即可提升推理能力。然而,其计算成本和对SLM自身能力的依赖仍是挑战。
208 7
公理训练让LLM学会因果推理:6700万参数模型比肩万亿参数级GPT-4
【8月更文挑战第3天】新论文提出“公理训练”法,使仅有6700万参数的语言模型掌握因果推理,性能媲美万亿级GPT-4。研究通过大量合成数据示例教授模型因果公理,实现有效推理并泛化至复杂图结构。尽管面临合成数据需求大及复杂关系处理限制,此法仍为语言模型的因果理解开辟新途径。[链接: https://arxiv.org/pdf/2407.07612]
109 1
|
6月前
长上下文能力只是吹牛?最强GPT-4o正确率仅55.8%,开源模型不如瞎蒙
【8月更文挑战第10天】新研究NoCha挑战显示,即使是顶级的大型语言模型GPT-4o,在处理长篇幅文本时正确率仅55.8%,低于人类直观水平。该挑战基于近作英文小说,检验模型对整本书信息的理解与推理能力。结果显示,模型在全局推理上的表现不佳,倾向于依赖局部信息而非整体上下文,尤其是在复杂推理需求高的科幻小说上表现更弱。这一发现揭示了当前模型在处理长上下文任务上的局限性。论文链接: [https://arxiv.org/pdf/2406.16264](https://arxiv.org/pdf/2406.16264)。
154 65
训练Sora模型,你可能需要这些(开源代码,模型,数据集及算力评估)
在之前的文章《复刻Sora有多难?一张图带你读懂Sora的技术路径》,《一文看Sora技术推演》我们总结了Sora模型上用到的一些核心技术和论文,今天这篇文章我们将整理和总结现有的一些开源代码、模型、数据集,以及初步训练的算力评估,希望可以帮助到国内的创业公司和个人开发者展开更深的研究。
可解释性研究新突破:OpenAI成功训练1600万个特征的自动编码器
【6月更文挑战第13天】OpenAI团队在可解释性研究上取得进展,训练出拥有1600万特征的自动编码器来解析GPT-4。此模型旨在揭示语言模型的工作原理,提高AI透明度。自动编码器从低维度特征空间重建输入数据,研究通过稀疏特征增强可解释性。虽然规模扩大带来解释性提升,但计算资源需求大,且评估指标的全面性仍受质疑。[论文链接](https://cdn.openai.com/papers/sparse-autoencoders.pdf)
100 1
论文介绍:InfLLM——揭示大型语言模型在无需训练的情况下处理极长序列的内在能力
【5月更文挑战第18天】InfLLM是一种新方法,无需额外训练即可增强大型语言模型处理极长序列的能力。通过使用记忆单元存储长序列的远距离上下文,InfLLM能更准确地捕捉长距离依赖,提高对长文本理解。实验表明,InfLLM使预训练在短序列上的模型在处理极长序列时表现媲美甚至超过专门训练的模型。尽管有挑战,如动态上下文分割和记忆单元效率,InfLLM为长序列处理提供了有效且未经训练的解决方案。论文链接:https://arxiv.org/abs/2402.04617
224 3
论文介绍:自我对弈微调——将弱语言模型转化为强语言模型的新方法
【5月更文挑战第17天】论文《自我对弈微调》提出了一种新方法,名为SPIN,用于在无需额外人工标注数据的情况下增强大型语言模型(LLM)。SPIN利用自我对弈机制,让模型通过与自身历史版本交互生成自我训练数据,实现性能提升。该方法在多个基准数据集上表现出色,超越了传统监督微调和直接偏好优化。SPIN还为生成对抗网络研究提供了新思路,展示了自我对弈在强化学习和深度学习中的潜力。实验表明,SPIN有效提升了模型性能,为未来研究奠定了基础。[[arxiv](https://arxiv.org/abs/2401.01335v1)]
93 3
【大模型】小样本学习的概念及其在微调 LLM 中的应用
【5月更文挑战第5天】【大模型】小样本学习的概念及其在微调 LLM 中的应用

热门文章

最新文章