LLM은 어떻게 학습할까?

Decoder 모델의 학습 최적화 방법

0. 개요

대규모 언어 모델(LLM)을 학습할 때는 단순히 데이터를 준비하고 학습 코드를 실행하는 것 이상의 이해가 필요하다. 모델의 크기, 데이터의 양, GPU 메모리 제약, 그리고 분산 환경 활용 방식이 서로 맞물려 학습 효율과 안정성을 결정한다. 학습을 잘 설계하려면 크게 병렬화 방식, 메모리 사용 구조, 메모리 최적화 방법, 학습 설정, 데이터와 모델 선택 전략을 종합적으로 이해해야 한다.

---

1. 병렬화 기법 (Parallelism)

Data Parallelism은 모든 GPU에 모델 전체를 중복 로드하고, 데이터 배치만 나눠 처리하는 방식이다. 각 GPU가 서로 다른 배치로 forward와 backward를 수행한 뒤, gradient를 모아 평균내고 동기화한다. 이렇게 하려면 모델은 각 GPU에 중복으로 로드되어야 하고, 데이터만 분산된다. Hugging Face에서는 accelerate를 사용하여 쉽게 설정할 수 있으며, 예를 들어 다음과 같이 실행한다.

accelerate launch --multi_gpu train.py

이 방식은 단순하고 확장성이 좋지만, 모델이 크면 각 GPU가 모든 파라미터와 옵티마이저 상태, 그래디언트, 액티베이션을 다 들고 있어야 하므로 메모리 한계에 부딪히기 쉽다.

Model Parallelism은 모델이 한 GPU에 다 들어가지 않는 경우 사용하는 방법이다. 파라미터를 여러 GPU에 나눠 담는 방식이며, 크게 Tensor Parallelism과 Pipeline Parallelism으로 나뉜다. Tensor Parallelism은 큰 행렬 연산을 잘라 여러 GPU에서 동시에 수행하고, Pipeline Parallelism은 모델 레이어를 순서대로 GPU에 배치해 파이프라인 형태로 데이터를 흘려보낸다. 이 경우 데이터는 각 GPU에 중복 로드되지 않고, 순서대로 GPU를 거쳐 처리된다. Hugging Face에서는 device_map="auto"로 레이어를 자동 분산 로드할 수 있다.

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("model_name", device_map="auto")

print(model.hf_device_map)

요즘처럼 LLM 크기가 커지는 환경에서는 모델이 GPU 메모리에 안 들어가는 경우가 많아 Model Parallelism이 필수적이며, 동시에 데이터 양이 많으면 Data Parallelism과 결합한 Hybrid Parallelism을 쓴다. 단일 노드에서 4개의 GPU를 사용하는 경우에도 Hybrid Parallelism은 가능하지만 NVLink나 PCIe 대역폭 한계로 효율이 제한적일 수 있다.

---

2. 학습 프레임워크와 주요 툴

Hugging Face Transformers는 모델 정의와 로딩(AutoModelForCausalLM), 토크나이저(AutoTokenizer), 학습 설정과 실행(TrainingArguments, Trainer)을 제공한다. Trainer 내부에서 accelerate를 사용하므로 별도 병렬화 설정이 필요 없다. accelerate는 Data Parallelism, FSDP, DeepSpeed 등을 지원하지만, Model Parallelism의 Tensor/Pipeline은 직접 제공하지 않는다. 다만 Hugging Face 모델의 device_map 설정을 통해 레이어 분산 로드는 가능하다.

DeepSpeed는 대규모 모델 학습 최적화에 특화된 프레임워크로 ZeRO(Zero Redundancy Optimizer)와 Offloading 기능을 제공한다. ZeRO Stage 1은 옵티마이저 상태를 분산, Stage 2는 옵티마이저 상태와 그래디언트를 분산, Stage 3는 파라미터까지 분산해 각 GPU의 메모리 부담을 줄인다. Stage 3에서도 파라미터는 연산 시점에 GPU로 로드되어 사용되므로 GPU 메모리에 전혀 안 올라가는 것은 아니지만, 상시 상주하지 않는다.

