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

fix: use already-quantized 4-bit model for training (no BnB needed)

Browse Source
This commit is contained in:
Christian Medina
2026-07-02 12:09:59 -04:00
parent 3e26d7d37e
commit f9c748706f
1 changed files with 4 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
train.py
View File

@@ -38,24 +38,18 @@ def train(config_path):
errors = []
# ------------------------------------------------------------------
# Strategy 1: QLoRA 4-bit with device_map="auto" (distributed across GPUs, no FSDP)
# Strategy 1: Load already-quantized 4-bit model (distributed across GPUs)
# ------------------------------------------------------------------
print("\n[1/4] Trying: 4-bit QLoRA (distributed across GPUs, no FSDP)...")
print("\n[1/4] Trying: 4-bit model AS-IS (distributed across GPUs)...")
try:
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16,
bnb_4bit_use_double_quant=True,
)
model = AutoModelForCausalLM.from_pretrained(
config["base_model"],
quantization_config=bnb_config,
device_map="auto", # Distribute layers across GPUs automatically
torch_dtype=torch.float16,
trust_remote_code=True,
low_cpu_mem_usage=True,
)
print("✓ Success: QLoRA 4-bit distributed across GPUs (no FSDP)")
print("✓ Success: 4-bit model distributed across GPUs (no FSDP)")
except Exception as e:
errors.append(("QLoRA 4-bit", e))
print(f"✗ Failed: {e}")