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

feat: add Strategy 1 - 4-bit QLoRA with device_map=auto (distributed across GPUs)

Browse Source
This commit is contained in:
Christian Medina
2026-07-02 11:35:09 -04:00
parent 93ac7391dc
commit 1b3f678b50
1 changed files with 34 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

45
train.py
View File

@@ -38,34 +38,57 @@ def train(config_path):
errors = [] errors = []
# ------------------------------------------------------------------ # ------------------------------------------------------------------
# Strategy 1: QLoRA 4-bit with FSDP (load to CPU, FSDP shards across GPUs) # Strategy 1: QLoRA 4-bit with device_map="auto" (distributed across GPUs, no FSDP)
# ------------------------------------------------------------------ # ------------------------------------------------------------------
print("\n[1/4] Trying: 4-bit QLoRA (FSDP, load to CPU)...") print("\n[1/4] Trying: 4-bit QLoRA (distributed across GPUs, no FSDP)...")
try: try:
bnb_config = BitsAndBytesConfig( bnb_config = BitsAndBytesConfig(
load_in_4bit=True, load_in_4bit=True,
bnb_4bit_quant_type="nf4", bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_compute_dtype=torch.bfloat16,
bnb_4bit_use_double_quant=True, bnb_4bit_use_double_quant=True,
bnb_4bit_quant_storage=torch.bfloat16, # Enable FSDP sharding of 4-bit weights
) )
model = AutoModelForCausalLM.from_pretrained( model = AutoModelForCausalLM.from_pretrained(
config["base_model"], config["base_model"],
quantization_config=bnb_config, quantization_config=bnb_config,
device_map="cpu", # Load to CPU, FSDP shards later device_map="auto", # Distribute layers across GPUs automatically
trust_remote_code=True, trust_remote_code=True,
low_cpu_mem_usage=True, low_cpu_mem_usage=True,
) )
print("✓ Success: QLoRA 4-bit loaded to CPU (FSDP will shard across GPUs)") print("✓ Success: QLoRA 4-bit distributed across GPUs (no FSDP)")
except Exception as e: except Exception as e:
errors.append(("QLoRA 4-bit", e)) errors.append(("QLoRA 4-bit", e))
print(f"✗ Failed: {e}") print(f"✗ Failed: {e}")
# -------------------------------------------------------------- # --------------------------------------------------------------
# Strategy 2: BF16 CPU (model too large for single GPU) # Strategy 2: QLoRA 4-bit with FSDP (load to CPU, FSDP shards across GPUs)
# -------------------------------------------------------------- # --------------------------------------------------------------
print("\n[2/4] Trying: bf16 CPU...") print("\n[2/4] Trying: 4-bit QLoRA (FSDP, load to CPU)...")
try: 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,
bnb_4bit_quant_storage=torch.bfloat16, # Enable FSDP sharding of 4-bit weights
)
model = AutoModelForCausalLM.from_pretrained(
config["base_model"],
quantization_config=bnb_config,
device_map="cpu", # Load to CPU, FSDP shards later
trust_remote_code=True,
low_cpu_mem_usage=True,
)
print("✓ Success: QLoRA 4-bit loaded to CPU (FSDP will shard across GPUs)")
except Exception as e:
errors.append(("QLoRA 4-bit FSDP", e))
print(f"✗ Failed: {e}")
# --------------------------------------------------------------
# Strategy 3: BF16 CPU (model too large for single GPU)
# --------------------------------------------------------------
print("\n[3/4] Trying: bf16 CPU...")
try:
model = AutoModelForCausalLM.from_pretrained( model = AutoModelForCausalLM.from_pretrained(
config["base_model"], config["base_model"],
torch_dtype=torch.bfloat16, torch_dtype=torch.bfloat16,
@@ -79,9 +102,9 @@ def train(config_path):
print(f"✗ Failed: {e}") print(f"✗ Failed: {e}")
# ---------------------------------------------------------- # ----------------------------------------------------------
# Strategy 3: FP16 CPU (fallback) # Strategy 4: FP16 CPU (fallback)
# ---------------------------------------------------------- # ----------------------------------------------------------
print("\n[3/4] Trying: fp16 CPU...") print("\n[4/4] Trying: fp16 CPU...")
try: try:
model = AutoModelForCausalLM.from_pretrained( model = AutoModelForCausalLM.from_pretrained(
config["base_model"], config["base_model"],
@@ -96,9 +119,9 @@ def train(config_path):
print(f"✗ Failed: {e}") print(f"✗ Failed: {e}")
# ------------------------------------------------------ # ------------------------------------------------------
# Strategy 4: 4-bit AS-IS (already quantized) # Strategy 5: 4-bit AS-IS (already quantized)
# ---------------------------------------------------------- # ----------------------------------------------------------
print("\n[4/4] Trying: 4-bit AS-IS...") print("\n[5/5] Trying: 4-bit AS-IS...")
try: try:
model = AutoModelForCausalLM.from_pretrained( model = AutoModelForCausalLM.from_pretrained(
config["base_model"], config["base_model"],