AI/LLM

[고급-코드] 병렬화(Parallelize)

moviewine 2023. 8. 28. 16:58

# 출처: https://github.com/AGI-Edgerunners/LLM-Adapters/blob/main/finetune.py


# Parallelize 관련 code line

import os

device_map=“auto”
world_size = int(os.environ.get(“WORLD_SIZE”, 1))
ddp = world_size != 1
if ddp:
    device_map = {“”: int(os.environ.get(“LOCAL_RANK”) or 0)}
    gradient_accumulation_steps = grandient_accumulation_steps // world_size

if not ddp and torch.cuda.device_count()