backup: training logs before server expiry

.gitattributes

@@ -902,3 +902,10 @@ checkpoints/star-coconut-k10L10-fixed-128_n1000_qwen3-0.6b-base_coconut_lr1e-4_n
 checkpoints/star-coconut-k10L10-fixed-128_n1000_qwen3-0.6b-base_coconut_lr1e-4_no-reset_epoch-334c5641/train_state_30 filter=lfs diff=lfs merge=lfs -text
 checkpoints/star-coconut-k10L10-fixed-128_n1000_qwen3-0.6b-base_coconut_lr1e-4_no-reset_epoch-334c5641/checkpoint_31 filter=lfs diff=lfs merge=lfs -text
 checkpoints/star-coconut-k10L10-fixed-128_n1000_qwen3-0.6b-base_coconut_lr1e-4_no-reset_epoch-334c5641/train_state_31 filter=lfs diff=lfs merge=lfs -text
+logs_backup/star_atc_8gpu.log filter=lfs diff=lfs merge=lfs -text
+logs_backup/star_atc_k1L100_adaptive.log filter=lfs diff=lfs merge=lfs -text
+logs_backup/star_coconut_4gpu.log filter=lfs diff=lfs merge=lfs -text
+logs_backup/star_coconut_k10L10_fixed.log filter=lfs diff=lfs merge=lfs -text
+logs_backup/star_nocot_dc_k1L100.log filter=lfs diff=lfs merge=lfs -text
+logs_backup/star_nocot_nocurr_4gpu.log filter=lfs diff=lfs merge=lfs -text
+logs_backup/star_nocot_nocurr_k1L100.log filter=lfs diff=lfs merge=lfs -text

logs_backup/star_atc_8gpu.log

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

logs_backup/star_atc_k1L100_adaptive.log

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

logs_backup/star_coconut_4gpu.log

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

logs_backup/star_coconut_k10L10_fixed.log

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

logs_backup/star_nocot_dc_k1L100.log

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

logs_backup/star_nocot_nocurr_4gpu.log

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

logs_backup/star_nocot_nocurr_k1L100.log

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

logs_backup/star_smoke.log

@@ -0,0 +1,119 @@
+Config: {'project': 'thoughtformer', 'name': 'star-smoke', 'only_eval': False, 'method': 'thoughtformer', 'data_curriculum': True, 'thought_curriculum': False, 'staging': 'two_pointers', 'init_thought_stage': 0, 'init_data_stage': 1, 'patience': 5, 'staging_threshold': 0.9, 'c_thought': 0, 'max_latent_stage': 0, 'save_only_improve': False, 'model_id': 'Qwen/Qwen3-0.6B-Base', 'load_model_path': 'None', 'seed': 0, 'resume': 0, 'bf16': False, 'train_path': 'data/star_k10_L10_1000_train.json', 'val_path': 'data/star_k10_L10_50_valid.json', 'reset_optimizer': False, 'lr_schedule': 'cosine', 'batch_size_training': 4, 'eval_only_trained': False, 'eval_batch_size': 8, 'eval_every': 1, 'debug': False, 'gradient_accumulation_steps': 1, 'num_epochs': 2, 'lr': 0.0001, 'weight_decay': 0.01, 'group': 'smoke', 'tags': ['star', 'k10', 'L10', 'no_cot'], 'run_type': 'pilot'}
+/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/c10d_logger.py:83: UserWarning: barrier(): using the device under current context. You can specify `device_id` in `init_process_group` to mute this warning.
+  return func(*args, **kwargs)
+[rank0]:[W528 07:02:16.327385348 ProcessGroupNCCL.cpp:5324] Guessing device ID based on global rank. This can cause a hang if rank to GPU mapping is heterogeneous. You can specify device_id in init_process_group()
+
+Running FSDP on rank = 0, world size = 1
+/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py:478: UserWarning: FSDP is switching to use `NO_SHARD` instead of ShardingStrategy.FULL_SHARD since the world size is 1.
+  _init_core_state(
+FullyShardedDataParallel(
+  (_fsdp_wrapped_module): Coconut(
+    (base_causallm): Qwen3ForCausalLM(
+      (model): Qwen3Model(
+        (embed_tokens): Embedding(151672, 1024)
+        (layers): ModuleList(
+          (0-27): 28 x FullyShardedDataParallel(
+            (_fsdp_wrapped_module): Qwen3DecoderLayer(
+              (self_attn): Qwen3Attention(
+                (q_proj): Linear(in_features=1024, out_features=2048, bias=False)
+                (k_proj): Linear(in_features=1024, out_features=1024, bias=False)
+                (v_proj): Linear(in_features=1024, out_features=1024, bias=False)
+                (o_proj): Linear(in_features=2048, out_features=1024, bias=False)
+                (q_norm): Qwen3RMSNorm((128,), eps=1e-06)
+                (k_norm): Qwen3RMSNorm((128,), eps=1e-06)
+              )
+              (mlp): Qwen3MLP(
+                (gate_proj): Linear(in_features=1024, out_features=3072, bias=False)
+                (up_proj): Linear(in_features=1024, out_features=3072, bias=False)
+                (down_proj): Linear(in_features=3072, out_features=1024, bias=False)
+                (act_fn): SiLUActivation()
+              )
+              (input_layernorm): Qwen3RMSNorm((1024,), eps=1e-06)
+              (post_attention_layernorm): Qwen3RMSNorm((1024,), eps=1e-06)
+            )
+          )
+        )
+        (norm): Qwen3RMSNorm((1024,), eps=1e-06)
+        (rotary_emb): Qwen3RotaryEmbedding()
+      )
+      (lm_head): Linear(in_features=1024, out_features=151936, bias=False)
+    )
+    (embedding): Embedding(151672, 1024)
+  )
+)
+
+
+wandb: [wandb.login()] Loaded credentials for https://api.wandb.ai from /home/ubuntu/.netrc.
+wandb: Currently logged in as: seyedparsa to https://api.wandb.ai. Use `wandb login --relogin` to force relogin
+wandb: setting up run kbjpqcuk
+wandb: Tracking run with wandb version 0.25.1
+wandb: Run data is saved locally in /home/ubuntu/thoughtformer/wandb/run-20260528_070253-kbjpqcuk
+wandb: Run `wandb offline` to turn off syncing.
+wandb: Syncing run star-smoke_n1000_qwen3-0.6b-base_thoughtformer_lr1e-4_no-reset_dc_no-tc_two_pointers-bdad6776
+wandb: ⭐️ View project at https://wandb.ai/seyedparsa/thoughtformer
+wandb: 🚀 View run at https://wandb.ai/seyedparsa/thoughtformer/runs/kbjpqcuk
+
+============================================================
+EPOCH 0/2 (thought_stage=0, data_stage=1)
+============================================================
+  thought_stage=0, c_thought=0, max_difficulty=1
+
+
+
+
+
+  LR schedule: cosine, warmup=125 steps, total=5000 steps (max_steps/epoch=2500)
+
+/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/autograd/graph.py:882: UserWarning: The AccumulateGrad node's stream does not match the stream of the node that produced the incoming gradient. This may incur unnecessary synchronization and break CUDA graph capture if the AccumulateGrad node's stream is the default stream. This mismatch is caused by an AccumulateGrad node created prior to the current iteration being kept alive. This can happen if the autograd graph is still being kept alive by tensors such as the loss, or if you are using DDP, which will stash a reference to the node. To resolve the mismatch, delete all references to the autograd graph or ensure that DDP initialization is performed under the same stream as subsequent forwards. If the mismatch is intentional, you can use torch.autograd.graph.set_warn_on_accumulate_grad_stream_mismatch(False) to suppress this warning. (Triggered internally at /pytorch/torch/csrc/autograd/input_buffer.cpp:240.)
+  return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
+rank 0 | gpu 0 | alloc=6.35 GB | reserved=27.88 GB | max=21.84 GB
+
+/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/c10d_logger.py:83: UserWarning: barrier(): using the device under current context. You can specify `device_id` in `init_process_group` to mute this warning.
+  return func(*args, **kwargs)
+/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py:821: FutureWarning: FSDP.state_dict_type() and FSDP.set_state_dict_type() are being deprecated. Please use APIs, get_state_dict() and set_state_dict(), which can support different parallelisms, FSDP1, FSDP2, DDP. API doc: https://pytorch.org/docs/stable/distributed.checkpoint.html#torch.distributed.checkpoint.state_dict.get_state_dict .Tutorial: https://pytorch.org/tutorials/recipes/distributed_checkpoint_recipe.html .
+  prev_state_dict_settings = FullyShardedDataParallel.set_state_dict_type(
+/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/utils/_contextlib.py:124: UserWarning: When using ``NO_SHARD`` for ``ShardingStrategy``, full_state_dict will be returned.
+  return func(*args, **kwargs)
+/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/fsdp/_optim_utils.py:1172: UserWarning: `_get_pg_default_device` will be deprecated, it only stays for backward-compatibility reason. If you need to find a device for object collectives, please use `_get_object_coll_device`. If you need to query the device types supported by group, please use `_device_capability(group)`. 
+  device = _get_pg_default_device(group)
+/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py:828: FutureWarning: FSDP.state_dict_type() and FSDP.set_state_dict_type() are being deprecated. Please use APIs, get_state_dict() and set_state_dict(), which can support different parallelisms, FSDP1, FSDP2, DDP. API doc: https://pytorch.org/docs/stable/distributed.checkpoint.html#torch.distributed.checkpoint.state_dict.get_state_dict .Tutorial: https://pytorch.org/tutorials/recipes/distributed_checkpoint_recipe.html .
+  FullyShardedDataParallel.set_state_dict_type(
+saving model + train state.
+Warning: failed to upload checkpoint_1 to HF: (Request ID: Root=1-6a17e919-1dce5dca1d2ddb050a1a0a4f;c5ffd1eb-6364-4939-8651-3bc2bb1d3fe2)
+
+403 Forbidden: You need to setup automatic credit recharge in order to upload more data. You can do so at /settings/billing..
+Cannot access content at: https://huggingface.co/seyedparsa/thoughtformer-checkpoints.git/info/lfs/objects/batch.
+Make sure your token has the correct permissions.
+Warning: failed to upload train_state_1 to HF: (Request ID: Root=1-6a17e91d-3685080228a0e1cb2da99c67;496f22d9-c5af-4e21-bb1b-4dc6beb52df0)
+
+403 Forbidden: You need to setup automatic credit recharge in order to upload more data. You can do so at /settings/billing..
+Cannot access content at: https://huggingface.co/seyedparsa/thoughtformer-checkpoints.git/info/lfs/objects/batch.
+Make sure your token has the correct permissions.
+Keeping all local checkpoints because HF upload failed; disk fallback until HF recovers.
+eval loss 0.896042138338089
+
+[WRONG] Example 0 (diff=3)
+  Expected : 'Patrick'
+  Extracted: 'Nancy'
+  Generated: '### Nancy'
+
+[CORRECT] Example 1 (diff=2)
+  Expected : 'Roman'
+  Extracted: 'Roman'
+  Generated: '### Roman'
+
+[WRONG] Example 2 (diff=2)
+  Expected : 'Harold'
+  Extracted: 'Stella'
+  Generated: '### Stella'
+
+[WRONG] Example 3 (diff=2)
+  Expected : 'Larry'
+  Extracted: 'Liam'
+  Generated: '### Liam'
+
+[CORRECT] Example 4 (diff=2)
+  Expected : 'Israel'
+  Extracted: 'Israel'
+  Generated: '### Israel'
+

logs_backup/star_sweep_bs128_lr1e-4.log

@@ -0,0 +1,296 @@
+Config: {'project': 'thoughtformer', 'name': 'star-sweep-bs128-lr1e-4', 'only_eval': False, 'method': 'thoughtformer', 'data_curriculum': True, 'thought_curriculum': True, 'staging': 'adaptive', 'init_thought_stage': 1, 'init_data_stage': 1, 'patience': 5, 'staging_threshold': 0.9, 'c_thought': 1, 'max_latent_stage': 10, 'uniform_prob': 0.1, 'save_only_improve': False, 'model_id': 'Qwen/Qwen3-0.6B-Base', 'load_model_path': 'None', 'seed': 0, 'resume': 0, 'bf16': False, 'train_path': 'data/star_k10_L10_1000_train.json', 'val_path': 'data/star_k10_L10_50_valid.json', 'reset_optimizer': False, 'lr_schedule': 'cosine', 'batch_size_training': 128, 'eval_only_trained': False, 'eval_batch_size': 8, 'eval_every': 1, 'debug': False, 'gradient_accumulation_steps': 1, 'num_epochs': 30, 'lr': 0.0001, 'weight_decay': 0.01, 'group': 'star_bs_lr_sweep', 'tags': ['star', 'k10', 'L10', 'atc', 'adaptive', 'lr1e-5', 'best', 'budget_aug'], 'run_type': 'pilot'}
+/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/c10d_logger.py:83: UserWarning: barrier(): using the device under current context. You can specify `device_id` in `init_process_group` to mute this warning.
+  return func(*args, **kwargs)
+[rank0]:[W530 03:31:55.296639842 ProcessGroupNCCL.cpp:5324] Guessing device ID based on global rank. This can cause a hang if rank to GPU mapping is heterogeneous. You can specify device_id in init_process_group()
+
+Running FSDP on rank = 0, world size = 1
+/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py:478: UserWarning: FSDP is switching to use `NO_SHARD` instead of ShardingStrategy.FULL_SHARD since the world size is 1.
+  _init_core_state(
+FullyShardedDataParallel(
+  (_fsdp_wrapped_module): Coconut(
+    (base_causallm): Qwen3ForCausalLM(
+      (model): Qwen3Model(
+        (embed_tokens): Embedding(151672, 1024)
+        (layers): ModuleList(
+          (0-27): 28 x FullyShardedDataParallel(
+            (_fsdp_wrapped_module): Qwen3DecoderLayer(
+              (self_attn): Qwen3Attention(
+                (q_proj): Linear(in_features=1024, out_features=2048, bias=False)
+                (k_proj): Linear(in_features=1024, out_features=1024, bias=False)
+                (v_proj): Linear(in_features=1024, out_features=1024, bias=False)
+                (o_proj): Linear(in_features=2048, out_features=1024, bias=False)
+                (q_norm): Qwen3RMSNorm((128,), eps=1e-06)
+                (k_norm): Qwen3RMSNorm((128,), eps=1e-06)
+              )
+              (mlp): Qwen3MLP(
+                (gate_proj): Linear(in_features=1024, out_features=3072, bias=False)
+                (up_proj): Linear(in_features=1024, out_features=3072, bias=False)
+                (down_proj): Linear(in_features=3072, out_features=1024, bias=False)
+                (act_fn): SiLUActivation()
+              )
+              (input_layernorm): Qwen3RMSNorm((1024,), eps=1e-06)
+              (post_attention_layernorm): Qwen3RMSNorm((1024,), eps=1e-06)
+            )
+          )
+        )
+        (norm): Qwen3RMSNorm((1024,), eps=1e-06)
+        (rotary_emb): Qwen3RotaryEmbedding()
+      )
+      (lm_head): Linear(in_features=1024, out_features=151936, bias=False)
+    )
+    (embedding): Embedding(151672, 1024)
+  )
+)
+
+
+wandb: [wandb.login()] Loaded credentials for https://api.wandb.ai from /home/ubuntu/.netrc.
+wandb: Currently logged in as: seyedparsa to https://api.wandb.ai. Use `wandb login --relogin` to force relogin
+wandb: Tracking run with wandb version 0.25.1
+wandb: Run data is saved locally in /home/ubuntu/thoughtformer/wandb/run-20260530_033222-2vtfied7
+wandb: Run `wandb offline` to turn off syncing.
+wandb: Syncing run star-sweep-bs128-lr1e-4_n1000_qwen3-0.6b-base_thoughtformer_lr1e-4_no-reset_dc_tc_adaptive-6c6fcb83
+wandb: ⭐️ View project at https://wandb.ai/seyedparsa/thoughtformer
+wandb: 🚀 View run at https://wandb.ai/seyedparsa/thoughtformer/runs/2vtfied7
+
+============================================================
+EPOCH 0/30 (thought_stage=1, data_stage=1)
+============================================================
+  thought_stage=1, c_thought=1, max_difficulty=1
+
+
+
+
+
+  LR schedule: cosine, warmup=3 steps, total=2340 steps (max_steps/epoch=78)
+
+wandb: WARNING Serializing object of type str that is 1676601 bytes
+Traceback (most recent call last):
+  File "/home/ubuntu/thoughtformer/run.py", line 1194, in <module>
+    main()
+  File "/home/ubuntu/thoughtformer/run.py", line 937, in main
+    outputs = parallel_model(**batch)
+              ^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+    return forward_call(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 857, in forward
+    output = self._fsdp_wrapped_module(*args, **kwargs)
+             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+    return forward_call(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/coconut.py", line 101, in forward
+    outputs = self.base_causallm(
+              ^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+    return forward_call(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/utils/generic.py", line 857, in wrapper
+    output = func(self, *args, **kwargs)
+             ^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 492, in forward
+    outputs: BaseModelOutputWithPast = self.model(
+                                       ^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+    return forward_call(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/utils/generic.py", line 931, in wrapper
+    output = func(self, *args, **kwargs)
+             ^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/utils/output_capturing.py", line 248, in wrapper
+    outputs = func(self, *args, **kwargs)
+              ^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 424, in forward
+    hidden_states = decoder_layer(
+                    ^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+    return forward_call(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 857, in forward
+    output = self._fsdp_wrapped_module(*args, **kwargs)
+             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/modeling_layers.py", line 93, in __call__
+    return super().__call__(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1884, in _call_impl
+    return inner()
+           ^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1832, in inner
+    result = forward_call(*args, **kwargs)
+             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 332, in forward
+    hidden_states = self.mlp(hidden_states)
+                    ^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+    return forward_call(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 82, in forward
+    down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+                               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+    return forward_call(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/activations.py", line 103, in forward
+    return nn.functional.silu(input)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/functional.py", line 2397, in silu
+    return torch._C._nn.silu(input)
+           ^^^^^^^^^^^^^^^^^^^^^^^^
+torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 1.19 GiB. GPU 0 has a total capacity of 139.80 GiB of which 583.06 MiB is free. Including non-PyTorch memory, this process has 139.22 GiB memory in use. Of the allocated memory 136.99 GiB is allocated by PyTorch, and 1.03 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://docs.pytorch.org/docs/stable/notes/cuda.html#optimizing-memory-usage-with-pytorch-cuda-alloc-conf)
+[rank0]: Traceback (most recent call last):
+[rank0]:   File "/home/ubuntu/thoughtformer/run.py", line 1194, in <module>
+[rank0]:     main()
+[rank0]:   File "/home/ubuntu/thoughtformer/run.py", line 937, in main
+[rank0]:     outputs = parallel_model(**batch)
+[rank0]:               ^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+[rank0]:     return forward_call(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 857, in forward
+[rank0]:     output = self._fsdp_wrapped_module(*args, **kwargs)
+[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+[rank0]:     return forward_call(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/coconut.py", line 101, in forward
+[rank0]:     outputs = self.base_causallm(
+[rank0]:               ^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+[rank0]:     return forward_call(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/utils/generic.py", line 857, in wrapper
+[rank0]:     output = func(self, *args, **kwargs)
+[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 492, in forward
+[rank0]:     outputs: BaseModelOutputWithPast = self.model(
+[rank0]:                                        ^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+[rank0]:     return forward_call(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/utils/generic.py", line 931, in wrapper
+[rank0]:     output = func(self, *args, **kwargs)
+[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/utils/output_capturing.py", line 248, in wrapper
+[rank0]:     outputs = func(self, *args, **kwargs)
+[rank0]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 424, in forward
+[rank0]:     hidden_states = decoder_layer(
+[rank0]:                     ^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+[rank0]:     return forward_call(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 857, in forward
+[rank0]:     output = self._fsdp_wrapped_module(*args, **kwargs)
+[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/modeling_layers.py", line 93, in __call__
+[rank0]:     return super().__call__(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1884, in _call_impl
+[rank0]:     return inner()
+[rank0]:            ^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1832, in inner
+[rank0]:     result = forward_call(*args, **kwargs)
+[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 332, in forward
+[rank0]:     hidden_states = self.mlp(hidden_states)
+[rank0]:                     ^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+[rank0]:     return forward_call(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 82, in forward
+[rank0]:     down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+[rank0]:                                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+[rank0]:     return forward_call(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/activations.py", line 103, in forward
+[rank0]:     return nn.functional.silu(input)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/functional.py", line 2397, in silu
+[rank0]:     return torch._C._nn.silu(input)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]: torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 1.19 GiB. GPU 0 has a total capacity of 139.80 GiB of which 583.06 MiB is free. Including non-PyTorch memory, this process has 139.22 GiB memory in use. Of the allocated memory 136.99 GiB is allocated by PyTorch, and 1.03 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://docs.pytorch.org/docs/stable/notes/cuda.html#optimizing-memory-usage-with-pytorch-cuda-alloc-conf)
+wandb: 
+wandb: 🚀 View run star-sweep-bs128-lr1e-4_n1000_qwen3-0.6b-base_thoughtformer_lr1e-4_no-reset_dc_tc_adaptive-6c6fcb83 at: https://wandb.ai/seyedparsa/thoughtformer/runs/2vtfied7
+wandb: Find logs at: wandb/run-20260530_033222-2vtfied7/logs
+E0530 03:32:33.929000 465640 torch/distributed/elastic/multiprocessing/api.py:988] failed (exitcode: 1) local_rank: 0 (pid: 466218) of binary: /home/ubuntu/thoughtformer/.venv/bin/python3
+Traceback (most recent call last):
+  File "/home/ubuntu/thoughtformer/.venv/bin/torchrun", line 8, in <module>
+    sys.exit(main())
+             ^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 367, in wrapper
+    return f(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/run.py", line 1016, in main
+    run(args)
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/run.py", line 1007, in run
+    elastic_launch(
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/launcher/api.py", line 191, in __call__
+    return launch_agent(self._config, self._entrypoint, list(args))
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/launcher/api.py", line 371, in launch_agent
+    raise ChildFailedError(
+torch.distributed.elastic.multiprocessing.errors.ChildFailedError: 
+============================================================
+run.py FAILED
+------------------------------------------------------------
+Failures:
+  <NO_OTHER_FAILURES>
+------------------------------------------------------------
+Root Cause (first observed failure):
+[0]:
+  time      : 2026-05-30_03:32:33
+  host      : ip-172-31-10-226.us-east-2.compute.internal
+  rank      : 0 (local_rank: 0)
+  exitcode  : 1 (pid: 466218) 
+  error_file: <N/A>
+  traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
+============================================================

logs_backup/star_sweep_bs128_lr1e-5.log

@@ -0,0 +1,296 @@
+Config: {'project': 'thoughtformer', 'name': 'star-sweep-bs128-lr1e-5', 'only_eval': False, 'method': 'thoughtformer', 'data_curriculum': True, 'thought_curriculum': True, 'staging': 'adaptive', 'init_thought_stage': 1, 'init_data_stage': 1, 'patience': 5, 'staging_threshold': 0.9, 'c_thought': 1, 'max_latent_stage': 10, 'uniform_prob': 0.1, 'save_only_improve': False, 'model_id': 'Qwen/Qwen3-0.6B-Base', 'load_model_path': 'None', 'seed': 0, 'resume': 0, 'bf16': False, 'train_path': 'data/star_k10_L10_1000_train.json', 'val_path': 'data/star_k10_L10_50_valid.json', 'reset_optimizer': False, 'lr_schedule': 'cosine', 'batch_size_training': 128, 'eval_only_trained': False, 'eval_batch_size': 8, 'eval_every': 1, 'debug': False, 'gradient_accumulation_steps': 1, 'num_epochs': 30, 'lr': 1e-05, 'weight_decay': 0.01, 'group': 'star_bs_lr_sweep', 'tags': ['star', 'k10', 'L10', 'atc', 'adaptive', 'lr1e-5', 'best', 'budget_aug'], 'run_type': 'pilot'}
+/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/c10d_logger.py:83: UserWarning: barrier(): using the device under current context. You can specify `device_id` in `init_process_group` to mute this warning.
+  return func(*args, **kwargs)
+[rank0]:[W530 03:31:53.943261833 ProcessGroupNCCL.cpp:5324] Guessing device ID based on global rank. This can cause a hang if rank to GPU mapping is heterogeneous. You can specify device_id in init_process_group()
+
+Running FSDP on rank = 0, world size = 1
+/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py:478: UserWarning: FSDP is switching to use `NO_SHARD` instead of ShardingStrategy.FULL_SHARD since the world size is 1.
+  _init_core_state(
+FullyShardedDataParallel(
+  (_fsdp_wrapped_module): Coconut(
+    (base_causallm): Qwen3ForCausalLM(
+      (model): Qwen3Model(
+        (embed_tokens): Embedding(151672, 1024)
+        (layers): ModuleList(
+          (0-27): 28 x FullyShardedDataParallel(
+            (_fsdp_wrapped_module): Qwen3DecoderLayer(
+              (self_attn): Qwen3Attention(
+                (q_proj): Linear(in_features=1024, out_features=2048, bias=False)
+                (k_proj): Linear(in_features=1024, out_features=1024, bias=False)
+                (v_proj): Linear(in_features=1024, out_features=1024, bias=False)
+                (o_proj): Linear(in_features=2048, out_features=1024, bias=False)
+                (q_norm): Qwen3RMSNorm((128,), eps=1e-06)
+                (k_norm): Qwen3RMSNorm((128,), eps=1e-06)
+              )
+              (mlp): Qwen3MLP(
+                (gate_proj): Linear(in_features=1024, out_features=3072, bias=False)
+                (up_proj): Linear(in_features=1024, out_features=3072, bias=False)
+                (down_proj): Linear(in_features=3072, out_features=1024, bias=False)
+                (act_fn): SiLUActivation()
+              )
+              (input_layernorm): Qwen3RMSNorm((1024,), eps=1e-06)
+              (post_attention_layernorm): Qwen3RMSNorm((1024,), eps=1e-06)
+            )
+          )
+        )
+        (norm): Qwen3RMSNorm((1024,), eps=1e-06)
+        (rotary_emb): Qwen3RotaryEmbedding()
+      )
+      (lm_head): Linear(in_features=1024, out_features=151936, bias=False)
+    )
+    (embedding): Embedding(151672, 1024)
+  )
+)
+
+
+wandb: [wandb.login()] Loaded credentials for https://api.wandb.ai from /home/ubuntu/.netrc.
+wandb: Currently logged in as: seyedparsa to https://api.wandb.ai. Use `wandb login --relogin` to force relogin
+wandb: Tracking run with wandb version 0.25.1
+wandb: Run data is saved locally in /home/ubuntu/thoughtformer/wandb/run-20260530_033220-efpd3uen
+wandb: Run `wandb offline` to turn off syncing.
+wandb: Syncing run star-sweep-bs128-lr1e-5_n1000_qwen3-0.6b-base_thoughtformer_lr1e-5_no-reset_dc_tc_adaptive-46d2d3b7
+wandb: ⭐️ View project at https://wandb.ai/seyedparsa/thoughtformer
+wandb: 🚀 View run at https://wandb.ai/seyedparsa/thoughtformer/runs/efpd3uen
+
+============================================================
+EPOCH 0/30 (thought_stage=1, data_stage=1)
+============================================================
+  thought_stage=1, c_thought=1, max_difficulty=1
+
+
+
+
+
+  LR schedule: cosine, warmup=3 steps, total=2340 steps (max_steps/epoch=78)
+
+wandb: WARNING Serializing object of type str that is 1676601 bytes
+Traceback (most recent call last):
+  File "/home/ubuntu/thoughtformer/run.py", line 1194, in <module>
+    main()
+  File "/home/ubuntu/thoughtformer/run.py", line 937, in main
+    outputs = parallel_model(**batch)
+              ^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+    return forward_call(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 857, in forward
+    output = self._fsdp_wrapped_module(*args, **kwargs)
+             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+    return forward_call(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/coconut.py", line 101, in forward
+    outputs = self.base_causallm(
+              ^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+    return forward_call(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/utils/generic.py", line 857, in wrapper
+    output = func(self, *args, **kwargs)
+             ^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 492, in forward
+    outputs: BaseModelOutputWithPast = self.model(
+                                       ^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+    return forward_call(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/utils/generic.py", line 931, in wrapper
+    output = func(self, *args, **kwargs)
+             ^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/utils/output_capturing.py", line 248, in wrapper
+    outputs = func(self, *args, **kwargs)
+              ^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 424, in forward
+    hidden_states = decoder_layer(
+                    ^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+    return forward_call(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 857, in forward
+    output = self._fsdp_wrapped_module(*args, **kwargs)
+             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/modeling_layers.py", line 93, in __call__
+    return super().__call__(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1884, in _call_impl
+    return inner()
+           ^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1832, in inner
+    result = forward_call(*args, **kwargs)
+             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 332, in forward
+    hidden_states = self.mlp(hidden_states)
+                    ^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+    return forward_call(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 82, in forward
+    down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+                               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+    return forward_call(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/activations.py", line 103, in forward
+    return nn.functional.silu(input)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/functional.py", line 2397, in silu
+    return torch._C._nn.silu(input)
+           ^^^^^^^^^^^^^^^^^^^^^^^^
+torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 1.19 GiB. GPU 0 has a total capacity of 139.80 GiB of which 583.06 MiB is free. Including non-PyTorch memory, this process has 139.22 GiB memory in use. Of the allocated memory 136.99 GiB is allocated by PyTorch, and 1.03 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://docs.pytorch.org/docs/stable/notes/cuda.html#optimizing-memory-usage-with-pytorch-cuda-alloc-conf)
+[rank0]: Traceback (most recent call last):
+[rank0]:   File "/home/ubuntu/thoughtformer/run.py", line 1194, in <module>
+[rank0]:     main()
+[rank0]:   File "/home/ubuntu/thoughtformer/run.py", line 937, in main
+[rank0]:     outputs = parallel_model(**batch)
+[rank0]:               ^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+[rank0]:     return forward_call(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 857, in forward
+[rank0]:     output = self._fsdp_wrapped_module(*args, **kwargs)
+[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+[rank0]:     return forward_call(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/coconut.py", line 101, in forward
+[rank0]:     outputs = self.base_causallm(
+[rank0]:               ^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+[rank0]:     return forward_call(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/utils/generic.py", line 857, in wrapper
+[rank0]:     output = func(self, *args, **kwargs)
+[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 492, in forward
+[rank0]:     outputs: BaseModelOutputWithPast = self.model(
+[rank0]:                                        ^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+[rank0]:     return forward_call(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/utils/generic.py", line 931, in wrapper
+[rank0]:     output = func(self, *args, **kwargs)
+[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/utils/output_capturing.py", line 248, in wrapper
+[rank0]:     outputs = func(self, *args, **kwargs)
+[rank0]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 424, in forward
+[rank0]:     hidden_states = decoder_layer(
+[rank0]:                     ^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+[rank0]:     return forward_call(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 857, in forward
+[rank0]:     output = self._fsdp_wrapped_module(*args, **kwargs)
+[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/modeling_layers.py", line 93, in __call__
+[rank0]:     return super().__call__(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1884, in _call_impl
+[rank0]:     return inner()
+[rank0]:            ^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1832, in inner
+[rank0]:     result = forward_call(*args, **kwargs)
+[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 332, in forward
+[rank0]:     hidden_states = self.mlp(hidden_states)
+[rank0]:                     ^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+[rank0]:     return forward_call(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 82, in forward
+[rank0]:     down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+[rank0]:                                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+[rank0]:     return forward_call(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/activations.py", line 103, in forward
+[rank0]:     return nn.functional.silu(input)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/functional.py", line 2397, in silu
+[rank0]:     return torch._C._nn.silu(input)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]: torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 1.19 GiB. GPU 0 has a total capacity of 139.80 GiB of which 583.06 MiB is free. Including non-PyTorch memory, this process has 139.22 GiB memory in use. Of the allocated memory 136.99 GiB is allocated by PyTorch, and 1.03 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://docs.pytorch.org/docs/stable/notes/cuda.html#optimizing-memory-usage-with-pytorch-cuda-alloc-conf)
+wandb: 
+wandb: 🚀 View run star-sweep-bs128-lr1e-5_n1000_qwen3-0.6b-base_thoughtformer_lr1e-5_no-reset_dc_tc_adaptive-46d2d3b7 at: https://wandb.ai/seyedparsa/thoughtformer/runs/efpd3uen
+wandb: Find logs at: wandb/run-20260530_033220-efpd3uen/logs
+E0530 03:32:31.760000 464943 torch/distributed/elastic/multiprocessing/api.py:988] failed (exitcode: 1) local_rank: 0 (pid: 465570) of binary: /home/ubuntu/thoughtformer/.venv/bin/python3
+Traceback (most recent call last):
+  File "/home/ubuntu/thoughtformer/.venv/bin/torchrun", line 8, in <module>
+    sys.exit(main())
+             ^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 367, in wrapper
+    return f(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/run.py", line 1016, in main
+    run(args)
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/run.py", line 1007, in run
+    elastic_launch(
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/launcher/api.py", line 191, in __call__
+    return launch_agent(self._config, self._entrypoint, list(args))
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/launcher/api.py", line 371, in launch_agent
+    raise ChildFailedError(
+torch.distributed.elastic.multiprocessing.errors.ChildFailedError: 
+============================================================
+run.py FAILED
+------------------------------------------------------------
+Failures:
+  <NO_OTHER_FAILURES>
+------------------------------------------------------------
+Root Cause (first observed failure):
+[0]:
+  time      : 2026-05-30_03:32:31
+  host      : ip-172-31-10-226.us-east-2.compute.internal
+  rank      : 0 (local_rank: 0)
+  exitcode  : 1 (pid: 465570) 
+  error_file: <N/A>
+  traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
+============================================================

logs_backup/star_sweep_bs256_lr1e-4.log

@@ -0,0 +1,290 @@
+Config: {'project': 'thoughtformer', 'name': 'star-sweep-bs256-lr1e-4', 'only_eval': False, 'method': 'thoughtformer', 'data_curriculum': True, 'thought_curriculum': True, 'staging': 'adaptive', 'init_thought_stage': 1, 'init_data_stage': 1, 'patience': 5, 'staging_threshold': 0.9, 'c_thought': 1, 'max_latent_stage': 10, 'uniform_prob': 0.1, 'save_only_improve': False, 'model_id': 'Qwen/Qwen3-0.6B-Base', 'load_model_path': 'None', 'seed': 0, 'resume': 0, 'bf16': False, 'train_path': 'data/star_k10_L10_1000_train.json', 'val_path': 'data/star_k10_L10_50_valid.json', 'reset_optimizer': False, 'lr_schedule': 'cosine', 'batch_size_training': 256, 'eval_only_trained': False, 'eval_batch_size': 8, 'eval_every': 1, 'debug': False, 'gradient_accumulation_steps': 1, 'num_epochs': 30, 'lr': 0.0001, 'weight_decay': 0.01, 'group': 'star_bs_lr_sweep', 'tags': ['star', 'k10', 'L10', 'atc', 'adaptive', 'lr1e-5', 'best', 'budget_aug'], 'run_type': 'pilot'}
+/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/c10d_logger.py:83: UserWarning: barrier(): using the device under current context. You can specify `device_id` in `init_process_group` to mute this warning.
+  return func(*args, **kwargs)
+[rank0]:[W530 03:31:59.134090576 ProcessGroupNCCL.cpp:5324] Guessing device ID based on global rank. This can cause a hang if rank to GPU mapping is heterogeneous. You can specify device_id in init_process_group()
+
+Running FSDP on rank = 0, world size = 1
+/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py:478: UserWarning: FSDP is switching to use `NO_SHARD` instead of ShardingStrategy.FULL_SHARD since the world size is 1.
+  _init_core_state(
+FullyShardedDataParallel(
+  (_fsdp_wrapped_module): Coconut(
+    (base_causallm): Qwen3ForCausalLM(
+      (model): Qwen3Model(
+        (embed_tokens): Embedding(151672, 1024)
+        (layers): ModuleList(
+          (0-27): 28 x FullyShardedDataParallel(
+            (_fsdp_wrapped_module): Qwen3DecoderLayer(
+              (self_attn): Qwen3Attention(
+                (q_proj): Linear(in_features=1024, out_features=2048, bias=False)
+                (k_proj): Linear(in_features=1024, out_features=1024, bias=False)
+                (v_proj): Linear(in_features=1024, out_features=1024, bias=False)
+                (o_proj): Linear(in_features=2048, out_features=1024, bias=False)
+                (q_norm): Qwen3RMSNorm((128,), eps=1e-06)
+                (k_norm): Qwen3RMSNorm((128,), eps=1e-06)
+              )
+              (mlp): Qwen3MLP(
+                (gate_proj): Linear(in_features=1024, out_features=3072, bias=False)
+                (up_proj): Linear(in_features=1024, out_features=3072, bias=False)
+                (down_proj): Linear(in_features=3072, out_features=1024, bias=False)
+                (act_fn): SiLUActivation()
+              )
+              (input_layernorm): Qwen3RMSNorm((1024,), eps=1e-06)
+              (post_attention_layernorm): Qwen3RMSNorm((1024,), eps=1e-06)
+            )
+          )
+        )
+        (norm): Qwen3RMSNorm((1024,), eps=1e-06)
+        (rotary_emb): Qwen3RotaryEmbedding()
+      )
+      (lm_head): Linear(in_features=1024, out_features=151936, bias=False)
+    )
+    (embedding): Embedding(151672, 1024)
+  )
+)
+
+
+wandb: [wandb.login()] Loaded credentials for https://api.wandb.ai from /home/ubuntu/.netrc.
+wandb: Currently logged in as: seyedparsa to https://api.wandb.ai. Use `wandb login --relogin` to force relogin
+wandb: Tracking run with wandb version 0.25.1
+wandb: Run data is saved locally in /home/ubuntu/thoughtformer/wandb/run-20260530_033226-wj48phx5
+wandb: Run `wandb offline` to turn off syncing.
+wandb: Syncing run star-sweep-bs256-lr1e-4_n1000_qwen3-0.6b-base_thoughtformer_lr1e-4_no-reset_dc_tc_adaptive-ff592230
+wandb: ⭐️ View project at https://wandb.ai/seyedparsa/thoughtformer
+wandb: 🚀 View run at https://wandb.ai/seyedparsa/thoughtformer/runs/wj48phx5
+
+============================================================
+EPOCH 0/30 (thought_stage=1, data_stage=1)
+============================================================
+  thought_stage=1, c_thought=1, max_difficulty=1
+
+
+
+
+
+  LR schedule: cosine, warmup=1 steps, total=1170 steps (max_steps/epoch=39)
+
+wandb: WARNING Serializing object of type str that is 3353666 bytes
+Traceback (most recent call last):
+  File "/home/ubuntu/thoughtformer/run.py", line 1194, in <module>
+    main()
+  File "/home/ubuntu/thoughtformer/run.py", line 937, in main
+    outputs = parallel_model(**batch)
+              ^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+    return forward_call(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 857, in forward
+    output = self._fsdp_wrapped_module(*args, **kwargs)
+             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+    return forward_call(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/coconut.py", line 101, in forward
+    outputs = self.base_causallm(
+              ^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+    return forward_call(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/utils/generic.py", line 857, in wrapper
+    output = func(self, *args, **kwargs)
+             ^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 492, in forward
+    outputs: BaseModelOutputWithPast = self.model(
+                                       ^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+    return forward_call(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/utils/generic.py", line 931, in wrapper
+    output = func(self, *args, **kwargs)
+             ^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/utils/output_capturing.py", line 248, in wrapper
+    outputs = func(self, *args, **kwargs)
+              ^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 424, in forward
+    hidden_states = decoder_layer(
+                    ^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+    return forward_call(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 857, in forward
+    output = self._fsdp_wrapped_module(*args, **kwargs)
+             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/modeling_layers.py", line 93, in __call__
+    return super().__call__(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1884, in _call_impl
+    return inner()
+           ^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1832, in inner
+    result = forward_call(*args, **kwargs)
+             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 332, in forward
+    hidden_states = self.mlp(hidden_states)
+                    ^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+    return forward_call(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 82, in forward
+    down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+                                                                ^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+    return forward_call(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/linear.py", line 134, in forward
+    return F.linear(input, self.weight, self.bias)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 2.38 GiB. GPU 0 has a total capacity of 139.80 GiB of which 1.40 GiB is free. Including non-PyTorch memory, this process has 138.39 GiB memory in use. Of the allocated memory 135.99 GiB is allocated by PyTorch, and 1.20 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://docs.pytorch.org/docs/stable/notes/cuda.html#optimizing-memory-usage-with-pytorch-cuda-alloc-conf)
+[rank0]: Traceback (most recent call last):
+[rank0]:   File "/home/ubuntu/thoughtformer/run.py", line 1194, in <module>
+[rank0]:     main()
+[rank0]:   File "/home/ubuntu/thoughtformer/run.py", line 937, in main
+[rank0]:     outputs = parallel_model(**batch)
+[rank0]:               ^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+[rank0]:     return forward_call(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 857, in forward
+[rank0]:     output = self._fsdp_wrapped_module(*args, **kwargs)
+[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+[rank0]:     return forward_call(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/coconut.py", line 101, in forward
+[rank0]:     outputs = self.base_causallm(
+[rank0]:               ^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+[rank0]:     return forward_call(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/utils/generic.py", line 857, in wrapper
+[rank0]:     output = func(self, *args, **kwargs)
+[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 492, in forward
+[rank0]:     outputs: BaseModelOutputWithPast = self.model(
+[rank0]:                                        ^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+[rank0]:     return forward_call(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/utils/generic.py", line 931, in wrapper
+[rank0]:     output = func(self, *args, **kwargs)
+[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/utils/output_capturing.py", line 248, in wrapper
+[rank0]:     outputs = func(self, *args, **kwargs)
+[rank0]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 424, in forward
+[rank0]:     hidden_states = decoder_layer(
+[rank0]:                     ^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+[rank0]:     return forward_call(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 857, in forward
+[rank0]:     output = self._fsdp_wrapped_module(*args, **kwargs)
+[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/modeling_layers.py", line 93, in __call__
+[rank0]:     return super().__call__(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1884, in _call_impl
+[rank0]:     return inner()
+[rank0]:            ^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1832, in inner
+[rank0]:     result = forward_call(*args, **kwargs)
+[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 332, in forward
+[rank0]:     hidden_states = self.mlp(hidden_states)
+[rank0]:                     ^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+[rank0]:     return forward_call(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 82, in forward
+[rank0]:     down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+[rank0]:                                                                 ^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+[rank0]:     return forward_call(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/linear.py", line 134, in forward
+[rank0]:     return F.linear(input, self.weight, self.bias)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]: torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 2.38 GiB. GPU 0 has a total capacity of 139.80 GiB of which 1.40 GiB is free. Including non-PyTorch memory, this process has 138.39 GiB memory in use. Of the allocated memory 135.99 GiB is allocated by PyTorch, and 1.20 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://docs.pytorch.org/docs/stable/notes/cuda.html#optimizing-memory-usage-with-pytorch-cuda-alloc-conf)
+wandb: 
+wandb: 🚀 View run star-sweep-bs256-lr1e-4_n1000_qwen3-0.6b-base_thoughtformer_lr1e-4_no-reset_dc_tc_adaptive-ff592230 at: https://wandb.ai/seyedparsa/thoughtformer/runs/wj48phx5
+wandb: Find logs at: wandb/run-20260530_033226-wj48phx5/logs
+E0530 03:32:40.146000 466881 torch/distributed/elastic/multiprocessing/api.py:988] failed (exitcode: 1) local_rank: 0 (pid: 467529) of binary: /home/ubuntu/thoughtformer/.venv/bin/python3
+Traceback (most recent call last):
+  File "/home/ubuntu/thoughtformer/.venv/bin/torchrun", line 8, in <module>
+    sys.exit(main())
+             ^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 367, in wrapper
+    return f(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/run.py", line 1016, in main
+    run(args)
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/run.py", line 1007, in run
+    elastic_launch(
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/launcher/api.py", line 191, in __call__
+    return launch_agent(self._config, self._entrypoint, list(args))
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/launcher/api.py", line 371, in launch_agent
+    raise ChildFailedError(
+torch.distributed.elastic.multiprocessing.errors.ChildFailedError: 
+============================================================
+run.py FAILED
+------------------------------------------------------------
+Failures:
+  <NO_OTHER_FAILURES>
+------------------------------------------------------------
+Root Cause (first observed failure):
+[0]:
+  time      : 2026-05-30_03:32:40
+  host      : ip-172-31-10-226.us-east-2.compute.internal
+  rank      : 0 (local_rank: 0)
+  exitcode  : 1 (pid: 467529) 
+  error_file: <N/A>
+  traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
+============================================================

logs_backup/star_sweep_bs256_lr1e-5.log

@@ -0,0 +1,290 @@
+Config: {'project': 'thoughtformer', 'name': 'star-sweep-bs256-lr1e-5', 'only_eval': False, 'method': 'thoughtformer', 'data_curriculum': True, 'thought_curriculum': True, 'staging': 'adaptive', 'init_thought_stage': 1, 'init_data_stage': 1, 'patience': 5, 'staging_threshold': 0.9, 'c_thought': 1, 'max_latent_stage': 10, 'uniform_prob': 0.1, 'save_only_improve': False, 'model_id': 'Qwen/Qwen3-0.6B-Base', 'load_model_path': 'None', 'seed': 0, 'resume': 0, 'bf16': False, 'train_path': 'data/star_k10_L10_1000_train.json', 'val_path': 'data/star_k10_L10_50_valid.json', 'reset_optimizer': False, 'lr_schedule': 'cosine', 'batch_size_training': 256, 'eval_only_trained': False, 'eval_batch_size': 8, 'eval_every': 1, 'debug': False, 'gradient_accumulation_steps': 1, 'num_epochs': 30, 'lr': 1e-05, 'weight_decay': 0.01, 'group': 'star_bs_lr_sweep', 'tags': ['star', 'k10', 'L10', 'atc', 'adaptive', 'lr1e-5', 'best', 'budget_aug'], 'run_type': 'pilot'}
+/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/c10d_logger.py:83: UserWarning: barrier(): using the device under current context. You can specify `device_id` in `init_process_group` to mute this warning.
+  return func(*args, **kwargs)
+[rank0]:[W530 03:31:57.089584891 ProcessGroupNCCL.cpp:5324] Guessing device ID based on global rank. This can cause a hang if rank to GPU mapping is heterogeneous. You can specify device_id in init_process_group()
+
+Running FSDP on rank = 0, world size = 1
+/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py:478: UserWarning: FSDP is switching to use `NO_SHARD` instead of ShardingStrategy.FULL_SHARD since the world size is 1.
+  _init_core_state(
+FullyShardedDataParallel(
+  (_fsdp_wrapped_module): Coconut(
+    (base_causallm): Qwen3ForCausalLM(
+      (model): Qwen3Model(
+        (embed_tokens): Embedding(151672, 1024)
+        (layers): ModuleList(
+          (0-27): 28 x FullyShardedDataParallel(
+            (_fsdp_wrapped_module): Qwen3DecoderLayer(
+              (self_attn): Qwen3Attention(
+                (q_proj): Linear(in_features=1024, out_features=2048, bias=False)
+                (k_proj): Linear(in_features=1024, out_features=1024, bias=False)
+                (v_proj): Linear(in_features=1024, out_features=1024, bias=False)
+                (o_proj): Linear(in_features=2048, out_features=1024, bias=False)
+                (q_norm): Qwen3RMSNorm((128,), eps=1e-06)
+                (k_norm): Qwen3RMSNorm((128,), eps=1e-06)
+              )
+              (mlp): Qwen3MLP(
+                (gate_proj): Linear(in_features=1024, out_features=3072, bias=False)
+                (up_proj): Linear(in_features=1024, out_features=3072, bias=False)
+                (down_proj): Linear(in_features=3072, out_features=1024, bias=False)
+                (act_fn): SiLUActivation()
+              )
+              (input_layernorm): Qwen3RMSNorm((1024,), eps=1e-06)
+              (post_attention_layernorm): Qwen3RMSNorm((1024,), eps=1e-06)
+            )
+          )
+        )
+        (norm): Qwen3RMSNorm((1024,), eps=1e-06)
+        (rotary_emb): Qwen3RotaryEmbedding()
+      )
+      (lm_head): Linear(in_features=1024, out_features=151936, bias=False)
+    )
+    (embedding): Embedding(151672, 1024)
+  )
+)
+
+
+wandb: [wandb.login()] Loaded credentials for https://api.wandb.ai from /home/ubuntu/.netrc.
+wandb: Currently logged in as: seyedparsa to https://api.wandb.ai. Use `wandb login --relogin` to force relogin
+wandb: Tracking run with wandb version 0.25.1
+wandb: Run data is saved locally in /home/ubuntu/thoughtformer/wandb/run-20260530_033224-tjrts1mq
+wandb: Run `wandb offline` to turn off syncing.
+wandb: Syncing run star-sweep-bs256-lr1e-5_n1000_qwen3-0.6b-base_thoughtformer_lr1e-5_no-reset_dc_tc_adaptive-a9a8c070
+wandb: ⭐️ View project at https://wandb.ai/seyedparsa/thoughtformer
+wandb: 🚀 View run at https://wandb.ai/seyedparsa/thoughtformer/runs/tjrts1mq
+
+============================================================
+EPOCH 0/30 (thought_stage=1, data_stage=1)
+============================================================
+  thought_stage=1, c_thought=1, max_difficulty=1
+
+
+
+
+
+  LR schedule: cosine, warmup=1 steps, total=1170 steps (max_steps/epoch=39)
+
+wandb: WARNING Serializing object of type str that is 3353666 bytes
+Traceback (most recent call last):
+  File "/home/ubuntu/thoughtformer/run.py", line 1194, in <module>
+    main()
+  File "/home/ubuntu/thoughtformer/run.py", line 937, in main
+    outputs = parallel_model(**batch)
+              ^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+    return forward_call(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 857, in forward
+    output = self._fsdp_wrapped_module(*args, **kwargs)
+             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+    return forward_call(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/coconut.py", line 101, in forward
+    outputs = self.base_causallm(
+              ^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+    return forward_call(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/utils/generic.py", line 857, in wrapper
+    output = func(self, *args, **kwargs)
+             ^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 492, in forward
+    outputs: BaseModelOutputWithPast = self.model(
+                                       ^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+    return forward_call(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/utils/generic.py", line 931, in wrapper
+    output = func(self, *args, **kwargs)
+             ^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/utils/output_capturing.py", line 248, in wrapper
+    outputs = func(self, *args, **kwargs)
+              ^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 424, in forward
+    hidden_states = decoder_layer(
+                    ^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+    return forward_call(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 857, in forward
+    output = self._fsdp_wrapped_module(*args, **kwargs)
+             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/modeling_layers.py", line 93, in __call__
+    return super().__call__(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1884, in _call_impl
+    return inner()
+           ^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1832, in inner
+    result = forward_call(*args, **kwargs)
+             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 332, in forward
+    hidden_states = self.mlp(hidden_states)
+                    ^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+    return forward_call(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 82, in forward
+    down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+                                                                ^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+    return forward_call(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/linear.py", line 134, in forward
+    return F.linear(input, self.weight, self.bias)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 2.38 GiB. GPU 0 has a total capacity of 139.80 GiB of which 1.40 GiB is free. Including non-PyTorch memory, this process has 138.39 GiB memory in use. Of the allocated memory 135.99 GiB is allocated by PyTorch, and 1.20 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://docs.pytorch.org/docs/stable/notes/cuda.html#optimizing-memory-usage-with-pytorch-cuda-alloc-conf)
+[rank0]: Traceback (most recent call last):
+[rank0]:   File "/home/ubuntu/thoughtformer/run.py", line 1194, in <module>
+[rank0]:     main()
+[rank0]:   File "/home/ubuntu/thoughtformer/run.py", line 937, in main
+[rank0]:     outputs = parallel_model(**batch)
+[rank0]:               ^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+[rank0]:     return forward_call(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 857, in forward
+[rank0]:     output = self._fsdp_wrapped_module(*args, **kwargs)
+[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+[rank0]:     return forward_call(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/coconut.py", line 101, in forward
+[rank0]:     outputs = self.base_causallm(
+[rank0]:               ^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+[rank0]:     return forward_call(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/utils/generic.py", line 857, in wrapper
+[rank0]:     output = func(self, *args, **kwargs)
+[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 492, in forward
+[rank0]:     outputs: BaseModelOutputWithPast = self.model(
+[rank0]:                                        ^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+[rank0]:     return forward_call(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/utils/generic.py", line 931, in wrapper
+[rank0]:     output = func(self, *args, **kwargs)
+[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/utils/output_capturing.py", line 248, in wrapper
+[rank0]:     outputs = func(self, *args, **kwargs)
+[rank0]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 424, in forward
+[rank0]:     hidden_states = decoder_layer(
+[rank0]:                     ^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+[rank0]:     return forward_call(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 857, in forward
+[rank0]:     output = self._fsdp_wrapped_module(*args, **kwargs)
+[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/modeling_layers.py", line 93, in __call__
+[rank0]:     return super().__call__(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1884, in _call_impl
+[rank0]:     return inner()
+[rank0]:            ^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1832, in inner
+[rank0]:     result = forward_call(*args, **kwargs)
+[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 332, in forward
+[rank0]:     hidden_states = self.mlp(hidden_states)
+[rank0]:                     ^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+[rank0]:     return forward_call(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 82, in forward
+[rank0]:     down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+[rank0]:                                                                 ^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+[rank0]:     return forward_call(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/linear.py", line 134, in forward
+[rank0]:     return F.linear(input, self.weight, self.bias)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]: torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 2.38 GiB. GPU 0 has a total capacity of 139.80 GiB of which 1.40 GiB is free. Including non-PyTorch memory, this process has 138.39 GiB memory in use. Of the allocated memory 135.99 GiB is allocated by PyTorch, and 1.20 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://docs.pytorch.org/docs/stable/notes/cuda.html#optimizing-memory-usage-with-pytorch-cuda-alloc-conf)
+wandb: 
+wandb: 🚀 View run star-sweep-bs256-lr1e-5_n1000_qwen3-0.6b-base_thoughtformer_lr1e-5_no-reset_dc_tc_adaptive-a9a8c070 at: https://wandb.ai/seyedparsa/thoughtformer/runs/tjrts1mq
+wandb: Find logs at: wandb/run-20260530_033224-tjrts1mq/logs
+E0530 03:32:37.865000 466226 torch/distributed/elastic/multiprocessing/api.py:988] failed (exitcode: 1) local_rank: 0 (pid: 466874) of binary: /home/ubuntu/thoughtformer/.venv/bin/python3
+Traceback (most recent call last):
+  File "/home/ubuntu/thoughtformer/.venv/bin/torchrun", line 8, in <module>
+    sys.exit(main())
+             ^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 367, in wrapper
+    return f(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/run.py", line 1016, in main
+    run(args)
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/run.py", line 1007, in run
+    elastic_launch(
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/launcher/api.py", line 191, in __call__
+    return launch_agent(self._config, self._entrypoint, list(args))
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/launcher/api.py", line 371, in launch_agent
+    raise ChildFailedError(
+torch.distributed.elastic.multiprocessing.errors.ChildFailedError: 
+============================================================
+run.py FAILED
+------------------------------------------------------------
+Failures:
+  <NO_OTHER_FAILURES>
+------------------------------------------------------------
+Root Cause (first observed failure):
+[0]:
+  time      : 2026-05-30_03:32:37
+  host      : ip-172-31-10-226.us-east-2.compute.internal
+  rank      : 0 (local_rank: 0)
+  exitcode  : 1 (pid: 466874) 
+  error_file: <N/A>
+  traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
+============================================================

logs_backup/star_sweep_bs32_lr1e-4.log

@@ -0,0 +1,152 @@
+Config: {'project': 'thoughtformer', 'name': 'star-sweep-bs32-lr1e-4', 'only_eval': False, 'method': 'thoughtformer', 'data_curriculum': True, 'thought_curriculum': True, 'staging': 'adaptive', 'init_thought_stage': 1, 'init_data_stage': 1, 'patience': 5, 'staging_threshold': 0.9, 'c_thought': 1, 'max_latent_stage': 10, 'uniform_prob': 0.1, 'save_only_improve': False, 'model_id': 'Qwen/Qwen3-0.6B-Base', 'load_model_path': 'None', 'seed': 0, 'resume': 0, 'bf16': False, 'train_path': 'data/star_k10_L10_1000_train.json', 'val_path': 'data/star_k10_L10_50_valid.json', 'reset_optimizer': False, 'lr_schedule': 'cosine', 'batch_size_training': 32, 'eval_only_trained': False, 'eval_batch_size': 8, 'eval_every': 1, 'debug': False, 'gradient_accumulation_steps': 1, 'num_epochs': 30, 'lr': 0.0001, 'weight_decay': 0.01, 'group': 'star_bs_lr_sweep', 'tags': ['star', 'k10', 'L10', 'atc', 'adaptive', 'lr1e-5', 'best', 'budget_aug'], 'run_type': 'pilot'}
+/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/c10d_logger.py:83: UserWarning: barrier(): using the device under current context. You can specify `device_id` in `init_process_group` to mute this warning.
+  return func(*args, **kwargs)
+[rank0]:[W530 03:31:47.851491967 ProcessGroupNCCL.cpp:5324] Guessing device ID based on global rank. This can cause a hang if rank to GPU mapping is heterogeneous. You can specify device_id in init_process_group()
+
+Running FSDP on rank = 0, world size = 1
+/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py:478: UserWarning: FSDP is switching to use `NO_SHARD` instead of ShardingStrategy.FULL_SHARD since the world size is 1.
+  _init_core_state(
+FullyShardedDataParallel(
+  (_fsdp_wrapped_module): Coconut(
+    (base_causallm): Qwen3ForCausalLM(
+      (model): Qwen3Model(
+        (embed_tokens): Embedding(151672, 1024)
+        (layers): ModuleList(
+          (0-27): 28 x FullyShardedDataParallel(
+            (_fsdp_wrapped_module): Qwen3DecoderLayer(
+              (self_attn): Qwen3Attention(
+                (q_proj): Linear(in_features=1024, out_features=2048, bias=False)
+                (k_proj): Linear(in_features=1024, out_features=1024, bias=False)
+                (v_proj): Linear(in_features=1024, out_features=1024, bias=False)
+                (o_proj): Linear(in_features=2048, out_features=1024, bias=False)
+                (q_norm): Qwen3RMSNorm((128,), eps=1e-06)
+                (k_norm): Qwen3RMSNorm((128,), eps=1e-06)
+              )
+              (mlp): Qwen3MLP(
+                (gate_proj): Linear(in_features=1024, out_features=3072, bias=False)
+                (up_proj): Linear(in_features=1024, out_features=3072, bias=False)
+                (down_proj): Linear(in_features=3072, out_features=1024, bias=False)
+                (act_fn): SiLUActivation()
+              )
+              (input_layernorm): Qwen3RMSNorm((1024,), eps=1e-06)
+              (post_attention_layernorm): Qwen3RMSNorm((1024,), eps=1e-06)
+            )
+          )
+        )
+        (norm): Qwen3RMSNorm((1024,), eps=1e-06)
+        (rotary_emb): Qwen3RotaryEmbedding()
+      )
+      (lm_head): Linear(in_features=1024, out_features=151936, bias=False)
+    )
+    (embedding): Embedding(151672, 1024)
+  )
+)
+
+
+wandb: [wandb.login()] Loaded credentials for https://api.wandb.ai from /home/ubuntu/.netrc.
+wandb: Currently logged in as: seyedparsa to https://api.wandb.ai. Use `wandb login --relogin` to force relogin
+wandb: Tracking run with wandb version 0.25.1
+wandb: Run data is saved locally in /home/ubuntu/thoughtformer/wandb/run-20260530_033214-gu4dqvfl
+wandb: Run `wandb offline` to turn off syncing.
+wandb: Syncing run star-sweep-bs32-lr1e-4_n1000_qwen3-0.6b-base_thoughtformer_lr1e-4_no-reset_dc_tc_adaptive-d2591099
+wandb: ⭐️ View project at https://wandb.ai/seyedparsa/thoughtformer
+wandb: 🚀 View run at https://wandb.ai/seyedparsa/thoughtformer/runs/gu4dqvfl
+
+============================================================
+EPOCH 0/30 (thought_stage=1, data_stage=1)
+============================================================
+  thought_stage=1, c_thought=1, max_difficulty=1
+
+
+
+
+
+  LR schedule: cosine, warmup=15 steps, total=9360 steps (max_steps/epoch=312)
+
+wandb: WARNING Serializing object of type str that is 419185 bytes
+Traceback (most recent call last):
+  File "/home/ubuntu/thoughtformer/run.py", line 1194, in <module>
+    main()
+  File "/home/ubuntu/thoughtformer/run.py", line 937, in main
+    outputs = parallel_model(**batch)
+              ^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+    return forward_call(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 857, in forward
+    output = self._fsdp_wrapped_module(*args, **kwargs)
+             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+    return forward_call(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/coconut.py", line 239, in forward
+    logits = torch.cat(logits, dim=-2)
+             ^^^^^^^^^^^^^^^^^^^^^^^^^
+torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 14.77 GiB. GPU 0 has a total capacity of 139.80 GiB of which 6.10 GiB is free. Including non-PyTorch memory, this process has 133.69 GiB memory in use. Of the allocated memory 131.74 GiB is allocated by PyTorch, and 767.15 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://docs.pytorch.org/docs/stable/notes/cuda.html#optimizing-memory-usage-with-pytorch-cuda-alloc-conf)
+[rank0]: Traceback (most recent call last):
+[rank0]:   File "/home/ubuntu/thoughtformer/run.py", line 1194, in <module>
+[rank0]:     main()
+[rank0]:   File "/home/ubuntu/thoughtformer/run.py", line 937, in main
+[rank0]:     outputs = parallel_model(**batch)
+[rank0]:               ^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+[rank0]:     return forward_call(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 857, in forward
+[rank0]:     output = self._fsdp_wrapped_module(*args, **kwargs)
+[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+[rank0]:     return forward_call(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/coconut.py", line 239, in forward
+[rank0]:     logits = torch.cat(logits, dim=-2)
+[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]: torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 14.77 GiB. GPU 0 has a total capacity of 139.80 GiB of which 6.10 GiB is free. Including non-PyTorch memory, this process has 133.69 GiB memory in use. Of the allocated memory 131.74 GiB is allocated by PyTorch, and 767.15 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://docs.pytorch.org/docs/stable/notes/cuda.html#optimizing-memory-usage-with-pytorch-cuda-alloc-conf)
+wandb: 
+wandb: 🚀 View run star-sweep-bs32-lr1e-4_n1000_qwen3-0.6b-base_thoughtformer_lr1e-4_no-reset_dc_tc_adaptive-d2591099 at: https://wandb.ai/seyedparsa/thoughtformer/runs/gu4dqvfl
+wandb: Find logs at: wandb/run-20260530_033214-gu4dqvfl/logs
+E0530 03:32:24.440000 464181 torch/distributed/elastic/multiprocessing/api.py:988] failed (exitcode: 1) local_rank: 0 (pid: 464254) of binary: /home/ubuntu/thoughtformer/.venv/bin/python3
+Traceback (most recent call last):
+  File "/home/ubuntu/thoughtformer/.venv/bin/torchrun", line 8, in <module>
+    sys.exit(main())
+             ^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 367, in wrapper
+    return f(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/run.py", line 1016, in main
+    run(args)
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/run.py", line 1007, in run
+    elastic_launch(
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/launcher/api.py", line 191, in __call__
+    return launch_agent(self._config, self._entrypoint, list(args))
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/launcher/api.py", line 371, in launch_agent
+    raise ChildFailedError(
+torch.distributed.elastic.multiprocessing.errors.ChildFailedError: 
+============================================================
+run.py FAILED
+------------------------------------------------------------
+Failures:
+  <NO_OTHER_FAILURES>
+------------------------------------------------------------
+Root Cause (first observed failure):
+[0]:
+  time      : 2026-05-30_03:32:24
+  host      : ip-172-31-10-226.us-east-2.compute.internal
+  rank      : 0 (local_rank: 0)
+  exitcode  : 1 (pid: 464254) 
+  error_file: <N/A>
+  traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
+============================================================

logs_backup/star_sweep_bs32_lr1e-5.log

@@ -0,0 +1,152 @@
+Config: {'project': 'thoughtformer', 'name': 'star-sweep-bs32-lr1e-5', 'only_eval': False, 'method': 'thoughtformer', 'data_curriculum': True, 'thought_curriculum': True, 'staging': 'adaptive', 'init_thought_stage': 1, 'init_data_stage': 1, 'patience': 5, 'staging_threshold': 0.9, 'c_thought': 1, 'max_latent_stage': 10, 'uniform_prob': 0.1, 'save_only_improve': False, 'model_id': 'Qwen/Qwen3-0.6B-Base', 'load_model_path': 'None', 'seed': 0, 'resume': 0, 'bf16': False, 'train_path': 'data/star_k10_L10_1000_train.json', 'val_path': 'data/star_k10_L10_50_valid.json', 'reset_optimizer': False, 'lr_schedule': 'cosine', 'batch_size_training': 32, 'eval_only_trained': False, 'eval_batch_size': 8, 'eval_every': 1, 'debug': False, 'gradient_accumulation_steps': 1, 'num_epochs': 30, 'lr': 1e-05, 'weight_decay': 0.01, 'group': 'star_bs_lr_sweep', 'tags': ['star', 'k10', 'L10', 'atc', 'adaptive', 'lr1e-5', 'best', 'budget_aug'], 'run_type': 'pilot'}
+/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/c10d_logger.py:83: UserWarning: barrier(): using the device under current context. You can specify `device_id` in `init_process_group` to mute this warning.
+  return func(*args, **kwargs)
+[rank0]:[W530 03:31:45.823267663 ProcessGroupNCCL.cpp:5324] Guessing device ID based on global rank. This can cause a hang if rank to GPU mapping is heterogeneous. You can specify device_id in init_process_group()
+
+Running FSDP on rank = 0, world size = 1
+/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py:478: UserWarning: FSDP is switching to use `NO_SHARD` instead of ShardingStrategy.FULL_SHARD since the world size is 1.
+  _init_core_state(
+FullyShardedDataParallel(
+  (_fsdp_wrapped_module): Coconut(
+    (base_causallm): Qwen3ForCausalLM(
+      (model): Qwen3Model(
+        (embed_tokens): Embedding(151672, 1024)
+        (layers): ModuleList(
+          (0-27): 28 x FullyShardedDataParallel(
+            (_fsdp_wrapped_module): Qwen3DecoderLayer(
+              (self_attn): Qwen3Attention(
+                (q_proj): Linear(in_features=1024, out_features=2048, bias=False)
+                (k_proj): Linear(in_features=1024, out_features=1024, bias=False)
+                (v_proj): Linear(in_features=1024, out_features=1024, bias=False)
+                (o_proj): Linear(in_features=2048, out_features=1024, bias=False)
+                (q_norm): Qwen3RMSNorm((128,), eps=1e-06)
+                (k_norm): Qwen3RMSNorm((128,), eps=1e-06)
+              )
+              (mlp): Qwen3MLP(
+                (gate_proj): Linear(in_features=1024, out_features=3072, bias=False)
+                (up_proj): Linear(in_features=1024, out_features=3072, bias=False)
+                (down_proj): Linear(in_features=3072, out_features=1024, bias=False)
+                (act_fn): SiLUActivation()
+              )
+              (input_layernorm): Qwen3RMSNorm((1024,), eps=1e-06)
+              (post_attention_layernorm): Qwen3RMSNorm((1024,), eps=1e-06)
+            )
+          )
+        )
+        (norm): Qwen3RMSNorm((1024,), eps=1e-06)
+        (rotary_emb): Qwen3RotaryEmbedding()
+      )
+      (lm_head): Linear(in_features=1024, out_features=151936, bias=False)
+    )
+    (embedding): Embedding(151672, 1024)
+  )
+)
+
+
+wandb: [wandb.login()] Loaded credentials for https://api.wandb.ai from /home/ubuntu/.netrc.
+wandb: Currently logged in as: seyedparsa to https://api.wandb.ai. Use `wandb login --relogin` to force relogin
+wandb: Tracking run with wandb version 0.25.1
+wandb: Run data is saved locally in /home/ubuntu/thoughtformer/wandb/run-20260530_033212-jek5luxr
+wandb: Run `wandb offline` to turn off syncing.
+wandb: Syncing run star-sweep-bs32-lr1e-5_n1000_qwen3-0.6b-base_thoughtformer_lr1e-5_no-reset_dc_tc_adaptive-c3e91928
+wandb: ⭐️ View project at https://wandb.ai/seyedparsa/thoughtformer
+wandb: 🚀 View run at https://wandb.ai/seyedparsa/thoughtformer/runs/jek5luxr
+
+============================================================
+EPOCH 0/30 (thought_stage=1, data_stage=1)
+============================================================
+  thought_stage=1, c_thought=1, max_difficulty=1
+
+
+
+
+
+  LR schedule: cosine, warmup=15 steps, total=9360 steps (max_steps/epoch=312)
+
+wandb: WARNING Serializing object of type str that is 419185 bytes
+Traceback (most recent call last):
+  File "/home/ubuntu/thoughtformer/run.py", line 1194, in <module>
+    main()
+  File "/home/ubuntu/thoughtformer/run.py", line 937, in main
+    outputs = parallel_model(**batch)
+              ^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+    return forward_call(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 857, in forward
+    output = self._fsdp_wrapped_module(*args, **kwargs)
+             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+    return forward_call(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/coconut.py", line 239, in forward
+    logits = torch.cat(logits, dim=-2)
+             ^^^^^^^^^^^^^^^^^^^^^^^^^
+torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 14.77 GiB. GPU 0 has a total capacity of 139.80 GiB of which 6.10 GiB is free. Including non-PyTorch memory, this process has 133.69 GiB memory in use. Of the allocated memory 131.74 GiB is allocated by PyTorch, and 767.15 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://docs.pytorch.org/docs/stable/notes/cuda.html#optimizing-memory-usage-with-pytorch-cuda-alloc-conf)
+[rank0]: Traceback (most recent call last):
+[rank0]:   File "/home/ubuntu/thoughtformer/run.py", line 1194, in <module>
+[rank0]:     main()
+[rank0]:   File "/home/ubuntu/thoughtformer/run.py", line 937, in main
+[rank0]:     outputs = parallel_model(**batch)
+[rank0]:               ^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+[rank0]:     return forward_call(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 857, in forward
+[rank0]:     output = self._fsdp_wrapped_module(*args, **kwargs)
+[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+[rank0]:     return forward_call(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/coconut.py", line 239, in forward
+[rank0]:     logits = torch.cat(logits, dim=-2)
+[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]: torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 14.77 GiB. GPU 0 has a total capacity of 139.80 GiB of which 6.10 GiB is free. Including non-PyTorch memory, this process has 133.69 GiB memory in use. Of the allocated memory 131.74 GiB is allocated by PyTorch, and 767.15 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://docs.pytorch.org/docs/stable/notes/cuda.html#optimizing-memory-usage-with-pytorch-cuda-alloc-conf)
+wandb: 
+wandb: 🚀 View run star-sweep-bs32-lr1e-5_n1000_qwen3-0.6b-base_thoughtformer_lr1e-5_no-reset_dc_tc_adaptive-c3e91928 at: https://wandb.ai/seyedparsa/thoughtformer/runs/jek5luxr
+wandb: Find logs at: wandb/run-20260530_033212-jek5luxr/logs
+E0530 03:32:23.853000 464041 torch/distributed/elastic/multiprocessing/api.py:988] failed (exitcode: 1) local_rank: 0 (pid: 464115) of binary: /home/ubuntu/thoughtformer/.venv/bin/python3
+Traceback (most recent call last):
+  File "/home/ubuntu/thoughtformer/.venv/bin/torchrun", line 8, in <module>
+    sys.exit(main())
+             ^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 367, in wrapper
+    return f(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/run.py", line 1016, in main
+    run(args)
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/run.py", line 1007, in run
+    elastic_launch(
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/launcher/api.py", line 191, in __call__
+    return launch_agent(self._config, self._entrypoint, list(args))
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/launcher/api.py", line 371, in launch_agent
+    raise ChildFailedError(
+torch.distributed.elastic.multiprocessing.errors.ChildFailedError: 
+============================================================
+run.py FAILED
+------------------------------------------------------------
+Failures:
+  <NO_OTHER_FAILURES>
+------------------------------------------------------------
+Root Cause (first observed failure):
+[0]:
+  time      : 2026-05-30_03:32:23
+  host      : ip-172-31-10-226.us-east-2.compute.internal
+  rank      : 0 (local_rank: 0)
+  exitcode  : 1 (pid: 464115) 
+  error_file: <N/A>
+  traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
+============================================================

logs_backup/star_sweep_bs64_lr1e-4.log

@@ -0,0 +1,285 @@
+Config: {'project': 'thoughtformer', 'name': 'star-sweep-bs64-lr1e-4', 'only_eval': False, 'method': 'thoughtformer', 'data_curriculum': True, 'thought_curriculum': True, 'staging': 'adaptive', 'init_thought_stage': 1, 'init_data_stage': 1, 'patience': 5, 'staging_threshold': 0.9, 'c_thought': 1, 'max_latent_stage': 10, 'uniform_prob': 0.1, 'save_only_improve': False, 'model_id': 'Qwen/Qwen3-0.6B-Base', 'load_model_path': 'None', 'seed': 0, 'resume': 0, 'bf16': False, 'train_path': 'data/star_k10_L10_1000_train.json', 'val_path': 'data/star_k10_L10_50_valid.json', 'reset_optimizer': False, 'lr_schedule': 'cosine', 'batch_size_training': 64, 'eval_only_trained': False, 'eval_batch_size': 8, 'eval_every': 1, 'debug': False, 'gradient_accumulation_steps': 1, 'num_epochs': 30, 'lr': 0.0001, 'weight_decay': 0.01, 'group': 'star_bs_lr_sweep', 'tags': ['star', 'k10', 'L10', 'atc', 'adaptive', 'lr1e-5', 'best', 'budget_aug'], 'run_type': 'pilot'}
+/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/c10d_logger.py:83: UserWarning: barrier(): using the device under current context. You can specify `device_id` in `init_process_group` to mute this warning.
+  return func(*args, **kwargs)
+[rank0]:[W530 03:31:51.090585317 ProcessGroupNCCL.cpp:5324] Guessing device ID based on global rank. This can cause a hang if rank to GPU mapping is heterogeneous. You can specify device_id in init_process_group()
+
+Running FSDP on rank = 0, world size = 1
+/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py:478: UserWarning: FSDP is switching to use `NO_SHARD` instead of ShardingStrategy.FULL_SHARD since the world size is 1.
+  _init_core_state(
+FullyShardedDataParallel(
+  (_fsdp_wrapped_module): Coconut(
+    (base_causallm): Qwen3ForCausalLM(
+      (model): Qwen3Model(
+        (embed_tokens): Embedding(151672, 1024)
+        (layers): ModuleList(
+          (0-27): 28 x FullyShardedDataParallel(
+            (_fsdp_wrapped_module): Qwen3DecoderLayer(
+              (self_attn): Qwen3Attention(
+                (q_proj): Linear(in_features=1024, out_features=2048, bias=False)
+                (k_proj): Linear(in_features=1024, out_features=1024, bias=False)
+                (v_proj): Linear(in_features=1024, out_features=1024, bias=False)
+                (o_proj): Linear(in_features=2048, out_features=1024, bias=False)
+                (q_norm): Qwen3RMSNorm((128,), eps=1e-06)
+                (k_norm): Qwen3RMSNorm((128,), eps=1e-06)
+              )
+              (mlp): Qwen3MLP(
+                (gate_proj): Linear(in_features=1024, out_features=3072, bias=False)
+                (up_proj): Linear(in_features=1024, out_features=3072, bias=False)
+                (down_proj): Linear(in_features=3072, out_features=1024, bias=False)
+                (act_fn): SiLUActivation()
+              )
+              (input_layernorm): Qwen3RMSNorm((1024,), eps=1e-06)
+              (post_attention_layernorm): Qwen3RMSNorm((1024,), eps=1e-06)
+            )
+          )
+        )
+        (norm): Qwen3RMSNorm((1024,), eps=1e-06)
+        (rotary_emb): Qwen3RotaryEmbedding()
+      )
+      (lm_head): Linear(in_features=1024, out_features=151936, bias=False)
+    )
+    (embedding): Embedding(151672, 1024)
+  )
+)
+
+
+wandb: [wandb.login()] Loaded credentials for https://api.wandb.ai from /home/ubuntu/.netrc.
+wandb: Currently logged in as: seyedparsa to https://api.wandb.ai. Use `wandb login --relogin` to force relogin
+wandb: setting up run edlb2i8x
+wandb: Tracking run with wandb version 0.25.1
+wandb: Run data is saved locally in /home/ubuntu/thoughtformer/wandb/run-20260530_033218-edlb2i8x
+wandb: Run `wandb offline` to turn off syncing.
+wandb: Syncing run star-sweep-bs64-lr1e-4_n1000_qwen3-0.6b-base_thoughtformer_lr1e-4_no-reset_dc_tc_adaptive-64cda263
+wandb: ⭐️ View project at https://wandb.ai/seyedparsa/thoughtformer
+wandb: 🚀 View run at https://wandb.ai/seyedparsa/thoughtformer/runs/edlb2i8x
+
+============================================================
+EPOCH 0/30 (thought_stage=1, data_stage=1)
+============================================================
+  thought_stage=1, c_thought=1, max_difficulty=1
+
+
+
+
+
+  LR schedule: cosine, warmup=7 steps, total=4680 steps (max_steps/epoch=156)
+
+wandb: WARNING Serializing object of type str that is 838345 bytes
+Traceback (most recent call last):
+  File "/home/ubuntu/thoughtformer/run.py", line 1194, in <module>
+    main()
+  File "/home/ubuntu/thoughtformer/run.py", line 937, in main
+    outputs = parallel_model(**batch)
+              ^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+    return forward_call(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 857, in forward
+    output = self._fsdp_wrapped_module(*args, **kwargs)
+             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+    return forward_call(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/coconut.py", line 101, in forward
+    outputs = self.base_causallm(
+              ^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+    return forward_call(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/utils/generic.py", line 857, in wrapper
+    output = func(self, *args, **kwargs)
+             ^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 492, in forward
+    outputs: BaseModelOutputWithPast = self.model(
+                                       ^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+    return forward_call(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/utils/generic.py", line 931, in wrapper
+    output = func(self, *args, **kwargs)
+             ^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/utils/output_capturing.py", line 248, in wrapper
+    outputs = func(self, *args, **kwargs)
+              ^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 424, in forward
+    hidden_states = decoder_layer(
+                    ^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+    return forward_call(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 857, in forward
+    output = self._fsdp_wrapped_module(*args, **kwargs)
+             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/modeling_layers.py", line 93, in __call__
+    return super().__call__(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1884, in _call_impl
+    return inner()
+           ^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1832, in inner
+    result = forward_call(*args, **kwargs)
+             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 318, in forward
+    hidden_states, _ = self.self_attn(
+                       ^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1884, in _call_impl
+    return inner()
+           ^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1832, in inner
+    result = forward_call(*args, **kwargs)
+             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 277, in forward
+    attn_output, attn_weights = attention_interface(
+                                ^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/integrations/sdpa_attention.py", line 92, in sdpa_attention_forward
+    attn_output = torch.nn.functional.scaled_dot_product_attention(
+                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 406.00 MiB. GPU 0 has a total capacity of 139.80 GiB of which 45.06 MiB is free. Including non-PyTorch memory, this process has 139.75 GiB memory in use. Of the allocated memory 137.33 GiB is allocated by PyTorch, and 1.22 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://docs.pytorch.org/docs/stable/notes/cuda.html#optimizing-memory-usage-with-pytorch-cuda-alloc-conf)
+[rank0]: Traceback (most recent call last):
+[rank0]:   File "/home/ubuntu/thoughtformer/run.py", line 1194, in <module>
+[rank0]:     main()
+[rank0]:   File "/home/ubuntu/thoughtformer/run.py", line 937, in main
+[rank0]:     outputs = parallel_model(**batch)
+[rank0]:               ^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+[rank0]:     return forward_call(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 857, in forward
+[rank0]:     output = self._fsdp_wrapped_module(*args, **kwargs)
+[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+[rank0]:     return forward_call(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/coconut.py", line 101, in forward
+[rank0]:     outputs = self.base_causallm(
+[rank0]:               ^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+[rank0]:     return forward_call(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/utils/generic.py", line 857, in wrapper
+[rank0]:     output = func(self, *args, **kwargs)
+[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 492, in forward
+[rank0]:     outputs: BaseModelOutputWithPast = self.model(
+[rank0]:                                        ^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+[rank0]:     return forward_call(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/utils/generic.py", line 931, in wrapper
+[rank0]:     output = func(self, *args, **kwargs)
+[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/utils/output_capturing.py", line 248, in wrapper
+[rank0]:     outputs = func(self, *args, **kwargs)
+[rank0]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 424, in forward
+[rank0]:     hidden_states = decoder_layer(
+[rank0]:                     ^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+[rank0]:     return forward_call(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 857, in forward
+[rank0]:     output = self._fsdp_wrapped_module(*args, **kwargs)
+[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/modeling_layers.py", line 93, in __call__
+[rank0]:     return super().__call__(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1884, in _call_impl
+[rank0]:     return inner()
+[rank0]:            ^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1832, in inner
+[rank0]:     result = forward_call(*args, **kwargs)
+[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 318, in forward
+[rank0]:     hidden_states, _ = self.self_attn(
+[rank0]:                        ^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1884, in _call_impl
+[rank0]:     return inner()
+[rank0]:            ^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1832, in inner
+[rank0]:     result = forward_call(*args, **kwargs)
+[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 277, in forward
+[rank0]:     attn_output, attn_weights = attention_interface(
+[rank0]:                                 ^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/integrations/sdpa_attention.py", line 92, in sdpa_attention_forward
+[rank0]:     attn_output = torch.nn.functional.scaled_dot_product_attention(
+[rank0]:                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]: torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 406.00 MiB. GPU 0 has a total capacity of 139.80 GiB of which 45.06 MiB is free. Including non-PyTorch memory, this process has 139.75 GiB memory in use. Of the allocated memory 137.33 GiB is allocated by PyTorch, and 1.22 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://docs.pytorch.org/docs/stable/notes/cuda.html#optimizing-memory-usage-with-pytorch-cuda-alloc-conf)
+wandb: 
+wandb: 🚀 View run star-sweep-bs64-lr1e-4_n1000_qwen3-0.6b-base_thoughtformer_lr1e-4_no-reset_dc_tc_adaptive-64cda263 at: https://wandb.ai/seyedparsa/thoughtformer/runs/edlb2i8x
+wandb: Find logs at: wandb/run-20260530_033218-edlb2i8x/logs
+E0530 03:32:29.958000 464466 torch/distributed/elastic/multiprocessing/api.py:988] failed (exitcode: 1) local_rank: 0 (pid: 464939) of binary: /home/ubuntu/thoughtformer/.venv/bin/python3
+Traceback (most recent call last):
+  File "/home/ubuntu/thoughtformer/.venv/bin/torchrun", line 8, in <module>
+    sys.exit(main())
+             ^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 367, in wrapper
+    return f(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/run.py", line 1016, in main
+    run(args)
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/run.py", line 1007, in run
+    elastic_launch(
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/launcher/api.py", line 191, in __call__
+    return launch_agent(self._config, self._entrypoint, list(args))
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/launcher/api.py", line 371, in launch_agent
+    raise ChildFailedError(
+torch.distributed.elastic.multiprocessing.errors.ChildFailedError: 
+============================================================
+run.py FAILED
+------------------------------------------------------------
+Failures:
+  <NO_OTHER_FAILURES>
+------------------------------------------------------------
+Root Cause (first observed failure):
+[0]:
+  time      : 2026-05-30_03:32:29
+  host      : ip-172-31-10-226.us-east-2.compute.internal
+  rank      : 0 (local_rank: 0)
+  exitcode  : 1 (pid: 464939) 
+  error_file: <N/A>
+  traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
+============================================================

logs_backup/star_sweep_bs64_lr1e-5.log

@@ -0,0 +1,284 @@
+Config: {'project': 'thoughtformer', 'name': 'star-sweep-bs64-lr1e-5', 'only_eval': False, 'method': 'thoughtformer', 'data_curriculum': True, 'thought_curriculum': True, 'staging': 'adaptive', 'init_thought_stage': 1, 'init_data_stage': 1, 'patience': 5, 'staging_threshold': 0.9, 'c_thought': 1, 'max_latent_stage': 10, 'uniform_prob': 0.1, 'save_only_improve': False, 'model_id': 'Qwen/Qwen3-0.6B-Base', 'load_model_path': 'None', 'seed': 0, 'resume': 0, 'bf16': False, 'train_path': 'data/star_k10_L10_1000_train.json', 'val_path': 'data/star_k10_L10_50_valid.json', 'reset_optimizer': False, 'lr_schedule': 'cosine', 'batch_size_training': 64, 'eval_only_trained': False, 'eval_batch_size': 8, 'eval_every': 1, 'debug': False, 'gradient_accumulation_steps': 1, 'num_epochs': 30, 'lr': 1e-05, 'weight_decay': 0.01, 'group': 'star_bs_lr_sweep', 'tags': ['star', 'k10', 'L10', 'atc', 'adaptive', 'lr1e-5', 'best', 'budget_aug'], 'run_type': 'pilot'}
+/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/c10d_logger.py:83: UserWarning: barrier(): using the device under current context. You can specify `device_id` in `init_process_group` to mute this warning.
+  return func(*args, **kwargs)
+[rank0]:[W530 03:31:49.899347241 ProcessGroupNCCL.cpp:5324] Guessing device ID based on global rank. This can cause a hang if rank to GPU mapping is heterogeneous. You can specify device_id in init_process_group()
+
+Running FSDP on rank = 0, world size = 1
+/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py:478: UserWarning: FSDP is switching to use `NO_SHARD` instead of ShardingStrategy.FULL_SHARD since the world size is 1.
+  _init_core_state(
+FullyShardedDataParallel(
+  (_fsdp_wrapped_module): Coconut(
+    (base_causallm): Qwen3ForCausalLM(
+      (model): Qwen3Model(
+        (embed_tokens): Embedding(151672, 1024)
+        (layers): ModuleList(
+          (0-27): 28 x FullyShardedDataParallel(
+            (_fsdp_wrapped_module): Qwen3DecoderLayer(
+              (self_attn): Qwen3Attention(
+                (q_proj): Linear(in_features=1024, out_features=2048, bias=False)
+                (k_proj): Linear(in_features=1024, out_features=1024, bias=False)
+                (v_proj): Linear(in_features=1024, out_features=1024, bias=False)
+                (o_proj): Linear(in_features=2048, out_features=1024, bias=False)
+                (q_norm): Qwen3RMSNorm((128,), eps=1e-06)
+                (k_norm): Qwen3RMSNorm((128,), eps=1e-06)
+              )
+              (mlp): Qwen3MLP(
+                (gate_proj): Linear(in_features=1024, out_features=3072, bias=False)
+                (up_proj): Linear(in_features=1024, out_features=3072, bias=False)
+                (down_proj): Linear(in_features=3072, out_features=1024, bias=False)
+                (act_fn): SiLUActivation()
+              )
+              (input_layernorm): Qwen3RMSNorm((1024,), eps=1e-06)
+              (post_attention_layernorm): Qwen3RMSNorm((1024,), eps=1e-06)
+            )
+          )
+        )
+        (norm): Qwen3RMSNorm((1024,), eps=1e-06)
+        (rotary_emb): Qwen3RotaryEmbedding()
+      )
+      (lm_head): Linear(in_features=1024, out_features=151936, bias=False)
+    )
+    (embedding): Embedding(151672, 1024)
+  )
+)
+
+
+wandb: [wandb.login()] Loaded credentials for https://api.wandb.ai from /home/ubuntu/.netrc.
+wandb: Currently logged in as: seyedparsa to https://api.wandb.ai. Use `wandb login --relogin` to force relogin
+wandb: Tracking run with wandb version 0.25.1
+wandb: Run data is saved locally in /home/ubuntu/thoughtformer/wandb/run-20260530_033216-nwpqaes3
+wandb: Run `wandb offline` to turn off syncing.
+wandb: Syncing run star-sweep-bs64-lr1e-5_n1000_qwen3-0.6b-base_thoughtformer_lr1e-5_no-reset_dc_tc_adaptive-c9ba8dcb
+wandb: ⭐️ View project at https://wandb.ai/seyedparsa/thoughtformer
+wandb: 🚀 View run at https://wandb.ai/seyedparsa/thoughtformer/runs/nwpqaes3
+
+============================================================
+EPOCH 0/30 (thought_stage=1, data_stage=1)
+============================================================
+  thought_stage=1, c_thought=1, max_difficulty=1
+
+
+
+
+
+  LR schedule: cosine, warmup=7 steps, total=4680 steps (max_steps/epoch=156)
+
+wandb: WARNING Serializing object of type str that is 838345 bytes
+Traceback (most recent call last):
+  File "/home/ubuntu/thoughtformer/run.py", line 1194, in <module>
+    main()
+  File "/home/ubuntu/thoughtformer/run.py", line 937, in main
+    outputs = parallel_model(**batch)
+              ^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+    return forward_call(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 857, in forward
+    output = self._fsdp_wrapped_module(*args, **kwargs)
+             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+    return forward_call(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/coconut.py", line 101, in forward
+    outputs = self.base_causallm(
+              ^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+    return forward_call(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/utils/generic.py", line 857, in wrapper
+    output = func(self, *args, **kwargs)
+             ^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 492, in forward
+    outputs: BaseModelOutputWithPast = self.model(
+                                       ^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+    return forward_call(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/utils/generic.py", line 931, in wrapper
+    output = func(self, *args, **kwargs)
+             ^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/utils/output_capturing.py", line 248, in wrapper
+    outputs = func(self, *args, **kwargs)
+              ^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 424, in forward
+    hidden_states = decoder_layer(
+                    ^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+    return forward_call(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 857, in forward
+    output = self._fsdp_wrapped_module(*args, **kwargs)
+             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/modeling_layers.py", line 93, in __call__
+    return super().__call__(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1884, in _call_impl
+    return inner()
+           ^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1832, in inner
+    result = forward_call(*args, **kwargs)
+             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 318, in forward
+    hidden_states, _ = self.self_attn(
+                       ^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+    return self._call_impl(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1884, in _call_impl
+    return inner()
+           ^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1832, in inner
+    result = forward_call(*args, **kwargs)
+             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 277, in forward
+    attn_output, attn_weights = attention_interface(
+                                ^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/integrations/sdpa_attention.py", line 92, in sdpa_attention_forward
+    attn_output = torch.nn.functional.scaled_dot_product_attention(
+                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 406.00 MiB. GPU 0 has a total capacity of 139.80 GiB of which 45.06 MiB is free. Including non-PyTorch memory, this process has 139.75 GiB memory in use. Of the allocated memory 137.33 GiB is allocated by PyTorch, and 1.22 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://docs.pytorch.org/docs/stable/notes/cuda.html#optimizing-memory-usage-with-pytorch-cuda-alloc-conf)
+[rank0]: Traceback (most recent call last):
+[rank0]:   File "/home/ubuntu/thoughtformer/run.py", line 1194, in <module>
+[rank0]:     main()
+[rank0]:   File "/home/ubuntu/thoughtformer/run.py", line 937, in main
+[rank0]:     outputs = parallel_model(**batch)
+[rank0]:               ^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+[rank0]:     return forward_call(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 857, in forward
+[rank0]:     output = self._fsdp_wrapped_module(*args, **kwargs)
+[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+[rank0]:     return forward_call(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/coconut.py", line 101, in forward
+[rank0]:     outputs = self.base_causallm(
+[rank0]:               ^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+[rank0]:     return forward_call(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/utils/generic.py", line 857, in wrapper
+[rank0]:     output = func(self, *args, **kwargs)
+[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 492, in forward
+[rank0]:     outputs: BaseModelOutputWithPast = self.model(
+[rank0]:                                        ^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+[rank0]:     return forward_call(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/utils/generic.py", line 931, in wrapper
+[rank0]:     output = func(self, *args, **kwargs)
+[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/utils/output_capturing.py", line 248, in wrapper
+[rank0]:     outputs = func(self, *args, **kwargs)
+[rank0]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 424, in forward
+[rank0]:     hidden_states = decoder_layer(
+[rank0]:                     ^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
+[rank0]:     return forward_call(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 857, in forward
+[rank0]:     output = self._fsdp_wrapped_module(*args, **kwargs)
+[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/modeling_layers.py", line 93, in __call__
+[rank0]:     return super().__call__(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1884, in _call_impl
+[rank0]:     return inner()
+[rank0]:            ^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1832, in inner
+[rank0]:     result = forward_call(*args, **kwargs)
+[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 318, in forward
+[rank0]:     hidden_states, _ = self.self_attn(
+[rank0]:                        ^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
+[rank0]:     return self._call_impl(*args, **kwargs)
+[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1884, in _call_impl
+[rank0]:     return inner()
+[rank0]:            ^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1832, in inner
+[rank0]:     result = forward_call(*args, **kwargs)
+[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/models/qwen3/modeling_qwen3.py", line 277, in forward
+[rank0]:     attn_output, attn_weights = attention_interface(
+[rank0]:                                 ^^^^^^^^^^^^^^^^^^^^
+[rank0]:   File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/transformers/integrations/sdpa_attention.py", line 92, in sdpa_attention_forward
+[rank0]:     attn_output = torch.nn.functional.scaled_dot_product_attention(
+[rank0]:                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+[rank0]: torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 406.00 MiB. GPU 0 has a total capacity of 139.80 GiB of which 45.06 MiB is free. Including non-PyTorch memory, this process has 139.75 GiB memory in use. Of the allocated memory 137.33 GiB is allocated by PyTorch, and 1.22 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://docs.pytorch.org/docs/stable/notes/cuda.html#optimizing-memory-usage-with-pytorch-cuda-alloc-conf)
+wandb: 
+wandb: 🚀 View run star-sweep-bs64-lr1e-5_n1000_qwen3-0.6b-base_thoughtformer_lr1e-5_no-reset_dc_tc_adaptive-c9ba8dcb at: https://wandb.ai/seyedparsa/thoughtformer/runs/nwpqaes3
+wandb: Find logs at: wandb/run-20260530_033216-nwpqaes3/logs
+E0530 03:32:27.283000 464320 torch/distributed/elastic/multiprocessing/api.py:988] failed (exitcode: 1) local_rank: 0 (pid: 464399) of binary: /home/ubuntu/thoughtformer/.venv/bin/python3
+Traceback (most recent call last):
+  File "/home/ubuntu/thoughtformer/.venv/bin/torchrun", line 8, in <module>
+    sys.exit(main())
+             ^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 367, in wrapper
+    return f(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/run.py", line 1016, in main
+    run(args)
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/run.py", line 1007, in run
+    elastic_launch(
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/launcher/api.py", line 191, in __call__
+    return launch_agent(self._config, self._entrypoint, list(args))
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/ubuntu/thoughtformer/.venv/lib/python3.12/site-packages/torch/distributed/launcher/api.py", line 371, in launch_agent
+    raise ChildFailedError(
+torch.distributed.elastic.multiprocessing.errors.ChildFailedError: 
+============================================================
+run.py FAILED
+------------------------------------------------------------
+Failures:
+  <NO_OTHER_FAILURES>
+------------------------------------------------------------
+Root Cause (first observed failure):
+[0]:
+  time      : 2026-05-30_03:32:27
+  host      : ip-172-31-10-226.us-east-2.compute.internal
+  rank      : 0 (local_rank: 0)
+  exitcode  : 1 (pid: 464399) 
+  error_file: <N/A>
+  traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
+============================================================