Repository: hpcaitech/ColossalAI
File: colossalai/booster/plugin/torch_ddp_plugin.py (Python)

Code before:
            super().save_unsharded_optimizer(optimizer, checkpoint, gather_dtensor)

    def save_lr_scheduler(self, lr_scheduler: LRScheduler, checkpoint: str):
        """
        Save model to checkpoint but only on master process.
        """
        if self.coordinator.is_master():
            super().save_lr_scheduler(lr_scheduler, checkpoint)

    def save_sharded_model(self,
                           model: nn.Module,
                           checkpoint_path: str,
                           gather_dtensor: bool = True,
                           prefix: Optional[str] = None,
                           max_shard_size: int = 1024,
                           use_safetensors: bool = False):
        """
        Save model to checkpoint but only on master process.
        """
        if self.coordinator.is_master():
            super().save_sharded_model(model, checkpoint_path, gather_dtensor, prefix, max_shard_size, use_safetensors)

    def save_sharded_optimier(self,
                              optimizer: Optimizer,
                              checkpoint: str,
                              gather_dtensor: bool = False,
                              prefix: str = None,
                              size_per_shard: int = 1024):
        """
        Save optimizer to checkpoint but only on master process.
        """
        if self.coordinator.is_master():
            super().save_sharded_optimizer(optimizer, checkpoint, gather_dtensor, prefix, size_per_shard)


class TorchDDPModel(ModelWrapper):

    def __init__(self, module: nn.Module, *args, **kwargs) -> None:
        super().__init__(module)
        self.module = DDP(module, *args, **kwargs)

    def unwrap(self):
        return self.module.module


class TorchDDPPlugin(DPPluginBase):
    """
    Plugin for PyTorch DDP.

    Example:
        >>> from colossalai.booster import Booster

Review comment (suggestion):
I find that the default value of `gather_dtensor` of difference methods are different. Please use the same value.

Code after:
            super().save_unsharded_optimizer(optimizer, checkpoint, gather_dtensor)

    def save_lr_scheduler(self, lr_scheduler: LRScheduler, checkpoint: str):
        """
        Save model to checkpoint but only on master process.
        """
        if self.coordinator.is_master():
            super().save_lr_scheduler(lr_scheduler, checkpoint)

    def save_sharded_model(self,
                           model: nn.Module,
                           checkpoint_path: str,
                           gather_dtensor: bool = True,
                           prefix: Optional[str] = None,
                           max_shard_size: int = 1024,
                           use_safetensors: bool = False):
        """
        Save model to checkpoint but only on master process.
        """
        if self.coordinator.is_master():
            super().save_sharded_model(model, checkpoint_path, gather_dtensor, prefix, max_shard_size, use_safetensors)

    def save_sharded_optimizer(self,
                               optimizer: Optimizer,
                               checkpoint: str,
                               gather_dtensor: bool = True,
                               prefix: Optional[str] = None,
                               size_per_shard: int = 1024):
        """
        Save optimizer to checkpoint but only on master process.
        """
        if self.coordinator.is_master():
            super().save_sharded_optimizer(optimizer, checkpoint, gather_dtensor, prefix, size_per_shard)


class TorchDDPModel(ModelWrapper):

    def __init__(self, module: nn.Module, *args, **kwargs) -> None:
        super().__init__(module)
        self.module = DDP(module, *args, **kwargs)

    def unwrap(self):
        return self.module.module


class TorchDDPPlugin(DPPluginBase):
    """
    Plugin for PyTorch DDP.

    Example:
        >>> from colossalai.booster import Booster