modeling_matriochka.py

"""
modeling_matriochka.py — Matriochka Stacked LLM
================================================
Custom HuggingFace model file enabling `trust_remote_code=True` loading.

The architecture is a nested cascade: sub-model i feeds its last hidden states
into sub-model i+1 as part of its input embeddings.  This means sub-model i+1
REQUIRES sub-model i to have run first — they cannot be used in isolation.

Each branch therefore contains the FULL CASCADE up to and including that size:

    revision="300M"  →  MatriochkaForCausalLM with sub-models [100M, 300M]
    revision="600M"  →  MatriochkaForCausalLM with sub-models [100M, 300M, 600M]
    revision="main"  →  MatriochkaForCausalLM with all sub-models

Usage
-----
    # Load the 300M cascade (downloads only 100M + 300M weights)
    model = AutoModelForCausalLM.from_pretrained(
        "your-org/matriochka-lm",
        revision="300M",
        trust_remote_code=True,
    )
    out = model(input_ids)      # .logits shape: [B, T, vocab]

    # Load the full stack
    model = AutoModelForCausalLM.from_pretrained(
        "your-org/matriochka-lm",
        trust_remote_code=True,
    )

See upload_to_hub() at the bottom for the upload workflow.
"""

from __future__ import annotations

from typing import Dict, List, Optional, Tuple

import torch
import torch.nn as nn
from transformers import (
    AutoConfig,
    AutoModelForCausalLM,
    LlamaConfig,
    LlamaForCausalLM,
    PretrainedConfig,
    PreTrainedModel,
)
from transformers.modeling_outputs import CausalLMOutputWithPast


# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------

class MatriochkaConfig(PretrainedConfig):
    """
    Config for a Matriochka cascade (any prefix of the full stack).

    `sub_model_configs` : list of LlamaConfig kwargs, one per sub-model in this cascade.
    `sub_model_tags` : matching human-readable labels, e.g. ["100M", "300M"].

    embed_tokens width per sub-model:
        index 0  →  own_hidden                      (standard embedding)
        index i  →  own_hidden_i - own_hidden_{i-1} (narrowed so that
                     cat([prev_hs, embed_out]) == own_hidden_i exactly)
    """

    model_type = "matriochka"

    def __init__(
        self,
        sub_model_configs: Optional[List[dict]] = None,  # noqa: keep distinct from HF sub-config scanning
        sub_model_tags: Optional[List[str]] = None,
        vocab_size: int = 49152,
        bos_token_id: int = 1,
        eos_token_id: int = 2,
        junction_type: str = "norm",
        **kwargs,
    ):
        super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)
        self.sub_model_configs = sub_model_configs or []
        self.sub_model_tags = sub_model_tags or []
        self.vocab_size = vocab_size
        self.junction_type = junction_type

    @classmethod
    def from_shape_list(
        cls,
        shapes: List[Tuple[int, int, int]],
        tags: List[str],
        base_model_id: str = "HuggingFaceTB/SmolLM2-135M",
        **kwargs,
    ) -> "MatriochkaConfig":
        """
        shapes : list of (num_layers, num_heads, head_dim)
        tags   : matching size labels, e.g. ["100M","300M","600M","1B"]
        """
        base = AutoConfig.from_pretrained(base_model_id)
        sub_configs = []
        for (layers, heads, head_dim) in shapes:
            hidden = heads * head_dim
            sub_configs.append({
                "num_hidden_layers": layers,
                "num_attention_heads": heads,
                "num_key_value_heads": heads,
                "head_dim": head_dim,
                "hidden_size": hidden,
                "intermediate_size": 4 * hidden,
                "vocab_size": base.vocab_size,
                "max_position_embeddings": base.max_position_embeddings,
                "rope_theta": getattr(base, "rope_theta", 100000.0),
                "rms_norm_eps": base.rms_norm_eps,
                "tie_word_embeddings": False,
            })
        return cls(
            sub_model_configs=sub_configs,
            sub_model_tags=tags,
            vocab_size=base.vocab_size,
            bos_token_id=base.bos_token_id,
            eos_token_id=base.eos_token_id,
            **kwargs,
        )


    def truncated(self, tag: str) -> "MatriochkaConfig":
        """Return a new config containing only sub-models up to and including `tag`."""
        idx = self.sub_model_tags.index(tag)
        return MatriochkaConfig(
            sub_model_configs=self.sub_model_configs[: idx + 1],
            sub_model_tags=self.sub_model_tags[: idx + 1],
            vocab_size=self.vocab_size,
            bos_token_id=self.bos_token_id,
            eos_token_id=self.eos_token_id,
            junction_type=self.junction_type,
        )


# ---------------------------------------------------------------------------
# Internal sub-model wrapper  (not registered with AutoModel)
# ---------------------------------------------------------------------------

class _MatriochkaSubModel(nn.Module):
    """
    Wraps a single LlamaForCausalLM inside the cascade.

    embed_tokens width:
        index 0  →  own_hidden          (normal)
        index i  →  own_hidden - prev_hidden   (narrowed; concat restores own_hidden)

    prev_hidden_size=None signals index 0 (no predecessor).
    """

    def __init__(self, llama_cfg: LlamaConfig, prev_hidden_size: Optional[int], junction_type: str = "norm"):
        super().__init__()
        self.prev_hidden_size = prev_hidden_size
        self.junction_type = junction_type
        self.backbone = LlamaForCausalLM(llama_cfg)

        if prev_hidden_size is not None:
            own_hidden = llama_cfg.hidden_size
            embed_dim = own_hidden - prev_hidden_size
            assert embed_dim > 0, (
                f"own_hidden ({own_hidden}) must be > prev_hidden ({prev_hidden_size})"
            )
            self.backbone.model.embed_tokens = nn.Embedding(
                llama_cfg.vocab_size, embed_dim
            )

    def forward(
        self,
        input_ids: torch.LongTensor,
        attention_mask: Optional[torch.Tensor],
        prev_hidden_states: Optional[torch.Tensor],
        labels: Optional[torch.LongTensor],
        output_hidden_states: bool,
        return_dict: bool,
        **kwargs,
    ):
        # embed_tokens produces (own - prev) dims, or own dims for index 0
        inputs_embeds = self.backbone.get_input_embeddings()(input_ids)

        if self.junction_type == "zero" and self.prev_hidden_size is not None:
            inputs_embeds = 0 * inputs_embeds

        if self.prev_hidden_size is not None and prev_hidden_states is not None:
            # Combine prev_hs with own embedding; result width: own_hidden.
            if self.junction_type == "norm":
                factor = (
                    inputs_embeds.pow(2).mean(-1, keepdim=True).sqrt()
                    / (1e-9 + prev_hidden_states.pow(2).mean(-1, keepdim=True).sqrt())
                )
            else:
                factor = 1.0
            inputs_embeds = torch.cat([prev_hidden_states * factor, inputs_embeds], dim=-1)

        return self.backbone(
            inputs_embeds=inputs_embeds,
            attention_mask=attention_mask,
            labels=labels,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            **kwargs,
        )


# ---------------------------------------------------------------------------
# Main model
# ---------------------------------------------------------------------------

