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

[NVIDIA] Added transformer engine support and GPU optimizations #1385

Closed
wants to merge 7 commits into from
Closed

[NVIDIA] Added transformer engine support and GPU optimizations #1385

wants to merge 7 commits into from

Conversation

terrykong
Copy link
Contributor

  • Added Transformer Engine + FP8 support
  • Updated T5x and jax version=0.4.11
  • A100 Perf gains!
    • 80% speedup - T5-small
    • 23% speedup - T5-large
    • 18% speedup - T5-xl
    • 40% speedup - T5-xxl
  • H100 support, with gains over A100
    • 2.08x faster - T5-large
    • 2.24x faster - T5-xl

@google-cla
Copy link

google-cla bot commented Aug 26, 2023

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

terrykong and others added 7 commits August 26, 2023 17:07
@terrykong @SahilJain314
@yhtang @mingxu1067 @nouiz @sharathts @xrennvidia @jeng1220 @zlsh80826
Added transformer engine support and GPU optimizations
0fc080d 
Co-authored-by: Sahil Jain <[email protected]>
Co-authored-by: Terry Kong <[email protected]>
Co-authored-by: Yu-Hang Tang <[email protected]>
Co-authored-by: Ming Huang <[email protected]>
Co-authored-by: Frederic Bastien <[email protected]>
Co-authored-by: Sharath Turuvekere Sreenivas <[email protected]>
Co-authored-by: Xiaowei Ren <[email protected]>
Co-authored-by: Ryan Jeng <[email protected]>
Co-authored-by: Reese Wang <[email protected]>
@terrykong
UNINSTALL_TE in fine-tuning scripts now defaults to no-action
336d640
@terrykong
remove use_gda from LegacyCheckpointManager in train.py for fp8
462b9fb
@terrykong
Allow singlenode scripts to tee to stdout for better indication of
ad28880 
training status
@zlsh80826 @terrykong
Explicit specify self_attn_mask_type
d250182
@terrykong
Disables check for packing by the te_helper util since not all dataset
985ff36 
configs use packing (CV/Multimodal)
@terrykong
Corrected T5x large baselines
1fa57af 
Updated T5x-large MNLI and SQUAD baselines
@terrykong terrykong force-pushed the patch/t5x_te_in_contrib_noindent branch from 80ae059 to 1fa57af Compare August 27, 2023 00:08
@jon-chuang
Copy link

Hello, out of curiosity (while I understand it may not be tested), would this in theory be able to support training/fine-tuning for models built on top of t5x like Flan-UL2?

I guess yes, as it is simply a t5x model with specific config?

@terrykong
Copy link
Contributor Author

@jon-chuang Yes, I believe that's correct given my understanding of the followup architectures to T5: UL2/Flan-T5/Flan-UL2. As long as the core model is the same and only the objective/inputs&targets change, those finetunings should also benefit.

@terrykong
Copy link
Contributor Author

Closing in favor of #1391

@terrykong terrykong closed this Sep 15, 2023
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
@terrykong @jon-chuang @zlsh80826