nextrec 0.4.33__py3-none-any.whl → 0.5.0__py3-none-any.whl

nextrec/loss/grad_norm.py

@@ -2,12 +2,40 @@
 GradNorm loss weighting for multi-task learning.
 
 Date: create on 27/10/2025
-Checkpoint: edit on 24/12/2025
+Checkpoint: edit on 22/01/2026
 Author: Yang Zhou,zyaztec@gmail.com
 
 Reference:
 Chen, Zhao, et al. "GradNorm: Gradient Normalization for Adaptive Loss Balancing
 in Deep Multitask Networks." ICML 2018.
+
+pseudocode:
+---
+Initialize w_i = 1
+Record L_i(0)
+
+for each step:
+    1. Forward: compute each task loss L_i
+    2. Compute G_i = ||∇_W (w_i * L_i)||
+    3. Compute r_i = (L_i / L_i(0)) / mean(...)
+    4. Compute target: Ĝ_i = mean(G) * r_i^α
+    5. L_grad = sum |G_i - Ĝ_i|
+    6. Update w_i using ∇ L_grad
+    7. Backprop with sum_i (w_i * L_i) to update model
+
+伪代码：
+---
+初始化 w_i = 1
+记录 L_i(0)
+
+for each step:
+    1. 前向算各 task loss: L_i
+    2. 计算 G_i = ||∇_W (w_i * L_i)||
+    3. 计算 r_i = (L_i / L_i(0)) / mean(...)
+    4. 计算 target: Ĝ_i = mean(G) * r_i^α
+    5. L_grad = sum |G_i - Ĝ_i|
+    6. 对 w_i 用 ∇L_grad 更新
+    7. 用 ∑ w_i * L_i 反传更新模型
 """
 
 from __future__ import annotations
@@ -15,6 +43,7 @@ from __future__ import annotations
 from typing import Iterable
 
 import torch
+import torch.distributed as dist
 import torch.nn as nn
 import torch.nn.functional as F
 
@@ -23,7 +52,15 @@ def get_grad_norm_shared_params(
     model,
     shared_modules=None,
 ):
+    """
+    Get shared parameters for GradNorm.
+
+    Args:
+        model: A pytorch model instance containing grad_norm_shared_modules attribute.
+        shared_modules: Optional list of module names to consider as shared.
+    """
     if not shared_modules:
+        # If no specific shared modules are provided, consider all parameters as shared
         return [p for p in model.parameters() if p.requires_grad]
     shared_params = []
     seen = set()
@@ -35,26 +72,10 @@ def get_grad_norm_shared_params(
             if param.requires_grad and id(param) not in seen:
                 shared_params.append(param)
                 seen.add(id(param))
-    if not shared_params:
-        return [p for p in model.parameters() if p.requires_grad]
     return shared_params
 
 
 class GradNormLossWeighting:
-    """
-    Adaptive multi-task loss weighting with GradNorm.
-
-    Args:
-        nums_task: Number of tasks.
-        alpha: GradNorm balancing strength.
-        lr: Learning rate for the weight optimizer.
-        init_weights: Optional initial weights per task.
-        device: Torch device for weights.
-        ema_decay: Optional EMA decay for smoothing loss ratios.
-        init_ema_steps: Number of steps to build EMA for initial losses.
-        init_ema_decay: EMA decay for initial losses when init_ema_steps > 0.
-        eps: Small value for numerical stability.
-    """
 
     def __init__(
         self,
@@ -63,58 +84,43 @@ class GradNormLossWeighting:
         lr: float = 0.025,
         init_weights: Iterable[float] | None = None,
         device: torch.device | str | None = None,
-        ema_decay: float | None = None,
-        init_ema_steps: int = 0,
-        init_ema_decay: float = 0.9,
         eps: float = 1e-8,
     ) -> None:
+        """
+        Adaptive multi-task loss weighting with GradNorm.
+
+        Args:
+            nums_task: Number of tasks.
+            alpha: GradNorm balancing strength.
+            lr: Learning rate for the weight optimizer.
+            init_weights: Optional initial weights per task.
+            device: Torch device for weights.
+            eps: Small value for numerical stability.
+
+
+
+        """
+
         if nums_task <= 1:
             raise ValueError("GradNorm requires nums_task > 1.")
+
         self.nums_task = nums_task
         self.alpha = alpha
         self.eps = eps
-        if ema_decay is not None:
-            ema_decay = ema_decay
-            if ema_decay < 0.0 or ema_decay >= 1.0:
-                raise ValueError("ema_decay must be in [0.0, 1.0).")
-        self.ema_decay = ema_decay
-        self.init_ema_steps = init_ema_steps
-        if self.init_ema_steps < 0:
-            raise ValueError("init_ema_steps must be >= 0.")
-        self.init_ema_decay = init_ema_decay
-        if self.init_ema_decay < 0.0 or self.init_ema_decay >= 1.0:
-            raise ValueError("init_ema_decay must be in [0.0, 1.0).")
-        self.init_ema_count = 0
 
         if init_weights is None:
             weights = torch.ones(self.nums_task, dtype=torch.float32)
         else:
             weights = torch.tensor(list(init_weights), dtype=torch.float32)
-            if weights.numel() != self.nums_task:
-                raise ValueError(
-                    "init_weights length must match nums_task for GradNorm."
-                )
+
         if device is not None:
             weights = weights.to(device)
         self.weights = nn.Parameter(weights)
         self.optimizer = torch.optim.Adam([self.weights], lr=float(lr))
 
         self.initial_losses = None
-        self.initial_losses_ema = None
-        self.loss_ema = None
         self.pending_grad = None
 
-    def to(self, device):
-        device = torch.device(device)
-        self.weights.data = self.weights.data.to(device)
-        if self.initial_losses is not None:
-            self.initial_losses = self.initial_losses.to(device)
-        if self.initial_losses_ema is not None:
-            self.initial_losses_ema = self.initial_losses_ema.to(device)
-        if self.loss_ema is not None:
-            self.loss_ema = self.loss_ema.to(device)
-        return self
-
     def compute_weighted_loss(
         self,
         task_losses: list[torch.Tensor],
@@ -122,6 +128,8 @@ class GradNormLossWeighting:
     ) -> torch.Tensor:
         """
         Return weighted total loss and update task weights with GradNorm.
