Use causal_padding instead of padding

Signed-off-by: Reese Wang <[email protected]>
google · Mar 11, 2024 · d2ba1bc · d2ba1bc
1 parent e3e785c
commit d2ba1bc
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/praxis/contrib/gpu/scripts_gpu/te_helper.py b/praxis/contrib/gpu/scripts_gpu/te_helper.py
@@ -208,7 +208,7 @@ def update_attn_te_tpl(te_tpl, attn_tpl):
         assert (transformer_layer_tpl.tr_fflayer_tpl.has_bias ==
             transformer_layer_tpl.tr_atten_tpl.use_bias), "TE only allows same bias settings."
         te_transformer_tpl.use_bias = transformer_layer_tpl.tr_fflayer_tpl.has_bias
-        te_transformer_tpl.self_attn_mask_type = 'causal' \
+        te_transformer_tpl.self_attn_mask_type = 'causal_padding' \
             if stacked_transformer_obj.mask_self_attention else 'padding'
 
         te_transformer_tpl.logical_axes_rules = te_flax.extend_logical_axis_rules(tuple())