Merge pull request #4334 from AI-Hypercomputer:fix-tunix-pr-1637-degen-groups

PiperOrigin-RevId: 944068388

Author: Google-ML-Automation

src/maxtext/configs/post_train/rl.yml

@@ -66,8 +66,6 @@ rl:
   off_policy_steps: 0
   # System prompt injected into the agent at rollout time.
   system_prompt: ''
-  # If true, mask degenerate groups (all-zero advantages) from contributing to the loss.
-  degenerate_group_masking: true
   # Upper-bound clipping epsilon for GRPO loss; defaults to grpo_epsilon when null.
   epsilon_high: null
   # Number of model keys to chunk for resharding tensors between trainer and rollout devices.

src/maxtext/configs/types.py

@@ -2156,10 +2156,6 @@ class RL(BaseModel):
       "",
       description="System prompt injected into the agent at rollout time (agentic only).",
   )
-  degenerate_group_masking: bool = Field(
-      True,
-      description="Mask degenerate groups (all-zero advantages) from contributing to loss (agentic only).",
-  )
   epsilon_high: Optional[float] = Field(
       None,
       description="Upper-bound clipping epsilon for GRPO loss. Defaults to epsilon when None (agentic only).",

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

@@ -580,7 +580,6 @@ def _reward_fn(**kwargs):
         max_concurrency=trainer_config.rl.max_concurrency,
         off_policy_steps=trainer_config.rl.off_policy_steps,
         system_prompt=trainer_config.rl.system_prompt,
-        degenerate_group_masking=trainer_config.rl.degenerate_group_masking,
         epsilon_high=trainer_config.rl.epsilon_high,
     )
     # Instantiate the custom MaxText chat parser