+
+        BaseModel will use this method to compute the weighted loss when self.grad_norm is enabled.
         """
         if len(task_losses) != self.nums_task:
             raise ValueError(
@@ -136,19 +144,7 @@ class GradNormLossWeighting:
                 [loss.item() for loss in task_losses], device=self.weights.device
             )
             if self.initial_losses is None:
-                if self.init_ema_steps > 0:
-                    if self.initial_losses_ema is None:
-                        self.initial_losses_ema = loss_values
-                    else:
-                        self.initial_losses_ema = (
-                            self.init_ema_decay * self.initial_losses_ema
-                            + (1.0 - self.init_ema_decay) * loss_values
-                        )
-                    self.init_ema_count += 1
-                    if self.init_ema_count >= self.init_ema_steps:
-                        self.initial_losses = self.initial_losses_ema.clone()
-                else:
-                    self.initial_losses = loss_values
+                self.initial_losses = loss_values.clone()
 
         weights_detached = self.weights.detach()
         weighted_losses = [
@@ -157,25 +153,14 @@ class GradNormLossWeighting:
         total_loss = torch.stack(weighted_losses).sum()
 
         grad_norms = self.compute_grad_norms(task_losses, shared_params)
+
+        # compute inverse training rate, inv rate = loss_ratio / mean(loss_ratio)
         with torch.no_grad():
-            if self.ema_decay is not None:
-                if self.loss_ema is None:
-                    self.loss_ema = loss_values
-                else:
-                    self.loss_ema = (
-                        self.ema_decay * self.loss_ema
-                        + (1.0 - self.ema_decay) * loss_values
-                    )
-                ratio_source = self.loss_ema
-            else:
-                ratio_source = loss_values
             if self.initial_losses is not None:
                 base_initial = self.initial_losses
-            elif self.initial_losses_ema is not None:
-                base_initial = self.initial_losses_ema
             else:
                 base_initial = loss_values
-            loss_ratios = ratio_source / (base_initial + self.eps)
+            loss_ratios = loss_values / (base_initial + self.eps)
             inv_rate = loss_ratios / (loss_ratios.mean() + self.eps)
             target = grad_norms.mean() * (inv_rate**self.alpha)
 
@@ -187,6 +172,7 @@ class GradNormLossWeighting:
 
     def compute_grad_norms(self, task_losses, shared_params):
         grad_norms = []
+        # compute gradient norms for each task, gradient norms = sqrt(sum(grad^2))
         for i, task_loss in enumerate(task_losses):
             grads = torch.autograd.grad(
                 self.weights[i] * task_loss,
@@ -230,3 +216,19 @@ class GradNormLossWeighting:
             self.weights.copy_(w)
 
         self.pending_grad = None
+
+    def sync(self) -> None:
+        """
+        Synchronize GradNorm buffers across DDP ranks.
+
+        - pending_grad: averaged so all ranks update weights consistently
+        - initial_losses: averaged so the baseline loss is consistent
+        """
+
+        world_size = dist.get_world_size()
+        if self.pending_grad is not None:
+            dist.all_reduce(self.pending_grad, op=dist.ReduceOp.SUM)
+            self.pending_grad /= world_size
+        if self.initial_losses is not None:
+            dist.all_reduce(self.initial_losses, op=dist.ReduceOp.SUM)
+            self.initial_losses /= world_size

nextrec/models/multi_task/ple.py

@@ -54,6 +54,7 @@ from nextrec.basic.model import BaseModel
 from nextrec.utils.model import get_mlp_output_dim
 from nextrec.utils.types import TaskTypeInput
 
+
 class CGCLayer(nn.Module):
     """
     CGC (Customized Gate Control) block used by PLE.

