Add optional process-pool evaluation of reward functions

[gutianyu-google]

Sequence-level reward functions are pure Python, evaluated serially over the whole batch on one core while the host's remaining cores idle. On a GRPO post-training benchmark (qwen3-0.6b, 2048 sequences/step, 3 reward functions on a 40+ core head node), this serial pass took 29.6 s of a 155.6 s global step — the second-largest component after rollout.

This PR adds AlgorithmConfig.reward_num_workers:

• Default 0 keeps the serial implementation and preserves current behavior bit-for-bit.
• When > 1, SequenceRewardManager evaluates each reward function over contiguous chunks of the batch in a fork-based process pool and concatenates results in order. Per-sequence reward functions therefore produce identical results to the serial path. kwargs that are per-example columns (len == num_prompts, e.g. ground-truth answer lists) are sliced in lockstep with the batch; everything else passes through unchanged.

Robustness, all covered by unit tests:

• Functions that cannot run in a worker — unpicklable, or spawning subprocesses of their own (which daemonic pool workers forbid) — are detected at runtime and evaluated in the parent process from then on, where their own internal parallelism keeps working.
• Pool-creation failure, worker crashes, chunk timeouts, or a platform without the fork start method all fall back to the serial implementation with one warning.
• TUNIX_REWARD_PARALLEL=0 disables the pool at runtime without a config change.

Measured (same benchmark, reward_num_workers=32): reward computation 29.6 s → ~3.5 s, global step 155.6 s → 129.0 s (−17%), with per-step rewards byte-identical to the serial run over all 20 steps (max abs diff 0.000000). One of the three reward functions dispatches to its own subprocess pool and ran in the parent via the fallback; widening that internal pool contributed part of the win.

Reference
Benchmark context: MaxText GRPO post-training on TPU v7x; reward functions are the match_format_* / check_numbers family from MaxText's RL utils.

Checklist

• I have added all the necessary unit tests for my change.
• I have verified that my change does not break existing code and all unit tests pass.
• I have added all appropriate doc-strings/documentation.
• My PR is based on the latest changes of the main branch (if unsure, rebase the code).
• I have signed the Contributor License Agreement.
• I have followed Contribution Guidelines.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.