[axolotl] Qwen3.5 성능 극대화: PEFT의 불필요한 Dtype 캐스팅 제거와 LoRA 커널 퓨전

PR 링크: axolotl-ai-cloud/axolotl#3732 상태: Merged | 변경: +784 / -6

들어가며

최근 Qwen3.5와 같은 최신 모델들은 성능 향상을 위해 GatedDeltaNet(Linear Attention) 구조를 적극적으로 도입하고 있습니다. 하지만 이러한 새로운 아키텍처를 기존의 PEFT(Parameter-Efficient Fine-Tuning) 라이브러리로 학습할 때, 예상치 못한 성능 병목이 발생하곤 합니다.

이번 글에서는 Axolotl 라이브러리에 적용된 최적화 PR을 통해, PEFT의 기본 동작 방식이 가진 오버헤드를 어떻게 해결하고, fused LoRA kernel을 통해 학습 속도를 어떻게 5% 이상 끌어올렸는지 심층 분석해 보겠습니다. 이 PR의 핵심은 불필요한 bf16 -> fp32 -> bf16 데이터 타입 변환(Round-trip)을 제거하고, 여러 프로젝션을 하나의 커널로 묶어 메모리 대역폭 효율을 극대화한 것입니다.

문제 상황: PEFT의 Dtype 캐스팅 세금

기본적으로 Hugging Face의 PEFT 라이브러리는 수치적 안정성을 위해 LoRA forward 시 _cast_input_dtype을 호출합니다. 이 함수는 입력 텐서([Batch, Seq, Hidden]) 전체를 bf16에서 fp32로 변환했다가 다시 bf16으로 되돌리는 작업을 수행합니다.

문제는 이 작업이 Forward, Gradient Checkpointing Recompute, Backward 과정마다 반복된다는 점입니다. 모델의 크기가 커지고 시퀀스 길이가 길어질수록, 이 '수치적으로는 아무 의미 없는' 복사 트래픽은 수백 MB의 메모리 대역폭을 잡아먹으며 성능을 저하시킵니다.

코드 분석: Fused LoRA 커널로의 전환

1. 개별 프로젝션 최적화 (apply_lora_linear)

기존에는 PEFT의 forward를 그대로 호출했다면, 이제는 Axolotl이 직접 구현한 LoRA_O 커널을 통해 불필요한 캐스팅 없이 연산을 수행합니다.

Before (PEFT Default):

# PEFT 내부적으로 발생하는 동작 (개념적)
def forward(self, x):
    x = x.to(torch.float32) # 불필요한 캐스팅
    result = self.base_layer(x.to(self.base_layer.weight.dtype)) + self.lora_A(x) * scaling
    return result.to(original_dtype)

After (Axolotl Optimized):

def apply_lora_linear(proj: nn.Module, X: torch.Tensor) -> torch.Tensor:
    """Routed through ``LoRA_O`` to avoid peft's bf16->fp32->bf16 dtype round-trip."""
    W, b, W_quant, A, B, s, lora_bias, dropout, magnitude = get_lora_parameters(proj)
    X_drop = _apply_dropout(dropout, X, proj.training)
    # 입력 X의 dtype을 유지한 채 커널 내부에서 필요한 부분만 캐스팅
    return LoRA_O.apply(X, X_drop, W, b, W_quant, A, B, s, lora_bias, magnitude)

2. GatedDeltaNet In-Projection 퓨전 (LoRA_FusedProj)

GatedDeltaNet의 in_proj 그룹(qkv, z, b, a)은 동일한 입력 X를 공유합니다. 이 PR은 이 4개의 프로젝션을 하나의 torch.autograd.Function으로 묶어(fused) 처리합니다.

핵심 변경 사항 (src/axolotl/kernels/lora.py):

class LoRA_FusedProj(torch.autograd.Function):
    @staticmethod
    @torch_amp_custom_fwd
    def forward(ctx, scales, quants, X, X_drop, *tensors):
        # ... 생략 ...
        for i, (W, b, A, B, lora_bias, magnitude) in enumerate(
            _chunk(tensors, _FUSED_PROJ_NUM_TENSORS)
        ):
            # 여러 프로젝션을 순회하며 연산하지만, 입력 X는 공유됨
            out = matmul_lora(X, W, b, quant, A, B, s, X_drop=X_drop, lora_bias=lora_bias)
            outs.append(out)
        # ...
        return tuple(outs)

이 구조를 통해 입력 X에 대한 메모리 읽기 횟수를 줄이고, Backward 시에도 공유된 X_lora_t를 재사용하여 연산 효율을 높였습니다.

3. 동적 패칭 로직 (src/axolotl/monkeypatch/lora_kernels.py)

리뷰어들의 피드백을 반영하여, in_proj는 lora_qkv_kernel 설정에, out_proj는 lora_o_kernel 설정에 따라 독립적으로 활성화되도록 구현되었습니다.

if cfg.lora_qkv_kernel and any(
    hasattr(getattr(linear_attn, name), "lora_A")
    for name in LINEAR_ATTN_IN_PROJS
):
    linear_attn.apply_in_proj_fused = types.MethodType(
        _make_apply_lora_gdn_in_proj(LINEAR_ATTN_IN_PROJS), linear_attn
    )

왜 이게 좋은가? (성능 및 교훈)

성능 수치

RTX 5090에서 시퀀스 길이에 따른 마이크로 벤치마크 결과는 놀랍습니다.

• Sequence 32768: PEFT(99.33ms) vs Fused(83.53ms) → 1.19배 속도 향상
• 메모리 절약: 레이어당 최대 0.25 GiB 절감

전체 모델 학습(Qwen3.5-9B)에서도 throughput이 5.1% 향상되는 결과를 보여주었습니다. 이는 단순한 연산 최적화가 아니라, 메모리 대역폭 병목을 해결했을 때 얻을 수 있는 전형적인 성능 그래프입니다.

일반적인 교훈

1. Activation-scaling Copy Traffic을 경계하라: 모델 파라미터보다 활성화 함수(Activation)의 크기가 훨씬 큰 긴 시퀀스 학습에서는, 데이터 타입 변환이나 단순 복사 작업이 연산(Compute)보다 더 큰 비용을 초래합니다.
2. 커널 퓨전의 힘: 동일한 입력을 사용하는 여러 연산을 하나로 묶으면 메모리 IO를 획기적으로 줄일 수 있습니다.
3. 수치적 동일성 검증: 이 PR은 최적화 과정에서 bit-identical 혹은 cos 0.99999 수준의 오차를 유지함을 입증했습니다. 성능을 위해 정확도를 포기하지 않았다는 점이 시니어 엔지니어링의 정수입니다.

마치며

이번 최적화는 Qwen3.5와 같은 최신 아키텍처를 효율적으로 학습시키기 위한 필수적인 단계였습니다. PEFT의 유연함과 커스텀 커널의 성능 사이에서 적절한 균형을 찾은 훌륭한 사례입니다. Axolotl 사용자라면 lora_qkv_kernel과 lora_o_kernel 옵션을 켜는 것만으로도 이 혜택을 누릴 수 있습니다.

참고 자료

• https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function
• https://huggingface.co/docs/peft/index
• https://pytorch.org/docs/stable/notes/amp_examples.html#custom-autograd-function

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