From c22b2d136dc116b03a2ab95f86daffedf322df6f Mon Sep 17 00:00:00 2001 From: Ludwig Schneider Date: Tue, 27 Jan 2026 10:38:15 -0600 Subject: [PATCH 01/13] instrumenting Signed-off-by: Ludwig Schneider --- cpp/tensorrt_llm/thop/allreduceOp.cpp | 86 ++++++++++++++++++++++++ tensorrt_llm/_torch/distributed/ops.py | 90 ++++++++++++++++++++++++++ 2 files changed, 176 insertions(+) diff --git a/cpp/tensorrt_llm/thop/allreduceOp.cpp b/cpp/tensorrt_llm/thop/allreduceOp.cpp index c75324251865..ef4deede5620 100644 --- a/cpp/tensorrt_llm/thop/allreduceOp.cpp +++ b/cpp/tensorrt_llm/thop/allreduceOp.cpp @@ -52,9 +52,12 @@ #include #include +#include #include #include +#include #include +#include #include // using namespace nvinfer1; @@ -83,6 +86,23 @@ struct overloaded : Ts... template overloaded(Ts...) -> overloaded; +constexpr char const* kDebugLogPath = "/Users/lschneider/git/TensorRT-LLM/.cursor/debug.log"; + +inline void write_debug_log(char const* location, char const* message, std::string const& data, + char const* hypothesisId, char const* runId = "pre-fix") +{ + using namespace std::chrono; + auto const ts = duration_cast(system_clock::now().time_since_epoch()).count(); + std::ofstream file(kDebugLogPath, std::ios::app); + if (!file.is_open()) + { + return; + } + file << "{\"sessionId\":\"debug-session\",\"runId\":\"" << runId << "\",\"hypothesisId\":\"" << hypothesisId + << "\",\"location\":\"" << location << "\",\"message\":\"" << message << "\",\"data\":" << data + << ",\"timestamp\":" << ts << "}\n"; +} + class NvmlManager { public: @@ -295,6 +315,20 @@ class AllreduceOp auto const rank = getRank(); TLLM_LOG_DEBUG( "AllReduceOp runtime strategy for rank %d: " + tensorrt_llm::kernels::toString(runtime_strategy), rank); + // #region agent log + { + std::ostringstream data; + data << "{" + << "\"rank\":" << rank << "," + << "\"seq_len\":" << seq_len << "," + << "\"hidden_size\":" << hidden_size << "," + << "\"numel\":" << size << "," + << "\"bytes_per_element\":" << bytes_per_element << "," + << "\"runtime_strategy\":" << static_cast(runtime_strategy) << "," + << "\"configured_strategy\":" << static_cast(mStrategy) << "}"; + write_debug_log("allreduceOp.cpp:AllreduceOp::run", "selected runtime strategy", data.str(), "H1"); + } + // #endregion // Dispatch to different allreduce implementations switch (runtime_strategy) @@ -453,6 +487,17 @@ class AllreduceOp // Use ProcessGroup's allreduce directly and return early if (mNcclComm.index() == 1) { + // #region agent log + { + std::ostringstream data; + data << "{" + << "\"comm_index\":1," + << "\"numel\":" << input.numel() << "," + << "\"element_size\":" << input.element_size() << "}"; + write_debug_log("allreduceOp.cpp:runNCCLAllReduceSymmetric", "process group path", data.str(), "H2"); + } + // #endregion + auto torchPg = std::get<1>(mNcclComm); torch::Tensor reduceOutput = input.clone(); @@ -551,6 +596,24 @@ class AllreduceOp void* outputPtr = windowBuffer1.ptr; // Perform allreduce + // #region agent log + { + std::ostringstream data; + data << "{" + << "\"comm_index\":0," + << "\"numel\":" << size << "," + << "\"buffer_bytes\":" << bufferSizeBytes << "," + << "\"min_registration_threshold\":" << minRegistrationThreshold << "," + << "\"n_ranks\":" << nRanks << "," + << "\"nvlink_supported\":" << (mIsNVLINKSupported ? "true" : "false") << "," + << "\"mnnvl_supported\":" << (mIsMNNVLSupported ? "true" : "false") << "," + << "\"env_threshold_set\":" << (envThreshold != nullptr ? "true" : "false") << "," + << "\"input_window_valid\":" << (windowBuffer0.isValid() ? "true" : "false") << "," + << "\"input_ptr_is_window\":" << (inputPtr == windowBuffer0.ptr ? "true" : "false") << "}"; + write_debug_log("allreduceOp.cpp:runNCCLAllReduceSymmetric", "pre ncclAllReduce", data.str(), "H3"); + } + // #endregion + NCCLCHECK_THROW(ncclAllReduce(inputPtr, outputPtr, size, (*getDtypeMap())[mType], ncclSum, comm, stream)); if (mOp == AllReduceFusionOp::NONE) @@ -1279,6 +1342,29 @@ std::vector allreduce_raw(torch::Tensor const& input, torch::opti bool const trigger_completion_at_end_) { #if ENABLE_MULTI_DEVICE + // #region agent log + { + std::ostringstream data; + data << "{" + << "\"input_numel\":" << input.numel() << "," + << "\"input_dim0\":" << input.size(0) << "," + << "\"input_dim_last\":" << input.size(-1) << "," + << "\"input_device_index\":" << input.get_device() << "," + << "\"input_is_cuda\":" << (input.is_cuda() ? "true" : "false") << "," + << "\"strategy_arg\":" << strategy_ << "," + << "\"fusion_op_arg\":" << fusion_op_ << "," + << "\"eps\":" << eps_ << "," + << "\"group_size\":" << group_.size() << "," + << "\"has_residual\":" << (residual.has_value() ? "true" : "false") << "," + << "\"has_norm_weight\":" << (norm_weight.has_value() ? "true" : "false") << "," + << "\"has_scale\":" << (scale.has_value() ? "true" : "false") << "," + << "\"has_bias\":" << (bias.has_value() ? "true" : "false") << "," + << "\"has_workspace\":" << (workspace.has_value() ? "true" : "false") << "," + << "\"trigger_completion_at_end\":" << (trigger_completion_at_end_ ? "true" : "false") << "}"; + write_debug_log("allreduceOp.cpp:allreduce_raw", "entry args", data.str(), "H1"); + } + // #endregion + auto const dtype = tensorrt_llm::runtime::TorchUtils::dataType(input.scalar_type()); auto const strategy = static_cast(int8_t(strategy_)); auto const fusion_op = static_cast(int8_t(fusion_op_)); diff --git a/tensorrt_llm/_torch/distributed/ops.py b/tensorrt_llm/_torch/distributed/ops.py index 84468dc612f0..8fee7f70ae6a 100644 --- a/tensorrt_llm/_torch/distributed/ops.py +++ b/tensorrt_llm/_torch/distributed/ops.py @@ -1,7 +1,9 @@ +import json import math import os import platform import threading +import time from typing import Dict, List, Optional, Tuple, Union import torch @@ -21,6 +23,29 @@ _thread_local = threading.local() +_DEBUG_LOG_PATH = "/Users/lschneider/git/TensorRT-LLM/.cursor/debug.log" + + +def _write_debug_log(location: str, + message: str, + data: Dict, + hypothesis_id: str, + run_id: str = "pre-fix") -> None: + payload = { + "sessionId": "debug-session", + "runId": run_id, + "hypothesisId": hypothesis_id, + "location": location, + "message": message, + "data": data, + "timestamp": int(time.time() * 1000), + } + try: + with open(_DEBUG_LOG_PATH, "a", encoding="utf-8") as log_file: + log_file.write(json.dumps(payload) + "\n") + except OSError: + return + def get_allreduce_workspace(mapping: Mapping) -> torch.LongTensor: if not hasattr(_thread_local, f'allreduce_workspaces_{mapping.pp_rank}'): @@ -790,6 +815,25 @@ def forward( if all_reduce_params is None: all_reduce_params = AllReduceParams() + # #region agent log + _write_debug_log( + "ops.py:AllReduce.forward", + "entry", + { + "tp_size": self.mapping.tp_size, + "tp_rank": self.mapping.tp_rank, + "strategy": int(allreduce_strategy), + "disable_mpi": self._disable_mpi, + "has_symm_mem": self.symm_mem_allreduce is not None, + "has_mnnvl": self.mnnvl_allreduce is not None, + "fusion_op": int(all_reduce_params.fusion_op), + "input_shape": list(input.shape), + "input_dtype": str(input.dtype), + }, + "H1", + ) + # #endregion + # Try Symmetric Memory AllReduce first if available # Note: Currently only supports NONE fusion op (plain allreduce) if self.symm_mem_allreduce and all_reduce_params.fusion_op == AllReduceFusionOp.NONE: @@ -798,6 +842,17 @@ def forward( logger.debug( f"Using SymmetricMemoryAllReduce (MULTIMEM) for input shape {input.shape}" ) + # #region agent log + _write_debug_log( + "ops.py:AllReduce.forward", + "used symm_mem allreduce", + { + "fusion_op": int(all_reduce_params.fusion_op), + "input_shape": list(input.shape), + }, + "H4", + ) + # #endregion return symm_mem_output elif self.symm_mem_allreduce and all_reduce_params.fusion_op != AllReduceFusionOp.NONE: # Log once per rank that we're skipping symm_mem due to fusion @@ -806,6 +861,17 @@ def forward( key=(self.mapping.tp_rank, all_reduce_params.fusion_op, "debug_fusion_skip"), ) + # #region agent log + _write_debug_log( + "ops.py:AllReduce.forward", + "skipped symm_mem due to fusion", + { + "fusion_op": int(all_reduce_params.fusion_op), + "input_shape": list(input.shape), + }, + "H4", + ) + # #endregion # Try MNNVL AllReduce if symm_mem didn't handle it if self.mnnvl_allreduce: @@ -836,6 +902,18 @@ def forward( "TLLM_DISABLE_ALLREDUCE_AUTOTUNE", "0") == "1" if allreduce_strategy == AllReduceStrategy.AUTO and not disable_allreduce_autotune and not self._disable_mpi: + # #region agent log + _write_debug_log( + "ops.py:AllReduce.forward", + "calling tunable_allreduce", + { + "strategy": int(allreduce_strategy), + "fusion_op": int(all_reduce_params.fusion_op), + "disable_autotune": disable_allreduce_autotune, + }, + "H5", + ) + # #endregion output = torch.ops.trtllm.tunable_allreduce( input=input, residual=all_reduce_params.residual, @@ -851,6 +929,18 @@ def forward( trigger_completion_at_end, ) else: + # #region agent log + _write_debug_log( + "ops.py:AllReduce.forward", + "calling all_reduce_op", + { + "strategy": int(allreduce_strategy), + "fusion_op": int(all_reduce_params.fusion_op), + "disable_mpi": self._disable_mpi, + }, + "H5", + ) + # #endregion output = self.all_reduce_op( input=input, residual=all_reduce_params.residual, From 50815e69629d80d411e26289f5a466cfce167be2 Mon Sep 17 00:00:00 2001 From: Ludwig Schneider Date: Tue, 27 Jan 2026 11:08:33 -0600 Subject: [PATCH 02/13] correct output Signed-off-by: Ludwig Schneider --- cpp/tensorrt_llm/thop/allreduceOp.cpp | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/cpp/tensorrt_llm/thop/allreduceOp.cpp b/cpp/tensorrt_llm/thop/allreduceOp.cpp index ef4deede5620..4a9802f33185 100644 --- a/cpp/tensorrt_llm/thop/allreduceOp.cpp +++ b/cpp/tensorrt_llm/thop/allreduceOp.cpp @@ -55,6 +55,7 @@ #include #include #include +#include #include #include #include @@ -93,14 +94,18 @@ inline void write_debug_log(char const* location, char const* message, std::stri { using namespace std::chrono; auto const ts = duration_cast(system_clock::now().time_since_epoch()).count(); + std::ostringstream line; + line << "{\"sessionId\":\"debug-session\",\"runId\":\"" << runId << "\",\"hypothesisId\":\"" << hypothesisId + << "\",\"location\":\"" << location << "\",\"message\":\"" << message << "\",\"data\":" << data + << ",\"timestamp\":" << ts << "}"; + std::fprintf(stderr, "%s\n", line.str().c_str()); + std::fflush(stderr); + std::ofstream file(kDebugLogPath, std::ios::app); - if (!file.is_open()) + if (file.is_open()) { - return; + file << line.str() << "\n"; } - file << "{\"sessionId\":\"debug-session\",\"runId\":\"" << runId << "\",\"hypothesisId\":\"" << hypothesisId - << "\",\"location\":\"" << location << "\",\"message\":\"" << message << "\",\"data\":" << data - << ",\"timestamp\":" << ts << "}\n"; } class NvmlManager From 9ea9503ed8015b0ebaab99257b284ce0e6abdc8c Mon Sep 17 00:00:00 2001 From: Ludwig Schneider Date: Tue, 27 Jan 2026 11:52:35 -0600 Subject: [PATCH 03/13] detailed log Signed-off-by: Ludwig Schneider --- cpp/tensorrt_llm/common/ncclUtils.h | 61 +++++++++++++++ cpp/tensorrt_llm/thop/allreduceOp.cpp | 102 +++++++++++++++++++++++++- 2 files changed, 160 insertions(+), 3 deletions(-) diff --git a/cpp/tensorrt_llm/common/ncclUtils.h b/cpp/tensorrt_llm/common/ncclUtils.h index c809394a482c..b71274fd7e4a 100644 --- a/cpp/tensorrt_llm/common/ncclUtils.h +++ b/cpp/tensorrt_llm/common/ncclUtils.h @@ -27,6 +27,8 @@ #include #include +#include +#include #include #include #include @@ -53,6 +55,19 @@ TRTLLM_NAMESPACE_BEGIN namespace common::nccl_util { +inline void write_nccl_window_debug_log(char const* location, char const* message, std::string const& data, + char const* hypothesisId, char const* runId = "pre-fix") +{ + using namespace std::chrono; + auto const ts = duration_cast(system_clock::now().time_since_epoch()).count(); + std::ostringstream line; + line << "{\"sessionId\":\"debug-session\",\"runId\":\"" << runId << "\",\"hypothesisId\":\"" << hypothesisId + << "\",\"location\":\"" << location << "\",\"message\":\"" << message << "\",\"data\":" << data + << ",\"timestamp\":" << ts << "}"; + std::fprintf(stderr, "%s\n", line.str().c_str()); + std::fflush(stderr); +} + //============================================================================== // NCCL Helper - Dynamic Library Loading //============================================================================== @@ -364,6 +379,28 @@ inline std::pair createNCCLWindowTensor( int64_t buffer_size = std::accumulate(shape.begin(), shape.end(), 1LL, std::multiplies()) * torch::elementSize(dtype); + // #region agent log + { + std::ostringstream data; + data << "{" + << "\"comm\":\"" << static_cast(comm) << "\"," + << "\"dtype\":" << static_cast(dtype) << "," + << "\"elem_size\":" << torch::elementSize(dtype) << "," + << "\"buffer_size\":" << buffer_size << "," + << "\"shape\":["; + for (size_t i = 0; i < shape.size(); ++i) + { + if (i != 0) + { + data << ","; + } + data << shape[i]; + } + data << "]}"; + write_nccl_window_debug_log("ncclUtils.h:createNCCLWindowTensor", "entry", data.str(), "H6"); + } + // #endregion + // Calculate strides std::vector strides_vec(shape.size()); if (!shape.empty()) @@ -388,12 +425,36 @@ inline std::pair createNCCLWindowTensor( throw std::runtime_error(oss.str()); } + // #region agent log + { + std::ostringstream data; + data << "{" + << "\"comm\":\"" << static_cast(comm) << "\"," + << "\"buffer_ptr\":\"" << buffer.ptr << "\"," + << "\"buffer_handle\":" << buffer.handle << "," + << "\"buffer_size\":" << buffer.size << "," + << "\"window\":\"" << buffer.window << "\"}"; + write_nccl_window_debug_log("ncclUtils.h:createNCCLWindowTensor", "buffer allocated", data.str(), "H6"); + } + // #endregion + // Create custom deleter that releases the buffer auto deleter = [comm, ptr = buffer.ptr](void*) { NCCLWindowAllocator::getInstance().releaseBuffer(comm, ptr); }; // Create tensor from the buffer auto tensor = torch::from_blob(buffer.ptr, shape, strides_vec, deleter, torch::dtype(dtype).device(torch::kCUDA)); + // #region agent log + { + std::ostringstream data; + data << "{" + << "\"tensor_ptr\":\"" << tensor.data_ptr() << "\"," + << "\"buffer_ptr\":\"" << buffer.ptr << "\"," + << "\"tensor_numel\":" << tensor.numel() << "}"; + write_nccl_window_debug_log("ncclUtils.h:createNCCLWindowTensor", "tensor created", data.str(), "H6"); + } + // #endregion + return std::make_pair(tensor, buffer); } diff --git a/cpp/tensorrt_llm/thop/allreduceOp.cpp b/cpp/tensorrt_llm/thop/allreduceOp.cpp index 4a9802f33185..8a0eafca0d6c 100644 --- a/cpp/tensorrt_llm/thop/allreduceOp.cpp +++ b/cpp/tensorrt_llm/thop/allreduceOp.cpp @@ -570,6 +570,17 @@ class AllreduceOp { if (bufferSizeBytes < minRegistrationThreshold) { + // #region agent log + { + std::ostringstream data; + data << "{" + << "\"buffer_bytes\":" << bufferSizeBytes << "," + << "\"min_registration_threshold\":" << minRegistrationThreshold << "}"; + write_debug_log("allreduceOp.cpp:runNCCLAllReduceSymmetric", "small buffer: skip registration", + data.str(), "H3"); + } + // #endregion + // Small buffer: use input directly without window registration TLLM_LOG_DEBUG( "[runNCCLAllReduceSymmetric] Buffer size %zu bytes < threshold %zu bytes, " @@ -579,11 +590,67 @@ class AllreduceOp } else { + // #region agent log + { + std::ostringstream data; + data << "{" + << "\"buffer_bytes\":" << bufferSizeBytes << "," + << "\"min_registration_threshold\":" << minRegistrationThreshold << "}"; + write_debug_log( + "allreduceOp.cpp:runNCCLAllReduceSymmetric", "large buffer: will register", data.str(), "H3"); + } + // #endregion + // Large buffer: create window buffer and copy input (can swap inputTensor reference) + // #region agent log + { + std::ostringstream data; + data << "{" + << "\"input_ptr\":\"" << input.data_ptr() << "\"," + << "\"buffer_bytes\":" << bufferSizeBytes << "}"; + write_debug_log("allreduceOp.cpp:runNCCLAllReduceSymmetric", + "before createNCCLWindowTensor (input)", data.str(), "H6"); + } + // #endregion + auto [symmetricInput, symmetricBuffer0] = createNCCLWindowTensor(comm, input.sizes(), input.scalar_type()); + // #region agent log + { + std::ostringstream data; + data << "{" + << "\"symmetric_ptr\":\"" << symmetricInput.data_ptr() << "\"," + << "\"buffer_ptr\":\"" << symmetricBuffer0.ptr << "\"," + << "\"buffer_size\":" << symmetricBuffer0.size << "}"; + write_debug_log("allreduceOp.cpp:runNCCLAllReduceSymmetric", "after createNCCLWindowTensor (input)", + data.str(), "H6"); + } + // #endregion + + // #region agent log + { + std::ostringstream data; + data << "{" + << "\"dst_ptr\":\"" << symmetricBuffer0.ptr << "\"," + << "\"src_ptr\":\"" << input.data_ptr() << "\"," + << "\"bytes\":" << bufferSizeBytes << "}"; + write_debug_log("allreduceOp.cpp:runNCCLAllReduceSymmetric", "before cudaMemcpyAsync (input)", + data.str(), "H6"); + } + // #endregion TLLM_CUDA_CHECK(cudaMemcpyAsync( symmetricBuffer0.ptr, input.data_ptr(), bufferSizeBytes, cudaMemcpyDeviceToDevice, stream)); + // #region agent log + { + std::ostringstream data; + data << "{" + << "\"dst_ptr\":\"" << symmetricBuffer0.ptr << "\"," + << "\"bytes\":" << bufferSizeBytes << "}"; + write_debug_log( + "allreduceOp.cpp:runNCCLAllReduceSymmetric", "after cudaMemcpyAsync (input)", data.str(), "H6"); + } + // #endregion + windowBuffer0 = symmetricBuffer0; inputTensor = symmetricInput; // Swap to window-backed tensor inputPtr = windowBuffer0.ptr; @@ -596,7 +663,27 @@ class AllreduceOp } // Use window-backed output buffer + // #region agent log + { + std::ostringstream data; + data << "{" + << "\"buffer_bytes\":" << bufferSizeBytes << "}"; + write_debug_log("allreduceOp.cpp:runNCCLAllReduceSymmetric", "before createNCCLWindowTensor (output)", + data.str(), "H6"); + } + // #endregion auto [normOut, windowBuffer1] = createNCCLWindowTensor(comm, input.sizes(), input.scalar_type()); + // #region agent log + { + std::ostringstream data; + data << "{" + << "\"output_ptr\":\"" << normOut.data_ptr() << "\"," + << "\"buffer_ptr\":\"" << windowBuffer1.ptr << "\"," + << "\"buffer_size\":" << windowBuffer1.size << "}"; + write_debug_log( + "allreduceOp.cpp:runNCCLAllReduceSymmetric", "after createNCCLWindowTensor (output)", data.str(), "H6"); + } + // #endregion torch::Tensor outputTensor = normOut; void* outputPtr = windowBuffer1.ptr; @@ -614,12 +701,21 @@ class AllreduceOp << "\"mnnvl_supported\":" << (mIsMNNVLSupported ? "true" : "false") << "," << "\"env_threshold_set\":" << (envThreshold != nullptr ? "true" : "false") << "," << "\"input_window_valid\":" << (windowBuffer0.isValid() ? "true" : "false") << "," - << "\"input_ptr_is_window\":" << (inputPtr == windowBuffer0.ptr ? "true" : "false") << "}"; - write_debug_log("allreduceOp.cpp:runNCCLAllReduceSymmetric", "pre ncclAllReduce", data.str(), "H3"); + << "\"input_ptr_is_window\":" << (inputPtr == windowBuffer0.ptr ? "true" : "false") << "," + << "\"input_ptr\":\"" << inputPtr << "\"," + << "\"output_ptr\":\"" << outputPtr << "\"}"; + write_debug_log("allreduceOp.cpp:runNCCLAllReduceSymmetric", "before ncclAllReduce", data.str(), "H3"); } // #endregion - NCCLCHECK_THROW(ncclAllReduce(inputPtr, outputPtr, size, (*getDtypeMap())[mType], ncclSum, comm, stream)); + // #region agent log + { + std::ostringstream data; + data << "{" + << "\"numel\":" << size << "}"; + write_debug_log("allreduceOp.cpp:runNCCLAllReduceSymmetric", "after ncclAllReduce", data.str(), "H3"); + } + // #endregion if (mOp == AllReduceFusionOp::NONE) { From 444ae24cc013e6154ba090a9039d82519e6a5c40 Mon Sep 17 00:00:00 2001 From: Ludwig Schneider Date: Tue, 27 Jan 2026 12:01:26 -0600 Subject: [PATCH 04/13] more debug info Signed-off-by: Ludwig Schneider --- cpp/tensorrt_llm/thop/allreduceOp.cpp | 138 ++++++++++++++++++++++---- 1 file changed, 117 insertions(+), 21 deletions(-) diff --git a/cpp/tensorrt_llm/thop/allreduceOp.cpp b/cpp/tensorrt_llm/thop/allreduceOp.cpp index 8a0eafca0d6c..2638449e4303 100644 --- a/cpp/tensorrt_llm/thop/allreduceOp.cpp +++ b/cpp/tensorrt_llm/thop/allreduceOp.cpp @@ -338,17 +338,64 @@ class AllreduceOp // Dispatch to different allreduce implementations switch (runtime_strategy) { - case AllReduceStrategyType::UB: return runUBAllReduce(input, residual, norm_weight, scale, bias); - case AllReduceStrategyType::NCCL: return runNCCLAllReduce(input, residual, norm_weight, scale, bias); + case AllReduceStrategyType::UB: + { + // #region agent log + write_debug_log("allreduceOp.cpp:AllreduceOp::run", "dispatch UB", "{}", "H8"); + // #endregion + auto result = runUBAllReduce(input, residual, norm_weight, scale, bias); + // #region agent log + write_debug_log("allreduceOp.cpp:AllreduceOp::run", "return UB", "{}", "H8"); + // #endregion + return result; + } + case AllReduceStrategyType::NCCL: + { + // #region agent log + write_debug_log("allreduceOp.cpp:AllreduceOp::run", "dispatch NCCL", "{}", "H8"); + // #endregion + auto result = runNCCLAllReduce(input, residual, norm_weight, scale, bias); + // #region agent log + write_debug_log("allreduceOp.cpp:AllreduceOp::run", "return NCCL", "{}", "H8"); + // #endregion + return result; + } case AllReduceStrategyType::NCCL_SYMMETRIC: - return runNCCLAllReduceSymmetric(input, residual, norm_weight, scale, bias); + { + // #region agent log + write_debug_log("allreduceOp.cpp:AllreduceOp::run", "dispatch NCCL_SYMMETRIC", "{}", "H8"); + // #endregion + auto result = runNCCLAllReduceSymmetric(input, residual, norm_weight, scale, bias); + // #region agent log + write_debug_log("allreduceOp.cpp:AllreduceOp::run", "return NCCL_SYMMETRIC", "{}", "H8"); + // #endregion + return result; + } case AllReduceStrategyType::MIN_LATENCY: case AllReduceStrategyType::ONESHOT: case AllReduceStrategyType::TWOSHOT: - return runFusionAllReduce( + { + // #region agent log + write_debug_log("allreduceOp.cpp:AllreduceOp::run", "dispatch FUSION", "{}", "H8"); + // #endregion + auto result = runFusionAllReduce( input, residual, norm_weight, scale, bias, trigger_completion_at_end, workspace, runtime_strategy); + // #region agent log + write_debug_log("allreduceOp.cpp:AllreduceOp::run", "return FUSION", "{}", "H8"); + // #endregion + return result; + } case AllReduceStrategyType::LOWPRECISION: - return runLowPrecisionAllReduce(input, residual, norm_weight, scale, bias); + { + // #region agent log + write_debug_log("allreduceOp.cpp:AllreduceOp::run", "dispatch LOWPRECISION", "{}", "H8"); + // #endregion + auto result = runLowPrecisionAllReduce(input, residual, norm_weight, scale, bias); + // #region agent log + write_debug_log("allreduceOp.cpp:AllreduceOp::run", "return LOWPRECISION", "{}", "H8"); + // #endregion + return result; + } default: TORCH_CHECK(false, "Invalid runtime strategy"); return {}; } } @@ -457,22 +504,61 @@ class AllreduceOp { torch::Tensor reduce_output; - std::visit(overloaded{[&](std::shared_ptr& rawComm) - { - auto stream = at::cuda::getCurrentCUDAStream(input.get_device()); - int size = input.numel(); - reduce_output = torch::empty_like(input); - NCCLCHECK_THROW(ncclAllReduce(input.data_ptr(), reduce_output.mutable_data_ptr(), size, - (*getDtypeMap())[mType], ncclSum, *rawComm, stream)); - }, - [&](c10::intrusive_ptr& torchPg) - { - reduce_output = input.clone(); - // TLLM_LOG_INFO("AllReduce Rank: %d, tensor numel: %d", torchPg->getRank(), - // reduce_output.numel()); - std::vector tensors{reduce_output}; - PGCHECK_THROW(torchPg->allreduce(tensors, {c10d::ReduceOp::SUM})); - }}, + std::visit( + overloaded{[&](std::shared_ptr& rawComm) + { + // #region agent log + { + std::ostringstream data; + data << "{" + << "\"path\":\"raw_comm\"," + << "\"numel\":" << input.numel() << "}"; + write_debug_log("allreduceOp.cpp:runNCCLAllReduce", "before ncclAllReduce", data.str(), "H7"); + } + // #endregion + auto stream = at::cuda::getCurrentCUDAStream(input.get_device()); + int size = input.numel(); + reduce_output = torch::empty_like(input); + NCCLCHECK_THROW(ncclAllReduce(input.data_ptr(), reduce_output.mutable_data_ptr(), size, + (*getDtypeMap())[mType], ncclSum, *rawComm, stream)); + // #region agent log + { + std::ostringstream data; + data << "{" + << "\"path\":\"raw_comm\"," + << "\"numel\":" << input.numel() << "}"; + write_debug_log("allreduceOp.cpp:runNCCLAllReduce", "after ncclAllReduce", data.str(), "H7"); + } + // #endregion + }, + [&](c10::intrusive_ptr& torchPg) + { + // #region agent log + { + std::ostringstream data; + data << "{" + << "\"path\":\"process_group\"," + << "\"numel\":" << input.numel() << "," + << "\"rank\":" << torchPg->getRank() << "}"; + write_debug_log("allreduceOp.cpp:runNCCLAllReduce", "before pg allreduce", data.str(), "H7"); + } + // #endregion + reduce_output = input.clone(); + // TLLM_LOG_INFO("AllReduce Rank: %d, tensor numel: %d", torchPg->getRank(), + // reduce_output.numel()); + std::vector tensors{reduce_output}; + PGCHECK_THROW(torchPg->allreduce(tensors, {c10d::ReduceOp::SUM})); + // #region agent log + { + std::ostringstream data; + data << "{" + << "\"path\":\"process_group\"," + << "\"numel\":" << input.numel() << "," + << "\"rank\":" << torchPg->getRank() << "}"; + write_debug_log("allreduceOp.cpp:runNCCLAllReduce", "after pg allreduce", data.str(), "H7"); + } + // #endregion + }}, mNcclComm); if (mOp == AllReduceFusionOp::NONE) @@ -488,6 +574,16 @@ class AllreduceOp torch::optional const& residual, torch::optional const& norm_weight, torch::optional const& scale, torch::optional const& bias) { + // #region agent log + { + std::ostringstream data; + data << "{" + << "\"comm_index\":" << mNcclComm.index() << "," + << "\"numel\":" << input.numel() << "}"; + write_debug_log("allreduceOp.cpp:runNCCLAllReduceSymmetric", "entry", data.str(), "H3"); + } + // #endregion + // Handle ProcessGroup path first - cannot extract NCCL comm for window registration // Use ProcessGroup's allreduce directly and return early if (mNcclComm.index() == 1) From 291765d34c6fcd58ee787defedd33a4110b53aba Mon Sep 17 00:00:00 2001 From: Ludwig Schneider Date: Tue, 27 Jan 2026 13:16:30 -0600 Subject: [PATCH 05/13] instrument NCCL failure Signed-off-by: Ludwig Schneider --- cpp/tensorrt_llm/common/ncclUtils.cpp | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/cpp/tensorrt_llm/common/ncclUtils.cpp b/cpp/tensorrt_llm/common/ncclUtils.cpp index 21b976d73d99..2c93c0834e62 100644 --- a/cpp/tensorrt_llm/common/ncclUtils.cpp +++ b/cpp/tensorrt_llm/common/ncclUtils.cpp @@ -518,6 +518,17 @@ NCCLWindowBuffer NCCLWindowAllocator::allocateAndRegisterBuffer(ncclComm_t comm, ncclResult_t allocResult = ncclMemAllocFunc(&buffer.ptr, size); if (allocResult != ncclSuccess) { + // #region agent log + { + std::ostringstream data; + data << "{" + << "\"alloc_result\":" << static_cast(allocResult) << "," + << "\"alloc_error\":\"" << ncclGetErrorString(allocResult) << "\"," + << "\"size\":" << size << "}"; + write_nccl_window_debug_log( + "ncclUtils.cpp:allocateAndRegisterBuffer", "ncclMemAlloc failed", data.str(), "H6"); + } + // #endregion TLLM_THROW("ncclMemAlloc failed with error: %d", allocResult); } buffer.size = size; @@ -527,6 +538,19 @@ NCCLWindowBuffer NCCLWindowAllocator::allocateAndRegisterBuffer(ncclComm_t comm, = ncclCommWindowRegisterFunc(comm, buffer.ptr, size, &buffer.window, NCCL_WIN_COLL_SYMMETRIC); if (regResult != ncclSuccess) { + // #region agent log + { + std::ostringstream data; + data << "{" + << "\"reg_result\":" << static_cast(regResult) << "," + << "\"reg_error\":\"" << ncclGetErrorString(regResult) << "\"," + << "\"comm\":\"" << static_cast(comm) << "\"," + << "\"ptr\":\"" << buffer.ptr << "\"," + << "\"size\":" << size << "}"; + write_nccl_window_debug_log( + "ncclUtils.cpp:allocateAndRegisterBuffer", "ncclCommWindowRegister failed", data.str(), "H6"); + } + // #endregion ncclMemFree(buffer.ptr); TLLM_THROW("ncclCommWindowRegister failed with error: %d", regResult); } From 6b3f5321437d9f96dd1efc30e001d2d319f66d06 Mon Sep 17 00:00:00 2001 From: Ludwig Schneider Date: Tue, 27 Jan 2026 14:19:00 -0600 Subject: [PATCH 06/13] avoid buffer registration during capture Signed-off-by: Ludwig Schneider --- cpp/tensorrt_llm/common/ncclUtils.cpp | 25 +++++ cpp/tensorrt_llm/common/ncclUtils.h | 41 +++++++- cpp/tensorrt_llm/thop/allreduceOp.cpp | 136 +++++++++++++++++++------- 3 files changed, 161 insertions(+), 41 deletions(-) diff --git a/cpp/tensorrt_llm/common/ncclUtils.cpp b/cpp/tensorrt_llm/common/ncclUtils.cpp index 2c93c0834e62..3f26f6e97ddf 100644 --- a/cpp/tensorrt_llm/common/ncclUtils.cpp +++ b/cpp/tensorrt_llm/common/ncclUtils.cpp @@ -383,6 +383,31 @@ NCCLWindowBuffer NCCLWindowAllocator::requestBuffer(ncclComm_t comm, size_t size return bestFit->buffer; } + // No available buffer found, avoid registration during CUDA graph capture + auto stream = at::cuda::getCurrentCUDAStream(); + cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone; + auto capture_err = cudaStreamIsCapturing(stream, &capture_status); + if (capture_err != cudaSuccess) + { + TLLM_LOG_WARNING("[NCCLUtil] cudaStreamIsCapturing failed: %s", cudaGetErrorString(capture_err)); + } + if (capture_err == cudaSuccess && capture_status != cudaStreamCaptureStatusNone) + { + TLLM_LOG_WARNING("[NCCLUtil] Skipping NCCL window allocation during capture for comm %p (requested: %zu)", + static_cast(comm), size); + // #region agent log + { + std::ostringstream data; + data << "{" + << "\"comm\":\"" << static_cast(comm) << "\"," + << "\"size\":" << size << "}"; + write_nccl_window_debug_log( + "ncclUtils.cpp:requestBuffer", "capture detected; no available buffer", data.str(), "H6"); + } + // #endregion + return NCCLWindowBuffer(); + } + // No available buffer found, allocate a new one TLLM_LOG_TRACE( "[NCCLUtil] Allocating new NCCL window buffer for comm %p, size=%zu", static_cast(comm), size); diff --git a/cpp/tensorrt_llm/common/ncclUtils.h b/cpp/tensorrt_llm/common/ncclUtils.h index b71274fd7e4a..c3400d9b3f1e 100644 --- a/cpp/tensorrt_llm/common/ncclUtils.h +++ b/cpp/tensorrt_llm/common/ncclUtils.h @@ -21,6 +21,8 @@ #include "tensorrt_llm/common/logger.h" #if ENABLE_MULTI_DEVICE +#include +#include #include #include #endif @@ -414,15 +416,44 @@ inline std::pair createNCCLWindowTensor( // Request buffer from allocator auto& allocator = NCCLWindowAllocator::getInstance(); - auto buffer = allocator.requestBuffer(comm, buffer_size); + NCCLWindowBuffer buffer; + + try + { + buffer = allocator.requestBuffer(comm, buffer_size); + } + catch (std::exception const& e) + { + TLLM_LOG_WARNING("[createNCCLWindowTensor] requestBuffer failed; returning invalid buffer: %s", e.what()); + // #region agent log + { + std::ostringstream data; + data << "{" + << "\"buffer_size\":" << buffer_size << "," + << "\"error\":\"" << e.what() << "\"}"; + write_nccl_window_debug_log("ncclUtils.h:createNCCLWindowTensor", + "requestBuffer failed; returning invalid buffer", data.str(), "H6"); + } + // #endregion + return std::make_pair(torch::Tensor(), NCCLWindowBuffer()); + } // Defensive validation: ensure buffer is valid before proceeding if (!buffer.isValid()) { - std::ostringstream oss; - oss << "Failed to allocate NCCL window buffer: invalid buffer returned from requestBuffer " - << "(comm=" << static_cast(comm) << ", buffer_size=" << buffer_size << ")"; - throw std::runtime_error(oss.str()); + TLLM_LOG_WARNING( + "[createNCCLWindowTensor] invalid buffer returned from requestBuffer; returning invalid buffer"); + // #region agent log + { + std::ostringstream data; + data << "{" + << "\"buffer_size\":" << buffer_size << "," + << "\"comm\":\"" << static_cast(comm) << "\"}"; + write_nccl_window_debug_log( + "ncclUtils.h:createNCCLWindowTensor", "invalid buffer from requestBuffer", data.str(), "H6"); + } + // #endregion + return std::make_pair(torch::Tensor(), NCCLWindowBuffer()); } // #region agent log diff --git a/cpp/tensorrt_llm/thop/allreduceOp.cpp b/cpp/tensorrt_llm/thop/allreduceOp.cpp index 2638449e4303..9f8fa3133c4d 100644 --- a/cpp/tensorrt_llm/thop/allreduceOp.cpp +++ b/cpp/tensorrt_llm/thop/allreduceOp.cpp @@ -42,6 +42,7 @@ #include #include #include +#include #include #include #include @@ -654,8 +655,9 @@ class AllreduceOp minRegistrationThreshold = static_cast(std::atoi(envThreshold)) * input.element_size(); } - // Search for existing buffer auto& allocator = NCCLWindowAllocator::getInstance(); + + // Search for existing buffer auto windowBuffer0 = allocator.searchBuffer(comm, input.data_ptr()); torch::Tensor inputTensor = input; @@ -711,45 +713,59 @@ class AllreduceOp auto [symmetricInput, symmetricBuffer0] = createNCCLWindowTensor(comm, input.sizes(), input.scalar_type()); - // #region agent log - { - std::ostringstream data; - data << "{" - << "\"symmetric_ptr\":\"" << symmetricInput.data_ptr() << "\"," - << "\"buffer_ptr\":\"" << symmetricBuffer0.ptr << "\"," - << "\"buffer_size\":" << symmetricBuffer0.size << "}"; - write_debug_log("allreduceOp.cpp:runNCCLAllReduceSymmetric", "after createNCCLWindowTensor (input)", - data.str(), "H6"); - } - // #endregion - - // #region agent log + if (!symmetricBuffer0.isValid()) { - std::ostringstream data; - data << "{" - << "\"dst_ptr\":\"" << symmetricBuffer0.ptr << "\"," - << "\"src_ptr\":\"" << input.data_ptr() << "\"," - << "\"bytes\":" << bufferSizeBytes << "}"; - write_debug_log("allreduceOp.cpp:runNCCLAllReduceSymmetric", "before cudaMemcpyAsync (input)", - data.str(), "H6"); + TLLM_LOG_WARNING( + "[runNCCLAllReduceSymmetric] NCCL window not available during capture; " + "using input buffer directly for symmetric allreduce"); + // #region agent log + write_debug_log("allreduceOp.cpp:runNCCLAllReduceSymmetric", + "invalid input window; using input buffer", "{}", "H6"); + // #endregion + // inputTensor and inputPtr remain pointing to original input } - // #endregion - TLLM_CUDA_CHECK(cudaMemcpyAsync( - symmetricBuffer0.ptr, input.data_ptr(), bufferSizeBytes, cudaMemcpyDeviceToDevice, stream)); - // #region agent log + else { - std::ostringstream data; - data << "{" - << "\"dst_ptr\":\"" << symmetricBuffer0.ptr << "\"," - << "\"bytes\":" << bufferSizeBytes << "}"; - write_debug_log( - "allreduceOp.cpp:runNCCLAllReduceSymmetric", "after cudaMemcpyAsync (input)", data.str(), "H6"); - } - // #endregion + // #region agent log + { + std::ostringstream data; + data << "{" + << "\"symmetric_ptr\":\"" << symmetricInput.data_ptr() << "\"," + << "\"buffer_ptr\":\"" << symmetricBuffer0.ptr << "\"," + << "\"buffer_size\":" << symmetricBuffer0.size << "}"; + write_debug_log("allreduceOp.cpp:runNCCLAllReduceSymmetric", + "after createNCCLWindowTensor (input)", data.str(), "H6"); + } + // #endregion + + // #region agent log + { + std::ostringstream data; + data << "{" + << "\"dst_ptr\":\"" << symmetricBuffer0.ptr << "\"," + << "\"src_ptr\":\"" << input.data_ptr() << "\"," + << "\"bytes\":" << bufferSizeBytes << "}"; + write_debug_log("allreduceOp.cpp:runNCCLAllReduceSymmetric", "before cudaMemcpyAsync (input)", + data.str(), "H6"); + } + // #endregion + TLLM_CUDA_CHECK(cudaMemcpyAsync( + symmetricBuffer0.ptr, input.data_ptr(), bufferSizeBytes, cudaMemcpyDeviceToDevice, stream)); + // #region agent log + { + std::ostringstream data; + data << "{" + << "\"dst_ptr\":\"" << symmetricBuffer0.ptr << "\"," + << "\"bytes\":" << bufferSizeBytes << "}"; + write_debug_log("allreduceOp.cpp:runNCCLAllReduceSymmetric", "after cudaMemcpyAsync (input)", + data.str(), "H6"); + } + // #endregion - windowBuffer0 = symmetricBuffer0; - inputTensor = symmetricInput; // Swap to window-backed tensor - inputPtr = windowBuffer0.ptr; + windowBuffer0 = symmetricBuffer0; + inputTensor = symmetricInput; // Swap to window-backed tensor + inputPtr = windowBuffer0.ptr; + } } } else @@ -769,6 +785,54 @@ class AllreduceOp } // #endregion auto [normOut, windowBuffer1] = createNCCLWindowTensor(comm, input.sizes(), input.scalar_type()); + if (!windowBuffer1.isValid()) + { + TLLM_LOG_WARNING( + "[runNCCLAllReduceSymmetric] NCCL window not available during capture; " + "using plain CUDA tensor for output"); + // #region agent log + write_debug_log("allreduceOp.cpp:runNCCLAllReduceSymmetric", + "invalid output window; using torch output tensor", "{}", "H6"); + // #endregion + torch::Tensor outputTensor = torch::empty_like(inputTensor); + void* outputPtr = outputTensor.data_ptr(); + // Perform allreduce + // #region agent log + { + std::ostringstream data; + data << "{" + << "\"comm_index\":0," + << "\"numel\":" << size << "," + << "\"buffer_bytes\":" << bufferSizeBytes << "," + << "\"min_registration_threshold\":" << minRegistrationThreshold << "," + << "\"n_ranks\":" << nRanks << "," + << "\"nvlink_supported\":" << (mIsNVLINKSupported ? "true" : "false") << "," + << "\"mnnvl_supported\":" << (mIsMNNVLSupported ? "true" : "false") << "," + << "\"env_threshold_set\":" << (envThreshold != nullptr ? "true" : "false") << "," + << "\"input_window_valid\":" << (windowBuffer0.isValid() ? "true" : "false") << "," + << "\"input_ptr_is_window\":" << (inputPtr == windowBuffer0.ptr ? "true" : "false") << "," + << "\"input_ptr\":\"" << inputPtr << "\"," + << "\"output_ptr\":\"" << outputPtr << "\"}"; + write_debug_log("allreduceOp.cpp:runNCCLAllReduceSymmetric", "before ncclAllReduce", data.str(), "H3"); + } + // #endregion + NCCLCHECK_THROW(ncclAllReduce(inputPtr, outputPtr, size, (*getDtypeMap())[mType], ncclSum, comm, stream)); + // #region agent log + { + std::ostringstream data; + data << "{" + << "\"numel\":" << size << "}"; + write_debug_log("allreduceOp.cpp:runNCCLAllReduceSymmetric", "after ncclAllReduce", data.str(), "H3"); + } + // #endregion + + if (mOp == AllReduceFusionOp::NONE) + { + return {outputTensor}; + } + + return fallbackRunSubsequentOps(input, residual, norm_weight, scale, bias, outputTensor); + } // #region agent log { std::ostringstream data; From c765f2c115fdd5d30a71761d2614703f96330ba2 Mon Sep 17 00:00:00 2001 From: Ludwig Schneider Date: Tue, 27 Jan 2026 14:50:40 -0600 Subject: [PATCH 07/13] removing instrumentation Signed-off-by: Ludwig Schneider --- cpp/tensorrt_llm/common/ncclUtils.cpp | 34 --- cpp/tensorrt_llm/common/ncclUtils.h | 82 ------ cpp/tensorrt_llm/thop/allreduceOp.cpp | 360 ++----------------------- tensorrt_llm/_torch/distributed/ops.py | 90 ------- 4 files changed, 21 insertions(+), 545 deletions(-) diff --git a/cpp/tensorrt_llm/common/ncclUtils.cpp b/cpp/tensorrt_llm/common/ncclUtils.cpp index 3f26f6e97ddf..48788691a525 100644 --- a/cpp/tensorrt_llm/common/ncclUtils.cpp +++ b/cpp/tensorrt_llm/common/ncclUtils.cpp @@ -395,16 +395,6 @@ NCCLWindowBuffer NCCLWindowAllocator::requestBuffer(ncclComm_t comm, size_t size { TLLM_LOG_WARNING("[NCCLUtil] Skipping NCCL window allocation during capture for comm %p (requested: %zu)", static_cast(comm), size); - // #region agent log - { - std::ostringstream data; - data << "{" - << "\"comm\":\"" << static_cast(comm) << "\"," - << "\"size\":" << size << "}"; - write_nccl_window_debug_log( - "ncclUtils.cpp:requestBuffer", "capture detected; no available buffer", data.str(), "H6"); - } - // #endregion return NCCLWindowBuffer(); } @@ -543,17 +533,6 @@ NCCLWindowBuffer NCCLWindowAllocator::allocateAndRegisterBuffer(ncclComm_t comm, ncclResult_t allocResult = ncclMemAllocFunc(&buffer.ptr, size); if (allocResult != ncclSuccess) { - // #region agent log - { - std::ostringstream data; - data << "{" - << "\"alloc_result\":" << static_cast(allocResult) << "," - << "\"alloc_error\":\"" << ncclGetErrorString(allocResult) << "\"," - << "\"size\":" << size << "}"; - write_nccl_window_debug_log( - "ncclUtils.cpp:allocateAndRegisterBuffer", "ncclMemAlloc failed", data.str(), "H6"); - } - // #endregion TLLM_THROW("ncclMemAlloc failed with error: %d", allocResult); } buffer.size = size; @@ -563,19 +542,6 @@ NCCLWindowBuffer NCCLWindowAllocator::allocateAndRegisterBuffer(ncclComm_t comm, = ncclCommWindowRegisterFunc(comm, buffer.ptr, size, &buffer.window, NCCL_WIN_COLL_SYMMETRIC); if (regResult != ncclSuccess) { - // #region agent log - { - std::ostringstream data; - data << "{" - << "\"reg_result\":" << static_cast(regResult) << "," - << "\"reg_error\":\"" << ncclGetErrorString(regResult) << "\"," - << "\"comm\":\"" << static_cast(comm) << "\"," - << "\"ptr\":\"" << buffer.ptr << "\"," - << "\"size\":" << size << "}"; - write_nccl_window_debug_log( - "ncclUtils.cpp:allocateAndRegisterBuffer", "ncclCommWindowRegister failed", data.str(), "H6"); - } - // #endregion ncclMemFree(buffer.ptr); TLLM_THROW("ncclCommWindowRegister failed with error: %d", regResult); } diff --git a/cpp/tensorrt_llm/common/ncclUtils.h b/cpp/tensorrt_llm/common/ncclUtils.h index c3400d9b3f1e..2c87ad49ee8f 100644 --- a/cpp/tensorrt_llm/common/ncclUtils.h +++ b/cpp/tensorrt_llm/common/ncclUtils.h @@ -29,14 +29,11 @@ #include #include -#include -#include #include #include #include #include #include -#include #include #include #include @@ -57,19 +54,6 @@ TRTLLM_NAMESPACE_BEGIN namespace common::nccl_util { -inline void write_nccl_window_debug_log(char const* location, char const* message, std::string const& data, - char const* hypothesisId, char const* runId = "pre-fix") -{ - using namespace std::chrono; - auto const ts = duration_cast(system_clock::now().time_since_epoch()).count(); - std::ostringstream line; - line << "{\"sessionId\":\"debug-session\",\"runId\":\"" << runId << "\",\"hypothesisId\":\"" << hypothesisId - << "\",\"location\":\"" << location << "\",\"message\":\"" << message << "\",\"data\":" << data - << ",\"timestamp\":" << ts << "}"; - std::fprintf(stderr, "%s\n", line.str().c_str()); - std::fflush(stderr); -} - //============================================================================== // NCCL Helper - Dynamic Library Loading //============================================================================== @@ -381,28 +365,6 @@ inline std::pair createNCCLWindowTensor( int64_t buffer_size = std::accumulate(shape.begin(), shape.end(), 1LL, std::multiplies()) * torch::elementSize(dtype); - // #region agent log - { - std::ostringstream data; - data << "{" - << "\"comm\":\"" << static_cast(comm) << "\"," - << "\"dtype\":" << static_cast(dtype) << "," - << "\"elem_size\":" << torch::elementSize(dtype) << "," - << "\"buffer_size\":" << buffer_size << "," - << "\"shape\":["; - for (size_t i = 0; i < shape.size(); ++i) - { - if (i != 0) - { - data << ","; - } - data << shape[i]; - } - data << "]}"; - write_nccl_window_debug_log("ncclUtils.h:createNCCLWindowTensor", "entry", data.str(), "H6"); - } - // #endregion - // Calculate strides std::vector strides_vec(shape.size()); if (!shape.empty()) @@ -425,16 +387,6 @@ inline std::pair createNCCLWindowTensor( catch (std::exception const& e) { TLLM_LOG_WARNING("[createNCCLWindowTensor] requestBuffer failed; returning invalid buffer: %s", e.what()); - // #region agent log - { - std::ostringstream data; - data << "{" - << "\"buffer_size\":" << buffer_size << "," - << "\"error\":\"" << e.what() << "\"}"; - write_nccl_window_debug_log("ncclUtils.h:createNCCLWindowTensor", - "requestBuffer failed; returning invalid buffer", data.str(), "H6"); - } - // #endregion return std::make_pair(torch::Tensor(), NCCLWindowBuffer()); } @@ -443,49 +395,15 @@ inline std::pair createNCCLWindowTensor( { TLLM_LOG_WARNING( "[createNCCLWindowTensor] invalid buffer returned from requestBuffer; returning invalid buffer"); - // #region agent log - { - std::ostringstream data; - data << "{" - << "\"buffer_size\":" << buffer_size << "," - << "\"comm\":\"" << static_cast(comm) << "\"}"; - write_nccl_window_debug_log( - "ncclUtils.h:createNCCLWindowTensor", "invalid buffer from requestBuffer", data.str(), "H6"); - } - // #endregion return std::make_pair(torch::Tensor(), NCCLWindowBuffer()); } - // #region agent log - { - std::ostringstream data; - data << "{" - << "\"comm\":\"" << static_cast(comm) << "\"," - << "\"buffer_ptr\":\"" << buffer.ptr << "\"," - << "\"buffer_handle\":" << buffer.handle << "," - << "\"buffer_size\":" << buffer.size << "," - << "\"window\":\"" << buffer.window << "\"}"; - write_nccl_window_debug_log("ncclUtils.h:createNCCLWindowTensor", "buffer allocated", data.str(), "H6"); - } - // #endregion - // Create custom deleter that releases the buffer auto deleter = [comm, ptr = buffer.ptr](void*) { NCCLWindowAllocator::getInstance().releaseBuffer(comm, ptr); }; // Create tensor from the buffer auto tensor = torch::from_blob(buffer.ptr, shape, strides_vec, deleter, torch::dtype(dtype).device(torch::kCUDA)); - // #region agent log - { - std::ostringstream data; - data << "{" - << "\"tensor_ptr\":\"" << tensor.data_ptr() << "\"," - << "\"buffer_ptr\":\"" << buffer.ptr << "\"," - << "\"tensor_numel\":" << tensor.numel() << "}"; - write_nccl_window_debug_log("ncclUtils.h:createNCCLWindowTensor", "tensor created", data.str(), "H6"); - } - // #endregion - return std::make_pair(tensor, buffer); } diff --git a/cpp/tensorrt_llm/thop/allreduceOp.cpp b/cpp/tensorrt_llm/thop/allreduceOp.cpp index 9f8fa3133c4d..1e5c83af19e7 100644 --- a/cpp/tensorrt_llm/thop/allreduceOp.cpp +++ b/cpp/tensorrt_llm/thop/allreduceOp.cpp @@ -53,13 +53,9 @@ #include #include -#include #include #include -#include -#include #include -#include #include // using namespace nvinfer1; @@ -88,27 +84,6 @@ struct overloaded : Ts... template overloaded(Ts...) -> overloaded; -constexpr char const* kDebugLogPath = "/Users/lschneider/git/TensorRT-LLM/.cursor/debug.log"; - -inline void write_debug_log(char const* location, char const* message, std::string const& data, - char const* hypothesisId, char const* runId = "pre-fix") -{ - using namespace std::chrono; - auto const ts = duration_cast(system_clock::now().time_since_epoch()).count(); - std::ostringstream line; - line << "{\"sessionId\":\"debug-session\",\"runId\":\"" << runId << "\",\"hypothesisId\":\"" << hypothesisId - << "\",\"location\":\"" << location << "\",\"message\":\"" << message << "\",\"data\":" << data - << ",\"timestamp\":" << ts << "}"; - std::fprintf(stderr, "%s\n", line.str().c_str()); - std::fflush(stderr); - - std::ofstream file(kDebugLogPath, std::ios::app); - if (file.is_open()) - { - file << line.str() << "\n"; - } -} - class NvmlManager { public: @@ -321,82 +296,20 @@ class AllreduceOp auto const rank = getRank(); TLLM_LOG_DEBUG( "AllReduceOp runtime strategy for rank %d: " + tensorrt_llm::kernels::toString(runtime_strategy), rank); - // #region agent log - { - std::ostringstream data; - data << "{" - << "\"rank\":" << rank << "," - << "\"seq_len\":" << seq_len << "," - << "\"hidden_size\":" << hidden_size << "," - << "\"numel\":" << size << "," - << "\"bytes_per_element\":" << bytes_per_element << "," - << "\"runtime_strategy\":" << static_cast(runtime_strategy) << "," - << "\"configured_strategy\":" << static_cast(mStrategy) << "}"; - write_debug_log("allreduceOp.cpp:AllreduceOp::run", "selected runtime strategy", data.str(), "H1"); - } - // #endregion - // Dispatch to different allreduce implementations switch (runtime_strategy) { - case AllReduceStrategyType::UB: - { - // #region agent log - write_debug_log("allreduceOp.cpp:AllreduceOp::run", "dispatch UB", "{}", "H8"); - // #endregion - auto result = runUBAllReduce(input, residual, norm_weight, scale, bias); - // #region agent log - write_debug_log("allreduceOp.cpp:AllreduceOp::run", "return UB", "{}", "H8"); - // #endregion - return result; - } - case AllReduceStrategyType::NCCL: - { - // #region agent log - write_debug_log("allreduceOp.cpp:AllreduceOp::run", "dispatch NCCL", "{}", "H8"); - // #endregion - auto result = runNCCLAllReduce(input, residual, norm_weight, scale, bias); - // #region agent log - write_debug_log("allreduceOp.cpp:AllreduceOp::run", "return NCCL", "{}", "H8"); - // #endregion - return result; - } + case AllReduceStrategyType::UB: return runUBAllReduce(input, residual, norm_weight, scale, bias); + case AllReduceStrategyType::NCCL: return runNCCLAllReduce(input, residual, norm_weight, scale, bias); case AllReduceStrategyType::NCCL_SYMMETRIC: - { - // #region agent log - write_debug_log("allreduceOp.cpp:AllreduceOp::run", "dispatch NCCL_SYMMETRIC", "{}", "H8"); - // #endregion - auto result = runNCCLAllReduceSymmetric(input, residual, norm_weight, scale, bias); - // #region agent log - write_debug_log("allreduceOp.cpp:AllreduceOp::run", "return NCCL_SYMMETRIC", "{}", "H8"); - // #endregion - return result; - } + return runNCCLAllReduceSymmetric(input, residual, norm_weight, scale, bias); case AllReduceStrategyType::MIN_LATENCY: case AllReduceStrategyType::ONESHOT: case AllReduceStrategyType::TWOSHOT: - { - // #region agent log - write_debug_log("allreduceOp.cpp:AllreduceOp::run", "dispatch FUSION", "{}", "H8"); - // #endregion - auto result = runFusionAllReduce( + return runFusionAllReduce( input, residual, norm_weight, scale, bias, trigger_completion_at_end, workspace, runtime_strategy); - // #region agent log - write_debug_log("allreduceOp.cpp:AllreduceOp::run", "return FUSION", "{}", "H8"); - // #endregion - return result; - } case AllReduceStrategyType::LOWPRECISION: - { - // #region agent log - write_debug_log("allreduceOp.cpp:AllreduceOp::run", "dispatch LOWPRECISION", "{}", "H8"); - // #endregion - auto result = runLowPrecisionAllReduce(input, residual, norm_weight, scale, bias); - // #region agent log - write_debug_log("allreduceOp.cpp:AllreduceOp::run", "return LOWPRECISION", "{}", "H8"); - // #endregion - return result; - } + return runLowPrecisionAllReduce(input, residual, norm_weight, scale, bias); default: TORCH_CHECK(false, "Invalid runtime strategy"); return {}; } } @@ -505,61 +418,22 @@ class AllreduceOp { torch::Tensor reduce_output; - std::visit( - overloaded{[&](std::shared_ptr& rawComm) - { - // #region agent log - { - std::ostringstream data; - data << "{" - << "\"path\":\"raw_comm\"," - << "\"numel\":" << input.numel() << "}"; - write_debug_log("allreduceOp.cpp:runNCCLAllReduce", "before ncclAllReduce", data.str(), "H7"); - } - // #endregion - auto stream = at::cuda::getCurrentCUDAStream(input.get_device()); - int size = input.numel(); - reduce_output = torch::empty_like(input); - NCCLCHECK_THROW(ncclAllReduce(input.data_ptr(), reduce_output.mutable_data_ptr(), size, - (*getDtypeMap())[mType], ncclSum, *rawComm, stream)); - // #region agent log - { - std::ostringstream data; - data << "{" - << "\"path\":\"raw_comm\"," - << "\"numel\":" << input.numel() << "}"; - write_debug_log("allreduceOp.cpp:runNCCLAllReduce", "after ncclAllReduce", data.str(), "H7"); - } - // #endregion - }, - [&](c10::intrusive_ptr& torchPg) - { - // #region agent log - { - std::ostringstream data; - data << "{" - << "\"path\":\"process_group\"," - << "\"numel\":" << input.numel() << "," - << "\"rank\":" << torchPg->getRank() << "}"; - write_debug_log("allreduceOp.cpp:runNCCLAllReduce", "before pg allreduce", data.str(), "H7"); - } - // #endregion - reduce_output = input.clone(); - // TLLM_LOG_INFO("AllReduce Rank: %d, tensor numel: %d", torchPg->getRank(), - // reduce_output.numel()); - std::vector tensors{reduce_output}; - PGCHECK_THROW(torchPg->allreduce(tensors, {c10d::ReduceOp::SUM})); - // #region agent log - { - std::ostringstream data; - data << "{" - << "\"path\":\"process_group\"," - << "\"numel\":" << input.numel() << "," - << "\"rank\":" << torchPg->getRank() << "}"; - write_debug_log("allreduceOp.cpp:runNCCLAllReduce", "after pg allreduce", data.str(), "H7"); - } - // #endregion - }}, + std::visit(overloaded{[&](std::shared_ptr& rawComm) + { + auto stream = at::cuda::getCurrentCUDAStream(input.get_device()); + int size = input.numel(); + reduce_output = torch::empty_like(input); + NCCLCHECK_THROW(ncclAllReduce(input.data_ptr(), reduce_output.mutable_data_ptr(), size, + (*getDtypeMap())[mType], ncclSum, *rawComm, stream)); + }, + [&](c10::intrusive_ptr& torchPg) + { + reduce_output = input.clone(); + // TLLM_LOG_INFO("AllReduce Rank: %d, tensor numel: %d", torchPg->getRank(), + // reduce_output.numel()); + std::vector tensors{reduce_output}; + PGCHECK_THROW(torchPg->allreduce(tensors, {c10d::ReduceOp::SUM})); + }}, mNcclComm); if (mOp == AllReduceFusionOp::NONE) @@ -575,31 +449,10 @@ class AllreduceOp torch::optional const& residual, torch::optional const& norm_weight, torch::optional const& scale, torch::optional const& bias) { - // #region agent log - { - std::ostringstream data; - data << "{" - << "\"comm_index\":" << mNcclComm.index() << "," - << "\"numel\":" << input.numel() << "}"; - write_debug_log("allreduceOp.cpp:runNCCLAllReduceSymmetric", "entry", data.str(), "H3"); - } - // #endregion - // Handle ProcessGroup path first - cannot extract NCCL comm for window registration // Use ProcessGroup's allreduce directly and return early if (mNcclComm.index() == 1) { - // #region agent log - { - std::ostringstream data; - data << "{" - << "\"comm_index\":1," - << "\"numel\":" << input.numel() << "," - << "\"element_size\":" << input.element_size() << "}"; - write_debug_log("allreduceOp.cpp:runNCCLAllReduceSymmetric", "process group path", data.str(), "H2"); - } - // #endregion - auto torchPg = std::get<1>(mNcclComm); torch::Tensor reduceOutput = input.clone(); @@ -668,17 +521,6 @@ class AllreduceOp { if (bufferSizeBytes < minRegistrationThreshold) { - // #region agent log - { - std::ostringstream data; - data << "{" - << "\"buffer_bytes\":" << bufferSizeBytes << "," - << "\"min_registration_threshold\":" << minRegistrationThreshold << "}"; - write_debug_log("allreduceOp.cpp:runNCCLAllReduceSymmetric", "small buffer: skip registration", - data.str(), "H3"); - } - // #endregion - // Small buffer: use input directly without window registration TLLM_LOG_DEBUG( "[runNCCLAllReduceSymmetric] Buffer size %zu bytes < threshold %zu bytes, " @@ -688,29 +530,7 @@ class AllreduceOp } else { - // #region agent log - { - std::ostringstream data; - data << "{" - << "\"buffer_bytes\":" << bufferSizeBytes << "," - << "\"min_registration_threshold\":" << minRegistrationThreshold << "}"; - write_debug_log( - "allreduceOp.cpp:runNCCLAllReduceSymmetric", "large buffer: will register", data.str(), "H3"); - } - // #endregion - // Large buffer: create window buffer and copy input (can swap inputTensor reference) - // #region agent log - { - std::ostringstream data; - data << "{" - << "\"input_ptr\":\"" << input.data_ptr() << "\"," - << "\"buffer_bytes\":" << bufferSizeBytes << "}"; - write_debug_log("allreduceOp.cpp:runNCCLAllReduceSymmetric", - "before createNCCLWindowTensor (input)", data.str(), "H6"); - } - // #endregion - auto [symmetricInput, symmetricBuffer0] = createNCCLWindowTensor(comm, input.sizes(), input.scalar_type()); if (!symmetricBuffer0.isValid()) @@ -718,49 +538,12 @@ class AllreduceOp TLLM_LOG_WARNING( "[runNCCLAllReduceSymmetric] NCCL window not available during capture; " "using input buffer directly for symmetric allreduce"); - // #region agent log - write_debug_log("allreduceOp.cpp:runNCCLAllReduceSymmetric", - "invalid input window; using input buffer", "{}", "H6"); - // #endregion // inputTensor and inputPtr remain pointing to original input } else { - // #region agent log - { - std::ostringstream data; - data << "{" - << "\"symmetric_ptr\":\"" << symmetricInput.data_ptr() << "\"," - << "\"buffer_ptr\":\"" << symmetricBuffer0.ptr << "\"," - << "\"buffer_size\":" << symmetricBuffer0.size << "}"; - write_debug_log("allreduceOp.cpp:runNCCLAllReduceSymmetric", - "after createNCCLWindowTensor (input)", data.str(), "H6"); - } - // #endregion - - // #region agent log - { - std::ostringstream data; - data << "{" - << "\"dst_ptr\":\"" << symmetricBuffer0.ptr << "\"," - << "\"src_ptr\":\"" << input.data_ptr() << "\"," - << "\"bytes\":" << bufferSizeBytes << "}"; - write_debug_log("allreduceOp.cpp:runNCCLAllReduceSymmetric", "before cudaMemcpyAsync (input)", - data.str(), "H6"); - } - // #endregion TLLM_CUDA_CHECK(cudaMemcpyAsync( symmetricBuffer0.ptr, input.data_ptr(), bufferSizeBytes, cudaMemcpyDeviceToDevice, stream)); - // #region agent log - { - std::ostringstream data; - data << "{" - << "\"dst_ptr\":\"" << symmetricBuffer0.ptr << "\"," - << "\"bytes\":" << bufferSizeBytes << "}"; - write_debug_log("allreduceOp.cpp:runNCCLAllReduceSymmetric", "after cudaMemcpyAsync (input)", - data.str(), "H6"); - } - // #endregion windowBuffer0 = symmetricBuffer0; inputTensor = symmetricInput; // Swap to window-backed tensor @@ -775,56 +558,16 @@ class AllreduceOp } // Use window-backed output buffer - // #region agent log - { - std::ostringstream data; - data << "{" - << "\"buffer_bytes\":" << bufferSizeBytes << "}"; - write_debug_log("allreduceOp.cpp:runNCCLAllReduceSymmetric", "before createNCCLWindowTensor (output)", - data.str(), "H6"); - } - // #endregion auto [normOut, windowBuffer1] = createNCCLWindowTensor(comm, input.sizes(), input.scalar_type()); if (!windowBuffer1.isValid()) { TLLM_LOG_WARNING( "[runNCCLAllReduceSymmetric] NCCL window not available during capture; " "using plain CUDA tensor for output"); - // #region agent log - write_debug_log("allreduceOp.cpp:runNCCLAllReduceSymmetric", - "invalid output window; using torch output tensor", "{}", "H6"); - // #endregion torch::Tensor outputTensor = torch::empty_like(inputTensor); void* outputPtr = outputTensor.data_ptr(); // Perform allreduce - // #region agent log - { - std::ostringstream data; - data << "{" - << "\"comm_index\":0," - << "\"numel\":" << size << "," - << "\"buffer_bytes\":" << bufferSizeBytes << "," - << "\"min_registration_threshold\":" << minRegistrationThreshold << "," - << "\"n_ranks\":" << nRanks << "," - << "\"nvlink_supported\":" << (mIsNVLINKSupported ? "true" : "false") << "," - << "\"mnnvl_supported\":" << (mIsMNNVLSupported ? "true" : "false") << "," - << "\"env_threshold_set\":" << (envThreshold != nullptr ? "true" : "false") << "," - << "\"input_window_valid\":" << (windowBuffer0.isValid() ? "true" : "false") << "," - << "\"input_ptr_is_window\":" << (inputPtr == windowBuffer0.ptr ? "true" : "false") << "," - << "\"input_ptr\":\"" << inputPtr << "\"," - << "\"output_ptr\":\"" << outputPtr << "\"}"; - write_debug_log("allreduceOp.cpp:runNCCLAllReduceSymmetric", "before ncclAllReduce", data.str(), "H3"); - } - // #endregion NCCLCHECK_THROW(ncclAllReduce(inputPtr, outputPtr, size, (*getDtypeMap())[mType], ncclSum, comm, stream)); - // #region agent log - { - std::ostringstream data; - data << "{" - << "\"numel\":" << size << "}"; - write_debug_log("allreduceOp.cpp:runNCCLAllReduceSymmetric", "after ncclAllReduce", data.str(), "H3"); - } - // #endregion if (mOp == AllReduceFusionOp::NONE) { @@ -833,49 +576,11 @@ class AllreduceOp return fallbackRunSubsequentOps(input, residual, norm_weight, scale, bias, outputTensor); } - // #region agent log - { - std::ostringstream data; - data << "{" - << "\"output_ptr\":\"" << normOut.data_ptr() << "\"," - << "\"buffer_ptr\":\"" << windowBuffer1.ptr << "\"," - << "\"buffer_size\":" << windowBuffer1.size << "}"; - write_debug_log( - "allreduceOp.cpp:runNCCLAllReduceSymmetric", "after createNCCLWindowTensor (output)", data.str(), "H6"); - } - // #endregion torch::Tensor outputTensor = normOut; void* outputPtr = windowBuffer1.ptr; // Perform allreduce - // #region agent log - { - std::ostringstream data; - data << "{" - << "\"comm_index\":0," - << "\"numel\":" << size << "," - << "\"buffer_bytes\":" << bufferSizeBytes << "," - << "\"min_registration_threshold\":" << minRegistrationThreshold << "," - << "\"n_ranks\":" << nRanks << "," - << "\"nvlink_supported\":" << (mIsNVLINKSupported ? "true" : "false") << "," - << "\"mnnvl_supported\":" << (mIsMNNVLSupported ? "true" : "false") << "," - << "\"env_threshold_set\":" << (envThreshold != nullptr ? "true" : "false") << "," - << "\"input_window_valid\":" << (windowBuffer0.isValid() ? "true" : "false") << "," - << "\"input_ptr_is_window\":" << (inputPtr == windowBuffer0.ptr ? "true" : "false") << "," - << "\"input_ptr\":\"" << inputPtr << "\"," - << "\"output_ptr\":\"" << outputPtr << "\"}"; - write_debug_log("allreduceOp.cpp:runNCCLAllReduceSymmetric", "before ncclAllReduce", data.str(), "H3"); - } - // #endregion NCCLCHECK_THROW(ncclAllReduce(inputPtr, outputPtr, size, (*getDtypeMap())[mType], ncclSum, comm, stream)); - // #region agent log - { - std::ostringstream data; - data << "{" - << "\"numel\":" << size << "}"; - write_debug_log("allreduceOp.cpp:runNCCLAllReduceSymmetric", "after ncclAllReduce", data.str(), "H3"); - } - // #endregion if (mOp == AllReduceFusionOp::NONE) { @@ -1603,29 +1308,6 @@ std::vector allreduce_raw(torch::Tensor const& input, torch::opti bool const trigger_completion_at_end_) { #if ENABLE_MULTI_DEVICE - // #region agent log - { - std::ostringstream data; - data << "{" - << "\"input_numel\":" << input.numel() << "," - << "\"input_dim0\":" << input.size(0) << "," - << "\"input_dim_last\":" << input.size(-1) << "," - << "\"input_device_index\":" << input.get_device() << "," - << "\"input_is_cuda\":" << (input.is_cuda() ? "true" : "false") << "," - << "\"strategy_arg\":" << strategy_ << "," - << "\"fusion_op_arg\":" << fusion_op_ << "," - << "\"eps\":" << eps_ << "," - << "\"group_size\":" << group_.size() << "," - << "\"has_residual\":" << (residual.has_value() ? "true" : "false") << "," - << "\"has_norm_weight\":" << (norm_weight.has_value() ? "true" : "false") << "," - << "\"has_scale\":" << (scale.has_value() ? "true" : "false") << "," - << "\"has_bias\":" << (bias.has_value() ? "true" : "false") << "," - << "\"has_workspace\":" << (workspace.has_value() ? "true" : "false") << "," - << "\"trigger_completion_at_end\":" << (trigger_completion_at_end_ ? "true" : "false") << "}"; - write_debug_log("allreduceOp.cpp:allreduce_raw", "entry args", data.str(), "H1"); - } - // #endregion - auto const dtype = tensorrt_llm::runtime::TorchUtils::dataType(input.scalar_type()); auto const strategy = static_cast(int8_t(strategy_)); auto const fusion_op = static_cast(int8_t(fusion_op_)); diff --git a/tensorrt_llm/_torch/distributed/ops.py b/tensorrt_llm/_torch/distributed/ops.py index 8fee7f70ae6a..84468dc612f0 100644 --- a/tensorrt_llm/_torch/distributed/ops.py +++ b/tensorrt_llm/_torch/distributed/ops.py @@ -1,9 +1,7 @@ -import json import math import os import platform import threading -import time from typing import Dict, List, Optional, Tuple, Union import torch @@ -23,29 +21,6 @@ _thread_local = threading.local() -_DEBUG_LOG_PATH = "/Users/lschneider/git/TensorRT-LLM/.cursor/debug.log" - - -def _write_debug_log(location: str, - message: str, - data: Dict, - hypothesis_id: str, - run_id: str = "pre-fix") -> None: - payload = { - "sessionId": "debug-session", - "runId": run_id, - "hypothesisId": hypothesis_id, - "location": location, - "message": message, - "data": data, - "timestamp": int(time.time() * 1000), - } - try: - with open(_DEBUG_LOG_PATH, "a", encoding="utf-8") as log_file: - log_file.write(json.dumps(payload) + "\n") - except OSError: - return - def get_allreduce_workspace(mapping: Mapping) -> torch.LongTensor: if not hasattr(_thread_local, f'allreduce_workspaces_{mapping.pp_rank}'): @@ -815,25 +790,6 @@ def forward( if all_reduce_params is None: all_reduce_params = AllReduceParams() - # #region agent log - _write_debug_log( - "ops.py:AllReduce.forward", - "entry", - { - "tp_size": self.mapping.tp_size, - "tp_rank": self.mapping.tp_rank, - "strategy": int(allreduce_strategy), - "disable_mpi": self._disable_mpi, - "has_symm_mem": self.symm_mem_allreduce is not None, - "has_mnnvl": self.mnnvl_allreduce is not None, - "fusion_op": int(all_reduce_params.fusion_op), - "input_shape": list(input.shape), - "input_dtype": str(input.dtype), - }, - "H1", - ) - # #endregion - # Try Symmetric Memory AllReduce first if available # Note: Currently only supports NONE fusion op (plain allreduce) if self.symm_mem_allreduce and all_reduce_params.fusion_op == AllReduceFusionOp.NONE: @@ -842,17 +798,6 @@ def forward( logger.debug( f"Using SymmetricMemoryAllReduce (MULTIMEM) for input shape {input.shape}" ) - # #region agent log - _write_debug_log( - "ops.py:AllReduce.forward", - "used symm_mem allreduce", - { - "fusion_op": int(all_reduce_params.fusion_op), - "input_shape": list(input.shape), - }, - "H4", - ) - # #endregion return symm_mem_output elif self.symm_mem_allreduce and all_reduce_params.fusion_op != AllReduceFusionOp.NONE: # Log once per rank that we're skipping symm_mem due to fusion @@ -861,17 +806,6 @@ def forward( key=(self.mapping.tp_rank, all_reduce_params.fusion_op, "debug_fusion_skip"), ) - # #region agent log - _write_debug_log( - "ops.py:AllReduce.forward", - "skipped symm_mem due to fusion", - { - "fusion_op": int(all_reduce_params.fusion_op), - "input_shape": list(input.shape), - }, - "H4", - ) - # #endregion # Try MNNVL AllReduce if symm_mem didn't handle it if self.mnnvl_allreduce: @@ -902,18 +836,6 @@ def forward( "TLLM_DISABLE_ALLREDUCE_AUTOTUNE", "0") == "1" if allreduce_strategy == AllReduceStrategy.AUTO and not disable_allreduce_autotune and not self._disable_mpi: - # #region agent log - _write_debug_log( - "ops.py:AllReduce.forward", - "calling tunable_allreduce", - { - "strategy": int(allreduce_strategy), - "fusion_op": int(all_reduce_params.fusion_op), - "disable_autotune": disable_allreduce_autotune, - }, - "H5", - ) - # #endregion output = torch.ops.trtllm.tunable_allreduce( input=input, residual=all_reduce_params.residual, @@ -929,18 +851,6 @@ def forward( trigger_completion_at_end, ) else: - # #region agent log - _write_debug_log( - "ops.py:AllReduce.forward", - "calling all_reduce_op", - { - "strategy": int(allreduce_strategy), - "fusion_op": int(all_reduce_params.fusion_op), - "disable_mpi": self._disable_mpi, - }, - "H5", - ) - # #endregion output = self.all_reduce_op( input=input, residual=all_reduce_params.residual, From a03e51dd0d1d926a0043105a56782df52805183c Mon Sep 17 00:00:00 2001 From: Ludwig Schneider Date: Tue, 27 Jan 2026 14:57:12 -0600 Subject: [PATCH 08/13] clarifying and streamlining Signed-off-by: Ludwig Schneider --- cpp/tensorrt_llm/thop/allreduceOp.cpp | 19 ++++--------------- 1 file changed, 4 insertions(+), 15 deletions(-) diff --git a/cpp/tensorrt_llm/thop/allreduceOp.cpp b/cpp/tensorrt_llm/thop/allreduceOp.cpp index 1e5c83af19e7..c8c9491b51a9 100644 --- a/cpp/tensorrt_llm/thop/allreduceOp.cpp +++ b/cpp/tensorrt_llm/thop/allreduceOp.cpp @@ -536,8 +536,8 @@ class AllreduceOp if (!symmetricBuffer0.isValid()) { TLLM_LOG_WARNING( - "[runNCCLAllReduceSymmetric] NCCL window not available during capture; " - "using input buffer directly for symmetric allreduce"); + "[runNCCLAllReduceSymmetric] No valid symmetric buffer available; " + "falling back to non-symmetric ncclAllReduce (input buffer)"); // inputTensor and inputPtr remain pointing to original input } else @@ -559,25 +559,14 @@ class AllreduceOp // Use window-backed output buffer auto [normOut, windowBuffer1] = createNCCLWindowTensor(comm, input.sizes(), input.scalar_type()); + torch::Tensor outputTensor = windowBuffer1.isValid() ? normOut : torch::empty_like(inputTensor); + void* outputPtr = windowBuffer1.isValid() ? windowBuffer1.ptr : outputTensor.data_ptr(); if (!windowBuffer1.isValid()) { TLLM_LOG_WARNING( "[runNCCLAllReduceSymmetric] NCCL window not available during capture; " "using plain CUDA tensor for output"); - torch::Tensor outputTensor = torch::empty_like(inputTensor); - void* outputPtr = outputTensor.data_ptr(); - // Perform allreduce - NCCLCHECK_THROW(ncclAllReduce(inputPtr, outputPtr, size, (*getDtypeMap())[mType], ncclSum, comm, stream)); - - if (mOp == AllReduceFusionOp::NONE) - { - return {outputTensor}; - } - - return fallbackRunSubsequentOps(input, residual, norm_weight, scale, bias, outputTensor); } - torch::Tensor outputTensor = normOut; - void* outputPtr = windowBuffer1.ptr; // Perform allreduce NCCLCHECK_THROW(ncclAllReduce(inputPtr, outputPtr, size, (*getDtypeMap())[mType], ncclSum, comm, stream)); From b1e932d06e445264c51a3ef9bd3066af3a22313a Mon Sep 17 00:00:00 2001 From: Ludwig Schneider Date: Tue, 27 Jan 2026 15:00:42 -0600 Subject: [PATCH 09/13] clarifying warning Signed-off-by: Ludwig Schneider --- cpp/tensorrt_llm/thop/allreduceOp.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/tensorrt_llm/thop/allreduceOp.cpp b/cpp/tensorrt_llm/thop/allreduceOp.cpp index c8c9491b51a9..f3dd040cc6ac 100644 --- a/cpp/tensorrt_llm/thop/allreduceOp.cpp +++ b/cpp/tensorrt_llm/thop/allreduceOp.cpp @@ -564,7 +564,7 @@ class AllreduceOp if (!windowBuffer1.isValid()) { TLLM_LOG_WARNING( - "[runNCCLAllReduceSymmetric] NCCL window not available during capture; " + "[runNCCLAllReduceSymmetric] No valid symmetric buffer available; " "using plain CUDA tensor for output"); } From 294693bccd5b60968e659165b2455a7e68d8760f Mon Sep 17 00:00:00 2001 From: Ludwig Schneider Date: Tue, 27 Jan 2026 20:43:17 -0600 Subject: [PATCH 10/13] changing to lower log level debug Signed-off-by: Ludwig Schneider --- cpp/tensorrt_llm/common/ncclUtils.cpp | 4 ++-- cpp/tensorrt_llm/common/ncclUtils.h | 5 ++--- cpp/tensorrt_llm/thop/allreduceOp.cpp | 4 ++-- 3 files changed, 6 insertions(+), 7 deletions(-) diff --git a/cpp/tensorrt_llm/common/ncclUtils.cpp b/cpp/tensorrt_llm/common/ncclUtils.cpp index 48788691a525..70e2a5dcb082 100644 --- a/cpp/tensorrt_llm/common/ncclUtils.cpp +++ b/cpp/tensorrt_llm/common/ncclUtils.cpp @@ -389,11 +389,11 @@ NCCLWindowBuffer NCCLWindowAllocator::requestBuffer(ncclComm_t comm, size_t size auto capture_err = cudaStreamIsCapturing(stream, &capture_status); if (capture_err != cudaSuccess) { - TLLM_LOG_WARNING("[NCCLUtil] cudaStreamIsCapturing failed: %s", cudaGetErrorString(capture_err)); + TLLM_LOG_DEBUG("[NCCLUtil] cudaStreamIsCapturing failed: %s", cudaGetErrorString(capture_err)); } if (capture_err == cudaSuccess && capture_status != cudaStreamCaptureStatusNone) { - TLLM_LOG_WARNING("[NCCLUtil] Skipping NCCL window allocation during capture for comm %p (requested: %zu)", + TLLM_LOG_DEBUG("[NCCLUtil] Skipping NCCL window allocation during capture for comm %p (requested: %zu)", static_cast(comm), size); return NCCLWindowBuffer(); } diff --git a/cpp/tensorrt_llm/common/ncclUtils.h b/cpp/tensorrt_llm/common/ncclUtils.h index 2c87ad49ee8f..506dcc555788 100644 --- a/cpp/tensorrt_llm/common/ncclUtils.h +++ b/cpp/tensorrt_llm/common/ncclUtils.h @@ -386,15 +386,14 @@ inline std::pair createNCCLWindowTensor( } catch (std::exception const& e) { - TLLM_LOG_WARNING("[createNCCLWindowTensor] requestBuffer failed; returning invalid buffer: %s", e.what()); + TLLM_LOG_DEBUG("[createNCCLWindowTensor] requestBuffer failed; returning invalid buffer: %s", e.what()); return std::make_pair(torch::Tensor(), NCCLWindowBuffer()); } // Defensive validation: ensure buffer is valid before proceeding if (!buffer.isValid()) { - TLLM_LOG_WARNING( - "[createNCCLWindowTensor] invalid buffer returned from requestBuffer; returning invalid buffer"); + TLLM_LOG_DEBUG("[createNCCLWindowTensor] invalid buffer returned from requestBuffer; returning invalid buffer"); return std::make_pair(torch::Tensor(), NCCLWindowBuffer()); } diff --git a/cpp/tensorrt_llm/thop/allreduceOp.cpp b/cpp/tensorrt_llm/thop/allreduceOp.cpp index f3dd040cc6ac..00fe82c6fd22 100644 --- a/cpp/tensorrt_llm/thop/allreduceOp.cpp +++ b/cpp/tensorrt_llm/thop/allreduceOp.cpp @@ -535,7 +535,7 @@ class AllreduceOp = createNCCLWindowTensor(comm, input.sizes(), input.scalar_type()); if (!symmetricBuffer0.isValid()) { - TLLM_LOG_WARNING( + TLLM_LOG_DEBUG( "[runNCCLAllReduceSymmetric] No valid symmetric buffer available; " "falling back to non-symmetric ncclAllReduce (input buffer)"); // inputTensor and inputPtr remain pointing to original input @@ -563,7 +563,7 @@ class AllreduceOp void* outputPtr = windowBuffer1.isValid() ? windowBuffer1.ptr : outputTensor.data_ptr(); if (!windowBuffer1.isValid()) { - TLLM_LOG_WARNING( + TLLM_LOG_DEBUG( "[runNCCLAllReduceSymmetric] No valid symmetric buffer available; " "using plain CUDA tensor for output"); } From f32ff2c539f9904a995902a805da2041834c7c53 Mon Sep 17 00:00:00 2001 From: Ludwig Schneider Date: Tue, 6 Jan 2026 09:45:25 -0800 Subject: [PATCH 11/13] activate NCCL_SYMMETRIC auto-tuning Signed-off-by: Ludwig Schneider --- tensorrt_llm/_torch/custom_ops/torch_custom_ops.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index 5b683637c6e3..90cf4b6b2b54 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -1695,8 +1695,7 @@ def get_valid_tactics( **kwargs, ) -> List[int]: valid_strategies = [ - # TODO: NCCL_SYMMETRIC will cause hang during tuning process - # AllReduceStrategy.NCCL_SYMMETRIC.value, + AllReduceStrategy.NCCL_SYMMETRIC.value, AllReduceStrategy.NCCL.value, ] # Fallback in allreduceOp is set to NCCL_SYMMETRIC as default @@ -1725,7 +1724,7 @@ def forward( input, residual, norm_weight, scale, bias, workspace = inputs if tactic == -1: # TODO: Use NCCL instead of NCCL_SYMMETRIC to avoid hanging during tuning process - tactic = AllReduceStrategy.NCCL.value + tactic = AllReduceStrategy.NCCL_SYMMETRIC.value return torch.ops.trtllm.allreduce( input, From 925cb925d58992550a0ad38443ddb3051eb1bfc9 Mon Sep 17 00:00:00 2001 From: Ludwig Schneider Date: Wed, 7 Jan 2026 12:44:01 -0600 Subject: [PATCH 12/13] removing todo Signed-off-by: Ludwig Schneider --- .../custom_ops/tensor_lifetime_workaround.py | 127 ++++++++++++++++++ .../_torch/custom_ops/torch_custom_ops.py | 1 - 2 files changed, 127 insertions(+), 1 deletion(-) create mode 100644 tensorrt_llm/_torch/custom_ops/tensor_lifetime_workaround.py diff --git a/tensorrt_llm/_torch/custom_ops/tensor_lifetime_workaround.py b/tensorrt_llm/_torch/custom_ops/tensor_lifetime_workaround.py new file mode 100644 index 000000000000..11637b0d972e --- /dev/null +++ b/tensorrt_llm/_torch/custom_ops/tensor_lifetime_workaround.py @@ -0,0 +1,127 @@ +""" +Workaround for PyTorch tensor lifetime issue during CUDA graph capture. + +This module provides a mechanism to keep tensors alive during CUDA graph capture +by storing Python references to them. This prevents premature destruction of tensors +returned from custom C++ operators that have custom deleters. + +Issue: During CUDA graph capture, tensors returned from custom C++ operators +may have use_count=1, causing them to be destroyed immediately before PyTorch +binding can increment the reference count. This causes custom deleters to be +called prematurely, releasing buffers while NCCL operations are still using them. + +Workaround: Store references to output tensors during graph capture, ensuring +they stay alive until graph execution completes. +""" + +import threading +from typing import List, Union + +import torch + + +class TensorLifetimeRegistry: + """ + Thread-safe registry to store tensor references during CUDA graph capture. + + This ensures tensors returned from custom operators stay alive during graph + capture and execution, preventing premature deleter calls. + """ + + def __init__(self): + self._lock = threading.Lock() + # Store references per thread to handle multi-threaded scenarios + self._thread_local = threading.local() + + def _get_storage(self) -> List[List[torch.Tensor]]: + """Get thread-local storage for tensor references.""" + if not hasattr(self._thread_local, "tensor_refs"): + self._thread_local.tensor_refs = [] + return self._thread_local.tensor_refs + + def register_tensors(self, tensors: Union[torch.Tensor, List[torch.Tensor], tuple]): + """ + Register tensor(s) to keep them alive during graph capture. + + Args: + tensors: Single tensor, list of tensors, or tuple of tensors to register + """ + with self._lock: + storage = self._get_storage() + + # Convert to list of tensors + if isinstance(tensors, torch.Tensor): + tensor_list = [tensors] + elif isinstance(tensors, (list, tuple)): + tensor_list = [t for t in tensors if isinstance(t, torch.Tensor)] + else: + return # Not a tensor, ignore + + # Only register if we're in graph capture + if self.is_capturing(): + storage.append(tensor_list) + print( + f"[TensorLifetimeRegistry] Registered {len(tensor_list)} tensor(s) " + f"during graph capture (total batches: {len(storage)})" + ) + + def is_capturing(self) -> bool: + """ + Check if we're currently in CUDA graph capture. + + Returns: + True if currently capturing a CUDA graph, False otherwise + """ + try: + # Check if any stream is currently capturing + # torch.cuda.is_current_stream_capturing() checks the current stream + return torch.cuda.is_current_stream_capturing() + except (AttributeError, RuntimeError): + # Fallback: if the function doesn't exist or there's an error, assume not capturing + return False + + def clear(self): + """Clear all registered tensor references (call after graph execution completes).""" + with self._lock: + if hasattr(self._thread_local, "tensor_refs"): + count = sum(len(batch) for batch in self._thread_local.tensor_refs) + self._thread_local.tensor_refs.clear() + print(f"[TensorLifetimeRegistry] Cleared {count} tensor reference(s)") + + def get_registered_count(self) -> int: + """Get the number of registered tensor batches.""" + with self._lock: + if hasattr(self._thread_local, "tensor_refs"): + return len(self._thread_local.tensor_refs) + return 0 + + +# Global singleton instance +_tensor_registry = TensorLifetimeRegistry() + + +def register_tensor_references(tensors: Union[torch.Tensor, List[torch.Tensor], tuple]): + """ + Register tensor(s) to keep them alive during CUDA graph capture. + + This is a convenience function that uses the global registry. + + Args: + tensors: Single tensor, list of tensors, or tuple of tensors to register + """ + _tensor_registry.register_tensors(tensors) + + +def clear_tensor_references(): + """Clear all registered tensor references.""" + _tensor_registry.clear() + + +def is_graph_capturing() -> bool: + """Check if we're currently in CUDA graph capture.""" + return _tensor_registry.is_capturing() + + +def get_registered_count() -> int: + """Get the number of registered tensor batches.""" + return _tensor_registry.get_registered_count() diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index 90cf4b6b2b54..0fbb5ceacf70 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -1723,7 +1723,6 @@ def forward( ) -> torch.Tensor: input, residual, norm_weight, scale, bias, workspace = inputs if tactic == -1: - # TODO: Use NCCL instead of NCCL_SYMMETRIC to avoid hanging during tuning process tactic = AllReduceStrategy.NCCL_SYMMETRIC.value return torch.ops.trtllm.allreduce( From d3cebb540cef4409588c14d92fe8ef9a18c9ca66 Mon Sep 17 00:00:00 2001 From: Ludwig Schneider Date: Tue, 27 Jan 2026 20:46:58 -0600 Subject: [PATCH 13/13] removing unused file Signed-off-by: Ludwig Schneider --- .../custom_ops/tensor_lifetime_workaround.py | 127 ------------------ 1 file changed, 127 deletions(-) delete mode 100644 tensorrt_llm/_torch/custom_ops/tensor_lifetime_workaround.py diff --git a/tensorrt_llm/_torch/custom_ops/tensor_lifetime_workaround.py b/tensorrt_llm/_torch/custom_ops/tensor_lifetime_workaround.py deleted file mode 100644 index 11637b0d972e..000000000000 --- a/tensorrt_llm/_torch/custom_ops/tensor_lifetime_workaround.py +++ /dev/null @@ -1,127 +0,0 @@ -""" -Workaround for PyTorch tensor lifetime issue during CUDA graph capture. - -This module provides a mechanism to keep tensors alive during CUDA graph capture -by storing Python references to them. This prevents premature destruction of tensors -returned from custom C++ operators that have custom deleters. - -Issue: During CUDA graph capture, tensors returned from custom C++ operators -may have use_count=1, causing them to be destroyed immediately before PyTorch -binding can increment the reference count. This causes custom deleters to be -called prematurely, releasing buffers while NCCL operations are still using them. - -Workaround: Store references to output tensors during graph capture, ensuring -they stay alive until graph execution completes. -""" - -import threading -from typing import List, Union - -import torch - - -class TensorLifetimeRegistry: - """ - Thread-safe registry to store tensor references during CUDA graph capture. - - This ensures tensors returned from custom operators stay alive during graph - capture and execution, preventing premature deleter calls. - """ - - def __init__(self): - self._lock = threading.Lock() - # Store references per thread to handle multi-threaded scenarios - self._thread_local = threading.local() - - def _get_storage(self) -> List[List[torch.Tensor]]: - """Get thread-local storage for tensor references.""" - if not hasattr(self._thread_local, "tensor_refs"): - self._thread_local.tensor_refs = [] - return self._thread_local.tensor_refs - - def register_tensors(self, tensors: Union[torch.Tensor, List[torch.Tensor], tuple]): - """ - Register tensor(s) to keep them alive during graph capture. - - Args: - tensors: Single tensor, list of tensors, or tuple of tensors to register - """ - with self._lock: - storage = self._get_storage() - - # Convert to list of tensors - if isinstance(tensors, torch.Tensor): - tensor_list = [tensors] - elif isinstance(tensors, (list, tuple)): - tensor_list = [t for t in tensors if isinstance(t, torch.Tensor)] - else: - return # Not a tensor, ignore - - # Only register if we're in graph capture - if self.is_capturing(): - storage.append(tensor_list) - print( - f"[TensorLifetimeRegistry] Registered {len(tensor_list)} tensor(s) " - f"during graph capture (total batches: {len(storage)})" - ) - - def is_capturing(self) -> bool: - """ - Check if we're currently in CUDA graph capture. - - Returns: - True if currently capturing a CUDA graph, False otherwise - """ - try: - # Check if any stream is currently capturing - # torch.cuda.is_current_stream_capturing() checks the current stream - return torch.cuda.is_current_stream_capturing() - except (AttributeError, RuntimeError): - # Fallback: if the function doesn't exist or there's an error, assume not capturing - return False - - def clear(self): - """Clear all registered tensor references (call after graph execution completes).""" - with self._lock: - if hasattr(self._thread_local, "tensor_refs"): - count = sum(len(batch) for batch in self._thread_local.tensor_refs) - self._thread_local.tensor_refs.clear() - print(f"[TensorLifetimeRegistry] Cleared {count} tensor reference(s)") - - def get_registered_count(self) -> int: - """Get the number of registered tensor batches.""" - with self._lock: - if hasattr(self._thread_local, "tensor_refs"): - return len(self._thread_local.tensor_refs) - return 0 - - -# Global singleton instance -_tensor_registry = TensorLifetimeRegistry() - - -def register_tensor_references(tensors: Union[torch.Tensor, List[torch.Tensor], tuple]): - """ - Register tensor(s) to keep them alive during CUDA graph capture. - - This is a convenience function that uses the global registry. - - Args: - tensors: Single tensor, list of tensors, or tuple of tensors to register - """ - _tensor_registry.register_tensors(tensors) - - -def clear_tensor_references(): - """Clear all registered tensor references.""" - _tensor_registry.clear() - - -def is_graph_capturing() -> bool: - """Check if we're currently in CUDA graph capture.""" - return _tensor_registry.is_capturing() - - -def get_registered_count() -> int: - """Get the number of registered tensor batches.""" - return _tensor_registry.get_registered_count()