Removing sequence from 'activation_vocab' sharding rule

[abhinavgoel95]

Description

Without this change, the automatic sharding resolution fails to pick up the correct sharding for activation_vocab when tensor_sequence is used.

Removing sequence from 'activation_vocab' sharding rule

Tests

Tested Llama3 suite of models using tensor_sequence. Leads to 10% improvement in step time, because the vocab embedding GeMM now gets correct sharded.

"python -m MaxText.train /opt/workspace/te-maxtext/src/MaxText/configs/base.yml run_name=logdir use_iota_embed=true scan_layers=False    steps=31 per_device_batch_size=0.5 model_name=${MODEL} remat_policy=minimal_flash enable_checkpointing=false logits_dot_in_fp32=false    base_output_directory=/opt/maxtext/local_train dataset_path=local dataset_type=synthetic attention=cudnn_flash_te max_target_length=8192    enable_goodput_recording=false monitor_goodput=false quantization="" tokenizer_path=tiktoken hardware=gpu    dcn_fsdp_parallelism=1 ici_data_parallelism=1    ici_fsdp_parallelism=2 dcn_data_parallelism=1    ici_tensor_sequence_parallelism=4 dcn_tensor_parallelism=1 ${PROFILING_SETTINGS}"

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[gobbleturk]

Thanks Abhinav!

Apologies for our terribly complex rules, we have a fix in the works!! (eventaully...)