[megatron] seq-mean-token-mean guard rejects CP=1 and appears overly restrictive for reconstructed THD outputs

System Info

• Observed in a custom training application using verl's v1 trainer / Megatron engine.
• Model: Qwen3.5-9B.
• Training topology: 16 GPUs, TP=8, PP=1, static CP=2 (DP=1).
• Python: 3.12.
• use_remove_padding=True, dynamic_context_parallel=False.
• Policy loss: GSPO; aggregation: seq-mean-token-mean.

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder
• My own task or dataset

Reproduction

Observed failure

The first actor update fails with the following effective configuration (relevant fields only, not a standalone launch command):

actor_rollout_ref:
  actor:
    policy_loss:
      loss_mode: gspo
    loss_agg_mode: seq-mean-token-mean
    megatron:
      tensor_model_parallel_size: 8
      pipeline_model_parallel_size: 1
      context_parallel_size: 2
      dynamic_context_parallel: false
      use_remove_padding: true
      use_fused_kernels: false

The engine reaches the calculate_per_token_loss=True branch and raises:

File "verl/workers/engine/megatron/transformer_impl.py", in postprocess_micro_batch_func
    raise ValueError(...)
ValueError: loss_agg_mode='seq-mean-token-mean' is incompatible with calculate_per_token_loss=True (auto-enabled by Megatron-Bridge under CP>1). The per-sequence inner division by n_s requires local-shard counts that diverge from global under CP. Use one of: 'token-mean', 'seq-mean-token-sum', 'seq-mean-token-sum-norm'.

1. The guard also rejects CP=1

The guard in postprocess_micro_batch_func effectively checks:

if self.tf_config.calculate_per_token_loss:
    if loss_agg_mode == "seq-mean-token-mean" and not dynamic_context_parallel:
        raise ValueError(...)

It does not check context_parallel_size > 1. Therefore, with the standard configured PPO loss callback, setting CP=1 while explicitly retaining calculate_per_token_loss=True also reaches this exception, despite there being no CP sharding.

2. The static-CP THD path appears to reconstruct full sequences before loss computation

The guard's comment states:

Static CP cannot compose per-sequence token means from local output shards.

However, tracing the non-fused THD path suggests that the policy loss does not receive only local output shards:

1. gptmodel_forward_model_engine applies postprocess_thd_engine to the log-probability outputs before returning them to the engine's loss callback.

2. postprocess_thd_engine selects the static CP group when local_cp_size is None, gathers the outputs, restores sequence order, and trims padding. Its gather uses:

   python复制

   torch.distributed.all_gather(output_list, output.detach(), group=cp_group)
   output_list[cp_rank] = output

   Thus, each rank sees full-sequence values while retaining autograd only for its local output shard.

3. ppo_loss extracts response log-probabilities and uses the full response_mask from the micro-batch. It does not CP-slice that mask before calling the policy loss.

4. compute_policy_loss_gspo and agg_loss therefore appear to use full-response lengths for both the sequence importance ratio and sequence-mean aggregation.

Under this path, the claimed CP-local sequence denominator does not appear to arise. Also, if local sequence means were actually being computed with incorrect local denominators, that would be a problem independently of calculate_per_token_loss; changing the outer gradient normalization would not fix those denominators.

3. Per-token gradient normalization does not inherently preclude sequence averaging

For a global mini-batch with B sequences, DP size D, and global routed-token count T, agg_loss returns each micro-batch's contribution as:

loss = (D / B) * sum(per_sequence_mean_losses_in_this_micro_batch)

The engine then returns:

local_sum = loss * routed_num_tokens / dp_size

Megatron accumulates gradients, sums them across the relevant parallel groups, and divides by the accumulated global token count. If that count is exactly T, the outer T factors cancel and the original sequence-mean objective is preserved. This assumes correct shard-gradient accounting and the existing global-batch normalization metadata.

Consequently, calculate_per_token_loss=True by itself does not establish incompatibility with seq-mean-token-mean.

Minimal CPU check

This single-case sample directly calls verl's compute_policy_loss_gspo. Edit CP_SIZE, CP_RANK, response lengths, advantages, and clipping parameters at the top. It fixes DP=1 and uses one micro-batch to focus on the disputed sequence denominator, without a parameter-sweep loop or source extraction.

After sequence reconstruction, every rank sees all log-probability values, while only its local shard retains autograd. The sample models precisely that state with torch.where(local_mask, x, x.detach()). It compares the full loss value and this rank's gradient contribution with an unsharded reference. Uniform-width zigzag ownership is used only to keep the example readable; this is not a reproduction of variable-length THD packing.

"""CPU loss-level check; requires an installed verl environment.

Assume full-sequence outputs have already been gathered. Check one CP rank's
local gradient against its portion of the unsharded reference. DP=1 and one
micro-batch; no model forward, actual gather, or token-counting test.
"""
import torch
from verl.trainer.ppo.core_algos import compute_policy_loss_gspo
from verl.workers.config import ActorConfig

# Edit these parameters and rerun this single case.
CP_SIZE = 2
CP_RANK = 0
RESPONSE_LENGTHS = [5, 9]
ADVANTAGES = [1.0, -1.0]
CLIP_LOW = 0.0003
CLIP_HIGH = 0.28
SEED = 31

assert CP_SIZE >= 1 and 0 <= CP_RANK < CP_SIZE
assert len(RESPONSE_LENGTHS) == len(ADVANTAGES)
assert min(RESPONSE_LENGTHS) > 0

torch.manual_seed(SEED)
dtype = torch.float64
batch_size = len(RESPONSE_LENGTHS)
# Use a common aligned width to keep the zigzag ownership example simple.
width = ((max(RESPONSE_LENGTHS) + 2 * CP_SIZE - 1) // (2 * CP_SIZE)) * (2 * CP_SIZE)
positions = torch.arange(width)
response_mask = positions[None, :] < torch.tensor(RESPONSE_LENGTHS)[:, None]
old_log_prob = torch.zeros(batch_size, width, dtype=dtype)
log_prob_values = 0.05 * torch.randn(batch_size, width, dtype=dtype)
advantages = torch.tensor(ADVANTAGES, dtype=dtype)[:, None].expand_as(old_log_prob)
config = ActorConfig(
    strategy="megatron", rollout_n=1, ppo_micro_batch_size_per_gpu=1,
    clip_ratio_low=CLIP_LOW, clip_ratio_high=CLIP_HIGH,
)
config.global_batch_info.update(dp_size=1, global_batch_size=batch_size)


def loss(log_prob):
    return compute_policy_loss_gspo(
        old_log_prob=old_log_prob,
        log_prob=log_prob,
        advantages=advantages,
        response_mask=response_mask,
        loss_agg_mode="seq-mean-token-mean",
        config=config,
    )[0]


# 1. Reference: every token retains its gradient.
reference = log_prob_values.clone().requires_grad_()
reference_loss = loss(reference)
reference_loss.backward()

# 2. After a CP gather: full values, but only this rank's shard is differentiable.
chunk_id = positions // (width // (2 * CP_SIZE))
local_mask = (chunk_id == CP_RANK) | (chunk_id == 2 * CP_SIZE - 1 - CP_RANK)
local_input = log_prob_values.clone().requires_grad_()
gathered = torch.where(local_mask[None, :], local_input, local_input.detach())
cp_loss = loss(gathered)

# Exact global normalization count is assumed here; counting is out of scope.
normalization_tokens = sum(RESPONSE_LENGTHS)
(cp_loss * normalization_tokens).backward()
cp_gradient = local_input.grad / normalization_tokens
expected_gradient = reference.grad * local_mask[None, :]

torch.testing.assert_close(cp_loss, reference_loss, rtol=1e-12, atol=1e-12)
torch.testing.assert_close(cp_gradient, expected_gradient, rtol=1e-12, atol=1e-12)
print(f"CP={CP_SIZE}, rank={CP_RANK}: PASS")
print(f"loss error: {(cp_loss - reference_loss).abs().item():.2e}")
print(f"local gradient error: {(cp_gradient - expected_gradient).abs().max().item():.2e}")

Actual output for the default parameters on the inspected downstream checkout:

CP=2, rank=0: PASS
loss error: 0.00e+00
local gradient error: 1.39e-17

This checks the GSPO loss and local token-log-probability gradient after an assumed full-sequence gather, using exact global normalization.

Expected behavior

• CP=1 with calculate_per_token_loss=True should not be rejected on the grounds of CP-local sequence denominators.
• For static CP with reconstructed full-sequence THD outputs, the restriction should be removed or narrowed after loss/gradient equivalence is verified, unless there is another unsupported interaction missing from this analysis.