add Fuji v3 405b and solve HBM OOMs for larger models #766
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
|
@kelvin-zou @hanzhi713 would appreciate your review to make sure this PR roughly matches 405B. Thank you! |
|
Getting this error: am I doing something wrong or is there a missing tokenizer? |
|
Fixed the issue after vocab model was uploaded. Now I'm hitting OOM issues. Here is the model config: |
ruomingp
reviewed
Oct 23, 2024
kelvin-zou
reviewed
Oct 23, 2024
axlearn/experiments/text/gpt/fuji.py
Outdated
| @@ -91,16 +91,19 @@ class Version(enum.Enum): | |||
| "test": 1 * (1024**4), # 1T tokens | |||
| "7B": 1 * (1024**4), # 1T tokens | |||
| "70B": int(1.4 * (1024**4)), # 1.4T tokens | |||
| "405B": int(1.4 * (1024**4)), # 1.4T tokens | |||
There was a problem hiding this comment.
There is no v1/v2 version for 405B right?
There was a problem hiding this comment.
correct but existing fuji code requires it to be defined otherwise threw an error. Let me clean it up though.
axlearn/experiments/text/gpt/fuji.py
Outdated
| hidden_dim=53248, | ||
| num_heads=128, | ||
| # No GQA support in V1 models, so num_kv_heads is the same as num_heads. | ||
| num_kv_heads=None if version == Version.V1 else 8, |
| ), | ||
| learner_kwargs=dict(peak_lr=8e-5, weight_decay=0.1), | ||
| max_sequence_length=max_sequence_length, | ||
| train_batch_size=train_batch_size, |
There was a problem hiding this comment.
Training batch is not 4M tokens for 400B model, it is 16M according to this
|
This PR is still in draft mode. I will be able to update it once I get 405B working on trillium. |
|
The implementation is wrong, hidden_dim should be 16384 and ffn_dim should be 53248 right? I will update this PR once I get my |
samos123
force-pushed
the
llama-3.1-405b
branch
from
a0142c1 to
497f90e
Compare
|
Unable to run golden_config_test: |
samos123
commented
Nov 19, 2024
| ChainConfigModifier.default_config().set( | ||
| config_modifiers=[ | ||
| MeshShapeModifier.default_config().set( | ||
| mesh_shape=mesh_shape_from_axes(data=-1, fsdp=256) |
There was a problem hiding this comment.
note this requires optimizer state weight only offloading PR to be merged: #789
|
@kelvin-zou could you give it another review? I added the TransformerLayer input checkpointing offload to host. This is required in order to run 405B. I did something similar in my full branch: main...samos123:axlearn:trillium-405b-offload |
samos123
changed the title
samos123
commented
Nov 21, 2024
| @@ -1598,7 +1598,7 @@ class Config(BaseLayer.Config): | |||
| @classmethod | |||
| def default_config(cls): | |||
| cfg = super().default_config() | |||
| cfg.param_partition_spec = (None, "model") | |||
| cfg.param_partition_spec = ("fsdp", "model") | |||
There was a problem hiding this comment.
@kelvin-zou unclear if this is safe to change by default for everyone. Please review this specifically. It was needed to fit 405b on HBM memory.
|
this needs to be rebased since shared_lm_head is a new thing in latest main. |
samos123
force-pushed
the
llama-3.1-405b
branch
from
408033d to
7532d9e
Compare
samos123
force-pushed
the
llama-3.1-405b
branch
from
7532d9e to
fe10772
Compare
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
Main things changed as part of this:
reference: https://github.com/AI-Hypercomputer/maxtext/blob/main/MaxText/configs/models/llama3.1-405b.yml
Requires optimizer state weight only to be merged for the fsdp=256 data=-1 config: #789