Skip to content

Conversation

@ashors1
Copy link
Contributor

@ashors1 ashors1 commented Nov 28, 2023

Adds support for NVIDIA's Transformer Engine. TE can be enabled by setting the environment variable ENABLE_TE=1. For more details about running Pax with Transformer Engine, refer to the JAX Toolbox README.

Requires Paxml PR #46.

@zhangqiaorjc
Copy link
Member

Keeping this open for Nvidia to build against. Not merging upstream

mingxu1067 and others added 15 commits August 26, 2024 18:29
Signed-off-by: Reese Wang <rewang@nvidia.com>
Signed-off-by: Reese Wang <rewang@nvidia.com>
Signed-off-by: Reese Wang <rewang@nvidia.com>
Signed-off-by: Reese Wang <rewang@nvidia.com>
Signed-off-by: Reese Wang <rewang@nvidia.com>
Signed-off-by: Reese Wang <rewang@nvidia.com>
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>
Signed-off-by: Ming Huang <mingh@nvidia.com>
Signed-off-by: Reese Wang <rewang@nvidia.com>
Signed-off-by: Reese Wang <rewang@nvidia.com>
Signed-off-by: Reese Wang <rewang@nvidia.com>
Signed-off-by: Ming Huang <mingh@nvidia.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants