Qwen36 routed-expert TWLA W1.58A16, group size 64
Qwen36 35B-A3B 계열 모델의 routed expert를 E2M-ATQ/TWLA로 ternary 양자화한 public 연구 체크포인트다. TWLA의 row scale과 offset 공유 범위를 전체 입력 차원 대신 연속된 64개 input column로 제한한 실험이다.
무엇이 양자화됐나
- 모든 40개 layer의 256개 routed expert
routed_experts.gate_up_projrouted_experts.down_proj- Weight alphabet:
{-1, 0, +1}, 이론상log2(3)=1.584963bits/weight - Activation: BF16 (A16)
- Router, attention, embedding, shared expert 및 나머지 파라미터: BF16
각 output row에서 연속된 64 input columns마다 독립적인 (mu, alpha)를 사용한다. Stage I은 Euclidean warm start, Stage II는 해당 group의 S=X^T X diagonal block에서 closed-form manifold relocation을 수행한다.
이 저장소는 실행 호환성을 위해 최종 ternary reconstruction을 BF16 tensor로 저장한 fake-quant checkpoint다. 즉 양자화 효과는 반영되어 있지만 디스크/메모리 파일 크기는 실제 1.58-bit packing 크기가 아니다.
Calibration 데이터와 목적함수
Gradient 학습은 하지 않았다. Calibration pool은 C4 160, code 32, math 32, science 32, long-context 8개 문서를 같은 텍스트로 재토큰화해 구성한 264개 혼합 corpus다. 양자화 실행에서는 seed에 따라 2,048-token chunk 128개를 사용했다.
각 expert의 실제 routed activation으로 두 선형 변환의 second moment를 별도로 누적했다. Stage II는 ternary code를 고정한 뒤 다음 reconstruction metric을 최소화하도록 group별 offset과 scale을 다시 구한다.
min_(mu, alpha) ||(W - mu - alpha*T) X||_F^2
실행 메타데이터는 twla_quant_stats.json에 있다.
추론
pip install -r requirements.txt
python inference.py --model minjaechoi/qwen36-twla-group64-w1p58a16 \
--prompt "Explain why the sky is blue." --max-new-tokens 512
체크포인트는 dense BF16 크기이므로 한 GPU에서 읽으려면 그에 맞는 VRAM이 필요하다. 메모리가 부족하면 device_map 또는 tensor parallel을 사용한다.
양자화 코드
code/에 실험에 사용한 실제 구현을 포함한다.
quantize_qwen36_experts_groupwise.py: group-size 진입점quantize_qwen36_experts.py: routed activation 수집과 checkpoint 저장quantize/E2M_ATQ_groupwise.py: groupwise Stage I/IIquantize/E2M_ATQ.py: 기본 TWLA 구현datautils.py: calibration loader
예시:
python code/quantize_qwen36_experts_groupwise.py \
--group_size 64 \
--model_dir /path/to/base-model \
--calibset /path/to/robust_mixed_c4.jsonl \
--nsamples 128 --seqlen 2048 \
--out_dir /path/to/output
코드의 로컬 기본 경로는 CLI 인자로 반드시 덮어쓴다.