[Qwen3.5] Onboard model to Checkpointing Util & Verify Correctness

[Rohan-Bierneni]

Description

This PR onboard the Qwen3.5 model to the checkpointing util. Also we now use the converted hf -> maxtext checkpoint to run forward_pass_logit checker and verify model correctness.

Tests

Results of forward_pass_logit_checker on 2x2 cases (scanned/unscanned) (to_maxtext/to_huggingface):

HF -> Maxtext Tests:

• Scanned:
  • Command: https://paste.googleplex.com/5199078444105728
  • Output: https://paste.googleplex.com/5416843419451392
• Unscanned:
  • Command: https://paste.googleplex.com/4965438162337792
  • Output: https://paste.googleplex.com/4676840485683200

MaxText -> HF Tests:
Note: Took the maxtext checkpoint above and ran to_huggingface then to_maxtext to finally run forward_pass_logit_checker against checkpoint directly from hf repo.

• Scanned:
  • Command:
    • to_huggingface: https://paste.googleplex.com/4880765002317824
    • to_maxtext: https://paste.googleplex.com/4838089334849536
    • forward_pass_logit_checker: https://paste.googleplex.com/5068151130816512
  • Output: https://paste.googleplex.com/6109510277136384
• Unscanned:
  • Command:
    • to_huggingface: https://paste.googleplex.com/4934408399355904
    • to_maxtext: https://paste.googleplex.com/6257034256318464
    • forward_pass_logit_checker: https://paste.googleplex.com/5068151130816512
  • Output: https://paste.googleplex.com/6321468832088064

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 3.46535% with 195 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...xtext/checkpoint_conversion/utils/param_mapping.py	1.57%	125 Missing ⚠️
...rc/maxtext/checkpoint_conversion/utils/hf_shape.py	0.00%	35 Missing ⚠️
src/maxtext/checkpoint_conversion/utils/utils.py	0.00%	23 Missing ⚠️
src/maxtext/checkpoint_conversion/to_maxtext.py	0.00%	9 Missing ⚠️
...xt/checkpoint_conversion/utils/hf_model_configs.py	62.50%	3 Missing ⚠️

📢 Thoughts on this report? Let us know!

[shuningjin]

Thanks for adding the Qwen 3.5 conversion! Great work figuring out the logic to handle this complex structure, along with comprehensive testing.

I see this adds support for composite_hf_key (multiple HF keys to one MT key). For context, we previously only had composite_mt_key (multiple MT keys to one HF key), so it's great to have this functionality. cc @hengtaoguo

To help maintain this going forward, I have two suggestions:

1. Centralize the logic: Currently, the scanned (list of HF key tuples) vs. unscanned (single HF key tuple) logic is handled via an if/else block in the model hooks. Would it be possible to incorporate this directly into the central framework (to_huggingface via utils.utils._process and to_maxtext)? Here is an example of what that might look like. This will help simplify the logic for future models.
2. Add inline documentation: Could we add more comments clarifying the usage of composite_hf_key, as well as composite_mt_key?

Finally, I left a few minor comments. Please also run a quick formatting pass on the files to address some inconsistent indentation.

[hengtaoguo]

Thanks for the great work! Approve to unblock.