enforce recompute flag on fsdpa quantization (#133)

HabanaAI · Apr 19, 2024 · b8c073a · b8c073a
1 parent b79036e
commit b8c073a
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/optimum/habana/transformers/models/llama/modeling_llama.py b/optimum/habana/transformers/models/llama/modeling_llama.py
@@ -1,3 +1,4 @@
+import os
 import math
 import warnings
 from typing import List, Optional, Tuple, Union
@@ -414,7 +415,8 @@ def pre_attn_forward(
 
             if q_len == 1:
                 # next token
-                with ht.sdp_kernel(enable_recompute=False):
+                use_recompute = True if os.getenv("QUANT_CONFIG", "") else False
+                with ht.sdp_kernel(enable_recompute=use_recompute):
                     attn_output = self.fused_scaled_dot_product_attention(
                         query_states, key_states, value_states, attention_mask, 0.0, False, None
                     )