fix test_collectives (#1934)

* fix: fix test_collectives.py

* feat: reenable test_collectives

.github/workflows/test.yml

@@ -162,7 +162,7 @@ jobs:
         name: Test multigpu
         run: |
           PYTHONPATH="." python test/external/dist/test_world.py
-          #PYTHONPATH="." python test/external/dist/test_collectives.py
+          PYTHONPATH="." python test/external/dist/test_collectives.py
       - if: ${{ matrix.task == 'realworld' }}
         name: Test KOPT
         run: PYTHONPATH="." KOPT=1 BUDGET=20 GPU=1 DEBUG=1 python -m pytest -rA -n=auto test/models/test_real_world.py

test/external/dist/test_collectives.py

@@ -30,6 +30,7 @@ def run():
 
   # reset jit
   allreduce_jit.cnt = 0
+  allreduce_jit.input_replace = {}
 
   # test uneven chunk sizes
   for _ in range(3):