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:
parent
a8a03154ea
commit
a1b4b9d431
2 changed files with 58 additions and 1 deletions
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue