Source code for kempnerforge.model.init

"""Weight initialization strategies for KempnerForge models."""

from __future__ import annotations

import math

import torch.nn as nn

from kempnerforge.config.schema import ModelConfig
from kempnerforge.model.moe import MoEMLP
from kempnerforge.model.mot import MoTAttention, MoTBlock
from kempnerforge.model.norm import RMSNorm
from kempnerforge.model.router import SigmoidTopKRouter

_RESETS_OWN_STATE = (RMSNorm, nn.LayerNorm, SigmoidTopKRouter)




def init_weights(model: nn.Module, config: ModelConfig) -> None:
    """Apply standard initialization to all parameters in a model.

    Strategy (following GPT-2/Llama conventions):

    - Linear layers: normal(0, 0.02)
    - Embedding layers: normal(0, 0.02)
    - Residual output projections (o_proj, down_proj): scaled by 1/sqrt(2 * n_layers)
    - Cross-attention block residuals: zero-initialized (identity-at-init warm-start)
    - MoT per-modality residual projections: zero-initialized (identity-at-construction)
    - Norm layers and the sigmoid router's own state: reset to their
      construction-time values, which a meta-device build does not preserve

    Only the modules in ``_RESETS_OWN_STATE`` are reset, and only with deterministic
    fills: a blanket ``reset_parameters()`` would also run ``nn.Linear``'s kaiming
    init, which draws from the RNG and changes every weight drawn after it for the
    same seed.
    """
    std = config.init_std
    residual_scale = 1.0 / math.sqrt(2.0 * config.n_layers)

    for module in model.modules():
        if isinstance(module, _RESETS_OWN_STATE) and not any(
            t.is_meta for t in (*module.parameters(recurse=False), *module.buffers(recurse=False))
        ):
            module.reset_parameters()

    for name, param in model.named_parameters():
        if param.is_meta:
            continue
        if param.dim() < 2:
            continue

        if name.startswith("cross_attention_layers.") and name.endswith(
            ("o_proj.weight", "down_proj.weight")
        ):
            nn.init.zeros_(param)
        elif name.endswith(("o_proj.weight", "down_proj.weight")):
            nn.init.normal_(param, mean=0.0, std=std * residual_scale)
        else:
            nn.init.normal_(param, mean=0.0, std=std)

    for module in model.modules():
        if isinstance(module, MoTAttention):
            for m in module.modalities:
                nn.init.zeros_(module.o_proj[m].weight)  # type: ignore[reportArgumentType]
        elif isinstance(module, MoTBlock):
            for m in module.modalities:
                mlp_m = module.mlp[m]
                if isinstance(mlp_m, MoEMLP):
                    continue
                if hasattr(mlp_m, "down_proj"):
                    nn.init.zeros_(mlp_m.down_proj.weight)  # type: ignore[union-attr]