PythonAI算法
大模型 Offload 技术实现低显存推理优化
本文介绍了大模型低显存推理优化的常用方法，重点讲解了 Offload 技术在 HuggingFace Transformers、DeepSpeed ZeRO-Inference、llama.cpp、vLLM 及 TensorRT-LLM 等框架中的应用。通过权重卸载和 KV Cache 卸载策略，实现在有限显存下运行大规模模型，涉及 Kernel 优化、模型压缩等技术细节及具体代码示例。
本文介绍了大模型低显存推理优化的常用方法，重点讲解了 Offload 技术在 HuggingFace Transformers、DeepSpeed ZeRO-Inference、llama.cpp、vLLM 及 TensorRT-LLM 等框架中的应用。通过权重卸载和 KV Cache 卸载策略，实现在有限显存下运行大规模模型，涉及 Kernel 优化、模型压缩等技术细节及具体代码示例。
近两年大模型火出天际；同时，也诞生了大量针对大模型的优化技术。由于 GPU 资源稀缺和昂贵，一定程度上使得大模型的应用和普及变得困难。因此，本文从大模型推理部署出发，介绍了一些低硬件成本（消费级 GPU、CPU 等）部署大模型的方法，并着重介绍了低显存推理优化技术 Offload 在各个 AI 框架的应用。
对于 LLM 推理场景而言，内存需求主要包括如下几个方面：
因此，低显存推理优化的对象通常为权重、KV Cache、激活等。
而要使用更低的显存进行推理，常用的优化手段有 Kernel 优化、模型压缩、Offload 技术等。
对于 Kernel 优化，主要有针对 Attention 的优化，比如：
对于推理场景，模型压缩主要的应用技术有模型量化和模型剪枝。
对于模型量化常用的技术有：
对于模型剪枝常用的技术有：
Offload 技术（也可以理解为混合部署或混合推理）将对象卸载到外部资源或将其分布在多种设备上以减少主设备的内存负载，常用用于在 GPU 显存有限的设备上优化内存使用。比如：将 GPU 显存中的权重卸载到 CPU 内存、NVMe/Disk。
对于推理场景下，Offload 的对象有以下两种：
目前，Offload 技术已经广泛用于各 AI 框架用于降低训练和推理的显存。下面来看看 Offload 在 AI 框架中的应用。
通常情况下，使用 Pytorch 来加载模型步骤如下：
state_dict 的字典对象）示例代码：
import torch
my_model = ModelClass(...)
state_dict = torch.load(checkpoint_file)
my_model.load_state_dict(state_dict)
对于超大模型进行推理，这种常规的加载方式存在一些明显的局限性。第一步，在 RAM 中加载模型的完整版本，并花一些时间随机初始化权重，第二步将模型的另一个完整预训练权重加载到 RAM 中。因此，对于一个 6B 的模型，使用 FP16 半精度加载，也需要 24G 内存。因此，这种常规的加载方式已经无法满足要求。
在 Transformers 库中利用 Accelerate 库来完成这些超大模型的加载。那么 Accelerate 如何利用 PyTorch 来加载和运行大模型进行推理呢？具体步骤如下：
示例代码如下：
import torch
from accelerate import init_empty_weights, load_checkpoint_and_dispatch
with init_empty_weights():
model = MyModel(...)
model = load_checkpoint_and_dispatch(
model, checkpoint=checkpoint_file, device_map="auto"
)
input = torch.randn(2,3)
input = input.to("cuda")
output = model(input)
上面的代码通过设置 device_map=auto，会根据可用资源自动确定模型每一层的放置位置。
在 Transformers 库中，也是通过在 from_pretrained() 或者 pipeline() 函数中设置 device_map=auto 完成同样的操作。
此外，device_map 有如下选项可供设置：
"auto" 或 "balanced": Accelerate 将会根据所有 GPU 均衡切分权重，尽量均匀的切分到各个 GPU 上；"balanced_low_0": Accelerate 均匀分割权重到各个 GPU 上，除了第一个 GPU（序号为 0）。在第一个 GPU 上会尽量节省显存（这种模式可以有效节省第一个 GPU 的显存，以便使用 generate 函数用于模型生成）；"sequential": Accelerate 按照 GPU 的顺序占用显存（后面的 GPU 可能根本不会使用）。当然也可以根据需求设置 device_map 参数，以决定各个部分权重应该放置的设备。
device_map = {"block1": 0, "block2.linear1": 0, "block2.linear2": 1, "block2.linear3": 1}
此外，还可以通过 max_memory 来控制各个设备使用的内存的大小。
max_memory={0: "10GiB", 1: "20GiB", 2: "20GiB", "cpu": "60GiB"}
下面展示了 Transformers 使用 3 张 4090 加载 Qwen1.5-72B 模型。其中，GPU0 显存限制为使用 10G、GPU1/2 显存限制为使用 20G、CPU 内存限制为使用 60G，其余使用 Disk 加载。
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
path = "/workspace/models/Qwen1.5-72B-Chat"
model = AutoModelForCausalLM.from_pretrained(
path,
device_map="auto",
#max_memory={0: "15GiB", 1: "20GiB", "cpu": "60GiB"},
#max_memory={0: "15GiB", 1: "20GiB", 2: "20GiB", "cpu": "60GiB"},
max_memory={0: "10GiB", 1: "20GiB", 2: "20GiB", "cpu": "60GiB"},
offload_folder="offload",
#offload_buffers=True,
offload_state_dict = True,
torch_dtype=torch.float16
)
model.hf_device_map
加载完成之后，可以看到模型中每一层放置的具体位置。
{
'model.embed_tokens': 0,
'model.layers.0': 0,
'model.layers.1': 0,
'model.layers.2': 0,
'model.layers.3': 1,
...
'model.layers.14': 1,
'model.layers.15': 2,
...
'model.layers.26': 2,
'model.layers.27': 'cpu',
...
'model.layers.61': 'cpu',
'model.layers.62': 'disk',
...
'model.layers.79': 'disk',
'model.norm': 'disk',
'lm_head': 'disk'
}
之后即可进行模型推理了。
此外，Transformers 为了减少 GPU VRAM 使用，实现了 Offload KV Cache 的功能，通过将大多数层的 KV 缓存移至 CPU 来实现这一点。
当模型的 forward() 方法迭代各层时，该策略会在 GPU 上维护当前层的缓存。同时，它异步预取下一层缓存，并将上一层缓存发送回 CPU。与 KV 缓存量化（可能损失模型精度）不同，此策略始终与默认 KV Cache 实现产生相同的结果。示例代码如下：
>>> import torch
>>> from transformers import AutoTokenizer, AutoModelForCausalLM
>>> ckpt = "microsoft/Phi-3-mini-4k-instruct"
>>> tokenizer = AutoTokenizer.from_pretrained(ckpt)
>>> model = AutoModelForCausalLM.from_pretrained(ckpt, torch_dtype=torch.float16).to("cuda:0")
>>> inputs = tokenizer("Fun fact: The shortest", return_tensors="pt").to(model.device)
>>> out = model.generate(**inputs, do_sample=False, max_new_tokens=23, cache_implementation="offloaded")
>>> print(tokenizer.batch_decode(out, skip_special_tokens=True)[0])
Fun fact: The shortest war in history was between Britain and Zanzibar on August 27, 1896.
>>> out = model.generate(**inputs, do_sample=False, max_new_tokens=23)
>>> print(tokenizer.batch_decode(out, skip_special_tokens=True)[0])
Fun fact: The shortest war in history was between Britain and Zanzibar on August 27, 1896.
因此，它可以作为默认 KV 缓存的直接替代品或备用方案，具体示例如下：
>>> import torch
>>> from transformers import AutoTokenizer, AutoModelForCausalLM
>>> def resilient_generate(model, *args, **kwargs):
... oom = False
... try:
... return model.generate(*args, **kwargs)
... except torch.cuda.OutOfMemoryError as e:
... print(e)
... print("retrying with cache_implementation='offloaded'")
... oom = True
... if oom:
... torch.cuda.empty_cache()
... kwargs["cache_implementation"] = "offloaded"
... return model.generate(*args, **kwargs)
...
...
>>> ckpt = "microsoft/Phi-3-mini-4k-instruct"
>>> tokenizer = AutoTokenizer.from_pretrained(ckpt)
>>> model = AutoModelForCausalLM.from_pretrained(ckpt, torch_dtype=torch.float16).to("cuda:0")
>>> prompt = ["okay "*1000 + "Fun fact: The most"]
>>> inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
>>> beams = { "num_beams": 40, "num_beam_groups": 40, : , : , : , : , }
out = resilient_generate(model, **inputs, **beams)
responses = tokenizer.batch_decode(out[:,-:], skip_special_tokens=)
2022 年 9 月 DeepSpeed 提出了 ZeRO-Inference 用于 Offload 权重。利用 ZeRO Stage 3 的数据并行特性，能够将模型分布到多张 GPU 上或者 Offload 到内存或者 NVMe 上，使单 GPU 无法加载的模型进行推理。
工作原理
大模型推理的大量计算需求意味着需要 GPU 等加速器才能高效执行。因此，在有限的 GPU 预算上进行大模型推理的一个重要设计决策是如何在模型权重、推理输入和中间结果之间分配 GPU 内存。Zero Inference 提供了两种方案。
第一种方案：卸载所有模型权重
将整个模型权重固定在 CPU 或 NVMe 中，并将权重逐层流式传输到 GPU 中以进行推理计算。计算完一层后，输出将保留在 GPU 内存中作为下一层的输入，而层权重消耗的内存将被释放以供下一层使用。因此，模型推理时间由在 GPU 上计算各层的时间和通过 PCIe 获取各层的时间组成。
ZeRO-Inference 通过两种方式提供了可扩展性的优势。
首先，通过只在 GPU 内存中保留一个（或几个）模型层，ZeRO-Inference 显著减少了推理大规模模型所需的 GPU 内存量。对于当前具有大约一百层的大模型（例如，GPT3-175B 和 Megatron-Turing-530B 中分别有 96 层和 105 层），ZeRO-Inference 将 GPU 内存需求降低了两个数量级。例如，通过 ZeRO-Inference，Megaton-Turing-530B 半精度（FP16）推理的 GPU 内存消耗从 1TB 下降到 10GB。
其次，通过将模型放置到比 GPU 内存便宜几个数量级的 CPU 或 NVMe 内存中，ZeRO-Inference 使得扩展到更大的模型（例如，具有数万亿或数十万亿参数）的成本比其他方法（将整个模型放入 GPU 内存中）更经济实惠。
尽管通过 PCIe 互连从 CPU 或 NVMe 获取模型权重存在延迟，但 ZeRO-Inference 仍可为面向吞吐量的推理应用程序提供高效计算。主要原因是通过将 GPU 内存使用限制为模型的一层或几层权重，ZeRO-Inference 可以使用大部分 GPU 内存来支持大量以长序列或大批量形式出现的输入 Token。大模型层需要大量计算，尤其是在处理具有许多输入 Token 的输入时。例如，一个 GPT3-175B 的 Transformer 层来处理批量大小为 1、序列长度为 2048 的输入需要大约 7 TFlops。因此，对于长序列长度和大批量大小的推理场景，计算时间在获取模型权重的延迟中占主导地位，这最终提高了效率。总而言之，ZeRO-Inference 利用 GPU 内存支持大量输入 Token 的策略可以为大型模型提供高性能推理。
此外，ZeRO-Inference 利用两项额外的优化来减少从 CPU 或 NVMe 获取层权重到 GPU 内存的延迟。
第二种方案： 在 GPU 内存中托管一些模型权重
将尽可能多的模型权重固定到 GPU 内存中，并在需要计算时获取剩余部分（从 CPU 或 NVMe）。这种方法的一个好处是避免了获取已固定在 GPU 内存中的权重的延迟。然而，这种方法有两个缺点：
补充：ZeRO-Inference 与 DeepSpeed-Inference 的区别
ZeRO-Inference 与 DeepSpeed-Inference 的示例代码如下：
# Launch with `deepspeed deepspeed-zero-inference.py`
import torch
import deepspeed
import os
import time
from transformers.deepspeed import HfDeepSpeedConfig
from transformers import AutoConfig, AutoTokenizer, AutoModelForCausalLM
local_rank = int(os.getenv("LOCAL_RANK", "0"))
world_size = int(os.getenv("WORLD_SIZE", "1"))
model_name = "/workspace/models/llama-7b-hf"
def run_zero_inference():
ds_config = {
"fp16": {"enabled": True},
"bf16": {"enabled": False},
"zero_optimization": {
"stage": 3,
"offload_param": {
"device": "cpu",
},
},
"train_micro_batch_size_per_gpu": 1,
}
# 与 HuggingFace 共享 DeepSpeed 配置，以便我们可以正确使用 zero stage 3 加载大模型
hfdsc = HfDeepSpeedConfig(ds_config)
# Load the model and tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
model_name, torch_dtype=torch.float16
)
# Initialize DeepSpeed
ds_engine = deepspeed.initialize(model=model, config_params=ds_config)[0]
ds_engine.module.eval()
model = ds_engine.module
# Run inference
start_time = time.time()
inputs = tokenizer.encode("DeepSpeed is", return_tensors="pt").to(
)
outputs = model.generate(inputs, max_new_tokens=)
output_str = tokenizer.decode(outputs[])
end_time = time.time()
(output_str)
(, end_time - start_time)
()
():
config = AutoConfig.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=)
deepspeed.OnDevice(dtype=torch.float16, device=, enabled=):
model = AutoModelForCausalLM.from_config(config, torch_dtype=torch.float16)
checkpoint_dir =
checkpoint_files = [
os.path.join(checkpoint_dir, )
i (, )
]
checkpoint_dict = {
: ,
: checkpoint_files,
: ,
}
(checkpoint_dict)
model = deepspeed.init_inference(
model,
replace_with_kernel_inject=,
tensor_parallel={: world_size},
dtype=torch.float16,
checkpoint=checkpoint_dict,
)
start_time = time.time()
inputs = tokenizer.encode(, return_tensors=).to(
)
outputs = model.generate(inputs, max_new_tokens=)
output_str = tokenizer.decode(outputs[])
end_time = time.time()
(output_str)
(, end_time - start_time)
__name__ == :
run_zero_inference()
run_deepspeed_inference()
2023 年 9 月 DeepSpeed 更新了 ZeRO-Inference 以支持 Offload KV Cache。相比之前的版本，通过利用两项内存优化（权重量化和KV cache offload），将推理吞吐量提升高达 20 倍。
模型主要的内存消耗来自模型权重和 KV Cache，这限制了可用于推理的最大批量大小（从而限制了吞吐量）。ZeRO-Inference 现在提供模型权重的 4 位量化，从而使其内存使用量和 PCIe 传输量减少大约 4×。同时，在 DeepSpeed 框架中集成了高效的量化/反量化 Kernel。
此外，KV Cache 作为提高系统吞吐量的另一个限制因素，现在可以将其卸载到更便宜的 CPU 内存，而不是使用更昂贵的 HBM 容量。官方提供了对 BLOOM、LLAMA2、OPT 模型的支持。如果需要支持其他模型进行 KV Cache Offload，则需要进行少许代码更改。具体步骤如下：
set_kv_cache_offload() 函数，为 HF 模型启用 KV 缓存卸载。该函数适当修改 HF 模型的 forward 函数，以触发注意力模块中的卸载行为。2 * num_layers * batch_size * max_seq_len * hidden_size，其中 2 表示 K 值和 V 值，num_layers 是 Transformer 块的数量，batch_size 是推理批量大小，max_seq_len 是提示和生成的 Token 的总长度，hidden_size 是模型的隐藏维度。在单个 NVIDIA A6000 GPU 上对各种模型大小执行面向吞吐量的推理。与之前版本相比，新版 ZeRO-Inference 的整体吞吐量提升明显。
此外，HF Transformers/Accelerate 对 DeepSpeed ZeRO-Inference 进行了集成，以便进行推理。
如果使用 Trainer，要运行 ZeRO Inference，将常用训练参数传递给类并添加 --do_eval 参数，示例如下。
deeepspeed --num_gpus=2 your_program.py <normal cl args> --do_eval --deepspeed ds_config.json
DeepSpeed 也可以在没有类的情况下与 Transformer 配合使用。由 HfDeepSpeedConfig 进行处理，当您调用时，它只负责收集（gathering） ZeRO-3 参数并在多个 GPU 上分割模型。
注意：如果您希望自动为您处理所有事情，请将 DeepSpeed 与结合使用！需遵循，并在配置文件中手动配置参数值（不能使用
"auto"值）。
要高效部署 ZeRO-3，必须在模型之前实例化 HfDeepSpeedConfig 对象并保持该对象处于活动状态。
不使用 Trainer 的情况下的示例如下所示：
#!/usr/bin/env python
# This script demonstrates how to use Deepspeed ZeRO in an inference mode when one can't fit a model
# into a single GPU
#
# 1. Use 1 GPU with CPU offload
# 2. Or use multiple GPUs instead
#
# First you need to install deepspeed: pip install deepspeed
#
# Here we use a 3B "bigscience/T0_3B" model which needs about 15GB GPU RAM - so 1 largish or 2
# small GPUs can handle it. or 1 small GPU and a lot of CPU memory.
#
# To use a larger model like "bigscience/T0" which needs about 50GB, unless you have an 80GB GPU -
# you will need 2-4 gpus. And then you can adapt the script to handle more gpus if you want to
# process multiple inputs at once.
#
# The provided deepspeed config also activates CPU memory offloading, so chances are that if you
# have a lot of available CPU memory and you don't mind a slowdown you should be able to load a
# model that doesn't normally fit into a single GPU. If you have enough GPU memory the program will
# run faster if you don't want offload to CPU - so disable that section then.
#
# To deploy on 1 gpu:
#
# deepspeed --num_gpus 1 t0.py
# or:
# python -m torch.distributed.run --nproc_per_node=1 t0.py
#
# To deploy on 2 gpus:
#
# deepspeed --num_gpus 2 deepspeed-inference-zero.py
# 既然指定了 GPU 的 ID，那么就不需要再设置'--num_nodes'、'--num_gpus'
# deepspeed --include=localhost:6,7 deepspeed-inference-zero.py
# or:
# python -m torch.distributed.run --nproc_per_node=2 t0.py
from transformers import AutoTokenizer, AutoConfig, AutoModelForSeq2SeqLM,AutoModelForCausalLM
transformers.integrations HfDeepSpeedConfig
deepspeed
os
torch
os.environ[] =
local_rank = (os.getenv(, ))
world_size = (os.getenv(, ))
torch.cuda.set_device(local_rank)
deespeed.init_distributed()
model_name =
config = AutoConfig.from_pretrained(model_name)
model_hidden_size = config.hidden_size
train_batch_size = * world_size
ds_config = {
: {
:
},
: {
:
},
: {
: ,
: {
: ,
:
},
: ,
: ,
: model_hidden_size * model_hidden_size,
: * model_hidden_size * model_hidden_size,
: * model_hidden_size
},
: ,
: train_batch_size,
: ,
:
}
deepspeed.zero.Init when model运行结果：
# deepspeed --num_gpus 2 t0.py
rank0:
in=Is this review positive or negative? Review: this is the best cast iron skillet you will ever buy
out=Positive
rank1:
in=Is this review positive or negative? Review: this is the worst restaurant ever
out=negative
这里需要注意的是采用多 GPU 与 ZeRO-3 组合使用进行推理生成时，需要通过在方法中设置 synced_gpus=True 同步 GPU。否则，如果一个 GPU 在另一个 GPU 之前完成推理生成，整个系统就会挂起（hangs），因为其余 GPU 不能从最先完成的 GPU 接收到权重分片。
对于 Transformers>=4.28，如果在推理生成过程中检测到多个 GPU，synced_gpus会自动设置为True。
llama.cpp 主要目标是在本地和云端的各种硬件上以最少的设置和最先进的性能实现 LLM 推理。其支持使用CPU+GPU 进行混合推理，以部分加速大于 VRAM 总容量的模型。使用 llama.cpp 之前需要将模型权重格式转换成GGUF格式。
补充：GGUF
GGUF（GPT-Generated Unified Format）是由 Georgi Gerganov（著名开源项目 llama.cpp 的创始人）定义发布的一种大模型文件格式。它是一种二进制格式文件的规范，原始的大模型预训练结果经过转换后变成 GGUF 格式可以更快地被载入使用，也会消耗更低的资源。原因在于 GGUF 采用了多种技术来保存大模型预训练结果，包括采用紧凑的二进制编码格式、优化的数据结构、内存映射等。
综上所述，GGUF 可以理解为一种格式定义，采用相应的工具将原始模型预训练结果转换成 GGUF 之后可以更加高效的使用。目前，GGUF 的发展势头强劲，HF 已经对 GGUF 提供支持，并且适配了 llama/qwen2 等模型架构。
llama.cpp 也有类似于 HuggingFace Transformers 和 DeepSpeed ZeRO-Inference 类似的功能，但是与其的稍微不同之处是：llama.cpp 是控制模型从 CPU 内存 offload 到 GPU 内存。而 HuggingFace Transformers 和 DeepSpeed ZeRO-Inference 的 offload 是从 GPU 内存 offload 到 CPU 内存。具体的控制参数如下：
--n-gpu-layers 或 -ngl：offload 模型权重层到 GPU 进行计算。0 表示仅加载到 CPU，999 表示模型全部加载到 GPU 等加速卡。--no-kv-offload 或 -nkvo：禁用 KV Cache offload 到 GPU，默认是允许 offload 到 GPU。直接使用 llama.cpp 的 Python 封装包部署模型，这里使用了 4 张 RTX 4090 部署 72B 模型，其中，将 30 个 Transoformer 层加载到 GPU 内存。具体示例如下：
from llama_cpp import Llama
import time
llm = Llama(
#model_path = "/workspace/models/Qwen2-7B-Instruct/ggml-model-f16.gguf",
model_path="/workspace/models/Qwen1.5-72B-Chat/ggml-model-f16.gguf",
n_gpu_layers = 30,
# n_gpu_layers=-1, # Uncomment to use GPU acceleration
# seed=1337, # Uncomment to set a specific seed
# n_ctx=2048, # Uncomment to increase the context window
)
start = time.time()
output = llm(
"Q:保持健康的秘诀有哪些？A: ", # Prompt
max_tokens=32, # Generate up to 32 tokens, set to None to generate up to the end of the context window
#stream=True,
stop=["Q:", "\n"], # Stop generating just before the model would generate a new question
echo=True # Echo the prompt back in the output
) # Generate a completion, can also call create_completion
print(output)
infer_time = time.time() - start
print("耗时：", infer_time)
vLLM 对于 offload 的支持相对较晚，目前 vLLM 已经在 0.5.3 版本提供了 offload 权重到 CPU 的支持，该实现与流水线并行相结合（因为 vllm 中的所有模型即将支持流水线并行，因此这应该是一个很好的选择）。通过 --cpu-offload-gb 参数控制，该参数表示每个 GPU 可以额外使用多少内存！例如，如果您有一个 24 GB GPU 并将其设置为 10，实际上您可以将其视为 34 GB GPU。然后您可以加载 BF16 权重的 13B 模型，这需要至少 26GB GPU 内存。具体示例如下：
from vllm import LLM, SamplingParams
# Sample prompts.
prompts = [
"Hello, my name is",
"The president of the United States is",
"The capital of France is",
"The future of AI is",
]
# Create a sampling params object.
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
# Create an LLM.
llm = LLM(model="meta-llama/Llama-2-13b-chat-hf", cpu_offload_gb=10)
# Generate texts from the prompts. The output is a list of RequestOutput objects
# that contain the prompt, generated text, and other information.
outputs = llm.generate(prompts, sampling_params)
# Print the outputs.
for output in outputs:
prompt = output.prompt
generated_text = output.outputs[0].text
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
TensorRT-LLM 中使用 Weight Streaming 可以将一些权重卸载到 CPU 内存，并在运行时将它们流式传输到 GPU 内存。以此来减少 GPU 内存中的权重大小，因此，可以在相同的 GPU 内存预算中运行更大的模型或更大的批量大小。
在构建期间，使用 --weight-streaming --strongly_typed --gemm_plugin disable 构建引擎，因为 Weight Streaming 仅支持强类型模型和 non-plugin 权重。
在运行时，使用 --gpu_weights_percent x 以配置 GPU 上保留的权重百分比。x 可以是 0.0 到 1.0 之间的值，默认值为 1.0。
以下是使用 Weight Streaming 运行 llama-7b 的示例：
# Convert model as normal. Assume hugging face model is in llama-7b-hf/
python3 examples/llama/convert_checkpoint.py \
--model_dir llama-7b-hf/ \
--output_dir /tmp/llama_7b/trt_ckpt/fp16/1-gpu/ \
--dtype float16
# Build engine that enabled Weight Streaming.
trtllm-build \
--checkpoint_dir /tmp/llama_7b/trt_ckpt/fp16/1-gpu/ \
--output_dir /tmp/llama_7b/trt_engines/fp16/1-gpu/ \
--weight_streaming \
--strongly_typed \
--gemm_plugin disable \
--max_batch_size 128 \
--max_input_len 512 \
--max_output_len 50
使用 20% 的权重在 GPU 内存运行引擎。
# Run the engine with 20% weights in GPU memory.
python3 examples/summarize.py \
--engine_dir /tmp/llama_7b/trt_engines/fp16/1-gpu/ \
--batch_size 1 \
--test_trt_llm \
--hf_model_dir llama-7b-hf/ \
--data_type fp16 \
--gpu_weights_percent 0.2
如果希望使用 Triton Server 部署，配置示例如下：
parameters: {
key: "gpu_weights_percent"
value: {
string_value: "0.2"
}
}
此外，TensorRT-LLM 也支持 Offload KV Cache 到 CPU 内存。Offload KV Cache 到 CPU 内存通常会与 KV Cache 重用配合使用。这增加了 KV Cache 重用的可能性。较高优先级任务所需的可重用块将被复制到 CPU 内存中的缓冲区，而不是被逐出。这极大地扩展了可重用的内存量，使块可以更长时间地保持可重用状态。
另一方面，KV Cache 块的卸载（以及 KV Cache 块被重用时的后续载入）会产生一些成本，因为块必须从 CPU 复制到 GPU 内存。但该成本在 Grace-Hopper GPU 上可以忽略不计，因为足够小，可以为具有 Hopper GPU 的 x86 机器上的许多场景产生净收益。由于 GPU 和主机（CPU）内存之间的（相对）较慢的连接，Offload 不太可能在旧架构上带来好处。
如果运行 Triton server 搭配 TensorRT-LLM 式，则可以使用 kv_cache_host_memory_bytes 参数启用主机内存卸载。下面展示了将其添加到 Triton 模型配置文件中，这将在主机内存中创建 45 GB Offload 缓冲区。
parameters: {
key: "kv_cache_host_memory_bytes"
value: {
string_value: "45000000000"
}
}
除了以上的 AI 框架中提供了像 fastllm、FlexGen、FlexFlow-Serve 中均提供了 Offload 功能，这里就不一一介绍了。
本文介绍了大模型低显存推理的一些常用方法，同时，重点介绍了 Offload 技术在各个 AI 框架中的应用。整个思想在各个 AI 框架中都大同小异，从技术实现上来说，各个 AI 框架的实现略有不同。从灵活性和可控制性上面来说，我更喜欢 Huggingface Transformers。从推理速度和易用性上来说，我觉得 LLama.cpp 更好。
微信公众号「极客日志」，在微信中扫描左侧二维码关注。展示文案：极客日志 zeeklog
使用加密算法（如AES、TripleDES、Rabbit或RC4）加密和解密文本明文。 在线工具，加密/解密文本在线工具，online
生成新的随机RSA私钥和公钥pem证书。 在线工具，RSA密钥对生成器在线工具，online
基于 Mermaid.js 实时预览流程图、时序图等图表，支持源码编辑与即时渲染。 在线工具，Mermaid 预览与可视化编辑在线工具，online
解析常见 curl 参数并生成 fetch、axios、PHP curl 或 Python requests 示例代码。 在线工具，curl 转代码在线工具，online
将字符串编码和解码为其 Base64 格式表示形式即可。 在线工具，Base64 字符串编码/解码在线工具，online
将字符串、文件或图像转换为其 Base64 表示形式。 在线工具，Base64 文件转换器在线工具，online