Upload Matriochka cascade up to main

config.json

@@ -9,6 +9,7 @@
   "bos_token_id": 0,
   "dtype": "float32",
   "eos_token_id": 0,
+  "junction_type": "norm",
   "model_type": "matriochka",
   "sub_model_configs": [
     {
@@ -20,7 +21,7 @@
       "num_hidden_layers": 24,
       "num_key_value_heads": 16,
       "rms_norm_eps": 1e-05,
-      "rope_theta": 100000,
+      "rope_theta": 100000.0,
       "tie_word_embeddings": false,
       "vocab_size": 49152
     },
@@ -33,7 +34,7 @@
       "num_hidden_layers": 10,
       "num_key_value_heads": 24,
       "rms_norm_eps": 1e-05,
-      "rope_theta": 100000,
+      "rope_theta": 100000.0,
       "tie_word_embeddings": false,
       "vocab_size": 49152
     },
@@ -46,16 +47,16 @@
       "num_hidden_layers": 5,
       "num_key_value_heads": 34,
       "rms_norm_eps": 1e-05,
-      "rope_theta": 100000,
+      "rope_theta": 100000.0,
       "tie_word_embeddings": false,
       "vocab_size": 49152
     }
   ],
   "sub_model_tags": [
-    "500M_5B",
-    "1-5B_5B",
-    "3B_5B"
+    "500M_60B_cd",
+    "1-5B_60B_cd",
+    "3B_60B_cd"
   ],
-  "transformers_version": "4.57.0",
+  "transformers_version": "5.6.2",
   "vocab_size": 49152
 }

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:856e925904135a66f0b60974c4dd153a9b1d6e1f1dc40d9a7abbadd94a8b0e70
+size 13434995600

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(
@@ -107,7 +109,7 @@ class MatriochkaConfig(PretrainedConfig):
                 "intermediate_size": 4 * hidden,
                 "vocab_size": base.vocab_size,
                 "max_position_embeddings": base.max_position_embeddings,
-                "rope_theta": getattr(base, "rope_theta", 10000.0),
+                "rope_theta": getattr(base, "rope_theta", 100000.0),
                 "rms_norm_eps": base.rms_norm_eps,
                 "tie_word_embeddings": False,
             })
@@ -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()