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/multigpu/patched/test_sp.py

@@ -69,6 +69,7 @@ class TestSequenceParallelism:
                 "use_tensorboard": True,
                 "sequence_parallel_degree": 2,
                 "ring_attn_func": ring_attn_func,
+                "save_first_step": False,
             }
         )