Skip to content

Add Qwen3 Audio Encoder #2952

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
aireenmei wants to merge 5 commits into
base: main
Choose a base branch
from
Open

Add Qwen3 Audio Encoder #2952

aireenmei wants to merge 5 commits into from

Conversation

@aireenmei
Copy link
Collaborator

@aireenmei aireenmei commented Jan 15, 2026 •
edited
Loading

Description

Original author @eitanporat in #2726

Revise according to the original PR comments:

  • rename flags for consistency and clean up usage
  • make AudioEncoder an NNX module
  • remove SinusoidsPositionEmbedding, replaced with expanded existing PositionalEmbedding module
  • remove Qwen3OmniAudioModel in qwen3.py since we use the AudioEncoder class in decoder.py instead
  • add precision related flags

Tests

tests/check_qwen3_omni_audio_vs_reference.py all pass

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 Jan 15, 2026 •
edited
Loading

Codecov Report

❌ Patch coverage is 19.83240% with 287 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
src/MaxText/multimodal_utils.py 11.66% 155 Missing and 4 partials ⚠️
src/MaxText/layers/qwen3.py 30.61% 68 Missing ⚠️
src/MaxText/layers/encoders.py 19.23% 21 Missing ⚠️
src/MaxText/layers/embeddings.py 18.18% 18 Missing ⚠️
src/MaxText/layers/models.py 23.52% 11 Missing and 2 partials ⚠️
src/MaxText/layers/decoders.py 0.00% 4 Missing and 1 partial ⚠️
src/MaxText/maxtext_utils.py 62.50% 2 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!

@aireenmei aireenmei force-pushed the aireen/qwen-audio branch 2 times, most recently from 580f4a2 to 875e674 Compare January 15, 2026 02:54
@aireenmei aireenmei marked this pull request as ready for review January 15, 2026 19:30
hengtaoguo
hengtaoguo approved these changes Jan 16, 2026
View reviewed changes
src/MaxText/layers/encoders.py
return embeddings


class AudioEncoder(nnx.Module):
Copy link
Collaborator

@hengtaoguo hengtaoguo Jan 15, 2026

Choose a reason for hiding this comment

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

Thanks for implementing it as NNX!

src/MaxText/layers/decoders.py

if audio_embeddings is not None and cfg.use_audio:
if cfg.model_name in ["qwen3-omni-30b-a3b"]:
y = multimodal_utils.merge_mm_embeddings(
Copy link
Collaborator

@hengtaoguo hengtaoguo Jan 15, 2026

Choose a reason for hiding this comment

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

Does it support "video/audio interleaving" at the moment? Or it's in following PRs?

src/MaxText/layers/models.py
image_embeddings=image_embeddings,
image_masks=encoder_image_masks,
audio_embeddings=audio_embeddings,
audio_masks=audio_masks,
Copy link
Collaborator

@hengtaoguo hengtaoguo Jan 16, 2026

Choose a reason for hiding this comment

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

I feel like we can consolidate all multimodal inputs into a class and pass this single variable to the model. But that can be a follow up later.

Copy link
Collaborator

@NicoGrande NicoGrande Jan 16, 2026

Choose a reason for hiding this comment

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

+1 to this!

src/MaxText/layers/qwen3.py
Comment on lines +1958 to +1960
config: Config
proj1: DenseGeneral
proj2: DenseGeneral
Copy link
Collaborator

@hengtaoguo hengtaoguo Jan 16, 2026

Choose a reason for hiding this comment

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

Are these three lines necessary?

src/MaxText/layers/qwen3.py
self,
hidden_states: Array,
deterministic: bool = False,
):
Copy link
Collaborator

@hengtaoguo hengtaoguo Jan 16, 2026

Choose a reason for hiding this comment

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

Can we also add an input/output shape comment here?

tests/check_qwen3_omni_audio_vs_reference.py
Copy link
Collaborator

@hengtaoguo hengtaoguo Jan 16, 2026

Choose a reason for hiding this comment

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

Would you prefer to create this standalone check, or merge it with https://github.com/AI-Hypercomputer/maxtext/blob/main/tests/check_qwen3_embedding_vs_reference.py?

@hengtaoguo
Copy link
Collaborator

hengtaoguo commented Jan 16, 2026

Thank you for the great work!

NicoGrande
NicoGrande approved these changes Jan 16, 2026
View reviewed changes
Copy link
Collaborator

@NicoGrande NicoGrande left a comment

Choose a reason for hiding this comment

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

LGTM

src/MaxText/layers/models.py
image_embeddings=image_embeddings,
image_masks=encoder_image_masks,
audio_embeddings=audio_embeddings,
audio_masks=audio_masks,
Copy link
Collaborator

@NicoGrande NicoGrande Jan 16, 2026

Choose a reason for hiding this comment

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

+1 to this!

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

Reviewers

@hengtaoguo hengtaoguo hengtaoguo approved these changes

@NicoGrande NicoGrande NicoGrande 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

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

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

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

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

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

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

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@aireenmei @hengtaoguo @NicoGrande @eitanporat