PyTorch 분산 학습 기초: 데이터 병렬화, 모델 병렬화, 파이프라인 병렬화

딥러닝 모델이 점점 커지고 데이터도 방대해지면서, 단일 GPU나 서버만으로는 학습 속도가 너무 느리거나 GPU 메모리가 부족해 학습이 불가능해진다. 이를 해결하기 위해 여러 GPU를 동시에 활용해 모델을 학습시키는 것이 바로 분산 학습이다.

1. 분산 학습 종류

1.1 데이터 병렬화(Data Parallelism)

[전체 데이터] → [분할된 미니배치1] → GPU0 (모델 복제)
            → [분할된 미니배치2] → GPU1 (모델 복제)
            → [분할된 미니배치3] → GPU2 (모델 복제)

[각 GPU] → forward & backward → all-reduce → 동기화 → 파라미터 업데이트

 

가장 보편적으로 사용되는 방식이다. 동일한 모델을 여러 GPU에 복제하고, 미니배치 데이터를 GPU별로 나누어 처리한다. 각 GPU에서 forward와 backward를 계산한 뒤, all-reduce 연산으로 그래디언트를 평균내고 동기화해 파라미터를 업데이트한다.

PyTorch에는 DataParallel(DP)과 DistributedDataParallel(DDP) 두 가지가 있다. DP는 단일 머신에서 여러 GPU에 모델을 복제해 데이터만 나눠 넣는 방식이며 구현이 간단하지만, Python GIL(Global Interpreter Lock)과 single-process 구조로 인해 통신 병목이 심하다. 그래서 실제로는 거의 항상 DDP를 사용한다. DDP는 각 GPU마다 별도의 프로세스를 생성해 통신 병목을 줄이고 효율적으로 all-reduce를 수행해 그래디언트를 동기화한다.

• 예시: PyTorch DDP, TensorFlow MirroredStrategy
• 장점: 구현이 단순하고 GPU를 늘리기 쉽다.
• 단점: 모델 전체가 각 GPU에 복제되므로, GPU 메모리가 부족하면 불가능하다.

1.2 모델 병렬화(Model Parallelism)

[입력 데이터] → GPU0 (Model Layer1-5) → GPU1 (Model Layer6-10) → GPU2 (Model Layer11-15)

forward → GPU 간 연속 전달 → backward → GPU 간 연속 전달

 

모델 자체를 여러 GPU에 나눠 저장하고 순차적으로 forward 연산을 GPU를 거쳐 진행한다. GPT-3, T5 같은 하나의 GPU에 올라가지 않을 정도로 거대한 모델에서 주로 사용한다. PyTorch에서는 특정 레이어를 수동으로 cuda:0, cuda:1에 올리거나, Megatron-LM과 같은 프레임워크가 Tensor Parallelism을 통해 자동으로 레이어를 나눠준다.

• 예시: Megatron-LM (Tensor Parallelism), manual PyTorch split
• 장점: 초대규모 모델을 여러 GPU에 나눠서 처리 가능하다.
• 단점: GPU 간 통신량이 많아 latency와 bandwidth 병목이 발생하기 쉽다. Layer간 종속으로 인해 병렬성이 제한된다.

1.3 파이프라인 병렬화(Pipeline Parallelism)

[미니배치1] → GPU0 (Stage1) → GPU1 (Stage2) → GPU2 (Stage3)
[미니배치2] →               → GPU0 (Stage1) → GPU1 (Stage2) → GPU2 (Stage3)
[미니배치3] →                               → GPU0 (Stage1) → GPU1 (Stage2) → GPU2 (Stage3)

→ 파이프라인 채워서 bubble 최소화

모델을 여러 stage로 나눠 GPU에 배치하고, 데이터를 연속적으로 흘려보내 처리하는 방식이다. 미니배치를 더 작게 쪼개 pipeline을 채워 idle time(bubble)을 줄인다. 구조 자체는 모델 병렬화처럼 layer를 여러 GPU에 나눠놓지만, GPU0에서 미니배치1을 연산하고 GPU1로 넘긴 뒤 GPU0은 바로 미니배치2를 연산하기 시작해 GPU들이 쉬지 않도록 만드는 방식이다. Deepspeed Pipeline이나 PyTorch torch.distributed.pipeline.sync.Pipe가 이를 지원한다.

