Draft: Feat/initialization component

config_files/training/config_lorem_ipsum.yaml

@@ -178,7 +178,25 @@ wrapped_model:
     sharding_strategy: FULL_SHARD
     block_names: [GPT2Block]
 
-model:
+model: 
+  component_key: model
+  variant_key: weight_initialized
+  config:
+    model:
+      instance_key: model_raw
+      pass_type: BY_REFERENCE
+    weight_initializer:
+      component_key: weight_initialization
+      variant_key: composed
+      config:
+        model_type: gpt2
+        weight_init_type: scaled_embed
+        mean: 0.0
+        plain_std: 0.02
+        hidden_dim: ${model_raw.config.n_embd}
+        num_layers: ${model_raw.config.n_layer}
+
+model_raw:
   component_key: model
   variant_key: gpt2
   config:
@@ -198,34 +216,30 @@ model:
       qkv_transforms:
         - type_hint: RotaryTransform
           config:
-            n_embd: ${model.config.n_embd}
-            n_head: ${model.config.n_head_q} #it has to be head_q here
+            n_embd: ${model_raw.config.n_embd}
+            n_head: ${model_raw.config.n_head_q} #it has to be head_q here
             seq_length_dim: -2
     attention_implementation: manual
     activation_type: gelu
-    weight_init:
-      mean: 0.0
-      std: 0.02
-      type: scaled
     attention_norm:
       component_key: layer_norm
       variant_key: rms_norm
       config:
-        ndim: ${model.config.n_embd}
+        ndim: ${model_raw.config.n_embd}
         bias: true
         epsilon: 1e-5
     ffn_norm:
       component_key: layer_norm
       variant_key: rms_norm
       config:
-        ndim: ${model.config.n_embd}
+        ndim: ${model_raw.config.n_embd}
         bias: true
         epsilon: 1e-5
     lm_head_norm:
       component_key: layer_norm
       variant_key: rms_norm
       config:
-        ndim: ${model.config.n_embd}
+        ndim: ${model_raw.config.n_embd}
         bias: true
         epsilon: 1e-5
 

src/modalities/config/config.py

@@ -22,6 +22,7 @@
     PydanticPytorchModuleType,
     PydanticSamplerIFType,
     PydanticTokenizerIFType,
+    PydanticWeightInitializationIFType,
 )
 from modalities.config.utils import parse_torch_device
 from modalities.running_env.env_utils import MixedPrecisionSettings, has_bfloat_support
@@ -228,6 +229,11 @@ def parse_sharding_strategy_by_name(cls, name):
         return parse_enum_by_name(name=name, enum_type=ShardingStrategy)
 
 
+class WeightInitializedModelConfig(BaseModel):
+    model: PydanticPytorchModuleType
+    weight_initializer: PydanticWeightInitializationIFType
+
+
 class PreTrainedHFTokenizerConfig(BaseModel):
     pretrained_model_name_or_path: str
     max_length: Annotated[int, Field(strict=True, ge=0)]

src/modalities/config/pydanctic_if_types.py

@@ -17,6 +17,7 @@
 from modalities.logging_broker.subscriber import MessageSubscriberIF
 from modalities.loss_functions import Loss
 from modalities.models.gpt2.collator import CollateFnIF
+from modalities.nn.weight_init.weight_init_if import WeightInitializationIF
 from modalities.tokenization.tokenizer_wrapper import TokenizerWrapper
 from modalities.training.gradient_clipping.gradient_clipper import GradientClipperIF
 
