[refactor] add generic LossConfig.weight, drop ListwiseRankLoss.alpha

[tiankongdeguiji]

Motivation

ListwiseRankLoss.alpha was documented as "weight of the list-wise term relative to the task's other losses" — a loss-mixing weight, not a hyperparameter of the InfoNCE objective. The tell is in the code: temperature_init / learnable_temperature reach the ListwiseRankLoss module's constructor, while alpha never did — it was a scalar multiply at the call site in RankModel._loss_impl. (BinaryFocalLoss.alpha and JRCLoss.alpha are module constructor args and stay untouched.)

It existed only because the weight ladder stopped at the task level. task_towers.weight / task_configs.weight scale every loss of a task alike, so they cannot trade one loss of a task against another — the test that pinned alpha said so in its own docstring. The consequence was that no other loss could be reweighted at all: a task carrying binary_cross_entropy + binary_focal_loss had no knob, and a SID model summed its reconstruction, commitment and contrastive terms 1:1:1.

What this does

Adds the missing rung, LossConfig.weight (default 1.0), mirroring EasyRec's Loss.weight, and deletes ListwiseRankLoss.alpha:

losses { binary_cross_entropy {} }
losses {
    listwise_rank_loss { temperature_init: 0.07 }
    weight: 0.1
}

The weight is honored at the single return tail of each of the three _loss_impl bodies — RankModel, MatchModel, BaseSidModel — which all seven loss() loops route through, so no consumer can silently ignore it. Two smaller decisions:

• The listwise_rank_loss branch no longer returns early. It returned early to skip the per-candidate loss_weight, which is sized off the candidate count rather than the request count; it now nulls loss_weight and falls through to the shared tail. That also stops MultiTaskRank's task weight (folded into the tensor loss_weight) from being silently dropped for this loss — unreachable today, since only DlrmHSTU publishes the candidate counts and applies task weight in its own loss(), but the trap is gone.
• The multiply is skipped when weight == 1.0, so every existing config produces an identical traced graph.

The reported per-loss value is the weighted one, as alpha was, which keeps the invariant that the logged terms sum to total_loss.

Compatibility

A config that still sets alpha now fails at load with a protobuf ParseError naming the field. This is deliberate: silently accepting it would move the effective default from 0.1 to 1.0, i.e. a 10x stronger list-wise term that would surface days later as a training-dynamics regression rather than as a config error. The fix is mechanical — alpha: 0.1 becomes weight: 0.1, one nesting level out. No in-repo .config carries it.

Test Plan

• New test_loss_weight_scales_loss in rank_model_test.py, over NORMAL and FX_TRACE graphs with and without sample weights. It reuses the exact batch and expected values of test_binary_classification_model, so a weight folded into the sample weights instead — which would renormalize away — would not reproduce them; the FX rows confirm the Python float constant-folds.
• New test_loss_weight_scales_loss in match_model_test.py and test_loss_weight_scales_only_its_own_term in sid_rqvae_test.py, covering the other two _loss_impl bodies. The SID one also pins that the weight touches its own term and not its sibling.
• test_listwise_loss_is_scaled_by_alpha renamed to test_listwise_loss_is_scaled_by_loss_weight; its docstring was the one place that documented why alpha existed, so it now documents why LossConfig.weight does.
• Full tzrec/models (238) and tzrec/loss (43) suites green on an A10, including the GPU-gated dlrm_hstu_onerank_test (18). pre-commit run and pyrefly check clean.

Docs

docs/source/models/loss.md gains a 损失权重 weight section (how it composes with the task-level weight, that the logged value is weighted, and a caution against tuning it alongside use_pareto_loss_weight) and a listwise_rank_loss section, which was missing entirely; its stale supported-loss list is refreshed. dlrm_hstu_onerank.md's example and parameter table move the knob out one level.

Alternatives considered

A learnable per-loss weight (EasyRec's learn_loss_weight) was rejected — pe_mtl_loss already covers adaptive weighting across tasks. Putting the knob on FusionSubTaskConfig was rejected as the wrong granularity: it stays OneRank-specific and leaves every other multi-loss task unserved.

🤖 Generated with Claude Code

https://claude.ai/code/session_01QBwVb17wQxFd5Kdb8mkwPS

[github-actions]

Consider reserving the removed field's number and name:

Suggested change

	message ListwiseRankLoss {
	message ListwiseRankLoss {
	reserved 1;
	reserved "alpha";

The repo already does this in train.proto (reserved 4, 16, 17; plus the names). The risk today is low — configs are persisted text-only (config_util.save_message → MessageToString), so no binary artifact can reinterpret field 1 — but reserving is zero-cost and prevents a future field from silently reusing the number, and prevents alpha from being reintroduced with a different meaning (a name the sibling JRCLoss / BinaryFocalLoss messages already use for something else).

[github-actions]

Minor: this composition formula isn't actually asserted. Every task here runs at the default task weight 1.0 — _task_configs only exposes a weight knob on is_like, which carries no listwise loss — so the task_weight × LossConfig.weight product is unpinned anywhere (the loss.md claim 它与任务级的weight相乘生效 also leans on it). Risk is low: both factors are individually pinned and the composition is two plain multiplies (dlrm_hstu.py:304 over the shared _loss_impl tail). But since this test is the natural place for it, consider building base/scaled with a non-1.0 weight on the is_click task and asserting the product on listwise_rank_loss_is_click — or soften the docstring to what's asserted.

[github-actions]

Nit: true for the two modules this test uses — SidReconLoss/SidCommitmentLoss are pure functions with no parameters/buffers — but SidContrastiveLoss has learnable logit_scale_* parameters, so "the sid loss modules are stateless" overstates the invariant the trick relies on. Extending this two-models/one-prediction pattern to a contrastive config would compare module-owned temperatures across models. Suggest narrowing the comment, e.g. "The recon/commitment loss modules are stateless, so ...".

[github-actions]

Review summary

Ran a five-area review (code quality, performance, test coverage, documentation accuracy, security/robustness) over the diff. Overall this is a clean, well-scoped refactor: the missing rung in the weight ladder is added at exactly the right level, and the migration story is sound. LGTM with three minor inline comments (proto reserved, an unasserted docstring formula, a test-comment nit).

Verified across the change:

• Single tail, no bypasses. Every LossConfig consumer routes through one of the three modified tails: RankModel / MultiTaskRank / RocketLaunching / DlrmHSTU → rank_model._loss_impl; MatchModel / HSTUMatch → match_model._loss_impl; SID models → _sid_loss_impl. No loss path skips the weight.
• Trace/compile safety. loss_cfg.weight != 1.0 is a Python-float config check that constant-folds under symbolic_trace (matching existing patterns like label_smoothing > 0); the new FX_TRACE test rows pin this. Default configs trace an identical graph — and actually drop one scalar multiply, since the old code applied alpha unconditionally. Tensor × Python-float promotion is mechanically identical to the old * alpha site.
• {loss_type: loss} in sid_model reproduces the three previously hardcoded keys byte-for-byte (WhichOneof("sid_loss") returns the oneof field names verbatim).
• Fail-closed removal holds. All training entry points parse with allow_unknown_field=False, so a stale alpha: raises at load; the only lenient path (predict from a scripted model) never consumes loss configs. Configs are persisted text-only (save_message → MessageToString), so the unreserved field number 1 has no binary artifact to misinterpret today — the inline reserved suggestion is hygiene, not an active bug.
• Compatibility of the 0.1 → 1.0 default shift. ListwiseRankLoss landed on master only in [feat] add DlrmHSTUOneRank (OneRank generative ranking) #661 (Sep 18, 2026) and is not in the latest tagged release (v1.4.0), so configs that omitted alpha and would silently get a 10× stronger list-wise term can only exist in the last hours of master adoption. The deliberate ParseError choice is well-founded.
• Docs match code. Weight composes multiplicatively with task weight on both the MultiTaskRank (folded into loss_weight) and DlrmHSTU (multiplied after _loss_impl) paths; logged values and loss metrics are the weighted ones (_update_loss_metric_impl reads the post-weight dict, and the new rank test's metric assertion pins it); the pareto caution is accurate — pe_mtl_loss consumes the already-weighted losses dict.
• Tests. The hardcoded constants (0.6356/0.6762, 0.71080) match the sibling tests they mirror, no assertion would pass with the weight ignored, and the new parameterization follows the repo's name_func / GPU-gating conventions.

Two non-blocking notes:

1. The PR body says the fall-through "stops MultiTaskRank's task weight (folded into the tensor loss_weight) from being silently dropped for this loss." As written, the code still drops it — the branch nulls loss_weight before the shared tail, exactly as the early return did (and this is numerically unobservable single-process, since RankModel.loss renormalizes sample weights to mean 1). The real improvements are that the listwise term now receives LossConfig.weight and the drop is explicit and commented. Unreachable today either way; consider rewording the body so it doesn't imply a behavior change there.
2. loss.md's "各损失项之和即为total_loss" is accurate for everything observable (_log_train and TensorBoard both compute total as the sum of the logged terms), but under use_pareto_loss_weight the optimized total is the pareto combination, not that sum. Item 3 directly below warns about pareto re-weighting, so this is fine as-is; a one-clause qualifier would make it airtight.