本文是大模型训练基础课程的笔记，主要对分布式训练策略、机器学习框架和GPU作一些简单的总结。

分布式训练¶

在BERT或更早一些的RNN、CNN时代，机器学习和深度学习对算力的要求还没有现在那么高，训练一个小的模型很少会遇到算力瓶颈，随着训练数据量的增加，我们逐渐开始使用数据并行（Data Parallelism）的方式，将训练数据切分为多份，每份数据交由一台设备（比如一张显卡）来训练，每个设备保存一份模型的副本，得到各个设备上的梯度之后，通过一种方式将所有设备上的梯度进行合并计算，然后再将新的模型参数发送到每个设备上，进行新一轮的迭代；数据并行常见的方式有Parameter server和Ring AllReduce两种，前者有一台master机器用于参数管理、分发、梯度同步等操作。而后者在PyTorch中的实现即是我们常用的DDP，它利用环状通信来降低数据传输的延迟。

随着模型参数量进一步增加，当单张卡无法再容纳一个完整的模型时，就需要按照模型的层进行切分了，如对于类GPT模型，我们可以将其Embedding层即后续的多个decoder层按照时序进行切分，这种并行方式即是流水线并行（Pipeline Parallelism），流水线并行中，需要先透过前向传播，逐层向后运算，运算完毕后再由后层向前层进行反向传播，这种方式会导致部分显卡在多数时间内都是空闲状态，另外，前序设备需要保存大量的中间激活值。GPipe和PipeDream是流水线并行的两种实现。

当流水线并行切分后，单层仍无法放在单卡上时，需要对层中的张量进行切分，如对于做矩阵乘的两个矩阵$AB$来说，可以对$A$按行切分，对$B$按列切分。即张量并行（Tensor Parallelism）。

在如今的大规模语言模型训练中，一般混合使用上述三种并行方式。

机器学习框架¶

计算图、自动微分技术和GPU算力规模化称得上机器学习框架发展的几大突破，深度学习发展到现在，模型的架构暂时趋于稳定，在机器学习框架层面做研发创新的不确定性变得更高，PyTorch、Tensorflow、PaddlePaddle等几个框架占据了市场的绝大部分。现今，无论是CV、NLP还是推荐领域，Transformer架构已经逐渐成为主流，类似DeepSpeed、megatron等框架基于Transformer做了很多深度的扩展和支持，比如算子融合等，这使得框架本身逐渐“边缘化”，能做的事情也越来越有限。

自动微分技术¶

自动微分技术可以使得算法工程师无需关心其使用的算子该如何完成反向传播梯度计算，在算子的定义阶段，就需要定义好一个算子的前向计算、反向计算方法、及backward时需要的变量。示例如下：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15

class MySquare(Function):

    @staticmethod
    def forward(x):
        return x ** 2

    @staticmethod
    def setup_context(ctx, inputs, output):
        x, = inputs
        ctx.save_for_backward(x)

    @staticmethod
    def backward(ctx, g0):
        x, = ctx.saved_tensors
        return g0 * 2 * x