minor fix

arealite/api/cli_args.py

@@ -322,7 +322,23 @@ class ExperimentSaveEvalControl:
     total_train_epochs: int = field(
         default=1, metadata={"help": "Total number of epochs to train the model."}
     )
-
+    # Save control
+    freq_epochs: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": "Save frequency in epochs. None disables epoch-based saving."
+        },
+    )
+    freq_steps: Optional[int] = field(
+        default=None,
+        metadata={"help": "Save frequency in steps. None disables step-based saving."},
+    )
+    freq_secs: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": "Save frequency in seconds. None disables time-based saving."
+        },
+    )
     # Checkpointing control
     ckpt_freq_epochs: Optional[int] = field(
         default=None,
@@ -379,30 +395,6 @@ class ExperimentSaveEvalControl:
     )
 
 
-@dataclass
-class SaverConfig:
-    experiment_name: str
-    trial_name: str
-    fileroot: str
-    # Save control
-    freq_epochs: Optional[int] = field(
-        default=None,
-        metadata={
-            "help": "Save frequency in epochs. None disables epoch-based saving."
-        },
-    )
-    freq_steps: Optional[int] = field(
-        default=None,
-        metadata={"help": "Save frequency in steps. None disables step-based saving."},
-    )
-    freq_secs: Optional[int] = field(
-        default=None,
-        metadata={
-            "help": "Save frequency in seconds. None disables time-based saving."
-        },
-    )
-
-
 @dataclass
 class WandBConfig:
     mode: str = "disabled"

examples/arealite/gsm8k_sft.py

@@ -85,6 +85,4 @@ def main_sft():
 
 
 if __name__ == "__main__":
-    args = parse_args()
-    trainer = SFTTrainer(args)
-    trainer.train()
+    main_sft()