Skip to content

Don't review Orbax configs: Regular and Emergency #1323

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

Don't review Orbax configs: Regular and Emergency #1323

wants to merge 58 commits into from

Conversation

samos123
Copy link
Contributor

No description provided.

samos123 and others added 30 commits June 20, 2025 11:28
@samos123
add test script
6005b6d
@samos123
Jun Orbax regular checkpointer fixes
40d1ed5
@samos123
Orbax emergency trainer config for Fuji
b9194c1
@samos123
update fuji configs to use regular orbax checkpointer
74f7522
@samos123
support keep_period
8e80931
@samos123
support run for orbax regular checkpointer
10d0576
@samos123
fix for A TypeHandler for "<class 'jax.Array'>" is already registered.
abbe0ac 
https://buganizer.corp.google.com/issues/419599840#comment18
@samos123
pdbs=1 and print every step
1292819
@samos123
checkpoint every 100 steps
073cbfa
@samos123
increase termination Grace Period to 300s
d723cc7
@samos123
termination grace period to 900s
383d0bc
@samos123
Revert "termination grace period to 900s"
90f1e27 
This reverts commit 383d0bc.
@samos123
save the data iterator
871d2db
@samos123
use fuji v3 8b for tiktoken
bd25cff
@samos123
use tokenizers instead of tokenizer
8923ca6
@samos123
disable saving of data iterator
03c3fb6
@samos123
use orbaxem
410134d
@samos123
add needed fix for orbax em
9051f90 
bigger buffers and newer orbax
@samos123
enable orbax debug logging
b10878f
@samos123
sort the to be assigned keys and available process indexes
bc3cbfc
@samos123
sort proc_infos as well
f8000d7
@samos123
Revert "sort proc_infos as well"
e752a12 
This reverts commit f8000d7.
@samos123
Revert "sort the to be assigned keys and available process indexes"
c7bdd39 
This reverts commit bc3cbfc.
@samos123
gemini fix?
13cb8ac
@findmyway
fail fast
ce3d7ad
@samos123
Revert "gemini fix?"
6223a3f 
This reverts commit 13cb8ac.
@samos123
Merge branch 'orbax-fuji-v2' of github.com:samos123/axlearn into orba…
879afcb 
…x-fuji-v2
@samos123
switch orbax with Jun's patch
1fcb62b
@samos123
add large scale config
ad4a498
@samos123
enable BlockingRecreate
3be33cd
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers
No reviews
Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
3 participants
@samos123 @findmyway @dipannita08