메모리 사용량 줄이기

학습 워크플로우는 메모리 사용량을 줄이는 방향으로 최적화할 수 있으며, TRL은 이를 돕기 위해 여러 내장 기능을 제공합니다.

아래에서는 이러한 기법을 소개하며, 특정 환경에 가장 적합한 구성을 찾기 위해 다양한 조합을 실험해 보기를 권장합니다.

각 메서드에는 지원되는 trainer에 대한 예제가 포함되어 있습니다. 특정 기법이 사용 중인 trainer와 호환되는지 확실하지 않다면, 해당 trainer의 문서를 참고해 주세요.

모든 trainer에서 지원되는 activation checkpointing과 같은 추가 전략에 대해서는 transformers 최적화 가이드를 참고하세요.

단축

데이터셋의 시퀀스 길이는 크게 달라질 수 있습니다. 데이터가 배치로 처리되면 배치 내 가장 긴 시퀀스에 맞춰 패딩이 수행되므로, 대부분의 시퀀스가 상대적으로 짧더라도 메모리 사용량이 높아질 수 있습니다.

메모리 사용량을 줄이려면 시퀀스를 적절한 길이로 잘라내는 것이 중요합니다. TRL 트레이너는 기본적으로 시퀀스를 잘라내지만, 특정 사용 사례에 더 잘 맞도록 기본 잘라내기 길이를 조정할 수 있습니다.

from trl import DPOConfig

training_args = DPOConfig(..., max_length=...)

from trl import SFTConfig

training_args = SFTConfig(..., max_length=...)

max_length 값을 어떻게 선택하나요?

max_length이 너무 작으면 토큰의 상당 부분이 폐기되어 학습에 기여하지 못합니다. 너무 크면 메모리 사용량이 급증하여 OOM(Out-Of-Memory) 오류가 발생할 수 있습니다. packing이나 padding-free를 사용하지 않는 경우, max_length가 크면 많은 토큰이 padding이 되어 비효율적인 학습이 될 수 있습니다.

적절한 값을 선택할 수 있도록, 데이터셋의 시퀀스 길이 분포를 시각화하는 유틸리티를 제공합니다.

패킹

이 기법은 SFT 학습 및 FlashAttention(또는 그 변형)을 사용하는 설정에서만 사용할 수 있습니다.

Truncation에는 몇 가지 단점이 있습니다:

1. 정보 손실: 시퀀스 끝부분의 중요한 토큰이 삭제될 수 있습니다.
2. 잘림 길이 선택: 너무 짧으면 데이터가 손실되고, 너무 길면 효율이 떨어집니다.

Packing은 여러 시퀀스를 동일한 학습 행에 묶어 각 행을 max_length까지 채움으로써 이러한 문제를 완화합니다.

TRL은 패딩을 최소화하면서 시퀀스를 효율적으로 그룹화하는 Best-Fit Decreasing (BFD) 빈 패킹을 사용하여 패킹을 구현합니다. 시퀀스가 max_length를 초과할 경우, 오버플로 토큰을 처리하는 방식은 다양한 전략에 따라 결정됩니다.

TRL은 세 가지 전략을 지원합니다:

• "bfd" (기본값): Best-Fit Decreasing packing을 사용합니다. 시퀀스가 max_length를 초과하면 초과된 토큰은 삭제됩니다.

• "bfd_split": Best-Fit Decreasing packing을 사용하지만, 긴 시퀀스는 패킹 전에 max_length 이하의 청크로 분할됩니다. 이를 통해 모든 토큰이 보존되며, Fewer Truncations Improve Language Modeling에서 제안한 접근 방식을 따릅니다.

• "wrapped": 모든 토큰이 하나의 스트림으로 연결된 후 고정 길이의 블록으로 분할됩니다. 이 방식은 패딩을 최소화하지만 관련이 없는 예시가 섞일 수 있습니다. 이 전략은 문헌에서 설명된 concatenate-then-split 전처리(예: Fewer Truncations Improve Language Modeling)에 해당합니다. Qwen3-Coder-Next Technical Report에서 논의된 바와 같이, 데이터셋의 상당 부분에서 시퀀스 연속성을 끊어 성능을 저하시킨다는 단점이 있습니다.

