Upload Matriochka cascade up to main

config.json

@@ -9,6 +9,7 @@
   "bos_token_id": 0,
   "dtype": "float32",
   "eos_token_id": 0,
+  "junction_type": "none",
   "model_type": "matriochka",
   "sub_model_configs": [
     {

modeling_matriochka.py

@@ -75,12 +75,14 @@ class MatriochkaConfig(PretrainedConfig):
         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(
@@ -120,6 +122,7 @@ class MatriochkaConfig(PretrainedConfig):
             **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)
@@ -129,6 +132,7 @@ class MatriochkaConfig(PretrainedConfig):
             vocab_size=self.vocab_size,
             bos_token_id=self.bos_token_id,
             eos_token_id=self.eos_token_id,
+            junction_type=self.junction_type,
         )
 
 
@@ -147,9 +151,10 @@ class _MatriochkaSubModel(nn.Module):
     prev_hidden_size=None signals index 0 (no predecessor).
     """
 
-    def __init__(self, llama_cfg: LlamaConfig, prev_hidden_size: Optional[int]):
+    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:
@@ -175,13 +180,18 @@ class _MatriochkaSubModel(nn.Module):
         # 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:
-            # Scale prev_hs norm to match own embedding norm, then concat.
-            # Result width: prev_hidden + (own - prev) = own_hidden  ✓
-            factor = (
-                inputs_embeds.pow(2).mean(-1, keepdim=True).sqrt()
-                / (1e-9 + prev_hidden_states.pow(2).mean(-1, keepdim=True).sqrt())
-            )
+            # 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(
@@ -227,7 +237,7 @@ class MatriochkaForCausalLM(PreTrainedModel):
         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)
+            self.lm_model_dict[tag] = _MatriochkaSubModel(llama_cfg, prev_hidden, config.junction_type)
             prev_hidden = sub_cfg_dict["hidden_size"]
 
         self.post_init()