Upload altered non lfs files from Github branch

attention.py

@@ -188,15 +188,15 @@ class GroupedQueryAttention(nn.Module):
     implementation enables user to also use additive bias.
     """
 
-    def __init__(self, d_model: int, n_heads: int, kv_n_heads: int, attn_impl: str='triton', clip_qkv: Optional[float]=None, qk_ln: bool=False, softmax_scale: Optional[float]=None, attn_pdrop: float=0.0, norm_type: str='low_precision_layernorm', fc_type: str='torch', device: Optional[str]=None):
+    def __init__(self, hidden_size: int, n_heads: int, kv_n_heads: int, attn_impl: str='triton', clip_qkv: Optional[float]=None, qk_ln: bool=False, softmax_scale: Optional[float]=None, attn_pdrop: float=0.0, norm_type: str='low_precision_layernorm', fc_type: str='torch', device: Optional[str]=None):
         super().__init__()
         self.attn_impl = attn_impl
         self.clip_qkv = clip_qkv
         self.qk_ln = qk_ln
-        self.d_model = d_model
+        self.hidden_size = hidden_size
         self.n_heads = n_heads
         self.kv_n_heads = kv_n_heads
-        self.head_dim = d_model // n_heads
+        self.head_dim = hidden_size // n_heads
         if self.kv_n_heads <= 0:
             raise ValueError('kv_n_heads should be greater than zero.')
         if self.kv_n_heads > self.n_heads:
@@ -205,17 +205,17 @@ class GroupedQueryAttention(nn.Module):
             raise ValueError('Each Q head should get the same number of KV heads, so n_heads must be divisible by kv_n_heads.')
         self.softmax_scale = softmax_scale
         if self.softmax_scale is None:
-            self.softmax_scale = 1 / math.sqrt(self.d_model / self.n_heads)
+            self.softmax_scale = 1 / math.sqrt(self.hidden_size / self.n_heads)
         self.attn_dropout_p = attn_pdrop
         fc_kwargs = {}
         if fc_type != 'te':
             fc_kwargs['device'] = device
-        self.Wqkv = FC_CLASS_REGISTRY[fc_type](self.d_model, self.d_model + 2 * self.kv_n_heads * self.head_dim, **fc_kwargs)
+        self.Wqkv = FC_CLASS_REGISTRY[fc_type](self.hidden_size, self.hidden_size + 2 * self.kv_n_heads * self.head_dim, **fc_kwargs)
         fuse_splits = [i * self.head_dim for i in range(1, self.n_heads + 2 * self.kv_n_heads)]
         self.Wqkv._fused = (0, fuse_splits)
         if self.qk_ln:
             norm_class = NORM_CLASS_REGISTRY[norm_type.lower()]
-            self.q_ln = norm_class(self.d_model, device=device)
+            self.q_ln = norm_class(self.hidden_size, device=device)
             self.k_ln = norm_class(self.kv_n_heads * self.head_dim, device=device)
         if self.attn_impl == 'flash':
             self.attn_fn = flash_attn_fn
@@ -225,14 +225,14 @@ class GroupedQueryAttention(nn.Module):
             self.attn_fn = scaled_multihead_dot_product_attention
         else:
             raise ValueError(f'attn_impl={attn_impl!r} is an invalid setting.')
-        self.out_proj = FC_CLASS_REGISTRY[fc_type](self.d_model, self.d_model, **fc_kwargs)
+        self.out_proj = FC_CLASS_REGISTRY[fc_type](self.hidden_size, self.hidden_size, **fc_kwargs)
         self.out_proj._is_residual = True
 
     def forward(self, x: torch.Tensor, past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]]=None, attn_bias: Optional[torch.Tensor]=None, attention_mask: Optional[torch.Tensor]=None, is_causal: bool=True, needs_weights: bool=False) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor, torch.Tensor]]]:
         qkv = self.Wqkv(x)
         if self.clip_qkv:
             qkv = qkv.clamp(min=-self.clip_qkv, max=self.clip_qkv)
-        (query, key, value) = qkv.split([self.d_model, self.kv_n_heads * self.head_dim, self.kv_n_heads * self.head_dim], dim=2)
+        (query, key, value) = qkv.split([self.hidden_size, self.kv_n_heads * self.head_dim, self.kv_n_heads * self.head_dim], dim=2)
         key_padding_mask = attention_mask
         if self.qk_ln:
             dtype = query.dtype
@@ -248,8 +248,8 @@ class MultiheadAttention(GroupedQueryAttention):
     additive bias.
     """
 
