使用Accelerate庫在多GPU上進行LLM推理
作者：Geronimo
大型語言模型(llm)已經徹底改變了自然語言處理領域。隨著這些模型在規模和復雜性上的增長，推理的計算需求也顯著增加。為了應對這一挑戰利用多個gpu變得至關重要。
大型語言模型(llm)已經徹底改變了自然語言處理領域。隨著這些模型在規模和復雜性上的增長，推理的計算需求也顯著增加。為了應對這一挑戰利用多個gpu變得至關重要。
所以本文將在多個gpu上并行執行推理，主要包括：Accelerate庫介紹，簡單的方法與工作代碼示例和使用多個gpu的性能基準測試。
本文將使用多個3090將llama2-7b的推理擴展在多個GPU上
基本示例
我們首先介紹一個簡單的示例來演示使用Accelerate進行多gpu“消息傳遞”。
from accelerate import Accelerator
from accelerate.utils import gather_object
accelerator = Accelerator()
# each GPU creates a string
message=[ f"Hello this is GPU {accelerator.process_index}" ]
# collect the messages from all GPUs
messages=gather_object(message)
# output the messages only on the main process with accelerator.print()
accelerator.print(messages)輸出如下：
['Hello this is GPU 0',
'Hello this is GPU 1',
'Hello this is GPU 2',
'Hello this is GPU 3',
'Hello this is GPU 4']多GPU推理
下面是一個簡單的、非批處理的推理方法。代碼很簡單，因為Accelerate庫已經幫我們做了很多工作，我們直接使用就可以：
from accelerate import Accelerator
from accelerate.utils import gather_object
from transformers import AutoModelForCausalLM, AutoTokenizer
from statistics import mean
import torch, time, json
accelerator = Accelerator()
# 10*10 Prompts. Source: https://www.penguin.co.uk/articles/2022/04/best-first-lines-in-books
prompts_all=[
"The King is dead. Long live the Queen.",
"Once there were four children whose names were Peter, Susan, Edmund, and Lucy.",
"The story so far: in the beginning, the universe was created.",
"It was a bright cold day in April, and the clocks were striking thirteen.",
"It is a truth universally acknowledged, that a single man in possession of a good fortune, must be in want of a wife.",
"The sweat wis lashing oafay Sick Boy; he wis trembling.",
"124 was spiteful. Full of Baby's venom.",
"As Gregor Samsa awoke one morning from uneasy dreams he found himself transformed in his bed into a gigantic insect.",
"I write this sitting in the kitchen sink.",
"We were somewhere around Barstow on the edge of the desert when the drugs began to take hold.",
] * 10
# load a base model and tokenizer
model_path="models/llama2-7b"
model = AutoModelForCausalLM.from_pretrained(
model_path,
device_map={"": accelerator.process_index},
torch_dtype=torch.bfloat16,
)
tokenizer = AutoTokenizer.from_pretrained(model_path)
# sync GPUs and start the timer
accelerator.wait_for_everyone()
start=time.time()
# divide the prompt list onto the available GPUs
with accelerator.split_between_processes(prompts_all) as prompts:
# store output of generations in dict
results=dict(outputs=[], num_tokens=0)
# have each GPU do inference, prompt by prompt
for prompt in prompts:
prompt_tokenized=tokenizer(prompt, return_tensors="pt").to("cuda")
output_tokenized = model.generate(**prompt_tokenized, max_new_tokens=100)[0]
# remove prompt from output
output_tokenized=output_tokenized[len(prompt_tokenized["input_ids"][0]):]
# store outputs and number of tokens in result{}
results["outputs"].append( tokenizer.decode(output_tokenized) )
results["num_tokens"] += len(output_tokenized)
results=[ results ] # transform to list, otherwise gather_object() will not collect correctly
# collect results from all the GPUs
results_gathered=gather_object(results)
if accelerator.is_main_process:
timediff=time.time()-start
num_tokens=sum([r["num_tokens"] for r in results_gathered ])
print(f"tokens/sec: {num_tokens//timediff}, time {timediff}, total tokens {num_tokens}, total prompts {len(prompts_all)}")使用多個gpu會導致一些通信開銷:性能在4個gpu時呈線性增長，然后在這種特定設置中趨于穩定。當然這里的性能取決于許多參數，如模型大小和量化、提示長度、生成的令牌數量和采樣策略，所以我們只討論一般的情況
1 GPU: 44個token /秒，時間:225.5s
2 gpu: 88個token /秒，時間:112.9s
3 gpu: 128個token /秒，時間:77.6s
4 gpu: 137個token /秒，時間:72.7s
5 gpu: 119個token /秒，時間:83.8s
在多GPU上進行批處理
現實世界中，我們可以使用批處理推理來加快速度。這會減少GPU之間的通訊，加快推理速度。我們只需要增加prepare_prompts函數將一批數據而不是單條數據輸入到模型即可：
from accelerate import Accelerator
from accelerate.utils import gather_object
from transformers import AutoModelForCausalLM, AutoTokenizer
from statistics import mean
import torch, time, json
accelerator = Accelerator()
def write_pretty_json(file_path, data):
import json
with open(file_path, "w") as write_file:
json.dump(data, write_file, indent=4)
# 10*10 Prompts. Source: https://www.penguin.co.uk/articles/2022/04/best-first-lines-in-books
prompts_all=[
"The King is dead. Long live the Queen.",
"Once there were four children whose names were Peter, Susan, Edmund, and Lucy.",
"The story so far: in the beginning, the universe was created.",
"It was a bright cold day in April, and the clocks were striking thirteen.",
"It is a truth universally acknowledged, that a single man in possession of a good fortune, must be in want of a wife.",
"The sweat wis lashing oafay Sick Boy; he wis trembling.",
"124 was spiteful. Full of Baby's venom.",
"As Gregor Samsa awoke one morning from uneasy dreams he found himself transformed in his bed into a gigantic insect.",
"I write this sitting in the kitchen sink.",
"We were somewhere around Barstow on the edge of the desert when the drugs began to take hold.",
] * 10
# load a base model and tokenizer
model_path="models/llama2-7b"
model = AutoModelForCausalLM.from_pretrained(
model_path,
device_map={"": accelerator.process_index},
torch_dtype=torch.bfloat16,
)
tokenizer = AutoTokenizer.from_pretrained(model_path)
tokenizer.pad_token = tokenizer.eos_token
# batch, left pad (for inference), and tokenize
def prepare_prompts(prompts, tokenizer, batch_size=16):
batches=[prompts[i:i + batch_size] for i in range(0, len(prompts), batch_size)]
batches_tok=[]
tokenizer.padding_side="left"
for prompt_batch in batches:
batches_tok.append(
tokenizer(
prompt_batch,
return_tensors="pt",
padding='longest',
truncatinotallow=False,
pad_to_multiple_of=8,
add_special_tokens=False).to("cuda")
)
tokenizer.padding_side="right"
return batches_tok
# sync GPUs and start the timer
accelerator.wait_for_everyone()
start=time.time()
# divide the prompt list onto the available GPUs
with accelerator.split_between_processes(prompts_all) as prompts:
results=dict(outputs=[], num_tokens=0)
# have each GPU do inference in batches
prompt_batches=prepare_prompts(prompts, tokenizer, batch_size=16)
for prompts_tokenized in prompt_batches:
outputs_tokenized=model.generate(**prompts_tokenized, max_new_tokens=100)
# remove prompt from gen. tokens
outputs_tokenized=[ tok_out[len(tok_in):]
for tok_in, tok_out in zip(prompts_tokenized["input_ids"], outputs_tokenized) ]
# count and decode gen. tokens
num_tokens=sum([ len(t) for t in outputs_tokenized ])
outputs=tokenizer.batch_decode(outputs_tokenized)
# store in results{} to be gathered by accelerate
results["outputs"].extend(outputs)
results["num_tokens"] += num_tokens
results=[ results ] # transform to list, otherwise gather_object() will not collect correctly
# collect results from all the GPUs
results_gathered=gather_object(results)
if accelerator.is_main_process:
timediff=time.time()-start
num_tokens=sum([r["num_tokens"] for r in results_gathered ])
print(f"tokens/sec: {num_tokens//timediff}, time elapsed: {timediff}, num_tokens {num_tokens}")可以看到批處理會大大加快速度。
1 GPU: 520 token /sec，時間:19.2s
2 gpu: 900 token /sec，時間:11.1s
3 gpu: 1205個token /秒，時間:8.2s
4 gpu: 1655 token /sec，時間:6.0s
5 gpu: 1658 token /sec，時間:6.0s
總結
截止到本文為止，llama.cpp，ctransformer還不支持多GPU推理，好像llama.cpp在6月有個多GPU的merge，但是我沒看到官方更新，所以這里暫時確定不支持多GPU。如果有小伙伴確認可以支持多GPU請留言。
huggingface的Accelerate包則為我們使用多GPU提供了一個很方便的選擇，使用多個GPU推理可以顯著提高性能，但gpu之間通信的開銷隨著gpu數量的增加而顯著增加。
責任編輯：華軒
來源：
DeepHub IMBA