Offloading은 CPU나 NVMe SSD를 활용해 파라미터나 옵티마이저 상태를 GPU 밖으로 옮겨 GPU 메모리를 절약하는 방법이다. ZeRO Stage만 설정해서는 오프로딩이 자동 적용되지 않으며, offload_optimizer, offload_param 옵션을 통해 위치를 지정해야 한다.

"zero_optimization": {

"stage": 3,

"offload_optimizer": { "device": "cpu" },

"offload_param": { "device": "nvme", "nvme_path": "/nvme_storage" }

}

위 설정은 옵티마이저 상태를 CPU, 파라미터를 NVMe에 두고 필요 시 GPU로 불러온다.

---

3. 메모리 사용 구조

파라미터 1개의 크기는 FP32일 때 4바이트, FP16이나 BF16일 때 2바이트다. 학습 시에는 파라미터 외에도 그래디언트(대개 1배), 옵티마이저 상태(AdamW의 경우 m, v 두 개로 약 2배), 액티베이션 메모리가 필요하다. 이를 합치면 파라미터 크기의 약 4~5배에 액티베이션 메모리가 추가된다.

Activation 메모리는 대략 다음과 같이 계산할 수 있다.

activation_memory ≈ batch_size × seq_length × hidden_size × num_layers × bytes_per_element × C

여기서 C는 구현 상수로, 중간 텐서나 캐시를 포함해 2~6배가 될 수 있다.

Data Parallelism에서는 각 GPU가 모델 전체를 복제하므로, 총 메모리를 GPU 수로 나누는 식의 계산은 맞지 않다. 예를 들어 32B 파라미터 모델(bf16 기준 64GB)을 A100 80GB에서 DP로 학습하려 하면, 파라미터(64GB) + 그래디언트(~64GB) + 옵티마이저 상태(~128GB) + 액티베이션으로 80GB를 훌쩍 넘기게 된다.

---

4. 메모리 최적화 방법

Mixed Precision은 FP16이나 BF16을 사용해 메모리 절약과 연산 속도 향상을 도모한다. BF16은 FP16보다 표현 범위가 넓어 gradient scaling 없이도 안정적이며, A100/H100 GPU에서 추천된다. LayerNorm, Bias, 일부 정규화나 Softmax, 일부 손실 계산 등은 수치 안정성을 위해 FP32로 유지한다. AdamW 옵티마이저의 m과 v 역시 FP32를 유지하는데, 이는 장기 누적 통계의 정밀도를 보장하기 위해서다.

Recompute(Activation Checkpointing)는 forward 시 모든 액티베이션을 저장하지 않고 일부만 저장한 뒤, backward 시 다시 계산해 메모리를 줄이는 방법이다. 메모리는 줄지만 연산량과 시간이 늘어난다.

ZeRO/FSDP는 파라미터, 그래디언트, 옵티마이저 상태를 샤딩해 각 GPU의 메모리 부담을 줄인다. 여기에 Offloading을 결합해 옵티마이저 상태는 CPU, 파라미터는 NVMe에 저장하면 GPU 메모리를 크게 줄일 수 있다.

Gradient Accumulation은 작은 마이크로 배치로 여러 번 forward/backward를 수행하고, gradient를 누적해 한 번에 업데이트하는 방식이다. 이렇게 하면 한 번의 forward/backward에 올라가는 데이터량이 줄어 액티베이션 메모리를 절약할 수 있다. 예를 들어 총 effective batch를 64로 만들고 싶을 때, batch_size=64, accumulation=1이면 한 번에 64개 샘플이 GPU에 올라가지만, batch_size=8, accumulation=8이면 한 번에 8개만 올라가 메모리가 절약된다.

Gradient Clipping은 그래디언트 폭주를 방지한다. Hugging Face Trainer에서는 다음과 같이 설정할 수 있다.

TrainingArguments(max_grad_norm=1.0)

PyTorch 직접 구현 시에는 clip_grad_norm_를 호출한다.

---

5. 학습 설정과 배치/시퀀스 길이 결정

학습 설정의 핵심 파라미터로는 max_seq_length, per_device_batch_size, gradient_accumulation_steps가 있다. effective_batch_size는 per_device_batch_size × num_gpus × gradient_accumulation_steps로 계산한다.