-    def __init__(self, d_model: int, n_heads: int, attn_impl: str='triton', clip_qkv: Optional[float]=None, qk_ln: bool=False, softmax_scale: Optional[float]=None, attn_pdrop: float=0.0, norm_type: str='low_precision_layernorm', fc_type: str='torch', device: Optional[str]=None):
-        super().__init__(d_model=d_model, n_heads=n_heads, kv_n_heads=n_heads, attn_impl=attn_impl, clip_qkv=clip_qkv, qk_ln=qk_ln, softmax_scale=softmax_scale, attn_pdrop=attn_pdrop, norm_type=norm_type, fc_type=fc_type, device=device)
+    def __init__(self, hidden_size: int, n_heads: int, attn_impl: str='triton', clip_qkv: Optional[float]=None, qk_ln: bool=False, softmax_scale: Optional[float]=None, attn_pdrop: float=0.0, norm_type: str='low_precision_layernorm', fc_type: str='torch', device: Optional[str]=None):
+        super().__init__(hidden_size=hidden_size, n_heads=n_heads, kv_n_heads=n_heads, attn_impl=attn_impl, clip_qkv=clip_qkv, qk_ln=qk_ln, softmax_scale=softmax_scale, attn_pdrop=attn_pdrop, norm_type=norm_type, fc_type=fc_type, device=device)
 
 class MultiQueryAttention(GroupedQueryAttention):
     """Multi-Query self attention.
@@ -258,8 +258,8 @@ class MultiQueryAttention(GroupedQueryAttention):
     additive bias.
     """
 
-    def __init__(self, d_model: int, n_heads: int, attn_impl: str='triton', clip_qkv: Optional[float]=None, qk_ln: bool=False, softmax_scale: Optional[float]=None, attn_pdrop: float=0.0, norm_type: str='low_precision_layernorm', fc_type: str='torch', device: Optional[str]=None):
-        super().__init__(d_model=d_model, n_heads=n_heads, kv_n_heads=1, attn_impl=attn_impl, clip_qkv=clip_qkv, qk_ln=qk_ln, softmax_scale=softmax_scale, attn_pdrop=attn_pdrop, norm_type=norm_type, fc_type=fc_type, device=device)
+    def __init__(self, hidden_size: int, n_heads: int, attn_impl: str='triton', clip_qkv: Optional[float]=None, qk_ln: bool=False, softmax_scale: Optional[float]=None, attn_pdrop: float=0.0, norm_type: str='low_precision_layernorm', fc_type: str='torch', device: Optional[str]=None):
+        super().__init__(hidden_size=hidden_size, n_heads=n_heads, kv_n_heads=1, attn_impl=attn_impl, clip_qkv=clip_qkv, qk_ln=qk_ln, softmax_scale=softmax_scale, attn_pdrop=attn_pdrop, norm_type=norm_type, fc_type=fc_type, device=device)
 
 def attn_bias_shape(attn_impl: str, n_heads: int, seq_len: int, alibi: bool, prefix_lm: bool, causal: bool, use_sequence_id: bool) -> Optional[Tuple[int, int, int, int]]:
     if attn_impl == 'flash':

blocks.py

@@ -8,7 +8,7 @@ from .norm import NORM_CLASS_REGISTRY
 
 class MPTBlock(nn.Module):
 
-    def __init__(self, d_model: int, n_heads: int, expansion_ratio: int, attn_config: Optional[Dict]=None, ffn_config: Optional[Dict]=None, resid_pdrop: float=0.0, norm_type: str='low_precision_layernorm', fc_type: str='torch', device: Optional[str]=None, **kwargs: Any):
+    def __init__(self, hidden_size: int, n_heads: int, expansion_ratio: int, attn_config: Optional[Dict]=None, ffn_config: Optional[Dict]=None, resid_pdrop: float=0.0, norm_type: str='low_precision_layernorm', fc_type: str='torch', device: Optional[str]=None, **kwargs: Any):
         if attn_config is None:
             attn_config = {'attn_type': 'multihead_attention', 'attn_pdrop': 0.0, 'attn_impl': 'triton', 'qk_ln': False, 'clip_qkv': None, 'softmax_scale': None, 'prefix_lm': False, 'attn_uses_sequence_id': False, 'alibi': False, 'alibi_bias_max': 8}
         if ffn_config is None:
@@ -20,12 +20,12 @@ class MPTBlock(nn.Module):
         attn_class = ATTN_CLASS_REGISTRY[attn_config['attn_type']]
         args_to_exclude_in_attn_class = {'attn_type', 'prefix_lm', 'alibi', 'attn_uses_sequence_id', 'alibi_bias_max'}
         attn_config_subset_for_attn_class = {k: v for (k, v) in attn_config.items() if k not in args_to_exclude_in_attn_class}
-        self.norm_1 = norm_class(d_model, device=device)
-        self.attn = attn_class(d_model=d_model, n_heads=n_heads, fc_type=fc_type, device=device, **attn_config_subset_for_attn_class)
+        self.norm_1 = norm_class(hidden_size, device=device)
+        self.attn = attn_class(hidden_size=hidden_size, n_heads=n_heads, fc_type=fc_type, device=device, **attn_config_subset_for_attn_class)
         self.norm_2 = None
         if not getattr(FFN_CLASS_REGISTRY[ffn_config['ffn_type']], '_has_norm', False):
-            self.norm_2 = norm_class(d_model, device=device)
-        self.ffn = build_ffn(d_model=d_model, expansion_ratio=expansion_ratio, device=device, **ffn_config)
+            self.norm_2 = norm_class(hidden_size, device=device)
+        self.ffn = build_ffn(hidden_size=hidden_size, expansion_ratio=expansion_ratio, device=device, **ffn_config)
         self.resid_attn_dropout = nn.Dropout(resid_pdrop)
         self.resid_ffn_dropout = nn.Dropout(resid_pdrop)
 

config.json

@@ -19,7 +19,7 @@
     "AutoConfig": "configuration_mpt.MPTConfig",
     "AutoModelForCausalLM": "modeling_mpt.MPTForCausalLM"
   },
-  "d_model": 3072,
+  "hidden_size": 3072,
   "emb_pdrop": 0.0,
   "embedding_fraction": 1.0,
   "expansion_ratio": 4,

configuration_mpt.py

