⚡ AI Optimization

파이프라인 병렬화

개요

파이프라인 병렬화(Pipeline Parallelism, PP)는 순차적인 신경망의 연속된 레이어를 여러 device의 stage로 나누고, 서로 다른 micro-batch가 서로 다른 stage를 동시에 통과하도록 실행하는 모델 병렬화 기법이다. 데이터 병렬화가 모델 복제본을 여러 device에 두는 방식이라면, PP는 하나의 모델을 깊이 방향으로 분할하므로 단일 device 메모리에 들어가지 않는 모델도 여러 device에 배치할 수 있다.

한 번에 전체 batch를 한 stage에 보내면 다른 stage가 기다리는 시간이 길어지므로, batch를 여러 micro-batch로 쪼개 pipeline을 채운다. Forward activation은 다음 stage로 전달하고 backward에서 계산된 activation gradient는 이전 stage로 되돌린다. Stage 사이의 통신은 주로 point-to-point이며, 성능은 stage별 계산량 균형, micro-batch 수, 통신 링크, schedule의 pipeline bubble에 의해 결정된다.

핵심 개념

Stage와 pipeline degree

모델의 forward graph를 순서가 보존되는 여러 partition으로 나누고, 각 partition을 stage에 배치한다. pipeline degree(p)는 stage 수이며 보통 하나의 pipeline group 안의 device 수와 같다. Stage i는 자기 레이어의 parameter와 optimizer state를 소유하고, 이전 stage에서 activation을 받아 다음 stage로 결과를 보낸다.

용어	의미	설계 시 확인할 점
stage	모델의 연속된 레이어 partition	parameter 수뿐 아니라 activation 크기와 연산 시간을 함께 균형화
pipeline degree (`p`)	pipeline group의 stage 수	다른 data/tensor parallel group과 process mapping을 구분
micro-batch	pipeline에 흘려보내는 작은 batch 단위	수를 늘리면 bubble이 줄지만 activation buffer가 증가
activation	한 stage의 forward 출력	다음 stage로 P2P 전송되며 backward까지 저장될 수 있음
activation gradient	다음 stage에서 역전파되어 돌아오는 gradient	이전 stage의 입력 activation에 대한 gradient
pipeline bubble	stage가 실행할 작업이 없어 유휴 상태인 시간	warmup, steady state, cooldown 경계에서 발생

Partitioning과 균형

가장 단순한 partition은 레이어 수를 동일하게 나누는 방식이지만, 레이어마다 parameter 수와 실행 시간이 다르므로 실제 stage time이 같아지지 않을 수 있다. Transformer에서는 embedding과 language-model head, attention/FFN의 비용, sequence length에 따른 activation 크기를 함께 측정해야 한다. 목표는 각 stage의 forward + backward 시간이 비슷하고 stage 사이 activation 전송도 과도하게 크지 않도록 경계를 정하는 것이다.

매우 큰 모델에서는 하나의 process가 여러 virtual stage를 맡는 interleaved pipeline을 사용해 stage의 granularity를 높일 수 있다. 다만 virtual stage 사이의 local scheduling과 activation buffer가 추가되므로, 단순히 partition 수를 늘리는 것이 항상 빠른 것은 아니다.

Micro-batch와 유효 batch

하나의 optimizer step에 포함되는 micro-batch 수를 m이라 하자. 각 micro-batch는 forward와 backward graph를 독립적으로 흘러가지만, 마지막에 계산된 gradient는 하나의 optimizer update에 누적된다. Pipeline 자체는 batch를 자동으로 쪼갤 수 있지만, loss scaling과 label chunking이 global batch 의미와 일치하는지 확인해야 한다.

global batch = micro_batch_size × micro_batches × data_parallel_degree

Pipeline stage 수가 p이고 stage 시간이 동일하다고 단순화하면, fill-drain schedule의 이상적인 pipeline utilization은 다음과 같이 근사할 수 있다.

utilization ≈ m / (m + p - 1)

이는 통신, 불균형, backward 비용 차이, input pipeline을 제외한 모델이다. 따라서 m을 키우면 bubble 비율은 줄지만, 동시에 더 많은 in-flight activation을 저장해야 하고 global batch와 optimizer semantics가 달라질 수 있다.

Schedule의 종류