• 예시: Deepspeed Pipeline, PyTorch Pipe
• 장점: layer를 stage별로 분리해 메모리 부담을 더 세밀히 분산할 수 있다.
• 단점: 순차 처리 구조라서 latency가 증가하며, bubble이 발생하지 않도록 careful하게 microbatch를 조절해야 한다.

어떤 분산 학습을 언제 쓰나?

• 작은~중간 크기 모델 (수천만~수억 파라미터): 단순히 DDP만으로 충분. GPU를 여러 개 쓸수록 학습 속도가 선형적으로 증가한다.
• 하나의 GPU 메모리에 올라가지 않는 모델 (수십억 파라미터): 모델 병렬 또는 파이프라인 병렬을 적용해 GPU 메모리를 분산.
• 수십억~수백억 파라미터 이상 초거대 모델: DDP + 모델 병렬 + 파이프라인 병렬을 함께 조합한다. Megatron-LM, Deepspeed ZeRO-Infinity 같은 프레임워크가 이를 자동화해준다.

따라서 모델 크기, GPU 메모리, 네트워크 대역폭을 고려해 적절한 분산 방식 또는 조합을 선택하는 것이 중요하다.

 

2. PyTorch DistributedDataParallel(DDP) 사용하기

PyTorch DistributedDataParallel(DDP)은 데이터 병렬(Data Parallelism)을 위한 가장 표준적인 방법이다. 여러 GPU에 동일한 모델을 올리고, 각 GPU가 서로 다른 데이터 미니배치를 학습한 뒤, all-reduce를 통해 그래디언트를 동기화해 파라미터를 업데이트한다.

단일 머신(여러 GPU)에서도 쓸 수 있고, 멀티 노드(여러 서버)에서도 동일한 코드 구조를 사용할 수 있다.

all-reduce는 각 GPU에서 계산한 그래디언트를 서로 주고받아 평균을 구하고 동기화하는 통신 연산이다. 이를 통해 모든 GPU가 동일한 파라미터를 유지하면서 학습을 진행할 수 있다. 주로 NCCL로 고속 수행되며, DDP에서 자동으로 수행된다.

2.1 사전 준비

• 모든 머신(노드)에 동일한 코드, 데이터, 라이브러리가 설치되어 있어야 한다.
• CUDA와 NCCL이 설치되어 있어야 한다. (PyTorch의 backend로 주로 nccl을 사용)
• 단일 머신이면 CUDA_VISIBLE_DEVICES로 GPU를 관리하면 되고, 멀티 노드라면 MASTER_ADDR, MASTER_PORT 환경변수를 설정해야 한다.

멀티 노드 환경에서는 보통 MASTER_ADDR을 rank=0 머신의 IP로 설정한다.

export MASTER_ADDR=192.168.1.10
export MASTER_PORT=12355

 

2.2 DDP 기본 코드 구조

import os
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler

def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'
    dist.init_process_group(backend='nccl', rank=rank, world_size=world_size)

def cleanup():
    dist.destroy_process_group()

def train(rank, world_size):
    setup(rank, world_size)
    device = torch.device(f'cuda:{rank}')
    
    model = MyModel().to(device)
    ddp_model = DDP(model, device_ids=[rank])

    dataset = MyDataset()
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True)
    dataloader = DataLoader(dataset, batch_size=64, sampler=sampler)

    optimizer = torch.optim.Adam(ddp_model.parameters(), lr=1e-4)
    criterion = torch.nn.CrossEntropyLoss()

    for epoch in range(10):
        sampler.set_epoch(epoch)
        for batch in dataloader:
            inputs, targets = batch
            inputs, targets = inputs.to(device), targets.to(device)
            outputs = ddp_model(inputs)
            loss = criterion(outputs, targets)
            
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

    cleanup()

이 스크립트는 GPU 0, GPU 1, GPU 2... 각각이 독립적인 Python 프로세스로 실행되어, 동일한 모델을 각 GPU에서 학습하고 all-reduce를 통해 그래디언트를 동기화한다.

 

