Jupyter Notebook에서 분산 학습 시작하기

Jupyter Notebook 안에서 분산 학습을 돌리려면 약간의 설정이 필요해요. Accelerate는 notebook_launcher를 제공해서, 노트북 셀 안에서도 멀티 GPU·멀티 노드 학습을 띄울 수 있게 해줘요. 이 페이지에서는 환경 설정부터 실제 학습 실행까지 흐름을 보여드릴게요.

출처: Launching distributed training from Jupyter Notebooks (공식)

환경 설정하기

학습 전에 Accelerate 설정 파일이 시스템에 있어야 해요. 보통 터미널에서 accelerate config를 실행하고 프롬프트에 답하면 만들어져요.

accelerate config

일반적인 기본값이면 충분하고 TPU를 쓰지 않는다면, write_basic_config() 유틸리티로 기기 설정을 빠르게 파일로 쓸 수 있어요. CUDA/XPU는 멀티 디바이스 시스템에서 한 번만 초기화할 수 있으므로, 설정 작성 뒤엔 노트북을 재시작해야 해요.

import os
from accelerate.utils import write_basic_config

write_basic_config()   # 설정 파일 작성
os._exit(00)           # 노트북 재시작

데이터셋과 모델 준비하기

데이터로더와 모델을 준비할 때는 어떤 것도 GPU에 올리지 않도록 주의해야 해요. GPU 관련 코드는 notebook_launcher에 넘길 함수 안에 넣는 걸 권장해요.

import os, re, torch, PIL
import numpy as np

from torch.optim.lr_scheduler import OneCycleLR
from torch.utils.data import DataLoader, Dataset
from torchvision.transforms import Compose, RandomResizedCrop, Resize, ToTensor

from accelerate import Accelerator
from accelerate.utils import set_seed
from timm import create_model

파일명에서 레이블을 추출하는 함수를 만들고, 이미지와 레이블을 반환하는 Dataset 클래스를 정의해요.

def extract_label(fname):
    stem = fname.split(os.path.sep)[-1]
    return re.search(r"^(.*)_\d+\.jpg$", stem).groups()[0]

class PetsDataset(Dataset):
    def __init__(self, file_names, image_transform=None, label_to_id=None):
        self.file_names = file_names
        self.image_transform = image_transform
        self.label_to_id = label_to_id

    def __len__(self):
        return len(self.file_names)

    def __getitem__(self, idx):
        fname = self.file_names[idx]
        raw_image = PIL.Image.open(fname)
        image = raw_image.convert("RGB")
        if self.image_transform is not None:
            image = self.image_transform(image)
        label = extract_label(fname)
        if self.label_to_id is not None:
            label = self.label_to_id[label]
        return {"image": image, "label": label}

학습 함수 작성하기

notebook_launcher()는 함수를 인자로 받아 분산 시스템 전체에서 실행해요. 학습 함수 안에서 시드를 설정하고 Accelerator를 가능한 한 일찍 만들어요.

def training_loop(mixed_precision="fp16", seed: int = 42, batch_size: int = 64):
    set_seed(seed)
    accelerator = Accelerator(mixed_precision=mixed_precision)
    train_dataloader, eval_dataloader = get_dataloaders(batch_size)
    model = create_model("resnet50d", pretrained=True, num_classes=len(label_to_id))

    # 전이 학습: 인코더는 고정하고 분류 헤드만 학습
    for param in model.parameters():
        param.requires_grad = False
    for param in model.get_classifier().parameters():
        param.requires_grad = True

    optimizer = torch.optim.Adam(params=model.parameters(), lr=3e-2 / 25)
    lr_scheduler = OneCycleLR(optimizer=optimizer, max_lr=3e-2, epochs=5,
                              steps_per_epoch=len(train_dataloader))

    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
    )

    for epoch in range(5):
        model.train()
        for batch in train_dataloader:
            inputs = (batch["image"] - mean) / std
            outputs = model(inputs)
            loss = torch.nn.functional.cross_entropy(outputs, batch["label"])
            accelerator.backward(loss)
            optimizer.step()
            lr_scheduler.step()
            optimizer.zero_grad()

prepare()에 넘긴 순서 그대로 객체를 풀어서 쓰면 돼요.

학습 실행하기

notebook_launcher를 import해서 학습 함수와 인자, 프로세스 수를 넘기면 돼요.

from accelerate import notebook_launcher

args = ("fp16", 42, 64)
notebook_launcher(training_loop, args, num_processes=2)

여러 노드에서 실행하려면 각 노드에서 Jupyter 세션을 열고 동시에 실행해야 해요. 예를 들어 2노드·각 8GPU, 메인 컴퓨터 IP가 172.31.43.8이라면 이렇게 실행해요.

# 첫 번째 노드
notebook_launcher(training_loop, args, master_addr="172.31.43.8",
                  node_rank=0, num_nodes=2, num_processes=8)

# 두 번째 노드 — node_rank만 바뀜
notebook_launcher(training_loop, args, master_addr="172.31.43.8",
                  node_rank=1, num_nodes=2, num_processes=8)

TPU에서는 모델을 학습 함수 밖에서 선언하고 매개변수로 넘기는 걸 기억하세요.

디버깅

Notebook에서 흔히 마주치는 문제가 "CUDA/XPU가 이미 초기화됨" 오류예요. 보통 노트북의 import나 앞선 코드가 torch.cuda·torch.xpu를 호출해서 생겨요. 환경변수 ACCELERATE_DEBUG_MODE=yes로 실행하면 스폰 시 추가 검사를 해서 원인을 좁히는 데 도움이 돼요.

Jupyter Notebook에서 분산 학습 시작하기