[bugfix] apply enable_global_average_loss to the list-wise rank loss

[tiankongdeguiji]

The bug

RankModel._loss_impl gated the list-wise request-count rescale on:

getattr(self._base_model_config, "enable_global_average_loss", False)

_base_model_config is the outer ModelConfig envelope (model.py:66), and the flag is defined only on the three HSTU messages in multi_task_rank.proto. Protobuf raises AttributeError for a field a message does not define, so the getattr default always yielded False. Verified on a live model:

master read  getattr(_base_model_config) -> False      # _base_model_config type: ModelConfig
PR read      _model_config.enable_...     -> <config>   # _model_config type: DlrmHSTUOneRank

So the rescale was dead code. Under DDP the list-wise term averaged over each rank's own request count and the cross-rank gradient average came out biased whenever those counts differ — which cost-based batching (sample_cost_field / batch_cost_size) makes routine, since it slices on cumulative cost rather than row count.

The fix

The root cause is a generic base class reaching for a model-specific config field, so the fix removes the reach rather than correcting the message name — that way the same mistake cannot recur in the branch.

DlrmHSTU.loss already loops over loss_cfg, so it picks the factor matching each loss's denominator — request count for the list-wise term, candidate count for the point-wise ones — and passes it through the existing loss_weight argument. Since the list-wise loss reduces to a scalar, the shared tail's torch.mean(scalar * loss_weight) is bit-identical to the multiply ListwiseRankLoss was doing internally, so the module's loss_weight parameter goes away (it now matches the other loss modules, which take no weights) and the branch in _loss_impl ends up shaped like its neighbours.

The activated math is the standard one: per-rank S_r/N_r weighted by N_r/N̄ makes DDP's averaged gradient the unbiased global request mean Σ S_r / Σ N_r.

Compatibility

Multi-rank training and eval with a listwise_rank_loss will see the list-wise loss scale change, since enable_global_average_loss defaults to true. One extra all_reduce per loss() call is added, gated on a task actually configuring a list-wise term so models without one pay nothing. Single-process runs are unaffected (every factor is exactly 1.0). The point-wise path is untouched — it always read the flag from the correct message.

Alternatives considered

Correcting the read to self._model_config would be a one-liner, but RankModel's inner config is MultiTower, DeepFM, … for most models, none of which define the flag — so it would still be a getattr probe in a base class for a field only some models own, i.e. the same bug class one rename away. Moving the decision to the model that owns the flag closes it.

Adding explicit sample_avg_weight / request_avg_weight parameters to _loss_impl was the other option, but it changes a signature shared by four callers to express what the existing loss_weight slot already carries.

Collective fusion

The rescale factors used to cost one all-reduce for the request axis plus one per task for the candidate axis — and every one of the per-task reductions carried the same number, because each task's jagged label spans the same candidates. Measured on the 3-task OneRank fixture:

all_reduce count : 4      values reduced : [2.0, 6.0, 6.0, 6.0]

Both counts are already available from the candidate lengths the model publishes (lengths.size(0) requests, lengths.sum() candidates), so fx_avg_batch_size becomes fx_avg_counts and reduces the pair together:

all_reduce count : 1      payload : [[2.0, 6.0]]

The saving is rank synchronizations, not bandwidth — a scalar collective costs a sync point, and three of them per step were buying nothing. It also removes the need to gate the request-axis reduction on whether any task configures a list-wise loss: one collective serves every loss of every task, so a task whose losses are all list-wise no longer pays for a candidate factor it never reads.

This leans on lengths.sum() == label.size(0) for every task. That is already required for the point-wise losses to run at all — logits_<task> and the jagged label are indexed against the same candidate rows, and a mismatch raises from BCEWithLogitsLoss (or from the list-wise segment ops) long before the weighting matters — so the invariant is enforced, not newly assumed.

Test Plan

• Added test_global_average_loss_rescales_each_term_by_its_own_denominator (dlrm_hstu_onerank_test.py). Runs loss() against a mocked peer rank holding 6 requests / 2 candidates versus this batch's 2 / 6, so the request ratio (0.5) and candidate ratio (1.5) differ and neither is 1.0, then asserts the list-wise term takes the former and the point-wise terms the latter. It also re-runs the flag-off model inside the same mock and asserts nothing moves, pinning the global_average gate.

  This closes a real hole: the enable_global_average_loss sweeps in dlrm_hstu_test.py, dlrm_hstu_onerank_test.py and ultra_hstu_test.py all run single-process, where both factors collapse to exactly 1.0, so nothing could observe the dead code. The test fails on the base commit, as a regression test for this bug should.

  It also asserts all_reduce.call_count == 1, which pins the fusion.

  Verified by mutation, three times: swapping the two factors in DlrmHSTU.loss fails the ratio assertions; dropping the global_average gate fails the flag-off assertion; splitting the request axis back into its own reduction fails the count assertion with 2 != 1. All pass after reverting.

• Rewrote FxAvgBatchSizeTest as FxAvgCountsTest, keeping the AVG-vs-SUM check and adding one that both axes share a single collective.

• Removed test_loss_weight_is_a_scalar_multiplier — it covered only the deleted module parameter. Every other test in that file, and the reference check in dlrm_hstu_onerank_test.py, already called the module with three arguments.

Run on 4x A10:

python -m unittest tzrec.models.dlrm_hstu_onerank_test tzrec.loss.listwise_rank_loss_test \
    tzrec.models.rank_model_test tzrec.utils.fx_util_test     # 49 tests, OK
python -m unittest tzrec.models.dlrm_hstu_test tzrec.models.ultra_hstu_test \
    tzrec.models.multi_task_rank_test                         # 12 tests, OK
pre-commit run --files <changed>                                              # all pass
pyrefly check                                                                 # 0 errors

🤖 Generated with Claude Code

https://claude.ai/code/session_01YSiYGTLEYMo8mq8hthZGJS

[github-actions]

Major — the refactor is not behavior-preserving: it activates a rescale that was dead code on master.

The removed block in RankModel._loss_impl gated the request-count factor on getattr(self._base_model_config, "enable_global_average_loss", False), but _base_model_config is the outer ModelConfig envelope (model.py:66), and model.proto does not define enable_global_average_loss on that message — it exists only on the three HSTU messages in multi_task_rank.proto. Protobuf raises AttributeError for undefined fields, so the getattr default always yielded False: on master the listwise term was never rescaled (the in-module factor was always None, and the caller's loss_weight was nulled).

Reading the flag from self._model_config here activates it. In multi-rank training and eval with a listwise_rank_loss configured and the proto default (enable_global_average_loss = true), the listwise term now gets multiplied by local_requests / avg_requests (≠ 1.0 on ragged batches) and one extra all_reduce runs per loss() call — loss values and the listwise/pointwise gradient balance shift relative to master. The new test corroborates the delta: on the base commit the listwise ratio would be 1.0, not 0.5, so test_global_average_loss_rescales_each_term_by_its_own_denominator would fail there.

The activated math itself is correct — per-rank S_r/N_r weighted by N_r/N̄ makes DDP's averaged gradient the unbiased global request mean Σ S_r / Σ N_r — and it finally makes the proto's documented promise ("enables loss averaging computation globally across all ranks") true for the listwise term too. So the change looks right; the framing is the issue. Suggest confirming the activation is intended and updating the PR/commit message: "Loss values are unchanged" holds only single-process, and per AGENTS.md this half is really a [bugfix] (root cause: the flag was read from the wrong config object) rather than a [refactor], so users upgrading mid-training know the listwise loss scale will change.

[github-actions]

Minor — the contract behind the deleted guard is now undocumented.

The removed loss_weight = None carried a comment explaining why ("it is sized off the candidate count, not the request count"). After this change the shared tail applies whatever loss_weight the caller passes to the request-level scalar. Nothing can go wrong today — only DlrmHSTU.loss reaches this branch (only the HSTU family publishes TARGET_REPEAT_INTERLEAVE_KEY) and it passes only None or the request-count scalar — but a future model that publishes the key and reuses RankModel.loss/MultiTaskRank.loss with per-sample weights would silently broadcast a candidate-sized vector against the request-level scalar instead of failing. A one-line # NOTE: at this branch preserving the deleted rationale would keep the invariant stated where it's enforced, at no structural cost.

[github-actions]

Minor / optional — two cheap hardenings for the new gate:

• The mock only exercises the flag-on model. If the global_average guard in DlrmHSTU.loss were dropped, every single-process factor still collapses to 1.0, so nothing in the suite would catch that regression. Recomputing base.loss inside the same mock and asserting it equals base_losses (or asserting dist_mock.all_reduce.call_count == 4) would pin the gate.
• _has_listwise_loss is an any() over all tasks' losses, but the listwise loss is wired only into is_click (as in every listwise test in this file). A variant with the listwise term on a later task — or on two tasks, which would also exercise the single request_avg_weight being reused across multiple listwise terms — would cover the actual scan.

[github-actions]

Review summary

Five parallel review passes (code quality, performance, test coverage, documentation accuracy, multi-process safety) plus manual cross-checks; all five completed.

The refactor itself is well executed: _loss_impl's listwise branch now matches its siblings, ListwiseRankLoss.forward matches the other loss modules, the dropped torch.fx.wrap(fx_avg_batch_size) in rank_model.py is safe (the only remaining caller is dlrm_hstu.py, which keeps its own module-level wrap — torch.fx.wrap registers per calling namespace, so that one is load-bearing), _has_listwise_loss is set for all three concrete HSTU classes before loss() can run, and the collective count stays config-derived so ranks cannot diverge. The new numeric test closes a real gap — every existing sweep is single-process, where both factors collapse to exactly 1.0 — and the mocked-peer construction is sound (every fx_avg_batch_size call hits a peer_of key; 0.5 / 1.5 are exact binary factors).

One major finding (inline on tzrec/models/dlrm_hstu.py): the change is not behavior-preserving in distributed training. On master the listwise request-count rescale was dead code — the getattr read enable_global_average_loss from the ModelConfig envelope, which does not define that field, so it always defaulted to False — and this PR activates it. The activated math is correct and matches the flag's documented intent, but multi-rank training and eval with a listwise term will see a different loss scale than on master (and one extra all_reduce per step), so "Loss values are unchanged" holds only single-process. Worth confirming the activation is intended and re-labelling that half as a latent bugfix in the PR/commit message.

Minor (inline): the request-level-weight contract in _loss_impl is now implicit (the explaining comment was deleted with the guard); optional test hardening for the flag-off gate and for listwise losses on tasks other than the first.

Optional follow-ups (non-blocking, in the spirit of the PR body's own defused perf question):

• A task whose losses are all listwise still pays a per-task sample_avg_weight all_reduce whose result is then discarded — pre-existing waste, but the _has_listwise_loss gating pattern introduced here is exactly what a follow-up would need to skip it.
• tzrec/utils/fx_util_test.py's docstring line "so every rank model now consumes it" is made stale by the rank_model.py import removal — after this PR only DlrmHSTU.loss consumes fx_avg_batch_size.