2.3 실제 실행하기

✅ 단일 머신

import torch.multiprocessing as mp

def main():
    world_size = 4
    mp.spawn(train, args=(world_size,), nprocs=world_size, join=True)

if __name__ == '__main__':
    main()

단일 머신에서 GPU가 4개라면 torch.multiprocessing.spawn을 사용해 위와 같이 실행할 수 있다.

 

✅ 멀티 머신

멀티 머신에서는 모든 머신에서 동일하게 export와 torchrun을 실행해야 한다. 즉, master(보통 rank=0) 머신에서만 환경 변수를 설정하거나 torchrun을 돌리는 것이 아니라, 모든 머신에서 환경 변수를 동일하게 설정하고 torchrun을 각각 실행해야 한다.

예를 들어 머신이 3대 있고, 각각 GPU가 4개씩 있는 경우 다음과 같이 한다.

export MASTER_ADDR=192.168.1.10
export MASTER_PORT=12355

# 머신1 (rank=0)
torchrun --nnodes=3 --nproc_per_node=4 --node_rank=0 train.py

# 머신2 (rank=1)
torchrun --nnodes=3 --nproc_per_node=4 --node_rank=1 train.py

# 머신3 (rank=2)
torchrun --nnodes=3 --nproc_per_node=4 --node_rank=2 train.py

• 모든 머신에서 동일하게 export 한다. MASTER_ADDR은 master 역할을 할 rank=0 머신의 IP를 적는다.
  • “통신을 위한 기준이 되는 마스터 IP(=rank=0 머신)”를 알려주는 것
• 그리고 각 머신에서 자신의 node_rank에 맞게 torchrun을 실행한다.
• 즉, 각 머신이 모두 동시에 이 명령을 실행해야 DDP가 정상적으로 통신을 시작한다.
• 마스터 머신만 실행하는 것이 아니며, 모든 머신이 자신의 node_rank를 지정해 동일한 train.py를 실행해야 한다.

2.4 주요 개념과 주의사항

 

• rank 는 전체 프로세스에서 고유 번호를 말한다. 예를 들어 8 GPU면 rank=0~7까지 있다.
• world_size 는 전체 GPU(=전체 프로세스) 수를 의미한다.
• DistributedSampler 는 각 프로세스가 같은 데이터셋을 서로 다른 순서/범위로 읽게 한다. epoch마다 set_epoch(epoch)를 호출해야 데이터 셔플링이 잘 동작한다.
• DDP는 DataParallel(DP)과 달리 Python GIL 병목 없이 각 프로세스가 GPU 하나씩을 전담하므로 훨씬 효율적이다.

PyTorch DDP는 단일 머신, 멀티 머신 모두에서 사용할 수 있는 강력하고 표준적인 데이터 병렬 프레임워크이다. setup → DDP로 모델 감싸기 → DistributedSampler → backward에서 all-reduce 순으로 동작한다. 이를 통해 동일한 모델을 GPU마다 독립적으로 학습하고, 그래디언트를 자동으로 통신·동기화해 효율적으로 대규모 모델을 학습할 수 있다.

 

 

3. Slurm에서 멀티 머신 PyTorch DDP 실행하기

멀티 노드 환경에서 Slurm을 사용하면 srun으로 여러 머신에서 동시에 torchrun을 실행할 수 있어 편리하다. Slurm이 자동으로 MASTER_ADDR, RANK, WORLD_SIZE 같은 환경 변수를 세팅해 주기 때문에 따로 export 할 필요도 없다.

 

하지만, 슬럼은 GPU 서버를 클러스터로 구성해 놓은 환경에서만 쓸 수 있는 도구라 그냥 EC2 여러 대를 띄운다고 곧장 사용할 수 있는 게 아니고, 따로 Slurm 생태계(=slurmd, slurmctld, config)를 설치 + 설정해야 한다.

 

3.1 Slurm 스크립트 예제 (multi-node DDP)

아래는 Slurm job 스크립트(train_job.sh) 예제이다.

#!/bin/bash
#SBATCH --job-name=ddp_train
#SBATCH --nodes=3              # 사용할 노드 수
#SBATCH --ntasks-per-node=4    # 노드당 GPU 개수
#SBATCH --gres=gpu:4           # 노드당 GPU 할당
#SBATCH --cpus-per-task=4
#SBATCH --time=24:00:00
#SBATCH --partition=gpu