배치 크기와 시퀀스 길이는 GPU 한 장의 사용 가능 메모리에서 모델 관련 메모리와 액티베이션 메모리를 뺀 나머지에 맞춰 정한다.

모델 관련 메모리(파라미터+그래디언트+옵티마이저 상태) + 액티베이션 메모리 ≤ GPU 메모리

모델 관련 메모리는 ZeRO-3/FSDP, Model Parallelism, Offloading, Recompute 등으로 줄이고, 남는 메모리에 맞춰 per_device_batch_size와 max_seq_length를 설정한 뒤, 부족하면 gradient_accumulation_steps로 effective batch를 맞춘다.

---

6. Optimizer, Learning Rate, Scheduler

Optimizer는 gradient를 받아 weight를 어떻게 업데이트할지 결정하는 방식이다. AdamW는 Adaptive Learning Rate와 Momentum, Weight Decay를 결합해 학습 속도와 일반화 성능을 높인다. Learning Rate는 한 번 업데이트할 때의 크기를 결정하며, 너무 크면 발산하고 너무 작으면 학습이 느리다. Scheduler는 학습 과정에서 Learning Rate를 시간에 따라 조절하는 역할을 한다. 예를 들어 Warmup, Cosine Decay, Step Decay 등이 있다.

하나의 step에서 순서는 forward → loss 계산 → backward → Scheduler가 이번 step의 learning rate 결정 → Optimizer가 해당 learning rate를 반영해 weight 업데이트 순으로 진행된다.

---

7. 데이터 포맷과 모델 선택

학습 데이터는 Instruction-style과 Base-style로 나뉜다. Instruction-style은 instruction, input, output 필드를 가지며, Instruct 모델 학습이나 Instruction-tuning에 쓰인다.

[

{

"instruction": "Summarize the following article.",

"input": "The stock market experienced a significant drop today due to...",

"output": "The stock market fell sharply today due to..."

}

]

Base-style은 text 필드 하나로 구성되며, 언어 모델 프리트레이닝이나 도메인 적응에 사용된다.

[

{ "text": "The quick brown fox jumps over the lazy dog." }

]

새로운 도메인 학습은 Base 모델에 Base-style로 먼저 도메인 적응을 하고 Instruction-tuning을 하는 것이 안정적이지만, Base 모델에 바로 Instruction-style로 학습하는 것도 가능하다. 다만 이 경우 충분한 데이터와 스텝을 주는 것이 좋다.

토크나이징 과정에서 padding과 truncation 정책은 액티베이션 메모리와 처리 속도에 직접적인 영향을 미친다. 데이터 길이 분포를 분석해 max_seq_length를 적절히 설정하는 것이 중요하다.

---

8. 추가 최적화와 아키텍처 팁

LoRA, QLoRA, PEFT는 파라미터 효율적 학습 기법으로, 전체 가중치를 업데이트하지 않고 일부 어댑터만 학습해 메모리와 시간을 절약한다. Flash Attention은 긴 시퀀스에서 Attention 연산을 메모리와 시간 면에서 최적화한다. 추론 단계에서는 KV 캐싱으로 길이 증가에 따른 연산 낭비를 줄일 수 있다.

---

마무리

LLM 학습 설계의 핵심은 GPU 메모리 한도 안에서 모델 크기, 데이터 길이, 배치 크기, 병렬화, 오프로딩, Mixed Precision, Recompute, Gradient Accumulation 등을 적절히 조합하는 것이다. 모델이 GPU에 들어가면 Data Parallelism이 단순하고 효율적이지만, 안 들어가면 Model Parallelism이 필수이고, 대부분은 Hybrid Parallelism을 쓴다. ZeRO-3/FSDP와 Offloading으로 복제를 제거하고, Mixed Precision과 Recompute로 메모리를 줄이며, Accumulation으로 배치를 확보하는 것이 일반적인 로드맵이다. Optimizer, Learning Rate, Scheduler의 역할을 이해하고, 데이터 포맷과 토크나이징 전략을 맞추는 것이 성능과 효율을 좌우한다.