diff --git a/CHANGELOG.md b/CHANGELOG.md index e79c814..6c460d2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,24 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added + +- **AMALIA-9B** (European Portuguese, EuroLLM-9B base) support, golden-gated + against mlx-lm (`reference/fixtures_amalia/`). The forward pass is the stock + Llama decoder, run in **bfloat16** via the new `ModelConfig::bf16_compute` + (this family's residual stream overflows fp16 from layer 9 — an fp16 forward + is all-NaN logits; selected by config fingerprint, all other checkpoints stay + on fp16 bit-for-bit). The other new pieces are tokenizer/prompt plumbing: the + SentencePiece-BPE backend honors the Metaspace `prepend_scheme: "always"` + pre-tokenizer and the `Strip` decode step (both parsed from `tokenizer.json`; + Gemma unchanged), a `ChatFormat::Amalia` ChatML template with the default + Portuguese system prompt (detected from the checkpoint's `chat_template.jinja` + via the new `chat_format_for_model_dir`, since AMALIA's `model_type` is + `"llama"`), and config-load rejection of non-affine quantization modes. + KV-cache quantization and the prefix cache are rejected for bf16-compute + models (no golden coverage yet); the golden compare harness now fails on + nonfinite values instead of passing NaN-vs-NaN vacuously. + ## [0.2.0] - 2026-06-12 This release adds a from-scratch **Qwen3-VL vision-language** model (image → text), diff --git a/CLAUDE.md b/CLAUDE.md index 3b178a9..7cdd997 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -138,11 +138,32 @@ reference/.venv/bin/python reference/dump_ref.py (`reference/fixtures*/tokenizer_corpus.json`, regenerated by `dump_ref.py`). Two backends: byte-level BPE (`tokenizer/bpe.{h,cpp}`, `BpeTokenizer`; Llama-3.2 / Qwen) and SentencePiece-BPE (`tokenizer/spm.{h,cpp}`, `SpmBpeTokenizer`; - metaspace + `byte_fallback`, e.g. Gemma — tokenizer-only, no Gemma model class). - Both are pure/`const`/thread-safe (no mutex). `Tokenizer::from_file` throws on - still-unimplemented families (e.g. Unigram/WordPiece). -- **Masks are additive fp16, never boolean** (avoids MLX bug #2894). See - `DecoderModel::batch_mask`. + metaspace + `byte_fallback` — Gemma (tokenizer-only, no Gemma model class) and + AMALIA-9B/EuroLLM). The SPM backend honors the Metaspace `prepend_scheme` per + checkpoint ("always" prepends `▁` to the *input-leading* segment only — never + after special tokens, never doubled onto a leading space; empirically pinned + by the AMALIA corpus fixture — Gemma never prepends; "first" throws) and a + decoder `Strip` step (drop N leading spaces on decode). Both are + pure/`const`/thread-safe (no mutex). `Tokenizer::from_file` throws on + still-unimplemented families (e.g. Unigram/WordPiece). AMALIA's chat format + can't come from `model_type` (it's `"llama"`): `chat_format_for_model_dir` + detects it from the checkpoint's `chat_template.jinja`, and because AMALIA's + BOS id 3 IS `<|im_start|>`, `apply_chat_template` encodes the fully-rendered + ChatML **without** the BOS prepend (HF `add_special_tokens=False`) — letting + the encoder prepend it would double the opening tag. +- **Masks are additive fp16 (or the model's compute dtype), never boolean** + (avoids MLX bug #2894). See `DecoderModel::batch_mask`. +- **The engine computes in fp16 — except checkpoints that overflow it.** + AMALIA-9B/EuroLLM-9B's residual stream exceeds fp16 range from layer 9, so an + fp16 forward is all-NaN logits and the fp16 mlx-lm reference is equally + degenerate (NaN goldens gate nothing — `compare_close` now fails on + nonfinite values for exactly this reason). `ModelConfig::bf16_compute` + (selected by the family's config fingerprint in `from_json`, since every MLX + conversion declares `dtype: bfloat16` and can't drive this) switches the + weight-load cast, the batched mask, and the prefill logits placeholders to + bf16; `skinny_mm` self-disables (fp16-only kernels) and kv-quant / prefix + cache are rejected at engine creation for bf16 models. All other checkpoints + keep fp16 bit-for-bit. - **RoPE scaling is precomputed and validated, never silent** — `compute_rope_setup` mirrors `mlx_lm` exactly (`Llama3RoPE`, `YarnRoPE`, `nn.RoPE(scale=1/factor)` for linear) and feeds `fast::rope` via `freqs` diff --git a/README.md b/README.md index 3bedf8b..7fc9383 100644 --- a/README.md +++ b/README.md @@ -81,10 +81,11 @@ the C-ABI / Node quickstart. ## Supported models The forward pass (the `DecoderModel` base in `model/`) is family-shared and runs -Llama-3.2, Qwen3 (dense / MoE), and Qwen3.5 hybrid models from safetensors (fp16 / -4-bit) or a single-file GGUF, plus **Qwen3-VL** vision-language (image → text). The -chat template and special-token handling are selected from `config.json` — loading -is org-agnostic. +Llama-3.2, AMALIA-9B (European Portuguese), Qwen3 (dense / MoE), and Qwen3.5 hybrid +models from safetensors (fp16 / 4-bit) or a single-file GGUF, plus **Qwen3-VL** +vision-language (image → text). The chat template and special-token handling are +selected from `config.json` and the checkpoint's chat template — loading is +org-agnostic. | Family | Example repo | Precision | | --- | --- | --- | @@ -95,6 +96,7 @@ is org-agnostic. | Qwen3 (MoE) | `mlx-community/Qwen3-30B-A3B-4bit` | 4-bit / fp16 | | Qwen3.5 (hybrid) | `mlx-community/Qwen3.5-0.8B-4bit` | 4-bit (text tower) | | Qwen3-VL (vision) | `mlx-community/Qwen3-VL-4B-Instruct-4bit` | 4-bit (image → text) | +| AMALIA-9B (pt-PT) | `layerx-labs/AMALIA-9B-0626-DPO-MLX-4bit` | 4-bit | See [`doc/supported-models.md`](./doc/supported-models.md) for the full compatibility matrix, the per-family deltas, and how to add a new family. diff --git a/apps/mlxforge_cli.cpp b/apps/mlxforge_cli.cpp index b74a7d8..39968e8 100644 --- a/apps/mlxforge_cli.cpp +++ b/apps/mlxforge_cli.cpp @@ -110,7 +110,7 @@ LoadedModel load_for_inference(const std::string& spec, const std::string& rope_ mlxforge::validate_rope_scaling(lm.cfg); lm.model = mlxforge::create_model(lm.cfg, mlxforge::load_weights(resolved, lm.cfg)); lm.tok = mlxforge::Tokenizer::from_file(resolved + "/tokenizer.json", lm.cfg.bos_token_id, - mlxforge::chat_format_from_model_type(lm.cfg.model_type)); + mlxforge::chat_format_for_model_dir(resolved, lm.cfg.model_type)); } return lm; } diff --git a/doc/supported-models.md b/doc/supported-models.md index 812f74b..a527f12 100644 --- a/doc/supported-models.md +++ b/doc/supported-models.md @@ -25,6 +25,7 @@ override — all selected automatically by `create_model` from `config.json` | Qwen3 (MoE, GGUF) | `Qwen/Qwen3-30B-A3B-GGUF` | Q4_0/Q4_1/Q8_0, Q4_K/Q5_K/Q6_K | ChatML (`<\|im_start\|>…`) | yes | | Qwen3.5 (hybrid) | `mlx-community/Qwen3.5-0.8B-4bit` | 4-bit (mixed bits ok) | ChatML (`<\|im_start\|>…`) | yes (text only) | | Qwen3-VL (dense) | `mlx-community/Qwen3-VL-4B-Instruct-4bit` | 4-bit (ViT in fp16) | ChatML + vision tokens | yes (image → text) | +| AMALIA-9B (pt-PT) | `layerx-labs/AMALIA-9B-0626-DPO-MLX-4bit` | 4-bit, bf16 compute | ChatML + PT default system | yes | **Qwen3 dense** models (0.6B/1.7B/4B/8B/14B/32B) run end-to-end. They add three deltas over Llama-3.2, all handled automatically: per-head **QK-Norm** (an @@ -111,6 +112,39 @@ is **not** bit-identical to HF's PIL bicubic, so results on non-aligned images m differ slightly from the reference. The **MoE-VL** variant (30B-A3B) is not yet a distinct model class. +**AMALIA-9B** (European Portuguese, built on EuroLLM-9B) runs end-to-end. Its +forward pass is the **stock Llama decoder** — plain `LlamaForCausalLM` (42 +layers, GQA with 8 KV heads, `rope_theta` 1e6, no rope scaling, untied +embeddings), so `create_model` falls through to `LlamaModel` with no per-family +hook — but it runs in **bfloat16, not the engine's usual fp16**: the family's +residual stream overflows fp16's range from layer 9, so an fp16 forward is +all-NaN logits (silent garbage — the fp16 mlx-lm reference is equally +degenerate). `ModelConfig::bf16_compute` selects bf16 from the family's config +fingerprint; KV-cache quantization and the prefix cache are rejected for +bf16-compute models (no golden coverage yet). The remaining deltas are in the +tokenizer and prompt format, all golden-gated against mlx-lm +(`reference/fixtures_amalia/`): + +- **SentencePiece-BPE with Metaspace prepend**: the Metaspace pre-tokenizer uses + `prepend_scheme: "always"` (a `▁` is prepended to the input-leading text + segment before BPE — never after special tokens, and never doubled onto a + leading space; the exact semantics are pinned by the golden corpus), and the + decoder chain carries a `Strip` step that removes the resulting leading space + on decode. Both are parsed from `tokenizer.json`, not hard-coded (Gemma never + prepends). +- **ChatML with a default Portuguese system prompt**: when the conversation has + no system turn, the template injects the "O teu nome é Amália…" preamble. + `model_type` is `"llama"`, so this format is detected from the checkpoint's + `chat_template.jinja` (via `chat_format_for_model_dir`), not from `model_type`. +- **`<|im_start|>` (id 3) doubles as BOS**: the tokenizer's post-processor + prepends it on plain encode; `apply_chat_template` therefore encodes the + (fully rendered) template *without* the BOS prepend, mirroring HF's + `add_special_tokens=False` — prepending would double the opening tag. EOS is + `[4, 2]` (`<|im_end|>` / ``). + +The SFT variant (`AMALIA-9B-0626-SFT-MLX-4bit`) shares the identical +architecture, tokenizer, and template, so it runs unchanged. + Other LLaMA-family models will be re-onboarded as needed; because the forward pass is shared, that work is mostly tokenizer/chat-format plus any small attention delta (see [Adding a new model family](#adding-a-new-model-family)). @@ -253,15 +287,20 @@ constants: - **`model_type`** selects the chat format via `chat_format_from_model_type`: the Llama-3.2 header format, or Qwen ChatML (`qwen3` / `qwen3_moe` / `qwen2`, `qwen3_5`, and `qwen3_vl` — the last rendering vision blocks for attached images). - For a VLM, the nested `vision_config` and the M-RoPE / vision-token fields are read - too. + Formats that `model_type` cannot identify are recognized from the checkpoint's + on-disk template (`chat_template.jinja`, else the `chat_template` string in + `tokenizer_config.json`) by `chat_format_for_model_dir` — AMALIA is `model_type` + `"llama"` but ChatML. For a VLM, the nested `vision_config` and the M-RoPE / + vision-token fields are read too. - **`tokenizer.json`** drives encode/decode and supplies the special-token ids (`added_tokens[*].special`) that are skipped on decode — there are no hard-coded token ids. Two backends sit behind one `EncoderBackend` interface, picked by inspecting the blob: **byte-level BPE** (`ByteLevel` decoder; Llama-3.2 / Qwen) - and **SentencePiece-BPE** (`byte_fallback` + metaspace; Gemma). BOS is resolved - from the fast tokenizer's post-processor first (Gemma prepends `` despite - `add_bos_token: false`), then the `tokenizer_config.json` flag. + and **SentencePiece-BPE** (`byte_fallback` + metaspace; Gemma, AMALIA — the + Metaspace `prepend_scheme` and the decoder's `Strip` step are honored per + checkpoint). BOS is resolved from the fast tokenizer's post-processor first + (Gemma prepends `` despite `add_bos_token: false`; AMALIA prepends + `<|im_start|>`), then the `tokenizer_config.json` flag. - A separate `lm_head.weight` is used if present; otherwise the embedding is tied. ## What is and isn't implemented @@ -283,9 +322,9 @@ GGUF `Q4_0`/`Q4_1`/`Q8_0` (group_size 32) all run; GGUF `Q4_K`/`Q5_K`/`Q6_K` - **Sliding-window attention.** Only plain causal attention is supported, so models that rely on a sliding window are not. - **Other tokenizer families / chat templates.** Byte-level BPE (Llama-3.2 / - Qwen) and SentencePiece-BPE (Gemma) are implemented; other families (e.g. - Unigram/WordPiece) make `Tokenizer::from_file` throw. Chat formats are limited - to the Llama-3.2 header format and Qwen ChatML. + Qwen) and SentencePiece-BPE (Gemma, AMALIA) are implemented; other families + (e.g. Unigram/WordPiece) make `Tokenizer::from_file` throw. Chat formats are + limited to the Llama-3.2 header format and ChatML (Qwen, AMALIA). - **Batched multimodal serving.** Qwen3-VL runs image → text, but multimodal requests are served **single-stream**; sharing a continuous-decode batch with text (per-row 3D M-RoPE in `BatchKVCache`) is not yet implemented. The **MoE-VL** diff --git a/doc/tokenizer.md b/doc/tokenizer.md index 389859d..450a53e 100644 --- a/doc/tokenizer.md +++ b/doc/tokenizer.md @@ -125,12 +125,15 @@ than BPE mechanics: configured `bos_id` (default `128000` = `<|begin_of_text|>`; `-1` adds none), matching `mlx-lm`'s `tok.encode`. Decode pre-filters special ids symmetrically. - **Chat template.** `apply_chat_template` / `render_chat_template` render the - Llama-3.2 header format (`<|start_header_id|>role<|end_header_id|>\n\n…<|eot_id|>`) - and inject the default "Cutting Knowledge Date / Today Date" system preamble. - The format is selected from `config.json`'s `model_type` via - `chat_format_from_model_type` → `ChatFormat`, keeping the (shared) forward pass - decoupled from the (per-family) prompt formatting. Only `Llama3` exists today; - the `enum` is the seam for new families. + model's format: the Llama-3.2 header format + (`<|start_header_id|>role<|end_header_id|>\n\n…<|eot_id|>` with the default + "Cutting Knowledge Date / Today Date" system preamble), Qwen ChatML + (`Qwen3`/`Qwen35`, with the thinking toggle), or AMALIA ChatML (`Amalia`, with + the default Portuguese system prompt). The format is selected from + `config.json`'s `model_type` via `chat_format_from_model_type` → `ChatFormat` + — or, for formats `model_type` cannot identify (AMALIA is `"llama"`), from the + checkpoint's on-disk chat template via `chat_format_for_model_dir` — keeping + the (shared) forward pass decoupled from the (per-family) prompt formatting. - **Streaming detokenization.** `StreamingDetokenizer` is fed one new id at a time and returns only the text that has become **complete UTF-8** — it never emits a broken multi-byte character or a partial byte-BPE sequence mid-stream @@ -140,10 +143,13 @@ than BPE mechanics: ## Supported families and the routing seam `BpeTokenizer::is_supported` returns `true` only for `model.type == "BPE"` **and** -a `ByteLevel` decoder. `Tokenizer::from_file` calls it and **throws** otherwise — -in particular a Metaspace / SentencePiece tokenizer is rejected rather than -silently mistokenized, because the byte-level pipeline and the hand-rolled -splitter are only correct for byte-level BPE. +a `ByteLevel` decoder. A SentencePiece-style BPE (`byte_fallback` + a +`ByteFallback` decoder — Gemma, AMALIA/EuroLLM) routes to `SpmBpeTokenizer` +(`tokenizer/spm.{h,cpp}`) instead, which also honors the Metaspace +pre-tokenizer's `prepend_scheme` ("always" prepends `▁` to every plain segment; +"first" is rejected as unimplemented) and a `Strip` decode step. +`Tokenizer::from_file` **throws** when no backend matches — an unsupported +family (e.g. Unigram/WordPiece) is rejected rather than silently mistokenized. Re-onboarding a non-Llama family therefore means (a) routing its tokenizer to an appropriate backend behind this same `is_supported` check, and (b) adding its diff --git a/reference/dump_ref.py b/reference/dump_ref.py index 11d8080..729a53a 100644 --- a/reference/dump_ref.py +++ b/reference/dump_ref.py @@ -152,6 +152,39 @@ "café ▁ leading metaspace", ], }, + # AMALIA-9B (European Portuguese, EuroLLM-9B base) is a plain LlamaForCausalLM + # served by the stock LlamaModel; what it uniquely exercises is the tokenizer: + # SentencePiece-BPE with a Metaspace pre-tokenizer (prepend_scheme "always" — + # a U+2581 is prepended to the input-leading plain segment before BPE, unlike + # Gemma) and a Strip decoder (one leading space removed on decode), plus + # ChatML with <|im_start|> (id 3) doubling as BOS via the post-processor. + # The 4bit repo is the only published MLX conversion. Compute dtype is + # **bfloat16**, NOT the usual fp16 cast: this family's residual stream + # overflows fp16 from layer 9, so an fp16 forward is all-NaN logits (and an + # fp16 reference is equally degenerate — NaN goldens gate nothing). The C++ + # engine mirrors this via ModelConfig::bf16_compute. The extra corpus pins + # the Metaspace prepend/Strip semantics where they are ambiguous: after + # special tokens, between adjacent specials, and around explicit spaces. + "amalia": { + "repo": "layerx-labs/AMALIA-9B-0626-DPO-MLX-4bit", + "fixtures": "fixtures_amalia", + "compute_dtype": mx.bfloat16, + # The AMALIA template injects a default Portuguese system message when the + # conversation has none — this chat prompt exercises exactly that path. + "chat_messages": [{"role": "user", "content": "Qual é a capital de Portugal?"}], + "thinking": False, + "extra_corpus": [ + "<|im_start|>user\nOlá, como estás?<|im_end|>", + "<|im_start|>system\nÉs a Amália.<|im_end|>\n<|im_start|>user\nOlá<|im_end|>\n" + "<|im_start|>assistant\n", + "a<|im_end|><|im_end|>b", + "<|im_end|> depois de especial com espaço", + "oi", + "não especial", + "português: ação, coração, à noite, pêssego", + "Amália Rodrigues cantava fado.", + ], + }, } # Fixed prompt set — committed so dumps are reproducible. Index 0 is the primary @@ -567,7 +600,9 @@ def main(): manifest = { "model_repo": MODEL_REPO, - "compute_dtype": "float16" if COMPUTE_DTYPE is mx.float16 else "quantized", + "compute_dtype": ("float16" if COMPUTE_DTYPE is mx.float16 + else "bfloat16" if COMPUTE_DTYPE is mx.bfloat16 + else "quantized"), "prompts": PROMPTS, "chat_messages": CHAT_MESSAGES, "greedy_max_new": GREEDY_MAX_NEW, @@ -579,6 +614,10 @@ def main(): def save(name, arr): """Eval an MLX array (or accept a numpy array) and write it as .npy.""" if isinstance(arr, mx.array): + # numpy has no bfloat16: store bf16 dumps as float32 (exact upcast; + # the C++ compare harness casts both sides to float32 anyway). + if arr.dtype == mx.bfloat16: + arr = arr.astype(mx.float32) mx.eval(arr) np_arr = np.array(arr) else: @@ -612,7 +651,8 @@ def save(name, arr): # byte-level BPE; tok.encode matches mlxforge::Tokenizer::encode for the family). # Dumped per-model so each tokenizer (Llama digit runs vs Qwen single digits) is # validated against its own oracle. - corpus = [{"text": s, "ids": [int(x) for x in tok.encode(s)]} for s in TOKENIZER_CORPUS] + corpus_texts = TOKENIZER_CORPUS + spec.get("extra_corpus", []) + corpus = [{"text": s, "ids": [int(x) for x in tok.encode(s)]} for s in corpus_texts] with open(os.path.join(FIXTURES_DIR, "tokenizer_corpus.json"), "w") as f: json.dump(corpus, f, ensure_ascii=False, indent=2) print(f" wrote tokenizer_corpus.json ({len(corpus)} strings)") diff --git a/reference/fixtures_amalia/argmax.npy b/reference/fixtures_amalia/argmax.npy new file mode 100644 index 0000000..3109b8c Binary files /dev/null and b/reference/fixtures_amalia/argmax.npy differ diff --git a/reference/fixtures_amalia/attn_norm0.npy b/reference/fixtures_amalia/attn_norm0.npy new file mode 100644 index 0000000..97b3746 Binary files /dev/null and b/reference/fixtures_amalia/attn_norm0.npy differ diff --git a/reference/fixtures_amalia/block0.npy b/reference/fixtures_amalia/block0.npy new file mode 100644 index 0000000..99a5dd8 Binary files /dev/null and b/reference/fixtures_amalia/block0.npy differ diff --git a/reference/fixtures_amalia/chat_ids.npy b/reference/fixtures_amalia/chat_ids.npy new file mode 100644 index 0000000..96b4d7f Binary files /dev/null and b/reference/fixtures_amalia/chat_ids.npy differ diff --git a/reference/fixtures_amalia/embeddings.npy b/reference/fixtures_amalia/embeddings.npy new file mode 100644 index 0000000..2255ef3 Binary files /dev/null and b/reference/fixtures_amalia/embeddings.npy differ diff --git a/reference/fixtures_amalia/greedy_gaps_kvq4.npy b/reference/fixtures_amalia/greedy_gaps_kvq4.npy new file mode 100644 index 0000000..5705d83 Binary files /dev/null and b/reference/fixtures_amalia/greedy_gaps_kvq4.npy differ diff --git a/reference/fixtures_amalia/greedy_gaps_kvq8.npy b/reference/fixtures_amalia/greedy_gaps_kvq8.npy new file mode 100644 index 0000000..9e2feea Binary files /dev/null and b/reference/fixtures_amalia/greedy_gaps_kvq8.npy differ diff --git a/reference/fixtures_amalia/greedy_tokens.npy b/reference/fixtures_amalia/greedy_tokens.npy new file mode 100644 index 0000000..63ecd5b Binary files /dev/null and b/reference/fixtures_amalia/greedy_tokens.npy differ diff --git a/reference/fixtures_amalia/greedy_tokens_kvq4.npy b/reference/fixtures_amalia/greedy_tokens_kvq4.npy new file mode 100644 index 0000000..63ecd5b Binary files /dev/null and b/reference/fixtures_amalia/greedy_tokens_kvq4.npy differ diff --git a/reference/fixtures_amalia/greedy_tokens_kvq8.npy b/reference/fixtures_amalia/greedy_tokens_kvq8.npy new file mode 100644 index 0000000..3ebda90 Binary files /dev/null and b/reference/fixtures_amalia/greedy_tokens_kvq8.npy differ diff --git a/reference/fixtures_amalia/k_rope0.npy b/reference/fixtures_amalia/k_rope0.npy new file mode 100644 index 0000000..1253652 Binary files /dev/null and b/reference/fixtures_amalia/k_rope0.npy differ diff --git a/reference/fixtures_amalia/logits_last.npy b/reference/fixtures_amalia/logits_last.npy new file mode 100644 index 0000000..e22bbb0 Binary files /dev/null and b/reference/fixtures_amalia/logits_last.npy differ diff --git a/reference/fixtures_amalia/manifest.json b/reference/fixtures_amalia/manifest.json new file mode 100644 index 0000000..cd7b5f1 --- /dev/null +++ b/reference/fixtures_amalia/manifest.json @@ -0,0 +1,149 @@ +{ + "model_repo": "layerx-labs/AMALIA-9B-0626-DPO-MLX-4bit", + "compute_dtype": "bfloat16", + "prompts": [ + "The capital of France is", + "Hello, world!", + "Once upon a time, in a land far away," + ], + "chat_messages": [ + { + "role": "user", + "content": "Qual \u00e9 a capital de Portugal?" + } + ], + "greedy_max_new": 20, + "arrays": { + "prompt_0_ids": { + "shape": [ + 6 + ], + "dtype": "int32" + }, + "prompt_1_ids": { + "shape": [ + 5 + ], + "dtype": "int32" + }, + "prompt_2_ids": { + "shape": [ + 12 + ], + "dtype": "int32" + }, + "chat_ids": { + "shape": [ + 71 + ], + "dtype": "int32" + }, + "embeddings": { + "shape": [ + 1, + 6, + 4096 + ], + "dtype": "float32" + }, + "attn_norm0": { + "shape": [ + 1, + 6, + 4096 + ], + "dtype": "float32" + }, + "q_pre0": { + "shape": [ + 1, + 32, + 6, + 128 + ], + "dtype": "float32" + }, + "q_rope0": { + "shape": [ + 1, + 32, + 6, + 128 + ], + "dtype": "float32" + }, + "k_rope0": { + "shape": [ + 1, + 8, + 6, + 128 + ], + "dtype": "float32" + }, + "v0": { + "shape": [ + 1, + 8, + 6, + 128 + ], + "dtype": "float32" + }, + "block0": { + "shape": [ + 1, + 6, + 4096 + ], + "dtype": "float32" + }, + "logits_last": { + "shape": [ + 1, + 128000 + ], + "dtype": "float32" + }, + "argmax": { + "shape": [ + 1 + ], + "dtype": "int32" + }, + "greedy_tokens": { + "shape": [ + 20 + ], + "dtype": "int32" + }, + "greedy_tokens_kvq8": { + "shape": [ + 20 + ], + "dtype": "int32" + }, + "greedy_gaps_kvq8": { + "shape": [ + 20 + ], + "dtype": "float32" + }, + "greedy_tokens_kvq4": { + "shape": [ + 20 + ], + "dtype": "int32" + }, + "greedy_gaps_kvq4": { + "shape": [ + 20 + ], + "dtype": "float32" + } + }, + "eos_token_ids": [ + 2, + 4 + ] +} \ No newline at end of file diff --git a/reference/fixtures_amalia/prompt_0_ids.npy b/reference/fixtures_amalia/prompt_0_ids.npy new file mode 100644 index 0000000..b968588 Binary files /dev/null and b/reference/fixtures_amalia/prompt_0_ids.npy differ diff --git a/reference/fixtures_amalia/prompt_1_ids.npy b/reference/fixtures_amalia/prompt_1_ids.npy new file mode 100644 index 0000000..ec36e01 Binary files /dev/null and b/reference/fixtures_amalia/prompt_1_ids.npy differ diff --git a/reference/fixtures_amalia/prompt_2_ids.npy b/reference/fixtures_amalia/prompt_2_ids.npy new file mode 100644 index 0000000..325c6d0 Binary files /dev/null and b/reference/fixtures_amalia/prompt_2_ids.npy differ diff --git a/reference/fixtures_amalia/q_pre0.npy b/reference/fixtures_amalia/q_pre0.npy new file mode 100644 index 0000000..ba33652 Binary files /dev/null and b/reference/fixtures_amalia/q_pre0.npy differ diff --git a/reference/fixtures_amalia/q_rope0.npy b/reference/fixtures_amalia/q_rope0.npy new file mode 100644 index 0000000..918454a Binary files /dev/null and b/reference/fixtures_amalia/q_rope0.npy differ diff --git a/reference/fixtures_amalia/tokenizer_corpus.json b/reference/fixtures_amalia/tokenizer_corpus.json new file mode 100644 index 0000000..9441077 --- /dev/null +++ b/reference/fixtures_amalia/tokenizer_corpus.json @@ -0,0 +1,597 @@ +[ + { + "text": "", + "ids": [ + 3 + ] + }, + { + "text": "The capital of France is Paris.", + "ids": [ + 3, + 806, + 5865, + 589, + 4472, + 656, + 5133, + 119735 + ] + }, + { + "text": "Hello, world!", + "ids": [ + 3, + 97849, + 119732, + 4437, + 119906 + ] + }, + { + "text": "don't I'll we've they're it's can't", + "ids": [ + 3, + 2072, + 119792, + 119721, + 606, + 119792, + 1539, + 950, + 119792, + 843, + 1538, + 119792, + 540, + 834, + 119792, + 119723, + 933, + 119792, + 119721 + ] + }, + { + "text": "DON'T SHOUT", + "ids": [ + 3, + 607, + 3419, + 119792, + 119766, + 577, + 21564, + 10722 + ] + }, + { + "text": "spaces here and more", + "ids": [ + 3, + 31078, + 668, + 8696, + 668, + 603, + 2511, + 1763 + ] + }, + { + "text": " leading and trailing ", + "ids": [ + 3, + 534, + 64442, + 603, + 924, + 13243, + 534 + ] + }, + { + "text": "tabs\tand\tmore\ttabs", + "ids": [ + 3, + 4710, + 119723, + 270, + 625, + 270, + 5699, + 270, + 5297, + 119723 + ] + }, + { + "text": "newlines\n\nand\r\nwindows\r\nendings", + "ids": [ + 3, + 2784, + 16002, + 271, + 271, + 625, + 274, + 271, + 14218, + 2252, + 274, + 271, + 664, + 1419 + ] + }, + { + "text": "mixed \n \t whitespace \n\n", + "ids": [ + 3, + 27899, + 119715, + 271, + 534, + 270, + 765, + 538, + 28711, + 119715, + 271, + 271 + ] + }, + { + "text": "1 12 123 1234 100000 3.14159", + "ids": [ + 3, + 119715, + 119749, + 119715, + 119749, + 119759, + 119715, + 119749, + 119759, + 119790, + 119715, + 119749, + 119759, + 119790, + 119799, + 119715, + 119749, + 119752, + 119752, + 119752, + 119752, + 119752, + 119715, + 119790, + 119735, + 119749, + 119799, + 119749, + 119791, + 119772 + ] + }, + { + "text": "snake_case camelCase kebab-case", + "ids": [ + 3, + 117435, + 119839, + 29638, + 4065, + 535, + 82449, + 566, + 885, + 675, + 119754, + 29638 + ] + }, + { + "text": "for (int i = 0; i < n; ++i) { sum += a[i]; }", + "ids": [ + 3, + 667, + 585, + 951, + 629, + 1631, + 119715, + 119752, + 119852, + 629, + 4594, + 544, + 119852, + 2568, + 119994, + 119718, + 119762, + 3812, + 3650, + 86018, + 520, + 119980, + 119718, + 41313, + 5522 + ] + }, + { + "text": "café naïve résumé Zürich", + "ids": [ + 3, + 24823, + 689, + 120354, + 843, + 5738, + 70260, + 39371 + ] + }, + { + "text": "Ünïcödé ßharp", + "ids": [ + 3, + 38291, + 88291, + 5884, + 119756, + 119715, + 120004, + 2901, + 119730 + ] + }, + { + "text": "你好世界，今天天气很好。", + "ids": [ + 3, + 119715, + 121037, + 120491, + 7594, + 500, + 449, + 401, + 43386, + 117046, + 120748, + 120491, + 119867 + ] + }, + { + "text": "こんにちは世界", + "ids": [ + 3, + 119715, + 120107, + 120257, + 119944, + 120522, + 119985, + 7594 + ] + }, + { + "text": "Привет мир", + "ids": [ + 3, + 4860, + 3328, + 22936 + ] + }, + { + "text": "emoji 😀 and 👨‍👩‍👧‍👦 family", + "ids": [ + 3, + 43917, + 1146, + 97129, + 603, + 119715, + 501, + 420, + 406, + 429, + 121125, + 501, + 420, + 406, + 430, + 121125, + 501, + 420, + 406, + 428, + 121125, + 501, + 420, + 406, + 427, + 6610 + ] + }, + { + "text": "math ∑∫√≠≤ symbols", + "ids": [ + 3, + 38076, + 119715, + 127701, + 127526, + 124416, + 125193, + 124006, + 50959 + ] + }, + { + "text": "<|begin_of_text|>hi<|eot_id|>", + "ids": [ + 3, + 4594, + 119920, + 10658, + 119839, + 894, + 119839, + 2309, + 119920, + 120057, + 3358, + 120155, + 119920, + 119716, + 602, + 119839, + 576, + 119920, + 120057 + ] + }, + { + "text": "<|start_header_id|>user<|end_header_id|>\n\nWhat?<|eot_id|>", + "ids": [ + 3, + 4594, + 119920, + 28300, + 119839, + 537, + 4841, + 119839, + 576, + 119920, + 120057, + 13676, + 120155, + 119920, + 664, + 119839, + 537, + 4841, + 119839, + 576, + 119920, + 120057, + 271, + 271, + 14913, + 119882, + 120155, + 119920, + 119716, + 602, + 119839, + 576, + 119920, + 120057 + ] + }, + { + "text": "a<|eot_id|><|eot_id|>b", + "ids": [ + 3, + 520, + 120155, + 119920, + 119716, + 602, + 119839, + 576, + 119920, + 37509, + 119920, + 119716, + 602, + 119839, + 576, + 119920, + 120057, + 119736 + ] + }, + { + "text": "URL: https://example.com/path?q=1&x=2#frag", + "ids": [ + 3, + 52800, + 119782, + 9301, + 3574, + 45298, + 119735, + 1456, + 119858, + 18488, + 119882, + 119764, + 119926, + 119749, + 119945, + 119778, + 119926, + 119759, + 120135, + 41328 + ] + }, + { + "text": "@user #hashtag $100 50% (parens) [brackets] {braces}", + "ids": [ + 3, + 7322, + 13676, + 3285, + 109878, + 2621, + 837, + 119749, + 119752, + 119752, + 119715, + 119791, + 119752, + 119941, + 585, + 1379, + 706, + 119762, + 1964, + 2695, + 1130, + 1849, + 119979, + 3812, + 2695, + 1122, + 119829 + ] + }, + { + "text": "<|im_start|>user\nOlá, como estás?<|im_end|>", + "ids": [ + 3, + 3, + 13676, + 271, + 119802, + 7188, + 119732, + 1227, + 68882, + 119882, + 4 + ] + }, + { + "text": "<|im_start|>system\nÉs a Amália.<|im_end|>\n<|im_start|>user\nOlá<|im_end|>\n<|im_start|>assistant\n", + "ids": [ + 3, + 3, + 21113, + 271, + 120105, + 119723, + 520, + 2098, + 23315, + 119735, + 4, + 271, + 3, + 13676, + 271, + 119802, + 7188, + 4, + 271, + 3, + 788, + 35441, + 271 + ] + }, + { + "text": "a<|im_end|><|im_end|>b", + "ids": [ + 3, + 520, + 4, + 4, + 119736 + ] + }, + { + "text": "<|im_end|> depois de especial com espaço", + "ids": [ + 3, + 4, + 9396, + 543, + 3936, + 643, + 22127 + ] + }, + { + "text": "oi", + "ids": [ + 3, + 1, + 2856, + 2 + ] + }, + { + "text": "não especial", + "ids": [ + 3, + 5, + 61922, + 3936 + ] + }, + { + "text": "português: ação, coração, à noite, pêssego", + "ids": [ + 3, + 31055, + 119782, + 28537, + 119732, + 45352, + 119732, + 898, + 25938, + 119732, + 90718, + 1536, + 1109 + ] + }, + { + "text": "Amália Rodrigues cantava fado.", + "ids": [ + 3, + 2098, + 23315, + 59636, + 933, + 14002, + 548, + 886, + 119735 + ] + } +] \ No newline at end of file diff --git a/reference/fixtures_amalia/v0.npy b/reference/fixtures_amalia/v0.npy new file mode 100644 index 0000000..42577cd Binary files /dev/null and b/reference/fixtures_amalia/v0.npy differ diff --git a/src/core/config.cpp b/src/core/config.cpp index 34007a9..a5a1004 100644 --- a/src/core/config.cpp +++ b/src/core/config.cpp @@ -203,8 +203,14 @@ ModelConfig ModelConfig::from_json(const nlohmann::json& j_top) { c.quantized = true; c.quant_group_size = it->value("group_size", c.quant_group_size); c.quant_bits = it->value("bits", c.quant_bits); + // mx::quantized_matmul assumes affine quantization; a non-affine checkpoint + // (e.g. mxfp4) would run to silent numerical garbage, so reject it at load. + const std::string mode = it->value("mode", "affine"); + if (mode != "affine") + throw std::runtime_error("config: unsupported quantization mode '" + mode + + "' (only 'affine' is implemented)"); for (const auto& [key, val] : it->items()) { - if (key == "group_size" || key == "bits") continue; // top-level defaults + if (key == "group_size" || key == "bits" || key == "mode") continue; // top-level defaults if (val.is_object() && val.contains("bits")) { QuantParams qp; qp.group_size = val.value("group_size", c.quant_group_size); @@ -213,6 +219,17 @@ ModelConfig ModelConfig::from_json(const nlohmann::json& j_top) { } } } + // Compute dtype. The engine runs fp16 by default (what every committed golden + // gates), but the EuroLLM-9B / AMALIA-9B family overflows fp16 — its residual + // stream exceeds fp16 range from layer 9, making an fp16 forward all-NaN — so + // it must run in bfloat16 (mlx-lm's native behavior for it). The config's + // declared dtype can't drive this (every MLX conversion says "bfloat16", + // including the fp16-proven Llama/Qwen checkpoints), so the family is + // identified by its architecture fingerprint. Gated end-to-end by + // reference/fixtures_amalia/. + c.bf16_compute = (c.vocab == 128000 && c.hidden == 4096 && c.intermediate_size == 12288 && + c.n_layers == 42); + // Multimodal vision tower (VLMs, e.g. Qwen3-VL). The ViT config and the vision // special-token ids live at the top level; M-RoPE parameters nest with the text // rope config (read from `j`). All absent for text-only models, which then keep diff --git a/src/core/config.h b/src/core/config.h index 750d75e..18396a5 100644 --- a/src/core/config.h +++ b/src/core/config.h @@ -126,6 +126,16 @@ struct ModelConfig { bool quantized = false; ///< Any weight is integer-quantized (logging/info). int quant_group_size = 64; ///< Default group size for weight quantization. int quant_bits = 4; ///< Default quantization bit-width (typically 4). + + // ----- Compute dtype ----- + /// True when the forward pass must run in bfloat16 instead of the engine's + /// fp16 default. Some checkpoints have activation magnitudes far beyond + /// fp16's range (65504) — EuroLLM-9B/AMALIA's residual stream blows up from + /// layer 9, so an fp16 forward is all-NaN logits (silent garbage, the exact + /// failure mode this engine exists to prevent). Set by from_json for the + /// known-overflowing family via its config fingerprint; every other + /// checkpoint keeps fp16, whose numerics all committed goldens gate. + bool bf16_compute = false; /// Per-module quant overrides, keyed by weight base (the key without the /// trailing ".weight", e.g. "model.layers.0.mlp.down_proj"). Empty => every /// quantized weight uses the defaults above. diff --git a/src/core/weights.cpp b/src/core/weights.cpp index 87ceb4c..ddb838d 100644 --- a/src/core/weights.cpp +++ b/src/core/weights.cpp @@ -117,21 +117,24 @@ std::string Weights::summary() const { } namespace { -// Merge one shard's tensors into `out`, applying sanitize + fp16 cast. +// Merge one shard's tensors into `out`, applying sanitize + compute-dtype cast. // `keep_vision` retains the ViT tower (VLMs) instead of dropping it. void absorb(std::unordered_map& out, - const std::unordered_map& shard, bool keep_vision) { + const std::unordered_map& shard, bool keep_vision, + mx::Dtype compute_dtype) { for (const auto& [raw, arr] : shard) { auto canon = sanitize_key(raw, keep_vision); if (!canon) continue; // dropped buffer - // Cast only floating tensors to fp16; packed 4-bit weights (uint32) and - // other integer tensors are kept as-is. Exception: Gated-DeltaNet's `A_log` - // (the per-head decay log-rate) is kept in its source fp32 — the recurrence - // exponentiates it twice (exp(-exp(A_log))) and fp16 there visibly drifts the - // decay. Mirrors mlx_lm's cast_predicate, which excludes A_log from the cast. + // Cast only floating tensors to the compute dtype (fp16, or bf16 for + // fp16-overflowing checkpoints — see ModelConfig::bf16_compute); packed + // 4-bit weights (uint32) and other integer tensors are kept as-is. + // Exception: Gated-DeltaNet's `A_log` (the per-head decay log-rate) is kept + // in its source fp32 — the recurrence exponentiates it twice + // (exp(-exp(A_log))) and fp16 there visibly drifts the decay. Mirrors + // mlx_lm's cast_predicate, which excludes A_log from the cast. const bool keep_dtype = ends_with(*canon, ".A_log"); mx::array value = (!keep_dtype && mx::issubdtype(arr.dtype(), mx::floating)) - ? mx::astype(arr, mx::float16) + ? mx::astype(arr, compute_dtype) : arr; out.emplace(*canon, value); } @@ -212,6 +215,9 @@ Weights load_weights(const std::string& model_dir, const ModelConfig& cfg) { // Keep the ViT tower only for vision-language checkpoints; text-only models // drop it so the load stays lean. const bool keep_vision = cfg.has_vision_tower(); + const mx::Dtype cast_dtype = compute_dtype(cfg); + if (cfg.bf16_compute) + log::info("weights: bf16 compute (fp16 overflows for this checkpoint family)"); // Prefer the sharded layout, but only when every shard the index names is // actually present: some mlx-community exports ship a single consolidated @@ -232,7 +238,8 @@ Weights load_weights(const std::string& model_dir, const ModelConfig& cfg) { log::debug("weights: sharded checkpoint, {} files", files.size()); for (const auto& file : files) { log::debug("weights: loading shard {}", file); - absorb(w.tensors, mx::load_safetensors(model_dir + "/" + file).first, keep_vision); + absorb(w.tensors, mx::load_safetensors(model_dir + "/" + file).first, keep_vision, + cast_dtype); } loaded = true; } else { @@ -245,16 +252,17 @@ Weights load_weights(const std::string& model_dir, const ModelConfig& cfg) { throw std::runtime_error("weights: no model.safetensors[.index.json] in '" + model_dir + "'"); } log::debug("weights: single-file checkpoint"); - absorb(w.tensors, mx::load_safetensors(single).first, keep_vision); + absorb(w.tensors, mx::load_safetensors(single).first, keep_vision, cast_dtype); } - std::size_t non_fp16 = 0; + std::size_t non_compute = 0; for (const auto& [_, a] : w.tensors) - if (a.dtype() != mx::float16) ++non_fp16; - log::info("weights: loaded {} tensors from '{}' ({} non-fp16)", w.tensors.size(), model_dir, - non_fp16); - if (non_fp16 > 0) - log::warn("weights: {} tensors are not fp16 (expected for quantized models)", non_fp16); + if (a.dtype() != cast_dtype) ++non_compute; + log::info("weights: loaded {} tensors from '{}' ({} not in the compute dtype)", + w.tensors.size(), model_dir, non_compute); + if (non_compute > 0) + log::warn("weights: {} tensors are not in the compute dtype (expected for quantized models)", + non_compute); normalize_backbone_root_keys(w); // embedding checkpoints: backbone-root -> model.* stack_moe_experts(w, cfg); // raw per-expert MoE tensors -> stacked switch_mlp diff --git a/src/core/weights.h b/src/core/weights.h index 6ecf1c1..047726d 100644 --- a/src/core/weights.h +++ b/src/core/weights.h @@ -56,10 +56,19 @@ struct Weights { std::string summary() const; }; +// The floating dtype the forward pass runs in — and everything that must match +// it: the on-load weight cast, the batched additive mask, and the prefill logit +// placeholders. bfloat16 for fp16-overflowing checkpoints +// (ModelConfig::bf16_compute, e.g. EuroLLM/AMALIA), fp16 for every other model. +inline mlx::core::Dtype compute_dtype(const ModelConfig& cfg) { + return cfg.bf16_compute ? mlx::core::bfloat16 : mlx::core::float16; +} + // Load every weight tensor from a model directory, applying sanitize and casting -// to fp16. `cfg` supplies the quant params (defaults + per-module overrides) used -// to populate Weights::quant for each quantized tensor. Throws if neither an -// index JSON nor model.safetensors is found. +// to the compute dtype (fp16, or bf16 — see compute_dtype). `cfg` supplies the +// quant params (defaults + per-module overrides) used to populate Weights::quant +// for each quantized tensor. Throws if neither an index JSON nor +// model.safetensors is found. Weights load_weights(const std::string& model_dir, const ModelConfig& cfg); } // namespace mlxforge diff --git a/src/model/decoder_model.cpp b/src/model/decoder_model.cpp index aac1de0..81a3aab 100644 --- a/src/model/decoder_model.cpp +++ b/src/model/decoder_model.cpp @@ -307,9 +307,11 @@ mx::array DecoderModel::batch_mask(int prev_idx, int n_query, const mx::array& l mx::reshape(kpos, {1, 1, 1, t_kv})); mx::array keep = mx::logical_and(causal, valid); // -> (B, 1, N, T_kv) - // Additive fp16 mask (avoid boolean masks; #2894). + // Additive mask in the compute dtype (avoid boolean masks, #2894; a dtype + // mismatch with the scores would silently promote the whole SDPA). + const mx::Dtype dt = compute_dtype(cfg_); const float ninf = -std::numeric_limits::infinity(); - return mx::where(keep, mx::array(0.0f, mx::float16), mx::array(ninf, mx::float16)); + return mx::where(keep, mx::array(0.0f, dt), mx::array(ninf, dt)); } mx::array DecoderModel::attention_batched(const mx::array& x, int layer, const mx::array& offset, diff --git a/src/runtime/batching.cpp b/src/runtime/batching.cpp index 4311e37..b2d40ac 100644 --- a/src/runtime/batching.cpp +++ b/src/runtime/batching.cpp @@ -33,7 +33,8 @@ PrefillResult prefill(const DecoderModel& model, const std::vector 0) throw std::runtime_error("KV-cache quantization is not supported for hybrid (Qwen3.5) models"); + // quantized_sdpa's mask/underflow conventions are proven for fp16 compute + // only; bf16-compute checkpoints (AMALIA) have no quantized-KV golden yet. + if (mc.bf16_compute) + throw std::runtime_error( + "KV-cache quantization is not supported for bf16-compute models (AMALIA)"); return {ec.kv_bits, ec.kv_group_size}; } @@ -83,6 +88,10 @@ PrefixCacheConfig validate_prefix_cache(const EngineConfig& ec, const ModelConfi throw std::runtime_error("the prefix cache is not supported for vision-language models"); if (mc.full_attention_interval > 0) throw std::runtime_error("the prefix cache is not supported for hybrid (Qwen3.5) models"); + // The warm==cold exact-token gate and the SSD serializer (fp16 storage codes) + // only cover fp16 compute; bf16-compute checkpoints (AMALIA) are ungated. + if (mc.bf16_compute) + throw std::runtime_error("the prefix cache is not supported for bf16-compute models (AMALIA)"); PrefixCacheConfig pc; pc.enabled = true; pc.block_size = bs; @@ -140,7 +149,7 @@ Engine::Loaded Engine::load_head(const std::string& spec, const std::string& rop out.tokenizer = Tokenizer::from_file( out.dir + "/tokenizer.json", out.config.bos_token_id, - chat_format_from_model_type(out.config.model_type) + chat_format_for_model_dir(out.dir, out.config.model_type) ); } diff --git a/src/tokenizer/spm.cpp b/src/tokenizer/spm.cpp index 6389022..6e335a4 100644 --- a/src/tokenizer/spm.cpp +++ b/src/tokenizer/spm.cpp @@ -66,14 +66,15 @@ void SpmBpeTokenizer::emit_symbol(const std::string& sym, std::vector& out) } } -void SpmBpeTokenizer::encode_plain(const std::string& segment, std::vector& out) const { +void SpmBpeTokenizer::encode_plain(const std::string& segment, bool at_text_start, + std::vector& out) const { if (segment.empty()) return; // Normalize: every space becomes the metaspace marker. (The vestigial space // pre-tokenizer then finds nothing to split, so the whole segment is one BPE // word — matching the HF pipeline for Gemma-style tokenizers.) std::string normalized; - normalized.reserve(segment.size()); + normalized.reserve(segment.size() + 3); for (char ch : segment) { if (ch == ' ') normalized += kMetaspace; @@ -81,6 +82,14 @@ void SpmBpeTokenizer::encode_plain(const std::string& segment, std::vector& normalized.push_back(ch); } + // prepend_scheme "always": a leading ▁ is added, but — settled empirically + // against the HF tokenizer (AMALIA corpus fixture) — only to the segment at + // the very start of the input (never to segments following a special token), + // and never doubled onto a segment that already starts with the marker + // (leading spaces become ▁ via the replace above). + if (prepend_metaspace_ && at_text_start && normalized.compare(0, 3, kMetaspace) != 0) + normalized.insert(0, kMetaspace); + std::vector symbols = utf8_chars(normalized); // Repeatedly merge the lowest-rank adjacent pair (merging all of its @@ -128,7 +137,7 @@ std::vector SpmBpeTokenizer::encode(const std::string& text) const { if (special_first_bytes_.count(static_cast(text[i]))) { for (const auto& [lit, id] : special_tokens_) { // sorted longest-first if (i + lit.size() <= text.size() && text.compare(i, lit.size(), lit) == 0) { - encode_plain(text.substr(run_start, i - run_start), out); + encode_plain(text.substr(run_start, i - run_start), /*at_text_start=*/run_start == 0, out); out.push_back(id); i += lit.size(); run_start = i; @@ -139,7 +148,7 @@ std::vector SpmBpeTokenizer::encode(const std::string& text) const { } if (!matched) ++i; } - encode_plain(text.substr(run_start), out); + encode_plain(text.substr(run_start), /*at_text_start=*/run_start == 0, out); return out; } @@ -171,7 +180,12 @@ std::string SpmBpeTokenizer::decode(const std::vector& ids) const { } } out += pending_bytes; - return out; + // The Strip decoder runs after the per-token pieces are fused into one string: + // drop up to N leading spaces, undoing the Metaspace prepend. + int strip = strip_leading_spaces_; + size_t start = 0; + while (strip-- > 0 && start < out.size() && out[start] == ' ') ++start; + return start == 0 ? out : out.substr(start); } bool SpmBpeTokenizer::is_supported(const std::string& tokenizer_json) { @@ -236,6 +250,31 @@ SpmBpeTokenizer SpmBpeTokenizer::from_blob(const std::string& tokenizer_json) { if (it != t.token_to_id_.end()) t.unk_id_ = it->second; } + // Metaspace pre-tokenizer prepend_scheme: "always" prepends ▁ to every plain + // segment (LlamaTokenizer-style, e.g. AMALIA/EuroLLM); "never"/absent matches + // Gemma. "first" (prepend only to the segment at offset 0) is rejected rather + // than approximated — a wrong prepend is silent id divergence, not a crash. + if (auto pt = j.find("pre_tokenizer"); pt != j.end() && pt->is_object() && + pt->value("type", std::string()) == "Metaspace") { + const std::string scheme = pt->value("prepend_scheme", std::string("never")); + if (scheme == "always") + t.prepend_metaspace_ = true; + else if (scheme != "never") + throw std::runtime_error("spm: Metaspace prepend_scheme '" + scheme + "' not implemented"); + } + + // A Strip entry in the decoder Sequence removes N leading spaces from the + // fused decode output (the inverse of the Metaspace prepend). + if (auto dec = j.find("decoder"); dec != j.end() && dec->is_object()) { + if (auto ds = dec->find("decoders"); ds != dec->end() && ds->is_array()) { + for (const auto& d : *ds) { + if (d.value("type", std::string()) == "Strip" && + d.value("content", std::string()) == " ") + t.strip_leading_spaces_ = d.value("start", 0); + } + } + } + // Added tokens: HF isolates EVERY added token atomically before BPE, whether or // not it is `special`. So segment on all of them; the `special` flag only // governs decode-skipping (special_ids_). diff --git a/src/tokenizer/spm.h b/src/tokenizer/spm.h index 8f70304..8dfda86 100644 --- a/src/tokenizer/spm.h +++ b/src/tokenizer/spm.h @@ -3,9 +3,13 @@ // a tokenizer.json with `model.type == "BPE"` and `byte_fallback == true`: // // special-token segmentation -> Metaspace normalization (every ' ' -> U+2581 -// "▁") -> BPE merges over the whole segment's Unicode characters -> vocab -// lookup, with byte_fallback (a character absent from the vocab is emitted as -// its UTF-8 bytes via the "<0xNN>" byte tokens). +// "▁"; under prepend_scheme "always" a ▁ is also prepended to the segment +// that opens the input, the LlamaTokenizer/EuroLLM/AMALIA convention — +// Gemma's scheme never prepends) -> BPE merges over the whole segment's +// Unicode characters -> vocab lookup, with byte_fallback (a character absent +// from the vocab is emitted as its UTF-8 bytes via the "<0xNN>" byte tokens). +// Decode honors a trailing Strip decoder (drop N leading spaces), undoing the +// prepend. // // This is a distinct family from the byte-level BPE (tokenizer/bpe.h): there is // no GPT-2 byte->unicode remapping and no regex pre-tokenizer; spaces become the @@ -53,9 +57,10 @@ class SpmBpeTokenizer : public EncoderBackend { const std::unordered_set& special_ids() const override { return special_ids_; } private: - // Normalize a special-token-free segment (space -> metaspace), then run the BPE - // merge loop over its Unicode characters and append the ids. - void encode_plain(const std::string& segment, std::vector& out) const; + // Normalize a special-token-free segment (space -> metaspace; under + // prepend_scheme "always" a leading ▁ when the segment opens the input), then + // run the BPE merge loop over its Unicode characters and append the ids. + void encode_plain(const std::string& segment, bool at_text_start, std::vector& out) const; // Map one final BPE symbol to ids: its vocab id, or byte_fallback to "<0xNN>" // byte tokens, or the unk id. void emit_symbol(const std::string& sym, std::vector& out) const; @@ -69,6 +74,14 @@ class SpmBpeTokenizer : public EncoderBackend { std::unordered_map byte_token_value_; // "<0xNN>" id -> byte value (decode) int unk_id_ = -1; + // Metaspace prepend_scheme "always": prepend ▁ to the input-leading plain + // segment before BPE (AMALIA/EuroLLM — the exact HF semantics are pinned by + // the corpus fixture); false for Gemma, whose scheme never prepends. + bool prepend_metaspace_ = false; + // Strip decoder: number of leading spaces removed from the fused decode output + // (undoes the prepended ▁; 0 when the decoder chain has no Strip, e.g. Gemma). + int strip_leading_spaces_ = 0; + // (literal, id) for every added-token, sorted by descending literal length so // longest-match wins during segmentation; first bytes are a cheap pre-filter. std::vector> special_tokens_; diff --git a/src/tokenizer/tokenizer.cpp b/src/tokenizer/tokenizer.cpp index a454384..98ec19b 100644 --- a/src/tokenizer/tokenizer.cpp +++ b/src/tokenizer/tokenizer.cpp @@ -35,6 +35,15 @@ std::string load_file(const std::string& path) { return ss.str(); } +// Read a file that may legitimately be absent; empty string when it is. +std::string read_file_or_empty(const std::string& path) { + std::ifstream f(path, std::ios::binary); + if (!f) return ""; + std::ostringstream ss; + ss << f.rdbuf(); + return ss.str(); +} + // Today's date as "01 Jun 2026" (the format the Llama-3.2 template uses). std::string current_date() { std::time_t t = std::time(nullptr); @@ -122,8 +131,39 @@ std::shared_ptr make_backend(const std::string& path, const std: "' is not a supported tokenizer (only byte-level BPE " "(Llama-3.2/Qwen) and SentencePiece-BPE (Gemma) are implemented)"); } + +// The default system message AMALIA's chat template injects when the +// conversation carries no system turn. Its opening words also serve as the +// template's identifying marker in chat_format_for_model_dir. +constexpr const char* kAmaliaSystemPrompt = + "O teu nome é Amália, e és um modelo avançado de linguagem útil. Responde sempre na língua " + "do utilizador, a menos que sejas instruído em contrário, e lembra-te que a tua língua " + "principal é o português europeu."; } // namespace +ChatFormat chat_format_for_model_dir(const std::string& model_dir, const std::string& model_type) { + // The on-disk chat template: the standalone chat_template.jinja (transformers + // >= 4.51 convention, what AMALIA ships) or the legacy "chat_template" string + // inside tokenizer_config.json. + std::string tmpl = read_file_or_empty(model_dir + "/chat_template.jinja"); + if (tmpl.empty()) { + const std::string cfg_blob = read_file_or_empty(model_dir + "/tokenizer_config.json"); + if (!cfg_blob.empty()) { + nlohmann::json j = nlohmann::json::parse(cfg_blob, /*cb=*/nullptr, /*allow_exceptions=*/false); + if (!j.is_discarded() && j.contains("chat_template") && j["chat_template"].is_string()) + tmpl = j["chat_template"].get(); + } + } + if (tmpl.find("O teu nome é Amália") != std::string::npos) { + log::info("tokenizer: AMALIA chat template detected in '{}'", model_dir); + return ChatFormat::Amalia; + } + // Every family ships some template (Llama-3.2's and Qwen3's are what the + // Llama3/Qwen3 renderers implement), so an unmatched marker is the normal + // case, not an error: defer to the model_type mapping as before. + return chat_format_from_model_type(model_type); +} + Tokenizer Tokenizer::from_file(const std::string& tokenizer_json_path, int bos_id, ChatFormat fmt) { const std::string blob = load_file(tokenizer_json_path); Tokenizer t; @@ -302,6 +342,28 @@ std::string render_qwen3(const std::vector& messages, return os.str(); } +// AMALIA ChatML template, rendered in full (mirrors the checkpoint's +// chat_template.jinja; no tools/thinking handling — the template has none). +// When the conversation has no leading system turn, the default Portuguese +// system message is injected. Unlike Llama-3.2 the leading special token is NOT +// left to the encoder: apply_chat_template encodes this WITHOUT the BOS prepend +// (HF uses add_special_tokens=False), because AMALIA's BOS id 3 IS +// `<|im_start|>` and prepending it would double the opening tag — and the +// tokenizer's ▁-prepend applies only to input-leading text, so the template +// must start with the special token for "system" to tokenize like the +// reference. Gated byte-exact by fixtures_amalia/chat_ids.npy. +std::string render_amalia(const std::vector& messages, + bool add_generation_prompt) { + std::ostringstream os; + const bool have_system = !messages.empty() && messages.front().role == "system"; + if (!have_system) + os << "<|im_start|>system\n" << kAmaliaSystemPrompt << "<|im_end|>\n"; + for (const auto& m : messages) + os << "<|im_start|>" << m.role << "\n" << m.content << "<|im_end|>\n"; + if (add_generation_prompt) os << "<|im_start|>assistant\n"; + return os.str(); +} + } // namespace std::string Tokenizer::render_chat_template(const std::vector& messages, @@ -312,6 +374,7 @@ std::string Tokenizer::render_chat_template(const std::vector& messages if (fmt == ChatFormat::Qwen3 || fmt == ChatFormat::Qwen35) return render_qwen3(messages, add_generation_prompt, enable_thinking, tools, /*open_think=*/fmt == ChatFormat::Qwen35); + if (fmt == ChatFormat::Amalia) return render_amalia(messages, add_generation_prompt); return render_llama3(messages, add_generation_prompt, today_date, tools); } @@ -320,8 +383,13 @@ std::vector Tokenizer::apply_chat_template(const std::vector& mess const std::string& today_date, const std::vector& tools, bool enable_thinking) const { - return encode(render_chat_template(messages, add_generation_prompt, today_date, chat_format_, - tools, enable_thinking)); + const std::string text = render_chat_template(messages, add_generation_prompt, today_date, + chat_format_, tools, enable_thinking); + // AMALIA's template opens with its own `<|im_start|>` — the very token the + // encoder would prepend as BOS (id 3) — so encode it bare, mirroring HF's + // apply_chat_template (add_special_tokens=False). See render_amalia. + if (chat_format_ == ChatFormat::Amalia) return impl_->encode(text); + return encode(text); } std::string StreamingDetokenizer::add(int id) { diff --git a/src/tokenizer/tokenizer.h b/src/tokenizer/tokenizer.h index dab46b5..e3d9629 100644 --- a/src/tokenizer/tokenizer.h +++ b/src/tokenizer/tokenizer.h @@ -20,12 +20,23 @@ namespace mlxforge { // forward pass (shared) and the prompt formatting (per-family) stay decoupled. // Qwen35 is ChatML like Qwen3 but, with thinking enabled, opens the reasoning // block (`\n`) in the generation prompt instead of leaving it to the model. -enum class ChatFormat { Llama3, Qwen3, Qwen35 }; +// Amalia (AMALIA-9B, model_type "llama") is plain ChatML with a default +// European-Portuguese system message — identified from the checkpoint's chat +// template on disk, since model_type cannot distinguish it from Llama. +enum class ChatFormat { Llama3, Qwen3, Qwen35, Amalia }; // Map a config.json model_type to a chat format: "qwen3"/"qwen3_moe"/"qwen2" -> // Qwen3, "qwen3_5" -> Qwen35 (both ChatML), everything else -> Llama3. ChatFormat chat_format_from_model_type(const std::string& model_type); +// Chat format for a resolved model directory: reads chat_template.jinja (or the +// "chat_template" string in tokenizer_config.json) and recognizes known +// templates that model_type alone cannot identify (AMALIA is model_type "llama" +// but ChatML). An unrecognized on-disk template is the normal case (Llama/Qwen +// ship their own, handled by their renderers), so it quietly falls back to +// chat_format_from_model_type. +ChatFormat chat_format_for_model_dir(const std::string& model_dir, const std::string& model_type); + class Tokenizer { public: struct Message { diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 6ff044b..67c12c7 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -40,6 +40,7 @@ add_executable(mlxforge_tests runtime/bucketing_test.cpp scheduler/validation_test.cpp model/qwen3_test.cpp + model/amalia_test.cpp model/rope_scaling_test.cpp model/qwen3_yarn_test.cpp model/qwen3_moe_test.cpp @@ -141,6 +142,15 @@ if(_mlxforge_qwen3_embedding_weights) "${_mlxforge_qwen3_embedding_weight0}" DIRECTORY) endif() +# AMALIA-9B model snapshot for the plain-Llama / SentencePiece-Metaspace / ChatML +# integration tests (empty -> those tests self-skip). The 4bit checkpoint is ~5GB. +file(GLOB _mlxforge_model_snapshots_amalia + "$ENV{HOME}/.cache/huggingface/hub/models--layerx-labs--AMALIA-9B-0626-DPO-MLX-4bit/snapshots/*") +set(MLXFORGE_MODEL_DIR_AMALIA "") +if(_mlxforge_model_snapshots_amalia) + list(GET _mlxforge_model_snapshots_amalia 0 MLXFORGE_MODEL_DIR_AMALIA) +endif() + # Gemma-2 tokenizer dir for the SentencePiece-BPE integration test (empty -> it # self-skips). Only the ungated mirror's tokenizer files are needed (no weights). file(GLOB _mlxforge_model_snapshots_gemma @@ -170,6 +180,7 @@ target_compile_definitions(mlxforge_tests PRIVATE MLXFORGE_REF_FIXTURES_DIR_QWEN3_VL="${CMAKE_SOURCE_DIR}/reference/fixtures_qwen3_vl" MLXFORGE_REF_FIXTURES_DIR_QWEN3_EMBEDDING="${CMAKE_SOURCE_DIR}/reference/fixtures_qwen3_embedding" MLXFORGE_REF_FIXTURES_DIR_GEMMA="${CMAKE_SOURCE_DIR}/reference/fixtures_gemma" + MLXFORGE_REF_FIXTURES_DIR_AMALIA="${CMAKE_SOURCE_DIR}/reference/fixtures_amalia" MLXFORGE_MODEL_DIR="${MLXFORGE_MODEL_DIR}" MLXFORGE_MODEL_DIR_4BIT="${MLXFORGE_MODEL_DIR_4BIT}" MLXFORGE_MODEL_DIR_QWEN3="${MLXFORGE_MODEL_DIR_QWEN3}" @@ -178,6 +189,7 @@ target_compile_definitions(mlxforge_tests PRIVATE MLXFORGE_MODEL_DIR_QWEN3_VL="${MLXFORGE_MODEL_DIR_QWEN3_VL}" MLXFORGE_MODEL_DIR_QWEN3_EMBEDDING="${MLXFORGE_MODEL_DIR_QWEN3_EMBEDDING}" MLXFORGE_MODEL_DIR_GEMMA="${MLXFORGE_MODEL_DIR_GEMMA}" + MLXFORGE_MODEL_DIR_AMALIA="${MLXFORGE_MODEL_DIR_AMALIA}" MLXFORGE_GGUF_MODEL="${MLXFORGE_GGUF_MODEL}") doctest_discover_tests(mlxforge_tests) diff --git a/tests/core/config_test.cpp b/tests/core/config_test.cpp index 66007bc..40704de 100644 --- a/tests/core/config_test.cpp +++ b/tests/core/config_test.cpp @@ -63,6 +63,41 @@ TEST_CASE("ModelConfig ignores unknown/extra keys") { CHECK(c.eos_token_ids.empty()); } +TEST_CASE("ModelConfig selects bf16 compute for the EuroLLM/AMALIA fingerprint only") { + // AMALIA-9B / EuroLLM-9B geometry -> bf16 (fp16 overflows from layer 9). + auto j = nlohmann::json::parse(R"({ + "num_hidden_layers": 42, "hidden_size": 4096, "num_attention_heads": 32, + "num_key_value_heads": 8, "vocab_size": 128000, "intermediate_size": 12288, + "rope_theta": 1000000.0, "rms_norm_eps": 1e-5 + })"); + CHECK(ModelConfig::from_json(j).bf16_compute); + + // Any other geometry (here: the Llama-3.2-1B reference config) stays on the + // engine's fp16 default, which every committed golden gates. + ModelConfig llama = ModelConfig::from_file(fixture("config_llama32_1b.json")); + CHECK_FALSE(llama.bf16_compute); +} + +TEST_CASE("ModelConfig accepts affine quantization mode and rejects others") { + // AMALIA-style block: an explicit "mode": "affine" parses like a bare block. + auto j = nlohmann::json::parse(R"({ + "num_hidden_layers": 2, "hidden_size": 8, "num_attention_heads": 4, + "num_key_value_heads": 2, "vocab_size": 100, "intermediate_size": 16, + "rope_theta": 10000.0, "rms_norm_eps": 1e-6, + "quantization": {"group_size": 64, "bits": 4, "mode": "affine"} + })"); + ModelConfig c = ModelConfig::from_json(j); + CHECK(c.quantized); + CHECK(c.quant_group_size == 64); + CHECK(c.quant_bits == 4); + + // A non-affine mode would run mx::quantized_matmul on the wrong math — reject. + j["quantization"]["mode"] = "mxfp4"; + CHECK_THROWS_WITH_AS(ModelConfig::from_json(j), + "config: unsupported quantization mode 'mxfp4' (only 'affine' is implemented)", + std::runtime_error); +} + TEST_CASE("ModelConfig parses the Qwen3.5 hybrid text_config") { // A representative subset of mlx-community/Qwen3.5-0.8B-4bit's config.json: the // text hyperparameters nest under "text_config", rope_theta lives in a diff --git a/tests/model/amalia_test.cpp b/tests/model/amalia_test.cpp new file mode 100644 index 0000000..c36f67b --- /dev/null +++ b/tests/model/amalia_test.cpp @@ -0,0 +1,185 @@ +// AMALIA-9B golden-reference checks. The forward pass is the plain Llama +// decoder run in **bfloat16** (the family's residual stream overflows fp16 +// from layer 9 — an fp16 forward is all-NaN logits, so bf16_compute is part of +// what these gates prove, at 42 layers / GQA 8 / rope_theta 1e6 through 4-bit +// quantized projections). What AMALIA otherwise uniquely exercises is the +// tokenizer: SentencePiece-BPE with a Metaspace prepend +// (`prepend_scheme: "always"`), a Strip decoder, and ChatML where <|im_start|> +// (id 3) doubles as BOS. All vs the fixtures dumped by +// `reference/dump_ref.py --model amalia`. Self-skips unless both the model +// (MLXFORGE_MODEL_DIR_AMALIA) and the committed fixtures are present. +// Quantized-KV fixtures are dumped for uniformity but not gated here — the +// engine rejects kv-quant for bf16-compute models. +#include + +#include +#include +#include + +#include + +#include "core/config.h" +#include "mlx/ops.h" +#include "support/model_fixture.h" +#include "support/reference.h" +#include "tokenizer/tokenizer.h" + +using namespace mlxforge::test; + +namespace { +bool amalia_fixtures_present() { return std::ifstream(amalia_ref_path("manifest.json")).good(); } +bool amalia_ready() { return amalia_model_available() && amalia_fixtures_present(); } + +std::string read_file(const std::string& path) { + std::ifstream f(path, std::ios::binary); + std::ostringstream ss; + ss << f.rdbuf(); + return ss.str(); +} +} // namespace + +TEST_CASE("AMALIA: plain-Llama attention front-half matches the reference") { + if (!amalia_ready()) { + MESSAGE("AMALIA model/fixtures not present; skipping golden-reference check"); + return; + } + mlxforge::LlamaModel& model = shared_amalia_model(); + + // The factory contract: AMALIA is the stock decoder (no QK-Norm), loaded 4-bit + // with bf16 compute (the fingerprint in ModelConfig::from_json — fp16 would be + // all-NaN from layer 9, and these gates would then fail on nonfinite values). + CHECK(model.config().bf16_compute); + CHECK(model.weights().at("model.layers.0.input_layernorm.weight").dtype() == mx::bfloat16); + CHECK_FALSE(model.weights().has("model.layers.0.self_attn.q_norm.weight")); + mlxforge::QuantParams qp; + CHECK(model.weights().is_quantized("model.layers.0.self_attn.q_proj.weight", qp)); + CHECK(qp.bits == 4); + CHECK(qp.group_size == 64); + + std::vector ids = load_amalia_token_ids("prompt_0_ids.npy"); + mx::array tokens(ids.data(), {1, static_cast(ids.size())}, mx::int32); + + mx::array emb = model.embed(tokens); + assert_close(emb, load_amalia_npy("embeddings.npy")); + + mx::array normed = + model.rms_norm(emb, model.weights().at("model.layers.0.input_layernorm.weight")); + assert_close(normed, load_amalia_npy("attn_norm0.npy")); + + // RoPE'd Q/K and un-roped V through the quantized projections (gates + // rope_theta 1e6 with no scaling, GQA 32/8, head_dim 128). + mlxforge::DecoderModel::QKV qkv = model.attn_qkv(emb, /*layer=*/0); + assert_close(qkv.q, load_amalia_npy("q_rope0.npy")); + assert_close(qkv.k, load_amalia_npy("k_rope0.npy")); + assert_close(qkv.v, load_amalia_npy("v0.npy")); +} + +TEST_CASE("AMALIA: decoder block 0 output matches the reference") { + if (!amalia_ready()) { + MESSAGE("AMALIA model/fixtures not present; skipping golden-reference check"); + return; + } + mlxforge::LlamaModel& model = shared_amalia_model(); + std::vector ids = load_amalia_token_ids("prompt_0_ids.npy"); + mx::array tokens(ids.data(), {1, static_cast(ids.size())}, mx::int32); + mx::array emb = model.embed(tokens); + + assert_close(model.decoder_block(emb, /*layer=*/0), load_amalia_npy("block0.npy")); +} + +TEST_CASE("AMALIA: full forward logits + first-token argmax match the reference") { + if (!amalia_ready()) { + MESSAGE("AMALIA model/fixtures not present; skipping golden-reference check"); + return; + } + mlxforge::LlamaModel& model = shared_amalia_model(); + std::vector ids = load_amalia_token_ids("prompt_0_ids.npy"); + const int T = static_cast(ids.size()); + mx::array tokens(ids.data(), {1, T}, mx::int32); + + mx::array logits = model.forward(tokens); // (1, T, vocab) + const int vocab = logits.shape()[2]; + mx::array last = mx::reshape(mx::slice(logits, {0, T - 1, 0}, {1, T, vocab}), {1, vocab}); + // Deliberately loose: 42 layers of bf16 (8 mantissa bits) over 4-bit + // quantized matmuls, which are fusion-context-sensitive (see the kv-quant + // notes — mlx-lm disagrees with itself across graph contexts), so a tight + // raw-logit bound would assert kernel fusion, not math. This is a + // sanity/finiteness check; the exact argmax below and the exact greedy + // stream in the next case are the correctness gates. + assert_close(last, load_amalia_npy("logits_last.npy"), /*rtol=*/1e-1f, /*atol=*/5e-1f); + + std::vector argmax = load_amalia_token_ids("argmax.npy"); + mx::array got = mx::astype(mx::argmax(last, /*axis=*/-1), mx::int32); + mx::eval(got); + CHECK(got.data()[0] == argmax[0]); +} + +TEST_CASE("AMALIA: greedy continuation reproduces the reference token stream") { + if (!amalia_ready()) { + MESSAGE("AMALIA model/fixtures not present; skipping golden-reference check"); + return; + } + mlxforge::LlamaModel& model = shared_amalia_model(); + std::vector ids = load_amalia_token_ids("prompt_0_ids.npy"); + std::vector expected = load_amalia_token_ids("greedy_tokens.npy"); + + // Full-recompute greedy loop (no cache), mirroring the reference oracle. + std::vector got; + for (size_t i = 0; i < expected.size(); ++i) { + mx::array tokens(ids.data(), {1, static_cast(ids.size())}, mx::int32); + mx::array logits = model.forward(tokens); + const int T = static_cast(ids.size()); + const int vocab = logits.shape()[2]; + mx::array last = mx::reshape(mx::slice(logits, {0, T - 1, 0}, {1, T, vocab}), {1, vocab}); + mx::array nxt = mx::astype(mx::argmax(last, /*axis=*/-1), mx::int32); + mx::eval(nxt); + const int tok = nxt.data()[0]; + got.push_back(tok); + ids.push_back(tok); + } + assert_tokens_equal(got, expected); +} + +TEST_CASE("AMALIA: tokenizer matches the golden ids (Metaspace prepend, Strip, ChatML+BOS)") { + if (!amalia_ready()) { + MESSAGE("AMALIA model/fixtures not present; skipping"); + return; + } + const std::string dir = amalia_model_dir(); + mlxforge::ModelConfig cfg = mlxforge::ModelConfig::from_file(dir + "/config.json"); + + // model_type is "llama"; the ChatML format must come from the on-disk + // chat_template.jinja, not the model_type mapping. + const mlxforge::ChatFormat fmt = mlxforge::chat_format_for_model_dir(dir, cfg.model_type); + CHECK(fmt == mlxforge::ChatFormat::Amalia); + mlxforge::Tokenizer tok = + mlxforge::Tokenizer::from_file(dir + "/tokenizer.json", cfg.bos_token_id, fmt); + + // Corpus byte-match: gates the Metaspace prepend (incl. after special tokens), + // byte_fallback, added-token isolation, and the post-processor BOS (id 3). + nlohmann::json corpus = + nlohmann::json::parse(read_file(amalia_ref_path("tokenizer_corpus.json"))); + REQUIRE(corpus.is_array()); + REQUIRE(corpus.size() > 0); + for (const auto& entry : corpus) { + const std::string text = entry["text"].get(); + const std::vector expected = entry["ids"].get>(); + INFO("input: " << text); + assert_tokens_equal(tok.encode(text), expected); + } + + // The Strip decoder must undo the Metaspace prepend: decode(encode(s)) == s + // for special-free text (mirrors HF decode with skip_special_tokens). + for (const char* s : {"Qual é a capital de Portugal?", "português: ação, coração", + "Amália Rodrigues cantava fado."}) { + CHECK(tok.decode(tok.encode(s)) == s); + } + + // One exact check gates the whole chat path: the Amalia template rendering + // (default PT system prompt), the BOS-completes-<|im_start|> convention (no + // doubled id 3), and the post-special ▁-prepend, vs mlx-lm's + // apply_chat_template on the checkpoint's own jinja. + std::vector messages = { + {"user", "Qual é a capital de Portugal?"}}; + CHECK(tok.apply_chat_template(messages) == load_amalia_token_ids("chat_ids.npy")); +} diff --git a/tests/support/model_fixture.h b/tests/support/model_fixture.h index 5763757..847e731 100644 --- a/tests/support/model_fixture.h +++ b/tests/support/model_fixture.h @@ -111,6 +111,25 @@ inline Qwen35Model& shared_qwen3_5_model() { return model; } +// Same, for AMALIA-9B (plain-Llama forward; what it uniquely exercises is the +// SentencePiece Metaspace-prepend tokenizer and the ChatML-with-BOS template, +// so the model reuses LlamaModel directly). +inline std::string amalia_model_dir() { return MLXFORGE_MODEL_DIR_AMALIA; } + +inline bool amalia_model_available() { + const std::string d = amalia_model_dir(); + return !d.empty() && std::ifstream(d + "/config.json").good(); +} + +inline LlamaModel& shared_amalia_model() { + static LlamaModel model = [] { + ModelConfig cfg = ModelConfig::from_file(amalia_model_dir() + "/config.json"); + Weights w = load_weights(amalia_model_dir(), cfg); + return LlamaModel(std::move(cfg), std::move(w)); + }(); + return model; +} + // Qwen3-VL vision-language model (ViT + multimodal integration tests). The // config and the (keep_vision) Weights are loaded once and shared: the ViT // encoder borrows the Weights, the language model takes a (cheap, handle-only) diff --git a/tests/support/reference.h b/tests/support/reference.h index d225530..33e67aa 100644 --- a/tests/support/reference.h +++ b/tests/support/reference.h @@ -110,6 +110,17 @@ inline std::vector load_qwen3_vl_token_ids(const std::string& name) { return load_token_ids_at(MLXFORGE_REF_FIXTURES_DIR_QWEN3_VL, name); } +// Same accessors against the AMALIA fixture set (reference/fixtures_amalia). +inline std::string amalia_ref_path(const std::string& name) { + return std::string(MLXFORGE_REF_FIXTURES_DIR_AMALIA) + "/" + name; +} +inline mx::array load_amalia_npy(const std::string& name) { + return mx::load(amalia_ref_path(name)); +} +inline std::vector load_amalia_token_ids(const std::string& name) { + return load_token_ids_at(MLXFORGE_REF_FIXTURES_DIR_AMALIA, name); +} + struct CompareResult { bool ok = true; std::string message; @@ -160,6 +171,13 @@ inline CompareResult compare_close(const mx::array& actual, const mx::array& exp for (int64_t i = 0; i < static_cast(a.size()); ++i) { const double av = static_cast(pa[i]); const double bv = static_cast(pb[i]); + // Nonfinite values are always a failure: NaN comparisons are false, so + // without this a NaN-degenerate pair (e.g. an fp16-overflowed forward + // compared against an equally overflowed reference) would pass vacuously. + if (!std::isfinite(av) || !std::isfinite(bv)) { + if (first_bad < 0) first_bad = i; + continue; + } const double diff = std::abs(av - bv); const double tol = atol + static_cast(rtol) * std::abs(bv); if (diff > tol && first_bad < 0) first_bad = i; diff --git a/tests/tokenizer/spm_test.cpp b/tests/tokenizer/spm_test.cpp index 4553d41..47c6d7c 100644 --- a/tests/tokenizer/spm_test.cpp +++ b/tests/tokenizer/spm_test.cpp @@ -26,6 +26,14 @@ bool gemma_available() { return !gemma_dir().empty() && std::ifstream(tokenizer_path()).good(); } +std::string amalia_tokenizer_path() { + return std::string(MLXFORGE_MODEL_DIR_AMALIA) + "/tokenizer.json"; +} +bool amalia_tokenizer_available() { + return !std::string(MLXFORGE_MODEL_DIR_AMALIA).empty() && + std::ifstream(amalia_tokenizer_path()).good(); +} + std::string read_file(const std::string& path) { std::ifstream f(path, std::ios::binary); std::ostringstream ss; @@ -47,6 +55,18 @@ TEST_CASE("SpmBpeTokenizer recognizes the Gemma tokenizer and rejects byte-level CHECK_FALSE(mlxforge::BpeTokenizer::is_supported(blob)); } +TEST_CASE("SpmBpeTokenizer recognizes the AMALIA tokenizer and rejects byte-level BPE") { + if (!amalia_tokenizer_available()) { + MESSAGE("MLXFORGE_MODEL_DIR_AMALIA not present; skipping"); + return; + } + // AMALIA (EuroLLM tokenizer) is SentencePiece-BPE with a Metaspace + // prepend_scheme "always" pre-tokenizer and a Strip decoder -> SPM backend. + const std::string blob = read_file(amalia_tokenizer_path()); + CHECK(mlxforge::SpmBpeTokenizer::is_supported(blob)); + CHECK_FALSE(mlxforge::BpeTokenizer::is_supported(blob)); +} + TEST_CASE("SpmBpeTokenizer matches the Gemma golden ids on a diverse corpus") { if (!gemma_available()) { MESSAGE("MLXFORGE_MODEL_DIR_GEMMA not present; skipping"); diff --git a/tests/tokenizer/tokenizer_test.cpp b/tests/tokenizer/tokenizer_test.cpp index ba46116..eb2116f 100644 --- a/tests/tokenizer/tokenizer_test.cpp +++ b/tests/tokenizer/tokenizer_test.cpp @@ -137,6 +137,34 @@ TEST_CASE("Qwen3 ChatML template renders thinking, tools, and tool turns (no mod std::string::npos); } +TEST_CASE("AMALIA ChatML template injects the default PT system prompt (no model needed)") { + using mlxforge::ChatFormat; + using mlxforge::Tokenizer; + + // No system message: the Portuguese default is injected. The template is + // rendered in full — apply_chat_template encodes it without the BOS prepend + // (AMALIA's BOS id 3 IS <|im_start|>; see render_amalia). + std::string basic = + Tokenizer::render_chat_template({{"user", "Olá"}}, /*add_generation_prompt=*/true, "", + ChatFormat::Amalia); + CHECK(basic.rfind("<|im_start|>system\nO teu nome é Amália", 0) == 0); + CHECK(basic.find("português europeu.<|im_end|>\n<|im_start|>user\nOlá<|im_end|>\n") != + std::string::npos); + CHECK(basic.find("<|im_start|>assistant\n") == basic.size() - 22); + + // An explicit system message replaces the default. + std::string with_sys = Tokenizer::render_chat_template( + {{"system", "sê breve"}, {"user", "Olá"}}, true, "", ChatFormat::Amalia); + CHECK(with_sys.rfind("<|im_start|>system\nsê breve<|im_end|>\n", 0) == 0); + CHECK(with_sys.find("Amália") == std::string::npos); + + // Without a generation prompt the render ends at the last turn. + std::string no_gen = Tokenizer::render_chat_template({{"user", "Olá"}}, false, "", + ChatFormat::Amalia); + CHECK(no_gen.find("assistant") == std::string::npos); + CHECK(no_gen.rfind("<|im_end|>\n") == no_gen.size() - 11); +} + TEST_CASE("streaming detokenizer never emits broken UTF-8") { if (!model_available()) { MESSAGE("MLXFORGE_MODEL_DIR not present; skipping");