@@ -9,11 +9,11 @@ init_config_defaults: Dict = {'name': 'kaiming_normal_', 'fan_mode': 'fan_in', '
 class MPTConfig(PretrainedConfig):
     model_type = 'mpt'
 
-    def __init__(self, d_model: int=2048, n_heads: int=16, n_layers: int=24, expansion_ratio: int=4, max_seq_len: int=2048, vocab_size: int=50368, resid_pdrop: float=0.0, emb_pdrop: float=0.0, learned_pos_emb: bool=True, attn_config: Dict=attn_config_defaults, ffn_config: Dict=ffn_config_defaults, init_device: str='cpu', logit_scale: Optional[Union[float, str]]=None, no_bias: bool=False, embedding_fraction: float=1.0, norm_type: str='low_precision_layernorm', use_cache: bool=False, init_config: Dict=init_config_defaults, fc_type: str='torch', verbose: Optional[int]=None, **kwargs: Any):
+    def __init__(self, hidden_size: int=2048, n_heads: int=16, n_layers: int=24, expansion_ratio: int=4, max_seq_len: int=2048, vocab_size: int=50368, resid_pdrop: float=0.0, emb_pdrop: float=0.0, learned_pos_emb: bool=True, attn_config: Dict=attn_config_defaults, ffn_config: Dict=ffn_config_defaults, init_device: str='cpu', logit_scale: Optional[Union[float, str]]=None, no_bias: bool=False, embedding_fraction: float=1.0, norm_type: str='low_precision_layernorm', use_cache: bool=False, init_config: Dict=init_config_defaults, fc_type: str='torch', verbose: Optional[int]=None, **kwargs: Any):
         """The MPT configuration class.
 
         Args:
-            d_model (int): The size of the embedding dimension of the model.
+            hidden_size (int): The size of the embedding dimension of the model.
             n_heads (int): The number of attention heads.
             n_layers (int): The number of layers in the model.
             expansion_ratio (int): The ratio of the up/down scale in the ffn.
@@ -67,7 +67,7 @@ class MPTConfig(PretrainedConfig):
                 See llmfoundry.models.utils.param_init_fns.py for info on other param init config options
             fc_type (str): choose fc layer implementation. Options: torch and te. te layers support fp8 when using H100 GPUs.
         """
-        self.d_model = d_model
+        self.hidden_size = hidden_size
         self.n_heads = n_heads
         self.n_layers = n_layers
         self.expansion_ratio = expansion_ratio
@@ -108,8 +108,8 @@ class MPTConfig(PretrainedConfig):
         self.attn_config = self._set_config_defaults(self.attn_config, attn_config_defaults)
         self.ffn_config = self._set_config_defaults(self.ffn_config, ffn_config_defaults)
         self.init_config = self._set_config_defaults(self.init_config, init_config_defaults)
-        if self.d_model % self.n_heads != 0:
-            raise ValueError('d_model must be divisible by n_heads')
+        if self.hidden_size % self.n_heads != 0:
+            raise ValueError('hidden_size must be divisible by n_heads')
         if any((prob < 0 or prob > 1 for prob in [self.attn_config['attn_pdrop'], self.resid_pdrop, self.emb_pdrop])):
             raise ValueError("self.attn_config['attn_pdrop'], resid_pdrop, emb_pdrop are probabilities and must be between 0 and 1")
         if self.attn_config['attn_impl'] not in ['torch', 'flash', 'triton']:
@@ -122,8 +122,8 @@ class MPTConfig(PretrainedConfig):
             raise NotImplementedError('attn_uses_sequence_id only implemented with torch and triton attention.')
         if self.embedding_fraction > 1 or self.embedding_fraction <= 0:
             raise ValueError('model.embedding_fraction must be between 0 (exclusive) and 1 (inclusive)!')
-        if isinstance(self.logit_scale, str) and self.logit_scale != 'inv_sqrt_d_model':
-            raise ValueError(f"self.logit_scale={self.logit_scale!r} is not recognized as an option; use numeric value or 'inv_sqrt_d_model'.")
+        if isinstance(self.logit_scale, str) and self.logit_scale != 'inv_sqrt_hidden_size':
+            raise ValueError(f"self.logit_scale={self.logit_scale!r} is not recognized as an option; use numeric value or 'inv_sqrt_hidden_size'.")
         if self.init_config.get('name', None) is None:
             raise ValueError(f"self.init_config={self.init_config!r} 'name' needs to be set.")
         if not self.learned_pos_emb and (not self.attn_config['alibi']):

ffn.py

@@ -10,14 +10,14 @@ except:
 
 class MPTMLP(nn.Module):
 
-    def __init__(self, d_model: int, expansion_ratio: int, fc_type: str='torch', device: Optional[str]=None):
+    def __init__(self, hidden_size: int, expansion_ratio: int, fc_type: str='torch', device: Optional[str]=None):
         super().__init__()
         fc_kwargs = {}
         if fc_type != 'te':
             fc_kwargs['device'] = device
-        self.up_proj = FC_CLASS_REGISTRY[fc_type](d_model, expansion_ratio * d_model, **fc_kwargs)
+        self.up_proj = FC_CLASS_REGISTRY[fc_type](hidden_size, expansion_ratio * hidden_size, **fc_kwargs)
         self.act = nn.GELU(approximate='none')
-        self.down_proj = FC_CLASS_REGISTRY[fc_type](expansion_ratio * d_model, d_model, **fc_kwargs)
+        self.down_proj = FC_CLASS_REGISTRY[fc_type](expansion_ratio * hidden_size, hidden_size, **fc_kwargs)
         self.down_proj._is_residual = True
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -27,13 +27,13 @@ if te is not None:
     te.LayerNormMLP._has_norm = True
     FFN_CLASS_REGISTRY['te_ln_mlp'] = te.LayerNormMLP
 
-def build_ffn(d_model: int, expansion_ratio: int, fc_type: str='torch', device: Optional[str]=None, **kwargs: Any) -> nn.Module:
+def build_ffn(hidden_size: int, expansion_ratio: int, fc_type: str='torch', device: Optional[str]=None, **kwargs: Any) -> nn.Module:
     ffn_type = kwargs.pop('ffn_type')
     if ffn_type == 'mptmlp':
         if len(kwargs) > 0:
             raise ValueError(f'MPTMLP got an unexpected keyword argument: {kwargs}')
-        return MPTMLP(d_model=d_model, expansion_ratio=expansion_ratio, fc_type=fc_type, device=device)
+        return MPTMLP(hidden_size=hidden_size, expansion_ratio=expansion_ratio, fc_type=fc_type, device=device)
     elif ffn_type == 'te_ln_mlp':
         assert te is not None
-        return te.LayerNormMLP(hidden_size=d_model, ffn_hidden_size=d_model * expansion_ratio, **kwargs)
+        return te.LayerNormMLP(hidden_size=hidden_size, ffn_hidden_size=hidden_size * expansion_ratio, **kwargs)
     raise ValueError(f'ffn_type={ffn_type!r} not recognized.')

modeling_mpt.py

@@ -56,12 +56,12 @@ class MPTModel(MPTPreTrainedModel):
             raise NotImplementedError(f'Requested norm type ({config.norm_type}) is not implemented within this repo (Options: {norm_options}).')
         norm_class = NORM_CLASS_REGISTRY[config.norm_type.lower()]
         self.embedding_fraction = config.embedding_fraction
-        self.wte = SharedEmbedding(config.vocab_size, config.d_model, device=config.init_device)
+        self.wte = SharedEmbedding(config.vocab_size, config.hidden_size, device=config.init_device)
         if self.learned_pos_emb:
-            self.wpe = torch.nn.Embedding(config.max_seq_len, config.d_model, device=config.init_device)
+            self.wpe = torch.nn.Embedding(config.max_seq_len, config.hidden_size, device=config.init_device)
         self.emb_drop = nn.Dropout(config.emb_pdrop)
         self.blocks = nn.ModuleList([MPTBlock(device=config.init_device, **config.to_dict()) for _ in range(config.n_layers)])
-        self.norm_f = norm_class(config.d_model, device=config.init_device)
+        self.norm_f = norm_class(config.hidden_size, device=config.init_device)
         if config.init_device != 'meta':
             log.info(f'We recommend using config.init_device="meta" with Composer + FSDP for faster initialization.')
             self.apply(self.param_init_fn)
@@ -213,7 +213,7 @@ class MPTModel(MPTPreTrainedModel):
 
     def param_init_fn(self, module: nn.Module) -> None:
         init_fn_name = self.config.init_config['name']
-        MODEL_INIT_REGISTRY[init_fn_name](module=module, n_layers=self.config.n_layers, d_model=self.config.d_model, **self.config.init_config)
+        MODEL_INIT_REGISTRY[init_fn_name](module=module, n_layers=self.config.n_layers, hidden_size=self.config.hidden_size, **self.config.init_config)
 
     def fsdp_wrap_fn(self, module: nn.Module) -> bool:
         return isinstance(module, MPTBlock)
@@ -238,10 +238,10 @@ class MPTForCausalLM(MPTPreTrainedModel):
         if config.logit_scale is not None:
             logit_scale = config.logit_scale
             if isinstance(logit_scale, str):
-                if logit_scale == 'inv_sqrt_d_model':
-                    logit_scale = 1 / math.sqrt(config.d_model)
+                if logit_scale == 'inv_sqrt_hidden_size':
+                    logit_scale = 1 / math.sqrt(config.hidden_size)
                 else:
-                    raise ValueError(f"logit_scale={logit_scale!r} is not recognized as an option; use numeric value or 'inv_sqrt_d_model'.")
+                    raise ValueError(f"logit_scale={logit_scale!r} is not recognized as an option; use numeric value or 'inv_sqrt_hidden_size'.")
             self.logit_scale = logit_scale
 
     def get_input_embeddings(self) -> nn.Embedding:
@@ -282,7 +282,7 @@ class MPTForCausalLM(MPTPreTrainedModel):
 
     def param_init_fn(self, module: nn.Module) -> None:
         init_fn_name = self.config.init_config['name']
-        MODEL_INIT_REGISTRY[init_fn_name](module=module, n_layers=self.config.n_layers, d_model=self.config.d_model, **self.config.init_config)
+        MODEL_INIT_REGISTRY[init_fn_name](module=module, n_layers=self.config.n_layers, hidden_size=self.config.hidden_size, **self.config.init_config)
 
     def fsdp_wrap_fn(self, module: nn.Module) -> bool:
         return isinstance(module, MPTBlock)

param_init_fns.py

@@ -29,7 +29,7 @@ def fused_init_helper_(module: nn.Module, init_fn_: Callable) -> None:
         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:
+def generic_param_init_fn_(module: nn.Module, init_fn_: Callable, n_layers: int, hidden_size: 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:
@@ -85,8 +85,8 @@ def generic_param_init_fn_(module: nn.Module, init_fn_: Callable, n_layers: int,
         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
+            assert hidden_size is not None
+            _d = hidden_size
             splits = (0, _d, 2 * _d, 3 * _d)
             for (s, e) in zip(splits[:-1], splits[1:]):
                 init_fn_(module.in_proj_weight[s:e])
@@ -130,23 +130,23 @@ def generic_param_init_fn_(module: nn.Module, init_fn_: Callable, n_layers: int,
 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:
+def _normal_param_init_fn_(module: nn.Module, std: float, n_layers: int, hidden_size: 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)
+    generic_param_init_fn_(module=module, init_fn_=init_fn_, hidden_size=hidden_size, 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:
+def baseline_param_init_fn_(module: nn.Module, init_std: Optional[float], n_layers: int, hidden_size: 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)
+    _normal_param_init_fn_(module=module, std=init_std, hidden_size=hidden_size, 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:
+def small_param_init_fn_(module: nn.Module, n_layers: int, hidden_size: 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)
+    std = math.sqrt(2 / (5 * hidden_size))
+    _normal_param_init_fn_(module=module, std=std, hidden_size=hidden_size, 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:
+def neox_param_init_fn_(module: nn.Module, n_layers: int, hidden_size: 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)
@@ -155,25 +155,25 @@ def neox_param_init_fn_(module: nn.Module, n_layers: int, d_model: int, emb_init
     """
     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)
+    small_param_init_fn_(module=module, hidden_size=hidden_size, 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:
+def kaiming_uniform_param_init_fn_(module: nn.Module, n_layers: int, hidden_size: 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)
+    generic_param_init_fn_(module=module, init_fn_=kaiming_uniform_, hidden_size=hidden_size, 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:
+def kaiming_normal_param_init_fn_(module: nn.Module, n_layers: int, hidden_size: 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)
+    generic_param_init_fn_(module=module, init_fn_=kaiming_normal_, hidden_size=hidden_size, 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:
+def xavier_uniform_param_init_fn_(module: nn.Module, n_layers: int, hidden_size: 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)
+    generic_param_init_fn_(module=module, init_fn_=xavier_uniform_, hidden_size=hidden_size, 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:
+def xavier_normal_param_init_fn_(module: nn.Module, n_layers: int, hidden_size: 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)
+    generic_param_init_fn_(module=module, init_fn_=xavier_normal_, hidden_size=hidden_size, 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_}