feat: Automated Post-Conversion Checkpoint Shape and Mock Tensor Vali…

[olufiyin19]

…dation for the Checkpoint Validation Agent

Description

This PR introduces the first two phases of the automated post checkpoint conversion validation pipeline. It ensures that converted ML checkpoints (e.g. Gemma, Llama) are structurally sound and computationally viable before they are deployed to TPU clusters.

The validation is broken into two sequential subtasks:

1. Static Shape Validation (checkpoint_shape_validator.py): Parses and compares the theoretical MaxText blueprint against the actual Orbax checkpoint layer by layer to ensure 1:1 parameter mapping.
2. Dynamic Mock Tensor Validation (mock_tensor_test.py): Performs a dynamic, "hollow" forward pass using jax.eval_shape. This ensures the converted weights perfectly align with the network's expected graph dimensions without requiring physical TPU memory.

Key Implementation Details

1. Static Shape Validator

• Implemented strict dictionary comparison logic to flag expected vs. actual dimension mismatches.
• Modularized the comparison logic into check_mismatches to support local CI unit testing without requiring temp files.
• Added test_checkpoint_shape_validator.py to unit test the dictionary comparison logic.

2. Mock Tensor Trace
   To enable single-pod Kubernetes execution, the following adjustments were made to the forward pass trace:

• NNX Initialization: Used the transformer_as_linen wrapper to preserve .init and .apply shape tracing.
• Distributed Bypass: Injected skip_jax_distributed_system=true into the pyconfig to allow the model to build without a TPU coordinator address.
• Hardware Mocking: Implemented a dummy 1-device Mesh (('fsdp', 'tensor', 'sequence')) to satisfy internal sharding layout requirements in utils/sharding.py.
• Positional Execution: Explicitly passed hollow decoder_positions and decoder_segment_ids tensors to the __call__ method, satisfying the strict positional arguments of the refactored Transformer.

Tests

Verified dynamic structural tracing across four distinct architectures. The script correctly parses Airflow overrides, auto-detects Hugging Face vocab sizes, and traces abstract zero-memory geometries.

1. Llama 3.1-8b

**Airflow Trigger JSON:** ```json { "checkpoint_gcs_path": "gs://maxtext-model-checkpoints/llama3.1-8b/2025-01-23-19-04/unscanned/checkpoints/0/items", "maxtext_model_name": "llama3.1-8b", "maxtext_overrides": { "attention": "dot_product", "max_target_length": 4096, "per_device_batch_size": 4, "scan_layers": false, "tokenizer_path": "meta-llama/Llama-3.1-8B-Instruct", "tokenizer_type": "huggingface" }, "run_name": "llama-context-stress-test" } ```

checkpoint_shape_validation Task Execution Log:
https://paste.googleplex.com/6204023324672000

mock_tensor_validation Task Execution Log:
https://paste.googleplex.com/4820299131518976

2. Deepseek3-671b

**Airflow Trigger JSON:** ```json { "checkpoint_gcs_path":"gs://maxtext-model-checkpoints/deepseek3-671b/0/items" "maxtext_model_name":"deepseek3-671b" "maxtext_overrides":{ "attention":"dot_product" "capacity_factor":-1 "load_balance_loss_weight":0.01 "max_target_length":1024 "per_device_batch_size":1 "scan_layers":true "sparse_matmul":true "tokenizer_path":"deepseek-ai/DeepSeek-V3" "tokenizer_type":"huggingface" } "run_name":"deepseek-moe-test" } ```

checkpoint_shape_validation Task Execution Log:
https://paste.googleplex.com/5227432637497344

mock_tensor_validation Task Execution Log:
https://paste.googleplex.com/6553731624730624

3. Mistral-7b

**Airflow Trigger JSON:** ```json { "checkpoint_gcs_path":"gs://maxtext-model-checkpoints/mistral-7b/2025-01-23-19-04/unscanned/checkpoints/0/items" "maxtext_model_name":"mistral-7b" "maxtext_overrides":{ "dtype":"bfloat16" "max_target_length":256 "per_device_batch_size":1 "remat_policy":"minimal" "scan_layers":false "tokenizer_path":"mistralai/Mistral-7B-v0.1" "tokenizer_type":"huggingface" } "run_name":"mistral-7b-dense-test" } ```

checkpoint_shape_validation Task Execution Log:
https://paste.googleplex.com/4872771216736256

mock_tensor_validation Task Execution Log:
https://paste.googleplex.com/6022800585392128

4. Gemma3-4b

**Airflow Trigger JSON:** ```json { "checkpoint_gcs_path":"gs://maxtext-model-checkpoints/gemma3-4b/2025-03-18-19-03/unscanned/checkpoints/0/items" "maxtext_model_name":"gemma3-4b" "maxtext_overrides":{ "max_target_length":128 "per_device_batch_size":1 "scan_layers":false "tokenizer_path":"google/gemma-3-4b-it" "tokenizer_type":"huggingface" } "run_name":"gemma3-4b-basic-test" } ```

checkpoint_shape_validation Task Execution Log:
https://paste.googleplex.com/6677126723010560

mock_tensor_validation Task Execution Log:
https://paste.googleplex.com/5595861106163712

5. Qwen3-8b

**Airflow Trigger JSON:** ```json { "checkpoint_gcs_path":"gs://maxtext-model-checkpoints/qwen3-8b/unscanned/0/items" "maxtext_model_name":"qwen3-8b" "maxtext_overrides":{ "attention":"dot_product" "max_target_length":2048 "per_device_batch_size":8 "scan_layers":false "tokenizer_path":"Qwen/Qwen3-8B-Instruct" "tokenizer_type":"huggingface" } "run_name":"qwen3-custom-shape-test" } ```

checkpoint_shape_validation Task Execution Log:
https://paste.googleplex.com/6621602107228160

mock_tensor_validation Task Execution Log:
https://paste.googleplex.com/5138689889009664

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!

[entrpn]

is this report saved locally or does it need to be uploaded to gcs?

[olufiyin19]

I've updated the report to get saved locally only if no gcs bucket is provided in the input json but to gcs otherwise