def save_checkpoint(self, filename, extra_state):
        """Save all training state in a checkpoint file."""
        logger.info(f"Saving checkpoint to {filename}")
        
        state_dict = self.state_dict(exclude_model_opt=True)
        state_dict["extra_state"].update(extra_state)

        self.model.save_checkpoint(save_dir=filename, client_state=state_dict)

        logger.info(f"Finished saving checkpoint to {filename}")
Save all training state in a checkpoint file.
kosmos-2/fairseq/fairseq/ds_trainer.py