From d7d2e5ba68c41142930538d4aedcc3876f15f5a2 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sat, 27 Jun 2026 07:25:45 +0000
Subject: [PATCH] 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
---
 unsloth/_compressed_quantize.py | 43 ++++++++++++++++++++++++---------
 unsloth/save.py                 | 24 +++++++++++-------
 2 files changed, 47 insertions(+), 20 deletions(-)

diff --git a/unsloth/_compressed_quantize.py b/unsloth/_compressed_quantize.py
index 85a09ec47d..a61c9a5d0d 100644
--- a/unsloth/_compressed_quantize.py
+++ b/unsloth/_compressed_quantize.py
@@ -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)
diff --git a/unsloth/save.py b/unsloth/save.py
index 2f48f5b432..12e8234df4 100644
--- a/unsloth/save.py
+++ b/unsloth/save.py
@@ -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,