[Axolotl] 가중치 동기 로딩으로 OOM 방지

PR 링크: axolotl-ai-cloud/axolotl#3477 상태: Merged | 변경: +10 / -0

들어가며

Mixture of Experts(MoE) 모델은 전문가(expert) 레이어가 수십~수백 개에 달한다. Axolotl은 이런 MoE 모델을 로딩할 때 on-load quantization을 적용하여 VRAM을 절약한다. bf16 가중치를 GPU에 올리는 즉시 4비트로 양자화하고 원본을 해제하는 방식이다.

문제는 Transformers 라이브러리의 convert_and_load_state_dict_in_model이 ThreadPoolExecutor를 사용해 텐서를 비동기로 GPU에 미리 전송한다는 것이다. 양자화 패치가 한 expert를 처리하는 동안, 백그라운드 스레드가 다른 expert들의 bf16 텐서를 GPU에 계속 올려놓는다. 전체 가중치의 5%도 로딩되지 않았는데 OOM이 발생하는 상황이 벌어진다.

핵심 코드 분석

비동기 로딩 비활성화

Before (moe_quant.py):

def patch_moe_quantization_on_load(cfg):
    _moe_load_state["quant_type"] = quant_type
    _moe_load_state["compress_statistics"] = compress_statistics

    # Disable caching_allocator_warmup ...

After:

def patch_moe_quantization_on_load(cfg):
    _moe_load_state["quant_type"] = quant_type
    _moe_load_state["compress_statistics"] = compress_statistics

    # Disable async tensor loading.  Transformers' convert_and_load_state_dict_in_model
    # uses a ThreadPoolExecutor to materialise tensors (move from safetensors → CUDA)
    # ahead of time.  With MoE models this pre-fetches many large bf16 expert tensors
    # onto the GPU simultaneously — long before our set_param_for_module patch can
    # quantise and free them one-by-one — causing OOM even at <5 % of weights loaded.
    # Sequential loading ensures only ONE bf16 expert tensor is on-GPU at a time.
    os.environ["HF_DEACTIVATE_ASYNC_LOAD"] = "1"

    # Disable caching_allocator_warmup ...

핵심은 os.environ["HF_DEACTIVATE_ASYNC_LOAD"] = "1" 한 줄이다. Hugging Face Transformers는 이 환경변수를 확인하여 비동기 텐서 전송을 건너뛴다.

왜 이게 좋은가

이 수정의 핵심을 메모리 타임라인으로 이해하면 명확하다:

비동기 로딩 (수정 전):

시간 →
GPU: [expert_1_bf16] [expert_2_bf16] [expert_3_bf16] [expert_4_bf16] ... OOM!
                      ↑ 백그라운드 pre-fetch
양자화:  [expert_1 → 4bit] ← 아직 1번만 처리 중

동기 로딩 (수정 후):

시간 →
GPU: [expert_1_bf16] → [expert_1_4bit] → [expert_2_bf16] → [expert_2_4bit] → ...
     ↑ 로드           ↑ 양자화+해제      ↑ 로드           ↑ 양자화+해제

동기 로딩에서는 한 번에 하나의 bf16 expert만 GPU에 존재하므로, 피크 VRAM이 단일 expert의 bf16 크기 + 양자화된 모든 expert로 제한된다.

• OOM 해결: Mixtral-8x7B 같은 대형 MoE 모델도 단일 GPU에서 양자화 로딩이 가능해진다
• 기존 기능 보존: 환경변수 방식이므로 non-MoE 모델의 로딩 속도에는 영향 없다
• 최소 침습: 코드 10줄 추가로 근본 원인을 해결했다

정리

비동기 최적화가 다른 최적화(on-load quantization)와 충돌하는 전형적인 사례다. Transformers의 비동기 로딩은 일반 모델에는 유리하지만, "로딩 즉시 양자화하여 메모리를 회수"하는 패턴과는 근본적으로 상충한다. 해결 방법이 환경변수 한 줄인 것이 인상적이다. 프레임워크 간 상호작용에서 발생하는 OOM 문제를 디버깅할 때, 비동기 동작을 의심해보라는 교훈을 준다.

참고 자료

• Hugging Face Accelerate: Large Model Loading — 대형 모델 로딩 메커니즘
• bitsandbytes Quantization — on-load quantization 구현체

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