[PyTorch] Centralize TE↔torch binary boundary in a single facade file - #21
Open
pggPL wants to merge 4 commits into
Open
[PyTorch] Centralize TE↔torch binary boundary in a single facade file#21pggPL wants to merge 4 commits into
pggPL wants to merge 4 commits into
Conversation
Introduce a single facade file that is the only place in the PyTorch extension allowed to include libtorch/ATen/c10 headers or name at::/c10::/ torch:: symbols. It exposes torch types as `using` aliases (Tensor, Device, ScalarType, Stream, ProcessGroup, PhiloxCudaState, ...) and torch ops as free functions (GetATenDType, GetTransformerEngineDType, new_cuda_tensor, getCurrentCUDAStream), with a compile switch (-DTE_WITH_STABLE_ABI) toward torch::stable::Tensor. - torch_backend.h / torch_backend.cpp: the facade (usings + methods). - common.h / common.cpp: migrated to route through the facade (torch-token-free). - qa/L0_pytorch_lint/check_torch_boundary.py + torch_boundary_allowlist.txt: lint guard enforcing that only the facade touches the torch ABI; unmigrated files are grandfathered via an allowlist that only shrinks. - wired the guard into qa/L0_pytorch_lint/test.sh. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Migrate every file under transformer_engine/pytorch/csrc to talk to the
PyTorch ABI only through the torch_backend.h facade: no file other than
torch_backend.{h,cpp} names at::/c10::/c10d::/torch:: symbols or includes
libtorch/ATen/c10 headers anymore.
- Expanded torch_backend.h with the full re-export vocabulary (type aliases,
dtype/device constants, factory/op re-exports via using-declarations with
identical semantics, CUDA helpers, c10d collective option structs, and a
torch::indexing alias).
- Migrated all 34 remaining csrc files (quantizer, type_converters, extensions.h,
pybind.h, and every extensions/*.cpp incl. multi_tensor/*) to the facade
spellings; tensor method calls and non-torch code unchanged.
- Emptied qa/L0_pytorch_lint/torch_boundary_allowlist.txt: nothing is
grandfathered; the boundary guard now passes with 0 pending files.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Follow-up fixes surfaced by building the pytorch extension: - common.h: std::to_string(ArrayRef<T>) lives in namespace std, so the ArrayRef alias must be fully qualified (transformer_engine::pytorch::ArrayRef). - torch_backend.h: drop the bare at::device()/at::dtype() re-exports -- their names collide with the ubiquitous `device`/`dtype` locals/params. Call sites now build options explicitly: TensorOptions().device(...).dtype(...). - Convert all free device()/dtype() option-builder calls accordingly (allocate, cast, comm_gemm_overlap, nvshmem_comm, permutation, quantizer, router). - Qualify facade aliases in the global-namespace CommOverlap* classes in extensions.h (they sit outside transformer_engine::pytorch). - Add `using namespace transformer_engine::pytorch;` to the TUs with global/anonymous-namespace code (attention.cpp mha_fill, comm_gemm_overlap.cpp method defs, pybind.cpp PYBIND11_MODULE) so the facade names resolve there. Verified on the workstation (te_pytorch): full `pip install -e .` compiles with 0 errors; the transformer_engine_torch extension loads (170 symbols) and a te.LayerNorm GPU forward runs correctly. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
GetATenDType, GetTransformerEngineDType and new_cuda_tensor sit on the per-tensor Python<->C++ marshalling path. Define them `inline` in torch_backend.h instead of out-of-line in torch_backend.cpp so they stay inlinable at call sites, matching the pre-facade codegen (the dtype maps were `inline` in common.h before this branch). torch_backend.cpp is now an empty TU. Verified on the workstation: full `pip install -e .` compiles with 0 errors; extension loads (170 symbols) and te.LayerNorm GPU forward runs. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> 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
Centralize the entire binary boundary between Transformer Engine and PyTorch in
the pytorch extension into a single facade file, so that:
at::Tensortotorch::stable::Tensor(LibTorch stable ABI) becomes a one-line compile-flagchange in one place instead of editing ~40 files, and
communication happens through that file.
This is the opposite direction to the dispersed stable-ABI migration: instead of
spreading the stable ABI across every file, it collapses the TE↔torch ABI surface
to one auditable file.
Type of change
Changes
csrc/torch_backend.{h,cpp}— the facade: the only place allowed to includelibtorch/ATen/c10 headers or name
at::/c10::/c10d::/torch::symbols. Itre-exports torch as
usingaliases (Tensor,ScalarType,Device,Stream,ProcessGroup,PhiloxCudaState, ...), dtype/device constants, and factory/opwrappers (via using-declarations with identical semantics), with a compile
switch (
-DTE_WITH_STABLE_ABI) towardtorch::stable::Tensor.pytorch/csrc(quantizer, type_converters, extensions.h, pybind.h,every
extensions/*.cppincl.multi_tensor/*, common.{h,cpp}) migrated to thefacade spellings; tensor method calls unchanged.
qa/L0_pytorch_lint/check_torch_boundary.py+torch_boundary_allowlist.txt— a lint guard (wired into
L0_pytorch_lint/test.sh) that fails if any fileother than the facade touches the torch ABI. The allowlist is empty: nothing is
grandfathered.
Verification
pip install -e .on an RTX-Adate_pytorchcontainer compiles with 0 errors.transformer_engine_torchextension loads (170 symbols) and ate.LayerNormGPU forward runs correctly.
Checklist:
🤖 Generated with Claude Code