모델을 여러 개의 GPU에 로드해야 할 때

큰 모델을 GPU에 올리다 보면 CUDA out of memory 에러가 빈번하게 날 수 있다. device_map="auto"로 사용 가능한 모든 GPU에 자동으로  모델을 나눠서 올리는 것인데, 내가 원하는 GPU를 적접 정하고 싶은 경우에는 어떻게 할까? 내 경우에는 랩실 사람들끼리 사용하기로 한 GPU를 나눠서 지정했기 때문에 아래 방법을 찾아보게 되었다. 

 

config.py에 device_map을 정의하고, 모델을 정의한 device_map으로 분산시켜 올린다. 

## opt-13b에 대한 레이어 분산 설정

device_map = {
    "model.decoder.embed_tokens": 0,      
    "model.decoder.embed_positions": 0,
    "model.decoder.final_layer_norm": 0,
    "lm_head": 0,
}

# 레이어 분산 설정 (0~19: cuda:0, 20~39: cuda:2)
for i in range(40):
    if i < 20:
        device_map[f"model.decoder.layers.{i}"] = 0
    else:
        device_map[f"model.decoder.layers.{i}"] = 2

 

## 모델 레이어 분산 설정

model = model_class.from_pretrained(args.model_name_or_path, device_map=config.device_map)
tokenizer = tokenizer_class.from_pretrained(args.model_name_or_path)