[PyTorch] Add torch_stable compatibility layer for incremental migration to the torch stable ABI - #3335
Draft
pggPL wants to merge 7 commits into
Draft
[PyTorch] Add torch_stable compatibility layer for incremental migration to the torch stable ABI#3335pggPL wants to merge 7 commits into
pggPL wants to merge 7 commits into
Conversation
… caster Start of the incremental migration to the torch stable ABI on pybind11: - csrc/torch_compat.h: dual-mode compat layer (stable shims with NVTE_WITH_TORCH_STABLE, full-ABI polyfills otherwise) - pybind11 type_caster<torch::stable::Tensor> built on it - makeTransformerEngineTensor/convertTorchShape overloads for stable tensors - extensions/recipe.cpp ported to torch::stable::Tensor - experimental NVTE_TORCH_STABLE_ABI build flag Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
torch_compat::Tensor is at::Tensor in the default build (torch's own pybind caster applies) and torch::stable::Tensor only under NVTE_WITH_TORCH_STABLE; the stable tensor caster and the stable overloads are compiled only in stable mode. Migrated code targets the torch_compat surface instead of torch::stable directly. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
The namespace is the torch::stable API surface TE migrates to, not a backwards-compat shim; the name states the contract: ported files use torch_stable::, never at::. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Replaces the import-torch + isinstance probe (~200ns/call, Python round trip) with the C shim (~20ns). Requires the torch PR adding torch_is_tensor_pyobject on top of torch 2.14. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Stable-ABI builds link shims that only exist in libtorch >= 2.14; on an older runtime the import would die with a raw dynamic-linker error. The build now records TORCH_STABLE_ABI in a generated _build_config.py and __init__.py fails with a clear message before loading the extension. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
for more information, see https://pre-commit.ci
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
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
First step of the incremental migration of
transformer_engine_torchto the torch stable ABI (so that one wheel can work across torch versions).Key observation: pybind11 itself does not depend on the torch ABI (only on CPython) — the torch-ABI dependency comes from
torch/extension.h/at::Tensor/ torch's pybind casters. So the extension can migrate totorch::stableincrementally, file by file, while staying on pybind11.This PR adds the foundation:
csrc/torch_compat.h— a thin compatibility layer. Migrated code is written against thetorch_compatsurface, which is restricted to whattorch::stableprovides. In the default build it maps to the full torch ABI (torch_compat::Tensorisat::Tensor, torch's own pybind caster applies) — no stable headers are touched, so torch >= 2.1 keeps working unchanged. WithNVTE_TORCH_STABLE_ABI=1(experimental, requires torch >= 2.14) it maps totorch::stable(includingtorch::stable::tensor_from_pyobject/tensor_to_pyobjectfrom Adding conversion from PyObject to torch::stable::tensor pytorch/pytorch#183323).csrc/extensions/stable_tensor_caster.h— a pybind11type_caster<torch::stable::Tensor>, compiled only in stable mode; atorch.Tensorargument is unwrapped into a stable tensor sharing the same TensorImpl.makeTransformerEngineTensor+convertTorchShapeoverloads for stable tensors incommon.cpp(stable mode only).csrc/extensions/recipe.cpp(compute_amax,fused_amax_and_scale_update_after_reduction) now targets thetorch_compatsurface. No behavior change in the default build.build_tools/pytorch.py(NVTE_TORCH_STABLE_ABI→-DNVTE_WITH_TORCH_STABLE -DTORCH_TARGET_VERSION=2.14).Follow-ups: port the remaining
csrcfiles in small PRs, growingtorch_compat.has needed (missing stable APIs: RNG/Philox, c10d, some ops — being upstreamed to pytorch separately), then flip the stable mode on for torch ≥ 2.14.Type of change
Changes
torch_compat.hdual-mode compatibility layer (full ABI by default — torch >= 2.1 unaffected; stable ABI opt-in)type_caster<torch::stable::Tensor>(stable mode only)NVTE_TORCH_STABLE_ABIbuild flag (experimental)extensions/recipe.cppto thetorch_compatsurfaceChecklist: