NVIDIA / NVIDIA/Megatron-LM

Sliding-window path in FusedScaleMaskSoftmax silently discards the caller-provided attention mask

Open
#7,039 0 comments 0 reactions 0 assignees View on GitHub
community-request
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 6h
Merged PRs (30d)
271

Description

## Summary

In `megatron/core/fusions/fused_softmax.py::forward_torch_softmax` the sliding-window branch **unconditionally overwrites** the caller-provided mask:

```python
# Generate causal mask if not given
sq, sk = input.size(2), input.size(3)
if self.window_size is not None:
mask = get_sliding_window_causal_mask(sq, sk, self.window_size) # <- discards the incoming mask
elif self.attn_mask_type == AttnMaskType.causal and mask is None and sq > 1:
...
mask = get_default_causal_mask(sq)
```

Note the asymmetry: the plain-causal branch is correctly guarded by `mask is None`, the sliding-window branch is not.

## Impact

For any model with sliding-window layers, a padding mask (micro_batch_size > 1) or a packed-sequence block-diagonal mask passed by the caller is silently ignored on the SWA layers: pad tokens / cross-document tokens receive attention, the loss is quietly polluted, and nothing raises.

We hit this while running full SFT of MiMo-V2.5 (39 of 48 layers are SWA with this torch-fallback path): the only safe configurations today are micro_batch_size=1 and no sequence packing, which is a significant efficiency loss for variable-length multimodal data.

## Repro sketch

Build `FusedScaleMaskSoftmax(..., attn_mask_type=AttnMaskType.causal, scaled_masked_softmax_fusion=False, window_size=(4, 0))`, pass zero scores of shape [2, 1, 8, 8] and a bool padding mask masking the last 3 keys of sample 1 — the returned probabilities at the padded keys are non-zero.

## Proposed fix

Compose instead of overwrite (both masks are bool with True = masked; `[b,1,sq,sk] | [sq,sk]` broadcasts):

```python
swa_mask = get_sliding_window_causal_mask(sq, sk, self.window_size)
mask = swa_mask if mask is None else (mask | swa_mask)
```

Behavior is unchanged for every caller that passes `mask=None`. PR with this fix + a unit test incoming (scope: the torch fallback path only; fused/TE kernel paths are untouched).

Contributor guide

Open the contributing guide

Research direction

Start in megatron/core/fusions/fused_softmax.py at forward_torch_softmax and reproduce the sliding-window case with a caller-provided padding mask. Check the existing causal-mask handling and add coverage for combining the caller mask with the sliding-window mask; done means padded or cross-document keys remain masked while mask=None behavior is unchanged.

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
Active
Clarity
Clearly specified
Newbie friendliness
72/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.