Nethermind/Mpt-Instruct-DotNet-XS

1# Copyright 2022 MosaicML Examples authors2# SPDX-License-Identifier: Apache-2.03import math4import warnings5from collections.abc import Sequence6from functools import partial7from typing import Optional, Tuple, Union8 9import torch10from torch import nn11 12 13def torch_default_param_init_fn_(14    module: nn.Module,15    verbose: int = 0,16    **kwargs,17):18    del kwargs  # unused, just to capture any extra args from the config19    if verbose > 1:20        warnings.warn(21            f"Initializing network using module's reset_parameters attribute")22 23    if hasattr(module, 'reset_parameters'):24        module.reset_parameters()  # type: ignore25 26 27def fused_init_helper_(module: nn.Module, init_fn_):28    # parameter initialization is often based on the parameters shape.29    # If a layer is fused, initialization should be based on the shapes30    # of the original tensor instead of the shape of the fused tensor.31    # Layers which are fused should have the _fused attibute defined.32    # The first element of _fused is the dimension along which the tensor is fused.33    # This is followed by an iterable of split indices."34 35    _fused = getattr(module, '_fused', None)36 37    if _fused is None:38        raise RuntimeError(f'Internal logic error')39 40    dim, splits = _fused41    splits = (0, *splits, module.weight.size(dim))  # type: ignore42    for s, e in zip(splits[:-1], splits[1:]):43        slice_indices = [slice(None)] * module.weight.ndim  # type: ignore44        slice_indices[dim] = slice(s, e)45        init_fn_(module.weight[slice_indices])  # type: ignore46 47 48def generic_param_init_fn_(49    module: nn.Module,50    init_fn_,51    n_layers: int,52    d_model: Optional[int] = None,53    init_div_is_residual: Union[int, float, str, bool] = True,54    emb_init_std: Optional[float] = None,55    emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]] = None,56    verbose: int = 0,57    **kwargs,58):59    del kwargs  # unused, just to capture any extra args from the config60    if verbose > 1:61        warnings.warn(62            f'If model has bias parameters they are initialized to 0.')63 64    # enable user to divide _is_residual weights by65    # a value which defaults to math.sqrt(2 * cfg.n_layers)66    init_div_is_residual = init_div_is_residual67 68    if init_div_is_residual is False:69        # not used, for pyright70        div_is_residual = 1.071    elif init_div_is_residual is True:72        div_is_residual = math.sqrt(2 * n_layers)73    elif isinstance(init_div_is_residual, float) or isinstance(74            init_div_is_residual, int):75        div_is_residual = init_div_is_residual76    elif isinstance(init_div_is_residual,77                    str) and init_div_is_residual.isnumeric():78        # do not trust YAML parsing to always convert numbers to numbers79        div_is_residual = float(init_div_is_residual)80    else:81        # not used, for pyright82        div_is_residual = 1.083        raise ValueError(84            f'Expected init_div_is_residual to be boolean or numeric, got {init_div_is_residual}'85        )86 87    if init_div_is_residual is not False:88        if verbose > 1:89            warnings.warn(90                f'Initializing _is_residual layers then dividing them by {div_is_residual}.' +\91                f'set `init_div_is_residual: false` in model config to disable this.'92            )93 94    if isinstance(module, nn.Linear):95        # Linear96        if hasattr(module, '_fused'):97            fused_init_helper_(module, init_fn_)98        else:99            init_fn_(module.weight)100        if module.bias is not None:101            torch.nn.init.zeros_(module.bias)102 103        if init_div_is_residual is not False and getattr(104                module, '_is_residual', False):105            with torch.no_grad():106                module.weight.div_(div_is_residual)107 108    elif isinstance(module, nn.Embedding):109        # Embedding110        if emb_init_std is not None:111            std = emb_init_std112            if std == 0:113                warnings.warn(f'Embedding layer initialized to 0.')114            emb_init_fn_ = partial(torch.nn.init.normal_, mean=0.0, std=std)115            if verbose > 1:116                warnings.warn(117                    f'Embedding layer initialized using normal distribution with mean=0 and {std=}.'118                )119        elif emb_init_uniform_lim is not None:120            lim = emb_init_uniform_lim121            if isinstance(lim, Sequence):122                if len(lim) > 2:123                    raise ValueError(124                        f'Uniform init requires a min and a max limit. User input: {lim}.'125                    )126                if lim[0] == lim[1]:127                    warnings.warn(f'Embedding layer initialized to {lim[0]}.')128            else:129                if lim == 0:130                    warnings.warn(f'Embedding layer initialized to 0.')131                lim = [-lim, lim]132            a, b = lim133            emb_init_fn_ = partial(torch.nn.init.uniform_, a=a, b=b)134            if verbose > 1:135                warnings.warn(136                    f'Embedding layer initialized using uniform distribution in range {lim}.'137                )138        else:139            emb_init_fn_ = init_fn_140 141        emb_init_fn_(module.weight)142 143    elif isinstance(module, nn.LayerNorm):144        # LayerNorm145        if verbose > 1:146            warnings.warn(147                f'LayerNorm gamma weights are set to 1. If the layer has a bias it is initialized to 0.'148            )149        torch.nn.init.ones_(module.weight)150        if module.bias is not None:151            torch.nn.init.zeros_(module.bias)152 153    elif isinstance(module, nn.MultiheadAttention):154        # torch's MultiheadAttention155        if module._qkv_same_embed_dim:156            assert module.in_proj_weight is not None157            assert module.q_proj_weight is None and module.k_proj_weight is None and module.v_proj_weight is None158            assert d_model is not None159            # in_proj_weight is actually 3 layers and should be split up for width based init160            _d = d_model161            splits = (0, _d, 2 * _d, 3 * _d)162            for s, e in zip(splits[:-1], splits[1:]):163                init_fn_(module.in_proj_weight[s:e])164        else:165            assert module.q_proj_weight is not None and module.k_proj_weight is not None and module.v_proj_weight is not None166            assert module.in_proj_weight is None167            init_fn_(module.q_proj_weight)168            init_fn_(module.k_proj_weight)169            init_fn_(module.v_proj_weight)170 171        # bias172        if module.in_proj_bias is not None:173            torch.nn.init.zeros_(module.in_proj_bias)174        if module.bias_k is not None:175            torch.nn.init.zeros_(module.bias_k)176        if module.bias_v is not None:177            torch.nn.init.zeros_(module.bias_v)178 179        # out proj180        init_fn_(module.out_proj.weight)181        if init_div_is_residual is not False and getattr(182                module.out_proj, '_is_residual', False):183            with torch.no_grad():184                module.out_proj.weight.div_(div_is_residual)185        if module.out_proj.bias is not None:186            torch.nn.init.zeros_(module.out_proj.bias)187 188    else:189        for _ in module.parameters(recurse=False):190            # raise error if uninitialized module has any parameters191            raise NotImplementedError(192                f'{module.__class__.__name__} parameters are not initialized by param_init_fn.'193            )194 195 196def _normal_init_(std, mean=0.0):197    return partial(torch.nn.init.normal_, mean=mean, std=std)198 199 200def _normal_param_init_fn_(201    module: nn.Module,202    std: float,203    n_layers: int,204    d_model: Optional[int] = None,205    init_div_is_residual: Union[int, float, str, bool] = True,206    emb_init_std: Optional[float] = None,207    emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]] = None,208    verbose: int = 0,209    **kwargs,210):211    del kwargs  # unused, just to capture any extra args from the config212    init_fn_ = _normal_init_(std=std)213 214    if verbose > 1:215        warnings.warn(216            f'Using torch.nn.init.normal_ init fn mean=0.0, std={std}')217 218    generic_param_init_fn_(219        module=module,220        init_fn_=init_fn_,221        d_model=d_model,222        n_layers=n_layers,223        init_div_is_residual=init_div_is_residual,224        emb_init_std=emb_init_std,225        emb_init_uniform_lim=emb_init_uniform_lim,226        verbose=verbose,227    )228 229 230def baseline_param_init_fn_(231    module: nn.Module,232    init_std: float,233    n_layers: int,234    d_model: Optional[int] = None,235    init_div_is_residual: Union[int, float, str, bool] = True,236    emb_init_std: Optional[float] = None,237    emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]] = None,238    verbose: int = 0,239    **kwargs,240):241    del kwargs  # unused, just to capture any extra args from the config242    if init_std is None:243        raise ValueError(244            'You must set model.init_std to a float value to use the default initialization scheme.'245        )246    _normal_param_init_fn_(247        module=module,248        std=init_std,249        d_model=d_model,250        n_layers=n_layers,251        init_div_is_residual=init_div_is_residual,252        emb_init_std=emb_init_std,253        emb_init_uniform_lim=emb_init_uniform_lim,254        verbose=verbose,255    )256 257 258def small_param_init_fn_(259    module: nn.Module,260    n_layers: int,261    d_model: int,262    init_div_is_residual: Union[int, float, str, bool] = True,263    emb_init_std: Optional[float] = None,264    emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]] = None,265    verbose: int = 0,266    **kwargs,267):268    del kwargs  # unused, just to capture any extra args from the config269    # very close to kaiming normal270    # from Transformers without Tears (2019) - Nguyen & Salazar271    std = math.sqrt(2 / (5 * d_model))272    _normal_param_init_fn_(273        module=module,274        std=std,275        d_model=d_model,276        n_layers=n_layers,277        init_div_is_residual=init_div_is_residual,278        emb_init_std=emb_init_std,279        emb_init_uniform_lim=emb_init_uniform_lim,280        verbose=verbose,281    )282 283 284def neox_param_init_fn_(285    module: nn.Module,286    n_layers: int,287    d_model: int,288    emb_init_std: Optional[float] = None,289    emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]] = None,290    verbose: int = 0,291    **kwargs,292):293    """From section 2.3.1 of GPT-NeoX-20B:294 295    An Open-Source AutoregressiveLanguage Model — Black et. al. (2022)296    see https://github.com/EleutherAI/gpt-neox/blob/9610391ab319403cef079b438edd016a2443af54/megatron/model/init_functions.py#L151297    and https://github.com/EleutherAI/gpt-neox/blob/main/megatron/model/transformer.py298    """299    del kwargs  # unused, just to capture any extra args from the config300    residual_div = n_layers / math.sqrt(10)  # small std / wang std301 302    if verbose > 1:303        warnings.warn(f'setting init_div_is_residual to {residual_div}')304 305    small_param_init_fn_(306        module=module,307        d_model=d_model,308        n_layers=n_layers,309        init_div_is_residual=residual_div,310        emb_init_std=emb_init_std,311        emb_init_uniform_lim=emb_init_uniform_lim,312        verbose=verbose,313    )314 315 316def kaiming_uniform_param_init_fn_(317    module: nn.Module,318    n_layers: int,319    d_model: Optional[int] = None,320    init_div_is_residual: Union[int, float, str, bool] = True,321    emb_init_std: Optional[float] = None,322    emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]] = None,323    init_gain: float = 0,324    fan_mode: str = 'fan_in',325    init_nonlinearity: str = 'leaky_relu',326    verbose: int = 0,327    **kwargs,328):329    del kwargs  # unused, just to capture any extra args from the config330 331    if verbose > 1:332        warnings.warn(333            f'Using nn.init.kaiming_uniform_ init fn with parameters: ' +\334            f'a={init_gain}, mode={fan_mode}, nonlinearity={init_nonlinearity}'335        )336 337    kaiming_uniform_ = partial(nn.init.kaiming_uniform_,338                               a=init_gain,339                               mode=fan_mode,340                               nonlinearity=init_nonlinearity)341 342    generic_param_init_fn_(343        module=module,344        init_fn_=kaiming_uniform_,345        d_model=d_model,346        n_layers=n_layers,347        init_div_is_residual=init_div_is_residual,348        emb_init_std=emb_init_std,349        emb_init_uniform_lim=emb_init_uniform_lim,350        verbose=verbose,351    )352 353 354def kaiming_normal_param_init_fn_(355    module: nn.Module,356    n_layers: int,357    d_model: Optional[int] = None,358    init_div_is_residual: Union[int, float, str, bool] = True,359    emb_init_std: Optional[float] = None,360    emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]] = None,361    init_gain: float = 0,362    fan_mode: str = 'fan_in',363    init_nonlinearity: str = 'leaky_relu',364    verbose: int = 0,365    **kwargs,366):367    del kwargs  # unused, just to capture any extra args from the config368 369    if verbose > 1:370        warnings.warn(371            f'Using nn.init.kaiming_normal_ init fn with parameters: ' +\372            f'a={init_gain}, mode={fan_mode}, nonlinearity={init_nonlinearity}'373        )374 375    kaiming_normal_ = partial(torch.nn.init.kaiming_normal_,376                              a=init_gain,377                              mode=fan_mode,378                              nonlinearity=init_nonlinearity)379 380    generic_param_init_fn_(381        module=module,382        init_fn_=kaiming_normal_,383        d_model=d_model,384        n_layers=n_layers,385        init_div_is_residual=init_div_is_residual,386        emb_init_std=emb_init_std,387        emb_init_uniform_lim=emb_init_uniform_lim,388        verbose=verbose,389    )390 391 392def xavier_uniform_param_init_fn_(393    module: nn.Module,394    n_layers: int,395    d_model: Optional[int] = None,396    init_div_is_residual: Union[int, float, str, bool] = True,397    emb_init_std: Optional[float] = None,398    emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]] = None,399    init_gain: float = 0,400    verbose: int = 0,401    **kwargs,402):403    del kwargs  # unused, just to capture any extra args from the config404    xavier_uniform_ = partial(torch.nn.init.xavier_uniform_, gain=init_gain)405 406    if verbose > 1:407        warnings.warn(408            f'Using torch.nn.init.xavier_uniform_ init fn with parameters: ' +\409            f'gain={init_gain}'410        )411 412    generic_param_init_fn_(413        module=module,414        init_fn_=xavier_uniform_,415        d_model=d_model,416        n_layers=n_layers,417        init_div_is_residual=init_div_is_residual,418        emb_init_std=emb_init_std,419        emb_init_uniform_lim=emb_init_uniform_lim,420        verbose=verbose,421    )422 423 424def xavier_normal_param_init_fn_(425    module: nn.Module,426    n_layers: int,427    d_model: Optional[int] = None,428    init_div_is_residual: Union[int, float, str, bool] = True,429    emb_init_std: Optional[float] = None,430    emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]] = None,431    init_gain: float = 0,432    verbose: int = 0,433    **kwargs,434):435    xavier_normal_ = partial(torch.nn.init.xavier_normal_, gain=init_gain)436 437    if verbose > 1:438        warnings.warn(439            f'Using torch.nn.init.xavier_normal_ init fn with parameters: ' +\440            f'gain={init_gain}'441        )442 443    generic_param_init_fn_(444        module=module,445        init_fn_=xavier_normal_,446        d_model=d_model,447        n_layers=n_layers,448        init_div_is_residual=init_div_is_residual,449        emb_init_std=emb_init_std,450        emb_init_uniform_lim=emb_init_uniform_lim,451        verbose=verbose,452    )453 454 455MODEL_INIT_REGISTRY = {456    'default_': torch_default_param_init_fn_,457    'baseline_': baseline_param_init_fn_,458    'kaiming_uniform_': kaiming_uniform_param_init_fn_,459    'kaiming_normal_': kaiming_normal_param_init_fn_,460    'neox_init_': neox_param_init_fn_,461    'small_init_': small_param_init_fn_,462    'xavier_uniform_': xavier_uniform_param_init_fn_,463    'xavier_normal_': xavier_normal_param_init_fn_,464}465