feat(models): add Arcee AFM-4.5B (arcee, dense ArceeForCausalLM) text model support
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
## Problem / Background
`models/afm-4.5b` (`arcee-ai/AFM-4.5B`, Apache-2.0, 4.6B params, bfloat16, 9.2 GB on disk) does not load: `mlxcel generate` exits with `Unsupported model type: arcee` (`src/models/detection.rs:798`), recorded as `FAIL:unsupported_arch_arcee_dense` in `benchmarks/metal_m5max_2026-09-09.csv`.
**This is not #843.** That issue is Apple's AFM-7 (`mlx_lm/models/afm7.py`, `Copyright Apple Inc.`) and already carries a correction saying so after the same vendor confusion happened once. The two share the "AFM" initials and nothing else.
## Current Behavior
`mlxcel arch` prints an "Arcee" heading (`src/main.rs:2871`) whose only entry is `ModelType::Afmoe`, the AFMoE / Trinity MoE variant (`src/models/mod.rs:980`), so the vendor heading alone suggests support that is not there.
## Why this is blocked
mlx-lm does not implement this family either: `from mlx_lm import load` on this checkpoint raises `ValueError: Model type arcee not supported.`, and `references/mlx-lm/mlx_lm/models/` ships only `afm7.py` and `afmoe.py`. Without a reference there is no parity target, and the only evidence of correctness would be that our own output looks plausible, which is not enough for an inference change. Revisit when mlx-lm adds `arcee`.
## Architecture (config.json)
36 layers, hidden 2560, intermediate 18432, 20 heads / 4 KV heads (GQA), vocab 128005, `max_position_embeddings` 65536. `hidden_act: relu2` with no `gate_proj` in the weight index (each layer ships `mlp.up_proj` and `mlp.down_proj` only), so the MLP is gate-free squared-ReLU and unusually wide at 7.2x hidden to compensate. `rope_scaling` is YaRN: `factor 20.0`, `original_max_position_embeddings 4096`, `beta_fast 32.0`, `beta_slow 1.0`, `mscale 1.0`.
## Implementation Notes
The building blocks exist, so this is gated on validation rather than difficulty. Squared ReLU: `mlxcel_core::utils::relu_squared`, used by `src/models/nemotron.rs:324`, `src/models/rwkv7.rs:356` and `src/models/falcon_ocr.rs:472`. YaRN: `RopeScheme::Yarn` in `src/models/rope_utils.rs:146`. GQA and the dense decoder are shared paths. `src/models/afmoe.rs` establishes where an Arcee-family loader goes and which registration sites it touches.
## Acceptance Criteria
- [ ] A reference implementation exists and its commit is recorded here before work starts.
- [ ] `mlxcel arch` lists a dense `arcee` entry under the Arcee heading, distinct from AFMoE.
- [ ] Logit-trace parity against the reference per `docs/benchmarks.md`.
- [ ] `docs/supported-models.md` updated and the benchmark row re-run so `FAIL:unsupported_arch_arcee_dense` clears.
## Verification
```bash
cargo test --release arcee
cargo clippy --all-targets -- -D warnings
./target/release/mlxcel generate -m models/afm-4.5b -p "Hello" -n 50
```
Contributor guide
Research direction
First confirm that mlx-lm has added an arcee reference implementation and record its commit, since the issue is explicitly blocked on that prerequisite. Then read src/models/afmoe.rs, src/models/detection.rs, src/models/rope_utils.rs, and docs/benchmarks.md; use cargo test --release arcee as the initial check. Done means distinct dense arcee discovery, reference logit-trace parity, updated docs/supported-models.md, and a cleared benchmark failure.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, rust
- Domain
- ai, machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100