test(rocm): correctness matrix against a Metal baseline and a verify-test-rocm gate
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
Part of #1801. Phase 3. Depends on #1803, #1805, #1806. Rows for mxfp8 and mxfp4 checkpoints are added when #1807 and #1808 land.
## Context
The spike validated MLX ops and mlx-lm generation, not mlxcel's own model code. mlxcel judges numerics with a teacher-forced logit trace and the decided-position mismatch rate (`docs/benchmarks.md`, "Judging a change that moves the numbers"), because byte equality is not expected across backends and perplexity can hide greedy flips. Real-model integration tests skip silently when a checkpoint is missing (`tests/common/mod.rs`, `repo_model_dir` searches `models/`, `models/mlx/`, `models/mlx-big/`) unless `MLXCEL_REQUIRE_MODELS=1`.
There is no ROCm test gate. `make verify` uses `--features metal,accelerate` and is not usable on Linux; the CUDA analogue is `make verify-test-cuda` (`Makefile:706`).
## Scope
A reproducible correctness matrix for ROCm against a Metal reference, and a local test gate for ROCm.
## Implementation plan
1. **Gate.** Add `make verify-test-rocm` = `cargo test --workspace --profile test-fast --features rocm --no-fail-fast -- --test-threads=1` (test binaries must not share the GPU; see the note on `verify-test-cuda`). Record the pass count; a run reporting far fewer tests than the workspace total (about 11k) is a failure.
2. **Model matrix** (affine 4-bit first, one per path mlxcel distinguishes):
- dense: `Meta-Llama-3.1-8B-Instruct-4bit`, `Qwen2.5-7B-Instruct-4bit`
- small smoke: `Qwen3-0.6B-4bit`
- sliding window: `gemma-3-4b-it-4bit`
- MoE: `Qwen3-30B-A3B-4bit`, `Mixtral-8x7B-Instruct-v0.1-4bit`
- SSM hybrid: `granite-4.0-h-tiny-4bit`, `NVIDIA-Nemotron-3-Nano-30B-A3B-4bit` (SSM and gated-delta kernels are Metal-only, so ROCm runs the graph paths)
- one VLM on the default path
- `Mixtral-8x7B-Instruct-v0.1-4bit` (E=8) doubles as the guard for the ROCm expert-batched gather path that the overlay keeps opt-in: its width-256 trace goes through sorted `gather_qmm` with at most 64 experts.
3. **Reference.** Metal logit traces from an M1 Ultra host (Apple GPU generation 13, so the Metal path without NAX; record this, since M5-class hosts gated by `hw.has_neural_accelerator && hw.macos_supports_na` take different kernels) with the same checkpoint, corpus and arguments, built from the same mlxcel commit. Corpus `tests/fixtures/wikitext2_excerpt.txt`; `logit_trace MODEL_DIR TEXT_FILE CHUNK_TOKENS MAX_CHUNKS TOPK PREFILL` at width 1 (`1 128 8 0`), width 8 after 512 tokens of context (`8 80 8 512`) and width 256 (`256 2 8 0`). Note the known `logit_trace` defect that chunk 0 gets no BOS (lablup/mlxcel#1785); it affects both sides equally.
4. **Output A/B.** `scripts/ab_output_equality.sh` with `--out-dir` set explicitly (its default uses BSD-style `mktemp -d -t`, which GNU `mktemp` may reject) for greedy smoke on each model.
5. **Server.** Start `mlxcel-server` with one dense and one MoE model and exercise `/v1/chat/completions` (streaming and non-streaming).
6. Store the matrix, commands and thresholds in `docs/benchmark_results/` next to the benchmark page from #1810.
## Judging cross-backend results
Metal and ROCm are different accelerators with different reduction orders, accumulation dtypes, kernel choices per width and transcendental implementations, so byte-identical logits or free-running outputs are not the criterion.
- **Metric.** Decided-position mismatch from `scripts/compare_logit_traces.py` (default `--decided 2.0`) at each of the three widths. Report top-1 mismatch, where the mismatched token ranks in the reference, perplexity difference and any NaN alongside it, as context rather than as the gate.
- **Threshold.** Not zero. Set it from a measured noise floor: the spread between backends that are already trusted on the same checkpoint, corpus and arguments (for example Metal on two GPU generations, and CUDA against Metal). Which combinations to measure is decided after the first ROCm results are in. For scale, the `81ba1c6a` pin bump alone moved top-1 by up to 15 of 512 positions on unchanged Metal hardware while decided-position mismatch stayed at 0.
- **Align known systematic differences first**, so only backend numerics remain: quantized checkpoints keep bf16 scales on both sides; unquantized checkpoints run on Metal with `MLXCEL_KEEP_BF16=1` (Metal otherwise converts bf16 weights to f16 at load); MoE models are traced with fused MoE on and off on Metal and compared with whichever path ROCm takes after lablup/mlxcel#1803; both sides use the same mlxcel commit.
- **Expected sources of difference** that should not by themselves block: GEMV/GEMM kernel selection per width M, reduction order and accumulation dtype, `exp`/`sigmoid` implementations (for example upstream's switch to `precise::exp`, ml-explore/mlx#4461), and MoE router logits whose last-digit differences select a different expert.
- **Free-running greedy output** is a smoke check only (fluency, no NaN, no empty output); it diverges after a few tokens even between healthy backends.
## Acceptance criteria
- [ ] `make verify-test-rocm` exists and passes on `gfx1151` with `MLXCEL_REQUIRE_MODELS=1` for the checkpoints present.
- [ ] Every model in the matrix generates on ROCm; decided-position mismatch rates against the Metal reference are reported per model and width, and fall within the measured noise floor described above, or each outlier has a filed follow-up.
- [ ] `mlxcel-server` serves chat completions on ROCm for one dense and one MoE model.
- [ ] The matrix document lists mlxcel commit, MLX pin, source fork commit of the mlxcelverse ROCm overlay, ROCm/HIP version, `gfx` target and host.
## Validation
```
make verify-test-rocm
cargo build --release --features rocm --example logit_trace
./target/release/examples/logit_trace models/mlx/Meta-Llama-3.1-8B-Instruct-4bit corpus.txt > rocm.tsv
python3 scripts/compare_logit_traces.py metal.tsv rocm.tsv
```
## References
- `docs/benchmarks.md`, `examples/logit_trace`, `scripts/compare_logit_traces.py`, `scripts/ab_output_equality.sh`
- `tests/common/mod.rs` (`repo_model_dir`), `Makefile:706`
Contributor guide
Research direction
Start at Makefile:706 and the CUDA gate, then read tests/common/mod.rs, examples/logit_trace, scripts/compare_logit_traces.py, and scripts/ab_output_equality.sh. Run the listed ROCm validation commands first; done means the gate passes, the model matrix and cross-backend measurements are recorded in docs/benchmark_results/, and dense and MoE server requests work on ROCm.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, rust, shell
- Domain
- backend, testing
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 25/100