Dump activation shardings

[charlesli640]

Description

To dump activation shardings to golden file for further comparison. It can include in unit test in case further code change touches activation shardings. This PR is the initial submission for sharding dump json files.

Output

The output format is readable and comparable by both human and machine. For exampletests deepseek2-16b/v5p-16/slice_1 activation dump as below

{
  "Activation Sharding Dump": [
    {
      "deepseek/inputs: bfloat16[96,2048,2048]": {
        "logic_axes": "('activation_batch', 'activation_norm_length', 'activation_embed')",
        "PartitionSpec": "P('fsdp', None, None)"
      }
    },
    {
      "deepseek/pre_attention_norm: bfloat16[96,2048,2048]": {
        "logic_axes": "('activation_batch', 'activation_norm_length', 'activation_embed')",
        "PartitionSpec": "P('fsdp', None, None)"
      }
    },
    {
      "attention_mla/inputs_q: bfloat16[96,2048,2048]": {
        "logic_axes": "('activation_batch', 'activation_length_no_exp', 'activation_embed')",
        "PartitionSpec": "P('fsdp', None, None)"
      }
    },
    {
      "attention_mla/inputs_kv: bfloat16[96,2048,2048]": {
        "logic_axes": "('activation_batch', 'activation_length_no_exp', 'activation_embed')",
        "PartitionSpec": "P('fsdp', None, None)"
      }
    },
    {
      "attention_mla/q_nope: bfloat16[96,2048,16,128]": {
        "logic_axes": "('activation_kv_batch', 'activation_length_no_exp', 'activation_kv_heads', 'activation_kv_head_dim')",
        "PartitionSpec": "P('fsdp', None, None, None)"
      }
    },
    {
      "attention_mla/q_pe: bfloat16[96,2048,16,64]": {
        "logic_axes": "('activation_kv_batch', 'activation_length_no_exp', 'activation_kv_heads', 'activation_kv_head_dim')",
        "PartitionSpec": "P('fsdp', None, None, None)"
      }
    },
    {
      "attention_mla/query: bfloat16[96,2048,16,192]": {
        "logic_axes": "('activation_kv_batch', 'activation_length_no_exp', 'activation_kv_heads', 'activation_kv_head_dim')",
        "PartitionSpec": "P('fsdp', None, None, None)"
      }
    },
    {
      "attention_mla/key_nope: bfloat16[96,2048,16,128]": {
        "logic_axes": "('activation_kv_batch', 'activation_length_no_exp', 'activation_kv_heads', 'activation_kv_head_dim')",
        "PartitionSpec": "P('fsdp', None, None, None)"
      }
    },
    {
      "attention_mla/key_rope: bfloat16[96,2048,16,64]": {
        "logic_axes": "('activation_kv_batch', 'activation_length_no_exp', 'activation_kv_heads', 'activation_kv_head_dim')",
        "PartitionSpec": "P('fsdp', None, None, None)"
      }
    },
    {
      "attention_mla/key: bfloat16[96,2048,16,192]": {
        "logic_axes": "('activation_kv_batch', 'activation_length_no_exp', 'activation_kv_heads', 'activation_kv_head_dim')",
        "PartitionSpec": "P('fsdp', None, None, None)"
      }
    },
    {
      "attention_mla/value: bfloat16[96,2048,16,128]": {
        "logic_axes": "('activation_kv_batch', 'activation_length_no_exp', 'activation_kv_heads', 'activation_kv_head_dim')",
        "PartitionSpec": "P('fsdp', None, None, None)"
      }
    },
    {
      "attention_op/query: bfloat16[96,16,2048,192]": {
        "logic_axes": "Unknown",
        "PartitionSpec": "P('fsdp', None, None, None)"
      }
    },
    {
      "attention_op/key: bfloat16[96,16,2048,192]": {
        "logic_axes": "Unknown",
        "PartitionSpec": "P('fsdp', None, None, None)"
      }
    },
    {
      "attention_op/value: bfloat16[96,16,2048,128]": {
        "logic_axes": "Unknown",
        "PartitionSpec": "P('fsdp', None, None, None)"
      }
    },
    {
      "attention_mla/out: bfloat16[96,2048,16,128]": {
        "logic_axes": "('activation_batch', 'activation_length_no_exp', 'activation_heads', 'activation_kv')",
        "PartitionSpec": "P('fsdp', None, None, None)"
      }
    },
    {
      "deepseek/attention_result: bfloat16[96,2048,2048]": {
        "logic_axes": "('activation_batch', 'activation_norm_length', 'activation_embed')",
        "PartitionSpec": "P('fsdp', None, None)"
      }
    },
    {
      "deepseek/post_attention_norm: bfloat16[96,2048,2048]": {
        "logic_axes": "('activation_batch', 'activation_norm_length', 'activation_embed')",
        "PartitionSpec": "P('fsdp', None, None)"
      }
    },
    {
      "linears/x: bfloat16[96,2048,10944]": {
        "logic_axes": "('activation_batch', 'activation_length_no_exp', 'activation_mlp')",
        "PartitionSpec": "P('fsdp', None, None)"
      }
    },
    {
      "deepseek/mlp: bfloat16[96,2048,2048]": {
        "logic_axes": "('activation_batch', 'activation_norm_length', 'activation_embed')",
        "PartitionSpec": "P('fsdp', None, None)"
      }
    },
    {
      "deepseek/x: bfloat16[96,2048,2048]": {
        "logic_axes": "('activation_batch', 'activation_norm_length', 'activation_embed')",
        "PartitionSpec": "P('fsdp', None, None)"
      }
    },
    {
      "moe/inputs: bfloat16[96,2048,2048]": {
        "logic_axes": "('activation_batch', 'activation_norm_length', None)",
        "PartitionSpec": "P('fsdp', None, None)"
      }
    },
    {
      "moe/gate_logits: bfloat16[96,2048,64]": {
        "logic_axes": "('activation_batch', 'activation_norm_length', None)",
        "PartitionSpec": "P('fsdp', None, None)"
      }
    },
    {
      "linears/x: bfloat16[96,2048,2816]": {
        "logic_axes": "('activation_batch', 'activation_length_no_exp', 'activation_mlp')",
        "PartitionSpec": "P('fsdp', None, None)"
      }
    },
    {
      "deepseek/mlp_lnx: bfloat16[96,2048,2048]": {
        "logic_axes": "('activation_batch', 'activation_norm_length', 'activation_embed')",
        "PartitionSpec": "P('fsdp', None, None)"
      }
    }
  ]
}

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.

[codecov]

Codecov Report

❌ Patch coverage is 68.42105% with 12 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/MaxText/sharding.py	60.00%	11 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!

[gobbleturk]

I think this LGTM although there are a lot of names to review! How did you generate these names?

[charlesli640]

I think this LGTM although there are a lot of names to review! How did you generate these names?

These names are generated from local <file_name>/<variable_name>. Sometimes it may not correctly reflect the actual model/layer, but it is basically serving as an identifier/key for logging/dumping/comparing purpose.

[NuojCheng]

Thanks Charles! Just some minor comments