Use causal_padding instead of padding

Signed-off-by: Reese Wang <[email protected]>

Author: zlsh80826

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())