지난 포스팅에서 저는 Gemma-4의 128-expert MoE (Mixture of Experts)를 inf2.24xlarge에서 실행하는 과정을 다루며, "2코어 박스에 맞추려면 fp4가 필요하며, 이는 별도의 탐험이 될 것"이라는 절벽 엔딩(cliffhanger)으로 마무리했습니다. 바로 그 탐험이 이번 글입니다.

결과는 제가 예상했던 곳과는 전혀 다른 곳에서 끝났습니다. fp4를 사용하지도 않았고, 목표로 했던 8xlarge에서 실행하지도 않았습니다. 대신 AWS에서 판매하는 가장 작고 저렴한 Inferentia2 인스턴스인 **단일 inf2.xlarge (16 GB의 호스트 RAM 탑재)**에서 26B-parameter 모델을 실행하는 데 성공했습니다. 막다른 길을 포함한 그 개선의 과정을 소개합니다.

완성된 포팅 버전(이전 기사 참조)은 24xlarge에서 26B-A4B를 서비스했습니다: 12개의 NeuronCores, 192 GB HBM, 시간당 약 $6.49(on-demand). 작동했고, 정확했으며, 한 번에 하나의 질문을 던지는 한 명의 사용자를 감당하기에는 과잉(overkill)이었습니다. "4B-active" MoE의 핵심은 실행 비용이 저렴해야 한다는 점입니다. 이를 호스팅하기 위해 데이터 센터급 박스가 필요해서는 안 됩니다. 목표는 시간당 약 $0.76인 최하위 계층 inf2.xlarge — 2개 코어, 32 GB HBM, 그리고 단 16 GB의 호스트 RAM — 에 모델을 올리는 것이었습니다. 이는 8.6배 더 저렴합니다.

두 가지 장벽이 가로막고 있었습니다: 모델이 bf16 형식으로는 32 GB HBM에 들어가지 않으며, 일반적인 방식으로는 16 GB 호스트 RAM에 들어가지 않는다는 점입니다. 저는 이 두 가지를 모두 깨뜨려야 했습니다.

장벽 1: 메모리 계산, 그리고 왜 "A4B"가 당신을 구원하지 못하는가

Expert(전문가)들이 전체 가중치의 약 93%를 차지합니다. bf16에서 128개의 expert는 약 45.6 GB입니다. 이를 TP=2(Tensor Parallelism)로 복제하면 16 GB 예산 대비 코어당 약 22.8 GB가 됩니다. Top-8 라우팅(routing)은 연산량(compute)을 줄여줄 뿐, 메모리 점유량(footprint)을 줄여주지는 않습니다. 128개의 expert가 모두 상주해야 하기 때문입니다. 따라서 expert만으로도 이미 인스턴스 용량을 초과합니다.

저의 첫 번째 시도는 당연하게도 **expert를 int8로 양자화(quantization)**하는 것이었습니다. NxD에는 QuantizedColumnParallel / QuantizedRowParallel이 있으며, (이전 기사에서 언급했듯이) 이 모델에 대한 per-channel int8 적용은 fp32와 토큰 단위로 동일할 만큼 수치적으로 완벽합니다. 이 경우 expert는 rank당 약 11.4 GB로 줄어듭니다. 이 정도면 들어갈 수 있을 것입니다.

하지만 되지 않았습니다. Could not load the model status=4 Allocation Failure — 16 GB 코어 용량을 약 3~4 GB 초과했습니다. 그리고 inf2에는 그 차이를 메울 수 있는 4-코어 SKU가 없습니다. 당시 제가 적어두었던 결론은 다음과 같습니다: 마지막 3~4 GB를 채우려면 fp4 전문가(experts)가 필요하다.

fp4라는 미궁 (두 번의 시도, 둘 다 막다른 길)

저는 fp4를 두 번 쫓았습니다.

