Skip to content

Adapt inspect_checkpoint.py for Automated Ckpt Validation Pipeline (NNX Support)#4567

Open
olufiyin19 wants to merge 1 commit into
mainfrom
feat/legacy-inspect-checkpoint-nnx-support
Open

Adapt inspect_checkpoint.py for Automated Ckpt Validation Pipeline (NNX Support)#4567
olufiyin19 wants to merge 1 commit into
mainfrom
feat/legacy-inspect-checkpoint-nnx-support

Conversation

@olufiyin19

@olufiyin19 olufiyin19 commented Jul 22, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

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 ~/maxtext/src/maxtext/checkpoint_conversion/inspect_checkpoint.py maxtext model_name=llama2-7b scan_layers=False enable_nnx=True

--- Inspecting MaxText Architecture: llama2-7b (scan_layers: False, enable_nnx: True) ---

Total Parameters: 6738415616 (~6.74 B)

=== Structure ===
key: params.params.decoder.decoder_norm.scale | shape: (4096,)
key: params.params.decoder.layers_0.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_0.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_0.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_0.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_0.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_0.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_0.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_0.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_0.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_1.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_1.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_1.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_1.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_1.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_1.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_1.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_1.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_1.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_2.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_2.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_2.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_2.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_2.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_2.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_2.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_2.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_2.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_3.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_3.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_3.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_3.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_3.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_3.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_3.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_3.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_3.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_4.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_4.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_4.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_4.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_4.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_4.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_4.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_4.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_4.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_5.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_5.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_5.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_5.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_5.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_5.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_5.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_5.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_5.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_6.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_6.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_6.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_6.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_6.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_6.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_6.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_6.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_6.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_7.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_7.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_7.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_7.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_7.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_7.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_7.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_7.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_7.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_8.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_8.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_8.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_8.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_8.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_8.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_8.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_8.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_8.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_9.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_9.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_9.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_9.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_9.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_9.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_9.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_9.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_9.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_10.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_10.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_10.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_10.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_10.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_10.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_10.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_10.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_10.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_11.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_11.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_11.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_11.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_11.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_11.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_11.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_11.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_11.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_12.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_12.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_12.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_12.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_12.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_12.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_12.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_12.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_12.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_13.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_13.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_13.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_13.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_13.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_13.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_13.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_13.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_13.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_14.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_14.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_14.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_14.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_14.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_14.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_14.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_14.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_14.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_15.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_15.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_15.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_15.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_15.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_15.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_15.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_15.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_15.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_16.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_16.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_16.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_16.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_16.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_16.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_16.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_16.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_16.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_17.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_17.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_17.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_17.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_17.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_17.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_17.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_17.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_17.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_18.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_18.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_18.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_18.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_18.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_18.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_18.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_18.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_18.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_19.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_19.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_19.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_19.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_19.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_19.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_19.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_19.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_19.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_20.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_20.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_20.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_20.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_20.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_20.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_20.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_20.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_20.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_21.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_21.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_21.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_21.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_21.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_21.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_21.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_21.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_21.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_22.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_22.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_22.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_22.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_22.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_22.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_22.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_22.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_22.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_23.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_23.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_23.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_23.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_23.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_23.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_23.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_23.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_23.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_24.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_24.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_24.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_24.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_24.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_24.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_24.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_24.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_24.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_25.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_25.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_25.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_25.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_25.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_25.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_25.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_25.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_25.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_26.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_26.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_26.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_26.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_26.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_26.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_26.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_26.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_26.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_27.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_27.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_27.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_27.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_27.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_27.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_27.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_27.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_27.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_28.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_28.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_28.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_28.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_28.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_28.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_28.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_28.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_28.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_29.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_29.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_29.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_29.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_29.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_29.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_29.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_29.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_29.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_30.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_30.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_30.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_30.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_30.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_30.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_30.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_30.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_30.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_31.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_31.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_31.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_31.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_31.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_31.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_31.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_31.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_31.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.logits_dense.kernel | shape: (4096, 32000)
key: params.params.token_embedder.embedding | shape: (32000, 4096)
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 ~/maxtext/src/maxtext/checkpoint_conversion/inspect_checkpoint.py maxtext model_name=llama2-7b scan_layers=False enable_nnx=False

--- Inspecting MaxText Architecture: llama2-7b (scan_layers: False, enable_nnx: False) ---

Total Parameters: 6738415616 (~6.74 B)

=== Structure ===
key: params.params.decoder.decoder_norm.scale | shape: (4096,)
key: params.params.decoder.layers_0.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_0.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_0.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_0.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_0.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_0.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_0.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_0.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_0.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_1.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_1.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_1.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_1.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_1.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_1.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_1.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_1.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_1.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_2.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_2.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_2.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_2.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_2.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_2.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_2.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_2.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_2.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_3.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_3.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_3.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_3.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_3.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_3.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_3.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_3.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_3.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_4.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_4.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_4.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_4.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_4.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_4.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_4.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_4.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_4.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_5.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_5.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_5.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_5.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_5.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_5.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_5.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_5.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_5.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_6.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_6.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_6.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_6.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_6.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_6.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_6.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_6.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_6.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_7.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_7.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_7.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_7.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_7.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_7.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_7.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_7.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_7.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_8.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_8.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_8.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_8.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_8.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_8.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_8.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_8.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_8.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_9.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_9.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_9.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_9.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_9.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_9.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_9.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_9.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_9.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_10.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_10.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_10.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_10.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_10.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_10.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_10.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_10.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_10.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_11.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_11.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_11.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_11.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_11.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_11.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_11.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_11.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_11.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_12.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_12.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_12.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_12.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_12.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_12.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_12.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_12.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_12.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_13.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_13.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_13.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_13.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_13.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_13.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_13.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_13.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_13.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_14.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_14.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_14.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_14.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_14.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_14.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_14.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_14.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_14.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_15.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_15.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_15.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_15.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_15.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_15.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_15.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_15.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_15.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_16.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_16.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_16.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_16.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_16.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_16.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_16.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_16.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_16.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_17.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_17.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_17.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_17.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_17.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_17.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_17.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_17.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_17.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_18.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_18.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_18.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_18.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_18.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_18.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_18.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_18.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_18.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_19.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_19.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_19.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_19.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_19.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_19.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_19.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_19.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_19.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_20.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_20.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_20.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_20.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_20.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_20.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_20.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_20.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_20.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_21.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_21.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_21.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_21.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_21.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_21.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_21.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_21.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_21.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_22.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_22.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_22.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_22.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_22.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_22.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_22.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_22.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_22.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_23.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_23.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_23.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_23.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_23.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_23.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_23.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_23.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_23.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_24.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_24.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_24.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_24.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_24.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_24.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_24.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_24.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_24.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_25.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_25.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_25.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_25.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_25.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_25.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_25.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_25.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_25.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_26.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_26.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_26.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_26.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_26.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_26.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_26.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_26.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_26.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_27.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_27.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_27.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_27.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_27.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_27.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_27.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_27.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_27.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_28.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_28.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_28.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_28.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_28.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_28.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_28.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_28.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_28.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_29.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_29.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_29.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_29.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_29.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_29.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_29.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_29.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_29.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_30.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_30.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_30.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_30.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_30.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_30.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_30.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_30.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_30.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_31.mlp.wi_0.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_31.mlp.wi_1.kernel | shape: (4096, 11008)
key: params.params.decoder.layers_31.mlp.wo.kernel | shape: (11008, 4096)
key: params.params.decoder.layers_31.post_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_31.pre_self_attention_layer_norm.scale | shape: (4096,)
key: params.params.decoder.layers_31.self_attention.key.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_31.self_attention.out.kernel | shape: (32, 128, 4096)
key: params.params.decoder.layers_31.self_attention.query.kernel | shape: (4096, 32, 128)
key: params.params.decoder.layers_31.self_attention.value.kernel | shape: (4096, 32, 128)
key: params.params.decoder.logits_dense.kernel | shape: (4096, 32000)
key: params.params.token_embedder.embedding | shape: (32000, 4096)
INFO:absl:Peak Memory: 7.37 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

gemini-code-assist Bot commented Jul 22, 2026

Copy link
Copy Markdown

Caution

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

@olufiyin19
olufiyin19 force-pushed the feat/legacy-inspect-checkpoint-nnx-support branch from 0453b56 to c31d3ae Compare July 22, 2026 18:32
@codecov

codecov Bot commented Jul 22, 2026 •
edited
Loading

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 0% with 16 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
...axtext/checkpoint_conversion/inspect_checkpoint.py 0.00% 16 Missing ⚠️

📢 Thoughts on this report? Let us know!

shuningjin
shuningjin reviewed Jul 22, 2026
View reviewed changes

@shuningjin shuningjin left a comment •
edited
Loading

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

@olufiyin19
olufiyin19 requested a review from parambole as a code owner July 22, 2026 23:36
@olufiyin19
olufiyin19 requested a review from shuningjin July 22, 2026 23:42
shuningjin
shuningjin reviewed Jul 24, 2026
View reviewed changes

@shuningjin shuningjin left a comment •
edited
Loading

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the update! Left some comments.

Do we expect inspect_checkpoint maxtext to have same results for nnx and linen?

