利用 Amazon FSx for Lustre 上的 GPUDirect 和 TurboQuant 加速 LLM 模型加载并扩大上下文窗口

在 AWS GPU 实例上迭代部署大型语言模型 (LLM) 时，模型加载到 GPU 高带宽内存 (HBM) 的时间会随着模型规模的增大而显著增加，导致冷启动首次令牌时间 (TTFT) 过长。本文介绍了如何通过 Amazon FSx for Lustre 与 NVIDIA GPUDirect Storage (GDS) 的结合，加上巧妙的规划，从根本改变冷启动 TTFT 的计算方式，将模型加载时间从分钟缩短到秒。同时，还介绍了新发布的 TurboQuant KV 缓存对增加上下文窗口的巨大作用。

背景：AWS 上的 NVIDIA Blackwell 架构 AWS 近期推出了基于 NVIDIA Blackwell 架构的 Amazon EC2 P6e 和 P6 实例系列。旗舰级 P6e UltraServer 将 72 个 NVIDIA Blackwell GPU 集成到一个 NVLink 域中，提供 130 TB/s 的双向带宽、13.4 TB 的 HBM3e 和 360 petaflops 的 FP8 计算能力。这些 UltraServer 通常用于万亿参数级前沿模型的大规模分布式训练。本文则聚焦于优化单个 P6 或 P5en 实例的冷启动 TTFT。

模型加载瓶颈 传统 CPU 模型加载将检查点流经 CPU 内存，并依次通过 PCIe 将权重复制到每个 GPU。而分片 GDS 加载则预分检查点，所有 GPU 直接通过 EFA 从 FSx for Lustre 并行读取分片到 HBM，完全绕过 CPU。对于 Llama 3.1 405B 模型（约 800 GB BF16 检查点），传统加载需要 10–20 分钟，而 GDS 几乎可以消除这一瓶颈。长时间加载不仅影响冷启动延迟，还影响自动扩缩响应速度、故障恢复以及成本效率。

直接路径：FSx for Lustre 与 GPUDirect Storage Amazon FSx for Lustre 是一种高性能并行文件系统，与 NVIDIA GDS 结合时，可以建立多条直达 GPU 内存的数据路径。此集成依赖两个关键技术：Amazon Elastic Fabric Adapter (EFA) 和 NVIDIA GPUDirect Storage。EFA 使用 SRD 协议绕过操作系统开销，P5en 实例具有 16 个 200 Gbps 的 EFA 接口，总聚合网络带宽达 3,200 Gbps。GDS 允许 DMA 直接从网络接口传输到 GPU HBM。测试中使用了 Persistent_2 EFA 文件系统，每 TiB 提供 1000 MBps 吞吐量，20 个 OST 可提供约 94 GiB/s 的文件系统吞吐量。

分片并行加载 以 P5en 实例（8 个 H200 GPU，每个 141 GB HBM3e）为例，对于 Llama 3.1 405B 模型（FP8 时约 400 GB），权重无法装入单个 GPU，需要张量并行。方法分为四个阶段：

• 阶段 0：配置基础设施：在同一个 VPC 和可用区中准备启用 EFA 的 FSx for Lustre 文件系统和配置了 GDS 的 GPU 实例。使用 AWS CloudFormation 模板和设置脚本可自动完成。
• 阶段 1：预分片和预量化：离线使用 vLLM 将模型分割为 8 个张量并行分片并量化为 FP8，保存到 FSx for Lustre。检查点大小从 ~800 GB 降至 ~400 GB。此步骤仅在模型更新时需要重复。
• 阶段 2：GDS 并行加载：启动推理服务器时，每个 GPU 通过 GDS 直接从 FSx for Lustre 读取自己的分片，并行加载权重到 HBM。
• 阶段 3：推理：模型就绪后立即开始服务推理请求。

TurboQuant KV 缓存：TurboQuant 是一种量化技术，以 4 位甚至 2 位精度存储 KV 缓存，在不影响质量的情况下将每个令牌的 KV 缓存内存占用减少 4–8 倍，从而支持更长的上下文窗口。

通过结合这些技术，用户可以在 AWS 上实现接近即时的大型模型冷启动，并支持更长的上下文。