Fix DDP parameter and gradient layouts

2026-07-09 16:19:49 +08:00
parent ed8537fb3e
commit 8078fcb835
3 changed files with 42 additions and 7 deletions
--- a/train_next_step.py
+++ b/train_next_step.py
@@ -901,4 +901,8 @@ def main() -> None:


 if __name__ == "__main__":
-    main()
+    try:
+        main()
+    finally:
+        if dist.is_initialized():
+            dist.destroy_process_group()