概念

• vLLM 推理加速

训练加速：

• Megatron：并行和流水线
• DeepSpeed：优化器。也有推理能力

更底层的

• NCCL：通信库

vLLM

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

KV Cache 机制

1. 大模型一个一个吐出 Token，然后把已经输出的所有的 Token 放一起继续计算下一个。
2. 这个过程中，前面的 Token 之间的 K 和 V 是不变的。
3. 如果每次都重新计算它们，就太浪费了。
4. 因此把前面的 K 和 V 存下来，就节省了大量计算资源
5. 也正是因为这个机制，在 LLM 输出时，你不会感觉到越来越慢

但 KV Cache 也是有浪费的

1. 大模型一开始不知道要生成多少 Token，因此按最大预分配
2. 大模型预测第一个 Token 时，也要预分配全部 Cache，随着时间才逐渐用完，这也浪费
3. 显存之间有碎片

vLLM 解决了以上问题，用了若干个方法：

1. 参考操作系统对内存的管理，用 “页” 来管理 KV Cache 集合，每页 N 个。也引入了逻辑显存、物理显存的概念
2. 在 Prompt 相同时，只存放一份关于 Prompt 的 KV Cache，做法如下
   • 对 Prompt 对应的 KV 块计数，计数
   • Beam Search

例子

# 安装，要注意 CUDA 版本
# pip install vllm

from vllm import LLM, SamplingParams
prompts = [
    "用 Python 写一个快排函数",
    "描述 AI 的未来"
]

sampling_params = SamplingParams(temperature=0.8, top_p=0.95, max_token=10)
outputs = llm.generate(prompts, sampling_params)

for output in outputs:
    print(output.promt)
    print(output.outputs[0].text)

Megatron

https://github.com/NVIDIA/Megatron-LM

为什么？

• NN 训练耗时 = 训练数据 X 单步计算量 / 计算速率
• 计算速率 = 单设备计算速率 x 设备数 x 多设备并行效率
• 扩展来讲，取决因素为：
• 计算速率 = 单设备计算速率（摩尔定律、混合精度、算子融合、梯度叠加） x 设备数（服务器架构、通讯拓扑优化） x 多设备并行效率（数据并行、模型并行、流水并行）

Megatron：从单卡到千卡测试，算力线形提升

原理：

• 流水线。NN 按照层切分到不同设备，实现并行。（类似计算机原理的流水线）

MoE

OLMoE: Open Mixture-of-Experts Language Models（2024）https://arxiv.org/abs/2409.02060

Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity（2021 https://arxiv.org/abs/2101.03961

混合专家模型。模型中的一部分参数拆分成多个“专家网络”，每次推理只激活若干专家

• 结构
  • 结构：子网络通常是 Transformer 的 FNN 层
  • 路由：小的路由决定如何分配
  • 稀疏激活：每个 token 只计算少数专家
• 优点：
  • 性能
  • 论文指出，从推理最终表现看，A1B-7B 模型 超过稠密模型 7B 和 1B。训练角度看，MoE 的收敛速度、训练效率，都有明显提高
• 缺点：
  • 结构更复杂，
  • 可能有专家负载不均衡的问题