Archive modified and untracked RCP SparseAtt working-tree files

epfl_final/SparseAtt/$script

@@ -0,0 +1,9 @@
+#!/usr/bin/env bash
+set -euo pipefail
+seed71_log=/mloscratch/homes/navasard/sparse/train_logs/p2_1p2b_mixed_seed71_4gpu_resume45250.log
+seed72_log=/mloscratch/homes/navasard/sparse/train_logs/p2_1p2b_mixed_seed72_4gpu_resume31000_after_seed71.log
+repo=/mloscratch/homes/navasard/sparse/SparseAtt
+while ! grep -q "Training complete" "$seed71_log"; do sleep 60; done
+if pgrep -f "torchrun.*p2_1p2b_mixed_seed72_4gpu_autoresume_current.json" >/dev/null; then exit 0; fi
+cd "$repo"
+nohup env USER=navasard LOGNAME=navasard HOME=/mloscratch/homes/navasard TORCHINDUCTOR_CACHE_DIR=/mloscratch/homes/navasard/.cache/torchinductor /mloscratch/homes/navasard/conda/bin/torchrun --standalone --nproc_per_node=4 train.py --config configs/p2_1p2b_mixed_seed72_4gpu_autoresume_current.json >> "$seed72_log" 2>&1 < /dev/null &

epfl_final/SparseAtt/closure24_tau1_eval.py

@@ -0,0 +1,60 @@
+import csv, math, os, sys
+from pathlib import Path
+import torch
+
+REPO = Path('/mloscratch/homes/navasard/sparse/SparseAtt')
+BASE = Path('/mloscratch/homes/navasard/sparse')
+DATA = BASE.parent / 'data'
+sys.path.insert(0, str(REPO))
+from eval.checkpoint import parse_dtype, resolve_device
+from eval.sparse_eval import load_model_for_deploy
+from analysis.r2_calibration_probe import apply_temperature
+from sparseatt import load_tokenizer
+from datasets import load_from_disk
+
+L = 24576
+CK = {
+ ('dense',70): 'checkpoints/l8k1p2b_dense_seed70/run_20260823_043810/checkpoint-50462',
+ ('mixed',70): 'checkpoints/l8k1p2b_mixed_seed70_repaired/run_20260909_054557/checkpoint-50462',
+ ('dense',72): 'checkpoints/l8k1p2b_dense_seed72_repaired/run_20260906_repair_from9250/checkpoint-50462',
+ ('mixed',72): 'checkpoints/l8k1p2b_mixed_seed72/run_20260827_071551/checkpoint-50462',
+}
+CORP = {'fineweb': DATA/'fineweb_val_seed42', 'pg19': DATA/'pg19_ood_windows', 'arxiv': DATA/'arxiv_ood_windows'}
+OUT = REPO/'eval/results/RESOLUTION_20260921/tau1_fix/closure24_rows.tsv'
+
+def stream(path):
+    ds = load_from_disk(str(path)); ds = ds['train'] if hasattr(ds,'keys') and 'train' in ds else ds
+    ids=[]
+    for x in ds: ids.extend(x['input_ids'])
+    return torch.tensor(ids)
+
+@torch.no_grad()
+def score(m, ids, dev, dt):
+    x=ids[None].to(dev)
+    with torch.autocast('cuda', dtype=dt): y=m(input_ids=x, attention_mask=None)['logits'][:,:-1]
+    t=x[:,1:]; total=0.0; n=0
+    for i in range(0,t.shape[1],512):
+        z=t[:,i:i+512]; q=y[:,i:i+512].float()
+        total += float(torch.nn.functional.cross_entropy(q.reshape(-1,q.shape[-1]), z.reshape(-1), reduction='sum')); n += z.numel()
+    return total/n, n
+
+def main():
+    os.environ.setdefault('USER','navasard'); os.environ.setdefault('HOME','/mloscratch/homes/navasard')
+    dev=resolve_device(None); dt=parse_dtype('bf16'); tok=load_tokenizer('/mloscratch/homes/navasard/qwen3_06b_base')
+    streams={c:stream(p) for c,p in CORP.items()}
+    rows=[]
+    for (arm,seed), rel in CK.items():
+        print('LOAD',arm,seed,flush=True)
+        m=load_model_for_deploy(checkpoint_dir=str(BASE/rel), deploy_sparsity='dense', eval_len=L, block_size=128, pattern='random', force_sparse_layers='all', mask_pool_size=1, mask_seed=42, h2o_recent=128, device=dev, dtype=dt)
+        apply_temperature(m,1.0)
+        for c,ids in streams.items():
+            for j,o in enumerate((2*L,3*L)):
+                nll,n=score(m,ids[o:o+L],dev,dt)
+                rows.append(dict(arm=arm,seed=seed,length=L,split='heldout_existing',corpus=c,window_id=f'closure{j}',start_offset=o,tokens=n,tau=1.0,avg_nll=nll,perplexity=math.exp(nll)))
+                print(arm,seed,c,j,nll,flush=True)
+        del m; torch.cuda.empty_cache()
+    OUT.parent.mkdir(parents=True,exist_ok=True)
+    with OUT.open('w',newline='') as f:
+        w=csv.DictWriter(f,fieldnames=rows[0],delimiter='\t'); w.writeheader(); w.writerows(rows)
+    print('WROTE',OUT,len(rows),flush=True)
+if __name__=='__main__': main()

epfl_final/SparseAtt/concat_tau1_fix.py

@@ -0,0 +1,25 @@
+import csv
+from pathlib import Path
+
+ROOT = Path('/mloscratch/homes/navasard/sparse/SparseAtt/eval/results/RESOLUTION_20260921/tau1_fix')
+OUT = ROOT / 'SPARSEATT_RESOLUTION_TAU1_FIX_20260921_ALL.tsv'
+
+def main():
+    files = sorted(ROOT.glob('*_heldout.tsv'))
+    if len(files) != 24:
+        raise SystemExit(f'expected 24 result TSVs, found {len(files)}')
+    with files[0].open(newline='') as f:
+        fields = list(csv.DictReader(f, delimiter='\t').fieldnames or [])
+    with OUT.open('w', newline='') as f:
+        w = csv.DictWriter(f, fieldnames=['source_file'] + fields, delimiter='\t')
+        w.writeheader()
+        total = 0
+        for p in files:
+            with p.open(newline='') as h:
+                for row in csv.DictReader(h, delimiter='\t'):
+                    w.writerow({'source_file': p.name, **row})
+                    total += 1
+    print(f'wrote {OUT} from {len(files)} files and {total} rows')
+
+if __name__ == '__main__':
+    main()

epfl_final/SparseAtt/configs/FINAL_20260924/365m_dense_qkgainfp32_s42.json

@@ -0,0 +1,8 @@
+{
+  "model": {"path": "../../qwen3_06b_base", "dtype": "bfloat16", "init_from_scratch": true, "hidden_size": 768, "intermediate_size": 3072, "num_hidden_layers": 28, "num_attention_heads": 6, "num_key_value_heads": 2, "rope_theta": 1000000.0, "qk_norm": "on", "qk_gain_fp32": true, "positional": "rope"},
+  "data": {"train_path": "../../data/fineweb_large/tokenized", "val_path": null, "seed": 42, "max_seq_length": 8192, "num_workers": 8},
+  "training": {"output_dir": "models/FINAL_20260924/part_a/365m_dense_qkgainfp32_s42", "resume_from_checkpoint": null, "force_new_run": false, "num_epochs": 1, "max_steps": 44000, "per_device_batch_size": 1, "gradient_accumulation_steps": 20, "learning_rate": 0.0005, "weight_decay": 0.01, "warmup_steps": 750, "lr_scheduler": "cosine", "max_grad_norm": 1.0, "save_steps": 1000, "eval_steps": 250, "eval_batches": 50, "logging_steps": 100, "seed": 42, "verbose": false, "checkpoint_activations": false},
+  "distributed": {"strategy": "ddp", "mixed_precision": "bf16"},
+  "wandb": {"enabled": false, "project": "stochastic_sparsity", "name": "FINAL partA 365m dense qkgainfp32 s42", "tags": ["final_20260924", "part_a", "dense", "qkgainfp32", "seed42"]},
+  "attention": {"type": "dense", "dense_mode": false}
+}

epfl_final/SparseAtt/configs/FINAL_20260924/365m_dense_qkgainfp32_s42_2gpu_latest_resume.json

@@ -0,0 +1,8 @@
+{
+  "model": {"path": "../../qwen3_06b_base", "dtype": "bfloat16", "init_from_scratch": true, "hidden_size": 768, "intermediate_size": 3072, "num_hidden_layers": 28, "num_attention_heads": 6, "num_key_value_heads": 2, "rope_theta": 1000000.0, "qk_norm": "on", "qk_gain_fp32": true, "positional": "rope"},
+  "data": {"train_path": "../../data/fineweb_large/tokenized", "val_path": null, "seed": 42, "max_seq_length": 8192, "num_workers": 8},
+  "training": {"output_dir": "models/FINAL_20260924/part_a/365m_dense_qkgainfp32_s42", "resume_from_checkpoint": "models/FINAL_20260924/part_a/365m_dense_qkgainfp32_s42/run_20260924_045342/checkpoint-26000", "force_new_run": false, "num_epochs": 1, "max_steps": 44000, "per_device_batch_size": 2, "gradient_accumulation_steps": 5, "resume_micro_step_override": 130000, "learning_rate": 0.0005, "weight_decay": 0.01, "warmup_steps": 750, "lr_scheduler": "cosine", "max_grad_norm": 1.0, "save_steps": 1000, "eval_steps": 250, "eval_batches": 50, "logging_steps": 100, "seed": 42, "verbose": false, "checkpoint_activations": false},
+  "distributed": {"strategy": "ddp", "mixed_precision": "bf16"},
+  "wandb": {"enabled": false, "project": "stochastic_sparsity", "name": "FINAL partA 365m dense qkgainfp32 s42", "tags": ["final_20260924", "part_a", "dense", "qkgainfp32", "seed42", "latest_resume"]},
+  "attention": {"type": "dense", "dense_mode": false}
+}

epfl_final/SparseAtt/configs/FINAL_20260924/365m_dense_qkgainfp32_s42_2gpu_restart_resume.json

@@ -0,0 +1,8 @@
+{
+  "model": {"path": "../../qwen3_06b_base", "dtype": "bfloat16", "init_from_scratch": true, "hidden_size": 768, "intermediate_size": 3072, "num_hidden_layers": 28, "num_attention_heads": 6, "num_key_value_heads": 2, "rope_theta": 1000000.0, "qk_norm": "on", "qk_gain_fp32": true, "positional": "rope"},
+  "data": {"train_path": "../../data/fineweb_large/tokenized", "val_path": null, "seed": 42, "max_seq_length": 8192, "num_workers": 8},
+  "training": {"output_dir": "models/FINAL_20260924/part_a/365m_dense_qkgainfp32_s42", "resume_from_checkpoint": "models/FINAL_20260924/part_a/365m_dense_qkgainfp32_s42/run_20260924_045342/checkpoint-9000", "force_new_run": false, "num_epochs": 1, "max_steps": 44000, "per_device_batch_size": 2, "gradient_accumulation_steps": 5, "resume_micro_step_override": 45000, "learning_rate": 0.0005, "weight_decay": 0.01, "warmup_steps": 750, "lr_scheduler": "cosine", "max_grad_norm": 1.0, "save_steps": 1000, "eval_steps": 250, "eval_batches": 50, "logging_steps": 100, "seed": 42, "verbose": false, "checkpoint_activations": false},
+  "distributed": {"strategy": "ddp", "mixed_precision": "bf16"},
+  "wandb": {"enabled": false, "project": "stochastic_sparsity", "name": "FINAL partA 365m dense qkgainfp32 s42", "tags": ["final_20260924", "part_a", "dense", "qkgainfp32", "seed42"]},
+  "attention": {"type": "dense", "dense_mode": false}
+}

epfl_final/SparseAtt/configs/FINAL_20260924/365m_dense_qkgainfp32_s42_2gpu_resume.json

@@ -0,0 +1,8 @@
+{
+  "model": {"path": "../../qwen3_06b_base", "dtype": "bfloat16", "init_from_scratch": true, "hidden_size": 768, "intermediate_size": 3072, "num_hidden_layers": 28, "num_attention_heads": 6, "num_key_value_heads": 2, "rope_theta": 1000000.0, "qk_norm": "on", "qk_gain_fp32": true, "positional": "rope"},
+  "data": {"train_path": "../../data/fineweb_large/tokenized", "val_path": null, "seed": 42, "max_seq_length": 8192, "num_workers": 8},
+  "training": {"output_dir": "models/FINAL_20260924/part_a/365m_dense_qkgainfp32_s42", "resume_from_checkpoint": "models/FINAL_20260924/part_a/365m_dense_qkgainfp32_s42/run_20260924_045342/checkpoint-9000", "force_new_run": false, "num_epochs": 1, "max_steps": 44000, "per_device_batch_size": 2, "gradient_accumulation_steps": 5, "resume_micro_step_override": 45000, "learning_rate": 0.0005, "weight_decay": 0.01, "warmup_steps": 750, "lr_scheduler": "cosine", "max_grad_norm": 1.0, "save_steps": 1000, "eval_steps": 250, "eval_batches": 50, "logging_steps": 100, "seed": 42, "verbose": false, "checkpoint_activations": false},
+  "distributed": {"strategy": "ddp", "mixed_precision": "bf16"},
+  "wandb": {"enabled": false, "project": "stochastic_sparsity", "name": "FINAL partA 365m dense qkgainfp32 s42", "tags": ["final_20260924", "part_a", "dense", "qkgainfp32", "seed42"]},
+  "attention": {"type": "dense", "dense_mode": false}
+}

epfl_final/SparseAtt/configs/FINAL_20260924/365m_dense_qkgainfp32_s42_5gpu_a100_resume.json

@@ -0,0 +1,8 @@
+{
+  "model": {"path": "../../qwen3_06b_base", "dtype": "bfloat16", "init_from_scratch": true, "hidden_size": 768, "intermediate_size": 3072, "num_hidden_layers": 28, "num_attention_heads": 6, "num_key_value_heads": 2, "rope_theta": 1000000.0, "qk_norm": "on", "qk_gain_fp32": true, "positional": "rope"},
+  "data": {"train_path": "../../data/fineweb_large/tokenized", "val_path": null, "seed": 42, "max_seq_length": 8192, "num_workers": 8},
+  "training": {"output_dir": "models/FINAL_20260924/part_a/365m_dense_qkgainfp32_s42", "resume_from_checkpoint": "models/FINAL_20260924/part_a/365m_dense_qkgainfp32_s42/run_20260924_045342/checkpoint-26000", "force_new_run": false, "num_epochs": 1, "max_steps": 44000, "per_device_batch_size": 1, "gradient_accumulation_steps": 4, "resume_micro_step_override": 104000, "learning_rate": 0.0005, "weight_decay": 0.01, "warmup_steps": 750, "lr_scheduler": "cosine", "max_grad_norm": 1.0, "save_steps": 1000, "eval_steps": 250, "eval_batches": 50, "logging_steps": 100, "seed": 42, "verbose": false, "checkpoint_activations": false},
+  "distributed": {"strategy": "ddp", "mixed_precision": "bf16"},
+  "wandb": {"enabled": false, "project": "stochastic_sparsity", "name": "FINAL partA 365m dense qkgainfp32 s42", "tags": ["final_20260924", "part_a", "dense", "qkgainfp32", "seed42", "a100_resume"]},
+  "attention": {"type": "dense", "dense_mode": false}
+}

epfl_final/SparseAtt/configs/FINAL_20260924/365m_dense_qkgainfp32_s43_2gpu.json

