고급 파이토치 실전 가이드

소개

PyTorch를 마스터한다는 것은 더 이상 기능에 익숙해지는 것을 의미하지 않습니다. 그것은 훈련 코드가 실제 프로덕션 워크로드에서도 빠르고, 확장 가능하며, 복구 가능한 반복 가능한 엔지니어링 워크플로를 운영하는 것을 의미합니다. torch.compile, torch.profiler, DDP/FSDP, 분산 체크포인팅과 같은 기능들은 ML 및 인프라 팀이 훈련을 빠르게 유지하기 위해 사용하는 강력한 도구입니다. 그러나 올바른 순서로 적용하고 철저히 검증하지 않으면 아무런 의미가 없습니다.

이 기사에서는 권장되는 워크플로 경로를 공유합니다: 기준선 → 컴파일 → 프로파일 → 확장 → 체크포인트. 최적화 전에 측정해야 할 사항, 컴파일러 및 프로파일러에서 피해야 할 일반적인 실수, DDP와 FSDP를 선택할 때 따라야 할 결정 규칙, 다중 노드 학습 실행을 위한 내결함성 체크포인트 구현 방법을 다룹니다.

핵심 요점

• PyTorch 성능 튜닝은 기능 체크리스트가 아니라 반복적인 엔지니어링 과정으로 접근하세요: 베이스라인 → 컴파일 → 프로파일 → 스케일 → 체크포인트를 반복적으로 적용하는 것이 최적화를 개별적으로 활성화하는 것보다 더 지속 가능한 성능을 제공합니다.
• 이미 처리량이 확인되고 정확성이 검증된 안정적인 단일 GPU 이거 베이스라인이 없다면, 성능을 효과적으로 비교하거나 문제를 디버깅할 수 없습니다.
• 올바른 기준선을 설정하기 전에 최적화하지 마십시오: 측정된 처리량과 검증된 정확성을 갖춘 안정적인 단일 GPU 즉시 실행 기준선은 의미 있는 성능 비교와 디버깅에 필수적입니다.
• torch.compile을 무작정 사용하지 말고 신중하게 사용하세요: 그래프 중단을 추적하고, 형태를 관리하며, 측정하기 전에 예열하고, 안정적인 상태에서의 성능이 실제로 Eager 실행보다 향상되는지 확인하세요.
• 가정 확인이 아니라 의사 결정을 안내하는 프로필: torch.profiler를 사용하여 명확한 결정을 내리세요, CPU 정체, 커널 핫스팟, 형태 재추적 문제를 해결하고 분산 실행 시 통신 오버헤드를 검증하세요.
• 첫날부터 장애를 대비한 설계 체크포인팅: 분산 체크포인팅(선택적 비동기 저장 지원 포함)은 전체 학습 상태를 캡처해야 하며, GPU 수에 따른 리샤딩을 지원하고, 실제 장애 대응력을 보장하기 위해 복원 연습을 통해 정기적으로 테스트되어야 합니다.

기준선: 기준 점 설정

먼저 작동하는 단일 GPU 학습 예제부터 시작하세요. 이것이 기능과 성능 모두에 대한 기준점이 될 것입니다. 모델, 데이터로더, 학습 루프를 정의하세요. 모든 것이 즉시 실행 모드(컴파일 없이, 한 프로세스)에서 끝까지 실행되는지 확인하세요. 학습 루프가 어떻게 생길 수 있는지에 대한 예시는 다음과 같습니다:

import torch
import torch.nn as nn
# Dummy model and data for illustration
model = nn.Sequential(nn.Linear(100, 50), nn.ReLU(), nn.Linear(50, 10))
data = torch.randn(32, 100)
targets = torch.randint(0, 10, (32,))
# Baseline forward + backward
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
outputs = model(data)
loss = nn.CrossEntropyLoss()(outputs, targets)
loss.backward()
optimizer.step()

몇 번의 학습 반복을 실행하고 처리량(예: 샘플/초)을 측정하세요. 올바른 기준선을 갖추는 것은 모델이 수렴하는지 확인하고 향후 최적화를 위한 비교 지점을 제공합니다. 이 시점에서는 정확성과 기준 성능이 우선입니다: GPU가 사용되고 있는지 확인하고(nvidia-smi 또는 Torch 로그 확인), 주요한 명백한 병목 현상이 없는지 확인하세요(데이터에서 로드 생성 중 지연, 과도한 CPU 연산 등). 안정적인 기준선을 확보하기 전에는 컴파일 및 기타 최적화로 넘어가지 마세요.

기본 체크리스트

• 기능적 정확성: 모델은 단일 GPU에서 훈련되고 예상 결과를 생성합니다.
• 기본 성능 기록: 예: 배치당 시간 또는 GPU 사용률 퍼센트.
• 뚜렷한 병목 현상 없음: 데이터 파이프라인이 GPU를 계속 사용하여 기준선에서 긴 유휴 시간이 없음.

좋은 기본 구성을 구현한 후에는 PyTorch 2.x의 주요 컴파일러를 사용하여 학습을 가속화하는 것부터 시작해 점진적으로 고급 기능을 추가할 수 있습니다.

컴파일: PyTorch에서 torch.compile로 가속화하기

PyTorch 2는 사용자가 torch.compile을 사용하여 성능을 위해 모델을 JIT 컴파일할 수 있도록 했습니다. 모델이나 함수를 torch.compile로 감싸면 PyTorch가 내부에서 최적화된 코드를 생성합니다.

아래에서 얼마나 쉬운지 확인해 보세요:

# Switch to compiled mode
model = torch.compile(model)  # uses default backend 'inductor'

model(data)를 호출하면 자동으로 최적화된 코드 경로가 사용됩니다. 처음 몇 번의 반복에서는 모델이 즉시 컴파일되며, 이후 반복에서는 최적화된 커널을 직접 사용하게 됩니다. PyTorch 2.9는 향후 실행 성능을 향상시키기 위해 (프로세스 간에도) 컴파일 결과를 자동으로 캐시합니다.

그러나 컴파일에서 최대 성능을 얻고 싶다면 그래프 분할과 동적 형태에 대해 알아야 할 몇 가지 사항이 있습니다.

그래프 중단

그래프 브레이크는 컴파일러가 코드의 일부를 단일 그래프로 캡처하지 못할 때 발생합니다(예: 데이터 의존형 Python 제어 흐름). 기본적으로 torch.compile은 지원되지 않는 부분(그래프 깨짐)에 대해서는 자동으로 'eager execution'로 전환하고 나머지는 컴파일합니다. 이렇게 하면 코드가 크래시되지 않습니다. 하지만 각 그래프 브레이크는 일부 코드가 최적화되는 것을 막습니다(해당 부분은 퓨즈드 그래프가 아닌 순수 파이썬에서만 실행됩니다). torch.compile(fullgraph=True)는 모델 중 어느 부분도 컴파일할 수 없을 경우 오류를 발생시키기 때문에 그래프 브레이크의 개발 시 디버깅에 유용합니다. 예를 들어:

model = torch.compile(model, fullgraph=True)
try:
    model(data)
except Exception as e:
    print("Graph break:", e)
model = torch.compile(model, fullgraph=True)
try:
    model(data)
except Exception as e:
    print("Graph break:", e)

이것은 첫 번째로 지원되지 않는 연산에서 오류를 발생시키고, 리팩토링이 필요한 부분을 가리켜 줄 것입니다. 오류 메시지는 일반적으로 문제를 일으킨 원인과 이를 피하는 방법에 대한 힌트(또는 URL)를 포함합니다. PyTorch 로그가 유용할 수 있습니다. 스크립트를 환경 변수 TORCH_LOGS=“graph_breaks ” 와 함께 실행하면 그래프가 깨진 이유와 위치를 출력합니다. 제공된 정보를 사용하여 컴파일을 방해하는 Python 측 연산을 다시 작성하거나 제거하세요(예: Python list() 연산 또는 데이터 의존적 if/else 코드를 텐서 연산으로 교체하거나 완전히 제거하는 것과 같이).

동적인 형태

기본적으로 torch.compile는 확인한 입력 형태에 맞춰 컴파일된 그래프를 최적화합니다. 모델이 다양한 크기의 입력으로 실행되는 경우, 새로운 형태를 만날 때마다 그래프를 다시 컴파일합니다. 이 재컴파일에는 약간의 오버헤드가 발생합니다. PyTorch 2.9에서는 동적 플래그를 추가했습니다. 사전에 dynamic=True를 지정하면 여러 형태(기호적 형태)를 처리할 수 있는 단일 커널을 생성하려고 시도합니다. 예를 들어:

model = torch.compile(model, dynamic=True)

이는 컴파일러에게 일반화된 그래프를 추적하도록 지시하며, 이는 크기 변경(예: 변수 시퀀스 길이) 시 재컴파일을 절충합니다. 반면에 dynamic=False를 전달하면 동적 커널이 생성되지 않으며 정확한 형태에 맞게 특수화가 강제됩니다(참일 경우 더 빠른 코드가 나올 수 있음). dynamic=None 설정은 2.9부터 기본값이며, 컴파일러에게 여러 번 재컴파일이 발생했는지 자동으로 감지하고 필요 시 동적 커널을 즉시 사용하도록 지시합니다.

• 동적 설정은 기본값으로 두거나 입력 형태가 일정하거나 거의 변하지 않는 경우 최대한 최적화를 위해 False로 설정하세요.
• 모양 변화(NLP 모델이 길이가 다른 문장을 처리하는 경우 등) 때문에 자주 재컴파일이 발생하고 재컴파일 한도(기본값은 8회) 때문에 eager 실행으로 되돌아가는 것을 피하고 싶다면 dynamic=True로 설정하세요.

컴파일된 예제

간단한 모델을 컴파일하고 그래프 끊김과 동적 형태를 처리하는 방법을 시연해 봅시다:

import torch
# Sample model with a potential graph break (data-dependent control flow)
class ToyModel(torch.nn.Module):
    def forward(self, x):
        # Example: data-dependent branching (not traceable by Dynamo)
        if x.sum() > 0:
            return x * 2
        else:
            return x
model = ToyModel()
# Attempt full-graph compilation to catch breaks
try:
    model = torch.compile(model, fullgraph=True)
    output = model(torch.randn(4, 4))
except Exception as e:
    print("Graph break detected:", e)
    # Rewrite model or accept partial graph compilation
    model = torch.compile(model, fullgraph=False)  # fallback to allow breaks

위 코드에서 데이터에 따라 달라지는 if 문은 fullgraph=True를 사용할 때 그래프 중단 오류를 발생시킵니다. 우리는 그 오류를 잡아서 출력한 후, 그래프 중단을 허용하는 기본 동작으로 다시 컴파일합니다. 실제로는 모델을 수정하여(이 경우 데이터에 따라 달라지는 분기를 제거) 예외 없이 fullgraph=True로 컴파일할 수 있을 때까지 작업합니다. 이는 모델 전체를 하나의 그래프로 정적으로 컴파일하여 최대 속도를 낼 수 있음을 의미합니다.

체크리스트 컴파일

• torch.compile()로 모델을 감싸세요: 작은 코드 변경으로 종종 상당한 속도 향상을 얻을 수 있습니다. 특별한 이유가 없다면 기본적으로 백엔드 “inductor”를 선택하세요.
• 워밍업: 시간을 측정하기 전에 몇 번 실행해 보세요 – 첫 실행에는 컴파일 시간이 포함됩니다.
• 그래프 중단 확인: 개발 시 fullgraph=True로 실행하고 환경 변수 TORCH_LOGS=“graph_breaks”를 설정하여 지원되지 않는 코드 부분을 확인하세요. 해당 사용을 제거하거나 리팩토링하세요.
• 동적 형태 조정: 많은 재컴파일이 발생하는 경우(stdout/stderr/logs에서 “recompiling ”을 확인), 형태가 유연한 그래프를 생성하기 위해 dynamic=True를 고려할 수 있습니다.
• 속도 모니터링: 기준(baseline) 대비 처리량을 측정하세요. 첫 번째 반복(iteration)은 일반적으로 느릴 수 있지만, 안정된 상태(steady-state) 속도는 초기 기준보다 빨라야 합니다. 그렇지 않다면 성능 힌트(performance hints)를 활성화하고 실행하세요(예: TORCH_LOGS=“perf_hints ” 설정).

모델을 컴파일하고 성능을 향상시킨 후에는 실행을 분석하여 남아 있는 병목 지점을 파악하고 추가로 조정할 영역을 결정해야 합니다.

프로필: torch.profiler로 병목 현상 진단하기

성능 병목 현상은 컴파일 후에도 여전히 존재할 수 있습니다. GPU 활용이 부족하거나, I/O 병목 현상이 발생하거나, 커널이 최적화되지 않은 경우가 이에 해당합니다. PyTorch에는 실행 추적을 수집하고 이러한 병목 현상을 식별할 수 있는 내장 Profiler가 포함되어 있습니다. PyTorch 2.9에서는 torch.profiler가 CPU와 GPU 활동을 추적하고, 텐서의 형태(shape)를 기록하며, 시각화를 위해 Chrome Trace Viewer 또는 TensorBoard와 연동할 수 있습니다. torch.profiler는 또한 비동기적으로 추적을 캡처할 수 있습니다. 이렇게 하면 프로그램 흐름을 방해하지 않고도 훈련의 일부를 프로파일링할 수 있습니다.

import torch
import torch.nn as nn
import torch.optim as optim
import torch.profiler
# -----------------------------
# Device setup
# -----------------------------
device = "cuda" if torch.cuda.is_available() else "cpu"
# -----------------------------
# Model definition
# -----------------------------
model = nn.Sequential(
    nn.Linear(100, 256),
    nn.ReLU(),
    nn.Linear(256, 10)
).to(device)
# -----------------------------
# Optimizer and loss
# -----------------------------
optimizer = optim.Adam(model.parameters(), lr=1e-3)
loss_fn = nn.CrossEntropyLoss()
# -----------------------------
# Data iterator (synthetic)
# -----------------------------
def data_generator(batch_size=32):
    while True:
        inputs = torch.randn(batch_size, 100, device=device)
        targets = torch.randint(0, 10, (batch_size,), device=device)
        yield inputs, targets
data_iter = data_generator()
# -----------------------------
# Training step
# -----------------------------
def train_step(batch):
    model.train()
    inputs, targets = batch
    optimizer.zero_grad(set_to_none=True)
    outputs = model(inputs)
    loss = loss_fn(outputs, targets)
    loss.backward()
    optimizer.step()
    return loss
# -----------------------------
# Warm-up phase
# -----------------------------
for _ in range(5):
    train_step(next(data_iter))
# -----------------------------
# Profiling phase
# -----------------------------
with torch.profiler.profile(
    activities=[
        torch.profiler.ProfilerActivity.CPU,
        torch.profiler.ProfilerActivity.CUDA,
    ],
    record_shapes=True,
    profile_memory=True,
) as prof:
    for _ in range(3):
        train_step(next(data_iter))
        prof.step()
# -----------------------------
# Export trace
# -----------------------------
prof.export_chrome_trace("trace.json")
print("Profiler trace saved to trace.json")

이 예제 스크립트는 최소한의(그러나 end-to-end) PyTorch 학습 루프를 정의하고 torch.profiler를 사용하여 안정적인 CPU 및 GPU 성능을 캡처하는 방법을 보여줍니다. 이 스크립트는 간단한 신경망을 구축하고, 반복자를 사용하여 합성 데이터를 생성하며, 학습 반복을 train_step() 함수로 감쌉니다.

우리는 먼저 몇 번의 워밍업 반복을 수행합니다(한 번만 발생하는 초기화/캐시 오버헤드를 프로파일링하지 않기 위해). 그런 다음 몇 번의 학습 단계를 프로파일링하여 연산자 모양과 메모리 사용량을 기록합니다. 또한 Chrome 트레이스(trace.json)를 내보내는데, 이 파일은 Chrome 브라우저에서 로드하여 GPU 사용량, 커널 실행, CPU–GPU 오버랩, 잠재적인 성능 병목 현상을 확인할 수 있습니다. 입력 모양과 메모리 할당을 추적하여 OOM 오류나 비효율 문제를 디버깅하기 위해 record_shapes=True 및 profile_memory=True를 활성화했습니다.

팁: Google Chrome에서 chrome://tracing을 열고 trace.json을 로드하여 타임라인 뷰를 확인하세요.

자동 병목 현상 감지를 위해 torch.utils.bottleneck 또는 torch.profiler.schedule를 사용하여 스냅샷을 캡처할 수 있습니다. 예를 들어, 스케줄러를 사용하여 매 에포크마다 몇 단계씩 프로파일링할 수 있습니다:

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
import torch.profiler
# -----------------------------
# Device
# -----------------------------
device = "cuda" if torch.cuda.is_available() else "cpu"
# -----------------------------
# Model
# -----------------------------
model = nn.Sequential(
    nn.Linear(100, 256),
    nn.ReLU(),
    nn.Linear(256, 10),
).to(device)
# -----------------------------
# Optimizer and loss
# -----------------------------
optimizer = optim.Adam(model.parameters(), lr=1e-3)
loss_fn = nn.CrossEntropyLoss()
# -----------------------------
# Dataset and DataLoader
# -----------------------------
num_samples = 10_000
batch_size = 32
inputs = torch.randn(num_samples, 100)
targets = torch.randint(0, 10, (num_samples,))
dataset = TensorDataset(inputs, targets)
data_loader = DataLoader(
    dataset,
    batch_size=batch_size,
    shuffle=True,
    num_workers=2,
)
# -----------------------------
# Training step
# -----------------------------
def train_step(batch):
    model.train()
    inputs, targets = batch
    inputs = inputs.to(device, non_blocking=True)
    targets = targets.to(device, non_blocking=True)
    optimizer.zero_grad(set_to_none=True)
    outputs = model(inputs)
    loss = loss_fn(outputs, targets)
    loss.backward()
    optimizer.step()
    return loss
# -----------------------------
# Warm-up (outside profiler)
# -----------------------------
for i, batch in enumerate(data_loader):
    train_step(batch)
    if i >= 5:
        break
# -----------------------------
# Profiling with schedule
# -----------------------------
with torch.profiler.profile(
    activities=[
        torch.profiler.ProfilerActivity.CPU,
        torch.profiler.ProfilerActivity.CUDA,
    ],
    schedule=torch.profiler.schedule(
        wait=1,
        warmup=1,
        active=3,
        repeat=2,
    ),
    on_trace_ready=torch.profiler.tensorboard_trace_handler("./prof_log"),
    record_shapes=True,
    profile_memory=True,
    with_stack=True,
) as prof:
    for step, batch in enumerate(data_loader):
        loss = train_step(batch)
        prof.step()

        if step >= 50:
            break

이것은 대기/워밍업 기간을 건너뛴 후, 비동기적으로 각각 3스텝 길이의 두 개 추적 기간을 시작합니다. 그 후 on_trace_ready 콜백은 학습을 중단하지 않고도 TensorBoard의 프로파일 탭에서 볼 수 있는 추적 데이터를 출력합니다. 비동기 추적 캡처를 사용하면 간헐적으로 프로파일링(예: 50스텝 창에서 3스텝마다)을 수행하여 오버헤드를 줄이면서도 성능 데이터를 수집할 수 있습니다.

프로파일링 후, 결과를 해석하여 조치를 취하세요:

• CPU 제한(많은 빈 시간이 GPU가 대기 중일 때 발생): 더 많은 작업을 GPU로 전달할 수 있는지, 아니면 비동기 데이터 로딩을 사용하여 데이터 로딩과 계산을 겹칠 수 있는지 고려해 보세요(예: DataLoader에 `workers`를 전달하거나 GPU에서 전처리 수행). 또한, 학습 중에 작업을 수행하는 Python 루프가 없는지 확인하세요. 만약 있다면, 이를 제거하거나 `compile`해 보세요.
• GPU 커널이 병목 현상입니다: 이러한 연산(fused ops, 저정밀 연산 등)을 최적화할 수 있습니다. PyTorch 성능 가이드는 라이브러리가 최적화할 수 있도록 고수준 연산(torch.nn.functional 호출 등)을 사용하는 것을 권장합니다. 그러나 특정 연산이 훨씬 느린 것을 발견하면, 그것이 예상된 것인지(큰 FFT를 불러오는 경우 느릴 수 있음) 아니면 다른 방법으로 계산할 수 있는지 판단하세요.
• 여러 형태로 인해 재추적 발생: 입력을 크기별로 나누어 처리하거나, 위에서 언급한 대로 dynamic=True를 사용하세요.
• 외부 병목 현상이 없는지 확인하세요: 예를 들어 분산 학습을 사용하는 경우, 통신(all-reduce)이 병목이 아닌지 확인해야 합니다. 이는 NCCL 커널이나 네트워크에서 대기 중인 CPU로 표시될 수 있습니다.

프로파일링 체크리스트

• 프로파일링 전에 준비 운동: 컴파일이나 지연 초기화 오버헤드는 측정하지 마세요.
• CPU와 GPU 모두 캡처: CPU와 GPU 모두(activities=[CPU, CUDA])를 캡처하여 그 상호작용을 이해합니다.
• 모양 및 메모리 기록: 모양 관련 문제와 메모리 사용량 급증을 식별하는 데 도움을 줍니다.
• 장시간 실행 시 스케줄링 사용: profiler.schedule를 사용하여 짧은 시간 창을 주기적으로 캡처하세요. 이렇게 하면 프로파일러 오버헤드를 줄이면서도 프로그램 동작을 확인할 수 있습니다.
• 활용도 분석: GPU가 충분히 활용되고 있는지 확인하세요. 그렇지 않다면 CPU나 I/O가 병목인지 진단하세요.
• 주요 연산 식별: 프로파일러는 자기 시간(self-time)과 같은 중요한 지표에 따라 정렬된 연산 목록을 출력할 수 있습니다. 시간이 가장 많이 소요되는 연산을 찾아 최적화하세요(더 나은 알고리즘 사용, GPU가 충분히 활용되지 않는 경우 배치 크기 증가 등).

애플리케이션이 단일 머신에서 최적화된 후, 다음 도전 과제는 여러 GPU나 노드를 활용하기 위해 수평 확장(scaling out)하는 것으로, 이는 PyTorch의 분산 학습 패러다임을 도입하게 됩니다.

스케일: DDP 또는 FSDP를 통한 분산 학습

작업량이나 모델이 단일 GPU를 초과할 경우, 여러 GPU로 학습을 확장해야 합니다. PyTorch는 GPU 간 모델 학습을 위한 두 가지 주요 방법을 제공합니다: 분산 데이터 병렬(Distributed Data Parallel)과 완전 분할 데이터 병렬(Fully Sharded Data Parallel)입니다. 두 방법 모두 내부적으로 데이터 병렬성을 활용하여 각 프로세스가 데이터의 서로 다른 하위 집합을 학습하지만, 메모리 내 모델 파라미터를 관리하는 접근 방식에는 차이가 있습니다. 각 방법을 선택해야 하는 시점과 torchrun을 사용한 단일 노드 및 다중 노드 학습 구성을 살펴보겠습니다.

분산 데이터 병렬 처리

DDP는 각 GPU에 전체 모델의 복사본을 유지하고 각 단계 후에 그래디언트를 모두 합산(all-reduce)하여 복사본을 동기화합니다. 단일 GPU의 메모리에 쉽게 맞출 수 있는 모델에 이상적으로 적합하며, DDP는 개념적으로 간단합니다: 먼저 프로세스 그룹을 초기화한 다음, 모델을 torch.nn.parallel.DistributedDataParallel로 감싸면 됩니다.

import os
import torch
import torch.nn as nn
import torch.optim as optim
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, TensorDataset
from torch.utils.data.distributed import DistributedSampler
# -----------------------------
# Setup distributed (safe)
# -----------------------------
def setup_distributed():
    world_size = int(os.environ.get("WORLD_SIZE", "1"))
    if world_size > 1:
        dist.init_process_group(backend="nccl")
        local_rank = int(os.environ["LOCAL_RANK"])
        torch.cuda.set_device(local_rank)
        device = torch.device(f"cuda:{local_rank}")
        distributed = True
    else:
        local_rank = 0
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        distributed = False

    return distributed, local_rank, device

def cleanup_distributed(distributed):
    if distributed:
        dist.destroy_process_group()

# -----------------------------
# Main training
# -----------------------------
def main():
    distributed, local_rank, device = setup_distributed()

    # -----------------------------
    # Model
    # -----------------------------
    model = nn.Sequential(
        nn.Linear(100, 256),
        nn.ReLU(),
        nn.Linear(256, 10),
    ).to(device)
    if distributed:
        model = DDP(model, device_ids=[local_rank])
    # -----------------------------
    # Optimizer and loss
    # -----------------------------
    optimizer = optim.Adam(model.parameters(), lr=1e-3)
    loss_fn = nn.CrossEntropyLoss()
    # -----------------------------
    # Dataset
    # -----------------------------
    num_samples = 20_000
    batch_size = 32
    inputs = torch.randn(num_samples, 100)
    targets = torch.randint(0, 10, (num_samples,))
    dataset = TensorDataset(inputs, targets)
    if distributed:
        sampler = DistributedSampler(dataset, shuffle=True)
    else:
        sampler = None
    train_loader = DataLoader(
        dataset,
        batch_size=batch_size,
        sampler=sampler,
        shuffle=(sampler is None),
        num_workers=2,
        pin_memory=True,
    )

    # -----------------------------
    # Training loop
    # -----------------------------
    epochs = 3
    for epoch in range(epochs):
        if distributed:
            sampler.set_epoch(epoch)

        for step, (x, y) in enumerate(train_loader):
            x = x.to(device, non_blocking=True)
            y = y.to(device, non_blocking=True)

            optimizer.zero_grad(set_to_none=True)
            outputs = model(x)
            loss = loss_fn(outputs, y)
            loss.backward()
            optimizer.step()
            if step % 50 == 0 and local_rank == 0:
                print(
                    f"[Epoch {epoch}] Step {step} | "
                    f"Loss {loss.item():.4f}"
                )
    cleanup_distributed(distributed)
# -----------------------------
# Entry point
# -----------------------------
if __name__ == "__main__":
    main()

DDP 작업을 시작하는 것은 torchrun으로 수행됩니다. 예를 들어, 4개의 GPU를 가진 단일 노드에서 작업을 시작하고자 할 경우:

torchrun --nproc_per_node=4 train.py

