Update test_fused_softmax.py (#1782)

eqy · web-flow · commit 810ffae374a2 · 2024-03-12T13:38:33.000+09:00
diff --git a/tests/L0/run_transformer/test_fused_softmax.py b/tests/L0/run_transformer/test_fused_softmax.py
@@ -346,6 +346,8 @@ def test_forward(self, allmasked: bool=False):
 
     def test_backward(self, allmasked: bool=False):
         import generic_scaled_masked_softmax_cuda
+        prev_thresh = self.thresh
+        self.thresh = {"atol": 1.5e-1, "rtol": 5e-3}
         for qlen, klen in self.q_k_lens:
             inputs = torch.normal(0, 2, (self.batch, self.attn, qlen, klen), dtype=self.dtype, device=self.device)
             backward = torch.rand_like(inputs, dtype=torch.float16, device=self.device)
@@ -359,6 +361,7 @@ def test_backward(self, allmasked: bool=False):
             softmax_results_torch = forward_torch_softmax(inputs, masks, self.scale_t)
             softmax_results_torch.backward(backward)
             self.assertEqual(back_grad, inputs.grad, **self.thresh, msg=f"(q, k) = ({qlen, klen})")
+        self.thresh = prev_thresh
 
     def test_allmasked(self):
         self.test_forward(True)
