EPFL ALiBi and training provenance 651-700 of 1184

epfl_final/logs/ALIBI_20260921/ddp4_resume_20260921_195123/seed42.log

@@ -0,0 +1,633 @@
+W0921 19:51:34.094000 148 site-packages/torch/distributed/run.py:851] 
+W0921 19:51:34.094000 148 site-packages/torch/distributed/run.py:851] *****************************************
+W0921 19:51:34.094000 148 site-packages/torch/distributed/run.py:851] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. 
+W0921 19:51:34.094000 148 site-packages/torch/distributed/run.py:851] *****************************************
+Auto-resuming from latest checkpoint: /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-17000
+Resuming: Using existing output directory /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951
+Running config:
+{
+  "model_path": "/mloscratch/homes/navasard/qwen3_06b_base",
+  "model_dtype": "bfloat16",
+  "init_from_scratch": true,
+  "train_path": "/mloscratch/homes/navasard/data/fineweb_large/tokenized",
+  "val_path": null,
+  "max_seq_length": 8192,
+  "num_workers": 8,
+  "output_dir": "/mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951",
+  "resume_from_checkpoint": "/mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-17000",
+  "force_new_run": false,
+  "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,
+  "resume_micro_step_override": null,
+  "strategy": "ddp",
+  "mixed_precision": "bf16",
+  "wandb_enabled": true,
+  "wandb_project": "stochastic_sparsity",
+  "wandb_name": "365m alibi 8k s42",
+  "wandb_tags": [
+    "qwen3",
+    "365m",
+    "8k",
+    "scaleladder",
+    "alibi",
+    "seed42"
+  ],
+  "hidden_size": 768,
+  "intermediate_size": 3072,
+  "num_attention_heads": 6,
+  "num_key_value_heads": 2,
+  "num_hidden_layers": 28,
+  "head_dim": null,
+  "attention_dropout": null,
+  "rope_train_position_scale_set": null,
+  "positional": "alibi",
+  "attention_type": "dense",
+  "sparse_block_size": 128,
+  "sparse_num_random_blocks": 31,
+  "sparse_mask_pool_size": 1,
+  "sparse_val_dense_compare": false,
+  "sparse_layer_indices": null,
+  "sparse_mask_cache_dir": null,
+  "sparse_mask_seed": 0,
+  "sparse_mixed_nrb_set": null,
+  "dense_mode": false,
+  "dense_step_intervals": null
+}
+Loading tokenizer from /mloscratch/homes/navasard/qwen3_06b_base...
+Checkpoint activations: False
+Checkpoint activations: False
+Resuming from checkpoint: /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-17000
+Checkpoint activations: False
+Checkpoint activations: False
+
+
+================================================================================
+MODEL ARCHITECTURE
+================================================================================
+Model config:
+  vocab_size: 151936
+  hidden_size: 768
+  num_hidden_layers: 28
+  num_attention_heads: 6
+  num_key_value_heads: 2
+  head_dim: 128
+  intermediate_size: 3072
+  max_position_embeddings: 32768
+  rope_theta: 1000000
+  rope_train_position_scale_set: None
+  attention_type: dense
+
+Parameter count:
+  Total parameters: 358,958,848
+  Trainable parameters: 358,958,848
+  Non-trainable parameters: 0
+  Model size (BF16): 0.72 GB
+
+--------------------------------------------------------------------------------
+LAYER-BY-LAYER SUMMARY
+--------------------------------------------------------------------------------
+/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/nn/attention/flex_attention.py:1704: UserWarning: flex_attention called without torch.compile() - this will use an unfused implementation that materializes the full scores matrix instead of generating a fused kernel.
+
+SOLUTION: Use torch.compile(flex_attention)(...)
+
+If you want to debug your score_mod/mask_mod, you can set:
+torch.nn.attention.flex_attention._FLEX_ATTENTION_DISABLE_COMPILE_DEBUG = True
+
+This will allow you to use print statements or breakpoints. Note: This doesn't work with the backwards pass and may produce incorrect results.
+  _warn_once(
+=================================================================================================================================================
+Layer (type:depth-idx)                        Input Shape               Output Shape              Param #                   Trainable
+=================================================================================================================================================
+Qwen3ForCausalLM                              [1, 8192]                 --                        --                        True
+├─Qwen3Model: 1-1                             --                        [1, 8192, 768]            --                        True
+│    └─Embedding: 2-1                         [1, 8192]                 [1, 8192, 768]            116,686,848               True
+│    └─ModuleList: 2-2                        --                        --                        --                        True
+│    │    └─Qwen3DecoderLayer: 3-1            [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-2            [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-3            [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-4            [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-5            [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-6            [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-7            [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-8            [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-9            [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-10           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-11           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-12           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-13           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-14           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-15           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-16           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-17           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-18           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-19           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-20           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-21           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-22           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-23           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-24           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-25           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-26           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-27           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-28           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    └─Qwen3RMSNorm: 2-3                      [1, 8192, 768]            [1, 8192, 768]            768                       True
+├─Linear: 1-2                                 [1, 8192, 768]            [1, 8192, 151936]         116,686,848               True
+=================================================================================================================================================
+Total params: 475,645,696
+Trainable params: 475,645,696
+Non-trainable params: 0
+Total mult-adds (Units.MEGABYTES): 475.65
+=================================================================================================================================================
+Input size (MB): 0.07
+Forward/backward pass size (MB): 15598.62
+Params size (MB): 951.29
+Estimated Total Size (MB): 16549.97
+=================================================================================================================================================
+  NOTE: torchinfo double-counts tied embed_tokens/lm_head weight (116,686,848 params). True unique parameter count: 358,958,848 (359.0M)
+================================================================================
+
+Loading dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Using pre-tokenized dataset (no runtime tokenization)
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+
+
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Loaded 12748481 pre-tokenized examples (stored_len=2048, pack_factor=4, packed_examples=3187120)
+Loaded 12748481 pre-tokenized examples (stored_len=2048, pack_factor=4, packed_examples=3187120)
+Loaded 12748481 pre-tokenized examples (stored_len=2048, pack_factor=4, packed_examples=3187120)
+Loaded 12748481 pre-tokenized examples (stored_len=2048, pack_factor=4, packed_examples=3187120)
+Dataset split: 3027764 train, 159356 validation
+Training examples (packed, train split):  3027764
+Validation examples (packed, val split):  159356
+Per-device batch size:                    1
+World size (GPUs):                        4
+Gradient accumulation steps:              5
+Effective batch size (seqs/step):         20
+Sequence length:                          8192
+Tokens per optimizer step:                163,840
+Micro-batches per epoch (per rank):       756941
+Optimizer steps per epoch:                151388
+Total optimizer steps:                    44000  (capped by max_steps; 0.29 epochs)
+Loading checkpoint from /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-17000...
+Preserving checkpoint LR state for exact resume: current=[0.0003451463250063146], schedule_base=[0.0005]
+Resumed from checkpoint at step 17000 (epoch 0, micro_step 85000)
+Resuming training from step 17000
+Resuming wandb run nihe6o4e
+wandb: WARNING `resume` will be ignored since W&B syncing is set to `offline`. Starting a new run with run id nihe6o4e.
+wandb: Tracking run with wandb version 0.26.1
+wandb: W&B syncing is set to `offline` in this directory. Run `wandb online` or set WANDB_MODE=online to enable cloud syncing.
+wandb: Run data is saved locally in /mloscratch/homes/navasard/sparse/SparseAtt/wandb/offline-run-20260921_195241-nihe6o4e
+Starting training...
+Validation at step 17000: Loss = 3.0619, Perplexity = 21.37
+Switching attention mode for step 17001: sparse_stochastic (configured sparse attention)
+Step 17100/44000 | Loss: 3.0356 | LR: 3.43e-04 | 2236.4ms/step | 8.9 samples/s
+Step 17200/44000 | Loss: 3.0205 | LR: 3.42e-04 | 434.3ms/step | 46.0 samples/s
+Validation at step 17250: Loss = 3.0589, Perplexity = 21.30
+Step 17300/44000 | Loss: 3.0223 | LR: 3.40e-04 | 452.7ms/step | 44.2 samples/s
+Step 17400/44000 | Loss: 3.0272 | LR: 3.38e-04 | 436.1ms/step | 45.9 samples/s
+Step 17500/44000 | Loss: 2.9948 | LR: 3.37e-04 | 436.2ms/step | 45.9 samples/s
+Validation at step 17500: Loss = 3.0545, Perplexity = 21.21
+Step 17600/44000 | Loss: 3.0312 | LR: 3.35e-04 | 453.1ms/step | 44.1 samples/s
+Step 17700/44000 | Loss: 3.0148 | LR: 3.33e-04 | 436.6ms/step | 45.8 samples/s
+Validation at step 17750: Loss = 3.0525, Perplexity = 21.17
+Step 17800/44000 | Loss: 3.0038 | LR: 3.32e-04 | 454.9ms/step | 44.0 samples/s
+Step 17900/44000 | Loss: 3.0169 | LR: 3.30e-04 | 439.2ms/step | 45.5 samples/s
+Step 18000/44000 | Loss: 3.0159 | LR: 3.28e-04 | 437.5ms/step | 45.7 samples/s
+Validation at step 18000: Loss = 3.0504, Perplexity = 21.12
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-18000
+Step 18100/44000 | Loss: 3.0117 | LR: 3.26e-04 | 484.5ms/step | 41.3 samples/s
+Step 18200/44000 | Loss: 2.9962 | LR: 3.25e-04 | 437.2ms/step | 45.7 samples/s
+Validation at step 18250: Loss = 3.0481, Perplexity = 21.08
+Step 18300/44000 | Loss: 3.0004 | LR: 3.23e-04 | 453.9ms/step | 44.1 samples/s
+Step 18400/44000 | Loss: 3.0141 | LR: 3.21e-04 | 437.1ms/step | 45.8 samples/s
+Step 18500/44000 | Loss: 3.0007 | LR: 3.19e-04 | 436.7ms/step | 45.8 samples/s
+Validation at step 18500: Loss = 3.0452, Perplexity = 21.01
+Step 18600/44000 | Loss: 2.9966 | LR: 3.18e-04 | 454.0ms/step | 44.1 samples/s
+Step 18700/44000 | Loss: 2.9863 | LR: 3.16e-04 | 436.7ms/step | 45.8 samples/s
+Validation at step 18750: Loss = 3.0422, Perplexity = 20.95
+Step 18800/44000 | Loss: 2.9931 | LR: 3.14e-04 | 453.8ms/step | 44.1 samples/s
+Step 18900/44000 | Loss: 2.9856 | LR: 3.12e-04 | 436.9ms/step | 45.8 samples/s
+Step 19000/44000 | Loss: 2.9980 | LR: 3.11e-04 | 436.5ms/step | 45.8 samples/s
+Validation at step 19000: Loss = 3.0399, Perplexity = 20.90
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-19000
+Step 19100/44000 | Loss: 2.9832 | LR: 3.09e-04 | 482.1ms/step | 41.5 samples/s
+Step 19200/44000 | Loss: 2.9938 | LR: 3.07e-04 | 436.2ms/step | 45.9 samples/s
+Validation at step 19250: Loss = 3.0368, Perplexity = 20.84
+Step 19300/44000 | Loss: 2.9943 | LR: 3.05e-04 | 453.2ms/step | 44.1 samples/s
+Step 19400/44000 | Loss: 2.9984 | LR: 3.04e-04 | 436.6ms/step | 45.8 samples/s
+Step 19500/44000 | Loss: 3.0033 | LR: 3.02e-04 | 437.0ms/step | 45.8 samples/s
+Validation at step 19500: Loss = 3.0332, Perplexity = 20.76
+Step 19600/44000 | Loss: 3.0076 | LR: 3.00e-04 | 453.7ms/step | 44.1 samples/s
+Step 19700/44000 | Loss: 2.9789 | LR: 2.98e-04 | 438.4ms/step | 45.6 samples/s
+Validation at step 19750: Loss = 3.0310, Perplexity = 20.72
+Step 19800/44000 | Loss: 2.9954 | LR: 2.96e-04 | 455.3ms/step | 43.9 samples/s
+Step 19900/44000 | Loss: 2.9913 | LR: 2.95e-04 | 438.0ms/step | 45.7 samples/s
+Step 20000/44000 | Loss: 2.9879 | LR: 2.93e-04 | 437.6ms/step | 45.7 samples/s
+Validation at step 20000: Loss = 3.0268, Perplexity = 20.63
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-20000
+Step 20100/44000 | Loss: 2.9844 | LR: 2.91e-04 | 485.1ms/step | 41.2 samples/s
+Step 20200/44000 | Loss: 2.9934 | LR: 2.89e-04 | 437.3ms/step | 45.7 samples/s
+Validation at step 20250: Loss = 3.0236, Perplexity = 20.57
+Step 20300/44000 | Loss: 3.0058 | LR: 2.88e-04 | 454.4ms/step | 44.0 samples/s
+Step 20400/44000 | Loss: 2.9815 | LR: 2.86e-04 | 437.3ms/step | 45.7 samples/s
+Step 20500/44000 | Loss: 2.9850 | LR: 2.84e-04 | 437.1ms/step | 45.8 samples/s
+Validation at step 20500: Loss = 3.0216, Perplexity = 20.52
+Step 20600/44000 | Loss: 2.9519 | LR: 2.82e-04 | 454.9ms/step | 44.0 samples/s
+Step 20700/44000 | Loss: 2.9983 | LR: 2.80e-04 | 437.8ms/step | 45.7 samples/s
+Validation at step 20750: Loss = 3.0200, Perplexity = 20.49
+Step 20800/44000 | Loss: 2.9799 | LR: 2.79e-04 | 454.8ms/step | 44.0 samples/s
+Step 20900/44000 | Loss: 2.9805 | LR: 2.77e-04 | 436.7ms/step | 45.8 samples/s
+Step 21000/44000 | Loss: 2.9590 | LR: 2.75e-04 | 436.4ms/step | 45.8 samples/s
+Validation at step 21000: Loss = 3.0179, Perplexity = 20.45
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-21000
+Step 21100/44000 | Loss: 2.9825 | LR: 2.73e-04 | 490.5ms/step | 40.8 samples/s
+Step 21200/44000 | Loss: 2.9731 | LR: 2.71e-04 | 437.0ms/step | 45.8 samples/s
+Validation at step 21250: Loss = 3.0155, Perplexity = 20.40
+Step 21300/44000 | Loss: 2.9784 | LR: 2.70e-04 | 454.1ms/step | 44.0 samples/s
+Step 21400/44000 | Loss: 2.9674 | LR: 2.68e-04 | 436.8ms/step | 45.8 samples/s
+Step 21500/44000 | Loss: 2.9688 | LR: 2.66e-04 | 437.3ms/step | 45.7 samples/s
+Validation at step 21500: Loss = 3.0139, Perplexity = 20.37
+Step 21600/44000 | Loss: 2.9522 | LR: 2.64e-04 | 455.5ms/step | 43.9 samples/s
+Step 21700/44000 | Loss: 2.9937 | LR: 2.62e-04 | 437.9ms/step | 45.7 samples/s
+Validation at step 21750: Loss = 3.0102, Perplexity = 20.29
+Step 21800/44000 | Loss: 2.9631 | LR: 2.60e-04 | 454.2ms/step | 44.0 samples/s
+Step 21900/44000 | Loss: 2.9742 | LR: 2.59e-04 | 437.0ms/step | 45.8 samples/s
+Step 22000/44000 | Loss: 2.9705 | LR: 2.57e-04 | 437.3ms/step | 45.7 samples/s
+Validation at step 22000: Loss = 3.0074, Perplexity = 20.23
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-22000
+Step 22100/44000 | Loss: 2.9864 | LR: 2.55e-04 | 485.0ms/step | 41.2 samples/s
+Step 22200/44000 | Loss: 2.9646 | LR: 2.53e-04 | 436.9ms/step | 45.8 samples/s
+Validation at step 22250: Loss = 3.0050, Perplexity = 20.19
+Step 22300/44000 | Loss: 2.9635 | LR: 2.51e-04 | 453.6ms/step | 44.1 samples/s
+Step 22400/44000 | Loss: 2.9777 | LR: 2.50e-04 | 436.5ms/step | 45.8 samples/s
+Step 22500/44000 | Loss: 2.9654 | LR: 2.48e-04 | 436.4ms/step | 45.8 samples/s
+Validation at step 22500: Loss = 3.0030, Perplexity = 20.15
+Step 22600/44000 | Loss: 2.9393 | LR: 2.46e-04 | 454.1ms/step | 44.0 samples/s
+Step 22700/44000 | Loss: 2.9572 | LR: 2.44e-04 | 437.2ms/step | 45.7 samples/s
+Validation at step 22750: Loss = 3.0006, Perplexity = 20.10
+Step 22800/44000 | Loss: 2.9668 | LR: 2.42e-04 | 453.8ms/step | 44.1 samples/s
+Step 22900/44000 | Loss: 2.9650 | LR: 2.40e-04 | 437.1ms/step | 45.8 samples/s
+Step 23000/44000 | Loss: 2.9784 | LR: 2.39e-04 | 436.5ms/step | 45.8 samples/s
+Validation at step 23000: Loss = 2.9986, Perplexity = 20.06
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-23000
+Step 23100/44000 | Loss: 2.9535 | LR: 2.37e-04 | 484.0ms/step | 41.3 samples/s
+Step 23200/44000 | Loss: 2.9648 | LR: 2.35e-04 | 436.4ms/step | 45.8 samples/s
+Validation at step 23250: Loss = 2.9962, Perplexity = 20.01
+Step 23300/44000 | Loss: 2.9451 | LR: 2.33e-04 | 453.7ms/step | 44.1 samples/s
+Step 23400/44000 | Loss: 2.9829 | LR: 2.31e-04 | 436.7ms/step | 45.8 samples/s
+Step 23500/44000 | Loss: 2.9454 | LR: 2.30e-04 | 436.2ms/step | 45.8 samples/s
+Validation at step 23500: Loss = 2.9933, Perplexity = 19.95
+Step 23600/44000 | Loss: 2.9508 | LR: 2.28e-04 | 454.1ms/step | 44.0 samples/s
+Step 23700/44000 | Loss: 2.9536 | LR: 2.26e-04 | 436.2ms/step | 45.9 samples/s
+Validation at step 23750: Loss = 2.9917, Perplexity = 19.92
+Step 23800/44000 | Loss: 2.9361 | LR: 2.24e-04 | 453.9ms/step | 44.1 samples/s
+Step 23900/44000 | Loss: 2.9423 | LR: 2.22e-04 | 436.6ms/step | 45.8 samples/s
+Step 24000/44000 | Loss: 2.9561 | LR: 2.21e-04 | 435.6ms/step | 45.9 samples/s
+Validation at step 24000: Loss = 2.9895, Perplexity = 19.88
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-24000
+Step 24100/44000 | Loss: 2.9405 | LR: 2.19e-04 | 484.5ms/step | 41.3 samples/s
+Step 24200/44000 | Loss: 2.9431 | LR: 2.17e-04 | 436.2ms/step | 45.8 samples/s
+Validation at step 24250: Loss = 2.9873, Perplexity = 19.83
+Step 24300/44000 | Loss: 2.9489 | LR: 2.15e-04 | 453.3ms/step | 44.1 samples/s
+Step 24400/44000 | Loss: 2.9427 | LR: 2.13e-04 | 436.0ms/step | 45.9 samples/s
+Step 24500/44000 | Loss: 2.9479 | LR: 2.12e-04 | 436.1ms/step | 45.9 samples/s
+Validation at step 24500: Loss = 2.9859, Perplexity = 19.80
+Step 24600/44000 | Loss: 2.9462 | LR: 2.10e-04 | 454.0ms/step | 44.1 samples/s
+Step 24700/44000 | Loss: 2.9486 | LR: 2.08e-04 | 436.8ms/step | 45.8 samples/s
+Validation at step 24750: Loss = 2.9832, Perplexity = 19.75
+Step 24800/44000 | Loss: 2.9388 | LR: 2.06e-04 | 453.8ms/step | 44.1 samples/s
+Step 24900/44000 | Loss: 2.9570 | LR: 2.04e-04 | 436.3ms/step | 45.8 samples/s
+Step 25000/44000 | Loss: 2.9629 | LR: 2.03e-04 | 436.6ms/step | 45.8 samples/s
+Validation at step 25000: Loss = 2.9809, Perplexity = 19.70
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-25000
+Step 25100/44000 | Loss: 2.9405 | LR: 2.01e-04 | 484.2ms/step | 41.3 samples/s
+Step 25200/44000 | Loss: 2.9285 | LR: 1.99e-04 | 436.8ms/step | 45.8 samples/s
+Validation at step 25250: Loss = 2.9801, Perplexity = 19.69
+Step 25300/44000 | Loss: 2.9436 | LR: 1.97e-04 | 453.4ms/step | 44.1 samples/s
+Step 25400/44000 | Loss: 2.9376 | LR: 1.96e-04 | 438.2ms/step | 45.6 samples/s
+Step 25500/44000 | Loss: 2.9267 | LR: 1.94e-04 | 439.2ms/step | 45.5 samples/s
+Validation at step 25500: Loss = 2.9779, Perplexity = 19.65
+Step 25600/44000 | Loss: 2.9379 | LR: 1.92e-04 | 455.8ms/step | 43.9 samples/s
+Step 25700/44000 | Loss: 2.9225 | LR: 1.90e-04 | 437.9ms/step | 45.7 samples/s
+Validation at step 25750: Loss = 2.9761, Perplexity = 19.61
+Step 25800/44000 | Loss: 2.9375 | LR: 1.88e-04 | 455.6ms/step | 43.9 samples/s
+Step 25900/44000 | Loss: 2.9144 | LR: 1.87e-04 | 438.3ms/step | 45.6 samples/s
+Step 26000/44000 | Loss: 2.9529 | LR: 1.85e-04 | 436.4ms/step | 45.8 samples/s
+Validation at step 26000: Loss = 2.9744, Perplexity = 19.58
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-26000
+Step 26100/44000 | Loss: 2.9262 | LR: 1.83e-04 | 484.3ms/step | 41.3 samples/s
+Step 26200/44000 | Loss: 2.9411 | LR: 1.81e-04 | 437.3ms/step | 45.7 samples/s
+Validation at step 26250: Loss = 2.9733, Perplexity = 19.56
+Step 26300/44000 | Loss: 2.9262 | LR: 1.80e-04 | 454.3ms/step | 44.0 samples/s
+Step 26400/44000 | Loss: 2.9284 | LR: 1.78e-04 | 438.1ms/step | 45.6 samples/s
+Step 26500/44000 | Loss: 2.9199 | LR: 1.76e-04 | 438.1ms/step | 45.6 samples/s
+Validation at step 26500: Loss = 2.9715, Perplexity = 19.52
+Step 26600/44000 | Loss: 2.9185 | LR: 1.74e-04 | 454.3ms/step | 44.0 samples/s
+Step 26700/44000 | Loss: 2.9165 | LR: 1.73e-04 | 437.9ms/step | 45.7 samples/s
+Validation at step 26750: Loss = 2.9702, Perplexity = 19.50
+Step 26800/44000 | Loss: 2.9141 | LR: 1.71e-04 | 454.0ms/step | 44.1 samples/s
+Step 26900/44000 | Loss: 2.9321 | LR: 1.69e-04 | 436.6ms/step | 45.8 samples/s
+Step 27000/44000 | Loss: 2.9215 | LR: 1.68e-04 | 437.6ms/step | 45.7 samples/s
+Validation at step 27000: Loss = 2.9684, Perplexity = 19.46
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-27000
+Step 27100/44000 | Loss: 2.9234 | LR: 1.66e-04 | 484.4ms/step | 41.3 samples/s
+Step 27200/44000 | Loss: 2.9353 | LR: 1.64e-04 | 437.6ms/step | 45.7 samples/s
+Validation at step 27250: Loss = 2.9673, Perplexity = 19.44
+Step 27300/44000 | Loss: 2.9169 | LR: 1.62e-04 | 454.0ms/step | 44.1 samples/s
+Step 27400/44000 | Loss: 2.9509 | LR: 1.61e-04 | 437.0ms/step | 45.8 samples/s
+Step 27500/44000 | Loss: 2.9056 | LR: 1.59e-04 | 437.0ms/step | 45.8 samples/s
+Validation at step 27500: Loss = 2.9652, Perplexity = 19.40
+Step 27600/44000 | Loss: 2.9206 | LR: 1.57e-04 | 454.8ms/step | 44.0 samples/s
+Step 27700/44000 | Loss: 2.9347 | LR: 1.56e-04 | 436.9ms/step | 45.8 samples/s
+Validation at step 27750: Loss = 2.9641, Perplexity = 19.38
+Step 27800/44000 | Loss: 2.9121 | LR: 1.54e-04 | 454.3ms/step | 44.0 samples/s
+Step 27900/44000 | Loss: 2.9319 | LR: 1.52e-04 | 436.4ms/step | 45.8 samples/s
+Step 28000/44000 | Loss: 2.9141 | LR: 1.51e-04 | 437.1ms/step | 45.8 samples/s
+Validation at step 28000: Loss = 2.9623, Perplexity = 19.34
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-28000
+Step 28100/44000 | Loss: 2.9271 | LR: 1.49e-04 | 483.5ms/step | 41.4 samples/s
+Step 28200/44000 | Loss: 2.9216 | LR: 1.47e-04 | 436.8ms/step | 45.8 samples/s
+Validation at step 28250: Loss = 2.9616, Perplexity = 19.33
+Step 28300/44000 | Loss: 2.9263 | LR: 1.46e-04 | 454.8ms/step | 44.0 samples/s
+Step 28400/44000 | Loss: 2.9113 | LR: 1.44e-04 | 437.8ms/step | 45.7 samples/s
+Step 28500/44000 | Loss: 2.9059 | LR: 1.42e-04 | 437.9ms/step | 45.7 samples/s
+Validation at step 28500: Loss = 2.9599, Perplexity = 19.30
+Step 28600/44000 | Loss: 2.9311 | LR: 1.41e-04 | 453.7ms/step | 44.1 samples/s
+Step 28700/44000 | Loss: 2.9267 | LR: 1.39e-04 | 436.9ms/step | 45.8 samples/s
+Validation at step 28750: Loss = 2.9588, Perplexity = 19.27
+Step 28800/44000 | Loss: 2.9419 | LR: 1.38e-04 | 453.3ms/step | 44.1 samples/s
+Step 28900/44000 | Loss: 2.9289 | LR: 1.36e-04 | 437.4ms/step | 45.7 samples/s
+Step 29000/44000 | Loss: 2.9095 | LR: 1.34e-04 | 437.8ms/step | 45.7 samples/s
+Validation at step 29000: Loss = 2.9578, Perplexity = 19.26
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-29000
+Step 29100/44000 | Loss: 2.9032 | LR: 1.33e-04 | 486.1ms/step | 41.1 samples/s
+Step 29200/44000 | Loss: 2.9121 | LR: 1.31e-04 | 438.3ms/step | 45.6 samples/s
+Validation at step 29250: Loss = 2.9573, Perplexity = 19.25
+Step 29300/44000 | Loss: 2.9168 | LR: 1.29e-04 | 455.1ms/step | 43.9 samples/s
+Step 29400/44000 | Loss: 2.8960 | LR: 1.28e-04 | 437.7ms/step | 45.7 samples/s
+Step 29500/44000 | Loss: 2.9227 | LR: 1.26e-04 | 437.9ms/step | 45.7 samples/s
+Validation at step 29500: Loss = 2.9568, Perplexity = 19.24
+Step 29600/44000 | Loss: 2.9112 | LR: 1.25e-04 | 454.9ms/step | 44.0 samples/s
+Step 29700/44000 | Loss: 2.8995 | LR: 1.23e-04 | 437.4ms/step | 45.7 samples/s
+Validation at step 29750: Loss = 2.9562, Perplexity = 19.22
+Step 29800/44000 | Loss: 2.9073 | LR: 1.22e-04 | 454.4ms/step | 44.0 samples/s
+Step 29900/44000 | Loss: 2.9137 | LR: 1.20e-04 | 437.1ms/step | 45.8 samples/s
+Step 30000/44000 | Loss: 2.9133 | LR: 1.19e-04 | 437.9ms/step | 45.7 samples/s
+Validation at step 30000: Loss = 2.9545, Perplexity = 19.19
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-30000
+Step 30100/44000 | Loss: 2.9091 | LR: 1.17e-04 | 484.7ms/step | 41.3 samples/s
+Step 30200/44000 | Loss: 2.9108 | LR: 1.15e-04 | 437.9ms/step | 45.7 samples/s
+Validation at step 30250: Loss = 2.9542, Perplexity = 19.19
+Step 30300/44000 | Loss: 2.9104 | LR: 1.14e-04 | 454.4ms/step | 44.0 samples/s
+Step 30400/44000 | Loss: 2.9291 | LR: 1.12e-04 | 437.2ms/step | 45.7 samples/s
+Step 30500/44000 | Loss: 2.9121 | LR: 1.11e-04 | 437.4ms/step | 45.7 samples/s
+Validation at step 30500: Loss = 2.9532, Perplexity = 19.17
+Step 30600/44000 | Loss: 2.9144 | LR: 1.09e-04 | 454.3ms/step | 44.0 samples/s
+Step 30700/44000 | Loss: 2.9138 | LR: 1.08e-04 | 436.6ms/step | 45.8 samples/s
+Validation at step 30750: Loss = 2.9524, Perplexity = 19.15
+Step 30800/44000 | Loss: 2.9148 | LR: 1.06e-04 | 454.5ms/step | 44.0 samples/s
+Step 30900/44000 | Loss: 2.9137 | LR: 1.05e-04 | 436.8ms/step | 45.8 samples/s
+Step 31000/44000 | Loss: 2.9124 | LR: 1.03e-04 | 436.9ms/step | 45.8 samples/s
+Validation at step 31000: Loss = 2.9516, Perplexity = 19.14
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-31000
+Step 31100/44000 | Loss: 2.9249 | LR: 1.02e-04 | 484.3ms/step | 41.3 samples/s
+Step 31200/44000 | Loss: 2.9104 | LR: 1.00e-04 | 437.0ms/step | 45.8 samples/s
+Validation at step 31250: Loss = 2.9510, Perplexity = 19.12
+Step 31300/44000 | Loss: 2.9000 | LR: 9.90e-05 | 454.7ms/step | 44.0 samples/s
+Step 31400/44000 | Loss: 2.9000 | LR: 9.76e-05 | 437.1ms/step | 45.8 samples/s
+Step 31500/44000 | Loss: 2.9047 | LR: 9.62e-05 | 437.8ms/step | 45.7 samples/s
+Validation at step 31500: Loss = 2.9502, Perplexity = 19.11
+Step 31600/44000 | Loss: 2.9292 | LR: 9.47e-05 | 454.6ms/step | 44.0 samples/s
+Step 31700/44000 | Loss: 2.8975 | LR: 9.33e-05 | 437.5ms/step | 45.7 samples/s
+Validation at step 31750: Loss = 2.9493, Perplexity = 19.09
+Step 31800/44000 | Loss: 2.8970 | LR: 9.19e-05 | 453.5ms/step | 44.1 samples/s
+Step 31900/44000 | Loss: 2.9086 | LR: 9.05e-05 | 437.5ms/step | 45.7 samples/s
+Step 32000/44000 | Loss: 2.9162 | LR: 8.91e-05 | 437.6ms/step | 45.7 samples/s
+Validation at step 32000: Loss = 2.9491, Perplexity = 19.09
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-32000
+Step 32100/44000 | Loss: 2.8867 | LR: 8.77e-05 | 484.2ms/step | 41.3 samples/s
+Step 32200/44000 | Loss: 2.9098 | LR: 8.63e-05 | 437.1ms/step | 45.8 samples/s
+Validation at step 32250: Loss = 2.9484, Perplexity = 19.08
+Step 32300/44000 | Loss: 2.9202 | LR: 8.50e-05 | 454.4ms/step | 44.0 samples/s
+Step 32400/44000 | Loss: 2.9078 | LR: 8.36e-05 | 437.7ms/step | 45.7 samples/s
+Step 32500/44000 | Loss: 2.9058 | LR: 8.23e-05 | 437.6ms/step | 45.7 samples/s
+Validation at step 32500: Loss = 2.9481, Perplexity = 19.07
+Step 32600/44000 | Loss: 2.9019 | LR: 8.09e-05 | 454.4ms/step | 44.0 samples/s
+Step 32700/44000 | Loss: 2.9039 | LR: 7.96e-05 | 437.8ms/step | 45.7 samples/s
+Validation at step 32750: Loss = 2.9473, Perplexity = 19.05
+Step 32800/44000 | Loss: 2.9323 | LR: 7.83e-05 | 454.6ms/step | 44.0 samples/s
+Step 32900/44000 | Loss: 2.9106 | LR: 7.70e-05 | 437.7ms/step | 45.7 samples/s
+Step 33000/44000 | Loss: 2.9091 | LR: 7.56e-05 | 439.2ms/step | 45.5 samples/s
+Validation at step 33000: Loss = 2.9470, Perplexity = 19.05
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-33000
+Step 33100/44000 | Loss: 2.8922 | LR: 7.44e-05 | 486.3ms/step | 41.1 samples/s
+Step 33200/44000 | Loss: 2.9036 | LR: 7.31e-05 | 439.0ms/step | 45.6 samples/s
+Validation at step 33250: Loss = 2.9466, Perplexity = 19.04
+Step 33300/44000 | Loss: 2.9246 | LR: 7.18e-05 | 456.0ms/step | 43.9 samples/s
+Step 33400/44000 | Loss: 2.9131 | LR: 7.05e-05 | 439.2ms/step | 45.5 samples/s
+Step 33500/44000 | Loss: 2.9039 | LR: 6.93e-05 | 437.9ms/step | 45.7 samples/s
+Validation at step 33500: Loss = 2.9464, Perplexity = 19.04
+Step 33600/44000 | Loss: 2.9110 | LR: 6.80e-05 | 454.0ms/step | 44.0 samples/s
+Step 33700/44000 | Loss: 2.9098 | LR: 6.68e-05 | 437.2ms/step | 45.8 samples/s
+Validation at step 33750: Loss = 2.9454, Perplexity = 19.02
+Step 33800/44000 | Loss: 2.8883 | LR: 6.55e-05 | 455.8ms/step | 43.9 samples/s
+Step 33900/44000 | Loss: 2.9103 | LR: 6.43e-05 | 438.2ms/step | 45.6 samples/s
+Step 34000/44000 | Loss: 2.8978 | LR: 6.31e-05 | 437.2ms/step | 45.7 samples/s
+Validation at step 34000: Loss = 2.9452, Perplexity = 19.01
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-34000
+Step 34100/44000 | Loss: 2.9269 | LR: 6.19e-05 | 484.3ms/step | 41.3 samples/s
+Step 34200/44000 | Loss: 2.9041 | LR: 6.07e-05 | 437.0ms/step | 45.8 samples/s
+Validation at step 34250: Loss = 2.9449, Perplexity = 19.01
+Step 34300/44000 | Loss: 2.9098 | LR: 5.95e-05 | 455.1ms/step | 43.9 samples/s
+Step 34400/44000 | Loss: 2.9094 | LR: 5.84e-05 | 438.3ms/step | 45.6 samples/s
+Step 34500/44000 | Loss: 2.9051 | LR: 5.72e-05 | 437.4ms/step | 45.7 samples/s
+Validation at step 34500: Loss = 2.9445, Perplexity = 19.00
+Step 34600/44000 | Loss: 2.8926 | LR: 5.60e-05 | 455.0ms/step | 44.0 samples/s
+Step 34700/44000 | Loss: 2.9188 | LR: 5.49e-05 | 437.7ms/step | 45.7 samples/s
+Validation at step 34750: Loss = 2.9444, Perplexity = 19.00
+Step 34800/44000 | Loss: 2.9080 | LR: 5.38e-05 | 454.8ms/step | 44.0 samples/s
+Step 34900/44000 | Loss: 2.9065 | LR: 5.27e-05 | 437.6ms/step | 45.7 samples/s
+Step 35000/44000 | Loss: 2.8951 | LR: 5.15e-05 | 437.4ms/step | 45.7 samples/s
+Validation at step 35000: Loss = 2.9442, Perplexity = 19.00
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-35000
+Step 35100/44000 | Loss: 2.8960 | LR: 5.04e-05 | 487.5ms/step | 41.0 samples/s
+Step 35200/44000 | Loss: 2.9044 | LR: 4.94e-05 | 438.2ms/step | 45.6 samples/s
+Validation at step 35250: Loss = 2.9438, Perplexity = 18.99
+Step 35300/44000 | Loss: 2.8990 | LR: 4.83e-05 | 454.4ms/step | 44.0 samples/s
+Step 35400/44000 | Loss: 2.9109 | LR: 4.72e-05 | 436.4ms/step | 45.8 samples/s
+Step 35500/44000 | Loss: 2.9150 | LR: 4.62e-05 | 436.5ms/step | 45.8 samples/s
+Validation at step 35500: Loss = 2.9436, Perplexity = 18.98
+Step 35600/44000 | Loss: 2.9011 | LR: 4.51e-05 | 454.7ms/step | 44.0 samples/s
+Step 35700/44000 | Loss: 2.9049 | LR: 4.41e-05 | 436.7ms/step | 45.8 samples/s
+Validation at step 35750: Loss = 2.9436, Perplexity = 18.98
+Step 35800/44000 | Loss: 2.8831 | LR: 4.31e-05 | 454.5ms/step | 44.0 samples/s
+Step 35900/44000 | Loss: 2.9003 | LR: 4.20e-05 | 437.5ms/step | 45.7 samples/s
+Step 36000/44000 | Loss: 2.9301 | LR: 4.10e-05 | 436.2ms/step | 45.8 samples/s
+Validation at step 36000: Loss = 2.9434, Perplexity = 18.98
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-36000
+Step 36100/44000 | Loss: 2.8998 | LR: 4.00e-05 | 486.2ms/step | 41.1 samples/s
+Step 36200/44000 | Loss: 2.9058 | LR: 3.91e-05 | 437.4ms/step | 45.7 samples/s
+Validation at step 36250: Loss = 2.9434, Perplexity = 18.98
+Step 36300/44000 | Loss: 2.9010 | LR: 3.81e-05 | 454.0ms/step | 44.1 samples/s
+Step 36400/44000 | Loss: 2.9119 | LR: 3.71e-05 | 437.0ms/step | 45.8 samples/s
+Step 36500/44000 | Loss: 2.9049 | LR: 3.62e-05 | 437.0ms/step | 45.8 samples/s
+Validation at step 36500: Loss = 2.9431, Perplexity = 18.98
+Step 36600/44000 | Loss: 2.8979 | LR: 3.53e-05 | 454.5ms/step | 44.0 samples/s
+Step 36700/44000 | Loss: 2.9048 | LR: 3.43e-05 | 437.6ms/step | 45.7 samples/s
+Validation at step 36750: Loss = 2.9431, Perplexity = 18.97
+Step 36800/44000 | Loss: 2.9167 | LR: 3.34e-05 | 454.3ms/step | 44.0 samples/s
+Step 36900/44000 | Loss: 2.9073 | LR: 3.25e-05 | 437.0ms/step | 45.8 samples/s
+Step 37000/44000 | Loss: 2.8956 | LR: 3.16e-05 | 437.7ms/step | 45.7 samples/s
+Validation at step 37000: Loss = 2.9430, Perplexity = 18.97
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-37000
+Step 37100/44000 | Loss: 2.9025 | LR: 3.07e-05 | 484.1ms/step | 41.3 samples/s
+Step 37200/44000 | Loss: 2.9125 | LR: 2.99e-05 | 436.1ms/step | 45.9 samples/s
+Validation at step 37250: Loss = 2.9429, Perplexity = 18.97
+Step 37300/44000 | Loss: 2.9050 | LR: 2.90e-05 | 453.8ms/step | 44.1 samples/s
+Step 37400/44000 | Loss: 2.8952 | LR: 2.82e-05 | 437.2ms/step | 45.7 samples/s
+Step 37500/44000 | Loss: 2.9260 | LR: 2.74e-05 | 437.3ms/step | 45.7 samples/s
+Validation at step 37500: Loss = 2.9429, Perplexity = 18.97
+Step 37600/44000 | Loss: 2.9103 | LR: 2.65e-05 | 454.4ms/step | 44.0 samples/s
+Step 37700/44000 | Loss: 2.8781 | LR: 2.57e-05 | 436.9ms/step | 45.8 samples/s
+Validation at step 37750: Loss = 2.9427, Perplexity = 18.97
+Step 37800/44000 | Loss: 2.9165 | LR: 2.49e-05 | 455.5ms/step | 43.9 samples/s
+Step 37900/44000 | Loss: 2.9013 | LR: 2.41e-05 | 436.7ms/step | 45.8 samples/s
+Step 38000/44000 | Loss: 2.9022 | LR: 2.34e-05 | 437.2ms/step | 45.7 samples/s
+Validation at step 38000: Loss = 2.9427, Perplexity = 18.97
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-38000
+Step 38100/44000 | Loss: 2.8820 | LR: 2.26e-05 | 485.4ms/step | 41.2 samples/s
+Step 38200/44000 | Loss: 2.8973 | LR: 2.19e-05 | 437.5ms/step | 45.7 samples/s
+Validation at step 38250: Loss = 2.9426, Perplexity = 18.96
+Step 38300/44000 | Loss: 2.9056 | LR: 2.11e-05 | 454.4ms/step | 44.0 samples/s
+Step 38400/44000 | Loss: 2.9050 | LR: 2.04e-05 | 437.6ms/step | 45.7 samples/s
+Step 38500/44000 | Loss: 2.9004 | LR: 1.97e-05 | 438.0ms/step | 45.7 samples/s
+Validation at step 38500: Loss = 2.9426, Perplexity = 18.96
+Step 38600/44000 | Loss: 2.8895 | LR: 1.90e-05 | 455.3ms/step | 43.9 samples/s
+Step 38700/44000 | Loss: 2.8986 | LR: 1.83e-05 | 438.6ms/step | 45.6 samples/s
+Validation at step 38750: Loss = 2.9425, Perplexity = 18.96
+Step 38800/44000 | Loss: 2.9114 | LR: 1.76e-05 | 455.1ms/step | 43.9 samples/s
+Step 38900/44000 | Loss: 2.8867 | LR: 1.70e-05 | 437.4ms/step | 45.7 samples/s
+Step 39000/44000 | Loss: 2.8865 | LR: 1.63e-05 | 437.8ms/step | 45.7 samples/s
+Validation at step 39000: Loss = 2.9425, Perplexity = 18.96
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-39000
+Step 39100/44000 | Loss: 2.8847 | LR: 1.57e-05 | 486.3ms/step | 41.1 samples/s
+Step 39200/44000 | Loss: 2.9052 | LR: 1.50e-05 | 437.7ms/step | 45.7 samples/s
+Validation at step 39250: Loss = 2.9425, Perplexity = 18.96
+Step 39300/44000 | Loss: 2.8853 | LR: 1.44e-05 | 455.5ms/step | 43.9 samples/s
+Step 39400/44000 | Loss: 2.9161 | LR: 1.38e-05 | 438.5ms/step | 45.6 samples/s
+Step 39500/44000 | Loss: 2.9164 | LR: 1.32e-05 | 438.5ms/step | 45.6 samples/s
+Validation at step 39500: Loss = 2.9424, Perplexity = 18.96
+Step 39600/44000 | Loss: 2.9057 | LR: 1.27e-05 | 455.5ms/step | 43.9 samples/s
+Step 39700/44000 | Loss: 2.9129 | LR: 1.21e-05 | 438.2ms/step | 45.6 samples/s
+Validation at step 39750: Loss = 2.9424, Perplexity = 18.96
+Step 39800/44000 | Loss: 2.9096 | LR: 1.15e-05 | 455.0ms/step | 44.0 samples/s
+Step 39900/44000 | Loss: 2.9188 | LR: 1.10e-05 | 437.8ms/step | 45.7 samples/s
+Step 40000/44000 | Loss: 2.9064 | LR: 1.05e-05 | 438.2ms/step | 45.6 samples/s
+Validation at step 40000: Loss = 2.9424, Perplexity = 18.96
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-40000
+Step 40100/44000 | Loss: 2.8944 | LR: 9.96e-06 | 486.4ms/step | 41.1 samples/s
+Step 40200/44000 | Loss: 2.8920 | LR: 9.46e-06 | 438.1ms/step | 45.6 samples/s
+Validation at step 40250: Loss = 2.9424, Perplexity = 18.96
+Step 40300/44000 | Loss: 2.9163 | LR: 8.97e-06 | 455.9ms/step | 43.9 samples/s
+Step 40400/44000 | Loss: 2.8962 | LR: 8.50e-06 | 438.4ms/step | 45.6 samples/s
+Step 40500/44000 | Loss: 2.8868 | LR: 8.04e-06 | 438.1ms/step | 45.7 samples/s
+Validation at step 40500: Loss = 2.9424, Perplexity = 18.96
+Step 40600/44000 | Loss: 2.8995 | LR: 7.59e-06 | 455.3ms/step | 43.9 samples/s
+Step 40700/44000 | Loss: 2.9103 | LR: 7.15e-06 | 438.3ms/step | 45.6 samples/s
+Validation at step 40750: Loss = 2.9423, Perplexity = 18.96
+Step 40800/44000 | Loss: 2.9073 | LR: 6.72e-06 | 455.4ms/step | 43.9 samples/s
+Step 40900/44000 | Loss: 2.9062 | LR: 6.31e-06 | 438.6ms/step | 45.6 samples/s
+Step 41000/44000 | Loss: 2.9028 | LR: 5.91e-06 | 438.1ms/step | 45.6 samples/s
+Validation at step 41000: Loss = 2.9424, Perplexity = 18.96
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-41000
+Step 41100/44000 | Loss: 2.9286 | LR: 5.53e-06 | 485.9ms/step | 41.2 samples/s
+Step 41200/44000 | Loss: 2.8938 | LR: 5.15e-06 | 437.5ms/step | 45.7 samples/s
+Validation at step 41250: Loss = 2.9424, Perplexity = 18.96
+Step 41300/44000 | Loss: 2.9098 | LR: 4.79e-06 | 454.7ms/step | 44.0 samples/s
+Step 41400/44000 | Loss: 2.9049 | LR: 4.45e-06 | 438.0ms/step | 45.7 samples/s
+Step 41500/44000 | Loss: 2.8782 | LR: 4.11e-06 | 438.2ms/step | 45.6 samples/s
+Validation at step 41500: Loss = 2.9424, Perplexity = 18.96
+Step 41600/44000 | Loss: 2.9010 | LR: 3.79e-06 | 455.3ms/step | 43.9 samples/s
+Step 41700/44000 | Loss: 2.9046 | LR: 3.48e-06 | 438.1ms/step | 45.6 samples/s
+Validation at step 41750: Loss = 2.9423, Perplexity = 18.96
+Step 41800/44000 | Loss: 2.9250 | LR: 3.19e-06 | 455.3ms/step | 43.9 samples/s
+Step 41900/44000 | Loss: 2.8708 | LR: 2.90e-06 | 438.1ms/step | 45.6 samples/s
+Step 42000/44000 | Loss: 2.9088 | LR: 2.63e-06 | 438.2ms/step | 45.6 samples/s
+Validation at step 42000: Loss = 2.9423, Perplexity = 18.96
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-42000
+Step 42100/44000 | Loss: 2.9261 | LR: 2.38e-06 | 488.0ms/step | 41.0 samples/s
+Step 42200/44000 | Loss: 2.8979 | LR: 2.13e-06 | 437.5ms/step | 45.7 samples/s
+Validation at step 42250: Loss = 2.9424, Perplexity = 18.96
+Step 42300/44000 | Loss: 2.8978 | LR: 1.90e-06 | 454.7ms/step | 44.0 samples/s
+Step 42400/44000 | Loss: 2.9129 | LR: 1.69e-06 | 438.1ms/step | 45.7 samples/s
+Step 42500/44000 | Loss: 2.9031 | LR: 1.48e-06 | 438.1ms/step | 45.7 samples/s
+Validation at step 42500: Loss = 2.9423, Perplexity = 18.96
+Step 42600/44000 | Loss: 2.9200 | LR: 1.29e-06 | 455.5ms/step | 43.9 samples/s
+Step 42700/44000 | Loss: 2.9030 | LR: 1.11e-06 | 438.3ms/step | 45.6 samples/s
+Validation at step 42750: Loss = 2.9424, Perplexity = 18.96
+Step 42800/44000 | Loss: 2.9052 | LR: 9.49e-07 | 455.5ms/step | 43.9 samples/s
+Step 42900/44000 | Loss: 2.8999 | LR: 7.98e-07 | 437.4ms/step | 45.7 samples/s
+Step 43000/44000 | Loss: 2.8856 | LR: 6.59e-07 | 437.5ms/step | 45.7 samples/s
+Validation at step 43000: Loss = 2.9423, Perplexity = 18.96
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-43000
+Step 43100/44000 | Loss: 2.9006 | LR: 5.34e-07 | 484.5ms/step | 41.3 samples/s
+Step 43200/44000 | Loss: 2.9036 | LR: 4.22e-07 | 436.5ms/step | 45.8 samples/s
+Validation at step 43250: Loss = 2.9423, Perplexity = 18.96
+Step 43300/44000 | Loss: 2.8914 | LR: 3.23e-07 | 454.3ms/step | 44.0 samples/s
+Step 43400/44000 | Loss: 2.9069 | LR: 2.37e-07 | 437.4ms/step | 45.7 samples/s
+Step 43500/44000 | Loss: 2.8928 | LR: 1.65e-07 | 437.4ms/step | 45.7 samples/s
+Validation at step 43500: Loss = 2.9424, Perplexity = 18.96
+Step 43600/44000 | Loss: 2.9161 | LR: 1.06e-07 | 454.1ms/step | 44.0 samples/s
+Step 43700/44000 | Loss: 2.9227 | LR: 5.94e-08 | 436.9ms/step | 45.8 samples/s
+Validation at step 43750: Loss = 2.9423, Perplexity = 18.96
+Step 43800/44000 | Loss: 2.9093 | LR: 2.64e-08 | 454.3ms/step | 44.0 samples/s
+Step 43900/44000 | Loss: 2.9149 | LR: 6.60e-09 | 437.4ms/step | 45.7 samples/s
+Step 44000/44000 | Loss: 2.9269 | LR: 0.00e+00 | 437.0ms/step | 45.8 samples/s
+Validation at step 44000: Loss = 2.9424, Perplexity = 18.96
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-44000
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-44000
+wandb: 
+wandb: Run history:
+wandb: eval/dense_interval_active ▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁
+wandb:    eval/dense_interval_end ▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁
+wandb:  eval/dense_interval_start ▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁
+wandb:                  eval/loss ████▇▆▆▅▅▅▅▄▄▄▃▃▃▃▂▂▂▂▂▂▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁
+wandb:            eval/perplexity ██▇▇▇▆▆▆▅▅▅▄▃▃▃▂▂▂▂▂▂▂▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁
+wandb:  eval/uses_dense_attention ▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁
+wandb: eval/uses_sparse_attention ▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁
+wandb:          speed/ms_per_step ▁▁▁▁▄█▄▄▁▁▄▁▄▁��▃▄▁▄▁▄▁▁▁▄▁▄▁▄▁▄▁▄▄█▄▄▁▁▁
+wandb:      speed/samples_per_sec ▆█▂▅▅█▅█▁████▅▅▁▅█▅███▅██▁▅██▁████████▅█
+wandb:        speed/steps_per_sec ██▇███▁▇▁▅███▅▁██▇█▇▁▇▅▇▁▅▅█████▁█▇▅▇▁▅█
+wandb:                         +4 ...
+wandb: 
+wandb: Run summary:
+wandb: eval/dense_interval_active 0
+wandb:    eval/dense_interval_end -1
+wandb:  eval/dense_interval_start -1
+wandb:                  eval/loss 2.94235
+wandb:            eval/perplexity 18.96044
+wandb:  eval/uses_dense_attention 0
+wandb: eval/uses_sparse_attention 1
+wandb:          speed/ms_per_step 436.9685
+wandb:      speed/samples_per_sec 45.76989
+wandb:        speed/steps_per_sec 2.28849
+wandb:                         +4 ...
+wandb: 
+wandb: You can sync this run to the cloud by running:
+wandb: wandb sync /mloscratch/homes/navasard/sparse/SparseAtt/wandb/offline-run-20260921_195241-nihe6o4e
+wandb: Find logs at: ./wandb/offline-run-20260921_195241-nihe6o4e/logs
+Training complete!

epfl_final/logs/ALIBI_20260921/ddp4_seq_20260921_163113/orchestrator.log

@@ -0,0 +1 @@
+bash: line 20: seed: unbound variable

epfl_final/logs/ALIBI_20260921/ddp4_seq_20260921_163148/orchestrator.log

@@ -0,0 +1 @@
+START seed=42 2026-09-21T16:31:48+00:00

epfl_final/logs/ALIBI_20260921/ddp4_seq_20260921_163148/seed42.log

@@ -0,0 +1,429 @@
+W0921 16:31:51.249000 193 site-packages/torch/distributed/run.py:851] 
+W0921 16:31:51.249000 193 site-packages/torch/distributed/run.py:851] *****************************************
+W0921 16:31:51.249000 193 site-packages/torch/distributed/run.py:851] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. 
+W0921 16:31:51.249000 193 site-packages/torch/distributed/run.py:851] *****************************************
+Auto-resuming from latest checkpoint: /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-2000
+Resuming: Using existing output directory /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951
+Running config:
+{
+  "model_path": "/mloscratch/homes/navasard/qwen3_06b_base",
+  "model_dtype": "bfloat16",
+  "init_from_scratch": true,
+  "train_path": "/mloscratch/homes/navasard/data/fineweb_large/tokenized",
+  "val_path": null,
+  "max_seq_length": 8192,
+  "num_workers": 8,
+  "output_dir": "/mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951",
+  "resume_from_checkpoint": "/mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-2000",
+  "force_new_run": false,
+  "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,
+  "resume_micro_step_override": 10000,
+  "strategy": "ddp",
+  "mixed_precision": "bf16",
+  "wandb_enabled": true,
+  "wandb_project": "stochastic_sparsity",
+  "wandb_name": "365m alibi 8k s42",
+  "wandb_tags": [
+    "qwen3",
+    "365m",
+    "8k",
+    "scaleladder",
+    "alibi",
+    "seed42"
+  ],
+  "hidden_size": 768,
+  "intermediate_size": 3072,
+  "num_attention_heads": 6,
+  "num_key_value_heads": 2,
+  "num_hidden_layers": 28,
+  "head_dim": null,
+  "attention_dropout": null,
+  "rope_train_position_scale_set": null,
+  "positional": "alibi",
+  "attention_type": "dense",
+  "sparse_block_size": 128,
+  "sparse_num_random_blocks": 31,
+  "sparse_mask_pool_size": 1,
+  "sparse_val_dense_compare": false,
+  "sparse_layer_indices": null,
+  "sparse_mask_cache_dir": null,
+  "sparse_mask_seed": 0,
+  "sparse_mixed_nrb_set": null,
+  "dense_mode": false,
+  "dense_step_intervals": null
+}
+Loading tokenizer from /mloscratch/homes/navasard/qwen3_06b_base...
+Resuming from checkpoint: /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-2000
+Checkpoint activations: False
+Checkpoint activations: False
+Checkpoint activations: False
+Checkpoint activations: False
+
+
+================================================================================
+MODEL ARCHITECTURE
+================================================================================
+Model config:
+  vocab_size: 151936
+  hidden_size: 768
+  num_hidden_layers: 28
+  num_attention_heads: 6
+  num_key_value_heads: 2
+  head_dim: 128
+  intermediate_size: 3072
+  max_position_embeddings: 32768
+  rope_theta: 1000000
+  rope_train_position_scale_set: None
+  attention_type: dense
+
+Parameter count:
+  Total parameters: 358,958,848
+  Trainable parameters: 358,958,848
+  Non-trainable parameters: 0
+  Model size (BF16): 0.72 GB
+
+--------------------------------------------------------------------------------
+LAYER-BY-LAYER SUMMARY
+--------------------------------------------------------------------------------
+/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/nn/attention/flex_attention.py:1704: UserWarning: flex_attention called without torch.compile() - this will use an unfused implementation that materializes the full scores matrix instead of generating a fused kernel.
+
+SOLUTION: Use torch.compile(flex_attention)(...)
+
+If you want to debug your score_mod/mask_mod, you can set:
+torch.nn.attention.flex_attention._FLEX_ATTENTION_DISABLE_COMPILE_DEBUG = True
+
+This will allow you to use print statements or breakpoints. Note: This doesn't work with the backwards pass and may produce incorrect results.
+  _warn_once(
+=================================================================================================================================================
+Layer (type:depth-idx)                        Input Shape               Output Shape              Param #                   Trainable
+=================================================================================================================================================
+Qwen3ForCausalLM                              [1, 8192]                 --                        --                        True
+���─Qwen3Model: 1-1                             --                        [1, 8192, 768]            --                        True
+│    └─Embedding: 2-1                         [1, 8192]                 [1, 8192, 768]            116,686,848               True
+│    └─ModuleList: 2-2                        --                        --                        --                        True
+│    │    └─Qwen3DecoderLayer: 3-1            [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-2            [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-3            [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-4            [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-5            [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-6            [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-7            [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-8            [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-9            [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-10           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-11           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-12           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-13           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-14           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-15           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-16           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-17           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-18           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-19           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-20           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-21           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-22           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-23           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-24           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-25           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-26           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-27           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-28           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    └─Qwen3RMSNorm: 2-3                      [1, 8192, 768]            [1, 8192, 768]            768                       True
+├─Linear: 1-2                                 [1, 8192, 768]            [1, 8192, 151936]         116,686,848               True
+=================================================================================================================================================
+Total params: 475,645,696
+Trainable params: 475,645,696
+Non-trainable params: 0
+Total mult-adds (Units.MEGABYTES): 475.65
+=================================================================================================================================================
+Input size (MB): 0.07
+Forward/backward pass size (MB): 15598.62
+Params size (MB): 951.29
+Estimated Total Size (MB): 16549.97
+=================================================================================================================================================
+  NOTE: torchinfo double-counts tied embed_tokens/lm_head weight (116,686,848 params). True unique parameter count: 358,958,848 (359.0M)
+================================================================================
+
+Loading dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Using pre-tokenized dataset (no runtime tokenization)
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Loaded 12748481 pre-tokenized examples (stored_len=2048, pack_factor=4, packed_examples=3187120)
+Loaded 12748481 pre-tokenized examples (stored_len=2048, pack_factor=4, packed_examples=3187120)
+Loaded 12748481 pre-tokenized examples (stored_len=2048, pack_factor=4, packed_examples=3187120)
+Loaded 12748481 pre-tokenized examples (stored_len=2048, pack_factor=4, packed_examples=3187120)
+Dataset split: 3027764 train, 159356 validation
+Training examples (packed, train split):  3027764
+Validation examples (packed, val split):  159356
+Per-device batch size:                    1
+World size (GPUs):                        4
+Gradient accumulation steps:              5
+Effective batch size (seqs/step):         20
+Sequence length:                          8192
+Tokens per optimizer step:                163,840
+Micro-batches per epoch (per rank):       756941
+Optimizer steps per epoch:                151388
+Total optimizer steps:                    44000  (capped by max_steps; 0.29 epochs)
+Loading checkpoint from /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-2000...
+Preserving checkpoint LR state for exact resume: current=[0.0004989701850946613], schedule_base=[0.0005]
+Resumed from checkpoint at step 2000 (epoch 0, micro_step 20000)
+Overriding resume micro-step 20000 -> 10000 for DDP migration
+Resuming training from step 2000
+Resuming wandb run nihe6o4e
+wandb: WARNING `resume` will be ignored since W&B syncing is set to `offline`. Starting a new run with run id nihe6o4e.
+wandb: Tracking run with wandb version 0.26.1
+wandb: W&B syncing is set to `offline` in this directory. Run `wandb online` or set WANDB_MODE=online to enable cloud syncing.
+wandb: Run data is saved locally in /mloscratch/homes/navasard/sparse/SparseAtt/wandb/offline-run-20260921_163219-nihe6o4e
+Starting training...
+Validation at step 2000: Loss = 3.6852, Perplexity = 39.85
+Switching attention mode for step 2001: sparse_stochastic (configured sparse attention)
+Step 2100/44000 | Loss: 3.6414 | LR: 4.99e-04 | 1840.7ms/step | 10.9 samples/s
+Step 2200/44000 | Loss: 3.6338 | LR: 4.99e-04 | 438.2ms/step | 45.6 samples/s
+Validation at step 2250: Loss = 3.6290, Perplexity = 37.67
+Step 2300/44000 | Loss: 3.5974 | LR: 4.98e-04 | 456.5ms/step | 43.8 samples/s
+Step 2400/44000 | Loss: 3.5602 | LR: 4.98e-04 | 440.6ms/step | 45.4 samples/s
+Step 2500/44000 | Loss: 3.5566 | LR: 4.98e-04 | 440.2ms/step | 45.4 samples/s
+Validation at step 2500: Loss = 3.5785, Perplexity = 35.82
+Step 2600/44000 | Loss: 3.5170 | LR: 4.98e-04 | 457.0ms/step | 43.8 samples/s
+Step 2700/44000 | Loss: 3.4957 | LR: 4.97e-04 | 438.8ms/step | 45.6 samples/s
+Validation at step 2750: Loss = 3.5411, Perplexity = 34.50
+Step 2800/44000 | Loss: 3.5299 | LR: 4.97e-04 | 455.1ms/step | 43.9 samples/s
+Step 2900/44000 | Loss: 3.4934 | LR: 4.97e-04 | 438.2ms/step | 45.6 samples/s
+Step 3000/44000 | Loss: 3.5026 | LR: 4.97e-04 | 438.9ms/step | 45.6 samples/s
+Validation at step 3000: Loss = 3.5080, Perplexity = 33.38
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-3000
+Step 3100/44000 | Loss: 3.4725 | LR: 4.96e-04 | 487.0ms/step | 41.1 samples/s
+Step 3200/44000 | Loss: 3.4509 | LR: 4.96e-04 | 439.2ms/step | 45.5 samples/s
+Validation at step 3250: Loss = 3.4794, Perplexity = 32.44
+Step 3300/44000 | Loss: 3.4351 | LR: 4.96e-04 | 455.9ms/step | 43.9 samples/s
+Step 3400/44000 | Loss: 3.4322 | LR: 4.95e-04 | 440.2ms/step | 45.4 samples/s
+Step 3500/44000 | Loss: 3.4099 | LR: 4.95e-04 | 438.7ms/step | 45.6 samples/s
+Validation at step 3500: Loss = 3.4544, Perplexity = 31.64
+Step 3600/44000 | Loss: 3.4246 | LR: 4.95e-04 | 455.9ms/step | 43.9 samples/s
+Step 3700/44000 | Loss: 3.4090 | LR: 4.94e-04 | 438.6ms/step | 45.6 samples/s
+Validation at step 3750: Loss = 3.4323, Perplexity = 30.95
+Step 3800/44000 | Loss: 3.3831 | LR: 4.94e-04 | 455.5ms/step | 43.9 samples/s
+Step 3900/44000 | Loss: 3.3802 | LR: 4.93e-04 | 438.9ms/step | 45.6 samples/s
+Step 4000/44000 | Loss: 3.3732 | LR: 4.93e-04 | 438.9ms/step | 45.6 samples/s
+Validation at step 4000: Loss = 3.4119, Perplexity = 30.32
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-4000
+Step 4100/44000 | Loss: 3.3711 | LR: 4.93e-04 | 484.4ms/step | 41.3 samples/s
+Step 4200/44000 | Loss: 3.3590 | LR: 4.92e-04 | 437.0ms/step | 45.8 samples/s
+Validation at step 4250: Loss = 3.3891, Perplexity = 29.64
+Step 4300/44000 | Loss: 3.3634 | LR: 4.92e-04 | 454.9ms/step | 44.0 samples/s
+Step 4400/44000 | Loss: 3.3420 | LR: 4.91e-04 | 437.8ms/step | 45.7 samples/s
+Step 4500/44000 | Loss: 3.3241 | LR: 4.91e-04 | 438.3ms/step | 45.6 samples/s
+Validation at step 4500: Loss = 3.3773, Perplexity = 29.29
+Step 4600/44000 | Loss: 3.3226 | LR: 4.90e-04 | 455.9ms/step | 43.9 samples/s
+Step 4700/44000 | Loss: 3.3287 | LR: 4.90e-04 | 438.9ms/step | 45.6 samples/s
+Validation at step 4750: Loss = 3.3632, Perplexity = 28.88
+Step 4800/44000 | Loss: 3.3254 | LR: 4.89e-04 | 454.5ms/step | 44.0 samples/s
+Step 4900/44000 | Loss: 3.3012 | LR: 4.89e-04 | 438.0ms/step | 45.7 samples/s
+Step 5000/44000 | Loss: 3.3060 | LR: 4.88e-04 | 437.8ms/step | 45.7 samples/s
+Validation at step 5000: Loss = 3.3488, Perplexity = 28.47
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-5000
+Step 5100/44000 | Loss: 3.2958 | LR: 4.88e-04 | 488.8ms/step | 40.9 samples/s
+Step 5200/44000 | Loss: 3.2990 | LR: 4.87e-04 | 440.8ms/step | 45.4 samples/s
+Validation at step 5250: Loss = 3.3333, Perplexity = 28.03
+Step 5300/44000 | Loss: 3.3205 | LR: 4.86e-04 | 456.0ms/step | 43.9 samples/s
+Step 5400/44000 | Loss: 3.2890 | LR: 4.86e-04 | 437.9ms/step | 45.7 samples/s
+Step 5500/44000 | Loss: 3.2792 | LR: 4.85e-04 | 437.9ms/step | 45.7 samples/s
+Validation at step 5500: Loss = 3.3190, Perplexity = 27.63
+Step 5600/44000 | Loss: 3.2905 | LR: 4.85e-04 | 454.7ms/step | 44.0 samples/s
+Step 5700/44000 | Loss: 3.2627 | LR: 4.84e-04 | 438.1ms/step | 45.7 samples/s
+Validation at step 5750: Loss = 3.3040, Perplexity = 27.22
+Step 5800/44000 | Loss: 3.2672 | LR: 4.83e-04 | 454.0ms/step | 44.1 samples/s
+Step 5900/44000 | Loss: 3.2572 | LR: 4.83e-04 | 436.5ms/step | 45.8 samples/s
+Step 6000/44000 | Loss: 3.2646 | LR: 4.82e-04 | 436.3ms/step | 45.8 samples/s
+Validation at step 6000: Loss = 3.2875, Perplexity = 26.78
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-6000
+Step 6100/44000 | Loss: 3.2554 | LR: 4.81e-04 | 485.1ms/step | 41.2 samples/s
+Step 6200/44000 | Loss: 3.2378 | LR: 4.81e-04 | 437.7ms/step | 45.7 samples/s
+Validation at step 6250: Loss = 3.2817, Perplexity = 26.62
+Step 6300/44000 | Loss: 3.2241 | LR: 4.80e-04 | 455.3ms/step | 43.9 samples/s
+Step 6400/44000 | Loss: 3.2477 | LR: 4.79e-04 | 437.1ms/step | 45.8 samples/s
+Step 6500/44000 | Loss: 3.2323 | LR: 4.79e-04 | 437.9ms/step | 45.7 samples/s
+Validation at step 6500: Loss = 3.2703, Perplexity = 26.32
+Step 6600/44000 | Loss: 3.2180 | LR: 4.78e-04 | 454.8ms/step | 44.0 samples/s
+Step 6700/44000 | Loss: 3.2344 | LR: 4.77e-04 | 437.9ms/step | 45.7 samples/s
+Validation at step 6750: Loss = 3.2625, Perplexity = 26.12
+Step 6800/44000 | Loss: 3.2201 | LR: 4.76e-04 | 454.7ms/step | 44.0 samples/s
+Step 6900/44000 | Loss: 3.2113 | LR: 4.75e-04 | 438.5ms/step | 45.6 samples/s
+Step 7000/44000 | Loss: 3.2210 | LR: 4.75e-04 | 438.5ms/step | 45.6 samples/s
+Validation at step 7000: Loss = 3.2546, Perplexity = 25.91
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-7000
+Step 7100/44000 | Loss: 3.1977 | LR: 4.74e-04 | 487.4ms/step | 41.0 samples/s
+Step 7200/44000 | Loss: 3.2093 | LR: 4.73e-04 | 438.1ms/step | 45.7 samples/s
+Validation at step 7250: Loss = 3.2443, Perplexity = 25.64
+Step 7300/44000 | Loss: 3.1957 | LR: 4.72e-04 | 455.1ms/step | 43.9 samples/s
+Step 7400/44000 | Loss: 3.2087 | LR: 4.71e-04 | 438.4ms/step | 45.6 samples/s
+Step 7500/44000 | Loss: 3.2125 | LR: 4.71e-04 | 438.8ms/step | 45.6 samples/s
+Validation at step 7500: Loss = 3.2374, Perplexity = 25.47
+Step 7600/44000 | Loss: 3.1824 | LR: 4.70e-04 | 455.2ms/step | 43.9 samples/s
+Step 7700/44000 | Loss: 3.2040 | LR: 4.69e-04 | 437.7ms/step | 45.7 samples/s
+Validation at step 7750: Loss = 3.2296, Perplexity = 25.27
+Step 7800/44000 | Loss: 3.1780 | LR: 4.68e-04 | 455.1ms/step | 44.0 samples/s
+Step 7900/44000 | Loss: 3.1836 | LR: 4.67e-04 | 437.5ms/step | 45.7 samples/s
+Step 8000/44000 | Loss: 3.1987 | LR: 4.66e-04 | 437.3ms/step | 45.7 samples/s
+Validation at step 8000: Loss = 3.2216, Perplexity = 25.07
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-8000
+Step 8100/44000 | Loss: 3.1784 | LR: 4.65e-04 | 484.7ms/step | 41.3 samples/s
+Step 8200/44000 | Loss: 3.1806 | LR: 4.64e-04 | 437.3ms/step | 45.7 samples/s
+Validation at step 8250: Loss = 3.2140, Perplexity = 24.88
+Step 8300/44000 | Loss: 3.1826 | LR: 4.63e-04 | 455.1ms/step | 43.9 samples/s
+Step 8400/44000 | Loss: 3.1689 | LR: 4.62e-04 | 438.4ms/step | 45.6 samples/s
+Step 8500/44000 | Loss: 3.1824 | LR: 4.61e-04 | 437.8ms/step | 45.7 samples/s
+Validation at step 8500: Loss = 3.2081, Perplexity = 24.73
+Step 8600/44000 | Loss: 3.1641 | LR: 4.60e-04 | 455.1ms/step | 44.0 samples/s
+Step 8700/44000 | Loss: 3.1513 | LR: 4.59e-04 | 438.1ms/step | 45.7 samples/s
+Validation at step 8750: Loss = 3.2015, Perplexity = 24.57
+Step 8800/44000 | Loss: 3.1708 | LR: 4.58e-04 | 455.6ms/step | 43.9 samples/s
+Step 8900/44000 | Loss: 3.1471 | LR: 4.57e-04 | 437.3ms/step | 45.7 samples/s
+Step 9000/44000 | Loss: 3.1585 | LR: 4.56e-04 | 437.4ms/step | 45.7 samples/s
+Validation at step 9000: Loss = 3.1946, Perplexity = 24.40
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-9000
+Step 9100/44000 | Loss: 3.1521 | LR: 4.55e-04 | 483.9ms/step | 41.3 samples/s
+Step 9200/44000 | Loss: 3.1565 | LR: 4.54e-04 | 437.4ms/step | 45.7 samples/s
+Validation at step 9250: Loss = 3.1909, Perplexity = 24.31
+Step 9300/44000 | Loss: 3.1727 | LR: 4.53e-04 | 453.7ms/step | 44.1 samples/s
+Step 9400/44000 | Loss: 3.1604 | LR: 4.52e-04 | 438.0ms/step | 45.7 samples/s
+Step 9500/44000 | Loss: 3.1477 | LR: 4.51e-04 | 437.3ms/step | 45.7 samples/s
+Validation at step 9500: Loss = 3.1854, Perplexity = 24.18
+Step 9600/44000 | Loss: 3.1458 | LR: 4.50e-04 | 453.8ms/step | 44.1 samples/s
+Step 9700/44000 | Loss: 3.1306 | LR: 4.49e-04 | 437.6ms/step | 45.7 samples/s
+Validation at step 9750: Loss = 3.1773, Perplexity = 23.98
+Step 9800/44000 | Loss: 3.1282 | LR: 4.48e-04 | 454.4ms/step | 44.0 samples/s
+Step 9900/44000 | Loss: 3.1320 | LR: 4.47e-04 | 436.8ms/step | 45.8 samples/s
+Step 10000/44000 | Loss: 3.1347 | LR: 4.46e-04 | 437.1ms/step | 45.8 samples/s
+Validation at step 10000: Loss = 3.1726, Perplexity = 23.87
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-10000
+Step 10100/44000 | Loss: 3.1067 | LR: 4.45e-04 | 485.1ms/step | 41.2 samples/s
+Step 10200/44000 | Loss: 3.1271 | LR: 4.43e-04 | 437.8ms/step | 45.7 samples/s
+Validation at step 10250: Loss = 3.1653, Perplexity = 23.70
+Step 10300/44000 | Loss: 3.1240 | LR: 4.42e-04 | 454.4ms/step | 44.0 samples/s
+Step 10400/44000 | Loss: 3.1232 | LR: 4.41e-04 | 438.6ms/step | 45.6 samples/s
+Step 10500/44000 | Loss: 3.1159 | LR: 4.40e-04 | 438.0ms/step | 45.7 samples/s
+Validation at step 10500: Loss = 3.1629, Perplexity = 23.64
+Step 10600/44000 | Loss: 3.1165 | LR: 4.39e-04 | 456.6ms/step | 43.8 samples/s
+Step 10700/44000 | Loss: 3.1098 | LR: 4.37e-04 | 438.2ms/step | 45.6 samples/s
+Validation at step 10750: Loss = 3.1567, Perplexity = 23.49
+Step 10800/44000 | Loss: 3.1029 | LR: 4.36e-04 | 454.9ms/step | 44.0 samples/s
+Step 10900/44000 | Loss: 3.1110 | LR: 4.35e-04 | 437.7ms/step | 45.7 samples/s
+Step 11000/44000 | Loss: 3.1161 | LR: 4.34e-04 | 438.4ms/step | 45.6 samples/s
+Validation at step 11000: Loss = 3.1528, Perplexity = 23.40
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-11000
+Step 11100/44000 | Loss: 3.1044 | LR: 4.33e-04 | 484.6ms/step | 41.3 samples/s
+Step 11200/44000 | Loss: 3.1291 | LR: 4.31e-04 | 437.1ms/step | 45.8 samples/s
+Validation at step 11250: Loss = 3.1471, Perplexity = 23.27
+Step 11300/44000 | Loss: 3.1105 | LR: 4.30e-04 | 454.8ms/step | 44.0 samples/s
+Step 11400/44000 | Loss: 3.1041 | LR: 4.29e-04 | 437.2ms/step | 45.7 samples/s
+Step 11500/44000 | Loss: 3.0933 | LR: 4.28e-04 | 436.9ms/step | 45.8 samples/s
+Validation at step 11500: Loss = 3.1421, Perplexity = 23.15
+Step 11600/44000 | Loss: 3.0865 | LR: 4.26e-04 | 454.4ms/step | 44.0 samples/s
+Step 11700/44000 | Loss: 3.1113 | LR: 4.25e-04 | 437.6ms/step | 45.7 samples/s
+Validation at step 11750: Loss = 3.1374, Perplexity = 23.04
+Step 11800/44000 | Loss: 3.0923 | LR: 4.24e-04 | 455.4ms/step | 43.9 samples/s
+Step 11900/44000 | Loss: 3.0900 | LR: 4.22e-04 | 437.7ms/step | 45.7 samples/s
+Step 12000/44000 | Loss: 3.0963 | LR: 4.21e-04 | 437.8ms/step | 45.7 samples/s
+Validation at step 12000: Loss = 3.1351, Perplexity = 22.99
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-12000
+Step 12100/44000 | Loss: 3.0952 | LR: 4.20e-04 | 484.0ms/step | 41.3 samples/s
+Step 12200/44000 | Loss: 3.0868 | LR: 4.18e-04 | 437.6ms/step | 45.7 samples/s
+Validation at step 12250: Loss = 3.1284, Perplexity = 22.84
+Step 12300/44000 | Loss: 3.0975 | LR: 4.17e-04 | 455.6ms/step | 43.9 samples/s
+Step 12400/44000 | Loss: 3.0845 | LR: 4.16e-04 | 438.6ms/step | 45.6 samples/s
+Step 12500/44000 | Loss: 3.0788 | LR: 4.14e-04 | 437.6ms/step | 45.7 samples/s
+Validation at step 12500: Loss = 3.1241, Perplexity = 22.74
+Step 12600/44000 | Loss: 3.0698 | LR: 4.13e-04 | 455.2ms/step | 43.9 samples/s
+Step 12700/44000 | Loss: 3.0854 | LR: 4.12e-04 | 438.2ms/step | 45.6 samples/s
+Validation at step 12750: Loss = 3.1169, Perplexity = 22.58
+Step 12800/44000 | Loss: 3.0874 | LR: 4.10e-04 | 455.2ms/step | 43.9 samples/s
+Step 12900/44000 | Loss: 3.0739 | LR: 4.09e-04 | 437.2ms/step | 45.7 samples/s
+Step 13000/44000 | Loss: 3.0768 | LR: 4.07e-04 | 437.5ms/step | 45.7 samples/s
+Validation at step 13000: Loss = 3.1155, Perplexity = 22.55
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-13000
+Step 13100/44000 | Loss: 3.0873 | LR: 4.06e-04 | 485.6ms/step | 41.2 samples/s
+Step 13200/44000 | Loss: 3.0853 | LR: 4.05e-04 | 437.3ms/step | 45.7 samples/s
+Validation at step 13250: Loss = 3.1105, Perplexity = 22.43
+Step 13300/44000 | Loss: 3.0821 | LR: 4.03e-04 | 453.6ms/step | 44.1 samples/s
+Step 13400/44000 | Loss: 3.0847 | LR: 4.02e-04 | 436.8ms/step | 45.8 samples/s
+Step 13500/44000 | Loss: 3.0735 | LR: 4.00e-04 | 437.2ms/step | 45.7 samples/s
+Validation at step 13500: Loss = 3.1069, Perplexity = 22.35
+Step 13600/44000 | Loss: 3.0673 | LR: 3.99e-04 | 455.0ms/step | 44.0 samples/s
+Step 13700/44000 | Loss: 3.0836 | LR: 3.97e-04 | 437.6ms/step | 45.7 samples/s
+Validation at step 13750: Loss = 3.1012, Perplexity = 22.23
+Step 13800/44000 | Loss: 3.0566 | LR: 3.96e-04 | 455.1ms/step | 43.9 samples/s
+Step 13900/44000 | Loss: 3.0440 | LR: 3.94e-04 | 437.0ms/step | 45.8 samples/s
+Step 14000/44000 | Loss: 3.0619 | LR: 3.93e-04 | 436.9ms/step | 45.8 samples/s
+Validation at step 14000: Loss = 3.1004, Perplexity = 22.21
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-14000
+Step 14100/44000 | Loss: 3.0430 | LR: 3.91e-04 | 484.3ms/step | 41.3 samples/s
+Step 14200/44000 | Loss: 3.0631 | LR: 3.90e-04 | 440.1ms/step | 45.4 samples/s
+Validation at step 14250: Loss = 3.0984, Perplexity = 22.16
+Step 14300/44000 | Loss: 3.0631 | LR: 3.88e-04 | 455.4ms/step | 43.9 samples/s
+Step 14400/44000 | Loss: 3.0579 | LR: 3.87e-04 | 437.0ms/step | 45.8 samples/s
+Step 14500/44000 | Loss: 3.0545 | LR: 3.85e-04 | 436.7ms/step | 45.8 samples/s
+Validation at step 14500: Loss = 3.0922, Perplexity = 22.02
+Step 14600/44000 | Loss: 3.0545 | LR: 3.84e-04 | 453.9ms/step | 44.1 samples/s
+Step 14700/44000 | Loss: 3.0573 | LR: 3.82e-04 | 437.4ms/step | 45.7 samples/s
+Validation at step 14750: Loss = 3.0897, Perplexity = 21.97
+Step 14800/44000 | Loss: 3.0694 | LR: 3.81e-04 | 454.2ms/step | 44.0 samples/s
+Step 14900/44000 | Loss: 3.0670 | LR: 3.79e-04 | 436.6ms/step | 45.8 samples/s
+Step 15000/44000 | Loss: 3.0720 | LR: 3.78e-04 | 437.7ms/step | 45.7 samples/s
+Validation at step 15000: Loss = 3.0858, Perplexity = 21.89
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-15000
+Step 15100/44000 | Loss: 3.0604 | LR: 3.76e-04 | 485.0ms/step | 41.2 samples/s
+Step 15200/44000 | Loss: 3.0517 | LR: 3.74e-04 | 437.9ms/step | 45.7 samples/s
+Validation at step 15250: Loss = 3.0830, Perplexity = 21.82
+Step 15300/44000 | Loss: 3.0399 | LR: 3.73e-04 | 453.9ms/step | 44.1 samples/s
+Step 15400/44000 | Loss: 3.0466 | LR: 3.71e-04 | 436.6ms/step | 45.8 samples/s
+Step 15500/44000 | Loss: 3.0513 | LR: 3.70e-04 | 437.1ms/step | 45.8 samples/s
+Validation at step 15500: Loss = 3.0820, Perplexity = 21.80
+Step 15600/44000 | Loss: 3.0313 | LR: 3.68e-04 | 454.8ms/step | 44.0 samples/s
+Step 15700/44000 | Loss: 3.0489 | LR: 3.67e-04 | 436.7ms/step | 45.8 samples/s
+Validation at step 15750: Loss = 3.0796, Perplexity = 21.75
+Step 15800/44000 | Loss: 3.0370 | LR: 3.65e-04 | 454.2ms/step | 44.0 samples/s
+Step 15900/44000 | Loss: 3.0526 | LR: 3.63e-04 | 437.3ms/step | 45.7 samples/s
+Step 16000/44000 | Loss: 3.0398 | LR: 3.62e-04 | 439.0ms/step | 45.6 samples/s
+Validation at step 16000: Loss = 3.0750, Perplexity = 21.65
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-16000
+Step 16100/44000 | Loss: 3.0418 | LR: 3.60e-04 | 487.2ms/step | 41.0 samples/s
+Step 16200/44000 | Loss: 3.0545 | LR: 3.58e-04 | 437.7ms/step | 45.7 samples/s
+Validation at step 16250: Loss = 3.0699, Perplexity = 21.54
+Step 16300/44000 | Loss: 3.0429 | LR: 3.57e-04 | 454.6ms/step | 44.0 samples/s
+Step 16400/44000 | Loss: 3.0364 | LR: 3.55e-04 | 437.3ms/step | 45.7 samples/s
+Step 16500/44000 | Loss: 3.0330 | LR: 3.53e-04 | 438.3ms/step | 45.6 samples/s
+Validation at step 16500: Loss = 3.0672, Perplexity = 21.48
+Step 16600/44000 | Loss: 3.0124 | LR: 3.52e-04 | 454.7ms/step | 44.0 samples/s
+Step 16700/44000 | Loss: 3.0507 | LR: 3.50e-04 | 437.3ms/step | 45.7 samples/s
+Validation at step 16750: Loss = 3.0645, Perplexity = 21.42
+Step 16800/44000 | Loss: 3.0211 | LR: 3.48e-04 | 454.9ms/step | 44.0 samples/s
+Step 16900/44000 | Loss: 3.0055 | LR: 3.47e-04 | 438.0ms/step | 45.7 samples/s
+Step 17000/44000 | Loss: 3.0231 | LR: 3.45e-04 | 438.0ms/step | 45.7 samples/s
+Validation at step 17000: Loss = 3.0619, Perplexity = 21.37
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-17000
+Step 17100/44000 | Loss: 3.0356 | LR: 3.43e-04 | 483.9ms/step | 41.3 samples/s
+Step 17200/44000 | Loss: 3.0205 | LR: 3.42e-04 | 438.0ms/step | 45.7 samples/s
+Validation at step 17250: Loss = 3.0589, Perplexity = 21.30
+Step 17300/44000 | Loss: 3.0222 | LR: 3.40e-04 | 454.3ms/step | 44.0 samples/s
+Step 17400/44000 | Loss: 3.0271 | LR: 3.38e-04 | 437.0ms/step | 45.8 samples/s
+Step 17500/44000 | Loss: 2.9947 | LR: 3.37e-04 | 437.0ms/step | 45.8 samples/s
+Validation at step 17500: Loss = 3.0545, Perplexity = 21.21

epfl_final/logs/ALIBI_20260921/ddp5_a10040_s43_20260921_164403/seed43.log

@@ -0,0 +1,862 @@
+W0921 16:44:15.953000 145 site-packages/torch/distributed/run.py:851] 
+W0921 16:44:15.953000 145 site-packages/torch/distributed/run.py:851] *****************************************
+W0921 16:44:15.953000 145 site-packages/torch/distributed/run.py:851] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. 
+W0921 16:44:15.953000 145 site-packages/torch/distributed/run.py:851] *****************************************
+Auto-resuming from latest checkpoint: /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-2000
+Resuming: Using existing output directory /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951
+Running config:
+{
+  "model_path": "/mloscratch/homes/navasard/qwen3_06b_base",
+  "model_dtype": "bfloat16",
+  "init_from_scratch": true,
+  "train_path": "/mloscratch/homes/navasard/data/fineweb_large/tokenized",
+  "val_path": null,
+  "max_seq_length": 8192,
+  "num_workers": 8,
+  "output_dir": "/mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951",
+  "resume_from_checkpoint": "/mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-2000",
+  "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": 43,
+  "verbose": true,
+  "checkpoint_activations": false,
+  "resume_micro_step_override": 8000,
+  "strategy": "ddp",
+  "mixed_precision": "bf16",
+  "wandb_enabled": true,
+  "wandb_project": "stochastic_sparsity",
+  "wandb_name": "365m alibi 8k s43",
+  "wandb_tags": [
+    "qwen3",
+    "365m",
+    "8k",
+    "scaleladder",
+    "alibi",
+    "seed43"
+  ],
+  "hidden_size": 768,
+  "intermediate_size": 3072,
+  "num_attention_heads": 6,
+  "num_key_value_heads": 2,
+  "num_hidden_layers": 28,
+  "head_dim": null,
+  "attention_dropout": null,
+  "rope_train_position_scale_set": null,
+  "positional": "alibi",
+  "attention_type": "dense",
+  "sparse_block_size": 128,
+  "sparse_num_random_blocks": 31,
+  "sparse_mask_pool_size": 1,
+  "sparse_val_dense_compare": false,
+  "sparse_layer_indices": null,
+  "sparse_mask_cache_dir": null,
+  "sparse_mask_seed": 0,
+  "sparse_mixed_nrb_set": null,
+  "dense_mode": false,
+  "dense_step_intervals": null
+}
+Loading tokenizer from /mloscratch/homes/navasard/qwen3_06b_base...
+Checkpoint activations: False
+Checkpoint activations: False
+Checkpoint activations: False
+Resuming from checkpoint: /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-2000
+Checkpoint activations: False
+Checkpoint activations: False
+
+
+================================================================================
+MODEL ARCHITECTURE
+================================================================================
+Model config:
+  vocab_size: 151936
+  hidden_size: 768
+  num_hidden_layers: 28
+  num_attention_heads: 6
+  num_key_value_heads: 2
+  head_dim: 128
+  intermediate_size: 3072
+  max_position_embeddings: 32768
+  rope_theta: 1000000
+  rope_train_position_scale_set: None
+  attention_type: dense
+
+Parameter count:
+  Total parameters: 358,958,848
+  Trainable parameters: 358,958,848
+  Non-trainable parameters: 0
+  Model size (BF16): 0.72 GB
+
+--------------------------------------------------------------------------------
+LAYER-BY-LAYER SUMMARY
+--------------------------------------------------------------------------------
+/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/nn/attention/flex_attention.py:1704: UserWarning: flex_attention called without torch.compile() - this will use an unfused implementation that materializes the full scores matrix instead of generating a fused kernel.
+
+SOLUTION: Use torch.compile(flex_attention)(...)
+
+If you want to debug your score_mod/mask_mod, you can set:
+torch.nn.attention.flex_attention._FLEX_ATTENTION_DISABLE_COMPILE_DEBUG = True
+
+This will allow you to use print statements or breakpoints. Note: This doesn't work with the backwards pass and may produce incorrect results.
+  _warn_once(
+=================================================================================================================================================
+Layer (type:depth-idx)                        Input Shape               Output Shape              Param #                   Trainable
+=================================================================================================================================================
+Qwen3ForCausalLM                              [1, 8192]                 --                        --                        True
+├─Qwen3Model: 1-1                             --                        [1, 8192, 768]            --                        True
+│    └─Embedding: 2-1                         [1, 8192]                 [1, 8192, 768]            116,686,848               True
+│    └─ModuleList: 2-2                        --                        --                        --                        True
+│    │    └─Qwen3DecoderLayer: 3-1            [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-2            [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-3            [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-4            [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-5            [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-6            [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-7            [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-8            [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-9            [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-10           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-11           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-12           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-13           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-14           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-15           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-16           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-17           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-18           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-19           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-20           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-21           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-22           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-23           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-24           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-25           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-26           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-27           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-28           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    └─Qwen3RMSNorm: 2-3                      [1, 8192, 768]            [1, 8192, 768]            768                       True
+├─Linear: 1-2                                 [1, 8192, 768]            [1, 8192, 151936]         116,686,848               True
+=================================================================================================================================================
+Total params: 475,645,696
+Trainable params: 475,645,696
+Non-trainable params: 0
+Total mult-adds (Units.MEGABYTES): 475.65
+=================================================================================================================================================
+Input size (MB): 0.07
+Forward/backward pass size (MB): 15598.62
+Params size (MB): 951.29
+Estimated Total Size (MB): 16549.97
+=================================================================================================================================================
+  NOTE: torchinfo double-counts tied embed_tokens/lm_head weight (116,686,848 params). True unique parameter count: 358,958,848 (359.0M)
+================================================================================
+
+Loading dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Using pre-tokenized dataset (no runtime tokenization)
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Loaded 12748481 pre-tokenized examples (stored_len=2048, pack_factor=4, packed_examples=3187120)
+Loaded 12748481 pre-tokenized examples (stored_len=2048, pack_factor=4, packed_examples=3187120)
+Loaded 12748481 pre-tokenized examples (stored_len=2048, pack_factor=4, packed_examples=3187120)
+Loaded 12748481 pre-tokenized examples (stored_len=2048, pack_factor=4, packed_examples=3187120)
+Loaded 12748481 pre-tokenized examples (stored_len=2048, pack_factor=4, packed_examples=3187120)
+Dataset split: 3027764 train, 159356 validation
+Training examples (packed, train split):  3027764
+Validation examples (packed, val split):  159356
+Per-device batch size:                    1
+World size (GPUs):                        5
+Gradient accumulation steps:              4
+Effective batch size (seqs/step):         20
+Sequence length:                          8192
+Tokens per optimizer step:                163,840
+Micro-batches per epoch (per rank):       605553
+Optimizer steps per epoch:                151388
+Total optimizer steps:                    44000  (capped by max_steps; 0.29 epochs)
+Loading checkpoint from /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-2000...
+Preserving checkpoint LR state for exact resume: current=[0.0004989701850946613], schedule_base=[0.0005]
+Resumed from checkpoint at step 2000 (epoch 0, micro_step 20000)
+Overriding resume micro-step 20000 -> 8000 for DDP migration
+Resuming training from step 2000
+Resuming wandb run zgrsvkbi
+wandb: WARNING `resume` will be ignored since W&B syncing is set to `offline`. Starting a new run with run id zgrsvkbi.
+wandb: Tracking run with wandb version 0.26.1
+wandb: W&B syncing is set to `offline` in this directory. Run `wandb online` or set WANDB_MODE=online to enable cloud syncing.
+wandb: Run data is saved locally in /mloscratch/homes/navasard/sparse/SparseAtt/wandb/offline-run-20260921_164537-zgrsvkbi
+Starting training...
+Validation at step 2000: Loss = 3.6518, Perplexity = 38.55
+Switching attention mode for step 2001: sparse_stochastic (configured sparse attention)
+Step 2100/44000 | Loss: 3.6382 | LR: 4.99e-04 | 2592.7ms/step | 7.7 samples/s
+Step 2200/44000 | Loss: 3.6223 | LR: 4.99e-04 | 952.6ms/step | 21.0 samples/s
+Validation at step 2250: Loss = 3.5968, Perplexity = 36.48
+Step 2300/44000 | Loss: 3.5842 | LR: 4.98e-04 | 989.1ms/step | 20.2 samples/s
+Step 2400/44000 | Loss: 3.5884 | LR: 4.98e-04 | 957.1ms/step | 20.9 samples/s
+Step 2500/44000 | Loss: 3.5434 | LR: 4.98e-04 | 958.2ms/step | 20.9 samples/s
+Validation at step 2500: Loss = 3.5530, Perplexity = 34.92
+Step 2600/44000 | Loss: 3.5227 | LR: 4.98e-04 | 983.6ms/step | 20.3 samples/s
+Step 2700/44000 | Loss: 3.5398 | LR: 4.97e-04 | 957.5ms/step | 20.9 samples/s
+Validation at step 2750: Loss = 3.5089, Perplexity = 33.41
+Step 2800/44000 | Loss: 3.4975 | LR: 4.97e-04 | 986.8ms/step | 20.3 samples/s
+Step 2900/44000 | Loss: 3.4855 | LR: 4.97e-04 | 957.9ms/step | 20.9 samples/s
+Step 3000/44000 | Loss: 3.4846 | LR: 4.97e-04 | 956.8ms/step | 20.9 samples/s
+Validation at step 3000: Loss = 3.4766, Perplexity = 32.35
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-3000
+Step 3100/44000 | Loss: 3.4808 | LR: 4.96e-04 | 1025.0ms/step | 19.5 samples/s
+Step 3200/44000 | Loss: 3.4663 | LR: 4.96e-04 | 957.0ms/step | 20.9 samples/s
+Validation at step 3250: Loss = 3.4460, Perplexity = 31.37
+Step 3300/44000 | Loss: 3.4399 | LR: 4.96e-04 | 983.5ms/step | 20.3 samples/s
+Step 3400/44000 | Loss: 3.4406 | LR: 4.95e-04 | 957.0ms/step | 20.9 samples/s
+Step 3500/44000 | Loss: 3.4149 | LR: 4.95e-04 | 957.5ms/step | 20.9 samples/s
+Validation at step 3500: Loss = 3.4237, Perplexity = 30.68
+Step 3600/44000 | Loss: 3.4054 | LR: 4.95e-04 | 984.2ms/step | 20.3 samples/s
+Step 3700/44000 | Loss: 3.4051 | LR: 4.94e-04 | 956.3ms/step | 20.9 samples/s
+Validation at step 3750: Loss = 3.3963, Perplexity = 29.85
+Step 3800/44000 | Loss: 3.3714 | LR: 4.94e-04 | 985.1ms/step | 20.3 samples/s
+Step 3900/44000 | Loss: 3.4144 | LR: 4.93e-04 | 958.7ms/step | 20.9 samples/s
+Step 4000/44000 | Loss: 3.3740 | LR: 4.93e-04 | 957.7ms/step | 20.9 samples/s
+Validation at step 4000: Loss = 3.3777, Perplexity = 29.30
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-4000
+Step 4100/44000 | Loss: 3.3732 | LR: 4.93e-04 | 1026.9ms/step | 19.5 samples/s
+Step 4200/44000 | Loss: 3.3720 | LR: 4.92e-04 | 957.0ms/step | 20.9 samples/s
+Validation at step 4250: Loss = 3.3555, Perplexity = 28.66
+Step 4300/44000 | Loss: 3.3449 | LR: 4.92e-04 | 985.0ms/step | 20.3 samples/s
+Step 4400/44000 | Loss: 3.3408 | LR: 4.91e-04 | 956.3ms/step | 20.9 samples/s
+Step 4500/44000 | Loss: 3.3476 | LR: 4.91e-04 | 957.0ms/step | 20.9 samples/s
+Validation at step 4500: Loss = 3.3412, Perplexity = 28.25
+Step 4600/44000 | Loss: 3.3251 | LR: 4.90e-04 | 983.6ms/step | 20.3 samples/s
+Step 4700/44000 | Loss: 3.3299 | LR: 4.90e-04 | 956.1ms/step | 20.9 samples/s
+Validation at step 4750: Loss = 3.3240, Perplexity = 27.77
+Step 4800/44000 | Loss: 3.3377 | LR: 4.89e-04 | 984.0ms/step | 20.3 samples/s
+Step 4900/44000 | Loss: 3.3186 | LR: 4.89e-04 | 957.4ms/step | 20.9 samples/s
+Step 5000/44000 | Loss: 3.3082 | LR: 4.88e-04 | 957.6ms/step | 20.9 samples/s
+Validation at step 5000: Loss = 3.3073, Perplexity = 27.31
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-5000
+Step 5100/44000 | Loss: 3.3120 | LR: 4.88e-04 | 1023.6ms/step | 19.5 samples/s
+Step 5200/44000 | Loss: 3.3036 | LR: 4.87e-04 | 955.4ms/step | 20.9 samples/s
+Validation at step 5250: Loss = 3.2933, Perplexity = 26.93
+Step 5300/44000 | Loss: 3.2727 | LR: 4.86e-04 | 982.0ms/step | 20.4 samples/s
+Step 5400/44000 | Loss: 3.2945 | LR: 4.86e-04 | 955.6ms/step | 20.9 samples/s
+Step 5500/44000 | Loss: 3.2988 | LR: 4.85e-04 | 955.5ms/step | 20.9 samples/s
+Validation at step 5500: Loss = 3.2841, Perplexity = 26.68
+Step 5600/44000 | Loss: 3.2750 | LR: 4.85e-04 | 982.2ms/step | 20.4 samples/s
+Step 5700/44000 | Loss: 3.2701 | LR: 4.84e-04 | 955.7ms/step | 20.9 samples/s
+Validation at step 5750: Loss = 3.2746, Perplexity = 26.43
+Step 5800/44000 | Loss: 3.2681 | LR: 4.83e-04 | 983.7ms/step | 20.3 samples/s
+Step 5900/44000 | Loss: 3.2665 | LR: 4.83e-04 | 956.0ms/step | 20.9 samples/s
+Step 6000/44000 | Loss: 3.2554 | LR: 4.82e-04 | 955.8ms/step | 20.9 samples/s
+Validation at step 6000: Loss = 3.2616, Perplexity = 26.09
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-6000
+Step 6100/44000 | Loss: 3.2583 | LR: 4.81e-04 | 1024.7ms/step | 19.5 samples/s
+Step 6200/44000 | Loss: 3.2471 | LR: 4.81e-04 | 955.9ms/step | 20.9 samples/s
+Validation at step 6250: Loss = 3.2441, Perplexity = 25.64
+Step 6300/44000 | Loss: 3.2665 | LR: 4.80e-04 | 981.9ms/step | 20.4 samples/s
+Step 6400/44000 | Loss: 3.2482 | LR: 4.79e-04 | 955.7ms/step | 20.9 samples/s
+Step 6500/44000 | Loss: 3.1944 | LR: 4.79e-04 | 956.3ms/step | 20.9 samples/s
+Validation at step 6500: Loss = 3.2364, Perplexity = 25.44
+Step 6600/44000 | Loss: 3.2361 | LR: 4.78e-04 | 982.0ms/step | 20.4 samples/s
+Step 6700/44000 | Loss: 3.2358 | LR: 4.77e-04 | 956.4ms/step | 20.9 samples/s
+Validation at step 6750: Loss = 3.2257, Perplexity = 25.17
+Step 6800/44000 | Loss: 3.2354 | LR: 4.76e-04 | 981.9ms/step | 20.4 samples/s
+Step 6900/44000 | Loss: 3.2297 | LR: 4.75e-04 | 954.7ms/step | 20.9 samples/s
+Step 7000/44000 | Loss: 3.2051 | LR: 4.75e-04 | 955.7ms/step | 20.9 samples/s
+Validation at step 7000: Loss = 3.2184, Perplexity = 24.99
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-7000
+Step 7100/44000 | Loss: 3.2300 | LR: 4.74e-04 | 1024.2ms/step | 19.5 samples/s
+Step 7200/44000 | Loss: 3.2021 | LR: 4.73e-04 | 956.8ms/step | 20.9 samples/s
+Validation at step 7250: Loss = 3.2103, Perplexity = 24.79
+Step 7300/44000 | Loss: 3.1894 | LR: 4.72e-04 | 985.5ms/step | 20.3 samples/s
+Step 7400/44000 | Loss: 3.2239 | LR: 4.71e-04 | 959.0ms/step | 20.9 samples/s
+Step 7500/44000 | Loss: 3.1957 | LR: 4.71e-04 | 958.5ms/step | 20.9 samples/s
+Validation at step 7500: Loss = 3.2034, Perplexity = 24.62
+Step 7600/44000 | Loss: 3.1794 | LR: 4.70e-04 | 984.8ms/step | 20.3 samples/s
+Step 7700/44000 | Loss: 3.1877 | LR: 4.69e-04 | 959.5ms/step | 20.8 samples/s
+Validation at step 7750: Loss = 3.1947, Perplexity = 24.40
+Step 7800/44000 | Loss: 3.1963 | LR: 4.68e-04 | 986.0ms/step | 20.3 samples/s
+Step 7900/44000 | Loss: 3.2053 | LR: 4.67e-04 | 959.6ms/step | 20.8 samples/s
+Step 8000/44000 | Loss: 3.2034 | LR: 4.66e-04 | 957.8ms/step | 20.9 samples/s
+Validation at step 8000: Loss = 3.1850, Perplexity = 24.17
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-8000
+Step 8100/44000 | Loss: 3.1877 | LR: 4.65e-04 | 1028.1ms/step | 19.5 samples/s
+Step 8200/44000 | Loss: 3.1918 | LR: 4.64e-04 | 958.3ms/step | 20.9 samples/s
+Validation at step 8250: Loss = 3.1792, Perplexity = 24.03
+Step 8300/44000 | Loss: 3.1822 | LR: 4.63e-04 | 984.2ms/step | 20.3 samples/s
+Step 8400/44000 | Loss: 3.1500 | LR: 4.62e-04 | 957.3ms/step | 20.9 samples/s
+Step 8500/44000 | Loss: 3.1783 | LR: 4.61e-04 | 958.2ms/step | 20.9 samples/s
+Validation at step 8500: Loss = 3.1702, Perplexity = 23.81
+Step 8600/44000 | Loss: 3.1904 | LR: 4.60e-04 | 983.7ms/step | 20.3 samples/s
+Step 8700/44000 | Loss: 3.1651 | LR: 4.59e-04 | 957.8ms/step | 20.9 samples/s
+Validation at step 8750: Loss = 3.1635, Perplexity = 23.65
+Step 8800/44000 | Loss: 3.1577 | LR: 4.58e-04 | 983.7ms/step | 20.3 samples/s
+Step 8900/44000 | Loss: 3.1533 | LR: 4.57e-04 | 956.4ms/step | 20.9 samples/s
+Step 9000/44000 | Loss: 3.1506 | LR: 4.56e-04 | 956.7ms/step | 20.9 samples/s
+Validation at step 9000: Loss = 3.1557, Perplexity = 23.47
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-9000
+Step 9100/44000 | Loss: 3.1629 | LR: 4.55e-04 | 1024.9ms/step | 19.5 samples/s
+Step 9200/44000 | Loss: 3.1528 | LR: 4.54e-04 | 957.8ms/step | 20.9 samples/s
+Validation at step 9250: Loss = 3.1471, Perplexity = 23.27
+Step 9300/44000 | Loss: 3.1757 | LR: 4.53e-04 | 983.9ms/step | 20.3 samples/s
+Step 9400/44000 | Loss: 3.1365 | LR: 4.52e-04 | 957.4ms/step | 20.9 samples/s
+Step 9500/44000 | Loss: 3.1633 | LR: 4.51e-04 | 957.3ms/step | 20.9 samples/s
+Validation at step 9500: Loss = 3.1434, Perplexity = 23.18
+Step 9600/44000 | Loss: 3.1469 | LR: 4.50e-04 | 988.1ms/step | 20.2 samples/s
+Step 9700/44000 | Loss: 3.1566 | LR: 4.49e-04 | 962.4ms/step | 20.8 samples/s
+Validation at step 9750: Loss = 3.1369, Perplexity = 23.03
+Step 9800/44000 | Loss: 3.1517 | LR: 4.48e-04 | 989.5ms/step | 20.2 samples/s
+Step 9900/44000 | Loss: 3.1425 | LR: 4.47e-04 | 959.9ms/step | 20.8 samples/s
+Step 10000/44000 | Loss: 3.1324 | LR: 4.46e-04 | 958.7ms/step | 20.9 samples/s
+Validation at step 10000: Loss = 3.1323, Perplexity = 22.93
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-10000
+Step 10100/44000 | Loss: 3.1400 | LR: 4.45e-04 | 1026.1ms/step | 19.5 samples/s
+Step 10200/44000 | Loss: 3.1607 | LR: 4.43e-04 | 957.3ms/step | 20.9 samples/s
+Validation at step 10250: Loss = 3.1279, Perplexity = 22.83
+Step 10300/44000 | Loss: 3.1363 | LR: 4.42e-04 | 984.0ms/step | 20.3 samples/s
+Step 10400/44000 | Loss: 3.1263 | LR: 4.41e-04 | 956.9ms/step | 20.9 samples/s
+Step 10500/44000 | Loss: 3.1215 | LR: 4.40e-04 | 956.1ms/step | 20.9 samples/s
+Validation at step 10500: Loss = 3.1227, Perplexity = 22.71
+Step 10600/44000 | Loss: 3.1116 | LR: 4.39e-04 | 983.7ms/step | 20.3 samples/s
+Step 10700/44000 | Loss: 3.1225 | LR: 4.37e-04 | 956.1ms/step | 20.9 samples/s
+Validation at step 10750: Loss = 3.1193, Perplexity = 22.63
+Step 10800/44000 | Loss: 3.1269 | LR: 4.36e-04 | 983.5ms/step | 20.3 samples/s
+Step 10900/44000 | Loss: 3.1149 | LR: 4.35e-04 | 956.5ms/step | 20.9 samples/s
+Step 11000/44000 | Loss: 3.1303 | LR: 4.34e-04 | 956.6ms/step | 20.9 samples/s
+Validation at step 11000: Loss = 3.1128, Perplexity = 22.48
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-11000
+Step 11100/44000 | Loss: 3.1113 | LR: 4.33e-04 | 1024.5ms/step | 19.5 samples/s
+Step 11200/44000 | Loss: 3.1217 | LR: 4.31e-04 | 956.5ms/step | 20.9 samples/s
+Validation at step 11250: Loss = 3.1093, Perplexity = 22.40
+Step 11300/44000 | Loss: 3.1055 | LR: 4.30e-04 | 984.0ms/step | 20.3 samples/s
+Step 11400/44000 | Loss: 3.0908 | LR: 4.29e-04 | 957.1ms/step | 20.9 samples/s
+Step 11500/44000 | Loss: 3.1050 | LR: 4.28e-04 | 957.6ms/step | 20.9 samples/s
+Validation at step 11500: Loss = 3.1042, Perplexity = 22.29
+Step 11600/44000 | Loss: 3.0982 | LR: 4.26e-04 | 984.2ms/step | 20.3 samples/s
+Step 11700/44000 | Loss: 3.0911 | LR: 4.25e-04 | 957.5ms/step | 20.9 samples/s
+Validation at step 11750: Loss = 3.1012, Perplexity = 22.22
+Step 11800/44000 | Loss: 3.0764 | LR: 4.24e-04 | 983.5ms/step | 20.3 samples/s
+Step 11900/44000 | Loss: 3.0912 | LR: 4.22e-04 | 956.9ms/step | 20.9 samples/s
+Step 12000/44000 | Loss: 3.1236 | LR: 4.21e-04 | 957.3ms/step | 20.9 samples/s
+Validation at step 12000: Loss = 3.0953, Perplexity = 22.09
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-12000
+Step 12100/44000 | Loss: 3.0994 | LR: 4.20e-04 | 1025.7ms/step | 19.5 samples/s
+Step 12200/44000 | Loss: 3.0803 | LR: 4.18e-04 | 956.6ms/step | 20.9 samples/s
+Validation at step 12250: Loss = 3.0919, Perplexity = 22.02
+Step 12300/44000 | Loss: 3.0899 | LR: 4.17e-04 | 983.2ms/step | 20.3 samples/s
+Step 12400/44000 | Loss: 3.0987 | LR: 4.16e-04 | 957.2ms/step | 20.9 samples/s
+Step 12500/44000 | Loss: 3.0731 | LR: 4.14e-04 | 957.0ms/step | 20.9 samples/s
+Validation at step 12500: Loss = 3.0869, Perplexity = 21.91
+Step 12600/44000 | Loss: 3.0911 | LR: 4.13e-04 | 983.6ms/step | 20.3 samples/s
+Step 12700/44000 | Loss: 3.0684 | LR: 4.12e-04 | 957.1ms/step | 20.9 samples/s
+Validation at step 12750: Loss = 3.0820, Perplexity = 21.80
+Step 12800/44000 | Loss: 3.0938 | LR: 4.10e-04 | 982.5ms/step | 20.4 samples/s
+Step 12900/44000 | Loss: 3.0708 | LR: 4.09e-04 | 955.9ms/step | 20.9 samples/s
+Step 13000/44000 | Loss: 3.0916 | LR: 4.07e-04 | 956.0ms/step | 20.9 samples/s
+Validation at step 13000: Loss = 3.0795, Perplexity = 21.75
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-13000
+Step 13100/44000 | Loss: 3.0670 | LR: 4.06e-04 | 1024.4ms/step | 19.5 samples/s
+Step 13200/44000 | Loss: 3.0763 | LR: 4.05e-04 | 955.9ms/step | 20.9 samples/s
+Validation at step 13250: Loss = 3.0755, Perplexity = 21.66
+Step 13300/44000 | Loss: 3.0618 | LR: 4.03e-04 | 982.5ms/step | 20.4 samples/s
+Step 13400/44000 | Loss: 3.0820 | LR: 4.02e-04 | 956.1ms/step | 20.9 samples/s
+Step 13500/44000 | Loss: 3.0685 | LR: 4.00e-04 | 955.8ms/step | 20.9 samples/s
+Validation at step 13500: Loss = 3.0688, Perplexity = 21.52
+Step 13600/44000 | Loss: 3.0625 | LR: 3.99e-04 | 982.5ms/step | 20.4 samples/s
+Step 13700/44000 | Loss: 3.0753 | LR: 3.97e-04 | 955.7ms/step | 20.9 samples/s
+Validation at step 13750: Loss = 3.0670, Perplexity = 21.48
+Step 13800/44000 | Loss: 3.0650 | LR: 3.96e-04 | 982.2ms/step | 20.4 samples/s
+Step 13900/44000 | Loss: 3.0625 | LR: 3.94e-04 | 956.0ms/step | 20.9 samples/s
+Step 14000/44000 | Loss: 3.0447 | LR: 3.93e-04 | 955.8ms/step | 20.9 samples/s
+Validation at step 14000: Loss = 3.0644, Perplexity = 21.42
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-14000
+Step 14100/44000 | Loss: 3.0699 | LR: 3.91e-04 | 1024.5ms/step | 19.5 samples/s
+Step 14200/44000 | Loss: 3.0580 | LR: 3.90e-04 | 956.3ms/step | 20.9 samples/s
+Validation at step 14250: Loss = 3.0572, Perplexity = 21.27
+Step 14300/44000 | Loss: 3.0667 | LR: 3.88e-04 | 982.5ms/step | 20.4 samples/s
+Step 14400/44000 | Loss: 3.0645 | LR: 3.87e-04 | 956.0ms/step | 20.9 samples/s
+Step 14500/44000 | Loss: 3.0630 | LR: 3.85e-04 | 956.2ms/step | 20.9 samples/s
+Validation at step 14500: Loss = 3.0553, Perplexity = 21.23
+Step 14600/44000 | Loss: 3.0709 | LR: 3.84e-04 | 982.1ms/step | 20.4 samples/s
+Step 14700/44000 | Loss: 3.0569 | LR: 3.82e-04 | 955.4ms/step | 20.9 samples/s
+Validation at step 14750: Loss = 3.0529, Perplexity = 21.18
+Step 14800/44000 | Loss: 3.0559 | LR: 3.81e-04 | 982.5ms/step | 20.4 samples/s
+Step 14900/44000 | Loss: 3.0592 | LR: 3.79e-04 | 955.9ms/step | 20.9 samples/s
+Step 15000/44000 | Loss: 3.0296 | LR: 3.78e-04 | 955.6ms/step | 20.9 samples/s
+Validation at step 15000: Loss = 3.0500, Perplexity = 21.12
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-15000
+Step 15100/44000 | Loss: 3.0570 | LR: 3.76e-04 | 1021.3ms/step | 19.6 samples/s
+Step 15200/44000 | Loss: 3.0340 | LR: 3.74e-04 | 954.9ms/step | 20.9 samples/s
+Validation at step 15250: Loss = 3.0474, Perplexity = 21.06
+Step 15300/44000 | Loss: 3.0453 | LR: 3.73e-04 | 981.3ms/step | 20.4 samples/s
+Step 15400/44000 | Loss: 3.0343 | LR: 3.71e-04 | 954.7ms/step | 20.9 samples/s
+Step 15500/44000 | Loss: 3.0400 | LR: 3.70e-04 | 954.7ms/step | 21.0 samples/s
+Validation at step 15500: Loss = 3.0424, Perplexity = 20.96
+Step 15600/44000 | Loss: 3.0385 | LR: 3.68e-04 | 981.8ms/step | 20.4 samples/s
+Step 15700/44000 | Loss: 3.0386 | LR: 3.67e-04 | 954.9ms/step | 20.9 samples/s
+Validation at step 15750: Loss = 3.0390, Perplexity = 20.89
+Step 15800/44000 | Loss: 3.0538 | LR: 3.65e-04 | 980.9ms/step | 20.4 samples/s
+Step 15900/44000 | Loss: 3.0191 | LR: 3.63e-04 | 954.7ms/step | 20.9 samples/s
+Step 16000/44000 | Loss: 3.0461 | LR: 3.62e-04 | 955.0ms/step | 20.9 samples/s
+Validation at step 16000: Loss = 3.0346, Perplexity = 20.79
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-16000
+Step 16100/44000 | Loss: 3.0331 | LR: 3.60e-04 | 1024.1ms/step | 19.5 samples/s
+Step 16200/44000 | Loss: 3.0119 | LR: 3.58e-04 | 957.0ms/step | 20.9 samples/s
+Validation at step 16250: Loss = 3.0343, Perplexity = 20.79
+Step 16300/44000 | Loss: 3.0587 | LR: 3.57e-04 | 982.8ms/step | 20.3 samples/s
+Step 16400/44000 | Loss: 3.0392 | LR: 3.55e-04 | 956.1ms/step | 20.9 samples/s
+Step 16500/44000 | Loss: 3.0134 | LR: 3.53e-04 | 956.3ms/step | 20.9 samples/s
+Validation at step 16500: Loss = 3.0301, Perplexity = 20.70
+Step 16600/44000 | Loss: 3.0343 | LR: 3.52e-04 | 982.9ms/step | 20.3 samples/s
+Step 16700/44000 | Loss: 3.0395 | LR: 3.50e-04 | 956.5ms/step | 20.9 samples/s
+Validation at step 16750: Loss = 3.0260, Perplexity = 20.61
+Step 16800/44000 | Loss: 3.0314 | LR: 3.48e-04 | 982.9ms/step | 20.3 samples/s
+Step 16900/44000 | Loss: 3.0141 | LR: 3.47e-04 | 956.5ms/step | 20.9 samples/s
+Step 17000/44000 | Loss: 3.0233 | LR: 3.45e-04 | 956.3ms/step | 20.9 samples/s
+Validation at step 17000: Loss = 3.0228, Perplexity = 20.55
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-17000
+Step 17100/44000 | Loss: 3.0460 | LR: 3.43e-04 | 1021.9ms/step | 19.6 samples/s
+Step 17200/44000 | Loss: 2.9961 | LR: 3.42e-04 | 954.5ms/step | 21.0 samples/s
+Validation at step 17250: Loss = 3.0199, Perplexity = 20.49
+Step 17300/44000 | Loss: 3.0271 | LR: 3.40e-04 | 982.3ms/step | 20.4 samples/s
+Step 17400/44000 | Loss: 3.0121 | LR: 3.38e-04 | 955.0ms/step | 20.9 samples/s
+Step 17500/44000 | Loss: 3.0256 | LR: 3.37e-04 | 955.2ms/step | 20.9 samples/s
+Validation at step 17500: Loss = 3.0171, Perplexity = 20.43
+Step 17600/44000 | Loss: 3.0172 | LR: 3.35e-04 | 981.3ms/step | 20.4 samples/s
+Step 17700/44000 | Loss: 3.0247 | LR: 3.33e-04 | 954.9ms/step | 20.9 samples/s
+Validation at step 17750: Loss = 3.0165, Perplexity = 20.42
+Step 17800/44000 | Loss: 3.0119 | LR: 3.32e-04 | 981.4ms/step | 20.4 samples/s
+Step 17900/44000 | Loss: 3.0091 | LR: 3.30e-04 | 955.3ms/step | 20.9 samples/s
+Step 18000/44000 | Loss: 3.0049 | LR: 3.28e-04 | 955.4ms/step | 20.9 samples/s
+Validation at step 18000: Loss = 3.0083, Perplexity = 20.25
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-18000
+Step 18100/44000 | Loss: 2.9995 | LR: 3.26e-04 | 1023.4ms/step | 19.5 samples/s
+Step 18200/44000 | Loss: 3.0244 | LR: 3.25e-04 | 956.2ms/step | 20.9 samples/s
+Validation at step 18250: Loss = 3.0076, Perplexity = 20.24
+Step 18300/44000 | Loss: 3.0217 | LR: 3.23e-04 | 982.9ms/step | 20.3 samples/s
+Step 18400/44000 | Loss: 2.9827 | LR: 3.21e-04 | 956.0ms/step | 20.9 samples/s
+Step 18500/44000 | Loss: 2.9867 | LR: 3.19e-04 | 956.3ms/step | 20.9 samples/s
+Validation at step 18500: Loss = 3.0048, Perplexity = 20.18
+Step 18600/44000 | Loss: 2.9997 | LR: 3.18e-04 | 982.7ms/step | 20.4 samples/s
+Step 18700/44000 | Loss: 2.9981 | LR: 3.16e-04 | 956.4ms/step | 20.9 samples/s
+Validation at step 18750: Loss = 3.0013, Perplexity = 20.11
+Step 18800/44000 | Loss: 2.9810 | LR: 3.14e-04 | 982.9ms/step | 20.3 samples/s
+Step 18900/44000 | Loss: 2.9782 | LR: 3.12e-04 | 956.2ms/step | 20.9 samples/s
+Step 19000/44000 | Loss: 3.0140 | LR: 3.11e-04 | 956.6ms/step | 20.9 samples/s
+Validation at step 19000: Loss = 3.0017, Perplexity = 20.12
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-19000
+Step 19100/44000 | Loss: 2.9835 | LR: 3.09e-04 | 1021.9ms/step | 19.6 samples/s
+Step 19200/44000 | Loss: 2.9979 | LR: 3.07e-04 | 954.5ms/step | 21.0 samples/s
+Validation at step 19250: Loss = 2.9977, Perplexity = 20.04
+Step 19300/44000 | Loss: 3.0075 | LR: 3.05e-04 | 980.7ms/step | 20.4 samples/s
+Step 19400/44000 | Loss: 2.9921 | LR: 3.04e-04 | 954.6ms/step | 21.0 samples/s
+Step 19500/44000 | Loss: 2.9724 | LR: 3.02e-04 | 955.1ms/step | 20.9 samples/s
+Validation at step 19500: Loss = 2.9947, Perplexity = 19.98
+Step 19600/44000 | Loss: 2.9988 | LR: 3.00e-04 | 981.4ms/step | 20.4 samples/s
+Step 19700/44000 | Loss: 2.9922 | LR: 2.98e-04 | 954.6ms/step | 21.0 samples/s
+Validation at step 19750: Loss = 2.9929, Perplexity = 19.94
+Step 19800/44000 | Loss: 2.9866 | LR: 2.96e-04 | 981.2ms/step | 20.4 samples/s
+Step 19900/44000 | Loss: 2.9876 | LR: 2.95e-04 | 955.2ms/step | 20.9 samples/s
+Step 20000/44000 | Loss: 2.9704 | LR: 2.93e-04 | 955.1ms/step | 20.9 samples/s
+Validation at step 20000: Loss = 2.9895, Perplexity = 19.88
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-20000
+Step 20100/44000 | Loss: 2.9958 | LR: 2.91e-04 | 1022.0ms/step | 19.6 samples/s
+Step 20200/44000 | Loss: 2.9630 | LR: 2.89e-04 | 955.0ms/step | 20.9 samples/s
+Validation at step 20250: Loss = 2.9865, Perplexity = 19.82
+Step 20300/44000 | Loss: 2.9767 | LR: 2.88e-04 | 981.5ms/step | 20.4 samples/s
+Step 20400/44000 | Loss: 2.9899 | LR: 2.86e-04 | 955.2ms/step | 20.9 samples/s
+Step 20500/44000 | Loss: 3.0059 | LR: 2.84e-04 | 955.3ms/step | 20.9 samples/s
+Validation at step 20500: Loss = 2.9836, Perplexity = 19.76
+Step 20600/44000 | Loss: 2.9946 | LR: 2.82e-04 | 981.9ms/step | 20.4 samples/s
+Step 20700/44000 | Loss: 2.9970 | LR: 2.80e-04 | 955.0ms/step | 20.9 samples/s
+Validation at step 20750: Loss = 2.9815, Perplexity = 19.72
+Step 20800/44000 | Loss: 2.9686 | LR: 2.79e-04 | 981.3ms/step | 20.4 samples/s
+Step 20900/44000 | Loss: 2.9701 | LR: 2.77e-04 | 955.8ms/step | 20.9 samples/s
+Step 21000/44000 | Loss: 2.9939 | LR: 2.75e-04 | 955.0ms/step | 20.9 samples/s
+Validation at step 21000: Loss = 2.9781, Perplexity = 19.65
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-21000
+Step 21100/44000 | Loss: 2.9968 | LR: 2.73e-04 | 1022.1ms/step | 19.6 samples/s
+Step 21200/44000 | Loss: 2.9705 | LR: 2.71e-04 | 954.9ms/step | 20.9 samples/s
+Validation at step 21250: Loss = 2.9754, Perplexity = 19.60
+Step 21300/44000 | Loss: 2.9698 | LR: 2.70e-04 | 981.9ms/step | 20.4 samples/s
+Step 21400/44000 | Loss: 2.9986 | LR: 2.68e-04 | 956.3ms/step | 20.9 samples/s
+Step 21500/44000 | Loss: 2.9692 | LR: 2.66e-04 | 956.9ms/step | 20.9 samples/s
+Validation at step 21500: Loss = 2.9729, Perplexity = 19.55
+Step 21600/44000 | Loss: 2.9726 | LR: 2.64e-04 | 983.1ms/step | 20.3 samples/s
+Step 21700/44000 | Loss: 2.9580 | LR: 2.62e-04 | 957.0ms/step | 20.9 samples/s
+Validation at step 21750: Loss = 2.9708, Perplexity = 19.51
+Step 21800/44000 | Loss: 2.9645 | LR: 2.60e-04 | 983.0ms/step | 20.3 samples/s
+Step 21900/44000 | Loss: 2.9573 | LR: 2.59e-04 | 956.2ms/step | 20.9 samples/s
+Step 22000/44000 | Loss: 2.9552 | LR: 2.57e-04 | 956.2ms/step | 20.9 samples/s
+Validation at step 22000: Loss = 2.9673, Perplexity = 19.44
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-22000
+Step 22100/44000 | Loss: 2.9430 | LR: 2.55e-04 | 1025.2ms/step | 19.5 samples/s
+Step 22200/44000 | Loss: 2.9660 | LR: 2.53e-04 | 956.7ms/step | 20.9 samples/s
+Validation at step 22250: Loss = 2.9652, Perplexity = 19.40
+Step 22300/44000 | Loss: 2.9641 | LR: 2.51e-04 | 983.0ms/step | 20.3 samples/s
+Step 22400/44000 | Loss: 2.9549 | LR: 2.50e-04 | 957.0ms/step | 20.9 samples/s
+Step 22500/44000 | Loss: 2.9714 | LR: 2.48e-04 | 957.4ms/step | 20.9 samples/s
+Validation at step 22500: Loss = 2.9627, Perplexity = 19.35
+Step 22600/44000 | Loss: 2.9415 | LR: 2.46e-04 | 983.9ms/step | 20.3 samples/s
+Step 22700/44000 | Loss: 2.9645 | LR: 2.44e-04 | 956.6ms/step | 20.9 samples/s
+Validation at step 22750: Loss = 2.9617, Perplexity = 19.33
+Step 22800/44000 | Loss: 2.9512 | LR: 2.42e-04 | 983.3ms/step | 20.3 samples/s
+Step 22900/44000 | Loss: 2.9777 | LR: 2.40e-04 | 956.6ms/step | 20.9 samples/s
+Step 23000/44000 | Loss: 2.9460 | LR: 2.39e-04 | 956.6ms/step | 20.9 samples/s
+Validation at step 23000: Loss = 2.9581, Perplexity = 19.26
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-23000
+Step 23100/44000 | Loss: 2.9421 | LR: 2.37e-04 | 1022.4ms/step | 19.6 samples/s
+Step 23200/44000 | Loss: 2.9967 | LR: 2.35e-04 | 955.5ms/step | 20.9 samples/s
+Validation at step 23250: Loss = 2.9568, Perplexity = 19.24
+Step 23300/44000 | Loss: 2.9426 | LR: 2.33e-04 | 981.7ms/step | 20.4 samples/s
+Step 23400/44000 | Loss: 2.9479 | LR: 2.31e-04 | 955.8ms/step | 20.9 samples/s
+Step 23500/44000 | Loss: 2.9707 | LR: 2.30e-04 | 955.8ms/step | 20.9 samples/s
+Validation at step 23500: Loss = 2.9538, Perplexity = 19.18
+Step 23600/44000 | Loss: 2.9191 | LR: 2.28e-04 | 981.9ms/step | 20.4 samples/s
+Step 23700/44000 | Loss: 2.9681 | LR: 2.26e-04 | 955.3ms/step | 20.9 samples/s
+Validation at step 23750: Loss = 2.9513, Perplexity = 19.13
+Step 23800/44000 | Loss: 2.9507 | LR: 2.24e-04 | 981.3ms/step | 20.4 samples/s
+Step 23900/44000 | Loss: 2.9509 | LR: 2.22e-04 | 954.6ms/step | 21.0 samples/s
+Step 24000/44000 | Loss: 2.9527 | LR: 2.21e-04 | 954.8ms/step | 20.9 samples/s
+Validation at step 24000: Loss = 2.9502, Perplexity = 19.11
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-24000
+Step 24100/44000 | Loss: 2.9604 | LR: 2.19e-04 | 1024.3ms/step | 19.5 samples/s
+Step 24200/44000 | Loss: 2.9366 | LR: 2.17e-04 | 956.9ms/step | 20.9 samples/s
+Validation at step 24250: Loss = 2.9484, Perplexity = 19.08
+Step 24300/44000 | Loss: 2.9470 | LR: 2.15e-04 | 983.0ms/step | 20.3 samples/s
+Step 24400/44000 | Loss: 2.9505 | LR: 2.13e-04 | 956.5ms/step | 20.9 samples/s
+Step 24500/44000 | Loss: 2.9249 | LR: 2.12e-04 | 956.7ms/step | 20.9 samples/s
+Validation at step 24500: Loss = 2.9470, Perplexity = 19.05
+Step 24600/44000 | Loss: 2.9230 | LR: 2.10e-04 | 983.1ms/step | 20.3 samples/s
+Step 24700/44000 | Loss: 2.9387 | LR: 2.08e-04 | 956.8ms/step | 20.9 samples/s
+Validation at step 24750: Loss = 2.9444, Perplexity = 19.00
+Step 24800/44000 | Loss: 2.9506 | LR: 2.06e-04 | 983.0ms/step | 20.3 samples/s
+Step 24900/44000 | Loss: 2.9335 | LR: 2.04e-04 | 956.5ms/step | 20.9 samples/s
+Step 25000/44000 | Loss: 2.9447 | LR: 2.03e-04 | 957.0ms/step | 20.9 samples/s
+Validation at step 25000: Loss = 2.9439, Perplexity = 18.99
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-25000
+Step 25100/44000 | Loss: 2.9273 | LR: 2.01e-04 | 1023.1ms/step | 19.5 samples/s
+Step 25200/44000 | Loss: 2.9465 | LR: 1.99e-04 | 956.7ms/step | 20.9 samples/s
+Validation at step 25250: Loss = 2.9408, Perplexity = 18.93
+Step 25300/44000 | Loss: 2.9377 | LR: 1.97e-04 | 983.2ms/step | 20.3 samples/s
+Step 25400/44000 | Loss: 2.9310 | LR: 1.96e-04 | 956.7ms/step | 20.9 samples/s
+Step 25500/44000 | Loss: 2.9255 | LR: 1.94e-04 | 956.8ms/step | 20.9 samples/s
+Validation at step 25500: Loss = 2.9402, Perplexity = 18.92
+Step 25600/44000 | Loss: 2.9340 | LR: 1.92e-04 | 982.9ms/step | 20.3 samples/s
+Step 25700/44000 | Loss: 2.9075 | LR: 1.90e-04 | 956.4ms/step | 20.9 samples/s
+Validation at step 25750: Loss = 2.9391, Perplexity = 18.90
+Step 25800/44000 | Loss: 2.9276 | LR: 1.88e-04 | 983.0ms/step | 20.3 samples/s
+Step 25900/44000 | Loss: 2.9216 | LR: 1.87e-04 | 956.8ms/step | 20.9 samples/s
+Step 26000/44000 | Loss: 2.9396 | LR: 1.85e-04 | 956.5ms/step | 20.9 samples/s
+Validation at step 26000: Loss = 2.9375, Perplexity = 18.87
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-26000
+Step 26100/44000 | Loss: 2.9103 | LR: 1.83e-04 | 1022.3ms/step | 19.6 samples/s
+Step 26200/44000 | Loss: 2.9384 | LR: 1.81e-04 | 955.2ms/step | 20.9 samples/s
+Validation at step 26250: Loss = 2.9358, Perplexity = 18.84
+Step 26300/44000 | Loss: 2.9284 | LR: 1.80e-04 | 981.6ms/step | 20.4 samples/s
+Step 26400/44000 | Loss: 2.9373 | LR: 1.78e-04 | 955.0ms/step | 20.9 samples/s
+Step 26500/44000 | Loss: 2.9269 | LR: 1.76e-04 | 955.2ms/step | 20.9 samples/s
+Validation at step 26500: Loss = 2.9344, Perplexity = 18.81
+Step 26600/44000 | Loss: 2.9004 | LR: 1.74e-04 | 981.4ms/step | 20.4 samples/s
+Step 26700/44000 | Loss: 2.9123 | LR: 1.73e-04 | 954.8ms/step | 20.9 samples/s
+Validation at step 26750: Loss = 2.9325, Perplexity = 18.77
+Step 26800/44000 | Loss: 2.9132 | LR: 1.71e-04 | 981.6ms/step | 20.4 samples/s
+Step 26900/44000 | Loss: 2.9373 | LR: 1.69e-04 | 955.3ms/step | 20.9 samples/s
+Step 27000/44000 | Loss: 2.9080 | LR: 1.68e-04 | 955.3ms/step | 20.9 samples/s
+Validation at step 27000: Loss = 2.9313, Perplexity = 18.75
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-27000
+Step 27100/44000 | Loss: 2.9245 | LR: 1.66e-04 | 1022.0ms/step | 19.6 samples/s
+Step 27200/44000 | Loss: 2.9370 | LR: 1.64e-04 | 955.0ms/step | 20.9 samples/s
+Validation at step 27250: Loss = 2.9306, Perplexity = 18.74
+Step 27300/44000 | Loss: 2.9112 | LR: 1.62e-04 | 981.7ms/step | 20.4 samples/s
+Step 27400/44000 | Loss: 2.9219 | LR: 1.61e-04 | 954.7ms/step | 20.9 samples/s
+Step 27500/44000 | Loss: 2.9132 | LR: 1.59e-04 | 954.5ms/step | 21.0 samples/s
+Validation at step 27500: Loss = 2.9285, Perplexity = 18.70
+Step 27600/44000 | Loss: 2.9314 | LR: 1.57e-04 | 981.1ms/step | 20.4 samples/s
+Step 27700/44000 | Loss: 2.9083 | LR: 1.56e-04 | 954.9ms/step | 20.9 samples/s
+Validation at step 27750: Loss = 2.9272, Perplexity = 18.67
+Step 27800/44000 | Loss: 2.9122 | LR: 1.54e-04 | 981.2ms/step | 20.4 samples/s
+Step 27900/44000 | Loss: 2.9302 | LR: 1.52e-04 | 955.0ms/step | 20.9 samples/s
+Step 28000/44000 | Loss: 2.9224 | LR: 1.51e-04 | 955.0ms/step | 20.9 samples/s
+Validation at step 28000: Loss = 2.9258, Perplexity = 18.65
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-28000
+Step 28100/44000 | Loss: 2.9228 | LR: 1.49e-04 | 1021.8ms/step | 19.6 samples/s
+Step 28200/44000 | Loss: 2.9144 | LR: 1.47e-04 | 955.0ms/step | 20.9 samples/s
+Validation at step 28250: Loss = 2.9248, Perplexity = 18.63
+Step 28300/44000 | Loss: 2.9146 | LR: 1.46e-04 | 981.5ms/step | 20.4 samples/s
+Step 28400/44000 | Loss: 2.9198 | LR: 1.44e-04 | 954.9ms/step | 20.9 samples/s
+Step 28500/44000 | Loss: 2.9261 | LR: 1.42e-04 | 955.2ms/step | 20.9 samples/s
+Validation at step 28500: Loss = 2.9233, Perplexity = 18.60
+Step 28600/44000 | Loss: 2.9086 | LR: 1.41e-04 | 981.1ms/step | 20.4 samples/s
+Step 28700/44000 | Loss: 2.9075 | LR: 1.39e-04 | 954.8ms/step | 20.9 samples/s
+Validation at step 28750: Loss = 2.9224, Perplexity = 18.59
+Step 28800/44000 | Loss: 2.8959 | LR: 1.38e-04 | 981.0ms/step | 20.4 samples/s
+Step 28900/44000 | Loss: 2.9222 | LR: 1.36e-04 | 955.2ms/step | 20.9 samples/s
+Step 29000/44000 | Loss: 2.9346 | LR: 1.34e-04 | 955.3ms/step | 20.9 samples/s
+Validation at step 29000: Loss = 2.9212, Perplexity = 18.56
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-29000
+Step 29100/44000 | Loss: 2.9077 | LR: 1.33e-04 | 1023.6ms/step | 19.5 samples/s
+Step 29200/44000 | Loss: 2.9065 | LR: 1.31e-04 | 956.6ms/step | 20.9 samples/s
+Validation at step 29250: Loss = 2.9204, Perplexity = 18.55
+Step 29300/44000 | Loss: 2.9028 | LR: 1.29e-04 | 983.4ms/step | 20.3 samples/s
+Step 29400/44000 | Loss: 2.9217 | LR: 1.28e-04 | 956.5ms/step | 20.9 samples/s
+Step 29500/44000 | Loss: 2.9037 | LR: 1.26e-04 | 956.5ms/step | 20.9 samples/s
+Validation at step 29500: Loss = 2.9196, Perplexity = 18.53
+Step 29600/44000 | Loss: 2.9031 | LR: 1.25e-04 | 982.5ms/step | 20.4 samples/s
+Step 29700/44000 | Loss: 2.8981 | LR: 1.23e-04 | 956.3ms/step | 20.9 samples/s
+Validation at step 29750: Loss = 2.9186, Perplexity = 18.52
+Step 29800/44000 | Loss: 2.9018 | LR: 1.22e-04 | 982.6ms/step | 20.4 samples/s
+Step 29900/44000 | Loss: 2.9061 | LR: 1.20e-04 | 955.7ms/step | 20.9 samples/s
+Step 30000/44000 | Loss: 2.9215 | LR: 1.19e-04 | 956.0ms/step | 20.9 samples/s
+Validation at step 30000: Loss = 2.9177, Perplexity = 18.50
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-30000
+Step 30100/44000 | Loss: 2.9074 | LR: 1.17e-04 | 1022.2ms/step | 19.6 samples/s
+Step 30200/44000 | Loss: 2.9182 | LR: 1.15e-04 | 956.1ms/step | 20.9 samples/s
+Validation at step 30250: Loss = 2.9170, Perplexity = 18.49
+Step 30300/44000 | Loss: 2.8947 | LR: 1.14e-04 | 982.6ms/step | 20.4 samples/s
+Step 30400/44000 | Loss: 2.8964 | LR: 1.12e-04 | 956.5ms/step | 20.9 samples/s
+Step 30500/44000 | Loss: 2.8844 | LR: 1.11e-04 | 956.9ms/step | 20.9 samples/s
+Validation at step 30500: Loss = 2.9160, Perplexity = 18.47
+Step 30600/44000 | Loss: 2.9059 | LR: 1.09e-04 | 983.0ms/step | 20.3 samples/s
+Step 30700/44000 | Loss: 2.8911 | LR: 1.08e-04 | 956.3ms/step | 20.9 samples/s
+Validation at step 30750: Loss = 2.9155, Perplexity = 18.46
+Step 30800/44000 | Loss: 2.9241 | LR: 1.06e-04 | 982.8ms/step | 20.4 samples/s
+Step 30900/44000 | Loss: 2.9059 | LR: 1.05e-04 | 956.0ms/step | 20.9 samples/s
+Step 31000/44000 | Loss: 2.8828 | LR: 1.03e-04 | 956.2ms/step | 20.9 samples/s
+Validation at step 31000: Loss = 2.9146, Perplexity = 18.44
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-31000
+Step 31100/44000 | Loss: 2.9211 | LR: 1.02e-04 | 1023.0ms/step | 19.5 samples/s
+Step 31200/44000 | Loss: 2.9118 | LR: 1.00e-04 | 956.1ms/step | 20.9 samples/s
+Validation at step 31250: Loss = 2.9138, Perplexity = 18.43
+Step 31300/44000 | Loss: 2.9110 | LR: 9.90e-05 | 982.8ms/step | 20.3 samples/s
+Step 31400/44000 | Loss: 2.9140 | LR: 9.76e-05 | 956.4ms/step | 20.9 samples/s
+Step 31500/44000 | Loss: 2.9093 | LR: 9.62e-05 | 956.6ms/step | 20.9 samples/s
+Validation at step 31500: Loss = 2.9130, Perplexity = 18.41
+Step 31600/44000 | Loss: 2.9151 | LR: 9.47e-05 | 982.9ms/step | 20.3 samples/s
+Step 31700/44000 | Loss: 2.9153 | LR: 9.33e-05 | 956.7ms/step | 20.9 samples/s
+Validation at step 31750: Loss = 2.9126, Perplexity = 18.40
+Step 31800/44000 | Loss: 2.9209 | LR: 9.19e-05 | 982.9ms/step | 20.3 samples/s
+Step 31900/44000 | Loss: 2.9020 | LR: 9.05e-05 | 956.6ms/step | 20.9 samples/s
+Step 32000/44000 | Loss: 2.9257 | LR: 8.91e-05 | 956.3ms/step | 20.9 samples/s
+Validation at step 32000: Loss = 2.9122, Perplexity = 18.40
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-32000
+Step 32100/44000 | Loss: 2.8984 | LR: 8.77e-05 | 1023.5ms/step | 19.5 samples/s
+Step 32200/44000 | Loss: 2.9190 | LR: 8.63e-05 | 956.2ms/step | 20.9 samples/s
+Validation at step 32250: Loss = 2.9119, Perplexity = 18.39
+Step 32300/44000 | Loss: 2.8999 | LR: 8.50e-05 | 982.7ms/step | 20.4 samples/s
+Step 32400/44000 | Loss: 2.8866 | LR: 8.36e-05 | 956.4ms/step | 20.9 samples/s
+Step 32500/44000 | Loss: 2.9144 | LR: 8.23e-05 | 956.3ms/step | 20.9 samples/s
+Validation at step 32500: Loss = 2.9113, Perplexity = 18.38
+Step 32600/44000 | Loss: 2.9176 | LR: 8.09e-05 | 982.8ms/step | 20.3 samples/s
+Step 32700/44000 | Loss: 2.8896 | LR: 7.96e-05 | 956.1ms/step | 20.9 samples/s
+Validation at step 32750: Loss = 2.9108, Perplexity = 18.37
+Step 32800/44000 | Loss: 2.9258 | LR: 7.83e-05 | 982.4ms/step | 20.4 samples/s
+Step 32900/44000 | Loss: 2.9141 | LR: 7.70e-05 | 956.1ms/step | 20.9 samples/s
+Step 33000/44000 | Loss: 2.9065 | LR: 7.56e-05 | 956.8ms/step | 20.9 samples/s
+Validation at step 33000: Loss = 2.9105, Perplexity = 18.37
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-33000
+Step 33100/44000 | Loss: 2.8869 | LR: 7.44e-05 | 1022.2ms/step | 19.6 samples/s
+Step 33200/44000 | Loss: 2.9058 | LR: 7.31e-05 | 955.8ms/step | 20.9 samples/s
+Validation at step 33250: Loss = 2.9102, Perplexity = 18.36
+Step 33300/44000 | Loss: 2.8882 | LR: 7.18e-05 | 981.8ms/step | 20.4 samples/s
+Step 33400/44000 | Loss: 2.8906 | LR: 7.05e-05 | 955.1ms/step | 20.9 samples/s
+Step 33500/44000 | Loss: 2.9183 | LR: 6.93e-05 | 955.1ms/step | 20.9 samples/s
+Validation at step 33500: Loss = 2.9097, Perplexity = 18.35
+Step 33600/44000 | Loss: 2.9172 | LR: 6.80e-05 | 983.3ms/step | 20.3 samples/s
+Step 33700/44000 | Loss: 2.9241 | LR: 6.68e-05 | 955.4ms/step | 20.9 samples/s
+Validation at step 33750: Loss = 2.9091, Perplexity = 18.34
+Step 33800/44000 | Loss: 2.8981 | LR: 6.55e-05 | 981.5ms/step | 20.4 samples/s
+Step 33900/44000 | Loss: 2.9291 | LR: 6.43e-05 | 955.2ms/step | 20.9 samples/s
+Step 34000/44000 | Loss: 2.9160 | LR: 6.31e-05 | 955.4ms/step | 20.9 samples/s
+Validation at step 34000: Loss = 2.9090, Perplexity = 18.34
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-34000
+Step 34100/44000 | Loss: 2.8999 | LR: 6.19e-05 | 1026.7ms/step | 19.5 samples/s
+Step 34200/44000 | Loss: 2.8906 | LR: 6.07e-05 | 957.5ms/step | 20.9 samples/s
+Validation at step 34250: Loss = 2.9087, Perplexity = 18.33
+Step 34300/44000 | Loss: 2.9100 | LR: 5.95e-05 | 983.2ms/step | 20.3 samples/s
+Step 34400/44000 | Loss: 2.9060 | LR: 5.84e-05 | 956.9ms/step | 20.9 samples/s
+Step 34500/44000 | Loss: 2.9152 | LR: 5.72e-05 | 956.3ms/step | 20.9 samples/s
+Validation at step 34500: Loss = 2.9085, Perplexity = 18.33
+Step 34600/44000 | Loss: 2.8988 | LR: 5.60e-05 | 982.6ms/step | 20.4 samples/s
+Step 34700/44000 | Loss: 2.9235 | LR: 5.49e-05 | 956.5ms/step | 20.9 samples/s
+Validation at step 34750: Loss = 2.9082, Perplexity = 18.32
+Step 34800/44000 | Loss: 2.8992 | LR: 5.38e-05 | 982.7ms/step | 20.4 samples/s
+Step 34900/44000 | Loss: 2.9080 | LR: 5.27e-05 | 956.9ms/step | 20.9 samples/s
+Step 35000/44000 | Loss: 2.9224 | LR: 5.15e-05 | 956.7ms/step | 20.9 samples/s
+Validation at step 35000: Loss = 2.9080, Perplexity = 18.32
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-35000
+Step 35100/44000 | Loss: 2.8947 | LR: 5.04e-05 | 1026.3ms/step | 19.5 samples/s
+Step 35200/44000 | Loss: 2.8861 | LR: 4.94e-05 | 956.8ms/step | 20.9 samples/s
+Validation at step 35250: Loss = 2.9078, Perplexity = 18.32
+Step 35300/44000 | Loss: 2.9092 | LR: 4.83e-05 | 983.3ms/step | 20.3 samples/s
+Step 35400/44000 | Loss: 2.8840 | LR: 4.72e-05 | 956.8ms/step | 20.9 samples/s
+Step 35500/44000 | Loss: 2.9026 | LR: 4.62e-05 | 956.6ms/step | 20.9 samples/s
+Validation at step 35500: Loss = 2.9076, Perplexity = 18.31
+Step 35600/44000 | Loss: 2.8955 | LR: 4.51e-05 | 982.7ms/step | 20.4 samples/s
+Step 35700/44000 | Loss: 2.9119 | LR: 4.41e-05 | 956.4ms/step | 20.9 samples/s
+Validation at step 35750: Loss = 2.9072, Perplexity = 18.31
+Step 35800/44000 | Loss: 2.9062 | LR: 4.31e-05 | 982.8ms/step | 20.3 samples/s
+Step 35900/44000 | Loss: 2.9021 | LR: 4.20e-05 | 957.3ms/step | 20.9 samples/s
+Step 36000/44000 | Loss: 2.8644 | LR: 4.10e-05 | 956.3ms/step | 20.9 samples/s
+Validation at step 36000: Loss = 2.9072, Perplexity = 18.30
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-36000
+Step 36100/44000 | Loss: 2.8916 | LR: 4.00e-05 | 1024.5ms/step | 19.5 samples/s
+Step 36200/44000 | Loss: 2.8997 | LR: 3.91e-05 | 956.4ms/step | 20.9 samples/s
+Validation at step 36250: Loss = 2.9070, Perplexity = 18.30
+Step 36300/44000 | Loss: 2.9009 | LR: 3.81e-05 | 983.0ms/step | 20.3 samples/s
+Step 36400/44000 | Loss: 2.8996 | LR: 3.71e-05 | 956.8ms/step | 20.9 samples/s
+Step 36500/44000 | Loss: 2.8793 | LR: 3.62e-05 | 956.2ms/step | 20.9 samples/s
+Validation at step 36500: Loss = 2.9068, Perplexity = 18.30
+Step 36600/44000 | Loss: 2.9007 | LR: 3.53e-05 | 982.3ms/step | 20.4 samples/s
+Step 36700/44000 | Loss: 2.8990 | LR: 3.43e-05 | 956.3ms/step | 20.9 samples/s
+Validation at step 36750: Loss = 2.9068, Perplexity = 18.30
+Step 36800/44000 | Loss: 2.8915 | LR: 3.34e-05 | 982.3ms/step | 20.4 samples/s
+Step 36900/44000 | Loss: 2.8977 | LR: 3.25e-05 | 956.0ms/step | 20.9 samples/s
+Step 37000/44000 | Loss: 2.9182 | LR: 3.16e-05 | 956.2ms/step | 20.9 samples/s
+Validation at step 37000: Loss = 2.9067, Perplexity = 18.30
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-37000
+Step 37100/44000 | Loss: 2.8979 | LR: 3.07e-05 | 1022.7ms/step | 19.6 samples/s
+Step 37200/44000 | Loss: 2.9071 | LR: 2.99e-05 | 955.5ms/step | 20.9 samples/s
+Validation at step 37250: Loss = 2.9066, Perplexity = 18.29
+Step 37300/44000 | Loss: 2.9064 | LR: 2.90e-05 | 981.4ms/step | 20.4 samples/s
+Step 37400/44000 | Loss: 2.8837 | LR: 2.82e-05 | 955.1ms/step | 20.9 samples/s
+Step 37500/44000 | Loss: 2.9241 | LR: 2.74e-05 | 954.8ms/step | 20.9 samples/s
+Validation at step 37500: Loss = 2.9064, Perplexity = 18.29
+Step 37600/44000 | Loss: 2.8786 | LR: 2.65e-05 | 981.0ms/step | 20.4 samples/s
+Step 37700/44000 | Loss: 2.9124 | LR: 2.57e-05 | 954.8ms/step | 20.9 samples/s
+Validation at step 37750: Loss = 2.9065, Perplexity = 18.29
+Step 37800/44000 | Loss: 2.8824 | LR: 2.49e-05 | 981.2ms/step | 20.4 samples/s
+Step 37900/44000 | Loss: 2.9065 | LR: 2.41e-05 | 955.1ms/step | 20.9 samples/s
+Step 38000/44000 | Loss: 2.9038 | LR: 2.34e-05 | 954.7ms/step | 20.9 samples/s
+Validation at step 38000: Loss = 2.9064, Perplexity = 18.29
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-38000
+Step 38100/44000 | Loss: 2.9123 | LR: 2.26e-05 | 1022.8ms/step | 19.6 samples/s
+Step 38200/44000 | Loss: 2.8908 | LR: 2.19e-05 | 955.0ms/step | 20.9 samples/s
+Validation at step 38250: Loss = 2.9062, Perplexity = 18.29
+Step 38300/44000 | Loss: 2.8900 | LR: 2.11e-05 | 980.8ms/step | 20.4 samples/s
+Step 38400/44000 | Loss: 2.9047 | LR: 2.04e-05 | 954.7ms/step | 20.9 samples/s
+Step 38500/44000 | Loss: 2.9184 | LR: 1.97e-05 | 954.4ms/step | 21.0 samples/s
+Validation at step 38500: Loss = 2.9062, Perplexity = 18.29
+Step 38600/44000 | Loss: 2.9094 | LR: 1.90e-05 | 982.2ms/step | 20.4 samples/s
+Step 38700/44000 | Loss: 2.9055 | LR: 1.83e-05 | 955.5ms/step | 20.9 samples/s
+Validation at step 38750: Loss = 2.9063, Perplexity = 18.29
+Step 38800/44000 | Loss: 2.8957 | LR: 1.76e-05 | 980.9ms/step | 20.4 samples/s
+Step 38900/44000 | Loss: 2.8822 | LR: 1.70e-05 | 954.5ms/step | 21.0 samples/s
+Step 39000/44000 | Loss: 2.8864 | LR: 1.63e-05 | 954.7ms/step | 20.9 samples/s
+Validation at step 39000: Loss = 2.9062, Perplexity = 18.29
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-39000
+Step 39100/44000 | Loss: 2.9083 | LR: 1.57e-05 | 1023.8ms/step | 19.5 samples/s
+Step 39200/44000 | Loss: 2.9020 | LR: 1.50e-05 | 955.5ms/step | 20.9 samples/s
+Validation at step 39250: Loss = 2.9062, Perplexity = 18.29
+Step 39300/44000 | Loss: 2.8970 | LR: 1.44e-05 | 982.3ms/step | 20.4 samples/s
+Step 39400/44000 | Loss: 2.9211 | LR: 1.38e-05 | 956.2ms/step | 20.9 samples/s
+Step 39500/44000 | Loss: 2.8937 | LR: 1.32e-05 | 956.3ms/step | 20.9 samples/s
+Validation at step 39500: Loss = 2.9061, Perplexity = 18.29
+Step 39600/44000 | Loss: 2.8850 | LR: 1.27e-05 | 982.3ms/step | 20.4 samples/s
+Step 39700/44000 | Loss: 2.9150 | LR: 1.21e-05 | 956.3ms/step | 20.9 samples/s
+Validation at step 39750: Loss = 2.9061, Perplexity = 18.29
+Step 39800/44000 | Loss: 2.8888 | LR: 1.15e-05 | 982.5ms/step | 20.4 samples/s
+Step 39900/44000 | Loss: 2.9141 | LR: 1.10e-05 | 956.1ms/step | 20.9 samples/s
+Step 40000/44000 | Loss: 2.9116 | LR: 1.05e-05 | 956.5ms/step | 20.9 samples/s
+Validation at step 40000: Loss = 2.9061, Perplexity = 18.29
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-40000
+Step 40100/44000 | Loss: 2.9035 | LR: 9.96e-06 | 1024.0ms/step | 19.5 samples/s
+Step 40200/44000 | Loss: 2.9151 | LR: 9.46e-06 | 956.2ms/step | 20.9 samples/s
+Validation at step 40250: Loss = 2.9061, Perplexity = 18.28
+Step 40300/44000 | Loss: 2.9058 | LR: 8.97e-06 | 982.4ms/step | 20.4 samples/s
+Step 40400/44000 | Loss: 2.8916 | LR: 8.50e-06 | 956.4ms/step | 20.9 samples/s
+Step 40500/44000 | Loss: 2.8939 | LR: 8.04e-06 | 956.0ms/step | 20.9 samples/s
+Validation at step 40500: Loss = 2.9060, Perplexity = 18.28
+Step 40600/44000 | Loss: 2.9029 | LR: 7.59e-06 | 982.4ms/step | 20.4 samples/s
+Step 40700/44000 | Loss: 2.8945 | LR: 7.15e-06 | 956.4ms/step | 20.9 samples/s
+Validation at step 40750: Loss = 2.9060, Perplexity = 18.28
+Step 40800/44000 | Loss: 2.8922 | LR: 6.72e-06 | 982.3ms/step | 20.4 samples/s
+Step 40900/44000 | Loss: 2.8892 | LR: 6.31e-06 | 956.3ms/step | 20.9 samples/s
+Step 41000/44000 | Loss: 2.8947 | LR: 5.91e-06 | 956.3ms/step | 20.9 samples/s
+Validation at step 41000: Loss = 2.9060, Perplexity = 18.28
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-41000
+Step 41100/44000 | Loss: 2.8857 | LR: 5.53e-06 | 1024.5ms/step | 19.5 samples/s
+Step 41200/44000 | Loss: 2.9097 | LR: 5.15e-06 | 956.4ms/step | 20.9 samples/s
+Validation at step 41250: Loss = 2.9060, Perplexity = 18.28
+Step 41300/44000 | Loss: 2.8894 | LR: 4.79e-06 | 982.5ms/step | 20.4 samples/s
+Step 41400/44000 | Loss: 2.8933 | LR: 4.45e-06 | 956.2ms/step | 20.9 samples/s
+Step 41500/44000 | Loss: 2.9148 | LR: 4.11e-06 | 956.2ms/step | 20.9 samples/s
+Validation at step 41500: Loss = 2.9060, Perplexity = 18.28
+Step 41600/44000 | Loss: 2.9025 | LR: 3.79e-06 | 982.8ms/step | 20.3 samples/s
+Step 41700/44000 | Loss: 2.9438 | LR: 3.48e-06 | 956.0ms/step | 20.9 samples/s
+Validation at step 41750: Loss = 2.9060, Perplexity = 18.28
+Step 41800/44000 | Loss: 2.9042 | LR: 3.19e-06 | 982.6ms/step | 20.4 samples/s
+Step 41900/44000 | Loss: 2.9083 | LR: 2.90e-06 | 956.1ms/step | 20.9 samples/s
+Step 42000/44000 | Loss: 2.9005 | LR: 2.63e-06 | 955.9ms/step | 20.9 samples/s
+Validation at step 42000: Loss = 2.9060, Perplexity = 18.28
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-42000
+Step 42100/44000 | Loss: 2.9003 | LR: 2.38e-06 | 1024.6ms/step | 19.5 samples/s
+Step 42200/44000 | Loss: 2.8982 | LR: 2.13e-06 | 956.4ms/step | 20.9 samples/s
+Validation at step 42250: Loss = 2.9060, Perplexity = 18.28
+Step 42300/44000 | Loss: 2.8951 | LR: 1.90e-06 | 982.8ms/step | 20.3 samples/s
+Step 42400/44000 | Loss: 2.9073 | LR: 1.69e-06 | 956.2ms/step | 20.9 samples/s
+Step 42500/44000 | Loss: 2.9089 | LR: 1.48e-06 | 956.0ms/step | 20.9 samples/s
+Validation at step 42500: Loss = 2.9060, Perplexity = 18.28
+Step 42600/44000 | Loss: 2.9029 | LR: 1.29e-06 | 982.5ms/step | 20.4 samples/s
+Step 42700/44000 | Loss: 2.8938 | LR: 1.11e-06 | 956.1ms/step | 20.9 samples/s
+Validation at step 42750: Loss = 2.9060, Perplexity = 18.28
+Step 42800/44000 | Loss: 2.8994 | LR: 9.49e-07 | 982.7ms/step | 20.4 samples/s
+Step 42900/44000 | Loss: 2.8944 | LR: 7.98e-07 | 956.2ms/step | 20.9 samples/s
+Step 43000/44000 | Loss: 2.9136 | LR: 6.59e-07 | 956.1ms/step | 20.9 samples/s
+Validation at step 43000: Loss = 2.9060, Perplexity = 18.28
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-43000
+Step 43100/44000 | Loss: 2.8979 | LR: 5.34e-07 | 1025.1ms/step | 19.5 samples/s
+Step 43200/44000 | Loss: 2.8984 | LR: 4.22e-07 | 956.3ms/step | 20.9 samples/s
+Validation at step 43250: Loss = 2.9060, Perplexity = 18.28
+Step 43300/44000 | Loss: 2.9087 | LR: 3.23e-07 | 982.4ms/step | 20.4 samples/s
+Step 43400/44000 | Loss: 2.8762 | LR: 2.37e-07 | 956.2ms/step | 20.9 samples/s
+Step 43500/44000 | Loss: 2.9012 | LR: 1.65e-07 | 956.4ms/step | 20.9 samples/s
+Validation at step 43500: Loss = 2.9060, Perplexity = 18.28
+Step 43600/44000 | Loss: 2.8828 | LR: 1.06e-07 | 982.4ms/step | 20.4 samples/s
+Step 43700/44000 | Loss: 2.8933 | LR: 5.94e-08 | 956.4ms/step | 20.9 samples/s
+Validation at step 43750: Loss = 2.9060, Perplexity = 18.28
+Step 43800/44000 | Loss: 2.8995 | LR: 2.64e-08 | 982.6ms/step | 20.4 samples/s
+Step 43900/44000 | Loss: 2.9133 | LR: 6.60e-09 | 956.4ms/step | 20.9 samples/s
+Step 44000/44000 | Loss: 2.8973 | LR: 0.00e+00 | 956.3ms/step | 20.9 samples/s
+Validation at step 44000: Loss = 2.9060, Perplexity = 18.28
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-44000
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-44000
+wandb: 
+wandb: Run history:
+wandb: eval/dense_interval_active ▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁
+wandb:    eval/dense_interval_end ▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁
+wandb:  eval/dense_interval_start ▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁
+wandb:                  eval/loss █▇▇▇▆▅▅▄▄▄▃▃▃▃▃▃▃▂▂▂▂▂▂▂▂▂▁▁▁▁▁▁▁▁▁▁▁▁▁▁
+wandb:            eval/perplexity █▅▅▄▄▄▃▃▃▃▂▂▂▂▂▂▂▂▂▂▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁
+wandb:  eval/uses_dense_attention ▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁
+wandb: eval/uses_sparse_attention ▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁
+wandb:          speed/ms_per_step ▁▁▁▁▄▅▁▁█▁▁▁█▁▁▁▁█▁▁▁▁▄█▁▁█▁▁▁▁▁▁▁█▄▁▄▁▁
+wandb:      speed/samples_per_sec █▅█▅██▇█▁▇▅████▅▅███▅██▁▅██▁██████▅▅██▅█
+wandb:        speed/steps_per_sec ▅█▁█▅▅█▅▅█▅█▅████▅██▅▁████▅█████▅▅█▅████
+wandb:                         +4 ...
+wandb: 
+wandb: Run summary:
+wandb: eval/dense_interval_active 0
+wandb:    eval/dense_interval_end -1
+wandb:  eval/dense_interval_start -1
+wandb:                  eval/loss 2.90601
+wandb:            eval/perplexity 18.28364
+wandb:  eval/uses_dense_attention 0
+wandb: eval/uses_sparse_attention 1
+wandb:          speed/ms_per_step 956.27475
+wandb:      speed/samples_per_sec 20.91449
+wandb:        speed/steps_per_sec 1.04572
+wandb:                         +4 ...
+wandb: 
+wandb: You can sync this run to the cloud by running:
+wandb: wandb sync /mloscratch/homes/navasard/sparse/SparseAtt/wandb/offline-run-20260921_164537-zgrsvkbi
+wandb: Find logs at: ./wandb/offline-run-20260921_164537-zgrsvkbi/logs
+Training complete!

epfl_final/logs/ALIBI_20260921/ddp5_a10040_s45_20260921_205236/orchestrator.log

@@ -0,0 +1 @@
+bash: line 4: LOGDIR: unbound variable

epfl_final/logs/ALIBI_20260921/ddp5_a10040_s45_20260921_205236/seed45.log

@@ -0,0 +1,538 @@
+W0921 20:53:23.594000 83 site-packages/torch/distributed/run.py:851] 
+W0921 20:53:23.594000 83 site-packages/torch/distributed/run.py:851] *****************************************
+W0921 20:53:23.594000 83 site-packages/torch/distributed/run.py:851] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. 
+W0921 20:53:23.594000 83 site-packages/torch/distributed/run.py:851] *****************************************
+Auto-resuming from latest checkpoint: /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s45/run_20260921_142951/checkpoint-2000
+Resuming: Using existing output directory /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s45/run_20260921_142951
+Running config:
+{
+  "model_path": "/mloscratch/homes/navasard/qwen3_06b_base",
+  "model_dtype": "bfloat16",
+  "init_from_scratch": true,
+  "train_path": "/mloscratch/homes/navasard/data/fineweb_large/tokenized",
+  "val_path": null,
+  "max_seq_length": 8192,
+  "num_workers": 8,
+  "output_dir": "/mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s45/run_20260921_142951",
+  "resume_from_checkpoint": "/mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s45/run_20260921_142951/checkpoint-2000",
+  "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": 45,
+  "verbose": true,
+  "checkpoint_activations": false,
+  "resume_micro_step_override": 8000,
+  "strategy": "ddp",
+  "mixed_precision": "bf16",
+  "wandb_enabled": true,
+  "wandb_project": "stochastic_sparsity",
+  "wandb_name": "365m alibi 8k s45",
+  "wandb_tags": [
+    "qwen3",
+    "365m",
+    "8k",
+    "scaleladder",
+    "alibi",
+    "seed45"
+  ],
+  "hidden_size": 768,
+  "intermediate_size": 3072,
+  "num_attention_heads": 6,
+  "num_key_value_heads": 2,
+  "num_hidden_layers": 28,
+  "head_dim": null,
+  "attention_dropout": null,
+  "rope_train_position_scale_set": null,
+  "positional": "alibi",
+  "attention_type": "dense",
+  "sparse_block_size": 128,
+  "sparse_num_random_blocks": 31,
+  "sparse_mask_pool_size": 1,
+  "sparse_val_dense_compare": false,
+  "sparse_layer_indices": null,
+  "sparse_mask_cache_dir": null,
+  "sparse_mask_seed": 0,
+  "sparse_mixed_nrb_set": null,
+  "dense_mode": false,
+  "dense_step_intervals": null
+}
+Loading tokenizer from /mloscratch/homes/navasard/qwen3_06b_base...
+Resuming from checkpoint: /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s45/run_20260921_142951/checkpoint-2000
+Checkpoint activations: FalseCheckpoint activations: False
+Checkpoint activations: False
+Checkpoint activations: False
+
+Checkpoint activations: False
+
+
+================================================================================
+MODEL ARCHITECTURE
+================================================================================
+Model config:
+  vocab_size: 151936
+  hidden_size: 768
+  num_hidden_layers: 28
+  num_attention_heads: 6
+  num_key_value_heads: 2
+  head_dim: 128
+  intermediate_size: 3072
+  max_position_embeddings: 32768
+  rope_theta: 1000000
+  rope_train_position_scale_set: None
+  attention_type: dense
+
+Parameter count:
+  Total parameters: 358,958,848
+  Trainable parameters: 358,958,848
+  Non-trainable parameters: 0
+  Model size (BF16): 0.72 GB
+
+--------------------------------------------------------------------------------
+LAYER-BY-LAYER SUMMARY
+--------------------------------------------------------------------------------
+/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/nn/attention/flex_attention.py:1704: UserWarning: flex_attention called without torch.compile() - this will use an unfused implementation that materializes the full scores matrix instead of generating a fused kernel.
+
+SOLUTION: Use torch.compile(flex_attention)(...)
+
+If you want to debug your score_mod/mask_mod, you can set:
+torch.nn.attention.flex_attention._FLEX_ATTENTION_DISABLE_COMPILE_DEBUG = True
+
+This will allow you to use print statements or breakpoints. Note: This doesn't work with the backwards pass and may produce incorrect results.
+  _warn_once(
+=================================================================================================================================================
+Layer (type:depth-idx)                        Input Shape               Output Shape              Param #                   Trainable
+=================================================================================================================================================
+Qwen3ForCausalLM                              [1, 8192]                 --                        --                        True
+├─Qwen3Model: 1-1                             --                        [1, 8192, 768]            --                        True
+│    └─Embedding: 2-1                         [1, 8192]                 [1, 8192, 768]            116,686,848               True
+│    └─ModuleList: 2-2                        --                        --                        --                        True
+│    │    └─Qwen3DecoderLayer: 3-1            [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-2            [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-3            [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-4            [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-5            [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-6            [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-7            [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-8            [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-9            [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-10           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-11           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-12           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-13           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-14           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-15           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-16           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-17           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-18           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-19           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-20           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-21           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-22           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-23           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-24           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-25           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-26           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-27           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-28           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    └─Qwen3RMSNorm: 2-3                      [1, 8192, 768]            [1, 8192, 768]            768                       True
+├─Linear: 1-2                                 [1, 8192, 768]            [1, 8192, 151936]         116,686,848               True
+=================================================================================================================================================
+Total params: 475,645,696
+Trainable params: 475,645,696
+Non-trainable params: 0
+Total mult-adds (Units.MEGABYTES): 475.65
+=================================================================================================================================================
+Input size (MB): 0.07
+Forward/backward pass size (MB): 15598.62
+Params size (MB): 951.29
+Estimated Total Size (MB): 16549.97
+=================================================================================================================================================
+  NOTE: torchinfo double-counts tied embed_tokens/lm_head weight (116,686,848 params). True unique parameter count: 358,958,848 (359.0M)
+================================================================================
+
+Loading dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Using pre-tokenized dataset (no runtime tokenization)
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Loaded 12748481 pre-tokenized examples (stored_len=2048, pack_factor=4, packed_examples=3187120)
+Loaded 12748481 pre-tokenized examples (stored_len=2048, pack_factor=4, packed_examples=3187120)
+Loaded 12748481 pre-tokenized examples (stored_len=2048, pack_factor=4, packed_examples=3187120)
+Loaded 12748481 pre-tokenized examples (stored_len=2048, pack_factor=4, packed_examples=3187120)
+Loaded 12748481 pre-tokenized examples (stored_len=2048, pack_factor=4, packed_examples=3187120)
+Dataset split: 3027764 train, 159356 validation
+Training examples (packed, train split):  3027764
+Validation examples (packed, val split):  159356
+Per-device batch size:                    1
+World size (GPUs):                        5
+Gradient accumulation steps:              4
+Effective batch size (seqs/step):         20
+Sequence length:                          8192
+Tokens per optimizer step:                163,840
+Micro-batches per epoch (per rank):       605553
+Optimizer steps per epoch:                151388
+Total optimizer steps:                    44000  (capped by max_steps; 0.29 epochs)
+Loading checkpoint from /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s45/run_20260921_142951/checkpoint-2000...
+Preserving checkpoint LR state for exact resume: current=[0.0004989701850946613], schedule_base=[0.0005]
+Resumed from checkpoint at step 2000 (epoch 0, micro_step 20000)
+Overriding resume micro-step 20000 -> 8000 for DDP migration
+Resuming training from step 2000
+Resuming wandb run ghfbqjah
+wandb: WARNING `resume` will be ignored since W&B syncing is set to `offline`. Starting a new run with run id ghfbqjah.
+wandb: Tracking run with wandb version 0.26.1
+wandb: W&B syncing is set to `offline` in this directory. Run `wandb online` or set WANDB_MODE=online to enable cloud syncing.
+wandb: Run data is saved locally in /mloscratch/homes/navasard/sparse/SparseAtt/wandb/offline-run-20260921_205453-ghfbqjah
+Starting training...
+Validation at step 2000: Loss = 3.5867, Perplexity = 36.11
+Switching attention mode for step 2001: sparse_stochastic (configured sparse attention)
+Step 2100/44000 | Loss: 3.6125 | LR: 4.99e-04 | 1895.2ms/step | 10.6 samples/s
+Step 2200/44000 | Loss: 3.6174 | LR: 4.99e-04 | 974.0ms/step | 20.5 samples/s
+Validation at step 2250: Loss = 3.5340, Perplexity = 34.26
+Step 2300/44000 | Loss: 3.5961 | LR: 4.98e-04 | 1003.3ms/step | 19.9 samples/s
+Step 2400/44000 | Loss: 3.5680 | LR: 4.98e-04 | 977.8ms/step | 20.5 samples/s
+Step 2500/44000 | Loss: 3.5607 | LR: 4.98e-04 | 978.0ms/step | 20.5 samples/s
+Validation at step 2500: Loss = 3.4869, Perplexity = 32.68
+Step 2600/44000 | Loss: 3.5426 | LR: 4.98e-04 | 1005.7ms/step | 19.9 samples/s
+Step 2700/44000 | Loss: 3.5307 | LR: 4.97e-04 | 978.9ms/step | 20.4 samples/s
+Validation at step 2750: Loss = 3.4438, Perplexity = 31.30
+Step 2800/44000 | Loss: 3.5142 | LR: 4.97e-04 | 1006.4ms/step | 19.9 samples/s
+Step 2900/44000 | Loss: 3.4805 | LR: 4.97e-04 | 980.1ms/step | 20.4 samples/s
+Step 3000/44000 | Loss: 3.4992 | LR: 4.97e-04 | 980.0ms/step | 20.4 samples/s
+Validation at step 3000: Loss = 3.4117, Perplexity = 30.32
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s45/run_20260921_142951/checkpoint-3000
+Step 3100/44000 | Loss: 3.4626 | LR: 4.96e-04 | 1046.6ms/step | 19.1 samples/s
+Step 3200/44000 | Loss: 3.4497 | LR: 4.96e-04 | 977.8ms/step | 20.5 samples/s
+Validation at step 3250: Loss = 3.3805, Perplexity = 29.38
+Step 3300/44000 | Loss: 3.4570 | LR: 4.96e-04 | 1004.4ms/step | 19.9 samples/s
+Step 3400/44000 | Loss: 3.4613 | LR: 4.95e-04 | 978.0ms/step | 20.4 samples/s
+Step 3500/44000 | Loss: 3.4060 | LR: 4.95e-04 | 978.2ms/step | 20.4 samples/s
+Validation at step 3500: Loss = 3.3584, Perplexity = 28.74
+Step 3600/44000 | Loss: 3.4011 | LR: 4.95e-04 | 1005.0ms/step | 19.9 samples/s
+Step 3700/44000 | Loss: 3.3820 | LR: 4.94e-04 | 978.0ms/step | 20.4 samples/s
+Validation at step 3750: Loss = 3.3301, Perplexity = 27.94
+Step 3800/44000 | Loss: 3.3811 | LR: 4.94e-04 | 1005.0ms/step | 19.9 samples/s
+Step 3900/44000 | Loss: 3.3672 | LR: 4.93e-04 | 978.3ms/step | 20.4 samples/s
+Step 4000/44000 | Loss: 3.3661 | LR: 4.93e-04 | 978.8ms/step | 20.4 samples/s
+Validation at step 4000: Loss = 3.3109, Perplexity = 27.41
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s45/run_20260921_142951/checkpoint-4000
+Step 4100/44000 | Loss: 3.3793 | LR: 4.93e-04 | 1047.2ms/step | 19.1 samples/s
+Step 4200/44000 | Loss: 3.3607 | LR: 4.92e-04 | 977.6ms/step | 20.5 samples/s
+Validation at step 4250: Loss = 3.2955, Perplexity = 26.99
+Step 4300/44000 | Loss: 3.3551 | LR: 4.92e-04 | 1004.7ms/step | 19.9 samples/s
+Step 4400/44000 | Loss: 3.3476 | LR: 4.91e-04 | 978.4ms/step | 20.4 samples/s
+Step 4500/44000 | Loss: 3.3333 | LR: 4.91e-04 | 978.4ms/step | 20.4 samples/s
+Validation at step 4500: Loss = 3.2713, Perplexity = 26.34
+Step 4600/44000 | Loss: 3.3427 | LR: 4.90e-04 | 1005.6ms/step | 19.9 samples/s
+Step 4700/44000 | Loss: 3.3218 | LR: 4.90e-04 | 978.1ms/step | 20.4 samples/s
+Validation at step 4750: Loss = 3.2576, Perplexity = 25.99
+Step 4800/44000 | Loss: 3.3376 | LR: 4.89e-04 | 1004.7ms/step | 19.9 samples/s
+Step 4900/44000 | Loss: 3.3232 | LR: 4.89e-04 | 978.0ms/step | 20.5 samples/s
+Step 5000/44000 | Loss: 3.3035 | LR: 4.88e-04 | 978.0ms/step | 20.5 samples/s
+Validation at step 5000: Loss = 3.2434, Perplexity = 25.62
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s45/run_20260921_142951/checkpoint-5000
+Step 5100/44000 | Loss: 3.2829 | LR: 4.88e-04 | 1046.4ms/step | 19.1 samples/s
+Step 5200/44000 | Loss: 3.2933 | LR: 4.87e-04 | 977.8ms/step | 20.5 samples/s
+Validation at step 5250: Loss = 3.2323, Perplexity = 25.34
+Step 5300/44000 | Loss: 3.2979 | LR: 4.86e-04 | 1004.7ms/step | 19.9 samples/s
+Step 5400/44000 | Loss: 3.2930 | LR: 4.86e-04 | 977.2ms/step | 20.5 samples/s
+Step 5500/44000 | Loss: 3.2855 | LR: 4.85e-04 | 977.9ms/step | 20.5 samples/s
+Validation at step 5500: Loss = 3.2179, Perplexity = 24.97
+Step 5600/44000 | Loss: 3.2644 | LR: 4.85e-04 | 1004.2ms/step | 19.9 samples/s
+Step 5700/44000 | Loss: 3.2746 | LR: 4.84e-04 | 977.7ms/step | 20.5 samples/s
+Validation at step 5750: Loss = 3.2050, Perplexity = 24.66
+Step 5800/44000 | Loss: 3.2552 | LR: 4.83e-04 | 1004.9ms/step | 19.9 samples/s
+Step 5900/44000 | Loss: 3.2679 | LR: 4.83e-04 | 978.0ms/step | 20.4 samples/s
+Step 6000/44000 | Loss: 3.2576 | LR: 4.82e-04 | 977.6ms/step | 20.5 samples/s
+Validation at step 6000: Loss = 3.1939, Perplexity = 24.38
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s45/run_20260921_142951/checkpoint-6000
+Step 6100/44000 | Loss: 3.2653 | LR: 4.81e-04 | 1046.0ms/step | 19.1 samples/s
+Step 6200/44000 | Loss: 3.2610 | LR: 4.81e-04 | 978.1ms/step | 20.4 samples/s
+Validation at step 6250: Loss = 3.1841, Perplexity = 24.15
+Step 6300/44000 | Loss: 3.2556 | LR: 4.80e-04 | 1004.7ms/step | 19.9 samples/s
+Step 6400/44000 | Loss: 3.2483 | LR: 4.79e-04 | 977.8ms/step | 20.5 samples/s
+Step 6500/44000 | Loss: 3.2267 | LR: 4.79e-04 | 978.0ms/step | 20.5 samples/s
+Validation at step 6500: Loss = 3.1729, Perplexity = 23.88
+Step 6600/44000 | Loss: 3.2366 | LR: 4.78e-04 | 1005.1ms/step | 19.9 samples/s
+Step 6700/44000 | Loss: 3.2353 | LR: 4.77e-04 | 978.1ms/step | 20.4 samples/s
+Validation at step 6750: Loss = 3.1666, Perplexity = 23.73
+Step 6800/44000 | Loss: 3.2341 | LR: 4.76e-04 | 1005.0ms/step | 19.9 samples/s
+Step 6900/44000 | Loss: 3.2123 | LR: 4.75e-04 | 977.9ms/step | 20.5 samples/s
+Step 7000/44000 | Loss: 3.2029 | LR: 4.75e-04 | 977.8ms/step | 20.5 samples/s
+Validation at step 7000: Loss = 3.1568, Perplexity = 23.50
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s45/run_20260921_142951/checkpoint-7000
+Step 7100/44000 | Loss: 3.2321 | LR: 4.74e-04 | 1045.8ms/step | 19.1 samples/s
+Step 7200/44000 | Loss: 3.1889 | LR: 4.73e-04 | 978.3ms/step | 20.4 samples/s
+Validation at step 7250: Loss = 3.1506, Perplexity = 23.35
+Step 7300/44000 | Loss: 3.1949 | LR: 4.72e-04 | 1005.3ms/step | 19.9 samples/s
+Step 7400/44000 | Loss: 3.1878 | LR: 4.71e-04 | 978.0ms/step | 20.4 samples/s
+Step 7500/44000 | Loss: 3.2185 | LR: 4.71e-04 | 978.5ms/step | 20.4 samples/s
+Validation at step 7500: Loss = 3.1435, Perplexity = 23.19
+Step 7600/44000 | Loss: 3.2002 | LR: 4.70e-04 | 1005.7ms/step | 19.9 samples/s
+Step 7700/44000 | Loss: 3.1809 | LR: 4.69e-04 | 978.7ms/step | 20.4 samples/s
+Validation at step 7750: Loss = 3.1344, Perplexity = 22.98
+Step 7800/44000 | Loss: 3.1655 | LR: 4.68e-04 | 1004.6ms/step | 19.9 samples/s
+Step 7900/44000 | Loss: 3.1969 | LR: 4.67e-04 | 978.2ms/step | 20.4 samples/s
+Step 8000/44000 | Loss: 3.1893 | LR: 4.66e-04 | 979.0ms/step | 20.4 samples/s
+Validation at step 8000: Loss = 3.1270, Perplexity = 22.81
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s45/run_20260921_142951/checkpoint-8000
+Step 8100/44000 | Loss: 3.1800 | LR: 4.65e-04 | 1046.3ms/step | 19.1 samples/s
+Step 8200/44000 | Loss: 3.2009 | LR: 4.64e-04 | 978.1ms/step | 20.4 samples/s
+Validation at step 8250: Loss = 3.1196, Perplexity = 22.64
+Step 8300/44000 | Loss: 3.1646 | LR: 4.63e-04 | 1005.0ms/step | 19.9 samples/s
+Step 8400/44000 | Loss: 3.1846 | LR: 4.62e-04 | 978.5ms/step | 20.4 samples/s
+Step 8500/44000 | Loss: 3.1805 | LR: 4.61e-04 | 978.4ms/step | 20.4 samples/s
+Validation at step 8500: Loss = 3.1088, Perplexity = 22.39
+Step 8600/44000 | Loss: 3.1711 | LR: 4.60e-04 | 1005.0ms/step | 19.9 samples/s
+Step 8700/44000 | Loss: 3.1907 | LR: 4.59e-04 | 978.8ms/step | 20.4 samples/s
+Validation at step 8750: Loss = 3.1007, Perplexity = 22.21
+Step 8800/44000 | Loss: 3.1662 | LR: 4.58e-04 | 1004.9ms/step | 19.9 samples/s
+Step 8900/44000 | Loss: 3.1539 | LR: 4.57e-04 | 978.0ms/step | 20.4 samples/s
+Step 9000/44000 | Loss: 3.1663 | LR: 4.56e-04 | 978.0ms/step | 20.5 samples/s
+Validation at step 9000: Loss = 3.0961, Perplexity = 22.11
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s45/run_20260921_142951/checkpoint-9000
+Step 9100/44000 | Loss: 3.1607 | LR: 4.55e-04 | 1046.5ms/step | 19.1 samples/s
+Step 9200/44000 | Loss: 3.1460 | LR: 4.54e-04 | 978.1ms/step | 20.4 samples/s
+Validation at step 9250: Loss = 3.0923, Perplexity = 22.03
+Step 9300/44000 | Loss: 3.1473 | LR: 4.53e-04 | 1005.1ms/step | 19.9 samples/s
+Step 9400/44000 | Loss: 3.1608 | LR: 4.52e-04 | 977.9ms/step | 20.5 samples/s
+Step 9500/44000 | Loss: 3.1467 | LR: 4.51e-04 | 977.7ms/step | 20.5 samples/s
+Validation at step 9500: Loss = 3.0849, Perplexity = 21.87
+Step 9600/44000 | Loss: 3.1350 | LR: 4.50e-04 | 1005.1ms/step | 19.9 samples/s
+Step 9700/44000 | Loss: 3.1496 | LR: 4.49e-04 | 977.8ms/step | 20.5 samples/s
+Validation at step 9750: Loss = 3.0798, Perplexity = 21.75
+Step 9800/44000 | Loss: 3.1330 | LR: 4.48e-04 | 1005.2ms/step | 19.9 samples/s
+Step 9900/44000 | Loss: 3.1449 | LR: 4.47e-04 | 978.3ms/step | 20.4 samples/s
+Step 10000/44000 | Loss: 3.1374 | LR: 4.46e-04 | 977.6ms/step | 20.5 samples/s
+Validation at step 10000: Loss = 3.0760, Perplexity = 21.67
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s45/run_20260921_142951/checkpoint-10000
+Step 10100/44000 | Loss: 3.1558 | LR: 4.45e-04 | 1046.0ms/step | 19.1 samples/s
+Step 10200/44000 | Loss: 3.1202 | LR: 4.43e-04 | 978.0ms/step | 20.5 samples/s
+Validation at step 10250: Loss = 3.0696, Perplexity = 21.53
+Step 10300/44000 | Loss: 3.1150 | LR: 4.42e-04 | 1005.0ms/step | 19.9 samples/s
+Step 10400/44000 | Loss: 3.1385 | LR: 4.41e-04 | 978.2ms/step | 20.4 samples/s
+Step 10500/44000 | Loss: 3.1108 | LR: 4.40e-04 | 978.1ms/step | 20.4 samples/s
+Validation at step 10500: Loss = 3.0669, Perplexity = 21.47
+Step 10600/44000 | Loss: 3.1104 | LR: 4.39e-04 | 1004.7ms/step | 19.9 samples/s
+Step 10700/44000 | Loss: 3.1087 | LR: 4.37e-04 | 978.0ms/step | 20.4 samples/s
+Validation at step 10750: Loss = 3.0610, Perplexity = 21.35
+Step 10800/44000 | Loss: 3.1119 | LR: 4.36e-04 | 1005.3ms/step | 19.9 samples/s
+Step 10900/44000 | Loss: 3.0995 | LR: 4.35e-04 | 977.6ms/step | 20.5 samples/s
+Step 11000/44000 | Loss: 3.1133 | LR: 4.34e-04 | 978.4ms/step | 20.4 samples/s
+Validation at step 11000: Loss = 3.0565, Perplexity = 21.25
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s45/run_20260921_142951/checkpoint-11000
+Step 11100/44000 | Loss: 3.0926 | LR: 4.33e-04 | 1046.8ms/step | 19.1 samples/s
+Step 11200/44000 | Loss: 3.1084 | LR: 4.31e-04 | 978.2ms/step | 20.4 samples/s
+Validation at step 11250: Loss = 3.0513, Perplexity = 21.14
+Step 11300/44000 | Loss: 3.0956 | LR: 4.30e-04 | 1005.1ms/step | 19.9 samples/s
+Step 11400/44000 | Loss: 3.1096 | LR: 4.29e-04 | 978.1ms/step | 20.4 samples/s
+Step 11500/44000 | Loss: 3.0854 | LR: 4.28e-04 | 978.4ms/step | 20.4 samples/s
+Validation at step 11500: Loss = 3.0469, Perplexity = 21.05
+Step 11600/44000 | Loss: 3.1177 | LR: 4.26e-04 | 1005.3ms/step | 19.9 samples/s
+Step 11700/44000 | Loss: 3.0957 | LR: 4.25e-04 | 978.4ms/step | 20.4 samples/s
+Validation at step 11750: Loss = 3.0448, Perplexity = 21.01
+Step 11800/44000 | Loss: 3.0877 | LR: 4.24e-04 | 1005.5ms/step | 19.9 samples/s
+Step 11900/44000 | Loss: 3.1049 | LR: 4.22e-04 | 978.5ms/step | 20.4 samples/s
+Step 12000/44000 | Loss: 3.0914 | LR: 4.21e-04 | 977.9ms/step | 20.5 samples/s
+Validation at step 12000: Loss = 3.0388, Perplexity = 20.88
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s45/run_20260921_142951/checkpoint-12000
+Step 12100/44000 | Loss: 3.0825 | LR: 4.20e-04 | 1045.5ms/step | 19.1 samples/s
+Step 12200/44000 | Loss: 3.0779 | LR: 4.18e-04 | 977.6ms/step | 20.5 samples/s
+Validation at step 12250: Loss = 3.0338, Perplexity = 20.78
+Step 12300/44000 | Loss: 3.1082 | LR: 4.17e-04 | 1004.6ms/step | 19.9 samples/s
+Step 12400/44000 | Loss: 3.0830 | LR: 4.16e-04 | 977.9ms/step | 20.5 samples/s
+Step 12500/44000 | Loss: 3.0684 | LR: 4.14e-04 | 978.0ms/step | 20.4 samples/s
+Validation at step 12500: Loss = 3.0300, Perplexity = 20.70
+Step 12600/44000 | Loss: 3.0815 | LR: 4.13e-04 | 1004.9ms/step | 19.9 samples/s
+Step 12700/44000 | Loss: 3.0857 | LR: 4.12e-04 | 978.1ms/step | 20.4 samples/s
+Validation at step 12750: Loss = 3.0270, Perplexity = 20.63
+Step 12800/44000 | Loss: 3.0511 | LR: 4.10e-04 | 1005.4ms/step | 19.9 samples/s
+Step 12900/44000 | Loss: 3.0802 | LR: 4.09e-04 | 977.9ms/step | 20.5 samples/s
+Step 13000/44000 | Loss: 3.0830 | LR: 4.07e-04 | 977.9ms/step | 20.5 samples/s
+Validation at step 13000: Loss = 3.0213, Perplexity = 20.52
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s45/run_20260921_142951/checkpoint-13000
+Step 13100/44000 | Loss: 3.0905 | LR: 4.06e-04 | 1045.1ms/step | 19.1 samples/s
+Step 13200/44000 | Loss: 3.0736 | LR: 4.05e-04 | 977.3ms/step | 20.5 samples/s
+Validation at step 13250: Loss = 3.0166, Perplexity = 20.42
+Step 13300/44000 | Loss: 3.0687 | LR: 4.03e-04 | 1004.5ms/step | 19.9 samples/s
+Step 13400/44000 | Loss: 3.0969 | LR: 4.02e-04 | 978.1ms/step | 20.4 samples/s
+Step 13500/44000 | Loss: 3.0804 | LR: 4.00e-04 | 977.7ms/step | 20.5 samples/s
+Validation at step 13500: Loss = 3.0146, Perplexity = 20.38
+Step 13600/44000 | Loss: 3.0835 | LR: 3.99e-04 | 1004.8ms/step | 19.9 samples/s
+Step 13700/44000 | Loss: 3.0728 | LR: 3.97e-04 | 977.6ms/step | 20.5 samples/s
+Validation at step 13750: Loss = 3.0119, Perplexity = 20.33
+Step 13800/44000 | Loss: 3.0642 | LR: 3.96e-04 | 1004.5ms/step | 19.9 samples/s
+Step 13900/44000 | Loss: 3.0632 | LR: 3.94e-04 | 977.6ms/step | 20.5 samples/s
+Step 14000/44000 | Loss: 3.0425 | LR: 3.93e-04 | 977.5ms/step | 20.5 samples/s
+Validation at step 14000: Loss = 3.0054, Perplexity = 20.19
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s45/run_20260921_142951/checkpoint-14000
+Step 14100/44000 | Loss: 3.0634 | LR: 3.91e-04 | 1046.4ms/step | 19.1 samples/s
+Step 14200/44000 | Loss: 3.0647 | LR: 3.90e-04 | 977.8ms/step | 20.5 samples/s
+Validation at step 14250: Loss = 3.0035, Perplexity = 20.16
+Step 14300/44000 | Loss: 3.0313 | LR: 3.88e-04 | 1004.7ms/step | 19.9 samples/s
+Step 14400/44000 | Loss: 3.0521 | LR: 3.87e-04 | 977.5ms/step | 20.5 samples/s
+Step 14500/44000 | Loss: 3.0604 | LR: 3.85e-04 | 977.8ms/step | 20.5 samples/s
+Validation at step 14500: Loss = 2.9993, Perplexity = 20.07
+Step 14600/44000 | Loss: 3.0691 | LR: 3.84e-04 | 1004.9ms/step | 19.9 samples/s
+Step 14700/44000 | Loss: 3.0365 | LR: 3.82e-04 | 977.8ms/step | 20.5 samples/s
+Validation at step 14750: Loss = 2.9954, Perplexity = 19.99
+Step 14800/44000 | Loss: 3.0538 | LR: 3.81e-04 | 1004.9ms/step | 19.9 samples/s
+Step 14900/44000 | Loss: 3.0338 | LR: 3.79e-04 | 978.4ms/step | 20.4 samples/s
+Step 15000/44000 | Loss: 3.0549 | LR: 3.78e-04 | 978.0ms/step | 20.4 samples/s
+Validation at step 15000: Loss = 2.9928, Perplexity = 19.94
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s45/run_20260921_142951/checkpoint-15000
+Step 15100/44000 | Loss: 3.0597 | LR: 3.76e-04 | 1046.6ms/step | 19.1 samples/s
+Step 15200/44000 | Loss: 3.0455 | LR: 3.74e-04 | 978.0ms/step | 20.5 samples/s
+Validation at step 15250: Loss = 2.9922, Perplexity = 19.93
+Step 15300/44000 | Loss: 3.0368 | LR: 3.73e-04 | 1004.8ms/step | 19.9 samples/s
+Step 15400/44000 | Loss: 3.0382 | LR: 3.71e-04 | 977.7ms/step | 20.5 samples/s
+Step 15500/44000 | Loss: 3.0468 | LR: 3.70e-04 | 978.2ms/step | 20.4 samples/s
+Validation at step 15500: Loss = 2.9885, Perplexity = 19.86
+Step 15600/44000 | Loss: 3.0455 | LR: 3.68e-04 | 1005.4ms/step | 19.9 samples/s
+Step 15700/44000 | Loss: 3.0284 | LR: 3.67e-04 | 977.7ms/step | 20.5 samples/s
+Validation at step 15750: Loss = 2.9837, Perplexity = 19.76
+Step 15800/44000 | Loss: 3.0415 | LR: 3.65e-04 | 1005.4ms/step | 19.9 samples/s
+Step 15900/44000 | Loss: 3.0516 | LR: 3.63e-04 | 978.2ms/step | 20.4 samples/s
+Step 16000/44000 | Loss: 3.0218 | LR: 3.62e-04 | 978.1ms/step | 20.4 samples/s
+Validation at step 16000: Loss = 2.9808, Perplexity = 19.70
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s45/run_20260921_142951/checkpoint-16000
+Step 16100/44000 | Loss: 3.0460 | LR: 3.60e-04 | 1047.4ms/step | 19.1 samples/s
+Step 16200/44000 | Loss: 3.0348 | LR: 3.58e-04 | 977.9ms/step | 20.5 samples/s
+Validation at step 16250: Loss = 2.9793, Perplexity = 19.67
+Step 16300/44000 | Loss: 3.0196 | LR: 3.57e-04 | 1005.4ms/step | 19.9 samples/s
+Step 16400/44000 | Loss: 3.0133 | LR: 3.55e-04 | 978.7ms/step | 20.4 samples/s
+Step 16500/44000 | Loss: 3.0230 | LR: 3.53e-04 | 978.6ms/step | 20.4 samples/s
+Validation at step 16500: Loss = 2.9772, Perplexity = 19.63
+Step 16600/44000 | Loss: 3.0100 | LR: 3.52e-04 | 1005.1ms/step | 19.9 samples/s
+Step 16700/44000 | Loss: 3.0215 | LR: 3.50e-04 | 978.3ms/step | 20.4 samples/s
+Validation at step 16750: Loss = 2.9745, Perplexity = 19.58
+Step 16800/44000 | Loss: 3.0353 | LR: 3.48e-04 | 1006.0ms/step | 19.9 samples/s
+Step 16900/44000 | Loss: 3.0345 | LR: 3.47e-04 | 978.6ms/step | 20.4 samples/s
+Step 17000/44000 | Loss: 3.0248 | LR: 3.45e-04 | 978.3ms/step | 20.4 samples/s
+Validation at step 17000: Loss = 2.9691, Perplexity = 19.47
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s45/run_20260921_142951/checkpoint-17000
+Step 17100/44000 | Loss: 3.0143 | LR: 3.43e-04 | 1046.3ms/step | 19.1 samples/s
+Step 17200/44000 | Loss: 3.0194 | LR: 3.42e-04 | 978.1ms/step | 20.4 samples/s
+Validation at step 17250: Loss = 2.9663, Perplexity = 19.42
+Step 17300/44000 | Loss: 3.0247 | LR: 3.40e-04 | 1006.0ms/step | 19.9 samples/s
+Step 17400/44000 | Loss: 2.9972 | LR: 3.38e-04 | 978.2ms/step | 20.4 samples/s
+Step 17500/44000 | Loss: 3.0160 | LR: 3.37e-04 | 978.2ms/step | 20.4 samples/s
+Validation at step 17500: Loss = 2.9639, Perplexity = 19.37
+Step 17600/44000 | Loss: 3.0132 | LR: 3.35e-04 | 1005.7ms/step | 19.9 samples/s
+Step 17700/44000 | Loss: 3.0188 | LR: 3.33e-04 | 978.6ms/step | 20.4 samples/s
+Validation at step 17750: Loss = 2.9610, Perplexity = 19.32
+Step 17800/44000 | Loss: 3.0168 | LR: 3.32e-04 | 1005.7ms/step | 19.9 samples/s
+Step 17900/44000 | Loss: 3.0116 | LR: 3.30e-04 | 978.5ms/step | 20.4 samples/s
+Step 18000/44000 | Loss: 3.0018 | LR: 3.28e-04 | 978.3ms/step | 20.4 samples/s
+Validation at step 18000: Loss = 2.9586, Perplexity = 19.27
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s45/run_20260921_142951/checkpoint-18000
+Step 18100/44000 | Loss: 3.0174 | LR: 3.26e-04 | 1047.8ms/step | 19.1 samples/s
+Step 18200/44000 | Loss: 3.0003 | LR: 3.25e-04 | 978.7ms/step | 20.4 samples/s
+Validation at step 18250: Loss = 2.9538, Perplexity = 19.18
+Step 18300/44000 | Loss: 3.0138 | LR: 3.23e-04 | 1005.2ms/step | 19.9 samples/s
+Step 18400/44000 | Loss: 2.9986 | LR: 3.21e-04 | 978.1ms/step | 20.4 samples/s
+Step 18500/44000 | Loss: 2.9832 | LR: 3.19e-04 | 978.4ms/step | 20.4 samples/s
+Validation at step 18500: Loss = 2.9510, Perplexity = 19.13
+Step 18600/44000 | Loss: 2.9914 | LR: 3.18e-04 | 1005.4ms/step | 19.9 samples/s
+Step 18700/44000 | Loss: 3.0116 | LR: 3.16e-04 | 978.2ms/step | 20.4 samples/s
+Validation at step 18750: Loss = 2.9497, Perplexity = 19.10
+Step 18800/44000 | Loss: 3.0067 | LR: 3.14e-04 | 1005.4ms/step | 19.9 samples/s
+Step 18900/44000 | Loss: 3.0077 | LR: 3.12e-04 | 978.2ms/step | 20.4 samples/s
+Step 19000/44000 | Loss: 3.0010 | LR: 3.11e-04 | 977.8ms/step | 20.5 samples/s
+Validation at step 19000: Loss = 2.9483, Perplexity = 19.07
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s45/run_20260921_142951/checkpoint-19000
+Step 19100/44000 | Loss: 2.9917 | LR: 3.09e-04 | 1047.3ms/step | 19.1 samples/s
+Step 19200/44000 | Loss: 3.0100 | LR: 3.07e-04 | 977.8ms/step | 20.5 samples/s
+Validation at step 19250: Loss = 2.9430, Perplexity = 18.97
+Step 19300/44000 | Loss: 3.0036 | LR: 3.05e-04 | 1004.9ms/step | 19.9 samples/s
+Step 19400/44000 | Loss: 2.9819 | LR: 3.04e-04 | 978.4ms/step | 20.4 samples/s
+Step 19500/44000 | Loss: 2.9944 | LR: 3.02e-04 | 977.9ms/step | 20.5 samples/s
+Validation at step 19500: Loss = 2.9409, Perplexity = 18.93
+Step 19600/44000 | Loss: 3.0054 | LR: 3.00e-04 | 1004.9ms/step | 19.9 samples/s
+Step 19700/44000 | Loss: 2.9904 | LR: 2.98e-04 | 977.8ms/step | 20.5 samples/s
+Validation at step 19750: Loss = 2.9369, Perplexity = 18.86
+Step 19800/44000 | Loss: 2.9894 | LR: 2.96e-04 | 1005.3ms/step | 19.9 samples/s
+Step 19900/44000 | Loss: 2.9729 | LR: 2.95e-04 | 978.2ms/step | 20.4 samples/s
+Step 20000/44000 | Loss: 2.9928 | LR: 2.93e-04 | 977.9ms/step | 20.5 samples/s
+Validation at step 20000: Loss = 2.9338, Perplexity = 18.80
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s45/run_20260921_142951/checkpoint-20000
+Step 20100/44000 | Loss: 2.9929 | LR: 2.91e-04 | 1047.9ms/step | 19.1 samples/s
+Step 20200/44000 | Loss: 2.9561 | LR: 2.89e-04 | 978.0ms/step | 20.4 samples/s
+Validation at step 20250: Loss = 2.9314, Perplexity = 18.75
+Step 20300/44000 | Loss: 2.9822 | LR: 2.88e-04 | 1005.0ms/step | 19.9 samples/s
+Step 20400/44000 | Loss: 3.0020 | LR: 2.86e-04 | 978.0ms/step | 20.4 samples/s
+Step 20500/44000 | Loss: 2.9735 | LR: 2.84e-04 | 978.6ms/step | 20.4 samples/s
+Validation at step 20500: Loss = 2.9291, Perplexity = 18.71
+Step 20600/44000 | Loss: 2.9568 | LR: 2.82e-04 | 1005.3ms/step | 19.9 samples/s
+Step 20700/44000 | Loss: 2.9963 | LR: 2.80e-04 | 978.4ms/step | 20.4 samples/s
+Validation at step 20750: Loss = 2.9255, Perplexity = 18.64
+Step 20800/44000 | Loss: 2.9744 | LR: 2.79e-04 | 1005.2ms/step | 19.9 samples/s
+Step 20900/44000 | Loss: 2.9749 | LR: 2.77e-04 | 977.9ms/step | 20.5 samples/s
+Step 21000/44000 | Loss: 2.9638 | LR: 2.75e-04 | 978.6ms/step | 20.4 samples/s
+Validation at step 21000: Loss = 2.9250, Perplexity = 18.63
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s45/run_20260921_142951/checkpoint-21000
+Step 21100/44000 | Loss: 2.9764 | LR: 2.73e-04 | 1047.8ms/step | 19.1 samples/s
+Step 21200/44000 | Loss: 2.9672 | LR: 2.71e-04 | 978.5ms/step | 20.4 samples/s
+Validation at step 21250: Loss = 2.9212, Perplexity = 18.56
+Step 21300/44000 | Loss: 2.9832 | LR: 2.70e-04 | 1005.2ms/step | 19.9 samples/s
+Step 21400/44000 | Loss: 2.9527 | LR: 2.68e-04 | 978.3ms/step | 20.4 samples/s
+Step 21500/44000 | Loss: 2.9714 | LR: 2.66e-04 | 978.3ms/step | 20.4 samples/s
+Validation at step 21500: Loss = 2.9195, Perplexity = 18.53
+Step 21600/44000 | Loss: 2.9696 | LR: 2.64e-04 | 1005.5ms/step | 19.9 samples/s
+Step 21700/44000 | Loss: 2.9749 | LR: 2.62e-04 | 977.8ms/step | 20.5 samples/s
+Validation at step 21750: Loss = 2.9167, Perplexity = 18.48
+Step 21800/44000 | Loss: 2.9751 | LR: 2.60e-04 | 1005.3ms/step | 19.9 samples/s
+Step 21900/44000 | Loss: 2.9652 | LR: 2.59e-04 | 977.8ms/step | 20.5 samples/s
+Step 22000/44000 | Loss: 2.9470 | LR: 2.57e-04 | 978.2ms/step | 20.4 samples/s
+Validation at step 22000: Loss = 2.9160, Perplexity = 18.47
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s45/run_20260921_142951/checkpoint-22000
+W0922 02:28:53.845000 83 site-packages/torch/distributed/elastic/agent/server/api.py:739] Received 15 death signal, shutting down workers
+W0922 02:28:53.847000 83 site-packages/torch/distributed/elastic/multiprocessing/api.py:1012] Sending process 151 closing signal SIGTERM
+W0922 02:28:53.849000 83 site-packages/torch/distributed/elastic/multiprocessing/api.py:1012] Sending process 152 closing signal SIGTERM
+W0922 02:28:53.850000 83 site-packages/torch/distributed/elastic/multiprocessing/api.py:1012] Sending process 153 closing signal SIGTERM
+W0922 02:28:53.851000 83 site-packages/torch/distributed/elastic/multiprocessing/api.py:1012] Sending process 154 closing signal SIGTERM
+W0922 02:28:53.855000 83 site-packages/torch/distributed/elastic/multiprocessing/api.py:1012] Sending process 155 closing signal SIGTERM
+Traceback (most recent call last):
+  File "/mloscratch/homes/navasard/conda/bin/torchrun", line 6, in <module>
+    sys.exit(main())
+             ~~~~^^
+  File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 362, in wrapper
+    return f(*args, **kwargs)
+  File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/distributed/run.py", line 990, in main
+    run(args)
+    ~~~^^^^^^
+  File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/distributed/run.py", line 981, in run
+    elastic_launch(
+    ~~~~~~~~~~~~~~~
+        config=config,
+        ~~~~~~~~~~~~~~
+        entrypoint=cmd,
+        ~~~~~~~~~~~~~~~
+    )(*cmd_args)
+    ~^^^^^^^^^^^
+  File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/distributed/launcher/api.py", line 170, in __call__
+    return launch_agent(self._config, self._entrypoint, list(args))
+  File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/distributed/launcher/api.py", line 308, in launch_agent
+    result = agent.run()
+  File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/distributed/elastic/metrics/api.py", line 134, in wrapper
+    result = f(*args, **kwargs)
+  File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/distributed/elastic/agent/server/api.py", line 731, in run
+    result = self._invoke_run(role)
+  File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/distributed/elastic/agent/server/api.py", line 908, in _invoke_run
+    time.sleep(monitor_interval)
+    ~~~~~~~~~~^^^^^^^^^^^^^^^^^^
+  File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/distributed/elastic/multiprocessing/api.py", line 86, in _terminate_process_handler
+    raise SignalException(f"Process {os.getpid()} got signal: {sigval}", sigval=sigval)
+torch.distributed.elastic.multiprocessing.api.SignalException: Process 83 got signal: 15

epfl_final/logs/ALIBI_20260921/ddp5_a10080_s44_20260921_164403/seed44.log

@@ -0,0 +1,862 @@
+W0921 16:44:15.919000 143 site-packages/torch/distributed/run.py:851] 
+W0921 16:44:15.919000 143 site-packages/torch/distributed/run.py:851] *****************************************
+W0921 16:44:15.919000 143 site-packages/torch/distributed/run.py:851] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. 
+W0921 16:44:15.919000 143 site-packages/torch/distributed/run.py:851] *****************************************
+Auto-resuming from latest checkpoint: /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-2000
+Resuming: Using existing output directory /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951
+Running config:
+{
+  "model_path": "/mloscratch/homes/navasard/qwen3_06b_base",
+  "model_dtype": "bfloat16",
+  "init_from_scratch": true,
+  "train_path": "/mloscratch/homes/navasard/data/fineweb_large/tokenized",
+  "val_path": null,
+  "max_seq_length": 8192,
+  "num_workers": 8,
+  "output_dir": "/mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951",
+  "resume_from_checkpoint": "/mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-2000",
+  "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": true,
+  "checkpoint_activations": false,
+  "resume_micro_step_override": 8000,
+  "strategy": "ddp",
+  "mixed_precision": "bf16",
+  "wandb_enabled": true,
+  "wandb_project": "stochastic_sparsity",
+  "wandb_name": "365m alibi 8k s44",
+  "wandb_tags": [
+    "qwen3",
+    "365m",
+    "8k",
+    "scaleladder",
+    "alibi",
+    "seed44"
+  ],
+  "hidden_size": 768,
+  "intermediate_size": 3072,
+  "num_attention_heads": 6,
+  "num_key_value_heads": 2,
+  "num_hidden_layers": 28,
+  "head_dim": null,
+  "attention_dropout": null,
+  "rope_train_position_scale_set": null,
+  "positional": "alibi",
+  "attention_type": "dense",
+  "sparse_block_size": 128,
+  "sparse_num_random_blocks": 31,
+  "sparse_mask_pool_size": 1,
+  "sparse_val_dense_compare": false,
+  "sparse_layer_indices": null,
+  "sparse_mask_cache_dir": null,
+  "sparse_mask_seed": 0,
+  "sparse_mixed_nrb_set": null,
+  "dense_mode": false,
+  "dense_step_intervals": null
+}
+Loading tokenizer from /mloscratch/homes/navasard/qwen3_06b_base...
+Checkpoint activations: FalseCheckpoint activations: False
+Checkpoint activations: False
+Checkpoint activations: False
+
+Resuming from checkpoint: /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-2000
+Checkpoint activations: False
+
+
+================================================================================
+MODEL ARCHITECTURE
+================================================================================
+Model config:
+  vocab_size: 151936
+  hidden_size: 768
+  num_hidden_layers: 28
+  num_attention_heads: 6
+  num_key_value_heads: 2
+  head_dim: 128
+  intermediate_size: 3072
+  max_position_embeddings: 32768
+  rope_theta: 1000000
+  rope_train_position_scale_set: None
+  attention_type: dense
+
+Parameter count:
+  Total parameters: 358,958,848
+  Trainable parameters: 358,958,848
+  Non-trainable parameters: 0
+  Model size (BF16): 0.72 GB
+
+--------------------------------------------------------------------------------
+LAYER-BY-LAYER SUMMARY
+--------------------------------------------------------------------------------
+/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/nn/attention/flex_attention.py:1704: UserWarning: flex_attention called without torch.compile() - this will use an unfused implementation that materializes the full scores matrix instead of generating a fused kernel.
+
+SOLUTION: Use torch.compile(flex_attention)(...)
+
+If you want to debug your score_mod/mask_mod, you can set:
+torch.nn.attention.flex_attention._FLEX_ATTENTION_DISABLE_COMPILE_DEBUG = True
+
+This will allow you to use print statements or breakpoints. Note: This doesn't work with the backwards pass and may produce incorrect results.
+  _warn_once(
+=================================================================================================================================================
+Layer (type:depth-idx)                        Input Shape               Output Shape              Param #                   Trainable
+=================================================================================================================================================
+Qwen3ForCausalLM                              [1, 8192]                 --                        --                        True
+├─Qwen3Model: 1-1                             --                        [1, 8192, 768]            --                        True
+│    └─Embedding: 2-1                         [1, 8192]                 [1, 8192, 768]            116,686,848               True
+│    └─ModuleList: 2-2                        --                        --                        --                        True
+│    │    └─Qwen3DecoderLayer: 3-1            [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-2            [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-3            [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-4            [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-5            [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-6            [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-7            [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-8            [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-9            [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-10           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-11           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-12           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-13           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-14           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-15           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-16           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-17           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-18           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-19           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-20           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-21           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-22           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-23           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-24           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-25           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-26           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-27           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-28           [1, 8192, 768]            [1, 8192, 768]            8,652,544                 True
+│    └─Qwen3RMSNorm: 2-3                      [1, 8192, 768]            [1, 8192, 768]            768                       True
+├─Linear: 1-2                                 [1, 8192, 768]            [1, 8192, 151936]         116,686,848               True
+=================================================================================================================================================
+Total params: 475,645,696
+Trainable params: 475,645,696
+Non-trainable params: 0
+Total mult-adds (Units.MEGABYTES): 475.65
+=================================================================================================================================================
+Input size (MB): 0.07
+Forward/backward pass size (MB): 15598.62
+Params size (MB): 951.29
+Estimated Total Size (MB): 16549.97
+=================================================================================================================================================
+  NOTE: torchinfo double-counts tied embed_tokens/lm_head weight (116,686,848 params). True unique parameter count: 358,958,848 (359.0M)
+================================================================================
+
+Loading dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Using pre-tokenized dataset (no runtime tokenization)
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Loaded 12748481 pre-tokenized examples (stored_len=2048, pack_factor=4, packed_examples=3187120)
+Loaded 12748481 pre-tokenized examples (stored_len=2048, pack_factor=4, packed_examples=3187120)
+Loaded 12748481 pre-tokenized examples (stored_len=2048, pack_factor=4, packed_examples=3187120)
+Loaded 12748481 pre-tokenized examples (stored_len=2048, pack_factor=4, packed_examples=3187120)
+Loaded 12748481 pre-tokenized examples (stored_len=2048, pack_factor=4, packed_examples=3187120)
+Dataset split: 3027764 train, 159356 validation
+Training examples (packed, train split):  3027764
+Validation examples (packed, val split):  159356
+Per-device batch size:                    1
+World size (GPUs):                        5
+Gradient accumulation steps:              4
+Effective batch size (seqs/step):         20
+Sequence length:                          8192
+Tokens per optimizer step:                163,840
+Micro-batches per epoch (per rank):       605553
+Optimizer steps per epoch:                151388
+Total optimizer steps:                    44000  (capped by max_steps; 0.29 epochs)
+Loading checkpoint from /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-2000...
+Preserving checkpoint LR state for exact resume: current=[0.0004989701850946613], schedule_base=[0.0005]
+Resumed from checkpoint at step 2000 (epoch 0, micro_step 20000)
+Overriding resume micro-step 20000 -> 8000 for DDP migration
+Resuming training from step 2000
+Resuming wandb run kbn66izc
+wandb: WARNING `resume` will be ignored since W&B syncing is set to `offline`. Starting a new run with run id kbn66izc.
+wandb: Tracking run with wandb version 0.26.1
+wandb: W&B syncing is set to `offline` in this directory. Run `wandb online` or set WANDB_MODE=online to enable cloud syncing.
+wandb: Run data is saved locally in /mloscratch/homes/navasard/sparse/SparseAtt/wandb/offline-run-20260921_164533-kbn66izc
+Starting training...
+Validation at step 2000: Loss = 3.6763, Perplexity = 39.50
+Switching attention mode for step 2001: sparse_stochastic (configured sparse attention)
+Step 2100/44000 | Loss: 3.6261 | LR: 4.99e-04 | 2585.1ms/step | 7.7 samples/s
+Step 2200/44000 | Loss: 3.6012 | LR: 4.99e-04 | 924.8ms/step | 21.6 samples/s
+Validation at step 2250: Loss = 3.6187, Perplexity = 37.29
+Step 2300/44000 | Loss: 3.5676 | LR: 4.98e-04 | 955.4ms/step | 20.9 samples/s
+Step 2400/44000 | Loss: 3.5955 | LR: 4.98e-04 | 930.2ms/step | 21.5 samples/s
+Step 2500/44000 | Loss: 3.5539 | LR: 4.98e-04 | 930.4ms/step | 21.5 samples/s
+Validation at step 2500: Loss = 3.5731, Perplexity = 35.63
+Step 2600/44000 | Loss: 3.5545 | LR: 4.98e-04 | 955.1ms/step | 20.9 samples/s
+Step 2700/44000 | Loss: 3.5484 | LR: 4.97e-04 | 929.0ms/step | 21.5 samples/s
+Validation at step 2750: Loss = 3.5317, Perplexity = 34.18
+Step 2800/44000 | Loss: 3.4846 | LR: 4.97e-04 | 955.7ms/step | 20.9 samples/s
+Step 2900/44000 | Loss: 3.4962 | LR: 4.97e-04 | 929.2ms/step | 21.5 samples/s
+Step 3000/44000 | Loss: 3.4766 | LR: 4.97e-04 | 929.3ms/step | 21.5 samples/s
+Validation at step 3000: Loss = 3.5015, Perplexity = 33.17
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-3000
+Step 3100/44000 | Loss: 3.4811 | LR: 4.96e-04 | 988.3ms/step | 20.2 samples/s
+Step 3200/44000 | Loss: 3.4465 | LR: 4.96e-04 | 929.3ms/step | 21.5 samples/s
+Validation at step 3250: Loss = 3.4699, Perplexity = 32.13
+Step 3300/44000 | Loss: 3.4349 | LR: 4.96e-04 | 954.6ms/step | 21.0 samples/s
+Step 3400/44000 | Loss: 3.4055 | LR: 4.95e-04 | 929.1ms/step | 21.5 samples/s
+Step 3500/44000 | Loss: 3.3978 | LR: 4.95e-04 | 929.6ms/step | 21.5 samples/s
+Validation at step 3500: Loss = 3.4436, Perplexity = 31.30
+Step 3600/44000 | Loss: 3.3898 | LR: 4.95e-04 | 955.0ms/step | 20.9 samples/s
+Step 3700/44000 | Loss: 3.3967 | LR: 4.94e-04 | 929.5ms/step | 21.5 samples/s
+Validation at step 3750: Loss = 3.4180, Perplexity = 30.51
+Step 3800/44000 | Loss: 3.3840 | LR: 4.94e-04 | 955.3ms/step | 20.9 samples/s
+Step 3900/44000 | Loss: 3.3804 | LR: 4.93e-04 | 929.1ms/step | 21.5 samples/s
+Step 4000/44000 | Loss: 3.3901 | LR: 4.93e-04 | 929.4ms/step | 21.5 samples/s
+Validation at step 4000: Loss = 3.4005, Perplexity = 29.98
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-4000
+Step 4100/44000 | Loss: 3.3621 | LR: 4.93e-04 | 988.6ms/step | 20.2 samples/s
+Step 4200/44000 | Loss: 3.3244 | LR: 4.92e-04 | 929.0ms/step | 21.5 samples/s
+Validation at step 4250: Loss = 3.3775, Perplexity = 29.30
+Step 4300/44000 | Loss: 3.3313 | LR: 4.92e-04 | 954.9ms/step | 20.9 samples/s
+Step 4400/44000 | Loss: 3.3437 | LR: 4.91e-04 | 929.0ms/step | 21.5 samples/s
+Step 4500/44000 | Loss: 3.3289 | LR: 4.91e-04 | 929.0ms/step | 21.5 samples/s
+Validation at step 4500: Loss = 3.3598, Perplexity = 28.78
+Step 4600/44000 | Loss: 3.3331 | LR: 4.90e-04 | 954.1ms/step | 21.0 samples/s
+Step 4700/44000 | Loss: 3.3081 | LR: 4.90e-04 | 928.2ms/step | 21.5 samples/s
+Validation at step 4750: Loss = 3.3423, Perplexity = 28.28
+Step 4800/44000 | Loss: 3.3183 | LR: 4.89e-04 | 955.1ms/step | 20.9 samples/s
+Step 4900/44000 | Loss: 3.3224 | LR: 4.89e-04 | 928.8ms/step | 21.5 samples/s
+Step 5000/44000 | Loss: 3.3070 | LR: 4.88e-04 | 928.9ms/step | 21.5 samples/s
+Validation at step 5000: Loss = 3.3296, Perplexity = 27.93
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-5000
+Step 5100/44000 | Loss: 3.3020 | LR: 4.88e-04 | 988.9ms/step | 20.2 samples/s
+Step 5200/44000 | Loss: 3.2865 | LR: 4.87e-04 | 929.7ms/step | 21.5 samples/s
+Validation at step 5250: Loss = 3.3136, Perplexity = 27.49
+Step 5300/44000 | Loss: 3.3092 | LR: 4.86e-04 | 954.8ms/step | 20.9 samples/s
+Step 5400/44000 | Loss: 3.2564 | LR: 4.86e-04 | 928.8ms/step | 21.5 samples/s
+Step 5500/44000 | Loss: 3.2838 | LR: 4.85e-04 | 929.2ms/step | 21.5 samples/s
+Validation at step 5500: Loss = 3.3036, Perplexity = 27.21
+Step 5600/44000 | Loss: 3.2705 | LR: 4.85e-04 | 955.4ms/step | 20.9 samples/s
+Step 5700/44000 | Loss: 3.2746 | LR: 4.84e-04 | 929.2ms/step | 21.5 samples/s
+Validation at step 5750: Loss = 3.2890, Perplexity = 26.82
+Step 5800/44000 | Loss: 3.2493 | LR: 4.83e-04 | 955.1ms/step | 20.9 samples/s
+Step 5900/44000 | Loss: 3.2661 | LR: 4.83e-04 | 929.1ms/step | 21.5 samples/s
+Step 6000/44000 | Loss: 3.2505 | LR: 4.82e-04 | 929.4ms/step | 21.5 samples/s
+Validation at step 6000: Loss = 3.2792, Perplexity = 26.55
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-6000
+Step 6100/44000 | Loss: 3.2413 | LR: 4.81e-04 | 988.6ms/step | 20.2 samples/s
+Step 6200/44000 | Loss: 3.2375 | LR: 4.81e-04 | 929.7ms/step | 21.5 samples/s
+Validation at step 6250: Loss = 3.2671, Perplexity = 26.24
+Step 6300/44000 | Loss: 3.2452 | LR: 4.80e-04 | 955.6ms/step | 20.9 samples/s
+Step 6400/44000 | Loss: 3.2548 | LR: 4.79e-04 | 930.2ms/step | 21.5 samples/s
+Step 6500/44000 | Loss: 3.2169 | LR: 4.79e-04 | 928.8ms/step | 21.5 samples/s
+Validation at step 6500: Loss = 3.2605, Perplexity = 26.06
+Step 6600/44000 | Loss: 3.2363 | LR: 4.78e-04 | 955.8ms/step | 20.9 samples/s
+Step 6700/44000 | Loss: 3.2210 | LR: 4.77e-04 | 929.7ms/step | 21.5 samples/s
+Validation at step 6750: Loss = 3.2483, Perplexity = 25.75
+Step 6800/44000 | Loss: 3.2389 | LR: 4.76e-04 | 955.5ms/step | 20.9 samples/s
+Step 6900/44000 | Loss: 3.1977 | LR: 4.75e-04 | 930.3ms/step | 21.5 samples/s
+Step 7000/44000 | Loss: 3.1962 | LR: 4.75e-04 | 929.8ms/step | 21.5 samples/s
+Validation at step 7000: Loss = 3.2388, Perplexity = 25.50
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-7000
+Step 7100/44000 | Loss: 3.2247 | LR: 4.74e-04 | 996.3ms/step | 20.1 samples/s
+Step 7200/44000 | Loss: 3.2017 | LR: 4.73e-04 | 929.2ms/step | 21.5 samples/s
+Validation at step 7250: Loss = 3.2322, Perplexity = 25.33
+Step 7300/44000 | Loss: 3.2054 | LR: 4.72e-04 | 954.7ms/step | 20.9 samples/s
+Step 7400/44000 | Loss: 3.2063 | LR: 4.71e-04 | 929.3ms/step | 21.5 samples/s
+Step 7500/44000 | Loss: 3.1921 | LR: 4.71e-04 | 929.1ms/step | 21.5 samples/s
+Validation at step 7500: Loss = 3.2204, Perplexity = 25.04
+Step 7600/44000 | Loss: 3.1927 | LR: 4.70e-04 | 955.2ms/step | 20.9 samples/s
+Step 7700/44000 | Loss: 3.1717 | LR: 4.69e-04 | 928.9ms/step | 21.5 samples/s
+Validation at step 7750: Loss = 3.2103, Perplexity = 24.79
+Step 7800/44000 | Loss: 3.1971 | LR: 4.68e-04 | 955.1ms/step | 20.9 samples/s
+Step 7900/44000 | Loss: 3.1891 | LR: 4.67e-04 | 928.8ms/step | 21.5 samples/s
+Step 8000/44000 | Loss: 3.1780 | LR: 4.66e-04 | 928.4ms/step | 21.5 samples/s
+Validation at step 8000: Loss = 3.2041, Perplexity = 24.63
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-8000
+Step 8100/44000 | Loss: 3.1894 | LR: 4.65e-04 | 986.4ms/step | 20.3 samples/s
+Step 8200/44000 | Loss: 3.1837 | LR: 4.64e-04 | 928.8ms/step | 21.5 samples/s
+Validation at step 8250: Loss = 3.1988, Perplexity = 24.50
+Step 8300/44000 | Loss: 3.1634 | LR: 4.63e-04 | 955.3ms/step | 20.9 samples/s
+Step 8400/44000 | Loss: 3.1796 | LR: 4.62e-04 | 930.0ms/step | 21.5 samples/s
+Step 8500/44000 | Loss: 3.1604 | LR: 4.61e-04 | 929.8ms/step | 21.5 samples/s
+Validation at step 8500: Loss = 3.1897, Perplexity = 24.28
+Step 8600/44000 | Loss: 3.1589 | LR: 4.60e-04 | 956.7ms/step | 20.9 samples/s
+Step 8700/44000 | Loss: 3.1841 | LR: 4.59e-04 | 930.0ms/step | 21.5 samples/s
+Validation at step 8750: Loss = 3.1844, Perplexity = 24.15
+Step 8800/44000 | Loss: 3.1538 | LR: 4.58e-04 | 955.3ms/step | 20.9 samples/s
+Step 8900/44000 | Loss: 3.1461 | LR: 4.57e-04 | 929.5ms/step | 21.5 samples/s
+Step 9000/44000 | Loss: 3.1615 | LR: 4.56e-04 | 929.2ms/step | 21.5 samples/s
+Validation at step 9000: Loss = 3.1799, Perplexity = 24.04
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-9000
+Step 9100/44000 | Loss: 3.1606 | LR: 4.55e-04 | 987.5ms/step | 20.3 samples/s
+Step 9200/44000 | Loss: 3.1368 | LR: 4.54e-04 | 929.5ms/step | 21.5 samples/s
+Validation at step 9250: Loss = 3.1725, Perplexity = 23.87
+Step 9300/44000 | Loss: 3.1662 | LR: 4.53e-04 | 957.2ms/step | 20.9 samples/s
+Step 9400/44000 | Loss: 3.1393 | LR: 4.52e-04 | 930.0ms/step | 21.5 samples/s
+Step 9500/44000 | Loss: 3.1583 | LR: 4.51e-04 | 930.2ms/step | 21.5 samples/s
+Validation at step 9500: Loss = 3.1688, Perplexity = 23.78
+Step 9600/44000 | Loss: 3.1543 | LR: 4.50e-04 | 954.7ms/step | 20.9 samples/s
+Step 9700/44000 | Loss: 3.1133 | LR: 4.49e-04 | 930.0ms/step | 21.5 samples/s
+Validation at step 9750: Loss = 3.1658, Perplexity = 23.71
+Step 9800/44000 | Loss: 3.1222 | LR: 4.48e-04 | 955.1ms/step | 20.9 samples/s
+Step 9900/44000 | Loss: 3.1334 | LR: 4.47e-04 | 928.9ms/step | 21.5 samples/s
+Step 10000/44000 | Loss: 3.1055 | LR: 4.46e-04 | 928.8ms/step | 21.5 samples/s
+Validation at step 10000: Loss = 3.1556, Perplexity = 23.47
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-10000
+Step 10100/44000 | Loss: 3.1251 | LR: 4.45e-04 | 988.1ms/step | 20.2 samples/s
+Step 10200/44000 | Loss: 3.1371 | LR: 4.43e-04 | 927.9ms/step | 21.6 samples/s
+Validation at step 10250: Loss = 3.1507, Perplexity = 23.35
+Step 10300/44000 | Loss: 3.1330 | LR: 4.42e-04 | 953.6ms/step | 21.0 samples/s
+Step 10400/44000 | Loss: 3.1409 | LR: 4.41e-04 | 928.9ms/step | 21.5 samples/s
+Step 10500/44000 | Loss: 3.1153 | LR: 4.40e-04 | 928.1ms/step | 21.5 samples/s
+Validation at step 10500: Loss = 3.1458, Perplexity = 23.24
+Step 10600/44000 | Loss: 3.1087 | LR: 4.39e-04 | 955.5ms/step | 20.9 samples/s
+Step 10700/44000 | Loss: 3.1440 | LR: 4.37e-04 | 928.6ms/step | 21.5 samples/s
+Validation at step 10750: Loss = 3.1422, Perplexity = 23.15
+Step 10800/44000 | Loss: 3.1086 | LR: 4.36e-04 | 955.0ms/step | 20.9 samples/s
+Step 10900/44000 | Loss: 3.1269 | LR: 4.35e-04 | 928.3ms/step | 21.5 samples/s
+Step 11000/44000 | Loss: 3.1198 | LR: 4.34e-04 | 929.1ms/step | 21.5 samples/s
+Validation at step 11000: Loss = 3.1373, Perplexity = 23.04
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-11000
+Step 11100/44000 | Loss: 3.1058 | LR: 4.33e-04 | 987.5ms/step | 20.3 samples/s
+Step 11200/44000 | Loss: 3.1282 | LR: 4.31e-04 | 928.2ms/step | 21.5 samples/s
+Validation at step 11250: Loss = 3.1342, Perplexity = 22.97
+Step 11300/44000 | Loss: 3.0840 | LR: 4.30e-04 | 954.4ms/step | 21.0 samples/s
+Step 11400/44000 | Loss: 3.1136 | LR: 4.29e-04 | 927.0ms/step | 21.6 samples/s
+Step 11500/44000 | Loss: 3.0744 | LR: 4.28e-04 | 928.5ms/step | 21.5 samples/s
+Validation at step 11500: Loss = 3.1224, Perplexity = 22.70
+Step 11600/44000 | Loss: 3.0897 | LR: 4.26e-04 | 955.2ms/step | 20.9 samples/s
+Step 11700/44000 | Loss: 3.0877 | LR: 4.25e-04 | 930.0ms/step | 21.5 samples/s
+Validation at step 11750: Loss = 3.1182, Perplexity = 22.61
+Step 11800/44000 | Loss: 3.0945 | LR: 4.24e-04 | 955.2ms/step | 20.9 samples/s
+Step 11900/44000 | Loss: 3.1167 | LR: 4.22e-04 | 929.4ms/step | 21.5 samples/s
+Step 12000/44000 | Loss: 3.0895 | LR: 4.21e-04 | 930.4ms/step | 21.5 samples/s
+Validation at step 12000: Loss = 3.1129, Perplexity = 22.49
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-12000
+Step 12100/44000 | Loss: 3.1150 | LR: 4.20e-04 | 986.6ms/step | 20.3 samples/s
+Step 12200/44000 | Loss: 3.0681 | LR: 4.18e-04 | 928.7ms/step | 21.5 samples/s
+Validation at step 12250: Loss = 3.1068, Perplexity = 22.35
+Step 12300/44000 | Loss: 3.0868 | LR: 4.17e-04 | 953.7ms/step | 21.0 samples/s
+Step 12400/44000 | Loss: 3.0925 | LR: 4.16e-04 | 927.9ms/step | 21.6 samples/s
+Step 12500/44000 | Loss: 3.0824 | LR: 4.14e-04 | 927.7ms/step | 21.6 samples/s
+Validation at step 12500: Loss = 3.1035, Perplexity = 22.28
+Step 12600/44000 | Loss: 3.0996 | LR: 4.13e-04 | 953.8ms/step | 21.0 samples/s
+Step 12700/44000 | Loss: 3.0903 | LR: 4.12e-04 | 928.8ms/step | 21.5 samples/s
+Validation at step 12750: Loss = 3.1025, Perplexity = 22.25
+Step 12800/44000 | Loss: 3.0857 | LR: 4.10e-04 | 952.5ms/step | 21.0 samples/s
+Step 12900/44000 | Loss: 3.0595 | LR: 4.09e-04 | 928.4ms/step | 21.5 samples/s
+Step 13000/44000 | Loss: 3.0728 | LR: 4.07e-04 | 928.2ms/step | 21.5 samples/s
+Validation at step 13000: Loss = 3.0982, Perplexity = 22.16
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-13000
+Step 13100/44000 | Loss: 3.0736 | LR: 4.06e-04 | 986.3ms/step | 20.3 samples/s
+Step 13200/44000 | Loss: 3.0718 | LR: 4.05e-04 | 928.6ms/step | 21.5 samples/s
+Validation at step 13250: Loss = 3.0961, Perplexity = 22.11
+Step 13300/44000 | Loss: 3.0789 | LR: 4.03e-04 | 953.9ms/step | 21.0 samples/s
+Step 13400/44000 | Loss: 3.0683 | LR: 4.02e-04 | 928.0ms/step | 21.6 samples/s
+Step 13500/44000 | Loss: 3.0720 | LR: 4.00e-04 | 928.7ms/step | 21.5 samples/s
+Validation at step 13500: Loss = 3.0912, Perplexity = 22.00
+Step 13600/44000 | Loss: 3.0731 | LR: 3.99e-04 | 953.5ms/step | 21.0 samples/s
+Step 13700/44000 | Loss: 3.0673 | LR: 3.97e-04 | 928.5ms/step | 21.5 samples/s
+Validation at step 13750: Loss = 3.0866, Perplexity = 21.90
+Step 13800/44000 | Loss: 3.0587 | LR: 3.96e-04 | 954.5ms/step | 21.0 samples/s
+Step 13900/44000 | Loss: 3.0682 | LR: 3.94e-04 | 929.5ms/step | 21.5 samples/s
+Step 14000/44000 | Loss: 3.0527 | LR: 3.93e-04 | 929.9ms/step | 21.5 samples/s
+Validation at step 14000: Loss = 3.0838, Perplexity = 21.84
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-14000
+Step 14100/44000 | Loss: 3.0447 | LR: 3.91e-04 | 987.3ms/step | 20.3 samples/s
+Step 14200/44000 | Loss: 3.0614 | LR: 3.90e-04 | 929.2ms/step | 21.5 samples/s
+Validation at step 14250: Loss = 3.0797, Perplexity = 21.75
+Step 14300/44000 | Loss: 3.0590 | LR: 3.88e-04 | 955.2ms/step | 20.9 samples/s
+Step 14400/44000 | Loss: 3.0597 | LR: 3.87e-04 | 929.9ms/step | 21.5 samples/s
+Step 14500/44000 | Loss: 3.0633 | LR: 3.85e-04 | 929.9ms/step | 21.5 samples/s
+Validation at step 14500: Loss = 3.0770, Perplexity = 21.69
+Step 14600/44000 | Loss: 3.0402 | LR: 3.84e-04 | 955.4ms/step | 20.9 samples/s
+Step 14700/44000 | Loss: 3.0519 | LR: 3.82e-04 | 930.8ms/step | 21.5 samples/s
+Validation at step 14750: Loss = 3.0739, Perplexity = 21.63
+Step 14800/44000 | Loss: 3.0364 | LR: 3.81e-04 | 955.1ms/step | 20.9 samples/s
+Step 14900/44000 | Loss: 3.0253 | LR: 3.79e-04 | 930.7ms/step | 21.5 samples/s
+Step 15000/44000 | Loss: 3.0373 | LR: 3.78e-04 | 930.2ms/step | 21.5 samples/s
+Validation at step 15000: Loss = 3.0708, Perplexity = 21.56
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-15000
+Step 15100/44000 | Loss: 3.0284 | LR: 3.76e-04 | 988.7ms/step | 20.2 samples/s
+Step 15200/44000 | Loss: 3.0392 | LR: 3.74e-04 | 931.7ms/step | 21.5 samples/s
+Validation at step 15250: Loss = 3.0662, Perplexity = 21.46
+Step 15300/44000 | Loss: 3.0568 | LR: 3.73e-04 | 956.7ms/step | 20.9 samples/s
+Step 15400/44000 | Loss: 3.0652 | LR: 3.71e-04 | 931.3ms/step | 21.5 samples/s
+Step 15500/44000 | Loss: 3.0406 | LR: 3.70e-04 | 930.8ms/step | 21.5 samples/s
+Validation at step 15500: Loss = 3.0652, Perplexity = 21.44
+Step 15600/44000 | Loss: 3.0318 | LR: 3.68e-04 | 956.8ms/step | 20.9 samples/s
+Step 15700/44000 | Loss: 3.0414 | LR: 3.67e-04 | 930.8ms/step | 21.5 samples/s
+Validation at step 15750: Loss = 3.0618, Perplexity = 21.36
+Step 15800/44000 | Loss: 3.0540 | LR: 3.65e-04 | 955.8ms/step | 20.9 samples/s
+Step 15900/44000 | Loss: 3.0344 | LR: 3.63e-04 | 930.1ms/step | 21.5 samples/s
+Step 16000/44000 | Loss: 3.0256 | LR: 3.62e-04 | 930.1ms/step | 21.5 samples/s
+Validation at step 16000: Loss = 3.0584, Perplexity = 21.29
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-16000
+Step 16100/44000 | Loss: 3.0493 | LR: 3.60e-04 | 988.7ms/step | 20.2 samples/s
+Step 16200/44000 | Loss: 3.0301 | LR: 3.58e-04 | 931.0ms/step | 21.5 samples/s
+Validation at step 16250: Loss = 3.0543, Perplexity = 21.21
+Step 16300/44000 | Loss: 3.0276 | LR: 3.57e-04 | 956.7ms/step | 20.9 samples/s
+Step 16400/44000 | Loss: 3.0300 | LR: 3.55e-04 | 929.6ms/step | 21.5 samples/s
+Step 16500/44000 | Loss: 3.0118 | LR: 3.53e-04 | 928.3ms/step | 21.5 samples/s
+Validation at step 16500: Loss = 3.0510, Perplexity = 21.14
+Step 16600/44000 | Loss: 3.0178 | LR: 3.52e-04 | 954.9ms/step | 20.9 samples/s
+Step 16700/44000 | Loss: 3.0351 | LR: 3.50e-04 | 930.2ms/step | 21.5 samples/s
+Validation at step 16750: Loss = 3.0492, Perplexity = 21.10
+Step 16800/44000 | Loss: 3.0145 | LR: 3.48e-04 | 954.3ms/step | 21.0 samples/s
+Step 16900/44000 | Loss: 3.0251 | LR: 3.47e-04 | 928.3ms/step | 21.5 samples/s
+Step 17000/44000 | Loss: 3.0175 | LR: 3.45e-04 | 929.4ms/step | 21.5 samples/s
+Validation at step 17000: Loss = 3.0459, Perplexity = 21.03
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-17000
+Step 17100/44000 | Loss: 3.0277 | LR: 3.43e-04 | 988.1ms/step | 20.2 samples/s
+Step 17200/44000 | Loss: 3.0112 | LR: 3.42e-04 | 929.6ms/step | 21.5 samples/s
+Validation at step 17250: Loss = 3.0428, Perplexity = 20.96
+Step 17300/44000 | Loss: 3.0242 | LR: 3.40e-04 | 954.9ms/step | 20.9 samples/s
+Step 17400/44000 | Loss: 3.0276 | LR: 3.38e-04 | 929.5ms/step | 21.5 samples/s
+Step 17500/44000 | Loss: 3.0284 | LR: 3.37e-04 | 929.6ms/step | 21.5 samples/s
+Validation at step 17500: Loss = 3.0410, Perplexity = 20.93
+Step 17600/44000 | Loss: 2.9899 | LR: 3.35e-04 | 956.1ms/step | 20.9 samples/s
+Step 17700/44000 | Loss: 3.0100 | LR: 3.33e-04 | 928.4ms/step | 21.5 samples/s
+Validation at step 17750: Loss = 3.0378, Perplexity = 20.86
+Step 17800/44000 | Loss: 3.0106 | LR: 3.32e-04 | 955.3ms/step | 20.9 samples/s
+Step 17900/44000 | Loss: 3.0014 | LR: 3.30e-04 | 931.8ms/step | 21.5 samples/s
+Step 18000/44000 | Loss: 3.0050 | LR: 3.28e-04 | 931.2ms/step | 21.5 samples/s
+Validation at step 18000: Loss = 3.0334, Perplexity = 20.77
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-18000
+Step 18100/44000 | Loss: 2.9994 | LR: 3.26e-04 | 997.1ms/step | 20.1 samples/s
+Step 18200/44000 | Loss: 3.0144 | LR: 3.25e-04 | 935.4ms/step | 21.4 samples/s
+Validation at step 18250: Loss = 3.0288, Perplexity = 20.67
+Step 18300/44000 | Loss: 3.0149 | LR: 3.23e-04 | 960.7ms/step | 20.8 samples/s
+Step 18400/44000 | Loss: 3.0064 | LR: 3.21e-04 | 931.5ms/step | 21.5 samples/s
+Step 18500/44000 | Loss: 2.9826 | LR: 3.19e-04 | 932.3ms/step | 21.5 samples/s
+Validation at step 18500: Loss = 3.0257, Perplexity = 20.61
+Step 18600/44000 | Loss: 3.0154 | LR: 3.18e-04 | 958.2ms/step | 20.9 samples/s
+Step 18700/44000 | Loss: 2.9939 | LR: 3.16e-04 | 932.8ms/step | 21.4 samples/s
+Validation at step 18750: Loss = 3.0231, Perplexity = 20.55
+Step 18800/44000 | Loss: 2.9815 | LR: 3.14e-04 | 958.4ms/step | 20.9 samples/s
+Step 18900/44000 | Loss: 2.9953 | LR: 3.12e-04 | 932.6ms/step | 21.4 samples/s
+Step 19000/44000 | Loss: 2.9955 | LR: 3.11e-04 | 931.9ms/step | 21.5 samples/s
+Validation at step 19000: Loss = 3.0204, Perplexity = 20.50
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-19000
+Step 19100/44000 | Loss: 3.0043 | LR: 3.09e-04 | 992.1ms/step | 20.2 samples/s
+Step 19200/44000 | Loss: 2.9754 | LR: 3.07e-04 | 930.8ms/step | 21.5 samples/s
+Validation at step 19250: Loss = 3.0182, Perplexity = 20.45
+Step 19300/44000 | Loss: 2.9713 | LR: 3.05e-04 | 956.2ms/step | 20.9 samples/s
+Step 19400/44000 | Loss: 2.9868 | LR: 3.04e-04 | 930.1ms/step | 21.5 samples/s
+Step 19500/44000 | Loss: 2.9612 | LR: 3.02e-04 | 930.0ms/step | 21.5 samples/s
+Validation at step 19500: Loss = 3.0143, Perplexity = 20.38
+Step 19600/44000 | Loss: 2.9974 | LR: 3.00e-04 | 957.0ms/step | 20.9 samples/s
+Step 19700/44000 | Loss: 2.9885 | LR: 2.98e-04 | 929.9ms/step | 21.5 samples/s
+Validation at step 19750: Loss = 3.0117, Perplexity = 20.32
+Step 19800/44000 | Loss: 3.0057 | LR: 2.96e-04 | 955.9ms/step | 20.9 samples/s
+Step 19900/44000 | Loss: 2.9779 | LR: 2.95e-04 | 929.6ms/step | 21.5 samples/s
+Step 20000/44000 | Loss: 3.0108 | LR: 2.93e-04 | 929.1ms/step | 21.5 samples/s
+Validation at step 20000: Loss = 3.0095, Perplexity = 20.28
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-20000
+Step 20100/44000 | Loss: 2.9604 | LR: 2.91e-04 | 989.8ms/step | 20.2 samples/s
+Step 20200/44000 | Loss: 2.9971 | LR: 2.89e-04 | 931.7ms/step | 21.5 samples/s
+Validation at step 20250: Loss = 3.0080, Perplexity = 20.25
+Step 20300/44000 | Loss: 2.9718 | LR: 2.88e-04 | 956.2ms/step | 20.9 samples/s
+Step 20400/44000 | Loss: 2.9780 | LR: 2.86e-04 | 930.0ms/step | 21.5 samples/s
+Step 20500/44000 | Loss: 2.9931 | LR: 2.84e-04 | 931.1ms/step | 21.5 samples/s
+Validation at step 20500: Loss = 3.0059, Perplexity = 20.20
+Step 20600/44000 | Loss: 3.0010 | LR: 2.82e-04 | 955.0ms/step | 20.9 samples/s
+Step 20700/44000 | Loss: 2.9754 | LR: 2.80e-04 | 929.1ms/step | 21.5 samples/s
+Validation at step 20750: Loss = 3.0035, Perplexity = 20.16
+Step 20800/44000 | Loss: 2.9792 | LR: 2.79e-04 | 954.5ms/step | 21.0 samples/s
+Step 20900/44000 | Loss: 2.9685 | LR: 2.77e-04 | 931.9ms/step | 21.5 samples/s
+Step 21000/44000 | Loss: 2.9522 | LR: 2.75e-04 | 929.5ms/step | 21.5 samples/s
+Validation at step 21000: Loss = 2.9995, Perplexity = 20.08
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-21000
+Step 21100/44000 | Loss: 2.9713 | LR: 2.73e-04 | 990.5ms/step | 20.2 samples/s
+Step 21200/44000 | Loss: 2.9918 | LR: 2.71e-04 | 930.9ms/step | 21.5 samples/s
+Validation at step 21250: Loss = 2.9983, Perplexity = 20.05
+Step 21300/44000 | Loss: 2.9915 | LR: 2.70e-04 | 955.9ms/step | 20.9 samples/s
+Step 21400/44000 | Loss: 2.9717 | LR: 2.68e-04 | 929.6ms/step | 21.5 samples/s
+Step 21500/44000 | Loss: 2.9679 | LR: 2.66e-04 | 930.5ms/step | 21.5 samples/s
+Validation at step 21500: Loss = 2.9954, Perplexity = 19.99
+Step 21600/44000 | Loss: 2.9760 | LR: 2.64e-04 | 955.6ms/step | 20.9 samples/s
+Step 21700/44000 | Loss: 2.9762 | LR: 2.62e-04 | 929.7ms/step | 21.5 samples/s
+Validation at step 21750: Loss = 2.9928, Perplexity = 19.94
+Step 21800/44000 | Loss: 2.9475 | LR: 2.60e-04 | 956.0ms/step | 20.9 samples/s
+Step 21900/44000 | Loss: 2.9720 | LR: 2.59e-04 | 929.9ms/step | 21.5 samples/s
+Step 22000/44000 | Loss: 2.9645 | LR: 2.57e-04 | 930.4ms/step | 21.5 samples/s
+Validation at step 22000: Loss = 2.9912, Perplexity = 19.91
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-22000
+Step 22100/44000 | Loss: 2.9659 | LR: 2.55e-04 | 990.3ms/step | 20.2 samples/s
+Step 22200/44000 | Loss: 2.9616 | LR: 2.53e-04 | 930.8ms/step | 21.5 samples/s
+Validation at step 22250: Loss = 2.9890, Perplexity = 19.86
+Step 22300/44000 | Loss: 2.9549 | LR: 2.51e-04 | 956.1ms/step | 20.9 samples/s
+Step 22400/44000 | Loss: 2.9759 | LR: 2.50e-04 | 930.8ms/step | 21.5 samples/s
+Step 22500/44000 | Loss: 2.9560 | LR: 2.48e-04 | 930.7ms/step | 21.5 samples/s
+Validation at step 22500: Loss = 2.9854, Perplexity = 19.79
+Step 22600/44000 | Loss: 2.9483 | LR: 2.46e-04 | 958.5ms/step | 20.9 samples/s
+Step 22700/44000 | Loss: 2.9690 | LR: 2.44e-04 | 934.6ms/step | 21.4 samples/s
+Validation at step 22750: Loss = 2.9836, Perplexity = 19.76
+Step 22800/44000 | Loss: 2.9792 | LR: 2.42e-04 | 962.5ms/step | 20.8 samples/s
+Step 22900/44000 | Loss: 2.9624 | LR: 2.40e-04 | 937.6ms/step | 21.3 samples/s
+Step 23000/44000 | Loss: 2.9584 | LR: 2.39e-04 | 938.8ms/step | 21.3 samples/s
+Validation at step 23000: Loss = 2.9814, Perplexity = 19.71
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-23000
+Step 23100/44000 | Loss: 2.9736 | LR: 2.37e-04 | 1007.3ms/step | 19.9 samples/s
+Step 23200/44000 | Loss: 2.9554 | LR: 2.35e-04 | 939.0ms/step | 21.3 samples/s
+Validation at step 23250: Loss = 2.9790, Perplexity = 19.67
+Step 23300/44000 | Loss: 2.9626 | LR: 2.33e-04 | 965.3ms/step | 20.7 samples/s
+Step 23400/44000 | Loss: 2.9702 | LR: 2.31e-04 | 938.7ms/step | 21.3 samples/s
+Step 23500/44000 | Loss: 2.9323 | LR: 2.30e-04 | 939.3ms/step | 21.3 samples/s
+Validation at step 23500: Loss = 2.9769, Perplexity = 19.63
+Step 23600/44000 | Loss: 2.9337 | LR: 2.28e-04 | 965.1ms/step | 20.7 samples/s
+Step 23700/44000 | Loss: 2.9606 | LR: 2.26e-04 | 938.9ms/step | 21.3 samples/s
+Validation at step 23750: Loss = 2.9745, Perplexity = 19.58
+Step 23800/44000 | Loss: 2.9647 | LR: 2.24e-04 | 966.8ms/step | 20.7 samples/s
+Step 23900/44000 | Loss: 2.9400 | LR: 2.22e-04 | 939.3ms/step | 21.3 samples/s
+Step 24000/44000 | Loss: 2.9229 | LR: 2.21e-04 | 939.3ms/step | 21.3 samples/s
+Validation at step 24000: Loss = 2.9721, Perplexity = 19.53
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-24000
+Step 24100/44000 | Loss: 2.9388 | LR: 2.19e-04 | 1005.6ms/step | 19.9 samples/s
+Step 24200/44000 | Loss: 2.9428 | LR: 2.17e-04 | 938.8ms/step | 21.3 samples/s
+Validation at step 24250: Loss = 2.9706, Perplexity = 19.50
+Step 24300/44000 | Loss: 2.9537 | LR: 2.15e-04 | 965.4ms/step | 20.7 samples/s
+Step 24400/44000 | Loss: 2.9368 | LR: 2.13e-04 | 938.6ms/step | 21.3 samples/s
+Step 24500/44000 | Loss: 2.9330 | LR: 2.12e-04 | 938.8ms/step | 21.3 samples/s
+Validation at step 24500: Loss = 2.9684, Perplexity = 19.46
+Step 24600/44000 | Loss: 2.9344 | LR: 2.10e-04 | 964.5ms/step | 20.7 samples/s
+Step 24700/44000 | Loss: 2.9430 | LR: 2.08e-04 | 940.1ms/step | 21.3 samples/s
+Validation at step 24750: Loss = 2.9657, Perplexity = 19.41
+Step 24800/44000 | Loss: 2.9493 | LR: 2.06e-04 | 966.1ms/step | 20.7 samples/s
+Step 24900/44000 | Loss: 2.9297 | LR: 2.04e-04 | 939.6ms/step | 21.3 samples/s
+Step 25000/44000 | Loss: 2.9206 | LR: 2.03e-04 | 939.3ms/step | 21.3 samples/s
+Validation at step 25000: Loss = 2.9654, Perplexity = 19.40
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-25000
+Step 25100/44000 | Loss: 2.9420 | LR: 2.01e-04 | 998.9ms/step | 20.0 samples/s
+Step 25200/44000 | Loss: 2.9245 | LR: 1.99e-04 | 929.0ms/step | 21.5 samples/s
+Validation at step 25250: Loss = 2.9637, Perplexity = 19.37
+Step 25300/44000 | Loss: 2.9358 | LR: 1.97e-04 | 953.6ms/step | 21.0 samples/s
+Step 25400/44000 | Loss: 2.9126 | LR: 1.96e-04 | 928.2ms/step | 21.5 samples/s
+Step 25500/44000 | Loss: 2.9145 | LR: 1.94e-04 | 927.5ms/step | 21.6 samples/s
+Validation at step 25500: Loss = 2.9619, Perplexity = 19.34
+Step 25600/44000 | Loss: 2.9301 | LR: 1.92e-04 | 953.8ms/step | 21.0 samples/s
+Step 25700/44000 | Loss: 2.9056 | LR: 1.90e-04 | 928.0ms/step | 21.6 samples/s
+Validation at step 25750: Loss = 2.9600, Perplexity = 19.30
+Step 25800/44000 | Loss: 2.9339 | LR: 1.88e-04 | 954.6ms/step | 21.0 samples/s
+Step 25900/44000 | Loss: 2.9463 | LR: 1.87e-04 | 928.5ms/step | 21.5 samples/s
+Step 26000/44000 | Loss: 2.9133 | LR: 1.85e-04 | 928.0ms/step | 21.6 samples/s
+Validation at step 26000: Loss = 2.9588, Perplexity = 19.27
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-26000
+Step 26100/44000 | Loss: 2.9317 | LR: 1.83e-04 | 987.0ms/step | 20.3 samples/s
+Step 26200/44000 | Loss: 2.9323 | LR: 1.81e-04 | 928.3ms/step | 21.5 samples/s
+Validation at step 26250: Loss = 2.9554, Perplexity = 19.21
+Step 26300/44000 | Loss: 2.9441 | LR: 1.80e-04 | 953.3ms/step | 21.0 samples/s
+Step 26400/44000 | Loss: 2.9238 | LR: 1.78e-04 | 928.0ms/step | 21.6 samples/s
+Step 26500/44000 | Loss: 2.9263 | LR: 1.76e-04 | 928.2ms/step | 21.5 samples/s
+Validation at step 26500: Loss = 2.9548, Perplexity = 19.20
+Step 26600/44000 | Loss: 2.9408 | LR: 1.74e-04 | 954.0ms/step | 21.0 samples/s
+Step 26700/44000 | Loss: 2.9442 | LR: 1.73e-04 | 928.0ms/step | 21.6 samples/s
+Validation at step 26750: Loss = 2.9534, Perplexity = 19.17
+Step 26800/44000 | Loss: 2.9483 | LR: 1.71e-04 | 953.4ms/step | 21.0 samples/s
+Step 26900/44000 | Loss: 2.9202 | LR: 1.69e-04 | 927.4ms/step | 21.6 samples/s
+Step 27000/44000 | Loss: 2.9278 | LR: 1.68e-04 | 927.9ms/step | 21.6 samples/s
+Validation at step 27000: Loss = 2.9519, Perplexity = 19.14
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-27000
+Step 27100/44000 | Loss: 2.9221 | LR: 1.66e-04 | 986.0ms/step | 20.3 samples/s
+Step 27200/44000 | Loss: 2.9470 | LR: 1.64e-04 | 927.9ms/step | 21.6 samples/s
+Validation at step 27250: Loss = 2.9502, Perplexity = 19.11
+Step 27300/44000 | Loss: 2.9277 | LR: 1.62e-04 | 953.5ms/step | 21.0 samples/s
+Step 27400/44000 | Loss: 2.9054 | LR: 1.61e-04 | 927.4ms/step | 21.6 samples/s
+Step 27500/44000 | Loss: 2.9130 | LR: 1.59e-04 | 926.9ms/step | 21.6 samples/s
+Validation at step 27500: Loss = 2.9495, Perplexity = 19.10
+Step 27600/44000 | Loss: 2.9025 | LR: 1.57e-04 | 953.6ms/step | 21.0 samples/s
+Step 27700/44000 | Loss: 2.9110 | LR: 1.56e-04 | 927.6ms/step | 21.6 samples/s
+Validation at step 27750: Loss = 2.9485, Perplexity = 19.08
+Step 27800/44000 | Loss: 2.9248 | LR: 1.54e-04 | 953.3ms/step | 21.0 samples/s
+Step 27900/44000 | Loss: 2.9369 | LR: 1.52e-04 | 927.3ms/step | 21.6 samples/s
+Step 28000/44000 | Loss: 2.9339 | LR: 1.51e-04 | 927.0ms/step | 21.6 samples/s
+Validation at step 28000: Loss = 2.9465, Perplexity = 19.04
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-28000
+Step 28100/44000 | Loss: 2.9117 | LR: 1.49e-04 | 984.7ms/step | 20.3 samples/s
+Step 28200/44000 | Loss: 2.9146 | LR: 1.47e-04 | 926.7ms/step | 21.6 samples/s
+Validation at step 28250: Loss = 2.9461, Perplexity = 19.03
+Step 28300/44000 | Loss: 2.9231 | LR: 1.46e-04 | 953.3ms/step | 21.0 samples/s
+Step 28400/44000 | Loss: 2.9446 | LR: 1.44e-04 | 926.9ms/step | 21.6 samples/s
+Step 28500/44000 | Loss: 2.9260 | LR: 1.42e-04 | 927.0ms/step | 21.6 samples/s
+Validation at step 28500: Loss = 2.9451, Perplexity = 19.01
+Step 28600/44000 | Loss: 2.8893 | LR: 1.41e-04 | 953.3ms/step | 21.0 samples/s
+Step 28700/44000 | Loss: 2.9010 | LR: 1.39e-04 | 927.3ms/step | 21.6 samples/s
+Validation at step 28750: Loss = 2.9437, Perplexity = 18.99
+Step 28800/44000 | Loss: 2.9225 | LR: 1.38e-04 | 952.8ms/step | 21.0 samples/s
+Step 28900/44000 | Loss: 2.9357 | LR: 1.36e-04 | 926.6ms/step | 21.6 samples/s
+Step 29000/44000 | Loss: 2.9387 | LR: 1.34e-04 | 926.9ms/step | 21.6 samples/s
+Validation at step 29000: Loss = 2.9428, Perplexity = 18.97
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-29000
+Step 29100/44000 | Loss: 2.9221 | LR: 1.33e-04 | 986.6ms/step | 20.3 samples/s
+Step 29200/44000 | Loss: 2.9033 | LR: 1.31e-04 | 927.4ms/step | 21.6 samples/s
+Validation at step 29250: Loss = 2.9419, Perplexity = 18.95
+Step 29300/44000 | Loss: 2.9168 | LR: 1.29e-04 | 953.1ms/step | 21.0 samples/s
+Step 29400/44000 | Loss: 2.9084 | LR: 1.28e-04 | 927.2ms/step | 21.6 samples/s
+Step 29500/44000 | Loss: 2.8965 | LR: 1.26e-04 | 927.7ms/step | 21.6 samples/s
+Validation at step 29500: Loss = 2.9404, Perplexity = 18.92
+Step 29600/44000 | Loss: 2.9560 | LR: 1.25e-04 | 952.6ms/step | 21.0 samples/s
+Step 29700/44000 | Loss: 2.9283 | LR: 1.23e-04 | 927.3ms/step | 21.6 samples/s
+Validation at step 29750: Loss = 2.9399, Perplexity = 18.91
+Step 29800/44000 | Loss: 2.9303 | LR: 1.22e-04 | 952.7ms/step | 21.0 samples/s
+Step 29900/44000 | Loss: 2.8983 | LR: 1.20e-04 | 927.4ms/step | 21.6 samples/s
+Step 30000/44000 | Loss: 2.9248 | LR: 1.19e-04 | 927.4ms/step | 21.6 samples/s
+Validation at step 30000: Loss = 2.9388, Perplexity = 18.89
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-30000
+Step 30100/44000 | Loss: 2.9132 | LR: 1.17e-04 | 985.6ms/step | 20.3 samples/s
+Step 30200/44000 | Loss: 2.9095 | LR: 1.15e-04 | 928.1ms/step | 21.5 samples/s
+Validation at step 30250: Loss = 2.9379, Perplexity = 18.88
+Step 30300/44000 | Loss: 2.9104 | LR: 1.14e-04 | 953.1ms/step | 21.0 samples/s
+Step 30400/44000 | Loss: 2.9166 | LR: 1.12e-04 | 927.7ms/step | 21.6 samples/s
+Step 30500/44000 | Loss: 2.9077 | LR: 1.11e-04 | 928.1ms/step | 21.5 samples/s
+Validation at step 30500: Loss = 2.9369, Perplexity = 18.86
+Step 30600/44000 | Loss: 2.9338 | LR: 1.09e-04 | 953.8ms/step | 21.0 samples/s
+Step 30700/44000 | Loss: 2.9095 | LR: 1.08e-04 | 927.7ms/step | 21.6 samples/s
+Validation at step 30750: Loss = 2.9366, Perplexity = 18.85
+Step 30800/44000 | Loss: 2.9205 | LR: 1.06e-04 | 953.1ms/step | 21.0 samples/s
+Step 30900/44000 | Loss: 2.9083 | LR: 1.05e-04 | 928.2ms/step | 21.5 samples/s
+Step 31000/44000 | Loss: 2.9010 | LR: 1.03e-04 | 927.9ms/step | 21.6 samples/s
+Validation at step 31000: Loss = 2.9357, Perplexity = 18.83
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-31000
+Step 31100/44000 | Loss: 2.9103 | LR: 1.02e-04 | 986.4ms/step | 20.3 samples/s
+Step 31200/44000 | Loss: 2.9131 | LR: 1.00e-04 | 927.9ms/step | 21.6 samples/s
+Validation at step 31250: Loss = 2.9348, Perplexity = 18.82
+Step 31300/44000 | Loss: 2.8852 | LR: 9.90e-05 | 953.4ms/step | 21.0 samples/s
+Step 31400/44000 | Loss: 2.9149 | LR: 9.76e-05 | 929.4ms/step | 21.5 samples/s
+Step 31500/44000 | Loss: 2.9110 | LR: 9.62e-05 | 933.5ms/step | 21.4 samples/s
+Validation at step 31500: Loss = 2.9344, Perplexity = 18.81
+Step 31600/44000 | Loss: 2.8978 | LR: 9.47e-05 | 961.3ms/step | 20.8 samples/s
+Step 31700/44000 | Loss: 2.9056 | LR: 9.33e-05 | 934.6ms/step | 21.4 samples/s
+Validation at step 31750: Loss = 2.9343, Perplexity = 18.81
+Step 31800/44000 | Loss: 2.8787 | LR: 9.19e-05 | 960.1ms/step | 20.8 samples/s
+Step 31900/44000 | Loss: 2.9150 | LR: 9.05e-05 | 934.7ms/step | 21.4 samples/s
+Step 32000/44000 | Loss: 2.9110 | LR: 8.91e-05 | 934.4ms/step | 21.4 samples/s
+Validation at step 32000: Loss = 2.9333, Perplexity = 18.79
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-32000
+Step 32100/44000 | Loss: 2.8934 | LR: 8.77e-05 | 993.7ms/step | 20.1 samples/s
+Step 32200/44000 | Loss: 2.9078 | LR: 8.63e-05 | 933.4ms/step | 21.4 samples/s
+Validation at step 32250: Loss = 2.9329, Perplexity = 18.78
+Step 32300/44000 | Loss: 2.9061 | LR: 8.50e-05 | 959.4ms/step | 20.8 samples/s
+Step 32400/44000 | Loss: 2.9198 | LR: 8.36e-05 | 932.7ms/step | 21.4 samples/s
+Step 32500/44000 | Loss: 2.8909 | LR: 8.23e-05 | 935.5ms/step | 21.4 samples/s
+Validation at step 32500: Loss = 2.9326, Perplexity = 18.78
+Step 32600/44000 | Loss: 2.8991 | LR: 8.09e-05 | 962.6ms/step | 20.8 samples/s
+Step 32700/44000 | Loss: 2.9095 | LR: 7.96e-05 | 935.1ms/step | 21.4 samples/s
+Validation at step 32750: Loss = 2.9319, Perplexity = 18.76
+Step 32800/44000 | Loss: 2.8881 | LR: 7.83e-05 | 962.2ms/step | 20.8 samples/s
+Step 32900/44000 | Loss: 2.9073 | LR: 7.70e-05 | 935.4ms/step | 21.4 samples/s
+Step 33000/44000 | Loss: 2.9061 | LR: 7.56e-05 | 934.8ms/step | 21.4 samples/s
+Validation at step 33000: Loss = 2.9312, Perplexity = 18.75
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-33000
+Step 33100/44000 | Loss: 2.9077 | LR: 7.44e-05 | 996.0ms/step | 20.1 samples/s
+Step 33200/44000 | Loss: 2.9283 | LR: 7.31e-05 | 935.2ms/step | 21.4 samples/s
+Validation at step 33250: Loss = 2.9312, Perplexity = 18.75
+Step 33300/44000 | Loss: 2.9126 | LR: 7.18e-05 | 961.8ms/step | 20.8 samples/s
+Step 33400/44000 | Loss: 2.9152 | LR: 7.05e-05 | 934.7ms/step | 21.4 samples/s
+Step 33500/44000 | Loss: 2.9007 | LR: 6.93e-05 | 934.0ms/step | 21.4 samples/s
+Validation at step 33500: Loss = 2.9311, Perplexity = 18.75
+Step 33600/44000 | Loss: 2.8966 | LR: 6.80e-05 | 960.1ms/step | 20.8 samples/s
+Step 33700/44000 | Loss: 2.9045 | LR: 6.68e-05 | 934.2ms/step | 21.4 samples/s
+Validation at step 33750: Loss = 2.9306, Perplexity = 18.74
+Step 33800/44000 | Loss: 2.8911 | LR: 6.55e-05 | 960.3ms/step | 20.8 samples/s
+Step 33900/44000 | Loss: 2.8846 | LR: 6.43e-05 | 934.7ms/step | 21.4 samples/s
+Step 34000/44000 | Loss: 2.9252 | LR: 6.31e-05 | 934.4ms/step | 21.4 samples/s
+Validation at step 34000: Loss = 2.9301, Perplexity = 18.73
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-34000
+Step 34100/44000 | Loss: 2.9045 | LR: 6.19e-05 | 996.9ms/step | 20.1 samples/s
+Step 34200/44000 | Loss: 2.9095 | LR: 6.07e-05 | 935.1ms/step | 21.4 samples/s
+Validation at step 34250: Loss = 2.9299, Perplexity = 18.73
+Step 34300/44000 | Loss: 2.9153 | LR: 5.95e-05 | 961.3ms/step | 20.8 samples/s
+Step 34400/44000 | Loss: 2.8917 | LR: 5.84e-05 | 935.2ms/step | 21.4 samples/s
+Step 34500/44000 | Loss: 2.9055 | LR: 5.72e-05 | 933.5ms/step | 21.4 samples/s
+Validation at step 34500: Loss = 2.9296, Perplexity = 18.72
+Step 34600/44000 | Loss: 2.9269 | LR: 5.60e-05 | 959.8ms/step | 20.8 samples/s
+Step 34700/44000 | Loss: 2.8946 | LR: 5.49e-05 | 934.3ms/step | 21.4 samples/s
+Validation at step 34750: Loss = 2.9293, Perplexity = 18.71
+Step 34800/44000 | Loss: 2.9152 | LR: 5.38e-05 | 961.0ms/step | 20.8 samples/s
+Step 34900/44000 | Loss: 2.9092 | LR: 5.27e-05 | 934.0ms/step | 21.4 samples/s
+Step 35000/44000 | Loss: 2.8999 | LR: 5.15e-05 | 934.4ms/step | 21.4 samples/s
+Validation at step 35000: Loss = 2.9289, Perplexity = 18.71
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-35000
+Step 35100/44000 | Loss: 2.9048 | LR: 5.04e-05 | 995.0ms/step | 20.1 samples/s
+Step 35200/44000 | Loss: 2.9160 | LR: 4.94e-05 | 934.1ms/step | 21.4 samples/s
+Validation at step 35250: Loss = 2.9287, Perplexity = 18.70
+Step 35300/44000 | Loss: 2.9055 | LR: 4.83e-05 | 961.2ms/step | 20.8 samples/s
+Step 35400/44000 | Loss: 2.8832 | LR: 4.72e-05 | 934.0ms/step | 21.4 samples/s
+Step 35500/44000 | Loss: 2.9128 | LR: 4.62e-05 | 934.6ms/step | 21.4 samples/s
+Validation at step 35500: Loss = 2.9284, Perplexity = 18.70
+Step 35600/44000 | Loss: 2.9118 | LR: 4.51e-05 | 961.4ms/step | 20.8 samples/s
+Step 35700/44000 | Loss: 2.8983 | LR: 4.41e-05 | 934.8ms/step | 21.4 samples/s
+Validation at step 35750: Loss = 2.9283, Perplexity = 18.70
+Step 35800/44000 | Loss: 2.8964 | LR: 4.31e-05 | 960.3ms/step | 20.8 samples/s
+Step 35900/44000 | Loss: 2.9130 | LR: 4.20e-05 | 934.2ms/step | 21.4 samples/s
+Step 36000/44000 | Loss: 2.8960 | LR: 4.10e-05 | 933.4ms/step | 21.4 samples/s
+Validation at step 36000: Loss = 2.9282, Perplexity = 18.69
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-36000
+Step 36100/44000 | Loss: 2.8903 | LR: 4.00e-05 | 997.2ms/step | 20.1 samples/s
+Step 36200/44000 | Loss: 2.9142 | LR: 3.91e-05 | 934.8ms/step | 21.4 samples/s
+Validation at step 36250: Loss = 2.9281, Perplexity = 18.69
+Step 36300/44000 | Loss: 2.8819 | LR: 3.81e-05 | 960.3ms/step | 20.8 samples/s
+Step 36400/44000 | Loss: 2.8859 | LR: 3.71e-05 | 934.2ms/step | 21.4 samples/s
+Step 36500/44000 | Loss: 2.9107 | LR: 3.62e-05 | 934.0ms/step | 21.4 samples/s
+Validation at step 36500: Loss = 2.9278, Perplexity = 18.69
+Step 36600/44000 | Loss: 2.9050 | LR: 3.53e-05 | 960.0ms/step | 20.8 samples/s
+Step 36700/44000 | Loss: 2.8765 | LR: 3.43e-05 | 934.7ms/step | 21.4 samples/s
+Validation at step 36750: Loss = 2.9277, Perplexity = 18.68
+Step 36800/44000 | Loss: 2.8936 | LR: 3.34e-05 | 960.7ms/step | 20.8 samples/s
+Step 36900/44000 | Loss: 2.9041 | LR: 3.25e-05 | 934.3ms/step | 21.4 samples/s
+Step 37000/44000 | Loss: 2.8850 | LR: 3.16e-05 | 935.0ms/step | 21.4 samples/s
+Validation at step 37000: Loss = 2.9276, Perplexity = 18.68
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-37000
+Step 37100/44000 | Loss: 2.9044 | LR: 3.07e-05 | 996.0ms/step | 20.1 samples/s
+Step 37200/44000 | Loss: 2.9091 | LR: 2.99e-05 | 934.9ms/step | 21.4 samples/s
+Validation at step 37250: Loss = 2.9275, Perplexity = 18.68
+Step 37300/44000 | Loss: 2.9077 | LR: 2.90e-05 | 961.3ms/step | 20.8 samples/s
+Step 37400/44000 | Loss: 2.9120 | LR: 2.82e-05 | 935.0ms/step | 21.4 samples/s
+Step 37500/44000 | Loss: 2.8890 | LR: 2.74e-05 | 935.7ms/step | 21.4 samples/s
+Validation at step 37500: Loss = 2.9275, Perplexity = 18.68
+Step 37600/44000 | Loss: 2.8934 | LR: 2.65e-05 | 962.4ms/step | 20.8 samples/s
+Step 37700/44000 | Loss: 2.9062 | LR: 2.57e-05 | 935.1ms/step | 21.4 samples/s
+Validation at step 37750: Loss = 2.9274, Perplexity = 18.68
+Step 37800/44000 | Loss: 2.8865 | LR: 2.49e-05 | 962.6ms/step | 20.8 samples/s
+Step 37900/44000 | Loss: 2.8871 | LR: 2.41e-05 | 937.1ms/step | 21.3 samples/s
+Step 38000/44000 | Loss: 2.8747 | LR: 2.34e-05 | 936.5ms/step | 21.4 samples/s
+Validation at step 38000: Loss = 2.9274, Perplexity = 18.68
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-38000
+Step 38100/44000 | Loss: 2.8928 | LR: 2.26e-05 | 997.2ms/step | 20.1 samples/s
+Step 38200/44000 | Loss: 2.8930 | LR: 2.19e-05 | 936.4ms/step | 21.4 samples/s
+Validation at step 38250: Loss = 2.9274, Perplexity = 18.68
+Step 38300/44000 | Loss: 2.9097 | LR: 2.11e-05 | 964.0ms/step | 20.7 samples/s
+Step 38400/44000 | Loss: 2.9045 | LR: 2.04e-05 | 937.2ms/step | 21.3 samples/s
+Step 38500/44000 | Loss: 2.9008 | LR: 1.97e-05 | 937.1ms/step | 21.3 samples/s
+Validation at step 38500: Loss = 2.9273, Perplexity = 18.68
+Step 38600/44000 | Loss: 2.8909 | LR: 1.90e-05 | 962.9ms/step | 20.8 samples/s
+Step 38700/44000 | Loss: 2.8889 | LR: 1.83e-05 | 936.6ms/step | 21.4 samples/s
+Validation at step 38750: Loss = 2.9272, Perplexity = 18.68
+Step 38800/44000 | Loss: 2.9015 | LR: 1.76e-05 | 962.5ms/step | 20.8 samples/s
+Step 38900/44000 | Loss: 2.8924 | LR: 1.70e-05 | 934.3ms/step | 21.4 samples/s
+Step 39000/44000 | Loss: 2.9024 | LR: 1.63e-05 | 934.4ms/step | 21.4 samples/s
+Validation at step 39000: Loss = 2.9272, Perplexity = 18.67
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-39000
+Step 39100/44000 | Loss: 2.9018 | LR: 1.57e-05 | 995.5ms/step | 20.1 samples/s
+Step 39200/44000 | Loss: 2.8936 | LR: 1.50e-05 | 935.1ms/step | 21.4 samples/s
+Validation at step 39250: Loss = 2.9271, Perplexity = 18.67
+Step 39300/44000 | Loss: 2.8826 | LR: 1.44e-05 | 961.5ms/step | 20.8 samples/s
+Step 39400/44000 | Loss: 2.9100 | LR: 1.38e-05 | 934.6ms/step | 21.4 samples/s
+Step 39500/44000 | Loss: 2.8862 | LR: 1.32e-05 | 935.4ms/step | 21.4 samples/s
+Validation at step 39500: Loss = 2.9271, Perplexity = 18.67
+Step 39600/44000 | Loss: 2.9115 | LR: 1.27e-05 | 963.8ms/step | 20.8 samples/s
+Step 39700/44000 | Loss: 2.9172 | LR: 1.21e-05 | 937.7ms/step | 21.3 samples/s
+Validation at step 39750: Loss = 2.9271, Perplexity = 18.67
+Step 39800/44000 | Loss: 2.8882 | LR: 1.15e-05 | 963.2ms/step | 20.8 samples/s
+Step 39900/44000 | Loss: 2.9041 | LR: 1.10e-05 | 936.6ms/step | 21.4 samples/s
+Step 40000/44000 | Loss: 2.8748 | LR: 1.05e-05 | 936.6ms/step | 21.4 samples/s
+Validation at step 40000: Loss = 2.9271, Perplexity = 18.67
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-40000
+Step 40100/44000 | Loss: 2.8959 | LR: 9.96e-06 | 994.6ms/step | 20.1 samples/s
+Step 40200/44000 | Loss: 2.9089 | LR: 9.46e-06 | 935.4ms/step | 21.4 samples/s
+Validation at step 40250: Loss = 2.9271, Perplexity = 18.67
+Step 40300/44000 | Loss: 2.9043 | LR: 8.97e-06 | 960.9ms/step | 20.8 samples/s
+Step 40400/44000 | Loss: 2.9075 | LR: 8.50e-06 | 934.0ms/step | 21.4 samples/s
+Step 40500/44000 | Loss: 2.8956 | LR: 8.04e-06 | 933.8ms/step | 21.4 samples/s
+Validation at step 40500: Loss = 2.9270, Perplexity = 18.67
+Step 40600/44000 | Loss: 2.8959 | LR: 7.59e-06 | 959.7ms/step | 20.8 samples/s
+Step 40700/44000 | Loss: 2.9027 | LR: 7.15e-06 | 934.0ms/step | 21.4 samples/s
+Validation at step 40750: Loss = 2.9270, Perplexity = 18.67
+Step 40800/44000 | Loss: 2.9031 | LR: 6.72e-06 | 960.4ms/step | 20.8 samples/s
+Step 40900/44000 | Loss: 2.8824 | LR: 6.31e-06 | 935.1ms/step | 21.4 samples/s
+Step 41000/44000 | Loss: 2.8955 | LR: 5.91e-06 | 934.6ms/step | 21.4 samples/s
+Validation at step 41000: Loss = 2.9270, Perplexity = 18.67
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-41000
+Step 41100/44000 | Loss: 2.9007 | LR: 5.53e-06 | 994.3ms/step | 20.1 samples/s
+Step 41200/44000 | Loss: 2.9193 | LR: 5.15e-06 | 934.5ms/step | 21.4 samples/s
+Validation at step 41250: Loss = 2.9270, Perplexity = 18.67
+Step 41300/44000 | Loss: 2.8862 | LR: 4.79e-06 | 962.8ms/step | 20.8 samples/s
+Step 41400/44000 | Loss: 2.9035 | LR: 4.45e-06 | 934.8ms/step | 21.4 samples/s
+Step 41500/44000 | Loss: 2.8926 | LR: 4.11e-06 | 934.1ms/step | 21.4 samples/s
+Validation at step 41500: Loss = 2.9270, Perplexity = 18.67
+Step 41600/44000 | Loss: 2.8982 | LR: 3.79e-06 | 960.0ms/step | 20.8 samples/s
+Step 41700/44000 | Loss: 2.9078 | LR: 3.48e-06 | 934.3ms/step | 21.4 samples/s
+Validation at step 41750: Loss = 2.9270, Perplexity = 18.67
+Step 41800/44000 | Loss: 2.8827 | LR: 3.19e-06 | 960.5ms/step | 20.8 samples/s
+Step 41900/44000 | Loss: 2.8953 | LR: 2.90e-06 | 934.0ms/step | 21.4 samples/s
+Step 42000/44000 | Loss: 2.8826 | LR: 2.63e-06 | 934.8ms/step | 21.4 samples/s
+Validation at step 42000: Loss = 2.9270, Perplexity = 18.67
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-42000
+Step 42100/44000 | Loss: 2.8993 | LR: 2.38e-06 | 995.5ms/step | 20.1 samples/s
+Step 42200/44000 | Loss: 2.9043 | LR: 2.13e-06 | 934.6ms/step | 21.4 samples/s
+Validation at step 42250: Loss = 2.9270, Perplexity = 18.67
+Step 42300/44000 | Loss: 2.8867 | LR: 1.90e-06 | 962.0ms/step | 20.8 samples/s
+Step 42400/44000 | Loss: 2.8991 | LR: 1.69e-06 | 935.4ms/step | 21.4 samples/s
+Step 42500/44000 | Loss: 2.8846 | LR: 1.48e-06 | 935.9ms/step | 21.4 samples/s
+Validation at step 42500: Loss = 2.9270, Perplexity = 18.67
+Step 42600/44000 | Loss: 2.8859 | LR: 1.29e-06 | 962.0ms/step | 20.8 samples/s
+Step 42700/44000 | Loss: 2.9037 | LR: 1.11e-06 | 936.5ms/step | 21.4 samples/s
+Validation at step 42750: Loss = 2.9270, Perplexity = 18.67
+Step 42800/44000 | Loss: 2.9030 | LR: 9.49e-07 | 962.6ms/step | 20.8 samples/s
+Step 42900/44000 | Loss: 2.9108 | LR: 7.98e-07 | 936.8ms/step | 21.3 samples/s
+Step 43000/44000 | Loss: 2.8970 | LR: 6.59e-07 | 936.7ms/step | 21.4 samples/s
+Validation at step 43000: Loss = 2.9270, Perplexity = 18.67
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-43000
+Step 43100/44000 | Loss: 2.9240 | LR: 5.34e-07 | 996.8ms/step | 20.1 samples/s
+Step 43200/44000 | Loss: 2.9003 | LR: 4.22e-07 | 936.8ms/step | 21.4 samples/s
+Validation at step 43250: Loss = 2.9270, Perplexity = 18.67
+Step 43300/44000 | Loss: 2.9115 | LR: 3.23e-07 | 962.3ms/step | 20.8 samples/s
+Step 43400/44000 | Loss: 2.8716 | LR: 2.37e-07 | 935.6ms/step | 21.4 samples/s
+Step 43500/44000 | Loss: 2.9060 | LR: 1.65e-07 | 936.7ms/step | 21.4 samples/s
+Validation at step 43500: Loss = 2.9270, Perplexity = 18.67
+Step 43600/44000 | Loss: 2.9181 | LR: 1.06e-07 | 961.9ms/step | 20.8 samples/s
+Step 43700/44000 | Loss: 2.9149 | LR: 5.94e-08 | 935.6ms/step | 21.4 samples/s
+Validation at step 43750: Loss = 2.9270, Perplexity = 18.67
+Step 43800/44000 | Loss: 2.9105 | LR: 2.64e-08 | 961.9ms/step | 20.8 samples/s
+Step 43900/44000 | Loss: 2.8927 | LR: 6.60e-09 | 936.3ms/step | 21.4 samples/s
+Step 44000/44000 | Loss: 2.9010 | LR: 0.00e+00 | 936.5ms/step | 21.4 samples/s
+Validation at step 44000: Loss = 2.9270, Perplexity = 18.67
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-44000
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-44000
+wandb: 
+wandb: Run history:
+wandb: eval/dense_interval_active ▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁
+wandb:    eval/dense_interval_end ▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁
+wandb:  eval/dense_interval_start ▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁
+wandb:                  eval/loss █▆▅▄▄▃▃▃▃▃▃▃▂▂▂▂▂▂▂▂▂▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁
+wandb:            eval/perplexity █▆▆▄▄▃▃▃▃▂▂▂▂▂▂▂▂▂▂▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁
+wandb:  eval/uses_dense_attention ▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁
+wandb: eval/uses_sparse_attention ▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁
+wandb:          speed/ms_per_step ▁▄▁▄▁▁▁▁▄▁▄▁▄▅▂▆▂▄▁▄▄▁▁▄█▂▂▂▂▂▅▂▅▂▅▅▂▅▂▂
+wandb:      speed/samples_per_sec ██▂▂██▃▃█▆▆▃▅█▂▅█▂▅████▇▁▆█▆█▆▂▅▇▇▅▅▇▇▂▇
+wandb:        speed/steps_per_sec ██▅█████████▇▁▇▂▆████████▆▆▇▅▅▇▇▇▇▇▂▇▇▇▇
+wandb:                         +4 ...
+wandb: 
+wandb: Run summary:
+wandb: eval/dense_interval_active 0
+wandb:    eval/dense_interval_end -1
+wandb:  eval/dense_interval_start -1
+wandb:                  eval/loss 2.927
+wandb:            eval/perplexity 18.67158
+wandb:  eval/uses_dense_attention 0
+wandb: eval/uses_sparse_attention 1
+wandb:          speed/ms_per_step 936.49953
+wandb:      speed/samples_per_sec 21.35612
+wandb:        speed/steps_per_sec 1.06781
+wandb:                         +4 ...
+wandb: 
+wandb: You can sync this run to the cloud by running:
+wandb: wandb sync /mloscratch/homes/navasard/sparse/SparseAtt/wandb/offline-run-20260921_164533-kbn66izc
+wandb: Find logs at: ./wandb/offline-run-20260921_164533-kbn66izc/logs
+Training complete!

epfl_final/logs/ALIBI_20260921/restart_20260921_131155/seed47.log

@@ -0,0 +1,187 @@
+New run: Checkpoints will be saved to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s47/run_20260921_131213
+Running config:
+{
+  "model_path": "/mloscratch/homes/navasard/qwen3_06b_base",
+  "model_dtype": "bfloat16",
+  "init_from_scratch": true,
+  "train_path": "/mloscratch/homes/navasard/data/fineweb_large/tokenized",
+  "val_path": null,
+  "max_seq_length": 8192,
+  "num_workers": 8,
+  "output_dir": "/mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s47/run_20260921_131213",
+  "resume_from_checkpoint": null,
+  "force_new_run": true,
+  "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,
+  "resume_micro_step_override": null,
+  "strategy": "ddp",
+  "mixed_precision": "bf16",
+  "wandb_enabled": true,
+  "wandb_project": "stochastic_sparsity",
+  "wandb_name": "365m alibi 8k s47",
+  "wandb_tags": [
+    "qwen3",
+    "365m",
+    "8k",
+    "scaleladder",
+    "alibi",
+    "seed47"
+  ],
+  "hidden_size": 768,
+  "intermediate_size": 3072,
+  "num_attention_heads": 6,
+  "num_key_value_heads": 2,
+  "num_hidden_layers": 28,
+  "head_dim": null,
+  "attention_dropout": null,
+  "rope_train_position_scale_set": null,
+  "positional": "alibi",
+  "attention_type": "dense",
+  "sparse_block_size": 128,
+  "sparse_num_random_blocks": 31,
+  "sparse_mask_pool_size": 1,
+  "sparse_val_dense_compare": false,
+  "sparse_layer_indices": null,
+  "sparse_mask_cache_dir": null,
+  "sparse_mask_seed": 0,
+  "sparse_mixed_nrb_set": null,
+  "dense_mode": false,
+  "dense_step_intervals": null
+}
+Loading tokenizer from /mloscratch/homes/navasard/qwen3_06b_base...
+Initializing model from scratch (random weights)...
+Checkpoint activations: False
+
+
+================================================================================
+MODEL ARCHITECTURE
+================================================================================
+Model config:
+  vocab_size: 151936
+  hidden_size: 768
+  num_hidden_layers: 28
+  num_attention_heads: 6
+  num_key_value_heads: 2
+  head_dim: 128
+  intermediate_size: 3072
+  max_position_embeddings: 32768
+  rope_theta: 1000000
+  rope_train_position_scale_set: None
+  attention_type: dense
+
+Parameter count:
+  Total parameters: 358,958,848
+  Trainable parameters: 358,958,848
+  Non-trainable parameters: 0
+  Model size (BF16): 0.72 GB
+
+--------------------------------------------------------------------------------
+LAYER-BY-LAYER SUMMARY
+--------------------------------------------------------------------------------
+/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/nn/attention/flex_attention.py:1704: UserWarning: flex_attention called without torch.compile() - this will use an unfused implementation that materializes the full scores matrix instead of generating a fused kernel.
+
+SOLUTION: Use torch.compile(flex_attention)(...)
+
+If you want to debug your score_mod/mask_mod, you can set:
+torch.nn.attention.flex_attention._FLEX_ATTENTION_DISABLE_COMPILE_DEBUG = True
+
+This will allow you to use print statements or breakpoints. Note: This doesn't work with the backwards pass and may produce incorrect results.
+  _warn_once(
+=================================================================================================================================================
+Layer (type:depth-idx)                        Input Shape               Output Shape              Param #                   Trainable
+=================================================================================================================================================
+Qwen3ForCausalLM                              [2, 8192]                 --                        --                        True
+├─Qwen3Model: 1-1                             --                        [2, 8192, 768]            --                        True
+│    └─Embedding: 2-1                         [2, 8192]                 [2, 8192, 768]            116,686,848               True
+│    └─ModuleList: 2-2                        --                        --                        --                        True
+│    │    └─Qwen3DecoderLayer: 3-1            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-2            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-3            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-4            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-5            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-6            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-7            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-8            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-9            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-10           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-11           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-12           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-13           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-14           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-15           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-16           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-17           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-18           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-19           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-20           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-21           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-22           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-23           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-24           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-25           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-26           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-27           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-28           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    └─Qwen3RMSNorm: 2-3                      [2, 8192, 768]            [2, 8192, 768]            768                       True
+├─Linear: 1-2                                 [2, 8192, 768]            [2, 8192, 151936]         116,686,848               True
+=================================================================================================================================================
+Total params: 475,645,696
+Trainable params: 475,645,696
+Non-trainable params: 0
+Total mult-adds (Units.MEGABYTES): 951.29
+=================================================================================================================================================
+Input size (MB): 0.13
+Forward/backward pass size (MB): 31197.23
+Params size (MB): 951.29
+Estimated Total Size (MB): 32148.66
+=================================================================================================================================================
+  NOTE: torchinfo double-counts tied embed_tokens/lm_head weight (116,686,848 params). True unique parameter count: 358,958,848 (359.0M)
+================================================================================
+
+Loading dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Using pre-tokenized dataset (no runtime tokenization)
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Loaded 12748481 pre-tokenized examples (stored_len=2048, pack_factor=4, packed_examples=3187120)
+Dataset split: 3027764 train, 159356 validation
+Training examples (packed, train split):  3027764
+Validation examples (packed, val split):  159356
+Per-device batch size:                    2
+World size (GPUs):                        1
+Gradient accumulation steps:              10
+Effective batch size (seqs/step):         20
+Sequence length:                          8192
+Tokens per optimizer step:                163,840
+Micro-batches per epoch (per rank):       1513882
+Optimizer steps per epoch:                151388
+Total optimizer steps:                    44000  (capped by max_steps; 0.29 epochs)
+wandb: Tracking run with wandb version 0.26.1
+wandb: W&B syncing is set to `offline` in this directory. Run `wandb online` or set WANDB_MODE=online to enable cloud syncing.
+wandb: Run data is saved locally in /mloscratch/homes/navasard/sparse/SparseAtt/wandb/offline-run-20260921_131249-xc6c08ps
+Starting training...
+Switching attention mode for step 1: sparse_stochastic (configured sparse attention)
+Step 100/44000 | Loss: 10.6289 | LR: 6.67e-05 | 1615.0ms/step | 12.4 samples/s
+Step 200/44000 | Loss: 7.3292 | LR: 1.33e-04 | 1475.3ms/step | 13.6 samples/s
+Validation at step 250: Loss = 6.3164, Perplexity = 553.56
+Step 300/44000 | Loss: 6.3471 | LR: 2.00e-04 | 1767.5ms/step | 11.3 samples/s
+Step 400/44000 | Loss: 5.8670 | LR: 2.67e-04 | 1484.2ms/step | 13.5 samples/s
+Step 500/44000 | Loss: 5.4703 | LR: 3.33e-04 | 1484.6ms/step | 13.5 samples/s
+Validation at step 500: Loss = 5.2970, Perplexity = 199.74
+Step 600/44000 | Loss: 5.1419 | LR: 4.00e-04 | 1722.4ms/step | 11.6 samples/s
+Step 700/44000 | Loss: 4.8627 | LR: 4.67e-04 | 1484.0ms/step | 13.5 samples/s
+Validation at step 750: Loss = 4.6010, Perplexity = 99.58
+Step 800/44000 | Loss: 4.5920 | LR: 5.00e-04 | 1722.0ms/step | 11.6 samples/s
+Step 900/44000 | Loss: 4.3801 | LR: 5.00e-04 | 1484.8ms/step | 13.5 samples/s

epfl_final/logs/ALIBI_20260921/restart_20260921_131155/seed47.pid

@@ -0,0 +1 @@
+358

epfl_final/logs/ALIBI_20260921/resume5_20260921_154738/seed42.log

@@ -0,0 +1,195 @@
+Auto-resuming from latest checkpoint: /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-2000
+Resuming: Using existing output directory /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951
+Running config:
+{
+  "model_path": "/mloscratch/homes/navasard/qwen3_06b_base",
+  "model_dtype": "bfloat16",
+  "init_from_scratch": true,
+  "train_path": "/mloscratch/homes/navasard/data/fineweb_large/tokenized",
+  "val_path": null,
+  "max_seq_length": 8192,
+  "num_workers": 8,
+  "output_dir": "/mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951",
+  "resume_from_checkpoint": "/mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-2000",
+  "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": 42,
+  "verbose": true,
+  "checkpoint_activations": false,
+  "resume_micro_step_override": null,
+  "strategy": "ddp",
+  "mixed_precision": "bf16",
+  "wandb_enabled": true,
+  "wandb_project": "stochastic_sparsity",
+  "wandb_name": "365m alibi 8k s42",
+  "wandb_tags": [
+    "qwen3",
+    "365m",
+    "8k",
+    "scaleladder",
+    "alibi",
+    "seed42"
+  ],
+  "hidden_size": 768,
+  "intermediate_size": 3072,
+  "num_attention_heads": 6,
+  "num_key_value_heads": 2,
+  "num_hidden_layers": 28,
+  "head_dim": null,
+  "attention_dropout": null,
+  "rope_train_position_scale_set": null,
+  "positional": "alibi",
+  "attention_type": "dense",
+  "sparse_block_size": 128,
+  "sparse_num_random_blocks": 31,
+  "sparse_mask_pool_size": 1,
+  "sparse_val_dense_compare": false,
+  "sparse_layer_indices": null,
+  "sparse_mask_cache_dir": null,
+  "sparse_mask_seed": 0,
+  "sparse_mixed_nrb_set": null,
+  "dense_mode": false,
+  "dense_step_intervals": null
+}
+Loading tokenizer from /mloscratch/homes/navasard/qwen3_06b_base...
+Resuming from checkpoint: /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-2000
+Checkpoint activations: False
+
+
+================================================================================
+MODEL ARCHITECTURE
+================================================================================
+Model config:
+  vocab_size: 151936
+  hidden_size: 768
+  num_hidden_layers: 28
+  num_attention_heads: 6
+  num_key_value_heads: 2
+  head_dim: 128
+  intermediate_size: 3072
+  max_position_embeddings: 32768
+  rope_theta: 1000000
+  rope_train_position_scale_set: None
+  attention_type: dense
+
+Parameter count:
+  Total parameters: 358,958,848
+  Trainable parameters: 358,958,848
+  Non-trainable parameters: 0
+  Model size (BF16): 0.72 GB
+
+--------------------------------------------------------------------------------
+LAYER-BY-LAYER SUMMARY
+--------------------------------------------------------------------------------
+/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/nn/attention/flex_attention.py:1704: UserWarning: flex_attention called without torch.compile() - this will use an unfused implementation that materializes the full scores matrix instead of generating a fused kernel.
+
+SOLUTION: Use torch.compile(flex_attention)(...)
+
+If you want to debug your score_mod/mask_mod, you can set:
+torch.nn.attention.flex_attention._FLEX_ATTENTION_DISABLE_COMPILE_DEBUG = True
+
+This will allow you to use print statements or breakpoints. Note: This doesn't work with the backwards pass and may produce incorrect results.
+  _warn_once(
+=================================================================================================================================================
+Layer (type:depth-idx)                        Input Shape               Output Shape              Param #                   Trainable
+=================================================================================================================================================
+Qwen3ForCausalLM                              [2, 8192]                 --                        --                        True
+├─Qwen3Model: 1-1                             --                        [2, 8192, 768]            --                        True
+│    └─Embedding: 2-1                         [2, 8192]                 [2, 8192, 768]            116,686,848               True
+│    └─ModuleList: 2-2                        --                        --                        --                        True
+│    │    └─Qwen3DecoderLayer: 3-1            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-2            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-3            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-4            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-5            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-6            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-7            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-8            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-9            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-10           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-11           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-12           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-13           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-14           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-15           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-16           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-17           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-18           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-19           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-20           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-21           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-22           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-23           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-24           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-25           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-26           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-27           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-28           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    └─Qwen3RMSNorm: 2-3                      [2, 8192, 768]            [2, 8192, 768]            768                       True
+├─Linear: 1-2                                 [2, 8192, 768]            [2, 8192, 151936]         116,686,848               True
+=================================================================================================================================================
+Total params: 475,645,696
+Trainable params: 475,645,696
+Non-trainable params: 0
+Total mult-adds (Units.MEGABYTES): 951.29
+=================================================================================================================================================
+Input size (MB): 0.13
+Forward/backward pass size (MB): 31197.23
+Params size (MB): 951.29
+Estimated Total Size (MB): 32148.66
+=================================================================================================================================================
+  NOTE: torchinfo double-counts tied embed_tokens/lm_head weight (116,686,848 params). True unique parameter count: 358,958,848 (359.0M)
+================================================================================
+
+Loading dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Using pre-tokenized dataset (no runtime tokenization)
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Loaded 12748481 pre-tokenized examples (stored_len=2048, pack_factor=4, packed_examples=3187120)
+Dataset split: 3027764 train, 159356 validation
+Training examples (packed, train split):  3027764
+Validation examples (packed, val split):  159356
+Per-device batch size:                    2
+World size (GPUs):                        1
+Gradient accumulation steps:              10
+Effective batch size (seqs/step):         20
+Sequence length:                          8192
+Tokens per optimizer step:                163,840
+Micro-batches per epoch (per rank):       1513882
+Optimizer steps per epoch:                151388
+Total optimizer steps:                    44000  (capped by max_steps; 0.29 epochs)
+Loading checkpoint from /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-2000...
+Preserving checkpoint LR state for exact resume: current=[0.0004989701850946613], schedule_base=[0.0005]
+Resumed from checkpoint at step 2000 (epoch 0, micro_step 20000)
+Resuming training from step 2000
+Resuming wandb run nihe6o4e
+wandb: WARNING `resume` will be ignored since W&B syncing is set to `offline`. Starting a new run with run id nihe6o4e.
+wandb: Tracking run with wandb version 0.26.1
+wandb: W&B syncing is set to `offline` in this directory. Run `wandb online` or set WANDB_MODE=online to enable cloud syncing.
+wandb: Run data is saved locally in /mloscratch/homes/navasard/sparse/SparseAtt/wandb/offline-run-20260921_154845-nihe6o4e
+Starting training...
+Validation at step 2000: Loss = 3.6610, Perplexity = 38.90
+Switching attention mode for step 2001: sparse_stochastic (configured sparse attention)
+Step 2100/44000 | Loss: 3.6362 | LR: 4.99e-04 | 2815.9ms/step | 7.1 samples/s
+Step 2200/44000 | Loss: 3.6120 | LR: 4.99e-04 | 1469.7ms/step | 13.6 samples/s
+Validation at step 2250: Loss = 3.6071, Perplexity = 36.86
+Step 2300/44000 | Loss: 3.5862 | LR: 4.98e-04 | 1706.8ms/step | 11.7 samples/s
+Step 2400/44000 | Loss: 3.5683 | LR: 4.98e-04 | 1469.8ms/step | 13.6 samples/s
+Step 2500/44000 | Loss: 3.5451 | LR: 4.98e-04 | 1470.3ms/step | 13.6 samples/s
+Validation at step 2500: Loss = 3.5592, Perplexity = 35.14
+Step 2600/44000 | Loss: 3.5446 | LR: 4.98e-04 | 1707.9ms/step | 11.7 samples/s
+Step 2700/44000 | Loss: 3.5179 | LR: 4.97e-04 | 1470.1ms/step | 13.6 samples/s
+Validation at step 2750: Loss = 3.5190, Perplexity = 33.75
+Step 2800/44000 | Loss: 3.5063 | LR: 4.97e-04 | 1710.3ms/step | 11.7 samples/s
+Step 2900/44000 | Loss: 3.5002 | LR: 4.97e-04 | 1474.0ms/step | 13.6 samples/s

epfl_final/logs/ALIBI_20260921/resume5_20260921_154738/seed42.pid

@@ -0,0 +1 @@
+174

epfl_final/logs/ALIBI_20260921/resume5_20260921_154738/seed43.log

@@ -0,0 +1,195 @@
+Auto-resuming from latest checkpoint: /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-2000
+Resuming: Using existing output directory /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951
+Running config:
+{
+  "model_path": "/mloscratch/homes/navasard/qwen3_06b_base",
+  "model_dtype": "bfloat16",
+  "init_from_scratch": true,
+  "train_path": "/mloscratch/homes/navasard/data/fineweb_large/tokenized",
+  "val_path": null,
+  "max_seq_length": 8192,
+  "num_workers": 8,
+  "output_dir": "/mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951",
+  "resume_from_checkpoint": "/mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-2000",
+  "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": 43,
+  "verbose": true,
+  "checkpoint_activations": false,
+  "resume_micro_step_override": null,
+  "strategy": "ddp",
+  "mixed_precision": "bf16",
+  "wandb_enabled": true,
+  "wandb_project": "stochastic_sparsity",
+  "wandb_name": "365m alibi 8k s43",
+  "wandb_tags": [
+    "qwen3",
+    "365m",
+    "8k",
+    "scaleladder",
+    "alibi",
+    "seed43"
+  ],
+  "hidden_size": 768,
+  "intermediate_size": 3072,
+  "num_attention_heads": 6,
+  "num_key_value_heads": 2,
+  "num_hidden_layers": 28,
+  "head_dim": null,
+  "attention_dropout": null,
+  "rope_train_position_scale_set": null,
+  "positional": "alibi",
+  "attention_type": "dense",
+  "sparse_block_size": 128,
+  "sparse_num_random_blocks": 31,
+  "sparse_mask_pool_size": 1,
+  "sparse_val_dense_compare": false,
+  "sparse_layer_indices": null,
+  "sparse_mask_cache_dir": null,
+  "sparse_mask_seed": 0,
+  "sparse_mixed_nrb_set": null,
+  "dense_mode": false,
+  "dense_step_intervals": null
+}
+Loading tokenizer from /mloscratch/homes/navasard/qwen3_06b_base...
+Resuming from checkpoint: /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-2000
+Checkpoint activations: False
+
+
+================================================================================
+MODEL ARCHITECTURE
+================================================================================
+Model config:
+  vocab_size: 151936
+  hidden_size: 768
+  num_hidden_layers: 28
+  num_attention_heads: 6
+  num_key_value_heads: 2
+  head_dim: 128
+  intermediate_size: 3072
+  max_position_embeddings: 32768
+  rope_theta: 1000000
+  rope_train_position_scale_set: None
+  attention_type: dense
+
+Parameter count:
+  Total parameters: 358,958,848
+  Trainable parameters: 358,958,848
+  Non-trainable parameters: 0
+  Model size (BF16): 0.72 GB
+
+--------------------------------------------------------------------------------
+LAYER-BY-LAYER SUMMARY
+--------------------------------------------------------------------------------
+/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/nn/attention/flex_attention.py:1704: UserWarning: flex_attention called without torch.compile() - this will use an unfused implementation that materializes the full scores matrix instead of generating a fused kernel.
+
+SOLUTION: Use torch.compile(flex_attention)(...)
+
+If you want to debug your score_mod/mask_mod, you can set:
+torch.nn.attention.flex_attention._FLEX_ATTENTION_DISABLE_COMPILE_DEBUG = True
+
+This will allow you to use print statements or breakpoints. Note: This doesn't work with the backwards pass and may produce incorrect results.
+  _warn_once(
+=================================================================================================================================================
+Layer (type:depth-idx)                        Input Shape               Output Shape              Param #                   Trainable
+=================================================================================================================================================
+Qwen3ForCausalLM                              [2, 8192]                 --                        --                        True
+├─Qwen3Model: 1-1                             --                        [2, 8192, 768]            --                        True
+│    └─Embedding: 2-1                         [2, 8192]                 [2, 8192, 768]            116,686,848               True
+│    └─ModuleList: 2-2                        --                        --                        --                        True
+│    │    └─Qwen3DecoderLayer: 3-1            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-2            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-3            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-4            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-5            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-6            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-7            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-8            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-9            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-10           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-11           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-12           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-13           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-14           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-15           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-16           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-17           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-18           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-19           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-20           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-21           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-22           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-23           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-24           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-25           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-26           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-27           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-28           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    └─Qwen3RMSNorm: 2-3                      [2, 8192, 768]            [2, 8192, 768]            768                       True
+├─Linear: 1-2                                 [2, 8192, 768]            [2, 8192, 151936]         116,686,848               True
+=================================================================================================================================================
+Total params: 475,645,696
+Trainable params: 475,645,696
+Non-trainable params: 0
+Total mult-adds (Units.MEGABYTES): 951.29
+=================================================================================================================================================
+Input size (MB): 0.13
+Forward/backward pass size (MB): 31197.23
+Params size (MB): 951.29
+Estimated Total Size (MB): 32148.66
+=================================================================================================================================================
+  NOTE: torchinfo double-counts tied embed_tokens/lm_head weight (116,686,848 params). True unique parameter count: 358,958,848 (359.0M)
+================================================================================
+
+Loading dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Using pre-tokenized dataset (no runtime tokenization)
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Loaded 12748481 pre-tokenized examples (stored_len=2048, pack_factor=4, packed_examples=3187120)
+Dataset split: 3027764 train, 159356 validation
+Training examples (packed, train split):  3027764
+Validation examples (packed, val split):  159356
+Per-device batch size:                    2
+World size (GPUs):                        1
+Gradient accumulation steps:              10
+Effective batch size (seqs/step):         20
+Sequence length:                          8192
+Tokens per optimizer step:                163,840
+Micro-batches per epoch (per rank):       1513882
+Optimizer steps per epoch:                151388
+Total optimizer steps:                    44000  (capped by max_steps; 0.29 epochs)
+Loading checkpoint from /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-2000...
+Preserving checkpoint LR state for exact resume: current=[0.0004989701850946613], schedule_base=[0.0005]
+Resumed from checkpoint at step 2000 (epoch 0, micro_step 20000)
+Resuming training from step 2000
+Resuming wandb run zgrsvkbi
+wandb: WARNING `resume` will be ignored since W&B syncing is set to `offline`. Starting a new run with run id zgrsvkbi.
+wandb: Tracking run with wandb version 0.26.1
+wandb: W&B syncing is set to `offline` in this directory. Run `wandb online` or set WANDB_MODE=online to enable cloud syncing.
+wandb: Run data is saved locally in /mloscratch/homes/navasard/sparse/SparseAtt/wandb/offline-run-20260921_154845-zgrsvkbi
+Starting training...
+Validation at step 2000: Loss = 3.6382, Perplexity = 38.02
+Switching attention mode for step 2001: sparse_stochastic (configured sparse attention)
+Step 2100/44000 | Loss: 3.6414 | LR: 4.99e-04 | 2865.2ms/step | 7.0 samples/s
+Step 2200/44000 | Loss: 3.6026 | LR: 4.99e-04 | 1519.4ms/step | 13.2 samples/s
+Validation at step 2250: Loss = 3.5841, Perplexity = 36.02
+Step 2300/44000 | Loss: 3.5932 | LR: 4.98e-04 | 1765.3ms/step | 11.3 samples/s
+Step 2400/44000 | Loss: 3.5771 | LR: 4.98e-04 | 1522.5ms/step | 13.1 samples/s
+Step 2500/44000 | Loss: 3.5544 | LR: 4.98e-04 | 1522.9ms/step | 13.1 samples/s
+Validation at step 2500: Loss = 3.5366, Perplexity = 34.35
+Step 2600/44000 | Loss: 3.5414 | LR: 4.98e-04 | 1766.7ms/step | 11.3 samples/s
+Step 2700/44000 | Loss: 3.5250 | LR: 4.97e-04 | 1520.7ms/step | 13.2 samples/s
+Validation at step 2750: Loss = 3.4984, Perplexity = 33.06
+Step 2800/44000 | Loss: 3.5066 | LR: 4.97e-04 | 1765.7ms/step | 11.3 samples/s
+Step 2900/44000 | Loss: 3.4950 | LR: 4.97e-04 | 1523.7ms/step | 13.1 samples/s

epfl_final/logs/ALIBI_20260921/resume5_20260921_154738/seed43.pid

@@ -0,0 +1 @@
+179

epfl_final/logs/ALIBI_20260921/resume5_20260921_154738/seed44.log

@@ -0,0 +1,195 @@
+Auto-resuming from latest checkpoint: /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-2000
+Resuming: Using existing output directory /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951
+Running config:
+{
+  "model_path": "/mloscratch/homes/navasard/qwen3_06b_base",
+  "model_dtype": "bfloat16",
+  "init_from_scratch": true,
+  "train_path": "/mloscratch/homes/navasard/data/fineweb_large/tokenized",
+  "val_path": null,
+  "max_seq_length": 8192,
+  "num_workers": 8,
+  "output_dir": "/mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951",
+  "resume_from_checkpoint": "/mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-2000",
+  "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": 44,
+  "verbose": true,
+  "checkpoint_activations": false,
+  "resume_micro_step_override": null,
+  "strategy": "ddp",
+  "mixed_precision": "bf16",
+  "wandb_enabled": true,
+  "wandb_project": "stochastic_sparsity",
+  "wandb_name": "365m alibi 8k s44",
+  "wandb_tags": [
+    "qwen3",
+    "365m",
+    "8k",
+    "scaleladder",
+    "alibi",
+    "seed44"
+  ],
+  "hidden_size": 768,
+  "intermediate_size": 3072,
+  "num_attention_heads": 6,
+  "num_key_value_heads": 2,
+  "num_hidden_layers": 28,
+  "head_dim": null,
+  "attention_dropout": null,
+  "rope_train_position_scale_set": null,
+  "positional": "alibi",
+  "attention_type": "dense",
+  "sparse_block_size": 128,
+  "sparse_num_random_blocks": 31,
+  "sparse_mask_pool_size": 1,
+  "sparse_val_dense_compare": false,
+  "sparse_layer_indices": null,
+  "sparse_mask_cache_dir": null,
+  "sparse_mask_seed": 0,
+  "sparse_mixed_nrb_set": null,
+  "dense_mode": false,
+  "dense_step_intervals": null
+}
+Loading tokenizer from /mloscratch/homes/navasard/qwen3_06b_base...
+Resuming from checkpoint: /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-2000
+Checkpoint activations: False
+
+
+================================================================================
+MODEL ARCHITECTURE
+================================================================================
+Model config:
+  vocab_size: 151936
+  hidden_size: 768
+  num_hidden_layers: 28
+  num_attention_heads: 6
+  num_key_value_heads: 2
+  head_dim: 128
+  intermediate_size: 3072
+  max_position_embeddings: 32768
+  rope_theta: 1000000
+  rope_train_position_scale_set: None
+  attention_type: dense
+
+Parameter count:
+  Total parameters: 358,958,848
+  Trainable parameters: 358,958,848
+  Non-trainable parameters: 0
+  Model size (BF16): 0.72 GB
+
+--------------------------------------------------------------------------------
+LAYER-BY-LAYER SUMMARY
+--------------------------------------------------------------------------------
+/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/nn/attention/flex_attention.py:1704: UserWarning: flex_attention called without torch.compile() - this will use an unfused implementation that materializes the full scores matrix instead of generating a fused kernel.
+
+SOLUTION: Use torch.compile(flex_attention)(...)
+
+If you want to debug your score_mod/mask_mod, you can set:
+torch.nn.attention.flex_attention._FLEX_ATTENTION_DISABLE_COMPILE_DEBUG = True
+
+This will allow you to use print statements or breakpoints. Note: This doesn't work with the backwards pass and may produce incorrect results.
+  _warn_once(
+=================================================================================================================================================
+Layer (type:depth-idx)                        Input Shape               Output Shape              Param #                   Trainable
+=================================================================================================================================================
+Qwen3ForCausalLM                              [2, 8192]                 --                        --                        True
+├─Qwen3Model: 1-1                             --                        [2, 8192, 768]            --                        True
+│    └─Embedding: 2-1                         [2, 8192]                 [2, 8192, 768]            116,686,848               True
+│    └─ModuleList: 2-2                        --                        --                        --                        True
+│    │    └─Qwen3DecoderLayer: 3-1            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-2            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-3            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-4            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-5            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-6            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-7            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-8            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-9            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-10           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-11           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-12           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-13           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-14           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-15           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-16           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-17           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-18           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-19           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-20           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-21           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-22           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-23           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-24           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-25           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-26           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-27           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-28           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    └─Qwen3RMSNorm: 2-3                      [2, 8192, 768]            [2, 8192, 768]            768                       True
+├─Linear: 1-2                                 [2, 8192, 768]            [2, 8192, 151936]         116,686,848               True
+=================================================================================================================================================
+Total params: 475,645,696
+Trainable params: 475,645,696
+Non-trainable params: 0
+Total mult-adds (Units.MEGABYTES): 951.29
+=================================================================================================================================================
+Input size (MB): 0.13
+Forward/backward pass size (MB): 31197.23
+Params size (MB): 951.29
+Estimated Total Size (MB): 32148.66
+=================================================================================================================================================
+  NOTE: torchinfo double-counts tied embed_tokens/lm_head weight (116,686,848 params). True unique parameter count: 358,958,848 (359.0M)
+================================================================================
+
+Loading dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Using pre-tokenized dataset (no runtime tokenization)
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Loaded 12748481 pre-tokenized examples (stored_len=2048, pack_factor=4, packed_examples=3187120)
+Dataset split: 3027764 train, 159356 validation
+Training examples (packed, train split):  3027764
+Validation examples (packed, val split):  159356
+Per-device batch size:                    2
+World size (GPUs):                        1
+Gradient accumulation steps:              10
+Effective batch size (seqs/step):         20
+Sequence length:                          8192
+Tokens per optimizer step:                163,840
+Micro-batches per epoch (per rank):       1513882
+Optimizer steps per epoch:                151388
+Total optimizer steps:                    44000  (capped by max_steps; 0.29 epochs)
+Loading checkpoint from /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-2000...
+Preserving checkpoint LR state for exact resume: current=[0.0004989701850946613], schedule_base=[0.0005]
+Resumed from checkpoint at step 2000 (epoch 0, micro_step 20000)
+Resuming training from step 2000
+Resuming wandb run kbn66izc
+wandb: WARNING `resume` will be ignored since W&B syncing is set to `offline`. Starting a new run with run id kbn66izc.
+wandb: Tracking run with wandb version 0.26.1
+wandb: W&B syncing is set to `offline` in this directory. Run `wandb online` or set WANDB_MODE=online to enable cloud syncing.
+wandb: Run data is saved locally in /mloscratch/homes/navasard/sparse/SparseAtt/wandb/offline-run-20260921_154846-kbn66izc
+Starting training...
+Validation at step 2000: Loss = 3.6519, Perplexity = 38.55
+Switching attention mode for step 2001: sparse_stochastic (configured sparse attention)
+Step 2100/44000 | Loss: 3.6295 | LR: 4.99e-04 | 2841.6ms/step | 7.0 samples/s
+Step 2200/44000 | Loss: 3.6040 | LR: 4.99e-04 | 1498.5ms/step | 13.3 samples/s
+Validation at step 2250: Loss = 3.5921, Perplexity = 36.31
+Step 2300/44000 | Loss: 3.5890 | LR: 4.98e-04 | 1740.5ms/step | 11.5 samples/s
+Step 2400/44000 | Loss: 3.5556 | LR: 4.98e-04 | 1500.3ms/step | 13.3 samples/s
+Step 2500/44000 | Loss: 3.5480 | LR: 4.98e-04 | 1500.4ms/step | 13.3 samples/s
+Validation at step 2500: Loss = 3.5482, Perplexity = 34.75
+Step 2600/44000 | Loss: 3.5346 | LR: 4.98e-04 | 1741.0ms/step | 11.5 samples/s
+Step 2700/44000 | Loss: 3.5145 | LR: 4.97e-04 | 1500.0ms/step | 13.3 samples/s
+Validation at step 2750: Loss = 3.5061, Perplexity = 33.32
+Step 2800/44000 | Loss: 3.5066 | LR: 4.97e-04 | 1741.3ms/step | 11.5 samples/s
+Step 2900/44000 | Loss: 3.4976 | LR: 4.97e-04 | 1500.2ms/step | 13.3 samples/s

epfl_final/logs/ALIBI_20260921/resume5_20260921_154738/seed44.pid

@@ -0,0 +1 @@
+184

epfl_final/logs/ALIBI_20260921/resume5_20260921_154738/seed45.log

@@ -0,0 +1,195 @@
+Auto-resuming from latest checkpoint: /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s45/run_20260921_142951/checkpoint-2000
+Resuming: Using existing output directory /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s45/run_20260921_142951
+Running config:
+{
+  "model_path": "/mloscratch/homes/navasard/qwen3_06b_base",
+  "model_dtype": "bfloat16",
+  "init_from_scratch": true,
+  "train_path": "/mloscratch/homes/navasard/data/fineweb_large/tokenized",
+  "val_path": null,
+  "max_seq_length": 8192,
+  "num_workers": 8,
+  "output_dir": "/mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s45/run_20260921_142951",
+  "resume_from_checkpoint": "/mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s45/run_20260921_142951/checkpoint-2000",
+  "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": 45,
+  "verbose": true,
+  "checkpoint_activations": false,
+  "resume_micro_step_override": null,
+  "strategy": "ddp",
+  "mixed_precision": "bf16",
+  "wandb_enabled": true,
+  "wandb_project": "stochastic_sparsity",
+  "wandb_name": "365m alibi 8k s45",
+  "wandb_tags": [
+    "qwen3",
+    "365m",
+    "8k",
+    "scaleladder",
+    "alibi",
+    "seed45"
+  ],
+  "hidden_size": 768,
+  "intermediate_size": 3072,
+  "num_attention_heads": 6,
+  "num_key_value_heads": 2,
+  "num_hidden_layers": 28,
+  "head_dim": null,
+  "attention_dropout": null,
+  "rope_train_position_scale_set": null,
+  "positional": "alibi",
+  "attention_type": "dense",
+  "sparse_block_size": 128,
+  "sparse_num_random_blocks": 31,
+  "sparse_mask_pool_size": 1,
+  "sparse_val_dense_compare": false,
+  "sparse_layer_indices": null,
+  "sparse_mask_cache_dir": null,
+  "sparse_mask_seed": 0,
+  "sparse_mixed_nrb_set": null,
+  "dense_mode": false,
+  "dense_step_intervals": null
+}
+Loading tokenizer from /mloscratch/homes/navasard/qwen3_06b_base...
+Resuming from checkpoint: /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s45/run_20260921_142951/checkpoint-2000
+Checkpoint activations: False
+
+
+================================================================================
+MODEL ARCHITECTURE
+================================================================================
+Model config:
+  vocab_size: 151936
+  hidden_size: 768
+  num_hidden_layers: 28
+  num_attention_heads: 6
+  num_key_value_heads: 2
+  head_dim: 128
+  intermediate_size: 3072
+  max_position_embeddings: 32768
+  rope_theta: 1000000
+  rope_train_position_scale_set: None
+  attention_type: dense
+
+Parameter count:
+  Total parameters: 358,958,848
+  Trainable parameters: 358,958,848
+  Non-trainable parameters: 0
+  Model size (BF16): 0.72 GB
+
+--------------------------------------------------------------------------------
+LAYER-BY-LAYER SUMMARY
+--------------------------------------------------------------------------------
+/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/nn/attention/flex_attention.py:1704: UserWarning: flex_attention called without torch.compile() - this will use an unfused implementation that materializes the full scores matrix instead of generating a fused kernel.
+
+SOLUTION: Use torch.compile(flex_attention)(...)
+
+If you want to debug your score_mod/mask_mod, you can set:
+torch.nn.attention.flex_attention._FLEX_ATTENTION_DISABLE_COMPILE_DEBUG = True
+
+This will allow you to use print statements or breakpoints. Note: This doesn't work with the backwards pass and may produce incorrect results.
+  _warn_once(
+=================================================================================================================================================
+Layer (type:depth-idx)                        Input Shape               Output Shape              Param #                   Trainable
+=================================================================================================================================================
+Qwen3ForCausalLM                              [2, 8192]                 --                        --                        True
+├─Qwen3Model: 1-1                             --                        [2, 8192, 768]            --                        True
+│    └─Embedding: 2-1                         [2, 8192]                 [2, 8192, 768]            116,686,848               True
+│    └─ModuleList: 2-2                        --                        --                        --                        True
+│    │    └─Qwen3DecoderLayer: 3-1            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-2            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-3            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-4            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-5            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-6            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-7            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-8            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-9            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-10           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-11           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-12           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-13           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-14           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-15           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-16           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-17           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-18           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-19           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-20           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-21           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-22           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-23           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-24           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-25           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-26           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-27           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-28           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    └─Qwen3RMSNorm: 2-3                      [2, 8192, 768]            [2, 8192, 768]            768                       True
+├─Linear: 1-2                                 [2, 8192, 768]            [2, 8192, 151936]         116,686,848               True
+=================================================================================================================================================
+Total params: 475,645,696
+Trainable params: 475,645,696
+Non-trainable params: 0
+Total mult-adds (Units.MEGABYTES): 951.29
+=================================================================================================================================================
+Input size (MB): 0.13
+Forward/backward pass size (MB): 31197.23
+Params size (MB): 951.29
+Estimated Total Size (MB): 32148.66
+=================================================================================================================================================
+  NOTE: torchinfo double-counts tied embed_tokens/lm_head weight (116,686,848 params). True unique parameter count: 358,958,848 (359.0M)
+================================================================================
+
+Loading dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Using pre-tokenized dataset (no runtime tokenization)
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Loaded 12748481 pre-tokenized examples (stored_len=2048, pack_factor=4, packed_examples=3187120)
+Dataset split: 3027764 train, 159356 validation
+Training examples (packed, train split):  3027764
+Validation examples (packed, val split):  159356
+Per-device batch size:                    2
+World size (GPUs):                        1
+Gradient accumulation steps:              10
+Effective batch size (seqs/step):         20
+Sequence length:                          8192
+Tokens per optimizer step:                163,840
+Micro-batches per epoch (per rank):       1513882
+Optimizer steps per epoch:                151388
+Total optimizer steps:                    44000  (capped by max_steps; 0.29 epochs)
+Loading checkpoint from /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s45/run_20260921_142951/checkpoint-2000...
+Preserving checkpoint LR state for exact resume: current=[0.0004989701850946613], schedule_base=[0.0005]
+Resumed from checkpoint at step 2000 (epoch 0, micro_step 20000)
+Resuming training from step 2000
+Resuming wandb run ghfbqjah
+wandb: WARNING `resume` will be ignored since W&B syncing is set to `offline`. Starting a new run with run id ghfbqjah.
+wandb: Tracking run with wandb version 0.26.1
+wandb: W&B syncing is set to `offline` in this directory. Run `wandb online` or set WANDB_MODE=online to enable cloud syncing.
+wandb: Run data is saved locally in /mloscratch/homes/navasard/sparse/SparseAtt/wandb/offline-run-20260921_154845-ghfbqjah
+Starting training...
+Validation at step 2000: Loss = 3.6535, Perplexity = 38.61
+Switching attention mode for step 2001: sparse_stochastic (configured sparse attention)
+Step 2100/44000 | Loss: 3.6382 | LR: 4.99e-04 | 2813.4ms/step | 7.1 samples/s
+Step 2200/44000 | Loss: 3.6122 | LR: 4.99e-04 | 1477.4ms/step | 13.5 samples/s
+Validation at step 2250: Loss = 3.5960, Perplexity = 36.45
+Step 2300/44000 | Loss: 3.5911 | LR: 4.98e-04 | 1715.2ms/step | 11.7 samples/s
+Step 2400/44000 | Loss: 3.5725 | LR: 4.98e-04 | 1475.7ms/step | 13.6 samples/s
+Step 2500/44000 | Loss: 3.5568 | LR: 4.98e-04 | 1476.1ms/step | 13.5 samples/s
+Validation at step 2500: Loss = 3.5512, Perplexity = 34.86
+Step 2600/44000 | Loss: 3.5353 | LR: 4.98e-04 | 1712.4ms/step | 11.7 samples/s
+Step 2700/44000 | Loss: 3.5181 | LR: 4.97e-04 | 1475.0ms/step | 13.6 samples/s
+Validation at step 2750: Loss = 3.5136, Perplexity = 33.57
+Step 2800/44000 | Loss: 3.5075 | LR: 4.97e-04 | 1714.3ms/step | 11.7 samples/s
+Step 2900/44000 | Loss: 3.4935 | LR: 4.97e-04 | 1476.0ms/step | 13.6 samples/s

epfl_final/logs/ALIBI_20260921/resume5_20260921_154738/seed45.pid

@@ -0,0 +1 @@
+189

epfl_final/logs/ALIBI_20260921/resume5_20260921_154738/seed46.log

@@ -0,0 +1,195 @@
+Auto-resuming from latest checkpoint: /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s46/run_20260921_142951/checkpoint-2000
+Resuming: Using existing output directory /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s46/run_20260921_142951
+Running config:
+{
+  "model_path": "/mloscratch/homes/navasard/qwen3_06b_base",
+  "model_dtype": "bfloat16",
+  "init_from_scratch": true,
+  "train_path": "/mloscratch/homes/navasard/data/fineweb_large/tokenized",
+  "val_path": null,
+  "max_seq_length": 8192,
+  "num_workers": 8,
+  "output_dir": "/mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s46/run_20260921_142951",
+  "resume_from_checkpoint": "/mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s46/run_20260921_142951/checkpoint-2000",
+  "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": 46,
+  "verbose": true,
+  "checkpoint_activations": false,
+  "resume_micro_step_override": null,
+  "strategy": "ddp",
+  "mixed_precision": "bf16",
+  "wandb_enabled": true,
+  "wandb_project": "stochastic_sparsity",
+  "wandb_name": "365m alibi 8k s46",
+  "wandb_tags": [
+    "qwen3",
+    "365m",
+    "8k",
+    "scaleladder",
+    "alibi",
+    "seed46"
+  ],
+  "hidden_size": 768,
+  "intermediate_size": 3072,
+  "num_attention_heads": 6,
+  "num_key_value_heads": 2,
+  "num_hidden_layers": 28,
+  "head_dim": null,
+  "attention_dropout": null,
+  "rope_train_position_scale_set": null,
+  "positional": "alibi",
+  "attention_type": "dense",
+  "sparse_block_size": 128,
+  "sparse_num_random_blocks": 31,
+  "sparse_mask_pool_size": 1,
+  "sparse_val_dense_compare": false,
+  "sparse_layer_indices": null,
+  "sparse_mask_cache_dir": null,
+  "sparse_mask_seed": 0,
+  "sparse_mixed_nrb_set": null,
+  "dense_mode": false,
+  "dense_step_intervals": null
+}
+Loading tokenizer from /mloscratch/homes/navasard/qwen3_06b_base...
+Resuming from checkpoint: /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s46/run_20260921_142951/checkpoint-2000
+Checkpoint activations: False
+
+
+================================================================================
+MODEL ARCHITECTURE
+================================================================================
+Model config:
+  vocab_size: 151936
+  hidden_size: 768
+  num_hidden_layers: 28
+  num_attention_heads: 6
+  num_key_value_heads: 2
+  head_dim: 128
+  intermediate_size: 3072
+  max_position_embeddings: 32768
+  rope_theta: 1000000
+  rope_train_position_scale_set: None
+  attention_type: dense
+
+Parameter count:
+  Total parameters: 358,958,848
+  Trainable parameters: 358,958,848
+  Non-trainable parameters: 0
+  Model size (BF16): 0.72 GB
+
+--------------------------------------------------------------------------------
+LAYER-BY-LAYER SUMMARY
+--------------------------------------------------------------------------------
+/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/nn/attention/flex_attention.py:1704: UserWarning: flex_attention called without torch.compile() - this will use an unfused implementation that materializes the full scores matrix instead of generating a fused kernel.
+
+SOLUTION: Use torch.compile(flex_attention)(...)
+
+If you want to debug your score_mod/mask_mod, you can set:
+torch.nn.attention.flex_attention._FLEX_ATTENTION_DISABLE_COMPILE_DEBUG = True
+
+This will allow you to use print statements or breakpoints. Note: This doesn't work with the backwards pass and may produce incorrect results.
+  _warn_once(
+=================================================================================================================================================
+Layer (type:depth-idx)                        Input Shape               Output Shape              Param #                   Trainable
+=================================================================================================================================================
+Qwen3ForCausalLM                              [2, 8192]                 --                        --                        True
+├─Qwen3Model: 1-1                             --                        [2, 8192, 768]            --                        True
+│    └─Embedding: 2-1                         [2, 8192]                 [2, 8192, 768]            116,686,848               True
+│    └─ModuleList: 2-2                        --                        --                        --                        True
+│    │    └─Qwen3DecoderLayer: 3-1            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-2            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-3            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-4            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-5            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-6            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-7            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-8            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-9            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-10           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-11           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-12           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-13           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-14           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-15           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-16           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-17           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-18           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-19           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-20           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-21           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-22           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-23           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-24           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-25           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-26           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-27           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-28           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    └─Qwen3RMSNorm: 2-3                      [2, 8192, 768]            [2, 8192, 768]            768                       True
+├─Linear: 1-2                                 [2, 8192, 768]            [2, 8192, 151936]         116,686,848               True
+=================================================================================================================================================
+Total params: 475,645,696
+Trainable params: 475,645,696
+Non-trainable params: 0
+Total mult-adds (Units.MEGABYTES): 951.29
+=================================================================================================================================================
+Input size (MB): 0.13
+Forward/backward pass size (MB): 31197.23
+Params size (MB): 951.29
+Estimated Total Size (MB): 32148.66
+=================================================================================================================================================
+  NOTE: torchinfo double-counts tied embed_tokens/lm_head weight (116,686,848 params). True unique parameter count: 358,958,848 (359.0M)
+================================================================================
+
+Loading dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Using pre-tokenized dataset (no runtime tokenization)
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Loaded 12748481 pre-tokenized examples (stored_len=2048, pack_factor=4, packed_examples=3187120)
+Dataset split: 3027764 train, 159356 validation
+Training examples (packed, train split):  3027764
+Validation examples (packed, val split):  159356
+Per-device batch size:                    2
+World size (GPUs):                        1
+Gradient accumulation steps:              10
+Effective batch size (seqs/step):         20
+Sequence length:                          8192
+Tokens per optimizer step:                163,840
+Micro-batches per epoch (per rank):       1513882
+Optimizer steps per epoch:                151388
+Total optimizer steps:                    44000  (capped by max_steps; 0.29 epochs)
+Loading checkpoint from /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s46/run_20260921_142951/checkpoint-2000...
+Preserving checkpoint LR state for exact resume: current=[0.0004989701850946613], schedule_base=[0.0005]
+Resumed from checkpoint at step 2000 (epoch 0, micro_step 20000)
+Resuming training from step 2000
+Resuming wandb run wk3rt7ae
+wandb: WARNING `resume` will be ignored since W&B syncing is set to `offline`. Starting a new run with run id wk3rt7ae.
+wandb: Tracking run with wandb version 0.26.1
+wandb: W&B syncing is set to `offline` in this directory. Run `wandb online` or set WANDB_MODE=online to enable cloud syncing.
+wandb: Run data is saved locally in /mloscratch/homes/navasard/sparse/SparseAtt/wandb/offline-run-20260921_154845-wk3rt7ae
+Starting training...
+Validation at step 2000: Loss = 3.6463, Perplexity = 38.33
+Switching attention mode for step 2001: sparse_stochastic (configured sparse attention)
+Step 2100/44000 | Loss: 3.6305 | LR: 4.99e-04 | 2833.0ms/step | 7.1 samples/s
+Step 2200/44000 | Loss: 3.6098 | LR: 4.99e-04 | 1495.0ms/step | 13.4 samples/s
+Validation at step 2250: Loss = 3.5927, Perplexity = 36.33
+Step 2300/44000 | Loss: 3.5900 | LR: 4.98e-04 | 1735.9ms/step | 11.5 samples/s
+Step 2400/44000 | Loss: 3.5647 | LR: 4.98e-04 | 1497.0ms/step | 13.4 samples/s
+Step 2500/44000 | Loss: 3.5608 | LR: 4.98e-04 | 1497.7ms/step | 13.4 samples/s
+Validation at step 2500: Loss = 3.5437, Perplexity = 34.59
+Step 2600/44000 | Loss: 3.5374 | LR: 4.98e-04 | 1736.7ms/step | 11.5 samples/s
+Step 2700/44000 | Loss: 3.5225 | LR: 4.97e-04 | 1497.3ms/step | 13.4 samples/s
+Validation at step 2750: Loss = 3.5048, Perplexity = 33.28
+Step 2800/44000 | Loss: 3.4900 | LR: 4.97e-04 | 1738.6ms/step | 11.5 samples/s
+Step 2900/44000 | Loss: 3.4879 | LR: 4.97e-04 | 1497.5ms/step | 13.4 samples/s

epfl_final/logs/ALIBI_20260921/resume5_20260921_154738/seed46.pid

@@ -0,0 +1 @@
+194

epfl_final/logs/ALIBI_20260921/resume_20260921_142934/seed42.log

@@ -0,0 +1,205 @@
+New run: Checkpoints will be saved to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951
+Running config:
+{
+  "model_path": "/mloscratch/homes/navasard/qwen3_06b_base",
+  "model_dtype": "bfloat16",
+  "init_from_scratch": true,
+  "train_path": "/mloscratch/homes/navasard/data/fineweb_large/tokenized",
+  "val_path": null,
+  "max_seq_length": 8192,
+  "num_workers": 8,
+  "output_dir": "/mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951",
+  "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": 42,
+  "verbose": true,
+  "checkpoint_activations": false,
+  "resume_micro_step_override": null,
+  "strategy": "ddp",
+  "mixed_precision": "bf16",
+  "wandb_enabled": true,
+  "wandb_project": "stochastic_sparsity",
+  "wandb_name": "365m alibi 8k s42",
+  "wandb_tags": [
+    "qwen3",
+    "365m",
+    "8k",
+    "scaleladder",
+    "alibi",
+    "seed42"
+  ],
+  "hidden_size": 768,
+  "intermediate_size": 3072,
+  "num_attention_heads": 6,
+  "num_key_value_heads": 2,
+  "num_hidden_layers": 28,
+  "head_dim": null,
+  "attention_dropout": null,
+  "rope_train_position_scale_set": null,
+  "positional": "alibi",
+  "attention_type": "dense",
+  "sparse_block_size": 128,
+  "sparse_num_random_blocks": 31,
+  "sparse_mask_pool_size": 1,
+  "sparse_val_dense_compare": false,
+  "sparse_layer_indices": null,
+  "sparse_mask_cache_dir": null,
+  "sparse_mask_seed": 0,
+  "sparse_mixed_nrb_set": null,
+  "dense_mode": false,
+  "dense_step_intervals": null
+}
+Loading tokenizer from /mloscratch/homes/navasard/qwen3_06b_base...
+Initializing model from scratch (random weights)...
+Checkpoint activations: False
+
+
+================================================================================
+MODEL ARCHITECTURE
+================================================================================
+Model config:
+  vocab_size: 151936
+  hidden_size: 768
+  num_hidden_layers: 28
+  num_attention_heads: 6
+  num_key_value_heads: 2
+  head_dim: 128
+  intermediate_size: 3072
+  max_position_embeddings: 32768
+  rope_theta: 1000000
+  rope_train_position_scale_set: None
+  attention_type: dense
+
+Parameter count:
+  Total parameters: 358,958,848
+  Trainable parameters: 358,958,848
+  Non-trainable parameters: 0
+  Model size (BF16): 0.72 GB
+
+--------------------------------------------------------------------------------
+LAYER-BY-LAYER SUMMARY
+--------------------------------------------------------------------------------
+/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/nn/attention/flex_attention.py:1704: UserWarning: flex_attention called without torch.compile() - this will use an unfused implementation that materializes the full scores matrix instead of generating a fused kernel.
+
+SOLUTION: Use torch.compile(flex_attention)(...)
+
+If you want to debug your score_mod/mask_mod, you can set:
+torch.nn.attention.flex_attention._FLEX_ATTENTION_DISABLE_COMPILE_DEBUG = True
+
+This will allow you to use print statements or breakpoints. Note: This doesn't work with the backwards pass and may produce incorrect results.
+  _warn_once(
+=================================================================================================================================================
+Layer (type:depth-idx)                        Input Shape               Output Shape              Param #                   Trainable
+=================================================================================================================================================
+Qwen3ForCausalLM                              [2, 8192]                 --                        --                        True
+├─Qwen3Model: 1-1                             --                        [2, 8192, 768]            --                        True
+│    └─Embedding: 2-1                         [2, 8192]                 [2, 8192, 768]            116,686,848               True
+│    └─ModuleList: 2-2                        --                        --                        --                        True
+│    │    └─Qwen3DecoderLayer: 3-1            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-2            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-3            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-4            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-5            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-6            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-7            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-8            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-9            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-10           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-11           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-12           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-13           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-14           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-15           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-16           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-17           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-18           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-19           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-20           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-21           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-22           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-23           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-24           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-25           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-26           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-27           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-28           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    └─Qwen3RMSNorm: 2-3                      [2, 8192, 768]            [2, 8192, 768]            768                       True
+├─Linear: 1-2                                 [2, 8192, 768]            [2, 8192, 151936]         116,686,848               True
+=================================================================================================================================================
+Total params: 475,645,696
+Trainable params: 475,645,696
+Non-trainable params: 0
+Total mult-adds (Units.MEGABYTES): 951.29
+=================================================================================================================================================
+Input size (MB): 0.13
+Forward/backward pass size (MB): 31197.23
+Params size (MB): 951.29
+Estimated Total Size (MB): 32148.66
+=================================================================================================================================================
+  NOTE: torchinfo double-counts tied embed_tokens/lm_head weight (116,686,848 params). True unique parameter count: 358,958,848 (359.0M)
+================================================================================
+
+Loading dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Using pre-tokenized dataset (no runtime tokenization)
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Loaded 12748481 pre-tokenized examples (stored_len=2048, pack_factor=4, packed_examples=3187120)
+Dataset split: 3027764 train, 159356 validation
+Training examples (packed, train split):  3027764
+Validation examples (packed, val split):  159356
+Per-device batch size:                    2
+World size (GPUs):                        1
+Gradient accumulation steps:              10
+Effective batch size (seqs/step):         20
+Sequence length:                          8192
+Tokens per optimizer step:                163,840
+Micro-batches per epoch (per rank):       1513882
+Optimizer steps per epoch:                151388
+Total optimizer steps:                    44000  (capped by max_steps; 0.29 epochs)
+wandb: Tracking run with wandb version 0.26.1
+wandb: W&B syncing is set to `offline` in this directory. Run `wandb online` or set WANDB_MODE=online to enable cloud syncing.
+wandb: Run data is saved locally in /mloscratch/homes/navasard/sparse/SparseAtt/wandb/offline-run-20260921_143024-nihe6o4e
+Starting training...
+Switching attention mode for step 1: sparse_stochastic (configured sparse attention)
+Step 100/44000 | Loss: 10.5675 | LR: 6.67e-05 | 1612.8ms/step | 12.4 samples/s
+Step 200/44000 | Loss: 7.3112 | LR: 1.33e-04 | 1503.3ms/step | 13.3 samples/s
+Validation at step 250: Loss = 6.3046, Perplexity = 547.07
+Step 300/44000 | Loss: 6.3235 | LR: 2.00e-04 | 1796.9ms/step | 11.1 samples/s
+Step 400/44000 | Loss: 5.8468 | LR: 2.67e-04 | 1511.2ms/step | 13.2 samples/s
+Step 500/44000 | Loss: 5.4582 | LR: 3.33e-04 | 1509.9ms/step | 13.2 samples/s
+Validation at step 500: Loss = 5.2754, Perplexity = 195.47
+Step 600/44000 | Loss: 5.1382 | LR: 4.00e-04 | 1754.4ms/step | 11.4 samples/s
+Step 700/44000 | Loss: 4.8505 | LR: 4.67e-04 | 1510.9ms/step | 13.2 samples/s
+Validation at step 750: Loss = 4.5974, Perplexity = 99.23
+Step 800/44000 | Loss: 4.5809 | LR: 5.00e-04 | 1754.1ms/step | 11.4 samples/s
+Step 900/44000 | Loss: 4.3776 | LR: 5.00e-04 | 1511.2ms/step | 13.2 samples/s
+Step 1000/44000 | Loss: 4.2240 | LR: 5.00e-04 | 1509.3ms/step | 13.3 samples/s
+Validation at step 1000: Loss = 4.1754, Perplexity = 65.07
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-1000
+Step 1100/44000 | Loss: 4.1065 | LR: 5.00e-04 | 1781.2ms/step | 11.2 samples/s
+Step 1200/44000 | Loss: 4.0142 | LR: 5.00e-04 | 1506.7ms/step | 13.3 samples/s
+Validation at step 1250: Loss = 3.9616, Perplexity = 52.54
+Step 1300/44000 | Loss: 3.9476 | LR: 5.00e-04 | 1749.0ms/step | 11.4 samples/s
+Step 1400/44000 | Loss: 3.8829 | LR: 5.00e-04 | 1507.4ms/step | 13.3 samples/s
+Step 1500/44000 | Loss: 3.8558 | LR: 5.00e-04 | 1507.2ms/step | 13.3 samples/s
+Validation at step 1500: Loss = 3.8283, Perplexity = 45.99
+Step 1600/44000 | Loss: 3.7996 | LR: 5.00e-04 | 1748.3ms/step | 11.4 samples/s
+Step 1700/44000 | Loss: 3.7480 | LR: 4.99e-04 | 1507.4ms/step | 13.3 samples/s
+Validation at step 1750: Loss = 3.7342, Perplexity = 41.85
+Step 1800/44000 | Loss: 3.7112 | LR: 4.99e-04 | 1748.0ms/step | 11.4 samples/s
+Step 1900/44000 | Loss: 3.6984 | LR: 4.99e-04 | 1507.5ms/step | 13.3 samples/s
+Step 2000/44000 | Loss: 3.6660 | LR: 4.99e-04 | 1506.5ms/step | 13.3 samples/s
+Validation at step 2000: Loss = 3.6610, Perplexity = 38.90
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_142951/checkpoint-2000

epfl_final/logs/ALIBI_20260921/resume_20260921_142934/seed42.pid

@@ -0,0 +1 @@
+67

epfl_final/logs/ALIBI_20260921/resume_20260921_142934/seed43.log

@@ -0,0 +1,206 @@
+New run: Checkpoints will be saved to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951
+Running config:
+{
+  "model_path": "/mloscratch/homes/navasard/qwen3_06b_base",
+  "model_dtype": "bfloat16",
+  "init_from_scratch": true,
+  "train_path": "/mloscratch/homes/navasard/data/fineweb_large/tokenized",
+  "val_path": null,
+  "max_seq_length": 8192,
+  "num_workers": 8,
+  "output_dir": "/mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951",
+  "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": 43,
+  "verbose": true,
+  "checkpoint_activations": false,
+  "resume_micro_step_override": null,
+  "strategy": "ddp",
+  "mixed_precision": "bf16",
+  "wandb_enabled": true,
+  "wandb_project": "stochastic_sparsity",
+  "wandb_name": "365m alibi 8k s43",
+  "wandb_tags": [
+    "qwen3",
+    "365m",
+    "8k",
+    "scaleladder",
+    "alibi",
+    "seed43"
+  ],
+  "hidden_size": 768,
+  "intermediate_size": 3072,
+  "num_attention_heads": 6,
+  "num_key_value_heads": 2,
+  "num_hidden_layers": 28,
+  "head_dim": null,
+  "attention_dropout": null,
+  "rope_train_position_scale_set": null,
+  "positional": "alibi",
+  "attention_type": "dense",
+  "sparse_block_size": 128,
+  "sparse_num_random_blocks": 31,
+  "sparse_mask_pool_size": 1,
+  "sparse_val_dense_compare": false,
+  "sparse_layer_indices": null,
+  "sparse_mask_cache_dir": null,
+  "sparse_mask_seed": 0,
+  "sparse_mixed_nrb_set": null,
+  "dense_mode": false,
+  "dense_step_intervals": null
+}
+Loading tokenizer from /mloscratch/homes/navasard/qwen3_06b_base...
+Initializing model from scratch (random weights)...
+Checkpoint activations: False
+
+
+================================================================================
+MODEL ARCHITECTURE
+================================================================================
+Model config:
+  vocab_size: 151936
+  hidden_size: 768
+  num_hidden_layers: 28
+  num_attention_heads: 6
+  num_key_value_heads: 2
+  head_dim: 128
+  intermediate_size: 3072
+  max_position_embeddings: 32768
+  rope_theta: 1000000
+  rope_train_position_scale_set: None
+  attention_type: dense
+
+Parameter count:
+  Total parameters: 358,958,848
+  Trainable parameters: 358,958,848
+  Non-trainable parameters: 0
+  Model size (BF16): 0.72 GB
+
+--------------------------------------------------------------------------------
+LAYER-BY-LAYER SUMMARY
+--------------------------------------------------------------------------------
+/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/nn/attention/flex_attention.py:1704: UserWarning: flex_attention called without torch.compile() - this will use an unfused implementation that materializes the full scores matrix instead of generating a fused kernel.
+
+SOLUTION: Use torch.compile(flex_attention)(...)
+
+If you want to debug your score_mod/mask_mod, you can set:
+torch.nn.attention.flex_attention._FLEX_ATTENTION_DISABLE_COMPILE_DEBUG = True
+
+This will allow you to use print statements or breakpoints. Note: This doesn't work with the backwards pass and may produce incorrect results.
+  _warn_once(
+=================================================================================================================================================
+Layer (type:depth-idx)                        Input Shape               Output Shape              Param #                   Trainable
+=================================================================================================================================================
+Qwen3ForCausalLM                              [2, 8192]                 --                        --                        True
+├─Qwen3Model: 1-1                             --                        [2, 8192, 768]            --                        True
+│    └─Embedding: 2-1                         [2, 8192]                 [2, 8192, 768]            116,686,848               True
+│    └─ModuleList: 2-2                        --                        --                        --                        True
+│    │    └─Qwen3DecoderLayer: 3-1            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-2            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-3            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-4            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-5            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-6            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-7            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-8            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-9            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-10           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-11           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-12           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-13           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-14           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-15           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-16           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-17           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-18           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-19           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-20           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-21           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-22           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-23           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-24           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-25           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-26           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-27           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-28           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    └─Qwen3RMSNorm: 2-3                      [2, 8192, 768]            [2, 8192, 768]            768                       True
+├─Linear: 1-2                                 [2, 8192, 768]            [2, 8192, 151936]         116,686,848               True
+=================================================================================================================================================
+Total params: 475,645,696
+Trainable params: 475,645,696
+Non-trainable params: 0
+Total mult-adds (Units.MEGABYTES): 951.29
+=================================================================================================================================================
+Input size (MB): 0.13
+Forward/backward pass size (MB): 31197.23
+Params size (MB): 951.29
+Estimated Total Size (MB): 32148.66
+=================================================================================================================================================
+  NOTE: torchinfo double-counts tied embed_tokens/lm_head weight (116,686,848 params). True unique parameter count: 358,958,848 (359.0M)
+================================================================================
+
+Loading dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Using pre-tokenized dataset (no runtime tokenization)
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Loaded 12748481 pre-tokenized examples (stored_len=2048, pack_factor=4, packed_examples=3187120)
+Dataset split: 3027764 train, 159356 validation
+Training examples (packed, train split):  3027764
+Validation examples (packed, val split):  159356
+Per-device batch size:                    2
+World size (GPUs):                        1
+Gradient accumulation steps:              10
+Effective batch size (seqs/step):         20
+Sequence length:                          8192
+Tokens per optimizer step:                163,840
+Micro-batches per epoch (per rank):       1513882
+Optimizer steps per epoch:                151388
+Total optimizer steps:                    44000  (capped by max_steps; 0.29 epochs)
+wandb: Tracking run with wandb version 0.26.1
+wandb: W&B syncing is set to `offline` in this directory. Run `wandb online` or set WANDB_MODE=online to enable cloud syncing.
+wandb: Run data is saved locally in /mloscratch/homes/navasard/sparse/SparseAtt/wandb/offline-run-20260921_143024-zgrsvkbi
+Starting training...
+Switching attention mode for step 1: sparse_stochastic (configured sparse attention)
+Step 100/44000 | Loss: 10.6180 | LR: 6.67e-05 | 1577.9ms/step | 12.7 samples/s
+Step 200/44000 | Loss: 7.3271 | LR: 1.33e-04 | 1469.2ms/step | 13.6 samples/s
+Validation at step 250: Loss = 6.2973, Perplexity = 543.10
+Step 300/44000 | Loss: 6.3191 | LR: 2.00e-04 | 1765.4ms/step | 11.3 samples/s
+Step 400/44000 | Loss: 5.8579 | LR: 2.67e-04 | 1466.3ms/step | 13.6 samples/s
+Step 500/44000 | Loss: 5.4560 | LR: 3.33e-04 | 1465.8ms/step | 13.6 samples/s
+Validation at step 500: Loss = 5.2849, Perplexity = 197.34
+Step 600/44000 | Loss: 5.1461 | LR: 4.00e-04 | 1705.8ms/step | 11.7 samples/s
+Step 700/44000 | Loss: 4.8569 | LR: 4.67e-04 | 1467.3ms/step | 13.6 samples/s
+Validation at step 750: Loss = 4.5889, Perplexity = 98.38
+Step 800/44000 | Loss: 4.6015 | LR: 5.00e-04 | 1705.1ms/step | 11.7 samples/s
+Step 900/44000 | Loss: 4.3807 | LR: 5.00e-04 | 1467.1ms/step | 13.6 samples/s
+Step 1000/44000 | Loss: 4.2271 | LR: 5.00e-04 | 1465.7ms/step | 13.6 samples/s
+Validation at step 1000: Loss = 4.1545, Perplexity = 63.72
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-1000
+Step 1100/44000 | Loss: 4.1170 | LR: 5.00e-04 | 1738.2ms/step | 11.5 samples/s
+Step 1200/44000 | Loss: 4.0256 | LR: 5.00e-04 | 1471.8ms/step | 13.6 samples/s
+Validation at step 1250: Loss = 3.9459, Perplexity = 51.72
+Step 1300/44000 | Loss: 3.9546 | LR: 5.00e-04 | 1708.8ms/step | 11.7 samples/s
+Step 1400/44000 | Loss: 3.8984 | LR: 5.00e-04 | 1471.8ms/step | 13.6 samples/s
+Step 1500/44000 | Loss: 3.8381 | LR: 5.00e-04 | 1471.2ms/step | 13.6 samples/s
+Validation at step 1500: Loss = 3.8081, Perplexity = 45.07
+Step 1600/44000 | Loss: 3.8055 | LR: 5.00e-04 | 1708.2ms/step | 11.7 samples/s
+Step 1700/44000 | Loss: 3.7553 | LR: 4.99e-04 | 1470.8ms/step | 13.6 samples/s
+Validation at step 1750: Loss = 3.7134, Perplexity = 40.99
+Step 1800/44000 | Loss: 3.7206 | LR: 4.99e-04 | 1709.1ms/step | 11.7 samples/s
+Step 1900/44000 | Loss: 3.6958 | LR: 4.99e-04 | 1471.3ms/step | 13.6 samples/s
+Step 2000/44000 | Loss: 3.6689 | LR: 4.99e-04 | 1471.9ms/step | 13.6 samples/s
+Validation at step 2000: Loss = 3.6382, Perplexity = 38.02
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_142951/checkpoint-2000
+Step 2100/44000 | Loss: 3.6335 | LR: 4.99e-04 | 1738.8ms/step | 11.5 samples/s

epfl_final/logs/ALIBI_20260921/resume_20260921_142934/seed43.pid

@@ -0,0 +1 @@
+68

epfl_final/logs/ALIBI_20260921/resume_20260921_142934/seed44.log

@@ -0,0 +1,205 @@
+New run: Checkpoints will be saved to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951
+Running config:
+{
+  "model_path": "/mloscratch/homes/navasard/qwen3_06b_base",
+  "model_dtype": "bfloat16",
+  "init_from_scratch": true,
+  "train_path": "/mloscratch/homes/navasard/data/fineweb_large/tokenized",
+  "val_path": null,
+  "max_seq_length": 8192,
+  "num_workers": 8,
+  "output_dir": "/mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951",
+  "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": 44,
+  "verbose": true,
+  "checkpoint_activations": false,
+  "resume_micro_step_override": null,
+  "strategy": "ddp",
+  "mixed_precision": "bf16",
+  "wandb_enabled": true,
+  "wandb_project": "stochastic_sparsity",
+  "wandb_name": "365m alibi 8k s44",
+  "wandb_tags": [
+    "qwen3",
+    "365m",
+    "8k",
+    "scaleladder",
+    "alibi",
+    "seed44"
+  ],
+  "hidden_size": 768,
+  "intermediate_size": 3072,
+  "num_attention_heads": 6,
+  "num_key_value_heads": 2,
+  "num_hidden_layers": 28,
+  "head_dim": null,
+  "attention_dropout": null,
+  "rope_train_position_scale_set": null,
+  "positional": "alibi",
+  "attention_type": "dense",
+  "sparse_block_size": 128,
+  "sparse_num_random_blocks": 31,
+  "sparse_mask_pool_size": 1,
+  "sparse_val_dense_compare": false,
+  "sparse_layer_indices": null,
+  "sparse_mask_cache_dir": null,
+  "sparse_mask_seed": 0,
+  "sparse_mixed_nrb_set": null,
+  "dense_mode": false,
+  "dense_step_intervals": null
+}
+Loading tokenizer from /mloscratch/homes/navasard/qwen3_06b_base...
+Initializing model from scratch (random weights)...
+Checkpoint activations: False
+
+
+================================================================================
+MODEL ARCHITECTURE
+================================================================================
+Model config:
+  vocab_size: 151936
+  hidden_size: 768
+  num_hidden_layers: 28
+  num_attention_heads: 6
+  num_key_value_heads: 2
+  head_dim: 128
+  intermediate_size: 3072
+  max_position_embeddings: 32768
+  rope_theta: 1000000
+  rope_train_position_scale_set: None
+  attention_type: dense
+
+Parameter count:
+  Total parameters: 358,958,848
+  Trainable parameters: 358,958,848
+  Non-trainable parameters: 0
+  Model size (BF16): 0.72 GB
+
+--------------------------------------------------------------------------------
+LAYER-BY-LAYER SUMMARY
+--------------------------------------------------------------------------------
+/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/nn/attention/flex_attention.py:1704: UserWarning: flex_attention called without torch.compile() - this will use an unfused implementation that materializes the full scores matrix instead of generating a fused kernel.
+
+SOLUTION: Use torch.compile(flex_attention)(...)
+
+If you want to debug your score_mod/mask_mod, you can set:
+torch.nn.attention.flex_attention._FLEX_ATTENTION_DISABLE_COMPILE_DEBUG = True
+
+This will allow you to use print statements or breakpoints. Note: This doesn't work with the backwards pass and may produce incorrect results.
+  _warn_once(
+=================================================================================================================================================
+Layer (type:depth-idx)                        Input Shape               Output Shape              Param #                   Trainable
+=================================================================================================================================================
+Qwen3ForCausalLM                              [2, 8192]                 --                        --                        True
+├─Qwen3Model: 1-1                             --                        [2, 8192, 768]            --                        True
+│    └─Embedding: 2-1                         [2, 8192]                 [2, 8192, 768]            116,686,848               True
+│    └─ModuleList: 2-2                        --                        --                        --                        True
+│    │    └─Qwen3DecoderLayer: 3-1            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-2            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-3            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-4            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-5            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-6            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-7            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-8            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-9            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-10           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-11           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-12           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-13           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-14           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-15           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-16           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-17           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-18           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-19           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-20           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-21           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-22           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-23           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-24           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-25           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-26           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-27           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-28           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    └─Qwen3RMSNorm: 2-3                      [2, 8192, 768]            [2, 8192, 768]            768                       True
+├─Linear: 1-2                                 [2, 8192, 768]            [2, 8192, 151936]         116,686,848               True
+=================================================================================================================================================
+Total params: 475,645,696
+Trainable params: 475,645,696
+Non-trainable params: 0
+Total mult-adds (Units.MEGABYTES): 951.29
+=================================================================================================================================================
+Input size (MB): 0.13
+Forward/backward pass size (MB): 31197.23
+Params size (MB): 951.29
+Estimated Total Size (MB): 32148.66
+=================================================================================================================================================
+  NOTE: torchinfo double-counts tied embed_tokens/lm_head weight (116,686,848 params). True unique parameter count: 358,958,848 (359.0M)
+================================================================================
+
+Loading dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Using pre-tokenized dataset (no runtime tokenization)
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Loaded 12748481 pre-tokenized examples (stored_len=2048, pack_factor=4, packed_examples=3187120)
+Dataset split: 3027764 train, 159356 validation
+Training examples (packed, train split):  3027764
+Validation examples (packed, val split):  159356
+Per-device batch size:                    2
+World size (GPUs):                        1
+Gradient accumulation steps:              10
+Effective batch size (seqs/step):         20
+Sequence length:                          8192
+Tokens per optimizer step:                163,840
+Micro-batches per epoch (per rank):       1513882
+Optimizer steps per epoch:                151388
+Total optimizer steps:                    44000  (capped by max_steps; 0.29 epochs)
+wandb: Tracking run with wandb version 0.26.1
+wandb: W&B syncing is set to `offline` in this directory. Run `wandb online` or set WANDB_MODE=online to enable cloud syncing.
+wandb: Run data is saved locally in /mloscratch/homes/navasard/sparse/SparseAtt/wandb/offline-run-20260921_143024-kbn66izc
+Starting training...
+Switching attention mode for step 1: sparse_stochastic (configured sparse attention)
+Step 100/44000 | Loss: 10.6003 | LR: 6.67e-05 | 1621.2ms/step | 12.3 samples/s
+Step 200/44000 | Loss: 7.3145 | LR: 1.33e-04 | 1513.5ms/step | 13.2 samples/s
+Validation at step 250: Loss = 6.3145, Perplexity = 552.50
+Step 300/44000 | Loss: 6.3199 | LR: 2.00e-04 | 1808.6ms/step | 11.1 samples/s
+Step 400/44000 | Loss: 5.8467 | LR: 2.67e-04 | 1522.0ms/step | 13.1 samples/s
+Step 500/44000 | Loss: 5.4660 | LR: 3.33e-04 | 1521.0ms/step | 13.1 samples/s
+Validation at step 500: Loss = 5.3073, Perplexity = 201.80
+Step 600/44000 | Loss: 5.1441 | LR: 4.00e-04 | 1769.2ms/step | 11.3 samples/s
+Step 700/44000 | Loss: 4.8768 | LR: 4.67e-04 | 1525.7ms/step | 13.1 samples/s
+Validation at step 750: Loss = 4.6389, Perplexity = 103.43
+Step 800/44000 | Loss: 4.6295 | LR: 5.00e-04 | 1771.4ms/step | 11.3 samples/s
+Step 900/44000 | Loss: 4.3996 | LR: 5.00e-04 | 1525.8ms/step | 13.1 samples/s
+Step 1000/44000 | Loss: 4.2347 | LR: 5.00e-04 | 1518.8ms/step | 13.2 samples/s
+Validation at step 1000: Loss = 4.1789, Perplexity = 65.29
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-1000
+Step 1100/44000 | Loss: 4.1037 | LR: 5.00e-04 | 1792.7ms/step | 11.2 samples/s
+Step 1200/44000 | Loss: 4.0177 | LR: 5.00e-04 | 1519.6ms/step | 13.2 samples/s
+Validation at step 1250: Loss = 3.9505, Perplexity = 51.96
+Step 1300/44000 | Loss: 3.9444 | LR: 5.00e-04 | 1767.2ms/step | 11.3 samples/s
+Step 1400/44000 | Loss: 3.8773 | LR: 5.00e-04 | 1521.1ms/step | 13.1 samples/s
+Step 1500/44000 | Loss: 3.8416 | LR: 5.00e-04 | 1521.4ms/step | 13.1 samples/s
+Validation at step 1500: Loss = 3.8165, Perplexity = 45.45
+Step 1600/44000 | Loss: 3.7900 | LR: 5.00e-04 | 1763.9ms/step | 11.3 samples/s
+Step 1700/44000 | Loss: 3.7546 | LR: 4.99e-04 | 1519.8ms/step | 13.2 samples/s
+Validation at step 1750: Loss = 3.7226, Perplexity = 41.37
+Step 1800/44000 | Loss: 3.7246 | LR: 4.99e-04 | 1764.4ms/step | 11.3 samples/s
+Step 1900/44000 | Loss: 3.6833 | LR: 4.99e-04 | 1521.2ms/step | 13.1 samples/s
+Step 2000/44000 | Loss: 3.6624 | LR: 4.99e-04 | 1518.4ms/step | 13.2 samples/s
+Validation at step 2000: Loss = 3.6519, Perplexity = 38.55
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_142951/checkpoint-2000

epfl_final/logs/ALIBI_20260921/resume_20260921_142934/seed44.pid

@@ -0,0 +1 @@
+69

epfl_final/logs/ALIBI_20260921/resume_20260921_142934/seed45.log

@@ -0,0 +1,206 @@
+New run: Checkpoints will be saved to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s45/run_20260921_142951
+Running config:
+{
+  "model_path": "/mloscratch/homes/navasard/qwen3_06b_base",
+  "model_dtype": "bfloat16",
+  "init_from_scratch": true,
+  "train_path": "/mloscratch/homes/navasard/data/fineweb_large/tokenized",
+  "val_path": null,
+  "max_seq_length": 8192,
+  "num_workers": 8,
+  "output_dir": "/mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s45/run_20260921_142951",
+  "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": 45,
+  "verbose": true,
+  "checkpoint_activations": false,
+  "resume_micro_step_override": null,
+  "strategy": "ddp",
+  "mixed_precision": "bf16",
+  "wandb_enabled": true,
+  "wandb_project": "stochastic_sparsity",
+  "wandb_name": "365m alibi 8k s45",
+  "wandb_tags": [
+    "qwen3",
+    "365m",
+    "8k",
+    "scaleladder",
+    "alibi",
+    "seed45"
+  ],
+  "hidden_size": 768,
+  "intermediate_size": 3072,
+  "num_attention_heads": 6,
+  "num_key_value_heads": 2,
+  "num_hidden_layers": 28,
+  "head_dim": null,
+  "attention_dropout": null,
+  "rope_train_position_scale_set": null,
+  "positional": "alibi",
+  "attention_type": "dense",
+  "sparse_block_size": 128,
+  "sparse_num_random_blocks": 31,
+  "sparse_mask_pool_size": 1,
+  "sparse_val_dense_compare": false,
+  "sparse_layer_indices": null,
+  "sparse_mask_cache_dir": null,
+  "sparse_mask_seed": 0,
+  "sparse_mixed_nrb_set": null,
+  "dense_mode": false,
+  "dense_step_intervals": null
+}
+Loading tokenizer from /mloscratch/homes/navasard/qwen3_06b_base...
+Initializing model from scratch (random weights)...
+Checkpoint activations: False
+
+
+================================================================================
+MODEL ARCHITECTURE
+================================================================================
+Model config:
+  vocab_size: 151936
+  hidden_size: 768
+  num_hidden_layers: 28
+  num_attention_heads: 6
+  num_key_value_heads: 2
+  head_dim: 128
+  intermediate_size: 3072
+  max_position_embeddings: 32768
+  rope_theta: 1000000
+  rope_train_position_scale_set: None
+  attention_type: dense
+
+Parameter count:
+  Total parameters: 358,958,848
+  Trainable parameters: 358,958,848
+  Non-trainable parameters: 0
+  Model size (BF16): 0.72 GB
+
+--------------------------------------------------------------------------------
+LAYER-BY-LAYER SUMMARY
+--------------------------------------------------------------------------------
+/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/nn/attention/flex_attention.py:1704: UserWarning: flex_attention called without torch.compile() - this will use an unfused implementation that materializes the full scores matrix instead of generating a fused kernel.
+
+SOLUTION: Use torch.compile(flex_attention)(...)
+
+If you want to debug your score_mod/mask_mod, you can set:
+torch.nn.attention.flex_attention._FLEX_ATTENTION_DISABLE_COMPILE_DEBUG = True
+
+This will allow you to use print statements or breakpoints. Note: This doesn't work with the backwards pass and may produce incorrect results.
+  _warn_once(
+=================================================================================================================================================
+Layer (type:depth-idx)                        Input Shape               Output Shape              Param #                   Trainable
+=================================================================================================================================================
+Qwen3ForCausalLM                              [2, 8192]                 --                        --                        True
+├─Qwen3Model: 1-1                             --                        [2, 8192, 768]            --                        True
+│    └─Embedding: 2-1                         [2, 8192]                 [2, 8192, 768]            116,686,848               True
+│    └─ModuleList: 2-2                        --                        --                        --                        True
+│    │    └─Qwen3DecoderLayer: 3-1            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-2            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-3            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-4            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-5            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-6            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-7            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-8            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-9            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-10           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-11           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-12           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-13           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-14           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-15           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-16           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-17           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-18           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-19           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-20           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-21           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-22           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-23           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-24           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-25           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-26           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-27           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-28           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    └─Qwen3RMSNorm: 2-3                      [2, 8192, 768]            [2, 8192, 768]            768                       True
+├─Linear: 1-2                                 [2, 8192, 768]            [2, 8192, 151936]         116,686,848               True
+=================================================================================================================================================
+Total params: 475,645,696
+Trainable params: 475,645,696
+Non-trainable params: 0
+Total mult-adds (Units.MEGABYTES): 951.29
+=================================================================================================================================================
+Input size (MB): 0.13
+Forward/backward pass size (MB): 31197.23
+Params size (MB): 951.29
+Estimated Total Size (MB): 32148.66
+=================================================================================================================================================
+  NOTE: torchinfo double-counts tied embed_tokens/lm_head weight (116,686,848 params). True unique parameter count: 358,958,848 (359.0M)
+================================================================================
+
+Loading dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Using pre-tokenized dataset (no runtime tokenization)
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Loaded 12748481 pre-tokenized examples (stored_len=2048, pack_factor=4, packed_examples=3187120)
+Dataset split: 3027764 train, 159356 validation
+Training examples (packed, train split):  3027764
+Validation examples (packed, val split):  159356
+Per-device batch size:                    2
+World size (GPUs):                        1
+Gradient accumulation steps:              10
+Effective batch size (seqs/step):         20
+Sequence length:                          8192
+Tokens per optimizer step:                163,840
+Micro-batches per epoch (per rank):       1513882
+Optimizer steps per epoch:                151388
+Total optimizer steps:                    44000  (capped by max_steps; 0.29 epochs)
+wandb: Tracking run with wandb version 0.26.1
+wandb: W&B syncing is set to `offline` in this directory. Run `wandb online` or set WANDB_MODE=online to enable cloud syncing.
+wandb: Run data is saved locally in /mloscratch/homes/navasard/sparse/SparseAtt/wandb/offline-run-20260921_143024-ghfbqjah
+Starting training...
+Switching attention mode for step 1: sparse_stochastic (configured sparse attention)
+Step 100/44000 | Loss: 10.6683 | LR: 6.67e-05 | 1599.8ms/step | 12.5 samples/s
+Step 200/44000 | Loss: 7.3481 | LR: 1.33e-04 | 1488.9ms/step | 13.4 samples/s
+Validation at step 250: Loss = 6.3176, Perplexity = 554.21
+Step 300/44000 | Loss: 6.3392 | LR: 2.00e-04 | 1784.0ms/step | 11.2 samples/s
+Step 400/44000 | Loss: 5.8523 | LR: 2.67e-04 | 1499.4ms/step | 13.3 samples/s
+Step 500/44000 | Loss: 5.4673 | LR: 3.33e-04 | 1499.0ms/step | 13.3 samples/s
+Validation at step 500: Loss = 5.2963, Perplexity = 199.59
+Step 600/44000 | Loss: 5.1331 | LR: 4.00e-04 | 1739.7ms/step | 11.5 samples/s
+Step 700/44000 | Loss: 4.8681 | LR: 4.67e-04 | 1499.0ms/step | 13.3 samples/s
+Validation at step 750: Loss = 4.6091, Perplexity = 100.40
+Step 800/44000 | Loss: 4.6111 | LR: 5.00e-04 | 1739.8ms/step | 11.5 samples/s
+Step 900/44000 | Loss: 4.3903 | LR: 5.00e-04 | 1499.7ms/step | 13.3 samples/s
+Step 1000/44000 | Loss: 4.2294 | LR: 5.00e-04 | 1497.9ms/step | 13.4 samples/s
+Validation at step 1000: Loss = 4.1733, Perplexity = 64.93
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s45/run_20260921_142951/checkpoint-1000
+Step 1100/44000 | Loss: 4.1142 | LR: 5.00e-04 | 1767.5ms/step | 11.3 samples/s
+Step 1200/44000 | Loss: 4.0208 | LR: 5.00e-04 | 1497.8ms/step | 13.4 samples/s
+Validation at step 1250: Loss = 3.9543, Perplexity = 52.16
+Step 1300/44000 | Loss: 3.9399 | LR: 5.00e-04 | 1738.3ms/step | 11.5 samples/s
+Step 1400/44000 | Loss: 3.8898 | LR: 5.00e-04 | 1497.9ms/step | 13.4 samples/s
+Step 1500/44000 | Loss: 3.8383 | LR: 5.00e-04 | 1499.1ms/step | 13.3 samples/s
+Validation at step 1500: Loss = 3.8165, Perplexity = 45.45
+Step 1600/44000 | Loss: 3.8006 | LR: 5.00e-04 | 1736.8ms/step | 11.5 samples/s
+Step 1700/44000 | Loss: 3.7584 | LR: 4.99e-04 | 1498.6ms/step | 13.3 samples/s
+Validation at step 1750: Loss = 3.7225, Perplexity = 41.37
+Step 1800/44000 | Loss: 3.7134 | LR: 4.99e-04 | 1737.2ms/step | 11.5 samples/s
+Step 1900/44000 | Loss: 3.6849 | LR: 4.99e-04 | 1497.1ms/step | 13.4 samples/s
+Step 2000/44000 | Loss: 3.6698 | LR: 4.99e-04 | 1497.2ms/step | 13.4 samples/s
+Validation at step 2000: Loss = 3.6535, Perplexity = 38.61
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s45/run_20260921_142951/checkpoint-2000
+Step 2100/44000 | Loss: 3.6373 | LR: 4.99e-04 | 1766.5ms/step | 11.3 samples/s

epfl_final/logs/ALIBI_20260921/resume_20260921_142934/seed45.pid

@@ -0,0 +1 @@
+70

epfl_final/logs/ALIBI_20260921/resume_20260921_142934/seed46.log

@@ -0,0 +1,206 @@
+New run: Checkpoints will be saved to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s46/run_20260921_142951
+Running config:
+{
+  "model_path": "/mloscratch/homes/navasard/qwen3_06b_base",
+  "model_dtype": "bfloat16",
+  "init_from_scratch": true,
+  "train_path": "/mloscratch/homes/navasard/data/fineweb_large/tokenized",
+  "val_path": null,
+  "max_seq_length": 8192,
+  "num_workers": 8,
+  "output_dir": "/mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s46/run_20260921_142951",
+  "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": 46,
+  "verbose": true,
+  "checkpoint_activations": false,
+  "resume_micro_step_override": null,
+  "strategy": "ddp",
+  "mixed_precision": "bf16",
+  "wandb_enabled": true,
+  "wandb_project": "stochastic_sparsity",
+  "wandb_name": "365m alibi 8k s46",
+  "wandb_tags": [
+    "qwen3",
+    "365m",
+    "8k",
+    "scaleladder",
+    "alibi",
+    "seed46"
+  ],
+  "hidden_size": 768,
+  "intermediate_size": 3072,
+  "num_attention_heads": 6,
+  "num_key_value_heads": 2,
+  "num_hidden_layers": 28,
+  "head_dim": null,
+  "attention_dropout": null,
+  "rope_train_position_scale_set": null,
+  "positional": "alibi",
+  "attention_type": "dense",
+  "sparse_block_size": 128,
+  "sparse_num_random_blocks": 31,
+  "sparse_mask_pool_size": 1,
+  "sparse_val_dense_compare": false,
+  "sparse_layer_indices": null,
+  "sparse_mask_cache_dir": null,
+  "sparse_mask_seed": 0,
+  "sparse_mixed_nrb_set": null,
+  "dense_mode": false,
+  "dense_step_intervals": null
+}
+Loading tokenizer from /mloscratch/homes/navasard/qwen3_06b_base...
+Initializing model from scratch (random weights)...
+Checkpoint activations: False
+
+
+================================================================================
+MODEL ARCHITECTURE
+================================================================================
+Model config:
+  vocab_size: 151936
+  hidden_size: 768
+  num_hidden_layers: 28
+  num_attention_heads: 6
+  num_key_value_heads: 2
+  head_dim: 128
+  intermediate_size: 3072
+  max_position_embeddings: 32768
+  rope_theta: 1000000
+  rope_train_position_scale_set: None
+  attention_type: dense
+
+Parameter count:
+  Total parameters: 358,958,848
+  Trainable parameters: 358,958,848
+  Non-trainable parameters: 0
+  Model size (BF16): 0.72 GB
+
+--------------------------------------------------------------------------------
+LAYER-BY-LAYER SUMMARY
+--------------------------------------------------------------------------------
+/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/nn/attention/flex_attention.py:1704: UserWarning: flex_attention called without torch.compile() - this will use an unfused implementation that materializes the full scores matrix instead of generating a fused kernel.
+
+SOLUTION: Use torch.compile(flex_attention)(...)
+
+If you want to debug your score_mod/mask_mod, you can set:
+torch.nn.attention.flex_attention._FLEX_ATTENTION_DISABLE_COMPILE_DEBUG = True
+
+This will allow you to use print statements or breakpoints. Note: This doesn't work with the backwards pass and may produce incorrect results.
+  _warn_once(
+=================================================================================================================================================
+Layer (type:depth-idx)                        Input Shape               Output Shape              Param #                   Trainable
+=================================================================================================================================================
+Qwen3ForCausalLM                              [2, 8192]                 --                        --                        True
+├─Qwen3Model: 1-1                             --                        [2, 8192, 768]            --                        True
+│    └─Embedding: 2-1                         [2, 8192]                 [2, 8192, 768]            116,686,848               True
+│    └─ModuleList: 2-2                        --                        --                        --                        True
+│    │    └─Qwen3DecoderLayer: 3-1            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-2            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-3            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-4            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-5            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-6            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-7            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-8            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-9            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-10           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-11           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-12           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-13           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-14           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-15           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-16           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-17           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-18           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-19           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-20           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-21           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-22           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-23           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-24           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-25           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-26           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-27           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-28           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    └─Qwen3RMSNorm: 2-3                      [2, 8192, 768]            [2, 8192, 768]            768                       True
+├─Linear: 1-2                                 [2, 8192, 768]            [2, 8192, 151936]         116,686,848               True
+=================================================================================================================================================
+Total params: 475,645,696
+Trainable params: 475,645,696
+Non-trainable params: 0
+Total mult-adds (Units.MEGABYTES): 951.29
+=================================================================================================================================================
+Input size (MB): 0.13
+Forward/backward pass size (MB): 31197.23
+Params size (MB): 951.29
+Estimated Total Size (MB): 32148.66
+=================================================================================================================================================
+  NOTE: torchinfo double-counts tied embed_tokens/lm_head weight (116,686,848 params). True unique parameter count: 358,958,848 (359.0M)
+================================================================================
+
+Loading dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Using pre-tokenized dataset (no runtime tokenization)
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Loaded 12748481 pre-tokenized examples (stored_len=2048, pack_factor=4, packed_examples=3187120)
+Dataset split: 3027764 train, 159356 validation
+Training examples (packed, train split):  3027764
+Validation examples (packed, val split):  159356
+Per-device batch size:                    2
+World size (GPUs):                        1
+Gradient accumulation steps:              10
+Effective batch size (seqs/step):         20
+Sequence length:                          8192
+Tokens per optimizer step:                163,840
+Micro-batches per epoch (per rank):       1513882
+Optimizer steps per epoch:                151388
+Total optimizer steps:                    44000  (capped by max_steps; 0.29 epochs)
+wandb: Tracking run with wandb version 0.26.1
+wandb: W&B syncing is set to `offline` in this directory. Run `wandb online` or set WANDB_MODE=online to enable cloud syncing.
+wandb: Run data is saved locally in /mloscratch/homes/navasard/sparse/SparseAtt/wandb/offline-run-20260921_143024-wk3rt7ae
+Starting training...
+Switching attention mode for step 1: sparse_stochastic (configured sparse attention)
+Step 100/44000 | Loss: 10.6054 | LR: 6.67e-05 | 1588.1ms/step | 12.6 samples/s
+Step 200/44000 | Loss: 7.3367 | LR: 1.33e-04 | 1467.4ms/step | 13.6 samples/s
+Validation at step 250: Loss = 6.3121, Perplexity = 551.21
+Step 300/44000 | Loss: 6.3284 | LR: 2.00e-04 | 1768.3ms/step | 11.3 samples/s
+Step 400/44000 | Loss: 5.8568 | LR: 2.67e-04 | 1474.3ms/step | 13.6 samples/s
+Step 500/44000 | Loss: 5.4691 | LR: 3.33e-04 | 1475.3ms/step | 13.6 samples/s
+Validation at step 500: Loss = 5.2977, Perplexity = 199.88
+Step 600/44000 | Loss: 5.1317 | LR: 4.00e-04 | 1722.6ms/step | 11.6 samples/s
+Step 700/44000 | Loss: 4.8783 | LR: 4.67e-04 | 1483.7ms/step | 13.5 samples/s
+Validation at step 750: Loss = 4.6162, Perplexity = 101.11
+Step 800/44000 | Loss: 4.6176 | LR: 5.00e-04 | 1721.4ms/step | 11.6 samples/s
+Step 900/44000 | Loss: 4.3929 | LR: 5.00e-04 | 1483.0ms/step | 13.5 samples/s
+Step 1000/44000 | Loss: 4.2150 | LR: 5.00e-04 | 1482.0ms/step | 13.5 samples/s
+Validation at step 1000: Loss = 4.1633, Perplexity = 64.29
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s46/run_20260921_142951/checkpoint-1000
+Step 1100/44000 | Loss: 4.1118 | LR: 5.00e-04 | 1750.2ms/step | 11.4 samples/s
+Step 1200/44000 | Loss: 4.0237 | LR: 5.00e-04 | 1479.6ms/step | 13.5 samples/s
+Validation at step 1250: Loss = 3.9478, Perplexity = 51.82
+Step 1300/44000 | Loss: 3.9589 | LR: 5.00e-04 | 1720.4ms/step | 11.6 samples/s
+Step 1400/44000 | Loss: 3.8765 | LR: 5.00e-04 | 1479.5ms/step | 13.5 samples/s
+Step 1500/44000 | Loss: 3.8295 | LR: 5.00e-04 | 1476.5ms/step | 13.5 samples/s
+Validation at step 1500: Loss = 3.8106, Perplexity = 45.18
+Step 1600/44000 | Loss: 3.7865 | LR: 5.00e-04 | 1715.6ms/step | 11.7 samples/s
+Step 1700/44000 | Loss: 3.7385 | LR: 4.99e-04 | 1474.9ms/step | 13.6 samples/s
+Validation at step 1750: Loss = 3.7148, Perplexity = 41.05
+Step 1800/44000 | Loss: 3.7060 | LR: 4.99e-04 | 1711.8ms/step | 11.7 samples/s
+Step 1900/44000 | Loss: 3.6922 | LR: 4.99e-04 | 1473.2ms/step | 13.6 samples/s
+Step 2000/44000 | Loss: 3.6616 | LR: 4.99e-04 | 1473.2ms/step | 13.6 samples/s
+Validation at step 2000: Loss = 3.6463, Perplexity = 38.33
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s46/run_20260921_142951/checkpoint-2000
+Step 2100/44000 | Loss: 3.6354 | LR: 4.99e-04 | 1748.4ms/step | 11.4 samples/s

epfl_final/logs/ALIBI_20260921/resume_20260921_142934/seed46.pid

@@ -0,0 +1 @@
+71

epfl_final/logs/ALIBI_20260921/resume_20260921_142934/seed47.log

@@ -0,0 +1,206 @@
+New run: Checkpoints will be saved to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s47/run_20260921_142951
+Running config:
+{
+  "model_path": "/mloscratch/homes/navasard/qwen3_06b_base",
+  "model_dtype": "bfloat16",
+  "init_from_scratch": true,
+  "train_path": "/mloscratch/homes/navasard/data/fineweb_large/tokenized",
+  "val_path": null,
+  "max_seq_length": 8192,
+  "num_workers": 8,
+  "output_dir": "/mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s47/run_20260921_142951",
+  "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,
+  "resume_micro_step_override": null,
+  "strategy": "ddp",
+  "mixed_precision": "bf16",
+  "wandb_enabled": true,
+  "wandb_project": "stochastic_sparsity",
+  "wandb_name": "365m alibi 8k s47",
+  "wandb_tags": [
+    "qwen3",
+    "365m",
+    "8k",
+    "scaleladder",
+    "alibi",
+    "seed47"
+  ],
+  "hidden_size": 768,
+  "intermediate_size": 3072,
+  "num_attention_heads": 6,
+  "num_key_value_heads": 2,
+  "num_hidden_layers": 28,
+  "head_dim": null,
+  "attention_dropout": null,
+  "rope_train_position_scale_set": null,
+  "positional": "alibi",
+  "attention_type": "dense",
+  "sparse_block_size": 128,
+  "sparse_num_random_blocks": 31,
+  "sparse_mask_pool_size": 1,
+  "sparse_val_dense_compare": false,
+  "sparse_layer_indices": null,
+  "sparse_mask_cache_dir": null,
+  "sparse_mask_seed": 0,
+  "sparse_mixed_nrb_set": null,
+  "dense_mode": false,
+  "dense_step_intervals": null
+}
+Loading tokenizer from /mloscratch/homes/navasard/qwen3_06b_base...
+Initializing model from scratch (random weights)...
+Checkpoint activations: False
+
+
+================================================================================
+MODEL ARCHITECTURE
+================================================================================
+Model config:
+  vocab_size: 151936
+  hidden_size: 768
+  num_hidden_layers: 28
+  num_attention_heads: 6
+  num_key_value_heads: 2
+  head_dim: 128
+  intermediate_size: 3072
+  max_position_embeddings: 32768
+  rope_theta: 1000000
+  rope_train_position_scale_set: None
+  attention_type: dense
+
+Parameter count:
+  Total parameters: 358,958,848
+  Trainable parameters: 358,958,848
+  Non-trainable parameters: 0
+  Model size (BF16): 0.72 GB
+
+--------------------------------------------------------------------------------
+LAYER-BY-LAYER SUMMARY
+--------------------------------------------------------------------------------
+/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/nn/attention/flex_attention.py:1704: UserWarning: flex_attention called without torch.compile() - this will use an unfused implementation that materializes the full scores matrix instead of generating a fused kernel.
+
+SOLUTION: Use torch.compile(flex_attention)(...)
+
+If you want to debug your score_mod/mask_mod, you can set:
+torch.nn.attention.flex_attention._FLEX_ATTENTION_DISABLE_COMPILE_DEBUG = True
+
+This will allow you to use print statements or breakpoints. Note: This doesn't work with the backwards pass and may produce incorrect results.
+  _warn_once(
+=================================================================================================================================================
+Layer (type:depth-idx)                        Input Shape               Output Shape              Param #                   Trainable
+=================================================================================================================================================
+Qwen3ForCausalLM                              [2, 8192]                 --                        --                        True
+├─Qwen3Model: 1-1                             --                        [2, 8192, 768]            --                        True
+│    └─Embedding: 2-1                         [2, 8192]                 [2, 8192, 768]            116,686,848               True
+│    └─ModuleList: 2-2                        --                        --                        --                        True
+│    │    └─Qwen3DecoderLayer: 3-1            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-2            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-3            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-4            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-5            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-6            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-7            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-8            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-9            [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-10           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-11           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-12           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-13           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-14           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-15           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-16           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-17           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-18           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-19           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-20           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-21           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-22           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-23           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-24           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-25           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-26           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-27           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    │    └─Qwen3DecoderLayer: 3-28           [2, 8192, 768]            [2, 8192, 768]            8,652,544                 True
+│    └─Qwen3RMSNorm: 2-3                      [2, 8192, 768]            [2, 8192, 768]            768                       True
+├─Linear: 1-2                                 [2, 8192, 768]            [2, 8192, 151936]         116,686,848               True
+=================================================================================================================================================
+Total params: 475,645,696
+Trainable params: 475,645,696
+Non-trainable params: 0
+Total mult-adds (Units.MEGABYTES): 951.29
+=================================================================================================================================================
+Input size (MB): 0.13
+Forward/backward pass size (MB): 31197.23
+Params size (MB): 951.29
+Estimated Total Size (MB): 32148.66
+=================================================================================================================================================
+  NOTE: torchinfo double-counts tied embed_tokens/lm_head weight (116,686,848 params). True unique parameter count: 358,958,848 (359.0M)
+================================================================================
+
+Loading dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Using pre-tokenized dataset (no runtime tokenization)
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Loaded 12748481 pre-tokenized examples (stored_len=2048, pack_factor=4, packed_examples=3187120)
+Dataset split: 3027764 train, 159356 validation
+Training examples (packed, train split):  3027764
+Validation examples (packed, val split):  159356
+Per-device batch size:                    2
+World size (GPUs):                        1
+Gradient accumulation steps:              10
+Effective batch size (seqs/step):         20
+Sequence length:                          8192
+Tokens per optimizer step:                163,840
+Micro-batches per epoch (per rank):       1513882
+Optimizer steps per epoch:                151388
+Total optimizer steps:                    44000  (capped by max_steps; 0.29 epochs)
+wandb: Tracking run with wandb version 0.26.1
+wandb: W&B syncing is set to `offline` in this directory. Run `wandb online` or set WANDB_MODE=online to enable cloud syncing.
+wandb: Run data is saved locally in /mloscratch/homes/navasard/sparse/SparseAtt/wandb/offline-run-20260921_143024-6ijtsdl5
+Starting training...
+Switching attention mode for step 1: sparse_stochastic (configured sparse attention)
+Step 100/44000 | Loss: 10.6289 | LR: 6.67e-05 | 1602.5ms/step | 12.5 samples/s
+Step 200/44000 | Loss: 7.3278 | LR: 1.33e-04 | 1485.1ms/step | 13.5 samples/s
+Validation at step 250: Loss = 6.3147, Perplexity = 552.63
+Step 300/44000 | Loss: 6.3426 | LR: 2.00e-04 | 1779.5ms/step | 11.2 samples/s
+Step 400/44000 | Loss: 5.8651 | LR: 2.67e-04 | 1491.6ms/step | 13.4 samples/s
+Step 500/44000 | Loss: 5.4717 | LR: 3.33e-04 | 1491.3ms/step | 13.4 samples/s
+Validation at step 500: Loss = 5.2942, Perplexity = 199.17
+Step 600/44000 | Loss: 5.1439 | LR: 4.00e-04 | 1731.1ms/step | 11.6 samples/s
+Step 700/44000 | Loss: 4.8670 | LR: 4.67e-04 | 1495.1ms/step | 13.4 samples/s
+Validation at step 750: Loss = 4.6078, Perplexity = 100.26
+Step 800/44000 | Loss: 4.5991 | LR: 5.00e-04 | 1737.2ms/step | 11.5 samples/s
+Step 900/44000 | Loss: 4.3847 | LR: 5.00e-04 | 1498.2ms/step | 13.3 samples/s
+Step 1000/44000 | Loss: 4.2247 | LR: 5.00e-04 | 1496.7ms/step | 13.4 samples/s
+Validation at step 1000: Loss = 4.1555, Perplexity = 63.78
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s47/run_20260921_142951/checkpoint-1000
+Step 1100/44000 | Loss: 4.0996 | LR: 5.00e-04 | 1769.0ms/step | 11.3 samples/s
+Step 1200/44000 | Loss: 4.0259 | LR: 5.00e-04 | 1498.1ms/step | 13.4 samples/s
+Validation at step 1250: Loss = 3.9434, Perplexity = 51.59
+Step 1300/44000 | Loss: 3.9467 | LR: 5.00e-04 | 1738.5ms/step | 11.5 samples/s
+Step 1400/44000 | Loss: 3.8872 | LR: 5.00e-04 | 1498.3ms/step | 13.3 samples/s
+Step 1500/44000 | Loss: 3.8288 | LR: 5.00e-04 | 1496.2ms/step | 13.4 samples/s
+Validation at step 1500: Loss = 3.8142, Perplexity = 45.34
+Step 1600/44000 | Loss: 3.8018 | LR: 5.00e-04 | 1734.8ms/step | 11.5 samples/s
+Step 1700/44000 | Loss: 3.7499 | LR: 4.99e-04 | 1492.7ms/step | 13.4 samples/s
+Validation at step 1750: Loss = 3.7169, Perplexity = 41.14
+Step 1800/44000 | Loss: 3.7212 | LR: 4.99e-04 | 1732.5ms/step | 11.5 samples/s
+Step 1900/44000 | Loss: 3.6873 | LR: 4.99e-04 | 1493.1ms/step | 13.4 samples/s
+Step 2000/44000 | Loss: 3.6613 | LR: 4.99e-04 | 1493.2ms/step | 13.4 samples/s
+Validation at step 2000: Loss = 3.6485, Perplexity = 38.42
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s47/run_20260921_142951/checkpoint-2000
+Step 2100/44000 | Loss: 3.6308 | LR: 4.99e-04 | 1763.2ms/step | 11.3 samples/s

epfl_final/logs/ALIBI_20260921/resume_20260921_142934/seed47.pid

@@ -0,0 +1 @@
+72

epfl_final/train_logs/eval30k_resamp.log

@@ -0,0 +1,39 @@
+windows gutenberg (16, 30720)
+windows arxiv (16, 30720)
+windows pubmed (16, 30720)
+Checkpoint activations: False
+built topk=240 attn=moba layer0.top_k=240 (compiling flex...)
+/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/_inductor/compile_fx.py:322: UserWarning: TensorFloat32 tensor cores for float32 matrix multiplication available but not enabled. Consider setting `torch.set_float32_matmul_precision('high')` for better performance.
+  warnings.warn(
+topk 240 gutenberg L30720 ppl    48.114 nll 3.8736 win 16  59s
+topk 240 arxiv     L30720 ppl   162.161 nll 5.0886 win 16  10s
+topk 240 pubmed    L30720 ppl    42.233 nll 3.7432 win 16  10s
+Checkpoint activations: False
+built topk=120 attn=moba layer0.top_k=120 (compiling flex...)
+topk 120 gutenberg L30720 ppl    45.042 nll 3.8076 win 16  60s
+topk 120 arxiv     L30720 ppl   143.233 nll 4.9645 win 16  9s
+topk 120 pubmed    L30720 ppl    37.151 nll 3.6150 win 16  9s
+Checkpoint activations: False
+built topk=60 attn=moba layer0.top_k=60 (compiling flex...)
+topk  60 gutenberg L30720 ppl    34.144 nll 3.5306 win 16  25s
+topk  60 arxiv     L30720 ppl   100.379 nll 4.6090 win 16  7s
+topk  60 pubmed    L30720 ppl    24.305 nll 3.1907 win 16  7s
+Checkpoint activations: False
+built topk=30 attn=moba layer0.top_k=30 (compiling flex...)
+topk  30 gutenberg L30720 ppl    26.310 nll 3.2700 win 16  53s
+topk  30 arxiv     L30720 ppl    73.115 nll 4.2920 win 16  7s
+topk  30 pubmed    L30720 ppl    16.861 nll 2.8250 win 16  7s
+Checkpoint activations: False
+built topk=15 attn=moba layer0.top_k=15 (compiling flex...)
+topk  15 gutenberg L30720 ppl    24.712 nll 3.2073 win 16  54s
+topk  15 arxiv     L30720 ppl    67.401 nll 4.2107 win 16  6s
+topk  15 pubmed    L30720 ppl    16.393 nll 2.7968 win 16  6s
+Checkpoint activations: False
+built topk=8 attn=moba layer0.top_k=8 (compiling flex...)
+topk   8 gutenberg L30720 ppl    25.779 nll 3.2496 win 16  54s
+topk   8 arxiv     L30720 ppl    73.672 nll 4.2996 win 16  6s
+topk   8 pubmed    L30720 ppl    18.760 nll 2.9317 win 16  6s
+wrote /mloscratch/homes/navasard/sparse/SparseAtt/eval/results/L365M_8K/ppl_resamp365_L30720_gutenberg.tsv
+wrote /mloscratch/homes/navasard/sparse/SparseAtt/eval/results/L365M_8K/ppl_resamp365_L30720_arxiv.tsv
+wrote /mloscratch/homes/navasard/sparse/SparseAtt/eval/results/L365M_8K/ppl_resamp365_L30720_pubmed.tsv
+DONE resamp365

epfl_final/train_logs/fixed_nrb1.log

@@ -0,0 +1,376 @@
+New run: Checkpoints will be saved to /mloscratch/homes/navasard/sparse/checkpoints/run_20260706_040422
+Running config:
+{
+  "model_path": "/mloscratch/homes/navasard/qwen3_06b_base",
+  "model_dtype": "bfloat16",
+  "init_from_scratch": true,
+  "train_path": "/mloscratch/homes/navasard/data/fineweb_large/tokenized",
+  "val_path": null,
+  "max_seq_length": 8192,
+  "num_workers": 8,
+  "output_dir": "/mloscratch/homes/navasard/sparse/checkpoints/run_20260706_040422",
+  "resume_from_checkpoint": null,
+  "force_new_run": true,
+  "num_epochs": 1,
+  "max_steps": null,
+  "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": 10,
+  "seed": 42,
+  "verbose": true,
+  "checkpoint_activations": false,
+  "strategy": "ddp",
+  "mixed_precision": "bf16",
+  "wandb_enabled": true,
+  "wandb_project": "stochastic_sparsity",
+  "wandb_name": "fixed nrb1 diag+1",
+  "wandb_tags": [
+    "qwen3",
+    "1.2b",
+    "sparse_stochastic",
+    "nrb1",
+    "fixed",
+    "extreme"
+  ],
+  "hidden_size": 1536,
+  "intermediate_size": 6144,
+  "num_attention_heads": 12,
+  "num_key_value_heads": 4,
+  "num_hidden_layers": null,
+  "head_dim": null,
+  "attention_type": "sparse_stochastic",
+  "sparse_block_size": 128,
+  "sparse_num_random_blocks": 1,
+  "sparse_mask_pool_size": 1,
+  "sparse_val_dense_compare": false,
+  "sparse_layer_indices": null,
+  "sparse_mask_cache_dir": null,
+  "dense_mode": false,
+  "dense_step_intervals": null
+}
+Loading tokenizer from /mloscratch/homes/navasard/qwen3_06b_base...
+Initializing model from scratch (random weights)...
+Checkpoint activations: False
+Checkpoint activations: False
+Checkpoint activations: False
+
+
+================================================================================
+MODEL ARCHITECTURE
+================================================================================
+Model config:
+  vocab_size: 151936
+  hidden_size: 1536
+  num_hidden_layers: 28
+  num_attention_heads: 12
+  num_key_value_heads: 4
+  head_dim: 128
+  intermediate_size: 6144
+  max_position_embeddings: 32768
+  rope_theta: 1000000
+  attention_type: sparse_stochastic
+  sparse_block_size: 128
+  sparse_num_random_blocks: 1
+  sparse_mask_pool_size: 1
+  sparse_max_seq_length: 8192
+
+Parameter count:
+  Total parameters: 1,202,352,640
+  Trainable parameters: 1,202,352,640
+  Non-trainable parameters: 0
+  Model size (BF16): 2.40 GB
+
+--------------------------------------------------------------------------------
+LAYER-BY-LAYER SUMMARY
+--------------------------------------------------------------------------------
+/mloscratch/homes/navasard/pylibs/lib/python3.10/site-packages/torch/nn/attention/flex_attention.py:2055: UserWarning: flex_attention called without torch.compile() - this will use an unfused implementation that materializes the full scores matrix instead of generating a fused kernel.
+
+SOLUTION: Use torch.compile(flex_attention)(...)
+
+If you want to debug your score_mod/mask_mod, you can set:
+torch.nn.attention.flex_attention._FLEX_ATTENTION_DISABLE_COMPILE_DEBUG = True
+
+This will allow you to use print statements or breakpoints. Note: This doesn't work with the backwards pass and may produce incorrect results.
+  _warn_once(
+================================================================================================================================================================
+Layer (type:depth-idx)                                       Input Shape               Output Shape              Param #                   Trainable
+================================================================================================================================================================
+Qwen3ForCausalLM                                             [2, 8192]                 --                        --                        True
+├─Qwen3Model: 1-1                                            --                        [2, 8192, 1536]           --                        True
+│    └─Embedding: 2-1                                        [2, 8192]                 [2, 8192, 1536]           233,373,696               True
+│    └─ModuleList: 2-2                                       --                        --                        --                        True
+│    │    └─Qwen3DecoderLayer: 3-1                           [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-2                           [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-3                           [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-4                           [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-5                           [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-6                           [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-7                           [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-8                           [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-9                           [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-10                          [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-11                          [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-12                          [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-13                          [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-14                          [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-15                          [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-16                          [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-17                          [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-18                          [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-19                          [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-20                          [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-21                          [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-22                          [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-23                          [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-24                          [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-25                          [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-26                          [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-27                          [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-28                          [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    └─Qwen3RMSNorm: 2-3                                     [2, 8192, 1536]           [2, 8192, 1536]           1,536                     True
+├─Linear: 1-2                                                [2, 8192, 1536]           [2, 8192, 151936]         233,373,696               True
+================================================================================================================================================================
+Total params: 1,435,726,336
+Trainable params: 1,435,726,336
+Non-trainable params: 0
+Total mult-adds (G): 2.87
+================================================================================================================================================================
+Input size (MB): 0.13
+Forward/backward pass size (MB): 52437.19
+Params size (MB): 2871.45
+Estimated Total Size (MB): 55308.77
+================================================================================================================================================================
+  NOTE: torchinfo double-counts tied embed_tokens/lm_head weight (233,373,696 params). True unique parameter count: 1,202,352,640 (1202.4M)
+================================================================================
+
+Loading dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Using pre-tokenized dataset (no runtime tokenization)
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Loaded 12748481 pre-tokenized examples (stored_len=2048, pack_factor=4, packed_examples=3187120)Loaded 12748481 pre-tokenized examples (stored_len=2048, pack_factor=4, packed_examples=3187120)
+
+Loaded 12748481 pre-tokenized examples (stored_len=2048, pack_factor=4, packed_examples=3187120)
+Dataset split: 3027764 train, 159356 validation
+Training examples (packed, train split):  3027764
+Validation examples (packed, val split):  159356
+Per-device batch size:                    2
+World size (GPUs):                        3
+Gradient accumulation steps:              10
+Effective batch size (seqs/step):         60
+Sequence length:                          8192
+Tokens per optimizer step:                491,520
+Micro-batches per epoch (per rank):       504627
+Optimizer steps per epoch:                50462
+Total optimizer steps:                    50462
+wandb: [wandb.login()] Loaded credentials for https://api.wandb.ai from WANDB_API_KEY.
+wandb: Currently logged in as: ruben-navasardyan4869 (ruben-navasardyan4869-epfl) to https://api.wandb.ai. Use `wandb login --relogin` to force relogin
+wandb: Tracking run with wandb version 0.28.0
+wandb: Run data is saved locally in /mloscratch/homes/navasard/sparse/SparseAtt/wandb/run-20260706_040451-i83whqph
+wandb: Run `wandb offline` to turn off syncing.
+wandb: Syncing run fixed nrb1 diag+1
+wandb: ⭐️ View project at https://wandb.ai/ruben-navasardyan4869-epfl/stochastic_sparsity
+wandb: 🚀 View run at https://wandb.ai/ruben-navasardyan4869-epfl/stochastic_sparsity/runs/i83whqph
+Starting training...
+Switching attention mode for step 1: sparse_stochastic (configured sparse attention)
+Step 10/50462 | Loss: 12.2484 | LR: 6.67e-06 | 7527.4ms/step | 8.0 samples/s
+Step 20/50462 | Loss: 12.1240 | LR: 1.33e-05 | 2683.0ms/step | 22.4 samples/s
+Step 30/50462 | Loss: 11.5407 | LR: 2.00e-05 | 2682.9ms/step | 22.4 samples/s
+Step 40/50462 | Loss: 10.4363 | LR: 2.67e-05 | 2681.6ms/step | 22.4 samples/s
+Step 50/50462 | Loss: 9.9149 | LR: 3.33e-05 | 2679.8ms/step | 22.4 samples/s
+Step 60/50462 | Loss: 9.4867 | LR: 4.00e-05 | 2679.1ms/step | 22.4 samples/s
+Step 70/50462 | Loss: 9.0238 | LR: 4.67e-05 | 2743.8ms/step | 21.9 samples/s
+Step 80/50462 | Loss: 8.4818 | LR: 5.33e-05 | 2675.4ms/step | 22.4 samples/s
+Step 90/50462 | Loss: 7.9989 | LR: 6.00e-05 | 2681.4ms/step | 22.4 samples/s
+Step 100/50462 | Loss: 7.5997 | LR: 6.67e-05 | 2682.5ms/step | 22.4 samples/s
+Step 110/50462 | Loss: 7.3385 | LR: 7.33e-05 | 2687.3ms/step | 22.3 samples/s
+Step 120/50462 | Loss: 7.1860 | LR: 8.00e-05 | 2690.4ms/step | 22.3 samples/s
+Step 130/50462 | Loss: 7.0646 | LR: 8.67e-05 | 2692.3ms/step | 22.3 samples/s
+Step 140/50462 | Loss: 6.9363 | LR: 9.33e-05 | 2695.4ms/step | 22.3 samples/s
+Step 150/50462 | Loss: 6.8000 | LR: 1.00e-04 | 2694.6ms/step | 22.3 samples/s
+Step 160/50462 | Loss: 6.6806 | LR: 1.07e-04 | 2696.4ms/step | 22.3 samples/s
+Step 170/50462 | Loss: 6.5531 | LR: 1.13e-04 | 2696.5ms/step | 22.3 samples/s
+Step 180/50462 | Loss: 6.4830 | LR: 1.20e-04 | 2699.6ms/step | 22.2 samples/s
+Step 190/50462 | Loss: 6.4046 | LR: 1.27e-04 | 2698.2ms/step | 22.2 samples/s
+Step 200/50462 | Loss: 6.3157 | LR: 1.33e-04 | 2700.1ms/step | 22.2 samples/s
+Step 210/50462 | Loss: 6.2217 | LR: 1.40e-04 | 2699.9ms/step | 22.2 samples/s
+Step 220/50462 | Loss: 6.1521 | LR: 1.47e-04 | 2700.2ms/step | 22.2 samples/s
+Step 230/50462 | Loss: 6.1420 | LR: 1.53e-04 | 2698.8ms/step | 22.2 samples/s
+Step 240/50462 | Loss: 6.0774 | LR: 1.60e-04 | 2698.4ms/step | 22.2 samples/s
+Step 250/50462 | Loss: 6.0141 | LR: 1.67e-04 | 2701.0ms/step | 22.2 samples/s
+Validation at step 250: Loss = 5.9621, Perplexity = 388.41
+Step 260/50462 | Loss: 5.9378 | LR: 1.73e-04 | 6740.0ms/step | 8.9 samples/s
+Step 270/50462 | Loss: 5.9469 | LR: 1.80e-04 | 2689.4ms/step | 22.3 samples/s
+Step 280/50462 | Loss: 5.8680 | LR: 1.87e-04 | 2697.2ms/step | 22.2 samples/s
+Step 290/50462 | Loss: 5.8074 | LR: 1.93e-04 | 2700.4ms/step | 22.2 samples/s
+Step 300/50462 | Loss: 5.7335 | LR: 2.00e-04 | 2701.3ms/step | 22.2 samples/s
+Step 310/50462 | Loss: 5.6945 | LR: 2.07e-04 | 2703.0ms/step | 22.2 samples/s
+Step 320/50462 | Loss: 5.6354 | LR: 2.13e-04 | 2703.4ms/step | 22.2 samples/s
+Step 330/50462 | Loss: 5.6109 | LR: 2.20e-04 | 2702.0ms/step | 22.2 samples/s
+Step 340/50462 | Loss: 5.5531 | LR: 2.27e-04 | 2698.8ms/step | 22.2 samples/s
+Step 350/50462 | Loss: 5.5068 | LR: 2.33e-04 | 2698.5ms/step | 22.2 samples/s
+Step 360/50462 | Loss: 5.4480 | LR: 2.40e-04 | 2702.7ms/step | 22.2 samples/s
+Step 370/50462 | Loss: 5.3898 | LR: 2.47e-04 | 2703.0ms/step | 22.2 samples/s
+Step 380/50462 | Loss: 5.3795 | LR: 2.53e-04 | 2706.3ms/step | 22.2 samples/s
+Step 390/50462 | Loss: 5.3246 | LR: 2.60e-04 | 2702.8ms/step | 22.2 samples/s
+Step 400/50462 | Loss: 5.2640 | LR: 2.67e-04 | 2702.4ms/step | 22.2 samples/s
+Step 410/50462 | Loss: 5.2326 | LR: 2.73e-04 | 2705.3ms/step | 22.2 samples/s
+Step 420/50462 | Loss: 5.2115 | LR: 2.80e-04 | 2702.7ms/step | 22.2 samples/s
+Step 430/50462 | Loss: 5.1535 | LR: 2.87e-04 | 2702.3ms/step | 22.2 samples/s
+Step 440/50462 | Loss: 5.1121 | LR: 2.93e-04 | 2701.3ms/step | 22.2 samples/s
+Step 450/50462 | Loss: 5.0737 | LR: 3.00e-04 | 2700.4ms/step | 22.2 samples/s
+Step 460/50462 | Loss: 5.0504 | LR: 3.07e-04 | 2768.2ms/step | 21.7 samples/s
+Step 470/50462 | Loss: 5.0068 | LR: 3.13e-04 | 2698.1ms/step | 22.2 samples/s
+Step 480/50462 | Loss: 4.9795 | LR: 3.20e-04 | 2701.8ms/step | 22.2 samples/s
+Step 490/50462 | Loss: 4.9861 | LR: 3.27e-04 | 2702.4ms/step | 22.2 samples/s
+Step 500/50462 | Loss: 4.9464 | LR: 3.33e-04 | 2700.1ms/step | 22.2 samples/s
+Validation at step 500: Loss = 4.8870, Perplexity = 132.56
+Step 510/50462 | Loss: 4.8661 | LR: 3.40e-04 | 4150.5ms/step | 14.5 samples/s
+Step 520/50462 | Loss: 4.8597 | LR: 3.47e-04 | 2699.1ms/step | 22.2 samples/s
+Step 530/50462 | Loss: 4.8921 | LR: 3.53e-04 | 2766.5ms/step | 21.7 samples/s
+Step 540/50462 | Loss: 4.7909 | LR: 3.60e-04 | 2703.6ms/step | 22.2 samples/s
+Step 550/50462 | Loss: 4.7424 | LR: 3.67e-04 | 2770.0ms/step | 21.7 samples/s
+Step 560/50462 | Loss: 4.7633 | LR: 3.73e-04 | 2704.4ms/step | 22.2 samples/s
+Step 570/50462 | Loss: 4.7767 | LR: 3.80e-04 | 2703.8ms/step | 22.2 samples/s
+Step 580/50462 | Loss: 4.6568 | LR: 3.87e-04 | 2703.5ms/step | 22.2 samples/s
+Step 590/50462 | Loss: 4.6375 | LR: 3.93e-04 | 2703.8ms/step | 22.2 samples/s
+Step 600/50462 | Loss: 4.6142 | LR: 4.00e-04 | 2701.1ms/step | 22.2 samples/s
+Step 610/50462 | Loss: 4.6426 | LR: 4.07e-04 | 2701.4ms/step | 22.2 samples/s
+Step 620/50462 | Loss: 4.6182 | LR: 4.13e-04 | 2702.1ms/step | 22.2 samples/s
+Step 630/50462 | Loss: 4.5529 | LR: 4.20e-04 | 2701.9ms/step | 22.2 samples/s
+Step 640/50462 | Loss: 4.5129 | LR: 4.27e-04 | 2704.7ms/step | 22.2 samples/s
+Step 650/50462 | Loss: 4.4876 | LR: 4.33e-04 | 2704.0ms/step | 22.2 samples/s
+Step 660/50462 | Loss: 4.4495 | LR: 4.40e-04 | 2702.7ms/step | 22.2 samples/s
+Step 670/50462 | Loss: 4.4391 | LR: 4.47e-04 | 2700.4ms/step | 22.2 samples/s
+Step 680/50462 | Loss: 4.4193 | LR: 4.53e-04 | 2700.4ms/step | 22.2 samples/s
+Step 690/50462 | Loss: 4.3732 | LR: 4.60e-04 | 2702.2ms/step | 22.2 samples/s
+Step 700/50462 | Loss: 4.3868 | LR: 4.67e-04 | 2703.1ms/step | 22.2 samples/s
+Step 710/50462 | Loss: 4.3452 | LR: 4.73e-04 | 2705.3ms/step | 22.2 samples/s
+Step 720/50462 | Loss: 4.3096 | LR: 4.80e-04 | 2702.9ms/step | 22.2 samples/s
+Step 730/50462 | Loss: 4.2990 | LR: 4.87e-04 | 2703.7ms/step | 22.2 samples/s
+Step 740/50462 | Loss: 4.2616 | LR: 4.93e-04 | 2704.9ms/step | 22.2 samples/s
+Step 750/50462 | Loss: 4.2783 | LR: 5.00e-04 | 2703.0ms/step | 22.2 samples/s
+Validation at step 750: Loss = 4.2548, Perplexity = 70.44
+Step 760/50462 | Loss: 4.2222 | LR: 5.00e-04 | 4148.5ms/step | 14.5 samples/s
+Step 770/50462 | Loss: 4.2413 | LR: 5.00e-04 | 2700.3ms/step | 22.2 samples/s
+Step 780/50462 | Loss: 4.2028 | LR: 5.00e-04 | 2704.8ms/step | 22.2 samples/s
+Step 790/50462 | Loss: 4.2029 | LR: 5.00e-04 | 2704.6ms/step | 22.2 samples/s
+Step 800/50462 | Loss: 4.1514 | LR: 5.00e-04 | 2756.9ms/step | 21.8 samples/s
+Step 810/50462 | Loss: 4.1582 | LR: 5.00e-04 | 2702.5ms/step | 22.2 samples/s
+Step 820/50462 | Loss: 4.1472 | LR: 5.00e-04 | 2704.7ms/step | 22.2 samples/s
+Step 830/50462 | Loss: 4.1297 | LR: 5.00e-04 | 2704.7ms/step | 22.2 samples/s
+Step 840/50462 | Loss: 4.0992 | LR: 5.00e-04 | 2705.1ms/step | 22.2 samples/s
+Step 850/50462 | Loss: 4.0617 | LR: 5.00e-04 | 2704.1ms/step | 22.2 samples/s
+Step 860/50462 | Loss: 4.0601 | LR: 5.00e-04 | 2702.3ms/step | 22.2 samples/s
+Step 870/50462 | Loss: 4.0459 | LR: 5.00e-04 | 2760.6ms/step | 21.7 samples/s
+Step 880/50462 | Loss: 4.0262 | LR: 5.00e-04 | 2757.3ms/step | 21.8 samples/s
+Step 890/50462 | Loss: 4.0214 | LR: 5.00e-04 | 2702.7ms/step | 22.2 samples/s
+Step 900/50462 | Loss: 4.0181 | LR: 5.00e-04 | 2696.9ms/step | 22.2 samples/s
+Step 910/50462 | Loss: 3.9917 | LR: 5.00e-04 | 2694.0ms/step | 22.3 samples/s
+Step 920/50462 | Loss: 3.9486 | LR: 5.00e-04 | 2697.5ms/step | 22.2 samples/s
+Step 930/50462 | Loss: 4.0074 | LR: 5.00e-04 | 2703.6ms/step | 22.2 samples/s
+Step 940/50462 | Loss: 3.9516 | LR: 5.00e-04 | 2706.2ms/step | 22.2 samples/s
+Step 950/50462 | Loss: 3.9381 | LR: 5.00e-04 | 2707.8ms/step | 22.2 samples/s
+Step 960/50462 | Loss: 3.9388 | LR: 5.00e-04 | 2708.8ms/step | 22.2 samples/s
+Step 970/50462 | Loss: 3.9318 | LR: 5.00e-04 | 2709.3ms/step | 22.1 samples/s
+Step 980/50462 | Loss: 3.8942 | LR: 5.00e-04 | 2710.0ms/step | 22.1 samples/s
+Step 990/50462 | Loss: 3.9123 | LR: 5.00e-04 | 2710.5ms/step | 22.1 samples/s
+Step 1000/50462 | Loss: 3.9161 | LR: 5.00e-04 | 2712.9ms/step | 22.1 samples/s
+Validation at step 1000: Loss = 3.9162, Perplexity = 50.21
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/run_20260706_040422/checkpoint-1000
+Step 1010/50462 | Loss: 3.8950 | LR: 5.00e-04 | 5064.2ms/step | 11.8 samples/s
+Step 1020/50462 | Loss: 3.9062 | LR: 5.00e-04 | 2700.0ms/step | 22.2 samples/s
+Step 1030/50462 | Loss: 3.8888 | LR: 5.00e-04 | 2712.5ms/step | 22.1 samples/s
+Step 1040/50462 | Loss: 3.8945 | LR: 5.00e-04 | 2714.8ms/step | 22.1 samples/s
+Step 1050/50462 | Loss: 3.8628 | LR: 5.00e-04 | 2715.6ms/step | 22.1 samples/s
+Step 1060/50462 | Loss: 3.8635 | LR: 5.00e-04 | 2713.4ms/step | 22.1 samples/s
+Step 1070/50462 | Loss: 3.8484 | LR: 5.00e-04 | 2713.7ms/step | 22.1 samples/s
+Step 1080/50462 | Loss: 3.8564 | LR: 5.00e-04 | 2713.0ms/step | 22.1 samples/s
+Step 1090/50462 | Loss: 3.8182 | LR: 5.00e-04 | 2714.4ms/step | 22.1 samples/s
+Step 1100/50462 | Loss: 3.8417 | LR: 5.00e-04 | 2712.7ms/step | 22.1 samples/s
+Step 1110/50462 | Loss: 3.8512 | LR: 5.00e-04 | 2711.5ms/step | 22.1 samples/s
+Step 1120/50462 | Loss: 3.8053 | LR: 5.00e-04 | 2715.3ms/step | 22.1 samples/s
+Step 1130/50462 | Loss: 3.8228 | LR: 5.00e-04 | 2769.0ms/step | 21.7 samples/s
+Step 1140/50462 | Loss: 3.7991 | LR: 5.00e-04 | 2712.7ms/step | 22.1 samples/s
+Step 1150/50462 | Loss: 3.7665 | LR: 5.00e-04 | 2712.0ms/step | 22.1 samples/s
+Step 1160/50462 | Loss: 3.7871 | LR: 5.00e-04 | 2713.2ms/step | 22.1 samples/s
+Step 1170/50462 | Loss: 3.7761 | LR: 5.00e-04 | 2713.2ms/step | 22.1 samples/s
+Step 1180/50462 | Loss: 3.7418 | LR: 5.00e-04 | 2710.2ms/step | 22.1 samples/s
+Step 1190/50462 | Loss: 3.7695 | LR: 5.00e-04 | 2707.3ms/step | 22.2 samples/s
+Step 1200/50462 | Loss: 3.7765 | LR: 5.00e-04 | 2763.9ms/step | 21.7 samples/s
+Step 1210/50462 | Loss: 3.7845 | LR: 5.00e-04 | 2708.6ms/step | 22.2 samples/s
+Step 1220/50462 | Loss: 3.7430 | LR: 5.00e-04 | 2766.0ms/step | 21.7 samples/s
+Step 1230/50462 | Loss: 3.7115 | LR: 5.00e-04 | 2713.7ms/step | 22.1 samples/s
+Step 1240/50462 | Loss: 3.7684 | LR: 5.00e-04 | 2713.6ms/step | 22.1 samples/s
+Step 1250/50462 | Loss: 3.7305 | LR: 5.00e-04 | 2711.0ms/step | 22.1 samples/s
+Validation at step 1250: Loss = 3.7457, Perplexity = 42.34
+Step 1260/50462 | Loss: 3.7549 | LR: 5.00e-04 | 4160.7ms/step | 14.4 samples/s
+Step 1270/50462 | Loss: 3.7178 | LR: 5.00e-04 | 2713.5ms/step | 22.1 samples/s
+Step 1280/50462 | Loss: 3.7033 | LR: 5.00e-04 | 2714.4ms/step | 22.1 samples/s
+Step 1290/50462 | Loss: 3.7069 | LR: 5.00e-04 | 2713.7ms/step | 22.1 samples/s
+Step 1300/50462 | Loss: 3.6999 | LR: 5.00e-04 | 2715.0ms/step | 22.1 samples/s
+Step 1310/50462 | Loss: 3.7032 | LR: 5.00e-04 | 2709.7ms/step | 22.1 samples/s
+Step 1320/50462 | Loss: 3.6988 | LR: 5.00e-04 | 2715.1ms/step | 22.1 samples/s
+Step 1330/50462 | Loss: 3.6967 | LR: 5.00e-04 | 2714.1ms/step | 22.1 samples/s
+Step 1340/50462 | Loss: 3.6827 | LR: 5.00e-04 | 2712.9ms/step | 22.1 samples/s
+Step 1350/50462 | Loss: 3.6665 | LR: 5.00e-04 | 2711.9ms/step | 22.1 samples/s
+Step 1360/50462 | Loss: 3.6848 | LR: 5.00e-04 | 2711.4ms/step | 22.1 samples/s
+Step 1370/50462 | Loss: 3.6631 | LR: 5.00e-04 | 2710.6ms/step | 22.1 samples/s
+Step 1380/50462 | Loss: 3.7003 | LR: 5.00e-04 | 2713.4ms/step | 22.1 samples/s
+Step 1390/50462 | Loss: 3.6871 | LR: 5.00e-04 | 2713.9ms/step | 22.1 samples/s
+Step 1400/50462 | Loss: 3.6462 | LR: 5.00e-04 | 2714.1ms/step | 22.1 samples/s
+Step 1410/50462 | Loss: 3.6726 | LR: 5.00e-04 | 2712.8ms/step | 22.1 samples/s
+Step 1420/50462 | Loss: 3.6439 | LR: 5.00e-04 | 2714.6ms/step | 22.1 samples/s
+Step 1430/50462 | Loss: 3.6282 | LR: 5.00e-04 | 2713.8ms/step | 22.1 samples/s
+Step 1440/50462 | Loss: 3.6609 | LR: 5.00e-04 | 2715.1ms/step | 22.1 samples/s
+Step 1450/50462 | Loss: 3.6282 | LR: 5.00e-04 | 2714.9ms/step | 22.1 samples/s
+Step 1460/50462 | Loss: 3.6319 | LR: 5.00e-04 | 2767.9ms/step | 21.7 samples/s
+Step 1470/50462 | Loss: 3.6127 | LR: 5.00e-04 | 2713.4ms/step | 22.1 samples/s
+Step 1480/50462 | Loss: 3.6521 | LR: 5.00e-04 | 2714.7ms/step | 22.1 samples/s
+Step 1490/50462 | Loss: 3.6304 | LR: 5.00e-04 | 2712.8ms/step | 22.1 samples/s
+Step 1500/50462 | Loss: 3.5960 | LR: 5.00e-04 | 2716.0ms/step | 22.1 samples/s
+Validation at step 1500: Loss = 3.6393, Perplexity = 38.06
+Step 1510/50462 | Loss: 3.6383 | LR: 5.00e-04 | 4173.6ms/step | 14.4 samples/s
+Step 1520/50462 | Loss: 3.6056 | LR: 5.00e-04 | 2716.1ms/step | 22.1 samples/s
+Step 1530/50462 | Loss: 3.6322 | LR: 5.00e-04 | 2718.6ms/step | 22.1 samples/s
+Step 1540/50462 | Loss: 3.5982 | LR: 5.00e-04 | 2774.0ms/step | 21.6 samples/s
+Step 1550/50462 | Loss: 3.6089 | LR: 5.00e-04 | 2712.1ms/step | 22.1 samples/s
+Step 1560/50462 | Loss: 3.5995 | LR: 5.00e-04 | 2768.0ms/step | 21.7 samples/s
+Step 1570/50462 | Loss: 3.6266 | LR: 5.00e-04 | 2714.6ms/step | 22.1 samples/s
+Step 1580/50462 | Loss: 3.5970 | LR: 5.00e-04 | 2714.8ms/step | 22.1 samples/s
+W0706 05:19:15.162000 210 torch/distributed/elastic/agent/server/api.py:753] Received Signals.SIGTERM death signal, shutting down workers
+W0706 05:19:15.168000 210 torch/distributed/elastic/multiprocessing/api.py:1014] Sending process 221 closing signal SIGTERM
+Traceback (most recent call last):
+  File "/usr/lib/python3.10/runpy.py", line 196, in _run_module_as_main
+    return _run_code(code, main_globals, None,
+  File "/usr/lib/python3.10/runpy.py", line 86, in _run_code
+    exec(code, run_globals)
+  File "/mloscratch/homes/navasard/pylibs/lib/python3.10/site-packages/torch/distributed/run.py", line 1020, in <module>
+    main()
+  File "/mloscratch/homes/navasard/pylibs/lib/python3.10/site-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 367, in wrapper
+    return f(*args, **kwargs)
+  File "/mloscratch/homes/navasard/pylibs/lib/python3.10/site-packages/torch/distributed/run.py", line 1016, in main
+    run(args)
+  File "/mloscratch/homes/navasard/pylibs/lib/python3.10/site-packages/torch/distributed/run.py", line 1007, in run
+    elastic_launch(
+  File "/mloscratch/homes/navasard/pylibs/lib/python3.10/site-packages/torch/distributed/launcher/api.py", line 191, in __call__
+    return launch_agent(self._config, self._entrypoint, list(args))
+  File "/mloscratch/homes/navasard/pylibs/lib/python3.10/site-packages/torch/distributed/launcher/api.py", line 362, in launch_agent
+    result = agent.run()
+  File "/mloscratch/homes/navasard/pylibs/lib/python3.10/site-packages/torch/distributed/elastic/metrics/api.py", line 134, in wrapper
+    result = f(*args, **kwargs)
+  File "/mloscratch/homes/navasard/pylibs/lib/python3.10/site-packages/torch/distributed/elastic/agent/server/api.py", line 745, in run
+    result = self._invoke_run(role)
+  File "/mloscratch/homes/navasard/pylibs/lib/python3.10/site-packages/torch/distributed/elastic/agent/server/api.py", line 923, in _invoke_run
+    time.sleep(monitor_interval)
+  File "/mloscratch/homes/navasard/pylibs/lib/python3.10/site-packages/torch/distributed/elastic/multiprocessing/api.py", line 86, in _terminate_process_handler
+    raise SignalException(f"Process {os.getpid()} got signal: {sigval}", sigval=sigval)
+torch.distributed.elastic.multiprocessing.api.SignalException: Process 210 got signal: 15

epfl_final/train_logs/resampled_nrb1.log

@@ -0,0 +1,401 @@
+New run: Checkpoints will be saved to /mloscratch/homes/navasard/sparse/checkpoints/run_20260706_035226
+Running config:
+{
+  "model_path": "/mloscratch/homes/navasard/qwen3_06b_base",
+  "model_dtype": "bfloat16",
+  "init_from_scratch": true,
+  "train_path": "/mloscratch/homes/navasard/data/fineweb_large/tokenized",
+  "val_path": null,
+  "max_seq_length": 8192,
+  "num_workers": 8,
+  "output_dir": "/mloscratch/homes/navasard/sparse/checkpoints/run_20260706_035226",
+  "resume_from_checkpoint": null,
+  "force_new_run": true,
+  "num_epochs": 1,
+  "max_steps": null,
+  "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": 10,
+  "seed": 42,
+  "verbose": true,
+  "checkpoint_activations": false,
+  "strategy": "ddp",
+  "mixed_precision": "bf16",
+  "wandb_enabled": true,
+  "wandb_project": "stochastic_sparsity",
+  "wandb_name": "resampled nrb1 diag+1",
+  "wandb_tags": [
+    "qwen3",
+    "1.2b",
+    "sparse_stochastic",
+    "nrb1",
+    "resampled",
+    "extreme"
+  ],
+  "hidden_size": 1536,
+  "intermediate_size": 6144,
+  "num_attention_heads": 12,
+  "num_key_value_heads": 4,
+  "num_hidden_layers": null,
+  "head_dim": null,
+  "attention_type": "sparse_stochastic",
+  "sparse_block_size": 128,
+  "sparse_num_random_blocks": 1,
+  "sparse_mask_pool_size": 1000,
+  "sparse_val_dense_compare": false,
+  "sparse_layer_indices": null,
+  "sparse_mask_cache_dir": "/mloscratch/homes/navasard/sparse/mask_cache_nrb1",
+  "dense_mode": false,
+  "dense_step_intervals": null
+}
+Loading tokenizer from /mloscratch/homes/navasard/qwen3_06b_base...
+Initializing model from scratch (random weights)...
+Checkpoint activations: False
+Checkpoint activations: False
+Checkpoint activations: False
+
+
+================================================================================
+MODEL ARCHITECTURE
+================================================================================
+Model config:
+  vocab_size: 151936
+  hidden_size: 1536
+  num_hidden_layers: 28
+  num_attention_heads: 12
+  num_key_value_heads: 4
+  head_dim: 128
+  intermediate_size: 6144
+  max_position_embeddings: 32768
+  rope_theta: 1000000
+  attention_type: sparse_stochastic
+  sparse_block_size: 128
+  sparse_num_random_blocks: 1
+  sparse_mask_pool_size: 1000
+  sparse_max_seq_length: 8192
+
+Parameter count:
+  Total parameters: 1,202,352,640
+  Trainable parameters: 1,202,352,640
+  Non-trainable parameters: 0
+  Model size (BF16): 2.40 GB
+
+--------------------------------------------------------------------------------
+LAYER-BY-LAYER SUMMARY
+--------------------------------------------------------------------------------
+/mloscratch/homes/navasard/pylibs/lib/python3.10/site-packages/torch/nn/attention/flex_attention.py:2055: UserWarning: flex_attention called without torch.compile() - this will use an unfused implementation that materializes the full scores matrix instead of generating a fused kernel.
+
+SOLUTION: Use torch.compile(flex_attention)(...)
+
+If you want to debug your score_mod/mask_mod, you can set:
+torch.nn.attention.flex_attention._FLEX_ATTENTION_DISABLE_COMPILE_DEBUG = True
+
+This will allow you to use print statements or breakpoints. Note: This doesn't work with the backwards pass and may produce incorrect results.
+  _warn_once(
+================================================================================================================================================================
+Layer (type:depth-idx)                                       Input Shape               Output Shape              Param #                   Trainable
+================================================================================================================================================================
+Qwen3ForCausalLM                                             [2, 8192]                 --                        --                        True
+├─Qwen3Model: 1-1                                            --                        [2, 8192, 1536]           --                        True
+│    └─Embedding: 2-1                                        [2, 8192]                 [2, 8192, 1536]           233,373,696               True
+│    └─ModuleList: 2-2                                       --                        --                        --                        True
+│    │    └─Qwen3DecoderLayer: 3-1                           [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-2                           [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-3                           [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-4                           [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-5                           [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-6                           [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-7                           [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-8                           [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-9                           [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-10                          [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-11                          [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-12                          [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-13                          [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-14                          [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-15                          [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-16                          [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-17                          [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-18                          [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-19                          [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-20                          [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-21                          [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-22                          [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-23                          [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-24                          [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-25                          [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-26                          [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-27                          [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-28                          [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    └─Qwen3RMSNorm: 2-3                                     [2, 8192, 1536]           [2, 8192, 1536]           1,536                     True
+├─Linear: 1-2                                                [2, 8192, 1536]           [2, 8192, 151936]         233,373,696               True
+================================================================================================================================================================
+Total params: 1,435,726,336
+Trainable params: 1,435,726,336
+Non-trainable params: 0
+Total mult-adds (G): 2.87
+================================================================================================================================================================
+Input size (MB): 0.13
+Forward/backward pass size (MB): 52437.19
+Params size (MB): 2871.45
+Estimated Total Size (MB): 55308.77
+================================================================================================================================================================
+  NOTE: torchinfo double-counts tied embed_tokens/lm_head weight (233,373,696 params). True unique parameter count: 1,202,352,640 (1202.4M)
+================================================================================
+
+Loading dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Using pre-tokenized dataset (no runtime tokenization)
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Loaded 12748481 pre-tokenized examples (stored_len=2048, pack_factor=4, packed_examples=3187120)
+Loaded 12748481 pre-tokenized examples (stored_len=2048, pack_factor=4, packed_examples=3187120)
+Loaded 12748481 pre-tokenized examples (stored_len=2048, pack_factor=4, packed_examples=3187120)
+Dataset split: 3027764 train, 159356 validation
+Training examples (packed, train split):  3027764
+Validation examples (packed, val split):  159356
+Per-device batch size:                    2
+World size (GPUs):                        3
+Gradient accumulation steps:              10
+Effective batch size (seqs/step):         60
+Sequence length:                          8192
+Tokens per optimizer step:                491,520
+Micro-batches per epoch (per rank):       504627
+Optimizer steps per epoch:                50462
+Total optimizer steps:                    50462
+wandb: [wandb.login()] Loaded credentials for https://api.wandb.ai from WANDB_API_KEY.
+wandb: Currently logged in as: ruben-navasardyan4869 (ruben-navasardyan4869-epfl) to https://api.wandb.ai. Use `wandb login --relogin` to force relogin
+wandb: Tracking run with wandb version 0.28.0
+wandb: Run data is saved locally in /mloscratch/homes/navasard/sparse/SparseAtt/wandb/run-20260706_035257-1pbbmh59
+wandb: Run `wandb offline` to turn off syncing.
+wandb: Syncing run resampled nrb1 diag+1
+wandb: ⭐️ View project at https://wandb.ai/ruben-navasardyan4869-epfl/stochastic_sparsity
+wandb: 🚀 View run at https://wandb.ai/ruben-navasardyan4869-epfl/stochastic_sparsity/runs/1pbbmh59
+Starting training...
+Switching attention mode for step 1: sparse_stochastic (configured sparse attention)
+Step 10/50462 | Loss: 12.2487 | LR: 6.67e-06 | 7849.5ms/step | 7.6 samples/s
+Step 20/50462 | Loss: 12.1281 | LR: 1.33e-05 | 2712.8ms/step | 22.1 samples/s
+Step 30/50462 | Loss: 11.5520 | LR: 2.00e-05 | 2712.9ms/step | 22.1 samples/s
+Step 40/50462 | Loss: 10.4438 | LR: 2.67e-05 | 2707.7ms/step | 22.2 samples/s
+Step 50/50462 | Loss: 9.9194 | LR: 3.33e-05 | 2706.0ms/step | 22.2 samples/s
+Step 60/50462 | Loss: 9.4912 | LR: 4.00e-05 | 2703.6ms/step | 22.2 samples/s
+Step 70/50462 | Loss: 9.0321 | LR: 4.67e-05 | 2704.8ms/step | 22.2 samples/s
+Step 80/50462 | Loss: 8.4877 | LR: 5.33e-05 | 2708.9ms/step | 22.1 samples/s
+Step 90/50462 | Loss: 8.0049 | LR: 6.00e-05 | 2711.8ms/step | 22.1 samples/s
+Step 100/50462 | Loss: 7.6035 | LR: 6.67e-05 | 2712.2ms/step | 22.1 samples/s
+Step 110/50462 | Loss: 7.3411 | LR: 7.33e-05 | 2711.7ms/step | 22.1 samples/s
+Step 120/50462 | Loss: 7.1885 | LR: 8.00e-05 | 2715.8ms/step | 22.1 samples/s
+Step 130/50462 | Loss: 7.0661 | LR: 8.67e-05 | 2718.7ms/step | 22.1 samples/s
+Step 140/50462 | Loss: 6.9361 | LR: 9.33e-05 | 2720.4ms/step | 22.1 samples/s
+Step 150/50462 | Loss: 6.8014 | LR: 1.00e-04 | 2719.8ms/step | 22.1 samples/s
+Step 160/50462 | Loss: 6.6800 | LR: 1.07e-04 | 2720.2ms/step | 22.1 samples/s
+Step 170/50462 | Loss: 6.5537 | LR: 1.13e-04 | 2722.9ms/step | 22.0 samples/s
+Step 180/50462 | Loss: 6.4709 | LR: 1.20e-04 | 2724.4ms/step | 22.0 samples/s
+Step 190/50462 | Loss: 6.3962 | LR: 1.27e-04 | 2723.7ms/step | 22.0 samples/s
+Step 200/50462 | Loss: 6.3103 | LR: 1.33e-04 | 2722.1ms/step | 22.0 samples/s
+Step 210/50462 | Loss: 6.2192 | LR: 1.40e-04 | 2722.3ms/step | 22.0 samples/s
+Step 220/50462 | Loss: 6.1495 | LR: 1.47e-04 | 2721.7ms/step | 22.0 samples/s
+Step 230/50462 | Loss: 6.1285 | LR: 1.53e-04 | 2722.6ms/step | 22.0 samples/s
+Step 240/50462 | Loss: 6.0880 | LR: 1.60e-04 | 2723.4ms/step | 22.0 samples/s
+Step 250/50462 | Loss: 6.0311 | LR: 1.67e-04 | 2723.0ms/step | 22.0 samples/s
+Validation at step 250: Loss = 5.9991, Perplexity = 403.06
+Step 260/50462 | Loss: 5.9490 | LR: 1.73e-04 | 5786.7ms/step | 10.4 samples/s
+Step 270/50462 | Loss: 5.9466 | LR: 1.80e-04 | 2718.8ms/step | 22.1 samples/s
+Step 280/50462 | Loss: 5.8604 | LR: 1.87e-04 | 2722.2ms/step | 22.0 samples/s
+Step 290/50462 | Loss: 5.8283 | LR: 1.93e-04 | 2719.7ms/step | 22.1 samples/s
+Step 300/50462 | Loss: 5.7623 | LR: 2.00e-04 | 2722.4ms/step | 22.0 samples/s
+Step 310/50462 | Loss: 5.7100 | LR: 2.07e-04 | 2726.4ms/step | 22.0 samples/s
+Step 320/50462 | Loss: 5.6344 | LR: 2.13e-04 | 2726.2ms/step | 22.0 samples/s
+Step 330/50462 | Loss: 5.6107 | LR: 2.20e-04 | 2729.6ms/step | 22.0 samples/s
+Step 340/50462 | Loss: 5.5687 | LR: 2.27e-04 | 2728.9ms/step | 22.0 samples/s
+Step 350/50462 | Loss: 5.5140 | LR: 2.33e-04 | 2729.8ms/step | 22.0 samples/s
+Step 360/50462 | Loss: 5.4525 | LR: 2.40e-04 | 2730.0ms/step | 22.0 samples/s
+Step 370/50462 | Loss: 5.4094 | LR: 2.47e-04 | 2730.2ms/step | 22.0 samples/s
+Step 380/50462 | Loss: 5.3872 | LR: 2.53e-04 | 2730.3ms/step | 22.0 samples/s
+Step 390/50462 | Loss: 5.3373 | LR: 2.60e-04 | 2730.0ms/step | 22.0 samples/s
+Step 400/50462 | Loss: 5.2693 | LR: 2.67e-04 | 2731.5ms/step | 22.0 samples/s
+Step 410/50462 | Loss: 5.2467 | LR: 2.73e-04 | 2729.4ms/step | 22.0 samples/s
+Step 420/50462 | Loss: 5.2083 | LR: 2.80e-04 | 2730.9ms/step | 22.0 samples/s
+Step 430/50462 | Loss: 5.1244 | LR: 2.87e-04 | 2731.3ms/step | 22.0 samples/s
+Step 440/50462 | Loss: 5.1650 | LR: 2.93e-04 | 2730.2ms/step | 22.0 samples/s
+Step 450/50462 | Loss: 5.1326 | LR: 3.00e-04 | 2732.6ms/step | 22.0 samples/s
+Step 460/50462 | Loss: 5.0389 | LR: 3.07e-04 | 2730.5ms/step | 22.0 samples/s
+Step 470/50462 | Loss: 5.0087 | LR: 3.13e-04 | 2732.8ms/step | 22.0 samples/s
+Step 480/50462 | Loss: 5.0167 | LR: 3.20e-04 | 2730.7ms/step | 22.0 samples/s
+Step 490/50462 | Loss: 5.0158 | LR: 3.27e-04 | 2731.9ms/step | 22.0 samples/s
+Step 500/50462 | Loss: 4.9535 | LR: 3.33e-04 | 2731.2ms/step | 22.0 samples/s
+Validation at step 500: Loss = 4.8987, Perplexity = 134.11
+Step 510/50462 | Loss: 4.8771 | LR: 3.40e-04 | 4211.5ms/step | 14.2 samples/s
+Step 520/50462 | Loss: 4.8483 | LR: 3.47e-04 | 2732.1ms/step | 22.0 samples/s
+Step 530/50462 | Loss: 4.9048 | LR: 3.53e-04 | 2733.9ms/step | 21.9 samples/s
+Step 540/50462 | Loss: 4.8072 | LR: 3.60e-04 | 2735.6ms/step | 21.9 samples/s
+Step 550/50462 | Loss: 4.7502 | LR: 3.67e-04 | 2843.2ms/step | 21.1 samples/s
+Step 560/50462 | Loss: 4.7328 | LR: 3.73e-04 | 2733.3ms/step | 22.0 samples/s
+Step 570/50462 | Loss: 4.7495 | LR: 3.80e-04 | 2844.7ms/step | 21.1 samples/s
+Step 580/50462 | Loss: 4.6493 | LR: 3.87e-04 | 2736.2ms/step | 21.9 samples/s
+Step 590/50462 | Loss: 4.6332 | LR: 3.93e-04 | 2737.0ms/step | 21.9 samples/s
+Step 600/50462 | Loss: 4.6240 | LR: 4.00e-04 | 2735.6ms/step | 21.9 samples/s
+Step 610/50462 | Loss: 4.5943 | LR: 4.07e-04 | 2736.3ms/step | 21.9 samples/s
+Step 620/50462 | Loss: 4.5894 | LR: 4.13e-04 | 2842.6ms/step | 21.1 samples/s
+Step 630/50462 | Loss: 4.5361 | LR: 4.20e-04 | 2724.5ms/step | 22.0 samples/s
+Step 640/50462 | Loss: 4.5136 | LR: 4.27e-04 | 2729.8ms/step | 22.0 samples/s
+Step 650/50462 | Loss: 4.4822 | LR: 4.33e-04 | 2732.6ms/step | 22.0 samples/s
+Step 660/50462 | Loss: 4.4492 | LR: 4.40e-04 | 2732.5ms/step | 22.0 samples/s
+Step 670/50462 | Loss: 4.4211 | LR: 4.47e-04 | 2731.7ms/step | 22.0 samples/s
+Step 680/50462 | Loss: 4.4125 | LR: 4.53e-04 | 2731.4ms/step | 22.0 samples/s
+Step 690/50462 | Loss: 4.3521 | LR: 4.60e-04 | 2734.3ms/step | 21.9 samples/s
+Step 700/50462 | Loss: 4.3790 | LR: 4.67e-04 | 2733.2ms/step | 22.0 samples/s
+Step 710/50462 | Loss: 4.3428 | LR: 4.73e-04 | 2733.6ms/step | 21.9 samples/s
+Step 720/50462 | Loss: 4.3330 | LR: 4.80e-04 | 2733.0ms/step | 22.0 samples/s
+Step 730/50462 | Loss: 4.2979 | LR: 4.87e-04 | 2733.4ms/step | 22.0 samples/s
+Step 740/50462 | Loss: 4.2526 | LR: 4.93e-04 | 2733.7ms/step | 21.9 samples/s
+Step 750/50462 | Loss: 4.2661 | LR: 5.00e-04 | 2734.8ms/step | 21.9 samples/s
+Validation at step 750: Loss = 4.2387, Perplexity = 69.32
+Step 760/50462 | Loss: 4.2168 | LR: 5.00e-04 | 4204.9ms/step | 14.3 samples/s
+Step 770/50462 | Loss: 4.2507 | LR: 5.00e-04 | 2732.1ms/step | 22.0 samples/s
+Step 780/50462 | Loss: 4.2006 | LR: 5.00e-04 | 2731.7ms/step | 22.0 samples/s
+Step 790/50462 | Loss: 4.1914 | LR: 5.00e-04 | 2734.0ms/step | 21.9 samples/s
+Step 800/50462 | Loss: 4.1470 | LR: 5.00e-04 | 2732.1ms/step | 22.0 samples/s
+Step 810/50462 | Loss: 4.1532 | LR: 5.00e-04 | 2732.8ms/step | 22.0 samples/s
+Step 820/50462 | Loss: 4.1431 | LR: 5.00e-04 | 2731.2ms/step | 22.0 samples/s
+Step 830/50462 | Loss: 4.1267 | LR: 5.00e-04 | 2731.2ms/step | 22.0 samples/s
+Step 840/50462 | Loss: 4.0930 | LR: 5.00e-04 | 2731.3ms/step | 22.0 samples/s
+Step 850/50462 | Loss: 4.0601 | LR: 5.00e-04 | 2730.8ms/step | 22.0 samples/s
+Step 860/50462 | Loss: 4.0583 | LR: 5.00e-04 | 2732.0ms/step | 22.0 samples/s
+Step 870/50462 | Loss: 4.0439 | LR: 5.00e-04 | 2731.0ms/step | 22.0 samples/s
+Step 880/50462 | Loss: 4.0228 | LR: 5.00e-04 | 2729.8ms/step | 22.0 samples/s
+Step 890/50462 | Loss: 4.0190 | LR: 5.00e-04 | 2732.7ms/step | 22.0 samples/s
+Step 900/50462 | Loss: 4.0151 | LR: 5.00e-04 | 2733.5ms/step | 21.9 samples/s
+Step 910/50462 | Loss: 3.9902 | LR: 5.00e-04 | 2731.2ms/step | 22.0 samples/s
+Step 920/50462 | Loss: 3.9476 | LR: 5.00e-04 | 2731.8ms/step | 22.0 samples/s
+Step 930/50462 | Loss: 4.0064 | LR: 5.00e-04 | 2732.4ms/step | 22.0 samples/s
+Step 940/50462 | Loss: 3.9505 | LR: 5.00e-04 | 2731.4ms/step | 22.0 samples/s
+Step 950/50462 | Loss: 3.9353 | LR: 5.00e-04 | 2731.6ms/step | 22.0 samples/s
+Step 960/50462 | Loss: 3.9362 | LR: 5.00e-04 | 2732.2ms/step | 22.0 samples/s
+Step 970/50462 | Loss: 3.9283 | LR: 5.00e-04 | 2732.3ms/step | 22.0 samples/s
+Step 980/50462 | Loss: 3.8905 | LR: 5.00e-04 | 2734.2ms/step | 21.9 samples/s
+Step 990/50462 | Loss: 3.9119 | LR: 5.00e-04 | 2734.7ms/step | 21.9 samples/s
+Step 1000/50462 | Loss: 3.9167 | LR: 5.00e-04 | 2734.8ms/step | 21.9 samples/s
+Validation at step 1000: Loss = 3.9140, Perplexity = 50.10
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/run_20260706_035226/checkpoint-1000
+Step 1010/50462 | Loss: 3.8936 | LR: 5.00e-04 | 5118.7ms/step | 11.7 samples/s
+Step 1020/50462 | Loss: 3.9031 | LR: 5.00e-04 | 2727.6ms/step | 22.0 samples/s
+Step 1030/50462 | Loss: 3.8876 | LR: 5.00e-04 | 2730.8ms/step | 22.0 samples/s
+Step 1040/50462 | Loss: 3.8949 | LR: 5.00e-04 | 2735.6ms/step | 21.9 samples/s
+Step 1050/50462 | Loss: 3.8628 | LR: 5.00e-04 | 2825.2ms/step | 21.2 samples/s
+Step 1060/50462 | Loss: 3.8628 | LR: 5.00e-04 | 2723.6ms/step | 22.0 samples/s
+Step 1070/50462 | Loss: 3.8466 | LR: 5.00e-04 | 2825.4ms/step | 21.2 samples/s
+Step 1080/50462 | Loss: 3.8539 | LR: 5.00e-04 | 2729.7ms/step | 22.0 samples/s
+Step 1090/50462 | Loss: 3.8183 | LR: 5.00e-04 | 2732.7ms/step | 22.0 samples/s
+Step 1100/50462 | Loss: 3.8414 | LR: 5.00e-04 | 2732.3ms/step | 22.0 samples/s
+Step 1110/50462 | Loss: 3.8495 | LR: 5.00e-04 | 2733.0ms/step | 22.0 samples/s
+Step 1120/50462 | Loss: 3.8048 | LR: 5.00e-04 | 2735.2ms/step | 21.9 samples/s
+Step 1130/50462 | Loss: 3.8228 | LR: 5.00e-04 | 2826.8ms/step | 21.2 samples/s
+Step 1140/50462 | Loss: 3.7977 | LR: 5.00e-04 | 2728.3ms/step | 22.0 samples/s
+Step 1150/50462 | Loss: 3.7653 | LR: 5.00e-04 | 2731.6ms/step | 22.0 samples/s
+Step 1160/50462 | Loss: 3.7864 | LR: 5.00e-04 | 2732.5ms/step | 22.0 samples/s
+Step 1170/50462 | Loss: 3.7751 | LR: 5.00e-04 | 2734.0ms/step | 21.9 samples/s
+Step 1180/50462 | Loss: 3.7397 | LR: 5.00e-04 | 2732.2ms/step | 22.0 samples/s
+Step 1190/50462 | Loss: 3.7687 | LR: 5.00e-04 | 2731.8ms/step | 22.0 samples/s
+Step 1200/50462 | Loss: 3.7776 | LR: 5.00e-04 | 2732.4ms/step | 22.0 samples/s
+Step 1210/50462 | Loss: 3.7841 | LR: 5.00e-04 | 2734.3ms/step | 21.9 samples/s
+Step 1220/50462 | Loss: 3.7423 | LR: 5.00e-04 | 2733.6ms/step | 21.9 samples/s
+Step 1230/50462 | Loss: 3.7118 | LR: 5.00e-04 | 2733.1ms/step | 22.0 samples/s
+Step 1240/50462 | Loss: 3.7669 | LR: 5.00e-04 | 2731.8ms/step | 22.0 samples/s
+Step 1250/50462 | Loss: 3.7293 | LR: 5.00e-04 | 2734.4ms/step | 21.9 samples/s
+Validation at step 1250: Loss = 3.7439, Perplexity = 42.26
+Step 1260/50462 | Loss: 3.7533 | LR: 5.00e-04 | 4210.8ms/step | 14.2 samples/s
+Step 1270/50462 | Loss: 3.7187 | LR: 5.00e-04 | 2734.1ms/step | 21.9 samples/s
+Step 1280/50462 | Loss: 3.7036 | LR: 5.00e-04 | 2735.7ms/step | 21.9 samples/s
+Step 1290/50462 | Loss: 3.7054 | LR: 5.00e-04 | 2734.4ms/step | 21.9 samples/s
+Step 1300/50462 | Loss: 3.6988 | LR: 5.00e-04 | 2734.4ms/step | 21.9 samples/s
+Step 1310/50462 | Loss: 3.7031 | LR: 5.00e-04 | 2733.2ms/step | 22.0 samples/s
+Step 1320/50462 | Loss: 3.6987 | LR: 5.00e-04 | 2733.7ms/step | 21.9 samples/s
+Step 1330/50462 | Loss: 3.6963 | LR: 5.00e-04 | 2728.7ms/step | 22.0 samples/s
+Step 1340/50462 | Loss: 3.6819 | LR: 5.00e-04 | 2727.2ms/step | 22.0 samples/s
+Step 1350/50462 | Loss: 3.6646 | LR: 5.00e-04 | 2730.2ms/step | 22.0 samples/s
+Step 1360/50462 | Loss: 3.6852 | LR: 5.00e-04 | 2731.3ms/step | 22.0 samples/s
+Step 1370/50462 | Loss: 3.6627 | LR: 5.00e-04 | 2732.5ms/step | 22.0 samples/s
+Step 1380/50462 | Loss: 3.6991 | LR: 5.00e-04 | 2733.8ms/step | 21.9 samples/s
+Step 1390/50462 | Loss: 3.6855 | LR: 5.00e-04 | 2732.2ms/step | 22.0 samples/s
+Step 1400/50462 | Loss: 3.6452 | LR: 5.00e-04 | 2730.6ms/step | 22.0 samples/s
+Step 1410/50462 | Loss: 3.6715 | LR: 5.00e-04 | 2732.2ms/step | 22.0 samples/s
+Step 1420/50462 | Loss: 3.6436 | LR: 5.00e-04 | 2732.0ms/step | 22.0 samples/s
+Step 1430/50462 | Loss: 3.6269 | LR: 5.00e-04 | 2734.1ms/step | 21.9 samples/s
+Step 1440/50462 | Loss: 3.6598 | LR: 5.00e-04 | 2733.6ms/step | 21.9 samples/s
+Step 1450/50462 | Loss: 3.6290 | LR: 5.00e-04 | 2732.4ms/step | 22.0 samples/s
+Step 1460/50462 | Loss: 3.6302 | LR: 5.00e-04 | 2732.5ms/step | 22.0 samples/s
+Step 1470/50462 | Loss: 3.6116 | LR: 5.00e-04 | 2734.3ms/step | 21.9 samples/s
+Step 1480/50462 | Loss: 3.6510 | LR: 5.00e-04 | 2732.8ms/step | 22.0 samples/s
+Step 1490/50462 | Loss: 3.6315 | LR: 5.00e-04 | 2733.5ms/step | 22.0 samples/s
+Step 1500/50462 | Loss: 3.5970 | LR: 5.00e-04 | 2733.9ms/step | 21.9 samples/s
+Validation at step 1500: Loss = 3.6397, Perplexity = 38.08
+Step 1510/50462 | Loss: 3.6399 | LR: 5.00e-04 | 4209.6ms/step | 14.3 samples/s
+Step 1520/50462 | Loss: 3.6057 | LR: 5.00e-04 | 2736.5ms/step | 21.9 samples/s
+Step 1530/50462 | Loss: 3.6327 | LR: 5.00e-04 | 2735.3ms/step | 21.9 samples/s
+Step 1540/50462 | Loss: 3.5982 | LR: 5.00e-04 | 2737.0ms/step | 21.9 samples/s
+Step 1550/50462 | Loss: 3.6094 | LR: 5.00e-04 | 2736.3ms/step | 21.9 samples/s
+Step 1560/50462 | Loss: 3.5990 | LR: 5.00e-04 | 2829.2ms/step | 21.2 samples/s
+Step 1570/50462 | Loss: 3.6280 | LR: 5.00e-04 | 2730.5ms/step | 22.0 samples/s
+Step 1580/50462 | Loss: 3.5972 | LR: 5.00e-04 | 2830.5ms/step | 21.2 samples/s
+Step 1590/50462 | Loss: 3.6143 | LR: 5.00e-04 | 2733.4ms/step | 22.0 samples/s
+Step 1600/50462 | Loss: 3.6102 | LR: 5.00e-04 | 2734.5ms/step | 21.9 samples/s
+Step 1610/50462 | Loss: 3.5729 | LR: 5.00e-04 | 2735.7ms/step | 21.9 samples/s
+Step 1620/50462 | Loss: 3.5985 | LR: 5.00e-04 | 2735.9ms/step | 21.9 samples/s
+Step 1630/50462 | Loss: 3.6093 | LR: 5.00e-04 | 2735.9ms/step | 21.9 samples/s
+Step 1640/50462 | Loss: 3.5807 | LR: 5.00e-04 | 2826.7ms/step | 21.2 samples/s
+Step 1650/50462 | Loss: 3.5394 | LR: 5.00e-04 | 2731.7ms/step | 22.0 samples/s
+Step 1660/50462 | Loss: 3.5663 | LR: 5.00e-04 | 2734.0ms/step | 21.9 samples/s
+Step 1670/50462 | Loss: 3.5889 | LR: 5.00e-04 | 2735.1ms/step | 21.9 samples/s
+Step 1680/50462 | Loss: 3.5619 | LR: 5.00e-04 | 2736.3ms/step | 21.9 samples/s
+Step 1690/50462 | Loss: 3.5610 | LR: 5.00e-04 | 2735.0ms/step | 21.9 samples/s
+Step 1700/50462 | Loss: 3.5697 | LR: 5.00e-04 | 2735.8ms/step | 21.9 samples/s
+Step 1710/50462 | Loss: 3.5828 | LR: 5.00e-04 | 2733.7ms/step | 21.9 samples/s
+Step 1720/50462 | Loss: 3.5525 | LR: 5.00e-04 | 2733.5ms/step | 21.9 samples/s
+Step 1730/50462 | Loss: 3.5511 | LR: 5.00e-04 | 2735.4ms/step | 21.9 samples/s
+Step 1740/50462 | Loss: 3.5757 | LR: 5.00e-04 | 2733.7ms/step | 21.9 samples/s
+Step 1750/50462 | Loss: 3.5921 | LR: 5.00e-04 | 2734.3ms/step | 21.9 samples/s
+Validation at step 1750: Loss = 3.5618, Perplexity = 35.23
+Step 1760/50462 | Loss: 3.5675 | LR: 4.99e-04 | 4209.5ms/step | 14.3 samples/s
+Step 1770/50462 | Loss: 3.5524 | LR: 4.99e-04 | 2733.5ms/step | 21.9 samples/s
+Step 1780/50462 | Loss: 3.5203 | LR: 4.99e-04 | 2732.2ms/step | 22.0 samples/s
+Step 1790/50462 | Loss: 3.5273 | LR: 4.99e-04 | 2734.6ms/step | 21.9 samples/s
+Step 1800/50462 | Loss: 3.5169 | LR: 4.99e-04 | 2734.1ms/step | 21.9 samples/s
+Step 1810/50462 | Loss: 3.5356 | LR: 4.99e-04 | 2734.7ms/step | 21.9 samples/s
+Step 1820/50462 | Loss: 3.5181 | LR: 4.99e-04 | 2732.2ms/step | 22.0 samples/s
+W0706 05:19:07.385000 1283 torch/distributed/elastic/agent/server/api.py:753] Received Signals.SIGTERM death signal, shutting down workers
+W0706 05:19:07.395000 1283 torch/distributed/elastic/multiprocessing/api.py:1014] Sending process 1296 closing signal SIGTERM
+Traceback (most recent call last):
+  File "/usr/lib/python3.10/runpy.py", line 196, in _run_module_as_main
+    return _run_code(code, main_globals, None,
+  File "/usr/lib/python3.10/runpy.py", line 86, in _run_code
+    exec(code, run_globals)
+  File "/mloscratch/homes/navasard/pylibs/lib/python3.10/site-packages/torch/distributed/run.py", line 1020, in <module>
+    main()
+  File "/mloscratch/homes/navasard/pylibs/lib/python3.10/site-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 367, in wrapper
+    return f(*args, **kwargs)
+  File "/mloscratch/homes/navasard/pylibs/lib/python3.10/site-packages/torch/distributed/run.py", line 1016, in main
+    run(args)
+  File "/mloscratch/homes/navasard/pylibs/lib/python3.10/site-packages/torch/distributed/run.py", line 1007, in run
+    elastic_launch(
+  File "/mloscratch/homes/navasard/pylibs/lib/python3.10/site-packages/torch/distributed/launcher/api.py", line 191, in __call__
+    return launch_agent(self._config, self._entrypoint, list(args))
+  File "/mloscratch/homes/navasard/pylibs/lib/python3.10/site-packages/torch/distributed/launcher/api.py", line 362, in launch_agent
+    result = agent.run()
+  File "/mloscratch/homes/navasard/pylibs/lib/python3.10/site-packages/torch/distributed/elastic/metrics/api.py", line 134, in wrapper
+    result = f(*args, **kwargs)
+  File "/mloscratch/homes/navasard/pylibs/lib/python3.10/site-packages/torch/distributed/elastic/agent/server/api.py", line 745, in run
+    result = self._invoke_run(role)
+  File "/mloscratch/homes/navasard/pylibs/lib/python3.10/site-packages/torch/distributed/elastic/agent/server/api.py", line 923, in _invoke_run
+    time.sleep(monitor_interval)
+  File "/mloscratch/homes/navasard/pylibs/lib/python3.10/site-packages/torch/distributed/elastic/multiprocessing/api.py", line 86, in _terminate_process_handler
+    raise SignalException(f"Process {os.getpid()} got signal: {sigval}", sigval=sigval)
+torch.distributed.elastic.multiprocessing.api.SignalException: Process 1283 got signal: 15

epfl_final/train_logs/resume_fixed30k.log

@@ -0,0 +1,639 @@
+Auto-resuming from latest checkpoint: /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_fixed/run_20260709_000503/checkpoint-1500
+Resuming: Using existing output directory /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_fixed/run_20260709_000503
+Running config:
+{
+  "model_path": "/mloscratch/homes/navasard/qwen3_06b_base",
+  "model_dtype": "bfloat16",
+  "init_from_scratch": true,
+  "train_path": "/mloscratch/homes/navasard/data/longctx30k/tokenized",
+  "val_path": null,
+  "max_seq_length": 30720,
+  "num_workers": 4,
+  "output_dir": "/mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_fixed/run_20260709_000503",
+  "resume_from_checkpoint": "/mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_fixed/run_20260709_000503/checkpoint-1500",
+  "force_new_run": false,
+  "num_epochs": 1,
+  "max_steps": null,
+  "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": 300,
+  "eval_steps": 500,
+  "eval_batches": 20,
+  "logging_steps": 10,
+  "seed": 42,
+  "verbose": true,
+  "checkpoint_activations": true,
+  "strategy": "ddp",
+  "mixed_precision": "bf16",
+  "wandb_enabled": true,
+  "wandb_project": "stochastic_sparsity",
+  "wandb_name": "1.2b fixed 30k longctx",
+  "wandb_tags": [
+    "qwen3",
+    "1.2b",
+    "fixed",
+    "30k",
+    "longctx",
+    "b128",
+    "scaleladder"
+  ],
+  "hidden_size": 1536,
+  "intermediate_size": 6144,
+  "num_attention_heads": 12,
+  "num_key_value_heads": 4,
+  "num_hidden_layers": null,
+  "head_dim": null,
+  "attention_type": "sparse_stochastic",
+  "sparse_block_size": 128,
+  "sparse_num_random_blocks": 120,
+  "sparse_mask_pool_size": 1,
+  "sparse_val_dense_compare": false,
+  "sparse_layer_indices": null,
+  "sparse_mask_cache_dir": "/mloscratch/homes/navasard/sparse/mask_cache_30k_b128_nrb120",
+  "sparse_mask_seed": 1,
+  "dense_mode": false,
+  "dense_step_intervals": null
+}
+Loading tokenizer from /mloscratch/homes/navasard/qwen3_06b_base...
+Checkpoint activations: True
+Checkpoint activations: True
+Checkpoint activations: True
+Resuming from checkpoint: /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_fixed/run_20260709_000503/checkpoint-1500
+Checkpoint activations: True
+
+
+================================================================================
+MODEL ARCHITECTURE
+================================================================================
+Model config:
+  vocab_size: 151936
+  hidden_size: 1536
+  num_hidden_layers: 28
+  num_attention_heads: 12
+  num_key_value_heads: 4
+  head_dim: 128
+  intermediate_size: 6144
+  max_position_embeddings: 32768
+  rope_theta: 1000000
+  attention_type: sparse_stochastic
+  sparse_block_size: 128
+  sparse_num_random_blocks: 120
+  sparse_mask_pool_size: 1
+  sparse_max_seq_length: 30720
+
+Parameter count:
+  Total parameters: 1,202,352,640
+  Trainable parameters: 1,202,352,640
+  Non-trainable parameters: 0
+  Model size (BF16): 2.40 GB
+
+--------------------------------------------------------------------------------
+LAYER-BY-LAYER SUMMARY
+--------------------------------------------------------------------------------
+/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/nn/attention/flex_attention.py:1704: UserWarning: flex_attention called without torch.compile() - this will use an unfused implementation that materializes the full scores matrix instead of generating a fused kernel.
+
+SOLUTION: Use torch.compile(flex_attention)(...)
+
+If you want to debug your score_mod/mask_mod, you can set:
+torch.nn.attention.flex_attention._FLEX_ATTENTION_DISABLE_COMPILE_DEBUG = True
+
+This will allow you to use print statements or breakpoints. Note: This doesn't work with the backwards pass and may produce incorrect results.
+  _warn_once(
+Could not generate detailed summary: Failed to run torchinfo. See above stack traces for more details. Executed layers up to: [Embedding: 2, Qwen3RMSNorm: 4, Linear: 5, Linear: 5, Linear: 5, Qwen3RMSNorm: 5, Qwen3RMSNorm: 5, Qwen3RotaryEmbedding: 5]
+  NOTE: torchinfo double-counts tied embed_tokens/lm_head weight (233,373,696 params). True unique parameter count: 1,202,352,640 (1202.4M)
+================================================================================
+
+Loading dataset from /mloscratch/homes/navasard/data/longctx30k/tokenized...
+Using pre-tokenized dataset (no runtime tokenization)
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/longctx30k/tokenized...
+Loaded 788736 pre-tokenized examples (stored_len=30720, pack_factor=1, packed_examples=788736)
+Dataset split: 749299 train, 39437 validation
+Training examples (packed, train split):  749299
+Validation examples (packed, val split):  39437
+Per-device batch size:                    1
+World size (GPUs):                        4
+Gradient accumulation steps:              5
+Effective batch size (seqs/step):         20
+Sequence length:                          30720
+Tokens per optimizer step:                614,400
+Micro-batches per epoch (per rank):       187325
+Optimizer steps per epoch:                37465
+Total optimizer steps:                    37465
+Loading checkpoint from /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_fixed/run_20260709_000503/checkpoint-1500...
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/longctx30k/tokenized...
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/longctx30k/tokenized...
+Loaded 788736 pre-tokenized examples (stored_len=30720, pack_factor=1, packed_examples=788736)
+Loaded 788736 pre-tokenized examples (stored_len=30720, pack_factor=1, packed_examples=788736)
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/longctx30k/tokenized...
+Loaded 788736 pre-tokenized examples (stored_len=30720, pack_factor=1, packed_examples=788736)
+Overriding checkpoint LR with config LR: 0.0005
+Resumed from checkpoint at step 1500 (epoch 0, micro_step 7500)
+Resuming training from step 1500
+Resuming wandb run 9blizq3m
+wandb: [wandb.login()] Loaded credentials for https://api.wandb.ai from WANDB_API_KEY.
+wandb: Currently logged in as: ruben-navasardyan4869 (ruben-navasardyan4869-epfl) to https://api.wandb.ai. Use `wandb login --relogin` to force relogin
+wandb: setting up run 9blizq3m
+wandb: Tracking run with wandb version 0.26.1
+wandb: Run data is saved locally in /mloscratch/homes/navasard/sparse/SparseAtt/wandb/run-20260709_223121-9blizq3m
+wandb: Run `wandb offline` to turn off syncing.
+wandb: Resuming run 1.2b fixed 30k longctx
+wandb: ⭐️ View project at https://wandb.ai/ruben-navasardyan4869-epfl/stochastic_sparsity
+wandb: 🚀 View run at https://wandb.ai/ruben-navasardyan4869-epfl/stochastic_sparsity/runs/9blizq3m
+Starting training...
+Validation at step 1500: Loss = 2.9803, Perplexity = 19.69
+wandb: WARNING Tried to log to step 1500 that is less than the current step 1561. Steps must be monotonically increasing, so this data will be ignored. See https://wandb.me/define-metric to log data out of order.
+Switching attention mode for step 1501: sparse_stochastic (configured sparse attention)
+Step 1510/37465 | Loss: 2.9913 | LR: 4.99e-04 | 44079.1ms/step | 0.5 samples/s
+wandb: WARNING Tried to log to step 1510 that is less than the current step 1561. Steps must be monotonically increasing, so this data will be ignored. See https://wandb.me/define-metric to log data out of order.
+Step 1520/37465 | Loss: 2.9083 | LR: 4.99e-04 | 28202.3ms/step | 0.7 samples/s
+wandb: WARNING Tried to log to step 1520 that is less than the current step 1561. Steps must be monotonically increasing, so this data will be ignored. See https://wandb.me/define-metric to log data out of order.
+Step 1530/37465 | Loss: 2.8394 | LR: 4.99e-04 | 28206.1ms/step | 0.7 samples/s
+wandb: WARNING Tried to log to step 1530 that is less than the current step 1561. Steps must be monotonically increasing, so this data will be ignored. See https://wandb.me/define-metric to log data out of order.
+Step 1540/37465 | Loss: 2.8274 | LR: 4.99e-04 | 28203.1ms/step | 0.7 samples/s
+wandb: WARNING Tried to log to step 1540 that is less than the current step 1561. Steps must be monotonically increasing, so this data will be ignored. See https://wandb.me/define-metric to log data out of order.
+Step 1550/37465 | Loss: 3.1272 | LR: 4.99e-04 | 28200.0ms/step | 0.7 samples/s
+wandb: WARNING Tried to log to step 1550 that is less than the current step 1561. Steps must be monotonically increasing, so this data will be ignored. See https://wandb.me/define-metric to log data out of order.
+Step 1560/37465 | Loss: 3.0044 | LR: 4.99e-04 | 28203.0ms/step | 0.7 samples/s
+wandb: WARNING Tried to log to step 1560 that is less than the current step 1561. Steps must be monotonically increasing, so this data will be ignored. See https://wandb.me/define-metric to log data out of order.
+Step 1570/37465 | Loss: 2.8018 | LR: 4.99e-04 | 28204.8ms/step | 0.7 samples/s
+Step 1580/37465 | Loss: 3.1594 | LR: 4.99e-04 | 28210.3ms/step | 0.7 samples/s
+Step 1590/37465 | Loss: 3.0262 | LR: 4.99e-04 | 28204.2ms/step | 0.7 samples/s
+Step 1600/37465 | Loss: 3.3198 | LR: 4.99e-04 | 28203.9ms/step | 0.7 samples/s
+Step 1610/37465 | Loss: 2.8294 | LR: 4.99e-04 | 28201.5ms/step | 0.7 samples/s
+Step 1620/37465 | Loss: 3.0865 | LR: 4.99e-04 | 28200.0ms/step | 0.7 samples/s
+Step 1630/37465 | Loss: 2.9460 | LR: 4.99e-04 | 28200.5ms/step | 0.7 samples/s
+Step 1640/37465 | Loss: 2.8964 | LR: 4.99e-04 | 28201.6ms/step | 0.7 samples/s
+Step 1650/37465 | Loss: 2.6442 | LR: 4.99e-04 | 28197.7ms/step | 0.7 samples/s
+Step 1660/37465 | Loss: 2.7789 | LR: 4.99e-04 | 28196.9ms/step | 0.7 samples/s
+Step 1670/37465 | Loss: 2.7841 | LR: 4.99e-04 | 28205.4ms/step | 0.7 samples/s
+Step 1680/37465 | Loss: 3.0626 | LR: 4.99e-04 | 28198.3ms/step | 0.7 samples/s
+Step 1690/37465 | Loss: 2.9370 | LR: 4.99e-04 | 28204.0ms/step | 0.7 samples/s
+Step 1700/37465 | Loss: 2.9952 | LR: 4.99e-04 | 28200.6ms/step | 0.7 samples/s
+Step 1710/37465 | Loss: 3.0111 | LR: 4.99e-04 | 28205.5ms/step | 0.7 samples/s
+Step 1720/37465 | Loss: 3.0586 | LR: 4.99e-04 | 28208.0ms/step | 0.7 samples/s
+Step 1730/37465 | Loss: 2.7984 | LR: 4.99e-04 | 28202.8ms/step | 0.7 samples/s
+Step 1740/37465 | Loss: 2.9516 | LR: 4.99e-04 | 28207.0ms/step | 0.7 samples/s
+Step 1750/37465 | Loss: 2.9898 | LR: 4.99e-04 | 28206.3ms/step | 0.7 samples/s
+Step 1760/37465 | Loss: 2.8837 | LR: 4.99e-04 | 28204.5ms/step | 0.7 samples/s
+Step 1770/37465 | Loss: 2.8924 | LR: 4.99e-04 | 28209.3ms/step | 0.7 samples/s
+Step 1780/37465 | Loss: 2.7415 | LR: 4.99e-04 | 28207.8ms/step | 0.7 samples/s
+Step 1790/37465 | Loss: 2.9827 | LR: 4.99e-04 | 28204.6ms/step | 0.7 samples/s
+Step 1800/37465 | Loss: 2.6493 | LR: 4.99e-04 | 28208.3ms/step | 0.7 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_fixed/run_20260709_000503/checkpoint-1800
+Step 1810/37465 | Loss: 2.9498 | LR: 4.99e-04 | 29218.1ms/step | 0.7 samples/s
+Step 1820/37465 | Loss: 2.9610 | LR: 4.99e-04 | 28203.1ms/step | 0.7 samples/s
+Step 1830/37465 | Loss: 2.8554 | LR: 4.99e-04 | 28200.7ms/step | 0.7 samples/s
+Step 1840/37465 | Loss: 2.9721 | LR: 4.99e-04 | 28198.2ms/step | 0.7 samples/s
+Step 1850/37465 | Loss: 2.7502 | LR: 4.99e-04 | 28196.5ms/step | 0.7 samples/s
+Step 1860/37465 | Loss: 2.7822 | LR: 4.99e-04 | 28197.3ms/step | 0.7 samples/s
+Step 1870/37465 | Loss: 3.1765 | LR: 4.99e-04 | 28197.7ms/step | 0.7 samples/s
+Step 1880/37465 | Loss: 2.9186 | LR: 4.99e-04 | 28201.0ms/step | 0.7 samples/s
+Step 1890/37465 | Loss: 2.9122 | LR: 4.99e-04 | 28201.0ms/step | 0.7 samples/s
+Step 1900/37465 | Loss: 2.7845 | LR: 4.99e-04 | 28197.3ms/step | 0.7 samples/s
+Step 1910/37465 | Loss: 2.7688 | LR: 4.99e-04 | 28201.1ms/step | 0.7 samples/s
+Step 1920/37465 | Loss: 2.9802 | LR: 4.99e-04 | 28200.1ms/step | 0.7 samples/s
+Step 1930/37465 | Loss: 2.8521 | LR: 4.99e-04 | 28200.7ms/step | 0.7 samples/s
+Step 1940/37465 | Loss: 2.8682 | LR: 4.99e-04 | 28201.3ms/step | 0.7 samples/s
+Step 1950/37465 | Loss: 2.7948 | LR: 4.99e-04 | 28202.5ms/step | 0.7 samples/s
+Step 1960/37465 | Loss: 2.7922 | LR: 4.99e-04 | 28205.8ms/step | 0.7 samples/s
+Step 1970/37465 | Loss: 2.9998 | LR: 4.99e-04 | 28205.9ms/step | 0.7 samples/s
+Step 1980/37465 | Loss: 2.7698 | LR: 4.99e-04 | 28196.7ms/step | 0.7 samples/s
+Step 1990/37465 | Loss: 2.7311 | LR: 4.99e-04 | 28201.8ms/step | 0.7 samples/s
+Step 2000/37465 | Loss: 2.8256 | LR: 4.99e-04 | 28198.0ms/step | 0.7 samples/s
+Validation at step 2000: Loss = 2.8067, Perplexity = 16.55
+Step 2010/37465 | Loss: 2.7373 | LR: 4.99e-04 | 29879.9ms/step | 0.7 samples/s
+Step 2020/37465 | Loss: 2.8007 | LR: 4.99e-04 | 28206.4ms/step | 0.7 samples/s
+Step 2030/37465 | Loss: 2.6649 | LR: 4.99e-04 | 28204.0ms/step | 0.7 samples/s
+Step 2040/37465 | Loss: 2.7121 | LR: 4.98e-04 | 28194.5ms/step | 0.7 samples/s
+Step 2050/37465 | Loss: 2.8025 | LR: 4.98e-04 | 28201.6ms/step | 0.7 samples/s
+Step 2060/37465 | Loss: 2.6587 | LR: 4.98e-04 | 28201.5ms/step | 0.7 samples/s
+Step 2070/37465 | Loss: 2.6982 | LR: 4.98e-04 | 28202.5ms/step | 0.7 samples/s
+Step 2080/37465 | Loss: 2.8682 | LR: 4.98e-04 | 28199.6ms/step | 0.7 samples/s
+Step 2090/37465 | Loss: 2.8524 | LR: 4.98e-04 | 28200.3ms/step | 0.7 samples/s
+Step 2100/37465 | Loss: 2.9557 | LR: 4.98e-04 | 28203.8ms/step | 0.7 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_fixed/run_20260709_000503/checkpoint-2100
+Step 2110/37465 | Loss: 2.8975 | LR: 4.98e-04 | 29229.2ms/step | 0.7 samples/s
+Step 2120/37465 | Loss: 2.8102 | LR: 4.98e-04 | 28211.3ms/step | 0.7 samples/s
+Step 2130/37465 | Loss: 2.5449 | LR: 4.98e-04 | 28213.3ms/step | 0.7 samples/s
+Step 2140/37465 | Loss: 2.8684 | LR: 4.98e-04 | 28213.9ms/step | 0.7 samples/s
+Step 2150/37465 | Loss: 2.8026 | LR: 4.98e-04 | 28211.1ms/step | 0.7 samples/s
+Step 2160/37465 | Loss: 2.7089 | LR: 4.98e-04 | 28211.8ms/step | 0.7 samples/s
+Step 2170/37465 | Loss: 2.8118 | LR: 4.98e-04 | 28211.3ms/step | 0.7 samples/s
+Step 2180/37465 | Loss: 2.8394 | LR: 4.98e-04 | 28214.3ms/step | 0.7 samples/s
+Step 2190/37465 | Loss: 2.7915 | LR: 4.98e-04 | 28214.8ms/step | 0.7 samples/s
+Step 2200/37465 | Loss: 2.7080 | LR: 4.98e-04 | 28211.5ms/step | 0.7 samples/s
+Step 2210/37465 | Loss: 2.6452 | LR: 4.98e-04 | 28214.0ms/step | 0.7 samples/s
+Step 2220/37465 | Loss: 2.8388 | LR: 4.98e-04 | 28213.8ms/step | 0.7 samples/s
+Step 2230/37465 | Loss: 2.8298 | LR: 4.98e-04 | 28214.9ms/step | 0.7 samples/s
+Step 2240/37465 | Loss: 2.8257 | LR: 4.98e-04 | 28214.3ms/step | 0.7 samples/s
+Step 2250/37465 | Loss: 2.7224 | LR: 4.98e-04 | 28218.1ms/step | 0.7 samples/s
+Step 2260/37465 | Loss: 2.8717 | LR: 4.98e-04 | 28211.6ms/step | 0.7 samples/s
+Step 2270/37465 | Loss: 2.7679 | LR: 4.98e-04 | 28208.5ms/step | 0.7 samples/s
+Step 2280/37465 | Loss: 2.6003 | LR: 4.98e-04 | 28206.0ms/step | 0.7 samples/s
+Step 2290/37465 | Loss: 2.7364 | LR: 4.98e-04 | 28204.7ms/step | 0.7 samples/s
+Step 2300/37465 | Loss: 2.8338 | LR: 4.98e-04 | 28204.1ms/step | 0.7 samples/s
+Step 2310/37465 | Loss: 2.8896 | LR: 4.98e-04 | 28204.5ms/step | 0.7 samples/s
+Step 2320/37465 | Loss: 2.7513 | LR: 4.98e-04 | 28209.5ms/step | 0.7 samples/s
+Step 2330/37465 | Loss: 3.0114 | LR: 4.98e-04 | 28213.7ms/step | 0.7 samples/s
+Step 2340/37465 | Loss: 2.5985 | LR: 4.98e-04 | 28207.2ms/step | 0.7 samples/s
+Step 2350/37465 | Loss: 2.9940 | LR: 4.98e-04 | 28206.7ms/step | 0.7 samples/s
+Step 2360/37465 | Loss: 2.6100 | LR: 4.98e-04 | 28206.4ms/step | 0.7 samples/s
+Step 2370/37465 | Loss: 2.8667 | LR: 4.98e-04 | 28213.1ms/step | 0.7 samples/s
+Step 2380/37465 | Loss: 2.6913 | LR: 4.98e-04 | 28207.0ms/step | 0.7 samples/s
+Step 2390/37465 | Loss: 2.5641 | LR: 4.98e-04 | 28207.6ms/step | 0.7 samples/s
+Step 2400/37465 | Loss: 2.5432 | LR: 4.98e-04 | 28208.6ms/step | 0.7 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_fixed/run_20260709_000503/checkpoint-2400
+Step 2410/37465 | Loss: 2.7141 | LR: 4.97e-04 | 29248.7ms/step | 0.7 samples/s
+Step 2420/37465 | Loss: 2.9042 | LR: 4.97e-04 | 28216.2ms/step | 0.7 samples/s
+Step 2430/37465 | Loss: 2.8914 | LR: 4.97e-04 | 28208.3ms/step | 0.7 samples/s
+Step 2440/37465 | Loss: 2.8282 | LR: 4.97e-04 | 28214.5ms/step | 0.7 samples/s
+Step 2450/37465 | Loss: 2.8363 | LR: 4.97e-04 | 28215.8ms/step | 0.7 samples/s
+Step 2460/37465 | Loss: 2.8236 | LR: 4.97e-04 | 28209.3ms/step | 0.7 samples/s
+Step 2470/37465 | Loss: 2.8401 | LR: 4.97e-04 | 28213.0ms/step | 0.7 samples/s
+Step 2480/37465 | Loss: 2.6442 | LR: 4.97e-04 | 28214.2ms/step | 0.7 samples/s
+Step 2490/37465 | Loss: 2.6397 | LR: 4.97e-04 | 28213.9ms/step | 0.7 samples/s
+Step 2500/37465 | Loss: 2.7731 | LR: 4.97e-04 | 28215.4ms/step | 0.7 samples/s
+Validation at step 2500: Loss = 2.7050, Perplexity = 14.95
+Step 2510/37465 | Loss: 2.8276 | LR: 4.97e-04 | 29861.6ms/step | 0.7 samples/s
+Step 2520/37465 | Loss: 2.5297 | LR: 4.97e-04 | 28216.5ms/step | 0.7 samples/s
+Step 2530/37465 | Loss: 2.6744 | LR: 4.97e-04 | 28215.5ms/step | 0.7 samples/s
+Step 2540/37465 | Loss: 2.9592 | LR: 4.97e-04 | 28219.8ms/step | 0.7 samples/s
+Step 2550/37465 | Loss: 2.4623 | LR: 4.97e-04 | 28217.6ms/step | 0.7 samples/s
+Step 2560/37465 | Loss: 2.7488 | LR: 4.97e-04 | 28211.1ms/step | 0.7 samples/s
+Step 2570/37465 | Loss: 2.7104 | LR: 4.97e-04 | 28214.9ms/step | 0.7 samples/s
+Step 2580/37465 | Loss: 2.4993 | LR: 4.97e-04 | 28209.6ms/step | 0.7 samples/s
+Step 2590/37465 | Loss: 2.7331 | LR: 4.97e-04 | 28213.1ms/step | 0.7 samples/s
+Step 2600/37465 | Loss: 2.8604 | LR: 4.97e-04 | 28211.5ms/step | 0.7 samples/s
+Step 2610/37465 | Loss: 2.5769 | LR: 4.97e-04 | 28213.4ms/step | 0.7 samples/s
+Step 2620/37465 | Loss: 2.8670 | LR: 4.97e-04 | 28213.5ms/step | 0.7 samples/s
+Step 2630/37465 | Loss: 2.6857 | LR: 4.97e-04 | 28212.1ms/step | 0.7 samples/s
+Step 2640/37465 | Loss: 2.6256 | LR: 4.97e-04 | 28211.1ms/step | 0.7 samples/s
+Step 2650/37465 | Loss: 2.6200 | LR: 4.97e-04 | 28209.7ms/step | 0.7 samples/s
+Step 2660/37465 | Loss: 2.5157 | LR: 4.97e-04 | 28213.0ms/step | 0.7 samples/s
+Step 2670/37465 | Loss: 2.5916 | LR: 4.97e-04 | 28210.6ms/step | 0.7 samples/s
+Step 2680/37465 | Loss: 2.8806 | LR: 4.97e-04 | 28213.7ms/step | 0.7 samples/s
+Step 2690/37465 | Loss: 2.6171 | LR: 4.97e-04 | 28211.8ms/step | 0.7 samples/s
+Step 2700/37465 | Loss: 2.7728 | LR: 4.97e-04 | 28206.6ms/step | 0.7 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_fixed/run_20260709_000503/checkpoint-2700
+Step 2710/37465 | Loss: 2.6148 | LR: 4.96e-04 | 29202.7ms/step | 0.7 samples/s
+Step 2720/37465 | Loss: 2.5995 | LR: 4.96e-04 | 28203.3ms/step | 0.7 samples/s
+Step 2730/37465 | Loss: 2.6339 | LR: 4.96e-04 | 28203.9ms/step | 0.7 samples/s
+Step 2740/37465 | Loss: 2.7182 | LR: 4.96e-04 | 28208.1ms/step | 0.7 samples/s
+Step 2750/37465 | Loss: 2.6234 | LR: 4.96e-04 | 28208.3ms/step | 0.7 samples/s
+Step 2760/37465 | Loss: 2.8483 | LR: 4.96e-04 | 28215.8ms/step | 0.7 samples/s
+Step 2770/37465 | Loss: 2.6882 | LR: 4.96e-04 | 28211.4ms/step | 0.7 samples/s
+Step 2780/37465 | Loss: 2.7589 | LR: 4.96e-04 | 28213.2ms/step | 0.7 samples/s
+Step 2790/37465 | Loss: 2.5082 | LR: 4.96e-04 | 28212.1ms/step | 0.7 samples/s
+Step 2800/37465 | Loss: 2.5738 | LR: 4.96e-04 | 28211.3ms/step | 0.7 samples/s
+Step 2810/37465 | Loss: 2.5457 | LR: 4.96e-04 | 28209.1ms/step | 0.7 samples/s
+Step 2820/37465 | Loss: 2.3763 | LR: 4.96e-04 | 28212.5ms/step | 0.7 samples/s
+Step 2830/37465 | Loss: 2.4791 | LR: 4.96e-04 | 28216.5ms/step | 0.7 samples/s
+Step 2840/37465 | Loss: 2.9267 | LR: 4.96e-04 | 28215.7ms/step | 0.7 samples/s
+Step 2850/37465 | Loss: 2.5846 | LR: 4.96e-04 | 28218.1ms/step | 0.7 samples/s
+Step 2860/37465 | Loss: 2.5474 | LR: 4.96e-04 | 28217.2ms/step | 0.7 samples/s
+Step 2870/37465 | Loss: 2.7486 | LR: 4.96e-04 | 28209.9ms/step | 0.7 samples/s
+Step 2880/37465 | Loss: 2.7018 | LR: 4.96e-04 | 28211.7ms/step | 0.7 samples/s
+Step 2890/37465 | Loss: 2.6266 | LR: 4.96e-04 | 28209.7ms/step | 0.7 samples/s
+Step 2900/37465 | Loss: 2.5835 | LR: 4.96e-04 | 28209.8ms/step | 0.7 samples/s
+Step 2910/37465 | Loss: 2.6910 | LR: 4.96e-04 | 28218.0ms/step | 0.7 samples/s
+Step 2920/37465 | Loss: 2.4776 | LR: 4.96e-04 | 28212.8ms/step | 0.7 samples/s
+Step 2930/37465 | Loss: 2.6225 | LR: 4.96e-04 | 28213.3ms/step | 0.7 samples/s
+Step 2940/37465 | Loss: 2.5352 | LR: 4.96e-04 | 28208.8ms/step | 0.7 samples/s
+Step 2950/37465 | Loss: 2.5628 | LR: 4.96e-04 | 28213.5ms/step | 0.7 samples/s
+Step 2960/37465 | Loss: 2.5977 | LR: 4.96e-04 | 28214.4ms/step | 0.7 samples/s
+Step 2970/37465 | Loss: 2.7209 | LR: 4.96e-04 | 28220.4ms/step | 0.7 samples/s
+Step 2980/37465 | Loss: 2.3571 | LR: 4.95e-04 | 28216.3ms/step | 0.7 samples/s
+Step 2990/37465 | Loss: 2.7625 | LR: 4.95e-04 | 28214.3ms/step | 0.7 samples/s
+Step 3000/37465 | Loss: 2.6710 | LR: 4.95e-04 | 28218.2ms/step | 0.7 samples/s
+Validation at step 3000: Loss = 2.6398, Perplexity = 14.01
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_fixed/run_20260709_000503/checkpoint-3000
+Step 3010/37465 | Loss: 2.6393 | LR: 4.95e-04 | 30857.8ms/step | 0.6 samples/s
+Step 3020/37465 | Loss: 2.6273 | LR: 4.95e-04 | 28282.4ms/step | 0.7 samples/s
+Step 3030/37465 | Loss: 2.8217 | LR: 4.95e-04 | 28229.8ms/step | 0.7 samples/s
+Step 3040/37465 | Loss: 2.7478 | LR: 4.95e-04 | 28230.7ms/step | 0.7 samples/s
+Step 3050/37465 | Loss: 2.6453 | LR: 4.95e-04 | 28228.6ms/step | 0.7 samples/s
+Step 3060/37465 | Loss: 2.6271 | LR: 4.95e-04 | 28228.1ms/step | 0.7 samples/s
+Step 3070/37465 | Loss: 2.4544 | LR: 4.95e-04 | 28228.1ms/step | 0.7 samples/s
+Step 3080/37465 | Loss: 2.4177 | LR: 4.95e-04 | 28231.8ms/step | 0.7 samples/s
+Step 3090/37465 | Loss: 2.5919 | LR: 4.95e-04 | 28219.5ms/step | 0.7 samples/s
+Step 3100/37465 | Loss: 2.6931 | LR: 4.95e-04 | 28222.5ms/step | 0.7 samples/s
+Step 3110/37465 | Loss: 2.6318 | LR: 4.95e-04 | 28227.8ms/step | 0.7 samples/s
+Step 3120/37465 | Loss: 2.5201 | LR: 4.95e-04 | 28224.0ms/step | 0.7 samples/s
+Step 3130/37465 | Loss: 2.6239 | LR: 4.95e-04 | 28225.4ms/step | 0.7 samples/s
+Step 3140/37465 | Loss: 2.4329 | LR: 4.95e-04 | 28225.0ms/step | 0.7 samples/s
+Step 3150/37465 | Loss: 2.3984 | LR: 4.95e-04 | 28234.1ms/step | 0.7 samples/s
+Step 3160/37465 | Loss: 2.5789 | LR: 4.95e-04 | 28229.6ms/step | 0.7 samples/s
+Step 3170/37465 | Loss: 2.8802 | LR: 4.95e-04 | 28236.7ms/step | 0.7 samples/s
+Step 3180/37465 | Loss: 2.8211 | LR: 4.95e-04 | 28231.0ms/step | 0.7 samples/s
+Step 3190/37465 | Loss: 2.3020 | LR: 4.95e-04 | 28231.1ms/step | 0.7 samples/s
+Step 3200/37465 | Loss: 2.4850 | LR: 4.95e-04 | 28227.5ms/step | 0.7 samples/s
+Step 3210/37465 | Loss: 2.4132 | LR: 4.94e-04 | 28220.2ms/step | 0.7 samples/s
+Step 3220/37465 | Loss: 2.6648 | LR: 4.94e-04 | 28215.0ms/step | 0.7 samples/s
+Step 3230/37465 | Loss: 2.6317 | LR: 4.94e-04 | 28218.2ms/step | 0.7 samples/s
+Step 3240/37465 | Loss: 2.5026 | LR: 4.94e-04 | 28214.9ms/step | 0.7 samples/s
+Step 3250/37465 | Loss: 2.5259 | LR: 4.94e-04 | 28215.7ms/step | 0.7 samples/s
+Step 3260/37465 | Loss: 2.6992 | LR: 4.94e-04 | 28213.7ms/step | 0.7 samples/s
+Step 3270/37465 | Loss: 2.6554 | LR: 4.94e-04 | 28211.2ms/step | 0.7 samples/s
+Step 3280/37465 | Loss: 2.6233 | LR: 4.94e-04 | 28215.3ms/step | 0.7 samples/s
+Step 3290/37465 | Loss: 2.5224 | LR: 4.94e-04 | 28217.7ms/step | 0.7 samples/s
+Step 3300/37465 | Loss: 2.4942 | LR: 4.94e-04 | 28214.1ms/step | 0.7 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_fixed/run_20260709_000503/checkpoint-3300
+Step 3310/37465 | Loss: 2.6969 | LR: 4.94e-04 | 29243.5ms/step | 0.7 samples/s
+Step 3320/37465 | Loss: 2.3859 | LR: 4.94e-04 | 28210.7ms/step | 0.7 samples/s
+Step 3330/37465 | Loss: 2.7094 | LR: 4.94e-04 | 28210.3ms/step | 0.7 samples/s
+Step 3340/37465 | Loss: 2.6680 | LR: 4.94e-04 | 28208.5ms/step | 0.7 samples/s
+Step 3350/37465 | Loss: 2.9503 | LR: 4.94e-04 | 28207.7ms/step | 0.7 samples/s
+Step 3360/37465 | Loss: 2.5454 | LR: 4.94e-04 | 28201.8ms/step | 0.7 samples/s
+Step 3370/37465 | Loss: 2.5691 | LR: 4.94e-04 | 28216.0ms/step | 0.7 samples/s
+Step 3380/37465 | Loss: 2.6116 | LR: 4.94e-04 | 28205.9ms/step | 0.7 samples/s
+Step 3390/37465 | Loss: 2.4950 | LR: 4.94e-04 | 28216.6ms/step | 0.7 samples/s
+Step 3400/37465 | Loss: 2.4059 | LR: 4.94e-04 | 28218.2ms/step | 0.7 samples/s
+Step 3410/37465 | Loss: 2.4635 | LR: 4.94e-04 | 28216.2ms/step | 0.7 samples/s
+Step 3420/37465 | Loss: 2.6051 | LR: 4.94e-04 | 28213.7ms/step | 0.7 samples/s
+Step 3430/37465 | Loss: 2.7656 | LR: 4.93e-04 | 28208.1ms/step | 0.7 samples/s
+Step 3440/37465 | Loss: 2.6520 | LR: 4.93e-04 | 28216.7ms/step | 0.7 samples/s
+Step 3450/37465 | Loss: 2.6154 | LR: 4.93e-04 | 28207.2ms/step | 0.7 samples/s
+Step 3460/37465 | Loss: 2.5061 | LR: 4.93e-04 | 28211.5ms/step | 0.7 samples/s
+Step 3470/37465 | Loss: 2.3746 | LR: 4.93e-04 | 28206.2ms/step | 0.7 samples/s
+Step 3480/37465 | Loss: 2.5236 | LR: 4.93e-04 | 28209.4ms/step | 0.7 samples/s
+Step 3490/37465 | Loss: 2.6571 | LR: 4.93e-04 | 28217.0ms/step | 0.7 samples/s
+Step 3500/37465 | Loss: 2.6145 | LR: 4.93e-04 | 28219.4ms/step | 0.7 samples/s
+Validation at step 3500: Loss = 2.5797, Perplexity = 13.19
+Step 3510/37465 | Loss: 2.8275 | LR: 4.93e-04 | 29891.8ms/step | 0.7 samples/s
+Step 3520/37465 | Loss: 2.4929 | LR: 4.93e-04 | 28215.1ms/step | 0.7 samples/s
+Step 3530/37465 | Loss: 2.5376 | LR: 4.93e-04 | 28211.3ms/step | 0.7 samples/s
+Step 3540/37465 | Loss: 2.6971 | LR: 4.93e-04 | 28209.5ms/step | 0.7 samples/s
+Step 3550/37465 | Loss: 2.5596 | LR: 4.93e-04 | 28214.3ms/step | 0.7 samples/s
+Step 3560/37465 | Loss: 2.4734 | LR: 4.93e-04 | 28221.7ms/step | 0.7 samples/s
+Step 3570/37465 | Loss: 2.4779 | LR: 4.93e-04 | 28217.3ms/step | 0.7 samples/s
+Step 3580/37465 | Loss: 2.8165 | LR: 4.93e-04 | 28216.4ms/step | 0.7 samples/s
+Step 3590/37465 | Loss: 2.5325 | LR: 4.93e-04 | 28216.4ms/step | 0.7 samples/s
+Step 3600/37465 | Loss: 2.6648 | LR: 4.93e-04 | 28219.9ms/step | 0.7 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_fixed/run_20260709_000503/checkpoint-3600
+Step 3610/37465 | Loss: 2.5532 | LR: 4.93e-04 | 29208.8ms/step | 0.7 samples/s
+Step 3620/37465 | Loss: 2.4811 | LR: 4.92e-04 | 28218.2ms/step | 0.7 samples/s
+Step 3630/37465 | Loss: 2.4093 | LR: 4.92e-04 | 28223.9ms/step | 0.7 samples/s
+Step 3640/37465 | Loss: 2.5256 | LR: 4.92e-04 | 28217.9ms/step | 0.7 samples/s
+Step 3650/37465 | Loss: 2.6203 | LR: 4.92e-04 | 28214.0ms/step | 0.7 samples/s
+Step 3660/37465 | Loss: 2.5007 | LR: 4.92e-04 | 28217.5ms/step | 0.7 samples/s
+Step 3670/37465 | Loss: 2.4079 | LR: 4.92e-04 | 28216.7ms/step | 0.7 samples/s
+Step 3680/37465 | Loss: 2.7110 | LR: 4.92e-04 | 28223.1ms/step | 0.7 samples/s
+Step 3690/37465 | Loss: 2.8531 | LR: 4.92e-04 | 28222.6ms/step | 0.7 samples/s
+Step 3700/37465 | Loss: 2.5256 | LR: 4.92e-04 | 28222.4ms/step | 0.7 samples/s
+Step 3710/37465 | Loss: 2.8063 | LR: 4.92e-04 | 28222.8ms/step | 0.7 samples/s
+Step 3720/37465 | Loss: 2.5781 | LR: 4.92e-04 | 28226.5ms/step | 0.7 samples/s
+Step 3730/37465 | Loss: 2.7449 | LR: 4.92e-04 | 28224.5ms/step | 0.7 samples/s
+Step 3740/37465 | Loss: 2.3901 | LR: 4.92e-04 | 28221.5ms/step | 0.7 samples/s
+Step 3750/37465 | Loss: 2.5318 | LR: 4.92e-04 | 28211.3ms/step | 0.7 samples/s
+Step 3760/37465 | Loss: 2.7665 | LR: 4.92e-04 | 28227.2ms/step | 0.7 samples/s
+Step 3770/37465 | Loss: 2.7299 | LR: 4.92e-04 | 28221.1ms/step | 0.7 samples/s
+Step 3780/37465 | Loss: 2.6101 | LR: 4.92e-04 | 28226.2ms/step | 0.7 samples/s
+Step 3790/37465 | Loss: 2.4713 | LR: 4.92e-04 | 28220.9ms/step | 0.7 samples/s
+Step 3800/37465 | Loss: 2.5710 | LR: 4.92e-04 | 28222.0ms/step | 0.7 samples/s
+Step 3810/37465 | Loss: 2.6657 | LR: 4.91e-04 | 28221.7ms/step | 0.7 samples/s
+Step 3820/37465 | Loss: 2.6067 | LR: 4.91e-04 | 28227.9ms/step | 0.7 samples/s
+Step 3830/37465 | Loss: 2.7359 | LR: 4.91e-04 | 28227.0ms/step | 0.7 samples/s
+Step 3840/37465 | Loss: 2.3054 | LR: 4.91e-04 | 28225.1ms/step | 0.7 samples/s
+Step 3850/37465 | Loss: 2.6879 | LR: 4.91e-04 | 28221.7ms/step | 0.7 samples/s
+Step 3860/37465 | Loss: 2.4230 | LR: 4.91e-04 | 28227.6ms/step | 0.7 samples/s
+Step 3870/37465 | Loss: 2.7475 | LR: 4.91e-04 | 28225.0ms/step | 0.7 samples/s
+Step 3880/37465 | Loss: 2.3435 | LR: 4.91e-04 | 28217.9ms/step | 0.7 samples/s
+Step 3890/37465 | Loss: 2.3320 | LR: 4.91e-04 | 28218.8ms/step | 0.7 samples/s
+Step 3900/37465 | Loss: 2.4114 | LR: 4.91e-04 | 28223.4ms/step | 0.7 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_fixed/run_20260709_000503/checkpoint-3900
+Step 3910/37465 | Loss: 2.6000 | LR: 4.91e-04 | 29216.0ms/step | 0.7 samples/s
+Step 3920/37465 | Loss: 2.5912 | LR: 4.91e-04 | 28228.0ms/step | 0.7 samples/s
+Step 3930/37465 | Loss: 2.6558 | LR: 4.91e-04 | 28232.9ms/step | 0.7 samples/s
+Step 3940/37465 | Loss: 2.8491 | LR: 4.91e-04 | 28230.9ms/step | 0.7 samples/s
+Step 3950/37465 | Loss: 2.3515 | LR: 4.91e-04 | 28229.1ms/step | 0.7 samples/s
+Step 3960/37465 | Loss: 2.5781 | LR: 4.91e-04 | 28219.6ms/step | 0.7 samples/s
+Step 3970/37465 | Loss: 2.3724 | LR: 4.91e-04 | 28214.0ms/step | 0.7 samples/s
+Step 3980/37465 | Loss: 2.5743 | LR: 4.91e-04 | 28203.7ms/step | 0.7 samples/s
+Step 3990/37465 | Loss: 2.5542 | LR: 4.90e-04 | 28218.5ms/step | 0.7 samples/s
+Step 4000/37465 | Loss: 2.5374 | LR: 4.90e-04 | 28220.0ms/step | 0.7 samples/s
+Validation at step 4000: Loss = 2.5376, Perplexity = 12.65
+Step 4010/37465 | Loss: 2.1420 | LR: 4.90e-04 | 29897.7ms/step | 0.7 samples/s
+Step 4020/37465 | Loss: 2.3239 | LR: 4.90e-04 | 28214.7ms/step | 0.7 samples/s
+Step 4030/37465 | Loss: 2.4654 | LR: 4.90e-04 | 28217.3ms/step | 0.7 samples/s
+Step 4040/37465 | Loss: 2.4849 | LR: 4.90e-04 | 28220.1ms/step | 0.7 samples/s
+Step 4050/37465 | Loss: 2.6970 | LR: 4.90e-04 | 28222.9ms/step | 0.7 samples/s
+Step 4060/37465 | Loss: 2.5109 | LR: 4.90e-04 | 28224.9ms/step | 0.7 samples/s
+Step 4070/37465 | Loss: 2.5021 | LR: 4.90e-04 | 28224.8ms/step | 0.7 samples/s
+Step 4080/37465 | Loss: 2.4815 | LR: 4.90e-04 | 28229.8ms/step | 0.7 samples/s
+Step 4090/37465 | Loss: 2.5819 | LR: 4.90e-04 | 28221.3ms/step | 0.7 samples/s
+Step 4100/37465 | Loss: 2.4240 | LR: 4.90e-04 | 28229.8ms/step | 0.7 samples/s
+Step 4110/37465 | Loss: 2.5416 | LR: 4.90e-04 | 28223.2ms/step | 0.7 samples/s
+Step 4120/37465 | Loss: 2.3295 | LR: 4.90e-04 | 28225.4ms/step | 0.7 samples/s
+Step 4130/37465 | Loss: 2.3013 | LR: 4.90e-04 | 28213.1ms/step | 0.7 samples/s
+Step 4140/37465 | Loss: 2.5953 | LR: 4.90e-04 | 28223.6ms/step | 0.7 samples/s
+Step 4150/37465 | Loss: 2.5789 | LR: 4.89e-04 | 28223.4ms/step | 0.7 samples/s
+Step 4160/37465 | Loss: 2.6216 | LR: 4.89e-04 | 28228.4ms/step | 0.7 samples/s
+Step 4170/37465 | Loss: 2.5202 | LR: 4.89e-04 | 28223.6ms/step | 0.7 samples/s
+Step 4180/37465 | Loss: 2.3104 | LR: 4.89e-04 | 28212.4ms/step | 0.7 samples/s
+Step 4190/37465 | Loss: 2.5843 | LR: 4.89e-04 | 28216.6ms/step | 0.7 samples/s
+Step 4200/37465 | Loss: 2.5538 | LR: 4.89e-04 | 28213.4ms/step | 0.7 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_fixed/run_20260709_000503/checkpoint-4200
+Step 4210/37465 | Loss: 2.1640 | LR: 4.89e-04 | 29200.8ms/step | 0.7 samples/s
+Step 4220/37465 | Loss: 2.5114 | LR: 4.89e-04 | 28218.8ms/step | 0.7 samples/s
+Step 4230/37465 | Loss: 2.4094 | LR: 4.89e-04 | 28217.5ms/step | 0.7 samples/s
+Step 4240/37465 | Loss: 2.4589 | LR: 4.89e-04 | 28235.4ms/step | 0.7 samples/s
+Step 4250/37465 | Loss: 2.3746 | LR: 4.89e-04 | 28217.0ms/step | 0.7 samples/s
+Step 4260/37465 | Loss: 2.4317 | LR: 4.89e-04 | 28215.3ms/step | 0.7 samples/s
+Step 4270/37465 | Loss: 2.5068 | LR: 4.89e-04 | 28222.7ms/step | 0.7 samples/s
+Step 4280/37465 | Loss: 2.4866 | LR: 4.89e-04 | 28221.9ms/step | 0.7 samples/s
+Step 4290/37465 | Loss: 2.5645 | LR: 4.89e-04 | 28216.4ms/step | 0.7 samples/s
+Step 4300/37465 | Loss: 2.4434 | LR: 4.89e-04 | 28214.7ms/step | 0.7 samples/s
+Step 4310/37465 | Loss: 2.6270 | LR: 4.88e-04 | 28210.4ms/step | 0.7 samples/s
+Step 4320/37465 | Loss: 2.3658 | LR: 4.88e-04 | 28217.5ms/step | 0.7 samples/s
+Step 4330/37465 | Loss: 2.2291 | LR: 4.88e-04 | 28214.0ms/step | 0.7 samples/s
+Step 4340/37465 | Loss: 2.5317 | LR: 4.88e-04 | 28215.9ms/step | 0.7 samples/s
+Step 4350/37465 | Loss: 2.6447 | LR: 4.88e-04 | 28218.3ms/step | 0.7 samples/s
+Step 4360/37465 | Loss: 2.6117 | LR: 4.88e-04 | 28216.3ms/step | 0.7 samples/s
+Step 4370/37465 | Loss: 2.4347 | LR: 4.88e-04 | 28216.9ms/step | 0.7 samples/s
+Step 4380/37465 | Loss: 2.4317 | LR: 4.88e-04 | 28215.4ms/step | 0.7 samples/s
+Step 4390/37465 | Loss: 2.5850 | LR: 4.88e-04 | 28219.5ms/step | 0.7 samples/s
+Step 4400/37465 | Loss: 2.3529 | LR: 4.88e-04 | 28220.1ms/step | 0.7 samples/s
+Step 4410/37465 | Loss: 2.6017 | LR: 4.88e-04 | 28222.5ms/step | 0.7 samples/s
+Step 4420/37465 | Loss: 2.3320 | LR: 4.88e-04 | 28217.0ms/step | 0.7 samples/s
+Step 4430/37465 | Loss: 2.5218 | LR: 4.88e-04 | 28217.7ms/step | 0.7 samples/s
+Step 4440/37465 | Loss: 2.3608 | LR: 4.88e-04 | 28214.5ms/step | 0.7 samples/s
+Step 4450/37465 | Loss: 2.4483 | LR: 4.88e-04 | 28217.7ms/step | 0.7 samples/s
+Step 4460/37465 | Loss: 2.3357 | LR: 4.88e-04 | 28219.1ms/step | 0.7 samples/s
+Step 4470/37465 | Loss: 2.3514 | LR: 4.87e-04 | 28219.1ms/step | 0.7 samples/s
+Step 4480/37465 | Loss: 2.3672 | LR: 4.87e-04 | 28222.4ms/step | 0.7 samples/s
+Step 4490/37465 | Loss: 2.4445 | LR: 4.87e-04 | 28224.0ms/step | 0.7 samples/s
+Step 4500/37465 | Loss: 2.6611 | LR: 4.87e-04 | 28224.9ms/step | 0.7 samples/s
+Validation at step 4500: Loss = 2.4969, Perplexity = 12.15
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_fixed/run_20260709_000503/checkpoint-4500
+Step 4510/37465 | Loss: 2.3760 | LR: 4.87e-04 | 30836.0ms/step | 0.6 samples/s
+Step 4520/37465 | Loss: 2.4521 | LR: 4.87e-04 | 28215.1ms/step | 0.7 samples/s
+Step 4530/37465 | Loss: 2.5622 | LR: 4.87e-04 | 28212.4ms/step | 0.7 samples/s
+Step 4540/37465 | Loss: 2.4836 | LR: 4.87e-04 | 28211.1ms/step | 0.7 samples/s
+Step 4550/37465 | Loss: 2.4312 | LR: 4.87e-04 | 28212.7ms/step | 0.7 samples/s
+Step 4560/37465 | Loss: 2.4430 | LR: 4.87e-04 | 28209.5ms/step | 0.7 samples/s
+Step 4570/37465 | Loss: 2.2183 | LR: 4.87e-04 | 28211.1ms/step | 0.7 samples/s
+Step 4580/37465 | Loss: 2.5472 | LR: 4.87e-04 | 28216.8ms/step | 0.7 samples/s
+Step 4590/37465 | Loss: 2.4428 | LR: 4.87e-04 | 28211.5ms/step | 0.7 samples/s
+Step 4600/37465 | Loss: 2.6602 | LR: 4.87e-04 | 28209.3ms/step | 0.7 samples/s
+Step 4610/37465 | Loss: 2.4210 | LR: 4.86e-04 | 28213.4ms/step | 0.7 samples/s
+Step 4620/37465 | Loss: 2.5968 | LR: 4.86e-04 | 28211.3ms/step | 0.7 samples/s
+Step 4630/37465 | Loss: 2.5574 | LR: 4.86e-04 | 28212.7ms/step | 0.7 samples/s
+Step 4640/37465 | Loss: 2.3346 | LR: 4.86e-04 | 28214.2ms/step | 0.7 samples/s
+Step 4650/37465 | Loss: 2.2910 | LR: 4.86e-04 | 28208.7ms/step | 0.7 samples/s
+Step 4660/37465 | Loss: 2.4726 | LR: 4.86e-04 | 28207.7ms/step | 0.7 samples/s
+Step 4670/37465 | Loss: 2.5162 | LR: 4.86e-04 | 28211.0ms/step | 0.7 samples/s
+Step 4680/37465 | Loss: 2.4988 | LR: 4.86e-04 | 28208.6ms/step | 0.7 samples/s
+Step 4690/37465 | Loss: 2.6269 | LR: 4.86e-04 | 28209.5ms/step | 0.7 samples/s
+Step 4700/37465 | Loss: 2.5082 | LR: 4.86e-04 | 28210.0ms/step | 0.7 samples/s
+Step 4710/37465 | Loss: 2.3155 | LR: 4.86e-04 | 28209.2ms/step | 0.7 samples/s
+Step 4720/37465 | Loss: 2.5762 | LR: 4.86e-04 | 28210.2ms/step | 0.7 samples/s
+Step 4730/37465 | Loss: 2.3441 | LR: 4.86e-04 | 28210.5ms/step | 0.7 samples/s
+Step 4740/37465 | Loss: 2.5053 | LR: 4.86e-04 | 28209.2ms/step | 0.7 samples/s
+Step 4750/37465 | Loss: 2.4690 | LR: 4.85e-04 | 28208.6ms/step | 0.7 samples/s
+Step 4760/37465 | Loss: 2.4078 | LR: 4.85e-04 | 28206.7ms/step | 0.7 samples/s
+Step 4770/37465 | Loss: 2.3711 | LR: 4.85e-04 | 28212.6ms/step | 0.7 samples/s
+Step 4780/37465 | Loss: 2.5204 | LR: 4.85e-04 | 28212.6ms/step | 0.7 samples/s
+Step 4790/37465 | Loss: 2.6356 | LR: 4.85e-04 | 28207.4ms/step | 0.7 samples/s
+Step 4800/37465 | Loss: 2.5186 | LR: 4.85e-04 | 28206.1ms/step | 0.7 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_fixed/run_20260709_000503/checkpoint-4800
+Step 4810/37465 | Loss: 2.5433 | LR: 4.85e-04 | 29221.6ms/step | 0.7 samples/s
+Step 4820/37465 | Loss: 2.5463 | LR: 4.85e-04 | 28208.8ms/step | 0.7 samples/s
+Step 4830/37465 | Loss: 2.5411 | LR: 4.85e-04 | 28209.5ms/step | 0.7 samples/s
+Step 4840/37465 | Loss: 2.3615 | LR: 4.85e-04 | 28208.7ms/step | 0.7 samples/s
+Step 4850/37465 | Loss: 2.5114 | LR: 4.85e-04 | 28214.8ms/step | 0.7 samples/s
+Step 4860/37465 | Loss: 2.2537 | LR: 4.85e-04 | 28210.8ms/step | 0.7 samples/s
+Step 4870/37465 | Loss: 2.5776 | LR: 4.85e-04 | 28207.8ms/step | 0.7 samples/s
+Step 4880/37465 | Loss: 2.8333 | LR: 4.85e-04 | 28207.0ms/step | 0.7 samples/s
+Step 4890/37465 | Loss: 2.6752 | LR: 4.84e-04 | 28211.2ms/step | 0.7 samples/s
+Step 4900/37465 | Loss: 2.3849 | LR: 4.84e-04 | 28204.7ms/step | 0.7 samples/s
+Step 4910/37465 | Loss: 2.4903 | LR: 4.84e-04 | 28208.2ms/step | 0.7 samples/s
+Step 4920/37465 | Loss: 2.6326 | LR: 4.84e-04 | 28208.8ms/step | 0.7 samples/s
+Step 4930/37465 | Loss: 2.1796 | LR: 4.84e-04 | 28206.6ms/step | 0.7 samples/s
+Step 4940/37465 | Loss: 2.3953 | LR: 4.84e-04 | 28205.5ms/step | 0.7 samples/s
+Step 4950/37465 | Loss: 2.4512 | LR: 4.84e-04 | 28205.5ms/step | 0.7 samples/s
+Step 4960/37465 | Loss: 2.4611 | LR: 4.84e-04 | 28211.1ms/step | 0.7 samples/s
+Step 4970/37465 | Loss: 2.4643 | LR: 4.84e-04 | 28206.5ms/step | 0.7 samples/s
+Step 4980/37465 | Loss: 2.3424 | LR: 4.84e-04 | 28209.6ms/step | 0.7 samples/s
+Step 4990/37465 | Loss: 2.4297 | LR: 4.84e-04 | 28205.5ms/step | 0.7 samples/s
+Step 5000/37465 | Loss: 2.6856 | LR: 4.84e-04 | 28205.9ms/step | 0.7 samples/s
+Validation at step 5000: Loss = 2.4646, Perplexity = 11.76
+Step 5010/37465 | Loss: 2.3841 | LR: 4.84e-04 | 29883.2ms/step | 0.7 samples/s
+Step 5020/37465 | Loss: 2.4298 | LR: 4.83e-04 | 28208.4ms/step | 0.7 samples/s
+Step 5030/37465 | Loss: 2.5368 | LR: 4.83e-04 | 28210.6ms/step | 0.7 samples/s
+Step 5040/37465 | Loss: 2.3500 | LR: 4.83e-04 | 28212.0ms/step | 0.7 samples/s
+Step 5050/37465 | Loss: 2.3595 | LR: 4.83e-04 | 28206.9ms/step | 0.7 samples/s
+Step 5060/37465 | Loss: 2.6008 | LR: 4.83e-04 | 28208.6ms/step | 0.7 samples/s
+Step 5070/37465 | Loss: 2.6280 | LR: 4.83e-04 | 28209.8ms/step | 0.7 samples/s
+Step 5080/37465 | Loss: 2.3638 | LR: 4.83e-04 | 28211.4ms/step | 0.7 samples/s
+Step 5090/37465 | Loss: 2.2817 | LR: 4.83e-04 | 28211.1ms/step | 0.7 samples/s
+Step 5100/37465 | Loss: 2.6281 | LR: 4.83e-04 | 28214.5ms/step | 0.7 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_fixed/run_20260709_000503/checkpoint-5100
+Step 5110/37465 | Loss: 2.4118 | LR: 4.83e-04 | 29238.3ms/step | 0.7 samples/s
+Step 5120/37465 | Loss: 2.2684 | LR: 4.83e-04 | 28211.2ms/step | 0.7 samples/s
+Step 5130/37465 | Loss: 2.3044 | LR: 4.83e-04 | 28206.2ms/step | 0.7 samples/s
+Step 5140/37465 | Loss: 2.3090 | LR: 4.83e-04 | 28206.4ms/step | 0.7 samples/s
+Step 5150/37465 | Loss: 2.4081 | LR: 4.82e-04 | 28210.1ms/step | 0.7 samples/s
+Step 5160/37465 | Loss: 2.2526 | LR: 4.82e-04 | 28208.9ms/step | 0.7 samples/s
+Step 5170/37465 | Loss: 2.7202 | LR: 4.82e-04 | 28206.1ms/step | 0.7 samples/s
+Step 5180/37465 | Loss: 2.3574 | LR: 4.82e-04 | 28207.3ms/step | 0.7 samples/s
+Step 5190/37465 | Loss: 2.5932 | LR: 4.82e-04 | 28202.3ms/step | 0.7 samples/s
+Step 5200/37465 | Loss: 2.5836 | LR: 4.82e-04 | 28204.9ms/step | 0.7 samples/s
+Step 5210/37465 | Loss: 2.3194 | LR: 4.82e-04 | 28200.2ms/step | 0.7 samples/s
+Step 5220/37465 | Loss: 2.3422 | LR: 4.82e-04 | 28204.2ms/step | 0.7 samples/s
+Step 5230/37465 | Loss: 2.4587 | LR: 4.82e-04 | 28209.2ms/step | 0.7 samples/s
+Step 5240/37465 | Loss: 2.4866 | LR: 4.82e-04 | 28205.2ms/step | 0.7 samples/s
+Step 5250/37465 | Loss: 2.6048 | LR: 4.82e-04 | 28202.8ms/step | 0.7 samples/s
+Step 5260/37465 | Loss: 2.5828 | LR: 4.82e-04 | 28205.2ms/step | 0.7 samples/s
+Step 5270/37465 | Loss: 2.4557 | LR: 4.82e-04 | 28205.5ms/step | 0.7 samples/s
+Step 5280/37465 | Loss: 2.3707 | LR: 4.81e-04 | 28206.2ms/step | 0.7 samples/s
+Step 5290/37465 | Loss: 2.5570 | LR: 4.81e-04 | 28204.7ms/step | 0.7 samples/s
+Step 5300/37465 | Loss: 2.6720 | LR: 4.81e-04 | 28205.3ms/step | 0.7 samples/s
+Step 5310/37465 | Loss: 2.3422 | LR: 4.81e-04 | 28207.0ms/step | 0.7 samples/s
+Step 5320/37465 | Loss: 2.4032 | LR: 4.81e-04 | 28200.2ms/step | 0.7 samples/s
+Step 5330/37465 | Loss: 2.1944 | LR: 4.81e-04 | 28200.5ms/step | 0.7 samples/s
+Step 5340/37465 | Loss: 2.4332 | LR: 4.81e-04 | 28201.8ms/step | 0.7 samples/s
+Step 5350/37465 | Loss: 2.6754 | LR: 4.81e-04 | 28202.4ms/step | 0.7 samples/s
+Step 5360/37465 | Loss: 2.5270 | LR: 4.81e-04 | 28202.3ms/step | 0.7 samples/s
+Step 5370/37465 | Loss: 2.5446 | LR: 4.81e-04 | 28204.6ms/step | 0.7 samples/s
+Step 5380/37465 | Loss: 2.4730 | LR: 4.81e-04 | 28203.4ms/step | 0.7 samples/s
+Step 5390/37465 | Loss: 2.3932 | LR: 4.81e-04 | 28200.6ms/step | 0.7 samples/s
+Step 5400/37465 | Loss: 2.5352 | LR: 4.80e-04 | 28203.8ms/step | 0.7 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_fixed/run_20260709_000503/checkpoint-5400
+Step 5410/37465 | Loss: 2.4252 | LR: 4.80e-04 | 29245.7ms/step | 0.7 samples/s
+Step 5420/37465 | Loss: 2.3871 | LR: 4.80e-04 | 28197.4ms/step | 0.7 samples/s
+Step 5430/37465 | Loss: 2.3216 | LR: 4.80e-04 | 28202.6ms/step | 0.7 samples/s
+Step 5440/37465 | Loss: 2.4534 | LR: 4.80e-04 | 28196.8ms/step | 0.7 samples/s
+Step 5450/37465 | Loss: 2.3695 | LR: 4.80e-04 | 28201.6ms/step | 0.7 samples/s
+Step 5460/37465 | Loss: 2.5923 | LR: 4.80e-04 | 28200.8ms/step | 0.7 samples/s
+Step 5470/37465 | Loss: 2.5585 | LR: 4.80e-04 | 28204.9ms/step | 0.7 samples/s
+Step 5480/37465 | Loss: 2.3787 | LR: 4.80e-04 | 28202.3ms/step | 0.7 samples/s
+Step 5490/37465 | Loss: 2.6341 | LR: 4.80e-04 | 28205.3ms/step | 0.7 samples/s
+Step 5500/37465 | Loss: 2.3661 | LR: 4.80e-04 | 28200.8ms/step | 0.7 samples/s
+Validation at step 5500: Loss = 2.4438, Perplexity = 11.52
+Step 5510/37465 | Loss: 2.2862 | LR: 4.80e-04 | 29861.9ms/step | 0.7 samples/s
+Step 5520/37465 | Loss: 2.3151 | LR: 4.79e-04 | 28202.6ms/step | 0.7 samples/s
+Step 5530/37465 | Loss: 2.4188 | LR: 4.79e-04 | 28204.7ms/step | 0.7 samples/s
+Step 5540/37465 | Loss: 2.2129 | LR: 4.79e-04 | 28204.2ms/step | 0.7 samples/s
+Step 5550/37465 | Loss: 2.2930 | LR: 4.79e-04 | 28201.1ms/step | 0.7 samples/s
+Step 5560/37465 | Loss: 2.6650 | LR: 4.79e-04 | 28203.4ms/step | 0.7 samples/s
+Step 5570/37465 | Loss: 2.2805 | LR: 4.79e-04 | 28199.7ms/step | 0.7 samples/s
+Step 5580/37465 | Loss: 2.4486 | LR: 4.79e-04 | 28199.6ms/step | 0.7 samples/s
+Step 5590/37465 | Loss: 2.4593 | LR: 4.79e-04 | 28210.4ms/step | 0.7 samples/s
+Step 5600/37465 | Loss: 2.3003 | LR: 4.79e-04 | 28202.3ms/step | 0.7 samples/s
+Step 5610/37465 | Loss: 2.4029 | LR: 4.79e-04 | 28202.0ms/step | 0.7 samples/s
+Step 5620/37465 | Loss: 2.4410 | LR: 4.79e-04 | 28199.9ms/step | 0.7 samples/s
+Step 5630/37465 | Loss: 2.5237 | LR: 4.79e-04 | 28198.6ms/step | 0.7 samples/s
+Step 5640/37465 | Loss: 2.3727 | LR: 4.78e-04 | 28205.3ms/step | 0.7 samples/s
+Step 5650/37465 | Loss: 2.2762 | LR: 4.78e-04 | 28206.5ms/step | 0.7 samples/s
+Step 5660/37465 | Loss: 2.4433 | LR: 4.78e-04 | 28205.9ms/step | 0.7 samples/s
+Step 5670/37465 | Loss: 2.5133 | LR: 4.78e-04 | 28204.4ms/step | 0.7 samples/s
+Step 5680/37465 | Loss: 2.4399 | LR: 4.78e-04 | 28201.1ms/step | 0.7 samples/s
+Step 5690/37465 | Loss: 2.3988 | LR: 4.78e-04 | 28202.2ms/step | 0.7 samples/s
+Step 5700/37465 | Loss: 2.5529 | LR: 4.78e-04 | 28200.1ms/step | 0.7 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_fixed/run_20260709_000503/checkpoint-5700
+Step 5710/37465 | Loss: 2.4231 | LR: 4.78e-04 | 29228.6ms/step | 0.7 samples/s
+Step 5720/37465 | Loss: 2.3970 | LR: 4.78e-04 | 28195.9ms/step | 0.7 samples/s
+Step 5730/37465 | Loss: 2.4189 | LR: 4.78e-04 | 28207.0ms/step | 0.7 samples/s
+Step 5740/37465 | Loss: 2.3991 | LR: 4.78e-04 | 28201.7ms/step | 0.7 samples/s
+W0711 07:56:50.458000 35834 site-packages/torch/distributed/elastic/agent/server/api.py:739] Received 15 death signal, shutting down workers
+W0711 07:56:50.462000 35834 site-packages/torch/distributed/elastic/multiprocessing/api.py:1012] Sending process 35847 closing signal SIGTERM
+W0711 07:56:50.464000 35834 site-packages/torch/distributed/elastic/multiprocessing/api.py:1012] Sending process 35848 closing signal SIGTERM
+W0711 07:56:50.465000 35834 site-packages/torch/distributed/elastic/multiprocessing/api.py:1012] Sending process 35849 closing signal SIGTERM
+W0711 07:56:50.465000 35834 site-packages/torch/distributed/elastic/multiprocessing/api.py:1012] Sending process 35850 closing signal SIGTERM
+Traceback (most recent call last):
+  File "/mloscratch/homes/navasard/conda/bin/torchrun", line 6, in <module>
+    sys.exit(main())
+             ~~~~^^
+  File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 362, in wrapper
+    return f(*args, **kwargs)
+  File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/distributed/run.py", line 990, in main
+    run(args)
+    ~~~^^^^^^
+  File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/distributed/run.py", line 981, in run
+    elastic_launch(
+    ~~~~~~~~~~~~~~~
+        config=config,
+        ~~~~~~~~~~~~~~
+        entrypoint=cmd,
+        ~~~~~~~~~~~~~~~
+    )(*cmd_args)
+    ~^^^^^^^^^^^
+  File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/distributed/launcher/api.py", line 170, in __call__
+    return launch_agent(self._config, self._entrypoint, list(args))
+  File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/distributed/launcher/api.py", line 308, in launch_agent
+    result = agent.run()
+  File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/distributed/elastic/metrics/api.py", line 134, in wrapper
+    result = f(*args, **kwargs)
+  File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/distributed/elastic/agent/server/api.py", line 731, in run
+    result = self._invoke_run(role)
+  File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/distributed/elastic/agent/server/api.py", line 908, in _invoke_run
+    time.sleep(monitor_interval)
+    ~~~~~~~~~~^^^^^^^^^^^^^^^^^^
+  File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/distributed/elastic/multiprocessing/api.py", line 86, in _terminate_process_handler
+    raise SignalException(f"Process {os.getpid()} got signal: {sigval}", sigval=sigval)
+torch.distributed.elastic.multiprocessing.api.SignalException: Process 35834 got signal: 15

epfl_final/train_logs/resume_fixed30k_h200.log

@@ -0,0 +1,436 @@
+W0711 07:57:24.347000 2190 site-packages/torch/distributed/run.py:851] 
+W0711 07:57:24.347000 2190 site-packages/torch/distributed/run.py:851] *****************************************
+W0711 07:57:24.347000 2190 site-packages/torch/distributed/run.py:851] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. 
+W0711 07:57:24.347000 2190 site-packages/torch/distributed/run.py:851] *****************************************
+Auto-resuming from latest checkpoint: /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_fixed/run_20260709_000503/checkpoint-5700
+Resuming: Using existing output directory /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_fixed/run_20260709_000503
+Running config:
+{
+  "model_path": "/mloscratch/homes/navasard/qwen3_06b_base",
+  "model_dtype": "bfloat16",
+  "init_from_scratch": true,
+  "train_path": "/mloscratch/homes/navasard/data/longctx30k/tokenized",
+  "val_path": null,
+  "max_seq_length": 30720,
+  "num_workers": 4,
+  "output_dir": "/mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_fixed/run_20260709_000503",
+  "resume_from_checkpoint": "/mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_fixed/run_20260709_000503/checkpoint-5700",
+  "force_new_run": false,
+  "num_epochs": 1,
+  "max_steps": null,
+  "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": 300,
+  "eval_steps": 500,
+  "eval_batches": 20,
+  "logging_steps": 10,
+  "seed": 42,
+  "verbose": true,
+  "checkpoint_activations": true,
+  "strategy": "ddp",
+  "mixed_precision": "bf16",
+  "wandb_enabled": true,
+  "wandb_project": "stochastic_sparsity",
+  "wandb_name": "1.2b fixed 30k longctx",
+  "wandb_tags": [
+    "qwen3",
+    "1.2b",
+    "fixed",
+    "30k",
+    "longctx",
+    "b128",
+    "scaleladder"
+  ],
+  "hidden_size": 1536,
+  "intermediate_size": 6144,
+  "num_attention_heads": 12,
+  "num_key_value_heads": 4,
+  "num_hidden_layers": null,
+  "head_dim": null,
+  "attention_type": "sparse_stochastic",
+  "sparse_block_size": 128,
+  "sparse_num_random_blocks": 120,
+  "sparse_mask_pool_size": 1,
+  "sparse_val_dense_compare": false,
+  "sparse_layer_indices": null,
+  "sparse_mask_cache_dir": "/mloscratch/homes/navasard/sparse/mask_cache_30k_b128_nrb120",
+  "sparse_mask_seed": 1,
+  "dense_mode": false,
+  "dense_step_intervals": null
+}
+Loading tokenizer from /mloscratch/homes/navasard/qwen3_06b_base...
+Checkpoint activations: True
+Checkpoint activations: True
+Checkpoint activations: True
+Resuming from checkpoint: /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_fixed/run_20260709_000503/checkpoint-5700
+Checkpoint activations: True
+
+
+================================================================================
+MODEL ARCHITECTURE
+================================================================================
+Model config:
+  vocab_size: 151936
+  hidden_size: 1536
+  num_hidden_layers: 28
+  num_attention_heads: 12
+  num_key_value_heads: 4
+  head_dim: 128
+  intermediate_size: 6144
+  max_position_embeddings: 32768
+  rope_theta: 1000000
+  attention_type: sparse_stochastic
+  sparse_block_size: 128
+  sparse_num_random_blocks: 120
+  sparse_mask_pool_size: 1
+  sparse_max_seq_length: 30720
+
+Parameter count:
+  Total parameters: 1,202,352,640
+  Trainable parameters: 1,202,352,640
+  Non-trainable parameters: 0
+  Model size (BF16): 2.40 GB
+
+--------------------------------------------------------------------------------
+LAYER-BY-LAYER SUMMARY
+--------------------------------------------------------------------------------
+/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/nn/attention/flex_attention.py:1704: UserWarning: flex_attention called without torch.compile() - this will use an unfused implementation that materializes the full scores matrix instead of generating a fused kernel.
+
+SOLUTION: Use torch.compile(flex_attention)(...)
+
+If you want to debug your score_mod/mask_mod, you can set:
+torch.nn.attention.flex_attention._FLEX_ATTENTION_DISABLE_COMPILE_DEBUG = True
+
+This will allow you to use print statements or breakpoints. Note: This doesn't work with the backwards pass and may produce incorrect results.
+  _warn_once(
+Could not generate detailed summary: Failed to run torchinfo. See above stack traces for more details. Executed layers up to: [Embedding: 2, Qwen3RMSNorm: 4, Linear: 5, Linear: 5, Linear: 5, Qwen3RMSNorm: 5, Qwen3RMSNorm: 5, Qwen3RotaryEmbedding: 5]
+  NOTE: torchinfo double-counts tied embed_tokens/lm_head weight (233,373,696 params). True unique parameter count: 1,202,352,640 (1202.4M)
+================================================================================
+
+Loading dataset from /mloscratch/homes/navasard/data/longctx30k/tokenized...
+Using pre-tokenized dataset (no runtime tokenization)
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/longctx30k/tokenized...
+Loaded 788736 pre-tokenized examples (stored_len=30720, pack_factor=1, packed_examples=788736)
+Dataset split: 749299 train, 39437 validation
+Training examples (packed, train split):  749299
+Validation examples (packed, val split):  39437
+Per-device batch size:                    1
+World size (GPUs):                        4
+Gradient accumulation steps:              5
+Effective batch size (seqs/step):         20
+Sequence length:                          30720
+Tokens per optimizer step:                614,400
+Micro-batches per epoch (per rank):       187325
+Optimizer steps per epoch:                37465
+Total optimizer steps:                    37465
+Loading checkpoint from /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_fixed/run_20260709_000503/checkpoint-5700...
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/longctx30k/tokenized...
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/longctx30k/tokenized...
+Loaded 788736 pre-tokenized examples (stored_len=30720, pack_factor=1, packed_examples=788736)
+Loaded 788736 pre-tokenized examples (stored_len=30720, pack_factor=1, packed_examples=788736)
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/longctx30k/tokenized...
+Loaded 788736 pre-tokenized examples (stored_len=30720, pack_factor=1, packed_examples=788736)
+Overriding checkpoint LR with config LR: 0.0005
+Resumed from checkpoint at step 5700 (epoch 0, micro_step 28500)
+Resuming training from step 5700
+Resuming wandb run 9blizq3m
+wandb: [wandb.login()] Loaded credentials for https://api.wandb.ai from WANDB_API_KEY.
+wandb: Currently logged in as: ruben-navasardyan4869 (ruben-navasardyan4869-epfl) to https://api.wandb.ai. Use `wandb login --relogin` to force relogin
+wandb: Tracking run with wandb version 0.26.1
+wandb: Run data is saved locally in /mloscratch/homes/navasard/sparse/SparseAtt/wandb/run-20260711_080025-9blizq3m
+wandb: Run `wandb offline` to turn off syncing.
+wandb: Resuming run 1.2b fixed 30k longctx
+wandb: ⭐️ View project at https://wandb.ai/ruben-navasardyan4869-epfl/stochastic_sparsity
+wandb: 🚀 View run at https://wandb.ai/ruben-navasardyan4869-epfl/stochastic_sparsity/runs/9blizq3m
+Starting training...
+Validation at step 5700: Loss = 2.4255, Perplexity = 11.31
+wandb: WARNING Tried to log to step 5700 that is less than the current step 5731. Steps must be monotonically increasing, so this data will be ignored. See https://wandb.me/define-metric to log data out of order.
+Switching attention mode for step 5701: sparse_stochastic (configured sparse attention)
+Step 5710/37465 | Loss: 2.4233 | LR: 4.78e-04 | 25615.7ms/step | 0.8 samples/s
+wandb: WARNING Tried to log to step 5710 that is less than the current step 5731. Steps must be monotonically increasing, so this data will be ignored. See https://wandb.me/define-metric to log data out of order.
+Step 5720/37465 | Loss: 2.3971 | LR: 4.78e-04 | 8106.0ms/step | 2.5 samples/s
+wandb: WARNING Tried to log to step 5720 that is less than the current step 5731. Steps must be monotonically increasing, so this data will be ignored. See https://wandb.me/define-metric to log data out of order.
+Step 5730/37465 | Loss: 2.4188 | LR: 4.78e-04 | 8148.0ms/step | 2.5 samples/s
+wandb: WARNING Tried to log to step 5730 that is less than the current step 5731. Steps must be monotonically increasing, so this data will be ignored. See https://wandb.me/define-metric to log data out of order.
+Step 5740/37465 | Loss: 2.3991 | LR: 4.78e-04 | 8130.2ms/step | 2.5 samples/s
+Step 5750/37465 | Loss: 2.4403 | LR: 4.77e-04 | 8157.3ms/step | 2.5 samples/s
+Step 5760/37465 | Loss: 2.1459 | LR: 4.77e-04 | 8150.0ms/step | 2.5 samples/s
+Step 5770/37465 | Loss: 2.4807 | LR: 4.77e-04 | 8176.4ms/step | 2.4 samples/s
+Step 5780/37465 | Loss: 2.4041 | LR: 4.77e-04 | 8155.3ms/step | 2.5 samples/s
+Step 5790/37465 | Loss: 2.2887 | LR: 4.77e-04 | 8185.4ms/step | 2.4 samples/s
+Step 5800/37465 | Loss: 2.4940 | LR: 4.77e-04 | 8195.5ms/step | 2.4 samples/s
+Step 5810/37465 | Loss: 2.4236 | LR: 4.77e-04 | 8181.2ms/step | 2.4 samples/s
+Step 5820/37465 | Loss: 2.1601 | LR: 4.77e-04 | 8181.0ms/step | 2.4 samples/s
+Step 5830/37465 | Loss: 2.3791 | LR: 4.77e-04 | 8183.4ms/step | 2.4 samples/s
+Step 5840/37465 | Loss: 2.4035 | LR: 4.77e-04 | 8200.5ms/step | 2.4 samples/s
+Step 5850/37465 | Loss: 2.2279 | LR: 4.77e-04 | 8146.1ms/step | 2.5 samples/s
+Step 5860/37465 | Loss: 2.1539 | LR: 4.76e-04 | 8175.4ms/step | 2.4 samples/s
+Step 5870/37465 | Loss: 2.4310 | LR: 4.76e-04 | 8168.1ms/step | 2.4 samples/s
+Step 5880/37465 | Loss: 2.4335 | LR: 4.76e-04 | 8099.7ms/step | 2.5 samples/s
+Step 5890/37465 | Loss: 2.5099 | LR: 4.76e-04 | 8151.9ms/step | 2.5 samples/s
+Step 5900/37465 | Loss: 2.5700 | LR: 4.76e-04 | 8134.2ms/step | 2.5 samples/s
+Step 5910/37465 | Loss: 2.2665 | LR: 4.76e-04 | 8145.1ms/step | 2.5 samples/s
+Step 5920/37465 | Loss: 2.4116 | LR: 4.76e-04 | 8143.3ms/step | 2.5 samples/s
+Step 5930/37465 | Loss: 2.3634 | LR: 4.76e-04 | 8128.2ms/step | 2.5 samples/s
+Step 5940/37465 | Loss: 2.6368 | LR: 4.76e-04 | 8131.6ms/step | 2.5 samples/s
+Step 5950/37465 | Loss: 2.5229 | LR: 4.76e-04 | 8116.3ms/step | 2.5 samples/s
+Step 5960/37465 | Loss: 2.2806 | LR: 4.76e-04 | 8141.0ms/step | 2.5 samples/s
+Step 5970/37465 | Loss: 2.2798 | LR: 4.75e-04 | 8138.6ms/step | 2.5 samples/s
+Step 5980/37465 | Loss: 2.3335 | LR: 4.75e-04 | 8157.1ms/step | 2.5 samples/s
+Step 5990/37465 | Loss: 2.3756 | LR: 4.75e-04 | 8141.2ms/step | 2.5 samples/s
+Step 6000/37465 | Loss: 2.2740 | LR: 4.75e-04 | 8117.4ms/step | 2.5 samples/s
+Validation at step 6000: Loss = 2.4133, Perplexity = 11.17
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_fixed/run_20260709_000503/checkpoint-6000
+Step 6010/37465 | Loss: 2.2714 | LR: 4.75e-04 | 9634.0ms/step | 2.1 samples/s
+Step 6020/37465 | Loss: 2.5321 | LR: 4.75e-04 | 8083.3ms/step | 2.5 samples/s
+Step 6030/37465 | Loss: 2.4918 | LR: 4.75e-04 | 8106.1ms/step | 2.5 samples/s
+Step 6040/37465 | Loss: 2.3314 | LR: 4.75e-04 | 8101.6ms/step | 2.5 samples/s
+Step 6050/37465 | Loss: 2.2412 | LR: 4.75e-04 | 8168.9ms/step | 2.4 samples/s
+Step 6060/37465 | Loss: 2.5855 | LR: 4.75e-04 | 8174.5ms/step | 2.4 samples/s
+Step 6070/37465 | Loss: 2.3413 | LR: 4.75e-04 | 8172.7ms/step | 2.4 samples/s
+Step 6080/37465 | Loss: 2.5651 | LR: 4.74e-04 | 8196.3ms/step | 2.4 samples/s
+Step 6090/37465 | Loss: 2.3406 | LR: 4.74e-04 | 8193.4ms/step | 2.4 samples/s
+Step 6100/37465 | Loss: 2.1750 | LR: 4.74e-04 | 8184.3ms/step | 2.4 samples/s
+Step 6110/37465 | Loss: 2.5999 | LR: 4.74e-04 | 8157.3ms/step | 2.5 samples/s
+Step 6120/37465 | Loss: 2.5244 | LR: 4.74e-04 | 8183.5ms/step | 2.4 samples/s
+Step 6130/37465 | Loss: 2.4220 | LR: 4.74e-04 | 8200.1ms/step | 2.4 samples/s
+Step 6140/37465 | Loss: 2.3484 | LR: 4.74e-04 | 8198.3ms/step | 2.4 samples/s
+Step 6150/37465 | Loss: 2.5200 | LR: 4.74e-04 | 8141.7ms/step | 2.5 samples/s
+Step 6160/37465 | Loss: 2.4814 | LR: 4.74e-04 | 8120.6ms/step | 2.5 samples/s
+Step 6170/37465 | Loss: 2.3418 | LR: 4.74e-04 | 8141.2ms/step | 2.5 samples/s
+Step 6180/37465 | Loss: 2.4501 | LR: 4.73e-04 | 8049.7ms/step | 2.5 samples/s
+Step 6190/37465 | Loss: 2.4962 | LR: 4.73e-04 | 8232.9ms/step | 2.4 samples/s
+Step 6200/37465 | Loss: 2.3876 | LR: 4.73e-04 | 8211.8ms/step | 2.4 samples/s
+Step 6210/37465 | Loss: 2.2009 | LR: 4.73e-04 | 8231.0ms/step | 2.4 samples/s
+Step 6220/37465 | Loss: 2.4581 | LR: 4.73e-04 | 8234.9ms/step | 2.4 samples/s
+Step 6230/37465 | Loss: 2.3382 | LR: 4.73e-04 | 8233.5ms/step | 2.4 samples/s
+Step 6240/37465 | Loss: 2.4628 | LR: 4.73e-04 | 8188.2ms/step | 2.4 samples/s
+Step 6250/37465 | Loss: 2.3168 | LR: 4.73e-04 | 8166.0ms/step | 2.4 samples/s
+Step 6260/37465 | Loss: 2.2641 | LR: 4.73e-04 | 8188.7ms/step | 2.4 samples/s
+Step 6270/37465 | Loss: 1.9939 | LR: 4.73e-04 | 8173.5ms/step | 2.4 samples/s
+Step 6280/37465 | Loss: 2.2060 | LR: 4.73e-04 | 8177.8ms/step | 2.4 samples/s
+Step 6290/37465 | Loss: 2.3970 | LR: 4.72e-04 | 8196.8ms/step | 2.4 samples/s
+Step 6300/37465 | Loss: 2.4667 | LR: 4.72e-04 | 8114.9ms/step | 2.5 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_fixed/run_20260709_000503/checkpoint-6300
+Step 6310/37465 | Loss: 2.2527 | LR: 4.72e-04 | 9896.9ms/step | 2.0 samples/s
+Step 6320/37465 | Loss: 2.3347 | LR: 4.72e-04 | 8162.6ms/step | 2.5 samples/s
+Step 6330/37465 | Loss: 2.4564 | LR: 4.72e-04 | 8161.8ms/step | 2.5 samples/s
+Step 6340/37465 | Loss: 2.2622 | LR: 4.72e-04 | 8169.7ms/step | 2.4 samples/s
+Step 6350/37465 | Loss: 2.6875 | LR: 4.72e-04 | 8166.6ms/step | 2.4 samples/s
+Step 6360/37465 | Loss: 2.2354 | LR: 4.72e-04 | 8183.3ms/step | 2.4 samples/s
+Step 6370/37465 | Loss: 2.1165 | LR: 4.72e-04 | 8210.6ms/step | 2.4 samples/s
+Step 6380/37465 | Loss: 2.5141 | LR: 4.72e-04 | 8191.8ms/step | 2.4 samples/s
+Step 6390/37465 | Loss: 2.5378 | LR: 4.71e-04 | 8218.6ms/step | 2.4 samples/s
+Step 6400/37465 | Loss: 2.3583 | LR: 4.71e-04 | 8170.3ms/step | 2.4 samples/s
+Step 6410/37465 | Loss: 2.4088 | LR: 4.71e-04 | 8177.7ms/step | 2.4 samples/s
+Step 6420/37465 | Loss: 2.2487 | LR: 4.71e-04 | 8226.4ms/step | 2.4 samples/s
+Step 6430/37465 | Loss: 2.2890 | LR: 4.71e-04 | 8239.9ms/step | 2.4 samples/s
+Step 6440/37465 | Loss: 2.2322 | LR: 4.71e-04 | 8236.4ms/step | 2.4 samples/s
+Step 6450/37465 | Loss: 2.3595 | LR: 4.71e-04 | 8172.2ms/step | 2.4 samples/s
+Step 6460/37465 | Loss: 2.3734 | LR: 4.71e-04 | 8105.1ms/step | 2.5 samples/s
+Step 6470/37465 | Loss: 2.3530 | LR: 4.71e-04 | 8126.2ms/step | 2.5 samples/s
+Step 6480/37465 | Loss: 2.3063 | LR: 4.71e-04 | 8204.0ms/step | 2.4 samples/s
+Step 6490/37465 | Loss: 2.3493 | LR: 4.70e-04 | 8162.5ms/step | 2.5 samples/s
+Step 6500/37465 | Loss: 2.5273 | LR: 4.70e-04 | 8223.7ms/step | 2.4 samples/s
+Validation at step 6500: Loss = 2.3911, Perplexity = 10.93
+Step 6510/37465 | Loss: 2.3985 | LR: 4.70e-04 | 8884.6ms/step | 2.3 samples/s
+Step 6520/37465 | Loss: 2.4548 | LR: 4.70e-04 | 8207.0ms/step | 2.4 samples/s
+Step 6530/37465 | Loss: 2.2737 | LR: 4.70e-04 | 8177.7ms/step | 2.4 samples/s
+Step 6540/37465 | Loss: 2.4096 | LR: 4.70e-04 | 8219.9ms/step | 2.4 samples/s
+Step 6550/37465 | Loss: 2.3655 | LR: 4.70e-04 | 8177.0ms/step | 2.4 samples/s
+Step 6560/37465 | Loss: 2.4906 | LR: 4.70e-04 | 8139.6ms/step | 2.5 samples/s
+Step 6570/37465 | Loss: 2.5839 | LR: 4.70e-04 | 8224.5ms/step | 2.4 samples/s
+Step 6580/37465 | Loss: 2.2572 | LR: 4.70e-04 | 7998.1ms/step | 2.5 samples/s
+Step 6590/37465 | Loss: 2.2486 | LR: 4.69e-04 | 8212.1ms/step | 2.4 samples/s
+Step 6600/37465 | Loss: 2.2569 | LR: 4.69e-04 | 8185.3ms/step | 2.4 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_fixed/run_20260709_000503/checkpoint-6600
+Step 6610/37465 | Loss: 2.3613 | LR: 4.69e-04 | 9094.6ms/step | 2.2 samples/s
+Step 6620/37465 | Loss: 2.5718 | LR: 4.69e-04 | 8179.7ms/step | 2.4 samples/s
+Step 6630/37465 | Loss: 2.3047 | LR: 4.69e-04 | 8214.3ms/step | 2.4 samples/s
+Step 6640/37465 | Loss: 2.3896 | LR: 4.69e-04 | 8201.9ms/step | 2.4 samples/s
+Step 6650/37465 | Loss: 2.5540 | LR: 4.69e-04 | 8178.4ms/step | 2.4 samples/s
+Step 6660/37465 | Loss: 2.4475 | LR: 4.69e-04 | 8090.6ms/step | 2.5 samples/s
+Step 6670/37465 | Loss: 2.2657 | LR: 4.69e-04 | 8114.4ms/step | 2.5 samples/s
+Step 6680/37465 | Loss: 2.4151 | LR: 4.69e-04 | 8184.5ms/step | 2.4 samples/s
+Step 6690/37465 | Loss: 2.6053 | LR: 4.68e-04 | 8201.4ms/step | 2.4 samples/s
+Step 6700/37465 | Loss: 2.1858 | LR: 4.68e-04 | 8195.0ms/step | 2.4 samples/s
+Step 6710/37465 | Loss: 2.5396 | LR: 4.68e-04 | 8219.8ms/step | 2.4 samples/s
+Step 6720/37465 | Loss: 2.2380 | LR: 4.68e-04 | 8210.9ms/step | 2.4 samples/s
+Step 6730/37465 | Loss: 2.4316 | LR: 4.68e-04 | 8185.2ms/step | 2.4 samples/s
+Step 6740/37465 | Loss: 2.2572 | LR: 4.68e-04 | 8208.6ms/step | 2.4 samples/s
+Step 6750/37465 | Loss: 2.2984 | LR: 4.68e-04 | 8211.7ms/step | 2.4 samples/s
+Step 6760/37465 | Loss: 2.4685 | LR: 4.68e-04 | 8202.9ms/step | 2.4 samples/s
+Step 6770/37465 | Loss: 2.2939 | LR: 4.68e-04 | 8199.0ms/step | 2.4 samples/s
+Step 6780/37465 | Loss: 2.5630 | LR: 4.67e-04 | 8211.3ms/step | 2.4 samples/s
+Step 6790/37465 | Loss: 2.4770 | LR: 4.67e-04 | 8196.7ms/step | 2.4 samples/s
+Step 6800/37465 | Loss: 2.3050 | LR: 4.67e-04 | 8169.9ms/step | 2.4 samples/s
+Step 6810/37465 | Loss: 2.5238 | LR: 4.67e-04 | 8170.5ms/step | 2.4 samples/s
+Step 6820/37465 | Loss: 2.5106 | LR: 4.67e-04 | 8176.7ms/step | 2.4 samples/s
+Step 6830/37465 | Loss: 2.6382 | LR: 4.67e-04 | 8194.0ms/step | 2.4 samples/s
+Step 6840/37465 | Loss: 2.4227 | LR: 4.67e-04 | 8168.3ms/step | 2.4 samples/s
+Step 6850/37465 | Loss: 2.3855 | LR: 4.67e-04 | 8183.9ms/step | 2.4 samples/s
+Step 6860/37465 | Loss: 2.3495 | LR: 4.67e-04 | 8175.4ms/step | 2.4 samples/s
+Step 6870/37465 | Loss: 2.4752 | LR: 4.66e-04 | 8130.8ms/step | 2.5 samples/s
+Step 6880/37465 | Loss: 2.3842 | LR: 4.66e-04 | 8118.0ms/step | 2.5 samples/s
+Step 6890/37465 | Loss: 2.5300 | LR: 4.66e-04 | 8087.4ms/step | 2.5 samples/s
+Step 6900/37465 | Loss: 2.3989 | LR: 4.66e-04 | 8137.5ms/step | 2.5 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_fixed/run_20260709_000503/checkpoint-6900
+Step 6910/37465 | Loss: 2.4433 | LR: 4.66e-04 | 8961.7ms/step | 2.2 samples/s
+Step 6920/37465 | Loss: 2.4076 | LR: 4.66e-04 | 8075.6ms/step | 2.5 samples/s
+Step 6930/37465 | Loss: 2.2864 | LR: 4.66e-04 | 8103.1ms/step | 2.5 samples/s
+Step 6940/37465 | Loss: 2.1097 | LR: 4.66e-04 | 8107.7ms/step | 2.5 samples/s
+Step 6950/37465 | Loss: 2.5191 | LR: 4.66e-04 | 8043.5ms/step | 2.5 samples/s
+Step 6960/37465 | Loss: 2.4665 | LR: 4.66e-04 | 8100.3ms/step | 2.5 samples/s
+Step 6970/37465 | Loss: 2.4891 | LR: 4.65e-04 | 8025.7ms/step | 2.5 samples/s
+Step 6980/37465 | Loss: 2.4306 | LR: 4.65e-04 | 8115.8ms/step | 2.5 samples/s
+Step 6990/37465 | Loss: 2.4929 | LR: 4.65e-04 | 8124.1ms/step | 2.5 samples/s
+Step 7000/37465 | Loss: 2.1510 | LR: 4.65e-04 | 8104.7ms/step | 2.5 samples/s
+Validation at step 7000: Loss = 2.3711, Perplexity = 10.71
+Step 7010/37465 | Loss: 2.2900 | LR: 4.65e-04 | 8746.1ms/step | 2.3 samples/s
+Step 7020/37465 | Loss: 2.2995 | LR: 4.65e-04 | 8105.8ms/step | 2.5 samples/s
+Step 7030/37465 | Loss: 2.3253 | LR: 4.65e-04 | 8143.7ms/step | 2.5 samples/s
+Step 7040/37465 | Loss: 2.1986 | LR: 4.65e-04 | 8143.9ms/step | 2.5 samples/s
+Step 7050/37465 | Loss: 2.3172 | LR: 4.65e-04 | 8174.8ms/step | 2.4 samples/s
+Step 7060/37465 | Loss: 2.4077 | LR: 4.64e-04 | 8164.4ms/step | 2.4 samples/s
+Step 7070/37465 | Loss: 2.2571 | LR: 4.64e-04 | 8165.3ms/step | 2.4 samples/s
+Step 7080/37465 | Loss: 2.4516 | LR: 4.64e-04 | 8167.4ms/step | 2.4 samples/s
+Step 7090/37465 | Loss: 2.2556 | LR: 4.64e-04 | 8026.7ms/step | 2.5 samples/s
+Step 7100/37465 | Loss: 2.3510 | LR: 4.64e-04 | 8176.6ms/step | 2.4 samples/s
+Step 7110/37465 | Loss: 2.1612 | LR: 4.64e-04 | 8077.1ms/step | 2.5 samples/s
+Step 7120/37465 | Loss: 1.9000 | LR: 4.64e-04 | 8154.8ms/step | 2.5 samples/s
+Step 7130/37465 | Loss: 2.3495 | LR: 4.64e-04 | 8238.0ms/step | 2.4 samples/s
+Step 7140/37465 | Loss: 2.4536 | LR: 4.64e-04 | 8211.4ms/step | 2.4 samples/s
+Step 7150/37465 | Loss: 2.5172 | LR: 4.63e-04 | 8118.8ms/step | 2.5 samples/s
+Step 7160/37465 | Loss: 2.4395 | LR: 4.63e-04 | 8237.3ms/step | 2.4 samples/s
+Step 7170/37465 | Loss: 2.3546 | LR: 4.63e-04 | 8177.9ms/step | 2.4 samples/s
+Step 7180/37465 | Loss: 2.5471 | LR: 4.63e-04 | 8185.2ms/step | 2.4 samples/s
+Step 7190/37465 | Loss: 2.6000 | LR: 4.63e-04 | 8225.6ms/step | 2.4 samples/s
+Step 7200/37465 | Loss: 2.3475 | LR: 4.63e-04 | 8132.2ms/step | 2.5 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_fixed/run_20260709_000503/checkpoint-7200
+Step 7210/37465 | Loss: 2.6240 | LR: 4.63e-04 | 9976.2ms/step | 2.0 samples/s
+Step 7220/37465 | Loss: 2.3894 | LR: 4.63e-04 | 8216.2ms/step | 2.4 samples/s
+Step 7230/37465 | Loss: 2.4416 | LR: 4.63e-04 | 8234.5ms/step | 2.4 samples/s
+Step 7240/37465 | Loss: 2.3953 | LR: 4.62e-04 | 8155.0ms/step | 2.5 samples/s
+Step 7250/37465 | Loss: 2.3746 | LR: 4.62e-04 | 8174.5ms/step | 2.4 samples/s
+Step 7260/37465 | Loss: 2.3994 | LR: 4.62e-04 | 8180.1ms/step | 2.4 samples/s
+Step 7270/37465 | Loss: 2.2348 | LR: 4.62e-04 | 8056.5ms/step | 2.5 samples/s
+Step 7280/37465 | Loss: 2.4378 | LR: 4.62e-04 | 8159.2ms/step | 2.5 samples/s
+Step 7290/37465 | Loss: 2.6233 | LR: 4.62e-04 | 8155.0ms/step | 2.5 samples/s
+Step 7300/37465 | Loss: 2.3456 | LR: 4.62e-04 | 8197.2ms/step | 2.4 samples/s
+Step 7310/37465 | Loss: 2.3139 | LR: 4.62e-04 | 8240.0ms/step | 2.4 samples/s
+Step 7320/37465 | Loss: 2.2889 | LR: 4.62e-04 | 8202.1ms/step | 2.4 samples/s
+Step 7330/37465 | Loss: 2.4655 | LR: 4.61e-04 | 8245.5ms/step | 2.4 samples/s
+Step 7340/37465 | Loss: 2.3131 | LR: 4.61e-04 | 8219.7ms/step | 2.4 samples/s
+Step 7350/37465 | Loss: 2.4320 | LR: 4.61e-04 | 8105.2ms/step | 2.5 samples/s
+Step 7360/37465 | Loss: 2.4561 | LR: 4.61e-04 | 8164.5ms/step | 2.4 samples/s
+Step 7370/37465 | Loss: 2.1081 | LR: 4.61e-04 | 8134.0ms/step | 2.5 samples/s
+Step 7380/37465 | Loss: 2.4054 | LR: 4.61e-04 | 8181.1ms/step | 2.4 samples/s
+Step 7390/37465 | Loss: 2.3375 | LR: 4.61e-04 | 8172.4ms/step | 2.4 samples/s
+Step 7400/37465 | Loss: 2.3454 | LR: 4.61e-04 | 8221.0ms/step | 2.4 samples/s
+Step 7410/37465 | Loss: 2.2249 | LR: 4.60e-04 | 8192.7ms/step | 2.4 samples/s
+Step 7420/37465 | Loss: 2.4995 | LR: 4.60e-04 | 8208.0ms/step | 2.4 samples/s
+Step 7430/37465 | Loss: 2.4542 | LR: 4.60e-04 | 8215.9ms/step | 2.4 samples/s
+Step 7440/37465 | Loss: 2.1935 | LR: 4.60e-04 | 8215.2ms/step | 2.4 samples/s
+Step 7450/37465 | Loss: 2.2351 | LR: 4.60e-04 | 8232.4ms/step | 2.4 samples/s
+Step 7460/37465 | Loss: 2.3569 | LR: 4.60e-04 | 8171.1ms/step | 2.4 samples/s
+Step 7470/37465 | Loss: 2.1997 | LR: 4.60e-04 | 8132.9ms/step | 2.5 samples/s
+Step 7480/37465 | Loss: 2.4366 | LR: 4.60e-04 | 8132.6ms/step | 2.5 samples/s
+Step 7490/37465 | Loss: 2.4346 | LR: 4.60e-04 | 8136.7ms/step | 2.5 samples/s
+Step 7500/37465 | Loss: 2.2755 | LR: 4.59e-04 | 8204.4ms/step | 2.4 samples/s
+Validation at step 7500: Loss = 2.3556, Perplexity = 10.54
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_fixed/run_20260709_000503/checkpoint-7500
+Step 7510/37465 | Loss: 2.3687 | LR: 4.59e-04 | 9753.3ms/step | 2.1 samples/s
+Step 7520/37465 | Loss: 2.4069 | LR: 4.59e-04 | 8148.1ms/step | 2.5 samples/s
+Step 7530/37465 | Loss: 2.4684 | LR: 4.59e-04 | 8205.8ms/step | 2.4 samples/s
+Step 7540/37465 | Loss: 2.2359 | LR: 4.59e-04 | 8167.7ms/step | 2.4 samples/s
+Step 7550/37465 | Loss: 2.4971 | LR: 4.59e-04 | 8180.7ms/step | 2.4 samples/s
+Step 7560/37465 | Loss: 2.3584 | LR: 4.59e-04 | 8203.0ms/step | 2.4 samples/s
+Step 7570/37465 | Loss: 2.3746 | LR: 4.59e-04 | 8166.4ms/step | 2.4 samples/s
+Step 7580/37465 | Loss: 2.5086 | LR: 4.59e-04 | 8169.6ms/step | 2.4 samples/s
+Step 7590/37465 | Loss: 2.3596 | LR: 4.58e-04 | 8188.5ms/step | 2.4 samples/s
+Step 7600/37465 | Loss: 2.5025 | LR: 4.58e-04 | 8177.6ms/step | 2.4 samples/s
+Step 7610/37465 | Loss: 2.4599 | LR: 4.58e-04 | 8200.8ms/step | 2.4 samples/s
+Step 7620/37465 | Loss: 2.1612 | LR: 4.58e-04 | 8218.8ms/step | 2.4 samples/s
+Step 7630/37465 | Loss: 2.3423 | LR: 4.58e-04 | 8224.4ms/step | 2.4 samples/s
+Step 7640/37465 | Loss: 2.2358 | LR: 4.58e-04 | 8214.2ms/step | 2.4 samples/s
+Step 7650/37465 | Loss: 2.2991 | LR: 4.58e-04 | 8186.7ms/step | 2.4 samples/s
+Step 7660/37465 | Loss: 2.3484 | LR: 4.58e-04 | 8181.6ms/step | 2.4 samples/s
+Step 7670/37465 | Loss: 2.2338 | LR: 4.57e-04 | 8178.5ms/step | 2.4 samples/s
+Step 7680/37465 | Loss: 2.3069 | LR: 4.57e-04 | 8180.8ms/step | 2.4 samples/s
+Step 7690/37465 | Loss: 2.1779 | LR: 4.57e-04 | 8201.0ms/step | 2.4 samples/s
+Step 7700/37465 | Loss: 2.2300 | LR: 4.57e-04 | 8185.0ms/step | 2.4 samples/s
+Step 7710/37465 | Loss: 2.6032 | LR: 4.57e-04 | 8193.9ms/step | 2.4 samples/s
+Step 7720/37465 | Loss: 2.1712 | LR: 4.57e-04 | 8201.8ms/step | 2.4 samples/s
+Step 7730/37465 | Loss: 2.3737 | LR: 4.57e-04 | 8178.0ms/step | 2.4 samples/s
+Step 7740/37465 | Loss: 2.4715 | LR: 4.57e-04 | 8200.8ms/step | 2.4 samples/s
+Step 7750/37465 | Loss: 2.4677 | LR: 4.56e-04 | 8207.4ms/step | 2.4 samples/s
+Step 7760/37465 | Loss: 2.3958 | LR: 4.56e-04 | 8205.5ms/step | 2.4 samples/s
+Step 7770/37465 | Loss: 2.1620 | LR: 4.56e-04 | 8206.5ms/step | 2.4 samples/s
+Step 7780/37465 | Loss: 2.1917 | LR: 4.56e-04 | 8202.3ms/step | 2.4 samples/s
+Step 7790/37465 | Loss: 2.1820 | LR: 4.56e-04 | 8177.7ms/step | 2.4 samples/s
+Step 7800/37465 | Loss: 2.1496 | LR: 4.56e-04 | 8176.2ms/step | 2.4 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_fixed/run_20260709_000503/checkpoint-7800
+Step 7810/37465 | Loss: 2.3445 | LR: 4.56e-04 | 9051.5ms/step | 2.2 samples/s
+Step 7820/37465 | Loss: 2.3243 | LR: 4.56e-04 | 8191.0ms/step | 2.4 samples/s
+Step 7830/37465 | Loss: 2.2762 | LR: 4.56e-04 | 8182.2ms/step | 2.4 samples/s
+Step 7840/37465 | Loss: 2.4698 | LR: 4.55e-04 | 8157.2ms/step | 2.5 samples/s
+Step 7850/37465 | Loss: 2.2608 | LR: 4.55e-04 | 8187.7ms/step | 2.4 samples/s
+Step 7860/37465 | Loss: 2.2268 | LR: 4.55e-04 | 8196.4ms/step | 2.4 samples/s
+Step 7870/37465 | Loss: 2.3566 | LR: 4.55e-04 | 8210.0ms/step | 2.4 samples/s
+Step 7880/37465 | Loss: 2.2290 | LR: 4.55e-04 | 8207.9ms/step | 2.4 samples/s
+Step 7890/37465 | Loss: 2.4396 | LR: 4.55e-04 | 8201.9ms/step | 2.4 samples/s
+Step 7900/37465 | Loss: 2.2428 | LR: 4.55e-04 | 8206.0ms/step | 2.4 samples/s
+Step 7910/37465 | Loss: 2.2175 | LR: 4.55e-04 | 8156.6ms/step | 2.5 samples/s
+Step 7920/37465 | Loss: 2.1634 | LR: 4.54e-04 | 8014.5ms/step | 2.5 samples/s
+Step 7930/37465 | Loss: 2.2800 | LR: 4.54e-04 | 8081.7ms/step | 2.5 samples/s
+Step 7940/37465 | Loss: 2.1433 | LR: 4.54e-04 | 8093.1ms/step | 2.5 samples/s
+Step 7950/37465 | Loss: 2.3007 | LR: 4.54e-04 | 8103.2ms/step | 2.5 samples/s
+Step 7960/37465 | Loss: 2.4438 | LR: 4.54e-04 | 8096.1ms/step | 2.5 samples/s
+Step 7970/37465 | Loss: 2.3276 | LR: 4.54e-04 | 8083.8ms/step | 2.5 samples/s
+Step 7980/37465 | Loss: 2.3070 | LR: 4.54e-04 | 8108.2ms/step | 2.5 samples/s
+Step 7990/37465 | Loss: 2.4770 | LR: 4.54e-04 | 8144.9ms/step | 2.5 samples/s
+Step 8000/37465 | Loss: 2.3272 | LR: 4.53e-04 | 8162.2ms/step | 2.5 samples/s
+Validation at step 8000: Loss = 2.3388, Perplexity = 10.37
+Step 8010/37465 | Loss: 2.0893 | LR: 4.53e-04 | 8718.1ms/step | 2.3 samples/s
+Step 8020/37465 | Loss: 2.3537 | LR: 4.53e-04 | 8087.2ms/step | 2.5 samples/s
+Step 8030/37465 | Loss: 2.3397 | LR: 4.53e-04 | 8111.2ms/step | 2.5 samples/s
+Step 8040/37465 | Loss: 2.3065 | LR: 4.53e-04 | 8110.2ms/step | 2.5 samples/s
+Step 8050/37465 | Loss: 2.1418 | LR: 4.53e-04 | 8139.4ms/step | 2.5 samples/s
+Step 8060/37465 | Loss: 2.4284 | LR: 4.53e-04 | 8095.5ms/step | 2.5 samples/s
+Step 8070/37465 | Loss: 2.4237 | LR: 4.53e-04 | 8058.4ms/step | 2.5 samples/s
+Step 8080/37465 | Loss: 2.3563 | LR: 4.52e-04 | 8163.7ms/step | 2.4 samples/s
+Step 8090/37465 | Loss: 2.2095 | LR: 4.52e-04 | 8144.3ms/step | 2.5 samples/s
+Step 8100/37465 | Loss: 2.4239 | LR: 4.52e-04 | 8176.8ms/step | 2.4 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_fixed/run_20260709_000503/checkpoint-8100
+Step 8110/37465 | Loss: 2.4102 | LR: 4.52e-04 | 9041.5ms/step | 2.2 samples/s
+Step 8120/37465 | Loss: 2.4006 | LR: 4.52e-04 | 8162.7ms/step | 2.5 samples/s
+Step 8130/37465 | Loss: 2.3514 | LR: 4.52e-04 | 8043.4ms/step | 2.5 samples/s
+Step 8140/37465 | Loss: 2.2464 | LR: 4.52e-04 | 8121.2ms/step | 2.5 samples/s
+Step 8150/37465 | Loss: 2.4491 | LR: 4.52e-04 | 8135.2ms/step | 2.5 samples/s
+Step 8160/37465 | Loss: 2.2588 | LR: 4.51e-04 | 8196.7ms/step | 2.4 samples/s
+Step 8170/37465 | Loss: 2.3349 | LR: 4.51e-04 | 8081.3ms/step | 2.5 samples/s
+Step 8180/37465 | Loss: 2.4152 | LR: 4.51e-04 | 7996.8ms/step | 2.5 samples/s
+Step 8190/37465 | Loss: 2.2468 | LR: 4.51e-04 | 8094.2ms/step | 2.5 samples/s
+Step 8200/37465 | Loss: 2.3353 | LR: 4.51e-04 | 8123.1ms/step | 2.5 samples/s
+Step 8210/37465 | Loss: 2.4674 | LR: 4.51e-04 | 8035.2ms/step | 2.5 samples/s
+Step 8220/37465 | Loss: 2.3570 | LR: 4.51e-04 | 8156.5ms/step | 2.5 samples/s
+Step 8230/37465 | Loss: 2.2833 | LR: 4.51e-04 | 8231.9ms/step | 2.4 samples/s
+Step 8240/37465 | Loss: 2.2904 | LR: 4.50e-04 | 8042.1ms/step | 2.5 samples/s
+Step 8250/37465 | Loss: 2.2846 | LR: 4.50e-04 | 8132.7ms/step | 2.5 samples/s
+Step 8260/37465 | Loss: 2.5686 | LR: 4.50e-04 | 8096.2ms/step | 2.5 samples/s
+Step 8270/37465 | Loss: 2.4240 | LR: 4.50e-04 | 8086.1ms/step | 2.5 samples/s
+Step 8280/37465 | Loss: 2.1655 | LR: 4.50e-04 | 8169.1ms/step | 2.4 samples/s
+Step 8290/37465 | Loss: 2.4233 | LR: 4.50e-04 | 8231.7ms/step | 2.4 samples/s
+Step 8300/37465 | Loss: 2.2389 | LR: 4.50e-04 | 8191.6ms/step | 2.4 samples/s
+Step 8310/37465 | Loss: 2.4158 | LR: 4.49e-04 | 8181.8ms/step | 2.4 samples/s
+Step 8320/37465 | Loss: 2.4746 | LR: 4.49e-04 | 8235.6ms/step | 2.4 samples/s
+Step 8330/37465 | Loss: 2.1171 | LR: 4.49e-04 | 8195.8ms/step | 2.4 samples/s
+Step 8340/37465 | Loss: 2.4831 | LR: 4.49e-04 | 8188.3ms/step | 2.4 samples/s
+Step 8350/37465 | Loss: 2.4076 | LR: 4.49e-04 | 8185.2ms/step | 2.4 samples/s
+Step 8360/37465 | Loss: 2.0461 | LR: 4.49e-04 | 8230.2ms/step | 2.4 samples/s
+Step 8370/37465 | Loss: 2.3678 | LR: 4.49e-04 | 8169.6ms/step | 2.4 samples/s

epfl_final/train_logs/resume_resampled30k.log

@@ -0,0 +1,209 @@
+Auto-resuming from latest checkpoint: /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_resampled/run_20260708_112209/checkpoint-5700
+Resuming: Using existing output directory /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_resampled/run_20260708_112209
+Running config:
+{
+  "model_path": "/mloscratch/homes/navasard/qwen3_06b_base",
+  "model_dtype": "bfloat16",
+  "init_from_scratch": true,
+  "train_path": "/mloscratch/homes/navasard/data/longctx30k/tokenized",
+  "val_path": null,
+  "max_seq_length": 30720,
+  "num_workers": 4,
+  "output_dir": "/mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_resampled/run_20260708_112209",
+  "resume_from_checkpoint": "/mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_resampled/run_20260708_112209/checkpoint-5700",
+  "force_new_run": false,
+  "num_epochs": 1,
+  "max_steps": null,
+  "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": 300,
+  "eval_steps": 500,
+  "eval_batches": 20,
+  "logging_steps": 10,
+  "seed": 42,
+  "verbose": true,
+  "checkpoint_activations": true,
+  "strategy": "ddp",
+  "mixed_precision": "bf16",
+  "wandb_enabled": true,
+  "wandb_project": "stochastic_sparsity",
+  "wandb_name": "1.2b resampled 30k longctx",
+  "wandb_tags": [
+    "qwen3",
+    "1.2b",
+    "resampled",
+    "30k",
+    "longctx",
+    "b128",
+    "scaleladder"
+  ],
+  "hidden_size": 1536,
+  "intermediate_size": 6144,
+  "num_attention_heads": 12,
+  "num_key_value_heads": 4,
+  "num_hidden_layers": null,
+  "head_dim": null,
+  "attention_type": "sparse_stochastic",
+  "sparse_block_size": 128,
+  "sparse_num_random_blocks": 120,
+  "sparse_mask_pool_size": 1000,
+  "sparse_val_dense_compare": false,
+  "sparse_layer_indices": null,
+  "sparse_mask_cache_dir": "/mloscratch/homes/navasard/sparse/mask_cache_30k_b128_nrb120",
+  "sparse_mask_seed": 0,
+  "dense_mode": false,
+  "dense_step_intervals": null
+}
+Loading tokenizer from /mloscratch/homes/navasard/qwen3_06b_base...
+Checkpoint activations: TrueCheckpoint activations: True
+
+Resuming from checkpoint: /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_resampled/run_20260708_112209/checkpoint-5700
+Checkpoint activations: True
+Checkpoint activations: True
+
+
+================================================================================
+MODEL ARCHITECTURE
+================================================================================
+Model config:
+  vocab_size: 151936
+  hidden_size: 1536
+  num_hidden_layers: 28
+  num_attention_heads: 12
+  num_key_value_heads: 4
+  head_dim: 128
+  intermediate_size: 6144
+  max_position_embeddings: 32768
+  rope_theta: 1000000
+  attention_type: sparse_stochastic
+  sparse_block_size: 128
+  sparse_num_random_blocks: 120
+  sparse_mask_pool_size: 1000
+  sparse_max_seq_length: 30720
+
+Parameter count:
+  Total parameters: 1,202,352,640
+  Trainable parameters: 1,202,352,640
+  Non-trainable parameters: 0
+  Model size (BF16): 2.40 GB
+
+--------------------------------------------------------------------------------
+LAYER-BY-LAYER SUMMARY
+--------------------------------------------------------------------------------
+/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/nn/attention/flex_attention.py:1704: UserWarning: flex_attention called without torch.compile() - this will use an unfused implementation that materializes the full scores matrix instead of generating a fused kernel.
+
+SOLUTION: Use torch.compile(flex_attention)(...)
+
+If you want to debug your score_mod/mask_mod, you can set:
+torch.nn.attention.flex_attention._FLEX_ATTENTION_DISABLE_COMPILE_DEBUG = True
+
+This will allow you to use print statements or breakpoints. Note: This doesn't work with the backwards pass and may produce incorrect results.
+  _warn_once(
+Could not generate detailed summary: Failed to run torchinfo. See above stack traces for more details. Executed layers up to: [Embedding: 2, Qwen3RMSNorm: 4, Linear: 5, Linear: 5, Linear: 5, Qwen3RMSNorm: 5, Qwen3RMSNorm: 5, Qwen3RotaryEmbedding: 5]
+  NOTE: torchinfo double-counts tied embed_tokens/lm_head weight (233,373,696 params). True unique parameter count: 1,202,352,640 (1202.4M)
+================================================================================
+
+Loading dataset from /mloscratch/homes/navasard/data/longctx30k/tokenized...
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/longctx30k/tokenized...
+Using pre-tokenized dataset (no runtime tokenization)
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/longctx30k/tokenized...
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/longctx30k/tokenized...
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/longctx30k/tokenized...
+Loaded 788736 pre-tokenized examples (stored_len=30720, pack_factor=1, packed_examples=788736)
+Loaded 788736 pre-tokenized examples (stored_len=30720, pack_factor=1, packed_examples=788736)
+Loaded 788736 pre-tokenized examples (stored_len=30720, pack_factor=1, packed_examples=788736)
+Loaded 788736 pre-tokenized examples (stored_len=30720, pack_factor=1, packed_examples=788736)
+Dataset split: 749299 train, 39437 validation
+Training examples (packed, train split):  749299
+Validation examples (packed, val split):  39437
+Per-device batch size:                    1
+World size (GPUs):                        4
+Gradient accumulation steps:              5
+Effective batch size (seqs/step):         20
+Sequence length:                          30720
+Tokens per optimizer step:                614,400
+Micro-batches per epoch (per rank):       187325
+Optimizer steps per epoch:                37465
+Total optimizer steps:                    37465
+Loading checkpoint from /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_resampled/run_20260708_112209/checkpoint-5700...
+Overriding checkpoint LR with config LR: 0.0005
+Resumed from checkpoint at step 5700 (epoch 0, micro_step 28500)
+Resuming training from step 5700
+Resuming wandb run hn868f80
+wandb: [wandb.login()] Loaded credentials for https://api.wandb.ai from WANDB_API_KEY.
+wandb: Currently logged in as: ruben-navasardyan4869 (ruben-navasardyan4869-epfl) to https://api.wandb.ai. Use `wandb login --relogin` to force relogin
+wandb: Tracking run with wandb version 0.26.1
+wandb: Run data is saved locally in /mloscratch/homes/navasard/sparse/SparseAtt/wandb/run-20260709_222958-hn868f80
+wandb: Run `wandb offline` to turn off syncing.
+wandb: Resuming run 1.2b resampled 30k longctx
+wandb: ⭐️ View project at https://wandb.ai/ruben-navasardyan4869-epfl/stochastic_sparsity
+wandb: 🚀 View run at https://wandb.ai/ruben-navasardyan4869-epfl/stochastic_sparsity/runs/hn868f80
+Starting training...
+Validation at step 5700: Loss = 2.4275, Perplexity = 11.33
+wandb: WARNING Tried to log to step 5700 that is less than the current step 5761. Steps must be monotonically increasing, so this data will be ignored. See https://wandb.me/define-metric to log data out of order.
+Switching attention mode for step 5701: sparse_stochastic (configured sparse attention)
+Step 5710/37465 | Loss: 2.4246 | LR: 4.78e-04 | 49011.5ms/step | 0.4 samples/s
+wandb: WARNING Tried to log to step 5710 that is less than the current step 5761. Steps must be monotonically increasing, so this data will be ignored. See https://wandb.me/define-metric to log data out of order.
+Step 5720/37465 | Loss: 2.3987 | LR: 4.78e-04 | 28273.4ms/step | 0.7 samples/s
+wandb: WARNING Tried to log to step 5720 that is less than the current step 5761. Steps must be monotonically increasing, so this data will be ignored. See https://wandb.me/define-metric to log data out of order.
+Step 5730/37465 | Loss: 2.4191 | LR: 4.78e-04 | 28275.3ms/step | 0.7 samples/s
+wandb: WARNING Tried to log to step 5730 that is less than the current step 5761. Steps must be monotonically increasing, so this data will be ignored. See https://wandb.me/define-metric to log data out of order.
+Step 5740/37465 | Loss: 2.3961 | LR: 4.78e-04 | 28271.8ms/step | 0.7 samples/s
+wandb: WARNING Tried to log to step 5740 that is less than the current step 5761. Steps must be monotonically increasing, so this data will be ignored. See https://wandb.me/define-metric to log data out of order.
+Step 5750/37465 | Loss: 2.4381 | LR: 4.77e-04 | 28275.6ms/step | 0.7 samples/s
+wandb: WARNING Tried to log to step 5750 that is less than the current step 5761. Steps must be monotonically increasing, so this data will be ignored. See https://wandb.me/define-metric to log data out of order.
+Step 5760/37465 | Loss: 2.1498 | LR: 4.77e-04 | 28277.9ms/step | 0.7 samples/s
+wandb: WARNING Tried to log to step 5760 that is less than the current step 5761. Steps must be monotonically increasing, so this data will be ignored. See https://wandb.me/define-metric to log data out of order.
+Step 5770/37465 | Loss: 2.4837 | LR: 4.77e-04 | 28273.9ms/step | 0.7 samples/s
+Step 5780/37465 | Loss: 2.4049 | LR: 4.77e-04 | 28276.1ms/step | 0.7 samples/s
+Step 5790/37465 | Loss: 2.2924 | LR: 4.77e-04 | 28272.8ms/step | 0.7 samples/s
+Step 5800/37465 | Loss: 2.4950 | LR: 4.77e-04 | 28272.1ms/step | 0.7 samples/s
+Step 5810/37465 | Loss: 2.4212 | LR: 4.77e-04 | 28284.5ms/step | 0.7 samples/s
+Step 5820/37465 | Loss: 2.1602 | LR: 4.77e-04 | 28279.2ms/step | 0.7 samples/s
+Step 5830/37465 | Loss: 2.3799 | LR: 4.77e-04 | 28282.1ms/step | 0.7 samples/s
+Step 5840/37465 | Loss: 2.4057 | LR: 4.77e-04 | 28278.5ms/step | 0.7 samples/s
+Step 5850/37465 | Loss: 2.2287 | LR: 4.77e-04 | 28280.3ms/step | 0.7 samples/s
+Step 5860/37465 | Loss: 2.1535 | LR: 4.76e-04 | 28277.0ms/step | 0.7 samples/s
+Step 5870/37465 | Loss: 2.4309 | LR: 4.76e-04 | 28279.3ms/step | 0.7 samples/s
+W0709 23:58:10.171000 142 site-packages/torch/distributed/elastic/agent/server/api.py:739] Received 15 death signal, shutting down workers
+W0709 23:58:10.173000 142 site-packages/torch/distributed/elastic/multiprocessing/api.py:1012] Sending process 155 closing signal SIGTERM
+W0709 23:58:10.174000 142 site-packages/torch/distributed/elastic/multiprocessing/api.py:1012] Sending process 156 closing signal SIGTERM
+W0709 23:58:10.175000 142 site-packages/torch/distributed/elastic/multiprocessing/api.py:1012] Sending process 157 closing signal SIGTERM
+W0709 23:58:10.176000 142 site-packages/torch/distributed/elastic/multiprocessing/api.py:1012] Sending process 158 closing signal SIGTERM
+Traceback (most recent call last):
+  File "/mloscratch/homes/navasard/conda/bin/torchrun", line 6, in <module>
+    sys.exit(main())
+             ~~~~^^
+  File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 362, in wrapper
+    return f(*args, **kwargs)
+  File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/distributed/run.py", line 990, in main
+    run(args)
+    ~~~^^^^^^
+  File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/distributed/run.py", line 981, in run
+    elastic_launch(
+    ~~~~~~~~~~~~~~~
+        config=config,
+        ~~~~~~~~~~~~~~
+        entrypoint=cmd,
+        ~~~~~~~~~~~~~~~
+    )(*cmd_args)
+    ~^^^^^^^^^^^
+  File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/distributed/launcher/api.py", line 170, in __call__
+    return launch_agent(self._config, self._entrypoint, list(args))
+  File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/distributed/launcher/api.py", line 308, in launch_agent
+    result = agent.run()
+  File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/distributed/elastic/metrics/api.py", line 134, in wrapper
+    result = f(*args, **kwargs)
+  File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/distributed/elastic/agent/server/api.py", line 731, in run
+    result = self._invoke_run(role)
+  File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/distributed/elastic/agent/server/api.py", line 908, in _invoke_run
+    time.sleep(monitor_interval)
+    ~~~~~~~~~~^^^^^^^^^^^^^^^^^^
+  File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/distributed/elastic/multiprocessing/api.py", line 86, in _terminate_process_handler
+    raise SignalException(f"Process {os.getpid()} got signal: {sigval}", sigval=sigval)
+torch.distributed.elastic.multiprocessing.api.SignalException: Process 142 got signal: 15

epfl_final/train_logs/resume_resampled30k_minenv.log

@@ -0,0 +1,550 @@
+W0709 23:59:19.429000 3736 site-packages/torch/distributed/run.py:851] 
+W0709 23:59:19.429000 3736 site-packages/torch/distributed/run.py:851] *****************************************
+W0709 23:59:19.429000 3736 site-packages/torch/distributed/run.py:851] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. 
+W0709 23:59:19.429000 3736 site-packages/torch/distributed/run.py:851] *****************************************
+Auto-resuming from latest checkpoint: /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_resampled/run_20260708_112209/checkpoint-5700
+Resuming: Using existing output directory /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_resampled/run_20260708_112209
+Running config:
+{
+  "model_path": "/mloscratch/homes/navasard/qwen3_06b_base",
+  "model_dtype": "bfloat16",
+  "init_from_scratch": true,
+  "train_path": "/mloscratch/homes/navasard/data/longctx30k/tokenized",
+  "val_path": null,
+  "max_seq_length": 30720,
+  "num_workers": 4,
+  "output_dir": "/mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_resampled/run_20260708_112209",
+  "resume_from_checkpoint": "/mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_resampled/run_20260708_112209/checkpoint-5700",
+  "force_new_run": false,
+  "num_epochs": 1,
+  "max_steps": null,
+  "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": 300,
+  "eval_steps": 500,
+  "eval_batches": 20,
+  "logging_steps": 10,
+  "seed": 42,
+  "verbose": true,
+  "checkpoint_activations": true,
+  "strategy": "ddp",
+  "mixed_precision": "bf16",
+  "wandb_enabled": true,
+  "wandb_project": "stochastic_sparsity",
+  "wandb_name": "1.2b resampled 30k longctx",
+  "wandb_tags": [
+    "qwen3",
+    "1.2b",
+    "resampled",
+    "30k",
+    "longctx",
+    "b128",
+    "scaleladder"
+  ],
+  "hidden_size": 1536,
+  "intermediate_size": 6144,
+  "num_attention_heads": 12,
+  "num_key_value_heads": 4,
+  "num_hidden_layers": null,
+  "head_dim": null,
+  "attention_type": "sparse_stochastic",
+  "sparse_block_size": 128,
+  "sparse_num_random_blocks": 120,
+  "sparse_mask_pool_size": 1000,
+  "sparse_val_dense_compare": false,
+  "sparse_layer_indices": null,
+  "sparse_mask_cache_dir": "/mloscratch/homes/navasard/sparse/mask_cache_30k_b128_nrb120",
+  "sparse_mask_seed": 0,
+  "dense_mode": false,
+  "dense_step_intervals": null
+}
+Loading tokenizer from /mloscratch/homes/navasard/qwen3_06b_base...
+Checkpoint activations: True
+Resuming from checkpoint: /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_resampled/run_20260708_112209/checkpoint-5700
+Checkpoint activations: True
+Checkpoint activations: True
+Checkpoint activations: True
+
+
+================================================================================
+MODEL ARCHITECTURE
+================================================================================
+Model config:
+  vocab_size: 151936
+  hidden_size: 1536
+  num_hidden_layers: 28
+  num_attention_heads: 12
+  num_key_value_heads: 4
+  head_dim: 128
+  intermediate_size: 6144
+  max_position_embeddings: 32768
+  rope_theta: 1000000
+  attention_type: sparse_stochastic
+  sparse_block_size: 128
+  sparse_num_random_blocks: 120
+  sparse_mask_pool_size: 1000
+  sparse_max_seq_length: 30720
+
+Parameter count:
+  Total parameters: 1,202,352,640
+  Trainable parameters: 1,202,352,640
+  Non-trainable parameters: 0
+  Model size (BF16): 2.40 GB
+
+--------------------------------------------------------------------------------
+LAYER-BY-LAYER SUMMARY
+--------------------------------------------------------------------------------
+/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/nn/attention/flex_attention.py:1704: UserWarning: flex_attention called without torch.compile() - this will use an unfused implementation that materializes the full scores matrix instead of generating a fused kernel.
+
+SOLUTION: Use torch.compile(flex_attention)(...)
+
+If you want to debug your score_mod/mask_mod, you can set:
+torch.nn.attention.flex_attention._FLEX_ATTENTION_DISABLE_COMPILE_DEBUG = True
+
+This will allow you to use print statements or breakpoints. Note: This doesn't work with the backwards pass and may produce incorrect results.
+  _warn_once(
+Could not generate detailed summary: Failed to run torchinfo. See above stack traces for more details. Executed layers up to: [Embedding: 2, Qwen3RMSNorm: 4, Linear: 5, Linear: 5, Linear: 5, Qwen3RMSNorm: 5, Qwen3RMSNorm: 5, Qwen3RotaryEmbedding: 5]
+  NOTE: torchinfo double-counts tied embed_tokens/lm_head weight (233,373,696 params). True unique parameter count: 1,202,352,640 (1202.4M)
+================================================================================
+
+Loading dataset from /mloscratch/homes/navasard/data/longctx30k/tokenized...
+Using pre-tokenized dataset (no runtime tokenization)
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/longctx30k/tokenized...
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/longctx30k/tokenized...
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/longctx30k/tokenized...
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/longctx30k/tokenized...
+Loaded 788736 pre-tokenized examples (stored_len=30720, pack_factor=1, packed_examples=788736)
+Loaded 788736 pre-tokenized examples (stored_len=30720, pack_factor=1, packed_examples=788736)
+Dataset split: 749299 train, 39437 validation
+Training examples (packed, train split):  749299
+Validation examples (packed, val split):  39437
+Per-device batch size:                    1
+World size (GPUs):                        4
+Gradient accumulation steps:              5
+Effective batch size (seqs/step):         20
+Sequence length:                          30720
+Tokens per optimizer step:                614,400
+Micro-batches per epoch (per rank):       187325
+Optimizer steps per epoch:                37465
+Total optimizer steps:                    37465
+Loading checkpoint from /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_resampled/run_20260708_112209/checkpoint-5700...
+Loaded 788736 pre-tokenized examples (stored_len=30720, pack_factor=1, packed_examples=788736)
+Loaded 788736 pre-tokenized examples (stored_len=30720, pack_factor=1, packed_examples=788736)
+Overriding checkpoint LR with config LR: 0.0005
+Resumed from checkpoint at step 5700 (epoch 0, micro_step 28500)
+Resuming training from step 5700
+Resuming wandb run hn868f80
+wandb: [wandb.login()] Loaded credentials for https://api.wandb.ai from WANDB_API_KEY.
+wandb: Currently logged in as: ruben-navasardyan4869 (ruben-navasardyan4869-epfl) to https://api.wandb.ai. Use `wandb login --relogin` to force relogin
+wandb: setting up run hn868f80
+wandb: Tracking run with wandb version 0.26.1
+wandb: Run data is saved locally in /mloscratch/homes/navasard/sparse/SparseAtt/wandb/run-20260710_000026-hn868f80
+wandb: Run `wandb offline` to turn off syncing.
+wandb: Resuming run 1.2b resampled 30k longctx
+wandb: ⭐️ View project at https://wandb.ai/ruben-navasardyan4869-epfl/stochastic_sparsity
+wandb: 🚀 View run at https://wandb.ai/ruben-navasardyan4869-epfl/stochastic_sparsity/runs/hn868f80
+Starting training...
+Validation at step 5700: Loss = 2.4275, Perplexity = 11.33
+wandb: WARNING Tried to log to step 5700 that is less than the current step 5861. Steps must be monotonically increasing, so this data will be ignored. See https://wandb.me/define-metric to log data out of order.
+Switching attention mode for step 5701: sparse_stochastic (configured sparse attention)
+Step 5710/37465 | Loss: 2.4245 | LR: 4.78e-04 | 41956.3ms/step | 0.5 samples/s
+wandb: WARNING Tried to log to step 5710 that is less than the current step 5861. Steps must be monotonically increasing, so this data will be ignored. See https://wandb.me/define-metric to log data out of order.
+Step 5720/37465 | Loss: 2.3987 | LR: 4.78e-04 | 28271.7ms/step | 0.7 samples/s
+wandb: WARNING Tried to log to step 5720 that is less than the current step 5861. Steps must be monotonically increasing, so this data will be ignored. See https://wandb.me/define-metric to log data out of order.
+Step 5730/37465 | Loss: 2.4191 | LR: 4.78e-04 | 28278.9ms/step | 0.7 samples/s
+wandb: WARNING Tried to log to step 5730 that is less than the current step 5861. Steps must be monotonically increasing, so this data will be ignored. See https://wandb.me/define-metric to log data out of order.
+Step 5740/37465 | Loss: 2.3961 | LR: 4.78e-04 | 28284.0ms/step | 0.7 samples/s
+wandb: WARNING Tried to log to step 5740 that is less than the current step 5861. Steps must be monotonically increasing, so this data will be ignored. See https://wandb.me/define-metric to log data out of order.
+Step 5750/37465 | Loss: 2.4381 | LR: 4.77e-04 | 28280.6ms/step | 0.7 samples/s
+wandb: WARNING Tried to log to step 5750 that is less than the current step 5861. Steps must be monotonically increasing, so this data will be ignored. See https://wandb.me/define-metric to log data out of order.
+Step 5760/37465 | Loss: 2.1498 | LR: 4.77e-04 | 28270.3ms/step | 0.7 samples/s
+wandb: WARNING Tried to log to step 5760 that is less than the current step 5861. Steps must be monotonically increasing, so this data will be ignored. See https://wandb.me/define-metric to log data out of order.
+Step 5770/37465 | Loss: 2.4837 | LR: 4.77e-04 | 28270.5ms/step | 0.7 samples/s
+wandb: WARNING Tried to log to step 5770 that is less than the current step 5861. Steps must be monotonically increasing, so this data will be ignored. See https://wandb.me/define-metric to log data out of order.
+Step 5780/37465 | Loss: 2.4049 | LR: 4.77e-04 | 28273.6ms/step | 0.7 samples/s
+wandb: WARNING Tried to log to step 5780 that is less than the current step 5861. Steps must be monotonically increasing, so this data will be ignored. See https://wandb.me/define-metric to log data out of order.
+Step 5790/37465 | Loss: 2.2924 | LR: 4.77e-04 | 28272.8ms/step | 0.7 samples/s
+wandb: WARNING Tried to log to step 5790 that is less than the current step 5861. Steps must be monotonically increasing, so this data will be ignored. See https://wandb.me/define-metric to log data out of order.
+Step 5800/37465 | Loss: 2.4950 | LR: 4.77e-04 | 28308.7ms/step | 0.7 samples/s
+wandb: WARNING Tried to log to step 5800 that is less than the current step 5861. Steps must be monotonically increasing, so this data will be ignored. See https://wandb.me/define-metric to log data out of order.
+Step 5810/37465 | Loss: 2.4212 | LR: 4.77e-04 | 28273.6ms/step | 0.7 samples/s
+wandb: WARNING Tried to log to step 5810 that is less than the current step 5861. Steps must be monotonically increasing, so this data will be ignored. See https://wandb.me/define-metric to log data out of order.
+Step 5820/37465 | Loss: 2.1602 | LR: 4.77e-04 | 28277.4ms/step | 0.7 samples/s
+wandb: WARNING Tried to log to step 5820 that is less than the current step 5861. Steps must be monotonically increasing, so this data will be ignored. See https://wandb.me/define-metric to log data out of order.
+Step 5830/37465 | Loss: 2.3799 | LR: 4.77e-04 | 28273.6ms/step | 0.7 samples/s
+wandb: WARNING Tried to log to step 5830 that is less than the current step 5861. Steps must be monotonically increasing, so this data will be ignored. See https://wandb.me/define-metric to log data out of order.
+Step 5840/37465 | Loss: 2.4058 | LR: 4.77e-04 | 28266.4ms/step | 0.7 samples/s
+wandb: WARNING Tried to log to step 5840 that is less than the current step 5861. Steps must be monotonically increasing, so this data will be ignored. See https://wandb.me/define-metric to log data out of order.
+Step 5850/37465 | Loss: 2.2288 | LR: 4.77e-04 | 28265.1ms/step | 0.7 samples/s
+wandb: WARNING Tried to log to step 5850 that is less than the current step 5861. Steps must be monotonically increasing, so this data will be ignored. See https://wandb.me/define-metric to log data out of order.
+Step 5860/37465 | Loss: 2.1536 | LR: 4.76e-04 | 28274.3ms/step | 0.7 samples/s
+wandb: WARNING Tried to log to step 5860 that is less than the current step 5861. Steps must be monotonically increasing, so this data will be ignored. See https://wandb.me/define-metric to log data out of order.
+Step 5870/37465 | Loss: 2.4310 | LR: 4.76e-04 | 28269.1ms/step | 0.7 samples/s
+Step 5880/37465 | Loss: 2.4376 | LR: 4.76e-04 | 28267.1ms/step | 0.7 samples/s
+Step 5890/37465 | Loss: 2.5116 | LR: 4.76e-04 | 28270.0ms/step | 0.7 samples/s
+Step 5900/37465 | Loss: 2.5715 | LR: 4.76e-04 | 28277.3ms/step | 0.7 samples/s
+Step 5910/37465 | Loss: 2.2715 | LR: 4.76e-04 | 28266.2ms/step | 0.7 samples/s
+Step 5920/37465 | Loss: 2.4137 | LR: 4.76e-04 | 28263.8ms/step | 0.7 samples/s
+Step 5930/37465 | Loss: 2.3639 | LR: 4.76e-04 | 28262.1ms/step | 0.7 samples/s
+Step 5940/37465 | Loss: 2.6373 | LR: 4.76e-04 | 28261.6ms/step | 0.7 samples/s
+Step 5950/37465 | Loss: 2.5219 | LR: 4.76e-04 | 28266.7ms/step | 0.7 samples/s
+Step 5960/37465 | Loss: 2.2818 | LR: 4.76e-04 | 28263.0ms/step | 0.7 samples/s
+Step 5970/37465 | Loss: 2.2845 | LR: 4.75e-04 | 28261.8ms/step | 0.7 samples/s
+Step 5980/37465 | Loss: 2.3351 | LR: 4.75e-04 | 28263.7ms/step | 0.7 samples/s
+Step 5990/37465 | Loss: 2.3731 | LR: 4.75e-04 | 28258.5ms/step | 0.7 samples/s
+Step 6000/37465 | Loss: 2.2704 | LR: 4.75e-04 | 28259.4ms/step | 0.7 samples/s
+Validation at step 6000: Loss = 2.4159, Perplexity = 11.20
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_resampled/run_20260708_112209/checkpoint-6000
+Step 6010/37465 | Loss: 2.2741 | LR: 4.75e-04 | 30800.3ms/step | 0.6 samples/s
+Step 6020/37465 | Loss: 2.5314 | LR: 4.75e-04 | 28268.3ms/step | 0.7 samples/s
+Step 6030/37465 | Loss: 2.4936 | LR: 4.75e-04 | 28257.7ms/step | 0.7 samples/s
+Step 6040/37465 | Loss: 2.3298 | LR: 4.75e-04 | 28253.7ms/step | 0.7 samples/s
+Step 6050/37465 | Loss: 2.2415 | LR: 4.75e-04 | 28257.3ms/step | 0.7 samples/s
+Step 6060/37465 | Loss: 2.5859 | LR: 4.75e-04 | 28261.4ms/step | 0.7 samples/s
+Step 6070/37465 | Loss: 2.3417 | LR: 4.75e-04 | 28253.8ms/step | 0.7 samples/s
+Step 6080/37465 | Loss: 2.5655 | LR: 4.74e-04 | 28255.7ms/step | 0.7 samples/s
+Step 6090/37465 | Loss: 2.3407 | LR: 4.74e-04 | 28247.4ms/step | 0.7 samples/s
+Step 6100/37465 | Loss: 2.1764 | LR: 4.74e-04 | 28285.3ms/step | 0.7 samples/s
+Step 6110/37465 | Loss: 2.5956 | LR: 4.74e-04 | 28252.5ms/step | 0.7 samples/s
+Step 6120/37465 | Loss: 2.5236 | LR: 4.74e-04 | 28252.8ms/step | 0.7 samples/s
+Step 6130/37465 | Loss: 2.4235 | LR: 4.74e-04 | 28254.4ms/step | 0.7 samples/s
+Step 6140/37465 | Loss: 2.3519 | LR: 4.74e-04 | 28254.3ms/step | 0.7 samples/s
+Step 6150/37465 | Loss: 2.5209 | LR: 4.74e-04 | 28248.1ms/step | 0.7 samples/s
+Step 6160/37465 | Loss: 2.4799 | LR: 4.74e-04 | 28251.5ms/step | 0.7 samples/s
+Step 6170/37465 | Loss: 2.3391 | LR: 4.74e-04 | 28250.4ms/step | 0.7 samples/s
+Step 6180/37465 | Loss: 2.4530 | LR: 4.73e-04 | 28250.1ms/step | 0.7 samples/s
+Step 6190/37465 | Loss: 2.4972 | LR: 4.73e-04 | 28254.7ms/step | 0.7 samples/s
+Step 6200/37465 | Loss: 2.3914 | LR: 4.73e-04 | 28251.3ms/step | 0.7 samples/s
+Step 6210/37465 | Loss: 2.2029 | LR: 4.73e-04 | 28255.2ms/step | 0.7 samples/s
+Step 6220/37465 | Loss: 2.4602 | LR: 4.73e-04 | 28243.4ms/step | 0.7 samples/s
+Step 6230/37465 | Loss: 2.3374 | LR: 4.73e-04 | 28246.3ms/step | 0.7 samples/s
+Step 6240/37465 | Loss: 2.4669 | LR: 4.73e-04 | 28241.0ms/step | 0.7 samples/s
+Step 6250/37465 | Loss: 2.3224 | LR: 4.73e-04 | 28250.6ms/step | 0.7 samples/s
+Step 6260/37465 | Loss: 2.2651 | LR: 4.73e-04 | 28249.7ms/step | 0.7 samples/s
+Step 6270/37465 | Loss: 1.9950 | LR: 4.73e-04 | 28250.4ms/step | 0.7 samples/s
+Step 6280/37465 | Loss: 2.2056 | LR: 4.73e-04 | 28249.2ms/step | 0.7 samples/s
+Step 6290/37465 | Loss: 2.3963 | LR: 4.72e-04 | 28248.1ms/step | 0.7 samples/s
+Step 6300/37465 | Loss: 2.4665 | LR: 4.72e-04 | 28250.6ms/step | 0.7 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_resampled/run_20260708_112209/checkpoint-6300
+Step 6310/37465 | Loss: 2.2523 | LR: 4.72e-04 | 29169.6ms/step | 0.7 samples/s
+Step 6320/37465 | Loss: 2.3348 | LR: 4.72e-04 | 28244.1ms/step | 0.7 samples/s
+Step 6330/37465 | Loss: 2.4598 | LR: 4.72e-04 | 28244.5ms/step | 0.7 samples/s
+Step 6340/37465 | Loss: 2.2694 | LR: 4.72e-04 | 28247.1ms/step | 0.7 samples/s
+Step 6350/37465 | Loss: 2.6894 | LR: 4.72e-04 | 28247.1ms/step | 0.7 samples/s
+Step 6360/37465 | Loss: 2.2401 | LR: 4.72e-04 | 28245.2ms/step | 0.7 samples/s
+Step 6370/37465 | Loss: 2.1168 | LR: 4.72e-04 | 28243.0ms/step | 0.7 samples/s
+Step 6380/37465 | Loss: 2.5137 | LR: 4.72e-04 | 28244.8ms/step | 0.7 samples/s
+Step 6390/37465 | Loss: 2.5300 | LR: 4.71e-04 | 28242.7ms/step | 0.7 samples/s
+Step 6400/37465 | Loss: 2.3572 | LR: 4.71e-04 | 28247.1ms/step | 0.7 samples/s
+Step 6410/37465 | Loss: 2.4121 | LR: 4.71e-04 | 28243.0ms/step | 0.7 samples/s
+Step 6420/37465 | Loss: 2.2517 | LR: 4.71e-04 | 28245.8ms/step | 0.7 samples/s
+Step 6430/37465 | Loss: 2.2926 | LR: 4.71e-04 | 28241.6ms/step | 0.7 samples/s
+Step 6440/37465 | Loss: 2.2358 | LR: 4.71e-04 | 28240.5ms/step | 0.7 samples/s
+Step 6450/37465 | Loss: 2.3600 | LR: 4.71e-04 | 28244.1ms/step | 0.7 samples/s
+Step 6460/37465 | Loss: 2.3751 | LR: 4.71e-04 | 28248.0ms/step | 0.7 samples/s
+Step 6470/37465 | Loss: 2.3544 | LR: 4.71e-04 | 28242.3ms/step | 0.7 samples/s
+Step 6480/37465 | Loss: 2.3160 | LR: 4.71e-04 | 28241.6ms/step | 0.7 samples/s
+Step 6490/37465 | Loss: 2.3512 | LR: 4.70e-04 | 28240.6ms/step | 0.7 samples/s
+Step 6500/37465 | Loss: 2.5265 | LR: 4.70e-04 | 28240.7ms/step | 0.7 samples/s
+Validation at step 6500: Loss = 2.3912, Perplexity = 10.93
+Step 6510/37465 | Loss: 2.3988 | LR: 4.70e-04 | 29924.4ms/step | 0.7 samples/s
+Step 6520/37465 | Loss: 2.4542 | LR: 4.70e-04 | 28243.2ms/step | 0.7 samples/s
+Step 6530/37465 | Loss: 2.2724 | LR: 4.70e-04 | 28245.8ms/step | 0.7 samples/s
+Step 6540/37465 | Loss: 2.4087 | LR: 4.70e-04 | 28247.0ms/step | 0.7 samples/s
+Step 6550/37465 | Loss: 2.3598 | LR: 4.70e-04 | 28244.5ms/step | 0.7 samples/s
+Step 6560/37465 | Loss: 2.4893 | LR: 4.70e-04 | 28237.0ms/step | 0.7 samples/s
+Step 6570/37465 | Loss: 2.5836 | LR: 4.70e-04 | 28239.9ms/step | 0.7 samples/s
+Step 6580/37465 | Loss: 2.2616 | LR: 4.70e-04 | 28240.5ms/step | 0.7 samples/s
+Step 6590/37465 | Loss: 2.2487 | LR: 4.69e-04 | 28241.5ms/step | 0.7 samples/s
+Step 6600/37465 | Loss: 2.2595 | LR: 4.69e-04 | 28241.9ms/step | 0.7 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_resampled/run_20260708_112209/checkpoint-6600
+Step 6610/37465 | Loss: 2.3638 | LR: 4.69e-04 | 29182.0ms/step | 0.7 samples/s
+Step 6620/37465 | Loss: 2.5765 | LR: 4.69e-04 | 28240.2ms/step | 0.7 samples/s
+Step 6630/37465 | Loss: 2.3085 | LR: 4.69e-04 | 28245.7ms/step | 0.7 samples/s
+Step 6640/37465 | Loss: 2.3932 | LR: 4.69e-04 | 28246.6ms/step | 0.7 samples/s
+Step 6650/37465 | Loss: 2.5582 | LR: 4.69e-04 | 28245.0ms/step | 0.7 samples/s
+Step 6660/37465 | Loss: 2.4505 | LR: 4.69e-04 | 28244.2ms/step | 0.7 samples/s
+Step 6670/37465 | Loss: 2.2805 | LR: 4.69e-04 | 28243.3ms/step | 0.7 samples/s
+Step 6680/37465 | Loss: 2.4219 | LR: 4.69e-04 | 28243.7ms/step | 0.7 samples/s
+Step 6690/37465 | Loss: 2.6080 | LR: 4.68e-04 | 28241.1ms/step | 0.7 samples/s
+Step 6700/37465 | Loss: 2.1873 | LR: 4.68e-04 | 28241.0ms/step | 0.7 samples/s
+Step 6710/37465 | Loss: 2.5408 | LR: 4.68e-04 | 28243.8ms/step | 0.7 samples/s
+Step 6720/37465 | Loss: 2.2383 | LR: 4.68e-04 | 28240.1ms/step | 0.7 samples/s
+Step 6730/37465 | Loss: 2.4406 | LR: 4.68e-04 | 28245.0ms/step | 0.7 samples/s
+Step 6740/37465 | Loss: 2.2592 | LR: 4.68e-04 | 28245.1ms/step | 0.7 samples/s
+Step 6750/37465 | Loss: 2.2993 | LR: 4.68e-04 | 28249.3ms/step | 0.7 samples/s
+Step 6760/37465 | Loss: 2.4677 | LR: 4.68e-04 | 28240.8ms/step | 0.7 samples/s
+Step 6770/37465 | Loss: 2.2949 | LR: 4.68e-04 | 28243.6ms/step | 0.7 samples/s
+Step 6780/37465 | Loss: 2.5679 | LR: 4.67e-04 | 28243.1ms/step | 0.7 samples/s
+Step 6790/37465 | Loss: 2.4837 | LR: 4.67e-04 | 28243.5ms/step | 0.7 samples/s
+Step 6800/37465 | Loss: 2.3041 | LR: 4.67e-04 | 28244.3ms/step | 0.7 samples/s
+Step 6810/37465 | Loss: 2.5263 | LR: 4.67e-04 | 28247.6ms/step | 0.7 samples/s
+Step 6820/37465 | Loss: 2.5097 | LR: 4.67e-04 | 28239.5ms/step | 0.7 samples/s
+Step 6830/37465 | Loss: 2.6388 | LR: 4.67e-04 | 28246.6ms/step | 0.7 samples/s
+Step 6840/37465 | Loss: 2.4213 | LR: 4.67e-04 | 28249.3ms/step | 0.7 samples/s
+Step 6850/37465 | Loss: 2.3859 | LR: 4.67e-04 | 28240.0ms/step | 0.7 samples/s
+Step 6860/37465 | Loss: 2.3492 | LR: 4.67e-04 | 28241.7ms/step | 0.7 samples/s
+Step 6870/37465 | Loss: 2.4770 | LR: 4.66e-04 | 28243.6ms/step | 0.7 samples/s
+Step 6880/37465 | Loss: 2.3874 | LR: 4.66e-04 | 28241.0ms/step | 0.7 samples/s
+Step 6890/37465 | Loss: 2.5344 | LR: 4.66e-04 | 28247.8ms/step | 0.7 samples/s
+Step 6900/37465 | Loss: 2.4020 | LR: 4.66e-04 | 28242.2ms/step | 0.7 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_resampled/run_20260708_112209/checkpoint-6900
+Step 6910/37465 | Loss: 2.4471 | LR: 4.66e-04 | 29236.3ms/step | 0.7 samples/s
+Step 6920/37465 | Loss: 2.4108 | LR: 4.66e-04 | 28246.4ms/step | 0.7 samples/s
+Step 6930/37465 | Loss: 2.2877 | LR: 4.66e-04 | 28249.4ms/step | 0.7 samples/s
+Step 6940/37465 | Loss: 2.1094 | LR: 4.66e-04 | 28247.6ms/step | 0.7 samples/s
+Step 6950/37465 | Loss: 2.5191 | LR: 4.66e-04 | 28250.8ms/step | 0.7 samples/s
+Step 6960/37465 | Loss: 2.4610 | LR: 4.66e-04 | 28248.5ms/step | 0.7 samples/s
+Step 6970/37465 | Loss: 2.4725 | LR: 4.65e-04 | 28250.0ms/step | 0.7 samples/s
+Step 6980/37465 | Loss: 2.4281 | LR: 4.65e-04 | 28249.6ms/step | 0.7 samples/s
+Step 6990/37465 | Loss: 2.4955 | LR: 4.65e-04 | 28249.3ms/step | 0.7 samples/s
+Step 7000/37465 | Loss: 2.1528 | LR: 4.65e-04 | 28246.5ms/step | 0.7 samples/s
+Validation at step 7000: Loss = 2.3725, Perplexity = 10.72
+Step 7010/37465 | Loss: 2.2906 | LR: 4.65e-04 | 29916.1ms/step | 0.7 samples/s
+Step 7020/37465 | Loss: 2.3027 | LR: 4.65e-04 | 28242.9ms/step | 0.7 samples/s
+Step 7030/37465 | Loss: 2.3279 | LR: 4.65e-04 | 28259.6ms/step | 0.7 samples/s
+Step 7040/37465 | Loss: 2.2016 | LR: 4.65e-04 | 28246.8ms/step | 0.7 samples/s
+Step 7050/37465 | Loss: 2.3227 | LR: 4.65e-04 | 28246.0ms/step | 0.7 samples/s
+Step 7060/37465 | Loss: 2.4083 | LR: 4.64e-04 | 28248.1ms/step | 0.7 samples/s
+Step 7070/37465 | Loss: 2.2603 | LR: 4.64e-04 | 28253.4ms/step | 0.7 samples/s
+Step 7080/37465 | Loss: 2.4541 | LR: 4.64e-04 | 28249.2ms/step | 0.7 samples/s
+Step 7090/37465 | Loss: 2.2565 | LR: 4.64e-04 | 28247.1ms/step | 0.7 samples/s
+Step 7100/37465 | Loss: 2.3519 | LR: 4.64e-04 | 28250.4ms/step | 0.7 samples/s
+Step 7110/37465 | Loss: 2.1629 | LR: 4.64e-04 | 28248.0ms/step | 0.7 samples/s
+Step 7120/37465 | Loss: 1.9042 | LR: 4.64e-04 | 28243.9ms/step | 0.7 samples/s
+Step 7130/37465 | Loss: 2.3516 | LR: 4.64e-04 | 28248.7ms/step | 0.7 samples/s
+Step 7140/37465 | Loss: 2.4558 | LR: 4.64e-04 | 28246.8ms/step | 0.7 samples/s
+Step 7150/37465 | Loss: 2.5282 | LR: 4.63e-04 | 28245.0ms/step | 0.7 samples/s
+Step 7160/37465 | Loss: 2.4421 | LR: 4.63e-04 | 28241.1ms/step | 0.7 samples/s
+Step 7170/37465 | Loss: 2.3576 | LR: 4.63e-04 | 28245.4ms/step | 0.7 samples/s
+Step 7180/37465 | Loss: 2.5474 | LR: 4.63e-04 | 28246.7ms/step | 0.7 samples/s
+Step 7190/37465 | Loss: 2.6020 | LR: 4.63e-04 | 28253.0ms/step | 0.7 samples/s
+Step 7200/37465 | Loss: 2.3393 | LR: 4.63e-04 | 28248.8ms/step | 0.7 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_resampled/run_20260708_112209/checkpoint-7200
+Step 7210/37465 | Loss: 2.6241 | LR: 4.63e-04 | 29209.9ms/step | 0.7 samples/s
+Step 7220/37465 | Loss: 2.3910 | LR: 4.63e-04 | 28256.3ms/step | 0.7 samples/s
+Step 7230/37465 | Loss: 2.4426 | LR: 4.63e-04 | 28249.9ms/step | 0.7 samples/s
+Step 7240/37465 | Loss: 2.3983 | LR: 4.62e-04 | 28251.8ms/step | 0.7 samples/s
+Step 7250/37465 | Loss: 2.3759 | LR: 4.62e-04 | 28248.1ms/step | 0.7 samples/s
+Step 7260/37465 | Loss: 2.3979 | LR: 4.62e-04 | 28247.6ms/step | 0.7 samples/s
+Step 7270/37465 | Loss: 2.2362 | LR: 4.62e-04 | 28253.1ms/step | 0.7 samples/s
+Step 7280/37465 | Loss: 2.4393 | LR: 4.62e-04 | 28257.1ms/step | 0.7 samples/s
+Step 7290/37465 | Loss: 2.6248 | LR: 4.62e-04 | 28248.5ms/step | 0.7 samples/s
+Step 7300/37465 | Loss: 2.3450 | LR: 4.62e-04 | 28247.9ms/step | 0.7 samples/s
+Step 7310/37465 | Loss: 2.3153 | LR: 4.62e-04 | 28245.4ms/step | 0.7 samples/s
+Step 7320/37465 | Loss: 2.2891 | LR: 4.62e-04 | 28249.1ms/step | 0.7 samples/s
+Step 7330/37465 | Loss: 2.4664 | LR: 4.61e-04 | 28243.0ms/step | 0.7 samples/s
+Step 7340/37465 | Loss: 2.3123 | LR: 4.61e-04 | 28304.2ms/step | 0.7 samples/s
+Step 7350/37465 | Loss: 2.4346 | LR: 4.61e-04 | 28247.8ms/step | 0.7 samples/s
+Step 7360/37465 | Loss: 2.4559 | LR: 4.61e-04 | 28248.1ms/step | 0.7 samples/s
+Step 7370/37465 | Loss: 2.1096 | LR: 4.61e-04 | 28244.8ms/step | 0.7 samples/s
+Step 7380/37465 | Loss: 2.4060 | LR: 4.61e-04 | 28244.2ms/step | 0.7 samples/s
+Step 7390/37465 | Loss: 2.3375 | LR: 4.61e-04 | 28243.7ms/step | 0.7 samples/s
+Step 7400/37465 | Loss: 2.3464 | LR: 4.61e-04 | 28248.4ms/step | 0.7 samples/s
+Step 7410/37465 | Loss: 2.2242 | LR: 4.60e-04 | 28244.9ms/step | 0.7 samples/s
+Step 7420/37465 | Loss: 2.5017 | LR: 4.60e-04 | 28238.8ms/step | 0.7 samples/s
+Step 7430/37465 | Loss: 2.4564 | LR: 4.60e-04 | 28242.3ms/step | 0.7 samples/s
+Step 7440/37465 | Loss: 2.1967 | LR: 4.60e-04 | 28241.7ms/step | 0.7 samples/s
+Step 7450/37465 | Loss: 2.2369 | LR: 4.60e-04 | 28237.4ms/step | 0.7 samples/s
+Step 7460/37465 | Loss: 2.3569 | LR: 4.60e-04 | 28235.5ms/step | 0.7 samples/s
+Step 7470/37465 | Loss: 2.2035 | LR: 4.60e-04 | 28239.7ms/step | 0.7 samples/s
+Step 7480/37465 | Loss: 2.4400 | LR: 4.60e-04 | 28242.8ms/step | 0.7 samples/s
+Step 7490/37465 | Loss: 2.4374 | LR: 4.60e-04 | 28241.5ms/step | 0.7 samples/s
+Step 7500/37465 | Loss: 2.2728 | LR: 4.59e-04 | 28238.1ms/step | 0.7 samples/s
+Validation at step 7500: Loss = 2.3558, Perplexity = 10.55
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_resampled/run_20260708_112209/checkpoint-7500
+Step 7510/37465 | Loss: 2.3687 | LR: 4.59e-04 | 30818.7ms/step | 0.6 samples/s
+Step 7520/37465 | Loss: 2.4084 | LR: 4.59e-04 | 28237.4ms/step | 0.7 samples/s
+Step 7530/37465 | Loss: 2.4761 | LR: 4.59e-04 | 28243.2ms/step | 0.7 samples/s
+Step 7540/37465 | Loss: 2.2384 | LR: 4.59e-04 | 28241.8ms/step | 0.7 samples/s
+Step 7550/37465 | Loss: 2.4994 | LR: 4.59e-04 | 28238.9ms/step | 0.7 samples/s
+Step 7560/37465 | Loss: 2.3619 | LR: 4.59e-04 | 28242.3ms/step | 0.7 samples/s
+Step 7570/37465 | Loss: 2.3680 | LR: 4.59e-04 | 28240.4ms/step | 0.7 samples/s
+Step 7580/37465 | Loss: 2.5105 | LR: 4.59e-04 | 28244.5ms/step | 0.7 samples/s
+Step 7590/37465 | Loss: 2.3612 | LR: 4.58e-04 | 28244.0ms/step | 0.7 samples/s
+Step 7600/37465 | Loss: 2.5033 | LR: 4.58e-04 | 28245.7ms/step | 0.7 samples/s
+Step 7610/37465 | Loss: 2.4621 | LR: 4.58e-04 | 28241.6ms/step | 0.7 samples/s
+Step 7620/37465 | Loss: 2.1627 | LR: 4.58e-04 | 28245.2ms/step | 0.7 samples/s
+Step 7630/37465 | Loss: 2.3425 | LR: 4.58e-04 | 28252.2ms/step | 0.7 samples/s
+Step 7640/37465 | Loss: 2.2341 | LR: 4.58e-04 | 28248.8ms/step | 0.7 samples/s
+Step 7650/37465 | Loss: 2.3031 | LR: 4.58e-04 | 28244.5ms/step | 0.7 samples/s
+Step 7660/37465 | Loss: 2.3506 | LR: 4.58e-04 | 28241.5ms/step | 0.7 samples/s
+Step 7670/37465 | Loss: 2.2373 | LR: 4.57e-04 | 28242.9ms/step | 0.7 samples/s
+Step 7680/37465 | Loss: 2.3048 | LR: 4.57e-04 | 28242.9ms/step | 0.7 samples/s
+Step 7690/37465 | Loss: 2.1781 | LR: 4.57e-04 | 28250.6ms/step | 0.7 samples/s
+Step 7700/37465 | Loss: 2.2315 | LR: 4.57e-04 | 28254.5ms/step | 0.7 samples/s
+Step 7710/37465 | Loss: 2.6047 | LR: 4.57e-04 | 28240.5ms/step | 0.7 samples/s
+Step 7720/37465 | Loss: 2.1713 | LR: 4.57e-04 | 28246.1ms/step | 0.7 samples/s
+Step 7730/37465 | Loss: 2.3746 | LR: 4.57e-04 | 28242.5ms/step | 0.7 samples/s
+Step 7740/37465 | Loss: 2.4747 | LR: 4.57e-04 | 28242.6ms/step | 0.7 samples/s
+Step 7750/37465 | Loss: 2.4702 | LR: 4.56e-04 | 28242.7ms/step | 0.7 samples/s
+Step 7760/37465 | Loss: 2.3955 | LR: 4.56e-04 | 28248.0ms/step | 0.7 samples/s
+Step 7770/37465 | Loss: 2.1627 | LR: 4.56e-04 | 28243.0ms/step | 0.7 samples/s
+Step 7780/37465 | Loss: 2.1919 | LR: 4.56e-04 | 28247.6ms/step | 0.7 samples/s
+Step 7790/37465 | Loss: 2.1817 | LR: 4.56e-04 | 28246.8ms/step | 0.7 samples/s
+Step 7800/37465 | Loss: 2.1525 | LR: 4.56e-04 | 28249.9ms/step | 0.7 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_resampled/run_20260708_112209/checkpoint-7800
+Step 7810/37465 | Loss: 2.3451 | LR: 4.56e-04 | 29179.2ms/step | 0.7 samples/s
+Step 7820/37465 | Loss: 2.3248 | LR: 4.56e-04 | 28250.1ms/step | 0.7 samples/s
+Step 7830/37465 | Loss: 2.2779 | LR: 4.56e-04 | 28251.5ms/step | 0.7 samples/s
+Step 7840/37465 | Loss: 2.4719 | LR: 4.55e-04 | 28255.2ms/step | 0.7 samples/s
+Step 7850/37465 | Loss: 2.2598 | LR: 4.55e-04 | 28250.4ms/step | 0.7 samples/s
+Step 7860/37465 | Loss: 2.2278 | LR: 4.55e-04 | 28248.6ms/step | 0.7 samples/s
+Step 7870/37465 | Loss: 2.3601 | LR: 4.55e-04 | 28253.1ms/step | 0.7 samples/s
+Step 7880/37465 | Loss: 2.2324 | LR: 4.55e-04 | 28250.3ms/step | 0.7 samples/s
+Step 7890/37465 | Loss: 2.4406 | LR: 4.55e-04 | 28249.0ms/step | 0.7 samples/s
+Step 7900/37465 | Loss: 2.2465 | LR: 4.55e-04 | 28250.4ms/step | 0.7 samples/s
+Step 7910/37465 | Loss: 2.2177 | LR: 4.55e-04 | 28258.2ms/step | 0.7 samples/s
+Step 7920/37465 | Loss: 2.1655 | LR: 4.54e-04 | 28254.0ms/step | 0.7 samples/s
+Step 7930/37465 | Loss: 2.2817 | LR: 4.54e-04 | 28253.9ms/step | 0.7 samples/s
+Step 7940/37465 | Loss: 2.1449 | LR: 4.54e-04 | 28255.3ms/step | 0.7 samples/s
+Step 7950/37465 | Loss: 2.3005 | LR: 4.54e-04 | 28252.8ms/step | 0.7 samples/s
+Step 7960/37465 | Loss: 2.4433 | LR: 4.54e-04 | 28247.2ms/step | 0.7 samples/s
+Step 7970/37465 | Loss: 2.3244 | LR: 4.54e-04 | 28251.7ms/step | 0.7 samples/s
+Step 7980/37465 | Loss: 2.3075 | LR: 4.54e-04 | 28249.7ms/step | 0.7 samples/s
+Step 7990/37465 | Loss: 2.4497 | LR: 4.54e-04 | 28259.1ms/step | 0.7 samples/s
+Step 8000/37465 | Loss: 2.3176 | LR: 4.53e-04 | 28251.1ms/step | 0.7 samples/s
+Validation at step 8000: Loss = 2.3356, Perplexity = 10.34
+Step 8010/37465 | Loss: 2.0908 | LR: 4.53e-04 | 29928.8ms/step | 0.7 samples/s
+Step 8020/37465 | Loss: 2.3554 | LR: 4.53e-04 | 28249.7ms/step | 0.7 samples/s
+Step 8030/37465 | Loss: 2.3430 | LR: 4.53e-04 | 28250.0ms/step | 0.7 samples/s
+Step 8040/37465 | Loss: 2.3104 | LR: 4.53e-04 | 28252.5ms/step | 0.7 samples/s
+Step 8050/37465 | Loss: 2.1452 | LR: 4.53e-04 | 28253.9ms/step | 0.7 samples/s
+Step 8060/37465 | Loss: 2.4289 | LR: 4.53e-04 | 28250.1ms/step | 0.7 samples/s
+Step 8070/37465 | Loss: 2.4227 | LR: 4.53e-04 | 28249.3ms/step | 0.7 samples/s
+Step 8080/37465 | Loss: 2.3569 | LR: 4.52e-04 | 28248.7ms/step | 0.7 samples/s
+Step 8090/37465 | Loss: 2.2148 | LR: 4.52e-04 | 28253.3ms/step | 0.7 samples/s
+Step 8100/37465 | Loss: 2.4227 | LR: 4.52e-04 | 28251.8ms/step | 0.7 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_resampled/run_20260708_112209/checkpoint-8100
+Step 8110/37465 | Loss: 2.4100 | LR: 4.52e-04 | 29172.6ms/step | 0.7 samples/s
+Step 8120/37465 | Loss: 2.4015 | LR: 4.52e-04 | 28258.9ms/step | 0.7 samples/s
+Step 8130/37465 | Loss: 2.3536 | LR: 4.52e-04 | 28257.8ms/step | 0.7 samples/s
+Step 8140/37465 | Loss: 2.2470 | LR: 4.52e-04 | 28252.8ms/step | 0.7 samples/s
+Step 8150/37465 | Loss: 2.4475 | LR: 4.52e-04 | 28251.1ms/step | 0.7 samples/s
+Step 8160/37465 | Loss: 2.2579 | LR: 4.51e-04 | 28252.0ms/step | 0.7 samples/s
+Step 8170/37465 | Loss: 2.3338 | LR: 4.51e-04 | 28252.2ms/step | 0.7 samples/s
+Step 8180/37465 | Loss: 2.4141 | LR: 4.51e-04 | 28251.5ms/step | 0.7 samples/s
+Step 8190/37465 | Loss: 2.2446 | LR: 4.51e-04 | 28250.5ms/step | 0.7 samples/s
+Step 8200/37465 | Loss: 2.3346 | LR: 4.51e-04 | 28245.6ms/step | 0.7 samples/s
+Step 8210/37465 | Loss: 2.4666 | LR: 4.51e-04 | 28254.4ms/step | 0.7 samples/s
+Step 8220/37465 | Loss: 2.3568 | LR: 4.51e-04 | 28252.4ms/step | 0.7 samples/s
+Step 8230/37465 | Loss: 2.2838 | LR: 4.51e-04 | 28244.7ms/step | 0.7 samples/s
+Step 8240/37465 | Loss: 2.2910 | LR: 4.50e-04 | 28249.7ms/step | 0.7 samples/s
+Step 8250/37465 | Loss: 2.2859 | LR: 4.50e-04 | 28248.8ms/step | 0.7 samples/s
+Step 8260/37465 | Loss: 2.5687 | LR: 4.50e-04 | 28255.1ms/step | 0.7 samples/s
+Step 8270/37465 | Loss: 2.4251 | LR: 4.50e-04 | 28249.1ms/step | 0.7 samples/s
+Step 8280/37465 | Loss: 2.1694 | LR: 4.50e-04 | 28248.8ms/step | 0.7 samples/s
+Step 8290/37465 | Loss: 2.4229 | LR: 4.50e-04 | 28247.3ms/step | 0.7 samples/s
+Step 8300/37465 | Loss: 2.2396 | LR: 4.50e-04 | 28252.5ms/step | 0.7 samples/s
+Step 8310/37465 | Loss: 2.4142 | LR: 4.49e-04 | 28246.5ms/step | 0.7 samples/s
+Step 8320/37465 | Loss: 2.4780 | LR: 4.49e-04 | 28250.6ms/step | 0.7 samples/s
+Step 8330/37465 | Loss: 2.1195 | LR: 4.49e-04 | 28249.1ms/step | 0.7 samples/s
+Step 8340/37465 | Loss: 2.4829 | LR: 4.49e-04 | 28245.3ms/step | 0.7 samples/s
+Step 8350/37465 | Loss: 2.4091 | LR: 4.49e-04 | 28251.2ms/step | 0.7 samples/s
+Step 8360/37465 | Loss: 2.0440 | LR: 4.49e-04 | 28249.2ms/step | 0.7 samples/s
+Step 8370/37465 | Loss: 2.3666 | LR: 4.49e-04 | 28245.9ms/step | 0.7 samples/s
+Step 8380/37465 | Loss: 2.1888 | LR: 4.49e-04 | 28247.0ms/step | 0.7 samples/s
+Step 8390/37465 | Loss: 2.2704 | LR: 4.48e-04 | 28254.2ms/step | 0.7 samples/s
+Step 8400/37465 | Loss: 2.2250 | LR: 4.48e-04 | 28250.0ms/step | 0.7 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_resampled/run_20260708_112209/checkpoint-8400
+Step 8410/37465 | Loss: 2.3778 | LR: 4.48e-04 | 29167.1ms/step | 0.7 samples/s
+Step 8420/37465 | Loss: 2.2435 | LR: 4.48e-04 | 28256.3ms/step | 0.7 samples/s
+Step 8430/37465 | Loss: 2.3137 | LR: 4.48e-04 | 28248.6ms/step | 0.7 samples/s
+Step 8440/37465 | Loss: 2.2650 | LR: 4.48e-04 | 28255.9ms/step | 0.7 samples/s
+Step 8450/37465 | Loss: 2.1394 | LR: 4.48e-04 | 28249.5ms/step | 0.7 samples/s
+Step 8460/37465 | Loss: 2.2003 | LR: 4.48e-04 | 28254.4ms/step | 0.7 samples/s
+Step 8470/37465 | Loss: 2.2522 | LR: 4.47e-04 | 28248.8ms/step | 0.7 samples/s
+Step 8480/37465 | Loss: 2.2528 | LR: 4.47e-04 | 28245.1ms/step | 0.7 samples/s
+Step 8490/37465 | Loss: 2.3905 | LR: 4.47e-04 | 28247.5ms/step | 0.7 samples/s
+Step 8500/37465 | Loss: 2.3215 | LR: 4.47e-04 | 28249.3ms/step | 0.7 samples/s
+Validation at step 8500: Loss = 2.3238, Perplexity = 10.21
+Step 8510/37465 | Loss: 2.2486 | LR: 4.47e-04 | 29917.7ms/step | 0.7 samples/s
+Step 8520/37465 | Loss: 2.2742 | LR: 4.47e-04 | 28242.0ms/step | 0.7 samples/s
+Step 8530/37465 | Loss: 2.1569 | LR: 4.47e-04 | 28245.5ms/step | 0.7 samples/s
+Step 8540/37465 | Loss: 2.2703 | LR: 4.46e-04 | 28243.5ms/step | 0.7 samples/s
+Step 8550/37465 | Loss: 2.2952 | LR: 4.46e-04 | 28243.9ms/step | 0.7 samples/s
+Step 8560/37465 | Loss: 2.5951 | LR: 4.46e-04 | 28242.0ms/step | 0.7 samples/s
+Step 8570/37465 | Loss: 2.2779 | LR: 4.46e-04 | 28243.8ms/step | 0.7 samples/s
+Step 8580/37465 | Loss: 2.3204 | LR: 4.46e-04 | 28249.0ms/step | 0.7 samples/s
+Step 8590/37465 | Loss: 2.1752 | LR: 4.46e-04 | 28244.7ms/step | 0.7 samples/s
+Step 8600/37465 | Loss: 2.3792 | LR: 4.46e-04 | 28246.3ms/step | 0.7 samples/s
+Step 8610/37465 | Loss: 2.6003 | LR: 4.46e-04 | 28243.3ms/step | 0.7 samples/s
+Step 8620/37465 | Loss: 2.1397 | LR: 4.45e-04 | 28239.0ms/step | 0.7 samples/s
+Step 8630/37465 | Loss: 2.3302 | LR: 4.45e-04 | 28244.1ms/step | 0.7 samples/s
+Step 8640/37465 | Loss: 2.4079 | LR: 4.45e-04 | 28243.8ms/step | 0.7 samples/s
+Step 8650/37465 | Loss: 2.2444 | LR: 4.45e-04 | 28243.3ms/step | 0.7 samples/s
+Step 8660/37465 | Loss: 2.2380 | LR: 4.45e-04 | 28240.7ms/step | 0.7 samples/s
+Step 8670/37465 | Loss: 2.3202 | LR: 4.45e-04 | 28240.5ms/step | 0.7 samples/s
+Step 8680/37465 | Loss: 2.4565 | LR: 4.45e-04 | 28241.6ms/step | 0.7 samples/s
+Step 8690/37465 | Loss: 2.2976 | LR: 4.44e-04 | 28242.9ms/step | 0.7 samples/s
+Step 8700/37465 | Loss: 2.4527 | LR: 4.44e-04 | 28244.4ms/step | 0.7 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_resampled/run_20260708_112209/checkpoint-8700
+Step 8710/37465 | Loss: 2.2391 | LR: 4.44e-04 | 29148.6ms/step | 0.7 samples/s
+Step 8720/37465 | Loss: 2.2357 | LR: 4.44e-04 | 28244.1ms/step | 0.7 samples/s
+Step 8730/37465 | Loss: 2.4762 | LR: 4.44e-04 | 28245.6ms/step | 0.7 samples/s
+Step 8740/37465 | Loss: 2.2624 | LR: 4.44e-04 | 28240.7ms/step | 0.7 samples/s
+Step 8750/37465 | Loss: 2.2417 | LR: 4.44e-04 | 28238.0ms/step | 0.7 samples/s
+Step 8760/37465 | Loss: 2.3755 | LR: 4.44e-04 | 28242.7ms/step | 0.7 samples/s
+Step 8770/37465 | Loss: 2.3873 | LR: 4.43e-04 | 28237.3ms/step | 0.7 samples/s
+Step 8780/37465 | Loss: 2.1722 | LR: 4.43e-04 | 28243.7ms/step | 0.7 samples/s
+Step 8790/37465 | Loss: 2.4269 | LR: 4.43e-04 | 28236.1ms/step | 0.7 samples/s
+Step 8800/37465 | Loss: 2.3189 | LR: 4.43e-04 | 28235.2ms/step | 0.7 samples/s
+Step 8810/37465 | Loss: 2.2542 | LR: 4.43e-04 | 28243.7ms/step | 0.7 samples/s
+Step 8820/37465 | Loss: 2.3955 | LR: 4.43e-04 | 28239.9ms/step | 0.7 samples/s
+Step 8830/37465 | Loss: 2.2250 | LR: 4.43e-04 | 28240.2ms/step | 0.7 samples/s
+Step 8840/37465 | Loss: 2.1999 | LR: 4.42e-04 | 28237.8ms/step | 0.7 samples/s
+Step 8850/37465 | Loss: 2.1886 | LR: 4.42e-04 | 28240.7ms/step | 0.7 samples/s
+Step 8860/37465 | Loss: 2.3696 | LR: 4.42e-04 | 28239.7ms/step | 0.7 samples/s
+Step 8870/37465 | Loss: 2.2762 | LR: 4.42e-04 | 28246.2ms/step | 0.7 samples/s
+Step 8880/37465 | Loss: 2.4358 | LR: 4.42e-04 | 28237.9ms/step | 0.7 samples/s
+Step 8890/37465 | Loss: 2.2100 | LR: 4.42e-04 | 28240.0ms/step | 0.7 samples/s
+Step 8900/37465 | Loss: 2.0125 | LR: 4.42e-04 | 28235.1ms/step | 0.7 samples/s
+Step 8910/37465 | Loss: 2.2455 | LR: 4.41e-04 | 28243.8ms/step | 0.7 samples/s
+Step 8920/37465 | Loss: 2.4135 | LR: 4.41e-04 | 28240.2ms/step | 0.7 samples/s
+Step 8930/37465 | Loss: 2.5532 | LR: 4.41e-04 | 28241.7ms/step | 0.7 samples/s
+Step 8940/37465 | Loss: 2.3636 | LR: 4.41e-04 | 28235.1ms/step | 0.7 samples/s
+Step 8950/37465 | Loss: 2.4623 | LR: 4.41e-04 | 28236.1ms/step | 0.7 samples/s
+Step 8960/37465 | Loss: 2.4028 | LR: 4.41e-04 | 28238.4ms/step | 0.7 samples/s
+Step 8970/37465 | Loss: 2.1772 | LR: 4.41e-04 | 28238.6ms/step | 0.7 samples/s
+W0711 01:45:57.049000 3736 site-packages/torch/distributed/elastic/agent/server/api.py:739] Received 15 death signal, shutting down workers
+W0711 01:45:57.052000 3736 site-packages/torch/distributed/elastic/multiprocessing/api.py:1012] Sending process 3805 closing signal SIGTERM
+W0711 01:45:57.054000 3736 site-packages/torch/distributed/elastic/multiprocessing/api.py:1012] Sending process 3806 closing signal SIGTERM
+W0711 01:45:57.057000 3736 site-packages/torch/distributed/elastic/multiprocessing/api.py:1012] Sending process 3807 closing signal SIGTERM
+W0711 01:45:57.059000 3736 site-packages/torch/distributed/elastic/multiprocessing/api.py:1012] Sending process 3808 closing signal SIGTERM
+Traceback (most recent call last):
+  File "/mloscratch/homes/navasard/conda/bin/torchrun", line 6, in <module>
+    sys.exit(main())
+             ~~~~^^
+  File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 362, in wrapper
+    return f(*args, **kwargs)
+  File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/distributed/run.py", line 990, in main
+    run(args)
+    ~~~^^^^^^
+  File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/distributed/run.py", line 981, in run
+    elastic_launch(
+    ~~~~~~~~~~~~~~~
+        config=config,
+        ~~~~~~~~~~~~~~
+        entrypoint=cmd,
+        ~~~~~~~~~~~~~~~
+    )(*cmd_args)
+    ~^^^^^^^^^^^
+  File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/distributed/launcher/api.py", line 170, in __call__
+    return launch_agent(self._config, self._entrypoint, list(args))
+  File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/distributed/launcher/api.py", line 308, in launch_agent
+    result = agent.run()
+  File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/distributed/elastic/metrics/api.py", line 134, in wrapper
+    result = f(*args, **kwargs)
+  File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/distributed/elastic/agent/server/api.py", line 731, in run
+    result = self._invoke_run(role)
+  File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/distributed/elastic/agent/server/api.py", line 908, in _invoke_run
+    time.sleep(monitor_interval)
+    ~~~~~~~~~~^^^^^^^^^^^^^^^^^^
+  File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/distributed/elastic/multiprocessing/api.py", line 86, in _terminate_process_handler
+    raise SignalException(f"Process {os.getpid()} got signal: {sigval}", sigval=sigval)
+torch.distributed.elastic.multiprocessing.api.SignalException: Process 3736 got signal: 15

epfl_final/train_logs/train_moba_1p2b_matched.log

@@ -0,0 +1,669 @@
+W0711 08:31:23.427000 151 site-packages/torch/distributed/run.py:851] 
+W0711 08:31:23.427000 151 site-packages/torch/distributed/run.py:851] *****************************************
+W0711 08:31:23.427000 151 site-packages/torch/distributed/run.py:851] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. 
+W0711 08:31:23.427000 151 site-packages/torch/distributed/run.py:851] *****************************************
+New run: Checkpoints will be saved to /mloscratch/homes/navasard/sparse/checkpoints/l8k1_2b_moba_matched/run_20260711_083136
+Running config:
+{
+  "model_path": "/mloscratch/homes/navasard/qwen3_06b_base",
+  "model_dtype": "bfloat16",
+  "init_from_scratch": true,
+  "train_path": "/mloscratch/homes/navasard/data/fineweb_large/tokenized",
+  "val_path": null,
+  "max_seq_length": 8192,
+  "num_workers": 8,
+  "output_dir": "/mloscratch/homes/navasard/sparse/checkpoints/l8k1_2b_moba_matched/run_20260711_083136",
+  "resume_from_checkpoint": null,
+  "force_new_run": true,
+  "num_epochs": 1,
+  "max_steps": null,
+  "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": 10,
+  "seed": 42,
+  "verbose": true,
+  "checkpoint_activations": false,
+  "strategy": "ddp",
+  "mixed_precision": "bf16",
+  "wandb_enabled": true,
+  "wandb_project": "stochastic_sparsity",
+  "wandb_name": "1.2b moba 8k MATCHED",
+  "wandb_tags": [
+    "qwen3",
+    "1.2b",
+    "8k",
+    "moba",
+    "matched"
+  ],
+  "hidden_size": 1536,
+  "intermediate_size": 6144,
+  "num_attention_heads": 12,
+  "num_key_value_heads": 4,
+  "num_hidden_layers": null,
+  "head_dim": null,
+  "attention_type": "moba",
+  "sparse_block_size": 128,
+  "sparse_num_random_blocks": 1,
+  "sparse_mask_pool_size": 1000,
+  "sparse_val_dense_compare": false,
+  "sparse_layer_indices": null,
+  "sparse_mask_cache_dir": null,
+  "sparse_mask_seed": 0,
+  "dense_mode": false,
+  "dense_step_intervals": null,
+  "moba_block_size": 128,
+  "moba_top_k": 32,
+  "asentmax_alpha": 1.5,
+  "asentmax_learn_temp": true,
+  "asentmax_init_temp": 1.0
+}
+Loading tokenizer from /mloscratch/homes/navasard/qwen3_06b_base...
+Checkpoint activations: False
+Initializing model from scratch (random weights)...
+Checkpoint activations: False
+Checkpoint activations: False
+Checkpoint activations: False
+
+
+================================================================================
+MODEL ARCHITECTURE
+================================================================================
+Model config:
+  vocab_size: 151936
+  hidden_size: 1536
+  num_hidden_layers: 28
+  num_attention_heads: 12
+  num_key_value_heads: 4
+  head_dim: 128
+  intermediate_size: 6144
+  max_position_embeddings: 32768
+  rope_theta: 1000000
+  attention_type: moba
+
+Parameter count:
+  Total parameters: 1,202,352,640
+  Trainable parameters: 1,202,352,640
+  Non-trainable parameters: 0
+  Model size (BF16): 2.40 GB
+
+--------------------------------------------------------------------------------
+LAYER-BY-LAYER SUMMARY
+--------------------------------------------------------------------------------
+/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/nn/attention/flex_attention.py:1704: UserWarning: flex_attention called without torch.compile() - this will use an unfused implementation that materializes the full scores matrix instead of generating a fused kernel.
+
+SOLUTION: Use torch.compile(flex_attention)(...)
+
+If you want to debug your score_mod/mask_mod, you can set:
+torch.nn.attention.flex_attention._FLEX_ATTENTION_DISABLE_COMPILE_DEBUG = True
+
+This will allow you to use print statements or breakpoints. Note: This doesn't work with the backwards pass and may produce incorrect results.
+  _warn_once(
+===========================================================================================================================================================
+Layer (type:depth-idx)                                  Input Shape               Output Shape              Param #                   Trainable
+===========================================================================================================================================================
+Qwen3ForCausalLM                                        [2, 8192]                 --                        --                        True
+├─Qwen3Model: 1-1                                       --                        [2, 8192, 1536]           --                        True
+│    └─Embedding: 2-1                                   [2, 8192]                 [2, 8192, 1536]           233,373,696               True
+│    └─ModuleList: 2-2                                  --                        --                        --                        True
+│    │    └─Qwen3DecoderLayer: 3-1                      [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-2                      [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-3                      [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-4                      [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-5                      [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-6                      [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-7                      [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-8                      [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-9                      [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-10                     [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-11                     [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-12                     [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-13                     [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-14                     [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-15                     [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-16                     [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-17                     [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-18                     [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-19                     [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-20                     [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-21                     [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-22                     [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-23                     [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-24                     [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-25                     [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-26                     [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-27                     [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-28                     [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    └─Qwen3RMSNorm: 2-3                                [2, 8192, 1536]           [2, 8192, 1536]           1,536                     True
+├─Linear: 1-2                                           [2, 8192, 1536]           [2, 8192, 151936]         233,373,696               True
+===========================================================================================================================================================
+Total params: 1,435,726,336
+Trainable params: 1,435,726,336
+Non-trainable params: 0
+Total mult-adds (Units.GIGABYTES): 2.87
+===========================================================================================================================================================
+Input size (MB): 0.13
+Forward/backward pass size (MB): 52437.19
+Params size (MB): 2871.45
+Estimated Total Size (MB): 55308.77
+===========================================================================================================================================================
+  NOTE: torchinfo double-counts tied embed_tokens/lm_head weight (233,373,696 params). True unique parameter count: 1,202,352,640 (1202.4M)
+================================================================================
+
+Loading dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Using pre-tokenized dataset (no runtime tokenization)
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Loaded 12748481 pre-tokenized examples (stored_len=2048, pack_factor=4, packed_examples=3187120)
+Loaded 12748481 pre-tokenized examples (stored_len=2048, pack_factor=4, packed_examples=3187120)
+Loaded 12748481 pre-tokenized examples (stored_len=2048, pack_factor=4, packed_examples=3187120)
+Loaded 12748481 pre-tokenized examples (stored_len=2048, pack_factor=4, packed_examples=3187120)
+Dataset split: 3027764 train, 159356 validation
+Training examples (packed, train split):  3027764
+Validation examples (packed, val split):  159356
+Per-device batch size:                    2
+World size (GPUs):                        4
+Gradient accumulation steps:              10
+Effective batch size (seqs/step):         80
+Sequence length:                          8192
+Tokens per optimizer step:                655,360
+Micro-batches per epoch (per rank):       378470
+Optimizer steps per epoch:                37847
+Total optimizer steps:                    37847
+wandb: [wandb.login()] Loaded credentials for https://api.wandb.ai from WANDB_API_KEY.
+wandb: Currently logged in as: ruben-navasardyan4869 (ruben-navasardyan4869-epfl) to https://api.wandb.ai. Use `wandb login --relogin` to force relogin
+wandb: Tracking run with wandb version 0.26.1
+wandb: Run data is saved locally in /mloscratch/homes/navasard/sparse/SparseAtt-baselines/wandb/run-20260711_083217-cjxy1m86
+wandb: Run `wandb offline` to turn off syncing.
+wandb: Syncing run 1.2b moba 8k MATCHED
+wandb: ⭐️ View project at https://wandb.ai/ruben-navasardyan4869-epfl/stochastic_sparsity
+wandb: 🚀 View run at https://wandb.ai/ruben-navasardyan4869-epfl/stochastic_sparsity/runs/cjxy1m86
+Starting training...
+Switching attention mode for step 1: sparse_stochastic (configured sparse attention)
+W0711 08:33:33.592000 151 site-packages/torch/distributed/elastic/agent/server/api.py:739] Received 15 death signal, shutting down workers
+W0711 08:33:33.593000 151 site-packages/torch/distributed/elastic/multiprocessing/api.py:1012] Sending process 220 closing signal SIGTERM
+W0711 08:33:33.608000 151 site-packages/torch/distributed/elastic/multiprocessing/api.py:1012] Sending process 221 closing signal SIGTERM
+W0711 08:33:33.623000 151 site-packages/torch/distributed/elastic/multiprocessing/api.py:1012] Sending process 222 closing signal SIGTERM
+W0711 08:33:33.623000 151 site-packages/torch/distributed/elastic/multiprocessing/api.py:1012] Sending process 223 closing signal SIGTERM
+Traceback (most recent call last):
+  File "/mloscratch/homes/navasard/conda/bin/torchrun", line 6, in <module>
+    sys.exit(main())
+             ~~~~^^
+  File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 362, in wrapper
+    return f(*args, **kwargs)
+  File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/distributed/run.py", line 990, in main
+    run(args)
+    ~~~^^^^^^
+  File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/distributed/run.py", line 981, in run
+    elastic_launch(
+    ~~~~~~~~~~~~~~~
+        config=config,
+        ~~~~~~~~~~~~~~
+        entrypoint=cmd,
+        ~~~~~~~~~~~~~~~
+    )(*cmd_args)
+    ~^^^^^^^^^^^
+  File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/distributed/launcher/api.py", line 170, in __call__
+    return launch_agent(self._config, self._entrypoint, list(args))
+  File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/distributed/launcher/api.py", line 308, in launch_agent
+    result = agent.run()
+  File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/distributed/elastic/metrics/api.py", line 134, in wrapper
+    result = f(*args, **kwargs)
+  File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/distributed/elastic/agent/server/api.py", line 731, in run
+    result = self._invoke_run(role)
+  File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/distributed/elastic/agent/server/api.py", line 908, in _invoke_run
+    time.sleep(monitor_interval)
+    ~~~~~~~~~~^^^^^^^^^^^^^^^^^^
+  File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/distributed/elastic/multiprocessing/api.py", line 86, in _terminate_process_handler
+    raise SignalException(f"Process {os.getpid()} got signal: {sigval}", sigval=sigval)
+torch.distributed.elastic.multiprocessing.api.SignalException: Process 151 got signal: 15
+W0711 08:33:54.614000 2252 site-packages/torch/distributed/run.py:851] 
+W0711 08:33:54.614000 2252 site-packages/torch/distributed/run.py:851] *****************************************
+W0711 08:33:54.614000 2252 site-packages/torch/distributed/run.py:851] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. 
+W0711 08:33:54.614000 2252 site-packages/torch/distributed/run.py:851] *****************************************
+New run: Checkpoints will be saved to /mloscratch/homes/navasard/sparse/checkpoints/l8k1_2b_moba_matched/run_20260711_083403
+Running config:
+{
+  "model_path": "/mloscratch/homes/navasard/qwen3_06b_base",
+  "model_dtype": "bfloat16",
+  "init_from_scratch": true,
+  "train_path": "/mloscratch/homes/navasard/data/fineweb_large/tokenized",
+  "val_path": null,
+  "max_seq_length": 8192,
+  "num_workers": 8,
+  "output_dir": "/mloscratch/homes/navasard/sparse/checkpoints/l8k1_2b_moba_matched/run_20260711_083403",
+  "resume_from_checkpoint": null,
+  "force_new_run": true,
+  "num_epochs": 1,
+  "max_steps": null,
+  "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": 10,
+  "seed": 42,
+  "verbose": true,
+  "checkpoint_activations": false,
+  "strategy": "ddp",
+  "mixed_precision": "bf16",
+  "wandb_enabled": true,
+  "wandb_project": "stochastic_sparsity",
+  "wandb_name": "1.2b moba 8k MATCHED",
+  "wandb_tags": [
+    "qwen3",
+    "1.2b",
+    "8k",
+    "moba",
+    "matched"
+  ],
+  "hidden_size": 1536,
+  "intermediate_size": 6144,
+  "num_attention_heads": 12,
+  "num_key_value_heads": 4,
+  "num_hidden_layers": null,
+  "head_dim": null,
+  "attention_type": "moba",
+  "sparse_block_size": 128,
+  "sparse_num_random_blocks": 1,
+  "sparse_mask_pool_size": 1000,
+  "sparse_val_dense_compare": false,
+  "sparse_layer_indices": null,
+  "sparse_mask_cache_dir": null,
+  "sparse_mask_seed": 0,
+  "dense_mode": false,
+  "dense_step_intervals": null,
+  "moba_block_size": 128,
+  "moba_top_k": 32,
+  "asentmax_alpha": 1.5,
+  "asentmax_learn_temp": true,
+  "asentmax_init_temp": 1.0
+}
+Loading tokenizer from /mloscratch/homes/navasard/qwen3_06b_base...
+Checkpoint activations: False
+Checkpoint activations: False
+Initializing model from scratch (random weights)...
+Checkpoint activations: False
+
+
+================================================================================
+MODEL ARCHITECTURE
+================================================================================
+Model config:
+  vocab_size: 151936
+  hidden_size: 1536
+  num_hidden_layers: 28
+  num_attention_heads: 12
+  num_key_value_heads: 4
+  head_dim: 128
+  intermediate_size: 6144
+  max_position_embeddings: 32768
+  rope_theta: 1000000
+  attention_type: moba
+
+Parameter count:
+  Total parameters: 1,202,352,640
+  Trainable parameters: 1,202,352,640
+  Non-trainable parameters: 0
+  Model size (BF16): 2.40 GB
+
+--------------------------------------------------------------------------------
+LAYER-BY-LAYER SUMMARY
+--------------------------------------------------------------------------------
+/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/torch/nn/attention/flex_attention.py:1704: UserWarning: flex_attention called without torch.compile() - this will use an unfused implementation that materializes the full scores matrix instead of generating a fused kernel.
+
+SOLUTION: Use torch.compile(flex_attention)(...)
+
+If you want to debug your score_mod/mask_mod, you can set:
+torch.nn.attention.flex_attention._FLEX_ATTENTION_DISABLE_COMPILE_DEBUG = True
+
+This will allow you to use print statements or breakpoints. Note: This doesn't work with the backwards pass and may produce incorrect results.
+  _warn_once(
+===========================================================================================================================================================
+Layer (type:depth-idx)                                  Input Shape               Output Shape              Param #                   Trainable
+===========================================================================================================================================================
+Qwen3ForCausalLM                                        [2, 8192]                 --                        --                        True
+├─Qwen3Model: 1-1                                       --                        [2, 8192, 1536]           --                        True
+│    └─Embedding: 2-1                                   [2, 8192]                 [2, 8192, 1536]           233,373,696               True
+│    └─ModuleList: 2-2                                  --                        --                        --                        True
+│    │    └─Qwen3DecoderLayer: 3-1                      [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-2                      [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-3                      [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-4                      [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-5                      [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-6                      [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-7                      [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-8                      [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-9                      [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-10                     [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-11                     [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-12                     [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-13                     [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-14                     [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-15                     [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-16                     [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-17                     [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-18                     [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-19                     [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-20                     [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-21                     [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-22                     [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-23                     [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-24                     [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-25                     [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-26                     [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-27                     [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-28                     [2, 8192, 1536]           [2, 8192, 1536]           34,606,336                True
+│    └─Qwen3RMSNorm: 2-3                                [2, 8192, 1536]           [2, 8192, 1536]           1,536                     True
+├─Linear: 1-2                                           [2, 8192, 1536]           [2, 8192, 151936]         233,373,696               True
+===========================================================================================================================================================
+Total params: 1,435,726,336
+Trainable params: 1,435,726,336
+Non-trainable params: 0
+Total mult-adds (Units.GIGABYTES): 2.87
+===========================================================================================================================================================
+Input size (MB): 0.13
+Forward/backward pass size (MB): 52437.19
+Params size (MB): 2871.45
+Estimated Total Size (MB): 55308.77
+===========================================================================================================================================================
+  NOTE: torchinfo double-counts tied embed_tokens/lm_head weight (233,373,696 params). True unique parameter count: 1,202,352,640 (1202.4M)
+================================================================================
+
+Loading dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Using pre-tokenized dataset (no runtime tokenization)
+Loading pre-tokenized HF dataset from /mloscratch/homes/navasard/data/fineweb_large/tokenized...
+Loaded 12748481 pre-tokenized examples (stored_len=2048, pack_factor=4, packed_examples=3187120)
+Loaded 12748481 pre-tokenized examples (stored_len=2048, pack_factor=4, packed_examples=3187120)
+Loaded 12748481 pre-tokenized examples (stored_len=2048, pack_factor=4, packed_examples=3187120)
+Dataset split: 3027764 train, 159356 validation
+Training examples (packed, train split):  3027764
+Validation examples (packed, val split):  159356
+Per-device batch size:                    2
+World size (GPUs):                        3
+Gradient accumulation steps:              10
+Effective batch size (seqs/step):         60
+Sequence length:                          8192
+Tokens per optimizer step:                491,520
+Micro-batches per epoch (per rank):       504627
+Optimizer steps per epoch:                50462
+Total optimizer steps:                    50462
+wandb: [wandb.login()] Loaded credentials for https://api.wandb.ai from WANDB_API_KEY.
+wandb: Currently logged in as: ruben-navasardyan4869 (ruben-navasardyan4869-epfl) to https://api.wandb.ai. Use `wandb login --relogin` to force relogin
+wandb: Tracking run with wandb version 0.26.1
+wandb: Run data is saved locally in /mloscratch/homes/navasard/sparse/SparseAtt-baselines/wandb/run-20260711_083448-np2dm5av
+wandb: Run `wandb offline` to turn off syncing.
+wandb: Syncing run 1.2b moba 8k MATCHED
+wandb: ⭐️ View project at https://wandb.ai/ruben-navasardyan4869-epfl/stochastic_sparsity
+wandb: 🚀 View run at https://wandb.ai/ruben-navasardyan4869-epfl/stochastic_sparsity/runs/np2dm5av
+Starting training...
+Switching attention mode for step 1: sparse_stochastic (configured sparse attention)
+Step 10/50462 | Loss: 12.2712 | LR: 6.67e-06 | 12968.6ms/step | 4.6 samples/s
+Step 20/50462 | Loss: 11.9930 | LR: 1.33e-05 | 3571.4ms/step | 16.8 samples/s
+Step 30/50462 | Loss: 11.1397 | LR: 2.00e-05 | 3579.3ms/step | 16.8 samples/s
+Step 40/50462 | Loss: 10.2173 | LR: 2.67e-05 | 3572.9ms/step | 16.8 samples/s
+Step 50/50462 | Loss: 9.7769 | LR: 3.33e-05 | 3566.9ms/step | 16.8 samples/s
+Step 60/50462 | Loss: 9.2959 | LR: 4.00e-05 | 3560.8ms/step | 16.9 samples/s
+Step 70/50462 | Loss: 8.7899 | LR: 4.67e-05 | 3570.3ms/step | 16.8 samples/s
+Step 80/50462 | Loss: 8.3037 | LR: 5.33e-05 | 3567.6ms/step | 16.8 samples/s
+Step 90/50462 | Loss: 7.8821 | LR: 6.00e-05 | 3566.0ms/step | 16.8 samples/s
+Step 100/50462 | Loss: 7.5261 | LR: 6.67e-05 | 3573.5ms/step | 16.8 samples/s
+Step 110/50462 | Loss: 7.2710 | LR: 7.33e-05 | 3579.1ms/step | 16.8 samples/s
+Step 120/50462 | Loss: 7.1058 | LR: 8.00e-05 | 3579.1ms/step | 16.8 samples/s
+Step 130/50462 | Loss: 6.9950 | LR: 8.67e-05 | 3582.0ms/step | 16.8 samples/s
+Step 140/50462 | Loss: 6.8790 | LR: 9.33e-05 | 3588.8ms/step | 16.7 samples/s
+Step 150/50462 | Loss: 6.7455 | LR: 1.00e-04 | 3597.5ms/step | 16.7 samples/s
+Step 160/50462 | Loss: 6.6286 | LR: 1.07e-04 | 3600.3ms/step | 16.7 samples/s
+Step 170/50462 | Loss: 6.5159 | LR: 1.13e-04 | 3599.2ms/step | 16.7 samples/s
+Step 180/50462 | Loss: 6.4271 | LR: 1.20e-04 | 3592.4ms/step | 16.7 samples/s
+Step 190/50462 | Loss: 6.3586 | LR: 1.27e-04 | 3593.7ms/step | 16.7 samples/s
+Step 200/50462 | Loss: 6.2810 | LR: 1.33e-04 | 3591.4ms/step | 16.7 samples/s
+Step 210/50462 | Loss: 6.1988 | LR: 1.40e-04 | 3591.1ms/step | 16.7 samples/s
+Step 220/50462 | Loss: 6.1356 | LR: 1.47e-04 | 3589.5ms/step | 16.7 samples/s
+Step 230/50462 | Loss: 6.1218 | LR: 1.53e-04 | 3591.2ms/step | 16.7 samples/s
+Step 240/50462 | Loss: 6.0742 | LR: 1.60e-04 | 3590.5ms/step | 16.7 samples/s
+Step 250/50462 | Loss: 6.0059 | LR: 1.67e-04 | 3586.2ms/step | 16.7 samples/s
+Validation at step 250: Loss = 5.9593, Perplexity = 387.33
+Step 260/50462 | Loss: 5.9316 | LR: 1.73e-04 | 14395.9ms/step | 4.2 samples/s
+Step 270/50462 | Loss: 5.9353 | LR: 1.80e-04 | 3580.0ms/step | 16.8 samples/s
+Step 280/50462 | Loss: 5.8424 | LR: 1.87e-04 | 3590.2ms/step | 16.7 samples/s
+Step 290/50462 | Loss: 5.7822 | LR: 1.93e-04 | 3594.1ms/step | 16.7 samples/s
+Step 300/50462 | Loss: 5.7254 | LR: 2.00e-04 | 3600.4ms/step | 16.7 samples/s
+Step 310/50462 | Loss: 5.6998 | LR: 2.07e-04 | 3595.1ms/step | 16.7 samples/s
+Step 320/50462 | Loss: 5.6249 | LR: 2.13e-04 | 3598.1ms/step | 16.7 samples/s
+Step 330/50462 | Loss: 5.5989 | LR: 2.20e-04 | 3597.8ms/step | 16.7 samples/s
+Step 340/50462 | Loss: 5.5498 | LR: 2.27e-04 | 3598.5ms/step | 16.7 samples/s
+Step 350/50462 | Loss: 5.4828 | LR: 2.33e-04 | 3601.6ms/step | 16.7 samples/s
+Step 360/50462 | Loss: 5.4261 | LR: 2.40e-04 | 3599.5ms/step | 16.7 samples/s
+Step 370/50462 | Loss: 5.3679 | LR: 2.47e-04 | 3601.4ms/step | 16.7 samples/s
+Step 380/50462 | Loss: 5.3380 | LR: 2.53e-04 | 3603.7ms/step | 16.6 samples/s
+Step 390/50462 | Loss: 5.2916 | LR: 2.60e-04 | 3597.7ms/step | 16.7 samples/s
+Step 400/50462 | Loss: 5.2466 | LR: 2.67e-04 | 3595.8ms/step | 16.7 samples/s
+Step 410/50462 | Loss: 5.2134 | LR: 2.73e-04 | 3600.0ms/step | 16.7 samples/s
+Step 420/50462 | Loss: 5.1941 | LR: 2.80e-04 | 3598.2ms/step | 16.7 samples/s
+Step 430/50462 | Loss: 5.1166 | LR: 2.87e-04 | 3598.1ms/step | 16.7 samples/s
+Step 440/50462 | Loss: 5.0694 | LR: 2.93e-04 | 3598.0ms/step | 16.7 samples/s
+Step 450/50462 | Loss: 5.0540 | LR: 3.00e-04 | 3600.0ms/step | 16.7 samples/s
+Step 460/50462 | Loss: 5.0197 | LR: 3.07e-04 | 3597.6ms/step | 16.7 samples/s
+Step 470/50462 | Loss: 4.9741 | LR: 3.13e-04 | 3602.0ms/step | 16.7 samples/s
+Step 480/50462 | Loss: 4.9311 | LR: 3.20e-04 | 3599.7ms/step | 16.7 samples/s
+Step 490/50462 | Loss: 4.8841 | LR: 3.27e-04 | 3598.3ms/step | 16.7 samples/s
+Step 500/50462 | Loss: 4.9094 | LR: 3.33e-04 | 3599.8ms/step | 16.7 samples/s
+Validation at step 500: Loss = 4.8570, Perplexity = 128.64
+Step 510/50462 | Loss: 4.8544 | LR: 3.40e-04 | 5608.8ms/step | 10.7 samples/s
+Step 520/50462 | Loss: 4.8026 | LR: 3.47e-04 | 3597.5ms/step | 16.7 samples/s
+Step 530/50462 | Loss: 4.7795 | LR: 3.53e-04 | 3601.8ms/step | 16.7 samples/s
+Step 540/50462 | Loss: 4.7071 | LR: 3.60e-04 | 3599.9ms/step | 16.7 samples/s
+Step 550/50462 | Loss: 4.7569 | LR: 3.67e-04 | 3600.3ms/step | 16.7 samples/s
+Step 560/50462 | Loss: 4.7164 | LR: 3.73e-04 | 3598.0ms/step | 16.7 samples/s
+Step 570/50462 | Loss: 4.6124 | LR: 3.80e-04 | 3600.2ms/step | 16.7 samples/s
+Step 580/50462 | Loss: 4.5380 | LR: 3.87e-04 | 3604.7ms/step | 16.6 samples/s
+Step 590/50462 | Loss: 4.6199 | LR: 3.93e-04 | 3602.9ms/step | 16.7 samples/s
+Step 600/50462 | Loss: 4.6552 | LR: 4.00e-04 | 3601.8ms/step | 16.7 samples/s
+Step 610/50462 | Loss: 4.5168 | LR: 4.07e-04 | 3600.2ms/step | 16.7 samples/s
+Step 620/50462 | Loss: 4.4920 | LR: 4.13e-04 | 3600.4ms/step | 16.7 samples/s
+Step 630/50462 | Loss: 4.4789 | LR: 4.20e-04 | 3599.4ms/step | 16.7 samples/s
+Step 640/50462 | Loss: 4.4302 | LR: 4.27e-04 | 3600.3ms/step | 16.7 samples/s
+Step 650/50462 | Loss: 4.3546 | LR: 4.33e-04 | 3599.6ms/step | 16.7 samples/s
+Step 660/50462 | Loss: 4.3508 | LR: 4.40e-04 | 3598.3ms/step | 16.7 samples/s
+Step 670/50462 | Loss: 4.4010 | LR: 4.47e-04 | 3600.3ms/step | 16.7 samples/s
+Step 680/50462 | Loss: 4.3196 | LR: 4.53e-04 | 3599.8ms/step | 16.7 samples/s
+Step 690/50462 | Loss: 4.2209 | LR: 4.60e-04 | 3600.0ms/step | 16.7 samples/s
+Step 700/50462 | Loss: 4.2721 | LR: 4.67e-04 | 3597.7ms/step | 16.7 samples/s
+Step 710/50462 | Loss: 4.2820 | LR: 4.73e-04 | 3600.0ms/step | 16.7 samples/s
+Step 720/50462 | Loss: 4.2123 | LR: 4.80e-04 | 3599.8ms/step | 16.7 samples/s
+Step 730/50462 | Loss: 4.1674 | LR: 4.87e-04 | 3600.5ms/step | 16.7 samples/s
+Step 740/50462 | Loss: 4.1267 | LR: 4.93e-04 | 3599.8ms/step | 16.7 samples/s
+Step 750/50462 | Loss: 4.1854 | LR: 5.00e-04 | 3599.9ms/step | 16.7 samples/s
+Validation at step 750: Loss = 4.1339, Perplexity = 62.42
+Step 760/50462 | Loss: 4.0912 | LR: 5.00e-04 | 5608.0ms/step | 10.7 samples/s
+Step 770/50462 | Loss: 4.0943 | LR: 5.00e-04 | 3597.8ms/step | 16.7 samples/s
+Step 780/50462 | Loss: 4.0769 | LR: 5.00e-04 | 3600.3ms/step | 16.7 samples/s
+Step 790/50462 | Loss: 4.0613 | LR: 5.00e-04 | 3608.5ms/step | 16.6 samples/s
+Step 800/50462 | Loss: 3.9934 | LR: 5.00e-04 | 3602.9ms/step | 16.7 samples/s
+Step 810/50462 | Loss: 4.0156 | LR: 5.00e-04 | 3600.6ms/step | 16.7 samples/s
+Step 820/50462 | Loss: 3.9963 | LR: 5.00e-04 | 3600.2ms/step | 16.7 samples/s
+Step 830/50462 | Loss: 3.9618 | LR: 5.00e-04 | 3599.6ms/step | 16.7 samples/s
+Step 840/50462 | Loss: 3.9328 | LR: 5.00e-04 | 3600.0ms/step | 16.7 samples/s
+Step 850/50462 | Loss: 3.9050 | LR: 5.00e-04 | 3604.2ms/step | 16.6 samples/s
+Step 860/50462 | Loss: 3.8963 | LR: 5.00e-04 | 3599.4ms/step | 16.7 samples/s
+Step 870/50462 | Loss: 3.8750 | LR: 5.00e-04 | 3600.5ms/step | 16.7 samples/s
+Step 880/50462 | Loss: 3.8397 | LR: 5.00e-04 | 3600.0ms/step | 16.7 samples/s
+Step 890/50462 | Loss: 3.8519 | LR: 5.00e-04 | 3599.8ms/step | 16.7 samples/s
+Step 900/50462 | Loss: 3.8371 | LR: 5.00e-04 | 3601.9ms/step | 16.7 samples/s
+Step 910/50462 | Loss: 3.8082 | LR: 5.00e-04 | 3600.8ms/step | 16.7 samples/s
+Step 920/50462 | Loss: 3.7720 | LR: 5.00e-04 | 3601.3ms/step | 16.7 samples/s
+Step 930/50462 | Loss: 3.8192 | LR: 5.00e-04 | 3599.8ms/step | 16.7 samples/s
+Step 940/50462 | Loss: 3.7628 | LR: 5.00e-04 | 3600.3ms/step | 16.7 samples/s
+Step 950/50462 | Loss: 3.7498 | LR: 5.00e-04 | 3597.6ms/step | 16.7 samples/s
+Step 960/50462 | Loss: 3.7532 | LR: 5.00e-04 | 3597.8ms/step | 16.7 samples/s
+Step 970/50462 | Loss: 3.7323 | LR: 5.00e-04 | 3598.0ms/step | 16.7 samples/s
+Step 980/50462 | Loss: 3.7033 | LR: 5.00e-04 | 3600.3ms/step | 16.7 samples/s
+Step 990/50462 | Loss: 3.7081 | LR: 5.00e-04 | 3598.0ms/step | 16.7 samples/s
+Step 1000/50462 | Loss: 3.7120 | LR: 5.00e-04 | 3599.2ms/step | 16.7 samples/s
+Validation at step 1000: Loss = 3.7141, Perplexity = 41.02
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l8k1_2b_moba_matched/run_20260711_083403/checkpoint-1000
+Step 1010/50462 | Loss: 3.6907 | LR: 5.00e-04 | 6598.2ms/step | 9.1 samples/s
+Step 1020/50462 | Loss: 3.6956 | LR: 5.00e-04 | 3595.8ms/step | 16.7 samples/s
+Step 1030/50462 | Loss: 3.6859 | LR: 5.00e-04 | 3597.8ms/step | 16.7 samples/s
+Step 1040/50462 | Loss: 3.6855 | LR: 5.00e-04 | 3597.9ms/step | 16.7 samples/s
+Step 1050/50462 | Loss: 3.6536 | LR: 5.00e-04 | 3596.3ms/step | 16.7 samples/s
+Step 1060/50462 | Loss: 3.6436 | LR: 5.00e-04 | 3599.9ms/step | 16.7 samples/s
+Step 1070/50462 | Loss: 3.6381 | LR: 5.00e-04 | 3596.5ms/step | 16.7 samples/s
+Step 1080/50462 | Loss: 3.6411 | LR: 5.00e-04 | 3598.0ms/step | 16.7 samples/s
+Step 1090/50462 | Loss: 3.6170 | LR: 5.00e-04 | 3599.9ms/step | 16.7 samples/s
+Step 1100/50462 | Loss: 3.6264 | LR: 5.00e-04 | 3595.3ms/step | 16.7 samples/s
+Step 1110/50462 | Loss: 3.6353 | LR: 5.00e-04 | 3594.7ms/step | 16.7 samples/s
+Step 1120/50462 | Loss: 3.5931 | LR: 5.00e-04 | 3600.1ms/step | 16.7 samples/s
+Step 1130/50462 | Loss: 3.6077 | LR: 5.00e-04 | 3601.6ms/step | 16.7 samples/s
+Step 1140/50462 | Loss: 3.5691 | LR: 5.00e-04 | 3598.5ms/step | 16.7 samples/s
+Step 1150/50462 | Loss: 3.5363 | LR: 5.00e-04 | 3599.9ms/step | 16.7 samples/s
+Step 1160/50462 | Loss: 3.5617 | LR: 5.00e-04 | 3597.8ms/step | 16.7 samples/s
+Step 1170/50462 | Loss: 3.5569 | LR: 5.00e-04 | 3594.9ms/step | 16.7 samples/s
+Step 1180/50462 | Loss: 3.5224 | LR: 5.00e-04 | 3594.2ms/step | 16.7 samples/s
+Step 1190/50462 | Loss: 3.5387 | LR: 5.00e-04 | 3590.8ms/step | 16.7 samples/s
+Step 1200/50462 | Loss: 3.5554 | LR: 5.00e-04 | 3591.6ms/step | 16.7 samples/s
+Step 1210/50462 | Loss: 3.5518 | LR: 5.00e-04 | 3591.6ms/step | 16.7 samples/s
+Step 1220/50462 | Loss: 3.5126 | LR: 5.00e-04 | 3597.1ms/step | 16.7 samples/s
+Step 1230/50462 | Loss: 3.4851 | LR: 5.00e-04 | 3599.2ms/step | 16.7 samples/s
+Step 1240/50462 | Loss: 3.5415 | LR: 5.00e-04 | 3598.6ms/step | 16.7 samples/s
+Step 1250/50462 | Loss: 3.5026 | LR: 5.00e-04 | 3599.4ms/step | 16.7 samples/s
+Validation at step 1250: Loss = 3.5093, Perplexity = 33.43
+Step 1260/50462 | Loss: 3.5206 | LR: 5.00e-04 | 5604.4ms/step | 10.7 samples/s
+Step 1270/50462 | Loss: 3.4876 | LR: 5.00e-04 | 3597.3ms/step | 16.7 samples/s
+Step 1280/50462 | Loss: 3.4640 | LR: 5.00e-04 | 3591.8ms/step | 16.7 samples/s
+Step 1290/50462 | Loss: 3.4695 | LR: 5.00e-04 | 3594.3ms/step | 16.7 samples/s
+Step 1300/50462 | Loss: 3.4633 | LR: 5.00e-04 | 3592.3ms/step | 16.7 samples/s
+Step 1310/50462 | Loss: 3.4739 | LR: 5.00e-04 | 3593.6ms/step | 16.7 samples/s
+Step 1320/50462 | Loss: 3.4637 | LR: 5.00e-04 | 3589.2ms/step | 16.7 samples/s
+Step 1330/50462 | Loss: 3.4476 | LR: 5.00e-04 | 3582.4ms/step | 16.7 samples/s
+Step 1340/50462 | Loss: 3.4441 | LR: 5.00e-04 | 3577.4ms/step | 16.8 samples/s
+Step 1350/50462 | Loss: 3.4234 | LR: 5.00e-04 | 3584.4ms/step | 16.7 samples/s
+Step 1360/50462 | Loss: 3.4439 | LR: 5.00e-04 | 3586.6ms/step | 16.7 samples/s
+Step 1370/50462 | Loss: 3.4225 | LR: 5.00e-04 | 3584.9ms/step | 16.7 samples/s
+Step 1380/50462 | Loss: 3.4582 | LR: 5.00e-04 | 3584.0ms/step | 16.7 samples/s
+Step 1390/50462 | Loss: 3.4324 | LR: 5.00e-04 | 3585.2ms/step | 16.7 samples/s
+Step 1400/50462 | Loss: 3.4035 | LR: 5.00e-04 | 3587.8ms/step | 16.7 samples/s
+Step 1410/50462 | Loss: 3.4300 | LR: 5.00e-04 | 3595.6ms/step | 16.7 samples/s
+Step 1420/50462 | Loss: 3.4085 | LR: 5.00e-04 | 3592.7ms/step | 16.7 samples/s
+Step 1430/50462 | Loss: 3.3840 | LR: 5.00e-04 | 3591.7ms/step | 16.7 samples/s
+Step 1440/50462 | Loss: 3.4178 | LR: 5.00e-04 | 3588.2ms/step | 16.7 samples/s
+Step 1450/50462 | Loss: 3.3829 | LR: 5.00e-04 | 3584.7ms/step | 16.7 samples/s
+Step 1460/50462 | Loss: 3.3911 | LR: 5.00e-04 | 3577.6ms/step | 16.8 samples/s
+Step 1470/50462 | Loss: 3.3637 | LR: 5.00e-04 | 3576.0ms/step | 16.8 samples/s
+Step 1480/50462 | Loss: 3.4023 | LR: 5.00e-04 | 3580.3ms/step | 16.8 samples/s
+Step 1490/50462 | Loss: 3.3860 | LR: 5.00e-04 | 3584.9ms/step | 16.7 samples/s
+Step 1500/50462 | Loss: 3.3518 | LR: 5.00e-04 | 3574.9ms/step | 16.8 samples/s
+Validation at step 1500: Loss = 3.3874, Perplexity = 29.59
+Step 1510/50462 | Loss: 3.3837 | LR: 5.00e-04 | 5570.7ms/step | 10.8 samples/s
+Step 1520/50462 | Loss: 3.3587 | LR: 5.00e-04 | 3581.8ms/step | 16.8 samples/s
+Step 1530/50462 | Loss: 3.3830 | LR: 5.00e-04 | 3582.0ms/step | 16.8 samples/s
+Step 1540/50462 | Loss: 3.3505 | LR: 5.00e-04 | 3583.8ms/step | 16.7 samples/s
+Step 1550/50462 | Loss: 3.3593 | LR: 5.00e-04 | 3590.7ms/step | 16.7 samples/s
+Step 1560/50462 | Loss: 3.3496 | LR: 5.00e-04 | 3589.5ms/step | 16.7 samples/s
+Step 1570/50462 | Loss: 3.3834 | LR: 5.00e-04 | 3577.9ms/step | 16.8 samples/s
+Step 1580/50462 | Loss: 3.3455 | LR: 5.00e-04 | 3568.6ms/step | 16.8 samples/s
+Step 1590/50462 | Loss: 3.3703 | LR: 5.00e-04 | 3576.1ms/step | 16.8 samples/s
+Step 1600/50462 | Loss: 3.3566 | LR: 5.00e-04 | 3582.4ms/step | 16.7 samples/s
+Step 1610/50462 | Loss: 3.3219 | LR: 5.00e-04 | 3586.3ms/step | 16.7 samples/s
+Step 1620/50462 | Loss: 3.3366 | LR: 5.00e-04 | 3590.1ms/step | 16.7 samples/s
+Step 1630/50462 | Loss: 3.3573 | LR: 5.00e-04 | 3593.1ms/step | 16.7 samples/s
+Step 1640/50462 | Loss: 3.3332 | LR: 5.00e-04 | 3593.8ms/step | 16.7 samples/s
+Step 1650/50462 | Loss: 3.2924 | LR: 5.00e-04 | 3598.6ms/step | 16.7 samples/s
+Step 1660/50462 | Loss: 3.3095 | LR: 5.00e-04 | 3598.0ms/step | 16.7 samples/s
+Step 1670/50462 | Loss: 3.3341 | LR: 5.00e-04 | 3597.0ms/step | 16.7 samples/s
+Step 1680/50462 | Loss: 3.3081 | LR: 5.00e-04 | 3598.8ms/step | 16.7 samples/s
+Step 1690/50462 | Loss: 3.3121 | LR: 5.00e-04 | 3600.0ms/step | 16.7 samples/s
+Step 1700/50462 | Loss: 3.3180 | LR: 5.00e-04 | 3595.9ms/step | 16.7 samples/s
+Step 1710/50462 | Loss: 3.3349 | LR: 5.00e-04 | 3595.7ms/step | 16.7 samples/s
+Step 1720/50462 | Loss: 3.2958 | LR: 5.00e-04 | 3598.0ms/step | 16.7 samples/s
+Step 1730/50462 | Loss: 3.2953 | LR: 5.00e-04 | 3596.2ms/step | 16.7 samples/s
+Step 1740/50462 | Loss: 3.3261 | LR: 5.00e-04 | 3597.7ms/step | 16.7 samples/s
+Step 1750/50462 | Loss: 3.3392 | LR: 5.00e-04 | 3601.0ms/step | 16.7 samples/s
+Validation at step 1750: Loss = 3.3021, Perplexity = 27.17
+Step 1760/50462 | Loss: 3.3127 | LR: 4.99e-04 | 5609.9ms/step | 10.7 samples/s
+Step 1770/50462 | Loss: 3.2974 | LR: 4.99e-04 | 3599.5ms/step | 16.7 samples/s
+Step 1780/50462 | Loss: 3.2681 | LR: 4.99e-04 | 3604.2ms/step | 16.6 samples/s
+Step 1790/50462 | Loss: 3.2669 | LR: 4.99e-04 | 3600.3ms/step | 16.7 samples/s
+Step 1800/50462 | Loss: 3.2636 | LR: 4.99e-04 | 3603.3ms/step | 16.7 samples/s
+Step 1810/50462 | Loss: 3.2759 | LR: 4.99e-04 | 3602.6ms/step | 16.7 samples/s
+Step 1820/50462 | Loss: 3.2543 | LR: 4.99e-04 | 3606.1ms/step | 16.6 samples/s
+Step 1830/50462 | Loss: 3.2597 | LR: 4.99e-04 | 3608.0ms/step | 16.6 samples/s
+Step 1840/50462 | Loss: 3.2552 | LR: 4.99e-04 | 3619.5ms/step | 16.6 samples/s
+Step 1850/50462 | Loss: 3.2674 | LR: 4.99e-04 | 3608.9ms/step | 16.6 samples/s
+Step 1860/50462 | Loss: 3.2454 | LR: 4.99e-04 | 3606.5ms/step | 16.6 samples/s
+Step 1870/50462 | Loss: 3.2896 | LR: 4.99e-04 | 3603.3ms/step | 16.7 samples/s
+Step 1880/50462 | Loss: 3.2655 | LR: 4.99e-04 | 3609.7ms/step | 16.6 samples/s
+Step 1890/50462 | Loss: 3.2332 | LR: 4.99e-04 | 3607.9ms/step | 16.6 samples/s
+Step 1900/50462 | Loss: 3.2347 | LR: 4.99e-04 | 3605.3ms/step | 16.6 samples/s
+Step 1910/50462 | Loss: 3.2692 | LR: 4.99e-04 | 3602.0ms/step | 16.7 samples/s
+Step 1920/50462 | Loss: 3.2690 | LR: 4.99e-04 | 3606.8ms/step | 16.6 samples/s
+Step 1930/50462 | Loss: 3.2722 | LR: 4.99e-04 | 3604.1ms/step | 16.6 samples/s
+Step 1940/50462 | Loss: 3.2589 | LR: 4.99e-04 | 3614.1ms/step | 16.6 samples/s
+Step 1950/50462 | Loss: 3.2590 | LR: 4.99e-04 | 3618.4ms/step | 16.6 samples/s
+Step 1960/50462 | Loss: 3.2497 | LR: 4.99e-04 | 3604.8ms/step | 16.6 samples/s
+Step 1970/50462 | Loss: 3.2488 | LR: 4.99e-04 | 3618.7ms/step | 16.6 samples/s
+Step 1980/50462 | Loss: 3.2423 | LR: 4.99e-04 | 3615.7ms/step | 16.6 samples/s
+Step 1990/50462 | Loss: 3.2648 | LR: 4.99e-04 | 3601.6ms/step | 16.7 samples/s
+Step 2000/50462 | Loss: 3.2666 | LR: 4.99e-04 | 3602.0ms/step | 16.7 samples/s
+Validation at step 2000: Loss = 3.2394, Perplexity = 25.52
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l8k1_2b_moba_matched/run_20260711_083403/checkpoint-2000
+Step 2010/50462 | Loss: 3.2475 | LR: 4.99e-04 | 6769.6ms/step | 8.9 samples/s
+Step 2020/50462 | Loss: 3.2483 | LR: 4.99e-04 | 3592.9ms/step | 16.7 samples/s
+Step 2030/50462 | Loss: 3.2196 | LR: 4.99e-04 | 3599.0ms/step | 16.7 samples/s
+Step 2040/50462 | Loss: 3.2304 | LR: 4.99e-04 | 3600.3ms/step | 16.7 samples/s
+Step 2050/50462 | Loss: 3.1998 | LR: 4.99e-04 | 3603.8ms/step | 16.6 samples/s
+Step 2060/50462 | Loss: 3.2041 | LR: 4.99e-04 | 3602.0ms/step | 16.7 samples/s
+Step 2070/50462 | Loss: 3.2204 | LR: 4.99e-04 | 3602.2ms/step | 16.7 samples/s
+Step 2080/50462 | Loss: 3.2100 | LR: 4.99e-04 | 3599.6ms/step | 16.7 samples/s
+Step 2090/50462 | Loss: 3.2129 | LR: 4.99e-04 | 3599.8ms/step | 16.7 samples/s
+Step 2100/50462 | Loss: 3.2001 | LR: 4.99e-04 | 3598.7ms/step | 16.7 samples/s
+Step 2110/50462 | Loss: 3.2484 | LR: 4.99e-04 | 3602.2ms/step | 16.7 samples/s
+Step 2120/50462 | Loss: 3.2341 | LR: 4.99e-04 | 3599.5ms/step | 16.7 samples/s
+Step 2130/50462 | Loss: 3.1982 | LR: 4.99e-04 | 3604.1ms/step | 16.6 samples/s
+Step 2140/50462 | Loss: 3.1818 | LR: 4.99e-04 | 3599.4ms/step | 16.7 samples/s
+Step 2150/50462 | Loss: 3.1864 | LR: 4.99e-04 | 3600.3ms/step | 16.7 samples/s
+Step 2160/50462 | Loss: 3.1998 | LR: 4.99e-04 | 3604.2ms/step | 16.6 samples/s
+Step 2170/50462 | Loss: 3.1867 | LR: 4.99e-04 | 3600.0ms/step | 16.7 samples/s
+Step 2180/50462 | Loss: 3.2105 | LR: 4.99e-04 | 3599.2ms/step | 16.7 samples/s
+Step 2190/50462 | Loss: 3.1839 | LR: 4.99e-04 | 3600.5ms/step | 16.7 samples/s
+Step 2200/50462 | Loss: 3.1821 | LR: 4.99e-04 | 3608.4ms/step | 16.6 samples/s
+Step 2210/50462 | Loss: 3.1918 | LR: 4.99e-04 | 3610.3ms/step | 16.6 samples/s
+Step 2220/50462 | Loss: 3.1735 | LR: 4.99e-04 | 3610.1ms/step | 16.6 samples/s
+Step 2230/50462 | Loss: 3.1899 | LR: 4.99e-04 | 3613.1ms/step | 16.6 samples/s
+Step 2240/50462 | Loss: 3.1827 | LR: 4.99e-04 | 3602.4ms/step | 16.7 samples/s
+Step 2250/50462 | Loss: 3.2222 | LR: 4.99e-04 | 3599.7ms/step | 16.7 samples/s
+Validation at step 2250: Loss = 3.1891, Perplexity = 24.27
+Step 2260/50462 | Loss: 3.1850 | LR: 4.99e-04 | 5615.7ms/step | 10.7 samples/s
+Step 2270/50462 | Loss: 3.1964 | LR: 4.99e-04 | 3600.5ms/step | 16.7 samples/s
+Step 2280/50462 | Loss: 3.1485 | LR: 4.99e-04 | 3601.5ms/step | 16.7 samples/s
+Step 2290/50462 | Loss: 3.1513 | LR: 4.99e-04 | 3606.4ms/step | 16.6 samples/s
+Step 2300/50462 | Loss: 3.1533 | LR: 4.99e-04 | 3602.5ms/step | 16.7 samples/s
+Step 2310/50462 | Loss: 3.1740 | LR: 4.99e-04 | 3604.0ms/step | 16.6 samples/s
+Step 2320/50462 | Loss: 3.1694 | LR: 4.99e-04 | 3599.5ms/step | 16.7 samples/s
+Step 2330/50462 | Loss: 3.1301 | LR: 4.99e-04 | 3601.6ms/step | 16.7 samples/s
+Step 2340/50462 | Loss: 3.1658 | LR: 4.99e-04 | 3600.2ms/step | 16.7 samples/s
+Step 2350/50462 | Loss: 3.1394 | LR: 4.99e-04 | 3603.7ms/step | 16.6 samples/s
+Step 2360/50462 | Loss: 3.1387 | LR: 4.99e-04 | 3600.6ms/step | 16.7 samples/s
+Step 2370/50462 | Loss: 3.1607 | LR: 4.99e-04 | 3601.2ms/step | 16.7 samples/s
+Step 2380/50462 | Loss: 3.1080 | LR: 4.99e-04 | 3600.1ms/step | 16.7 samples/s
+Step 2390/50462 | Loss: 3.1344 | LR: 4.99e-04 | 3600.0ms/step | 16.7 samples/s
+Step 2400/50462 | Loss: 3.1579 | LR: 4.99e-04 | 3599.9ms/step | 16.7 samples/s
+Step 2410/50462 | Loss: 3.1229 | LR: 4.99e-04 | 3598.2ms/step | 16.7 samples/s
+Step 2420/50462 | Loss: 3.1621 | LR: 4.99e-04 | 3599.7ms/step | 16.7 samples/s

epfl_final/train_logs/watchdog_generic.sh

@@ -0,0 +1,79 @@
+#!/bin/bash
+# Persistent in-pod training watchdog.
+# Relaunches torchrun (auto-resumes from latest checkpoint) ONLY when the run is
+# truly dead: GPU compute apps == 0 AND no matching train procs, sustained.
+# Exits (stops watching) once the run has completed.
+# Args: JOB CONFIG OUTPUT_DIR PROC_PAT TOTAL [NPROC]
+set -u
+JOB="$1"; CFG="$2"; OUTDIR="$3"; PROC_PAT="$4"; TOTAL="$5"; NPROC="${6:-4}"
+BASE=/mloscratch/homes/navasard/sparse
+SPA="$BASE/SparseAtt"
+TR=/mloscratch/homes/navasard/conda/bin/torchrun
+export HF_HOME=/mloscratch/homes/navasard/hf_cache
+WLOG="$BASE/train_logs/wd_${JOB}.watch.log"
+RUNLOG="$BASE/train_logs/wd_${JOB}.run.log"
+LOCK="$BASE/train_logs/wd_${JOB}.lock"
+DOWN_THRESHOLD=10   # x30s = 5 min sustained-dead before relaunch
+
+log(){ echo "$(date -u +%Y-%m-%dT%H:%M:%SZ) $*" >> "$WLOG"; }
+
+# ---- singleton guard ----
+if [ -f "$LOCK" ]; then
+  oldpid=$(cat "$LOCK" 2>/dev/null || true)
+  if [ -n "${oldpid:-}" ] && kill -0 "$oldpid" 2>/dev/null; then
+    log "duplicate start: watchdog already running pid=$oldpid; exiting"; exit 0
+  fi
+fi
+echo $$ > "$LOCK"
+log "=== watchdog start pid=$$ job=$JOB cfg=$CFG total=$TOTAL nproc=$NPROC ==="
+
+latest_ckpt_step(){
+  ls -1d "$OUTDIR"/run_*/checkpoint-* 2>/dev/null \
+    | sed -E 's#.*/checkpoint-##' | grep -E '^[0-9]+$' | sort -n | tail -1
+}
+
+is_complete(){
+  # 1) explicit completion marker in any of this job's run logs
+  if grep -qs "Training complete!" "$BASE"/train_logs/*"${JOB}"*.log 2>/dev/null; then return 0; fi
+  # 2) a checkpoint at/after the final step exists
+  local s; s=$(latest_ckpt_step)
+  if [ -n "${s:-}" ] && [ "$s" -ge "$TOTAL" ] 2>/dev/null; then return 0; fi
+  return 1
+}
+
+n_procs(){ pgrep -f "$PROC_PAT" 2>/dev/null | grep -c '[0-9]'; }
+n_gpu_apps(){ nvidia-smi --query-compute-apps=pid --format=csv,noheader 2>/dev/null | grep -c '[0-9]'; }
+
+relaunch(){
+  cd "$SPA" || { log "ERR cd $SPA failed"; return 1; }
+  log "RELAUNCH torchrun nproc=$NPROC cfg=$CFG (latest ckpt step=$(latest_ckpt_step))"
+  setsid nohup "$TR" --standalone --nproc_per_node="$NPROC" train.py --config "$CFG" \
+      >> "$RUNLOG" 2>&1 < /dev/null &
+  sleep 240
+  if [ "$(n_procs)" -ge 1 ] && { [ "$(n_gpu_apps)" -ge 1 ] \
+       || grep -qE "Step [0-9]+/|Resuming from checkpoint|Starting training|New run:" "$RUNLOG" 2>/dev/null; }; then
+    log "RELAUNCH OK (procs=$(n_procs) gpu_apps=$(n_gpu_apps))"; return 0
+  fi
+  log "RELAUNCH FAILED (procs=$(n_procs) gpu_apps=$(n_gpu_apps)); runlog tail:"
+  tail -n 8 "$RUNLOG" >> "$WLOG" 2>/dev/null
+  return 1
+}
+
+zero=0
+while true; do
+  if is_complete; then log "COMPLETE detected (ckpt step=$(latest_ckpt_step)); watchdog exiting"; rm -f "$LOCK"; exit 0; fi
+  procs=$(n_procs); apps=$(n_gpu_apps)
+  if [ "$procs" -eq 0 ] && [ "$apps" -eq 0 ]; then zero=$((zero+1)); else zero=0; fi
+  if [ "$zero" -ge "$DOWN_THRESHOLD" ]; then
+    log "SUSTAINED DEAD (${DOWN_THRESHOLD}x30s: procs=0 gpu_apps=0)"
+    if is_complete; then log "actually complete; exit"; rm -f "$LOCK"; exit 0; fi
+    ok=0
+    for a in 1 2 3; do
+      if relaunch; then ok=1; break; fi
+      log "relaunch attempt $a failed; retry in 60s"; sleep 60
+    done
+    [ "$ok" -eq 1 ] && log "recovered" || log "!!! RELAUNCH EXHAUSTED - manual action needed !!!"
+    zero=0
+  fi
+  sleep 30
+done

epfl_final/train_logs/wd_dense30k.watch.log

@@ -0,0 +1,2 @@
+2026-07-11T17:46:42Z === watchdog start pid=3542 job=dense30k cfg=configs/train_config_30k_dense_b128.json total=37465 nproc=4 ===
+2026-07-12T20:22:03Z COMPLETE detected (ckpt step=37465); watchdog exiting

epfl_final/train_logs/wd_fixed30k.watch.log

@@ -0,0 +1,2 @@
+2026-07-13T03:17:36Z === watchdog start pid=252 job=fixed30k cfg=configs/train_config_30k_fixed_b128.json total=37465 nproc=4 ===
+2026-07-15T21:44:30Z COMPLETE detected (ckpt step=37465); watchdog exiting

epfl_final/train_logs/wd_resampled30k.watch.log

@@ -0,0 +1,2 @@
+2026-07-11T17:47:03Z === watchdog start pid=6813 job=resampled30k cfg=configs/train_config_30k_resampled_b128.json total=37465 nproc=4 ===
+2026-07-13T19:12:57Z COMPLETE detected (ckpt step=37465); watchdog exiting