Revert "checkpoint model on first step callback (#2906)"

This reverts commit 10ba1622f7.

tests/e2e/multigpu/solo/test_flex.py

@@ -61,7 +61,6 @@ class TestPackedFlex:
                 "max_steps": 2,
                 "use_tensorboard": True,
                 "save_strategy": "no",
-                "save_first_step": False,
             }
         )
         if is_torch_bf16_gpu_available():

tests/e2e/multigpu/solo/test_grpo.py

@@ -223,7 +223,6 @@ def oai_gsm8k_transform(cfg, *args, **kwargs):
                 "save_safetensors": True,
                 "bf16": "auto",
                 "use_tensorboard": True,
-                "save_first_step": False,
             }
         )
 
@@ -318,7 +317,6 @@ def oai_gsm8k_transform(cfg, *args, **kwargs):
                 "save_safetensors": True,
                 "bf16": "auto",
                 "use_tensorboard": True,
-                "save_first_step": False,
             }
         )
 
@@ -411,7 +409,6 @@ def oai_gsm8k_transform(cfg, *args, **kwargs):
                 "save_safetensors": True,
                 "bf16": "auto",
                 "use_tensorboard": True,
-                "save_first_step": False,
             }
         )