FSDP: 하나의 GPU에 안 들어가는 모델 학습

모델이 GPU 하나에 들어가지 않을 때

DDP는 모든 GPU에 모델 전체를 복제해. 모델 하나가 GPU 메모리에 들어가지 않으면 이 방식을 쓸 수 없어. FullyShardedDataParallel(FSDP)는 매개변수, 기울기, 옵티마이저 상태를 GPU마다 나눠 보관해. 각 GPU는 전체의 1/N만 계속 들고 있고, 순전파와 역전파에 필요한 조각만 잠시 모았다가 다시 해제해. 그래서 충분한 수의 24GB GPU를 묶으면 700억 매개변수 규모의 모델도 학습할 수 있어.

샤딩 전략

FULL_SHARD: 매개변수, 기울기, 옵티마이저 상태를 모두 나눠. 메모리를 가장 아끼므로 아주 큰 모델의 기본 선택이야.
SHARD_GRAD_OP: 기울기와 옵티마이저 상태만 나누고 매개변수는 복제해. 매개변수 all-gather가 적어서 FULL_SHARD보다 빠를 수 있지만 메모리는 더 써.
NO_SHARD: DDP와 같은 방식이야. 같은 학습 스크립트의 단위 테스트에 유용해.

FSDP1과 FSDP2

FSDP1은 모듈을 감싸는 기존 API야. FSDP2는 fully_shard()를 사용하는 매개변수별 방식으로, 인터페이스가 더 깔끔하고 torch.compile과도 잘 조합돼. 2026년의 최신 PyTorch를 대상으로 새 코드를 작성한다면 FSDP2를 우선할 수 있지만, FSDP1은 여전히 운영 환경에서 안정적이고 문서도 더 성숙해.

학습 반복문은 거의 그대로야

학습 반복문 안에서는 FSDP와 DDP가 거의 같아 보여. 모델을 감싸고 평소처럼 학습하면 돼. 복잡한 부분은 반복문이 아니라 중첩 모듈의 자동 감싸기 정책과 혼합 정밀도 같은 구성에 있어.

Code

FSDP1: 모델을 감싸는 API·python

import torch
import torch.distributed as dist
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import ShardingStrategy

dist.init_process_group("nccl")
rank = int(os.environ['LOCAL_RANK'])
torch.cuda.set_device(rank)

model = MyHugeModel().cuda(rank)
model = FSDP(
    model,
    sharding_strategy=ShardingStrategy.FULL_SHARD,
    device_id=rank,
)

# The training loop is the SAME as DDP
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
for epoch in range(num_epochs):
    for x, y in loader:
        optimizer.zero_grad()
        loss = criterion(model(x.cuda(rank)), y.cuda(rank))
        loss.backward()
        optimizer.step()

Auto-wrap: 중첩 모듈 FSDP가 처리·python

import torch
import functools
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy
import torch.nn as nn

# Wrap any nn.Module submodule with > 100M params automatically
auto_wrap_policy = functools.partial(
    size_based_auto_wrap_policy,
    min_num_params=100_000_000,
)

model = FSDP(
    MyHugeModel(),
    auto_wrap_policy=auto_wrap_policy,
    sharding_strategy=ShardingStrategy.FULL_SHARD,
    device_id=rank,
)

# For Transformer-shaped models, prefer transformer_auto_wrap_policy
# which wraps each block — much better for FSDP performance

FSDP2: 현대적인 매개변수별 API·python

import torch.nn as nn
from torch.distributed._composable.fsdp import fully_shard

# FSDP2 is per-parameter and per-block, not whole-model
class TransformerStack(nn.Module):
    def __init__(self):
        super().__init__()
        self.blocks = nn.ModuleList([TransformerBlock() for _ in range(24)])
    def forward(self, x):
        for b in self.blocks:
            x = b(x)
        return x

model = TransformerStack()

# Apply FSDP2 to each block — finer control than FSDP1's wrap policy
for layer in model.blocks:
    fully_shard(layer)
fully_shard(model)

# FSDP2 advantages:
#   - Cleaner interaction with torch.compile
#   - Per-parameter sharding (mix sharded / non-sharded freely)
#   - More composable with other parallelism strategies (TP, PP)

Exercise

GPU 하나에 여유 있게 들어가는 모델을 골라 다중 GPU 컴퓨터에서 FULL_SHARD 방식의 FSDP1으로 실행해 봐. GPU별로 샤딩 대상 상태의 메모리 사용량이 단일 GPU 사용량보다 크게 줄어드는지 확인해. 전체 메모리가 정확히 1/N이 되는 건 활성화와 임시 버퍼 때문에 기대하지 않아도 돼. 이 규모에서는 DDP보다 빨라지길 기대하지 말고, 모델이 다른 방법으로 들어가지 않을 때 FSDP의 가치가 생긴다는 점을 확인해.