Fix trackio auth and buffer issues

train.py

@@ -43,6 +43,13 @@ from smolomni.svd_init import initialize_mla_from_pretrained
 
 import trackio
 
+# Safe trackio wrapper
+def safe_trackio_log(metrics):
+    try:
+        trackio.log(metrics)
+    except Exception:
+        pass
+
 
 # ===== Stage 1: KL Distillation Dataset =====
 class TextDistillationDataset(IterableDataset):
@@ -163,12 +170,14 @@ def train_stage1(args, config):
     )
     
     if accelerator.is_main_process:
-        trackio.init(
-            project="SmolOmni-MLA",
-            name="Stage1-KD",
-            space_id="TinmanLabSL/SmolOmni-MLA-Tracking",
-            config=vars(args),
-        )
+        try:
+            trackio.init(
+                project="SmolOmni-MLA",
+                name="Stage1-KD",
+                config=vars(args),
+            )
+        except Exception as e:
+            print(f"[WARN] Trackio init failed: {e}. Continuing without remote tracking.")
     
     set_seed(args.seed)
     
@@ -305,7 +314,7 @@ def train_stage1(args, config):
                 print(f"Step {global_step}/{args.max_steps} | Loss: {avg_loss:.4f} | "
                       f"LR: {scheduler.get_last_lr()[0]:.2e} | "
                       f"Speed: {steps_per_sec:.1f} steps/s")
-                trackio.log(metrics)
+                safe_trackio_log(metrics)
             
             total_loss = 0.0
         
@@ -345,12 +354,14 @@ def train_stage2(args, config):
     )
     
     if accelerator.is_main_process:
-        trackio.init(
-            project="SmolOmni-MLA",
-            name="Stage2-Joint",
-            space_id="TinmanLabSL/SmolOmni-MLA-Tracking",
-            config=vars(args),
-        )
+        try:
+            trackio.init(
+                project="SmolOmni-MLA",
+                name="Stage2-Joint",
+                config=vars(args),
+            )
+        except Exception as e:
+            print(f"[WARN] Trackio init failed: {e}. Continuing without remote tracking.")
     
     set_seed(args.seed)
     
@@ -469,7 +480,7 @@ def train_stage2(args, config):
                       f"Loss: {total_loss/n:.4f} | "
                       f"AR: {total_ar_loss/n:.4f} | "
                       f"Flow: {total_flow_loss/n:.4f}")
-                trackio.log(metrics)
+                safe_trackio_log(metrics)
             total_loss = total_ar_loss = total_flow_loss = 0.0
         
         if global_step % args.save_every == 0 and accelerator.is_main_process: