decoder_loss.py

# Generated by ml.integrations.export.runtime_packager.write_remote_code_bundle.
# Exported for HuggingFace trust_remote_code loading.
# This file is intentionally self-contained.

from __future__ import annotations

from collections.abc import Callable
from functools import lru_cache
import importlib

import torch
from torch import nn

from .loss_stats import mean_loss_from_sum_and_count
from .loss_stats import shifted_loss_sum_and_count
from .decoder_types import DecoderConfig

def loss_stats(
    logits: torch.Tensor,
    labels: torch.Tensor,
    *,
    label_offset: int = 0,
) -> tuple[torch.Tensor, torch.Tensor]:
    return shifted_loss_sum_and_count(
        logits,
        labels,
        label_offset=label_offset,
        ignore_index=-100,
    )


def mean_cross_entropy_loss(
    logits: torch.Tensor,
    labels: torch.Tensor,
    *,
    label_offset: int = 0,
) -> torch.Tensor:
    loss_sum, count = loss_stats(logits, labels, label_offset=label_offset)
    return mean_loss_from_sum_and_count(
        loss_sum=loss_sum,
        count=count,
        reference=logits,
    )


@lru_cache(maxsize=1)
def _liger_fused_linear_ce_func() -> Callable[..., torch.Tensor]:
    module_name = "liger_kernel.transformers.functional"
    try:
        module = importlib.import_module(module_name)
    except ModuleNotFoundError as exc:
        if exc.name == "liger_kernel":
            raise RuntimeError(
                "Liger fused linear cross entropy requires liger_kernel with "
                "liger_fused_linear_cross_entropy"
            ) from exc
        raise
    fn = getattr(module, "liger_fused_linear_cross_entropy", None)
    if fn is None:
        raise RuntimeError(f"{module_name}.liger_fused_linear_cross_entropy is unavailable")
    return fn


@lru_cache(maxsize=1)
def _liger_graph_safe_ops():
    module_name = "liger_kernel.ops.fused_linear_cross_entropy"
    try:
        module = importlib.import_module(module_name)
    except ModuleNotFoundError as exc:
        if exc.name == "liger_kernel":
            raise RuntimeError(
                "Liger fused linear cross entropy requires liger_kernel"
            ) from exc
        raise
    required = (
        "MAX_FUSED_SIZE",
        "element_mul_kernel",
        "is_hip",
        "liger_cross_entropy_kernel",
        "triton",
    )
    missing = [name for name in required if not hasattr(module, name)]
    if missing:
        raise RuntimeError(
            f"{module_name} is missing required graph-safe operators: {missing}"
        )
    return module


def _liger_linear_ce_loss_chunk_stats(
    hidden_chunk: torch.Tensor,
    chunk_labels: torch.Tensor,
    *,
    norm: nn.Module,
    output: nn.Module,
) -> tuple[torch.Tensor, torch.Tensor]:
    weight = getattr(output, "weight", None)
    bias = getattr(output, "bias", None)
    if not torch.is_tensor(weight):
        raise TypeError("liger_linear_ce loss backend requires output.weight")

    flat_labels = chunk_labels.reshape(-1).to(dtype=torch.long)
    keep = flat_labels != -100
    count = keep.sum().to(dtype=torch.float32)

    hidden_norm = norm(hidden_chunk).reshape(-1, int(hidden_chunk.size(-1)))
    loss_sum = _call_liger_fused_linear_ce(
        hidden_norm,
        weight,
        flat_labels,
        bias=bias,
    )
    if not torch.is_tensor(loss_sum):
        raise RuntimeError("liger_fused_linear_cross_entropy must return a tensor loss")
    return loss_sum.float(), count


