Upload mythos-fineweb-moe.py with huggingface_hub

mythos-fineweb-moe.py

@@ -958,6 +958,7 @@ def main():
     micro_batch = int(os.environ.get("MICRO_BATCH", "32"))
     target_tokens = int(os.environ.get("TARGET_TOKENS", "50_000_000"))
     grad_accum = int(os.environ.get("GRAD_ACCUM", "1"))
+    n_loops = int(os.environ.get("N_LOOPS", "6"))
     global_batch_tok = world_size * micro_batch * grad_accum * seq_len
     total_steps = target_tokens // global_batch_tok
     warmup_steps = 200
@@ -971,8 +972,8 @@ def main():
 
     if master:
         logger.info(
-            f"[MOE MLA+Engram] hidden=2048 | layers=6 | experts=32 | top-2 | "
-            f"seq_len={seq_len} | micro_batch={micro_batch} | grad_accum={grad_accum} | "
+            f"[MOE MLA+Engram] hidden=2048 | layers=6 | experts=16 | top-1 | "
+            f"n_loops={n_loops} | seq_len={seq_len} | micro_batch={micro_batch} | grad_accum={grad_accum} | "
             f"global_batch_tokens={global_batch_tok:,} | total_steps={total_steps:,}"
         )
         logger.info(
@@ -986,9 +987,9 @@ def main():
     # ------------------------------------------------------------------
     cfg = SpiderPortalConfig(
         hidden_size=2048, num_hidden_layers=6, num_attention_heads=16,
-        num_key_value_heads=4, intermediate_size=4096,
-        num_experts=32, num_experts_per_tok=2, num_shared_experts=1,
-        router_aux_loss_coef=0.05, max_loop_iters=2,
+        num_key_value_heads=4, intermediate_size=1024,
+        num_experts=16, num_experts_per_tok=1, num_shared_experts=1,
+        router_aux_loss_coef=0.05, max_loop_iters=16,
         prelude_layers=2, coda_layers=2, lora_rank=128,
         rope_theta=10000000.0,
         rope_scaling=None,
@@ -1172,7 +1173,7 @@ def main():
                 else model.no_sync()
             )
             with sync, amp_ctx, sdpa_ctx:
-                output = model(x)
+                output = model(x, n_loops=n_loops)
                 if master and step == start_step and micro_step == 0:
                     peak_vram = torch.cuda.max_memory_allocated() / 1024**3
                     logger.info(f"Reached first model forward  |  Peak VRAM: {peak_vram:.1f}GB")
@@ -1182,8 +1183,8 @@ def main():
                 else:
                     logits = output
                     aux_loss = 0.0
-                loss = nn.functional.cross_entropy(
-                    logits.view(-1, vocab_size), y.view(-1)
+                loss = F.nll_loss(
+                    logits.view(-1, vocab_size).log_softmax(dim=-1), y.view(-1)
                 )
                 loss = loss + cfg.router_aux_loss_coef * aux_loss
                 loss = loss / grad_accum