lablup / lablup/mlxcel

feat(models): add Arcee AFM-4.5B (arcee, dense ArceeForCausalLM) text model support

Open
#1,719 1 comment 0 reactions 0 assignees View on GitHub
arch:dense area:models modelsize:small modeltype:text priority:medium status:blocked type:enhancement
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.