LLM 학습 인프라의 메모리 관리와 분산 병렬화

모델 크기가 바꾸는 메모리 설계

GPT-4, Claude, LLaMA 같은 대규모 언어 모델은 수백억~수천억 파라미터 규모로 확장되면서 단일 GPU 메모리의 경계를 넘는다. 학습 환경에서는 파라미터 자체뿐 아니라 옵티마이저 상태, 그래디언트, 활성화 값이 동시에 자리를 차지한다. 따라서 메모리 관리와 멀티노드 분산 전략은 모델을 올리기 위한 부가 작업이 아니라 학습 설계의 일부가 된다.

GPT-3 175B는 FP32 기준 파라미터만 700GB이고, LLaMA 65B는 260GB다. FP16을 쓰면 175B의 파라미터 메모리는 350GB로 줄어든다. 다만 학습 중에는 Optimizer State가 2-3배, 그래디언트와 Activation이 추가된다.

⤢✕

추론의 메모리 양상도 단순하지 않다. 자기회귀 생성에서는 과거 키와 값을 KV Cache에 보관하며, Batch Size를 높이면 처리량은 늘지만 메모리는 선형 증가한다. INT8·INT4 양자화는 메모리를 1/4~1/8 수준으로 줄일 수 있다. LLaMA 13B를 FP16으로 추론하면서 Batch=32, Seq=2048을 사용하면 약 60GB가 필요하다.

메모리 벽을 낮추는 방법들

ZeRO(Zero Redundancy Optimizer)는 DeepSpeed가 제공하는 메모리 분할 방식이다. Stage 1은 Optimizer State를 분할해 4배 메모리를 절감하고, Stage 2는 Gradient Partitioning으로 8배 절감한다. Stage 3는 Parameter Partitioning까지 포함해 Linear scaling을 목표로 한다. ZeRO-Offload는 상태를 CPU 또는 NVMe로 옮긴다.

⤢✕

Gradient Checkpointing은 중간 활성화 값을 계속 저장하는 대신 필요한 시점에 다시 계산한다. 체크포인트 간격이 k이면 메모리는 1/k로 감소하고, 연산은 1회 forward 추가로 약 33% 증가한다. Transformer 블록마다 체크포인트를 두는 방식이 적용 예다.

Flash Attention은 시퀀스 길이 N에서 어텐션의 메모리 사용량을 O(N²)에서 O(N)으로 낮춘다. 연산을 블록으로 나누어 SRAM에서 처리하는 Tiling 기법을 사용하며, 메모리 대역폭 최적화로 속도도 2-4배 향상된다.

양자화는 가중치 표현을 줄이는 선택지다. Post-Training Quantization은 학습 이후 가중치를 변환하고, Quantization-Aware Training은 학습 과정에서 양자화를 시뮬레이션한다. INT8은 2배 메모리 절감과 Tensor Core 활용을, INT4/GPTQ는 4배 절감과 정확도 손실 최소화를 목표로 한다. Mixed Precision에서는 중요 레이어만 FP16으로 유지한다.

CPU/NVMe Offloading은 GPU에 모든 상태를 상주시킬 수 없을 때 계층형 저장 공간을 활용한다.

⤢✕

ZeRO-Infinity는 1조 파라미터 모델 학습을 가능하게 하며, DeepSpeed Offload는 NVMe 대역폭 3.5GB/s를 활용한다. 이때 연산과 데이터 이동을 파이프라이닝해 겹치는 방식이 중요하다.

레이어와 노드에 작업을 나누는 방식

Tensor Parallelism은 하나의 레이어를 여러 GPU에 나누는 Intra-Layer 병렬화다. Megatron-LM은 Transformer의 어텐션과 FFN을 분할해 구현한다. QKV 프로젝션에는 Column-wise split, Output 프로젝션에는 Row-wise split을 적용하며, 레이어 내부 통신에는 All-Reduce가 쓰인다.

Pipeline Parallelism은 레이어를 스테이지별로 분리하는 Inter-Layer 병렬화다. GPipe는 마이크로 배치로 파이프라인 버블을 줄이고, PipeDream은 비동기 파이프라인으로 버블 최소화를 노린다. 1F1B(One-Forward-One-Backward)는 메모리 효율을 고려한 스케줄링 방식이다.

⤢✕

