Add d10 and LFM post-hoc checkpoints

README.md

@@ -0,0 +1,43 @@
+---
+license: other
+license_name: mixed
+tags:
+- interpretability
+- sparsity
+- swiglu
+- top-k
+---
+
+# mask-align checkpoints
+
+Checkpoints from a study of top-K masks on SwiGLU gates and the **V2 similarity-alignment loss**. V2 pushes tokens whose MLP outputs are similar to also share similar sets of active gate units. The aim is to make individual units and groups of units easier to interpret.
+
+All files are raw PyTorch `state_dict`s (`torch.load`, pickle). Load only if you trust this repo.
+
+## `d10/`: nanochat depth-10 models trained from scratch
+
+nanochat (commit `92d63d4`, patched `nanochat/gpt.py` and `scripts/base_train.py`): 10 layers, d_model 640, MLP hidden size 1792, vocab 16384, sequence length 2048, 1402 steps × 524k tokens on ClimbMix. `topk_*` models keep the top 25% of gates per token (`topk_frac 0.25`) and add `λ · (1 − V2)` to the loss. Each folder has `model_001402.pt` and `meta_001402.json` (full config). `d10/tokenizer/` is the nanochat tokenizer.
+
+| Model | λ (V2) | val bpb |
+|---|---|---|
+| swiglu (dense baseline) | – | 0.913 |
+| topk_l0 | 0 | 0.943 |
+| topk_l0p1 | 0.1 | 0.947 |
+| topk_l0p3 | 0.3 | 0.953 |
+| topk_l1 | 1 | 0.972 |
+| topk_l3 | 3 | 0.999 |
+| topk_l10 | 10 | 1.033 |
+
+## `lfm_posthoc/`: post-hoc sparsified LFM2.5-230M-Base
+
+These are fine-tunes of [LiquidAI/LFM2.5-230M-Base](https://huggingface.co/LiquidAI/LFM2.5-230M-Base), so they inherit the **LFM Open License v1.0**. See the base model's [LICENSE](https://huggingface.co/LiquidAI/LFM2.5-230M-Base/blob/main/LICENSE). The weights were modified by the training described below.
+
+In layers 2–13 the MLP keeps the top 25% of units by |gate| (silu·u values are kept; layers 0–1 stay dense). Each model is trained for 20M ClimbMix tokens by KL distillation from the frozen dense model, plus `λ · (1 − V2)`. The V2 target is relu(cosine similarity) of the teacher's batch-centred MLP outputs, and the mask is sign-split. Settings: lr 5e-5, batch 8×1024. bpb is measured on 64 held-out documents: the dense model scores 0.835, and the sparsified model before training scores 1.116.
+
+| Model | λ | bpb | V2 |
+|---|---|---|---|
+| lam0 | 0 | 0.901 | 0.525 |
+| lam10 | 10 | 0.949 | 0.743 |
+| lam30 | 30 | 1.009 | 0.757 |
+
+`student.pt` is a full `LFM2ForCausalLM` state dict in fp32. Load it into `AutoModelForCausalLM.from_pretrained("LiquidAI/LFM2.5-230M-Base")`, then patch the sparse MLPs: see `Sparse` / `load_student` in `lfm_posthoc.py`. Without the patch the model runs dense, which these weights were not trained for. `log.json` is the training log.

d10/swiglu/meta_001402.json

@@ -0,0 +1,67 @@
+{
+  "step": 1402,
+  "val_bpb": 0.9134131077908676,
+  "model_config": {
+    "sequence_len": 2048,
+    "vocab_size": 16384,
+    "n_layer": 10,
+    "n_head": 5,
+    "n_kv_head": 5,
+    "n_embd": 640,
+    "window_pattern": "SSSL",
+    "mlp_type": "swiglu",
+    "mlp_hidden": 1792,
+    "topk_frac": 0.25,
+    "v2_lambda": 0.0,
+    "v2_tokens": 1024
+  },
+  "user_config": {
+    "run": "dummy",
+    "device_type": "",
+    "fp8": false,
+    "fp8_recipe": "tensorwise",
+    "depth": 10,
+    "aspect_ratio": 64,
+    "head_dim": 128,
+    "max_seq_len": 2048,
+    "window_pattern": "SSSL",
+    "mlp_type": "swiglu",
+    "topk_frac": 0.25,
+    "v2_lambda": 0.0,
+    "v2_tokens": 1024,
+    "num_iterations": -1,
+    "target_flops": -1.0,
+    "target_param_data_ratio": 12,
+    "device_batch_size": 32,
+    "total_batch_size": -1,
+    "embedding_lr": 0.3,
+    "unembedding_lr": 0.008,
+    "weight_decay": 0.28,
+    "matrix_lr": 0.02,
+    "scalar_lr": 0.5,
+    "warmup_steps": 40,
+    "warmdown_ratio": 0.65,
+    "final_lr_frac": 0.05,
+    "resume_from_step": -1,
+    "eval_every": 250,
+    "eval_tokens": 41943040,
+    "core_metric_every": -1,
+    "core_metric_max_per_task": 500,
+    "sample_every": 2000,
+    "save_every": -1,
+    "model_tag": "swiglu"
+  },
+  "device_batch_size": 32,
+  "max_seq_len": 2048,
+  "total_batch_size": 524288,
+  "dataloader_state_dict": {
+    "pq_idx": 15,
+    "rg_idx": 53,
+    "epoch": 1
+  },
+  "loop_state": {
+    "min_val_bpb": 0.9134131077908676,
+    "smooth_train_loss": 2.801379575935762,
+    "total_training_time": 1574.414921283722
+  }
+}

d10/swiglu/model_001402.pt

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

d10/tokenizer/token_bytes.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9b4a42091962c62c7bbd110da0468a6748de5d9fa2fbda1567abde88790cd1bb
+size 67113

d10/tokenizer/tokenizer.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1529a0969445dc411f8f395e65e541047d9492d43804e0d56343b25752eab424
+size 198922

d10/topk_l0/meta_001402.json

@@ -0,0 +1,67 @@
+{
+  "step": 1402,
+  "val_bpb": 0.9427678492299922,
+  "model_config": {
+    "sequence_len": 2048,
+    "vocab_size": 16384,
+    "n_layer": 10,
+    "n_head": 5,
+    "n_kv_head": 5,
+    "n_embd": 640,
+    "window_pattern": "SSSL",
+    "mlp_type": "topk",
+    "mlp_hidden": 1792,
+    "topk_frac": 0.25,
+    "v2_lambda": 0.0,
+    "v2_tokens": 1024
+  },
+  "user_config": {
+    "run": "dummy",
+    "device_type": "",
+    "fp8": false,
+    "fp8_recipe": "tensorwise",
+    "depth": 10,
+    "aspect_ratio": 64,
+    "head_dim": 128,
+    "max_seq_len": 2048,
+    "window_pattern": "SSSL",
+    "mlp_type": "topk",
+    "topk_frac": 0.25,
+    "v2_lambda": 0.0,
+    "v2_tokens": 1024,
+    "num_iterations": -1,
+    "target_flops": -1.0,
+    "target_param_data_ratio": 12,
+    "device_batch_size": 32,
+    "total_batch_size": -1,
+    "embedding_lr": 0.3,
+    "unembedding_lr": 0.008,
+    "weight_decay": 0.28,
+    "matrix_lr": 0.02,
+    "scalar_lr": 0.5,
+    "warmup_steps": 40,
+    "warmdown_ratio": 0.65,
+    "final_lr_frac": 0.05,
+    "resume_from_step": -1,
+    "eval_every": 250,
+    "eval_tokens": 41943040,
+    "core_metric_every": -1,
+    "core_metric_max_per_task": 500,
+    "sample_every": 2000,
+    "save_every": -1,
+    "model_tag": "topk_l0"
+  },
+  "device_batch_size": 32,
+  "max_seq_len": 2048,
+  "total_batch_size": 524288,
+  "dataloader_state_dict": {
+    "pq_idx": 15,
+    "rg_idx": 53,
+    "epoch": 1
+  },
+  "loop_state": {
+    "min_val_bpb": 0.9427678492299922,
+    "smooth_train_loss": 2.8899907907043128,
+    "total_training_time": 1948.8127439022064
+  }
+}

d10/topk_l0/model_001402.pt

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

d10/topk_l0p1/meta_001402.json

@@ -0,0 +1,67 @@
+{
+  "step": 1402,
+  "val_bpb": 0.9466192382981357,
+  "model_config": {
+    "sequence_len": 2048,
+    "vocab_size": 16384,
+    "n_layer": 10,
+    "n_head": 5,
+    "n_kv_head": 5,
+    "n_embd": 640,
+    "window_pattern": "SSSL",
+    "mlp_type": "topk",
+    "mlp_hidden": 1792,
+    "topk_frac": 0.25,
+    "v2_lambda": 0.1,
+    "v2_tokens": 1024
+  },
+  "user_config": {
+    "run": "dummy",
+    "device_type": "",
+    "fp8": false,
+    "fp8_recipe": "tensorwise",
+    "depth": 10,
+    "aspect_ratio": 64,
+    "head_dim": 128,
+    "max_seq_len": 2048,
+    "window_pattern": "SSSL",
+    "mlp_type": "topk",
+    "topk_frac": 0.25,
+    "v2_lambda": 0.1,
+    "v2_tokens": 1024,
+    "num_iterations": -1,
+    "target_flops": -1.0,
+    "target_param_data_ratio": 12,
+    "device_batch_size": 32,
+    "total_batch_size": -1,
+    "embedding_lr": 0.3,
+    "unembedding_lr": 0.008,
+    "weight_decay": 0.28,
+    "matrix_lr": 0.02,
+    "scalar_lr": 0.5,
+    "warmup_steps": 40,
+    "warmdown_ratio": 0.65,
+    "final_lr_frac": 0.05,
+    "resume_from_step": -1,
+    "eval_every": 250,
+    "eval_tokens": 41943040,
+    "core_metric_every": -1,
+    "core_metric_max_per_task": 500,
+    "sample_every": 2000,
+    "save_every": -1,
+    "model_tag": "topk_l0p1"
+  },
+  "device_batch_size": 32,
+  "max_seq_len": 2048,
+  "total_batch_size": 524288,
+  "dataloader_state_dict": {
+    "pq_idx": 15,
+    "rg_idx": 53,
+    "epoch": 1
+  },
+  "loop_state": {
+    "min_val_bpb": 0.9466192382981357,
+    "smooth_train_loss": 2.9303368415033866,
+    "total_training_time": 1977.5749444961548
+  }
+}

d10/topk_l0p1/model_001402.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:482cc4d399298959a50b6b924f79c6a4471e36c872ae23170b23ef7c9819db44
+size 370969391

d10/topk_l0p3/meta_001402.json

@@ -0,0 +1,67 @@
+{
+  "step": 1402,
+  "val_bpb": 0.9532622899084581,
+  "model_config": {
+    "sequence_len": 2048,
+    "vocab_size": 16384,
+    "n_layer": 10,
+    "n_head": 5,
+    "n_kv_head": 5,
+    "n_embd": 640,
+    "window_pattern": "SSSL",
+    "mlp_type": "topk",
+    "mlp_hidden": 1792,
+    "topk_frac": 0.25,
+    "v2_lambda": 0.3,
+    "v2_tokens": 1024
+  },
+  "user_config": {
+    "run": "dummy",
+    "device_type": "",
+    "fp8": false,
+    "fp8_recipe": "tensorwise",
+    "depth": 10,
+    "aspect_ratio": 64,
+    "head_dim": 128,
+    "max_seq_len": 2048,
+    "window_pattern": "SSSL",
+    "mlp_type": "topk",
+    "topk_frac": 0.25,
+    "v2_lambda": 0.3,
+    "v2_tokens": 1024,
+    "num_iterations": -1,
+    "target_flops": -1.0,
+    "target_param_data_ratio": 12,
+    "device_batch_size": 32,
+    "total_batch_size": -1,
+    "embedding_lr": 0.3,
+    "unembedding_lr": 0.008,
+    "weight_decay": 0.28,
+    "matrix_lr": 0.02,
+    "scalar_lr": 0.5,
+    "warmup_steps": 40,
+    "warmdown_ratio": 0.65,
+    "final_lr_frac": 0.05,
+    "resume_from_step": -1,
+    "eval_every": 250,
+    "eval_tokens": 41943040,
+    "core_metric_every": -1,
+    "core_metric_max_per_task": 500,
+    "sample_every": 2000,
+    "save_every": -1,
+    "model_tag": "topk_l0p3"
+  },
+  "device_batch_size": 32,
+  "max_seq_len": 2048,
+  "total_batch_size": 524288,
+  "dataloader_state_dict": {
+    "pq_idx": 15,
+    "rg_idx": 53,
+    "epoch": 1
+  },
+  "loop_state": {
+    "min_val_bpb": 0.9532622899084581,
+    "smooth_train_loss": 2.996994468281134,
+    "total_training_time": 1978.3972182273865
+  }
+}

d10/topk_l0p3/model_001402.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7eb2c11b79176e6cbdc94d99fc50141062e1621194af071e56ebf223a9ce5302
+size 370969391

d10/topk_l1/meta_001402.json

@@ -0,0 +1,67 @@
+{
+  "step": 1402,
+  "val_bpb": 0.9721558720552959,
+  "model_config": {
+    "sequence_len": 2048,
+    "vocab_size": 16384,
+    "n_layer": 10,
+    "n_head": 5,
+    "n_kv_head": 5,
+    "n_embd": 640,
+    "window_pattern": "SSSL",
+    "mlp_type": "topk",
+    "mlp_hidden": 1792,
+    "topk_frac": 0.25,
+    "v2_lambda": 1.0,
+    "v2_tokens": 1024
+  },
+  "user_config": {
+    "run": "dummy",
+    "device_type": "",
+    "fp8": false,
+    "fp8_recipe": "tensorwise",
+    "depth": 10,
+    "aspect_ratio": 64,
+    "head_dim": 128,
+    "max_seq_len": 2048,
+    "window_pattern": "SSSL",
+    "mlp_type": "topk",
+    "topk_frac": 0.25,
+    "v2_lambda": 1.0,
+    "v2_tokens": 1024,
+    "num_iterations": -1,
+    "target_flops": -1.0,
+    "target_param_data_ratio": 12,
+    "device_batch_size": 32,
+    "total_batch_size": -1,
+    "embedding_lr": 0.3,
+    "unembedding_lr": 0.008,
+    "weight_decay": 0.28,
+    "matrix_lr": 0.02,
+    "scalar_lr": 0.5,
+    "warmup_steps": 40,
+    "warmdown_ratio": 0.65,
+    "final_lr_frac": 0.05,
+    "resume_from_step": -1,
+    "eval_every": 250,
+    "eval_tokens": 41943040,
+    "core_metric_every": -1,
+    "core_metric_max_per_task": 500,
+    "sample_every": 2000,
+    "save_every": -1,
+    "model_tag": "topk_l1"
+  },
+  "device_batch_size": 32,
+  "max_seq_len": 2048,
+  "total_batch_size": 524288,
+  "dataloader_state_dict": {
+    "pq_idx": 15,
+    "rg_idx": 53,
+    "epoch": 1
+  },
+  "loop_state": {
+    "min_val_bpb": 0.9721558720552959,
+    "smooth_train_loss": 3.1739991677565245,
+    "total_training_time": 1976.9823906421661
+  }
+}

d10/topk_l1/model_001402.pt

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

d10/topk_l10/meta_001402.json

@@ -0,0 +1,67 @@
+{
+  "step": 1402,
+  "val_bpb": 1.0326348566934875,
+  "model_config": {
+    "sequence_len": 2048,
+    "vocab_size": 16384,
+    "n_layer": 10,
+    "n_head": 5,
+    "n_kv_head": 5,
+    "n_embd": 640,
+    "window_pattern": "SSSL",
+    "mlp_type": "topk",
+    "mlp_hidden": 1792,
+    "topk_frac": 0.25,
+    "v2_lambda": 10.0,
+    "v2_tokens": 1024
+  },
+  "user_config": {
+    "run": "dummy",
+    "device_type": "",
+    "fp8": false,
+    "fp8_recipe": "tensorwise",
+    "depth": 10,
+    "aspect_ratio": 64,
+    "head_dim": 128,
+    "max_seq_len": 2048,
+    "window_pattern": "SSSL",
+    "mlp_type": "topk",
+    "topk_frac": 0.25,
+    "v2_lambda": 10.0,
+    "v2_tokens": 1024,
+    "num_iterations": -1,
+    "target_flops": -1.0,
+    "target_param_data_ratio": 12,
+    "device_batch_size": 32,
+    "total_batch_size": -1,
+    "embedding_lr": 0.3,
+    "unembedding_lr": 0.008,
+    "weight_decay": 0.28,
+    "matrix_lr": 0.02,
+    "scalar_lr": 0.5,
+    "warmup_steps": 40,
+    "warmdown_ratio": 0.65,
+    "final_lr_frac": 0.05,
+    "resume_from_step": -1,
+    "eval_every": 250,
+    "eval_tokens": 41943040,
+    "core_metric_every": -1,
+    "core_metric_max_per_task": 500,
+    "sample_every": 2000,
+    "save_every": -1,
+    "model_tag": "topk_l10"
+  },
+  "device_batch_size": 32,
+  "max_seq_len": 2048,
+  "total_batch_size": 524288,
+  "dataloader_state_dict": {
+    "pq_idx": 15,
+    "rg_idx": 53,
+    "epoch": 1
+  },
+  "loop_state": {
+    "min_val_bpb": 1.0326348566934875,
+    "smooth_train_loss": 3.700921912464884,
+    "total_training_time": 2103.8987390995026
+  }
+}

d10/topk_l10/model_001402.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:093e902b3570127dcfeb0903aeaf36921b245a5dfcf5aad7272b3db2dc040e7f
+size 370969391

d10/topk_l3/meta_001402.json

@@ -0,0 +1,67 @@
+{
+  "step": 1402,
+  "val_bpb": 0.9986484763751647,
+  "model_config": {
+    "sequence_len": 2048,
+    "vocab_size": 16384,
+    "n_layer": 10,
+    "n_head": 5,
+    "n_kv_head": 5,
+    "n_embd": 640,
+    "window_pattern": "SSSL",
+    "mlp_type": "topk",
+    "mlp_hidden": 1792,
+    "topk_frac": 0.25,
+    "v2_lambda": 3.0,
+    "v2_tokens": 1024
+  },
+  "user_config": {
+    "run": "dummy",
+    "device_type": "",
+    "fp8": false,
+    "fp8_recipe": "tensorwise",
+    "depth": 10,
+    "aspect_ratio": 64,
+    "head_dim": 128,
+    "max_seq_len": 2048,
+    "window_pattern": "SSSL",
+    "mlp_type": "topk",
+    "topk_frac": 0.25,
+    "v2_lambda": 3.0,
+    "v2_tokens": 1024,
+    "num_iterations": -1,
+    "target_flops": -1.0,
+    "target_param_data_ratio": 12,
+    "device_batch_size": 32,
+    "total_batch_size": -1,
+    "embedding_lr": 0.3,
+    "unembedding_lr": 0.008,
+    "weight_decay": 0.28,
+    "matrix_lr": 0.02,
+    "scalar_lr": 0.5,
+    "warmup_steps": 40,
+    "warmdown_ratio": 0.65,
+    "final_lr_frac": 0.05,
+    "resume_from_step": -1,
+    "eval_every": 250,
+    "eval_tokens": 41943040,
+    "core_metric_every": -1,
+    "core_metric_max_per_task": 500,
+    "sample_every": 2000,
+    "save_every": -1,
+    "model_tag": "topk_l3"
+  },
+  "device_batch_size": 32,
+  "max_seq_len": 2048,
+  "total_batch_size": 524288,
+  "dataloader_state_dict": {
+    "pq_idx": 15,
+    "rg_idx": 53,
+    "epoch": 1
+  },
+  "loop_state": {
+    "min_val_bpb": 0.9986484763751647,
+    "smooth_train_loss": 3.2499740512602906,
+    "total_training_time": 1977.9597957134247
+  }
+}

d10/topk_l3/model_001402.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5261e371ce30bc79d2acb38f897218e24983db1d7dd3b2cc06cf167089060f08
+size 370969391

lfm_posthoc.py

@@ -0,0 +1,140 @@
+"""Post-hoc V2 mask alignment on LFM2.5-230M-Base (local MPS pilot). Student = LFM with sparse MLPs, trained by
+distillation from the frozen original (teacher).
+Student MLP: g = w1 x, u = w3 x; mask = top-K of |g| (K = FRAC of the hidden size, fixed for the whole run);
+h = mask * silu(g) * u (values kept, not binarised). Layers in --dense stay dense.
+V2 loss (as nanochat_patch/gpt.py v2_score, sign-split): units = [+ halves | - halves] of the neurons; forward mask =
+hard signed mask, backward through sigmoid(+g) / sigmoid(-g) (straight-through). Target T = relu(cos of batch-centred
+MLP outputs of the *teacher* at the same layer and positions), so T cannot collapse. Same V2_TOKENS random positions for
+every layer per micro-batch (--micro rows; --bs rows per optimizer step via gradient accumulation). Loss = KL(teacher || student) per token + lam * mean over sparse layers of (1 - v2).
+Data: documents from a ClimbMix parquet shard (--data), BOS-aligned rows: every row starts with BOS and a new document.
+Eval: bpb on lfm_topk_probe.texts(EVAL_DOCS) (held-out shard; disjoint from training docs when --data is a train shard).
+  python lfm_posthoc.py --name lam0p1 --lam 0.1 --tokens 20e6 [--data val_text/shard_00000.parquet]
+"""
+import argparse, json, math, os, time
+import numpy as np, pyarrow.parquet as pq, torch, torch.nn.functional as F
+from huggingface_hub import hf_hub_download
+from transformers import AutoModelForCausalLM, PreTrainedTokenizerFast
+import lfm_topk_probe as P
+from contextlib import nullcontext
+
+ap = argparse.ArgumentParser()
+ap.add_argument("--name", required=True); ap.add_argument("--lam", type=float, default=0.0)
+ap.add_argument("--tokens", type=float, default=20e6); ap.add_argument("--seq", type=int, default=1024)
+ap.add_argument("--bs", type=int, default=8); ap.add_argument("--micro", type=int, default=2); ap.add_argument("--lr", type=float, default=5e-5)
+ap.add_argument("--warmup", type=int, default=100); ap.add_argument("--frac", type=float, default=0.25)
+ap.add_argument("--dense", default="0,1"); ap.add_argument("--v2-tokens", type=int, default=1024)
+ap.add_argument("--data", default=os.path.join(P.ROOT, "val_text", "shard_00000.parquet"))
+ap.add_argument("--doc-start", type=int, default=0); ap.add_argument("--eval-every", type=int, default=500)
+ap.add_argument("--eval-docs", type=int, default=64); ap.add_argument("--max-steps", type=int, default=0)
+args = ap.parse_args() if __name__ == "__main__" else None
+dev = P.dev; torch.manual_seed(0)
+amp = (lambda: torch.autocast("cuda", torch.bfloat16)) if dev == "cuda" else nullcontext   # bf16 compute, fp32 weights
+
+
+def v2_score(m, y):
+    """Uncentred cosine of masks vs relu(cosine of batch-centred outputs), off-diagonal pairs (nanochat_patch/gpt.py)."""
+    m, y = m.float(), y.float()
+    Km = F.normalize(m, dim=-1, eps=1e-8); Km = Km @ Km.T
+    Ky = F.normalize(y - y.mean(0, keepdim=True), dim=-1, eps=1e-8); Ky = torch.relu(Ky @ Ky.T)
+    off = 1.0 - torch.eye(m.size(0), device=m.device); Km, Ky = Km * off, Ky * off
+    return (Km * Ky).sum() / (Km.norm() * Ky.norm() + 1e-8)
+
+
+class Sparse:
+    """Patches the student's MLPs. self.idx = sampled flat positions for V2 (None: no V2); self.target[l] = teacher y."""
+    def __init__(self, model, frac, dense):
+        self.idx, self.target, self.v2 = None, {}, {}
+        self.sparse = [l for l in range(len(model.model.layers)) if l not in dense]
+        for l in self.sparse:
+            m = model.model.layers[l].feed_forward; m.forward = self.make(l, m, int(frac * m.w1.out_features))
+
+    def make(self, l, m, K):
+        def fwd(x):
+            g = m.w1(x); u = m.w3(x); s = F.silu(g)
+            hard = torch.zeros_like(g).scatter_(-1, g.abs().topk(K, -1).indices, 1.0)
+            y = m.w2(hard * s * u)
+            if self.idx is not None:
+                gi = g.flatten(0, 1)[self.idx]; hi = hard.flatten(0, 1)[self.idx]
+                hp, hn = hi * (gi > 0), hi * (gi < 0); sp, sn = torch.sigmoid(gi), torch.sigmoid(-gi)
+                ms = torch.cat([hp + sp - sp.detach(), hn + sn - sn.detach()], -1)   # forward hard, backward sigmoid
+                self.v2[l] = v2_score(ms, self.target[l])
+            return y
+        return fwd
+
+
+def rows(tok, path, seq, start):
+    """BOS-aligned rows of seq+1 tokens: each row starts with BOS + a new document; documents are cut at the row end."""
+    texts = pq.read_table(path, columns=["text"]).column("text").to_pylist()[start:]
+    bos = tok.bos_token_id; buf = []
+    for t in texts:
+        ids = tok(t, add_special_tokens=False).input_ids
+        if not buf: buf = [bos]
+        else: buf += [bos]
+        buf += ids
+        if len(buf) >= seq + 1: yield buf[:seq + 1]; buf = []
+
+
+@torch.no_grad()
+def eval_bpb(model, tok, docs, nbytes):
+    tot = 0.0
+    for t in docs:
+        x = torch.cat([torch.tensor([[tok.bos_token_id]]), tok(t, return_tensors="pt", add_special_tokens=False).input_ids], 1).to(dev)
+        with amp(): lg = model(x).logits
+        tot += F.cross_entropy(lg[0, :-1].float(), x[0, 1:], reduction="sum").item()
+    return tot / math.log(2) / nbytes
+
+
+def load_student(name, frac=0.25, dense=(0, 1)):
+    """Trained student from nanochat_runs/lfm_posthoc/<name>/student.pt, sparse MLPs patched (V2 off)."""
+    m = AutoModelForCausalLM.from_pretrained(P.MODEL, dtype=torch.float32)
+    m.load_state_dict(torch.load(os.path.join(P.RUNS, "lfm_posthoc", name, "student.pt"), map_location="cpu"))
+    m = m.to(P.dev).eval(); Sparse(m, frac, set(dense)); return m
+
+
+if __name__ == "__main__":
+    OUT = os.path.join(P.RUNS, "lfm_posthoc", args.name); os.makedirs(OUT, exist_ok=True)
+    tok = PreTrainedTokenizerFast(tokenizer_file=hf_hub_download(P.MODEL, "tokenizer.json"), bos_token="<|startoftext|>")
+    teacher = AutoModelForCausalLM.from_pretrained(P.MODEL, dtype=torch.bfloat16).to(dev).eval().requires_grad_(False)
+    student = AutoModelForCausalLM.from_pretrained(P.MODEL, dtype=torch.float32).to(dev).train()
+    S = Sparse(student, args.frac, {int(x) for x in args.dense.split(",") if x != ""})
+    tgt = {}
+    def thook(l):
+        def f(mod, a, o):
+            if S.idx is not None: tgt[l] = o.flatten(0, 1)[S.idx].detach()
+        return f
+    for l in S.sparse: teacher.model.layers[l].feed_forward.register_forward_hook(thook(l))
+    docs = P.texts(args.eval_docs); nbytes = sum(len(t.encode()) for t in docs)
+    opt = torch.optim.AdamW(student.parameters(), lr=args.lr, betas=(0.9, 0.95), weight_decay=0.0)
+    steps = args.max_steps or int(args.tokens // (args.bs * args.seq))
+    sched = lambda s: min(1, (s + 1) / args.warmup) * 0.5 * (1 + math.cos(math.pi * min(1, s / steps)))
+    log = []; student.eval(); b0 = eval_bpb(student, tok, docs, nbytes); student.train()
+    print(f"{args.name}: {steps} steps x {args.bs * args.seq} tokens, lam {args.lam}, sparse layers {S.sparse}; "
+          f"bpb at start {b0:.4f}", flush=True); log.append(dict(step=0, bpb=b0))
+    it = rows(tok, args.data, args.seq, args.doc_start); t0 = time.time()
+    for step in range(1, steps + 1):
+        for gp in opt.param_groups: gp["lr"] = args.lr * sched(step)
+        opt.zero_grad(set_to_none=True); n_acc = args.bs // args.micro; kl_s = v2_s = 0.0
+        for _ in range(n_acc):                                   # micro-batches: logits are B x T x 65536
+            x = torch.tensor([next(it) for _ in range(args.micro)], device=dev)[:, :-1]
+            S.idx = torch.randint(0, x.numel(), (args.v2_tokens,), device=dev)
+            with torch.no_grad(): tl = F.log_softmax(teacher(x).logits.float(), -1)
+            S.target = tgt
+            with amp(): sl = student(x).logits
+            sl = F.log_softmax(sl.float(), -1)
+            kl = F.kl_div(sl, tl, log_target=True, reduction="none").sum(-1).mean(); del tl
+            v2 = torch.stack([S.v2[l] for l in S.sparse]).mean()
+            loss = kl + args.lam * (1 - v2) if args.lam > 0 else kl
+            (loss / n_acc).backward(); kl_s += kl.item() / n_acc; v2_s += v2.item() / n_acc
+        torch.nn.utils.clip_grad_norm_(student.parameters(), 1.0); opt.step()
+        if step % 25 == 0 or step == 1:
+            el = time.time() - t0
+            print(f"step {step}/{steps} kl {kl_s:.4f} v2 {v2_s:.4f} | {step * args.bs * args.seq / el:.0f} tok/s, "
+                  f"eta {(steps - step) * el / step / 60:.0f} min", flush=True)
+            log.append(dict(step=step, kl=kl_s, v2=v2_s))
+        if step % args.eval_every == 0 or step == steps:
+            S.idx = None; student.eval(); b = eval_bpb(student, tok, docs, nbytes); student.train()
+            print(f"  eval step {step}: bpb {b:.4f}", flush=True); log.append(dict(step=step, bpb=b))
+            json.dump(dict(args=vars(args), log=log), open(os.path.join(OUT, "log.json"), "w"), indent=1)
+    S.idx = None
+    torch.save(student.state_dict(), os.path.join(OUT, "student.pt"))
+    print("saved ->", OUT)

lfm_posthoc/lam0/log.json

@@ -0,0 +1,536 @@
+{
+ "args": {
+  "name": "lam0",
+  "lam": 0.0,
+  "tokens": 20000000.0,
+  "seq": 1024,
+  "bs": 8,
+  "micro": 8,
+  "lr": 5e-05,
+  "warmup": 100,
+  "frac": 0.25,
+  "dense": "0,1",
+  "v2_tokens": 1024,
+  "data": "/workspace/lfm/val_text/shard_00000.parquet",
+  "doc_start": 0,
+  "eval_every": 500,
+  "eval_docs": 64,
+  "max_steps": 0
+ },
+ "log": [
+  {
+   "step": 0,
+   "bpb": 1.1158697211586375
+  },
+  {
+   "step": 1,
+   "kl": 0.9567757844924927,
+   "v2": 0.5121780633926392
+  },
+  {
+   "step": 25,
+   "kl": 0.6885182857513428,
+   "v2": 0.514909029006958
+  },
+  {
+   "step": 50,
+   "kl": 0.58332359790802,
+   "v2": 0.5146559476852417
+  },
+  {
+   "step": 75,
+   "kl": 0.5070444941520691,
+   "v2": 0.525088906288147
+  },
+  {
+   "step": 100,
+   "kl": 0.5131533145904541,
+   "v2": 0.523467481136322
+  },
+  {
+   "step": 125,
+   "kl": 0.5119062662124634,
+   "v2": 0.5204391479492188
+  },
+  {
+   "step": 150,
+   "kl": 0.4638305902481079,
+   "v2": 0.5317404866218567
+  },
+  {
+   "step": 175,
+   "kl": 0.4514150619506836,
+   "v2": 0.5199365615844727
+  },
+  {
+   "step": 200,
+   "kl": 0.4031620919704437,
+   "v2": 0.5212535858154297
+  },
+  {
+   "step": 225,
+   "kl": 0.48501044511795044,
+   "v2": 0.5195940732955933
+  },
+  {
+   "step": 250,
+   "kl": 0.4423227310180664,
+   "v2": 0.5215173959732056
+  },
+  {
+   "step": 275,
+   "kl": 0.3924141526222229,
+   "v2": 0.529947817325592
+  },
+  {
+   "step": 300,
+   "kl": 0.4181095361709595,
+   "v2": 0.5186288356781006
+  },
+  {
+   "step": 325,
+   "kl": 0.4091278314590454,
+   "v2": 0.5205795764923096
+  },
+  {
+   "step": 350,
+   "kl": 0.37230756878852844,
+   "v2": 0.5182966589927673
+  },
+  {
+   "step": 375,
+   "kl": 0.425018310546875,
+   "v2": 0.520012617111206
+  },
+  {
+   "step": 400,
+   "kl": 0.4136752486228943,
+   "v2": 0.5219123959541321
+  },
+  {
+   "step": 425,
+   "kl": 0.4318866729736328,
+   "v2": 0.5299752354621887
+  },
+  {
+   "step": 450,
+   "kl": 0.3470580577850342,
+   "v2": 0.5208612680435181
+  },
+  {
+   "step": 475,
+   "kl": 0.3659360408782959,
+   "v2": 0.52242112159729
+  },
+  {
+   "step": 500,
+   "kl": 0.38594430685043335,
+   "v2": 0.5255236625671387
+  },
+  {
+   "step": 500,
+   "bpb": 0.9362901667566185
+  },
+  {
+   "step": 525,
+   "kl": 0.36414384841918945,
+   "v2": 0.5246122479438782
+  },
+  {
+   "step": 550,
+   "kl": 0.3833131790161133,
+   "v2": 0.5197753310203552
+  },
+  {
+   "step": 575,
+   "kl": 0.36855340003967285,
+   "v2": 0.521148145198822
+  },
+  {
+   "step": 600,
+   "kl": 0.3700331747531891,
+   "v2": 0.5210527777671814
+  },
+  {
+   "step": 625,
+   "kl": 0.37408578395843506,
+   "v2": 0.5178028345108032
+  },
+  {
+   "step": 650,
+   "kl": 0.3581323027610779,
+   "v2": 0.5182215571403503
+  },
+  {
+   "step": 675,
+   "kl": 0.36711782217025757,
+   "v2": 0.5178583860397339
+  },
+  {
+   "step": 700,
+   "kl": 0.3608144223690033,
+   "v2": 0.5191864967346191
+  },
+  {
+   "step": 725,
+   "kl": 0.3568350374698639,
+   "v2": 0.5205861330032349
+  },
+  {
+   "step": 750,
+   "kl": 0.3533484935760498,
+   "v2": 0.5241600275039673
+  },
+  {
+   "step": 775,
+   "kl": 0.31049370765686035,
+   "v2": 0.5191895961761475
+  },
+  {
+   "step": 800,
+   "kl": 0.3350827693939209,
+   "v2": 0.5257576704025269
+  },
+  {
+   "step": 825,
+   "kl": 0.31789565086364746,
+   "v2": 0.5225245356559753
+  },
+  {
+   "step": 850,
+   "kl": 0.34591662883758545,
+   "v2": 0.5281481742858887
+  },
+  {
+   "step": 875,
+   "kl": 0.33013397455215454,
+   "v2": 0.5243662595748901
+  },
+  {
+   "step": 900,
+   "kl": 0.32055214047431946,
+   "v2": 0.5224398374557495
+  },
+  {
+   "step": 925,
+   "kl": 0.318074107170105,
+   "v2": 0.5229991674423218
+  },
+  {
+   "step": 950,
+   "kl": 0.30721181631088257,
+   "v2": 0.5217205286026001
+  },
+  {
+   "step": 975,
+   "kl": 0.35029369592666626,
+   "v2": 0.5179663896560669
+  },
+  {
+   "step": 1000,
+   "kl": 0.2959432005882263,
+   "v2": 0.526552140712738
+  },
+  {
+   "step": 1000,
+   "bpb": 0.9194461306085423
+  },
+  {
+   "step": 1025,
+   "kl": 0.29746514558792114,
+   "v2": 0.5173307657241821
+  },
+  {
+   "step": 1050,
+   "kl": 0.32962122559547424,
+   "v2": 0.5228357911109924
+  },
+  {
+   "step": 1075,
+   "kl": 0.3026488721370697,
+   "v2": 0.5189902186393738
+  },
+  {
+   "step": 1100,
+   "kl": 0.3141504228115082,
+   "v2": 0.5239812135696411
+  },
+  {
+   "step": 1125,
+   "kl": 0.31310582160949707,
+   "v2": 0.5218887329101562
+  },
+  {
+   "step": 1150,
+   "kl": 0.3075018525123596,
+   "v2": 0.5181629061698914
+  },
+  {
+   "step": 1175,
+   "kl": 0.2963619828224182,
+   "v2": 0.5199472904205322
+  },
+  {
+   "step": 1200,
+   "kl": 0.3094053268432617,
+   "v2": 0.5184849500656128
+  },
+  {
+   "step": 1225,
+   "kl": 0.3153882324695587,
+   "v2": 0.522022008895874
+  },
+  {
+   "step": 1250,
+   "kl": 0.3018542528152466,
+   "v2": 0.5238970518112183
+  },
+  {
+   "step": 1275,
+   "kl": 0.27067339420318604,
+   "v2": 0.5240085124969482
+  },
+  {
+   "step": 1300,
+   "kl": 0.30779898166656494,
+   "v2": 0.5223900079727173
+  },
+  {
+   "step": 1325,
+   "kl": 0.29199105501174927,
+   "v2": 0.5210978984832764
+  },
+  {
+   "step": 1350,
+   "kl": 0.3226560652256012,
+   "v2": 0.5234858989715576
+  },
+  {
+   "step": 1375,
+   "kl": 0.2660519778728485,
+   "v2": 0.5188547968864441
+  },
+  {
+   "step": 1400,
+   "kl": 0.2909112572669983,
+   "v2": 0.5231382846832275
+  },
+  {
+   "step": 1425,
+   "kl": 0.29638445377349854,
+   "v2": 0.5226170420646667
+  },
+  {
+   "step": 1450,
+   "kl": 0.27889877557754517,
+   "v2": 0.5281535387039185
+  },
+  {
+   "step": 1475,
+   "kl": 0.28541916608810425,
+   "v2": 0.5213003754615784
+  },
+  {
+   "step": 1500,
+   "kl": 0.27059319615364075,
+   "v2": 0.5176469683647156
+  },
+  {
+   "step": 1500,
+   "bpb": 0.907340418818327
+  },
+  {
+   "step": 1525,
+   "kl": 0.25855404138565063,
+   "v2": 0.5191320180892944
+  },
+  {
+   "step": 1550,
+   "kl": 0.2625178396701813,
+   "v2": 0.5207079648971558
+  },
+  {
+   "step": 1575,
+   "kl": 0.2738356590270996,
+   "v2": 0.5189085006713867
+  },
+  {
+   "step": 1600,
+   "kl": 0.263283908367157,
+   "v2": 0.5242109894752502
+  },
+  {
+   "step": 1625,
+   "kl": 0.2942907214164734,
+   "v2": 0.515937089920044
+  },
+  {
+   "step": 1650,
+   "kl": 0.26394811272621155,
+   "v2": 0.5181228518486023
+  },
+  {
+   "step": 1675,
+   "kl": 0.2513970136642456,
+   "v2": 0.5181796550750732
+  },
+  {
+   "step": 1700,
+   "kl": 0.27342689037323,
+   "v2": 0.5190655589103699
+  },
+  {
+   "step": 1725,
+   "kl": 0.27730390429496765,
+   "v2": 0.530895471572876
+  },
+  {
+   "step": 1750,
+   "kl": 0.27800506353378296,
+   "v2": 0.5248849987983704
+  },
+  {
+   "step": 1775,
+   "kl": 0.2694709002971649,
+   "v2": 0.5175774097442627
+  },
+  {
+   "step": 1800,
+   "kl": 0.263372540473938,
+   "v2": 0.5212461352348328
+  },
+  {
+   "step": 1825,
+   "kl": 0.24625633656978607,
+   "v2": 0.5317647457122803
+  },
+  {
+   "step": 1850,
+   "kl": 0.2899714708328247,
+   "v2": 0.5262032151222229
+  },
+  {
+   "step": 1875,
+   "kl": 0.2597591280937195,
+   "v2": 0.518336832523346
+  },
+  {
+   "step": 1900,
+   "kl": 0.2807551622390747,
+   "v2": 0.5238334536552429
+  },
+  {
+   "step": 1925,
+   "kl": 0.2533589005470276,
+   "v2": 0.5188561677932739
+  },
+  {
+   "step": 1950,
+   "kl": 0.2727213203907013,
+   "v2": 0.5232084393501282
+  },
+  {
+   "step": 1975,
+   "kl": 0.2687305510044098,
+   "v2": 0.5206012725830078
+  },
+  {
+   "step": 2000,
+   "kl": 0.2753087282180786,
+   "v2": 0.5226956605911255
+  },
+  {
+   "step": 2000,
+   "bpb": 0.9012767352630817
+  },
+  {
+   "step": 2025,
+   "kl": 0.26929429173469543,
+   "v2": 0.5180810689926147
+  },
+  {
+   "step": 2050,
+   "kl": 0.2385600209236145,
+   "v2": 0.5212873220443726
+  },
+  {
+   "step": 2075,
+   "kl": 0.2434336394071579,
+   "v2": 0.5218642950057983
+  },
+  {
+   "step": 2100,
+   "kl": 0.2579026222229004,
+   "v2": 0.5206228494644165
+  },
+  {
+   "step": 2125,
+   "kl": 0.2450314313173294,
+   "v2": 0.5237295031547546
+  },
+  {
+   "step": 2150,
+   "kl": 0.23844994604587555,
+   "v2": 0.5228844881057739
+  },
+  {
+   "step": 2175,
+   "kl": 0.29961922764778137,
+   "v2": 0.5233830809593201
+  },
+  {
+   "step": 2200,
+   "kl": 0.2664281725883484,
+   "v2": 0.5224109888076782
+  },
+  {
+   "step": 2225,
+   "kl": 0.2812555432319641,
+   "v2": 0.5225793719291687
+  },
+  {
+   "step": 2250,
+   "kl": 0.2587583065032959,
+   "v2": 0.5195565223693848
+  },
+  {
+   "step": 2275,
+   "kl": 0.26953667402267456,
+   "v2": 0.5209571123123169
+  },
+  {
+   "step": 2300,
+   "kl": 0.24559755623340607,
+   "v2": 0.5220211744308472
+  },
+  {
+   "step": 2325,
+   "kl": 0.24561521410942078,
+   "v2": 0.5179486870765686
+  },
+  {
+   "step": 2350,
+   "kl": 0.26217955350875854,
+   "v2": 0.5261898040771484
+  },
+  {
+   "step": 2375,
+   "kl": 0.2720721960067749,
+   "v2": 0.5240718126296997
+  },
+  {
+   "step": 2400,
+   "kl": 0.25979429483413696,
+   "v2": 0.525181233882904
+  },
+  {
+   "step": 2425,
+   "kl": 0.27663499116897583,
+   "v2": 0.5254063606262207
+  },
+  {
+   "step": 2441,
+   "bpb": 0.9008672009858338
+  }
+ ]
+}

lfm_posthoc/lam0/student.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7f2b8039fce59aee7ad178642e1c2115d640ff3eea63beee24e5973396cf3c57
+size 918825335

lfm_posthoc/lam10/log.json

@@ -0,0 +1,536 @@
+{
+ "args": {
+  "name": "lam10",
+  "lam": 10.0,
+  "tokens": 20000000.0,
+  "seq": 1024,
+  "bs": 8,
+  "micro": 8,
+  "lr": 5e-05,
+  "warmup": 100,
+  "frac": 0.25,
+  "dense": "0,1",
+  "v2_tokens": 1024,
+  "data": "/workspace/lfm/val_text/shard_00000.parquet",
+  "doc_start": 0,
+  "eval_every": 500,
+  "eval_docs": 64,
+  "max_steps": 0
+ },
+ "log": [
+  {
+   "step": 0,
+   "bpb": 1.1158697211586375
+  },
+  {
+   "step": 1,
+   "kl": 0.9567757844924927,
+   "v2": 0.5121780633926392
+  },
+  {
+   "step": 25,
+   "kl": 0.6805148124694824,
+   "v2": 0.5166142582893372
+  },
+  {
+   "step": 50,
+   "kl": 0.5758242607116699,
+   "v2": 0.5218320488929749
+  },
+  {
+   "step": 75,
+   "kl": 0.5063725709915161,
+   "v2": 0.5415690541267395
+  },
+  {
+   "step": 100,
+   "kl": 0.510666012763977,
+   "v2": 0.5490570068359375
+  },
+  {
+   "step": 125,
+   "kl": 0.5222418308258057,
+   "v2": 0.5583083629608154
+  },
+  {
+   "step": 150,
+   "kl": 0.49504542350769043,
+   "v2": 0.5874360799789429
+  },
+  {
+   "step": 175,
+   "kl": 0.49410390853881836,
+   "v2": 0.5849936008453369
+  },
+  {
+   "step": 200,
+   "kl": 0.4617905616760254,
+   "v2": 0.5984220504760742
+  },
+  {
+   "step": 225,
+   "kl": 0.5620356202125549,
+   "v2": 0.5975736379623413
+  },
+  {
+   "step": 250,
+   "kl": 0.5122577548027039,
+   "v2": 0.6126584410667419
+  },
+  {
+   "step": 275,
+   "kl": 0.4663079082965851,
+   "v2": 0.6313234567642212
+  },
+  {
+   "step": 300,
+   "kl": 0.4986872673034668,
+   "v2": 0.6250436902046204
+  },
+  {
+   "step": 325,
+   "kl": 0.49636000394821167,
+   "v2": 0.6161631345748901
+  },
+  {
+   "step": 350,
+   "kl": 0.46386951208114624,
+   "v2": 0.6398818492889404
+  },
+  {
+   "step": 375,
+   "kl": 0.5316632986068726,
+   "v2": 0.6414067149162292
+  },
+  {
+   "step": 400,
+   "kl": 0.5013736486434937,
+   "v2": 0.637249231338501
+  },
+  {
+   "step": 425,
+   "kl": 0.5416275262832642,
+   "v2": 0.6509120464324951
+  },
+  {
+   "step": 450,
+   "kl": 0.44913533329963684,
+   "v2": 0.6528211832046509
+  },
+  {
+   "step": 475,
+   "kl": 0.4779491722583771,
+   "v2": 0.6567109823226929
+  },
+  {
+   "step": 500,
+   "kl": 0.5010685324668884,
+   "v2": 0.662144660949707
+  },
+  {
+   "step": 500,
+   "bpb": 0.9717341480500024
+  },
+  {
+   "step": 525,
+   "kl": 0.4771064221858978,
+   "v2": 0.6661219596862793
+  },
+  {
+   "step": 550,
+   "kl": 0.5050444602966309,
+   "v2": 0.6608487367630005
+  },
+  {
+   "step": 575,
+   "kl": 0.49432283639907837,
+   "v2": 0.674696147441864
+  },
+  {
+   "step": 600,
+   "kl": 0.4970431923866272,
+   "v2": 0.6805545091629028
+  },
+  {
+   "step": 625,
+   "kl": 0.5078287124633789,
+   "v2": 0.678322970867157
+  },
+  {
+   "step": 650,
+   "kl": 0.4857408106327057,
+   "v2": 0.6887492537498474
+  },
+  {
+   "step": 675,
+   "kl": 0.5085928440093994,
+   "v2": 0.6892139315605164
+  },
+  {
+   "step": 700,
+   "kl": 0.5089031457901001,
+   "v2": 0.6878249049186707
+  },
+  {
+   "step": 725,
+   "kl": 0.4961130619049072,
+   "v2": 0.6971300840377808
+  },
+  {
+   "step": 750,
+   "kl": 0.5042147040367126,
+   "v2": 0.6996113657951355
+  },
+  {
+   "step": 775,
+   "kl": 0.45396888256073,
+   "v2": 0.7082163691520691
+  },
+  {
+   "step": 800,
+   "kl": 0.48240041732788086,
+   "v2": 0.7139182090759277
+  },
+  {
+   "step": 825,
+   "kl": 0.461306095123291,
+   "v2": 0.7144457101821899
+  },
+  {
+   "step": 850,
+   "kl": 0.49483346939086914,
+   "v2": 0.7098020315170288
+  },
+  {
+   "step": 875,
+   "kl": 0.4782881736755371,
+   "v2": 0.7162332534790039
+  },
+  {
+   "step": 900,
+   "kl": 0.4713932275772095,
+   "v2": 0.7105610370635986
+  },
+  {
+   "step": 925,
+   "kl": 0.48280078172683716,
+   "v2": 0.7109979391098022
+  },
+  {
+   "step": 950,
+   "kl": 0.450606107711792,
+   "v2": 0.7168129682540894
+  },
+  {
+   "step": 975,
+   "kl": 0.5144466161727905,
+   "v2": 0.7079391479492188
+  },
+  {
+   "step": 1000,
+   "kl": 0.45295941829681396,
+   "v2": 0.7235746383666992
+  },
+  {
+   "step": 1000,
+   "bpb": 0.9680976517991156
+  },
+  {
+   "step": 1025,
+   "kl": 0.44200605154037476,
+   "v2": 0.7186775207519531
+  },
+  {
+   "step": 1050,
+   "kl": 0.48910245299339294,
+   "v2": 0.7191118597984314
+  },
+  {
+   "step": 1075,
+   "kl": 0.4604094624519348,
+   "v2": 0.7230839133262634
+  },
+  {
+   "step": 1100,
+   "kl": 0.46954652667045593,
+   "v2": 0.7238879203796387
+  },
+  {
+   "step": 1125,
+   "kl": 0.46877187490463257,
+   "v2": 0.7172483205795288
+  },
+  {
+   "step": 1150,
+   "kl": 0.46962401270866394,
+   "v2": 0.7197257876396179
+  },
+  {
+   "step": 1175,
+   "kl": 0.4423319697380066,
+   "v2": 0.7244020700454712
+  },
+  {
+   "step": 1200,
+   "kl": 0.4539686441421509,
+   "v2": 0.7232769727706909
+  },
+  {
+   "step": 1225,
+   "kl": 0.48193496465682983,
+   "v2": 0.7089299559593201
+  },
+  {
+   "step": 1250,
+   "kl": 0.45825693011283875,
+   "v2": 0.7220684885978699
+  },
+  {
+   "step": 1275,
+   "kl": 0.4153714179992676,
+   "v2": 0.7335607409477234
+  },
+  {
+   "step": 1300,
+   "kl": 0.47854793071746826,
+   "v2": 0.7152472734451294
+  },
+  {
+   "step": 1325,
+   "kl": 0.45901647210121155,
+   "v2": 0.7186559438705444
+  },
+  {
+   "step": 1350,
+   "kl": 0.4785187244415283,
+   "v2": 0.7257732152938843
+  },
+  {
+   "step": 1375,
+   "kl": 0.3946749269962311,
+   "v2": 0.7312217950820923
+  },
+  {
+   "step": 1400,
+   "kl": 0.44470900297164917,
+   "v2": 0.7313439846038818
+  },
+  {
+   "step": 1425,
+   "kl": 0.44638675451278687,
+   "v2": 0.7234344482421875
+  },
+  {
+   "step": 1450,
+   "kl": 0.4295235276222229,
+   "v2": 0.733188271522522
+  },
+  {
+   "step": 1475,
+   "kl": 0.43276941776275635,
+   "v2": 0.725287914276123
+  },
+  {
+   "step": 1500,
+   "kl": 0.40881919860839844,
+   "v2": 0.7202348709106445
+  },
+  {
+   "step": 1500,
+   "bpb": 0.9551572504817758
+  },
+  {
+   "step": 1525,
+   "kl": 0.40365666151046753,
+   "v2": 0.7291638255119324
+  },
+  {
+   "step": 1550,
+   "kl": 0.4154655635356903,
+   "v2": 0.7356778979301453
+  },
+  {
+   "step": 1575,
+   "kl": 0.4188691973686218,
+   "v2": 0.7242900729179382
+  },
+  {
+   "step": 1600,
+   "kl": 0.4097726047039032,
+   "v2": 0.7322882413864136
+  },
+  {
+   "step": 1625,
+   "kl": 0.46545079350471497,
+   "v2": 0.7162374258041382
+  },
+  {
+   "step": 1650,
+   "kl": 0.4153555631637573,
+   "v2": 0.7328672409057617
+  },
+  {
+   "step": 1675,
+   "kl": 0.3911184072494507,
+   "v2": 0.7268070578575134
+  },
+  {
+   "step": 1700,
+   "kl": 0.41736191511154175,
+   "v2": 0.7315605878829956
+  },
+  {
+   "step": 1725,
+   "kl": 0.4216347932815552,
+   "v2": 0.7403724193572998
+  },
+  {
+   "step": 1750,
+   "kl": 0.42835384607315063,
+   "v2": 0.7339169979095459
+  },
+  {
+   "step": 1775,
+   "kl": 0.41845932602882385,
+   "v2": 0.7234562635421753
+  },
+  {
+   "step": 1800,
+   "kl": 0.41512590646743774,
+   "v2": 0.7330150008201599
+  },
+  {
+   "step": 1825,
+   "kl": 0.4038649797439575,
+   "v2": 0.7267810702323914
+  },
+  {
+   "step": 1850,
+   "kl": 0.44040095806121826,
+   "v2": 0.7361396551132202
+  },
+  {
+   "step": 1875,
+   "kl": 0.4136882722377777,
+   "v2": 0.7325267791748047
+  },
+  {
+   "step": 1900,
+   "kl": 0.4516754746437073,
+   "v2": 0.7292265295982361
+  },
+  {
+   "step": 1925,
+   "kl": 0.40279334783554077,
+   "v2": 0.7277811765670776
+  },
+  {
+   "step": 1950,
+   "kl": 0.4331173896789551,
+   "v2": 0.7397186160087585
+  },
+  {
+   "step": 1975,
+   "kl": 0.42194291949272156,
+   "v2": 0.7219225764274597
+  },
+  {
+   "step": 2000,
+   "kl": 0.42822903394699097,
+   "v2": 0.7371947169303894
+  },
+  {
+   "step": 2000,
+   "bpb": 0.9496056799436946
+  },
+  {
+   "step": 2025,
+   "kl": 0.41823434829711914,
+   "v2": 0.7225486636161804
+  },
+  {
+   "step": 2050,
+   "kl": 0.3759946823120117,
+   "v2": 0.7401331663131714
+  },
+  {
+   "step": 2075,
+   "kl": 0.3947470188140869,
+   "v2": 0.7413791418075562
+  },
+  {
+   "step": 2100,
+   "kl": 0.40479955077171326,
+   "v2": 0.7358704805374146
+  },
+  {
+   "step": 2125,
+   "kl": 0.3906734585762024,
+   "v2": 0.7366509437561035
+  },
+  {
+   "step": 2150,
+   "kl": 0.3809944987297058,
+   "v2": 0.7408315539360046
+  },
+  {
+   "step": 2175,
+   "kl": 0.46800005435943604,
+   "v2": 0.7267584800720215
+  },
+  {
+   "step": 2200,
+   "kl": 0.443926602602005,
+   "v2": 0.7301518321037292
+  },
+  {
+   "step": 2225,
+   "kl": 0.44637081027030945,
+   "v2": 0.738400936126709
+  },
+  {
+   "step": 2250,
+   "kl": 0.4049396514892578,
+   "v2": 0.7364320158958435
+  },
+  {
+   "step": 2275,
+   "kl": 0.43352243304252625,
+   "v2": 0.7247334718704224
+  },
+  {
+   "step": 2300,
+   "kl": 0.39259201288223267,
+   "v2": 0.7407164573669434
+  },
+  {
+   "step": 2325,
+   "kl": 0.39357244968414307,
+   "v2": 0.7396523356437683
+  },
+  {
+   "step": 2350,
+   "kl": 0.4133729934692383,
+   "v2": 0.7367298007011414
+  },
+  {
+   "step": 2375,
+   "kl": 0.43429800868034363,
+   "v2": 0.7416977882385254
+  },
+  {
+   "step": 2400,
+   "kl": 0.41854673624038696,
+   "v2": 0.7408771514892578
+  },
+  {
+   "step": 2425,
+   "kl": 0.4388637840747833,
+   "v2": 0.7434670925140381
+  },
+  {
+   "step": 2441,
+   "bpb": 0.9488953864726722
+  }
+ ]
+}

lfm_posthoc/lam10/student.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7895c892173d42bcab648865c6463fc4e3fd10df2204de542574162d6cb05470
+size 918825335

lfm_posthoc/lam30/log.json

@@ -0,0 +1,536 @@
+{
+ "args": {
+  "name": "lam30",
+  "lam": 30.0,
+  "tokens": 20000000.0,
+  "seq": 1024,
+  "bs": 8,
+  "micro": 8,
+  "lr": 5e-05,
+  "warmup": 100,
+  "frac": 0.25,
+  "dense": "0,1",
+  "v2_tokens": 1024,
+  "data": "/workspace/lfm/val_text/shard_00000.parquet",
+  "doc_start": 0,
+  "eval_every": 500,
+  "eval_docs": 64,
+  "max_steps": 0
+ },
+ "log": [
+  {
+   "step": 0,
+   "bpb": 1.1158697211586375
+  },
+  {
+   "step": 1,
+   "kl": 0.9567757844924927,
+   "v2": 0.5121780633926392
+  },
+  {
+   "step": 25,
+   "kl": 0.6852031350135803,
+   "v2": 0.519575834274292
+  },
+  {
+   "step": 50,
+   "kl": 0.612939715385437,
+   "v2": 0.536693811416626
+  },
+  {
+   "step": 75,
+   "kl": 0.6628031730651855,
+   "v2": 0.5855995416641235
+  },
+  {
+   "step": 100,
+   "kl": 0.7301273941993713,
+   "v2": 0.6085925102233887
+  },
+  {
+   "step": 125,
+   "kl": 0.7643752694129944,
+   "v2": 0.6244039535522461
+  },
+  {
+   "step": 150,
+   "kl": 0.6954302191734314,
+   "v2": 0.6539231538772583
+  },
+  {
+   "step": 175,
+   "kl": 0.712317168712616,
+   "v2": 0.6509629487991333
+  },
+  {
+   "step": 200,
+   "kl": 0.6909964084625244,
+   "v2": 0.6611611843109131
+  },
+  {
+   "step": 225,
+   "kl": 0.7976746559143066,
+   "v2": 0.6576403975486755
+  },
+  {
+   "step": 250,
+   "kl": 0.7557328939437866,
+   "v2": 0.6770698428153992
+  },
+  {
+   "step": 275,
+   "kl": 0.6881502270698547,
+   "v2": 0.7029839754104614
+  },
+  {
+   "step": 300,
+   "kl": 0.7569065093994141,
+   "v2": 0.7006992101669312
+  },
+  {
+   "step": 325,
+   "kl": 0.765669584274292,
+   "v2": 0.6898188591003418
+  },
+  {
+   "step": 350,
+   "kl": 0.6862740516662598,
+   "v2": 0.7034606337547302
+  },
+  {
+   "step": 375,
+   "kl": 0.7587317824363708,
+   "v2": 0.7113478183746338
+  },
+  {
+   "step": 400,
+   "kl": 0.7582862973213196,
+   "v2": 0.7102171778678894
+  },
+  {
+   "step": 425,
+   "kl": 0.7752246260643005,
+   "v2": 0.7032108306884766
+  },
+  {
+   "step": 450,
+   "kl": 0.6610547304153442,
+   "v2": 0.7193240523338318
+  },
+  {
+   "step": 475,
+   "kl": 0.6940382719039917,
+   "v2": 0.710284948348999
+  },
+  {
+   "step": 500,
+   "kl": 0.723612904548645,
+   "v2": 0.7197579741477966
+  },
+  {
+   "step": 500,
+   "bpb": 1.039170600789881
+  },
+  {
+   "step": 525,
+   "kl": 0.6996949911117554,
+   "v2": 0.7240695953369141
+  },
+  {
+   "step": 550,
+   "kl": 0.7239395380020142,
+   "v2": 0.7148751020431519
+  },
+  {
+   "step": 575,
+   "kl": 0.6885330677032471,
+   "v2": 0.725610613822937
+  },
+  {
+   "step": 600,
+   "kl": 0.6794240474700928,
+   "v2": 0.7286481857299805
+  },
+  {
+   "step": 625,
+   "kl": 0.7052067518234253,
+   "v2": 0.7196173667907715
+  },
+  {
+   "step": 650,
+   "kl": 0.6813843250274658,
+   "v2": 0.737176775932312
+  },
+  {
+   "step": 675,
+   "kl": 0.7123562097549438,
+   "v2": 0.729498565196991
+  },
+  {
+   "step": 700,
+   "kl": 0.6975288391113281,
+   "v2": 0.7263223528862
+  },
+  {
+   "step": 725,
+   "kl": 0.6707978248596191,
+   "v2": 0.7262341380119324
+  },
+  {
+   "step": 750,
+   "kl": 0.681881308555603,
+   "v2": 0.7310070395469666
+  },
+  {
+   "step": 775,
+   "kl": 0.6304110884666443,
+   "v2": 0.73924720287323
+  },
+  {
+   "step": 800,
+   "kl": 0.6538665294647217,
+   "v2": 0.7397634387016296
+  },
+  {
+   "step": 825,
+   "kl": 0.6433805227279663,
+   "v2": 0.7406503558158875
+  },
+  {
+   "step": 850,
+   "kl": 0.686224102973938,
+   "v2": 0.7293993830680847
+  },
+  {
+   "step": 875,
+   "kl": 0.660400390625,
+   "v2": 0.7434613108634949
+  },
+  {
+   "step": 900,
+   "kl": 0.6423088312149048,
+   "v2": 0.7325193285942078
+  },
+  {
+   "step": 925,
+   "kl": 0.6635453701019287,
+   "v2": 0.7307019233703613
+  },
+  {
+   "step": 950,
+   "kl": 0.6423885822296143,
+   "v2": 0.7362762689590454
+  },
+  {
+   "step": 975,
+   "kl": 0.697304368019104,
+   "v2": 0.731419563293457
+  },
+  {
+   "step": 1000,
+   "kl": 0.6290038228034973,
+   "v2": 0.7452073097229004
+  },
+  {
+   "step": 1000,
+   "bpb": 1.0275161617816257
+  },
+  {
+   "step": 1025,
+   "kl": 0.6236750483512878,
+   "v2": 0.7377729415893555
+  },
+  {
+   "step": 1050,
+   "kl": 0.6637212038040161,
+   "v2": 0.7397140860557556
+  },
+  {
+   "step": 1075,
+   "kl": 0.6306295394897461,
+   "v2": 0.7423157095909119
+  },
+  {
+   "step": 1100,
+   "kl": 0.6457698345184326,
+   "v2": 0.7419235706329346
+  },
+  {
+   "step": 1125,
+   "kl": 0.6500687599182129,
+   "v2": 0.7352355718612671
+  },
+  {
+   "step": 1150,
+   "kl": 0.6646006107330322,
+   "v2": 0.731339156627655
+  },
+  {
+   "step": 1175,
+   "kl": 0.6234750747680664,
+   "v2": 0.7429856061935425
+  },
+  {
+   "step": 1200,
+   "kl": 0.6179300546646118,
+   "v2": 0.7386311292648315
+  },
+  {
+   "step": 1225,
+   "kl": 0.6625158786773682,
+   "v2": 0.7238904237747192
+  },
+  {
+   "step": 1250,
+   "kl": 0.6490756273269653,
+   "v2": 0.7379249334335327
+  },
+  {
+   "step": 1275,
+   "kl": 0.588732123374939,
+   "v2": 0.7491040229797363
+  },
+  {
+   "step": 1300,
+   "kl": 0.6630625128746033,
+   "v2": 0.7297951579093933
+  },
+  {
+   "step": 1325,
+   "kl": 0.6420645713806152,
+   "v2": 0.7347722053527832
+  },
+  {
+   "step": 1350,
+   "kl": 0.6723450422286987,
+   "v2": 0.7428981065750122
+  },
+  {
+   "step": 1375,
+   "kl": 0.5604075789451599,
+   "v2": 0.745333194732666
+  },
+  {
+   "step": 1400,
+   "kl": 0.6222208738327026,
+   "v2": 0.7469989657402039
+  },
+  {
+   "step": 1425,
+   "kl": 0.625604510307312,
+   "v2": 0.7394630312919617
+  },
+  {
+   "step": 1450,
+   "kl": 0.6182290315628052,
+   "v2": 0.7485737204551697
+  },
+  {
+   "step": 1475,
+   "kl": 0.6214414834976196,
+   "v2": 0.7396450042724609
+  },
+  {
+   "step": 1500,
+   "kl": 0.5841304063796997,
+   "v2": 0.737494170665741
+  },
+  {
+   "step": 1500,
+   "bpb": 1.0141624442718846
+  },
+  {
+   "step": 1525,
+   "kl": 0.5751022100448608,
+   "v2": 0.7439262270927429
+  },
+  {
+   "step": 1550,
+   "kl": 0.5885058641433716,
+   "v2": 0.7491492033004761
+  },
+  {
+   "step": 1575,
+   "kl": 0.6095647811889648,
+   "v2": 0.7392956614494324
+  },
+  {
+   "step": 1600,
+   "kl": 0.5844666957855225,
+   "v2": 0.7473974227905273
+  },
+  {
+   "step": 1625,
+   "kl": 0.6698254942893982,
+   "v2": 0.7298187017440796
+  },
+  {
+   "step": 1650,
+   "kl": 0.5906268358230591,
+   "v2": 0.7461249232292175
+  },
+  {
+   "step": 1675,
+   "kl": 0.5622061491012573,
+   "v2": 0.7406004667282104
+  },
+  {
+   "step": 1700,
+   "kl": 0.5837856531143188,
+   "v2": 0.7463316917419434
+  },
+  {
+   "step": 1725,
+   "kl": 0.5965714454650879,
+   "v2": 0.7548431158065796
+  },
+  {
+   "step": 1750,
+   "kl": 0.6182277202606201,
+   "v2": 0.7483779788017273
+  },
+  {
+   "step": 1775,
+   "kl": 0.5947036743164062,
+   "v2": 0.7394563555717468
+  },
+  {
+   "step": 1800,
+   "kl": 0.5840776562690735,
+   "v2": 0.7466443777084351
+  },
+  {
+   "step": 1825,
+   "kl": 0.5858625173568726,
+   "v2": 0.7399744987487793
+  },
+  {
+   "step": 1850,
+   "kl": 0.6211203336715698,
+   "v2": 0.7482191324234009
+  },
+  {
+   "step": 1875,
+   "kl": 0.5869969725608826,
+   "v2": 0.74660325050354
+  },
+  {
+   "step": 1900,
+   "kl": 0.6490370035171509,
+   "v2": 0.7442652583122253
+  },
+  {
+   "step": 1925,
+   "kl": 0.569796621799469,
+   "v2": 0.7425372004508972
+  },
+  {
+   "step": 1950,
+   "kl": 0.606709361076355,
+   "v2": 0.7533676028251648
+  },
+  {
+   "step": 1975,
+   "kl": 0.5977227687835693,
+   "v2": 0.7364615797996521
+  },
+  {
+   "step": 2000,
+   "kl": 0.6061283349990845,
+   "v2": 0.7509064674377441
+  },
+  {
+   "step": 2000,
+   "bpb": 1.0094860592225428
+  },
+  {
+   "step": 2025,
+   "kl": 0.5938690304756165,
+   "v2": 0.7389944791793823
+  },
+  {
+   "step": 2050,
+   "kl": 0.5432311296463013,
+   "v2": 0.7529944777488708
+  },
+  {
+   "step": 2075,
+   "kl": 0.5748882293701172,
+   "v2": 0.7557010054588318
+  },
+  {
+   "step": 2100,
+   "kl": 0.5797919034957886,
+   "v2": 0.7504448294639587
+  },
+  {
+   "step": 2125,
+   "kl": 0.570436418056488,
+   "v2": 0.7510671019554138
+  },
+  {
+   "step": 2150,
+   "kl": 0.5468230247497559,
+   "v2": 0.7531040906906128
+  },
+  {
+   "step": 2175,
+   "kl": 0.683279812335968,
+   "v2": 0.7414388656616211
+  },
+  {
+   "step": 2200,
+   "kl": 0.6222787499427795,
+   "v2": 0.7441861033439636
+  },
+  {
+   "step": 2225,
+   "kl": 0.6235721111297607,
+   "v2": 0.7518894076347351
+  },
+  {
+   "step": 2250,
+   "kl": 0.5736346244812012,
+   "v2": 0.7490965723991394
+  },
+  {
+   "step": 2275,
+   "kl": 0.6147838234901428,
+   "v2": 0.7375831604003906
+  },
+  {
+   "step": 2300,
+   "kl": 0.5597199201583862,
+   "v2": 0.7532729506492615
+  },
+  {
+   "step": 2325,
+   "kl": 0.5606181621551514,
+   "v2": 0.7539691925048828
+  },
+  {
+   "step": 2350,
+   "kl": 0.5969135761260986,
+   "v2": 0.7491833567619324
+  },
+  {
+   "step": 2375,
+   "kl": 0.622349739074707,
+   "v2": 0.7543088793754578
+  },
+  {
+   "step": 2400,
+   "kl": 0.6025652289390564,
+   "v2": 0.7545272707939148
+  },
+  {
+   "step": 2425,
+   "kl": 0.6382248401641846,
+   "v2": 0.7574693560600281
+  },
+  {
+   "step": 2441,
+   "bpb": 1.008766128718237
+  }
+ ]
+}

lfm_posthoc/lam30/student.pt

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