From 41329b95049cd0499e0b6f157fdeca6d13cf682d Mon Sep 17 00:00:00 2001 From: Xianduo Li Date: Mon, 14 Sep 2026 16:00:03 +0800 Subject: [PATCH 01/19] fix: patch cuda hardcodes for plugin compatibility (complete debug directory coverage) This commit completes the CUDA plugin compatibility patches across the entire codebase: 1. Debug directory (transformer_engine/debug/): - fake_quant.py: Replace is_cuda checks with te_device_type() - log_fp8_tensor_stats.py: Replace torch.device('cuda') and torch.cuda.get_device_capability() - per_tensor_scaling.py: Replace is_cuda checks with te_device_type() - utils/stats_buffer.py: Replace .cuda() calls with .to(te_device_type()) 2. Additional missing files from cuda_patch_rc2.14: - permutation.py: Replace all is_cuda checks with device.type comparisons - tensor/float8_tensor.py: Replace is_cuda checks with device.type comparisons - tensor/float8_blockwise_tensor.py: Replace is_cuda checks with device.type comparisons - tensor/mxfp8_tensor.py: Replace is_cuda checks with device.type comparisons 3. Core infrastructure: - transformer_engine/__init__.py: Export te_platform and te_device_type from pytorch module - transformer_engine/pytorch/__init__.py: Define te_platform() and te_device_type() functions All runtime CUDA hardcodes have been replaced with plugin-compatible abstractions. Total: 64 files modified (+405/-289 lines) Signed-off-by: Xianduo Li --- transformer_engine/__init__.py | 21 +++++ .../debug/features/fake_quant.py | 5 +- .../debug/features/log_fp8_tensor_stats.py | 7 +- .../debug/features/per_tensor_scaling.py | 5 +- .../debug/features/utils/stats_buffer.py | 5 +- transformer_engine/pytorch/__init__.py | 19 +++++ .../dot_product_attention/backends.py | 11 +-- .../dot_product_attention/context_parallel.py | 55 ++++++------- .../dot_product_attention.py | 5 +- .../dot_product_attention/flex_attention.py | 15 ++-- .../dot_product_attention/softmax.py | 3 +- .../attention/dot_product_attention/utils.py | 27 ++++--- .../pytorch/attention/inference.py | 21 ++--- .../pytorch/attention/multi_head_attention.py | 5 +- transformer_engine/pytorch/attention/rope.py | 5 +- .../pytorch/cpp_extensions/gemm.py | 79 ++++++++++++++----- transformer_engine/pytorch/cpu_offload.py | 31 ++++---- transformer_engine/pytorch/cpu_offload_v1.py | 31 ++++---- .../custom_recipes/quantizer_factories.py | 3 +- .../custom_recipes/quantizer_factory_zoo.py | 3 +- .../reference_current_scaling.py | 3 +- transformer_engine/pytorch/distributed.py | 33 ++++---- .../pytorch/dynamo/tensor_spec.py | 5 +- transformer_engine/pytorch/ep.py | 9 ++- transformer_engine/pytorch/graph.py | 39 ++++----- transformer_engine/pytorch/jit.py | 68 ++++++++++------ transformer_engine/pytorch/module/base.py | 13 +-- .../pytorch/module/grouped_linear.py | 7 +- .../pytorch/module/layernorm_linear.py | 7 +- .../pytorch/module/layernorm_mlp.py | 19 ++--- transformer_engine/pytorch/module/linear.py | 9 ++- transformer_engine/pytorch/ops/_common.py | 3 +- .../pytorch/ops/basic/activation.py | 5 +- .../pytorch/ops/basic/basic_linear.py | 3 +- transformer_engine/pytorch/ops/basic/bias.py | 3 +- .../pytorch/ops/basic/grouped_linear.py | 5 +- .../pytorch/ops/basic/l2normalization.py | 3 +- .../pytorch/ops/basic/swiglu.py | 7 +- .../fused/forward_linear_bias_activation.py | 3 +- .../ops/fused/forward_linear_bias_add.py | 3 +- .../ops/fused/forward_linear_scale_add.py | 3 +- .../pytorch/ops/fused/grouped_mlp.py | 5 +- .../ops/fused/userbuffers_backward_linear.py | 5 +- .../ops/fused/userbuffers_forward_linear.py | 5 +- transformer_engine/pytorch/ops/op.py | 5 +- .../pytorch/optimizers/fused_adam.py | 5 +- .../pytorch/optimizers/newton_schulz.py | 5 +- transformer_engine/pytorch/permutation.py | 35 +++++--- transformer_engine/pytorch/quantization.py | 17 ++-- .../pytorch/quantized_tensor.py | 7 +- transformer_engine/pytorch/router.py | 5 +- .../pytorch/tensor/float8_blockwise_tensor.py | 6 +- .../pytorch/tensor/float8_tensor.py | 3 +- .../pytorch/tensor/grouped_tensor.py | 3 +- .../pytorch/tensor/identity_tensor.py | 3 +- .../pytorch/tensor/mxfp8_tensor.py | 3 +- .../pytorch/tensor/nvfp4_tensor.py | 5 +- .../tensor/storage/grouped_tensor_storage.py | 5 +- .../tensor/storage/mxfp8_tensor_storage.py | 5 +- .../tensor/storage/nvfp4_tensor_storage.py | 5 +- transformer_engine/pytorch/transformer.py | 5 +- transformer_engine/pytorch/triton/mhc.py | 5 +- .../pytorch/triton/permutation.py | 25 +++--- transformer_engine/pytorch/utils.py | 33 ++++---- 64 files changed, 488 insertions(+), 318 deletions(-) diff --git a/transformer_engine/__init__.py b/transformer_engine/__init__.py index 1c35b1262b8..e442b2f8614 100644 --- a/transformer_engine/__init__.py +++ b/transformer_engine/__init__.py @@ -15,6 +15,26 @@ from typing import Optional, Tuple import transformer_engine.common + +# Plugin system: set NVTE_PLUGIN to the plugin module name to enable. +# e.g. NVTE_PLUGIN=transformer_engine_plugin_fl +_nvte_plugin = os.environ.get("NVTE_PLUGIN") +if _nvte_plugin: + try: + from importlib import import_module + + _patches = import_module(f"{_nvte_plugin}.patches") + _patches.apply_patches() + except Exception as e: + import warnings + + warnings.warn( + f"NVTE_PLUGIN={_nvte_plugin} but plugin patch apply failed: {e}", + RuntimeWarning, + stacklevel=1, + ) + + # Minimum NCCL version for the runtime-loaded NCCL EP backend. _NCCL_EP_MIN_VERSION = (2, 30, 4) @@ -83,6 +103,7 @@ def require_nccl_ep() -> None: try: from . import pytorch + from .pytorch import te_platform, te_device_type except ImportError: pass except FileNotFoundError as e: diff --git a/transformer_engine/debug/features/fake_quant.py b/transformer_engine/debug/features/fake_quant.py index 48112abe4fe..570c0aadf73 100644 --- a/transformer_engine/debug/features/fake_quant.py +++ b/transformer_engine/debug/features/fake_quant.py @@ -13,6 +13,7 @@ from nvdlfw_inspect.utils import append_parent_docstring +from transformer_engine import te_device_type from transformer_engine.debug.features.api import TEConfigAPIMapper from transformer_engine.common.recipe import Format from transformer_engine.pytorch import DType @@ -30,7 +31,9 @@ def fake_quantize(tensor: torch.Tensor, fp8_format: str, out=None): torch.float16, torch.bfloat16, ), "[NVTORCH INSPECT ERROR] Unsupported tensor type." - assert tensor.is_cuda, "[NVTORCH INSPECT ERROR] Must be a GPU tensor." + assert ( + tensor.device.type == te_device_type() + ), f"[NVTORCH INSPECT ERROR] Must be a {te_device_type()} tensor." assert fp8_format in { "FP8E4M3", "FP8E5M2", diff --git a/transformer_engine/debug/features/log_fp8_tensor_stats.py b/transformer_engine/debug/features/log_fp8_tensor_stats.py index 96f1b644cfb..3e0f9babee5 100644 --- a/transformer_engine/debug/features/log_fp8_tensor_stats.py +++ b/transformer_engine/debug/features/log_fp8_tensor_stats.py @@ -18,6 +18,7 @@ from transformer_engine.debug.features.utils import get_reduction_params, next_enabled_iter from transformer_engine.pytorch.tensor import Quantizer, QuantizedTensor from transformer_engine.pytorch.tensor.float8_tensor import ( +from transformer_engine import te_device_type, te_platform Float8Quantizer, Float8CurrentScalingQuantizer, ) @@ -57,7 +58,7 @@ def _get_new_quantizer(recipe_name, fp8_dtype): return Float8BlockQuantizer(fp8_dtype=fp8_dtype, rowwise=True, columnwise=True) if recipe_name == "fp8_current_scaling": return Float8CurrentScalingQuantizer( - fp8_dtype=fp8_dtype, device=torch.device("cuda"), rowwise=True, columnwise=True + fp8_dtype=fp8_dtype, device=torch.device(te_device_type()), rowwise=True, columnwise=True ) if recipe_name == "mxfp8": return MXFP8Quantizer(fp8_dtype=fp8_dtype, rowwise=True, columnwise=True) @@ -210,10 +211,10 @@ def check_if_stat_is_supported( if recipe_from_stat == "fp8_delayed_scaling" and stat_without_recipe == "overflows%": return True - if recipe_from_stat in ["fp8_block_scaling"] and torch.cuda.get_device_capability()[0] < 9: + if recipe_from_stat in ["fp8_block_scaling"] and te_platform().get_device_capability()[0] < 9: raise ValueError(f"Stat {stat} needs Hopper or later GPU.") - if recipe_from_stat == "mxfp8" and torch.cuda.get_device_capability()[0] < 10: + if recipe_from_stat == "mxfp8" and te_platform().get_device_capability()[0] < 10: raise ValueError(f"Stat {stat} needs Blackwell or later GPU.") supported_stats = ["underflows%", "scale_inv_min", "scale_inv_max", "mse"] diff --git a/transformer_engine/debug/features/per_tensor_scaling.py b/transformer_engine/debug/features/per_tensor_scaling.py index e0fc44bd08f..f83b2795042 100644 --- a/transformer_engine/debug/features/per_tensor_scaling.py +++ b/transformer_engine/debug/features/per_tensor_scaling.py @@ -11,6 +11,7 @@ import nvdlfw_inspect.api as debug_api from nvdlfw_inspect.registry import Registry, api_method +from transformer_engine import te_device_type from transformer_engine.pytorch import DType from transformer_engine.pytorch.tensor import Quantizer from transformer_engine.pytorch.tensor.float8_tensor import ( @@ -33,7 +34,9 @@ def per_tensor_cast( torch.float16, torch.bfloat16, ), "[NVTORCH INSPECT ERROR] Unsupported tensor type for per tensor current scaling" - assert tensor.is_cuda, "[NVTORCH INSPECT ERROR] Must be a GPU tensor." + assert ( + tensor.device.type == te_device_type() + ), f"[NVTORCH INSPECT ERROR] Must be a {te_device_type()} tensor." assert fp8_dtype in { DType.kFloat8E4M3, DType.kFloat8E5M2, diff --git a/transformer_engine/debug/features/utils/stats_buffer.py b/transformer_engine/debug/features/utils/stats_buffer.py index ca7f22e2de5..51cc5a0c1f2 100644 --- a/transformer_engine/debug/features/utils/stats_buffer.py +++ b/transformer_engine/debug/features/utils/stats_buffer.py @@ -16,6 +16,7 @@ from nvdlfw_inspect.utils import gather_along_first_dim from nvdlfw_inspect.logging import MetricLogger +from transformer_engine import te_device_type from transformer_engine.debug.features.utils.stats_computation import ( STATS, DEPENDENCIES, @@ -41,14 +42,14 @@ def __init__(self, layer_name, tensor_name, stats, reduction_group, reduce_withi for stat in stats: self.stats_to_compute = self.stats_to_compute | DEPENDENCIES[stat] - self._buffer = torch.zeros(len(STATS), dtype=torch.float32).cuda() + self._buffer = torch.zeros(len(STATS), dtype=torch.float32).to(te_device_type()) self._new_buffer = self._buffer.clone() self._tmp_buffer = self._buffer.clone() # in case of data parallelism it is possible that layer will not be run on one node # modified is set to True if node is run # we do not take not run nodes into account - self.modified = torch.tensor([False], dtype=torch.bool).cuda() + self.modified = torch.tensor([False], dtype=torch.bool).to(te_device_type()) self.iteration = None self.skip_reduction = False diff --git a/transformer_engine/pytorch/__init__.py b/transformer_engine/pytorch/__init__.py index 576fb57c5e5..70bafdd936e 100644 --- a/transformer_engine/pytorch/__init__.py +++ b/transformer_engine/pytorch/__init__.py @@ -9,6 +9,7 @@ import functools import torch +import transformer_engine from transformer_engine.common import load_framework_extension from transformer_engine.pytorch.torch_version import torch_version @@ -16,6 +17,24 @@ assert torch_version() >= (2, 1), f"Minimum torch version 2.1 required. Found {torch_version()}." load_framework_extension("torch") + +# Device type and platform globals — live here because torch cannot be imported +# at the top-level transformer_engine package (shared with JAX). +# Use plugin-set values if available, otherwise fall back to CUDA defaults. +transformer_engine.TE_DEVICE_TYPE = getattr(transformer_engine, "TE_DEVICE_TYPE", "cuda") +transformer_engine.TE_PLATFORM = getattr(transformer_engine, "TE_PLATFORM", torch.cuda) + + +def te_device_type() -> str: + return transformer_engine.TE_DEVICE_TYPE + + +def te_platform(): + return transformer_engine.TE_PLATFORM + + +transformer_engine.te_device_type = te_device_type +transformer_engine.te_platform = te_platform from transformer_engine.pytorch import constants from transformer_engine.pytorch.constants import CPLoadBalancingStrategy, DType from transformer_engine.pytorch.module import LayerNormLinear diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 895881fced1..2bd515b6a71 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -17,6 +17,7 @@ import torch import torch.nn.functional as F from transformer_engine.pytorch.utils import ( +from transformer_engine import te_platform get_device_compute_capability, split_tensor_along_dim, ) @@ -106,7 +107,7 @@ except PackageNotFoundError: pass # only print warning if use_flash_attention_2 = True in get_attention_backend else: - if torch.cuda.is_available() and get_device_compute_capability() >= (10, 0): + if te_platform().is_available() and get_device_compute_capability() >= (10, 0): if fa_utils.is_version_supported(fa_utils.version, fa_utils.version_required_blackwell): fa_utils.is_installed = True elif fa_utils.is_version_supported(fa_utils.version, fa_utils.version_required): @@ -127,7 +128,7 @@ # Setup Flash attention utils fa_utils.set_flash_attention_version() elif ( - torch.cuda.is_available() + te_platform().is_available() and get_device_compute_capability() >= (8, 0) and dpa_utils._NVTE_FLASH_ATTN ): @@ -577,7 +578,7 @@ def _forward( output_size[2], output_size[3], dtype=query_layer.dtype, - device=torch.cuda.current_device(), + device=te_platform().current_device(), ) scale = self.softmax_scale @@ -596,10 +597,10 @@ def _forward( # S/dP are forced to use DS quantizers in DPA.init_fp8_metadata; revert them here for true CS emulation if fp8_recipe.float8_current_scaling(): S_quantizer = Float8CurrentScalingQuantizer( - fp8_dtype=S_quantizer.dtype, device="cuda" + fp8_dtype=S_quantizer.dtype, device=te_device_type() ) dP_quantizer = Float8CurrentScalingQuantizer( - fp8_dtype=dP_quantizer.dtype, device="cuda" + fp8_dtype=dP_quantizer.dtype, device=te_device_type() ) # disable swizzle for MXFP8Quantizer for quantizer in [ diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index 510d14ac638..d3e26a6614c 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -9,6 +9,7 @@ import transformer_engine_torch as tex from transformer_engine.pytorch.utils import ( +from transformer_engine import te_platform get_cudnn_version, nvtx_range_pop, nvtx_range_push, @@ -286,7 +287,7 @@ def get_seq_chunk_ids_for_reordering_after_attn(cp_size, device): def _get_thd_partition_cu_seqlens(cu_seqlens_padded, device=None): """Return physical boundaries used by the THD partition CUDA kernels.""" target_device = torch.device(device if device is not None else cu_seqlens_padded.device) - target_dtype = torch.int32 if target_device.type == "cuda" else cu_seqlens_padded.dtype + target_dtype = torch.int32 if target_device.type == te_device_type() else cu_seqlens_padded.dtype return cu_seqlens_padded.to(device=target_device, dtype=target_dtype) @@ -310,7 +311,7 @@ def get_thd_partitioned_indices( cp_size, ) target_device = torch.device(device if device is not None else cu_seqlens_padded.device) - target_dtype = torch.int32 if target_device.type == "cuda" else cu_seqlens_padded.dtype + target_dtype = torch.int32 if target_device.type == te_device_type() else cu_seqlens_padded.dtype chunk_size = total_tokens // cp_size return torch.arange( cp_rank * chunk_size, @@ -320,7 +321,7 @@ def get_thd_partitioned_indices( ) cu_seqlens_padded = _get_thd_partition_cu_seqlens(cu_seqlens_padded, device) assert cu_seqlens_padded.is_cuda, ( - "Per-document THD partitioning requires CUDA cu_seqlens; pass device='cuda' " + "Per-document THD partitioning requires CUDA cu_seqlens; pass device=te_device_type() " "when the source metadata is on CPU." ) if cu_seqlens_padded.dtype != torch.int32: @@ -566,7 +567,7 @@ def flash_attn_a2a_communicate( a2a_outputs[i - 1], a2a_inputs[i - 1], group=cp_group, async_op=True ) if i > 1: - with torch.cuda.stream(cp_stream): + with te_platform().stream(cp_stream): a2a_reqs[i - 2].wait() x = a2a_outputs[i - 2] if qkv_format in ["bshd", "sbhd", "bhsd"]: @@ -639,7 +640,7 @@ def flash_attn_a2a_communicate( # [cp*t, h//cp, d] -> [cp, t, h//cp, d] a2a_inputs[i] = x.view(cp_size, -1, *x.shape[-2:]) if i > 1: - with torch.cuda.stream(cp_stream): + with te_platform().stream(cp_stream): a2a_reqs[i - 2].wait() x = a2a_outputs[i - 2] # [cp, 2, b, s//2, h//cp, d] -> [2, b, s//2, cp, h//cp, d] @@ -671,7 +672,7 @@ def flash_attn_a2a_communicate( # [b, cp, h//cp, 2, s//2, d] -> [b*h, s, d] # [t, cp, h//cp, d] -> [t, h, d] a2a_outputs[i - 2] = x.view(-1, x.shape[-3] * x.shape[-2], x.shape[-1]) - torch.cuda.current_stream().wait_stream(cp_stream) + te_platform().current_stream().wait_stream(cp_stream) return a2a_outputs[0] if len(a2a_inputs) == 1 else a2a_outputs @@ -710,14 +711,14 @@ def flash_attn_a2a_communicate_softmax_offset( # [1, h//cp, 1, 1] -> [1, h, 1, 1] inp = tensor.view(-1) output = torch.empty(cp_size * inp.shape[0], dtype=tensor.dtype, device=device) - with torch.cuda.stream(cp_stream): + with te_platform().stream(cp_stream): torch.distributed.all_gather_into_tensor( output, inp, group=cp_group, async_op=False, ) - torch.cuda.current_stream().wait_stream(cp_stream) + te_platform().current_stream().wait_stream(cp_stream) output = output.view( *tensor.shape[:h_dim], cp_size * tensor.shape[h_dim], *tensor.shape[h_dim + 1 :] ) @@ -1918,9 +1919,9 @@ def forward( attn_biases = [None for _ in range(cp_size)] # create two streams to resolve wave quantization issue of Flash Attn in each step - flash_attn_streams = [torch.cuda.current_stream(), cp_stream] + flash_attn_streams = [te_platform().current_stream(), cp_stream] # synchronize fwd results correction across steps - fwd_results_correction_done = torch.cuda.Event() + fwd_results_correction_done = te_platform().Event() # q, k, v, o: # causal: [b, 2, s//2, h, d] or [2, s//2, b, h, d] @@ -1940,7 +1941,7 @@ def forward( o_format = qkv_format for i in range(cp_size + 1): if i < cp_size: - with torch.cuda.stream(flash_attn_streams[i % 2]): + with te_platform().stream(flash_attn_streams[i % 2]): # wait until KV is received for req in send_recv_reqs[(i + 1) % 2]: req.wait() @@ -2160,7 +2161,7 @@ def forward( if i > 1: flash_attn_streams[(i - 1) % 2].wait_event(fwd_results_correction_done) - with torch.cuda.stream(flash_attn_streams[(i - 1) % 2]): + with te_platform().stream(flash_attn_streams[(i - 1) % 2]): if use_fused_attention: # [b, h, sq, 1] -> [b, h, sq] # [t, h, 1] -> [t, h] @@ -2211,7 +2212,7 @@ def forward( if i < cp_size: flash_attn_streams[(i - 1) % 2].record_event(fwd_results_correction_done) - torch.cuda.current_stream().wait_stream(flash_attn_streams[1]) + te_platform().current_stream().wait_stream(flash_attn_streams[1]) if return_max_logit: torch.distributed.all_reduce( max_logit, op=torch.distributed.ReduceOp.MAX, group=cp_group @@ -3500,7 +3501,7 @@ def forward( v_ag = v_ag.view(-1, *v.shape[1:]) # Non-THD cp_stream inputs are ready after K/V reorder, so wait here to # preserve overlap with output initialization below. - cp_stream.wait_stream(torch.cuda.current_stream()) + cp_stream.wait_stream(te_platform().current_stream()) # Shapes before per-step slicing and FP8 metadata wrapping. # q: [b, 2, s//2, h, d] or [2, s//2, b, h, d] @@ -3514,7 +3515,7 @@ def forward( out_f16 = torch.empty(o_shape, dtype=fwd_nominal_dtype, device=q.device) # create two streams to resolve wave quantization issue of Flash Attn in each step - flash_attn_streams = [torch.cuda.current_stream(), cp_stream] + flash_attn_streams = [te_platform().current_stream(), cp_stream] # prepare per-step tensors local_seq_chunk_ids = ( [rank] @@ -3630,7 +3631,7 @@ def forward( if qkv_format == "thd": # Delay the THD wait so one dependency covers both restored K/V and the # per-step metadata produced above on the current stream. - cp_stream.wait_stream(torch.cuda.current_stream()) + cp_stream.wait_stream(te_platform().current_stream()) for i in range(len(local_seq_chunk_ids) + 1): if i < len(local_seq_chunk_ids): @@ -3639,7 +3640,7 @@ def forward( # do not overlap. FusedAttention keeps the existing per-step overlap. if i > 0 and (use_flash_attn_3 or use_flash_attn_4): flash_attn_streams[i].wait_stream(flash_attn_streams[i - 1]) - with torch.cuda.stream(flash_attn_streams[i]): + with te_platform().stream(flash_attn_streams[i]): new_qkv_layout = qkv_layout qkv_scale_inv_format = None if qkv_format in ["bshd", "sbhd"]: @@ -3844,7 +3845,7 @@ def forward( if return_max_logit and i == 0: max_logit = torch.clone(max_logit_per_step[0]) if i > 0: - with torch.cuda.stream(flash_attn_streams[i - 1]): + with te_platform().stream(flash_attn_streams[i - 1]): if o_format == "bshd": out_f16[:, i - 1].copy_(out_per_step[i - 1]) elif o_format == "sbhd": @@ -3865,10 +3866,10 @@ def forward( # default stream, so without this wait the read can race with # the write. The post-loop wait_stream(cp_stream) is too late. # No-op when flash_attn_streams[i-1] is current_stream(). - torch.cuda.current_stream().wait_stream(flash_attn_streams[i - 1]) + te_platform().current_stream().wait_stream(flash_attn_streams[i - 1]) max_logit = torch.maximum(max_logit, max_logit_per_step[i - 1]) - torch.cuda.current_stream().wait_stream(cp_stream) + te_platform().current_stream().wait_stream(cp_stream) # all reduce max_logit across ranks if return_max_logit: @@ -4118,9 +4119,9 @@ def backward(ctx, dout, *_args): dv_per_step = [None, None] # create two streams to resolve wave quantization issue of Flash Attn in each step - flash_attn_streams = [torch.cuda.current_stream(), ctx.cp_stream] + flash_attn_streams = [te_platform().current_stream(), ctx.cp_stream] # synchronize dkv update across steps - dkv_update_done = torch.cuda.Event() + dkv_update_done = te_platform().Event() # gather k and v along s or t: [s, b, h, d] -> [cp, s, b, h, d] or [t, h, d] -> [cp*t, h, d] k_ag, _ = gather_along_first_dim(k, ctx.cp_group) @@ -4149,7 +4150,7 @@ def backward(ctx, dout, *_args): # flatten: [cp*2, s//2, b, h, d] -> [cp*s, b, h, d] k_ag = k_ag.view(-1, *k.shape[1:]) v_ag = v_ag.view(-1, *v.shape[1:]) - ctx.cp_stream.wait_stream(torch.cuda.current_stream()) + ctx.cp_stream.wait_stream(te_platform().current_stream()) # set up flash_attn_bwd flash_attn_bwd = None @@ -4210,7 +4211,7 @@ def backward(ctx, dout, *_args): # per-step overlap. if i > 0 and (ctx.use_flash_attn_3 or ctx.use_flash_attn_4): flash_attn_streams[i].wait_stream(flash_attn_streams[i - 1]) - with torch.cuda.stream(flash_attn_streams[i]): + with te_platform().stream(flash_attn_streams[i]): if ctx.qkv_format == "thd": # THD passes full Q/dout; per-step cu_seqlens select chunks. q_part = q @@ -4458,7 +4459,7 @@ def backward(ctx, dout, *_args): if i > 0: # dq/dk/dv, dq_per_step/dk_per_step/dv_per_step: ctx.fwd_nominal_dtype - with torch.cuda.stream(flash_attn_streams[i - 1]): + with te_platform().stream(flash_attn_streams[i - 1]): if ctx.qkv_format == "thd": # dQ: copy every sequence's valid token range from this split's dQ. tex.thd_copy_valid_tokens_from_per_split_to_rank_local( @@ -4503,7 +4504,7 @@ def backward(ctx, dout, *_args): if i < len(local_seq_chunk_ids): flash_attn_streams[i - 1].record_event(dkv_update_done) - torch.cuda.current_stream().wait_stream(ctx.cp_stream) + te_platform().current_stream().wait_stream(ctx.cp_stream) if ctx.qkv_format == "thd": # Reorder dK/dV from sequence order back to dual-chunk CP rank order, @@ -5962,7 +5963,7 @@ def build_rank_indices(device): slice_sizes = (cu_seqlens_padded[1:] - cu_seqlens_padded[:-1]) // total_slices def build_rank_indices(device): - if device.type == "cuda": + if device.type == te_device_type(): return get_thd_partitioned_indices( cu_seqlens_padded, seq_len_val, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index 658dab5d88d..8ec209551b3 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -16,6 +16,7 @@ from torch.nn.parameter import Parameter from transformer_engine.common.recipe import ( +from transformer_engine import te_device_type Format, Recipe, DelayedScaling, @@ -962,12 +963,12 @@ def __init__( self.softmax_offset = None if self.softmax_type == "off-by-one": self.softmax_offset = torch.zeros( - self.num_attention_heads // self.tp_size, device="cuda" + self.num_attention_heads // self.tp_size, device=te_device_type() ) if self.softmax_type == "learnable": self.register_parameter( "softmax_offset", - Parameter(torch.zeros(self.num_attention_heads // self.tp_size, device="cuda")), + Parameter(torch.zeros(self.num_attention_heads // self.tp_size, device=te_device_type())), get_rng_state_tracker=get_rng_state_tracker, ) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py index b9593b42d9b..967e983a60a 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py @@ -10,6 +10,7 @@ from typing import Any, Callable, Dict, Optional, Tuple import torch +from transformer_engine import te_platform _cudnn_score_mod_handles: Dict[torch.device, Any] = {} _cudnn_score_mod_graph_cache: Dict[Tuple[Any, ...], Any] = {} @@ -140,10 +141,10 @@ def _score_mod_callback_cache_key(callback: Optional[Callable]) -> Any: def _score_mod_device_key(device: torch.device) -> Tuple[Any, ...]: """Normalize a tensor device for graph cache keys.""" - if device.type == "cuda": + if device.type == te_device_type(): index = device.index if index is None: - index = torch.cuda.current_device() + index = te_platform().current_device() return (device.type, index) return (device.type, device.index) @@ -194,13 +195,13 @@ def _wrapped_score_mod(sdpa_graph, score_tensor): def _get_cudnn_current_stream_handle(cudnn, device: torch.device): """Return a cuDNN handle for device, bound to PyTorch's current stream.""" - if device.type != "cuda": + if device.type != te_device_type(): raise ValueError(f"Flex Attention only supports CUDA tensors, got device {device}.") if device.index is None: - device = torch.device("cuda", torch.cuda.current_device()) + device = torch.device(te_device_type(), te_platform().current_device()) handle = _cudnn_score_mod_handles.get(device) - with torch.cuda.device(device): + with te_platform().device(device): if handle is None: handle = cudnn.create_handle() _cudnn_score_mod_handles[device] = handle @@ -287,8 +288,8 @@ def _execute_cudnn_graph( """Execute a built cuDNN frontend Python graph.""" cudnn = _import_cudnn_frontend() - if device.type == "cuda" and device.index is None: - device = torch.device("cuda", torch.cuda.current_device()) + if device.type == te_device_type() and device.index is None: + device = torch.device(te_device_type(), te_platform().current_device()) workspace = torch.empty( workspace_size, device=device, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/softmax.py b/transformer_engine/pytorch/attention/dot_product_attention/softmax.py index 6e8a14402f0..b55a960481d 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/softmax.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/softmax.py @@ -10,6 +10,7 @@ from torch import nn import transformer_engine_torch as tex from transformer_engine.pytorch.export import is_in_onnx_export_mode +from transformer_engine import te_device_type THREADS_PER_WARP = 32 THREADS_PER_BLOCK = 128 @@ -236,7 +237,7 @@ def _get_default_causal_mask(mask_type: str, sq: int, sk: int) -> torch.Tensor: def _get_mask(): diagonal_offset = sk - sq + 1 if "bottom_right" in mask_type else 1 return torch.triu( - torch.ones(sq, sk, dtype=torch.bool, device="cuda"), diagonal=diagonal_offset + torch.ones(sq, sk, dtype=torch.bool, device=te_device_type()), diagonal=diagonal_offset ) if is_in_onnx_export_mode(): diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 23b89287dfd..9514cf7149e 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -23,6 +23,7 @@ import transformer_engine_torch as tex import transformer_engine as te from transformer_engine.pytorch.cpp_extensions.fused_attn import ( +from transformer_engine import te_platform QKVLayout, QKVFormat, AttnBiasType, @@ -1828,7 +1829,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt use_flash_attention_4 and not FlashAttentionUtils.v4_is_installed and not FlashAttentionUtils.v4_warning_printed - and torch.cuda.current_device() == 0 + and te_platform().current_device() == 0 ): logger.warning( "flash-attn v4 may provide important feature support or performance improvement." @@ -1840,7 +1841,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt use_flash_attention_3 and not FlashAttentionUtils.v3_is_installed and not FlashAttentionUtils.v3_warning_printed - and torch.cuda.current_device() == 0 + and te_platform().current_device() == 0 ): logger.warning( "flash-attn v3 may provide important feature support or performance improvement." @@ -1852,7 +1853,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt use_flash_attention_2 and not FlashAttentionUtils.is_installed and not FlashAttentionUtils.warning_printed - and torch.cuda.current_device() == 0 + and te_platform().current_device() == 0 ): logger.warning( "flash-attn may provide important feature support or performance improvement." @@ -1965,7 +1966,7 @@ def _mask(cu_seqlens: torch.Tensor, max_seqlen: int) -> torch.Tensor: return ( (positions.unsqueeze(0) >= seqlens.unsqueeze(1)) .view(batch_size, 1, 1, max_seqlen) - .to(device="cuda") + .to(device=te_device_type()) ) attention_mask_q = _mask(cu_seqlens_q, max_seqlen_q) @@ -2091,9 +2092,9 @@ def get_full_mask( actual_seqlens_kv = m[:, 0, 0, :].sum(dim=1) # apply SWA mask - mask = torch.arange(max_seqlen_q, dtype=torch.int32, device="cuda").view( + mask = torch.arange(max_seqlen_q, dtype=torch.int32, device=te_device_type()).view( 1, 1, max_seqlen_q, 1 - ) - torch.arange(max_seqlen_kv, dtype=torch.int32, device="cuda").view(1, 1, 1, max_seqlen_kv) + ) - torch.arange(max_seqlen_kv, dtype=torch.int32, device=te_device_type()).view(1, 1, 1, max_seqlen_kv) swa_left = None swa_right = None if attn_mask_type == "causal_bottom_right" or ( @@ -2189,7 +2190,7 @@ def get_alibi( m_hat = torch.pow(m_hat_0, torch.arange(1, 1 + 2 * (num_heads - n), 2)) m = torch.cat([m, m_hat]) - _alibi_cache["_alibi_slopes"] = m.to(dtype=torch.float32, device="cuda") + _alibi_cache["_alibi_slopes"] = m.to(dtype=torch.float32, device=te_device_type()) _alibi_cache["_num_heads"] = num_heads _alibi_cache["_alibi_slopes_require_update"] = False @@ -2202,9 +2203,9 @@ def get_alibi( else: raise ValueError("ALiBi slopes cannot exceed 2 dimensions.") - bias = torch.arange(max_seqlen_q, dtype=torch.int32, device="cuda").view( + bias = torch.arange(max_seqlen_q, dtype=torch.int32, device=te_device_type()).view( 1, 1, max_seqlen_q, 1 - ) - torch.arange(max_seqlen_kv, dtype=torch.int32, device="cuda").view( + ) - torch.arange(max_seqlen_kv, dtype=torch.int32, device=te_device_type()).view( 1, 1, 1, max_seqlen_kv ) if actual_seqlens_q is None and actual_seqlens_kv is None: @@ -2224,7 +2225,7 @@ def get_alibi( _alibi_cache["_max_seqlen_q"], _alibi_cache["_max_seqlen_kv"] = max_seqlen_q, max_seqlen_kv _alibi_cache["_bottom_right_alignment"] = bottom_right_alignment bias_dtype = torch.float32 if bias_dtype is None else bias_dtype - _alibi_cache["_alibi_bias"] = bias.contiguous().to(dtype=bias_dtype, device="cuda") + _alibi_cache["_alibi_bias"] = bias.contiguous().to(dtype=bias_dtype, device=te_device_type()) _alibi_cache["_alibi_bias_require_update"] = False return _alibi_cache["_alibi_slopes"], _alibi_cache["_alibi_bias"] @@ -2239,7 +2240,7 @@ def get_cu_seqlens(mask: torch.Tensor) -> torch.Tensor: mask = mask.squeeze(1).squeeze(1) reduced_mask = mask.logical_not().sum(dim=1) cu_seqlens = reduced_mask.cumsum(dim=0).to(torch.int32) - zero = torch.zeros(1, dtype=torch.int32, device="cuda") + zero = torch.zeros(1, dtype=torch.int32, device=te_device_type()) cu_seqlens = torch.cat((zero, cu_seqlens)) return cu_seqlens @@ -2257,7 +2258,7 @@ def get_cu_seqlens_and_indices(mask: torch.Tensor) -> Tuple[torch.Tensor, torch. reduced_mask = mask.logical_not().sum(dim=1) cu_seqlens = reduced_mask.cumsum(dim=0).to(torch.int32) - zero = torch.zeros(1, dtype=torch.int32, device="cuda") + zero = torch.zeros(1, dtype=torch.int32, device=te_device_type()) cu_seqlens = torch.cat((zero, cu_seqlens)) mask = mask.reshape(-1) @@ -2293,7 +2294,7 @@ def get_indices(max_seqlen: int, cu_seqlens: torch.Tensor) -> torch.Tensor: ordered, bs * max_seqlen, ) - return indices.to(dtype=torch.int64, device="cuda").unsqueeze(1).unsqueeze(1) + return indices.to(dtype=torch.int64, device=te_device_type()).unsqueeze(1).unsqueeze(1) def get_full_cu_seqlens( diff --git a/transformer_engine/pytorch/attention/inference.py b/transformer_engine/pytorch/attention/inference.py index f98d1693a65..b9849806c44 100644 --- a/transformer_engine/pytorch/attention/inference.py +++ b/transformer_engine/pytorch/attention/inference.py @@ -11,6 +11,7 @@ import torch from transformer_engine.pytorch.attention.custom_ops import ( +from transformer_engine import te_platform copy_to_kv_cache, QKV_FORMAT_VALUE, ) @@ -209,12 +210,12 @@ def __init__( self.cu_seqlens_q = torch.zeros( self.max_batch_size + 1, dtype=torch.int32, - device=torch.cuda.current_device(), + device=te_platform().current_device(), ) self.cu_seqlens_kv = torch.zeros( self.max_batch_size + 1, dtype=torch.int32, - device=torch.cuda.current_device(), + device=te_platform().current_device(), ) # This internal buffer holds the running length of each @@ -225,7 +226,7 @@ def __init__( self.pre_step_seqlens = torch.zeros( self.max_batch_size, dtype=torch.int32, - device=torch.cuda.current_device(), + device=te_platform().current_device(), ) def reset(self): @@ -431,14 +432,14 @@ def __init__( self.batch_indices = torch.zeros( self.max_batch_size, dtype=torch.int32, - device=torch.cuda.current_device(), + device=te_platform().current_device(), ) # after re-indexing, batch indices are always [0, ..., b-1] self.batch_indices_post_step = torch.range( 0, self.max_batch_size - 1, dtype=torch.int32, - device=torch.cuda.current_device(), + device=te_platform().current_device(), ) # whether reindexing is needed, i.e. when batch seq_ids have changed self.need_reindex = True @@ -451,7 +452,7 @@ def allocate_memory(self, layer_number): self.num_heads, self.head_dim_k, dtype=self.dtype, - device=torch.cuda.current_device(), + device=te_platform().current_device(), ) v_cache = torch.zeros( self.max_batch_size, @@ -459,7 +460,7 @@ def allocate_memory(self, layer_number): self.num_heads, self.head_dim_v, dtype=self.dtype, - device=torch.cuda.current_device(), + device=te_platform().current_device(), ) # Mutated in place, so CUDA graphs are only captured if the addresses are fixed. torch._dynamo.mark_static_address(k_cache) @@ -631,7 +632,7 @@ def __init__( self.allocated_pages = defaultdict(list) # page table, [batch_size, max_pages_per_seq] self.page_table = torch.zeros( - self.max_batch_size, self.max_pages_per_seq, dtype=torch.int32, device="cuda" + self.max_batch_size, self.max_pages_per_seq, dtype=torch.int32, device=te_device_type() ) def reset(self): @@ -651,7 +652,7 @@ def allocate_memory(self, layer_number): self.num_heads, self.head_dim_k, dtype=self.dtype, - device=torch.cuda.current_device(), + device=te_platform().current_device(), ) v_cache = torch.zeros( self.total_num_pages, @@ -659,7 +660,7 @@ def allocate_memory(self, layer_number): self.num_heads, self.head_dim_v, dtype=self.dtype, - device=torch.cuda.current_device(), + device=te_platform().current_device(), ) # Mutated in place, so CUDA graphs are only captured if the addresses are fixed. torch._dynamo.mark_static_address(k_cache) diff --git a/transformer_engine/pytorch/attention/multi_head_attention.py b/transformer_engine/pytorch/attention/multi_head_attention.py index e07e203ca7e..3fd67038070 100644 --- a/transformer_engine/pytorch/attention/multi_head_attention.py +++ b/transformer_engine/pytorch/attention/multi_head_attention.py @@ -15,6 +15,7 @@ from transformer_engine.pytorch.module import LayerNormLinear, Linear, RMSNorm, LayerNorm from transformer_engine.pytorch.ops.basic.l2normalization import L2Normalization from transformer_engine.pytorch.utils import ( +from transformer_engine import te_device_type SplitAlongDim, divide, get_default_init_method, @@ -141,7 +142,7 @@ class MultiheadAttention(torch.nn.Module): whether to use interleaved rotary position embeddings. bias : bool, default = True if set to ``False``, the transformer layer will not learn any additive biases. - device : Union[torch.device, str], default = "cuda" + device : Union[torch.device, str], default=te_device_type() The device on which the parameters of the model will be allocated. It is the user's responsibility to ensure all parameters are moved to the GPU before running the forward pass. @@ -285,7 +286,7 @@ def __init__( ub_bulk_wgrad: bool = False, bias: bool = True, normalization: str = "LayerNorm", - device: Union[torch.device, str] = "cuda", + device: Union[torch.device, str] = te_device_type(), qkv_format: str = "sbhd", name: str = None, qk_norm_type: Optional[str] = None, diff --git a/transformer_engine/pytorch/attention/rope.py b/transformer_engine/pytorch/attention/rope.py index 77ad57ed8f7..3462793dcf7 100644 --- a/transformer_engine/pytorch/attention/rope.py +++ b/transformer_engine/pytorch/attention/rope.py @@ -10,6 +10,7 @@ import transformer_engine_torch as tex from transformer_engine.pytorch.cpp_extensions.fused_attn import QKVFormat +from transformer_engine import te_platform __all__ = ["RotaryPositionEmbedding", "apply_rotary_pos_emb", "apply_fused_qkv_rotary_pos_emb"] @@ -54,7 +55,7 @@ def __init__( inv_freq = 1.0 / ( self.rotary_base ** ( - torch.arange(0, dim, 2, dtype=torch.float32, device=torch.cuda.current_device()) + torch.arange(0, dim, 2, dtype=torch.float32, device=te_platform().current_device()) / dim ) ) @@ -76,7 +77,7 @@ def forward(self, max_seq_len: int, offset: int = 0): offset: int, default = 0 Fixed offset for frequencies. """ - with torch.autocast(enabled=False, device_type="cuda"): + with torch.autocast(enabled=False, device_type=te_device_type()): seq = ( torch.arange(max_seq_len, device=self.inv_freq.device, dtype=self.inv_freq.dtype) + offset diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index ab51f9e5147..a82e588acb4 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -17,7 +17,9 @@ from ..tensor.float8_tensor import Float8CurrentScalingQuantizer, Float8Quantizer from ..tensor.mxfp8_tensor import MXFP8Quantizer from ..tensor.nvfp4_tensor import NVFP4Quantizer -from ..tensor.storage.float8_blockwise_tensor_storage import Float8BlockwiseQTensorStorage +from ..tensor.storage.float8_blockwise_tensor_storage import ( + Float8BlockwiseQTensorStorage, +) from ..tensor.storage.float8_tensor_storage import Float8TensorStorage from ..tensor.storage.grouped_tensor_storage import GroupedTensorStorage from ..tensor.storage.hybrid_tensor_storage import HybridQuantizedTensorStorage @@ -26,6 +28,7 @@ from ..tensor.utils import is_custom from ..custom_recipes.gemm import custom_gemm from ...debug.pytorch.debug_quantization import DebugQuantizedTensor, DebugQuantizer +from transformer_engine import te_platform __all__ = [ "general_gemm", @@ -40,7 +43,7 @@ def get_cublas_workspace_size_bytes() -> None: """Return 32 MiB if using hopper, 4 MiB for all other architectures.""" - if torch.cuda.get_device_properties(torch.cuda.current_device()).major >= 9: + if torch.cuda.get_device_properties(te_platform().current_device()).major >= 9: # 32 MiB for NVFP4 GEMM, plus additional 1024 B for alignment and misc scales return 32 * 1024 * 1024 + 1024 return 4_194_304 @@ -61,11 +64,15 @@ def get_cublas_workspace(device: int, ub: bool, grouped_gemm: bool) -> torch.Ten _multi_stream_cublas_workspace = [] for _ in range(tex.get_num_cublas_streams()): _multi_stream_cublas_workspace.append( - torch.empty(get_cublas_workspace_size_bytes(), dtype=torch.uint8, device=device) + torch.empty( + get_cublas_workspace_size_bytes(), dtype=torch.uint8, device=device + ) ) return _multi_stream_cublas_workspace - return torch.empty(get_cublas_workspace_size_bytes(), dtype=torch.uint8, device=device) + return torch.empty( + get_cublas_workspace_size_bytes(), dtype=torch.uint8, device=device + ) def validate_gemm_scale(scale: Optional[float], required: bool) -> float: @@ -104,7 +111,10 @@ def _nvfp4_row_scaled_gemm_inputs( B_metadata[b_amax_key] = output_row_scales.new_ones(1) B_metadata["row_scaled_nvfp4"] = False - assert output_row_scales.dtype == torch.float32 and output_col_scales.dtype == torch.float32 + assert ( + output_row_scales.dtype == torch.float32 + and output_col_scales.dtype == torch.float32 + ) return ( NVFP4TensorStorage(**A_metadata), NVFP4TensorStorage(**B_metadata), @@ -151,7 +161,9 @@ def _unwrap_tensor( # Select the direction of a hybrid tensor, then process its sub-storage. if isinstance(tensor, HybridQuantizedTensorStorage): sub_storage = ( - tensor.rowwise_sub_storage if usage == "rowwise" else tensor.columnwise_sub_storage + tensor.rowwise_sub_storage + if usage == "rowwise" + else tensor.columnwise_sub_storage ) return _unwrap_tensor(sub_storage, usage) @@ -234,10 +246,14 @@ def general_gemm( ) if ub is not None: - assert ub_type is not None, "Comm+GEMM overlap requires a valid `comm_type` argument." + assert ( + ub_type is not None + ), "Comm+GEMM overlap requires a valid `comm_type` argument." if ub_type == tex.CommOverlapType.RS: if not (bulk_overlap and not ub.is_fp8_ubuf()): - assert extra_output is not None, "GEMM+RS overlap requires extra output tensor." + assert ( + extra_output is not None + ), "GEMM+RS overlap requires extra output tensor." if out is not None: if not out.is_contiguous(): @@ -266,7 +282,9 @@ def general_gemm( # Use bfloat16 as default bias_dtype bias_dtype = TE_DType[torch.bfloat16 if bias is None else bias.dtype] - if isinstance(A, Float8BlockwiseQTensorStorage) or isinstance(B, Float8BlockwiseQTensorStorage): + if isinstance(A, Float8BlockwiseQTensorStorage) or isinstance( + B, Float8BlockwiseQTensorStorage + ): # FP8 block-scaling requires split accumulator use_split_accumulator = True @@ -304,11 +322,15 @@ def general_gemm( assert ( quantization_params is None ), "Row-scaled NVFP4 GEMM currently does not support output quantization." - assert ub is None, "Row-scaled NVFP4 GEMM currently does not support CommOverlap." + assert ( + ub is None + ), "Row-scaled NVFP4 GEMM currently does not support CommOverlap." assert ( extra_output is None ), "Row-scaled NVFP4 GEMM currently does not support extra output." - assert not bulk_overlap, "Row-scaled NVFP4 GEMM currently does not support bulk overlap." + assert ( + not bulk_overlap + ), "Row-scaled NVFP4 GEMM currently does not support bulk overlap." assert out is None or ( isinstance(out, torch.Tensor) and not is_custom(out) ), "Row-scaled NVFP4 GEMM currently supports only plain torch.Tensor outputs." @@ -320,8 +342,8 @@ def general_gemm( ), "Row-scaled NVFP4 GEMM currently requires NVFP4 B." # Reuse the per-tensor GEMM and apply selected row/column global scales # to the FP32 output. This extends #2931 without a dedicated GEMM kernel. - gemm_A, gemm_B, output_row_scales, output_col_scales = _nvfp4_row_scaled_gemm_inputs( - A, B, transa=transa, transb=transb + gemm_A, gemm_B, output_row_scales, output_col_scales = ( + _nvfp4_row_scaled_gemm_inputs(A, B, transa=transa, transb=transb) ) requested_out, requested_out_dtype = out, out_dtype @@ -340,7 +362,9 @@ def general_gemm( gemm_args[14] = False # accumulate after applying the outer scales gemm_kwargs = dict(kwargs) gemm_kwargs["beta"] = 0.0 - out, bias_grad, gelu_input, extra_output = tex.generic_gemm(*gemm_args, **gemm_kwargs) + out, bias_grad, gelu_input, extra_output = tex.generic_gemm( + *gemm_args, **gemm_kwargs + ) out_2d = out.reshape(-1, out.shape[-1]) assert output_row_scales.numel() in (1, out_2d.shape[0]) @@ -359,7 +383,9 @@ def general_gemm( out_2d.mul_(output_row_scales) out_2d.mul_(output_col_scales) if bias is not None: - assert not grad, "Row-scaled NVFP4 backward does not support fused bias gradient." + assert ( + not grad + ), "Row-scaled NVFP4 backward does not support fused bias gradient." out_2d.add_(bias.to(dtype=torch.float32)) if requested_out is not None: @@ -417,7 +443,8 @@ def general_grouped_gemm( if grad and use_bias: grad_bias = [ - torch.empty(B[i].size(1), dtype=out[0].dtype, device="cuda") for i in range(num_gemms) + torch.empty(B[i].size(1), dtype=out[0].dtype, device=te_device_type()) + for i in range(num_gemms) ] else: grad_bias = empty_tensors @@ -430,7 +457,9 @@ def general_grouped_gemm( if any(_is_nvfp4_row_scaled_tensor(tensor) for tensor in A) or any( _is_nvfp4_row_scaled_tensor(tensor) for tensor in B ): - assert D_dtype is None, "Row-scaled NVFP4 grouped GEMM currently does not support D_dtype." + assert ( + D_dtype is None + ), "Row-scaled NVFP4 grouped GEMM currently does not support D_dtype." if single_output: assert ( m_splits is not None @@ -563,7 +592,9 @@ def _get_grouped_cublas_workspace(device: int, layout: str) -> torch.Tensor: kernels). """ assert layout in ("TN", "NN", "NT"), f"unexpected grouped GEMM layout {layout}" - return torch.empty(get_cublas_workspace_size_bytes(), dtype=torch.uint8, device=device) + return torch.empty( + get_cublas_workspace_size_bytes(), dtype=torch.uint8, device=device + ) def general_grouped_gemm_for_grouped_tensor( @@ -599,11 +630,17 @@ def general_grouped_gemm_for_grouped_tensor( raise ValueError("Both A and out are discrete. This is not supported yet.") if isinstance(A, GroupedTensorStorage) and A.row_scaled_nvfp4: - raise NotImplementedError("Row-scaled NVFP4 GroupedTensor GEMM is not supported yet.") + raise NotImplementedError( + "Row-scaled NVFP4 GroupedTensor GEMM is not supported yet." + ) if isinstance(B, GroupedTensorStorage) and B.row_scaled_nvfp4: - raise NotImplementedError("Row-scaled NVFP4 GroupedTensor GEMM is not supported yet.") + raise NotImplementedError( + "Row-scaled NVFP4 GroupedTensor GEMM is not supported yet." + ) if isinstance(out, GroupedTensorStorage) and out.row_scaled_nvfp4: - raise NotImplementedError("Row-scaled NVFP4 GroupedTensor GEMM is not supported yet.") + raise NotImplementedError( + "Row-scaled NVFP4 GroupedTensor GEMM is not supported yet." + ) def _is_fp8_blockwise(operand) -> bool: if isinstance(operand, (list, tuple)): diff --git a/transformer_engine/pytorch/cpu_offload.py b/transformer_engine/pytorch/cpu_offload.py index c81c18e64fc..f9e6aba73e3 100644 --- a/transformer_engine/pytorch/cpu_offload.py +++ b/transformer_engine/pytorch/cpu_offload.py @@ -17,6 +17,7 @@ import transformer_engine.pytorch as te import transformer_engine.pytorch.cpu_offload_v1 as v1_code_path from .quantized_tensor import ( +from transformer_engine import te_platform, te_device_type restore_from_saved, prepare_for_saving, QuantizedTensor, @@ -73,8 +74,8 @@ def _mark_tensor_for_offload(t): if t is None: return # Attach an event to mark when the tensor is ready for reload. - t.start_reload_event = torch.cuda.Event() - t.start_reload_event.record(torch.cuda.current_stream()) + t.start_reload_event = te_platform().Event() + t.start_reload_event.record(te_platform().current_stream()) if offload_base_tensor and t._base is not None: setattr(t, "offload_base_tensor", True) @@ -286,7 +287,7 @@ def start_offload(self): # Wait for the moment the tensor is ready to be offloaded. self.offload_stream.wait_event(self.fwd_gpu_tensor_group.events[tensor_id]) # type: ignore[arg-type] - with torch.cuda.stream(self.offload_stream): + with te_platform().stream(self.offload_stream): if allocate_cpu_buffers: offloaded_tensor = torch.empty_like( tensor, device=torch.device("cpu"), pin_memory=True @@ -308,7 +309,7 @@ def start_offload(self): self.aux = aux if len(self.fwd_gpu_tensor_group.tensor_list) > 0: - self.finish_offload_event = torch.cuda.Event() + self.finish_offload_event = te_platform().Event() self.finish_offload_event.record(self.offload_stream) def release_activation_forward_gpu_memory(self): @@ -321,7 +322,7 @@ def release_activation_forward_gpu_memory(self): ) self.state = "offload_finished" if len(self.fwd_gpu_tensor_group.tensor_list) > 0: - torch.cuda.current_stream().wait_event(self.finish_offload_event) # type: ignore[arg-type] + te_platform().current_stream().wait_event(self.finish_offload_event) # type: ignore[arg-type] # GPU memory can be released safely after the offload. # Notice that the memory needs to be kept alive when GPU->CPU copy is performed. @@ -347,13 +348,13 @@ def start_reload(self): # cannot move tensors from pool of one stream to another without # calling cudaFree and cudaMalloc again. - reloaded_tensor = torch.empty_like(tensor, device=torch.device("cuda")) - self.offload_stream.wait_stream(torch.cuda.current_stream()) + reloaded_tensor = torch.empty_like(tensor, device=torch.device(te_device_type())) + self.offload_stream.wait_stream(te_platform().current_stream()) - with torch.cuda.stream(self.offload_stream): + with te_platform().stream(self.offload_stream): reloaded_tensor.copy_(tensor, non_blocking=True) - reload_tensor_event = torch.cuda.Event() + reload_tensor_event = te_platform().Event() reload_tensor_event.record(self.offload_stream) self.bwd_gpu_tensor_group.events.append(reload_tensor_event) self.bwd_gpu_tensor_group.tensor_list.append(reloaded_tensor) @@ -394,8 +395,8 @@ def push_tensor(self, tensor: torch.Tensor) -> int | torch.Tensor | tuple[list, if hasattr(tensor, "start_reload_event"): self.fwd_gpu_tensor_group.events.append(tensor.start_reload_event) else: - self.fwd_gpu_tensor_group.events.append(torch.cuda.Event()) - self.fwd_gpu_tensor_group.events[-1].record(torch.cuda.current_stream()) + self.fwd_gpu_tensor_group.events.append(te_platform().Event()) + self.fwd_gpu_tensor_group.events[-1].record(te_platform().current_stream()) return len(self.fwd_gpu_tensor_group.tensor_list) - 1 return tensor @@ -438,7 +439,7 @@ def pop_tensor( f"but got state='{self.state}' for tensor={tensor_or_tensor_id}" ) # wait for the tensor to be reloaded - torch.cuda.current_stream().wait_event( + te_platform().current_stream().wait_event( self.bwd_gpu_tensor_group.events[tensor_or_tensor_id] ) return self.bwd_gpu_tensor_group.tensor_list[tensor_or_tensor_id] @@ -463,7 +464,7 @@ def _check_if_offload(self, t: torch.Tensor) -> bool: not isinstance(t, torch.nn.Parameter) and not getattr(t, "_TE_do_not_offload", False) and not isinstance(t, torch._subclasses.FakeTensor) - and t.device.type == "cuda" + and t.device.type == te_device_type() ): if not t.is_contiguous() and not getattr(t, "offload_base_tensor", False): warnings.warn( @@ -506,7 +507,7 @@ def __init__( offload_stream: Optional[torch.cuda.Stream] = None, ): self.num_layers = num_layers - self.offload_stream = offload_stream if offload_stream is not None else torch.cuda.Stream() + self.offload_stream = offload_stream if offload_stream is not None else te_platform().Stream() self.layer_states = { i: OffloadableLayerState(self.offload_stream, retain_pinned_cpu_buffers) @@ -775,7 +776,7 @@ def get_cpu_offload_context( .. code-block:: python - offload_stream = torch.cuda.Stream() + offload_stream = te_platform().Stream() cpu_offload_context, sync_function, manual_controller = get_cpu_offload_context( enabled=True, model_layers=num_layers, manual_synchronization=True, offload_stream=offload_stream) diff --git a/transformer_engine/pytorch/cpu_offload_v1.py b/transformer_engine/pytorch/cpu_offload_v1.py index fb62546cc0a..30ddd57da2c 100644 --- a/transformer_engine/pytorch/cpu_offload_v1.py +++ b/transformer_engine/pytorch/cpu_offload_v1.py @@ -12,6 +12,7 @@ from transformer_engine.debug.pytorch.debug_state import TEDebugState from .quantized_tensor import QuantizedTensorStorage from .tensor.float8_tensor import Float8Tensor +from transformer_engine import te_platform __all__ = ["get_cpu_offload_context"] @@ -363,8 +364,8 @@ def __init__( self.layer_window_map[i] += constant # allocate streams and events for synchronization - self.d2h_stream = torch.cuda.Stream() - self.h2d_stream = torch.cuda.Stream() + self.d2h_stream = te_platform().Stream() + self.h2d_stream = te_platform().Stream() def tensor_push(self, tensor: torch.Tensor, **kwargs) -> Any: global CPUOffloadedLayer @@ -463,7 +464,7 @@ def tensor_pop(self, tensor_tag, **kwargs): def bulk_offload_group(self, group_to_offload): """Bulk offload group.""" - with torch.cuda.stream(self.d2h_stream): + with te_platform().stream(self.d2h_stream): for tensor_tag, state in self.tensor_tag_to_state.items(): group_id, _ = tensor_tag if group_id == group_to_offload: @@ -507,7 +508,7 @@ def synchronize_on_group_commit_forward(self, current_group): # For the first group, kickstart the offload after we have # the first compute completion if current_group == 0: - self.d2h_stream.wait_stream(torch.cuda.current_stream()) + self.d2h_stream.wait_stream(te_platform().current_stream()) if not self.double_buffer_created: # Creating the first copy of double buffer for tensors that are offloaded @@ -529,8 +530,8 @@ def synchronize_on_group_commit_forward(self, current_group): if self.layer_window_map[self.offloaded_group_count] == current_group: # Stream synchronization both ways - self.d2h_stream.wait_stream(torch.cuda.current_stream()) - torch.cuda.current_stream().wait_stream(self.d2h_stream) + self.d2h_stream.wait_stream(te_platform().current_stream()) + te_platform().current_stream().wait_stream(self.d2h_stream) # Time to free the activation memory after usage for tensor_tag, tensor_buf in self.tensor_tag_to_buf.items(): @@ -577,9 +578,9 @@ def bulk_reload_group(self, group_to_reload): buffer_idx = 0 double_buffer_idx = group_to_reload % 2 - main_stream = torch.cuda.current_stream() + main_stream = te_platform().current_stream() - with torch.cuda.stream(self.h2d_stream): + with te_platform().stream(self.h2d_stream): # move back tensors for tensor_label, state in self.tensor_tag_to_state.items(): group_id, _ = tensor_label @@ -589,9 +590,9 @@ def bulk_reload_group(self, group_to_reload): if self.double_buffering: reload_buffer = self.reload_double_buffer[double_buffer_idx][buffer_idx] else: - with torch.cuda.stream(main_stream): + with te_platform().stream(main_stream): reload_buffer = torch.empty_like( - state[1], device=torch.cuda.current_device() + state[1], device=te_platform().current_device() ) recovered_tensor = SynchronizedGroupOffloadHandler.reload( @@ -609,9 +610,9 @@ def bulk_reload_group(self, group_to_reload): buffer_idx ] else: - with torch.cuda.stream(main_stream): + with te_platform().stream(main_stream): reload_buffer = torch.empty_like( - state_tuple[1], device=torch.cuda.current_device() + state_tuple[1], device=te_platform().current_device() ) tensor_list.append( @@ -654,8 +655,8 @@ def on_group_commit_backward(self): if self.layer_window_map[self.offloaded_group_count - 1] == self.current_group: # Stream synchronization both ways - self.h2d_stream.wait_stream(torch.cuda.current_stream()) - torch.cuda.current_stream().wait_stream(self.h2d_stream) + self.h2d_stream.wait_stream(te_platform().current_stream()) + te_platform().current_stream().wait_stream(self.h2d_stream) # Time to reload the next group self.bulk_reload_group(self.offloaded_group_count - 1) @@ -665,7 +666,7 @@ def on_group_commit_backward(self): # Last group computation needs to wait till all the reloads complete if self.current_group == 0: - torch.cuda.current_stream().wait_stream(self.h2d_stream) + te_platform().current_stream().wait_stream(self.h2d_stream) self.offloaded_group_count = 0 diff --git a/transformer_engine/pytorch/custom_recipes/quantizer_factories.py b/transformer_engine/pytorch/custom_recipes/quantizer_factories.py index 250ce8c9393..fbbfe6ef00e 100644 --- a/transformer_engine/pytorch/custom_recipes/quantizer_factories.py +++ b/transformer_engine/pytorch/custom_recipes/quantizer_factories.py @@ -30,6 +30,7 @@ import torch from transformer_engine.pytorch.quantization import QuantizerRole from ..constants import DType +from transformer_engine import te_device_type def high_precision_factory( @@ -86,7 +87,7 @@ def current_scaling_factory( return Float8CurrentScalingQuantizer( fp8_dtype=fp8_dtype, - device=torch.device("cuda"), + device=torch.device(te_device_type()), force_pow_2_scales=False, amax_epsilon=0.0, ) diff --git a/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py b/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py index 0d5fc0c4419..95adb53998e 100644 --- a/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py +++ b/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py @@ -74,6 +74,7 @@ from transformer_engine.pytorch.quantization import QuantizerRole from ..constants import DType from .quantizer_factories import mxfp8_factory, nvfp4_factory +from transformer_engine import te_device_type # ----------------------------------------------------------------------------- # Pre-training-Oriented Recipes @@ -426,6 +427,6 @@ def nvfp4_linear_fp8_dpa_factory( is_dpa_boundary and "dpa_grad_input" in role.name ) fp8_dtype = DType.kFloat8E5M2 if is_bwd_role else DType.kFloat8E4M3 - return Float8CurrentScalingQuantizer(fp8_dtype=fp8_dtype, device="cuda") + return Float8CurrentScalingQuantizer(fp8_dtype=fp8_dtype, device=te_device_type()) return nvfp4_factory(role) diff --git a/transformer_engine/pytorch/custom_recipes/reference_current_scaling.py b/transformer_engine/pytorch/custom_recipes/reference_current_scaling.py index 4e5292019ac..a68fd9489b9 100644 --- a/transformer_engine/pytorch/custom_recipes/reference_current_scaling.py +++ b/transformer_engine/pytorch/custom_recipes/reference_current_scaling.py @@ -13,6 +13,7 @@ from transformer_engine.pytorch.custom_recipes import gemm from transformer_engine.pytorch.custom_recipes import reference_utils from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage, Quantizer +from transformer_engine import te_device_type def current_scaling_ref_factory(role): @@ -501,7 +502,7 @@ def make_empty( # Canonicalize tensor attributes if device is None: - device = torch.device("cuda") + device = torch.device(te_device_type()) # Allocate quantized data qx = torch.empty(shape, dtype=self.dtype, device=device) diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index 55b9d4d914c..31442d973b5 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -20,6 +20,7 @@ from torch.distributed.fsdp import FullyShardedDataParallel as FSDP from torch.distributed.fsdp._common_utils import _get_module_fsdp_state from torch.distributed.fsdp._traversal_utils import _get_fsdp_states_with_modules +from transformer_engine import te_platform, te_device_type try: import torch.distributed._symmetric_memory as symm_mem @@ -96,7 +97,7 @@ def is_graph_safe_rng_state(state: Union[torch.Tensor, torch.Generator]) -> bool def _get_cuda_rng_state( - device: Union[int, str, torch.device] = "cuda", + device: Union[int, str, torch.device] = te_device_type(), clone: bool = False, graph_safe: bool = True, ) -> torch.Tensor: @@ -106,10 +107,10 @@ def _get_cuda_rng_state( if isinstance(device, str): device = torch.device(device) elif isinstance(device, int): - device = torch.device("cuda", device) + device = torch.device(te_device_type(), device) idx = device.index if idx is None: - idx = torch.cuda.current_device() + idx = te_platform().current_device() default_generator = torch.cuda.default_generators[idx] if graph_safe_rng_available() and graph_safe: if clone: @@ -128,16 +129,16 @@ def _set_cuda_rng_state( """Sets the random number generator state of the current GPU.""" if device == -1: - device = torch.device("cuda") + device = torch.device(te_device_type()) elif isinstance(device, str): device = torch.device(device) elif isinstance(device, int): - device = torch.device("cuda", device) + device = torch.device(te_device_type(), device) def cb() -> None: idx = device.index if idx is None: - idx = torch.cuda.current_device() + idx = te_platform().current_device() default_generator = torch.cuda.default_generators[idx] if graph_safe_rng_available() and graph_safe: default_generator.graphsafe_set_state(new_state) @@ -216,7 +217,7 @@ def split_tensor_into_1d_equal_chunks( data = torch.empty( partition_size, dtype=tensor.dtype, - device=torch.cuda.current_device(), + device=te_platform().current_device(), requires_grad=False, ) data.copy_(tensor.view(-1)[start_index:end_index]) @@ -231,7 +232,7 @@ def gather_split_1d_tensor(tensor: torch.Tensor, tp_group: dist_group_type) -> t gathered = torch.empty( numel_gathered, dtype=tensor.dtype, - device=torch.cuda.current_device(), + device=te_platform().current_device(), requires_grad=False, ) torch.distributed.all_gather_into_tensor(gathered, tensor, group=tp_group) @@ -294,10 +295,10 @@ def _get_active_autocast_contexts(): autocast_cached = torch.is_autocast_cache_enabled() if torch_version() >= (2, 4, 0): - gpu_autocast_enabled = torch.is_autocast_enabled("cuda") - gpu_autocast_dtype = torch.get_autocast_dtype("cuda") + gpu_autocast_enabled = torch.is_autocast_enabled(te_device_type()) + gpu_autocast_dtype = torch.get_autocast_dtype(te_device_type()) gpu_autocast_ctx = torch.amp.autocast( - "cuda", + te_device_type(), enabled=gpu_autocast_enabled, dtype=gpu_autocast_dtype, cache_enabled=autocast_cached, @@ -314,7 +315,7 @@ def _get_active_autocast_contexts(): else: gpu_autocast_enabled = torch.is_autocast_enabled() gpu_autocast_dtype = torch.get_autocast_gpu_dtype() - gpu_autocast_ctx = torch.cuda.amp.autocast( + gpu_autocast_ctx = te_platform().amp.autocast( gpu_autocast_enabled, gpu_autocast_dtype, autocast_cached ) @@ -963,7 +964,7 @@ def reduce_scatter_along_first_dim( dim_size[0] = dim_size[0] // world_size if output is None: - output = torch.empty(dim_size, dtype=inp.dtype, device=torch.cuda.current_device()) + output = torch.empty(dim_size, dtype=inp.dtype, device=te_platform().current_device()) handle = torch.distributed.reduce_scatter_tensor( output, inp.contiguous(), group=tp_group, async_op=async_op ) @@ -1045,7 +1046,7 @@ def _all_gather_fp8( out: Float8TensorStorage if quantizer is not None: dtype = torch.float32 - device = "cuda" + device = te_device_type() if isinstance(inp, Float8Tensor): dtype = inp.dtype device = inp.device @@ -1964,7 +1965,7 @@ def release_symm_mem_pool() -> None: "release_symm_mem_pool()." ) pools.pop(device, None) - torch.cuda.empty_cache() + te_platform().empty_cache() def symm_mem_alloc( @@ -1985,7 +1986,7 @@ def symm_mem_alloc( be a shared static buffer. ``backend`` selects the symm-mem backend (default NCCL; for the pool it is captured at pool creation).""" if device is None: - device = torch.device("cuda", torch.cuda.current_device()) + device = torch.device(te_device_type(), te_platform().current_device()) if not HAS_TORCH_SYMMETRIC: raise RuntimeError( "torch.distributed._symmetric_memory is unavailable; symm_mem_alloc " diff --git a/transformer_engine/pytorch/dynamo/tensor_spec.py b/transformer_engine/pytorch/dynamo/tensor_spec.py index 4cfe225952a..55511c2ff51 100644 --- a/transformer_engine/pytorch/dynamo/tensor_spec.py +++ b/transformer_engine/pytorch/dynamo/tensor_spec.py @@ -11,6 +11,7 @@ import torch from torch._prims_common import make_contiguous_strides_for +from transformer_engine import te_device_type @dataclass @@ -104,7 +105,7 @@ def create_inner_tensors(self) -> List[torch.Tensor]: Under ``register_fake`` the ``torch.empty`` calls produce ``FakeTensor``s; ``requires_grad`` is left default (managed by ``register_autograd``). """ - device = self.device if self.device is not None else torch.device("cuda") + device = self.device if self.device is not None else torch.device(te_device_type()) if self.quantizer is None: return [torch.empty(tuple(self.shape), dtype=self.dtype, device=device)] inner = self.quantizer.alloc_tensors(tuple(self.shape), device=device) @@ -135,7 +136,7 @@ def create_tensor(self) -> torch.Tensor: inner tensors via :meth:`assemble`. """ if self.quantizer is None: - device = self.device if self.device is not None else torch.device("cuda") + device = self.device if self.device is not None else torch.device(te_device_type()) return torch.empty( tuple(self.shape), dtype=self.dtype, diff --git a/transformer_engine/pytorch/ep.py b/transformer_engine/pytorch/ep.py index 4464036fc78..0508d5d2653 100644 --- a/transformer_engine/pytorch/ep.py +++ b/transformer_engine/pytorch/ep.py @@ -18,6 +18,7 @@ from .cpu_offload import mark_not_offload from .distributed import symm_mem_alloc, release_symm_mem_pool from .quantized_tensor import QuantizedTensor +from transformer_engine import te_platform, te_device_type # Type-hint-only import; keeps the ``Recipe`` annotation without a runtime import of # common.recipe (the concrete recipe classes are imported lazily where used). @@ -158,8 +159,8 @@ def ep_bootstrap( ) # Materialize the PG's NCCL comm before borrowing its raw handle. - dist.barrier(group=ep_group, device_ids=[torch.cuda.current_device()]) - comm_ptr = ep_group._get_backend(torch.device("cuda"))._comm_ptr() + dist.barrier(group=ep_group, device_ids=[te_platform().current_device()]) + comm_ptr = ep_group._get_backend(torch.device(te_device_type()))._comm_ptr() tex.ep_initialize( int(comm_ptr), @@ -274,7 +275,7 @@ def __init__( if not _BOOTSTRAPPED: raise RuntimeError("EpBuffer requires ep_bootstrap() to be called first.") if device is None: - device = torch.device("cuda", torch.cuda.current_device()) + device = torch.device(te_device_type(), te_platform().current_device()) alignment = int(alignment) if alignment > 1 and (alignment & (alignment - 1)) != 0: raise ValueError(f"alignment must be 0, 1, or a power of two (got {alignment}).") @@ -897,7 +898,7 @@ def _alloc_io(shape, dtype: torch.dtype, device, zero_copy: bool) -> torch.Tenso The zero-copy pool path is not CUDA-graph capturable; supply persistent recv_tokens / grad_out buffers to capture a graph.""" if zero_copy: - if torch.cuda.is_current_stream_capturing(): + if te_platform().is_current_stream_capturing(): raise RuntimeError( "EP zero-copy pool allocation is not CUDA-graph capturable; supply persistent " "recv_tokens / grad_out buffers (allocated once via symm_mem_alloc) before capture." diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index 04fa56721d9..f518c943b92 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -18,6 +18,7 @@ from transformer_engine.common.recipe import DelayedScaling, Recipe from transformer_engine.pytorch.constants import dist_group_type from .quantization import ( +from transformer_engine import te_platform autocast, FP8GlobalStateManager, get_default_fp8_recipe, @@ -103,7 +104,7 @@ def _graph_context_wrapper(*args, **kwargs): gc_is_enabled = gc.isenabled() if gc_is_enabled: gc.disable() - with torch.cuda.graph(*args, **kwargs): + with te_platform().graph(*args, **kwargs): yield if gc_is_enabled: gc.enable() @@ -458,9 +459,9 @@ def _make_graphed_callables( for i in range(len(flatten_sample_args)) ] - fwd_graphs = [torch.cuda.CUDAGraph() for _ in range(len(flatten_sample_args))] - bwd_graphs = [torch.cuda.CUDAGraph() for _ in range(len(flatten_sample_args))] - bwd_dw_graphs = [torch.cuda.CUDAGraph() for _ in range(len(flatten_sample_args))] + fwd_graphs = [te_platform().CUDAGraph() for _ in range(len(flatten_sample_args))] + bwd_graphs = [te_platform().CUDAGraph() for _ in range(len(flatten_sample_args))] + bwd_dw_graphs = [te_platform().CUDAGraph() for _ in range(len(flatten_sample_args))] graph_callables = [None for _ in range(len(flatten_sample_args))] def _returned_param_grad_clone_slots(static_grad_inputs, module_params): @@ -494,7 +495,7 @@ def _returned_param_grad_clone_slots(static_grad_inputs, module_params): # Warmup # Hopefully prevents cudnn benchmarking and other lazy-initialization cuda work # from ending up in any captures. - torch.cuda.synchronize() + te_platform().synchronize() # Get warmup func and func_idx. warmup_func_idx = [] @@ -630,7 +631,7 @@ def _run_warmup_backward(func_idx, func, outputs, warmup_iter, callable_idx): need_bwd_dw_graph[func_idx] = need_backward_dw # Run warmup and do the above filtering. - with torch.cuda.stream(torch.cuda.Stream()): + with te_platform().stream(te_platform().Stream()): if pre_warmup_hook is not None: pre_warmup_hook() @@ -688,7 +689,7 @@ def _run_warmup_backward(func_idx, func, outputs, warmup_iter, callable_idx): if post_warmup_hook is not None: post_warmup_hook() - torch.cuda.synchronize() + te_platform().synchronize() # All captures here share a mempool. To avoid replays corrupting each other's memory, # the safest approach is to capture all passes in the same order they'll run: @@ -1081,14 +1082,14 @@ def forward(ctx, skip_fp8_weight_update, cuda_graph_stream, cuda_graph_event, *i static_input_surface[i].copy_(inputs[i]) # Replay forward graph - if cuda_graph_stream != torch.cuda.current_stream(): - cuda_graph_stream.wait_stream(torch.cuda.current_stream()) + if cuda_graph_stream != te_platform().current_stream(): + cuda_graph_stream.wait_stream(te_platform().current_stream()) with cuda_graph_stream: fwd_graph.replay() if cuda_graph_event is not None: - torch.cuda.current_stream().wait_event(cuda_graph_event) + te_platform().current_stream().wait_event(cuda_graph_event) else: - torch.cuda.current_stream().wait_stream(cuda_graph_stream) + te_platform().current_stream().wait_stream(cuda_graph_stream) else: fwd_graph.replay() if not isinstance(static_outputs, tuple): @@ -1117,14 +1118,14 @@ def backward(ctx, *grads): # incoming grad is already in the right place if g.data_ptr() != grad.data_ptr(): g.copy_(grad) - if ctx.cuda_graph_stream != torch.cuda.current_stream(): - ctx.cuda_graph_stream.wait_stream(torch.cuda.current_stream()) + if ctx.cuda_graph_stream != te_platform().current_stream(): + ctx.cuda_graph_stream.wait_stream(te_platform().current_stream()) with ctx.cuda_graph_stream: bwd_graph.replay() if ctx.cuda_graph_event is not None: - torch.cuda.current_stream().wait_event(ctx.cuda_graph_event) + te_platform().current_stream().wait_event(ctx.cuda_graph_event) else: - torch.cuda.current_stream().wait_stream(ctx.cuda_graph_stream) + te_platform().current_stream().wait_stream(ctx.cuda_graph_stream) else: bwd_graph.replay() @@ -1177,7 +1178,7 @@ def functionalized(*user_args, **user_kwargs): cuda_graph_stream = user_kwargs["cuda_graph_stream"] user_kwargs.pop("cuda_graph_stream") else: - cuda_graph_stream = torch.cuda.current_stream() + cuda_graph_stream = te_platform().current_stream() if "cuda_graph_event" in user_kwargs: cuda_graph_event = user_kwargs["cuda_graph_event"] user_kwargs.pop("cuda_graph_event") @@ -1689,12 +1690,12 @@ def call_func(self, *args, **kwargs): # Save RNG state. if graph_safe_rng_available(): generators = [ - torch.cuda.default_generators[torch.cuda.current_device()], + torch.cuda.default_generators[te_platform().current_device()], *get_all_rng_states().values(), ] original_rng_states = [state.get_state() for state in generators] else: - original_rng_states = torch.cuda.get_rng_state() + original_rng_states = te_platform().get_rng_state() graphed_callables = _make_graphed_callables( forward_funcs, @@ -1719,7 +1720,7 @@ def call_func(self, *args, **kwargs): for gen, state in zip(generators, original_rng_states): gen.set_state(state) else: - torch.cuda.set_rng_state(original_rng_states) + te_platform().set_rng_state(original_rng_states) # Remove FP8 wrapper. for module_cls, old_call in old_call_funcs.items(): diff --git a/transformer_engine/pytorch/jit.py b/transformer_engine/pytorch/jit.py index 9bfd87b4808..145eb2c45a7 100644 --- a/transformer_engine/pytorch/jit.py +++ b/transformer_engine/pytorch/jit.py @@ -13,6 +13,7 @@ from .torch_version import torch_version from .export import is_in_onnx_export_mode from .utils import gpu_autocast_ctx +from transformer_engine import te_platform # pylint: disable=unnecessary-lambda-assignment @@ -81,7 +82,8 @@ def wrapper(*args, **kwargs): # Arguments are matched to names by position, which only holds # without a *args in between. assert not any( - p.kind is inspect.Parameter.VAR_POSITIONAL for p in parameters.values() + p.kind is inspect.Parameter.VAR_POSITIONAL + for p in parameters.values() ), f"no_torch_dynamo(when=...) does not support *args, which {f.__name__} takes" parameter_names = list(parameters) @@ -168,9 +170,9 @@ def bgrad_dgelu_fused_( x = inp + bias tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 - ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) + ff = 0.5 * x * ( + (1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x) + ) + 0.5 * (1 + tanh_out) dgelu = ff * grad_output bgrad = dgelu.sum(dim=0) return bgrad, dgelu @@ -184,9 +186,9 @@ def dgelu_fused_(grad_output: torch.Tensor, inp: torch.Tensor) -> torch.Tensor: x = inp tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 - ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( - 1 + tanh_out - ) + ff = 0.5 * x * ( + (1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x) + ) + 0.5 * (1 + tanh_out) dgelu = ff * grad_output return dgelu @@ -203,7 +205,9 @@ def l2normalization_fused_(x: torch.Tensor, eps: float) -> torch.Tensor: @jit_fuser -def l2normalization_fwd_fused_(x: torch.Tensor, eps: float) -> tuple[torch.Tensor, torch.Tensor]: +def l2normalization_fwd_fused_( + x: torch.Tensor, eps: float +) -> tuple[torch.Tensor, torch.Tensor]: """L2 normalization fused - training version that returns intermediate values""" x_fp32 = x.float() x_squared = x_fp32.pow(2) @@ -257,7 +261,9 @@ def l2normalization_fused(x: torch.Tensor, eps: float) -> torch.Tensor: return l2normalization_fused_(x, eps) -def l2normalization_fwd_fused(x: torch.Tensor, eps: float) -> tuple[torch.Tensor, torch.Tensor]: +def l2normalization_fwd_fused( + x: torch.Tensor, eps: float +) -> tuple[torch.Tensor, torch.Tensor]: """Disable native AMP for l2normalization_fwd_fused_ - training version""" with gpu_autocast_ctx(enabled=False): return l2normalization_fwd_fused_(x, eps) @@ -335,15 +341,25 @@ def warmup_jit_bias_dropout_add( """Compile BDA JIT function before the main training steps""" # Save cuda RNG state to ensure warmup does not affect reproducibility. - rng_state = torch.cuda.get_rng_state() + rng_state = te_platform().get_rng_state() - inp = torch.rand((seq_length, micro_batch_size, hidden_size), dtype=dtype, device="cuda") - residual = torch.rand((seq_length, micro_batch_size, hidden_size), dtype=dtype, device="cuda") - bias = torch.rand((hidden_size), dtype=dtype, device="cuda") + inp = torch.rand( + (seq_length, micro_batch_size, hidden_size), + dtype=dtype, + device=te_device_type(), + ) + residual = torch.rand( + (seq_length, micro_batch_size, hidden_size), + dtype=dtype, + device=te_device_type(), + ) + bias = torch.rand((hidden_size), dtype=dtype, device=te_device_type()) dropout_rate = 0.1 # Warmup JIT fusions with the input grad_enable state of both forward # prop and recomputation - for input_grad, bias_grad, residual_grad in zip([False, True], [True, True], [True, True]): + for input_grad, bias_grad, residual_grad in zip( + [False, True], [True, True], [True, True] + ): inp.requires_grad = input_grad bias.requires_grad = bias_grad residual.requires_grad = residual_grad @@ -351,8 +367,8 @@ def warmup_jit_bias_dropout_add( output = bias_dropout_add_fused_train(inp, bias, residual, dropout_rate) del bias, inp, residual, output - torch.cuda.empty_cache() - torch.cuda.set_rng_state(rng_state) + te_platform().empty_cache() + te_platform().set_rng_state(rng_state) def warmup_jit_bias_dropout_add_all_dtypes( @@ -372,13 +388,15 @@ def warmup_jit_bias_gelu( """Compile bias-gelu JIT function before the main training steps""" # Save cuda RNG state to ensure warmup does not affect reproducibility. - rng_state = torch.cuda.get_rng_state() + rng_state = te_platform().get_rng_state() - bias = torch.rand(ffn_hidden_size_per_partition, dtype=dtype, device="cuda") + bias = torch.rand( + ffn_hidden_size_per_partition, dtype=dtype, device=te_device_type() + ) inp = torch.rand( (seq_length * micro_batch_size, ffn_hidden_size_per_partition), dtype=dtype, - device="cuda", + device=te_device_type(), ) # Warmup JIT fusions with the input grad_enable state of both forward # prop and recomputation @@ -389,8 +407,8 @@ def warmup_jit_bias_gelu( _ = gelu_fused_(inp) del bias, inp - torch.cuda.empty_cache() - torch.cuda.set_rng_state(rng_state) + te_platform().empty_cache() + te_platform().set_rng_state(rng_state) def warmup_jit_bias_gelu_all_dtypes( @@ -407,12 +425,12 @@ def warmup_jit_l2normalization( """Compile L2Normalization JIT function before the main training steps""" # Save cuda RNG state to ensure warmup does not affect reproducibility. - rng_state = torch.cuda.get_rng_state() + rng_state = te_platform().get_rng_state() inp = torch.rand( (seq_length * micro_batch_size, hidden_size), dtype=dtype, - device="cuda", + device=te_device_type(), ) eps = 1e-6 # Warmup JIT fusions with the input grad_enable state of both forward @@ -431,8 +449,8 @@ def warmup_jit_l2normalization( output = l2normalization_fused_(inp, eps) del inp, output - torch.cuda.empty_cache() - torch.cuda.set_rng_state(rng_state) + te_platform().empty_cache() + te_platform().set_rng_state(rng_state) def warmup_jit_l2normalization_all_dtypes( diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index a3131f7436d..15963252c9d 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -22,6 +22,7 @@ from ._common import _ParameterInitMeta, noop_cat from .._extra_state import ( +from transformer_engine import te_platform, te_device_type extra_state_pickle_advisory, is_stateless_recipe, should_load_extra_state_pickle, @@ -140,7 +141,7 @@ def get_dummy_wgrad(shape: list, dtype: torch.dtype, zero=False) -> torch.Tensor _dummy_wgrads[key] = torch.empty( shape, dtype=dtype, - device="cuda", + device=te_device_type(), requires_grad=False, ) if zero: @@ -892,7 +893,7 @@ class TransformerEngineBaseModule(torch.nn.Module, ABC): def __init__(self, name: Optional[str] = None) -> None: super().__init__() - if not torch.cuda.is_available(): + if not te_platform().is_available(): raise RuntimeError("TransformerEngine needs CUDA.") self.name = name self.next_iter_when_debug_should_be_run = 0 @@ -1411,7 +1412,7 @@ def to_cpu(src: torch.Tensor) -> torch.Tensor: state["extra_fp8_variables"] = extra # Serialize state into byte tensor - torch.cuda.synchronize() + te_platform().synchronize() state_serialized = bytearray(pickle.dumps(state)) state_serialized = torch.frombuffer(state_serialized, dtype=torch.uint8) return state_serialized @@ -1438,7 +1439,7 @@ def set_extra_state(self, state: torch.Tensor) -> None: if not unsafe_pickle_extra_state_enabled(): raise RuntimeError(extra_state_pickle_advisory(context)) state.seek(0) - state = torch.load(state, map_location="cuda", weights_only=False) + state = torch.load(state, map_location=te_device_type(), weights_only=False) else: raise RuntimeError("Unsupported checkpoint format.") @@ -1477,7 +1478,7 @@ def copy_tensor(src: torch.Tensor, dst: torch.Tensor) -> None: copy_tensor(state["amax_history_fwd"], self.fp8_meta["scaling_fwd"].amax_history) copy_tensor(state["scale_bwd"], self.fp8_meta["scaling_bwd"].scale) copy_tensor(state["amax_history_bwd"], self.fp8_meta["scaling_bwd"].amax_history) - torch.cuda.synchronize() + te_platform().synchronize() def set_activation_dtype(self, inp: torch.Tensor) -> None: """Get activation data type for AMP.""" @@ -1818,7 +1819,7 @@ def reset_parameters(self, defer_init: Optional[bool] = False) -> None: param = param._local_tensor if is_dtensor else param # Ensure parameter is on a real device if param.device == torch.device("meta"): - param = torch.empty_like(param, device="cuda") + param = torch.empty_like(param, device=te_device_type()) # Initialize the parameter values on device init_fn = self.param_init_meta[name].init_fn get_rng_state_tracker = self.param_init_meta[name].get_rng_state_tracker diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 8b078944f17..9d22e30530f 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -17,6 +17,7 @@ from transformer_engine.common.recipe import Recipe from transformer_engine.pytorch.tensor.grouped_tensor import ( +from transformer_engine import te_device_type GroupedTensor, GroupedTensorStorage, ) @@ -782,7 +783,7 @@ def forward( "single_grouped_bias to allow the split-quantize fallback." ) if grouped_tensor_supported: - if m_splits.device.type != "cuda": + if m_splits.device.type != te_device_type(): raise ValueError( "The native grouped_tensor path requires CUDA m_splits. Pass a CUDA int64 " "tensor, or set use_grouped_tensor=False." @@ -1556,7 +1557,7 @@ class GroupedLinear(TransformerEngineBaseModule): used to get the random number generator state tracker for initializing weights. rng_tracker_name : str, default = None the param passed to get_rng_state_tracker to get the specific rng tracker. - device : Union[torch.device, str], default = "cuda" + device : Union[torch.device, str], default=te_device_type() The device on which the parameters of the model will be allocated. It is the user's responsibility to ensure all parameters are moved to the GPU before running the forward pass. @@ -1633,7 +1634,7 @@ def __init__( return_bias: bool = False, params_dtype: Optional[torch.dtype] = None, parallel_mode: Optional[str] = None, - device: Union[torch.device, str] = "cuda", + device: Union[torch.device, str] = te_device_type(), ub_overlap_rs: bool = False, ub_overlap_ag: bool = False, ub_name: Optional[str] = None, diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index 561e813348c..bc6e257c5f4 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -19,6 +19,7 @@ from transformer_engine.pytorch.torch_version import torch_version from transformer_engine.pytorch.tensor.utils import clear_columnwise_cache, is_custom from .base import ( +from transformer_engine import te_platform, te_device_type fill_userbuffers_buffer_for_all_gather, get_ub, get_ub_is_fp8, @@ -966,7 +967,7 @@ def backward( # We use the send stream to copy into the userbuffers. # This is the same stream that we will use to access the data in the AG, # so we dont need to add any syncs yet. - with torch.cuda.stream(dgrad_send_stream): + with te_platform().stream(dgrad_send_stream): grad_output, _ = fill_userbuffers_buffer_for_all_gather( ub_obj_overlap_wgrad, grad_outputs[0], @@ -1249,7 +1250,7 @@ class LayerNormLinear(TransformerEngineBaseModule): .. math:: y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \varepsilon}} * (1 + \gamma) + \beta - device : Union[torch.device, str], default = "cuda" + device : Union[torch.device, str], default=te_device_type() The device on which the parameters of the model will be allocated. It is the user's responsibility to ensure all parameters are moved to the GPU before running the forward pass. @@ -1323,7 +1324,7 @@ def __init__( return_layernorm_output_gathered: bool = False, parameters_split: Optional[Union[Tuple[str, ...], Dict[str, int]]] = None, zero_centered_gamma: bool = False, - device: Union[torch.device, str] = "cuda", + device: Union[torch.device, str] = te_device_type(), ub_overlap_ag: bool = False, ub_overlap_rs: bool = False, ub_overlap_rs_dgrad: bool = False, diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index 3ee0cda50c5..de96326804f 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -20,6 +20,7 @@ from transformer_engine.pytorch.torch_version import torch_version from transformer_engine.pytorch.tensor.utils import clear_columnwise_cache, is_custom from .base import ( +from transformer_engine import te_platform, te_device_type fill_userbuffers_buffer_for_all_gather, _ub_communicators, get_ub, @@ -1332,7 +1333,7 @@ def backward( # We use the send stream to copy into the userbuffers. # This is the same stream that we will use to access the data in the AG, # so we dont need to add any syncs yet. - with torch.cuda.stream(dgrad_send_stream): + with te_platform().stream(dgrad_send_stream): grad_output, _ = fill_userbuffers_buffer_for_all_gather( ub_obj_fc2_wgrad, grad_outputs[0], @@ -1541,7 +1542,7 @@ def fc2_wgrad_gemm( reduce_scatter_out = None if ctx.ub_overlap_rs_dgrad: reduce_scatter_out = torch.empty( - fc1_dgrad_shape, dtype=ctx.activation_dtype, device="cuda" + fc1_dgrad_shape, dtype=ctx.activation_dtype, device=te_device_type() ) if ctx.ub_bulk_wgrad: gemm_out = ub_obj_fc1_wgrad.get_buffer(local_chunk=False) @@ -1637,7 +1638,7 @@ def fc2_wgrad_gemm( reduce_scatter_out = None if ctx.ub_bulk_wgrad and ub_obj_fc1_wgrad.is_fp8_ubuf(): reduce_scatter_out = torch.empty( - fc1_dgrad_shape, dtype=ctx.activation_dtype, device="cuda" + fc1_dgrad_shape, dtype=ctx.activation_dtype, device=te_device_type() ) # Arguments to include in wgrad GEMM closure @@ -1766,14 +1767,14 @@ def fc1_wgrad_gemm( fc1_wgrad = torch.zeros( fc1_weight_main_grad.shape, dtype=fc1_weight_python_object.dtype, - device=torch.cuda.current_device(), + device=te_platform().current_device(), requires_grad=False, ) else: fc1_wgrad = torch.empty( fc1_weight_main_grad.shape, dtype=fc1_weight_python_object.dtype, - device=torch.cuda.current_device(), + device=te_platform().current_device(), requires_grad=False, ) elif ctx.fuse_wgrad_accumulation: @@ -1791,14 +1792,14 @@ def fc1_wgrad_gemm( fc2_wgrad = torch.zeros( fc2_weight_main_grad.shape, dtype=fc2_weight_python_object.dtype, - device=torch.cuda.current_device(), + device=te_platform().current_device(), requires_grad=False, ) else: fc2_wgrad = torch.empty( fc2_weight_main_grad.shape, dtype=fc2_weight_python_object.dtype, - device=torch.cuda.current_device(), + device=te_platform().current_device(), requires_grad=False, ) elif ctx.fuse_wgrad_accumulation: @@ -1881,7 +1882,7 @@ class LayerNormMLP(TransformerEngineBaseModule): .. math:: y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \varepsilon}} * (1 + \gamma) + \beta - device : Union[torch.device, str], default = "cuda" + device : Union[torch.device, str], default=te_device_type() The device on which the parameters of the model will be allocated. It is the user's responsibility to ensure all parameters are moved to the GPU before running the forward pass. @@ -1970,7 +1971,7 @@ def __init__( micro_batch_size: Optional[int] = None, set_parallel_mode: bool = False, zero_centered_gamma: bool = False, - device: Union[torch.device, str] = "cuda", + device: Union[torch.device, str] = te_device_type(), ub_overlap_ag: bool = False, name: Optional[str] = None, ub_overlap_rs: bool = False, diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 94de69e9759..c931b32a8a4 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -19,6 +19,7 @@ from transformer_engine.pytorch.torch_version import torch_version from .base import ( +from transformer_engine import te_platform, te_device_type fill_userbuffers_buffer_for_all_gather, get_dummy_wgrad, get_ub, @@ -1539,7 +1540,7 @@ def _linear_backward_impl(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None # We use the send stream to copy into the userbuffers. # This is the same stream that we will use to access the data in the AG, # so we dont need to add any syncs yet. - with torch.cuda.stream(dgrad_send_stream): + with te_platform().stream(dgrad_send_stream): grad_output, _ = fill_userbuffers_buffer_for_all_gather( ub_obj_overlap_wgrad, grad_output_arg, @@ -1932,7 +1933,7 @@ class Linear(TransformerEngineBaseModule): values as split sizes along dim 0. The resulting parameters will have names that end in ``_weight`` or ``_bias``, so trailing underscores are stripped from any provided names. - device : Union[torch.device, str], default = "cuda" + device : Union[torch.device, str], default=te_device_type() The device on which the parameters of the model will be allocated. It is the user's responsibility to ensure all parameters are moved to the GPU before running the forward pass. @@ -2008,7 +2009,7 @@ def __init__( params_dtype: Optional[torch.dtype] = None, parallel_mode: Optional[str] = None, parameters_split: Optional[Union[Tuple[str, ...], Dict[str, int]]] = None, - device: Union[torch.device, str] = "cuda", + device: Union[torch.device, str] = te_device_type(), ub_overlap_ag: bool = False, ub_overlap_rs: bool = False, ub_overlap_rs_dgrad: bool = False, @@ -2319,7 +2320,7 @@ def _prealloc_cublas_workspace(self) -> None: from torch._subclasses.fake_tensor import is_fake weight = getattr(self, self.weight_names[0], None) - if weight is None or weight.device.type != "cuda": + if weight is None or weight.device.type != te_device_type(): return if is_fake(weight) or detect_fake_mode(): return diff --git a/transformer_engine/pytorch/ops/_common.py b/transformer_engine/pytorch/ops/_common.py index f39115c4c62..f48abeecbe6 100644 --- a/transformer_engine/pytorch/ops/_common.py +++ b/transformer_engine/pytorch/ops/_common.py @@ -10,6 +10,7 @@ import torch +from transformer_engine import te_device_type from transformer_engine_torch import FP8TensorMeta from ..torch_version import torch_version from ..quantization import FP8GlobalStateManager @@ -93,7 +94,7 @@ def maybe_dequantize( def maybe_autocast_dtype( *, - device_type: str = "cuda", + device_type: str = te_device_type(), default_dtype: Optional[torch.dtype] = None, ) -> torch.dtype: """Get autocast dtype if enabled""" diff --git a/transformer_engine/pytorch/ops/basic/activation.py b/transformer_engine/pytorch/ops/basic/activation.py index 26d5261b139..8bc825958ff 100644 --- a/transformer_engine/pytorch/ops/basic/activation.py +++ b/transformer_engine/pytorch/ops/basic/activation.py @@ -19,6 +19,7 @@ from ...utils import _compile_safe_warn, clear_tensor_data from ..op import BasicOperation, OperationContext from .._common import maybe_dequantize +from transformer_engine import te_device_type __all__ = [ "GELU", @@ -96,7 +97,7 @@ def op_forward( # Compute dtype dtype: torch.dtype if torch.is_autocast_enabled(): - dtype = torch.get_autocast_dtype("cuda") + dtype = torch.get_autocast_dtype(te_device_type()) else: dtype = input_.dtype if dtype not in (torch.float32, torch.float16, torch.bfloat16): @@ -413,7 +414,7 @@ def fuser_forward( ) if torch.is_autocast_enabled(): - dtype = torch.get_autocast_dtype("cuda") + dtype = torch.get_autocast_dtype(te_device_type()) elif isinstance(input_, torch.Tensor): dtype = input_.dtype else: diff --git a/transformer_engine/pytorch/ops/basic/basic_linear.py b/transformer_engine/pytorch/ops/basic/basic_linear.py index cb429055a4f..a04086d5eae 100644 --- a/transformer_engine/pytorch/ops/basic/basic_linear.py +++ b/transformer_engine/pytorch/ops/basic/basic_linear.py @@ -15,6 +15,7 @@ from ...cpp_extensions import general_gemm from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload from ...distributed import ( +from transformer_engine import te_device_type CudaRNGStatesTracker, gather_along_first_dim, reduce_scatter_along_first_dim, @@ -1016,7 +1017,7 @@ def op_forward( # Get autocast dtype if needed if torch.is_autocast_enabled(): - dtype = torch.get_autocast_dtype("cuda") + dtype = torch.get_autocast_dtype(te_device_type()) else: dtype = self.weight.dtype diff --git a/transformer_engine/pytorch/ops/basic/bias.py b/transformer_engine/pytorch/ops/basic/bias.py index 88f563b2c5b..c87800278d8 100644 --- a/transformer_engine/pytorch/ops/basic/bias.py +++ b/transformer_engine/pytorch/ops/basic/bias.py @@ -14,6 +14,7 @@ from ..op import BasicOperation, OperationContext from ...utils import canonicalize_device, canonicalize_dtype from ...tensor import Quantizer +from transformer_engine import te_device_type class Bias(BasicOperation): @@ -95,7 +96,7 @@ def reset_parameters(self) -> None: # Make sure parameter is initialized bias = self.bias - if bias.device.type != "cuda": + if bias.device.type != te_device_type(): bias = torch.empty_like(bias, device=self.device) else: bias = bias.to(device=self.device) diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index 95516500450..40587e5525b 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -20,6 +20,7 @@ from ...distributed import CudaRNGStatesTracker from ...module._common import WeightGradStore from ...module.base import ( +from transformer_engine import te_device_type _2X_ACC_FPROP, _2X_ACC_DGRAD, _2X_ACC_WGRAD, @@ -1030,7 +1031,7 @@ def fuser_forward( # Get autocast dtype if needed if torch.is_autocast_enabled(): - dtype = torch.get_autocast_dtype("cuda") + dtype = torch.get_autocast_dtype(te_device_type()) else: dtype = weight_param.dtype @@ -1182,7 +1183,7 @@ def fuser_forward_save_ctx( # through ``save_for_backward`` (see ``_fuser_forward_split_quantize`` # and ``_fuser_forward_grouped_tensor`` for the saved-tensor layout). if torch.is_autocast_enabled(): - ctx.dtype = torch.get_autocast_dtype("cuda") + ctx.dtype = torch.get_autocast_dtype(te_device_type()) else: ctx.dtype = weight_param.dtype ctx.input_requires_grad = requires_grad[0] diff --git a/transformer_engine/pytorch/ops/basic/l2normalization.py b/transformer_engine/pytorch/ops/basic/l2normalization.py index be155c93565..27d3a221844 100644 --- a/transformer_engine/pytorch/ops/basic/l2normalization.py +++ b/transformer_engine/pytorch/ops/basic/l2normalization.py @@ -13,6 +13,7 @@ from ...torch_version import torch_version from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload from ...jit import ( +from transformer_engine import te_platform l2normalization_fused, l2normalization_fwd_fused, l2normalization_backward_fused, @@ -64,7 +65,7 @@ def __init__( # JIT warmup for L2Normalization fused operations if seq_length and micro_batch_size: if ( - torch.cuda.is_available() + te_platform().is_available() and torch_version() >= (2, 0, 0) and bool(int(os.getenv("NVTE_TORCH_COMPILE", "1"))) ): diff --git a/transformer_engine/pytorch/ops/basic/swiglu.py b/transformer_engine/pytorch/ops/basic/swiglu.py index c19a09c2f6c..558ee7ce163 100644 --- a/transformer_engine/pytorch/ops/basic/swiglu.py +++ b/transformer_engine/pytorch/ops/basic/swiglu.py @@ -18,6 +18,7 @@ from ...utils import clear_tensor_data from ..op import BasicOperation, OperationContext from .._common import maybe_dequantize +from transformer_engine import te_device_type __all__ = [ "SwiGLU", @@ -116,7 +117,7 @@ def op_forward( # Compute dtype dtype: torch.dtype if torch.is_autocast_enabled(): - dtype = torch.get_autocast_dtype("cuda") + dtype = torch.get_autocast_dtype(te_device_type()) else: dtype = input_.dtype if dtype not in (torch.float32, torch.float16, torch.bfloat16): @@ -363,7 +364,7 @@ def op_forward( # Compute dtype dtype: torch.dtype if torch.is_autocast_enabled(): - dtype = torch.get_autocast_dtype("cuda") + dtype = torch.get_autocast_dtype(te_device_type()) else: dtype = input_.dtype if dtype not in (torch.float32, torch.float16, torch.bfloat16): @@ -530,7 +531,7 @@ def fuser_forward( # Determine compute dtype if torch.is_autocast_enabled(): - dtype = torch.get_autocast_dtype("cuda") + dtype = torch.get_autocast_dtype(te_device_type()) elif isinstance(input_, torch.Tensor): dtype = input_.dtype else: diff --git a/transformer_engine/pytorch/ops/fused/forward_linear_bias_activation.py b/transformer_engine/pytorch/ops/fused/forward_linear_bias_activation.py index 6fa63675b4c..d58eb5b6f69 100644 --- a/transformer_engine/pytorch/ops/fused/forward_linear_bias_activation.py +++ b/transformer_engine/pytorch/ops/fused/forward_linear_bias_activation.py @@ -15,6 +15,7 @@ from ...tensor import Quantizer from ..basic import BasicLinear, Bias from ..op import FusedOperation, FusibleOperation, OperationContext +from transformer_engine import te_device_type class ForwardLinearBiasActivation(FusedOperation): @@ -99,7 +100,7 @@ def fuser_forward( # Get autocast dtype if needed if torch.is_autocast_enabled(): - dtype = torch.get_autocast_dtype("cuda") + dtype = torch.get_autocast_dtype(te_device_type()) else: dtype = linear_op.weight.dtype diff --git a/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py b/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py index 28586360f59..25594c63008 100644 --- a/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py +++ b/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py @@ -15,6 +15,7 @@ from ...tensor import Quantizer from ..basic import AddExtraInput, BasicLinear, Bias from ..op import FusedOperation, FusibleOperation, OperationContext +from transformer_engine import te_device_type class ForwardLinearBiasAdd(FusedOperation): @@ -93,7 +94,7 @@ def fuser_forward( # Get autocast dtype if needed if torch.is_autocast_enabled(): - dtype = torch.get_autocast_dtype("cuda") + dtype = torch.get_autocast_dtype(te_device_type()) else: dtype = linear_op.weight.dtype diff --git a/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py b/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py index 277263e0ecb..2d28f628c3e 100644 --- a/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py +++ b/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py @@ -15,6 +15,7 @@ from ...tensor import Quantizer from ..basic import AddExtraInput, BasicLinear, ConstantScale from ..op import ( +from transformer_engine import te_device_type FusedOperation, FusibleOperation, OperationContext, @@ -75,7 +76,7 @@ def fuser_forward( # Get autocast dtype if needed if torch.is_autocast_enabled(): - dtype = torch.get_autocast_dtype("cuda") + dtype = torch.get_autocast_dtype(te_device_type()) else: dtype = linear_op.weight.dtype diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index bd8b82d3faa..59bfc895958 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -21,6 +21,7 @@ from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload, start_offload from ...cpp_extensions import general_gemm, general_grouped_gemm_for_grouped_tensor from ...distributed_weight import ( +from transformer_engine import te_platform is_distributed_weight, materialize_weight_for_forward, materialize_weight_for_backward, @@ -548,7 +549,7 @@ def _cudnn_compute_wgrad( b = X = (total_tokens, in_features) column-major. """ if current_stream is None: - current_stream = torch.cuda.current_stream().cuda_stream + current_stream = te_platform().current_stream().cuda_stream out_features, in_features = weight_shape total_tokens = grouped_dy.logical_shape[0] @@ -1149,7 +1150,7 @@ def fuser_forward( device = fc1_weight_param.device if torch.is_autocast_enabled(): - dtype = torch.get_autocast_dtype("cuda") + dtype = torch.get_autocast_dtype(te_device_type()) else: dtype = fc1_weight_param.dtype diff --git a/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py b/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py index c1070e38a62..69a20a8dbc7 100644 --- a/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py +++ b/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py @@ -14,6 +14,7 @@ from ...cpp_extensions import general_gemm from ...distributed import gather_along_first_dim, get_distributed_world_size from ...module.base import ( +from transformer_engine import te_platform _2X_ACC_DGRAD, _2X_ACC_WGRAD, fill_userbuffers_buffer_for_all_gather, @@ -181,7 +182,7 @@ def _functional_backward( else: device = grad_output.device device = canonicalize_device(device) - if device.type != "cuda": + if device.type != te_device_type(): raise ValueError(f"Only CUDA devices are supported (got {device})") # Check datatype @@ -481,7 +482,7 @@ def _functional_backward( # We use the send stream to copy into the userbuffers. # This is the same stream that we will use to access the data in the AG, # so we dont need to add any syncs yet. - with torch.cuda.stream(dgrad_send_stream): + with te_platform().stream(dgrad_send_stream): dy, _ = fill_userbuffers_buffer_for_all_gather( ub_obj_overlap_wgrad, dy_local, diff --git a/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py b/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py index cb5f8a16e61..8750aed87e8 100644 --- a/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py +++ b/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py @@ -16,6 +16,7 @@ from ...distributed import get_distributed_world_size from ...quantization import FP8GlobalStateManager from ...module.base import ( +from transformer_engine import te_device_type fill_userbuffers_buffer_for_all_gather, get_ub, _2X_ACC_FPROP, @@ -155,7 +156,7 @@ def _functional_forward( """ # Check device - if device.type != "cuda": + if device.type != te_device_type(): raise ValueError(f"Only CUDA devices are supported (got {device})") # Check datatype @@ -322,7 +323,7 @@ def fuser_forward( # Get autocast dtype if needed if torch.is_autocast_enabled(): - dtype = torch.get_autocast_dtype("cuda") + dtype = torch.get_autocast_dtype(te_device_type()) else: dtype = linear_op.weight.dtype diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py index d057d46816e..57b52398e45 100644 --- a/transformer_engine/pytorch/ops/op.py +++ b/transformer_engine/pytorch/ops/op.py @@ -16,6 +16,7 @@ from transformer_engine.common.recipe import Recipe from .._extra_state import is_stateless_recipe, should_load_extra_state_pickle from ..quantization import ( +from transformer_engine import te_platform FP8GlobalStateManager, QuantizerRole, RecipeState, @@ -703,7 +704,7 @@ def to_cpu(src: torch.Tensor) -> torch.Tensor: return torch.empty(0, dtype=torch.uint8) # Serialize state into byte tensor - torch.cuda.synchronize() + te_platform().synchronize() state_serialized = bytearray(pickle.dumps(state)) state_serialized = torch.frombuffer(state_serialized, dtype=torch.uint8) return state_serialized @@ -766,7 +767,7 @@ def copy_tensor(src: torch.Tensor, dst: torch.Tensor) -> None: ) # Finish CPU-GPU memory transfers - torch.cuda.synchronize() + te_platform().synchronize() def _load_from_state_dict(self, *args, **kwargs) -> None: """Load state""" diff --git a/transformer_engine/pytorch/optimizers/fused_adam.py b/transformer_engine/pytorch/optimizers/fused_adam.py index b4b7f707a83..fd7eb24774a 100644 --- a/transformer_engine/pytorch/optimizers/fused_adam.py +++ b/transformer_engine/pytorch/optimizers/fused_adam.py @@ -17,6 +17,7 @@ from transformer_engine.pytorch.quantized_tensor import QuantizedTensor from ..constants import DType from .multi_tensor_apply import multi_tensor_applier +from transformer_engine import te_platform # Bound temporary NVTETensor handles created before the CUDA launcher chunks its metadata. @@ -190,7 +191,7 @@ def __init__( self._step_supports_amp_scaling = True # Skip buffer - self._dummy_overflow_buf = torch.tensor([0], dtype=torch.int, device="cuda") + self._dummy_overflow_buf = torch.tensor([0], dtype=torch.int, device=te_device_type()) self.multi_tensor_adam = tex.multi_tensor_adam self.multi_tensor_adam_param_remainder = tex.multi_tensor_adam_param_remainder self.multi_tensor_adam_fp8 = tex.multi_tensor_adam_fp8 @@ -553,7 +554,7 @@ def step(self, closure=None, grad_scaler=None): Arguments: closure (callable, optional): A closure that reevaluates the model and returns the loss. - grad_scaler (torch.cuda.amp.GradScaler, optional): + grad_scaler (te_platform().amp.GradScaler, optional): gradient scaler (default: None) """ loss = None diff --git a/transformer_engine/pytorch/optimizers/newton_schulz.py b/transformer_engine/pytorch/optimizers/newton_schulz.py index 4f868f1c127..7d2edd6cd6c 100644 --- a/transformer_engine/pytorch/optimizers/newton_schulz.py +++ b/transformer_engine/pytorch/optimizers/newton_schulz.py @@ -11,6 +11,7 @@ import torch.distributed as dist import transformer_engine_torch as tex +from transformer_engine import te_platform, te_device_type _COEFFICIENT_SETS = { @@ -181,8 +182,8 @@ def _get_nccl_comm_ptr(group: dist.ProcessGroup) -> int: # ProcessGroupNCCL creates communicators lazily. This device-specific # collective ensures that the communicator returned by _comm_ptr() exists # and is ready on every group rank before cuSolverMp borrows it. - dist.barrier(group=group, device_ids=[torch.cuda.current_device()]) - nccl_backend = group._get_backend(torch.device("cuda")) + dist.barrier(group=group, device_ids=[te_platform().current_device()]) + nccl_backend = group._get_backend(torch.device(te_device_type())) comm_ptr = nccl_backend._comm_ptr() if not isinstance(comm_ptr, int) or comm_ptr == 0: raise RuntimeError("ProcessGroupNCCL returned an invalid communicator pointer") diff --git a/transformer_engine/pytorch/permutation.py b/transformer_engine/pytorch/permutation.py index 5e6bdb4dfcd..50880bcd7f7 100644 --- a/transformer_engine/pytorch/permutation.py +++ b/transformer_engine/pytorch/permutation.py @@ -7,6 +7,7 @@ from typing import Optional, Tuple import torch import transformer_engine_torch as tex +from transformer_engine import te_device_type import transformer_engine.pytorch.triton.permutation as triton_permutation from transformer_engine.pytorch.constants import TE_DType from transformer_engine.pytorch.quantized_tensor import ( @@ -44,10 +45,14 @@ def moe_permute_index_map_forward( if not inp.numel(): return inp.clone(), torch.empty(0, dtype=torch.int32, device=inp.device) - if not inp.is_cuda: - raise ValueError(f"inp must be a CUDA tensor, but got tensor on {inp.device}.") - if not index.is_cuda: - raise ValueError(f"index must be a CUDA tensor, but got tensor on {index.device}.") + if inp.device.type != te_device_type(): + raise ValueError( + f"inp must be a {te_device_type()} tensor, but got tensor on {inp.device}." + ) + if index.device.type != te_device_type(): + raise ValueError( + f"index must be a {te_device_type()} tensor, but got tensor on {index.device}." + ) if inp.size(0) != index.size(0): raise ValueError( f"Permute not possible: inp.size(0) ({inp.size(0)}) must match " @@ -299,19 +304,25 @@ def moe_permute_mask_map_forward( ) return inp.clone(), row_id_map, permuted_probs - if not inp.is_cuda: - raise ValueError(f"inp must be a CUDA tensor, but got tensor on {inp.device}.") - if not routing_map.is_cuda: + if inp.device.type != te_device_type(): raise ValueError( - f"routing_map must be a CUDA tensor, but got tensor on {routing_map.device}." + f"inp must be a {te_device_type()} tensor, but got tensor on {inp.device}." + ) + if routing_map.device.type != te_device_type(): + raise ValueError( + f"routing_map must be a {te_device_type()} tensor, but got tensor on" + f" {routing_map.device}." ) if probs is not None: - if not probs.is_cuda: - raise ValueError(f"probs must be a CUDA tensor, but got tensor on {probs.device}.") + if probs.device.type != te_device_type(): + raise ValueError( + f"probs must be a {te_device_type()} tensor, but got tensor on {probs.device}." + ) if pad_offsets is not None: - if not pad_offsets.is_cuda: + if pad_offsets.device.type != te_device_type(): raise ValueError( - f"pad_offsets must be a CUDA tensor, but got tensor on {pad_offsets.device}." + f"pad_offsets must be a {te_device_type()} tensor, but got tensor on" + f" {pad_offsets.device}." ) if inp.size(0) != routing_map.size(0): raise ValueError( diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index 34ae20b4984..f097aed0bc5 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -17,6 +17,7 @@ import torch import transformer_engine_torch as tex from transformer_engine.common.recipe import ( +from transformer_engine import te_platform, te_device_type Recipe, DelayedScaling, Format, @@ -424,7 +425,7 @@ def set_skip_fp8_weight_update_tensor(cls, skip: bool) -> None: """Set the skip fp8 weight update tensor""" if cls.quantization_state.skip_fp8_weight_update_tensor is None: cls.quantization_state.skip_fp8_weight_update_tensor = torch.empty( - 1, dtype=torch.float32, device="cuda" + 1, dtype=torch.float32, device=te_device_type() ) cls.quantization_state.skip_fp8_weight_update_tensor.fill_(skip) @@ -591,7 +592,7 @@ def fp8_graph_capturing(cls) -> bool: assert not cls.quantization_state.fp8_graph_capturing return False return ( - cls.quantization_state.fp8_graph_capturing or torch.cuda.is_current_stream_capturing() + cls.quantization_state.fp8_graph_capturing or te_platform().is_current_stream_capturing() ) @classmethod @@ -1432,7 +1433,7 @@ def __init__( # Allocate buffers if device is None: - device = torch.device("cuda") + device = torch.device(te_device_type()) self.scale = torch.ones(num_quantizers, dtype=torch.float32, device=device) self.amax_history = torch.zeros( recipe.amax_history_len, @@ -1481,7 +1482,7 @@ def __init__( # Allocate buffers if device is None: - device = torch.device("cuda") + device = torch.device(te_device_type()) self.device = device def make_quantizers(self) -> list: @@ -1524,7 +1525,7 @@ def __init__( # Allocate buffers if device is None: - device = torch.device("cuda") + device = torch.device(te_device_type()) def make_quantizers(self) -> list: # TODO(ksivamani); Find better design for this, adding here to avoid circular import. @@ -1582,7 +1583,7 @@ def __init__( # Allocate buffers if device is None: - device = torch.device("cuda") + device = torch.device(te_device_type()) self.device = device def make_quantizers(self) -> list: @@ -1665,7 +1666,7 @@ def __init__( # Allocate buffers if device is None: - device = torch.device("cuda") + device = torch.device(te_device_type()) def make_quantizers(self) -> list: """Build one ``NVFP4Quantizer`` per slot, dispatched by tensor type. @@ -1896,7 +1897,7 @@ def __init__( self.num_quantizers = num_quantizers self.roles = roles if device is None: - device = torch.device("cuda") + device = torch.device(te_device_type()) self.device = device # -- Stateful sub-state slots (initialized empty) -- diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index 7149a5a1633..b1f6633d54f 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -18,6 +18,7 @@ from transformer_engine.common.recipe import Recipe from transformer_engine.pytorch.constants import dist_group_type from transformer_engine.pytorch.tensor._quantization_helpers import ( +from transformer_engine import te_platform, te_device_type _QuantizeFunc, _IdentityFunc, _stride_from_shape, @@ -484,7 +485,7 @@ def make_empty( ) if device is None: - device = torch.device("cuda") + device = torch.device(te_device_type()) # Handle the device passed as string device = torch.device(device) result = tex.create_empty_quantized_tensor( @@ -540,7 +541,7 @@ def alloc_tensors( Returns ``{attr_name: torch.Tensor}`` suitable as the ``inner_tensors`` argument of the storage's ``__tensor_unflatten__``. """ - device = torch.device(device if device is not None else "cuda") + device = torch.device(device if device is not None else te_device_type()) return { attr: torch.empty(buf_shape, dtype=buf_dtype, device=device) for attr, (buf_shape, buf_dtype) in self.inner_tensor_specs(tuple(shape)).items() @@ -740,7 +741,7 @@ def __new__( dtype=dtype, layout=torch.strided, requires_grad=requires_grad, - device=torch.cuda.current_device() if device is None else device, + device=te_platform().current_device() if device is None else device, ) instance._requires_grad = requires_grad instance._dtype = dtype diff --git a/transformer_engine/pytorch/router.py b/transformer_engine/pytorch/router.py index ac38f7d1cf9..ba156cfc6d8 100644 --- a/transformer_engine/pytorch/router.py +++ b/transformer_engine/pytorch/router.py @@ -19,6 +19,7 @@ import torch import transformer_engine_torch as tex +from transformer_engine import te_platform # Re-export the C++ enum NVTERoutingMapFormat under a friendlier Python name. # Members: @@ -57,8 +58,8 @@ def _validate_qb_bin_bounds(bin_bounds: torch.Tensor) -> bool: version = bin_bounds._version if getattr(bin_bounds, _QB_BOUNDS_VALIDATED_VERSION_ATTR, None) == version: return True - with torch.cuda.device(bin_bounds.device): - if torch.cuda.is_current_stream_capturing(): + with te_platform().device(bin_bounds.device): + if te_platform().is_current_stream_capturing(): raise RuntimeError( "QB bin_bounds must be validated by an eager router call before CUDA graph capture" ) diff --git a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py index 6105b18a736..c85577fe2f8 100644 --- a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py @@ -11,6 +11,8 @@ import torch import transformer_engine_torch as tex + +from transformer_engine import te_device_type from transformer_engine.common.recipe import Float8BlockScaling, Recipe from .storage.float8_blockwise_tensor_storage import ( Float8BlockwiseQTensorStorage, @@ -441,7 +443,7 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): qt._rowwise_scale_inv, qt._columnwise_scale_inv, ): - if t is not None and t.is_cuda: + if t is not None and t.device.type == te_device_type(): t.record_stream(stream) return None @@ -538,7 +540,7 @@ def _set_data(self, tensor: torch.Tensor) -> None: """ # Tensor device - new_device = tensor.device if tensor.is_cuda else self.device + new_device = tensor.device if tensor.device.type == te_device_type() else self.device def _set_from_tensor(dst: Float8BlockwiseQTensor, src: Float8BlockwiseQTensor): dst._rowwise_data = src._rowwise_data diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index cf37c36c598..d37b8e2800e 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -10,6 +10,7 @@ from torch.distributed.fsdp._fully_shard._fsdp_common import TrainingState import transformer_engine_torch as tex +from transformer_engine import te_device_type from transformer_engine.common.recipe import ( DelayedScaling, Float8CurrentScaling, @@ -1139,7 +1140,7 @@ def _set_data(self, tensor: torch.Tensor) -> None: """ # Tensor device - new_device = tensor.device if tensor.is_cuda else self.device + new_device = tensor.device if tensor.device.type == te_device_type() else self.device if not devices_match(new_device, tensor.device): tensor = tensor.to(device=new_device) diff --git a/transformer_engine/pytorch/tensor/grouped_tensor.py b/transformer_engine/pytorch/tensor/grouped_tensor.py index 0cc03602a1f..bcb6dee6b69 100644 --- a/transformer_engine/pytorch/tensor/grouped_tensor.py +++ b/transformer_engine/pytorch/tensor/grouped_tensor.py @@ -12,6 +12,7 @@ from ..quantized_tensor import QuantizedTensorStorage, Quantizer from .storage.grouped_tensor_storage import GroupedTensorStorage +from transformer_engine import te_device_type def _stride_from_shape(shape: Tuple[int, ...]) -> Tuple[int, ...]: @@ -131,7 +132,7 @@ def __new__( device = maybe_tensor.device break if device is None: - device = torch.device("cuda") + device = torch.device(te_device_type()) # Match QuantizedTensor __new__: accept externally-computed stride to # avoid Python-side stride computation overhead for C++ construction. diff --git a/transformer_engine/pytorch/tensor/identity_tensor.py b/transformer_engine/pytorch/tensor/identity_tensor.py index ec171564fe2..b4fda2d00d2 100644 --- a/transformer_engine/pytorch/tensor/identity_tensor.py +++ b/transformer_engine/pytorch/tensor/identity_tensor.py @@ -21,6 +21,7 @@ from .storage.identity_tensor_storage import IdentityTensorStorage from ..quantized_tensor import QuantizedTensor, QuantizedTensorStorage, Quantizer +from transformer_engine import te_device_type class IdentityQuantizer(Quantizer): @@ -114,7 +115,7 @@ def make_empty( pin_memory: bool = False, ) -> Union["IdentityTensor", IdentityTensorStorage]: if device is None: - device = torch.device("cuda") + device = torch.device(te_device_type()) device = torch.device(device) data_dtype = self.dtype if self.dtype is not None else dtype data = torch.empty(tuple(shape), dtype=data_dtype, device=device, pin_memory=pin_memory) diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py index 267806e43ee..c4081b75c64 100644 --- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py +++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py @@ -13,6 +13,7 @@ from torch.distributed.fsdp._fully_shard._fsdp_common import TrainingState import transformer_engine_torch as tex +from transformer_engine import te_device_type from transformer_engine.common.recipe import MXFP8BlockScaling, Recipe from ..constants import MXFP8_BLOCK_SCALING_SIZE, DType from ..utils import devices_match, round_up_to_nearest_multiple @@ -823,7 +824,7 @@ def _set_data(self, tensor: torch.Tensor) -> None: """ # Tensor device - new_device = tensor.device if tensor.is_cuda else self.device + new_device = tensor.device if tensor.device.type == te_device_type() else self.device if not devices_match(new_device, tensor.device): tensor = tensor.to(device=new_device) diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index 5e537c3ff0e..dbcc016e40b 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -16,6 +16,7 @@ from transformer_engine.common.recipe import NVFP4BlockScaling, Recipe from ..constants import NVFP4_BLOCK_SCALING_SIZE, dist_group_type, DType from ..utils import ( +from transformer_engine import te_platform canonicalize_process_group, devices_match, round_up_to_nearest_multiple, @@ -173,7 +174,7 @@ def __init__( if self.nvfp4_4over6_err_mode not in ("MAE", "MSE"): raise ValueError("nvfp4_4over6_err_mode must be 'MAE' or 'MSE'.") self.rht_matrix_random_sign_mask_t = get_random_sign_mask_for_rht( - with_random_sign_mask, torch.cuda.current_device() + with_random_sign_mask, te_platform().current_device() ) self._rebuild_derived_state() @@ -193,7 +194,7 @@ def _rebuild_derived_state(self) -> None: cheap hit. """ self.rht_matrix = get_rht_matrix( - self.rht_matrix_random_sign_mask_t != 0, torch.cuda.current_device() + self.rht_matrix_random_sign_mask_t != 0, te_platform().current_device() ) def update_quantized( diff --git a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py index 3473024c03f..bea825bf124 100644 --- a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py @@ -19,6 +19,7 @@ from .float8_blockwise_tensor_storage import Float8BlockwiseQTensorStorage from .nvfp4_tensor_storage import NVFP4TensorStorage from ...utils import is_non_tn_fp8_gemm_supported +from transformer_engine import te_platform class GroupedTensorStorage: @@ -668,7 +669,7 @@ def make_grouped_tensor( # Set device if device is None: - device = torch.cuda.current_device() + device = te_platform().current_device() # Shape patterns and validation. all_same_first = first_dims is None @@ -713,7 +714,7 @@ def make_grouped_tensor( compatible_recipe = None if no_quantization else quantizer._get_compatible_recipe() if not shape: - if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): + if te_platform().is_available() and te_platform().is_current_stream_capturing(): raise ValueError( "Varying-dimension grouped tensor construction is not graph-safe: it" " requires device-to-host copies of per-tensor shapes/offsets." diff --git a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py index f4a95607d48..df4b8ae0e38 100644 --- a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py @@ -16,6 +16,7 @@ from .._quantization_helpers import safe_quantized_repr from ...constants import ( +from transformer_engine import te_device_type TE_DType as torch_to_transformer_engine_dtype, MXFP8_BLOCK_SCALING_SIZE, DType, @@ -44,8 +45,8 @@ def forward( te_dtype = torch_to_transformer_engine_dtype[dtype] # ``tex.dequantize`` requires CUDA-resident buffers. src_device = tensor.device - if src_device.type != "cuda": - cuda_tensor = tensor.to(device=torch.device("cuda")) + if src_device.type != te_device_type(): + cuda_tensor = tensor.to(device=torch.device(te_device_type())) result = tex.dequantize(cuda_tensor, te_dtype) return result.to(device=src_device) return tex.dequantize(tensor, te_dtype) diff --git a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py index 7e0861c967b..83b5380ce8e 100644 --- a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py @@ -20,6 +20,7 @@ from ...constants import TE_DType as torch_to_transformer_engine_dtype, DType from ...utils import _empty_tensor +from transformer_engine import te_device_type @functools.lru_cache(maxsize=None) @@ -54,8 +55,8 @@ def forward( # ``tex.dequantize`` requires CUDA-resident buffers. If the tensor has src_device = tensor.device - if src_device.type != "cuda": - cuda_tensor = tensor.to(device=torch.device("cuda")) + if src_device.type != te_device_type(): + cuda_tensor = tensor.to(device=torch.device(te_device_type())) result = tex.dequantize(cuda_tensor, torch_to_transformer_engine_dtype[dtype]) return result.to(device=src_device) return tex.dequantize(tensor, torch_to_transformer_engine_dtype[dtype]) diff --git a/transformer_engine/pytorch/transformer.py b/transformer_engine/pytorch/transformer.py index db9a352881c..cc728f2d5e1 100644 --- a/transformer_engine/pytorch/transformer.py +++ b/transformer_engine/pytorch/transformer.py @@ -15,6 +15,7 @@ from transformer_engine.pytorch.attention.multi_head_attention import MultiheadAttention from transformer_engine.pytorch.attention.inference import InferenceParams from transformer_engine.pytorch.jit import ( +from transformer_engine import te_device_type set_jit_fusion_options, warmup_jit_bias_dropout_add_all_dtypes, get_bias_dropout_add, @@ -199,7 +200,7 @@ class TransformerLayer(torch.nn.Module): supports ``'limit'``, ``'alpha'``, and ``'glu_linear_offset'`` parameters. You can set these as ``activation_params={'limit': 7.0, 'alpha': 1.702, 'glu_linear_offset': 1.0}``. - device : Union[torch.device, str], default = "cuda" + device : Union[torch.device, str], default=te_device_type() The device on which the parameters of the model will be allocated. It is the user's responsibility to ensure all parameters are moved to the GPU before running the forward pass. @@ -352,7 +353,7 @@ def __init__( activation: str = "gelu", activation_params: Optional[dict] = None, normalization: str = "LayerNorm", - device: Union[torch.device, str] = "cuda", + device: Union[torch.device, str] = te_device_type(), attn_input_format: str = "sbhd", name: str = None, qk_norm_type: Optional[str] = None, diff --git a/transformer_engine/pytorch/triton/mhc.py b/transformer_engine/pytorch/triton/mhc.py index 41f71966a23..59ee8d0b68f 100644 --- a/transformer_engine/pytorch/triton/mhc.py +++ b/transformer_engine/pytorch/triton/mhc.py @@ -10,6 +10,7 @@ import triton from transformer_engine.common.triton.mhc import ( +from transformer_engine import te_platform _mhc_projection_bwd_fused_dphi, _mhc_projection_bwd_fused_dx, _mhc_scale_fwd_fused, @@ -31,7 +32,7 @@ def _support_tma(x: torch.Tensor): # get_device_capability returns a (major, minor) tuple; TMA needs Hopper+ (major >= 9) - return torch.cuda.get_device_capability(x.device)[0] >= 9 + return te_platform().get_device_capability(x.device)[0] >= 9 def _tma_aligned(t): @@ -50,7 +51,7 @@ def _init_tma_allocator(): def alloc_fn( size: int, alignment: int, stream: Optional[int] ): # pylint: disable=unused-argument - return torch.empty(size, device="cuda", dtype=torch.int8) + return torch.empty(size, device=te_device_type(), dtype=torch.int8) triton.set_allocator(alloc_fn) _tma_allocator_initialized = True diff --git a/transformer_engine/pytorch/triton/permutation.py b/transformer_engine/pytorch/triton/permutation.py index c155d73e1e3..d4c28f56d01 100644 --- a/transformer_engine/pytorch/triton/permutation.py +++ b/transformer_engine/pytorch/triton/permutation.py @@ -10,6 +10,7 @@ import triton from transformer_engine.common.triton.permutation import ( +from transformer_engine import te_device_type _row_id_map_pass_1_kernel, _row_id_map_pass_2_kernel, _row_id_map_pass_3_kernel, @@ -49,10 +50,10 @@ def make_row_id_map( The [num_experts, num_experts + n_routed) items are the indices of the experts corresponding to the first n_routed row indices above. """ - row_id_map = torch.empty((num_tokens, num_experts * 2 + 1), dtype=torch.int32, device="cuda") + row_id_map = torch.empty((num_tokens, num_experts * 2 + 1), dtype=torch.int32, device=te_device_type()) block_size = 1024 grid = (num_experts, triton.cdiv(num_tokens, block_size)) - workspace_tensor = torch.empty(grid, dtype=torch.int32, device="cuda") + workspace_tensor = torch.empty(grid, dtype=torch.int32, device=te_device_type()) # supposing num_tokens == 5, num_experts == 3, block_size == 3 # and we have a routing_map like this: @@ -160,12 +161,12 @@ def permute_with_mask_map( # Use torch.zeros when pad_offsets is provided to ensure padding regions are zeroed. # The kernel writes only to valid positions, leaving padding positions at zero. alloc = torch.zeros if pad_offsets is not None else torch.empty - output = alloc((num_out_tokens, hidden_size), dtype=inp.dtype, device="cuda") + output = alloc((num_out_tokens, hidden_size), dtype=inp.dtype, device=te_device_type()) permuted_probs = ( - alloc((num_out_tokens,), dtype=probs.dtype, device="cuda") if probs is not None else None + alloc((num_out_tokens,), dtype=probs.dtype, device=te_device_type()) if probs is not None else None ) permuted_scale = ( - alloc((num_out_tokens, scale_hidden_dim), dtype=scale.dtype, device="cuda") + alloc((num_out_tokens, scale_hidden_dim), dtype=scale.dtype, device=te_device_type()) if scale is not None else None ) @@ -243,10 +244,10 @@ def unpermute_with_mask_map( hidden_size : int Hidden size of the permuted tensor. """ - output = torch.empty((num_tokens, hidden_size), dtype=inp.dtype, device="cuda") + output = torch.empty((num_tokens, hidden_size), dtype=inp.dtype, device=te_device_type()) if permuted_probs is not None: unpermuted_probs = torch.empty( - (num_tokens, num_experts), dtype=permuted_probs.dtype, device="cuda" + (num_tokens, num_experts), dtype=permuted_probs.dtype, device=te_device_type() ) else: unpermuted_probs = None @@ -325,9 +326,9 @@ def unpermute_with_mask_map_bwd_with_merging_probs( # by the kernel. This matches the behavior of Fp8Unpadding.backward which zeros # out the padding slots. alloc = torch.zeros if pad_offsets is not None else torch.empty - act_grad = alloc((num_out_tokens, hidden_size), dtype=fwd_output_grad.dtype, device="cuda") + act_grad = alloc((num_out_tokens, hidden_size), dtype=fwd_output_grad.dtype, device=te_device_type()) merging_probs_grad = torch.empty( - (num_tokens, num_experts), dtype=merging_probs.dtype, device="cuda" + (num_tokens, num_experts), dtype=merging_probs.dtype, device=te_device_type() ) grid = (num_tokens,) _unpermute_bwd_with_merging_probs_kernel[grid]( @@ -378,7 +379,7 @@ def make_chunk_sort_map( num_splits : int Number of splits of split_sizes and sorted_indices. """ - row_id_map = torch.empty((num_tokens,), dtype=torch.int32, device="cuda") + row_id_map = torch.empty((num_tokens,), dtype=torch.int32, device=te_device_type()) grid = (num_tokens,) _make_chunk_sort_map_kernel[grid]( split_sizes, @@ -416,9 +417,9 @@ def sort_chunks_by_map( is_forward : bool Whether the sort is for forward or backward. """ - output = torch.empty((num_tokens, hidden_size), dtype=inp.dtype, device="cuda") + output = torch.empty((num_tokens, hidden_size), dtype=inp.dtype, device=te_device_type()) if probs is not None: - permuted_probs = torch.empty((num_tokens,), dtype=probs.dtype, device="cuda") + permuted_probs = torch.empty((num_tokens,), dtype=probs.dtype, device=te_device_type()) else: permuted_probs = None # pylint: disable=unnecessary-lambda-assignment diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py index 513632e095a..09548f1d395 100644 --- a/transformer_engine/pytorch/utils.py +++ b/transformer_engine/pytorch/utils.py @@ -16,6 +16,7 @@ from .torch_version import torch_version from ..debug.pytorch.debug_quantization import DebugQuantizedTensor +from transformer_engine import te_platform __all__ = [ @@ -157,7 +158,7 @@ def _get_device_compute_capability(device: torch.device) -> Tuple[int, int]: @torch.compiler.assume_constant_result def get_device_compute_capability() -> Tuple[int, int]: """CUDA compute capability of current GPU""" - return _get_device_compute_capability(torch.cuda.current_device()) + return _get_device_compute_capability(te_platform().current_device()) def deinterleave_glu_tensor(tensor: torch.Tensor, interleave_size: int) -> torch.Tensor: @@ -346,7 +347,7 @@ def all_close(a: torch.Tensor, b: torch.Tensor) -> bool: def print_rank_0(*args: Any) -> None: """print on rank 0""" - if torch.cuda.current_device() == 0: + if te_platform().current_device() == 0: print(*args) @@ -801,12 +802,12 @@ def canonicalize_device(device: Optional[torch.device | str]) -> torch.device: if device is None: # Use default CUDA device device = torch.get_default_device() - if device.type != "cuda": - device = torch.device("cuda", torch.cuda.current_device()) + if device.type != te_device_type(): + device = torch.device(te_device_type(), te_platform().current_device()) elif not isinstance(device, torch.device): device = torch.device(device) - if device.type == "cuda" and device.index is None: - device = torch.device("cuda", torch.cuda.current_device()) + if device.type == te_device_type() and device.index is None: + device = torch.device(te_device_type(), te_platform().current_device()) return device @@ -828,15 +829,15 @@ def devices_match(device1: torch.device, device2: torch.device) -> bool: device2 = torch.device(device2) if device1.type != device2.type: return False - if device1.type == "cuda": + if device1.type == te_device_type(): index1 = device1.index index2 = device2.index if index1 == index2: return True if index1 is None: - index1 = torch.cuda.current_device() + index1 = te_platform().current_device() if index2 is None: - index2 = torch.cuda.current_device() + index2 = te_platform().current_device() return index1 == index2 return device1 == device2 @@ -844,7 +845,7 @@ def devices_match(device1: torch.device, device2: torch.device) -> bool: @functools.lru_cache def get_sm_count() -> int: """Returns the number of streaming multiprocessors in the current device.""" - return torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count + return torch.cuda.get_device_properties(te_platform().current_device()).multi_processor_count def ceil_div(numerator, denominator): @@ -898,7 +899,7 @@ def get_nvtx_range_context(msg: str): """ if _nvtx_enabled(): - return torch.cuda.nvtx.range(msg) + return te_platform().nvtx.range(msg) return nullcontext() @@ -917,7 +918,7 @@ def nvtx_range_push(msg: str) -> None: if not _nvtx_enabled(): return _nvtx_range_messages.append(msg) - torch.cuda.nvtx.range_push(msg) + te_platform().nvtx.range_push(msg) def nvtx_range_pop(msg: Optional[str] = None) -> None: @@ -948,7 +949,7 @@ def nvtx_range_pop(msg: Optional[str] = None) -> None: ) # Pop NVTX range - torch.cuda.nvtx.range_pop() + te_platform().nvtx.range_pop() def canonicalize_process_group( @@ -967,14 +968,14 @@ def canonicalize_process_group( def torch_get_autocast_gpu_dtype() -> torch.dtype: """Get PyTorch autocast GPU dtype.""" if torch_version() >= (2, 4, 0): - return torch.get_autocast_dtype("cuda") + return torch.get_autocast_dtype(te_device_type()) return torch.get_autocast_gpu_dtype() if torch_version() >= (2, 4, 0): - gpu_autocast_ctx = functools.partial(torch.amp.autocast, device_type="cuda") + gpu_autocast_ctx = functools.partial(torch.amp.autocast, device_type=te_device_type()) else: - gpu_autocast_ctx = torch.cuda.amp.autocast + gpu_autocast_ctx = te_platform().amp.autocast _torch_dtype_to_np_typestr_dict = { From c8feb1a0451c9c5853b4d31d21e1246d1688514d Mon Sep 17 00:00:00 2001 From: Xianduo Li Date: Mon, 14 Sep 2026 16:07:03 +0800 Subject: [PATCH 02/19] fix: remove explicit import to avoid black parsing error The te_platform and te_device_type functions are dynamically assigned to the transformer_engine module namespace via: transformer_engine.te_platform = te_platform transformer_engine.te_platform = te_platform in transformer_engine/pytorch/__init__.py, so explicit imports in transformer_engine/__init__.py cause black to fail parsing due to circular import detection. Signed-off-by: Xianduo Li --- transformer_engine/__init__.py | 1 - 1 file changed, 1 deletion(-) diff --git a/transformer_engine/__init__.py b/transformer_engine/__init__.py index e442b2f8614..ee4a69e9467 100644 --- a/transformer_engine/__init__.py +++ b/transformer_engine/__init__.py @@ -103,7 +103,6 @@ def require_nccl_ep() -> None: try: from . import pytorch - from .pytorch import te_platform, te_device_type except ImportError: pass except FileNotFoundError as e: From e8bf31caca961a8ae8307842a61f91681a8357cd Mon Sep 17 00:00:00 2001 From: Xianduo Li Date: Mon, 14 Sep 2026 16:14:08 +0800 Subject: [PATCH 03/19] fix: correct malformed imports that prevent module loading Move te_platform and te_device_type imports outside of parenthesized import lists to fix SyntaxError. The imports were incorrectly placed inside multi-line import statements, preventing the modules from loading. Fixed files: - transformer_engine/debug/features/log_fp8_tensor_stats.py - transformer_engine/pytorch/attention/dot_product_attention/backends.py - transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py - transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py - transformer_engine/pytorch/attention/inference.py - transformer_engine/pytorch/graph.py - transformer_engine/pytorch/quantization.py - transformer_engine/pytorch/quantized_tensor.py - transformer_engine/pytorch/transformer.py Signed-off-by: Xianduo Li --- .../debug/features/log_fp8_tensor_stats.py | 13 +++- .../dot_product_attention/backends.py | 3 +- .../dot_product_attention/context_parallel.py | 10 ++- .../dot_product_attention.py | 6 +- .../pytorch/attention/inference.py | 2 +- .../pytorch/cpp_extensions/gemm.py | 69 +++++-------------- transformer_engine/pytorch/graph.py | 3 +- transformer_engine/pytorch/jit.py | 31 +++------ transformer_engine/pytorch/quantization.py | 6 +- .../pytorch/quantized_tensor.py | 3 +- transformer_engine/pytorch/transformer.py | 3 +- 11 files changed, 63 insertions(+), 86 deletions(-) diff --git a/transformer_engine/debug/features/log_fp8_tensor_stats.py b/transformer_engine/debug/features/log_fp8_tensor_stats.py index 3e0f9babee5..f20accf1491 100644 --- a/transformer_engine/debug/features/log_fp8_tensor_stats.py +++ b/transformer_engine/debug/features/log_fp8_tensor_stats.py @@ -17,8 +17,9 @@ from transformer_engine.debug.features.utils.stats_buffer import STATS_BUFFERS from transformer_engine.debug.features.utils import get_reduction_params, next_enabled_iter from transformer_engine.pytorch.tensor import Quantizer, QuantizedTensor -from transformer_engine.pytorch.tensor.float8_tensor import ( + from transformer_engine import te_device_type, te_platform +from transformer_engine.pytorch.tensor.float8_tensor import ( Float8Quantizer, Float8CurrentScalingQuantizer, ) @@ -58,7 +59,10 @@ def _get_new_quantizer(recipe_name, fp8_dtype): return Float8BlockQuantizer(fp8_dtype=fp8_dtype, rowwise=True, columnwise=True) if recipe_name == "fp8_current_scaling": return Float8CurrentScalingQuantizer( - fp8_dtype=fp8_dtype, device=torch.device(te_device_type()), rowwise=True, columnwise=True + fp8_dtype=fp8_dtype, + device=torch.device(te_device_type()), + rowwise=True, + columnwise=True, ) if recipe_name == "mxfp8": return MXFP8Quantizer(fp8_dtype=fp8_dtype, rowwise=True, columnwise=True) @@ -211,7 +215,10 @@ def check_if_stat_is_supported( if recipe_from_stat == "fp8_delayed_scaling" and stat_without_recipe == "overflows%": return True - if recipe_from_stat in ["fp8_block_scaling"] and te_platform().get_device_capability()[0] < 9: + if ( + recipe_from_stat in ["fp8_block_scaling"] + and te_platform().get_device_capability()[0] < 9 + ): raise ValueError(f"Stat {stat} needs Hopper or later GPU.") if recipe_from_stat == "mxfp8" and te_platform().get_device_capability()[0] < 10: diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 2bd515b6a71..18a8341d93d 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -16,8 +16,9 @@ import torch import torch.nn.functional as F -from transformer_engine.pytorch.utils import ( + from transformer_engine import te_platform +from transformer_engine.pytorch.utils import ( get_device_compute_capability, split_tensor_along_dim, ) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index d3e26a6614c..8237d1ac291 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -8,8 +8,8 @@ import torch import transformer_engine_torch as tex -from transformer_engine.pytorch.utils import ( from transformer_engine import te_platform +from transformer_engine.pytorch.utils import ( get_cudnn_version, nvtx_range_pop, nvtx_range_push, @@ -287,7 +287,9 @@ def get_seq_chunk_ids_for_reordering_after_attn(cp_size, device): def _get_thd_partition_cu_seqlens(cu_seqlens_padded, device=None): """Return physical boundaries used by the THD partition CUDA kernels.""" target_device = torch.device(device if device is not None else cu_seqlens_padded.device) - target_dtype = torch.int32 if target_device.type == te_device_type() else cu_seqlens_padded.dtype + target_dtype = ( + torch.int32 if target_device.type == te_device_type() else cu_seqlens_padded.dtype + ) return cu_seqlens_padded.to(device=target_device, dtype=target_dtype) @@ -311,7 +313,9 @@ def get_thd_partitioned_indices( cp_size, ) target_device = torch.device(device if device is not None else cu_seqlens_padded.device) - target_dtype = torch.int32 if target_device.type == te_device_type() else cu_seqlens_padded.dtype + target_dtype = ( + torch.int32 if target_device.type == te_device_type() else cu_seqlens_padded.dtype + ) chunk_size = total_tokens // cp_size return torch.arange( cp_rank * chunk_size, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index 8ec209551b3..32a9b2ed2ee 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -15,8 +15,8 @@ from torch.fx.experimental.symbolic_shapes import guard_scalar from torch.nn.parameter import Parameter -from transformer_engine.common.recipe import ( from transformer_engine import te_device_type +from transformer_engine.common.recipe import ( Format, Recipe, DelayedScaling, @@ -968,7 +968,9 @@ def __init__( if self.softmax_type == "learnable": self.register_parameter( "softmax_offset", - Parameter(torch.zeros(self.num_attention_heads // self.tp_size, device=te_device_type())), + Parameter( + torch.zeros(self.num_attention_heads // self.tp_size, device=te_device_type()) + ), get_rng_state_tracker=get_rng_state_tracker, ) diff --git a/transformer_engine/pytorch/attention/inference.py b/transformer_engine/pytorch/attention/inference.py index b9849806c44..9ba385a420b 100644 --- a/transformer_engine/pytorch/attention/inference.py +++ b/transformer_engine/pytorch/attention/inference.py @@ -10,8 +10,8 @@ import torch -from transformer_engine.pytorch.attention.custom_ops import ( from transformer_engine import te_platform +from transformer_engine.pytorch.attention.custom_ops import ( copy_to_kv_cache, QKV_FORMAT_VALUE, ) diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index a82e588acb4..a27f5998a13 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -64,15 +64,11 @@ def get_cublas_workspace(device: int, ub: bool, grouped_gemm: bool) -> torch.Ten _multi_stream_cublas_workspace = [] for _ in range(tex.get_num_cublas_streams()): _multi_stream_cublas_workspace.append( - torch.empty( - get_cublas_workspace_size_bytes(), dtype=torch.uint8, device=device - ) + torch.empty(get_cublas_workspace_size_bytes(), dtype=torch.uint8, device=device) ) return _multi_stream_cublas_workspace - return torch.empty( - get_cublas_workspace_size_bytes(), dtype=torch.uint8, device=device - ) + return torch.empty(get_cublas_workspace_size_bytes(), dtype=torch.uint8, device=device) def validate_gemm_scale(scale: Optional[float], required: bool) -> float: @@ -111,10 +107,7 @@ def _nvfp4_row_scaled_gemm_inputs( B_metadata[b_amax_key] = output_row_scales.new_ones(1) B_metadata["row_scaled_nvfp4"] = False - assert ( - output_row_scales.dtype == torch.float32 - and output_col_scales.dtype == torch.float32 - ) + assert output_row_scales.dtype == torch.float32 and output_col_scales.dtype == torch.float32 return ( NVFP4TensorStorage(**A_metadata), NVFP4TensorStorage(**B_metadata), @@ -161,9 +154,7 @@ def _unwrap_tensor( # Select the direction of a hybrid tensor, then process its sub-storage. if isinstance(tensor, HybridQuantizedTensorStorage): sub_storage = ( - tensor.rowwise_sub_storage - if usage == "rowwise" - else tensor.columnwise_sub_storage + tensor.rowwise_sub_storage if usage == "rowwise" else tensor.columnwise_sub_storage ) return _unwrap_tensor(sub_storage, usage) @@ -246,14 +237,10 @@ def general_gemm( ) if ub is not None: - assert ( - ub_type is not None - ), "Comm+GEMM overlap requires a valid `comm_type` argument." + assert ub_type is not None, "Comm+GEMM overlap requires a valid `comm_type` argument." if ub_type == tex.CommOverlapType.RS: if not (bulk_overlap and not ub.is_fp8_ubuf()): - assert ( - extra_output is not None - ), "GEMM+RS overlap requires extra output tensor." + assert extra_output is not None, "GEMM+RS overlap requires extra output tensor." if out is not None: if not out.is_contiguous(): @@ -282,9 +269,7 @@ def general_gemm( # Use bfloat16 as default bias_dtype bias_dtype = TE_DType[torch.bfloat16 if bias is None else bias.dtype] - if isinstance(A, Float8BlockwiseQTensorStorage) or isinstance( - B, Float8BlockwiseQTensorStorage - ): + if isinstance(A, Float8BlockwiseQTensorStorage) or isinstance(B, Float8BlockwiseQTensorStorage): # FP8 block-scaling requires split accumulator use_split_accumulator = True @@ -322,15 +307,11 @@ def general_gemm( assert ( quantization_params is None ), "Row-scaled NVFP4 GEMM currently does not support output quantization." - assert ( - ub is None - ), "Row-scaled NVFP4 GEMM currently does not support CommOverlap." + assert ub is None, "Row-scaled NVFP4 GEMM currently does not support CommOverlap." assert ( extra_output is None ), "Row-scaled NVFP4 GEMM currently does not support extra output." - assert ( - not bulk_overlap - ), "Row-scaled NVFP4 GEMM currently does not support bulk overlap." + assert not bulk_overlap, "Row-scaled NVFP4 GEMM currently does not support bulk overlap." assert out is None or ( isinstance(out, torch.Tensor) and not is_custom(out) ), "Row-scaled NVFP4 GEMM currently supports only plain torch.Tensor outputs." @@ -342,8 +323,8 @@ def general_gemm( ), "Row-scaled NVFP4 GEMM currently requires NVFP4 B." # Reuse the per-tensor GEMM and apply selected row/column global scales # to the FP32 output. This extends #2931 without a dedicated GEMM kernel. - gemm_A, gemm_B, output_row_scales, output_col_scales = ( - _nvfp4_row_scaled_gemm_inputs(A, B, transa=transa, transb=transb) + gemm_A, gemm_B, output_row_scales, output_col_scales = _nvfp4_row_scaled_gemm_inputs( + A, B, transa=transa, transb=transb ) requested_out, requested_out_dtype = out, out_dtype @@ -362,9 +343,7 @@ def general_gemm( gemm_args[14] = False # accumulate after applying the outer scales gemm_kwargs = dict(kwargs) gemm_kwargs["beta"] = 0.0 - out, bias_grad, gelu_input, extra_output = tex.generic_gemm( - *gemm_args, **gemm_kwargs - ) + out, bias_grad, gelu_input, extra_output = tex.generic_gemm(*gemm_args, **gemm_kwargs) out_2d = out.reshape(-1, out.shape[-1]) assert output_row_scales.numel() in (1, out_2d.shape[0]) @@ -383,9 +362,7 @@ def general_gemm( out_2d.mul_(output_row_scales) out_2d.mul_(output_col_scales) if bias is not None: - assert ( - not grad - ), "Row-scaled NVFP4 backward does not support fused bias gradient." + assert not grad, "Row-scaled NVFP4 backward does not support fused bias gradient." out_2d.add_(bias.to(dtype=torch.float32)) if requested_out is not None: @@ -457,9 +434,7 @@ def general_grouped_gemm( if any(_is_nvfp4_row_scaled_tensor(tensor) for tensor in A) or any( _is_nvfp4_row_scaled_tensor(tensor) for tensor in B ): - assert ( - D_dtype is None - ), "Row-scaled NVFP4 grouped GEMM currently does not support D_dtype." + assert D_dtype is None, "Row-scaled NVFP4 grouped GEMM currently does not support D_dtype." if single_output: assert ( m_splits is not None @@ -592,9 +567,7 @@ def _get_grouped_cublas_workspace(device: int, layout: str) -> torch.Tensor: kernels). """ assert layout in ("TN", "NN", "NT"), f"unexpected grouped GEMM layout {layout}" - return torch.empty( - get_cublas_workspace_size_bytes(), dtype=torch.uint8, device=device - ) + return torch.empty(get_cublas_workspace_size_bytes(), dtype=torch.uint8, device=device) def general_grouped_gemm_for_grouped_tensor( @@ -630,17 +603,11 @@ def general_grouped_gemm_for_grouped_tensor( raise ValueError("Both A and out are discrete. This is not supported yet.") if isinstance(A, GroupedTensorStorage) and A.row_scaled_nvfp4: - raise NotImplementedError( - "Row-scaled NVFP4 GroupedTensor GEMM is not supported yet." - ) + raise NotImplementedError("Row-scaled NVFP4 GroupedTensor GEMM is not supported yet.") if isinstance(B, GroupedTensorStorage) and B.row_scaled_nvfp4: - raise NotImplementedError( - "Row-scaled NVFP4 GroupedTensor GEMM is not supported yet." - ) + raise NotImplementedError("Row-scaled NVFP4 GroupedTensor GEMM is not supported yet.") if isinstance(out, GroupedTensorStorage) and out.row_scaled_nvfp4: - raise NotImplementedError( - "Row-scaled NVFP4 GroupedTensor GEMM is not supported yet." - ) + raise NotImplementedError("Row-scaled NVFP4 GroupedTensor GEMM is not supported yet.") def _is_fp8_blockwise(operand) -> bool: if isinstance(operand, (list, tuple)): diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index f518c943b92..40374a62437 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -17,8 +17,9 @@ from transformer_engine.common.recipe import DelayedScaling, Recipe from transformer_engine.pytorch.constants import dist_group_type -from .quantization import ( + from transformer_engine import te_platform +from .quantization import ( autocast, FP8GlobalStateManager, get_default_fp8_recipe, diff --git a/transformer_engine/pytorch/jit.py b/transformer_engine/pytorch/jit.py index 145eb2c45a7..7ca52b0aa75 100644 --- a/transformer_engine/pytorch/jit.py +++ b/transformer_engine/pytorch/jit.py @@ -82,8 +82,7 @@ def wrapper(*args, **kwargs): # Arguments are matched to names by position, which only holds # without a *args in between. assert not any( - p.kind is inspect.Parameter.VAR_POSITIONAL - for p in parameters.values() + p.kind is inspect.Parameter.VAR_POSITIONAL for p in parameters.values() ), f"no_torch_dynamo(when=...) does not support *args, which {f.__name__} takes" parameter_names = list(parameters) @@ -170,9 +169,9 @@ def bgrad_dgelu_fused_( x = inp + bias tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 - ff = 0.5 * x * ( - (1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x) - ) + 0.5 * (1 + tanh_out) + ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( + 1 + tanh_out + ) dgelu = ff * grad_output bgrad = dgelu.sum(dim=0) return bgrad, dgelu @@ -186,9 +185,9 @@ def dgelu_fused_(grad_output: torch.Tensor, inp: torch.Tensor) -> torch.Tensor: x = inp tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 - ff = 0.5 * x * ( - (1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x) - ) + 0.5 * (1 + tanh_out) + ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( + 1 + tanh_out + ) dgelu = ff * grad_output return dgelu @@ -205,9 +204,7 @@ def l2normalization_fused_(x: torch.Tensor, eps: float) -> torch.Tensor: @jit_fuser -def l2normalization_fwd_fused_( - x: torch.Tensor, eps: float -) -> tuple[torch.Tensor, torch.Tensor]: +def l2normalization_fwd_fused_(x: torch.Tensor, eps: float) -> tuple[torch.Tensor, torch.Tensor]: """L2 normalization fused - training version that returns intermediate values""" x_fp32 = x.float() x_squared = x_fp32.pow(2) @@ -261,9 +258,7 @@ def l2normalization_fused(x: torch.Tensor, eps: float) -> torch.Tensor: return l2normalization_fused_(x, eps) -def l2normalization_fwd_fused( - x: torch.Tensor, eps: float -) -> tuple[torch.Tensor, torch.Tensor]: +def l2normalization_fwd_fused(x: torch.Tensor, eps: float) -> tuple[torch.Tensor, torch.Tensor]: """Disable native AMP for l2normalization_fwd_fused_ - training version""" with gpu_autocast_ctx(enabled=False): return l2normalization_fwd_fused_(x, eps) @@ -357,9 +352,7 @@ def warmup_jit_bias_dropout_add( dropout_rate = 0.1 # Warmup JIT fusions with the input grad_enable state of both forward # prop and recomputation - for input_grad, bias_grad, residual_grad in zip( - [False, True], [True, True], [True, True] - ): + for input_grad, bias_grad, residual_grad in zip([False, True], [True, True], [True, True]): inp.requires_grad = input_grad bias.requires_grad = bias_grad residual.requires_grad = residual_grad @@ -390,9 +383,7 @@ def warmup_jit_bias_gelu( # Save cuda RNG state to ensure warmup does not affect reproducibility. rng_state = te_platform().get_rng_state() - bias = torch.rand( - ffn_hidden_size_per_partition, dtype=dtype, device=te_device_type() - ) + bias = torch.rand(ffn_hidden_size_per_partition, dtype=dtype, device=te_device_type()) inp = torch.rand( (seq_length * micro_batch_size, ffn_hidden_size_per_partition), dtype=dtype, diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index f097aed0bc5..7c765479188 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -16,8 +16,9 @@ import torch import transformer_engine_torch as tex -from transformer_engine.common.recipe import ( + from transformer_engine import te_platform, te_device_type +from transformer_engine.common.recipe import ( Recipe, DelayedScaling, Format, @@ -592,7 +593,8 @@ def fp8_graph_capturing(cls) -> bool: assert not cls.quantization_state.fp8_graph_capturing return False return ( - cls.quantization_state.fp8_graph_capturing or te_platform().is_current_stream_capturing() + cls.quantization_state.fp8_graph_capturing + or te_platform().is_current_stream_capturing() ) @classmethod diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index b1f6633d54f..6c80d6652c3 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -17,8 +17,9 @@ from transformer_engine.common.recipe import Recipe from transformer_engine.pytorch.constants import dist_group_type -from transformer_engine.pytorch.tensor._quantization_helpers import ( + from transformer_engine import te_platform, te_device_type +from transformer_engine.pytorch.tensor._quantization_helpers import ( _QuantizeFunc, _IdentityFunc, _stride_from_shape, diff --git a/transformer_engine/pytorch/transformer.py b/transformer_engine/pytorch/transformer.py index cc728f2d5e1..8f8e4f90381 100644 --- a/transformer_engine/pytorch/transformer.py +++ b/transformer_engine/pytorch/transformer.py @@ -14,8 +14,9 @@ from transformer_engine.pytorch.module import LayerNormMLP, LayerNorm, RMSNorm from transformer_engine.pytorch.attention.multi_head_attention import MultiheadAttention from transformer_engine.pytorch.attention.inference import InferenceParams -from transformer_engine.pytorch.jit import ( + from transformer_engine import te_device_type +from transformer_engine.pytorch.jit import ( set_jit_fusion_options, warmup_jit_bias_dropout_add_all_dtypes, get_bias_dropout_add, From 100803e4267e2ff58205e0755d13c59c74c1bab5 Mon Sep 17 00:00:00 2001 From: Xianduo Li Date: Mon, 14 Sep 2026 16:15:46 +0800 Subject: [PATCH 04/19] fix: add missing te_device_type import to files using it Several files imported only te_platform but also called te_device_type(), causing NameError at runtime. Add te_device_type to imports. Fixed files: - transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py - transformer_engine/pytorch/attention/rope.py - transformer_engine/pytorch/cpp_extensions/gemm.py - transformer_engine/pytorch/optimizers/fused_adam.py - transformer_engine/pytorch/utils.py Signed-off-by: Xianduo Li --- .../pytorch/attention/dot_product_attention/flex_attention.py | 2 +- transformer_engine/pytorch/attention/rope.py | 2 +- transformer_engine/pytorch/cpp_extensions/gemm.py | 2 +- transformer_engine/pytorch/optimizers/fused_adam.py | 2 +- transformer_engine/pytorch/utils.py | 2 +- 5 files changed, 5 insertions(+), 5 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py index 967e983a60a..7de387a5d96 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py @@ -10,7 +10,7 @@ from typing import Any, Callable, Dict, Optional, Tuple import torch -from transformer_engine import te_platform +from transformer_engine import te_device_type, te_platform _cudnn_score_mod_handles: Dict[torch.device, Any] = {} _cudnn_score_mod_graph_cache: Dict[Tuple[Any, ...], Any] = {} diff --git a/transformer_engine/pytorch/attention/rope.py b/transformer_engine/pytorch/attention/rope.py index 3462793dcf7..d60673f0a0c 100644 --- a/transformer_engine/pytorch/attention/rope.py +++ b/transformer_engine/pytorch/attention/rope.py @@ -10,7 +10,7 @@ import transformer_engine_torch as tex from transformer_engine.pytorch.cpp_extensions.fused_attn import QKVFormat -from transformer_engine import te_platform +from transformer_engine import te_device_type, te_platform __all__ = ["RotaryPositionEmbedding", "apply_rotary_pos_emb", "apply_fused_qkv_rotary_pos_emb"] diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index a27f5998a13..4e1fce4f4ac 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -28,7 +28,7 @@ from ..tensor.utils import is_custom from ..custom_recipes.gemm import custom_gemm from ...debug.pytorch.debug_quantization import DebugQuantizedTensor, DebugQuantizer -from transformer_engine import te_platform +from transformer_engine import te_device_type, te_platform __all__ = [ "general_gemm", diff --git a/transformer_engine/pytorch/optimizers/fused_adam.py b/transformer_engine/pytorch/optimizers/fused_adam.py index fd7eb24774a..6b14cc89568 100644 --- a/transformer_engine/pytorch/optimizers/fused_adam.py +++ b/transformer_engine/pytorch/optimizers/fused_adam.py @@ -17,7 +17,7 @@ from transformer_engine.pytorch.quantized_tensor import QuantizedTensor from ..constants import DType from .multi_tensor_apply import multi_tensor_applier -from transformer_engine import te_platform +from transformer_engine import te_device_type, te_platform # Bound temporary NVTETensor handles created before the CUDA launcher chunks its metadata. diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py index 09548f1d395..461d6857d69 100644 --- a/transformer_engine/pytorch/utils.py +++ b/transformer_engine/pytorch/utils.py @@ -16,7 +16,7 @@ from .torch_version import torch_version from ..debug.pytorch.debug_quantization import DebugQuantizedTensor -from transformer_engine import te_platform +from transformer_engine import te_device_type, te_platform __all__ = [ From 9e0b0e2f24be8735393603d52daf8a7dfcf7cb99 Mon Sep 17 00:00:00 2001 From: Xianduo Li Date: Mon, 14 Sep 2026 16:23:07 +0800 Subject: [PATCH 05/19] fix: correct malformed imports in additional files Move te_platform and te_device_type imports outside of parenthesized import lists to fix SyntaxError across 20 more files. Fixed files: - transformer_engine/pytorch/attention/dot_product_attention/utils.py - transformer_engine/pytorch/attention/multi_head_attention.py - transformer_engine/pytorch/cpu_offload.py - transformer_engine/pytorch/module/base.py - transformer_engine/pytorch/module/grouped_linear.py - transformer_engine/pytorch/module/layernorm_linear.py - transformer_engine/pytorch/module/layernorm_mlp.py - transformer_engine/pytorch/module/linear.py - transformer_engine/pytorch/ops/basic/basic_linear.py - transformer_engine/pytorch/ops/basic/grouped_linear.py - transformer_engine/pytorch/ops/basic/l2normalization.py - transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py - transformer_engine/pytorch/ops/fused/grouped_mlp.py - transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py - transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py - transformer_engine/pytorch/ops/op.py - transformer_engine/pytorch/tensor/nvfp4_tensor.py - transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py - transformer_engine/pytorch/triton/mhc.py - transformer_engine/pytorch/triton/permutation.py Signed-off-by: Xianduo Li --- .../attention/dot_product_attention/utils.py | 10 +++++++--- .../pytorch/attention/multi_head_attention.py | 2 +- transformer_engine/pytorch/cpu_offload.py | 6 ++++-- transformer_engine/pytorch/module/base.py | 2 +- .../pytorch/module/grouped_linear.py | 2 +- .../pytorch/module/layernorm_linear.py | 2 +- transformer_engine/pytorch/module/layernorm_mlp.py | 2 +- transformer_engine/pytorch/module/linear.py | 2 +- .../pytorch/ops/basic/basic_linear.py | 2 +- .../pytorch/ops/basic/grouped_linear.py | 2 +- .../pytorch/ops/basic/l2normalization.py | 2 +- .../pytorch/ops/fused/forward_linear_scale_add.py | 2 +- .../pytorch/ops/fused/grouped_mlp.py | 2 +- .../ops/fused/userbuffers_backward_linear.py | 2 +- .../ops/fused/userbuffers_forward_linear.py | 2 +- transformer_engine/pytorch/ops/op.py | 2 +- transformer_engine/pytorch/tensor/nvfp4_tensor.py | 2 +- .../pytorch/tensor/storage/mxfp8_tensor_storage.py | 2 +- transformer_engine/pytorch/triton/mhc.py | 2 +- transformer_engine/pytorch/triton/permutation.py | 14 ++++++++++---- 20 files changed, 38 insertions(+), 26 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 9514cf7149e..8e954ef9d5d 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -22,8 +22,8 @@ import torch.nn.functional as F import transformer_engine_torch as tex import transformer_engine as te -from transformer_engine.pytorch.cpp_extensions.fused_attn import ( from transformer_engine import te_platform +from transformer_engine.pytorch.cpp_extensions.fused_attn import ( QKVLayout, QKVFormat, AttnBiasType, @@ -2094,7 +2094,9 @@ def get_full_mask( # apply SWA mask mask = torch.arange(max_seqlen_q, dtype=torch.int32, device=te_device_type()).view( 1, 1, max_seqlen_q, 1 - ) - torch.arange(max_seqlen_kv, dtype=torch.int32, device=te_device_type()).view(1, 1, 1, max_seqlen_kv) + ) - torch.arange(max_seqlen_kv, dtype=torch.int32, device=te_device_type()).view( + 1, 1, 1, max_seqlen_kv + ) swa_left = None swa_right = None if attn_mask_type == "causal_bottom_right" or ( @@ -2225,7 +2227,9 @@ def get_alibi( _alibi_cache["_max_seqlen_q"], _alibi_cache["_max_seqlen_kv"] = max_seqlen_q, max_seqlen_kv _alibi_cache["_bottom_right_alignment"] = bottom_right_alignment bias_dtype = torch.float32 if bias_dtype is None else bias_dtype - _alibi_cache["_alibi_bias"] = bias.contiguous().to(dtype=bias_dtype, device=te_device_type()) + _alibi_cache["_alibi_bias"] = bias.contiguous().to( + dtype=bias_dtype, device=te_device_type() + ) _alibi_cache["_alibi_bias_require_update"] = False return _alibi_cache["_alibi_slopes"], _alibi_cache["_alibi_bias"] diff --git a/transformer_engine/pytorch/attention/multi_head_attention.py b/transformer_engine/pytorch/attention/multi_head_attention.py index 3fd67038070..93ed14ad413 100644 --- a/transformer_engine/pytorch/attention/multi_head_attention.py +++ b/transformer_engine/pytorch/attention/multi_head_attention.py @@ -14,8 +14,8 @@ from transformer_engine.pytorch.module.base import TransformerEngineBaseModule from transformer_engine.pytorch.module import LayerNormLinear, Linear, RMSNorm, LayerNorm from transformer_engine.pytorch.ops.basic.l2normalization import L2Normalization -from transformer_engine.pytorch.utils import ( from transformer_engine import te_device_type +from transformer_engine.pytorch.utils import ( SplitAlongDim, divide, get_default_init_method, diff --git a/transformer_engine/pytorch/cpu_offload.py b/transformer_engine/pytorch/cpu_offload.py index f9e6aba73e3..42415a07ac0 100644 --- a/transformer_engine/pytorch/cpu_offload.py +++ b/transformer_engine/pytorch/cpu_offload.py @@ -16,8 +16,8 @@ from transformer_engine.debug.pytorch.debug_state import TEDebugState import transformer_engine.pytorch as te import transformer_engine.pytorch.cpu_offload_v1 as v1_code_path -from .quantized_tensor import ( from transformer_engine import te_platform, te_device_type +from .quantized_tensor import ( restore_from_saved, prepare_for_saving, QuantizedTensor, @@ -507,7 +507,9 @@ def __init__( offload_stream: Optional[torch.cuda.Stream] = None, ): self.num_layers = num_layers - self.offload_stream = offload_stream if offload_stream is not None else te_platform().Stream() + self.offload_stream = ( + offload_stream if offload_stream is not None else te_platform().Stream() + ) self.layer_states = { i: OffloadableLayerState(self.offload_stream, retain_pinned_cpu_buffers) diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 15963252c9d..0074aa87b84 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -21,8 +21,8 @@ import transformer_engine_torch as tex from ._common import _ParameterInitMeta, noop_cat -from .._extra_state import ( from transformer_engine import te_platform, te_device_type +from .._extra_state import ( extra_state_pickle_advisory, is_stateless_recipe, should_load_extra_state_pickle, diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 9d22e30530f..bec0b401288 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -16,8 +16,8 @@ import transformer_engine_torch as tex from transformer_engine.common.recipe import Recipe -from transformer_engine.pytorch.tensor.grouped_tensor import ( from transformer_engine import te_device_type +from transformer_engine.pytorch.tensor.grouped_tensor import ( GroupedTensor, GroupedTensorStorage, ) diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index bc6e257c5f4..a09027068eb 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -18,8 +18,8 @@ from transformer_engine.common.recipe import Recipe from transformer_engine.pytorch.torch_version import torch_version from transformer_engine.pytorch.tensor.utils import clear_columnwise_cache, is_custom -from .base import ( from transformer_engine import te_platform, te_device_type +from .base import ( fill_userbuffers_buffer_for_all_gather, get_ub, get_ub_is_fp8, diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index de96326804f..5952560ffa0 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -19,8 +19,8 @@ from transformer_engine.common.recipe import Recipe from transformer_engine.pytorch.torch_version import torch_version from transformer_engine.pytorch.tensor.utils import clear_columnwise_cache, is_custom -from .base import ( from transformer_engine import te_platform, te_device_type +from .base import ( fill_userbuffers_buffer_for_all_gather, _ub_communicators, get_ub, diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index c931b32a8a4..ecf80bf41fd 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -18,8 +18,8 @@ from transformer_engine.common.recipe import Recipe from transformer_engine.pytorch.torch_version import torch_version -from .base import ( from transformer_engine import te_platform, te_device_type +from .base import ( fill_userbuffers_buffer_for_all_gather, get_dummy_wgrad, get_ub, diff --git a/transformer_engine/pytorch/ops/basic/basic_linear.py b/transformer_engine/pytorch/ops/basic/basic_linear.py index a04086d5eae..31d71d2e1ae 100644 --- a/transformer_engine/pytorch/ops/basic/basic_linear.py +++ b/transformer_engine/pytorch/ops/basic/basic_linear.py @@ -14,8 +14,8 @@ from ...cpp_extensions import general_gemm from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload -from ...distributed import ( from transformer_engine import te_device_type +from ...distributed import ( CudaRNGStatesTracker, gather_along_first_dim, reduce_scatter_along_first_dim, diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index 40587e5525b..b754b66dd64 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -19,8 +19,8 @@ from ...cpp_extensions import general_grouped_gemm, general_grouped_gemm_for_grouped_tensor from ...distributed import CudaRNGStatesTracker from ...module._common import WeightGradStore -from ...module.base import ( from transformer_engine import te_device_type +from ...module.base import ( _2X_ACC_FPROP, _2X_ACC_DGRAD, _2X_ACC_WGRAD, diff --git a/transformer_engine/pytorch/ops/basic/l2normalization.py b/transformer_engine/pytorch/ops/basic/l2normalization.py index 27d3a221844..05106c7833c 100644 --- a/transformer_engine/pytorch/ops/basic/l2normalization.py +++ b/transformer_engine/pytorch/ops/basic/l2normalization.py @@ -12,8 +12,8 @@ from ...torch_version import torch_version from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload -from ...jit import ( from transformer_engine import te_platform +from ...jit import ( l2normalization_fused, l2normalization_fwd_fused, l2normalization_backward_fused, diff --git a/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py b/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py index 2d28f628c3e..cae659f4fac 100644 --- a/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py +++ b/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py @@ -14,8 +14,8 @@ from ...quantization import FP8GlobalStateManager from ...tensor import Quantizer from ..basic import AddExtraInput, BasicLinear, ConstantScale -from ..op import ( from transformer_engine import te_device_type +from ..op import ( FusedOperation, FusibleOperation, OperationContext, diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 59bfc895958..0fffb566ba5 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -20,8 +20,8 @@ from ...constants import DType, MXFP8_BLOCK_SCALING_SIZE, NVFP4_BLOCK_SCALING_SIZE, TE_DType from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload, start_offload from ...cpp_extensions import general_gemm, general_grouped_gemm_for_grouped_tensor -from ...distributed_weight import ( from transformer_engine import te_platform +from ...distributed_weight import ( is_distributed_weight, materialize_weight_for_forward, materialize_weight_for_backward, diff --git a/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py b/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py index 69a20a8dbc7..fb04ee8629f 100644 --- a/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py +++ b/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py @@ -13,8 +13,8 @@ from transformer_engine_torch import CommOverlapType, bulk_overlap_ag_with_external_gemm from ...cpp_extensions import general_gemm from ...distributed import gather_along_first_dim, get_distributed_world_size -from ...module.base import ( from transformer_engine import te_platform +from ...module.base import ( _2X_ACC_DGRAD, _2X_ACC_WGRAD, fill_userbuffers_buffer_for_all_gather, diff --git a/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py b/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py index 8750aed87e8..f8e23e3c730 100644 --- a/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py +++ b/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py @@ -15,8 +15,8 @@ from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload from ...distributed import get_distributed_world_size from ...quantization import FP8GlobalStateManager -from ...module.base import ( from transformer_engine import te_device_type +from ...module.base import ( fill_userbuffers_buffer_for_all_gather, get_ub, _2X_ACC_FPROP, diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py index 57b52398e45..6ffddd6e225 100644 --- a/transformer_engine/pytorch/ops/op.py +++ b/transformer_engine/pytorch/ops/op.py @@ -15,8 +15,8 @@ from transformer_engine.common.recipe import Recipe from .._extra_state import is_stateless_recipe, should_load_extra_state_pickle -from ..quantization import ( from transformer_engine import te_platform +from ..quantization import ( FP8GlobalStateManager, QuantizerRole, RecipeState, diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index dbcc016e40b..e3bdf5852a6 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -15,8 +15,8 @@ from transformer_engine.common.recipe import NVFP4BlockScaling, Recipe from ..constants import NVFP4_BLOCK_SCALING_SIZE, dist_group_type, DType -from ..utils import ( from transformer_engine import te_platform +from ..utils import ( canonicalize_process_group, devices_match, round_up_to_nearest_multiple, diff --git a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py index df4b8ae0e38..9d40ac5b48b 100644 --- a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py @@ -15,8 +15,8 @@ from ...quantized_tensor import InnerTensor, QuantizedTensorStorage, Quantizer from .._quantization_helpers import safe_quantized_repr -from ...constants import ( from transformer_engine import te_device_type +from ...constants import ( TE_DType as torch_to_transformer_engine_dtype, MXFP8_BLOCK_SCALING_SIZE, DType, diff --git a/transformer_engine/pytorch/triton/mhc.py b/transformer_engine/pytorch/triton/mhc.py index 59ee8d0b68f..36a0ee5a3df 100644 --- a/transformer_engine/pytorch/triton/mhc.py +++ b/transformer_engine/pytorch/triton/mhc.py @@ -9,8 +9,8 @@ import torch import triton -from transformer_engine.common.triton.mhc import ( from transformer_engine import te_platform +from transformer_engine.common.triton.mhc import ( _mhc_projection_bwd_fused_dphi, _mhc_projection_bwd_fused_dx, _mhc_scale_fwd_fused, diff --git a/transformer_engine/pytorch/triton/permutation.py b/transformer_engine/pytorch/triton/permutation.py index d4c28f56d01..7739b34edfa 100644 --- a/transformer_engine/pytorch/triton/permutation.py +++ b/transformer_engine/pytorch/triton/permutation.py @@ -9,8 +9,8 @@ import torch import triton -from transformer_engine.common.triton.permutation import ( from transformer_engine import te_device_type +from transformer_engine.common.triton.permutation import ( _row_id_map_pass_1_kernel, _row_id_map_pass_2_kernel, _row_id_map_pass_3_kernel, @@ -50,7 +50,9 @@ def make_row_id_map( The [num_experts, num_experts + n_routed) items are the indices of the experts corresponding to the first n_routed row indices above. """ - row_id_map = torch.empty((num_tokens, num_experts * 2 + 1), dtype=torch.int32, device=te_device_type()) + row_id_map = torch.empty( + (num_tokens, num_experts * 2 + 1), dtype=torch.int32, device=te_device_type() + ) block_size = 1024 grid = (num_experts, triton.cdiv(num_tokens, block_size)) workspace_tensor = torch.empty(grid, dtype=torch.int32, device=te_device_type()) @@ -163,7 +165,9 @@ def permute_with_mask_map( alloc = torch.zeros if pad_offsets is not None else torch.empty output = alloc((num_out_tokens, hidden_size), dtype=inp.dtype, device=te_device_type()) permuted_probs = ( - alloc((num_out_tokens,), dtype=probs.dtype, device=te_device_type()) if probs is not None else None + alloc((num_out_tokens,), dtype=probs.dtype, device=te_device_type()) + if probs is not None + else None ) permuted_scale = ( alloc((num_out_tokens, scale_hidden_dim), dtype=scale.dtype, device=te_device_type()) @@ -326,7 +330,9 @@ def unpermute_with_mask_map_bwd_with_merging_probs( # by the kernel. This matches the behavior of Fp8Unpadding.backward which zeros # out the padding slots. alloc = torch.zeros if pad_offsets is not None else torch.empty - act_grad = alloc((num_out_tokens, hidden_size), dtype=fwd_output_grad.dtype, device=te_device_type()) + act_grad = alloc( + (num_out_tokens, hidden_size), dtype=fwd_output_grad.dtype, device=te_device_type() + ) merging_probs_grad = torch.empty( (num_tokens, num_experts), dtype=merging_probs.dtype, device=te_device_type() ) From d7b7fe970a2a3617091d11eb89ce1b35dc30a5ae Mon Sep 17 00:00:00 2001 From: Xianduo Li Date: Mon, 14 Sep 2026 16:30:17 +0800 Subject: [PATCH 06/19] fix: use routing_map device instead of hardcoded device type In make_row_id_map, use routing_map.device for tensor allocation instead of te_device_type() to respect the input tensor's device. Signed-off-by: Xianduo Li --- transformer_engine/pytorch/triton/permutation.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/transformer_engine/pytorch/triton/permutation.py b/transformer_engine/pytorch/triton/permutation.py index 7739b34edfa..744ab751058 100644 --- a/transformer_engine/pytorch/triton/permutation.py +++ b/transformer_engine/pytorch/triton/permutation.py @@ -51,11 +51,11 @@ def make_row_id_map( to the first n_routed row indices above. """ row_id_map = torch.empty( - (num_tokens, num_experts * 2 + 1), dtype=torch.int32, device=te_device_type() + (num_tokens, num_experts * 2 + 1), dtype=torch.int32, device=routing_map.device ) block_size = 1024 grid = (num_experts, triton.cdiv(num_tokens, block_size)) - workspace_tensor = torch.empty(grid, dtype=torch.int32, device=te_device_type()) + workspace_tensor = torch.empty(grid, dtype=torch.int32, device=routing_map.device) # supposing num_tokens == 5, num_experts == 3, block_size == 3 # and we have a routing_map like this: From 6180739c7e24f80af6b0c2e492db9206fe4b3e7f Mon Sep 17 00:00:00 2001 From: Xianduo Li Date: Mon, 14 Sep 2026 16:43:18 +0800 Subject: [PATCH 07/19] revert: remove te_platform() replacements, keep te_device_type() Revert all te_platform() method calls back to torch.cuda.* calls. Keep te_device_type() string replacements for device type parameters. Preserve te_platform definition in transformer_engine/pytorch/__init__.py. Changes: - Reverted 146 te_platform() calls back to torch.cuda.* - Removed 'from transformer_engine import te_platform' imports - Kept all te_device_type() replacements intact - Kept te_platform definition in __init__.py for future use Signed-off-by: Xianduo Li --- .../debug/features/log_fp8_tensor_stats.py | 9 ++-- .../dot_product_attention/backends.py | 7 ++- .../dot_product_attention/context_parallel.py | 47 +++++++++---------- .../dot_product_attention/flex_attention.py | 10 ++-- .../attention/dot_product_attention/utils.py | 7 ++- .../pytorch/attention/inference.py | 19 ++++---- transformer_engine/pytorch/attention/rope.py | 4 +- .../pytorch/cpp_extensions/gemm.py | 4 +- transformer_engine/pytorch/cpu_offload.py | 30 ++++++------ transformer_engine/pytorch/cpu_offload_v1.py | 31 ++++++------ transformer_engine/pytorch/distributed.py | 16 +++---- transformer_engine/pytorch/ep.py | 8 ++-- transformer_engine/pytorch/graph.py | 39 ++++++++------- transformer_engine/pytorch/jit.py | 19 ++++---- transformer_engine/pytorch/module/base.py | 8 ++-- .../pytorch/module/layernorm_linear.py | 4 +- .../pytorch/module/layernorm_mlp.py | 12 ++--- transformer_engine/pytorch/module/linear.py | 4 +- .../pytorch/ops/basic/l2normalization.py | 3 +- .../pytorch/ops/fused/grouped_mlp.py | 3 +- .../ops/fused/userbuffers_backward_linear.py | 3 +- transformer_engine/pytorch/ops/op.py | 5 +- .../pytorch/optimizers/fused_adam.py | 2 +- .../pytorch/optimizers/newton_schulz.py | 4 +- transformer_engine/pytorch/quantization.py | 5 +- .../pytorch/quantized_tensor.py | 4 +- transformer_engine/pytorch/router.py | 5 +- .../pytorch/tensor/nvfp4_tensor.py | 5 +- .../tensor/storage/grouped_tensor_storage.py | 5 +- transformer_engine/pytorch/triton/mhc.py | 3 +- transformer_engine/pytorch/utils.py | 16 +++---- 31 files changed, 160 insertions(+), 181 deletions(-) diff --git a/transformer_engine/debug/features/log_fp8_tensor_stats.py b/transformer_engine/debug/features/log_fp8_tensor_stats.py index f20accf1491..efc0d8793b3 100644 --- a/transformer_engine/debug/features/log_fp8_tensor_stats.py +++ b/transformer_engine/debug/features/log_fp8_tensor_stats.py @@ -18,7 +18,7 @@ from transformer_engine.debug.features.utils import get_reduction_params, next_enabled_iter from transformer_engine.pytorch.tensor import Quantizer, QuantizedTensor -from transformer_engine import te_device_type, te_platform +from transformer_engine import te_device_type from transformer_engine.pytorch.tensor.float8_tensor import ( Float8Quantizer, Float8CurrentScalingQuantizer, @@ -215,13 +215,10 @@ def check_if_stat_is_supported( if recipe_from_stat == "fp8_delayed_scaling" and stat_without_recipe == "overflows%": return True - if ( - recipe_from_stat in ["fp8_block_scaling"] - and te_platform().get_device_capability()[0] < 9 - ): + if recipe_from_stat in ["fp8_block_scaling"] and torch.cuda.get_device_capability()[0] < 9: raise ValueError(f"Stat {stat} needs Hopper or later GPU.") - if recipe_from_stat == "mxfp8" and te_platform().get_device_capability()[0] < 10: + if recipe_from_stat == "mxfp8" and torch.cuda.get_device_capability()[0] < 10: raise ValueError(f"Stat {stat} needs Blackwell or later GPU.") supported_stats = ["underflows%", "scale_inv_min", "scale_inv_max", "mse"] diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 18a8341d93d..8d1be7a3093 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -17,7 +17,6 @@ import torch import torch.nn.functional as F -from transformer_engine import te_platform from transformer_engine.pytorch.utils import ( get_device_compute_capability, split_tensor_along_dim, @@ -108,7 +107,7 @@ except PackageNotFoundError: pass # only print warning if use_flash_attention_2 = True in get_attention_backend else: - if te_platform().is_available() and get_device_compute_capability() >= (10, 0): + if torch.cuda.is_available() and get_device_compute_capability() >= (10, 0): if fa_utils.is_version_supported(fa_utils.version, fa_utils.version_required_blackwell): fa_utils.is_installed = True elif fa_utils.is_version_supported(fa_utils.version, fa_utils.version_required): @@ -129,7 +128,7 @@ # Setup Flash attention utils fa_utils.set_flash_attention_version() elif ( - te_platform().is_available() + torch.cuda.is_available() and get_device_compute_capability() >= (8, 0) and dpa_utils._NVTE_FLASH_ATTN ): @@ -579,7 +578,7 @@ def _forward( output_size[2], output_size[3], dtype=query_layer.dtype, - device=te_platform().current_device(), + device=torch.cuda.current_device(), ) scale = self.softmax_scale diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index 8237d1ac291..f426790c3cb 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -8,7 +8,6 @@ import torch import transformer_engine_torch as tex -from transformer_engine import te_platform from transformer_engine.pytorch.utils import ( get_cudnn_version, nvtx_range_pop, @@ -571,7 +570,7 @@ def flash_attn_a2a_communicate( a2a_outputs[i - 1], a2a_inputs[i - 1], group=cp_group, async_op=True ) if i > 1: - with te_platform().stream(cp_stream): + with torch.cuda.stream(cp_stream): a2a_reqs[i - 2].wait() x = a2a_outputs[i - 2] if qkv_format in ["bshd", "sbhd", "bhsd"]: @@ -644,7 +643,7 @@ def flash_attn_a2a_communicate( # [cp*t, h//cp, d] -> [cp, t, h//cp, d] a2a_inputs[i] = x.view(cp_size, -1, *x.shape[-2:]) if i > 1: - with te_platform().stream(cp_stream): + with torch.cuda.stream(cp_stream): a2a_reqs[i - 2].wait() x = a2a_outputs[i - 2] # [cp, 2, b, s//2, h//cp, d] -> [2, b, s//2, cp, h//cp, d] @@ -676,7 +675,7 @@ def flash_attn_a2a_communicate( # [b, cp, h//cp, 2, s//2, d] -> [b*h, s, d] # [t, cp, h//cp, d] -> [t, h, d] a2a_outputs[i - 2] = x.view(-1, x.shape[-3] * x.shape[-2], x.shape[-1]) - te_platform().current_stream().wait_stream(cp_stream) + torch.cuda.current_stream().wait_stream(cp_stream) return a2a_outputs[0] if len(a2a_inputs) == 1 else a2a_outputs @@ -715,14 +714,14 @@ def flash_attn_a2a_communicate_softmax_offset( # [1, h//cp, 1, 1] -> [1, h, 1, 1] inp = tensor.view(-1) output = torch.empty(cp_size * inp.shape[0], dtype=tensor.dtype, device=device) - with te_platform().stream(cp_stream): + with torch.cuda.stream(cp_stream): torch.distributed.all_gather_into_tensor( output, inp, group=cp_group, async_op=False, ) - te_platform().current_stream().wait_stream(cp_stream) + torch.cuda.current_stream().wait_stream(cp_stream) output = output.view( *tensor.shape[:h_dim], cp_size * tensor.shape[h_dim], *tensor.shape[h_dim + 1 :] ) @@ -1923,9 +1922,9 @@ def forward( attn_biases = [None for _ in range(cp_size)] # create two streams to resolve wave quantization issue of Flash Attn in each step - flash_attn_streams = [te_platform().current_stream(), cp_stream] + flash_attn_streams = [torch.cuda.current_stream(), cp_stream] # synchronize fwd results correction across steps - fwd_results_correction_done = te_platform().Event() + fwd_results_correction_done = torch.cuda.Event() # q, k, v, o: # causal: [b, 2, s//2, h, d] or [2, s//2, b, h, d] @@ -1945,7 +1944,7 @@ def forward( o_format = qkv_format for i in range(cp_size + 1): if i < cp_size: - with te_platform().stream(flash_attn_streams[i % 2]): + with torch.cuda.stream(flash_attn_streams[i % 2]): # wait until KV is received for req in send_recv_reqs[(i + 1) % 2]: req.wait() @@ -2165,7 +2164,7 @@ def forward( if i > 1: flash_attn_streams[(i - 1) % 2].wait_event(fwd_results_correction_done) - with te_platform().stream(flash_attn_streams[(i - 1) % 2]): + with torch.cuda.stream(flash_attn_streams[(i - 1) % 2]): if use_fused_attention: # [b, h, sq, 1] -> [b, h, sq] # [t, h, 1] -> [t, h] @@ -2216,7 +2215,7 @@ def forward( if i < cp_size: flash_attn_streams[(i - 1) % 2].record_event(fwd_results_correction_done) - te_platform().current_stream().wait_stream(flash_attn_streams[1]) + torch.cuda.current_stream().wait_stream(flash_attn_streams[1]) if return_max_logit: torch.distributed.all_reduce( max_logit, op=torch.distributed.ReduceOp.MAX, group=cp_group @@ -3505,7 +3504,7 @@ def forward( v_ag = v_ag.view(-1, *v.shape[1:]) # Non-THD cp_stream inputs are ready after K/V reorder, so wait here to # preserve overlap with output initialization below. - cp_stream.wait_stream(te_platform().current_stream()) + cp_stream.wait_stream(torch.cuda.current_stream()) # Shapes before per-step slicing and FP8 metadata wrapping. # q: [b, 2, s//2, h, d] or [2, s//2, b, h, d] @@ -3519,7 +3518,7 @@ def forward( out_f16 = torch.empty(o_shape, dtype=fwd_nominal_dtype, device=q.device) # create two streams to resolve wave quantization issue of Flash Attn in each step - flash_attn_streams = [te_platform().current_stream(), cp_stream] + flash_attn_streams = [torch.cuda.current_stream(), cp_stream] # prepare per-step tensors local_seq_chunk_ids = ( [rank] @@ -3635,7 +3634,7 @@ def forward( if qkv_format == "thd": # Delay the THD wait so one dependency covers both restored K/V and the # per-step metadata produced above on the current stream. - cp_stream.wait_stream(te_platform().current_stream()) + cp_stream.wait_stream(torch.cuda.current_stream()) for i in range(len(local_seq_chunk_ids) + 1): if i < len(local_seq_chunk_ids): @@ -3644,7 +3643,7 @@ def forward( # do not overlap. FusedAttention keeps the existing per-step overlap. if i > 0 and (use_flash_attn_3 or use_flash_attn_4): flash_attn_streams[i].wait_stream(flash_attn_streams[i - 1]) - with te_platform().stream(flash_attn_streams[i]): + with torch.cuda.stream(flash_attn_streams[i]): new_qkv_layout = qkv_layout qkv_scale_inv_format = None if qkv_format in ["bshd", "sbhd"]: @@ -3849,7 +3848,7 @@ def forward( if return_max_logit and i == 0: max_logit = torch.clone(max_logit_per_step[0]) if i > 0: - with te_platform().stream(flash_attn_streams[i - 1]): + with torch.cuda.stream(flash_attn_streams[i - 1]): if o_format == "bshd": out_f16[:, i - 1].copy_(out_per_step[i - 1]) elif o_format == "sbhd": @@ -3870,10 +3869,10 @@ def forward( # default stream, so without this wait the read can race with # the write. The post-loop wait_stream(cp_stream) is too late. # No-op when flash_attn_streams[i-1] is current_stream(). - te_platform().current_stream().wait_stream(flash_attn_streams[i - 1]) + torch.cuda.current_stream().wait_stream(flash_attn_streams[i - 1]) max_logit = torch.maximum(max_logit, max_logit_per_step[i - 1]) - te_platform().current_stream().wait_stream(cp_stream) + torch.cuda.current_stream().wait_stream(cp_stream) # all reduce max_logit across ranks if return_max_logit: @@ -4123,9 +4122,9 @@ def backward(ctx, dout, *_args): dv_per_step = [None, None] # create two streams to resolve wave quantization issue of Flash Attn in each step - flash_attn_streams = [te_platform().current_stream(), ctx.cp_stream] + flash_attn_streams = [torch.cuda.current_stream(), ctx.cp_stream] # synchronize dkv update across steps - dkv_update_done = te_platform().Event() + dkv_update_done = torch.cuda.Event() # gather k and v along s or t: [s, b, h, d] -> [cp, s, b, h, d] or [t, h, d] -> [cp*t, h, d] k_ag, _ = gather_along_first_dim(k, ctx.cp_group) @@ -4154,7 +4153,7 @@ def backward(ctx, dout, *_args): # flatten: [cp*2, s//2, b, h, d] -> [cp*s, b, h, d] k_ag = k_ag.view(-1, *k.shape[1:]) v_ag = v_ag.view(-1, *v.shape[1:]) - ctx.cp_stream.wait_stream(te_platform().current_stream()) + ctx.cp_stream.wait_stream(torch.cuda.current_stream()) # set up flash_attn_bwd flash_attn_bwd = None @@ -4215,7 +4214,7 @@ def backward(ctx, dout, *_args): # per-step overlap. if i > 0 and (ctx.use_flash_attn_3 or ctx.use_flash_attn_4): flash_attn_streams[i].wait_stream(flash_attn_streams[i - 1]) - with te_platform().stream(flash_attn_streams[i]): + with torch.cuda.stream(flash_attn_streams[i]): if ctx.qkv_format == "thd": # THD passes full Q/dout; per-step cu_seqlens select chunks. q_part = q @@ -4463,7 +4462,7 @@ def backward(ctx, dout, *_args): if i > 0: # dq/dk/dv, dq_per_step/dk_per_step/dv_per_step: ctx.fwd_nominal_dtype - with te_platform().stream(flash_attn_streams[i - 1]): + with torch.cuda.stream(flash_attn_streams[i - 1]): if ctx.qkv_format == "thd": # dQ: copy every sequence's valid token range from this split's dQ. tex.thd_copy_valid_tokens_from_per_split_to_rank_local( @@ -4508,7 +4507,7 @@ def backward(ctx, dout, *_args): if i < len(local_seq_chunk_ids): flash_attn_streams[i - 1].record_event(dkv_update_done) - te_platform().current_stream().wait_stream(ctx.cp_stream) + torch.cuda.current_stream().wait_stream(ctx.cp_stream) if ctx.qkv_format == "thd": # Reorder dK/dV from sequence order back to dual-chunk CP rank order, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py index 7de387a5d96..2c8198881de 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py @@ -10,7 +10,7 @@ from typing import Any, Callable, Dict, Optional, Tuple import torch -from transformer_engine import te_device_type, te_platform +from transformer_engine import te_device_type _cudnn_score_mod_handles: Dict[torch.device, Any] = {} _cudnn_score_mod_graph_cache: Dict[Tuple[Any, ...], Any] = {} @@ -144,7 +144,7 @@ def _score_mod_device_key(device: torch.device) -> Tuple[Any, ...]: if device.type == te_device_type(): index = device.index if index is None: - index = te_platform().current_device() + index = torch.cuda.current_device() return (device.type, index) return (device.type, device.index) @@ -198,10 +198,10 @@ def _get_cudnn_current_stream_handle(cudnn, device: torch.device): if device.type != te_device_type(): raise ValueError(f"Flex Attention only supports CUDA tensors, got device {device}.") if device.index is None: - device = torch.device(te_device_type(), te_platform().current_device()) + device = torch.device(te_device_type(), torch.cuda.current_device()) handle = _cudnn_score_mod_handles.get(device) - with te_platform().device(device): + with torch.cuda.device(device): if handle is None: handle = cudnn.create_handle() _cudnn_score_mod_handles[device] = handle @@ -289,7 +289,7 @@ def _execute_cudnn_graph( cudnn = _import_cudnn_frontend() if device.type == te_device_type() and device.index is None: - device = torch.device(te_device_type(), te_platform().current_device()) + device = torch.device(te_device_type(), torch.cuda.current_device()) workspace = torch.empty( workspace_size, device=device, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 8e954ef9d5d..90dbf57521e 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -22,7 +22,6 @@ import torch.nn.functional as F import transformer_engine_torch as tex import transformer_engine as te -from transformer_engine import te_platform from transformer_engine.pytorch.cpp_extensions.fused_attn import ( QKVLayout, QKVFormat, @@ -1829,7 +1828,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt use_flash_attention_4 and not FlashAttentionUtils.v4_is_installed and not FlashAttentionUtils.v4_warning_printed - and te_platform().current_device() == 0 + and torch.cuda.current_device() == 0 ): logger.warning( "flash-attn v4 may provide important feature support or performance improvement." @@ -1841,7 +1840,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt use_flash_attention_3 and not FlashAttentionUtils.v3_is_installed and not FlashAttentionUtils.v3_warning_printed - and te_platform().current_device() == 0 + and torch.cuda.current_device() == 0 ): logger.warning( "flash-attn v3 may provide important feature support or performance improvement." @@ -1853,7 +1852,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt use_flash_attention_2 and not FlashAttentionUtils.is_installed and not FlashAttentionUtils.warning_printed - and te_platform().current_device() == 0 + and torch.cuda.current_device() == 0 ): logger.warning( "flash-attn may provide important feature support or performance improvement." diff --git a/transformer_engine/pytorch/attention/inference.py b/transformer_engine/pytorch/attention/inference.py index 9ba385a420b..ef842c04caa 100644 --- a/transformer_engine/pytorch/attention/inference.py +++ b/transformer_engine/pytorch/attention/inference.py @@ -10,7 +10,6 @@ import torch -from transformer_engine import te_platform from transformer_engine.pytorch.attention.custom_ops import ( copy_to_kv_cache, QKV_FORMAT_VALUE, @@ -210,12 +209,12 @@ def __init__( self.cu_seqlens_q = torch.zeros( self.max_batch_size + 1, dtype=torch.int32, - device=te_platform().current_device(), + device=torch.cuda.current_device(), ) self.cu_seqlens_kv = torch.zeros( self.max_batch_size + 1, dtype=torch.int32, - device=te_platform().current_device(), + device=torch.cuda.current_device(), ) # This internal buffer holds the running length of each @@ -226,7 +225,7 @@ def __init__( self.pre_step_seqlens = torch.zeros( self.max_batch_size, dtype=torch.int32, - device=te_platform().current_device(), + device=torch.cuda.current_device(), ) def reset(self): @@ -432,14 +431,14 @@ def __init__( self.batch_indices = torch.zeros( self.max_batch_size, dtype=torch.int32, - device=te_platform().current_device(), + device=torch.cuda.current_device(), ) # after re-indexing, batch indices are always [0, ..., b-1] self.batch_indices_post_step = torch.range( 0, self.max_batch_size - 1, dtype=torch.int32, - device=te_platform().current_device(), + device=torch.cuda.current_device(), ) # whether reindexing is needed, i.e. when batch seq_ids have changed self.need_reindex = True @@ -452,7 +451,7 @@ def allocate_memory(self, layer_number): self.num_heads, self.head_dim_k, dtype=self.dtype, - device=te_platform().current_device(), + device=torch.cuda.current_device(), ) v_cache = torch.zeros( self.max_batch_size, @@ -460,7 +459,7 @@ def allocate_memory(self, layer_number): self.num_heads, self.head_dim_v, dtype=self.dtype, - device=te_platform().current_device(), + device=torch.cuda.current_device(), ) # Mutated in place, so CUDA graphs are only captured if the addresses are fixed. torch._dynamo.mark_static_address(k_cache) @@ -652,7 +651,7 @@ def allocate_memory(self, layer_number): self.num_heads, self.head_dim_k, dtype=self.dtype, - device=te_platform().current_device(), + device=torch.cuda.current_device(), ) v_cache = torch.zeros( self.total_num_pages, @@ -660,7 +659,7 @@ def allocate_memory(self, layer_number): self.num_heads, self.head_dim_v, dtype=self.dtype, - device=te_platform().current_device(), + device=torch.cuda.current_device(), ) # Mutated in place, so CUDA graphs are only captured if the addresses are fixed. torch._dynamo.mark_static_address(k_cache) diff --git a/transformer_engine/pytorch/attention/rope.py b/transformer_engine/pytorch/attention/rope.py index d60673f0a0c..18397f954b3 100644 --- a/transformer_engine/pytorch/attention/rope.py +++ b/transformer_engine/pytorch/attention/rope.py @@ -10,7 +10,7 @@ import transformer_engine_torch as tex from transformer_engine.pytorch.cpp_extensions.fused_attn import QKVFormat -from transformer_engine import te_device_type, te_platform +from transformer_engine import te_device_type __all__ = ["RotaryPositionEmbedding", "apply_rotary_pos_emb", "apply_fused_qkv_rotary_pos_emb"] @@ -55,7 +55,7 @@ def __init__( inv_freq = 1.0 / ( self.rotary_base ** ( - torch.arange(0, dim, 2, dtype=torch.float32, device=te_platform().current_device()) + torch.arange(0, dim, 2, dtype=torch.float32, device=torch.cuda.current_device()) / dim ) ) diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 4e1fce4f4ac..2fd43b7a11b 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -28,7 +28,7 @@ from ..tensor.utils import is_custom from ..custom_recipes.gemm import custom_gemm from ...debug.pytorch.debug_quantization import DebugQuantizedTensor, DebugQuantizer -from transformer_engine import te_device_type, te_platform +from transformer_engine import te_device_type __all__ = [ "general_gemm", @@ -43,7 +43,7 @@ def get_cublas_workspace_size_bytes() -> None: """Return 32 MiB if using hopper, 4 MiB for all other architectures.""" - if torch.cuda.get_device_properties(te_platform().current_device()).major >= 9: + if torch.cuda.get_device_properties(torch.cuda.current_device()).major >= 9: # 32 MiB for NVFP4 GEMM, plus additional 1024 B for alignment and misc scales return 32 * 1024 * 1024 + 1024 return 4_194_304 diff --git a/transformer_engine/pytorch/cpu_offload.py b/transformer_engine/pytorch/cpu_offload.py index 42415a07ac0..73604634721 100644 --- a/transformer_engine/pytorch/cpu_offload.py +++ b/transformer_engine/pytorch/cpu_offload.py @@ -16,7 +16,7 @@ from transformer_engine.debug.pytorch.debug_state import TEDebugState import transformer_engine.pytorch as te import transformer_engine.pytorch.cpu_offload_v1 as v1_code_path -from transformer_engine import te_platform, te_device_type +from transformer_engine import te_device_type from .quantized_tensor import ( restore_from_saved, prepare_for_saving, @@ -74,8 +74,8 @@ def _mark_tensor_for_offload(t): if t is None: return # Attach an event to mark when the tensor is ready for reload. - t.start_reload_event = te_platform().Event() - t.start_reload_event.record(te_platform().current_stream()) + t.start_reload_event = torch.cuda.Event() + t.start_reload_event.record(torch.cuda.current_stream()) if offload_base_tensor and t._base is not None: setattr(t, "offload_base_tensor", True) @@ -287,7 +287,7 @@ def start_offload(self): # Wait for the moment the tensor is ready to be offloaded. self.offload_stream.wait_event(self.fwd_gpu_tensor_group.events[tensor_id]) # type: ignore[arg-type] - with te_platform().stream(self.offload_stream): + with torch.cuda.stream(self.offload_stream): if allocate_cpu_buffers: offloaded_tensor = torch.empty_like( tensor, device=torch.device("cpu"), pin_memory=True @@ -309,7 +309,7 @@ def start_offload(self): self.aux = aux if len(self.fwd_gpu_tensor_group.tensor_list) > 0: - self.finish_offload_event = te_platform().Event() + self.finish_offload_event = torch.cuda.Event() self.finish_offload_event.record(self.offload_stream) def release_activation_forward_gpu_memory(self): @@ -322,7 +322,7 @@ def release_activation_forward_gpu_memory(self): ) self.state = "offload_finished" if len(self.fwd_gpu_tensor_group.tensor_list) > 0: - te_platform().current_stream().wait_event(self.finish_offload_event) # type: ignore[arg-type] + torch.cuda.current_stream().wait_event(self.finish_offload_event) # type: ignore[arg-type] # GPU memory can be released safely after the offload. # Notice that the memory needs to be kept alive when GPU->CPU copy is performed. @@ -349,12 +349,12 @@ def start_reload(self): # calling cudaFree and cudaMalloc again. reloaded_tensor = torch.empty_like(tensor, device=torch.device(te_device_type())) - self.offload_stream.wait_stream(te_platform().current_stream()) + self.offload_stream.wait_stream(torch.cuda.current_stream()) - with te_platform().stream(self.offload_stream): + with torch.cuda.stream(self.offload_stream): reloaded_tensor.copy_(tensor, non_blocking=True) - reload_tensor_event = te_platform().Event() + reload_tensor_event = torch.cuda.Event() reload_tensor_event.record(self.offload_stream) self.bwd_gpu_tensor_group.events.append(reload_tensor_event) self.bwd_gpu_tensor_group.tensor_list.append(reloaded_tensor) @@ -395,8 +395,8 @@ def push_tensor(self, tensor: torch.Tensor) -> int | torch.Tensor | tuple[list, if hasattr(tensor, "start_reload_event"): self.fwd_gpu_tensor_group.events.append(tensor.start_reload_event) else: - self.fwd_gpu_tensor_group.events.append(te_platform().Event()) - self.fwd_gpu_tensor_group.events[-1].record(te_platform().current_stream()) + self.fwd_gpu_tensor_group.events.append(torch.cuda.Event()) + self.fwd_gpu_tensor_group.events[-1].record(torch.cuda.current_stream()) return len(self.fwd_gpu_tensor_group.tensor_list) - 1 return tensor @@ -439,7 +439,7 @@ def pop_tensor( f"but got state='{self.state}' for tensor={tensor_or_tensor_id}" ) # wait for the tensor to be reloaded - te_platform().current_stream().wait_event( + torch.cuda.current_stream().wait_event( self.bwd_gpu_tensor_group.events[tensor_or_tensor_id] ) return self.bwd_gpu_tensor_group.tensor_list[tensor_or_tensor_id] @@ -507,9 +507,7 @@ def __init__( offload_stream: Optional[torch.cuda.Stream] = None, ): self.num_layers = num_layers - self.offload_stream = ( - offload_stream if offload_stream is not None else te_platform().Stream() - ) + self.offload_stream = offload_stream if offload_stream is not None else torch.cuda.Stream() self.layer_states = { i: OffloadableLayerState(self.offload_stream, retain_pinned_cpu_buffers) @@ -778,7 +776,7 @@ def get_cpu_offload_context( .. code-block:: python - offload_stream = te_platform().Stream() + offload_stream = torch.cuda.Stream() cpu_offload_context, sync_function, manual_controller = get_cpu_offload_context( enabled=True, model_layers=num_layers, manual_synchronization=True, offload_stream=offload_stream) diff --git a/transformer_engine/pytorch/cpu_offload_v1.py b/transformer_engine/pytorch/cpu_offload_v1.py index 30ddd57da2c..fb62546cc0a 100644 --- a/transformer_engine/pytorch/cpu_offload_v1.py +++ b/transformer_engine/pytorch/cpu_offload_v1.py @@ -12,7 +12,6 @@ from transformer_engine.debug.pytorch.debug_state import TEDebugState from .quantized_tensor import QuantizedTensorStorage from .tensor.float8_tensor import Float8Tensor -from transformer_engine import te_platform __all__ = ["get_cpu_offload_context"] @@ -364,8 +363,8 @@ def __init__( self.layer_window_map[i] += constant # allocate streams and events for synchronization - self.d2h_stream = te_platform().Stream() - self.h2d_stream = te_platform().Stream() + self.d2h_stream = torch.cuda.Stream() + self.h2d_stream = torch.cuda.Stream() def tensor_push(self, tensor: torch.Tensor, **kwargs) -> Any: global CPUOffloadedLayer @@ -464,7 +463,7 @@ def tensor_pop(self, tensor_tag, **kwargs): def bulk_offload_group(self, group_to_offload): """Bulk offload group.""" - with te_platform().stream(self.d2h_stream): + with torch.cuda.stream(self.d2h_stream): for tensor_tag, state in self.tensor_tag_to_state.items(): group_id, _ = tensor_tag if group_id == group_to_offload: @@ -508,7 +507,7 @@ def synchronize_on_group_commit_forward(self, current_group): # For the first group, kickstart the offload after we have # the first compute completion if current_group == 0: - self.d2h_stream.wait_stream(te_platform().current_stream()) + self.d2h_stream.wait_stream(torch.cuda.current_stream()) if not self.double_buffer_created: # Creating the first copy of double buffer for tensors that are offloaded @@ -530,8 +529,8 @@ def synchronize_on_group_commit_forward(self, current_group): if self.layer_window_map[self.offloaded_group_count] == current_group: # Stream synchronization both ways - self.d2h_stream.wait_stream(te_platform().current_stream()) - te_platform().current_stream().wait_stream(self.d2h_stream) + self.d2h_stream.wait_stream(torch.cuda.current_stream()) + torch.cuda.current_stream().wait_stream(self.d2h_stream) # Time to free the activation memory after usage for tensor_tag, tensor_buf in self.tensor_tag_to_buf.items(): @@ -578,9 +577,9 @@ def bulk_reload_group(self, group_to_reload): buffer_idx = 0 double_buffer_idx = group_to_reload % 2 - main_stream = te_platform().current_stream() + main_stream = torch.cuda.current_stream() - with te_platform().stream(self.h2d_stream): + with torch.cuda.stream(self.h2d_stream): # move back tensors for tensor_label, state in self.tensor_tag_to_state.items(): group_id, _ = tensor_label @@ -590,9 +589,9 @@ def bulk_reload_group(self, group_to_reload): if self.double_buffering: reload_buffer = self.reload_double_buffer[double_buffer_idx][buffer_idx] else: - with te_platform().stream(main_stream): + with torch.cuda.stream(main_stream): reload_buffer = torch.empty_like( - state[1], device=te_platform().current_device() + state[1], device=torch.cuda.current_device() ) recovered_tensor = SynchronizedGroupOffloadHandler.reload( @@ -610,9 +609,9 @@ def bulk_reload_group(self, group_to_reload): buffer_idx ] else: - with te_platform().stream(main_stream): + with torch.cuda.stream(main_stream): reload_buffer = torch.empty_like( - state_tuple[1], device=te_platform().current_device() + state_tuple[1], device=torch.cuda.current_device() ) tensor_list.append( @@ -655,8 +654,8 @@ def on_group_commit_backward(self): if self.layer_window_map[self.offloaded_group_count - 1] == self.current_group: # Stream synchronization both ways - self.h2d_stream.wait_stream(te_platform().current_stream()) - te_platform().current_stream().wait_stream(self.h2d_stream) + self.h2d_stream.wait_stream(torch.cuda.current_stream()) + torch.cuda.current_stream().wait_stream(self.h2d_stream) # Time to reload the next group self.bulk_reload_group(self.offloaded_group_count - 1) @@ -666,7 +665,7 @@ def on_group_commit_backward(self): # Last group computation needs to wait till all the reloads complete if self.current_group == 0: - te_platform().current_stream().wait_stream(self.h2d_stream) + torch.cuda.current_stream().wait_stream(self.h2d_stream) self.offloaded_group_count = 0 diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index 31442d973b5..d12c43d608e 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -20,7 +20,7 @@ from torch.distributed.fsdp import FullyShardedDataParallel as FSDP from torch.distributed.fsdp._common_utils import _get_module_fsdp_state from torch.distributed.fsdp._traversal_utils import _get_fsdp_states_with_modules -from transformer_engine import te_platform, te_device_type +from transformer_engine import te_device_type try: import torch.distributed._symmetric_memory as symm_mem @@ -110,7 +110,7 @@ def _get_cuda_rng_state( device = torch.device(te_device_type(), device) idx = device.index if idx is None: - idx = te_platform().current_device() + idx = torch.cuda.current_device() default_generator = torch.cuda.default_generators[idx] if graph_safe_rng_available() and graph_safe: if clone: @@ -138,7 +138,7 @@ def _set_cuda_rng_state( def cb() -> None: idx = device.index if idx is None: - idx = te_platform().current_device() + idx = torch.cuda.current_device() default_generator = torch.cuda.default_generators[idx] if graph_safe_rng_available() and graph_safe: default_generator.graphsafe_set_state(new_state) @@ -217,7 +217,7 @@ def split_tensor_into_1d_equal_chunks( data = torch.empty( partition_size, dtype=tensor.dtype, - device=te_platform().current_device(), + device=torch.cuda.current_device(), requires_grad=False, ) data.copy_(tensor.view(-1)[start_index:end_index]) @@ -232,7 +232,7 @@ def gather_split_1d_tensor(tensor: torch.Tensor, tp_group: dist_group_type) -> t gathered = torch.empty( numel_gathered, dtype=tensor.dtype, - device=te_platform().current_device(), + device=torch.cuda.current_device(), requires_grad=False, ) torch.distributed.all_gather_into_tensor(gathered, tensor, group=tp_group) @@ -964,7 +964,7 @@ def reduce_scatter_along_first_dim( dim_size[0] = dim_size[0] // world_size if output is None: - output = torch.empty(dim_size, dtype=inp.dtype, device=te_platform().current_device()) + output = torch.empty(dim_size, dtype=inp.dtype, device=torch.cuda.current_device()) handle = torch.distributed.reduce_scatter_tensor( output, inp.contiguous(), group=tp_group, async_op=async_op ) @@ -1965,7 +1965,7 @@ def release_symm_mem_pool() -> None: "release_symm_mem_pool()." ) pools.pop(device, None) - te_platform().empty_cache() + torch.cuda.empty_cache() def symm_mem_alloc( @@ -1986,7 +1986,7 @@ def symm_mem_alloc( be a shared static buffer. ``backend`` selects the symm-mem backend (default NCCL; for the pool it is captured at pool creation).""" if device is None: - device = torch.device(te_device_type(), te_platform().current_device()) + device = torch.device(te_device_type(), torch.cuda.current_device()) if not HAS_TORCH_SYMMETRIC: raise RuntimeError( "torch.distributed._symmetric_memory is unavailable; symm_mem_alloc " diff --git a/transformer_engine/pytorch/ep.py b/transformer_engine/pytorch/ep.py index 0508d5d2653..a4ea00e3dbc 100644 --- a/transformer_engine/pytorch/ep.py +++ b/transformer_engine/pytorch/ep.py @@ -18,7 +18,7 @@ from .cpu_offload import mark_not_offload from .distributed import symm_mem_alloc, release_symm_mem_pool from .quantized_tensor import QuantizedTensor -from transformer_engine import te_platform, te_device_type +from transformer_engine import te_device_type # Type-hint-only import; keeps the ``Recipe`` annotation without a runtime import of # common.recipe (the concrete recipe classes are imported lazily where used). @@ -159,7 +159,7 @@ def ep_bootstrap( ) # Materialize the PG's NCCL comm before borrowing its raw handle. - dist.barrier(group=ep_group, device_ids=[te_platform().current_device()]) + dist.barrier(group=ep_group, device_ids=[torch.cuda.current_device()]) comm_ptr = ep_group._get_backend(torch.device(te_device_type()))._comm_ptr() tex.ep_initialize( @@ -275,7 +275,7 @@ def __init__( if not _BOOTSTRAPPED: raise RuntimeError("EpBuffer requires ep_bootstrap() to be called first.") if device is None: - device = torch.device(te_device_type(), te_platform().current_device()) + device = torch.device(te_device_type(), torch.cuda.current_device()) alignment = int(alignment) if alignment > 1 and (alignment & (alignment - 1)) != 0: raise ValueError(f"alignment must be 0, 1, or a power of two (got {alignment}).") @@ -898,7 +898,7 @@ def _alloc_io(shape, dtype: torch.dtype, device, zero_copy: bool) -> torch.Tenso The zero-copy pool path is not CUDA-graph capturable; supply persistent recv_tokens / grad_out buffers to capture a graph.""" if zero_copy: - if te_platform().is_current_stream_capturing(): + if torch.cuda.is_current_stream_capturing(): raise RuntimeError( "EP zero-copy pool allocation is not CUDA-graph capturable; supply persistent " "recv_tokens / grad_out buffers (allocated once via symm_mem_alloc) before capture." diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index 40374a62437..2e4b3c99acf 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -18,7 +18,6 @@ from transformer_engine.common.recipe import DelayedScaling, Recipe from transformer_engine.pytorch.constants import dist_group_type -from transformer_engine import te_platform from .quantization import ( autocast, FP8GlobalStateManager, @@ -105,7 +104,7 @@ def _graph_context_wrapper(*args, **kwargs): gc_is_enabled = gc.isenabled() if gc_is_enabled: gc.disable() - with te_platform().graph(*args, **kwargs): + with torch.cuda.graph(*args, **kwargs): yield if gc_is_enabled: gc.enable() @@ -460,9 +459,9 @@ def _make_graphed_callables( for i in range(len(flatten_sample_args)) ] - fwd_graphs = [te_platform().CUDAGraph() for _ in range(len(flatten_sample_args))] - bwd_graphs = [te_platform().CUDAGraph() for _ in range(len(flatten_sample_args))] - bwd_dw_graphs = [te_platform().CUDAGraph() for _ in range(len(flatten_sample_args))] + fwd_graphs = [torch.cuda.CUDAGraph() for _ in range(len(flatten_sample_args))] + bwd_graphs = [torch.cuda.CUDAGraph() for _ in range(len(flatten_sample_args))] + bwd_dw_graphs = [torch.cuda.CUDAGraph() for _ in range(len(flatten_sample_args))] graph_callables = [None for _ in range(len(flatten_sample_args))] def _returned_param_grad_clone_slots(static_grad_inputs, module_params): @@ -496,7 +495,7 @@ def _returned_param_grad_clone_slots(static_grad_inputs, module_params): # Warmup # Hopefully prevents cudnn benchmarking and other lazy-initialization cuda work # from ending up in any captures. - te_platform().synchronize() + torch.cuda.synchronize() # Get warmup func and func_idx. warmup_func_idx = [] @@ -632,7 +631,7 @@ def _run_warmup_backward(func_idx, func, outputs, warmup_iter, callable_idx): need_bwd_dw_graph[func_idx] = need_backward_dw # Run warmup and do the above filtering. - with te_platform().stream(te_platform().Stream()): + with torch.cuda.stream(torch.cuda.Stream()): if pre_warmup_hook is not None: pre_warmup_hook() @@ -690,7 +689,7 @@ def _run_warmup_backward(func_idx, func, outputs, warmup_iter, callable_idx): if post_warmup_hook is not None: post_warmup_hook() - te_platform().synchronize() + torch.cuda.synchronize() # All captures here share a mempool. To avoid replays corrupting each other's memory, # the safest approach is to capture all passes in the same order they'll run: @@ -1083,14 +1082,14 @@ def forward(ctx, skip_fp8_weight_update, cuda_graph_stream, cuda_graph_event, *i static_input_surface[i].copy_(inputs[i]) # Replay forward graph - if cuda_graph_stream != te_platform().current_stream(): - cuda_graph_stream.wait_stream(te_platform().current_stream()) + if cuda_graph_stream != torch.cuda.current_stream(): + cuda_graph_stream.wait_stream(torch.cuda.current_stream()) with cuda_graph_stream: fwd_graph.replay() if cuda_graph_event is not None: - te_platform().current_stream().wait_event(cuda_graph_event) + torch.cuda.current_stream().wait_event(cuda_graph_event) else: - te_platform().current_stream().wait_stream(cuda_graph_stream) + torch.cuda.current_stream().wait_stream(cuda_graph_stream) else: fwd_graph.replay() if not isinstance(static_outputs, tuple): @@ -1119,14 +1118,14 @@ def backward(ctx, *grads): # incoming grad is already in the right place if g.data_ptr() != grad.data_ptr(): g.copy_(grad) - if ctx.cuda_graph_stream != te_platform().current_stream(): - ctx.cuda_graph_stream.wait_stream(te_platform().current_stream()) + if ctx.cuda_graph_stream != torch.cuda.current_stream(): + ctx.cuda_graph_stream.wait_stream(torch.cuda.current_stream()) with ctx.cuda_graph_stream: bwd_graph.replay() if ctx.cuda_graph_event is not None: - te_platform().current_stream().wait_event(ctx.cuda_graph_event) + torch.cuda.current_stream().wait_event(ctx.cuda_graph_event) else: - te_platform().current_stream().wait_stream(ctx.cuda_graph_stream) + torch.cuda.current_stream().wait_stream(ctx.cuda_graph_stream) else: bwd_graph.replay() @@ -1179,7 +1178,7 @@ def functionalized(*user_args, **user_kwargs): cuda_graph_stream = user_kwargs["cuda_graph_stream"] user_kwargs.pop("cuda_graph_stream") else: - cuda_graph_stream = te_platform().current_stream() + cuda_graph_stream = torch.cuda.current_stream() if "cuda_graph_event" in user_kwargs: cuda_graph_event = user_kwargs["cuda_graph_event"] user_kwargs.pop("cuda_graph_event") @@ -1691,12 +1690,12 @@ def call_func(self, *args, **kwargs): # Save RNG state. if graph_safe_rng_available(): generators = [ - torch.cuda.default_generators[te_platform().current_device()], + torch.cuda.default_generators[torch.cuda.current_device()], *get_all_rng_states().values(), ] original_rng_states = [state.get_state() for state in generators] else: - original_rng_states = te_platform().get_rng_state() + original_rng_states = torch.cuda.get_rng_state() graphed_callables = _make_graphed_callables( forward_funcs, @@ -1721,7 +1720,7 @@ def call_func(self, *args, **kwargs): for gen, state in zip(generators, original_rng_states): gen.set_state(state) else: - te_platform().set_rng_state(original_rng_states) + torch.cuda.set_rng_state(original_rng_states) # Remove FP8 wrapper. for module_cls, old_call in old_call_funcs.items(): diff --git a/transformer_engine/pytorch/jit.py b/transformer_engine/pytorch/jit.py index 7ca52b0aa75..ece60737138 100644 --- a/transformer_engine/pytorch/jit.py +++ b/transformer_engine/pytorch/jit.py @@ -13,7 +13,6 @@ from .torch_version import torch_version from .export import is_in_onnx_export_mode from .utils import gpu_autocast_ctx -from transformer_engine import te_platform # pylint: disable=unnecessary-lambda-assignment @@ -336,7 +335,7 @@ def warmup_jit_bias_dropout_add( """Compile BDA JIT function before the main training steps""" # Save cuda RNG state to ensure warmup does not affect reproducibility. - rng_state = te_platform().get_rng_state() + rng_state = torch.cuda.get_rng_state() inp = torch.rand( (seq_length, micro_batch_size, hidden_size), @@ -360,8 +359,8 @@ def warmup_jit_bias_dropout_add( output = bias_dropout_add_fused_train(inp, bias, residual, dropout_rate) del bias, inp, residual, output - te_platform().empty_cache() - te_platform().set_rng_state(rng_state) + torch.cuda.empty_cache() + torch.cuda.set_rng_state(rng_state) def warmup_jit_bias_dropout_add_all_dtypes( @@ -381,7 +380,7 @@ def warmup_jit_bias_gelu( """Compile bias-gelu JIT function before the main training steps""" # Save cuda RNG state to ensure warmup does not affect reproducibility. - rng_state = te_platform().get_rng_state() + rng_state = torch.cuda.get_rng_state() bias = torch.rand(ffn_hidden_size_per_partition, dtype=dtype, device=te_device_type()) inp = torch.rand( @@ -398,8 +397,8 @@ def warmup_jit_bias_gelu( _ = gelu_fused_(inp) del bias, inp - te_platform().empty_cache() - te_platform().set_rng_state(rng_state) + torch.cuda.empty_cache() + torch.cuda.set_rng_state(rng_state) def warmup_jit_bias_gelu_all_dtypes( @@ -416,7 +415,7 @@ def warmup_jit_l2normalization( """Compile L2Normalization JIT function before the main training steps""" # Save cuda RNG state to ensure warmup does not affect reproducibility. - rng_state = te_platform().get_rng_state() + rng_state = torch.cuda.get_rng_state() inp = torch.rand( (seq_length * micro_batch_size, hidden_size), @@ -440,8 +439,8 @@ def warmup_jit_l2normalization( output = l2normalization_fused_(inp, eps) del inp, output - te_platform().empty_cache() - te_platform().set_rng_state(rng_state) + torch.cuda.empty_cache() + torch.cuda.set_rng_state(rng_state) def warmup_jit_l2normalization_all_dtypes( diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 0074aa87b84..30a01bd6d11 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -21,7 +21,7 @@ import transformer_engine_torch as tex from ._common import _ParameterInitMeta, noop_cat -from transformer_engine import te_platform, te_device_type +from transformer_engine import te_device_type from .._extra_state import ( extra_state_pickle_advisory, is_stateless_recipe, @@ -893,7 +893,7 @@ class TransformerEngineBaseModule(torch.nn.Module, ABC): def __init__(self, name: Optional[str] = None) -> None: super().__init__() - if not te_platform().is_available(): + if not torch.cuda.is_available(): raise RuntimeError("TransformerEngine needs CUDA.") self.name = name self.next_iter_when_debug_should_be_run = 0 @@ -1412,7 +1412,7 @@ def to_cpu(src: torch.Tensor) -> torch.Tensor: state["extra_fp8_variables"] = extra # Serialize state into byte tensor - te_platform().synchronize() + torch.cuda.synchronize() state_serialized = bytearray(pickle.dumps(state)) state_serialized = torch.frombuffer(state_serialized, dtype=torch.uint8) return state_serialized @@ -1478,7 +1478,7 @@ def copy_tensor(src: torch.Tensor, dst: torch.Tensor) -> None: copy_tensor(state["amax_history_fwd"], self.fp8_meta["scaling_fwd"].amax_history) copy_tensor(state["scale_bwd"], self.fp8_meta["scaling_bwd"].scale) copy_tensor(state["amax_history_bwd"], self.fp8_meta["scaling_bwd"].amax_history) - te_platform().synchronize() + torch.cuda.synchronize() def set_activation_dtype(self, inp: torch.Tensor) -> None: """Get activation data type for AMP.""" diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index a09027068eb..666e44dd294 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -18,7 +18,7 @@ from transformer_engine.common.recipe import Recipe from transformer_engine.pytorch.torch_version import torch_version from transformer_engine.pytorch.tensor.utils import clear_columnwise_cache, is_custom -from transformer_engine import te_platform, te_device_type +from transformer_engine import te_device_type from .base import ( fill_userbuffers_buffer_for_all_gather, get_ub, @@ -967,7 +967,7 @@ def backward( # We use the send stream to copy into the userbuffers. # This is the same stream that we will use to access the data in the AG, # so we dont need to add any syncs yet. - with te_platform().stream(dgrad_send_stream): + with torch.cuda.stream(dgrad_send_stream): grad_output, _ = fill_userbuffers_buffer_for_all_gather( ub_obj_overlap_wgrad, grad_outputs[0], diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index 5952560ffa0..9d2584035c5 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -19,7 +19,7 @@ from transformer_engine.common.recipe import Recipe from transformer_engine.pytorch.torch_version import torch_version from transformer_engine.pytorch.tensor.utils import clear_columnwise_cache, is_custom -from transformer_engine import te_platform, te_device_type +from transformer_engine import te_device_type from .base import ( fill_userbuffers_buffer_for_all_gather, _ub_communicators, @@ -1333,7 +1333,7 @@ def backward( # We use the send stream to copy into the userbuffers. # This is the same stream that we will use to access the data in the AG, # so we dont need to add any syncs yet. - with te_platform().stream(dgrad_send_stream): + with torch.cuda.stream(dgrad_send_stream): grad_output, _ = fill_userbuffers_buffer_for_all_gather( ub_obj_fc2_wgrad, grad_outputs[0], @@ -1767,14 +1767,14 @@ def fc1_wgrad_gemm( fc1_wgrad = torch.zeros( fc1_weight_main_grad.shape, dtype=fc1_weight_python_object.dtype, - device=te_platform().current_device(), + device=torch.cuda.current_device(), requires_grad=False, ) else: fc1_wgrad = torch.empty( fc1_weight_main_grad.shape, dtype=fc1_weight_python_object.dtype, - device=te_platform().current_device(), + device=torch.cuda.current_device(), requires_grad=False, ) elif ctx.fuse_wgrad_accumulation: @@ -1792,14 +1792,14 @@ def fc1_wgrad_gemm( fc2_wgrad = torch.zeros( fc2_weight_main_grad.shape, dtype=fc2_weight_python_object.dtype, - device=te_platform().current_device(), + device=torch.cuda.current_device(), requires_grad=False, ) else: fc2_wgrad = torch.empty( fc2_weight_main_grad.shape, dtype=fc2_weight_python_object.dtype, - device=te_platform().current_device(), + device=torch.cuda.current_device(), requires_grad=False, ) elif ctx.fuse_wgrad_accumulation: diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index ecf80bf41fd..6dad718eded 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -18,7 +18,7 @@ from transformer_engine.common.recipe import Recipe from transformer_engine.pytorch.torch_version import torch_version -from transformer_engine import te_platform, te_device_type +from transformer_engine import te_device_type from .base import ( fill_userbuffers_buffer_for_all_gather, get_dummy_wgrad, @@ -1540,7 +1540,7 @@ def _linear_backward_impl(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None # We use the send stream to copy into the userbuffers. # This is the same stream that we will use to access the data in the AG, # so we dont need to add any syncs yet. - with te_platform().stream(dgrad_send_stream): + with torch.cuda.stream(dgrad_send_stream): grad_output, _ = fill_userbuffers_buffer_for_all_gather( ub_obj_overlap_wgrad, grad_output_arg, diff --git a/transformer_engine/pytorch/ops/basic/l2normalization.py b/transformer_engine/pytorch/ops/basic/l2normalization.py index 05106c7833c..be155c93565 100644 --- a/transformer_engine/pytorch/ops/basic/l2normalization.py +++ b/transformer_engine/pytorch/ops/basic/l2normalization.py @@ -12,7 +12,6 @@ from ...torch_version import torch_version from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload -from transformer_engine import te_platform from ...jit import ( l2normalization_fused, l2normalization_fwd_fused, @@ -65,7 +64,7 @@ def __init__( # JIT warmup for L2Normalization fused operations if seq_length and micro_batch_size: if ( - te_platform().is_available() + torch.cuda.is_available() and torch_version() >= (2, 0, 0) and bool(int(os.getenv("NVTE_TORCH_COMPILE", "1"))) ): diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 0fffb566ba5..6286b7622a2 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -20,7 +20,6 @@ from ...constants import DType, MXFP8_BLOCK_SCALING_SIZE, NVFP4_BLOCK_SCALING_SIZE, TE_DType from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload, start_offload from ...cpp_extensions import general_gemm, general_grouped_gemm_for_grouped_tensor -from transformer_engine import te_platform from ...distributed_weight import ( is_distributed_weight, materialize_weight_for_forward, @@ -549,7 +548,7 @@ def _cudnn_compute_wgrad( b = X = (total_tokens, in_features) column-major. """ if current_stream is None: - current_stream = te_platform().current_stream().cuda_stream + current_stream = torch.cuda.current_stream().cuda_stream out_features, in_features = weight_shape total_tokens = grouped_dy.logical_shape[0] diff --git a/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py b/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py index fb04ee8629f..d37d9d88098 100644 --- a/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py +++ b/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py @@ -13,7 +13,6 @@ from transformer_engine_torch import CommOverlapType, bulk_overlap_ag_with_external_gemm from ...cpp_extensions import general_gemm from ...distributed import gather_along_first_dim, get_distributed_world_size -from transformer_engine import te_platform from ...module.base import ( _2X_ACC_DGRAD, _2X_ACC_WGRAD, @@ -482,7 +481,7 @@ def _functional_backward( # We use the send stream to copy into the userbuffers. # This is the same stream that we will use to access the data in the AG, # so we dont need to add any syncs yet. - with te_platform().stream(dgrad_send_stream): + with torch.cuda.stream(dgrad_send_stream): dy, _ = fill_userbuffers_buffer_for_all_gather( ub_obj_overlap_wgrad, dy_local, diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py index 6ffddd6e225..d057d46816e 100644 --- a/transformer_engine/pytorch/ops/op.py +++ b/transformer_engine/pytorch/ops/op.py @@ -15,7 +15,6 @@ from transformer_engine.common.recipe import Recipe from .._extra_state import is_stateless_recipe, should_load_extra_state_pickle -from transformer_engine import te_platform from ..quantization import ( FP8GlobalStateManager, QuantizerRole, @@ -704,7 +703,7 @@ def to_cpu(src: torch.Tensor) -> torch.Tensor: return torch.empty(0, dtype=torch.uint8) # Serialize state into byte tensor - te_platform().synchronize() + torch.cuda.synchronize() state_serialized = bytearray(pickle.dumps(state)) state_serialized = torch.frombuffer(state_serialized, dtype=torch.uint8) return state_serialized @@ -767,7 +766,7 @@ def copy_tensor(src: torch.Tensor, dst: torch.Tensor) -> None: ) # Finish CPU-GPU memory transfers - te_platform().synchronize() + torch.cuda.synchronize() def _load_from_state_dict(self, *args, **kwargs) -> None: """Load state""" diff --git a/transformer_engine/pytorch/optimizers/fused_adam.py b/transformer_engine/pytorch/optimizers/fused_adam.py index 6b14cc89568..ff7eb39bab6 100644 --- a/transformer_engine/pytorch/optimizers/fused_adam.py +++ b/transformer_engine/pytorch/optimizers/fused_adam.py @@ -17,7 +17,7 @@ from transformer_engine.pytorch.quantized_tensor import QuantizedTensor from ..constants import DType from .multi_tensor_apply import multi_tensor_applier -from transformer_engine import te_device_type, te_platform +from transformer_engine import te_device_type # Bound temporary NVTETensor handles created before the CUDA launcher chunks its metadata. diff --git a/transformer_engine/pytorch/optimizers/newton_schulz.py b/transformer_engine/pytorch/optimizers/newton_schulz.py index 7d2edd6cd6c..d4696393eb1 100644 --- a/transformer_engine/pytorch/optimizers/newton_schulz.py +++ b/transformer_engine/pytorch/optimizers/newton_schulz.py @@ -11,7 +11,7 @@ import torch.distributed as dist import transformer_engine_torch as tex -from transformer_engine import te_platform, te_device_type +from transformer_engine import te_device_type _COEFFICIENT_SETS = { @@ -182,7 +182,7 @@ def _get_nccl_comm_ptr(group: dist.ProcessGroup) -> int: # ProcessGroupNCCL creates communicators lazily. This device-specific # collective ensures that the communicator returned by _comm_ptr() exists # and is ready on every group rank before cuSolverMp borrows it. - dist.barrier(group=group, device_ids=[te_platform().current_device()]) + dist.barrier(group=group, device_ids=[torch.cuda.current_device()]) nccl_backend = group._get_backend(torch.device(te_device_type())) comm_ptr = nccl_backend._comm_ptr() if not isinstance(comm_ptr, int) or comm_ptr == 0: diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index 7c765479188..10335745dbc 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -17,7 +17,7 @@ import torch import transformer_engine_torch as tex -from transformer_engine import te_platform, te_device_type +from transformer_engine import te_device_type from transformer_engine.common.recipe import ( Recipe, DelayedScaling, @@ -593,8 +593,7 @@ def fp8_graph_capturing(cls) -> bool: assert not cls.quantization_state.fp8_graph_capturing return False return ( - cls.quantization_state.fp8_graph_capturing - or te_platform().is_current_stream_capturing() + cls.quantization_state.fp8_graph_capturing or torch.cuda.is_current_stream_capturing() ) @classmethod diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index 6c80d6652c3..3b91a72decc 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -18,7 +18,7 @@ from transformer_engine.common.recipe import Recipe from transformer_engine.pytorch.constants import dist_group_type -from transformer_engine import te_platform, te_device_type +from transformer_engine import te_device_type from transformer_engine.pytorch.tensor._quantization_helpers import ( _QuantizeFunc, _IdentityFunc, @@ -742,7 +742,7 @@ def __new__( dtype=dtype, layout=torch.strided, requires_grad=requires_grad, - device=te_platform().current_device() if device is None else device, + device=torch.cuda.current_device() if device is None else device, ) instance._requires_grad = requires_grad instance._dtype = dtype diff --git a/transformer_engine/pytorch/router.py b/transformer_engine/pytorch/router.py index ba156cfc6d8..ac38f7d1cf9 100644 --- a/transformer_engine/pytorch/router.py +++ b/transformer_engine/pytorch/router.py @@ -19,7 +19,6 @@ import torch import transformer_engine_torch as tex -from transformer_engine import te_platform # Re-export the C++ enum NVTERoutingMapFormat under a friendlier Python name. # Members: @@ -58,8 +57,8 @@ def _validate_qb_bin_bounds(bin_bounds: torch.Tensor) -> bool: version = bin_bounds._version if getattr(bin_bounds, _QB_BOUNDS_VALIDATED_VERSION_ATTR, None) == version: return True - with te_platform().device(bin_bounds.device): - if te_platform().is_current_stream_capturing(): + with torch.cuda.device(bin_bounds.device): + if torch.cuda.is_current_stream_capturing(): raise RuntimeError( "QB bin_bounds must be validated by an eager router call before CUDA graph capture" ) diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index e3bdf5852a6..5e537c3ff0e 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -15,7 +15,6 @@ from transformer_engine.common.recipe import NVFP4BlockScaling, Recipe from ..constants import NVFP4_BLOCK_SCALING_SIZE, dist_group_type, DType -from transformer_engine import te_platform from ..utils import ( canonicalize_process_group, devices_match, @@ -174,7 +173,7 @@ def __init__( if self.nvfp4_4over6_err_mode not in ("MAE", "MSE"): raise ValueError("nvfp4_4over6_err_mode must be 'MAE' or 'MSE'.") self.rht_matrix_random_sign_mask_t = get_random_sign_mask_for_rht( - with_random_sign_mask, te_platform().current_device() + with_random_sign_mask, torch.cuda.current_device() ) self._rebuild_derived_state() @@ -194,7 +193,7 @@ def _rebuild_derived_state(self) -> None: cheap hit. """ self.rht_matrix = get_rht_matrix( - self.rht_matrix_random_sign_mask_t != 0, te_platform().current_device() + self.rht_matrix_random_sign_mask_t != 0, torch.cuda.current_device() ) def update_quantized( diff --git a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py index bea825bf124..3473024c03f 100644 --- a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py @@ -19,7 +19,6 @@ from .float8_blockwise_tensor_storage import Float8BlockwiseQTensorStorage from .nvfp4_tensor_storage import NVFP4TensorStorage from ...utils import is_non_tn_fp8_gemm_supported -from transformer_engine import te_platform class GroupedTensorStorage: @@ -669,7 +668,7 @@ def make_grouped_tensor( # Set device if device is None: - device = te_platform().current_device() + device = torch.cuda.current_device() # Shape patterns and validation. all_same_first = first_dims is None @@ -714,7 +713,7 @@ def make_grouped_tensor( compatible_recipe = None if no_quantization else quantizer._get_compatible_recipe() if not shape: - if te_platform().is_available() and te_platform().is_current_stream_capturing(): + if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): raise ValueError( "Varying-dimension grouped tensor construction is not graph-safe: it" " requires device-to-host copies of per-tensor shapes/offsets." diff --git a/transformer_engine/pytorch/triton/mhc.py b/transformer_engine/pytorch/triton/mhc.py index 36a0ee5a3df..7288b4af923 100644 --- a/transformer_engine/pytorch/triton/mhc.py +++ b/transformer_engine/pytorch/triton/mhc.py @@ -9,7 +9,6 @@ import torch import triton -from transformer_engine import te_platform from transformer_engine.common.triton.mhc import ( _mhc_projection_bwd_fused_dphi, _mhc_projection_bwd_fused_dx, @@ -32,7 +31,7 @@ def _support_tma(x: torch.Tensor): # get_device_capability returns a (major, minor) tuple; TMA needs Hopper+ (major >= 9) - return te_platform().get_device_capability(x.device)[0] >= 9 + return torch.cuda.get_device_capability(x.device)[0] >= 9 def _tma_aligned(t): diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py index 461d6857d69..de58541a7a3 100644 --- a/transformer_engine/pytorch/utils.py +++ b/transformer_engine/pytorch/utils.py @@ -16,7 +16,7 @@ from .torch_version import torch_version from ..debug.pytorch.debug_quantization import DebugQuantizedTensor -from transformer_engine import te_device_type, te_platform +from transformer_engine import te_device_type __all__ = [ @@ -158,7 +158,7 @@ def _get_device_compute_capability(device: torch.device) -> Tuple[int, int]: @torch.compiler.assume_constant_result def get_device_compute_capability() -> Tuple[int, int]: """CUDA compute capability of current GPU""" - return _get_device_compute_capability(te_platform().current_device()) + return _get_device_compute_capability(torch.cuda.current_device()) def deinterleave_glu_tensor(tensor: torch.Tensor, interleave_size: int) -> torch.Tensor: @@ -347,7 +347,7 @@ def all_close(a: torch.Tensor, b: torch.Tensor) -> bool: def print_rank_0(*args: Any) -> None: """print on rank 0""" - if te_platform().current_device() == 0: + if torch.cuda.current_device() == 0: print(*args) @@ -803,11 +803,11 @@ def canonicalize_device(device: Optional[torch.device | str]) -> torch.device: # Use default CUDA device device = torch.get_default_device() if device.type != te_device_type(): - device = torch.device(te_device_type(), te_platform().current_device()) + device = torch.device(te_device_type(), torch.cuda.current_device()) elif not isinstance(device, torch.device): device = torch.device(device) if device.type == te_device_type() and device.index is None: - device = torch.device(te_device_type(), te_platform().current_device()) + device = torch.device(te_device_type(), torch.cuda.current_device()) return device @@ -835,9 +835,9 @@ def devices_match(device1: torch.device, device2: torch.device) -> bool: if index1 == index2: return True if index1 is None: - index1 = te_platform().current_device() + index1 = torch.cuda.current_device() if index2 is None: - index2 = te_platform().current_device() + index2 = torch.cuda.current_device() return index1 == index2 return device1 == device2 @@ -845,7 +845,7 @@ def devices_match(device1: torch.device, device2: torch.device) -> bool: @functools.lru_cache def get_sm_count() -> int: """Returns the number of streaming multiprocessors in the current device.""" - return torch.cuda.get_device_properties(te_platform().current_device()).multi_processor_count + return torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count def ceil_div(numerator, denominator): From a0f9487533804b79ece2d862cfa651258e3d9892 Mon Sep 17 00:00:00 2001 From: Xianduo Li Date: Mon, 14 Sep 2026 16:47:00 +0800 Subject: [PATCH 08/19] fix: remove remaining te_platform() calls in utils, fused_adam, and distributed Replace remaining te_platform() calls with torch.cuda equivalents: - torch.cuda.nvtx.range() / range_push() / range_pop() for NVTX profiling - torch.cuda.amp.autocast for autocast context manager - Update docstring reference in fused_adam.py All te_platform() usage has been removed except for the definition in transformer_engine/pytorch/__init__.py. Signed-off-by: Xianduo Li --- transformer_engine/pytorch/distributed.py | 2 +- transformer_engine/pytorch/optimizers/fused_adam.py | 2 +- transformer_engine/pytorch/utils.py | 8 ++++---- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index d12c43d608e..14e781536e0 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -315,7 +315,7 @@ def _get_active_autocast_contexts(): else: gpu_autocast_enabled = torch.is_autocast_enabled() gpu_autocast_dtype = torch.get_autocast_gpu_dtype() - gpu_autocast_ctx = te_platform().amp.autocast( + gpu_autocast_ctx = torch.cuda.amp.autocast( gpu_autocast_enabled, gpu_autocast_dtype, autocast_cached ) diff --git a/transformer_engine/pytorch/optimizers/fused_adam.py b/transformer_engine/pytorch/optimizers/fused_adam.py index ff7eb39bab6..0a5d8173fb7 100644 --- a/transformer_engine/pytorch/optimizers/fused_adam.py +++ b/transformer_engine/pytorch/optimizers/fused_adam.py @@ -554,7 +554,7 @@ def step(self, closure=None, grad_scaler=None): Arguments: closure (callable, optional): A closure that reevaluates the model and returns the loss. - grad_scaler (te_platform().amp.GradScaler, optional): + grad_scaler (torch.cuda.amp.GradScaler, optional): gradient scaler (default: None) """ loss = None diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py index de58541a7a3..2ab9f1d7bd3 100644 --- a/transformer_engine/pytorch/utils.py +++ b/transformer_engine/pytorch/utils.py @@ -899,7 +899,7 @@ def get_nvtx_range_context(msg: str): """ if _nvtx_enabled(): - return te_platform().nvtx.range(msg) + return torch.cuda.nvtx.range(msg) return nullcontext() @@ -918,7 +918,7 @@ def nvtx_range_push(msg: str) -> None: if not _nvtx_enabled(): return _nvtx_range_messages.append(msg) - te_platform().nvtx.range_push(msg) + torch.cuda.nvtx.range_push(msg) def nvtx_range_pop(msg: Optional[str] = None) -> None: @@ -949,7 +949,7 @@ def nvtx_range_pop(msg: Optional[str] = None) -> None: ) # Pop NVTX range - te_platform().nvtx.range_pop() + torch.cuda.nvtx.range_pop() def canonicalize_process_group( @@ -975,7 +975,7 @@ def torch_get_autocast_gpu_dtype() -> torch.dtype: if torch_version() >= (2, 4, 0): gpu_autocast_ctx = functools.partial(torch.amp.autocast, device_type=te_device_type()) else: - gpu_autocast_ctx = te_platform().amp.autocast + gpu_autocast_ctx = torch.cuda.amp.autocast _torch_dtype_to_np_typestr_dict = { From 1bad7eea0f9d46516835deef6722a9164ebb4fbd Mon Sep 17 00:00:00 2001 From: Xianduo Li Date: Tue, 22 Sep 2026 11:29:45 +0800 Subject: [PATCH 09/19] keep sync with torch.cuda api usage Signed-off-by: Xianduo Li --- transformer_engine/pytorch/cpu_offload.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/transformer_engine/pytorch/cpu_offload.py b/transformer_engine/pytorch/cpu_offload.py index 73604634721..c81c18e64fc 100644 --- a/transformer_engine/pytorch/cpu_offload.py +++ b/transformer_engine/pytorch/cpu_offload.py @@ -16,7 +16,6 @@ from transformer_engine.debug.pytorch.debug_state import TEDebugState import transformer_engine.pytorch as te import transformer_engine.pytorch.cpu_offload_v1 as v1_code_path -from transformer_engine import te_device_type from .quantized_tensor import ( restore_from_saved, prepare_for_saving, @@ -348,7 +347,7 @@ def start_reload(self): # cannot move tensors from pool of one stream to another without # calling cudaFree and cudaMalloc again. - reloaded_tensor = torch.empty_like(tensor, device=torch.device(te_device_type())) + reloaded_tensor = torch.empty_like(tensor, device=torch.device("cuda")) self.offload_stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(self.offload_stream): @@ -464,7 +463,7 @@ def _check_if_offload(self, t: torch.Tensor) -> bool: not isinstance(t, torch.nn.Parameter) and not getattr(t, "_TE_do_not_offload", False) and not isinstance(t, torch._subclasses.FakeTensor) - and t.device.type == te_device_type() + and t.device.type == "cuda" ): if not t.is_contiguous() and not getattr(t, "offload_base_tensor", False): warnings.warn( From 2d63707dc4fb3ccaa76103cd5c8a2da018daa585 Mon Sep 17 00:00:00 2001 From: Xianduo Li Date: Thu, 8 Oct 2026 14:55:41 +0800 Subject: [PATCH 10/19] format Signed-off-by: Xianduo Li --- transformer_engine/pytorch/ep.py | 1 - transformer_engine/pytorch/optimizers/newton_schulz.py | 1 - 2 files changed, 2 deletions(-) diff --git a/transformer_engine/pytorch/ep.py b/transformer_engine/pytorch/ep.py index 474cbae7a1f..ed42c2ecfeb 100644 --- a/transformer_engine/pytorch/ep.py +++ b/transformer_engine/pytorch/ep.py @@ -18,7 +18,6 @@ from .constants import DType from .cpu_offload import mark_not_offload from .distributed import symm_mem_alloc, release_symm_mem_pool - from .quantized_tensor import QuantizedTensor, QuantizedTensorStorage from .tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage from transformer_engine import te_device_type diff --git a/transformer_engine/pytorch/optimizers/newton_schulz.py b/transformer_engine/pytorch/optimizers/newton_schulz.py index 0cdebc44a01..e1b1a0bcbde 100644 --- a/transformer_engine/pytorch/optimizers/newton_schulz.py +++ b/transformer_engine/pytorch/optimizers/newton_schulz.py @@ -11,7 +11,6 @@ import torch.distributed as dist import transformer_engine_torch as tex -from transformer_engine import te_device_type _COEFFICIENT_SETS = { From 90d6d322792a343dbc9e8b05c5c5a828bdca98b9 Mon Sep 17 00:00:00 2001 From: Xianduo Li Date: Mon, 19 Oct 2026 13:09:40 +0800 Subject: [PATCH 11/19] Fix missing te_device_type imports and import ordering Add the missing `te_device_type` import to eight modules that called it without importing it, which would raise NameError at runtime: jit, attention/inference, attention/dot_product_attention/{backends, context_parallel,utils}, ops/fused/{grouped_mlp, userbuffers_backward_linear}, and triton/mhc. Move the `te_device_type` import ahead of relative imports where pylint classifies it as a first-party import that must precede local ones. In jit and distributed, place it so it does not split the torch import group. Order `transformer_engine_torch` before `transformer_engine` in ops/_common. Add docstrings to `te_device_type` and `te_platform`. Signed-off-by: Xianduo Li --- transformer_engine/pytorch/__init__.py | 2 ++ .../pytorch/attention/dot_product_attention/backends.py | 1 + .../pytorch/attention/dot_product_attention/context_parallel.py | 1 + .../pytorch/attention/dot_product_attention/utils.py | 1 + transformer_engine/pytorch/attention/inference.py | 1 + transformer_engine/pytorch/cpp_extensions/gemm.py | 2 +- .../pytorch/custom_recipes/quantizer_factories.py | 2 +- .../pytorch/custom_recipes/quantizer_factory_zoo.py | 2 +- transformer_engine/pytorch/distributed.py | 2 +- transformer_engine/pytorch/ep.py | 2 +- transformer_engine/pytorch/jit.py | 1 + transformer_engine/pytorch/module/base.py | 2 +- transformer_engine/pytorch/ops/_common.py | 2 +- transformer_engine/pytorch/ops/basic/activation.py | 2 +- transformer_engine/pytorch/ops/basic/bias.py | 2 +- transformer_engine/pytorch/ops/basic/grouped_linear.py | 2 +- transformer_engine/pytorch/ops/basic/swiglu.py | 2 +- transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py | 2 +- .../pytorch/ops/fused/forward_linear_scale_add.py | 2 +- transformer_engine/pytorch/ops/fused/grouped_mlp.py | 1 + .../pytorch/ops/fused/userbuffers_backward_linear.py | 1 + .../pytorch/ops/fused/userbuffers_forward_linear.py | 2 +- transformer_engine/pytorch/optimizers/fused_adam.py | 2 +- transformer_engine/pytorch/tensor/grouped_tensor.py | 2 +- transformer_engine/pytorch/tensor/identity_tensor.py | 2 +- .../pytorch/tensor/storage/mxfp8_tensor_storage.py | 2 +- .../pytorch/tensor/storage/nvfp4_tensor_storage.py | 2 +- transformer_engine/pytorch/triton/mhc.py | 1 + transformer_engine/pytorch/utils.py | 2 +- 29 files changed, 30 insertions(+), 20 deletions(-) diff --git a/transformer_engine/pytorch/__init__.py b/transformer_engine/pytorch/__init__.py index cb53a18f893..fcd0a1d3ce8 100644 --- a/transformer_engine/pytorch/__init__.py +++ b/transformer_engine/pytorch/__init__.py @@ -26,10 +26,12 @@ def te_device_type() -> str: + """Torch device type used by Transformer Engine, e.g. `cuda`.""" return transformer_engine.TE_DEVICE_TYPE def te_platform(): + """Torch device module used by Transformer Engine, e.g. `torch.cuda`.""" return transformer_engine.TE_PLATFORM diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 4692710a5bb..aefe8507c26 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -18,6 +18,7 @@ import torch import torch.nn.functional as F +from transformer_engine import te_device_type from transformer_engine.pytorch.utils import ( get_device_compute_capability, split_tensor_along_dim, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index e30f4ae155a..a45183a358b 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -8,6 +8,7 @@ import torch import transformer_engine_torch as tex +from transformer_engine import te_device_type from transformer_engine.pytorch.utils import ( get_cudnn_version, nvtx_range_pop, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index d04621b061b..608efd92b41 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -22,6 +22,7 @@ import torch.nn.functional as F import transformer_engine_torch as tex import transformer_engine as te +from transformer_engine import te_device_type from transformer_engine.pytorch.cpp_extensions.fused_attn import ( QKVLayout, QKVFormat, diff --git a/transformer_engine/pytorch/attention/inference.py b/transformer_engine/pytorch/attention/inference.py index ef842c04caa..6746ceccaba 100644 --- a/transformer_engine/pytorch/attention/inference.py +++ b/transformer_engine/pytorch/attention/inference.py @@ -10,6 +10,7 @@ import torch +from transformer_engine import te_device_type from transformer_engine.pytorch.attention.custom_ops import ( copy_to_kv_cache, QKV_FORMAT_VALUE, diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index bf273772533..51e446945af 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -11,6 +11,7 @@ import functools import torch import transformer_engine_torch as tex +from transformer_engine import te_device_type from ..constants import MXFP8_BLOCK_SCALING_SIZE, NVFP4_BLOCK_SCALING_SIZE, TE_DType, DType from ..utils import ceil_div, get_cached_ones_tensor, get_sm_count, _empty_tensor @@ -30,7 +31,6 @@ from ..tensor.utils import is_custom from ..custom_recipes.gemm import custom_gemm from ...debug.pytorch.debug_quantization import DebugQuantizedTensor, DebugQuantizer -from transformer_engine import te_device_type __all__ = [ "general_gemm", diff --git a/transformer_engine/pytorch/custom_recipes/quantizer_factories.py b/transformer_engine/pytorch/custom_recipes/quantizer_factories.py index fbbfe6ef00e..b6d22901dda 100644 --- a/transformer_engine/pytorch/custom_recipes/quantizer_factories.py +++ b/transformer_engine/pytorch/custom_recipes/quantizer_factories.py @@ -29,8 +29,8 @@ import torch from transformer_engine.pytorch.quantization import QuantizerRole -from ..constants import DType from transformer_engine import te_device_type +from ..constants import DType def high_precision_factory( diff --git a/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py b/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py index 95adb53998e..5cf573d61e4 100644 --- a/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py +++ b/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py @@ -72,9 +72,9 @@ from typing import Optional from transformer_engine.pytorch.quantization import QuantizerRole +from transformer_engine import te_device_type from ..constants import DType from .quantizer_factories import mxfp8_factory, nvfp4_factory -from transformer_engine import te_device_type # ----------------------------------------------------------------------------- # Pre-training-Oriented Recipes diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index cc67c5eacdc..ec43bf088d7 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -20,7 +20,6 @@ from torch.distributed.fsdp import FullyShardedDataParallel as FSDP from torch.distributed.fsdp._common_utils import _get_module_fsdp_state from torch.distributed.fsdp._traversal_utils import _get_fsdp_states_with_modules -from transformer_engine import te_device_type try: import torch.distributed._symmetric_memory as symm_mem @@ -32,6 +31,7 @@ import transformer_engine_torch as tex from transformer_engine.pytorch.triton.pad import pad_columnwise_scale_inv +from transformer_engine import te_device_type from .torch_version import torch_version from .utils import ( is_non_tn_fp8_gemm_supported, diff --git a/transformer_engine/pytorch/ep.py b/transformer_engine/pytorch/ep.py index ed42c2ecfeb..0ac7dc56f04 100644 --- a/transformer_engine/pytorch/ep.py +++ b/transformer_engine/pytorch/ep.py @@ -15,12 +15,12 @@ import transformer_engine_torch as tex +from transformer_engine import te_device_type from .constants import DType from .cpu_offload import mark_not_offload from .distributed import symm_mem_alloc, release_symm_mem_pool from .quantized_tensor import QuantizedTensor, QuantizedTensorStorage from .tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage -from transformer_engine import te_device_type # Type-hint-only import; keeps the ``Recipe`` annotation without a runtime import of # common.recipe (the concrete recipe classes are imported lazily where used). diff --git a/transformer_engine/pytorch/jit.py b/transformer_engine/pytorch/jit.py index 9c98c5d2189..8052f65bc2b 100644 --- a/transformer_engine/pytorch/jit.py +++ b/transformer_engine/pytorch/jit.py @@ -10,6 +10,7 @@ from typing import Callable, Optional, Tuple import torch +from .. import te_device_type from .torch_version import torch_version from .export import is_in_onnx_export_mode from .utils import gpu_autocast_ctx diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 8e1835b23e1..1123e181d10 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -20,8 +20,8 @@ import transformer_engine_torch as tex -from ._common import _ParameterInitMeta, noop_cat, sum_bias_grad from transformer_engine import te_device_type +from ._common import _ParameterInitMeta, noop_cat, sum_bias_grad from .._extra_state import ( extra_state_pickle_advisory, is_stateless_recipe, diff --git a/transformer_engine/pytorch/ops/_common.py b/transformer_engine/pytorch/ops/_common.py index 3a33ce161d4..58ec503f172 100644 --- a/transformer_engine/pytorch/ops/_common.py +++ b/transformer_engine/pytorch/ops/_common.py @@ -10,8 +10,8 @@ import torch -from transformer_engine import te_device_type from transformer_engine_torch import FP8TensorMeta +from transformer_engine import te_device_type from ..dynamo.tensor_spec import TensorSpec from ..ep import ( EpBuffer, diff --git a/transformer_engine/pytorch/ops/basic/activation.py b/transformer_engine/pytorch/ops/basic/activation.py index 968cb42fac1..5e0b547ad74 100644 --- a/transformer_engine/pytorch/ops/basic/activation.py +++ b/transformer_engine/pytorch/ops/basic/activation.py @@ -13,13 +13,13 @@ import torch import transformer_engine_torch as tex +from transformer_engine import te_device_type from ...constants import DType from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload from ...tensor.float8_tensor import Float8CurrentScalingQuantizer, Quantizer from ...utils import _compile_safe_warn, clear_tensor_data from ..op import BasicOperation, OperationContext from .._common import maybe_dequantize -from transformer_engine import te_device_type __all__ = [ "GELU", diff --git a/transformer_engine/pytorch/ops/basic/bias.py b/transformer_engine/pytorch/ops/basic/bias.py index d163426dd72..3d5be8fbf2c 100644 --- a/transformer_engine/pytorch/ops/basic/bias.py +++ b/transformer_engine/pytorch/ops/basic/bias.py @@ -11,12 +11,12 @@ import torch import transformer_engine_torch as tex +from transformer_engine import te_device_type from ...dynamo import TensorOrQuantized, TensorSpec from ...quantization import FP8GlobalStateManager from ..op import BasicOperation from ...utils import canonicalize_device, canonicalize_dtype from ...tensor import Quantizer -from transformer_engine import te_device_type @dataclass(slots=True) diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index 3f1bad91f5e..8f4d9bf1e33 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -15,11 +15,11 @@ import transformer_engine_torch as tex from transformer_engine.common.recipe import Recipe, Format as RecipeFormat +from transformer_engine import te_device_type from ...constants import DType, TE_DType from ...cpp_extensions import general_grouped_gemm, general_grouped_gemm_for_grouped_tensor from ...distributed import CudaRNGStatesTracker from ...module._common import WeightGradStore -from transformer_engine import te_device_type from ...module.base import ( _2X_ACC_FPROP, _2X_ACC_DGRAD, diff --git a/transformer_engine/pytorch/ops/basic/swiglu.py b/transformer_engine/pytorch/ops/basic/swiglu.py index ae9a31be3e7..9af5ad760d7 100644 --- a/transformer_engine/pytorch/ops/basic/swiglu.py +++ b/transformer_engine/pytorch/ops/basic/swiglu.py @@ -12,13 +12,13 @@ import torch import transformer_engine_torch as tex +from transformer_engine import te_device_type from ...constants import DType from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload from ...tensor import Float8CurrentScalingQuantizer, Quantizer from ...utils import clear_tensor_data from ..op import BasicOperation, OperationContext from .._common import maybe_dequantize -from transformer_engine import te_device_type __all__ = [ "SwiGLU", diff --git a/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py b/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py index 25594c63008..04a79ac6660 100644 --- a/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py +++ b/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py @@ -10,12 +10,12 @@ import torch +from transformer_engine import te_device_type from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload from ...quantization import FP8GlobalStateManager from ...tensor import Quantizer from ..basic import AddExtraInput, BasicLinear, Bias from ..op import FusedOperation, FusibleOperation, OperationContext -from transformer_engine import te_device_type class ForwardLinearBiasAdd(FusedOperation): diff --git a/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py b/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py index cae659f4fac..b7408c0379b 100644 --- a/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py +++ b/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py @@ -10,11 +10,11 @@ import torch +from transformer_engine import te_device_type from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload from ...quantization import FP8GlobalStateManager from ...tensor import Quantizer from ..basic import AddExtraInput, BasicLinear, ConstantScale -from transformer_engine import te_device_type from ..op import ( FusedOperation, FusibleOperation, diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index b4f62427a7a..71d08ebe459 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -17,6 +17,7 @@ from packaging.version import Version as PkgVersion import transformer_engine_torch as tex +from transformer_engine import te_device_type from ....common.recipe import Format as RecipeFormat from ...constants import DType, MXFP8_BLOCK_SCALING_SIZE, NVFP4_BLOCK_SCALING_SIZE, TE_DType from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload, start_offload diff --git a/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py b/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py index d37d9d88098..3bcf9ea8cbd 100644 --- a/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py +++ b/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py @@ -11,6 +11,7 @@ import torch from transformer_engine_torch import CommOverlapType, bulk_overlap_ag_with_external_gemm +from transformer_engine import te_device_type from ...cpp_extensions import general_gemm from ...distributed import gather_along_first_dim, get_distributed_world_size from ...module.base import ( diff --git a/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py b/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py index f8e23e3c730..606bf7418fc 100644 --- a/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py +++ b/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py @@ -11,11 +11,11 @@ import torch from transformer_engine_torch import CommOverlapType +from transformer_engine import te_device_type from ...cpp_extensions import general_gemm from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload from ...distributed import get_distributed_world_size from ...quantization import FP8GlobalStateManager -from transformer_engine import te_device_type from ...module.base import ( fill_userbuffers_buffer_for_all_gather, get_ub, diff --git a/transformer_engine/pytorch/optimizers/fused_adam.py b/transformer_engine/pytorch/optimizers/fused_adam.py index 0a5d8173fb7..2cde7780af0 100644 --- a/transformer_engine/pytorch/optimizers/fused_adam.py +++ b/transformer_engine/pytorch/optimizers/fused_adam.py @@ -15,9 +15,9 @@ import transformer_engine_torch as tex from transformer_engine.pytorch.tensor.float8_tensor import Float8Tensor, Float8Quantizer from transformer_engine.pytorch.quantized_tensor import QuantizedTensor +from transformer_engine import te_device_type from ..constants import DType from .multi_tensor_apply import multi_tensor_applier -from transformer_engine import te_device_type # Bound temporary NVTETensor handles created before the CUDA launcher chunks its metadata. diff --git a/transformer_engine/pytorch/tensor/grouped_tensor.py b/transformer_engine/pytorch/tensor/grouped_tensor.py index fa55e2ec3e7..474d932165d 100644 --- a/transformer_engine/pytorch/tensor/grouped_tensor.py +++ b/transformer_engine/pytorch/tensor/grouped_tensor.py @@ -10,10 +10,10 @@ import torch from torch.utils._pytree import tree_map +from transformer_engine import te_device_type from ..quantized_tensor import QuantizedTensorStorage, Quantizer from .storage.grouped_tensor_storage import GroupedTensorStorage from ..constants import DType -from transformer_engine import te_device_type def _stride_from_shape(shape: Tuple[int, ...]) -> Tuple[int, ...]: diff --git a/transformer_engine/pytorch/tensor/identity_tensor.py b/transformer_engine/pytorch/tensor/identity_tensor.py index b4fda2d00d2..8f8bda44cbc 100644 --- a/transformer_engine/pytorch/tensor/identity_tensor.py +++ b/transformer_engine/pytorch/tensor/identity_tensor.py @@ -19,9 +19,9 @@ import torch from torch.ops import aten +from transformer_engine import te_device_type from .storage.identity_tensor_storage import IdentityTensorStorage from ..quantized_tensor import QuantizedTensor, QuantizedTensorStorage, Quantizer -from transformer_engine import te_device_type class IdentityQuantizer(Quantizer): diff --git a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py index 9d40ac5b48b..cccd048cb83 100644 --- a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py @@ -12,10 +12,10 @@ import transformer_engine_torch as tex +from transformer_engine import te_device_type from ...quantized_tensor import InnerTensor, QuantizedTensorStorage, Quantizer from .._quantization_helpers import safe_quantized_repr -from transformer_engine import te_device_type from ...constants import ( TE_DType as torch_to_transformer_engine_dtype, MXFP8_BLOCK_SCALING_SIZE, diff --git a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py index 9209442dab5..a39a0732851 100644 --- a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py @@ -15,12 +15,12 @@ import transformer_engine_torch as tex +from transformer_engine import te_device_type from ...quantized_tensor import InnerTensor, QuantizedTensorStorage, Quantizer from .._quantization_helpers import safe_quantized_repr from ...constants import TE_DType as torch_to_transformer_engine_dtype, DType from ...utils import _empty_tensor -from transformer_engine import te_device_type @functools.lru_cache(maxsize=None) diff --git a/transformer_engine/pytorch/triton/mhc.py b/transformer_engine/pytorch/triton/mhc.py index 7288b4af923..ce9a507b448 100644 --- a/transformer_engine/pytorch/triton/mhc.py +++ b/transformer_engine/pytorch/triton/mhc.py @@ -25,6 +25,7 @@ _mhc_sinkhorn_bwd_fused, ) from transformer_engine.pytorch.cpp_extensions.gemm import general_gemm +from transformer_engine import te_device_type ENFORCE_DETERMINISTIC = os.environ.get("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1") == "0" diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py index 76ebb6efb55..734d7f12a2a 100644 --- a/transformer_engine/pytorch/utils.py +++ b/transformer_engine/pytorch/utils.py @@ -14,10 +14,10 @@ import numpy as np import torch +from transformer_engine import te_device_type from .torch_version import torch_version from ..common import decode_cudnn_version from ..debug.pytorch.debug_quantization import DebugQuantizedTensor -from transformer_engine import te_device_type __all__ = [ From 011d005b3063b91c5b3b0ec21ff14998f7dab1ca Mon Sep 17 00:00:00 2001 From: Xianduo Li Date: Mon, 19 Oct 2026 14:19:29 +0800 Subject: [PATCH 12/19] Query autocast state and dtype on the same device The device-type patch had changed the autocast dtype lookups to the configured device while leaving the enabled checks on CUDA. With a non-CUDA device type, vendor autocast enabled, and CUDA autocast disabled, these sites took the non-autocast branch and kept the input dtype instead of the requested autocast dtype. Route the operation sites through `maybe_autocast_dtype`, which already queries one device for both calls and carries the torch version guard. This covers the activation, SwiGLU, grouped linear, basic linear, and fused linear forward operations. Add `torch_is_autocast_gpu_enabled` next to `torch_get_autocast_gpu_dtype` for the two sites outside the operation layer that had the same mismatch: `set_activation_dtype`, which sets the activation dtype for every module, and the TransformerLayer forward cast. Signed-off-by: Xianduo Li --- transformer_engine/pytorch/module/base.py | 3 ++- .../pytorch/ops/basic/activation.py | 17 ++++---------- .../pytorch/ops/basic/basic_linear.py | 7 ++---- .../pytorch/ops/basic/grouped_linear.py | 12 +++------- .../pytorch/ops/basic/swiglu.py | 23 ++++--------------- .../ops/fused/forward_linear_bias_add.py | 7 ++---- .../ops/fused/forward_linear_scale_add.py | 7 ++---- .../pytorch/ops/fused/grouped_mlp.py | 7 ++---- .../ops/fused/userbuffers_forward_linear.py | 7 ++---- transformer_engine/pytorch/transformer.py | 3 ++- transformer_engine/pytorch/utils.py | 7 ++++++ 11 files changed, 33 insertions(+), 67 deletions(-) diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 5a933dcc1fa..aa5883c2dd8 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -63,6 +63,7 @@ from ..utils import ( is_non_tn_fp8_gemm_supported, torch_get_autocast_gpu_dtype, + torch_is_autocast_gpu_enabled, get_device_compute_capability, get_nvtx_range_context, nvtx_range_push, @@ -1541,7 +1542,7 @@ def copy_tensor(src: torch.Tensor, dst: torch.Tensor) -> None: def set_activation_dtype(self, inp: torch.Tensor) -> None: """Get activation data type for AMP.""" # Native AMP (`torch.autocast`) gets highest priority - if torch.is_autocast_enabled(): + if torch_is_autocast_gpu_enabled(): self.fast_setattr("activation_dtype", torch_get_autocast_gpu_dtype()) return dtype = inp.dtype diff --git a/transformer_engine/pytorch/ops/basic/activation.py b/transformer_engine/pytorch/ops/basic/activation.py index 5e0b547ad74..21a376ca3d1 100644 --- a/transformer_engine/pytorch/ops/basic/activation.py +++ b/transformer_engine/pytorch/ops/basic/activation.py @@ -13,13 +13,12 @@ import torch import transformer_engine_torch as tex -from transformer_engine import te_device_type from ...constants import DType from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload from ...tensor.float8_tensor import Float8CurrentScalingQuantizer, Quantizer from ...utils import _compile_safe_warn, clear_tensor_data from ..op import BasicOperation, OperationContext -from .._common import maybe_dequantize +from .._common import maybe_autocast_dtype, maybe_dequantize __all__ = [ "GELU", @@ -95,11 +94,7 @@ def op_forward( ) -> torch.Tensor: # Compute dtype - dtype: torch.dtype - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_dtype(te_device_type()) - else: - dtype = input_.dtype + dtype: torch.dtype = maybe_autocast_dtype(default_dtype=input_.dtype) if dtype not in (torch.float32, torch.float16, torch.bfloat16): raise RuntimeError(f"Unsupported dtype ({dtype})") @@ -413,12 +408,8 @@ def fuser_forward( "in the fused grouped MLP path." ) - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_dtype(te_device_type()) - extra_input_dtype = dtype - else: - dtype = input_.dtype - extra_input_dtype = extra_input.dtype + dtype = maybe_autocast_dtype(default_dtype=input_.dtype) + extra_input_dtype = maybe_autocast_dtype(default_dtype=extra_input.dtype) x = maybe_dequantize(input_.contiguous(), dtype) scales = maybe_dequantize(extra_input.contiguous(), extra_input_dtype) diff --git a/transformer_engine/pytorch/ops/basic/basic_linear.py b/transformer_engine/pytorch/ops/basic/basic_linear.py index df8e13604a9..e60662226d3 100644 --- a/transformer_engine/pytorch/ops/basic/basic_linear.py +++ b/transformer_engine/pytorch/ops/basic/basic_linear.py @@ -43,6 +43,7 @@ get_dummy_wgrads_for_params, get_main_grad_from_param, is_quantized_tensor, + maybe_autocast_dtype, maybe_dequantize, ) @@ -1108,11 +1109,7 @@ def pack_forward_args( return BasicLinearFwdArgs( input_=input_, weight=self.weight, - dtype=( - torch.get_autocast_dtype("cuda") - if torch.is_autocast_enabled() - else self.weight.dtype - ), + dtype=maybe_autocast_dtype(default_dtype=self.weight.dtype), input_requires_grad=basic_op_ctxs[0].requires_grad, weight_requires_grad=basic_op_ctxs[0].requires_grad and self.weight.requires_grad, with_quantized_compute=with_quantized_compute, diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index c16fe95cbab..faceb58a666 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -17,7 +17,6 @@ import transformer_engine_torch as tex from transformer_engine.common.recipe import Recipe, Format as RecipeFormat -from transformer_engine import te_device_type from ...constants import DType, TE_DType from ...cpp_extensions import general_grouped_gemm, general_grouped_gemm_for_grouped_tensor from ...distributed import CudaRNGStatesTracker @@ -50,6 +49,7 @@ get_dummy_wgrads_for_params, get_main_grad_from_param, is_quantized_tensor, + maybe_autocast_dtype, maybe_dequantize, validate_or_alloc_output, view_main_grad_as_grouped_buffer, @@ -1049,10 +1049,7 @@ def fuser_forward( weight_quantizers[group_idx] = self.get_quantizer("forward", 2 * group_idx + 1) # Get autocast dtype if needed - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_dtype(te_device_type()) - else: - dtype = weight_param.dtype + dtype = maybe_autocast_dtype(default_dtype=weight_param.dtype) # Extract split sizes from extra input. Keep on GPU for graph safety. split_sizes = basic_op_extra_inputs[0][0] @@ -1203,10 +1200,7 @@ def fuser_forward_save_ctx( # ``split_sizes``, offset metadata, and related tensors are routed # through ``save_for_backward`` (see ``_fuser_forward_split_quantize`` # and ``_fuser_forward_grouped_tensor`` for the saved-tensor layout). - if torch.is_autocast_enabled(): - ctx.dtype = torch.get_autocast_dtype(te_device_type()) - else: - ctx.dtype = weight_param.dtype + ctx.dtype = maybe_autocast_dtype(default_dtype=weight_param.dtype) ctx.input_requires_grad = requires_grad[0] ctx.weight_requires_grad = requires_grad[0] and weight_param.requires_grad # Caller-provided backward grad-input buffer. diff --git a/transformer_engine/pytorch/ops/basic/swiglu.py b/transformer_engine/pytorch/ops/basic/swiglu.py index 9af5ad760d7..71ff010ae53 100644 --- a/transformer_engine/pytorch/ops/basic/swiglu.py +++ b/transformer_engine/pytorch/ops/basic/swiglu.py @@ -12,13 +12,12 @@ import torch import transformer_engine_torch as tex -from transformer_engine import te_device_type from ...constants import DType from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload from ...tensor import Float8CurrentScalingQuantizer, Quantizer from ...utils import clear_tensor_data from ..op import BasicOperation, OperationContext -from .._common import maybe_dequantize +from .._common import maybe_autocast_dtype, maybe_dequantize __all__ = [ "SwiGLU", @@ -115,11 +114,7 @@ def op_forward( ) -> torch.Tensor: # Compute dtype - dtype: torch.dtype - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_dtype(te_device_type()) - else: - dtype = input_.dtype + dtype: torch.dtype = maybe_autocast_dtype(default_dtype=input_.dtype) if dtype not in (torch.float32, torch.float16, torch.bfloat16): raise RuntimeError(f"Unsupported dtype ({dtype})") @@ -362,11 +357,7 @@ def op_forward( ) -> torch.Tensor: # Compute dtype - dtype: torch.dtype - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_dtype(te_device_type()) - else: - dtype = input_.dtype + dtype: torch.dtype = maybe_autocast_dtype(default_dtype=input_.dtype) if dtype not in (torch.float32, torch.float16, torch.bfloat16): raise RuntimeError(f"Unsupported dtype ({dtype})") @@ -529,12 +520,8 @@ def fuser_forward( extra_input = basic_op_extra_inputs[0][0] # Determine compute dtype - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_dtype(te_device_type()) - extra_input_dtype = dtype - else: - dtype = input_.dtype - extra_input_dtype = extra_input.dtype + dtype = maybe_autocast_dtype(default_dtype=input_.dtype) + extra_input_dtype = maybe_autocast_dtype(default_dtype=extra_input.dtype) # Prepare the activation input in the compute dtype. input_ = maybe_dequantize(input_, dtype) diff --git a/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py b/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py index 04a79ac6660..7f7a6bd97e3 100644 --- a/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py +++ b/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py @@ -10,10 +10,10 @@ import torch -from transformer_engine import te_device_type from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload from ...quantization import FP8GlobalStateManager from ...tensor import Quantizer +from .._common import maybe_autocast_dtype from ..basic import AddExtraInput, BasicLinear, Bias from ..op import FusedOperation, FusibleOperation, OperationContext @@ -93,10 +93,7 @@ def fuser_forward( backward_override = None # Get autocast dtype if needed - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_dtype(te_device_type()) - else: - dtype = linear_op.weight.dtype + dtype = maybe_autocast_dtype(default_dtype=linear_op.weight.dtype) # Linear forward output = basic_op_extra_inputs[self._op_idxs["add"]][0] diff --git a/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py b/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py index b7408c0379b..da78bec23b7 100644 --- a/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py +++ b/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py @@ -10,10 +10,10 @@ import torch -from transformer_engine import te_device_type from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload from ...quantization import FP8GlobalStateManager from ...tensor import Quantizer +from .._common import maybe_autocast_dtype from ..basic import AddExtraInput, BasicLinear, ConstantScale from ..op import ( FusedOperation, @@ -75,10 +75,7 @@ def fuser_forward( extra_input = basic_op_extra_inputs[2][0] # Get autocast dtype if needed - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_dtype(te_device_type()) - else: - dtype = linear_op.weight.dtype + dtype = maybe_autocast_dtype(default_dtype=linear_op.weight.dtype) # Linear forward output, x_local, w = BasicLinear._functional_forward( diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index a5818e5e617..0f555536f96 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -18,7 +18,6 @@ from packaging.version import Version as PkgVersion import transformer_engine_torch as tex -from transformer_engine import te_device_type from ....common.recipe import Format as RecipeFormat from ...constants import DType, MXFP8_BLOCK_SCALING_SIZE, NVFP4_BLOCK_SCALING_SIZE, TE_DType from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload, start_offload @@ -60,6 +59,7 @@ get_dummy_wgrads_for_params, get_main_grad_from_param, is_quantized_tensor, + maybe_autocast_dtype, maybe_dequantize, validate_or_alloc_output, view_main_grad_as_grouped_buffer, @@ -1344,10 +1344,7 @@ def fuser_forward( assert fc2_op.weight0.is_routed_expert and fc2_op.weight0.weight_list is not None device = fc1_weight_param.device - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_dtype(te_device_type()) - else: - dtype = fc1_weight_param.dtype + dtype = maybe_autocast_dtype(default_dtype=fc1_weight_param.dtype) # Check which grads are required requires_grad = any(ctx.requires_grad for ctx in basic_op_ctxs) diff --git a/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py b/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py index 606bf7418fc..9bc2cfc1e75 100644 --- a/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py +++ b/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py @@ -24,7 +24,7 @@ from ...quantized_tensor import Quantizer from ...tensor.float8_tensor import Float8Quantizer, Float8CurrentScalingQuantizer from ...tensor.storage.float8_tensor_storage import Float8TensorStorage -from .._common import maybe_dequantize, is_quantized_tensor +from .._common import maybe_autocast_dtype, maybe_dequantize, is_quantized_tensor from ..basic import BasicLinear, Bias, ReduceScatter from ..op import ( FusedOperation, @@ -322,10 +322,7 @@ def fuser_forward( ) # Get autocast dtype if needed - if torch.is_autocast_enabled(): - dtype = torch.get_autocast_dtype(te_device_type()) - else: - dtype = linear_op.weight.dtype + dtype = maybe_autocast_dtype(default_dtype=linear_op.weight.dtype) # Userbuffers options if linear_op._userbuffers_options is None: diff --git a/transformer_engine/pytorch/transformer.py b/transformer_engine/pytorch/transformer.py index 8f8e4f90381..cc361454a4c 100644 --- a/transformer_engine/pytorch/transformer.py +++ b/transformer_engine/pytorch/transformer.py @@ -27,6 +27,7 @@ cast_if_needed, get_default_init_method, torch_get_autocast_gpu_dtype, + torch_is_autocast_gpu_enabled, ) from transformer_engine.pytorch.constants import ( AttnMaskTypes, @@ -909,7 +910,7 @@ def forward( ) # For AMP - if torch.is_autocast_enabled(): + if torch_is_autocast_gpu_enabled(): hidden_states = cast_if_needed(hidden_states, torch_get_autocast_gpu_dtype()) # Self attention. diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py index 309c3556dc5..a3208db5026 100644 --- a/transformer_engine/pytorch/utils.py +++ b/transformer_engine/pytorch/utils.py @@ -953,6 +953,13 @@ def canonicalize_process_group( return group +def torch_is_autocast_gpu_enabled() -> bool: + """Check whether PyTorch autocast is enabled for the GPU device.""" + if torch_version() >= (2, 4, 0): + return torch.is_autocast_enabled(te_device_type()) + return torch.is_autocast_enabled() + + def torch_get_autocast_gpu_dtype() -> torch.dtype: """Get PyTorch autocast GPU dtype.""" if torch_version() >= (2, 4, 0): From a2baa208c218b87868a9fc6195665acd6b11437b Mon Sep 17 00:00:00 2001 From: Xianduo Li Date: Mon, 19 Oct 2026 14:20:07 +0800 Subject: [PATCH 13/19] Gate per-document THD partitioning on CUDA `tex.thd_get_partitioned_indices` allocates a CUDA output tensor and launches a CUDA kernel on the current CUDA stream, so it runs only on CUDA. The device-type patch had changed its dispatch condition to the configured device type, which sent packed sequences into that kernel on a non-CUDA device instead of the device-capable fallback, failing the `is_cuda` assertion before the fallback could run. Restore the CUDA condition in `build_rank_indices` and in `_get_thd_partition_cu_seqlens`, which had the same problem when normalizing the boundaries that feed the kernel, and record why both are gated on CUDA rather than on the configured device. Keep the configured device in the NO_LOAD_BALANCE branch, which builds indices with plain tensor operations and runs no CUDA kernel. Restore the assertion message, which named the device type that triggers it. Signed-off-by: Xianduo Li --- .../dot_product_attention/context_parallel.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index a45183a358b..68eb7b1a129 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -276,9 +276,9 @@ def get_seq_chunk_ids_for_reordering_after_attn(cp_size, device): def _get_thd_partition_cu_seqlens(cu_seqlens_padded, device=None): """Return physical boundaries used by the THD partition CUDA kernels.""" target_device = torch.device(device if device is not None else cu_seqlens_padded.device) - target_dtype = ( - torch.int32 if target_device.type == te_device_type() else cu_seqlens_padded.dtype - ) + # These boundaries feed CUDA kernels that require int32, so the cast is gated on CUDA + # rather than on the configured device type. + target_dtype = torch.int32 if target_device.type == "cuda" else cu_seqlens_padded.dtype return cu_seqlens_padded.to(device=target_device, dtype=target_dtype) @@ -302,6 +302,8 @@ def get_thd_partitioned_indices( cp_size, ) target_device = torch.device(device if device is not None else cu_seqlens_padded.device) + # This branch builds indices with plain tensor ops rather than a CUDA kernel, so the + # configured device gets the same int32 normalization that CUDA does. target_dtype = ( torch.int32 if target_device.type == te_device_type() else cu_seqlens_padded.dtype ) @@ -314,7 +316,7 @@ def get_thd_partitioned_indices( ) cu_seqlens_padded = _get_thd_partition_cu_seqlens(cu_seqlens_padded, device) assert cu_seqlens_padded.is_cuda, ( - "Per-document THD partitioning requires CUDA cu_seqlens; pass device=te_device_type() " + "Per-document THD partitioning requires CUDA cu_seqlens; pass device='cuda' " "when the source metadata is on CPU." ) if cu_seqlens_padded.dtype != torch.int32: @@ -5960,7 +5962,11 @@ def build_rank_indices(device): slice_sizes = (cu_seqlens_padded[1:] - cu_seqlens_padded[:-1]) // total_slices def build_rank_indices(device): - if device.type == te_device_type(): + # Per-document partitioning runs a CUDA kernel + # (``tex.thd_get_partitioned_indices``), so it is gated on CUDA rather + # than on the configured device type. Other devices, including CPU and + # non-CUDA plugin devices, use the fallback below. + if device.type == "cuda": return get_thd_partitioned_indices( cu_seqlens_padded, seq_len_val, From 171148f0354a44e59c9366d6dc9035febacf7aff Mon Sep 17 00:00:00 2001 From: Xianduo Li Date: Mon, 19 Oct 2026 14:36:23 +0800 Subject: [PATCH 14/19] Drop incidental changes from the device-type patch Narrow the diff against main to the device-type substitutions themselves: - context_parallel.py: remove the comments on the THD partitioning device gates. - cpp_extensions/gemm.py: restore the single-line Float8BlockwiseQTensorStorage import, which did not need wrapping. - graph.py: remove an added blank line, leaving the file unchanged from main. Signed-off-by: Xianduo Li --- .../attention/dot_product_attention/context_parallel.py | 8 -------- transformer_engine/pytorch/cpp_extensions/gemm.py | 4 +--- transformer_engine/pytorch/graph.py | 1 - 3 files changed, 1 insertion(+), 12 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index 68eb7b1a129..376f8dc90be 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -276,8 +276,6 @@ def get_seq_chunk_ids_for_reordering_after_attn(cp_size, device): def _get_thd_partition_cu_seqlens(cu_seqlens_padded, device=None): """Return physical boundaries used by the THD partition CUDA kernels.""" target_device = torch.device(device if device is not None else cu_seqlens_padded.device) - # These boundaries feed CUDA kernels that require int32, so the cast is gated on CUDA - # rather than on the configured device type. target_dtype = torch.int32 if target_device.type == "cuda" else cu_seqlens_padded.dtype return cu_seqlens_padded.to(device=target_device, dtype=target_dtype) @@ -302,8 +300,6 @@ def get_thd_partitioned_indices( cp_size, ) target_device = torch.device(device if device is not None else cu_seqlens_padded.device) - # This branch builds indices with plain tensor ops rather than a CUDA kernel, so the - # configured device gets the same int32 normalization that CUDA does. target_dtype = ( torch.int32 if target_device.type == te_device_type() else cu_seqlens_padded.dtype ) @@ -5962,10 +5958,6 @@ def build_rank_indices(device): slice_sizes = (cu_seqlens_padded[1:] - cu_seqlens_padded[:-1]) // total_slices def build_rank_indices(device): - # Per-document partitioning runs a CUDA kernel - # (``tex.thd_get_partitioned_indices``), so it is gated on CUDA rather - # than on the configured device type. Other devices, including CPU and - # non-CUDA plugin devices, use the fallback below. if device.type == "cuda": return get_thd_partitioned_indices( cu_seqlens_padded, diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 51e446945af..00b452fe585 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -20,9 +20,7 @@ from ..tensor.float8_tensor import Float8CurrentScalingQuantizer, Float8Quantizer from ..tensor.mxfp8_tensor import MXFP8Quantizer from ..tensor.nvfp4_tensor import NVFP4Quantizer -from ..tensor.storage.float8_blockwise_tensor_storage import ( - Float8BlockwiseQTensorStorage, -) +from ..tensor.storage.float8_blockwise_tensor_storage import Float8BlockwiseQTensorStorage from ..tensor.storage.float8_tensor_storage import Float8TensorStorage from ..tensor.storage.grouped_tensor_storage import GroupedTensorStorage from ..tensor.storage.hybrid_tensor_storage import HybridQuantizedTensorStorage diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index 705c5bf018c..bc89a815af1 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -17,7 +17,6 @@ from transformer_engine.common.recipe import DelayedScaling, Recipe from transformer_engine.pytorch.constants import dist_group_type - from .quantization import ( autocast, FP8GlobalStateManager, From 748ed354209fadb05259d4a45c610fa54ab8f2ab Mon Sep 17 00:00:00 2001 From: Xianduo Li Date: Mon, 19 Oct 2026 14:59:26 +0800 Subject: [PATCH 15/19] Cover device-type checks missed in recently merged code Two paths reached from main used CUDA device checks where the work below them is not CUDA-specific: - tensor/nvfp4_tensor.py: `NVFP4Tensor` kept CUDA checks in its `record_stream` dispatch and in `copy_`, while the sibling `Float8Tensor`, `MXFP8Tensor`, and `Float8BlockwiseQTensor` classes use the configured device type for the same logic. On another device type `copy_` would move the operand to the destination device and `record_stream` would skip stream tracking. - cross_entropy.py: the input validation rejected non-CUDA tensors even though the Triton backend derives every allocation from the input device and contains nothing CUDA-specific. Leave the checks that guard compiled CUDA kernels, the `device_types` dispatch keys on custom operations, and the `is_cuda` properties that implement the PyTorch tensor contract. Signed-off-by: Xianduo Li --- transformer_engine/pytorch/cross_entropy.py | 7 ++++--- transformer_engine/pytorch/tensor/nvfp4_tensor.py | 5 +++-- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/transformer_engine/pytorch/cross_entropy.py b/transformer_engine/pytorch/cross_entropy.py index 8ae3b853801..2318969963b 100644 --- a/transformer_engine/pytorch/cross_entropy.py +++ b/transformer_engine/pytorch/cross_entropy.py @@ -10,6 +10,7 @@ import torch import transformer_engine.pytorch.triton.cross_entropy as triton_cross_entropy +from transformer_engine import te_device_type from transformer_engine.pytorch.jit import no_torch_dynamo __all__ = ["parallel_cross_entropy"] @@ -114,10 +115,10 @@ def _validate_inputs( raise ValueError(f"inp dimensions must be non-zero, but got shape {tuple(inp.shape)}") if inp.dtype not in (torch.bfloat16, torch.float32): raise TypeError(f"inp must have BF16 or FP32 dtype, but got {inp.dtype}") - if not inp.is_cuda: - raise ValueError("inp must be a CUDA tensor") + if inp.device.type != te_device_type(): + raise ValueError(f"inp must be a {te_device_type()} tensor, but got {inp.device}") if target.device != inp.device: - raise ValueError("target must be on the same CUDA device as inp") + raise ValueError(f"target must be on the same {te_device_type()} device as inp") if target.dtype != torch.int64: raise TypeError(f"target must have torch.int64 dtype, but got {target.dtype}") if target.numel() != inp.shape[0] * inp.shape[1]: diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index 3b4c7fd2d89..55a99e0f5c6 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -13,6 +13,7 @@ import torch import transformer_engine_torch as tex +from transformer_engine import te_device_type from transformer_engine.common.recipe import NVFP4BlockScaling, Recipe from ..constants import NVFP4_BLOCK_SCALING_SIZE, dist_group_type, DType from ..utils import ( @@ -813,7 +814,7 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): qt._amax_rowwise, qt._amax_columnwise, ): - if t is not None and t.is_cuda: + if t is not None and t.device.type == te_device_type(): t.record_stream(stream) return None @@ -971,7 +972,7 @@ def _set_data(self, tensor: torch.Tensor) -> None: """ # Tensor device - new_device = tensor.device if tensor.is_cuda else self.device + new_device = tensor.device if tensor.device.type == te_device_type() else self.device if not devices_match(new_device, tensor.device): tensor = tensor.to(device=new_device) From 0c7cf397684463bb980dfc8e6a43e0e80e3ae673 Mon Sep 17 00:00:00 2001 From: Xianduo Li Date: Wed, 21 Oct 2026 15:28:10 +0800 Subject: [PATCH 16/19] Apply plugin device patches after the _nv alias is registered The top-level hook in transformer_engine/__init__.py imported .patches before load_framework_extension() registered transformer_engine_torch_nv. A plugin that imports that alias at package level - a case the loader explicitly supports - failed there, and the enclosing except only warned, so apply_patches() never ran. The loader later imported the plugin successfully but never retried the patches, leaving device patches unapplied while vendor calls still went through CUDA. Move the call into load_framework_extension(), after the alias is registered and before the PyTorch submodules capture references to torch.cuda entry points. Probe the submodule with find_spec so a plugin that needs no patches may omit it, without masking real import failures inside an existing patches module. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Xianduo Li --- transformer_engine/__init__.py | 19 ------------------- transformer_engine/common/__init__.py | 5 +++++ 2 files changed, 5 insertions(+), 19 deletions(-) diff --git a/transformer_engine/__init__.py b/transformer_engine/__init__.py index ee4a69e9467..baaf416a23b 100644 --- a/transformer_engine/__init__.py +++ b/transformer_engine/__init__.py @@ -16,25 +16,6 @@ import transformer_engine.common -# Plugin system: set NVTE_PLUGIN to the plugin module name to enable. -# e.g. NVTE_PLUGIN=transformer_engine_plugin_fl -_nvte_plugin = os.environ.get("NVTE_PLUGIN") -if _nvte_plugin: - try: - from importlib import import_module - - _patches = import_module(f"{_nvte_plugin}.patches") - _patches.apply_patches() - except Exception as e: - import warnings - - warnings.warn( - f"NVTE_PLUGIN={_nvte_plugin} but plugin patch apply failed: {e}", - RuntimeWarning, - stacklevel=1, - ) - - # Minimum NCCL version for the runtime-loaded NCCL EP backend. _NCCL_EP_MIN_VERSION = (2, 30, 4) diff --git a/transformer_engine/common/__init__.py b/transformer_engine/common/__init__.py index e4a5f9256b8..103ca6e8f59 100644 --- a/transformer_engine/common/__init__.py +++ b/transformer_engine/common/__init__.py @@ -231,6 +231,11 @@ def load_framework_extension(framework: str) -> None: # plugin module may import transformer_engine_torch_nv at top level. sys.modules[module_name + "_nv"] = solib _plugin = importlib.import_module(_nvte_plugin) + # Apply device patches before Transformer Engine's PyTorch submodules + # capture references to torch.cuda entry points. The submodule is + # optional: a plugin that needs no patches may omit it. + if importlib.util.find_spec(f"{_nvte_plugin}.patches") is not None: + importlib.import_module(f"{_nvte_plugin}.patches").apply_patches() _plugin.load_plugins() except Exception as e: # pylint: disable=broad-exception-caught # Rollback to pre-plugin state if plugin failed to fully initialize From 87ff4f299e9dad58d15f6e77bd3698bbede4d1ef Mon Sep 17 00:00:00 2001 From: Xianduo Li Date: Wed, 21 Oct 2026 15:45:19 +0800 Subject: [PATCH 17/19] Accept the configured device type in the shared forward input check The modules whose `device` argument now defaults to `te_device_type()` allocate their parameters on the configured device, but the `prepare_forward()` input check still rejected anything whose `is_cuda` was false. That property is a tensor attribute, so patching the `torch.cuda` namespace does not change it, and forward raised before a plugin's native operation could run. Compare `device.type` against `te_device_type()` instead, and report the configured device type in the error message. One check covers Linear, LayerNormLinear, GroupedLinear, and LayerNormMLP, which share this base method. Behavior on CUDA is unchanged. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Xianduo Li --- transformer_engine/pytorch/module/base.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 6b85b1fd4d0..2e998761154 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -1720,9 +1720,9 @@ def prepare_forward( delayed_scaling_recipe = _has_delayed_scaling_state(self.fp8_meta) FP8GlobalStateManager.get_old_fp8_meta_tensors_for_recompute(self.fp8_meta) else: - if not inp.is_cuda: + if inp.device.type != te_device_type(): raise RuntimeError( - f"TransformerEngine needs CUDA. Got input on device: {inp.device}" + f"TransformerEngine needs {te_device_type()}. Got input on device: {inp.device}" ) if self.tp_size > 1: From 15d350242c19e88e9a12b63edbe43b2b7832de24 Mon Sep 17 00:00:00 2001 From: Xianduo Li Date: Wed, 21 Oct 2026 15:55:11 +0800 Subject: [PATCH 18/19] Restore the original spacing in the top-level package Removing the plugin hook left an extra blank line behind. Drop it so the file matches main and no longer appears in the diff. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Xianduo Li --- transformer_engine/__init__.py | 1 - 1 file changed, 1 deletion(-) diff --git a/transformer_engine/__init__.py b/transformer_engine/__init__.py index baaf416a23b..1c35b1262b8 100644 --- a/transformer_engine/__init__.py +++ b/transformer_engine/__init__.py @@ -15,7 +15,6 @@ from typing import Optional, Tuple import transformer_engine.common - # Minimum NCCL version for the runtime-loaded NCCL EP backend. _NCCL_EP_MIN_VERSION = (2, 30, 4) From d92f20ccd09915e1c4c626ffbaaae2379ebf32b6 Mon Sep 17 00:00:00 2001 From: Xianduo Li Date: Wed, 21 Oct 2026 16:02:38 +0800 Subject: [PATCH 19/19] Skip patch discovery for single-file plugins find_spec() raises ModuleNotFoundError instead of returning None when the parent module is not a package, so a single-file plugin that provides load_plugins() failed the optional-patches probe. The exception handler then disabled the plugin and restored the original extension. Check for __path__ before searching for the submodule. Genuine import failures inside an existing patches module still propagate. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Xianduo Li --- transformer_engine/common/__init__.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/transformer_engine/common/__init__.py b/transformer_engine/common/__init__.py index 103ca6e8f59..ee50d14b7a6 100644 --- a/transformer_engine/common/__init__.py +++ b/transformer_engine/common/__init__.py @@ -233,8 +233,11 @@ def load_framework_extension(framework: str) -> None: _plugin = importlib.import_module(_nvte_plugin) # Apply device patches before Transformer Engine's PyTorch submodules # capture references to torch.cuda entry points. The submodule is - # optional: a plugin that needs no patches may omit it. - if importlib.util.find_spec(f"{_nvte_plugin}.patches") is not None: + # optional: a plugin that needs no patches may omit it, and a + # single-file plugin has no submodules to search. + if hasattr(_plugin, "__path__") and ( + importlib.util.find_spec(f"{_nvte_plugin}.patches") is not None + ): importlib.import_module(f"{_nvte_plugin}.patches").apply_patches() _plugin.load_plugins() except Exception as e: # pylint: disable=broad-exception-caught