첫 번째, QuantizedColumnParallel은 F4E2M1FN_X4를 광고합니다. 하지만 그 state_dict 가중치는 from_float/preshard_hook에 의해 생성된 packed uint16 마이크로스케일링 (microscaling) 형식(32 fp4 → 8 uint16)입니다. 이는 int8의 5줄짜리 양자화기(quantizer)와는 전혀 다릅니다. 그 후 AWS 문서를 명확히 읽었습니다: NxD Inference 양자화는 INT8과 FP8만 지원합니다. FP8은 8비트이므로 int8과 크기가 같아 도움이 되지 않습니다. QuantizedDtype.F4E2M1FN_X4는 상수로서 존재하지만, 연결된 추론 경로(inference path)는 아닙니다.

두 번째, 최신 Neuron 2.31 스택(neuronx-cc 2.26, NxD 0.19)에서는 fp4가 덜 부재하는 것처럼 보였습니다 — BLOCKWISE_SYMMETRIC, block_size, from_float가 있었습니다. 그래서 다시 시도했지만, 두 번 모두 바닥을 쳤습니다:

• 깔끔한 QuantizedColumnParallel fp4는 blockwise_scale_dequantize를 호출하는데, 이는 CPU 전용 torch 레퍼런스(assert device == cpu)입니다 — NeuronCore로 트레이스(trace)될 수 없습니다.
• 실제 디바이스 경로는 blockwise_mm NKI 커널을 요구하며, NxD는 이를 neuronxcc.nki._private.blockwise_mm에서 임포트합니다 — 하지만 2.26 버전은 이를 _pre_prod_kernels 아래에 배포합니다. 경로가 잘못되었고, 프리프로덕션(pre-prod) 단계이므로 임포트가 실패합니다.
• 저의 DenseExperts를 fp4 패킹 기능이 있는 NxD의 ExpertMLPs로 교체하는 것도 도움이 되지 않았습니다: 해당 모듈의 blockwise NKI 경로는 kernel_assert에서 걸리고, use_torch_block_wise=True 경로는 데이터 의존적 인덱스 연산(data-dependent index op)에서 torch-xla의 형상 추론(shape inference)을 충돌(crash)시킵니다. 트레이스되지 않는 MoE 라우팅(routing) — 연산을 위해 dense 트릭이 해결했던 바로 그 문제가, 이제 양자화를 위해 다시 돌아온 것입니다.

이 SDK에서 inf2의 fp4는 벽입니다. 저는 그것이 유일한 문이라고 확신하며 실제 많은 시간을 소비했습니다. 그것은 문조차 아니었습니다.

돌파구: 저는 잘못된 것을 양자화하고 있었습니다

그것을 해결해 준 재정의(reframe)는 다음과 같습니다. 저는 전문가(experts)들이 거대하기 때문에 그들에게만 매달려 있었습니다. 하지만 제가 OOM(Out of Memory)을 겪었을 때 전문가들은 _이미 int8 상태_였습니다. 제게 필요했던 3~4 GB는 전문가들에게 있었던 것이 아니라, **bf16으로 남겨두었던 가중치(weights)**에 있었습니다. 그중 핵심은 하나였습니다: **rank당 1.48 GB씩 복제된 tied lm_head**였습니다.

해결책은 _더 작은 전문가(smaller experts)_를 만드는 것이 아니었습니다. 여전히 비대했던 모든 요소를 **전부 int8로 압축(all-int8 squeeze)**하는 것이었습니다:

1. lm_head를 int8로 양자화하고 샤딩(shard)하기. QuantizedColumnParallel(gather_output=True)는 262k 행의 head를 rank에 걸쳐 슬라이스하고 양자화합니다: 1.48 GB → rank당 ~0.37 GB. 이것이 단일 항목 중 가장 큰 영향력을 발휘한 요소였으며, 전문가(experts)와는 아무런 관련이 없었습니다.
2. 공유된 dense MLP(dual-path FFN의 나머지 절반)를 int8로 양자화하기: rank당 약 0.27 GB를 추가로 절감합니다.
3. int8 전문가(experts)를 그대로 유지하기.

