DoRA 확장: 인수 분해 노름 및 퓨즈드 커널을 통한 고순위 적응 기법

기존 DoRA 기법은 고성능 파인튜닝이 가능하지만 계산 과정에서 거대한 행렬을 생성해야 하므로 메모리 소모가 극심해 고순위(High-Rank) 설정이 어려웠다. 이 논문은 수학적 분해와 커널 최적화를 통해 메모리 사용량을 획기적으로 줄여 일반적인 GPU에서도 대규모 모델의 고성능 학습을 가능하게 한다.

DoRA는 가중치를 크기(magnitude)와 방향(direction)으로 나누어 학습하는 기법으로 LoRA보다 성능이 뛰어나지만 계산 과정에서 가중치 노름(Norm) 값을 구해야 한다. 기존 프레임워크는 이를 위해 저순위 행렬들을 곱해 원래 가중치 크기의 거대한 행렬을 복원하는데, 이 과정에서 엄청난 메모리가 낭비되어 고순위 학습이 불가능했다.

이 논문은 중학교 수학의 완전제곱식 전개와 유사하게 노름의 제곱을 세 가지 구성 요소로 분해하는 아이디어를 적용했다. 즉, 거대한 행렬을 먼저 만들지 않고도 기존 가중치, 저순위 행렬, 그리고 이들의 상호작용 항을 각각 계산한 뒤 합치는 방식으로 동일한 결과를 얻는다. 이는 메모리 사용량을 행렬 크기의 제곱에서 순위(rank)에 비례하는 수준으로 낮춘다.

또한 파이썬 수준에서 여러 번 호출되던 연산들을 GPU 전용 언어인 Triton을 이용해 하나의 통합 연산으로 묶었다. 이를 통해 데이터를 메모리에서 읽고 쓰는 횟수를 최소화하고 부동 소수점 연산에서 발생할 수 있는 오차를 방지하여 정확도와 속도를 동시에 확보했다.

인수 분해 노름(Factored Norm)은 가중치 노름의 제곱을 기저(base), 교차(cross), Gram 항으로 전개한다. [입력: 가중치 W, 저순위 행렬 A, B, 스케일 s] → [연산: 행렬 분해 법칙을 적용해 각 항을 청크 단위로 제곱합 및 내적 계산] → [출력: 행별 노름 값] → [의미: 거대 중간 행렬 생성 없이 메모리 효율적으로 노름 획득].

퓨즈드 트리톤 커널은 DoRA의 최종 업데이트 수식을 단일 패스로 처리한다. [입력: 기저 가중치, LoRA 출력, 크기 조절 인자 g] → [연산: 단일 CUDA 커널 내에서 요소별 연산 통합 및 fp32 정밀도 유지] → [출력: 최종 델타 가중치] → [의미: 메모리 읽기/쓰기 횟수를 4배 감소시켜 대역폭 병목 해결].

런타임 디스패치 시스템은 requires_grad 상태에 따라 실행 경로를 결정한다. 학습 시에는 역전파를 위한 중간 텐서 저장을 최적화하는 Tier 1 경로를, 추론 시에는 메모리 절약형 Tier 2 경로를 사용하며 사전에 정의된 임계값에 따라 최적의 커널을 자동 선택한다.

8B에서 32B 규모의 시각-언어 모델(VLM) 6종을 대상으로 실험한 결과, 기존 Hugging Face PEFT 구현 대비 추론 속도는 1.52.0배, 그래디언트 계산 속도는 1.51.9배 향상되었다. 피크 VRAM 사용량은 모델에 따라 최대 6.7GB까지 감소하여 메모리 효율성을 입증했다.

RTX 6000 PRO, H200, B200 등 다양한 NVIDIA GPU 아키텍처에서 일관된 성능 향상을 보였다. 특히 B200 GPU에서는 퓨즈드 커널을 통해 이론적 최대 메모리 대역폭의 약 53%를 활용하는 높은 효율을 기록했으며, 이는 기존 방식(20%) 대비 2.5배 이상 높은 수치이다.

수치적 정밀도 측면에서 최종 로짓의 코사인 유사도가 0.9999 이상을 유지했다. 2000 스텝 이상의 학습 실험에서도 기존 방식과 거의 동일한 손실(Loss) 변화를 보여 최적화로 인한 모델 성능 저하가 없음을 확인했다.

본 연구는 DoRA의 노름 계산 복잡도를 O(d_in^2)에서 O(d_out*r + r^2)로 낮추는 대수적 재구성을 핵심으로 한다. 기존 방식이 단위 행렬을 생성하여 전체 가중치를 복원하던 것과 달리, 저순위 행렬의 특성을 이용해 중간 생성물만을 계산함으로써 메모리 점유율을 획기적으로 줄였다.

수치적 안정성을 위해 (g-1) * base + g * s * lora 공식을 채택하여 g가 1에 가까울 때 발생하는 정밀도 손실을 방지했다. Triton 커널 내부에서는 모든 누적 연산을 fp32로 수행하며, 하드웨어별 최적 설정을 찾기 위해 자동 튜닝(Autotuning) 프로세스를 거친다.

분산 학습 환경인 DeepSpeed ZeRO-2/3 및 FSDP1을 지원하여 대규모 클러스터에서도 활용 가능하다. 다만 가중치 W 전체에 대한 행별 접근이 필요한 구조적 특성상 가중치를 샤딩하는 FSDP2/DTensor는 현재 지원 대상에서 제외되어 있다.

FSDP2 및 DTensor는 현재 지원되지 않으며 가중치가 샤딩된 환경에서는 추가적인 통신 오버헤드가 발생할 수 있다. 또한 현재의 디스패치 임계값은 실험적인 수치로 향후 하드웨어 변화에 따른 재조정이 필요할 수 있다.