AI Agent· 第五讲 多张GPU训练大语言模型

常规Pytorch训练

1
2
3
4
5
6
7
8
9

for epoch in range(num_epochs):
    for batch in dataloader:
        inputs, targets = batch
        outputs = model(inputs)
        loss = loss_fn(outputs, targets)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

LLM训练

• activations：模型某一层对于input的输出

• batch size：通常1个batch要有4-60M

• gradiaent accumulation：

  把batch size分为若干个mini batch size（GPU一次可以装的下的forward backward的大小），然后每一次loss保留，要保留gradient accumulation轮，接着把这些一起拿到global model里面backward

  global batch size = mini batch size * gradient accumulation

需要的内存

随着context增长指数级增长

多GPU训练问题

• 多张gpu一起跑：按照每张GPU大小切片除了inputs以外的内容。

  • DeepSpeed - Zero Redundancy Optimizer（ZeRO）

    

    • ZeRO-1：切片最大的optimizer

      

    • Zero-2：切片第二大的

      

    • Zero-3：所有的部分都要切片

      

    • Zero-offload：一部分工作转移到CPU上

      optimizer offload & optimizer offload+model offload

  长上下文（128K以上）问题

  1. Kernel（跑在GPU上面的方式）

     

     • 效果

       

  2. Flash Attention

     • 速度

       

     • Fused Kernel：把原来多个 GPU 计算步骤合并成一个 GPU kernel 一次性执行。

       

     • 记忆

       通过把大部分计算放到CPU上面，需要的时候调回的时候GPU来加快训练速度

       

  3. Liger Kernel

     • 用Triton code替换模型的一些部分来优化原本计算的所以对于

       Triton 是 OpenAI 开源的一个 GPU 编程工具。

     • 使用方法

       

     • 结果

       

  Quantization

  1. 压缩大小

     

     • GGML family
     • GPTQ
     • AWQ
     • BitsAndBytes