class MatriochkaForCausalLM(PreTrainedModel):
    """
    Matriochka cascade of LLMs.

    forward() runs the full cascade and returns the last sub-model's
    CausalLMOutputWithPast, so .logits / .loss work directly and the model
    is compatible with lm-eval and standard HF inference pipelines.

    forward_all() returns a dict[tag → CausalLMOutputWithPast] for all
    sub-models, useful when computing per-sub-model losses.
    """

    config_class = MatriochkaConfig
    base_model_prefix = "lm_model_dict"
    
    _supports_flash_attn = True
    _supports_sdpa = True
    _supports_flex_attn = True

    _can_compile_fullgraph = True
    _supports_attention_backend = True

    def __init__(self, config: MatriochkaConfig):
        super().__init__(config)
        self.lm_model_dict = nn.ModuleDict()

        prev_hidden: Optional[int] = None
        for tag, sub_cfg_dict in zip(config.sub_model_tags, config.sub_model_configs):
            llama_cfg = LlamaConfig(**sub_cfg_dict)
            self.lm_model_dict[tag] = _MatriochkaSubModel(llama_cfg, prev_hidden, config.junction_type)
            prev_hidden = sub_cfg_dict["hidden_size"]

        self.post_init()

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        labels: Optional[torch.LongTensor] = None,
        output_hidden_states: bool = True,
        return_dict: bool = True,
        **kwargs,
    ) -> CausalLMOutputWithPast:
        """
        Run the full cascade. Returns the last sub-model's CausalLMOutputWithPast
        so that .logits and .loss are directly accessible.
        Use forward_all() to get outputs for every sub-model.
        """
        prev_hs: Optional[torch.Tensor] = None
        out = None
        for sub_model in self.lm_model_dict.values():
            out = sub_model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                prev_hidden_states=prev_hs,
                labels=labels,
                output_hidden_states=True,
                return_dict=True,
                **kwargs,
            )
            prev_hs = out.hidden_states[-1]
        return out

    def forward_all(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        labels: Optional[torch.LongTensor] = None,
        **kwargs,
    ) -> Dict[str, CausalLMOutputWithPast]:
        """Run the full cascade and return outputs for every sub-model."""
        results: Dict[str, CausalLMOutputWithPast] = {}
        prev_hs: Optional[torch.Tensor] = None
        for tag, sub_model in self.lm_model_dict.items():
            out = sub_model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                prev_hidden_states=prev_hs,
                labels=labels,
                output_hidden_states=True,
                return_dict=True,
                **kwargs,
            )
            results[tag] = out
            prev_hs = out.hidden_states[-1]
        return results

    def forward_up_to(
        self,
        tag: str,
        input_ids: torch.LongTensor,
        **kwargs,
    ) -> CausalLMOutputWithPast:
        """Run the cascade up to and including `tag`."""
        prev_hs = None
        out = None
        for t, sub_model in self.lm_model_dict.items():
            out = sub_model(
                input_ids=input_ids,
                prev_hidden_states=prev_hs,
                output_hidden_states=True,
                return_dict=True,
                **kwargs,
            )
            prev_hs = out.hidden_states[-1]
            if t == tag:
                break
        return out


# ---------------------------------------------------------------------------
# Auto-class registration
# ---------------------------------------------------------------------------

AutoConfig.register("matriochka", MatriochkaConfig)
AutoModelForCausalLM.register(MatriochkaConfig, MatriochkaForCausalLM)


# ---------------------------------------------------------------------------
# Smoke-test  (python modeling_matriochka.py)
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    SHAPES = [(16, 8, 64), (11, 14, 64), (8, 22, 64), (8, 18, 96)]
    TAGS   = ["100M", "300M", "600M", "1B"]

    hiddens = [h * d for _, h, d in SHAPES]
    print("hidden sizes:         ", hiddens)
    print("embed_tokens out dims:", [hiddens[0]] + [hiddens[i] - hiddens[i-1] for i in range(1, len(hiddens))])

    cfg = MatriochkaConfig(
        sub_model_configs=[
            {
                "num_hidden_layers": l, "num_attention_heads": h,
                "num_key_value_heads": h, "head_dim": d,
                "hidden_size": h * d, "intermediate_size": 4 * h * d,
                "vocab_size": 49152, "max_position_embeddings": 2048,
                "tie_word_embeddings": False,
            }
            for l, h, d in SHAPES
        ],
        sub_model_tags=TAGS,
    )

    # Verify all cascade prefixes work independently
    for i, tag in enumerate(TAGS):
        trunc_cfg = cfg.truncated(tag)
        m = MatriochkaForCausalLM(trunc_cfg)
        n = sum(p.numel() for p in m.parameters()) / 1e6
        dummy = torch.randint(0, 49152, (1, 32))
        out = m(dummy)
        print(f"  [{', '.join(trunc_cfg.sub_model_tags)}]  {n:.1f}M params  |  logits: {out.logits.shape}")