1.3.llm_archs

训练框架

mfu(Model Flops Utilization)模型算力利用率是分布式训练效率的优化目标。

一个模型定义好之后，前向和反向的计算量就是固定（不考虑动态图的话）的，除以每个step的latency就是mfu。以nanoGPT中的代码为例：

def estimate_mfu(self, fwdbwd_per_iter, dt):
    """ estimate model flops utilization (MFU) in units of A100 bfloat16 peak FLOPS """
    # first estimate the number of flops we do per iteration.
    # see PaLM paper Appendix B as ref: https://arxiv.org/abs/2204.02311
    N = self.get_num_params()
    cfg = self.config
    L, H, Q, T = cfg.n_layer, cfg.n_head, cfg.n_embd//cfg.n_head, cfg.block_size
    flops_per_token = 6*N + 12*L*H*Q*T
    flops_per_fwdbwd = flops_per_token * T # 计算量（T是序列长度）
    flops_per_iter = flops_per_fwdbwd * fwdbwd_per_iter # 每个step的flops * 每一次更新梯度要多少个step
    # express our flops throughput as ratio of A100 bfloat16 peak flops
    flops_achieved = flops_per_iter * (1.0/dt) # per second 一轮的计算量/一轮的耗时
    flops_promised = 312e12 # A100 GPU bfloat16 peak flops is 312 TFLOPS
    mfu = flops_achieved / flops_promised
    return mfu

##...
def xxx():
    # timing and logging
    t1 = time.time()
    dt = t1 - t0 ## 一次gradient_accumulation_steps后 更新梯度
    t0 = t1
    if iter_num % log_interval == 0 and master_process:
        # get loss as float. note: this is a CPU-GPU sync point
        # scale up to undo the division above, approximating the true total loss
        # (exact would have been a sum)
        lossf = loss.item() * gradient_accumulation_steps
        if local_iter_num >= 5: # let the training loop settle a bit
            mfu = raw_model.estimate_mfu(batch_size * gradient_accumulation_steps, dt)
            running_mfu = mfu if running_mfu == -1.0 else 0.9*running_mfu + 0.1*mfu
        print(f"iter {iter_num}: loss {lossf:.4f}, time {dt*1000:.2f}ms, mfu {running_mfu*100:.2f}%")
    iter_num += 1
    local_iter_num += 1

一般分布式训练参数量越多->卡数越多->通信占比越高->MFU越低，所以要优化通信效率。

优化设置

batchsize：通常用比较大的batchsize，提高训练稳定性和吞吐量。GPT-3和PaLM在训练时动态增加batshzie，最终达到百万级别，batchsize从3.2w逐渐增加到320w个token。
优化器：
学习率：
- 衰减（decay）：后续steps中余弦衰减，逐渐降低到最大值的约10%，直到收敛
稳定训练：
- 权重衰减和梯度裁剪：权重衰减率设为0.1，梯度裁剪阈值设为1.0
- 梯度检查点：容易出现loss突增，PaLM和OPT从发生突增前的一个ckpt重新开始训练，并跳过可能有问题的数据
- 缩减emb梯度：GLM发现emb的异常梯度通常会导致loss突增，故缩减emb梯度以缓解

混合精度训练

FP16

Mixed precision training提出了用16位float（FP16）训练，减少内存使用和通信开销。A100等GPU具有的FP16计算单元是FP32的两倍，故FP16的计算效率能进一步提高。

推理（预测）：所有参数都是fp16，相对fp32，存储变成一半，速度提升1倍。
训练：参数和梯度用fp32存储，但是在计算前会转成fp16，计算后再转回fp32。主要为了防止溢出，loss要乘一个scale，然后在fp16的梯度要除以scale。

以adam优化器为例，对于每1个参数来说，fp16的训练占用20bytes显存，包括（详见：https://zhuanlan.zhihu.com/p/519264636）

fp16的参数：2bytes
fp16的梯度：2bytes
优化器状态（optimizer state）：16bytes
- fp32参数（4bytes）
- fp32梯度（4bytes）(zero论文里提到的，用于reduce之类操作时需要的fp32内存，以1.5B的gpt2为例，需要6GB内存，倒推回来，就需要6/1.5byte=4byte)
- fp32 variance【历史梯度平方和】（4bytes）
- fp32 momentum【历史梯度滑动平均】（4bytes）

BF16

FP16可能导致计算精度的损失从而影响模型性能，BLOOM里用BF16(brain floating point)比FP16分配更多指数位和更少的有效位，在准确性方面更好

https://blog.csdn.net/orangerfun/article/details/133106913

bf16的指数位和fp32一样多

可扩展的训练

需要提高训练吞吐量和加载更大模型到显存中

3D并行

如下三种并行（数据并行、流水线并行、张量并行）的组合

数据并行（Data Parallelism）

将模型参数和优化器状态复制到多个GPU上，每个GPU只处理分给它的数据，不同GPU算出的梯度进行聚合得到batch的梯度，再更新所有GPU上的模型。高度可扩展，增加GPU数就能提高训练吞吐。

torch的ddp

from torch.nn.parallel import DistributedDataParallel as DDP

ddp = int(os.environ.get('RANK', -1)) != -1 # is this a ddp run?
if ddp:
    init_process_group(backend=backend)
    ddp_rank = int(os.environ['RANK'])
    ddp_local_rank = int(os.environ['LOCAL_RANK'])
    device = f'cuda:{ddp_local_rank}'
    torch.cuda.set_device(device)
    # this process will do logging, checkpointing etc.
    master_process = ddp_rank == 0
    seed_offset = ddp_rank # each process gets a different seed

# wrap model into DDP container
if ddp:
    model = DDP(model, device_ids=[ddp_local_rank])

可以一起搞的技巧——梯度累积，当显存不够跑较大的batchsize时，训练效果可能会很差，可以先跑多个mini-batch的前向和反向，把梯度累积起来，再更新一次参数，在数学上等价于跑一个较大的batchsize。

# forward backward update, with optional gradient accumulation to simulate larger batch size
# and using the GradScaler if data type is float16
for micro_step in range(gradient_accumulation_steps):
    if ddp:
        # in DDP training we only need to sync gradients at the last micro step.
        # 最后一个micro step才要sync梯度
        model.require_backward_grad_sync = (micro_step == gradient_accumulation_steps - 1)
    with ctx:
        logits, loss = model(X, Y)
    loss.backward() # 只是计算梯度，并不真的更新
optimizer.step()
optimizer.zero_grad(set_to_none=True)

也可以用torch的no_sync()：

ddp = torch.nn.parallel.DistributedDataParallel(model, pg)
with ddp.no_sync():
    for input in inputs:
        ddp(input).backward()  # no synchronization, accumulate grads
ddp(another_input).backward()  # synchronize grads

流水线并行（Pipeline Parallelism）

将LLM的不同层分配到多个GPU上，一般Transformer模型中会将连续的层加载到同一GPU上，以减少在GPU间传输已计算的隐层状态或梯度的成本。简单的实现会导致GPU利用率降低，因为每个GPU要等前一个完成计算，导致不必要的气泡开销，如下方法可以提高流水线效率：

GPipe：Gpipe: Efficient training of giant neural networks using pipeline parallelism
PipeDream：PipeDream: Fast and Efficient Pipeline Parallel DNN Training，填充多个数据batch+异步梯度更新？看下paper先。。。

1) GPipe

Gpipe主要思想：

图a：把模型不同layers顺序放在4张卡上，0->3卡流水线前向计算loss，3->0再反向计算gradients
图b：从时间顺序上看，每张卡有3/4时间是空闲的，GPU利用率非常低
图c：配合梯度累积，多个mini-batch可以同时跑在流水线里面，每张卡则有3/(3+4)的时间空闲（Bubble）

GPT里用Weight Tying提升效果，输入和输出共享vocab embedding

2) 重计算

即gradient checkpointing，重计算(recomputation)是对于pipeline parallelism非常重要的一个优化，最开始在Training Deep Nets with Sublinear Memory Cost一文中提到，在flash attention中也用了。

因为要做pipeline+梯度累积，前向过程中的激活值要保存，以留给反向过程使用，保存很多份的激活值对显存造成了很大压力。recomputation(也叫checkpointing)用时间来换空间（反向的时候进行一次激活值的重计算），可以缓解显存问题。

pytorch的实现。大致逻辑是包了一个autograd.Function，前向时保存一些inputs/rng_state(RNG state是Random Number Generator state的缩写，随机数生成器的状态。在深度学习和其他计算任务中，随机数生成器用于初始化参数、决定正则化技术如dropout的行为，以及在训练过程中选择样本等。RNG状态是指随机数生成器当前的内部状态，它可以用来在需要时重现或恢复特定的随机数序列，确保实验或模型训练的可重复性)，反向时重新计算

深度更深的时候，一般效果更好

cpu offload：层数很深的时候，可能可以把一些计算挪到cpu上去，再搞回gpu

张量并行（Tensor Parallelism）

Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism

Megatron-LM：能扩展到更高维度的张量
Colossal-AI：
- 为更高维度的张量实现了张量并行，An efficient 2d method for training super-large deep learning models、Tesseract: Parallelize the tensor parallelism efficiently和Maximizing Parallelism in Distributed Training for Huge Neural Networks
- 特别针对序列数据提出序列并行(Sequence Parallelism: Long Sequence Training from System Perspective)，详见下一节

参考https://zhuanlan.zhihu.com/p/622036840

原始矩阵乘法是[m,k], [k, n] -> [m, n]，有如下两种矩阵分解的等效：

- [m,k], [k, n/2] -> [m, n/2]
- concat([m, n/2], [m, n/2]) -> [m, n]
- [m, k/2], [k/2, n] -> [m, n]
- elemwise_add([m, n], [m, n]) -> [m, n]

行并行还可以扩展到推荐里，假设user有k/2维，item也是k/2维，concat在一起，然后过一个k*d的mlp，即[1,k] * [k, d] -->[1,d]，那么可以按行并行的方法，拆成2个[1, k/2]和[k/2,d]相乘，再相加。这样item侧的[k/2,d]可以把全库缓存过来，在线实时算user，排序时把对应item向量抽出来，和user加起来就行

megatron对transformer进行了如下优化：

MLP第一个nn按列分割，第二个nn按行分割，中间省了一次通信
Attention按照head来分割(类似列分割)，后面接的nn按行分割，中间也省了一次通信

图里面的通信算子

综合来看，一层transformer layer如下

具体的计算量可以参考https://colossalai.org/docs/features/1D_tensor_parallel/#introduction：

ZeRO

ZeRO: Memory Optimization Towards Training A Trillion Parameter Models

fp16那一节中，optimizer state的显存占用，在前向和反向的时候都不用，只有最后optimizer step的时候才用。

===>zero的思想：把optimizer state分shard存在不同的卡上，只在最后gather时才用。

ZeRO（Zero Redundancy Optimizer）在DeepSpeed库中提出，解决数据并行中的内存冗余问题。数据并行其实并不需要每个GPU都存整个模型、梯度和优化器参数，ZeRO在每个GPU仅保存部分数据，当需要其余数据时从其他GPU检索。3种解决方案：

优化器状态分区：zero1，对显存最大开销的部分进行shard
梯度分区：zero2
参数分区：zero3

前两种方案不会增加通信开销，第三种方案增加约50%通信开销，但能节省和gpu数成比例的内存。

详见官方博客：ZeRO & DeepSpeed: New system optimizations enable training models with over 100 billion parameters

import deepspeed
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", type=int, default=0)
deepspeed.add_config_arguments(parser)
args = parser.parse_args()

model, optimizer, _, _ = deepspeed.initialize(args=args,
                                              model=model,
                                              model_parameters=model.parameters())
X, Y = get_batch('train')
logits, loss = model(X, Y)
model.backward(loss)
model.step()

需要指定deepspeed的配置：

