Flash Attention & 分布式训练

Flash Attention 优化点

• 通过将大矩阵 的计算拆分为小块来进行，每次只计算和存储一个小块的内容，每次计算都在 SRAM 中存放需要的部分数据，从而最大限度地减少对 HBM 的访问。

• 使用了分块并行计算的方法，在计算过程中实时计算 Softmax。

• 大部分中间结果只需要在较小的、延迟低的 SRAM 中缓存，而不必频繁访问 HBM，因此在处理长上下文时优势明显。

GPU中，SRAM 用于存储计算的中间结果和临时数据，访问速度非常快，HBM存储空间大，但访问延迟高。

分布式训练相关

deepspeed

• Stage 1：分割优化器状态 ZeRO 的第一个阶段将优化器状态（如 Adam 优化器中的一阶和二阶动量）分布到多个 GPU 上，减少每个 GPU 上的冗余内存占用。

• Stage 2：分割梯度 在此阶段，不仅优化器状态被分割，模型的梯度也会在多个设备之间分割和存储，进一步减少显存开销。

• Stage 3：分割参数 在最后的阶段，模型的所有参数也会被分布到各个 GPU 上，这样每个 GPU 只保存一部分参数，从而显著节省显存。这一阶段使得单个 GPU 能够训练数千亿参数的模型。

DP、DDP

• DP在多个GPU上复制模型，模型的参数存储在主GPU上，主GPU负责将模型参数广播到其他GPU上，多个GPU计算得到的梯度在主GPU上汇总，然后更新模型参数。仅支持单机多卡，不支持跨节点的分布式训练。

• DDP通过每个进程负责一个GPU的方式进行分布式训练，每个进程独立维护模型副本，并行计算梯度。梯度计算完成后，进程间通过全局通信（如all-reduce）同步梯度，然后各自更新模型参数。DDP不仅支持单机多卡，还支持跨多节点的分布式训练，非常适合大规模训练任务。