""" 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}")