RoPE fp32 + gradient checkpointing + SDPA attention backend

#1
by deshwalmahesh - opened

Three fixes to modeling_domynedge.py. Default CUDA varlen path untouched (SDPA is opt-in): no GPU regression. Checks: CPU, tiny random config, no weights (torch 2.15).

1. RoPE precision β€” resolves the TODO in DomynEdgeRotaryEmbedding.forward

  • Found: angles round to bf16 at large positions β€” under autocast(bfloat16) the matmul runs bf16, and after model.bfloat16() the inv_freq buffer is bf16 (both noted in the TODO).
  • Why: cos/sin drift at long context. Real val rows (bf16): prob shift vs fp32 ~8% at <1k, ~53% at 4–8k; accuracy 53% β†’ 27% at 4–8k.
  • Fix: disable autocast, recompute inv_freq in fp32 in forward (as the TODO proposes).
  • Proved: cos error at pos 7000 = 0.0 (fix) vs 2.0 (old).

2. Gradient checkpointing β€” supports_gradient_checkpointing was False

  • Found: flag False β†’ gradient_checkpointing_enable() is a no-op.
  • Why: 16,384-token training needs ~161 GB of activations β€” won't fit one 141 GB GPU without recompute.
  • Fix: layers are stateless β†’ set flag, route layer loop through _gradient_checkpointing_func when training.
  • Proved: gradients identical to non-checkpointed (max diff 0.0).

3. SDPA attention β€” _supports_sdpa was False

  • Found: attention hard-wired to CUDA torch.nn.attention.varlen; attn_implementation ignored, so no portable path.
  • Why: can't run without the flash/varlen kernel; no fp32-capable attention.
  • Fix: opt-in attn_implementation="sdpa" branch β€” same GQA + per-layer sliding window + padding mask via scaled_dot_product_attention; default path unchanged.
  • Proved: matches independent masked-attention reference to 1.5e-8; full forward runs on CPU with finite logits (impossible on varlen). GPU varlen-vs-SDPA parity left to confirm on GPU.

Hardware

before after
H200 / A100 (flash varlen) βœ… βœ… unchanged + gradient checkpointing
T4 / older GPU (no varlen) ❌ βœ… sdpa
CPU ❌ βœ… sdpa (separately tested)

Thank you @deshwalmahesh for the PR, by pure coincidence, we had been working on the same fixes a few days earlier in the original Domyn Edge training repo https://github.com/DomynAI/domyn-forge/pull/12.
The problems with RoPE precision and SDPA attention should be fixed, while you are more than welcome to add support on Gradient Checkpointing! I will close this PR now, but feel free to open a PR on the original training repo: https://github.com/DomynAI/domyn-forge/

veneres-domyn changed pull request status to closed

Sign up or log in to comment