nextrec/models/multi_task/share_bottom.py

@@ -45,6 +45,7 @@ from nextrec.basic.heads import TaskHead
 from nextrec.basic.model import BaseModel
 from nextrec.utils.types import TaskTypeInput
 
+
 class ShareBottom(BaseModel):
     @property
     def model_name(self):

nextrec/models/ranking/afm.py

@@ -156,7 +156,7 @@ class AFM(BaseModel):
         # First-order dense part
         if self.linear_dense is not None:
             dense_inputs = [
-                x[f.name].float().view(batch_size, -1) for f in self.dense_features
+                x[f.name].float().reshape(batch_size, -1) for f in self.dense_features
             ]
             dense_stack = torch.cat(dense_inputs, dim=1) if dense_inputs else None
             if dense_stack is not None:
@@ -170,7 +170,7 @@ class AFM(BaseModel):
                 term = emb(x[feature.name].long())  # [B, 1]
             else:  # SequenceFeature
                 seq_input = x[feature.name].long()  # [B, 1]
-                if feature.max_len is not None and seq_input.size(1) > feature.max_len:
+                if feature.max_len is not None:
                     seq_input = seq_input[:, -feature.max_len :]
                 mask = self.input_mask(x, feature, seq_input).squeeze(1)  # [B, 1]
                 seq_weight = emb(seq_input).squeeze(-1)  # [B, L]
@@ -186,16 +186,11 @@ class AFM(BaseModel):
         for feature in self.fm_features:
             value = x.get(f"{feature.name}_value")
             if value is not None:
-                value = value.float()
-                if value.dim() == 1:
-                    value = value.unsqueeze(-1)
+                value = value.float().reshape(batch_size, -1)
             else:
                 if isinstance(feature, SequenceFeature):
                     seq_input = x[feature.name].long()
-                    if (
-                        feature.max_len is not None
-                        and seq_input.size(1) > feature.max_len
-                    ):
+                    if feature.max_len is not None:
                         seq_input = seq_input[:, -feature.max_len :]
                     value = self.input_mask(x, feature, seq_input).sum(dim=2)  # [B, 1]
                 else:

nextrec/models/ranking/dien.py

@@ -390,13 +390,13 @@ class DIEN(BaseModel):
                 dim=-1,
             )
             score_t = self.attention_layer.attention_net(concat_feat)  # [B, 1]
-            att_scores_list.append(score_t)
+            att_scores_list.append(score_t.unsqueeze(1))
 
         # [B, seq_len, 1]
         att_scores = torch.cat(att_scores_list, dim=1)
 
-        scores_flat = att_scores.squeeze(-1)  # [B, seq_len]
-        mask_flat = mask.squeeze(-1)  # [B, seq_len]
+        scores_flat = att_scores[..., 0]  # [B, seq_len]
+        mask_flat = mask[..., 0]  # [B, seq_len]
 
         scores_flat = scores_flat.masked_fill(mask_flat == 0, -1e9)
         att_weights = torch.softmax(scores_flat, dim=1)  # [B, seq_len]
@@ -437,8 +437,7 @@ class DIEN(BaseModel):
 
         for feat in self.dense_features:
             val = x[feat.name].float()
-            if val.dim() == 1:
-                val = val.unsqueeze(1)
+            val = val.view(val.size(0), -1)
             other_embeddings.append(val)
 
         concat_input = torch.cat(other_embeddings, dim=-1)  # [B, total_dim]
@@ -460,15 +459,15 @@ class DIEN(BaseModel):
         interest_states = interest_states[:, :-1, :]
         pos_seq = behavior_emb[:, 1:, :]
         neg_seq = neg_behavior_emb[:, 1:, :]
-        aux_mask = mask[:, 1:, :].squeeze(-1)
+        aux_mask = mask[:, 1:, 0]
 
         if aux_mask.sum() == 0:
             return torch.tensor(0.0, device=self.device)
 
         pos_input = torch.cat([interest_states, pos_seq], dim=-1)
         neg_input = torch.cat([interest_states, neg_seq], dim=-1)
-        pos_logits = self.auxiliary_net(pos_input).squeeze(-1)
-        neg_logits = self.auxiliary_net(neg_input).squeeze(-1)
+        pos_logits = self.auxiliary_net(pos_input)[..., 0]
+        neg_logits = self.auxiliary_net(neg_input)[..., 0]
 
         pos_loss = F.binary_cross_entropy_with_logits(
             pos_logits, torch.ones_like(pos_logits), reduction="none"

nextrec/models/ranking/ffm.py

@@ -190,7 +190,7 @@ class FFM(BaseModel):
             return emb(x[feature.name].long())
 
         seq_input = x[feature.name].long()
-        if feature.max_len is not None and seq_input.size(1) > feature.max_len:
+        if feature.max_len is not None:
             seq_input = seq_input[:, -feature.max_len :]
         seq_emb = emb(seq_input)  # [B, L, D]
         mask = self.input_mask(x, feature, seq_input)
@@ -224,7 +224,7 @@ class FFM(BaseModel):
                 term = emb(x[feature.name].long())  # [B, 1]
             else:
                 seq_input = x[feature.name].long()
-                if feature.max_len is not None and seq_input.size(1) > feature.max_len:
+                if feature.max_len is not None:
                     seq_input = seq_input[:, -feature.max_len :]
                 mask = self.input_mask(x, feature, seq_input).squeeze(1)  # [B, L]
                 seq_weight = emb(seq_input).squeeze(-1)  # [B, L]

nextrec/models/retrieval/sdm.py

@@ -223,8 +223,7 @@ class SDM(BaseMatchModel):
             for feat in self.user_dense_features:
                 if feat.name in user_input:
                     val = user_input[feat.name].float()
-                    if val.dim() == 1:
-                        val = val.unsqueeze(1)
+                    val = val.reshape(val.size(0), -1)
                     dense_features.append(val)
             if dense_features:
                 features_list.append(torch.cat(dense_features, dim=1))

nextrec/models/sequential/hstu.py

@@ -438,8 +438,6 @@ class HSTU(BaseModel):
         return self.causal_mask[:seq_len, :seq_len]
 
     def trim_sequence(self, seq: torch.Tensor) -> torch.Tensor:
-        if seq.size(1) <= self.max_seq_len:
-            return seq
         return seq[:, -self.max_seq_len :]
 
     def forward(self, x: dict[str, torch.Tensor]) -> torch.Tensor:

nextrec/models/tree_base/base.py

@@ -29,7 +29,7 @@ from nextrec.data.dataloader import RecDataLoader
 from nextrec.data.data_processing import get_column_data
 from nextrec.utils.console import display_metrics_table
 from nextrec.utils.data import FILE_FORMAT_CONFIG, check_streaming_support
-from nextrec.utils.feature import to_list
+from nextrec.utils.torch_utils import to_list
 from nextrec.utils.torch_utils import to_numpy
 
 

nextrec/utils/__init__.py

@@ -36,7 +36,7 @@ from .data import (
     resolve_file_paths,
 )
 from .embedding import get_auto_embedding_dim
-from .feature import to_list
+from .torch_utils import as_float, to_list
 from .model import (
     compute_pair_scores,
     get_mlp_output_dim,
@@ -90,6 +90,7 @@ __all__ = [
     "normalize_task_loss",
     # Feature utilities
     "to_list",
+    "as_float",
     # Config utilities
     "resolve_path",
     "safe_value",

nextrec/utils/config.py

@@ -21,7 +21,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Tuple
 import pandas as pd
 import torch
 
-from nextrec.utils.feature import to_list
+from nextrec.utils.torch_utils import to_list
 
 if TYPE_CHECKING:
     from nextrec.basic.features import DenseFeature, SequenceFeature, SparseFeature

nextrec/utils/console.py

@@ -36,7 +36,7 @@ from rich.progress import (
 from rich.table import Table
 from rich.text import Text
 
-from nextrec.utils.feature import as_float, to_list
+from nextrec.utils.torch_utils import as_float, to_list
 
 T = TypeVar("T")