toolkit, environment, 

GPT-NeoX - DeepSpeed Inference

GPT-NeoX - DeepSpeed Inference

GPT-NeoX - DeepSpeed Inference

DeepSpeed Inference를 사용하면 간단하게 모델 추론 성능을 끌어올릴 수 있다.
Tensor Parallel 등의 최적화를 쉽게 제공하는데, injection이라는 기능을 사용하면 더욱 빠르게 할 수 있다.

replace_with_kernel_inject=True로 넣어주면 Bert, GPT2, GPT-Neo and GPT-J 모델들은 알아서 최적화를 해주는데, NeoX는 아직 자동으로 최적화가 안 된다.

그래서 조금 커스텀하게 다음과 같이 적용해서 사용 가능하다.

import deepspeed
import torch
from transformers import AutoModelForCausalLM, GPTNeoXLayer

with deepspeed.OnDevice(dtype=torch.float16, device="cuda"):
    model = AutoModelForCausalLM.from_pretrained('EleutherAI/polyglot-ko-5.8b', torch_dtype=torch.float16, cache_dir='/data/.cache')

model = deepspeed.init_inference(
    injection_policy={GPTNeoXLayer: ('attention.dense','mlp.dense_4h_to_h')}

Subscribe to SOOFTWARE

Get the latest posts delivered right to your inbox