Update modeling_chatglm.py

model/modeling_chatglm.py

@@ -277,7 +277,7 @@ def attention_fn(
         
         q, k, v = [rearrange(x, 'b s h d -> b h s d') for x in [query_layer, key_layer, value_layer]]
         ctx_lens1 = ctx_lens.to(q.device)
-        output = flash_attn_func(q, k, v,query_key_layer_scaling_coeff,ctx_lens1)
+        output = flash_attn_func(q, k, v,query_key_layer_scaling_coeff,ctx_lens1,3)
         context_layer = output.permute(2, 0, 1, 3)
         new_context_layer_shape = context_layer.size()[:-2] + (hidden_size_per_partition,)
         context_layer = context_layer.view(*new_context_layer_shape)