Skip to content

Qwen3-Omni Checkpoint Mapping #3096

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Open
entrpn wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Qwen3-Omni Checkpoint Mapping #3096

entrpn wants to merge 1 commit into from
+364 −10

Conversation

@entrpn
Copy link
Collaborator

@entrpn entrpn commented Feb 5, 2026 •
edited
Loading

Original author @eitanporat in #2728

Description

Checkpoint conversion mapping from hf

Tests

  • Checked that the weights are correct by comparing forward passes of the model.
  • Converted HF weights to orbax checkpoint.

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
Copy link

codecov bot commented Feb 5, 2026 •
edited
Loading

Codecov Report

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

Files with missing lines Patch % Lines
...xText/utils/ckpt_conversion/utils/param_mapping.py 0.00% 185 Missing ⚠️
src/MaxText/layers/encoders.py 0.00% 7 Missing ⚠️
src/MaxText/layers/models.py 0.00% 2 Missing ⚠️
src/MaxText/utils/ckpt_conversion/utils/utils.py 0.00% 1 Missing ⚠️

📢 Thoughts on this report? Let us know!

hengtaoguo
hengtaoguo reviewed Feb 9, 2026
View reviewed changes
Copy link
Collaborator

@hengtaoguo hengtaoguo left a comment

Choose a reason for hiding this comment

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

I wonder if you have tried the full checkpoint conversion? Are all layers filled with weights? Not sure if its ready to decode now, seems we are pending two more PRs.

src/MaxText/layers/models.py
if self.config.use_multimodal and encoder_images is not None:
image_embeddings = self.vision_encoder(input_images=encoder_images, deterministic=not enable_dropout)
# qwen3-omni-30b-a3b returns deep features from the vision encoder.
image_embeddings, _ = self.vision_encoder(input_images=encoder_images, deterministic=not enable_dropout)
Copy link
Collaborator

@hengtaoguo hengtaoguo Feb 9, 2026

Choose a reason for hiding this comment

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

Since you changed the output format of encoder, could you verify if all other references have been updated too?

Copy link
Collaborator

@aireenmei aireenmei Feb 9, 2026

Choose a reason for hiding this comment

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

Should we update return var number where vision_encoder is called, such as https://github.com/AI-Hypercomputer/maxtext/blob/main/src/MaxText/layers/models.py#L462

Copy link
Collaborator

@aireenmei aireenmei Feb 9, 2026

Choose a reason for hiding this comment

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

Also the returned deep_feats should not be _ ?

Copy link
Collaborator Author

@entrpn entrpn Feb 10, 2026

Choose a reason for hiding this comment

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

@aireenmei it cannot be used right now, it requires the next PR in our list: #2729 thus I'm conforming to the order.

entrpn
entrpn commented Feb 9, 2026
View reviewed changes
src/MaxText/layers/encoders.py
embeddings = encoder(input_images, deterministic=deterministic)
encoder_output = encoder(input_images, deterministic=deterministic)

deep_feats = None
Copy link
Collaborator Author

@entrpn entrpn Feb 9, 2026

Choose a reason for hiding this comment

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

@hengtaoguo instead of modifying all references, this is how its done to return a encoder_output and deep_feats which will be None for the other models.

@entrpn @eitanporat
qwen3 omni checkpoint conversion
77b061e 
Co-authored-by: Eitan Porat <eporat@lightricks.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@aireenmei aireenmei aireenmei left review comments

@hengtaoguo hengtaoguo hengtaoguo approved these changes

@parambole parambole Awaiting requested review from parambole parambole is a code owner

@shuningjin shuningjin Awaiting requested review from shuningjin shuningjin is a code owner

@RissyRan RissyRan Awaiting requested review from RissyRan RissyRan is a code owner

@suexu1025 suexu1025 Awaiting requested review from suexu1025 suexu1025 is a code owner

@jiangjy1982 jiangjy1982 Awaiting requested review from jiangjy1982 jiangjy1982 is a code owner

@gobbleturk gobbleturk Awaiting requested review from gobbleturk gobbleturk is a code owner

@bvandermoon bvandermoon Awaiting requested review from bvandermoon bvandermoon 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

@richjames0 richjames0 Awaiting requested review from richjames0 richjames0 is a code owner

@NicoGrande NicoGrande Awaiting requested review from NicoGrande NicoGrande is a code owner

@jesselu-google jesselu-google Awaiting requested review from jesselu-google jesselu-google is a code owner

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

Assignees

No one assigned

Labels

gemini-review

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@entrpn @aireenmei @hengtaoguo