|
import random |
|
from typing import Optional |
|
|
|
import numpy as np |
|
import torch |
|
from colossalai.booster.plugin import LowLevelZeroPlugin |
|
from colossalai.cluster import ProcessGroupMesh |
|
from torch.utils.data import DataLoader |
|
from torch.utils.data.distributed import DistributedSampler |
|
|
|
DP_AXIS, SP_AXIS = 0, 1 |
|
|
|
|
|
class ZeroSeqParallelPlugin(LowLevelZeroPlugin): |
|
def __init__( |
|
self, |
|
sp_size: int = 1, |
|
stage: int = 2, |
|
precision: str = "fp16", |
|
initial_scale: float = 2**32, |
|
min_scale: float = 1, |
|
growth_factor: float = 2, |
|
backoff_factor: float = 0.5, |
|
growth_interval: int = 1000, |
|
hysteresis: int = 2, |
|
max_scale: float = 2**32, |
|
max_norm: float = 0.0, |
|
norm_type: float = 2.0, |
|
reduce_bucket_size_in_m: int = 12, |
|
communication_dtype: Optional[torch.dtype] = None, |
|
overlap_communication: bool = True, |
|
cpu_offload: bool = False, |
|
master_weights: bool = True, |
|
verbose: bool = False, |
|
) -> None: |
|
super().__init__( |
|
stage=stage, |
|
precision=precision, |
|
initial_scale=initial_scale, |
|
min_scale=min_scale, |
|
growth_factor=growth_factor, |
|
backoff_factor=backoff_factor, |
|
growth_interval=growth_interval, |
|
hysteresis=hysteresis, |
|
max_scale=max_scale, |
|
max_norm=max_norm, |
|
norm_type=norm_type, |
|
reduce_bucket_size_in_m=reduce_bucket_size_in_m, |
|
communication_dtype=communication_dtype, |
|
overlap_communication=overlap_communication, |
|
cpu_offload=cpu_offload, |
|
master_weights=master_weights, |
|
verbose=verbose, |
|
) |
|
self.sp_size = sp_size |
|
assert self.world_size % sp_size == 0, "world_size must be divisible by sp_size" |
|
self.dp_size = self.world_size // sp_size |
|
self.pg_mesh = ProcessGroupMesh(self.dp_size, self.sp_size) |
|
self.dp_group = self.pg_mesh.get_group_along_axis(DP_AXIS) |
|
self.sp_group = self.pg_mesh.get_group_along_axis(SP_AXIS) |
|
self.dp_rank = self.pg_mesh.coordinate(DP_AXIS) |
|
self.sp_rank = self.pg_mesh.coordinate(SP_AXIS) |
|
|
|
def __del__(self): |
|
"""Destroy the prcess groups in ProcessGroupMesh""" |
|
self.pg_mesh.destroy_mesh_process_groups() |
|
|
|
def prepare_dataloader( |
|
self, |
|
dataset, |
|
batch_size, |
|
shuffle=False, |
|
seed=1024, |
|
drop_last=False, |
|
pin_memory=False, |
|
num_workers=0, |
|
distributed_sampler_cls=None, |
|
**kwargs, |
|
): |
|
_kwargs = kwargs.copy() |
|
distributed_sampler_cls = distributed_sampler_cls or DistributedSampler |
|
sampler = distributed_sampler_cls(dataset, num_replicas=self.dp_size, rank=self.dp_rank, shuffle=shuffle) |
|
|
|
|
|
def seed_worker(worker_id): |
|
worker_seed = seed |
|
np.random.seed(worker_seed) |
|
torch.manual_seed(worker_seed) |
|
random.seed(worker_seed) |
|
|
|
return DataLoader( |
|
dataset, |
|
batch_size=batch_size, |
|
sampler=sampler, |
|
worker_init_fn=seed_worker, |
|
drop_last=drop_last, |
|
pin_memory=pin_memory, |
|
num_workers=num_workers, |
|
**_kwargs, |
|
) |
|
|
