add 9-config hparam sweep + new_layer_lr_mul param-groups support

configs/base.toml

@@ -33,6 +33,7 @@ student_dtype        = "bfloat16"
 teacher_dtype        = "bfloat16"
 mixed_precision      = "bf16"
 kl_chunk_size        = 0
+new_layer_lr_mul     = 1.0
 
 [eval]
 every_steps = 5

configs/grow40_simple.toml

@@ -34,6 +34,7 @@ student_dtype        = "bfloat16"
 teacher_dtype        = "bfloat16"
 mixed_precision      = "bf16"
 kl_chunk_size        = 0
+new_layer_lr_mul     = 1.0
 
 [eval]
 every_steps = 50

configs/grow40_winning.toml

@@ -35,6 +35,7 @@ student_dtype        = "bfloat16"
 teacher_dtype        = "bfloat16"
 mixed_precision      = "bf16"
 kl_chunk_size        = 256
+new_layer_lr_mul     = 1.0
 
 [eval]
 every_steps = 50

configs/grow40_winning_v2.toml

@@ -34,6 +34,7 @@ student_dtype        = "bfloat16"
 teacher_dtype        = "bfloat16"
 mixed_precision      = "bf16"
 kl_chunk_size        = 256
+new_layer_lr_mul     = 1.0
 
 [eval]
 every_steps = 50

configs/replicate_zero4.toml

@@ -33,6 +33,7 @@ student_dtype        = "float32"
 teacher_dtype        = "bfloat16"
 mixed_precision      = "bf16"
 kl_chunk_size        = 256
+new_layer_lr_mul     = 1.0
 
 [eval]
 every_steps = 50

configs/sweep/A_resume_lr1e7_cos.toml

@@ -0,0 +1,52 @@
+# Resume from grow40_winning best (eval kl 0.2219). Lower peak LR to avoid the
+# overshoot we saw at 5e-7. Cosine warmup 100, 1500 steps.
+
+[model]
+teacher    = "Qwen/Qwen3.5-35B-A3B"
+student    = "./out/grow40_winning/best"
+tokenizer  = "Qwen/Qwen3.5-35B-A3B"
+
+[data]
+dataset        = "karpathy/climbmix-400b-shuffle"
+text_field     = "text"
+min_chars      = 2560
+max_seq_len    = 2048
+kl_start_pos   = 128
+seed           = 6767
+shuffle_buffer = 10000
+
+[train]
+seed                 = 6767
+lr                   = 1.0e-7
+schedule             = "cosine"
+warmup_steps         = 100
+weight_decay         = 0.0
+grad_clip            = 1.0
+betas                = [0.9, 0.999]
+eps                  = 1.0e-3
+samples_per_step     = 4
+micro_batch_size     = 4
+max_steps            = 1500
+grad_checkpointing   = true
+attn_implementation  = "flash_attention_2"
+student_dtype        = "bfloat16"
+teacher_dtype        = "bfloat16"
+mixed_precision      = "bf16"
+kl_chunk_size        = 256
+new_layer_lr_mul     = 1.0
+
+[eval]
+every_steps = 50
+samples     = 500
+seed        = 4242
+
+[log]
+wandb         = true
+wandb_project = "distil-subnet97"
+wandb_run     = "A_resume_lr1e7_cos"
+log_every     = 1
+output_dir    = "./out/sweep/A_resume_lr1e7_cos"
+
+[init]
+zero_layers        = []
+target_num_layers  = 40

configs/sweep/B_resume_lr5e8_cos.toml

