Gemma-4 12B (인코더 프리 멀티모달 모델)을 AWS Inferentia2로 포팅하기
요약
Gemma-4 12B (인코더 프리) 모델을 AWS Inferentia2 하드웨어에 포팅하는 과정을 설명합니다. 이 과정은 멀티모달 아키텍처의 복잡성에도 불구하고, 인코더를 제외하고 핵심 텍스트 경로만 추출하여 성공적으로 컴파일했습니다. 특히 글로벌 레이어 샤딩 문제와 SBUF 초과 문제를 해결하기 위해 'eager 어텐션'을 강제하는 세 가지 주요 수정 사항이 필요했음을 강조합니다.
핵심 포인트
- 인코더 프리 멀티모달 모델 구조를 이해하고 텍스트 경로만 추출하여 포팅 가능.
- 글로벌 레이어 샤딩 시, 쿼리 헤드와 단일 KV 헤드의 그룹화 처리가 중요함.
- SBUF 초과 문제 해결을 위해 'eager 어텐션' 구현을 강제하는 것이 효과적임.
시리즈의 중간 아이입니다. 12B는 MatFormer "효율적" 모델(E2B/E4B)도 아니고 거대 모델(31B/26B)도 아닌 최초의 Gemma-4입니다. 하지만 이 모델은 두 가지 놀라운 점을 가지고 있습니다. 첫째, 인코더가 없는
**멀티모달* 모델로 패키징되어 있고, 둘째, 어텐션이 작은 모델들이 보여주지 않았던 방식으로 Neuron 하드웨어 버퍼를 오버플로우합니다. 이 포팅은 E4B 레시피에 단 세 가지 수정만 추가하면 되는 짧고 깔끔한 작업입니다.*
| Model | google/gemma-4-12B-it — dense, model_type: gemma4_unified |
| ... |
놀라운 점 #1: 인코더가 없는 멀티모달 클래스
google/gemma-4-12B-it은 E 계열처럼 Gemma4ForConditionalGeneration으로 로드되지 않고, Gemma4UnifiedForConditionalGeneration, model_type: gemma4_unified입니다. "Unified"는 이것이 멀티모달 아키텍처라는 것을 의미합니다... 다만 이 체크포인트에는 오디오 또는 비전 인코더가 없습니다. 즉, 인코더 프리 멀티모달 모델입니다: 클래스 내의 모든 멀티모달 _배관(plumbing)_은 있지만, 인코더는 전혀 없습니다.
Inferentia 텍스트 포팅에 있어 실제 좋은 소식은 — 다른 크기들과 마찬가지로 **model.language_model**만 가져오고 나머지는 무시하면 됩니다. 하지만 통합된 Wrapper를 지나쳐야 한다는 것을 알아야 하며, 설정의 비전/오디오 필드는 존재하지만 작동하지 않는(inert) 상태이므로, 텍스트 경로를 인스턴스화하기 전에
- Per-Layer 임베딩 없음.
hidden_size_per_layer_input = 0으로 설정하고, plaininputs_embeds를 공급하며 PLE 경로를 삭제합니다 (31B 모델에서 사용된 것과 같은 단순화). - Cross-layer KV 공유 없음. 각 레이어가 자체적인 KV를 가집니다.
- 혼합 슬라이딩 + 글로벌 어텐션,
sliding_window = 1024, GQA (쿼리 16개 / 키-값 8개), softcap 30을 사용합니다.
이후 컴파일 사이클을 거쳐 찾아낸 세 가지 수정 사항이 있습니다.
수정 #1: 글로벌 레이어의 nkv=1 샤딩 문제
글로벌 (full_attention) 레이어는 1개의 KV 헤드(num_key_value_heads- ext{per-global} = 1)를 가집니다. TP=2일 때 이를 샤딩할 수 있지만, 이때 쿼리 헤드만 샤딩하고 단일 KV 헤드는 올바르게 그룹화해야 합니다 (groups = nq // TP). 그룹화가 잘못되면 글로벌 레이어가 잘못된 키에 조용히 어텐션하게 됩니다. (이것은 31B 모델에서 발생했던
이것이 실제로 막히는 부분입니다. Neuron 컴파일러는 속도를 위해 어텐션(flash-style)을 융합(fuse)하려고 합니다. sliding_window = 1024에서, 이 융합된 커널의 작업 집합(working set)이 SBUF(온칩 스크래치 버퍼)를 초과하여 컴파일에 실패합니다. 해결책은 간단합니다: eager 어텐션 구현(attn_implementation="eager")을 강제하여 어텐션을 일반적인 matmul로 낮추는 것입니다. 융합된 커널의 속도는 포기하지만, 이는 컴파일되고 올바르게 실행되므로, 존재하지 않는 빠른 그래프보다 낫습니다. (더 작은 창 크기의 E-family 모델들은 이 문제를 겪지 않습니다. 이것은 구체적으로 1024 창 레이어에서 버퍼를 초과하는 경우입니다.)
결과(Result)
이 세 가지 설정으로, 12B 모델은 컴파일되고 장치 그리디 디코드 시 토큰별로 CPU fp32 참조 값과 동일합니다. 즉, SEQ_MATCH True, `
AI 자동 생성 콘텐츠
본 콘텐츠는 Dev.to AI tag의 원문을 AI가 자동으로 요약·번역·분석한 것입니다. 원 저작권은 원저작자에게 있으며, 정확한 내용은 반드시 원문을 확인해 주세요.
원문 바로가기