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

Temporary pinning of TE to unblock nightlies #378

Merged
merged 2 commits into from
Nov 17, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 5 additions & 1 deletion .github/container/Dockerfile.pax.amd64
Original file line number Diff line number Diff line change
Expand Up @@ -15,10 +15,14 @@ ARG REPO_PAXML=https://github.com/google/paxml.git
ARG REPO_PRAXIS=https://github.com/google/praxis.git
ARG REF_PAXML=main
ARG REF_PRAXIS=main
ARG REPO_TE=https://github.com/NVIDIA/TransformerEngine.git
# TODO: This is a temporary pinning of TE as the API in TE no longer matches the TE patch
# This should be reverted to main ASAP
ARG REF_TE=7976bd003fcf084dd068069b92a9a79b1743316a
RUN <<"EOF" bash -ex
install-pax.sh --defer --from_paxml ${REPO_PAXML} --from_praxis ${REPO_PRAXIS} --ref_paxml ${REF_PAXML} --ref_praxis ${REF_PRAXIS}
install-flax.sh --defer
install-te.sh --defer
install-te.sh --defer --from ${REPO_TE} --ref ${REF_TE}

if [[ -f /opt/requirements-defer.txt ]]; then
# SKIP_HEAD_INSTALLS avoids having to install jax from Github source so that
Expand Down
4 changes: 3 additions & 1 deletion .github/container/Dockerfile.t5x
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,9 @@ ENV NVTE_FRAMEWORK=jax
ARG REPO_T5X=https://github.com/google-research/t5x.git
ARG REF_T5X=main
ARG REPO_TE=https://github.com/NVIDIA/TransformerEngine.git
ARG REF_TE=main
# TODO: This is a temporary pinning of TE as the API in TE no longer matches the TE patch
# This should be reverted to main ASAP
ARG REF_TE=7976bd003fcf084dd068069b92a9a79b1743316a
RUN <<"EOF" bash -ex
install-t5x.sh --defer --from ${REPO_T5X} --ref ${REF_T5X}
install-te.sh --defer --from ${REPO_TE} --ref ${REF_TE}
Expand Down
4 changes: 3 additions & 1 deletion .github/workflows/_build_t5x.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,9 @@ on:
type: string
description: Git commit, tag, or branch for TE
required: false
default: main
# TODO: This is a temporary pinning of TE as the API in TE no longer matches the TE patch
# This should be reverted to main ASAP
default: 7976bd003fcf084dd068069b92a9a79b1743316a
outputs:
DOCKER_TAGS:
description: "Tags of the image built"
Expand Down
8 changes: 6 additions & 2 deletions .github/workflows/ci.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,9 @@ on:
description: 'TE source: <repo>#<branch|tag|commit>'
type: string
required: true
default: 'https://github.com/NVIDIA/TransformerEngine.git#main'
# TODO: This is a temporary pinning of TE as the API in TE no longer matches the TE patch
# This should be reverted to main ASAP
default: 'https://github.com/NVIDIA/TransformerEngine.git#7976bd003fcf084dd068069b92a9a79b1743316a'
SRC_T5X:
description: 'T5X source: <repo>#<branch|tag|commit>'
type: string
Expand Down Expand Up @@ -94,7 +96,9 @@ jobs:
# default values are for `pull_request`` event types
parse_git_src JAX "${{ inputs.SRC_JAX }}" "https://github.com/google/jax.git#main"
parse_git_src XLA "${{ inputs.SRC_XLA }}" "https://github.com/openxla/xla.git#main"
parse_git_src TE "${{ inputs.SRC_TE }}" "https://github.com/NVIDIA/TransformerEngine.git#main"
# TODO: This is a temporary pinning of TE as the API in TE no longer matches the TE patch
# This should be reverted to main ASAP
parse_git_src TE "${{ inputs.SRC_TE }}" "https://github.com/NVIDIA/TransformerEngine.git#7976bd003fcf084dd068069b92a9a79b1743316a"
parse_git_src T5X "${{ inputs.SRC_T5X }}" "https://github.com/google-research/t5x.git#main"
parse_git_src PAXML "${{ inputs.SRC_PAXML }}" "https://github.com/google/paxml.git#main"
parse_git_src PRAXIS "${{ inputs.SRC_PRAXIS }}" "https://github.com/google/praxis.git#main"
Expand Down