fix: load already-quantized 4-bit model directly (no BnB)

This commit is contained in:
Christian Medina
2026-07-02 18:20:43 -04:00
parent 199176c8d0
commit 1c7e73e927

View File

@@ -33,23 +33,16 @@ def train(config_path):
print(f"Loading model: {config['base_model']}")
# Load model with 4-bit quantization (fit in single GPU)
print(f"\n[INFO] Loading {config['base_model']} with BnB 4-bit...")
from transformers import BitsAndBytesConfig
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.float16,
)
# Load model (already 4-bit on disk, just move to GPU)
print(f"\n[INFO] Loading {config['base_model']} to GPU 0...")
model = AutoModelForCausalLM.from_pretrained(
config["base_model"],
device_map="cuda:0",
quantization_config=bnb_config,
torch_dtype=torch.float16,
trust_remote_code=True,
low_cpu_mem_usage=True,
)
print("✓ Success: Model loaded to GPU 0 (4-bit)")
print("✓ Success: Model loaded to GPU 0")
print(f" GPU 0: {torch.cuda.memory_allocated(0) / 1e9:.2f} GB")
print(f" Free VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9 - torch.cuda.memory_allocated(0) / 1e9:.2f} GB")