Move some functional APIs to fairseq2.ops (#1165)

cbalioglu · web-flow · commit 5136e5ad03f7 · 2025-05-05T12:16:18.000-04:00
diff --git a/src/fairseq2/generation/_sampling/_generator.py b/src/fairseq2/generation/_sampling/_generator.py
@@ -23,8 +23,8 @@
 from fairseq2.models.encoder_decoder import EncoderDecoderModel
 from fairseq2.models.sequence import SequenceModelOutput
 from fairseq2.nn import BatchLayout, IncrementalStateBag
+from fairseq2.ops import repeat_interleave
 from fairseq2.utils.stopwatch import Stopwatch
-from fairseq2.utils.tensor import repeat_interleave
 
 # isort: split
 
diff --git a/src/fairseq2/models/sequence.py b/src/fairseq2/models/sequence.py
@@ -8,13 +8,13 @@
 
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
-from typing import Literal, final
+from typing import Literal, Protocol, final
 
 from torch import Tensor
 from torch.nn import Module
 
 from fairseq2.nn import BatchLayout
-from fairseq2.nn.ops import CrossEntropy, cross_entropy
+from fairseq2.nn.functional import cross_entropy
 
 
 class SequenceModel(Module, ABC):
@@ -36,6 +36,18 @@ def forward(
     ) -> SequenceModelOutput: ...
 
 
+class CrossEntropy(Protocol):
+    def __call__(
+        self,
+        logits: Tensor,
+        targets: Tensor,
+        pad_idx: int | None,
+        *,
+        label_smoothing: float = 0.0,
+        reduction: Literal["sum", "mean", "none"] = "sum",
+    ) -> Tensor: ...
+
+
 @final
 @dataclass
 class SequenceModelOutput:
diff --git a/src/fairseq2/models/transformer/_multihead_attention.py b/src/fairseq2/models/transformer/_multihead_attention.py
@@ -29,7 +29,7 @@
     PositionEncoder,
     Projection,
 )
-from fairseq2.utils.tensor import repeat_interleave
+from fairseq2.ops import repeat_interleave
 
 # isort: split
 
diff --git a/src/fairseq2/models/transformer_lm/_model.py b/src/fairseq2/models/transformer_lm/_model.py
@@ -13,10 +13,10 @@
 from typing_extensions import override
 
 from fairseq2.models.decoder import DecoderModel
-from fairseq2.models.sequence import SequenceModelOutput
+from fairseq2.models.sequence import CrossEntropy, SequenceModelOutput
 from fairseq2.models.transformer import TransformerFrontend
 from fairseq2.nn import BatchLayout, IncrementalStateBag, Projection
-from fairseq2.nn.ops import CrossEntropy, cross_entropy
+from fairseq2.nn.functional import cross_entropy
 
 # isort: split
 
diff --git a/src/fairseq2/models/wav2vec2/_model.py b/src/fairseq2/models/wav2vec2/_model.py
@@ -19,7 +19,7 @@
 from fairseq2.error import InternalError
 from fairseq2.models.transformer import TransformerEncoder
 from fairseq2.nn import BatchLayout, Linear
-from fairseq2.utils.tensor import repeat_interleave
+from fairseq2.ops import repeat_interleave
 
 # isort: split
 
diff --git a/src/fairseq2/nn/_embedding.py b/src/fairseq2/nn/_embedding.py
@@ -23,7 +23,7 @@
 from fairseq2.error import InternalError
 from fairseq2.gang import Gang
 from fairseq2.nn.utils.module import to_empty
-from fairseq2.tensor_parallel import gather, reduce, reduce_on_backward
+from fairseq2.ops.tensor_parallel import gather, reduce, reduce_on_backward
 
 
 class Embedding(Module, ABC):
diff --git a/src/fairseq2/nn/_position_encoder.py b/src/fairseq2/nn/_position_encoder.py
@@ -22,7 +22,7 @@
 from fairseq2.data_type import DataType
 from fairseq2.device import Device
 from fairseq2.error import InternalError
-from fairseq2.utils.tensor import unsqueeze
+from fairseq2.ops import unsqueeze
 
 # isort: split
 
diff --git a/src/fairseq2/nn/_projection.py b/src/fairseq2/nn/_projection.py
@@ -24,7 +24,7 @@
 from fairseq2.error import InternalError
 from fairseq2.gang import Gang
 from fairseq2.nn.utils.module import to_empty
-from fairseq2.tensor_parallel import gather, reduce, reduce_on_backward, scatter
+from fairseq2.ops.tensor_parallel import gather, reduce, reduce_on_backward, scatter
 
 
 class Projection(Module, ABC):
diff --git a/src/fairseq2/nn/functional.py b/src/fairseq2/nn/functional.py
@@ -6,26 +6,14 @@
 
 from __future__ import annotations
 
-from typing import Literal, Protocol
+from typing import Literal
 
 import torch
 from torch import Tensor
 from torch.nn.functional import log_softmax
 from torch.nn.functional import nll_loss as torch_nll_loss
 
 
