Skip to content
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

fix GemmaBackbone.get_layout_map + test #1669

Merged
merged 5 commits into from
Jun 21, 2024
Merged

fix GemmaBackbone.get_layout_map + test #1669

merged 5 commits into from
Jun 21, 2024

Conversation

martin-gorner
Copy link
Contributor

This PR fixes several issues:

  1. The default Gemma layout map was still missing the fixes from issue 19496.
  2. The PR also adds a test for Gemma + get_layout_map + LoRA. This test fails without the get_layout_map fix in 1).
  3. The PR also adds a layout map for ffw_gating_2, which was apparently forgotten (but there was a test for it!). The value chosen has been tested to be the one offering the fastest training on TPU_v3_8.
  4. The PR also fixes a typo in the tests "ffw_linearl" => "ffw_linear"
  5. Finally, the PR changes the number of heads in the test Gemma config from 4 to 8, otherwise the tests won't pass on TPU

Both tests test_distribution and test_distribution_with_lora now pass on TPU_v3_8.
However, I had to manually change device detection from 'CPU' to 'TPU' in file gemma_backbone_test.py, or the distribution tests would be skipped. I suspect these tests are NOT run nightly and I kindly ask a team member to check this.

@martin-gorner
fix to default GemmaBackbone.get_layout_map() to use fixed regexes as…
9aaa915 
… per keras-team/keras#19496 (comment)
@martin-gorner
fix to default GemmaBackbone.get_layout_map() to use fixed regexes as…
2fe48e0 
… per keras-team/keras#19496 (comment)
@martin-gorner
Also fixing forgotten ffw_gating_2 in GemmaBackbone.get_layout_map. T…
12b222b 
…he sharding spec ("batch", "model") is the one that provides the best training performance. ("batch", "model") and (None, None) are slower (the first one by 40%, the second by 2%).

Fixing test too, including typo ffw_linearl => ffw_linear
@github-actions github-actions bot added the Gemma Gemma model specific issues label Jun 19, 2024
@martin-gorner
changed test_architecture_characteristics test to follow the 4->8 hea…
12fc70c 
…ds change necessary for the test to work on TPUs.

Also fixed formatting.
mattdangerw
mattdangerw approved these changes Jun 20, 2024
View reviewed changes
Copy link
Member

@mattdangerw mattdangerw left a comment

Choose a reason for hiding this comment

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

Looks good! Just minor comments.

keras_nlp/src/models/gemma/gemma_backbone_test.py Outdated Show resolved Hide resolved
@mattdangerw mattdangerw added the kokoro:force-run Runs Tests on GPU label Jun 20, 2024
@kokoro-team kokoro-team removed the kokoro:force-run Runs Tests on GPU label Jun 20, 2024
@mattdangerw
Update gemma_backbone_test.py
2001a3d 
Better test messages
@mattdangerw mattdangerw merged commit b58b56e into keras-team:master Jun 21, 2024
5 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@mattdangerw mattdangerw mattdangerw approved these changes

@qlzh727 qlzh727 qlzh727 approved these changes

Assignees
No one assigned
Labels
Gemma Gemma model specific issues
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.

None yet
4 participants
@martin-gorner @mattdangerw @qlzh727 @kokoro-team