Add LLaMA weight initialization

Author: cbalioglu

src/fairseq2/models/jepa/_factory.py

@@ -6,7 +6,6 @@
 
 from __future__ import annotations
 
-import math
 from typing import cast
 
 import torch
@@ -183,12 +182,12 @@ def create_encoder(self) -> TransformerEncoder:
             layer_norm_factory=self.create_layer_norm,
         )
 
-    def create_encoder_layer(self, idx: int) -> TransformerEncoderLayer:
+    def create_encoder_layer(self, layer_idx: int) -> TransformerEncoderLayer:
         config = self._config
 
-        self_attn = self.create_attention(idx)
+        self_attn = self.create_attention(layer_idx)
 
-        ffn = self.create_ffn(idx)
+        ffn = self.create_ffn(layer_idx)
 
         drop_path = DropPathResidualConnect(drop_p=config.droppath_p)
 
@@ -226,7 +225,7 @@ def init_projection(proj: Linear) -> None:
             _init_truncated_normal(proj.weight, proj.bias, std=init_std)
 
             with torch.no_grad():
-                proj.weight.div_(math.sqrt(2.0 * (layer_idx + 1)))
+                proj.weight.div_((2.0 * (layer_idx + 1)) ** 0.5)
 
         return Linear(
             config.model_dim, config.model_dim, bias=True, init_fn=init_projection
@@ -241,7 +240,7 @@ def init_projection(proj: Linear) -> None:
             _init_truncated_normal(proj.weight, proj.bias, std=init_std)
 
             with torch.no_grad():
-                proj.weight.div_(math.sqrt(2.0 * (layer_idx + 1)))
+                proj.weight.div_((2.0 * (layer_idx + 1)) ** 0.5)
 
         inner_dim = int(config.model_dim * config.ffn_inner_dim_ratio)
 
@@ -282,7 +281,7 @@ def init_layer_norm(m: LayerNorm) -> None:
 def _init_truncated_normal(
     weight: Tensor, bias: Tensor | None, *, std: float = 1.0
 ) -> None:
-    nn.init.trunc_normal_(weight, std=std)
+    nn.init.trunc_normal_(weight, mean=0.0, std=std)
 
     if bias is not None:
         nn.init.zeros_(bias)

src/fairseq2/models/llama/_config.py

@@ -7,7 +7,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
-from typing import Final
+from typing import Final, Literal
 
 from fairseq2.context import RuntimeContext
 from fairseq2.data import VocabularyInfo
@@ -81,6 +81,20 @@ class LLaMAConfig:
     encoder, aiming to increase the context length.
     """
 
+    init_std: float | None = None
+    """
+    If not ``None``, the standard deviation to initialize input embeddings and
+    projection weights; otherwise, ``model_dim ** -0.5`` will be used instead.
+    """
+
+    init_std_scale: Literal["none", "layer", "stack"] = "layer"
+    """
+    The method to use to scale ``init_std`` per layer. If 'none', no scaling
+    will be applied. If 'layer', ``init_std`` will be scaled by the depth of
+    the layer. If 'stack', ``init_std`` will be scaled by the total depth of
+    the decoder.
+    """
+
     dropout_p: float = 0.1
     """The dropout probability on outputs of Transformer layers."""
 

src/fairseq2/models/llama/_factory.py

@@ -10,13 +10,13 @@
 from functools import partial
 
 import torch
+import torch.nn as nn
 from torch import Tensor
 
 from fairseq2.models.llama._config import LLaMAConfig, LLaMARopeScalingConfig
 from fairseq2.models.transformer import (
     TransformerEmbeddingFrontend,
     TransformerFrontend,
-    init_final_projection,
 )
 from fairseq2.models.transformer_decoder import TransformerDecoderModel
 from fairseq2.nn import (
@@ -77,8 +77,19 @@ def create_model(self) -> TransformerDecoderModel:
     def create_embedding(self) -> Embedding:
         config = self._config
 
+        init_std = config.init_std
+
+        def init_embed(embed: StandardEmbedding) -> None:
+            embed_dim = embed.weight.shape[1]
+
+            std = init_std or (embed_dim**-0.5)
+
+            _init_truncated_normal(embed.weight, bias=None, std=std)
+
         return StandardEmbedding(
-            num_embeddings=config.vocab_info.size, embedding_dim=config.model_dim
+            num_embeddings=config.vocab_info.size,
+            embedding_dim=config.model_dim,
+            init_fn=init_embed,
         )
 
     def create_decoder_frontend(self, embed: Embedding) -> TransformerFrontend:
@@ -95,8 +106,8 @@ def create_decoder(self) -> TransformerDecoder:
 
         layers = []
 
-        for _ in range(config.num_layers):
-            layer = self.create_decoder_layer(pos_encoder)
+        for idx in range(config.num_layers):
+            layer = self.create_decoder_layer(idx, pos_encoder)
 
             layers.append(layer)
 
@@ -125,11 +136,11 @@ def create_position_encoder(self) -> PositionEncoder:
         )
 
     def create_decoder_layer(
-        self, pos_encoder: PositionEncoder
+        self, layer_idx: int, pos_encoder: PositionEncoder
     ) -> TransformerDecoderLayer:
-        self_attn = self.create_attention(pos_encoder)
+        self_attn = self.create_attention(layer_idx, pos_encoder)
 
-        ffn = self.create_ffn()
+        ffn = self.create_ffn(layer_idx)
 
         return StandardTransformerDecoderLayer(
             self_attn,
@@ -139,23 +150,49 @@ def create_decoder_layer(
             layer_norm_factory=self.create_layer_norm,
         )
 
-    def create_attention(self, pos_encoder: PositionEncoder) -> MultiheadAttention:
+    def create_attention(
+        self, layer_idx: int, pos_encoder: PositionEncoder
+    ) -> MultiheadAttention:
         config = self._config
 
+        init_std = config.init_std
+
+        std_scale_factor = self._get_std_scale_factor(layer_idx)
+
+        def init_projection(proj: Linear) -> None:
+            input_dim = proj.weight.shape[1]
+
+            std = init_std or (input_dim**-0.5)
+
+            _init_truncated_normal(proj.weight, proj.bias, std=std / std_scale_factor)
+
         sdpa = create_default_sdpa(attn_dropout_p=config.dropout_p)
 
         return StandardMultiheadAttention(
             config.model_dim,
             config.num_attn_heads,
             num_key_value_heads=config.num_key_value_heads,
+            qkv_proj_init_fn=init_projection,
             sdpa=sdpa,
             pos_encoder=pos_encoder,
+            output_proj_init_fn=init_projection,
             bias=False,
         )
 
-    def create_ffn(self) -> FeedForwardNetwork:
+    def create_ffn(self, layer_idx: int) -> FeedForwardNetwork:
         config = self._config
 
+        init_std = config.init_std
+
+        std_scale_factor = self._get_std_scale_factor(layer_idx)
+
+        def init_projection(proj: Linear) -> None:
+            input_dim = proj.weight.shape[1]
+
+            std = init_std or (input_dim**-0.5)
+
+            _init_truncated_normal(proj.weight, proj.bias, std=std / std_scale_factor)
+
         ffn_inner_dim = int(config.ffn_inner_dim * config.ffn_inner_dim_multiplier)
 
         return GLUFeedForwardNetwork(
@@ -165,8 +202,26 @@ def create_ffn(self) -> FeedForwardNetwork:
             inner_dim_scale=config.ffn_inner_dim_scale,
             inner_dim_to_multiple=config.ffn_inner_dim_to_multiple,
             inner_dropout_p=config.dropout_p,
+            proj_init_fn=init_projection,
         )
 
+    def _get_std_scale_factor(self, layer_idx: int) -> float:
+        config = self._config
+
+        match config.init_std_scale:
+            case "layer":
+                n = layer_idx
+            case "stack":
+                n = config.num_layers
+            case "none":
+                return 1.0
+            case _:
+                raise ValueError(
+                    f"`config.init_std_scale` must be 'none', 'layer', or 'stack', but is '{config.init_std_scale}' instead."
+                )
+
+        return (2 * (n + 1)) ** 0.5  # type: ignore[no-any-return]
+
     def create_final_proj(self, embed: Embedding) -> Projection:
         config = self._config
 
@@ -178,11 +233,20 @@ def create_final_proj(self, embed: Embedding) -> Projection:
 
             return TiedProjection(embed.weight, bias=None)
 
+        init_std = config.init_std
+
+        def init_projection(proj: Linear) -> None:
+            input_dim = proj.weight.shape[1]
+
+            std = init_std or (input_dim**-0.5)
+
+            _init_truncated_normal(proj.weight, proj.bias, std=std)
+
         return Linear(
             config.model_dim,
             config.vocab_info.size,
             bias=False,
-            init_fn=init_final_projection,
+            init_fn=init_projection,
         )
 
     @staticmethod
@@ -192,6 +256,15 @@ def create_layer_norm(
         return RMSNorm(model_dim, bias=False, device=device, dtype=dtype)
 
 
+def _init_truncated_normal(
+    weight: Tensor, bias: Tensor | None, *, std: float = 1.0
+) -> None:
+    nn.init.trunc_normal_(weight, mean=0.0, std=std, a=-3 * std, b=3 * std)
+
+    if bias is not None:
+        nn.init.zeros_(bias)
+
+
 def init_llama_scaled_freqs(
     pos_encoder: RotaryEncoder, rope_scaling: LLaMARopeScalingConfig
 ) -> Tensor:

src/fairseq2/models/transformer/__init__.py

@@ -22,16 +22,16 @@
 from fairseq2.models.transformer._factory import (
     create_transformer_model as create_transformer_model,
 )
+from fairseq2.models.transformer._factory import (
+    init_final_projection as init_final_projection,
+)
 from fairseq2.models.transformer._frontend import (
     TransformerEmbeddingFrontend as TransformerEmbeddingFrontend,
 )
 from fairseq2.models.transformer._frontend import (
     TransformerFrontend as TransformerFrontend,
 )
 from fairseq2.models.transformer._model import TransformerModel as TransformerModel
-from fairseq2.models.transformer._model import (
-    init_final_projection as init_final_projection,
-)
 
 # isort: split
 

src/fairseq2/models/transformer/_factory.py

@@ -6,6 +6,8 @@
 
 from __future__ import annotations
 
+import torch.nn as nn
+
 from fairseq2.models.transformer._config import TransformerConfig
 from fairseq2.models.transformer._frontend import (
     TransformerEmbeddingFrontend,
@@ -174,3 +176,10 @@ def create_final_proj(self, embed: Embedding) -> Projection:
             return TiedProjection(embed.weight, bias=None)
 
         return Linear(config.model_dim, config.vocab_info.size, bias=False)
+
+
+def init_final_projection(proj: Linear) -> None:
+    nn.init.normal_(proj.weight, std=proj.input_dim**-0.5)
+
+    if proj.bias is not None:
+        nn.init.zeros_(proj.bias)

src/fairseq2/models/transformer/_model.py

@@ -8,15 +8,14 @@
 
 from typing import final
 
-import torch.nn as nn
 from torch import Tensor
 from typing_extensions import override
 
 from fairseq2.data import VocabularyInfo
 from fairseq2.models.encoder_decoder import EncoderDecoderModel
 from fairseq2.models.sequence import SequenceModelOutput
 from fairseq2.models.transformer._frontend import TransformerFrontend
-from fairseq2.nn import IncrementalStateBag, Linear, Projection
+from fairseq2.nn import IncrementalStateBag, Projection
 from fairseq2.nn.padding import PaddingMask
 from fairseq2.nn.transformer import TransformerDecoder, TransformerEncoder
 
@@ -106,10 +105,3 @@ def project(
         logits = self.final_proj(decoder_output)
 
         return SequenceModelOutput(logits, self.target_vocab_info.pad_idx)
-
-
-def init_final_projection(proj: Linear) -> None:
-    nn.init.normal_(proj.weight, std=proj.input_dim**-0.5)
-
-    if proj.bias is not None:
-        nn.init.zeros_(proj.bias)

src/fairseq2/models/wav2vec2/asr/_model.py

@@ -74,7 +74,7 @@ def __init__(
             self.model_dim,
             target_vocab_info.size,
             bias=True,
-            init_fn=init_final_projection,
+            init_fn=_init_final_projection,
             device=device,
             dtype=dtype,
         )
@@ -105,7 +105,7 @@ def forward(self, batch: SequenceBatch) -> AsrModelOutput:
         return AsrModelOutput(logits, padding_mask)
 
 
-def init_final_projection(proj: Linear) -> None:
+def _init_final_projection(proj: Linear) -> None:
     """Initialize ``proj`` as the final projection of a wav2vec 2.0 ASR model."""
     nn.init.xavier_uniform_(proj.weight)