What do I need to edit to get MiMo 2.5 Non-Pro's Dflash working?

#1
by Downtown-Case - opened

I'm trying to convert non-Pro MiMo 2.5 dflash myself, but I'm getting horrible acceptance rates with the built-in benchmark (less than 5%) with any settings I try.

...But I noticed that you seem to manually change some metadata in the converted GGUF. For example, the backbone_rotary_base and rope.dimension_count didn't match the config to your Pro Dflash GGUF, so I edited them to appropriate values (5000000 for backbone_rotary_base and 32 for rope.dimension_count).

Precisely what did you do to correct the GGUF? As I still can't get acceptance rates to anything sane.

Thanks for flagging this. The model card had a stale metadata row from an intermediate conversion, and I’ve corrected it.
The final GGUF was produced by reconverting with the converter from ik_llama PR #2048. I did not manually patch the finished GGUF.
The important distinction is that the draft uses two separate RoPE bases:

  • dflash-draft.rope.freq_base = 10000 for the draft’s own Q/K RoPE
  • dflash-draft.dflash.backbone_rotary_base = 5000000 for captured target-context K
  • dflash-draft.rope.dimension_count = 64, from head_dim 128 × partial_rotary_factor 0.5
    Do not set the normal rope.freq_base to 5,000,000. That recreates an earlier converter bug.
    The conversion command was:
python3 ../ik_llama.cpp/convert_hf_to_gguf.py dflash \
  --target-model-dir . \
  --outtype bf16 \
  --outfile mimo-v25-pro-dflash-draft-bf16.gguf

A correct output has 63 tensors, including five blk.N.attn_sinks.weight tensors, and no token_embd or output tensors because those are shared with the target. It also needs attention.value_scale=0.612, SWA 1024, target_layer_ids=[0,15,31,47,69], and n_target_features=30720.
Editing two metadata fields on an older conversion is not sufficient. PR #2048 also added attention-sink conversion, value-scale handling, and a MiMo2 target-feature capture fix in the runtime.
As a diagnostic, try your target with the published draft using an ik_llama build containing PR #2048 and start with dflash:n_max=1,p_min=0.0,cross_ctx=16. Also confirm the target generates coherently without DFlash. If the published draft also remains below 5%, please send the target and draft metadata dumps, exact ik_llama commit, and benchmark command. That will tell us whether the problem is the target conversion, runtime, or draft conversion.
As noted in #2048, this was a validation enabling DFlash for MiMo v2.5 Pro, and I didn't manage a speedup over no-spec, with acceptance in the 50s.

Okay, with a cross_ctx value around 8, and n_max=1, AND your GGUF meta tweaks acceptance rate is rising to ~15%. Still a bit broken, but dramatically better.

Do not set the normal rope.freq_base to 5,000,000. That recreates an earlier converter bug.
The conversion command was:

Ah I see, I just saw this reply.

I will correct this back in the GGUF and see if it improves, thanks for the response.

Though with this comment:

dflash-draft.rope.dimension_count = 64, from head_dim 128 × partial_rotary_factor 0.5
Do not set the normal rope.freq_base to 5,000,000. That recreates an earlier converter bug.
The conversion command was:

It appears this value is set incorrectly for non-pro MiMo? It has 64 heads, partial_rotary_factor 0.5, so its rope dimension should be 32, no?

But the rope dimension 64 in the unedited GGUF. See here:

https://huggingface.co/Downtown-Case/dflash-temp-debug/blob/main/Dflash-1.5B-BF16.gguf

Is this correct? Or should it be 32?

And the original config is here, as a quick link: https://huggingface.co/XiaomiMiMo/MiMo-V2.5-DFlash/blob/main/dflash/config.json

And the GGUF was made yesterday, so it should have your PR in already.

Though with this comment:

It appears this value is set incorrectly for non-pro MiMo? It has 64 heads, partial_rotary_factor 0.5, so its rope dimension should be 32, no?

But the rope dimension 64 in the unedited GGUF. See here:

https://huggingface.co/Downtown-Case/dflash-temp-debug/blob/main/Dflash-1.5B-BF16.gguf

Is this correct? Or should it be 32?

That's entirely possible. #2048 was specific to Pro, and I didn't work with non-Pro. #2048 investigation and validation was all on rented hardware, so I can't recreate it on my local hardware for non-Pro.

I checked the published non-Pro DFlash config, and it looks like the RoPE dimension should be 64 according to both the config and the converter:

head_dim 128 × partial_rotary_factor 0.5 = 64

rope.dimension_count is the number of rotated dimensions in the #2048 converter. Dividing by two again to get 32 would be counting RoPE pairs instead, and that is not how this converter writes the GGUF key. The Pro and non-Pro draft configs have the same head_dim=128 and partial_rotary_factor=0.5, so I do not see config evidence for using 32.

The non-Pro draft does have several values that must not be copied from Pro:

  • attention.head_count = 64
  • attention.head_count_kv = 8
  • rope.freq_base = 10000
  • dflash.backbone_rotary_base = 5000000
  • attention.value_scale = 0.612
  • mask_token_id = 151675
  • target_layer_ids = [0, 11, 23, 35, 47]
  • n_target_features = 20480
  • attention.layer_norm_rms_epsilon = 1e-6
  • sliding window 1024 across all five draft layers

One easy point of confusion is the value scale. The non-Pro target config has attention_value_scale=0.707, but the DFlash draft’s nested config has attention_value_scale=0.612. The target and draft GGUFs should retain their respective values.

I would start with 64 and verify the non-Pro-specific fields above with a build containing #2048.

I hope that's helpful, and wouldn't mind hearing how your experiment turns out.

Thanks! Yes, I was indeed mixing up fields, heh.

I hope that's helpful, and wouldn't mind hearing how your experiment turns out.

Unsuccessful so far, but I'm going to check the main quant's KLD vs the BF16, make sure its alright, then try again.

Sign up or log in to comment