[feat] support session_name in listwise_rank_loss and jrc_loss - #684
Conversation
4ddd9bd to
55f1317
Compare
tiankongdeguiji
changed the title
55f1317 to
5d61131
Compare
5be9687 to
4b38761
Compare
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
| elif loss_type == "listwise_rank_loss": | ||
| # The module averages over requests itself, so the per-sample | ||
| # `reduction` of the surrounding losses does not apply. | ||
| # The module averages over lists, so a per-sample weight (which | ||
| # the callers signal with reduction="none") has nowhere to apply. | ||
| if reduction == "none": |
There was a problem hiding this comment.
This guard misfires for multi-task towers that set only a scalar weight: multi_task_rank.py:84 derives reduction = "none" if self.has_weight(...), and has_weight returns True for HasField("weight") (the scalar tower weight) alone. So task_towers { losses { listwise_rank_loss { session_name: "request_id" } } weight: 0.5 } — a natural config on the multi-tower/DBMTL models this PR newly enables — fails at model init with an error claiming per-sample weights are configured when none are.
Since a scalar weight is numerically equivalent under either reduction for the pointwise losses, multi_task_rank.py could derive the reduction from sample_weight_name / task_space_indicator_label only, leaving has_weight to drive the loss_weight branch as today. A regression test with tower weight + listwise_rank_loss would pin it.
There was a problem hiding this comment.
Fixed. MultiTaskRank.init_loss now derives reduction="none" from sample_weight_name / task_space_indicator_label only; has_weight still drives the loss_weight branch (a scalar tower weight is identical under either reduction). Regression test: multi_task_rank_test.test_listwise_loss_with_tower_weight (tower weight: 0.5 + listwise builds and scales both losses; sample_weight_name on the same tower is still refused).
| """ | ||
| ids = batch.sparse_features[BASE_DATA_GROUP][name].to_padded_dense(1)[:, 0] | ||
| if TARGET_REPEAT_INTERLEAVE_KEY in predictions: | ||
| ids = ids.repeat_interleave(predictions[TARGET_REPEAT_INTERLEAVE_KEY]) |
There was a problem hiding this comment.
repeat_interleave with tensor repeats and no output_size must read sum(repeats) back to the host to allocate the output — a hidden device→host sync on every call. This is the exact pitfall scatter_ops.lengths_to_index's docstring calls out, and this PR fixed every other call site to pass output_size; this helper (which feeds the new session_name loss path and grouped_auc/grouped_xauc on exactly the models that publish TARGET_REPEAT_INTERLEAVE_KEY) was missed. The output size is statically known at both call sites (pred.size(0) in _loss_impl, label.size(0) in _update_metric_impl), so threading it through as a parameter is a small change. Pattern is carried over from the old jrc code, but it's now shared and newly reachable from listwise_rank_loss, so worth fixing here.
There was a problem hiding this comment.
Fixed. _group_ids(..., num_samples) passes output_size to repeat_interleave; _list_index threads it through, and the call sites pass pred.size(0) / label.size(0).
| # per-candidate weight the sibling losses take. | ||
| losses[loss_name] = self._loss_modules[loss_name](pred, label, lengths) | ||
| lengths, index = _list_index( | ||
| predictions, batch, loss_cfg.listwise_rank_loss.session_name |
There was a problem hiding this comment.
With session_name set on a DlrmHSTU-family model, lengths here is now the session count (from torch.unique), but DlrmHSTU.loss (dlrm_hstu.py:303-307) still applies request_avg_weight — built from the per-request TARGET_REPEAT_INTERLEAVE_KEY counts — to any listwise_rank_loss, by loss type alone. enable_global_average_loss defaults to true, and loss.md explicitly advertises the session_name-override-on-DlrmHSTU combination, so this is reachable: the loss is (sum over lists) / num_lists * (num_requests / avg_requests), and when the sessions/requests ratio varies across ranks the DDP-averaged gradient is no longer the unbiased global mean that listwise_rank_loss.py:137-145 documents. The rescale factor would need to be session-count-based when session_name is set. This interaction is also untested (dlrm_hstu_onerank_test only covers the per-request default); one case with session_name spanning two requests of one user would pin both the repeat branch of _group_ids and the rescale.
There was a problem hiding this comment.
Resolved by narrowing the feature rather than adding a second rescale: listwise_rank_loss.session_name is now only for models that do not publish per-request candidate counts — on the DlrmHSTU family each request is the list and setting the field raises at loss time — so the loss's denominator and the request_avg_weight rescale always agree. jrc_loss keeps its pre-existing DlrmHSTU behaviour (request id repeated over candidates); it averages over candidates, so sample_avg_weight stays correct. Tests: rank_model_test.test_published_counts_define_the_lists (uses the published counts; refuses session_name) and test_jrc_session_spans_requests_on_published_counts. Proto comment and docs updated.
| predictions[TARGET_REPEAT_INTERLEAVE_KEY] | ||
| ) | ||
| losses[loss_name] = self._loss_modules[loss_name](pred, label, session_id) | ||
| lengths, index = _list_index( |
There was a problem hiding this comment.
The jrc branch was rewired in this PR (session_ids → _list_index lengths/index) but has no CPU model-level test: no jrc case exists in rank_model_test.py, and the only coverage is the GPU-lane integration config (tzrec/tests/configs/dlrm_hstu_kuairand_1k.config), which asserts training completes, not loss values. A regression in this 8-line branch (swapped lengths/index, wrong suffix) would only surface as an AUC drift. test_listwise_rank_loss_with_session is the exact template; the _reference_session_loss helper in jrc_loss_test.py is directly reusable for the expected value, and it would also cover jrc's FX_TRACE lane, which lost its old @torch.fx.wrap helpers.
There was a problem hiding this comment.
Added rank_model_test.test_jrc_loss_with_session (NORMAL + FX_TRACE, against a direct JRCLoss call on the torch.unique grouping) and test_jrc_session_spans_requests_on_published_counts (a test model that publishes per-request counts: one shared id → one 6-row session, two ids → two sessions).
| Returns: | ||
| torch.Tensor: 1-d id tensor aligned with the samples. | ||
| """ | ||
| ids = batch.sparse_features[BASE_DATA_GROUP][name].to_padded_dense(1)[:, 0] |
There was a problem hiding this comment.
Two validation gaps in the user-facing session_name path, both silent or opaque:
- A typo'd name surfaces as a bare
KeyError: 'user_is'from torchrec'sKeyedTensor.__getitem__, with no pointer tolosses.listwise_rank_loss.session_name— and only after the full data pipeline and distributed process-group startup. to_padded_dense(1)[:, 0]silently takes the first value of a multi-valued feature (merging rows that share only the first tag) and pads empty rows with 0 (coalescing all id-less rows into one bogus session together with real id 0).torch.uniquethen builds wrong lists with no error anywhere.
Since init_loss already holds the model's features, an init-time check that the feature exists (and is single-valued per sample) would turn both into clear config-time errors. The pattern is pre-existing for jrc/grouped_auc, but this PR makes session_name the primary — and for non-DlrmHSTU models the only — list source, so the reach is new.
There was a problem hiding this comment.
Not in this PR. The gap is pre-existing and shared by jrc_loss.session_name and grouped_auc/grouped_xauc.grouping_key; a name check belongs to all three together (follow-up), and single-valuedness is a property of the data, not checkable at init.
| @@ -12,16 +12,13 @@ | |||
| """Per-request list-wise InfoNCE over a jagged candidate list (paper 2.5).""" | |||
There was a problem hiding this comment.
The module and class docstrings still describe only the per-request mode ("Per-request list-wise InfoNCE over a jagged candidate list", "over the candidates of one request", "negatives of a candidate are the other candidates of the same request", "this stays inside the jagged layout"). After this PR a list can also be a session spanning multiple requests of one user, in arbitrary row order — the forward docstring was updated but the class-level summary now misdescribes the general behavior. Worth rewording to "one list (a request's candidates or a session's samples)".
There was a problem hiding this comment.
Fixed — module and class docstrings now describe a list as a request's candidates or a session's samples.
| ``` | ||
|
|
||
| 对于该损失函数，要求同一个session_id的样本尽量在一个batch中进行训练，在一个session中尽量要求样本保持有序。 | ||
| 对于该损失函数，要求同一个session_id的样本落在同一个batch中，在一个session中尽量要求样本保持有序。 |
There was a problem hiding this comment.
这条"在一个session中尽量要求样本保持有序"已经过时：新的实现（torch.unique 分组 + 顺序无关的 segment 归约，jrc_loss_test.py 的 test_index_in_any_order 也验证了任意顺序不变性）使得 session 内样本顺序不影响损失值。同一段第 106 行也已说明"不要求相邻"。建议删去排序子句，只保留"落在同一个batch中"的要求（SQL 示例中按 time_stamp 排序对损失值也不再有意义）。
There was a problem hiding this comment.
已修复，删掉了该子句。
| 该损失通常与逐点损失（如binary_cross_entropy）搭配使用，用同级的`weight`调节相对权重，经验值0.1量级 | ||
| 注意事项： | ||
|
|
||
| 1. 一个list需要同时含有至少一个正样本和一个负样本才会计入该项损失，被屏蔽的list仍计入分母。使用`session_name`时要求同一个session的样本落在同一个batch内（不要求相邻），对于非DlrmHSTU系列模型样本构造方式与[jrc_loss](#jrc_loss)相同（`DISTRIBUTE BY session_id SORT BY session_id`），并尽量加大`batch_size`；若样本未按session分组，绝大多数list只含一条样本而被屏蔽，表现为该项损失接近0 |
There was a problem hiding this comment.
这里建议补一条：listwise_rank_loss 不能与 sample_weight_name / task_space_indicator_label 同任务同时配置 —— 本 PR 在 _init_loss_impl 中将其改为模型构建期的硬性 ValueError，而这两者在 mmoe/dbmtl/ple/pepnet 文档中都有介绍，且本文档正好推荐在这些多塔模型上启用 listwise。用户按文档组合配置会在构建期报错，而文档没有预警。
There was a problem hiding this comment.
有意不写入文档（维护者的决定）。修掉 tower 标量 weight 的误报后，该限制只覆盖真正的逐样本权重（sample_weight_name / task_space_indicator_label），并且构建期的报错信息会直接点名这两个字段。
Review summaryOverall this is a high-quality change. Verified independently during review: the Findings (details inline, priority order):
None of these block the core rewrite; 1 and 2 deserve a fix (or at least a decision) before merge. 🤖 Generated with Claude Code |
Group list-wise losses by a session feature on rank models that do not publish per-request candidate counts (the DlrmHSTU family keeps each request as the list). jrc_loss computes its session term with scatter reductions instead of a batch-by-batch mask, so it no longer scales quadratically with batch size; the reductions move to tzrec/ops/ scatter_ops.py under torch_scatter names. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01WsLRKMqJQHRmUKrGy7ivdJ
2fc18cf to
f8963ff
Compare
2 participants
Summary
listwise_rank_lossis model-agnostic in code, but its only source of lists was the per-request candidate count that the DlrmHSTU family publishes, so every other rank model hit aValueError. This PR addsListwiseRankLoss.session_name(mirroringjrc_lossand EasyRec'sLISTWISE_RANK_LOSS), and makes both list-wise losses share one list-grouping helper:listwise_rank_loss.session_nameis for models that do not publish per-request candidate counts: samples of the batch sharing the feature value form one list, in any order (torch.unique(return_inverse=True, return_counts=True)gives per-rowindexand counts; the reductions are scatter-based so nothing is permuted). On the DlrmHSTU family each request is the list and settingsession_nameis refused, which keeps theenable_global_average_lossrescale (per-request counts) consistent with the loss's denominator.jrc_loss.session_namekeeps its pre-existing DlrmHSTU behaviour (the request-level id is repeated over the candidates, so a session may span a user's requests); jrc averages over candidates, so its rescale is unaffected.jrc_lossis rewritten on the same segment ops. Its session term is, per sample,softplus(logsumexp(competitors) − own logit)where the log-sum-exp is one value per session (negatives' positive-logits for a positive, positives' negative-logits for a negative), so it takes two segment reductions instead of the previouseye(B)/B×Bmask /[P,B]+[N,B]tiles. Same values and gradients as before (tested against the mask formulation), reductionnoneandmeanunchanged, and the cost stops growing quadratically with batch size — which matters since the doc tells users to pushbatch_sizeas high as memory allows.Other notes:
jrc_loss,grouped_aucandgrouped_xaucis factored into_group_ids;_list_indexbuilds on it. GAUC stays id-based (it groups across the eval set).jagged_tensors.pytotzrec/ops/scatter_ops.pyasscatter_sum / scatter_max / scatter_logsumexp(src, index, dim_size)(torch_scatter naming). They reduce by a per-row index and never needed a jagged layout; now that they take indices fromtorch.uniqueon non-contiguous batches, thejagged_prefix and the vestigiallengthsargument (only read for.size(0)) promised a layout that isn't there. In PyTorchsegment_*(torch.segment_reduce) means contiguous-by-lengths, so that name would have been wrong the other way.jagged_segment_idsbecomeslengths_to_indexin the same module (jaggedlengths→ torch_scatterindex, PyG'sptr2indexanalogue), and the per-row map is calledindexat every call site and in both losses'forward; pure rename acrossjrc_loss,listwise_rank_loss,onerank_sd,task_tower.sample_weight_name/task_space_indicator_label) are rejected atinit_lossforlistwise_rank_loss(a mean over lists has nowhere to apply them; EasyRec silently ignores them).MultiTaskRank.init_lossnow derivesreduction="none"from those two fields only, so a scalar towerweight(numerically identical under either reduction) no longer trips the guard.jrc_losskeeps supporting per-sample weights via reductionnone._group_idspassesoutput_sizetorepeat_interleave, sparing the device-to-host sync on the models that publish candidate counts.JRCLoss.session_namebecomesoptional(DlrmHSTU can rely on per-request sessions). Existing configs that set it keep their behaviour.label_is_logits/transform_fn/scale_logits/listwise_distill_loss(its distillation surface; the multiplicative half ofscale_logitsis already the learnable temperature).Test Plan
tzrec.loss.jrc_loss_test: the two existing cases (same expected value 0.7199) on the newlengthsAPI; equivalence of values and gradients against the original batch-mask formulation on sessions with both classes / all positives / all negatives / a single row; permutation invariance with arbitraryindex.tzrec.loss.listwise_rank_loss_test: existing masking-branch table plus permutation invariance withindex.tzrec.models.rank_model_test: listwise on an unsorted 3-session batch underNORMALandFX_TRACE(one mixed list counted, an all-positive list and a singleton masked but in the denominator); jrc withsession_nameunder both graph types against a directJRCLosscall; a test model that publishes per-request counts: listwise uses them and refusessession_name, jrc'ssession_namerepeats the request id over the candidates (one 6-row session vs. two sessions);ValueErrorat construction with sample weights;ValueErrorat loss time with neither list source.tzrec.models.multi_task_rank_test: towerweight: 0.5+listwise_rank_lossbuilds and scales both losses by 0.5;sample_weight_nameon the same tower is refused.scatter_logsumexpandJRCLoss(incl. one-class sessions) track fp32 within bf16 tolerance.tzrec.ops.scatter_ops_test: empty groups (sum → 0, max → 0), bf16 round-trip, index in arbitrary order,scatter_logsumexpvalues and gradients against per-grouptorch.logsumexp/ softmax.tzrec.models.multi_tower_test,dbmtl_test,modules.task_tower_test,ops.jagged_tensors_test,modules.gr.onerank_sd_test,dlrm_hstu_test,dlrm_hstu_onerank_test(A10) — pass.pre-commitandpyrefly checkclean.Benchmark
Old (batch mask) vs new (segment ops)
JRCLoss, forward + backward, one A10, sessions of 20 samples, 30% positives, 20 timed iterations after warm-up. "any order" is thesession_namepath:torch.unique(return_inverse=True, return_counts=True)plus the loss with explicitsegment_ids. Loss values agree to 5 decimals at every size.lengths)(time / peak allocated memory of the loss step)
🤖 Generated with Claude Code
https://claude.ai/code/session_01WsLRKMqJQHRmUKrGy7ivdJ