Apple Silicon에서 구동하는 Command A+ (218B MoE) — MLX 포트, PR 오픈
Command A+ (218B MoE) running on Apple Silicon — MLX port, PR open
핵심 요약
Cohere의 218B MoE 모델인 Command A+를 Apple Silicon에서 구동하기 위한 MLX 구현체와 아키텍처 상세 정보를 공유함.
- MLX 포트 구현 — Apple Silicon 환경에서 Command A+ 모델을 구동할 수 있도록 mlx-lm 구현체 작성.
- 모델 아키텍처 상세 — Sigmoid 라우팅, 슬라이딩 윈도우, 병렬 attn+MLP 블록 등 기술적 특징 분석.
- 하드웨어 요구사항 — W4A4 양자화 모델 구동 시 약 132GB의 메모리가 필요하며, 128GB M3 Max로는 검증 불가.
- PR 진행 상황 — ml-explore/mlx-lm에 PR이 오픈되었으며, 고사양 환경에서의 테스트와 피드백을 기다리는 중.
Cohere가 20일에 Command A+를 공개함(총 218B / 활성 25B, 128개 전문가 모델 중 상위 8개 선택, Apache 2.0). Apple Silicon에서 구동하기 위해 mlx-lm용 cohere2_moe 구현체를 작성함.
이 모델을 파고드는 분들을 위한 아키텍처 노트:
-
더 큰 중간층(16384 = 4096×4)을 가진 단일 공유 전문가 모델과 (라우팅된 + 공유)/2를 통한 라우팅된 출력의 결합.
-
Sigmoid 라우팅(softmax 아님), 정규화된 상위 8개 선택.
-
3:1 슬라이딩 윈도우(3개 슬라이딩 + 1개 전체), 슬라이딩 레이어에만 인터리브된 RoPE 적용.
-
동일한 LayerNorm에서 분기되는 병렬 attn+MLP 블록.
-
몇 번의 반복 끝에 알아낸 주의사항: W4A4 체크포인트의 편향값(bias)은 NVFP4 양자화 아티팩트임. BF16 모델은 편향값이 전혀 없음. sanitize()가 두 형식 모두 처리함.
로컬에서 검증할 수 없었음(W4A4는 약 132GB가 필요한데 내 M3 Max는 128GB임). https://github.com/vlbosch 가 더 큰 장비에서 실행함: BF16→Q8 변환 + 깔끔한 생성, 도구 호출, KV-캐시 연속성을 포함한 멀티턴, 22.9 tok/s 생성 / 57.6 tok/s 프롬프트, 241GB 피크 메모리 사용.
ml-explore/mlx-lm에 PR이 오픈됨(검토 중). 피드백이나 수정 사항은 언제든 환영함. 192GB 이상의 메모리를 가진 분이 W4A4 경로를 직접 테스트하고 에러 로그를 공유해 주면 정말 좋겠음. https://github.com/ml-explore/mlx-lm/pull/1294
