Use sequential calibration pipeline and validate Hub access early

- nvfp4 calibration no longer forces the memory-hungry "basic" pipeline. The
  quantization runs in a clean subprocess, so llm-compressor's default
  sequential pipeline (layer-by-layer onloading) works and lets large models
  that do not fit at once still calibrate; fall back to "basic" only if tracing
  fails
- For push_to_hub compressed exports, create/validate the repo up front so a bad
  token or denied repo fails before the merge and quantization instead of after
This commit is contained in:
Daniel Han 2026-06-27 07:25:45 +00:00
commit d7d2e5ba68
2 changed files with 47 additions and 20 deletions

View file

@ -165,7 +165,9 @@ def main():
)
tokenizer = None
recipe = QuantizationModifier(targets = "Linear", scheme = args.scheme, ignore = ["lm_head"])
def _make_recipe():
return QuantizationModifier(targets = "Linear", scheme = args.scheme, ignore = ["lm_head"])
if args.needs_calibration:
ds = _build_calibration_dataset(
tokenizer,
@ -174,17 +176,36 @@ def main():
args.num_calibration_samples,
args.max_seq_length,
)
# "basic" pipeline runs a normal forward (no AST tracing / sequential splitting).
oneshot(
model = model,
dataset = ds,
recipe = recipe,
max_seq_length = args.max_seq_length,
num_calibration_samples = args.num_calibration_samples,
pipeline = "basic",
)
# Let llm-compressor pick its default (sequential) pipeline: it onloads layer-by-layer,
# so models that do not fit in memory at once can still calibrate. Running here in a clean
# process (Unsloth's attention patches are absent) means tracing works; fall back to the
# memory-hungry "basic" pipeline only if tracing fails.
try:
oneshot(
model = model,
dataset = ds,
recipe = _make_recipe(),
max_seq_length = args.max_seq_length,
num_calibration_samples = args.num_calibration_samples,
)
except Exception as e:
print(
f"Unsloth: sequential calibration pipeline failed ({type(e).__name__}: {e}); "
"retrying with the 'basic' pipeline (needs the full model to fit in memory).",
flush = True,
)
model = _from_pretrained(auto_model, args.model, args.trust_remote_code)
model.eval()
oneshot(
model = model,
dataset = ds,
recipe = _make_recipe(),
max_seq_length = args.max_seq_length,
num_calibration_samples = args.num_calibration_samples,
pipeline = "basic",
)
else:
oneshot(model = model, recipe = recipe)
oneshot(model = model, recipe = _make_recipe())
os.makedirs(args.out, exist_ok = True)
model.save_pretrained(args.out, save_compressed = True)

View file

@ -3859,7 +3859,21 @@ def _unsloth_save_compressed_tensors(
# Wrap the body so the isolated temp dirs are always cleaned up, even when the merge,
# quantization, validation, or hub upload raises.
api = None
try:
# Validate Hub access up front (a bad token / denied repo should fail before the expensive
# merge and quantization, matching the normal push path). create_repo is idempotent.
if push_to_hub:
from huggingface_hub import HfApi
api = HfApi(token = token)
api.create_repo(
repo_id = repo_id,
repo_type = "model",
private = merge_kwargs.get("private", None),
exist_ok = True,
)
# 3) Merge to 16bit at local_dir (kept for local saves) via unsloth_generic_save, so LoRA
# adapters are merged and full-finetuned models written in 16bit consistently. Extra
# save kwargs (state_dict, max_shard_size, ...) flow through merge_kwargs.
@ -3986,17 +4000,9 @@ def _unsloth_save_compressed_tensors(
)
# 8) Optional hub upload of the compressed artifact (not the intermediate 16bit one).
# The repo was already created/validated up front, so just upload here.
if push_to_hub:
print(f"Unsloth: Uploading {scheme} checkpoint to '{repo_id}' ...")
from huggingface_hub import HfApi
api = HfApi(token = token)
api.create_repo(
repo_id = repo_id,
repo_type = "model",
private = merge_kwargs.get("private", None),
exist_ok = True,
)
api.upload_folder(
folder_path = out_dir,
repo_id = repo_id,