Plumb rl.loss_agg_mode to tunix GrpoConfig

tunix's GrpoConfig defaults loss_agg_mode to 'sequence-mean-token-mean',
but GPU NeMo-RL stacks use 'token-mean'. With group-normalized advantages
the two modes produce materially different losses, breaking GPU<->TPU
recipe parity.

Adds the field to the RL Pydantic schema + rl.yml default + passes it
through to GrpoConfig construction so users can override via cmdline:
'rl.loss_agg_mode=token-mean'.

Author: Pooya Moradi

src/maxtext/configs/post_train/rl.yml

@@ -54,6 +54,12 @@ rl:
   grpo_epsilon: 0.2
   loss_algo: 'grpo' # grpo or gspo-token
 
+  # Loss aggregation mode passed to tunix's GrpoConfig. tunix defaults to
+  # 'sequence-mean-token-mean'; set 'token-mean' for parity with GPU NeMo-RL
+  # stacks that use token-mean aggregation. With group-normalized advantages
+  # the two modes produce materially different losses.
+  loss_agg_mode: 'sequence-mean-token-mean'  # 'token-mean' | 'sequence-mean' | 'sequence-mean-token-mean'
+
   # ====== Agentic Rollout ======
   # If true, uses the async AgenticGRPOLearner, which overlaps rollout generation
   # with training for faster throughput via online vLLM inference.

src/maxtext/configs/types.py

@@ -1963,6 +1963,16 @@ class RL(BaseModel):
   grpo_beta: float = Field(0.08, description="Coefficient for the KL divergence penalty (β).")
   grpo_epsilon: float = Field(0.2, description="Epsilon value for clipping in the GRPO loss.")
   loss_algo: Literal["grpo", "gspo-token"] = Field("grpo", description="Loss algorithm, i.e., 'grpo' or 'gspo-token'.")
+  loss_agg_mode: Literal["token-mean", "sequence-mean", "sequence-mean-token-mean"] = Field(
+      "sequence-mean-token-mean",
+      description=(
+          "Loss aggregation mode passed to tunix's GrpoConfig. tunix defaults"
+          " to 'sequence-mean-token-mean'; set 'token-mean' for parity with"
+          " GPU NeMo-RL stacks that use token-mean aggregation. With"
+          " group-normalized advantages the two modes produce materially"
+          " different losses."
+      ),
+  )
   use_agentic_rollout: bool = Field(
       False,
       description="If True, uses the asynchronous AgenticGRPOLearner for online vLLM rollouts.",

src/maxtext/trainers/post_train/rl/train_rl.py

@@ -571,6 +571,7 @@ def _reward_fn(**kwargs):
         beta=trainer_config.rl.grpo_beta,
         epsilon=trainer_config.rl.grpo_epsilon,
         loss_algo=trainer_config.rl.loss_algo,
+        loss_agg_mode=trainer_config.rl.loss_agg_mode,
     )
     rl_trainer = GrpoLearner(
         rl_cluster=rl_cluster,