FSDP, ZeRO 예제

1. FSDP (Fully Sharded Data Parallel) 예제

핵심

• PyTorch의 torch.distributed.fsdp

• 모델 파라미터, 옵티마 상태, 그래디언트를 GPU에 분산 저장

Details

• model = FSDP(model)를 통해 모델 전체를 샤딩
  • 모든 GPU가 전체 모델을 저장하지 않고, 필요할 때만 해당 가중치를 가져와 사용.
• PyTorch 기본 분산 학습(torch.distributed)을 활용하여 DeepSpeed 없이도 실행 가능!

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp.wrap import auto_wrap

from transformers import AutoModelForCausalLM

# 분산 학습 초기화
dist.init_process_group("nccl")
torch.cuda.set_device(dist.get_rank())

# 모델 로드
model = AutoModelForCausalLM.from_pretrained("gpt2").cuda()
model = FSDP(model)  # 모델을 Fully Sharded Data Parallel로 감싸기

# 옵티마이저 정의
optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5)

# 더미 입력
inputs = torch.randint(0, 50257, (1, 128)).cuda()

# 학습 루프
for epoch in range(1):
    optimizer.zero_grad()
    outputs = model(inputs, labels=inputs)
    loss = outputs.loss
    loss.backward()
    optimizer.step()

2. ZeRO (Zero Redundancy)

2-1. Zero-1 (Optimizer State Sharding) 예제

ZeRO Stage 1은 옵티마이저 상태를 샤딩하여 메모리 사용량을 줄임.

• 옵티마이저 상태(Optimizer State)를 분산 저장하여 GPU 메모리 절약.

• 모델 전체 가중치는 모든 GPU가 가지고 있음.

• deepspeed.initialize()를 통해 자동으로 ZeRO-1을 활성화

import deepspeed
from transformers import AutoModelForCausalLM

# 모델 로드
model = AutoModelForCausalLM.from_pretrained("gpt2")

# DeepSpeed ZeRO-1 설정
ds_config = {
    "train_batch_size": 8,
    "zero_optimization": {
        "stage": 1  # ZeRO Stage 1 활성화
    }
}

# DeepSpeed 초기화
model, optimizer, _, _ = deepspeed.initialize(model=model, config=ds_config)

# 학습 예제
inputs = torch.randint(0, 50257, (1, 128)).cuda()
outputs = model(inputs, labels=inputs)
loss = outputs.loss
model.backward(loss)  # DeepSpeed가 제공하는 backward()
model.step()

2-2. ZeRO-2 (Optimizer State + Gradient Sharding) 예제

ZeRO Stage 2는 옵티마이저 상태 + 그래디언트를 샤딩하여 추가적인 메모리 절약 가능.

• 옵티마이저 상태를 CPU로 옮기는 offload_optimizer 옵션을 사용 가능.

모델 가중치는 여전히 모든 GPU가 저장.

import deepspeed
from transformers import AutoModelForCausalLM

# 모델 로드
model = AutoModelForCausalLM.from_pretrained("gpt2")

# DeepSpeed ZeRO-2 설정
ds_config = {
    "train_batch_size": 8,
    "zero_optimization": {
        "stage": 2,  # ZeRO Stage 2 활성화
        "offload_optimizer": {"device": "cpu"}  # 옵티마이저를 CPU로 오프로딩 가능
    }
}

# DeepSpeed 초기화
model, optimizer, _, _ = deepspeed.initialize(model=model, config=ds_config)

# 학습 예제
inputs = torch.randint(0, 50257, (1, 128)).cuda()
outputs = model(inputs, labels=inputs)
loss = outputs.loss
model.backward(loss)
model.step()

2-3. ZeRO-3 (Full Model Sharding) 예제

ZeRO Stage 3는 모델 파라미터까지 포함하여 모든 요소를 샤딩.

• 모델 파라미터(Weights), 옵티마 상태, 그래디언트까지 모든 요소를 샤딩.

• 모든 GPU가 모델의 일부만 저장하며, 필요할 때 해당 부분만 로딩하여 연산.

• 메모리를 가장 많이 절약할 수 있는 방식

import deepspeed
from transformers import AutoModelForCausalLM

# 모델 로드
model = AutoModelForCausalLM.from_pretrained("gpt2")

# DeepSpeed ZeRO-3 설정
ds_config = {
    "train_batch_size": 8,
    "zero_optimization": {
        "stage": 3  # ZeRO Stage 3 활성화 (모델 전체 샤딩)
    }
}

# DeepSpeed 초기화
model, optimizer, _, _ = deepspeed.initialize(model=model, config=ds_config)

# 학습 예제
inputs = torch.randint(0, 50257, (1, 128)).cuda()
outputs = model(inputs, labels=inputs)
loss = outputs.loss
model.backward(loss)
model.step()

3. 요약

	FSDP	ZeRO-1	ZeRO-2	ZeRO-3
프레임워크	PyTorch	DeepSpeed	DeepSpeed	DeepSpeed
Optimizer State Sharding	✔	✔	✔	✔
Gradient Sharding	✔	❌	✔	✔
Weight Sharding	✔	❌	❌	✔
활성화 체크포인팅	❌ (직접 설정 필요)	❌	❌	✔
사용 편의성	PyTorch 네이티브	DeepSpeed 설정 필요	DeepSpeed 설정 필요	DeepSpeed 설정 필요
모델 크기 축소 효과	높음	낮음	중간	최고

4. (Optimizer state) Offloading

	ZeRO-1	ZeRO-2	ZeRO-3
옵티마이저 상태(Optimizer State) 샤딩	✔	✔	✔
옵티마이저 오프로딩 가능 여부	❌	✔	✔
모델 가중치(Weights) 샤딩	❌	❌	✔

(1) ZeRO-1

• 옵티마이저 상태(Optimizer State)를 GPU별로 샤딩하지만, 옵티마이저를 CPU로 옮기는 기능은 없음.
• 즉, 옵티마이저는 모든 GPU에서 유지되며 CPU 오프로딩이 지원되지 않음.

(2) ZeRO-2

• 옵티마이저 상태(Optimizer State)와 그래디언트(Gradients)를 샤딩.
• 추가적으로 "offload_optimizer": {"device": "cpu"}를 설정하면 옵티마이저 상태를 CPU로 옮길 수 있음.
• 옵티마이저의 메모리 사용량이 줄어 VRAM이 더 절약됨.
• 하지만 CPU ↔ GPU 간 데이터 이동이 필요하기 때문에 훈련 속도가 저하될 수 있음.

(3) ZeRO-3

• 모델의 모든 요소(모델 파라미터, 그래디언트, 옵티마이저 상태)를 샤딩.
• 옵티마이저가 이미 샤딩되므로 CPU 오프로딩 없이도 VRAM 사용이 최소화됨.
• "offload_optimizer": {"device": "cpu"} 옵션 없이도 VRAM 절약 효과가 크기 때문에, 보통 ZeRO-3에서는 따로 사용하지 않음.
• 다만, "offload_optimizer": {"device": "cpu"}를 추가하면 CPU 메모리를 더 활용하여 VRAM을 극단적으로 줄일 수 있지만, 속도 저하가 심해지는 단점이 있음.

(4) Summary

• ZeRO-2에서는 옵티마이저 상태가 GPU에 저장되므로 CPU로 옮기는 옵션이 존재
• ZeRO-3에서는 모델 파라미터까지 샤딩하므로, 기본적으로 CPU 오프로딩 없이도 VRAM을 충분히 절약 가능
• ZeRO-1은 옵티마이저 오프로딩을 지원하지 않음