2023-05-18 12:05:59 +00:00
|
|
|
from typing import Callable, Iterator, List, Optional, Tuple, Union
|
2023-03-21 09:39:30 +00:00
|
|
|
|
|
|
|
import torch.nn as nn
|
|
|
|
from torch.nn.parallel import DistributedDataParallel as DDP
|
|
|
|
from torch.optim import Optimizer
|
|
|
|
from torch.optim.lr_scheduler import _LRScheduler as LRScheduler
|
|
|
|
from torch.utils.data import DataLoader
|
|
|
|
|
2023-03-27 02:24:14 +00:00
|
|
|
from colossalai.checkpoint_io import CheckpointIO, GeneralCheckpointIO
|
|
|
|
from colossalai.cluster import DistCoordinator
|
|
|
|
from colossalai.interface import ModelWrapper, OptimizerWrapper
|
2023-03-21 09:39:30 +00:00
|
|
|
|
2023-05-05 11:36:10 +00:00
|
|
|
from .dp_plugin_base import DPPluginBase
|
2023-03-21 09:39:30 +00:00
|
|
|
|
|
|
|
__all__ = ['TorchDDPPlugin']
|
|
|
|
|
|
|
|
|
2023-03-27 02:24:14 +00:00
|
|
|
class TorchDDPCheckpointIO(GeneralCheckpointIO):
|
|
|
|
|
|
|
|
def __init__(self) -> None:
|
|
|
|
super().__init__()
|
|
|
|
self.coordinator = DistCoordinator()
|
|
|
|
|
|
|
|
def load_unsharded_model(self, model: nn.Module, checkpoint: str, strict: bool = True):
|
|
|
|
"""
|
|
|
|
Load model from checkpoint with automatic unwrapping.
|
|
|
|
"""
|
|
|
|
# the model should be unwrapped in self.load_model via ModelWrapper.unwrap
|
|
|
|
return super().load_unsharded_model(model, checkpoint, strict=strict)
|
|
|
|
|
2023-04-06 01:43:51 +00:00
|
|
|
def save_unsharded_model(self, model: nn.Module, checkpoint: str, gather_dtensor: bool, use_safetensors: bool):
|
2023-03-27 02:24:14 +00:00
|
|
|
"""
|
|
|
|
Save model to checkpoint but only on master process.
|
|
|
|
"""
|
|
|
|
# the model should be unwrapped in self.load_model via ModelWrapper.unwrap
|
|
|
|
if self.coordinator.is_master():
|
2023-04-06 01:43:51 +00:00
|
|
|
super().save_unsharded_model(model, checkpoint, gather_dtensor, use_safetensors)
|
2023-03-27 02:24:14 +00:00
|
|
|
|
2023-04-04 07:23:01 +00:00
|
|
|
def save_unsharded_optimizer(self, optimizer: Optimizer, checkpoint: str, gather_dtensor: bool):
|
2023-03-27 02:24:14 +00:00
|
|
|
"""
|
|
|
|
Save optimizer to checkpoint but only on master process.
|
|
|
|
"""
|
|
|
|
if self.coordinator.is_master():
|
2023-04-06 01:43:51 +00:00
|
|
|
super().save_unsharded_optimizer(optimizer, checkpoint, gather_dtensor)
|
2023-03-27 02:24:14 +00:00
|
|
|
|
|
|
|
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)
|
|
|
|
|
2023-05-18 12:05:59 +00:00
|
|
|
def save_sharded_model(self,
|
|
|
|
model: nn.Module,
|
|
|
|
checkpoint_path: str,
|
|
|
|
gather_dtensor: bool = False,
|
|
|
|
variant: Optional[str] = None,
|
|
|
|
max_shard_size: int = 1024,
|
|
|
|
use_safetensors: bool = False):
|
|
|
|
if self.coordinator.is_master():
|
|
|
|
super().save_sharded_model(model, checkpoint_path, gather_dtensor, variant, max_shard_size, use_safetensors)
|
|
|
|
|
2023-03-27 02:24:14 +00:00
|
|
|
|
|
|
|
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
|
|
|
|
|
|
|
|
|
2023-05-05 11:36:10 +00:00
|
|
|
class TorchDDPPlugin(DPPluginBase):
|
2023-03-21 09:39:30 +00:00
|
|
|
"""
|
|
|
|
Plugin for PyTorch DDP.
|
|
|
|
|
|
|
|
Example:
|
|
|
|
>>> from colossalai.booster import Booster
|
|
|
|
>>> from colossalai.booster.plugin import TorchDDPPlugin
|
|
|
|
>>>
|
|
|
|
>>> model, train_dataset, optimizer, criterion = ...
|
|
|
|
>>> plugin = TorchDDPPlugin()
|
|
|
|
|
2023-05-08 07:44:03 +00:00
|
|
|
>>> train_dataloader = plugin.prepare_dataloader(train_dataset, batch_size=8)
|
2023-03-21 09:39:30 +00:00
|
|
|
>>> booster = Booster(plugin=plugin)
|
|
|
|
>>> model, optimizer, train_dataloader, criterion = booster.boost(model, optimizer, train_dataloader, criterion)
|
|
|
|
|
|
|
|
Args:
|
|
|
|
broadcast_buffers (bool, optional): Whether to broadcast buffers in the beginning of training. Defaults to True.
|
|
|
|
bucket_cap_mb (int, optional): The bucket size in MB. Defaults to 25.
|
|
|
|
find_unused_parameters (bool, optional): Whether to find unused parameters. Defaults to False.
|
|
|
|
check_reduction (bool, optional): Whether to check reduction. Defaults to False.
|
|
|
|
gradient_as_bucket_view (bool, optional): Whether to use gradient as bucket view. Defaults to False.
|
|
|
|
static_graph (bool, optional): Whether to use static graph. Defaults to False.
|
|
|
|
"""
|
|
|
|
|
|
|
|
def __init__(self,
|
|
|
|
broadcast_buffers: bool = True,
|
|
|
|
bucket_cap_mb: int = 25,
|
|
|
|
find_unused_parameters: bool = False,
|
|
|
|
check_reduction: bool = False,
|
|
|
|
gradient_as_bucket_view: bool = False,
|
|
|
|
static_graph: bool = False) -> None:
|
2023-05-05 11:36:10 +00:00
|
|
|
super().__init__()
|
2023-03-21 09:39:30 +00:00
|
|
|
self.ddp_kwargs = dict(broadcast_buffers=broadcast_buffers,
|
|
|
|
bucket_cap_mb=bucket_cap_mb,
|
|
|
|
find_unused_parameters=find_unused_parameters,
|
|
|
|
check_reduction=check_reduction,
|
|
|
|
gradient_as_bucket_view=gradient_as_bucket_view,
|
|
|
|
static_graph=static_graph)
|
|
|
|
|
|
|
|
def support_no_sync(self) -> bool:
|
|
|
|
return True
|
|
|
|
|
|
|
|
def control_precision(self) -> bool:
|
|
|
|
return False
|
|
|
|
|
|
|
|
def supported_precisions(self) -> List[str]:
|
|
|
|
return ['fp16', 'fp16_apex', 'bf16', 'fp8']
|
|
|
|
|
|
|
|
def control_device(self) -> bool:
|
|
|
|
return True
|
|
|
|
|
|
|
|
def supported_devices(self) -> List[str]:
|
|
|
|
return ['cuda']
|
|
|
|
|
|
|
|
def configure(
|
|
|
|
self,
|
|
|
|
model: nn.Module,
|
|
|
|
optimizer: Optimizer,
|
|
|
|
criterion: Callable = None,
|
|
|
|
dataloader: DataLoader = None,
|
|
|
|
lr_scheduler: LRScheduler = None,
|
|
|
|
) -> Tuple[Union[nn.Module, OptimizerWrapper, LRScheduler, DataLoader]]:
|
|
|
|
# cast model to cuda
|
|
|
|
model = model.cuda()
|
|
|
|
|
2023-03-27 02:24:14 +00:00
|
|
|
# convert model to sync bn
|
|
|
|
model = nn.SyncBatchNorm.convert_sync_batchnorm(model, None)
|
|
|
|
|
2023-03-21 09:39:30 +00:00
|
|
|
# wrap the model with PyTorch DDP
|
2023-03-27 02:24:14 +00:00
|
|
|
model = TorchDDPModel(model, **self.ddp_kwargs)
|
2023-03-21 09:39:30 +00:00
|
|
|
|
|
|
|
if not isinstance(optimizer, OptimizerWrapper):
|
|
|
|
optimizer = OptimizerWrapper(optimizer)
|
|
|
|
|
|
|
|
return model, optimizer, criterion, dataloader, lr_scheduler
|
2023-03-27 02:24:14 +00:00
|
|
|
|
|
|
|
def control_checkpoint_io(self) -> bool:
|
|
|
|
return True
|
|
|
|
|
|
|
|
def get_checkpoint_io(self) -> CheckpointIO:
|
|
|
|
return TorchDDPCheckpointIO()
|
2023-05-09 03:10:02 +00:00
|
|
|
|
|
|
|
def no_sync(self, model: nn.Module) -> Iterator[None]:
|
|
|
|
assert isinstance(model, TorchDDPModel), 'Model is not boosted by TorchDDPPlugin.'
|
|
|
|
return model.module.no_sync()
|