모든 시퀀스가 max_length보다 짧다면, 절단이나 분할이 필요 없으므로 bfd와 bfd_split은 동일한 방식으로 동작합니다.

from trl import SFTConfig

training_args = SFTConfig(
    ...,
    packing=True,
    packing_strategy="bfd",
    max_length=512,
)

매개변수 효율적 미세 조정용 PEFT

LoRA와 같은 Parameter-Efficient Fine-Tuning (PEFT) 방법은 학습 중 메모리 사용량을 줄이는 가장 효과적인 기법 중 하나입니다. 모델의 모든 파라미터를 학습하는 대신, PEFT 방법은 소수의 어댑터 파라미터만 학습하여 메모리 요구 사항을 크게 줄이고, 제한된 하드웨어에서 더 큰 모델을 파인튜닝할 수 있게 합니다.

PEFT를 TRL과 함께 사용하는 방법에 대한 자세한 내용, 다양한 어댑터 방법, 양자화 옵션 및 고급 구성에 대해서는 PEFT Integration을 참조하세요.

메모리 사용량을 줄이기 위해 PEFT를 사용하려면:

from datasets import load_dataset
from peft import LoraConfig
from trl import SFTTrainer

dataset = load_dataset("trl-lib/Capybara", split="train")

peft_config = LoraConfig()

trainer = SFTTrainer(
    model="Qwen/Qwen2.5-0.5B",
    train_dataset=dataset,
    peft_config=peft_config,
)

PEFT는 양자화(4비트 또는 8비트)와 같은 다른 메모리 절감 기법과 결합하여 더 큰 메모리 절감 효과를 얻을 수 있습니다. 양자화 예제는 PEFT Integration을 참고하세요.

패딩 없음

Padding-free batching은 메모리 사용량을 줄이기 위한 대안적 접근 방식입니다. 이 방법에서는 먼저 배치(batch)를 샘플링한 후 패딩(padding)을 피하기 위해 단일 시퀀스로 평탄화(flatten)합니다. 여러 샘플의 부분을 결합하여 불완전한 시퀀스가 생성될 수 있는 packing과 달리, padding-free batching은 모든 시퀀스가 완전하고 온전한 상태로 유지되도록 보장합니다.

from trl import DPOConfig

training_args = DPOConfig(..., padding_free=True, model_init_kwargs={"attn_implementation": "kernels-community/flash-attn2"})

from trl import SFTConfig

training_args = SFTConfig(..., padding_free=True, model_init_kwargs={"attn_implementation": "kernels-community/flash-attn2"})

Activation 오프로딩

Activation offloading은 forward pass 동안 activation 텐서를 CPU RAM으로 일시적으로 이동하고, backward pass에 필요할 때만 다시 가져와서 GPU VRAM 사용량을 줄이는 메모리 효율성 기법입니다. 이로 인해 학습 시간이 약간 증가하는 대가로 피크 메모리 사용량이 크게 감소합니다.

SFT 학습 구성에서 activation offloading을 활성화하려면:

from trl import SFTConfig

training_args = SFTConfig(..., activation_offloading=True)

내부적으로 activation offloading은 PyTorch의 saved_tensors_hooks를 구현하여 forward pass 중 activation을 가로채 처리합니다. 크기와 컨텍스트를 기반으로 어떤 텐서를 offload할지 지능적으로 관리하며, 비효율적인 output 텐서의 offloading을 피합니다. 성능 최적화를 위해, 플래그(기본값은 true)를 통해 CUDA streams를 사용하여 계산과 CPU-GPU 전송을 오버랩할 수 있습니다.

Padding 시퀀스를 배수로

이 기법은 현재 SFT 및 Reward 트레이너에서 지원됩니다.

이 옵션을 활성화하면 모든 시퀀스가 지정된 값의 배수로 패딩되도록 보장됩니다.
이는 시퀀스 길이를 메모리에 유리한 경계에 정렬하여 일부 하드웨어에서 계산 효율성을 높일 수 있습니다.

from trl import SFTConfig

training_args = SFTConfig(..., pad_to_multiple_of=2048)

from trl import RewardConfig

training_args = RewardConfig(..., pad_to_multiple_of=2048)

PyTorch 캐싱 할당기

