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

fix: use 4-bit model AS-IS + fix FSDP config

Browse Source
This commit is contained in:
Christian Medina
2026-07-02 00:09:03 -04:00
parent d5ccd28e39
commit 11c635b0e0
2 changed files with 25 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files

53
train.py
View File

@@ -37,77 +37,71 @@ def train(config_path):
errors = [] errors = []
# ------------------------------------------------------------------ # ------------------------------------------------------------------
# Strategy 1: QLoRA with FSDP (preferred) # Strategy 1: 4-bit model AS-IS (already quantized)
# ------------------------------------------------------------------ # ------------------------------------------------------------------
print("\n[1/4] Trying: 4-bit QLoRA (FSDP)...") print("\n[1/4] Trying: 4-bit AS-IS (FSDP)...")
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,
)
model = AutoModelForCausalLM.from_pretrained( model = AutoModelForCausalLM.from_pretrained(
config["base_model"], config["base_model"],
quantization_config=bnb_config, torch_dtype=torch.float16,
device_map="auto", device_map="auto",
trust_remote_code=True, trust_remote_code=True,
low_cpu_mem_usage=True, low_cpu_mem_usage=True,
) )
print("✓ Success: QLoRA 4-bit") print("✓ Success: 4-bit AS-IS (FSDP)")
except Exception as e: except Exception as e:
errors.append(("QLoRA 4-bit", e)) errors.append(("4-bit AS-IS", e))
print(f"✗ Failed: {e}") print(f"✗ Failed: {e}")
# -------------------------------------------------------------- # --------------------------------------------------------------
# Strategy 2: BF16 GPU # Strategy 2: 4-bit to CPU
# -------------------------------------------------------------- # --------------------------------------------------------------
print("\n[2/4] Trying: bf16 GPU...") print("\n[2/4] Trying: 4-bit to CPU...")
try: try:
model = AutoModelForCausalLM.from_pretrained( model = AutoModelForCausalLM.from_pretrained(
config["base_model"], config["base_model"],
torch_dtype=torch.bfloat16, torch_dtype=torch.float16,
device_map="auto", device_map="cpu",
trust_remote_code=True, trust_remote_code=True,
low_cpu_mem_usage=True, low_cpu_mem_usage=True,
) )
print("✓ Success: bf16 GPU") print("✓ Success: 4-bit to CPU")
except Exception as e: except Exception as e:
errors.append(("bf16 GPU", e)) errors.append(("4-bit CPU", e))
print(f"✗ Failed: {e}") print(f"✗ Failed: {e}")
# ---------------------------------------------------------- # ----------------------------------------------------------
# Strategy 3: BF16 CPU # Strategy 3: bf16 auto
# ---------------------------------------------------------- # ----------------------------------------------------------
print("\n[3/4] Trying: bf16 CPU...") print("\n[3/4] Trying: bf16 auto...")
try: try:
model = AutoModelForCausalLM.from_pretrained( model = AutoModelForCausalLM.from_pretrained(
config["base_model"], config["base_model"],
torch_dtype=torch.bfloat16, torch_dtype=torch.bfloat16,
device_map="cpu", device_map="auto",
trust_remote_code=True, trust_remote_code=True,
low_cpu_mem_usage=True, low_cpu_mem_usage=True,
) )
print("✓ Success: bf16 CPU") print("✓ Success: bf16 auto")
except Exception as e: except Exception as e:
errors.append(("bf16 CPU", e)) errors.append(("bf16 auto", e))
print(f"✗ Failed: {e}") print(f"✗ Failed: {e}")
# ------------------------------------------------------ # ------------------------------------------------------
# Strategy 4: FP16 GPU # Strategy 4: bf16 CPU
# ------------------------------------------------------ # ------------------------------------------------------
print("\n[4/4] Trying: fp16 GPU...") print("\n[4/4] Trying: bf16 CPU...")
try: try:
model = AutoModelForCausalLM.from_pretrained( model = AutoModelForCausalLM.from_pretrained(
config["base_model"], config["base_model"],
torch_dtype=torch.float16, torch_dtype=torch.bfloat16,
device_map="auto", device_map="cpu",
trust_remote_code=True, trust_remote_code=True,
low_cpu_mem_usage=True, low_cpu_mem_usage=True,
) )
print("✓ Success: fp16 GPU") print("✓ Success: bf16 CPU")
except Exception as e: except Exception as e:
errors.append(("fp16 GPU", e)) errors.append(("bf16 CPU", e))
print(f"✗ Failed: {e}") print(f"✗ Failed: {e}")
msg = "\n".join( msg = "\n".join(
f"{name}: {err}" for name, err in errors f"{name}: {err}" for name, err in errors
@@ -167,13 +161,14 @@ def train(config_path):
eval_steps=config.get("eval_steps", 100), eval_steps=config.get("eval_steps", 100),
bf16=True, bf16=True,
gradient_checkpointing=config.get("gradient_checkpointing", True), gradient_checkpointing=config.get("gradient_checkpointing", True),
fsdp=["full_shard", "auto_wrap"], fsdp=True,
fsdp_config={ fsdp_config={
"backward_prefetch": "backward_pre", "backward_prefetch": "backward_pre",
"forward_prefetch": "true", "forward_prefetch": "true",
"cpu_offload": "false", "cpu_offload": "false",
"auto_wrap_policy": "TRANSFORMER_BASED_WRAP", "auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
"transformer_layer_cls_to_wrap": "Qwen3_5MoeDecoderLayer", "transformer_layer_cls_to_wrap": "Qwen3_5MoeDecoderLayer",
"activation_checkpointing": True,
}, },
) )

2
training/configs/ornith-35b-lora.yaml
View File

@@ -1,7 +1,7 @@
# LoRA Training Configuration for Llama-2-7b # LoRA Training Configuration for Llama-2-7b
# Dataset: cyron_summary_lora_dataset (20k examples) # Dataset: cyron_summary_lora_dataset (20k examples)
base_model: /data/models/Ornith-1.0-35B base_model: /data/models/Ornith-1.0-35B-4bit
model_type: LlamaForCausalLM model_type: LlamaForCausalLM
tokenizer_type: LlamaTokenizer tokenizer_type: LlamaTokenizer