facebookresearch / facebookresearch/sam3
Reduced Triton sigmoid focal loss returns NaN gradients for gamma=0 when logits saturate
- Dominant language
- Python
- Stars
- 11.7k
- Forks
- 1.8k
- PR merge metrics
- No merged PRs in 30d
Description
### Summary
`triton_sigmoid_focal_loss_reduce` can return NaN gradients for finite inputs when `gamma=0` and logits saturate in the correct direction. This matters for SAM3 training because `IABCEMdetr` defaults `presence_gamma=0.0`, so the presence loss can hit this reduced Triton focal path.
For `gamma=0`, sigmoid focal loss should reduce to alpha-weighted BCE, whose gradient is finite:
```text
dL/dx = alpha_t * (sigmoid(x) - target)
```
However, the Triton backward computes `(1 - p_t) ** (gamma - 1)` before simplification. When `gamma=0` and `p_t == 1`, this becomes `0 ** -1 = inf`, and later terms produce `0 * inf = NaN`.
### Minimal Reproduction
```python
import torch
from sam3.train.loss.sigmoid_focal_loss import triton_sigmoid_focal_loss_reduce
x = torch.tensor([[18.0]], device="cuda", requires_grad=True)
y = torch.tensor([[1.0]], device="cuda")
loss = triton_sigmoid_focal_loss_reduce(x, y, 0.5, 0.0)
loss.backward()
print(loss, x.grad)
```
Observed:
```text
loss = tensor(0., device='cuda:0', grad_fn=)
x.grad = tensor([[nan]], device='cuda:0')
```
The same issue occurs for `x=-18.0, y=0.0`.
### Expected Behavior
The gradient should be finite. With `gamma=0`, the loss is equivalent to alpha-weighted BCE, so saturated correct predictions should produce a zero or very small finite gradient, not NaN.
### Related Presence Logit Clamp Issue
There is also an ineffective clamp in `sam3/model/decoder.py`:
```python
intermediate_layer_presence_logits.clamp(...)
```
Since the return value is not assigned and this is not `clamp_`, the intended presence-logit clamp is not applied. This can let presence logits reach the saturation range that triggers the `gamma=0` backward issue.
### Environment
```text
torch 2.12.0+cu132
triton 3.7.0
CUDA 13.2
GPU: NVIDIA RTX 6000 Ada Generation
```
### Proposed Fix
A small fix is to route `sigmoid_focal_loss(..., gamma=0)` through the existing PyTorch fallback path, where the expression is numerically stable and mathematically equivalent to BCE, and to assign the presence clamp result.
Contributor guide
Research direction
Start with sam3/train/loss/sigmoid_focal_loss.py and reproduce the issue using the CUDA example with gamma=0 and saturated logits. Compare the Triton path with the existing PyTorch fallback, then inspect sam3/model/decoder.py for the ineffective presence-logit clamp. Done means finite gradients for the reported cases and an applied presence clamp.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 65/100