@@ -0,0 +1,51 @@
+# Resume from grow40_winning best (eval kl 0.2219). Even lower peak LR.
+
+[model]
+teacher    = "Qwen/Qwen3.5-35B-A3B"
+student    = "./out/grow40_winning/best"
+tokenizer  = "Qwen/Qwen3.5-35B-A3B"
+
+[data]
+dataset        = "karpathy/climbmix-400b-shuffle"
+text_field     = "text"
+min_chars      = 2560
+max_seq_len    = 2048
+kl_start_pos   = 128
+seed           = 6767
+shuffle_buffer = 10000
+
+[train]
+seed                 = 6767
+lr                   = 5.0e-8
+schedule             = "cosine"
+warmup_steps         = 100
+weight_decay         = 0.0
+grad_clip            = 1.0
+betas                = [0.9, 0.999]
+eps                  = 1.0e-3
+samples_per_step     = 4
+micro_batch_size     = 4
+max_steps            = 1500
+grad_checkpointing   = true
+attn_implementation  = "flash_attention_2"
+student_dtype        = "bfloat16"
+teacher_dtype        = "bfloat16"
+mixed_precision      = "bf16"
+kl_chunk_size        = 256
+new_layer_lr_mul     = 1.0
+
+[eval]
+every_steps = 50
+samples     = 500
+seed        = 4242
+
+[log]
+wandb         = true
+wandb_project = "distil-subnet97"
+wandb_run     = "B_resume_lr5e8_cos"
+log_every     = 1
+output_dir    = "./out/sweep/B_resume_lr5e8_cos"
+
+[init]
+zero_layers        = []
+target_num_layers  = 40

configs/sweep/C_resume_lr2e8_cos.toml

@@ -0,0 +1,51 @@
+# Resume from grow40_winning best. Very small LR - basically a fine-tune.
+
+[model]
+teacher    = "Qwen/Qwen3.5-35B-A3B"
+student    = "./out/grow40_winning/best"
+tokenizer  = "Qwen/Qwen3.5-35B-A3B"
+
+[data]
+dataset        = "karpathy/climbmix-400b-shuffle"
+text_field     = "text"
+min_chars      = 2560
+max_seq_len    = 2048
+kl_start_pos   = 128
+seed           = 6767
+shuffle_buffer = 10000
+
+[train]
+seed                 = 6767
+lr                   = 2.0e-8
+schedule             = "cosine"
+warmup_steps         = 100
+weight_decay         = 0.0
+grad_clip            = 1.0
+betas                = [0.9, 0.999]
+eps                  = 1.0e-3
+samples_per_step     = 4
+micro_batch_size     = 4
+max_steps            = 1500
+grad_checkpointing   = true
+attn_implementation  = "flash_attention_2"
+student_dtype        = "bfloat16"
+teacher_dtype        = "bfloat16"
+mixed_precision      = "bf16"
+kl_chunk_size        = 256
+new_layer_lr_mul     = 1.0
+
+[eval]
+every_steps = 50
+samples     = 500
+seed        = 4242
+
+[log]
+wandb         = true
+wandb_project = "distil-subnet97"
+wandb_run     = "C_resume_lr2e8_cos"
+log_every     = 1
+output_dir    = "./out/sweep/C_resume_lr2e8_cos"
+
+[init]
+zero_layers        = []
+target_num_layers  = 40

configs/sweep/D_resume_lr1e7_const.toml

@@ -0,0 +1,51 @@
+# Resume from grow40_winning best. Constant LR (no schedule overshoot at all).
+
+[model]
+teacher    = "Qwen/Qwen3.5-35B-A3B"
+student    = "./out/grow40_winning/best"
+tokenizer  = "Qwen/Qwen3.5-35B-A3B"
+
+[data]
+dataset        = "karpathy/climbmix-400b-shuffle"
+text_field     = "text"
+min_chars      = 2560
+max_seq_len    = 2048
+kl_start_pos   = 128
+seed           = 6767
+shuffle_buffer = 10000
+
+[train]
+seed                 = 6767
+lr                   = 1.0e-7
+schedule             = "constant"
+warmup_steps         = 0
+weight_decay         = 0.0
+grad_clip            = 1.0
+betas                = [0.9, 0.999]
+eps                  = 1.0e-3
+samples_per_step     = 4
+micro_batch_size     = 4
+max_steps            = 1500
+grad_checkpointing   = true
+attn_implementation  = "flash_attention_2"
+student_dtype        = "bfloat16"
+teacher_dtype        = "bfloat16"
+mixed_precision      = "bf16"
+kl_chunk_size        = 256
+new_layer_lr_mul     = 1.0
+
+[eval]
+every_steps = 50
+samples     = 500
+seed        = 4242
+
+[log]
+wandb         = true
+wandb_project = "distil-subnet97"
+wandb_run     = "D_resume_lr1e7_const"
+log_every     = 1
+output_dir    = "./out/sweep/D_resume_lr1e7_const"
+
+[init]
+zero_layers        = []
+target_num_layers  = 40