Comment thread src/maxtext/experimental/agent/ckpt_validation_pipeline/inspect_checkpoint.py Outdated
Comment thread src/maxtext/checkpoint_conversion/inspect_checkpoint.py Outdated
Comment thread src/maxtext/checkpoint_conversion/inspect_checkpoint.py Outdated
Comment thread src/maxtext/checkpoint_conversion/inspect_checkpoint.py
Comment on lines +271 to +275
key_str = ".".join(key_parts)
if getattr(config, "enable_nnx", False) and not key_str.startswith("params"):
param_key = "params.params." + key_str
else:
param_key = "params." + key_str

@shuningjin shuningjin Jul 24, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The name looks a bit inconsistent from prev test. Maybe rerun?

I expect both linen and nnx to print like

key: params.Tid2EidVar.decoder.layers_0.mlp.MoeBlock_0.tid2eid | shape: (129280, 6)
key: params.Tid2EidVar.decoder.layers_1.mlp.MoeBlock_0.tid2eid | shape: (129280, 6)
key: params.Tid2EidVar.decoder.layers_2.mlp.MoeBlock_0.tid2eid | shape: (129280, 6)
key: params.params.decoder.decoder_norm.scale | shape: (4096,)
key: params.params.decoder.hc_head.hc_base | shape: (4,)
key: params.params.decoder.hc_head.hc_fn | shape: (16384, 4)
key: params.params.decoder.hc_head.hc_scale | shape: (1,)
key: params.params.decoder.layers_0.mhc_attention.mhc_norm.scale | shape: (16384,)
key: params.params.decoder.layers_0.mhc_attention.post_alpha | shape: (16384, 4)

@olufiyin19 olufiyin19 Jul 24, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I used param_key_parts_from_path to fix the underscores so output now resembles layers_0. and updated the prefix logic to explicitly ignore Tid2EidVar so MoE layers no longer get the accidental params.params. double prefix. Both outputs now match exactly and I've updated them in my description

Comment thread src/maxtext/checkpoint_conversion/inspect_checkpoint.py
Comment thread src/maxtext/checkpoint_conversion/inspect_checkpoint.py Outdated
@olufiyin19
olufiyin19 force-pushed the feat/legacy-inspect-checkpoint-nnx-support branch 2 times, most recently from 5c5d7e1 to bf2d3dc Compare July 24, 2026 22:30
@olufiyin19
olufiyin19 force-pushed the feat/legacy-inspect-checkpoint-nnx-support branch from bf2d3dc to fe82cb9 Compare July 24, 2026 23:22
@olufiyin19 olufiyin19 changed the title Adapt inspect_checkpoint.py for Automated Ckpt Validation Pipeline (Agent Integration) Adapt inspect_checkpoint.py for Automated Ckpt Validation Pipeline (NNX Support) Jul 25, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@shuningjin shuningjin shuningjin left review comments
@abhinavclemson abhinavclemson Awaiting requested review from abhinavclemson
@gobbleturk gobbleturk Awaiting requested review from gobbleturk gobbleturk is a code owner
@khatwanimohit khatwanimohit Awaiting requested review from khatwanimohit
@bvandermoon bvandermoon Awaiting requested review from bvandermoon bvandermoon is a code owner
@vipannalla vipannalla Awaiting requested review from vipannalla
@RissyRan RissyRan Awaiting requested review from RissyRan RissyRan is a code owner
@richjames0 richjames0 Awaiting requested review from richjames0 richjames0 is a code owner
@gagika gagika Awaiting requested review from gagika gagika is a code owner
@shralex shralex Awaiting requested review from shralex shralex is a code owner
@SurbhiJainUSC SurbhiJainUSC Awaiting requested review from SurbhiJainUSC
@hengtaoguo hengtaoguo Awaiting requested review from hengtaoguo hengtaoguo is a code owner
@A9isha A9isha Awaiting requested review from A9isha
@aireenmei aireenmei Awaiting requested review from aireenmei aireenmei is a code owner
@NuojCheng NuojCheng Awaiting requested review from NuojCheng
@jiangjy1982 jiangjy1982 Awaiting requested review from jiangjy1982 jiangjy1982 is a code owner
@suexu1025 suexu1025 Awaiting requested review from suexu1025 suexu1025 is a code owner
@dipannita08 dipannita08 Awaiting requested review from dipannita08
@igorts-git igorts-git Awaiting requested review from igorts-git igorts-git is a code owner
@darisoy darisoy Awaiting requested review from darisoy
@xibinliu xibinliu Awaiting requested review from xibinliu
@huytransformer huytransformer Awaiting requested review from huytransformer
@entrpn entrpn Awaiting requested review from entrpn
@Rohan-Bierneni Rohan-Bierneni Awaiting requested review from Rohan-Bierneni
@parambole parambole Awaiting requested review from parambole parambole is a code owner

At least 2 approving reviews are required to merge this pull request.

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@olufiyin19 @shuningjin