Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -262,9 +262,9 @@ Real checkpoints on Apple M5 Max, release binary rebuilt from this tree after th
## 11. Follow-up

- `--draft-block-size` remains unbounded on the CLI while `split-mtp --block-size` is capped at 16. A warning rather than a gate is the right shape, and #1751 left the same note for DSpark, so one change can answer both.
- The `MAX_BLOCK_SIZE` refusal message carries two runs of 14 literal spaces: a wrapped string literal lost its `\` continuations and rustfmt joined the line. The test asserts only that the message contains `20000` and `maximum`, so it does not catch the formatting.
- `--force` removes the stale index but not stale `model-XXXXX-of-YYYYY.safetensors` shards. With the index gone, `collect_shard_paths` falls through to a glob and would read the drafter's `model.safetensors` alongside the orphans. Refusing a forced overwrite of a sharded directory outright is the cheaper fix than deleting shards on the user's behalf.
- The `--q-bits` help text says `(2, 3, 4, 5, 6 or 8)` but the only validation is `validate_quantization_params`, which accepts `1..=32`. A depth MLX's affine packing does not support passes the tool and fails at load.
- Resolved by #1763: the `MAX_BLOCK_SIZE` refusal message carried two runs of 14 literal spaces because a wrapped string literal had lost its `\` continuations and rustfmt joined the line; the message now renders as one sentence and its test asserts no run of two or more spaces.
- Resolved by #1763: `--force` now refuses rather than deletes when the output directory holds any `*.safetensors` other than the drafter's own `model.safetensors` (a previous checkpoint's `model-*.safetensors` shards, or a `consolidated.safetensors`), naming the file and leaving every file untouched; it still clears a stale index when no such file remains.
- Resolved by #1763: `split-mtp` now validates `--q-bits` against `SUPPORTED_AFFINE_BITS` (`{2, 3, 4, 5, 6, 8}`) and `--q-group-size` against `SUPPORTED_AFFINE_GROUP_SIZES` (`{32, 64, 128}`) before any tensor work or directory change, so the help text and the enforced set agree; the shared load-time `validate_quantization_params` bounds check (`1..=32`) is unchanged.
- Batched `B > 1` MTP for this family stays out of scope, and so does drafting deeper than `num_nextn_predict_layers` by default.
- The hidden-tap acceptance gap (1.753 against 1.716) rests on one prompt. If this pairing ever gets a benchmark entry, measure the tap across several prompts there rather than re-deriving it from the module doc.
- Issue #1289 (order-preserving streamed `qmv`) is the route to passing the probe on the wide kernel rather than pinning the process narrow. Until then the pairing is byte-identical only under `MLXCEL_QMV_WIDE=0`, which the gate applies inside a server process and which a standalone `mlxcel generate` comparison has to set by hand.
Original file line number Diff line number Diff line change
Expand Up @@ -262,9 +262,9 @@ Apple M5 Max의 실제 체크포인트, 리베이스 뒤 이 트리에서 다시
## 11. 후속 조치

- `--draft-block-size`는 CLI에서 여전히 무경계인데 `split-mtp --block-size`는 16에서 묶인다. 게이트보다 경고가 어울리고 #1751이 DSpark에 같은 메모를 남겼으니 한 변경이 둘 다 답할 수 있다.
- `MAX_BLOCK_SIZE` 거부 메시지에 14칸짜리 공백 덩어리가 두 개 들어 있다. 줄바꿈된 문자열 리터럴이 `\` 연속 표시를 잃었고 rustfmt가 줄을 합쳤다. 테스트는 메시지가 `20000`과 `maximum`을 포함하는지만 단언하므로 이 서식은 잡히지 않는다.
- `--force`는 낡은 인덱스를 지우지만 낡은 `model-XXXXX-of-YYYYY.safetensors` 샤드는 지우지 않는다. 인덱스가 사라지면 `collect_shard_paths`가 glob으로 떨어져 drafter의 `model.safetensors`와 고아 샤드를 함께 읽는다. 사용자 대신 샤드를 지우는 것보다 샤드로 나뉜 디렉터리에 대한 강제 덮어쓰기를 아예 거부하는 편이 싸다.
- `--q-bits` 도움말은 `(2, 3, 4, 5, 6 or 8)`이라고 적혀 있지만 실제 검증은 `1..=32`를 받는 `validate_quantization_params` 하나다. MLX의 affine 패킹이 지원하지 않는 깊이가 도구를 통과해 로드에서 실패한다.
- #1763에서 해결: `MAX_BLOCK_SIZE` 거부 메시지의 14칸짜리 공백 덩어리 두 개는 줄바꿈된 문자열 리터럴이 `\` 연속 표시를 잃고 rustfmt가 줄을 합치면서 생긴 것이었다. 이제 메시지는 한 문장으로 렌더링되고, 테스트는 공백이 두 칸 이상 연속하지 않음을 단언한다.
- #1763에서 해결: `--force`는 이제 출력 디렉터리에 드래프터 자신의 `model.safetensors`가 아닌 `*.safetensors`(이전 체크포인트의 `model-*.safetensors` 샤드나 `consolidated.safetensors`)가 하나라도 있으면 삭제 대신 거부하며, 그 파일을 이름으로 지목하고 모든 파일을 그대로 둔다. 그런 파일이 없을 때는 여전히 낡은 인덱스를 지운다.
- #1763에서 해결: `split-mtp`는 이제 어떤 텐서 작업이나 디렉터리 변경도 하기 전에 `--q-bits`를 `SUPPORTED_AFFINE_BITS`(`{2, 3, 4, 5, 6, 8}`)로, `--q-group-size`를 `SUPPORTED_AFFINE_GROUP_SIZES`(`{32, 64, 128}`)로 검증하여 도움말과 강제되는 집합이 일치한다. 공유 로드 시점 검증인 `validate_quantization_params`의 범위 검사(`1..=32`)는 그대로다.
- 이 계열의 배치 `B > 1` MTP는 범위 밖으로 남고, 기본값으로 `num_nextn_predict_layers`보다 깊이 draft하는 것도 그렇다.
- 은닉 탭 수락률 차이(1.753 대 1.716)는 프롬프트 하나에 기대고 있다. 이 페어링이 언젠가 벤치마크 항목을 갖게 되면 모듈 문서에서 다시 끌어오는 대신 거기서 여러 프롬프트에 걸쳐 탭을 재는 편이 낫다.
- 이슈 #1289(순서 보존 스트리밍 `qmv`)가 프로세스를 narrow에 고정하는 대신 wide 커널에서 프로브를 통과시키는 경로다. 그전까지 이 페어링은 `MLXCEL_QMV_WIDE=0` 아래에서만 바이트 동일하고 게이트는 서버 프로세스 안에서 그것을 적용하지만 단독 `mlxcel generate` 비교는 손으로 설정해야 한다.
83 changes: 83 additions & 0 deletions TECHNICAL_REPORTS/1778-split-mtp-output-guards-20260911.en.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,83 @@
# Technical Report: PR #1778 - fix(split-mtp): fix refusal text, --force shards, --q-bits validation

**Date**: 2026-09-11
**Author**: mlxcel maintainers
**Reviewer**: implementation review cycle (implementation review, security and performance review)
**Status**: Completed
**Languages**: Rust
**Risk Level**: Low (CLI guards and one shared validator; the drafter written by `split-mtp` is byte-identical to the pre-PR output)

---

## Executive Summary

The PR #1753 report listed three defects in `mlxcel split-mtp` among its follow-ups. The block-size refusal printed two runs of 14 spaces. `--force` left another checkpoint's weight shards next to the drafter. `--q-bits` accepted widths MLX cannot quantize, and the process aborted after the tensor work was done. This PR fixes all three. Review of the first version found five more gaps in the same guards, all fixed here too. The guard now refuses to write beside any safetensors file it did not write. The group size is validated as well as the bit width. Checks that change nothing on disk run before the first check that does.

On the real `glm-4.7-flash-bf16` checkpoint, the drafter this PR writes is byte-identical to the one from the pre-PR binary. The two refusal paths exit 1 without changing a file.

---

## Problem Statement

- **Refusal text.** A wrapped string literal in the `MAX_BLOCK_SIZE` refusal had lost its `\` continuations, and rustfmt then joined the lines, so the indentation became part of the message.
- **`--force` and foreign weight files.** `--force` removed a stale `model.safetensors.index.json` and nothing else. The loader's `glob_safetensors` reads every `*.safetensors` in a directory when there is no index. Shards of a previous checkpoint, left beside the drafter's single `model.safetensors`, would load as part of the drafter.
- **Unsupported widths.** The only check before quantizing was the shared load-time bounds check, `validate_quantization_params`, which accepts `1..=32`. MLX's affine quantize supports 2, 3, 4, 5, 6 and 8, and rejects anything else with `std::invalid_argument`. That call goes through a cxx bridge that does not return `Result`, so the throw was a `std::terminate`: `--q-bits 1` aborted the process after the layer had been read and partly processed.

Review of the first fix found:

- Gemma 4 kept its own copy of the bit set (`SUPPORTED_OVERRIDE_BITS`) instead of the new shared constant.
- `--q-group-size` had no check, so `--q-bits 4 --q-group-size 16` aborted the same way.
- The shard predicate matched only `model-*.safetensors`, so `consolidated.safetensors` bypassed it.
- An unreadable output directory was treated as having no shards.
- `prepare_output_dir` ran before the `--q-bits` check and before the output-equals-source refusal. With `-m X -o X --force` on a single-file source, it deleted the source's own stale index before the refusal ran.

---

## Change Summary

- **Shared validators in `mlxcel-core::layers`.** Adds `SUPPORTED_AFFINE_BITS` (2, 3, 4, 5, 6, 8) and `SUPPORTED_AFFINE_GROUP_SIZES` (32, 64, 128), with `validate_affine_quantization_bits` and `validate_affine_quantization_group_size`. The group sizes are the ones the pinned MLX `mlx/ops.cpp` accepts. `infer_quantization_bits` and Gemma 4's override path read the shared bit set. `validate_quantization_params` is unchanged (#929).
- **Producer-side check in `mlxcel-surgery`.** `split_mtp` validates both parameters before it touches the weight map. `refuse_output_is_source` is public so the CLI can run it early.
- **CLI ordering.** `preflight_checks` (bits, group size, output equals source) changes nothing on disk and runs first. `prepare_output_dir` runs only after it passes.
- **`--force` contract.** `prepare_output_dir` refuses and names the file when the output directory holds any `*.safetensors` other than `model.safetensors`. `find_foreign_safetensors_file` returns an `io::Result`, so a listing error stops the run instead of reading as "clear". A stale index is still removed when nothing foreign remains.
- **Messages.** The block-size refusal is one sentence. Messages from the core validators name the parameter ("quantization bits"), and split-mtp adds the flag name. An earlier doc claim that the six widths are the only ones affine packing can represent was wrong: 7 packs as well, but MLX has no kernel for it. The wording now says only what MLX implements.
- **PR #1753 report.** The three follow-up bullets are marked resolved, with the final scope of the guard.

---

## Technical Decisions

**Refuse, do not delete.** `--force` means "overwrite the files this tool writes", not "empty the directory". A shard in the output directory may be the only copy of a checkpoint that the user pointed at by mistake. Deleting it would be irreversible. Refusing costs one manual removal. The refusal names the file so the user can see what is at stake.

**Screen every foreign safetensors file, not a name pattern.** The loader decides what it reads by extension, so the guard uses the same rule. A name-based list would have to guess every other converter's naming, and `consolidated.safetensors` shows the guess fails.

**Validate at the producer, in a shared module.** The supported sets are MLX's, not split-mtp's, and Gemma 4 already carried a second copy. One definition in `layers.rs` serves both callers. Its tests accept every value in each set and refuse sampled values outside it: 1, 7, 9, 16 and 32 for bits, and 16, 256 and 1 for group size. The producer is the only place a bad width can fail cleanly: once the value reaches MLX's quantize call through a non-`Result` bridge, the only outcome is a process abort.

**Leave the load-time bound alone.** `validate_quantization_params` guards against out-of-range values in configs the loader reads (#929). Narrowing it to MLX's current kernel set would reject checkpoints before MLX is asked. That is a separate decision with its own blast radius, and this PR does not make it.

---

## Validation

- **Gate.** On 059ea43e the workspace gate passed: exit 0, 123 binaries, 11,038 passed, 0 failed. Clippy, fmt, `verify-versions`, `verify-kernel-dtype-keys` and `verify-llama-compat` are clean, and CI is green.
- **Revert arms.** Every new or strengthened test was run with its fix reverted and failed by name. With only the load-time `1..=32` bound left, `--q-bits 1` does not fail with a message. It aborts the test binary: `std::invalid_argument: [quantize] The requested number of bits 1 is not supported`, SIGABRT. The refusal test uses an empty `WeightMap`, which proves the check runs before any tensor work.
- **Real run.** A release binary at 059ea43e against `glm-4.7-flash-bf16` (62.5 GB):
- `--q-bits 4` writes the drafter: 54 tensors, 0.72 GB, `glm4_moe_lite_mtp`, 4-bit affine, group size 64. `model.safetensors` and `config.json` are byte-identical to the output of the #1772 merge binary.
- `--force` over a planted `model-00001-of-00003.safetensors` exits 1 and names the file. Every file's name, size, mtime and inode is unchanged.
- `--q-bits 7` exits 1 in under a second with `quantization bits (7) must be one of 2, 3, 4, 5, 6, 8`, and `drafter2` is never created.

---

## Learning Points

- **Test the rendered message, not the literal.** rustfmt joining a string whose `\` continuation was lost produces valid code and a broken message. The block-size test now asserts there is no run of two spaces.
- **A precondition MLX enforces by throwing is a process abort here.** Through a `UniquePtr`-returning bridge, `std::invalid_argument` becomes `std::terminate`. The only graceful place to reject the value is before the call. The same pattern turned a config value into an abort at first inference in earlier issues.
- **Order checks by side effect.** A guard that mutates must come after every guard that does not. Otherwise a refusal that looks correct in isolation runs after the damage it exists to prevent.
- **Match the guard to the consumer's rule.** The loader reads by extension. A guard that screens by name protects against the names someone thought of.

---

## Follow-ups

- The output writes follow symlinks, and the marker check does not look at tokenizer file names. `exists()` also treats a broken symlink as absent. In a shared directory someone else created, a planted `tokenizer.json` symlink would be truncated even without `--force`. Checking and then writing also leaves minutes in which a second concurrent run passes the same guard. Writing into a temporary directory and renaming it into place closes both. This predates PR #1778.
- The Qwen 3.5 MTP drafter (`drafter/qwen3_5_mtp/model.rs`) passes `bits` and `group_size` from its config to `quantize_weights_with_mode` without the new validators.
- The Markov drafter's supported bit list omits 5.
Loading