module load cuda/12.1
module load python/3.11

srun torchrun \
    --nnodes=$SLURM_JOB_NUM_NODES \
    --nproc_per_node=$SLURM_NTASKS_PER_NODE \
    --node_rank=$SLURM_NODEID \
    train.py

#SBATCH --nodes=3, #SBATCH --ntasks-per-node=4 이라고 설정했으므로

슬럼 변수	실제 예시 값
$SLURM_JOB_NUM_NODES	3
$SLURM_NTASKS_PER_NODE	4
$SLURM_NODEID	각 노드에서 0,1,2

슬럼 변수는 위와 같이 세팅되고, 그러면 srun이 각 노드에서 이렇게 실행하게 됩니다.

# 노드1에서
torchrun --nnodes=3 --nproc_per_node=4 --node_rank=0 train.py

# 노드2에서
torchrun --nnodes=3 --nproc_per_node=4 --node_rank=1 train.py

# 노드3에서
torchrun --nnodes=3 --nproc_per_node=4 --node_rank=2 train.py

특정 노드를 사용해서 job을 실행하고 싶으면 --nodelist=node05,node06,node07 또는 --nodelist=node[05-07] 와 같이 세팅하면 된다.

 

3.2 단계별 동작 및 특징

1. srun이 Slurm에서 각 노드에 작업을 자동으로 배포한다.
2. Slurm이 MASTER_ADDR, MASTER_PORT, RANK, WORLD_SIZE, NODEID 등을 자동으로 설정해 각 노드가 동일하게 공유한다.
3. torchrun은 이 환경 변수를 기반으로 각 노드가 통신할 마스터를 자동으로 찾고 init_process_group()을 통해 연결을 맺는다.
4. 모든 노드에서 동시에 forward, backward를 수행하고, all-reduce를 통해 그래디언트를 동기화한다.

또한,

• 표준 출력과 에러 로그를 slurm-<jobid>.out에 자동 저장한다.
• squeue, scontrol show job <jobid> 같은 명령어로 상태(대기, 실행 중, 종료, 실패)를 쉽게 확인할 수 있다.
• Slurm Array Job, Checkpointing, Preemption 같은 기능을 통해 장기 학습에서 서버가 다운되더라도 복구하거나 재시작할 수 있는 기능을 지원한다.
• sbatch를 여러 번 제출하거나, Array Job(--array=0-9)을 사용해 하이퍼파라미터 튜닝 실험을 한꺼번에 10개, 100개도 동시에 돌릴 수 있다.

 

3.3 Slurm에서 멀티 노드 PyTorch DDP 학습하기

즉 위의 스크립트를 작성한 뒤, Slurm에 제출하면 된다.

sbatch train_job.sh

이렇게 하면 Slurm이 자동으로 3대 머신에 각각 torchrun을 올바르게 실행시켜 멀티 노드 DDP 학습을 시작한다.

즉, 원래는 각 머신에 직접 들어가서 torchrun을 실행해야 하는데, Slurm이 이를 대신해서 각 노드에 들어가 torchrun을 자동으로 실행해 주기 때문에 사용자는 메인 서버(=Slurm 컨트롤 노드)에서 sbatch train_job.sh 한 번만 실행하면 된다.

 

💡 정리

• Slurm + srun을 사용하면 여러 머신에서 rank, master ip를 직접 지정하지 않아도 된다.
• Slurm이 --node_rank에 필요한 환경 변수를 자동으로 잡아주므로 매우 편리하다.
• 따라서 대규모 GPU 클러스터 환경에서는 Slurm과 PyTorch DDP를 결합해 멀티 노드 학습을 안정적이고 쉽게 구현할 수 있다.

 

4. PyTorch로 모델 병렬화 구현하기

딥러닝 모델이 하나의 GPU에 올리기에는 너무 커지면, 모델 병렬(Model Parallelism) 이 필요하다. PyTorch는 DataParallel, DistributedDataParallel처럼 자동으로 모델을 여러 GPU에 복제해 데이터만 분산하는 데이터 병렬 방식과 달리, 모델 병렬은 사용자가 직접 모델의 레이어를 GPU에 나눠 배치해 forward, backward를 GPU 간에 순차적으로 흘려보내는 방식이다.