def _reference_linear_ce_loss_chunk_stats(
    hidden_chunk: torch.Tensor,
    chunk_labels: torch.Tensor,
    *,
    norm: nn.Module,
    output: nn.Module,
    tile: int = 1024,
) -> tuple[torch.Tensor, torch.Tensor]:
    """The same statistic as the fused kernel, in plain PyTorch.

    Liger's fused linear cross entropy is a Triton kernel, so on a device
    without one the model cannot score its own loss at all -- which is exactly
    what a host wants to do when the accelerator is busy training. Every other
    device-specific path in the model already has a reference implementation to
    fall back to; this is the one that did not.

    Logits are materialised a tile of rows at a time because the vocabulary is
    65,536 wide and a whole chunk at once is gigabytes for no reason.
    """
    weight = getattr(output, "weight", None)
    bias = getattr(output, "bias", None)
    if not torch.is_tensor(weight):
        raise TypeError("reference linear_ce loss backend requires output.weight")

    flat_labels = chunk_labels.reshape(-1).to(dtype=torch.long)
    keep = flat_labels != -100
    count = keep.sum().to(dtype=torch.float32)

    hidden_norm = norm(hidden_chunk).reshape(-1, int(hidden_chunk.size(-1)))
    weight_f32 = weight.float()
    bias_f32 = None if bias is None else bias.float()
    loss_sum = hidden_norm.new_zeros((), dtype=torch.float32)
    rows = int(hidden_norm.size(0))
    for start in range(0, rows, int(tile)):
        end = min(start + int(tile), rows)
        logits = torch.nn.functional.linear(
            hidden_norm[start:end].float(), weight_f32, bias_f32
        )
        loss_sum = loss_sum + torch.nn.functional.cross_entropy(
            logits,
            flat_labels[start:end],
            ignore_index=-100,
            reduction="sum",
        )
    return loss_sum.float(), count