@@ -0,0 +1,8 @@
+{
+  "model": {"path": "../../qwen3_06b_base", "dtype": "bfloat16", "init_from_scratch": true, "hidden_size": 768, "intermediate_size": 3072, "num_hidden_layers": 28, "num_attention_heads": 6, "num_key_value_heads": 2, "rope_theta": 1000000.0, "qk_norm": "on", "qk_gain_fp32": true, "positional": "rope"},
+  "data": {"train_path": "../../data/fineweb_large/tokenized", "val_path": null, "seed": 43, "max_seq_length": 8192, "num_workers": 8},
+  "training": {"output_dir": "models/FINAL_20260924/part_a/365m_dense_qkgainfp32_s43", "resume_from_checkpoint": null, "force_new_run": false, "num_epochs": 1, "max_steps": 44000, "per_device_batch_size": 2, "gradient_accumulation_steps": 5, "learning_rate": 0.0005, "weight_decay": 0.01, "warmup_steps": 750, "lr_scheduler": "cosine", "max_grad_norm": 1.0, "save_steps": 1000, "eval_steps": 250, "eval_batches": 50, "logging_steps": 100, "seed": 43, "verbose": false, "checkpoint_activations": false},
+  "distributed": {"strategy": "ddp", "mixed_precision": "bf16"},
+  "wandb": {"enabled": false, "project": "stochastic_sparsity", "name": "FINAL partA 365m dense qkgainfp32 s43", "tags": ["final_20260924", "part_a", "dense", "qkgainfp32", "seed43"]},
+  "attention": {"type": "dense", "dense_mode": false}
+}

epfl_final/SparseAtt/configs/FINAL_20260924/365m_dense_qkgainfp32_s43_2gpu_resume.json

@@ -0,0 +1,8 @@
+{
+  "model": {"path": "../../qwen3_06b_base", "dtype": "bfloat16", "init_from_scratch": true, "hidden_size": 768, "intermediate_size": 3072, "num_hidden_layers": 28, "num_attention_heads": 6, "num_key_value_heads": 2, "rope_theta": 1000000.0, "qk_norm": "on", "qk_gain_fp32": true, "positional": "rope"},
+  "data": {"train_path": "../../data/fineweb_large/tokenized", "val_path": null, "seed": 43, "max_seq_length": 8192, "num_workers": 8},
+  "training": {"output_dir": "models/FINAL_20260924/part_a/365m_dense_qkgainfp32_s43", "resume_from_checkpoint": "models/FINAL_20260924/part_a/365m_dense_qkgainfp32_s43/run_20260924_152136/checkpoint-4000", "force_new_run": false, "num_epochs": 1, "max_steps": 44000, "per_device_batch_size": 2, "gradient_accumulation_steps": 5, "learning_rate": 0.0005, "weight_decay": 0.01, "warmup_steps": 750, "lr_scheduler": "cosine", "max_grad_norm": 1.0, "save_steps": 1000, "eval_steps": 250, "eval_batches": 50, "logging_steps": 100, "seed": 43, "verbose": false, "checkpoint_activations": false},
+  "distributed": {"strategy": "ddp", "mixed_precision": "bf16"},
+  "wandb": {"enabled": false, "project": "stochastic_sparsity", "name": "FINAL partA 365m dense qkgainfp32 s43", "tags": ["final_20260924", "part_a", "dense", "qkgainfp32", "seed43"]},
+  "attention": {"type": "dense", "dense_mode": false}
+}

epfl_final/SparseAtt/configs/FINAL_20260924/365m_dense_qkgainfp32_s44_5gpu.json

@@ -0,0 +1,8 @@
+{
+  "model": {"path": "../../qwen3_06b_base", "dtype": "bfloat16", "init_from_scratch": true, "hidden_size": 768, "intermediate_size": 3072, "num_hidden_layers": 28, "num_attention_heads": 6, "num_key_value_heads": 2, "rope_theta": 1000000.0, "qk_norm": "on", "qk_gain_fp32": true, "positional": "rope"},
+  "data": {"train_path": "../../data/fineweb_large/tokenized", "val_path": null, "seed": 44, "max_seq_length": 8192, "num_workers": 8},
+  "training": {"output_dir": "models/FINAL_20260924/part_a/365m_dense_qkgainfp32_s44", "resume_from_checkpoint": null, "force_new_run": false, "num_epochs": 1, "max_steps": 44000, "per_device_batch_size": 1, "gradient_accumulation_steps": 4, "learning_rate": 0.0005, "weight_decay": 0.01, "warmup_steps": 750, "lr_scheduler": "cosine", "max_grad_norm": 1.0, "save_steps": 1000, "eval_steps": 250, "eval_batches": 50, "logging_steps": 100, "seed": 44, "verbose": false, "checkpoint_activations": false},
+  "distributed": {"strategy": "ddp", "mixed_precision": "bf16"},
+  "wandb": {"enabled": false, "project": "stochastic_sparsity", "name": "FINAL partA 365m dense qkgainfp32 s44", "tags": ["final_20260924", "part_a", "dense", "qkgainfp32", "seed44"]},
+  "attention": {"type": "dense", "dense_mode": false}
+}

epfl_final/SparseAtt/configs/FINAL_20260924/365m_mixed_qkgainfp32_s42.json

@@ -0,0 +1,8 @@
+{
+  "model": {"path": "../../qwen3_06b_base", "dtype": "bfloat16", "init_from_scratch": true, "hidden_size": 768, "intermediate_size": 3072, "num_hidden_layers": 28, "num_attention_heads": 6, "num_key_value_heads": 2, "rope_theta": 1000000.0, "qk_norm": "on", "qk_gain_fp32": true, "positional": "rope"},
+  "data": {"train_path": "../../data/fineweb_large/tokenized", "val_path": null, "seed": 42, "max_seq_length": 8192, "num_workers": 8},
+  "training": {"output_dir": "models/FINAL_20260924/part_a/365m_mixed_qkgainfp32_s42", "resume_from_checkpoint": null, "force_new_run": false, "num_epochs": 1, "max_steps": 44000, "per_device_batch_size": 1, "gradient_accumulation_steps": 20, "learning_rate": 0.0005, "weight_decay": 0.01, "warmup_steps": 750, "lr_scheduler": "cosine", "max_grad_norm": 1.0, "save_steps": 1000, "eval_steps": 250, "eval_batches": 50, "logging_steps": 100, "seed": 42, "verbose": false, "checkpoint_activations": false},
+  "distributed": {"strategy": "ddp", "mixed_precision": "bf16"},
+  "wandb": {"enabled": false, "project": "stochastic_sparsity", "name": "FINAL partA 365m mixed qkgainfp32 s42", "tags": ["final_20260924", "part_a", "mixed", "qkgainfp32", "seed42"]},
+  "attention": {"type": "sparse_stochastic", "dense_mode": false, "sparse_block_size": 128, "sparse_num_random_blocks": 63, "sparse_mask_pool_size": 10000, "sparse_mixed_nrb_set": [63, 47, 31, 21, 11], "sparse_mask_seed": 0, "sparse_mask_cache_dir": "../mask_cache_8k_mixed"}
+}

epfl_final/SparseAtt/configs/FINAL_20260924/365m_mixed_qkgainfp32_s42_2gpu_h100_restart_resume.json

@@ -0,0 +1,8 @@
+{
+  "model": {"path": "../../qwen3_06b_base", "dtype": "bfloat16", "init_from_scratch": true, "hidden_size": 768, "intermediate_size": 3072, "num_hidden_layers": 28, "num_attention_heads": 6, "num_key_value_heads": 2, "rope_theta": 1000000.0, "qk_norm": "on", "qk_gain_fp32": true, "positional": "rope"},
+  "data": {"train_path": "../../data/fineweb_large/tokenized", "val_path": null, "seed": 42, "max_seq_length": 8192, "num_workers": 8},
+  "training": {"output_dir": "models/FINAL_20260924/part_a/365m_mixed_qkgainfp32_s42", "resume_from_checkpoint": "models/FINAL_20260924/part_a/365m_mixed_qkgainfp32_s42/run_20260924_045342/checkpoint-9000", "force_new_run": false, "num_epochs": 1, "max_steps": 44000, "per_device_batch_size": 2, "gradient_accumulation_steps": 5, "resume_micro_step_override": 45000, "learning_rate": 0.0005, "weight_decay": 0.01, "warmup_steps": 750, "lr_scheduler": "cosine", "max_grad_norm": 1.0, "save_steps": 1000, "eval_steps": 250, "eval_batches": 50, "logging_steps": 100, "seed": 42, "verbose": false, "checkpoint_activations": false},
+  "distributed": {"strategy": "ddp", "mixed_precision": "bf16"},
+  "wandb": {"enabled": false, "project": "stochastic_sparsity", "name": "FINAL partA 365m mixed qkgainfp32 s42", "tags": ["final_20260924", "part_a", "mixed", "qkgainfp32", "seed42", "h100_restart_resume"]},
+  "attention": {"type": "sparse_stochastic", "dense_mode": false, "sparse_block_size": 128, "sparse_num_random_blocks": 63, "sparse_mask_pool_size": 10000, "sparse_mixed_nrb_set": [63, 47, 31, 21, 11], "sparse_mask_seed": 0, "sparse_mask_cache_dir": "../mask_cache_8k_mixed"}
+}

epfl_final/SparseAtt/configs/FINAL_20260924/365m_mixed_qkgainfp32_s42_2gpu_h100_resume.json

@@ -0,0 +1,8 @@
+{
+  "model": {"path": "../../qwen3_06b_base", "dtype": "bfloat16", "init_from_scratch": true, "hidden_size": 768, "intermediate_size": 3072, "num_hidden_layers": 28, "num_attention_heads": 6, "num_key_value_heads": 2, "rope_theta": 1000000.0, "qk_norm": "on", "qk_gain_fp32": true, "positional": "rope"},
+  "data": {"train_path": "../../data/fineweb_large/tokenized", "val_path": null, "seed": 42, "max_seq_length": 8192, "num_workers": 8},
+  "training": {"output_dir": "models/FINAL_20260924/part_a/365m_mixed_qkgainfp32_s42", "resume_from_checkpoint": "models/FINAL_20260924/part_a/365m_mixed_qkgainfp32_s42/run_20260924_045342/checkpoint-26000", "force_new_run": false, "num_epochs": 1, "max_steps": 44000, "per_device_batch_size": 2, "gradient_accumulation_steps": 5, "resume_micro_step_override": 130000, "learning_rate": 0.0005, "weight_decay": 0.01, "warmup_steps": 750, "lr_scheduler": "cosine", "max_grad_norm": 1.0, "save_steps": 1000, "eval_steps": 250, "eval_batches": 50, "logging_steps": 100, "seed": 42, "verbose": false, "checkpoint_activations": false},
+  "distributed": {"strategy": "ddp", "mixed_precision": "bf16"},
+  "wandb": {"enabled": false, "project": "stochastic_sparsity", "name": "FINAL partA 365m mixed qkgainfp32 s42", "tags": ["final_20260924", "part_a", "mixed", "qkgainfp32", "seed42", "h100_resume"]},
+  "attention": {"type": "sparse_stochastic", "dense_mode": false, "sparse_block_size": 128, "sparse_num_random_blocks": 63, "sparse_mask_pool_size": 10000, "sparse_mixed_nrb_set": [63, 47, 31, 21, 11], "sparse_mask_seed": 0, "sparse_mask_cache_dir": "../mask_cache_8k_mixed"}
+}

epfl_final/SparseAtt/configs/FINAL_20260924/365m_mixed_qkgainfp32_s42_5gpu_latest_resume.json

@@ -0,0 +1,8 @@
+{
+  "model": {"path": "../../qwen3_06b_base", "dtype": "bfloat16", "init_from_scratch": true, "hidden_size": 768, "intermediate_size": 3072, "num_hidden_layers": 28, "num_attention_heads": 6, "num_key_value_heads": 2, "rope_theta": 1000000.0, "qk_norm": "on", "qk_gain_fp32": true, "positional": "rope"},
+  "data": {"train_path": "../../data/fineweb_large/tokenized", "val_path": null, "seed": 42, "max_seq_length": 8192, "num_workers": 8},
+  "training": {"output_dir": "models/FINAL_20260924/part_a/365m_mixed_qkgainfp32_s42", "resume_from_checkpoint": "models/FINAL_20260924/part_a/365m_mixed_qkgainfp32_s42/run_20260924_045342/checkpoint-32000", "force_new_run": false, "num_epochs": 1, "max_steps": 44000, "per_device_batch_size": 1, "gradient_accumulation_steps": 4, "resume_micro_step_override": 128000, "learning_rate": 0.0005, "weight_decay": 0.01, "warmup_steps": 750, "lr_scheduler": "cosine", "max_grad_norm": 1.0, "save_steps": 1000, "eval_steps": 250, "eval_batches": 50, "logging_steps": 100, "seed": 42, "verbose": false, "checkpoint_activations": false},
+  "distributed": {"strategy": "ddp", "mixed_precision": "bf16"},
+  "wandb": {"enabled": false, "project": "stochastic_sparsity", "name": "FINAL partA 365m mixed qkgainfp32 s42", "tags": ["final_20260924", "part_a", "mixed", "qkgainfp32", "seed42", "latest_resume"]},
+  "attention": {"type": "sparse_stochastic", "dense_mode": false, "sparse_block_size": 128, "sparse_num_random_blocks": 63, "sparse_mask_pool_size": 10000, "sparse_mixed_nrb_set": [63, 47, 31, 21, 11], "sparse_mask_seed": 0, "sparse_mask_cache_dir": "../mask_cache_8k_mixed"}
+}

epfl_final/SparseAtt/configs/FINAL_20260924/365m_mixed_qkgainfp32_s42_5gpu_restart_resume.json

@@ -0,0 +1,8 @@
+{
+  "model": {"path": "../../qwen3_06b_base", "dtype": "bfloat16", "init_from_scratch": true, "hidden_size": 768, "intermediate_size": 3072, "num_hidden_layers": 28, "num_attention_heads": 6, "num_key_value_heads": 2, "rope_theta": 1000000.0, "qk_norm": "on", "qk_gain_fp32": true, "positional": "rope"},
+  "data": {"train_path": "../../data/fineweb_large/tokenized", "val_path": null, "seed": 42, "max_seq_length": 8192, "num_workers": 8},
+  "training": {"output_dir": "models/FINAL_20260924/part_a/365m_mixed_qkgainfp32_s42", "resume_from_checkpoint": "models/FINAL_20260924/part_a/365m_mixed_qkgainfp32_s42/run_20260924_045342/checkpoint-9000", "force_new_run": false, "num_epochs": 1, "max_steps": 44000, "per_device_batch_size": 1, "gradient_accumulation_steps": 4, "resume_micro_step_override": 36000, "learning_rate": 0.0005, "weight_decay": 0.01, "warmup_steps": 750, "lr_scheduler": "cosine", "max_grad_norm": 1.0, "save_steps": 1000, "eval_steps": 250, "eval_batches": 50, "logging_steps": 100, "seed": 42, "verbose": false, "checkpoint_activations": false},
+  "distributed": {"strategy": "ddp", "mixed_precision": "bf16"},
+  "wandb": {"enabled": false, "project": "stochastic_sparsity", "name": "FINAL partA 365m mixed qkgainfp32 s42", "tags": ["final_20260924", "part_a", "mixed", "qkgainfp32", "seed42", "a100_restart_resume"]},
+  "attention": {"type": "sparse_stochastic", "dense_mode": false, "sparse_block_size": 128, "sparse_num_random_blocks": 63, "sparse_mask_pool_size": 10000, "sparse_mixed_nrb_set": [63, 47, 31, 21, 11], "sparse_mask_seed": 0, "sparse_mask_cache_dir": "../mask_cache_8k_mixed"}
+}

epfl_final/SparseAtt/configs/FINAL_20260924/365m_mixed_qkgainfp32_s43.json

@@ -0,0 +1,8 @@
+{
+  "model": {"path": "../../qwen3_06b_base", "dtype": "bfloat16", "init_from_scratch": true, "hidden_size": 768, "intermediate_size": 3072, "num_hidden_layers": 28, "num_attention_heads": 6, "num_key_value_heads": 2, "rope_theta": 1000000.0, "qk_norm": "on", "qk_gain_fp32": true, "positional": "rope"},
+  "data": {"train_path": "../../data/fineweb_large/tokenized", "val_path": null, "seed": 43, "max_seq_length": 8192, "num_workers": 8},
+  "training": {"output_dir": "models/FINAL_20260924/part_a/365m_mixed_qkgainfp32_s43", "resume_from_checkpoint": null, "force_new_run": false, "num_epochs": 1, "max_steps": 44000, "per_device_batch_size": 1, "gradient_accumulation_steps": 20, "learning_rate": 0.0005, "weight_decay": 0.01, "warmup_steps": 750, "lr_scheduler": "cosine", "max_grad_norm": 1.0, "save_steps": 1000, "eval_steps": 250, "eval_batches": 50, "logging_steps": 100, "seed": 43, "verbose": false, "checkpoint_activations": false},
+  "distributed": {"strategy": "ddp", "mixed_precision": "bf16"},
+  "wandb": {"enabled": false, "project": "stochastic_sparsity", "name": "FINAL partA 365m mixed qkgainfp32 s43", "tags": ["final_20260924", "part_a", "mixed", "qkgainfp32", "seed43"]},
+  "attention": {"type": "sparse_stochastic", "dense_mode": false, "sparse_block_size": 128, "sparse_num_random_blocks": 63, "sparse_mask_pool_size": 10000, "sparse_mixed_nrb_set": [63, 47, 31, 21, 11], "sparse_mask_seed": 0, "sparse_mask_cache_dir": "../mask_cache_8k_mixed"}
+}

epfl_final/SparseAtt/configs/FINAL_20260924/365m_mixed_qkgainfp32_s43_2gpu_resume.json

@@ -0,0 +1,8 @@
+{
+  "model": {"path": "../../qwen3_06b_base", "dtype": "bfloat16", "init_from_scratch": true, "hidden_size": 768, "intermediate_size": 3072, "num_hidden_layers": 28, "num_attention_heads": 6, "num_key_value_heads": 2, "rope_theta": 1000000.0, "qk_norm": "on", "qk_gain_fp32": true, "positional": "rope"},
+  "data": {"train_path": "../../data/fineweb_large/tokenized", "val_path": null, "seed": 43, "max_seq_length": 8192, "num_workers": 8},
+  "training": {"output_dir": "models/FINAL_20260924/part_a/365m_mixed_qkgainfp32_s43", "resume_from_checkpoint": "models/FINAL_20260924/part_a/365m_mixed_qkgainfp32_s43/run_20260924_050049/checkpoint-5000", "force_new_run": false, "num_epochs": 1, "max_steps": 44000, "per_device_batch_size": 1, "gradient_accumulation_steps": 10, "resume_micro_step_override": 50000, "learning_rate": 0.0005, "weight_decay": 0.01, "warmup_steps": 750, "lr_scheduler": "cosine", "max_grad_norm": 1.0, "save_steps": 1000, "eval_steps": 250, "eval_batches": 50, "logging_steps": 100, "seed": 43, "verbose": false, "checkpoint_activations": false},
+  "distributed": {"strategy": "ddp", "mixed_precision": "bf16"},
+  "wandb": {"enabled": false, "project": "stochastic_sparsity", "name": "FINAL partA 365m mixed qkgainfp32 s43", "tags": ["final_20260924", "part_a", "mixed", "qkgainfp32", "seed43"]},
+  "attention": {"type": "sparse_stochastic", "dense_mode": false, "sparse_block_size": 128, "sparse_num_random_blocks": 63, "sparse_mask_pool_size": 10000, "sparse_mixed_nrb_set": [63, 47, 31, 21, 11], "sparse_mask_seed": 0, "sparse_mask_cache_dir": "../mask_cache_8k_mixed"}
+}

epfl_final/SparseAtt/configs/FINAL_20260924/365m_mixed_qkgainfp32_s43_5gpu_resume.json

@@ -0,0 +1,8 @@
+{
+  "model": {"path": "../../qwen3_06b_base", "dtype": "bfloat16", "init_from_scratch": true, "hidden_size": 768, "intermediate_size": 3072, "num_hidden_layers": 28, "num_attention_heads": 6, "num_key_value_heads": 2, "rope_theta": 1000000.0, "qk_norm": "on", "qk_gain_fp32": true, "positional": "rope"},
+  "data": {"train_path": "../../data/fineweb_large/tokenized", "val_path": null, "seed": 43, "max_seq_length": 8192, "num_workers": 8},
+  "training": {"output_dir": "models/FINAL_20260924/part_a/365m_mixed_qkgainfp32_s43", "resume_from_checkpoint": "models/FINAL_20260924/part_a/365m_mixed_qkgainfp32_s43/run_20260924_050049/checkpoint-5000", "force_new_run": false, "num_epochs": 1, "max_steps": 44000, "per_device_batch_size": 1, "gradient_accumulation_steps": 4, "resume_micro_step_override": 20000, "learning_rate": 0.0005, "weight_decay": 0.01, "warmup_steps": 750, "lr_scheduler": "cosine", "max_grad_norm": 1.0, "save_steps": 1000, "eval_steps": 250, "eval_batches": 50, "logging_steps": 100, "seed": 43, "verbose": false, "checkpoint_activations": false},
+  "distributed": {"strategy": "ddp", "mixed_precision": "bf16"},
+  "wandb": {"enabled": false, "project": "stochastic_sparsity", "name": "FINAL partA 365m mixed qkgainfp32 s43", "tags": ["final_20260924", "part_a", "mixed", "qkgainfp32", "seed43"]},
+  "attention": {"type": "sparse_stochastic", "dense_mode": false, "sparse_block_size": 128, "sparse_num_random_blocks": 63, "sparse_mask_pool_size": 10000, "sparse_mixed_nrb_set": [63, 47, 31, 21, 11], "sparse_mask_seed": 0, "sparse_mask_cache_dir": "../mask_cache_8k_mixed"}
+}

epfl_final/SparseAtt/configs/FINAL_20260924/365m_mixed_qkgainfp32_s44_2gpu_h100_latest_resume.json

@@ -0,0 +1,8 @@
+{
+  "model": {"path": "../../qwen3_06b_base", "dtype": "bfloat16", "init_from_scratch": true, "hidden_size": 768, "intermediate_size": 3072, "num_hidden_layers": 28, "num_attention_heads": 6, "num_key_value_heads": 2, "rope_theta": 1000000.0, "qk_norm": "on", "qk_gain_fp32": true, "positional": "rope"},
+  "data": {"train_path": "../../data/fineweb_large/tokenized", "val_path": null, "seed": 44, "max_seq_length": 8192, "num_workers": 8},
+  "training": {"output_dir": "models/FINAL_20260924/part_a/365m_mixed_qkgainfp32_s44", "resume_from_checkpoint": "models/FINAL_20260924/part_a/365m_mixed_qkgainfp32_s44/run_20260924_144318/checkpoint-22000", "force_new_run": false, "num_epochs": 1, "max_steps": 44000, "per_device_batch_size": 2, "gradient_accumulation_steps": 5, "learning_rate": 0.0005, "weight_decay": 0.01, "warmup_steps": 750, "lr_scheduler": "cosine", "max_grad_norm": 1.0, "save_steps": 1000, "eval_steps": 250, "eval_batches": 50, "logging_steps": 100, "seed": 44, "verbose": false, "checkpoint_activations": false},
+  "distributed": {"strategy": "ddp", "mixed_precision": "bf16"},
+  "wandb": {"enabled": false, "project": "stochastic_sparsity", "name": "FINAL partA 365m mixed qkgainfp32 s44", "tags": ["final_20260924", "part_a", "mixed", "qkgainfp32", "seed44", "latest_resume"]},
+  "attention": {"type": "sparse_stochastic", "dense_mode": false, "sparse_block_size": 128, "sparse_num_random_blocks": 63, "sparse_mask_pool_size": 10000, "sparse_mixed_nrb_set": [63, 47, 31, 21, 11], "sparse_mask_seed": 0, "sparse_mask_cache_dir": "../mask_cache_8k_mixed"}
+}

epfl_final/SparseAtt/configs/FINAL_20260924/365m_mixed_qkgainfp32_s44_2gpu_h100_resume.json

@@ -0,0 +1,8 @@
+{
+  "model": {"path": "../../qwen3_06b_base", "dtype": "bfloat16", "init_from_scratch": true, "hidden_size": 768, "intermediate_size": 3072, "num_hidden_layers": 28, "num_attention_heads": 6, "num_key_value_heads": 2, "rope_theta": 1000000.0, "qk_norm": "on", "qk_gain_fp32": true, "positional": "rope"},
+  "data": {"train_path": "../../data/fineweb_large/tokenized", "val_path": null, "seed": 44, "max_seq_length": 8192, "num_workers": 8},
+  "training": {"output_dir": "models/FINAL_20260924/part_a/365m_mixed_qkgainfp32_s44", "resume_from_checkpoint": "models/FINAL_20260924/part_a/365m_mixed_qkgainfp32_s44/run_20260924_144318/checkpoint-21000", "force_new_run": false, "num_epochs": 1, "max_steps": 44000, "per_device_batch_size": 2, "gradient_accumulation_steps": 5, "resume_micro_step_override": 105000, "learning_rate": 0.0005, "weight_decay": 0.01, "warmup_steps": 750, "lr_scheduler": "cosine", "max_grad_norm": 1.0, "save_steps": 1000, "eval_steps": 250, "eval_batches": 50, "logging_steps": 100, "seed": 44, "verbose": false, "checkpoint_activations": false},
+  "distributed": {"strategy": "ddp", "mixed_precision": "bf16"},
+  "wandb": {"enabled": false, "project": "stochastic_sparsity", "name": "FINAL partA 365m mixed qkgainfp32 s44", "tags": ["final_20260924", "part_a", "mixed", "qkgainfp32", "seed44", "h100_resume"]},
+  "attention": {"type": "sparse_stochastic", "dense_mode": false, "sparse_block_size": 128, "sparse_num_random_blocks": 63, "sparse_mask_pool_size": 10000, "sparse_mixed_nrb_set": [63, 47, 31, 21, 11], "sparse_mask_seed": 0, "sparse_mask_cache_dir": "../mask_cache_8k_mixed"}
+}

epfl_final/SparseAtt/configs/FINAL_20260924/365m_mixed_qkgainfp32_s44_5gpu.json

@@ -0,0 +1,8 @@
+{
+  "model": {"path": "../../qwen3_06b_base", "dtype": "bfloat16", "init_from_scratch": true, "hidden_size": 768, "intermediate_size": 3072, "num_hidden_layers": 28, "num_attention_heads": 6, "num_key_value_heads": 2, "rope_theta": 1000000.0, "qk_norm": "on", "qk_gain_fp32": true, "positional": "rope"},
+  "data": {"train_path": "../../data/fineweb_large/tokenized", "val_path": null, "seed": 44, "max_seq_length": 8192, "num_workers": 8},
+  "training": {"output_dir": "models/FINAL_20260924/part_a/365m_mixed_qkgainfp32_s44", "resume_from_checkpoint": null, "force_new_run": false, "num_epochs": 1, "max_steps": 44000, "per_device_batch_size": 1, "gradient_accumulation_steps": 4, "learning_rate": 0.0005, "weight_decay": 0.01, "warmup_steps": 750, "lr_scheduler": "cosine", "max_grad_norm": 1.0, "save_steps": 1000, "eval_steps": 250, "eval_batches": 50, "logging_steps": 100, "seed": 44, "verbose": false, "checkpoint_activations": false},
+  "distributed": {"strategy": "ddp", "mixed_precision": "bf16"},
+  "wandb": {"enabled": false, "project": "stochastic_sparsity", "name": "FINAL partA 365m mixed qkgainfp32 s44", "tags": ["final_20260924", "part_a", "mixed", "qkgainfp32", "seed44"]},
+  "attention": {"type": "sparse_stochastic", "dense_mode": false, "sparse_block_size": 128, "sparse_num_random_blocks": 63, "sparse_mask_pool_size": 10000, "sparse_mixed_nrb_set": [63, 47, 31, 21, 11], "sparse_mask_seed": 0, "sparse_mask_cache_dir": "../mask_cache_8k_mixed"}
+}

epfl_final/SparseAtt/configs/REGIME_20260922/365m_dense_noqknorm_s42.json

@@ -0,0 +1,51 @@
+{
+  "model": {
+    "path": "../../qwen3_06b_base",
+    "dtype": "bfloat16",
+    "init_from_scratch": true,
+    "hidden_size": 768,
+    "intermediate_size": 3072,
+    "num_hidden_layers": 28,
+    "num_attention_heads": 6,
+    "num_key_value_heads": 2,
+    "rope_theta": 1000000.0,
+    "qk_norm": "off",
+    "positional": "rope"
+  },
+  "data": {
+    "train_path": "../../data/fineweb_large/tokenized",
+    "val_path": null,
+    "seed": 0,
+    "max_seq_length": 8192,
+    "num_workers": 8
+  },
+  "training": {
+    "output_dir": "models/REGIME_20260922/365m_dense_noqknorm_s42",
+    "resume_from_checkpoint": null,
+    "force_new_run": false,
+    "num_epochs": 1,
+    "max_steps": 44000,
+    "per_device_batch_size": 1,
+    "gradient_accumulation_steps": 4,
+    "learning_rate": 0.0005,
+    "weight_decay": 0.01,
+    "warmup_steps": 750,
+    "lr_scheduler": "cosine",
+    "max_grad_norm": 1.0,
+    "save_steps": 1000,
+    "eval_steps": 250,
+    "eval_batches": 50,
+    "logging_steps": 100,
+    "seed": 42,
+    "verbose": true,
+    "checkpoint_activations": false
+  },
+  "distributed": {"strategy": "ddp", "mixed_precision": "bf16"},
+  "wandb": {
+    "enabled": true,
+    "project": "stochastic_sparsity",
+    "name": "REGIME 365m dense noqknorm s42",
+    "tags": ["regime_20260922", "365m", "dense", "noqknorm", "seed42"]
+  },
+  "attention": {"type": "dense", "dense_mode": false}
+}

epfl_final/SparseAtt/configs/REGIME_20260922/365m_dense_rope1e4_s42.json

@@ -0,0 +1,51 @@
+{
+  "model": {
+    "path": "../../qwen3_06b_base",
+    "dtype": "bfloat16",
+    "init_from_scratch": true,
+    "hidden_size": 768,
+    "intermediate_size": 3072,
+    "num_hidden_layers": 28,
+    "num_attention_heads": 6,
+    "num_key_value_heads": 2,
+    "rope_theta": 10000.0,
+    "qk_norm": "on",
+    "positional": "rope"
+  },
+  "data": {
+    "train_path": "../../data/fineweb_large/tokenized",
+    "val_path": null,
+    "seed": 0,
+    "max_seq_length": 8192,
+    "num_workers": 8
+  },
+  "training": {
+    "output_dir": "models/REGIME_20260922/365m_dense_rope1e4_s42",
+    "resume_from_checkpoint": null,
+    "force_new_run": false,
+    "num_epochs": 1,
+    "max_steps": 44000,
+    "per_device_batch_size": 1,
+    "gradient_accumulation_steps": 4,
+    "learning_rate": 0.0005,
+    "weight_decay": 0.01,
+    "warmup_steps": 750,
+    "lr_scheduler": "cosine",
+    "max_grad_norm": 1.0,
+    "save_steps": 1000,
+    "eval_steps": 250,
+    "eval_batches": 50,
+    "logging_steps": 100,
+    "seed": 42,
+    "verbose": true,
+    "checkpoint_activations": false
+  },
+  "distributed": {"strategy": "ddp", "mixed_precision": "bf16"},
+  "wandb": {
+    "enabled": true,
+    "project": "stochastic_sparsity",
+    "name": "REGIME 365m dense rope1e4 s42",
+    "tags": ["regime_20260922", "365m", "dense", "rope1e4", "seed42"]
+  },
+  "attention": {"type": "dense", "dense_mode": false}
+}

epfl_final/SparseAtt/configs/REGIME_20260922/365m_dense_rope1e4_s42_1gpu.json

@@ -0,0 +1,51 @@
+{
+  "model": {
+    "path": "../../qwen3_06b_base",
+    "dtype": "bfloat16",
+    "init_from_scratch": true,
+    "hidden_size": 768,
+    "intermediate_size": 3072,
+    "num_hidden_layers": 28,
+    "num_attention_heads": 6,
+    "num_key_value_heads": 2,
+    "rope_theta": 10000.0,
+    "qk_norm": "on",
+    "positional": "rope"
+  },
+  "data": {
+    "train_path": "../../data/fineweb_large/tokenized",
+    "val_path": null,
+    "seed": 0,
+    "max_seq_length": 8192,
+    "num_workers": 8
+  },
+  "training": {
+    "output_dir": "models/REGIME_20260922/365m_dense_rope1e4_s42",
+    "resume_from_checkpoint": null,
+    "force_new_run": false,
+    "num_epochs": 1,
+    "max_steps": 44000,
+    "per_device_batch_size": 1,
+    "gradient_accumulation_steps": 20,
+    "learning_rate": 0.0005,
+    "weight_decay": 0.01,
+    "warmup_steps": 750,
+    "lr_scheduler": "cosine",
+    "max_grad_norm": 1.0,
+    "save_steps": 1000,
+    "eval_steps": 250,
+    "eval_batches": 50,
+    "logging_steps": 100,
+    "seed": 42,
+    "verbose": true,
+    "checkpoint_activations": false
+  },
+  "distributed": {"strategy": "ddp", "mixed_precision": "bf16"},
+  "wandb": {
+    "enabled": true,
+    "project": "stochastic_sparsity",
+    "name": "REGIME 365m dense rope1e4 s42 1gpu",
+    "tags": ["regime_20260922", "365m", "dense", "rope1e4", "seed42"]
+  },
+  "attention": {"type": "dense", "dense_mode": false}
+}

epfl_final/SparseAtt/configs/REGIME_20260922/365m_dense_rope1e4_s42_resume.json

@@ -0,0 +1,64 @@
+{
+  "model": {
+    "path": "../../qwen3_06b_base",
+    "dtype": "bfloat16",
+    "init_from_scratch": true,
+    "hidden_size": 768,
+    "intermediate_size": 3072,
+    "num_hidden_layers": 28,
+    "num_attention_heads": 6,
+    "num_key_value_heads": 2,
+    "rope_theta": 10000.0,
+    "qk_norm": "on",
+    "positional": "rope"
+  },
+  "data": {
+    "train_path": "../../data/fineweb_large/tokenized",
+    "val_path": null,
+    "seed": 0,
+    "max_seq_length": 8192,
+    "num_workers": 8
+  },
+  "training": {
+    "output_dir": "models/REGIME_20260922/365m_dense_rope1e4_s42",
+    "resume_from_checkpoint": "models/REGIME_20260922/365m_dense_rope1e4_s42/run_20260922_225417/checkpoint-29000",
+    "force_new_run": false,
+    "num_epochs": 1,
+    "max_steps": 44000,
+    "per_device_batch_size": 1,
+    "gradient_accumulation_steps": 10,
+    "learning_rate": 0.0005,
+    "weight_decay": 0.01,
+    "warmup_steps": 750,
+    "lr_scheduler": "cosine",
+    "max_grad_norm": 1.0,
+    "save_steps": 1000,
+    "eval_steps": 250,
+    "eval_batches": 50,
+    "logging_steps": 100,
+    "seed": 42,
+    "verbose": true,
+    "checkpoint_activations": false,
+    "resume_micro_step_override": 290000
+  },
+  "distributed": {
+    "strategy": "ddp",
+    "mixed_precision": "bf16"
+  },
+  "wandb": {
+    "enabled": true,
+    "project": "stochastic_sparsity",
+    "name": "REGIME 365m dense rope1e4 s42",
+    "tags": [
+      "regime_20260922",
+      "365m",
+      "dense",
+      "rope1e4",
+      "seed42"
+    ]
+  },
+  "attention": {
+    "type": "dense",
+    "dense_mode": false
+  }
+}

epfl_final/SparseAtt/configs/REGIME_20260922/365m_mixed_noqknorm_s42.json

@@ -0,0 +1,60 @@
+{
+  "model": {
+    "path": "../../qwen3_06b_base",
+    "dtype": "bfloat16",
+    "init_from_scratch": true,
+    "hidden_size": 768,
+    "intermediate_size": 3072,
+    "num_hidden_layers": 28,
+    "num_attention_heads": 6,
+    "num_key_value_heads": 2,
+    "rope_theta": 1000000.0,
+    "qk_norm": "off",
+    "positional": "rope"
+  },
+  "data": {
+    "train_path": "../../data/fineweb_large/tokenized",
+    "val_path": null,
+    "seed": 0,
+    "max_seq_length": 8192,
+    "num_workers": 8
+  },
+  "training": {
+    "output_dir": "models/REGIME_20260922/365m_mixed_noqknorm_s42",
+    "resume_from_checkpoint": null,
+    "force_new_run": false,
+    "num_epochs": 1,
+    "max_steps": 44000,
+    "per_device_batch_size": 1,
+    "gradient_accumulation_steps": 4,
+    "learning_rate": 0.0005,
+    "weight_decay": 0.01,
+    "warmup_steps": 750,
+    "lr_scheduler": "cosine",
+    "max_grad_norm": 1.0,
+    "save_steps": 1000,
+    "eval_steps": 250,
+    "eval_batches": 50,
+    "logging_steps": 100,
+    "seed": 42,
+    "verbose": true,
+    "checkpoint_activations": false
+  },
+  "distributed": {"strategy": "ddp", "mixed_precision": "bf16"},
+  "wandb": {
+    "enabled": true,
+    "project": "stochastic_sparsity",
+    "name": "REGIME 365m mixed noqknorm s42",
+    "tags": ["regime_20260922", "365m", "mixed", "noqknorm", "seed42"]
+  },
+  "attention": {
+    "type": "sparse_stochastic",
+    "dense_mode": false,
+    "sparse_block_size": 128,
+    "sparse_num_random_blocks": 63,
+    "sparse_mask_pool_size": 10000,
+    "sparse_mixed_nrb_set": [63, 47, 31, 21, 11],
+    "sparse_mask_seed": 0,
+    "sparse_mask_cache_dir": "../mask_cache_8k_mixed"
+  }
+}

epfl_final/SparseAtt/configs/REGIME_20260922/365m_mixed_noqknorm_s42_resume.json

@@ -0,0 +1,60 @@
+{
+  "model": {
+    "path": "../../qwen3_06b_base",
+    "dtype": "bfloat16",
+    "init_from_scratch": true,
+    "hidden_size": 768,
+    "intermediate_size": 3072,
+    "num_hidden_layers": 28,
+    "num_attention_heads": 6,
+    "num_key_value_heads": 2,
+    "rope_theta": 1000000.0,
+    "qk_norm": "off",
+    "positional": "rope"
+  },
+  "data": {
+    "train_path": "../../data/fineweb_large/tokenized",
+    "val_path": null,
+    "seed": 0,
+    "max_seq_length": 8192,
+    "num_workers": 8
+  },
+  "training": {
+    "output_dir": "models/REGIME_20260922/365m_mixed_noqknorm_s42",
+    "resume_from_checkpoint": "models/REGIME_20260922/365m_mixed_noqknorm_s42/run_20260922_225813/checkpoint-40000",
+    "force_new_run": false,
+    "num_epochs": 1,
+    "max_steps": 44000,
+    "per_device_batch_size": 1,
+    "gradient_accumulation_steps": 4,
+    "learning_rate": 0.0005,
+    "weight_decay": 0.01,
+    "warmup_steps": 750,
+    "lr_scheduler": "cosine",
+    "max_grad_norm": 1.0,
+    "save_steps": 1000,
+    "eval_steps": 250,
+    "eval_batches": 50,
+    "logging_steps": 100,
+    "seed": 42,
+    "verbose": true,
+    "checkpoint_activations": false
+  },
+  "distributed": {"strategy": "ddp", "mixed_precision": "bf16"},
+  "wandb": {
+    "enabled": true,
+    "project": "stochastic_sparsity",
+    "name": "REGIME 365m mixed noqknorm s42",
+    "tags": ["regime_20260922", "365m", "mixed", "noqknorm", "seed42"]
+  },
+  "attention": {
+    "type": "sparse_stochastic",
+    "dense_mode": false,
+    "sparse_block_size": 128,
+    "sparse_num_random_blocks": 63,
+    "sparse_mask_pool_size": 10000,
+    "sparse_mixed_nrb_set": [63, 47, 31, 21, 11],
+    "sparse_mask_seed": 0,
+    "sparse_mask_cache_dir": "../mask_cache_8k_mixed"
+  }
+}

epfl_final/SparseAtt/configs/REGIME_20260922/365m_mixed_rope1e4_s42.json

@@ -0,0 +1,60 @@
+{
+  "model": {
+    "path": "../../qwen3_06b_base",
+    "dtype": "bfloat16",
+    "init_from_scratch": true,
+    "hidden_size": 768,
+    "intermediate_size": 3072,
+    "num_hidden_layers": 28,
+    "num_attention_heads": 6,
+    "num_key_value_heads": 2,
+    "rope_theta": 10000.0,
+    "qk_norm": "on",
+    "positional": "rope"
+  },
+  "data": {
+    "train_path": "../../data/fineweb_large/tokenized",
+    "val_path": null,
+    "seed": 0,
+    "max_seq_length": 8192,
+    "num_workers": 8
+  },
+  "training": {
+    "output_dir": "models/REGIME_20260922/365m_mixed_rope1e4_s42",
+    "resume_from_checkpoint": null,
+    "force_new_run": false,
+    "num_epochs": 1,
+    "max_steps": 44000,
+    "per_device_batch_size": 1,
+    "gradient_accumulation_steps": 4,
+    "learning_rate": 0.0005,
+    "weight_decay": 0.01,
+    "warmup_steps": 750,
+    "lr_scheduler": "cosine",
+    "max_grad_norm": 1.0,
+    "save_steps": 1000,
+    "eval_steps": 250,
+    "eval_batches": 50,
+    "logging_steps": 100,
+    "seed": 42,
+    "verbose": true,
+    "checkpoint_activations": false
+  },
+  "distributed": {"strategy": "ddp", "mixed_precision": "bf16"},
+  "wandb": {
+    "enabled": true,
+    "project": "stochastic_sparsity",
+    "name": "REGIME 365m mixed rope1e4 s42",
+    "tags": ["regime_20260922", "365m", "mixed", "rope1e4", "seed42"]
+  },
+  "attention": {
+    "type": "sparse_stochastic",
+    "dense_mode": false,
+    "sparse_block_size": 128,
+    "sparse_num_random_blocks": 63,
+    "sparse_mask_pool_size": 10000,
+    "sparse_mixed_nrb_set": [63, 47, 31, 21, 11],
+    "sparse_mask_seed": 0,
+    "sparse_mask_cache_dir": "../mask_cache_8k_mixed"
+  }
+}

epfl_final/SparseAtt/configs/REGIME_20260922/365m_mixed_rope1e4_s42_2gpu.json

@@ -0,0 +1,60 @@
+{
+  "model": {
+    "path": "../../qwen3_06b_base",
+    "dtype": "bfloat16",
+    "init_from_scratch": true,
+    "hidden_size": 768,
+    "intermediate_size": 3072,
+    "num_hidden_layers": 28,
+    "num_attention_heads": 6,
+    "num_key_value_heads": 2,
+    "rope_theta": 10000.0,
+    "qk_norm": "on",
+    "positional": "rope"
+  },
+  "data": {
+    "train_path": "../../data/fineweb_large/tokenized",
+    "val_path": null,
+    "seed": 0,
+    "max_seq_length": 8192,
+    "num_workers": 8
+  },
+  "training": {
+    "output_dir": "models/REGIME_20260922/365m_mixed_rope1e4_s42",
+    "resume_from_checkpoint": null,
+    "force_new_run": false,
+    "num_epochs": 1,
+    "max_steps": 44000,
+    "per_device_batch_size": 1,
+    "gradient_accumulation_steps": 10,
+    "learning_rate": 0.0005,
+    "weight_decay": 0.01,
+    "warmup_steps": 750,
+    "lr_scheduler": "cosine",
+    "max_grad_norm": 1.0,
+    "save_steps": 1000,
+    "eval_steps": 250,
+    "eval_batches": 50,
+    "logging_steps": 100,
+    "seed": 42,
+    "verbose": true,
+    "checkpoint_activations": false
+  },
+  "distributed": {"strategy": "ddp", "mixed_precision": "bf16"},
+  "wandb": {
+    "enabled": true,
+    "project": "stochastic_sparsity",
+    "name": "REGIME 365m mixed rope1e4 s42 2gpu",
+    "tags": ["regime_20260922", "365m", "mixed", "rope1e4", "seed42"]
+  },
+  "attention": {
+    "type": "sparse_stochastic",
+    "dense_mode": false,
+    "sparse_block_size": 128,
+    "sparse_num_random_blocks": 63,
+    "sparse_mask_pool_size": 10000,
+    "sparse_mixed_nrb_set": [63, 47, 31, 21, 11],
+    "sparse_mask_seed": 0,
+    "sparse_mask_cache_dir": "../mask_cache_8k_mixed"
+  }
+}

epfl_final/SparseAtt/configs/alibi_20260921/1p2b_alibi_s42.json

@@ -0,0 +1,62 @@
+{
+  "model": {
+    "path": "/mloscratch/homes/navasard/qwen3_06b_base",
+    "dtype": "bfloat16",
+    "init_from_scratch": true,
+    "hidden_size": 1536,
+    "intermediate_size": 6144,
+    "num_hidden_layers": 28,
+    "num_attention_heads": 12,
+    "num_key_value_heads": 4,
+    "positional": "alibi"
+  },
+  "data": {
+    "train_path": "/mloscratch/homes/navasard/data/fineweb_large/tokenized",
+    "val_path": null,
+    "max_seq_length": 8192,
+    "num_workers": 8
+  },
+  "training": {
+    "output_dir": "/mloscratch/homes/navasard/sparse/models/ALIBI_20260921/1p2b_alibi_s42",
+    "resume_from_checkpoint": null,
+    "force_new_run": false,
+    "num_epochs": 1,
+    "max_steps": 50462,
+    "per_device_batch_size": 1,
+    "gradient_accumulation_steps": 10,
+    "learning_rate": 0.0005,
+    "weight_decay": 0.01,
+    "warmup_steps": 750,
+    "lr_scheduler": "cosine",
+    "max_grad_norm": 1.0,
+    "save_steps": 1000,
+    "eval_steps": 1000,
+    "eval_batches": 50,
+    "logging_steps": 100,
+    "seed": 42,
+    "verbose": true,
+    "checkpoint_activations": false
+  },
+  "distributed": {
+    "strategy": "ddp",
+    "mixed_precision": "bf16"
+  },
+  "wandb": {
+    "enabled": true,
+    "project": "stochastic_sparsity",
+    "name": "1.2b alibi 8k s42",
+    "tags": ["qwen3", "1.2b", "8k", "alibi", "dense", "seed42"]
+  },
+  "attention": {
+    "type": "dense",
+    "dense_mode": false,
+    "sparse_block_size": 128,
+    "sparse_num_random_blocks": 31,
+    "sparse_mask_pool_size": 1,
+    "sparse_val_dense_compare": false,
+    "dense_step_intervals": null,
+    "sparse_layer_indices": null,
+    "sparse_mask_cache_dir": null,
+    "sparse_mask_seed": 42
+  }
+}

epfl_final/SparseAtt/configs/alibi_20260921/1p2b_alibi_s42_bs12_resume17000.json

@@ -0,0 +1,59 @@
+{
+  "model": {
+    "path": "../../qwen3_06b_base",
+    "dtype": "bfloat16",
+    "init_from_scratch": true,
+    "hidden_size": 1536,
+    "intermediate_size": 6144,
+    "num_hidden_layers": 28,
+    "num_attention_heads": 12,
+    "num_key_value_heads": 4,
+    "positional": "alibi"
+  },
+  "data": {
+    "train_path": "../../data/fineweb_large/tokenized",
+    "val_path": null,
+    "max_seq_length": 8192,
+    "num_workers": 8
+  },
+  "training": {
+    "output_dir": "../models/ALIBI_20260921/1p2b_alibi_s42",
+    "resume_from_checkpoint": "../models/ALIBI_20260921/1p2b_alibi_s42/run_20260922_004519/checkpoint-17000",
+    "force_new_run": false,
+    "num_epochs": 1,
+    "max_steps": 50462,
+    "per_device_batch_size": 1,
+    "gradient_accumulation_steps": 12,
+    "learning_rate": 0.0005,
+    "weight_decay": 0.01,
+    "warmup_steps": 750,
+    "lr_scheduler": "cosine",
+    "max_grad_norm": 1.0,
+    "save_steps": 1000,
+    "eval_steps": 1000,
+    "eval_batches": 50,
+    "logging_steps": 100,
+    "seed": 42,
+    "verbose": true,
+    "checkpoint_activations": false
+  },
+  "distributed": {"strategy": "ddp", "mixed_precision": "bf16"},
+  "wandb": {
+    "enabled": true,
+    "project": "stochastic_sparsity",
+    "name": "1.2b alibi 8k s42 resume17000 5gpu",
+    "tags": ["qwen3", "1.2b", "8k", "alibi", "dense", "seed42", "resume17000"]
+  },
+  "attention": {
+    "type": "dense",
+    "dense_mode": false,
+    "sparse_block_size": 128,
+    "sparse_num_random_blocks": 31,
+    "sparse_mask_pool_size": 1,
+    "sparse_val_dense_compare": false,
+    "dense_step_intervals": null,
+    "sparse_layer_indices": null,
+    "sparse_mask_cache_dir": null,
+    "sparse_mask_seed": 42
+  }
+}

epfl_final/SparseAtt/configs/alibi_20260921/1p2b_alibi_s42_bs15.json

@@ -0,0 +1,62 @@
+{
+  "model": {
+    "path": "/mloscratch/homes/navasard/qwen3_06b_base",
+    "dtype": "bfloat16",
+    "init_from_scratch": true,
+    "hidden_size": 1536,
+    "intermediate_size": 6144,
+    "num_hidden_layers": 28,
+    "num_attention_heads": 12,
+    "num_key_value_heads": 4,
+    "positional": "alibi"
+  },
+  "data": {
+    "train_path": "/mloscratch/homes/navasard/data/fineweb_large/tokenized",
+    "val_path": null,
+    "max_seq_length": 8192,
+    "num_workers": 8
+  },
+  "training": {
+    "output_dir": "/mloscratch/homes/navasard/sparse/models/ALIBI_20260921/1p2b_alibi_s42",
+    "resume_from_checkpoint": null,
+    "force_new_run": false,
+    "num_epochs": 1,
+    "max_steps": 50462,
+    "per_device_batch_size": 1,
+    "gradient_accumulation_steps": 15,
+    "learning_rate": 0.0005,
+    "weight_decay": 0.01,
+    "warmup_steps": 750,
+    "lr_scheduler": "cosine",
+    "max_grad_norm": 1.0,
+    "save_steps": 1000,
+    "eval_steps": 1000,
+    "eval_batches": 50,
+    "logging_steps": 100,
+    "seed": 42,
+    "verbose": true,
+    "checkpoint_activations": false
+  },
+  "distributed": {
+    "strategy": "ddp",
+    "mixed_precision": "bf16"
+  },
+  "wandb": {
+    "enabled": true,
+    "project": "stochastic_sparsity",
+    "name": "1.2b alibi 8k s42",
+    "tags": ["qwen3", "1.2b", "8k", "alibi", "dense", "seed42"]
+  },
+  "attention": {
+    "type": "dense",
+    "dense_mode": false,
+    "sparse_block_size": 128,
+    "sparse_num_random_blocks": 31,
+    "sparse_mask_pool_size": 1,
+    "sparse_val_dense_compare": false,
+    "dense_step_intervals": null,
+    "sparse_layer_indices": null,
+    "sparse_mask_cache_dir": null,
+    "sparse_mask_seed": 42
+  }
+}

epfl_final/SparseAtt/configs/alibi_20260921/1p2b_alibi_s42_bs20.json

@@ -0,0 +1,62 @@
+{
+  "model": {
+    "path": "/mloscratch/homes/navasard/qwen3_06b_base",
+    "dtype": "bfloat16",
+    "init_from_scratch": true,
+    "hidden_size": 1536,
+    "intermediate_size": 6144,
+    "num_hidden_layers": 28,
+    "num_attention_heads": 12,
+    "num_key_value_heads": 4,
+    "positional": "alibi"
+  },
+  "data": {
+    "train_path": "/mloscratch/homes/navasard/data/fineweb_large/tokenized",
+    "val_path": null,
+    "max_seq_length": 8192,
+    "num_workers": 8
+  },
+  "training": {
+    "output_dir": "/mloscratch/homes/navasard/sparse/models/ALIBI_20260921/1p2b_alibi_s42",
+    "resume_from_checkpoint": null,
+    "force_new_run": false,
+    "num_epochs": 1,
+    "max_steps": 50462,
+    "per_device_batch_size": 1,
+    "gradient_accumulation_steps": 20,
+    "learning_rate": 0.0005,
+    "weight_decay": 0.01,
+    "warmup_steps": 750,
+    "lr_scheduler": "cosine",
+    "max_grad_norm": 1.0,
+    "save_steps": 1000,
+    "eval_steps": 1000,
+    "eval_batches": 50,
+    "logging_steps": 100,
+    "seed": 42,
+    "verbose": true,
+    "checkpoint_activations": false
+  },
+  "distributed": {
+    "strategy": "ddp",
+    "mixed_precision": "bf16"
+  },
+  "wandb": {
+    "enabled": true,
+    "project": "stochastic_sparsity",
+    "name": "1.2b alibi 8k s42",
+    "tags": ["qwen3", "1.2b", "8k", "alibi", "dense", "seed42"]
+  },
+  "attention": {
+    "type": "dense",
+    "dense_mode": false,
+    "sparse_block_size": 128,
+    "sparse_num_random_blocks": 31,
+    "sparse_mask_pool_size": 1,
+    "sparse_val_dense_compare": false,
+    "dense_step_intervals": null,
+    "sparse_layer_indices": null,
+    "sparse_mask_cache_dir": null,
+    "sparse_mask_seed": 42
+  }
+}

epfl_final/SparseAtt/configs/alibi_20260921/365m_alibi_s42.json

@@ -0,0 +1,70 @@
+{
+  "model": {
+    "path": "/mloscratch/homes/navasard/qwen3_06b_base",
+    "dtype": "bfloat16",
+    "init_from_scratch": true,
+    "hidden_size": 768,
+    "intermediate_size": 3072,
+    "num_hidden_layers": 28,
+    "num_attention_heads": 6,
+    "num_key_value_heads": 2,
+    "positional": "alibi"
+  },
+  "data": {
+    "train_path": "/mloscratch/homes/navasard/data/fineweb_large/tokenized",
+    "val_path": null,
+    "max_seq_length": 8192,
+    "num_workers": 8
+  },
+  "training": {
+    "output_dir": "/mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42",
+    "resume_from_checkpoint": null,
+    "force_new_run": false,
+    "resume_micro_step_override": null,
+    "num_epochs": 1,
+    "max_steps": 44000,
+    "per_device_batch_size": 1,
+    "gradient_accumulation_steps": 5,
+    "learning_rate": 0.0005,
+    "weight_decay": 0.01,
+    "warmup_steps": 750,
+    "lr_scheduler": "cosine",
+    "max_grad_norm": 1.0,
+    "save_steps": 1000,
+    "eval_steps": 250,
+    "eval_batches": 50,
+    "logging_steps": 100,
+    "seed": 42,
+    "verbose": true,
+    "checkpoint_activations": false
+  },
+  "distributed": {
+    "strategy": "ddp",
+    "mixed_precision": "bf16"
+  },
+  "wandb": {
+    "enabled": true,
+    "project": "stochastic_sparsity",
+    "name": "365m alibi 8k s42",
+    "tags": [
+      "qwen3",
+      "365m",
+      "8k",
+      "scaleladder",
+      "alibi",
+      "seed42"
+    ]
+  },
+  "attention": {
+    "type": "dense",
+    "dense_mode": false,
+    "sparse_block_size": 128,
+    "sparse_num_random_blocks": 31,
+    "sparse_mask_pool_size": 1,
+    "sparse_val_dense_compare": false,
+    "dense_step_intervals": null,
+    "sparse_layer_indices": null,
+    "sparse_mask_cache_dir": null,
+    "sparse_mask_seed": 0
+  }
+}

epfl_final/SparseAtt/configs/alibi_20260921/365m_alibi_s43.json

@@ -0,0 +1,70 @@
+{
+  "model": {
+    "path": "/mloscratch/homes/navasard/qwen3_06b_base",
+    "dtype": "bfloat16",
+    "init_from_scratch": true,
+    "hidden_size": 768,
+    "intermediate_size": 3072,
+    "num_hidden_layers": 28,
+    "num_attention_heads": 6,
+    "num_key_value_heads": 2,
+    "positional": "alibi"
+  },
+  "data": {
+    "train_path": "/mloscratch/homes/navasard/data/fineweb_large/tokenized",
+    "val_path": null,
+    "max_seq_length": 8192,
+    "num_workers": 8
+  },
+  "training": {
+    "output_dir": "/mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43",
+    "resume_from_checkpoint": null,
+    "force_new_run": false,
+    "resume_micro_step_override": null,
+    "num_epochs": 1,
+    "max_steps": 44000,
+    "per_device_batch_size": 1,
+    "gradient_accumulation_steps": 4,
+    "learning_rate": 0.0005,
+    "weight_decay": 0.01,
+    "warmup_steps": 750,
+    "lr_scheduler": "cosine",
+    "max_grad_norm": 1.0,
+    "save_steps": 1000,
+    "eval_steps": 250,
+    "eval_batches": 50,
+    "logging_steps": 100,
+    "seed": 43,
+    "verbose": true,
+    "checkpoint_activations": false
+  },
+  "distributed": {
+    "strategy": "ddp",
+    "mixed_precision": "bf16"
+  },
+  "wandb": {
+    "enabled": true,
+    "project": "stochastic_sparsity",
+    "name": "365m alibi 8k s43",
+    "tags": [
+      "qwen3",
+      "365m",
+      "8k",
+      "scaleladder",
+      "alibi",
+      "seed43"
+    ]
+  },
+  "attention": {
+    "type": "dense",
+    "dense_mode": false,
+    "sparse_block_size": 128,
+    "sparse_num_random_blocks": 31,
+    "sparse_mask_pool_size": 1,
+    "sparse_val_dense_compare": false,
+    "dense_step_intervals": null,
+    "sparse_layer_indices": null,
+    "sparse_mask_cache_dir": null,
+    "sparse_mask_seed": 0
+  }
+}

epfl_final/SparseAtt/configs/alibi_20260921/365m_alibi_s44.json

@@ -0,0 +1,70 @@
+{
+  "model": {
+    "path": "/mloscratch/homes/navasard/qwen3_06b_base",
+    "dtype": "bfloat16",
+    "init_from_scratch": true,
+    "hidden_size": 768,
+    "intermediate_size": 3072,
+    "num_hidden_layers": 28,
+    "num_attention_heads": 6,
+    "num_key_value_heads": 2,
+    "positional": "alibi"
+  },
+  "data": {
+    "train_path": "/mloscratch/homes/navasard/data/fineweb_large/tokenized",
+    "val_path": null,
+    "max_seq_length": 8192,
+    "num_workers": 8
+  },
+  "training": {
+    "output_dir": "/mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44",
+    "resume_from_checkpoint": null,
+    "force_new_run": false,
+    "resume_micro_step_override": null,
+    "num_epochs": 1,
+    "max_steps": 44000,
+    "per_device_batch_size": 1,
+    "gradient_accumulation_steps": 4,
+    "learning_rate": 0.0005,
+    "weight_decay": 0.01,
+    "warmup_steps": 750,
+    "lr_scheduler": "cosine",
+    "max_grad_norm": 1.0,
+    "save_steps": 1000,
+    "eval_steps": 250,
+    "eval_batches": 50,
+    "logging_steps": 100,
+    "seed": 44,
+    "verbose": true,
+    "checkpoint_activations": false
+  },
+  "distributed": {
+    "strategy": "ddp",
+    "mixed_precision": "bf16"
+  },
+  "wandb": {
+    "enabled": true,
+    "project": "stochastic_sparsity",
+    "name": "365m alibi 8k s44",
+    "tags": [
+      "qwen3",
+      "365m",
+      "8k",
+      "scaleladder",
+      "alibi",
+      "seed44"
+    ]
+  },
+  "attention": {
+    "type": "dense",
+    "dense_mode": false,
+    "sparse_block_size": 128,
+    "sparse_num_random_blocks": 31,
+    "sparse_mask_pool_size": 1,
+    "sparse_val_dense_compare": false,
+    "dense_step_intervals": null,
+    "sparse_layer_indices": null,
+    "sparse_mask_cache_dir": null,
+    "sparse_mask_seed": 0
+  }
+}

epfl_final/SparseAtt/configs/alibi_20260921/365m_alibi_s45.json

@@ -0,0 +1,70 @@
+{
+  "model": {
+    "path": "/mloscratch/homes/navasard/qwen3_06b_base",
+    "dtype": "bfloat16",
+    "init_from_scratch": true,
+    "hidden_size": 768,
+    "intermediate_size": 3072,
+    "num_hidden_layers": 28,
+    "num_attention_heads": 6,
+    "num_key_value_heads": 2,
+    "positional": "alibi"
+  },
+  "data": {
+    "train_path": "/mloscratch/homes/navasard/data/fineweb_large/tokenized",
+    "val_path": null,
+    "max_seq_length": 8192,
+    "num_workers": 8
+  },
+  "training": {
+    "output_dir": "/mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s45",
+    "resume_from_checkpoint": null,
+    "force_new_run": false,
+    "resume_micro_step_override": null,
+    "num_epochs": 1,
+    "max_steps": 44000,
+    "per_device_batch_size": 1,
+    "gradient_accumulation_steps": 4,
+    "learning_rate": 0.0005,
+    "weight_decay": 0.01,
+    "warmup_steps": 750,
+    "lr_scheduler": "cosine",
+    "max_grad_norm": 1.0,
+    "save_steps": 1000,
+    "eval_steps": 250,
+    "eval_batches": 50,
+    "logging_steps": 100,
+    "seed": 45,
+    "verbose": true,
+    "checkpoint_activations": false
+  },
+  "distributed": {
+    "strategy": "ddp",
+    "mixed_precision": "bf16"
+  },
+  "wandb": {
+    "enabled": true,
+    "project": "stochastic_sparsity",
+    "name": "365m alibi 8k s45",
+    "tags": [
+      "qwen3",
+      "365m",
+      "8k",
+      "scaleladder",
+      "alibi",
+      "seed45"
+    ]
+  },
+  "attention": {
+    "type": "dense",
+    "dense_mode": false,
+    "sparse_block_size": 128,
+    "sparse_num_random_blocks": 31,
+    "sparse_mask_pool_size": 1,
+    "sparse_val_dense_compare": false,
+    "dense_step_intervals": null,
+    "sparse_layer_indices": null,
+    "sparse_mask_cache_dir": null,
+    "sparse_mask_seed": 0
+  }
+}

epfl_final/SparseAtt/configs/alibi_20260921/365m_alibi_s46.json

@@ -0,0 +1,70 @@
+{
+  "model": {
+    "path": "/mloscratch/homes/navasard/qwen3_06b_base",
+    "dtype": "bfloat16",
+    "init_from_scratch": true,
+    "hidden_size": 768,
+    "intermediate_size": 3072,
+    "num_hidden_layers": 28,
+    "num_attention_heads": 6,
+    "num_key_value_heads": 2,
+    "positional": "alibi"
+  },
+  "data": {
+    "train_path": "/mloscratch/homes/navasard/data/fineweb_large/tokenized",
+    "val_path": null,
+    "max_seq_length": 8192,
+    "num_workers": 8
+  },
+  "training": {
+    "output_dir": "/mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s46",
+    "resume_from_checkpoint": null,
+    "force_new_run": false,
+    "resume_micro_step_override": null,
+    "num_epochs": 1,
+    "max_steps": 44000,
+    "per_device_batch_size": 1,
+    "gradient_accumulation_steps": 5,
+    "learning_rate": 0.0005,
+    "weight_decay": 0.01,
+    "warmup_steps": 750,
+    "lr_scheduler": "cosine",
+    "max_grad_norm": 1.0,
+    "save_steps": 1000,
+    "eval_steps": 250,
+    "eval_batches": 50,
+    "logging_steps": 100,
+    "seed": 46,
+    "verbose": true,
+    "checkpoint_activations": false
+  },
+  "distributed": {
+    "strategy": "ddp",
+    "mixed_precision": "bf16"
+  },
+  "wandb": {
+    "enabled": true,
+    "project": "stochastic_sparsity",
+    "name": "365m alibi 8k s46",
+    "tags": [
+      "qwen3",
+      "365m",
+      "8k",
+      "scaleladder",
+      "alibi",
+      "seed46"
+    ]
+  },
+  "attention": {
+    "type": "dense",
+    "dense_mode": false,
+    "sparse_block_size": 128,
+    "sparse_num_random_blocks": 31,
+    "sparse_mask_pool_size": 1,
+    "sparse_val_dense_compare": false,
+    "dense_step_intervals": null,
+    "sparse_layer_indices": null,
+    "sparse_mask_cache_dir": null,
+    "sparse_mask_seed": 0
+  }
+}

epfl_final/SparseAtt/configs/alibi_20260921/365m_alibi_s47.json

@@ -0,0 +1,69 @@
+{
+  "model": {
+    "path": "/mloscratch/homes/navasard/qwen3_06b_base",
+    "dtype": "bfloat16",
+    "init_from_scratch": true,
+    "hidden_size": 768,
+    "intermediate_size": 3072,
+    "num_hidden_layers": 28,
+    "num_attention_heads": 6,
+    "num_key_value_heads": 2,
+    "positional": "alibi"
+  },
+  "data": {
+    "train_path": "/mloscratch/homes/navasard/data/fineweb_large/tokenized",
+    "val_path": null,
+    "max_seq_length": 8192,
+    "num_workers": 8
+  },
+  "training": {
+    "output_dir": "/mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s47",
+    "resume_from_checkpoint": null,
+    "force_new_run": false,
+    "num_epochs": 1,
+    "max_steps": 44000,
+    "per_device_batch_size": 2,
+    "gradient_accumulation_steps": 10,
+    "learning_rate": 0.0005,
+    "weight_decay": 0.01,
+    "warmup_steps": 750,
+    "lr_scheduler": "cosine",
+    "max_grad_norm": 1.0,
+    "save_steps": 1000,
+    "eval_steps": 250,
+    "eval_batches": 50,
+    "logging_steps": 100,
+    "seed": 47,
+    "verbose": true,
+    "checkpoint_activations": false
+  },
+  "distributed": {
+    "strategy": "ddp",
+    "mixed_precision": "bf16"
+  },
+  "wandb": {
+    "enabled": true,
+    "project": "stochastic_sparsity",
+    "name": "365m alibi 8k s47",
+    "tags": [
+      "qwen3",
+      "365m",
+      "8k",
+      "scaleladder",
+      "alibi",
+      "seed47"
+    ]
+  },
+  "attention": {
+    "type": "dense",
+    "dense_mode": false,
+    "sparse_block_size": 128,
+    "sparse_num_random_blocks": 31,
+    "sparse_mask_pool_size": 1,
+    "sparse_val_dense_compare": false,
+    "dense_step_intervals": null,
+    "sparse_layer_indices": null,
+    "sparse_mask_cache_dir": null,
+    "sparse_mask_seed": 0
+  }
+}

epfl_final/SparseAtt/configs/alibi_20260921/365m_alibi_s47_restart_20260922.json

@@ -0,0 +1,69 @@
+{
+  "model": {
+    "path": "/mloscratch/homes/navasard/qwen3_06b_base",
+    "dtype": "bfloat16",
+    "init_from_scratch": true,
+    "hidden_size": 768,
+    "intermediate_size": 3072,
+    "num_hidden_layers": 28,
+    "num_attention_heads": 6,
+    "num_key_value_heads": 2,
+    "positional": "alibi"
+  },
+  "data": {
+    "train_path": "/mloscratch/homes/navasard/data/fineweb_large/tokenized",
+    "val_path": null,
+    "max_seq_length": 8192,
+    "num_workers": 8
+  },
+  "training": {
+    "output_dir": "/mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s47_restart_20260922",
+    "resume_from_checkpoint": null,
+    "force_new_run": false,
+    "num_epochs": 1,
+    "max_steps": 44000,
+    "per_device_batch_size": 1,
+    "gradient_accumulation_steps": 4,
+    "learning_rate": 0.0005,
+    "weight_decay": 0.01,
+    "warmup_steps": 750,
+    "lr_scheduler": "cosine",
+    "max_grad_norm": 1.0,
+    "save_steps": 1000,
+    "eval_steps": 250,
+    "eval_batches": 50,
+    "logging_steps": 100,
+    "seed": 47,
+    "verbose": true,
+    "checkpoint_activations": false
+  },
+  "distributed": {
+    "strategy": "ddp",
+    "mixed_precision": "bf16"
+  },
+  "wandb": {
+    "enabled": true,
+    "project": "stochastic_sparsity",
+    "name": "365m alibi 8k s47 restart 20260922",
+    "tags": [
+      "qwen3",
+      "365m",
+      "8k",
+      "scaleladder",
+      "alibi",
+      "seed47"
+    ]
+  },
+  "attention": {
+    "type": "dense",
+    "dense_mode": false,
+    "sparse_block_size": 128,
+    "sparse_num_random_blocks": 31,
+    "sparse_mask_pool_size": 1,
+    "sparse_val_dense_compare": false,
+    "dense_step_intervals": null,
+    "sparse_layer_indices": null,
+    "sparse_mask_cache_dir": null,
+    "sparse_mask_seed": 0
+  }
+}

epfl_final/SparseAtt/configs/alibi_20260921/365m_alibi_s47_resume_2xh100.json

@@ -0,0 +1,69 @@
+{
+  "model": {
+    "path": "/mloscratch/homes/navasard/qwen3_06b_base",
+    "dtype": "bfloat16",
+    "init_from_scratch": true,
+    "hidden_size": 768,
+    "intermediate_size": 3072,
+    "num_hidden_layers": 28,
+    "num_attention_heads": 6,
+    "num_key_value_heads": 2,
+    "positional": "alibi"
+  },
+  "data": {
+    "train_path": "/mloscratch/homes/navasard/data/fineweb_large/tokenized",
+    "val_path": null,
+    "max_seq_length": 8192,
+    "num_workers": 8
+  },
+  "training": {
+    "output_dir": "/mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s47_restart_20260922",
+    "resume_from_checkpoint": "/mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s47_restart_20260922/run_20260922_192827/checkpoint-41000",
+    "force_new_run": false,
+    "num_epochs": 1,
+    "max_steps": 44000,
+    "per_device_batch_size": 2,
+    "gradient_accumulation_steps": 5,
+    "learning_rate": 0.0005,
+    "weight_decay": 0.01,
+    "warmup_steps": 750,
+    "lr_scheduler": "cosine",
+    "max_grad_norm": 1.0,
+    "save_steps": 1000,
+    "eval_steps": 250,
+    "eval_batches": 50,
+    "logging_steps": 100,
+    "seed": 47,
+    "verbose": true,
+    "checkpoint_activations": false
+  },
+  "distributed": {
+    "strategy": "ddp",
+    "mixed_precision": "bf16"
+  },
+  "wandb": {
+    "enabled": false,
+    "project": "stochastic_sparsity",
+    "name": "365m alibi 8k s47 restart 20260922",
+    "tags": [
+      "qwen3",
+      "365m",
+      "8k",
+      "scaleladder",
+      "alibi",
+      "seed47"
+    ]
+  },
+  "attention": {
+    "type": "dense",
+    "dense_mode": false,
+    "sparse_block_size": 128,
+    "sparse_num_random_blocks": 31,
+    "sparse_mask_pool_size": 1,
+    "sparse_val_dense_compare": false,
+    "dense_step_intervals": null,
+    "sparse_layer_indices": null,
+    "sparse_mask_cache_dir": null,
+    "sparse_mask_seed": 0
+  }
+}

epfl_final/SparseAtt/finalize_tau1_fix.py

@@ -0,0 +1,95 @@
+#!/usr/bin/env python3
+import csv, hashlib, json, math, subprocess, time
+from pathlib import Path
+
+B=Path('/mloscratch/homes/navasard/sparse'); R=B/'SparseAtt'; SRC=R/'eval/results/RESOLUTION_20260921'; OUT=SRC/'tau1_fix'
+REV=subprocess.check_output(['git','rev-parse','HEAD'],cwd=R,text=True).strip()
+REF={
+('365m',42,'dense',24576):(3.65148,3.29270,4.49264),('365m',42,'mixed',24576):(3.60371,3.15984,4.41666),
+('365m',42,'dense',30720):(3.76818,3.56380,4.78925),('365m',42,'mixed',30720):(3.74426,3.36130,4.68366),
+('1p2b',42,'dense',24576):(3.28561,2.96720,4.16685),('1p2b',42,'mixed',24576):(3.21926,2.80155,4.07777),
+('1p2b',42,'dense',30720):(3.48349,3.36662,4.56361),('1p2b',42,'mixed',30720):(3.36643,3.10652,4.43694),
+('1p2b',70,'dense',24576):(3.32724,3.05613,4.18805),('1p2b',70,'mixed',24576):(3.10813,2.74541,4.09636),
+('1p2b',70,'dense',30720):(3.75455,3.61592,4.71141),('1p2b',70,'mixed',30720):(3.30011,3.14062,4.50743),
+('1p2b',71,'dense',30720):(3.53217,3.49435,4.70393),('1p2b',71,'mixed',30720):(3.31447,3.24131,4.45906),
+('1p2b',72,'dense',24576):(2.98509,3.01634,4.24704),('1p2b',72,'mixed',24576):(3.19117,2.76665,4.08030),
+('1p2b',72,'dense',30720):(3.46737,3.53236,4.61552),('1p2b',72,'mixed',30720):(3.40891,3.19377,4.49819),
+('3b',42,'dense',24576):(3.13416,2.72716,4.06614),('3b',42,'mixed',24576):(2.96409,2.55597,3.87343),
+('3b',42,'dense',30720):(3.38926,3.14997,4.60384),('3b',42,'mixed',30720):(3.14371,2.85993,4.24324)}
+
+# Canonical tau=1 held-out means come from the original evaluator outputs,
+# not the resolution brief's tau* table.  The closure campaigns provide the
+# seed-70/72 1.2B cells; seed-71 30k and seed-72 dense are retained from the
+# compendium because no canonical temperature TSV was archived for them.
+CANONICAL_FILES = {
+    ('365m',42,'dense'): 'TEMPERATURE_PROTOCOL_20260906/365m_dense_s42/L{L}_dense.tsv',
+    ('365m',42,'mixed'): 'TEMPERATURE_PROTOCOL_20260906/365m_mixed_s42/L{L}_dense.tsv',
+    ('1p2b',42,'dense'): 'TEMPERATURE_PROTOCOL_20260906/1p2b_dense_s42/L{L}_dense.tsv',
+    ('1p2b',42,'mixed'): 'TEMPERATURE_PROTOCOL_20260906/1p2b_mixed_s42/L{L}_dense.tsv',
+    ('3b',42,'dense'): 'TEMPERATURE_PROTOCOL_20260906/3b_dense/L{L}_dense.tsv',
+    ('3b',42,'mixed'): 'TEMPERATURE_PROTOCOL_20260906/3b_mixed/L{L}_dense.tsv',
+    ('1p2b',70,'dense'): 'REVIEWER_CLOSURES_20260915/temperature/1p2b_dense_s70_L{L}_dense.tsv',
+    ('1p2b',70,'mixed'): 'GAP_CLOSURE_20260915/temperature/1p2b_mixed_s70_L{L}_dense.tsv',
+    ('1p2b',72,'mixed'): 'GAP_CLOSURE_20260915/temperature/1p2b_mixed_s72_L{L}_dense.tsv',
+}
+CANONICAL_OVERRIDES = {
+    # The archive has no dense-72 temperature TSV; these are the independent
+    # canonical-path scores on its closure held-out offsets 2L and 3L.
+    ('1p2b',72,'dense',24576): (3.128708458839331, 2.945860953345546, 4.109781368717415),
+}
+
+def sha(p):
+ h=hashlib.sha256();
+ with p.open('rb') as f:
+  for b in iter(lambda:f.read(1024*1024),b''): h.update(b)
+ return h.hexdigest()
+def means(path, split, tau):
+ d={}
+ with path.open(newline='') as h:
+  for r in csv.DictReader(h,delimiter='\t'):
+   if r['split']==split and float(r['tau'])==tau: d.setdefault(r['corpus'],[]).append(float(r['avg_nll']))
+ return {c:sum(v)/len(v) for c,v in d.items()}
+
+def canonical(key):
+ if key in CANONICAL_OVERRIDES:
+  return dict(zip(('arxiv','fineweb','pg19'), CANONICAL_OVERRIDES[key]))
+ tier,seed,arm,L=key
+ rel=CANONICAL_FILES.get((tier,seed,arm))
+ if rel:
+  p=R/'eval/results'/rel.format(L=L)
+  if p.exists():
+   m=means(p,'heldout',1.0)
+   if len(m)==3:return m
+ return dict(zip(('arxiv','fineweb','pg19'), REF[key]))
+
+def main():
+ src_files=sorted(SRC.glob('*_heldout.tsv'))
+ out_files=sorted(OUT.glob('*_heldout.tsv'))
+ if len(out_files)!=24: raise SystemExit(f'expected 24 corrected TSVs, found {len(out_files)}')
+ rows=[]
+ for key,ref in REF.items():
+  tier,seed,arm,L=key; sf=SRC/f'{tier}_{arm}_s{seed}_L{L}_heldout.tsv'; of=OUT/f'{tier}_{arm}_s{seed}_L{L}_heldout.tsv'
+  refm=canonical(key)
+  bm=means(sf,'heldout_existing',1.0); am=means(of,'heldout_existing',1.0)
+  bef=[None if c not in bm else math.exp(bm[c]-refm[c])-1 for c in ('arxiv','fineweb','pg19')]
+  aft=[None if c not in am else math.exp(am[c]-refm[c])-1 for c in ('arxiv','fineweb','pg19')]
+  rows.append((tier,seed,arm,L,bef,aft,all(x is not None and abs(x)<=.001 for x in aft)))
+ drifts={}
+ compares={}
+ for p in OUT.glob('pristine_drift_shard*.json'): drifts.update(json.loads(p.read_text()))
+ for p in OUT.glob('tau_star_compare_shard*.json'): compares.update(json.loads(p.read_text()))
+ lines=['# Resolution τ=1 fix 20260921','',f'Git: `{REV}`.',f'Wall-clock is recorded by the shard logs.', '', '## Cause','',
+ 'The original resolution driver applies temperature cumulatively in place: `temp(m, q/cur); cur=q` followed by `q_norm.weight.div_(t)`. It does not restore q-norm from a pristine value before each τ.',
+ 'The canonical helper stores `_pristine_qnorm` and uses `q_norm.weight.copy_(_pristine_qnorm / tau)` for every τ. Thus τ=1 after prior temperatures was not the canonical model state.',
+ 'For the 1.2B seed-70/72 24k closure cells, the closure evaluator used held-out offsets 2L and 3L. The first draft used the resolution driver\'s 3L/4L/5L set; those rows were replaced with the canonical closure-window scores before this note was regenerated.',
+ 'Both evaluators use the dense path: `load_model_for_deploy` sets `config.attention_type = "dense"`; dense attention calls `F.scaled_dot_product_attention`, while only the sparse branch compiles FlexAttention. There is no explicit 3B/30k backend branch; its agreement is not evidence of a separate τ=1 implementation.',
+ '', 'Quoted implementation lines:', '', '```text', 'resolution driver: temp(m, q/cur); cur=q', 'resolution driver: for l in m.model.layers: l.self_attn.q_norm.weight.div_(t)', 'canonical helper: attn._pristine_qnorm = attn.q_norm.weight.detach().clone()', 'canonical helper: attn.q_norm.weight.copy_(attn._pristine_qnorm / tau)', '```', '', 'Maximum q-norm drift before corrected τ=1 scoring:', '', '```json',json.dumps(drifts,indent=2,sort_keys=True),'```','', '## Acceptance','', 'Relative perplexity error is `exp(mean_new_nll - canonical_nll) - 1`; PASS requires absolute error ≤0.1%.','', '| Tier | Seed | Arm | Length | Before arXiv/FineWeb/PG-19 | After arXiv/FineWeb/PG-19 | Result |','|---|---:|---|---:|---|---|---|']
+ for tier,seed,arm,L,b,a,p in rows:
+  fmt=lambda xs:', '.join('NA' if x is None else f'{100*x:.3f}%' for x in xs)
+  lines.append(f'| {tier} | {seed} | {arm} | {L} | {fmt(b)} | {fmt(a)} | {"PASS" if p else "FAIL"} |')
+ lines += ['', '## τ* path check','', '365M Dense seed42 at 30k was re-scored at τ*=0.85 with the canonical helper; these are mean NLL values on the existing windows.', '', '```json',json.dumps(compares,indent=2,sort_keys=True),'```','', '## Result files','', '| File | SHA-256 |','|---|---|']
+ for p in out_files: lines.append(f'| `{p.name}` | `{sha(p)}` |')
+ lines += ['',f'Total result TSVs: {len(out_files)}.']
+ (OUT/'SPARSEATT_RESOLUTION_TAU1_FIX_20260921.md').write_text('\n'.join(lines)+'\n')
+ print(json.dumps({'result_files':len(out_files),'all_pass':all(r[-1] for r in rows),'note_sha256':sha(OUT/'SPARSEATT_RESOLUTION_TAU1_FIX_20260921.md')},indent=2))
+if __name__=='__main__': main()

