并行运行 70B 大模型
我们已经看到，量化已经成为在低端 GPU（比如 Colab、Kaggle 等）上加载大型语言模型（LLMs）的最常见方法了，但这会降低准确性并增加幻觉现象。
那如果你和你的朋友们把一个大型语言模型分着用呢？
每台笔记本的 GPU 负责一部分，这样大家一起推理和处理任务就容易多了。
我们要用 Petals 来做到这一点，把我们的大模型分布式加载到你朋友或家人那里的多个 GPU 上，它们会一起托管我们的模型。
LLMs 是怎么分布式的
在我们开始写代码之前，需要先理解一下分布是怎么进行的，要理解这个，用一张图来直观地展示是最合适的方式。
70B 大模型并行示意图（由 Fareed Khan 绘制）
每个 LLM 都是由很多 block（模块）组成的，它们彼此是独立的。
比如说 LLaMA 3 70B，它由好几个 block 组成，假设总共有 10 个 block。我们可以让一台笔记本托管其中几个 block，另一台托管另外几个，以此类推。
这样，多个 block 被分别托管了起来，就可以一起拿来推理使用了。
我们还能更进一步，比如用量化的方法来托管更大的 LLM，比如 405B 参数的 LLaMA，这样就能访问更大型的模型，而不需要付费 API 或 GPU 使用时间。
目录
• 设置环境
• 检查可用 GPU
• 创建分布式模型
• 在分布式模型上贪婪推理
• 正确地生成 Token
• 模型长什么样
• Prompt Tuning（提示词微调）
• 微调可训练的适配器
• 采样方法
• 私人 Swarm 网络
• 关键结论
设置环境
确保你的环境里安装了支持 CUDA 的 PyTorch。你可以从这里查看最新版。我们先安装带 CUDA 支持的 PyTorch。
# 安装带 CUDA 支持的 PyTorch
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
之后我们要安装 Petals。你可以从 GitHub 查看它的官方文档。我们先安装一下。
# 从源代码安装 Petals
pip install git+https://github.com/bigscience-workshop/petals
安装会花点时间，等依赖装完之后，我们就来检查一下可用的 GPU。
检查可用 GPU
为了模拟一群人一起合作的场景，我用 Colab、Kaggle 和 Lightning.ai 创建了好几个基于 GPU 的 Jupyter 笔记本，包括我本地的 GPU，用来做推理用。
来看一下我这边运行 (FP16) 70B LLM 的可用 GPU：
可用 GPU（由 Fareed Khan 绘制）
我用不同的临时账号，在同一台服务器上开了多个笔记本。所以如果把总 GPU 内存加在一起，大概有 (~170GB RAM)。
当然了，如果你目标是 8B 未量化的 LLM，肯定用不了这么多笔记本。
接下来，咱们就要开始把 LLM 的 workload（block）分布到不同的 GPU 上了。
创建分布式模型
想要把一些 block 推送到某个 GPU 上，只要在对应 GPU 的笔记本上运行这个命令，就能开始分布式加载 LLM。我们用 70B LLM 来操作：
# 把 70B LLaMA 的 block 推送到一个运行着的 GPU 上
python -m petals.cli.run_server meta-llama/Llama-3.1-70B-Instruct
执行这个命令后，它会先检查该笔记本的可用 GPU。
输出 = 检查可用 GPU
less
Mar 26 12:43:46.586 [INFO] 测量网络和计算吞吐量，大约需要一分钟，结果会缓存以供以后使用
Mar 26 12:43:52.596 [INFO] 推理吞吐量：每 block 每秒 478.6 tokens（1 tokens/batch，Tesla T4 GPU，bfloat16）
因为我用的是 Colab 的 T4 GPU 笔记本，所以它首先检测到了 GPU，还有每个 block 处理的 tokens/sec 等有用信息。
检测完 GPU 后，它会识别出我们 LLM 共有多少个 block（大约 80 个 block）。
输出 = 检查总 Blocks
less
Mar 26 12:43:58.461 [INFO] 前向传递吞吐量：每 block 每秒 5874.5 tokens（1024 tokens/batch，Tesla T4 GPU，bfloat16，量化到 nf4）
Mar 26 12:44:06.030 [INFO] 网络吞吐量：每秒 6320.9 tokens（下载 1158.43 Mbit/s，上传 414.25 Mbit/s）
Mar 26 12:44:06.031 [INFO] 汇报吞吐量：80 blocks 总体每秒 356.0 tokens
Mar 26 12:44:12.768 [INFO] 公告：blocks [14, 11, 6, 12, 14, 4, 7] 加入
它决定把 7 个 block 上传到我第一个 Colab T4 GPU 笔记本上。
所有这些都是通过 Petals Swarm 的后端管理的，自动处理了分布和内存管理，我们不需要自己手动管理。
不过 Petals 的文档非常强大，想深入了解的话可以去读（前面发过链接了）。
一旦对应的 block safetensors 下载完成，它们就会托管在 Petals 的公共 Swarm 服务器上。可以通过 health.petals.dev 监控你托管的 block 健康状况。
我们的 LLaMA 70B 托管的 block
可以看到，一旦那 7 个 block 成功托管在某个笔记本上，它们就出现在健康监控页面上了。
虽然单凭这一部分我们就能开始像 API 一样生成文本了，但如果托管更多的 block，生成效果肯定会更好更准。
所以，我们可以在另一个笔记本上用同样的命令继续操作，所有笔记本上流程都一样，Petals 会自动选择还没有被托管的 block。
托管 LLaMA 70B
等到所有笔记本都托管好之后，你可以看到具体托管了哪些 block，哪些 block 是在哪个服务器上。
一旦我们的 LLM 成功托管（虽然不是全部，但大部分 block 都有了），就可以像平常一样拿来推理和生成输出了。
在分布式模型上贪婪推理
我们来创建一个分布式模型，用来生成文本。这个设置下，我们本地 GPU（2GB 的机器）会下载一小部分模型权重，网络上的其他电脑负责剩下的部分。
Petals 让这一切变得很简单，它能直接和 PyTorch 还有 HF Transformers 配合使用，所以用起来就像跑本地模型一样轻松。
import torch
from transformers import AutoTokenizer
from petals import AutoDistributedModelForCausalLM
# 指定 Hugging Face Hub 上的模型名字（这里是 LLaMA 70B）
model_name = "meta-llama/Llama-3.1-70B-Instruct"
# 加载 tokenizer，并配置一些参数
tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False, add_bos_token=False)
# 加载用于因果语言建模的分布式模型
model = AutoDistributedModelForCausalLM.from_pretrained(model_name)
# 把模型移动到 GPU 上，加速计算
model = model.cuda()
如果你以前在本地用过 LLM，这段代码基本不用解释了。
这里唯一的新东西就是 AutoDistributedModelForCausalLM，它负责把本地的 LLM 接到分布式托管的模型层网络上。
接下来，我们就可以用 model.generate() 方法来生成文本了。
第一次调用时可能要花几秒钟连接到 Petals，之后生成速度大约是每秒 5–6 个 token，比 CPU 或 offloading 要快。
# 将输入文本进行 token 化并移动到 GPU
inputs = tokenizer('A boy is riding "', return_tensors="pt")["input_ids"].cuda()
# 生成输出 tokens，限制最多生成 3 个新 token
outputs = model.generate(inputs, max_new_tokens=3)
# 解码并打印生成的输出
print(tokenizer.decode(outputs[0]))
输出
A boy is riding house and
如果输出看着不太对劲，不用担心。默认情况下，model.generate() 用的是贪婪生成（greedy generation）。
你可以尝试用 top-p、top-k 采样或者束搜索（beam search）等其他方法，通过调整 Transformers 的 .generate() 方法参数实现。
正确地生成 Token
如果想让模型像聊天机器人一样实时互动，我们可以用推理会话（inference session）接口。这样生成的 token 会实时显示，非常适合做聊天机器人。
推理会话工作流（由 Fareed Khan 绘制）
这个会话会找到服务器来执行每一步，并存储注意力缓存（attention caches），所以你不需要每次都重新处理之前的 token。
如果有服务器掉线，Petals 会快速找到替代节点，只会重新生成一小部分缓存，非常高效。
来看一下，怎么一边生成一边实时显示：
# 创建一个假的 token，用来在解码时保留前导空格
fake_token = tokenizer("^")["input_ids"][0]
# 定义初始的文本提示词
text = "What is a good chatbot? Answer:"
# 把提示词 token 化并移动到 GPU
prefix = tokenizer(text, return_tensors="pt")["input_ids"].cuda()
# 开始一个推理会话，最多生成 30 个 token
with model.inference_session(max_length=30) as sess:
for i in range(20): # 迭代生成最多 20 个 token
# 只在第一次生成时提供 prefix
inputs = prefix if i == 0 else None
# 生成 1 个 token，使用采样（temperature=0.9，top-p=0.6 增加多样性）
outputs = model.generate(
inputs, max_new_tokens=1, session=sess,
do_sample=True, temperature=0.9, top_p=0.6
)
# 解码新生成的 token，并追加到文本后面
text += tokenizer.decode([fake_token, outputs[0, -1].item()])[1:]
# 每生成一个 token 就打印一次文本
print(text)
总结一下，我们刚才做了什么：
• 做了一个小技巧，确保解码时保留前导空格。
• 把提示词 token 化并移到 GPU。
• 开启推理会话，最多生成 20 个 token。
• 用采样（temperature 和 top-p）方式提高输出多样性。
• 每生成一个 token 都追加到文本，并实时打印。
接下来，我们来创建一个简单的聊天机器人！加个循环，接受用户输入，生成回复，遇到换行符（\n）就停止。
# 开启一个推理会话，最多生成 512 个 token
with model.inference_session(max_length=512) as sess:
while True:
# 获取用户输入作为聊天提示
prompt = input('Human: ')
# 如果输入为空，退出循环
if prompt == "":
break
# 格式化输入，模拟对话
prefix = f"Human: {prompt}\nFriendly AI:"
# 把输入 token 化并移动到 GPU
prefix = tokenizer(prefix, return_tensors="pt")["input_ids"].cuda()
print("Friendly AI:", end="", flush=True)
while True:
# 使用采样生成一个 token，保证回答多样性
outputs = model.generate(prefix, max_new_tokens=1, session=sess,
do_sample=True, temperature=0.9, top_p=0.6)
# 解码新生成的 token，同时保留前导空格
outputs = tokenizer.decode([fake_token, outputs[0, -1].item()])[1:]
# 立即打印新生成的 token
print(outputs, end="", flush=True)
# 如果检测到换行或结束标记，则退出
if "\n" in outputs or "</s>" in outputs:
break
# 之后不再提供 prefix，继续从会话状态生成
prefix = None
输出
Human: Hi, how are you?
Friendly AI: I am fine, thanks. And you?
你可以看到，实时推理的输出更加复杂和正式，跟我们原本的 LLM 保持一致。
如果托管的 block 更多，效果肯定会更接近原版 LLM。
我们的模型长什么样
我们用的模型其实和原版是一样的，只是部分加载到了本地 GPU 上。来看看它的结构：
# 打印模型架构
print(model)
输出
DistributedLlamaForCausalLM(
(model): DistributedLlamaModel(
(embed_tokens): Embedding(32000, 8192, padding_idx=0)
(layers): RemoteSequential(modules=llama-3-70b-intstruct.0..llama-3-70b-intstruct.79)
(norm): LlamaRMSNorm()
)
(lm_head): LMHead()
)
词嵌入（word embeddings）和一部分层是作为普通的 PyTorch 模块在本地跑的，剩下的模型部分（像 Transformer blocks）是通过 RemoteSequential 类跑在其他机器上的。
RemoteSequential 是一个特别的 PyTorch 模块，可以让模型分布式执行。
我们还能访问到单独的层，查看它们的输出，甚至能单独执行前向（forward）或反向（backward）传播。
比如提取前 5 层：
# 从模型内部的层堆栈中提取前五层
first_five_layers = model.model.layers[0:5]
# 显示提取出来的层（打印它们的细节或结构）
first_five_layers
输出
RemoteSequential(modules=llama-3-70b-intstruct.0..llama-3-70b-intstruct.4)
它打印出来了托管在其他 GPU 上的前 5 个 block，你可以直接看到。
Prompt Tuning（提示词微调）
远程托管的 transformer blocks 是冻结的（frozen），为了保证预训练模型在所有用户之间保持一致。
不过我们还是可以通过参数高效的方法（比如可训练的 prompt 或适配器 LoRA）来微调。
因为所有可训练参数和优化器都保存在本地，所以不会影响到其他用户。
在这个例子里，我们要用可训练的 prompt 来做一个简单任务：让模型学会把一句话变成它的反义句。
比如：
从 “A small white rabbit hops across the green field.”
变成
“A small white rabbit did not hop across the green field”。
首先来看一下在没有微调前，模型怎么回答：
# token 化输入文本并移动到 GPU
inputs = tokenizer("A small white rabbit ", return_tensors="pt")["input_ids"].cuda()
# 用模型生成最多 7 个新 token
outputs = model.generate(inputs, max_new_tokens=7)
# 解码并打印生成的文本
print("Generated:", tokenizer.decode(outputs[0]))
输出
A small white rabbit hop across the green field
如果你对 Prompt Tuning 不熟悉，简单说就是：在输入前面加一些可训练的“token”，这些 token 也是模型参数的一部分。
更进一步的 Deep Prompt Tuning（深度提示词微调）是在每一个 Transformer block 前都加可训练的 token，这样可以增加可训练参数数量，提高性能。
Deep Prompt Tuning（图片来自 Liu, Xiao 等人）
Petals 内置了 prompt tuning 的支持，参数 tuning_mode="ptune"；也支持 deep prompt tuning，参数 tuning_mode="deep_ptune"。
我们这里用 deep prompt tuning，每层加 3 个 token（pre_seq_len=3）：
# 用 deep prompt tuning 加载预训练的因果语言模型
model = AutoDistributedModelForCausalLM.from_pretrained(model_name, tuning_mode='deep_ptune', pre_seq_len=3)
# 把模型移到 GPU 上
model = model.cuda()
接下来，我们可以像本地模型一样设置 Adam 优化器，开始微调分布式模型：
# 用学习率 0.001 初始化 Adam 优化器
opt = torch.optim.Adam(model.parameters(), lr=1e-3)
# token 化目标文本并移到 GPU
the_rabbit_did_not_hop = tokenizer("A small white rabbit did not hop across the green field", return_tensors="pt")["input_ids"].cuda()
# 训练 12 步
for i in range(12):
# 前向传递，计算 loss（自监督学习）
loss = model(input_ids=the_rabbit_did_not_hop, labels=the_rabbit_did_not_hop).loss
print(f"loss[{i}] = {loss.item():.3f}")
# 清空前一次的梯度
opt.zero_grad()
# 反向传播计算梯度
loss.backward()
# 优化器更新参数
opt.step()
print("opt.step()")
输出
loss[0] = 5.324
opt.step()
loss[1] = 4.983
opt.step()
loss[2] = 4.512
opt.step()
loss[3] = 4.217
opt.step()
...
loss[11] = 0.014
opt.step()
当 loss 几乎降到 0 时，模型就学会了我们要它生成的句子。
再来测试一下：
# token 化输入 "A small white rabbit" 并移动到 GPU
inputs = tokenizer("A small white rabbit", return_tensors="pt")["input_ids"].cuda()
# 生成最多 7 个新 token
outputs = model.generate(inputs, max_new_tokens=7)
# 解码并打印生成的文本
print("Generated:", tokenizer.decode(outputs[0]))
输出
generated: A small white rabbit did not hop across the green field.
可以看到，它成功生成了反义句：“小白兔没有跳过绿色草地”。
微调可训练适配器
接下来，我们要试一下另一种微调方法，叫做可训练适配器。
适配器就是在预训练 Transformer block 之间或者旁边加的小型层。
我们会在模型里插入一个简单的可训练线性层，还会替换掉模型的头部（head），做分类任务。
像之前一样，适配器的权重和优化器状态也都是保存在本地的。
import torch.nn as nn
import torch.nn.functional as F
# 加载一个预训练的因果语言模型，并移动到 GPU
model = AutoDistributedModelForCausalLM.from_pretrained(model_name)
model = model.cuda()
# 定义一个基于 LLM 的分类器
class LLMBasedClassifier(nn.Module):
def __init__(self, model):
super().__init__()
# 提取模型的 transformer 层，用于分布式处理
self.distributed_layers = model.transformer.h
# 添加一个适配器模块：两个线性层，用来降维再升维
self.adapter = nn.Sequential(
nn.Linear(model.config.hidden_size, 32),
nn.Linear(32, model.config.hidden_size)
)
# 分类头，把隐藏状态映射成 2 个类别
self.head = nn.Linear(model.config.hidden_size, 2)
def forward(self, embeddings):
# 把模型一分为二，先通过前半部分层处理 embedding
mid_block = len(self.distributed_layers) // 2
hidden_states = self.distributed_layers[:mid_block](embeddings)
# 经过适配器进行转换
hidden_states = self.adapter(hidden_states)
# 再通过后半部分层处理
hidden_states = self.distributed_layers[mid_block:](hidden_states)
# 对序列做均值池化（mean pooling），得到固定大小的表示
pooled_states = torch.mean(hidden_states, dim=1)
# 最后通过分类头输出类别预测
return self.head(pooled_states)
然后，我们用 Adam 优化器训练这个模型，用常见的交叉熵损失（cross-entropy loss）来做分类任务：
# 初始化 LLMBasedClassifier 并移动到 GPU
classifier = LLMBasedClassifier(model).cuda()
# 用 Adam 优化器，学习率 3e-5
opt = torch.optim.Adam(classifier.parameters(), 3e-5)
# 创建虚拟输入数据（3 个样本，每个 2 个 token，隐藏层大小）
inputs = torch.randn(3, 2, model.config.hidden_size, device='cuda')
# 定义这 3 个样本的真实标签（0 和 1）
labels = torch.tensor([1, 0, 1], device='cuda')
# 训练 5 次迭代
for i in range(5):
# 计算损失：交叉熵
loss = F.cross_entropy(classifier(inputs), labels)
print(f"loss[{i}] = {loss.item():.3f}")
# 清空之前的梯度
opt.zero_grad()
# 反向传播计算梯度
loss.backward()
# 用优化器更新参数
opt.step()
# 打印预测的类别（通过 argmax 得到）
print('Predicted:', classifier(inputs).argmax(-1))
当我们开始训练时，它会在每次迭代后打印 loss：
loss[0] = 16.236
...
loss[4] = 1.254
predicted: tensor([1, 0, 1], device='cuda:0')
可以看到，loss 在不断下降，说明模型正在拟合我们的虚拟数据集！
就像 Petals 官方文档里说的那样，你还可以去看看他们的示例笔记本，在里面用 Llama 在著名的 SST2 数据集上做微调。
采样方法
之前你已经看到怎么用普通 PyTorch 代码来跟分布式模型交互了。
这种方式允许你实现更高级的微调和采样方法，很多时候是托管 API 做不到的。
现在，我们来从零开始自己写一个采样方法。
下面，我们要重新实现标准的 model.generate() 接口，手动一步步在所有层上做前向传递（forward pass）：
from hivemind import get_logger
import torch
# 初始化 logger，用来在每一步打印日志
logger = get_logger()
# 小技巧，确保 tokenizer.decode() 保留前导空格
fake_token = tokenizer("^")["input_ids"][0]
# 定义输入提示词
text = "How can I improve my writing skills? Answer:"
token_ids = tokenizer(text, return_tensors="pt")["input_ids"].cuda() # 把输入文本转成 token ID 并移到 GPU
# 设置生成文本的最大长度
max_length = 100
# 禁用梯度计算，加速推理
with torch.inference_mode():
# 开启推理会话
with model.inference_session(max_length=max_length) as sess:
# 一直生成 token，直到达到最大长度
while len(text) < max_length:
# 获取 token 的词嵌入
embs = model.transformer.word_embeddings(token_ids)
embs = model.transformer.word_embeddings_layernorm(embs)
# 前向传递
h = sess.step(embs)
h_last = model.transformer.ln_f(h[:, -1]) # 取最后一个隐藏状态
logits = model.lm_head(h_last) # 计算下一步的 logits
# 贪婪地选择概率最高的 token
next_token = logits.argmax(dim=-1)
# 解码新 token，追加到输出文本里
text += tokenizer.decode([fake_token, next_token.item()])[1:]
# 更新 token_ids，准备下一步生成
token_ids = next_token.reshape(1, 1)
# 在每一步打印生成的文本
logger.info(text)
运行上面这段代码时，它会开始一边生成一边打印输出：
[INFO] How can I improve my writing skills? Answer: P
[INFO] How can I improve my writing skills? Answer: Practice
[INFO] How can I improve my writing skills? Answer: Practice writing
[INFO] How can I improve my writing skills? Answer: Practice writing regularly
[INFO] How can I improve my writing skills? Answer: Practice writing regularly and
[INFO] How can I improve my writing skills? Answer: Practice writing regularly and read
[INFO] How can I improve my writing skills? Answer: Practice writing regularly and read more
[INFO] How can I improve my writing skills? Answer: Practice writing regularly and read more books
[INFO] How can I improve my writing skills? Answer: Practice writing regularly and read more books to
[INFO] How can I improve my writing skills? Answer: Practice writing regularly and read more books to improve
[INFO] How can I improve my writing skills? Answer: Practice writing regularly and read more books to improve your
[INFO] How can I improve my writing skills? Answer: Practice writing regularly and read more books to improve your vocabulary
...
Petals 的 model.inference_session() 接口让你可以自己写定制的推理逻辑。
通过它，你可以实现任何采样算法，比如自定义的束搜索（beam search），比如避免生成脏话之类的。
私人 Swarm 网络（Private Swarm）
Petals 的公共 Swarm 允许你跟别人共享计算资源，一起跑模型。
但有时候出于隐私或者输出正确性的要求，公共网络就不太合适了。
那如果你想搭建自己的私有托管网络呢？
只让自己和朋友访问大模型的 blocks？
这个很简单，我们来搞一个私人 Swarm。
首先，你需要一台可靠的机器（或者几台）作为私有网络的“入口点”或者“地址簿”。
这个机器不需要 GPU，我们叫它 bootstrap peer。
你可以用 tmux 或 screen 把它跑在后台。命令如下：
# 启动一个 bootstrap peer，监听特定端口（比如 31337）
# 把它的唯一 ID 保存到 bootstrap1.id
python -m petals.cli.run_dht --host_maddrs /ip4/0.0.0.0/tcp/31337 --identity_path bootstrap1.id
运行后注意看输出！
它会打印出一行提示你完整的地址，比如：
输出
[INFO] 正在运行 DHT 实例。要让其他节点连接到这个节点，
请使用 --initial_peers /ip4/YOUR_PUBLIC_IP_OR_LAN_IP/tcp/31337/p2p/QmTPAIfTh1sIsMyUnique1DDontCopyThisPart...
一定要复制这个完整的 /ip4/.../p2p/... 地址！
这就是你私人 Swarm 的钥匙。（确保你的公网 IP 或局域网 IP 是其他机器能访问到的。）
然后，在你的 GPU 机器上（笔记本、Colab、其他设备上），像之前一样启动 Petals server，
但是这次要告诉它，去连接你自己的私人网络：
# 把 bootstrap peer 的地址存到变量里
export MY_INITIAL_PEERS="/ip4/YOUR_PUBLIC_IP_OR_LAN_IP/tcp/31337/p2p/QmTPAIfTh1sIsMyUnique1DDontCopyThisPart..."
# 启动 LLaMA 3.1 70B 的服务器，连接到你的私人 Swarm
python -m petals.cli.run_server meta-llama/Llama-3.1-70B-Instruct --initial_peers $MY_INITIAL_PEERS
每台你想贡献 GPU 的机器，都要这么操作。
它们会像之前那样加载 blocks，但只在你的私有网络里互相通信，不会连接到公网上。
最后，当你在 Python 脚本里推理时，也要告诉它去连接你的私人 Swarm：
import torch
from transformers import AutoTokenizer
from petals import AutoDistributedModelForCausalLM
# 你的 bootstrap peer 的地址（可以是列表）
INITIAL_PEERS = ["/ip4/YOUR_PUBLIC_IP_OR_LAN_IP/tcp/31337/p2p/QmTPAIfTh1sIsMyUnique1DDontCopyThisPart..."]
# 指定模型名
model_name = "meta-llama/Llama-3.1-70B-Instruct"
# 正常加载 tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False, add_bos_token=False)
# 加载分布式模型，指定连接你的私人 Swarm
model = AutoDistributedModelForCausalLM.from_pretrained(
model_name,
initial_peers=INITIAL_PEERS # <-- 重点就在这！
)
# 如果本地有 GPU，可以把模型移过去
model = model.cuda()
# 然后就可以像之前一样用 model.generate() 或推理会话啦！
# inputs = tokenizer('A question: "', return_tensors="pt")["input_ids"].cuda()
# outputs = model.generate(inputs, max_new_tokens=5)
# print(tokenizer.decode(outputs[0]))
就这样！
你现在就在自己的一小圈电脑里，私下运行超大 LLM 了。
不再依赖公共网络，完全自主可控。
甚至可以搭建自己的健康监控页面，实时查看你们 Swarm 的状态！
你可以通过阅读 Petals 的官方文档来进一步掌握更多高级用法哦。
关键结论
你已经看到了，prompt 模板和托管 block 的数量是影响最终效果的两个关键因素；
否则 LLM 的性能可能达不到预期水平。
你还可以用量化技术来加载更大的 LLM。
Petals 支持很多参数配置，比如指定要托管哪些 block，避免重复上传。
我个人建议你可以先托管一个 8B 的 LLM，测试一下它的表现，
肯定会让你大吃一惊的！