如何手动执行BLOOM模型推理？（手动加载BLOOM块至VRAM）

我们项目的目标之一是将BLOOM模型分布式部署在多台计算机上，目前已解决计算机间的通信问题，但不清楚如何手动执行推理。掌握该方法后，我们就能构建计算机“链”来加速整体计算。

手动执行推理对我们至关重要，实现后我们会基于主机GPU优化链节点的选择流程。

简化假设（后续版本会调整）：

• 分布式架构接收的提示词按顺序处理
• 计算机数量及参数无限制
• 所有计算机的GPU型号一致
• 每台计算机仅配备一块GPU
• 每块GPU加载一个BLOOM块

计算机配置： #

• 操作系统：Windows 10 Home
• 处理器：Intel i7-11700K
• 内存：32GB
• 显卡：NVIDIA RTX 3060 12GB VRAM

已安装包（pip freeze）： #

accelerate==0.16.0
asttokens==2.2.1
backcall==0.2.0 
certifi==2022.12.7
charset-normalizer==3.0.1
colorama==0.4.6
comm==0.1.2
debugpy==1.6.6
decorator==5.1.1
executing==1.2.0
filelock==3.9.0
huggingface-hub==0.12.0
idna==3.4
ipykernel==6.21.1
ipython==8.9.0
jedi==0.18.2
jupyter_client==8.0.2
jupyter_core==5.2.0
matplotlib-inline==0.1.6
nest-asyncio==1.5.6
numpy==1.24.1
packaging==23.0
pandas==1.5.3
parso==0.8.3
pickleshare==0.7.5
Pillow==9.4.0
platformdirs==2.6.2
prompt-toolkit==3.0.36
psutil==5.9.4
pure-eval==0.2.2
Pygments==2.14.0
python-dateutil==2.8.2
pytz==2022.7.1
pywin32==305
PyYAML==6.0
pyzmq==25.0.0
regex==2022.10.31
requests==2.28.2
safetensors==0.2.8
six==1.16.0
stack-data==0.6.2
tokenizers==0.13.2
torch==1.13.1+cu117
torchaudio==0.13.1+cu117
torchvision==0.14.1+cu117
tornado==6.2
tqdm==4.64.1
traitlets==5.9.0
transformers==4.26.0
typing_extensions==4.4.0
urllib3==1.26.14
wcwidth==0.2.6

仓库结构（最简形式）： #

bloom
   ↳ bigscience/bloom仓库文件
bloom-7b1
   ↳ bigscience/bloom-7b1仓库文件
venv
   ↳ ...（虚拟环境文件）
main.py

代码（main.py）： #

from pprint import pprint

import torch

from transformers import AutoTokenizer, BloomConfig
from transformers.models.bloom.modeling_bloom import BloomBlock as BloomBlock

tokenizer = AutoTokenizer.from_pretrained(r'E:\programming\python\bloom-playground\models\bloom')

prompt = "Hi, I am TapMadl and I want to "
input_ids = tokenizer(prompt, return_tensors="pt").to(0)

configuration = BloomConfig('bloom\config.json')
pprint(configuration.num_hidden_layers)

modules = torch.nn.ModuleList([BloomBlock(configuration) for _ in range(configuration.num_hidden_layers)])

pprint(type(modules[0]))
pprint(len(modules))

module1 = modules[0]

module1.forward(input_ids)

代码输出： #

2
<class 'transformers.models.bloom.modeling_bloom.BloomBlock'>
2
Traceback (most recent call last):
  File "c:\programming\python\bloom-prototyping\main.py", line 23, in <module>
    module1.forward(input_ids)
TypeError: BloomBlock.forward() missing 2 required positional arguments: 'alibi' and 'attention_mask'

代码说明： #

最初我使用bigscience/bloom仓库示例代码中的AutoTokenizer和AutoModelForCausalLM流水线，希望找到一种可逐块执行推理的简便方法，但未能实现。

随后我通过调试追踪函数调用，研究手动推理的最小可行方案（MVP），发现AutoModelForCausalLM类实际逻辑复杂。通过在Jupyter Notebook中调试和修改代码，得到了上述代码。

希望有人能指导我实现手动推理，感谢帮助。

内容的提问来源于stack exchange，提问作者TapMadl