DeepSpeed Pipeline Parallelism — 레이어를 스테이지로 나눠 병렬 처리
DeepSpeed Pipeline Parallelism — 레이어를 스테이지로 나눠 병렬 처리
DeepSpeed의 파이프라인 병렬은 모델의 레이어를 스테이지로 분할하여 각기 다른 GPU에서 처리하게 해요. 메모리와 계산 효율을 동시에 높여요.
동작 원리
- 그래디언트 누적(gradient accumulation): 하나의 배치를 마이크로배치로 나누고, 파이프라인 각 스테이지가 순방향(F)·역방향(B)을 번갈아 수행해요.
- 스테이지가 마이크로배치의 순방향을 끝내면 활성화 메모리를 다음 스테이지로 전달하고, 역방향 그래디언트는 뒤로 전파돼요.
- 이후 데이터 병렬 그룹이 그래디언트를 all-reduce로 합쳐요.
PipelineModule
PipelineModule(layers=net, num_stages=2)처럼 torch.nn.Sequential을 감싸면 파이프라인 모델이 돼요. partition_method로 param 기준·layer 타입·균등 분할을 고를 수 있어요.
net = nn.Sequential(nn.Linear(in_features, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, out_features))
net = PipelineModule(layers=net, num_stages=2)
대규모 모델에서는 LayerSpec으로 레이어 생성을 지연해 CPU 메모리 중복을 피할 수 있어요.