From 03c10ab3b615608dd89382005213910c12ac2dbf Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 9 May 2026 08:00:50 +0000 Subject: [PATCH] Fix model loading: use modelopt get_model() instead of raw AutoModelForCausalLM MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Raw from_pretrained OOMs during weight conversion — torch.cat on expert gate_up_proj tries to allocate 31.5GB on a GPU with only 25.9GB free. modelopt's get_model() handles max_memory/device_map properly for models that need sequential device mapping. --- scripts/quantize_nvfp4.py | 25 +++++++++++++------------ 1 file changed, 13 insertions(+), 12 deletions(-) diff --git a/scripts/quantize_nvfp4.py b/scripts/quantize_nvfp4.py index d18d045..5286ea2 100644 --- a/scripts/quantize_nvfp4.py +++ b/scripts/quantize_nvfp4.py @@ -259,7 +259,7 @@ def run_calibration(model_path, export_dir, calib_save_path, amax_snapshot_path, os.environ["HUGGING_FACE_HUB_TOKEN"] = HF_TOKEN from hf_ptq import ( - get_model, get_tokenizer, make_calib_dataloader, + get_model as modelopt_get_model, get_tokenizer, make_calib_dataloader, build_quant_cfg, load_mtp_weights, copy_custom_model_files, QUANT_CFG_CHOICES, ) @@ -267,22 +267,24 @@ def run_calibration(model_path, export_dir, calib_save_path, amax_snapshot_path, from modelopt.torch.quantization.config import need_calibration from modelopt.torch.utils.dataset_utils import get_max_batch_size from modelopt.torch.export import export_hf_checkpoint - from transformers import AutoModelForCausalLM, AutoTokenizer + from transformers import AutoTokenizer apply_patches() # ── Load model ── + # Use modelopt's get_model() instead of raw AutoModelForCausalLM.from_pretrained. + # The raw call OOMs during weight conversion (torch.cat on experts needs 31.5GB, + # only 25.9GB free). modelopt's loader handles max_memory/device_map properly. print(f"\nLoading model from {model_path}...") t0 = time.time() - model = AutoModelForCausalLM.from_pretrained( + model = modelopt_get_model( model_path, + gpu_mem_percentage=GPU_MEM_PCT, trust_remote_code=True, - torch_dtype=torch.bfloat16, - device_map="sequential", - offload_folder="offload", + use_seq_device_map=True, ) - tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) + tokenizer = get_tokenizer(model_path, trust_remote_code=True) print(f"✓ Model loaded in {time.time()-t0:.0f}s") # ── Setup quantization config ── @@ -411,16 +413,15 @@ def run_export_only(calib_save_path, amax_snapshot_path, model_path, export_dir) apply_patches() - from transformers import AutoModelForCausalLM, AutoTokenizer + from hf_ptq import get_model as modelopt_get_model, get_tokenizer print(f"Loading model skeleton from {model_path}...") - model = AutoModelForCausalLM.from_pretrained( + model = modelopt_get_model( model_path, + device="cpu", trust_remote_code=True, - torch_dtype=torch.bfloat16, - device_map="cpu", ) - tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) + tokenizer = get_tokenizer(model_path, trust_remote_code=True) print(f"Loading calibrated state from {calib_save_path}...") state = torch.load(calib_save_path, map_location='cpu')