class _GraphSafeLinearCE(torch.autograd.Function):
    _TILE = 128
    _CUDA_CHUNK_SIZE = 4096

    @staticmethod
    def forward(
        ctx,
        hidden: torch.Tensor,
        weight: torch.Tensor,
        labels: torch.Tensor,
        bias: torch.Tensor,
    ) -> torch.Tensor:
        if hidden.is_cuda:
            ops = _liger_graph_safe_ops()
            token_count = int(hidden.size(0))
            vocab_size = int(weight.shape[0])
            block_size = min(
                int(ops.MAX_FUSED_SIZE),
                int(ops.triton.next_power_of_2(vocab_size)),
            )
            # The release graph is bounded by the 4096-token context. A single
            # full-width logits GEMM maximizes M without introducing a second
            # sequence-length capability or an unbounded workspace.
            chunk_size = min(int(token_count), _GraphSafeLinearCE._CUDA_CHUNK_SIZE)
            chunk_count = int(ops.triton.cdiv(int(token_count), chunk_size))
            has_bias = bool(bias.numel())
            grad_hidden = torch.empty_like(hidden)
            grad_weight = torch.empty_like(weight)
            grad_bias = torch.empty_like(bias) if has_bias else None
            loss_per_token = torch.zeros(
                int(token_count),
                dtype=torch.float32,
                device=hidden.device,
            )

            for chunk_index in range(chunk_count):
                start = chunk_index * chunk_size
                end = min(start + chunk_size, int(token_count))
                hidden_chunk = hidden[start:end]
                logits = hidden_chunk @ weight.t()
                if has_bias:
                    logits = logits + bias
                logits = logits.contiguous()
                label_chunk = labels[start:end].contiguous()
                loss_chunk = loss_per_token[start:end]
                row_count = int(logits.shape[0])
                ops.liger_cross_entropy_kernel[(row_count,)](
                    X_ptr=logits,
                    X_stride=logits.stride(-2),
                    Y_ptr=label_chunk,
                    Y_stride=label_chunk.stride(-1),
                    weight_ptr=None,
                    loss_ptr=loss_chunk,
                    z_loss_ptr=None,
                    loss_stride=loss_chunk.stride(-1),
                    token_accuracy_ptr=None,
                    token_accuracy_stride=0,
                    predicted_tokens_ptr=None,
                    predicted_tokens_stride=0,
                    n_cols=vocab_size,
                    n_non_ignore=int(token_count),
                    sum_non_ignore_weight=int(token_count),
                    weight_sum=0.0,
                    ignore_index=-100,
                    lse_square_scale=0.0,
                    label_smoothing=0.0,
                    reduction="sum",
                    softcap=None,
                    RETURN_Z_LOSS=False,
                    RETURN_TOKEN_ACCURACY=False,
                    RETURN_PREDICTED_TOKENS=False,
                    HAS_WEIGHT=False,
                    HAS_SOFTCAPPING=False,
                    HAS_GRADIENTS=True,
                    BLOCK_SIZE=block_size,
                    num_warps=32 if not ops.is_hip() else 16,
                )
                torch.mm(
                    logits,
                    weight,
                    out=grad_hidden[start:end],
                )
                if chunk_index == 0:
                    torch.mm(logits.transpose(0, 1), hidden_chunk, out=grad_weight)
                else:
                    torch.addmm(
                        grad_weight,
                        logits.transpose(0, 1),
                        hidden_chunk,
                        out=grad_weight,
                    )
                if grad_bias is not None:
                    bias_grad = logits.sum(dim=0)
                    if chunk_index == 0:
                        grad_bias.copy_(bias_grad)
                    else:
                        grad_bias.add_(bias_grad)

            saved_bias_grad = (
                grad_bias if grad_bias is not None else hidden.new_empty((0,))
            )
            ctx.save_for_backward(grad_hidden, grad_weight, saved_bias_grad)
            ctx.has_bias = has_bias
            ctx.liger_cuda = True
            return loss_per_token.sum()

        ctx.save_for_backward(hidden, weight, labels, bias)
        ctx.has_bias = bool(bias.numel())
        ctx.liger_cuda = False
        total = hidden.new_zeros((), dtype=torch.float32)
        with torch.no_grad():
            for start in range(0, int(hidden.size(0)), _GraphSafeLinearCE._TILE):
                logits = torch.nn.functional.linear(
                    hidden[start : start + _GraphSafeLinearCE._TILE],
                    weight,
                    bias if ctx.has_bias else None,
                )
                total = total + torch.nn.functional.cross_entropy(
                    logits,
                    labels[start : start + _GraphSafeLinearCE._TILE],
                    ignore_index=-100,
                    reduction="sum",
                ).float()
        return total

    @staticmethod
    def backward(ctx, grad_output: torch.Tensor):
        if ctx.liger_cuda:
            saved_hidden_grad, saved_weight_grad, saved_bias_grad = ctx.saved_tensors
            grad_hidden = torch.empty_like(saved_hidden_grad)
            grad_weight = torch.empty_like(saved_weight_grad)
            grad_hidden.copy_(saved_hidden_grad)
            grad_weight.copy_(saved_weight_grad)
            ops = _liger_graph_safe_ops()
            block_size = min(
                int(ops.MAX_FUSED_SIZE),
                int(ops.triton.next_power_of_2(int(grad_hidden.shape[-1]))),
            )
            num_warps = 32 if not ops.is_hip() else 16
            ops.element_mul_kernel[(int(grad_hidden.shape[0]),)](
                grad_hidden,
                grad_hidden.stride(-2),
                grad_output,
                int(grad_hidden.shape[-1]),
                BLOCK_SIZE=block_size,
                num_warps=num_warps,
            )
            ops.element_mul_kernel[(int(grad_weight.shape[0]),)](
                grad_weight,
                grad_weight.stride(-2),
                grad_output,
                int(grad_weight.shape[-1]),
                BLOCK_SIZE=block_size,
                num_warps=num_warps,
            )
            grad_bias = None
            if ctx.has_bias:
                grad_bias = torch.empty_like(saved_bias_grad)
                grad_bias.copy_(saved_bias_grad)
                ops.element_mul_kernel[(int(grad_bias.shape[0]),)](
                    grad_bias,
                    grad_bias.stride(-1),
                    grad_output,
                    1,
                    BLOCK_SIZE=block_size,
                    num_warps=num_warps,
                )
            return grad_hidden, grad_weight, None, grad_bias

        hidden, weight, labels, bias = ctx.saved_tensors
        grad_hidden = torch.zeros_like(hidden)
        grad_weight = torch.zeros_like(weight)
        grad_bias = torch.zeros_like(bias) if ctx.has_bias else None
        for start in range(0, int(hidden.size(0)), _GraphSafeLinearCE._TILE):
            hidden_tile = hidden[start : start + _GraphSafeLinearCE._TILE]
            labels_tile = labels[start : start + _GraphSafeLinearCE._TILE]
            logits = torch.nn.functional.linear(
                hidden_tile,
                weight,
                bias if ctx.has_bias else None,
            )
            valid = labels_tile != -100
            probs = torch.softmax(logits, dim=-1)
            safe_labels = labels_tile.clamp_min(0).unsqueeze(1)
            probs.scatter_add_(
                1,
                safe_labels,
                -valid.to(dtype=probs.dtype).unsqueeze(1),
            )
            probs.mul_(valid.unsqueeze(1))
            torch.mm(
                probs,
                weight,
                out=grad_hidden[start : start + _GraphSafeLinearCE._TILE],
            )
            torch.addmm(
                grad_weight,
                probs.transpose(0, 1),
                hidden_tile,
                out=grad_weight,
            )
            if grad_bias is not None:
                grad_bias.add_(probs.sum(dim=0))
        scale = grad_output.to(dtype=grad_hidden.dtype)
        return grad_hidden * scale, grad_weight * scale, None, (
            None if grad_bias is None else grad_bias * scale
        )


