fix: guard softmax_offset.grad.zero_() by is_training

[andrewwhitecdw]

This PR addresses the following issue in tests/pytorch/attention/run_attention_with_cp.py: guard softmax_offset.grad.zero_() by is_training.

Changes

• tests/pytorch/attention/run_attention_with_cp.py: guard softmax_offset.grad.zero_() by is_training.

Details

--- a/tests/pytorch/attention/run_attention_with_cp.py
+++ b/tests/pytorch/attention/run_attention_with_cp.py
@@ -1,3 +1,3 @@
-    if config.softmax_type != "vanilla":
-        core_attn.softmax_offset.grad.zero_()
-    if dtype == "fp8":
+    if is_training and config.softmax_type != "vanilla":
+        core_attn.softmax_offset.grad.zero_()
+    if dtype == "fp8":

Tests

Let me know if you want tests added for this fix or not.

[greptile-apps]

Greptile Summary

The change prevents inference-mode context-parallel attention tests from calling zero_() on a nonexistent softmax-offset gradient.

• Adds the is_training condition to the softmax-offset gradient reset.
• Preserves gradient clearing between the non-context-parallel and context-parallel backward passes during training.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
tests/pytorch/attention/run_attention_with_cp.py	Correctly limits softmax-offset gradient clearing to training, the only mode in which this harness performs backward propagation and creates that gradient.

_{Reviews (2): Last reviewed commit: "Merge branch 'main' into bugfix/run-atte..." | Re-trigger Greptile}