diff --git a/docs/en/examples/algorithms.md b/docs/en/examples/algorithms.md index 370e5adda..2333111cd 100644 --- a/docs/en/examples/algorithms.md +++ b/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 diff --git a/docs/zh/examples/algorithms.md b/docs/zh/examples/algorithms.md index 76a8b8231..b314fb845 100644 --- a/docs/zh/examples/algorithms.md +++ b/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, 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. diff --git a/relax/backends/megatron/loss.py b/relax/backends/megatron/loss.py index 170265a24..fcca9f356 100644 --- a/relax/backends/megatron/loss.py +++ b/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 diff --git a/relax/backends/megatron/model.py b/relax/backends/megatron/model.py index 077e461d9..dd2c75f93 100644 --- a/relax/backends/megatron/model.py +++ b/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 diff --git a/tests/backends/megatron/test_rloo_cp_reduction.py b/tests/backends/megatron/test_rloo_cp_reduction.py index 0d80b1eec..e2139bf7b 100644 --- a/tests/backends/megatron/test_rloo_cp_reduction.py +++ b/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 diff --git a/tests/backends/megatron/test_rloo_policy_loss_dispatch.py b/tests/backends/megatron/test_rloo_policy_loss_dispatch.py index a06d45c7b..44b36bdd2 100644 --- a/tests/backends/megatron/test_rloo_policy_loss_dispatch.py +++ b/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]) diff --git a/tests/backends/megatron/test_zero_token_step.py b/tests/backends/megatron/test_zero_token_step.py new file mode 100644 index 000000000..005cd5672 --- /dev/null +++ b/tests/backends/megatron/test_zero_token_step.py @@ -0,0 +1,208 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Zero-token no-signal step handling for the shared Megatron trainer.""" + +from types import SimpleNamespace + +import pytest + + +torch = pytest.importorskip("torch") +pytest.importorskip("megatron.core") + + +@pytest.fixture() +def model_module(): + from relax.backends.megatron import model as model_module + + return model_module + + +def _two_microbatch_losses(zero: bool) -> list[dict[str, object]]: + """Two microbatches; ``values[0]`` is the (CP-local) token count.""" + return ( + [ + {"values": torch.tensor([0.0, 1.0, 2.0])}, + {"values": torch.tensor([0.0, 3.0])}, + ] + if zero + else [ + {"values": torch.tensor([4.0, 1.0, 2.0])}, + {"values": torch.tensor([6.0, 3.0])}, + ] + ) + + +def _patch_mpu( + monkeypatch, + model_module, + *, + pp_size: int, + is_last_stage: bool, + all_reduce_impl, + broadcast_impl, +): + mpu = model_module.mpu + dist = model_module.torch.distributed + monkeypatch.setattr(mpu, "is_pipeline_last_stage", lambda ignore_virtual=False: is_last_stage) + monkeypatch.setattr(mpu, "get_data_parallel_group", lambda with_context_parallel=False: object()) + monkeypatch.setattr(mpu, "get_pipeline_model_parallel_world_size", lambda: pp_size) + monkeypatch.setattr(mpu, "get_pipeline_model_parallel_group", lambda: object()) + monkeypatch.setattr(dist, "all_reduce", all_reduce_impl) + monkeypatch.setattr(dist, "broadcast", broadcast_impl) + + +def test_is_global_zero_token_step_true_pp1_no_broadcast(model_module, monkeypatch): + """PP=1: every rank is the last stage, so the all-reduced count is already + consistent locally and no pipeline broadcast is required.""" + calls = {"all_reduce": 0, "broadcast": 0} + + def all_reduce(tensor, group=None): + calls["all_reduce"] += 1 + # Zero token count: leave the tensor as-is. + + def broadcast(tensor, src=0, group=None): + calls["broadcast"] += 1 + + _patch_mpu( + monkeypatch, + model_module, + pp_size=1, + is_last_stage=True, + all_reduce_impl=all_reduce, + broadcast_impl=broadcast, + ) + assert model_module._is_global_zero_token_step(_two_microbatch_losses(zero=True)) is True + assert calls == {"all_reduce": 1, "broadcast": 0} + + +def test_is_global_zero_token_step_false_pp1_no_broadcast(model_module, monkeypatch): + """PP=1 with a nonzero token count: the local all-reduced count is + nonzero.""" + calls = {"all_reduce": 0, "broadcast": 0} + + def all_reduce(tensor, group=None): + calls["all_reduce"] += 1 + tensor.fill_(10) # nonzero global token count + + def broadcast(tensor, src=0, group=None): + calls["broadcast"] += 1 + + _patch_mpu( + monkeypatch, + model_module, + pp_size=1, + is_last_stage=True, + all_reduce_impl=all_reduce, + broadcast_impl=broadcast, + ) + assert model_module._is_global_zero_token_step(_two_microbatch_losses(zero=False)) is False + assert calls == {"all_reduce": 1, "broadcast": 0} + + +def test_is_global_zero_token_step_last_stage_pp2_broadcasts(model_module, monkeypatch): + """PP>1: the last stage reduces the count then broadcasts the decision.""" + calls = {"all_reduce": 0, "broadcast": 0} + + def all_reduce(tensor, group=None): + calls["all_reduce"] += 1 + + def broadcast(tensor, src=0, group=None): + calls["broadcast"] += 1 + assert src == 1 # pp_size - 1 + + _patch_mpu( + monkeypatch, + model_module, + pp_size=2, + is_last_stage=True, + all_reduce_impl=all_reduce, + broadcast_impl=broadcast, + ) + assert model_module._is_global_zero_token_step(_two_microbatch_losses(zero=True)) is True + assert calls == {"all_reduce": 1, "broadcast": 1} + + +def test_is_global_zero_token_step_non_last_stage_pp2_enters_broadcast(model_module, monkeypatch): + """PP>1: a non-last stage skips the all-reduce but must join the broadcast + so the collective completes.""" + calls = {"all_reduce": 0, "broadcast": 0} + + def all_reduce(tensor, group=None): + calls["all_reduce"] += 1 + + def broadcast(tensor, src=0, group=None): + calls["broadcast"] += 1 + assert src == 1 + + _patch_mpu( + monkeypatch, + model_module, + pp_size=2, + is_last_stage=False, + all_reduce_impl=all_reduce, + broadcast_impl=broadcast, + ) + assert model_module._is_global_zero_token_step(_two_microbatch_losses(zero=True)) is False + assert calls == {"all_reduce": 0, "broadcast": 1} + + +def _make_args() -> SimpleNamespace: + return SimpleNamespace( + custom_megatron_before_train_step_hook_path=None, + ci_test=False, + enable_mtp_training=False, + check_for_nan_in_loss_and_grad=True, + global_batch_size=32, + dynamic_context_parallel=False, + use_dynamic_batch_size=False, + fully_async=False, + seq_length=512, + micro_batch_size=1, + decoder_seq_length=512, + attention_backend="flash", + ) + + +def test_train_one_step_skips_optimizer_and_scheduler_on_zero_token_step(model_module, monkeypatch): + """A global zero-token step must not update parameters or the LR + scheduler.""" + from megatron.core import mpu + + calls = {"optimizer_step": 0, "scheduler_step": 0} + + optimizer = SimpleNamespace( + step=lambda: calls.__setitem__("optimizer_step", calls["optimizer_step"] + 1) or (True, 0.0, 0), + zero_grad=lambda: None, + get_loss_scale=lambda: SimpleNamespace(item=lambda: 1.0), + param_groups=[], + ) + scheduler = SimpleNamespace( + step=lambda increment=None: calls.__setitem__("scheduler_step", calls["scheduler_step"] + 1) + ) + model = [SimpleNamespace(zero_grad_buffer=lambda: None)] + + losses_reduced = _two_microbatch_losses(zero=True) + + monkeypatch.setattr(model_module, "_is_global_zero_token_step", lambda losses: True) + monkeypatch.setattr(model_module, "get_args", lambda: _make_args()) + monkeypatch.setattr(model_module, "get_forward_backward_func", lambda: lambda **_kwargs: losses_reduced) + monkeypatch.setattr(model_module, "maybe_verify_critic_value_head_movement", lambda *a, **k: None) + monkeypatch.setattr(mpu, "is_pipeline_last_stage", lambda ignore_virtual=False: False) + monkeypatch.setattr(mpu, "get_virtual_pipeline_model_parallel_world_size", lambda: None) + + loss_reduced, grad_norm = model_module.train_one_step( + args=_make_args(), + rollout_id=0, + step_id=3, + data_iterator=[[]], + model=model, + optimizer=optimizer, + opt_param_scheduler=scheduler, + num_microbatches=1, + ) + + assert calls["optimizer_step"] == 0 + assert calls["scheduler_step"] == 0 + assert grad_norm == 0.0 + assert loss_reduced == {}