Repository navigation
Conversation
) * [PyTorch] Reuse cuBLASMp's all-gathered grad_output for wgrad With the cuBLASMp comm+GEMM overlap backend, the dgrad AllGather + GEMM gathers grad_output and the wgrad GEMM gathered it a second time, while Userbuffers wgrad reads the tensor that its dgrad all-gather left in the communication buffer. cuBLASMp 0.11.0 leaves the gathered B at the start of its workspace (a documented cublasMpMatmul workspace property), so wgrad now reuses it the same way: - nvte_comm_gemm_workspace() (comm_gemm.h, TE-internal) returns the comm-gemm context's workspace. - With the cuBLASMp backend, CommOverlap/CommOverlapP2P::get_buffer() returns a view of that workspace with the communicator's buffer shape and dtype, like the Userbuffers buffer view. - get_cublasmp_all_gather_output() builds the gathered tensor from it for unquantized and per-tensor FP8 data. Linear, LayerNormLinear, LayerNormMLP (fc2) and the fused te.ops backward use it instead of gather_along_first_dim; MXFP8 and other formats, a delayed wgrad and a skipped dgrad GEMM still gather. This needs cuBLASMp 0.11.0 or newer, since older versions keep the gathered B elsewhere in the workspace; the next commit checks the version. Tested on 4x H100 with cuBLASMp 0.11.0: the layer and multi-layer tests of test_comm_gemm_overlap.py with the cuBLASMp backend pass (39 passed, 13 MXFP8 cases skipped on H100), and test_fusible_ops_with_userbuffers.py --use-cublasmp passes in BF16 and FP8 delayed and current scaling. All-gathers per backward: row-parallel Linear 1 -> 0, LayerNormMLP 2 -> 1 (fc1 still re-gathers its input, as with Userbuffers). Signed-off-by: Almog Segal <asegal@nvidia.com> * [PyTorch] Fall back to the all-gather with cuBLASMp before 0.11.0 Reusing the gathered grad_output relies on cuBLASMp 0.11.0 leaving the gathered B at the start of its workspace. Earlier versions keep it after their cuBLASLt workspaces, so the view held other data. nvte_comm_gemm_workspace() now checks the loaded library with cublasMpGetVersion() and returns nullptr before 0.11.0; get_buffer() then reports that no gathered input is available, and get_cublasmp_all_gather_output() returns None, so the callers all-gather as before. Tested on 4x H100: with cuBLASMp 0.11.0, the row-parallel Linear backward issues no all-gather and the Linear and LayerNormMLP overlap tests with the cuBLASMp backend pass (9 passed). With cuBLASMp 0.10.0 it all-gathers once per backward and the same tests pass, as they do on TE main. Signed-off-by: Almog Segal <asegal@nvidia.com> * [PyTorch] Check the cuBLASMp version instead of catching the get_buffer() error get_cublasmp_all_gather_output() fell back to the all-gather for cuBLASMp versions before 0.11.0 by catching the error that get_buffer() raises and matching its message. That breaks silently if the message changes, and with an older cuBLASMp it raises and catches an exception in every backward pass of every layer that overlaps an all-gather. nvte_cublasmp_version() now returns the version of the loaded cuBLASMp library, read once with cublasMpGetVersion() (0 without cuBLASMp), and Python exposes it as tex.nvte_cublasmp_version(). get_cublasmp_all_gather_output() checks a cached version >= 0.11.0 before calling get_buffer(), so get_buffer() errors are no longer caught. nvte_comm_gemm_workspace() uses the same function. Signed-off-by: Almog Segal <asegal@nvidia.com> * [PyTorch] Test the fused te.ops comm+GEMM overlap with cuBLASMp test_fuser_ops_with_userbuffers only ran the Userbuffers backend, although the script supports --use-cublasmp and the fused te.ops backward reuses cuBLASMp's gathered grad_output. It is now parametrized over both backends and skips the cuBLASMp cases when Transformer Engine is built without cuBLASMp. Signed-off-by: Almog Segal <asegal@nvidia.com> --------- Signed-off-by: Almog Segal <asegal@nvidia.com>
* [PyTorch] Resolve linear attention operators during initialization Cache the cuDNN operator on the shared adapter during construction so Dynamo does not trace importlib from forward. Keep backend loading deferred until a linear attention module is instantiated. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * [PyTorch] Test full-graph compilation of linear attention Exercise the existing GDN, GDN2, and DeltaProduct numerical tests with fullgraph compilation, including dense FP16/BF16 layouts and THD forward, recurrent state, and backward checks against FP64 references. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * [PyTorch] Simplify linear attention operator dispatch Resolve the cuDNN operator directly in the adapter constructor and remove the cached import helper. Share the operator call across variants, keeping only Householder layout preparation in DeltaProduct. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * [PyTorch] Keep linear attention state arguments explicit Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * [PyTorch] Preserve variant-specific linear attention dispatch Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * [PyTorch] Mock linear attention backends in module tests Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Fix cudnnFE init time requirement Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> * [PyTorch] Use the cached operator in KDA dispatch Align the KDA adapter with the shared linear-attention call signature after merging main. Drop the obsolete op argument and invoke self._op so valid KDA calls reach the backend. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> --------- Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* [Common] Support NVFP4 in the cuBLASMp comm+GEMM overlap NVFP4 GEMMs apply the tensors' global scales through alpha, which TE computes on the GPU from the amaxes (nvte_nvfp4_compute_per_tensor_scale). Without a host sync this needs device pointer mode for alpha and beta, which cuBLASMp added in 0.11.0 (CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_POINTER_MODE), so the cuBLASMp backend rejected NVFP4 until now. The cuBLASMp GEMM now mirrors the cuBLASLt NVFP4 path: it uses the TN layout (the transposed column-wise data and its amax replace the row-wise data where needed), passes the FP4 data with VEC16_UE4M3 block scales, and computes alpha on the stream into a small per-context device buffer that also holds the betas 0 and 1. Row-scaled NVFP4, UE5M3 block scales and quantized outputs are rejected with an error. NVFP4 needs cuBLASMp 0.11.0 or newer both at build time and at run time; the run-time check uses cublasMpGetVersion(), so an older loaded library reports an error instead of misreading the device pointers. When n / tp_size is not a multiple of 128, NVFP4 GEMM + ReduceScatter and GEMM + AllReduce use cuBLASMp's default algorithm and warn once, since the explicitly selected split algorithms need a separate B scale tensor for each output-column chunk. With cuBLASMp, the wgrad GEMM after the dgrad AllGather + GEMM all-gathers grad_output again. An NVFP4 grad output is now quantized once, with both usages, before the dgrad GEMM, which uses the row-wise data, and the wgrad all-gather takes the column-wise data, as the non-overlapped all-gather does. Quantizing it a second time would also draw different stochastic-rounding numbers than the non-overlapped path. Signed-off-by: Almog Segal <asegal@nvidia.com> * [PyTorch] Test NVFP4 comm+GEMM overlap with cuBLASMp Add NVFP4 to the cuBLASMp AllGather + GEMM and GEMM + ReduceScatter cases and to the layer overlap tests. Userbuffers does not support NVFP4, so those layer cases are skipped. The layer runner creates BF16 inputs for NVFP4, since the NVFP4 random Hadamard transform supports only BF16 input. The GEMM overlap check compares NVFP4 with the GEMM of its dequantized operands, since NVFP4 quantization error would dominate a comparison with the unquantized operands; this also checks the global scales applied through alpha. The all-gathered input reduces its amax across ranks, as TE's sequence-parallel layers do, while the reduce-scatter inputs keep per-rank global scales. Signed-off-by: Almog Segal <asegal@nvidia.com> * [PyTorch] Skip the NVFP4 cuBLASMp tests before cuBLASMp 0.11.0 NVFP4 with cuBLASMp needs cuBLASMp 0.11.0 or newer, and the GEMM reports an error with an older library. The GEMM and layer overlap tests now read the version of the cuBLASMp library that TE loaded (cublasMpGetVersion()) and skip their NVFP4 cuBLASMp cases when it is older, as they skip on GPUs without NVFP4 support. Signed-off-by: Almog Segal <asegal@nvidia.com> * [PyTorch] Compare quantized comm+GEMM tests with the dequantized operands The GEMM overlap test compared FP8 and MXFP8 outputs with a GEMM of the unquantized operands, so quantization error set its tolerances. For every quantized format, compute the reference from the dequantized test operands instead, as was already done for NVFP4. Quantization then cancels out, the test checks only the distributed GEMM, and it uses the unquantized tolerances unless the GEMM output itself is quantized. Signed-off-by: Almog Segal <asegal@nvidia.com> * Refactor UB/cuBLASMp GEMM tests Apply consistent quant-dequant pattern for reference tensors. Remove dead FP8 logic. Signed-off-by: Tim Moon <tmoon@nvidia.com> * Use BF16 input data in UB/cuBLASMp linear tests Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Almog Segal <asegal@nvidia.com> Signed-off-by: Tim Moon <tmoon@nvidia.com> Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Co-authored-by: Tim Moon <tmoon@nvidia.com> Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Signed-off-by: Deep Patel <deepatel@nvidia.com>
* [PyTorch] Support FP8 weight caching in compiled Linear Create a fresh FP8 workspace on the first microbatch so the functional Linear op does not mutate a cached input. Reuse the workspace on the remaining microbatches and preserve the fused weight-gradient accumulation fallback. Add repeated optimizer-step comparisons, deferred backward coverage, and CUDA Graph replay assertions. GPU validation is pending in this draft. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * [PyTorch] Remove redundant microbatch cache identity test Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * [PyTorch] Measure microbatch cache memory and mark CUDA Graph minibatches Add an isolated-process memory example covering eager, functional cache refresh, compile and CUDA Graph replay. Mark each complete minibatch as one CUDA Graph iteration so cached weights survive the intervening backward calls. Linear tests: 18 passed, 12 skipped, 1 existing XPASS on RTX Ada. Four 4096-wide layers peak at 452 MiB in eager and 539 MiB in compile after warmup; creating fresh caches in eager does not raise that global peak. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * [PyTorch] Keep memory diagnostics outside the microbatch PR Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> --------- Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
…IDIA#3347) * [PyTorch] Make the Triton mask-map permute respect num_out_tokens The Triton _permute_kernel received num_out_tokens explicitly marked unused, and nothing clamped destination rows against it, while the host sizes the output at exactly num_out_tokens rows: a routing map that routes more tokens than num_out_tokens wrote the excess rows past the end of the allocation. The CUDA index-map path behind the same public API drops over-capacity entries by marking them -1 in the row map (moe_permute_row_map); do the same in row-id-map pass 2, so every consumer of the map (permute, unpermute, backward) skips them via the existing -1 handling. The JAX lowering passes an explicit no-drop sentinel, keeping its current behavior unchanged. Signed-off-by: Truong Vu <truongvu0911nd@gmail.com> * [PyTorch] Harden the mask-map capacity-drop test Add a 4096-token case so the row-map kernels span several 1024-token blocks, and compare the unpermuted output with the test-utils assert_close, which scales tolerances with dtype. Signed-off-by: Alp Dener <adener@nvidia.com> --------- Signed-off-by: Truong Vu <truongvu0911nd@gmail.com> Signed-off-by: Alp Dener <adener@nvidia.com> Co-authored-by: Alp Dener <adener@nvidia.com>
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Remove unused
torch::CustomClassHolderinheritance from the three C++ Userbuffers wrappers:CommOverlapHelper,CommOverlap, andCommOverlapP2P.These classes are exposed through pybind11, with
std::shared_ptrholders for the overlap classes. They are not registered as TorchBind classes and do not use the intrusive-pointer functionality supplied byCustomClassHolder. Removing the bases eliminates an unnecessary PyTorch dependency; the existing TE base classes and pybind11 bindings are retained.Type of change
Changes
torch::CustomClassHolderbases intransformer_engine/pytorch/csrc/extensions.h.Validation
2.15.0a0+875d815502, CUDA 13.4. Reused the previously built common CUDA library; MPI, cuBLASMp, NVSHMEM and NCCL EP disabled.CPP_ONLY=1), license check andgit diff --checkpassed.Multi-GPU communication/GEMM overlap, optional backends, and the full L0 runtime suite were not run. This is not a complete stable-ABI migration. C++ consumers must rebuild against the updated class declarations.
Checklist