Handle DatasetDict calibration, MoE routers, and MTP models in compressed export

- Reduce an in-memory DatasetDict calibration set to a single split before row
  subsampling, so save_to_disk does not copy every split to the temp dir
- For MoE models, keep the router/gate unquantized and pass
  moe_calibrate_all_experts so every expert is calibrated
- Warn when a model carries MTP / speculative-decoding tensors that the
  compressed export does not include
This commit is contained in:
Daniel Han 2026-06-27 12:19:34 +00:00
commit a1b4b9d431
2 changed files with 58 additions and 1 deletions

View file

@ -25,6 +25,34 @@ import os
import sys
def _is_moe(config):
"""True if the model config looks like a sparse Mixture-of-Experts model."""
if config is None:
return False
for cfg in (config, getattr(config, "text_config", None)):
if cfg is None:
continue
for attr in ("num_experts", "num_local_experts", "n_routed_experts", "moe_num_experts"):
v = getattr(cfg, attr, None)
if isinstance(v, int) and v > 1:
return True
return "moe" in (getattr(config, "model_type", "") or "").lower()
def _has_mtp(config):
"""True if the model carries MTP / speculative-decoding layers (e.g. Qwen3-Next, DeepSeek)."""
if config is None:
return False
mt = (getattr(config, "model_type", "") or "").lower()
if "qwen3_next" in mt or "mtp" in mt:
return True
for attr in ("num_nextn_predict_layers", "num_mtp_layers", "mtp_num_layers"):
v = getattr(config, attr, None)
if isinstance(v, int) and v > 0:
return True
return False
def _build_calibration_dataset(tokenizer, kind, value, num_samples, max_seq_length):
from datasets import DatasetDict, load_dataset, load_from_disk
@ -198,8 +226,16 @@ def main():
)
tokenizer = None
# MoE models: keep the router/gate unquantized (it decides expert routing) and calibrate every
# expert even if the sample set does not route tokens to all of them.
is_moe = _is_moe(getattr(model, "config", None))
ignore = ["lm_head"]
if is_moe:
ignore.append("re:.*\\.gate$")
moe_kwargs = {"moe_calibrate_all_experts": True} if is_moe else {}
def _make_recipe():
return QuantizationModifier(targets = "Linear", scheme = args.scheme, ignore = ["lm_head"])
return QuantizationModifier(targets = "Linear", scheme = args.scheme, ignore = ignore)
if args.needs_calibration:
ds = _build_calibration_dataset(
@ -221,6 +257,7 @@ def main():
max_seq_length = args.max_seq_length,
num_calibration_samples = args.num_calibration_samples,
pipeline = "sequential",
**moe_kwargs,
)
except Exception as e:
print(
@ -254,6 +291,7 @@ def main():
max_seq_length = args.max_seq_length,
num_calibration_samples = args.num_calibration_samples,
pipeline = "basic",
**moe_kwargs,
)
else:
oneshot(model = model, recipe = _make_recipe())
@ -263,6 +301,14 @@ def main():
if tokenizer is not None:
tokenizer.save_pretrained(args.out)
if _has_mtp(getattr(model, "config", None)):
print(
"Unsloth: WARNING - this model has MTP / speculative-decoding tensors that are not "
"included in the compressed export (only the main model is quantized and saved). Use "
"the non-compressed save path if you need the MTP weights.",
flush = True,
)
cfg_path = os.path.join(args.out, "config.json")
cfg = {}
if os.path.exists(cfg_path):

View file

@ -3929,6 +3929,17 @@ def _unsloth_save_compressed_tensors(
elif hasattr(calibration_dataset, "save_to_disk"):
# Only persist the samples we need, so multi-GB training sets are not fully copied.
ds_to_save = calibration_dataset
# A DatasetDict's len() is the split count, not rows; pick one split first so the
# row subsample below applies and we do not save every split to the temp dir.
try:
from datasets import DatasetDict
if isinstance(ds_to_save, DatasetDict):
ds_to_save = ds_to_save.get("train", None) or next(
iter(ds_to_save.values())
)
except Exception:
pass
try:
if (
num_calibration_samples