[Fix] Avoid fp16-overflow NaNs in TTT fused-chunk eta gradient (#1597)

Author: siyudong-syd

fla/ops/ttt/fused_chunk.py

@@ -425,8 +425,8 @@ def fused_chunk_ttt_linear_bwd_kernel_dh(
         mask = (o_i == i_last)
         b_dk -= b_e_last * tl.dot(b_v2, tl.trans(b_dh).to(b_v2.dtype))
         b_dk -= tl.dot(tl.trans(b_ds), tl.trans(b_q) * b_e[:, None])
-        b_de = mask * tl.sum(- b_dh * tl.trans(tl.dot(tl.trans(b_v2), b_k))).to(b_k.dtype)
-        b_de -= mask * tl.sum(b_dhb * tl.sum(b_v2, axis=0)).to(b_k.dtype)
+        b_de = tl.where(mask, tl.sum(- b_dh * tl.trans(tl.dot(tl.trans(b_v2), b_k))), 0.0)
+        b_de -= tl.where(mask, tl.sum(b_dhb * tl.sum(b_v2, axis=0)), 0.0)
         b_de -= tl.sum(tl.dot(b_ds, b_k) * tl.trans(b_q).to(b_k.dtype), axis=1)
         b_de -= tl.sum(b_ds, axis=1)
         b_dh += tl.dot(b_q, b_do.to(b_q.dtype)) + tl.dot(tl.trans(b_k).to(b_dkh.dtype), b_dkh)

tests/ops/test_ttt.py

@@ -44,6 +44,7 @@ def test_chunk(
         pytest.skip(reason="Current CI do not support this config")
     if T > 1000:
         pytest.skip(reason="Current CI do not support this config")
+    torch.manual_seed(42)
     eta_base = 5e-3
     q = torch.randn(B, T, H, D, dtype=dtype)
     k = F.normalize(torch.randn(B, T, H, D, dtype=torch.float32), p=2, dim=-1).to(dtype)
@@ -136,6 +137,7 @@ def test_fused_chunk(
         pytest.skip(reason="Current CI do not support this config")
     if T > 1000:
         pytest.skip(reason="Current CI do not support this config")
+    torch.manual_seed(42)
     eta_base = 5e-3
     q = torch.randn(B, T, H, D, dtype=dtype)
     k = F.normalize(torch.randn(B, T, H, D, dtype=torch.float32), p=2, dim=-1).to(dtype)