Fix: batch_size=16, 10K samples, unbuffered output, 2 epochs

train_on_hf.py

@@ -13,8 +13,13 @@ Usage (on HF with GPU):
 
 import argparse
 import json
+import sys
 from pathlib import Path
 
+# Force unbuffered output so HF Spaces logs show immediately
+sys.stdout.reconfigure(line_buffering=True)
+sys.stderr.reconfigure(line_buffering=True)
+
 
 def setup(hf_token: str):
     """Login and download data from HF."""
@@ -47,7 +52,7 @@ def setup(hf_token: str):
     return data_dir
 
 
-def train_world_model(data_dir: Path, max_samples: int = 50000):
+def train_world_model(data_dir: Path, max_samples: int = 10000):
     """Stage 2: Train World Model via SFT."""
     from datasets import Dataset
     from transformers import AutoModelForCausalLM, AutoTokenizer
@@ -104,20 +109,22 @@ def train_world_model(data_dir: Path, max_samples: int = 50000):
 
     training_args = SFTConfig(
         output_dir="./world_model_checkpoints",
-        num_train_epochs=3,
-        per_device_train_batch_size=8,
+        num_train_epochs=2,
+        per_device_train_batch_size=16,
         gradient_accumulation_steps=2,
         learning_rate=2e-4,
         lr_scheduler_type="cosine",
         warmup_ratio=0.1,
-        logging_steps=10,
+        logging_steps=5,
         eval_strategy="steps",
-        eval_steps=200,
+        eval_steps=100,
         save_steps=500,
         save_total_limit=2,
         fp16=True,
         max_length=512,
         report_to="none",
+        disable_tqdm=False,
+        dataloader_num_workers=0,
     )
 
     trainer = SFTTrainer(