Hybrid Parallelism은 Data, Tensor, Pipeline Parallelism을 동시에 쓰는 3D Parallelism을 포함한다. PyTorch 2.0은 FSDP + Tensor Parallelism 조합을 지원한다. 노드 안에서는 NVLink를 활용한 Tensor Parallelism을, 노드 사이에서는 네트워크 통신을 줄이기 위한 Pipeline Parallelism을 두고, 데이터 복제는 최외곽 레이어에 둔다.

모델이 GPU 메모리에 들어가면 Data Parallelism만 사용한다. 13B~70B 중간 규모에서는 Tensor Parallelism + Data Parallelism을 조합하고, 100B+ 초대규모 모델에는 3D Parallelism + ZeRO Stage 3가 대상이다. 추론에서는 지연 시간을 낮추기 위해 Tensor Parallelism을 사용한다.

멀티노드 통신 병목 다루기

노드 내부 GPU 연결과 노드 간 연결은 서로 다른 특성을 가진다. 다음 구성에서는 노드 내부에 NVLink 600GB/s, 노드 간에는 InfiniBand HDR 200Gbps가 사용된다.

⤢✕

NCCL에서는 Ring All-Reduce가 대역폭을, Double Binary Tree가 지연 시간을 각각 최적화한다. 홉 수는 Ring All-Reduce가 O(N), Double Binary Tree가 O(log N)이다. 점검할 환경 변수는 NCCL_DEBUG=INFO, NCCL_IB_DISABLE=0이며, InfiniBand에서는 ib_timeout, ib_retry_count를 튜닝한다.

Gradient Accumulation은 작은 마이크로 배치로 계산한 그래디언트를 누적한 뒤 한 번에 업데이트한다. N배 누적하면 통신 횟수는 1/N이 되고, 큰 배치 크기의 효과를 유지하면서 메모리를 아낄 수 있다.

통신을 backward 연산과 겹치게 하는 것도 중요하다.

⤢✕

Mixed Precision 학습에서는 동적 Loss Scaling으로 언더플로우를 막는다. 그래디언트 크기에 맞춰 스케일을 조절하고, Overflow가 검출되면 값을 낮춘다. 일정 스텝 동안 Overflow가 없으면 스케일을 높인다.

프레임워크가 제공하는 분산 학습 기능

DeepSpeed는 Microsoft가 개발한 오픈소스 프레임워크다. ZeRO Optimizer의 Stage 1/2/3 및 Offload, 3D Parallelism 통합 API, Sparse Attention, 1-bit Adam을 제공한다. 1-bit Adam은 통신량을 1/16로 감소시킨다.

Megatron-LM은 NVIDIA가 개발했으며 GPT/BERT 학습에 특화됐다. Tensor Parallelism을 효율적으로 구현하고, Pipeline Parallelism에서는 1F1B 스케줄링을 사용한다. 긴 시퀀스를 위한 Sequence Parallelism과 BERT 어텐션 연산 커널 최적화도 포함한다.

FSDP(Fully Sharded Data Parallel)는 PyTorch 네이티브의 ZeRO Stage 3 구현이다. Mixed Precision Sharding은 FP16 통신과 FP32 업데이트를 사용하며, CPU Offload는 Optimizer State를 CPU로 옮긴다. 모델 분할에는 auto_wrap_policy를 사용하는 Auto Wrap이 있다.

Alpa는 JAX 기반 자동 병렬화 프레임워크다. Inter/Intra-Operator Parallelism을 자동 선택하고, ILP(Integer Linear Programming) 기반으로 Compile 타임에 최적 전략을 탐색한다.

OS가 학습 프로세스에 미치는 영향

메모리 설정에서는 Huge Pages가 2MB 페이지로 TLB 미스를 50% 낮춘다. vm.swappiness=0은 스왑을 비활성화하고 OOM 시 명시적으로 종료하게 하며, vm.overcommit_memory=1은 메모리 오버커밋을 허용한다. cgroup memory limit은 컨테이너별 메모리 격리에 사용한다.

GPU와 CPU의 NUMA 거리를 확인하고 프로세스를 가까운 노드에 바인딩한다.

# GPU 0과 가까운 CPU 노드 확인
nvidia-smi topo -m

# NUMA 노드 바인딩
numactl --cpunodebind=0 --membind=0 python train.py

# GPU와 CPU 친화성 매핑
taskset -c 0-15 python train.py  # CPU 0-15번 사용