@@ -61,3 +62,4 @@ def __get_pydantic_core_schema__(
 PydanticPytorchDeviceType = Annotated[torch.device, PydanticThirdPartyTypeIF(torch.device)]
 PydanticTextInferenceComponentType = Annotated[TextInferenceComponent, PydanticThirdPartyTypeIF(TextInferenceComponent)]
 PydanticGradientClipperIFType = Annotated[GradientClipperIF, PydanticThirdPartyTypeIF(GradientClipperIF)]
+PydanticWeightInitializationIFType = Annotated[WeightInitializationIF, PydanticThirdPartyTypeIF(WeightInitializationIF)]

src/modalities/models/coca/coca_model.py

@@ -8,7 +8,7 @@
 from modalities.models.coca.attention_pooling import AttentionPooling
 from modalities.models.coca.multi_modal_decoder import MultiModalTextDecoder
 from modalities.models.coca.text_decoder import TextDecoder
-from modalities.models.model import ActivationType, NNModel, WeightInitializationConfig
+from modalities.models.model import ActivationType, NNModel
 from modalities.models.vision_transformer.vision_transformer_model import VisionTransformer, VisionTransformerConfig
 from modalities.nn.attention import AttentionConfig
 
@@ -42,7 +42,6 @@ class CoCaConfig(BaseModel):
     n_vision_queries: Annotated[int, Field(ge=1)]
     bias_attn_pool: bool
     epsilon_attn_pool: Annotated[float, Field(ge=0.0)]
-    weight_init: WeightInitializationConfig
 
 
 class CoCa(NNModel):
@@ -68,7 +67,6 @@ def __init__(
         epsilon_attn_pool: float,
         vision_encoder_config: VisionTransformerConfig,
         text_decoder_config: TextDecoderConfig,
-        weight_init: WeightInitializationConfig,
     ) -> None:
         super().__init__()
         self.prediction_key = prediction_key
@@ -123,14 +121,6 @@ def __init__(
             attention_config=text_decoder_config.attention_config,
         )
 
-        # init all weights
-        assert weight_init.type in ["plain", "scaled"], f"ERROR! weight_init.type = {weight_init.type} not implemented."
-        self.initialize_weights(
-            weight_init,
-            number_of_layers=text_decoder_config.n_layer_text + text_decoder_config.n_layer_multimodal_text,
-            hidden_dim=None,  # not well-defined as hidden_dim can be different for the text and multimodal decoder
-        )
-
     def forward(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
         vision_embd, vision_cls_token = self._forward_encode_vision(inputs)
         text_embd, text_cls_token = self._forward_encode_text(inputs)

src/modalities/models/gpt2/gpt2_model.py

@@ -15,7 +15,7 @@
 
 from modalities.config.pydanctic_if_types import PydanticPytorchModuleType
 from modalities.config.utils import convert_base_model_config_to_dict
-from modalities.models.model import ActivationType, NNModel, WeightInitializationConfig
+from modalities.models.model import ActivationType, NNModel
 from modalities.util import parse_enum_by_name
 
 # GPT2 implementation taken from nanogpt https://github.com/karpathy/nanoGPT
@@ -154,7 +154,6 @@ class GPT2LLMConfig(BaseModel):
     attention_norm: PydanticPytorchModuleType
     ffn_norm: PydanticPytorchModuleType
     lm_head_norm: PydanticPytorchModuleType
-    weight_init: WeightInitializationConfig
 
     @model_validator(mode="after")
     def check_divisibility(self) -> "GPT2LLMConfig":
@@ -419,7 +418,6 @@ def __init__(
         bias: bool,
         activation_type: ActivationType,
         attention_implementation: AttentionImplementation,
-        weight_init: WeightInitializationConfig,
         attention_config: AttentionConfig,
         attention_norm: nn.Module,
         ffn_norm: nn.Module,
@@ -489,9 +487,6 @@ def __init__(
         # not 100% sure what this is, so far seems to be harmless. TODO investigate
         self.transformer.wte.weight = self.lm_head.weight  # https://paperswithcode.com/method/weight-tying
 
-        # init all weights
-        self.initialize_weights(weight_init, number_of_layers=n_layer, hidden_dim=n_embd)
-
     def forward_impl(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
         input_ids = inputs[self.sample_key]
         device = input_ids.device

src/modalities/models/model.py

@@ -1,12 +1,9 @@
-import math
 from abc import abstractmethod
 from enum import Enum
-from functools import partial
-from typing import Annotated, Dict, List, Optional
+from typing import Dict, List, Optional
 
 import torch
 import torch.nn as nn
-from pydantic import BaseModel, Field
 
 from modalities.batch import DatasetBatch, InferenceResultBatch
 
@@ -18,12 +15,6 @@ class ActivationType(str, Enum):
     FUSED_SWIGLU = "fused_swiglu"
 
 
-class WeightInitializationConfig(BaseModel):
-    mean: Annotated[float, Field(strict=True, ge=0.0)]
-    std: Annotated[float, Field(strict=True, ge=0.0)] | str  # can be float or "auto"
-    type: str
-
-
 class NNModel(nn.Module):
     def __init__(self, seed: int = None, weight_decay_groups: Optional[WeightDecayGroups] = None):
         if seed is not None:
@@ -42,42 +33,6 @@ def forward(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
     def get_parameters(self) -> Dict[str, torch.Tensor]:
         return {name: param for name, param in self.named_parameters()}
 
-    def _init_weights(self, module: nn.Module, weight_init: WeightInitializationConfig):
-        if isinstance(module, nn.Linear):
-            torch.nn.init.normal_(module.weight, mean=weight_init.mean, std=weight_init.std)
-            if module.bias is not None:
-                torch.nn.init.zeros_(module.bias)
-        elif isinstance(module, nn.Embedding):
-            torch.nn.init.normal_(module.weight, mean=weight_init.mean, std=weight_init.std)
-
-    def initialize_weights(
-        self, weight_init: WeightInitializationConfig, number_of_layers: int, hidden_dim: Optional[int] = None
-    ):
-        # auto: choose std automatically
-        if weight_init.std == "auto":
-            assert hidden_dim is not None, "ERROR! weight_init.std = auto not implemented"
-            weight_init.std = math.sqrt(2 / (5 * hidden_dim))
-
-        # initialize weights
-        self.apply(partial(self._init_weights, weight_init=weight_init))
-
-        if weight_init.type == "plain":
-            pass  # nothing more to do
-        elif weight_init.type in ["scaled", "scaled_embed"]:
-            # apply special scaled init to the residual projections, per GPT-2 paper
-            for pn, p in self.named_parameters():
-                if pn.endswith("c_proj.weight"):
-                    torch.nn.init.normal_(
-                        p, mean=weight_init.mean, std=weight_init.std / math.sqrt(2 * number_of_layers)
-                    )
-            if weight_init.type == "scaled_embed":
-                # apply special init to embeddings, see https://arxiv.org/abs/2312.16903
-                for pn, p in self.named_parameters():
-                    if pn.endswith("wte.weight") or pn.endswith("wpe.weight"):
-                        torch.nn.init.normal_(p, mean=weight_init.mean, std=0.4)
-        else:
-            raise Exception(f"ERROR! weight_init.type = {weight_init.type} not implemented.")
-
 
 def model_predict_batch(model: nn.Module, batch: DatasetBatch) -> InferenceResultBatch:
     forward_result = model.forward(batch.samples)

src/modalities/models/model_factory.py

@@ -8,6 +8,7 @@
 from torch.distributed.fsdp import ShardingStrategy
 
 from modalities.checkpointing.checkpoint_loading import CheckpointLoadingIF
+from modalities.nn.weight_init.weight_init_if import WeightInitializationIF
 from modalities.running_env.env_utils import MixedPrecisionSettings
 from modalities.running_env.fsdp.fsdp_auto_wrapper import FSDPTransformerAutoWrapPolicyFactory
 from modalities.util import compute_number_of_trainable_parameters
@@ -55,3 +56,8 @@ def get_fsdp_wrapped_model(
         )
 
         return fsdp_model
+
+    @staticmethod
+    def get_weight_initalized_model(model: nn.Module, weight_initializer: WeightInitializationIF) -> nn.Module:
+        weight_initializer.initialize_in_place(model)
+        return model

src/modalities/nn/weight_init/__init__.py

@@ -0,0 +1 @@
+