fix fsdp2d ds logical rule

Author: NuojCheng

src/MaxText/configs/models/deepseek3-671b-2dfsdp.yml

@@ -61,6 +61,8 @@ data_sharding: [['data', 'fsdp', 'fsdp_transpose', 'expert', 'context']]
 logical_axis_rules: [
     ['activation_batch', ['data', 'fsdp', 'fsdp_transpose', 'expert', 'context']],
     ['activation_kv_batch', ['data', 'fsdp', 'fsdp_transpose', 'expert', 'context']],
+    ['activation_embed_and_logits_batch', ['data', 'fsdp', 'fsdp_transpose', 'expert']],
+    ['activation_norm_length', ['context']],
     ['activation_heads', []],
     ['embed', ['fsdp']],
     ['embed_no_exp', ['fsdp']],