I/O 스케줄러는 저장장치 특성에 맞춘다. NVMe SSD에는 멀티큐를 사용하는 none 또는 kyber, HDD에는 순차 접근에 최적화된 mq-deadline이 해당한다. Read-ahead는 blockdev --setra 8192 /dev/nvme0n1로 설정하며, Direct I/O는 대용량 순차 읽기에서 효율적이다.

네트워크 쪽에서는 net.core.rmem_max=134217728로 TCP 버퍼 크기를 조정하고, net.ipv4.tcp_congestion_control=bbr로 Congestion Control을 설정한다. RDMA를 위해 mlx5_core, ib_core 모듈을 로드하며, IRQ Affinity로 네트워크 카드 인터럽트를 CPU 코어에 고정한다.

요청 처리 경로의 추론 최적화

KV Cache를 페이지 단위로 관리하는 vLLM의 Paged Attention은 동적 할당·해제에서 생기는 메모리 단편화를 막는다. Prefix Caching은 공통 프롬프트를 재사용하고, Continuous Batching은 요청마다 독립적으로 처리해 처리량을 높인다.

Speculative Decoding은 작은 모델이 초안(Draft)을 만들고 큰 모델이 이를 병렬 검증하는 방식이다. 출력 품질은 동일하게 유지하면서 지연 시간을 2-3배 줄인다. LLaMA 7B → LLaMA 70B 조합이 예시다.

추론에서도 Tensor Parallelism은 작은 배치에서 효과가 있다. All-Reduce 오버헤드가 연산 시간보다 작고, NVLink로 통신 지연을 낮출 수 있다. 적용 예로 8x A100에서 175B 모델이 1초 이하로 응답한다.

관측 데이터로 병목 찾기

GPU 사용률과 메모리, 온도를 실시간으로 확인하고 CUDA 프로파일을 수집할 수 있다.

# 실시간 GPU 사용률
nvidia-smi dmon -s pucvmet

# DCGM (Data Center GPU Manager)
dcgmi dmon -e 155,203,204  # SM 활성도, 메모리 사용률, 온도

# Profiling
nsys profile --trace=cuda,nvtx python train.py

PyTorch 메모리 스냅샷은 학습 중 메모리 상태를 기록하는 데 사용한다.

# PyTorch 메모리 스냅샷
torch.cuda.memory._record_memory_history()
# ... training ...
snapshot = torch.cuda.memory._snapshot()
torch.cuda.memory._dump_snapshot("memory.pickle")

통신 병목은 nccl-tests/build/all_reduce_perf -b 1G -e 8G로 NCCL Tests를 실행해 확인한다. 이론치 200Gbps와 실제 대역폭을 비교하고, 병목이 노드 내부인지 노드 간인지 구분한다. nsys 프로파일에서는 통신과 연산의 중첩 여부를 분석한다.

운영 중 확인할 항목과 복구 수단

학습 시작 전에는 GPU 드라이버, CUDA, cuDNN 버전의 일치 여부를 확인한다. NCCL은 최신 안정 버전인 2.18+를 사용하고, InfiniBand 연결 상태는 ibstatus, ibv_devinfo로 점검한다. 파일 시스템은 Lustre/BeeGFS 같은 병렬 FS를 권장하며, 체크포인트는 공유 스토리지에 두고 복제본을 유지한다.

학습 중에는 GPU 사용률 95% 이상, 메모리 사용률 90% 이상을 목표로 삼되 OOM을 주의한다. 통신 대역폭은 이론치의 80% 이상을 기준으로 보고, NaN/Inf가 발생하면 Loss Scaling을 조정한다. 처리량(Tokens/sec)은 기준치와 비교해 추적한다.

장애 대응에는 SLURM --requeue 옵션을 이용한 Automatic Restart, 노드 추가·제거에 대응하는 Elastic Training이 있다. 체크포인트는 1-2시간마다 저장하고, grad_clip=1.0의 그래디언트 클리핑으로 발산을 막는다. GPU 메모리 에러는 Health Check로 주기적으로 검사한다.

대규모 LLM 인프라는 파라미터를 여러 GPU에 나누는 것만으로 끝나지 않는다. ZeRO, Megatron-LM, FSDP로 메모리 제약을 다루고, 파이프라인 병렬화와 NCCL·InfiniBand 설정으로 통신 비용을 조절한다. 여기에 OS 메모리 관리와 네트워크 운영을 결합해야 수백 개 GPU를 하나의 가속기처럼 활용할 수 있다.