mirror of https://github.com/hpcaitech/ColossalAI
210 lines
9.0 KiB
Python
210 lines
9.0 KiB
Python
import warnings
|
|
from typing import Optional
|
|
|
|
import torch.nn as nn
|
|
|
|
import colossalai
|
|
from colossalai.booster.plugin import GeminiPlugin, LowLevelZeroPlugin
|
|
from colossalai.booster.plugin.low_level_zero_plugin import LowLevelZeroModel
|
|
from colossalai.zero.gemini.gemini_ddp import GeminiDDP
|
|
|
|
from .ddp import DDPStrategy
|
|
|
|
|
|
class LowLevelZeroStrategy(DDPStrategy):
|
|
"""
|
|
The strategy for training with ColossalAI.
|
|
|
|
Args:
|
|
stage(int): The stage to use in ZeRO. Choose in (1, 2)
|
|
precision(str): The precision to use. Choose in ('fp32', 'fp16').
|
|
seed(int): The seed for the random number generator.
|
|
placement_policy(str): The placement policy for gemini. Choose in ('cpu', 'cuda')
|
|
If it is “cpu”, parameters, gradients and optimizer states will be offloaded to CPU,
|
|
If it is “cuda”, they will not be offloaded, which means max CUDA memory will be used. It is the fastest.
|
|
reduce_bucket_size(int): The reduce bucket size in bytes. Only for ZeRO-1 and ZeRO-2.
|
|
overlap_communication(bool): Whether to overlap communication and computation. Only for ZeRO-1 and ZeRO-2.
|
|
initial_scale(float): The initial scale for the optimizer.
|
|
growth_factor(float): The growth factor for the optimizer.
|
|
backoff_factor(float): The backoff factor for the optimizer.
|
|
growth_interval(int): The growth interval for the optimizer.
|
|
hysteresis(int): The hysteresis for the optimizer.
|
|
min_scale(float): The minimum scale for the optimizer.
|
|
max_scale(float): The maximum scale for the optimizer.
|
|
max_norm(float): The maximum norm for the optimizer.
|
|
norm_type(float): The norm type for the optimizer.
|
|
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
stage: int = 2,
|
|
precision: str = "fp16",
|
|
seed: int = 42,
|
|
placement_policy: str = "cuda",
|
|
reduce_bucket_size: int = 12 * 1024**2, # only for stage 1&2
|
|
overlap_communication: bool = True, # only for stage 1&2
|
|
initial_scale: float = 2**16,
|
|
growth_factor: float = 2,
|
|
backoff_factor: float = 0.5,
|
|
growth_interval: int = 1000,
|
|
hysteresis: int = 2,
|
|
min_scale: float = 1,
|
|
max_scale: float = 2**32,
|
|
max_norm: float = 0.0,
|
|
norm_type: float = 2.0,
|
|
) -> None:
|
|
assert stage in (1, 2), f'Unsupported stage "{stage}"'
|
|
assert placement_policy in ("cpu", "cuda"), f'Unsupported placement policy "{placement_policy}"'
|
|
assert precision in ("fp32", "fp16"), f'Unsupported precision "{precision}"'
|
|
|
|
plugin_initializer = lambda: LowLevelZeroPlugin(
|
|
stage=stage,
|
|
precision=precision,
|
|
reduce_bucket_size_in_m=reduce_bucket_size,
|
|
overlap_communication=overlap_communication,
|
|
cpu_offload=(placement_policy == "cpu"),
|
|
initial_scale=initial_scale,
|
|
growth_factor=growth_factor,
|
|
backoff_factor=backoff_factor,
|
|
growth_interval=growth_interval,
|
|
hysteresis=hysteresis,
|
|
min_scale=min_scale,
|
|
max_scale=max_scale,
|
|
max_norm=max_norm,
|
|
norm_type=norm_type,
|
|
)
|
|
|
|
super().__init__(seed, plugin_initializer)
|
|
|
|
def _post_init(self) -> None:
|
|
assert isinstance(
|
|
self.plugin, LowLevelZeroPlugin
|
|
), f"{type(self).__name__}'s plugin is not initialized properly."
|
|
|
|
def setup_distributed(self) -> None:
|
|
colossalai.launch_from_torch({}, seed=self.seed)
|
|
|
|
def unwrap_model(self, model: nn.Module) -> nn.Module:
|
|
assert isinstance(model, LowLevelZeroModel)
|
|
return model.module
|
|
|
|
def get_model_state_dict_shard(self, model: nn.Module, **config):
|
|
assert isinstance(model, LowLevelZeroModel)
|
|
yield from model.state_dict_shard(max_shard_size=1024, only_rank_0=False)
|
|
|
|
|
|
class GeminiStrategy(DDPStrategy):
|
|
"""
|
|
The strategy for training with ColossalAI.
|
|
|
|
Args:
|
|
seed(int): The seed for the random number generator.
|
|
shard_init(bool): Whether to shard the model parameters during initialization. Only for ZeRO-3.
|
|
This is not compatible with `from_pretrained()`. We temporarily disable this and will support it in the future.
|
|
placement_policy(str): The placement policy for gemini. Choose in ('cpu', 'cuda')
|
|
If it is “cpu”, parameters, gradients and optimizer states will be offloaded to CPU,
|
|
If it is “cuda”, they will not be offloaded, which means max CUDA memory will be used. It is the fastest.
|
|
pin_memory(bool): Whether to pin the memory for the data loader. Only for ZeRO-3.
|
|
force_outputs_fp32(bool): Whether to force the outputs to be fp32. Only for ZeRO-3.
|
|
search_range_m(int): The number of search range for the chunk size, divided by 2^20. Only for ZeRO-3.
|
|
hidden_dim(optional, int): The hidden dimension for the gemini. Only for ZeRO-3.
|
|
min_chunk_size_m(float): The minimum chunk size divided by 2^20. Only for ZeRO-3.
|
|
gpu_margin_mem_ratio(float): The margin memory ratio for the GPU. Only for ZeRO-3.
|
|
initial_scale(float): The initial scale for the optimizer.
|
|
growth_factor(float): The growth factor for the optimizer.
|
|
backoff_factor(float): The backoff factor for the optimizer.
|
|
growth_interval(int): The growth interval for the optimizer.
|
|
hysteresis(int): The hysteresis for the optimizer.
|
|
min_scale(float): The minimum scale for the optimizer.
|
|
max_scale(float): The maximum scale for the optimizer.
|
|
max_norm(float): The maximum norm for the optimizer.
|
|
norm_type(float): The norm type for the optimizer.
|
|
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
seed: int = 42,
|
|
shard_init: bool = False, # only for stage 3
|
|
placement_policy: str = "auto",
|
|
shard_param_frac: float = 1.0, # only for static placement
|
|
offload_optim_frac: float = 0.0, # only for static placement
|
|
offload_param_frac: float = 0.0, # only for static placement
|
|
pin_memory: bool = True, # only for stage 3
|
|
force_outputs_fp32: bool = False, # only for stage 3
|
|
search_range_m: int = 32, # only for stage 3
|
|
hidden_dim: Optional[int] = None, # only for stage 3
|
|
min_chunk_size_m: float = 32, # only for stage 3
|
|
gpu_margin_mem_ratio: float = 0.0, # only for stage 3
|
|
initial_scale: float = 2**16,
|
|
growth_factor: float = 2,
|
|
backoff_factor: float = 0.5,
|
|
growth_interval: int = 1000,
|
|
hysteresis: int = 2,
|
|
min_scale: float = 1,
|
|
max_scale: float = 2**32,
|
|
max_norm: float = 0.0,
|
|
norm_type: float = 2.0,
|
|
) -> None:
|
|
# TODO(ver217): support shard_init when using from_pretrained()
|
|
if shard_init:
|
|
warnings.warn(
|
|
f"Shard init is not supported model.from_pretrained() yet. "
|
|
"Please load weights after strategy.prepare()"
|
|
)
|
|
self.shard_init = shard_init
|
|
|
|
warnings.warn(f"Stage 3 only supports fp16. Precision is set to fp16.")
|
|
|
|
# colossalai has changed api for get_current_device in 0.3.4 version or newer
|
|
try:
|
|
from colossalai.accelerator import get_accelerator
|
|
|
|
chunk_init_device = get_accelerator().get_current_device()
|
|
except:
|
|
from colossalai.utils import get_current_device
|
|
|
|
chunk_init_device = get_current_device()
|
|
|
|
# NOTE: dist should be initialized before calling get_current_device()
|
|
plugin_initializer = lambda: GeminiPlugin(
|
|
chunk_init_device=chunk_init_device,
|
|
placement_policy=placement_policy,
|
|
shard_param_frac=shard_param_frac,
|
|
offload_optim_frac=offload_optim_frac,
|
|
offload_param_frac=offload_param_frac,
|
|
precision="fp16",
|
|
pin_memory=pin_memory,
|
|
force_outputs_fp32=force_outputs_fp32,
|
|
strict_ddp_mode=shard_init,
|
|
search_range_m=search_range_m,
|
|
hidden_dim=hidden_dim,
|
|
min_chunk_size_m=min_chunk_size_m,
|
|
gpu_margin_mem_ratio=gpu_margin_mem_ratio,
|
|
initial_scale=initial_scale,
|
|
growth_factor=growth_factor,
|
|
backoff_factor=backoff_factor,
|
|
growth_interval=growth_interval,
|
|
hysteresis=hysteresis,
|
|
min_scale=min_scale,
|
|
max_scale=max_scale,
|
|
max_norm=max_norm,
|
|
norm_type=norm_type,
|
|
)
|
|
|
|
super().__init__(seed, plugin_initializer)
|
|
|
|
def _post_init(self) -> None:
|
|
assert isinstance(self.plugin, GeminiPlugin), f"{type(self).__name__}'s plugin is not initialized properly."
|
|
|
|
def setup_distributed(self) -> None:
|
|
colossalai.launch_from_torch({}, seed=self.seed)
|
|
|
|
def model_init_context(self):
|
|
return super().model_init_context()
|
|
|
|
def unwrap_model(self, model: nn.Module) -> nn.Module:
|
|
assert isinstance(model, GeminiDDP)
|
|
return model.module
|