Upload train.py with huggingface_hub

train.py

@@ -10,8 +10,8 @@
 # ///
 """Lexwell SFT — Qwen2.5-3B-Instruct fine-tune on IRAC contract-review corpus.
 
-Runs on HF Jobs. Reads 80 hand-curated SFT rows from Curious-PM/lexwell-contract-irac
-and pushes the LoRA adapter to Curious-PM/lexwell-contract-irac-qwen2.5-3b-lora.
+Memory-tuned for a10g-large (24 GB): gradient checkpointing + batch 2 + grad-accum 2
+keeps effective batch at 4 and total steps at 200 across 10 epochs.
 """
 from datasets import load_dataset
 from peft import LoraConfig
@@ -24,7 +24,6 @@ OUT   = "Curious-PM/lexwell-contract-irac-qwen2.5-3b-lora"
 print(f"Loading dataset {DS}...")
 ds = load_dataset(DS, data_files="lexwell_v2.jsonl", split="train")
 print(f"Loaded {len(ds)} rows")
-print(f"First row keys: {list(ds[0].keys())}")
 
 trainer = SFTTrainer(
     model=BASE,
@@ -38,20 +37,24 @@ trainer = SFTTrainer(
     args=SFTConfig(
         output_dir="lexwell-lora",
         num_train_epochs=10,
-        per_device_train_batch_size=4,
-        gradient_accumulation_steps=1,
+        per_device_train_batch_size=2,
+        gradient_accumulation_steps=2,
         learning_rate=2e-4,
         max_length=2048,
         logging_steps=10,
         eval_strategy="no",
         save_strategy="no",
         bf16=True,
+        gradient_checkpointing=True,
+        gradient_checkpointing_kwargs={"use_reentrant": False},
+        optim="adamw_torch_fused",
+        dataloader_pin_memory=False,
         push_to_hub=True,
         hub_model_id=OUT,
         hub_strategy="end",
         report_to="trackio",
         project="lexwell-irac",
-        run_name="qwen2.5-3b-sft-r16-e10",
+        run_name="qwen2.5-3b-sft-r16-e10-gc",
     ),
 )