Schedule	실행 패턴	장점	주요 비용/주의점
GPipe / fill-drain	모든 micro-batch forward 후 모든 backward	이해와 구현이 단순하고 gradient accumulation 경계가 명확	steady state 전후 bubble이 있고 activation을 많이 보관
1F1B	warmup 뒤 forward 1개와 backward 1개를 번갈아 실행	in-flight activation과 메모리 사용량을 줄임	schedule 의존성이 커지고 stage 간 순서 검증 필요
Interleaved 1F1B	한 rank가 여러 virtual stage를 맡아 1F1B 실행	stage granularity와 장치 활용도를 개선할 수 있음	virtual stage mapping, 통신, buffer 관리가 복잡
Zero-bubble 계열	weight-gradient 등 작업을 bubble에 배치	계산 분해가 맞으면 bubble을 더 줄일 수 있음	backward를 분리하고 schedule 제약을 만족해야 함

GPipe의 fill-drain은 pipeline을 먼저 채운 뒤 forward가 끝난 micro-batch를 역순으로 backward하는 형태다. 1F1B는 warmup 이후 각 stage가 forward와 backward를 교차 실행하므로 장치당 저장 activation 수가 줄어드는 대신, micro-batch의 의존성과 update 시점을 정확하게 관리해야 한다. 최신 runtime은 이들 schedule과 사용자 정의 schedule을 함께 제공하지만, API 안정성은 프레임워크 버전에 따라 다르다.

비교/분석

병렬화 방식 비교

방식	분할 단위	주 통신	메모리 효과	적합한 상황
Data Parallelism	입력 batch	gradient AllReduce	모델 state가 device마다 복제	모델이 한 device에 들어가고 throughput을 키울 때
Tensor Parallelism	한 레이어의 tensor/행렬	layer 내부 AllReduce 또는 AllGather	한 레이어의 계산과 state를 분산	한 레이어가 크거나 빠른 scale-up 링크가 있을 때
Pipeline Parallelism	연속된 레이어 stage	activation/gradient P2P	모델 깊이 방향으로 parameter 분산	모델이 여러 레이어에 걸쳐 너무 클 때
FSDP / ZeRO	parameter, gradient, optimizer state shard	AllGather/ReduceScatter	모델 state 복제량을 절감	state 메모리가 병목이고 통신을 감당할 때
Hybrid 3D parallelism	data + tensor + pipeline	각 축의 collective/P2P	각 축의 장점을 결합	대규모 Transformer training

PP는 레이어 경계의 activation만 통신하므로, 각 레이어 내부를 분할하는 tensor parallelism보다 통신 빈도가 낮을 수 있다. 대신 stage 사이에 순차 의존성이 있어 pipeline bubble이 발생하고, 느린 stage 하나가 전체 clock을 결정한다. 실제 대규모 학습에서는 보통 pipeline group 안에서 PP를 적용하고, 같은 stage 위치의 여러 replica 사이에는 data parallelism을 적용하며, stage 내부에는 tensor parallelism을 결합한다.

메모리와 throughput의 trade-off

PP의 가장 직접적인 이점은 parameter와 그에 대응하는 gradient/optimizer state를 stage별로 나누는 것이다. 그러나 stage가 backward를 수행하려면 각 in-flight micro-batch의 activation 또는 activation checkpoint를 유지해야 한다. 따라서 parameter 메모리가 줄어도 micro-batch 수와 sequence length가 커지면 activation 메모리가 새로운 병목이 된다.

최적화	줄이는 대상	대가
micro-batch 수 조정	bubble 또는 activation buffer	bubble과 global batch 사이의 절충
activation checkpointing	저장 activation	backward에서 일부 forward 재계산
interleaved schedule	stage 불균형과 bubble	더 복잡한 virtual stage/buffer
activation 통신 압축/overlap	P2P 통신 시간	추가 연산, 정확도·구현 검증
PP + FSDP/ZeRO	parameter/state 메모리	collective 통신과 checkpoint 복잡도

동작 원리

1. 모델 분할과 process mapping

먼저 모델의 forward 실행 순서를 기준으로 split point를 정한다. Stage 0은 입력 embedding 또는 초기 block을, 마지막 stage는 출력 head와 loss를 담당하는 식이다. 각 rank는 자신의 stage만 materialize하고, pipeline group 내 이웃 rank와 P2P communicator를 만든다. Data parallelism을 함께 사용한다면 같은 stage index를 가진 rank끼리 별도의 data-parallel group을 구성한다.