{
  "train_batch_size": 64,
  "gradient_accumulation_steps": 1,
  "optimizer": {
    "type": "Adam",
    "params": {
      "lr": 6e-4,
      "weight_decay": 1e-2,
      "betas": [0.9, 0.95]
    }zer
  },
  "scheduler": {
    "type": "WarmupLR",
    "params": {
        "warmup_min_lr": 6e-5,
        "warmup_max_lr": 6e-4,
        "warmup_num_steps": 2000
    }
  },
  "bf16": {
    "enabled": true
  },
  "zero_optimization": {
    "stage": 1
  }
}

启动：

deepspeed --num_gpus=8 train.py --deepspeed_config xx.json

facebook的开源库FSDP(full sharded data parallel)(Fairscale: A general purpose modular pytorch library for high performance and large scale training)里基于pytorch实现了类似ZeRO的技术。

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP, ShardingStrategy
model = FSDP(model)  #, sharding_strategy=ShardingStrategy.SHARD_GRAD_OP)

还有一些paper也能降低内存，如

Reducing activation recomputation in large transformer models

Training deep nets with sublinear memory cost

序列并行

序列并行(Sequence Parallelism: Long Sequence Training from System Perspective)，可以进一步分解Transformer的注意力操作。

Reducing Activation Recomputation in Large Transformer Models这个也是

对比TP：

SP：

综合对比各种并行

几个缩写：params（p）/gradients(g)/optimizer states(os)/activation(a)

整体对比可以看

一般这么整合：

把机器分成N组，不同组之间用DP
一组机器有M台机器，不同机器之间用PP
一台机器有K张卡，不同卡之间用TP

编译优化

pytorch的TorchDynamo

https://pytorch.org/docs/stable/torch.compiler_deepdive.html

最简单的用法torch.compile()

flash attention

Flashattention: Fast and memory-efficient exact attention with io-awareness

减少Global Memory的大小
加速attenttion的计算，因为读cache比访问Global Memory快多了。

flexattention

新PyTorch API：几行代码实现不同注意力变体，兼具FlashAttention性能和PyTorch灵活性

训练稳定性

学习率+batchsize的一些经验：

https://zhuanlan.zhihu.com/p/64864995

另外，swiglu会让act_norm变大，得加一些方式让模型能训动，一种方法是把weight_decay调小：

loss会变高==>https://poe.com/s/Cv6lODy94INz1ozQKJYJ，正则项变大，L+norm中的L相对影响变小了，所以会变大，即为了防止过拟合，但可能eval的时候会更好
act_norm会降低20-30%，因为正则项更重要了，会让权重更接近0

层数加深时的稳定性调优：Transformers Get Stable: An End-to-End Signal Propagation Theory for Language Models

数据/训练策略

综述

A Survey on Data Selection for Language Models

高质量数据

Llama架构比不上GPT2？神奇token提升10倍记忆？

Physics of Language Models: Part 3.3, Knowledge Capacity Scaling Laws

制造人工合成数据，通过控制数据中知识的数量和类型，来严格调控数据中的知识比特数 (bits)。使用不同大小和架构的 LLM 在人工合成数据上进行训练，并给出数学定理，来精确计算训练好的模型从数据中学到了多少比特的知识。有如下几个发现：

如果训练时间充足，不论使用何种模型架构，模型的存储效率均可以达到2bit/param（即平均每个模型参数可以存储2比特的信息）。而且发现transformer中的知识并非主要存储在MLP层，因为即便移除所有MLP层，模型仍能达到 2bit/param 的存储效率。
如果训练时间不充足，GPT2模型能比LlaMA/Mistral存储超过30%的知识，主要是因为GatedMLP(MoE)会导致训练不稳定，因此对同样的知识，需要更长的训练时间。
压缩/量化的影响：将训练好的模型从float32/16压缩到int8，对知识的存储毫无影响。LLM可以达到“信息论极限”的1/4——因为int8只有8比特，但平均每个参数可以存储2比特的知识。
高质量数据的影响：如果我们的预训练数据中，有1/8来自高质量知识库（如百度百科），7/8来自低质量数据（如common crawl或论坛对话，甚至是完全随机的垃圾数据），会发现：
- 即使对高质量数据的训练时间保持一致，低质量数据的存在本身可能会让模型对高质量知识的存储量下降20倍，即便将高质量数据的训练时间延长 3倍，知识储量仍会降低3倍
- 解法：只需给所有的预训练数据加上自己的网站域名token即可，模型的知识存储量可以立即回升10倍，模型不需要任何先验知识来识别哪些网站上的知识是金子，而可以在预训练过程中，自动发现高质量知识的网站，并自动为这些高质量数据腾出存储空间。

