diff --git a/epfl_final/SparseAtt/$script b/epfl_final/SparseAtt/$script new file mode 100644 index 0000000000000000000000000000000000000000..98711919fc176414a2c97f32bd92d11bc2e27388 --- /dev/null +++ b/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 & diff --git a/epfl_final/SparseAtt/closure24_tau1_eval.py b/epfl_final/SparseAtt/closure24_tau1_eval.py new file mode 100644 index 0000000000000000000000000000000000000000..7374442e86a24e3596ee15414931c0c4bd9fc88b --- /dev/null +++ b/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() diff --git a/epfl_final/SparseAtt/concat_tau1_fix.py b/epfl_final/SparseAtt/concat_tau1_fix.py new file mode 100644 index 0000000000000000000000000000000000000000..15b1445bee11ad7a5506e69eca13d6cf43129de8 --- /dev/null +++ b/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() diff --git a/epfl_final/SparseAtt/configs/FINAL_20260924/365m_dense_qkgainfp32_s42.json b/epfl_final/SparseAtt/configs/FINAL_20260924/365m_dense_qkgainfp32_s42.json new file mode 100644 index 0000000000000000000000000000000000000000..8f73d815fb5479ae8bad721b27988431f84d3014 --- /dev/null +++ b/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} +} diff --git a/epfl_final/SparseAtt/configs/FINAL_20260924/365m_dense_qkgainfp32_s42_2gpu_latest_resume.json b/epfl_final/SparseAtt/configs/FINAL_20260924/365m_dense_qkgainfp32_s42_2gpu_latest_resume.json new file mode 100644 index 0000000000000000000000000000000000000000..2558f8993a9fbd7b77f67d646a0e1b6ddf0da1c5 --- /dev/null +++ b/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} +} diff --git a/epfl_final/SparseAtt/configs/FINAL_20260924/365m_dense_qkgainfp32_s42_2gpu_restart_resume.json b/epfl_final/SparseAtt/configs/FINAL_20260924/365m_dense_qkgainfp32_s42_2gpu_restart_resume.json new file mode 100644 index 0000000000000000000000000000000000000000..36c3091f8279e71c0676f16f9734e5e76f1ee4d3 --- /dev/null +++ b/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} +} diff --git a/epfl_final/SparseAtt/configs/FINAL_20260924/365m_dense_qkgainfp32_s42_2gpu_resume.json b/epfl_final/SparseAtt/configs/FINAL_20260924/365m_dense_qkgainfp32_s42_2gpu_resume.json new file mode 100644 index 0000000000000000000000000000000000000000..36c3091f8279e71c0676f16f9734e5e76f1ee4d3 --- /dev/null +++ b/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} +} diff --git a/epfl_final/SparseAtt/configs/FINAL_20260924/365m_dense_qkgainfp32_s42_5gpu_a100_resume.json b/epfl_final/SparseAtt/configs/FINAL_20260924/365m_dense_qkgainfp32_s42_5gpu_a100_resume.json new file mode 100644 index 0000000000000000000000000000000000000000..981b49429c922b81b4ba66dcb906e477d4759835 --- /dev/null +++ b/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} +} diff --git a/epfl_final/SparseAtt/configs/FINAL_20260924/365m_dense_qkgainfp32_s43_2gpu.json b/epfl_final/SparseAtt/configs/FINAL_20260924/365m_dense_qkgainfp32_s43_2gpu.json new file mode 100644 index 0000000000000000000000000000000000000000..cc39f0fd6bacbb82abc77dc0eed6b1e71c21d4ee --- /dev/null +++ b/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} +} diff --git a/epfl_final/SparseAtt/configs/FINAL_20260924/365m_dense_qkgainfp32_s43_2gpu_resume.json b/epfl_final/SparseAtt/configs/FINAL_20260924/365m_dense_qkgainfp32_s43_2gpu_resume.json new file mode 100644 index 0000000000000000000000000000000000000000..63d9b8f1542eb1b96802bec0b4ad8218596585c6 --- /dev/null +++ b/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} +} diff --git a/epfl_final/SparseAtt/configs/FINAL_20260924/365m_dense_qkgainfp32_s44_5gpu.json b/epfl_final/SparseAtt/configs/FINAL_20260924/365m_dense_qkgainfp32_s44_5gpu.json new file mode 100644 index 0000000000000000000000000000000000000000..5a627b6585f776eef2cd1731b142363eeaa74aeb --- /dev/null +++ b/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} +} diff --git a/epfl_final/SparseAtt/configs/FINAL_20260924/365m_mixed_qkgainfp32_s42.json b/epfl_final/SparseAtt/configs/FINAL_20260924/365m_mixed_qkgainfp32_s42.json new file mode 100644 index 0000000000000000000000000000000000000000..e9bfebb79e1c597a9924c93c834920dd5dd1a49d --- /dev/null +++ b/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"} +} diff --git a/epfl_final/SparseAtt/configs/FINAL_20260924/365m_mixed_qkgainfp32_s42_2gpu_h100_restart_resume.json b/epfl_final/SparseAtt/configs/FINAL_20260924/365m_mixed_qkgainfp32_s42_2gpu_h100_restart_resume.json new file mode 100644 index 0000000000000000000000000000000000000000..1e780e1e3654c6f7947bfb5974d4d857a3d61ef4 --- /dev/null +++ b/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"} +} diff --git a/epfl_final/SparseAtt/configs/FINAL_20260924/365m_mixed_qkgainfp32_s42_2gpu_h100_resume.json b/epfl_final/SparseAtt/configs/FINAL_20260924/365m_mixed_qkgainfp32_s42_2gpu_h100_resume.json new file mode 100644 index 0000000000000000000000000000000000000000..aaa4615331ea2de79f426899375a350161b93fa8 --- /dev/null +++ b/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"} +} diff --git a/epfl_final/SparseAtt/configs/FINAL_20260924/365m_mixed_qkgainfp32_s42_5gpu_latest_resume.json b/epfl_final/SparseAtt/configs/FINAL_20260924/365m_mixed_qkgainfp32_s42_5gpu_latest_resume.json new file mode 100644 index 0000000000000000000000000000000000000000..621455ebcbfc2487a965ea184ea9006cb5bc1a09 --- /dev/null +++ b/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"} +} diff --git a/epfl_final/SparseAtt/configs/FINAL_20260924/365m_mixed_qkgainfp32_s42_5gpu_restart_resume.json b/epfl_final/SparseAtt/configs/FINAL_20260924/365m_mixed_qkgainfp32_s42_5gpu_restart_resume.json new file mode 100644 index 0000000000000000000000000000000000000000..bd5524e01b65ccc2d4aaa7869070945e69123505 --- /dev/null +++ b/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"} +} diff --git a/epfl_final/SparseAtt/configs/FINAL_20260924/365m_mixed_qkgainfp32_s43.json b/epfl_final/SparseAtt/configs/FINAL_20260924/365m_mixed_qkgainfp32_s43.json new file mode 100644 index 0000000000000000000000000000000000000000..4c2d0ba4572460929ff640ef7049484df0c004ae --- /dev/null +++ b/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"} +} diff --git a/epfl_final/SparseAtt/configs/FINAL_20260924/365m_mixed_qkgainfp32_s43_2gpu_resume.json b/epfl_final/SparseAtt/configs/FINAL_20260924/365m_mixed_qkgainfp32_s43_2gpu_resume.json new file mode 100644 index 0000000000000000000000000000000000000000..c15ca1ccc9d9621a1f075052ffb4bd0c4a765782 --- /dev/null +++ b/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"} +} diff --git a/epfl_final/SparseAtt/configs/FINAL_20260924/365m_mixed_qkgainfp32_s43_5gpu_resume.json b/epfl_final/SparseAtt/configs/FINAL_20260924/365m_mixed_qkgainfp32_s43_5gpu_resume.json new file mode 100644 index 0000000000000000000000000000000000000000..d7f15e0fafdb92dade2687d993ec2758b2848d55 --- /dev/null +++ b/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"} +} diff --git a/epfl_final/SparseAtt/configs/FINAL_20260924/365m_mixed_qkgainfp32_s44_2gpu_h100_latest_resume.json b/epfl_final/SparseAtt/configs/FINAL_20260924/365m_mixed_qkgainfp32_s44_2gpu_h100_latest_resume.json new file mode 100644 index 0000000000000000000000000000000000000000..3c4bb30728c9b9cbd8e29d3aac08010c875a4b56 --- /dev/null +++ b/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"} +} diff --git a/epfl_final/SparseAtt/configs/FINAL_20260924/365m_mixed_qkgainfp32_s44_2gpu_h100_resume.json b/epfl_final/SparseAtt/configs/FINAL_20260924/365m_mixed_qkgainfp32_s44_2gpu_h100_resume.json new file mode 100644 index 0000000000000000000000000000000000000000..81311d795f11d5aef3f34b6446f079dea4ac1e81 --- /dev/null +++ b/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"} +} diff --git a/epfl_final/SparseAtt/configs/FINAL_20260924/365m_mixed_qkgainfp32_s44_5gpu.json b/epfl_final/SparseAtt/configs/FINAL_20260924/365m_mixed_qkgainfp32_s44_5gpu.json new file mode 100644 index 0000000000000000000000000000000000000000..2afdcfb4a9fb82a1253c43eb114d92fd70365df3 --- /dev/null +++ b/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"} +} diff --git a/epfl_final/SparseAtt/configs/REGIME_20260922/365m_dense_noqknorm_s42.json b/epfl_final/SparseAtt/configs/REGIME_20260922/365m_dense_noqknorm_s42.json new file mode 100644 index 0000000000000000000000000000000000000000..228f6ad0b42a4e65a42894bd0d1a78a19ae594db --- /dev/null +++ b/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} +} diff --git a/epfl_final/SparseAtt/configs/REGIME_20260922/365m_dense_rope1e4_s42.json b/epfl_final/SparseAtt/configs/REGIME_20260922/365m_dense_rope1e4_s42.json new file mode 100644 index 0000000000000000000000000000000000000000..e39e7943e579048ef7bde930c7caa8d970e12300 --- /dev/null +++ b/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} +} diff --git a/epfl_final/SparseAtt/configs/REGIME_20260922/365m_dense_rope1e4_s42_1gpu.json b/epfl_final/SparseAtt/configs/REGIME_20260922/365m_dense_rope1e4_s42_1gpu.json new file mode 100644 index 0000000000000000000000000000000000000000..751972e5608f36ed7c41aca8b2b42b7aa7360ec3 --- /dev/null +++ b/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} +} diff --git a/epfl_final/SparseAtt/configs/REGIME_20260922/365m_dense_rope1e4_s42_resume.json b/epfl_final/SparseAtt/configs/REGIME_20260922/365m_dense_rope1e4_s42_resume.json new file mode 100644 index 0000000000000000000000000000000000000000..1fcbf86e2b5c91eae2dd951ae2566a6949de5c50 --- /dev/null +++ b/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 + } +} \ No newline at end of file diff --git a/epfl_final/SparseAtt/configs/REGIME_20260922/365m_mixed_noqknorm_s42.json b/epfl_final/SparseAtt/configs/REGIME_20260922/365m_mixed_noqknorm_s42.json new file mode 100644 index 0000000000000000000000000000000000000000..34830dd3e3c00ebfd8d9ff0880c7f90d1ee68836 --- /dev/null +++ b/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" + } +} diff --git a/epfl_final/SparseAtt/configs/REGIME_20260922/365m_mixed_noqknorm_s42_resume.json b/epfl_final/SparseAtt/configs/REGIME_20260922/365m_mixed_noqknorm_s42_resume.json new file mode 100644 index 0000000000000000000000000000000000000000..cc6d33233633770cf4e7be37bda116c474e7e85d --- /dev/null +++ b/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" + } +} diff --git a/epfl_final/SparseAtt/configs/REGIME_20260922/365m_mixed_rope1e4_s42.json b/epfl_final/SparseAtt/configs/REGIME_20260922/365m_mixed_rope1e4_s42.json new file mode 100644 index 0000000000000000000000000000000000000000..23261a43f2cf938506bead67b119b1ea768b13b5 --- /dev/null +++ b/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" + } +} diff --git a/epfl_final/SparseAtt/configs/REGIME_20260922/365m_mixed_rope1e4_s42_2gpu.json b/epfl_final/SparseAtt/configs/REGIME_20260922/365m_mixed_rope1e4_s42_2gpu.json new file mode 100644 index 0000000000000000000000000000000000000000..c0af2b4116036d86408798f72485a7654f58bb86 --- /dev/null +++ b/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" + } +} diff --git a/epfl_final/SparseAtt/configs/alibi_20260921/1p2b_alibi_s42.json b/epfl_final/SparseAtt/configs/alibi_20260921/1p2b_alibi_s42.json new file mode 100644 index 0000000000000000000000000000000000000000..bf4df5a630c10581186c37fac620f900ad24d0f2 --- /dev/null +++ b/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 + } +} diff --git a/epfl_final/SparseAtt/configs/alibi_20260921/1p2b_alibi_s42_bs12_resume17000.json b/epfl_final/SparseAtt/configs/alibi_20260921/1p2b_alibi_s42_bs12_resume17000.json new file mode 100644 index 0000000000000000000000000000000000000000..23cd3f18826ccf4e497b0d3cb4ca02663c2b0edf --- /dev/null +++ b/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 + } +} diff --git a/epfl_final/SparseAtt/configs/alibi_20260921/1p2b_alibi_s42_bs15.json b/epfl_final/SparseAtt/configs/alibi_20260921/1p2b_alibi_s42_bs15.json new file mode 100644 index 0000000000000000000000000000000000000000..23a313bd6f7450c093e25d0342a4c31e4634fbbd --- /dev/null +++ b/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 + } +} diff --git a/epfl_final/SparseAtt/configs/alibi_20260921/1p2b_alibi_s42_bs20.json b/epfl_final/SparseAtt/configs/alibi_20260921/1p2b_alibi_s42_bs20.json new file mode 100644 index 0000000000000000000000000000000000000000..12267b9b6991da31a4c9720ef0a9b40d091223a1 --- /dev/null +++ b/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 + } +} diff --git a/epfl_final/SparseAtt/configs/alibi_20260921/365m_alibi_s42.json b/epfl_final/SparseAtt/configs/alibi_20260921/365m_alibi_s42.json new file mode 100644 index 0000000000000000000000000000000000000000..94ce0edb3a36fb6492f34e9648a21f5d89211eea --- /dev/null +++ b/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 + } +} diff --git a/epfl_final/SparseAtt/configs/alibi_20260921/365m_alibi_s43.json b/epfl_final/SparseAtt/configs/alibi_20260921/365m_alibi_s43.json new file mode 100644 index 0000000000000000000000000000000000000000..037511de6f07093a098337fa380a63eda11ab2b5 --- /dev/null +++ b/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 + } +} diff --git a/epfl_final/SparseAtt/configs/alibi_20260921/365m_alibi_s44.json b/epfl_final/SparseAtt/configs/alibi_20260921/365m_alibi_s44.json new file mode 100644 index 0000000000000000000000000000000000000000..d3c003eb7ed386c4e973d3a17c180e70091696fd --- /dev/null +++ b/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 + } +} diff --git a/epfl_final/SparseAtt/configs/alibi_20260921/365m_alibi_s45.json b/epfl_final/SparseAtt/configs/alibi_20260921/365m_alibi_s45.json new file mode 100644 index 0000000000000000000000000000000000000000..ad5afcebb6c68c7c8c0c350ef798225ae24e58d7 --- /dev/null +++ b/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 + } +} diff --git a/epfl_final/SparseAtt/configs/alibi_20260921/365m_alibi_s46.json b/epfl_final/SparseAtt/configs/alibi_20260921/365m_alibi_s46.json new file mode 100644 index 0000000000000000000000000000000000000000..35e3b37d6e320ff5f950bf3d8873e99646ae1fba --- /dev/null +++ b/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 + } +} diff --git a/epfl_final/SparseAtt/configs/alibi_20260921/365m_alibi_s47.json b/epfl_final/SparseAtt/configs/alibi_20260921/365m_alibi_s47.json new file mode 100644 index 0000000000000000000000000000000000000000..3575f5775c8f5217ff9e9003a3f32ebe4648fda7 --- /dev/null +++ b/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 + } +} diff --git a/epfl_final/SparseAtt/configs/alibi_20260921/365m_alibi_s47_restart_20260922.json b/epfl_final/SparseAtt/configs/alibi_20260921/365m_alibi_s47_restart_20260922.json new file mode 100644 index 0000000000000000000000000000000000000000..6918d60122b3c4ee8d978722483a82a6889d17e6 --- /dev/null +++ b/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 + } +} diff --git a/epfl_final/SparseAtt/configs/alibi_20260921/365m_alibi_s47_resume_2xh100.json b/epfl_final/SparseAtt/configs/alibi_20260921/365m_alibi_s47_resume_2xh100.json new file mode 100644 index 0000000000000000000000000000000000000000..e35cbfbdbde1b86e889356938c50c79c88707b18 --- /dev/null +++ b/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 + } +} diff --git a/epfl_final/SparseAtt/finalize_tau1_fix.py b/epfl_final/SparseAtt/finalize_tau1_fix.py new file mode 100644 index 0000000000000000000000000000000000000000..f0a5ad9e6ce06740a54757f5a77fcb636cf1910b --- /dev/null +++ b/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() diff --git a/epfl_final/SparseAtt/launch_1p2b_alibi_after_eval.sh b/epfl_final/SparseAtt/launch_1p2b_alibi_after_eval.sh new file mode 100644 index 0000000000000000000000000000000000000000..aba3870f63b5de369ee0b00ebeccb93f5a423ce2 --- /dev/null +++ b/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 diff --git a/epfl_final/SparseAtt/launch_alibi_20260921.sh b/epfl_final/SparseAtt/launch_alibi_20260921.sh new file mode 100644 index 0000000000000000000000000000000000000000..b6155b81739f41ae90b4f0e4306ffdcfc861cb5a --- /dev/null +++ b/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 diff --git a/epfl_final/SparseAtt/merge_closure24.py b/epfl_final/SparseAtt/merge_closure24.py new file mode 100644 index 0000000000000000000000000000000000000000..b9cf616f81f7e8324cc2730283f2211dbc937bdf --- /dev/null +++ b/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() diff --git a/epfl_final/SparseAtt/rebuild_closure24_from_log.py b/epfl_final/SparseAtt/rebuild_closure24_from_log.py new file mode 100644 index 0000000000000000000000000000000000000000..08c0c42aade6d72038d7d937db340c8eef497d3a --- /dev/null +++ b/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() diff --git a/epfl_final/SparseAtt/sparseatt/asentmax_attention.py b/epfl_final/SparseAtt/sparseatt/asentmax_attention.py new file mode 100644 index 0000000000000000000000000000000000000000..2fb6cce07c7b687463d2cf6d6252aa17bb32f246 --- /dev/null +++ b/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 diff --git a/epfl_final/SparseAtt/sparseatt/attention.py b/epfl_final/SparseAtt/sparseatt/attention.py new file mode 100644 index 0000000000000000000000000000000000000000..25a8cba69dc768060d2824a3ce8cdf4525cb1b45 --- /dev/null +++ b/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) + diff --git a/epfl_final/SparseAtt/sparseatt/config.py b/epfl_final/SparseAtt/sparseatt/config.py new file mode 100644 index 0000000000000000000000000000000000000000..2ecc4532d354ce286457080fbab05d4e8dda0018 --- /dev/null +++ b/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) diff --git a/epfl_final/SparseAtt/sparseatt/moba_attention.py b/epfl_final/SparseAtt/sparseatt/moba_attention.py new file mode 100644 index 0000000000000000000000000000000000000000..321d0cc7255eddf70b15f51025667ccdb9d22885 --- /dev/null +++ b/epfl_final/SparseAtt/sparseatt/moba_attention.py @@ -0,0 +1,160 @@ +""" +MoBA (Mixture of Block Attention) baseline — arm #2, FlexAttention version. + +Content-AWARE block routing: each query-block attends to the top-k key-blocks by +pooled-Q x pooled-K affinity (+ own block, + block-causal). Identical selection to +the dense version, but realised through a FlexAttention block_mask (flash-style +block-sparse kernel) instead of a materialised N^2 matrix -> no O(T^2) memory, no +activation checkpointing, ~same cost as the resampled arm's kernel. + +The BlockMask is built per-forward directly from the content-dependent top-k +block indices. Past selected blocks are marked full, while diagonal blocks use a +plain causal mask. Direct metadata construction avoids expanding the routing +lookup over every token pair. + +Ref: MoBA — Mixture of Block Attention for Long-Context LLMs (arXiv:2502.13189). +""" +from typing import Optional, Tuple +import math + +import torch +import torch.nn.functional as F + +try: + from torch.nn.attention.flex_attention import BlockMask, flex_attention +except ImportError: # pragma: no cover + BlockMask = None + flex_attention = None + +from .attention import Qwen3Attention, apply_rotary_pos_emb + + +def _causal_mask_mod(b, h, q_idx, kv_idx): + """Plain causal for runtime partial (diagonal) blocks — the block-level pattern + is already encoded in the BlockMask structure. Avoids an index gather Inductor + can't lower in a pointwise subgraph (same reason as the sparse-stochastic arm).""" + return q_idx >= kv_idx + + +class Qwen3MoBAAttention(Qwen3Attention): + """Block-sparse attention with content-aware top-k block routing via FlexAttention.""" + + def __init__(self, config, layer_idx: int): + super().__init__(config, layer_idx) + if flex_attention is None or BlockMask is None: + raise ImportError("FlexAttention unavailable (needs torch>=2.5)") + self.block_size = int(getattr(config, "moba_block_size", + getattr(config, "sparse_block_size", 128))) + self.top_k = int(getattr(config, "moba_top_k", 32)) + # ``block_pooled`` is the historical evaluator. ``causal_prev`` is an + # opt-in diagnostic: block b routes from the pooled query of block b-1, + # so routing never reads tokens later than a query in b. + self.routing = str(getattr(config, "moba_routing", "block_pooled")) + if self.routing not in {"block_pooled", "causal_prev"}: + raise ValueError(f"unknown MoBA routing mode: {self.routing}") + self.capture_routing = False + self.last_routing_indices = None + self.dense_mode = False + self._flex_attention = flex_attention + + 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) + + # At a full block budget, dispatching through SDPA is both exactly the + # same operator and avoids harmless-but-measurable FlexAttention + # accumulation-order differences. This is used by the causal-routing + # diagnostic's full-budget equivalence check. + full_budget = self.top_k >= (T // self.block_size) + if self.dense_mode or (self.routing == "causal_prev" and full_budget): + out = F.scaled_dot_product_attention(q, k, v, is_causal=True) + else: + block_mask = self._build_block_mask(q, k) + out = self._flex_attention(q, k, v, block_mask=block_mask, enable_gqa=False) + + out = out.transpose(1, 2).contiguous().reshape(B, T, self.attention_hidden_size) + return self.o_proj(out), None, None + + def _build_block_mask(self, q: torch.Tensor, k: torch.Tensor): + B, H, T, d = q.shape + bs = self.block_size + if T % bs != 0: + raise ValueError(f"MoBA requires seq_len ({T}) divisible by block_size ({bs})") + nb = T // bs + scale = 1.0 / math.sqrt(d) + + # --- content-aware top-k block selection (per batch & head) --- + qb = q.reshape(B, H, nb, bs, d).mean(dim=3) # [B,H,nb,d] + if self.routing == "causal_prev": + # Block 0 has no predecessor and retains only its diagonal block. + # For b >= 1 this avoids the current block's later query tokens. + previous_qb = torch.zeros_like(qb) + previous_qb[:, :, 1:] = qb[:, :, :-1] + qb = previous_qb + kb = k.reshape(B, H, nb, bs, d).mean(dim=3) # [B,H,nb,d] + gate = torch.matmul(qb, kb.transpose(-2, -1)) * scale # [B,H,nb,nb] + idx = torch.arange(nb, device=q.device) + gate = gate.masked_fill(idx[None, :] > idx[:, None], float("-inf")) # block-causal + kk = min(self.top_k, nb) + topv, topi = gate.topk(kk, dim=-1) # [B,H,nb,kk] + if self.capture_routing: + # Diagnostic-only; disabled by default so compiled evaluator paths + # retain their existing behavior. + self.last_routing_indices = topi.detach().cpu() + + # Build FlexAttention's block metadata directly. Calling create_block_mask + # with a content-dependent tensor lookup vmaps that lookup over every token + # pair, creating an O(T^2) temporary (84 GiB at T=30720). Every selected + # past block is fully causal; only the diagonal block needs token-level + # causal masking. + q_blocks = idx.view(1, 1, nb, 1) + full_valid = torch.isfinite(topv) & (topi < q_blocks) + full_kv_num_blocks = full_valid.sum(dim=-1, dtype=torch.int32) + packed_order = torch.argsort( + full_valid.to(torch.int8), dim=-1, descending=True, stable=True + ) + packed_full_indices = topi.gather(-1, packed_order).to(torch.int32) + full_kv_indices = torch.zeros( + B, H, nb, nb, dtype=torch.int32, device=q.device + ) + full_kv_indices[..., :kk] = packed_full_indices + + kv_num_blocks = torch.ones( + B, H, nb, dtype=torch.int32, device=q.device + ) + kv_indices = torch.zeros( + B, H, nb, nb, dtype=torch.int32, device=q.device + ) + kv_indices[..., 0] = idx.view(1, 1, nb).to(torch.int32) + + return BlockMask.from_kv_blocks( + kv_num_blocks=kv_num_blocks, + kv_indices=kv_indices, + full_kv_num_blocks=full_kv_num_blocks, + full_kv_indices=full_kv_indices, + BLOCK_SIZE=bs, + mask_mod=_causal_mask_mod, + seq_lengths=(T, T), + ) diff --git a/epfl_final/SparseAtt/sparseatt/sparse_attention.py b/epfl_final/SparseAtt/sparseatt/sparse_attention.py new file mode 100644 index 0000000000000000000000000000000000000000..c22b7d11e97da208e03a185f049b78671dde5252 --- /dev/null +++ b/epfl_final/SparseAtt/sparseatt/sparse_attention.py @@ -0,0 +1,496 @@ +"""Stochastic sparse attention module separated from dense attention for prototyping.""" + +import os +from typing import Optional, Tuple + +import torch +import torch.nn.functional as F + +try: + from torch.nn.attention.flex_attention import create_block_mask, flex_attention +except ImportError: + create_block_mask = None + flex_attention = None + +from .config import Qwen3Config +from .attention import Qwen3Attention, apply_rotary_pos_emb + + +def sample_rt_row( + num_blocks: int, + num_random_blocks: int, + qb: int, + device: torch.device, +) -> torch.Tensor: + r"""Sample ``num_random_blocks`` KV-block indices uniformly from + ``[0, num_blocks) \ {qb}``. Causality is enforced separately by the + mask_mod; both past and future blocks are candidates here. + """ + perm = torch.randperm(num_blocks - 1, device=device)[:num_random_blocks].clone() + perm[perm >= qb] += 1 + return perm + + +def _causal_mask_mod(b, h, q_idx, kv_idx): + """Simple causal mask stored in BlockMask.mask_mod for within-block token masking. + + The block-level sparsity pattern is already encoded in the BlockMask structure + (built via create_block_mask + _MaskMod). At flex_attention runtime, mask_mod is + only called for *partial* blocks — exclusively the local causal diagonal blocks + where q_block == kv_block. In those blocks is_local is always True so the random- + block lookup in _MaskMod is dead code. Using a plain causal function avoids the + aten.index.Tensor gather that Inductor cannot lower inside a pointwise subgraph. + """ + return q_idx >= kv_idx + + +class _MaskMod: + """Picklable mask callable for create_block_mask and BlockMask serialisation. + + Defined at module level so it survives pickle round-trips (e.g. multiprocessing + queues, torch.save / torch.load across different entry-points). + + ``rt`` holds indices sampled uniformly from all non-diagonal blocks (past *and* + future). Any entry with ``kv_block > qb`` is zeroed by the ``q_idx >= kv_idx`` + clause in ``__call__`` and never materialised as an attended block. + """ + + def __init__(self, rt: torch.Tensor, block_size: int, num_random_blocks: int): + # rt: (num_blocks, num_random_blocks) int16, on any device + self.rt = rt + self.block_size = block_size + self.num_random_blocks = num_random_blocks + + def __call__(self, b, h, q_idx, kv_idx): + q_block = q_idx // self.block_size + kv_block = kv_idx // self.block_size + is_local = q_block == kv_block + if self.num_random_blocks == 0: + return is_local & (q_idx >= kv_idx) + is_global = kv_block == self.rt[q_block, 0] + for j in range(1, self.num_random_blocks): + is_global = is_global | (kv_block == self.rt[q_block, j]) + return (is_local | is_global) & (q_idx >= kv_idx) + + +class Qwen3SparseStochasticAttention(Qwen3Attention): + """Block-sparse stochastic attention with a prebuilt mask pool.""" + + def __init__(self, config: Qwen3Config, layer_idx: int, mask_cache_dir: Optional[str] = None): + if create_block_mask is None or flex_attention is None: + raise ImportError( + "Sparse attention requires torch.nn.attention.flex_attention support." + ) + + self._create_block_mask = create_block_mask + self._flex_attention = flex_attention + + super().__init__(config, layer_idx) + + # Stash the layer index for per-layer mask seeding (see _build_mask_pool). + self._mask_layer_idx = int(layer_idx) + self._mask_seed = int(getattr(config, "sparse_mask_seed", 0)) + + self.seq_len = int(config.sparse_max_seq_length) + self.block_size = int(config.sparse_block_size) + self.num_blocks = self.seq_len // self.block_size + self.mask_pool_size = int(config.sparse_mask_pool_size) + + if self.block_size <= 0: + raise ValueError("sparse_block_size must be > 0") + if self.block_size % 64 != 0: + raise ValueError( + "sparse_block_size must be divisible by 64 for stable torch.compile flex-attention kernels" + ) + if self.seq_len <= 0: + raise ValueError("sparse_max_seq_length must be > 0") + if self.seq_len % self.block_size != 0: + raise ValueError( + f"sparse_max_seq_length ({self.seq_len}) must be divisible by sparse_block_size ({self.block_size})" + ) + if self.num_blocks < 2: + raise ValueError("Sparse attention requires at least 2 blocks") + if self.mask_pool_size <= 0: + raise ValueError("sparse_mask_pool_size must be > 0") + + self.num_random_blocks = max( + 0, + min(int(config.sparse_num_random_blocks), self.num_blocks - 1), + ) + + # Mixed-sparsity: each pool entry samples its OWN nrb from this set, so masks picked + # at random per step also resample the sparsity LEVEL. The _rt_pool buffer is sized to + # the largest nrb; entries with a smaller nrb use only the leading columns. + raw_mixed = getattr(config, "sparse_mixed_nrb_set", None) + if raw_mixed: + self.mixed_nrb_set = sorted({ + max(0, min(int(x), self.num_blocks - 1)) for x in raw_mixed + }) + self.num_random_blocks = max(self.mixed_nrb_set) + else: + self.mixed_nrb_set = None + + # When True, forward() uses standard causal SDPA instead of flex_attention. + # Flip this flag after loading a sparse checkpoint to continue training in dense mode. + self.dense_mode: bool = False + + buffer_device = "cuda" if torch.cuda.is_available() else "cpu" + self.register_buffer( + "_rt_pool", + torch.zeros( + (self.mask_pool_size, self.num_blocks, self.num_random_blocks), + dtype=torch.int16, + device=buffer_device, + ), + persistent=False, # not a learned weight — rebuilt each run, never saved to checkpoint + ) + + cache_path = ( + os.path.join(mask_cache_dir, f"masks_layer_{layer_idx}.pt") + if mask_cache_dir is not None + else None + ) + # Masks are always kept on CPU to avoid occupying GPU VRAM for the + # entire pool (~700 MB/layer × 28 layers = ~20 GB). Only the one mask + # selected per forward pass is moved to the compute device on-the-fly. + mask_cpu = torch.device("cpu") + if cache_path is not None and os.path.isfile(cache_path): + mask_pool, rt = self._load_masks(cache_path, mask_cpu) + if len(mask_pool) < self.mask_pool_size: + # Cached pool is smaller than requested — rebuild. + self._mask_pool = self._build_mask_pool(mask_cpu) + os.makedirs(mask_cache_dir, exist_ok=True) + self.save_masks(cache_path) + else: + self._mask_pool = mask_pool[: self.mask_pool_size] + self._rt_pool.copy_(rt[: self.mask_pool_size].to(self._rt_pool.device, dtype=self._rt_pool.dtype)) + else: + # Eval / no-cache path: build on the GPU (a BlockMask is small block + # metadata, not a dense seq^2 tensor). CPU build is ~2.8s/mask, which + # at pool_size x 28 layers cost minutes; GPU build is ~0.3s/mask. + build_device = torch.device(buffer_device) + self._mask_pool = self._build_mask_pool(build_device) + if cache_path is not None: + os.makedirs(mask_cache_dir, exist_ok=True) + self.save_masks(cache_path) + + def _build_mask_pool(self, device: torch.device): + # Keep mask creation RNG use isolated from model/init RNG streams so that + # building masks never perturbs weight-init reproducibility. + cpu_rng_state = torch.random.get_rng_state() + cuda_rng_state = torch.cuda.get_rng_state() if torch.cuda.is_available() else None + # Seed PER-LAYER (deterministically) so every layer builds its OWN distinct + # mask pool. Previously this restored the RNG to the *same* state for every + # layer, which made all layers' masks byte-identical (a latent bug on the + # in-model / no-cache path: a `pool_size=1` model ended up with one shared + # mask across all layers instead of one distinct mask per layer). + torch.manual_seed(1234567 + 1000003 * int(getattr(self, "_mask_layer_idx", 0)) + 2000003 * int(getattr(self, "_mask_seed", 0))) + pool = self._build_mask_pool_inner(device) + torch.random.set_rng_state(cpu_rng_state) + if cuda_rng_state is not None: + torch.cuda.set_rng_state(cuda_rng_state) + return pool + + def _build_mask_pool_inner(self, device: torch.device): + pool = [] + block_size = self.block_size + num_blocks = self.num_blocks + mixed_nrb_set = self.mixed_nrb_set + + for i in range(self.mask_pool_size): + if mixed_nrb_set: + # Per-entry sparsity level: one nrb drawn from the set for this whole mask. + nrb_i = int(mixed_nrb_set[torch.randint(len(mixed_nrb_set), (1,)).item()]) + else: + nrb_i = self.num_random_blocks + + if nrb_i > 0: + self._rt_pool[i].zero_() + for qb in range(num_blocks): + perm = sample_rt_row(num_blocks, nrb_i, qb, self._rt_pool.device) + self._rt_pool[i, qb, :nrb_i] = perm.to(dtype=self._rt_pool.dtype) + + bm = self._create_block_mask( + _MaskMod(self._rt_pool[i].to(device), block_size, nrb_i), + B=None, + H=None, + Q_LEN=self.seq_len, + KV_LEN=self.seq_len, + device=device, + BLOCK_SIZE=block_size, + ) + # Replace the stored mask_mod with a simple causal function. + # _MaskMod is only needed above to build the correct block structure; + # flex_attention's within-block masking only ever needs q_idx >= kv_idx. + bm.mask_mod = _causal_mask_mod + pool.append(bm) + + return pool + + def save_masks(self, path: str) -> None: + """Persist the mask pool and random-block table to disk. + + Stores only raw index tensors from each BlockMask (no callables) so + the file is portable across entry-points and Python processes. + """ + mask_tensors = [ + { + "seq_lengths": bm.seq_lengths, + "kv_num_blocks": bm.kv_num_blocks.cpu(), + "kv_indices": bm.kv_indices.cpu(), + "full_kv_num_blocks": bm.full_kv_num_blocks.cpu() if bm.full_kv_num_blocks is not None else None, + "full_kv_indices": bm.full_kv_indices.cpu() if bm.full_kv_indices is not None else None, + "q_num_blocks": bm.q_num_blocks.cpu() if bm.q_num_blocks is not None else None, + "q_indices": bm.q_indices.cpu() if bm.q_indices is not None else None, + "full_q_num_blocks": bm.full_q_num_blocks.cpu() if bm.full_q_num_blocks is not None else None, + "full_q_indices": bm.full_q_indices.cpu() if bm.full_q_indices is not None else None, + "BLOCK_SIZE": bm.BLOCK_SIZE, + } + for bm in self._mask_pool + ] + torch.save({"rt_pool": self._rt_pool.cpu(), "mask_tensors": mask_tensors}, path) + + def _load_masks(self, path: str, device: torch.device): + """Load a previously saved mask pool from disk.""" + from torch.nn.attention.flex_attention import BlockMask # type: ignore + data = torch.load(path, map_location=device, weights_only=False) + rt_pool = data["rt_pool"].to(device=device, dtype=self._rt_pool.dtype) + mask_pool = [ + BlockMask( + seq_lengths=td["seq_lengths"], + kv_num_blocks=td["kv_num_blocks"].to(device), + kv_indices=td["kv_indices"].to(device), + full_kv_num_blocks=td["full_kv_num_blocks"].to(device) if td["full_kv_num_blocks"] is not None else None, + full_kv_indices=td["full_kv_indices"].to(device) if td["full_kv_indices"] is not None else None, + q_num_blocks=td["q_num_blocks"].to(device) if td["q_num_blocks"] is not None else None, + q_indices=td["q_indices"].to(device) if td["q_indices"] is not None else None, + full_q_num_blocks=td["full_q_num_blocks"].to(device) if td["full_q_num_blocks"] is not None else None, + full_q_indices=td["full_q_indices"].to(device) if td["full_q_indices"] is not None else None, + BLOCK_SIZE=td["BLOCK_SIZE"], + mask_mod=_causal_mask_mod, + ) + for i, td in enumerate(data["mask_tensors"]) + ] + return mask_pool, rt_pool + + 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, + block_mask=None, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor, torch.Tensor]]]: + batch_size, seq_len, _ = hidden_states.size() + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + 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) + + 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) + + cos, sin = self.rotary_emb(value_states, position_ids) + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + 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.dense_mode: + # Dense fallback: standard causal SDPA — weights are identical to the + # sparse checkpoint, only the attention pattern changes. + 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, + ) + else: + if seq_len != self.seq_len: + raise ValueError( + f"Sparse attention requires fixed seq_len={self.seq_len}, got {seq_len}" + ) + if attention_mask is not None: + raise ValueError( + "Sparse attention expects packed fixed-length batches and no additive attention mask." + ) + if past_key_value is not None or use_cache: + raise ValueError("Sparse generation/KV-cache is not enabled in this iteration") + + if block_mask is None: + idx = torch.randint(len(self._mask_pool), (1,)).item() + block_mask = self._mask_pool[idx] + block_mask = block_mask.to(query_states.device) + attn_output = self._flex_attention( + query_states, + key_states, + value_states, + block_mask=block_mask, + enable_gqa=False, + kernel_options={"BACKEND": "TRITON"}, + ) + + attn_weights = None + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.reshape(batch_size, seq_len, self.attention_hidden_size) + attn_output = self.o_proj(attn_output) + + if not output_attentions: + attn_weights = None + + return attn_output, attn_weights, None + + def pick_mask(self, idx: int): + """Return the BlockMask at position ``idx`` (wraps around pool size).""" + return self._mask_pool[idx % len(self._mask_pool)] + + # ------------------------------------------------------------------ + # Deploy-time (eval-only) reconfiguration helpers. + # + # Everything below is used exclusively by the offline deploy-sparsity + # evaluator (``eval/sparse_eval.py``). It is never called on the training + # hot path, so it must not change any behaviour for existing checkpoints. + # ------------------------------------------------------------------ + def realized_block_density(self) -> float: + r"""Mean attended-block fraction over the current mask pool. + + For each mask and query block ``qb`` the attended causal KV blocks are + the diagonal block plus any sampled random block ``r < qb`` (random + picks with ``r > qb`` are dropped by the ``q_idx >= kv_idx`` clause). + Density is ``sum(attended) / sum(qb + 1)`` averaged over the pool; a + fully-dense causal model returns ``1.0``. + """ + rt = self._rt_pool # (pool, num_blocks, num_random_blocks) int + num_blocks = self.num_blocks + attended = 0 + causal_total = 0 + for qb in range(num_blocks): + causal_total += (qb + 1) * rt.shape[0] + if self.num_random_blocks == 0: + attended += 1 * rt.shape[0] + continue + row = rt[:, qb, :] # (pool, num_random_blocks) + for i in range(rt.shape[0]): + picks = {int(v) for v in row[i].tolist() if int(v) < qb} + attended += 1 + len(picks) # diagonal + distinct past randoms + return float(attended) / float(causal_total) if causal_total else 1.0 + + def set_deploy_masks( + self, + *, + pattern: str = "random", + num_random_blocks: Optional[int] = None, + pool_size: Optional[int] = None, + seed: Optional[int] = None, + ) -> None: + """Rebuild the eval-time mask pool for a chosen DEPLOY pattern. + + ``pattern``: + - ``"random"``: canonical uniform random-block sampling (same algorithm + as training) at the requested ``num_random_blocks`` deploy budget. + - ``"streaming"``: a *training-free* StreamingLLM-style fixed pattern — + every query block attends to its own (local) block plus the first + ``num_random_blocks`` sink blocks. Deterministic, so ``pool_size`` is + forced to 1. + - ``"sliding"``: a deterministic local-window control where + ``num_random_blocks`` is the total block budget, including the + query's diagonal block. + + Only instance state is touched (``num_random_blocks``, ``mask_pool_size``, + ``_rt_pool``, ``_mask_pool``); weights are never modified. + """ + if num_random_blocks is not None: + maximum = self.num_blocks if pattern == "sliding" else self.num_blocks - 1 + minimum = 1 if pattern == "sliding" else 0 + self.num_random_blocks = max( + minimum, min(int(num_random_blocks), maximum) + ) + if pattern in {"streaming", "sliding"}: + pool_size = 1 + if pool_size is None: + pool_size = self.mask_pool_size + self.mask_pool_size = int(pool_size) + + nrb = max(self.num_random_blocks, 1) + rt_device = self._rt_pool.device + self._rt_pool = torch.zeros( + (self.mask_pool_size, self.num_blocks, nrb), + dtype=torch.int16, + device=rt_device, + ) + + if pattern == "streaming": + # Sink blocks 0..num_random_blocks-1 for every query block. Above- + # diagonal sinks (rare, only for the earliest blocks) are dropped by + # the causal clause at mask-build time. + for j in range(self.num_random_blocks): + self._rt_pool[:, :, j] = j + elif pattern == "sliding": + past_budget = self.num_random_blocks - 1 + for qb in range(self.num_blocks): + past_blocks = list( + range(max(0, qb - past_budget), qb) + ) + padded_blocks = past_blocks + [qb] * ( + self.num_random_blocks - len(past_blocks) + ) + if padded_blocks: + self._rt_pool[:, qb, : self.num_random_blocks] = torch.tensor( + padded_blocks, + dtype=self._rt_pool.dtype, + device=rt_device, + ) + elif pattern == "random": + cpu_rng_state = torch.random.get_rng_state() + cuda_rng_state = ( + torch.cuda.get_rng_state() if torch.cuda.is_available() else None + ) + if seed is not None: + # Seed PER-LAYER so each layer samples its OWN distinct mask pool + # (matches training, where layers use distinct per-layer masks). + # Without the per-layer offset, the save/restore below would make + # every layer's pool identical for a shared seed. + torch.manual_seed(int(seed) + 1000003 * int(getattr(self, "_mask_layer_idx", 0))) + if self.num_random_blocks > 0: + for i in range(self.mask_pool_size): + for qb in range(self.num_blocks): + perm = sample_rt_row( + self.num_blocks, self.num_random_blocks, qb, rt_device + ) + self._rt_pool[i, qb, : self.num_random_blocks] = perm.to( + dtype=self._rt_pool.dtype + ) + torch.random.set_rng_state(cpu_rng_state) + if cuda_rng_state is not None: + torch.cuda.set_rng_state(cuda_rng_state) + else: + raise ValueError(f"Unknown deploy pattern '{pattern}'") + + mask_dev = self._rt_pool.device + pool = [] + for i in range(self.mask_pool_size): + bm = self._create_block_mask( + _MaskMod( + self._rt_pool[i].to(mask_dev), self.block_size, self.num_random_blocks + ), + B=None, + H=None, + Q_LEN=self.seq_len, + KV_LEN=self.seq_len, + device=mask_dev, + BLOCK_SIZE=self.block_size, + ) + bm.mask_mod = _causal_mask_mod + pool.append(bm) + self._mask_pool = pool diff --git a/epfl_final/SparseAtt/tests/test_qk_norm_off.py b/epfl_final/SparseAtt/tests/test_qk_norm_off.py new file mode 100644 index 0000000000000000000000000000000000000000..d366fa94a9aeb33f8112e43df222a2e8b75ecab5 --- /dev/null +++ b/epfl_final/SparseAtt/tests/test_qk_norm_off.py @@ -0,0 +1,79 @@ +import torch + +from sparseatt.config import Qwen3Config +from sparseatt.model import Qwen3ForCausalLM + + +def tiny_config(**overrides): + values = dict( + vocab_size=97, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + max_position_embeddings=128, + attention_type="dense", + positional="rope", + qk_norm="on", + use_cache=False, + ) + values.update(overrides) + return Qwen3Config(**values) + + +def logits(model, ids): + model.eval() + with torch.no_grad(): + return model(input_ids=ids, attention_mask=None)["logits"] + + +def test_qk_norm_on_explicit_flag_matches_default(): + torch.manual_seed(0) + default = Qwen3ForCausalLM(tiny_config()) + explicit = Qwen3ForCausalLM(tiny_config(qk_norm="on")) + explicit.load_state_dict(default.state_dict()) + ids = torch.randint(0, 97, (1, 16)) + assert torch.equal(logits(default, ids), logits(explicit, ids)) + + +def test_qk_norm_off_removes_modules_and_has_finite_logits(): + torch.manual_seed(1) + model = Qwen3ForCausalLM(tiny_config(qk_norm="off")) + assert all(not hasattr(layer.self_attn, "q_norm") for layer in model.model.layers) + assert all(not hasattr(layer.self_attn, "k_norm") for layer in model.model.layers) + out = logits(model, torch.randint(0, 97, (1, 32))) + assert torch.isfinite(out).all() + + +def test_qk_norm_off_tau_one_is_exact_noop_and_tau_scales_queries(): + torch.manual_seed(2) + model = Qwen3ForCausalLM(tiny_config(qk_norm="off")) + ids = torch.randint(0, 97, (1, 16)) + baseline = logits(model, ids) + for layer in model.model.layers: + layer.self_attn.set_attention_temperature(1.0) + assert torch.equal(baseline, logits(model, ids)) + for layer in model.model.layers: + layer.self_attn.set_attention_temperature(0.85) + scaled = logits(model, ids) + assert torch.isfinite(scaled).all() + assert not torch.equal(baseline, scaled) + + +def test_qk_norm_temperature_paths_match(): + torch.manual_seed(3) + reference = Qwen3ForCausalLM(tiny_config(qk_norm="on")) + gain_path = Qwen3ForCausalLM(tiny_config(qk_norm="on")) + query_path = Qwen3ForCausalLM(tiny_config(qk_norm="on")) + gain_path.load_state_dict(reference.state_dict()) + query_path.load_state_dict(reference.state_dict()) + tau = 0.85 + with torch.no_grad(): + for layer in gain_path.model.layers: + layer.self_attn.q_norm.weight.div_(tau) + for layer in query_path.model.layers: + layer.self_attn.set_attention_temperature(tau) + ids = torch.randint(0, 97, (1, 16)) + assert torch.allclose(logits(gain_path, ids), logits(query_path, ids), atol=2e-5, rtol=2e-5) diff --git a/epfl_final/SparseAtt/train.py b/epfl_final/SparseAtt/train.py new file mode 100644 index 0000000000000000000000000000000000000000..4bb49798a503426bbe0392dc1001e649ef33cbd4 --- /dev/null +++ b/epfl_final/SparseAtt/train.py @@ -0,0 +1,1593 @@ +import argparse +import contextlib +import copy +import json +import os +import math +import time + +# Persist torch.compile()'d kernels across runs so startup is fast after the first run. +os.environ.setdefault( + "TORCHINDUCTOR_CACHE_DIR", + os.path.expanduser("~/.cache/torch_inductor"), +) +import wandb +import numpy as np +from pathlib import Path +from datetime import datetime +from typing import List, Optional, Sequence, Tuple, Union +from dataclasses import dataclass, asdict +from datasets import load_from_disk + +import torch +import torch.distributed as dist +from torch.utils.data import Dataset, DataLoader +from torch.utils.data.distributed import DistributedSampler +from torch.nn.parallel import DistributedDataParallel as DDP +from torch.optim.lr_scheduler import LambdaLR +from sparseatt.config import Qwen3Config +from sparseatt.model import Qwen3ForCausalLM, load_tokenizer +import torch._inductor.config as inductor_config + +try: + from torchinfo import summary +except ImportError: + summary = None + +torch.set_float32_matmul_precision('high') + + +@dataclass +class TrainConfig: + """Training configuration loaded from JSON.""" + # Model + model_path: str = "../qwen3_06b_base" + model_dtype: str = "bfloat16" + init_from_scratch: bool = False + + # Data + train_path: str = "./data/sample/sample.jsonl" + val_path: Optional[str] = None + max_seq_length: int = 512 + num_workers: int = 8 + # Optional data/sampler seed. When omitted, the sampler follows the + # global training seed; legacy-compatible configs may set this explicitly. + data_seed: Optional[int] = None + + # Training + output_dir: str = "./checkpoints" + resume_from_checkpoint: Optional[str] = None # Path to checkpoint dir to resume from + force_new_run: bool = False # Set True to start fresh even if a previous checkpoint exists + num_epochs: int = 1 + max_steps: Optional[int] = None + per_device_batch_size: int = 2 + gradient_accumulation_steps: int = 8 + learning_rate: float = 1e-5 + weight_decay: float = 0.01 + warmup_steps: int = 100 + lr_scheduler: str = "cosine" + max_grad_norm: float = 1.0 + save_steps: int = 500 + eval_steps: int = 500 + eval_batches: int = 100 + logging_steps: int = 50 + seed: int = 42 + verbose: bool = False + checkpoint_activations: bool = False + # Explicit per-rank DataLoader cursor for a deliberate DDP world-size migration. + resume_micro_step_override: Optional[int] = None + + # Distributed + strategy: str = "ddp" + mixed_precision: str = "bf16" + + # Wandb + wandb_enabled: bool = True + wandb_project: str = "qwen3-training" + wandb_name: Optional[str] = None + wandb_tags: Union[list, None] = None + + # Architecture overrides (applied on top of the base model config.json) + hidden_size: Optional[int] = None + intermediate_size: Optional[int] = None + num_attention_heads: Optional[int] = None + num_key_value_heads: Optional[int] = None + num_hidden_layers: Optional[int] = None + head_dim: Optional[int] = None + attention_dropout: Optional[float] = None + rope_theta: Optional[float] = None + rope_train_position_scale_set: Optional[List[float]] = None + positional: str = "rope" + qk_norm: str = "on" + qk_gain_fp32: bool = False + + # Attention mode + attention_type: str = "dense" + sparse_block_size: int = 128 + sparse_num_random_blocks: int = 1 + sparse_mask_pool_size: int = 1000 + sparse_val_dense_compare: bool = False + sparse_layer_indices: Optional[List[int]] = None + sparse_mask_cache_dir: Optional[str] = None + sparse_mask_seed: int = 0 + sparse_mixed_nrb_set: Optional[List[int]] = None + dense_mode: bool = False + dense_step_intervals: Optional[List[List[int]]] = None + + @classmethod + def from_json(cls, path: str) -> "TrainConfig": + with open(path, "r") as f: + data = json.load(f) + + config = cls() + + # Flatten nested config + if "model" in data: + config.model_path = data["model"].get("path", config.model_path) + config.model_dtype = data["model"].get("dtype", config.model_dtype) + config.init_from_scratch = data["model"].get("init_from_scratch", config.init_from_scratch) + config.hidden_size = data["model"].get("hidden_size", config.hidden_size) + config.intermediate_size = data["model"].get("intermediate_size", config.intermediate_size) + config.num_attention_heads = data["model"].get("num_attention_heads", config.num_attention_heads) + config.num_key_value_heads = data["model"].get("num_key_value_heads", config.num_key_value_heads) + config.num_hidden_layers = data["model"].get("num_hidden_layers", config.num_hidden_layers) + config.head_dim = data["model"].get("head_dim", config.head_dim) + config.attention_dropout = data["model"].get("attention_dropout", config.attention_dropout) + config.rope_theta = data["model"].get("rope_theta", config.rope_theta) + config.rope_train_position_scale_set = data["model"].get( + "rope_train_position_scale_set", config.rope_train_position_scale_set + ) + config.positional = data["model"].get("positional", config.positional) + config.qk_norm = data["model"].get("qk_norm", config.qk_norm) + config.qk_gain_fp32 = data["model"].get("qk_gain_fp32", config.qk_gain_fp32) + + if "data" in data: + config.train_path = data["data"].get("train_path", config.train_path) + config.val_path = data["data"].get("val_path", config.val_path) + config.max_seq_length = data["data"].get("max_seq_length", config.max_seq_length) + config.num_workers = data["data"].get("num_workers", config.num_workers) + config.data_seed = data["data"].get("seed", config.data_seed) + + if "training" in data: + t = data["training"] + config.output_dir = t.get("output_dir", config.output_dir) + config.resume_from_checkpoint = t.get("resume_from_checkpoint", config.resume_from_checkpoint) + config.force_new_run = t.get("force_new_run", config.force_new_run) + config.num_epochs = t.get("num_epochs", config.num_epochs) + config.max_steps = t.get("max_steps", config.max_steps) + config.per_device_batch_size = t.get("per_device_batch_size", config.per_device_batch_size) + config.gradient_accumulation_steps = t.get("gradient_accumulation_steps", config.gradient_accumulation_steps) + config.learning_rate = t.get("learning_rate", config.learning_rate) + config.weight_decay = t.get("weight_decay", config.weight_decay) + config.warmup_steps = t.get("warmup_steps", config.warmup_steps) + config.lr_scheduler = t.get("lr_scheduler", config.lr_scheduler) + config.max_grad_norm = t.get("max_grad_norm", config.max_grad_norm) + config.save_steps = t.get("save_steps", config.save_steps) + config.eval_steps = t.get("eval_steps", config.eval_steps) + config.eval_batches = t.get("eval_batches", config.eval_batches) + config.logging_steps = t.get("logging_steps", config.logging_steps) + config.seed = t.get("seed", config.seed) + config.verbose = t.get("verbose", config.verbose) + config.checkpoint_activations = t.get("checkpoint_activations", config.checkpoint_activations) + config.resume_micro_step_override = t.get("resume_micro_step_override", config.resume_micro_step_override) + if "distributed" in data: + config.strategy = data["distributed"].get("strategy", config.strategy) + config.mixed_precision = data["distributed"].get("mixed_precision", config.mixed_precision) + + if "wandb" in data: + config.wandb_enabled = data["wandb"].get("enabled", config.wandb_enabled) + config.wandb_project = data["wandb"].get("project", config.wandb_project) + config.wandb_name = data["wandb"].get("name", config.wandb_name) + config.wandb_tags = data["wandb"].get("tags", config.wandb_tags) + + attention_cfg = data.get("attention", {}) + config.attention_type = attention_cfg.get("type", data.get("attention_type", config.attention_type)) + config.sparse_block_size = attention_cfg.get("sparse_block_size", data.get("sparse_block_size", config.sparse_block_size)) + config.sparse_num_random_blocks = attention_cfg.get( + "sparse_num_random_blocks", + data.get("sparse_num_random_blocks", config.sparse_num_random_blocks), + ) + config.sparse_mask_pool_size = attention_cfg.get( + "sparse_mask_pool_size", + data.get("sparse_mask_pool_size", config.sparse_mask_pool_size), + ) + config.sparse_mask_seed = attention_cfg.get( + "sparse_mask_seed", data.get("sparse_mask_seed", config.sparse_mask_seed) + ) + config.sparse_mixed_nrb_set = attention_cfg.get( + "sparse_mixed_nrb_set", + data.get("sparse_mixed_nrb_set", config.sparse_mixed_nrb_set), + ) + config.sparse_val_dense_compare = attention_cfg.get( + "sparse_val_dense_compare", + attention_cfg.get( + "sparse_validate_dense_compare", + data.get("sparse_val_dense_compare", config.sparse_val_dense_compare), + ), + ) + config.sparse_layer_indices = attention_cfg.get( + "sparse_layer_indices", + data.get("sparse_layer_indices", config.sparse_layer_indices), + ) + config.sparse_mask_cache_dir = attention_cfg.get( + "sparse_mask_cache_dir", + data.get("sparse_mask_cache_dir", config.sparse_mask_cache_dir), + ) + config.dense_mode = attention_cfg.get( + "dense_mode", + data.get("dense_mode", config.dense_mode), + ) + config.dense_step_intervals = attention_cfg.get( + "dense_step_intervals", + attention_cfg.get( + "dense_mode_step_intervals", + data.get("dense_step_intervals", data.get("dense_mode_step_intervals", config.dense_step_intervals)), + ), + ) + + return config + + +def _normalize_dense_step_intervals( + raw_intervals: Optional[Sequence[Sequence[int]]], +) -> List[Tuple[int, int]]: + if raw_intervals in (None, []): + return [] + if not isinstance(raw_intervals, Sequence) or isinstance(raw_intervals, (str, bytes)): + raise ValueError("dense_step_intervals must be a list of [start, end] pairs") + + normalized: List[Tuple[int, int]] = [] + for idx, interval in enumerate(raw_intervals): + if not isinstance(interval, Sequence) or isinstance(interval, (str, bytes)) or len(interval) != 2: + raise ValueError( + f"dense_step_intervals[{idx}] must be a [start, end] pair" + ) + start, end = interval + if not isinstance(start, int) or not isinstance(end, int): + raise ValueError( + f"dense_step_intervals[{idx}] must contain integer step indices" + ) + if start < 0 or end < 0: + raise ValueError( + f"dense_step_intervals[{idx}] must contain non-negative step indices" + ) + if start > end: + raise ValueError( + f"dense_step_intervals[{idx}] start ({start}) must be <= end ({end})" + ) + normalized.append((start, end)) + + normalized.sort(key=lambda pair: pair[0]) + for idx in range(1, len(normalized)): + prev_start, prev_end = normalized[idx - 1] + start, end = normalized[idx] + if start <= prev_end: + raise ValueError( + "dense_step_intervals must be non-overlapping when endpoints are inclusive; " + f"got [{prev_start}, {prev_end}] and [{start}, {end}]" + ) + return normalized + + +def _step_uses_dense_attention(step: int, dense_step_intervals: Sequence[Tuple[int, int]]) -> bool: + if step <= 0: + return False + return any(start <= step <= end for start, end in dense_step_intervals) + + +def _active_dense_step_interval( + step: int, + dense_step_intervals: Sequence[Tuple[int, int]], +) -> Optional[Tuple[int, int]]: + if step <= 0: + return None + for start, end in dense_step_intervals: + if start <= step <= end: + return start, end + return None + + +def _set_model_dense_mode(model, enabled: bool) -> None: + base_model = _unwrap_base_model(model) + base_model.enable_dense_mode(enabled) + + +def _apply_dense_mode_for_step( + model, + config: TrainConfig, + step: int, + dense_step_intervals: Sequence[Tuple[int, int]], +) -> bool: + if config.attention_type != "sparse_stochastic": + return False + enabled = config.dense_mode or _step_uses_dense_attention(step, dense_step_intervals) + _set_model_dense_mode(model, enabled) + return enabled + + +class TextDataset(Dataset): + """Simple text dataset for language model training.""" + + def __init__( + self, + data_path: str, + tokenizer, + max_seq_length: int = 512, + ): + self.tokenizer = tokenizer + self.max_seq_length = max_seq_length + self.examples = [] + + # Load data + if data_path.endswith(".jsonl"): + self._load_jsonl(data_path) + else: + # Assume it's a directory with HF dataset + self._load_hf_dataset(data_path) + + def _load_jsonl(self, path: str): + + with open(path, "r") as f: + for line in f: + data = json.loads(line) + text = data.get("text", data.get("content", "")) + if text.strip(): + self.examples.append(text) + + def _load_hf_dataset(self, path: str): + + dataset = load_from_disk(path) + for example in dataset: + text = example.get("text", example.get("content", "")) + if text.strip(): + self.examples.append(text) + + def __len__(self): + return len(self.examples) + + def __getitem__(self, idx): + text = self.examples[idx] + + # Tokenize + encodings = self.tokenizer( + text, + truncation=True, + max_length=self.max_seq_length, + padding="max_length", + return_tensors="pt", + ) + + input_ids = encodings["input_ids"].squeeze(0) + attention_mask = encodings["attention_mask"].squeeze(0) + + # Labels are same as input_ids for causal LM + labels = input_ids.clone() + # Mask padding tokens in labels + labels[attention_mask == 0] = -100 + + return { + "input_ids": input_ids, + "attention_mask": attention_mask, + "labels": labels, + } + + +class PreTokenizedDataset(Dataset): + """Dataset for pre-tokenized data (no tokenization at runtime).""" + + def __init__(self, data_path: str, max_seq_length: int = 2048): + """ + Load pre-tokenized data. + + Args: + data_path: Path to directory containing tokenized_data.pt or + path to HuggingFace dataset directory with tokenized data + """ + self.data_path = data_path + self.max_seq_length = max_seq_length + # Check if it's a .pt file or directory + if os.path.isfile(data_path) and data_path.endswith(".pt"): + # Load from .pt file + print(f"Loading pre-tokenized data from {data_path}...") + self.data = torch.load(data_path) + self.is_hf_dataset = False + elif os.path.isdir(data_path): + # Check for tokenized_data.pt in directory + pt_path = os.path.join(data_path, "tokenized_data.pt") + if os.path.exists(pt_path): + print(f"Loading pre-tokenized data from {pt_path}...") + self.data = torch.load(pt_path) + self.is_hf_dataset = False + else: + # Try loading as HuggingFace dataset + print(f"Loading pre-tokenized HF dataset from {data_path}...") + self.data = load_from_disk(data_path) + # Set format to torch to avoid repeated tensor conversions + self.data.set_format(type="torch", columns=["input_ids", "attention_mask"]) + self.is_hf_dataset = True + else: + raise ValueError(f"Invalid data path: {data_path}") + + # Compute pack factor: how many stored sequences to concatenate per example. + # Uses ceil so any max_seq_length >= stored_len works, not just exact multiples. + self.pack_factor = 1 + self.stored_seq_length = max_seq_length + if self.is_hf_dataset and len(self.data) > 0: + stored_len = len(self.data[0]["input_ids"]) + if stored_len > 0 and stored_len < max_seq_length: + self.pack_factor = math.ceil(max_seq_length / stored_len) + self.stored_seq_length = stored_len + + n_packed = len(self.data) // self.pack_factor + print(f"Loaded {len(self.data)} pre-tokenized examples " + f"(stored_len={self.stored_seq_length}, pack_factor={self.pack_factor}, " + f"packed_examples={n_packed})") + + def __len__(self): + return len(self.data) // self.pack_factor + + def __getitem__(self, idx): + if self.is_hf_dataset: + if self.pack_factor > 1: + # Concatenate pack_factor consecutive stored sequences into one + start = idx * self.pack_factor + ids_list, mask_list = [], [] + for i in range(self.pack_factor): + ex = self.data[start + i] + ids_list.append(ex["input_ids"]) + mask_list.append(ex["attention_mask"]) + input_ids = torch.cat(ids_list, dim=0)[:self.max_seq_length] + attention_mask = torch.cat(mask_list, dim=0)[:self.max_seq_length] + else: + example = self.data[idx] + # Data is already in tensor format due to set_format() + input_ids = example["input_ids"][:self.max_seq_length] + attention_mask = example["attention_mask"][:self.max_seq_length] + + labels = input_ids.clone() + # Mask padding tokens in labels + labels[attention_mask == 0] = -100 + + return { + "input_ids": input_ids, + "attention_mask": attention_mask, + "labels": labels, + } + else: + # Data is already in tensor format + return self.data[idx][:self.max_seq_length] + + +def get_lr_scheduler(optimizer, config: TrainConfig, num_training_steps: int): + """Create learning rate scheduler.""" + + num_warmup_steps = config.warmup_steps + + if config.lr_scheduler == "cosine": + def lr_lambda(current_step): + if current_step < num_warmup_steps: + return float(current_step) / float(max(1, num_warmup_steps)) + progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps)) + return max(0.0, 0.5 * (1.0 + math.cos(math.pi * progress))) + elif config.lr_scheduler == "linear": + def lr_lambda(current_step): + if current_step < num_warmup_steps: + return float(current_step) / float(max(1, num_warmup_steps)) + return max(0.0, float(num_training_steps - current_step) / float(max(1, num_training_steps - num_warmup_steps))) + else: # constant + def lr_lambda(current_step): + if current_step < num_warmup_steps: + return float(current_step) / float(max(1, num_warmup_steps)) + return 1.0 + + return LambdaLR(optimizer, lr_lambda) + + +def setup_distributed(): + """Setup distributed training.""" + if "RANK" in os.environ: + rank = int(os.environ["RANK"]) + local_rank = int(os.environ["LOCAL_RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + + dist.init_process_group(backend="nccl") + torch.cuda.set_device(local_rank) + + return rank, local_rank, world_size + else: + return 0, 0, 1 + + +def is_main_process(rank: int) -> bool: + return rank == 0 + + +def _unwrap_model(model): + return model.module if hasattr(model, "module") else model + + +def _unwrap_base_model(model): + unwrapped = _unwrap_model(model) + return unwrapped._orig_mod if hasattr(unwrapped, "_orig_mod") else unwrapped + + +def _model_attention_mask(attention_mask): + """Map a batch attention mask to what the model should receive. + + Returns ``None`` for fully-unpadded (packed) batches so dense / dense_mode + attention layers take SDPA's ``is_causal`` fast path (keeps the Flash backend + eligible and applies the causal upper-triangle skip). Returns the mask + unchanged when padding is present so the explicit additive mask is built. + + The padding check runs here in eager mode — deliberately outside the compiled + model — so it never introduces a graph break / recompile in the hot forward. + Loss masking is unaffected: labels already carry ``-100`` for padded positions, + so dropping an all-ones attention mask does not change the loss. + """ + if attention_mask is None: + return None + return attention_mask if bool((attention_mask == 0).any()) else None + + +def _sample_sparse_indices(model, config: TrainConfig, mask_rng: Optional[np.random.RandomState]): + if config.attention_type != "sparse_stochastic": + return None + if mask_rng is None: + raise ValueError("mask_rng must be provided for sparse_stochastic attention") + + base_model = _unwrap_base_model(model) + layers = base_model.model.layers + pool_size = int(base_model.config.sparse_mask_pool_size) + indices = mask_rng.randint(0, pool_size, size=len(layers)) + # Pre-select BlockMask objects outside the compiled graph so Dynamo never + # sees list-indexing with a dynamic integer. flex_attention's HOP + # registration guards only on BlockMask structural metadata (Q_LEN, + # KV_LEN, BLOCK_SIZE) — not on which random blocks are chosen — so + # different masks across steps cause zero recompilations. + return [ + layer.self_attn.pick_mask(int(idx)) if layer.is_sparse else None + for layer, idx in zip(layers, indices) + ] + + +def _build_dense_eval_model(train_model, device: torch.device, dtype: torch.dtype): + train_base = _unwrap_base_model(train_model) + dense_cfg = copy.deepcopy(train_base.config) + dense_cfg.attention_type = "dense" + return Qwen3ForCausalLM(dense_cfg).to(device=device, dtype=dtype) + + +def _sync_dense_eval_model_weights(train_model, dense_eval_model): + train_base = _unwrap_base_model(train_model) + missing, unexpected = dense_eval_model.load_state_dict(train_base.state_dict(), strict=False) + if missing or unexpected: + print(f"[sync_dense_eval] missing={len(missing)} unexpected={len(unexpected)}") + if missing: print(" missing[:10]:", missing[:10]) + if unexpected: print(" unexpected[:10]:", unexpected[:10]) + + +def evaluate( + model, + eval_loader, + device, + config: TrainConfig, + rank: int, + mask_rng: Optional[np.random.RandomState] = None, + use_sparse_indices: bool = True, + dense_step_intervals: Sequence[Tuple[int, int]] = (), + step_for_mode: Optional[int] = None, +): + """Evaluate model on validation set.""" + model.eval() + total_loss = 0.0 + total_tokens = 0 + num_batches = 0 + + # Divide the target by world_size so the work is distributed evenly + world_size = dist.get_world_size() if dist.is_initialized() else 1 + + # Use math.ceil to ensure you don't round down to 0 on small batch settings + max_eval_batches = math.ceil( + (config.eval_batches * config.gradient_accumulation_steps) / world_size + ) + + autocast_dtype = torch.bfloat16 if config.mixed_precision == "bf16" else torch.float16 + dense_active = False + if step_for_mode is not None: + dense_active = _apply_dense_mode_for_step(model, config, step_for_mode, dense_step_intervals) + + with torch.no_grad(): + for batch in eval_loader: + if num_batches >= max_eval_batches: + break + + input_ids = batch["input_ids"].to(device) + attention_mask = batch["attention_mask"].to(device) + labels = batch["labels"].to(device) + sparse_indices = ( + _sample_sparse_indices(model, config, mask_rng) + if use_sparse_indices and not dense_active + else None + ) + + with torch.amp.autocast("cuda", dtype=autocast_dtype): + outputs = model( + input_ids=input_ids, + attention_mask=_model_attention_mask(attention_mask), + labels=labels, + sparse_indices=sparse_indices, + ) + loss = outputs["loss"] + + total_loss += loss.item() + total_tokens += attention_mask.sum().item() + num_batches += 1 + + avg_loss = total_loss / num_batches if num_batches > 0 else 0.0 + + # evaluate() is called only on rank 0 (gated by is_main_process in the train loop). + # All-reduce across ranks is therefore unreachable and has been removed to avoid + # potential deadlocks if the gating is ever adjusted. + + # Calculate perplexity + perplexity = math.exp(avg_loss) if avg_loss < 100 else float('inf') + + model.train() + return avg_loss, perplexity + + +def run_validation_step( + model, + val_loader, + device, + config: TrainConfig, + rank: int, + mask_rng: Optional[np.random.RandomState], + global_step: int, + dense_eval_model=None, + dense_step_intervals: Sequence[Tuple[int, int]] = (), +): + """Run validation across all ranks to prevent DDP deadlocks, logging only on rank 0.""" + active_dense_interval = _active_dense_step_interval(global_step, dense_step_intervals) + validation_uses_dense_attention = config.dense_mode or active_dense_interval is not None + + # CRITICAL: Every GPU must run this so they stay synchronized in the timeline + eval_loss, eval_perplexity = evaluate( + model, + val_loader, + device, + config, + rank, + mask_rng, + dense_step_intervals=dense_step_intervals, + step_for_mode=global_step, + ) + + dense_eval_loss = None + dense_eval_perplexity = None + if dense_eval_model is not None: + _sync_dense_eval_model_weights(model, dense_eval_model) + dense_eval_loss, dense_eval_perplexity = evaluate( + dense_eval_model, + val_loader, + device, + config, + rank, + mask_rng=None, + use_sparse_indices=False, + ) + + # GATING ONLY THE OUTPUTS: Only rank 0 prints and logs to WandB + if rank == 0: + print(f"Validation at step {global_step}: Loss = {eval_loss:.4f}, Perplexity = {eval_perplexity:.2f}") + + log_dict = { + "eval/loss": eval_loss, + "eval/perplexity": eval_perplexity, + "eval/uses_dense_attention": int(validation_uses_dense_attention), + "eval/uses_sparse_attention": int(not validation_uses_dense_attention), + } + + if active_dense_interval is not None: + interval_start, interval_end = active_dense_interval + log_dict.update({ + "eval/dense_interval_active": 1, + "eval/dense_interval_start": interval_start, + "eval/dense_interval_end": interval_end, + }) + else: + log_dict.update({ + "eval/dense_interval_active": 0, + "eval/dense_interval_start": -1, + "eval/dense_interval_end": -1, + }) + + if dense_eval_model is not None: + print( + f"Dense validation compare at step {global_step}: " + f"Loss = {dense_eval_loss:.4f}, Perplexity = {dense_eval_perplexity:.2f}" + ) + log_dict.update({ + "eval_dense_compare/loss": dense_eval_loss, + "eval_dense_compare/perplexity": dense_eval_perplexity, + }) + + if config.wandb_enabled: + wandb.log(log_dict, step=global_step) + + return log_dict + + return None + + +def load_checkpoint(checkpoint_path: str, model, optimizer, scheduler, device, rank: int, config: TrainConfig | None = None): + """Load model checkpoint and training state. + + Returns + ------- + (global_step, resume_epoch, resume_micro_step) + resume_epoch – epoch index to start the training loop from + resume_micro_step – number of micro-batches already consumed in that epoch; + the loop should skip this many batches before training. + """ + if is_main_process(rank): + print(f"Loading checkpoint from {checkpoint_path}...") + + checkpoint_dir = Path(checkpoint_path) + + # Load model weights + model_path = checkpoint_dir / "model.pt" + if not model_path.exists(): + raise FileNotFoundError(f"Model checkpoint not found: {model_path}") + + model_state = torch.load(model_path, map_location=device, weights_only=True) + # Normalise checkpoint keys: strip "_orig_mod." prefix if present (checkpoint + # was saved from a torch.compile()-wrapped model). + if any(k.startswith("_orig_mod.") for k in model_state): + model_state = {k.removeprefix("_orig_mod."): v for k, v in model_state.items()} + # Drop _rt_pool entries — they are non-learned buffers rebuilt each run and + # their shape depends on sparse_mask_pool_size which can differ across configs. + model_state = {k: v for k, v in model_state.items() if not k.endswith("._rt_pool")} + # Unwrap DDP, then torch.compile's OptimizedModule, so plain keys always match. + target = model.module if hasattr(model, "module") else model + if hasattr(target, "_orig_mod"): + target = target._orig_mod + target.load_state_dict(model_state) + + # Load training state (optimizer, scheduler, step, epoch, micro_step) + training_state_path = checkpoint_dir / "training_state.pt" + if not training_state_path.exists(): + raise FileNotFoundError(f"Training state not found: {training_state_path}") + + training_state = torch.load(training_state_path, map_location=device) + optimizer.load_state_dict(training_state["optimizer"]) + scheduler.load_state_dict(training_state["scheduler"]) + # Override both optimizer and scheduler LR state only when the configured + # schedule base really changed. For an exact resume, the config contains + # the original peak LR while the checkpoint optimizer contains the current + # (scheduled) LR. Unconditionally assigning the peak LR here would create + # a one-step spike before scheduler.step() restores the schedule. + if config is not None: + new_lr = config.learning_rate + checkpoint_base_lrs = list(getattr(scheduler, "base_lrs", [])) + same_schedule_base = bool(checkpoint_base_lrs) and all( + math.isclose(base_lr, new_lr, rel_tol=1e-12, abs_tol=0.0) + for base_lr in checkpoint_base_lrs + ) + if same_schedule_base: + if is_main_process(rank): + current_lrs = [pg["lr"] for pg in optimizer.param_groups] + print( + "Preserving checkpoint LR state for exact resume: " + f"current={current_lrs}, schedule_base={checkpoint_base_lrs}" + ) + else: + optimizer.defaults["lr"] = new_lr + for pg in optimizer.param_groups: + pg["lr"] = new_lr + pg["initial_lr"] = new_lr + if hasattr(scheduler, "base_lrs"): + scheduler.base_lrs = [new_lr for _ in scheduler.base_lrs] + if hasattr(scheduler, "_last_lr"): + scheduler._last_lr = [new_lr for _ in scheduler._last_lr] + if is_main_process(rank): + print(f"Overriding checkpoint LR with config LR: {new_lr}") + step = training_state["step"] + resume_epoch = training_state.get("epoch", 0) + resume_micro_step = training_state.get("micro_step", 0) + + # Restore per-rank RNG states so stochastic ops are exactly reproducible. + rng_path = checkpoint_dir / f"rng_state_rank{rank}.pt" + if rng_path.exists(): + # weights_only=False is required because the file contains numpy arrays + # (torch CPU RNG state and numpy RandomState). The file is always written + # by this codebase so the source is trusted. + rng_state = torch.load(rng_path, map_location="cpu", weights_only=False) + torch.set_rng_state(rng_state["torch_cpu"]) + if torch.cuda.is_available() and rng_state.get("torch_cuda") is not None: + torch.cuda.set_rng_state(rng_state["torch_cuda"]) + np.random.set_state(rng_state["numpy"]) + # mask_rng is a local RandomState; its state is returned so the caller + # can apply it after construction. + else: + if is_main_process(rank): + print(f"Warning: no RNG state file found for rank {rank}; RNG will not be restored.") + rng_state = None + + if is_main_process(rank): + print(f"Resumed from checkpoint at step {step} (epoch {resume_epoch}, micro_step {resume_micro_step})") + + return step, resume_epoch, resume_micro_step, rng_state + + +def save_checkpoint( + model, + optimizer, + scheduler, + step: int, + config: TrainConfig, + rank: int, + epoch: int = 0, + micro_step: int = 0, + mask_rng: Optional[np.random.RandomState] = None, + loader_generator: Optional[torch.Generator] = None, + sampler_seed: Optional[int] = None, +): + """Save model checkpoint, dataloader position, and per-rank RNG states.""" + checkpoint_dir = Path(config.output_dir) / f"checkpoint-{step}" + checkpoint_dir.mkdir(parents=True, exist_ok=True) + + # Every rank saves its own RNG state and DataLoader generator state for exact resume. + rng_state = { + "torch_cpu": torch.get_rng_state(), + "torch_cuda": torch.cuda.get_rng_state() if torch.cuda.is_available() else None, + "numpy": np.random.get_state(), + "mask_rng": mask_rng.get_state() if mask_rng is not None else None, + } + torch.save(rng_state, checkpoint_dir / f"rng_state_rank{rank}.pt") + if loader_generator is not None: + torch.save(loader_generator.get_state(), checkpoint_dir / f"loader_gen_state_rank{rank}.pt") + + if not is_main_process(rank): + return + + # Get model state dict (handle DDP) + model_to_save = model.module if hasattr(model, "module") else model + + # Save model + torch.save(model_to_save.state_dict(), checkpoint_dir / "model.pt") + + # Save optimizer, scheduler, and dataloader position + torch.save({ + "optimizer": optimizer.state_dict(), + "scheduler": scheduler.state_dict(), + "step": step, + "epoch": epoch, + "micro_step": micro_step, + "sampler_seed": sampler_seed, + }, checkpoint_dir / "training_state.pt") + + # Save config + with open(checkpoint_dir / "config.json", "w") as f: + json.dump(vars(model_to_save.config), f, indent=2, default=str) + + print(f"Saved checkpoint to {checkpoint_dir}") + + +def find_latest_checkpoint(output_dir: str) -> Optional[str]: + """Return the path of the most recent checkpoint inside output_dir/run_*/, or None.""" + output_path = Path(output_dir) + if not output_path.exists(): + return None + required_files = ("model.pt", "training_state.pt", "config.json") + run_dirs = sorted(output_path.glob("run_*"), reverse=True) # newest name first + for run_dir in run_dirs: + checkpoints = sorted( + ( + checkpoint + for checkpoint in run_dir.glob("checkpoint-*") + if all( + (checkpoint / filename).is_file() + and (checkpoint / filename).stat().st_size > 0 + for filename in required_files + ) + ), + key=lambda p: int(p.name.split("-")[-1]), + reverse=True, + ) + if checkpoints: + return str(checkpoints[0]) + return None + + +def resolve_sampler_seed(config: TrainConfig) -> int: + """Choose the sampler seed without changing legacy resume order. + + Checkpoints written before the sampler-seed refactor have no marker and + used DistributedSampler's default seed (0). New runs use config.seed and + persist that choice in training_state.pt for subsequent exact resumes. + """ + if not config.resume_from_checkpoint: + return int(config.data_seed if config.data_seed is not None else config.seed) + state_path = Path(config.resume_from_checkpoint) / "training_state.pt" + try: + state = torch.load(state_path, map_location="cpu", weights_only=True) + value = state.get("sampler_seed") + if value is not None: + return int(value) + except Exception: + pass + return int(config.data_seed if config.data_seed is not None else 0) + + +def train(config: TrainConfig): + """Main training function.""" + # Setup distributed + rank, local_rank, world_size = setup_distributed() + device = torch.device(f"cuda:{local_rank}" if torch.cuda.is_available() else "cpu") + + # Set seed + torch.manual_seed(config.seed + rank) + mask_rng = np.random.RandomState(config.seed + 1 + rank) + # mask_rng state is restored after load_checkpoint (below) once resume_mask_rng_state is known. + + if config.attention_type not in {"dense", "sparse_stochastic"}: + raise ValueError( + f"Unsupported attention_type={config.attention_type}. Expected 'dense' or 'sparse_stochastic'." + ) + if config.dense_mode and config.attention_type != "sparse_stochastic": + raise ValueError( + "dense_mode=True requires attention_type='sparse_stochastic' so the sparse layers exist to switch over." + ) + dense_step_intervals = _normalize_dense_step_intervals(config.dense_step_intervals) + config.dense_step_intervals = [list(interval) for interval in dense_step_intervals] or None + if dense_step_intervals and config.attention_type != "sparse_stochastic": + raise ValueError( + "dense_step_intervals requires attention_type='sparse_stochastic'." + ) + if dense_step_intervals and config.dense_mode: + raise ValueError( + "dense_step_intervals cannot be combined with dense_mode=True; use one dense fallback mechanism." + ) + + # Auto-detect latest checkpoint to resume from if not explicitly specified + if not config.resume_from_checkpoint and not config.force_new_run: + latest_ckpt = find_latest_checkpoint(config.output_dir) + if latest_ckpt is not None: + config.resume_from_checkpoint = latest_ckpt + if is_main_process(rank): + print(f"Auto-resuming from latest checkpoint: {latest_ckpt}") + + # Generate unique run ID and update output directory + if config.resume_from_checkpoint and not config.force_new_run: + # Truly resuming the same run — reuse its output directory so wandb + # and optimizer state continue from where they left off. + config.output_dir = str(Path(config.resume_from_checkpoint).parent) + if is_main_process(rank): + print(f"Resuming: Using existing output directory {config.output_dir}") + else: + # Either a completely fresh run or a fork (force_new_run=True with a + # resume_from_checkpoint for weights only) — always create a new dir + # so a new wandb run is created instead of hijacking the old one. + run_id = datetime.now().strftime("%Y%m%d_%H%M%S") + config.output_dir = os.path.join(config.output_dir, f"run_{run_id}") + if is_main_process(rank): + print(f"New run: Checkpoints will be saved to {config.output_dir}") + Path(config.output_dir).mkdir(parents=True, exist_ok=True) + + # Persist wandb run ID early (before heavy setup) so if the process crashes + # during model/dataset loading, recovery will continue in the same wandb run. + # Actual wandb.init() is deferred to just before training loop to ensure + # wall-clock timing only measures training time, not setup time. + wandb_run_id_path = None + prior_wandb_id = None + if config.wandb_enabled and is_main_process(rank): + wandb_run_id_path = Path(config.output_dir) / "wandb_run_id.txt" + if wandb_run_id_path.exists(): + prior_wandb_id = wandb_run_id_path.read_text().strip() + + if is_main_process(rank): + print("Running config:") + print(json.dumps(asdict(config), indent=2, default=str)) + + # Load tokenizer + if is_main_process(rank): + print(f"Loading tokenizer from {config.model_path}...") + tokenizer = load_tokenizer(config.model_path) + + # Set padding token if not set + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + + # Load model + dtype = torch.bfloat16 if config.model_dtype == "bfloat16" else torch.float16 + attention_overrides = { + "attention_type": config.attention_type, + "positional": config.positional, + "qk_norm": config.qk_norm, + "qk_gain_fp32": config.qk_gain_fp32, + "sparse_block_size": int(config.sparse_block_size), + "sparse_num_random_blocks": int(config.sparse_num_random_blocks), + "sparse_mask_pool_size": int(config.sparse_mask_pool_size), + "sparse_mask_seed": int(config.sparse_mask_seed), + "sparse_mixed_nrb_set": config.sparse_mixed_nrb_set, + "sparse_max_seq_length": int(config.max_seq_length), + "sparse_layer_indices": config.sparse_layer_indices, + "sparse_mask_cache_dir": config.sparse_mask_cache_dir, + } + # Add any architecture overrides specified in the train config + arch_override_keys = [ + "hidden_size", "intermediate_size", "num_attention_heads", + "num_key_value_heads", "num_hidden_layers", "head_dim", + "attention_dropout", "rope_theta", "qk_norm", "qk_gain_fp32", + "rope_train_position_scale_set", + ] + for key in arch_override_keys: + val = getattr(config, key, None) + if val is not None: + attention_overrides[key] = val + + # Determine model loading strategy + if config.resume_from_checkpoint: + # When resuming, load config from checkpoint + if is_main_process(rank): + print(f"Resuming from checkpoint: {config.resume_from_checkpoint}") + checkpoint_config_path = os.path.join(config.resume_from_checkpoint, "config.json") + model_config = Qwen3Config.from_json(checkpoint_config_path) + model_config.checkpoint_activations = config.checkpoint_activations + for key, value in attention_overrides.items(): + if hasattr(model_config, key): + setattr(model_config, key, value) + model = Qwen3ForCausalLM(model_config) + model = model.to(device=device, dtype=dtype) + + # Model weights will be loaded later along with optimizer/scheduler + elif config.init_from_scratch: + if is_main_process(rank): + print(f"Initializing model from scratch (random weights)...") + # Load config from pretrained path but initialize with random weights + config_path = os.path.join(config.model_path, "config.json") + model_config = Qwen3Config.from_json(config_path) + model_config.checkpoint_activations = config.checkpoint_activations + for key, value in attention_overrides.items(): + if hasattr(model_config, key): + setattr(model_config, key, value) + model = Qwen3ForCausalLM(model_config) + model = model.to(device=device, dtype=dtype) + else: + if is_main_process(rank): + print(f"Loading pretrained model from {config.model_path}...") + model = Qwen3ForCausalLM.from_pretrained( + config.model_path, + device=device, + dtype=dtype, + config_overrides=attention_overrides, + ) + model.config.checkpoint_activations = config.checkpoint_activations + model.model.checkpoint_activations = config.checkpoint_activations + + if config.dense_mode: + model.enable_dense_mode() + if is_main_process(rank): + print("dense_mode=True: sparse attention layers will use standard causal SDPA.") + + # ``model.to(..., dtype=bf16)`` above intentionally converts the whole + # model. Restore only Q/K RMSNorm gains to FP32 for this experiment; their + # forward method casts the value back to the activation dtype. + if config.qk_gain_fp32: + qk_gain_params = [] + for module in model.modules(): + if getattr(module, "weight_fp32", False): + module.weight.data = module.weight.data.float() + qk_gain_params.append(module.weight) + if len(qk_gain_params) != 56: + raise RuntimeError(f"expected 56 Q/K gain tensors, found {len(qk_gain_params)}") + else: + qk_gain_params = [] + + # Print model architecture BEFORE torch.compile so torchinfo hooks don't + # survive into the compiled region (which would cause recompile storms). + if is_main_process(rank): + print() + print("\n" + "=" * 80) + print("MODEL ARCHITECTURE") + print("=" * 80) + print(f"Model config:") + print(f" vocab_size: {model.config.vocab_size}") + print(f" hidden_size: {model.config.hidden_size}") + print(f" num_hidden_layers: {model.config.num_hidden_layers}") + print(f" num_attention_heads: {model.config.num_attention_heads}") + print(f" num_key_value_heads: {model.config.num_key_value_heads}") + print(f" head_dim: {model.config.head_dim}") + print(f" intermediate_size: {model.config.intermediate_size}") + print(f" max_position_embeddings: {model.config.max_position_embeddings}") + print(f" rope_theta: {model.config.rope_theta}") + print(f" rope_train_position_scale_set: {model.config.rope_train_position_scale_set}") + print(f" attention_type: {model.config.attention_type}") + if model.config.attention_type == "sparse_stochastic": + print(f" sparse_block_size: {model.config.sparse_block_size}") + print(f" sparse_num_random_blocks: {model.config.sparse_num_random_blocks}") + print(f" sparse_mask_pool_size: {model.config.sparse_mask_pool_size}") + print(f" sparse_max_seq_length: {model.config.sparse_max_seq_length}") + + # Count parameters + total_params = sum(p.numel() for p in model.parameters()) + trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad) + print(f"\nParameter count:") + print(f" Total parameters: {total_params:,}") + print(f" Trainable parameters: {trainable_params:,}") + print(f" Non-trainable parameters: {total_params - trainable_params:,}") + print(f" Model size (BF16): {total_params * 2 / 1e9:.2f} GB") + + # Print detailed layer summary if torchinfo is available + if summary is not None and config.verbose: + print("\n" + "-" * 80) + print("LAYER-BY-LAYER SUMMARY") + print("-" * 80) + try: + summary( + model, + input_size=(config.per_device_batch_size, config.max_seq_length), + dtypes=[torch.long], + col_names=["input_size", "output_size", "num_params", "trainable"], + depth=3, + verbose=config.verbose, + ) + except Exception as e: + print(f"Could not generate detailed summary: {e}") + if getattr(model.config, "tie_word_embeddings", False): + print(f" NOTE: torchinfo double-counts tied embed_tokens/lm_head weight " + f"({model.config.vocab_size * model.config.hidden_size:,} params). " + f"True unique parameter count: {total_params:,} ({total_params/1e6:.1f}M)") + + print("=" * 80 + "\n") + + if config.attention_type == "sparse_stochastic": + try: + backends = getattr(inductor_config, "max_autotune_gemm_backends", "") + if isinstance(backends, str): + backend_set = {token.strip() for token in backends.split(",") if token.strip()} + if "ATEN" not in backend_set: + backend_set.add("ATEN") + inductor_config.max_autotune_gemm_backends = ",".join(sorted(backend_set)) + except Exception as exc: + if is_main_process(rank): + print(f"Warning: could not set sparse compile backend fallback: {exc}") + + if os.environ.get("SPARSEATT_DISABLE_COMPILE", "0") != "1": + model = torch.compile(model) + elif is_main_process(rank): + print("torch.compile disabled for this run") + + dense_eval_model = None + if is_main_process(rank) and config.attention_type == "sparse_stochastic" and config.sparse_val_dense_compare: + dense_eval_model = _build_dense_eval_model(model, device, dtype) + print("Sparse validation compare enabled: running both sparse and dense validation passes") + + # Wrap with DDP + if world_size > 1: + # broadcast_buffers=False: _rt_pool is int16 (NCCL unsupported) and is + # rebuilt independently per rank each run, so broadcasting it is wrong anyway. + model = DDP(model, device_ids=[local_rank], output_device=local_rank, broadcast_buffers=False) + + # Create dataset and dataloader + if is_main_process(rank): + print(f"Loading dataset from {config.train_path}...") + + # Check if data is pre-tokenized + is_pretokenized = False + if os.path.isdir(config.train_path): + # Check for tokenized_data.pt or HF dataset with tokenized columns + if os.path.exists(os.path.join(config.train_path, "tokenized_data.pt")): + is_pretokenized = True + elif os.path.exists(os.path.join(config.train_path, "dataset_info.json")): + # HF dataset - check if it has tokenized columns + try: + dataset_test = load_from_disk(config.train_path) + if "input_ids" in dataset_test.column_names: + is_pretokenized = True + except: + pass + elif config.train_path.endswith(".pt"): + is_pretokenized = True + + if is_pretokenized: + if is_main_process(rank): + print("Using pre-tokenized dataset (no runtime tokenization)") + full_dataset = PreTokenizedDataset(config.train_path, max_seq_length=config.max_seq_length) + else: + if is_main_process(rank): + print("Using raw text dataset (tokenizing at runtime)") + full_dataset = TextDataset( + config.train_path, + tokenizer, + max_seq_length=config.max_seq_length, + ) + + # Do 95/5 train/val split + train_size = int(0.95 * len(full_dataset)) + val_size = len(full_dataset) - train_size + + # Set seed for reproducible split + generator = torch.Generator().manual_seed(config.seed) + train_dataset, val_dataset = torch.utils.data.random_split( + full_dataset, [train_size, val_size], generator=generator + ) + + if is_main_process(rank): + print(f"Dataset split: {train_size} train, {val_size} validation") + + # Create train dataloader + # Use a dedicated generator (seeded deterministically) for the DataLoader shuffle + # so that its RNG is independent from the global training RNG and can be + # saved/restored cleanly across checkpoints. + loader_generator = torch.Generator() + loader_generator.manual_seed(config.seed + rank) + sampler_seed = resolve_sampler_seed(config) + if is_main_process(rank): + source = "checkpoint marker" if config.resume_from_checkpoint and sampler_seed != 0 else ("legacy checkpoint fallback" if config.resume_from_checkpoint else "run seed") + print(f"DistributedSampler seed: {sampler_seed} ({source})") + + if world_size > 1: + # Tie the sampler permutation to the same run seed used for model + # initialization and the train/validation split. This makes fixed-seed + # Dense/Mixed/ALiBi runs consume the same data order when their DDP + # world size and loader geometry match. + train_sampler = DistributedSampler( + train_dataset, + num_replicas=world_size, + rank=rank, + shuffle=True, + seed=sampler_seed, + ) + else: + train_sampler = None + + train_loader = DataLoader( + train_dataset, + batch_size=config.per_device_batch_size, + sampler=train_sampler, + shuffle=(train_sampler is None), + generator=loader_generator if train_sampler is None else None, + num_workers=config.num_workers, + pin_memory=True, + drop_last=True, + persistent_workers=True if config.num_workers > 0 else False, + prefetch_factor=2 if config.num_workers > 0 else None, + ) + + # Create validation dataloader + if world_size > 1: + val_sampler = DistributedSampler( + val_dataset, + num_replicas=world_size, + rank=rank, + shuffle=False, + seed=sampler_seed, + ) + else: + val_sampler = None + + val_loader = DataLoader( + val_dataset, + batch_size=config.per_device_batch_size, + sampler=val_sampler, + shuffle=False, + num_workers=config.num_workers, + pin_memory=True, + drop_last=False, + persistent_workers=True if config.num_workers > 0 else False, + prefetch_factor=2 if config.num_workers > 0 else None, + ) + + # Calculate training steps + # len(train_loader) is per-rank micro-batches/epoch (DistributedSampler divides data; + # drop_last=True in the loader already floors to full batches). + micro_batches_per_epoch = len(train_loader) + num_update_steps_per_epoch = micro_batches_per_epoch // config.gradient_accumulation_steps + num_training_steps = config.max_steps or (num_update_steps_per_epoch * config.num_epochs) + effective_batch_size = config.per_device_batch_size * world_size * config.gradient_accumulation_steps + tokens_per_step = effective_batch_size * config.max_seq_length + + if is_main_process(rank): + print(f"Training examples (packed, train split): {len(train_dataset)}") + print(f"Validation examples (packed, val split): {len(val_dataset)}") + print(f"Per-device batch size: {config.per_device_batch_size}") + print(f"World size (GPUs): {world_size}") + print(f"Gradient accumulation steps: {config.gradient_accumulation_steps}") + print(f"Effective batch size (seqs/step): {effective_batch_size}") + print(f"Sequence length: {config.max_seq_length}") + print(f"Tokens per optimizer step: {tokens_per_step:,}") + print(f"Micro-batches per epoch (per rank): {micro_batches_per_epoch}") + print(f"Optimizer steps per epoch: {num_update_steps_per_epoch}") + print(f"Total optimizer steps: {num_training_steps}" + + (f" (capped by max_steps; {num_training_steps/num_update_steps_per_epoch:.2f} epochs)" if config.max_steps else "")) + + # Create optimizer + normal_params = [p for p in model.parameters() if all(p is not q for q in qk_gain_params)] + optimizer = torch.optim.AdamW( + ([{"params": normal_params}, {"params": qk_gain_params}] if qk_gain_params else normal_params), + lr=config.learning_rate, + weight_decay=config.weight_decay, + betas=(0.9, 0.95), + fused=True + ) + + # Create scheduler + scheduler = get_lr_scheduler(optimizer, config, num_training_steps) + + # Load checkpoint if resuming + global_step = 0 + resume_epoch = 0 + resume_micro_step = 0 + resume_mask_rng_state = None + if config.resume_from_checkpoint: + global_step, resume_epoch, resume_micro_step, _rng_state = load_checkpoint( + config.resume_from_checkpoint, + model, + optimizer, + scheduler, + device, + rank, + config=config, + ) + if _rng_state is not None: + resume_mask_rng_state = _rng_state.get("mask_rng") + if resume_mask_rng_state is not None: + mask_rng.set_state(resume_mask_rng_state) + if config.resume_micro_step_override is not None: + if config.resume_micro_step_override < 0: + raise ValueError("resume_micro_step_override must be non-negative") + if is_main_process(rank): + print(f"Overriding resume micro-step {resume_micro_step} -> {config.resume_micro_step_override} for DDP migration") + resume_micro_step = config.resume_micro_step_override + if is_main_process(rank): + print(f"Resuming training from step {global_step}") + + # Mixed precision + scaler = torch.amp.GradScaler("cuda") if config.mixed_precision == "fp16" else None + autocast_dtype = torch.bfloat16 if config.mixed_precision == "bf16" else torch.float16 + + # Initialize wandb just before training to ensure wall-clock timing + # only measures actual training time (not mask creation, model loading, etc.) + if config.wandb_enabled and is_main_process(rank): + if prior_wandb_id: + print(f"Resuming wandb run {prior_wandb_id}") + wandb.init( + project=config.wandb_project, + name=config.wandb_name, + tags=config.wandb_tags, + config=vars(config), + id=prior_wandb_id, + resume="must", + ) + else: + wandb.init( + project=config.wandb_project, + name=config.wandb_name, + tags=config.wandb_tags, + config=vars(config), + ) + wandb_run_id_path.write_text(wandb.run.id) + + # Training loop + model.train() + total_loss = 0.0 + start_time = time.time() + last_log_time = start_time + last_log_step = global_step + _resume_initial_global_step = global_step + + if is_main_process(rank): + print("Starting training...") + + # Restore loader_generator state from checkpoint so the DataLoader shuffle + # on resume is identical to what the original run would have produced. + resume_loader_gen_state = None + if config.resume_from_checkpoint: + _gen_path = Path(config.resume_from_checkpoint) / f"loader_gen_state_rank{rank}.pt" + if _gen_path.exists(): + resume_loader_gen_state = torch.load(_gen_path, map_location="cpu", weights_only=True) + loader_generator.set_state(resume_loader_gen_state) + else: + if is_main_process(rank): + print("Warning: no loader_gen_state file found; DataLoader shuffle order on resume may differ from original run.") + + # On resume, run one validation pass at the restored global_step before any + # new training work, so the baseline metric is logged at the correct step. + if config.resume_from_checkpoint and global_step > 0: + run_validation_step( + model=model, + val_loader=val_loader, + device=device, + config=config, + rank=rank, + mask_rng=mask_rng, + global_step=global_step, + dense_eval_model=dense_eval_model, + dense_step_intervals=dense_step_intervals, + ) + + last_dense_interval_mode = None + + for epoch in range(config.num_epochs): + # Skip epochs that were fully completed before the checkpoint. + if epoch < resume_epoch: + continue + + if train_sampler is not None: + train_sampler.set_epoch(epoch) + + # How many micro-batches to skip at the start of the resume epoch. + skip_micro_batches = resume_micro_step if epoch == resume_epoch else 0 + + for step, batch in enumerate(train_loader): + if step < skip_micro_batches: + # Fast-forward: advance the DataLoader without training. + continue + active_step = global_step + 1 + dense_interval_active = _apply_dense_mode_for_step( + model, + config, + active_step, + dense_step_intervals, + ) + if dense_interval_active != last_dense_interval_mode: + if is_main_process(rank): + mode_name = "dense" if dense_interval_active else "sparse_stochastic" + source = "dense step interval" if dense_interval_active else "configured sparse attention" + print(f"Switching attention mode for step {active_step}: {mode_name} ({source})") + last_dense_interval_mode = dense_interval_active + # Move batch to device with non-blocking transfer + input_ids = batch["input_ids"].to(device, non_blocking=True) + attention_mask = batch["attention_mask"].to(device, non_blocking=True) + labels = batch["labels"].to(device, non_blocking=True) + sparse_indices = None if dense_interval_active else _sample_sparse_indices(model, config, mask_rng) + + # Forward pass with autocast + with torch.amp.autocast("cuda", dtype=autocast_dtype): + outputs = model( + input_ids=input_ids, + attention_mask=_model_attention_mask(attention_mask), + labels=labels, + sparse_indices=sparse_indices, + ) + loss = outputs["loss"] / config.gradient_accumulation_steps + + # Backward pass — suppress DDP gradient sync on non-final accumulation steps + # (avoids 7/8 wasted all-reduces when gradient_accumulation_steps=8) + is_last_accum_step = (step + 1) % config.gradient_accumulation_steps == 0 + sync_ctx = contextlib.nullcontext() if is_last_accum_step else ( + model.no_sync() if hasattr(model, "no_sync") else contextlib.nullcontext() + ) + with sync_ctx: + if scaler is not None: + scaler.scale(loss).backward() + else: + loss.backward() + + # Accumulate UNSCALED loss for logging + total_loss += outputs["loss"].item() + + # Gradient accumulation + if (step + 1) % config.gradient_accumulation_steps == 0: + # Gradient clipping + if scaler is not None: + scaler.unscale_(optimizer) + + torch.nn.utils.clip_grad_norm_(model.parameters(), config.max_grad_norm) + + # Optimizer step + if scaler is not None: + scaler.step(optimizer) + scaler.update() + else: + optimizer.step() + + scheduler.step() + optimizer.zero_grad(set_to_none=True) + + global_step += 1 + + # Logging + if global_step % config.logging_steps == 0 and is_main_process(rank): + avg_loss = total_loss / (config.logging_steps * config.gradient_accumulation_steps) + lr = scheduler.get_last_lr()[0] + + # Calculate speed metrics + current_time = time.time() + + # Instantaneous throughput (since last log) + time_since_last_log = current_time - last_log_time + steps_since_last_log = global_step - last_log_step + steps_per_sec = steps_since_last_log / time_since_last_log if time_since_last_log > 0 else 0 + samples_since_last_log = steps_since_last_log * config.gradient_accumulation_steps * config.per_device_batch_size * world_size + samples_per_sec = samples_since_last_log / time_since_last_log if time_since_last_log > 0 else 0 + ms_per_step = (time_since_last_log / steps_since_last_log * 1000) if steps_since_last_log > 0 else 0 + + # Update last log tracking + last_log_time = current_time + last_log_step = global_step + + print(f"Step {global_step}/{num_training_steps} | Loss: {avg_loss:.4f} | LR: {lr:.2e} | " + f"{ms_per_step:.1f}ms/step | {samples_per_sec:.1f} samples/s") + + if config.wandb_enabled: + wandb.log({ + "train/loss": avg_loss, + "train/learning_rate": lr, + "train/epoch": global_step / num_update_steps_per_epoch, + "train/global_step": global_step, + "speed/steps_per_sec": steps_per_sec, + "speed/samples_per_sec": samples_per_sec, + "speed/ms_per_step": ms_per_step, + }, step=global_step) + + total_loss = 0.0 + + # Evaluation + if global_step % config.eval_steps == 0: + run_validation_step( + model=model, + val_loader=val_loader, + device=device, + config=config, + rank=rank, + mask_rng=mask_rng, + global_step=global_step, + dense_eval_model=dense_eval_model, + dense_step_intervals=dense_step_intervals, + ) + + # Save checkpoint + if global_step % config.save_steps == 0: + save_checkpoint( + model, optimizer, scheduler, global_step, config, rank, + epoch=epoch, micro_step=step + 1, mask_rng=mask_rng, + loader_generator=loader_generator, + sampler_seed=sampler_seed, + ) + + # Check max steps + if config.max_steps and global_step >= config.max_steps: + break + + if config.max_steps and global_step >= config.max_steps: + break + + # Do not overwrite the source checkpoint if a resume exits without a new optimizer step. + if not config.resume_from_checkpoint or global_step > _resume_initial_global_step: + save_checkpoint( + model, optimizer, scheduler, global_step, config, rank, + epoch=epoch, micro_step=step + 1, mask_rng=mask_rng, + loader_generator=loader_generator, + sampler_seed=sampler_seed, + ) + elif is_main_process(rank): + print(f"Skipping final save: no new optimizer step after resume at {global_step}") + + if config.wandb_enabled and is_main_process(rank): + wandb.finish() + + if world_size > 1: + dist.destroy_process_group() + + if is_main_process(rank): + print("Training complete!") + + +def main(): + parser = argparse.ArgumentParser(description="Train Qwen3 model") + parser.add_argument( + "--config", + type=str, + default="/mloscratch/homes/navasard/sparse/SparseAtt/configs/train_config_dense_resume.json", + help="Path to training config JSON", + ) + args = parser.parse_args() + + config = TrainConfig.from_json(args.config) + train(config) + +if __name__ == "__main__": + main()