Expert-axis Compression of 250B-class MoE Models

문제

Problem

250B급 MoE는 대부분 멀티 노드 GPU 클러스터를 전제로 배포됩니다. 국내 온프레미스나 소규모 GPU 환경에서는 그대로 올릴 수 없고, 레이어를 잘라내거나 distillation을 하면 원 모델의 언어 특성과 reasoning 형태가 함께 바뀝니다.

250B-class MoE checkpoints assume multi-node GPU clusters. On-premise and small-cluster environments cannot host them as shipped, while layer truncation or distillation changes the source model's language behavior and reasoning shape along with its size.

MoE에서 파라미터의 대부분은 routed expert에 있습니다. 구조를 유지한 채 줄이려면 레이어가 아니라 expert 축을 건드려야 한다는 것이 이 작업의 출발점입니다.

In MoE models most parameters live in the routed experts, so the structure-preserving place to cut is the expert axis, not the layer stack.

접근

Approach

1. REAP 프루닝 — 라우팅 가중 expert 출력 norm으로 saliency를 계산해 레이어별 하위 expert를 제거했습니다. K-EXAONE-236B는 128 → 80/76 expert, Solar-Open2-250B는 184 expert로 줄였습니다.

1. REAP pruning — Ranked experts by routing-weighted output norm and dropped the low-saliency tail per layer: K-EXAONE-236B from 128 to 80/76 experts, Solar-Open2-250B down to 184.

2. REAM 병합 — 버리는 대신 합쳤습니다. 상위 saliency expert를 seed로 pseudo-group을 만들고, activation PCA와 weight feature에 대한 Hungarian matching으로 뉴런을 정렬한 뒤 saliency 가중 평균으로 병합했습니다. DeepSeek-V4-Flash는 hash-routing 레이어의 token-to-expert 테이블까지 재매핑했습니다.

2. REAM merging — Instead of discarding, merged: high-saliency experts seed bounded pseudo-groups, neurons are aligned by Hungarian matching over activation-PCA and weight features, then averaged with saliency weights. For DeepSeek-V4-Flash the hash-routing layers' token-to-expert tables were remapped as well.

3. Healing — 병합 후 손실은 router를 full-rank로, attention/shared expert에는 low-rank residual로 회복시켰습니다. routed expert 텐서는 동결해 병합 결과를 그대로 보존했습니다.

3. Healing — Recovered post-merge loss by training routers at full rank with low-rank residuals on attention and shared-expert modules, keeping the merged 3-D expert tensors frozen.

4. NVFP4 양자화 — routed expert 투영만 FP4로 내리고 router, attention, shared expert, embedding, LM head는 BF16으로 남겼습니다. 라우터가 양자화되지 않으므로 양자화가 라우팅 결정 자체를 바꾸지 못합니다.

4. NVFP4 quantization — Only routed expert projections go to FP4; router, attention, shared expert, embedding, and LM head stay BF16. Because the router is never quantized, quantization cannot alter routing decisions.

기여와 검증

Contribution and verification

REAM 파이프라인을 Solar/DeepSeek 아키텍처에 맞게 직접 구현했고(Baekpica/solar-ream), 런타임이 gate_proj와 up_proj를 하나로 fuse해 로드하는 문제를 해결하기 위해 llm-compressor에 Solar 전용 fused mapping을 등록해 7,104개 gate/up 페어가 동일한 NVFP4 global scale을 공유하도록 했습니다.

Implemented the REAM pipeline for the Solar and DeepSeek architectures (Baekpica/solar-ream) and registered a Solar-specific fused mapping in llm-compressor so that all 7,104 gate/up pairs share one NVFP4 global scale, matching how the runtime loads them as a single fused tensor.

모든 릴리스는 동일한 고정 프롬프트, 동일 엔진, 동일 샘플링(vLLM, TP4, temperature=0, seed=42)으로 부모 모델과 비교했고, per-layer saliency 진단, 구조 감사, 데이터셋·어댑터 체크섬을 함께 공개했습니다. 산술 오류나 반복 루프처럼 회복되지 않은 한계도 모델 카드에 그대로 남겼습니다.

Every release was compared against its parent with the same fixed prompts, engine, and sampling (vLLM, TP4, temperature=0, seed=42), and ships per-layer saliency diagnostics, structural audits, and dataset/adapter checksums. Limitations that healing did not recover — arithmetic errors, repetition loops — are documented in the model cards rather than omitted.