Skip to content

Updates in gemma example - #4913

Merged
copybara-service[bot] merged 3 commits into
mainfrom
gemma-example-using-grain
Apr 13, 2026
Merged

Updates in gemma example#4913
copybara-service[bot] merged 3 commits into
mainfrom
gemma-example-using-grain

Conversation

@vfdev-5

@vfdev-5 vfdev-5 commented Aug 27, 2025 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Updates in gemma example:

  • grain dataflow,
  • add config for sharding rules
  • removed TrainingState
  • fixed mixed precision
  • updated sampler logic
  • added instructions for multi-host TPU
  • benchmark configs and results on GPU

Fixes #4740

@vfdev-5 vfdev-5 changed the title Updates in gemma example: grain dataflow, simplified sharding rules Updates in gemma example Oct 1, 2025
@vfdev-5
vfdev-5 force-pushed the gemma-example-using-grain branch from 5ad27b5 to 3beac99 Compare October 17, 2025 08:03
@vfdev-5
vfdev-5 force-pushed the gemma-example-using-grain branch from 3beac99 to aa1683d Compare October 24, 2025 09:28
@vfdev-5
vfdev-5 force-pushed the gemma-example-using-grain branch from aa1683d to e4e8597 Compare November 6, 2025 13:05
@vfdev-5
vfdev-5 force-pushed the gemma-example-using-grain branch 3 times, most recently from cc785f6 to a4a7c90 Compare November 20, 2025 12:54
@vfdev-5
vfdev-5 force-pushed the gemma-example-using-grain branch 3 times, most recently from 79ec970 to 1049e5f Compare November 27, 2025 11:21
@vfdev-5 vfdev-5 mentioned this pull request Jan 19, 2026
Closed
6 tasks
@vfdev-5
vfdev-5 force-pushed the gemma-example-using-grain branch from 1049e5f to b7dda72 Compare January 19, 2026 10:16
@vfdev-5
vfdev-5 force-pushed the gemma-example-using-grain branch from b7dda72 to 1b42250 Compare March 11, 2026 14:40
@vfdev-5
vfdev-5 force-pushed the gemma-example-using-grain branch 3 times, most recently from 9c711dd to 9b2b5dd Compare March 18, 2026 15:01
@vfdev-5
vfdev-5 marked this pull request as ready for review March 18, 2026 15:35
@vfdev-5

vfdev-5 commented Mar 19, 2026

Copy link
Copy Markdown
Collaborator Author

Let's use nnx.set_view and try to use sow under a config

@vfdev-5
vfdev-5 force-pushed the gemma-example-using-grain branch from 9b2b5dd to e22ece9 Compare March 25, 2026 15:32
@vfdev-5
vfdev-5 force-pushed the gemma-example-using-grain branch from e22ece9 to a860938 Compare March 25, 2026 17:36
IvyZX
IvyZX requested changes Mar 25, 2026
View reviewed changes
Comment thread examples/gemma/input_pipeline_grain.py Resolved
@vfdev-5
More Code updates
5260a1e 
Added benchmarks
Added config for gemma3 270m
Added config to benchmark jax and nnx transforms
@vfdev-5
vfdev-5 force-pushed the gemma-example-using-grain branch from a860938 to 5260a1e Compare March 26, 2026 14:57
@vfdev-5
vfdev-5 requested a review from IvyZX March 26, 2026 14:58
@vfdev-5 vfdev-5 mentioned this pull request Apr 13, 2026
Merged
@copybara-service
copybara-service Bot merged commit 9bed435 into main Apr 13, 2026
24 checks passed
@copybara-service
copybara-service Bot deleted the gemma-example-using-grain branch April 13, 2026 16:10
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@cgarciae cgarciae cgarciae approved these changes
@IvyZX IvyZX Awaiting requested review from IvyZX

Assignees

No one assigned

Labels

pull ready

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

Improve Gemma example

3 participants

@vfdev-5 @cgarciae @IvyZX