| import math | |
| import warnings | |
| from collections.abc import Sequence | |
| from functools import partial | |
| from typing import Any, Callable, Optional, Tuple, Union | |
| import torch | |
| from torch import nn | |
| from .fc import FC_CLASS_REGISTRY | |
| from .norm import NORM_CLASS_REGISTRY | |
| try: | |
| import transformer_engine.pytorch as te | |
| except: | |
| te = None | |
| def torch_default_param_init_fn_(module: nn.Module, **kwargs: Any) -> None: | |
| del kwargs | |
| if hasattr(module, 'reset_parameters') and isinstance(module.reset_parameters, Callable): | |
| module.reset_parameters() | |
| def fused_init_helper_(module: nn.Module, init_fn_: Callable) -> None: | |
| _fused = getattr(module, '_fused', None) | |
| if _fused is None: | |
| raise RuntimeError(f'Internal logic error') | |
| assert isinstance(module.weight, torch.Tensor) | |
| (dim, splits) = _fused | |
| splits = (0, *splits, module.weight.size(dim)) | |
| for (s, e) in zip(splits[:-1], splits[1:]): | |
| slice_indices = [slice(None)] * module.weight.ndim | |
| slice_indices[dim] = slice(s, e) | |
| init_fn_(module.weight[slice_indices]) | |
| def generic_param_init_fn_(module: nn.Module, init_fn_: Callable, n_layers: int, d_model: Optional[int]=None, init_div_is_residual: Union[int, float, str, bool]=True, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]]=None, **kwargs: Any) -> None: | |
| del kwargs | |
| init_div_is_residual = init_div_is_residual | |
| if init_div_is_residual is False: | |
| div_is_residual = 1.0 | |
| elif init_div_is_residual is True: | |
| div_is_residual = math.sqrt(2 * n_layers) | |
| elif isinstance(init_div_is_residual, float) or isinstance(init_div_is_residual, int): | |
| div_is_residual = init_div_is_residual | |
| elif init_div_is_residual.isnumeric(): | |
| div_is_residual = float(init_div_is_residual) | |
| else: | |
| div_is_residual = 1.0 | |
| raise ValueError(f'Expected init_div_is_residual to be boolean or numeric, got {init_div_is_residual}') | |
| if isinstance(module, tuple(set(FC_CLASS_REGISTRY.values()))): | |
| if hasattr(module, '_fused'): | |
| fused_init_helper_(module, init_fn_) | |
| else: | |
| init_fn_(module.weight) | |
| if module.bias is not None: | |
| assert isinstance(module.bias, torch.Tensor) | |
| torch.nn.init.zeros_(module.bias) | |
| if init_div_is_residual is not False and getattr(module, '_is_residual', False): | |
| with torch.no_grad(): | |
| module.weight.div_(div_is_residual) | |
| elif isinstance(module, nn.Embedding): | |
| if emb_init_std is not None: | |
| std = emb_init_std | |
| if std == 0: | |
| warnings.warn(f'Embedding layer initialized to 0.') | |
| emb_init_fn_ = partial(torch.nn.init.normal_, mean=0.0, std=std) | |
| elif emb_init_uniform_lim is not None: | |
| lim = emb_init_uniform_lim | |
| if isinstance(lim, Sequence): | |
| if len(lim) > 2: | |
| raise ValueError(f'Uniform init requires a min and a max limit. User input: {lim}.') | |
| if lim[0] == lim[1]: | |
| warnings.warn(f'Embedding layer initialized to {lim[0]}.') | |
| else: | |
| if lim == 0: | |
| warnings.warn(f'Embedding layer initialized to 0.') | |
| lim = [-lim, lim] | |
| (a, b) = lim | |
| emb_init_fn_ = partial(torch.nn.init.uniform_, a=a, b=b) | |
| else: | |
| emb_init_fn_ = init_fn_ | |
| emb_init_fn_(module.weight) | |
| elif isinstance(module, tuple(set(NORM_CLASS_REGISTRY.values()))): | |
| if hasattr(module, 'weight') and isinstance(module.weight, torch.Tensor): | |
| torch.nn.init.ones_(module.weight) | |
| if hasattr(module, 'bias') and isinstance(module.bias, torch.Tensor): | |
| torch.nn.init.zeros_(module.bias) | |
| elif isinstance(module, nn.MultiheadAttention): | |
| if module._qkv_same_embed_dim: | |
| assert module.in_proj_weight is not None | |
| assert module.q_proj_weight is None and module.k_proj_weight is None and (module.v_proj_weight is None) | |
| assert d_model is not None | |
| _d = d_model | |
| splits = (0, _d, 2 * _d, 3 * _d) | |
| for (s, e) in zip(splits[:-1], splits[1:]): | |
| init_fn_(module.in_proj_weight[s:e]) | |
| else: | |
| assert module.q_proj_weight is not None and module.k_proj_weight is not None and (module.v_proj_weight is not None) | |
| assert module.in_proj_weight is None | |
| init_fn_(module.q_proj_weight) | |
| init_fn_(module.k_proj_weight) | |
| init_fn_(module.v_proj_weight) | |
| if module.in_proj_bias is not None: | |
| torch.nn.init.zeros_(module.in_proj_bias) | |
| if module.bias_k is not None: | |
| torch.nn.init.zeros_(module.bias_k) | |
| if module.bias_v is not None: | |
| torch.nn.init.zeros_(module.bias_v) | |
| init_fn_(module.out_proj.weight) | |
| if init_div_is_residual is not False and getattr(module.out_proj, '_is_residual', False): | |
| with torch.no_grad(): | |
| module.out_proj.weight.div_(div_is_residual) | |
| if module.out_proj.bias is not None: | |
| torch.nn.init.zeros_(module.out_proj.bias) | |
| elif te is not None and isinstance(module, te.LayerNormMLP): | |
| if isinstance(module.layer_norm_weight, torch.Tensor): | |
| torch.nn.init.ones_(module.layer_norm_weight) | |
| if isinstance(module.layer_norm_bias, torch.Tensor): | |
| torch.nn.init.zeros_(module.layer_norm_bias) | |
| init_fn_(module.fc1_weight) | |
| if module.fc1_bias is not None: | |
| assert isinstance(module.fc1_bias, torch.Tensor) | |
| torch.nn.init.zeros_(module.fc1_bias) | |
| init_fn_(module.fc2_weight) | |
| if module.fc2_bias is not None: | |
| assert isinstance(module.fc2_bias, torch.Tensor) | |
| torch.nn.init.zeros_(module.fc2_bias) | |
| with torch.no_grad(): | |
| module.fc2_weight.div_(div_is_residual) | |
| else: | |
| for _ in module.parameters(recurse=False): | |
| raise NotImplementedError(f'{module.__class__.__name__} parameters are not initialized by param_init_fn.') | |
| def _normal_init_(std: float, mean: float=0.0) -> Callable: | |
| return partial(torch.nn.init.normal_, mean=mean, std=std) | |
| def _normal_param_init_fn_(module: nn.Module, std: float, n_layers: int, d_model: Optional[int]=None, init_div_is_residual: Union[int, float, str, bool]=True, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]]=None, **kwargs: Any) -> None: | |
| del kwargs | |
| init_fn_ = _normal_init_(std=std) | |
| generic_param_init_fn_(module=module, init_fn_=init_fn_, d_model=d_model, n_layers=n_layers, init_div_is_residual=init_div_is_residual, emb_init_std=emb_init_std, emb_init_uniform_lim=emb_init_uniform_lim) | |
| def baseline_param_init_fn_(module: nn.Module, init_std: Optional[float], n_layers: int, d_model: Optional[int]=None, init_div_is_residual: Union[int, float, str, bool]=True, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]]=None, **kwargs: Any) -> None: | |
| del kwargs | |
| if init_std is None: | |
| raise ValueError("You must set model.init_config['init_std'] to a float value to use the default initialization scheme.") | |
| _normal_param_init_fn_(module=module, std=init_std, d_model=d_model, n_layers=n_layers, init_div_is_residual=init_div_is_residual, emb_init_std=emb_init_std, emb_init_uniform_lim=emb_init_uniform_lim) | |
| def small_param_init_fn_(module: nn.Module, n_layers: int, d_model: int, init_div_is_residual: Union[int, float, str, bool]=True, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]]=None, **kwargs: Any) -> None: | |
| del kwargs | |
| std = math.sqrt(2 / (5 * d_model)) | |
| _normal_param_init_fn_(module=module, std=std, d_model=d_model, n_layers=n_layers, init_div_is_residual=init_div_is_residual, emb_init_std=emb_init_std, emb_init_uniform_lim=emb_init_uniform_lim) | |
| def neox_param_init_fn_(module: nn.Module, n_layers: int, d_model: int, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]]=None, **kwargs: Any) -> None: | |
| """From section 2.3.1 of GPT-NeoX-20B: | |
| An Open-Source AutoregressiveLanguage Model — Black et. al. (2022) | |
| see https://github.com/EleutherAI/gpt-neox/blob/9610391ab319403cef079b438edd016a2443af54/megatron/model/init_functions.py#L151 | |
| and https://github.com/EleutherAI/gpt-neox/blob/main/megatron/model/transformer.py | |
| """ | |
| del kwargs | |
| residual_div = n_layers / math.sqrt(10) | |
| small_param_init_fn_(module=module, d_model=d_model, n_layers=n_layers, init_div_is_residual=residual_div, emb_init_std=emb_init_std, emb_init_uniform_lim=emb_init_uniform_lim) | |
| def kaiming_uniform_param_init_fn_(module: nn.Module, n_layers: int, d_model: Optional[int]=None, init_div_is_residual: Union[int, float, str, bool]=True, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]]=None, init_gain: float=0, fan_mode: str='fan_in', init_nonlinearity: str='leaky_relu', **kwargs: Any) -> None: | |
| del kwargs | |
| kaiming_uniform_ = partial(nn.init.kaiming_uniform_, a=init_gain, mode=fan_mode, nonlinearity=init_nonlinearity) | |
| generic_param_init_fn_(module=module, init_fn_=kaiming_uniform_, d_model=d_model, n_layers=n_layers, init_div_is_residual=init_div_is_residual, emb_init_std=emb_init_std, emb_init_uniform_lim=emb_init_uniform_lim) | |
| def kaiming_normal_param_init_fn_(module: nn.Module, n_layers: int, d_model: Optional[int]=None, init_div_is_residual: Union[int, float, str, bool]=True, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]]=None, init_gain: float=0, fan_mode: str='fan_in', init_nonlinearity: str='leaky_relu', **kwargs: Any) -> None: | |
| del kwargs | |
| kaiming_normal_ = partial(torch.nn.init.kaiming_normal_, a=init_gain, mode=fan_mode, nonlinearity=init_nonlinearity) | |
| generic_param_init_fn_(module=module, init_fn_=kaiming_normal_, d_model=d_model, n_layers=n_layers, init_div_is_residual=init_div_is_residual, emb_init_std=emb_init_std, emb_init_uniform_lim=emb_init_uniform_lim) | |
| def xavier_uniform_param_init_fn_(module: nn.Module, n_layers: int, d_model: Optional[int]=None, init_div_is_residual: Union[int, float, str, bool]=True, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]]=None, init_gain: float=0, **kwargs: Any) -> None: | |
| del kwargs | |
| xavier_uniform_ = partial(torch.nn.init.xavier_uniform_, gain=init_gain) | |
| generic_param_init_fn_(module=module, init_fn_=xavier_uniform_, d_model=d_model, n_layers=n_layers, init_div_is_residual=init_div_is_residual, emb_init_std=emb_init_std, emb_init_uniform_lim=emb_init_uniform_lim) | |
| def xavier_normal_param_init_fn_(module: nn.Module, n_layers: int, d_model: Optional[int]=None, init_div_is_residual: Union[int, float, str, bool]=True, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]]=None, init_gain: float=0, **kwargs: Any) -> None: | |
| del kwargs | |
| xavier_normal_ = partial(torch.nn.init.xavier_normal_, gain=init_gain) | |
| generic_param_init_fn_(module=module, init_fn_=xavier_normal_, d_model=d_model, n_layers=n_layers, init_div_is_residual=init_div_is_residual, emb_init_std=emb_init_std, emb_init_uniform_lim=emb_init_uniform_lim) | |
| MODEL_INIT_REGISTRY = {'default_': torch_default_param_init_fn_, 'baseline_': baseline_param_init_fn_, 'kaiming_uniform_': kaiming_uniform_param_init_fn_, 'kaiming_normal_': kaiming_normal_param_init_fn_, 'neox_init_': neox_param_init_fn_, 'small_init_': small_param_init_fn_, 'xavier_uniform_': xavier_uniform_param_init_fn_, 'xavier_normal_': xavier_normal_param_init_fn_} |