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:
parent
b13ebdf905
commit
d7d2e5ba68
2 changed files with 47 additions and 20 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue