Upload checkpoints_vlm_gym_colorization_one_image_lr2e_5_mse_only_ins/checkpoints_vlm_gym_colorization_one_image_lr2e_5_mse_only_ins

checkpoints_vlm_gym_colorization_one_image_lr2e_5_mse_only_ins/checkpoints_vlm_gym_colorization_one_image_lr2e_5_mse_only_ins/wandb/offline-run-20260125_170309-vlm_gym_colorization_one_img_lr2e_5_mse_only_ins-run0/files/output.log

@@ -1,189 +1,3 @@
-FullyShardedDataParallel(
-  (_fsdp_wrapped_module): Bagel(
-    (language_model): Qwen2ForCausalLM(
-      (model): Qwen2Model(
-        (embed_tokens): Embedding(152064, 3584)
-        (layers): ModuleList(
-          (0-27): 28 x FullyShardedDataParallel(
-            (_fsdp_wrapped_module): CheckpointWrapper(
-              (_checkpoint_wrapped_module): Qwen2MoTDecoderLayer(
-                (self_attn): PackedAttentionMoT(
-                  (q_proj): Linear(in_features=3584, out_features=3584, bias=True)
-                  (k_proj): Linear(in_features=3584, out_features=512, bias=True)
-                  (v_proj): Linear(in_features=3584, out_features=512, bias=True)
-                  (o_proj): Linear(in_features=3584, out_features=3584, bias=False)
-                  (q_norm): Qwen2RMSNorm((128,), eps=1e-06)
-                  (k_norm): Qwen2RMSNorm((128,), eps=1e-06)
-                  (q_norm_moe_gen): Qwen2RMSNorm((128,), eps=1e-06)
-                  (k_norm_moe_gen): Qwen2RMSNorm((128,), eps=1e-06)
-                  (q_proj_moe_gen): Linear(in_features=3584, out_features=3584, bias=True)
-                  (k_proj_moe_gen): Linear(in_features=3584, out_features=512, bias=True)
-                  (v_proj_moe_gen): Linear(in_features=3584, out_features=512, bias=True)
-                  (o_proj_moe_gen): Linear(in_features=3584, out_features=3584, bias=False)
-                )
-                (mlp): Qwen2MLP(
-                  (gate_proj): Linear(in_features=3584, out_features=18944, bias=False)
-                  (up_proj): Linear(in_features=3584, out_features=18944, bias=False)
-                  (down_proj): Linear(in_features=18944, out_features=3584, bias=False)
-                  (act_fn): SiLU()
-                )
-                (mlp_moe_gen): Qwen2MLP(
-                  (gate_proj): Linear(in_features=3584, out_features=18944, bias=False)
-                  (up_proj): Linear(in_features=3584, out_features=18944, bias=False)
-                  (down_proj): Linear(in_features=18944, out_features=3584, bias=False)
-                  (act_fn): SiLU()
-                )
-                (input_layernorm): Qwen2RMSNorm((3584,), eps=1e-06)
-                (input_layernorm_moe_gen): Qwen2RMSNorm((3584,), eps=1e-06)
-                (post_attention_layernorm): Qwen2RMSNorm((3584,), eps=1e-06)
-                (post_attention_layernorm_moe_gen): Qwen2RMSNorm((3584,), eps=1e-06)
-              )
-            )
-          )
-        )
-        (norm): Qwen2RMSNorm((3584,), eps=1e-06)
-        (norm_moe_gen): Qwen2RMSNorm((3584,), eps=1e-06)
-        (rotary_emb): Qwen2RotaryEmbedding()
-      )
-      (lm_head): Linear(in_features=3584, out_features=152064, bias=False)
-    )
-    (time_embedder): FullyShardedDataParallel(
-      (_fsdp_wrapped_module): TimestepEmbedder(
-        (mlp): Sequential(
-          (0): Linear(in_features=256, out_features=3584, bias=True)
-          (1): SiLU()
-          (2): Linear(in_features=3584, out_features=3584, bias=True)
-        )
-      )
-    )
-    (vae2llm): Linear(in_features=64, out_features=3584, bias=True)
-    (llm2vae): Linear(in_features=3584, out_features=64, bias=True)
-    (latent_pos_embed): FullyShardedDataParallel(
-      (_fsdp_wrapped_module): PositionEmbedding()
-    )
-    (vit_model): SiglipVisionModel(
-      (vision_model): FullyShardedDataParallel(
-        (_fsdp_wrapped_module): SiglipVisionTransformer(
-          (embeddings): SiglipVisionEmbeddings(
-            (position_embedding): Embedding(4900, 1152)
-            (patch_embedding): Linear(in_features=588, out_features=1152, bias=True)
-          )
-          (encoder): SiglipEncoder(
-            (layers): ModuleList(
-              (0-25): 26 x FullyShardedDataParallel(
-                (_fsdp_wrapped_module): CheckpointWrapper(
-                  (_checkpoint_wrapped_module): SiglipEncoderLayer(
-                    (self_attn): SiglipFlashAttention2(
-                      (k_proj): Linear(in_features=1152, out_features=1152, bias=True)
-                      (v_proj): Linear(in_features=1152, out_features=1152, bias=True)
-                      (q_proj): Linear(in_features=1152, out_features=1152, bias=True)
-                      (out_proj): Linear(in_features=1152, out_features=1152, bias=True)
-                    )
-                    (layer_norm1): LayerNorm((1152,), eps=1e-06, elementwise_affine=True)
-                    (mlp): SiglipMLP(
-                      (activation_fn): PytorchGELUTanh()
-                      (fc1): Linear(in_features=1152, out_features=4304, bias=True)
-                      (fc2): Linear(in_features=4304, out_features=1152, bias=True)
-                    )
-                    (layer_norm2): LayerNorm((1152,), eps=1e-06, elementwise_affine=True)
-                  )
-                )
-              )
-            )
-          )
-          (post_layernorm): LayerNorm((1152,), eps=1e-06, elementwise_affine=True)
-        )
-      )
-    )
-    (connector): FullyShardedDataParallel(
-      (_fsdp_wrapped_module): CheckpointWrapper(
-        (_checkpoint_wrapped_module): MLPconnector(
-          (activation_fn): PytorchGELUTanh()
-          (fc1): Linear(in_features=1152, out_features=3584, bias=True)
-          (fc2): Linear(in_features=3584, out_features=3584, bias=True)
-        )
-      )
-    )
-    (vit_pos_embed): FullyShardedDataParallel(
-      (_fsdp_wrapped_module): PositionEmbedding()
-    )
-  )
-)
-_flat_param True
-language_model.model.layers.0._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.1._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.2._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.3._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.4._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.5._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.6._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.7._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.8._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.9._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.10._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.11._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.12._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.13._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.14._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.15._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.16._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.17._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.18._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.19._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.20._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.21._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.22._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.23._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.24._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.25._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.26._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.27._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-time_embedder._fsdp_wrapped_module._flat_param True
-latent_pos_embed._fsdp_wrapped_module._flat_param False
-vit_model.vision_model._fsdp_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.0._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.1._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.2._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.3._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.4._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.5._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.6._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.7._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.8._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.9._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.10._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.11._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.12._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.13._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.14._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.15._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.16._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.17._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.18._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.19._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.20._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.21._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.22._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.23._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.24._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.25._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-connector._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_pos_embed._fsdp_wrapped_module._flat_param False
-Preparing Dataset vlm_gym_colorization_mse_loss_only/vlm_gym_colorization_train
-base_dir is /dev/shm/models/checkpoints_vlm_gym_colorization_one_image_lr2e_5_mse_only_ins/eval_used_rows, step_tag is vlm_gym_colorization_one_img_lr2e_5_mse_only_ins_step0
-Preparing Dataset vlm_gym_colorization_mse_loss_only_evalonce/vlm_gym_colorization_val
-[eval debug] first 3 batch fingerprints:
-  fp[0]: [{'data_indexes': [0], 'worker_id': 0, 'dataset_name': 'vlm_gym_colorization_mse_loss_only_evalonce'}]
-  fp[1]: [{'data_indexes': [8], 'worker_id': 0, 'dataset_name': 'vlm_gym_colorization_mse_loss_only_evalonce'}]
-  fp[2]: [{'data_indexes': [16], 'worker_id': 0, 'dataset_name': 'vlm_gym_colorization_mse_loss_only_evalonce'}]
-ce_avg: 0.0, mse_avg: 0.05326032266020775
-base_dir is /dev/shm/models/checkpoints_vlm_gym_colorization_one_image_lr2e_5_mse_only_ins/eval_used_rows, step_tag is vlm_gym_colorization_one_img_lr2e_5_mse_only_ins_step500
-Preparing Dataset vlm_gym_colorization_mse_loss_only_evalonce/vlm_gym_colorization_val
-[eval debug] first 3 batch fingerprints:
-  fp[0]: [{'data_indexes': [0], 'worker_id': 0, 'dataset_name': 'vlm_gym_colorization_mse_loss_only_evalonce'}]
-  fp[1]: [{'data_indexes': [8], 'worker_id': 0, 'dataset_name': 'vlm_gym_colorization_mse_loss_only_evalonce'}]
-  fp[2]: [{'data_indexes': [16], 'worker_id': 0, 'dataset_name': 'vlm_gym_colorization_mse_loss_only_evalonce'}]
-ce_avg: 0.0, mse_avg: 0.007997258566319942
 wandb: Detected [huggingface_hub.inference] in use.
 wandb: Use W&B Weave for improved LLM call tracing. Install Weave with `pip install weave` then add `import weave` to the top of your script.
 wandb: For more information, check out the docs at: https://weave-docs.wandb.ai/
@@ -781,4 +595,244 @@ wandb: For more information, check out the docs at: https://weave-docs.wandb.ai/
 [2026-01-25 20:42:14] (step=0000584) Train Loss mse: 0.0077, Train Loss ce: 0.0000, Train Steps/Sec: 0.04,
 [2026-01-25 20:42:37] (step=0000585) Train Loss mse: 0.0083, Train Loss ce: 0.0000, Train Steps/Sec: 0.04,
 [2026-01-25 20:42:58] (step=0000586) Train Loss mse: 0.0079, Train Loss ce: 0.0000, Train Steps/Sec: 0.05,
-[2026-01-25 20:43:20] (step=0000587) Train Loss mse: 0.0078, Train Loss ce: 0.0000, Train Steps/Sec: 0.04,
+[2026-01-25 20:43:20] (step=0000587) Train Loss mse: 0.0078, Train Loss ce: 0.0000, Train Steps/Sec: 0.04,
+[2026-01-25 20:43:44] (step=0000588) Train Loss mse: 0.0073, Train Loss ce: 0.0000, Train Steps/Sec: 0.04,
+[2026-01-25 20:44:05] (step=0000589) Train Loss mse: 0.0080, Train Loss ce: 0.0000, Train Steps/Sec: 0.05,
+[2026-01-25 20:44:26] (step=0000590) Train Loss mse: 0.0075, Train Loss ce: 0.0000, Train Steps/Sec: 0.05,
+[2026-01-25 20:44:46] (step=0000591) Train Loss mse: 0.0082, Train Loss ce: 0.0000, Train Steps/Sec: 0.05,
+[2026-01-25 20:45:10] (step=0000592) Train Loss mse: 0.0075, Train Loss ce: 0.0000, Train Steps/Sec: 0.04,
+[2026-01-25 20:45:32] (step=0000593) Train Loss mse: 0.0071, Train Loss ce: 0.0000, Train Steps/Sec: 0.04,
+[2026-01-25 20:45:54] (step=0000594) Train Loss mse: 0.0077, Train Loss ce: 0.0000, Train Steps/Sec: 0.05,
+[2026-01-25 20:46:13] (step=0000595) Train Loss mse: 0.0072, Train Loss ce: 0.0000, Train Steps/Sec: 0.05,
+[2026-01-25 20:46:33] (step=0000596) Train Loss mse: 0.0071, Train Loss ce: 0.0000, Train Steps/Sec: 0.05,
+[2026-01-25 20:46:55] (step=0000597) Train Loss mse: 0.0076, Train Loss ce: 0.0000, Train Steps/Sec: 0.05,
+[2026-01-25 20:47:12] (step=0000598) Train Loss mse: 0.0088, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-25 20:47:33] (step=0000599) Train Loss mse: 0.0069, Train Loss ce: 0.0000, Train Steps/Sec: 0.05,
+[2026-01-25 20:47:53] (step=0000600) Train Loss mse: 0.0074, Train Loss ce: 0.0000, Train Steps/Sec: 0.05,
+[2026-01-25 20:48:14] (step=0000601) Train Loss mse: 0.0085, Train Loss ce: 0.0000, Train Steps/Sec: 0.05,
+[2026-01-25 20:48:35] (step=0000602) Train Loss mse: 0.0079, Train Loss ce: 0.0000, Train Steps/Sec: 0.05,
+[2026-01-25 20:48:58] (step=0000603) Train Loss mse: 0.0073, Train Loss ce: 0.0000, Train Steps/Sec: 0.04,
+[2026-01-25 20:49:25] (step=0000604) Train Loss mse: 0.0071, Train Loss ce: 0.0000, Train Steps/Sec: 0.04,
+[2026-01-25 20:49:46] (step=0000605) Train Loss mse: 0.0083, Train Loss ce: 0.0000, Train Steps/Sec: 0.05,
+[2026-01-25 20:50:07] (step=0000606) Train Loss mse: 0.0076, Train Loss ce: 0.0000, Train Steps/Sec: 0.05,
+[2026-01-25 20:50:28] (step=0000607) Train Loss mse: 0.0080, Train Loss ce: 0.0000, Train Steps/Sec: 0.05,
+[2026-01-25 20:50:47] (step=0000608) Train Loss mse: 0.0075, Train Loss ce: 0.0000, Train Steps/Sec: 0.05,
+[2026-01-25 20:51:07] (step=0000609) Train Loss mse: 0.0082, Train Loss ce: 0.0000, Train Steps/Sec: 0.05,
+[2026-01-25 20:51:29] (step=0000610) Train Loss mse: 0.0088, Train Loss ce: 0.0000, Train Steps/Sec: 0.05,
+[2026-01-25 20:51:49] (step=0000611) Train Loss mse: 0.0083, Train Loss ce: 0.0000, Train Steps/Sec: 0.05,
+[2026-01-25 20:52:22] (step=0000612) Train Loss mse: 0.0068, Train Loss ce: 0.0000, Train Steps/Sec: 0.03,
+[2026-01-25 20:52:45] (step=0000613) Train Loss mse: 0.0076, Train Loss ce: 0.0000, Train Steps/Sec: 0.04,
+[2026-01-25 20:53:07] (step=0000614) Train Loss mse: 0.0073, Train Loss ce: 0.0000, Train Steps/Sec: 0.05,
+[2026-01-25 20:53:28] (step=0000615) Train Loss mse: 0.0090, Train Loss ce: 0.0000, Train Steps/Sec: 0.05,
+[2026-01-25 20:53:49] (step=0000616) Train Loss mse: 0.0081, Train Loss ce: 0.0000, Train Steps/Sec: 0.05,
+FullyShardedDataParallel(
+  (_fsdp_wrapped_module): Bagel(
+    (language_model): Qwen2ForCausalLM(
+      (model): Qwen2Model(
+        (embed_tokens): Embedding(152064, 3584)
+        (layers): ModuleList(
+          (0-27): 28 x FullyShardedDataParallel(
+            (_fsdp_wrapped_module): CheckpointWrapper(
+              (_checkpoint_wrapped_module): Qwen2MoTDecoderLayer(
+                (self_attn): PackedAttentionMoT(
+                  (q_proj): Linear(in_features=3584, out_features=3584, bias=True)
+                  (k_proj): Linear(in_features=3584, out_features=512, bias=True)
+                  (v_proj): Linear(in_features=3584, out_features=512, bias=True)
+                  (o_proj): Linear(in_features=3584, out_features=3584, bias=False)
+                  (q_norm): Qwen2RMSNorm((128,), eps=1e-06)
+                  (k_norm): Qwen2RMSNorm((128,), eps=1e-06)
+                  (q_norm_moe_gen): Qwen2RMSNorm((128,), eps=1e-06)
+                  (k_norm_moe_gen): Qwen2RMSNorm((128,), eps=1e-06)
+                  (q_proj_moe_gen): Linear(in_features=3584, out_features=3584, bias=True)
+                  (k_proj_moe_gen): Linear(in_features=3584, out_features=512, bias=True)
+                  (v_proj_moe_gen): Linear(in_features=3584, out_features=512, bias=True)
+                  (o_proj_moe_gen): Linear(in_features=3584, out_features=3584, bias=False)
+                )
+                (mlp): Qwen2MLP(
+                  (gate_proj): Linear(in_features=3584, out_features=18944, bias=False)
+                  (up_proj): Linear(in_features=3584, out_features=18944, bias=False)
+                  (down_proj): Linear(in_features=18944, out_features=3584, bias=False)
+                  (act_fn): SiLU()
+                )
+                (mlp_moe_gen): Qwen2MLP(
+                  (gate_proj): Linear(in_features=3584, out_features=18944, bias=False)
+                  (up_proj): Linear(in_features=3584, out_features=18944, bias=False)
+                  (down_proj): Linear(in_features=18944, out_features=3584, bias=False)
+                  (act_fn): SiLU()
+                )
+                (input_layernorm): Qwen2RMSNorm((3584,), eps=1e-06)
+                (input_layernorm_moe_gen): Qwen2RMSNorm((3584,), eps=1e-06)
+                (post_attention_layernorm): Qwen2RMSNorm((3584,), eps=1e-06)
+                (post_attention_layernorm_moe_gen): Qwen2RMSNorm((3584,), eps=1e-06)
+              )
+            )
+          )
+        )
+        (norm): Qwen2RMSNorm((3584,), eps=1e-06)
+        (norm_moe_gen): Qwen2RMSNorm((3584,), eps=1e-06)
+        (rotary_emb): Qwen2RotaryEmbedding()
+      )
+      (lm_head): Linear(in_features=3584, out_features=152064, bias=False)
+    )
+    (time_embedder): FullyShardedDataParallel(
+      (_fsdp_wrapped_module): TimestepEmbedder(
+        (mlp): Sequential(
+          (0): Linear(in_features=256, out_features=3584, bias=True)
+          (1): SiLU()
+          (2): Linear(in_features=3584, out_features=3584, bias=True)
+        )
+      )
+    )
+    (vae2llm): Linear(in_features=64, out_features=3584, bias=True)
+    (llm2vae): Linear(in_features=3584, out_features=64, bias=True)
+    (latent_pos_embed): FullyShardedDataParallel(
+      (_fsdp_wrapped_module): PositionEmbedding()
+    )
+    (vit_model): SiglipVisionModel(
+      (vision_model): FullyShardedDataParallel(
+        (_fsdp_wrapped_module): SiglipVisionTransformer(
+          (embeddings): SiglipVisionEmbeddings(
+            (position_embedding): Embedding(4900, 1152)
+            (patch_embedding): Linear(in_features=588, out_features=1152, bias=True)
+          )
+          (encoder): SiglipEncoder(
+            (layers): ModuleList(
+              (0-25): 26 x FullyShardedDataParallel(
+                (_fsdp_wrapped_module): CheckpointWrapper(
+                  (_checkpoint_wrapped_module): SiglipEncoderLayer(
+                    (self_attn): SiglipFlashAttention2(
+                      (k_proj): Linear(in_features=1152, out_features=1152, bias=True)
+                      (v_proj): Linear(in_features=1152, out_features=1152, bias=True)
+                      (q_proj): Linear(in_features=1152, out_features=1152, bias=True)
+                      (out_proj): Linear(in_features=1152, out_features=1152, bias=True)
+                    )
+                    (layer_norm1): LayerNorm((1152,), eps=1e-06, elementwise_affine=True)
+                    (mlp): SiglipMLP(
+                      (activation_fn): PytorchGELUTanh()
+                      (fc1): Linear(in_features=1152, out_features=4304, bias=True)
+                      (fc2): Linear(in_features=4304, out_features=1152, bias=True)
+                    )
+                    (layer_norm2): LayerNorm((1152,), eps=1e-06, elementwise_affine=True)
+                  )
+                )
+              )
+            )
+          )
+          (post_layernorm): LayerNorm((1152,), eps=1e-06, elementwise_affine=True)
+        )
+      )
+    )
+    (connector): FullyShardedDataParallel(
+      (_fsdp_wrapped_module): CheckpointWrapper(
+        (_checkpoint_wrapped_module): MLPconnector(
+          (activation_fn): PytorchGELUTanh()
+          (fc1): Linear(in_features=1152, out_features=3584, bias=True)
+          (fc2): Linear(in_features=3584, out_features=3584, bias=True)
+        )
+      )
+    )
+    (vit_pos_embed): FullyShardedDataParallel(
+      (_fsdp_wrapped_module): PositionEmbedding()
+    )
+  )
+)
+_flat_param True
+language_model.model.layers.0._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.1._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.2._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.3._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.4._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.5._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.6._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.7._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.8._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.9._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.10._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.11._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.12._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.13._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.14._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.15._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.16._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.17._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.18._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.19._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.20._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.21._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.22._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.23._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.24._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.25._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.26._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.27._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+time_embedder._fsdp_wrapped_module._flat_param True
+latent_pos_embed._fsdp_wrapped_module._flat_param False
+vit_model.vision_model._fsdp_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.0._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.1._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.2._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.3._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.4._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.5._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.6._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.7._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.8._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.9._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.10._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.11._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.12._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.13._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.14._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.15._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.16._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.17._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.18._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.19._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.20._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.21._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.22._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.23._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.24._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.25._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+connector._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_pos_embed._fsdp_wrapped_module._flat_param False
+Preparing Dataset vlm_gym_colorization_mse_loss_only/vlm_gym_colorization_train
+base_dir is /dev/shm/models/checkpoints_vlm_gym_colorization_one_image_lr2e_5_mse_only_ins/eval_used_rows, step_tag is vlm_gym_colorization_one_img_lr2e_5_mse_only_ins_step0
+Preparing Dataset vlm_gym_colorization_mse_loss_only_evalonce/vlm_gym_colorization_val
+[eval debug] first 3 batch fingerprints:
+  fp[0]: [{'data_indexes': [0], 'worker_id': 0, 'dataset_name': 'vlm_gym_colorization_mse_loss_only_evalonce'}]
+  fp[1]: [{'data_indexes': [8], 'worker_id': 0, 'dataset_name': 'vlm_gym_colorization_mse_loss_only_evalonce'}]
+  fp[2]: [{'data_indexes': [16], 'worker_id': 0, 'dataset_name': 'vlm_gym_colorization_mse_loss_only_evalonce'}]
+ce_avg: 0.0, mse_avg: 0.05326032266020775
+base_dir is /dev/shm/models/checkpoints_vlm_gym_colorization_one_image_lr2e_5_mse_only_ins/eval_used_rows, step_tag is vlm_gym_colorization_one_img_lr2e_5_mse_only_ins_step500
+Preparing Dataset vlm_gym_colorization_mse_loss_only_evalonce/vlm_gym_colorization_val
+[eval debug] first 3 batch fingerprints:
+  fp[0]: [{'data_indexes': [0], 'worker_id': 0, 'dataset_name': 'vlm_gym_colorization_mse_loss_only_evalonce'}]
+  fp[1]: [{'data_indexes': [8], 'worker_id': 0, 'dataset_name': 'vlm_gym_colorization_mse_loss_only_evalonce'}]
+  fp[2]: [{'data_indexes': [16], 'worker_id': 0, 'dataset_name': 'vlm_gym_colorization_mse_loss_only_evalonce'}]
+ce_avg: 0.0, mse_avg: 0.007997258566319942
+[2026-01-25 20:54:14] (step=0000617) Train Loss mse: 0.0085, Train Loss ce: 0.0000, Train Steps/Sec: 0.04,
+[2026-01-25 20:54:35] (step=0000618) Train Loss mse: 0.0074, Train Loss ce: 0.0000, Train Steps/Sec: 0.05,
+[2026-01-25 20:54:54] (step=0000619) Train Loss mse: 0.0086, Train Loss ce: 0.0000, Train Steps/Sec: 0.05,
+[2026-01-25 20:55:14] (step=0000620) Train Loss mse: 0.0076, Train Loss ce: 0.0000, Train Steps/Sec: 0.05,
+[2026-01-25 20:55:37] (step=0000621) Train Loss mse: 0.0076, Train Loss ce: 0.0000, Train Steps/Sec: 0.04,
+[2026-01-25 20:55:59] (step=0000622) Train Loss mse: 0.0078, Train Loss ce: 0.0000, Train Steps/Sec: 0.05,
+[2026-01-25 20:56:22] (step=0000623) Train Loss mse: 0.0089, Train Loss ce: 0.0000, Train Steps/Sec: 0.04,
+[2026-01-25 20:56:45] (step=0000624) Train Loss mse: 0.0079, Train Loss ce: 0.0000, Train Steps/Sec: 0.04,
+[2026-01-25 20:57:06] (step=0000625) Train Loss mse: 0.0091, Train Loss ce: 0.0000, Train Steps/Sec: 0.05,
+[2026-01-25 20:57:30] (step=0000626) Train Loss mse: 0.0080, Train Loss ce: 0.0000, Train Steps/Sec: 0.04,
+[2026-01-25 20:57:51] (step=0000627) Train Loss mse: 0.0073, Train Loss ce: 0.0000, Train Steps/Sec: 0.05,
+[2026-01-25 20:58:14] (step=0000628) Train Loss mse: 0.0086, Train Loss ce: 0.0000, Train Steps/Sec: 0.04,
+[2026-01-25 20:58:34] (step=0000629) Train Loss mse: 0.0078, Train Loss ce: 0.0000, Train Steps/Sec: 0.05,
+[2026-01-25 20:58:58] (step=0000630) Train Loss mse: 0.0082, Train Loss ce: 0.0000, Train Steps/Sec: 0.04,
+[2026-01-25 20:59:18] (step=0000631) Train Loss mse: 0.0086, Train Loss ce: 0.0000, Train Steps/Sec: 0.05,
+[2026-01-25 20:59:40] (step=0000632) Train Loss mse: 0.0079, Train Loss ce: 0.0000, Train Steps/Sec: 0.05,
+[2026-01-25 20:59:59] (step=0000633) Train Loss mse: 0.0090, Train Loss ce: 0.0000, Train Steps/Sec: 0.05,
+[2026-01-25 21:00:21] (step=0000634) Train Loss mse: 0.0082, Train Loss ce: 0.0000, Train Steps/Sec: 0.04,
+[2026-01-25 21:00:44] (step=0000635) Train Loss mse: 0.0091, Train Loss ce: 0.0000, Train Steps/Sec: 0.04,
+[2026-01-25 21:01:04] (step=0000636) Train Loss mse: 0.0080, Train Loss ce: 0.0000, Train Steps/Sec: 0.05,
+[2026-01-25 21:01:25] (step=0000637) Train Loss mse: 0.0079, Train Loss ce: 0.0000, Train Steps/Sec: 0.05,
+[2026-01-25 21:01:50] (step=0000638) Train Loss mse: 0.0064, Train Loss ce: 0.0000, Train Steps/Sec: 0.04,
+[2026-01-25 21:02:12] (step=0000639) Train Loss mse: 0.0084, Train Loss ce: 0.0000, Train Steps/Sec: 0.04,
+[2026-01-25 21:02:34] (step=0000640) Train Loss mse: 0.0087, Train Loss ce: 0.0000, Train Steps/Sec: 0.05,
+[2026-01-25 21:02:56] (step=0000641) Train Loss mse: 0.0079, Train Loss ce: 0.0000, Train Steps/Sec: 0.04,

checkpoints_vlm_gym_colorization_one_image_lr2e_5_mse_only_ins/checkpoints_vlm_gym_colorization_one_image_lr2e_5_mse_only_ins/wandb/offline-run-20260125_170309-vlm_gym_colorization_one_img_lr2e_5_mse_only_ins-run0/files/wandb-summary.json

@@ -1 +1 @@
-{"_runtime": 13210.631898663, "lr": 1.981747237555166e-05, "total_ce_tokens": 0, "total_norm": 0.12721827626228333, "total_samples": 16, "eval/ce": 0, "ce": 0, "total_mse_tokens": 87744, "mem_allocated": 59549.859375, "mem_cache": 77350, "_step": 586, "eval/mse": 0.007997258566319942, "_timestamp": 1769373778.5024066, "mse": 0.0079260915517807}
+{"_runtime": 14386.023717632, "lr": 1.9742649624205542e-05, "total_ce_tokens": 0, "total_norm": 0.13556422293186188, "total_samples": 16, "eval/ce": 0, "ce": 0, "total_mse_tokens": 76256, "mem_allocated": 59549.859375, "mem_cache": 77430, "_step": 640, "eval/mse": 0.007997258566319942, "_timestamp": 1769374954.148881, "mse": 0.008700057864189148}

checkpoints_vlm_gym_colorization_one_image_lr2e_5_mse_only_ins/checkpoints_vlm_gym_colorization_one_image_lr2e_5_mse_only_ins/wandb/offline-run-20260125_170309-vlm_gym_colorization_one_img_lr2e_5_mse_only_ins-run0/run-vlm_gym_colorization_one_img_lr2e_5_mse_only_ins-run0.wandb

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:4281b748958d71bb7fd2acb4ba10fe4a6c796f5027202b05e70bbd8ef2e16891
-size 4030464
+oid sha256:7d2a6ef470a0f28027b3f864fd146368428ba578fc86e600d4d011f46973322a
+size 4358144