Integrate recommended LR and warmup/cosine scheduler into training loop

Author: haasonsaas

config_schema.py

@@ -64,6 +64,18 @@ class TrainingConfig(BaseModel):
     lora_rank: int = Field(
         default=16, ge=1, le=256, description="LoRA rank (adapter dimension)"
     )
+    warmup_steps: int = Field(
+        default=100, ge=0, description="Learning rate warmup steps"
+    )
+    max_steps: int = Field(
+        default=1000, ge=1, description="Total training steps across all rounds"
+    )
+    min_lr: float = Field(
+        default=1e-6, gt=0.0, description="Minimum learning rate floor"
+    )
+    use_recommended_lr: bool = Field(
+        default=False, description="Use Tinker's recommended LR formula instead of manual LR"
+    )
 
     @field_validator("train_file")
     @classmethod

eval_loop_config.json

@@ -14,5 +14,6 @@
   "evalops_test_suite_id": "your-test-suite-id-here",
   "steps_per_round": 1,
   "batch_size": 8,
-  "max_seq_length": 2048
+  "max_seq_length": 2048,
+  "lora_rank": 16
 }

trainer_with_eval.py

@@ -54,10 +54,13 @@
     from config_schema import TrainingConfig, load_and_validate_config
     from data_loader import DataLoader
     from simple_eval import run_simple_evaluation
+    from hyperparam_utils import get_recommended_lr, get_lr_with_warmup
 except ImportError:
     TrainingConfig = None
     DataLoader = None
     run_simple_evaluation = None
+    get_recommended_lr = None
+    get_lr_with_warmup = None
 
 
 def prepare_training_data(
@@ -232,7 +235,8 @@ async def async_main(config_path: str) -> None:
             run_training_round(training_client, datums, learning_rate)
 
             print("Saving model checkpoint...")
-            state_uri = training_client.save_state()
+            weights_uri = training_client.save_weights_for_sampler(name=f"round_{round_idx}")
+            state_uri = weights_uri.result().path if hasattr(weights_uri, 'result') else weights_uri
             print(f"Checkpoint saved at {state_uri}")
 
             print("Running evaluations...")