Adapt inspect_checkpoint.py for Automated Ckpt Validation Pipeline (Agent Integration)

[olufiyin19]

Description

This PR modifies inspect_checkpoint.py to introduce native NNX architecture support, while preserving Linen compatibility as a fallback. The previous core inspect_checkpoint.py utility only supported Linen-based models. As our automated validation pipeline transitions to an NNX-first architecture, a way to inspect checkpoints against NNX abstract models is required.

• Specific implementation details:

1. Integrated create_nnx_abstract_model and flax.nnx.split to dynamically generate abstract parameters for NNX models when enable_nnx=True is passed in the config.
2. Implemented parameter key restructuring to correctly map NNX trees (prepending params.params.) vs Linen trees (params.).
3. Retained the legacy Linen Transformer initialization logic strictly as a fallback for older models (like Deepseekv4) that have not yet migrated to NNX.

Tests

1. Manually verified that passing `enable_nnx=True` successfully initializes the NNX abstract model and accurately resolves the params.params. key structure.

Command:

python3 src/maxtext/experimental/agent/ckpt_validation_pipeline/inspect_checkpoint.py maxtext base_config=src/maxtext/configs/models/deepseek4-tiny.yml run_name=test load_parameters_path=gs://dummy_path enable_nnx=True 

Total Parameters: 1104218112 (~1.10 B)

=== Structure ===
key: params.params.decoder.decoder_norm.scale | shape: (2048,)
key: params.params.decoder.layers.mlp.wi_0.kernel | shape: (2048, 16, 7168)
key: params.params.decoder.layers.mlp.wi_1.kernel | shape: (2048, 16, 7168)
key: params.params.decoder.layers.mlp.wo.kernel | shape: (7168, 16, 2048)
key: params.params.decoder.layers.post_self_attention_layer_norm.scale | shape: (2048, 16)
key: params.params.decoder.layers.pre_self_attention_layer_norm.scale | shape: (2048, 16)
key: params.params.decoder.layers.self_attention.key.kernel | shape: (2048, 16, 16, 128)
key: params.params.decoder.layers.self_attention.out.kernel | shape: (16, 16, 128, 2048)
key: params.params.decoder.layers.self_attention.query.kernel | shape: (2048, 16, 16, 128)
key: params.params.decoder.layers.self_attention.value.kernel | shape: (2048, 16, 16, 128)
key: params.params.decoder.logits_dense.kernel | shape: (2048, 32000)
key: params.params.token_embedder.embedding | shape: (32000, 2048)
INFO:absl:Peak Memory: 7.34 GB

2. Verified that falling back to older configs successfully routes to the legacy Linen Transformer initialization without failure.

Command:

python3 src/maxtext/experimental/agent/ckpt_validation_pipeline/inspect_checkpoint.py maxtext base_config=src/maxtext/configs/models/deepseek4-tiny.yml run_name=test load_parameters_path=gs://dummy_path enable_nnx=False

Total Parameters: 1104218112 (~1.10 B)

=== Structure ===
key: params.decoder.decoder_norm.scale | shape: (2048,)
key: params.decoder.layers.mlp.wi_0.kernel | shape: (2048, 16, 7168)
key: params.decoder.layers.mlp.wi_1.kernel | shape: (2048, 16, 7168)
key: params.decoder.layers.mlp.wo.kernel | shape: (7168, 16, 2048)
key: params.decoder.layers.post_self_attention_layer_norm.scale | shape: (2048, 16)
key: params.decoder.layers.pre_self_attention_layer_norm.scale | shape: (2048, 16)
key: params.decoder.layers.self_attention.key.kernel | shape: (2048, 16, 16, 128)
key: params.decoder.layers.self_attention.out.kernel | shape: (16, 16, 128, 2048)
key: params.decoder.layers.self_attention.query.kernel | shape: (2048, 16, 16, 128)
key: params.decoder.layers.self_attention.value.kernel | shape: (2048, 16, 16, 128)
key: params.decoder.logits_dense.kernel | shape: (2048, 32000)
key: params.token_embedder.embedding | shape: (32000, 2048)
INFO:absl:Peak Memory: 7.36 GB

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.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[shuningjin]

Thanks for adding the support for inspecting NNX checkpoint! By comparing the code with checkpoint_conversion/inspect_checkpoint.py, it seems the change is relatively safe (diff). Would you consider directly incorporating to that script for broad use?