Repository: allenai/OLMo
File: olmo/checkpoint.py (Python)

Code before:
        optim: Optimizer,
        trainer_state: Dict[str, Any],
        *,
        upload_to: Optional[str] = None,
    ) -> None:
        with self._temporary_wd(dir) as checkpoint_dir:
            if isinstance(dist_model, FSDP):
                with FSDP.state_dict_type(
                    dist_model,
                    state_dict_type=StateDictType.FULL_STATE_DICT,
                    state_dict_config=FullStateDictConfig(rank0_only=True, offload_to_cpu=True),
                    optim_state_dict_config=FullOptimStateDictConfig(rank0_only=True, offload_to_cpu=True),
                ):
                    # We'll write the model and optimizer state dicts individually to reduce (CPU) memory consumption.
                    # First the model state.
                    model_state_dict = dist_model.state_dict()
                    self._write_model_dict(
                        model_state_dict, checkpoint_dir, upload_to, save_overwrite=self.cfg.save_overwrite
                    )

                    # Then the optimizer state.
                    optim_state_dict = FSDP.optim_state_dict(dist_model, optim)
                    self._write_optim_dict(
                        optim_state_dict, checkpoint_dir, upload_to, save_overwrite=self.cfg.save_overwrite
                    )
            elif isinstance(dist_model, DDP) or isinstance(dist_model, SingleAccelerator):
                # _write_model_dict and _write_optim_dict only write checkpoints for rank 0
                # First, get the model state dict from DDP wrapped model
                model_state_dict = dist_model.module.state_dict()
                self._write_model_dict(
                    model_state_dict, checkpoint_dir, upload_to, save_overwrite=self.cfg.save_overwrite
                )

                # Then get the optimizer state dict
                optim_state_dict = optim.state_dict()
                self._write_optim_dict(
                    optim_state_dict, checkpoint_dir, upload_to, save_overwrite=self.cfg.save_overwrite
                )
            else:
                log.info(
                    "`FullCheckpointer.save_checkpoint` only supported for FSDP, DDP, and SingleAccelerator distributed strategies!"
                )

            # Save trainer state.
            if get_global_rank() == 0:
                log.info("Saving trainer state...")
                save_state_dict(
                    checkpoint_dir,
                    "train.pt",
                    trainer_state,
                    upload_to=upload_to,

Review comment (suggestion):
```suggestion
            elif isinstance(dist_model, (DDP, SingleAccelerator)):
```

Code after:
        optim: Optimizer,
        trainer_state: Dict[str, Any],
        *,
        upload_to: Optional[str] = None,
    ) -> None:
        with self._temporary_wd(dir) as checkpoint_dir:
            if isinstance(dist_model, FSDP):
                with FSDP.state_dict_type(
                    dist_model,
                    state_dict_type=StateDictType.FULL_STATE_DICT,
                    state_dict_config=FullStateDictConfig(rank0_only=True, offload_to_cpu=True),
                    optim_state_dict_config=FullOptimStateDictConfig(rank0_only=True, offload_to_cpu=True),
                ):
                    # We'll write the model and optimizer state dicts individually to reduce (CPU) memory consumption.
                    # First the model state.
                    model_state_dict = dist_model.state_dict()
                    self._write_model_dict(
                        model_state_dict, checkpoint_dir, upload_to, save_overwrite=self.cfg.save_overwrite
                    )

                    # Then the optimizer state.
                    optim_state_dict = FSDP.optim_state_dict(dist_model, optim)
                    self._write_optim_dict(
                        optim_state_dict, checkpoint_dir, upload_to, save_overwrite=self.cfg.save_overwrite
                    )
            elif isinstance(dist_model, (DDP, SingleAccelerator)):
                # _write_model_dict and _write_optim_dict only write checkpoints for rank 0
                # First, get the model state dict from DDP wrapped model
                model_state_dict = dist_model.module.state_dict()
                self._write_model_dict(
                    model_state_dict, checkpoint_dir, upload_to, save_overwrite=self.cfg.save_overwrite
                )

                # Then get the optimizer state dict
                optim_state_dict = optim.state_dict()
                self._write_optim_dict(
                    optim_state_dict, checkpoint_dir, upload_to, save_overwrite=self.cfg.save_overwrite
                )
            else:
                log.info(
                    "`FullCheckpointer.save_checkpoint` only supported for FSDP, DDP, and SingleAccelerator distributed strategies!"
                )

            # Save trainer state.
            if get_global_rank() == 0:
                log.info("Saving trainer state...")
                save_state_dict(
                    checkpoint_dir,
                    "train.pt",
                    trainer_state,
                    upload_to=upload_to,