[JAX] CI fixes: dynamic FSDP_SIZE + accept exit 5 for pre-Blackwell skip

Two failures from the same PR run:

* CI B200 (8 GPUs) -- the test hard-coded EP_SIZE=2, FSDP_SIZE=2 ->
  4-device mesh, while jax.device_count()==8. mesh_utils.create_device_mesh
  rejected with 'Number of devices 8 must equal the product of mesh_shape
  (2, 2)'. FSDP_SIZE now derives from jax.device_count() // EP_SIZE, so
  the test scales to whatever the launcher gives us (4 GPUs -> FSDP=2,
  8 GPUs -> FSDP=4). BATCH formula already depends on EP_SIZE*FSDP_SIZE
  so it scales automatically and stays MXFP8-aligned.

* CI H100 (4 GPUs) -- the file emits pytest.skip(allow_module_level=True)
  when get_device_compute_capability(0) < 100 (Blackwell-only). pytest
  reports 'no tests collected' with exit code 5, which the multiprocess
  launcher was treating as a failure. Accept 5 alongside 0; anything
  else is still a failure.

Signed-off-by: tdophung <tdophung@nvidia.com>

Author: tdophung

tests/jax/run_multiprocess_moe_vjp.sh

@@ -105,10 +105,12 @@ done
 # Final pass/fail. Any non-zero in any process fails the suite, but
 # we tolerate non-zero on the non-zero processes only if proc 0
 # reports PASS (this matches the encoder launcher's logic). Simplest
-# strict rule: any non-zero is a failure.
+# Treat exit 0 (pass) and exit 5 (pytest "no tests collected", which
+# the file emits via ``pytest.skip(allow_module_level=True)`` on
+# pre-Blackwell GPUs) as success. Anything else is a failure.
 FAILED=0
 for e in "${EXITS[@]}"; do
-    if [ "$e" != "0" ]; then
+    if [ "$e" != "0" ] && [ "$e" != "5" ]; then
         FAILED=1
         break
     fi

tests/jax/test_multiprocess_moe_vjp.py

@@ -129,11 +129,17 @@ def _read_mp_options():
 ]
 
 
-NUM_DEVICES_REQUIRED = 4
 EP_AXIS = "ep"
 FSDP_AXIS = "fsdp"
 EP_SIZE = 2
-FSDP_SIZE = 2
+# FSDP_SIZE adapts to whatever the launcher gave us: dlcluster GB200
+# gives 4 GPUs (FSDP=2), CI B200 gives 8 GPUs (FSDP=4). Both stay
+# 128-aligned for MXFP8 and divide num_experts/topk cleanly.
+assert (
+    jax.device_count() % EP_SIZE == 0
+), f"device_count {jax.device_count()} must be divisible by EP_SIZE={EP_SIZE}"
+FSDP_SIZE = jax.device_count() // EP_SIZE
+NUM_DEVICES_REQUIRED = EP_SIZE * FSDP_SIZE
 
 LOGICAL_AXIS_RULES = (
     ("exp", EP_AXIS),
@@ -280,7 +286,7 @@ def _local_shard(x):
 # 4-way data-parallel shard of a Mixtral-8 block).
 # -----------------------------------------------------------------------------
 
-BATCH = EP_SIZE * FSDP_SIZE * 4  # 16
+BATCH = EP_SIZE * FSDP_SIZE * 4  # 16 on 4-GPU, 32 on 8-GPU
 SEQ = 2048
 HIDDEN = 1024
 INTER = 4096