EPFL ALiBi and training provenance 601-650 of 1184

epfl_final/logs/ALIBI_20260921/1p2b_s42_deferred.log

@@ -0,0 +1,327 @@
+2026-09-21T23:59:52+00:00 waiting for evaluation PID 1229
+2026-09-22T00:23:22+00:00 evaluation exited; waiting 300 seconds
+2026-09-22T00:36:58+00:00 waiting for evaluation PID 3741
+2026-09-22T00:36:58+00:00 evaluation exited; waiting 300 seconds
+2026-09-22T00:41:58+00:00 launching 1.2B ALiBi seed 42 on 6 GPUs
+W0922 00:42:00.258000 4327 site-packages/torch/distributed/run.py:851] 
+W0922 00:42:00.258000 4327 site-packages/torch/distributed/run.py:851] *****************************************
+W0922 00:42:00.258000 4327 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. 
+W0922 00:42:00.258000 4327 site-packages/torch/distributed/run.py:851] *****************************************
+New run: Checkpoints will be saved to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/1p2b_alibi_s42/run_20260922_004208
+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/1p2b_alibi_s42/run_20260922_004208",
+  "resume_from_checkpoint": null,
+  "force_new_run": false,
+  "num_epochs": 1,
+  "max_steps": 50462,
+  "per_device_batch_size": 1,
+  "gradient_accumulation_steps": 10,
+  "learning_rate": 0.0005,
+  "weight_decay": 0.01,
+  "warmup_steps": 750,
+  "lr_scheduler": "cosine",
+  "max_grad_norm": 1.0,
+  "save_steps": 1000,
+  "eval_steps": 1000,
+  "eval_batches": 50,
+  "logging_steps": 100,
+  "seed": 42,
+  "verbose": true,
+  "checkpoint_activations": false,
+  "resume_micro_step_override": null,
+  "strategy": "ddp",
+  "mixed_precision": "bf16",
+  "wandb_enabled": true,
+  "wandb_project": "stochastic_sparsity",
+  "wandb_name": "1.2b alibi 8k s42",
+  "wandb_tags": [
+    "qwen3",
+    "1.2b",
+    "8k",
+    "alibi",
+    "dense",
+    "seed42"
+  ],
+  "hidden_size": 1536,
+  "intermediate_size": 6144,
+  "num_attention_heads": 12,
+  "num_key_value_heads": 4,
+  "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": 42,
+  "sparse_mixed_nrb_set": null,
+  "dense_mode": false,
+  "dense_step_intervals": null
+}
+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
+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
+  rope_train_position_scale_set: None
+  attention_type: dense
+
+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                              [1, 8192]                 --                        --                        True
+├─Qwen3Model: 1-1                             --                        [1, 8192, 1536]           --                        True
+│    └─Embedding: 2-1                         [1, 8192]                 [1, 8192, 1536]           233,373,696               True
+│    └─ModuleList: 2-2                        --                        --                        --                        True
+│    │    └─Qwen3DecoderLayer: 3-1            [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-2            [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-3            [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-4            [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-5            [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-6            [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-7            [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-8            [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-9            [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-10           [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-11           [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-12           [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-13           [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-14           [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-15           [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-16           [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-17           [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-18           [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-19           [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-20           [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-21           [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-22           [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-23           [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-24           [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-25           [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-26           [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-27           [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-28           [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    └─Qwen3RMSNorm: 2-3                      [1, 8192, 1536]           [1, 8192, 1536]           1,536                     True
+├─Linear: 1-2                                 [1, 8192, 1536]           [1, 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): 1.44
+=================================================================================================================================================
+Input size (MB): 0.07
+Forward/backward pass size (MB): 26218.59
+Params size (MB): 2871.45
+Estimated Total Size (MB): 29090.11
+=================================================================================================================================================
+  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...
+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)
+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):                        6
+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):       504628
+Optimizer steps per epoch:                50462
+Total optimizer steps:                    50462  (capped by max_steps; 1.00 epochs)
+[rank0]: Traceback (most recent call last):
+[rank0]:   File "/mloscratch/homes/navasard/sparse/SparseAtt/train.py", line 1522, in <module>
+[rank0]:     main()
+[rank0]:     ~~~~^^
+[rank0]:   File "/mloscratch/homes/navasard/sparse/SparseAtt/train.py", line 1519, in main
+[rank0]:     train(config)
+[rank0]:     ~~~~~^^^^^^^^
+[rank0]:   File "/mloscratch/homes/navasard/sparse/SparseAtt/train.py", line 1296, in train
+[rank0]:     wandb.init(
+[rank0]:     ~~~~~~~~~~^
+[rank0]:         project=config.wandb_project,
+[rank0]:         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:     ...<2 lines>...
+[rank0]:         config=vars(config),
+[rank0]:         ^^^^^^^^^^^^^^^^^^^^
+[rank0]:     )
+[rank0]:     ^
+[rank0]:   File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/wandb/sdk/wandb_init.py", line 1603, in init
+[rank0]:     get_sentry().reraise(e)
+[rank0]:     ~~~~~~~~~~~~~~~~~~~~^^^
+[rank0]:   File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/wandb/analytics/sentry.py", line 190, in reraise
+[rank0]:     raise exc.with_traceback(tb)
+[rank0]:   File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/wandb/sdk/wandb_init.py", line 1524, in init
+[rank0]:     wi.maybe_login(init_settings)
+[rank0]:     ~~~~~~~~~~~~~~^^^^^^^^^^^^^^^
+[rank0]:   File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/wandb/sdk/wandb_init.py", line 200, in maybe_login
+[rank0]:     wandb_login._login(
+[rank0]:     ~~~~~~~~~~~~~~~~~~^
+[rank0]:         host=run_settings.base_url,
+[rank0]:         ^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:     ...<3 lines>...
+[rank0]:         update_api_key=explicit_key is None,
+[rank0]:         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:     )
+[rank0]:     ^
+[rank0]:   File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/wandb/sdk/wandb_login.py", line 190, in _login
+[rank0]:     auth = _find_or_prompt_for_key(
+[rank0]:         settings,
+[rank0]:     ...<4 lines>...
+[rank0]:         input_timeout=timeout,
+[rank0]:     )
+[rank0]:   File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/wandb/sdk/wandb_login.py", line 266, in _find_or_prompt_for_key
+[rank0]:     auth = wbauth.authenticate_session(
+[rank0]:         host=host,
+[rank0]:     ...<5 lines>...
+[rank0]:         relogin=relogin,
+[rank0]:     )
+[rank0]:   File "/mloscratch/homes/navasard/conda/lib/python3.13/site-packages/wandb/sdk/lib/wbauth/authenticate.py", line 139, in authenticate_session
+[rank0]:     raise UsageError(
+[rank0]:         "No API key configured. Use `wandb login` to log in."
+[rank0]:     ) from None
+[rank0]: wandb.errors.errors.UsageError: No API key configured. Use `wandb login` to log in.
+[rank0]:[W922 00:42:37.829994451 ProcessGroupNCCL.cpp:1575] Warning: WARNING: destroy_process_group() was not called before program exit, which can leak resources. For more info, please see https://pytorch.org/docs/stable/distributed.html#shutdown (function operator())
+W0922 00:42:38.771000 4327 site-packages/torch/distributed/elastic/multiprocessing/api.py:1012] Sending process 4471 closing signal SIGTERM
+W0922 00:42:38.773000 4327 site-packages/torch/distributed/elastic/multiprocessing/api.py:1012] Sending process 4472 closing signal SIGTERM
+W0922 00:42:38.774000 4327 site-packages/torch/distributed/elastic/multiprocessing/api.py:1012] Sending process 4473 closing signal SIGTERM
+W0922 00:42:38.777000 4327 site-packages/torch/distributed/elastic/multiprocessing/api.py:1012] Sending process 4474 closing signal SIGTERM
+W0922 00:42:38.778000 4327 site-packages/torch/distributed/elastic/multiprocessing/api.py:1012] Sending process 4475 closing signal SIGTERM
+E0922 00:42:38.782000 4327 site-packages/torch/distributed/elastic/multiprocessing/api.py:986] failed (exitcode: 1) local_rank: 0 (pid: 4470) of binary: /mloscratch/homes/navasard/conda/bin/python3.13
+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 317, in launch_agent
+    raise ChildFailedError(
+    ...<2 lines>...
+    )
+torch.distributed.elastic.multiprocessing.errors.ChildFailedError: 
+============================================================
+train.py FAILED
+------------------------------------------------------------
+Failures:
+[1]:
+  time      : 2026-09-22_00:42:38
+  host      : h100-workspace-6-gpu-0-59
+  rank      : 1 (local_rank: 1)
+  exitcode  : -15 (pid: 4471)
+  error_file: <N/A>
+  traceback : Signal 15 (SIGTERM) received by PID 4471
+[2]:
+  time      : 2026-09-22_00:42:38
+  host      : h100-workspace-6-gpu-0-59
+  rank      : 2 (local_rank: 2)
+  exitcode  : -15 (pid: 4472)
+  error_file: <N/A>
+  traceback : Signal 15 (SIGTERM) received by PID 4472
+[3]:
+  time      : 2026-09-22_00:42:38
+  host      : h100-workspace-6-gpu-0-59
+  rank      : 3 (local_rank: 3)
+  exitcode  : -15 (pid: 4473)
+  error_file: <N/A>
+  traceback : Signal 15 (SIGTERM) received by PID 4473
+[4]:
+  time      : 2026-09-22_00:42:38
+  host      : h100-workspace-6-gpu-0-59
+  rank      : 4 (local_rank: 4)
+  exitcode  : -15 (pid: 4474)
+  error_file: <N/A>
+  traceback : Signal 15 (SIGTERM) received by PID 4474
+[5]:
+  time      : 2026-09-22_00:42:38
+  host      : h100-workspace-6-gpu-0-59
+  rank      : 5 (local_rank: 5)
+  exitcode  : -15 (pid: 4475)
+  error_file: <N/A>
+  traceback : Signal 15 (SIGTERM) received by PID 4475
+------------------------------------------------------------
+Root Cause (first observed failure):
+[0]:
+  time      : 2026-09-22_00:42:38
+  host      : h100-workspace-6-gpu-0-59
+  rank      : 0 (local_rank: 0)
+  exitcode  : 1 (pid: 4470)
+  error_file: <N/A>
+  traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
+============================================================

epfl_final/logs/ALIBI_20260921/1p2b_s42_training.log

@@ -0,0 +1,384 @@
+W0922 00:45:12.575000 5167 site-packages/torch/distributed/run.py:851] 
+W0922 00:45:12.575000 5167 site-packages/torch/distributed/run.py:851] *****************************************
+W0922 00:45:12.575000 5167 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. 
+W0922 00:45:12.575000 5167 site-packages/torch/distributed/run.py:851] *****************************************
+New run: Checkpoints will be saved to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/1p2b_alibi_s42/run_20260922_004519
+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/1p2b_alibi_s42/run_20260922_004519",
+  "resume_from_checkpoint": null,
+  "force_new_run": false,
+  "num_epochs": 1,
+  "max_steps": 50462,
+  "per_device_batch_size": 1,
+  "gradient_accumulation_steps": 10,
+  "learning_rate": 0.0005,
+  "weight_decay": 0.01,
+  "warmup_steps": 750,
+  "lr_scheduler": "cosine",
+  "max_grad_norm": 1.0,
+  "save_steps": 1000,
+  "eval_steps": 1000,
+  "eval_batches": 50,
+  "logging_steps": 100,
+  "seed": 42,
+  "verbose": true,
+  "checkpoint_activations": false,
+  "resume_micro_step_override": null,
+  "strategy": "ddp",
+  "mixed_precision": "bf16",
+  "wandb_enabled": true,
+  "wandb_project": "stochastic_sparsity",
+  "wandb_name": "1.2b alibi 8k s42",
+  "wandb_tags": [
+    "qwen3",
+    "1.2b",
+    "8k",
+    "alibi",
+    "dense",
+    "seed42"
+  ],
+  "hidden_size": 1536,
+  "intermediate_size": 6144,
+  "num_attention_heads": 12,
+  "num_key_value_heads": 4,
+  "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": 42,
+  "sparse_mixed_nrb_set": null,
+  "dense_mode": false,
+  "dense_step_intervals": null
+}
+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
+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
+  rope_train_position_scale_set: None
+  attention_type: dense
+
+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                              [1, 8192]                 --                        --                        True
+├─Qwen3Model: 1-1                             --                        [1, 8192, 1536]           --                        True
+│    └─Embedding: 2-1                         [1, 8192]                 [1, 8192, 1536]           233,373,696               True
+│    └─ModuleList: 2-2                        --                        --                        --                        True
+│    │    └─Qwen3DecoderLayer: 3-1            [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-2            [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-3            [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-4            [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-5            [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-6            [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-7            [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-8            [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-9            [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-10           [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-11           [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-12           [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-13           [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-14           [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-15           [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-16           [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-17           [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-18           [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-19           [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-20           [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-21           [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-22           [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-23           [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-24           [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-25           [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-26           [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-27           [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-28           [1, 8192, 1536]           [1, 8192, 1536]           34,606,336                True
+│    └─Qwen3RMSNorm: 2-3                      [1, 8192, 1536]           [1, 8192, 1536]           1,536                     True
+├─Linear: 1-2                                 [1, 8192, 1536]           [1, 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): 1.44
+=================================================================================================================================================
+Input size (MB): 0.07
+Forward/backward pass size (MB): 26218.59
+Params size (MB): 2871.45
+Estimated Total Size (MB): 29090.11
+=================================================================================================================================================
+  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...
+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)
+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):                        6
+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):       504628
+Optimizer steps per epoch:                50462
+Total optimizer steps:                    50462  (capped by max_steps; 1.00 epochs)
+Starting training...
+Switching attention mode for step 1: sparse_stochastic (configured sparse attention)
+Step 100/50462 | Loss: 9.7564 | LR: 6.67e-05 | 2933.7ms/step | 20.5 samples/s
+Step 200/50462 | Loss: 6.6029 | LR: 1.33e-04 | 2002.0ms/step | 30.0 samples/s
+Step 300/50462 | Loss: 5.8119 | LR: 2.00e-04 | 2003.7ms/step | 29.9 samples/s
+Step 400/50462 | Loss: 5.2235 | LR: 2.67e-04 | 2004.4ms/step | 29.9 samples/s
+Step 500/50462 | Loss: 4.8063 | LR: 3.33e-04 | 2005.9ms/step | 29.9 samples/s
+Step 600/50462 | Loss: 4.4823 | LR: 4.00e-04 | 2005.3ms/step | 29.9 samples/s
+Step 700/50462 | Loss: 4.2181 | LR: 4.67e-04 | 2008.4ms/step | 29.9 samples/s
+Step 800/50462 | Loss: 3.9682 | LR: 5.00e-04 | 2005.8ms/step | 29.9 samples/s
+Step 900/50462 | Loss: 3.7799 | LR: 5.00e-04 | 2007.3ms/step | 29.9 samples/s
+Step 1000/50462 | Loss: 3.6488 | LR: 5.00e-04 | 2008.0ms/step | 29.9 samples/s
+Validation at step 1000: Loss = 3.5951, Perplexity = 36.42
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/1p2b_alibi_s42/run_20260922_004519/checkpoint-1000
+Step 1100/50462 | Loss: 3.5806 | LR: 5.00e-04 | 2365.9ms/step | 25.4 samples/s
+Step 1200/50462 | Loss: 3.4931 | LR: 5.00e-04 | 2010.4ms/step | 29.8 samples/s
+Step 1300/50462 | Loss: 3.4498 | LR: 5.00e-04 | 2008.8ms/step | 29.9 samples/s
+Step 1400/50462 | Loss: 3.3829 | LR: 5.00e-04 | 2007.8ms/step | 29.9 samples/s
+Step 1500/50462 | Loss: 3.3443 | LR: 5.00e-04 | 2007.4ms/step | 29.9 samples/s
+Step 1600/50462 | Loss: 3.3211 | LR: 5.00e-04 | 2010.8ms/step | 29.8 samples/s
+Step 1700/50462 | Loss: 3.2849 | LR: 5.00e-04 | 2010.9ms/step | 29.8 samples/s
+Step 1800/50462 | Loss: 3.2642 | LR: 4.99e-04 | 2008.3ms/step | 29.9 samples/s
+Step 1900/50462 | Loss: 3.2232 | LR: 4.99e-04 | 2009.6ms/step | 29.9 samples/s
+Step 2000/50462 | Loss: 3.2203 | LR: 4.99e-04 | 2011.5ms/step | 29.8 samples/s
+Validation at step 2000: Loss = 3.1844, Perplexity = 24.15
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/1p2b_alibi_s42/run_20260922_004519/checkpoint-2000
+Step 2100/50462 | Loss: 3.1919 | LR: 4.99e-04 | 2151.3ms/step | 27.9 samples/s
+Step 2200/50462 | Loss: 3.1773 | LR: 4.99e-04 | 2010.7ms/step | 29.8 samples/s
+Step 2300/50462 | Loss: 3.1579 | LR: 4.99e-04 | 2008.3ms/step | 29.9 samples/s
+Step 2400/50462 | Loss: 3.1279 | LR: 4.99e-04 | 2007.3ms/step | 29.9 samples/s
+Step 2500/50462 | Loss: 3.1264 | LR: 4.98e-04 | 2014.0ms/step | 29.8 samples/s
+Step 2600/50462 | Loss: 3.1103 | LR: 4.98e-04 | 2009.4ms/step | 29.9 samples/s
+Step 2700/50462 | Loss: 3.1026 | LR: 4.98e-04 | 2011.8ms/step | 29.8 samples/s
+Step 2800/50462 | Loss: 3.0749 | LR: 4.98e-04 | 2012.6ms/step | 29.8 samples/s
+Step 2900/50462 | Loss: 3.0640 | LR: 4.98e-04 | 2009.5ms/step | 29.9 samples/s
+Step 3000/50462 | Loss: 3.0504 | LR: 4.97e-04 | 2012.5ms/step | 29.8 samples/s
+Validation at step 3000: Loss = 3.0370, Perplexity = 20.84
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/1p2b_alibi_s42/run_20260922_004519/checkpoint-3000
+Step 3100/50462 | Loss: 3.0520 | LR: 4.97e-04 | 2151.4ms/step | 27.9 samples/s
+Step 3200/50462 | Loss: 3.0462 | LR: 4.97e-04 | 2012.0ms/step | 29.8 samples/s
+Step 3300/50462 | Loss: 3.0178 | LR: 4.97e-04 | 2012.8ms/step | 29.8 samples/s
+Step 3400/50462 | Loss: 3.0189 | LR: 4.97e-04 | 2009.2ms/step | 29.9 samples/s
+Step 3500/50462 | Loss: 3.0149 | LR: 4.96e-04 | 2012.1ms/step | 29.8 samples/s
+Step 3600/50462 | Loss: 3.0019 | LR: 4.96e-04 | 2014.0ms/step | 29.8 samples/s
+Step 3700/50462 | Loss: 2.9922 | LR: 4.96e-04 | 2014.4ms/step | 29.8 samples/s
+Step 3800/50462 | Loss: 2.9908 | LR: 4.95e-04 | 2011.4ms/step | 29.8 samples/s
+Step 3900/50462 | Loss: 2.9756 | LR: 4.95e-04 | 2014.8ms/step | 29.8 samples/s
+Step 4000/50462 | Loss: 2.9827 | LR: 4.95e-04 | 2012.8ms/step | 29.8 samples/s
+Validation at step 4000: Loss = 2.9485, Perplexity = 19.08
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/1p2b_alibi_s42/run_20260922_004519/checkpoint-4000
+Step 4100/50462 | Loss: 2.9640 | LR: 4.94e-04 | 2150.9ms/step | 27.9 samples/s
+Step 4200/50462 | Loss: 2.9522 | LR: 4.94e-04 | 2010.7ms/step | 29.8 samples/s
+Step 4300/50462 | Loss: 2.9387 | LR: 4.94e-04 | 2011.5ms/step | 29.8 samples/s
+Step 4400/50462 | Loss: 2.9428 | LR: 4.93e-04 | 2012.8ms/step | 29.8 samples/s
+Step 4500/50462 | Loss: 2.9382 | LR: 4.93e-04 | 2010.6ms/step | 29.8 samples/s
+Step 4600/50462 | Loss: 2.9436 | LR: 4.93e-04 | 2014.4ms/step | 29.8 samples/s
+Step 4700/50462 | Loss: 2.9131 | LR: 4.92e-04 | 2012.2ms/step | 29.8 samples/s
+Step 4800/50462 | Loss: 2.9221 | LR: 4.92e-04 | 2011.4ms/step | 29.8 samples/s
+Step 4900/50462 | Loss: 2.9140 | LR: 4.91e-04 | 2012.3ms/step | 29.8 samples/s
+Step 5000/50462 | Loss: 2.9217 | LR: 4.91e-04 | 2012.5ms/step | 29.8 samples/s
+Validation at step 5000: Loss = 2.8901, Perplexity = 18.00
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/1p2b_alibi_s42/run_20260922_004519/checkpoint-5000
+Step 5100/50462 | Loss: 2.9076 | LR: 4.91e-04 | 2156.6ms/step | 27.8 samples/s
+Step 5200/50462 | Loss: 2.8801 | LR: 4.90e-04 | 2012.9ms/step | 29.8 samples/s
+Step 5300/50462 | Loss: 2.9009 | LR: 4.90e-04 | 2012.3ms/step | 29.8 samples/s
+Step 5400/50462 | Loss: 2.9030 | LR: 4.89e-04 | 2013.2ms/step | 29.8 samples/s
+Step 5500/50462 | Loss: 2.8833 | LR: 4.89e-04 | 2014.0ms/step | 29.8 samples/s
+Step 5600/50462 | Loss: 2.8814 | LR: 4.88e-04 | 2015.0ms/step | 29.8 samples/s
+Step 5700/50462 | Loss: 2.8608 | LR: 4.88e-04 | 2014.2ms/step | 29.8 samples/s
+Step 5800/50462 | Loss: 2.8817 | LR: 4.87e-04 | 2013.1ms/step | 29.8 samples/s
+Step 5900/50462 | Loss: 2.8580 | LR: 4.87e-04 | 2016.1ms/step | 29.8 samples/s
+Step 6000/50462 | Loss: 2.8499 | LR: 4.86e-04 | 2012.9ms/step | 29.8 samples/s
+Validation at step 6000: Loss = 2.8466, Perplexity = 17.23
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/1p2b_alibi_s42/run_20260922_004519/checkpoint-6000
+Step 6100/50462 | Loss: 2.8600 | LR: 4.86e-04 | 2152.0ms/step | 27.9 samples/s
+Step 6200/50462 | Loss: 2.8410 | LR: 4.85e-04 | 2015.2ms/step | 29.8 samples/s
+Step 6300/50462 | Loss: 2.8497 | LR: 4.85e-04 | 2015.4ms/step | 29.8 samples/s
+Step 6400/50462 | Loss: 2.8409 | LR: 4.84e-04 | 2014.0ms/step | 29.8 samples/s
+Step 6500/50462 | Loss: 2.8485 | LR: 4.84e-04 | 2014.3ms/step | 29.8 samples/s
+Step 6600/50462 | Loss: 2.8411 | LR: 4.83e-04 | 2015.5ms/step | 29.8 samples/s
+Step 6700/50462 | Loss: 2.8474 | LR: 4.83e-04 | 2014.6ms/step | 29.8 samples/s
+Step 6800/50462 | Loss: 2.8414 | LR: 4.82e-04 | 2015.2ms/step | 29.8 samples/s
+Step 6900/50462 | Loss: 2.8228 | LR: 4.81e-04 | 2014.4ms/step | 29.8 samples/s
+Step 7000/50462 | Loss: 2.8189 | LR: 4.81e-04 | 2013.6ms/step | 29.8 samples/s
+Validation at step 7000: Loss = 2.8150, Perplexity = 16.69
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/1p2b_alibi_s42/run_20260922_004519/checkpoint-7000
+Step 7100/50462 | Loss: 2.8240 | LR: 4.80e-04 | 2153.6ms/step | 27.9 samples/s
+Step 7200/50462 | Loss: 2.8270 | LR: 4.80e-04 | 2012.6ms/step | 29.8 samples/s
+Step 7300/50462 | Loss: 2.8123 | LR: 4.79e-04 | 2014.2ms/step | 29.8 samples/s
+Step 7400/50462 | Loss: 2.8236 | LR: 4.78e-04 | 2011.4ms/step | 29.8 samples/s
+Step 7500/50462 | Loss: 2.8136 | LR: 4.78e-04 | 2013.6ms/step | 29.8 samples/s
+Step 7600/50462 | Loss: 2.8113 | LR: 4.77e-04 | 2013.4ms/step | 29.8 samples/s
+Step 7700/50462 | Loss: 2.8088 | LR: 4.76e-04 | 2015.6ms/step | 29.8 samples/s
+Step 7800/50462 | Loss: 2.8166 | LR: 4.76e-04 | 2014.6ms/step | 29.8 samples/s
+Step 7900/50462 | Loss: 2.8127 | LR: 4.75e-04 | 2015.6ms/step | 29.8 samples/s
+Step 8000/50462 | Loss: 2.7885 | LR: 4.74e-04 | 2016.4ms/step | 29.8 samples/s
+Validation at step 8000: Loss = 2.7844, Perplexity = 16.19
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/1p2b_alibi_s42/run_20260922_004519/checkpoint-8000
+Step 8100/50462 | Loss: 2.7861 | LR: 4.74e-04 | 2155.0ms/step | 27.8 samples/s
+Step 8200/50462 | Loss: 2.7968 | LR: 4.73e-04 | 2015.1ms/step | 29.8 samples/s
+Step 8300/50462 | Loss: 2.7881 | LR: 4.72e-04 | 2017.7ms/step | 29.7 samples/s
+Step 8400/50462 | Loss: 2.7924 | LR: 4.71e-04 | 2016.6ms/step | 29.8 samples/s
+Step 8500/50462 | Loss: 2.7857 | LR: 4.71e-04 | 2015.3ms/step | 29.8 samples/s
+Step 8600/50462 | Loss: 2.7857 | LR: 4.70e-04 | 2015.8ms/step | 29.8 samples/s
+Step 8700/50462 | Loss: 2.7900 | LR: 4.69e-04 | 2014.5ms/step | 29.8 samples/s
+Step 8800/50462 | Loss: 2.7865 | LR: 4.68e-04 | 2013.8ms/step | 29.8 samples/s
+Step 8900/50462 | Loss: 2.7675 | LR: 4.68e-04 | 2014.1ms/step | 29.8 samples/s
+Step 9000/50462 | Loss: 2.7817 | LR: 4.67e-04 | 2013.4ms/step | 29.8 samples/s
+Validation at step 9000: Loss = 2.7560, Perplexity = 15.74
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/1p2b_alibi_s42/run_20260922_004519/checkpoint-9000
+Step 9100/50462 | Loss: 2.7640 | LR: 4.66e-04 | 2147.0ms/step | 27.9 samples/s
+Step 9200/50462 | Loss: 2.7598 | LR: 4.65e-04 | 2012.3ms/step | 29.8 samples/s
+Step 9300/50462 | Loss: 2.7720 | LR: 4.64e-04 | 2012.9ms/step | 29.8 samples/s
+Step 9400/50462 | Loss: 2.7586 | LR: 4.64e-04 | 2012.4ms/step | 29.8 samples/s
+Step 9500/50462 | Loss: 2.7563 | LR: 4.63e-04 | 2011.9ms/step | 29.8 samples/s
+Step 9600/50462 | Loss: 2.7666 | LR: 4.62e-04 | 2012.8ms/step | 29.8 samples/s
+Step 9700/50462 | Loss: 2.7536 | LR: 4.61e-04 | 2012.5ms/step | 29.8 samples/s
+Step 9800/50462 | Loss: 2.7472 | LR: 4.60e-04 | 2012.5ms/step | 29.8 samples/s
+Step 9900/50462 | Loss: 2.7469 | LR: 4.59e-04 | 2012.2ms/step | 29.8 samples/s
+Step 10000/50462 | Loss: 2.7431 | LR: 4.58e-04 | 2013.0ms/step | 29.8 samples/s
+Validation at step 10000: Loss = 2.7340, Perplexity = 15.40
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/1p2b_alibi_s42/run_20260922_004519/checkpoint-10000
+Step 10100/50462 | Loss: 2.7630 | LR: 4.58e-04 | 2152.0ms/step | 27.9 samples/s
+Step 10200/50462 | Loss: 2.7424 | LR: 4.57e-04 | 2013.2ms/step | 29.8 samples/s
+Step 10300/50462 | Loss: 2.7425 | LR: 4.56e-04 | 2012.8ms/step | 29.8 samples/s
+Step 10400/50462 | Loss: 2.7413 | LR: 4.55e-04 | 2011.6ms/step | 29.8 samples/s
+Step 10500/50462 | Loss: 2.7390 | LR: 4.54e-04 | 2011.8ms/step | 29.8 samples/s
+Step 10600/50462 | Loss: 2.7433 | LR: 4.53e-04 | 2012.6ms/step | 29.8 samples/s
+Step 10700/50462 | Loss: 2.7340 | LR: 4.52e-04 | 2011.0ms/step | 29.8 samples/s
+Step 10800/50462 | Loss: 2.7416 | LR: 4.51e-04 | 2015.2ms/step | 29.8 samples/s
+Step 10900/50462 | Loss: 2.7322 | LR: 4.50e-04 | 2011.9ms/step | 29.8 samples/s
+Step 11000/50462 | Loss: 2.7405 | LR: 4.49e-04 | 2010.7ms/step | 29.8 samples/s
+Validation at step 11000: Loss = 2.7164, Perplexity = 15.13
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/1p2b_alibi_s42/run_20260922_004519/checkpoint-11000
+Step 11100/50462 | Loss: 2.7182 | LR: 4.48e-04 | 2148.6ms/step | 27.9 samples/s
+Step 11200/50462 | Loss: 2.7324 | LR: 4.47e-04 | 2010.7ms/step | 29.8 samples/s
+Step 11300/50462 | Loss: 2.7300 | LR: 4.46e-04 | 2009.9ms/step | 29.9 samples/s
+Step 11400/50462 | Loss: 2.7300 | LR: 4.45e-04 | 2010.2ms/step | 29.8 samples/s
+Step 11500/50462 | Loss: 2.7269 | LR: 4.44e-04 | 2010.6ms/step | 29.8 samples/s
+Step 11600/50462 | Loss: 2.7091 | LR: 4.43e-04 | 2008.2ms/step | 29.9 samples/s
+Step 11700/50462 | Loss: 2.7219 | LR: 4.42e-04 | 2012.0ms/step | 29.8 samples/s
+Step 11800/50462 | Loss: 2.7147 | LR: 4.41e-04 | 2010.0ms/step | 29.9 samples/s
+Step 11900/50462 | Loss: 2.7198 | LR: 4.40e-04 | 2008.2ms/step | 29.9 samples/s
+Step 12000/50462 | Loss: 2.7209 | LR: 4.39e-04 | 2010.6ms/step | 29.8 samples/s
+Validation at step 12000: Loss = 2.6993, Perplexity = 14.87
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/1p2b_alibi_s42/run_20260922_004519/checkpoint-12000
+Step 12100/50462 | Loss: 2.7052 | LR: 4.38e-04 | 2145.5ms/step | 28.0 samples/s
+Step 12200/50462 | Loss: 2.7075 | LR: 4.37e-04 | 2012.9ms/step | 29.8 samples/s
+Step 12300/50462 | Loss: 2.7010 | LR: 4.36e-04 | 2011.8ms/step | 29.8 samples/s
+Step 12400/50462 | Loss: 2.7102 | LR: 4.35e-04 | 2011.6ms/step | 29.8 samples/s
+Step 12500/50462 | Loss: 2.7180 | LR: 4.34e-04 | 2011.2ms/step | 29.8 samples/s
+Step 12600/50462 | Loss: 2.7074 | LR: 4.33e-04 | 2012.7ms/step | 29.8 samples/s
+Step 12700/50462 | Loss: 2.7135 | LR: 4.32e-04 | 2012.6ms/step | 29.8 samples/s
+Step 12800/50462 | Loss: 2.6982 | LR: 4.31e-04 | 2014.4ms/step | 29.8 samples/s
+Step 12900/50462 | Loss: 2.7084 | LR: 4.30e-04 | 2012.4ms/step | 29.8 samples/s
+Step 13000/50462 | Loss: 2.6955 | LR: 4.29e-04 | 2012.0ms/step | 29.8 samples/s
+Validation at step 13000: Loss = 2.6845, Perplexity = 14.65
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/1p2b_alibi_s42/run_20260922_004519/checkpoint-13000
+Step 13100/50462 | Loss: 2.7026 | LR: 4.28e-04 | 2152.9ms/step | 27.9 samples/s
+Step 13200/50462 | Loss: 2.7005 | LR: 4.27e-04 | 2013.6ms/step | 29.8 samples/s
+Step 13300/50462 | Loss: 2.7067 | LR: 4.25e-04 | 2014.4ms/step | 29.8 samples/s
+Step 13400/50462 | Loss: 2.6800 | LR: 4.24e-04 | 2012.3ms/step | 29.8 samples/s
+Step 13500/50462 | Loss: 2.6955 | LR: 4.23e-04 | 2010.7ms/step | 29.8 samples/s
+Step 13600/50462 | Loss: 2.6819 | LR: 4.22e-04 | 2013.4ms/step | 29.8 samples/s
+Step 13700/50462 | Loss: 2.6938 | LR: 4.21e-04 | 2013.7ms/step | 29.8 samples/s
+Step 13800/50462 | Loss: 2.6936 | LR: 4.20e-04 | 2015.2ms/step | 29.8 samples/s
+Step 13900/50462 | Loss: 2.6767 | LR: 4.19e-04 | 2013.5ms/step | 29.8 samples/s
+Step 14000/50462 | Loss: 2.6829 | LR: 4.17e-04 | 2013.9ms/step | 29.8 samples/s
+Validation at step 14000: Loss = 2.6721, Perplexity = 14.47
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/1p2b_alibi_s42/run_20260922_004519/checkpoint-14000
+Step 14100/50462 | Loss: 2.6794 | LR: 4.16e-04 | 2150.6ms/step | 27.9 samples/s
+Step 14200/50462 | Loss: 2.6904 | LR: 4.15e-04 | 2011.7ms/step | 29.8 samples/s
+Step 14300/50462 | Loss: 2.6939 | LR: 4.14e-04 | 2011.4ms/step | 29.8 samples/s
+Step 14400/50462 | Loss: 2.6725 | LR: 4.13e-04 | 2013.7ms/step | 29.8 samples/s
+Step 14500/50462 | Loss: 2.6714 | LR: 4.11e-04 | 2012.1ms/step | 29.8 samples/s
+Step 14600/50462 | Loss: 2.6864 | LR: 4.10e-04 | 2012.9ms/step | 29.8 samples/s
+Step 14700/50462 | Loss: 2.6831 | LR: 4.09e-04 | 2011.2ms/step | 29.8 samples/s
+Step 14800/50462 | Loss: 2.6668 | LR: 4.08e-04 | 2012.0ms/step | 29.8 samples/s
+Step 14900/50462 | Loss: 2.6653 | LR: 4.07e-04 | 2014.2ms/step | 29.8 samples/s
+Step 15000/50462 | Loss: 2.6649 | LR: 4.05e-04 | 2013.8ms/step | 29.8 samples/s
+Validation at step 15000: Loss = 2.6590, Perplexity = 14.28
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/1p2b_alibi_s42/run_20260922_004519/checkpoint-15000
+Step 15100/50462 | Loss: 2.6754 | LR: 4.04e-04 | 2154.4ms/step | 27.9 samples/s
+Step 15200/50462 | Loss: 2.6706 | LR: 4.03e-04 | 2012.8ms/step | 29.8 samples/s
+Step 15300/50462 | Loss: 2.6642 | LR: 4.02e-04 | 2014.3ms/step | 29.8 samples/s
+Step 15400/50462 | Loss: 2.6754 | LR: 4.00e-04 | 2015.3ms/step | 29.8 samples/s
+Step 15500/50462 | Loss: 2.6529 | LR: 3.99e-04 | 2010.9ms/step | 29.8 samples/s
+Step 15600/50462 | Loss: 2.6503 | LR: 3.98e-04 | 2013.6ms/step | 29.8 samples/s
+Step 15700/50462 | Loss: 2.6575 | LR: 3.96e-04 | 2011.7ms/step | 29.8 samples/s
+Step 15800/50462 | Loss: 2.6482 | LR: 3.95e-04 | 2010.7ms/step | 29.8 samples/s
+Step 15900/50462 | Loss: 2.6506 | LR: 3.94e-04 | 2011.5ms/step | 29.8 samples/s
+Step 16000/50462 | Loss: 2.6573 | LR: 3.93e-04 | 2012.5ms/step | 29.8 samples/s
+Validation at step 16000: Loss = 2.6454, Perplexity = 14.09
+Saved checkpoint to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/1p2b_alibi_s42/run_20260922_004519/checkpoint-16000
+Step 16100/50462 | Loss: 2.6677 | LR: 3.91e-04 | 2148.3ms/step | 27.9 samples/s

epfl_final/logs/ALIBI_20260921/launcher.log

@@ -0,0 +1,7 @@
+launch seed=42 gpu=0 config=/mloscratch/homes/navasard/sparse/SparseAtt/configs/alibi_20260921/365m_alibi_s42.json log=/mloscratch/homes/navasard/sparse/logs/ALIBI_20260921/seed42.log
+launch seed=43 gpu=1 config=/mloscratch/homes/navasard/sparse/SparseAtt/configs/alibi_20260921/365m_alibi_s43.json log=/mloscratch/homes/navasard/sparse/logs/ALIBI_20260921/seed43.log
+launch seed=44 gpu=2 config=/mloscratch/homes/navasard/sparse/SparseAtt/configs/alibi_20260921/365m_alibi_s44.json log=/mloscratch/homes/navasard/sparse/logs/ALIBI_20260921/seed44.log
+launch seed=45 gpu=3 config=/mloscratch/homes/navasard/sparse/SparseAtt/configs/alibi_20260921/365m_alibi_s45.json log=/mloscratch/homes/navasard/sparse/logs/ALIBI_20260921/seed45.log
+launch seed=46 gpu=4 config=/mloscratch/homes/navasard/sparse/SparseAtt/configs/alibi_20260921/365m_alibi_s46.json log=/mloscratch/homes/navasard/sparse/logs/ALIBI_20260921/seed46.log
+launch seed=47 gpu=5 config=/mloscratch/homes/navasard/sparse/SparseAtt/configs/alibi_20260921/365m_alibi_s47.json log=/mloscratch/homes/navasard/sparse/logs/ALIBI_20260921/seed47.log
+launched 2026-09-21T03:04:45+00:00

epfl_final/logs/ALIBI_20260921/launcher_restart_20260921_131155.log

@@ -0,0 +1,7 @@
+launch seed=42 gpu=0 config=/mloscratch/homes/navasard/sparse/SparseAtt/configs/alibi_20260921/365m_alibi_s42.json log=/mloscratch/homes/navasard/sparse/logs/ALIBI_20260921/restart_20260921_131155/seed42.log
+launch seed=43 gpu=1 config=/mloscratch/homes/navasard/sparse/SparseAtt/configs/alibi_20260921/365m_alibi_s43.json log=/mloscratch/homes/navasard/sparse/logs/ALIBI_20260921/restart_20260921_131155/seed43.log
+launch seed=44 gpu=2 config=/mloscratch/homes/navasard/sparse/SparseAtt/configs/alibi_20260921/365m_alibi_s44.json log=/mloscratch/homes/navasard/sparse/logs/ALIBI_20260921/restart_20260921_131155/seed44.log
+launch seed=45 gpu=3 config=/mloscratch/homes/navasard/sparse/SparseAtt/configs/alibi_20260921/365m_alibi_s45.json log=/mloscratch/homes/navasard/sparse/logs/ALIBI_20260921/restart_20260921_131155/seed45.log
+launch seed=46 gpu=4 config=/mloscratch/homes/navasard/sparse/SparseAtt/configs/alibi_20260921/365m_alibi_s46.json log=/mloscratch/homes/navasard/sparse/logs/ALIBI_20260921/restart_20260921_131155/seed46.log
+launch seed=47 gpu=5 config=/mloscratch/homes/navasard/sparse/SparseAtt/configs/alibi_20260921/365m_alibi_s47.json log=/mloscratch/homes/navasard/sparse/logs/ALIBI_20260921/restart_20260921_131155/seed47.log
+launched 2026-09-21T13:11:55+00:00

epfl_final/logs/ALIBI_20260921/launcher_resume_20260921_142934.log

@@ -0,0 +1,7 @@
+launch seed=42 gpu=0 config=/mloscratch/homes/navasard/sparse/SparseAtt/configs/alibi_20260921/365m_alibi_s42.json log=/mloscratch/homes/navasard/sparse/logs/ALIBI_20260921/resume_20260921_142934/seed42.log
+launch seed=43 gpu=1 config=/mloscratch/homes/navasard/sparse/SparseAtt/configs/alibi_20260921/365m_alibi_s43.json log=/mloscratch/homes/navasard/sparse/logs/ALIBI_20260921/resume_20260921_142934/seed43.log
+launch seed=44 gpu=2 config=/mloscratch/homes/navasard/sparse/SparseAtt/configs/alibi_20260921/365m_alibi_s44.json log=/mloscratch/homes/navasard/sparse/logs/ALIBI_20260921/resume_20260921_142934/seed44.log
+launch seed=45 gpu=3 config=/mloscratch/homes/navasard/sparse/SparseAtt/configs/alibi_20260921/365m_alibi_s45.json log=/mloscratch/homes/navasard/sparse/logs/ALIBI_20260921/resume_20260921_142934/seed45.log
+launch seed=46 gpu=4 config=/mloscratch/homes/navasard/sparse/SparseAtt/configs/alibi_20260921/365m_alibi_s46.json log=/mloscratch/homes/navasard/sparse/logs/ALIBI_20260921/resume_20260921_142934/seed46.log
+launch seed=47 gpu=5 config=/mloscratch/homes/navasard/sparse/SparseAtt/configs/alibi_20260921/365m_alibi_s47.json log=/mloscratch/homes/navasard/sparse/logs/ALIBI_20260921/resume_20260921_142934/seed47.log
+launched 2026-09-21T14:29:34+00:00

epfl_final/logs/ALIBI_20260921/restart_20260921_131155/seed42.log

@@ -0,0 +1,187 @@
+New run: Checkpoints will be saved to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/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_s42/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": 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_131249-3uan03m3
+Starting training...
+Switching attention mode for step 1: sparse_stochastic (configured sparse attention)
+Step 100/44000 | Loss: 10.5675 | LR: 6.67e-05 | 1620.5ms/step | 12.3 samples/s
+Step 200/44000 | Loss: 7.3103 | LR: 1.33e-04 | 1478.1ms/step | 13.5 samples/s
+Validation at step 250: Loss = 6.3045, Perplexity = 547.02
+Step 300/44000 | Loss: 6.3237 | LR: 2.00e-04 | 1774.0ms/step | 11.3 samples/s
+Step 400/44000 | Loss: 5.8414 | LR: 2.67e-04 | 1493.3ms/step | 13.4 samples/s
+Step 500/44000 | Loss: 5.4582 | LR: 3.33e-04 | 1494.7ms/step | 13.4 samples/s
+Validation at step 500: Loss = 5.2807, Perplexity = 196.50
+Step 600/44000 | Loss: 5.1359 | LR: 4.00e-04 | 1731.4ms/step | 11.6 samples/s
+Step 700/44000 | Loss: 4.8503 | LR: 4.67e-04 | 1490.9ms/step | 13.4 samples/s
+Validation at step 750: Loss = 4.5971, Perplexity = 99.20
+Step 800/44000 | Loss: 4.5806 | LR: 5.00e-04 | 1731.5ms/step | 11.6 samples/s
+Step 900/44000 | Loss: 4.3781 | LR: 5.00e-04 | 1491.6ms/step | 13.4 samples/s

epfl_final/logs/ALIBI_20260921/restart_20260921_131155/seed42.pid

@@ -0,0 +1 @@
+353

epfl_final/logs/ALIBI_20260921/restart_20260921_131155/seed43.log

@@ -0,0 +1,187 @@
+New run: Checkpoints will be saved to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/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_s43/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": 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_131249-n6s56yfr
+Starting training...
+Switching attention mode for step 1: sparse_stochastic (configured sparse attention)
+Step 100/44000 | Loss: 10.6180 | LR: 6.67e-05 | 1630.6ms/step | 12.3 samples/s
+Step 200/44000 | Loss: 7.3259 | LR: 1.33e-04 | 1491.2ms/step | 13.4 samples/s
+Validation at step 250: Loss = 6.3016, Perplexity = 545.45
+Step 300/44000 | Loss: 6.3205 | LR: 2.00e-04 | 1788.2ms/step | 11.2 samples/s
+Step 400/44000 | Loss: 5.8558 | LR: 2.67e-04 | 1492.4ms/step | 13.4 samples/s
+Step 500/44000 | Loss: 5.4556 | LR: 3.33e-04 | 1491.8ms/step | 13.4 samples/s
+Validation at step 500: Loss = 5.2855, Perplexity = 197.45
+Step 600/44000 | Loss: 5.1474 | LR: 4.00e-04 | 1732.4ms/step | 11.5 samples/s
+Step 700/44000 | Loss: 4.8551 | LR: 4.67e-04 | 1492.0ms/step | 13.4 samples/s
+Validation at step 750: Loss = 4.5870, Perplexity = 98.20
+Step 800/44000 | Loss: 4.5995 | LR: 5.00e-04 | 1731.6ms/step | 11.5 samples/s
+Step 900/44000 | Loss: 4.3816 | LR: 5.00e-04 | 1492.2ms/step | 13.4 samples/s

epfl_final/logs/ALIBI_20260921/restart_20260921_131155/seed43.pid

@@ -0,0 +1 @@
+354

epfl_final/logs/ALIBI_20260921/restart_20260921_131155/seed44.log

@@ -0,0 +1,187 @@
+New run: Checkpoints will be saved to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/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_s44/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": 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_131249-ifgmkpzt
+Starting training...
+Switching attention mode for step 1: sparse_stochastic (configured sparse attention)
+Step 100/44000 | Loss: 10.6002 | LR: 6.67e-05 | 1625.9ms/step | 12.3 samples/s
+Step 200/44000 | Loss: 7.3141 | LR: 1.33e-04 | 1494.6ms/step | 13.4 samples/s
+Validation at step 250: Loss = 6.3097, Perplexity = 549.87
+Step 300/44000 | Loss: 6.3228 | LR: 2.00e-04 | 1786.9ms/step | 11.2 samples/s
+Step 400/44000 | Loss: 5.8534 | LR: 2.67e-04 | 1504.9ms/step | 13.3 samples/s
+Step 500/44000 | Loss: 5.4672 | LR: 3.33e-04 | 1505.7ms/step | 13.3 samples/s
+Validation at step 500: Loss = 5.3115, Perplexity = 202.64
+Step 600/44000 | Loss: 5.1448 | LR: 4.00e-04 | 1743.6ms/step | 11.5 samples/s
+Step 700/44000 | Loss: 4.8769 | LR: 4.67e-04 | 1502.6ms/step | 13.3 samples/s
+Validation at step 750: Loss = 4.6322, Perplexity = 102.74
+Step 800/44000 | Loss: 4.6297 | LR: 5.00e-04 | 1745.7ms/step | 11.5 samples/s
+Step 900/44000 | Loss: 4.3989 | LR: 5.00e-04 | 1504.3ms/step | 13.3 samples/s

epfl_final/logs/ALIBI_20260921/restart_20260921_131155/seed44.pid

@@ -0,0 +1 @@
+355

epfl_final/logs/ALIBI_20260921/restart_20260921_131155/seed45.log

@@ -0,0 +1,187 @@
+New run: Checkpoints will be saved to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s45/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_s45/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": 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_131249-3la36nzy
+Starting training...
+Switching attention mode for step 1: sparse_stochastic (configured sparse attention)
+Step 100/44000 | Loss: 10.6683 | LR: 6.67e-05 | 1620.6ms/step | 12.3 samples/s
+Step 200/44000 | Loss: 7.3445 | LR: 1.33e-04 | 1481.8ms/step | 13.5 samples/s
+Validation at step 250: Loss = 6.3185, Perplexity = 554.73
+Step 300/44000 | Loss: 6.3372 | LR: 2.00e-04 | 1773.1ms/step | 11.3 samples/s
+Step 400/44000 | Loss: 5.8480 | LR: 2.67e-04 | 1489.5ms/step | 13.4 samples/s
+Step 500/44000 | Loss: 5.4653 | LR: 3.33e-04 | 1489.4ms/step | 13.4 samples/s
+Validation at step 500: Loss = 5.2802, Perplexity = 196.41
+Step 600/44000 | Loss: 5.1333 | LR: 4.00e-04 | 1728.2ms/step | 11.6 samples/s
+Step 700/44000 | Loss: 4.8685 | LR: 4.67e-04 | 1489.3ms/step | 13.4 samples/s
+Validation at step 750: Loss = 4.6110, Perplexity = 100.58
+Step 800/44000 | Loss: 4.6093 | LR: 5.00e-04 | 1728.6ms/step | 11.6 samples/s
+Step 900/44000 | Loss: 4.3876 | LR: 5.00e-04 | 1489.7ms/step | 13.4 samples/s

epfl_final/logs/ALIBI_20260921/restart_20260921_131155/seed45.pid

@@ -0,0 +1 @@
+356

epfl_final/logs/ALIBI_20260921/restart_20260921_131155/seed46.log

@@ -0,0 +1,187 @@
+New run: Checkpoints will be saved to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s46/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_s46/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": 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_131249-meicy3md
+Starting training...
+Switching attention mode for step 1: sparse_stochastic (configured sparse attention)
+Step 100/44000 | Loss: 10.6054 | LR: 6.67e-05 | 1648.3ms/step | 12.1 samples/s
+Step 200/44000 | Loss: 7.3376 | LR: 1.33e-04 | 1517.3ms/step | 13.2 samples/s
+Validation at step 250: Loss = 6.3202, Perplexity = 555.69
+Step 300/44000 | Loss: 6.3291 | LR: 2.00e-04 | 1814.0ms/step | 11.0 samples/s
+Step 400/44000 | Loss: 5.8539 | LR: 2.67e-04 | 1528.0ms/step | 13.1 samples/s
+Step 500/44000 | Loss: 5.4670 | LR: 3.33e-04 | 1528.6ms/step | 13.1 samples/s
+Validation at step 500: Loss = 5.2901, Perplexity = 198.35
+Step 600/44000 | Loss: 5.1335 | LR: 4.00e-04 | 1770.0ms/step | 11.3 samples/s
+Step 700/44000 | Loss: 4.8768 | LR: 4.67e-04 | 1526.3ms/step | 13.1 samples/s
+Validation at step 750: Loss = 4.6087, Perplexity = 100.36
+Step 800/44000 | Loss: 4.6136 | LR: 5.00e-04 | 1774.0ms/step | 11.3 samples/s
+Step 900/44000 | Loss: 4.3924 | LR: 5.00e-04 | 1528.4ms/step | 13.1 samples/s

epfl_final/logs/ALIBI_20260921/restart_20260921_131155/seed46.pid

@@ -0,0 +1 @@
+357

epfl_final/logs/ALIBI_20260921/seed42.log

@@ -0,0 +1,373 @@
+New run: Checkpoints will be saved to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s42/run_20260921_030452
+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_030452",
+  "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": 22000,
+  "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_030515-ao5edbbv
+Starting training...
+Switching attention mode for step 1: sparse_stochastic (configured sparse attention)
+Step 100/44000 | Loss: 10.5675 | LR: 6.67e-05 | 1972.1ms/step | 10.1 samples/s
+Step 200/44000 | Loss: 7.3113 | LR: 1.33e-04 | 1503.8ms/step | 13.3 samples/s
+Validation at step 250: Loss = 6.3071, Perplexity = 548.47
+Step 300/44000 | Loss: 6.3249 | LR: 2.00e-04 | 1897.6ms/step | 10.5 samples/s
+Step 400/44000 | Loss: 5.8445 | LR: 2.67e-04 | 1512.6ms/step | 13.2 samples/s
+Step 500/44000 | Loss: 5.4630 | LR: 3.33e-04 | 1513.7ms/step | 13.2 samples/s
+Validation at step 500: Loss = 5.2931, Perplexity = 198.95
+Step 600/44000 | Loss: 5.1375 | LR: 4.00e-04 | 1756.8ms/step | 11.4 samples/s
+Step 700/44000 | Loss: 4.8494 | LR: 4.67e-04 | 1514.6ms/step | 13.2 samples/s
+Validation at step 750: Loss = 4.5929, Perplexity = 98.78
+Step 800/44000 | Loss: 4.5821 | LR: 5.00e-04 | 1756.6ms/step | 11.4 samples/s
+Step 900/44000 | Loss: 4.3793 | LR: 5.00e-04 | 1513.0ms/step | 13.2 samples/s
+Step 1000/44000 | Loss: 4.2242 | LR: 5.00e-04 | 1510.9ms/step | 13.2 samples/s
+Validation at step 1000: Loss = 4.1745, Perplexity = 65.00
+Step 1100/44000 | Loss: 4.1056 | LR: 5.00e-04 | 1748.4ms/step | 11.4 samples/s
+Step 1200/44000 | Loss: 4.0142 | LR: 5.00e-04 | 1500.3ms/step | 13.3 samples/s
+Validation at step 1250: Loss = 3.9599, Perplexity = 52.45
+Step 1300/44000 | Loss: 3.9476 | LR: 5.00e-04 | 1745.2ms/step | 11.5 samples/s
+Step 1400/44000 | Loss: 3.8822 | LR: 5.00e-04 | 1504.6ms/step | 13.3 samples/s
+Step 1500/44000 | Loss: 3.8561 | LR: 5.00e-04 | 1500.8ms/step | 13.3 samples/s
+Validation at step 1500: Loss = 3.8278, Perplexity = 45.96
+Step 1600/44000 | Loss: 3.8000 | LR: 5.00e-04 | 1741.5ms/step | 11.5 samples/s
+Step 1700/44000 | Loss: 3.7480 | LR: 4.99e-04 | 1502.5ms/step | 13.3 samples/s
+Validation at step 1750: Loss = 3.7341, Perplexity = 41.85
+Step 1800/44000 | Loss: 3.7115 | LR: 4.99e-04 | 1746.7ms/step | 11.5 samples/s
+Step 1900/44000 | Loss: 3.6990 | LR: 4.99e-04 | 1501.7ms/step | 13.3 samples/s
+Step 2000/44000 | Loss: 3.6665 | LR: 4.99e-04 | 1503.2ms/step | 13.3 samples/s
+Validation at step 2000: Loss = 3.6616, Perplexity = 38.92
+Step 2100/44000 | Loss: 3.6322 | LR: 4.99e-04 | 1743.2ms/step | 11.5 samples/s
+Step 2200/44000 | Loss: 3.6143 | LR: 4.99e-04 | 1500.8ms/step | 13.3 samples/s
+Validation at step 2250: Loss = 3.6065, Perplexity = 36.84
+Step 2300/44000 | Loss: 3.5900 | LR: 4.98e-04 | 1746.1ms/step | 11.5 samples/s
+Step 2400/44000 | Loss: 3.5640 | LR: 4.98e-04 | 1502.7ms/step | 13.3 samples/s
+Step 2500/44000 | Loss: 3.5592 | LR: 4.98e-04 | 1500.4ms/step | 13.3 samples/s
+Validation at step 2500: Loss = 3.5598, Perplexity = 35.15
+Step 2600/44000 | Loss: 3.5447 | LR: 4.98e-04 | 1743.3ms/step | 11.5 samples/s
+Step 2700/44000 | Loss: 3.5305 | LR: 4.97e-04 | 1500.5ms/step | 13.3 samples/s
+Validation at step 2750: Loss = 3.5247, Perplexity = 33.94
+Step 2800/44000 | Loss: 3.5067 | LR: 4.97e-04 | 1744.8ms/step | 11.5 samples/s
+Step 2900/44000 | Loss: 3.4870 | LR: 4.97e-04 | 1501.4ms/step | 13.3 samples/s
+Step 3000/44000 | Loss: 3.4748 | LR: 4.97e-04 | 1500.6ms/step | 13.3 samples/s
+Validation at step 3000: Loss = 3.4869, Perplexity = 32.68
+Step 3100/44000 | Loss: 3.4692 | LR: 4.96e-04 | 1745.2ms/step | 11.5 samples/s
+Step 3200/44000 | Loss: 3.4570 | LR: 4.96e-04 | 1501.9ms/step | 13.3 samples/s
+Validation at step 3250: Loss = 3.4576, Perplexity = 31.74
+Step 3300/44000 | Loss: 3.4487 | LR: 4.96e-04 | 1744.1ms/step | 11.5 samples/s
+Step 3400/44000 | Loss: 3.4405 | LR: 4.95e-04 | 1502.2ms/step | 13.3 samples/s
+Step 3500/44000 | Loss: 3.4268 | LR: 4.95e-04 | 1502.0ms/step | 13.3 samples/s
+Validation at step 3500: Loss = 3.4310, Perplexity = 30.91
+Step 3600/44000 | Loss: 3.4135 | LR: 4.95e-04 | 1744.2ms/step | 11.5 samples/s
+Step 3700/44000 | Loss: 3.4041 | LR: 4.94e-04 | 1503.6ms/step | 13.3 samples/s
+Validation at step 3750: Loss = 3.4070, Perplexity = 30.18
+Step 3800/44000 | Loss: 3.3954 | LR: 4.94e-04 | 1743.4ms/step | 11.5 samples/s
+Step 3900/44000 | Loss: 3.3866 | LR: 4.93e-04 | 1502.8ms/step | 13.3 samples/s
+Step 4000/44000 | Loss: 3.3775 | LR: 4.93e-04 | 1501.4ms/step | 13.3 samples/s
+Validation at step 4000: Loss = 3.3862, Perplexity = 29.55
+Step 4100/44000 | Loss: 3.3811 | LR: 4.93e-04 | 1744.1ms/step | 11.5 samples/s
+Step 4200/44000 | Loss: 3.3671 | LR: 4.92e-04 | 1503.3ms/step | 13.3 samples/s
+Validation at step 4250: Loss = 3.3703, Perplexity = 29.09
+Step 4300/44000 | Loss: 3.3644 | LR: 4.92e-04 | 1745.1ms/step | 11.5 samples/s
+Step 4400/44000 | Loss: 3.3603 | LR: 4.91e-04 | 1502.7ms/step | 13.3 samples/s
+Step 4500/44000 | Loss: 3.3492 | LR: 4.91e-04 | 1503.4ms/step | 13.3 samples/s
+Validation at step 4500: Loss = 3.3517, Perplexity = 28.55
+Step 4600/44000 | Loss: 3.3282 | LR: 4.90e-04 | 1744.6ms/step | 11.5 samples/s
+Step 4700/44000 | Loss: 3.3278 | LR: 4.90e-04 | 1501.3ms/step | 13.3 samples/s
+Validation at step 4750: Loss = 3.3343, Perplexity = 28.06
+Step 4800/44000 | Loss: 3.3213 | LR: 4.89e-04 | 1746.1ms/step | 11.5 samples/s
+Step 4900/44000 | Loss: 3.3126 | LR: 4.89e-04 | 1502.6ms/step | 13.3 samples/s
+Step 5000/44000 | Loss: 3.3031 | LR: 4.88e-04 | 1502.7ms/step | 13.3 samples/s
+Validation at step 5000: Loss = 3.3192, Perplexity = 27.64
+Step 5100/44000 | Loss: 3.3089 | LR: 4.88e-04 | 1745.7ms/step | 11.5 samples/s
+Step 5200/44000 | Loss: 3.3100 | LR: 4.87e-04 | 1499.2ms/step | 13.3 samples/s
+Validation at step 5250: Loss = 3.3052, Perplexity = 27.25
+Step 5300/44000 | Loss: 3.2863 | LR: 4.86e-04 | 1744.7ms/step | 11.5 samples/s
+Step 5400/44000 | Loss: 3.2940 | LR: 4.86e-04 | 1502.5ms/step | 13.3 samples/s
+Step 5500/44000 | Loss: 3.2883 | LR: 4.85e-04 | 1503.5ms/step | 13.3 samples/s
+Validation at step 5500: Loss = 3.2924, Perplexity = 26.91
+Step 5600/44000 | Loss: 3.2723 | LR: 4.85e-04 | 1744.4ms/step | 11.5 samples/s
+Step 5700/44000 | Loss: 3.2723 | LR: 4.84e-04 | 1501.3ms/step | 13.3 samples/s
+Validation at step 5750: Loss = 3.2807, Perplexity = 26.59
+Step 5800/44000 | Loss: 3.2600 | LR: 4.83e-04 | 1744.5ms/step | 11.5 samples/s
+Step 5900/44000 | Loss: 3.2678 | LR: 4.83e-04 | 1502.7ms/step | 13.3 samples/s
+Step 6000/44000 | Loss: 3.2496 | LR: 4.82e-04 | 1501.6ms/step | 13.3 samples/s
+Validation at step 6000: Loss = 3.2681, Perplexity = 26.26
+Step 6100/44000 | Loss: 3.2571 | LR: 4.81e-04 | 1744.6ms/step | 11.5 samples/s
+Step 6200/44000 | Loss: 3.2435 | LR: 4.81e-04 | 1501.8ms/step | 13.3 samples/s
+Validation at step 6250: Loss = 3.2596, Perplexity = 26.04
+Step 6300/44000 | Loss: 3.2543 | LR: 4.80e-04 | 1742.7ms/step | 11.5 samples/s
+Step 6400/44000 | Loss: 3.2393 | LR: 4.79e-04 | 1502.8ms/step | 13.3 samples/s
+Step 6500/44000 | Loss: 3.2266 | LR: 4.79e-04 | 1502.4ms/step | 13.3 samples/s
+Validation at step 6500: Loss = 3.2486, Perplexity = 25.75
+Step 6600/44000 | Loss: 3.2433 | LR: 4.78e-04 | 1744.3ms/step | 11.5 samples/s
+Step 6700/44000 | Loss: 3.2293 | LR: 4.77e-04 | 1502.8ms/step | 13.3 samples/s
+Validation at step 6750: Loss = 3.2414, Perplexity = 25.57
+Step 6800/44000 | Loss: 3.2254 | LR: 4.76e-04 | 1744.5ms/step | 11.5 samples/s
+Step 6900/44000 | Loss: 3.2247 | LR: 4.75e-04 | 1502.4ms/step | 13.3 samples/s
+Step 7000/44000 | Loss: 3.2203 | LR: 4.75e-04 | 1499.6ms/step | 13.3 samples/s
+Validation at step 7000: Loss = 3.2317, Perplexity = 25.32
+Step 7100/44000 | Loss: 3.2162 | LR: 4.74e-04 | 1745.1ms/step | 11.5 samples/s
+Step 7200/44000 | Loss: 3.2167 | LR: 4.73e-04 | 1504.3ms/step | 13.3 samples/s
+Validation at step 7250: Loss = 3.2221, Perplexity = 25.08
+Step 7300/44000 | Loss: 3.2219 | LR: 4.72e-04 | 1745.8ms/step | 11.5 samples/s
+Step 7400/44000 | Loss: 3.2144 | LR: 4.71e-04 | 1503.0ms/step | 13.3 samples/s
+Step 7500/44000 | Loss: 3.2099 | LR: 4.71e-04 | 1503.9ms/step | 13.3 samples/s
+Validation at step 7500: Loss = 3.2146, Perplexity = 24.89
+Step 7600/44000 | Loss: 3.2063 | LR: 4.70e-04 | 1745.8ms/step | 11.5 samples/s
+Step 7700/44000 | Loss: 3.2042 | LR: 4.69e-04 | 1502.6ms/step | 13.3 samples/s
+Validation at step 7750: Loss = 3.2066, Perplexity = 24.70
+Step 7800/44000 | Loss: 3.1960 | LR: 4.68e-04 | 1745.8ms/step | 11.5 samples/s
+Step 7900/44000 | Loss: 3.1865 | LR: 4.67e-04 | 1503.4ms/step | 13.3 samples/s
+Step 8000/44000 | Loss: 3.1890 | LR: 4.66e-04 | 1504.6ms/step | 13.3 samples/s
+Validation at step 8000: Loss = 3.1980, Perplexity = 24.48
+Step 8100/44000 | Loss: 3.1743 | LR: 4.65e-04 | 1747.9ms/step | 11.4 samples/s
+Step 8200/44000 | Loss: 3.1829 | LR: 4.64e-04 | 1504.8ms/step | 13.3 samples/s
+Validation at step 8250: Loss = 3.1917, Perplexity = 24.33
+Step 8300/44000 | Loss: 3.1795 | LR: 4.63e-04 | 1746.3ms/step | 11.5 samples/s
+Step 8400/44000 | Loss: 3.1815 | LR: 4.62e-04 | 1503.7ms/step | 13.3 samples/s
+Step 8500/44000 | Loss: 3.1765 | LR: 4.61e-04 | 1501.7ms/step | 13.3 samples/s
+Validation at step 8500: Loss = 3.1842, Perplexity = 24.15
+Step 8600/44000 | Loss: 3.1726 | LR: 4.60e-04 | 1743.7ms/step | 11.5 samples/s
+Step 8700/44000 | Loss: 3.1608 | LR: 4.59e-04 | 1502.5ms/step | 13.3 samples/s
+Validation at step 8750: Loss = 3.1772, Perplexity = 23.98
+Step 8800/44000 | Loss: 3.1588 | LR: 4.58e-04 | 1745.5ms/step | 11.5 samples/s
+Step 8900/44000 | Loss: 3.1598 | LR: 4.57e-04 | 1501.4ms/step | 13.3 samples/s
+Step 9000/44000 | Loss: 3.1605 | LR: 4.56e-04 | 1501.3ms/step | 13.3 samples/s
+Validation at step 9000: Loss = 3.1706, Perplexity = 23.82
+Step 9100/44000 | Loss: 3.1585 | LR: 4.55e-04 | 1746.7ms/step | 11.5 samples/s
+Step 9200/44000 | Loss: 3.1446 | LR: 4.54e-04 | 1502.5ms/step | 13.3 samples/s
+Validation at step 9250: Loss = 3.1649, Perplexity = 23.69
+Step 9300/44000 | Loss: 3.1586 | LR: 4.53e-04 | 1743.7ms/step | 11.5 samples/s
+Step 9400/44000 | Loss: 3.1465 | LR: 4.52e-04 | 1503.3ms/step | 13.3 samples/s
+Step 9500/44000 | Loss: 3.1441 | LR: 4.51e-04 | 1501.0ms/step | 13.3 samples/s
+Validation at step 9500: Loss = 3.1589, Perplexity = 23.55
+Step 9600/44000 | Loss: 3.1510 | LR: 4.50e-04 | 1743.2ms/step | 11.5 samples/s
+Step 9700/44000 | Loss: 3.1472 | LR: 4.49e-04 | 1502.9ms/step | 13.3 samples/s
+Validation at step 9750: Loss = 3.1531, Perplexity = 23.41
+Step 9800/44000 | Loss: 3.1383 | LR: 4.48e-04 | 1743.5ms/step | 11.5 samples/s
+Step 9900/44000 | Loss: 3.1445 | LR: 4.47e-04 | 1504.0ms/step | 13.3 samples/s
+Step 10000/44000 | Loss: 3.1389 | LR: 4.46e-04 | 1503.8ms/step | 13.3 samples/s
+Validation at step 10000: Loss = 3.1468, Perplexity = 23.26
+Step 10100/44000 | Loss: 3.1305 | LR: 4.45e-04 | 1744.7ms/step | 11.5 samples/s
+Step 10200/44000 | Loss: 3.1405 | LR: 4.43e-04 | 1502.9ms/step | 13.3 samples/s
+Validation at step 10250: Loss = 3.1416, Perplexity = 23.14
+Step 10300/44000 | Loss: 3.1159 | LR: 4.42e-04 | 1742.2ms/step | 11.5 samples/s
+Step 10400/44000 | Loss: 3.1339 | LR: 4.41e-04 | 1500.8ms/step | 13.3 samples/s
+Step 10500/44000 | Loss: 3.1309 | LR: 4.40e-04 | 1501.2ms/step | 13.3 samples/s
+Validation at step 10500: Loss = 3.1365, Perplexity = 23.02
+Step 10600/44000 | Loss: 3.1209 | LR: 4.39e-04 | 1744.0ms/step | 11.5 samples/s
+Step 10700/44000 | Loss: 3.1139 | LR: 4.37e-04 | 1501.4ms/step | 13.3 samples/s
+Validation at step 10750: Loss = 3.1326, Perplexity = 22.93
+Step 10800/44000 | Loss: 3.1086 | LR: 4.36e-04 | 1744.4ms/step | 11.5 samples/s
+Step 10900/44000 | Loss: 3.1249 | LR: 4.35e-04 | 1504.1ms/step | 13.3 samples/s
+Step 11000/44000 | Loss: 3.1126 | LR: 4.34e-04 | 1504.8ms/step | 13.3 samples/s
+Validation at step 11000: Loss = 3.1267, Perplexity = 22.80
+Step 11100/44000 | Loss: 3.1125 | LR: 4.33e-04 | 1748.6ms/step | 11.4 samples/s
+Step 11200/44000 | Loss: 3.1057 | LR: 4.31e-04 | 1504.3ms/step | 13.3 samples/s
+Validation at step 11250: Loss = 3.1227, Perplexity = 22.71
+Step 11300/44000 | Loss: 3.1067 | LR: 4.30e-04 | 1849.6ms/step | 10.8 samples/s
+Step 11400/44000 | Loss: 3.1069 | LR: 4.29e-04 | 1502.1ms/step | 13.3 samples/s
+Step 11500/44000 | Loss: 3.1053 | LR: 4.28e-04 | 1503.0ms/step | 13.3 samples/s
+Validation at step 11500: Loss = 3.1172, Perplexity = 22.58
+Step 11600/44000 | Loss: 3.1020 | LR: 4.26e-04 | 1748.7ms/step | 11.4 samples/s
+Step 11700/44000 | Loss: 3.0952 | LR: 4.25e-04 | 1504.7ms/step | 13.3 samples/s
+Validation at step 11750: Loss = 3.1133, Perplexity = 22.50
+Step 11800/44000 | Loss: 3.1153 | LR: 4.24e-04 | 1748.0ms/step | 11.4 samples/s
+Step 11900/44000 | Loss: 3.1003 | LR: 4.22e-04 | 1506.6ms/step | 13.3 samples/s
+Step 12000/44000 | Loss: 3.0935 | LR: 4.21e-04 | 1508.0ms/step | 13.3 samples/s
+Validation at step 12000: Loss = 3.1081, Perplexity = 22.38
+Step 12100/44000 | Loss: 3.1009 | LR: 4.20e-04 | 1748.9ms/step | 11.4 samples/s
+Step 12200/44000 | Loss: 3.0923 | LR: 4.18e-04 | 1504.8ms/step | 13.3 samples/s
+Validation at step 12250: Loss = 3.1043, Perplexity = 22.29
+Step 12300/44000 | Loss: 3.0913 | LR: 4.17e-04 | 1749.0ms/step | 11.4 samples/s
+Step 12400/44000 | Loss: 3.0919 | LR: 4.16e-04 | 1505.1ms/step | 13.3 samples/s
+Step 12500/44000 | Loss: 3.0814 | LR: 4.14e-04 | 1505.5ms/step | 13.3 samples/s
+Validation at step 12500: Loss = 3.1012, Perplexity = 22.22
+Step 12600/44000 | Loss: 3.0923 | LR: 4.13e-04 | 1749.2ms/step | 11.4 samples/s
+Step 12700/44000 | Loss: 3.0961 | LR: 4.12e-04 | 1507.2ms/step | 13.3 samples/s
+Validation at step 12750: Loss = 3.0952, Perplexity = 22.09
+Step 12800/44000 | Loss: 3.0742 | LR: 4.10e-04 | 1749.2ms/step | 11.4 samples/s
+Step 12900/44000 | Loss: 3.0780 | LR: 4.09e-04 | 1506.4ms/step | 13.3 samples/s
+Step 13000/44000 | Loss: 3.0868 | LR: 4.07e-04 | 1506.5ms/step | 13.3 samples/s
+Validation at step 13000: Loss = 3.0925, Perplexity = 22.03
+Step 13100/44000 | Loss: 3.0775 | LR: 4.06e-04 | 1750.2ms/step | 11.4 samples/s
+Step 13200/44000 | Loss: 3.0868 | LR: 4.05e-04 | 1504.6ms/step | 13.3 samples/s
+Validation at step 13250: Loss = 3.0881, Perplexity = 21.94
+Step 13300/44000 | Loss: 3.0754 | LR: 4.03e-04 | 1748.7ms/step | 11.4 samples/s
+Step 13400/44000 | Loss: 3.0752 | LR: 4.02e-04 | 1506.9ms/step | 13.3 samples/s
+Step 13500/44000 | Loss: 3.0687 | LR: 4.00e-04 | 1506.2ms/step | 13.3 samples/s
+Validation at step 13500: Loss = 3.0844, Perplexity = 21.85
+Step 13600/44000 | Loss: 3.0677 | LR: 3.99e-04 | 1749.1ms/step | 11.4 samples/s
+Step 13700/44000 | Loss: 3.0661 | LR: 3.97e-04 | 1506.4ms/step | 13.3 samples/s
+Validation at step 13750: Loss = 3.0808, Perplexity = 21.78
+Step 13800/44000 | Loss: 3.0762 | LR: 3.96e-04 | 1748.9ms/step | 11.4 samples/s
+Step 13900/44000 | Loss: 3.0644 | LR: 3.94e-04 | 1506.7ms/step | 13.3 samples/s
+Step 14000/44000 | Loss: 3.0605 | LR: 3.93e-04 | 1500.1ms/step | 13.3 samples/s
+Validation at step 14000: Loss = 3.0766, Perplexity = 21.69
+Step 14100/44000 | Loss: 3.0565 | LR: 3.91e-04 | 1742.1ms/step | 11.5 samples/s
+Step 14200/44000 | Loss: 3.0685 | LR: 3.90e-04 | 1501.2ms/step | 13.3 samples/s

epfl_final/logs/ALIBI_20260921/seed42.pid

@@ -0,0 +1 @@
+4774

epfl_final/logs/ALIBI_20260921/seed43.log

@@ -0,0 +1,376 @@
+New run: Checkpoints will be saved to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s43/run_20260921_030452
+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_030452",
+  "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": 22000,
+  "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_030515-ut49ok16
+Starting training...
+Switching attention mode for step 1: sparse_stochastic (configured sparse attention)
+Step 100/44000 | Loss: 10.6180 | LR: 6.67e-05 | 1940.2ms/step | 10.3 samples/s
+Step 200/44000 | Loss: 7.3253 | LR: 1.33e-04 | 1467.4ms/step | 13.6 samples/s
+Validation at step 250: Loss = 6.2983, Perplexity = 543.67
+Step 300/44000 | Loss: 6.3205 | LR: 2.00e-04 | 1863.0ms/step | 10.7 samples/s
+Step 400/44000 | Loss: 5.8527 | LR: 2.67e-04 | 1474.6ms/step | 13.6 samples/s
+Step 500/44000 | Loss: 5.4548 | LR: 3.33e-04 | 1476.1ms/step | 13.5 samples/s
+Validation at step 500: Loss = 5.2838, Perplexity = 197.11
+Step 600/44000 | Loss: 5.1464 | LR: 4.00e-04 | 1715.1ms/step | 11.7 samples/s
+Step 700/44000 | Loss: 4.8593 | LR: 4.67e-04 | 1476.9ms/step | 13.5 samples/s
+Validation at step 750: Loss = 4.5986, Perplexity = 99.34
+Step 800/44000 | Loss: 4.6039 | LR: 5.00e-04 | 1714.5ms/step | 11.7 samples/s
+Step 900/44000 | Loss: 4.3826 | LR: 5.00e-04 | 1475.5ms/step | 13.6 samples/s
+Step 1000/44000 | Loss: 4.2278 | LR: 5.00e-04 | 1475.9ms/step | 13.6 samples/s
+Validation at step 1000: Loss = 4.1560, Perplexity = 63.82
+Step 1100/44000 | Loss: 4.1174 | LR: 5.00e-04 | 1714.5ms/step | 11.7 samples/s
+Step 1200/44000 | Loss: 4.0258 | LR: 5.00e-04 | 1475.8ms/step | 13.6 samples/s
+Validation at step 1250: Loss = 3.9436, Perplexity = 51.60
+Step 1300/44000 | Loss: 3.9551 | LR: 5.00e-04 | 1714.1ms/step | 11.7 samples/s
+Step 1400/44000 | Loss: 3.8988 | LR: 5.00e-04 | 1476.4ms/step | 13.5 samples/s
+Step 1500/44000 | Loss: 3.8390 | LR: 5.00e-04 | 1475.1ms/step | 13.6 samples/s
+Validation at step 1500: Loss = 3.8080, Perplexity = 45.06
+Step 1600/44000 | Loss: 3.8067 | LR: 5.00e-04 | 1712.5ms/step | 11.7 samples/s
+Step 1700/44000 | Loss: 3.7561 | LR: 4.99e-04 | 1474.2ms/step | 13.6 samples/s
+Validation at step 1750: Loss = 3.7151, Perplexity = 41.06
+Step 1800/44000 | Loss: 3.7215 | LR: 4.99e-04 | 1714.5ms/step | 11.7 samples/s
+Step 1900/44000 | Loss: 3.6968 | LR: 4.99e-04 | 1476.3ms/step | 13.5 samples/s
+Step 2000/44000 | Loss: 3.6696 | LR: 4.99e-04 | 1476.1ms/step | 13.5 samples/s
+Validation at step 2000: Loss = 3.6386, Perplexity = 38.04
+Step 2100/44000 | Loss: 3.6341 | LR: 4.99e-04 | 1713.7ms/step | 11.7 samples/s
+Step 2200/44000 | Loss: 3.6144 | LR: 4.99e-04 | 1475.7ms/step | 13.6 samples/s
+Validation at step 2250: Loss = 3.5822, Perplexity = 35.95
+Step 2300/44000 | Loss: 3.5909 | LR: 4.98e-04 | 1714.7ms/step | 11.7 samples/s
+Step 2400/44000 | Loss: 3.5714 | LR: 4.98e-04 | 1476.1ms/step | 13.5 samples/s
+Step 2500/44000 | Loss: 3.5600 | LR: 4.98e-04 | 1474.7ms/step | 13.6 samples/s
+Validation at step 2500: Loss = 3.5375, Perplexity = 34.38
+Step 2600/44000 | Loss: 3.5371 | LR: 4.98e-04 | 1713.9ms/step | 11.7 samples/s
+Step 2700/44000 | Loss: 3.5336 | LR: 4.97e-04 | 1475.5ms/step | 13.6 samples/s
+Validation at step 2750: Loss = 3.4989, Perplexity = 33.08
+Step 2800/44000 | Loss: 3.5104 | LR: 4.97e-04 | 1713.9ms/step | 11.7 samples/s
+Step 2900/44000 | Loss: 3.4923 | LR: 4.97e-04 | 1476.4ms/step | 13.5 samples/s
+Step 3000/44000 | Loss: 3.4826 | LR: 4.97e-04 | 1475.9ms/step | 13.6 samples/s
+Validation at step 3000: Loss = 3.4634, Perplexity = 31.93
+Step 3100/44000 | Loss: 3.4701 | LR: 4.96e-04 | 1714.1ms/step | 11.7 samples/s
+Step 3200/44000 | Loss: 3.4575 | LR: 4.96e-04 | 1475.1ms/step | 13.6 samples/s
+Validation at step 3250: Loss = 3.4348, Perplexity = 31.03
+Step 3300/44000 | Loss: 3.4432 | LR: 4.96e-04 | 1713.6ms/step | 11.7 samples/s
+Step 3400/44000 | Loss: 3.4324 | LR: 4.95e-04 | 1476.4ms/step | 13.5 samples/s
+Step 3500/44000 | Loss: 3.4300 | LR: 4.95e-04 | 1476.3ms/step | 13.5 samples/s
+Validation at step 3500: Loss = 3.4103, Perplexity = 30.27
+Step 3600/44000 | Loss: 3.4195 | LR: 4.95e-04 | 1713.9ms/step | 11.7 samples/s
+Step 3700/44000 | Loss: 3.3984 | LR: 4.94e-04 | 1475.6ms/step | 13.6 samples/s
+Validation at step 3750: Loss = 3.3852, Perplexity = 29.52
+Step 3800/44000 | Loss: 3.3968 | LR: 4.94e-04 | 1714.2ms/step | 11.7 samples/s
+Step 3900/44000 | Loss: 3.3933 | LR: 4.93e-04 | 1476.0ms/step | 13.6 samples/s
+Step 4000/44000 | Loss: 3.3835 | LR: 4.93e-04 | 1475.6ms/step | 13.6 samples/s
+Validation at step 4000: Loss = 3.3641, Perplexity = 28.91
+Step 4100/44000 | Loss: 3.3577 | LR: 4.93e-04 | 1714.5ms/step | 11.7 samples/s
+Step 4200/44000 | Loss: 3.3607 | LR: 4.92e-04 | 1476.2ms/step | 13.5 samples/s
+Validation at step 4250: Loss = 3.3456, Perplexity = 28.38
+Step 4300/44000 | Loss: 3.3509 | LR: 4.92e-04 | 1715.0ms/step | 11.7 samples/s
+Step 4400/44000 | Loss: 3.3449 | LR: 4.91e-04 | 1476.9ms/step | 13.5 samples/s
+Step 4500/44000 | Loss: 3.3409 | LR: 4.91e-04 | 1475.9ms/step | 13.6 samples/s
+Validation at step 4500: Loss = 3.3292, Perplexity = 27.92
+Step 4600/44000 | Loss: 3.3290 | LR: 4.90e-04 | 1713.9ms/step | 11.7 samples/s
+Step 4700/44000 | Loss: 3.3310 | LR: 4.90e-04 | 1475.3ms/step | 13.6 samples/s
+Validation at step 4750: Loss = 3.3107, Perplexity = 27.40
+Step 4800/44000 | Loss: 3.3248 | LR: 4.89e-04 | 1714.2ms/step | 11.7 samples/s
+Step 4900/44000 | Loss: 3.3212 | LR: 4.89e-04 | 1476.8ms/step | 13.5 samples/s
+Step 5000/44000 | Loss: 3.2953 | LR: 4.88e-04 | 1476.4ms/step | 13.5 samples/s
+Validation at step 5000: Loss = 3.2953, Perplexity = 26.99
+Step 5100/44000 | Loss: 3.3119 | LR: 4.88e-04 | 1714.6ms/step | 11.7 samples/s
+Step 5200/44000 | Loss: 3.2905 | LR: 4.87e-04 | 1475.4ms/step | 13.6 samples/s
+Validation at step 5250: Loss = 3.2818, Perplexity = 26.62
+Step 5300/44000 | Loss: 3.2814 | LR: 4.86e-04 | 1711.7ms/step | 11.7 samples/s
+Step 5400/44000 | Loss: 3.2842 | LR: 4.86e-04 | 1475.2ms/step | 13.6 samples/s
+Step 5500/44000 | Loss: 3.2809 | LR: 4.85e-04 | 1475.4ms/step | 13.6 samples/s
+Validation at step 5500: Loss = 3.2676, Perplexity = 26.25
+Step 5600/44000 | Loss: 3.2786 | LR: 4.85e-04 | 1713.8ms/step | 11.7 samples/s
+Step 5700/44000 | Loss: 3.2780 | LR: 4.84e-04 | 1476.4ms/step | 13.5 samples/s
+Validation at step 5750: Loss = 3.2552, Perplexity = 25.93
+Step 5800/44000 | Loss: 3.2575 | LR: 4.83e-04 | 1713.9ms/step | 11.7 samples/s
+Step 5900/44000 | Loss: 3.2616 | LR: 4.83e-04 | 1475.9ms/step | 13.6 samples/s
+Step 6000/44000 | Loss: 3.2546 | LR: 4.82e-04 | 1475.8ms/step | 13.6 samples/s
+Validation at step 6000: Loss = 3.2447, Perplexity = 25.65
+Step 6100/44000 | Loss: 3.2601 | LR: 4.81e-04 | 1713.3ms/step | 11.7 samples/s
+Step 6200/44000 | Loss: 3.2521 | LR: 4.81e-04 | 1475.6ms/step | 13.6 samples/s
+Validation at step 6250: Loss = 3.2337, Perplexity = 25.37
+Step 6300/44000 | Loss: 3.2551 | LR: 4.80e-04 | 1714.6ms/step | 11.7 samples/s
+Step 6400/44000 | Loss: 3.2349 | LR: 4.79e-04 | 1475.1ms/step | 13.6 samples/s
+Step 6500/44000 | Loss: 3.2342 | LR: 4.79e-04 | 1476.2ms/step | 13.5 samples/s
+Validation at step 6500: Loss = 3.2260, Perplexity = 25.18
+Step 6600/44000 | Loss: 3.2311 | LR: 4.78e-04 | 1714.9ms/step | 11.7 samples/s
+Step 6700/44000 | Loss: 3.2358 | LR: 4.77e-04 | 1476.1ms/step | 13.5 samples/s
+Validation at step 6750: Loss = 3.2151, Perplexity = 24.90
+Step 6800/44000 | Loss: 3.2363 | LR: 4.76e-04 | 1715.0ms/step | 11.7 samples/s
+Step 6900/44000 | Loss: 3.2318 | LR: 4.75e-04 | 1474.8ms/step | 13.6 samples/s
+Step 7000/44000 | Loss: 3.2119 | LR: 4.75e-04 | 1476.5ms/step | 13.5 samples/s
+Validation at step 7000: Loss = 3.2052, Perplexity = 24.66
+Step 7100/44000 | Loss: 3.2124 | LR: 4.74e-04 | 1713.8ms/step | 11.7 samples/s
+Step 7200/44000 | Loss: 3.2218 | LR: 4.73e-04 | 1476.7ms/step | 13.5 samples/s
+Validation at step 7250: Loss = 3.1978, Perplexity = 24.48
+Step 7300/44000 | Loss: 3.2093 | LR: 4.72e-04 | 1715.4ms/step | 11.7 samples/s
+Step 7400/44000 | Loss: 3.1988 | LR: 4.71e-04 | 1475.8ms/step | 13.6 samples/s
+Step 7500/44000 | Loss: 3.1996 | LR: 4.71e-04 | 1476.8ms/step | 13.5 samples/s
+Validation at step 7500: Loss = 3.1883, Perplexity = 24.25
+Step 7600/44000 | Loss: 3.1994 | LR: 4.70e-04 | 1714.5ms/step | 11.7 samples/s
+Step 7700/44000 | Loss: 3.1922 | LR: 4.69e-04 | 1475.4ms/step | 13.6 samples/s
+Validation at step 7750: Loss = 3.1814, Perplexity = 24.08
+Step 7800/44000 | Loss: 3.1948 | LR: 4.68e-04 | 1714.0ms/step | 11.7 samples/s
+Step 7900/44000 | Loss: 3.1895 | LR: 4.67e-04 | 1476.2ms/step | 13.5 samples/s
+Step 8000/44000 | Loss: 3.1872 | LR: 4.66e-04 | 1477.4ms/step | 13.5 samples/s
+Validation at step 8000: Loss = 3.1733, Perplexity = 23.89
+Step 8100/44000 | Loss: 3.1819 | LR: 4.65e-04 | 1714.9ms/step | 11.7 samples/s
+Step 8200/44000 | Loss: 3.1919 | LR: 4.64e-04 | 1476.4ms/step | 13.5 samples/s
+Validation at step 8250: Loss = 3.1658, Perplexity = 23.71
+Step 8300/44000 | Loss: 3.1801 | LR: 4.63e-04 | 1715.3ms/step | 11.7 samples/s
+Step 8400/44000 | Loss: 3.1789 | LR: 4.62e-04 | 1476.5ms/step | 13.5 samples/s
+Step 8500/44000 | Loss: 3.1733 | LR: 4.61e-04 | 1476.7ms/step | 13.5 samples/s
+Validation at step 8500: Loss = 3.1582, Perplexity = 23.53
+Step 8600/44000 | Loss: 3.1758 | LR: 4.60e-04 | 1712.9ms/step | 11.7 samples/s
+Step 8700/44000 | Loss: 3.1537 | LR: 4.59e-04 | 1476.0ms/step | 13.5 samples/s
+Validation at step 8750: Loss = 3.1525, Perplexity = 23.39
+Step 8800/44000 | Loss: 3.1675 | LR: 4.58e-04 | 1714.7ms/step | 11.7 samples/s
+Step 8900/44000 | Loss: 3.1559 | LR: 4.57e-04 | 1476.3ms/step | 13.5 samples/s
+Step 9000/44000 | Loss: 3.1599 | LR: 4.56e-04 | 1474.9ms/step | 13.6 samples/s
+Validation at step 9000: Loss = 3.1472, Perplexity = 23.27
+Step 9100/44000 | Loss: 3.1515 | LR: 4.55e-04 | 1712.9ms/step | 11.7 samples/s
+Step 9200/44000 | Loss: 3.1619 | LR: 4.54e-04 | 1475.7ms/step | 13.6 samples/s
+Validation at step 9250: Loss = 3.1393, Perplexity = 23.09
+Step 9300/44000 | Loss: 3.1502 | LR: 4.53e-04 | 1715.4ms/step | 11.7 samples/s
+Step 9400/44000 | Loss: 3.1578 | LR: 4.52e-04 | 1474.9ms/step | 13.6 samples/s
+Step 9500/44000 | Loss: 3.1459 | LR: 4.51e-04 | 1475.6ms/step | 13.6 samples/s
+Validation at step 9500: Loss = 3.1331, Perplexity = 22.95
+Step 9600/44000 | Loss: 3.1550 | LR: 4.50e-04 | 1714.4ms/step | 11.7 samples/s
+Step 9700/44000 | Loss: 3.1496 | LR: 4.49e-04 | 1475.1ms/step | 13.6 samples/s
+Validation at step 9750: Loss = 3.1282, Perplexity = 22.83
+Step 9800/44000 | Loss: 3.1465 | LR: 4.48e-04 | 1714.0ms/step | 11.7 samples/s
+Step 9900/44000 | Loss: 3.1365 | LR: 4.47e-04 | 1475.9ms/step | 13.6 samples/s
+Step 10000/44000 | Loss: 3.1348 | LR: 4.46e-04 | 1475.4ms/step | 13.6 samples/s
+Validation at step 10000: Loss = 3.1215, Perplexity = 22.68
+Step 10100/44000 | Loss: 3.1426 | LR: 4.45e-04 | 1714.6ms/step | 11.7 samples/s
+Step 10200/44000 | Loss: 3.1363 | LR: 4.43e-04 | 1476.8ms/step | 13.5 samples/s
+Validation at step 10250: Loss = 3.1172, Perplexity = 22.58
+Step 10300/44000 | Loss: 3.1341 | LR: 4.42e-04 | 1714.5ms/step | 11.7 samples/s
+Step 10400/44000 | Loss: 3.1369 | LR: 4.41e-04 | 1476.1ms/step | 13.5 samples/s
+Step 10500/44000 | Loss: 3.1289 | LR: 4.40e-04 | 1474.6ms/step | 13.6 samples/s
+Validation at step 10500: Loss = 3.1116, Perplexity = 22.46
+Step 10600/44000 | Loss: 3.1137 | LR: 4.39e-04 | 1713.8ms/step | 11.7 samples/s
+Step 10700/44000 | Loss: 3.1205 | LR: 4.37e-04 | 1474.4ms/step | 13.6 samples/s
+Validation at step 10750: Loss = 3.1063, Perplexity = 22.34
+Step 10800/44000 | Loss: 3.1210 | LR: 4.36e-04 | 1714.0ms/step | 11.7 samples/s
+Step 10900/44000 | Loss: 3.1163 | LR: 4.35e-04 | 1475.5ms/step | 13.6 samples/s
+Step 11000/44000 | Loss: 3.1236 | LR: 4.34e-04 | 1476.4ms/step | 13.5 samples/s
+Validation at step 11000: Loss = 3.1006, Perplexity = 22.21
+Step 11100/44000 | Loss: 3.1182 | LR: 4.33e-04 | 1714.7ms/step | 11.7 samples/s
+Step 11200/44000 | Loss: 3.1122 | LR: 4.31e-04 | 1476.3ms/step | 13.5 samples/s
+Validation at step 11250: Loss = 3.0962, Perplexity = 22.11
+Step 11300/44000 | Loss: 3.1026 | LR: 4.30e-04 | 1714.8ms/step | 11.7 samples/s
+Step 11400/44000 | Loss: 3.1060 | LR: 4.29e-04 | 1475.4ms/step | 13.6 samples/s
+Step 11500/44000 | Loss: 3.1111 | LR: 4.28e-04 | 1852.6ms/step | 10.8 samples/s
+Validation at step 11500: Loss = 3.0905, Perplexity = 21.99
+Step 11600/44000 | Loss: 3.1106 | LR: 4.26e-04 | 1730.5ms/step | 11.6 samples/s
+Step 11700/44000 | Loss: 3.1032 | LR: 4.25e-04 | 1483.4ms/step | 13.5 samples/s
+Validation at step 11750: Loss = 3.0861, Perplexity = 21.89
+Step 11800/44000 | Loss: 3.0917 | LR: 4.24e-04 | 1716.3ms/step | 11.7 samples/s
+Step 11900/44000 | Loss: 3.0956 | LR: 4.22e-04 | 1476.0ms/step | 13.6 samples/s
+Step 12000/44000 | Loss: 3.0914 | LR: 4.21e-04 | 1477.3ms/step | 13.5 samples/s
+Validation at step 12000: Loss = 3.0818, Perplexity = 21.80
+Step 12100/44000 | Loss: 3.0998 | LR: 4.20e-04 | 1715.4ms/step | 11.7 samples/s
+Step 12200/44000 | Loss: 3.0901 | LR: 4.18e-04 | 1476.9ms/step | 13.5 samples/s
+Validation at step 12250: Loss = 3.0782, Perplexity = 21.72
+Step 12300/44000 | Loss: 3.0781 | LR: 4.17e-04 | 1715.2ms/step | 11.7 samples/s
+Step 12400/44000 | Loss: 3.0901 | LR: 4.16e-04 | 1476.7ms/step | 13.5 samples/s
+Step 12500/44000 | Loss: 3.0875 | LR: 4.14e-04 | 1477.4ms/step | 13.5 samples/s
+Validation at step 12500: Loss = 3.0733, Perplexity = 21.61
+Step 12600/44000 | Loss: 3.0728 | LR: 4.13e-04 | 1714.9ms/step | 11.7 samples/s
+Step 12700/44000 | Loss: 3.0866 | LR: 4.12e-04 | 1475.3ms/step | 13.6 samples/s
+Validation at step 12750: Loss = 3.0692, Perplexity = 21.52
+Step 12800/44000 | Loss: 3.0855 | LR: 4.10e-04 | 1714.7ms/step | 11.7 samples/s
+Step 12900/44000 | Loss: 3.0831 | LR: 4.09e-04 | 1476.9ms/step | 13.5 samples/s
+Step 13000/44000 | Loss: 3.0697 | LR: 4.07e-04 | 1475.9ms/step | 13.6 samples/s
+Validation at step 13000: Loss = 3.0646, Perplexity = 21.43
+Step 13100/44000 | Loss: 3.0776 | LR: 4.06e-04 | 1714.9ms/step | 11.7 samples/s
+Step 13200/44000 | Loss: 3.0756 | LR: 4.05e-04 | 1476.7ms/step | 13.5 samples/s
+Validation at step 13250: Loss = 3.0625, Perplexity = 21.38
+Step 13300/44000 | Loss: 3.0782 | LR: 4.03e-04 | 1715.6ms/step | 11.7 samples/s
+Step 13400/44000 | Loss: 3.0706 | LR: 4.02e-04 | 1475.5ms/step | 13.6 samples/s
+Step 13500/44000 | Loss: 3.0710 | LR: 4.00e-04 | 1476.6ms/step | 13.5 samples/s
+Validation at step 13500: Loss = 3.0560, Perplexity = 21.24
+Step 13600/44000 | Loss: 3.0656 | LR: 3.99e-04 | 1715.2ms/step | 11.7 samples/s
+Step 13700/44000 | Loss: 3.0691 | LR: 3.97e-04 | 1477.2ms/step | 13.5 samples/s
+Validation at step 13750: Loss = 3.0564, Perplexity = 21.25
+Step 13800/44000 | Loss: 3.0752 | LR: 3.96e-04 | 1714.7ms/step | 11.7 samples/s
+Step 13900/44000 | Loss: 3.0639 | LR: 3.94e-04 | 1477.2ms/step | 13.5 samples/s
+Step 14000/44000 | Loss: 3.0624 | LR: 3.93e-04 | 1476.6ms/step | 13.5 samples/s
+Validation at step 14000: Loss = 3.0502, Perplexity = 21.12
+Step 14100/44000 | Loss: 3.0652 | LR: 3.91e-04 | 1716.1ms/step | 11.7 samples/s
+Step 14200/44000 | Loss: 3.0639 | LR: 3.90e-04 | 1476.7ms/step | 13.5 samples/s
+Validation at step 14250: Loss = 3.0456, Perplexity = 21.02
+Step 14300/44000 | Loss: 3.0598 | LR: 3.88e-04 | 1715.0ms/step | 11.7 samples/s
+Step 14400/44000 | Loss: 3.0619 | LR: 3.87e-04 | 1476.3ms/step | 13.5 samples/s

epfl_final/logs/ALIBI_20260921/seed43.pid

@@ -0,0 +1 @@
+4775

epfl_final/logs/ALIBI_20260921/seed44.log

@@ -0,0 +1,372 @@
+New run: Checkpoints will be saved to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s44/run_20260921_030452
+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_030452",
+  "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": 22000,
+  "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_030515-6k91ldtc
+Starting training...
+Switching attention mode for step 1: sparse_stochastic (configured sparse attention)
+Step 100/44000 | Loss: 10.6002 | LR: 6.67e-05 | 1950.1ms/step | 10.3 samples/s
+Step 200/44000 | Loss: 7.3142 | LR: 1.33e-04 | 1497.1ms/step | 13.4 samples/s
+Validation at step 250: Loss = 6.3137, Perplexity = 552.09
+Step 300/44000 | Loss: 6.3192 | LR: 2.00e-04 | 1885.5ms/step | 10.6 samples/s
+Step 400/44000 | Loss: 5.8496 | LR: 2.67e-04 | 1504.5ms/step | 13.3 samples/s
+Step 500/44000 | Loss: 5.4638 | LR: 3.33e-04 | 1505.7ms/step | 13.3 samples/s
+Validation at step 500: Loss = 5.2985, Perplexity = 200.04
+Step 600/44000 | Loss: 5.1429 | LR: 4.00e-04 | 1749.0ms/step | 11.4 samples/s
+Step 700/44000 | Loss: 4.8755 | LR: 4.67e-04 | 1507.4ms/step | 13.3 samples/s
+Validation at step 750: Loss = 4.6363, Perplexity = 103.16
+Step 800/44000 | Loss: 4.6284 | LR: 5.00e-04 | 1756.7ms/step | 11.4 samples/s
+Step 900/44000 | Loss: 4.3980 | LR: 5.00e-04 | 1511.9ms/step | 13.2 samples/s
+Step 1000/44000 | Loss: 4.2338 | LR: 5.00e-04 | 1511.1ms/step | 13.2 samples/s
+Validation at step 1000: Loss = 4.1763, Perplexity = 65.12
+Step 1100/44000 | Loss: 4.1035 | LR: 5.00e-04 | 1757.8ms/step | 11.4 samples/s
+Step 1200/44000 | Loss: 4.0171 | LR: 5.00e-04 | 1512.0ms/step | 13.2 samples/s
+Validation at step 1250: Loss = 3.9508, Perplexity = 51.98
+Step 1300/44000 | Loss: 3.9450 | LR: 5.00e-04 | 1755.8ms/step | 11.4 samples/s
+Step 1400/44000 | Loss: 3.8777 | LR: 5.00e-04 | 1512.0ms/step | 13.2 samples/s
+Step 1500/44000 | Loss: 3.8412 | LR: 5.00e-04 | 1509.9ms/step | 13.2 samples/s
+Validation at step 1500: Loss = 3.8162, Perplexity = 45.43
+Step 1600/44000 | Loss: 3.7897 | LR: 5.00e-04 | 1750.5ms/step | 11.4 samples/s
+Step 1700/44000 | Loss: 3.7536 | LR: 4.99e-04 | 1509.0ms/step | 13.3 samples/s
+Validation at step 1750: Loss = 3.7217, Perplexity = 41.34
+Step 1800/44000 | Loss: 3.7243 | LR: 4.99e-04 | 1753.3ms/step | 11.4 samples/s
+Step 1900/44000 | Loss: 3.6826 | LR: 4.99e-04 | 1507.1ms/step | 13.3 samples/s
+Step 2000/44000 | Loss: 3.6626 | LR: 4.99e-04 | 1508.5ms/step | 13.3 samples/s
+Validation at step 2000: Loss = 3.6522, Perplexity = 38.56
+Step 2100/44000 | Loss: 3.6381 | LR: 4.99e-04 | 1747.4ms/step | 11.4 samples/s
+Step 2200/44000 | Loss: 3.6075 | LR: 4.99e-04 | 1506.6ms/step | 13.3 samples/s
+Validation at step 2250: Loss = 3.5934, Perplexity = 36.36
+Step 2300/44000 | Loss: 3.5938 | LR: 4.98e-04 | 1751.8ms/step | 11.4 samples/s
+Step 2400/44000 | Loss: 3.5725 | LR: 4.98e-04 | 1507.0ms/step | 13.3 samples/s
+Step 2500/44000 | Loss: 3.5596 | LR: 4.98e-04 | 1506.6ms/step | 13.3 samples/s
+Validation at step 2500: Loss = 3.5459, Perplexity = 34.67
+Step 2600/44000 | Loss: 3.5404 | LR: 4.98e-04 | 1750.3ms/step | 11.4 samples/s
+Step 2700/44000 | Loss: 3.5205 | LR: 4.97e-04 | 1506.5ms/step | 13.3 samples/s
+Validation at step 2750: Loss = 3.5054, Perplexity = 33.29
+Step 2800/44000 | Loss: 3.5100 | LR: 4.97e-04 | 1752.4ms/step | 11.4 samples/s
+Step 2900/44000 | Loss: 3.4864 | LR: 4.97e-04 | 1507.7ms/step | 13.3 samples/s
+Step 3000/44000 | Loss: 3.4682 | LR: 4.97e-04 | 1505.7ms/step | 13.3 samples/s
+Validation at step 3000: Loss = 3.4721, Perplexity = 32.20
+Step 3100/44000 | Loss: 3.4670 | LR: 4.96e-04 | 1751.8ms/step | 11.4 samples/s
+Step 3200/44000 | Loss: 3.4540 | LR: 4.96e-04 | 1506.4ms/step | 13.3 samples/s
+Validation at step 3250: Loss = 3.4467, Perplexity = 31.40
+Step 3300/44000 | Loss: 3.4417 | LR: 4.96e-04 | 1751.6ms/step | 11.4 samples/s
+Step 3400/44000 | Loss: 3.4350 | LR: 4.95e-04 | 1508.0ms/step | 13.3 samples/s
+Step 3500/44000 | Loss: 3.4254 | LR: 4.95e-04 | 1507.2ms/step | 13.3 samples/s
+Validation at step 3500: Loss = 3.4185, Perplexity = 30.52
+Step 3600/44000 | Loss: 3.4129 | LR: 4.95e-04 | 1751.3ms/step | 11.4 samples/s
+Step 3700/44000 | Loss: 3.3913 | LR: 4.94e-04 | 1508.2ms/step | 13.3 samples/s
+Validation at step 3750: Loss = 3.3936, Perplexity = 29.77
+Step 3800/44000 | Loss: 3.3967 | LR: 4.94e-04 | 1750.4ms/step | 11.4 samples/s
+Step 3900/44000 | Loss: 3.3904 | LR: 4.93e-04 | 1508.9ms/step | 13.3 samples/s
+Step 4000/44000 | Loss: 3.3750 | LR: 4.93e-04 | 1506.7ms/step | 13.3 samples/s
+Validation at step 4000: Loss = 3.3706, Perplexity = 29.10
+Step 4100/44000 | Loss: 3.3636 | LR: 4.93e-04 | 1751.3ms/step | 11.4 samples/s
+Step 4200/44000 | Loss: 3.3561 | LR: 4.92e-04 | 1508.7ms/step | 13.3 samples/s
+Validation at step 4250: Loss = 3.3545, Perplexity = 28.63
+Step 4300/44000 | Loss: 3.3606 | LR: 4.92e-04 | 1751.9ms/step | 11.4 samples/s
+Step 4400/44000 | Loss: 3.3483 | LR: 4.91e-04 | 1508.8ms/step | 13.3 samples/s
+Step 4500/44000 | Loss: 3.3424 | LR: 4.91e-04 | 1508.4ms/step | 13.3 samples/s
+Validation at step 4500: Loss = 3.3353, Perplexity = 28.09
+Step 4600/44000 | Loss: 3.3269 | LR: 4.90e-04 | 1752.2ms/step | 11.4 samples/s
+Step 4700/44000 | Loss: 3.3173 | LR: 4.90e-04 | 1507.8ms/step | 13.3 samples/s
+Validation at step 4750: Loss = 3.3188, Perplexity = 27.63
+Step 4800/44000 | Loss: 3.3176 | LR: 4.89e-04 | 1752.4ms/step | 11.4 samples/s
+Step 4900/44000 | Loss: 3.3119 | LR: 4.89e-04 | 1508.4ms/step | 13.3 samples/s
+Step 5000/44000 | Loss: 3.2989 | LR: 4.88e-04 | 1509.2ms/step | 13.3 samples/s
+Validation at step 5000: Loss = 3.3046, Perplexity = 27.24
+Step 5100/44000 | Loss: 3.3001 | LR: 4.88e-04 | 1751.1ms/step | 11.4 samples/s
+Step 5200/44000 | Loss: 3.2960 | LR: 4.87e-04 | 1502.9ms/step | 13.3 samples/s
+Validation at step 5250: Loss = 3.2908, Perplexity = 26.86
+Step 5300/44000 | Loss: 3.2901 | LR: 4.86e-04 | 1752.2ms/step | 11.4 samples/s
+Step 5400/44000 | Loss: 3.2808 | LR: 4.86e-04 | 1507.6ms/step | 13.3 samples/s
+Step 5500/44000 | Loss: 3.2723 | LR: 4.85e-04 | 1507.6ms/step | 13.3 samples/s
+Validation at step 5500: Loss = 3.2798, Perplexity = 26.57
+Step 5600/44000 | Loss: 3.2735 | LR: 4.85e-04 | 1749.6ms/step | 11.4 samples/s
+Step 5700/44000 | Loss: 3.2644 | LR: 4.84e-04 | 1506.3ms/step | 13.3 samples/s
+Validation at step 5750: Loss = 3.2660, Perplexity = 26.21
+Step 5800/44000 | Loss: 3.2721 | LR: 4.83e-04 | 1750.7ms/step | 11.4 samples/s
+Step 5900/44000 | Loss: 3.2507 | LR: 4.83e-04 | 1507.2ms/step | 13.3 samples/s
+Step 6000/44000 | Loss: 3.2536 | LR: 4.82e-04 | 1507.5ms/step | 13.3 samples/s
+Validation at step 6000: Loss = 3.2549, Perplexity = 25.92
+Step 6100/44000 | Loss: 3.2559 | LR: 4.81e-04 | 1750.4ms/step | 11.4 samples/s
+Step 6200/44000 | Loss: 3.2505 | LR: 4.81e-04 | 1508.5ms/step | 13.3 samples/s
+Validation at step 6250: Loss = 3.2469, Perplexity = 25.71
+Step 6300/44000 | Loss: 3.2333 | LR: 4.80e-04 | 1750.3ms/step | 11.4 samples/s
+Step 6400/44000 | Loss: 3.2453 | LR: 4.79e-04 | 1508.2ms/step | 13.3 samples/s
+Step 6500/44000 | Loss: 3.2320 | LR: 4.79e-04 | 1507.0ms/step | 13.3 samples/s
+Validation at step 6500: Loss = 3.2352, Perplexity = 25.41
+Step 6600/44000 | Loss: 3.2414 | LR: 4.78e-04 | 1751.0ms/step | 11.4 samples/s
+Step 6700/44000 | Loss: 3.2280 | LR: 4.77e-04 | 1507.5ms/step | 13.3 samples/s
+Validation at step 6750: Loss = 3.2251, Perplexity = 25.16
+Step 6800/44000 | Loss: 3.2169 | LR: 4.76e-04 | 1751.5ms/step | 11.4 samples/s
+Step 6900/44000 | Loss: 3.2159 | LR: 4.75e-04 | 1506.4ms/step | 13.3 samples/s
+Step 7000/44000 | Loss: 3.2147 | LR: 4.75e-04 | 1507.2ms/step | 13.3 samples/s
+Validation at step 7000: Loss = 3.2163, Perplexity = 24.94
+Step 7100/44000 | Loss: 3.2075 | LR: 4.74e-04 | 1752.3ms/step | 11.4 samples/s
+Step 7200/44000 | Loss: 3.2151 | LR: 4.73e-04 | 1508.3ms/step | 13.3 samples/s
+Validation at step 7250: Loss = 3.2121, Perplexity = 24.83
+Step 7300/44000 | Loss: 3.2007 | LR: 4.72e-04 | 1751.3ms/step | 11.4 samples/s
+Step 7400/44000 | Loss: 3.1952 | LR: 4.71e-04 | 1507.5ms/step | 13.3 samples/s
+Step 7500/44000 | Loss: 3.2047 | LR: 4.71e-04 | 1508.9ms/step | 13.3 samples/s
+Validation at step 7500: Loss = 3.2008, Perplexity = 24.55
+Step 7600/44000 | Loss: 3.1950 | LR: 4.70e-04 | 1750.8ms/step | 11.4 samples/s
+Step 7700/44000 | Loss: 3.1895 | LR: 4.69e-04 | 1507.5ms/step | 13.3 samples/s
+Validation at step 7750: Loss = 3.1927, Perplexity = 24.35
+Step 7800/44000 | Loss: 3.1953 | LR: 4.68e-04 | 1751.9ms/step | 11.4 samples/s
+Step 7900/44000 | Loss: 3.1966 | LR: 4.67e-04 | 1507.9ms/step | 13.3 samples/s
+Step 8000/44000 | Loss: 3.1881 | LR: 4.66e-04 | 1508.0ms/step | 13.3 samples/s
+Validation at step 8000: Loss = 3.1851, Perplexity = 24.17
+Step 8100/44000 | Loss: 3.1745 | LR: 4.65e-04 | 1753.0ms/step | 11.4 samples/s
+Step 8200/44000 | Loss: 3.1779 | LR: 4.64e-04 | 1509.1ms/step | 13.3 samples/s
+Validation at step 8250: Loss = 3.1794, Perplexity = 24.03
+Step 8300/44000 | Loss: 3.1783 | LR: 4.63e-04 | 1751.9ms/step | 11.4 samples/s
+Step 8400/44000 | Loss: 3.1626 | LR: 4.62e-04 | 1505.7ms/step | 13.3 samples/s
+Step 8500/44000 | Loss: 3.1763 | LR: 4.61e-04 | 1507.5ms/step | 13.3 samples/s
+Validation at step 8500: Loss = 3.1710, Perplexity = 23.83
+Step 8600/44000 | Loss: 3.1566 | LR: 4.60e-04 | 1750.3ms/step | 11.4 samples/s
+Step 8700/44000 | Loss: 3.1579 | LR: 4.59e-04 | 1506.9ms/step | 13.3 samples/s
+Validation at step 8750: Loss = 3.1659, Perplexity = 23.71
+Step 8800/44000 | Loss: 3.1480 | LR: 4.58e-04 | 1751.3ms/step | 11.4 samples/s
+Step 8900/44000 | Loss: 3.1589 | LR: 4.57e-04 | 1505.4ms/step | 13.3 samples/s
+Step 9000/44000 | Loss: 3.1486 | LR: 4.56e-04 | 1506.5ms/step | 13.3 samples/s
+Validation at step 9000: Loss = 3.1588, Perplexity = 23.54
+Step 9100/44000 | Loss: 3.1540 | LR: 4.55e-04 | 1753.3ms/step | 11.4 samples/s
+Step 9200/44000 | Loss: 3.1469 | LR: 4.54e-04 | 1506.2ms/step | 13.3 samples/s
+Validation at step 9250: Loss = 3.1522, Perplexity = 23.39
+Step 9300/44000 | Loss: 3.1414 | LR: 4.53e-04 | 1750.9ms/step | 11.4 samples/s
+Step 9400/44000 | Loss: 3.1474 | LR: 4.52e-04 | 1506.9ms/step | 13.3 samples/s
+Step 9500/44000 | Loss: 3.1425 | LR: 4.51e-04 | 1504.5ms/step | 13.3 samples/s
+Validation at step 9500: Loss = 3.1457, Perplexity = 23.24
+Step 9600/44000 | Loss: 3.1385 | LR: 4.50e-04 | 1750.7ms/step | 11.4 samples/s
+Step 9700/44000 | Loss: 3.1399 | LR: 4.49e-04 | 1506.3ms/step | 13.3 samples/s
+Validation at step 9750: Loss = 3.1393, Perplexity = 23.09
+Step 9800/44000 | Loss: 3.1351 | LR: 4.48e-04 | 1749.1ms/step | 11.4 samples/s
+Step 9900/44000 | Loss: 3.1312 | LR: 4.47e-04 | 1506.3ms/step | 13.3 samples/s
+Step 10000/44000 | Loss: 3.1300 | LR: 4.46e-04 | 1508.7ms/step | 13.3 samples/s
+Validation at step 10000: Loss = 3.1349, Perplexity = 22.99
+Step 10100/44000 | Loss: 3.1234 | LR: 4.45e-04 | 1749.3ms/step | 11.4 samples/s
+Step 10200/44000 | Loss: 3.1345 | LR: 4.43e-04 | 1505.6ms/step | 13.3 samples/s
+Validation at step 10250: Loss = 3.1289, Perplexity = 22.85
+Step 10300/44000 | Loss: 3.1177 | LR: 4.42e-04 | 1745.7ms/step | 11.5 samples/s
+Step 10400/44000 | Loss: 3.1315 | LR: 4.41e-04 | 1504.5ms/step | 13.3 samples/s
+Step 10500/44000 | Loss: 3.1231 | LR: 4.40e-04 | 1504.4ms/step | 13.3 samples/s
+Validation at step 10500: Loss = 3.1249, Perplexity = 22.76
+Step 10600/44000 | Loss: 3.1186 | LR: 4.39e-04 | 1747.2ms/step | 11.4 samples/s
+Step 10700/44000 | Loss: 3.1241 | LR: 4.37e-04 | 1505.1ms/step | 13.3 samples/s
+Validation at step 10750: Loss = 3.1193, Perplexity = 22.63
+Step 10800/44000 | Loss: 3.1179 | LR: 4.36e-04 | 1750.4ms/step | 11.4 samples/s
+Step 10900/44000 | Loss: 3.1129 | LR: 4.35e-04 | 1508.0ms/step | 13.3 samples/s
+Step 11000/44000 | Loss: 3.1182 | LR: 4.34e-04 | 1508.2ms/step | 13.3 samples/s
+Validation at step 11000: Loss = 3.1141, Perplexity = 22.51
+Step 11100/44000 | Loss: 3.1054 | LR: 4.33e-04 | 1751.5ms/step | 11.4 samples/s
+Step 11200/44000 | Loss: 3.1061 | LR: 4.31e-04 | 1508.0ms/step | 13.3 samples/s
+Validation at step 11250: Loss = 3.1095, Perplexity = 22.41
+Step 11300/44000 | Loss: 3.1068 | LR: 4.30e-04 | 1822.9ms/step | 11.0 samples/s
+Step 11400/44000 | Loss: 3.1021 | LR: 4.29e-04 | 1668.9ms/step | 12.0 samples/s
+Step 11500/44000 | Loss: 3.1019 | LR: 4.28e-04 | 1507.9ms/step | 13.3 samples/s
+Validation at step 11500: Loss = 3.1043, Perplexity = 22.29
+Step 11600/44000 | Loss: 3.1109 | LR: 4.26e-04 | 1750.9ms/step | 11.4 samples/s
+Step 11700/44000 | Loss: 3.0964 | LR: 4.25e-04 | 1507.4ms/step | 13.3 samples/s
+Validation at step 11750: Loss = 3.0994, Perplexity = 22.18
+Step 11800/44000 | Loss: 3.1015 | LR: 4.24e-04 | 1754.5ms/step | 11.4 samples/s
+Step 11900/44000 | Loss: 3.0962 | LR: 4.22e-04 | 1510.1ms/step | 13.2 samples/s
+Step 12000/44000 | Loss: 3.0991 | LR: 4.21e-04 | 1509.3ms/step | 13.3 samples/s
+Validation at step 12000: Loss = 3.0954, Perplexity = 22.10
+Step 12100/44000 | Loss: 3.0912 | LR: 4.20e-04 | 1753.4ms/step | 11.4 samples/s
+Step 12200/44000 | Loss: 3.0914 | LR: 4.18e-04 | 1507.8ms/step | 13.3 samples/s
+Validation at step 12250: Loss = 3.0917, Perplexity = 22.01
+Step 12300/44000 | Loss: 3.0834 | LR: 4.17e-04 | 1752.0ms/step | 11.4 samples/s
+Step 12400/44000 | Loss: 3.0821 | LR: 4.16e-04 | 1506.9ms/step | 13.3 samples/s
+Step 12500/44000 | Loss: 3.0807 | LR: 4.14e-04 | 1507.4ms/step | 13.3 samples/s
+Validation at step 12500: Loss = 3.0869, Perplexity = 21.91
+Step 12600/44000 | Loss: 3.0742 | LR: 4.13e-04 | 1752.0ms/step | 11.4 samples/s
+Step 12700/44000 | Loss: 3.0904 | LR: 4.12e-04 | 1508.4ms/step | 13.3 samples/s
+Validation at step 12750: Loss = 3.0827, Perplexity = 21.82
+Step 12800/44000 | Loss: 3.0731 | LR: 4.10e-04 | 1752.0ms/step | 11.4 samples/s
+Step 12900/44000 | Loss: 3.0827 | LR: 4.09e-04 | 1509.2ms/step | 13.3 samples/s
+Step 13000/44000 | Loss: 3.0801 | LR: 4.07e-04 | 1507.7ms/step | 13.3 samples/s
+Validation at step 13000: Loss = 3.0792, Perplexity = 21.74
+Step 13100/44000 | Loss: 3.0617 | LR: 4.06e-04 | 1751.4ms/step | 11.4 samples/s
+Step 13200/44000 | Loss: 3.0763 | LR: 4.05e-04 | 1506.7ms/step | 13.3 samples/s
+Validation at step 13250: Loss = 3.0762, Perplexity = 21.68
+Step 13300/44000 | Loss: 3.0799 | LR: 4.03e-04 | 1751.7ms/step | 11.4 samples/s
+Step 13400/44000 | Loss: 3.0754 | LR: 4.02e-04 | 1508.0ms/step | 13.3 samples/s
+Step 13500/44000 | Loss: 3.0705 | LR: 4.00e-04 | 1507.0ms/step | 13.3 samples/s
+Validation at step 13500: Loss = 3.0717, Perplexity = 21.58
+Step 13600/44000 | Loss: 3.0666 | LR: 3.99e-04 | 1751.7ms/step | 11.4 samples/s
+Step 13700/44000 | Loss: 3.0644 | LR: 3.97e-04 | 1507.6ms/step | 13.3 samples/s
+Validation at step 13750: Loss = 3.0681, Perplexity = 21.50
+Step 13800/44000 | Loss: 3.0719 | LR: 3.96e-04 | 1752.0ms/step | 11.4 samples/s
+Step 13900/44000 | Loss: 3.0641 | LR: 3.94e-04 | 1508.3ms/step | 13.3 samples/s
+Step 14000/44000 | Loss: 3.0599 | LR: 3.93e-04 | 1509.2ms/step | 13.3 samples/s
+Validation at step 14000: Loss = 3.0631, Perplexity = 21.39
+Step 14100/44000 | Loss: 3.0586 | LR: 3.91e-04 | 1753.5ms/step | 11.4 samples/s

epfl_final/logs/ALIBI_20260921/seed44.pid

@@ -0,0 +1 @@
+4776

epfl_final/logs/ALIBI_20260921/seed45.log

@@ -0,0 +1,375 @@
+New run: Checkpoints will be saved to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s45/run_20260921_030452
+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_030452",
+  "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": 22000,
+  "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_030515-hmdx78bp
+Starting training...
+Switching attention mode for step 1: sparse_stochastic (configured sparse attention)
+Step 100/44000 | Loss: 10.6682 | LR: 6.67e-05 | 1936.7ms/step | 10.3 samples/s
+Step 200/44000 | Loss: 7.3451 | LR: 1.33e-04 | 1475.5ms/step | 13.6 samples/s
+Validation at step 250: Loss = 6.3177, Perplexity = 554.27
+Step 300/44000 | Loss: 6.3378 | LR: 2.00e-04 | 1871.0ms/step | 10.7 samples/s
+Step 400/44000 | Loss: 5.8509 | LR: 2.67e-04 | 1491.5ms/step | 13.4 samples/s
+Step 500/44000 | Loss: 5.4661 | LR: 3.33e-04 | 1494.3ms/step | 13.4 samples/s
+Validation at step 500: Loss = 5.2966, Perplexity = 199.67
+Step 600/44000 | Loss: 5.1352 | LR: 4.00e-04 | 1736.4ms/step | 11.5 samples/s
+Step 700/44000 | Loss: 4.8670 | LR: 4.67e-04 | 1492.9ms/step | 13.4 samples/s
+Validation at step 750: Loss = 4.6064, Perplexity = 100.12
+Step 800/44000 | Loss: 4.6074 | LR: 5.00e-04 | 1733.1ms/step | 11.5 samples/s
+Step 900/44000 | Loss: 4.3872 | LR: 5.00e-04 | 1492.7ms/step | 13.4 samples/s
+Step 1000/44000 | Loss: 4.2273 | LR: 5.00e-04 | 1492.1ms/step | 13.4 samples/s
+Validation at step 1000: Loss = 4.1695, Perplexity = 64.68
+Step 1100/44000 | Loss: 4.1123 | LR: 5.00e-04 | 1734.0ms/step | 11.5 samples/s
+Step 1200/44000 | Loss: 4.0200 | LR: 5.00e-04 | 1492.3ms/step | 13.4 samples/s
+Validation at step 1250: Loss = 3.9552, Perplexity = 52.20
+Step 1300/44000 | Loss: 3.9399 | LR: 5.00e-04 | 1732.9ms/step | 11.5 samples/s
+Step 1400/44000 | Loss: 3.8900 | LR: 5.00e-04 | 1492.7ms/step | 13.4 samples/s
+Step 1500/44000 | Loss: 3.8386 | LR: 5.00e-04 | 1491.6ms/step | 13.4 samples/s
+Validation at step 1500: Loss = 3.8155, Perplexity = 45.40
+Step 1600/44000 | Loss: 3.8002 | LR: 5.00e-04 | 1730.4ms/step | 11.6 samples/s
+Step 1700/44000 | Loss: 3.7587 | LR: 4.99e-04 | 1491.2ms/step | 13.4 samples/s
+Validation at step 1750: Loss = 3.7223, Perplexity = 41.36
+Step 1800/44000 | Loss: 3.7128 | LR: 4.99e-04 | 1732.6ms/step | 11.5 samples/s
+Step 1900/44000 | Loss: 3.6837 | LR: 4.99e-04 | 1491.7ms/step | 13.4 samples/s
+Step 2000/44000 | Loss: 3.6690 | LR: 4.99e-04 | 1492.2ms/step | 13.4 samples/s
+Validation at step 2000: Loss = 3.6523, Perplexity = 38.56
+Step 2100/44000 | Loss: 3.6367 | LR: 4.99e-04 | 1731.6ms/step | 11.6 samples/s
+Step 2200/44000 | Loss: 3.6079 | LR: 4.99e-04 | 1491.4ms/step | 13.4 samples/s
+Validation at step 2250: Loss = 3.5964, Perplexity = 36.47
+Step 2300/44000 | Loss: 3.5834 | LR: 4.98e-04 | 1732.5ms/step | 11.5 samples/s
+Step 2400/44000 | Loss: 3.5690 | LR: 4.98e-04 | 1491.8ms/step | 13.4 samples/s
+Step 2500/44000 | Loss: 3.5540 | LR: 4.98e-04 | 1490.9ms/step | 13.4 samples/s
+Validation at step 2500: Loss = 3.5487, Perplexity = 34.77
+Step 2600/44000 | Loss: 3.5427 | LR: 4.98e-04 | 1732.3ms/step | 11.5 samples/s
+Step 2700/44000 | Loss: 3.5218 | LR: 4.97e-04 | 1491.7ms/step | 13.4 samples/s
+Validation at step 2750: Loss = 3.5110, Perplexity = 33.48
+Step 2800/44000 | Loss: 3.5066 | LR: 4.97e-04 | 1733.5ms/step | 11.5 samples/s
+Step 2900/44000 | Loss: 3.4890 | LR: 4.97e-04 | 1492.4ms/step | 13.4 samples/s
+Step 3000/44000 | Loss: 3.4834 | LR: 4.97e-04 | 1490.6ms/step | 13.4 samples/s
+Validation at step 3000: Loss = 3.4757, Perplexity = 32.32
+Step 3100/44000 | Loss: 3.4572 | LR: 4.96e-04 | 1732.7ms/step | 11.5 samples/s
+Step 3200/44000 | Loss: 3.4575 | LR: 4.96e-04 | 1492.4ms/step | 13.4 samples/s
+Validation at step 3250: Loss = 3.4469, Perplexity = 31.40
+Step 3300/44000 | Loss: 3.4458 | LR: 4.96e-04 | 1732.3ms/step | 11.5 samples/s
+Step 3400/44000 | Loss: 3.4308 | LR: 4.95e-04 | 1491.7ms/step | 13.4 samples/s
+Step 3500/44000 | Loss: 3.4299 | LR: 4.95e-04 | 1491.9ms/step | 13.4 samples/s
+Validation at step 3500: Loss = 3.4207, Perplexity = 30.59
+Step 3600/44000 | Loss: 3.4231 | LR: 4.95e-04 | 1732.7ms/step | 11.5 samples/s
+Step 3700/44000 | Loss: 3.4001 | LR: 4.94e-04 | 1492.3ms/step | 13.4 samples/s
+Validation at step 3750: Loss = 3.3970, Perplexity = 29.88
+Step 3800/44000 | Loss: 3.3917 | LR: 4.94e-04 | 1732.3ms/step | 11.5 samples/s
+Step 3900/44000 | Loss: 3.3819 | LR: 4.93e-04 | 1492.3ms/step | 13.4 samples/s
+Step 4000/44000 | Loss: 3.3806 | LR: 4.93e-04 | 1491.8ms/step | 13.4 samples/s
+Validation at step 4000: Loss = 3.3763, Perplexity = 29.26
+Step 4100/44000 | Loss: 3.3595 | LR: 4.93e-04 | 1731.6ms/step | 11.6 samples/s
+Step 4200/44000 | Loss: 3.3745 | LR: 4.92e-04 | 1492.7ms/step | 13.4 samples/s
+Validation at step 4250: Loss = 3.3583, Perplexity = 28.74
+Step 4300/44000 | Loss: 3.3454 | LR: 4.92e-04 | 1733.4ms/step | 11.5 samples/s
+Step 4400/44000 | Loss: 3.3476 | LR: 4.91e-04 | 1492.1ms/step | 13.4 samples/s
+Step 4500/44000 | Loss: 3.3449 | LR: 4.91e-04 | 1492.9ms/step | 13.4 samples/s
+Validation at step 4500: Loss = 3.3398, Perplexity = 28.21
+Step 4600/44000 | Loss: 3.3269 | LR: 4.90e-04 | 1733.8ms/step | 11.5 samples/s
+Step 4700/44000 | Loss: 3.3223 | LR: 4.90e-04 | 1492.6ms/step | 13.4 samples/s
+Validation at step 4750: Loss = 3.3243, Perplexity = 27.78
+Step 4800/44000 | Loss: 3.3194 | LR: 4.89e-04 | 1733.0ms/step | 11.5 samples/s
+Step 4900/44000 | Loss: 3.3095 | LR: 4.89e-04 | 1493.0ms/step | 13.4 samples/s
+Step 5000/44000 | Loss: 3.3182 | LR: 4.88e-04 | 1492.7ms/step | 13.4 samples/s
+Validation at step 5000: Loss = 3.3067, Perplexity = 27.29
+Step 5100/44000 | Loss: 3.2957 | LR: 4.88e-04 | 1733.2ms/step | 11.5 samples/s
+Step 5200/44000 | Loss: 3.2880 | LR: 4.87e-04 | 1492.1ms/step | 13.4 samples/s
+Validation at step 5250: Loss = 3.2950, Perplexity = 26.98
+Step 5300/44000 | Loss: 3.2821 | LR: 4.86e-04 | 1731.9ms/step | 11.5 samples/s
+Step 5400/44000 | Loss: 3.2784 | LR: 4.86e-04 | 1492.3ms/step | 13.4 samples/s
+Step 5500/44000 | Loss: 3.2753 | LR: 4.85e-04 | 1491.9ms/step | 13.4 samples/s
+Validation at step 5500: Loss = 3.2808, Perplexity = 26.60
+Step 5600/44000 | Loss: 3.2699 | LR: 4.85e-04 | 1732.4ms/step | 11.5 samples/s
+Step 5700/44000 | Loss: 3.2648 | LR: 4.84e-04 | 1491.8ms/step | 13.4 samples/s
+Validation at step 5750: Loss = 3.2694, Perplexity = 26.30
+Step 5800/44000 | Loss: 3.2696 | LR: 4.83e-04 | 1732.4ms/step | 11.5 samples/s
+Step 5900/44000 | Loss: 3.2740 | LR: 4.83e-04 | 1492.1ms/step | 13.4 samples/s
+Step 6000/44000 | Loss: 3.2598 | LR: 4.82e-04 | 1491.6ms/step | 13.4 samples/s
+Validation at step 6000: Loss = 3.2590, Perplexity = 26.02
+Step 6100/44000 | Loss: 3.2458 | LR: 4.81e-04 | 1732.3ms/step | 11.5 samples/s
+Step 6200/44000 | Loss: 3.2470 | LR: 4.81e-04 | 1492.1ms/step | 13.4 samples/s
+Validation at step 6250: Loss = 3.2509, Perplexity = 25.81
+Step 6300/44000 | Loss: 3.2557 | LR: 4.80e-04 | 1732.7ms/step | 11.5 samples/s
+Step 6400/44000 | Loss: 3.2417 | LR: 4.79e-04 | 1492.3ms/step | 13.4 samples/s
+Step 6500/44000 | Loss: 3.2359 | LR: 4.79e-04 | 1492.8ms/step | 13.4 samples/s
+Validation at step 6500: Loss = 3.2372, Perplexity = 25.46
+Step 6600/44000 | Loss: 3.2394 | LR: 4.78e-04 | 1732.1ms/step | 11.5 samples/s
+Step 6700/44000 | Loss: 3.2256 | LR: 4.77e-04 | 1491.8ms/step | 13.4 samples/s
+Validation at step 6750: Loss = 3.2296, Perplexity = 25.27
+Step 6800/44000 | Loss: 3.2257 | LR: 4.76e-04 | 1732.4ms/step | 11.5 samples/s
+Step 6900/44000 | Loss: 3.2198 | LR: 4.75e-04 | 1493.2ms/step | 13.4 samples/s
+Step 7000/44000 | Loss: 3.2159 | LR: 4.75e-04 | 1491.5ms/step | 13.4 samples/s
+Validation at step 7000: Loss = 3.2196, Perplexity = 25.02
+Step 7100/44000 | Loss: 3.2219 | LR: 4.74e-04 | 1733.4ms/step | 11.5 samples/s
+Step 7200/44000 | Loss: 3.2042 | LR: 4.73e-04 | 1492.6ms/step | 13.4 samples/s
+Validation at step 7250: Loss = 3.2100, Perplexity = 24.78
+Step 7300/44000 | Loss: 3.2099 | LR: 4.72e-04 | 1733.6ms/step | 11.5 samples/s
+Step 7400/44000 | Loss: 3.1926 | LR: 4.71e-04 | 1492.5ms/step | 13.4 samples/s
+Step 7500/44000 | Loss: 3.2007 | LR: 4.71e-04 | 1492.2ms/step | 13.4 samples/s
+Validation at step 7500: Loss = 3.2024, Perplexity = 24.59
+Step 7600/44000 | Loss: 3.1956 | LR: 4.70e-04 | 1732.5ms/step | 11.5 samples/s
+Step 7700/44000 | Loss: 3.1992 | LR: 4.69e-04 | 1491.3ms/step | 13.4 samples/s
+Validation at step 7750: Loss = 3.1945, Perplexity = 24.40
+Step 7800/44000 | Loss: 3.1899 | LR: 4.68e-04 | 1732.5ms/step | 11.5 samples/s
+Step 7900/44000 | Loss: 3.1902 | LR: 4.67e-04 | 1492.8ms/step | 13.4 samples/s
+Step 8000/44000 | Loss: 3.1774 | LR: 4.66e-04 | 1492.6ms/step | 13.4 samples/s
+Validation at step 8000: Loss = 3.1875, Perplexity = 24.23
+Step 8100/44000 | Loss: 3.1845 | LR: 4.65e-04 | 1732.5ms/step | 11.5 samples/s
+Step 8200/44000 | Loss: 3.1800 | LR: 4.64e-04 | 1491.9ms/step | 13.4 samples/s
+Validation at step 8250: Loss = 3.1792, Perplexity = 24.03
+Step 8300/44000 | Loss: 3.1714 | LR: 4.63e-04 | 1733.1ms/step | 11.5 samples/s
+Step 8400/44000 | Loss: 3.1771 | LR: 4.62e-04 | 1492.2ms/step | 13.4 samples/s
+Step 8500/44000 | Loss: 3.1627 | LR: 4.61e-04 | 1491.7ms/step | 13.4 samples/s
+Validation at step 8500: Loss = 3.1737, Perplexity = 23.90
+Step 8600/44000 | Loss: 3.1696 | LR: 4.60e-04 | 1732.4ms/step | 11.5 samples/s
+Step 8700/44000 | Loss: 3.1686 | LR: 4.59e-04 | 1492.2ms/step | 13.4 samples/s
+Validation at step 8750: Loss = 3.1659, Perplexity = 23.71
+Step 8800/44000 | Loss: 3.1711 | LR: 4.58e-04 | 1732.8ms/step | 11.5 samples/s
+Step 8900/44000 | Loss: 3.1583 | LR: 4.57e-04 | 1492.8ms/step | 13.4 samples/s
+Step 9000/44000 | Loss: 3.1525 | LR: 4.56e-04 | 1491.5ms/step | 13.4 samples/s
+Validation at step 9000: Loss = 3.1614, Perplexity = 23.60
+Step 9100/44000 | Loss: 3.1567 | LR: 4.55e-04 | 1732.1ms/step | 11.5 samples/s
+Step 9200/44000 | Loss: 3.1525 | LR: 4.54e-04 | 1493.2ms/step | 13.4 samples/s
+Validation at step 9250: Loss = 3.1530, Perplexity = 23.41
+Step 9300/44000 | Loss: 3.1593 | LR: 4.53e-04 | 1732.6ms/step | 11.5 samples/s
+Step 9400/44000 | Loss: 3.1539 | LR: 4.52e-04 | 1492.3ms/step | 13.4 samples/s
+Step 9500/44000 | Loss: 3.1466 | LR: 4.51e-04 | 1492.8ms/step | 13.4 samples/s
+Validation at step 9500: Loss = 3.1479, Perplexity = 23.29
+Step 9600/44000 | Loss: 3.1470 | LR: 4.50e-04 | 1732.2ms/step | 11.5 samples/s
+Step 9700/44000 | Loss: 3.1416 | LR: 4.49e-04 | 1492.0ms/step | 13.4 samples/s
+Validation at step 9750: Loss = 3.1430, Perplexity = 23.17
+Step 9800/44000 | Loss: 3.1334 | LR: 4.48e-04 | 1731.8ms/step | 11.5 samples/s
+Step 9900/44000 | Loss: 3.1390 | LR: 4.47e-04 | 1491.9ms/step | 13.4 samples/s
+Step 10000/44000 | Loss: 3.1338 | LR: 4.46e-04 | 1492.8ms/step | 13.4 samples/s
+Validation at step 10000: Loss = 3.1362, Perplexity = 23.02
+Step 10100/44000 | Loss: 3.1313 | LR: 4.45e-04 | 1733.9ms/step | 11.5 samples/s
+Step 10200/44000 | Loss: 3.1417 | LR: 4.43e-04 | 1492.6ms/step | 13.4 samples/s
+Validation at step 10250: Loss = 3.1312, Perplexity = 22.90
+Step 10300/44000 | Loss: 3.1271 | LR: 4.42e-04 | 1732.3ms/step | 11.5 samples/s
+Step 10400/44000 | Loss: 3.1360 | LR: 4.41e-04 | 1490.9ms/step | 13.4 samples/s
+Step 10500/44000 | Loss: 3.1300 | LR: 4.40e-04 | 1492.0ms/step | 13.4 samples/s
+Validation at step 10500: Loss = 3.1248, Perplexity = 22.76
+Step 10600/44000 | Loss: 3.1200 | LR: 4.39e-04 | 1731.8ms/step | 11.5 samples/s
+Step 10700/44000 | Loss: 3.1121 | LR: 4.37e-04 | 1491.6ms/step | 13.4 samples/s
+Validation at step 10750: Loss = 3.1218, Perplexity = 22.69
+Step 10800/44000 | Loss: 3.1065 | LR: 4.36e-04 | 1732.5ms/step | 11.5 samples/s
+Step 10900/44000 | Loss: 3.1226 | LR: 4.35e-04 | 1492.6ms/step | 13.4 samples/s
+Step 11000/44000 | Loss: 3.1158 | LR: 4.34e-04 | 1493.3ms/step | 13.4 samples/s
+Validation at step 11000: Loss = 3.1185, Perplexity = 22.61
+Step 11100/44000 | Loss: 3.1164 | LR: 4.33e-04 | 1735.5ms/step | 11.5 samples/s
+Step 11200/44000 | Loss: 3.1070 | LR: 4.31e-04 | 1494.1ms/step | 13.4 samples/s
+Validation at step 11250: Loss = 3.1110, Perplexity = 22.44
+Step 11300/44000 | Loss: 3.1110 | LR: 4.30e-04 | 1735.1ms/step | 11.5 samples/s
+Step 11400/44000 | Loss: 3.1081 | LR: 4.29e-04 | 1491.6ms/step | 13.4 samples/s
+Step 11500/44000 | Loss: 3.1066 | LR: 4.28e-04 | 1492.8ms/step | 13.4 samples/s
+Validation at step 11500: Loss = 3.1064, Perplexity = 22.34
+Step 11600/44000 | Loss: 3.1027 | LR: 4.26e-04 | 1734.0ms/step | 11.5 samples/s
+Step 11700/44000 | Loss: 3.1099 | LR: 4.25e-04 | 1494.0ms/step | 13.4 samples/s
+Validation at step 11750: Loss = 3.1023, Perplexity = 22.25
+Step 11800/44000 | Loss: 3.0959 | LR: 4.24e-04 | 1735.2ms/step | 11.5 samples/s
+Step 11900/44000 | Loss: 3.0957 | LR: 4.22e-04 | 1494.7ms/step | 13.4 samples/s
+Step 12000/44000 | Loss: 3.0950 | LR: 4.21e-04 | 1494.3ms/step | 13.4 samples/s
+Validation at step 12000: Loss = 3.0970, Perplexity = 22.13
+Step 12100/44000 | Loss: 3.0878 | LR: 4.20e-04 | 1730.1ms/step | 11.6 samples/s
+Step 12200/44000 | Loss: 3.0883 | LR: 4.18e-04 | 1490.8ms/step | 13.4 samples/s
+Validation at step 12250: Loss = 3.0935, Perplexity = 22.05
+Step 12300/44000 | Loss: 3.0874 | LR: 4.17e-04 | 1732.0ms/step | 11.5 samples/s
+Step 12400/44000 | Loss: 3.0855 | LR: 4.16e-04 | 1493.9ms/step | 13.4 samples/s
+Step 12500/44000 | Loss: 3.0854 | LR: 4.14e-04 | 1494.6ms/step | 13.4 samples/s
+Validation at step 12500: Loss = 3.0921, Perplexity = 22.02
+Step 12600/44000 | Loss: 3.0870 | LR: 4.13e-04 | 1734.9ms/step | 11.5 samples/s
+Step 12700/44000 | Loss: 3.0766 | LR: 4.12e-04 | 1494.4ms/step | 13.4 samples/s
+Validation at step 12750: Loss = 3.0857, Perplexity = 21.88
+Step 12800/44000 | Loss: 3.0871 | LR: 4.10e-04 | 1735.0ms/step | 11.5 samples/s
+Step 12900/44000 | Loss: 3.0843 | LR: 4.09e-04 | 1493.8ms/step | 13.4 samples/s
+Step 13000/44000 | Loss: 3.0849 | LR: 4.07e-04 | 1494.5ms/step | 13.4 samples/s
+Validation at step 13000: Loss = 3.0816, Perplexity = 21.79
+Step 13100/44000 | Loss: 3.0787 | LR: 4.06e-04 | 1735.4ms/step | 11.5 samples/s
+Step 13200/44000 | Loss: 3.0748 | LR: 4.05e-04 | 1495.3ms/step | 13.4 samples/s
+Validation at step 13250: Loss = 3.0775, Perplexity = 21.71
+Step 13300/44000 | Loss: 3.0641 | LR: 4.03e-04 | 1735.0ms/step | 11.5 samples/s
+Step 13400/44000 | Loss: 3.0712 | LR: 4.02e-04 | 1494.2ms/step | 13.4 samples/s
+Step 13500/44000 | Loss: 3.0609 | LR: 4.00e-04 | 1494.8ms/step | 13.4 samples/s
+Validation at step 13500: Loss = 3.0733, Perplexity = 21.61
+Step 13600/44000 | Loss: 3.0643 | LR: 3.99e-04 | 1735.5ms/step | 11.5 samples/s
+Step 13700/44000 | Loss: 3.0719 | LR: 3.97e-04 | 1494.2ms/step | 13.4 samples/s
+Validation at step 13750: Loss = 3.0695, Perplexity = 21.53
+Step 13800/44000 | Loss: 3.0647 | LR: 3.96e-04 | 1734.8ms/step | 11.5 samples/s
+Step 13900/44000 | Loss: 3.0637 | LR: 3.94e-04 | 1494.6ms/step | 13.4 samples/s
+Step 14000/44000 | Loss: 3.0712 | LR: 3.93e-04 | 1494.0ms/step | 13.4 samples/s
+Validation at step 14000: Loss = 3.0662, Perplexity = 21.46
+Step 14100/44000 | Loss: 3.0585 | LR: 3.91e-04 | 1736.4ms/step | 11.5 samples/s
+Step 14200/44000 | Loss: 3.0503 | LR: 3.90e-04 | 1487.2ms/step | 13.4 samples/s
+Validation at step 14250: Loss = 3.0622, Perplexity = 21.37
+Step 14300/44000 | Loss: 3.0604 | LR: 3.88e-04 | 1727.6ms/step | 11.6 samples/s

epfl_final/logs/ALIBI_20260921/seed45.pid

@@ -0,0 +1 @@
+4777

epfl_final/logs/ALIBI_20260921/seed46.log

@@ -0,0 +1,372 @@
+New run: Checkpoints will be saved to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s46/run_20260921_030452
+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_030452",
+  "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": 22000,
+  "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_030515-1savcbqu
+Starting training...
+Switching attention mode for step 1: sparse_stochastic (configured sparse attention)
+Step 100/44000 | Loss: 10.6054 | LR: 6.67e-05 | 1953.8ms/step | 10.2 samples/s
+Step 200/44000 | Loss: 7.3366 | LR: 1.33e-04 | 1501.5ms/step | 13.3 samples/s
+Validation at step 250: Loss = 6.3209, Perplexity = 556.09
+Step 300/44000 | Loss: 6.3333 | LR: 2.00e-04 | 1893.5ms/step | 10.6 samples/s
+Step 400/44000 | Loss: 5.8512 | LR: 2.67e-04 | 1511.2ms/step | 13.2 samples/s
+Step 500/44000 | Loss: 5.4640 | LR: 3.33e-04 | 1513.4ms/step | 13.2 samples/s
+Validation at step 500: Loss = 5.2887, Perplexity = 198.09
+Step 600/44000 | Loss: 5.1317 | LR: 4.00e-04 | 1755.8ms/step | 11.4 samples/s
+Step 700/44000 | Loss: 4.8770 | LR: 4.67e-04 | 1513.9ms/step | 13.2 samples/s
+Validation at step 750: Loss = 4.6088, Perplexity = 100.36
+Step 800/44000 | Loss: 4.6148 | LR: 5.00e-04 | 1755.0ms/step | 11.4 samples/s
+Step 900/44000 | Loss: 4.3939 | LR: 5.00e-04 | 1511.6ms/step | 13.2 samples/s
+Step 1000/44000 | Loss: 4.2161 | LR: 5.00e-04 | 1510.4ms/step | 13.2 samples/s
+Validation at step 1000: Loss = 4.1646, Perplexity = 64.37
+Step 1100/44000 | Loss: 4.1135 | LR: 5.00e-04 | 1757.0ms/step | 11.4 samples/s
+Step 1200/44000 | Loss: 4.0255 | LR: 5.00e-04 | 1512.2ms/step | 13.2 samples/s
+Validation at step 1250: Loss = 3.9475, Perplexity = 51.81
+Step 1300/44000 | Loss: 3.9594 | LR: 5.00e-04 | 1756.7ms/step | 11.4 samples/s
+Step 1400/44000 | Loss: 3.8776 | LR: 5.00e-04 | 1512.5ms/step | 13.2 samples/s
+Step 1500/44000 | Loss: 3.8304 | LR: 5.00e-04 | 1509.5ms/step | 13.2 samples/s
+Validation at step 1500: Loss = 3.8112, Perplexity = 45.20
+Step 1600/44000 | Loss: 3.7871 | LR: 5.00e-04 | 1750.4ms/step | 11.4 samples/s
+Step 1700/44000 | Loss: 3.7392 | LR: 4.99e-04 | 1511.4ms/step | 13.2 samples/s
+Validation at step 1750: Loss = 3.7160, Perplexity = 41.10
+Step 1800/44000 | Loss: 3.7067 | LR: 4.99e-04 | 1756.2ms/step | 11.4 samples/s
+Step 1900/44000 | Loss: 3.6916 | LR: 4.99e-04 | 1510.3ms/step | 13.2 samples/s
+Step 2000/44000 | Loss: 3.6616 | LR: 4.99e-04 | 1511.5ms/step | 13.2 samples/s
+Validation at step 2000: Loss = 3.6466, Perplexity = 38.35
+Step 2100/44000 | Loss: 3.6359 | LR: 4.99e-04 | 1752.9ms/step | 11.4 samples/s
+Step 2200/44000 | Loss: 3.6159 | LR: 4.99e-04 | 1510.2ms/step | 13.2 samples/s
+Validation at step 2250: Loss = 3.5907, Perplexity = 36.26
+Step 2300/44000 | Loss: 3.5912 | LR: 4.98e-04 | 1755.5ms/step | 11.4 samples/s
+Step 2400/44000 | Loss: 3.5657 | LR: 4.98e-04 | 1510.4ms/step | 13.2 samples/s
+Step 2500/44000 | Loss: 3.5506 | LR: 4.98e-04 | 1510.1ms/step | 13.2 samples/s
+Validation at step 2500: Loss = 3.5438, Perplexity = 34.60
+Step 2600/44000 | Loss: 3.5387 | LR: 4.98e-04 | 1754.1ms/step | 11.4 samples/s
+Step 2700/44000 | Loss: 3.5123 | LR: 4.97e-04 | 1510.4ms/step | 13.2 samples/s
+Validation at step 2750: Loss = 3.5051, Perplexity = 33.28
+Step 2800/44000 | Loss: 3.5094 | LR: 4.97e-04 | 1756.3ms/step | 11.4 samples/s
+Step 2900/44000 | Loss: 3.4957 | LR: 4.97e-04 | 1511.2ms/step | 13.2 samples/s
+Step 3000/44000 | Loss: 3.4858 | LR: 4.97e-04 | 1509.7ms/step | 13.2 samples/s
+Validation at step 3000: Loss = 3.4684, Perplexity = 32.08
+Step 3100/44000 | Loss: 3.4605 | LR: 4.96e-04 | 1755.9ms/step | 11.4 samples/s
+Step 3200/44000 | Loss: 3.4553 | LR: 4.96e-04 | 1510.0ms/step | 13.2 samples/s
+Validation at step 3250: Loss = 3.4403, Perplexity = 31.20
+Step 3300/44000 | Loss: 3.4380 | LR: 4.96e-04 | 1755.6ms/step | 11.4 samples/s
+Step 3400/44000 | Loss: 3.4337 | LR: 4.95e-04 | 1511.9ms/step | 13.2 samples/s
+Step 3500/44000 | Loss: 3.4122 | LR: 4.95e-04 | 1511.1ms/step | 13.2 samples/s
+Validation at step 3500: Loss = 3.4169, Perplexity = 30.47
+Step 3600/44000 | Loss: 3.4008 | LR: 4.95e-04 | 1755.2ms/step | 11.4 samples/s
+Step 3700/44000 | Loss: 3.4068 | LR: 4.94e-04 | 1512.4ms/step | 13.2 samples/s
+Validation at step 3750: Loss = 3.3925, Perplexity = 29.74
+Step 3800/44000 | Loss: 3.4028 | LR: 4.94e-04 | 1754.1ms/step | 11.4 samples/s
+Step 3900/44000 | Loss: 3.3822 | LR: 4.93e-04 | 1511.0ms/step | 13.2 samples/s
+Step 4000/44000 | Loss: 3.3726 | LR: 4.93e-04 | 1509.1ms/step | 13.3 samples/s
+Validation at step 4000: Loss = 3.3691, Perplexity = 29.05
+Step 4100/44000 | Loss: 3.3756 | LR: 4.93e-04 | 1754.6ms/step | 11.4 samples/s
+Step 4200/44000 | Loss: 3.3534 | LR: 4.92e-04 | 1511.9ms/step | 13.2 samples/s
+Validation at step 4250: Loss = 3.3522, Perplexity = 28.57
+Step 4300/44000 | Loss: 3.3485 | LR: 4.92e-04 | 1757.0ms/step | 11.4 samples/s
+Step 4400/44000 | Loss: 3.3418 | LR: 4.91e-04 | 1512.3ms/step | 13.2 samples/s
+Step 4500/44000 | Loss: 3.3371 | LR: 4.91e-04 | 1512.5ms/step | 13.2 samples/s
+Validation at step 4500: Loss = 3.3351, Perplexity = 28.08
+Step 4600/44000 | Loss: 3.3235 | LR: 4.90e-04 | 1756.1ms/step | 11.4 samples/s
+Step 4700/44000 | Loss: 3.3284 | LR: 4.90e-04 | 1511.3ms/step | 13.2 samples/s
+Validation at step 4750: Loss = 3.3225, Perplexity = 27.73
+Step 4800/44000 | Loss: 3.3170 | LR: 4.89e-04 | 1757.5ms/step | 11.4 samples/s
+Step 4900/44000 | Loss: 3.3104 | LR: 4.89e-04 | 1511.9ms/step | 13.2 samples/s
+Step 5000/44000 | Loss: 3.3102 | LR: 4.88e-04 | 1512.1ms/step | 13.2 samples/s
+Validation at step 5000: Loss = 3.3046, Perplexity = 27.24
+Step 5100/44000 | Loss: 3.3055 | LR: 4.88e-04 | 1755.2ms/step | 11.4 samples/s
+Step 5200/44000 | Loss: 3.2973 | LR: 4.87e-04 | 1506.3ms/step | 13.3 samples/s
+Validation at step 5250: Loss = 3.2893, Perplexity = 26.82
+Step 5300/44000 | Loss: 3.2833 | LR: 4.86e-04 | 1756.8ms/step | 11.4 samples/s
+Step 5400/44000 | Loss: 3.2768 | LR: 4.86e-04 | 1510.8ms/step | 13.2 samples/s
+Step 5500/44000 | Loss: 3.2783 | LR: 4.85e-04 | 1512.2ms/step | 13.2 samples/s
+Validation at step 5500: Loss = 3.2798, Perplexity = 26.57
+Step 5600/44000 | Loss: 3.2814 | LR: 4.85e-04 | 1754.4ms/step | 11.4 samples/s
+Step 5700/44000 | Loss: 3.2602 | LR: 4.84e-04 | 1512.5ms/step | 13.2 samples/s
+Validation at step 5750: Loss = 3.2658, Perplexity = 26.20
+Step 5800/44000 | Loss: 3.2650 | LR: 4.83e-04 | 1755.6ms/step | 11.4 samples/s
+Step 5900/44000 | Loss: 3.2670 | LR: 4.83e-04 | 1510.1ms/step | 13.2 samples/s
+Step 6000/44000 | Loss: 3.2540 | LR: 4.82e-04 | 1511.6ms/step | 13.2 samples/s
+Validation at step 6000: Loss = 3.2550, Perplexity = 25.92
+Step 6100/44000 | Loss: 3.2587 | LR: 4.81e-04 | 1754.0ms/step | 11.4 samples/s
+Step 6200/44000 | Loss: 3.2394 | LR: 4.81e-04 | 1512.1ms/step | 13.2 samples/s
+Validation at step 6250: Loss = 3.2436, Perplexity = 25.63
+Step 6300/44000 | Loss: 3.2488 | LR: 4.80e-04 | 1754.2ms/step | 11.4 samples/s
+Step 6400/44000 | Loss: 3.2437 | LR: 4.79e-04 | 1512.1ms/step | 13.2 samples/s
+Step 6500/44000 | Loss: 3.2300 | LR: 4.79e-04 | 1509.8ms/step | 13.2 samples/s
+Validation at step 6500: Loss = 3.2343, Perplexity = 25.39
+Step 6600/44000 | Loss: 3.2325 | LR: 4.78e-04 | 1754.1ms/step | 11.4 samples/s
+Step 6700/44000 | Loss: 3.2328 | LR: 4.77e-04 | 1510.8ms/step | 13.2 samples/s
+Validation at step 6750: Loss = 3.2243, Perplexity = 25.14
+Step 6800/44000 | Loss: 3.2285 | LR: 4.76e-04 | 1755.7ms/step | 11.4 samples/s
+Step 6900/44000 | Loss: 3.2157 | LR: 4.75e-04 | 1510.0ms/step | 13.2 samples/s
+Step 7000/44000 | Loss: 3.2209 | LR: 4.75e-04 | 1512.1ms/step | 13.2 samples/s
+Validation at step 7000: Loss = 3.2149, Perplexity = 24.90
+Step 7100/44000 | Loss: 3.2077 | LR: 4.74e-04 | 1756.6ms/step | 11.4 samples/s
+Step 7200/44000 | Loss: 3.2027 | LR: 4.73e-04 | 1512.8ms/step | 13.2 samples/s
+Validation at step 7250: Loss = 3.2056, Perplexity = 24.67
+Step 7300/44000 | Loss: 3.2042 | LR: 4.72e-04 | 1755.2ms/step | 11.4 samples/s
+Step 7400/44000 | Loss: 3.1941 | LR: 4.71e-04 | 1511.6ms/step | 13.2 samples/s
+Step 7500/44000 | Loss: 3.1937 | LR: 4.71e-04 | 1512.6ms/step | 13.2 samples/s
+Validation at step 7500: Loss = 3.1983, Perplexity = 24.49
+Step 7600/44000 | Loss: 3.1950 | LR: 4.70e-04 | 1754.6ms/step | 11.4 samples/s
+Step 7700/44000 | Loss: 3.1852 | LR: 4.69e-04 | 1511.7ms/step | 13.2 samples/s
+Validation at step 7750: Loss = 3.1898, Perplexity = 24.28
+Step 7800/44000 | Loss: 3.1997 | LR: 4.68e-04 | 1755.8ms/step | 11.4 samples/s
+Step 7900/44000 | Loss: 3.1784 | LR: 4.67e-04 | 1512.2ms/step | 13.2 samples/s
+Step 8000/44000 | Loss: 3.1860 | LR: 4.66e-04 | 1511.7ms/step | 13.2 samples/s
+Validation at step 8000: Loss = 3.1814, Perplexity = 24.08
+Step 8100/44000 | Loss: 3.1797 | LR: 4.65e-04 | 1757.1ms/step | 11.4 samples/s
+Step 8200/44000 | Loss: 3.1862 | LR: 4.64e-04 | 1513.0ms/step | 13.2 samples/s
+Validation at step 8250: Loss = 3.1744, Perplexity = 23.91
+Step 8300/44000 | Loss: 3.1648 | LR: 4.63e-04 | 1756.6ms/step | 11.4 samples/s
+Step 8400/44000 | Loss: 3.1789 | LR: 4.62e-04 | 1509.1ms/step | 13.3 samples/s
+Step 8500/44000 | Loss: 3.1745 | LR: 4.61e-04 | 1511.2ms/step | 13.2 samples/s
+Validation at step 8500: Loss = 3.1686, Perplexity = 23.78
+Step 8600/44000 | Loss: 3.1650 | LR: 4.60e-04 | 1754.2ms/step | 11.4 samples/s
+Step 8700/44000 | Loss: 3.1645 | LR: 4.59e-04 | 1510.4ms/step | 13.2 samples/s
+Validation at step 8750: Loss = 3.1609, Perplexity = 23.59
+Step 8800/44000 | Loss: 3.1662 | LR: 4.58e-04 | 1756.2ms/step | 11.4 samples/s
+Step 8900/44000 | Loss: 3.1568 | LR: 4.57e-04 | 1511.7ms/step | 13.2 samples/s
+Step 9000/44000 | Loss: 3.1494 | LR: 4.56e-04 | 1511.3ms/step | 13.2 samples/s
+Validation at step 9000: Loss = 3.1545, Perplexity = 23.44
+Step 9100/44000 | Loss: 3.1530 | LR: 4.55e-04 | 1758.0ms/step | 11.4 samples/s
+Step 9200/44000 | Loss: 3.1558 | LR: 4.54e-04 | 1510.8ms/step | 13.2 samples/s
+Validation at step 9250: Loss = 3.1502, Perplexity = 23.34
+Step 9300/44000 | Loss: 3.1427 | LR: 4.53e-04 | 1755.2ms/step | 11.4 samples/s
+Step 9400/44000 | Loss: 3.1401 | LR: 4.52e-04 | 1510.2ms/step | 13.2 samples/s
+Step 9500/44000 | Loss: 3.1418 | LR: 4.51e-04 | 1509.8ms/step | 13.2 samples/s
+Validation at step 9500: Loss = 3.1439, Perplexity = 23.19
+Step 9600/44000 | Loss: 3.1385 | LR: 4.50e-04 | 1754.7ms/step | 11.4 samples/s
+Step 9700/44000 | Loss: 3.1397 | LR: 4.49e-04 | 1511.9ms/step | 13.2 samples/s
+Validation at step 9750: Loss = 3.1372, Perplexity = 23.04
+Step 9800/44000 | Loss: 3.1443 | LR: 4.48e-04 | 1755.7ms/step | 11.4 samples/s
+Step 9900/44000 | Loss: 3.1379 | LR: 4.47e-04 | 1512.8ms/step | 13.2 samples/s
+Step 10000/44000 | Loss: 3.1274 | LR: 4.46e-04 | 1513.6ms/step | 13.2 samples/s
+Validation at step 10000: Loss = 3.1313, Perplexity = 22.90
+Step 10100/44000 | Loss: 3.1261 | LR: 4.45e-04 | 1756.2ms/step | 11.4 samples/s
+Step 10200/44000 | Loss: 3.1226 | LR: 4.43e-04 | 1508.4ms/step | 13.3 samples/s
+Validation at step 10250: Loss = 3.1270, Perplexity = 22.81
+Step 10300/44000 | Loss: 3.1250 | LR: 4.42e-04 | 1752.7ms/step | 11.4 samples/s
+Step 10400/44000 | Loss: 3.1120 | LR: 4.41e-04 | 1509.7ms/step | 13.2 samples/s
+Step 10500/44000 | Loss: 3.1280 | LR: 4.40e-04 | 1510.0ms/step | 13.2 samples/s
+Validation at step 10500: Loss = 3.1204, Perplexity = 22.66
+Step 10600/44000 | Loss: 3.1213 | LR: 4.39e-04 | 1754.2ms/step | 11.4 samples/s
+Step 10700/44000 | Loss: 3.1114 | LR: 4.37e-04 | 1510.3ms/step | 13.2 samples/s
+Validation at step 10750: Loss = 3.1159, Perplexity = 22.55
+Step 10800/44000 | Loss: 3.1070 | LR: 4.36e-04 | 1755.5ms/step | 11.4 samples/s
+Step 10900/44000 | Loss: 3.1112 | LR: 4.35e-04 | 1513.1ms/step | 13.2 samples/s
+Step 11000/44000 | Loss: 3.1097 | LR: 4.34e-04 | 1512.4ms/step | 13.2 samples/s
+Validation at step 11000: Loss = 3.1116, Perplexity = 22.46
+Step 11100/44000 | Loss: 3.1115 | LR: 4.33e-04 | 1757.1ms/step | 11.4 samples/s
+Step 11200/44000 | Loss: 3.1028 | LR: 4.31e-04 | 1578.0ms/step | 12.7 samples/s
+Validation at step 11250: Loss = 3.1091, Perplexity = 22.40
+Step 11300/44000 | Loss: 3.1011 | LR: 4.30e-04 | 1748.0ms/step | 11.4 samples/s
+Step 11400/44000 | Loss: 3.1023 | LR: 4.29e-04 | 1509.3ms/step | 13.3 samples/s
+Step 11500/44000 | Loss: 3.1075 | LR: 4.28e-04 | 1513.3ms/step | 13.2 samples/s
+Validation at step 11500: Loss = 3.1022, Perplexity = 22.25
+Step 11600/44000 | Loss: 3.1060 | LR: 4.26e-04 | 1756.3ms/step | 11.4 samples/s
+Step 11700/44000 | Loss: 3.1019 | LR: 4.25e-04 | 1513.0ms/step | 13.2 samples/s
+Validation at step 11750: Loss = 3.0969, Perplexity = 22.13
+Step 11800/44000 | Loss: 3.0986 | LR: 4.24e-04 | 1758.1ms/step | 11.4 samples/s
+Step 11900/44000 | Loss: 3.0857 | LR: 4.22e-04 | 1515.0ms/step | 13.2 samples/s
+Step 12000/44000 | Loss: 3.0893 | LR: 4.21e-04 | 1513.7ms/step | 13.2 samples/s
+Validation at step 12000: Loss = 3.0927, Perplexity = 22.04
+Step 12100/44000 | Loss: 3.1023 | LR: 4.20e-04 | 1758.4ms/step | 11.4 samples/s
+Step 12200/44000 | Loss: 3.0822 | LR: 4.18e-04 | 1513.8ms/step | 13.2 samples/s
+Validation at step 12250: Loss = 3.0875, Perplexity = 21.92
+Step 12300/44000 | Loss: 3.0835 | LR: 4.17e-04 | 1756.7ms/step | 11.4 samples/s
+Step 12400/44000 | Loss: 3.0780 | LR: 4.16e-04 | 1512.3ms/step | 13.2 samples/s
+Step 12500/44000 | Loss: 3.0921 | LR: 4.14e-04 | 1513.4ms/step | 13.2 samples/s
+Validation at step 12500: Loss = 3.0837, Perplexity = 21.84
+Step 12600/44000 | Loss: 3.0888 | LR: 4.13e-04 | 1759.3ms/step | 11.4 samples/s
+Step 12700/44000 | Loss: 3.0814 | LR: 4.12e-04 | 1513.9ms/step | 13.2 samples/s
+Validation at step 12750: Loss = 3.0793, Perplexity = 21.74
+Step 12800/44000 | Loss: 3.0797 | LR: 4.10e-04 | 1759.8ms/step | 11.4 samples/s
+Step 12900/44000 | Loss: 3.0827 | LR: 4.09e-04 | 1515.1ms/step | 13.2 samples/s
+Step 13000/44000 | Loss: 3.0774 | LR: 4.07e-04 | 1514.1ms/step | 13.2 samples/s
+Validation at step 13000: Loss = 3.0747, Perplexity = 21.64
+Step 13100/44000 | Loss: 3.0727 | LR: 4.06e-04 | 1756.0ms/step | 11.4 samples/s
+Step 13200/44000 | Loss: 3.0786 | LR: 4.05e-04 | 1512.6ms/step | 13.2 samples/s
+Validation at step 13250: Loss = 3.0723, Perplexity = 21.59
+Step 13300/44000 | Loss: 3.0790 | LR: 4.03e-04 | 1757.7ms/step | 11.4 samples/s
+Step 13400/44000 | Loss: 3.0751 | LR: 4.02e-04 | 1513.6ms/step | 13.2 samples/s
+Step 13500/44000 | Loss: 3.0670 | LR: 4.00e-04 | 1512.9ms/step | 13.2 samples/s
+Validation at step 13500: Loss = 3.0691, Perplexity = 21.52
+Step 13600/44000 | Loss: 3.0535 | LR: 3.99e-04 | 1757.3ms/step | 11.4 samples/s
+Step 13700/44000 | Loss: 3.0690 | LR: 3.97e-04 | 1512.4ms/step | 13.2 samples/s
+Validation at step 13750: Loss = 3.0657, Perplexity = 21.45
+Step 13800/44000 | Loss: 3.0728 | LR: 3.96e-04 | 1759.1ms/step | 11.4 samples/s
+Step 13900/44000 | Loss: 3.0622 | LR: 3.94e-04 | 1513.3ms/step | 13.2 samples/s
+Step 14000/44000 | Loss: 3.0596 | LR: 3.93e-04 | 1513.3ms/step | 13.2 samples/s
+Validation at step 14000: Loss = 3.0597, Perplexity = 21.32
+Step 14100/44000 | Loss: 3.0605 | LR: 3.91e-04 | 1760.3ms/step | 11.4 samples/s

epfl_final/logs/ALIBI_20260921/seed46.pid

@@ -0,0 +1 @@
+4778

epfl_final/logs/ALIBI_20260921/seed47.log

@@ -0,0 +1,375 @@
+New run: Checkpoints will be saved to /mloscratch/homes/navasard/sparse/models/ALIBI_20260921/365m_alibi_s47/run_20260921_030452
+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_030452",
+  "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": 22000,
+  "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_030515-8km2sc23
+Starting training...
+Switching attention mode for step 1: sparse_stochastic (configured sparse attention)
+Step 100/44000 | Loss: 10.6289 | LR: 6.67e-05 | 1934.2ms/step | 10.3 samples/s
+Step 200/44000 | Loss: 7.3290 | LR: 1.33e-04 | 1471.9ms/step | 13.6 samples/s
+Validation at step 250: Loss = 6.3193, Perplexity = 555.20
+Step 300/44000 | Loss: 6.3461 | LR: 2.00e-04 | 1866.4ms/step | 10.7 samples/s
+Step 400/44000 | Loss: 5.8661 | LR: 2.67e-04 | 1482.2ms/step | 13.5 samples/s
+Step 500/44000 | Loss: 5.4685 | LR: 3.33e-04 | 1480.6ms/step | 13.5 samples/s
+Validation at step 500: Loss = 5.2935, Perplexity = 199.05
+Step 600/44000 | Loss: 5.1430 | LR: 4.00e-04 | 1728.0ms/step | 11.6 samples/s
+Step 700/44000 | Loss: 4.8626 | LR: 4.67e-04 | 1488.7ms/step | 13.4 samples/s
+Validation at step 750: Loss = 4.6031, Perplexity = 99.80
+Step 800/44000 | Loss: 4.5941 | LR: 5.00e-04 | 1727.9ms/step | 11.6 samples/s
+Step 900/44000 | Loss: 4.3813 | LR: 5.00e-04 | 1488.2ms/step | 13.4 samples/s
+Step 1000/44000 | Loss: 4.2221 | LR: 5.00e-04 | 1487.8ms/step | 13.4 samples/s
+Validation at step 1000: Loss = 4.1545, Perplexity = 63.72
+Step 1100/44000 | Loss: 4.0983 | LR: 5.00e-04 | 1731.7ms/step | 11.5 samples/s
+Step 1200/44000 | Loss: 4.0244 | LR: 5.00e-04 | 1488.9ms/step | 13.4 samples/s
+Validation at step 1250: Loss = 3.9445, Perplexity = 51.65
+Step 1300/44000 | Loss: 3.9451 | LR: 5.00e-04 | 1729.0ms/step | 11.6 samples/s
+Step 1400/44000 | Loss: 3.8860 | LR: 5.00e-04 | 1489.7ms/step | 13.4 samples/s
+Step 1500/44000 | Loss: 3.8279 | LR: 5.00e-04 | 1488.0ms/step | 13.4 samples/s
+Validation at step 1500: Loss = 3.8131, Perplexity = 45.29
+Step 1600/44000 | Loss: 3.8018 | LR: 5.00e-04 | 1726.6ms/step | 11.6 samples/s
+Step 1700/44000 | Loss: 3.7492 | LR: 4.99e-04 | 1488.3ms/step | 13.4 samples/s
+Validation at step 1750: Loss = 3.7166, Perplexity = 41.13
+Step 1800/44000 | Loss: 3.7205 | LR: 4.99e-04 | 1728.5ms/step | 11.6 samples/s
+Step 1900/44000 | Loss: 3.6866 | LR: 4.99e-04 | 1487.7ms/step | 13.4 samples/s
+Step 2000/44000 | Loss: 3.6607 | LR: 4.99e-04 | 1488.4ms/step | 13.4 samples/s
+Validation at step 2000: Loss = 3.6467, Perplexity = 38.35
+Step 2100/44000 | Loss: 3.6309 | LR: 4.99e-04 | 1727.7ms/step | 11.6 samples/s
+Step 2200/44000 | Loss: 3.6095 | LR: 4.99e-04 | 1487.3ms/step | 13.4 samples/s
+Validation at step 2250: Loss = 3.5903, Perplexity = 36.25
+Step 2300/44000 | Loss: 3.5873 | LR: 4.98e-04 | 1727.5ms/step | 11.6 samples/s
+Step 2400/44000 | Loss: 3.5730 | LR: 4.98e-04 | 1487.5ms/step | 13.4 samples/s
+Step 2500/44000 | Loss: 3.5536 | LR: 4.98e-04 | 1486.7ms/step | 13.5 samples/s
+Validation at step 2500: Loss = 3.5475, Perplexity = 34.73
+Step 2600/44000 | Loss: 3.5317 | LR: 4.98e-04 | 1727.6ms/step | 11.6 samples/s
+Step 2700/44000 | Loss: 3.5320 | LR: 4.97e-04 | 1487.9ms/step | 13.4 samples/s
+Validation at step 2750: Loss = 3.5097, Perplexity = 33.44
+Step 2800/44000 | Loss: 3.5089 | LR: 4.97e-04 | 1728.2ms/step | 11.6 samples/s
+Step 2900/44000 | Loss: 3.4845 | LR: 4.97e-04 | 1489.2ms/step | 13.4 samples/s
+Step 3000/44000 | Loss: 3.4731 | LR: 4.97e-04 | 1487.1ms/step | 13.4 samples/s
+Validation at step 3000: Loss = 3.4710, Perplexity = 32.17
+Step 3100/44000 | Loss: 3.4657 | LR: 4.96e-04 | 1728.6ms/step | 11.6 samples/s
+Step 3200/44000 | Loss: 3.4491 | LR: 4.96e-04 | 1488.4ms/step | 13.4 samples/s
+Validation at step 3250: Loss = 3.4424, Perplexity = 31.26
+Step 3300/44000 | Loss: 3.4490 | LR: 4.96e-04 | 1723.4ms/step | 11.6 samples/s
+Step 3400/44000 | Loss: 3.4294 | LR: 4.95e-04 | 1480.7ms/step | 13.5 samples/s
+Step 3500/44000 | Loss: 3.4234 | LR: 4.95e-04 | 1480.9ms/step | 13.5 samples/s
+Validation at step 3500: Loss = 3.4172, Perplexity = 30.48
+Step 3600/44000 | Loss: 3.4089 | LR: 4.95e-04 | 1719.1ms/step | 11.6 samples/s
+Step 3700/44000 | Loss: 3.4091 | LR: 4.94e-04 | 1480.8ms/step | 13.5 samples/s
+Validation at step 3750: Loss = 3.3927, Perplexity = 29.75
+Step 3800/44000 | Loss: 3.3873 | LR: 4.94e-04 | 1719.0ms/step | 11.6 samples/s
+Step 3900/44000 | Loss: 3.3831 | LR: 4.93e-04 | 1480.4ms/step | 13.5 samples/s
+Step 4000/44000 | Loss: 3.3658 | LR: 4.93e-04 | 1480.3ms/step | 13.5 samples/s
+Validation at step 4000: Loss = 3.3714, Perplexity = 29.12
+Step 4100/44000 | Loss: 3.3614 | LR: 4.93e-04 | 1718.9ms/step | 11.6 samples/s
+Step 4200/44000 | Loss: 3.3675 | LR: 4.92e-04 | 1480.4ms/step | 13.5 samples/s
+Validation at step 4250: Loss = 3.3523, Perplexity = 28.57
+Step 4300/44000 | Loss: 3.3592 | LR: 4.92e-04 | 1718.8ms/step | 11.6 samples/s
+Step 4400/44000 | Loss: 3.3519 | LR: 4.91e-04 | 1480.1ms/step | 13.5 samples/s
+Step 4500/44000 | Loss: 3.3339 | LR: 4.91e-04 | 1480.8ms/step | 13.5 samples/s
+Validation at step 4500: Loss = 3.3344, Perplexity = 28.06
+Step 4600/44000 | Loss: 3.3308 | LR: 4.90e-04 | 1728.8ms/step | 11.6 samples/s
+Step 4700/44000 | Loss: 3.3192 | LR: 4.90e-04 | 1488.3ms/step | 13.4 samples/s
+Validation at step 4750: Loss = 3.3186, Perplexity = 27.62
+Step 4800/44000 | Loss: 3.3074 | LR: 4.89e-04 | 1731.4ms/step | 11.6 samples/s
+Step 4900/44000 | Loss: 3.3113 | LR: 4.89e-04 | 1497.2ms/step | 13.4 samples/s
+Step 5000/44000 | Loss: 3.2987 | LR: 4.88e-04 | 1495.4ms/step | 13.4 samples/s
+Validation at step 5000: Loss = 3.3034, Perplexity = 27.21
+Step 5100/44000 | Loss: 3.3005 | LR: 4.88e-04 | 1737.3ms/step | 11.5 samples/s
+Step 5200/44000 | Loss: 3.2957 | LR: 4.87e-04 | 1494.7ms/step | 13.4 samples/s
+Validation at step 5250: Loss = 3.2896, Perplexity = 26.83
+Step 5300/44000 | Loss: 3.2950 | LR: 4.86e-04 | 1726.8ms/step | 11.6 samples/s
+Step 5400/44000 | Loss: 3.2756 | LR: 4.86e-04 | 1488.5ms/step | 13.4 samples/s
+Step 5500/44000 | Loss: 3.2747 | LR: 4.85e-04 | 1488.0ms/step | 13.4 samples/s
+Validation at step 5500: Loss = 3.2784, Perplexity = 26.53
+Step 5600/44000 | Loss: 3.2689 | LR: 4.85e-04 | 1727.1ms/step | 11.6 samples/s
+Step 5700/44000 | Loss: 3.2657 | LR: 4.84e-04 | 1487.4ms/step | 13.4 samples/s
+Validation at step 5750: Loss = 3.2651, Perplexity = 26.18
+Step 5800/44000 | Loss: 3.2595 | LR: 4.83e-04 | 1728.0ms/step | 11.6 samples/s
+Step 5900/44000 | Loss: 3.2643 | LR: 4.83e-04 | 1488.4ms/step | 13.4 samples/s
+Step 6000/44000 | Loss: 3.2495 | LR: 4.82e-04 | 1487.9ms/step | 13.4 samples/s
+Validation at step 6000: Loss = 3.2533, Perplexity = 25.88
+Step 6100/44000 | Loss: 3.2597 | LR: 4.81e-04 | 1727.8ms/step | 11.6 samples/s
+Step 6200/44000 | Loss: 3.2413 | LR: 4.81e-04 | 1487.9ms/step | 13.4 samples/s
+Validation at step 6250: Loss = 3.2424, Perplexity = 25.59
+Step 6300/44000 | Loss: 3.2479 | LR: 4.80e-04 | 1728.0ms/step | 11.6 samples/s
+Step 6400/44000 | Loss: 3.2334 | LR: 4.79e-04 | 1487.4ms/step | 13.4 samples/s
+Step 6500/44000 | Loss: 3.2360 | LR: 4.79e-04 | 1488.1ms/step | 13.4 samples/s
+Validation at step 6500: Loss = 3.2331, Perplexity = 25.36
+Step 6600/44000 | Loss: 3.2277 | LR: 4.78e-04 | 1727.7ms/step | 11.6 samples/s
+Step 6700/44000 | Loss: 3.2211 | LR: 4.77e-04 | 1488.1ms/step | 13.4 samples/s
+Validation at step 6750: Loss = 3.2235, Perplexity = 25.12
+Step 6800/44000 | Loss: 3.2188 | LR: 4.76e-04 | 1727.8ms/step | 11.6 samples/s
+Step 6900/44000 | Loss: 3.2095 | LR: 4.75e-04 | 1488.9ms/step | 13.4 samples/s
+Step 7000/44000 | Loss: 3.2118 | LR: 4.75e-04 | 1487.5ms/step | 13.4 samples/s
+Validation at step 7000: Loss = 3.2149, Perplexity = 24.90
+Step 7100/44000 | Loss: 3.2093 | LR: 4.74e-04 | 1728.1ms/step | 11.6 samples/s
+Step 7200/44000 | Loss: 3.2133 | LR: 4.73e-04 | 1488.3ms/step | 13.4 samples/s
+Validation at step 7250: Loss = 3.2052, Perplexity = 24.66
+Step 7300/44000 | Loss: 3.1981 | LR: 4.72e-04 | 1728.3ms/step | 11.6 samples/s
+Step 7400/44000 | Loss: 3.1948 | LR: 4.71e-04 | 1488.4ms/step | 13.4 samples/s
+Step 7500/44000 | Loss: 3.1945 | LR: 4.71e-04 | 1488.1ms/step | 13.4 samples/s
+Validation at step 7500: Loss = 3.1969, Perplexity = 24.46
+Step 7600/44000 | Loss: 3.1925 | LR: 4.70e-04 | 1727.6ms/step | 11.6 samples/s
+Step 7700/44000 | Loss: 3.2028 | LR: 4.69e-04 | 1488.5ms/step | 13.4 samples/s
+Validation at step 7750: Loss = 3.1907, Perplexity = 24.31
+Step 7800/44000 | Loss: 3.1919 | LR: 4.68e-04 | 1727.9ms/step | 11.6 samples/s
+Step 7900/44000 | Loss: 3.1907 | LR: 4.67e-04 | 1487.6ms/step | 13.4 samples/s
+Step 8000/44000 | Loss: 3.1840 | LR: 4.66e-04 | 1488.4ms/step | 13.4 samples/s
+Validation at step 8000: Loss = 3.1825, Perplexity = 24.11
+Step 8100/44000 | Loss: 3.1791 | LR: 4.65e-04 | 1727.9ms/step | 11.6 samples/s
+Step 8200/44000 | Loss: 3.1672 | LR: 4.64e-04 | 1489.0ms/step | 13.4 samples/s
+Validation at step 8250: Loss = 3.1748, Perplexity = 23.92
+Step 8300/44000 | Loss: 3.1711 | LR: 4.63e-04 | 1728.4ms/step | 11.6 samples/s
+Step 8400/44000 | Loss: 3.1686 | LR: 4.62e-04 | 1488.0ms/step | 13.4 samples/s
+Step 8500/44000 | Loss: 3.1747 | LR: 4.61e-04 | 1488.2ms/step | 13.4 samples/s
+Validation at step 8500: Loss = 3.1681, Perplexity = 23.76
+Step 8600/44000 | Loss: 3.1668 | LR: 4.60e-04 | 1727.3ms/step | 11.6 samples/s
+Step 8700/44000 | Loss: 3.1631 | LR: 4.59e-04 | 1487.9ms/step | 13.4 samples/s
+Validation at step 8750: Loss = 3.1616, Perplexity = 23.61
+Step 8800/44000 | Loss: 3.1462 | LR: 4.58e-04 | 1727.7ms/step | 11.6 samples/s
+Step 8900/44000 | Loss: 3.1657 | LR: 4.57e-04 | 1488.3ms/step | 13.4 samples/s
+Step 9000/44000 | Loss: 3.1542 | LR: 4.56e-04 | 1487.7ms/step | 13.4 samples/s
+Validation at step 9000: Loss = 3.1545, Perplexity = 23.44
+Step 9100/44000 | Loss: 3.1472 | LR: 4.55e-04 | 1727.4ms/step | 11.6 samples/s
+Step 9200/44000 | Loss: 3.1596 | LR: 4.54e-04 | 1488.8ms/step | 13.4 samples/s
+Validation at step 9250: Loss = 3.1495, Perplexity = 23.32
+Step 9300/44000 | Loss: 3.1543 | LR: 4.53e-04 | 1727.7ms/step | 11.6 samples/s
+Step 9400/44000 | Loss: 3.1423 | LR: 4.52e-04 | 1487.0ms/step | 13.5 samples/s
+Step 9500/44000 | Loss: 3.1473 | LR: 4.51e-04 | 1488.5ms/step | 13.4 samples/s
+Validation at step 9500: Loss = 3.1428, Perplexity = 23.17
+Step 9600/44000 | Loss: 3.1352 | LR: 4.50e-04 | 1727.1ms/step | 11.6 samples/s
+Step 9700/44000 | Loss: 3.1349 | LR: 4.49e-04 | 1487.6ms/step | 13.4 samples/s
+Validation at step 9750: Loss = 3.1371, Perplexity = 23.04
+Step 9800/44000 | Loss: 3.1275 | LR: 4.48e-04 | 1727.7ms/step | 11.6 samples/s
+Step 9900/44000 | Loss: 3.1296 | LR: 4.47e-04 | 1486.4ms/step | 13.5 samples/s
+Step 10000/44000 | Loss: 3.1308 | LR: 4.46e-04 | 1488.0ms/step | 13.4 samples/s
+Validation at step 10000: Loss = 3.1326, Perplexity = 22.93
+Step 10100/44000 | Loss: 3.1270 | LR: 4.45e-04 | 1727.6ms/step | 11.6 samples/s
+Step 10200/44000 | Loss: 3.1351 | LR: 4.43e-04 | 1487.9ms/step | 13.4 samples/s
+Validation at step 10250: Loss = 3.1267, Perplexity = 22.80
+Step 10300/44000 | Loss: 3.1321 | LR: 4.42e-04 | 1728.2ms/step | 11.6 samples/s
+Step 10400/44000 | Loss: 3.1200 | LR: 4.41e-04 | 1486.9ms/step | 13.5 samples/s
+Step 10500/44000 | Loss: 3.1247 | LR: 4.40e-04 | 1487.5ms/step | 13.4 samples/s
+Validation at step 10500: Loss = 3.1204, Perplexity = 22.65
+Step 10600/44000 | Loss: 3.1239 | LR: 4.39e-04 | 1726.8ms/step | 11.6 samples/s
+Step 10700/44000 | Loss: 3.1202 | LR: 4.37e-04 | 1488.4ms/step | 13.4 samples/s
+Validation at step 10750: Loss = 3.1169, Perplexity = 22.58
+Step 10800/44000 | Loss: 3.1140 | LR: 4.36e-04 | 1727.4ms/step | 11.6 samples/s
+Step 10900/44000 | Loss: 3.1101 | LR: 4.35e-04 | 1487.5ms/step | 13.4 samples/s
+Step 11000/44000 | Loss: 3.1156 | LR: 4.34e-04 | 1488.9ms/step | 13.4 samples/s
+Validation at step 11000: Loss = 3.1118, Perplexity = 22.46
+Step 11100/44000 | Loss: 3.1133 | LR: 4.33e-04 | 1729.5ms/step | 11.6 samples/s
+Step 11200/44000 | Loss: 3.0973 | LR: 4.31e-04 | 1489.9ms/step | 13.4 samples/s
+Validation at step 11250: Loss = 3.1065, Perplexity = 22.34
+Step 11300/44000 | Loss: 3.1079 | LR: 4.30e-04 | 1730.9ms/step | 11.6 samples/s
+Step 11400/44000 | Loss: 3.1055 | LR: 4.29e-04 | 1553.3ms/step | 12.9 samples/s
+Step 11500/44000 | Loss: 3.1044 | LR: 4.28e-04 | 1481.4ms/step | 13.5 samples/s
+Validation at step 11500: Loss = 3.1022, Perplexity = 22.25
+Step 11600/44000 | Loss: 3.0918 | LR: 4.26e-04 | 1719.6ms/step | 11.6 samples/s
+Step 11700/44000 | Loss: 3.1004 | LR: 4.25e-04 | 1482.7ms/step | 13.5 samples/s
+Validation at step 11750: Loss = 3.0974, Perplexity = 22.14
+Step 11800/44000 | Loss: 3.0951 | LR: 4.24e-04 | 1720.9ms/step | 11.6 samples/s
+Step 11900/44000 | Loss: 3.0961 | LR: 4.22e-04 | 1481.1ms/step | 13.5 samples/s
+Step 12000/44000 | Loss: 3.0925 | LR: 4.21e-04 | 1481.3ms/step | 13.5 samples/s
+Validation at step 12000: Loss = 3.0923, Perplexity = 22.03
+Step 12100/44000 | Loss: 3.0883 | LR: 4.20e-04 | 1721.9ms/step | 11.6 samples/s
+Step 12200/44000 | Loss: 3.0970 | LR: 4.18e-04 | 1481.8ms/step | 13.5 samples/s
+Validation at step 12250: Loss = 3.0895, Perplexity = 21.97
+Step 12300/44000 | Loss: 3.0874 | LR: 4.17e-04 | 1721.6ms/step | 11.6 samples/s
+Step 12400/44000 | Loss: 3.0851 | LR: 4.16e-04 | 1481.1ms/step | 13.5 samples/s
+Step 12500/44000 | Loss: 3.0817 | LR: 4.14e-04 | 1481.0ms/step | 13.5 samples/s
+Validation at step 12500: Loss = 3.0853, Perplexity = 21.87
+Step 12600/44000 | Loss: 3.0790 | LR: 4.13e-04 | 1720.6ms/step | 11.6 samples/s
+Step 12700/44000 | Loss: 3.0757 | LR: 4.12e-04 | 1480.5ms/step | 13.5 samples/s
+Validation at step 12750: Loss = 3.0805, Perplexity = 21.77
+Step 12800/44000 | Loss: 3.0765 | LR: 4.10e-04 | 1722.1ms/step | 11.6 samples/s
+Step 12900/44000 | Loss: 3.0793 | LR: 4.09e-04 | 1482.7ms/step | 13.5 samples/s
+Step 13000/44000 | Loss: 3.0834 | LR: 4.07e-04 | 1482.7ms/step | 13.5 samples/s
+Validation at step 13000: Loss = 3.0756, Perplexity = 21.66
+Step 13100/44000 | Loss: 3.0649 | LR: 4.06e-04 | 1721.5ms/step | 11.6 samples/s
+Step 13200/44000 | Loss: 3.0809 | LR: 4.05e-04 | 1481.7ms/step | 13.5 samples/s
+Validation at step 13250: Loss = 3.0711, Perplexity = 21.57
+Step 13300/44000 | Loss: 3.0772 | LR: 4.03e-04 | 1720.8ms/step | 11.6 samples/s
+Step 13400/44000 | Loss: 3.0721 | LR: 4.02e-04 | 1481.9ms/step | 13.5 samples/s
+Step 13500/44000 | Loss: 3.0807 | LR: 4.00e-04 | 1481.2ms/step | 13.5 samples/s
+Validation at step 13500: Loss = 3.0678, Perplexity = 21.49
+Step 13600/44000 | Loss: 3.0726 | LR: 3.99e-04 | 1721.2ms/step | 11.6 samples/s
+Step 13700/44000 | Loss: 3.0658 | LR: 3.97e-04 | 1481.2ms/step | 13.5 samples/s
+Validation at step 13750: Loss = 3.0639, Perplexity = 21.41
+Step 13800/44000 | Loss: 3.0615 | LR: 3.96e-04 | 1721.3ms/step | 11.6 samples/s
+Step 13900/44000 | Loss: 3.0590 | LR: 3.94e-04 | 1481.4ms/step | 13.5 samples/s
+Step 14000/44000 | Loss: 3.0578 | LR: 3.93e-04 | 1482.4ms/step | 13.5 samples/s
+Validation at step 14000: Loss = 3.0608, Perplexity = 21.34
+Step 14100/44000 | Loss: 3.0468 | LR: 3.91e-04 | 1724.1ms/step | 11.6 samples/s
+Step 14200/44000 | Loss: 3.0597 | LR: 3.90e-04 | 1486.4ms/step | 13.5 samples/s
+Validation at step 14250: Loss = 3.0565, Perplexity = 21.25
+Step 14300/44000 | Loss: 3.0588 | LR: 3.88e-04 | 1726.9ms/step | 11.6 samples/s

epfl_final/logs/ALIBI_20260921/seed47.pid

@@ -0,0 +1 @@
+4779

epfl_final/logs/build_4b_cache.log

@@ -0,0 +1,12 @@
+=== 4b cache build start Wed Jul 15 11:00:53 PM UTC 2026 ===
+mask_cache_32 layer0 pool=10000
+hardlinked 0-27 -> 28 files
+[gpu 0] layer 28: built+saved pool=10000 in 564s
+[gpu 1] layer 29: built+saved pool=10000 in 564s
+[gpu 2] layer 30: built+saved pool=10000 in 568s
+[gpu 3] layer 31: built+saved pool=10000 in 573s
+[gpu 1] layer 33: built+saved pool=10000 in 556s
+[gpu 0] layer 32: built+saved pool=10000 in 558s
+[gpu 2] layer 34: built+saved pool=10000 in 558s
+[gpu 3] layer 35: built+saved pool=10000 in 561s
+=== 4b cache build done Wed Jul 15 11:19:55 PM UTC 2026 rc=0 files=36 ===

epfl_final/logs/build_masks_4b.log

@@ -0,0 +1,50 @@
+[build_masks] Config:
+  sparse layers : [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35]
+  seq_len       : 30720
+  block_size    : 128
+  num_blocks    : 240
+  num_random    : 120
+  mask_pool_size: 1000
+  num_workers   : 192
+  cache_dir     : /mloscratch/homes/navasard/sparse/mask_cache_30k_b128_nrb120
+  force_rebuild : False
+
+[build_masks] Layer 0: cache exists, skipping.
+[build_masks] Layer 1: cache exists, skipping.
+[build_masks] Layer 2: cache exists, skipping.
+[build_masks] Layer 3: cache exists, skipping.
+[build_masks] Layer 4: cache exists, skipping.
+[build_masks] Layer 5: cache exists, skipping.
+[build_masks] Layer 6: cache exists, skipping.
+[build_masks] Layer 7: cache exists, skipping.
+[build_masks] Layer 8: cache exists, skipping.
+[build_masks] Layer 9: cache exists, skipping.
+[build_masks] Layer 10: cache exists, skipping.
+[build_masks] Layer 11: cache exists, skipping.
+[build_masks] Layer 12: cache exists, skipping.
+[build_masks] Layer 13: cache exists, skipping.
+[build_masks] Layer 14: cache exists, skipping.
+[build_masks] Layer 15: cache exists, skipping.
+[build_masks] Layer 16: cache exists, skipping.
+[build_masks] Layer 17: cache exists, skipping.
+[build_masks] Layer 18: cache exists, skipping.
+[build_masks] Layer 19: cache exists, skipping.
+[build_masks] Layer 20: cache exists, skipping.
+[build_masks] Layer 21: cache exists, skipping.
+[build_masks] Layer 22: cache exists, skipping.
+[build_masks] Layer 23: cache exists, skipping.
+[build_masks] Layer 24: cache exists, skipping.
+[build_masks] Layer 25: cache exists, skipping.
+[build_masks] Layer 26: cache exists, skipping.
+[build_masks] Layer 27: cache exists, skipping.
+[build_masks] Building 1000 masks × 8 layers on 4 GPU(s). Est. 70s/layer on GPU.
+[build_masks] GPU 1: layer 29 done.
+[build_masks] GPU 3: layer 31 done.
+[build_masks] GPU 2: layer 30 done.
+[build_masks] GPU 0: layer 28 done.
+[build_masks] GPU 1: layer 33 done.
+[build_masks] GPU 3: layer 35 done.
+[build_masks] GPU 2: layer 34 done.
+[build_masks] GPU 0: layer 32 done.
+[build_masks] Done. Total wall-time: 1477.3s
+[verify] PASS — layer 0: 1000 masks, rt_pool shape (1000, 240, 120)

epfl_final/logs/build_mixed_cache.log

@@ -0,0 +1,30 @@
+=== build start Wed Jul 15 10:29:43 PM UTC 2026 ===
+[gpu 0] layer 0: built+saved pool=10000 in 606s
+[gpu 3] layer 3: built+saved pool=10000 in 608s
+[gpu 1] layer 1: built+saved pool=10000 in 610s
+[gpu 2] layer 2: built+saved pool=10000 in 611s
+[gpu 0] layer 4: built+saved pool=10000 in 602s
+[gpu 1] layer 5: built+saved pool=10000 in 606s
+[gpu 2] layer 6: built+saved pool=10000 in 606s
+[gpu 3] layer 7: built+saved pool=10000 in 609s
+[gpu 0] layer 8: built+saved pool=10000 in 599s
+[gpu 1] layer 9: built+saved pool=10000 in 609s
+[gpu 3] layer 11: built+saved pool=10000 in 608s
+[gpu 2] layer 10: built+saved pool=10000 in 610s
+[gpu 0] layer 12: built+saved pool=10000 in 604s
+[gpu 3] layer 15: built+saved pool=10000 in 601s
+[gpu 1] layer 13: built+saved pool=10000 in 609s
+[gpu 2] layer 14: built+saved pool=10000 in 607s
+[gpu 0] layer 16: built+saved pool=10000 in 603s
+[gpu 3] layer 19: built+saved pool=10000 in 604s
+[gpu 2] layer 18: built+saved pool=10000 in 606s
+[gpu 1] layer 17: built+saved pool=10000 in 609s
+[gpu 0] layer 20: built+saved pool=10000 in 599s
+[gpu 3] layer 23: built+saved pool=10000 in 600s
+[gpu 2] layer 22: built+saved pool=10000 in 606s
+[gpu 1] layer 21: built+saved pool=10000 in 607s
+[gpu 0] layer 24: built+saved pool=10000 in 597s
+[gpu 3] layer 27: built+saved pool=10000 in 605s
+[gpu 2] layer 26: built+saved pool=10000 in 601s
+[gpu 1] layer 25: built+saved pool=10000 in 605s
+=== build done Wed Jul 15 11:40:39 PM UTC 2026 rc=0 layer_files=28 ===

epfl_final/logs/eval_1p2b_battery.log

@@ -0,0 +1,182 @@
+[05:51:19] total jobs=78 already-done=0 gpus=[0, 1, 2, 3] n=100
+[05:51:24] START gpu0 passkey_8k_dense try1 (~78 remaining)
+[05:51:24] START gpu1 passkey_8k_fixed50 try1 (~78 remaining)
+[06:16:25] DONE gpu0 passkey_8k_dense rc=0 ok
+[06:16:25] START gpu0 passkey_8k_resampled50 try1 (~77 remaining)
+[06:16:40] DONE gpu1 passkey_8k_fixed50 rc=0 ok
+[06:16:41] START gpu1 passkey_8k_resampled65 try1 (~76 remaining)
+[06:37:47] DONE gpu0 passkey_8k_resampled50 rc=0 ok
+[06:37:47] DONE gpu1 passkey_8k_resampled65 rc=0 ok
+[06:37:47] START gpu0 ppl_8k_dense_resampled_fineweb try1 (~74 remaining)
+[06:37:47] START gpu1 ppl_8k_dense_resampled_pg19 try1 (~74 remaining)
+[06:38:38] START gpu2 ppl_8k_dense_resampled_arxiv try1 (~74 remaining)
+[06:38:43] START gpu3 ppl_8k_dense_fixed_fineweb try1 (~74 remaining)
+[06:53:04] DONE gpu0 ppl_8k_dense_resampled_fineweb rc=0 ok
+[06:53:04] START gpu0 ppl_8k_dense_fixed_pg19 try1 (~73 remaining)
+[06:53:14] DONE gpu1 ppl_8k_dense_resampled_pg19 rc=0 ok
+[06:53:14] START gpu1 ppl_8k_dense_fixed_arxiv try1 (~72 remaining)
+[06:53:54] DONE gpu2 ppl_8k_dense_resampled_arxiv rc=0 ok
+[06:53:54] START gpu2 ppl_8k_fixed50_resampled_fineweb try1 (~71 remaining)
+[06:54:04] DONE gpu3 ppl_8k_dense_fixed_fineweb rc=0 ok
+[06:54:04] START gpu3 ppl_8k_fixed50_resampled_pg19 try1 (~70 remaining)
+[07:08:14] DONE gpu1 ppl_8k_dense_fixed_arxiv rc=0 ok
+[07:08:14] START gpu1 ppl_8k_fixed50_resampled_arxiv try1 (~69 remaining)
+[07:08:29] DONE gpu0 ppl_8k_dense_fixed_pg19 rc=0 ok
+[07:08:29] START gpu0 ppl_8k_fixed50_fixed_fineweb try1 (~68 remaining)
+[07:09:04] DONE gpu2 ppl_8k_fixed50_resampled_fineweb rc=0 ok
+[07:09:04] START gpu2 ppl_8k_fixed50_fixed_pg19 try1 (~67 remaining)
+[07:09:29] DONE gpu3 ppl_8k_fixed50_resampled_pg19 rc=0 ok
+[07:09:29] START gpu3 ppl_8k_fixed50_fixed_arxiv try1 (~66 remaining)
+[07:23:24] DONE gpu1 ppl_8k_fixed50_resampled_arxiv rc=0 ok
+[07:23:24] START gpu1 ppl_8k_resampled50_resampled_fineweb try1 (~65 remaining)
+[07:23:54] DONE gpu0 ppl_8k_fixed50_fixed_fineweb rc=0 ok
+[07:23:54] START gpu0 ppl_8k_resampled50_resampled_pg19 try1 (~64 remaining)
+[07:24:14] DONE gpu2 ppl_8k_fixed50_fixed_pg19 rc=0 ok
+[07:24:14] START gpu2 ppl_8k_resampled50_resampled_arxiv try1 (~63 remaining)
+[07:24:44] DONE gpu3 ppl_8k_fixed50_fixed_arxiv rc=0 ok
+[07:24:44] START gpu3 ppl_8k_resampled50_fixed_fineweb try1 (~62 remaining)
+[07:38:34] DONE gpu1 ppl_8k_resampled50_resampled_fineweb rc=0 ok
+[07:38:35] START gpu1 ppl_8k_resampled50_fixed_pg19 try1 (~61 remaining)
+[07:39:10] DONE gpu0 ppl_8k_resampled50_resampled_pg19 rc=0 ok
+[07:39:10] START gpu0 ppl_8k_resampled50_fixed_arxiv try1 (~60 remaining)
+[07:39:40] DONE gpu2 ppl_8k_resampled50_resampled_arxiv rc=0 ok
+[07:39:40] START gpu2 ppl_8k_resampled65_resampled_fineweb try1 (~59 remaining)
+[07:40:00] DONE gpu3 ppl_8k_resampled50_fixed_fineweb rc=0 ok
+[07:40:00] START gpu3 ppl_8k_resampled65_resampled_pg19 try1 (~58 remaining)
+[07:54:00] DONE gpu1 ppl_8k_resampled50_fixed_pg19 rc=0 ok
+[07:54:00] START gpu1 ppl_8k_resampled65_resampled_arxiv try1 (~57 remaining)
+[07:54:25] DONE gpu0 ppl_8k_resampled50_fixed_arxiv rc=0 ok
+[07:54:25] START gpu0 ppl_8k_resampled65_fixed_fineweb try1 (~56 remaining)
+[07:54:55] DONE gpu2 ppl_8k_resampled65_resampled_fineweb rc=0 ok
+[07:54:55] START gpu2 ppl_8k_resampled65_fixed_pg19 try1 (~55 remaining)
+[07:55:15] DONE gpu3 ppl_8k_resampled65_resampled_pg19 rc=0 ok
+[07:55:15] START gpu3 ppl_8k_resampled65_fixed_arxiv try1 (~54 remaining)
+[08:09:20] DONE gpu1 ppl_8k_resampled65_resampled_arxiv rc=0 ok
+[08:09:20] START gpu1 ruler_8k_niah_multikey_dense try1 (~53 remaining)
+[08:09:25] DONE gpu1 ruler_8k_niah_multikey_dense rc=2 NO-TSV
+[08:09:25] START gpu1 ruler_8k_niah_multikey_dense try2 (~52 remaining)
+[08:09:30] DONE gpu0 ppl_8k_resampled65_fixed_fineweb rc=0 ok
+[08:09:30] DONE gpu1 ruler_8k_niah_multikey_dense rc=2 NO-TSV
+[08:09:30] START gpu0 ruler_8k_niah_multikey_dense try3 (~51 remaining)
+[08:09:31] START gpu1 ruler_8k_niah_multikey_fixed50 try1 (~51 remaining)
+[08:09:36] DONE gpu0 ruler_8k_niah_multikey_dense rc=2 NO-TSV
+[08:09:36] START gpu0 ruler_8k_niah_multikey_resampled50 try1 (~51 remaining)
+[08:10:11] DONE gpu2 ppl_8k_resampled65_fixed_pg19 rc=0 ok
+[08:10:11] START gpu2 ruler_8k_niah_multikey_resampled65 try1 (~50 remaining)
+[08:10:41] DONE gpu3 ppl_8k_resampled65_fixed_arxiv rc=0 ok
+[08:10:41] START gpu3 ruler_8k_niah_multivalue_dense try1 (~49 remaining)
+[08:10:46] DONE gpu3 ruler_8k_niah_multivalue_dense rc=2 NO-TSV
+[08:10:46] START gpu3 ruler_8k_niah_multivalue_dense try2 (~49 remaining)
+[08:10:51] DONE gpu3 ruler_8k_niah_multivalue_dense rc=2 NO-TSV
+[08:10:51] START gpu3 ruler_8k_niah_multivalue_dense try3 (~48 remaining)
+[08:10:56] DONE gpu3 ruler_8k_niah_multivalue_dense rc=2 NO-TSV
+[08:10:56] START gpu3 ruler_8k_niah_multivalue_fixed50 try1 (~48 remaining)
+[08:30:26] DONE gpu1 ruler_8k_niah_multikey_fixed50 rc=0 ok
+[08:30:26] START gpu1 ruler_8k_niah_multivalue_resampled50 try1 (~47 remaining)
+[08:30:41] DONE gpu0 ruler_8k_niah_multikey_resampled50 rc=0 ok
+[08:30:41] START gpu0 ruler_8k_niah_multivalue_resampled65 try1 (~46 remaining)
+[08:31:16] DONE gpu2 ruler_8k_niah_multikey_resampled65 rc=0 ok
+[08:31:16] START gpu2 ruler_8k_vt_dense try1 (~45 remaining)
+[08:31:21] DONE gpu2 ruler_8k_vt_dense rc=2 NO-TSV
+[08:31:21] START gpu2 ruler_8k_vt_dense try2 (~45 remaining)
+[08:31:26] DONE gpu2 ruler_8k_vt_dense rc=2 NO-TSV
+[08:31:26] START gpu2 ruler_8k_vt_dense try3 (~44 remaining)
+[08:31:31] DONE gpu2 ruler_8k_vt_dense rc=2 NO-TSV
+[08:31:32] START gpu2 ruler_8k_vt_fixed50 try1 (~44 remaining)
+[08:31:42] DONE gpu3 ruler_8k_niah_multivalue_fixed50 rc=0 ok
+[08:31:42] START gpu3 ruler_8k_vt_resampled50 try1 (~43 remaining)
+[08:51:32] DONE gpu0 ruler_8k_niah_multivalue_resampled65 rc=0 ok
+[08:51:32] START gpu0 ruler_8k_vt_resampled65 try1 (~42 remaining)
+[08:51:42] DONE gpu1 ruler_8k_niah_multivalue_resampled50 rc=0 ok
+[08:51:42] START gpu1 kv_h2o_8k_dense try1 (~41 remaining)
+[08:52:22] DONE gpu2 ruler_8k_vt_fixed50 rc=0 ok
+[08:52:22] START gpu2 kv_stream_8k_dense try1 (~40 remaining)
+[08:52:42] DONE gpu3 ruler_8k_vt_resampled50 rc=0 ok
+[08:52:42] START gpu3 kv_h2o_8k_fixed50 try1 (~39 remaining)
+[08:54:12] DONE gpu2 kv_stream_8k_dense rc=0 ok
+[08:54:12] START gpu2 kv_stream_8k_fixed50 try1 (~38 remaining)
+[08:56:07] DONE gpu2 kv_stream_8k_fixed50 rc=0 ok
+[08:56:07] START gpu2 kv_h2o_8k_resampled50 try1 (~37 remaining)
+[08:56:17] DONE gpu1 kv_h2o_8k_dense rc=0 ok
+[08:56:17] START gpu1 kv_stream_8k_resampled50 try1 (~36 remaining)
+[08:57:17] DONE gpu3 kv_h2o_8k_fixed50 rc=0 ok
+[08:57:17] START gpu3 kv_h2o_8k_resampled65 try1 (~35 remaining)
+[08:58:12] DONE gpu1 kv_stream_8k_resampled50 rc=0 ok
+[08:58:12] START gpu1 kv_stream_8k_resampled65 try1 (~34 remaining)
+[09:00:12] DONE gpu1 kv_stream_8k_resampled65 rc=0 ok
+[09:00:12] START gpu1 passkey_30k_dense try1 (~33 remaining)
+[09:00:42] DONE gpu2 kv_h2o_8k_resampled50 rc=0 ok
+[09:00:42] START gpu2 passkey_30k_fixed try1 (~32 remaining)
+[09:01:52] DONE gpu3 kv_h2o_8k_resampled65 rc=0 ok
+[09:01:52] START gpu3 passkey_30k_resampled try1 (~31 remaining)
+[09:12:32] DONE gpu0 ruler_8k_vt_resampled65 rc=0 ok
+[09:12:33] START gpu0 ppl_30k_dense_resampled_pg19 try1 (~30 remaining)
+[11:32:53] DONE gpu0 ppl_30k_dense_resampled_pg19 rc=0 ok
+[11:32:53] START gpu0 ppl_30k_dense_resampled_arxiv try1 (~29 remaining)
+[11:50:03] DONE gpu2 passkey_30k_fixed rc=0 ok
+[11:50:03] START gpu2 ppl_30k_dense_fixed_pg19 try1 (~28 remaining)
+[11:50:13] DONE gpu1 passkey_30k_dense rc=0 ok
+[11:50:13] START gpu1 ppl_30k_dense_fixed_arxiv try1 (~27 remaining)
+[11:52:28] DONE gpu3 passkey_30k_resampled rc=0 ok
+[11:52:28] START gpu3 ppl_30k_fixed_resampled_pg19 try1 (~26 remaining)
+[13:54:23] DONE gpu0 ppl_30k_dense_resampled_arxiv rc=0 ok
+[13:54:23] START gpu0 ppl_30k_fixed_resampled_arxiv try1 (~25 remaining)
+[14:10:58] DONE gpu1 ppl_30k_dense_fixed_arxiv rc=0 ok
+[14:10:59] START gpu1 ppl_30k_fixed_fixed_pg19 try1 (~24 remaining)
+[14:11:54] DONE gpu2 ppl_30k_dense_fixed_pg19 rc=0 ok
+[14:11:54] START gpu2 ppl_30k_fixed_fixed_arxiv try1 (~23 remaining)
+[14:13:29] DONE gpu3 ppl_30k_fixed_resampled_pg19 rc=0 ok
+[14:13:29] START gpu3 ppl_30k_resampled_resampled_pg19 try1 (~22 remaining)
+[16:15:39] DONE gpu0 ppl_30k_fixed_resampled_arxiv rc=0 ok
+[16:15:39] START gpu0 ppl_30k_resampled_resampled_arxiv try1 (~21 remaining)
+[16:32:14] DONE gpu1 ppl_30k_fixed_fixed_pg19 rc=0 ok
+[16:32:14] START gpu1 ppl_30k_resampled_fixed_pg19 try1 (~20 remaining)
+[16:32:44] DONE gpu2 ppl_30k_fixed_fixed_arxiv rc=0 ok
+[16:32:44] START gpu2 ppl_30k_resampled_fixed_arxiv try1 (~19 remaining)
+[16:34:29] DONE gpu3 ppl_30k_resampled_resampled_pg19 rc=0 ok
+[16:34:30] START gpu3 ruler_30k_niah_multikey_dense try1 (~18 remaining)
+[16:34:35] DONE gpu3 ruler_30k_niah_multikey_dense rc=2 NO-TSV
+[16:34:35] START gpu3 ruler_30k_niah_multikey_dense try2 (~18 remaining)
+[16:34:40] DONE gpu3 ruler_30k_niah_multikey_dense rc=2 NO-TSV
+[16:34:40] START gpu3 ruler_30k_niah_multikey_dense try3 (~17 remaining)
+[16:34:45] DONE gpu3 ruler_30k_niah_multikey_dense rc=2 NO-TSV
+[16:34:45] START gpu3 ruler_30k_niah_multikey_fixed try1 (~17 remaining)
+[18:37:05] DONE gpu0 ppl_30k_resampled_resampled_arxiv rc=0 ok
+[18:37:05] START gpu0 ruler_30k_niah_multikey_resampled try1 (~16 remaining)
+[18:53:55] DONE gpu1 ppl_30k_resampled_fixed_pg19 rc=0 ok
+[18:53:55] START gpu1 ruler_30k_niah_multivalue_dense try1 (~15 remaining)
+[18:54:00] DONE gpu2 ppl_30k_resampled_fixed_arxiv rc=0 ok
+[18:54:00] DONE gpu1 ruler_30k_niah_multivalue_dense rc=2 NO-TSV
+[18:54:00] START gpu1 ruler_30k_niah_multivalue_dense try2 (~14 remaining)
+[18:54:00] START gpu2 ruler_30k_niah_multivalue_fixed try1 (~14 remaining)
+[18:54:05] DONE gpu1 ruler_30k_niah_multivalue_dense rc=2 NO-TSV
+[18:54:05] START gpu1 ruler_30k_niah_multivalue_dense try3 (~13 remaining)
+[18:54:10] DONE gpu1 ruler_30k_niah_multivalue_dense rc=2 NO-TSV
+[18:54:11] START gpu1 ruler_30k_niah_multivalue_resampled try1 (~13 remaining)
+[19:24:46] DONE gpu3 ruler_30k_niah_multikey_fixed rc=0 ok
+[19:24:46] START gpu3 ruler_30k_vt_dense try1 (~12 remaining)
+[19:24:51] DONE gpu3 ruler_30k_vt_dense rc=2 NO-TSV
+[19:24:51] START gpu3 ruler_30k_vt_dense try2 (~12 remaining)
+[19:24:56] DONE gpu3 ruler_30k_vt_dense rc=2 NO-TSV
+[19:24:56] START gpu3 ruler_30k_vt_dense try3 (~11 remaining)
+[19:25:01] DONE gpu3 ruler_30k_vt_dense rc=2 NO-TSV
+[19:25:01] START gpu3 ruler_30k_vt_fixed try1 (~11 remaining)
+[21:26:01] DONE gpu0 ruler_30k_niah_multikey_resampled rc=0 ok
+[21:26:01] START gpu0 ruler_30k_vt_resampled try1 (~10 remaining)
+[21:43:16] DONE gpu1 ruler_30k_niah_multivalue_resampled rc=0 ok
+[21:43:16] START gpu1 kv_h2o_30k_dense try1 (~9 remaining)
+[21:43:31] DONE gpu2 ruler_30k_niah_multivalue_fixed rc=0 ok
+[21:43:31] START gpu2 kv_stream_30k_dense try1 (~8 remaining)
+[21:45:21] DONE gpu2 kv_stream_30k_dense rc=0 ok
+[21:45:21] START gpu2 kv_h2o_30k_fixed try1 (~7 remaining)
+[21:47:51] DONE gpu1 kv_h2o_30k_dense rc=0 ok
+[21:47:52] START gpu1 kv_stream_30k_fixed try1 (~6 remaining)
+[21:49:52] DONE gpu2 kv_h2o_30k_fixed rc=0 ok
+[21:49:52] DONE gpu1 kv_stream_30k_fixed rc=0 ok
+[21:49:52] START gpu1 kv_h2o_30k_resampled try1 (~4 remaining)
+[21:49:52] START gpu2 kv_stream_30k_resampled try1 (~4 remaining)
+[21:51:47] DONE gpu2 kv_stream_30k_resampled rc=0 ok
+[21:54:24] DONE gpu1 kv_h2o_30k_resampled rc=0 ok
+[22:14:38] DONE gpu3 ruler_30k_vt_fixed rc=0 ok
+[00:15:44] DONE gpu0 ruler_30k_vt_resampled rc=0 ok
+[00:15:44] === 1.2B BATTERY DONE. done=72/78 failed_or_skipped=['ruler_8k_niah_multikey_dense', 'ruler_8k_niah_multivalue_dense', 'ruler_8k_vt_dense', 'ruler_30k_niah_multikey_dense', 'ruler_30k_niah_multivalue_dense', 'ruler_30k_vt_dense'] ===

epfl_final/logs/eval_30k_battery_a100.log

@@ -0,0 +1,98 @@
+[22:43:10] ckpt resampled -> /mloscratch/homes/navasard/sparse/checkpoints/l30k365m_resampled/run_20260707_035753/checkpoint-11700
+[22:43:10] ckpt fixed_s1 -> /mloscratch/homes/navasard/sparse/checkpoints/l30k365m_fixed_s1/run_20260707_035756/checkpoint-11700
+[22:43:10] ckpt fixed_s2 -> /mloscratch/homes/navasard/sparse/checkpoints/l30k365m_fixed_s2/run_20260707_035759/checkpoint-11700
+[22:43:10] ckpt dense -> /mloscratch/homes/navasard/sparse/checkpoints/l30k365m_dense/run_20260707_035752/checkpoint-11700
+[22:43:10] total jobs=104  already-done=10  gpus=[0, 1, 2, 3]  n=100
+[22:43:10] START gpu0 RULER_niah_multikey_fixed_s1_d120  (~94 remaining)
+[22:43:10] START gpu1 RULER_niah_multikey_fixed_s2_d120  (~94 remaining)
+[22:43:10] START gpu2 RULER_niah_multikey_fixed_s2_d60  (~94 remaining)
+[22:43:10] START gpu3 RULER_niah_multikey_fixed_s2_d31  (~94 remaining)
+[23:07:15] DONE  gpu3 RULER_niah_multikey_fixed_s2_d31 rc=0 ok
+[23:07:15] START gpu3 RULER_niah_multikey_dense_d16  (~92 remaining)
+[23:22:24] DONE  gpu2 RULER_niah_multikey_fixed_s2_d60 rc=0 ok
+[23:22:24] START gpu2 RULER_niah_multivalue_resampled_d31  (~88 remaining)
+[23:22:56] DONE  gpu3 RULER_niah_multikey_dense_d16 rc=0 ok
+[23:22:56] START gpu3 RULER_niah_multivalue_resampled_d16  (~87 remaining)
+[23:38:37] DONE  gpu3 RULER_niah_multivalue_resampled_d16 rc=0 ok
+[23:38:37] START gpu3 RULER_niah_multivalue_fixed_s1_d60  (~84 remaining)
+[23:46:02] DONE  gpu2 RULER_niah_multivalue_resampled_d31 rc=0 ok
+[23:46:02] START gpu2 RULER_niah_multivalue_fixed_s1_d31  (~83 remaining)
+[23:53:58] DONE  gpu1 RULER_niah_multikey_fixed_s2_d120 rc=0 ok
+[23:53:58] START gpu1 RULER_niah_multivalue_fixed_s1_d16  (~82 remaining)
+[23:54:06] DONE  gpu0 RULER_niah_multikey_fixed_s1_d120 rc=0 ok
+[23:54:06] START gpu0 RULER_niah_multivalue_fixed_s2_ddense  (~81 remaining)
+[23:58:06] DONE  gpu0 RULER_niah_multivalue_fixed_s2_ddense rc=0 ok
+[23:58:06] START gpu0 RULER_niah_multivalue_fixed_s2_d120  (~80 remaining)
+[00:09:43] DONE  gpu1 RULER_niah_multivalue_fixed_s1_d16 rc=0 ok
+[00:09:43] START gpu1 RULER_niah_multivalue_fixed_s2_d16  (~76 remaining)
+[00:09:47] DONE  gpu2 RULER_niah_multivalue_fixed_s1_d31 rc=0 ok
+[00:09:47] START gpu2 RULER_niah_multivalue_dense_ddense  (~76 remaining)
+[00:13:47] DONE  gpu2 RULER_niah_multivalue_dense_ddense rc=0 ok
+[00:13:47] START gpu2 RULER_niah_multivalue_dense_d120  (~75 remaining)
+[00:17:32] DONE  gpu3 RULER_niah_multivalue_fixed_s1_d60 rc=0 ok
+[00:17:32] START gpu3 RULER_niah_multivalue_dense_d60  (~74 remaining)
+[00:25:28] DONE  gpu1 RULER_niah_multivalue_fixed_s2_d16 rc=0 ok
+[00:25:28] START gpu1 RULER_niah_multivalue_dense_d16  (~72 remaining)
+[00:41:09] DONE  gpu1 RULER_niah_multivalue_dense_d16 rc=0 ok
+[00:41:09] START gpu1 RULER_vt_resampled_d31  (~68 remaining)
+[00:57:02] DONE  gpu3 RULER_niah_multivalue_dense_d60 rc=0 ok
+[00:57:02] START gpu3 RULER_vt_fixed_s1_d60  (~64 remaining)
+[01:05:22] DONE  gpu1 RULER_vt_resampled_d31 rc=0 ok
+[01:05:22] START gpu1 RULER_vt_fixed_s1_d16  (~62 remaining)
+[01:09:35] DONE  gpu0 RULER_niah_multivalue_fixed_s2_d120 rc=0 ok
+[01:09:35] START gpu0 RULER_vt_fixed_s2_ddense  (~61 remaining)
+[01:13:35] DONE  gpu0 RULER_vt_fixed_s2_ddense rc=0 ok
+[01:13:35] START gpu0 RULER_vt_fixed_s2_d120  (~60 remaining)
+[01:21:23] DONE  gpu1 RULER_vt_fixed_s1_d16 rc=0 ok
+[01:21:23] START gpu1 RULER_vt_fixed_s2_d31  (~58 remaining)
+[01:25:12] DONE  gpu2 RULER_niah_multivalue_dense_d120 rc=0 ok
+[01:25:12] START gpu2 RULER_vt_fixed_s2_d16  (~57 remaining)
+[01:36:48] DONE  gpu3 RULER_vt_fixed_s1_d60 rc=0 ok
+[01:36:48] START gpu3 RULER_vt_dense_d60  (~54 remaining)
+[01:41:17] DONE  gpu2 RULER_vt_fixed_s2_d16 rc=0 ok
+[01:41:17] START gpu2 RULER_vt_dense_d31  (~53 remaining)
+[01:45:33] DONE  gpu1 RULER_vt_fixed_s2_d31 rc=0 ok
+[01:45:33] START gpu1 PASSKEY_resampled_ddense  (~51 remaining)
+[01:49:33] DONE  gpu1 PASSKEY_resampled_ddense rc=0 ok
+[01:49:33] START gpu1 PASSKEY_resampled_d120  (~50 remaining)
+[02:05:18] DONE  gpu2 RULER_vt_dense_d31 rc=0 ok
+[02:05:18] START gpu2 PASSKEY_fixed_s1_ddense  (~46 remaining)
+[02:09:18] DONE  gpu2 PASSKEY_fixed_s1_ddense rc=0 ok
+[02:09:18] START gpu2 PASSKEY_fixed_s1_d120  (~45 remaining)
+[02:16:15] DONE  gpu3 RULER_vt_dense_d60 rc=0 ok
+[02:16:15] START gpu3 PASSKEY_fixed_s1_d60  (~44 remaining)
+[02:25:43] DONE  gpu0 RULER_vt_fixed_s2_d120 rc=0 ok
+[02:25:43] START gpu0 PASSKEY_fixed_s2_ddense  (~41 remaining)
+[02:29:43] DONE  gpu0 PASSKEY_fixed_s2_ddense rc=0 ok
+[02:29:43] START gpu0 PASSKEY_fixed_s2_d120  (~40 remaining)
+[02:55:45] DONE  gpu3 PASSKEY_fixed_s1_d60 rc=0 ok
+[02:55:45] START gpu3 PASSKEY_dense_d60  (~34 remaining)
+[03:01:09] DONE  gpu1 PASSKEY_resampled_d120 rc=0 ok
+[03:01:09] START gpu1 PPL_resampled_resampled_pg19  (~31 remaining)
+[03:20:34] DONE  gpu2 PASSKEY_fixed_s1_d120 rc=0 ok
+[03:20:34] START gpu2 PPL_fixed_s1_resampled_pg19  (~27 remaining)
+[03:34:55] DONE  gpu3 PASSKEY_dense_d60 rc=0 ok
+[03:34:55] START gpu3 PPL_fixed_s1_resampled_arxiv  (~26 remaining)
+[03:40:35] DONE  gpu0 PASSKEY_fixed_s2_d120 rc=0 ok
+[03:40:35] START gpu0 PPL_fixed_s1_fixed_pg19  (~25 remaining)
+[05:25:38] DONE  gpu1 PPL_resampled_resampled_pg19 rc=0 ok
+[05:25:38] START gpu1 PPL_fixed_s2_resampled_pg19  (~23 remaining)
+[05:46:35] DONE  gpu2 PPL_fixed_s1_resampled_pg19 rc=0 ok
+[05:46:35] START gpu2 PPL_dense_resampled_pg19  (~19 remaining)
+[06:04:32] DONE  gpu3 PPL_fixed_s1_resampled_arxiv rc=0 ok
+[06:04:32] START gpu3 PPL_dense_resampled_arxiv  (~18 remaining)
+[06:09:16] DONE  gpu0 PPL_fixed_s1_fixed_pg19 rc=0 ok
+[06:09:16] START gpu0 PPL_dense_fixed_pg19  (~17 remaining)
+[07:50:10] DONE  gpu1 PPL_fixed_s2_resampled_pg19 rc=0 ok
+[07:50:10] START gpu1 KVh2o_resampled  (~15 remaining)
+[07:52:22] DONE  gpu1 KVh2o_resampled rc=0 ok
+[07:52:22] START gpu1 KVstream_resampled  (~14 remaining)
+[07:53:22] DONE  gpu1 KVstream_resampled rc=0 ok
+[07:53:22] START gpu1 KVh2o_fixed_s1  (~13 remaining)
+[07:55:34] DONE  gpu1 KVh2o_fixed_s1 rc=0 ok
+[07:55:34] START gpu1 KVh2o_dense  (~9 remaining)
+[07:57:46] DONE  gpu1 KVh2o_dense rc=0 ok
+[08:09:48] DONE  gpu2 PPL_dense_resampled_pg19 rc=0 ok
+[08:25:38] DONE  gpu3 PPL_dense_resampled_arxiv rc=0 ok
+[08:29:30] DONE  gpu0 PPL_dense_fixed_pg19 rc=0 ok
+[08:29:30] === THIS WORKER DONE (no claimable jobs left) ===

epfl_final/logs/eval_30k_battery_a100_pod2.log

@@ -0,0 +1,102 @@
+[22:51:43] ckpt resampled -> /mloscratch/homes/navasard/sparse/checkpoints/l30k365m_resampled/run_20260707_035753/checkpoint-11700
+[22:51:43] ckpt fixed_s1 -> /mloscratch/homes/navasard/sparse/checkpoints/l30k365m_fixed_s1/run_20260707_035756/checkpoint-11700
+[22:51:43] ckpt fixed_s2 -> /mloscratch/homes/navasard/sparse/checkpoints/l30k365m_fixed_s2/run_20260707_035759/checkpoint-11700
+[22:51:43] ckpt dense -> /mloscratch/homes/navasard/sparse/checkpoints/l30k365m_dense/run_20260707_035752/checkpoint-11700
+[22:51:43] total jobs=104  already-done=10  gpus=[0, 1, 2, 3]  n=100
+[22:51:43] START gpu0 RULER_niah_multikey_fixed_s2_d16  (~94 remaining)
+[22:51:43] START gpu1 RULER_niah_multikey_dense_ddense  (~94 remaining)
+[22:51:43] START gpu2 RULER_niah_multikey_dense_d120  (~94 remaining)
+[22:51:43] START gpu3 RULER_niah_multikey_dense_d60  (~94 remaining)
+[22:56:00] DONE  gpu1 RULER_niah_multikey_dense_ddense rc=0 ok
+[22:56:00] START gpu1 RULER_niah_multikey_dense_d31  (~93 remaining)
+[23:07:48] DONE  gpu0 RULER_niah_multikey_fixed_s2_d16 rc=0 ok
+[23:07:48] START gpu0 RULER_niah_multivalue_resampled_ddense  (~91 remaining)
+[23:11:53] DONE  gpu0 RULER_niah_multivalue_resampled_ddense rc=0 ok
+[23:11:53] START gpu0 RULER_niah_multivalue_resampled_d120  (~90 remaining)
+[23:19:41] DONE  gpu1 RULER_niah_multikey_dense_d31 rc=0 ok
+[23:19:41] START gpu1 RULER_niah_multivalue_resampled_d60  (~89 remaining)
+[23:30:58] DONE  gpu3 RULER_niah_multikey_dense_d60 rc=0 ok
+[23:30:58] START gpu3 RULER_niah_multivalue_fixed_s1_ddense  (~86 remaining)
+[23:35:02] DONE  gpu3 RULER_niah_multivalue_fixed_s1_ddense rc=0 ok
+[23:35:02] START gpu3 RULER_niah_multivalue_fixed_s1_d120  (~85 remaining)
+[23:58:47] DONE  gpu1 RULER_niah_multivalue_resampled_d60 rc=0 ok
+[23:58:47] START gpu1 RULER_niah_multivalue_fixed_s2_d60  (~79 remaining)
+[00:03:32] DONE  gpu2 RULER_niah_multikey_dense_d120 rc=0 ok
+[00:03:32] START gpu2 RULER_niah_multivalue_fixed_s2_d31  (~78 remaining)
+[00:22:57] DONE  gpu0 RULER_niah_multivalue_resampled_d120 rc=0 ok
+[00:22:57] START gpu0 RULER_niah_multivalue_dense_d31  (~73 remaining)
+[00:27:25] DONE  gpu2 RULER_niah_multivalue_fixed_s2_d31 rc=0 ok
+[00:27:25] START gpu2 RULER_vt_resampled_ddense  (~71 remaining)
+[00:31:25] DONE  gpu2 RULER_vt_resampled_ddense rc=0 ok
+[00:31:25] START gpu2 RULER_vt_resampled_d120  (~70 remaining)
+[00:37:42] DONE  gpu1 RULER_niah_multivalue_fixed_s2_d60 rc=0 ok
+[00:37:42] START gpu1 RULER_vt_resampled_d60  (~69 remaining)
+[00:46:06] DONE  gpu3 RULER_niah_multivalue_fixed_s1_d120 rc=0 ok
+[00:46:06] START gpu3 RULER_vt_resampled_d16  (~67 remaining)
+[00:46:34] DONE  gpu0 RULER_niah_multivalue_dense_d31 rc=0 ok
+[00:46:34] START gpu0 RULER_vt_fixed_s1_ddense  (~66 remaining)
+[00:50:34] DONE  gpu0 RULER_vt_fixed_s1_ddense rc=0 ok
+[00:50:34] START gpu0 RULER_vt_fixed_s1_d120  (~65 remaining)
+[01:01:51] DONE  gpu3 RULER_vt_resampled_d16 rc=0 ok
+[01:01:51] START gpu3 RULER_vt_fixed_s1_d31  (~63 remaining)
+[01:16:52] DONE  gpu1 RULER_vt_resampled_d60 rc=0 ok
+[01:16:52] START gpu1 RULER_vt_fixed_s2_d60  (~59 remaining)
+[01:25:40] DONE  gpu3 RULER_vt_fixed_s1_d31 rc=0 ok
+[01:25:40] START gpu3 RULER_vt_dense_ddense  (~56 remaining)
+[01:29:41] DONE  gpu3 RULER_vt_dense_ddense rc=0 ok
+[01:29:41] START gpu3 RULER_vt_dense_d120  (~55 remaining)
+[01:42:21] DONE  gpu2 RULER_vt_resampled_d120 rc=0 ok
+[01:42:21] START gpu2 RULER_vt_dense_d16  (~52 remaining)
+[01:55:54] DONE  gpu1 RULER_vt_fixed_s2_d60 rc=0 ok
+[01:55:54] START gpu1 PASSKEY_resampled_d60  (~49 remaining)
+[01:58:06] DONE  gpu2 RULER_vt_dense_d16 rc=0 ok
+[01:58:06] START gpu2 PASSKEY_resampled_d31  (~48 remaining)
+[02:01:43] DONE  gpu0 RULER_vt_fixed_s1_d120 rc=0 ok
+[02:01:43] START gpu0 PASSKEY_resampled_d16  (~47 remaining)
+[02:17:20] DONE  gpu0 PASSKEY_resampled_d16 rc=0 ok
+[02:17:20] START gpu0 PASSKEY_fixed_s1_d31  (~43 remaining)
+[02:21:44] DONE  gpu2 PASSKEY_resampled_d31 rc=0 ok
+[02:21:44] START gpu2 PASSKEY_fixed_s1_d16  (~42 remaining)
+[02:34:57] DONE  gpu1 PASSKEY_resampled_d60 rc=0 ok
+[02:34:57] START gpu1 PASSKEY_fixed_s2_d60  (~39 remaining)
+[02:37:25] DONE  gpu2 PASSKEY_fixed_s1_d16 rc=0 ok
+[02:37:25] START gpu2 PASSKEY_fixed_s2_d31  (~38 remaining)
+[02:40:57] DONE  gpu3 RULER_vt_dense_d120 rc=0 ok
+[02:40:57] START gpu3 PASSKEY_fixed_s2_d16  (~36 remaining)
+[02:41:01] DONE  gpu0 PASSKEY_fixed_s1_d31 rc=0 ok
+[02:41:01] START gpu0 PASSKEY_dense_ddense  (~36 remaining)
+[02:45:01] DONE  gpu0 PASSKEY_dense_ddense rc=0 ok
+[02:45:01] START gpu0 PASSKEY_dense_d120  (~35 remaining)
+[02:56:38] DONE  gpu3 PASSKEY_fixed_s2_d16 rc=0 ok
+[02:56:38] START gpu3 PASSKEY_dense_d31  (~33 remaining)
+[03:01:06] DONE  gpu2 PASSKEY_fixed_s2_d31 rc=0 ok
+[03:01:07] START gpu2 PASSKEY_dense_d16  (~32 remaining)
+[03:13:55] DONE  gpu1 PASSKEY_fixed_s2_d60 rc=0 ok
+[03:13:55] START gpu1 PPL_resampled_resampled_arxiv  (~30 remaining)
+[03:16:51] DONE  gpu2 PASSKEY_dense_d16 rc=0 ok
+[03:16:51] START gpu2 PPL_resampled_fixed_pg19  (~29 remaining)
+[03:20:12] DONE  gpu3 PASSKEY_dense_d31 rc=0 ok
+[03:20:12] START gpu3 PPL_resampled_fixed_arxiv  (~28 remaining)
+[03:55:54] DONE  gpu0 PASSKEY_dense_d120 rc=0 ok
+[03:55:54] START gpu0 PPL_fixed_s1_fixed_arxiv  (~24 remaining)
+[05:33:19] DONE  gpu1 PPL_resampled_resampled_arxiv rc=0 ok
+[05:33:19] START gpu1 PPL_fixed_s2_resampled_arxiv  (~22 remaining)
+[05:35:43] DONE  gpu2 PPL_resampled_fixed_pg19 rc=0 ok
+[05:35:44] START gpu2 PPL_fixed_s2_fixed_pg19  (~21 remaining)
+[05:39:28] DONE  gpu3 PPL_resampled_fixed_arxiv rc=0 ok
+[05:39:28] START gpu3 PPL_fixed_s2_fixed_arxiv  (~20 remaining)
+[06:15:50] DONE  gpu0 PPL_fixed_s1_fixed_arxiv rc=0 ok
+[06:15:50] START gpu0 PPL_dense_fixed_arxiv  (~16 remaining)
+[07:53:27] DONE  gpu1 PPL_fixed_s2_resampled_arxiv rc=0 ok
+[07:53:27] START gpu1 KVstream_fixed_s1  (~12 remaining)
+[07:54:27] DONE  gpu1 KVstream_fixed_s1 rc=0 ok
+[07:54:27] START gpu1 KVh2o_fixed_s2  (~11 remaining)
+[07:55:15] DONE  gpu2 PPL_fixed_s2_fixed_pg19 rc=0 ok
+[07:55:15] START gpu2 KVstream_fixed_s2  (~10 remaining)
+[07:56:15] DONE  gpu2 KVstream_fixed_s2 rc=0 ok
+[07:56:15] START gpu2 KVstream_dense  (~8 remaining)
+[07:56:39] DONE  gpu1 KVh2o_fixed_s2 rc=0 ok
+[07:57:16] DONE  gpu2 KVstream_dense rc=0 ok
+[07:59:48] DONE  gpu3 PPL_fixed_s2_fixed_arxiv rc=0 ok
+[08:35:12] DONE  gpu0 PPL_dense_fixed_arxiv rc=0 ok
+[08:35:12] === THIS WORKER DONE (no claimable jobs left) ===

epfl_final/logs/eval_mixed_1p2b_waiter.log

@@ -0,0 +1 @@
+[2026-07-18T01:59:03Z] waiting for current eval jobs

epfl_final/logs/eval_moba_matched_waiter.log

@@ -0,0 +1,38 @@
+Checkpoint activations: False
+[moba1p2b_matched] loaded top_k=64
+/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(
+[moba1p2b_matched] top_k=64 depth=0.1 exact=0.00 tokacc=0.47
+[moba1p2b_matched] top_k=64 depth=0.3 exact=0.00 tokacc=0.65
+[moba1p2b_matched] top_k=64 depth=0.5 exact=0.32 tokacc=0.90
+[moba1p2b_matched] top_k=64 depth=0.7 exact=0.68 tokacc=0.97
+[moba1p2b_matched] top_k=64 depth=0.9 exact=0.84 tokacc=0.98
+Checkpoint activations: False
+[moba1p2b_matched] loaded top_k=31
+[moba1p2b_matched] top_k=31 depth=0.1 exact=0.00 tokacc=0.34
+[moba1p2b_matched] top_k=31 depth=0.3 exact=0.00 tokacc=0.32
+[moba1p2b_matched] top_k=31 depth=0.5 exact=0.00 tokacc=0.32
+[moba1p2b_matched] top_k=31 depth=0.7 exact=0.26 tokacc=0.88
+[moba1p2b_matched] top_k=31 depth=0.9 exact=0.86 tokacc=0.99
+Checkpoint activations: False
+[moba1p2b_matched] loaded top_k=16
+[moba1p2b_matched] top_k=16 depth=0.1 exact=0.00 tokacc=0.35
+[moba1p2b_matched] top_k=16 depth=0.3 exact=0.00 tokacc=0.33
+[moba1p2b_matched] top_k=16 depth=0.5 exact=0.00 tokacc=0.32
+[moba1p2b_matched] top_k=16 depth=0.7 exact=0.00 tokacc=0.33
+[moba1p2b_matched] top_k=16 depth=0.9 exact=0.21 tokacc=0.88
+Checkpoint activations: False
+[moba1p2b_matched] loaded top_k=8
+[moba1p2b_matched] top_k=8 depth=0.1 exact=0.00 tokacc=0.34
+[moba1p2b_matched] top_k=8 depth=0.3 exact=0.00 tokacc=0.34
+[moba1p2b_matched] top_k=8 depth=0.5 exact=0.00 tokacc=0.32
+[moba1p2b_matched] top_k=8 depth=0.7 exact=0.00 tokacc=0.34
+[moba1p2b_matched] top_k=8 depth=0.9 exact=0.00 tokacc=0.47
+Checkpoint activations: False
+[moba1p2b_matched] loaded top_k=4
+[moba1p2b_matched] top_k=4 depth=0.1 exact=0.00 tokacc=0.33
+[moba1p2b_matched] top_k=4 depth=0.3 exact=0.00 tokacc=0.34
+[moba1p2b_matched] top_k=4 depth=0.5 exact=0.00 tokacc=0.33
+[moba1p2b_matched] top_k=4 depth=0.7 exact=0.00 tokacc=0.34
+[moba1p2b_matched] top_k=4 depth=0.9 exact=0.00 tokacc=0.33
+[moba1p2b_matched] wrote /mloscratch/homes/navasard/sparse/SparseAtt/eval/results/L1_2B_BATTERY/passkey_8k_moba_matched_topk.tsv

epfl_final/logs/eval_moba_matched_waiter.pid

@@ -0,0 +1 @@
+12232

epfl_final/logs/l30k1_2b_dense_h200.log

@@ -0,0 +1,595 @@
+W0708 23:51:56.633000 767 site-packages/torch/distributed/run.py:851] 
+W0708 23:51:56.633000 767 site-packages/torch/distributed/run.py:851] *****************************************
+W0708 23:51:56.633000 767 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. 
+W0708 23:51:56.633000 767 site-packages/torch/distributed/run.py:851] *****************************************
+Auto-resuming from latest checkpoint: /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_dense/run_20260708_065119/checkpoint-4800
+Resuming: Using existing output directory /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_dense/run_20260708_065119
+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_dense/run_20260708_065119",
+  "resume_from_checkpoint": "/mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_dense/run_20260708_065119/checkpoint-4800",
+  "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": false,
+  "strategy": "ddp",
+  "mixed_precision": "bf16",
+  "wandb_enabled": true,
+  "wandb_project": "stochastic_sparsity",
+  "wandb_name": "1.2b dense 30k longctx",
+  "wandb_tags": [
+    "qwen3",
+    "1.2b",
+    "dense",
+    "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": null,
+  "sparse_mask_seed": 0,
+  "dense_mode": true,
+  "dense_step_intervals": null
+}
+Loading tokenizer from /mloscratch/homes/navasard/qwen3_06b_base...
+Resuming from checkpoint: /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_dense/run_20260708_065119/checkpoint-4800
+Checkpoint activations: FalseCheckpoint activations: False
+Checkpoint activations: FalseCheckpoint activations: False
+
+
+dense_mode=True: sparse attention layers will use standard causal SDPA.
+
+
+================================================================================
+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
+--------------------------------------------------------------------------------
+================================================================================================================================================================
+Layer (type:depth-idx)                                       Input Shape               Output Shape              Param #                   Trainable
+================================================================================================================================================================
+Qwen3ForCausalLM                                             [1, 30720]                --                        --                        True
+├─Qwen3Model: 1-1                                            --                        [1, 30720, 1536]          --                        True
+│    └─Embedding: 2-1                                        [1, 30720]                [1, 30720, 1536]          233,373,696               True
+│    └─ModuleList: 2-2                                       --                        --                        --                        True
+│    │    └─Qwen3DecoderLayer: 3-1                           [1, 30720, 1536]          [1, 30720, 1536]          34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-2                           [1, 30720, 1536]          [1, 30720, 1536]          34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-3                           [1, 30720, 1536]          [1, 30720, 1536]          34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-4                           [1, 30720, 1536]          [1, 30720, 1536]          34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-5                           [1, 30720, 1536]          [1, 30720, 1536]          34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-6                           [1, 30720, 1536]          [1, 30720, 1536]          34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-7                           [1, 30720, 1536]          [1, 30720, 1536]          34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-8                           [1, 30720, 1536]          [1, 30720, 1536]          34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-9                           [1, 30720, 1536]          [1, 30720, 1536]          34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-10                          [1, 30720, 1536]          [1, 30720, 1536]          34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-11                          [1, 30720, 1536]          [1, 30720, 1536]          34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-12                          [1, 30720, 1536]          [1, 30720, 1536]          34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-13                          [1, 30720, 1536]          [1, 30720, 1536]          34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-14                          [1, 30720, 1536]          [1, 30720, 1536]          34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-15                          [1, 30720, 1536]          [1, 30720, 1536]          34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-16                          [1, 30720, 1536]          [1, 30720, 1536]          34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-17                          [1, 30720, 1536]          [1, 30720, 1536]          34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-18                          [1, 30720, 1536]          [1, 30720, 1536]          34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-19                          [1, 30720, 1536]          [1, 30720, 1536]          34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-20                          [1, 30720, 1536]          [1, 30720, 1536]          34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-21                          [1, 30720, 1536]          [1, 30720, 1536]          34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-22                          [1, 30720, 1536]          [1, 30720, 1536]          34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-23                          [1, 30720, 1536]          [1, 30720, 1536]          34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-24                          [1, 30720, 1536]          [1, 30720, 1536]          34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-25                          [1, 30720, 1536]          [1, 30720, 1536]          34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-26                          [1, 30720, 1536]          [1, 30720, 1536]          34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-27                          [1, 30720, 1536]          [1, 30720, 1536]          34,606,336                True
+│    │    └─Qwen3DecoderLayer: 3-28                          [1, 30720, 1536]          [1, 30720, 1536]          34,606,336                True
+│    └─Qwen3RMSNorm: 2-3                                     [1, 30720, 1536]          [1, 30720, 1536]          1,536                     True
+├─Linear: 1-2                                                [1, 30720, 1536]          [1, 30720, 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): 1.44
+================================================================================================================================================================
+Input size (MB): 0.25
+Forward/backward pass size (MB): 98319.73
+Params size (MB): 2871.45
+Estimated Total Size (MB): 101191.43
+================================================================================================================================================================
+  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)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_dense/run_20260708_065119/checkpoint-4800...
+Overriding checkpoint LR with config LR: 0.0005
+Resumed from checkpoint at step 4800 (epoch 0, micro_step 24000)
+Resuming training from step 4800
+Resuming wandb run cd66e13p
+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 cd66e13p
+wandb: Tracking run with wandb version 0.26.1
+wandb: Run data is saved locally in /mloscratch/homes/navasard/sparse/SparseAtt/wandb/run-20260708_235322-cd66e13p
+wandb: Run `wandb offline` to turn off syncing.
+wandb: Resuming run 1.2b dense 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/cd66e13p
+Starting training...
+Validation at step 4800: Loss = 2.4541, Perplexity = 11.64
+wandb: WARNING Tried to log to step 4800 that is less than the current step 4801. 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 4801: dense (dense step interval)
+Step 4810/37465 | Loss: 2.5212 | LR: 4.85e-04 | 11880.1ms/step | 1.7 samples/s
+Step 4820/37465 | Loss: 2.5226 | LR: 4.85e-04 | 6521.1ms/step | 3.1 samples/s
+Step 4830/37465 | Loss: 2.5182 | LR: 4.85e-04 | 6551.5ms/step | 3.1 samples/s
+Step 4840/37465 | Loss: 2.3355 | LR: 4.85e-04 | 6555.2ms/step | 3.1 samples/s
+Step 4850/37465 | Loss: 2.4902 | LR: 4.85e-04 | 6552.2ms/step | 3.1 samples/s
+Step 4860/37465 | Loss: 2.2283 | LR: 4.85e-04 | 6548.1ms/step | 3.1 samples/s
+Step 4870/37465 | Loss: 2.5547 | LR: 4.85e-04 | 6532.2ms/step | 3.1 samples/s
+Step 4880/37465 | Loss: 2.8112 | LR: 4.85e-04 | 6572.8ms/step | 3.0 samples/s
+Step 4890/37465 | Loss: 2.6531 | LR: 4.84e-04 | 6559.7ms/step | 3.0 samples/s
+Step 4900/37465 | Loss: 2.3631 | LR: 4.84e-04 | 6551.0ms/step | 3.1 samples/s
+Step 4910/37465 | Loss: 2.4640 | LR: 4.84e-04 | 6560.3ms/step | 3.0 samples/s
+Step 4920/37465 | Loss: 2.6113 | LR: 4.84e-04 | 6547.7ms/step | 3.1 samples/s
+Step 4930/37465 | Loss: 2.1574 | LR: 4.84e-04 | 6542.7ms/step | 3.1 samples/s
+Step 4940/37465 | Loss: 2.3762 | LR: 4.84e-04 | 6550.6ms/step | 3.1 samples/s
+Step 4950/37465 | Loss: 2.4277 | LR: 4.84e-04 | 6550.9ms/step | 3.1 samples/s
+Step 4960/37465 | Loss: 2.4383 | LR: 4.84e-04 | 6544.4ms/step | 3.1 samples/s
+Step 4970/37465 | Loss: 2.4342 | LR: 4.84e-04 | 6542.3ms/step | 3.1 samples/s
+Step 4980/37465 | Loss: 2.3146 | LR: 4.84e-04 | 6538.2ms/step | 3.1 samples/s
+Step 4990/37465 | Loss: 2.4049 | LR: 4.84e-04 | 6534.6ms/step | 3.1 samples/s
+Step 5000/37465 | Loss: 2.6615 | LR: 4.84e-04 | 6537.4ms/step | 3.1 samples/s
+Validation at step 5000: Loss = 2.4427, Perplexity = 11.50
+Step 5010/37465 | Loss: 2.3587 | LR: 4.84e-04 | 7486.4ms/step | 2.7 samples/s
+Step 5020/37465 | Loss: 2.4034 | LR: 4.83e-04 | 6539.4ms/step | 3.1 samples/s
+Step 5030/37465 | Loss: 2.5192 | LR: 4.83e-04 | 6537.5ms/step | 3.1 samples/s
+Step 5040/37465 | Loss: 2.3445 | LR: 4.83e-04 | 6552.4ms/step | 3.1 samples/s
+Step 5050/37465 | Loss: 2.3393 | LR: 4.83e-04 | 6553.0ms/step | 3.1 samples/s
+Step 5060/37465 | Loss: 2.5758 | LR: 4.83e-04 | 6553.5ms/step | 3.1 samples/s
+Step 5070/37465 | Loss: 2.6069 | LR: 4.83e-04 | 6551.4ms/step | 3.1 samples/s
+Step 5080/37465 | Loss: 2.3409 | LR: 4.83e-04 | 6552.7ms/step | 3.1 samples/s
+Step 5090/37465 | Loss: 2.2568 | LR: 4.83e-04 | 6551.1ms/step | 3.1 samples/s
+Step 5100/37465 | Loss: 2.6086 | LR: 4.83e-04 | 6555.9ms/step | 3.1 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_dense/run_20260708_065119/checkpoint-5100
+Step 5110/37465 | Loss: 2.3859 | LR: 4.83e-04 | 7465.1ms/step | 2.7 samples/s
+Step 5120/37465 | Loss: 2.2411 | LR: 4.83e-04 | 6518.4ms/step | 3.1 samples/s
+Step 5130/37465 | Loss: 2.2797 | LR: 4.83e-04 | 6530.5ms/step | 3.1 samples/s
+Step 5140/37465 | Loss: 2.2900 | LR: 4.83e-04 | 6526.0ms/step | 3.1 samples/s
+Step 5150/37465 | Loss: 2.3875 | LR: 4.82e-04 | 6534.8ms/step | 3.1 samples/s
+Step 5160/37465 | Loss: 2.2290 | LR: 4.82e-04 | 6555.0ms/step | 3.1 samples/s
+Step 5170/37465 | Loss: 2.7020 | LR: 4.82e-04 | 6555.2ms/step | 3.1 samples/s
+Step 5180/37465 | Loss: 2.3337 | LR: 4.82e-04 | 6557.8ms/step | 3.0 samples/s
+Step 5190/37465 | Loss: 2.5715 | LR: 4.82e-04 | 6565.2ms/step | 3.0 samples/s
+Step 5200/37465 | Loss: 2.5633 | LR: 4.82e-04 | 6564.5ms/step | 3.0 samples/s
+Step 5210/37465 | Loss: 2.2979 | LR: 4.82e-04 | 6564.3ms/step | 3.0 samples/s
+Step 5220/37465 | Loss: 2.3175 | LR: 4.82e-04 | 6560.3ms/step | 3.0 samples/s
+Step 5230/37465 | Loss: 2.4364 | LR: 4.82e-04 | 6554.6ms/step | 3.1 samples/s
+Step 5240/37465 | Loss: 2.4634 | LR: 4.82e-04 | 6556.4ms/step | 3.1 samples/s
+Step 5250/37465 | Loss: 2.5823 | LR: 4.82e-04 | 6551.7ms/step | 3.1 samples/s
+Step 5260/37465 | Loss: 2.5625 | LR: 4.82e-04 | 6561.4ms/step | 3.0 samples/s
+Step 5270/37465 | Loss: 2.4340 | LR: 4.82e-04 | 6554.5ms/step | 3.1 samples/s
+Step 5280/37465 | Loss: 2.3518 | LR: 4.81e-04 | 6557.6ms/step | 3.0 samples/s
+Step 5290/37465 | Loss: 2.5347 | LR: 4.81e-04 | 6558.6ms/step | 3.0 samples/s
+Step 5300/37465 | Loss: 2.6521 | LR: 4.81e-04 | 6560.6ms/step | 3.0 samples/s
+Step 5310/37465 | Loss: 2.3147 | LR: 4.81e-04 | 6557.4ms/step | 3.0 samples/s
+Step 5320/37465 | Loss: 2.3823 | LR: 4.81e-04 | 6559.9ms/step | 3.0 samples/s
+Step 5330/37465 | Loss: 2.1705 | LR: 4.81e-04 | 6557.8ms/step | 3.0 samples/s
+Step 5340/37465 | Loss: 2.4096 | LR: 4.81e-04 | 6562.3ms/step | 3.0 samples/s
+Step 5350/37465 | Loss: 2.6535 | LR: 4.81e-04 | 6565.9ms/step | 3.0 samples/s
+Step 5360/37465 | Loss: 2.5062 | LR: 4.81e-04 | 6564.4ms/step | 3.0 samples/s
+Step 5370/37465 | Loss: 2.5248 | LR: 4.81e-04 | 6560.1ms/step | 3.0 samples/s
+Step 5380/37465 | Loss: 2.4481 | LR: 4.81e-04 | 6564.9ms/step | 3.0 samples/s
+Step 5390/37465 | Loss: 2.3757 | LR: 4.81e-04 | 6558.4ms/step | 3.0 samples/s
+Step 5400/37465 | Loss: 2.5138 | LR: 4.80e-04 | 6560.9ms/step | 3.0 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_dense/run_20260708_065119/checkpoint-5400
+Step 5410/37465 | Loss: 2.3998 | LR: 4.80e-04 | 7501.5ms/step | 2.7 samples/s
+Step 5420/37465 | Loss: 2.3600 | LR: 4.80e-04 | 6558.0ms/step | 3.0 samples/s
+Step 5430/37465 | Loss: 2.2952 | LR: 4.80e-04 | 6558.1ms/step | 3.0 samples/s
+Step 5440/37465 | Loss: 2.4305 | LR: 4.80e-04 | 6556.2ms/step | 3.1 samples/s
+Step 5450/37465 | Loss: 2.3481 | LR: 4.80e-04 | 6559.5ms/step | 3.0 samples/s
+Step 5460/37465 | Loss: 2.5696 | LR: 4.80e-04 | 6569.7ms/step | 3.0 samples/s
+Step 5470/37465 | Loss: 2.5299 | LR: 4.80e-04 | 6569.2ms/step | 3.0 samples/s
+Step 5480/37465 | Loss: 2.3464 | LR: 4.80e-04 | 6569.9ms/step | 3.0 samples/s
+Step 5490/37465 | Loss: 2.6104 | LR: 4.80e-04 | 6566.9ms/step | 3.0 samples/s
+Step 5500/37465 | Loss: 2.3538 | LR: 4.80e-04 | 6567.7ms/step | 3.0 samples/s
+Validation at step 5500: Loss = 2.4256, Perplexity = 11.31
+Step 5510/37465 | Loss: 2.2641 | LR: 4.80e-04 | 7527.4ms/step | 2.7 samples/s
+Step 5520/37465 | Loss: 2.2908 | LR: 4.79e-04 | 6565.3ms/step | 3.0 samples/s
+Step 5530/37465 | Loss: 2.3937 | LR: 4.79e-04 | 6570.7ms/step | 3.0 samples/s
+Step 5540/37465 | Loss: 2.1922 | LR: 4.79e-04 | 6566.4ms/step | 3.0 samples/s
+Step 5550/37465 | Loss: 2.2693 | LR: 4.79e-04 | 6569.0ms/step | 3.0 samples/s
+Step 5560/37465 | Loss: 2.6477 | LR: 4.79e-04 | 6571.2ms/step | 3.0 samples/s
+Step 5570/37465 | Loss: 2.2594 | LR: 4.79e-04 | 6567.9ms/step | 3.0 samples/s
+Step 5580/37465 | Loss: 2.4192 | LR: 4.79e-04 | 6565.9ms/step | 3.0 samples/s
+Step 5590/37465 | Loss: 2.4362 | LR: 4.79e-04 | 6568.9ms/step | 3.0 samples/s
+Step 5600/37465 | Loss: 2.2760 | LR: 4.79e-04 | 6566.3ms/step | 3.0 samples/s
+Step 5610/37465 | Loss: 2.3761 | LR: 4.79e-04 | 6570.7ms/step | 3.0 samples/s
+Step 5620/37465 | Loss: 2.4175 | LR: 4.79e-04 | 6568.8ms/step | 3.0 samples/s
+Step 5630/37465 | Loss: 2.4994 | LR: 4.79e-04 | 6571.3ms/step | 3.0 samples/s
+Step 5640/37465 | Loss: 2.3467 | LR: 4.78e-04 | 6569.4ms/step | 3.0 samples/s
+Step 5650/37465 | Loss: 2.2520 | LR: 4.78e-04 | 6571.0ms/step | 3.0 samples/s
+Step 5660/37465 | Loss: 2.4173 | LR: 4.78e-04 | 6570.9ms/step | 3.0 samples/s
+Step 5670/37465 | Loss: 2.4902 | LR: 4.78e-04 | 6569.9ms/step | 3.0 samples/s
+Step 5680/37465 | Loss: 2.4185 | LR: 4.78e-04 | 6574.8ms/step | 3.0 samples/s
+Step 5690/37465 | Loss: 2.3752 | LR: 4.78e-04 | 6575.5ms/step | 3.0 samples/s
+Step 5700/37465 | Loss: 2.5392 | LR: 4.78e-04 | 6571.3ms/step | 3.0 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_dense/run_20260708_065119/checkpoint-5700
+Step 5710/37465 | Loss: 2.3992 | LR: 4.78e-04 | 7504.5ms/step | 2.7 samples/s
+Step 5720/37465 | Loss: 2.3711 | LR: 4.78e-04 | 6559.7ms/step | 3.0 samples/s
+Step 5730/37465 | Loss: 2.3929 | LR: 4.78e-04 | 6558.6ms/step | 3.0 samples/s
+Step 5740/37465 | Loss: 2.3737 | LR: 4.78e-04 | 6561.3ms/step | 3.0 samples/s
+Step 5750/37465 | Loss: 2.4154 | LR: 4.77e-04 | 6562.0ms/step | 3.0 samples/s
+Step 5760/37465 | Loss: 2.1260 | LR: 4.77e-04 | 6561.1ms/step | 3.0 samples/s
+Step 5770/37465 | Loss: 2.4578 | LR: 4.77e-04 | 6559.0ms/step | 3.0 samples/s
+Step 5780/37465 | Loss: 2.3798 | LR: 4.77e-04 | 6557.7ms/step | 3.0 samples/s
+Step 5790/37465 | Loss: 2.2646 | LR: 4.77e-04 | 6554.3ms/step | 3.1 samples/s
+Step 5800/37465 | Loss: 2.4722 | LR: 4.77e-04 | 6554.7ms/step | 3.1 samples/s
+Step 5810/37465 | Loss: 2.4013 | LR: 4.77e-04 | 6568.4ms/step | 3.0 samples/s
+Step 5820/37465 | Loss: 2.1587 | LR: 4.77e-04 | 6562.6ms/step | 3.0 samples/s
+Step 5830/37465 | Loss: 2.3607 | LR: 4.77e-04 | 6562.7ms/step | 3.0 samples/s
+Step 5840/37465 | Loss: 2.3873 | LR: 4.77e-04 | 6560.7ms/step | 3.0 samples/s
+Step 5850/37465 | Loss: 2.2026 | LR: 4.77e-04 | 6564.8ms/step | 3.0 samples/s
+Step 5860/37465 | Loss: 2.1258 | LR: 4.76e-04 | 6560.7ms/step | 3.0 samples/s
+Step 5870/37465 | Loss: 2.4055 | LR: 4.76e-04 | 6558.7ms/step | 3.0 samples/s
+Step 5880/37465 | Loss: 2.4101 | LR: 4.76e-04 | 6558.8ms/step | 3.0 samples/s
+Step 5890/37465 | Loss: 2.4857 | LR: 4.76e-04 | 6559.3ms/step | 3.0 samples/s
+Step 5900/37465 | Loss: 2.5452 | LR: 4.76e-04 | 6562.3ms/step | 3.0 samples/s
+Step 5910/37465 | Loss: 2.2432 | LR: 4.76e-04 | 6564.1ms/step | 3.0 samples/s
+Step 5920/37465 | Loss: 2.3863 | LR: 4.76e-04 | 6561.5ms/step | 3.0 samples/s
+Step 5930/37465 | Loss: 2.3389 | LR: 4.76e-04 | 6563.5ms/step | 3.0 samples/s
+Step 5940/37465 | Loss: 2.6159 | LR: 4.76e-04 | 6561.0ms/step | 3.0 samples/s
+Step 5950/37465 | Loss: 2.4951 | LR: 4.76e-04 | 6562.2ms/step | 3.0 samples/s
+Step 5960/37465 | Loss: 2.2540 | LR: 4.76e-04 | 6564.7ms/step | 3.0 samples/s
+Step 5970/37465 | Loss: 2.2551 | LR: 4.75e-04 | 6562.5ms/step | 3.0 samples/s
+Step 5980/37465 | Loss: 2.3123 | LR: 4.75e-04 | 6559.4ms/step | 3.0 samples/s
+Step 5990/37465 | Loss: 2.3489 | LR: 4.75e-04 | 6559.6ms/step | 3.0 samples/s
+Step 6000/37465 | Loss: 2.2466 | LR: 4.75e-04 | 6565.2ms/step | 3.0 samples/s
+Validation at step 6000: Loss = 2.3910, Perplexity = 10.92
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_dense/run_20260708_065119/checkpoint-6000
+Step 6010/37465 | Loss: 2.2487 | LR: 4.75e-04 | 8437.7ms/step | 2.4 samples/s
+Step 6020/37465 | Loss: 2.5089 | LR: 4.75e-04 | 6553.4ms/step | 3.1 samples/s
+Step 6030/37465 | Loss: 2.4710 | LR: 4.75e-04 | 6555.7ms/step | 3.1 samples/s
+Step 6040/37465 | Loss: 2.3070 | LR: 4.75e-04 | 6561.9ms/step | 3.0 samples/s
+Step 6050/37465 | Loss: 2.2184 | LR: 4.75e-04 | 6554.5ms/step | 3.1 samples/s
+Step 6060/37465 | Loss: 2.5625 | LR: 4.75e-04 | 6558.8ms/step | 3.0 samples/s
+Step 6070/37465 | Loss: 2.3175 | LR: 4.75e-04 | 6557.2ms/step | 3.1 samples/s
+Step 6080/37465 | Loss: 2.5400 | LR: 4.74e-04 | 6554.6ms/step | 3.1 samples/s
+Step 6090/37465 | Loss: 2.3212 | LR: 4.74e-04 | 6557.7ms/step | 3.0 samples/s
+Step 6100/37465 | Loss: 2.1518 | LR: 4.74e-04 | 6552.0ms/step | 3.1 samples/s
+Step 6110/37465 | Loss: 2.5764 | LR: 4.74e-04 | 6556.7ms/step | 3.1 samples/s
+Step 6120/37465 | Loss: 2.5020 | LR: 4.74e-04 | 6555.7ms/step | 3.1 samples/s
+Step 6130/37465 | Loss: 2.3996 | LR: 4.74e-04 | 6554.5ms/step | 3.1 samples/s
+Step 6140/37465 | Loss: 2.3238 | LR: 4.74e-04 | 6556.5ms/step | 3.1 samples/s
+Step 6150/37465 | Loss: 2.4932 | LR: 4.74e-04 | 6551.1ms/step | 3.1 samples/s
+Step 6160/37465 | Loss: 2.4616 | LR: 4.74e-04 | 6552.6ms/step | 3.1 samples/s
+Step 6170/37465 | Loss: 2.3157 | LR: 4.74e-04 | 6552.4ms/step | 3.1 samples/s
+Step 6180/37465 | Loss: 2.4260 | LR: 4.73e-04 | 6557.7ms/step | 3.0 samples/s
+Step 6190/37465 | Loss: 2.4665 | LR: 4.73e-04 | 6559.5ms/step | 3.0 samples/s
+Step 6200/37465 | Loss: 2.3667 | LR: 4.73e-04 | 6559.9ms/step | 3.0 samples/s
+Step 6210/37465 | Loss: 2.1747 | LR: 4.73e-04 | 6563.3ms/step | 3.0 samples/s
+Step 6220/37465 | Loss: 2.4364 | LR: 4.73e-04 | 6564.4ms/step | 3.0 samples/s
+Step 6230/37465 | Loss: 2.3129 | LR: 4.73e-04 | 6559.6ms/step | 3.0 samples/s
+Step 6240/37465 | Loss: 2.4389 | LR: 4.73e-04 | 6562.2ms/step | 3.0 samples/s
+Step 6250/37465 | Loss: 2.2895 | LR: 4.73e-04 | 6556.4ms/step | 3.1 samples/s
+Step 6260/37465 | Loss: 2.2363 | LR: 4.73e-04 | 6558.7ms/step | 3.0 samples/s
+Step 6270/37465 | Loss: 1.9668 | LR: 4.73e-04 | 6556.2ms/step | 3.1 samples/s
+Step 6280/37465 | Loss: 2.1777 | LR: 4.73e-04 | 6553.0ms/step | 3.1 samples/s
+Step 6290/37465 | Loss: 2.3693 | LR: 4.72e-04 | 6551.1ms/step | 3.1 samples/s
+Step 6300/37465 | Loss: 2.4401 | LR: 4.72e-04 | 6554.8ms/step | 3.1 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_dense/run_20260708_065119/checkpoint-6300
+Step 6310/37465 | Loss: 2.2242 | LR: 4.72e-04 | 7501.0ms/step | 2.7 samples/s
+Step 6320/37465 | Loss: 2.3058 | LR: 4.72e-04 | 6546.0ms/step | 3.1 samples/s
+Step 6330/37465 | Loss: 2.4286 | LR: 4.72e-04 | 6563.9ms/step | 3.0 samples/s
+Step 6340/37465 | Loss: 2.2399 | LR: 4.72e-04 | 6568.7ms/step | 3.0 samples/s
+Step 6350/37465 | Loss: 2.6668 | LR: 4.72e-04 | 6563.4ms/step | 3.0 samples/s
+Step 6360/37465 | Loss: 2.2110 | LR: 4.72e-04 | 6560.9ms/step | 3.0 samples/s
+Step 6370/37465 | Loss: 2.0913 | LR: 4.72e-04 | 6562.4ms/step | 3.0 samples/s
+Step 6380/37465 | Loss: 2.4919 | LR: 4.72e-04 | 6566.5ms/step | 3.0 samples/s
+Step 6390/37465 | Loss: 2.5078 | LR: 4.71e-04 | 6567.1ms/step | 3.0 samples/s
+Step 6400/37465 | Loss: 2.3348 | LR: 4.71e-04 | 6566.9ms/step | 3.0 samples/s
+Step 6410/37465 | Loss: 2.3876 | LR: 4.71e-04 | 6569.1ms/step | 3.0 samples/s
+Step 6420/37465 | Loss: 2.2263 | LR: 4.71e-04 | 6566.6ms/step | 3.0 samples/s
+Step 6430/37465 | Loss: 2.2657 | LR: 4.71e-04 | 6567.2ms/step | 3.0 samples/s
+Step 6440/37465 | Loss: 2.2086 | LR: 4.71e-04 | 6565.4ms/step | 3.0 samples/s
+Step 6450/37465 | Loss: 2.3352 | LR: 4.71e-04 | 6559.1ms/step | 3.0 samples/s
+Step 6460/37465 | Loss: 2.3505 | LR: 4.71e-04 | 6566.2ms/step | 3.0 samples/s
+Step 6470/37465 | Loss: 2.3278 | LR: 4.71e-04 | 6564.2ms/step | 3.0 samples/s
+Step 6480/37465 | Loss: 2.2884 | LR: 4.71e-04 | 6559.8ms/step | 3.0 samples/s
+Step 6490/37465 | Loss: 2.3285 | LR: 4.70e-04 | 6561.9ms/step | 3.0 samples/s
+Step 6500/37465 | Loss: 2.5061 | LR: 4.70e-04 | 6563.8ms/step | 3.0 samples/s
+Validation at step 6500: Loss = 2.3665, Perplexity = 10.66
+Step 6510/37465 | Loss: 2.3768 | LR: 4.70e-04 | 7520.4ms/step | 2.7 samples/s
+Step 6520/37465 | Loss: 2.4297 | LR: 4.70e-04 | 6562.4ms/step | 3.0 samples/s
+Step 6530/37465 | Loss: 2.2470 | LR: 4.70e-04 | 6563.6ms/step | 3.0 samples/s
+Step 6540/37465 | Loss: 2.3848 | LR: 4.70e-04 | 6560.3ms/step | 3.0 samples/s
+Step 6550/37465 | Loss: 2.3375 | LR: 4.70e-04 | 6556.7ms/step | 3.1 samples/s
+Step 6560/37465 | Loss: 2.4665 | LR: 4.70e-04 | 6546.3ms/step | 3.1 samples/s
+Step 6570/37465 | Loss: 2.5655 | LR: 4.70e-04 | 6547.4ms/step | 3.1 samples/s
+Step 6580/37465 | Loss: 2.2351 | LR: 4.70e-04 | 6549.4ms/step | 3.1 samples/s
+Step 6590/37465 | Loss: 2.2198 | LR: 4.69e-04 | 6561.6ms/step | 3.0 samples/s
+Step 6600/37465 | Loss: 2.2341 | LR: 4.69e-04 | 6563.1ms/step | 3.0 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_dense/run_20260708_065119/checkpoint-6600
+Step 6610/37465 | Loss: 2.3424 | LR: 4.69e-04 | 7495.2ms/step | 2.7 samples/s
+Step 6620/37465 | Loss: 2.5599 | LR: 4.69e-04 | 6529.5ms/step | 3.1 samples/s
+Step 6630/37465 | Loss: 2.2833 | LR: 4.69e-04 | 6537.4ms/step | 3.1 samples/s
+Step 6640/37465 | Loss: 2.3685 | LR: 4.69e-04 | 6544.1ms/step | 3.1 samples/s
+Step 6650/37465 | Loss: 2.5342 | LR: 4.69e-04 | 6545.4ms/step | 3.1 samples/s
+Step 6660/37465 | Loss: 2.4262 | LR: 4.69e-04 | 6543.5ms/step | 3.1 samples/s
+Step 6670/37465 | Loss: 2.2469 | LR: 4.69e-04 | 6542.1ms/step | 3.1 samples/s
+Step 6680/37465 | Loss: 2.3957 | LR: 4.69e-04 | 6543.4ms/step | 3.1 samples/s
+Step 6690/37465 | Loss: 2.5835 | LR: 4.68e-04 | 6550.6ms/step | 3.1 samples/s
+Step 6700/37465 | Loss: 2.1616 | LR: 4.68e-04 | 6553.4ms/step | 3.1 samples/s
+Step 6710/37465 | Loss: 2.5178 | LR: 4.68e-04 | 6551.7ms/step | 3.1 samples/s
+Step 6720/37465 | Loss: 2.2121 | LR: 4.68e-04 | 6546.7ms/step | 3.1 samples/s
+Step 6730/37465 | Loss: 2.4016 | LR: 4.68e-04 | 6552.8ms/step | 3.1 samples/s
+Step 6740/37465 | Loss: 2.2324 | LR: 4.68e-04 | 6556.1ms/step | 3.1 samples/s
+Step 6750/37465 | Loss: 2.2727 | LR: 4.68e-04 | 6556.2ms/step | 3.1 samples/s
+Step 6760/37465 | Loss: 2.4407 | LR: 4.68e-04 | 6556.1ms/step | 3.1 samples/s
+Step 6770/37465 | Loss: 2.2672 | LR: 4.68e-04 | 6552.9ms/step | 3.1 samples/s
+Step 6780/37465 | Loss: 2.5406 | LR: 4.67e-04 | 6554.7ms/step | 3.1 samples/s
+Step 6790/37465 | Loss: 2.4610 | LR: 4.67e-04 | 6554.1ms/step | 3.1 samples/s
+Step 6800/37465 | Loss: 2.2732 | LR: 4.67e-04 | 6548.0ms/step | 3.1 samples/s
+Step 6810/37465 | Loss: 2.5056 | LR: 4.67e-04 | 6550.0ms/step | 3.1 samples/s
+Step 6820/37465 | Loss: 2.4836 | LR: 4.67e-04 | 6552.2ms/step | 3.1 samples/s
+Step 6830/37465 | Loss: 2.6152 | LR: 4.67e-04 | 6545.7ms/step | 3.1 samples/s
+Step 6840/37465 | Loss: 2.3958 | LR: 4.67e-04 | 6553.9ms/step | 3.1 samples/s
+Step 6850/37465 | Loss: 2.3590 | LR: 4.67e-04 | 6548.3ms/step | 3.1 samples/s
+Step 6860/37465 | Loss: 2.3221 | LR: 4.67e-04 | 6544.1ms/step | 3.1 samples/s
+Step 6870/37465 | Loss: 2.4467 | LR: 4.66e-04 | 6540.9ms/step | 3.1 samples/s
+Step 6880/37465 | Loss: 2.3592 | LR: 4.66e-04 | 6544.6ms/step | 3.1 samples/s
+Step 6890/37465 | Loss: 2.5089 | LR: 4.66e-04 | 6545.0ms/step | 3.1 samples/s
+Step 6900/37465 | Loss: 2.3767 | LR: 4.66e-04 | 6547.9ms/step | 3.1 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_dense/run_20260708_065119/checkpoint-6900
+Step 6910/37465 | Loss: 2.4214 | LR: 4.66e-04 | 7432.7ms/step | 2.7 samples/s
+Step 6920/37465 | Loss: 2.3828 | LR: 4.66e-04 | 6536.1ms/step | 3.1 samples/s
+Step 6930/37465 | Loss: 2.2631 | LR: 4.66e-04 | 6539.2ms/step | 3.1 samples/s
+Step 6940/37465 | Loss: 2.0882 | LR: 4.66e-04 | 6548.1ms/step | 3.1 samples/s
+Step 6950/37465 | Loss: 2.4975 | LR: 4.66e-04 | 6549.3ms/step | 3.1 samples/s
+Step 6960/37465 | Loss: 2.4436 | LR: 4.66e-04 | 6552.7ms/step | 3.1 samples/s
+Step 6970/37465 | Loss: 2.4489 | LR: 4.65e-04 | 6551.9ms/step | 3.1 samples/s
+Step 6980/37465 | Loss: 2.4082 | LR: 4.65e-04 | 6555.1ms/step | 3.1 samples/s
+Step 6990/37465 | Loss: 2.4742 | LR: 4.65e-04 | 6541.0ms/step | 3.1 samples/s
+Step 7000/37465 | Loss: 2.1327 | LR: 4.65e-04 | 6551.0ms/step | 3.1 samples/s
+Validation at step 7000: Loss = 2.3494, Perplexity = 10.48
+Step 7010/37465 | Loss: 2.2668 | LR: 4.65e-04 | 7504.4ms/step | 2.7 samples/s
+Step 7020/37465 | Loss: 2.2757 | LR: 4.65e-04 | 6553.7ms/step | 3.1 samples/s
+Step 7030/37465 | Loss: 2.3008 | LR: 4.65e-04 | 6549.3ms/step | 3.1 samples/s
+Step 7040/37465 | Loss: 2.1730 | LR: 4.65e-04 | 6543.4ms/step | 3.1 samples/s
+Step 7050/37465 | Loss: 2.2929 | LR: 4.65e-04 | 6547.8ms/step | 3.1 samples/s
+Step 7060/37465 | Loss: 2.3877 | LR: 4.64e-04 | 6553.3ms/step | 3.1 samples/s
+Step 7070/37465 | Loss: 2.2295 | LR: 4.64e-04 | 6548.5ms/step | 3.1 samples/s
+Step 7080/37465 | Loss: 2.4255 | LR: 4.64e-04 | 6551.8ms/step | 3.1 samples/s
+Step 7090/37465 | Loss: 2.2283 | LR: 4.64e-04 | 6549.5ms/step | 3.1 samples/s
+Step 7100/37465 | Loss: 2.3275 | LR: 4.64e-04 | 6551.7ms/step | 3.1 samples/s
+Step 7110/37465 | Loss: 2.1359 | LR: 4.64e-04 | 6544.7ms/step | 3.1 samples/s
+Step 7120/37465 | Loss: 1.8726 | LR: 4.64e-04 | 6549.3ms/step | 3.1 samples/s
+Step 7130/37465 | Loss: 2.3255 | LR: 4.64e-04 | 6547.3ms/step | 3.1 samples/s
+Step 7140/37465 | Loss: 2.4325 | LR: 4.64e-04 | 6549.0ms/step | 3.1 samples/s
+Step 7150/37465 | Loss: 2.5015 | LR: 4.63e-04 | 6549.3ms/step | 3.1 samples/s
+Step 7160/37465 | Loss: 2.4244 | LR: 4.63e-04 | 6540.9ms/step | 3.1 samples/s
+Step 7170/37465 | Loss: 2.3303 | LR: 4.63e-04 | 6551.2ms/step | 3.1 samples/s
+Step 7180/37465 | Loss: 2.5237 | LR: 4.63e-04 | 6544.2ms/step | 3.1 samples/s
+Step 7190/37465 | Loss: 2.5751 | LR: 4.63e-04 | 6548.8ms/step | 3.1 samples/s
+Step 7200/37465 | Loss: 2.3200 | LR: 4.63e-04 | 6546.0ms/step | 3.1 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_dense/run_20260708_065119/checkpoint-7200
+Step 7210/37465 | Loss: 2.6006 | LR: 4.63e-04 | 7477.4ms/step | 2.7 samples/s
+Step 7220/37465 | Loss: 2.3671 | LR: 4.63e-04 | 6553.6ms/step | 3.1 samples/s
+Step 7230/37465 | Loss: 2.4178 | LR: 4.63e-04 | 6550.5ms/step | 3.1 samples/s
+Step 7240/37465 | Loss: 2.3708 | LR: 4.62e-04 | 6560.0ms/step | 3.0 samples/s
+Step 7250/37465 | Loss: 2.3538 | LR: 4.62e-04 | 6561.3ms/step | 3.0 samples/s
+Step 7260/37465 | Loss: 2.3715 | LR: 4.62e-04 | 6563.3ms/step | 3.0 samples/s
+Step 7270/37465 | Loss: 2.2095 | LR: 4.62e-04 | 6564.2ms/step | 3.0 samples/s
+Step 7280/37465 | Loss: 2.4128 | LR: 4.62e-04 | 6563.6ms/step | 3.0 samples/s
+Step 7290/37465 | Loss: 2.6005 | LR: 4.62e-04 | 6563.5ms/step | 3.0 samples/s
+Step 7300/37465 | Loss: 2.3187 | LR: 4.62e-04 | 6563.5ms/step | 3.0 samples/s
+Step 7310/37465 | Loss: 2.2895 | LR: 4.62e-04 | 6562.5ms/step | 3.0 samples/s
+Step 7320/37465 | Loss: 2.2645 | LR: 4.62e-04 | 6562.1ms/step | 3.0 samples/s
+Step 7330/37465 | Loss: 2.4428 | LR: 4.61e-04 | 6562.4ms/step | 3.0 samples/s
+Step 7340/37465 | Loss: 2.2917 | LR: 4.61e-04 | 6562.1ms/step | 3.0 samples/s
+Step 7350/37465 | Loss: 2.4093 | LR: 4.61e-04 | 6561.4ms/step | 3.0 samples/s
+Step 7360/37465 | Loss: 2.4316 | LR: 4.61e-04 | 6561.8ms/step | 3.0 samples/s
+Step 7370/37465 | Loss: 2.0862 | LR: 4.61e-04 | 6561.9ms/step | 3.0 samples/s
+Step 7380/37465 | Loss: 2.3805 | LR: 4.61e-04 | 6560.8ms/step | 3.0 samples/s
+Step 7390/37465 | Loss: 2.3124 | LR: 4.61e-04 | 6559.0ms/step | 3.0 samples/s
+Step 7400/37465 | Loss: 2.3205 | LR: 4.61e-04 | 6562.1ms/step | 3.0 samples/s
+Step 7410/37465 | Loss: 2.1981 | LR: 4.60e-04 | 6559.8ms/step | 3.0 samples/s
+Step 7420/37465 | Loss: 2.4750 | LR: 4.60e-04 | 6560.1ms/step | 3.0 samples/s
+Step 7430/37465 | Loss: 2.4320 | LR: 4.60e-04 | 6561.9ms/step | 3.0 samples/s
+Step 7440/37465 | Loss: 2.1724 | LR: 4.60e-04 | 6561.2ms/step | 3.0 samples/s
+Step 7450/37465 | Loss: 2.2110 | LR: 4.60e-04 | 6562.9ms/step | 3.0 samples/s
+Step 7460/37465 | Loss: 2.3355 | LR: 4.60e-04 | 6562.3ms/step | 3.0 samples/s
+Step 7470/37465 | Loss: 2.1773 | LR: 4.60e-04 | 6563.5ms/step | 3.0 samples/s
+Step 7480/37465 | Loss: 2.4105 | LR: 4.60e-04 | 6559.2ms/step | 3.0 samples/s
+Step 7490/37465 | Loss: 2.4115 | LR: 4.60e-04 | 6563.8ms/step | 3.0 samples/s
+Step 7500/37465 | Loss: 2.2460 | LR: 4.59e-04 | 6562.6ms/step | 3.0 samples/s
+Validation at step 7500: Loss = 2.3312, Perplexity = 10.29
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_dense/run_20260708_065119/checkpoint-7500
+Step 7510/37465 | Loss: 2.3408 | LR: 4.59e-04 | 8379.8ms/step | 2.4 samples/s
+Step 7520/37465 | Loss: 2.3837 | LR: 4.59e-04 | 6547.2ms/step | 3.1 samples/s
+Step 7530/37465 | Loss: 2.4462 | LR: 4.59e-04 | 6557.8ms/step | 3.0 samples/s
+Step 7540/37465 | Loss: 2.2093 | LR: 4.59e-04 | 6560.5ms/step | 3.0 samples/s
+Step 7550/37465 | Loss: 2.4724 | LR: 4.59e-04 | 6558.9ms/step | 3.0 samples/s
+Step 7560/37465 | Loss: 2.3339 | LR: 4.59e-04 | 6560.4ms/step | 3.0 samples/s
+Step 7570/37465 | Loss: 2.3446 | LR: 4.59e-04 | 6557.8ms/step | 3.0 samples/s
+Step 7580/37465 | Loss: 2.4865 | LR: 4.59e-04 | 6561.1ms/step | 3.0 samples/s
+Step 7590/37465 | Loss: 2.3561 | LR: 4.58e-04 | 6561.2ms/step | 3.0 samples/s
+Step 7600/37465 | Loss: 2.4925 | LR: 4.58e-04 | 6560.8ms/step | 3.0 samples/s
+Step 7610/37465 | Loss: 2.4419 | LR: 4.58e-04 | 6563.1ms/step | 3.0 samples/s
+Step 7620/37465 | Loss: 2.1357 | LR: 4.58e-04 | 6561.3ms/step | 3.0 samples/s
+Step 7630/37465 | Loss: 2.3136 | LR: 4.58e-04 | 6564.3ms/step | 3.0 samples/s
+Step 7640/37465 | Loss: 2.2075 | LR: 4.58e-04 | 6559.0ms/step | 3.0 samples/s
+Step 7650/37465 | Loss: 2.2752 | LR: 4.58e-04 | 6561.5ms/step | 3.0 samples/s
+Step 7660/37465 | Loss: 2.3247 | LR: 4.58e-04 | 6564.2ms/step | 3.0 samples/s
+Step 7670/37465 | Loss: 2.2081 | LR: 4.57e-04 | 6563.0ms/step | 3.0 samples/s
+Step 7680/37465 | Loss: 2.2805 | LR: 4.57e-04 | 6561.9ms/step | 3.0 samples/s
+Step 7690/37465 | Loss: 2.1501 | LR: 4.57e-04 | 6560.1ms/step | 3.0 samples/s
+Step 7700/37465 | Loss: 2.2059 | LR: 4.57e-04 | 6557.7ms/step | 3.0 samples/s
+Step 7710/37465 | Loss: 2.5812 | LR: 4.57e-04 | 6550.8ms/step | 3.1 samples/s
+Step 7720/37465 | Loss: 2.1455 | LR: 4.57e-04 | 6547.1ms/step | 3.1 samples/s
+Step 7730/37465 | Loss: 2.3490 | LR: 4.57e-04 | 6545.7ms/step | 3.1 samples/s
+Step 7740/37465 | Loss: 2.4469 | LR: 4.57e-04 | 6558.7ms/step | 3.0 samples/s
+Step 7750/37465 | Loss: 2.4455 | LR: 4.56e-04 | 6559.1ms/step | 3.0 samples/s
+Step 7760/37465 | Loss: 2.3704 | LR: 4.56e-04 | 6549.6ms/step | 3.1 samples/s
+Step 7770/37465 | Loss: 2.1382 | LR: 4.56e-04 | 6549.5ms/step | 3.1 samples/s
+Step 7780/37465 | Loss: 2.1656 | LR: 4.56e-04 | 6547.9ms/step | 3.1 samples/s
+Step 7790/37465 | Loss: 2.1569 | LR: 4.56e-04 | 6553.0ms/step | 3.1 samples/s
+Step 7800/37465 | Loss: 2.1258 | LR: 4.56e-04 | 6578.3ms/step | 3.0 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_dense/run_20260708_065119/checkpoint-7800
+Step 7810/37465 | Loss: 2.3208 | LR: 4.56e-04 | 7514.6ms/step | 2.7 samples/s
+Step 7820/37465 | Loss: 2.2984 | LR: 4.56e-04 | 6552.2ms/step | 3.1 samples/s
+Step 7830/37465 | Loss: 2.2539 | LR: 4.56e-04 | 6560.4ms/step | 3.0 samples/s
+Step 7840/37465 | Loss: 2.4441 | LR: 4.55e-04 | 6561.7ms/step | 3.0 samples/s
+Step 7850/37465 | Loss: 2.2327 | LR: 4.55e-04 | 6557.9ms/step | 3.0 samples/s
+Step 7860/37465 | Loss: 2.2017 | LR: 4.55e-04 | 6562.4ms/step | 3.0 samples/s
+Step 7870/37465 | Loss: 2.3336 | LR: 4.55e-04 | 6557.0ms/step | 3.1 samples/s
+Step 7880/37465 | Loss: 2.2059 | LR: 4.55e-04 | 6558.7ms/step | 3.0 samples/s
+Step 7890/37465 | Loss: 2.4152 | LR: 4.55e-04 | 6558.1ms/step | 3.0 samples/s
+Step 7900/37465 | Loss: 2.2202 | LR: 4.55e-04 | 6559.8ms/step | 3.0 samples/s
+Step 7910/37465 | Loss: 2.1934 | LR: 4.55e-04 | 6559.1ms/step | 3.0 samples/s
+Step 7920/37465 | Loss: 2.1396 | LR: 4.54e-04 | 6555.9ms/step | 3.1 samples/s
+Step 7930/37465 | Loss: 2.2589 | LR: 4.54e-04 | 6555.0ms/step | 3.1 samples/s
+Step 7940/37465 | Loss: 2.1211 | LR: 4.54e-04 | 6554.8ms/step | 3.1 samples/s
+Step 7950/37465 | Loss: 2.2761 | LR: 4.54e-04 | 6555.3ms/step | 3.1 samples/s
+Step 7960/37465 | Loss: 2.4197 | LR: 4.54e-04 | 6554.6ms/step | 3.1 samples/s
+Step 7970/37465 | Loss: 2.3031 | LR: 4.54e-04 | 6556.1ms/step | 3.1 samples/s
+Step 7980/37465 | Loss: 2.2820 | LR: 4.54e-04 | 6559.1ms/step | 3.0 samples/s
+Step 7990/37465 | Loss: 2.4216 | LR: 4.54e-04 | 6558.9ms/step | 3.0 samples/s
+Step 8000/37465 | Loss: 2.2923 | LR: 4.53e-04 | 6559.1ms/step | 3.0 samples/s
+Validation at step 8000: Loss = 2.3092, Perplexity = 10.07
+Step 8010/37465 | Loss: 2.0611 | LR: 4.53e-04 | 7511.0ms/step | 2.7 samples/s
+Step 8020/37465 | Loss: 2.3299 | LR: 4.53e-04 | 6558.6ms/step | 3.0 samples/s
+Step 8030/37465 | Loss: 2.3175 | LR: 4.53e-04 | 6559.0ms/step | 3.0 samples/s
+Step 8040/37465 | Loss: 2.2853 | LR: 4.53e-04 | 6563.2ms/step | 3.0 samples/s
+Step 8050/37465 | Loss: 2.1186 | LR: 4.53e-04 | 6562.8ms/step | 3.0 samples/s
+Step 8060/37465 | Loss: 2.4040 | LR: 4.53e-04 | 6560.7ms/step | 3.0 samples/s
+Step 8070/37465 | Loss: 2.3971 | LR: 4.53e-04 | 6558.0ms/step | 3.0 samples/s
+Step 8080/37465 | Loss: 2.3314 | LR: 4.52e-04 | 6555.5ms/step | 3.1 samples/s
+Step 8090/37465 | Loss: 2.1851 | LR: 4.52e-04 | 6560.9ms/step | 3.0 samples/s
+Step 8100/37465 | Loss: 2.3967 | LR: 4.52e-04 | 6559.4ms/step | 3.0 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_dense/run_20260708_065119/checkpoint-8100
+Step 8110/37465 | Loss: 2.3850 | LR: 4.52e-04 | 7492.5ms/step | 2.7 samples/s
+Step 8120/37465 | Loss: 2.3767 | LR: 4.52e-04 | 6548.9ms/step | 3.1 samples/s
+Step 8130/37465 | Loss: 2.3291 | LR: 4.52e-04 | 6557.1ms/step | 3.1 samples/s
+Step 8140/37465 | Loss: 2.2218 | LR: 4.52e-04 | 6558.2ms/step | 3.0 samples/s
+Step 8150/37465 | Loss: 2.4259 | LR: 4.52e-04 | 6559.1ms/step | 3.0 samples/s
+Step 8160/37465 | Loss: 2.2300 | LR: 4.51e-04 | 6558.9ms/step | 3.0 samples/s
+Step 8170/37465 | Loss: 2.3107 | LR: 4.51e-04 | 6556.8ms/step | 3.1 samples/s
+Step 8180/37465 | Loss: 2.3898 | LR: 4.51e-04 | 6558.0ms/step | 3.0 samples/s
+Step 8190/37465 | Loss: 2.2195 | LR: 4.51e-04 | 6558.0ms/step | 3.0 samples/s
+Step 8200/37465 | Loss: 2.3088 | LR: 4.51e-04 | 6561.0ms/step | 3.0 samples/s
+Step 8210/37465 | Loss: 2.4433 | LR: 4.51e-04 | 6558.1ms/step | 3.0 samples/s
+Step 8220/37465 | Loss: 2.3324 | LR: 4.51e-04 | 6558.0ms/step | 3.0 samples/s
+Step 8230/37465 | Loss: 2.2587 | LR: 4.51e-04 | 6560.9ms/step | 3.0 samples/s
+Step 8240/37465 | Loss: 2.2621 | LR: 4.50e-04 | 6561.1ms/step | 3.0 samples/s
+Step 8250/37465 | Loss: 2.2610 | LR: 4.50e-04 | 6559.2ms/step | 3.0 samples/s
+Step 8260/37465 | Loss: 2.5453 | LR: 4.50e-04 | 6554.8ms/step | 3.1 samples/s
+Step 8270/37465 | Loss: 2.3990 | LR: 4.50e-04 | 6558.0ms/step | 3.0 samples/s
+Step 8280/37465 | Loss: 2.1385 | LR: 4.50e-04 | 6557.9ms/step | 3.0 samples/s
+Step 8290/37465 | Loss: 2.3973 | LR: 4.50e-04 | 6559.1ms/step | 3.0 samples/s
+Step 8300/37465 | Loss: 2.2114 | LR: 4.50e-04 | 6557.8ms/step | 3.0 samples/s
+Step 8310/37465 | Loss: 2.3886 | LR: 4.49e-04 | 6560.0ms/step | 3.0 samples/s
+Step 8320/37465 | Loss: 2.4536 | LR: 4.49e-04 | 6562.2ms/step | 3.0 samples/s
+Step 8330/37465 | Loss: 2.0917 | LR: 4.49e-04 | 6561.8ms/step | 3.0 samples/s
+Step 8340/37465 | Loss: 2.4583 | LR: 4.49e-04 | 6559.9ms/step | 3.0 samples/s
+Step 8350/37465 | Loss: 2.3892 | LR: 4.49e-04 | 6559.1ms/step | 3.0 samples/s
+Step 8360/37465 | Loss: 2.0192 | LR: 4.49e-04 | 6561.4ms/step | 3.0 samples/s
+Step 8370/37465 | Loss: 2.3429 | LR: 4.49e-04 | 6560.7ms/step | 3.0 samples/s
+Step 8380/37465 | Loss: 2.1619 | LR: 4.49e-04 | 6558.1ms/step | 3.0 samples/s
+Step 8390/37465 | Loss: 2.2445 | LR: 4.48e-04 | 6557.5ms/step | 3.0 samples/s
+Step 8400/37465 | Loss: 2.1991 | LR: 4.48e-04 | 6558.9ms/step | 3.0 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_dense/run_20260708_065119/checkpoint-8400
+Step 8410/37465 | Loss: 2.3588 | LR: 4.48e-04 | 7498.1ms/step | 2.7 samples/s
+Step 8420/37465 | Loss: 2.2198 | LR: 4.48e-04 | 6550.2ms/step | 3.1 samples/s
+Step 8430/37465 | Loss: 2.2904 | LR: 4.48e-04 | 6560.3ms/step | 3.0 samples/s
+Step 8440/37465 | Loss: 2.2391 | LR: 4.48e-04 | 6559.5ms/step | 3.0 samples/s
+Step 8450/37465 | Loss: 2.1156 | LR: 4.48e-04 | 6565.7ms/step | 3.0 samples/s
+Step 8460/37465 | Loss: 2.1749 | LR: 4.48e-04 | 6560.3ms/step | 3.0 samples/s
+Step 8470/37465 | Loss: 2.2245 | LR: 4.47e-04 | 6563.1ms/step | 3.0 samples/s
+Step 8480/37465 | Loss: 2.2289 | LR: 4.47e-04 | 6561.2ms/step | 3.0 samples/s
+Step 8490/37465 | Loss: 2.3633 | LR: 4.47e-04 | 6561.1ms/step | 3.0 samples/s
+Step 8500/37465 | Loss: 2.2933 | LR: 4.47e-04 | 6562.6ms/step | 3.0 samples/s
+Validation at step 8500: Loss = 2.2973, Perplexity = 9.95
+Step 8510/37465 | Loss: 2.2235 | LR: 4.47e-04 | 7522.2ms/step | 2.7 samples/s
+Step 8520/37465 | Loss: 2.2465 | LR: 4.47e-04 | 6563.8ms/step | 3.0 samples/s
+Step 8530/37465 | Loss: 2.1272 | LR: 4.47e-04 | 6559.6ms/step | 3.0 samples/s
+Step 8540/37465 | Loss: 2.2456 | LR: 4.46e-04 | 6568.0ms/step | 3.0 samples/s
+Step 8550/37465 | Loss: 2.2713 | LR: 4.46e-04 | 6559.3ms/step | 3.0 samples/s
+Step 8560/37465 | Loss: 2.5739 | LR: 4.46e-04 | 6564.2ms/step | 3.0 samples/s
+Step 8570/37465 | Loss: 2.2546 | LR: 4.46e-04 | 6562.4ms/step | 3.0 samples/s
+Step 8580/37465 | Loss: 2.2964 | LR: 4.46e-04 | 6560.7ms/step | 3.0 samples/s
+Step 8590/37465 | Loss: 2.1473 | LR: 4.46e-04 | 6550.9ms/step | 3.1 samples/s
+Step 8600/37465 | Loss: 2.3553 | LR: 4.46e-04 | 6564.6ms/step | 3.0 samples/s
+Step 8610/37465 | Loss: 2.5808 | LR: 4.46e-04 | 6558.7ms/step | 3.0 samples/s
+Step 8620/37465 | Loss: 2.1172 | LR: 4.45e-04 | 6554.4ms/step | 3.1 samples/s

epfl_final/logs/l30k1_2b_fixed.log

@@ -0,0 +1,343 @@
+W0709 00:04:52.727000 128 site-packages/torch/distributed/run.py:851] 
+W0709 00:04:52.727000 128 site-packages/torch/distributed/run.py:851] *****************************************
+W0709 00:04:52.727000 128 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 00:04:52.727000 128 site-packages/torch/distributed/run.py:851] *****************************************
+New run: Checkpoints will be saved to /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": null,
+  "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...
+Initializing model from scratch (random weights)...
+Checkpoint activations: True
+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: 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...
+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
+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...
+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)
+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_000650-9blizq3m
+wandb: Run `wandb offline` to turn off syncing.
+wandb: Syncing 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...
+Switching attention mode for step 1: sparse_stochastic (configured sparse attention)
+Step 10/37465 | Loss: 12.2539 | LR: 6.67e-06 | 16051.2ms/step | 1.2 samples/s
+Step 20/37465 | Loss: 12.1179 | LR: 1.33e-05 | 8082.2ms/step | 2.5 samples/s
+Step 30/37465 | Loss: 11.6617 | LR: 2.00e-05 | 8081.1ms/step | 2.5 samples/s
+Step 40/37465 | Loss: 10.6154 | LR: 2.67e-05 | 8080.6ms/step | 2.5 samples/s
+Step 50/37465 | Loss: 10.0327 | LR: 3.33e-05 | 8091.3ms/step | 2.5 samples/s
+Step 60/37465 | Loss: 9.5505 | LR: 4.00e-05 | 8065.8ms/step | 2.5 samples/s
+Step 70/37465 | Loss: 9.0958 | LR: 4.67e-05 | 8053.5ms/step | 2.5 samples/s
+Step 80/37465 | Loss: 8.6200 | LR: 5.33e-05 | 8049.3ms/step | 2.5 samples/s
+Step 90/37465 | Loss: 8.0455 | LR: 6.00e-05 | 8086.4ms/step | 2.5 samples/s
+Step 100/37465 | Loss: 7.7043 | LR: 6.67e-05 | 8064.0ms/step | 2.5 samples/s
+Step 110/37465 | Loss: 7.3621 | LR: 7.33e-05 | 8057.3ms/step | 2.5 samples/s
+Step 120/37465 | Loss: 7.0990 | LR: 8.00e-05 | 8081.9ms/step | 2.5 samples/s
+Step 130/37465 | Loss: 6.9080 | LR: 8.67e-05 | 8076.1ms/step | 2.5 samples/s
+Step 140/37465 | Loss: 6.7992 | LR: 9.33e-05 | 8100.9ms/step | 2.5 samples/s
+Step 150/37465 | Loss: 6.8092 | LR: 1.00e-04 | 8105.2ms/step | 2.5 samples/s
+Step 160/37465 | Loss: 6.6568 | LR: 1.07e-04 | 8122.5ms/step | 2.5 samples/s
+Step 170/37465 | Loss: 6.3303 | LR: 1.13e-04 | 8124.4ms/step | 2.5 samples/s
+Step 180/37465 | Loss: 6.1705 | LR: 1.20e-04 | 8111.3ms/step | 2.5 samples/s
+Step 190/37465 | Loss: 6.2347 | LR: 1.27e-04 | 8110.7ms/step | 2.5 samples/s
+Step 200/37465 | Loss: 5.9983 | LR: 1.33e-04 | 8095.5ms/step | 2.5 samples/s
+Step 210/37465 | Loss: 5.9896 | LR: 1.40e-04 | 8104.3ms/step | 2.5 samples/s
+Step 220/37465 | Loss: 5.9340 | LR: 1.47e-04 | 8108.6ms/step | 2.5 samples/s
+Step 230/37465 | Loss: 5.8768 | LR: 1.53e-04 | 8127.6ms/step | 2.5 samples/s
+Step 240/37465 | Loss: 5.8841 | LR: 1.60e-04 | 8082.9ms/step | 2.5 samples/s
+Step 250/37465 | Loss: 5.8065 | LR: 1.67e-04 | 8097.2ms/step | 2.5 samples/s
+Step 260/37465 | Loss: 5.6594 | LR: 1.73e-04 | 8141.2ms/step | 2.5 samples/s
+Step 270/37465 | Loss: 5.6777 | LR: 1.80e-04 | 8118.9ms/step | 2.5 samples/s
+Step 280/37465 | Loss: 5.6578 | LR: 1.87e-04 | 8143.6ms/step | 2.5 samples/s
+Step 290/37465 | Loss: 5.4694 | LR: 1.93e-04 | 8120.3ms/step | 2.5 samples/s
+Step 300/37465 | Loss: 5.3341 | LR: 2.00e-04 | 8087.3ms/step | 2.5 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_fixed/run_20260709_000503/checkpoint-300
+Step 310/37465 | Loss: 5.2074 | LR: 2.07e-04 | 8974.8ms/step | 2.2 samples/s
+Step 320/37465 | Loss: 5.3182 | LR: 2.13e-04 | 8119.7ms/step | 2.5 samples/s
+Step 330/37465 | Loss: 5.5086 | LR: 2.20e-04 | 8092.3ms/step | 2.5 samples/s
+Step 340/37465 | Loss: 5.3333 | LR: 2.27e-04 | 8092.1ms/step | 2.5 samples/s
+Step 350/37465 | Loss: 5.1660 | LR: 2.33e-04 | 8092.8ms/step | 2.5 samples/s
+Step 360/37465 | Loss: 5.2507 | LR: 2.40e-04 | 8090.9ms/step | 2.5 samples/s
+Step 370/37465 | Loss: 5.0838 | LR: 2.47e-04 | 8098.4ms/step | 2.5 samples/s
+Step 380/37465 | Loss: 5.3012 | LR: 2.53e-04 | 8117.0ms/step | 2.5 samples/s
+Step 390/37465 | Loss: 5.2305 | LR: 2.60e-04 | 8124.3ms/step | 2.5 samples/s
+Step 400/37465 | Loss: 5.2236 | LR: 2.67e-04 | 8104.6ms/step | 2.5 samples/s
+Step 410/37465 | Loss: 5.0358 | LR: 2.73e-04 | 8094.3ms/step | 2.5 samples/s
+Step 420/37465 | Loss: 4.8083 | LR: 2.80e-04 | 8099.8ms/step | 2.5 samples/s
+Step 430/37465 | Loss: 4.9145 | LR: 2.87e-04 | 8130.3ms/step | 2.5 samples/s
+Step 440/37465 | Loss: 4.8570 | LR: 2.93e-04 | 8099.9ms/step | 2.5 samples/s
+Step 450/37465 | Loss: 4.8512 | LR: 3.00e-04 | 8108.0ms/step | 2.5 samples/s
+Step 460/37465 | Loss: 4.6672 | LR: 3.07e-04 | 8100.2ms/step | 2.5 samples/s
+Step 470/37465 | Loss: 4.7526 | LR: 3.13e-04 | 8116.4ms/step | 2.5 samples/s
+Step 480/37465 | Loss: 4.8818 | LR: 3.20e-04 | 8105.0ms/step | 2.5 samples/s
+Step 490/37465 | Loss: 4.5789 | LR: 3.27e-04 | 8094.9ms/step | 2.5 samples/s
+Step 500/37465 | Loss: 4.8251 | LR: 3.33e-04 | 8097.8ms/step | 2.5 samples/s
+Validation at step 500: Loss = 4.5909, Perplexity = 98.58
+Step 510/37465 | Loss: 4.4435 | LR: 3.40e-04 | 12250.0ms/step | 1.6 samples/s
+Step 520/37465 | Loss: 4.6284 | LR: 3.47e-04 | 8154.5ms/step | 2.5 samples/s
+Step 530/37465 | Loss: 4.7687 | LR: 3.53e-04 | 8139.5ms/step | 2.5 samples/s
+Step 540/37465 | Loss: 4.9029 | LR: 3.60e-04 | 8135.7ms/step | 2.5 samples/s
+Step 550/37465 | Loss: 4.3930 | LR: 3.67e-04 | 8118.1ms/step | 2.5 samples/s
+Step 560/37465 | Loss: 4.4000 | LR: 3.73e-04 | 8118.4ms/step | 2.5 samples/s
+Step 570/37465 | Loss: 4.3734 | LR: 3.80e-04 | 8126.1ms/step | 2.5 samples/s
+Step 580/37465 | Loss: 4.0479 | LR: 3.87e-04 | 8117.9ms/step | 2.5 samples/s
+Step 590/37465 | Loss: 4.4079 | LR: 3.93e-04 | 8092.1ms/step | 2.5 samples/s
+Step 600/37465 | Loss: 4.5036 | LR: 4.00e-04 | 8142.3ms/step | 2.5 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_fixed/run_20260709_000503/checkpoint-600
+Step 610/37465 | Loss: 4.2265 | LR: 4.07e-04 | 8970.4ms/step | 2.2 samples/s
+Step 620/37465 | Loss: 4.0874 | LR: 4.13e-04 | 8141.6ms/step | 2.5 samples/s
+Step 630/37465 | Loss: 4.1582 | LR: 4.20e-04 | 8190.3ms/step | 2.4 samples/s
+Step 640/37465 | Loss: 4.2591 | LR: 4.27e-04 | 8199.7ms/step | 2.4 samples/s
+Step 650/37465 | Loss: 4.1802 | LR: 4.33e-04 | 8135.7ms/step | 2.5 samples/s
+Step 660/37465 | Loss: 4.2430 | LR: 4.40e-04 | 8193.1ms/step | 2.4 samples/s
+Step 670/37465 | Loss: 4.3667 | LR: 4.47e-04 | 8171.4ms/step | 2.4 samples/s
+Step 680/37465 | Loss: 4.1284 | LR: 4.53e-04 | 8108.1ms/step | 2.5 samples/s
+Step 690/37465 | Loss: 3.9607 | LR: 4.60e-04 | 8130.9ms/step | 2.5 samples/s
+Step 700/37465 | Loss: 4.0995 | LR: 4.67e-04 | 8119.4ms/step | 2.5 samples/s
+Step 710/37465 | Loss: 4.0221 | LR: 4.73e-04 | 8139.8ms/step | 2.5 samples/s
+Step 720/37465 | Loss: 3.9483 | LR: 4.80e-04 | 8150.2ms/step | 2.5 samples/s
+Step 730/37465 | Loss: 3.5534 | LR: 4.87e-04 | 8139.9ms/step | 2.5 samples/s
+Step 740/37465 | Loss: 4.0179 | LR: 4.93e-04 | 8110.7ms/step | 2.5 samples/s
+Step 750/37465 | Loss: 3.7463 | LR: 5.00e-04 | 8171.7ms/step | 2.4 samples/s
+Step 760/37465 | Loss: 3.6854 | LR: 5.00e-04 | 8151.6ms/step | 2.5 samples/s
+Step 770/37465 | Loss: 3.6698 | LR: 5.00e-04 | 8120.2ms/step | 2.5 samples/s
+Step 780/37465 | Loss: 3.7517 | LR: 5.00e-04 | 8120.2ms/step | 2.5 samples/s
+Step 790/37465 | Loss: 3.3773 | LR: 5.00e-04 | 8130.0ms/step | 2.5 samples/s
+Step 800/37465 | Loss: 3.7659 | LR: 5.00e-04 | 8161.8ms/step | 2.5 samples/s
+Step 810/37465 | Loss: 3.7977 | LR: 5.00e-04 | 8139.7ms/step | 2.5 samples/s
+Step 820/37465 | Loss: 3.6062 | LR: 5.00e-04 | 8139.7ms/step | 2.5 samples/s
+Step 830/37465 | Loss: 3.6142 | LR: 5.00e-04 | 8130.6ms/step | 2.5 samples/s
+Step 840/37465 | Loss: 3.7947 | LR: 5.00e-04 | 8130.6ms/step | 2.5 samples/s
+Step 850/37465 | Loss: 3.4291 | LR: 5.00e-04 | 8149.5ms/step | 2.5 samples/s
+Step 860/37465 | Loss: 3.6450 | LR: 5.00e-04 | 8110.0ms/step | 2.5 samples/s
+Step 870/37465 | Loss: 3.5762 | LR: 5.00e-04 | 8140.7ms/step | 2.5 samples/s
+Step 880/37465 | Loss: 3.4716 | LR: 5.00e-04 | 8127.4ms/step | 2.5 samples/s
+Step 890/37465 | Loss: 3.5228 | LR: 5.00e-04 | 8138.8ms/step | 2.5 samples/s
+Step 900/37465 | Loss: 3.2153 | LR: 5.00e-04 | 8181.3ms/step | 2.4 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_fixed/run_20260709_000503/checkpoint-900
+Step 910/37465 | Loss: 3.2410 | LR: 5.00e-04 | 9024.2ms/step | 2.2 samples/s
+Step 920/37465 | Loss: 3.3828 | LR: 5.00e-04 | 8132.0ms/step | 2.5 samples/s
+Step 930/37465 | Loss: 3.2360 | LR: 5.00e-04 | 8155.9ms/step | 2.5 samples/s
+Step 940/37465 | Loss: 3.5005 | LR: 5.00e-04 | 8169.7ms/step | 2.4 samples/s
+Step 950/37465 | Loss: 3.7573 | LR: 5.00e-04 | 8119.2ms/step | 2.5 samples/s
+Step 960/37465 | Loss: 3.0777 | LR: 5.00e-04 | 8131.2ms/step | 2.5 samples/s
+Step 970/37465 | Loss: 3.4576 | LR: 5.00e-04 | 8110.0ms/step | 2.5 samples/s
+Step 980/37465 | Loss: 3.2568 | LR: 5.00e-04 | 8109.8ms/step | 2.5 samples/s
+Step 990/37465 | Loss: 3.2432 | LR: 5.00e-04 | 8100.4ms/step | 2.5 samples/s
+Step 1000/37465 | Loss: 3.0835 | LR: 5.00e-04 | 8110.2ms/step | 2.5 samples/s
+Validation at step 1000: Loss = 3.2898, Perplexity = 26.84
+Step 1010/37465 | Loss: 3.1444 | LR: 5.00e-04 | 8817.9ms/step | 2.3 samples/s
+Step 1020/37465 | Loss: 3.4166 | LR: 5.00e-04 | 8124.0ms/step | 2.5 samples/s
+Step 1030/37465 | Loss: 2.9929 | LR: 5.00e-04 | 8162.4ms/step | 2.5 samples/s
+Step 1040/37465 | Loss: 3.4101 | LR: 5.00e-04 | 8167.8ms/step | 2.4 samples/s
+Step 1050/37465 | Loss: 3.2110 | LR: 5.00e-04 | 8129.6ms/step | 2.5 samples/s
+Step 1060/37465 | Loss: 3.2282 | LR: 5.00e-04 | 8160.7ms/step | 2.5 samples/s
+Step 1070/37465 | Loss: 3.2118 | LR: 5.00e-04 | 8131.9ms/step | 2.5 samples/s
+Step 1080/37465 | Loss: 3.2191 | LR: 5.00e-04 | 8129.6ms/step | 2.5 samples/s
+Step 1090/37465 | Loss: 3.2135 | LR: 5.00e-04 | 8138.3ms/step | 2.5 samples/s
+Step 1100/37465 | Loss: 3.0780 | LR: 5.00e-04 | 8112.0ms/step | 2.5 samples/s
+Step 1110/37465 | Loss: 3.0225 | LR: 5.00e-04 | 8120.4ms/step | 2.5 samples/s
+Step 1120/37465 | Loss: 2.9652 | LR: 5.00e-04 | 8142.1ms/step | 2.5 samples/s
+Step 1130/37465 | Loss: 3.2001 | LR: 5.00e-04 | 8139.7ms/step | 2.5 samples/s
+Step 1140/37465 | Loss: 3.1594 | LR: 5.00e-04 | 8161.0ms/step | 2.5 samples/s
+Step 1150/37465 | Loss: 2.9924 | LR: 5.00e-04 | 8184.9ms/step | 2.4 samples/s
+Step 1160/37465 | Loss: 3.3070 | LR: 5.00e-04 | 8132.1ms/step | 2.5 samples/s
+Step 1170/37465 | Loss: 2.8632 | LR: 5.00e-04 | 8132.0ms/step | 2.5 samples/s
+Step 1180/37465 | Loss: 3.2380 | LR: 5.00e-04 | 8133.9ms/step | 2.5 samples/s
+Step 1190/37465 | Loss: 3.0606 | LR: 5.00e-04 | 8150.3ms/step | 2.5 samples/s
+Step 1200/37465 | Loss: 2.7986 | LR: 5.00e-04 | 8111.2ms/step | 2.5 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_fixed/run_20260709_000503/checkpoint-1200
+Step 1210/37465 | Loss: 3.1805 | LR: 5.00e-04 | 8998.0ms/step | 2.2 samples/s
+Step 1220/37465 | Loss: 3.2459 | LR: 5.00e-04 | 8201.2ms/step | 2.4 samples/s
+Step 1230/37465 | Loss: 2.8701 | LR: 5.00e-04 | 8187.8ms/step | 2.4 samples/s
+Step 1240/37465 | Loss: 3.1206 | LR: 5.00e-04 | 8169.0ms/step | 2.4 samples/s
+Step 1250/37465 | Loss: 2.9399 | LR: 5.00e-04 | 8131.4ms/step | 2.5 samples/s
+Step 1260/37465 | Loss: 2.9571 | LR: 5.00e-04 | 8158.8ms/step | 2.5 samples/s
+Step 1270/37465 | Loss: 3.0485 | LR: 5.00e-04 | 8159.8ms/step | 2.5 samples/s
+Step 1280/37465 | Loss: 2.9158 | LR: 5.00e-04 | 8145.6ms/step | 2.5 samples/s
+Step 1290/37465 | Loss: 3.0296 | LR: 5.00e-04 | 8112.0ms/step | 2.5 samples/s
+Step 1300/37465 | Loss: 3.3326 | LR: 5.00e-04 | 8156.0ms/step | 2.5 samples/s
+Step 1310/37465 | Loss: 3.1125 | LR: 5.00e-04 | 8127.6ms/step | 2.5 samples/s
+Step 1320/37465 | Loss: 2.9565 | LR: 5.00e-04 | 8182.8ms/step | 2.4 samples/s
+Step 1330/37465 | Loss: 3.1525 | LR: 5.00e-04 | 8106.6ms/step | 2.5 samples/s
+Step 1340/37465 | Loss: 3.0781 | LR: 5.00e-04 | 8150.0ms/step | 2.5 samples/s
+Step 1350/37465 | Loss: 3.0035 | LR: 5.00e-04 | 8165.1ms/step | 2.4 samples/s
+Step 1360/37465 | Loss: 2.9374 | LR: 5.00e-04 | 8159.3ms/step | 2.5 samples/s
+Step 1370/37465 | Loss: 3.2396 | LR: 5.00e-04 | 8188.6ms/step | 2.4 samples/s
+Step 1380/37465 | Loss: 2.9515 | LR: 5.00e-04 | 8159.7ms/step | 2.5 samples/s
+Step 1390/37465 | Loss: 2.9292 | LR: 5.00e-04 | 8154.5ms/step | 2.5 samples/s
+Step 1400/37465 | Loss: 2.9672 | LR: 5.00e-04 | 8157.0ms/step | 2.5 samples/s
+Step 1410/37465 | Loss: 3.1546 | LR: 5.00e-04 | 8171.0ms/step | 2.4 samples/s
+Step 1420/37465 | Loss: 2.7728 | LR: 5.00e-04 | 8138.6ms/step | 2.5 samples/s
+Step 1430/37465 | Loss: 3.3204 | LR: 5.00e-04 | 8182.4ms/step | 2.4 samples/s
+Step 1440/37465 | Loss: 2.9759 | LR: 5.00e-04 | 8191.3ms/step | 2.4 samples/s
+Step 1450/37465 | Loss: 3.0437 | LR: 5.00e-04 | 8212.5ms/step | 2.4 samples/s
+Step 1460/37465 | Loss: 3.0972 | LR: 5.00e-04 | 8171.7ms/step | 2.4 samples/s
+Step 1470/37465 | Loss: 2.9715 | LR: 5.00e-04 | 8155.6ms/step | 2.5 samples/s
+Step 1480/37465 | Loss: 2.8303 | LR: 5.00e-04 | 8259.0ms/step | 2.4 samples/s
+Step 1490/37465 | Loss: 2.9437 | LR: 4.99e-04 | 8222.2ms/step | 2.4 samples/s
+Step 1500/37465 | Loss: 2.9037 | LR: 4.99e-04 | 8238.2ms/step | 2.4 samples/s
+Validation at step 1500: Loss = 2.9803, Perplexity = 19.69
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_fixed/run_20260709_000503/checkpoint-1500
+Step 1510/37465 | Loss: 2.9913 | LR: 4.99e-04 | 9668.6ms/step | 2.1 samples/s
+Step 1520/37465 | Loss: 2.9082 | LR: 4.99e-04 | 8174.8ms/step | 2.4 samples/s
+Step 1530/37465 | Loss: 2.8394 | LR: 4.99e-04 | 8172.6ms/step | 2.4 samples/s
+Step 1540/37465 | Loss: 2.8274 | LR: 4.99e-04 | 8160.0ms/step | 2.5 samples/s
+Step 1550/37465 | Loss: 3.1272 | LR: 4.99e-04 | 8199.0ms/step | 2.4 samples/s
+Step 1560/37465 | Loss: 3.0043 | LR: 4.99e-04 | 8150.2ms/step | 2.5 samples/s
+Step 1570/37465 | Loss: 2.8017 | LR: 4.99e-04 | 8121.9ms/step | 2.5 samples/s
+W0709 03:43:21.449000 128 site-packages/torch/distributed/elastic/agent/server/api.py:739] Received 15 death signal, shutting down workers
+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 128 got signal: 15

epfl_final/logs/l30k1_2b_resampled.log.reclaim_20260708_233805

@@ -0,0 +1,611 @@
+W0708 11:22:00.732000 30896 site-packages/torch/distributed/run.py:851] 
+W0708 11:22:00.732000 30896 site-packages/torch/distributed/run.py:851] *****************************************
+W0708 11:22:00.732000 30896 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. 
+W0708 11:22:00.732000 30896 site-packages/torch/distributed/run.py:851] *****************************************
+New run: Checkpoints will be saved to /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": null,
+  "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
+Checkpoint activations: True
+Checkpoint activations: True
+Initializing model from scratch (random weights)...
+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...
+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...
+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
+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-20260708_112402-hn868f80
+wandb: Run `wandb offline` to turn off syncing.
+wandb: Syncing 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...
+Switching attention mode for step 1: sparse_stochastic (configured sparse attention)
+Step 10/37465 | Loss: 12.2529 | LR: 6.67e-06 | 16108.6ms/step | 1.2 samples/s
+Step 20/37465 | Loss: 12.1155 | LR: 1.33e-05 | 8131.3ms/step | 2.5 samples/s
+Step 30/37465 | Loss: 11.6636 | LR: 2.00e-05 | 8123.6ms/step | 2.5 samples/s
+Step 40/37465 | Loss: 10.6239 | LR: 2.67e-05 | 8092.5ms/step | 2.5 samples/s
+Step 50/37465 | Loss: 10.0364 | LR: 3.33e-05 | 8084.0ms/step | 2.5 samples/s
+Step 60/37465 | Loss: 9.5529 | LR: 4.00e-05 | 8086.0ms/step | 2.5 samples/s
+Step 70/37465 | Loss: 9.0991 | LR: 4.67e-05 | 8086.6ms/step | 2.5 samples/s
+Step 80/37465 | Loss: 8.6253 | LR: 5.33e-05 | 8098.9ms/step | 2.5 samples/s
+Step 90/37465 | Loss: 8.0489 | LR: 6.00e-05 | 8077.6ms/step | 2.5 samples/s
+Step 100/37465 | Loss: 7.7035 | LR: 6.67e-05 | 8081.7ms/step | 2.5 samples/s
+Step 110/37465 | Loss: 7.3558 | LR: 7.33e-05 | 8080.7ms/step | 2.5 samples/s
+Step 120/37465 | Loss: 7.0939 | LR: 8.00e-05 | 8095.0ms/step | 2.5 samples/s
+Step 130/37465 | Loss: 6.9100 | LR: 8.67e-05 | 8109.2ms/step | 2.5 samples/s
+Step 140/37465 | Loss: 6.8013 | LR: 9.33e-05 | 8203.9ms/step | 2.4 samples/s
+Step 150/37465 | Loss: 6.8181 | LR: 1.00e-04 | 8287.3ms/step | 2.4 samples/s
+Step 160/37465 | Loss: 6.6680 | LR: 1.07e-04 | 8388.0ms/step | 2.4 samples/s
+Step 170/37465 | Loss: 6.3391 | LR: 1.13e-04 | 8137.4ms/step | 2.5 samples/s
+Step 180/37465 | Loss: 6.1800 | LR: 1.20e-04 | 8334.3ms/step | 2.4 samples/s
+Step 190/37465 | Loss: 6.2515 | LR: 1.27e-04 | 8359.5ms/step | 2.4 samples/s
+Step 200/37465 | Loss: 6.0149 | LR: 1.33e-04 | 8217.3ms/step | 2.4 samples/s
+Step 210/37465 | Loss: 5.9924 | LR: 1.40e-04 | 8219.1ms/step | 2.4 samples/s
+Step 220/37465 | Loss: 5.9197 | LR: 1.47e-04 | 8339.7ms/step | 2.4 samples/s
+Step 230/37465 | Loss: 5.8678 | LR: 1.53e-04 | 8380.4ms/step | 2.4 samples/s
+Step 240/37465 | Loss: 5.8344 | LR: 1.60e-04 | 8227.7ms/step | 2.4 samples/s
+Step 250/37465 | Loss: 5.7789 | LR: 1.67e-04 | 8311.3ms/step | 2.4 samples/s
+Step 260/37465 | Loss: 5.6418 | LR: 1.73e-04 | 8387.2ms/step | 2.4 samples/s
+Step 270/37465 | Loss: 5.6752 | LR: 1.80e-04 | 8305.4ms/step | 2.4 samples/s
+Step 280/37465 | Loss: 5.6590 | LR: 1.87e-04 | 8264.8ms/step | 2.4 samples/s
+Step 290/37465 | Loss: 5.4748 | LR: 1.93e-04 | 8364.3ms/step | 2.4 samples/s
+Step 300/37465 | Loss: 5.3536 | LR: 2.00e-04 | 8401.3ms/step | 2.4 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_resampled/run_20260708_112209/checkpoint-300
+Step 310/37465 | Loss: 5.2326 | LR: 2.07e-04 | 9109.2ms/step | 2.2 samples/s
+Step 320/37465 | Loss: 5.3250 | LR: 2.13e-04 | 8331.6ms/step | 2.4 samples/s
+Step 330/37465 | Loss: 5.5214 | LR: 2.20e-04 | 8361.3ms/step | 2.4 samples/s
+Step 340/37465 | Loss: 5.3371 | LR: 2.27e-04 | 8293.6ms/step | 2.4 samples/s
+Step 350/37465 | Loss: 5.1644 | LR: 2.33e-04 | 8315.2ms/step | 2.4 samples/s
+Step 360/37465 | Loss: 5.2510 | LR: 2.40e-04 | 8461.1ms/step | 2.4 samples/s
+Step 370/37465 | Loss: 5.0816 | LR: 2.47e-04 | 8194.5ms/step | 2.4 samples/s
+Step 380/37465 | Loss: 5.3100 | LR: 2.53e-04 | 8402.5ms/step | 2.4 samples/s
+Step 390/37465 | Loss: 5.2247 | LR: 2.60e-04 | 8362.2ms/step | 2.4 samples/s
+Step 400/37465 | Loss: 5.2271 | LR: 2.67e-04 | 8184.1ms/step | 2.4 samples/s
+Step 410/37465 | Loss: 5.0511 | LR: 2.73e-04 | 8313.9ms/step | 2.4 samples/s
+Step 420/37465 | Loss: 4.8262 | LR: 2.80e-04 | 8418.4ms/step | 2.4 samples/s
+Step 430/37465 | Loss: 4.9368 | LR: 2.87e-04 | 8264.9ms/step | 2.4 samples/s
+Step 440/37465 | Loss: 4.8697 | LR: 2.93e-04 | 8337.8ms/step | 2.4 samples/s
+Step 450/37465 | Loss: 4.8496 | LR: 3.00e-04 | 8381.9ms/step | 2.4 samples/s
+Step 460/37465 | Loss: 4.6666 | LR: 3.07e-04 | 8243.3ms/step | 2.4 samples/s
+Step 470/37465 | Loss: 4.6707 | LR: 3.13e-04 | 8268.4ms/step | 2.4 samples/s
+Step 480/37465 | Loss: 4.8633 | LR: 3.20e-04 | 8348.5ms/step | 2.4 samples/s
+Step 490/37465 | Loss: 4.5865 | LR: 3.27e-04 | 8334.4ms/step | 2.4 samples/s
+Step 500/37465 | Loss: 4.8181 | LR: 3.33e-04 | 8301.3ms/step | 2.4 samples/s
+Validation at step 500: Loss = 4.6234, Perplexity = 101.84
+Step 510/37465 | Loss: 4.4567 | LR: 3.40e-04 | 12530.7ms/step | 1.6 samples/s
+Step 520/37465 | Loss: 4.6110 | LR: 3.47e-04 | 8414.2ms/step | 2.4 samples/s
+Step 530/37465 | Loss: 4.7611 | LR: 3.53e-04 | 8301.6ms/step | 2.4 samples/s
+Step 540/37465 | Loss: 4.8849 | LR: 3.60e-04 | 8380.6ms/step | 2.4 samples/s
+Step 550/37465 | Loss: 4.3841 | LR: 3.67e-04 | 8311.8ms/step | 2.4 samples/s
+Step 560/37465 | Loss: 4.3966 | LR: 3.73e-04 | 8342.6ms/step | 2.4 samples/s
+Step 570/37465 | Loss: 4.3736 | LR: 3.80e-04 | 8409.2ms/step | 2.4 samples/s
+Step 580/37465 | Loss: 4.0383 | LR: 3.87e-04 | 8353.1ms/step | 2.4 samples/s
+Step 590/37465 | Loss: 4.4026 | LR: 3.93e-04 | 8452.4ms/step | 2.4 samples/s
+Step 600/37465 | Loss: 4.4882 | LR: 4.00e-04 | 8403.1ms/step | 2.4 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_resampled/run_20260708_112209/checkpoint-600
+Step 610/37465 | Loss: 4.2267 | LR: 4.07e-04 | 9073.7ms/step | 2.2 samples/s
+Step 620/37465 | Loss: 4.0985 | LR: 4.13e-04 | 8221.0ms/step | 2.4 samples/s
+Step 630/37465 | Loss: 4.1605 | LR: 4.20e-04 | 8221.4ms/step | 2.4 samples/s
+Step 640/37465 | Loss: 4.2581 | LR: 4.27e-04 | 8149.4ms/step | 2.5 samples/s
+Step 650/37465 | Loss: 4.1771 | LR: 4.33e-04 | 8182.6ms/step | 2.4 samples/s
+Step 660/37465 | Loss: 4.2448 | LR: 4.40e-04 | 8229.3ms/step | 2.4 samples/s
+Step 670/37465 | Loss: 4.3723 | LR: 4.47e-04 | 8152.3ms/step | 2.5 samples/s
+Step 680/37465 | Loss: 4.1398 | LR: 4.53e-04 | 8190.6ms/step | 2.4 samples/s
+Step 690/37465 | Loss: 3.9610 | LR: 4.60e-04 | 8376.1ms/step | 2.4 samples/s
+Step 700/37465 | Loss: 4.0882 | LR: 4.67e-04 | 8297.2ms/step | 2.4 samples/s
+Step 710/37465 | Loss: 4.0269 | LR: 4.73e-04 | 8247.7ms/step | 2.4 samples/s
+Step 720/37465 | Loss: 3.9731 | LR: 4.80e-04 | 8152.7ms/step | 2.5 samples/s
+Step 730/37465 | Loss: 3.5717 | LR: 4.87e-04 | 8128.0ms/step | 2.5 samples/s
+Step 740/37465 | Loss: 4.0327 | LR: 4.93e-04 | 8136.5ms/step | 2.5 samples/s
+Step 750/37465 | Loss: 3.7517 | LR: 5.00e-04 | 8156.7ms/step | 2.5 samples/s
+Step 760/37465 | Loss: 3.6823 | LR: 5.00e-04 | 8179.4ms/step | 2.4 samples/s
+Step 770/37465 | Loss: 3.6981 | LR: 5.00e-04 | 8187.6ms/step | 2.4 samples/s
+Step 780/37465 | Loss: 3.7687 | LR: 5.00e-04 | 8195.8ms/step | 2.4 samples/s
+Step 790/37465 | Loss: 3.4020 | LR: 5.00e-04 | 8201.2ms/step | 2.4 samples/s
+Step 800/37465 | Loss: 3.7875 | LR: 5.00e-04 | 8244.4ms/step | 2.4 samples/s
+Step 810/37465 | Loss: 3.8170 | LR: 5.00e-04 | 8204.8ms/step | 2.4 samples/s
+Step 820/37465 | Loss: 3.6359 | LR: 5.00e-04 | 8149.9ms/step | 2.5 samples/s
+Step 830/37465 | Loss: 3.6302 | LR: 5.00e-04 | 8159.7ms/step | 2.5 samples/s
+Step 840/37465 | Loss: 3.8067 | LR: 5.00e-04 | 8152.7ms/step | 2.5 samples/s
+Step 850/37465 | Loss: 3.4456 | LR: 5.00e-04 | 8219.3ms/step | 2.4 samples/s
+Step 860/37465 | Loss: 3.6688 | LR: 5.00e-04 | 8306.0ms/step | 2.4 samples/s
+Step 870/37465 | Loss: 3.6052 | LR: 5.00e-04 | 8273.8ms/step | 2.4 samples/s
+Step 880/37465 | Loss: 3.4955 | LR: 5.00e-04 | 8288.3ms/step | 2.4 samples/s
+Step 890/37465 | Loss: 3.5420 | LR: 5.00e-04 | 8239.0ms/step | 2.4 samples/s
+Step 900/37465 | Loss: 3.2254 | LR: 5.00e-04 | 8169.4ms/step | 2.4 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_resampled/run_20260708_112209/checkpoint-900
+Step 910/37465 | Loss: 3.2560 | LR: 5.00e-04 | 9184.3ms/step | 2.2 samples/s
+Step 920/37465 | Loss: 3.4002 | LR: 5.00e-04 | 8214.5ms/step | 2.4 samples/s
+Step 930/37465 | Loss: 3.2559 | LR: 5.00e-04 | 8258.3ms/step | 2.4 samples/s
+Step 940/37465 | Loss: 3.5204 | LR: 5.00e-04 | 8186.9ms/step | 2.4 samples/s
+Step 950/37465 | Loss: 3.7731 | LR: 5.00e-04 | 8245.3ms/step | 2.4 samples/s
+Step 960/37465 | Loss: 3.1069 | LR: 5.00e-04 | 8158.7ms/step | 2.5 samples/s
+Step 970/37465 | Loss: 3.4975 | LR: 5.00e-04 | 8265.2ms/step | 2.4 samples/s
+Step 980/37465 | Loss: 3.3028 | LR: 5.00e-04 | 8266.4ms/step | 2.4 samples/s
+Step 990/37465 | Loss: 3.2725 | LR: 5.00e-04 | 8200.5ms/step | 2.4 samples/s
+Step 1000/37465 | Loss: 3.1093 | LR: 5.00e-04 | 8170.5ms/step | 2.4 samples/s
+Validation at step 1000: Loss = 3.3107, Perplexity = 27.40
+Step 1010/37465 | Loss: 3.1640 | LR: 5.00e-04 | 8912.6ms/step | 2.2 samples/s
+Step 1020/37465 | Loss: 3.4299 | LR: 5.00e-04 | 8212.7ms/step | 2.4 samples/s
+Step 1030/37465 | Loss: 3.0263 | LR: 5.00e-04 | 8212.9ms/step | 2.4 samples/s
+Step 1040/37465 | Loss: 3.4330 | LR: 5.00e-04 | 8177.0ms/step | 2.4 samples/s
+Step 1050/37465 | Loss: 3.2475 | LR: 5.00e-04 | 8203.6ms/step | 2.4 samples/s
+Step 1060/37465 | Loss: 3.2549 | LR: 5.00e-04 | 8239.7ms/step | 2.4 samples/s
+Step 1070/37465 | Loss: 3.2350 | LR: 5.00e-04 | 8174.6ms/step | 2.4 samples/s
+Step 1080/37465 | Loss: 3.2388 | LR: 5.00e-04 | 8144.2ms/step | 2.5 samples/s
+Step 1090/37465 | Loss: 3.2410 | LR: 5.00e-04 | 8155.1ms/step | 2.5 samples/s
+Step 1100/37465 | Loss: 3.0985 | LR: 5.00e-04 | 8152.8ms/step | 2.5 samples/s
+Step 1110/37465 | Loss: 3.0416 | LR: 5.00e-04 | 8235.7ms/step | 2.4 samples/s
+Step 1120/37465 | Loss: 2.9873 | LR: 5.00e-04 | 8241.8ms/step | 2.4 samples/s
+Step 1130/37465 | Loss: 3.2195 | LR: 5.00e-04 | 8187.2ms/step | 2.4 samples/s
+Step 1140/37465 | Loss: 3.2064 | LR: 5.00e-04 | 8172.4ms/step | 2.4 samples/s
+Step 1150/37465 | Loss: 3.0180 | LR: 5.00e-04 | 8228.4ms/step | 2.4 samples/s
+Step 1160/37465 | Loss: 3.3100 | LR: 5.00e-04 | 8210.0ms/step | 2.4 samples/s
+Step 1170/37465 | Loss: 2.8792 | LR: 5.00e-04 | 8227.6ms/step | 2.4 samples/s
+Step 1180/37465 | Loss: 3.2582 | LR: 5.00e-04 | 8248.9ms/step | 2.4 samples/s
+Step 1190/37465 | Loss: 3.0823 | LR: 5.00e-04 | 8269.8ms/step | 2.4 samples/s
+Step 1200/37465 | Loss: 2.8300 | LR: 5.00e-04 | 8179.5ms/step | 2.4 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_resampled/run_20260708_112209/checkpoint-1200
+Step 1210/37465 | Loss: 3.1854 | LR: 5.00e-04 | 9095.3ms/step | 2.2 samples/s
+Step 1220/37465 | Loss: 3.2507 | LR: 5.00e-04 | 8178.1ms/step | 2.4 samples/s
+Step 1230/37465 | Loss: 2.8841 | LR: 5.00e-04 | 8179.4ms/step | 2.4 samples/s
+Step 1240/37465 | Loss: 3.1349 | LR: 5.00e-04 | 8155.6ms/step | 2.5 samples/s
+Step 1250/37465 | Loss: 2.9573 | LR: 5.00e-04 | 8202.6ms/step | 2.4 samples/s
+Step 1260/37465 | Loss: 2.9792 | LR: 5.00e-04 | 8185.9ms/step | 2.4 samples/s
+Step 1270/37465 | Loss: 3.0829 | LR: 5.00e-04 | 8252.0ms/step | 2.4 samples/s
+Step 1280/37465 | Loss: 2.9335 | LR: 5.00e-04 | 8206.7ms/step | 2.4 samples/s
+Step 1290/37465 | Loss: 3.0606 | LR: 5.00e-04 | 8187.6ms/step | 2.4 samples/s
+Step 1300/37465 | Loss: 3.3668 | LR: 5.00e-04 | 8204.1ms/step | 2.4 samples/s
+Step 1310/37465 | Loss: 3.1303 | LR: 5.00e-04 | 8189.0ms/step | 2.4 samples/s
+Step 1320/37465 | Loss: 2.9707 | LR: 5.00e-04 | 8211.9ms/step | 2.4 samples/s
+Step 1330/37465 | Loss: 3.1654 | LR: 5.00e-04 | 8237.1ms/step | 2.4 samples/s
+Step 1340/37465 | Loss: 3.0821 | LR: 5.00e-04 | 8188.5ms/step | 2.4 samples/s
+Step 1350/37465 | Loss: 3.0063 | LR: 5.00e-04 | 8214.2ms/step | 2.4 samples/s
+Step 1360/37465 | Loss: 2.9233 | LR: 5.00e-04 | 8179.1ms/step | 2.4 samples/s
+Step 1370/37465 | Loss: 3.2369 | LR: 5.00e-04 | 8193.5ms/step | 2.4 samples/s
+Step 1380/37465 | Loss: 2.9580 | LR: 5.00e-04 | 8154.0ms/step | 2.5 samples/s
+Step 1390/37465 | Loss: 2.9403 | LR: 5.00e-04 | 8206.3ms/step | 2.4 samples/s
+Step 1400/37465 | Loss: 2.9784 | LR: 5.00e-04 | 8174.0ms/step | 2.4 samples/s
+Step 1410/37465 | Loss: 3.1715 | LR: 5.00e-04 | 8226.6ms/step | 2.4 samples/s
+Step 1420/37465 | Loss: 2.7800 | LR: 5.00e-04 | 8163.5ms/step | 2.4 samples/s
+Step 1430/37465 | Loss: 3.3297 | LR: 5.00e-04 | 8170.5ms/step | 2.4 samples/s
+Step 1440/37465 | Loss: 2.9857 | LR: 5.00e-04 | 8143.5ms/step | 2.5 samples/s
+Step 1450/37465 | Loss: 3.0472 | LR: 5.00e-04 | 8186.8ms/step | 2.4 samples/s
+Step 1460/37465 | Loss: 3.1017 | LR: 5.00e-04 | 8201.9ms/step | 2.4 samples/s
+Step 1470/37465 | Loss: 2.9672 | LR: 5.00e-04 | 8281.7ms/step | 2.4 samples/s
+Step 1480/37465 | Loss: 2.8356 | LR: 5.00e-04 | 8175.8ms/step | 2.4 samples/s
+Step 1490/37465 | Loss: 2.9505 | LR: 4.99e-04 | 8182.0ms/step | 2.4 samples/s
+Step 1500/37465 | Loss: 2.9152 | LR: 4.99e-04 | 8204.8ms/step | 2.4 samples/s
+Validation at step 1500: Loss = 2.9929, Perplexity = 19.94
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_resampled/run_20260708_112209/checkpoint-1500
+Step 1510/37465 | Loss: 2.9940 | LR: 4.99e-04 | 9734.2ms/step | 2.1 samples/s
+Step 1520/37465 | Loss: 2.9082 | LR: 4.99e-04 | 8144.7ms/step | 2.5 samples/s
+Step 1530/37465 | Loss: 2.8439 | LR: 4.99e-04 | 8170.1ms/step | 2.4 samples/s
+Step 1540/37465 | Loss: 2.8328 | LR: 4.99e-04 | 8234.3ms/step | 2.4 samples/s
+Step 1550/37465 | Loss: 3.1183 | LR: 4.99e-04 | 8230.8ms/step | 2.4 samples/s
+Step 1560/37465 | Loss: 2.9935 | LR: 4.99e-04 | 8171.8ms/step | 2.4 samples/s
+Step 1570/37465 | Loss: 2.7985 | LR: 4.99e-04 | 8197.9ms/step | 2.4 samples/s
+Step 1580/37465 | Loss: 3.1618 | LR: 4.99e-04 | 8187.2ms/step | 2.4 samples/s
+Step 1590/37465 | Loss: 3.0263 | LR: 4.99e-04 | 8168.4ms/step | 2.4 samples/s
+Step 1600/37465 | Loss: 3.3189 | LR: 4.99e-04 | 8201.7ms/step | 2.4 samples/s
+Step 1610/37465 | Loss: 2.8287 | LR: 4.99e-04 | 8183.2ms/step | 2.4 samples/s
+Step 1620/37465 | Loss: 3.0913 | LR: 4.99e-04 | 8221.1ms/step | 2.4 samples/s
+Step 1630/37465 | Loss: 2.9591 | LR: 4.99e-04 | 8224.3ms/step | 2.4 samples/s
+Step 1640/37465 | Loss: 2.8990 | LR: 4.99e-04 | 8187.6ms/step | 2.4 samples/s
+Step 1650/37465 | Loss: 2.6485 | LR: 4.99e-04 | 8256.0ms/step | 2.4 samples/s
+Step 1660/37465 | Loss: 2.7825 | LR: 4.99e-04 | 8186.0ms/step | 2.4 samples/s
+Step 1670/37465 | Loss: 2.7878 | LR: 4.99e-04 | 8184.7ms/step | 2.4 samples/s
+Step 1680/37465 | Loss: 3.0618 | LR: 4.99e-04 | 8154.6ms/step | 2.5 samples/s
+Step 1690/37465 | Loss: 2.9331 | LR: 4.99e-04 | 8154.0ms/step | 2.5 samples/s
+Step 1700/37465 | Loss: 2.9924 | LR: 4.99e-04 | 8170.4ms/step | 2.4 samples/s
+Step 1710/37465 | Loss: 3.0073 | LR: 4.99e-04 | 8154.1ms/step | 2.5 samples/s
+Step 1720/37465 | Loss: 3.0590 | LR: 4.99e-04 | 8182.2ms/step | 2.4 samples/s
+Step 1730/37465 | Loss: 2.8025 | LR: 4.99e-04 | 8192.3ms/step | 2.4 samples/s
+Step 1740/37465 | Loss: 2.9592 | LR: 4.99e-04 | 8182.5ms/step | 2.4 samples/s
+Step 1750/37465 | Loss: 2.9930 | LR: 4.99e-04 | 8149.3ms/step | 2.5 samples/s
+Step 1760/37465 | Loss: 2.9122 | LR: 4.99e-04 | 8202.2ms/step | 2.4 samples/s
+Step 1770/37465 | Loss: 2.9019 | LR: 4.99e-04 | 8199.3ms/step | 2.4 samples/s
+Step 1780/37465 | Loss: 2.7550 | LR: 4.99e-04 | 8217.9ms/step | 2.4 samples/s
+Step 1790/37465 | Loss: 2.9807 | LR: 4.99e-04 | 8218.4ms/step | 2.4 samples/s
+Step 1800/37465 | Loss: 2.6345 | LR: 4.99e-04 | 8250.8ms/step | 2.4 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_resampled/run_20260708_112209/checkpoint-1800
+Step 1810/37465 | Loss: 2.9508 | LR: 4.99e-04 | 9148.1ms/step | 2.2 samples/s
+Step 1820/37465 | Loss: 2.9640 | LR: 4.99e-04 | 8218.2ms/step | 2.4 samples/s
+Step 1830/37465 | Loss: 2.8576 | LR: 4.99e-04 | 8240.7ms/step | 2.4 samples/s
+Step 1840/37465 | Loss: 2.9748 | LR: 4.99e-04 | 8187.0ms/step | 2.4 samples/s
+Step 1850/37465 | Loss: 2.7509 | LR: 4.99e-04 | 8185.5ms/step | 2.4 samples/s
+Step 1860/37465 | Loss: 2.7869 | LR: 4.99e-04 | 8164.2ms/step | 2.4 samples/s
+Step 1870/37465 | Loss: 3.1765 | LR: 4.99e-04 | 8173.6ms/step | 2.4 samples/s
+Step 1880/37465 | Loss: 2.9232 | LR: 4.99e-04 | 8160.3ms/step | 2.5 samples/s
+Step 1890/37465 | Loss: 2.9119 | LR: 4.99e-04 | 8150.2ms/step | 2.5 samples/s
+Step 1900/37465 | Loss: 2.7851 | LR: 4.99e-04 | 8160.2ms/step | 2.5 samples/s
+Step 1910/37465 | Loss: 2.7634 | LR: 4.99e-04 | 8163.8ms/step | 2.4 samples/s
+Step 1920/37465 | Loss: 2.9797 | LR: 4.99e-04 | 8163.2ms/step | 2.5 samples/s
+Step 1930/37465 | Loss: 2.8600 | LR: 4.99e-04 | 8142.1ms/step | 2.5 samples/s
+Step 1940/37465 | Loss: 2.8724 | LR: 4.99e-04 | 8167.8ms/step | 2.4 samples/s
+Step 1950/37465 | Loss: 2.7951 | LR: 4.99e-04 | 8183.9ms/step | 2.4 samples/s
+Step 1960/37465 | Loss: 2.7944 | LR: 4.99e-04 | 8172.5ms/step | 2.4 samples/s
+Step 1970/37465 | Loss: 3.0025 | LR: 4.99e-04 | 8171.6ms/step | 2.4 samples/s
+Step 1980/37465 | Loss: 2.7713 | LR: 4.99e-04 | 8160.0ms/step | 2.5 samples/s
+Step 1990/37465 | Loss: 2.7367 | LR: 4.99e-04 | 8149.8ms/step | 2.5 samples/s
+Step 2000/37465 | Loss: 2.8296 | LR: 4.99e-04 | 8178.1ms/step | 2.4 samples/s
+Validation at step 2000: Loss = 2.8157, Perplexity = 16.71
+Step 2010/37465 | Loss: 2.7407 | LR: 4.99e-04 | 8857.0ms/step | 2.3 samples/s
+Step 2020/37465 | Loss: 2.8025 | LR: 4.99e-04 | 8144.6ms/step | 2.5 samples/s
+Step 2030/37465 | Loss: 2.6655 | LR: 4.99e-04 | 8163.9ms/step | 2.4 samples/s
+Step 2040/37465 | Loss: 2.7126 | LR: 4.98e-04 | 8212.2ms/step | 2.4 samples/s
+Step 2050/37465 | Loss: 2.8032 | LR: 4.98e-04 | 8265.4ms/step | 2.4 samples/s
+Step 2060/37465 | Loss: 2.6603 | LR: 4.98e-04 | 8265.1ms/step | 2.4 samples/s
+Step 2070/37465 | Loss: 2.6939 | LR: 4.98e-04 | 8219.6ms/step | 2.4 samples/s
+Step 2080/37465 | Loss: 2.8709 | LR: 4.98e-04 | 8153.7ms/step | 2.5 samples/s
+Step 2090/37465 | Loss: 2.8547 | LR: 4.98e-04 | 8164.6ms/step | 2.4 samples/s
+Step 2100/37465 | Loss: 2.9556 | LR: 4.98e-04 | 8154.4ms/step | 2.5 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_resampled/run_20260708_112209/checkpoint-2100
+Step 2110/37465 | Loss: 2.8995 | LR: 4.98e-04 | 9140.8ms/step | 2.2 samples/s
+Step 2120/37465 | Loss: 2.8264 | LR: 4.98e-04 | 8196.8ms/step | 2.4 samples/s
+Step 2130/37465 | Loss: 2.5502 | LR: 4.98e-04 | 8167.0ms/step | 2.4 samples/s
+Step 2140/37465 | Loss: 2.8676 | LR: 4.98e-04 | 8158.7ms/step | 2.5 samples/s
+Step 2150/37465 | Loss: 2.8053 | LR: 4.98e-04 | 8175.0ms/step | 2.4 samples/s
+Step 2160/37465 | Loss: 2.7136 | LR: 4.98e-04 | 8182.5ms/step | 2.4 samples/s
+Step 2170/37465 | Loss: 2.8158 | LR: 4.98e-04 | 8178.2ms/step | 2.4 samples/s
+Step 2180/37465 | Loss: 2.8378 | LR: 4.98e-04 | 8164.8ms/step | 2.4 samples/s
+Step 2190/37465 | Loss: 2.7915 | LR: 4.98e-04 | 8185.7ms/step | 2.4 samples/s
+Step 2200/37465 | Loss: 2.7099 | LR: 4.98e-04 | 8155.4ms/step | 2.5 samples/s
+Step 2210/37465 | Loss: 2.6477 | LR: 4.98e-04 | 8213.5ms/step | 2.4 samples/s
+Step 2220/37465 | Loss: 2.8399 | LR: 4.98e-04 | 8167.4ms/step | 2.4 samples/s
+Step 2230/37465 | Loss: 2.8308 | LR: 4.98e-04 | 8169.0ms/step | 2.4 samples/s
+Step 2240/37465 | Loss: 2.8296 | LR: 4.98e-04 | 8170.1ms/step | 2.4 samples/s
+Step 2250/37465 | Loss: 2.7188 | LR: 4.98e-04 | 8209.2ms/step | 2.4 samples/s
+Step 2260/37465 | Loss: 2.8725 | LR: 4.98e-04 | 8153.3ms/step | 2.5 samples/s
+Step 2270/37465 | Loss: 2.7671 | LR: 4.98e-04 | 8187.0ms/step | 2.4 samples/s
+Step 2280/37465 | Loss: 2.5925 | LR: 4.98e-04 | 8195.9ms/step | 2.4 samples/s
+Step 2290/37465 | Loss: 2.7323 | LR: 4.98e-04 | 8184.2ms/step | 2.4 samples/s
+Step 2300/37465 | Loss: 2.8295 | LR: 4.98e-04 | 8156.4ms/step | 2.5 samples/s
+Step 2310/37465 | Loss: 2.8868 | LR: 4.98e-04 | 8166.3ms/step | 2.4 samples/s
+Step 2320/37465 | Loss: 2.7510 | LR: 4.98e-04 | 8204.4ms/step | 2.4 samples/s
+Step 2330/37465 | Loss: 3.0182 | LR: 4.98e-04 | 8176.3ms/step | 2.4 samples/s
+Step 2340/37465 | Loss: 2.6166 | LR: 4.98e-04 | 8154.8ms/step | 2.5 samples/s
+Step 2350/37465 | Loss: 3.0000 | LR: 4.98e-04 | 8207.7ms/step | 2.4 samples/s
+Step 2360/37465 | Loss: 2.6088 | LR: 4.98e-04 | 8220.1ms/step | 2.4 samples/s
+Step 2370/37465 | Loss: 2.8639 | LR: 4.98e-04 | 8221.8ms/step | 2.4 samples/s
+Step 2380/37465 | Loss: 2.6826 | LR: 4.98e-04 | 8219.5ms/step | 2.4 samples/s
+Step 2390/37465 | Loss: 2.5604 | LR: 4.98e-04 | 8159.0ms/step | 2.5 samples/s
+Step 2400/37465 | Loss: 2.5405 | LR: 4.98e-04 | 8160.4ms/step | 2.5 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_resampled/run_20260708_112209/checkpoint-2400
+Step 2410/37465 | Loss: 2.7121 | LR: 4.97e-04 | 9147.8ms/step | 2.2 samples/s
+Step 2420/37465 | Loss: 2.8993 | LR: 4.97e-04 | 8209.1ms/step | 2.4 samples/s
+Step 2430/37465 | Loss: 2.8897 | LR: 4.97e-04 | 8201.6ms/step | 2.4 samples/s
+Step 2440/37465 | Loss: 2.8287 | LR: 4.97e-04 | 8208.0ms/step | 2.4 samples/s
+Step 2450/37465 | Loss: 2.8289 | LR: 4.97e-04 | 8194.5ms/step | 2.4 samples/s
+Step 2460/37465 | Loss: 2.8179 | LR: 4.97e-04 | 8186.6ms/step | 2.4 samples/s
+Step 2470/37465 | Loss: 2.8386 | LR: 4.97e-04 | 8189.0ms/step | 2.4 samples/s
+Step 2480/37465 | Loss: 2.6495 | LR: 4.97e-04 | 8177.5ms/step | 2.4 samples/s
+Step 2490/37465 | Loss: 2.6424 | LR: 4.97e-04 | 8143.2ms/step | 2.5 samples/s
+Step 2500/37465 | Loss: 2.7739 | LR: 4.97e-04 | 8146.2ms/step | 2.5 samples/s
+Validation at step 2500: Loss = 2.7044, Perplexity = 14.95
+Step 2510/37465 | Loss: 2.8365 | LR: 4.97e-04 | 8863.9ms/step | 2.3 samples/s
+Step 2520/37465 | Loss: 2.5305 | LR: 4.97e-04 | 8186.7ms/step | 2.4 samples/s
+Step 2530/37465 | Loss: 2.6727 | LR: 4.97e-04 | 8181.0ms/step | 2.4 samples/s
+Step 2540/37465 | Loss: 2.9635 | LR: 4.97e-04 | 8212.1ms/step | 2.4 samples/s
+Step 2550/37465 | Loss: 2.4608 | LR: 4.97e-04 | 8267.2ms/step | 2.4 samples/s
+Step 2560/37465 | Loss: 2.7541 | LR: 4.97e-04 | 8212.9ms/step | 2.4 samples/s
+Step 2570/37465 | Loss: 2.7161 | LR: 4.97e-04 | 8190.1ms/step | 2.4 samples/s
+Step 2580/37465 | Loss: 2.5075 | LR: 4.97e-04 | 8142.6ms/step | 2.5 samples/s
+Step 2590/37465 | Loss: 2.7348 | LR: 4.97e-04 | 8212.6ms/step | 2.4 samples/s
+Step 2600/37465 | Loss: 2.8617 | LR: 4.97e-04 | 8269.4ms/step | 2.4 samples/s
+Step 2610/37465 | Loss: 2.5875 | LR: 4.97e-04 | 8246.0ms/step | 2.4 samples/s
+Step 2620/37465 | Loss: 2.8683 | LR: 4.97e-04 | 8259.7ms/step | 2.4 samples/s
+Step 2630/37465 | Loss: 2.6856 | LR: 4.97e-04 | 8280.1ms/step | 2.4 samples/s
+Step 2640/37465 | Loss: 2.6261 | LR: 4.97e-04 | 8257.1ms/step | 2.4 samples/s
+Step 2650/37465 | Loss: 2.6231 | LR: 4.97e-04 | 8230.0ms/step | 2.4 samples/s
+Step 2660/37465 | Loss: 2.5173 | LR: 4.97e-04 | 8192.3ms/step | 2.4 samples/s
+Step 2670/37465 | Loss: 2.5940 | LR: 4.97e-04 | 8198.6ms/step | 2.4 samples/s
+Step 2680/37465 | Loss: 2.8824 | LR: 4.97e-04 | 8178.3ms/step | 2.4 samples/s
+Step 2690/37465 | Loss: 2.6206 | LR: 4.97e-04 | 8214.1ms/step | 2.4 samples/s
+Step 2700/37465 | Loss: 2.7812 | LR: 4.97e-04 | 8207.4ms/step | 2.4 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_resampled/run_20260708_112209/checkpoint-2700
+Step 2710/37465 | Loss: 2.6251 | LR: 4.96e-04 | 9090.7ms/step | 2.2 samples/s
+Step 2720/37465 | Loss: 2.5969 | LR: 4.96e-04 | 8188.6ms/step | 2.4 samples/s
+Step 2730/37465 | Loss: 2.6373 | LR: 4.96e-04 | 8201.9ms/step | 2.4 samples/s
+Step 2740/37465 | Loss: 2.7168 | LR: 4.96e-04 | 8243.4ms/step | 2.4 samples/s
+Step 2750/37465 | Loss: 2.6317 | LR: 4.96e-04 | 8148.0ms/step | 2.5 samples/s
+Step 2760/37465 | Loss: 2.8555 | LR: 4.96e-04 | 8174.9ms/step | 2.4 samples/s
+Step 2770/37465 | Loss: 2.6921 | LR: 4.96e-04 | 8191.4ms/step | 2.4 samples/s
+Step 2780/37465 | Loss: 2.7705 | LR: 4.96e-04 | 8183.9ms/step | 2.4 samples/s
+Step 2790/37465 | Loss: 2.5200 | LR: 4.96e-04 | 8201.0ms/step | 2.4 samples/s
+Step 2800/37465 | Loss: 2.5776 | LR: 4.96e-04 | 8170.4ms/step | 2.4 samples/s
+Step 2810/37465 | Loss: 2.5499 | LR: 4.96e-04 | 8221.1ms/step | 2.4 samples/s
+Step 2820/37465 | Loss: 2.3869 | LR: 4.96e-04 | 8160.0ms/step | 2.5 samples/s
+Step 2830/37465 | Loss: 2.4888 | LR: 4.96e-04 | 8190.5ms/step | 2.4 samples/s
+Step 2840/37465 | Loss: 2.9341 | LR: 4.96e-04 | 8169.4ms/step | 2.4 samples/s
+Step 2850/37465 | Loss: 2.5894 | LR: 4.96e-04 | 8177.4ms/step | 2.4 samples/s
+Step 2860/37465 | Loss: 2.5536 | LR: 4.96e-04 | 8209.8ms/step | 2.4 samples/s
+Step 2870/37465 | Loss: 2.7532 | LR: 4.96e-04 | 8209.5ms/step | 2.4 samples/s
+Step 2880/37465 | Loss: 2.7046 | LR: 4.96e-04 | 8178.5ms/step | 2.4 samples/s
+Step 2890/37465 | Loss: 2.6350 | LR: 4.96e-04 | 8177.0ms/step | 2.4 samples/s
+Step 2900/37465 | Loss: 2.5862 | LR: 4.96e-04 | 8166.6ms/step | 2.4 samples/s
+Step 2910/37465 | Loss: 2.6961 | LR: 4.96e-04 | 8188.5ms/step | 2.4 samples/s
+Step 2920/37465 | Loss: 2.4846 | LR: 4.96e-04 | 8231.4ms/step | 2.4 samples/s
+Step 2930/37465 | Loss: 2.6273 | LR: 4.96e-04 | 8167.2ms/step | 2.4 samples/s
+Step 2940/37465 | Loss: 2.5379 | LR: 4.96e-04 | 8220.1ms/step | 2.4 samples/s
+Step 2950/37465 | Loss: 2.5631 | LR: 4.96e-04 | 8201.4ms/step | 2.4 samples/s
+Step 2960/37465 | Loss: 2.5970 | LR: 4.96e-04 | 8254.7ms/step | 2.4 samples/s
+Step 2970/37465 | Loss: 2.7225 | LR: 4.96e-04 | 8209.2ms/step | 2.4 samples/s
+Step 2980/37465 | Loss: 2.3602 | LR: 4.95e-04 | 8186.2ms/step | 2.4 samples/s
+Step 2990/37465 | Loss: 2.7677 | LR: 4.95e-04 | 8197.3ms/step | 2.4 samples/s
+Step 3000/37465 | Loss: 2.6670 | LR: 4.95e-04 | 8146.9ms/step | 2.5 samples/s
+Validation at step 3000: Loss = 2.6412, Perplexity = 14.03
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_resampled/run_20260708_112209/checkpoint-3000
+Step 3010/37465 | Loss: 2.6457 | LR: 4.95e-04 | 9731.0ms/step | 2.1 samples/s
+Step 3020/37465 | Loss: 2.6357 | LR: 4.95e-04 | 8155.9ms/step | 2.5 samples/s
+Step 3030/37465 | Loss: 2.8313 | LR: 4.95e-04 | 8174.5ms/step | 2.4 samples/s
+Step 3040/37465 | Loss: 2.7540 | LR: 4.95e-04 | 8178.7ms/step | 2.4 samples/s
+Step 3050/37465 | Loss: 2.6483 | LR: 4.95e-04 | 8186.1ms/step | 2.4 samples/s
+Step 3060/37465 | Loss: 2.6350 | LR: 4.95e-04 | 8176.8ms/step | 2.4 samples/s
+Step 3070/37465 | Loss: 2.4596 | LR: 4.95e-04 | 8184.4ms/step | 2.4 samples/s
+Step 3080/37465 | Loss: 2.4220 | LR: 4.95e-04 | 8187.9ms/step | 2.4 samples/s
+Step 3090/37465 | Loss: 2.5939 | LR: 4.95e-04 | 8211.1ms/step | 2.4 samples/s
+Step 3100/37465 | Loss: 2.6991 | LR: 4.95e-04 | 8172.4ms/step | 2.4 samples/s
+Step 3110/37465 | Loss: 2.6370 | LR: 4.95e-04 | 8190.8ms/step | 2.4 samples/s
+Step 3120/37465 | Loss: 2.5206 | LR: 4.95e-04 | 8185.5ms/step | 2.4 samples/s
+Step 3130/37465 | Loss: 2.6252 | LR: 4.95e-04 | 8191.0ms/step | 2.4 samples/s
+Step 3140/37465 | Loss: 2.4345 | LR: 4.95e-04 | 8203.0ms/step | 2.4 samples/s
+Step 3150/37465 | Loss: 2.4018 | LR: 4.95e-04 | 8216.1ms/step | 2.4 samples/s
+Step 3160/37465 | Loss: 2.5793 | LR: 4.95e-04 | 8243.0ms/step | 2.4 samples/s
+Step 3170/37465 | Loss: 2.8797 | LR: 4.95e-04 | 8210.6ms/step | 2.4 samples/s
+Step 3180/37465 | Loss: 2.8203 | LR: 4.95e-04 | 8221.0ms/step | 2.4 samples/s
+Step 3190/37465 | Loss: 2.2988 | LR: 4.95e-04 | 8160.3ms/step | 2.5 samples/s
+Step 3200/37465 | Loss: 2.4824 | LR: 4.95e-04 | 8160.1ms/step | 2.5 samples/s
+Step 3210/37465 | Loss: 2.4128 | LR: 4.94e-04 | 8168.4ms/step | 2.4 samples/s
+Step 3220/37465 | Loss: 2.6652 | LR: 4.94e-04 | 8158.0ms/step | 2.5 samples/s
+Step 3230/37465 | Loss: 2.6337 | LR: 4.94e-04 | 8203.8ms/step | 2.4 samples/s
+Step 3240/37465 | Loss: 2.5037 | LR: 4.94e-04 | 8205.8ms/step | 2.4 samples/s
+Step 3250/37465 | Loss: 2.5130 | LR: 4.94e-04 | 8188.6ms/step | 2.4 samples/s
+Step 3260/37465 | Loss: 2.6979 | LR: 4.94e-04 | 8187.4ms/step | 2.4 samples/s
+Step 3270/37465 | Loss: 2.6510 | LR: 4.94e-04 | 8178.9ms/step | 2.4 samples/s
+Step 3280/37465 | Loss: 2.6239 | LR: 4.94e-04 | 8159.1ms/step | 2.5 samples/s
+Step 3290/37465 | Loss: 2.5189 | LR: 4.94e-04 | 8208.9ms/step | 2.4 samples/s
+Step 3300/37465 | Loss: 2.4943 | LR: 4.94e-04 | 8171.3ms/step | 2.4 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_resampled/run_20260708_112209/checkpoint-3300
+Step 3310/37465 | Loss: 2.6931 | LR: 4.94e-04 | 9073.2ms/step | 2.2 samples/s
+Step 3320/37465 | Loss: 2.3834 | LR: 4.94e-04 | 8187.4ms/step | 2.4 samples/s
+Step 3330/37465 | Loss: 2.7048 | LR: 4.94e-04 | 8187.4ms/step | 2.4 samples/s
+Step 3340/37465 | Loss: 2.6655 | LR: 4.94e-04 | 8175.8ms/step | 2.4 samples/s
+Step 3350/37465 | Loss: 2.9431 | LR: 4.94e-04 | 8177.7ms/step | 2.4 samples/s
+Step 3360/37465 | Loss: 2.5442 | LR: 4.94e-04 | 8161.1ms/step | 2.5 samples/s
+Step 3370/37465 | Loss: 2.5679 | LR: 4.94e-04 | 8210.9ms/step | 2.4 samples/s
+Step 3380/37465 | Loss: 2.6058 | LR: 4.94e-04 | 8218.3ms/step | 2.4 samples/s
+Step 3390/37465 | Loss: 2.4915 | LR: 4.94e-04 | 8171.6ms/step | 2.4 samples/s
+Step 3400/37465 | Loss: 2.4060 | LR: 4.94e-04 | 8274.7ms/step | 2.4 samples/s
+Step 3410/37465 | Loss: 2.4611 | LR: 4.94e-04 | 8204.0ms/step | 2.4 samples/s
+Step 3420/37465 | Loss: 2.5989 | LR: 4.94e-04 | 8171.0ms/step | 2.4 samples/s
+Step 3430/37465 | Loss: 2.7648 | LR: 4.93e-04 | 8210.8ms/step | 2.4 samples/s
+Step 3440/37465 | Loss: 2.6462 | LR: 4.93e-04 | 8200.1ms/step | 2.4 samples/s
+Step 3450/37465 | Loss: 2.6194 | LR: 4.93e-04 | 8196.0ms/step | 2.4 samples/s
+Step 3460/37465 | Loss: 2.4992 | LR: 4.93e-04 | 8168.0ms/step | 2.4 samples/s
+Step 3470/37465 | Loss: 2.3703 | LR: 4.93e-04 | 8201.8ms/step | 2.4 samples/s
+Step 3480/37465 | Loss: 2.5212 | LR: 4.93e-04 | 8209.2ms/step | 2.4 samples/s
+Step 3490/37465 | Loss: 2.6557 | LR: 4.93e-04 | 8226.1ms/step | 2.4 samples/s
+Step 3500/37465 | Loss: 2.6127 | LR: 4.93e-04 | 8160.5ms/step | 2.5 samples/s
+Validation at step 3500: Loss = 2.5816, Perplexity = 13.22
+Step 3510/37465 | Loss: 2.8043 | LR: 4.93e-04 | 8885.6ms/step | 2.3 samples/s
+Step 3520/37465 | Loss: 2.4846 | LR: 4.93e-04 | 8180.8ms/step | 2.4 samples/s
+Step 3530/37465 | Loss: 2.5388 | LR: 4.93e-04 | 8201.8ms/step | 2.4 samples/s
+Step 3540/37465 | Loss: 2.6978 | LR: 4.93e-04 | 8211.1ms/step | 2.4 samples/s
+Step 3550/37465 | Loss: 2.5573 | LR: 4.93e-04 | 8159.3ms/step | 2.5 samples/s
+Step 3560/37465 | Loss: 2.4730 | LR: 4.93e-04 | 8189.3ms/step | 2.4 samples/s
+Step 3570/37465 | Loss: 2.4706 | LR: 4.93e-04 | 8171.7ms/step | 2.4 samples/s
+Step 3580/37465 | Loss: 2.9061 | LR: 4.93e-04 | 8210.1ms/step | 2.4 samples/s
+Step 3590/37465 | Loss: 2.5361 | LR: 4.93e-04 | 8176.7ms/step | 2.4 samples/s
+Step 3600/37465 | Loss: 2.6657 | LR: 4.93e-04 | 8179.6ms/step | 2.4 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_resampled/run_20260708_112209/checkpoint-3600
+Step 3610/37465 | Loss: 2.5515 | LR: 4.93e-04 | 9090.3ms/step | 2.2 samples/s
+Step 3620/37465 | Loss: 2.4754 | LR: 4.92e-04 | 8207.8ms/step | 2.4 samples/s
+Step 3630/37465 | Loss: 2.4045 | LR: 4.92e-04 | 8200.0ms/step | 2.4 samples/s
+Step 3640/37465 | Loss: 2.5239 | LR: 4.92e-04 | 8183.5ms/step | 2.4 samples/s
+Step 3650/37465 | Loss: 2.6176 | LR: 4.92e-04 | 8192.5ms/step | 2.4 samples/s
+Step 3660/37465 | Loss: 2.4970 | LR: 4.92e-04 | 8157.9ms/step | 2.5 samples/s
+Step 3670/37465 | Loss: 2.3967 | LR: 4.92e-04 | 8211.1ms/step | 2.4 samples/s
+Step 3680/37465 | Loss: 2.7111 | LR: 4.92e-04 | 8202.3ms/step | 2.4 samples/s
+Step 3690/37465 | Loss: 2.8511 | LR: 4.92e-04 | 8151.4ms/step | 2.5 samples/s
+Step 3700/37465 | Loss: 2.5237 | LR: 4.92e-04 | 8154.5ms/step | 2.5 samples/s
+Step 3710/37465 | Loss: 2.8043 | LR: 4.92e-04 | 8200.6ms/step | 2.4 samples/s
+Step 3720/37465 | Loss: 2.5724 | LR: 4.92e-04 | 8169.6ms/step | 2.4 samples/s
+Step 3730/37465 | Loss: 2.7419 | LR: 4.92e-04 | 8174.8ms/step | 2.4 samples/s
+Step 3740/37465 | Loss: 2.3886 | LR: 4.92e-04 | 8216.2ms/step | 2.4 samples/s
+Step 3750/37465 | Loss: 2.5297 | LR: 4.92e-04 | 8188.2ms/step | 2.4 samples/s
+Step 3760/37465 | Loss: 2.7628 | LR: 4.92e-04 | 8224.9ms/step | 2.4 samples/s
+Step 3770/37465 | Loss: 2.7325 | LR: 4.92e-04 | 8192.5ms/step | 2.4 samples/s
+Step 3780/37465 | Loss: 2.6115 | LR: 4.92e-04 | 8203.5ms/step | 2.4 samples/s
+Step 3790/37465 | Loss: 2.4737 | LR: 4.92e-04 | 8170.1ms/step | 2.4 samples/s
+Step 3800/37465 | Loss: 2.5705 | LR: 4.92e-04 | 8233.0ms/step | 2.4 samples/s
+Step 3810/37465 | Loss: 2.6723 | LR: 4.91e-04 | 8211.9ms/step | 2.4 samples/s
+Step 3820/37465 | Loss: 2.6029 | LR: 4.91e-04 | 8161.2ms/step | 2.5 samples/s
+Step 3830/37465 | Loss: 2.7360 | LR: 4.91e-04 | 8201.3ms/step | 2.4 samples/s
+Step 3840/37465 | Loss: 2.3041 | LR: 4.91e-04 | 8176.9ms/step | 2.4 samples/s
+Step 3850/37465 | Loss: 2.6950 | LR: 4.91e-04 | 8144.0ms/step | 2.5 samples/s
+Step 3860/37465 | Loss: 2.4204 | LR: 4.91e-04 | 8207.5ms/step | 2.4 samples/s
+Step 3870/37465 | Loss: 2.7476 | LR: 4.91e-04 | 8175.0ms/step | 2.4 samples/s
+Step 3880/37465 | Loss: 2.3562 | LR: 4.91e-04 | 8189.6ms/step | 2.4 samples/s
+Step 3890/37465 | Loss: 2.3554 | LR: 4.91e-04 | 8200.8ms/step | 2.4 samples/s
+Step 3900/37465 | Loss: 2.4197 | LR: 4.91e-04 | 8250.1ms/step | 2.4 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_resampled/run_20260708_112209/checkpoint-3900
+Step 3910/37465 | Loss: 2.6054 | LR: 4.91e-04 | 9080.9ms/step | 2.2 samples/s
+Step 3920/37465 | Loss: 2.5965 | LR: 4.91e-04 | 8191.7ms/step | 2.4 samples/s
+Step 3930/37465 | Loss: 2.6539 | LR: 4.91e-04 | 8170.0ms/step | 2.4 samples/s
+Step 3940/37465 | Loss: 2.8478 | LR: 4.91e-04 | 8180.8ms/step | 2.4 samples/s
+Step 3950/37465 | Loss: 2.3553 | LR: 4.91e-04 | 8138.7ms/step | 2.5 samples/s
+Step 3960/37465 | Loss: 2.5775 | LR: 4.91e-04 | 8147.7ms/step | 2.5 samples/s
+Step 3970/37465 | Loss: 2.3752 | LR: 4.91e-04 | 8177.8ms/step | 2.4 samples/s
+Step 3980/37465 | Loss: 2.5759 | LR: 4.91e-04 | 8183.7ms/step | 2.4 samples/s
+Step 3990/37465 | Loss: 2.5541 | LR: 4.90e-04 | 8148.6ms/step | 2.5 samples/s
+Step 4000/37465 | Loss: 2.5379 | LR: 4.90e-04 | 8175.9ms/step | 2.4 samples/s
+Validation at step 4000: Loss = 2.5357, Perplexity = 12.63
+Step 4010/37465 | Loss: 2.1452 | LR: 4.90e-04 | 8859.4ms/step | 2.3 samples/s
+Step 4020/37465 | Loss: 2.3252 | LR: 4.90e-04 | 8164.0ms/step | 2.4 samples/s
+Step 4030/37465 | Loss: 2.4662 | LR: 4.90e-04 | 8179.1ms/step | 2.4 samples/s
+Step 4040/37465 | Loss: 2.4781 | LR: 4.90e-04 | 8182.5ms/step | 2.4 samples/s
+Step 4050/37465 | Loss: 2.6928 | LR: 4.90e-04 | 8178.8ms/step | 2.4 samples/s
+Step 4060/37465 | Loss: 2.5012 | LR: 4.90e-04 | 8156.5ms/step | 2.5 samples/s
+Step 4070/37465 | Loss: 2.4975 | LR: 4.90e-04 | 8180.8ms/step | 2.4 samples/s
+Step 4080/37465 | Loss: 2.4839 | LR: 4.90e-04 | 8149.3ms/step | 2.5 samples/s
+Step 4090/37465 | Loss: 2.5882 | LR: 4.90e-04 | 8141.3ms/step | 2.5 samples/s
+Step 4100/37465 | Loss: 2.4278 | LR: 4.90e-04 | 8172.1ms/step | 2.4 samples/s
+Step 4110/37465 | Loss: 2.5256 | LR: 4.90e-04 | 8153.0ms/step | 2.5 samples/s
+Step 4120/37465 | Loss: 2.3251 | LR: 4.90e-04 | 8208.3ms/step | 2.4 samples/s
+Step 4130/37465 | Loss: 2.3016 | LR: 4.90e-04 | 8167.1ms/step | 2.4 samples/s
+Step 4140/37465 | Loss: 2.5975 | LR: 4.90e-04 | 8171.8ms/step | 2.4 samples/s
+Step 4150/37465 | Loss: 2.5831 | LR: 4.89e-04 | 8212.5ms/step | 2.4 samples/s
+Step 4160/37465 | Loss: 2.6225 | LR: 4.89e-04 | 8244.9ms/step | 2.4 samples/s
+Step 4170/37465 | Loss: 2.5197 | LR: 4.89e-04 | 8204.5ms/step | 2.4 samples/s
+Step 4180/37465 | Loss: 2.3115 | LR: 4.89e-04 | 8180.7ms/step | 2.4 samples/s
+Step 4190/37465 | Loss: 2.5864 | LR: 4.89e-04 | 8196.8ms/step | 2.4 samples/s
+Step 4200/37465 | Loss: 2.5527 | LR: 4.89e-04 | 8226.0ms/step | 2.4 samples/s
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l30k1_2b_resampled/run_20260708_112209/checkpoint-4200
+Step 4210/37465 | Loss: 2.1629 | LR: 4.89e-04 | 9110.9ms/step | 2.2 samples/s
+Step 4220/37465 | Loss: 2.5146 | LR: 4.89e-04 | 8213.3ms/step | 2.4 samples/s
+Step 4230/37465 | Loss: 2.4114 | LR: 4.89e-04 | 8219.3ms/step | 2.4 samples/s
+Step 4240/37465 | Loss: 2.4552 | LR: 4.89e-04 | 8197.1ms/step | 2.4 samples/s
+Step 4250/37465 | Loss: 2.3725 | LR: 4.89e-04 | 8180.0ms/step | 2.4 samples/s
+Step 4260/37465 | Loss: 2.4324 | LR: 4.89e-04 | 8227.9ms/step | 2.4 samples/s
+Step 4270/37465 | Loss: 2.5080 | LR: 4.89e-04 | 8160.4ms/step | 2.5 samples/s
+Step 4280/37465 | Loss: 2.4876 | LR: 4.89e-04 | 8213.9ms/step | 2.4 samples/s
+Step 4290/37465 | Loss: 2.5656 | LR: 4.89e-04 | 8184.7ms/step | 2.4 samples/s
+Step 4300/37465 | Loss: 2.4473 | LR: 4.89e-04 | 8167.4ms/step | 2.4 samples/s
+Step 4310/37465 | Loss: 2.6317 | LR: 4.88e-04 | 8177.3ms/step | 2.4 samples/s
+Step 4320/37465 | Loss: 2.3662 | LR: 4.88e-04 | 8154.0ms/step | 2.5 samples/s
+Step 4330/37465 | Loss: 2.2315 | LR: 4.88e-04 | 8174.0ms/step | 2.4 samples/s
+Step 4340/37465 | Loss: 2.5357 | LR: 4.88e-04 | 8260.1ms/step | 2.4 samples/s
+Step 4350/37465 | Loss: 2.6435 | LR: 4.88e-04 | 8219.4ms/step | 2.4 samples/s
+Step 4360/37465 | Loss: 2.6113 | LR: 4.88e-04 | 8226.3ms/step | 2.4 samples/s
+Step 4370/37465 | Loss: 2.4357 | LR: 4.88e-04 | 8260.2ms/step | 2.4 samples/s
+Step 4380/37465 | Loss: 2.4313 | LR: 4.88e-04 | 8229.2ms/step | 2.4 samples/s
+Step 4390/37465 | Loss: 2.5864 | LR: 4.88e-04 | 8291.7ms/step | 2.4 samples/s
+Step 4400/37465 | Loss: 2.3635 | LR: 4.88e-04 | 8218.0ms/step | 2.4 samples/s
+Step 4410/37465 | Loss: 2.6025 | LR: 4.88e-04 | 8207.8ms/step | 2.4 samples/s
+Step 4420/37465 | Loss: 2.3309 | LR: 4.88e-04 | 8289.2ms/step | 2.4 samples/s
+Step 4430/37465 | Loss: 2.5159 | LR: 4.88e-04 | 8263.9ms/step | 2.4 samples/s

epfl_final/logs/l8k365m_asentmax.log

@@ -0,0 +1,547 @@
+W0709 03:49:56.282000 4336 site-packages/torch/distributed/run.py:851] 
+W0709 03:49:56.282000 4336 site-packages/torch/distributed/run.py:851] *****************************************
+W0709 03:49:56.282000 4336 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 03:49:56.282000 4336 site-packages/torch/distributed/run.py:851] *****************************************
+New run: Checkpoints will be saved to /mloscratch/homes/navasard/sparse/checkpoints/l8k365m_asentmax/run_20260709_035002
+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/l8k365m_asentmax/run_20260709_035002",
+  "resume_from_checkpoint": null,
+  "force_new_run": true,
+  "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": 10,
+  "seed": 42,
+  "verbose": true,
+  "checkpoint_activations": true,
+  "strategy": "ddp",
+  "mixed_precision": "bf16",
+  "wandb_enabled": true,
+  "wandb_project": "stochastic_sparsity",
+  "wandb_name": "365m asentmax 8k",
+  "wandb_tags": [
+    "qwen3",
+    "365m",
+    "8k",
+    "scaleladder",
+    "asentmax"
+  ],
+  "hidden_size": 768,
+  "intermediate_size": 3072,
+  "num_attention_heads": 6,
+  "num_key_value_heads": 2,
+  "num_hidden_layers": 28,
+  "head_dim": null,
+  "attention_type": "asentmax",
+  "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: True
+Initializing model from scratch (random weights)...
+Checkpoint activations: True
+Checkpoint activations: True
+Checkpoint activations: True
+
+
+================================================================================
+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
+  attention_type: asentmax
+
+Parameter count:
+  Total parameters: 358,959,016
+  Trainable parameters: 358,959,016
+  Non-trainable parameters: 0
+  Model size (BF16): 0.72 GB
+
+--------------------------------------------------------------------------------
+LAYER-BY-LAYER SUMMARY
+--------------------------------------------------------------------------------
+===========================================================================================================================================================
+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,550                 True
+│    │    └─Qwen3DecoderLayer: 3-2                      [1, 8192, 768]            [1, 8192, 768]            8,652,550                 True
+│    │    └─Qwen3DecoderLayer: 3-3                      [1, 8192, 768]            [1, 8192, 768]            8,652,550                 True
+│    │    └─Qwen3DecoderLayer: 3-4                      [1, 8192, 768]            [1, 8192, 768]            8,652,550                 True
+│    │    └─Qwen3DecoderLayer: 3-5                      [1, 8192, 768]            [1, 8192, 768]            8,652,550                 True
+│    │    └─Qwen3DecoderLayer: 3-6                      [1, 8192, 768]            [1, 8192, 768]            8,652,550                 True
+│    │    └─Qwen3DecoderLayer: 3-7                      [1, 8192, 768]            [1, 8192, 768]            8,652,550                 True
+│    │    └─Qwen3DecoderLayer: 3-8                      [1, 8192, 768]            [1, 8192, 768]            8,652,550                 True
+│    │    └─Qwen3DecoderLayer: 3-9                      [1, 8192, 768]            [1, 8192, 768]            8,652,550                 True
+│    │    └─Qwen3DecoderLayer: 3-10                     [1, 8192, 768]            [1, 8192, 768]            8,652,550                 True
+│    │    └─Qwen3DecoderLayer: 3-11                     [1, 8192, 768]            [1, 8192, 768]            8,652,550                 True
+│    │    └─Qwen3DecoderLayer: 3-12                     [1, 8192, 768]            [1, 8192, 768]            8,652,550                 True
+│    │    └─Qwen3DecoderLayer: 3-13                     [1, 8192, 768]            [1, 8192, 768]            8,652,550                 True
+│    │    └─Qwen3DecoderLayer: 3-14                     [1, 8192, 768]            [1, 8192, 768]            8,652,550                 True
+│    │    └─Qwen3DecoderLayer: 3-15                     [1, 8192, 768]            [1, 8192, 768]            8,652,550                 True
+│    │    └─Qwen3DecoderLayer: 3-16                     [1, 8192, 768]            [1, 8192, 768]            8,652,550                 True
+│    │    └─Qwen3DecoderLayer: 3-17                     [1, 8192, 768]            [1, 8192, 768]            8,652,550                 True
+│    │    └─Qwen3DecoderLayer: 3-18                     [1, 8192, 768]            [1, 8192, 768]            8,652,550                 True
+│    │    └─Qwen3DecoderLayer: 3-19                     [1, 8192, 768]            [1, 8192, 768]            8,652,550                 True
+│    │    └─Qwen3DecoderLayer: 3-20                     [1, 8192, 768]            [1, 8192, 768]            8,652,550                 True
+│    │    └─Qwen3DecoderLayer: 3-21                     [1, 8192, 768]            [1, 8192, 768]            8,652,550                 True
+│    │    └─Qwen3DecoderLayer: 3-22                     [1, 8192, 768]            [1, 8192, 768]            8,652,550                 True
+│    │    └─Qwen3DecoderLayer: 3-23                     [1, 8192, 768]            [1, 8192, 768]            8,652,550                 True
+│    │    └─Qwen3DecoderLayer: 3-24                     [1, 8192, 768]            [1, 8192, 768]            8,652,550                 True
+│    │    └─Qwen3DecoderLayer: 3-25                     [1, 8192, 768]            [1, 8192, 768]            8,652,550                 True
+│    │    └─Qwen3DecoderLayer: 3-26                     [1, 8192, 768]            [1, 8192, 768]            8,652,550                 True
+│    │    └─Qwen3DecoderLayer: 3-27                     [1, 8192, 768]            [1, 8192, 768]            8,652,550                 True
+│    │    └─Qwen3DecoderLayer: 3-28                     [1, 8192, 768]            [1, 8192, 768]            8,652,550                 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,864
+Trainable params: 475,645,864
+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,959,016 (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)
+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)
+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)
+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-20260709_035020-omomq98g
+wandb: Run `wandb offline` to turn off syncing.
+wandb: Syncing run 365m asentmax 8k
+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/omomq98g
+Starting training...
+Switching attention mode for step 1: sparse_stochastic (configured sparse attention)
+Step 10/44000 | Loss: 12.0724 | LR: 6.67e-06 | 20928.3ms/step | 1.0 samples/s
+Step 20/44000 | Loss: 12.0316 | LR: 1.33e-05 | 20904.8ms/step | 1.0 samples/s
+Step 30/44000 | Loss: 11.8886 | LR: 2.00e-05 | 20934.2ms/step | 1.0 samples/s
+Step 40/44000 | Loss: 11.3860 | LR: 2.67e-05 | 20965.4ms/step | 1.0 samples/s
+Step 50/44000 | Loss: 10.9035 | LR: 3.33e-05 | 20972.8ms/step | 1.0 samples/s
+Step 60/44000 | Loss: 10.6336 | LR: 4.00e-05 | 20971.5ms/step | 1.0 samples/s
+Step 70/44000 | Loss: 10.4780 | LR: 4.67e-05 | 20976.4ms/step | 1.0 samples/s
+Step 80/44000 | Loss: 10.2118 | LR: 5.33e-05 | 20975.2ms/step | 1.0 samples/s
+Step 90/44000 | Loss: 9.9761 | LR: 6.00e-05 | 20974.6ms/step | 1.0 samples/s
+Step 100/44000 | Loss: 9.6347 | LR: 6.67e-05 | 20975.4ms/step | 1.0 samples/s
+Step 110/44000 | Loss: 9.0751 | LR: 7.33e-05 | 20974.9ms/step | 1.0 samples/s
+Step 120/44000 | Loss: 8.5446 | LR: 8.00e-05 | 20974.0ms/step | 1.0 samples/s
+Step 130/44000 | Loss: 8.0868 | LR: 8.67e-05 | 20977.0ms/step | 1.0 samples/s
+Step 140/44000 | Loss: 7.7200 | LR: 9.33e-05 | 20978.8ms/step | 1.0 samples/s
+Step 150/44000 | Loss: 7.4798 | LR: 1.00e-04 | 20981.3ms/step | 1.0 samples/s
+Step 160/44000 | Loss: 7.3115 | LR: 1.07e-04 | 20988.2ms/step | 1.0 samples/s
+Step 170/44000 | Loss: 7.2109 | LR: 1.13e-04 | 20990.0ms/step | 1.0 samples/s
+Step 180/44000 | Loss: 7.1714 | LR: 1.20e-04 | 20989.9ms/step | 1.0 samples/s
+Step 190/44000 | Loss: 7.0794 | LR: 1.27e-04 | 20993.5ms/step | 1.0 samples/s
+Step 200/44000 | Loss: 6.9238 | LR: 1.33e-04 | 20988.8ms/step | 1.0 samples/s
+Step 210/44000 | Loss: 6.9067 | LR: 1.40e-04 | 20991.1ms/step | 1.0 samples/s
+Step 220/44000 | Loss: 6.8194 | LR: 1.47e-04 | 20991.6ms/step | 1.0 samples/s
+Step 230/44000 | Loss: 6.6844 | LR: 1.53e-04 | 20990.6ms/step | 1.0 samples/s
+Step 240/44000 | Loss: 6.6439 | LR: 1.60e-04 | 20992.0ms/step | 1.0 samples/s
+Step 250/44000 | Loss: 6.5700 | LR: 1.67e-04 | 20988.9ms/step | 1.0 samples/s
+Validation at step 250: Loss = 6.5532, Perplexity = 701.48
+Step 260/44000 | Loss: 6.5711 | LR: 1.73e-04 | 32658.0ms/step | 0.6 samples/s
+Step 270/44000 | Loss: 6.5103 | LR: 1.80e-04 | 20991.4ms/step | 1.0 samples/s
+Step 280/44000 | Loss: 6.4252 | LR: 1.87e-04 | 20990.7ms/step | 1.0 samples/s
+Step 290/44000 | Loss: 6.3790 | LR: 1.93e-04 | 20989.9ms/step | 1.0 samples/s
+Step 300/44000 | Loss: 6.3915 | LR: 2.00e-04 | 20991.4ms/step | 1.0 samples/s
+Step 310/44000 | Loss: 6.3493 | LR: 2.07e-04 | 20989.3ms/step | 1.0 samples/s
+Step 320/44000 | Loss: 6.2548 | LR: 2.13e-04 | 20990.9ms/step | 1.0 samples/s
+Step 330/44000 | Loss: 6.2176 | LR: 2.20e-04 | 20989.4ms/step | 1.0 samples/s
+Step 340/44000 | Loss: 6.1828 | LR: 2.27e-04 | 20987.6ms/step | 1.0 samples/s
+Step 350/44000 | Loss: 6.1505 | LR: 2.33e-04 | 20992.8ms/step | 1.0 samples/s
+Step 360/44000 | Loss: 6.1303 | LR: 2.40e-04 | 20992.2ms/step | 1.0 samples/s
+Step 370/44000 | Loss: 5.9488 | LR: 2.47e-04 | 20990.6ms/step | 1.0 samples/s
+Step 380/44000 | Loss: 6.0514 | LR: 2.53e-04 | 20989.3ms/step | 1.0 samples/s
+Step 390/44000 | Loss: 6.0025 | LR: 2.60e-04 | 20992.4ms/step | 1.0 samples/s
+Step 400/44000 | Loss: 5.9410 | LR: 2.67e-04 | 20995.0ms/step | 1.0 samples/s
+Step 410/44000 | Loss: 5.9264 | LR: 2.73e-04 | 20989.3ms/step | 1.0 samples/s
+Step 420/44000 | Loss: 5.8832 | LR: 2.80e-04 | 20990.7ms/step | 1.0 samples/s
+Step 430/44000 | Loss: 5.8555 | LR: 2.87e-04 | 20988.5ms/step | 1.0 samples/s
+Step 440/44000 | Loss: 5.8668 | LR: 2.93e-04 | 20987.4ms/step | 1.0 samples/s
+Step 450/44000 | Loss: 5.7854 | LR: 3.00e-04 | 20987.8ms/step | 1.0 samples/s
+Step 460/44000 | Loss: 5.7228 | LR: 3.07e-04 | 20992.8ms/step | 1.0 samples/s
+Step 470/44000 | Loss: 5.7205 | LR: 3.13e-04 | 20990.8ms/step | 1.0 samples/s
+Step 480/44000 | Loss: 5.7083 | LR: 3.20e-04 | 20992.8ms/step | 1.0 samples/s
+Step 490/44000 | Loss: 5.6942 | LR: 3.27e-04 | 20990.3ms/step | 1.0 samples/s
+Step 500/44000 | Loss: 5.6374 | LR: 3.33e-04 | 20989.6ms/step | 1.0 samples/s
+Validation at step 500: Loss = 5.6238, Perplexity = 276.94
+Step 510/44000 | Loss: 5.5292 | LR: 3.40e-04 | 32585.6ms/step | 0.6 samples/s
+Step 520/44000 | Loss: 5.5300 | LR: 3.47e-04 | 20989.2ms/step | 1.0 samples/s
+Step 530/44000 | Loss: 5.5476 | LR: 3.53e-04 | 20996.9ms/step | 1.0 samples/s
+Step 540/44000 | Loss: 5.5214 | LR: 3.60e-04 | 20999.5ms/step | 1.0 samples/s
+Step 550/44000 | Loss: 5.5411 | LR: 3.67e-04 | 20990.5ms/step | 1.0 samples/s
+Step 560/44000 | Loss: 5.4989 | LR: 3.73e-04 | 20996.4ms/step | 1.0 samples/s
+Step 570/44000 | Loss: 5.4973 | LR: 3.80e-04 | 20993.1ms/step | 1.0 samples/s
+Step 580/44000 | Loss: 5.4088 | LR: 3.87e-04 | 20990.7ms/step | 1.0 samples/s
+Step 590/44000 | Loss: 5.3881 | LR: 3.93e-04 | 20990.3ms/step | 1.0 samples/s
+Step 600/44000 | Loss: 5.3192 | LR: 4.00e-04 | 20995.5ms/step | 1.0 samples/s
+Step 610/44000 | Loss: 5.2660 | LR: 4.07e-04 | 20992.4ms/step | 1.0 samples/s
+Step 620/44000 | Loss: 5.2317 | LR: 4.13e-04 | 20987.8ms/step | 1.0 samples/s
+Step 630/44000 | Loss: 5.2560 | LR: 4.20e-04 | 20990.4ms/step | 1.0 samples/s
+Step 640/44000 | Loss: 5.2241 | LR: 4.27e-04 | 20992.1ms/step | 1.0 samples/s
+Step 650/44000 | Loss: 5.1576 | LR: 4.33e-04 | 21001.2ms/step | 1.0 samples/s
+Step 660/44000 | Loss: 5.1848 | LR: 4.40e-04 | 20995.1ms/step | 1.0 samples/s
+Step 670/44000 | Loss: 5.0861 | LR: 4.47e-04 | 20993.9ms/step | 1.0 samples/s
+Step 680/44000 | Loss: 5.1299 | LR: 4.53e-04 | 20996.7ms/step | 1.0 samples/s
+Step 690/44000 | Loss: 5.0573 | LR: 4.60e-04 | 20993.9ms/step | 1.0 samples/s
+Step 700/44000 | Loss: 5.0336 | LR: 4.67e-04 | 20992.0ms/step | 1.0 samples/s
+Step 710/44000 | Loss: 5.0315 | LR: 4.73e-04 | 20993.1ms/step | 1.0 samples/s
+Step 720/44000 | Loss: 4.9879 | LR: 4.80e-04 | 20993.8ms/step | 1.0 samples/s
+Step 730/44000 | Loss: 4.9228 | LR: 4.87e-04 | 20991.8ms/step | 1.0 samples/s
+Step 740/44000 | Loss: 4.8911 | LR: 4.93e-04 | 20990.9ms/step | 1.0 samples/s
+Step 750/44000 | Loss: 4.8897 | LR: 5.00e-04 | 20991.5ms/step | 1.0 samples/s
+Validation at step 750: Loss = 4.8996, Perplexity = 134.24
+Step 760/44000 | Loss: 4.7838 | LR: 5.00e-04 | 32581.2ms/step | 0.6 samples/s
+Step 770/44000 | Loss: 4.8421 | LR: 5.00e-04 | 20993.2ms/step | 1.0 samples/s
+Step 780/44000 | Loss: 4.7923 | LR: 5.00e-04 | 20991.7ms/step | 1.0 samples/s
+Step 790/44000 | Loss: 4.7844 | LR: 5.00e-04 | 20991.5ms/step | 1.0 samples/s
+Step 800/44000 | Loss: 4.7864 | LR: 5.00e-04 | 20989.4ms/step | 1.0 samples/s
+Step 810/44000 | Loss: 4.7202 | LR: 5.00e-04 | 20987.6ms/step | 1.0 samples/s
+Step 820/44000 | Loss: 4.7077 | LR: 5.00e-04 | 20984.2ms/step | 1.0 samples/s
+Step 830/44000 | Loss: 4.6744 | LR: 5.00e-04 | 20988.1ms/step | 1.0 samples/s
+Step 840/44000 | Loss: 4.5939 | LR: 5.00e-04 | 20988.8ms/step | 1.0 samples/s
+Step 850/44000 | Loss: 4.6230 | LR: 5.00e-04 | 20984.2ms/step | 1.0 samples/s
+Step 860/44000 | Loss: 4.6334 | LR: 5.00e-04 | 20987.5ms/step | 1.0 samples/s
+Step 870/44000 | Loss: 4.6032 | LR: 5.00e-04 | 20986.1ms/step | 1.0 samples/s
+Step 880/44000 | Loss: 4.5534 | LR: 5.00e-04 | 20987.7ms/step | 1.0 samples/s
+Step 890/44000 | Loss: 4.5531 | LR: 5.00e-04 | 20987.4ms/step | 1.0 samples/s
+Step 900/44000 | Loss: 4.5278 | LR: 5.00e-04 | 20985.6ms/step | 1.0 samples/s
+Step 910/44000 | Loss: 4.5598 | LR: 5.00e-04 | 20984.6ms/step | 1.0 samples/s
+Step 920/44000 | Loss: 4.4118 | LR: 5.00e-04 | 20987.1ms/step | 1.0 samples/s
+Step 930/44000 | Loss: 4.4779 | LR: 5.00e-04 | 20983.0ms/step | 1.0 samples/s
+Step 940/44000 | Loss: 4.4986 | LR: 5.00e-04 | 20987.6ms/step | 1.0 samples/s
+Step 950/44000 | Loss: 4.4309 | LR: 5.00e-04 | 20984.6ms/step | 1.0 samples/s
+Step 960/44000 | Loss: 4.3400 | LR: 5.00e-04 | 20985.7ms/step | 1.0 samples/s
+Step 970/44000 | Loss: 4.4057 | LR: 5.00e-04 | 20985.2ms/step | 1.0 samples/s
+Step 980/44000 | Loss: 4.3519 | LR: 5.00e-04 | 20985.5ms/step | 1.0 samples/s
+Step 990/44000 | Loss: 4.4208 | LR: 5.00e-04 | 20987.4ms/step | 1.0 samples/s
+Step 1000/44000 | Loss: 4.3428 | LR: 5.00e-04 | 20983.3ms/step | 1.0 samples/s
+Validation at step 1000: Loss = 4.3815, Perplexity = 79.95
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l8k365m_asentmax/run_20260709_035002/checkpoint-1000
+Step 1010/44000 | Loss: 4.3564 | LR: 5.00e-04 | 32880.9ms/step | 0.6 samples/s
+Step 1020/44000 | Loss: 4.3157 | LR: 5.00e-04 | 20982.3ms/step | 1.0 samples/s
+Step 1030/44000 | Loss: 4.3328 | LR: 5.00e-04 | 20985.0ms/step | 1.0 samples/s
+Step 1040/44000 | Loss: 4.2260 | LR: 5.00e-04 | 20987.8ms/step | 1.0 samples/s
+Step 1050/44000 | Loss: 4.2697 | LR: 5.00e-04 | 20982.3ms/step | 1.0 samples/s
+Step 1060/44000 | Loss: 4.2886 | LR: 5.00e-04 | 20984.0ms/step | 1.0 samples/s
+Step 1070/44000 | Loss: 4.2226 | LR: 5.00e-04 | 20986.3ms/step | 1.0 samples/s
+Step 1080/44000 | Loss: 4.2932 | LR: 5.00e-04 | 20982.4ms/step | 1.0 samples/s
+Step 1090/44000 | Loss: 4.2030 | LR: 5.00e-04 | 20985.5ms/step | 1.0 samples/s
+Step 1100/44000 | Loss: 4.2363 | LR: 5.00e-04 | 20989.3ms/step | 1.0 samples/s
+Step 1110/44000 | Loss: 4.2310 | LR: 5.00e-04 | 20983.1ms/step | 1.0 samples/s
+Step 1120/44000 | Loss: 4.1465 | LR: 5.00e-04 | 20981.5ms/step | 1.0 samples/s
+Step 1130/44000 | Loss: 4.1680 | LR: 5.00e-04 | 20980.4ms/step | 1.0 samples/s
+Step 1140/44000 | Loss: 4.1155 | LR: 5.00e-04 | 20979.1ms/step | 1.0 samples/s
+Step 1150/44000 | Loss: 4.1772 | LR: 5.00e-04 | 20977.5ms/step | 1.0 samples/s
+Step 1160/44000 | Loss: 4.1410 | LR: 5.00e-04 | 20980.0ms/step | 1.0 samples/s
+Step 1170/44000 | Loss: 4.1056 | LR: 5.00e-04 | 20977.6ms/step | 1.0 samples/s
+Step 1180/44000 | Loss: 4.1768 | LR: 5.00e-04 | 20978.8ms/step | 1.0 samples/s
+Step 1190/44000 | Loss: 4.1269 | LR: 5.00e-04 | 20979.7ms/step | 1.0 samples/s
+Step 1200/44000 | Loss: 4.0907 | LR: 5.00e-04 | 20981.6ms/step | 1.0 samples/s
+Step 1210/44000 | Loss: 4.1294 | LR: 5.00e-04 | 20979.1ms/step | 1.0 samples/s
+Step 1220/44000 | Loss: 4.0995 | LR: 5.00e-04 | 20986.3ms/step | 1.0 samples/s
+Step 1230/44000 | Loss: 4.0976 | LR: 5.00e-04 | 20982.1ms/step | 1.0 samples/s
+Step 1240/44000 | Loss: 4.0701 | LR: 5.00e-04 | 20979.0ms/step | 1.0 samples/s
+Step 1250/44000 | Loss: 4.0568 | LR: 5.00e-04 | 20980.9ms/step | 1.0 samples/s
+Validation at step 1250: Loss = 4.1005, Perplexity = 60.37
+Step 1260/44000 | Loss: 4.1190 | LR: 5.00e-04 | 32590.9ms/step | 0.6 samples/s
+Step 1270/44000 | Loss: 3.9740 | LR: 5.00e-04 | 20980.7ms/step | 1.0 samples/s
+Step 1280/44000 | Loss: 4.0749 | LR: 5.00e-04 | 20976.7ms/step | 1.0 samples/s
+Step 1290/44000 | Loss: 4.0348 | LR: 5.00e-04 | 20980.2ms/step | 1.0 samples/s
+Step 1300/44000 | Loss: 4.0276 | LR: 5.00e-04 | 20976.2ms/step | 1.0 samples/s
+Step 1310/44000 | Loss: 3.9915 | LR: 5.00e-04 | 20977.4ms/step | 1.0 samples/s
+Step 1320/44000 | Loss: 3.9625 | LR: 5.00e-04 | 20975.2ms/step | 1.0 samples/s
+Step 1330/44000 | Loss: 3.9922 | LR: 5.00e-04 | 20992.7ms/step | 1.0 samples/s
+Step 1340/44000 | Loss: 4.0284 | LR: 5.00e-04 | 20981.9ms/step | 1.0 samples/s
+Step 1350/44000 | Loss: 3.9891 | LR: 5.00e-04 | 20981.4ms/step | 1.0 samples/s
+Step 1360/44000 | Loss: 3.9757 | LR: 5.00e-04 | 20977.6ms/step | 1.0 samples/s
+Step 1370/44000 | Loss: 3.9666 | LR: 5.00e-04 | 20980.1ms/step | 1.0 samples/s
+Step 1380/44000 | Loss: 3.9873 | LR: 5.00e-04 | 20983.7ms/step | 1.0 samples/s
+Step 1390/44000 | Loss: 3.9505 | LR: 5.00e-04 | 20982.7ms/step | 1.0 samples/s
+Step 1400/44000 | Loss: 4.0039 | LR: 5.00e-04 | 20985.0ms/step | 1.0 samples/s
+Step 1410/44000 | Loss: 3.9275 | LR: 5.00e-04 | 20987.6ms/step | 1.0 samples/s
+Step 1420/44000 | Loss: 3.9238 | LR: 5.00e-04 | 20985.9ms/step | 1.0 samples/s
+Step 1430/44000 | Loss: 3.9574 | LR: 5.00e-04 | 20987.1ms/step | 1.0 samples/s
+Step 1440/44000 | Loss: 3.8529 | LR: 5.00e-04 | 20981.0ms/step | 1.0 samples/s
+Step 1450/44000 | Loss: 3.8782 | LR: 5.00e-04 | 20987.2ms/step | 1.0 samples/s
+Step 1460/44000 | Loss: 3.9892 | LR: 5.00e-04 | 20980.0ms/step | 1.0 samples/s
+Step 1470/44000 | Loss: 3.9132 | LR: 5.00e-04 | 20983.6ms/step | 1.0 samples/s
+Step 1480/44000 | Loss: 3.8515 | LR: 5.00e-04 | 20983.3ms/step | 1.0 samples/s
+Step 1490/44000 | Loss: 3.8504 | LR: 5.00e-04 | 20981.6ms/step | 1.0 samples/s
+Step 1500/44000 | Loss: 3.9244 | LR: 5.00e-04 | 20982.7ms/step | 1.0 samples/s
+Validation at step 1500: Loss = 3.9269, Perplexity = 50.75
+Step 1510/44000 | Loss: 3.8838 | LR: 5.00e-04 | 32591.0ms/step | 0.6 samples/s
+Step 1520/44000 | Loss: 3.9142 | LR: 5.00e-04 | 20986.7ms/step | 1.0 samples/s
+Step 1530/44000 | Loss: 3.9495 | LR: 5.00e-04 | 20982.1ms/step | 1.0 samples/s
+Step 1540/44000 | Loss: 3.8108 | LR: 5.00e-04 | 20985.1ms/step | 1.0 samples/s
+Step 1550/44000 | Loss: 3.8586 | LR: 5.00e-04 | 20981.4ms/step | 1.0 samples/s
+Step 1560/44000 | Loss: 3.8401 | LR: 5.00e-04 | 20976.4ms/step | 1.0 samples/s
+Step 1570/44000 | Loss: 3.8809 | LR: 5.00e-04 | 20983.7ms/step | 1.0 samples/s
+Step 1580/44000 | Loss: 3.8293 | LR: 5.00e-04 | 20993.8ms/step | 1.0 samples/s
+Step 1590/44000 | Loss: 3.8578 | LR: 5.00e-04 | 20983.3ms/step | 1.0 samples/s
+Step 1600/44000 | Loss: 3.8406 | LR: 5.00e-04 | 20984.1ms/step | 1.0 samples/s
+Step 1610/44000 | Loss: 3.8554 | LR: 5.00e-04 | 20978.7ms/step | 1.0 samples/s
+Step 1620/44000 | Loss: 3.8300 | LR: 5.00e-04 | 20980.9ms/step | 1.0 samples/s
+Step 1630/44000 | Loss: 3.8049 | LR: 4.99e-04 | 20982.5ms/step | 1.0 samples/s
+Step 1640/44000 | Loss: 3.7415 | LR: 4.99e-04 | 20981.7ms/step | 1.0 samples/s
+Step 1650/44000 | Loss: 3.8732 | LR: 4.99e-04 | 20980.5ms/step | 1.0 samples/s
+Step 1660/44000 | Loss: 3.8299 | LR: 4.99e-04 | 20985.4ms/step | 1.0 samples/s
+Step 1670/44000 | Loss: 3.8277 | LR: 4.99e-04 | 20981.4ms/step | 1.0 samples/s
+Step 1680/44000 | Loss: 3.7691 | LR: 4.99e-04 | 20975.0ms/step | 1.0 samples/s
+Step 1690/44000 | Loss: 3.7753 | LR: 4.99e-04 | 20985.8ms/step | 1.0 samples/s
+Step 1700/44000 | Loss: 3.7924 | LR: 4.99e-04 | 20981.4ms/step | 1.0 samples/s
+Step 1710/44000 | Loss: 3.8039 | LR: 4.99e-04 | 20981.7ms/step | 1.0 samples/s
+Step 1720/44000 | Loss: 3.8377 | LR: 4.99e-04 | 20979.9ms/step | 1.0 samples/s
+Step 1730/44000 | Loss: 3.7867 | LR: 4.99e-04 | 20981.4ms/step | 1.0 samples/s
+Step 1740/44000 | Loss: 3.7984 | LR: 4.99e-04 | 20979.1ms/step | 1.0 samples/s
+Step 1750/44000 | Loss: 3.7478 | LR: 4.99e-04 | 20980.1ms/step | 1.0 samples/s
+Validation at step 1750: Loss = 3.8100, Perplexity = 45.15
+Step 1760/44000 | Loss: 3.7911 | LR: 4.99e-04 | 32580.4ms/step | 0.6 samples/s
+Step 1770/44000 | Loss: 3.8000 | LR: 4.99e-04 | 20972.1ms/step | 1.0 samples/s
+Step 1780/44000 | Loss: 3.7718 | LR: 4.99e-04 | 20968.0ms/step | 1.0 samples/s
+Step 1790/44000 | Loss: 3.7545 | LR: 4.99e-04 | 20975.4ms/step | 1.0 samples/s
+Step 1800/44000 | Loss: 3.7023 | LR: 4.99e-04 | 20973.2ms/step | 1.0 samples/s
+Step 1810/44000 | Loss: 3.7222 | LR: 4.99e-04 | 20972.1ms/step | 1.0 samples/s
+Step 1820/44000 | Loss: 3.7229 | LR: 4.99e-04 | 20971.0ms/step | 1.0 samples/s
+Step 1830/44000 | Loss: 3.7199 | LR: 4.99e-04 | 20968.9ms/step | 1.0 samples/s
+Step 1840/44000 | Loss: 3.7150 | LR: 4.99e-04 | 20970.7ms/step | 1.0 samples/s
+Step 1850/44000 | Loss: 3.7085 | LR: 4.99e-04 | 20973.5ms/step | 1.0 samples/s
+Step 1860/44000 | Loss: 3.7856 | LR: 4.99e-04 | 20970.8ms/step | 1.0 samples/s
+Step 1870/44000 | Loss: 3.7453 | LR: 4.99e-04 | 20970.6ms/step | 1.0 samples/s
+Step 1880/44000 | Loss: 3.7322 | LR: 4.99e-04 | 20971.7ms/step | 1.0 samples/s
+Step 1890/44000 | Loss: 3.6779 | LR: 4.99e-04 | 20971.2ms/step | 1.0 samples/s
+Step 1900/44000 | Loss: 3.6995 | LR: 4.99e-04 | 20968.9ms/step | 1.0 samples/s
+Step 1910/44000 | Loss: 3.7893 | LR: 4.99e-04 | 20972.9ms/step | 1.0 samples/s
+Step 1920/44000 | Loss: 3.7021 | LR: 4.99e-04 | 20969.2ms/step | 1.0 samples/s
+Step 1930/44000 | Loss: 3.7550 | LR: 4.99e-04 | 20971.4ms/step | 1.0 samples/s
+Step 1940/44000 | Loss: 3.7057 | LR: 4.99e-04 | 20972.2ms/step | 1.0 samples/s
+Step 1950/44000 | Loss: 3.6347 | LR: 4.99e-04 | 20969.8ms/step | 1.0 samples/s
+Step 1960/44000 | Loss: 3.6976 | LR: 4.99e-04 | 20970.5ms/step | 1.0 samples/s
+Step 1970/44000 | Loss: 3.6654 | LR: 4.99e-04 | 20967.6ms/step | 1.0 samples/s
+Step 1980/44000 | Loss: 3.6457 | LR: 4.99e-04 | 20968.6ms/step | 1.0 samples/s
+Step 1990/44000 | Loss: 3.7303 | LR: 4.99e-04 | 20971.0ms/step | 1.0 samples/s
+Step 2000/44000 | Loss: 3.6999 | LR: 4.99e-04 | 20971.8ms/step | 1.0 samples/s
+Validation at step 2000: Loss = 3.7277, Perplexity = 41.58
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l8k365m_asentmax/run_20260709_035002/checkpoint-2000
+Step 2010/44000 | Loss: 3.6757 | LR: 4.99e-04 | 32890.4ms/step | 0.6 samples/s
+Step 2020/44000 | Loss: 3.6951 | LR: 4.99e-04 | 20970.1ms/step | 1.0 samples/s
+Step 2030/44000 | Loss: 3.6851 | LR: 4.99e-04 | 20971.3ms/step | 1.0 samples/s
+Step 2040/44000 | Loss: 3.6652 | LR: 4.99e-04 | 20968.8ms/step | 1.0 samples/s
+Step 2050/44000 | Loss: 3.7102 | LR: 4.99e-04 | 20971.8ms/step | 1.0 samples/s
+Step 2060/44000 | Loss: 3.7074 | LR: 4.99e-04 | 20970.9ms/step | 1.0 samples/s
+Step 2070/44000 | Loss: 3.6658 | LR: 4.99e-04 | 20970.0ms/step | 1.0 samples/s
+Step 2080/44000 | Loss: 3.6761 | LR: 4.99e-04 | 20969.1ms/step | 1.0 samples/s
+Step 2090/44000 | Loss: 3.6928 | LR: 4.99e-04 | 20969.8ms/step | 1.0 samples/s
+Step 2100/44000 | Loss: 3.6080 | LR: 4.99e-04 | 20967.8ms/step | 1.0 samples/s
+Step 2110/44000 | Loss: 3.7178 | LR: 4.99e-04 | 20971.8ms/step | 1.0 samples/s
+Step 2120/44000 | Loss: 3.6932 | LR: 4.99e-04 | 20970.4ms/step | 1.0 samples/s
+Step 2130/44000 | Loss: 3.6617 | LR: 4.99e-04 | 20969.8ms/step | 1.0 samples/s
+Step 2140/44000 | Loss: 3.6685 | LR: 4.99e-04 | 20968.5ms/step | 1.0 samples/s
+Step 2150/44000 | Loss: 3.6519 | LR: 4.99e-04 | 20969.6ms/step | 1.0 samples/s
+Step 2160/44000 | Loss: 3.6223 | LR: 4.99e-04 | 20966.7ms/step | 1.0 samples/s
+Step 2170/44000 | Loss: 3.6344 | LR: 4.99e-04 | 20968.7ms/step | 1.0 samples/s
+Step 2180/44000 | Loss: 3.6332 | LR: 4.99e-04 | 20970.1ms/step | 1.0 samples/s
+Step 2190/44000 | Loss: 3.6732 | LR: 4.99e-04 | 20971.1ms/step | 1.0 samples/s
+Step 2200/44000 | Loss: 3.6500 | LR: 4.99e-04 | 20974.9ms/step | 1.0 samples/s
+Step 2210/44000 | Loss: 3.6043 | LR: 4.99e-04 | 20971.5ms/step | 1.0 samples/s
+Step 2220/44000 | Loss: 3.5832 | LR: 4.99e-04 | 20972.6ms/step | 1.0 samples/s
+Step 2230/44000 | Loss: 3.6871 | LR: 4.99e-04 | 20975.7ms/step | 1.0 samples/s
+Step 2240/44000 | Loss: 3.6899 | LR: 4.99e-04 | 20972.5ms/step | 1.0 samples/s
+Step 2250/44000 | Loss: 3.6088 | LR: 4.99e-04 | 20969.7ms/step | 1.0 samples/s
+Validation at step 2250: Loss = 3.6535, Perplexity = 38.61
+Step 2260/44000 | Loss: 3.5874 | LR: 4.98e-04 | 32553.3ms/step | 0.6 samples/s
+Step 2270/44000 | Loss: 3.6395 | LR: 4.98e-04 | 20984.3ms/step | 1.0 samples/s
+Step 2280/44000 | Loss: 3.6519 | LR: 4.98e-04 | 20972.1ms/step | 1.0 samples/s
+Step 2290/44000 | Loss: 3.6179 | LR: 4.98e-04 | 20966.8ms/step | 1.0 samples/s
+Step 2300/44000 | Loss: 3.5735 | LR: 4.98e-04 | 20969.2ms/step | 1.0 samples/s
+Step 2310/44000 | Loss: 3.5915 | LR: 4.98e-04 | 20964.1ms/step | 1.0 samples/s
+Step 2320/44000 | Loss: 3.6280 | LR: 4.98e-04 | 20965.1ms/step | 1.0 samples/s
+Step 2330/44000 | Loss: 3.6137 | LR: 4.98e-04 | 20973.1ms/step | 1.0 samples/s
+Step 2340/44000 | Loss: 3.5620 | LR: 4.98e-04 | 20968.9ms/step | 1.0 samples/s
+Step 2350/44000 | Loss: 3.5864 | LR: 4.98e-04 | 20968.2ms/step | 1.0 samples/s
+Step 2360/44000 | Loss: 3.5876 | LR: 4.98e-04 | 20966.3ms/step | 1.0 samples/s
+Step 2370/44000 | Loss: 3.5795 | LR: 4.98e-04 | 20968.2ms/step | 1.0 samples/s
+Step 2380/44000 | Loss: 3.5352 | LR: 4.98e-04 | 20966.3ms/step | 1.0 samples/s
+Step 2390/44000 | Loss: 3.5356 | LR: 4.98e-04 | 20964.5ms/step | 1.0 samples/s
+Step 2400/44000 | Loss: 3.6032 | LR: 4.98e-04 | 20967.4ms/step | 1.0 samples/s
+Step 2410/44000 | Loss: 3.5932 | LR: 4.98e-04 | 20966.9ms/step | 1.0 samples/s
+Step 2420/44000 | Loss: 3.5512 | LR: 4.98e-04 | 20968.1ms/step | 1.0 samples/s
+Step 2430/44000 | Loss: 3.6120 | LR: 4.98e-04 | 20966.0ms/step | 1.0 samples/s
+Step 2440/44000 | Loss: 3.6034 | LR: 4.98e-04 | 20967.8ms/step | 1.0 samples/s
+Step 2450/44000 | Loss: 3.5318 | LR: 4.98e-04 | 20966.0ms/step | 1.0 samples/s
+Step 2460/44000 | Loss: 3.5723 | LR: 4.98e-04 | 20988.1ms/step | 1.0 samples/s
+Step 2470/44000 | Loss: 3.5486 | LR: 4.98e-04 | 20980.4ms/step | 1.0 samples/s
+Step 2480/44000 | Loss: 3.6036 | LR: 4.98e-04 | 20962.4ms/step | 1.0 samples/s
+Step 2490/44000 | Loss: 3.5721 | LR: 4.98e-04 | 20965.0ms/step | 1.0 samples/s
+Step 2500/44000 | Loss: 3.5531 | LR: 4.98e-04 | 20965.2ms/step | 1.0 samples/s
+Validation at step 2500: Loss = 3.5958, Perplexity = 36.44
+Step 2510/44000 | Loss: 3.5825 | LR: 4.98e-04 | 32545.3ms/step | 0.6 samples/s
+Step 2520/44000 | Loss: 3.5574 | LR: 4.98e-04 | 20971.0ms/step | 1.0 samples/s
+Step 2530/44000 | Loss: 3.5345 | LR: 4.98e-04 | 20968.0ms/step | 1.0 samples/s
+Step 2540/44000 | Loss: 3.5436 | LR: 4.98e-04 | 20966.8ms/step | 1.0 samples/s
+Step 2550/44000 | Loss: 3.4974 | LR: 4.98e-04 | 20965.5ms/step | 1.0 samples/s
+Step 2560/44000 | Loss: 3.4953 | LR: 4.98e-04 | 20968.8ms/step | 1.0 samples/s
+Step 2570/44000 | Loss: 3.5563 | LR: 4.98e-04 | 20974.1ms/step | 1.0 samples/s
+Step 2580/44000 | Loss: 3.5388 | LR: 4.98e-04 | 20968.9ms/step | 1.0 samples/s
+Step 2590/44000 | Loss: 3.5290 | LR: 4.98e-04 | 20967.1ms/step | 1.0 samples/s
+Step 2600/44000 | Loss: 3.4965 | LR: 4.98e-04 | 20968.0ms/step | 1.0 samples/s
+Step 2610/44000 | Loss: 3.4744 | LR: 4.98e-04 | 20984.6ms/step | 1.0 samples/s
+Step 2620/44000 | Loss: 3.5145 | LR: 4.98e-04 | 20982.5ms/step | 1.0 samples/s
+Step 2630/44000 | Loss: 3.5248 | LR: 4.98e-04 | 20979.2ms/step | 1.0 samples/s
+Step 2640/44000 | Loss: 3.5079 | LR: 4.98e-04 | 20967.0ms/step | 1.0 samples/s
+Step 2650/44000 | Loss: 3.5395 | LR: 4.98e-04 | 20974.9ms/step | 1.0 samples/s
+Step 2660/44000 | Loss: 3.4869 | LR: 4.98e-04 | 20971.7ms/step | 1.0 samples/s
+Step 2670/44000 | Loss: 3.4808 | LR: 4.98e-04 | 20978.5ms/step | 1.0 samples/s
+Step 2680/44000 | Loss: 3.5431 | LR: 4.98e-04 | 20968.3ms/step | 1.0 samples/s
+Step 2690/44000 | Loss: 3.5217 | LR: 4.98e-04 | 20969.2ms/step | 1.0 samples/s
+Step 2700/44000 | Loss: 3.4927 | LR: 4.97e-04 | 20974.4ms/step | 1.0 samples/s
+Step 2710/44000 | Loss: 3.5948 | LR: 4.97e-04 | 20967.6ms/step | 1.0 samples/s
+Step 2720/44000 | Loss: 3.5314 | LR: 4.97e-04 | 20978.1ms/step | 1.0 samples/s
+Step 2730/44000 | Loss: 3.5454 | LR: 4.97e-04 | 20990.6ms/step | 1.0 samples/s
+Step 2740/44000 | Loss: 3.5148 | LR: 4.97e-04 | 20965.7ms/step | 1.0 samples/s
+Step 2750/44000 | Loss: 3.5002 | LR: 4.97e-04 | 20964.2ms/step | 1.0 samples/s
+Validation at step 2750: Loss = 3.5515, Perplexity = 34.86
+Step 2760/44000 | Loss: 3.5015 | LR: 4.97e-04 | 32541.7ms/step | 0.6 samples/s
+Step 2770/44000 | Loss: 3.5851 | LR: 4.97e-04 | 20962.7ms/step | 1.0 samples/s
+Step 2780/44000 | Loss: 3.5420 | LR: 4.97e-04 | 20987.1ms/step | 1.0 samples/s
+Step 2790/44000 | Loss: 3.5680 | LR: 4.97e-04 | 20975.8ms/step | 1.0 samples/s
+Step 2800/44000 | Loss: 3.5334 | LR: 4.97e-04 | 20979.2ms/step | 1.0 samples/s
+Step 2810/44000 | Loss: 3.4931 | LR: 4.97e-04 | 20987.1ms/step | 1.0 samples/s
+Step 2820/44000 | Loss: 3.5117 | LR: 4.97e-04 | 20973.2ms/step | 1.0 samples/s
+Step 2830/44000 | Loss: 3.5372 | LR: 4.97e-04 | 20976.6ms/step | 1.0 samples/s
+Step 2840/44000 | Loss: 3.4655 | LR: 4.97e-04 | 20972.4ms/step | 1.0 samples/s
+Step 2850/44000 | Loss: 3.5326 | LR: 4.97e-04 | 20977.5ms/step | 1.0 samples/s
+Step 2860/44000 | Loss: 3.5398 | LR: 4.97e-04 | 20966.9ms/step | 1.0 samples/s
+Step 2870/44000 | Loss: 3.5139 | LR: 4.97e-04 | 20968.3ms/step | 1.0 samples/s
+Step 2880/44000 | Loss: 3.4544 | LR: 4.97e-04 | 20968.2ms/step | 1.0 samples/s
+Step 2890/44000 | Loss: 3.4998 | LR: 4.97e-04 | 20967.7ms/step | 1.0 samples/s
+Step 2900/44000 | Loss: 3.4792 | LR: 4.97e-04 | 20977.9ms/step | 1.0 samples/s
+Step 2910/44000 | Loss: 3.5079 | LR: 4.97e-04 | 20972.3ms/step | 1.0 samples/s
+Step 2920/44000 | Loss: 3.4640 | LR: 4.97e-04 | 20966.3ms/step | 1.0 samples/s
+Step 2930/44000 | Loss: 3.4932 | LR: 4.97e-04 | 20965.2ms/step | 1.0 samples/s
+Step 2940/44000 | Loss: 3.4723 | LR: 4.97e-04 | 20975.6ms/step | 1.0 samples/s
+Step 2950/44000 | Loss: 3.4691 | LR: 4.97e-04 | 21001.6ms/step | 1.0 samples/s
+Step 2960/44000 | Loss: 3.5163 | LR: 4.97e-04 | 20994.9ms/step | 1.0 samples/s
+Step 2970/44000 | Loss: 3.5253 | LR: 4.97e-04 | 20967.1ms/step | 1.0 samples/s
+Step 2980/44000 | Loss: 3.6491 | LR: 4.97e-04 | 20971.5ms/step | 1.0 samples/s
+Step 2990/44000 | Loss: 3.4502 | LR: 4.97e-04 | 20976.5ms/step | 1.0 samples/s
+Step 3000/44000 | Loss: 3.5199 | LR: 4.97e-04 | 20972.1ms/step | 1.0 samples/s
+Validation at step 3000: Loss = 3.5135, Perplexity = 33.57
+Saved checkpoint to /mloscratch/homes/navasard/sparse/checkpoints/l8k365m_asentmax/run_20260709_035002/checkpoint-3000
+Step 3010/44000 | Loss: 3.4805 | LR: 4.97e-04 | 32856.6ms/step | 0.6 samples/s
+Step 3020/44000 | Loss: 3.4664 | LR: 4.97e-04 | 20972.7ms/step | 1.0 samples/s
+Step 3030/44000 | Loss: 3.5003 | LR: 4.97e-04 | 20967.7ms/step | 1.0 samples/s
+Step 3040/44000 | Loss: 3.5301 | LR: 4.97e-04 | 20973.5ms/step | 1.0 samples/s
+Step 3050/44000 | Loss: 3.5501 | LR: 4.97e-04 | 20975.9ms/step | 1.0 samples/s
+Step 3060/44000 | Loss: 3.3909 | LR: 4.96e-04 | 20966.9ms/step | 1.0 samples/s
+Step 3070/44000 | Loss: 3.4825 | LR: 4.96e-04 | 20967.3ms/step | 1.0 samples/s
+Step 3080/44000 | Loss: 3.4758 | LR: 4.96e-04 | 20978.8ms/step | 1.0 samples/s
+Step 3090/44000 | Loss: 3.4524 | LR: 4.96e-04 | 20976.1ms/step | 1.0 samples/s
+Step 3100/44000 | Loss: 3.4422 | LR: 4.96e-04 | 20966.2ms/step | 1.0 samples/s
+Step 3110/44000 | Loss: 3.4104 | LR: 4.96e-04 | 20988.5ms/step | 1.0 samples/s
+Step 3120/44000 | Loss: 3.4620 | LR: 4.96e-04 | 20969.1ms/step | 1.0 samples/s
+Step 3130/44000 | Loss: 3.4514 | LR: 4.96e-04 | 20970.8ms/step | 1.0 samples/s
+W0709 22:30:51.721000 4336 site-packages/torch/distributed/elastic/agent/server/api.py:739] Received 15 death signal, shutting down workers
+W0709 22:30:51.728000 4336 site-packages/torch/distributed/elastic/multiprocessing/api.py:1012] Sending process 4402 closing signal SIGTERM
+W0709 22:30:51.730000 4336 site-packages/torch/distributed/elastic/multiprocessing/api.py:1012] Sending process 4403 closing signal SIGTERM
+W0709 22:30:51.734000 4336 site-packages/torch/distributed/elastic/multiprocessing/api.py:1012] Sending process 4404 closing signal SIGTERM
+W0709 22:30:51.736000 4336 site-packages/torch/distributed/elastic/multiprocessing/api.py:1012] Sending process 4405 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 4336 got signal: 15

epfl_final/logs/l8k365m_fixed_s1_seed43.log

@@ -0,0 +1,205 @@
+/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(
+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-20260802_194752-oxm8wv25
+wandb: Run `wandb offline` to turn off syncing.
+wandb: Syncing run 365m fixed s1 8k seed43
+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/oxm8wv25
+New run: Checkpoints will be saved to /mloscratch/homes/navasard/sparse/checkpoints/l8k365m_fixed_s1_seed43/run_20260802_194730
+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/l8k365m_fixed_s1_seed43/run_20260802_194730",
+  "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": 10,
+  "seed": 43,
+  "verbose": true,
+  "checkpoint_activations": false,
+  "strategy": "ddp",
+  "mixed_precision": "bf16",
+  "wandb_enabled": true,
+  "wandb_project": "stochastic_sparsity",
+  "wandb_name": "365m fixed s1 8k seed43",
+  "wandb_tags": [
+    "qwen3",
+    "365m",
+    "8k",
+    "scaleladder",
+    "fixed",
+    "seed1",
+    "seed43",
+    "seedvariance"
+  ],
+  "hidden_size": 768,
+  "intermediate_size": 3072,
+  "num_attention_heads": 6,
+  "num_key_value_heads": 2,
+  "num_hidden_layers": 28,
+  "head_dim": null,
+  "attention_type": "sparse_stochastic",
+  "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": "/mloscratch/homes/navasard/sparse/mask_cache_8k_fixed_s1",
+  "sparse_mask_seed": 1,
+  "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
+  attention_type: sparse_stochastic
+  sparse_block_size: 128
+  sparse_num_random_blocks: 31
+  sparse_mask_pool_size: 1
+  sparse_max_seq_length: 8192
+
+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
+--------------------------------------------------------------------------------
+================================================================================================================================================================
+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)
+Starting training...
+Switching attention mode for step 1: sparse_stochastic (configured sparse attention)
+Step 10/44000 | Loss: 12.0467 | LR: 6.67e-06 | 10774.2ms/step | 1.9 samples/s
+Step 20/44000 | Loss: 11.9512 | LR: 1.33e-05 | 3230.4ms/step | 6.2 samples/s
+Step 30/44000 | Loss: 11.6272 | LR: 2.00e-05 | 3236.4ms/step | 6.2 samples/s
+Step 40/44000 | Loss: 11.0980 | LR: 2.67e-05 | 3221.2ms/step | 6.2 samples/s
+Step 50/44000 | Loss: 10.7668 | LR: 3.33e-05 | 3221.7ms/step | 6.2 samples/s
+Step 60/44000 | Loss: 10.5338 | LR: 4.00e-05 | 3213.8ms/step | 6.2 samples/s
+Step 70/44000 | Loss: 10.2891 | LR: 4.67e-05 | 3221.2ms/step | 6.2 samples/s
+Step 80/44000 | Loss: 9.9842 | LR: 5.33e-05 | 3216.4ms/step | 6.2 samples/s
+Step 90/44000 | Loss: 9.6242 | LR: 6.00e-05 | 3232.4ms/step | 6.2 samples/s
+Step 100/44000 | Loss: 9.1368 | LR: 6.67e-05 | 3228.7ms/step | 6.2 samples/s
+Step 110/44000 | Loss: 8.6305 | LR: 7.33e-05 | 3232.4ms/step | 6.2 samples/s
+Step 120/44000 | Loss: 8.1316 | LR: 8.00e-05 | 3238.5ms/step | 6.2 samples/s
+Step 130/44000 | Loss: 7.7702 | LR: 8.67e-05 | 3229.8ms/step | 6.2 samples/s
+Step 140/44000 | Loss: 7.4971 | LR: 9.33e-05 | 3239.4ms/step | 6.2 samples/s
+Step 150/44000 | Loss: 7.3383 | LR: 1.00e-04 | 3237.1ms/step | 6.2 samples/s
+Step 160/44000 | Loss: 7.2020 | LR: 1.07e-04 | 3237.8ms/step | 6.2 samples/s
+Step 170/44000 | Loss: 7.1126 | LR: 1.13e-04 | 3235.5ms/step | 6.2 samples/s
+Step 180/44000 | Loss: 7.0300 | LR: 1.20e-04 | 3238.1ms/step | 6.2 samples/s
+Step 190/44000 | Loss: 6.9173 | LR: 1.27e-04 | 3236.0ms/step | 6.2 samples/s
+Step 200/44000 | Loss: 6.8137 | LR: 1.33e-04 | 3240.9ms/step | 6.2 samples/s
+Step 210/44000 | Loss: 6.7266 | LR: 1.40e-04 | 3235.9ms/step | 6.2 samples/s
+Step 220/44000 | Loss: 6.6680 | LR: 1.47e-04 | 3245.9ms/step | 6.2 samples/s
+Step 230/44000 | Loss: 6.5990 | LR: 1.53e-04 | 3236.4ms/step | 6.2 samples/s
+Step 240/44000 | Loss: 6.5207 | LR: 1.60e-04 | 3243.0ms/step | 6.2 samples/s
+Step 250/44000 | Loss: 6.4579 | LR: 1.67e-04 | 3234.2ms/step | 6.2 samples/s

epfl_final/logs/l8k365m_fixed_s1_seed44.log

@@ -0,0 +1,205 @@
+/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(
+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-20260802_194757-rd5nxubb
+wandb: Run `wandb offline` to turn off syncing.
+wandb: Syncing run 365m fixed s1 8k seed44
+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/rd5nxubb
+New run: Checkpoints will be saved to /mloscratch/homes/navasard/sparse/checkpoints/l8k365m_fixed_s1_seed44/run_20260802_194739
+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/l8k365m_fixed_s1_seed44/run_20260802_194739",
+  "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": 10,
+  "seed": 44,
+  "verbose": true,
+  "checkpoint_activations": false,
+  "strategy": "ddp",
+  "mixed_precision": "bf16",
+  "wandb_enabled": true,
+  "wandb_project": "stochastic_sparsity",
+  "wandb_name": "365m fixed s1 8k seed44",
+  "wandb_tags": [
+    "qwen3",
+    "365m",
+    "8k",
+    "scaleladder",
+    "fixed",
+    "seed1",
+    "seed44",
+    "seedvariance"
+  ],
+  "hidden_size": 768,
+  "intermediate_size": 3072,
+  "num_attention_heads": 6,
+  "num_key_value_heads": 2,
+  "num_hidden_layers": 28,
+  "head_dim": null,
+  "attention_type": "sparse_stochastic",
+  "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": "/mloscratch/homes/navasard/sparse/mask_cache_8k_fixed_s1",
+  "sparse_mask_seed": 1,
+  "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
+  attention_type: sparse_stochastic
+  sparse_block_size: 128
+  sparse_num_random_blocks: 31
+  sparse_mask_pool_size: 1
+  sparse_max_seq_length: 8192
+
+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
+--------------------------------------------------------------------------------
+================================================================================================================================================================
+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)
+Starting training...
+Switching attention mode for step 1: sparse_stochastic (configured sparse attention)
+Step 10/44000 | Loss: 12.0407 | LR: 6.67e-06 | 10352.4ms/step | 1.9 samples/s
+Step 20/44000 | Loss: 11.9467 | LR: 1.33e-05 | 3275.8ms/step | 6.1 samples/s
+Step 30/44000 | Loss: 11.6121 | LR: 2.00e-05 | 3282.8ms/step | 6.1 samples/s
+Step 40/44000 | Loss: 11.0799 | LR: 2.67e-05 | 3276.7ms/step | 6.1 samples/s
+Step 50/44000 | Loss: 10.7621 | LR: 3.33e-05 | 3281.5ms/step | 6.1 samples/s
+Step 60/44000 | Loss: 10.5119 | LR: 4.00e-05 | 3278.9ms/step | 6.1 samples/s
+Step 70/44000 | Loss: 10.2858 | LR: 4.67e-05 | 3277.0ms/step | 6.1 samples/s
+Step 80/44000 | Loss: 9.9771 | LR: 5.33e-05 | 3281.0ms/step | 6.1 samples/s
+Step 90/44000 | Loss: 9.5920 | LR: 6.00e-05 | 3277.6ms/step | 6.1 samples/s
+Step 100/44000 | Loss: 9.1223 | LR: 6.67e-05 | 3288.2ms/step | 6.1 samples/s
+Step 110/44000 | Loss: 8.6026 | LR: 7.33e-05 | 3285.8ms/step | 6.1 samples/s
+Step 120/44000 | Loss: 8.1154 | LR: 8.00e-05 | 3292.5ms/step | 6.1 samples/s
+Step 130/44000 | Loss: 7.7540 | LR: 8.67e-05 | 3284.9ms/step | 6.1 samples/s
+Step 140/44000 | Loss: 7.5392 | LR: 9.33e-05 | 3293.1ms/step | 6.1 samples/s
+Step 150/44000 | Loss: 7.3258 | LR: 1.00e-04 | 3290.5ms/step | 6.1 samples/s
+Step 160/44000 | Loss: 7.2296 | LR: 1.07e-04 | 3308.4ms/step | 6.0 samples/s
+Step 170/44000 | Loss: 7.1082 | LR: 1.13e-04 | 3304.1ms/step | 6.1 samples/s
+Step 180/44000 | Loss: 7.0111 | LR: 1.20e-04 | 3308.9ms/step | 6.0 samples/s
+Step 190/44000 | Loss: 6.8883 | LR: 1.27e-04 | 3306.2ms/step | 6.0 samples/s
+Step 200/44000 | Loss: 6.8449 | LR: 1.33e-04 | 3315.0ms/step | 6.0 samples/s
+Step 210/44000 | Loss: 6.7680 | LR: 1.40e-04 | 3307.8ms/step | 6.0 samples/s
+Step 220/44000 | Loss: 6.6906 | LR: 1.47e-04 | 3313.7ms/step | 6.0 samples/s
+Step 230/44000 | Loss: 6.6219 | LR: 1.53e-04 | 3307.3ms/step | 6.0 samples/s
+Step 240/44000 | Loss: 6.5194 | LR: 1.60e-04 | 3315.4ms/step | 6.0 samples/s
+Step 250/44000 | Loss: 6.4661 | LR: 1.67e-04 | 3308.4ms/step | 6.0 samples/s

epfl_final/logs/l8k365m_resampled_seed43.log

@@ -0,0 +1,204 @@
+/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(
+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-20260802_195007-r5no0m0m
+wandb: Run `wandb offline` to turn off syncing.
+wandb: Syncing run 365m resampled 8k nrb31 seed43
+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/r5no0m0m
+New run: Checkpoints will be saved to /mloscratch/homes/navasard/sparse/checkpoints/l8k365m_resampled_seed43/run_20260802_194724
+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/l8k365m_resampled_seed43/run_20260802_194724",
+  "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": 10,
+  "seed": 43,
+  "verbose": true,
+  "checkpoint_activations": false,
+  "strategy": "ddp",
+  "mixed_precision": "bf16",
+  "wandb_enabled": true,
+  "wandb_project": "stochastic_sparsity",
+  "wandb_name": "365m resampled 8k nrb31 seed43",
+  "wandb_tags": [
+    "qwen3",
+    "365m",
+    "8k",
+    "scaleladder",
+    "resampled",
+    "seed43",
+    "seedvariance"
+  ],
+  "hidden_size": 768,
+  "intermediate_size": 3072,
+  "num_attention_heads": 6,
+  "num_key_value_heads": 2,
+  "num_hidden_layers": 28,
+  "head_dim": null,
+  "attention_type": "sparse_stochastic",
+  "sparse_block_size": 128,
+  "sparse_num_random_blocks": 31,
+  "sparse_mask_pool_size": 10000,
+  "sparse_val_dense_compare": false,
+  "sparse_layer_indices": null,
+  "sparse_mask_cache_dir": "/mloscratch/homes/navasard/sparse/mask_cache_32",
+  "sparse_mask_seed": 0,
+  "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
+  attention_type: sparse_stochastic
+  sparse_block_size: 128
+  sparse_num_random_blocks: 31
+  sparse_mask_pool_size: 10000
+  sparse_max_seq_length: 8192
+
+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
+--------------------------------------------------------------------------------
+================================================================================================================================================================
+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)
+Starting training...
+Switching attention mode for step 1: sparse_stochastic (configured sparse attention)
+Step 10/44000 | Loss: 12.0457 | LR: 6.67e-06 | 4433.4ms/step | 4.5 samples/s
+Step 20/44000 | Loss: 11.9514 | LR: 1.33e-05 | 3334.6ms/step | 6.0 samples/s
+Step 30/44000 | Loss: 11.6280 | LR: 2.00e-05 | 3300.2ms/step | 6.1 samples/s
+Step 40/44000 | Loss: 11.0986 | LR: 2.67e-05 | 3294.7ms/step | 6.1 samples/s
+Step 50/44000 | Loss: 10.7667 | LR: 3.33e-05 | 3287.7ms/step | 6.1 samples/s
+Step 60/44000 | Loss: 10.5338 | LR: 4.00e-05 | 3284.6ms/step | 6.1 samples/s
+Step 70/44000 | Loss: 10.2889 | LR: 4.67e-05 | 3285.4ms/step | 6.1 samples/s
+Step 80/44000 | Loss: 9.9837 | LR: 5.33e-05 | 3285.5ms/step | 6.1 samples/s
+Step 90/44000 | Loss: 9.6224 | LR: 6.00e-05 | 3297.3ms/step | 6.1 samples/s
+Step 100/44000 | Loss: 9.1339 | LR: 6.67e-05 | 3292.8ms/step | 6.1 samples/s
+Step 110/44000 | Loss: 8.6280 | LR: 7.33e-05 | 3293.0ms/step | 6.1 samples/s
+Step 120/44000 | Loss: 8.1296 | LR: 8.00e-05 | 3294.0ms/step | 6.1 samples/s
+Step 130/44000 | Loss: 7.7676 | LR: 8.67e-05 | 3297.3ms/step | 6.1 samples/s
+Step 140/44000 | Loss: 7.4919 | LR: 9.33e-05 | 3295.7ms/step | 6.1 samples/s
+Step 150/44000 | Loss: 7.3330 | LR: 1.00e-04 | 3297.8ms/step | 6.1 samples/s
+Step 160/44000 | Loss: 7.1982 | LR: 1.07e-04 | 3298.9ms/step | 6.1 samples/s
+Step 170/44000 | Loss: 7.1078 | LR: 1.13e-04 | 3299.9ms/step | 6.1 samples/s
+Step 180/44000 | Loss: 7.0317 | LR: 1.20e-04 | 3305.9ms/step | 6.0 samples/s
+Step 190/44000 | Loss: 6.9290 | LR: 1.27e-04 | 3302.7ms/step | 6.1 samples/s
+Step 200/44000 | Loss: 6.8254 | LR: 1.33e-04 | 3301.5ms/step | 6.1 samples/s
+Step 210/44000 | Loss: 6.7316 | LR: 1.40e-04 | 3300.7ms/step | 6.1 samples/s
+Step 220/44000 | Loss: 6.6872 | LR: 1.47e-04 | 3302.3ms/step | 6.1 samples/s
+Step 230/44000 | Loss: 6.5993 | LR: 1.53e-04 | 3309.1ms/step | 6.0 samples/s
+Step 240/44000 | Loss: 6.5255 | LR: 1.60e-04 | 3301.5ms/step | 6.1 samples/s
+Step 250/44000 | Loss: 6.4576 | LR: 1.67e-04 | 3301.8ms/step | 6.1 samples/s

epfl_final/logs/l8k365m_resampled_seed44.log

@@ -0,0 +1,204 @@
+/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(
+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-20260802_195008-c1ekdpek
+wandb: Run `wandb offline` to turn off syncing.
+wandb: Syncing run 365m resampled 8k nrb31 seed44
+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/c1ekdpek
+New run: Checkpoints will be saved to /mloscratch/homes/navasard/sparse/checkpoints/l8k365m_resampled_seed44/run_20260802_194724
+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/l8k365m_resampled_seed44/run_20260802_194724",
+  "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": 10,
+  "seed": 44,
+  "verbose": true,
+  "checkpoint_activations": false,
+  "strategy": "ddp",
+  "mixed_precision": "bf16",
+  "wandb_enabled": true,
+  "wandb_project": "stochastic_sparsity",
+  "wandb_name": "365m resampled 8k nrb31 seed44",
+  "wandb_tags": [
+    "qwen3",
+    "365m",
+    "8k",
+    "scaleladder",
+    "resampled",
+    "seed44",
+    "seedvariance"
+  ],
+  "hidden_size": 768,
+  "intermediate_size": 3072,
+  "num_attention_heads": 6,
+  "num_key_value_heads": 2,
+  "num_hidden_layers": 28,
+  "head_dim": null,
+  "attention_type": "sparse_stochastic",
+  "sparse_block_size": 128,
+  "sparse_num_random_blocks": 31,
+  "sparse_mask_pool_size": 10000,
+  "sparse_val_dense_compare": false,
+  "sparse_layer_indices": null,
+  "sparse_mask_cache_dir": "/mloscratch/homes/navasard/sparse/mask_cache_32",
+  "sparse_mask_seed": 0,
+  "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
+  attention_type: sparse_stochastic
+  sparse_block_size: 128
+  sparse_num_random_blocks: 31
+  sparse_mask_pool_size: 10000
+  sparse_max_seq_length: 8192
+
+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
+--------------------------------------------------------------------------------
+================================================================================================================================================================
+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)
+Starting training...
+Switching attention mode for step 1: sparse_stochastic (configured sparse attention)
+Step 10/44000 | Loss: 12.0415 | LR: 6.67e-06 | 4317.3ms/step | 4.6 samples/s
+Step 20/44000 | Loss: 11.9497 | LR: 1.33e-05 | 3281.8ms/step | 6.1 samples/s
+Step 30/44000 | Loss: 11.6134 | LR: 2.00e-05 | 3286.0ms/step | 6.1 samples/s
+Step 40/44000 | Loss: 11.0815 | LR: 2.67e-05 | 3285.1ms/step | 6.1 samples/s
+Step 50/44000 | Loss: 10.7627 | LR: 3.33e-05 | 3284.4ms/step | 6.1 samples/s
+Step 60/44000 | Loss: 10.5120 | LR: 4.00e-05 | 3284.5ms/step | 6.1 samples/s
+Step 70/44000 | Loss: 10.2859 | LR: 4.67e-05 | 3289.5ms/step | 6.1 samples/s
+Step 80/44000 | Loss: 9.9781 | LR: 5.33e-05 | 3298.4ms/step | 6.1 samples/s
+Step 90/44000 | Loss: 9.5917 | LR: 6.00e-05 | 3298.3ms/step | 6.1 samples/s
+Step 100/44000 | Loss: 9.1188 | LR: 6.67e-05 | 3296.9ms/step | 6.1 samples/s
+Step 110/44000 | Loss: 8.5995 | LR: 7.33e-05 | 3296.2ms/step | 6.1 samples/s
+Step 120/44000 | Loss: 8.1135 | LR: 8.00e-05 | 3295.5ms/step | 6.1 samples/s
+Step 130/44000 | Loss: 7.7521 | LR: 8.67e-05 | 3299.3ms/step | 6.1 samples/s
+Step 140/44000 | Loss: 7.5390 | LR: 9.33e-05 | 3296.5ms/step | 6.1 samples/s
+Step 150/44000 | Loss: 7.3258 | LR: 1.00e-04 | 3298.0ms/step | 6.1 samples/s
+Step 160/44000 | Loss: 7.2192 | LR: 1.07e-04 | 3300.1ms/step | 6.1 samples/s
+Step 170/44000 | Loss: 7.0858 | LR: 1.13e-04 | 3303.0ms/step | 6.1 samples/s
+Step 180/44000 | Loss: 6.9943 | LR: 1.20e-04 | 3304.5ms/step | 6.1 samples/s
+Step 190/44000 | Loss: 6.8804 | LR: 1.27e-04 | 3305.2ms/step | 6.1 samples/s
+Step 200/44000 | Loss: 6.8362 | LR: 1.33e-04 | 3304.2ms/step | 6.1 samples/s
+Step 210/44000 | Loss: 6.7623 | LR: 1.40e-04 | 3306.0ms/step | 6.0 samples/s
+Step 220/44000 | Loss: 6.6758 | LR: 1.47e-04 | 3308.6ms/step | 6.0 samples/s
+Step 230/44000 | Loss: 6.6009 | LR: 1.53e-04 | 3307.5ms/step | 6.0 samples/s
+Step 240/44000 | Loss: 6.4949 | LR: 1.60e-04 | 3303.8ms/step | 6.1 samples/s
+Step 250/44000 | Loss: 6.4545 | LR: 1.67e-04 | 3302.2ms/step | 6.1 samples/s

epfl_final/logs/launch_fr.out

@@ -0,0 +1,5 @@
+launched l8k365m_resampled_seed43 on gpu0 pid=557 log=/mloscratch/homes/navasard/sparse/logs/l8k365m_resampled_seed43.log
+launched l8k365m_resampled_seed44 on gpu1 pid=623 log=/mloscratch/homes/navasard/sparse/logs/l8k365m_resampled_seed44.log
+launched l8k365m_fixed_s1_seed43 on gpu2 pid=691 log=/mloscratch/homes/navasard/sparse/logs/l8k365m_fixed_s1_seed43.log
+launched l8k365m_fixed_s1_seed44 on gpu3 pid=758 log=/mloscratch/homes/navasard/sparse/logs/l8k365m_fixed_s1_seed44.log
+all launched

epfl_final/logs/watch_4b.log

@@ -0,0 +1,4 @@
+Wed Jul 15 11:18:06 PM UTC 2026 watcher[4b] start; waiting for /mloscratch/homes/navasard/sparse/mask_cache_8k_36L/.BUILD_DONE
+Wed Jul 15 11:20:06 PM UTC 2026 watcher[4b] cache ready
+Wed Jul 15 11:20:09 PM UTC 2026 watcher[4b] LAUNCHING nproc=4 port=29632 -> /mloscratch/homes/navasard/sparse/logs/l8k4b_resampled.log
+Wed Jul 15 11:20:09 PM UTC 2026 watcher[4b] launched pid 190929

epfl_final/logs/watch_mixed.log

@@ -0,0 +1,4 @@
+Wed Jul 15 11:18:22 PM UTC 2026 watcher[mixed] start; waiting for /mloscratch/homes/navasard/sparse/mask_cache_8k_mixed/.BUILD_DONE
+Wed Jul 15 11:40:52 PM UTC 2026 watcher[mixed] cache ready
+Wed Jul 15 11:40:55 PM UTC 2026 watcher[mixed] LAUNCHING nproc=3 port=29631 -> /mloscratch/homes/navasard/sparse/logs/l8k1_2b_mixed.log
+Wed Jul 15 11:40:55 PM UTC 2026 watcher[mixed] launched pid 123526