[axolotl] Axolotl: Long-Context 학습을 위한 Hidden State Offloading 최적화 (Non-Reentrant Checkpointing 지원)

PR 링크: axolotl-ai-cloud/axolotl#3776 상태: Merged | 변경: +568 / -53

들어가며

대규모 언어 모델(LLM)을 학습시키거나 파인튜닝할 때, 특히 긴 시퀀스 길이(long context)를 다룰 때 GPU 메모리 부족은 흔히 마주하는 문제입니다. Gradient Checkpointing은 이러한 문제를 완화하기 위한 강력한 기술 중 하나로, 역전파 시 필요한 중간 활성화 값(activations)을 모두 저장하는 대신 일부만 저장하고 나머지는 필요할 때 재계산하여 메모리 사용량을 줄입니다. 하지만 Gradient Checkpointing만으로는 여전히 메모리 제약이 발생할 수 있으며, 특히 모든 활성화 값을 CPU로 오프로드하는 방식은 CPU 메모리 및 대역폭 병목 현상을 유발할 수 있습니다.

이번 Axolotl PR은 이러한 문제를 해결하기 위해 decoder layer hidden state를 선택적으로 CPU로 오프로드하는 새로운 메커니즘을 도입했습니다. 기존 ALST(Activation Layer Saving Technique) 스타일의 reentrant checkpointing 방식과 유사하지만, torch.autograd.graph.saved_tensors_hooks를 활용하여 use_reentrant: false 설정에서도 동작하도록 개선되었습니다. 이로써 약 5%의 성능 향상과 함께 더욱 유연하고 효율적인 메모리 관리가 가능해졌습니다.

코드 분석: 무엇이 어떻게 개선되었나?

이 PR의 핵심은 saved_tensors_hooks를 사용하여 GradientCheckpointingLayer의 입력인 hidden_states만을 선택적으로 오프로드하는 새로운 CheckpointHiddenStatesOffload 클래스를 도입하고, 이를 Axolotl의 트레이닝 파이프라인에 통합한 것입니다. 파일별 주요 변경사항을 살펴보겠습니다.

1. src/axolotl/monkeypatch/checkpoint_activation_offload.py (새 파일)

이 파일은 새로운 hidden_states 오프로딩 로직의 심장부입니다. torch.autograd.graph.saved_tensors_hooks를 상속받는 CheckpointHiddenStatesOffload 클래스를 정의하여, PyTorch의 자동 미분(autograd) 시스템이 텐서를 저장하고 복원하는 과정에 개입합니다.

핵심 변경사항:

--- /dev/null
+++ b/src/axolotl/monkeypatch/checkpoint_activation_offload.py
@@ -0,0 +1,356 @@
+... (생략)
+
+def patch_gradient_checkpointing_layer_marker() -> None:
+    global _ORIG_GRADIENT_CHECKPOINTING_LAYER_CALL
+
+    if _ORIG_GRADIENT_CHECKPOINTING_LAYER_CALL is not None:
+        return
+
+    _ORIG_GRADIENT_CHECKPOINTING_LAYER_CALL = GradientCheckpointingLayer.__call__
+
+    def _checkpoint_offload_call(self, *args, **kwargs):
+        manager = _current_manager()
+        if (
+            manager is not None
+            and self.training
+            and torch.is_grad_enabled()
+            and getattr(self, "gradient_checkpointing", False)
+        ):
+            hidden_states = args[0] if args else kwargs.get("hidden_states")
+            if torch.is_tensor(hidden_states):
+                manager.mark(hidden_states)
+        return _ORIG_GRADIENT_CHECKPOINTING_LAYER_CALL(self, *args, **kwargs)
+
+    GradientCheckpointingLayer.__call__ = _checkpoint_offload_call
+
+class CheckpointHiddenStatesOffload(saved_tensors_hooks):
+    """Offload only marked checkpoint inputs.
+
+    The marker is installed on ``GradientCheckpointingLayer.__call__`` before
+    non-reentrant checkpointing saves positional layer inputs. All other saved
+    tensors pass through unchanged, including tensors from the final norm/head
+    and any non-checkpointed modules.
+    """
+    ...

patch_gradient_checkpointing_layer_marker 함수는 transformers.GradientCheckpointingLayer의 __call__ 메서드를 몽키패치합니다. 이 패치된 메서드는 hidden_states 텐서(레이어의 입력)를 _current_manager() (즉, CheckpointHiddenStatesOffload 인스턴스)에 mark하도록 지시합니다. 이렇게 mark된 텐서만이 CheckpointHiddenStatesOffload의 _retrieve_for_save 및 _retrieve_for_restore 훅을 통해 CPU로 오프로드되고 다시 GPU로 복원됩니다. 다른 모든 활성화 값은 GPU에 유지됩니다.