DSIR

Data Selection for Language Models via Importance Resampling

https://github.com/p-lambda/dsir

大致思路：

根据权重进行无放回的采样，兼顾多样性和分布一致性

DoReMi

DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining NeurIPS2023的spotlight

https://github.com/sangmichaelxie/doremi

用reference mnodel来指导一个小的proxy model的训练，proxy model使用group DRO（group distributionally robust optimization）来得到domain weights，即让最大的loss gap最小化

使用tune完的domain weights来训练大模型
可以重复这个过程，用新的domain weights重新训练ref_model，迭代下去

dataset decomposition

Dataset Decomposition: Faster LLM Training with Variable Sequence Length Curriculum

LLM的训练语料大都是相同长度的sequence，一般是多篇文档concat到一起，然后再切分成等长的sequence，有可能一个sequence里有不同的毫不相关的文档，这样算attention就不太适合了。方法

训练时可以给定一个固定的batch_len，即直接从某个桶里采样，也可以用特定长度策略，从多个桶里采样，组合为特定长度

RHO-1

RHO-1: Not All Tokens Are What You Need

https://github.com/microsoft/rho

把语料中的无用token从loss里删了

训练ref_model，挑选少量高质量语料，建模语料的整体loss分布情况
拿ref_model在整个预训练语料上计算每个token的ppl
训练LLM，只关注得分比较高的tokens

infinit lr

Simple and Scalable Strategies to Continually Pre-train Large Language Models

linear warmup and cosine decay schedule：

infinit lr decay：

linear warmup阶段：同上
常数阶段：学习率保持为这个常数，该阶段结束时的ckpt可以用于新数据集的继续pretrain
annealing阶段：逐渐减小到最小值

JEST

DeepMind新方法：训练时间减少13倍，算力降低90%

Data curation via joint example selection further accelerates multimodal learning

现有的大规模预训练数据筛选方法速度慢、成本高，并且没有考虑到批次组成或训练过程中数据相关性的变化，这限制了多模态学习中的效率提升。因此，DeepMind团队研究了联合选择数据批次而非单个样本是否能够加速多模态学习。

挑选好的数据批次比单独挑选数据点更为有效
在线模型近似可用于更高效地过滤数据
可以引导小型高质量数据集以利用更大的非精选数据集

JEST能够在仅使用10%的FLOP预算的情况下超越之前的最先进水平。

硬件的可能影响

https://zhuanlan.zhihu.com/p/701623664?utm_psn=1784500156948938753

大模型训练（如InternLM-7B）实践中，曾经遇到过在A100集群上表现正常的代码和数据，迁移到A800集群却出现了模型准确度下降和梯度范数爆炸的问题。经过调查，我们发现这与A800和A100 GPU的NVLink带宽差异有关。通过在两个集群上使用nanoGPT模型进行的对照实验，我们确认了精度差异的原因在于NCCL的Ring all-reduce算法实现。进一步实验表明，设置环境变量NCCL_ALGO=Tree或使用gloo作为backend可以解决精度对齐问题。最终，我们提出了一个解决方案：在A800集群上设置NCCL_ALGO=Tree，强制使用Tree算法进行all-reduce操作，从而避免了Ring算法带来的精度问题，使得A800集群的模型能够正常收敛，并且与A100集群的训练精度对齐。

fastpersist

DeepSpeed 最新力作：大模型CKPT速度提升116倍

FastPersist: Accelerating Model Checkpointing in Deep Learning

pathways

