最新大规模语言模型加速实战教程：利用FSx for Lustre和TurboQuant GPUDirect实现高效训练

如果你正在AWS GPU实例上迭代部署大语言模型，大概率已经遇到过这个让人抓狂的场景：模型尺寸越大，加载到GPU显存里的等待时间就越长。当模型膨胀到千亿参数级别、GPU集群规模也越来越大时，模型加载时间会直接拖累端到端的首Token延迟表现。

这篇文章想聊聊如何利用Amazon FSx for Lustre，配合NVIDIA GPUDirect Storage（GDS），再加上一点聪明的方案规划，从根本上改写冷启动场景下的首Token延迟公式。简单来说，这套组合拳能把模型启动时那几分钟的无产出加载时间，压缩到几秒钟。既然聊到了优化，文章最后还会分析新发布的TurboQuant KV缓存技术，看看它对上下文窗口尺寸的巨大提升。

背景：AWS上的NVIDIA Blackwell架构

AWS最近发布了Amazon EC2 P6e和P6实例系列，背后是NVIDIA的Blackwell架构（官方公告在此）。旗舰级的P6e UltraServer将72块NVIDIA Blackwell GPU整合到一个NVLink域中，提供130 TB/s的对分带宽、13.4 TB的HBM3e显存，以及360 petaflops的FP8算力（FP4下为720 petaflops）。这类UltraServer通常用于大规模分布式训练，处理万亿参数级别的前沿模型。

这篇文章的重点是改善单节点P6或P5en实例的冷启动首Token延迟。具体来说，就是如何以最快速度将模型权重以正确格式加载到HBM显存中。对于多节点的UltraCluster，同一流程会在集群内所有节点上并行执行。UltraServer中的每个节点都独立从共享的FSx for Lustre文件系统中加载模型，充分利用FSx for Lustre在GDS支持下提供的海量可扩展吞吐能力。

模型加载的瓶颈在哪

先来看GPUDirect Storage和传统CPU加载方式的根本区别。传统的CPU加载方式（图左）会将检查点流经CPU内存，然后通过PCIe依次将权重拷贝到每块GPU。而分片的GPUDirect Storage加载方式（图右）则预先将检查点按张量并行分片拆分在FSx for Lustre上，8块GPU通过EFA并行、直接地从存储层读取各自的分片到HBM中，完全绕开CPU。具体见图1。

图1：8-GPU实例上，CPU加载与分片GPUDirect Storage加载的对比示意。

以Llama 3.1 405B这个4050亿参数的模型为例，BF16格式的检查点数据大约800 GB。传统的加载路径通常是这样的：

1. 从存储介质中读取检查点文件到CPU系统内存
2. 反序列化权重（例如通过torch.load或safetensors解析）
3. 可选地在CPU上进行量化（BF16→FP8或INT4），以降低HBM需求
4. 通过PCIe将权重从CPU内存逐块依次拷贝到各GPU

这个流程受限于单线程、串行和CPU瓶颈的特性。在典型的非GDS部署中，对于Llama 3.1 405B，单线程模型加载配合CPU端量化，需要耗费10到20分钟。即使使用了预分片的检查点，也仍需几分钟时间。而在这整个过程中，你的GPU——整个架构中最昂贵的资源——完全处于空闲状态，冷启动到首Token的时间就这么被白白拖长了。

值得一提的是，vLLM等推理框架的近期版本在模型加载上做出了显著改进。vLLM V1引擎（自vLLM 0.19起为默认）引入了跨GPU的并行权重加载，相比旧版本大幅缩短了加载时间。但即便如此，数据依然需要流经CPU内存和PCIe总线。而GDS直接移除了这个瓶颈。

对于推理服务来说，这不仅仅是体验上的不便。漫长的模型加载时间会直接影响：

• 冷启动延迟——新实例必须等模型加载完成才能提供服务
• 自动扩缩容的响应速度——扩容事件会延迟几分钟，而非几秒
• 故障恢复——如果服务实例故障，替代实例需要数分钟才能上线
• 成本效率——加载过程中消耗的GPU小时数，本应用来处理推理请求

一条直达路径：FSx for Lustre + GPUDirect Storage

Amazon FSx for Lustre是一个全托管的、高性能的并行文件系统，专为计算密集型工作负载设计。当它和NVIDIA GPUDirect Storage结合后，可以在存储层和GPU显存之间建立多条直接的数据通道，完全绕开CPU和系统内存。

这种集成依赖于两项关键技术的协同：

• Amazon Elastic Fabric Adapter (EFA) 使用AWS的可扩展可靠数据报（SRD）协议来绕过操作系统层面的开销。P5en实例配备了16个200 Gbps的EFA接口，总网络带宽达到3,200 Gbps（400 GB/s）。FSx for Lustre可以使用其中8个或更多的EFA接口来实现从存储到GPU的直接数据传输。
• NVIDIA GPUDirect Storage (GDS) 允许DMA传输直接从网络接口写入GPU HBM，消除了原本在CPU内存中暂存数据导致的传统瓶颈。

这带来的不仅仅是吞吐量的显著提升。真正的关键在于，这条直接路径在架构上解锁了新的可能性：当使用分片并行模型加载时，完全可以绕过CPU。

在我们的测试配置中，使用的是Persistent_2 EFA文件系统，吞吐设置为1000 MBps/TiB，包含20个对象存储目标（OST），容量为96 TiB，文件系统吞吐约94 GiB/s。文件系统吞吐随容量线性增长——更大的文件系统意味着更多的OST和更多的并行I/O通道。关于如何构建高吞吐FSx for Lustre文件系统的详细指南，可以参考之前的博客文章《在不到一小时内构建和部署1 TB/s文件系统》。具体设置命令见下一部分的阶段0。

要将P5en客户端配置为支持GDS访问FSx for Lustre，请遵循FSx for Lustre用户指南中的“配置EFA客户端”步骤。AWS提供的配置脚本（setup.sh –optimized-for-gds）会自动检测你的实例类型，配置正确的EFA接口并进行NUMA感知的CPU分区，并创建一个systemd服务以确保配置重启后依然生效。默认情况下，P5en只使用其中8个EFA接口来连接FSx for Lustre，这为从存储到各GPU的HBM提供了直接的GDS数据路径。

P5en（8x H200）上的分片并行加载

为了演示分片加载方法，我们使用P5en实例（p5en.48xlarge）——8块NVIDIA H200 GPU，每块141 GB HBM3e显存，通过NVSwitch互联，提供3.6 TB/s的对分带宽。P5en支持GDS，并具备与P6-B200相同的3,200 Gbps SRD EFA网络能力，因此非常适合展示这一模式。对于存储读取阶段，其性能特征与P6节点完全一致，因为每个实例的网络带宽相同。

对于任何足够大的模型（比如FP8格式下400 GB的Llama 3.1 405B），单个GPU的HBM（H200为141 GB）是无法容纳的。这时必须使用张量并行（TP）技术，将模型拆分到多个GPU上。这意味着推理过程中，GPU之间需要通过NVLink进行低延迟通信（例如在注意力层和多层感知机层之后执行all-reduce操作）。这也解释了为什么像P5en和P6这样配备NVSwitch的实例对于服务这类模型至关重要。

整个方案分为四个阶段。虽然我们以Llama 3.1 405B作为参考模型，但这种模式适用于任何支持张量并行分片的模型，包括Mixtral、DeepSeek及其他自定义架构。核心要求是推理框架能够将模型拆分为每个GPU所需的分片。

阶段0：基础设施准备。在加载任何模型之前，需在同一个Amazon VPC和可用区内准备两样东西：一个启用了EFA的FSx for Lustre文件系统，和一个配置了GPUDirect Storage的P5en（或P6）GPU实例。随附的aws-samples仓库中提供了AWS CloudFormation模板和设置脚本，可自动完成下面描述的基础设施部署和GDS配置。

FSx for Lustre文件系统应部署为Persistent_2 SSD，启用EFA，吞吐设置为1000 MBps/TiB（如前一部分所述）。文件系统的容量决定了总读取吞吐——容量越大，OST越多，并行I/O路径越丰富。

P5en实例需要配置所有EFA接口、安装EFA驱动、构建并加载NVIDIA GDS内核模块nvidia-fs.ko、根据实例的NUMA拓扑配置NUMA感知的Lustre客户端网络，并放置好GDS运行时配置文件（cufile.json）。这涉及几个步骤：构建NVIDIA nvidia-fs.ko模块、将EFA接口与基于NUMA拓扑的CPU分区对齐、调优Lustre客户端参数。完整的设置步骤请参考FSx for Lustre用户指南中的“配置EFA客户端”指南。

基础设施就绪后，需要在输出目录上设置Lustre条带化，以确保来自多个GPU的并发GDS读取能够分布到所有OST上：

# 跨所有OST条带化，条带大小16MB（匹配最佳GDS块大小）
mkdir -p /fsx/model_shards/Llama-3.1-405B-FP8-8way
lfs setstripe -c -1 -S 16M /fsx/model_shards/Llama-3.1-405B-FP8-8way

阶段1：离线预分片与预量化模型权重。在离线环境下，使用vLLM将模型拆分为8个张量并行分片并执行FP8量化，然后保存到FSx for Lustre。源检查点为BF16格式（标准Llama 3.1 HuggingFace格式）。预分片步骤将权重量化到FP8，数据量因此减半，后续需通过GDS加载的数据也更少：

python sa ve_sharded_state.py --model /fsx/models/Llama-3.1-405B --quantization fp8 --tensor-parallel-size 8 --output /fsx/model_shards/Llama-3.1-405B-FP8-8way

这一步会生成8个TP感知的分片文件，每个分片精确包含该GPU推理所需的权重切片。注意力头和MLP的列被正确地拆分到各GPU上。预量化为FP8后，总检查点大小从约800 GB降至约400 GB。输出目录结构如下：

/fsx/model_shards/Llama-3.1-405B-FP8-8way/
├── model-rank-0-part-0.safetensors # ~51 GB — GPU 0的分片
├── model-rank-1-part-0.safetensors # ~51 GB — GPU 1的分片
├── ...
├── model-rank-7-part-0.safetensors # ~51 GB — GPU 7的分片
├── config.json
├── tokenizer.json
└── tokenizer_config.json

每次更新基础模型检查点后（例如微调后或切换到新模型版本），都需要重复此预分片步骤。但由于它在离线完成，且之后每次加载都复用此输出，摊销成本非常低。

如果训练流程可控，这个步骤甚至可以完全省去。像Megatron-LM和NeMo这样的框架可以直接以推理栈所需的张量并行布局和精度保存检查点，例如8路TP分片、FP8格式的safetensors。当训练输出直接就是服务格式时，检查点无需任何后处理即可用于GDS并行加载。需要注意的是，Megatron-LM默认以torch_dist格式保存检查点，需要使用Megatron Bridge进行转换才能用于基于safetensors的加载器——这是一次性的转换步骤，额外开销很小。

之所以选择FP8，是因为它是H200和B200上的一等数据类型，原生支持Tensor Core；在vLLM中只需一个参数（–quantization fp8）即可启用；且对Llama类模型而言，精度损失几乎可以忽略不计。更激进的4-bit权重量化方法（AWQ、GPTQ、HQQ）配合vLLM中优化的W4A16推理内核，可以再将检查点大小减半，并在带宽受限的工作负载上（例如H100/H200上小批量场景下）将生成吞吐提升2-3倍。AWQ和GPTQ需要对代表性数据进行一次简短的校准，HQQ则无需数据。本文中的GDS加载模式与量化方法无关——无论服务框架支持何种格式，从FSx for Lustre进行并行分片加载的方式都同样适用。一个4-bit的405B检查点，加载时间大约会是下文报告FP8结果的一半。

FP8量化将检查点大小减半，从而缩短加载时间，但权重量化并非唯一的杠杆。新兴技术如TurboQuant（Google Research，ICLR 2026）及其底层的PolarQuant方法，目标锁定在Key-Value（KV）缓存——推理过程中随上下文长度增长而消耗的内存。根据论文描述，TurboQuant将KV缓存压缩到约每值3比特（6倍压缩），在NVIDIA H100 GPU上可实现注意力计算高达8倍的加速，且无需微调、零精度损失。虽然这些方法不直接减少检查点大小或模型加载时间，但它们显著降低了推理过程中的HBM消耗，为更长的上下文窗口或更大的批次腾出GPU内存。结合本文使用的FP8权重量化，KV缓存压缩进一步降低了总体HBM需求，助你在相同硬件上服务更大的模型或处理更多并发请求。

阶段2：并行GDS读取。加载时，所有8块GPU通过GDS同时从FSx for Lustre直接读取各自分片到HBM中。由于分片已经预量化到FP8，总读取数据约400 GB，每块GPU约50 GB。GDS完全绕开CPU内存，数据不需要经过主机串行化，8次读取完全并行进行。

在GDS读取路径上，我们使用fastsafetensors这个开源库，它利用NVIDIA cuFile（GDS API）直接将safetensors文件读取到GPU内存中，并重构PyTorch张量，全程无需任何CPU端反序列化。每块GPU打开自己的分片文件，通过一次大型GDS读取将数据放入GPU缓冲区，然后根据safetensors头文件中的元数据从中提取所有张量。提取过程仅需不到1毫秒，本质上是对已加载到GPU缓冲区的数据进行指针运算。

