Switch to packing=True, batch=4, grad_acc=4, step-based checkpoints for speed

train.py

@@ -87,8 +87,6 @@ output_dir = os.environ.get("OUTPUT_DIR", "ssdataanalysis/gemma-4-E4B-hebrew-fir
 print(f"=== Training {model_id} -> {output_dir} ===")
 
 train_dataset = prepare_dataset(hebrew_ratio=0.5, max_total=120000)
-# No eval dataset to avoid OOM during evaluation on A10G 24GB
-# We will rely on training loss and periodic checkpointing
 
 print("Loading tokenizer...")
 tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
@@ -112,28 +110,30 @@ peft_config = LoraConfig(
     exclude_modules=["vision_tower", "multi_modal_projector"],
 )
 
+# Optimized: packing=True, larger batch, step-based checkpoints
 training_args = SFTConfig(
     output_dir=output_dir,
     num_train_epochs=3,
-    per_device_train_batch_size=1,
-    gradient_accumulation_steps=16,
+    per_device_train_batch_size=4,
+    gradient_accumulation_steps=4,
     learning_rate=2e-4,
     lr_scheduler_type="cosine",
     warmup_steps=500,
     weight_decay=0.01,
     max_length=2048,
-    packing=False,
+    packing=True,
     bf16=True,
     logging_strategy="steps",
     logging_steps=10,
     logging_first_step=True,
     eval_strategy="no",
-    save_strategy="epoch",
-    save_total_limit=2,
+    save_strategy="steps",
+    save_steps=500,
+    save_total_limit=3,
     push_to_hub=True,
     hub_model_id=output_dir,
     report_to="trackio",
-    run_name=output_dir,
+    run_name=output_dir + "-fast",
     remove_unused_columns=False,
     disable_tqdm=True,
     dataset_num_proc=4,