def load_checkpoint(
        self,
        filename,
        reset_optimizer=False,
        reset_lr_scheduler=False,
        optimizer_overrides=None,
        reset_meters=False,
    ):
        """
        Load all training state from a checkpoint file.
        rank = 0 will load the checkpoint, and then broadcast it to all
        other ranks.
        """
        extra_state, self._optim_history, last_optim_state = None, [], None

        logger.info(f"Preparing to load checkpoint {filename}")
        is_distributed = self.data_parallel_world_size > 1
        bexists = os.path.isdir(filename) #PathManager.isfile(filename)
        if not bexists:
            logger.info("No existing checkpoint found {}".format(filename))
            return None

        def load_model(src, dst):
            if torch.distributed.get_rank() == 0:
                print(self.cfg.model)
            dst.load_state_dict(src, strict=False, model_cfg=self.cfg.model)

        load_path, client_states = self.model.load_checkpoint(load_dir=filename, load_optimizer_states=not reset_optimizer, model_f=load_model)

        print(f'[{torch.distributed.get_rank()}] ckpt client states={client_states}')

        assert not utils.has_parameters(self.get_criterion()), "criterion w. params not supported yet"
        extra_state = client_states["extra_state"]

        if not reset_optimizer and not reset_lr_scheduler:
            self.lr_scheduler.load_state_dict(client_states["lr_scheduler_state"])

        self.set_num_updates(client_states["num_updates"])

        self.scaler.loss_scale = client_states["loss_scale"]

        if extra_state is not None:
            itr_state = extra_state["train_iterator"]
            epoch =  itr_state.get("epoch", 1)

            if "previous_training_time" in extra_state:
                self._previous_training_time = extra_state["previous_training_time"]
                self._start_time = time.time()

            self.lr_step(epoch)

            if itr_state.get("version", 1) >= 2 and itr_state["iterations_in_epoch"] == 0:
                # reset meters at start of epoch
                reset_meters = True

            if "metrics" in extra_state and not reset_meters:
                metrics.load_state_dict(extra_state["metrics"])

                # reset TimeMeters, since their start times don't make sense anymore
                for meter in metrics.get_meters("default"):
                    if isinstance(meter, meters.TimeMeter):
                        meter.reset()

            logger.info(
                "Loaded checkpoint {} (epoch {} @ {} updates)".format(
                    filename, epoch, self.get_num_updates()
                )
            )
        return extra_state

        Load all training state from a checkpoint file.
        rank = 0 will load the checkpoint, and then broadcast it to all
        other ranks.
        
kosmos-2/fairseq/fairseq/ds_trainer.py