configs/sweep/E_resume_lr5e8_b95.toml

@@ -0,0 +1,52 @@
+# Resume from grow40_winning best. Smaller second-moment memory (beta2=0.95)
+# so Adam stabilizes faster. Same low LR.
+
+[model]
+teacher    = "Qwen/Qwen3.5-35B-A3B"
+student    = "./out/grow40_winning/best"
+tokenizer  = "Qwen/Qwen3.5-35B-A3B"
+
+[data]
+dataset        = "karpathy/climbmix-400b-shuffle"
+text_field     = "text"
+min_chars      = 2560
+max_seq_len    = 2048
+kl_start_pos   = 128
+seed           = 6767
+shuffle_buffer = 10000
+
+[train]
+seed                 = 6767
+lr                   = 5.0e-8
+schedule             = "cosine"
+warmup_steps         = 100
+weight_decay         = 0.0
+grad_clip            = 1.0
+betas                = [0.9, 0.95]
+eps                  = 1.0e-8
+samples_per_step     = 4
+micro_batch_size     = 4
+max_steps            = 1500
+grad_checkpointing   = true
+attn_implementation  = "flash_attention_2"
+student_dtype        = "bfloat16"
+teacher_dtype        = "bfloat16"
+mixed_precision      = "bf16"
+kl_chunk_size        = 256
+new_layer_lr_mul     = 1.0
+
+[eval]
+every_steps = 50
+samples     = 500
+seed        = 4242
+
+[log]
+wandb         = true
+wandb_project = "distil-subnet97"
+wandb_run     = "E_resume_lr5e8_b95"
+log_every     = 1
+output_dir    = "./out/sweep/E_resume_lr5e8_b95"
+
+[init]
+zero_layers        = []
+target_num_layers  = 40

configs/sweep/F_cold_lr1e7_grow40.toml

@@ -0,0 +1,51 @@
+# Cold start, 40 layers, lower peak LR than the original winning recipe.
+
+[model]
+teacher    = "Qwen/Qwen3.5-35B-A3B"
+student    = "Troiaaa/m-6a3lnzvb"
+tokenizer  = "Qwen/Qwen3.5-35B-A3B"
+
+[data]
+dataset        = "karpathy/climbmix-400b-shuffle"
+text_field     = "text"
+min_chars      = 2560
+max_seq_len    = 2048
+kl_start_pos   = 128
+seed           = 6767
+shuffle_buffer = 10000
+
+[train]
+seed                 = 6767
+lr                   = 1.0e-7
+schedule             = "cosine"
+warmup_steps         = 100
+weight_decay         = 0.0
+grad_clip            = 1.0
+betas                = [0.9, 0.999]
+eps                  = 1.0e-3
+samples_per_step     = 4
+micro_batch_size     = 4
+max_steps            = 2000
+grad_checkpointing   = true
+attn_implementation  = "flash_attention_2"
+student_dtype        = "bfloat16"
+teacher_dtype        = "bfloat16"
+mixed_precision      = "bf16"
+kl_chunk_size        = 256
+new_layer_lr_mul     = 1.0
+
+[eval]
+every_steps = 50
+samples     = 500
+seed        = 4242
+
+[log]
+wandb         = true
+wandb_project = "distil-subnet97"
+wandb_run     = "F_cold_lr1e7_grow40"
+log_every     = 1
+output_dir    = "./out/sweep/F_cold_lr1e7_grow40"
+
+[init]
+zero_layers        = []
+target_num_layers  = 40

configs/sweep/G_cold_lr2e7_grow40.toml

