Merge branch 'main' into data-fixes

Author: daviswer

code-of-conduct.md

@@ -0,0 +1,3 @@
+# Foundation Model Stack Community Code of Conduct
+
+Please refer to [Foundation Model Stack Community Code of Conduct](https://github.com/foundation-model-stack/foundation-model-stack/blob/main/code-of-conduct.md).

fms_fsdp/utils/train_utils.py

@@ -166,7 +166,7 @@ def train(
             ddp_stats.zero_()
         torch.cuda.reset_peak_memory_stats(device=torch.cuda.current_device())
 
-        if batch_idx % cfg.checkpoint_interval == 0:
+        if batch_idx % cfg.checkpoint_interval == 0 or batch_idx == cfg.num_steps:
             checkpointer.save(
                 batch_idx,
                 model,

main_training_llama.py

@@ -169,8 +169,6 @@ def main(**kwargs):
         tokens_seen,
     )
 
-    checkpointer.save_single_file(cfg.num_steps, model)
-
     dist.barrier()
     dist.destroy_process_group()
 

main_training_mamba.py

@@ -169,8 +169,6 @@ def main(**kwargs):
         tokens_seen,
     )
 
-    checkpointer.save_single_file(cfg.num_steps, model)
-
     dist.barrier()
     dist.destroy_process_group()
 

speculator/train_speculator_utils.py

@@ -412,6 +412,7 @@ def train_speculator(
 
         if (
             batch_idx % cfg.checkpoint_interval == 0
+            or batch_idx == cfg.num_steps
             or do_ckpt(cfg.ckpt_save_path) is True
         ):
             torch.cuda.empty_cache()
@@ -425,13 +426,6 @@ def train_speculator(
             torch.cuda.empty_cache()
             do_ckpt(cfg.ckpt_save_path, reset=True)
 
-    checkpointer.save_single_file(
-        batch_idx,
-        speculator,
-        tokens_seen=elapsed_tokens + n_tok,
-        is_compiled=cfg.use_torch_compile,
-    )
-
 
 class EmbedGPTBigCode(GPTBigCode):
     # Overrides the forward function of GPTBigCode to allow returning embedding vectors