Skip to content

feat: add HF safetensors export for qwen3.5_moe - #2954

Draft
mrinalghoshh wants to merge 4 commits into
keras-team:masterfrom
mrinalghoshh:qwen3.5-moe-kerashub-to-safetensors-checkpoint-conversion
Draft

mrinalghoshh wants to merge 4 commits into
keras-team:masterfrom
mrinalghoshh:qwen3.5-moe-kerashub-to-safetensors-checkpoint-conversion

Conversation

@mrinalghoshh

@mrinalghoshh mrinalghoshh commented Aug 13, 2026

Copy link
Copy Markdown

Description of the change

Reference

Colab Notebook

https://colab.research.google.com/gist/mrinalghoshh/2457e93a4c4ffc437893f3af06db7dbd/qwen3_5_moe_export_verification.ipynb

Checklist

  • I have added all the necessary unit tests for my change.
  • I have verified that my change does not break existing code and works with all backends (TensorFlow, JAX, and PyTorch).
  • My PR is based on the latest changes of the main branch (if unsure, rebase the code).
  • I have followed the Keras Hub Model contribution guidelines in making these changes.
  • I have followed the Keras Hub API design guidelines in making these changes.
  • I have signed the Contributor License Agreement.

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Code Review

This pull request adds support for exporting Qwen 3.5 MoE models to Hugging Face Transformers format, introducing configuration mapping, weight conversion utilities, end-to-end tests, and a verification script. Feedback on the changes points out potential AttributeErrors in the configuration exporter due to missing backbone attributes, and suggests using the mean logits difference instead of the maximum absolute difference for numerical verification to comply with repository guidelines.

Comment on lines +24 to +29
"decoder_sparse_step": backbone.decoder_sparse_step,
"moe_intermediate_size": backbone.moe_intermediate_dim,
"shared_expert_intermediate_size": backbone.intermediate_dim,
"num_experts_per_tok": backbone.top_k,
"num_experts": backbone.num_experts,
"norm_top_k_prob": backbone.norm_top_k_prob,

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

high

The attributes decoder_sparse_step, norm_top_k_prob, and intermediate_dim (accessed as backbone.intermediate_dim) do not exist on Qwen3_5MoeBackbone. Accessing them directly will raise an AttributeError during export.

  • intermediate_dim should be shared_expert_intermediate_size.
  • decoder_sparse_step and norm_top_k_prob should be safely retrieved using getattr with appropriate defaults.
Suggested change
"decoder_sparse_step": backbone.decoder_sparse_step,
"moe_intermediate_size": backbone.moe_intermediate_dim,
"shared_expert_intermediate_size": backbone.intermediate_dim,
"num_experts_per_tok": backbone.top_k,
"num_experts": backbone.num_experts,
"norm_top_k_prob": backbone.norm_top_k_prob,
"decoder_sparse_step": getattr(backbone, "decoder_sparse_step", 1),
"moe_intermediate_size": backbone.moe_intermediate_dim,
"shared_expert_intermediate_size": backbone.shared_expert_intermediate_size,
"num_experts_per_tok": backbone.top_k,
"num_experts": backbone.num_experts,
"norm_top_k_prob": getattr(backbone, "norm_top_k_prob", True),


keras_outputs = keras_model.predict(input_ids.numpy(), verbose=0)

diff = np.max(np.abs(hf_outputs.numpy() - keras_outputs))

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

medium

According to the repository's general rules, when performing numerical verification for converted model checkpoints, we should validate using the mean logits difference tolerance rather than the maximum absolute difference.

Suggested change
diff = np.max(np.abs(hf_outputs.numpy() - keras_outputs))
diff = np.mean(np.abs(hf_outputs.numpy() - keras_outputs))
References
  1. When performing numerical verification for converted model checkpoints, validate using the mean logits difference tolerance rather than the maximum absolute difference.

@github-actions

Copy link
Copy Markdown

❌ Approved issue check failed. This PR stays in draft until it links an approved issue that is assigned to you.

To fix this:

  1. Find or open an issue for this change and ask a maintainer to approve it and assign it to you.
  2. Link it under the "Approved issue link" section of this PR's description, e.g. Fixes #123.

The PR will be marked Ready for review automatically once the check passes.

No issue reference was found in the description.

@mrinalghoshh mrinalghoshh added the kokoro:force-run Runs Tests on GPU label Sep 22, 2026
@kokoro-team kokoro-team removed the kokoro:force-run Runs Tests on GPU label Sep 22, 2026

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants