feat(megatron): handle zero-token no-signal steps in the shared trainer

docs/en/examples/algorithms.md

@@ -67,7 +67,7 @@ Each sample's scalar advantage is broadcast to its response tokens. Relax masks
 
 $$L_{\mathrm{RLOO}} = -\frac{1}{N_{\mathrm{eff}}}\sum_i\sum_t m_{i,t}\operatorname{stopgrad}(A_i)\log\pi_{i,t}$$
 
-This global-token reduction does not apply a separate $1/T_i$ weight to each response. `train/pg_clipfrac` is always `0` because RLOO uses no clipping.
+This global-token reduction does not apply a separate $1/T_i$ weight to each response. An empty or fully masked response still participates in its prompt group's leave-one-out reward baseline, but contributes zero policy-loss tokens and zero to $N_{\mathrm{eff}}$. If the entire global batch has no effective tokens, it is treated as a no-signal step and reports zero loss metrics rather than dividing by zero. `train/pg_clipfrac` is always `0` because RLOO uses no clipping.
 
 ### Requirements and Parameters
 

docs/zh/examples/algorithms.md

@@ -64,7 +64,7 @@ $$L_{i,t} = -\operatorname{stopgrad}(A_i)\log\pi_\theta(y_{i,t}\mid x,y_{i,<t})$
 
 $$L_{\mathrm{RLOO}} = -\frac{1}{N_{\mathrm{eff}}}\sum_i\sum_t m_{i,t}\operatorname{stopgrad}(A_i)\log\pi_{i,t}$$
 
-这种 global-token reduction 不会为每条 response 单独附加 $1/T_i$ 权重。RLOO 不使用 clipping，因此 `train/pg_clipfrac` 始终为 `0`。
+这种 global-token reduction 不会为每条 response 单独附加 $1/T_i$ 权重。空 response 或全 mask response 仍参与所属 prompt group 的 leave-one-out reward baseline，但贡献零个 policy-loss token，且不会计入 $N_{\mathrm{eff}}$。如果整个 global batch 都没有有效 token，则将其作为 no-signal step 并报告零 loss 指标，而不是执行除零。RLOO 不使用 clipping，因此 `train/pg_clipfrac` 始终为 `0`。
 
 ### 约束与参数
 

relax/backends/megatron/cp_utils.py

