Repository navigation
fix(split-mtp): tidy the block-size refusal text, clear stale shards under --force, and align the --q-bits help with the validator #1763
Description
Summary
The PR #1753 technical report recorded three small defects in its own follow-up section that were out of scope for the report pass (TECHNICAL_REPORTS/1753-glm4-moe-lite-mtp-drafter-20260910.en.md:265-267). All three sit in the same file pair, none touches inference arithmetic, so they are one PR. Verified against e3ca8afc.
Current behavior
1. The block-size refusal carries two runs of 14 literal spaces. The MAX_BLOCK_SIZE refusal at src/lib/mlxcel-surgery/src/ops/split_mtp.rs:196-199 is one string literal that lost its \ continuations, so rustfmt joined the wrapped source lines and left the indentation inside the message. On line 197 the runs sit at columns 94 and 190, 14 spaces each, and the user reads ...this family verifies one query row at a time. The covering test split_refuses_an_oversized_block_size (src/lib/mlxcel-surgery/src/ops/split_mtp_tests.rs:564-575) asserts only contains("20000") and contains("maximum") at lines 573-574, so the formatting is invisible to the gate.
2. --force clears the stale index but leaves stale shards. src/commands/split_mtp.rs:87-95 removes model.safetensors.index.json and nothing else, while existing_checkpoint_marker already recognises a model-*.safetensors shard as a checkpoint marker (src/commands/split_mtp.rs:164-169). Forcing over a previously sharded directory leaves the orphan shards beside the drafter's single model.safetensors, which is the case the index removal was meant to guard: with no index, collect_shard_paths takes the no-index branch (src/lib/mlxcel-core/src/weights.rs:607-609) and calls glob_safetensors, which returns every *.safetensors in the directory (src/lib/mlxcel-core/src/weights.rs:711). The loader then reads the drafter and the victim's orphans as one checkpoint.
3. The --q-bits help and the validation disagree. The help at src/commands/split_mtp.rs:53-56 says (2, 3, 4, 5, 6 or 8), but the only check before quantize_weights_with_mode is validate_quantization_params (called at src/lib/mlxcel-surgery/src/ops/split_mtp.rs:379-381), which accepts 1..=32 (src/lib/mlxcel-core/src/layers.rs:1328). --q-bits 7 is accepted, quantizes, writes a directory, and fails later at load. The repository already treats {2, 3, 4, 5, 6, 8} as the authoritative affine set, as a bare literal in infer_quantization_bits at src/lib/mlxcel-core/src/layers.rs:1227.
Expected behavior
The refusal reads as one sentence with single spaces. A forced overwrite of a previously sharded output directory fails naming the shard, deleting nothing and writing nothing. --q-bits outside {2, 3, 4, 5, 6, 8} is refused before any tensor work, naming the accepted widths, so the help and the enforced set agree.
Implementation plan
1. Restore the \ continuations in the literal at src/lib/mlxcel-surgery/src/ops/split_mtp.rs:196-199 so the rendered message holds no run of two or more spaces, and add assert!(!msg.contains(" "), "{msg}") to the test at src/lib/mlxcel-surgery/src/ops/split_mtp_tests.rs:573-574. The message is a single sentence with no intentional alignment, so a plain double-space assertion is the right shape.
2. Refuse, do not delete. This is the report's own conclusion at line 266 and it stands: a flag named --force does not promise to delete weight files, and the orphan shards may be the only surviving copy. Extract src/commands/split_mtp.rs:80-96 into fn prepare_output_dir(dir: &Path, force: bool) -> Result<()>, because run_split_mtp calls initialize_runtime_checked() and does GPU work so the guard cannot be exercised through it. Contract: no marker, Ok(()); marker and !force, today's already holds a checkpoint ({marker}); pass --force to overwrite error unchanged; marker, force, and at least one model-*.safetensors present, a new refusal naming one shard and telling the user to remove the directory by hand; marker, force, no shards, remove a stale index as today and return Ok(()). Factor the shard-name predicate inlined at src/commands/split_mtp.rs:164-169 into one named helper both callers share.
3. Enforce the affine set on the producer, not on the shared load-time validator. Do not tighten validate_quantization_params: it is deliberately a bounds check rather than an allowlist (issue #929), its test says so and asserts (1, 1) and (1, 32) are Ok (src/lib/mlxcel-core/src/layers.rs:5869-5879), and 20 files call it on the load path where a narrower set would refuse checkpoints that load today. split-mtp is a producer: it picks the width and hands it to quantize_weights_with_mode (src/lib/mlxcel-surgery/src/ops/split_mtp.rs:388-389). Promote the literal at src/lib/mlxcel-core/src/layers.rs:1227 to pub const SUPPORTED_AFFINE_BITS: [i32; 6] = [2, 3, 4, 5, 6, 8]; beside SUPPORTED_QUANTIZATION_MODES (src/lib/mlxcel-core/src/layers.rs:1362), have infer_quantization_bits read it, and add pub fn validate_affine_quantization_bits(bits: i32) -> Result<(), String> next to validate_quantization_params. Call it from split_mtp alongside resolve_block_size (src/lib/mlxcel-surgery/src/ops/split_mtp.rs:308-310), not at the existing call site inside the quantization branch, so a typo fails before the tensor work as the block-size ceiling does. The help text is already correct and does not change; the code moves to meet it.
Validation
Unit tests, all under the workspace gate:
split_refuses_an_oversized_block_size(src/lib/mlxcel-surgery/src/ops/split_mtp_tests.rs), extended with the no-double-space assertion.force_refuses_a_previously_sharded_output_directory(new,src/commands/split_mtp.rstests): a tempdir holdingmodel-00001-of-00003.safetensors,model.safetensors.index.jsonandconfig.jsonwithforce = truemust error naming the shard, and the test must assert both files are still on disk afterwards, so the guard is proven non-destructive.force_clears_a_stale_index_when_no_shards_remain(new, same module): a tempdir withmodel.safetensorsand an index,force = true, returnsOk, removes the index, leavesmodel.safetensors.split_refuses_an_unsupported_q_bits(new,split_mtp_tests.rs): for[1, 7, 9, 16, 32],split_mtprefuses naming the width and the accepted set.split_accepts_every_supported_affine_bit_width(new, same file): for[2, 3, 4, 5, 6, 8],split_mtpsucceeds and recordsconfig["quantization"]["bits"], followingsplit_accepts_block_sizes_up_to_the_ceiling.
One real run, so the guard is confirmed against a genuine checkpoint and not only a fixture:
cargo build --release --features metal,accelerate
OUT=$(mktemp -d)
./target/release/mlxcel split-mtp -m models/glm-4.7-flash-bf16 -o "$OUT/drafter" --q-bits 4
touch "$OUT/drafter/model-00001-of-00003.safetensors"
./target/release/mlxcel split-mtp -m models/glm-4.7-flash-bf16 -o "$OUT/drafter" --force
ls "$OUT/drafter"
./target/release/mlxcel split-mtp -m models/glm-4.7-flash-bf16 -o "$OUT/drafter2" --q-bits 7Pass looks like: the first run writes the drafter; the second exits non-zero naming model-00001-of-00003.safetensors and the ls still shows both that file and model.safetensors; the third exits non-zero naming 7 and the accepted widths without creating drafter2.
Acceptance criteria
- The
MAX_BLOCK_SIZErefusal renders with no run of two or more spaces, asserted by its test. -
--forceover a directory holdingmodel-*.safetensorsexits non-zero naming a shard and leaves every shard and index untouched on disk. -
--forceover a stale index with no shards still removes the index and proceeds, as today. -
--q-bitsof 1, 7, 9, 16 or 32 is refused before any tensor work naming the width and{2, 3, 4, 5, 6, 8}; each of 2, 3, 4, 5, 6 and 8 is accepted. -
SUPPORTED_AFFINE_BITSis the single definition of the set; the literal atsrc/lib/mlxcel-core/src/layers.rs:1227is gone. -
validate_quantization_paramsstill accepts1..=32and its issue fix(core): an unvalidated config.json quantization block reaches MLX and terminates the process #929 test is unchanged. - The real run above behaves as described against
models/glm-4.7-flash-bf16, with its output pasted into the PR body. -
cargo test --workspace --profile test-fast --features metal,acceleratepasses. -
cargo clippy --workspace --all-targets -- -D warningsandcargo fmt --checkpass. - The three follow-up bullets at
TECHNICAL_REPORTS/1753-glm4-moe-lite-mtp-drafter-20260910.en.md:265-267are marked resolved in both the.en.mdand.ko.mdreports.
Out of scope
The fourth follow-up bullet, --draft-block-size unbounded on the CLI while split-mtp --block-size caps at 16 (report line 264), is a warning-versus-gate decision that #1751 left the same note for on DSpark, so it belongs with that change. Batched B > 1 MTP and drafting deeper than num_nextn_predict_layers stay out of scope as the report recorded.
Activity
fix(split-mtp): fix refusal text, --force shards, --q-bits validation (…