Rank당 상주 메모리(resident)가 ~12–13 GB로 떨어졌고, 이는 16 GB보다 여유 있게 낮았습니다. 그리고 결과물은 여전히 정확했습니다:

Compiler status PASS · MB_TRACED
DEVICE_PARIS True
DEV GEN: 'The capital of France is Paris.'
...

26B MoE는 **2-core inf2.8xlarge**에서 실행되었습니다. "fp4 없이는 불가능하다"라는 결론은 fp4가 불가능하다는 점에서는 맞았지만, 그 외의 모든 것에 대해서는 틀렸습니다.

벽 2: 8xlarge와 xlarge는 동일한 HBM을 갖지만 — 호스트 RAM은 동일하지 않다

모두를 혼란에 빠뜨리는 미묘한 차이가 여기 있습니다. inf2.8xlarge와 inf2.xlarge는 동일한 2개의 NeuronCores와 32 GB HBM을 갖습니다. 차이점은 **호스트 RAM(host RAM)**입니다: 128 GB 대 16 GB입니다. 저의 압축 방식은 _장치(device)_에는 맞았습니다. 그렇다면 가장 저렴한 박스의 _호스트(host)_에도 맞을까요?

컴파일(Compiling) 단계에서는 불가능합니다. ModelBuilder trace는 fp32 discovery model + 구조 모델(structure model) + fp32 체크포인트 딕셔너리(checkpoint dict)를 보유하며, 피크(peak) 시점에 ~180 GB를 차지합니다. 이는 8xlarge의 128 GB조차 OOM을 발생시킵니다(저는 _컴파일_을 통과시키기 위해 100 GB 스왑 파일(swapfile)을 추가했으며, 약 40분이 소요되었습니다). 하지만 컴파일과 배포(deploying)는 서로 다른 작업입니다. 저장된 neff를 배포할 때는 모델이 호스트 RAM에 전혀 필요하지 않습니다. 트랜스포머 가중치(transformer weights)는 장치(device)에 상주하기 때문입니다.

따라서 xlarge 배포는 의도적으로 **슬림(slim)**하게 설계되었습니다 (deploy_sqz.py / optb_server_sqz.py):

• 구조 (Structure) (어떤 레이어가 KV-shared인지, head dims, sliding window 등)는 meta-device 모델 인스턴스화(accelerate.init_empty_weights())로부터 생성되며, 호스트 RAM을 거의 사용하지 않습니다 (~0 host RAM).
• **워드 임베딩 (Word embeddings)**은 호스트에서 한 번 로드되는 독립적인 1.48 GB 크기의 embed_tokens.pt 테이블에서 가져옵니다.
• neff는 디바이스 상의 모든 트랜스포머 가중치를 포함합니다 (torch.jit.load + 단 한 번의 initialize_with_saved_weights).
• **40 GB 스왑 파일 (swapfile)**이 일회성 neff 로드 시 발생하는 피크(peak)를 흡수합니다.

기본 사양의 16 GB inf2.xlarge에서의 결과:

NEFF_LOADED in 112s
DEV GEN: 'The capital of France is Paris.'
PREFILL 250 ms

호스트 RSS(Resident Set Size)는 약 11 GB로 피크를 찍었으며, 스왑은 거의 사용되지 않았습니다. 컴파일에는 180 GB가 필요하지만, 배포에는 노트북 한 대 정도의 사양만 있으면 됩니다.
7B 모델을 돌리기에도 망설여지는 사양의 박스에서 26B 모델을 돌린 것입니다.

파리를 공백의 벽으로 만들어버린 버그

하지만 첫 번째 슬림(slim) 실행에서는 파리(Paris)를 말하지 않았습니다. 다음과 같이 출력되었습니다:

DEV GEN: '                                                              '
DEVICE_PARIS False

neff는 검증되었습니다 (컴파일 시 파리를 생성했기 때문). 따라서 배포 과정에서 무언가 잘못된 값을 입력하고 있었던 것입니다. 동일한 토큰이 반복되는 현상은 **퇴화된, 거의 균등한 로짓 (degenerate, near-uniform logits)**의 특징이며, 이는 임베딩의 **크기 (magnitude)**가 잘못되었을 때 나타나는 전형적인 증상입니다.

