tests: gradient accum steps = 1, get checkpoint path

Signed-off-by: Will Johnson <mwjohnson728@gmail.com>

Author: willmj

tests/build/test_launch_script.py

@@ -46,7 +46,7 @@
     "num_train_epochs": 5,
     "per_device_train_batch_size": 4,
     "per_device_eval_batch_size": 4,
-    "gradient_accumulation_steps": 4,
+    "gradient_accumulation_steps": 1,
     "learning_rate": 0.00001,
     "weight_decay": 0,
     "warmup_ratio": 0.03,

tests/test_sft_trainer.py

@@ -22,6 +22,7 @@
 import copy
 import json
 import os
+import re
 import tempfile
 
 # Third Party
@@ -87,7 +88,7 @@
     num_train_epochs=5,
     per_device_train_batch_size=4,
     per_device_eval_batch_size=4,
-    gradient_accumulation_steps=4,
+    gradient_accumulation_steps=1,
     learning_rate=0.00001,
     weight_decay=0,
     warmup_ratio=0.03,
@@ -1142,7 +1143,13 @@ def _validate_hf_resource_scanner_file(tempdir):
 
 
 def _get_checkpoint_path(dir_path):
-    return os.path.join(dir_path, "checkpoint-5")
+    checkpoint_dirs = [
+        d
+        for d in os.listdir(dir_path)
+        if os.path.isdir(os.path.join(dir_path, d)) and re.match(r"^checkpoint-\d+$", d)
+    ]
+    checkpoint_dirs.sort(key=lambda name: int(name.split("-")[-1]))
+    return os.path.join(dir_path, checkpoint_dirs[-1])
 
 
 def _get_adapter_config(dir_path):