Draft: WIP Training checkpointing

Pull Request Title

Training checkpointing

Description

Wants to merge: training_checkpointing into main

Training checkpointing: Save / resume training from checkpoints (DDP / FSDP aware; DDP should only save the model on the main rank, while FSDP will most likely save sharded checkpoints). Save the model + tokenizer at the end of training (this should save the model in an easily loadable format, using either pytorch or safe-tensors, regardless of DDP / FSDP).

Type of change

• Bug fix
• New feature
• Enhancement
• Documentation update
• Other (specify right below)

Related Issues

See issue #21.

Screenshots or GIFs

Checklist

• I have tested the code with the changes manually.
• My code follows the project's style guidelines.
• I have documented my code for others to understand.
• I have updated documentation as needed (including README.md, code comments and doc strings).

Reviewer Guidelines

Additional Notes

@mentions

optimus/trainer.py

@@ -300,13 +300,14 @@ class Trainer():
                     self.progress.update(1)
 
                     if steps_trained_in_current_epoch % self.args.save_steps:
-                        # save checkpoint
-                        print(f'saving checkpoint at step {steps_trained_in_current_epoch}')
-                        os.mkdir(os.path.join(self.args.checkpoints_dir, f'checkpoint-{self.state.global_steps}'), parent=True)
-                        self.state.to_json_file(os.path.join(self.args.checkpoints_dir, os.path.join(f'checkpoint-{self.state.global_steps}', 'trainer_state.json')))
-                        torch.save(self.model.state_dict(),
-                                   os.path.join(self.args.checkpoints_dir,
-                                   os.path.join(f'checkpoint-{self.state.global_steps}', 'model.pt')))
+                        if dist_utils.get_rank() == 0: # TODO this only works for DDP
+                            # save checkpoint
+                            print(f'saving checkpoint at step {steps_trained_in_current_epoch}')
+                            os.makedirs(os.path.join(self.args.checkpoints_dir, f'checkpoint-{self.state.global_steps}'))
+                            self.state.to_json_file(os.path.join(self.args.checkpoints_dir, os.path.join(f'checkpoint-{self.state.global_steps}', 'trainer_state.json')))
+                            torch.save(self.model.state_dict(),
+                                       os.path.join(self.args.checkpoints_dir,
+                                       os.path.join(f'checkpoint-{self.state.global_steps}', 'model.pt')))
 
     def save_model(self, save_dir: Path) -> None:
         """