epfl_final/SparseAtt/launch_1p2b_alibi_after_eval.sh

@@ -0,0 +1,23 @@
+#!/bin/sh
+set -eu
+
+R=/mloscratch/homes/navasard/sparse/SparseAtt
+EVAL_PID="${1:?evaluation PID required}"
+LOG="$R/../logs/ALIBI_20260921/1p2b_s42_deferred.log"
+CFG="$R/configs/alibi_20260921/1p2b_alibi_s42.json"
+mkdir -p "$(dirname "$LOG")"
+{
+  echo "$(date -Is) waiting for evaluation PID $EVAL_PID"
+  while :; do
+    stat=$(ps -p "$EVAL_PID" -o stat= 2>/dev/null || true)
+    [ -z "$stat" ] && break
+    case "$stat" in Z*) break ;; esac
+    sleep 30
+  done
+  echo "$(date -Is) evaluation exited; waiting 300 seconds"
+  sleep 300
+  echo "$(date -Is) launching 1.2B ALiBi seed 42 on 6 GPUs"
+  cd "$R"
+  export USER=navasard LOGNAME=navasard PYTHONPATH="$R"
+  exec /mloscratch/homes/navasard/conda/bin/torchrun --standalone --nproc_per_node=6 train.py --config "$CFG"
+} >>"$LOG" 2>&1

epfl_final/SparseAtt/launch_alibi_20260921.sh

@@ -0,0 +1,33 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+REPO=/mloscratch/homes/navasard/sparse/SparseAtt
+RUN_TAG=${ALIBI_RUN_TAG:-restart_$(date -u +%Y%m%d_%H%M%S)}
+LOGDIR=/mloscratch/homes/navasard/sparse/logs/ALIBI_20260921/${RUN_TAG}
+PYTHON=/mloscratch/homes/navasard/conda/bin/python3
+export USER=navasard
+export HOME=/mloscratch/homes/navasard
+export HF_HOME=/mloscratch/homes/navasard/hf_cache
+export WANDB_MODE=offline
+mkdir -p "$LOGDIR"
+
+cd "$REPO"
+for seed in 42 43 44 45 46 47; do
+  cfg="$REPO/configs/alibi_20260921/365m_alibi_s${seed}.json"
+  test -f "$cfg"
+done
+
+for seed in 42 43 44 45 46 47; do
+  gpu=$((seed - 42))
+  cfg="$REPO/configs/alibi_20260921/365m_alibi_s${seed}.json"
+  log="$LOGDIR/seed${seed}.log"
+  echo "launch seed=${seed} gpu=${gpu} config=${cfg} log=${log}"
+  CUDA_VISIBLE_DEVICES="$gpu" \
+    TORCHINDUCTOR_CACHE_DIR="/mloscratch/homes/navasard/sparse/torch_cache/alibi_s${seed}" \
+    TRITON_CACHE_DIR="/mloscratch/homes/navasard/sparse/torch_cache/triton_s${seed}" \
+    "$PYTHON" -u train.py --config "$cfg" >"$log" 2>&1 &
+  echo $! > "$LOGDIR/seed${seed}.pid"
+done
+
+echo "launched $(date -Is)"
+wait

epfl_final/SparseAtt/merge_closure24.py

@@ -0,0 +1,28 @@
+import csv, os
+from pathlib import Path
+
+ROOT = Path('/mloscratch/homes/navasard/sparse/SparseAtt/eval/results/RESOLUTION_20260921/tau1_fix')
+affected = [('dense',70),('mixed',70),('dense',72),('mixed',72)]
+
+def main():
+    rows = {}
+    with (ROOT/'closure24_rows.tsv').open(newline='') as f:
+        for r in csv.DictReader(f, delimiter='\t'):
+            rows[(r['arm'], int(r['seed']), r['corpus'])] = r
+    for arm, seed in affected:
+        p = ROOT/f'1p2b_{arm}_s{seed}_L24576_heldout.tsv'
+        with p.open(newline='') as f: old = list(csv.DictReader(f, delimiter='\t'))
+        keep = [r for r in old if not (r['split']=='heldout_existing' and float(r['tau'])==1.0)]
+        for c in ('arxiv','fineweb','pg19'):
+            keep.append(dict(rows[(arm,seed,c)]))
+        fields = list(old[0])
+        for r in keep:
+            r['tier']='1p2b'; r['arm']=arm; r['seed']=str(seed); r['length']='24576'; r['git_rev']=old[0].get('git_rev','')
+            r['window_id'] = r.get('window_id', 'closure0')
+        tmp = p.with_suffix('.tmp')
+        with tmp.open('w', newline='') as f:
+            w=csv.DictWriter(f, fieldnames=fields, delimiter='\t'); w.writeheader(); w.writerows(keep)
+        os.replace(tmp,p)
+        print('merged',p)
+    (ROOT/'closure24_rows.tsv').unlink()
+if __name__=='__main__': main()

epfl_final/SparseAtt/rebuild_closure24_from_log.py

@@ -0,0 +1,26 @@
+import csv, math, os, re
+from pathlib import Path
+
+ROOT=Path('/mloscratch/homes/navasard/sparse/SparseAtt/eval/results/RESOLUTION_20260921/tau1_fix')
+LOG=ROOT/'closure24_eval.log'
+pat=re.compile(r'^(dense|mixed) (70|72) (fineweb|pg19|arxiv) ([01]) ([0-9.eE+-]+)$')
+
+def main():
+    vals=[]
+    for line in LOG.read_text().splitlines():
+        m=pat.match(line.strip())
+        if m:
+            arm,seed,c,j,nll=m.groups(); seed=int(seed); j=int(j); nll=float(nll)
+            vals.append(dict(tier='1p2b',arm=arm,seed=str(seed),length='24576',split='heldout_existing',corpus=c,window_id=f'closure{j}',start_offset=str((2+j)*24576),tokens='24575',tau='1.0',avg_nll=str(nll),perplexity=str(math.exp(nll)),git_rev='6ea672d1163c84c6d6c74c4dda9cd6073e49856e'))
+    if len(vals)!=24: raise SystemExit(f'expected 24 parsed rows, got {len(vals)}')
+    for arm,seed in [('dense',70),('mixed',70),('dense',72),('mixed',72)]:
+        p=ROOT/f'1p2b_{arm}_s{seed}_L24576_heldout.tsv'
+        with p.open(newline='') as f: old=list(csv.DictReader(f,delimiter='\t'))
+        keep=[r for r in old if not (r['split']=='heldout_existing' and float(r['tau'])==1.0)]
+        keep.extend(r for r in vals if r['arm']==arm and int(r['seed'])==seed)
+        fields=list(old[0]); tmp=p.with_suffix('.tmp')
+        with tmp.open('w',newline='') as f:
+            w=csv.DictWriter(f,fieldnames=fields,delimiter='\t'); w.writeheader(); w.writerows(keep)
+        os.replace(tmp,p); print('rebuilt',p)
+    LOG.unlink(missing_ok=True)
+if __name__=='__main__': main()

epfl_final/SparseAtt/sparseatt/asentmax_attention.py

@@ -0,0 +1,91 @@
+"""
+ASEntmax baseline — arm #3 of the elasticity study.
+
+Adaptive-Scalable Entmax: alpha-entmax attention (exact zeros -> content-adaptive
+token-level sparsity) with a learnable per-head temperature that interpolates the
+distribution between sparse and dense regimes. This is NOT block-sparse and has no
+dial-able block-density budget; it is included as a *trained-sparse* alternative,
+chiefly for the length-generalization comparison.
+
+Dense manual path (entmax cannot be fused into SDPA/FlexAttention).
+
+Ref: Long-Context Generalization with Sparse Attention (ASEntmax, arXiv:2506.16640,
+ICLR 2026).
+"""
+from typing import Optional, Tuple
+import math
+
+import torch
+import torch.nn as nn
+
+from .attention import Qwen3Attention, apply_rotary_pos_emb
+
+try:
+    from entmax import entmax15, entmax_bisect
+except Exception:  # pragma: no cover
+    entmax15 = None
+    entmax_bisect = None
+
+
+class Qwen3ASEntmaxAttention(Qwen3Attention):
+    """alpha-entmax attention with a learnable temperature (ASEntmax)."""
+
+    def __init__(self, config, layer_idx: int):
+        super().__init__(config, layer_idx)
+        if entmax15 is None:
+            raise ImportError("ASEntmax requires the `entmax` package (pip install entmax)")
+        self.alpha = float(getattr(config, "asentmax_alpha", 1.5))
+        learn = bool(getattr(config, "asentmax_learn_temp", True))
+        init_temp = float(getattr(config, "asentmax_init_temp", 1.0))
+        log_t = torch.full((self.num_heads, 1, 1), math.log(init_temp), dtype=torch.float32)
+        if learn:
+            self.log_temp = nn.Parameter(log_t)
+        else:
+            self.register_buffer("log_temp", log_t, persistent=False)
+        self.dense_mode = False
+
+    def _entmax(self, scores: torch.Tensor) -> torch.Tensor:
+        if abs(self.alpha - 1.5) < 1e-6:
+            return entmax15(scores, dim=-1)
+        return entmax_bisect(scores, alpha=self.alpha, dim=-1)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        use_cache: bool = False,
+        output_attentions: bool = False,
+    ):
+        B, T, _ = hidden_states.size()
+
+        q = self.q_proj(hidden_states).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(hidden_states).view(B, T, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(hidden_states).view(B, T, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+        if self.qk_norm == "on":
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+        if self.qk_norm == "off" or self._use_query_temperature:
+            q = q * self._query_temperature_scale.to(dtype=q.dtype)
+        cos, sin = self.rotary_emb(v, position_ids)
+        q, k = apply_rotary_pos_emb(q, k, cos, sin)
+
+        k = self._repeat_kv(k, self.num_key_value_groups)
+        v = self._repeat_kv(v, self.num_key_value_groups)
+
+        scale = 1.0 / math.sqrt(self.head_dim)
+        scores = torch.matmul(q, k.transpose(-2, -1)) * scale        # [B,H,T,T]
+        temp = torch.exp(self.log_temp)                              # [H,1,1]
+        scores = scores * temp                                      # learnable temperature
+
+        qpos = torch.arange(T, device=q.device)
+        causal = qpos[None, :] <= qpos[:, None]                     # [T,T]
+        # large finite negative (not -inf) so entmax stays numerically safe
+        scores = scores.float().masked_fill(~causal[None, None], -1e9)
+
+        p = self._entmax(scores).to(v.dtype)                        # exact zeros -> sparse
+        out = torch.matmul(p, v)
+        out = out.transpose(1, 2).contiguous().reshape(B, T, self.attention_hidden_size)
+        return self.o_proj(out), None, None

epfl_final/SparseAtt/sparseatt/attention.py

@@ -0,0 +1,338 @@
+"""
+Qwen3 Attention Module.
+
+This file is kept separate for easy modification during token pruning experiments.
+"""
+
+from typing import Optional, Tuple
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+try:
+    from torch.nn.attention.flex_attention import create_block_mask, flex_attention
+except ImportError:
+    create_block_mask = flex_attention = None
+
+from .config import Qwen3Config
+
+
+def alibi_slopes(num_heads: int) -> torch.Tensor:
+    """Press et al. slopes: 1/2, 1/4, ..., 2**(-8) for eight heads."""
+    return torch.pow(2.0, -8.0 * torch.arange(1, num_heads + 1, dtype=torch.float32) / num_heads)
+
+
+def _alibi_causal_mask(b, h, q_idx, kv_idx):
+    return q_idx >= kv_idx
+
+
+def _alibi_score_mod(num_heads: int):
+    # q/k indices and h are scalar symbolic tensors under FlexAttention.  This
+    # produces no LxL bias tensor; tau scales Q only and therefore never this term.
+    def score_mod(score, b, h, q_idx, kv_idx):
+        slope = 2.0 ** (-8.0 * (h + 1) / num_heads)
+        return score + slope * (kv_idx - q_idx)
+    return score_mod
+
+
+class Qwen3RMSNorm(nn.Module):
+    """RMS Normalization used in Qwen3."""
+    
+    def __init__(self, hidden_size: int, eps: float = 1e-6, weight_fp32: bool = False):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.eps = eps
+        self.weight_fp32 = weight_fp32
+        self._eval_temperature: Optional[float] = None
+
+    def set_eval_temperature(self, tau: float) -> None:
+        """Apply an eval-only Q gain scale before its activation-dtype cast."""
+        self._eval_temperature = float(tau)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
+        # QK gains can be deliberate FP32 master parameters while all model
+        # activations stay bf16. Casting here preserves fp32 optimizer updates
+        # without promoting the attention activations.
+        weight = self.weight if self._eval_temperature is None else self.weight / self._eval_temperature
+        return weight.to(input_dtype) * hidden_states.to(input_dtype)
+
+
+class Qwen3RotaryEmbedding(nn.Module):
+    """Rotary Position Embedding for Qwen3."""
+    
+    def __init__(
+        self,
+        dim: int,
+        max_position_embeddings: int = 32768,
+        base: float = 1000000.0,
+        train_position_scale_set: Optional[list[float]] = None,
+        device: Optional[torch.device] = None,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.max_position_embeddings = max_position_embeddings
+        self.base = base
+
+        raw_scales = train_position_scale_set or []
+        scales = [float(scale) for scale in raw_scales]
+        if any(not torch.isfinite(torch.tensor(scale)) or scale <= 0.0 for scale in scales):
+            raise ValueError("rope_train_position_scale_set must contain finite positive values")
+        self.register_buffer(
+            "train_position_scales",
+            torch.tensor(scales, dtype=torch.float32, device=device),
+            persistent=False,
+        )
+        
+        inv_freq = 1.0 / (
+            self.base ** (torch.arange(0, self.dim, 2, dtype=torch.float32, device=device) / self.dim)
+        )
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+    @torch.no_grad()
+    def forward(
+        self, x: torch.Tensor, position_ids: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        # Mixed sparse training resamples a mask (and therefore its realised
+        # key-distance distribution) independently in every layer.  This dense
+        # control mirrors only that positional exposure: each layer samples one
+        # scalar for the whole micro-batch, while all attention links remain
+        # present.  Eval is deliberately untouched.
+        if self.training and self.train_position_scales.numel() > 0:
+            scale_idx = torch.randint(
+                self.train_position_scales.numel(),
+                (1,),
+                device=position_ids.device,
+            )
+            position_ids = position_ids.float() * self.train_position_scales[scale_idx]
+        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
+        position_ids_expanded = position_ids[:, None, :].float()
+        
+        device_type = x.device.type
+        with torch.autocast(device_type=device_type, enabled=False):
+            freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            cos = emb.cos()
+            sin = emb.sin()
+        
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+def rotate_half(x: torch.Tensor) -> torch.Tensor:
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    unsqueeze_dim: int = 1,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Applies Rotary Position Embedding to the query and key tensors."""
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+class Qwen3Attention(nn.Module):
+    """Multi-head attention with Grouped Query Attention (GQA) support."""
+    
+    def __init__(self, config: Qwen3Config, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = config.head_dim
+        self.num_key_value_heads = config.num_key_value_heads
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.attention_dropout = config.attention_dropout
+        self.positional = config.positional
+        self.qk_norm = str(getattr(config, "qk_norm", "on")).lower()
+        self.qk_gain_fp32 = bool(getattr(config, "qk_gain_fp32", False))
+        if self.qk_norm not in ("on", "off"):
+            raise ValueError(f"unsupported qk_norm mode: {self.qk_norm}")
+        if self.positional not in ("rope", "alibi"):
+            raise ValueError(f"unsupported positional mode: {self.positional}")
+        
+        # Attention output dimension is num_heads * head_dim
+        self.attention_hidden_size = self.num_heads * self.head_dim
+        
+        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias)
+        self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)
+        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.attention_bias)
+        
+        self.rotary_emb = Qwen3RotaryEmbedding(
+            self.head_dim,
+            max_position_embeddings=config.max_position_embeddings,
+            base=config.rope_theta,
+            train_position_scale_set=config.rope_train_position_scale_set,
+        )
+        self.register_buffer("alibi_slopes", alibi_slopes(self.num_heads), persistent=False)
+        self._alibi_masks = {}
+        self._alibi_score_mod = _alibi_score_mod(self.num_heads) if self.positional == "alibi" else None
+        
+        # For Qwen3, q_norm and k_norm are per layer.  The explicit ``off``
+        # regime removes the modules entirely, rather than leaving identity
+        # modules in the state dict.  The default remains the legacy path.
+        if self.qk_norm == "on":
+            self.q_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps, weight_fp32=self.qk_gain_fp32)
+            self.k_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps, weight_fp32=self.qk_gain_fp32)
+        # Evaluation-only temperature for qk_norm=off.  It is a non-persistent
+        # buffer so no checkpoint weights are touched and compiled graphs can
+        # reuse the same shape/static graph across temperatures.
+        self.register_buffer("_query_temperature_scale", torch.ones((), dtype=torch.float32), persistent=False)
+        # Eval-only diagnostic. When positive, sample this many causal query
+        # positions and retain head-averaged attention entropy. The default is
+        # disabled and therefore leaves training/evaluation behavior unchanged.
+        self.entropy_probe_queries = 0
+        self.last_entropy_probe = None
+        self._use_query_temperature = False
+
+    @torch.no_grad()
+    def set_attention_temperature(self, tau: float) -> None:
+        """Set the query-side temperature without changing learned weights.
+
+        This is the temperature path for qk_norm=off and is also exposed for
+        equivalence checks on qk_norm=on checkpoints.
+        """
+        tau = float(tau)
+        if not math.isfinite(tau) or tau <= 0.0:
+            raise ValueError("attention temperature must be finite and > 0")
+        if self.qk_norm == "on":
+            # This is algebraically query-side scaling, but applying it before
+            # the bf16 activation cast matches the learned-gain division path
+            # exactly without mutating a checkpoint parameter.
+            self.q_norm.set_eval_temperature(tau)
+            self._use_query_temperature = False
+        else:
+            self._query_temperature_scale.fill_(1.0 / tau)
+            self._use_query_temperature = True
+
+    @torch.no_grad()
+    def _measure_attention_entropy(self, query_states, key_states):
+        seq_len = query_states.shape[2]
+        count = min(int(self.entropy_probe_queries), seq_len)
+        # Include short and long prefixes while avoiding the degenerate first token.
+        positions = torch.linspace(1, seq_len - 1, count, device=query_states.device).round().long().unique()
+        kt = key_states.transpose(-1, -2)
+        records = []
+        for chunk in positions.split(4):
+            q = query_states.index_select(2, chunk)
+            scores = torch.matmul(q, kt) / math.sqrt(self.head_dim)
+            keys = torch.arange(seq_len, device=scores.device)
+            allowed = keys[None, None, None, :] <= chunk[None, None, :, None]
+            scores = scores.float().masked_fill(~allowed, float('-inf'))
+            logz = torch.logsumexp(scores, dim=-1)
+            probs = torch.softmax(scores, dim=-1)
+            entropy = logz - torch.sum(probs * scores.masked_fill(~allowed, 0.0), dim=-1)
+            mean_heads = entropy.mean(dim=(0, 1)).cpu().tolist()
+            records.extend((int(pos), int(pos)+1, float(ent)) for pos,ent in zip(chunk.cpu().tolist(),mean_heads))
+            del q,scores,allowed,logz,probs,entropy
+        self.last_entropy_probe = records
+        
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        use_cache: bool = False,
+        output_attentions: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor, torch.Tensor]]]:
+        batch_size, seq_len, _ = hidden_states.size()
+        
+        # Project Q, K, V
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+        
+        # Reshape for multi-head attention
+        query_states = query_states.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(batch_size, seq_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(batch_size, seq_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+        
+        # Apply Q/K normalisation when enabled.  In the no-QK-norm regime the
+        # projections go straight to RoPE/logits, with only the evaluation-time
+        # query-side temperature scale below.
+        if self.qk_norm == "on":
+            query_states = self.q_norm(query_states)
+            key_states = self.k_norm(key_states)
+        if self.qk_norm == "off" or self._use_query_temperature:
+            query_states = query_states * self._query_temperature_scale.to(dtype=query_states.dtype)
+        
+        # RoPE and ALiBi are mutually exclusive.  ALiBi is added by FlexAttention
+        # after QK scaling, so changing q_norm for temperature rescales content only.
+        if self.positional == "rope":
+            cos, sin = self.rotary_emb(value_states, position_ids)
+            query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+        
+        # Handle KV cache - store keys/values BEFORE repeating for GQA
+        if past_key_value is not None:
+            past_key, past_value = past_key_value
+            key_states = torch.cat([past_key, key_states], dim=2)
+            value_states = torch.cat([past_value, value_states], dim=2)
+        
+        if use_cache:
+            present_key_value = (key_states, value_states)
+        else:
+            present_key_value = None
+        
+        key_states = self._repeat_kv(key_states, self.num_key_value_groups)
+        value_states = self._repeat_kv(value_states, self.num_key_value_groups)
+
+        if self.entropy_probe_queries:
+            self._measure_attention_entropy(query_states, key_states)
+        
+        if self.positional == "alibi":
+            if flex_attention is None or create_block_mask is None:
+                raise ImportError("ALiBi requires torch FlexAttention support")
+            if attention_mask is not None or past_key_value is not None:
+                raise ValueError("ALiBi currently requires packed causal batches without a KV cache")
+            key = (seq_len, query_states.device)
+            block_mask = self._alibi_masks.get(key)
+            if block_mask is None:
+                block_mask = create_block_mask(_alibi_causal_mask, B=None, H=None, Q_LEN=seq_len, KV_LEN=seq_len, device=query_states.device)
+                self._alibi_masks[key] = block_mask
+            attn_output = flex_attention(query_states, key_states, value_states, score_mod=self._alibi_score_mod, block_mask=block_mask, enable_gqa=False, kernel_options={"BACKEND": "TRITON"})
+        else:
+            attn_output = F.scaled_dot_product_attention(
+                query_states, key_states, value_states, attn_mask=attention_mask,
+                dropout_p=self.attention_dropout if self.training else 0.0,
+                is_causal=attention_mask is None)
+        
+        attn_weights = None # SDPA doesn't materialize weights, saving massive memory
+        
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.reshape(batch_size, seq_len, self.attention_hidden_size)
+        
+        # Output projection
+        attn_output = self.o_proj(attn_output)
+        
+        if not output_attentions:
+            attn_weights = None
+        
+        return attn_output, attn_weights, present_key_value
+    
+    @staticmethod
+    def _repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+        """Repeat KV heads to match the number of query heads for GQA."""
+        if n_rep == 1:
+            return hidden_states
+        batch, num_kv_heads, slen, head_dim = hidden_states.shape
+        hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_kv_heads, n_rep, slen, head_dim)
+        return hidden_states.reshape(batch, num_kv_heads * n_rep, slen, head_dim)
+    

epfl_final/SparseAtt/sparseatt/config.py

@@ -0,0 +1,90 @@
+"""
+Qwen3 Model Configuration.
+"""
+
+from dataclasses import dataclass, field
+from typing import List, Optional
+
+
+@dataclass
+class Qwen3Config:
+    """Configuration class for Qwen3 model."""
+    vocab_size: int = 151936
+    hidden_size: int = 1024
+    intermediate_size: int = 3072
+    num_hidden_layers: int = 28
+    num_attention_heads: int = 16
+    num_key_value_heads: int = 8
+    head_dim: int = 128
+    hidden_act: str = "silu"
+    max_position_embeddings: int = 32768
+    initializer_range: float = 0.02
+    rms_norm_eps: float = 1e-6
+    use_cache: bool = True
+    tie_word_embeddings: bool = True
+    rope_theta: float = 1000000.0
+    rope_scaling: Optional[dict] = None
+    # ``rope`` (default) preserves legacy behavior; ``alibi`` disables RoPE
+    # and applies Press et al. head-specific causal distance biases.
+    positional: str = "rope"
+    # QK normalisation is part of the published Qwen3-style attention path.
+    # Keep it on by default so checkpoints/configs written before this flag are
+    # bit-identical; ``off`` removes both per-head RMSNorm modules.
+    qk_norm: str = "on"
+    # Keep Q/K RMSNorm gains and their AdamW state in FP32 during training.
+    qk_gain_fp32: bool = False
+    # Training-only RoPE position scales. When non-empty, every attention layer
+    # independently samples one scale per forward pass. Evaluation always uses
+    # the standard scale 1.0. This provides a dense positional-augmentation
+    # control without deleting attention links.
+    rope_train_position_scale_set: Optional[List[float]] = None
+    attention_bias: bool = False
+    attention_dropout: float = 0.0
+    bos_token_id: int = 151643
+    eos_token_id: int = 151643
+    checkpoint_activations: bool = False
+    attention_type: str = "dense"
+    sparse_block_size: int = 128
+    sparse_num_random_blocks: int = 1
+    sparse_mask_pool_size: int = 1000
+    sparse_max_seq_length: int = 1024
+    # When set, only the listed layer indices use sparse attention; all others use dense.
+    # When None (default), attention_type applies to all layers uniformly.
+    sparse_layer_indices: Optional[List[int]] = None
+    # Directory to cache pre-built block masks (one file per layer). If None, masks are rebuilt every run.
+    sparse_mask_cache_dir: Optional[str] = None
+    # Seed offset for the frozen (pool_size=1) fixed mask. 0 = legacy (layer-index only).
+    # Distinct values yield distinct fixed masks (fixed-mask seeds).
+    sparse_mask_seed: int = 0
+    # Mixed-sparsity training: if set to a non-empty list, each mask-pool entry draws its
+    # OWN num_random_blocks uniformly from this set, so the per-step mask (picked at random
+    # from the pool) also resamples the sparsity LEVEL. None/empty = legacy fixed-nrb pool.
+    # This is the "Mixed" training arm of the paper; without it the arm silently degenerates
+    # to a fixed-nrb model.
+    sparse_mixed_nrb_set: Optional[List[int]] = None
+    # --- deployment-time baselines (learned MoBA + ASEntmax) ---
+    # moba_top_k is the per-query block budget used by the top-k deployment ladder.
+    moba_block_size: int = 128
+    moba_top_k: int = 32
+    asentmax_alpha: float = 1.5
+    asentmax_learn_temp: bool = True
+    asentmax_init_temp: float = 1.0
+
+    def is_layer_sparse(self, layer_idx: int) -> bool:
+        """Return True if this layer should use sparse stochastic attention."""
+        if self.attention_type != "sparse_stochastic":
+            return False
+        if self.sparse_layer_indices is None:
+            return True  # all layers sparse (existing behaviour)
+        return layer_idx in self.sparse_layer_indices
+
+    @classmethod
+    def from_dict(cls, config_dict: dict) -> "Qwen3Config":
+        return cls(**{k: v for k, v in config_dict.items() if k in cls.__dataclass_fields__})
+    
+    @classmethod
+    def from_json(cls, path: str) -> "Qwen3Config":
+        import json
+        with open(path, "r") as f:
+            config_dict = json.load(f)
+        return cls.from_dict(config_dict)