[CI] Add workflow verifying build against minimum supported PyTorch and JAX - #3396
Draft
pggPL wants to merge 15 commits into
Draft
[CI] Add workflow verifying build against minimum supported PyTorch and JAX#3396pggPL wants to merge 15 commits into
pggPL wants to merge 15 commits into
Conversation
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
… 8.9 Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
…1 needs c10::optional Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
…Make Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
… c10::Stream caster Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
…nd lacks c10::Stream caster" This reverts commit 04b4b0f. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
…torch 2.1 needs c10::optional" This reverts commit 061834f. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Contributor
Greptile SummaryThe PR adds minimum-version PyTorch and JAX build jobs and aligns PyTorch NCCL EP compilation with framework capability checks.
Confidence Score: 5/5The PR appears safe to merge because no blocking failure remains. No blocking failure remains. Important Files Changed
Reviews (4): Last reviewed commit: "Skip NCCL EP in torch extension when tor..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| pip install torch==${MIN_TORCH_VERSION} | ||
| pip install cmake pybind11[global] ninja pydantic importlib-metadata>=1.0 packaging einops onnxscript "nvidia-cudnn-frontend>=1.25.0" | ||
| - name: 'Checkout' | ||
| uses: actions/checkout@v3 |
Contributor
There was a problem hiding this comment.
The new workflow executes actions/checkout@v3 as root before recursively fetching and building the repository. Pinning the action to a full commit would prevent the reviewed CI dependency from changing when its upstream tag moves or is compromised.
Suggested change
| uses: actions/checkout@v3 | |
| uses: actions/checkout@<full-commit-sha> |
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
pggPL
changed the title
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
…ild min-torch job on sm90 Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
pggPL
marked this pull request as draft
August 18, 2026 15:12
… import Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Adds a GitHub Actions workflow that builds TE against pinned minimum framework versions and runs the sanity imports. Versions are set in a single place (
MIN_TORCH_VERSION: 2.8.0,MIN_JAX_VERSION: 0.5.3). This guards the declared lower bounds of the framework requirements, which are currently not exercised by CI.Context — both declared floors are stale:
build_tools/pytorch.pydeclarestorch>=2.1, but current main does not compile against torch 2.1 (std::optionalvsc10::optionalin generator APIs, missingc10::Streampybind caster) and does not import against anything older than the FSDP2_fully_shardinternals (~2.4). 2.8.0 is pinned as a verified-working floor.pyproject.tomldeclaresjax>=0.5.0, buttransformer_engine/jax/sharding.pyimportsjax.sharding.get_abstract_mesh, which only exists since jax 0.5.3. The build compiles against 0.5.0 but the import fails; 0.5.3 is the verified floor.The declared requirements may deserve a separate bump to match.
To keep the jobs small they build a single old arch (
NVTE_CUDA_ARCHS=70) withMAX_JOBS=2(the 7GB public runner OOMs at full parallelism) and sccache backed by the GitHub Actions cache.The arch-70-only build surfaced a real build bug, fixed here as well: the PyTorch extension enabled
NVTE_WITH_NCCL_EPbased only on the env var (default on), while the common library gates NCCL EP on having an arch >= 90 in the target list. A build targeting only archs < 90 compiled the EP extension against a core library with nonvte_ep_*symbols (and on older torch failed earlier on the missing SymmetricMemory header). The extension now uses the samenccl_ep_enabled()gate as the common CMake side; builds that include arch >= 90 are unaffected.Type of change
Changes
.github/workflows/minimum_versions.ymlwith two jobs, onpull_request/workflow_dispatch:torch==2.8.0on a single arch (sm70) + sanity importjax==0.5.3(CPU jaxlib) on a single arch (sm70) + sanity importbuild_tools/pytorch.py: gateNVTE_WITH_NCCL_EP/USE_NCCLdefines onnccl_ep_enabled()(matching the common CMake gate) instead of the raw env varChecklist:
🤖 Generated with Claude Code