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

Pinned build mode for GPU, with prebuilt Transformer Engine bdist #547

Merged
merged 19 commits into from
Apr 12, 2024

Conversation

chajath
Copy link
Collaborator

@chajath chajath commented Mar 25, 2024

Also add constraints to GPU stable build which are cherry-picked from #522 and #525

To update the wheel, take a look at maxtext_transformerengine_builder.Docker. Once the image is built, you can extract the wheel file under /root/TransformerEngine/dist of the image. Unfortunately, there's no one-liner to copy a file from offline docker image, but there's no shortage of workarounds: https://stackoverflow.com/q/25292198

To upload the wheel, builder who is authorized to upload to the artifact registry can do:

pip install twine keyrings.google-artifactregistry-auth wheel
twine upload --repository-url=https://us-python.pkg.dev/gce-ai-infra/maxtext-build-support-packages/ ./transformer_engine-1.4.0+$GITHASH-cp310-cp310-linux_x86_64.whl

Where $GITHASH is the commit hash version of the Transformer Engine

For nightly builds, we don't apply constraints so the image can be built with the latest and greatest packages.

We also copy setup.sh, requirements.txt and constraints.txt first before executing setup.sh, to avoid running redundant setup script when there is no dependency change. We also mount pip cache from host to speed up the build process.

@chajath chajath requested a review from rwitten as a code owner March 25, 2024 16:20
constraints.txt Outdated Show resolved Hide resolved
@chajath chajath changed the title Prebuild and install Transformer Engine package Pinned build mode for GPU, with prebuilt Transformer Engine bdist Apr 1, 2024
@chajath chajath requested a review from gobbleturk as a code owner April 2, 2024 22:39
Also undo some unnecessary change
@chajath
Copy link
Collaborator Author

chajath commented Apr 5, 2024

TE version has been updated in the anticipation of #555 landing soon. Once it lands, I will rebase and make sure to do some sanity check on all builds.

Also update constraints
setup.sh Outdated
Copy link
Collaborator

Choose a reason for hiding this comment

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

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

This is so that pip will know where to find prebuilt wheel, which is currently stored in the referenced artifact registry path

Copy link
Collaborator

Choose a reason for hiding this comment

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

Got it. Is it accessible publicly?

Copy link
Collaborator Author

@chajath chajath Apr 11, 2024

Choose a reason for hiding this comment

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

Yes - any developer will be able to download wheel. Although uploading to the registry will require credential.

@chajath chajath requested review from rwitten and NinaCai April 11, 2024 16:50
Copy link
Collaborator

@yangyuwei yangyuwei left a comment

Choose a reason for hiding this comment

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

Thanks In-Ho! This is really helpful!

@copybara-service copybara-service bot merged commit ebd39aa into main Apr 12, 2024
6 checks passed
@copybara-service copybara-service bot deleted the yiinho-prebuilt-te branch April 12, 2024 22:57
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

4 participants