利用 tensor-parallel 把模型訓練與推理的 workload 平均分布到多塊 GPU，一方面可以提高推理速度，另一方面 vram 的負載平衡也讓復雜的 prompt 能被輕松處理。
import 相關的 libs：
# torch version 2.0.0
import torch
# tensor-parallel version 1.0.22
from tensor_parallel import TensorParallelPreTrainedModel
# transformer version 4.28.0.dev0
from transformers import LlamaTokenizer, LlamaForCausalLM, GenerationConfig
加載 LLaMA-7B 并轉化為 TensorParallelPreTrainedModel：
model = LlamaForCausalLM.from_pretrained("./llama-7b-hf", torch_dtype=torch.float16)
model = TensorParallelPreTrainedModel(model, ["cuda:0", "cuda:1"])
加載 tokenizer 并進行推理：
tokenizer = LlamaTokenizer.from_pretrained("./llama-7b-hf")tokens = tokenizer("Hi, how are you?", return_tensors="pt")
tokenizer.decode(model.generate(tokens["input_ids"].cuda(0), attention_mask=tokens["attention_mask"].cuda(0))[0])# 輸出：
# 'Hi, how are you? I'm a 20 year old girl from the Netherlands'tokens = tokenizer("Once upon a time, there was a lonely computer ", return_tensors="pt")
tokenizer.decode(model.generate(tokens["input_ids"].cuda(0), attention_mask=tokens["attention_mask"].cuda(0), max_length=256)[0])# 輸出：
# 'Once upon a time, there was a lonely computer. It was a very old computer, and it had been sitting in a box for a long time. It was very sad, because it had no friends.\nOne day, a little girl came to the computer. She was very nice, and she said, “Hello, computer. I’m going to be your friend.”\nThe computer was very happy. It said, “Thank you, little girl. I’m very happy to have you as my friend.”\nThe little girl said, “I’m going to call you ‘Computer.’”\n“That’s a good name,” said Computer.\nThe little girl said, “I’m going to teach you how to play games.”\n“That’s a good idea,” said Computer.\nThe little girl said, “I’m going to teach you how to do math.”\nThe little girl said, “I’m going to teach you how to write stories.”\nThe little girl said, “I’m going to teach you how to draw pictures.”\nThe little girl said, “I’m going to teach you how to play music.”\nThe little girl said, “I’m'
在這里，我們把我們的推理邏輯平均分布到了兩塊 GPU 上。
tensor parallel 在主流的推理框架已經很好的支持了，vLLM 和 lightllm 都是很好的選擇。
)
)
