$FIX$ - Incompatible torch.lgamma and torch.distributions.Dirichlet for certain HPC clusters

[franeknowak]

Environment

• PyTorch: 2.5.1+cu118
• CUDA: 11.8
• Python: 3.11.0
• GPU: NVIDIA RTX A6000
• Platform: HPC cluster (SGE scheduler), Linux

---

Description

On HPC nodes, any call to torch.lgamma on a CUDA tensor, or any use of torch.distributions.Dirichlet in a process that also initialises a CUDA context, causes an Illegal instruction (SIGILL) crash. The process terminates silently with no Python traceback — the SIGILL is raised at the OS level before Python can catch it.

The same code runs without issue on a local workstation with a modern CPU. The HPC compute nodes have older CPUs that do not support AVX-512 (or the specific extended instruction set that PyTorch 2.5.1's prebuilt CUDA kernels and distribution code were compiled against).

Why it is silent

SIGILL is a CPU signal, not a Python exception. It cannot be caught by try/except and produces no traceback, making it appear as though training simply never started.

---

Root causes identified (two separate issues)

1. torch.lgamma CUDA kernel

torch.lgamma dispatches to a CUDA kernel whose surrounding CPU-side dispatch/registration code uses AVX-512 instructions. On HPC nodes without AVX-512 support, calling torch.lgamma on any CUDA tensor raises SIGILL.

2. torch.distributions.Dirichlet at module import time

Importing from torch.distributions import Dirichlet at the top of a module causes PyTorch to eagerly register CUDA kernels for digamma/lgamma when the CUDA context is first initialised (first .to('cuda') call in that process). This registration path also uses AVX-512 instructions and raises SIGILL — even before any Dirichlet object is constructed or used.

Importing Dirichlet in isolation (no CUDA context in the process) works fine. Initialising CUDA in isolation (without Dirichlet imported) also works fine. The crash only occurs when both happen in the same process, regardless of order.

---

Diagnosis steps

The following confirmed the root cause:

# Works fine — no lgamma, no Dirichlet
python3 -c "import torch; torch.zeros(1).cuda(); print('OK')"

# Works fine — import only, no CUDA context
python3 -c "from torch.distributions import Dirichlet; print('OK')"

# SIGILL — lgamma on GPU tensor
python3 -c "
import torch
a = torch.rand(4, 2).cuda()
print(torch.lgamma(a))  # crashes here
"

# SIGILL — Dirichlet imported at module level, CUDA context created later
python3 -c "
from torch.distributions import Dirichlet
import torch
torch.zeros(1).cuda()  # crashes here
"

---

Solution

Both issues are resolved by keeping all torch.lgamma calls and all Dirichlet construction on CPU, moving tensors explicitly before the operation and back afterwards. Gradients flow correctly through PyTorch's device transfer autograd ops. Performance impact is negligible — the tensors involved are small (B, 2) and the backbone forward/backward pass dominates training time.

The lazy import of Dirichlet (inside the function rather than at module level) ensures it is never imported before CUDA is initialised.

import torch


def beta_binomial_nll_per_sample(alpha: torch.Tensor,
                                 y_soft: torch.Tensor,
                                 n_annotators: int = 3) -> torch.Tensor:
    device = alpha.device
    alpha = alpha.cpu()
    y_soft = y_soft.cpu()

    a0 = alpha[:, 0]
    a1 = alpha[:, 1]

    y_soft = y_soft.view(-1)
    k = torch.round(y_soft * n_annotators).to(dtype=torch.long)

    kf = k.to(dtype=a0.dtype)
    nf = torch.tensor(float(n_annotators), dtype=a0.dtype)

    log_comb = torch.lgamma(nf + 1.0) - torch.lgamma(kf + 1.0) - torch.lgamma((nf - kf) + 1.0)

    log_B_num = (
        torch.lgamma(kf + a1)
        + torch.lgamma((nf - kf) + a0)
        - torch.lgamma(nf + a0 + a1)
    )
    log_B_den = (
        torch.lgamma(a1)
        + torch.lgamma(a0)
        - torch.lgamma(a0 + a1)
    )

    log_p = log_comb + log_B_num - log_B_den
    return (-log_p).to(device)


def dirichlet_kl_per_sample(alpha: torch.Tensor,
                            prior_alpha: torch.Tensor | None = None) -> torch.Tensor:
    from torch.distributions import Dirichlet  # lazy import — must not be at module level

    if prior_alpha is None:
        prior_alpha = torch.ones_like(alpha)
    else:
        prior_alpha = torch.as_tensor(prior_alpha, device=alpha.device, dtype=alpha.dtype)
        if prior_alpha.ndim == 1:
            prior_alpha = prior_alpha.view(1, 2).expand_as(alpha)

    device = alpha.device
    posterior = Dirichlet(alpha.cpu())
    prior = Dirichlet(prior_alpha.cpu())
    return torch.distributions.kl.kl_divergence(posterior, prior).to(device)


def evidential_bb_loss(alpha: torch.Tensor,
                       y_soft: torch.Tensor,
                       w_pos: float,
                       w_neg: float,
                       lambda_reg: float = 0.05,
                       use_kl: bool = True,
                       prior_alpha: torch.Tensor | None = None,
                       n_annotators: int = 3) -> torch.Tensor:
    """
    Final composed loss for one head:
        per-sample: w(y_hard) * L_BB + lambda_reg * KL   (KL optional)
    where y_hard is majority vote derived from k.

    w_pos, w_neg are precomputed scalars (e.g., w+ = N/(2N+), w- = N/(2N-)).
    """
    bb = beta_binomial_nll_per_sample(alpha, y_soft, n_annotators=n_annotators)  # (B,)

    k = torch.round(y_soft.view(-1) * n_annotators).to(dtype=torch.long)
    y_hard = (k >= (n_annotators // 2 + 1))  # bool (B,)

    w_pos_t = torch.tensor(float(w_pos), device=alpha.device, dtype=bb.dtype)
    w_neg_t = torch.tensor(float(w_neg), device=alpha.device, dtype=bb.dtype)
    w = torch.where(y_hard, w_pos_t, w_neg_t)  # (B,)

    loss = w * bb  # (B,)

    if use_kl:
        kl = dirichlet_kl_per_sample(alpha, prior_alpha=prior_alpha)  # (B,)
        loss = loss + lambda_reg * kl

    return loss.mean()


#weights = {'C1': (w_pos, w_neg),
#           'C2': (w_pos, w_neg),
#           'C3': (w_pos, w_neg)}
# w_pos e.g. 1.76
# w_neg e.g. 0.54
#prior_alpha = { 'C1': (prior_a0, prior_a1),
#                'C2': (prior_a0, prior_a1),
#                'C3': (prior_a0, prior_a1)}
# prior_a0 = (1-pi)*v e.g. (1-0.05)*2 = 0.95*2 = 1.9
# prior_a1 = pi*v, e.g. 0.05*2 = 0.1

def total_bb_loss(x: torch.Tensor,
                  y: torch.Tensor,
                  weights: dict,
                  lambda_reg: float = 0.05,
                  use_kl: bool = True,
                  prior_alpha: dict | None = None,
                  n_annotators: int = 3) -> torch.Tensor:

    total_loss = 0.0
    for i, key in enumerate(['C1', 'C2', 'C3']):
        w_pos, w_neg = weights[key]

        pa = None
        if use_kl and prior_alpha is not None:
            pa = torch.tensor(prior_alpha[key], device=x[i].device, dtype=x[i].dtype)

        total_loss += evidential_bb_loss(x[i],
                                         y[:, i],
                                         w_pos=w_pos,
                                         w_neg=w_neg,
                                         lambda_reg=lambda_reg,
                                         use_kl=use_kl,
                                         prior_alpha=pa,
                                         n_annotators=n_annotators)
    return total_loss

---

Notes

• This affects any HPC or cloud environment where PyTorch's prebuilt wheels target a newer CPU microarchitecture than the compute nodes.
• The symptom (silent SIGILL) is especially confusing because training appears to start and then vanishes with no error output.
• BCE and other standard losses are unaffected because they do not use torch.lgamma or torch.distributions on GPU.
• Confirmed working with PyTorch 2.5.1+cu118 on NVIDIA RTX A6000 after applying the fix.