Pathways: Asynchronous Distributed Dataflow for ML

下载了，pdf

这个回答分析得不错 https://www.zhihu.com/question/524596983/answer/2420225275

Google的大规模稀疏模型设计

DESIGNING EFFECTIVE SPARSE EXPERT MODELS

代码：https://github.com/tensorflow/mesh/blob/master/mesh_tensorflow/transformer/moe.py

Google的几大LLM加速工具

maxdiffusion

https://github.com/google/maxdiffusion

JetStream

https://github.com/google/JetStream

maxtext

https://github.com/google/maxtext

Fire-Flyer AI-HPC

用60%成本干80%的事，DeepSeek分享沉淀多年的高性能深度学习架构

Fire-Flyer AI-HPC: A Cost-Effective Software-Hardware Co-Design for Deep Learning

megatron-kwai

万字干货！手把手教你如何训练超大规模集群下的大语言模型

Accelerating the Training of Large Language Models using Efficient Activation Rematerialization and Optimal Hybrid Parallelism

https://github.com/kwai/Megatron-Kwai

DiPaCo

DiPaCo | Google分布式训练和分布式MoE 新范式（上）

DiPaCo | Google分布式训练和分布式MoE 新范式（下）

DiPaCo: Distributed Path Composition

推理框架

KVCache

kvcache的量化：https://zhuanlan.zhihu.com/p/691537237

一些结论：

浅层KV cache相比于深层KV cache，对模型的重要性更大。

量化

ZeroQuant-V2: Exploring Post-training Quantization in LLMs from Comprehensive Study to Low Rank Compensation和Compression of generative pre- trained language models via quantization

int8量化：LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scale
int4量化：GLM中用了

PPL.LLM

高性能 LLM 推理框架的设计与实现

https://github.com/openppl-public/ppl.llm.serving

vLLM

https://github.com/vllm-project/vllm

Efficient Memory Management for Large Language Model Serving with PagedAttention

对比sglang/trt-llm：https://blog.vllm.ai/2024/09/05/perf-update.html，看着和trt-llm差不多，sglang更菜

并行推理方法

Efficiently scaling transformer inference

3万字详细解析清华大学最新综述工作：大模型高效推理综述

万字综述大模型高效推理：无问芯穹与清华、上交最新联合研究全面解析大模型推理优化

A Survey on Efficient Inference for Large Language Models

LLM后端推理引擎性能大比拼

LayerSkip

LayerSkip: Enabling Early Exit Inference and Self-Speculative Decoding

https://github.com/facebookresearch/LayerSkip

speculative-decoding

Accelerating Large Language Model Decoding with Speculative Sampling

medusa

decoder的并行化： https://zhuanlan.zhihu.com/p/368592551

https://sites.google.com/view/medusa-llm

用了tree-attention

https://github.com/FasterDecoding/Medusa

CLLM

3倍生成速度还降内存成本，超越Medusa2的高效解码框架终于来了

CLLMs：Consistency Large Language Models

fasterTransformer/TensorRTLLM

https://github.com/NVIDIA/FasterTransformer

https://github.com/NVIDIA/TensorRT-LLM/

remove padding的逻辑如下，把整个batch的数据变成一行数据，加上offset标注是哪一条样本的

直接有这么个脚本：

https://github.com/NVIDIA/TensorRT-LLM/blob/main/examples/recurrentgemma/convert_checkpoint.py

https://github.com/daiwk/TensorRT-LLM/blob/main/tensorrt_llm/models/recurrentgemma/model.py

关于plugin：

自己加plugins：https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/openai_triton
已有plugins：https://github.com/NVIDIA/TensorRT-LLM/tree/main/cpp/tensorrt_llm/plugins

一些参数设置：

https://www.bentoml.com/blog/tuning-tensor-rt-llm-for-optimal-serving-with-bentoml

--multiple_profiles enables multiple TensorRT optimization profiles in the built engines, it will benefits the performance especially when GEMM plugin is disabled, because more optimization profiles help TensorRT have more chances to select better kernels. However, this feature will increase the engine build time.