@@ -189,12 +189,13 @@ def get_cp_local_num_tokens(
     ``* cp_size`` on the loss/metric — a coupling that only cancels when CP is
     uniform across the step.
 
-    For ``cp_size == 1`` this reduces to the total number of unmasked tokens
-    (preserving the historical per-sample ``clamp_min(., 1)``).
+    For ``cp_size == 1`` this is exactly the total number of unmasked tokens.
+    Empty responses contribute zero, rather than being counted as one token; this
+    matches the CP-partitioned path and the global-valid-token loss definition.
     """
     cp_size = dynamic_cp_size if dynamic_cp_size is not None else mpu.get_context_parallel_world_size()
     if cp_size == 1:
-        return sum([torch.clamp_min(loss_mask.sum(), 1) for loss_mask in loss_masks])
+        return sum([loss_mask.sum() for loss_mask in loss_masks])
 
     # cp_size > 1: mirror the chunk slicing done in get_sum_of_sample_mean so the
     # counted tokens exactly match the ones sum_of_token contributes on this rank.

relax/backends/megatron/loss.py

@@ -45,6 +45,29 @@
 )
 
 
+def normalize_reduced_loss_metrics(keys: list[str], values: list[float]) -> dict[str, float]:
+    """Normalize all-reduced metric numerators without dividing by zero.
+
+    Per-token training can legitimately produce a batch with no effective loss
+    tokens (for example, all responses are empty or fully masked). Such a batch
+    has zero metric numerators and represents a no-signal step, so report
+    zeros. A nonzero numerator with a zero denominator indicates an
+    inconsistent reducer and must fail loudly.
+    """
+    if len(keys) + 1 != len(values):
+        raise ValueError(f"Expected one denominator plus {len(keys)} metric values, got {len(values)} values.")
+
+    denominator = values[0]
+    numerators = values[1:]
+    if denominator == 0:
+        nonzero_keys = [key for key, value in zip(keys, numerators, strict=True) if value != 0]
+        if nonzero_keys:
+            raise RuntimeError(f"Zero loss-metric denominator with nonzero numerator(s): {nonzero_keys}.")
+        return dict.fromkeys(keys, 0.0)
+
+    return {key: value / denominator for key, value in zip(keys, numerators, strict=True)}
+
+
 def get_responses(
     logits: torch.Tensor,
     *,
@@ -114,7 +137,11 @@ def get_responses(
             else:
                 end += total_length
                 start = end - response_length
-            if response_length == total_length:
+            if response_length == 0:
+                # ``tokens[-0:]`` is the full prompt, not an empty slice.
+                logits_chunk = logits[0:0]
+                tokens_chunk = tokens[0:0]
+            elif response_length == total_length:
                 # SFT branch; see relax.utils.sft_utils.compute_sft_response_chunk.
                 from relax.utils.sft_utils import compute_sft_response_chunk
 

relax/backends/megatron/model.py

@@ -46,7 +46,7 @@
 
 from .checkpoint import load_checkpoint, save_checkpoint
 from .data import DataIterator, get_batch
-from .loss import loss_function
+from .loss import loss_function, normalize_reduced_loss_metrics
 from .model_provider import get_model_provider_func, wrap_model_provider_with_freeze
 
 
@@ -943,6 +943,48 @@ def _inject_meta(logits, *, _orig=f_partial, _meta=dcp_meta):
     return rollout_data
 
 
+def _is_global_zero_token_step(losses_reduced: list[dict[str, object]]) -> bool:
+    """Return True when the whole step has zero effective loss tokens.
+
+    An empty or fully-masked global batch connects a zero loss through
+    ``0 * logits.sum()``, so gradients are exactly zero. Running the optimizer
+    anyway would still move parameters through Adam momentum / weight decay and
+    advance the LR scheduler, despite there being no training signal. This
+    computes the globally reduced token count on the last pipeline stage and
+    broadcasts the decision to every rank so they agree on skipping the
+    optimizer and scheduler updates.
+    """
+    signal = torch.zeros(1, dtype=torch.int64, device=torch.cuda.current_device())
+    pp_size = mpu.get_pipeline_model_parallel_world_size()
+    if mpu.is_pipeline_last_stage(ignore_virtual=True):
+        # Sum the per-microbatch CP-local token counts, then reduce over DP+CP so
+        # every last-stage TP rank observes the same global count (mirrors the
+        # metric all-reduce below, which also uses the DP+CP group).
+        num_tokens_local = sum(x["values"][0] for x in losses_reduced)  # type: ignore[index]
+        torch.distributed.all_reduce(
+            num_tokens_local,
+            group=mpu.get_data_parallel_group(with_context_parallel=True),
+        )
+        signal[0] = 1 if num_tokens_local.item() == 0 else 0
+        if pp_size > 1:
+            # Non-last stages did not join the all-reduce; propagate the
+            # decision across the pipeline group so every rank agrees.
+            torch.distributed.broadcast(
+                signal,
+                src=pp_size - 1,
+                group=mpu.get_pipeline_model_parallel_group(),
+            )
+    elif pp_size > 1:
+        # Non-last stages must enter the broadcast so the collective completes;
+        # they keep the sentinel value overwritten by the last-stage decision.
+        torch.distributed.broadcast(
+            signal,
+            src=pp_size - 1,
+            group=mpu.get_pipeline_model_parallel_group(),
+        )
+    return bool(signal.item())
+
+
 def train_one_step(
     args: Namespace,
     rollout_id: int,
@@ -1190,33 +1232,47 @@ def forward_step(
     # double grad_scaler.update that the previous external prepare_grads() flow caused.
     # In fp16 with dynamic loss scaling, step() returns (False, None, None) on overflow.
     valid_step = True
-    update_successful, grad_norm, num_zeros_in_grad = optimizer.step()
-
-    if not getattr(args, "check_for_nan_in_loss_and_grad", True):
-        # fp16 with dynamic loss scaling auto-disables this flag (see Megatron arguments.py).
-        # Detect overflow via the documented (False, None, None) return signature.
-        found_inf_flag = not update_successful and grad_norm is None and num_zeros_in_grad is None
-        if found_inf_flag:
-            valid_step = False
-            current_scale = optimizer.get_loss_scale().item()
-            logger.warning(
-                "Inf found in gradients (step_id=%d, loss_scale=%s), skipping parameter "
-                "update (dynamic loss scaling will reduce scale)",
-                step_id,
-                current_scale,
-            )
-        else:
-            if isinstance(grad_norm, torch.Tensor):
-                valid_step = not (torch.isnan(grad_norm) or torch.isinf(grad_norm))
-            else:
-                valid_step = not (math.isnan(grad_norm) or math.isinf(grad_norm))
-
-    if valid_step:
-        # Update learning rate.
-        assert update_successful
-        opt_param_scheduler.step(increment=args.global_batch_size)
+    if _is_global_zero_token_step(losses_reduced):
+        # No effective loss tokens anywhere in the global batch. Gradients are
+        # exactly zero (the loss is zero-connected), so skip both the parameter
+        # and LR scheduler updates; momentum / weight decay must not move the
+        # model on a no-signal step. Every rank agrees because the decision is
+        # reduced over DP+CP and broadcast across the pipeline.
+        logger.warning(
+            "Training step %d has zero effective loss tokens globally; skipping optimizer and LR scheduler updates.",
+            step_id,
+        )
+        update_successful = True
+        grad_norm = 0.0
+        num_zeros_in_grad = None
     else:
-        grad_norm = float("nan")
+        update_successful, grad_norm, num_zeros_in_grad = optimizer.step()
+
+        if not getattr(args, "check_for_nan_in_loss_and_grad", True):
+            # fp16 with dynamic loss scaling auto-disables this flag (see Megatron arguments.py).
+            # Detect overflow via the documented (False, None, None) return signature.
+            found_inf_flag = not update_successful and grad_norm is None and num_zeros_in_grad is None
+            if found_inf_flag:
+                valid_step = False
+                current_scale = optimizer.get_loss_scale().item()
+                logger.warning(
+                    "Inf found in gradients (step_id=%d, loss_scale=%s), skipping parameter "
+                    "update (dynamic loss scaling will reduce scale)",
+                    step_id,
+                    current_scale,
+                )
+            else:
+                if isinstance(grad_norm, torch.Tensor):
+                    valid_step = not (torch.isnan(grad_norm) or torch.isinf(grad_norm))
+                else:
+                    valid_step = not (math.isnan(grad_norm) or math.isinf(grad_norm))
+
+        if valid_step:
+            # Update learning rate.
+            assert update_successful
+            opt_param_scheduler.step(increment=args.global_batch_size)
+        else:
+            grad_norm = float("nan")
 
     if critic_value_head_snapshot is not None:
         from relax.backends.megatron.ci_utils import assert_critic_value_head_updated
@@ -1246,16 +1302,17 @@ def forward_step(
         assert len(keys) + 1 == values.numel()
         torch.distributed.all_reduce(values, group=mpu.get_data_parallel_group(with_context_parallel=True))
 
-        loss_reduced = {}
         values = values.tolist()
         num_samples_or_tokens = values[0]
-        for key, value in zip(keys, values[1:], strict=False):
-            # No cp_size factor: num_samples_or_tokens is the all-reduced CP-local
-            # token count (per-token) or sample count, so each token/sample is
-            # already counted once. A `* cp_size` here would over-weight metrics by
-            # CP degree under dynamic CP (and is a no-op under static CP, where the
-            # count previously carried the cancelling cp factor).
-            loss_reduced[key] = value / num_samples_or_tokens
+        if num_samples_or_tokens == 0:
+            logger.warning(
+                "Training step %d has zero effective loss tokens; reporting zero loss metrics for this no-signal step.",
+                step_id,
+            )
+        # No cp_size factor: num_samples_or_tokens is the all-reduced CP-local
+        # token count (per-token) or sample count, so each token/sample is already
+        # counted once. A `* cp_size` here would over-weight metrics by CP degree.
+        loss_reduced = normalize_reduced_loss_metrics(keys, values)
         return loss_reduced, grad_norm
     return {}, grad_norm
 

tests/backends/megatron/test_rloo_cp_reduction.py

@@ -44,7 +44,7 @@ def _rloo_cp_worker(rank, world_size, init_file):
     try:
         torch.manual_seed(0)
         base_length = 2 * world_size * 4
-        response_lengths = [base_length, base_length * 2, base_length * 3]
+        response_lengths = [base_length, base_length * 2, 0]
         prompt_lengths = [2 * world_size * (sample_index + 1) for sample_index in range(3)]
         total_lengths = [
             prompt_length + response_length
@@ -71,7 +71,8 @@ def _rloo_cp_worker(rank, world_size, init_file):
         local_ownership = []
         for total_length, response_length, full_loss in zip(total_lengths, response_lengths, full_losses, strict=True):
             response_slices = _response_slices(total_length, response_length, rank, world_size)
-            local_losses.append(torch.cat([full_loss[response_slice] for response_slice in response_slices]))
+            local_parts = [full_loss[response_slice] for response_slice in response_slices]
+            local_losses.append(torch.cat(local_parts) if local_parts else full_loss.new_empty(0))
             ownership = torch.zeros(response_length, dtype=torch.int64)
             for response_slice in response_slices:
                 ownership[response_slice] += 1

tests/backends/megatron/test_rloo_policy_loss_dispatch.py

@@ -95,14 +95,36 @@ def test_policy_loss_function_dispatches_rloo_objective(monkeypatch):
     assert torch.allclose(log_probs.grad, -advantages)
 
 
-def test_rloo_unequal_lengths_use_global_token_scalar_and_gradient_oracle(monkeypatch):
-    """Exercise the production reducer and returned Megatron token normalizer
-    with unequal non-empty responses."""
-    response_lengths = [2, 4]
-    masks = [
-        torch.tensor([1.0, 1.0], dtype=torch.float64),
-        torch.tensor([1.0, 1.0, 1.0, 0.0], dtype=torch.float64),
-    ]
+@pytest.mark.parametrize(
+    ("response_lengths", "masks"),
+    [
+        (
+            [2, 4],
+            [
+                torch.tensor([1.0, 1.0], dtype=torch.float64),
+                torch.tensor([1.0, 1.0, 1.0, 0.0], dtype=torch.float64),
+            ],
+        ),
+        (
+            [2, 4],
+            [
+                torch.tensor([1.0, 0.0], dtype=torch.float64),
+                torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float64),
+            ],
+        ),
+        (
+            [2, 0],
+            [
+                torch.tensor([1.0, 1.0], dtype=torch.float64),
+                torch.empty(0, dtype=torch.float64),
+            ],
+        ),
+    ],
+    ids=["unequal-nonempty-responses", "unequal-with-fully-masked-response", "unequal-with-empty-response"],
+)
+def test_rloo_unequal_lengths_use_global_token_scalar_and_gradient_oracle(monkeypatch, response_lengths, masks):
+    """Exercise the production reducer and returned Megatron token
+    normalizer."""
     num_tokens = sum(response_lengths)
     log_probs = torch.tensor(
         [-0.2, -0.4, -0.1, -0.3, -0.5, -0.7][:num_tokens],
@@ -185,3 +207,99 @@ def test_rloo_unequal_lengths_use_global_token_scalar_and_gradient_oracle(monkey
     final_loss.backward()
     expected_gradient = -(advantages * flat_mask) / expected_num_tokens
     assert torch.allclose(log_probs.grad, expected_gradient)
+
+
+def test_get_responses_cp1_returns_matching_empty_chunks_for_empty_response():
+    args = SimpleNamespace(
+        qkv_format="thd",
+        allgather_cp=False,
+        rollout_temperature=1.0,
+        loss_type="policy_loss",
+    )
+    logits = torch.randn(1, 3, 11, dtype=torch.float32)
+    tokens = torch.tensor([3, 5, 7])
+
+    chunks = list(
+        loss_module.get_responses(
+            logits,
+            args=args,
+            unconcat_tokens=[tokens],
+            total_lengths=[3],
+            response_lengths=[0],
+            dynamic_cp_size=1,
+            dynamic_cp_rank=0,
+        )
+    )
+
+    assert len(chunks) == 1
+    logits_chunk, tokens_chunk = chunks[0]
+    assert logits_chunk.shape == (0, 11)
+    assert tokens_chunk.shape == (0,)
+
+
+def test_all_zero_token_loss_returns_zero_normalizer_and_numerator(monkeypatch):
+    log_probs = torch.tensor([-0.2, -0.4], dtype=torch.float64, requires_grad=True)
+    mask = torch.zeros(2, dtype=torch.float64)
+    monkeypatch.setattr(
+        loss_module,
+        "get_log_probs_and_entropy",
+        lambda *_args, **_kwargs: (
+            None,
+            {"log_probs": [log_probs], "entropy": [torch.zeros_like(log_probs)]},
+        ),
+    )
+    monkeypatch.setattr(loss_module, "resolve_opd_gather_topk_token_ids", lambda *_args, **_kwargs: None)
+    monkeypatch.setattr(loss_module, "compute_policy_opd_loss", lambda **_kwargs: (None, {}))
+    args = SimpleNamespace(
+        loss_type="policy_loss",
+        advantage_estimator="rloo",
+        calculate_per_token_loss=True,
+        qkv_format="thd",
+        recompute_loss_function=False,
+        allgather_cp=False,
+        global_batch_size=1,
+        true_on_policy_mode=False,
+        use_rollout_logprobs=False,
+        use_opsm=False,
+        get_mismatch_metrics=False,
+        use_tis=False,
+        custom_pg_loss_reducer_function_path=None,
+        entropy_coef=0.0,
+        use_kl_loss=False,
+    )
+    batch = {
+        "advantages": torch.tensor([1.0, 1.0], dtype=torch.float64),
+        "log_probs": [torch.zeros_like(log_probs)],
+        "response_lengths": [2],
+        "total_lengths": [3],
+        "unconcat_tokens": [torch.arange(3)],
+        "loss_masks": [mask],
+        "dynamic_cp_size": 1,
+        "dynamic_cp_rank": 0,
+    }
+
+    token_sum_loss, normalizer, logging = loss_module.loss_function(
+        args,
+        batch,
+        num_microbatches=1,
+        logits=torch.empty(1, 1, 1, dtype=torch.float64),
+    )
+
+    assert token_sum_loss.item() == 0.0
+    assert normalizer.item() == 0.0
+    assert logging["values"][0].item() == 0.0
+    assert torch.count_nonzero(logging["values"][1:]).item() == 0
+
+
+def test_zero_token_metrics_are_reported_as_zero_without_division():
+    keys = ["loss", "pg_loss", "ppo_kl"]
+    assert loss_module.normalize_reduced_loss_metrics(keys, [0.0, 0.0, -0.0, 0.0]) == {
+        "loss": 0.0,
+        "pg_loss": 0.0,
+        "ppo_kl": 0.0,
+    }
+
+
+def test_zero_token_metrics_reject_nonzero_numerator():
+    with pytest.raises(RuntimeError, match="nonzero numerator.*pg_loss"):
+        loss_module.normalize_reduced_loss_metrics(["loss", "pg_loss"], [0.0, 0.0, 1.0])