4.1 기본 수동 모델 병렬화 (Manual Model Parallelism)

가장 기본적인 모델 병렬화는 모델의 일부 레이어를 cuda:0에, 나머지를 cuda:1에 올려서 forward 계산 시 데이터를 GPU 간에 전송하도록 만드는 것이다.

 

✅ 단계별 구현

1. 모델의 각 파트를 원하는 GPU에 올린다.
2. forward 함수에서 to('cuda:x')를 통해 출력을 다음 GPU로 보낸다.
3. backward는 PyTorch의 Autograd가 자동으로 GPU 간 통신을 통해 gradient를 계산한다.

✅ 예시 코드

import torch
import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.layer1 = nn.Linear(1024, 2048).to('cuda:0')
        self.relu = nn.ReLU()
        self.layer2 = nn.Linear(2048, 1024).to('cuda:1')

    def forward(self, x):
        x = x.to('cuda:0')
        x = self.layer1(x)
        x = self.relu(x)
        x = x.to('cuda:1')
        x = self.layer2(x)
        return x

# 사용 예시
model = MyModel()
input_data = torch.randn(64, 1024).to('cuda:0')
output = model(input_data)

이 구조에서는 layer1은 GPU0에서, layer2는 GPU1에서 수행된다. 데이터는 forward 시 GPU0 → GPU1, backward 시 GPU1 → GPU0 순서로 통신하며 자동으로 gradient가 계산된다.

 

4.2 Pipeline Parallelism 사용하기 (torch.distributed.pipeline.sync.Pipe)

PyTorch는 Pipe 모듈을 통해 모델을 stage 단위로 나누어 여러 GPU에 배치하고, 입력을 micro-batch로 잘게 나눠 파이프라인을 채워 throughput을 극대화하는 기능을 제공한다.

 

✅ 단계별 구현

1. 모델을 torch.nn.Sequential로 정의한다.
2. Pipe를 사용해 각 stage를 지정한 GPU에 자동으로 올리도록 한다.
3. chunks 파라미터를 사용해 입력을 잘게 쪼개 pipeline bubble을 줄인다.

✅ 예시 코드

import torch
import torch.nn as nn
from torch.distributed.pipeline.sync import Pipe

# 모델을 stage로 나눌 Sequential 정의
model = nn.Sequential(
    nn.Linear(1024, 2048),
    nn.ReLU(),
    nn.Linear(2048, 1024)
)

# Pipe를 통해 GPU 2개에 나누어 배치, 입력을 8개로 쪼개 pipeline 처리
model = Pipe(model, devices=['cuda:0', 'cuda:1'], chunks=8)

# forward 실행
input_data = torch.randn(64, 1024).to('cuda:0')
output = model(input_data)

이렇게 하면 PyTorch가 자동으로 stage1은 cuda:0, stage2는 cuda:1에 올리고, 입력을 8개로 나눠 파이프라인을 돌리며 GPU들이 쉬지 않도록 학습을 수행한다.

 

4.3 모델 병렬화를 쓸 때 주의할 점

• 모델 병렬은 데이터 병렬(DDP)처럼 각 GPU가 같은 모델을 복제하는 방식이 아니므로, GPU 간 통신이 잦아 네트워크 병목이 발생하기 쉽다.
• 파라미터가 GPU마다 나뉘어 있으므로, 모델 저장(Checkpoint) 시 각 GPU의 파라미터를 따로 저장해야 한다.
• Pipe를 쓸 경우 chunks를 적절히 조절해 bubble(빈 시간)을 최소화하는 것이 중요하다.

 

PyTorch는 모델 병렬을 수동으로 구현하거나 Pipe를 이용해 파이프라인 병렬을 구현할 수 있다. 이는 하나의 GPU에 담기 어려운 초대형 모델을 학습할 때 유용하며, 필요에 따라 Megatron-LM, Deepspeed 같은 더 발전된 프레임워크로 넘어갈 수 있다.