긴 실행, 특히 온라인 RL(GRPO, RLOO, Online DPO)에서는 GPU 메모리가 조각날 수 있습니다. expandable_segments:True를 설정하면 PyTorch의 캐싱 할당기가 기존 세그먼트를 확장하도록 하여, 할당된 메모리와 예약된 메모리 사이의 격차를 줄일 수 있습니다:

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
export PYTORCH_ALLOC_CONF=expandable_segments:True  # canonical name since PyTorch 2.10

vLLM sleep mode(vllm_enable_sleep_mode=True)를 사용할 경우 vLLM은 이 설정을 우회하기 위해 PYTORCH_CUDA_ALLOC_CONF만 읽습니다. 따라서 PYTORCH_ALLOC_CONF만 설정하지 마세요.

온라인 방법에서 생성을 위한 모델 수집 비활성화

DeepSpeed ZeRO-3을 사용할 때 모델 가중치는 여러 GPU에 걸쳐 샤드됩니다. 온라인 방식은 학습 과정의 일부로 모델에서 완성본을 생성하는 것을 포함합니다. 이 단계에서 모델 가중치는 생성을 위해 일시적으로 단일 GPU에 모아집니다. 매우 큰 모델의 경우, 이 수집 과정으로 인해 OOM 오류가 발생할 수 있으며, 이는 #2250 이슈에서 설명되어 있습니다.

이 문제를 겪는 경우, 다음 매개변수를 설정하여 생성을 위한 모델 가중치 수집을 비활성화할 수 있습니다:

from trl import GRPOConfig

training_args = GRPOConfig(..., ds3_gather_for_generation=False)

from trl.experimental.online_dpo import OnlineDPOConfig

training_args = OnlineDPOConfig(..., ds3_gather_for_generation=False)

from trl import RLOOConfig

training_args = RLOOConfig(..., ds3_gather_for_generation=False)

이 조정은 모델 가중치가 수집되는 것을 방지하여 OOM 오류를 피할 수 있지만, 생성 속도가 느려질 수 있습니다.

vLLM sleep 모드

온라인 학습 방법에서 vLLM을 생성 백엔드로 사용할 때, sleep mode를 활성화하여 최적화 단계 동안 vLLM의 매개변수와 캐시를 CPU RAM으로 오프로드하고, 가중치 동기화 및 생성이 필요할 때 다시 GPU VRAM으로 로드할 수 있습니다.

from trl import GRPOConfig

training_args = GRPOConfig(..., vllm_enable_sleep_mode=True)

from trl import RLOOConfig

training_args = RLOOConfig(..., vllm_enable_sleep_mode=True)

vLLM의 가중치와 캐시를 오프로딩하면 GPU 메모리 사용량을 낮게 유지하는 데 도움이 되며, 이는 대규모 모델을 학습하거나 GPU 자원이 제한된 환경에서 특히 유용합니다. 그러나 sleep 모드에서 vLLM 엔진을 깨우는 과정에서는 host-device 간 전송 지연이 발생하여 학습 속도에 약간의 영향을 줄 수 있습니다.

Activation checkpointing

Activation checkpointing은 forward pass 동안 모든 중간 activation을 저장하지 않고, 대신 backward pass 동안 다시 계산함으로써 계산량을 메모리와 교환합니다.

from trl import SFTConfig

training_args = SFTConfig(..., activation_checkpointing=True)

메모리 사용량을 최적화하기 위해 모든 trainer에서 activation checkpointing이 기본적으로 활성화되어 있습니다. 필요할 경우 activation_checkpointing=False를 설정하여 비활성화할 수 있습니다.

선택적 활성화 체크포인트

SFTTrainer를 사용하면 forward pass 중에 attention 출력을 저장하고 backward pass에서 다시 계산하지 않을 수 있습니다. 이는 레이어당 hidden-state 크기의 텐서 하나를 추가로 사용하는 대가로, 긴 컨텍스트에서 checkpointing으로 인한 대부분의 성능 저하를 회복시킵니다. 이 기능은 non-reentrant checkpointing을 강제합니다.

from trl import SFTConfig

training_args = SFTConfig(
    ...,
    activation_checkpointing=True,
    activation_checkpointing_kwargs={"selective": True},
)

추가 메모리 최적화 기법은 Transformers Performance Guide를 참고하세요.