@@ -0,0 +1,51 @@
+# Cold start, 40 layers, lr=2e-7 (between 1e-7 and the failing 5e-7).
+
+[model]
+teacher    = "Qwen/Qwen3.5-35B-A3B"
+student    = "Troiaaa/m-6a3lnzvb"
+tokenizer  = "Qwen/Qwen3.5-35B-A3B"
+
+[data]
+dataset        = "karpathy/climbmix-400b-shuffle"
+text_field     = "text"
+min_chars      = 2560
+max_seq_len    = 2048
+kl_start_pos   = 128
+seed           = 6767
+shuffle_buffer = 10000
+
+[train]
+seed                 = 6767
+lr                   = 2.0e-7
+schedule             = "cosine"
+warmup_steps         = 100
+weight_decay         = 0.0
+grad_clip            = 1.0
+betas                = [0.9, 0.999]
+eps                  = 1.0e-3
+samples_per_step     = 4
+micro_batch_size     = 4
+max_steps            = 2000
+grad_checkpointing   = true
+attn_implementation  = "flash_attention_2"
+student_dtype        = "bfloat16"
+teacher_dtype        = "bfloat16"
+mixed_precision      = "bf16"
+kl_chunk_size        = 256
+new_layer_lr_mul     = 1.0
+
+[eval]
+every_steps = 50
+samples     = 500
+seed        = 4242
+
+[log]
+wandb         = true
+wandb_project = "distil-subnet97"
+wandb_run     = "G_cold_lr2e7_grow40"
+log_every     = 1
+output_dir    = "./out/sweep/G_cold_lr2e7_grow40"
+
+[init]
+zero_layers        = []
+target_num_layers  = 40

configs/sweep/H_cold_lr1e7_32L.toml

@@ -0,0 +1,52 @@
+# Cold start, 32 layers (NO grow), lower LR. Tests whether the +8 layers were
+# helping at all once we use the right LR.
+
+[model]
+teacher    = "Qwen/Qwen3.5-35B-A3B"
+student    = "Troiaaa/m-6a3lnzvb"
+tokenizer  = "Qwen/Qwen3.5-35B-A3B"
+
+[data]
+dataset        = "karpathy/climbmix-400b-shuffle"
+text_field     = "text"
+min_chars      = 2560
+max_seq_len    = 2048
+kl_start_pos   = 128
+seed           = 6767
+shuffle_buffer = 10000
+
+[train]
+seed                 = 6767
+lr                   = 1.0e-7
+schedule             = "cosine"
+warmup_steps         = 100
+weight_decay         = 0.0
+grad_clip            = 1.0
+betas                = [0.9, 0.999]
+eps                  = 1.0e-3
+samples_per_step     = 4
+micro_batch_size     = 4
+max_steps            = 2000
+grad_checkpointing   = true
+attn_implementation  = "flash_attention_2"
+student_dtype        = "bfloat16"
+teacher_dtype        = "bfloat16"
+mixed_precision      = "bf16"
+kl_chunk_size        = 256
+new_layer_lr_mul     = 1.0
+
+[eval]
+every_steps = 50
+samples     = 500
+seed        = 4242
+
+[log]
+wandb         = true
+wandb_project = "distil-subnet97"
+wandb_run     = "H_cold_lr1e7_32L"
+log_every     = 1
+output_dir    = "./out/sweep/H_cold_lr1e7_32L"
+
+[init]
+zero_layers        = []
+target_num_layers  = 32

configs/sweep/I_cold_paramgroups_grow40.toml

@@ -0,0 +1,52 @@
+# Cold start, 40 layers, low LR for original layers + 5x for the new ones.
+# Lets the new layers wake up faster without disturbing the trained layers.
+
+[model]
+teacher    = "Qwen/Qwen3.5-35B-A3B"
+student    = "Troiaaa/m-6a3lnzvb"
+tokenizer  = "Qwen/Qwen3.5-35B-A3B"
+
+[data]
+dataset        = "karpathy/climbmix-400b-shuffle"
+text_field     = "text"
+min_chars      = 2560
+max_seq_len    = 2048
+kl_start_pos   = 128
+seed           = 6767
+shuffle_buffer = 10000
+
+[train]
+seed                 = 6767
+lr                   = 1.0e-7
+schedule             = "cosine"
+warmup_steps         = 100
+weight_decay         = 0.0
+grad_clip            = 1.0
+betas                = [0.9, 0.999]
+eps                  = 1.0e-3
+samples_per_step     = 4
+micro_batch_size     = 4
+max_steps            = 2000
+grad_checkpointing   = true
+attn_implementation  = "flash_attention_2"
+student_dtype        = "bfloat16"
+teacher_dtype        = "bfloat16"
+mixed_precision      = "bf16"
+kl_chunk_size        = 256
+new_layer_lr_mul     = 5.0
+
+[eval]
+every_steps = 50
+samples     = 500
+seed        = 4242
+
+[log]
+wandb         = true
+wandb_project = "distil-subnet97"
+wandb_run     = "I_cold_paramgroups_grow40"
+log_every     = 1
+output_dir    = "./out/sweep/I_cold_paramgroups_grow40"
+
+[init]
+zero_layers        = []
+target_num_layers  = 40