@torch.compiler.disable
def _call_liger_fused_linear_ce(
    hidden_norm: torch.Tensor,
    weight: torch.Tensor,
    flat_labels: torch.Tensor,
    *,
    bias: torch.Tensor | None,
) -> torch.Tensor:
    if hidden_norm.is_cuda and torch.cuda.is_current_stream_capturing():
        bias_arg = (
            bias if bias is not None else hidden_norm.new_empty((0,))
        )
        return _GraphSafeLinearCE.apply(hidden_norm, weight, flat_labels, bias_arg)
    fn = _liger_fused_linear_ce_func()
    try:
        return fn(
            hidden_norm,
            weight,
            flat_labels,
            bias=bias,
            ignore_index=-100,
            reduction="sum",
        )
    except TypeError as exc:
        raise RuntimeError(
            "installed liger_fused_linear_cross_entropy has an unsupported signature"
        ) from exc


def _loss_chunk_stats(
    hidden_chunk: torch.Tensor,
    chunk_labels: torch.Tensor,
    *,
    norm: nn.Module,
    output: nn.Module,
) -> tuple[torch.Tensor, torch.Tensor]:
    if bool(
        getattr(output, "_sophia_release_cuda_cache_before_linear_ce", False)
    ):
        output._sophia_release_cuda_cache_before_linear_ce = False
        torch.cuda.empty_cache()
    if not hidden_chunk.is_cuda:
        return _reference_linear_ce_loss_chunk_stats(
            hidden_chunk,
            chunk_labels,
            norm=norm,
            output=output,
        )
    return _liger_linear_ce_loss_chunk_stats(
        hidden_chunk,
        chunk_labels,
        norm=norm,
        output=output,
    )


def chunked_loss_stats_from_hidden(
    hidden: torch.Tensor,
    labels: torch.Tensor,
    *,
    label_offset: int,
    norm: nn.Module,
    output: nn.Module,
    config: DecoderConfig,
) -> tuple[torch.Tensor, torch.Tensor]:
    max_pred_tokens = min(
        max(int(hidden.size(1)) - 1, 0),
        max(int(labels.size(1)) - int(label_offset) - 1, 0),
    )
    zero = hidden.new_zeros(())
    if max_pred_tokens <= 0:
        return zero, zero

    loss_sum = zero
    count = zero
    chunk_size = max(int(config.loss_chunk_size), 0)
    if chunk_size <= 0:
        chunk_size = int(max_pred_tokens)
    target_base = int(label_offset) + 1
    for token_start in range(0, max_pred_tokens, chunk_size):
        token_end = min(token_start + chunk_size, max_pred_tokens)
        chunk_labels = labels[
            :,
            target_base + token_start : target_base + token_end,
        ].contiguous()
        chunk_loss_sum, chunk_count = _loss_chunk_stats(
            hidden[:, token_start:token_end],
            chunk_labels,
            norm=norm,
            output=output,
        )
        loss_sum = loss_sum + chunk_loss_sum
        count = count + chunk_count
    return loss_sum, count


__all__ = [
    "chunked_loss_stats_from_hidden",
    "loss_stats",
    "mean_cross_entropy_loss",
]