Expand sharding dump

[Shuang-cnt]

Description

Expand sharding_dump.py to output logical axes and update unit test sharding_compare_test.py
Currently, two json files (logical_shardings.json and named_shardings.json ) will be generated per model/device_type/slice_number.

Command:

1. Get baseline sharding info: python3 -m tests.utils.run_sharding_dump
2. Compare sharding info: python3 -m pytest tests/unit/sharding_compare_test.py -s -v -k "llama3.1-70b" 2>&1 | tee test_output.log

Example

Content in logical_shardings.json

".params/['params']/['decoder']/['layers']/['mlp']/['wi_0']/['kernel']": {
    "partition_spec": [
      "embed",
      "layers",
      "mlp"
    ],
    "shape": [
      8192,
      80,
      28672
    ]
  },

Content in named_shardings.json

".params/['params']/['decoder']/['layers']/['mlp']/['wi_0']/['kernel']": {
    "mesh": {
      "axis_names": [
        "data",
        "stage",
        "fsdp",
        "fsdp_transpose",
        "sequence",
        "context",
        "context_autoregressive",
        "tensor",
        "tensor_transpose",
        "tensor_sequence",
        "expert",
        "autoregressive"
      ],
      "shape": {
        "data": 1,
        "stage": 1,
        "fsdp": 16,
        "fsdp_transpose": 1,
        "sequence": 1,
        "context": 1,
        "context_autoregressive": 1,
        "tensor": 1,
        "tensor_transpose": 1,
        "tensor_sequence": 1,
        "expert": 1,
        "autoregressive": 1
      }
    },
    "partition_spec": [
      [
        "fsdp",
        "sequence",
        "tensor_transpose",
        "context",
        "expert"
      ],
      "stage",
      [
        "fsdp_transpose",
        "tensor",
        "tensor_sequence",
        "autoregressive"
      ]
    ],
    "shape": [
      8192,
      80,
      28672
    ]
  },

Tests

UT for sharding dump comparison Failed (physical weight) : https://paste.googleplex.com/6032726044049408

[FAIL] Physical Sharding Mismatch: llama3.1-70b v5e-16 slice 1

PartitionSpec mismatch at '.params/['params']/['decoder']/['decoder_norm']/['scale']':
  Expected (Physical): [['fsdp', 'tensor_transpose']]
  Actual (Physical): [['tensor', 'tensor_transpose']]

UT for sharding dump comparison Failed (logical): https://paste.googleplex.com/5855428334452736

[FAIL] Logical Sharding Mismatch: llama3.1-70b v5e-16 slice 1

Shape mismatch at '.params/['params']/['decoder']/['decoder_norm']/['scale']':
  Expected (Logical): [80]
  Actual (Logical): [8192]

UT for sharding dump comparison Successed: https://paste.googleplex.com/6737857618247680

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

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[richjames0]

Thanks!

[NuojCheng]

nit: 2023-2026

[NuojCheng]

it looks like this test only verifies get_shaped_inputs and get_topology_mesh from train_compile. Please explicitly specify that this test is for train_compile only if this is the intention. Otherwise, I think it is better to test the get_abstract_state from maxtext_utils.