Gemma-4는 일반적인 임베딩을 사용하지 않습니다. Gemma4TextScaledWordEmbedding을 사용하며, 이의 순전파(forward)는 다음과 같습니다:

return super().forward(input_ids) * self.embed_scale   # embed_scale = hidden_size ** 0.5

×√hidden_size (= √2816 ≈ 53) 연산이 임베딩 내부에서 일어나며, 모델의 순전파(forward)는 이를 다시 정규화(re-normalize)하지 않고 단순히 hidden_states = inputs_embeds를 수행합니다. neff는 스케일링된 (scaled) 임베딩을 기준으로 추적되었습니다. 하지만 저의 슬림 호스트 경로는 일반적인 torch.nn.Embedding을 사용했고, 스케일링되지 않은 (unscaled) 임베딩을 입력했습니다 — 즉, 53배나 작았습니다. 이로 인해 다운스트림(downstream)의 모든 정보가 노이즈로 씻겨 내려갔습니다.

한 줄 요약:

ie = emb(ids) * (hidden_size ** 0.5)   # Gemma4TextScaledWordEmbedding과 일치시킴

DEV GEN: 'The capital of France is Paris.'. 이는 E2B slim 서버가 빠졌던 것과 동일한 함정입니다. 임베딩 조회(embedding lookup)를 호스트로 옮기면, 모델의 임베딩 클래스가 수행하던 모든 동작을 그대로 상속받게 됩니다.

배포: 512/128, 게시, 그리고 클린 풀(clean-pull) 증명

저는 프로덕션용 512/128 버킷 빌드(총 512 / 프롬프트 토큰 128)를 다시 컴파일했습니다. 동일한 레시피를 사용했으며, 스왑(swap)을 포함한 8xlarge 인스턴스에서 약 40분이 소요되었고, neff 크기는 26.8 GB였습니다. 이후 xlarge에서 재검증을 완료했습니다(112초 만에 로드, 호스트 피크 메모리 약 10 GB, Paris).

그 다음, 이를 슬림한 OpenAI 호환 서버로 감싸서 게시했습니다:

• Docker Hub xbill9/gemma4-optb-26b:xlarge — 가벼운 이미지입니다. 엔트리포인트(entrypoint)가 첫 실행 시 HF 리포지토리에서 약 28 GB의 아티팩트(artifacts)를 가져옵니다.
• HF xbill9/gemma-4-26B-A4B-it-inferentia2-xlarge (공개) — neff + 임베딩 테이블(embedding table) + 설정(config) + 서버 + Dockerfile.

마지막 증명은 저 이외의 누구에게나 중요한 것이었습니다. 바로 **콜드 클린 풀(cold, clean pull)**입니다. 새로운 스팟(spot) inf2.xlarge 인스턴스를 생성하고, 스왑을 추가한 뒤, docker run을 실행하고 자리를 비웠습니다.

