lablup / lablup/mlxcel

feat(models): add Arcee Foundation Model (afm7) text model support

Open
#843 2 comments 0 reactions 1 assignee Claimed by @inureyes View on GitHub
arch:dense area:models modeltype:text priority:medium status:blocked type:enhancement
Dominant language
Rust
Stars
467
Forks
54
Avg merge
4h 25m
Merged PRs (30d)
310

Description

> **Corrected against `references/mlx-lm` (commit `15b522f`, 2026-07-11) on 2026-08-05.**
> The original body below misidentified the vendor and the porting work. Both are corrected here; the rest of the original text is kept for history.

## Corrections

**This is not an Arcee model.** `mlx_lm/models/afm7.py` carries `# Copyright © 2025 Apple Inc.`, and AFM here is Apple's foundation model. `arcee-ai/AFM-4.5B` declares `model_type: arcee`, which is a different family with a different `ModelArgs`, so the checkpoint the original body names cannot load this implementation.

**Fused LoRA is not the work.** `FusedLoRALinear` is reachable only through `FusedLinear.to_lora()`, an adapter/training construct. A published checkpoint ships a plain fused `qkv_proj.weight`. `FusedLinear` is simply one weight with an uneven multi-output split, not a merged low-rank delta.

## What upstream actually implements

Config (`ModelArgs`): `hidden_dim`, `num_layers`, `num_kv_reuse_layers`, `num_heads`, `num_kv_heads`, `hidden_dim_scale_factor = 3.25`, `rope_theta = 50000`, `rms_norm_eps = 1e-5`. There is no `intermediate_size`: MLP width is `int(hidden_dim * hidden_dim_scale_factor)`.

The distinctive pieces, in rough order of porting risk:

- **KV-reuse layers.** The stack is `num_layers - num_kv_reuse_layers` ordinary `TransformerBlock`s followed by `num_kv_reuse_layers` `KVReuseTransformerBlock`s. The last ordinary layer uses a `ConcatenateKVCache`; its `keys`/`values` are then read once and reused by *every* KV-reuse block. Those blocks have a query projection only (no k_proj/v_proj/cache of their own) and rotate queries at `offset = S - L`.
- **Simulated 8-bit KV.** `fake_8bit_quant(x, scale)` rounds K and V through a per-tensor scale (`quant_key_scale`, `quant_value_scale`, both learned scalars in the checkpoint) before they enter the cache. Skipping it leaves output finite and plausible.
- **QK-norm applied AFTER RoPE**: `self.q_norm(self.rope(queries, offset=...))`. Most families in this tree norm before rotating; reversing the order here is silent.
- **Traditional RoPE**: `initialize_rope(head_dim, rope_theta, True)`, the third positional argument being `traditional`.
- **Fused uneven QKV split** at `[n_heads * head_dim, n_kv_heads * head_dim, n_kv_heads * head_dim]`.
- **Always-tied embeddings**: the model ends in `self.model.embedding.as_linear(out)` and defines no `lm_head`.

Weight prefixes: `model.embedding`, `model.layers.{i}`, `model.kv_reuse_layers.{i}`, `model.output_norm`.

## Checkpoint status: none public

Searches of the HuggingFace API for `afm7`, `apple/AFM`, `AFM-Text` and a `model_type: afm7` filter return zero results (2026-08-05). Apple ships AFM on-device rather than on the Hub. **The real-checkpoint gate this project requires cannot be met today**, so this issue stays open pending a published checkpoint rather than being implemented against synthetic weights alone.

---

Original issue body (retained for history; contains the vendor and fused-LoRA errors corrected above)

## Summary

Add support for the Arcee **Arcee Foundation Model (AFM)** dense family to the mlxcel Rust inference runtime, ported from mlx-lm. AFM is a Llama-style dense decoder (RMSNorm, RoPE via `initialize_rope`). The distinctive work is the load-time weight layout: upstream defines `FusedLoRALinear` / `FusedLinear` / `FusedQuantizedLinear` wrappers, meaning some linear layers may ship with fused LoRA deltas baked into the checkpoint. The `sanitize`/load path must fold or map the fused-LoRA linear layout correctly so a plain forward matches the reference.

## Upstream reference

- mlx-lm: https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/models/afm7.py

## Public checkpoint

- `arcee-ai/AFM-4.5B`. Verify an `mlx-community` conversion exists (or convert from the Arcee release) before closing. State the exact checkpoint id used for validation in the PR.

## Architecture notes

- Llama-style dense decoder: token embedding, N decoder layers, final RMSNorm, `lm_head`.
- RMSNorm; RoPE through `initialize_rope`.
- **Fused-linear weight layout**: upstream `FusedLinear` / `FusedLoRALinear` / `FusedQuantizedLinear` wrappers imply linear weights may carry fused LoRA deltas (base weight plus low-rank `A`/`B` factors, or a pre-merged weight). The load path must resolve these to an effective weight so the forward matches. Decide at load time whether the checkpoint ships pre-merged weights or separate LoRA factors, and handle both if the family varies.

## Implementation plan

- Reuse the dense decoder path in `src/models/llama3.rs` as the structural template.
- The distinctive effort is in `sanitize` / `from_weights`: detect the fused-LoRA linear layout and fold the LoRA deltas into the base weight (or map the factors onto the existing LoRA support) so downstream construction sees a plain linear. Validate that a merged-weight forward matches the reference numerically.
- Keep bf16 to f16 conversion for non-quantized weights at the standard load boundary; leave quantized scales/biases as bf16 per the project precision rules.

## Touchpoints & acceptance criteria

Follow `docs/adding-models.md`. Done only when the model loads and generates from a real checkpoint.

- [ ] Config struct + serde parse for the AFM config.
- [ ] `from_weights` constructs the model; fused-LoRA linear layers resolve to effective weights.
- [ ] `sanitize` handles the fused-linear weight layout (fold LoRA deltas or map factors) plus the standard bf16 to f16 conversion and any key remapping.
- [ ] Detection arm added in `src/models/detection.rs` keyed on the `afm7` model_type.
- [ ] Registration in `src/model_metadata.rs` via `for_each_model_registration!`.
- [ ] TP/distributed arch-string wired if AFM is a target for tensor/pipeline parallel.
- [ ] Unit tests in a `_tests.rs` file beside the implementation, covering the fused-linear fold path.
- [ ] `docs/supported-models.md` updated with the new family.
- [ ] Real-checkpoint validation: `./target/release/mlxcel generate` produces coherent output and `mlxcel list` reports the architecture.

## Effort

MEDIUM. The decoder is stock Llama; the risk concentrates in getting the fused-LoRA linear weight layout folded correctly at load time.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.