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.

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

updated lm1b example with jit #3302

Merged
merged 1 commit into from Oct 27, 2023
Merged

updated lm1b example with jit #3302

merged 1 commit into from Oct 27, 2023

Conversation

chiamp
Copy link
Collaborator

@chiamp chiamp commented Aug 31, 2023

Updated lm1b example to use jit instead of pmap.
lm1b now uses both data and model parallelism instead of just data parallelism.

examples/lm1b/train.py Outdated Show resolved Hide resolved
examples/lm1b/train.py Outdated Show resolved Hide resolved
examples/lm1b/configs/default.py Show resolved Hide resolved
examples/lm1b/configs/default.py Outdated Show resolved Hide resolved

# One axis for each parallelism type may hold a placeholder (-1)
# value to auto-shard based on available slices and devices.
# By default, product of the DCN axes should equal number of slices
Copy link
Collaborator

Choose a reason for hiding this comment

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

Maybe explain what DCN and ICI are, I suppose not everyone knows this.

examples/lm1b/train.py Show resolved Hide resolved
examples/lm1b/train.py Show resolved Hide resolved
examples/lm1b/train.py Show resolved Hide resolved
examples/lm1b/utils.py Outdated Show resolved Hide resolved
examples/lm1b/utils.py Show resolved Hide resolved
@codecov-commenter
Copy link

codecov-commenter commented Oct 11, 2023

Codecov Report

Merging #3302 (558be01) into main (fd8fd76) will not change coverage.
The diff coverage is n/a.

@@           Coverage Diff           @@
##             main    #3302   +/-   ##
=======================================
  Coverage   83.68%   83.68%           
=======================================
  Files          56       56           
  Lines        6808     6808           
=======================================
  Hits         5697     5697           
  Misses       1111     1111           

@marcvanzee
Copy link
Collaborator

Looks great, thanks a lot!

@copybara-service copybara-service bot merged commit 8db32ae into google:main Oct 27, 2023
19 checks passed
@chiamp chiamp deleted the lm1b branch October 27, 2023 21:08
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.

None yet

3 participants