checkpoint model on first step callback (#2906)

* checkpoint model on first step callback

* remove debug

* add test cases; update existing tests not to save on first step

* move test out of solo

* delete

* default to False

* typo

tests/e2e/test_qat.py

@@ -57,6 +57,7 @@ class TestQATLlama:
                 "max_steps": 5,
                 "save_safetensors": True,
                 "bf16": True,
+                "save_first_step": False,
             }
         )
         cfg = validate_config(cfg)
@@ -115,6 +116,7 @@ class TestQATLlama:
                     "weight_dtype": "int8",
                     "group_size": 8,
                 },
+                "save_first_step": False,
             }
         )
         cfg = validate_config(cfg)