Apple Silicon에서 실행되는 Command A+ (218B MoE) — MLX 포트, PR 오픈

Cohere가 20일에 Command A+를 출시했습니다 (총 218B / 활성 25B, 128명의 전문가 중 top-8 선택, Apache 2.0). Apple Silicon에서 실행할 수 있도록 mlx-lm을 위한 cohere2_moe 구현을 작성했습니다.

이 모델을 깊게 파고드는 분들을 위한 아키텍처 참고 사항:

• 더 큰 중간층(16384 = 4096×4)을 가진 단일 공유 전문가(Single shared expert)가 (routed + shared)/2를 통해 라우팅된 출력과 결합됩니다.

• Sigmoid 라우팅 (softmax가 아님), 정규화된 top-8

• 3:1 슬라이딩 윈도우 (3개의 슬라이딩 + 1개의 전체), 슬라이딩 레이어에만 인터리브된 RoPE (interleaved RoPE) 적용

• 동일한 LayerNorm에서 분리된 병렬 attn+MLP 블록

• 제가 몇 번의 시행착오를 겪었던 주의 사항: W4A4 체크포인트의 편향(biases)은 NVFP4 양자화 아티팩트(quantization artifacts)입니다 — BF16 모델은 편향이 전혀 없습니다. sanitize() 함수가 두 형식을 모두 처리합니다.

로컬에서 검증할 수는 없었습니다 (W4A4는 약 132GB가 필요하지만, 제 M3 Max는 128GB입니다). https://github.com/vlbosch가 더 큰 장비에서 실행했습니다: BF16→Q8 변환 + 깨끗한 생성, 도구 호출(tool calling), KV-cache 지속을 통한 멀티턴(multi-turn), 생성 속도 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