이미지 로딩 중...
AI Generated
2025. 11. 11. · 4 Views
바닥부터 만드는 ChatGPT 9편 대규모 사전 학습 파이프라인 완벽 가이드
ChatGPT와 같은 대규모 언어 모델을 학습시키는 파이프라인을 처음부터 구축하는 방법을 다룹니다. 데이터 전처리부터 분산 학습, 체크포인트 관리까지 실무에서 바로 활용할 수 있는 구체적인 구현 방법을 소개합니다.
목차
- 대규모 데이터셋 전처리 파이프라인 - 효율적인 토크나이징과 청킹
- 분산 데이터 병렬 학습 설정 - PyTorch DDP 구현
- 효율적인 체크포인트 관리 시스템 - 증분 저장과 복구
- 학습률 스케줄링과 워밍업 - Cosine Annealing with Linear Warmup
- 그래디언트 클리핑과 누적 - 메모리 효율적인 대규모 배치 학습
- 효율적인 어텐션 메커니즘 - Flash Attention과 메모리 최적화
- 혼합 정밀도 학습 - FP16/BF16으로 학습 가속화
- 데이터 로더 최적화 - 병렬 전처리와 프리페칭
1. 대규모 데이터셋 전처리 파이프라인 - 효율적인 토크나이징과 청킹
시작하며
여러분이 수십 GB의 텍스트 데이터로 언어 모델을 학습시키려고 할 때 이런 상황을 겪어본 적 있나요? 메모리가 부족해서 데이터를 한 번에 로드할 수 없고, 토크나이징에만 며칠이 걸리는 상황 말이죠.
이런 문제는 실제 개발 현장에서 자주 발생합니다. 특히 GPT나 BERT 같은 대규모 언어 모델을 처음 학습시켜보려는 팀에서는 데이터 전처리 단계에서부터 막혀버리는 경우가 많습니다.
잘못 설계된 파이프라인은 전체 학습 시간의 80% 이상을 데이터 로딩에 소비하게 만들 수 있습니다. 바로 이럴 때 필요한 것이 스트리밍 방식의 데이터 전처리 파이프라인입니다.
데이터를 작은 청크로 나누어 처리하고, 멀티프로세싱으로 병렬화하며, 전처리된 결과를 효율적으로 캐싱하는 방식으로 이 문제를 해결할 수 있습니다.
개요
간단히 말해서, 대규모 데이터셋 전처리 파이프라인은 수십~수백 GB의 텍스트 데이터를 메모리에 모두 올리지 않고도 효율적으로 토크나이징하고 모델 학습에 적합한 형태로 변환하는 시스템입니다. 왜 이 파이프라인이 필요한지 실무 관점에서 설명하자면, ChatGPT 같은 모델은 최소 수백 GB의 텍스트 데이터로 학습됩니다.
이를 일반적인 방법으로 처리하려면 수 TB의 RAM이 필요하지만, 스트리밍 파이프라인을 사용하면 16GB RAM으로도 충분히 처리할 수 있습니다. 예를 들어, 위키피디아 전체 덤프(약 20GB)를 토크나이징할 때 일반 방식은 40GB 이상의 메모리를 필요로 하지만, 청킹 방식은 2GB 이하로 처리 가능합니다.
전통적인 방법과의 비교를 해보면, 기존에는 모든 데이터를 메모리에 로드한 후 한 번에 토크나이징했다면, 이제는 파일을 청크 단위로 읽어가며 점진적으로 처리하고 디스크에 저장할 수 있습니다. 이 파이프라인의 핵심 특징은 첫째, 메모리 효율성(스트리밍 처리), 둘째, 병렬 처리(멀티프로세싱), 셋째, 재사용성(전처리 결과 캐싱)입니다.
이러한 특징들이 중요한 이유는 실제 프로덕션 환경에서 제한된 자원으로 대규모 데이터를 다뤄야 하기 때문입니다.
코드 예제
import multiprocessing as mp
from transformers import GPT2Tokenizer
from datasets import load_dataset
import numpy as np
# 토크나이저 초기화
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
tokenizer.pad_token = tokenizer.eos_token
# 청크 단위로 데이터 처리하는 함수
def tokenize_function(examples, max_length=512):
# 배치 단위로 토크나이징 - 메모리 효율적
tokens = tokenizer(
examples['text'],
truncation=True,
padding='max_length',
max_length=max_length,
return_tensors='np'
)
return tokens
# 스트리밍 방식으로 데이터셋 로드 및 전처리
dataset = load_dataset('openwebtext', split='train', streaming=True)
tokenized_dataset = dataset.map(
tokenize_function,
batched=True, # 배치 처리로 속도 향상
batch_size=1000, # 청크 크기 조절
remove_columns=['text'] # 원본 텍스트 제거로 메모리 절약
)
# 전처리된 데이터를 디스크에 저장
tokenized_dataset.save_to_disk('./preprocessed_data')
설명
이것이 하는 일: 이 파이프라인은 수십 GB의 텍스트 데이터를 작은 청크로 나누어 순차적으로 읽고, 각 청크를 토크나이징한 후 디스크에 저장하는 방식으로 동작합니다. 전체 데이터를 메모리에 올리지 않고도 대규모 데이터셋을 처리할 수 있는 핵심 메커니즘입니다.
첫 번째로, load_dataset에서 streaming=True 옵션을 사용하여 데이터를 스트리밍 방식으로 로드합니다. 이렇게 하면 전체 데이터셋이 아닌 현재 처리 중인 배치만 메모리에 올라가므로, 20GB 데이터셋도 2GB 메모리로 처리할 수 있습니다.
왜 이렇게 하는지 이해하려면, 일반적인 load_dataset은 전체 파일을 Arrow 포맷으로 메모리에 로드하는 반면, 스트리밍은 파일 포인터만 유지하고 필요한 부분만 읽어온다는 점을 알아야 합니다. 그 다음으로, map 함수에 batched=True와 batch_size=1000을 설정하여 1000개씩 묶어서 토크나이징을 실행합니다.
내부에서는 멀티프로세싱을 통해 여러 배치를 병렬로 처리하므로, 단일 프로세스 대비 CPU 코어 수만큼 속도가 빨라집니다. 예를 들어 8코어 CPU에서는 약 6-7배 빠른 처리가 가능합니다.
토크나이징 함수 내부에서는 truncation=True로 최대 길이를 초과하는 텍스트를 잘라내고, padding='max_length'로 모든 시퀀스를 동일한 길이로 맞춥니다. 마지막으로, remove_columns=['text']로 원본 텍스트를 제거하여 저장 공간을 절약합니다.
토큰 ID만 저장하면 텍스트 대비 약 1/4 크기로 줄어듭니다. 여러분이 이 코드를 사용하면 100GB 데이터셋도 일반 워크스테이션에서 처리할 수 있고, 전처리 결과를 재사용하여 실험 반복 시간을 크게 단축할 수 있습니다.
실무에서는 전처리에 하루가 걸려도 이후 수백 번의 실험에서 재사용되므로 초기 투자 대비 엄청난 시간 절약 효과를 얻게 됩니다.
실전 팁
💡 batch_size는 메모리와 속도의 트레이드오프입니다. 메모리가 충분하다면 1000-5000으로 설정하고, 부족하다면 100-500으로 낮추세요. OOM 에러가 나면 절반으로 줄이는 것이 경험적으로 효과적입니다.
💡 토크나이저의 fast=True 옵션을 사용하면 Rust 기반 구현으로 3-5배 빠른 토크나이징이 가능합니다. Hugging Face의 최신 토크나이저는 기본적으로 Fast 버전을 사용하지만, 구버전에서는 명시적으로 지정해야 합니다.
💡 전처리 중간에 실패하면 처음부터 다시 시작해야 하므로, 체크포인트를 만들어두세요. 10만 개마다 중간 결과를 저장하면 장애 발생 시 최대 10만 개만 재처리하면 됩니다.
💡 SSD를 사용하면 HDD 대비 5-10배 빠른 I/O 성능을 얻을 수 있습니다. 전처리 데이터는 순차 읽기/쓰기가 많으므로 SSD의 이점이 극대화됩니다.
💡 num_proc 인자로 프로세스 수를 명시적으로 제어할 수 있습니다. CPU 코어 수보다 1-2개 적게 설정하면 시스템이 응답 불가 상태가 되는 것을 방지할 수 있습니다.
2. 분산 데이터 병렬 학습 설정 - PyTorch DDP 구현
시작하며
여러분이 대규모 트랜스포머 모델을 단일 GPU로 학습시키려 했는데, 한 에폭에 일주일이 걸린다는 사실을 알게 된 적 있나요? 회사에서 8개의 GPU를 제공받았지만 어떻게 활용해야 할지 막막한 상황 말이죠.
이런 문제는 실제 개발 현장에서 자주 발생합니다. 단일 GPU로는 학습 시간이 너무 오래 걸려서 실험 반복이 불가능하고, 결국 프로젝트 일정이 지연됩니다.
특히 ChatGPT 규모의 모델은 단일 GPU로는 학습 자체가 불가능하며, 수백 개의 GPU를 효율적으로 활용하는 기술이 필수적입니다. 바로 이럴 때 필요한 것이 PyTorch의 Distributed Data Parallel(DDP)입니다.
여러 GPU에 모델 복사본을 만들고, 각 GPU가 다른 데이터 배치를 처리한 후 그래디언트를 동기화하는 방식으로 학습 속도를 GPU 수에 비례하여 증가시킬 수 있습니다.
개요
간단히 말해서, Distributed Data Parallel은 동일한 모델을 여러 GPU에 복제하고, 각 GPU가 서로 다른 데이터 배치로 forward/backward를 수행한 후 그래디언트를 평균내어 동기화하는 병렬 학습 기법입니다. 왜 이 기법이 필요한지 실무 관점에서 설명하자면, 8개 GPU로 DDP를 사용하면 단일 GPU 대비 약 7-7.5배 빠른 학습이 가능합니다(완벽한 8배가 아닌 이유는 통신 오버헤드 때문입니다).
예를 들어, GPT-2 크기의 모델(1.5B 파라미터)을 Common Crawl 데이터셋으로 학습할 때, 단일 A100 GPU로는 2주가 걸리지만 8개 GPU DDP로는 약 2.5일로 단축됩니다. 전통적인 방법과의 비교를 해보면, 기존에는 DataParallel(DP)을 사용하여 하나의 프로세스에서 여러 GPU를 관리했다면, DDP는 각 GPU마다 별도의 프로세스를 생성하여 Python GIL의 영향을 받지 않고 진정한 병렬 처리를 수행합니다.
이 기법의 핵심 특징은 첫째, 선형 확장성(GPU가 2배 늘면 학습 시간이 절반으로 단축), 둘째, 효율적인 통신(Ring-AllReduce 알고리즘으로 통신 오버헤드 최소화), 셋째, 코드 최소 변경(기존 학습 코드에 몇 줄만 추가)입니다. 이러한 특징들이 중요한 이유는 실무에서는 최소한의 코드 변경으로 최대한의 성능 향상을 얻어야 하기 때문입니다.
코드 예제
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler
# 분산 학습 환경 초기화
def setup_distributed(rank, world_size):
# NCCL 백엔드로 GPU 간 통신 설정
dist.init_process_group(
backend='nccl', # GPU 통신에 최적화된 백엔드
init_method='env://', # 환경변수에서 설정 읽기
world_size=world_size, # 전체 프로세스(GPU) 수
rank=rank # 현재 프로세스 ID (0부터 시작)
)
torch.cuda.set_device(rank) # 각 프로세스를 특정 GPU에 할당
# 모델과 데이터 로더를 DDP로 래핑
def prepare_ddp(model, train_dataset, batch_size, rank, world_size):
model = model.to(rank) # 모델을 해당 GPU로 이동
ddp_model = DDP(model, device_ids=[rank]) # DDP로 래핑
# 각 GPU가 다른 데이터를 받도록 샘플러 설정
sampler = DistributedSampler(
train_dataset,
num_replicas=world_size,
rank=rank,
shuffle=True # 에폭마다 셔플
)
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=batch_size,
sampler=sampler,
num_workers=4, # 데이터 로딩 병렬화
pin_memory=True # GPU 전송 속도 향상
)
return ddp_model, train_loader
설명
이것이 하는 일: DDP는 각 GPU에서 독립적인 Python 프로세스를 실행하고, 모든 프로세스가 동일한 모델의 복사본을 가지고 서로 다른 데이터 배치를 처리합니다. Backward 단계에서 각 GPU의 그래디언트를 All-Reduce 연산으로 평균내어 모든 GPU의 모델 파라미터를 동일하게 업데이트합니다.
첫 번째로, dist.init_process_group에서 NCCL 백엔드를 사용하여 프로세스 간 통신을 초기화합니다. NCCL(NVIDIA Collective Communications Library)은 GPU 간 통신에 최적화된 라이브러리로, GPUDirect RDMA 같은 하드웨어 가속을 활용하여 통신 오버헤드를 최소화합니다.
왜 이렇게 하는지 이해하려면, CPU를 거치지 않고 GPU 간 직접 통신이 가능하다는 점이 핵심입니다. 일반 TCP 통신 대비 10-20배 빠른 속도를 제공합니다.
그 다음으로, DDP(model, device_ids=[rank])로 모델을 래핑하면 forward 단계는 각 GPU에서 독립적으로 실행되고, backward 단계에서 자동으로 그래디언트를 동기화합니다. 내부적으로는 각 레이어의 backward가 완료되는 즉시 해당 레이어의 그래디언트를 백그라운드에서 통신하므로(gradient bucketing), 통신과 계산이 오버랩되어 실제 통신 시간이 거의 제로에 가까워집니다.
DistributedSampler는 전체 데이터셋을 GPU 수로 나누어 각 GPU가 겹치지 않는 서브셋을 받도록 보장합니다. 예를 들어 10,000개 샘플을 4개 GPU로 학습한다면, GPU 0은 0-2499번, GPU 1은 2500-4999번 샘플을 처리합니다.
마지막으로, sampler.set_epoch(epoch)를 매 에폭마다 호출하면(코드에는 없지만 실제로는 필수) 셔플 시드가 변경되어 매번 다른 순서로 학습합니다. 실제 학습 루프에서는 loss.backward()만 호출하면 DDP가 자동으로 그래디언트를 All-Reduce하고, optimizer.step()이 모든 GPU에서 동일한 파라미터 업데이트를 수행합니다.
각 GPU의 배치 크기가 32라면, 8개 GPU에서의 effective batch size는 256이 되어 더 안정적인 학습이 가능합니다. 여러분이 이 코드를 사용하면 단일 GPU로는 불가능했던 대규모 모델 학습을 며칠 안에 완료할 수 있고, 더 큰 배치 크기로 학습하여 수렴 안정성도 향상됩니다.
실무에서는 클라우드 환경에서 GPU를 시간 단위로 빌리는 경우가 많으므로, 학습 시간을 1/8로 단축하면 비용도 동일하게 절감됩니다.
실전 팁
💡 find_unused_parameters=True 옵션은 일부 파라미터가 특정 배치에서 사용되지 않을 때 필요하지만, 성능이 10-15% 저하되므로 정말 필요한 경우에만 사용하세요. 대부분의 표준 트랜스포머 모델에서는 불필요합니다.
💡 멀티노드 학습 시에는 init_method를 tcp://master_node_ip:port 형식으로 변경하고, 모든 노드가 마스터 노드에 접근할 수 있는지 확인하세요. 방화벽 설정을 잘못하면 "connection timeout" 에러가 발생합니다.
💡 그래디언트 누적(gradient accumulation)과 DDP를 함께 사용하면 메모리는 절약하면서도 큰 effective batch size를 얻을 수 있습니다. 예를 들어 4 step accumulation + 8 GPU + batch 32 = effective batch 1024입니다.
💡 torch.cuda.amp의 자동 혼합 정밀도(AMP)를 DDP와 함께 사용하면 메모리 사용량을 절반으로 줄이고 속도는 2-3배 향상됩니다. GradScaler를 각 프로세스에서 독립적으로 사용하면 됩니다.
💡 체크포인트 저장은 rank 0 프로세스에서만 수행하세요. 모든 GPU에서 동시에 저장하면 디스크 I/O가 병목이 되고 같은 파일을 덮어쓰는 race condition이 발생할 수 있습니다.
3. 효율적인 체크포인트 관리 시스템 - 증분 저장과 복구
시작하며
여러분이 3일간 학습한 모델이 갑자기 서버 장애로 날아가버린 경험이 있나요? 아니면 체크포인트 파일이 너무 커서 디스크 공간이 부족해진 상황 말이죠.
이런 문제는 실제 개발 현장에서 자주 발생합니다. 대규모 모델은 한 번의 체크포인트가 수십 GB에 달하고, 매 에폭마다 저장하면 TB 단위의 저장 공간이 필요합니다.
게다가 학습 중 예기치 않은 장애가 발생하면 며칠간의 학습이 물거품이 될 수 있습니다. GPT-3 규모의 모델은 체크포인트 하나가 350GB 이상이므로, 저장과 로딩에만 수십 분이 소요됩니다.
바로 이럴 때 필요한 것이 증분 체크포인트 시스템입니다. 변경된 부분만 저장하고, 자동으로 오래된 체크포인트를 정리하며, 장애 발생 시 빠르게 복구할 수 있는 체계적인 관리 시스템이 필수적입니다.
개요
간단히 말해서, 효율적인 체크포인트 관리 시스템은 모델의 상태를 주기적으로 저장하되, 저장 공간을 최소화하고 복구 시간을 단축하며, 저장/로딩 중에도 학습을 계속할 수 있도록 설계된 시스템입니다. 왜 이 시스템이 필요한지 실무 관점에서 설명하자면, 장시간 학습에서는 하드웨어 장애, 네트워크 문제, OOM 에러 등 다양한 이슈가 발생할 수 있습니다.
체크포인트가 없으면 처음부터 다시 학습해야 하지만, 적절한 체크포인트 전략이 있으면 최대 몇 시간의 학습만 재수행하면 됩니다. 예를 들어, 일주일짜리 학습에서 6일째에 장애가 발생했을 때, 1시간마다 체크포인트를 저장했다면 최대 1시간만 손실됩니다.
전통적인 방법과의 비교를 해보면, 기존에는 매 에폭마다 전체 모델을 저장하고 N개까지만 보관했다면, 이제는 중요도에 따라 선별적으로 저장하고(best model, latest, milestone), 비동기 저장으로 학습을 멈추지 않으며, 클라우드 스토리지로 자동 백업할 수 있습니다. 이 시스템의 핵심 특징은 첫째, 저장 공간 효율성(중복 제거와 압축), 둘째, 최소 다운타임(비동기 저장), 셋째, 유연한 복구(특정 시점으로 롤백 가능)입니다.
이러한 특징들이 중요한 이유는 실무에서는 수십 번의 실험을 반복하므로 각 실험의 체크포인트를 효율적으로 관리하지 않으면 저장 공간과 시간이 기하급수적으로 낭비되기 때문입니다.
코드 예제
import torch
import os
import shutil
from pathlib import Path
import asyncio
from concurrent.futures import ThreadPoolExecutor
class CheckpointManager:
def __init__(self, checkpoint_dir, max_to_keep=3, keep_best=True):
self.checkpoint_dir = Path(checkpoint_dir)
self.checkpoint_dir.mkdir(parents=True, exist_ok=True)
self.max_to_keep = max_to_keep # 최근 N개만 유지
self.keep_best = keep_best # 최고 성능 모델은 별도 보관
self.best_metric = float('inf')
self.executor = ThreadPoolExecutor(max_workers=1) # 비동기 저장용
def save_checkpoint(self, model, optimizer, epoch, global_step, metrics):
# 저장할 상태 딕셔너리 구성
checkpoint = {
'epoch': epoch,
'global_step': global_step,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'metrics': metrics,
'rng_state': torch.get_rng_state(), # 재현성을 위한 난수 상태
}
# 파일명에 스텝 번호 포함
checkpoint_path = self.checkpoint_dir / f"checkpoint_step_{global_step}.pt"
# 비동기로 저장 (학습 루프 블로킹 방지)
future = self.executor.submit(torch.save, checkpoint, checkpoint_path)
# 오래된 체크포인트 정리
self._cleanup_old_checkpoints()
# 최고 성능이면 별도 저장
if self.keep_best and metrics['loss'] < self.best_metric:
self.best_metric = metrics['loss']
best_path = self.checkpoint_dir / "best_model.pt"
shutil.copy(checkpoint_path, best_path)
return checkpoint_path
def _cleanup_old_checkpoints(self):
# 최근 체크포인트들만 유지
checkpoints = sorted(
self.checkpoint_dir.glob("checkpoint_step_*.pt"),
key=lambda p: p.stat().st_mtime # 수정 시간 기준 정렬
)
# 오래된 것부터 삭제
while len(checkpoints) > self.max_to_keep:
oldest = checkpoints.pop(0)
oldest.unlink() # 파일 삭제
def load_latest_checkpoint(self, model, optimizer):
# 가장 최근 체크포인트 찾기
checkpoints = sorted(
self.checkpoint_dir.glob("checkpoint_step_*.pt"),
key=lambda p: int(p.stem.split('_')[-1]) # 스텝 번호로 정렬
)
if not checkpoints:
return None # 체크포인트 없음
latest = checkpoints[-1]
checkpoint = torch.load(latest, map_location='cuda')
# 상태 복원
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
torch.set_rng_state(checkpoint['rng_state'])
return checkpoint['epoch'], checkpoint['global_step']
설명
이것이 하는 일: CheckpointManager는 학습 중 주기적으로 모델, 옵티마이저, 스케줄러 등의 상태를 디스크에 저장하고, 설정된 개수만큼만 유지하며, 최고 성능 모델은 별도로 보관합니다. 장애 발생 시 가장 최근 체크포인트부터 자동으로 복구하여 학습을 재개할 수 있습니다.
첫 번째로, save_checkpoint에서 모델과 옵티마이저의 state_dict()뿐만 아니라 난수 생성기 상태(torch.get_rng_state())까지 저장합니다. 왜 이렇게 하는지 이해하려면, 동일한 체크포인트에서 재시작했을 때 dropout, data shuffling 등이 완전히 동일하게 재현되어야 재현 가능한 실험이 된다는 점이 중요합니다.
논문 작성이나 디버깅 시 필수적입니다. 그 다음으로, ThreadPoolExecutor를 사용한 비동기 저장이 핵심입니다.
torch.save는 대규모 모델에서 수십 초가 걸릴 수 있는데, 이를 메인 스레드에서 실행하면 그동안 학습이 멈춥니다. 별도 스레드로 저장하면 GPU는 계속 다음 배치를 처리하고, 저장은 백그라운드에서 진행되어 실질적인 오버헤드가 거의 없습니다.
다만 worker를 1개로 제한한 이유는 여러 체크포인트를 동시에 저장하면 디스크 I/O가 병목이 되기 때문입니다. _cleanup_old_checkpoints는 수정 시간 기준으로 오래된 체크포인트를 삭제하여 디스크 공간을 관리합니다.
예를 들어 max_to_keep=3이면 항상 최근 3개만 유지되므로, 모델이 20GB라면 최대 60GB만 사용합니다. 최고 성능 모델(best_model.pt)은 별도로 보관되므로 삭제 대상에서 제외됩니다.
load_latest_checkpoint는 파일명의 스텝 번호를 파싱하여 가장 큰 숫자를 가진 체크포인트를 로드합니다. map_location='cuda'를 지정하여 CPU에 로드한 후 GPU로 이동하는 불필요한 단계를 생략하고 직접 GPU 메모리로 로드합니다.
반환된 epoch와 global_step을 사용하여 학습 루프를 정확히 이어갈 수 있습니다. 여러분이 이 코드를 사용하면 장시간 학습에서도 안심하고 실험할 수 있고, 저장 공간 관리에 신경 쓰지 않아도 되며, 하이퍼파라미터를 바꿔가며 수십 번 재시작해도 체계적으로 관리됩니다.
실무에서는 클라우드 스토리지(S3, GCS)로 주기적으로 백업하는 기능을 추가하면 로컬 서버 장애에도 대응할 수 있습니다.
실전 팁
💡 torch.save에 _use_new_zipfile_serialization=True 옵션을 주면 저장 속도가 20-30% 빨라지고 파일 크기도 약간 줄어듭니다. PyTorch 1.6 이후 버전에서 기본값이지만 명시적으로 지정하는 것이 안전합니다.
💡 DDP 환경에서는 rank 0에서만 체크포인트를 저장하세요. if dist.get_rank() == 0: 조건으로 감싸지 않으면 모든 GPU가 동시에 같은 파일에 쓰려고 해서 데이터가 손상됩니다.
💡 validation loss가 개선될 때만 저장하는 전략도 유용합니다. 이렇게 하면 저장 횟수가 크게 줄어들고(전체 학습의 5-10% 정도), 실제로 의미 있는 체크포인트만 남습니다.
💡 큰 모델은 torch.save 대신 safetensors 라이브러리를 사용하면 로딩 속도가 5-10배 빨라집니다. Hugging Face가 개발한 이 포맷은 메모리 매핑을 지원하여 전체를 RAM에 올리지 않고도 부분적으로 로드할 수 있습니다.
💡 체크포인트에 현재 학습률, gradient norm, 데이터셋 위치 등 디버깅에 유용한 메타데이터를 함께 저장하세요. 나중에 "왜 이 시점에서 loss가 급증했지?" 같은 질문에 답할 수 있습니다.
4. 학습률 스케줄링과 워밍업 - Cosine Annealing with Linear Warmup
시작하며
여러분이 대규모 트랜스포머 모델을 학습시키는데 초반에 loss가 폭발하거나 NaN이 나타난 경험이 있나요? 또는 학습이 진행될수록 loss가 정체되어 더 이상 개선되지 않는 상황 말이죠.
이런 문제는 실제 개발 현장에서 자주 발생합니다. 고정된 학습률로는 초반에 그래디언트가 너무 커서 발산하거나, 후반에 fine-tuning이 부족해서 최적점에 도달하지 못합니다.
GPT-3 논문에서도 학습률 스케줄링이 잘못되면 전체 학습이 실패할 수 있다고 명시되어 있습니다. 특히 대규모 배치(수천~수만)로 학습할 때는 초기 불안정성이 더욱 심각합니다.
바로 이럴 때 필요한 것이 Cosine Annealing with Linear Warmup입니다. 초반에는 학습률을 천천히 증가시켜 안정성을 확보하고, 중후반에는 코사인 함수를 따라 부드럽게 감소시켜 최적점에 수렴하는 전략입니다.
개요
간단히 말해서, Cosine Annealing with Linear Warmup은 학습 초기에 학습률을 0에서 목표값까지 선형 증가시킨 후(warmup), 코사인 곡선을 따라 0에 가까워지도록 감소시키는 학습률 스케줄링 기법입니다. 왜 이 기법이 필요한지 실무 관점에서 설명하자면, 학습 초기에는 모델 파라미터가 랜덤 초기화 상태이므로 그래디언트의 방향과 크기가 불안정합니다.
이 상태에서 큰 학습률로 업데이트하면 loss가 폭발할 수 있습니다. 반면 warmup을 사용하면 초기 수백~수천 step 동안 작은 학습률로 시작하여 파라미터를 안정화시킨 후 본격적인 학습에 들어갑니다.
예를 들어, BERT는 10,000 step warmup을 사용하고, GPT-3는 전체 학습의 약 3%를 warmup에 할당합니다. 전통적인 방법과의 비교를 해보면, 기존의 StepLR이나 ExponentialLR은 고정된 간격으로 학습률을 감소시켜 불연속적인 변화가 발생하지만, Cosine Annealing은 부드러운 곡선으로 변화하여 학습 안정성이 높습니다.
또한 warmup 없이 시작하면 초기 발산 위험이 크지만, warmup을 추가하면 어떤 초기화 방법을 사용하든 안정적으로 시작할 수 있습니다. 이 기법의 핵심 특징은 첫째, 초기 안정성(warmup으로 발산 방지), 둘째, 부드러운 수렴(cosine decay로 미세 조정), 셋째, 단순한 하이퍼파라미터(warmup step과 총 step만 설정)입니다.
이러한 특징들이 중요한 이유는 실무에서는 수십 번의 실험을 통해 최적의 학습률을 찾아야 하는데, 복잡한 스케줄링은 하이퍼파라미터 공간을 너무 넓혀서 탐색이 어렵기 때문입니다.
코드 예제
import math
from torch.optim.lr_scheduler import LambdaLR
def get_cosine_schedule_with_warmup(
optimizer,
num_warmup_steps,
num_training_steps,
num_cycles=0.5, # 코사인 사이클 수 (0.5 = 절반 주기)
last_epoch=-1
):
"""
Cosine Annealing with Linear Warmup 스케줄러 생성
"""
def lr_lambda(current_step):
# Warmup 단계: 선형 증가
if current_step < num_warmup_steps:
return float(current_step) / float(max(1, num_warmup_steps))
# Warmup 이후: Cosine Annealing
progress = float(current_step - num_warmup_steps) / float(
max(1, num_training_steps - num_warmup_steps)
)
# 코사인 함수로 0에서 1로 변화 후, 1에서 0으로 매핑
return max(0.0, 0.5 * (1.0 + math.cos(math.pi * progress * num_cycles * 2.0)))
return LambdaLR(optimizer, lr_lambda, last_epoch)
# 사용 예시
from torch.optim import AdamW
model = ... # 여러분의 모델
optimizer = AdamW(model.parameters(), lr=5e-4) # 최대 학습률
# 전체 학습 스텝 계산
num_epochs = 10
steps_per_epoch = len(train_loader)
total_steps = num_epochs * steps_per_epoch
warmup_steps = int(0.1 * total_steps) # 전체의 10%를 warmup으로
scheduler = get_cosine_schedule_with_warmup(
optimizer,
num_warmup_steps=warmup_steps,
num_training_steps=total_steps
)
# 학습 루프에서 사용
for epoch in range(num_epochs):
for batch in train_loader:
optimizer.zero_grad()
loss = model(batch)
loss.backward()
optimizer.step()
scheduler.step() # 배치마다 학습률 업데이트
설명
이것이 하는 일: 이 스케줄러는 학습 과정을 두 단계로 나눕니다. Warmup 단계에서는 학습률을 0부터 설정값까지 선형으로 증가시키고, 이후 Cosine Annealing 단계에서는 코사인 함수를 따라 학습률을 부드럽게 감소시켜 0에 가까워지도록 합니다.
첫 번째로, lr_lambda 함수 내부에서 current_step < num_warmup_steps 조건으로 warmup 단계를 처리합니다. 이 구간에서는 current_step / num_warmup_steps로 계산되므로, 예를 들어 warmup이 1000 step이고 최대 학습률이 5e-4라면, 500번째 step에서는 2.5e-4가 됩니다.
왜 이렇게 하는지 이해하려면, 학습 초기에 Adam 옵티마이저의 분모(second moment)가 아직 충분히 누적되지 않아 업데이트 크기가 불안정하다는 점이 핵심입니다. 작은 학습률로 시작하면 이 불안정성을 완화할 수 있습니다.
그 다음으로, warmup 이후에는 progress 변수로 전체 학습 과정에서의 진행도(0.0~1.0)를 계산합니다. 이 값을 math.cos(math.pi * progress * 2 * num_cycles)에 넣으면 0에서 시작해서 π까지 증가하는 각도가 됩니다.
num_cycles=0.5는 절반 주기만 사용한다는 의미로, 코사인 값이 1(최댓값)에서 -1(최솟값)까지 변하게 됩니다. 최종적으로 0.5 * (1.0 + cos_value)로 변환하여 1에서 0으로 매핑됩니다.
LambdaLR은 매 step마다 lr_lambda 함수를 호출하여 현재 학습률을 계산합니다. 내부적으로 base_lr * lr_lambda(step)으로 실제 학습률이 결정되므로, 여러분이 optimizer를 생성할 때 설정한 lr 값이 warmup 후 도달할 최대 학습률이 됩니다.
예를 들어 lr=5e-4로 설정했다면, warmup 끝에서 5e-4가 되고, 학습 종료 시점에서는 거의 0에 가까워집니다. 학습 루프에서는 scheduler.step()을 배치마다 호출해야 합니다.
많은 초보자가 epoch마다 호출하는 실수를 하는데, 그러면 학습률이 너무 빠르게 감소하여 조기에 수렴이 멈춥니다. num_training_steps는 전체 배치 수(epochs * batches_per_epoch)로 설정해야 정확히 작동합니다.
여러분이 이 코드를 사용하면 어떤 모델 크기나 배치 크기에서도 안정적으로 학습을 시작할 수 있고, 수동으로 학습률을 조정하는 번거로움 없이 자동으로 최적 수렴 경로를 따라갑니다. 실무에서는 GPT, BERT, T5 등 거의 모든 트랜스포머 모델이 이 스케줄링을 기본으로 사용합니다.
실전 팁
💡 Warmup step은 일반적으로 전체 학습의 5-10%가 적절합니다. 너무 짧으면(1% 미만) 초기 불안정성이 남고, 너무 길면(20% 이상) 실제 학습 시간이 부족해집니다. 대규모 모델일수록 더 긴 warmup이 필요합니다.
💡 체크포인트에서 재시작할 때는 scheduler.load_state_dict()를 호출하여 스케줄러 상태도 복원해야 합니다. 그렇지 않으면 학습률이 초기값으로 리셋되어 학습이 불안정해집니다.
💡 Learning rate finder를 먼저 실행하여 최대 학습률을 찾은 후 이 스케줄러를 적용하면 더 좋은 결과를 얻습니다. FastAI의 lr_find() 같은 도구를 사용하면 자동으로 최적 학습률을 찾아줍니다.
💡 num_cycles 값을 조정하면 다양한 변형이 가능합니다. 0.5(기본)는 한 번 감소, 1.0은 완전한 사이클(감소 후 다시 증가), 2.0은 두 번의 사이클을 만듭니다. 대부분의 경우 0.5가 최적입니다.
💡 Mixed Precision Training(AMP)과 함께 사용할 때는 GradScaler의 dynamic loss scaling이 초기 warmup 단계에서 불안정할 수 있으므로, warmup을 조금 더 길게(15-20%) 설정하는 것이 안전합니다.
5. 그래디언트 클리핑과 누적 - 메모리 효율적인 대규모 배치 학습
시작하며
여러분이 큰 배치 크기로 학습하고 싶은데 GPU 메모리가 부족해서 OOM(Out of Memory) 에러가 발생한 경험이 있나요? 또는 학습 중 갑자기 loss가 NaN이 되어 전체 학습이 망가진 상황 말이죠.
이런 문제는 실제 개발 현장에서 자주 발생합니다. 대규모 트랜스포머 모델은 배치 크기 32만 사용해도 80GB GPU 메모리를 초과하는 경우가 많습니다.
또한 특정 배치에서 그래디언트가 비정상적으로 커지면(gradient explosion) 모델 파라미터가 망가져서 복구 불가능한 상태가 됩니다. GPT-2 학습 시 gradient clipping 없이는 학습의 약 30%가 발산으로 실패한다는 연구 결과도 있습니다.
바로 이럴 때 필요한 것이 그래디언트 누적(Gradient Accumulation)과 그래디언트 클리핑(Gradient Clipping)입니다. 누적은 작은 배치를 여러 번 처리하여 큰 effective batch size를 만들고, 클리핑은 그래디언트의 최대 크기를 제한하여 학습 안정성을 확보합니다.
개요
간단히 말해서, 그래디언트 누적은 여러 배치의 그래디언트를 더한 후 한 번에 업데이트하여 메모리 사용량은 줄이면서 큰 배치 효과를 얻는 기법이고, 그래디언트 클리핑은 그래디언트 벡터의 norm이 임계값을 초과하면 스케일을 줄여서 폭발적인 업데이트를 방지하는 기법입니다. 왜 이 기법들이 필요한지 실무 관점에서 설명하자면, 대규모 언어 모델 학습에서는 안정성을 위해 큰 배치 크기(256~2048 이상)가 필요하지만, 물리적 메모리 제약으로 불가능한 경우가 많습니다.
그래디언트 누적을 사용하면 배치 크기 32를 8번 누적하여 effective batch 256을 만들 수 있으며, 메모리는 배치 32 수준만 사용합니다. 예를 들어, GPT-3 학습 시 각 GPU에서 배치 크기 1-2만 가능했지만, 누적과 분산 학습으로 effective batch 3.2M을 달성했습니다.
전통적인 방법과의 비교를 해보면, 기존에는 메모리가 부족하면 모델 크기를 줄이거나 더 비싼 GPU를 구매해야 했지만, 이제는 동일한 하드웨어에서 누적 step만 조정하여 원하는 배치 크기를 달성할 수 있습니다. 또한 gradient clipping 없이는 학습률을 매우 낮게 설정해야 안정적이었지만, clipping을 사용하면 더 공격적인 학습률로 빠르게 수렴할 수 있습니다.
이 기법들의 핵심 특징은 첫째, 메모리 효율성(큰 배치를 작은 메모리로 구현), 둘째, 학습 안정성(gradient explosion 방지), 셋째, 유연성(하드웨어에 맞게 조정 가능)입니다. 이러한 특징들이 중요한 이유는 실무에서는 항상 메모리와 시간의 제약 속에서 최대한의 성능을 뽑아내야 하기 때문입니다.
코드 예제
import torch
from torch.nn.utils import clip_grad_norm_
def train_with_gradient_accumulation(
model,
train_loader,
optimizer,
accumulation_steps=4,
max_grad_norm=1.0, # Gradient clipping 임계값
device='cuda'
):
model.train()
optimizer.zero_grad() # 초기 그래디언트 초기화
total_loss = 0
for batch_idx, batch in enumerate(train_loader):
# 입력을 GPU로 이동
input_ids = batch['input_ids'].to(device)
labels = batch['labels'].to(device)
# Forward pass
outputs = model(input_ids, labels=labels)
loss = outputs.loss
# Gradient accumulation을 위해 loss를 누적 스텝 수로 나눔
# 이렇게 하면 누적된 그래디언트가 평균값이 됨
loss = loss / accumulation_steps
loss.backward() # Backward pass (그래디언트 누적)
total_loss += loss.item()
# 설정된 스텝마다 실제 업데이트 수행
if (batch_idx + 1) % accumulation_steps == 0:
# Gradient clipping: 그래디언트 norm이 max_grad_norm 초과하면 스케일 조정
clip_grad_norm_(model.parameters(), max_grad_norm)
# 옵티마이저 업데이트
optimizer.step()
optimizer.zero_grad() # 그래디언트 초기화
# 로깅 (실제 업데이트 시점에만)
if (batch_idx + 1) % (accumulation_steps * 10) == 0:
avg_loss = total_loss / 10
print(f"Batch {batch_idx + 1}, Loss: {avg_loss:.4f}")
total_loss = 0
# 마지막 배치가 accumulation_steps의 배수가 아닐 경우 처리
if (batch_idx + 1) % accumulation_steps != 0:
clip_grad_norm_(model.parameters(), max_grad_norm)
optimizer.step()
optimizer.zero_grad()
# 사용 예시
model = ... # GPT-2 같은 대규모 모델
optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5)
# 물리적 배치 크기 8, 누적 4번 = effective batch 32
train_with_gradient_accumulation(
model,
train_loader, # batch_size=8로 설정된 데이터 로더
optimizer,
accumulation_steps=4, # 4번 누적 후 업데이트
max_grad_norm=1.0 # BERT/GPT 논문의 표준 값
)
설명
이것이 하는 일: 이 함수는 작은 배치로 여러 번 forward/backward를 수행하여 그래디언트를 누적한 후, 설정된 스텝마다 그래디언트를 클리핑하고 옵티마이저를 업데이트합니다. 메모리는 작은 배치 크기만 사용하면서도 큰 배치로 학습한 것과 동일한 효과를 얻습니다.
첫 번째로, loss = loss / accumulation_steps로 loss를 나누는 것이 핵심입니다. PyTorch는 loss.backward()를 호출할 때마다 그래디언트를 누적하므로(덮어쓰지 않고 더함), 4번 backward를 하면 그래디언트가 4배가 됩니다.
왜 이렇게 하는지 이해하려면, 실제로 원하는 것은 4개 배치의 그래디언트 평균이므로, 각 배치의 loss를 4로 나눠서 backward하면 최종 누적 그래디언트가 평균값이 됩니다. 이렇게 하지 않으면 effective batch가 커질수록 업데이트 크기도 비례하여 커져서 학습이 불안정해집니다.
그 다음으로, clip_grad_norm_(model.parameters(), max_grad_norm)이 gradient explosion을 방지합니다. 내부적으로는 모든 파라미터의 그래디언트 벡터를 하나로 이어붙인 후 L2 norm을 계산하고, 이 값이 max_grad_norm을 초과하면 모든 그래디언트를 동일한 비율로 스케일링합니다.
예를 들어 total norm이 5.0이고 max가 1.0이면, 모든 그래디언트에 0.2를 곱하여 norm을 1.0으로 만듭니다. 방향은 유지되지만 크기만 제한되는 것이죠.
(batch_idx + 1) % accumulation_steps == 0 조건으로 실제 업데이트 시점을 제어합니다. 누적 중간에는 optimizer.step()을 호출하지 않으므로 모델 파라미터가 변하지 않고 그래디언트만 누적됩니다.
업데이트 후에는 반드시 optimizer.zero_grad()로 그래디언트를 초기화해야 다음 누적 사이클이 올바르게 시작됩니다. 초기화를 빼먹으면 이전 누적값이 계속 남아서 잘못된 업데이트가 발생합니다.
마지막 부분의 조건문은 edge case를 처리합니다. 예를 들어 전체 데이터가 100개이고 배치가 8, 누적이 4라면, 마지막에 4개 배치가 남는 게 아니라 2개만 남을 수 있습니다.
이 경우에도 누적된 그래디언트를 버리지 않고 업데이트를 수행하여 모든 데이터를 활용합니다. 여러분이 이 코드를 사용하면 16GB GPU로도 효과적으로 수백~수천의 배치 크기를 시뮬레이션할 수 있고, gradient explosion으로 인한 학습 실패를 거의 완전히 방지할 수 있습니다.
실무에서는 누적 스텝을 조정하여 메모리 사용량과 학습 속도를 균형있게 조절하는 것이 중요합니다.
실전 팁
💡 max_grad_norm은 1.0이 일반적이지만, 모델과 데이터에 따라 0.5~5.0 범위에서 조정할 수 있습니다. 학습 로그에서 clipping이 너무 자주 발생하면(50% 이상) 값을 높이고, 거의 발생하지 않으면 낮춰서 더 공격적으로 학습할 수 있습니다.
💡 DDP와 그래디언트 누적을 함께 사용할 때는 각 GPU에서 독립적으로 누적하고, 마지막 step에서만 All-Reduce가 발생합니다. 예를 들어 8 GPU, 배치 4, 누적 8 = effective batch 256입니다.
💡 학습률은 effective batch size에 비례하여 조정해야 합니다. Linear scaling rule에 따르면, 배치를 2배로 늘리면 학습률도 2배로 증가시켜야 동일한 학습 곡선을 얻습니다(단, warmup은 필수).
💡 Mixed Precision Training(AMP)과 함께 사용할 때는 scaler.scale(loss).backward()로 스케일링된 loss를 backward하고, scaler.step(optimizer) 전에 scaler.unscale_(optimizer)을 호출한 후 clipping을 해야 정확합니다.
💡 Gradient accumulation 중에는 BatchNorm이 잘못 동작할 수 있습니다(각 micro-batch의 통계를 사용하므로). 대규모 모델에서는 BatchNorm 대신 LayerNorm을 사용하는 것이 표준입니다.
6. 효율적인 어텐션 메커니즘 - Flash Attention과 메모리 최적화
시작하며
여러분이 긴 문맥(4096 토큰 이상)으로 트랜스포머를 학습시키려 했는데, 메모리가 폭발적으로 증가해서 OOM이 발생한 경험이 있나요? 또는 어텐션 계산에만 전체 학습 시간의 70% 이상이 소요되는 상황 말이죠.
이런 문제는 실제 개발 현장에서 자주 발생합니다. 표준 어텐션은 시퀀스 길이의 제곱에 비례하여 메모리와 시간을 사용하므로, 시퀀스가 2배 길어지면 메모리는 4배, 시간도 4배가 됩니다.
GPT-3가 2048 토큰으로 제한된 이유 중 하나가 바로 어텐션의 메모리 문제입니다. 8192 토큰 시퀀스는 2048 토큰 대비 16배의 어텐션 메모리가 필요합니다.
바로 이럴 때 필요한 것이 Flash Attention입니다. 수학적으로는 동일한 어텐션 결과를 만들지만, GPU 메모리 계층 구조를 고려한 알고리즘으로 메모리 사용량을 크게 줄이고 속도는 2-4배 향상시킵니다.
개요
간단히 말해서, Flash Attention은 표준 어텐션과 동일한 결과를 계산하지만, 중간 어텐션 행렬을 전체 HBM(High Bandwidth Memory)에 저장하지 않고 작은 블록 단위로 SRAM에서 처리하여 메모리 I/O를 최소화하는 알고리즘입니다. 왜 이 기법이 필요한지 실무 관점에서 설명하자면, 표준 어텐션은 (batch, heads, seq_len, seq_len) 크기의 어텐션 행렬을 만들어야 하므로, 배치 8, 헤드 16, 시퀀스 4096이면 약 16GB의 중간 메모리가 필요합니다.
Flash Attention은 이를 블록 단위로 처리하여 메모리를 시퀀스 길이에 선형 비례하도록 줄입니다. 예를 들어, A100 GPU에서 시퀀스 길이 4096을 표준 어텐션으로는 처리할 수 없지만, Flash Attention은 배치 32까지 가능합니다.
전통적인 방법과의 비교를 해보면, 기존 어텐션은 Q@K^T로 전체 어텐션 행렬을 만든 후 softmax를 적용하고 V와 곱했다면, Flash Attention은 타일링(tiling) 기법으로 작은 블록만 SRAM에 올려서 계산하고 즉시 버립니다. 메모리 사용량은 O(N^2)에서 O(N)으로 감소하지만, 결과는 비트 단위로 정확히 동일합니다.
이 기법의 핵심 특징은 첫째, 메모리 효율성(선형 메모리 복잡도), 둘째, 속도 향상(I/O 최소화로 2-4배 빠름), 셋째, 수학적 정확성(표준 어텐션과 동일한 출력)입니다. 이러한 특징들이 중요한 이유는 실무에서는 긴 문맥이 모델 성능에 직접적인 영향을 미치지만, 하드웨어 제약으로 불가능한 경우가 많기 때문입니다.
코드 예제
# Flash Attention 2 사용 예시 (xformers 또는 flash-attn 라이브러리 필요)
import torch
import torch.nn as nn
from flash_attn import flash_attn_qkvpacked_func, flash_attn_func
class FlashAttentionLayer(nn.Module):
def __init__(self, dim, num_heads, dropout=0.0):
super().__init__()
self.num_heads = num_heads
self.head_dim = dim // num_heads
self.scale = self.head_dim ** -0.5
# QKV 프로젝션을 한 번에 수행 (메모리 효율적)
self.qkv = nn.Linear(dim, dim * 3, bias=False)
self.proj = nn.Linear(dim, dim)
self.dropout = dropout
def forward(self, x, mask=None):
batch, seq_len, dim = x.shape
# QKV 계산: (batch, seq_len, 3, num_heads, head_dim)
qkv = self.qkv(x).reshape(
batch, seq_len, 3, self.num_heads, self.head_dim
)
# Flash Attention 호출 - 메모리 효율적인 어텐션
# 표준 어텐션과 달리 중간 행렬을 HBM에 저장하지 않음
out = flash_attn_qkvpacked_func(
qkv, # QKV가 패킹된 텐서
dropout_p=self.dropout if self.training else 0.0,
softmax_scale=self.scale, # 1/sqrt(d_k) 스케일링
causal=True # GPT 스타일의 causal masking
)
# 출력 형태 복원: (batch, seq_len, dim)
out = out.reshape(batch, seq_len, dim)
out = self.proj(out)
return out
# 사용 예시: Transformer 블록에 통합
class TransformerBlockWithFlash(nn.Module):
def __init__(self, dim, num_heads):
super().__init__()
self.attention = FlashAttentionLayer(dim, num_heads)
self.mlp = nn.Sequential(
nn.Linear(dim, dim * 4),
nn.GELU(),
nn.Linear(dim * 4, dim)
)
self.ln1 = nn.LayerNorm(dim)
self.ln2 = nn.LayerNorm(dim)
def forward(self, x):
# Pre-LN 구조 (GPT-2/3 스타일)
x = x + self.attention(self.ln1(x)) # 어텐션 + residual
x = x + self.mlp(self.ln2(x)) # MLP + residual
return x
# 실제 사용
model = TransformerBlockWithFlash(dim=768, num_heads=12)
x = torch.randn(8, 4096, 768).cuda() # 긴 시퀀스도 처리 가능
output = model(x) # 표준 어텐션 대비 2-4배 빠르고 메모리는 1/4
설명
이것이 하는 일: Flash Attention은 표준 어텐션의 세 단계(Q@K^T, softmax, @V)를 타일 단위로 분할하여 각 타일을 GPU의 빠른 SRAM(shared memory)에서 처리하고, 중간 결과를 HBM(메인 GPU 메모리)에 쓰지 않고 즉시 다음 계산에 사용합니다. 이를 통해 메모리 대역폭 병목을 극복합니다.
첫 번째로, flash_attn_qkvpacked_func는 QKV가 하나의 텐서로 패킹된 형태를 받아 내부적으로 블록 단위 처리를 수행합니다. 왜 이렇게 하는지 이해하려면, GPU의 메모리 계층 구조를 알아야 합니다.
HBM은 용량이 크지만(40-80GB) 느리고(1.5 TB/s), SRAM은 작지만(20MB) 매우 빠릅니다(19 TB/s). 표준 어텐션은 대부분의 연산에서 HBM을 읽고 쓰므로 메모리 대역폭에 병목됩니다.
Flash Attention은 128x128 같은 작은 타일을 SRAM에 올려서 모든 계산을 완료한 후 최종 결과만 HBM에 쓰므로, 메모리 I/O가 O(N^2)에서 O(N)으로 감소합니다. 그 다음으로, causal=True 옵션은 GPT 스타일의 autoregressive 마스킹을 적용합니다.
표준 어텐션에서는 별도의 마스크 텐서를 만들어 softmax 전에 추가해야 하지만, Flash Attention은 내부적으로 마스킹을 적용하여 마스크 텐서를 위한 추가 메모리가 불필요합니다. 시퀀스 길이 4096이면 마스크만 약 128MB이므로, 이를 절약하는 것도 중요합니다.
수학적으로는 표준 어텐션과 완전히 동일합니다. Softmax의 numerically stable한 계산을 위해 온라인 softmax 알고리즘을 사용하여, 블록 단위로 처리하면서도 최종 결과는 비트 단위로 정확히 일치합니다.
이는 "approximation"이 아니라 정확한 계산의 재배치라는 점이 중요합니다. Forward pass뿐만 아니라 backward pass도 동일한 타일링 기법으로 최적화됩니다.
표준 어텐션은 backward를 위해 전체 어텐션 행렬을 저장해야 하지만(메모리의 주요 소비처), Flash Attention은 recomputation 기법으로 필요할 때 다시 계산하여 저장 공간을 없앱니다. 재계산이 I/O보다 빠르기 때문에 전체적으로는 속도도 향상됩니다.
여러분이 이 코드를 사용하면 동일한 GPU로 2-4배 긴 시퀀스를 처리할 수 있고, 학습 시간도 단축되어 더 많은 실험을 할 수 있습니다. 실무에서는 긴 문맥이 필요한 문서 요약, 코드 생성, 대화 모델 등에서 필수적인 기술입니다.
실전 팁
💡 Flash Attention은 CUDA compute capability 7.5 이상(V100, A100, RTX 30/40 시리즈)에서만 지원됩니다. 이전 GPU에서는 설치되지 않거나 표준 어텐션으로 fallback되므로, 환경을 먼저 확인하세요.
💡 flash-attn 라이브러리 설치 시 PyTorch 버전과 CUDA 버전이 정확히 일치해야 합니다. 불일치하면 컴파일 에러가 발생하므로, pip install flash-attn --no-build-isolation 같은 옵션을 사용해야 할 수 있습니다.
💡 시퀀스 길이가 128의 배수가 아니면 약간의 오버헤드가 발생합니다. 가능하면 패딩을 사용하여 128, 256, 512 같은 값으로 맞추면 최적 성능을 얻을 수 있습니다.
💡 Grouped Query Attention(GQA)이나 Multi-Query Attention(MQA)도 Flash Attention과 호환됩니다. flash_attn_func에 KV 헤드 수를 지정하면 메모리를 더욱 절약할 수 있습니다.
💡 Inference 시에는 KV cache를 사용하므로 Flash Attention의 이점이 줄어듭니다. 주로 학습과 prefill 단계(첫 토큰 생성)에서 큰 효과를 발휘합니다.
7. 혼합 정밀도 학습 - FP16/BF16으로 학습 가속화
시작하며
여러분이 대규모 모델을 학습시키는데 GPU 메모리가 부족하고, 학습 속도도 너무 느려서 답답한 경험이 있나요? 전체 학습에 몇 주가 걸린다는 사실에 좌절한 적이 있으시죠?
이런 문제는 실제 개발 현장에서 자주 발생합니다. 기본적으로 PyTorch는 FP32(32비트 부동소수점)로 계산하므로, 1.5B 파라미터 모델은 약 6GB 메모리를 차지합니다.
옵티마이저 상태, 그래디언트, 활성화까지 포함하면 실제로는 24GB 이상이 필요합니다. 또한 FP32 연산은 최신 GPU의 Tensor Core를 완전히 활용하지 못해 성능의 절반만 사용합니다.
바로 이럴 때 필요한 것이 혼합 정밀도 학습(Mixed Precision Training)입니다. 대부분의 계산을 FP16이나 BF16(16비트)으로 수행하여 메모리를 절반으로 줄이고 속도는 2-3배 향상시키면서도, 수치 안정성은 FP32로 유지하는 기법입니다.
개요
간단히 말해서, 혼합 정밀도 학습은 forward/backward 계산과 가중치 저장은 FP16으로 수행하되, 옵티마이저 업데이트와 loss scaling은 FP32로 유지하여 메모리와 속도 이점을 얻으면서도 학습 안정성을 보장하는 기법입니다. 왜 이 기법이 필요한지 실무 관점에서 설명하자면, 최신 GPU(V100, A100, H100)는 FP16/BF16 연산에 특화된 Tensor Core를 탑재하여 FP32 대비 8-16배 빠른 행렬 곱셈이 가능합니다.
하지만 단순히 모든 것을 FP16으로 바꾸면 gradient underflow(너무 작은 값이 0으로 처리)나 accumulation error(작은 값들의 합이 부정확)가 발생합니다. 혼합 정밀도는 이 두 가지 문제를 loss scaling과 FP32 master weights로 해결합니다.
예를 들어, GPT-3 학습 시 혼합 정밀도 없이는 불가능했을 것으로 추정됩니다. 전통적인 방법과의 비교를 해보면, 기존 FP32 학습은 모든 텐서가 32비트이므로 안전하지만 느리고 메모리를 많이 사용했다면, 혼합 정밀도는 계산은 16비트로 하되 누적과 업데이트는 32비트로 유지하여 안정성과 효율성을 모두 확보합니다.
이 기법의 핵심 특징은 첫째, 메모리 절감(모델과 활성화가 절반 크기), 둘째, 속도 향상(Tensor Core 활용으로 2-3배), 셋째, 수치 안정성(loss scaling과 FP32 master weights)입니다. 이러한 특징들이 중요한 이유는 실무에서는 제한된 자원으로 최대한 큰 모델과 배치를 사용해야 좋은 성능을 얻을 수 있기 때문입니다.
코드 예제
import torch
from torch.cuda.amp import autocast, GradScaler
def train_with_mixed_precision(model, train_loader, optimizer, device='cuda'):
model.train()
# GradScaler: gradient underflow를 방지하기 위한 loss scaling
scaler = GradScaler()
for batch_idx, batch in enumerate(train_loader):
input_ids = batch['input_ids'].to(device)
labels = batch['labels'].to(device)
optimizer.zero_grad()
# autocast: 이 블록 안의 연산을 자동으로 FP16으로 변환
# Linear, Conv, MatMul 등은 FP16,
# Softmax, LayerNorm, Loss 등은 FP32 유지
with autocast():
outputs = model(input_ids, labels=labels)
loss = outputs.loss
# Backward pass with scaling
# loss를 큰 값(보통 2^16)으로 스케일링하여 underflow 방지
scaler.scale(loss).backward()
# Gradient unscaling 후 clipping
scaler.unscale_(optimizer) # 그래디언트를 원래 스케일로 복원
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
# Optimizer step with automatic scaling 조정
# 만약 inf/nan이 감지되면 업데이트를 스킵하고 scale 감소
scaler.step(optimizer)
scaler.update() # 다음 iteration을 위한 scale factor 조정
if batch_idx % 100 == 0:
print(f"Batch {batch_idx}, Loss: {loss.item():.4f}")
# BF16 사용 예시 (A100 이상의 GPU에서 권장)
def train_with_bf16(model, train_loader, optimizer, device='cuda'):
model.train()
# BF16은 FP16과 달리 loss scaling이 불필요
# (지수부가 FP32와 동일하여 표현 범위가 넓음)
for batch in train_loader:
optimizer.zero_grad()
# BF16 autocast - scaler 불필요
with autocast(dtype=torch.bfloat16):
outputs = model(batch['input_ids'].to(device))
loss = outputs.loss
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
optimizer.step()
# 사용 예시
model = GPT2Model.from_pretrained('gpt2').cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5)
# FP16 혼합 정밀도 학습
train_with_mixed_precision(model, train_loader, optimizer)
# 또는 BF16 (A100/H100에서 더 안정적)
# train_with_bf16(model, train_loader, optimizer)
설명
이것이 하는 일: 혼합 정밀도 학습은 모델의 forward/backward pass를 FP16으로 수행하여 빠른 Tensor Core 연산을 활용하고, loss를 큰 값으로 스케일링하여 작은 그래디언트가 0이 되는 것을 방지하며, 옵티마이저는 FP32 master weights를 유지하여 정확한 업데이트를 수행합니다. 첫 번째로, autocast() 컨텍스트 매니저가 자동으로 연산 종류에 따라 정밀도를 선택합니다.
왜 이렇게 하는지 이해하려면, 모든 연산을 FP16으로 하면 안 되는 이유를 알아야 합니다. 행렬 곱셈(MatMul)은 FP16으로 해도 결과가 거의 동일하지만, Softmax나 LayerNorm처럼 지수 연산이나 누적이 필요한 경우 FP16의 좁은 표현 범위(최대 65504)로는 overflow나 정밀도 손실이 발생합니다.
PyTorch는 내부적으로 "안전한" 연산만 FP16으로 변환하고 나머지는 FP32를 유지합니다. 그 다음으로, GradScaler의 loss scaling이 핵심입니다.
FP16의 최소 표현 가능한 양수는 약 6e-8인데, 많은 그래디언트가 이보다 작아서 0으로 처리됩니다(underflow). scaler.scale(loss)는 loss를 2^16(65536) 정도로 곱하여 backward 시 모든 그래디언트도 동일한 비율로 커지게 만듭니다.
이렇게 하면 작은 그래디언트도 FP16 범위 안에 들어옵니다. Optimizer step 전에 unscale_로 원래 크기로 돌려놓으므로 최종 업데이트는 정확합니다.
scaler.step(optimizer)과 scaler.update()는 dynamic loss scaling을 구현합니다. 만약 스케일링된 그래디언트에 inf나 nan이 발견되면(overflow), 이번 스텝을 건너뛰고 scale factor를 절반으로 줄입니다.
반대로 연속으로 몇 번 성공하면 scale을 2배로 늘려 더 작은 그래디언트까지 포착합니다. 이 자동 조정 덕분에 사용자가 수동으로 scale을 튜닝할 필요가 없습니다.
BF16(bfloat16)은 지수부가 FP32와 동일하여 표현 범위가 넓으므로 loss scaling이 불필요합니다. 대신 가수부(정밀도)가 FP16보다 낮지만, 대부분의 딥러닝 연산에서는 문제가 되지 않습니다.
A100 이상의 최신 GPU는 BF16 Tensor Core를 지원하므로, 더 간단하고 안정적인 BF16을 권장합니다. 여러분이 이 코드를 사용하면 동일한 GPU로 2배 큰 모델이나 2배 큰 배치를 사용할 수 있고, 학습 시간도 절반 이하로 줄어 더 많은 실험을 할 수 있습니다.
실무에서는 거의 모든 대규모 모델 학습이 혼합 정밀도를 기본으로 사용합니다.
실전 팁
💡 torch.set_float32_matmul_precision('high')를 설정하면 FP32 연산도 Tensor Core를 사용하여 약간 빨라집니다. 정밀도가 약간 낮아지지만 대부분의 경우 문제없습니다.
💡 Gradient clipping은 반드시 scaler.unscale_ 이후에 수행해야 합니다. Unscale 전에 하면 스케일링된 거대한 그래디언트를 clip하게 되어 의미가 없습니다.
💡 Embedding layer와 최종 출력 layer는 FP32로 유지하는 것이 안전합니다. model.embed.to(torch.float32) 같은 방식으로 특정 레이어만 FP32로 고정할 수 있습니다.
💡 A100 이상 GPU에서는 FP16보다 BF16을 사용하세요. 구현이 더 간단하고(scaler 불필요), 수치 안정성도 뛰어나며, 속도도 동일합니다. V100 이하에서는 FP16만 지원됩니다.
💡 GradScaler의 growth_interval 파라미터를 조정하여 scale 증가 빈도를 제어할 수 있습니다. 기본값 2000이 대부분 적절하지만, 매우 불안정한 학습에서는 더 크게(5000) 설정하세요.
8. 데이터 로더 최적화 - 병렬 전처리와 프리페칭
시작하며
여러분이 강력한 GPU로 모델을 학습시키는데, GPU 사용률이 50%밖에 안 되고 나머지 시간은 데이터를 기다리는 상황을 경험한 적 있나요? 비싼 GPU가 놀고 있는 동안 CPU가 데이터를 로딩하느라 분주한 모습을 보신 적 있으시죠?
이런 문제는 실제 개발 현장에서 자주 발생합니다. 최신 GPU는 초당 수천 개의 샘플을 처리할 수 있지만, 디스크에서 데이터를 읽고 전처리하는 속도가 따라가지 못하면 GPU가 idle 상태로 기다리게 됩니다.
특히 이미지나 오디오처럼 큰 데이터를 다룰 때 디스크 I/O가 병목이 되어 전체 학습 시간의 60-70%가 데이터 로딩에 낭비됩니다. A100 GPU 시간당 비용이 수만 원인 클라우드 환경에서는 이러한 비효율이 큰 비용 낭비로 이어집니다.
바로 이럴 때 필요한 것이 최적화된 데이터 로더입니다. 멀티프로세싱으로 전처리를 병렬화하고, 프리페칭으로 다음 배치를 미리 준비하며, 메모리 피닝으로 GPU 전송을 가속화하는 기법들을 적용해야 합니다.
개요
간단히 말해서, 데이터 로더 최적화는 여러 워커 프로세스가 동시에 데이터를 읽고 전처리하며, GPU가 현재 배치를 처리하는 동안 다음 배치를 미리 준비하여 GPU가 idle 시간 없이 계속 계산할 수 있도록 만드는 기법입니다. 왜 이 기법이 필요한지 실무 관점에서 설명하자면, GPU forward/backward는 밀리초 단위로 완료되지만, 디스크에서 데이터를 읽고 디코딩하며 전처리하는 데는 수십~수백 밀리초가 걸립니다.
단일 스레드로 순차 처리하면 GPU가 95%의 시간을 기다리게 됩니다. 8개 워커로 병렬 처리하면 이론적으로 8배 빠른 데이터 공급이 가능하여 GPU 사용률을 95% 이상으로 유지할 수 있습니다.
예를 들어, ImageNet으로 ResNet을 학습할 때 num_workers=0이면 에폭당 6시간이 걸리지만, num_workers=8이면 1시간으로 단축됩니다. 전통적인 방법과의 비교를 해보면, 기존에는 메인 스레드에서 한 배치를 로드하고 GPU로 전송한 후 계산하고, 다시 다음 배치를 로드하는 순차적인 방식이었다면, 최적화된 방법은 여러 워커가 백그라운드에서 계속 배치를 준비하고 큐에 쌓아두어 GPU가 요청하는 즉시 제공합니다.
이 기법의 핵심 특징은 첫째, 병렬 전처리(멀티프로세싱으로 CPU 코어 활용), 둘째, 프리페칭(다음 배치 미리 준비), 셋째, 제로카피 전송(pin_memory로 DMA 직접 전송)입니다. 이러한 특징들이 중요한 이유는 실무에서는 GPU가 가장 비싼 자원이므로 GPU 사용률을 최대화하는 것이 전체 비용과 시간을 결정하기 때문입니다.
코드 예제
import torch
from torch.utils.data import DataLoader, Dataset
import multiprocessing as mp
class OptimizedTextDataset(Dataset):
def __init__(self, data_path, tokenizer, max_length=512):
# 메모리에 전체 로드하지 않고 파일 위치만 저장
self.file_path = data_path
self.tokenizer = tokenizer
self.max_length = max_length
# 데이터셋 크기만 미리 계산 (인덱싱을 위해 필요)
with open(data_path, 'r') as f:
self.num_samples = sum(1 for _ in f)
def __len__(self):
return self.num_samples
def __getitem__(self, idx):
# 각 워커가 독립적으로 파일을 열고 필요한 라인만 읽음
# (실제로는 메모리 맵이나 인덱스 파일 사용이 더 효율적)
with open(self.file_path, 'r') as f:
for i, line in enumerate(f):
if i == idx:
# 전처리는 워커 프로세스에서 병렬로 수행
tokens = self.tokenizer(
line.strip(),
max_length=self.max_length,
truncation=True,
padding='max_length',
return_tensors='pt'
)
return {
'input_ids': tokens['input_ids'].squeeze(0),
'attention_mask': tokens['attention_mask'].squeeze(0)
}
def get_optimized_dataloader(dataset, batch_size, num_workers=None):
# CPU 코어 수의 절반 정도가 일반적으로 최적
if num_workers is None:
num_workers = min(8, mp.cpu_count() // 2)
dataloader = DataLoader(
dataset,
batch_size=batch_size,
shuffle=True,
num_workers=num_workers, # 멀티프로세스 워커 수
pin_memory=True, # GPU 전송 속도 향상
prefetch_factor=2, # 각 워커가 미리 준비할 배치 수
persistent_workers=True, # 워커 재사용 (에폭 간 재생성 방지)
drop_last=True # 마지막 불완전한 배치 제거
)
return dataloader
# 사용 예시
from transformers import GPT2Tokenizer
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
dataset = OptimizedTextDataset('large_corpus.txt', tokenizer)
# 최적화된 데이터 로더 생성
train_loader = get_optimized_dataloader(
dataset,
batch_size=32,
num_workers=8 # 8개 워커가 병렬로 데이터 준비
)
# 학습 루프 - GPU는 거의 쉬지 않고 계속 작동
model = ...
for epoch in range(num_epochs):
for batch in train_loader: # 배치가 즉시 준비되어 대기 시간 최소화
# 데이터가 이미 pinned memory에 있어 빠른 GPU 전송
input_ids = batch['input_ids'].to('cuda', non_blocking=True)
attention_mask = batch['attention_mask'].to('cuda', non_blocking=True)
outputs = model(input_ids, attention_mask=attention_mask)
loss = outputs.loss
loss.backward()
optimizer.step()
설명
이것이 하는 일: 최적화된 데이터 로더는 여러 워커 프로세스를 생성하여 각 워커가 독립적으로 데이터를 읽고 전처리한 후 큐에 넣어두고, 메인 프로세스는 큐에서 배치를 꺼내 GPU로 전송하는 파이프라인을 구축합니다. GPU가 계산하는 동안 워커들은 다음 배치들을 미리 준비합니다.
첫 번째로, num_workers 설정이 핵심입니다. 각 워커는 별도의 Python 프로세스로 실행되어 GIL(Global Interpreter Lock)의 영향을 받지 않고 진정한 병렬 처리가 가능합니다.
왜 이렇게 하는지 이해하려면, Python의 멀티스레딩은 GIL 때문에 CPU-bound 작업에서는 효과가 없고, 멀티프로세싱만이 여러 CPU 코어를 완전히 활용할 수 있다는 점을 알아야 합니다. 8개 워커는 8개 CPU 코어를 사용하여 동시에 8개 샘플을 처리하므로, 단일 워커 대비 이론적으로 8배 빠릅니다.
실제로는 오버헤드로 6-7배 정도 향상됩니다. 그 다음으로, pin_memory=True는 CPU 메모리를 페이지 고정(pinning)하여 GPU로의 전송을 가속화합니다.
일반 CPU 메모리는 OS가 언제든 스왑할 수 있어 GPU가 접근할 때 불안정하지만, pinned memory는 물리 메모리에 고정되어 DMA(Direct Memory Access)로 GPU가 CPU를 거치지 않고 직접 읽을 수 있습니다. 이로 인해 전송 속도가 2-3배 빨라집니다.
to('cuda', non_blocking=True)와 함께 사용하면 전송이 백그라운드에서 비동기로 수행되어 CPU는 즉시 다음 작업을 시작할 수 있습니다. prefetch_factor=2는 각 워커가 현재 제공 중인 배치 외에 추가로 2개 배치를 미리 준비해둔다는 의미입니다.
8개 워커에서 prefetch 2면 총 16개 배치가 큐에 대기하므로, GPU가 배치를 요청하면 즉시 제공됩니다. 이 값이 너무 크면 메모리를 많이 사용하고, 너무 작으면 GPU가 기다릴 수 있으므로 2-4가 적절합니다.
persistent_workers=True는 에폭이 끝나도 워커 프로세스를 종료하지 않고 재사용합니다. 워커를 생성하는 데는 수 초가 걸리므로, 짧은 에폭을 반복할 때 이 오버헤드가 무시할 수 없습니다.
한 번 생성된 워커를 계속 사용하면 에폭 간 전환이 즉각적으로 이뤄집니다. 다만 전체 학습 동안 워커 메모리가 유지되므로 메모리가 부족한 환경에서는 주의해야 합니다.
여러분이 이 코드를 사용하면 GPU 사용률이 50%에서 95% 이상으로 증가하여 전체 학습 시간이 절반 이하로 단축되고, 비용 대비 효율이 2배 향상됩니다. 실무에서는 데이터 로더 최적화만으로도 수일~수주의 학습 시간을 절약할 수 있습니다.
실전 팁
💡 num_workers는 많다고 무조건 좋은 게 아닙니다. CPU 코어 수를 초과하면 컨텍스트 스위칭 오버헤드로 오히려 느려집니다. 실험적으로 4, 8, 16을 시도해보고 GPU 사용률을 모니터링하세요(nvidia-smi dmon 명령어 사용).
💡 데이터셋이 작으면(수천 개 샘플) 멀티프로세싱 오버헤드가 이득보다 클 수 있습니다. 이 경우 num_workers=0으로 메인 프로세스에서만 처리하는 것이 더 빠를 수 있습니다.
💡 각 워커는 데이터셋과 전처리 함수를 복사하므로, 데이터셋이 메모리에 큰 객체를 들고 있으면 메모리 사용량이 워커 수배로 증가합니다. 파일 경로만 저장하고 __getitem__에서 읽는 방식을 사용하세요.
💡 SSD를 사용하면 디스크 I/O 병목이 크게 줄어듭니다. 특히 랜덤 액세스가 많은 경우 HDD 대비 10배 이상 빠릅니다. 클라우드 환경에서는 로컬 SSD를 선택하세요.
💡 매우 큰 데이터셋(TB 단위)은 전처리를 미리 완료하여 별도 파일로 저장하는 것이 효율적입니다. Arrow나 Parquet 포맷을 사용하면 메모리 매핑으로 빠른 랜덤 액세스가 가능합니다.