Move RL batch-invariant tests to A10G + Fix RL CI

.github/workflows/integration_test_4gpu_rl.yaml → .github/workflows/integration_test_8gpu_rl.yaml

@@ -5,12 +5,12 @@ on:
     branches: [ main ]
     paths:
       - 'torchtitan/experiments/rl/**'
-      - '.github/workflows/integration_test_4gpu_rl.yaml'
+      - '.github/workflows/integration_test_8gpu_rl.yaml'
   pull_request:
     types: [opened, synchronize, reopened, ready_for_review]
     paths:
       - 'torchtitan/experiments/rl/**'
-      - '.github/workflows/integration_test_4gpu_rl.yaml'
+      - '.github/workflows/integration_test_8gpu_rl.yaml'
   schedule:
     # Runs every 12 hours
     - cron: '0 */12 * * *'
@@ -66,8 +66,8 @@ jobs:
         pip install uv
 
         # 1. Install Monarch and TorchStore
-        uv pip install torchmonarch==0.3.0
-        uv pip install --no-deps "git+https://github.com/meta-pytorch/torchstore.git@main"
+        uv pip install torchmonarch==0.4.1
+        uv pip install --no-deps "git+https://github.com/meta-pytorch/torchstore.git@8b41ecc52bccff73e762cbe38c6871efdfc5db36"
         uv pip install pygtrie portpicker
 
         # 2. Install batch-invariant ops
@@ -91,7 +91,21 @@ jobs:
         sudo mkdir -p "$RUNNER_TEMP/artifacts-to-be-uploaded"
         sudo chown -R $(id -u):$(id -g) "$RUNNER_TEMP/artifacts-to-be-uploaded"
 
-        # Run E2E RL integration tests (TP=2 on 4 GPUs)
+        # Disable FlashInfer JIT sampler — CI image lacks nvcc
+        export VLLM_USE_FLASHINFER_SAMPLER=0
+
+        # Run default E2E RL integration tests (TP=4 on 8 GPUs)
+        python -m torchtitan.experiments.rl.tests.integration_tests \
+          $RUNNER_TEMP/artifacts-to-be-uploaded --ngpu 8 \
+          --hf_assets_path "$MODEL_PATH"
+
+        # Run bitwise parity tests (TP=4, batch-invariant mode)
+        HF_ASSETS_PATH="$MODEL_PATH" torchrun --nproc-per-node=4 -m pytest \
+          torchtitan/experiments/rl/tests/test_bitwise_parity.py -v
+
+        # Run batch-invariant E2E RL integration test (TP=4 on 8 GPUs)
         python -m torchtitan.experiments.rl.tests.integration_tests \
-          $RUNNER_TEMP/artifacts-to-be-uploaded --ngpu 4 \
+          $RUNNER_TEMP/artifacts-to-be-uploaded --test_suite h100 --ngpu 8 \
           --hf_assets_path "$MODEL_PATH"
+
+        rm -rf $RUNNER_TEMP/artifacts-to-be-uploaded/*

torchtitan/experiments/rl/tests/integration_tests.py

@@ -34,34 +34,34 @@ def build_rl_test_list() -> list[OverrideDefinitions]:
                 [
                     "--module rl",
                     "--config rl_grpo_qwen3_0_6b",
-                    "--trainer.parallelism.tensor_parallel_degree 2",
-                    "--generator.parallelism.tensor_parallel_degree 2",
+                    "--trainer.parallelism.tensor_parallel_degree 4",
+                    "--generator.parallelism.tensor_parallel_degree 4",
                     "--generator.sampling.n 2",
                     "--trainer.debug.no_batch_invariant",
                     "--generator.debug.no_batch_invariant",
                     "--compile.no-enable",
                     "--generator.cudagraph.no-enable",
                 ],
             ],
-            "RL GRPO TP=2 no compile",
-            "rl_grpo_tp2_no_compile",
-            ngpu=4,
+            "RL GRPO TP=4 no compile",
+            "rl_grpo_tp4_no_compile",
+            ngpu=8,
         ),
         OverrideDefinitions(
             [
                 [
                     "--module rl",
                     "--config rl_grpo_qwen3_0_6b",
-                    "--trainer.parallelism.tensor_parallel_degree 2",
-                    "--generator.parallelism.tensor_parallel_degree 2",
+                    "--trainer.parallelism.tensor_parallel_degree 4",
+                    "--generator.parallelism.tensor_parallel_degree 4",
                     "--generator.sampling.n 2",
                     "--trainer.debug.no_batch_invariant",
                     "--generator.debug.no_batch_invariant",
                 ],
             ],
-            "RL GRPO TP=2 compile",
-            "rl_grpo_tp2_compile",
-            ngpu=4,
+            "RL GRPO TP=4 compile",
+            "rl_grpo_tp4_compile",
+            ngpu=8,
         ),
     ]
 
@@ -75,9 +75,9 @@ def build_rl_h100_test_list() -> list[OverrideDefinitions]:
                     "--config rl_grpo_qwen3_0_6b_batch_invariant",
                 ],
             ],
-            "RL GRPO TP=2 batch-invariant + deterministic",
-            "rl_grpo_tp2_batch_invariant",
-            ngpu=4,
+            "RL GRPO TP=4 batch-invariant + deterministic",
+            "rl_grpo_tp4_batch_invariant",
+            ngpu=8,
         ),
     ]