Insights
AI가 자동으로 큐레이션·번역·정리하는 기술 동향 피드입니다.
Xiaomi-Lens-0 모델의 마스크 생성 과정에서 flash-attn이 적용되지 않는 대체 경로(mha)가 과도한 VRAM을 소모하는 문제가 발견되었습니다. 이 문제를 수정하여 추가 메모리 사용량을 18GB 수준에서 2.28GB로 크게 절감했습니다. 이는 AI 추론 시 옵션 의존성을 파악하고 성능 병목 지점을 측정하는 것이 중요함을 보여줍니다.
AI 자동 생성 콘텐츠
본 콘텐츠는 Qiita AI의 원문을 AI가 자동으로 요약·번역·분석한 것입니다. 원 저작권은 원저작자에게 있으며, 정확한 내용은 반드시 원문을 확인해 주세요.
원문 바로가기flash-attn이 적용되지 않는 환경에서, 자신의 모델에 대한 대체 경로가 보조 메모리(secondary memory)를 소모하고 있는지 판별하고, 이를 수정하거나 우회할지 결정할 수 있다.
nn.MultiheadAttention이 need_weights=False임에도 어텐션 행렬 전체를 생성하는 조건과 그 근본적인 정보
추론 중인 정방형 텐서를 찾아내는 약 20줄 분량의 탐지 함수
flash-attn 자체의 빌드 방법
Xiaomi-Lens-0의 정확도 평가 및 모델 구조 상세
전체 앱의 VRAM 최적화 (양자화 또는 오프로딩)
최근, 이미지 세그멘테이션 대응 4B 모델인 Xiaomi-Lens-0에 단 하나의 파일 수정 사항을 보내 다음 날 병합되었습니다. flash-attn이 없는 환경에서 마스크를 생성하는 단계의 추가 메모리가 3 메가픽셀(MP) 이미지 기준으로 18.02 GB에서 2.28 GB로 줄었습니다.
글을 쓰게 된 동기는 문제 발생 지점에 need_weights=False가 정확하게 기재되어 있었기 때문입니다. 공식 문서가 권장하는 설정이었기에, 저 역시 처음에는 지나쳤습니다. 만약 올바른 설정 상태에서도 18GB가 사라진다면, 같은 함정은 많은 코드에 숨어 있을 것입니다.
읽으신 분들께 생각해 보셨으면 하는 점은, 옵션 의존성이 결여되었을 때 자신의 모델이 어떤 경로를 거치는지 한 번이라도 측정해 본 경험이 있는지 여부입니다.
GPU 수치는 PR에 올린 RTX 5090 32 GB (소비자용 하이엔드 구성의 대표적인 예)에서 실측했으며, 메커니즘 재현은 CPU 2 코어・메모리 7 GB의 Linux 환경에서 진행했습니다.
본 기사에서는 하나의 비유를 끝까지 사용합니다. Attention 계산은 모든 참가자가 모든 참가자와 만나는 총대결입니다. 주인공인 Flash Attention은 대결 결과를 암산으로 흘려보내는 명인 같지만, 결석하는 경우가 있습니다. 그때 무대에 오르는 대체역이 전체 총대결표를 전부 종이에 인쇄하는 타입이었다는 이야기입니다.
서두에서 주인공이 결석하면 대체역이 무대에 선다는 구도를 설정했습니다. 이 대체역이 얼마나 많은 비용을 발생시키는지 수치로 파악하지 못한다면, 수정할 가치가 있는지 판단할 수 없습니다.
이 섹션에서 알 수 있는 것: 문제의 규모와 측정된 수치
Xiaomi-Lens-0은 Qwen3.5-4B 기반의 멀티모달 모델이며, 생성한 <|mask|> 토큰을 마스크로 변환합니다.
마스크를 만드는 디코더인 PatchSelfAttention 레이어는 Self-Attention 구현체로 flash / sdpa / mha 세 가지를 가지고 있으며, 동일한 가중치를 공유합니다. 수정 전에는 flash-attn이 없거나, GPU가 sm80보다 오래되었거나, CPU에서 구동하거나 등 어느 하나에 해당하면 mha가 자동으로 선택되었습니다. 이 자동 대체 경로를 폴백(fallback)이라고 부릅니다. sdpa는 환경 변수 XIAOMI_LENS_DECODER_ATTN=sdpa를 아는 사람만이 선택할 수 있었습니다.
PR 실측 결과입니다. 수치는 '모델 로드 이후, 마스크 생성 단계인 run_forward가 추가로 사용한 VRAM의 피크'이며, 앱 전체 사용량은 아닙니다.
| 입력 이미지 | 패치 수 P | 수정 전 mha | 수정 후 sdpa |
|---|---|---|---|
| 0.5 MP | 1,944 | 0.86 GB | 0.46 GB |
| ... | |||
환경은 Windows 11, PyTorch 2.14.1+cu130, Transformers 5.4.0, flash-attn 없음. scikit-image에 포함된 커피 사진을 각 크기로 축소하여 'the coffee cup'을 찾게 했습니다. 동일한 양을 측정하기 위한 골격입니다 (infer.py와 같은 순서로 호출합니다). |
def mask_pass_extra_gb(model, processor, mask_id, image, text, device):
inputs, _ = build_inputs(processor, image, text, device)
ids, _, n_masks = generate_completion(
...
8 MP (3840×2560)의 경우, mha 경로는 30.75 GiB를 한 번에 확보하려다 OOM(Out Of Memory)이 발생했고, sdpa 경로는 정상적으로 완료되어 요청 전체의 피크 메모리는 15.05 GB였으며 약 6초가 소요되었습니다. 이는 ComfyUI를 통한 별도의 테스트(PatchSelfAttention은 바이트 단위로 동일한 복사본임)와 측정 범위가 다르기 때문에 표와 비교하지 마십시오.
이전 섹션에서 이미지를 6배로 늘리면 mha 경로의 추가 메모리가 약 21배 증가한다는 것을 확인했습니다. 이렇게 증가하는 추세를 파악하면 어떤 해상도에서 시스템이 한계에 도달할지 사전에 예측할 수 있습니다.
본 섹션에서 알게 될 것: 자신의 GPU가 몇 MP까지 처리 가능한지 추정하는 방법
기호를 하나씩 도입하겠습니다.
$P$ (패치 수, 정수)는 전체 테스트 참가자 수를 의미합니다. 이미지를 16×16 화소로 나눈 구획의 개수로, 대략적으로 화소 수 ÷ 256입니다.
$H$ (헤드 수, 정수)는 전체 테스트 표의 권수를 의미합니다. 이 디코더에서는 16입니다.
$b$ (요소당 바이트 수, 정수)는 칸 하나에 들어가는 잉크 양입니다. bf16이면 2, fp32이면 4입니다.
M_{ ext{표}} = H \cdot P^2 \cdot b
이 공식은 참가자가 2배가 되면 인쇄해야 할 종이는 4배가 된다는 것을 의미합니다. 3 MP의 경우 $16 \times 11,616^2 \times 2 \approx 4.32$ GB가 표 한 묶음 분량이며, 실제로 측정된 18.02 GB는 대략 4묶음이 동시에 책상 위에 놓여 있는 계산입니다. 여기에는 세부 내역(스코어 행렬, softmax 출력, 타입 변환의 임시 영역 등)은 포함되어 있지 않습니다.
| 기호 | 비유적 역할 | 3 MP에서의 값 |
|---|---|---|
| $P$ | 참가자 수 | 11,616 |
| ... |
8 MP의 OOM 역시 같은 공식으로 설명할 수 있습니다. $P$가 약 3.2만 명이 되면 표 한 묶음만으로도 30.75 GiB에 달하며, 32 GB GPU에는 모델 본체와 함께 올라갈 수 없습니다. PR의 네 지점에 mha는 $aP^2 + bP$, sdpa는 $cP + d$를 최소제곱법(least squares)을 이용해 추정하여 예측했습니다.
import numpy as np
P = np.array([1944, 3952, 7776, 11616], dtype=float)
mha_gb = np.array([0.86, 2.58, 8.53, 18.02])
...
| 이미지 | P | mha 예측 | sdpa 예측 |
|---|---|---|---|
| 3 MP | 11,719 | 18.3 GB | 2.30 GB |
| ... |
4 MP 이상은 측정 범위를 벗어난 외삽(extrapolation)입니다. 8 MP의 mha는 실제로는 처음 30.75 GiB를 확보하는 과정에서 OOM이 발생하므로 120 GB에는 도달하지 않습니다.
sdpa의 8 MP 예측에 모델 본체 (bf16으로 약 8 GB)를 더하면 약 14 GB가 되어 전체 피크 메모리 15.05 GB와 대략적으로 일치합니다. 따라서 32 GB GPU라면, mha 경로는 실용적인 상한선이 3 MP 전후이며, sdpa 경로는 상한인 8.4 MP까지 추가로 약 6 GB 정도의 여유가 있다고 판단할 수 있습니다.
이전 섹션에서 대리(代役)는 전체 테스트 표를 통째로 인쇄하고 있다는 것을 알았습니다. 하지만 문제의 원인이 되는 need_weights=False 설정은 공식 문서상으로는 인쇄하지 않는 경로에 올라야 할 설정입니다. 이 부분을 해결하지 않으면 'MultiheadAttention이 나쁘다'라는 잘못된 결론을 내리게 됩니다.
본 섹션에서 알게 될 것: nn.MultiheadAttention이 어텐션 행렬을 생성하는 정확한 조건
수정 전의 mha 경로는 self.attn(normed, normed, normed, need_weights=False, key_padding_mask=...)라고 불렸습니다. PyTorch 2.14 문서에는 다음 문장이 있습니다.
Set need_weights=False to use the optimized scaled_dot_product_attention
해석: 최적화된 scaled_dot_product_attention을 사용하려면 need_weights=False를 설정하십시오.
같은 페이지에는 추론용 고속 경로(fast path)에 진입하는 조건도 나열되어 있습니다. 요약하자면, 자기 어텐션(self-attention), batch_first의 3차원 입력, autograd 비활성화, 그리고 .eval() 호출입니다.
이와 같은 방식입니다. infer.py는 model.eval()과 torch.inference_mode()를 사용하여 추론을 수행하므로, 모든 요구 사항을 충족합니다.
문제는 고속 경로(fast path)의 내부 구현에 있습니다. CUDA 구현인 native_multi_head_attention_cuda는 SDPA에 전달할 조건 중 마스크 유무를 확인하고 있습니다.
if (!mask.has_value() && no_seq_len_1_nested &&
마스크가 전달되면 SDPA로 진행하지 않고, bmm_nt를 사용하여 스코어 행렬을 만들고 masked_softmax를 적용하는 경로로 떨어집니다. 주석에는 형태가 [B, num_head, T, T]라고 적혀 있습니다. 이 전체 메모리에 배치하는 움직임을 실체화(materialization)라고 부릅니다. 그리고 Xiaomi-Lens-0의 디코더는 패딩을 나타내는 key_padding_mask를 항상 전달하고 있었습니다.
CPU 상에서 4가지 조건을 각각 한 번씩 실행하여 호출된 aten 연산을 프로파일러로 포착했습니다.
import torch
import torch.nn as nn
from torch.profiler import ProfilerActivity, profile
...
| 모드 | 마스크 | 주로 호출된 연산 |
|---|---|---|
| eval | 없음 | _native_multi_head_attention, bmm, _softmax |
| eval | 있음 | _native_multi_head_attention, bmm, _masked_softmax |
| train | 없음 | scaled_dot_product_attention |
| train | 있음 | scaled_dot_product_attention |
eval 모드, 마스크 있음의 _masked_softmax에는 [1, 16, 512, 512]가 전달되었습니다. CPU는 마스크가 없어도 행렬을 생성한다는 점에서 CUDA와 다르지만, 마스크가 있는 추론에서 표를 출력하는 점은 공통적입니다.
가장 놀라웠던 것은 학습 모드에서는 SDPA를 사용하고 있다는 점입니다. 같은 레이어를 동일한 인수로 호출해도, 표를 출력하는 것은 추론 시에만 발생합니다. 학습 코드나 그래디언트를 활성화한 단일 테스트에서는 발견되지 않고, 추론 본방에서만 모습을 드러내는 함정이었습니다.
앞선 절에서 마스크가 있는 추론에서는 고속 경로가 스코어 행렬을 실체화한다는 것을 확인했습니다. 경로를 파악했더라도 그것이 메모리의 주범인지는 양을 측정하지 않으면 모른다고 할 수 없습니다.
이 섹션에서 알 수 있는 것: mha와 sdpa의 피크 메모리가 P에 따라 어떻게 증가하는가
forward를 1회씩 별도의 프로세스로 실행하고, /proc/self/clear_refs로 피크 상주 메모리(VmHWM)를 리셋한 후 증분을 측정했습니다.
import sys
import torch
import torch.nn as nn
...
| P | 표 1 묶음의 이론값 | mha 고속 경로 | sdpa |
|---|---|---|
| 1,024 | 64 MiB | 172 MiB | 44 MiB |
| ... |
이론값은 $H P^2 b$ (fp32에서 $b = 4$)이며, PyTorch 2.14.1에서는 두 경우 모두 같은 값이었습니다. mha는 P가 4배가 되면 약 13배로 증가하며, 이론값의 약 2배(스코어 행렬과 softmax 출력의 2묶음)에 근접합니다. sdpa는 P가 1,024 증가할 때마다 20 MiB의 직선적인 증가를 보입니다. 이는 표를 출력하지 않고 블록별로 암산하기 때문입니다.
앞선 절에서 sdpa라면 전체 표를 출력하지 않는다는 것을 양으로 확인했습니다. 남은 질문은 어떻게 최소한으로 안전하게 수정하느냐입니다.
이 섹션에서 알 수 있는 것: 수정안의 비교와 실제 차분
| 案 | 내용 | 판단 |
|---|---|---|
| A | mha의 호출 이름을 변경하여 고속 경로를 우회함 | 그래디언트(gradient) 유무에 의존하는 트릭이라, 버전에 따라 동작이 달라질 수 있음 |
| ... | ||
| C를 선택했습니다. 대기실에는 원래 인쇄하지 않는 대체재가 기다리고 있었습니다. 환경 변수를 모르는 사람도 자동으로 작동하도록 만든 것뿐입니다. 차이점은 다음과 같습니다 (문서 문자열은 생략하고, 시그니처는 1줄로 요약했습니다). |
def _self_attention(self, normed, key_padding_mask):
- if _DECODER_ATTN == "sdpa":
- return self._sdpa_self_attention(normed, key_padding_mask)
...
경고 메시지도 “falling back to SDPA”로 수정했습니다. XIAOMI_LENS_DECODER_ATTN=mha를 통해 이전 동작도 선택할 수 있습니다.
이전 절에서 대체재를 mha에서 sdpa로 교체했습니다. 대체재가 바뀌면 성능도 바뀝니다. 메모리가 줄어도 출력이 망가지면 의미가 없습니다.
이 섹션에서 알게 되는 것: 수정 전후에 무엇이 동일하고, 무엇이 미묘하게 다른지
PR에서는 6가지 케이스(고양이, 커피잔, 동전 23개, 로켓, 깃발, 추론을 수반하는 요청)를 사용하여 생성 토큰 ID, 박스, 모든 마스크 로짓을 비교했습니다.
| 비교 | 결과 |
|---|---|
| 수정 후 기본값 vs 수정 전 sdpa 명시 | 6/6에서 비트 단위 일치 |
| ... | |
| 위 2줄은 배선 확인이고, 아래 2줄은 경로 차이로 인한 차이입니다. 마스크의 차이는 경계에 미묘하게 나타날 정도였고, PR에서는 반올림으로 인한 차이라고 판단했습니다. 6가지 케이스가 포괄적인 품질 평가는 아니지만, 차이가 발생하는 위치와 크기는 확인할 수 있습니다. |
mha에서 SDPA로 이식할 때 부울(boolean) 마스크의 의미가 반대입니다. key_padding_mask는 True를 무시하고, SDPA의 attn_mask는 True를 포함합니다.
import torch
import torch.nn.functional as F
torch.manual_seed(0)
...
출력은 “반전하여 전달한 0.0”과 “그대로 전달한 3.46”이었습니다. Xiaomi-Lens-0는 -inf를 더하는 가산 마스크로 변환하여 전달하기 때문에, 이 함정에 빠지지 않았습니다.
이전 절에서 대체재 교체가 출력을 거의 바꾸지 않음을 확인했습니다. 이 함정은 Xiaomi-Lens-0에만 국한된 것이 아니며, nn.MultiheadAttention을 추론(inference) 시 마스크와 함께 호출하는 모든 코드에 숨어 있습니다.
이 섹션에서 알게 되는 것: 자신의 코드가 인쇄되는 대체재를 호출하고 있는지 판별하는 방법
추론과 동일한 inference_mode로 한 번 실행하여, 마지막 2차원이 같은 큰 텐서를 포착하는 함수입니다.
import torch
from torch.profiler import ProfilerActivity, profile
WATCH = ("softmax", "bmm", "baddbmm", "matmul")
...
| 대상 | 감지 결과 |
|---|---|
| MultiheadAttention | _masked_softmax (1, 16, 1024, 1024) 등 |
| TransformerEncoderLayer | _masked_softmax (1, 16, 1024, 1024) 등 |
| SDPA 경로 | 없음 |
표준 Transformer 레이어에서도 같은 형태가 나왔습니다. 대체재의 습관은 하나의 모델이 아니라 부품 자체에 있습니다.
이 함수를 만드는 과정에서 한 번 실패했습니다. .eval()을 호출하는 것을 잊은 레이어에서 테스트하여 아무것도 감지되지 않자 “문제없다”고 착각했습니다. 학습 모드에서는 고속 경로로 진입하지 않기 때문입니다. 검사는 반드시 실제와 동일한 eval 및 inference_mode 조합으로 수행하는 것이 이 함정의 가장 큰 교훈이었습니다.
이전 절까지로, 원인 특정, 수정, 출력 확인, 전파 도구들이 갖춰졌습니다. 발견한 수정 사항은 유지보자가 받을 수 있는 형태로 만들지 않으면 다른 사용자에게 전달되지 않습니다.
이 섹션에서 알게 되는 것: 작은 수정을 단시간에 머지(merge)하게 만드는 PR 구성 방법
PR 본문은 '문제 / 변경 / 측정 / 출력 비교'의 4단계로 구성했습니다. 측정 부분에서는 측정한 범위(추가적인 피크일 뿐 전체는 아님)를 명확히 기재했으며, 출력 비교에서는 배선 확인(비트 단위 일치)과 경로 차이로 인한 수치적 차이(IoU)를 나누어 작성했습니다.
10월 19일에 제출했고 다음 날인 20일에 메인테이너가 병합했습니다. 변경 사항은 파일 1개, 내용 추가 10줄 및 삭제 5줄입니다. 차이의 크기가 작았기 때문에, 그 차이가 올바르다는 것을 단시간에 판단할 수 있는 자료를 첨부한 것이 효과적이었다고 생각합니다.
| 증상 | 원인 | 대처 |
|---|---|---|
| 큰 이미지에서만 OOM 발생 | 주의 행렬이 $P^2$로 증가함 | 제7절의 함수에서 _masked_softmax를 찾기 |
need_weights=False인데 무거움 | 추론의 고속 경로가 마스크 처리되어 SDPA를 사용하지 않음 | SDPA로 대체하기 |
| 탐지 함수가 아무것도 반환하지 않음 | 레이어가 학습 모드 상태로 남아 있음 | .eval()을 호출한 후 검사하기 |
| SDPA로 대체했더니 출력이 깨짐 | 참/거짓 값 마스크의 의미가 반대임 | ~mask로 반전시키거나 가산 마스크로 처리하기 |
| 대체 후 추론 결과가 매번 흔들림 | dropout_p가 추론 시에도 유효함 | 추론 시에는 dropout_p=0.0을 전달하기 |
| 용어 | 의미 | 예시 대응 |
|---|---|---|
| 폴백(Fallback) | 본래 구현이 사용 불가능할 때의 대체 경로 | 주인공이 결석했을 때의 대역 |
| ... |
주인공이 결석하는 환경은 Windows, 구형 GPU, CPU 등 생각보다 넓게 존재합니다. 옵션 의존성을 제거한 상태에서 한 번만 추론을 실행하여 대역 경로를 측정해 보세요.
PyTorch 고속 경로 심층 분석을 중심으로 한 버전은 Zenn에도 작성했습니다.