torch.distributed.pipelining — GPipe/1F1B 스케줄의 표준 API
torch.distributed.pipelining — GPipe/1F1B 스케줄의 표준 API
PyTorch가 제공하는 분산 파이프라인 병렬 API예요. GPT 스타일 트랜스포머를 여러 스테이지로 나눠 마이크로배치 단위로 스케줄링하는 방법을 보여줘요.
모델 분할 방법
- 수동 분할: 스테이지마다 필요한 레이어만 남기고 나머지를 직접 제거.
- 트레이서 기반 분할:
split_spec={"layers.4": SplitPoint.BEGINNING}처럼 중단점을 지정하면 자동으로 나눠요.
스케줄
단일 스테이지-per-랭크 스케줄로 GPipe와 1F1B를, 다중 스테이지-per-랭크로 Interleaved1F1B, LoopedBFS를 지원해요. 투토리얼에서는 ScheduleGPipe를 사용해 2개의 프로세스(랭크 0,1)로 나눠 실행해요.
pipe = pipeline(module=model, mb_args=(example_input_microbatch,), split_spec={"layers.4": SplitPoint.BEGINNING})
stage = pipe.build_stage(stage_index, device, pp_group)
schedule = ScheduleGPipe(stage, n_microbatches=num_microbatches, loss_fn=tokenwise_loss_fn)
torchrun --nnodes 1 --nproc_per_node 2 pipelining_tutorial.py로 실행해요.