Skip to content

fix(split-mtp): tidy the block-size refusal text, clear stale shards under --force, and align the --q-bits help with the validator #1763

Closed

Description

@inureyes

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.rs tests): a tempdir holding model-00001-of-00003.safetensors, model.safetensors.index.json and config.json with force = true must 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 with model.safetensors and an index, force = true, returns Ok, removes the index, leaves model.safetensors.
  • split_refuses_an_unsupported_q_bits (new, split_mtp_tests.rs): for [1, 7, 9, 16, 32], split_mtp refuses 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_mtp succeeds and records config["quantization"]["bits"], following split_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 7

Pass 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_SIZE refusal renders with no run of two or more spaces, asserted by its test.
  • --force over a directory holding model-*.safetensors exits non-zero naming a shard and leaves every shard and index untouched on disk.
  • --force over a stale index with no shards still removes the index and proceeds, as today.
  • --q-bits of 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_BITS is the single definition of the set; the literal at src/lib/mlxcel-core/src/layers.rs:1227 is gone.
  • validate_quantization_params still accepts 1..=32 and 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,accelerate passes.
  • cargo clippy --workspace --all-targets -- -D warnings and cargo fmt --check pass.
  • The three follow-up bullets at TECHNICAL_REPORTS/1753-glm4-moe-lite-mtp-drafter-20260910.en.md:265-267 are marked resolved in both the .en.md and .ko.md reports.

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.

Dependencies

Follow-up of #1326 (PR #1753).

Activity

  1. added
    type:bugBug fixes, error corrections, or issue resolutions
    area:cliCommand-line interface / CLI flags
    area:coremlxcel-core: MLX FFI, primitives, KV cache, layers
    and removed on Sep 10, 2026
  2. added a commit that references this issue on Sep 15, 2026
    33cd41d
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    area:cliCommand-line interface / CLI flagsarea:coremlxcel-core: MLX FFI, primitives, KV cache, layerspriority:lowLow prioritystatus:doneCompletedtype:bugBug fixes, error corrections, or issue resolutions

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions