GRPO training fails in multi-process mode

[yhtang]

Expected Behavior
GRPO training with Tunix in multi-process SPMD (e.g., 2 GPU nodes) should generate rollouts and compute advantages without failures.

Actual Behavior
Crash during rollout sampling when iterating a distributed jax.Array:

• AssertionError at jax/_src/array.py:380: assert self.is_fully_replicated or self.is_fully_addressable
• Triggered in tunix/generate/sampler.py while doing zip(out_tokens, lengths)

Steps to Reproduce the Problem

1. Add jax.distributed.initialize() at the beginning of the GRPO example script.
2. Run the GRPO demo

srun -u --label --ntasks=2 --ntasks-per-node=1 -c${SLURM_CPUS_ON_NODE} python <grpo-script>

Checklist

• I have searched the existing issues for a similar bug report.
• I have provided all the required information in the "Environment" section.
• I have provided a minimal, reproducible example.

Would you like to help us fix it?

[wang2yn84]

Hi @yhtang , we didn't specifically tested multi process setup because we are leveraging Pathways for multihost support. But we are happy to solve all the issues you found. Can you provide more details about the crash you see during rollout?

[yhtang]

So the direct cause here is tokenizer decoding iterates over sharded token arrays. Specifically, the vanilla sampler decodes by iterating the out_tokens and lengths array that are FSDP-sharded, which is illegal unless arrays are fully-replicated or fully-addressable.

The root cause is that reward function evaluation is handled assuming a single-controller model, hence breaking multi-process SPMD.

Detailed error analysis:

• Stack:

    Traceback (most recent call last):
      ...
      File "/opt/tunix/tunix/rl/rollout/vanilla_rollout.py", line 55, in generate
        output = self._sampler(
      ...
      File "/opt/tunix/tunix/generate/sampler.py", line 770, in __call__
        for tokens, length in zip(out_tokens, lengths)
      File ".../jax/_src/array.py", line 380, in __iter__
        assert self.is_fully_replicated or self.is_fully_addressable
  

• Observed sharding:

    out_tokens shape: (4, 512)
    sharding: NamedSharding(mesh=Mesh('fsdp': 2, 'tp': 2, axis_types=(Auto, Auto)),
                            spec=PartitionSpec('fsdp',), memory_kind=device)
  

• Call path:

  GRPOLearner._generate_and_compute_advantage(...)
  --> RLCluster.generate(...)
  --> VanillaRollout.generate(...) (calls Sampler.__call__(..., pad_output=True))
  --> Sampler.__call__ builds decoded_outputs = [tokenizer.decode(...)] by iterating out_tokens/lengths.
  

• Downstream: .text is passed into RLLearner._compute_rewards(prompts, completions, ...), which executes host-side reward functions on every process and returns a small unsharded jnp.array of rewards.

Ideal solution: shard_map + callback

Users should not feel this complexity; it should be encapsulated.

• Keep completions on device as sharded arrays:
• Use shard_map with PartitionSpec('fsdp', ...) to run per-shard reward computation:
  • Inside the mapped function, use a single io_callback/pure_callback per shard to evaluate the reward functions.
• Return numeric rewards as a sharded JAX array.
• Continue GRPO training on-device.

Benefits:

• Avoid allgather of rollouts or reward.
• Reward computation is distributed.
• Minimal disruption to existing workflow and API.

Fallback solution for fast unblocking

All-gather out_tokens and lengths, evaluate reward on all processes. However, this will cause replicated CPU work and may produce duplicated metrics/logs, and may become bottleneck for large global batches or expensive rewards.

[wang2yn84]

Thank you very much for the detailed analysis! I feel like it's fine to all gather here, mainly because we need to all reduce here anyway for the rewards calculation (calculate mean and std).