[megatron] Fix TE sliding-window attention crash at micro_batch>1 on the non-packed forward path

[john-hahn]

[megatron] Fix TE sliding-window attention crash at micro_batch>1 on the non-packed forward path

Summary

On the non-packed Megatron forward path (remove_microbatch_padding=False),
megatron_model_wrapper.py passes the 2-D [batch, seq] keep-mask returned by
remove_left_padding() straight to model(...). For models whose attention uses a Transformer
Engine sliding-window mask, TE's get_full_mask combines that mask with the SWA mask via
torch.logical_or(swa_mask, attention_mask). A 2-D [batch, seq] mask is not broadcastable against
the [..., seq, seq] SWA mask, so at micro_batch_size > 1 the batch dimension collides with the
sequence dimension and the forward crashes.

This adds a pure-torch helper to_te_attention_mask that converts the 2-D keep-mask to a
[batch, 1, 1, seq] padding mask (True marks padding), and routes the model(...) call sites
through it. The original 2-D mask is left untouched for the downstream recover_left_padding().

Reproduction (current main, 0a18dd72)

Real model Qwen/Qwen3.6-35B-A3B (hybrid GatedDeltaNet + MoE; its full-attention layers use TE
sliding-window attention), tp=2, pp=1, micro_batch_size=2, remove_microbatch_padding=False,
left-padded batch:

• Stock main crashes in the megatron forward:

  transformer_engine/pytorch/attention/dot_product_attention/utils.py:1363, in get_full_mask
      attention_mask = torch.logical_or(swa_mask, attention_mask)
  RuntimeError: The size of tensor a (24) must match the size of tensor b (2) at non-singleton dimension 2
  

  (24 = sequence length, 2 = micro_batch_size.)
• With this patch the same forward completes and returns logprobs — test passes.
• Regression: a standard non-SWA model (Qwen/Qwen3-0.6B, non-packed, micro_batch_size=2,
  left-padded) passes the HF-vs-megatron parity test identically on stock and patched — the
  helper is a no-op there.

Lifting micro_batch_size from 1 to 4 on this path (once unblocked) increased training throughput
~3.3× at unchanged reward in our multi-LoRA RL runs.

Changes

• Add skyrl/backends/skyrl_train/distributed/megatron/mask_utils.py with to_te_attention_mask.
• Route both model(...) call sites in megatron_model_wrapper.py (4 invocations after [megatron] fused linear cross-entropy for 23x memory savings at 262k context #1841)
  through the helper.
• Add tests/backends/skyrl_train/distributed/test_mask_utils.py.

Design notes

• The reshape lives in a standalone pure-torch helper so it is unit-testable on CPU without the
  megatron extra.
• No-op on the packed path: there new_attention_mask is None and packed_seq_params drives
  masking, so the helper returns the input unchanged. An already higher-rank mask is also returned
  unchanged, so non-SWA TE attention (which already accepts the 2-D mask) is unaffected.
• Not folded into remove_left_padding() because that function's 2-D return is also consumed by
  recover_left_padding(), which needs the 2-D form.
• True = padding matches TE's mask convention. GatedDeltaNet / linear-attention layers ignore
  attention_mask. No config or API change; behavior is identical at micro_batch_size = 1.

Testing

• Unit (CPU): uv run --extra dev pytest tests/backends/skyrl_train/distributed/test_mask_utils.py
  — 5 passed.
• Integration (GPU, run on 2× B300): test_megatron_forward[tp2_pp1_policy] with
  Qwen/Qwen3.6-35B-A3B — fails on stock at get_full_mask, passes with this patch;
  Qwen/Qwen3-0.6B passes on both (regression-safe).

Checklist

• Follows Google Python style; comments describe what the code does.
• CPU unit test added; GPU reproduction + fix verified on Qwen3.6-35B-A3B.
• No paths/naming changes affecting .claude/ docs; no example-script or doc changes needed.
• bash format.sh (ruff 0.11.9 + black 24.10.0 + gitleaks) clean — verified, zero auto-fixes.

[gemini-code-assist]

Code Review

This pull request introduces a helper function to_te_attention_mask to convert 2-D keep-masks into 4-D padding masks, preventing dimension collisions in Transformer Engine attention when the micro-batch size is greater than one. The helper is integrated into the Megatron model wrapper's forward step, and comprehensive unit tests have been added to verify its behavior. There are no review comments, and I have no additional feedback to provide.

[erictang000]

hey @john-hahn, tested and the fix makes sense, running gpu ci now, we can likely get this merged tomorrow pending CI

just checking is there a reason you can't use the packed path?

[john-hahn]

hey @erictang000, just pulled the latest changes and saw your PR #1769 for qwen packing

when i first hit this issue, GDN couldnt pack, which is why i went the unpacked route. for some context, i was hitting an error like:

NotImplementedError: GDN does not support packed sequence

i would say this PR fix is still relevant because anyone who uses remove_microbatch_padding=False config will silently crash at micro_batch>1 for any TE-SWA model

thanks for adding the packed support!