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

feat: load model with BnB 4-bit quantization (fit in single GPU)

Browse Source
This commit is contained in:
Christian Medina
2026-07-02 18:05:27 -04:00
parent 5a84b0eee5
commit 199176c8d0
1 changed files with 10 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
train.py
View File

@@ -33,16 +33,23 @@ def train(config_path):
print(f"Loading model: {config['base_model']}") print(f"Loading model: {config['base_model']}")
# Load model to single GPU (MoE needs all layers on same device) # Load model with 4-bit quantization (fit in single GPU)
print(f"\n[INFO] Loading {config['base_model']} to 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,
)
model = AutoModelForCausalLM.from_pretrained( model = AutoModelForCausalLM.from_pretrained(
config["base_model"], config["base_model"],
device_map="cuda:0", device_map="cuda:0",
quantization_config=bnb_config,
torch_dtype=torch.float16, torch_dtype=torch.float16,
trust_remote_code=True, trust_remote_code=True,
low_cpu_mem_usage=True, low_cpu_mem_usage=True,
) )
print("✓ Success: Model loaded to GPU 0") print("✓ Success: Model loaded to GPU 0 (4-bit)")
print(f" GPU 0: {torch.cuda.memory_allocated(0) / 1e9:.2f} GB") 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") print(f" Free VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9 - torch.cuda.memory_allocated(0) / 1e9:.2f} GB")