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
2025-07-15 15:00:48 -04:00
parent d320ef6199
commit 10ba1622f7
146 changed files with 419 additions and 9 deletions
--- a/tests/e2e/patched/test_fa_xentropy.py
+++ b/tests/e2e/patched/test_fa_xentropy.py
@@ -62,6 +62,7 @@ class TestFAXentropyLlama:
                "optimizer": "adamw_8bit",
                "lr_scheduler": "cosine",
                "use_tensorboard": True,
+                "save_first_step": False,
            }
        )
        if is_torch_bf16_gpu_available():