READY in 490.9s — slim int8-squeeze, ModelBuilder TP=2, MAX=512 BUCKET=128
{"role":"assistant","content":"The capital of France is Paris."}
{"response":"AWS Inferentia is a purpose-built machine learning accelerator
...

(HF_TOKEN은 필요하지 않습니다 — 리포지토리가 공개되어 있기 때문입니다. 490초는 최초 1회의 콜드 EBS neff 로드 시간이며, 그 이후에는 바로 서비스가 가능합니다.)
인스턴스를 종료했습니다. 지속적인 순 지출은 0입니다.

솔직한 주의사항

디코딩(Decode) 속도는 **초당 약 6개 토큰(6 tok/s)**입니다. 이것은 지난 기사에서 언급한 DenseExperts 트릭의 대가입니다. 희소한(sparse) 상위 8개 전문가를 사용하는 대신, 매 토큰마다 128개의 전문가를 모두 계산합니다. 이는 2개의 코어에서 필요한 전문가 FLOPs보다 약 16배 더 많은 연산을 수행합니다. 결과는 정확하며, 호스팅 비용이 저렴합니다. 하지만 빠르지는 않습니다. ModelBuilder 환경에서 희소 라우팅(sparse routing)이 실제로 트레이스(trace)되도록 만드는 것이 진정한 다음 과제이며, (fp4 우회 경로를 참조하면) "벤더 프리미티브(vendor primitive)가 존재한다"는 것이 곧 "트레이스된다"는 것과 같지는 않습니다. 시간당 $0.76인 인스턴스에서 26B 모델을 단일 사용자용으로, 비용에 민감하고 지연 시간(latency)에 관대한 방식으로 서빙하기에는 적절한 트레이드오프(tradeoff)입니다. 하지만 부하가 걸린 챗봇에게는 아직 적합하지 않습니다.

요점

1. 3GB가 초과되었다고 해서, 그 3GB가 당연히 예상 가능한 부분일 것이라고 가정하지 마세요. 저는 fp4 전문가(experts) 레이어 때문에 며칠을 허비했습니다. 해결책은 제가 전혀 들여다보지 않았던 _복제된 lm_head (replicated lm_head)_를 양자화(quantizing)하는 것이었습니다. 아키텍처를 눈으로만 훑지 말고, 상주 메모리(residency)를 프로파일링(profile)하세요.
2. "지원되지 않음"은 "문서화가 잘 안 됨"보다 낫습니다. NxD 0.19 버전의 fp4에는 마치 커밋 한 번이면 될 것처럼 보이는 _상수(constants)_와 _부분적인 코드 경로(partial code paths)_가 존재합니다. 하지만 실제로는 그렇지 않습니다. 커널(kernels)이 CPU 전용 참조(refs)이거나 프리 프로덕션(pre-prod) 단계이기 때문입니다. 기능 가이드(feature guide)의 지원 목록을 읽고 그대로 믿으세요.
3. 컴파일 적합성(Compile-fit)과 배포 적합성(deploy-fit)은 서로 다른 예산 문제입니다. _트레이스(trace)_를 위해 180GB의 호스트 RAM(host RAM)이 필요한 모델이라도, 가중치(weights)가 디바이스(device)에 상주한다면 16GB에서도 배포할 수 있습니다. 슬림한 호스트 로딩(Slim host loading)과 스왑(swap)이 24xlarge급 모델을 xlarge급 제품으로 탈바꿈시킵니다.
4. 임베딩 룩업(embedding lookup)을 호스트로 옮긴다는 것은 임베딩 _클래스(class)_를 그대로 상속받는다는 것을 의미합니다. Gemma의 ×√H 스케일링은 Gemma4TextScaledWordEmbedding 내부에 있습니다. 이를 놓치면 확신에 찬 공백의 벽을 마주하게 될 것입니다.
5. 깔끔한 풀(clean-pull) 증명을 제출하세요. "검증된 내 박스에서는 잘 돌아간다"는 결과물이 아닙니다. "새로운 스팟 인스턴스(spot instance)에서 docker run을 했을 때 파리(Paris)라고 나온다"가 결과물입니다.

Artifacts

• Docker Hub: xbill9/gemma4-optb-26b:xlarge
• HF: xbill9/gemma-4-26B-A4B-it-inferentia2-xlarge (공개)
• 레시피(Recipe): tp_mb_moe_sqz.py (전체 int8 압축: int8 experts + sharded int8 lm_head + int8 shared MLP) · deploy_sqz.py / optb_server_sqz.py (슬림 호스트-임베딩 배포)

AWS가 대여하는 가장 저렴한 가속기 인스턴스에서, 하루에 커피 한 잔 가격으로 실행되는 26B Mixture-of-Experts 모델입니다. Claude Code 세션에서 AI의 도움을 받아 작성되었으며, 인용된 모든 로그 라인은 실제 하드웨어에서의 실제 실행 결과입니다.