|
|
import datetime |
|
|
import pathlib |
|
|
from typing import Optional |
|
|
|
|
|
import torch |
|
|
from diffusers.utils import is_accelerate_available |
|
|
|
|
|
from ..logging import get_logger |
|
|
from ..utils import get_device_info |
|
|
from .base import BaseParallelBackend |
|
|
from .utils import apply_ddp_accelerate |
|
|
|
|
|
|
|
|
if not is_accelerate_available(): |
|
|
raise ImportError( |
|
|
"Please install the accelerate package using `pip install accelerate` to use the AccelerateParallelBackend." |
|
|
) |
|
|
|
|
|
from accelerate import Accelerator |
|
|
from accelerate.data_loader import DataLoader |
|
|
from accelerate.utils import ( |
|
|
DataLoaderConfiguration, |
|
|
DistributedDataParallelKwargs, |
|
|
InitProcessGroupKwargs, |
|
|
ProjectConfiguration, |
|
|
) |
|
|
|
|
|
|
|
|
logger = get_logger() |
|
|
_device_type, _device_module = get_device_info() |
|
|
|
|
|
|
|
|
class AccelerateParallelBackend(BaseParallelBackend): |
|
|
def __init__( |
|
|
self, |
|
|
world_size: int, |
|
|
pp_degree: int = 1, |
|
|
dp_degree: int = 1, |
|
|
dp_shards: int = -1, |
|
|
cp_degree: int = 1, |
|
|
tp_degree: int = 1, |
|
|
backend: str = "nccl", |
|
|
timeout: int = 180, |
|
|
logging_dir: Optional[str] = None, |
|
|
output_dir: Optional[str] = None, |
|
|
gradient_accumulation_steps: Optional[int] = None, |
|
|
) -> None: |
|
|
super().__init__() |
|
|
|
|
|
self._world_size = world_size |
|
|
self._pp_degree = pp_degree |
|
|
self._dp_degree = dp_degree |
|
|
self._dp_shards = dp_shards |
|
|
self._cp_degree = cp_degree |
|
|
self._tp_degree = tp_degree |
|
|
self._output_dir = pathlib.Path(output_dir) if output_dir is not None else None |
|
|
self._logging_dir = ( |
|
|
self._output_dir / logging_dir if output_dir is not None and logging_dir is not None else None |
|
|
) |
|
|
self._backend = backend |
|
|
self._timeout = timeout |
|
|
self._gradient_accumulation_steps = gradient_accumulation_steps |
|
|
|
|
|
if pp_degree > 1 or dp_shards > 1 or cp_degree > 1 or tp_degree > 1: |
|
|
raise ValueError( |
|
|
"AccelerateParallelBackend does not support anything but Distributed Data Parallelism at the moment." |
|
|
) |
|
|
if dp_degree != world_size: |
|
|
raise ValueError("Data parallel degree must be equal to world size.") |
|
|
|
|
|
self._accelerator: Accelerator = None |
|
|
self._mesh: torch.distributed.DeviceMesh = None |
|
|
|
|
|
def apply_ddp(self, model: torch.nn.Module, *args, **kwargs) -> torch.nn.Module: |
|
|
project_config = None |
|
|
ddp_kwargs = None |
|
|
init_process_group_kwargs = None |
|
|
if self._accelerator is None: |
|
|
project_config = ProjectConfiguration(project_dir=self._output_dir, logging_dir=self._logging_dir) |
|
|
ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=False) |
|
|
dataloader_config = DataLoaderConfiguration( |
|
|
split_batches=False, dispatch_batches=False, use_stateful_dataloader=True |
|
|
) |
|
|
init_process_group_kwargs = InitProcessGroupKwargs( |
|
|
backend=self._backend, timeout=datetime.timedelta(seconds=self._timeout) |
|
|
) |
|
|
self._accelerator, model = apply_ddp_accelerate( |
|
|
model, |
|
|
project_config, |
|
|
ddp_kwargs, |
|
|
init_process_group_kwargs, |
|
|
dataloader_config, |
|
|
self._gradient_accumulation_steps, |
|
|
accelerator=self._accelerator, |
|
|
) |
|
|
logger.debug("Applied AccelerateParallel::apply_ddp to model.") |
|
|
return model |
|
|
|
|
|
def prepare_dataset(self, dataset: torch.utils.data.IterableDataset) -> torch.utils.data.IterableDataset: |
|
|
logger.debug("AccelerateParallelBackend::prepare_dataset completed!") |
|
|
return dataset |
|
|
|
|
|
def prepare_dataloader( |
|
|
self, |
|
|
dataset: torch.utils.data.IterableDataset, |
|
|
batch_size: int = 1, |
|
|
num_workers: int = 0, |
|
|
pin_memory: bool = False, |
|
|
) -> DataLoader: |
|
|
dataloader = torch.utils.data.DataLoader( |
|
|
dataset, batch_size=batch_size, num_workers=num_workers, pin_memory=pin_memory |
|
|
) |
|
|
dataloader = self._accelerator.prepare_data_loader(dataloader) |
|
|
logger.debug("AccelerateParallelBackend::prepare_dataloader completed!") |
|
|
return dataloader |
|
|
|
|
|
def prepare_optimizer(self, optimizer, lr_scheduler): |
|
|
optimizer = self._accelerator.prepare_optimizer(optimizer) |
|
|
lr_scheduler = self._accelerator.prepare_scheduler(lr_scheduler) |
|
|
return optimizer, lr_scheduler |
|
|
|
|
|
def get_mesh(self, name: Optional[str] = None) -> torch.distributed.DeviceMesh: |
|
|
def _get_mesh(): |
|
|
if name is None: |
|
|
return self._mesh |
|
|
try: |
|
|
return self._mesh[name] |
|
|
except (KeyError, RuntimeError): |
|
|
return self._mesh |
|
|
|
|
|
if self._mesh is not None: |
|
|
return _get_mesh() |
|
|
|
|
|
mesh_list = [("dp_replicate", self._dp_degree), ("dp_shard", self._dp_shards)] |
|
|
mesh_list = [(name, degree) for name, degree in mesh_list if degree > 1] |
|
|
names = [x[0] for x in mesh_list] |
|
|
degrees = [x[1] for x in mesh_list] |
|
|
mesh = torch.distributed.device_mesh.init_device_mesh(_device_type, mesh_shape=degrees, mesh_dim_names=names) |
|
|
|
|
|
dp_mesh_names, dp_cp_mesh_names, dp_shard_cp_mesh_names = [], [], [] |
|
|
|
|
|
if self.data_replication_enabled: |
|
|
dp_mesh_names.append("dp_replicate") |
|
|
dp_cp_mesh_names.append("dp_replicate") |
|
|
if self.data_sharding_enabled: |
|
|
dp_mesh_names.append("dp_shard") |
|
|
dp_cp_mesh_names.append("dp_shard") |
|
|
dp_shard_cp_mesh_names.append("dp_shard") |
|
|
if self.context_parallel_enabled: |
|
|
dp_cp_mesh_names.append("cp") |
|
|
dp_shard_cp_mesh_names.append("cp") |
|
|
|
|
|
if len(dp_mesh_names) > 0: |
|
|
mesh[tuple(dp_mesh_names)]._flatten(mesh_dim_name="dp") |
|
|
if len(dp_cp_mesh_names) > 0: |
|
|
mesh[tuple(dp_cp_mesh_names)]._flatten(mesh_dim_name="dp_cp") |
|
|
if len(dp_shard_cp_mesh_names) > 0: |
|
|
mesh[tuple(dp_shard_cp_mesh_names)]._flatten(mesh_dim_name="dp_shard_cp") |
|
|
|
|
|
logger.debug(f"Device mesh: {mesh}") |
|
|
self._mesh = mesh |
|
|
return _get_mesh() |
|
|
|
|
|
@property |
|
|
def world_size(self): |
|
|
return self._accelerator.num_processes |
|
|
|
|
|
@property |
|
|
def rank(self): |
|
|
return self._accelerator.process_index |
|
|
|
|
|
@property |
|
|
def local_rank(self): |
|
|
return self._accelerator.local_process_index |
|
|
|
|
|
@property |
|
|
def is_main_process(self): |
|
|
r"""Returns `True` if the current process is the main process on the master node.""" |
|
|
return self._accelerator.is_main_process |
|
|
|
|
|
@property |
|
|
def is_local_main_process(self): |
|
|
r"""Returns `True` if the current process is the main process on local node.""" |
|
|
return self._accelerator.is_local_main_process |
|
|
|
|
|
@property |
|
|
def device(self): |
|
|
return self._accelerator.device |
|
|
|
|
|
def wait_for_everyone(self): |
|
|
self._accelerator.wait_for_everyone() |
|
|
|
|
|
def destroy(self): |
|
|
self._accelerator.end_training() |
|
|
|
|
|
@property |
|
|
def pipeline_parallel_enabled(self): |
|
|
return self._pp_degree > 1 |
|
|
|
|
|
@property |
|
|
def data_parallel_enabled(self): |
|
|
return self._dp_degree > 1 or self._dp_shards > 1 |
|
|
|
|
|
@property |
|
|
def data_replication_enabled(self): |
|
|
return self._dp_degree > 1 |
|
|
|
|
|
@property |
|
|
def data_sharding_enabled(self): |
|
|
return self._dp_shards > 1 |
|
|
|
|
|
@property |
|
|
def context_parallel_enabled(self): |
|
|
return self._cp_degree > 1 |
|
|
|
|
|
@property |
|
|
def tensor_parallel_enabled(self): |
|
|
return self._tp_degree > 1 |
|
|
|
