혼합 정밀도와 기울기 체크포인팅

항상 기억할 메모리·속도 도구 두 가지

트랙 2에서 AMP의 기본을 배웠어. 이 레슨에서는 현대적인 AMP 구성을 다시 정리하고, 계산량을 늘리는 대신 메모리를 아끼는 기울기 체크포인팅을 더할게.

AMP 복습: 현대적인 bf16 경로

Ampere 이후 GPU(A100, RTX 30/40, H100)와 Apple Silicon에서는 bf16 자동 형변환을 우선해. CUDA라면 autocast(device_type='cuda', dtype=torch.bfloat16)를 사용하고 GradScaler는 쓰지 않아. bf16은 fp16과 비슷한 속도를 내면서 지수 범위가 넓어 언더플로 위험이 훨씬 작아. V100이나 T4 같은 이전 하드웨어에서는 fp16과 GradScaler가 여전히 필요할 수 있어.

기울기 체크포인팅

일반적으로 순전파는 역전파에 필요한 모든 중간 활성화를 저장해. 시퀀스 길이가 4K인 24계층 Transformer 같은 깊은 신경망에서는 이 활성화가 메모리의 대부분을 차지해. 기울기 체크포인팅은 일부 활성화를 저장하지 않고 역전파 중 필요할 때 다시 계산해. 활성화 메모리를 보통 50~70% 절감하는 대신 학습 시간이 약 30% 늘 수 있어. 메모리 부족을 피할 수 있다면 충분히 가치 있는 절충이야.

도구 조합

AMP, 기울기 체크포인팅, torch.compile, FSDP는 함께 사용할 수 있어. 규모를 키울 때는 다음 순서로 추가해:

즉시 실행 fp32에서 시작해 정확성을 먼저 맞춰.
bf16 자동 형변환을 추가해 1.5~2배의 속도 향상을 노려.
torch.compile을 추가해 약 1.5~2배의 속도 향상을 다시 노려.
원하는 배치 크기에서 여전히 메모리가 부족하면 기울기 체크포인팅을 넣어.
모델 자체가 GPU 하나에 들어가지 않으면 FSDP를 사용해.

Code

bf16 혼합 정밀도: 현대적인 구성법·python

import torch
from torch.amp import autocast

# bf16 — preferred on Ampere+ and Apple Silicon. NO GradScaler.
for x, y in loader:
    x, y = x.cuda(non_blocking=True), y.cuda(non_blocking=True)
    optimizer.zero_grad()

    with autocast(device_type='cuda', dtype=torch.bfloat16):
        out = model(x)
        loss = criterion(out, y)

    loss.backward()
    optimizer.step()

fp16 혼합 정밀도: 옛 GPU, GradScaler 필수·python

import torch
from torch.amp import autocast, GradScaler

scaler = GradScaler('cuda')

for x, y in loader:
    x, y = x.cuda(non_blocking=True), y.cuda(non_blocking=True)
    optimizer.zero_grad()

    with autocast(device_type='cuda', dtype=torch.float16):
        out = model(x)
        loss = criterion(out, y)

    scaler.scale(loss).backward()    # scale loss to avoid fp16 underflow
    scaler.step(optimizer)            # unscale + step (skip on inf/nan)
    scaler.update()                   # adjust scale dynamically

기울기 체크포인팅: 저장하는 대신 다시 계산·python

import torch.nn as nn
from torch.utils.checkpoint import checkpoint

class CheckpointedTransformer(nn.Module):
    def __init__(self, num_layers):
        super().__init__()
        self.layers = nn.ModuleList([TransformerBlock() for _ in range(num_layers)])

    def forward(self, x):
        for layer in self.layers:
            # Don't store this layer's activations; recompute in backward.
            x = checkpoint(layer, x, use_reentrant=False)
        return x

# For HuggingFace transformers, there's a one-liner:
# model.gradient_checkpointing_enable()

bf16 + 컴파일 + 체크포인팅 결합: 운영 환경 묶음·python

import torch
import torch.nn as nn
from torch.amp import autocast

model = MyTransformer().cuda()
model.gradient_checkpointing_enable()       # if available
model = torch.compile(model, mode='default')

optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

for epoch in range(num_epochs):
    for x, y in loader:
        x, y = x.cuda(non_blocking=True), y.cuda(non_blocking=True)
        optimizer.zero_grad(set_to_none=True)

        with autocast(device_type='cuda', dtype=torch.bfloat16):
            out = model(x)
            loss = criterion(out, y)

        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()

Exercise

Transformer 계열 모델에서 세 설정의 최대 GPU 메모리와 단계별 시간을 재 봐. (a) fp32 즉시 실행, (b) bf16 자동 형변환, (c) bf16과 기울기 체크포인팅을 함께 사용해. 보통 bf16은 메모리를 약 50% 줄이면서 더 빨라지고, 체크포인팅은 활성화 메모리를 추가로 30~50% 줄이는 대신 약 30% 느려질 수 있어.