configs/zero_14_17.toml

@@ -34,6 +34,7 @@ student_dtype        = "bfloat16"
 teacher_dtype        = "bfloat16"
 mixed_precision      = "bf16"
 kl_chunk_size        = 0
+new_layer_lr_mul     = 1.0
 
 [eval]
 every_steps = 50

distill.py

@@ -72,6 +72,7 @@ REQUIRED_KEYS = {
         "mixed_precision",
         "kl_chunk_size",
         "micro_batch_size",
+        "new_layer_lr_mul",
     ),
     "eval": ("every_steps", "samples", "seed"),
     "log": ("wandb", "wandb_project", "wandb_run", "log_every", "output_dir"),
@@ -403,15 +404,48 @@ def kl_loss_masked(student_logits, teacher_logits, attention_mask, start_pos, ch
 # Optimizer / scheduler
 # ----------------------------------------------------------------------------
 
-def make_optimizer(model, train_cfg):
-    return AdamW(
-        [p for p in model.parameters() if p.requires_grad],
-        lr=train_cfg["lr"],
+def make_optimizer(model, train_cfg, new_layer_indices=None):
+    """Create AdamW. If `new_layer_lr_mul != 1.0` and we know which layers are
+    'new' (returned from grow_layers), put their params in a separate group with
+    a multiplied LR. Useful for the 'wake up new layers without disturbing the
+    old ones' regime."""
+    base_lr = train_cfg["lr"]
+    mul = train_cfg["new_layer_lr_mul"]
+    common = dict(
         weight_decay=train_cfg["weight_decay"],
         betas=tuple(train_cfg["betas"]),
         eps=train_cfg["eps"],
     )
 
+    if not new_layer_indices or mul == 1.0:
+        return AdamW(
+            [p for p in model.parameters() if p.requires_grad],
+            lr=base_lr,
+            **common,
+        )
+
+    inner = get_inner_with_layers(model)
+    new_pids = set()
+    for idx in new_layer_indices:
+        for p in inner.layers[idx].parameters():
+            if p.requires_grad:
+                new_pids.add(id(p))
+
+    new_params = []
+    rest_params = []
+    for p in model.parameters():
+        if not p.requires_grad:
+            continue
+        (new_params if id(p) in new_pids else rest_params).append(p)
+
+    return AdamW(
+        [
+            {"params": rest_params, "lr": base_lr},
+            {"params": new_params, "lr": base_lr * mul},
+        ],
+        **common,
+    )
+
 
 def make_scheduler(optimizer, train_cfg):
     schedule = train_cfg["schedule"]
@@ -522,8 +556,10 @@ def main():
     # ---- Layer modifications: grow first, then zero (composable)
     target_n = cfg["init"]["target_num_layers"]
     cur_n = len(get_inner_with_layers(student).layers)
+    new_layer_indices = []
     if target_n != cur_n:
         new_n, new_zeroed = grow_layers(student, target_n)
+        new_layer_indices = [idx for idx, _ in new_zeroed]
         if accelerator.is_main_process:
             log.info(f"Grew student from {cur_n} -> {new_n} layers")
             for idx, names in new_zeroed:
@@ -538,8 +574,14 @@ def main():
     teacher = teacher.to(accelerator.device)
 
     # ---- Optimizer / scheduler
-    optimizer = make_optimizer(student, cfg["train"])
+    optimizer = make_optimizer(student, cfg["train"], new_layer_indices=new_layer_indices)
     scheduler = make_scheduler(optimizer, cfg["train"])
+    if accelerator.is_main_process and len(optimizer.param_groups) > 1:
+        log.info(
+            f"Param groups: rest lr={optimizer.param_groups[0]['lr']:.2e}, "
+            f"new lr={optimizer.param_groups[1]['lr']:.2e} "
+            f"({len(new_layer_indices)} layers grown)"
+        )
 
     # NB: do NOT pass `scheduler` to accelerator.prepare. When prepared, accelerate
     # advances the scheduler by `num_processes` steps per call (to match the

scripts/backup_to_hf.py

@@ -22,10 +22,20 @@ INCLUDE = [
     "configs/grow40_winning.toml",
     "configs/grow40_simple.toml",
     "configs/grow40_winning_v2.toml",
+    "configs/sweep/A_resume_lr1e7_cos.toml",
+    "configs/sweep/B_resume_lr5e8_cos.toml",
+    "configs/sweep/C_resume_lr2e8_cos.toml",
+    "configs/sweep/D_resume_lr1e7_const.toml",
+    "configs/sweep/E_resume_lr5e8_b95.toml",
+    "configs/sweep/F_cold_lr1e7_grow40.toml",
+    "configs/sweep/G_cold_lr2e7_grow40.toml",
+    "configs/sweep/H_cold_lr1e7_32L.toml",
+    "configs/sweep/I_cold_paramgroups_grow40.toml",
     "configs/accelerate.yaml",
     "scripts/backup_to_hf.py",
     "scripts/run_sweep.sh",
     "scripts/run_sweep_rerun.sh",
+    "scripts/run_hparam_sweep.sh",
     "pyproject.toml",
     "requirements.lock.txt",
 ]

scripts/run_hparam_sweep.sh

@@ -0,0 +1,54 @@
+#!/usr/bin/env bash
+# Hyperparameter sweep over 9 configs that try to push past grow40_winning's 0.2219.
+#
+# Each config grabs all 8 GPUs via accelerate, so they run sequentially.
+# Output goes to logs/<run>.log; the master log goes to logs/sweep_hparam_master.log.
+# Reads HF_TOKEN, HUGGING_FACE_HUB_TOKEN, WANDB_API_KEY from the calling env.
+#
+# Launch in the background with:
+#   nohup ./scripts/run_hparam_sweep.sh > logs/sweep_hparam_master.log 2>&1 &
+
+set -uo pipefail
+cd "$(dirname "$0")/.."
+
+CONFIGS=(
+    "configs/sweep/A_resume_lr1e7_cos.toml"
+    "configs/sweep/B_resume_lr5e8_cos.toml"
+    "configs/sweep/C_resume_lr2e8_cos.toml"
+    "configs/sweep/D_resume_lr1e7_const.toml"
+    "configs/sweep/E_resume_lr5e8_b95.toml"
+    "configs/sweep/F_cold_lr1e7_grow40.toml"
+    "configs/sweep/G_cold_lr2e7_grow40.toml"
+    "configs/sweep/H_cold_lr1e7_32L.toml"
+    "configs/sweep/I_cold_paramgroups_grow40.toml"
+)
+
+LOG_DIR="logs"
+mkdir -p "$LOG_DIR"
+
+for cfg in "${CONFIGS[@]}"; do
+    name="$(basename "$cfg" .toml)"
+    log="$LOG_DIR/$name.log"
+    echo ">>> [$(date '+%F %T')] starting $name -> $log"
+    .venv/bin/accelerate launch \
+        --config_file configs/accelerate.yaml \
+        distill.py \
+        --config "$cfg" \
+        > "$log" 2>&1
+    rc=$?
+    best_line=$(grep -E "Best eval KL" "$log" | tail -1)
+    echo "<<< [$(date '+%F %T')] finished $name (exit=$rc) ${best_line}"
+    if [[ $rc -ne 0 ]]; then
+        echo "    last 12 lines of $log:"
+        tail -12 "$log" | sed 's/^/      /'
+    fi
+done
+
+echo ">>> [$(date '+%F %T')] hparam sweep complete"
+echo ">>> summary of best eval KLs:"
+for cfg in "${CONFIGS[@]}"; do
+    name="$(basename "$cfg" .toml)"
+    log="$LOG_DIR/$name.log"
+    best=$(grep -E "Best eval KL" "$log" | tail -1 | sed 's/.*Best eval KL = //')
+    printf "    %-32s %s\n" "$name" "${best:-FAILED}"
+done