Speed optimization: set DDP find_unused_parameters=False

model/scripts/train.py

@@ -1476,10 +1476,10 @@ def main():
                                               str(C("compile_mode", "reduce-overhead")), args.smoke, R0,
                                               device=device, vocab_size=margs.vocab_size)
     if use_dist and not args.smoke:
-        # MoE top-2 routing me kuch experts kabhi-kabhi empty rehte hain -> find_unused_parameters=True (safe)
-        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[device.index], find_unused_parameters=True)
+        # All 242M parameters are 100% active in vectorized MoE -> find_unused_parameters=False (maximum throughput)
+        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[device.index], find_unused_parameters=False)
         if R0:
-            print(f"[dist] DDP wrapped (find_unused=True, micro={batch} accum={accum} synced)")
+            print(f"[dist] DDP wrapped (find_unused=False, micro={batch} accum={accum} synced)")
     ckpt_model = model.module if (use_dist and not args.smoke) else model  # ckpt bina module. prefix ke (resume-safe)
     total, active = ckpt_model.param_counts()
     if R0: