[startju]

What does this PR do?

optimal_token_baseline / tir_optimal_token_baseline advantage estimators assert on data.batch["sum_pi_squared"], but no worker code populated it — actor_config.calculate_sum_pi_squared was a dead flag. Any run with algorithm.adv_estimator=optimal_token_baseline crashed in compute_advantage.

This PR wires up the full producer pipeline for sum_pi_squared (Σπ², the sum of squared token probabilities used to compute the Logit-Gradient Norm proxy W(τ) = Σ_t[1 - 2π_t + Σπ²]).

No similar open PR found:

• https://github.com/verl-project/verl/pulls?q=is%3Apr+is%3Aopen+sum_pi_squared
• https://github.com/verl-project/verl/pulls?q=is%3Apr+is%3Aopen+optimal_token_baseline

Checklist Before Starting

• Search for similar PRs. Paste at least one query link here: https://github.com/verl-project/verl/pulls?q=is%3Apr+is%3Aopen+sum_pi_squared
• Format the PR title as [{modules}] {type}: {description}

Test

• CPU parametrised tests for calculate_sum_pi_squared_from_logits (shapes/dtypes + extreme-magnitude stability check): pytest tests/utils/test_torch_functional.py -k sum_pi_squared -v — all pass.
• Distributed test for vocab_parallel_sum_pi_squared mirroring test_vocab_parallel_entropy: TP=2 sharding vs single-rank reference, with non-destructiveness assertion.
• End-to-end: full FSDP training run with algorithm.adv_estimator=optimal_token_baseline on GSM8K (Qwen3-0.6B, 2 GPUs) completes without crash and shows decreasing loss.

API and Usage Example

Enable via config (now a shared base field, available to all backends):

actor_rollout_ref.actor.calculate_sum_pi_squared=True \
algorithm.adv_estimator=optimal_token_baseline

Design & Code Changes

Config

• Promote calculate_sum_pi_squared from per-backend (FSDPActorConfig, McoreActorConfig) to the shared base ActorConfig, alongside calculate_entropy. Update actor.yaml, dp_actor.yaml, megatron_actor.yaml and regenerate all _generated_*.yaml.

FSDP engine (verl/workers/engine/fsdp/transformer_impl.py)

• Compute Σπ² from logits via calculate_sum_pi_squared_from_logits, mirroring entropy through both rmpad and no-rmpad branches and the ulysses-sp gather.
• Respects sum_pi_squared_checkpointing (same pattern as entropy_checkpointing).
• Raises NotImplementedError on use_fused_kernels=True.

Megatron engine (verl/workers/engine/megatron/transformer_impl.py, verl/utils/megatron/tensor_parallel.py)

• Add vocab_parallel_sum_pi_squared (TP-sharded, non-destructive) and call it before the destructive vocab_parallel_entropy.
• Raises NotImplementedError on use_fused_kernels=True.

Trainer (verl/trainer/ppo/ray_trainer.py)

• _compute_old_log_prob reads calculate_sum_pi_squared, passes it as a non-tensor flag, and unions sum_pi_squared into old_log_prob so it reaches compute_advantage.

Checklist Before Submitting

• Read the Contribute Guide.
• Apply pre-commit checks: all hooks pass.
• Add / Update the documentation. Not required for a bug fix with no new user-facing API surface.
• Add unit or end-to-end test(s): CPU math tests + distributed TP test added. CI workflow not modified as the distributed test requires GPU allocation.

AI assistance was used (Claude Sonnet 4.6) to help implement and review this change. Every changed line has been reviewed and the end-to-end test was run by the human submitter.

[CLAassistant]

All committers have signed the CLA.

[gemini-code-assist]

Code Review

This pull request introduces the calculation of sum of squared probabilities (Σπ²), which is required for the Logit-Gradient Norm proxy in advantage estimation. The changes include the implementation of vocab_parallel_sum_pi_squared for Megatron-based tensor parallelism, integration into the FSDP and Megatron transformer engines, and updates to the PPO trainer and configuration files. Feedback highlights a potential AttributeError in the FSDP engine where sum_pi_squared_checkpointing is accessed from the wrong configuration object and provides a technical correction regarding the in-place mutation behavior of entropy calculations to ensure safe gradient handling.

[gemini-code-assist]

The field sum_pi_squared_checkpointing is defined in FSDPActorConfig but accessed here via self.engine_config. In verl/workers/config/actor.py, self.engine is assigned to self.fsdp_config (which is an instance of FSDPEngineConfig). If FSDPEngineConfig does not contain this field, this will result in an AttributeError at runtime. Please ensure that sum_pi_squared_checkpointing (and similarly entropy_checkpointing) are either moved to FSDPEngineConfig or that the engine is initialized with the full FSDPActorConfig object.

[gemini-code-assist]

The comment stating that vocab_parallel_entropy mutates logits in place via exp_() in the forward pass appears to be inaccurate based on the implementation in verl/utils/megatron/tensor_parallel.py. In that file, exp_() is called on normalized_vocab_parallel_logits, which is a new tensor created by vocab_parallel_logits - logits_max. However, vocab_parallel_entropy DOES mutate the input in place during the backward pass. While the current order is safe for inference (where compute_loss=False), it is important to maintain this order or use clones if gradients are ever required for both metrics simultaneously.