Use save_weights_for_sampler for eval and add LoRA rank to client creation

Author: haasonsaas

config_schema.py

@@ -61,6 +61,9 @@ class TrainingConfig(BaseModel):
     max_seq_length: int = Field(
         default=2048, ge=128, le=32768, description="Maximum sequence length"
     )
+    lora_rank: int = Field(
+        default=16, ge=1, le=256, description="LoRA rank (adapter dimension)"
+    )
 
     @field_validator("train_file")
     @classmethod

trainer_with_eval.py

@@ -208,8 +208,11 @@ async def async_main(config_path: str) -> None:
             except ValueError as e:
                 print(f"Warning: Could not initialize EvalOps client: {e}")
 
-    print(f"Creating LoRA training client for {base_model}...")
-    training_client = service_client.create_lora_training_client(base_model=base_model)
+    print(f"Creating LoRA training client for {base_model} (rank={config.lora_rank})...")
+    training_client = service_client.create_lora_training_client(
+        base_model=base_model,
+        rank=config.lora_rank,
+    )
 
     tokenizer = training_client.get_tokenizer()