이 명령은 4개의 프로세스(각 GPU마다 하나)를 생성하고 적절한 환경 변수(랭크, 월드 사이즈 등)를 초기화합니다. 여러 노드를 지정하려면 --nnodes와 네트워크 세부 정보를 전달해야 합니다.

# On node 0:
torchrun --nnodes=2 --nproc_per_node=4 --node_rank=0 --master_addr="<IP of node0>" --master_port=12345 train.py
# On node 1:
torchrun --nnodes=2 --nproc_per_node=4 --node_rank=1 --master_addr="<IP of node0>" --master_port=12345 train.py

완전 분할 데이터 병렬(FSDP)

FSDP는 모델 매개변수와 옵티마이저 상태를 각 GPU에 완전히 복제하는 대신 GPU 전반에 분할(sharding)합니다. PyTorch 2.9에서 FSDP를 사용하는 것은 FullyShardedDataParallel 래퍼를 통해 간단하게 수행할 수 있습니다.

핵심 사항

• 모델(또는 서브모듈)을 옵티마이저 생성 전에 FSDP로 감싸세요.
• DDP와 마찬가지로 torch.cuda.set_device(rank)를 사용하고 모델을 GPU로 이동하세요.
• 선택 사항: 모듈 단위로 샤딩하고 싶다면 자동 래핑 정책을 제공하세요(예: 각 트랜스포머 블록을 개별적으로 래핑). 기본적으로 FSDP(model)는 모델 전체를 하나의 샤드 단위(기술적으로 하나의 FSDP 단위)로 샤딩하며, 이는 매우 깊은 모델에서는 최적이 아닙니다.

이것은 하위 모듈 중 10만 개 이상의 파라미터를 가진 모듈을 재귀적으로 자체 FSDP 샤드로 나눕니다. 샤드가 너무 적으면(예: 모델 전체가 단일 샤드인 경우) 레이어 간 메모리 절약 효과를 얻을 수 없습니다. 그러나 샤드가 너무 많으면 통신 오버헤드가 증가합니다. 이제 DDP와 FSDP의 차이점을 요약해 보겠습니다:

예시 목적으로 DDP부터 시작할 수 있습니다. 메모리 부족 문제가 발생하거나 더 큰 모델이 필요하면 FSDP 사용을 고려해 보세요. PyTorch 2.9에서 FSDP는 발전했음을 주의하세요. 기본적으로 모든 것을 샤딩(ZeRO Stage 3)하며, 예상치 못한 메모리 급증을 방지하기 위해 limit_all_gathers=True와 같은 다른 기본값을 활용합니다.

확장 체크리스트

• 결정론 보장: 재현성을 위해 모든 랭크에서 동일한 랜덤 시드를 설정하고 있는지 확인하세요 (필요하다면 torch.manual_seed(랜덤 오프셋용 시드) 사용).
• DistributedSampler 사용: 데이터를 여러 랭크로 나눈 뒤, 매 에포크마다 set_epoch()를 호출하여 각 에포크마다 데이터를 다르게 섞이도록 합니다.
• DDP: DDP로 래핑된 모델에서 backward()를 호출하면 GPU 전체에서 그래디언트가 자동으로 합산됩니다. GPU 수를 늘릴 때 학습률이나 배치 크기를 늘리는 것을 잊지 마세요!
• FSDP: 옵티마이저를 만들기 전에 모듈을 래핑하세요. 매우 깊은 모델을 사용하고 있다면 하나의 거대한 샤드를 피하기 위해 auto_wrap_policy를 조정해보세요. GPU를 주시하며 샤드 크기가 실제로 맞는지 확인하세요.
• 통신 백엔드: 다중 노드의 경우, 네트워킹이 설정되어 있는지 확인하세요(일반적으로 NCCL). NVLink가 있는 경우 등 NCCL을 최적화하려면 환경 변수 NCCL_P2P_LEVEL=NVL을 설정하세요.
• 그래디언트 누적: 이를 사용하는 경우, FSDP와 함께 특정 반복에서 그래디언트 동기화를 건너뛸 때 model.no_sync()를 호출해야 합니다(DDP 사용과 유사).

체크포인트: 분산 체크포인트로 안정적으로 학습 복구

장기 또는 대규모 학습 실행에서는 체크포인트를 저장하는 것이 중요합니다. 학습 중에 오류가 발생한 후 다시 학습을 이어가거나, 학습 중간에서 미세 조정을 하거나, 중간 결과를 확인하고 싶을 수 있습니다. PyTorch 2.9에는 분산 체크포인트(DCP)를 위한 편리한 기본 기능이 포함되어 있습니다. 이를 통해 사용자는 기존의 torch.save 방식보다 더 효율적이고 안정적으로 저장/불러오기를 할 수 있습니다. 권장 접근 방식으로 들어가기 전에 두 가지 접근 방식을 비교해 보겠습니다.

전통적인 torch.save/torch.load

단일 GPU 또는 DDP 훈련 스크립트에서 전통적인 방법은 모델의 state_dict()를 파일에 저장합니다(일반적으로 모든 프로세스가 이 파일에 접근할 수 있으므로 주로 랭크 0에서만 저장). 예를 들어:

# On rank 0 only:
torch.save(model.state_dict(), "checkpoint.pt")
and to load:
model.load_state_dict(torch.load("checkpoint.pt"))

