FairScale FSDP
FairScale FSDP (Fully Sharded Data Parallel)
FairScale의 FullyShardedDataParallel(FSDP)은 모듈 파라미터를 데이터 병렬 워커들에 걸쳐 샤딩하는 래퍼예요. Xu et al.의 연구와 DeepSpeed의 ZeRO Stage 3에서 영감을 받았어요. FSDP는 큰 NN 모델로 확장하는 데 권장되는 방법으로, 이후 PyTorch로 업스트림되어 공식 API가 되었어요.
기본 사용법
FSDP로 모듈을 감싸고, 옵티마이저를 초기화한 뒤 평소처럼 학습하면 돼요.
import torch
from fairscale.nn.data_parallel import FullyShardedDataParallel as FSDP
torch.cuda.set_device(device_id)
sharded_module = FSDP(my_module)
optim = torch.optim.Adam(sharded_module.parameters(), lr=0.0001)
x = sharded_module(x, y=3, z=torch.Tensor([1]))
loss = x.sum()
loss.backward()
optim.step()
주의: 옵티마이저는 모듈이 래핑된 후에 초기화해야 해요. FSDP가 파라미터를 제자리에서 샤딩하기 때문에, 그 전에 만든 옵티마이저는 깨질 수 있어요.
개별 레이어 샤딩
개별 레이어를 따로 샤딩하고 바깥쪽 래퍼가 남은 파라미터를 처리하게 할 수도 있어요. 이렇게 하면 GPU 메모리 사용을 더 줄이고, 큰 모델 초기화 시 시스템 메모리 사용을 줄이며, forward 패스에 걸쳐 all-gather 단계를 겹쳐 학습 속도를 높일 수 있어요.
import torch
from fairscale.nn.wrap import wrap, enable_wrap, auto_wrap
from fairscale.nn.data_parallel import FullyShardedDataParallel as FSDP
fsdp_params = dict(wrapper_cls=FSDP, mixed_precision=True, flatten_parameters=True)
with enable_wrap(**fsdp_params):
l1 = wrap(torch.nn.Linear(5, 5)) # 컨텍스트 안에선 기본으로 FSDP로 래핑
large_tfmr = torch.nn.Transformer(
d_model=2048, num_encoder_layers=12, num_decoder_layers=12
)
l2 = auto_wrap(large_tfmr) # 자식 모듈을 자동으로 FSDP로 래핑
auto_wrap은 1e8개 이상의 파라미터를 가진 자식 모듈을 따로 FSDP로 래핑해 줘요.
FSDP가 해결하는 것
일반 DDP는 각 GPU에 모델 전체 복사본을 두고 배치만 나눠요. 하지만 모델이 커지면 메모리가 부족해져요. FSDP는 파라미터·그래디언트·옵티마이저 상태를 샤딩해서, 제한된 자원으로 훨씬 큰 모델을 학습할 수 있게 해줘요.