加载模型前，需确认GDS内核模块已加载、Lustre客户端已配置为使用EFA：

# 检查GDS内核模块是否加载
lsmod | grep nvidia_fs
# 检查EFA接口是否已为Lustre激活
sudo lnetctl net show | grep -c "nid:.*@efa"
# 应显示16（为Lustre配置的EFA接口数量——并非实例的全部网卡）
# 单节点部署中，设置脚本会为Lustre配置全部16个EFA接口。
# 在UltraCluster多节点配置中，你可能需要配置8个接口用于FSx，
# 保留其余8个用于节点间的NCCL集合通信。

如果nvidia_fs未加载，GDS读取会静默回退到CPU bounce buffer路径。此时结果依然正确，但性能大打折扣。使用sudo modprobe nvidia_fs加载该模块。

使用fastsafetensors，每块GPU并行加载其分片：

from fastsafetensors import SafeTensorsFileLoader
# 每块GPU通过GDS加载自己的分片
loader = SafeTensorsFileLoader(pg=None, device=f"cuda:{rank}", nogds=False)
loader.add_filenames({0: [f"/fsx/model_shards/Llama-3.1-405B-FP8-8way/model-rank-{rank}-part-0.safetensors"]})
fbuf = loader.copy_files_to_device()
# GDS读取：存储 → GPU HBM，直接传输
# 提取张量——不到一毫秒，仅对GPU缓冲区的指针操作
tensors = {name: fbuf.get_tensor(name) for name in loader.get_keys()}

阶段3：验证与服务。通过GDS将张量加载到GPU HBM后，所有推理计算完全在GPU内存中进行。FSx for Lustre文件系统仅在初始加载阶段参与。阶段2获得的tensors字典包含了当前rank的所有权重，已位于正确的GPU设备上，并已采用正确的张量并行布局。整个过程未接触任何CPU内存，也没有反序列化操作。权重直接从FSx for Lustre存储直达GPU HBM。

这些张量已准备好供任何接受预加载权重字典的张量并行推理引擎使用。具体的集成点取决于框架：vLLM、TensorRT-LLM和SGLang各自有内部API用于将权重注入其TP感知的模型计算图中。随着这些框架原生支持GDS感知的权重加载（fastsafetensors正是为此类集成而生，由Foundation Model Stack团队构建），阶段2展示的完整GDS路径将简化为一个命令。

对于当前生产环境中的推理服务，vLLM可以从FSx for Lustre加载同样的预分片检查点：

vllm serve /fsx/model_shards/Llama-3.1-405B-FP8-8way --load-format sharded_state --quantization fp8 --tensor-parallel-size 8

vLLM的内置权重加载器使用的是标准CPU读取路径，而非GDS。但预分片格式仍然消除了反序列化和逐GPU权重拆分这些在标准检查点加载中占主导的开销，将Llama 3.1 405B的加载时间从约18分钟降至约2分钟（详见下一节性能对照表）。而阶段2中的GDS并行路径更近一步，将约2分钟缩短到约6秒。

性能差异

以下是实测结果。下面的表格比较了在P5en实例（8x H200）上，使用96 TiB Persistent_2 EFA文件系统（20个OST，约94 GiB/s文件系统吞吐）时，不同加载方式下的模型加载时间。

实测：Llama 3.1 70B Instruct（8路TP，冷缓存）

GDS加载时间为每rank的耗时（8块GPU并行加载），使用fastsafetensors直接从GPU内存中读取safetensors文件，无CPU bounce buffer和反序列化开销。每个rank加载其分片，所有1124个张量在1秒多的时间内即就绪于对应GPU设备上。基线时间包含了vLLM引擎初始化和预热，而GDS行代表的是存储到GPU的传输时间。

实测：Llama 3.1 405B Instruct（8路TP，冷缓存）

这些结果基于96 TiB文件系统（20个OST）。由于GDS吞吐随OST数量线性扩展，更大的文件系统会成比例地缩短加载时间。例如，一个342 TiB文件系统（57个OST），理论上可将FP8加载时间压缩到3秒以下。粗略估算：94 GiB/s × 57/20 OST ≈ 268 GiB/s的理论上限，保守考虑OST开销后约190 GiB/s，408 GB ÷ 190 GiB/s ≈ 2.0秒。

在P6节点上，由于每实例带宽相同，加载性能将完全一致。

这种模式对于需要跨多GPU张量并行的大模型效果最为显著——因为并行读取的优势真正发挥出来了。对于能够单GPU承载的小模型，传统瓶颈在CPU端（反序列化、量化、串行PCIe传输），GDS和并行读取的优势有限。一旦模型需要分布在多个GPU上，每个GPU就可以通过GDS独立读取其分片，CPU不再是关键路径上的瓶颈。

即使使用预量化检查点，单线程路径仍需耗费大量时间在反序列化和通过PCIe跨8块GPU串行进行主机到设备的数据拷贝上。这解释了原始存储读取时间与总加载时间之间的差距。

分片GDS方法在6.4秒内将408 GB的FP8模型权重加载到8块GPU上。相比标准vLLM非GDS加载，提升了约169倍，而且这只是基于96 TiB文件系统的表现。更大的文件系统将进一步推高加速比。这意味着一件事：从“先去吃个午饭吧”到“已经加载完了”的体验飞跃。

加速来源于同时消除了传统路径上的每一个瓶颈：

• 无CPU bounce buffer——GDS直接将数据读入GPU HBM
• 无串行反序列化——预分片权重加载即用
• 无CPU量化——离线预量化，而非加载时进行
• 无顺序GPU加载——8块GPU并行读取
• 无需all-gather——TP感知的分片确保每块GPU已获得其所需权重

但快速加载只是故事的一半。FP8量化不仅缩短了加载时间，还将模型权重的HBM占用量减半，为KV缓存和推理服务腾出了显存。再结合TurboQuant的KV缓存压缩（3-4比特/值），可用于推理的GPU内存大幅增加：

FP8权重加上TurboQuant KV缓存压缩的组合意味着，你可以在单台P5en上以超过40万Token的上下文窗口服务Llama 3.1 405B，在P6节点上则接近50万Token。这意味着从处理几份文档到单次请求处理整本书的质的飞跃。

还值得一提的是，FSx for Lustre文件系统并非单一用途的基础设施。同一个加速模型加载的文件系统，也可以作为团队训练数据、检查点、数据集和模型制品的共享存储。通过减少模型加载期间的GPU空闲时间，可以提升模型测试和评估的迭代速率——几秒内加载一个新模型变体，而不是几分钟，意味着每天可以运行更多实验。最新定价请参考FSx for Lustre定价页面。

与推理框架的集成

这套模式可以与你正在使用的推理框架无缝配合。

TensorRT-LLM 原生支持张量并行检查点加载。你可以转换并构建8路张量并行的引擎，然后通过mpirun在所有GPU上启动。当底层文件系统支持GDS时，TensorRT-LLM可自动利用从存储到GPU的直接数据路径。

vLLM 支持跨GPU的张量并行，可通过--tensor-parallel-size 8配置模型服务。虽然vLLM的默认加载路径是基于CPU的，但预分片检查点方法配合GDS存储，可以在文件系统层面提供I/O加速。

总结

通过将Amazon FSx for Lustre与NVIDIA GPUDirect Storage以及预分片、预量化的模型检查点相结合，我们将Llama 3.1 405B的加载时间从非GDS环境下的10-20分钟，缩短到基于96 TiB文件系统GDS环境下的6秒。扩展到更大的文件系统还可以获得进一步增益。此外，通过应用TurboQuant KV缓存压缩（3-4比特/值），Llama 3.1 405B的可用上下文窗口从约8.2万Token增加到P5en上的超40万Token，或从约13.1万增加到P6上的约66万Token。这是在相同硬件上实现的5倍提升。

这套方法的核心优势：

• 冷启动速度大幅提升——新推理实例在几秒内就绪，而非几分钟
• 自动扩缩容响应更敏捷——扩容事件近乎实时应对流量高峰
• 每Token成本更低——GPU时间花在推理上，而非等待权重加载
• 故障恢复更快——故障实例被替换后，几秒内即可恢复服务
• 随集群扩展而线性扩展——UltraCluster中每个节点独立从同一共享文件系统并行加载
• 上下文窗口极大扩展——TurboQuant KV缓存压缩在相同硬件上实现5倍上下文长度

这套模式现已可用于P5en和P6实例，配合FSx for Lustre Persistent_2 EFA文件系统，以及vLLM和TensorRT-LLM等标准推理框架。立即开始，更快地加载你的大模型吧。

来源：互联网