Upload Matriochka cascade up to main

modeling_matriochka.py

@@ -107,7 +107,7 @@ class MatriochkaConfig(PretrainedConfig):
                 "intermediate_size": 4 * hidden,
                 "vocab_size": base.vocab_size,
                 "max_position_embeddings": base.max_position_embeddings,
-                "rope_theta": base.rope_theta,
+                "rope_theta": getattr(base, "rope_theta", 10000.0),
                 "rms_norm_eps": base.rms_norm_eps,
                 "tie_word_embeddings": False,
             })
@@ -147,11 +147,9 @@ class _MatriochkaSubModel(nn.Module):
     prev_hidden_size=None signals index 0 (no predecessor).
     """
 
-    def __init__(self, llama_cfg: LlamaConfig, prev_hidden_size: Optional[int], attn_implementation: Optional[str] = None):
+    def __init__(self, llama_cfg: LlamaConfig, prev_hidden_size: Optional[int]):
         super().__init__()
         self.prev_hidden_size = prev_hidden_size
-        if attn_implementation is not None:
-            llama_cfg._attn_implementation = attn_implementation
         self.backbone = LlamaForCausalLM(llama_cfg)
 
         if prev_hidden_size is not None:
@@ -214,17 +212,15 @@ class MatriochkaForCausalLM(PreTrainedModel):
 
     config_class = MatriochkaConfig
     base_model_prefix = "lm_model_dict"
-    _supports_sdpa = True
 
     def __init__(self, config: MatriochkaConfig):
         super().__init__(config)
         self.lm_model_dict = nn.ModuleDict()
 
-        attn_implementation = getattr(config, "_attn_implementation", None)
         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, attn_implementation=attn_implementation)
+            self.lm_model_dict[tag] = _MatriochkaSubModel(llama_cfg, prev_hidden)
             prev_hidden = sub_cfg_dict["hidden_size"]
 
         self.post_init()
@@ -313,116 +309,6 @@ AutoConfig.register("matriochka", MatriochkaConfig)
 AutoModelForCausalLM.register(MatriochkaConfig, MatriochkaForCausalLM)
 
 
-# ---------------------------------------------------------------------------
-# Upload helper
-# ---------------------------------------------------------------------------
-
-def upload_to_hub(
-    checkpoint_path: str,
-    repo_id: str,
-    shapes_config: List[Tuple[int, int, int]],
-    tags: List[str],
-    base_model_id: str = "HuggingFaceTB/SmolLM2-135M",
-    token: Optional[str] = None,
-):
-    """
-    Upload workflow.
-
-    Pushes one branch per tag, each containing the FULL CASCADE up to that tag:
-        main  →  [100M, 300M, 600M, 1B]   (all sub-models)
-        1B    →  [100M, 300M, 600M, 1B]   (same as main)
-        600M  →  [100M, 300M, 600M]
-        300M  →  [100M, 300M]
-        100M  →  [100M]
-
-    This means `from_pretrained(..., revision="300M")` downloads exactly the
-    weights needed to run the 300M cascade and nothing more.
-
-    Example
-    -------
-        from modeling_matriochka import upload_to_hub
-        upload_to_hub(
-            checkpoint_path="/data/ckpts/.../epoch=0-step=200000.ckpt",
-            repo_id="your-org/matriochka-lm",
-            shapes_config=[(16,8,64),(11,14,64),(8,22,64),(8,18,96)],
-            tags=["100M","300M","600M","1B"],
-        )
-    """
-    import os, re, shutil, tempfile
-    from huggingface_hub import HfApi
-
-    api = HfApi(token=token)
-    this_file = os.path.abspath(__file__)
-
-    # ── 1. Build full config & model ───────────────────────────────────────
-    full_cfg = MatriochkaConfig.from_shape_list(shapes_config, tags, base_model_id)
-    full_model = MatriochkaForCausalLM(full_cfg)
-
-    # ── 2. Load Lightning checkpoint ───────────────────────────────────────
-    print(f"Loading checkpoint: {checkpoint_path}")
-    ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
-    raw_sd = ckpt["state_dict"]
-
-    # Lightning prepends "model." — strip it.
-    # Checkpoint keys:  lm_model_dict.100M.model.layers.0…
-    # Our keys:         lm_model_dict.100M.backbone.model.layers.0…
-    def remap_key(k: str) -> str:
-        k = k.removeprefix("model.")
-        return re.sub(
-            r"^(lm_model_dict\.[^.]+\.)(?!backbone\.)",
-            r"\1backbone.",
-            k,
-        )
-    cleaned = {remap_key(k): v for k, v in raw_sd.items()}
-
-    missing, unexpected = full_model.load_state_dict(cleaned, strict=False)
-    if missing:
-        print(f"[WARN] {len(missing)} missing keys, e.g.: {missing[:3]}")
-    if unexpected:
-        print(f"[WARN] {len(unexpected)} unexpected keys, e.g.: {unexpected[:3]}")
-
-    # ── 3. Push each branch ────────────────────────────────────────────────
-    api.create_repo(repo_id, exist_ok=True, repo_type="model")
-
-    def _push(model: MatriochkaForCausalLM, cfg: MatriochkaConfig, branch: str):
-        cfg.auto_map = {
-            "AutoConfig": "modeling_matriochka.MatriochkaConfig",
-            "AutoModelForCausalLM": "modeling_matriochka.MatriochkaForCausalLM",
-        }
-        with tempfile.TemporaryDirectory() as tmpdir:
-            model.save_pretrained(tmpdir)
-            cfg.save_pretrained(tmpdir)
-            shutil.copy(this_file, os.path.join(tmpdir, "modeling_matriochka.py"))
-            if branch != "main":
-                try:
-                    api.create_branch(repo_id, branch=branch, repo_type="model")
-                except Exception:
-                    pass
-            api.upload_folder(
-                folder_path=tmpdir, repo_id=repo_id, repo_type="model",
-                revision=branch,
-                commit_message=f"Upload Matriochka cascade up to {branch}",
-            )
-        print(f"✓ {branch}  →  {repo_id} [{cfg.sub_model_tags}]")
-
-    # main = full stack
-    _push(full_model, full_cfg, "main")
-
-    # per-tag branches: truncated cascade
-    for tag in tags:
-        trunc_cfg = full_cfg.truncated(tag)
-        # Build a fresh truncated model and copy the relevant sub-model weights
-        trunc_model = MatriochkaForCausalLM(trunc_cfg)
-        trunc_sd = {
-            k: v for k, v in full_model.state_dict().items()
-            if any(k.startswith(f"lm_model_dict.{t}.") for t in trunc_cfg.sub_model_tags)
-        }
-        trunc_model.load_state_dict(trunc_sd, strict=True)
-        _push(trunc_model, trunc_cfg, tag)
-
-    print("All done!")
-
-
 # ---------------------------------------------------------------------------
 # Smoke-test  (python modeling_matriochka.py)
 # ---------------------------------------------------------------------------