2. src/axolotl/core/builders/base.py

이 파일은 Axolotl의 트레이닝 인자(training arguments)를 구성하는 부분입니다. activation_offloading: hidden_states 모드가 선택되었을 때의 동작이 변경되었습니다.

Before:

        if self.cfg.activation_offloading == "hidden_states":
            # The checkpoint monkeypatch offloads the per-layer input, so HF
            # checkpointing stays on and must be reentrant for the patch to apply.
            training_args_kwargs["gradient_checkpointing"] = True
            gc_kwargs = dict(self.cfg.gradient_checkpointing_kwargs or {})
            gc_kwargs["use_reentrant"] = True # <-- 강제로 True로 설정
            training_args_kwargs["gradient_checkpointing_kwargs"] = gc_kwargs

After:

--- a/src/axolotl/core/builders/base.py
+++ b/src/axolotl/core/builders/base.py
@@ -571,12 +571,13 @@ def _configure_gradient_checkpointing(self, training_args_kwargs: dict):
         if self.cfg.layer_offloading:
             training_args_kwargs["layer_offloading"] = True
         if self.cfg.activation_offloading == "hidden_states":
-            # The checkpoint monkeypatch offloads the per-layer input, so HF
-            # checkpointing stays on and must be reentrant for the patch to apply.
             training_args_kwargs["gradient_checkpointing"] = True
             gc_kwargs = dict(self.cfg.gradient_checkpointing_kwargs or {})
-            gc_kwargs["use_reentrant"] = True
             training_args_kwargs["gradient_checkpointing_kwargs"] = gc_kwargs
+            if gc_kwargs["use_reentrant"] is False:
+                training_args_kwargs["activation_offloading"] = (
+                    self.cfg.activation_offloading
+                )
         elif self.cfg.activation_offloading:
             # TRL offloader replaces HF recompute (re-added for full finetune in the
             # model loader), so disable HF checkpointing and pass the mode through.

이 변경으로 activation_offloading: hidden_states가 gradient_checkpointing_kwargs.use_reentrant: false와 함께 사용될 수 있게 되었습니다. 기존에는 use_reentrant가 강제로 True로 설정되었지만, 이제는 사용자가 설정한 값을 존중하며, False일 경우 새로운 saved_tensors_hook 기반의 오프로딩 메커니즘이 활성화되도록 activation_offloading 인자를 트레이너로 전달합니다. (리뷰어 NanoCode012의 is False 사용 제안이 반영되었습니다.)

3. src/axolotl/core/trainers/mixins/activation_checkpointing.py

트레이너 믹스인에 새로운 activation_offload_context를 초기화하는 로직이 추가되었습니다.

Before:

            use_streams = self.args.activation_offloading != "legacy"
            if use_streams:
                _patch_trl_offload_current_stream()
            if isinstance(self.model, PeftModel):
                self.activation_offload_context = get_lora_act_offloading_ctx_manager(
                    self.model, use_streams=use_streams
                )

After:

--- a/src/axolotl/core/trainers/mixins/activation_checkpointing.py
+++ b/src/axolotl/core/trainers/mixins/activation_checkpointing.py
@@ -74,7 +74,17 @@ def __init__(self, *args, **kwargs):
             use_streams = self.args.activation_offloading != "legacy"
             if use_streams:
                 _patch_trl_offload_current_stream()
-            if isinstance(self.model, PeftModel):
+            if self.args.activation_offloading == "hidden_states":
+                from axolotl.monkeypatch.checkpoint_activation_offload import (
+                    get_checkpoint_hidden_states_offloading_ctx_manager,
+                )
+
+                self.activation_offload_context = (
+                    get_checkpoint_hidden_states_offloading_ctx_manager(
+                        use_streams=use_streams
+                    )
+                )
+            elif isinstance(self.model, PeftModel):
                 self.activation_offload_context = get_lora_act_offloading_ctx_manager(
                     self.model, use_streams=use_streams
                 )

activation_offloading이 `

참고 자료

• https://pytorch.org/docs/stable/autograd.html
• https://huggingface.co/docs/transformers/main_classes/model#transformers.GradientCheckpointingLayer
• https://pytorch.org/docs/stable/generated/torch.Tensor.pin_memory.html

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.