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

This reverts commit 10ba1622f7.

tests/e2e/integrations/test_cut_cross_entropy.py

@@ -44,7 +44,6 @@ def min_cfg(temp_dir):
         "save_safetensors": True,
         "max_steps": 10,
         "bf16": "auto",
-        "save_first_step": False,
     }
 
 
@@ -99,7 +98,6 @@ class TestCutCrossEntropyIntegration:
                 "save_safetensors": True,
                 "max_steps": 10,
                 "bf16": "auto",
-                "save_first_step": False,
             }
         )
         cfg = validate_config(cfg)

tests/e2e/integrations/test_hooks.py

@@ -153,7 +153,6 @@ class TestPluginHooks:
                 "max_steps": 5,
                 "flash_attention": True,
                 "bf16": "auto",
-                "save_first_step": False,
             }
         )
 

tests/e2e/integrations/test_kd.py

@@ -67,7 +67,6 @@ def min_cfg(temp_dir):
         "output_dir": temp_dir,
         "save_safetensors": True,
         "use_tensorboard": True,
-        "save_first_step": False,
     }
 
 

tests/e2e/integrations/test_liger.py

@@ -50,7 +50,6 @@ class LigerIntegrationTestCase:
                 "save_safetensors": True,
                 "bf16": "auto",
                 "max_steps": 5,
-                "save_first_step": False,
             }
         )
         # pylint: disable=duplicate-code
@@ -97,7 +96,6 @@ class LigerIntegrationTestCase:
                 "save_safetensors": True,
                 "bf16": "auto",
                 "max_steps": 5,
-                "save_first_step": False,
             }
         )
         # pylint: disable=duplicate-code

tests/e2e/integrations/test_llm_compressor.py

@@ -81,7 +81,6 @@ class TestLLMCompressorIntegration:
                     },
                     "save_compressed": save_compressed,
                 },
-                "save_first_step": False,
             }
         )