fix GemmaBackbone.get_layout_map + test #1669
Merged
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This PR fixes several issues:
Both tests
test_distribution
andtest_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.