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.

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

V6e support #912

Draft
wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

V6e support #912

wants to merge 1 commit into from

Conversation

kelvin-zou
Copy link
Contributor

Add v6e support, including a few changes.

  1. Add a specific meshrule for v6e 1k and 2k.
  2. Fix the global tokens/batch to 16m for llama-v3 model.
  3. Also refactored and moved _save_and_offload_only_these_names_regex into util class for external access.
  4. Tuned the compiler_option a bit for better performance.
  5. Add an additional flag (megascale_grpc_enable_xor_tracer=False) to workaround the megascale OOM issue.

samos123
samos123 reviewed Jan 8, 2025
View reviewed changes
axlearn/common/compiler_options.py
xla_sc_disjoint_spmem="false",
xla_tpu_enable_sparse_core_collective_offload_all_reduce="true",
# TODO(kelvinzou): temporary workaround to avoid memory leak in megascale.
megascale_grpc_enable_xor_tracer="false",
Copy link
Contributor

Choose a reason for hiding this comment

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

current plan is to release a fix in jax 0.4.39 which is planned for Jan 15. The fix is in libtpu.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@samos123 samos123 samos123 left review comments

@ruomingp ruomingp Awaiting requested review from ruomingp ruomingp will be requested when the pull request is marked ready for review ruomingp is a code owner

@markblee markblee Awaiting requested review from markblee markblee will be requested when the pull request is marked ready for review markblee is a code owner

At least 1 approving review is required to merge this pull request.

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
2 participants
@kelvin-zou @samos123