PyTorch의 torch.distributed.pipelining은 수동으로 stage module을 만들거나 torch.export 기반 frontend로 실행 graph를 분할할 수 있다. PipelineStage는 stage의 input/output shape과 dtype에 맞춰 통신 buffer를 준비하므로 dynamic shape 모델은 runtime 제약을 확인해야 한다. DeepSpeed PipelineModule도 forward가 순서형 layer sequence로 표현되고 stage 사이의 interface가 단순해야 한다는 제약을 둔다.

2. Micro-batch forward

입력 batch를 m개 micro-batch로 나눈다. 첫 번째 micro-batch가 Stage 0에서 계산되는 동안 다른 stage는 아직 입력을 기다리지만, 이후에는 다음 흐름이 겹친다.

Stage 0: F(m0) -> F(m1) -> F(m2) -> ...
Stage 1:          F(m0) -> F(m1) -> ...
Stage 2:                   F(m0) -> ...

각 경계에서 activation_send(stage_i -> stage_i+1)와 activation_recv가 대응해야 한다. P2P 통신이 blocking이면 send/recv 순서가 맞지 않을 때 deadlock이 발생할 수 있으므로, runtime은 통신 순서와 buffer slot을 schedule에 맞춰 관리한다.

3. Warmup, steady state, cooldown

Fill-drain 또는 1F1B의 전체 실행은 세 구간으로 이해할 수 있다.

Warmup: 앞 stage가 여러 micro-batch를 forward하여 뒤 stage를 채운다. 뒤쪽 stage는 아직 backward할 출력이 없으므로 activation을 buffer에 저장한다.
Steady state: 앞쪽 stage의 새 forward와 뒤쪽 stage의 backward가 동시에 진행된다. 1F1B에서는 각 stage가 가능한 한 하나의 forward와 하나의 backward를 교대로 처리한다.
Cooldown: 더 이상 새 forward가 없을 때 남은 micro-batch의 backward를 끝내고 gradient를 누적한다.

4. Backward와 gradient 전달

마지막 stage가 loss를 계산하면 backward가 시작된다. Stage i+1은 자신이 받은 activation에 대한 gradient를 계산하고 이를 Stage i로 보낸다. Stage i는 저장해 둔 forward activation과 수신한 gradient로 local backward를 수행하고 parameter gradient를 누적한다. 이 과정은 첫 stage까지 반복된다.

forward:  h_i = stage_i(h_{i-1})
backward: d h_{i-1} = backward_i(d h_i, saved_activation_i)

한 optimizer step 동안 여러 micro-batch의 gradient를 누적할 때는 loss를 micro-batch 수에 맞게 scale해야 한다. Data parallel replica가 있다면 stage별 parameter gradient를 해당 data-parallel group에서 Reduce 또는 AllReduce한 후 optimizer를 실행한다. Pipeline group 자체의 stage 사이에서 gradient를 AllReduce하는 것은 PP의 기본 통신이 아니다.

5. Optimizer step과 weight version

동기식 GPipe/1F1B는 정해진 micro-batch의 backward가 모두 끝난 뒤 update하는 방식으로 구현할 수 있다. 반면 PipeDream 계열의 일부 schedule은 pipeline을 계속 채우기 위해 forward와 backward가 서로 다른 weight version을 사용할 수 있다. 이 경우 weight staleness, optimizer semantics, checkpoint 일관성을 별도로 분석해야 하므로 단순한 synchronous gradient accumulation과 같은 결과라고 가정하면 안 된다.

PyTorch pipeline API의 개념 예시

아래 코드는 PyTorch torch.distributed.pipelining의 stage와 schedule 연결 방식을 보여주는 축약 예시다. 실제 모델에서는 각 rank가 자신에게 해당하는 stage_module을 구성하고, stage 간 tensor shape/dtype과 loss 처리 방식을 맞춰야 한다.

import os
import torch
import torch.distributed as dist
from torch.distributed.pipelining import PipelineStage, Schedule1F1B

dist.init_process_group(backend="nccl")
rank = dist.get_rank()
world_size = dist.get_world_size()
device = torch.device("cuda", int(os.environ["LOCAL_RANK"]))

# rank별로 모델의 연속된 partition을 구성한다.
stage_module = build_stage_module(stage_index=rank, num_stages=world_size)
stage = PipelineStage(stage_module, rank, world_size, device)
schedule = Schedule1F1B(stage, n_microbatches=8, loss_fn=loss_fn)

if rank == 0:
    output = schedule.step(input_batch, target=target_batch)
else:
    output = schedule.step()

PyTorch 문서의 schedule은 whole batch 입력을 받아 micro-batch로 나누고, rank 0이 입력을 공급하며 마지막 stage가 출력과 loss를 처리하는 형태를 지원한다. ScheduleGPipe는 fill-drain, Schedule1F1B는 steady state에서 forward/backward를 교차 실행한다. PipelineStage가 기대하는 shape과 실제 micro-batch shape이 다르면 통신 buffer 오류가 발생할 수 있으므로 sequence length를 동적으로 바꾸는 모델은 별도 설계가 필요하다.

장단점

장점

모델 용량 확장: 연속된 레이어와 optimizer state를 여러 device에 나눠 단일 GPU 한계를 넘을 수 있다.
통신량의 예측 가능성: 일반적인 PP는 stage 경계의 activation과 gradient를 P2P로 교환하므로 layer 내부 collective보다 통신 경계가 명확하다.
대역폭 제한 환경에 유리: 충분한 micro-batch와 균형 잡힌 stage를 사용하면 통신과 계산을 overlap할 수 있다.
다른 병렬화와 결합 가능: data parallelism 및 tensor parallelism과 함께 3D parallelism을 구성할 수 있다.
추론에도 적용 가능: batch request 또는 token micro-batch를 pipeline에 흘려 여러 stage를 동시에 사용할 수 있다.

단점

Pipeline bubble: stage 수가 많고 micro-batch 수가 작으면 warmup/cooldown 유휴 시간이 커진다.
Partition 불균형: 한 stage의 forward/backward가 느리면 모든 stage가 그 속도에 맞춰 기다린다.
Activation 메모리: in-flight micro-batch 수가 증가할수록 저장 activation과 통신 buffer가 커진다.
Schedule 복잡도: forward/backward 의존성, buffer slot, send/recv 순서를 틀리면 hang이나 잘못된 gradient가 생긴다.
모델 구조 제약: 임의의 skip connection, dynamic control flow, stage 경계를 넘는 tied weight는 분할과 통신을 복잡하게 만든다.
Checkpoint와 장애 복구 부담: stage별 state, optimizer, schedule 위치를 일관되게 저장하고 복구해야 한다.
추론 latency와 throughput의 절충: 단일 요청은 여러 stage를 순차로 통과하므로 PP가 latency를 줄인다고 보장할 수 없고, 충분한 동시 micro-batch가 있어야 throughput 이점이 나타난다.

기술	역할	PP와의 관계
PyTorch `torch.distributed.pipelining`	model split frontend와 distributed runtime	GPipe, 1F1B, interleaved 및 custom schedule 제공
DeepSpeed `PipelineModule`	sequence layer 기반 pipeline training	layer partition, topology, activation checkpointing 지원
Megatron-LM	Transformer 대규모 학습 프레임워크	tensor/pipeline/data parallelism과 interleaved schedule 결합
GPipe	synchronous batch-splitting pipeline	fill-drain과 gradient accumulation의 대표 방식
PipeDream	pipeline execution과 weight versioning	비동기성 및 schedule의 throughput/정확도 trade-off 제시
NCCL / point-to-point transport	GPU 간 데이터 전송	PP는 주로 이웃 stage 간 activation/gradient 전달에 사용

핵심 정리

파이프라인 병렬화는 순차 모델을 stage로 나누고 여러 micro-batch를 겹쳐 실행해 모델 용량과 처리량을 확장한다.
pipeline degree, micro-batch 수, stage 계산량 균형이 pipeline bubble과 실제 utilization을 결정한다.
GPipe는 단순한 fill-drain, 1F1B는 activation 메모리와 steady-state 효율을 고려한 대표 schedule이다.
PP는 parameter를 분산하지만 activation buffer와 stage 경계 P2P 통신을 추가하므로 FSDP, tensor parallelism과 함께 전체 메모리·통신 예산을 설계해야 한다.
대규모 학습에서는 data + tensor + pipeline parallelism을 조합하되, process group, weight update 시점, checkpoint와 shape 제약을 함께 검증해야 한다.