이 방법은 하나의 프로세스가 전체 모델 상태를 메모리에 유지할 수 있는 작은 모델/설정(보통 rank 0)의 경우에는 괜찮지만, 모델이 여러 프로세스에 분산되어 있을 수 있는 분산 학습에서는 이 접근 방식을 사용하는 데 문제가 있습니다.

분산 체크포인트(DCP): PyTorch의 torch.distributed.checkpoint 모듈(v2.1 이상에서 추가, 2.9에서 안정화)은 이러한 문제를 해결합니다. 이 모듈은 저장을 병렬화하여 각 랭크가 모델 상태의 일부를 작성하도록 하며, 체크포인트를 구성하는 여러 파일(최소한 랭크별 하나)을 생성합니다. 또한 로드 시 리샤딩(resharding)을 수행할 수 있어, N개의 랭크로 저장하고 M개의 랭크로 로드할 때 DCP가 적절히 자동으로 데이터를 수집/분할합니다.

일반적인 사용 패턴은 모델(필요한 경우 옵티마이저 포함)을 state_dict를 제공하는 상태 저장 컨테이너에 감싸는 것을 포함합니다. 이는 FSDP를 인식하는 제공된 유틸리티 함수 get_state_dict와 set_state_dict를 사용하여 수행할 수 있습니다:

import torch.distributed.checkpoint as dcp
from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict
from torch.distributed.checkpoint.stateful import Stateful

# Define a stateful container for model & optimizer
class AppState(Stateful):
    def __init__(self, model, optimizer=None):
        self.model = model
        self.optimizer = optimizer
    def state_dict(self):
        model_sd, opt_sd = get_state_dict(self.model, self.optimizer)
        return {"model": model_sd, "optim": opt_sd}
    def load_state_dict(self, state):
        set_state_dict(self.model, self.optimizer, state["model"], state.get("optim"))

여기서 AppState는 Stateful 인터페이스를 구현하므로 DCP는 상태를 가져오고 설정하는 방법을 알 수 있습니다. 모델이 샤딩된 경우(get_state_dict)는 각 랭크가 모델의 일부만 가져오므로 FSDP의 샤드 상태 사전을 자동으로 사용한다는 점에 유의하세요.

DCP로 절약하기

모든 계급 실행:

app_state = AppState(model, optimizer)
dcp.save(app_state, checkpoint_id="mycheckpoint")

이렇게 하면 'mycheckpoint '라는 문자열이 접두사로 붙은 파일이 들어 있는 디렉토리가 생성됩니다 (mycheckpoint/mpirank_0_0.pdparams). 각 랭크는 동시에 자신의 샤드를 기록합니다. 이는 한 랭크가 모든 작업을 수행하는 것에 비해 시간을 크게 줄일 수 있습니다.

DCP로 로딩 중

체크포인트에서 다시 시작하려면 동일한 토폴로지를 가진 모델이 필요합니다(순위 수가 다를 수도 있음). 모델(및 옵티마이저)을 구성하고, AppState에 래핑한 후, 다음을 수행합니다:

app_state = AppState(model, optimizer)
dcp.load(app_state, checkpoint_id="mycheckpoint");

저장할 때와 랭크 수가 다를 경우, DCP는 내부적으로 샤드를 재배치합니다(예: 8개 GPU에서 4개 GPU로 줄이면 각 랭크가 두 개의 저장된 샤드를 불러오거나, 랭크를 늘리면 이에 맞게 분배합니다). 이 로드 시 리샤딩 단계는 수동 체크포인트 변환의 필요성을 없애줍니다.

비동기 체크포인트

PyTorch 2.9의 가장 멋진 기능 중 하나는 비동기 체크포인트 저장으로, 체크포인트 쓰기를 학습 과정과 겹치게 할 수 있습니다. dcp.async_save를 사용하면, API가 미래 객체(future)를 반환하고 무거운 I/O를 백그라운드 스레드에서 수행합니다. 기본 아이디어는:

save_future = dcp.async_save(app_state, checkpoint_id="chkpt_epoch10")
# ... training can continue immediately ...
save_future.wait()  # later, wait for completion (or periodically check)

async_save는 먼저 데이터를 CPU에서 로컬로 스테이징합니다(모델/옵티마이저 상태를 고정 메모리 버퍼로 복사). 그 다음 해당 버퍼를 비동기적으로 디스크나 네트워크에 기록합니다. 체크포인트가 플러시되는 동안 GPU는 연산 작업(즉, 학습)을 계속 진행할 수 있습니다.

비동기 저장은 일시적으로 모델 상태를 저장하는 데 필요한 메모리 요구량을 두 배로 필요로 합니다. 각 랭크는 체크포인트의 샤드 크기와 거의 동일한 CPU 버퍼를 할당합니다. 따라서 많은 랭크에 걸쳐 매우 큰 모델을 비동기적으로 저장하면 메모리 사용량이 상당히 증가할 수 있습니다(최악의 경우 CPU RAM + 고정 메모리 사용량 ≈ 랭크당 체크포인트 크기 * 랭크 수).

이제 각 방법을 언제 사용해야 하는지에 대한 지침과 함께 체크포인트 방법의 비교를 정리해 보겠습니다:

체크포인트 확인 목록

• 모든 필요한 상태 포함: 나중에 학습을 완전히 재개할 수 있도록 모델 가중치와 옵티마이저 상태(필요한 경우 스케줄러, RNG 상태 등)를 저장하도록 하세요. DCP의 Stateful 방식을 사용하는 경우, get_state_dict(…, model, optimizer)를 통해 쉽게 옵티마이저를 전달할 수 있습니다.
• 복원 테스트: 체크포인트를 저장하자마자, 깨끗한 환경에서 다시 불러오세요(가능하다면 동일한 GPU 수와 분산 시 다른 GPU 수에서도). 이렇게 하면 모든 것이 예상대로 복원되는지 확인할 수 있습니다. 이는 학습 상태의 일부를 초기 단계에서 저장하지 않은 것과 같은 오류를 찾는 데 도움이 됩니다.
• 스토리지 관리: 체크포인트는 매우 빠르게 커질 수 있습니다(특히 분할된 분산 체크포인트의 경우). 디스크 공간 부족을 피하기 위해 일부 정리 전략(마지막 k개의 체크포인트만 유지 등)을 구현하세요.
• 비동기 사용: async_save는 유용한 기능이지만, 한 번에 여러 개의 비동기 저장이 진행되지 않도록 주의해야 합니다(메모리 사용 방식과 중첩 허용 방법을 매우 신중하게 다루는 경우를 제외하고). 결국에는 결과로 나온 future에 대해 .wait()를 호출하여 예외를 발생시키거나 성공 여부를 확인해야 합니다.
• 일관성: 체크포인트에서 다시 시작할 때, 모든 랭크가 동일한 dcp.load.를 호출하도록 하세요. 이는 FSDP의 경우 특히 중요한데, 내부 로드를 동기화해야 하기 때문입니다. 또한, CPU 오프로드나 기타 메커니즘을 사용하는 경우 로드하기 전에 모델이 예상 장치에 있는지 확인하세요.

결론

이 워크플로를 따르면 성능이 좋고 오류에 강한 PyTorch 학습 코드를 개발할 수 있습니다:

• 작동하는 기준점에서 시작하세요.
• 최대 속도로 컴파일하십시오.
• 최적화 기회를 식별하기 위한 프로필.
• 적절한 데이터 또는 모델 병렬 전략을 사용하여 확장하세요.
• 체크포인트를 효과적으로 사용하여 실패에서 복구할 수 있습니다.

각 단계별로 사용할 수 있는 플레이북이 있습니다. 이를 반복적으로 사용하세요(예: 4단계에서 성능 특성이 변경된 후 다시 프로파일링하거나, 새로운 환경에 맞게 2단계에서 컴파일 설정을 다시 조정하는 경우). PyTorch는 빠르게 발전해 왔기 때문에 각 단계에서 도움이 되는 기능들이 있습니다. 각각을 언제 어떻게 사용할지 배우면, 빠르고 대규모 모델을 훈련하면서도 효과적인 복구 전략을 갖출 수 있습니다.

자주 묻는 질문

최적화하기 전에 기준선을 설정하는 것이 왜 그렇게 중요한가요?

최적화를 하기 전에 정확한 단일 GPU 이거(Eager) 기준 성능을 확보해야 합니다. 그래야 정확성과 성능에 대해 신뢰할 수 있는 기준이 생깁니다. 그렇지 않으면 최적화가 잘못된 결과를 초래하는 버그를 숨길 수 있습니다. 정확한 기준이 있어야만 속도 향상이나 성능 저하를 특정 코드 변경에 귀속시킬 수 있으며, 알 수 없는 버그나 측정 잡음 때문이라고 볼 수 없습니다.

torch.compile를 언제 사용해야 하고, 주의할 점은 무엇인가요?

torch.compile은 안정적인 기준선을 확보한 후 정상 상태 실행을 가속화하기 위해 사용해야 합니다. 그래프 중단을 모니터링하고, 측정 전에 워밍업하며, 과도한 재컴파일을 피하기 위해 동적 형태가 신중하게 처리되도록 하세요.

torch.profiler가 단순히 가정을 확인하는 것 이상의 도움을 주는 방법은 무엇인가요?

torch.profiler는 CPU 정체, 비효율적인 GPU 커널, 형태 재추적, 통신 오버헤드 등 실제 병목 현상을 보여줍니다. 이를 통해 직관이 아닌 증거를 기반으로 최적화 노력을 중요한 부분에 집중할 수 있습니다.

분산 학습에서 DDP와 FSDP 중 어떻게 선택하나요?

모델이 사용하려는 각 GPU에 무리 없이 들어간다면, 간단하고 상대적으로 속도가 빠르기 때문에 DDP를 선택하세요. GPU 메모리에 제한이 있거나 매우 큰 모델을 다루고 있다면, 복잡성이 증가하는 대신 FSDP가 더 높은 확장성을 제공합니다.

대규모 실행에서 왜 torch.save보다 분산 체크포인팅이 선호되나요?

분산 체크포인팅은 저장 작업을 여러 랭크에 병렬화하고, 로드 시 리샤딩을 지원하며, 비동기 체크포인팅을 가능하게 하여 대규모 다중 GPU 또는 다중 노드 학습 작업에서 복구를 신뢰성 있고 효율적으로 만듭니다.

참조

• 파이썬에서 PyTorch 라이브러리에 대한 궁극적인 가이드
• 여러 GPU에서 모델 학습 확장하기: PyTorch DDP와 FSDP를 활용한 효율적인 전략
• 완전 샤딩 데이터 병렬
• 토치 분산 체크포인트
• torch.compile 튜토리얼

Linux-Console.net 커뮤니티와 함께 학습해 주셔서 감사합니다.