[Inference] Support GPT-J-6B

Is your feature request related to a problem? Please describe. With the new release of transformers, the gpt-j-6b model will be available for the public: https://github.com/huggingface/transformers/pull/13022

Currently,

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

print(local_rank)
print(world_size)

pipeline.model = deepspeed.init_inference(
    pipeline.model,
    mp_size=1,
    dtype=torch.float16,
    replace_method='auto',
)

will only return

0
1
[2021-08-28 15:05:58,839] [INFO] [logging.py:68:log_dist] [Rank -1] DeepSpeed info: version=0.4.4, git-hash=unknown, git-branch=unknown

Deepspeed already supports the smaller gpt-neo variants, so the addition of gpt-j-6b would make sense.

Additional context If there is anything I could do (create a PR) with some guidance I’d be happy to work on the issue and contribute as well.

[RezaYazdaniAminabadi]

Hi @joehoover,

I am going to be more focused on this through next week. I would say it is ready by early December. Thanks, Reza

[yovizzle]

I’m also very interested in this one.