Skip to content

fix dtype mismatch when using vocab_tiling#3727

Open
aireenmei wants to merge 1 commit intomainfrom
aireen/vocab_dtype
Open

fix dtype mismatch when using vocab_tiling#3727
aireenmei wants to merge 1 commit intomainfrom
aireen/vocab_dtype

Conversation

@aireenmei
Copy link
Copy Markdown
Collaborator

@aireenmei aireenmei commented Apr 22, 2026

Description

Encountered this error when testing gemma4 with bfloat16 and num_vocab_tiling=8.

The root cause is the incompatibility of weight_dtype=bfloat16 and vocab_tiling. In my test, I was using weight_dtype=float32 dtype=bfloat16, but gemma4 has a bug that uses config.dtype (should use config.weight_dtype) to initiate layer scalar, resulting in a bfloat16 weight that causes this error.

This PR fix:

  • add dtype cast to allow using weight_dtype=bfloat16 with vocab_tiling. Although weight_dtype=bfloat16 is not common, I see some test scripts use it, so it's still good to support
  • gemma4 layer scalar initiate with config.weight_dtype

Tests

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

codecov Bot commented Apr 22, 2026

Codecov Report

❌ Patch coverage is 50.00000% with 1 line in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
src/maxtext/models/gemma4.py 0.00% 1 Missing ⚠️

📢 Thoughts on this report? Let us know!

Comment thread src/maxtext/utils/vocabulary_tiling.py Outdated
Copy link
Copy Markdown
Collaborator

@NuojCheng NuojCheng left a comment

Choose a reason for hiding this comment

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

thank you for the fix!

Copy link
Copy Markdown
Collaborator

@gobbleturk gobbleturk left a comment

Choose a reason for hiding this comment

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

Is there a compile test using vocab tiling that should have caught this? Can we add some test protecting this?

@aireenmei aireenmei force-pushed the aireen/vocab_dtype branch from e8d8140 to a912927 Compare April 23, 2026 00:26
@aireenmei
Copy link
Copy Markdown
Collaborator Author

Is there a compile test using vocab tiling that should have caught this? Can we add some test protecting this?

Added and updated PR decription

Copy link
Copy Markdown
Collaborator

@gobbleturk gobbleturk left a comment

Choose a reason for hiding this comment

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

Awesome, thank you for the test!

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

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants