Skip to content

[PyTorch] Add torch_stable compatibility layer for incremental migration to the torch stable ABI - #3335

Draft
pggPL wants to merge 7 commits into
NVIDIA:mainfrom
pggPL:te_stable_tensor_caster
Draft

[PyTorch] Add torch_stable compatibility layer for incremental migration to the torch stable ABI#3335
pggPL wants to merge 7 commits into
NVIDIA:mainfrom
pggPL:te_stable_tensor_caster

Conversation

@pggPL

@pggPL pggPL commented Aug 10, 2026

Copy link
Copy Markdown
Collaborator

Description

First step of the incremental migration of transformer_engine_torch to 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 to torch::stable incrementally, 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 the torch_compat surface, which is restricted to what torch::stable provides. In the default build it maps to the full torch ABI (torch_compat::Tensor is at::Tensor, torch's own pybind caster applies) — no stable headers are touched, so torch >= 2.1 keeps working unchanged. With NVTE_TORCH_STABLE_ABI=1 (experimental, requires torch >= 2.14) it maps to torch::stable (including torch::stable::tensor_from_pyobject / tensor_to_pyobject from Adding conversion from PyObject to torch::stable::tensor pytorch/pytorch#183323).
  • csrc/extensions/stable_tensor_caster.h — a pybind11 type_caster<torch::stable::Tensor>, compiled only in stable mode; a torch.Tensor argument is unwrapped into a stable tensor sharing the same TensorImpl.
  • makeTransformerEngineTensor + convertTorchShape overloads for stable tensors in common.cpp (stable mode only).
  • First ported file: csrc/extensions/recipe.cpp (compute_amax, fused_amax_and_scale_update_after_reduction) now targets the torch_compat surface. No behavior change in the default build.
  • Build flag plumbing in build_tools/pytorch.py (NVTE_TORCH_STABLE_ABI-DNVTE_WITH_TORCH_STABLE -DTORCH_TARGET_VERSION=2.14).

Follow-ups: port the remaining csrc files in small PRs, growing torch_compat.h as 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

  • Documentation change (change only to the documentation, either a fix or a new content)
  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • Infra/Build change
  • Code refactoring

Changes

  • Add torch_compat.h dual-mode compatibility layer (full ABI by default — torch >= 2.1 unaffected; stable ABI opt-in)
  • Add pybind11 type_caster<torch::stable::Tensor> (stable mode only)
  • Add NVTE_TORCH_STABLE_ABI build flag (experimental)
  • Port extensions/recipe.cpp to the torch_compat surface

Checklist:

  • I have read and followed the contributing guidelines
  • The functionality is complete
  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes

pggPL and others added 7 commits August 10, 2026 11:00
… 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>
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.

1 participant