From f8963ff9f04d925ffd1e7f9829be37de940c0070 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=A4=A9=E9=82=91?= Date: Tue, 22 Sep 2026 11:08:40 +0800 Subject: [PATCH] [feat] support session_name in listwise_rank_loss and jrc_loss 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 Claude-Session: https://claude.ai/code/session_01WsLRKMqJQHRmUKrGy7ivdJ --- docs/source/models/loss.md | 19 ++- tzrec/loss/jrc_loss.py | 103 +++++------- tzrec/loss/jrc_loss_test.py | 129 ++++++++++---- tzrec/loss/listwise_rank_loss.py | 57 ++++--- tzrec/loss/listwise_rank_loss_test.py | 14 ++ tzrec/models/multi_task_rank.py | 9 +- tzrec/models/multi_task_rank_test.py | 84 +++++++++- tzrec/models/rank_model.py | 138 ++++++++++----- tzrec/models/rank_model_test.py | 233 ++++++++++++++++++++++++++ tzrec/modules/gr/onerank_sd.py | 20 +-- tzrec/modules/task_tower.py | 6 +- tzrec/ops/jagged_tensors.py | 67 -------- tzrec/ops/jagged_tensors_test.py | 47 ------ tzrec/ops/scatter_ops.py | 77 +++++++++ tzrec/ops/scatter_ops_test.py | 89 ++++++++++ tzrec/protos/loss.proto | 25 ++- 16 files changed, 819 insertions(+), 298 deletions(-) create mode 100644 tzrec/ops/scatter_ops.py create mode 100644 tzrec/ops/scatter_ops_test.py diff --git a/docs/source/models/loss.md b/docs/source/models/loss.md index 9c62331a6..7ccacc7e7 100644 --- a/docs/source/models/loss.md +++ b/docs/source/models/loss.md @@ -103,6 +103,8 @@ model_config { 适用二分类任务的损失函数，其对应的任务num_class必须是2。该损失函数除了关注样本目标自身分类的正确性，还会关注在同一个batch的同一个session中，所有正样本的概率要尽可能的大于所有负样本的概率。 https://arxiv.org/abs/2208.06164 +session的来源与[listwise_rank_loss](#listwise_rank_loss)相同：设置`session_name`时，同一个batch内该字段取值相同的样本构成一个session（不要求相邻）；不设置时，DlrmHSTU系列模型以每个请求的候选序列为一个session，其他排序模型必须设置`session_name` + 配置如下 ``` @@ -115,7 +117,7 @@ model_config { } ``` -对于该损失函数，要求同一个session_id的样本尽量在一个batch中进行训练，在一个session中尽量要求样本保持有序。 +对于该损失函数，要求同一个session_id的样本落在同一个batch中。 我们使用sql如下方式构造样本,该数据集的session_name是user_id @@ -131,14 +133,21 @@ SORT BY user_id asc,time_stamp asc ## listwise_rank_loss -请求粒度的listwise排序损失（InfoNCE），同一个请求内的其他候选互为负样本。该损失函数要求模型发布每个请求的候选数以及该任务的`logits_`预测，目前只有DlrmHSTU系列满足，详见[dlrm_hstu_onerank](dlrm_hstu_onerank.md)。 一个请求需要同时含有至少一个正样本和一个负样本才会计入该项损失 +listwise排序损失（InfoNCE），同一个list内的其他样本互为负样本。list的来源有两种： + +- DlrmHSTU系列模型：模型会发布每个请求的候选数，一个请求的候选序列即为一个list，无需额外配置（设置`session_name`会报错），详见[dlrm_hstu_onerank](dlrm_hstu_onerank.md) +- 其他排序模型（多塔、DBMTL等）：必须通过`session_name`指定分组字段（如request_id或user_id），同一个batch内该字段取值相同的样本构成一个list。该字段需为模型的稀疏特征 配置如下 ``` model_config { + losses { + binary_cross_entropy {} + } losses { listwise_rank_loss { + session_name: "user_id" temperature_init: 0.07 learnable_temperature: true } @@ -149,10 +158,14 @@ model_config { 参数说明： +1. session_name: list分组的字段名，仅用于非DlrmHSTU系列模型（DlrmHSTU系列模型按请求分组，不可设置） 1. temperature_init: softmax温度初始值，logits会乘以 1 / temperature，默认值0.07 1. learnable_temperature: 温度是否可学习（训练中会被clamp防止溢出），默认值true -该损失通常与逐点损失（如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 +1. 该损失只约束list内的相对打分，不保证概率校准，通常与逐点损失（如binary_cross_entropy）搭配使用，用同级的`weight`调节相对权重，经验值0.1量级 ## pe_mtl_loss diff --git a/tzrec/loss/jrc_loss.py b/tzrec/loss/jrc_loss.py index a74635ab6..7be61c7fe 100644 --- a/tzrec/loss/jrc_loss.py +++ b/tzrec/loss/jrc_loss.py @@ -10,20 +10,19 @@ # limitations under the License. +from typing import Optional + import torch +import torch.nn.functional as F from torch import Tensor from torch.nn import CrossEntropyLoss from torch.nn.modules.loss import _Loss +from tzrec.ops.scatter_ops import lengths_to_index, scatter_logsumexp -@torch.fx.wrap -def _label_mask(labels: torch.Tensor) -> torch.Tensor: - return torch.eye(labels.size(0), dtype=torch.int64, device=labels.device) - - -@torch.fx.wrap -def _diag_index(labels: torch.Tensor) -> torch.Tensor: - return torch.arange(0, labels.size(0), dtype=torch.int64, device=labels.device) +# Logit assigned to samples that must not compete: exp underflows to exactly +# 0 after the segment max shift, so they drop out of the log-sum-exp. +_MASKED_LOGIT = -1e9 class JRCLoss(_Loss): @@ -31,6 +30,13 @@ class JRCLoss(_Loss): https://arxiv.org/abs/2208.06164 + The session term of the paper is a softmax of each positive against the + negatives of its session (on the positive logit) and of each negative + against the positives of its session (on the negative logit). Both + reduce to ``softplus(logsumexp(competitors) - own logit)``, where the + log-sum-exp is one value per session, so the term costs two segment + reductions instead of a batch-by-batch session mask. + Args: alpha (float): cross entropy loss weight. reduction (str, optional): Specifies the reduction to apply to the @@ -52,14 +58,22 @@ def forward( self, logits: Tensor, labels: Tensor, - session_ids: Tensor, + lengths: Tensor, + index: Optional[Tensor] = None, ) -> Tensor: """JRC loss. + Without ``index``, ``logits`` and ``labels`` are laid out + session by session in ``lengths`` order, each session contiguous. + Args: logits: a `Tensor` with shape [batch_size, 2]. labels: a `Tensor` with shape [batch_size]. - session_ids: a `Tensor` with shape [batch_size]. + lengths: a `Tensor` with shape [num_sessions], samples per + session, summing to batch_size. + index: a `Tensor` with shape [batch_size], session of each + sample (torch_scatter's ``index``), for samples in arbitrary + order. Return: loss: a `Tensor` with shape [batch_size] if reduction is 'none', @@ -67,54 +81,27 @@ def forward( """ ce_loss = self._ce_loss(logits, labels) - batch_size = labels.shape[0] - mask = torch.eq(session_ids.unsqueeze(1), session_ids.unsqueeze(0)).float() - diag_index = _diag_index(labels) + if index is None: + index = lengths_to_index(lengths, output_size=logits.size(0)) logits_neg, logits_pos = logits[:, 0], logits[:, 1] - diag = _label_mask(labels) - pos_num = torch.sum(labels) - neg_num = batch_size - pos_num - - # first, we calculate pos sample loss in during the session. - pos_mask_index = torch.where(labels == 1.0)[0] - pos_diag_label = torch.index_select(diag_index, 0, pos_mask_index) - # pyre-ignore [6] - logits_pos = logits_pos.unsqueeze(0).tile([pos_num, 1]) - pos_session_mask = torch.index_select(mask, 0, pos_mask_index) - # pyre-ignore [6] - y_pos = labels.unsqueeze(0).tile([pos_num, 1]) - diag_pos = torch.index_select(diag, 0, pos_mask_index) - # we mask not in the same session, is diagonal and is positive. - logits_pos = ( - logits_pos + ((1 - pos_session_mask) + (1 - diag_pos) * y_pos) * -1e9 - ) - loss_pos = self._ce_loss(logits_pos, pos_diag_label) - - # next, we calculate neg sample loss in during the session. - neg_mask_index = torch.where(labels == 0.0)[0] - neg_diag_label = torch.index_select(diag_index, 0, neg_mask_index) - # neg_num is a 0-d integer tensor, which tile accepts via __index__. - # pyrefly: ignore[no-matching-overload] - logits_neg = logits_neg.unsqueeze(0).tile([neg_num, 1]) - neg_session_mask = torch.index_select(mask, 0, neg_mask_index) - # pyrefly: ignore[no-matching-overload] - y_neg = (1 - labels).unsqueeze(0).tile([neg_num, 1]) - diag_neg = torch.index_select(diag, 0, neg_mask_index) - # we mask not in the same session, is diagonal and is negative. - logits_neg = ( - logits_neg + ((1 - neg_session_mask) + (1 - diag_neg) * y_neg) * -1e9 + is_pos = labels == 1 + # Competitors of a positive: the session's negatives on the positive + # logit; of a negative: the session's positives on the negative logit. + neg_competitors = torch.where(is_pos, _MASKED_LOGIT, logits_pos) + pos_competitors = torch.where(is_pos, logits_neg, _MASKED_LOGIT) + num_sessions = lengths.size(0) + lse_neg = scatter_logsumexp( + neg_competitors.unsqueeze(-1), index, num_sessions + ).squeeze(-1) + lse_pos = scatter_logsumexp( + pos_competitors.unsqueeze(-1), index, num_sessions + ).squeeze(-1) + ge_loss = torch.where( + is_pos, + F.softplus(lse_neg.index_select(0, index) - logits_pos), + F.softplus(lse_pos.index_select(0, index) - logits_neg), ) - loss_neg = self._ce_loss(logits_neg, neg_diag_label) - if self._reduction != "none": - loss_pos = loss_pos * pos_num / batch_size - loss_neg = loss_neg * neg_num / batch_size - ge_loss = loss_pos + loss_neg - else: - ge_loss = torch.zeros_like(labels, dtype=torch.float) - ge_loss.index_put_(torch.where(labels == 1.0), loss_pos) - ge_loss.index_put_(torch.where(labels == 0.0), loss_neg) - - loss = self._alpha * ce_loss + (1 - self._alpha) * ge_loss - # pyre-ignore [7] - return loss + ge_loss = ge_loss.mean() + + return self._alpha * ce_loss + (1 - self._alpha) * ge_loss diff --git a/tzrec/loss/jrc_loss_test.py b/tzrec/loss/jrc_loss_test.py index 401bf83ff..43201ec39 100644 --- a/tzrec/loss/jrc_loss_test.py +++ b/tzrec/loss/jrc_loss_test.py @@ -13,53 +13,114 @@ import unittest import torch +from parameterized import parameterized from tzrec.loss.jrc_loss import JRCLoss +from tzrec.utils.test_util import parameterized_name_func + +_LOGITS = torch.tensor( + [ + [0.9, 0.1], + [0.5, 0.5], + [0.3, 0.7], + [0.2, 0.8], + [0.8, 0.2], + [0.55, 0.45], + [0.33, 0.67], + [0.55, 0.45], + ], + dtype=torch.float32, +) +_LABELS = torch.tensor([0, 0, 1, 1, 0, 0, 1, 1]) +_LENGTHS = torch.tensor([4, 4]) + + +def _reference_session_loss( + logits: torch.Tensor, labels: torch.Tensor, session_ids: torch.Tensor +) -> torch.Tensor: + """Per-sample session term as the original batch-mask formulation. + + Each positive competes with the negatives of its session on the positive + logit, each negative with the positives of its session on the negative + logit, the sample itself being the softmax target. + """ + mask = session_ids.unsqueeze(1) == session_ids.unsqueeze(0) + is_pos = labels == 1 + same_and_other_class = mask & (is_pos.unsqueeze(1) != is_pos.unsqueeze(0)) + keep = same_and_other_class | torch.eye(labels.numel(), dtype=torch.bool) + own_channel = torch.where(is_pos, 1, 0) + rows = logits[:, own_channel].T + rows = rows.masked_fill(~keep, -1e9) + return torch.nn.functional.cross_entropy( + rows, torch.arange(labels.numel()), reduction="none" + ) class JRCLossTest(unittest.TestCase): def test_jrc_loss(self) -> None: - loss_class = JRCLoss() - logits = torch.tensor( - [ - [0.9, 0.1], - [0.5, 0.5], - [0.3, 0.7], - [0.2, 0.8], - [0.8, 0.2], - [0.55, 0.45], - [0.33, 0.67], - [0.55, 0.45], - ], - dtype=torch.float32, - ) - labels = torch.tensor([0, 0, 1, 1, 0, 0, 1, 1]) - session_ids = torch.tensor([1, 1, 1, 1, 2, 2, 2, 2], dtype=torch.int8) - loss = loss_class(logits, labels, session_ids) + loss = JRCLoss()(_LOGITS, _LABELS, _LENGTHS) self.assertEqual(0.7199, round(loss.item(), 4)) - -class JRCLossTestReduceNone(unittest.TestCase): def test_jrc_loss_reduce_none(self) -> None: - loss_class = JRCLoss(reduction="none") - logits = torch.tensor( + loss = JRCLoss(reduction="none")(_LOGITS, _LABELS, _LENGTHS) + self.assertEqual((8,), tuple(loss.shape)) + self.assertEqual(0.7199, round(torch.mean(loss).item(), 4)) + + @parameterized.expand( + [ + # sessions with both classes, all positives, all negatives, one row [ - [0.9, 0.1], - [0.5, 0.5], - [0.3, 0.7], - [0.2, 0.8], - [0.8, 0.2], - [0.55, 0.45], - [0.33, 0.67], - [0.55, 0.45], + [5, 3, 4, 1, 6], + [1, 0, 0, 1, 0, 1, 1, 1, 0, 0, 0, 0, 1, 0, 1, 1, 0, 0, 1], ], - dtype=torch.float32, + [[19], [0] * 18 + [1]], + [[7, 7, 5], [1] * 7 + [0] * 7 + [1, 0, 1, 0, 1]], + ], + name_func=parameterized_name_func, + ) + def test_matches_batch_mask_formulation(self, lengths, labels) -> None: + torch.manual_seed(0) + lengths_t = torch.tensor(lengths) + labels_t = torch.tensor(labels) + logits = torch.randn(labels_t.numel(), 2, requires_grad=True) + session_ids = torch.repeat_interleave(torch.arange(len(lengths)), lengths_t) + + loss = JRCLoss(alpha=0.3, reduction="none")(logits, labels_t, lengths_t) + ce = torch.nn.functional.cross_entropy(logits, labels_t, reduction="none") + expected = 0.3 * ce + 0.7 * _reference_session_loss( + logits, labels_t, session_ids ) - labels = torch.tensor([0, 0, 1, 1, 0, 0, 1, 1]) - session_ids = torch.tensor([1, 1, 1, 1, 2, 2, 2, 2], dtype=torch.int8) - loss = loss_class(logits, labels, session_ids) + torch.testing.assert_close(loss, expected) - self.assertEqual(0.7199, round(torch.mean(loss).item(), 4)) + (grad,) = torch.autograd.grad(loss.sum(), logits) + (expected_grad,) = torch.autograd.grad(expected.sum(), logits) + torch.testing.assert_close(grad, expected_grad) + + def test_index_in_any_order(self) -> None: + torch.manual_seed(0) + lengths = torch.tensor([5, 3, 4]) + labels = torch.tensor([1, 0, 0, 1, 0, 1, 1, 1, 0, 0, 0, 1]) + logits = torch.randn(labels.numel(), 2) + perm = torch.randperm(labels.numel()) + index = torch.repeat_interleave(torch.arange(3), lengths)[perm] + + loss = JRCLoss(reduction="none")(logits[perm], labels[perm], lengths, index) + expected = JRCLoss(reduction="none")(logits, labels, lengths)[perm] + torch.testing.assert_close(loss, expected) + + def test_bf16_logits_track_fp32(self) -> None: + """bf16 keeps the masked logit finite, incl. one-class sessions.""" + torch.manual_seed(0) + lengths = torch.tensor([5, 3, 4, 1]) + labels = torch.tensor([1, 0, 0, 1, 0, 1, 1, 1, 0, 0, 0, 0, 1]) + logits = torch.randn(labels.numel(), 2) + loss_fn = JRCLoss(reduction="none") + loss = loss_fn(logits.bfloat16(), labels, lengths) + self.assertEqual(loss.dtype, torch.bfloat16) + self.assertTrue(torch.isfinite(loss).all()) + torch.testing.assert_close( + loss.float(), loss_fn(logits, labels, lengths), atol=5e-2, rtol=2e-2 + ) if __name__ == "__main__": diff --git a/tzrec/loss/listwise_rank_loss.py b/tzrec/loss/listwise_rank_loss.py index 002c4b662..7d16e6604 100644 --- a/tzrec/loss/listwise_rank_loss.py +++ b/tzrec/loss/listwise_rank_loss.py @@ -9,19 +9,20 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Per-request list-wise InfoNCE over a jagged candidate list (paper 2.5).""" +"""List-wise InfoNCE over the samples of one list (paper 2.5). + +A list is a request's candidates (models that publish per-request counts) +or the samples of one session (grouped by a session feature). +""" import math +from typing import Optional import torch from torch import nn from torch.nn.modules.loss import _Loss -from tzrec.ops.jagged_tensors import ( - jagged_segment_ids, - jagged_segment_max, - jagged_segment_sum, -) +from tzrec.ops.scatter_ops import lengths_to_index, scatter_max, scatter_sum from tzrec.utils.fx_util import fx_size0_max1 # `torch.fx.wrap` registers by name in the *calling* module's globals; the @@ -35,16 +36,16 @@ class ListwiseRankLoss(_Loss): - """Softmax cross-entropy over the candidates of one request. + """Softmax cross-entropy over the samples of one list. - The negatives of a candidate are the other candidates of the *same* - request, so this stays inside the jagged layout and needs no cross-rank - all-gather -- unlike an in-batch contrastive loss such as + The negatives of a sample are the other samples of the *same* list, and + a list never crosses a rank, so this needs no cross-rank all-gather -- + unlike an in-batch contrastive loss such as :class:`~tzrec.loss.sid_contrastive_loss.SidContrastiveLoss`, gathering - other ranks here would add candidates that are not negatives of this - list at all. + other ranks here would add samples that are not negatives of this list + at all. - Two kinds of request contribute nothing and are masked out rather than + Two kinds of list contribute nothing and are masked out rather than branched on, so the shapes stay data-independent (``torch.compile`` friendly): @@ -52,7 +53,7 @@ class ListwiseRankLoss(_Loss): * no negative -- the objective degenerates to "make all scores equal", which is a gradient with no ranking information in it. - A request with several positives uses the multi-positive form, i.e. the + A list with several positives uses the multi-positive form, i.e. the mean of the positives' log-probabilities. Consumed through ``LossConfig.listwise_rank_loss`` in a task's @@ -87,12 +88,13 @@ def forward( logits: torch.Tensor, labels: torch.Tensor, lengths: torch.Tensor, + index: Optional[torch.Tensor] = None, ) -> torch.Tensor: """Compute the list-wise InfoNCE term. - ``logits``, ``labels`` and ``lengths`` must all be in the same - request order, with the candidates of a request laid out - contiguously. + Without ``index``, ``logits`` and ``labels`` are laid out + request by request in ``lengths`` order, each request's candidates + contiguous. Args: logits (torch.Tensor): ``(total,)`` per-candidate score. @@ -100,31 +102,34 @@ def forward( non-zero value counts as a positive. lengths (torch.Tensor): ``(B,)`` candidates per request, summing to ``total``. + index (torch.Tensor, optional): ``(total,)`` request of each + candidate (torch_scatter's ``index``), for candidates in + arbitrary order. Returns: torch.Tensor: scalar loss -- the mean of per-request losses over *all* ``B`` requests (masked-out requests contribute 0), so the denominator matches the ``enable_global_average_loss`` rescale. """ - segment_ids = jagged_segment_ids(lengths, output_size=logits.size(0)) + if index is None: + index = lengths_to_index(lengths, output_size=logits.size(0)) # Clamp before exp so a large temperature can't overflow to +Inf. scale = self.logit_scale.clamp(max=_LOGIT_SCALE_MAX).exp() scaled = (logits * scale).unsqueeze(-1) - maxes = jagged_segment_max(scaled.detach(), lengths, segment_ids) - shifted = scaled - maxes.index_select(0, segment_ids) + num_lists = lengths.size(0) + maxes = scatter_max(scaled.detach(), index, num_lists) + shifted = scaled - maxes.index_select(0, index) # A non-empty segment always sums to >= 1, because the row holding # the segment max contributes exp(0) == 1. So the clamp is exact # where it matters and only rewrites empty segments, where log(0) # would otherwise leak -inf into the tensor. - denom = jagged_segment_sum(torch.exp(shifted), lengths, segment_ids).clamp( - min=1.0 - ) - log_probs = shifted - torch.log(denom).index_select(0, segment_ids) + denom = scatter_sum(torch.exp(shifted), index, num_lists).clamp(min=1.0) + log_probs = shifted - torch.log(denom).index_select(0, index) positives = (labels != 0).to(log_probs.dtype).unsqueeze(-1) - num_pos = jagged_segment_sum(positives, lengths, segment_ids) - pos_log_prob = jagged_segment_sum(log_probs * positives, lengths, segment_ids) + num_pos = scatter_sum(positives, index, num_lists) + pos_log_prob = scatter_sum(log_probs * positives, index, num_lists) num_candidates = lengths.to(num_pos.dtype).unsqueeze(-1) valid = ((num_pos > 0) & (num_pos < num_candidates)).to(log_probs.dtype) diff --git a/tzrec/loss/listwise_rank_loss_test.py b/tzrec/loss/listwise_rank_loss_test.py index c8263fd77..3daafb155 100644 --- a/tzrec/loss/listwise_rank_loss_test.py +++ b/tzrec/loss/listwise_rank_loss_test.py @@ -246,6 +246,20 @@ def test_integer_labels_are_accepted(self) -> None: module(logits, torch.tensor([0.0, 1.0, 0.0, 0.0]), lengths), ) + def test_index_in_any_order(self) -> None: + torch.manual_seed(0) + lengths = torch.tensor([2, 5, 3, 4]) + labels = torch.tensor([1, 0] + [0, 1, 1, 0, 0] + [1, 1, 1] + [0, 0, 0, 1]) + logits = torch.randn(labels.numel()) + perm = torch.randperm(labels.numel()) + index = torch.repeat_interleave(torch.arange(4), lengths)[perm] + + loss = ListwiseRankLoss() + torch.testing.assert_close( + loss(logits[perm], labels[perm], lengths, index), + loss(logits, labels, lengths), + ) + if __name__ == "__main__": unittest.main() diff --git a/tzrec/models/multi_task_rank.py b/tzrec/models/multi_task_rank.py index 8cb5bfd51..3b7432451 100644 --- a/tzrec/models/multi_task_rank.py +++ b/tzrec/models/multi_task_rank.py @@ -81,7 +81,14 @@ def init_loss(self) -> None: """Initialize loss modules.""" for task_tower_cfg in self._task_tower_cfgs: tower_name = task_tower_cfg.tower_name - reduction = "none" if self.has_weight(task_tower_cfg) else "mean" + # Only a per-sample weight needs the unreduced loss; the scalar + # tower weight scales a mean just the same. + reduction = ( + "none" + if task_tower_cfg.HasField("sample_weight_name") + or task_tower_cfg.HasField("task_space_indicator_label") + else "mean" + ) for loss_cfg in task_tower_cfg.losses: self._init_loss_impl( loss_cfg, diff --git a/tzrec/models/multi_task_rank_test.py b/tzrec/models/multi_task_rank_test.py index b68cbe3cb..951bb427b 100644 --- a/tzrec/models/multi_task_rank_test.py +++ b/tzrec/models/multi_task_rank_test.py @@ -14,10 +14,11 @@ import torch from parameterized import parameterized -from torchrec import KeyedTensor +from torchrec import KeyedJaggedTensor, KeyedTensor from tzrec.datasets.utils import BASE_DATA_GROUP, Batch from tzrec.features.feature import BaseFeature +from tzrec.loss.listwise_rank_loss import ListwiseRankLoss from tzrec.models.model import TrainWrapper from tzrec.models.multi_task_rank import MultiTaskRank from tzrec.protos import loss_pb2, metric_pb2, model_pb2 @@ -167,6 +168,87 @@ def test_multi_task_rank_model(self, graph_type, t2_loss_weight, task_space): atol=1e-4, ) + def test_listwise_loss_with_tower_weight(self): + """A scalar tower weight is not a per-sample weight. + + It must neither trip the listwise guard at construction nor change + the reduction the point-wise sibling sees; a genuine per-sample + weight on the same tower is still refused. + """ + + def model_config(**tower_kwargs): + return model_pb2.ModelConfig( + simple_multi_task=multi_task_rank_pb2.SimpleMultiTask( + task_towers=[ + TaskTower( + tower_name="t1", + label_name="label1", + weight=0.5, + losses=[ + loss_pb2.LossConfig( + binary_cross_entropy=loss_pb2.BinaryCrossEntropy() + ), + loss_pb2.LossConfig( + listwise_rank_loss=loss_pb2.ListwiseRankLoss( + session_name="id_a", + learnable_temperature=False, + ) + ), + ], + **tower_kwargs, + ) + ] + ) + ) + + model = TrainWrapper( + _TestMultiTaskRankModel( + model_config=model_config(), features=[], labels=["label1"] + ) + ) + ids = torch.tensor([1, 2, 1, 2, 1, 3]) + torch.manual_seed(0) + logits = torch.randn(6) + label = torch.tensor([0, 1, 1, 1, 0, 0]) + batch = Batch( + dense_features={ + BASE_DATA_GROUP: KeyedTensor.from_tensor_list( + keys=["int_a"], tensors=[logits.unsqueeze(1)] + ) + }, + sparse_features={ + BASE_DATA_GROUP: KeyedJaggedTensor.from_lengths_sync( + keys=["id_a"], values=ids, lengths=torch.ones(6, dtype=torch.int64) + ) + }, + labels={"label1": label}, + ) + _, (losses, _, _) = model(batch) + + _, index, lengths = torch.unique(ids, return_inverse=True, return_counts=True) + expected_listwise = ListwiseRankLoss(learnable_temperature=False)( + logits, label, lengths, index + ) + expected_bce = torch.nn.functional.binary_cross_entropy_with_logits( + logits, label.float() + ) + torch.testing.assert_close( + losses["listwise_rank_loss_t1"], 0.5 * expected_listwise + ) + torch.testing.assert_close( + losses["binary_cross_entropy_t1"], 0.5 * expected_bce + ) + + with self.assertRaisesRegex(ValueError, "per-sample weights"): + TrainWrapper( + _TestMultiTaskRankModel( + model_config=model_config(sample_weight_name="w"), + features=[], + labels=["label1"], + sample_weights=["w"], + ) + ) + if __name__ == "__main__": unittest.main() diff --git a/tzrec/models/rank_model.py b/tzrec/models/rank_model.py index 5280076ae..af6a7c499 100644 --- a/tzrec/models/rank_model.py +++ b/tzrec/models/rank_model.py @@ -10,7 +10,7 @@ # limitations under the License. from collections import OrderedDict -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Optional, Tuple import torch import torchmetrics @@ -45,6 +45,75 @@ def _update_tensor_dict( tensor_dict[key] = new_tensor +def _group_ids( + predictions: Dict[str, torch.Tensor], batch: Batch, name: str, num_samples: int +) -> torch.Tensor: + """Per-sample value of the id feature ``name`` (a session or user id). + + A model that publishes per-request candidate counts carries the feature + once per request, so it is repeated to line up with the candidate rows. + + Args: + predictions (dict): the model predictions. + batch (Batch): the input batch. + name (str): the sparse feature holding the id. + num_samples (int): number of samples, the statically known size of + the result (spares ``repeat_interleave`` a device-to-host sync). + + Returns: + torch.Tensor: ``(num_samples,)`` id tensor aligned with the samples. + """ + 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], output_size=num_samples + ) + return ids + + +def _list_index( + predictions: Dict[str, torch.Tensor], + batch: Batch, + session_name: str, + num_samples: int, +) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """Per-list sample counts and the list index of each row for a list-wise loss. + + With ``session_name`` set, the samples of the batch that share that + feature value form one list, in any order; on a model that publishes + per-request candidate counts this widens the list beyond one request + (``jrc_loss`` allows that, ``listwise_rank_loss`` rejects it). Without + it, the published per-request counts define the lists. + + Args: + predictions (dict): the model predictions. + batch (Batch): the input batch. + session_name (str): the list-grouping feature, or empty. + num_samples (int): number of samples in the batch. + + Returns: + tuple: ``(lengths, index)`` -- ``(num_lists,)`` sample count of + each list and the ``(num_samples,)`` list of each row in the + torch_scatter sense, the latter ``None`` when the lists are + contiguous in row order. + """ + if session_name: + _, index, lengths = torch.unique( + _group_ids(predictions, batch, session_name, num_samples), + return_inverse=True, + return_counts=True, + ) + return lengths, index + lengths = predictions.get(TARGET_REPEAT_INTERLEAVE_KEY) + if lengths is None: + raise ValueError( + "a list-wise loss needs per-request candidate counts " + "(predictions[TARGET_REPEAT_INTERLEAVE_KEY]), which this model does " + "not publish; set session_name instead." + ) + return lengths, None + + def _is_classification_loss(loss_cfg: LossConfig) -> bool: loss_type = loss_cfg.WhichOneof("loss") return loss_type in [ @@ -217,8 +286,14 @@ def _init_loss_impl( elif loss_type == "l2_loss": self._loss_modules[loss_name] = nn.MSELoss(reduction=reduction) 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": + raise ValueError( + "listwise_rank_loss averages over lists and does not support " + "per-sample weights; drop sample_weight_name / " + "task_space_indicator_label from its task." + ) self._loss_modules[loss_name] = ListwiseRankLoss( temperature_init=loss_cfg.listwise_rank_loss.temperature_init, learnable_temperature=loss_cfg.listwise_rank_loss.learnable_temperature, @@ -260,30 +335,28 @@ def _loss_impl( elif loss_type == "jrc_loss": assert num_class == 2, f"num_class must be 2 when loss type is {loss_type}" pred = predictions["logits" + suffix] - session_id = batch.sparse_features[BASE_DATA_GROUP][ - loss_cfg.jrc_loss.session_name - ].to_padded_dense(1)[:, 0] - if TARGET_REPEAT_INTERLEAVE_KEY in predictions: - session_id = session_id.repeat_interleave( - predictions[TARGET_REPEAT_INTERLEAVE_KEY] - ) - losses[loss_name] = self._loss_modules[loss_name](pred, label, session_id) + lengths, index = _list_index( + predictions, batch, loss_cfg.jrc_loss.session_name, pred.size(0) + ) + losses[loss_name] = self._loss_modules[loss_name]( + pred, label, lengths, index + ) elif loss_type == "l2_loss": pred = predictions["y" + suffix] losses[loss_name] = self._loss_modules[loss_name](pred, label) elif loss_type == "listwise_rank_loss": pred = predictions["logits" + suffix] - lengths = predictions.get(TARGET_REPEAT_INTERLEAVE_KEY) - if lengths is None: + session_name = loss_cfg.listwise_rank_loss.session_name + if session_name and TARGET_REPEAT_INTERLEAVE_KEY in predictions: raise ValueError( - "listwise_rank_loss needs per-request candidate counts " - "(predictions[TARGET_REPEAT_INTERLEAVE_KEY]), which " - "this model does not publish." + "listwise_rank_loss.session_name is only for models that do " + "not publish per-request candidate counts; this model's lists " + "are its requests." ) - # NOTE: this loss is a mean over requests, so a loss_weight - # reaching the tail below must be request-level too, not the - # per-candidate weight the sibling losses take. - losses[loss_name] = self._loss_modules[loss_name](pred, label, lengths) + lengths, index = _list_index(predictions, batch, session_name, pred.size(0)) + losses[loss_name] = self._loss_modules[loss_name]( + pred, label, lengths, index + ) else: raise ValueError(f"loss[{loss_type}] is not supported yet.") if loss_weight is not None: @@ -416,10 +489,6 @@ def _update_metric_impl( oneof_metric_cfg = getattr(metric_cfg, metric_type) metric_name = metric_type + suffix - base_sparse_feat = None - if metric_type in ["grouped_auc", "grouped_xauc"]: - base_sparse_feat = batch.sparse_features[BASE_DATA_GROUP].to_dict() - if metric_type == "auc": pred = ( predictions["probs" + suffix] @@ -445,27 +514,18 @@ def _update_metric_impl( if num_class == 1 else predictions["probs1" + suffix] ) - # pyre-ignore [16] - grouping_key = base_sparse_feat[ - oneof_metric_cfg.grouping_key - ].to_padded_dense(1)[:, 0] - if TARGET_REPEAT_INTERLEAVE_KEY in predictions: - grouping_key = grouping_key.repeat_interleave( - predictions[TARGET_REPEAT_INTERLEAVE_KEY] - ) + grouping_key = _group_ids( + predictions, batch, oneof_metric_cfg.grouping_key, label.size(0) + ) self._metric_modules[metric_name].update(pred, label, grouping_key) elif metric_type == "xauc": pred = predictions["y" + suffix] self._metric_modules[metric_name].update(pred, label) elif metric_type == "grouped_xauc": pred = predictions["y" + suffix] - grouping_key = base_sparse_feat[ - oneof_metric_cfg.grouping_key - ].to_padded_dense(1)[:, 0] - if TARGET_REPEAT_INTERLEAVE_KEY in predictions: - grouping_key = grouping_key.repeat_interleave( - predictions[TARGET_REPEAT_INTERLEAVE_KEY] - ) + grouping_key = _group_ids( + predictions, batch, oneof_metric_cfg.grouping_key, label.size(0) + ) self._metric_modules[metric_name].update(pred, label, grouping_key) elif metric_type == "normalized_entropy": pred = predictions["probs" + suffix] diff --git a/tzrec/models/rank_model_test.py b/tzrec/models/rank_model_test.py index 579b553d5..aeebc800e 100644 --- a/tzrec/models/rank_model_test.py +++ b/tzrec/models/rank_model_test.py @@ -16,8 +16,11 @@ from parameterized import param, parameterized from torchrec import JaggedTensor, KeyedJaggedTensor, KeyedTensor +from tzrec.constant import TARGET_REPEAT_INTERLEAVE_KEY from tzrec.datasets.utils import BASE_DATA_GROUP, Batch from tzrec.features.feature import BaseFeature +from tzrec.loss.jrc_loss import JRCLoss +from tzrec.loss.listwise_rank_loss import ListwiseRankLoss from tzrec.models.model import TrainWrapper from tzrec.models.rank_model import RankModel from tzrec.protos import loss_pb2, metric_pb2, model_pb2 @@ -45,6 +48,19 @@ def predict(self, batch: Batch) -> Dict[str, torch.Tensor]: return self._output_to_prediction(y) +class _TestJaggedModel(_TestClassficationModel): + """Publishes per-request candidate counts like the DlrmHSTU family. + + Sparse features are carried once per request (2 rows), dense logits and + labels once per candidate (6 rows). + """ + + def predict(self, batch: Batch) -> Dict[str, torch.Tensor]: + predictions = super().predict(batch) + predictions[TARGET_REPEAT_INTERLEAVE_KEY] = torch.tensor([2, 4]) + return predictions + + class _TestRegressionModel(RankModel): def __init__( self, @@ -450,6 +466,223 @@ def test_multi_classification_model(self, graph_type): metric_result["accuracy"], expected_acc, rtol=1e-4, atol=1e-4 ) + @parameterized.expand( + [[TestGraphType.NORMAL], [TestGraphType.FX_TRACE]], + name_func=parameterized_name_func, + ) + def test_listwise_rank_loss_with_session(self, graph_type): + """``session_name`` groups the rows of a flat batch into lists. + + Rows are deliberately not sorted by session. Session 1 holds one + positive and two negatives and is the only list that counts; session + 2 is all-positive and session 3 has a single row, so both are masked + out but still count in the mean's denominator. + """ + model_config = model_pb2.ModelConfig( + losses=[ + loss_pb2.LossConfig(binary_cross_entropy=loss_pb2.BinaryCrossEntropy()), + loss_pb2.LossConfig( + listwise_rank_loss=loss_pb2.ListwiseRankLoss( + session_name="id_a", learnable_temperature=False + ), + weight=0.5, + ), + ], + ) + model = _TestClassficationModel( + model_config=model_config, features=[], labels=["label"] + ) + model = TrainWrapper(model) + model = create_test_model(model, graph_type) + + sparse_feature = KeyedJaggedTensor.from_lengths_sync( + keys=["id_a"], + values=torch.tensor([1, 2, 1, 2, 1, 3]), + lengths=torch.tensor([1, 1, 1, 1, 1, 1]), + ) + logits = torch.tensor([0.2, 0.3, -0.1, 0.5, 0.1, 0.4]) + dense_feature = KeyedTensor.from_tensor_list( + keys=["int_a"], tensors=[logits.unsqueeze(1)] + ) + label = torch.tensor([0, 1, 1, 1, 0, 0]) + batch = Batch( + dense_features={BASE_DATA_GROUP: dense_feature}, + sparse_features={BASE_DATA_GROUP: sparse_feature}, + labels={"label": label}, + ) + total_loss, (losses, predictions, batch) = model(batch) + + session1 = torch.log_softmax(logits[[0, 2, 4]] / 0.07, dim=0)[1] + expected_listwise = -session1 / 3 * 0.5 + expected_bce = torch.nn.functional.binary_cross_entropy_with_logits( + logits, label.float() + ) + torch.testing.assert_close( + losses["listwise_rank_loss"], expected_listwise, rtol=1e-4, atol=1e-4 + ) + torch.testing.assert_close( + losses["binary_cross_entropy"], expected_bce, rtol=1e-4, atol=1e-4 + ) + torch.testing.assert_close( + total_loss, expected_bce + expected_listwise, rtol=1e-4, atol=1e-4 + ) + + def test_listwise_rank_loss_rejects_sample_weight(self): + model_config = model_pb2.ModelConfig( + losses=[ + loss_pb2.LossConfig( + listwise_rank_loss=loss_pb2.ListwiseRankLoss(session_name="id_a") + ) + ], + ) + model = _TestClassficationModel( + model_config=model_config, + features=[], + labels=["label"], + sample_weights=["weight"], + ) + with self.assertRaisesRegex(ValueError, "per-sample weights"): + TrainWrapper(model) + + def test_listwise_rank_loss_needs_list_source(self): + model_config = model_pb2.ModelConfig( + losses=[ + loss_pb2.LossConfig(listwise_rank_loss=loss_pb2.ListwiseRankLoss()) + ], + ) + model = TrainWrapper( + _TestClassficationModel( + model_config=model_config, features=[], labels=["label"] + ) + ) + dense_feature = KeyedTensor.from_tensor_list( + keys=["int_a"], tensors=[torch.tensor([[0.2], [0.3]])] + ) + batch = Batch( + dense_features={BASE_DATA_GROUP: dense_feature}, + sparse_features={}, + labels={"label": torch.tensor([0, 1])}, + ) + with self.assertRaisesRegex(ValueError, "session_name"): + model(batch) + + @parameterized.expand( + [[TestGraphType.NORMAL], [TestGraphType.FX_TRACE]], + name_func=parameterized_name_func, + ) + def test_jrc_loss_with_session(self, graph_type): + model_config = model_pb2.ModelConfig( + num_class=2, + losses=[ + loss_pb2.LossConfig( + jrc_loss=loss_pb2.JRCLoss(session_name="id_a", alpha=0.3) + ) + ], + ) + model = TrainWrapper( + _TestClassficationModel( + model_config=model_config, features=[], labels=["label"] + ) + ) + model = create_test_model(model, graph_type) + + ids = torch.tensor([1, 2, 1, 2, 1, 3]) + sparse_feature = KeyedJaggedTensor.from_lengths_sync( + keys=["id_a"], values=ids, lengths=torch.ones(6, dtype=torch.int64) + ) + torch.manual_seed(0) + logits = torch.randn(6, 2) + dense_feature = KeyedTensor.from_tensor_list(keys=["int_a"], tensors=[logits]) + label = torch.tensor([0, 1, 1, 1, 0, 0]) + batch = Batch( + dense_features={BASE_DATA_GROUP: dense_feature}, + sparse_features={BASE_DATA_GROUP: sparse_feature}, + labels={"label": label}, + ) + total_loss, (losses, predictions, batch) = model(batch) + + _, index, lengths = torch.unique(ids, return_inverse=True, return_counts=True) + expected = JRCLoss(alpha=0.3)(logits, label, lengths, index) + torch.testing.assert_close(losses["jrc_loss"], expected) + torch.testing.assert_close(total_loss, expected) + + def test_published_counts_define_the_lists(self): + """A model with per-request counts uses them; session_name is refused.""" + sparse_feature = KeyedJaggedTensor.from_lengths_sync( + keys=["id_a"], values=torch.tensor([7, 7]), lengths=torch.tensor([1, 1]) + ) + torch.manual_seed(0) + logits = torch.randn(6) + label = torch.tensor([0, 1, 1, 0, 0, 1]) + batch = Batch( + dense_features={ + BASE_DATA_GROUP: KeyedTensor.from_tensor_list( + keys=["int_a"], tensors=[logits.unsqueeze(1)] + ) + }, + sparse_features={BASE_DATA_GROUP: sparse_feature}, + labels={"label": label}, + ) + + def build(**listwise_kwargs): + model_config = model_pb2.ModelConfig( + losses=[ + loss_pb2.LossConfig( + listwise_rank_loss=loss_pb2.ListwiseRankLoss( + learnable_temperature=False, **listwise_kwargs + ) + ) + ], + ) + return TrainWrapper( + _TestJaggedModel( + model_config=model_config, features=[], labels=["label"] + ) + ) + + _, (losses, _, _) = build()(batch) + expected = ListwiseRankLoss(learnable_temperature=False)( + logits, label, torch.tensor([2, 4]) + ) + torch.testing.assert_close(losses["listwise_rank_loss"], expected) + + with self.assertRaisesRegex(ValueError, "only for models"): + build(session_name="id_a")(batch) + + def test_jrc_session_spans_requests_on_published_counts(self): + """Jrc's session_name repeats the request-level id over its candidates.""" + torch.manual_seed(0) + logits = torch.randn(6, 2) + label = torch.tensor([0, 1, 1, 0, 0, 1]) + model_config = model_pb2.ModelConfig( + num_class=2, + losses=[ + loss_pb2.LossConfig(jrc_loss=loss_pb2.JRCLoss(session_name="id_a")) + ], + ) + model = TrainWrapper( + _TestJaggedModel(model_config=model_config, features=[], labels=["label"]) + ) + for ids, lengths in (([7, 7], [6]), ([7, 8], [2, 4])): + batch = Batch( + dense_features={ + BASE_DATA_GROUP: KeyedTensor.from_tensor_list( + keys=["int_a"], tensors=[logits] + ) + }, + sparse_features={ + BASE_DATA_GROUP: KeyedJaggedTensor.from_lengths_sync( + keys=["id_a"], + values=torch.tensor(ids), + lengths=torch.tensor([1, 1]), + ) + }, + labels={"label": label}, + ) + _, (losses, _, _) = model(batch) + expected = JRCLoss()(logits, label, torch.tensor(lengths)) + torch.testing.assert_close(losses["jrc_loss"], expected, msg=str(ids)) + if __name__ == "__main__": unittest.main() diff --git a/tzrec/modules/gr/onerank_sd.py b/tzrec/modules/gr/onerank_sd.py index af5e5c03c..500dc304f 100644 --- a/tzrec/modules/gr/onerank_sd.py +++ b/tzrec/modules/gr/onerank_sd.py @@ -40,11 +40,7 @@ from tzrec.modules.norm import LayerNorm from tzrec.modules.utils import BaseModule -from tzrec.ops.jagged_tensors import ( - jagged_segment_ids, - jagged_segment_max, - jagged_segment_sum, -) +from tzrec.ops.scatter_ops import lengths_to_index, scatter_max, scatter_sum def _jagged_softmax(logits: torch.Tensor, lengths: torch.Tensor) -> torch.Tensor: @@ -54,10 +50,10 @@ def _jagged_softmax(logits: torch.Tensor, lengths: torch.Tensor) -> torch.Tensor far-below-max segment to an all-zero row and a 0/0 denominator) and detached, because it is a constant of the softmax identity. """ - segment_ids = jagged_segment_ids(lengths, output_size=logits.size(0)) - maxes = jagged_segment_max(logits.detach(), lengths, segment_ids) - exp = torch.exp(logits - maxes.index_select(0, segment_ids)) - denom = jagged_segment_sum(exp, lengths, segment_ids) + index = lengths_to_index(lengths, output_size=logits.size(0)) + maxes = scatter_max(logits.detach(), index, lengths.size(0)) + exp = torch.exp(logits - maxes.index_select(0, index)) + denom = scatter_sum(exp, index, lengths.size(0)) return exp / torch.repeat_interleave(denom, lengths, dim=0, output_size=exp.size(0)) @@ -130,15 +126,15 @@ def _jagged_single_query_attn( return _varlen_single_query_attn(q, k, v, lengths, attn_scale) num_heads = q.size(1) head_dim = q.size(2) - segment_ids = jagged_segment_ids(lengths, output_size=k.size(0)) + index = lengths_to_index(lengths, output_size=k.size(0)) # Broadcast the query onto its own rows instead of padding the pool # to a dense (B, N_max, D) block. - q_rows = q.index_select(0, segment_ids) + q_rows = q.index_select(0, index) logits = (q_rows * k).sum(dim=-1) * attn_scale attn = _jagged_softmax(logits, lengths) attn = F.dropout(attn, p=dropout_ratio, training=training) weighted = (attn.unsqueeze(-1) * v).reshape(-1, num_heads * head_dim) - return jagged_segment_sum(weighted, lengths, segment_ids) + return scatter_sum(weighted, index, lengths.size(0)) torch.fx.wrap(_jagged_single_query_attn) diff --git a/tzrec/modules/task_tower.py b/tzrec/modules/task_tower.py index dd7041f03..e9987978a 100644 --- a/tzrec/modules/task_tower.py +++ b/tzrec/modules/task_tower.py @@ -20,7 +20,7 @@ from tzrec.modules.gr.onerank_sd import OneRankSituationDiscernment from tzrec.modules.mlp import MLP from tzrec.modules.utils import BaseModule -from tzrec.ops.jagged_tensors import jagged_segment_ids, jagged_segment_sum +from tzrec.ops.scatter_ops import lengths_to_index, scatter_sum class TaskTower(nn.Module): @@ -109,8 +109,8 @@ def _jagged_mean(values: torch.Tensor, lengths: torch.Tensor) -> torch.Tensor: num_tasks = values.size(1) dim = values.size(2) flat = values.reshape(values.size(0), num_tasks * dim) - sums = jagged_segment_sum( - flat, lengths, jagged_segment_ids(lengths, output_size=flat.size(0)) + sums = scatter_sum( + flat, lengths_to_index(lengths, output_size=flat.size(0)), batch_size ) denom = lengths.clamp(min=1).to(flat.dtype).unsqueeze(-1) return (sums / denom).reshape(batch_size, num_tasks, dim) diff --git a/tzrec/ops/jagged_tensors.py b/tzrec/ops/jagged_tensors.py index 91269e703..b484a512b 100644 --- a/tzrec/ops/jagged_tensors.py +++ b/tzrec/ops/jagged_tensors.py @@ -170,70 +170,3 @@ def jagged_dense_bmm_broadcast_add( dense=dense, bias=bias, ) - - -def jagged_segment_ids( - lengths: torch.Tensor, output_size: Optional[int] = None -) -> torch.Tensor: - """Map each jagged row to its segment index. - - ``output_size`` -- the statically known ``sum(lengths)`` -- avoids the - hidden device->host sync (and data-dependent graph break) that - ``repeat_interleave`` with tensor repeats otherwise performs. - """ - return torch.repeat_interleave( - torch.arange(lengths.size(0), device=lengths.device), - lengths, - output_size=output_size, - ) - - -def jagged_segment_sum( - values: torch.Tensor, - lengths: torch.Tensor, - segment_ids: torch.Tensor, -) -> torch.Tensor: - """Sum a jagged ``(total, C)`` tensor within each segment; empties -> 0. - - Reduced-precision inputs (fp16/bf16) accumulate in fp32 and cast back: - ``index_add_`` is not on autocast's promote list, so bf16 inputs would - otherwise add through bf16 atomics whose reorder noise sits at bf16 - rounding scale. ``promote_types`` keeps the dtype choice a graph node - rather than Python control flow, so fx tracing still inlines this. - """ - acc_dtype = torch.promote_types(values.dtype, torch.float32) - sums = torch.zeros( - (lengths.size(0), values.size(-1)), - dtype=acc_dtype, - device=values.device, - ) - sums.index_add_(0, segment_ids, values.to(acc_dtype)) - return sums.to(values.dtype) - - -def jagged_segment_max( - values: torch.Tensor, - lengths: torch.Tensor, - segment_ids: torch.Tensor, -) -> torch.Tensor: - """Max-reduce a jagged ``(total, C)`` tensor within each segment. - - ``scatter_reduce_`` rather than the still-beta ``index_reduce_`` (which - warns on every call); it has no amax backward either, so callers using - the result as a shift constant must detach ``values``. Empty segments - keep ``-inf``, which would make downstream broadcasts NaN under - torch.compile, so they are rewritten to 0 -- nothing reads them back. - """ - maxes = torch.full( - (lengths.size(0), values.size(-1)), - float("-inf"), - dtype=values.dtype, - device=values.device, - ).scatter_reduce_( - 0, - segment_ids.unsqueeze(-1).expand_as(values), - values, - "amax", - include_self=False, - ) - return torch.nan_to_num(maxes, neginf=0.0) diff --git a/tzrec/ops/jagged_tensors_test.py b/tzrec/ops/jagged_tensors_test.py index 60a009057..2add6bb52 100644 --- a/tzrec/ops/jagged_tensors_test.py +++ b/tzrec/ops/jagged_tensors_test.py @@ -17,11 +17,6 @@ from hypothesis import strategies as st from tzrec.ops import Kernel -from tzrec.ops.jagged_tensors import ( - jagged_segment_ids, - jagged_segment_max, - jagged_segment_sum, -) from tzrec.utils.test_util import ( cleanup_cuda_memory, generate_sparse_seq_len, @@ -494,47 +489,5 @@ def _test_jagged_dense_bmm_broadcast_add( ) -class JaggedSegmentOpsTest(unittest.TestCase): - """Value-level tests for the shared jagged segment reductions. - - Deterministic CPU cases pinning the contracts the three consumers - (listwise loss, SD, task tower) rely on: empty segments sum to 0, - max empties are rewritten to 0 (not -inf), and reduced-precision - inputs round-trip through fp32 accumulation. - """ - - def test_segment_ids_maps_rows_and_skips_empties(self) -> None: - lengths = torch.tensor([2, 0, 3]) - ids = jagged_segment_ids(lengths) - self.assertEqual(ids.tolist(), [0, 0, 2, 2, 2]) - - def test_segment_ids_honors_output_size(self) -> None: - lengths = torch.tensor([2, 0, 3]) - ids = jagged_segment_ids(lengths, output_size=5) - self.assertEqual(ids.tolist(), [0, 0, 2, 2, 2]) - - def test_segment_sum_within_segments_empty_is_zero(self) -> None: - values = torch.tensor([[1.0], [2.0], [3.0], [4.0], [5.0]]) - lengths = torch.tensor([2, 0, 3]) - ids = jagged_segment_ids(lengths, output_size=5) - sums = jagged_segment_sum(values, lengths, ids) - self.assertEqual(sums.tolist(), [[3.0], [0.0], [12.0]]) - - def test_segment_sum_round_trips_reduced_precision(self) -> None: - values = torch.tensor([[1.0], [2.0], [3.0]], dtype=torch.bfloat16) - lengths = torch.tensor([3]) - ids = jagged_segment_ids(lengths, output_size=3) - sums = jagged_segment_sum(values, lengths, ids) - self.assertEqual(sums.dtype, torch.bfloat16) - self.assertEqual(sums.tolist(), [[6.0]]) - - def test_segment_max_rewrites_empty_segments_to_zero(self) -> None: - values = torch.tensor([[1.0], [5.0], [2.0], [7.0], [3.0]]) - lengths = torch.tensor([2, 0, 3]) - ids = jagged_segment_ids(lengths, output_size=5) - maxes = jagged_segment_max(values, lengths, ids) - self.assertEqual(maxes.tolist(), [[5.0], [0.0], [7.0]]) - - if __name__ == "__main__": unittest.main() diff --git a/tzrec/ops/scatter_ops.py b/tzrec/ops/scatter_ops.py new file mode 100644 index 000000000..778404899 --- /dev/null +++ b/tzrec/ops/scatter_ops.py @@ -0,0 +1,77 @@ +# Copyright (c) 2026, Alibaba Group; +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Optional + +import torch + + +def lengths_to_index( + lengths: torch.Tensor, output_size: Optional[int] = None +) -> torch.Tensor: + """Per-row list index (torch_scatter's ``index``) of a jagged ``lengths`` layout. + + ``output_size`` -- the statically known ``sum(lengths)`` -- avoids the + hidden device->host sync (and data-dependent graph break) that + ``repeat_interleave`` with tensor repeats otherwise performs. + """ + return torch.repeat_interleave( + torch.arange(lengths.size(0), device=lengths.device), + lengths, + output_size=output_size, + ) + + +def scatter_sum(src: torch.Tensor, index: torch.Tensor, dim_size: int) -> torch.Tensor: + """Sum the rows of ``src`` ``(N, C)`` into ``dim_size`` groups by ``index``. + + Groups without a row are 0. Reduced-precision inputs (fp16/bf16) + accumulate in fp32 and cast back: ``index_add_`` is not on autocast's + promote list, so bf16 inputs would otherwise add through bf16 atomics + whose reorder noise sits at bf16 rounding scale. ``promote_types`` + keeps the dtype choice a graph node rather than Python control flow, + so fx tracing still inlines this. + """ + acc_dtype = torch.promote_types(src.dtype, torch.float32) + out = torch.zeros((dim_size, src.size(-1)), dtype=acc_dtype, device=src.device) + out.index_add_(0, index, src.to(acc_dtype)) + return out.to(src.dtype) + + +def scatter_max(src: torch.Tensor, index: torch.Tensor, dim_size: int) -> torch.Tensor: + """Max-reduce the rows of ``src`` ``(N, C)`` into ``dim_size`` groups. + + ``scatter_reduce_`` rather than the still-beta ``index_reduce_`` (which + warns on every call); it has no amax backward either, so callers using + the result as a shift constant must detach ``src``. Groups without a + row keep ``-inf``, which would make downstream broadcasts NaN under + torch.compile, so they are rewritten to 0 -- nothing reads them back. + """ + out = torch.full( + (dim_size, src.size(-1)), float("-inf"), dtype=src.dtype, device=src.device + ).scatter_reduce_( + 0, index.unsqueeze(-1).expand_as(src), src, "amax", include_self=False + ) + return torch.nan_to_num(out, neginf=0.0) + + +def scatter_logsumexp( + src: torch.Tensor, index: torch.Tensor, dim_size: int +) -> torch.Tensor: + """Log-sum-exp of the rows of ``src`` ``(N, C)`` within each group. + + Shifts by the detached group max, which the gradient does not depend + on, so the result is exact while ``exp`` cannot overflow. A group + without a row yields ``-inf``; nothing that has rows reads it back. + """ + maxes = scatter_max(src.detach(), index, dim_size) + sums = scatter_sum(torch.exp(src - maxes.index_select(0, index)), index, dim_size) + return maxes + torch.log(sums) diff --git a/tzrec/ops/scatter_ops_test.py b/tzrec/ops/scatter_ops_test.py new file mode 100644 index 000000000..3e4565295 --- /dev/null +++ b/tzrec/ops/scatter_ops_test.py @@ -0,0 +1,89 @@ +# Copyright (c) 2026, Alibaba Group; +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest + +import torch + +from tzrec.ops.scatter_ops import ( + lengths_to_index, + scatter_logsumexp, + scatter_max, + scatter_sum, +) + + +class ScatterOpsTest(unittest.TestCase): + """Contracts the consumers (losses, SD, task tower) rely on. + + Groups without a row sum to 0, max rewrites them to 0 (not -inf), + reduced-precision inputs round-trip through fp32 accumulation, and the + index may arrive in any order. + """ + + def test_scatter_sum_empty_group_is_zero(self) -> None: + values = torch.tensor([[1.0], [2.0], [3.0], [4.0], [5.0]]) + index = torch.tensor([0, 0, 2, 2, 2]) + self.assertEqual(scatter_sum(values, index, 3).tolist(), [[3.0], [0.0], [12.0]]) + + def test_scatter_sum_round_trips_reduced_precision(self) -> None: + values = torch.tensor([[1.0], [2.0], [3.0]], dtype=torch.bfloat16) + sums = scatter_sum(values, torch.tensor([0, 0, 0]), 1) + self.assertEqual(sums.dtype, torch.bfloat16) + self.assertEqual(sums.tolist(), [[6.0]]) + + def test_scatter_max_rewrites_empty_group_to_zero(self) -> None: + values = torch.tensor([[1.0], [5.0], [2.0], [7.0], [3.0]]) + index = torch.tensor([2, 0, 2, 0, 2]) + self.assertEqual(scatter_max(values, index, 3).tolist(), [[7.0], [0.0], [3.0]]) + + def test_scatter_logsumexp_matches_torch_per_group(self) -> None: + torch.manual_seed(0) + values = torch.randn(7, 2, requires_grad=True) + index = torch.tensor([1, 0, 1, 1, 0, 3, 3]) + out = scatter_logsumexp(values, index, 4) + for group in (0, 1, 3): + torch.testing.assert_close( + out[group], torch.logsumexp(values[index == group], dim=0) + ) + self.assertTrue(torch.isinf(out[2]).all() and (out[2] < 0).all()) + (grad,) = torch.autograd.grad(out[[0, 1, 3]].sum(), values) + expected = torch.zeros_like(grad) + for group in (0, 1, 3): + m = index == group + expected[m] = torch.softmax(values.detach()[m], dim=0) + torch.testing.assert_close(grad, expected) + + +class LengthsToIndexTest(unittest.TestCase): + def test_index_maps_rows_and_skips_empties(self) -> None: + lengths = torch.tensor([2, 0, 3]) + ids = lengths_to_index(lengths) + self.assertEqual(ids.tolist(), [0, 0, 2, 2, 2]) + + def test_index_honors_output_size(self) -> None: + lengths = torch.tensor([2, 0, 3]) + ids = lengths_to_index(lengths, output_size=5) + self.assertEqual(ids.tolist(), [0, 0, 2, 2, 2]) + + def test_scatter_logsumexp_bf16_tracks_fp32(self) -> None: + torch.manual_seed(0) + values = torch.randn(9, 3) * 4 + index = torch.tensor([0, 2, 2, 0, 1, 2, 0, 1, 1]) + out = scatter_logsumexp(values.bfloat16(), index, 3) + self.assertEqual(out.dtype, torch.bfloat16) + torch.testing.assert_close( + out.float(), scatter_logsumexp(values, index, 3), atol=5e-2, rtol=2e-2 + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tzrec/protos/loss.proto b/tzrec/protos/loss.proto index 9c78145c5..02eadab95 100644 --- a/tzrec/protos/loss.proto +++ b/tzrec/protos/loss.proto @@ -50,13 +50,17 @@ message BinaryCrossEntropy { optional float label_smoothing = 1 [default = 0.0]; } -// Per-request list-wise ranking loss: softmax cross-entropy over the -// candidates of one request, where the negatives of a candidate are the -// other candidates of the same request. Requires the model to publish -// per-request candidate counts and the task's `logits_` -// prediction (carry it next to a logit loss such as -// binary_cross_entropy); currently the DlrmHSTU family satisfies both. +// List-wise ranking loss: softmax cross-entropy over the samples of one +// list, where the negatives of a sample are the other samples of the same +// list. The list is the per-request candidate sequence a model publishes +// (the DlrmHSTU family) or the samples of a batch that share the +// `session_name` feature value. message ListwiseRankLoss { + // Feature (e.g. request_id or user_id) whose value groups the samples of + // a batch into lists, for models that do not publish per-request + // candidate counts. The DlrmHSTU family uses each request as the list + // and rejects this field. + optional string session_name = 1; // initial softmax temperature; the logits are multiplied by // 1 / temperature. optional float temperature_init = 2 [default = 0.07]; @@ -71,8 +75,15 @@ message SoftmaxCrossEntropy { message L2Loss { } +// Joint ranking and calibration loss: cross-entropy plus a per-session +// term in which each positive competes with the negatives of its session +// and vice versa. The session is defined like ListwiseRankLoss's list. message JRCLoss { - required string session_name = 1; + // Feature (e.g. user_id) whose value groups the samples of a batch into + // sessions. Required unless the model publishes per-request candidate + // counts (the DlrmHSTU family), where it overrides the per-request + // session. + optional string session_name = 1; optional float alpha = 2 [default = 0.5]; }