使用Tokenizer时Deepspeed张量并行出现张量对齐问题

我有多台小型GPU，单台无法容纳完整的Starcoder模型，因此尝试使用Deepspeed对其进行张量并行处理。执行以下代码：

from transformers import AutoModelForCausalLM, AutoTokenizer
import os
import torch
import deepspeed

local_rank = int(os.getenv('LOCAL_RANK', '0'))
world_size = int(os.getenv('WORLD_SIZE', '1'))

cache_dir = '/llm-benchmark/starcoder-cache'

os.environ['TRANSFORMERS_CACHE'] = cache_dir

checkpoint = "bigcode/starcoder"
device = "cuda" # for GPU usage or "cpu" for CPU usage

tokenizer = AutoTokenizer.from_pretrained(checkpoint, cache_dir=cache_dir)

# Load model without moving it to device
model = AutoModelForCausalLM.from_pretrained(checkpoint, cache_dir=cache_dir)

ds_engine = deepspeed.init_inference(model, tensor_parallel={'enabled': True, 'tp_size': world_size})
model = ds_engine.module

print('before tokenizing')
inputs = tokenizer.encode("def print_hello_world():", return_tensors="pt").to(f"{device}")
print('before generation')
outputs = model.generate(inputs)
print('after generation')
print(tokenizer.decode(outputs[0]))
print('full result')

运行时模型看似已成功拆分，但出现如下错误：

Traceback (most recent call last):
  File "/root/code/starcoder/generate.py", line 29, in <module>
    outputs = model.generate(inputs)
  File "/root/code/starcoder/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/root/code/starcoder/lib/python3.10/site-packages/transformers/generation/utils.py", line 1437, in generate
    return self.greedy_search(
  File "/root/code/starcoder/lib/python3.10/site-packages/transformers/generation/utils.py", line 2248, in greedy_search
    outputs = self(
  File "/root/code/starcoder/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/root/code/starcoder/lib/python3.10/site-packages/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py", line 808, in forward
    transformer_outputs = self.transformer(
  File "/root/code/starcoder/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/root/code/starcoder/lib/python3.10/site-packages/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py", line 673, in forward
    outputs = block(
  File "/root/code/starcoder/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/root/code/starcoder/lib/python3.10/site-packages/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py", line 316, in forward
    attn_outputs = self.attn(
  File "/root/code/starcoder/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/root/code/starcoder/lib/python3.10/site-packages/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py", line 230, in forward
    query, key_value = self.c_attn(hidden_states).split((self.embed_dim, 2 * self.kv_dim), dim=2)
  File "/root/code/starcoder/lib/python3.10/site-packages/torch/_tensor.py", line 803, in split
    return torch._VF.split_with_sizes(self, split_size, dim)
RuntimeError: split_with_sizes expects split_sizes to sum exactly to 1600 (input tensor's size at dimension 2), but got split_sizes=[1536, 256]

看起来Tokenizer与模型未对齐，请问该问题的成因是什么？

• 张量并行后模型属性未同步更新：Starcoder原始模型的embed_dim为1536，kv_dim为128，2*kv_dim是256，两者总和1792。但Deepspeed做张量并行时会将模型层的张量拆分到多个GPU上，导致实际输出的张量维度变为1600，而模型实例的embed_dim、kv_dim还是原始值，forward阶段用这些旧值计算的拆分尺寸和实际张量维度不匹配，触发报错。
• 模型加载与并行初始化的顺序问题：你先加载完整模型再初始化Deepspeed推理并行，这种流程下Deepspeed仅修改模型层的参数结构，但不会自动更新模型的属性字段，导致属性值和实际并行后的张量维度脱节。
• 和Tokenizer无关：错误本质是模型并行后的维度计算冲突，和Tokenizer的编码、解码逻辑完全无关，你之前的猜测方向有误。

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