diff --git a/.gitignore b/.gitignore index 8a627a7e76..a467b895a0 100644 --- a/.gitignore +++ b/.gitignore @@ -43,3 +43,4 @@ tensor_dumps/ artifacts/ .DS_Store .claude/ +transformer_engine/pytorch/_build_config.py diff --git a/build_tools/pytorch.py b/build_tools/pytorch.py index 2bb238c522..0fd111d07b 100644 --- a/build_tools/pytorch.py +++ b/build_tools/pytorch.py @@ -111,6 +111,24 @@ def setup_pytorch_extension( if bool(int(os.getenv("NVTE_WITH_CUBLASMP", 0))): cxx_flags.append("-DNVTE_WITH_CUBLASMP") + # Experimental: build the torch_stable layer against the torch stable ABI + # (requires torch >= 2.14). Without the flag the same code compiles against + # the full torch ABI. See csrc/torch_stable.h. + torch_stable_abi = bool(int(os.getenv("NVTE_TORCH_STABLE_ABI", "0"))) + if torch_stable_abi: + cxx_flags.append("-DNVTE_WITH_TORCH_STABLE") + cxx_flags.append("-DTORCH_TARGET_VERSION=0x020e000000000000") + + # Record build configuration for runtime checks (stable-ABI builds link + # shims that only exist in libtorch >= TORCH_TARGET_VERSION, so the Python + # package must reject older runtime torch before loading the extension). + build_config = Path(csrc_header_files).parent / "_build_config.py" + build_config.write_text( + '"""Build configuration. Generated by build_tools/pytorch.py, do not edit."""\n' + f"TORCH_STABLE_ABI = {torch_stable_abi}\n" + "TORCH_STABLE_ABI_MIN_TORCH = (2, 14)\n" + ) + # Construct PyTorch CUDA extension sources = [str(path) for path in sources] include_dirs = [str(path) for path in include_dirs] diff --git a/transformer_engine/pytorch/__init__.py b/transformer_engine/pytorch/__init__.py index 4c6b7fc67a..15b1069ca7 100644 --- a/transformer_engine/pytorch/__init__.py +++ b/transformer_engine/pytorch/__init__.py @@ -15,6 +15,23 @@ assert torch_version() >= (2, 1), f"Minimum torch version 2.1 required. Found {torch_version()}." +try: + from transformer_engine.pytorch._build_config import ( + TORCH_STABLE_ABI, + TORCH_STABLE_ABI_MIN_TORCH, + ) +except ImportError: + TORCH_STABLE_ABI = False + TORCH_STABLE_ABI_MIN_TORCH = None + +if TORCH_STABLE_ABI: + # Stable-ABI builds link shims that only exist in newer libtorch; loading + # them on an older runtime would fail with a raw dynamic-linker error. + assert torch_version() >= TORCH_STABLE_ABI_MIN_TORCH, ( + "This Transformer Engine build uses the torch stable ABI and requires torch >=" + f" {'.'.join(map(str, TORCH_STABLE_ABI_MIN_TORCH))} at runtime. Found {torch_version()}." + ) + load_framework_extension("torch") from transformer_engine.pytorch import constants from transformer_engine.pytorch.constants import DType diff --git a/transformer_engine/pytorch/csrc/common.cpp b/transformer_engine/pytorch/csrc/common.cpp index d85dcda159..114e11e275 100644 --- a/transformer_engine/pytorch/csrc/common.cpp +++ b/transformer_engine/pytorch/csrc/common.cpp @@ -62,6 +62,21 @@ NVTEShape convertTorchShape(const c10::IntArrayRef torch_shape) { return ret; } +#ifdef NVTE_WITH_TORCH_STABLE +NVTEShape convertTorchShape(const torch::headeronly::IntHeaderOnlyArrayRef torch_shape) { + NVTEShape ret; + ret.ndim = torch_shape.size(); + constexpr int max_dimensions = sizeof(ret.data) / sizeof(size_t); + NVTE_CHECK(ret.ndim < max_dimensions, + "Torch tensor has too many dimensions. Max supported: ", max_dimensions, " and got ", + ret.ndim, "."); + for (size_t i = 0; i < ret.ndim; ++i) { + ret.data[i] = static_cast(torch_shape[i]); + } + return ret; +} +#endif + std::unique_ptr convert_quantizer(py::handle quantizer) { init_extension(); if (quantizer.is_none()) { @@ -165,6 +180,15 @@ transformer_engine::TensorWrapper makeTransformerEngineTensor(at::Tensor tensor) return makeTransformerEngineTensor(tensor.data_ptr(), shape, dtype); } +#ifdef NVTE_WITH_TORCH_STABLE +transformer_engine::TensorWrapper makeTransformerEngineTensor(const torch_stable::Tensor& tensor) { + transformer_engine::DType dtype = GetTransformerEngineDType(tensor.scalar_type()); + const auto sizes = tensor.sizes(); + std::vector shape(sizes.begin(), sizes.end()); + return makeTransformerEngineTensor(tensor.data_ptr(), shape, dtype); +} +#endif + std::tuple, std::vector>, std::vector, size_t, size_t> makeTransformerEngineTensorList(std::vector> at_tensor_lists) { diff --git a/transformer_engine/pytorch/csrc/common.h b/transformer_engine/pytorch/csrc/common.h index aa0e0c87fe..3d33d2ac11 100644 --- a/transformer_engine/pytorch/csrc/common.h +++ b/transformer_engine/pytorch/csrc/common.h @@ -56,6 +56,8 @@ #include "c10/util/ArrayRef.h" #include "common/util/logging.h" #include "extensions/pybind_dtype_caster.h" +#include "extensions/stable_tensor_caster.h" +#include "torch_stable.h" namespace transformer_engine::pytorch { @@ -544,6 +546,10 @@ transformer_engine::TensorWrapper makeTransformerEngineTensor(void* data_ptr, transformer_engine::TensorWrapper makeTransformerEngineTensor(at::Tensor tensor); +#ifdef NVTE_WITH_TORCH_STABLE +transformer_engine::TensorWrapper makeTransformerEngineTensor(const torch_stable::Tensor& tensor); +#endif + std::tuple, std::vector>, std::vector, size_t, size_t> makeTransformerEngineTensorList(std::vector> at_tensor_lists); @@ -581,6 +587,10 @@ size_t ceildiv(size_t numer, size_t denom); NVTEShape convertTorchShape(const c10::IntArrayRef torch_shape); +#ifdef NVTE_WITH_TORCH_STABLE +NVTEShape convertTorchShape(const torch::headeronly::IntHeaderOnlyArrayRef torch_shape); +#endif + std::vector convert_shape_back_from_fp4(const std::vector& shape, bool transpose); // Flatten an N-D shape to 2D: {product(shape[:-1]), shape[-1]}. diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 95f02c64f0..b6c58593b3 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -437,11 +437,11 @@ at::Tensor scaled_aligned_causal_masked_softmax_backward(at::Tensor output_grads * FP8 recipe **************************************************************************************************/ -void compute_amax(const at::Tensor &tensor, at::Tensor &amax); +void compute_amax(const torch_stable::Tensor &tensor, torch_stable::Tensor &amax); -void fused_amax_and_scale_update_after_reduction(const at::Tensor &amax_reduction_buffer, - std::vector amax_histories, - std::vector scales, +void fused_amax_and_scale_update_after_reduction(const torch_stable::Tensor &amax_reduction_buffer, + std::vector amax_histories, + std::vector scales, const std::string &amax_compute_algo, DType fp8_dtype, float margin); diff --git a/transformer_engine/pytorch/csrc/extensions/recipe.cpp b/transformer_engine/pytorch/csrc/extensions/recipe.cpp index 9be288d2f7..158dcbfa70 100644 --- a/transformer_engine/pytorch/csrc/extensions/recipe.cpp +++ b/transformer_engine/pytorch/csrc/extensions/recipe.cpp @@ -4,34 +4,32 @@ * See LICENSE for license information. ************************************************************************/ -#include -#include - #include #include "../extensions.h" +#include "../torch_stable.h" #include "transformer_engine/transformer_engine.h" namespace transformer_engine::pytorch { -void compute_amax(const at::Tensor& tensor, at::Tensor& amax) { - auto input_tensor = tensor.contiguous(); +void compute_amax(const torch_stable::Tensor& tensor, torch_stable::Tensor& amax) { + auto input_tensor = torch_stable::contiguous(tensor); const TensorWrapper& te_input = makeTransformerEngineTensor(input_tensor); - TORCH_CHECK(amax.scalar_type() == at::kFloat, "amax must be a float tensor"); - TORCH_CHECK(amax.numel() == 1, "amax must have exactly one element"); - auto* amax_ptr = amax.data_ptr(); + NVTE_CHECK(amax.scalar_type() == torch_stable::ScalarType::Float, "amax must be a float tensor"); + NVTE_CHECK(amax.numel() == 1, "amax must have exactly one element"); + auto* amax_ptr = static_cast(amax.data_ptr()); TensorWrapper fake_te_output( /*dptr=*/nullptr, te_input.shape(), DType::kFloat32, // It doesn't matter because we only compute amax. amax_ptr); - nvte_compute_amax(te_input.data(), fake_te_output.data(), at::cuda::getCurrentCUDAStream()); + nvte_compute_amax(te_input.data(), fake_te_output.data(), torch_stable::getCurrentCUDAStream()); } -void fused_amax_and_scale_update_after_reduction(const at::Tensor& amax_reduction_buffer, - std::vector amax_histories, - std::vector scales, +void fused_amax_and_scale_update_after_reduction(const torch_stable::Tensor& amax_reduction_buffer, + std::vector amax_histories, + std::vector scales, const std::string& amax_compute_algo, DType fp8_dtype, float margin) { size_t num_tensors = amax_histories.size(); @@ -58,7 +56,7 @@ void fused_amax_and_scale_update_after_reduction(const at::Tensor& amax_reductio makeTransformerEngineTensor(amax_reduction_buffer).data(), std::vector(te_amax_histories.begin(), te_amax_histories.end()), std::vector(te_scales.begin(), te_scales.end()), amax_compute_algo.c_str(), - static_cast(fp8_dtype), margin, at::cuda::getCurrentCUDAStream()); + static_cast(fp8_dtype), margin, torch_stable::getCurrentCUDAStream()); } } // namespace transformer_engine::pytorch diff --git a/transformer_engine/pytorch/csrc/extensions/stable_tensor_caster.h b/transformer_engine/pytorch/csrc/extensions/stable_tensor_caster.h new file mode 100644 index 0000000000..b4f6ff82be --- /dev/null +++ b/transformer_engine/pytorch/csrc/extensions/stable_tensor_caster.h @@ -0,0 +1,56 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#ifndef TRANSFORMER_ENGINE_PYTORCH_CSRC_EXTENSIONS_STABLE_TENSOR_CASTER_H_ +#define TRANSFORMER_ENGINE_PYTORCH_CSRC_EXTENSIONS_STABLE_TENSOR_CASTER_H_ + +/* In the default (non-stable) build torch_stable::Tensor is at::Tensor and + * torch's own pybind caster applies; this caster exists only in stable mode. */ +#ifdef NVTE_WITH_TORCH_STABLE + +#include + +#include "../torch_stable.h" + +namespace pybind11 { +namespace detail { + +/*! @brief Custom type caster for ``torch::stable::Tensor``. + * + * Lets pybind-bound functions take/return ``torch::stable::Tensor`` directly: + * a ``torch.Tensor`` argument is unwrapped into a stable tensor sharing the + * same TensorImpl, and a returned stable tensor is wrapped back into a + * ``torch.Tensor``. + * + * NOTE: As a compile-time specialization this must be visible in every + * translation unit that converts ``torch::stable::Tensor`` (it is pulled in + * via the PyTorch extension's ``common.h``), otherwise different TUs would + * instantiate different casters for the same type (ODR violation). + */ +template <> +struct type_caster { + public: + PYBIND11_TYPE_CASTER(torch::stable::Tensor, const_name("torch.Tensor")); + + bool load(handle src, bool) { + if (!src || !torch::stable::is_tensor_pyobject(src.ptr())) { + return false; + } + value = torch::stable::tensor_from_pyobject(src.ptr()); + return true; + } + + static handle cast(const torch::stable::Tensor &src, return_value_policy, handle) { + return handle(static_cast(torch::stable::tensor_to_pyobject(src))); + } +}; + +} // namespace detail +} // namespace pybind11 + +#endif // NVTE_WITH_TORCH_STABLE + +#endif // TRANSFORMER_ENGINE_PYTORCH_CSRC_EXTENSIONS_STABLE_TENSOR_CASTER_H_ diff --git a/transformer_engine/pytorch/csrc/torch_stable.h b/transformer_engine/pytorch/csrc/torch_stable.h new file mode 100644 index 0000000000..b8319051ec --- /dev/null +++ b/transformer_engine/pytorch/csrc/torch_stable.h @@ -0,0 +1,57 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#ifndef TRANSFORMER_ENGINE_PYTORCH_CSRC_TORCH_STABLE_H_ +#define TRANSFORMER_ENGINE_PYTORCH_CSRC_TORCH_STABLE_H_ + +#include + +#ifdef NVTE_WITH_TORCH_STABLE +#include +#include +#include +#include +#else +#include +#include +#endif + +/* Compatibility layer for the incremental migration to the torch stable ABI. + * Migrated code is written against this surface, which is restricted to what + * torch::stable provides. With NVTE_WITH_TORCH_STABLE it maps to torch::stable + * (requires torch >= 2.14); without it (the default) it maps to the full torch + * ABI, keeping support for older torch versions intact. */ +namespace transformer_engine::pytorch::torch_stable { + +#ifdef NVTE_WITH_TORCH_STABLE +using Tensor = torch::stable::Tensor; +using ScalarType = torch::headeronly::ScalarType; +#else +using Tensor = at::Tensor; +using ScalarType = at::ScalarType; +#endif + +inline Tensor contiguous(const Tensor &tensor) { +#ifdef NVTE_WITH_TORCH_STABLE + return torch::stable::contiguous(tensor); +#else + return tensor.contiguous(); +#endif +} + +inline cudaStream_t getCurrentCUDAStream() { +#ifdef NVTE_WITH_TORCH_STABLE + return static_cast(torch::stable::accelerator::getCurrentStream( + torch::stable::accelerator::getCurrentDeviceIndex()) + .nativeHandle()); +#else + return at::cuda::getCurrentCUDAStream(); +#endif +} + +} // namespace transformer_engine::pytorch::torch_stable + +#endif // TRANSFORMER_ENGINE_PYTORCH_CSRC_TORCH_STABLE_H_