-class CrossEntropy(Protocol):
-    def __call__(
-        self,
-        logits: Tensor,
-        targets: Tensor,
-        pad_idx: int | None,
-        *,
-        label_smoothing: float = 0.0,
-        reduction: Literal["sum", "mean", "none"] = "sum",
-    ) -> Tensor: ...
-
-
 def cross_entropy(
     logits: Tensor,
     targets: Tensor,
diff --git a/src/fairseq2/nn/functional/__init__.py b/src/fairseq2/nn/functional/__init__.py
@@ -0,0 +1,10 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from __future__ import annotations
+
+from fairseq2.nn.functional._cross_entropy import cross_entropy as cross_entropy
+from fairseq2.nn.functional._nll_loss import nll_loss as nll_loss
diff --git a/src/fairseq2/nn/functional/_cross_entropy.py b/src/fairseq2/nn/functional/_cross_entropy.py
@@ -0,0 +1,63 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from __future__ import annotations
+
+from typing import Literal
+
+import torch
+from torch import Tensor
+from torch.nn.functional import log_softmax
+from torch.nn.functional import nll_loss as torch_nll_loss
+
+from fairseq2.nn.functional._nll_loss import nll_loss
+
+
+def cross_entropy(
+    logits: Tensor,
+    targets: Tensor,
+    pad_idx: int | None,
+    *,
+    label_smoothing: float = 0.0,
+    reduction: Literal["sum", "mean", "none"] = "sum",
+) -> Tensor:
+    """
+    Computes the cross entropy loss between ``logits`` and ``targets``.
+
+    This function differs from :func:`torch.nn.functional.cross_entropy` in
+    two ways:
+        - It uses a fused kernel based on ``log_sofmax()`` and ``nll_loss()``
+          which typically consumes less memory while having the same performance.
+        - Its label smoothing implementation is slightly different and has
+          parity with the original fairseq implementation.
+
+    :param logits: The logits. *Shape:* :math:`(S,T)`, where :math:`S` is the
+        sequence length and :math:`T` is the size of the vocabulary.
+    :param targets: The target indices. *Shape:* :math:`(S)`, where :math:`S` is
+        the sequence length.
+    :param pad_idx: The index of the PAD symbol in the target vocabulary.
+    :param label_smoothing: The amount of label smoothing to apply while
+        computing the loss.
+    :param reduction: The reduction to apply to the output.
+    """
+    # For numerical stability run in single precision.
+    # (S, T) -> (S, T)
+    lprobs = log_softmax(logits, dim=-1, dtype=torch.float32)
+
+    if label_smoothing == 0.0:
+        if pad_idx is None:
+            pad_idx = -100
+
+        # Unless we need label smoothing, use PyTorch `nll_loss()` as it can be
+        # fused with `log_softmax()` when compiled.
+        return torch_nll_loss(
+            lprobs, targets, ignore_index=pad_idx, reduction=reduction
+        )
+
+    # sum/mean: (), none: (S)
+    return nll_loss(
+        lprobs, targets, pad_idx, label_smoothing=label_smoothing, reduction=reduction
+    )
diff --git a/src/fairseq2/nn/functional/_nll_loss.py b/src/fairseq2/nn/functional/_nll_loss.py
@@ -0,0 +1,76 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from __future__ import annotations
+
+from typing import Literal
+
+from torch import Tensor
+
+
+def nll_loss(
+    lprobs: Tensor,
+    targets: Tensor,
+    pad_idx: int | None,
+    *,
+    label_smoothing: float = 0.0,
+    reduction: Literal["sum", "mean", "none"] = "sum",
+) -> Tensor:
+    """
+    Computes the negative log-likelihood loss.
+
+    In addition to :func:`torch.nn.functional.nll_loss`, this function accepts
+    a loss smoothing parameter that is compatible with the original fairseq.
+
+    :param lprobs: The log probabilities. *Shape:* :math:`(S,T)`, where :math:`S`
+        is the sequence length and :math:`T` is the size of the vocabulary.
+    :param targets: The target indices. *Shape:* :math:`(S)`, where :math:`S` is
+        the sequence length.
+    :param pad_idx: The index of the PAD symbol in the target vocabulary.
+    :param label_smoothing: The amount of label smoothing to apply while
+        computing the loss.
+    :param reduction: The reduction to apply to the output.
+    """
+    # (S) -> (S, 1)
+    targets = targets.unsqueeze(-1)
+
+    loss = -lprobs.gather(dim=-1, index=targets)
+
+    if label_smoothing > 0.0:
+        smooth_loss = -lprobs.sum(dim=-1, keepdim=True)
+    else:
+        smooth_loss = None
+
+    if pad_idx is not None:
+        padding_mask = targets.eq(pad_idx)
+
+        loss.masked_fill_(padding_mask, 0.0)
+
+        if smooth_loss is not None:
+            smooth_loss.masked_fill_(padding_mask, 0.0)
+
+    if reduction == "sum":
+        loss = loss.sum()
+
+        if smooth_loss is not None:
+            smooth_loss = smooth_loss.sum()
+    elif reduction == "mean":
+        loss = loss.mean()
+
+        if smooth_loss is not None:
+            smooth_loss = smooth_loss.mean()
+
+    if smooth_loss is not None:
+        # Our label smoothing implementation varies slightly from PyTorch's
+        # implementation in `cross_entropy`.
+        eps = label_smoothing / (lprobs.size(-1) - 1)
+
+        loss = (1.0 - label_smoothing - eps) * loss + eps * smooth_loss
+
+    if reduction == "none":
+        loss = loss.squeeze(-1)
+
+    return loss
diff --git a/src/fairseq2/nn/utils/mask.py b/src/fairseq2/nn/utils/mask.py
@@ -12,7 +12,7 @@
 from torch import Tensor
 
 from fairseq2.device import Device
-from fairseq2.utils.tensor import repeat_interleave
+from fairseq2.ops import repeat_interleave, unsqueeze
 
 
 def apply_mask(
@@ -28,8 +28,7 @@ def apply_mask(
 
     :returns: The input sequences with mask applied. *Shape:* Same as ``seqs``.
     """
-    for _ in range(seqs.ndim - mask.ndim):
-        mask = mask.unsqueeze(-1)
+    unsqueeze(mask, dim=-1, count=seqs.ndim - mask.ndim)
 
     return seqs.where(mask, fill_value)
 
diff --git a/src/fairseq2/ops/__init__.py b/src/fairseq2/ops/__init__.py
@@ -0,0 +1,10 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from __future__ import annotations
+
+from fairseq2.ops._shape import repeat_interleave as repeat_interleave
+from fairseq2.ops._shape import unsqueeze as unsqueeze
diff --git a/src/fairseq2/ops/_shape.py b/src/fairseq2/ops/_shape.py
@@ -0,0 +1,44 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from __future__ import annotations
+
+from torch import Tensor
+
+
+def unsqueeze(x: Tensor, dim: int, count: int = 1) -> Tensor:
+    for _ in range(count):
+        x = x.unsqueeze(dim=dim)
+
+    return x
+
+
+def repeat_interleave(x: Tensor, dim: int, repeat: int) -> Tensor:
+    """
+    Repeats elements of a tensor.
+
+    :param x: The input tensor.
+    :param dim: The dimension along which to repeat values.
+    :param repeat: The number of repetitions.
+
+    :returns: The repeated tensor which has the same shape as input, except
+        along the given axis.
+
+    .. note::
+        This is a lightweight version of :func:`torch.repeat_interleave` that
+        is faster for repetitions along a single dimension.
+    """
+    if repeat == 1:
+        return x
+
+    shape = [-1] * (x.ndim + 1)
+
+    if dim < 0:
+        dim += x.ndim
+
+    shape[dim + 1] = repeat
+
+    return x.unsqueeze(dim + 1).expand(shape).flatten(dim, dim + 1)
diff --git a/src/fairseq2/ops/tensor_parallel.py b/src/fairseq2/ops/tensor_parallel.py
diff --git a/src/fairseq2/utils/tensor.py b/src/fairseq2/utils/tensor.py
@@ -27,38 +27,3 @@ def to_tensor(
     t = torch.tensor(data, dtype=dtype, device=CPU, pin_memory=True)
 
     return t.to(device, non_blocking=True)
-
-
-def unsqueeze(x: Tensor, dim: int, count: int = 1) -> Tensor:
-    for _ in range(count):
-        x = x.unsqueeze(dim=dim)
-
-    return x
-
-
-def repeat_interleave(x: Tensor, dim: int, repeat: int) -> Tensor:
-    """
-    Repeats elements of a tensor.
-
-    :param x: The input tensor.
-    :param dim: The dimension along which to repeat values.
-    :param repeat: The number of repetitions.
-
-    :returns: The repeated tensor which has the same shape as input, except
-        along the given axis.
-
-    .. note::
-        This is a lightweight version of :func:`torch.repeat_interleave` that
-        is faster for repetitions along a single dimension.
-    """
-    if repeat == 1:
-        return x
-
-    shape = [-1] * (x.ndim + 1)
-
-    if dim < 0:
-        dim += x.ndim
-
-    shape[dim + 1] = repeat
-
-    return x.unsqueeze(dim + 1).expand(shape).flatten(dim, dim + 1)
diff --git a/tests/unit/nn/test_functional.py b/tests/unit/nn/test_functional.py

Original file line number	Diff line number	Diff line change
`@@ -29,7 +29,7 @@`
`29`	`29`	`PositionEncoder,`
`30`	`30`	`Projection,`
`31`	`31`	`)`
`32`		`-from fairseq2.utils.tensor import repeat_interleave`
	`32`	`+from fairseq2.ops import repeat_interleave`
`33`	`33`
`34`	`34`	`# isort: split`
`35`	`35`