def save_checkpoint(self, filename, extra_state):
        """Save all training state in a checkpoint file."""

        logger.info(f"Saving checkpoint to {os.path.abspath(filename)}")
        # call state_dict on all ranks in case it needs internal communication
        state_dict = utils.move_to_cpu(self.state_dict())
        state_dict["extra_state"].update(extra_state)
        if self.should_save_checkpoint_on_current_rank:
            checkpoint_utils.torch_persistent_save(
                state_dict,
                filename,
                async_write=self.cfg.checkpoint.write_checkpoints_asynchronously,
            )
        logger.info(f"Finished saving checkpoint to {os.path.abspath(filename)}")
Save all training state in a checkpoint file.
understand_icl/fairseq/fairseq/trainer.py