Support for GPT2-XL

[snoop2head]

Thank you for the great project!

How to reproduce

https://github.com/snoop2head/Language_Model_Memorization/blob/2c5db6f9bdd0206cba87d13b158d8c27ce0e55a7/parallel_inference.py#L39-L82

• Tested and works for gpt2, gpt2-medium, gpt2-large
• If AutoModelForCausalLM is changed into gpt2-xl, it yields the following error message

  File "/home/ubuntu/anaconda3/lib/python3.9/site-packages/parallelformers/parallel/process.py", line 193, in inference
  outputs = function_(
  File "/home/ubuntu/anaconda3/lib/python3.9/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
  return func(*args, **kwargs)
  File "/home/ubuntu/anaconda3/lib/python3.9/site-packages/transformers/generation_utils.py", line 1294, in generate
  return self.greedy_search(
  File "/home/ubuntu/anaconda3/lib/python3.9/site-packages/transformers/generation_utils.py", line 1689, in greedy_search
  outputs = self(
  File "/home/ubuntu/anaconda3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
  return forward_call(*input, **kwargs)
  File "/home/ubuntu/anaconda3/lib/python3.9/site-packages/transformers/models/gpt2/modeling_gpt2.py", line 1058, in forward
  transformer_outputs = self.transformer(
  File "/home/ubuntu/anaconda3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
  return forward_call(*input, **kwargs)
  File "/home/ubuntu/anaconda3/lib/python3.9/site-packages/transformers/models/gpt2/modeling_gpt2.py", line 901, in forward
  outputs = block(
  File "/home/ubuntu/anaconda3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
  return forward_call(*input, **kwargs)
  File "/home/ubuntu/anaconda3/lib/python3.9/site-packages/transformers/models/gpt2/modeling_gpt2.py", line 401, in forward
  attn_outputs = self.attn(
  File "/home/ubuntu/anaconda3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
  return forward_call(*input, **kwargs)
  File "/home/ubuntu/anaconda3/lib/python3.9/site-packages/transformers/models/gpt2/modeling_gpt2.py", line 325, in forward
  query = self._split_heads(query, self.num_heads, self.head_dim)
  File "/home/ubuntu/anaconda3/lib/python3.9/site-packages/transformers/models/gpt2/modeling_gpt2.py", line 290, in _split_heads
  tensor = tensor.view(new_shape)
  RuntimeError: shape '[116, 5, 12, 64]' is invalid for input of size 464000
  Traceback (most recent call last):
  File "/home/ubuntu/anaconda3/lib/python3.9/site-packages/parallelformers/parallel/process.py", line 193, in inference
  outputs = function_(
  File "/home/ubuntu/anaconda3/lib/python3.9/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
  return func(*args, **kwargs)
  File "/home/ubuntu/anaconda3/lib/python3.9/site-packages/transformers/generation_utils.py", line 1294, in generate
  return self.greedy_search(
  File "/home/ubuntu/anaconda3/lib/python3.9/site-packages/transformers/generation_utils.py", line 1689, in greedy_search
  outputs = self(
  File "/home/ubuntu/anaconda3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
  return forward_call(*input, **kwargs)
  File "/home/ubuntu/anaconda3/lib/python3.9/site-packages/transformers/models/gpt2/modeling_gpt2.py", line 1058, in forward
  transformer_outputs = self.transformer(
  File "/home/ubuntu/anaconda3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
  return forward_call(*input, **kwargs)
  File "/home/ubuntu/anaconda3/lib/python3.9/site-packages/transformers/models/gpt2/modeling_gpt2.py", line 901, in forward
  outputs = block(
  File "/home/ubuntu/anaconda3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
  return forward_call(*input, **kwargs)
  File "/home/ubuntu/anaconda3/lib/python3.9/site-packages/transformers/models/gpt2/modeling_gpt2.py", line 401, in forward
  attn_outputs = self.attn(
  File "/home/ubuntu/anaconda3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
  return forward_call(*input, **kwargs)
  File "/home/ubuntu/anaconda3/lib/python3.9/site-packages/transformers/models/gpt2/modeling_gpt2.py", line 325, in forward
  query = self._split_heads(query, self.num_heads, self.head_dim)
  File "/home/ubuntu/anaconda3/lib/python3.9/site-packages/transformers/models/gpt2/modeling_gpt2.py", line 290, in _split_heads
  tensor = tensor.view(new_shape)
  RuntimeError: shape '[116, 5, 12, 64]' is invalid for input of size 464000

[hyunwoongko]

yes that model can't be parallelized because the number of head is not divisible by 2.

[hyunwoongko]

for tensor model parallelism, the number of head must be divisible by tensor model parallel size (num_gpus in parallelformers). it's limitation of tensor model parallelism.

when gpt2 was published, model parallelism techniques were not developed as now. so they didn't consider this aspect. please use gpt-neo-1.3b model for 1~2b size causal modeling.

thank you!

[snoop2head]

I see! Thank you for the detailed reply and model recommendation 🤗