cmedina/agenx-lora-training
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: load BnB 4-bit model to single GPU (no distribution)

Browse Source
This commit is contained in:
Christian Medina
2026-07-02 20:42:05 -04:00
parent a23ecc49f0
commit fa2f21562f
1 changed files with 5 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
train.py
View File

@@ -33,26 +33,18 @@ def train(config_path):
print(f"Loading model: {config['base_model']}")
# Load BnB 4-bit model (already quantized)
print(f"\n[INFO] Loading {config['base_model']} (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 BnB 4-bit model to single GPU
print(f"\n[INFO] Loading {config['base_model']} (BnB 4-bit) to GPU 0...")
model = AutoModelForCausalLM.from_pretrained(
config["base_model"],
quantization_config=bnb_config,
device_map="auto",
device_map="cuda:0",
torch_dtype=torch.float16,
trust_remote_code=True,
low_cpu_mem_usage=True,
)
print("✓ Success: Model loaded (BnB 4-bit)")
print("✓ Success: Model loaded to GPU 0 (BnB 4-bit)")
print(f" GPU 0: {torch.cuda.memory_allocated(0) / 1e9:.2f} GB")
if torch.cuda.device_count() > 1:
print(f" GPU 1: {torch.cuda.memory_allocated(1) / 1e9:.2f} GB")
print(f" Free VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9 - torch.cuda.memory_allocated(0) / 1e9:.2f} GB")
# Add LoRA
lora_config = LoraConfig(