From 3ebb6c8351f8e045fe690d22145919e4ab77b7b5 Mon Sep 17 00:00:00 2001 From: Jinyang Yuan <154768711+jinyangyuan-nvidia@users.noreply.github.com> Date: Thu, 10 Apr 2025 00:15:13 -0700 Subject: [PATCH 1/4] perf: Eliminate the need for attention DP padding when possible Co-authored-by: raccoonliukai Signed-off-by: Jinyang Yuan <154768711+jinyangyuan-nvidia@users.noreply.github.com> --- cpp/tensorrt_llm/thop/allgatherOp.cpp | 87 +++++-- cpp/tensorrt_llm/thop/reducescatterOp.cpp | 96 ++++++-- .../_torch/auto_deploy/custom_ops/dist.py | 10 +- .../_torch/auto_deploy/distributed/trtllm.py | 6 +- .../_torch/custom_ops/cpp_custom_ops.py | 7 +- tensorrt_llm/_torch/distributed/ops.py | 122 +++++++--- .../_torch/models/modeling_deepseekv3.py | 33 ++- .../_torch/models/modeling_mixtral.py | 37 ++- tensorrt_llm/_torch/modules/fused_moe.py | 225 +++++++++--------- .../test_lists/test-db/l0_dgx_h200.yml | 1 + .../multi_gpu_modeling/test_deepseek.py | 12 +- 11 files changed, 424 insertions(+), 212 deletions(-) diff --git a/cpp/tensorrt_llm/thop/allgatherOp.cpp b/cpp/tensorrt_llm/thop/allgatherOp.cpp index b41628899473..f6b1feffccef 100644 --- a/cpp/tensorrt_llm/thop/allgatherOp.cpp +++ b/cpp/tensorrt_llm/thop/allgatherOp.cpp @@ -40,9 +40,8 @@ namespace class AllgatherOp { public: - AllgatherOp(std::set group, nvinfer1::DataType type) + AllgatherOp(std::set group) : mGroup(std::move(group)) - , mType(type) { } @@ -56,22 +55,62 @@ class AllgatherOp return 0; } - torch::Tensor run(torch::Tensor input) noexcept + torch::Tensor run(torch::Tensor input, torch::optional> all_rank_split_size) noexcept { + TLLM_CHECK_WITH_INFO(mNcclComm.get() != nullptr, "mNcclComm should be initialized before used"); auto stream = at::cuda::getCurrentCUDAStream(input.get_device()); + auto type = tensorrt_llm::runtime::TorchUtils::dataType(input.scalar_type()); std::vector outputShape = input.sizes().vec(); - outputShape.insert(outputShape.begin(), mGroup.size()); + if (all_rank_split_size.has_value()) + { + outputShape[0] = std::accumulate( + all_rank_split_size.value().begin(), all_rank_split_size.value().end(), 0, std::plus<>{}); + } + else + { + outputShape[0] *= mGroup.size(); + } auto output = torch::empty(outputShape, input.options()); - size_t size = input.numel(); - TLLM_CHECK_WITH_INFO(mNcclComm.get() != nullptr, "mNcclComm should be initialized before used"); - NCCLCHECK(ncclAllGather( - input.data_ptr(), output.mutable_data_ptr(), size, (*getDtypeMap())[mType], *mNcclComm, stream)); + if (all_rank_split_size.has_value()) + { + size_t numel_base = std::accumulate(outputShape.cbegin() + 1, outputShape.cend(), 1, std::multiplies<>{}); + int64_t split_offset = 0; + ncclGroupStart(); + for (int root = 0; root < static_cast(mGroup.size()); ++root) + { + auto split_size = all_rank_split_size.value()[root]; + NCCLCHECK(ncclBroadcast(input.data_ptr(), + output.index({torch::indexing::Slice(split_offset, torch::indexing::None)}).mutable_data_ptr(), + numel_base * split_size, (*getDtypeMap())[type], root, *mNcclComm, stream)); + split_offset += split_size; + } + ncclGroupEnd(); + } + else + { + NCCLCHECK(ncclAllGather(input.data_ptr(), output.mutable_data_ptr(), input.numel(), (*getDtypeMap())[type], + *mNcclComm, stream)); + } return output; } + std::vector run_list( + torch::TensorList input_list, torch::optional> all_rank_split_size) noexcept + { + std::vector output_list; + output_list.reserve(input_list.size()); + ncclGroupStart(); + for (auto const& input : input_list) + { + auto output = run(input, all_rank_split_size); + output_list.push_back(output); + } + ncclGroupEnd(); + return output_list; + } + private: std::set mGroup; - nvinfer1::DataType mType; std::shared_ptr mNcclComm; }; @@ -79,32 +118,52 @@ class AllgatherOp #endif // ENABLE_MULTI_DEVICE -torch::Tensor allgather(torch::Tensor input, torch::List group_) +torch::Tensor allgather( + torch::Tensor input, torch::optional> all_rank_split_size, torch::List group_) { #if ENABLE_MULTI_DEVICE - auto const type = tensorrt_llm::runtime::TorchUtils::dataType(input.scalar_type()); std::set group; for (int64_t rank : group_) { group.insert(static_cast(rank)); } - AllgatherOp op(group, type); + AllgatherOp op(group); op.initialize(); - auto output = op.run(input); + auto output = op.run(input, all_rank_split_size); return output; #else return input; #endif // ENABLE_MULTI_DEVICE } +std::vector allgather_list(torch::TensorList input_list, + torch::optional> all_rank_split_size, torch::List group_) +{ +#if ENABLE_MULTI_DEVICE + std::set group; + for (int64_t rank : group_) + { + group.insert(static_cast(rank)); + } + AllgatherOp op(group); + op.initialize(); + auto output_list = op.run_list(input_list, all_rank_split_size); + return output_list; +#else + return input_list.vec(); +#endif // ENABLE_MULTI_DEVICE +} + } // namespace torch_ext TORCH_LIBRARY_FRAGMENT(trtllm, m) { - m.def("allgather(Tensor input, int[] group) -> Tensor"); + m.def("allgather(Tensor input, int[]? all_rank_split_size, int[] group) -> Tensor"); + m.def("allgather_list(Tensor[] input_list, int[]? all_rank_split_size, int[] group) -> Tensor[]"); } TORCH_LIBRARY_IMPL(trtllm, CUDA, m) { m.impl("allgather", &torch_ext::allgather); + m.impl("allgather_list", &torch_ext::allgather_list); } diff --git a/cpp/tensorrt_llm/thop/reducescatterOp.cpp b/cpp/tensorrt_llm/thop/reducescatterOp.cpp index 054d5e2b4260..942e8f8482b3 100644 --- a/cpp/tensorrt_llm/thop/reducescatterOp.cpp +++ b/cpp/tensorrt_llm/thop/reducescatterOp.cpp @@ -40,9 +40,8 @@ namespace class ReducescatterOp { public: - ReducescatterOp(std::set group, nvinfer1::DataType type) + ReducescatterOp(std::set group) : mGroup(std::move(group)) - , mType(type) { } @@ -56,22 +55,71 @@ class ReducescatterOp return 0; } - torch::Tensor run(torch::Tensor const& input) noexcept + torch::Tensor run(torch::Tensor const& input, torch::optional> all_rank_split_size) noexcept { + TLLM_CHECK_WITH_INFO(mNcclComm.get() != nullptr, "mNcclComm should be initialized before used"); auto stream = at::cuda::getCurrentCUDAStream(input.get_device()); + auto type = tensorrt_llm::runtime::TorchUtils::dataType(input.scalar_type()); std::vector outputShape = input.sizes().vec(); - outputShape[0] = outputShape[0] / mGroup.size(); + if (all_rank_split_size.has_value()) + { + auto rank = COMM_SESSION.getRank(); + int groupRank = 0; + for (auto const& currentRank : mGroup) + { + if (rank == currentRank) + break; + ++groupRank; + } + TLLM_CHECK(static_cast(groupRank) < mGroup.size()); + outputShape[0] = all_rank_split_size.value()[groupRank]; + } + else + { + outputShape[0] = outputShape[0] / mGroup.size(); + } auto output = torch::empty(outputShape, input.options()); - size_t const size = output.numel(); - TLLM_CHECK_WITH_INFO(mNcclComm.get() != nullptr, "mNcclComm should be initialized before used"); - NCCLCHECK(ncclReduceScatter( - input.data_ptr(), output.mutable_data_ptr(), size, (*getDtypeMap())[mType], ncclSum, *mNcclComm, stream)); + if (all_rank_split_size.has_value()) + { + size_t numel_base = std::accumulate(outputShape.cbegin() + 1, outputShape.cend(), 1, std::multiplies<>{}); + int64_t split_offset = 0; + ncclGroupStart(); + for (int root = 0; root < static_cast(mGroup.size()); ++root) + { + auto split_size = all_rank_split_size.value()[root]; + NCCLCHECK( + ncclReduce(input.index({torch::indexing::Slice(split_offset, torch::indexing::None)}).data_ptr(), + output.mutable_data_ptr(), numel_base * split_size, (*getDtypeMap())[type], ncclSum, root, + *mNcclComm, stream)); + split_offset += split_size; + } + ncclGroupEnd(); + } + else + { + NCCLCHECK(ncclReduceScatter(input.data_ptr(), output.mutable_data_ptr(), output.numel(), + (*getDtypeMap())[type], ncclSum, *mNcclComm, stream)); + } return output; } + std::vector run_list( + torch::TensorList input_list, torch::optional> all_rank_split_size) noexcept + { + std::vector output_list; + output_list.reserve(input_list.size()); + ncclGroupStart(); + for (auto const& input : input_list) + { + auto output = run(input, all_rank_split_size); + output_list.push_back(output); + } + ncclGroupEnd(); + return output_list; + } + private: std::set mGroup; - nvinfer1::DataType mType; std::shared_ptr mNcclComm; }; @@ -79,32 +127,52 @@ class ReducescatterOp #endif // ENABLE_MULTI_DEVICE -extern torch::Tensor reducescatter(torch::Tensor input, torch::List group_) +extern torch::Tensor reducescatter( + torch::Tensor input, torch::optional> all_rank_split_size, torch::List group_) { #if ENABLE_MULTI_DEVICE - auto const type = tensorrt_llm::runtime::TorchUtils::dataType(input.scalar_type()); std::set group; for (int64_t rank : group_) { group.insert(static_cast(rank)); } - ReducescatterOp op(group, type); + ReducescatterOp op(group); op.initialize(); - auto output = op.run(input); + auto output = op.run(input, all_rank_split_size); return output; #else return input; #endif // ENABLE_MULTI_DEVICE } +extern std::vector reducescatter_list(torch::TensorList input_list, + torch::optional> all_rank_split_size, torch::List group_) +{ +#if ENABLE_MULTI_DEVICE + std::set group; + for (int64_t rank : group_) + { + group.insert(static_cast(rank)); + } + ReducescatterOp op(group); + op.initialize(); + auto output_list = op.run_list(input_list, all_rank_split_size); + return output_list; +#else + return input_list.vec(); +#endif // ENABLE_MULTI_DEVICE +} + } // namespace torch_ext TORCH_LIBRARY_FRAGMENT(trtllm, m) { - m.def("reducescatter(Tensor input, int[] group) -> Tensor"); + m.def("reducescatter(Tensor input, int[]? all_rank_split_size, int[] group) -> Tensor"); + m.def("reducescatter_list(Tensor[] input_list, int[]? all_rank_split_size, int[] group) -> Tensor[]"); } TORCH_LIBRARY_IMPL(trtllm, CUDA, m) { m.impl("reducescatter", &torch_ext::reducescatter); + m.impl("reducescatter_list", &torch_ext::reducescatter_list); } diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/dist.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/dist.py index 27902f322ec8..d95ada5e7cf3 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/dist.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/dist.py @@ -1,5 +1,7 @@ """Custom ops required for implementing tensor parallelism.""" +from typing import List, Optional + import torch from ..distributed import common as dist @@ -7,10 +9,14 @@ @torch.library.custom_op("dist::all_gather", mutates_args=(), device_types="cuda") -def all_gather(tensor: torch.Tensor, dim: int = 0) -> torch.Tensor: +def all_gather( + tensor: torch.Tensor, dim: int = 0, all_rank_split_size: Optional[List[int]] = None +) -> torch.Tensor: """All gather followed by concat in dim = 0. This is the default nccl behavior.""" if trtllm_dist.is_trtllm_op_available(): - return trtllm_dist.trtllm_allgather(tensor, dim=dim) + return trtllm_dist.trtllm_allgather( + tensor, dim=dim, all_rank_split_size=all_rank_split_size + ) tl = [torch.zeros_like(tensor) for _ in range(dist.get_world_size())] dist.all_gather(tl, tensor) return torch.cat(tl, dim=dim) diff --git a/tensorrt_llm/_torch/auto_deploy/distributed/trtllm.py b/tensorrt_llm/_torch/auto_deploy/distributed/trtllm.py index cada12f2100c..ef0c1edcbadd 100644 --- a/tensorrt_llm/_torch/auto_deploy/distributed/trtllm.py +++ b/tensorrt_llm/_torch/auto_deploy/distributed/trtllm.py @@ -8,10 +8,10 @@ from ...distributed import AllReduce, allgather from ...modules.linear import AllReduceFusionOp, AllReduceParams - def trtllm_allgather(tensor, dim): + def trtllm_allgather(tensor, dim, all_rank_split_size=None): rank, world_size = get_rank_world_size() p_config = Mapping(world_size=world_size, tp_size=world_size, rank=rank) - return allgather(tensor, p_config, gather_dim=dim) + return allgather(tensor, p_config, gather_dim=dim, all_rank_split_size=all_rank_split_size) def trtllm_allreduce(tensor, op, all_reduce_params=None): rank, world_size = get_rank_world_size() @@ -45,7 +45,7 @@ def fused_allreduce_residual_rmsnorm_fake( TRTLLM_OP_AVAILABLE = True except ImportError: - def trtllm_allgather(tensor, dim): + def trtllm_allgather(tensor, dim, all_rank_split_size=None): raise ImportError("TRT-LLM is not available.") def trtllm_allreduce(tensor, op): diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index a3121a385597..872446bd2c76 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -61,8 +61,11 @@ def _(residual, norm_weight, device_num_experts, scale_input, return [norm_out, residual_out] @torch.library.register_fake("trtllm::allgather") - def _(input, group): - output_shape = (len(group), *input.shape) + def _(input, all_rank_split_size, group): + if all_rank_split_size is None: + output_shape = (len(group) * input.shape[0], *input.shape[1:]) + else: + output_shape = (sum(all_rank_split_size), *input.shape[1:]) return input.new_empty(output_shape) @torch.library.register_fake("trtllm::cublas_scaled_mm") diff --git a/tensorrt_llm/_torch/distributed/ops.py b/tensorrt_llm/_torch/distributed/ops.py index 070fc7649722..5f49df8bdd0c 100644 --- a/tensorrt_llm/_torch/distributed/ops.py +++ b/tensorrt_llm/_torch/distributed/ops.py @@ -1,5 +1,5 @@ import threading -from typing import Optional, Tuple, Union +from typing import List, Optional, Tuple, Union import torch from torch import nn @@ -35,70 +35,118 @@ def userbuffers_allreduce_finalize( return output -def allgather(input: torch.Tensor, - mapping: Mapping, - gather_dim: int = -1) -> torch.Tensor: +def allgather( + input: Union[torch.Tensor, List[torch.Tensor]], + mapping: Mapping, + gather_dim: int = -1, + all_rank_split_size: Optional[List[int]] = None, +) -> Union[torch.Tensor, List[torch.Tensor]]: ''' Add an operation that performs a collective all-gather. - The input tensors in the different ranks must have the same shape. - The output tensor will be replicated among the TP group. + If 'all_rank_split_size' is 'None', the input tensors in the different ranks must have the same shape. + Otherwise, 'all_rank_split_size[i]' must be 'input.shape[gather_dim]' at rank i, and the input tensors in + the different ranks can only differ in shape at dimension `gather_dim`. - Given the 'section_size = input.shape[gather_dim]', each rank - contributes a section of its input tensor that correspond to - 'rank*section_size:(rank+1)*section_size', - and 'output.shape[gather_dim] = input.shape[gather_dim] * tp_group_size'. + The input tensors in the same TP group are concatenated at dimension 'gather_dim' to produce the output tensor. + If 'all_rank_split_size' is 'None', 'output.shape[gather_dim] = input.shape[gather_dim] * tp_group_size'. + Otherwise, 'output.shape[gather_dim] = sum(all_rank_split_size)'. - That operation is implemented using a torch op that wraps the NCCL all-gather - collective operation. See - https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/collectives.html#allgather - for details. + That operation is implemented using a torch op that wraps the NCCL all-gather collective operation or + the NCCL group call of a series of NCCL broadcast collective operations. See the following materials for details. + https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/collectives.html#allgather, + https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/collectives.html#broadcast, + https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/api/group.html. Args: - input (Tensor): The input tensor. + input (Union[Tensor, List[Tensor]]): The input tensor or tensor list. mapping (Mapping): The parallel mapping. gather_dim (int): Gather along given dimension. By default -1. + all_rank_split_size(Optional[List[int]]): An optional list indicating 'input.shape[gather_dim]' in all ranks. By default None. Returns: - The gathered tensor. + The gathered tensor or tensor list. ''' if mapping.tp_size == 1: return input - output = torch.ops.trtllm.allgather( + if all_rank_split_size is not None: + assert len(all_rank_split_size) == len(mapping.tp_group) + if isinstance(input, torch.Tensor): + assert input.shape[gather_dim] == all_rank_split_size[ + mapping.tp_rank] + else: + assert all([ + val.shape[gather_dim] == all_rank_split_size[mapping.tp_rank] + for val in input + ]) + # 'all_rank_split_size' is not needed if all inputs in the same TP group have the same shape + for split_size in all_rank_split_size[1:]: + if split_size != all_rank_split_size[0]: + break + else: + all_rank_split_size = None + + if isinstance(input, torch.Tensor): + torch_op = torch.ops.trtllm.allgather + input = input.movedim(gather_dim, 0).contiguous() + else: + torch_op = torch.ops.trtllm.allgather_list + input = [val.movedim(gather_dim, 0).contiguous() for val in input] + + output = torch_op( input, + all_rank_split_size, mapping.tp_group, ) - if gather_dim < 0: - gather_dim += input.ndim - - output = torch.movedim(output, 0, gather_dim) - input_shape = input.size() - output = output.reshape(input_shape[:gather_dim] + - (mapping.tp_size * input_shape[gather_dim], ) + - input_shape[gather_dim + 1:]) + if isinstance(input, torch.Tensor): + output = output.movedim(0, gather_dim).contiguous() + else: + output = [val.movedim(0, gather_dim).contiguous() for val in output] return output -def reducescatter(input: torch.Tensor, - mapping: Mapping, - scatter_dim: int = -1) -> torch.Tensor: +def reducescatter( + input: Union[torch.Tensor, List[torch.Tensor]], + mapping: Mapping, + scatter_dim: int = -1, + all_rank_split_size: Optional[List[int]] = None, +) -> Union[torch.Tensor, List[torch.Tensor]]: if mapping.tp_size == 1: return input - output = torch.ops.trtllm.reducescatter( + if all_rank_split_size is not None: + assert len(all_rank_split_size) == len(mapping.tp_group) + sum_split_size = sum(all_rank_split_size) + if isinstance(input, torch.Tensor): + assert input.shape[scatter_dim] == sum_split_size + else: + assert all( + [val.shape[scatter_dim] == sum_split_size for val in input]) + # 'all_rank_split_size' is not needed if all outputs in the same TP group have the same shape + for split_size in all_rank_split_size[1:]: + if split_size != all_rank_split_size[0]: + break + else: + all_rank_split_size = None + + if isinstance(input, torch.Tensor): + torch_op = torch.ops.trtllm.reducescatter + input = input.movedim(scatter_dim, 0).contiguous() + else: + torch_op = torch.ops.trtllm.reducescatter_list + input = [val.movedim(scatter_dim, 0).contiguous() for val in input] + + output = torch_op( input, + all_rank_split_size, mapping.tp_group, ) - if scatter_dim < 0: - scatter_dim += input.ndim - - output = torch.movedim(output, 0, scatter_dim) - input_shape = input.size() - output = output.reshape(input_shape[:scatter_dim] + - (input_shape[scatter_dim] // mapping.tp_size, ) + - input_shape[scatter_dim + 1:]) + if isinstance(input, torch.Tensor): + output = output.movedim(0, scatter_dim).contiguous() + else: + output = [val.movedim(0, scatter_dim).contiguous() for val in output] return output diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv3.py b/tensorrt_llm/_torch/models/modeling_deepseekv3.py index 7f40ba3a7678..b38fc6ee1beb 100644 --- a/tensorrt_llm/_torch/models/modeling_deepseekv3.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv3.py @@ -457,17 +457,24 @@ def should_enable_alltoall(model_config: ModelConfig, top_k: int) -> bool: def compute_routed_output(self, hidden_states, hidden_states_fp4, all_rank_num_tokens, cutlass_min_latency_mode): # max-throughput + use_dp_padding = False if self.use_dp and self.mapping.tp_size > 1: - max_num_token = max(all_rank_num_tokens) - hidden_states = torch.nn.functional.pad( - hidden_states, - (0, 0, 0, max_num_token - hidden_states.shape[0])) # FP4 all_gather moves this bf16 allgather in to after topk and fp4 quantization # to reduce allreduce BW if disable_fp4_allgather() and not self.enable_alltoall: - hidden_states = allgather(hidden_states, - self.mapping, - gather_dim=0) + hidden_states = allgather( + hidden_states, + self.mapping, + gather_dim=0, + all_rank_split_size=all_rank_num_tokens) + elif not self.experts.is_cutlass() or (not self.experts.has_fp8_qdq + and self.experts.has_nvfp4): + # Use padding when not using the cutlass path or when x_sf in self.experts is not None + use_dp_padding = True + max_num_token = max(all_rank_num_tokens) + hidden_states = torch.nn.functional.pad( + hidden_states, + (0, 0, 0, max_num_token - hidden_states.shape[0])) router_logits = self.gate(hidden_states) @@ -475,7 +482,8 @@ def compute_routed_output(self, hidden_states, hidden_states_fp4, router_logits, cutlass_min_latency_mode, output_dtype=hidden_states.dtype, - all_rank_num_tokens=all_rank_num_tokens) + all_rank_num_tokens=all_rank_num_tokens, + use_dp_padding=use_dp_padding) return routed_output @@ -928,12 +936,11 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig]): self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size self.num_hidden_layers = config.num_hidden_layers + aux_stream_list = [torch.cuda.Stream() for _ in range(2)] self.aux_stream_dict = { - key: torch.cuda.Stream() - for key in [ - AuxStreamType.Attention, AuxStreamType.MoeShared, - AuxStreamType.MoeChunkingOverlap - ] + AuxStreamType.Attention: aux_stream_list[0], + AuxStreamType.MoeShared: aux_stream_list[0], + AuxStreamType.MoeChunkingOverlap: aux_stream_list[1], } self.embed_tokens = Embedding( diff --git a/tensorrt_llm/_torch/models/modeling_mixtral.py b/tensorrt_llm/_torch/models/modeling_mixtral.py index eb7899dd58aa..39ea39c53e2f 100644 --- a/tensorrt_llm/_torch/models/modeling_mixtral.py +++ b/tensorrt_llm/_torch/models/modeling_mixtral.py @@ -8,7 +8,7 @@ from ..attention_backend import AttentionMetadata from ..attention_backend.interface import PositionalEmbeddingParams, RopeParams -from ..model_config import ModelConfig +from ..distributed import allgather from ..models.modeling_utils import ModelConfig from ..modules.attention import Attention from ..modules.decoder_layer import DecoderLayer @@ -16,6 +16,7 @@ from ..modules.fused_moe import FusedMoE, RenormalizeMoeRoutingMethod from ..modules.linear import Linear from ..modules.rms_norm import RMSNorm +from ..utils import disable_fp4_allgather from .modeling_utils import (DecoderModel, DecoderModelForCausalLM, register_auto_model) @@ -33,7 +34,7 @@ def __init__( self.ffn_dim = config.intermediate_size self.num_experts = config.num_local_experts self.top_k = config.num_experts_per_tok - self.enable_attention_dp = model_config.mapping.enable_attention_dp + self.use_dp = model_config.mapping.enable_attention_dp # moe gate (linear layer) only runs in half/full precision for now self.gate = Linear(self.hidden_dim, @@ -54,20 +55,38 @@ def __init__( reduce_results=reduce_results, model_config=model_config) + self.mapping = model_config.mapping + def forward( self, hidden_states: torch.Tensor, attn_metadata: AttentionMetadata, ) -> torch.Tensor: all_rank_num_tokens = attn_metadata.all_rank_num_tokens - if self.enable_attention_dp and len(all_rank_num_tokens) > 1: - max_num_token = max(all_rank_num_tokens) - hidden_states = torch.nn.functional.pad( - hidden_states, - (0, 0, 0, max_num_token - hidden_states.shape[0])) + use_dp_padding = False + if self.use_dp and self.mapping.tp_size > 1: + # FP4 all_gather moves this bf16 allgather in to after topk and fp4 quantization + # to reduce allreduce BW + if disable_fp4_allgather(): + hidden_states = allgather( + hidden_states, + self.mapping, + gather_dim=0, + all_rank_split_size=all_rank_num_tokens) + elif not self.experts.is_cutlass() or (not self.experts.has_fp8_qdq + and self.experts.has_nvfp4): + # Use padding when not using the cutlass path or when x_sf in self.experts is not None + use_dp_padding = True + max_num_token = max(all_rank_num_tokens) + hidden_states = torch.nn.functional.pad( + hidden_states, + (0, 0, 0, max_num_token - hidden_states.shape[0])) router_logits = self.gate(hidden_states) - final_hidden_states = self.experts(hidden_states, router_logits, - all_rank_num_tokens) + final_hidden_states = self.experts( + hidden_states, + router_logits, + all_rank_num_tokens=all_rank_num_tokens, + use_dp_padding=use_dp_padding) return final_hidden_states diff --git a/tensorrt_llm/_torch/modules/fused_moe.py b/tensorrt_llm/_torch/modules/fused_moe.py index f5d7a15d3440..38d46a68d2b2 100755 --- a/tensorrt_llm/_torch/modules/fused_moe.py +++ b/tensorrt_llm/_torch/modules/fused_moe.py @@ -229,7 +229,7 @@ class FusedMoE(nn.Module): top_k (int): Number of top experts to select for each input token. hidden_size (int): Size of the hidden state. intermediate_size (int): Size of the intermediate state. - aux_stream (torch.cuda.Stream): Auxiliary CUDA stream to overlap chunks. + aux_stream (Optional[torch.cuda.Stream]): Auxiliary CUDA stream to overlap chunks. dtype (Optional[torch.dtype]): Data type for the weights. reduce_results (bool): Whether to reduce the results across devices. model_config (ModelConfig): Configuration object for the model. @@ -285,7 +285,7 @@ def __init__( dtype: Optional[torch.dtype] = None, reduce_results: bool = False, model_config: ModelConfig = ModelConfig(), - aux_stream: torch.cuda.Stream = None, + aux_stream: Optional[torch.cuda.Stream] = None, weight_loading_mode: MoEWeightLoadingMode = MoEWeightLoadingMode. VANILLA, apply_router_weight_on_input: bool = False, @@ -300,15 +300,6 @@ def __init__( self.intermediate_size = intermediate_size self.weight_loading_mode = weight_loading_mode - if aux_stream is None: - self.aux_stream = torch.cuda.Stream() - else: - self.aux_stream = aux_stream - self.event_dict = { - key: torch.cuda.Event() - for key in [EventType.Main, EventType.MoeChunkingOverlap] - } - self.dtype = dtype self.reduce_results = reduce_results # could be modified later @@ -318,6 +309,8 @@ def __init__( self.cluster_size = model_config.mapping.moe_cluster_size self.smart_router = True if self.cluster_size > 1 else False + self.rank = model_config.mapping.rank + self.tp_rank = model_config.mapping.moe_tp_rank self.tp_size = model_config.mapping.moe_tp_size @@ -340,11 +333,23 @@ def __init__( self.expert_start + self.expert_size_per_partition, self.num_experts) - self.moe_max_num_tokens = model_config.moe_max_num_tokens - if self.moe_max_num_tokens is None: - self.moe_max_num_tokens = model_config.max_num_tokens - if self.use_dp: - self.moe_max_num_tokens *= model_config.mapping.world_size + max_num_tokens = model_config.max_num_tokens + # The maximum number of tokens in MoE are multiplied by DP size when attention DP is enabled + if self.use_dp: + max_num_tokens *= model_config.mapping.world_size + self.moe_max_num_tokens = model_config.moe_max_num_tokens if model_config.moe_max_num_tokens is not None else max_num_tokens + # The auxiliary CUDA stream and CUDA events are only used when MoE chunking is applied + if self.moe_max_num_tokens < max_num_tokens: + self.aux_stream = aux_stream if aux_stream is not None else torch.cuda.Stream( + ) + self.event_dict = { + key: torch.cuda.Event() + for key in [EventType.Main, EventType.MoeChunkingOverlap] + } + else: + self.aux_stream = None + self.event_dict = None + # The profiler converges on the same best tactic when the number of tokens is large enough. # To avoid long profiling time, the max number of tokens used in the profiling is capped to # around 16k tokens per expert, which is well into the compute bound domain. @@ -719,52 +724,20 @@ def create_weights(self): self.register_parameter("w2_weight", w2_weight) self._weights_created = True - def all_gather(self, input_tensors): - flatten_inputs = [] - shapes = [] - dtypes = [] - lengths = [] - start_indices = [] - start_idx = 0 - for input_tensor in input_tensors: - if input_tensor is None: - continue - shapes.append(input_tensor.shape) - dtypes.append(input_tensor.dtype) - lengths.append(input_tensor.nbytes) - start_indices.append(start_idx) - start_idx += input_tensor.nbytes - flatten_input = input_tensor.view(-1).view(torch.uint8) - flatten_inputs.append(flatten_input) - - if len(flatten_inputs) == 0: - return input_tensors - - flatten_outputs = allgather( - torch.cat(flatten_inputs), - self.mapping, - gather_dim=0, - ).view(self.parallel_size, -1) - - outputs = [] - for input_tensor in input_tensors: - if input_tensor is None: - output = None - else: - dtype = dtypes.pop(0) - nbytes = lengths.pop(0) - start_idx = start_indices.pop(0) - shape = [self.parallel_size, *shapes.pop(0)] - output = flatten_outputs[:, start_idx:start_idx + - nbytes].view(dtype).view(*shape) - outputs.append(output) - return outputs - - def reducescatter_or_allreduce(self, inputs): + def reducescatter_or_allreduce( + self, + inputs, + all_rank_num_tokens: Optional[List[int]] = None, + use_dp_padding: Optional[bool] = None, + ): outputs = inputs if self.parallel_size > 1 and not self.enable_alltoall: if self.use_dp: - outputs = reducescatter(inputs, self.mapping, scatter_dim=0) + outputs = reducescatter(inputs, + self.mapping, + scatter_dim=0, + all_rank_split_size=None if + use_dp_padding else all_rank_num_tokens) elif self.reduce_results: outputs = self.all_reduce(inputs) return outputs @@ -775,7 +748,8 @@ def forward_chunk( router_logits: torch.Tensor, cutlass_min_latency_mode: bool = False, output_dtype: Optional[torch.dtype] = None, - all_rank_num_tokens=None, + all_rank_num_tokens: Optional[List[int]] = None, + use_dp_padding: Optional[bool] = None, ) -> torch.Tensor: if isinstance(x, Fp4QuantizedTensor): assert output_dtype is not None @@ -844,21 +818,24 @@ def forward_chunk( if self.use_dp and self.parallel_size > 1 and not disable_fp4_allgather( ) and not self.enable_alltoall: - # Fp4 gemm has extra scaling factor - x_sf, token_selected_experts, token_final_scales = self.all_gather( - [x_sf, token_selected_experts, token_final_scales]) - x = allgather(x, self.mapping, gather_dim=0) - if x_sf is not None: + if x_sf is None: + x, token_selected_experts, token_final_scales = allgather( + [x, token_selected_experts, token_final_scales], + self.mapping, + gather_dim=0, + all_rank_split_size=None + if use_dp_padding else all_rank_num_tokens) + else: + # Fp4 gemm has extra scaling factor + x, x_sf, token_selected_experts, token_final_scales = allgather( + [x, x_sf, token_selected_experts, token_final_scales], + self.mapping, + gather_dim=0, + all_rank_split_size=None + if use_dp_padding else all_rank_num_tokens) x_sf = reswizzle_sf(x_sf, x_row, x_col, self.scaling_vector_size) - # llama4 token final scales are already multiplied with input x - if not self.apply_router_weight_on_input: - token_final_scales = token_final_scales.flatten(0, - 1).contiguous() - token_selected_experts = token_selected_experts.flatten( - 0, 1).contiguous() - if self.smart_router and not cutlass_min_latency_mode: ep_size = self.cluster_size ep_rank = self.cluster_rank @@ -927,6 +904,7 @@ def forward( cutlass_min_latency_mode: bool = False, output_dtype: Optional[torch.dtype] = None, all_rank_num_tokens: Optional[List[int]] = None, + use_dp_padding: Optional[bool] = None, ) -> torch.Tensor: """ cutlass_min_latency_mode has no effect when trtllm_gen backend is enabled. @@ -934,7 +912,7 @@ def forward( if self.is_cutlass(): return self.forward_cutlass(x, router_logits, cutlass_min_latency_mode, output_dtype, - all_rank_num_tokens) + all_rank_num_tokens, use_dp_padding) elif self.is_trtllm(): return self.forward_trtllmgen(x, router_logits) else: @@ -949,18 +927,19 @@ def forward_cutlass( cutlass_min_latency_mode: bool = False, output_dtype: Optional[torch.dtype] = None, all_rank_num_tokens: Optional[List[int]] = None, + use_dp_padding: Optional[bool] = None, ) -> torch.Tensor: assert self.is_cutlass() - max_chunk_size = self.moe_max_num_tokens if self.use_dp: assert all_rank_num_tokens is not None - if not disable_fp4_allgather(): - max_chunk_size //= len(all_rank_num_tokens) - - num_rows = x.shape[0] + assert use_dp_padding is not None + num_rows = sum(all_rank_num_tokens) + else: + num_rows = x.shape[0] # in case of num_rows is larger than max_chunk_size, we need to split the input into multiple chunks - num_chunks = (num_rows + max_chunk_size - 1) // max_chunk_size + num_chunks = (num_rows + self.moe_max_num_tokens - + 1) // self.moe_max_num_tokens if cutlass_min_latency_mode: assert num_chunks == 1 and ( @@ -973,8 +952,12 @@ def forward_cutlass( router_logits, cutlass_min_latency_mode, output_dtype, - all_rank_num_tokens=all_rank_num_tokens) - outputs = self.reducescatter_or_allreduce(outputs) + all_rank_num_tokens=all_rank_num_tokens, + use_dp_padding=use_dp_padding) + outputs = self.reducescatter_or_allreduce( + outputs, + all_rank_num_tokens=all_rank_num_tokens, + use_dp_padding=use_dp_padding) else: def split_chunk(split_token_num: int, split_num_chunks: int): @@ -984,34 +967,31 @@ def split_chunk(split_token_num: int, split_num_chunks: int): split_num_chunks - val_mod) return split_chunk_size_list - chunk_size_list = split_chunk(x.shape[0], num_chunks) + if self.use_dp: + all_rank_chunk_size_list = [ + split_chunk(val, num_chunks) for val in all_rank_num_tokens + ] + all_rank_num_tokens_list = [[ + val[idx_chunk] for val in all_rank_chunk_size_list + ] for idx_chunk in range(num_chunks)] + chunk_size_list = all_rank_chunk_size_list[self.rank] + if self.enable_alltoall: + all_rank_num_tokens_list = [[ + 1 if val == 0 else val for val in val_list + ] for val_list in all_rank_num_tokens_list] + else: + all_rank_num_tokens_list = [None] * num_chunks + chunk_size_list = split_chunk(x.shape[0], num_chunks) x_list = x.split(chunk_size_list) router_logits_list = router_logits.split(chunk_size_list) + + if not self.enable_alltoall: + self.event_dict[EventType.Main].record() + with torch.cuda.stream(self.aux_stream): + self.event_dict[EventType.Main].wait() + outputs_list = [] - all_rank_num_tokens_list = [None] * num_chunks - if self.use_dp and self.enable_alltoall: - all_rank_chunk_size_list = [] - for single_rank_num_tokens in all_rank_num_tokens: - single_rank_num_chunks = num_chunks - single_rank_chunk_size_list = split_chunk( - single_rank_num_tokens, single_rank_num_chunks) - single_rank_chunk_size_list = [ - 1 if x == 0 else x for x in single_rank_chunk_size_list - ] - all_rank_chunk_size_list.append(single_rank_chunk_size_list) - - for chunk_id in range(num_chunks): - chunk_all_rank_num_tokens = [ - all_rank_chunk_size_list[r][chunk_id] - for r in range(len(all_rank_num_tokens)) - ] - all_rank_num_tokens_list[ - chunk_id] = chunk_all_rank_num_tokens - - self.event_dict[EventType.Main].record() - with torch.cuda.stream(self.aux_stream): - self.event_dict[EventType.Main].wait() # Postpone reduce-scatter/all-reduce to the next iteration to achieve better overlap for idx_chunk, (x, router_logits) in enumerate( zip(x_list, router_logits_list)): @@ -1022,37 +1002,50 @@ def split_chunk(split_token_num: int, split_num_chunks: int): x, router_logits, all_rank_num_tokens=all_rank_num_tokens_list[ - idx_chunk]) + idx_chunk] if self.use_dp else None, + use_dp_padding=use_dp_padding) if idx_chunk > 0: outputs_list[-1] = self.reducescatter_or_allreduce( - outputs_list[-1]) + outputs_list[-1], + all_rank_num_tokens=all_rank_num_tokens_list[ + idx_chunk - 1], + use_dp_padding=use_dp_padding) else: outputs = self.forward_chunk( x, router_logits, all_rank_num_tokens=all_rank_num_tokens_list[ - idx_chunk]) + idx_chunk] if self.use_dp else None, + use_dp_padding=use_dp_padding) with torch.cuda.stream(self.aux_stream): outputs_list[-1] = self.reducescatter_or_allreduce( - outputs_list[-1]) + outputs_list[-1], + all_rank_num_tokens=all_rank_num_tokens_list[ + idx_chunk - 1], + use_dp_padding=use_dp_padding) else: outputs = self.forward_chunk( x, router_logits, - all_rank_num_tokens=all_rank_num_tokens_list[idx_chunk]) + all_rank_num_tokens=all_rank_num_tokens_list[idx_chunk] + if self.use_dp else None) outputs_list.append(outputs) if not self.enable_alltoall: if num_chunks % 2 == 0: outputs_list[-1] = self.reducescatter_or_allreduce( - outputs_list[-1]) + outputs_list[-1], + all_rank_num_tokens=all_rank_num_tokens_list[-1], + use_dp_padding=use_dp_padding) else: with torch.cuda.stream(self.aux_stream): outputs_list[-1] = self.reducescatter_or_allreduce( - outputs_list[-1]) + outputs_list[-1], + all_rank_num_tokens=all_rank_num_tokens_list[-1], + use_dp_padding=use_dp_padding) with torch.cuda.stream(self.aux_stream): self.event_dict[EventType.MoeChunkingOverlap].record() - self.event_dict[EventType.MoeChunkingOverlap].wait() + self.event_dict[EventType.MoeChunkingOverlap].wait() outputs = torch.cat(outputs_list) if self.use_dp: rank = self.mapping.tp_rank @@ -1151,8 +1144,10 @@ def alltoall_prepare_maybe_dispatch(self, all_rank_num_tokens: list, token_final_scales = torch.nn.functional.pad( token_final_scales, (0, 0, 0, max_num_token - token_final_scales.shape[0])) - gathered_token_selected_experts, gathered_token_final_scales = self.all_gather( - [token_selected_experts, token_final_scales]) + gathered_token_selected_experts, gathered_token_final_scales = allgather( + [token_selected_experts, token_final_scales], + self.mapping, + gather_dim=0) gathered_token_selected_experts = torch.flatten( gathered_token_selected_experts.contiguous(), start_dim=0, diff --git a/tests/integration/test_lists/test-db/l0_dgx_h200.yml b/tests/integration/test_lists/test-db/l0_dgx_h200.yml index f291305fa767..cc83f48679f2 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_h200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_h200.yml @@ -18,3 +18,4 @@ l0_dgx_h200: # - accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_fp8_blockscale[throughput] # OOM - accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_fp8_blockscale[latency] # 1h - unittest/_torch/multi_gpu_modeling/test_llama4.py::test_llama4[pp1-ep1-enable_graph-tp8-trtllm-scout] + - unittest/_torch/multi_gpu_modeling -k "deepseek" diff --git a/tests/unittest/_torch/multi_gpu_modeling/test_deepseek.py b/tests/unittest/_torch/multi_gpu_modeling/test_deepseek.py index 4ce920e574e7..f4baeb32e7f8 100644 --- a/tests/unittest/_torch/multi_gpu_modeling/test_deepseek.py +++ b/tests/unittest/_torch/multi_gpu_modeling/test_deepseek.py @@ -23,8 +23,13 @@ def similar(a, b, threshold=0.9): ids=["deepseekv3_lite"]) @pytest.mark.parametrize("backend", ["TRTLLM"], ids=["trtllm"]) @pytest.mark.parametrize("quant", ["bf16"]) -@pytest.mark.parametrize("tp_size", [1], ids=["tp1"]) -def test_deepseek_streaming(model_name, backend, quant, tp_size): +@pytest.mark.parametrize("tp_size", [1, 4], ids=["tp1", "tp4"]) +@pytest.mark.parametrize("enable_attention_dp", [False, True], + ids=["adp_off", "adp_on"]) +@pytest.mark.parametrize("moe_max_num_tokens", [None, 64], + ids=["moe_chunk_off", "moe_chunk_on"]) +def test_deepseek_streaming(model_name, backend, quant, tp_size, + enable_attention_dp, moe_max_num_tokens): model_path = { "bf16": "bf16", "fp8": "fp8", @@ -61,6 +66,7 @@ def test_deepseek_streaming(model_name, backend, quant, tp_size): use_cuda_graph=False, kv_cache_dtype="auto", attn_backend=backend, + moe_max_num_tokens=moe_max_num_tokens, ) model_dir = str(llm_models_root() / model_name / model_path[quant]) @@ -73,7 +79,7 @@ def test_deepseek_streaming(model_name, backend, quant, tp_size): pytorch_backend_config=pytorch_config, moe_expert_parallel_size=-1, moe_tensor_parallel_size=-1, - enable_attention_dp=False, + enable_attention_dp=enable_attention_dp, kv_cache_config=KvCacheConfig(enable_block_reuse=False)) sampling_params = SamplingParams(max_tokens=10) From e22606c3676b9a4b57cfd0e426a91910c8e802bc Mon Sep 17 00:00:00 2001 From: Jinyang Yuan <154768711+jinyangyuan-nvidia@users.noreply.github.com> Date: Tue, 13 May 2025 06:51:07 -0700 Subject: [PATCH 2/4] Minor modifications Signed-off-by: Jinyang Yuan <154768711+jinyangyuan-nvidia@users.noreply.github.com> --- cpp/tensorrt_llm/thop/allgatherOp.cpp | 45 ++++------ cpp/tensorrt_llm/thop/reducescatterOp.cpp | 43 ++++------ .../_torch/auto_deploy/custom_ops/dist.py | 6 +- .../_torch/auto_deploy/distributed/trtllm.py | 6 +- .../_torch/custom_ops/cpp_custom_ops.py | 6 +- tensorrt_llm/_torch/distributed/ops.py | 84 +++++++++---------- .../_torch/models/modeling_deepseekv3.py | 9 +- .../_torch/models/modeling_mixtral.py | 9 +- tensorrt_llm/_torch/modules/fused_moe.py | 24 +++--- .../test_lists/test-db/l0_dgx_h100.yml | 1 + .../test_lists/test-db/l0_dgx_h200.yml | 1 - .../test_lists/test-db/l0_h100.yml | 1 - .../multi_gpu_modeling/test_deepseek.py | 14 ++-- 13 files changed, 105 insertions(+), 144 deletions(-) diff --git a/cpp/tensorrt_llm/thop/allgatherOp.cpp b/cpp/tensorrt_llm/thop/allgatherOp.cpp index f6b1feffccef..c2d9a5de4e09 100644 --- a/cpp/tensorrt_llm/thop/allgatherOp.cpp +++ b/cpp/tensorrt_llm/thop/allgatherOp.cpp @@ -22,17 +22,14 @@ #include #include #include +#include #include #include #include #include -#if ENABLE_MULTI_DEVICE -#include -#endif // ENABLE_MULTI_DEVICE namespace torch_ext { -#if ENABLE_MULTI_DEVICE namespace { @@ -55,30 +52,29 @@ class AllgatherOp return 0; } - torch::Tensor run(torch::Tensor input, torch::optional> all_rank_split_size) noexcept + torch::Tensor run(torch::Tensor input, torch::optional> sizes) noexcept { TLLM_CHECK_WITH_INFO(mNcclComm.get() != nullptr, "mNcclComm should be initialized before used"); auto stream = at::cuda::getCurrentCUDAStream(input.get_device()); auto type = tensorrt_llm::runtime::TorchUtils::dataType(input.scalar_type()); std::vector outputShape = input.sizes().vec(); - if (all_rank_split_size.has_value()) + if (sizes.has_value()) { - outputShape[0] = std::accumulate( - all_rank_split_size.value().begin(), all_rank_split_size.value().end(), 0, std::plus<>{}); + outputShape[0] = std::accumulate(sizes.value().begin(), sizes.value().end(), 0, std::plus<>{}); } else { outputShape[0] *= mGroup.size(); } auto output = torch::empty(outputShape, input.options()); - if (all_rank_split_size.has_value()) + if (sizes.has_value()) { size_t numel_base = std::accumulate(outputShape.cbegin() + 1, outputShape.cend(), 1, std::multiplies<>{}); int64_t split_offset = 0; ncclGroupStart(); for (int root = 0; root < static_cast(mGroup.size()); ++root) { - auto split_size = all_rank_split_size.value()[root]; + auto split_size = sizes.value()[root]; NCCLCHECK(ncclBroadcast(input.data_ptr(), output.index({torch::indexing::Slice(split_offset, torch::indexing::None)}).mutable_data_ptr(), numel_base * split_size, (*getDtypeMap())[type], root, *mNcclComm, stream)); @@ -95,14 +91,14 @@ class AllgatherOp } std::vector run_list( - torch::TensorList input_list, torch::optional> all_rank_split_size) noexcept + torch::TensorList input_list, torch::optional> sizes) noexcept { std::vector output_list; output_list.reserve(input_list.size()); ncclGroupStart(); for (auto const& input : input_list) { - auto output = run(input, all_rank_split_size); + auto output = run(input, sizes); output_list.push_back(output); } ncclGroupEnd(); @@ -116,12 +112,8 @@ class AllgatherOp } // namespace -#endif // ENABLE_MULTI_DEVICE - -torch::Tensor allgather( - torch::Tensor input, torch::optional> all_rank_split_size, torch::List group_) +torch::Tensor allgather(torch::Tensor input, torch::optional> sizes, torch::List group_) { -#if ENABLE_MULTI_DEVICE std::set group; for (int64_t rank : group_) { @@ -129,17 +121,13 @@ torch::Tensor allgather( } AllgatherOp op(group); op.initialize(); - auto output = op.run(input, all_rank_split_size); + auto output = op.run(input, sizes); return output; -#else - return input; -#endif // ENABLE_MULTI_DEVICE } -std::vector allgather_list(torch::TensorList input_list, - torch::optional> all_rank_split_size, torch::List group_) +std::vector allgather_list( + torch::TensorList input_list, torch::optional> sizes, torch::List group_) { -#if ENABLE_MULTI_DEVICE std::set group; for (int64_t rank : group_) { @@ -147,19 +135,16 @@ std::vector allgather_list(torch::TensorList input_list, } AllgatherOp op(group); op.initialize(); - auto output_list = op.run_list(input_list, all_rank_split_size); + auto output_list = op.run_list(input_list, sizes); return output_list; -#else - return input_list.vec(); -#endif // ENABLE_MULTI_DEVICE } } // namespace torch_ext TORCH_LIBRARY_FRAGMENT(trtllm, m) { - m.def("allgather(Tensor input, int[]? all_rank_split_size, int[] group) -> Tensor"); - m.def("allgather_list(Tensor[] input_list, int[]? all_rank_split_size, int[] group) -> Tensor[]"); + m.def("allgather(Tensor input, int[]? sizes, int[] group) -> Tensor"); + m.def("allgather_list(Tensor[] input_list, int[]? sizes, int[] group) -> Tensor[]"); } TORCH_LIBRARY_IMPL(trtllm, CUDA, m) diff --git a/cpp/tensorrt_llm/thop/reducescatterOp.cpp b/cpp/tensorrt_llm/thop/reducescatterOp.cpp index 942e8f8482b3..36a354fdb28f 100644 --- a/cpp/tensorrt_llm/thop/reducescatterOp.cpp +++ b/cpp/tensorrt_llm/thop/reducescatterOp.cpp @@ -21,10 +21,8 @@ #include #include -#include -#if ENABLE_MULTI_DEVICE #include -#endif // ENABLE_MULTI_DEVICE +#include #include #include @@ -32,7 +30,6 @@ namespace torch_ext { -#if ENABLE_MULTI_DEVICE namespace { @@ -55,13 +52,13 @@ class ReducescatterOp return 0; } - torch::Tensor run(torch::Tensor const& input, torch::optional> all_rank_split_size) noexcept + torch::Tensor run(torch::Tensor const& input, torch::optional> sizes) noexcept { TLLM_CHECK_WITH_INFO(mNcclComm.get() != nullptr, "mNcclComm should be initialized before used"); auto stream = at::cuda::getCurrentCUDAStream(input.get_device()); auto type = tensorrt_llm::runtime::TorchUtils::dataType(input.scalar_type()); std::vector outputShape = input.sizes().vec(); - if (all_rank_split_size.has_value()) + if (sizes.has_value()) { auto rank = COMM_SESSION.getRank(); int groupRank = 0; @@ -72,21 +69,21 @@ class ReducescatterOp ++groupRank; } TLLM_CHECK(static_cast(groupRank) < mGroup.size()); - outputShape[0] = all_rank_split_size.value()[groupRank]; + outputShape[0] = sizes.value()[groupRank]; } else { outputShape[0] = outputShape[0] / mGroup.size(); } auto output = torch::empty(outputShape, input.options()); - if (all_rank_split_size.has_value()) + if (sizes.has_value()) { size_t numel_base = std::accumulate(outputShape.cbegin() + 1, outputShape.cend(), 1, std::multiplies<>{}); int64_t split_offset = 0; ncclGroupStart(); for (int root = 0; root < static_cast(mGroup.size()); ++root) { - auto split_size = all_rank_split_size.value()[root]; + auto split_size = sizes.value()[root]; NCCLCHECK( ncclReduce(input.index({torch::indexing::Slice(split_offset, torch::indexing::None)}).data_ptr(), output.mutable_data_ptr(), numel_base * split_size, (*getDtypeMap())[type], ncclSum, root, @@ -104,14 +101,14 @@ class ReducescatterOp } std::vector run_list( - torch::TensorList input_list, torch::optional> all_rank_split_size) noexcept + torch::TensorList input_list, torch::optional> sizes) noexcept { std::vector output_list; output_list.reserve(input_list.size()); ncclGroupStart(); for (auto const& input : input_list) { - auto output = run(input, all_rank_split_size); + auto output = run(input, sizes); output_list.push_back(output); } ncclGroupEnd(); @@ -125,12 +122,9 @@ class ReducescatterOp } // namespace -#endif // ENABLE_MULTI_DEVICE - extern torch::Tensor reducescatter( - torch::Tensor input, torch::optional> all_rank_split_size, torch::List group_) + torch::Tensor input, torch::optional> sizes, torch::List group_) { -#if ENABLE_MULTI_DEVICE std::set group; for (int64_t rank : group_) { @@ -138,17 +132,13 @@ extern torch::Tensor reducescatter( } ReducescatterOp op(group); op.initialize(); - auto output = op.run(input, all_rank_split_size); + auto output = op.run(input, sizes); return output; -#else - return input; -#endif // ENABLE_MULTI_DEVICE } -extern std::vector reducescatter_list(torch::TensorList input_list, - torch::optional> all_rank_split_size, torch::List group_) +extern std::vector reducescatter_list( + torch::TensorList input_list, torch::optional> sizes, torch::List group_) { -#if ENABLE_MULTI_DEVICE std::set group; for (int64_t rank : group_) { @@ -156,19 +146,16 @@ extern std::vector reducescatter_list(torch::TensorList input_lis } ReducescatterOp op(group); op.initialize(); - auto output_list = op.run_list(input_list, all_rank_split_size); + auto output_list = op.run_list(input_list, sizes); return output_list; -#else - return input_list.vec(); -#endif // ENABLE_MULTI_DEVICE } } // namespace torch_ext TORCH_LIBRARY_FRAGMENT(trtllm, m) { - m.def("reducescatter(Tensor input, int[]? all_rank_split_size, int[] group) -> Tensor"); - m.def("reducescatter_list(Tensor[] input_list, int[]? all_rank_split_size, int[] group) -> Tensor[]"); + m.def("reducescatter(Tensor input, int[]? sizes, int[] group) -> Tensor"); + m.def("reducescatter_list(Tensor[] input_list, int[]? sizes, int[] group) -> Tensor[]"); } TORCH_LIBRARY_IMPL(trtllm, CUDA, m) diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/dist.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/dist.py index d95ada5e7cf3..755638852e40 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/dist.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/dist.py @@ -10,13 +10,11 @@ @torch.library.custom_op("dist::all_gather", mutates_args=(), device_types="cuda") def all_gather( - tensor: torch.Tensor, dim: int = 0, all_rank_split_size: Optional[List[int]] = None + tensor: torch.Tensor, dim: int = 0, sizes: Optional[List[int]] = None ) -> torch.Tensor: """All gather followed by concat in dim = 0. This is the default nccl behavior.""" if trtllm_dist.is_trtllm_op_available(): - return trtllm_dist.trtllm_allgather( - tensor, dim=dim, all_rank_split_size=all_rank_split_size - ) + return trtllm_dist.trtllm_allgather(tensor, dim=dim, sizes=sizes) tl = [torch.zeros_like(tensor) for _ in range(dist.get_world_size())] dist.all_gather(tl, tensor) return torch.cat(tl, dim=dim) diff --git a/tensorrt_llm/_torch/auto_deploy/distributed/trtllm.py b/tensorrt_llm/_torch/auto_deploy/distributed/trtllm.py index ef0c1edcbadd..e0ac0db1b8ed 100644 --- a/tensorrt_llm/_torch/auto_deploy/distributed/trtllm.py +++ b/tensorrt_llm/_torch/auto_deploy/distributed/trtllm.py @@ -8,10 +8,10 @@ from ...distributed import AllReduce, allgather from ...modules.linear import AllReduceFusionOp, AllReduceParams - def trtllm_allgather(tensor, dim, all_rank_split_size=None): + def trtllm_allgather(tensor, dim, sizes=None): rank, world_size = get_rank_world_size() p_config = Mapping(world_size=world_size, tp_size=world_size, rank=rank) - return allgather(tensor, p_config, gather_dim=dim, all_rank_split_size=all_rank_split_size) + return allgather(tensor, p_config, dim=dim, sizes=sizes) def trtllm_allreduce(tensor, op, all_reduce_params=None): rank, world_size = get_rank_world_size() @@ -45,7 +45,7 @@ def fused_allreduce_residual_rmsnorm_fake( TRTLLM_OP_AVAILABLE = True except ImportError: - def trtllm_allgather(tensor, dim, all_rank_split_size=None): + def trtllm_allgather(tensor, dim, sizes=None): raise ImportError("TRT-LLM is not available.") def trtllm_allreduce(tensor, op): diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index 872446bd2c76..6f913da2a1f9 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -61,11 +61,11 @@ def _(residual, norm_weight, device_num_experts, scale_input, return [norm_out, residual_out] @torch.library.register_fake("trtllm::allgather") - def _(input, all_rank_split_size, group): - if all_rank_split_size is None: + def _(input, sizes, group): + if sizes is None: output_shape = (len(group) * input.shape[0], *input.shape[1:]) else: - output_shape = (sum(all_rank_split_size), *input.shape[1:]) + output_shape = (sum(sizes), *input.shape[1:]) return input.new_empty(output_shape) @torch.library.register_fake("trtllm::cublas_scaled_mm") diff --git a/tensorrt_llm/_torch/distributed/ops.py b/tensorrt_llm/_torch/distributed/ops.py index 5f49df8bdd0c..e977202e07ee 100644 --- a/tensorrt_llm/_torch/distributed/ops.py +++ b/tensorrt_llm/_torch/distributed/ops.py @@ -38,19 +38,19 @@ def userbuffers_allreduce_finalize( def allgather( input: Union[torch.Tensor, List[torch.Tensor]], mapping: Mapping, - gather_dim: int = -1, - all_rank_split_size: Optional[List[int]] = None, + dim: int = -1, + sizes: Optional[List[int]] = None, ) -> Union[torch.Tensor, List[torch.Tensor]]: ''' Add an operation that performs a collective all-gather. - If 'all_rank_split_size' is 'None', the input tensors in the different ranks must have the same shape. - Otherwise, 'all_rank_split_size[i]' must be 'input.shape[gather_dim]' at rank i, and the input tensors in - the different ranks can only differ in shape at dimension `gather_dim`. + If 'sizes' is 'None', the input tensors in the different ranks must have the same shape. + Otherwise, 'sizes[i]' must be 'input.shape[dim]' at rank i, and the input tensors in + the different ranks can only differ in shape at dimension `dim`. - The input tensors in the same TP group are concatenated at dimension 'gather_dim' to produce the output tensor. - If 'all_rank_split_size' is 'None', 'output.shape[gather_dim] = input.shape[gather_dim] * tp_group_size'. - Otherwise, 'output.shape[gather_dim] = sum(all_rank_split_size)'. + The input tensors in the same TP group are concatenated at dimension 'dim' to produce the output tensor. + If 'sizes' is 'None', 'output.shape[dim] = input.shape[dim] * tp_group_size'. + Otherwise, 'output.shape[dim] = sum(sizes)'. That operation is implemented using a torch op that wraps the NCCL all-gather collective operation or the NCCL group call of a series of NCCL broadcast collective operations. See the following materials for details. @@ -61,92 +61,88 @@ def allgather( Args: input (Union[Tensor, List[Tensor]]): The input tensor or tensor list. mapping (Mapping): The parallel mapping. - gather_dim (int): Gather along given dimension. By default -1. - all_rank_split_size(Optional[List[int]]): An optional list indicating 'input.shape[gather_dim]' in all ranks. By default None. + dim (int): Gather along given dimension. By default -1. + sizes(Optional[List[int]]): An optional list indicating 'input.shape[dim]' in all ranks. By default None. Returns: The gathered tensor or tensor list. ''' if mapping.tp_size == 1: return input - if all_rank_split_size is not None: - assert len(all_rank_split_size) == len(mapping.tp_group) + if sizes is not None: + assert len(sizes) == len(mapping.tp_group) if isinstance(input, torch.Tensor): - assert input.shape[gather_dim] == all_rank_split_size[ - mapping.tp_rank] + assert input.shape[dim] == sizes[mapping.tp_rank] else: - assert all([ - val.shape[gather_dim] == all_rank_split_size[mapping.tp_rank] - for val in input - ]) - # 'all_rank_split_size' is not needed if all inputs in the same TP group have the same shape - for split_size in all_rank_split_size[1:]: - if split_size != all_rank_split_size[0]: + assert all( + [val.shape[dim] == sizes[mapping.tp_rank] for val in input]) + # 'sizes' is not needed if all inputs in the same TP group have the same shape + for split_size in sizes[1:]: + if split_size != sizes[0]: break else: - all_rank_split_size = None + sizes = None if isinstance(input, torch.Tensor): torch_op = torch.ops.trtllm.allgather - input = input.movedim(gather_dim, 0).contiguous() + input = input.movedim(dim, 0).contiguous() else: torch_op = torch.ops.trtllm.allgather_list - input = [val.movedim(gather_dim, 0).contiguous() for val in input] + input = [val.movedim(dim, 0).contiguous() for val in input] output = torch_op( input, - all_rank_split_size, + sizes, mapping.tp_group, ) if isinstance(input, torch.Tensor): - output = output.movedim(0, gather_dim).contiguous() + output = output.movedim(0, dim).contiguous() else: - output = [val.movedim(0, gather_dim).contiguous() for val in output] + output = [val.movedim(0, dim).contiguous() for val in output] return output def reducescatter( input: Union[torch.Tensor, List[torch.Tensor]], mapping: Mapping, - scatter_dim: int = -1, - all_rank_split_size: Optional[List[int]] = None, + dim: int = -1, + sizes: Optional[List[int]] = None, ) -> Union[torch.Tensor, List[torch.Tensor]]: if mapping.tp_size == 1: return input - if all_rank_split_size is not None: - assert len(all_rank_split_size) == len(mapping.tp_group) - sum_split_size = sum(all_rank_split_size) + if sizes is not None: + assert len(sizes) == len(mapping.tp_group) + sum_split_size = sum(sizes) if isinstance(input, torch.Tensor): - assert input.shape[scatter_dim] == sum_split_size + assert input.shape[dim] == sum_split_size else: - assert all( - [val.shape[scatter_dim] == sum_split_size for val in input]) - # 'all_rank_split_size' is not needed if all outputs in the same TP group have the same shape - for split_size in all_rank_split_size[1:]: - if split_size != all_rank_split_size[0]: + assert all([val.shape[dim] == sum_split_size for val in input]) + # 'sizes' is not needed if all outputs in the same TP group have the same shape + for split_size in sizes[1:]: + if split_size != sizes[0]: break else: - all_rank_split_size = None + sizes = None if isinstance(input, torch.Tensor): torch_op = torch.ops.trtllm.reducescatter - input = input.movedim(scatter_dim, 0).contiguous() + input = input.movedim(dim, 0).contiguous() else: torch_op = torch.ops.trtllm.reducescatter_list - input = [val.movedim(scatter_dim, 0).contiguous() for val in input] + input = [val.movedim(dim, 0).contiguous() for val in input] output = torch_op( input, - all_rank_split_size, + sizes, mapping.tp_group, ) if isinstance(input, torch.Tensor): - output = output.movedim(0, scatter_dim).contiguous() + output = output.movedim(0, dim).contiguous() else: - output = [val.movedim(0, scatter_dim).contiguous() for val in output] + output = [val.movedim(0, dim).contiguous() for val in output] return output diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv3.py b/tensorrt_llm/_torch/models/modeling_deepseekv3.py index b38fc6ee1beb..687b5ab5ad01 100644 --- a/tensorrt_llm/_torch/models/modeling_deepseekv3.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv3.py @@ -462,11 +462,10 @@ def compute_routed_output(self, hidden_states, hidden_states_fp4, # FP4 all_gather moves this bf16 allgather in to after topk and fp4 quantization # to reduce allreduce BW if disable_fp4_allgather() and not self.enable_alltoall: - hidden_states = allgather( - hidden_states, - self.mapping, - gather_dim=0, - all_rank_split_size=all_rank_num_tokens) + hidden_states = allgather(hidden_states, + self.mapping, + dim=0, + sizes=all_rank_num_tokens) elif not self.experts.is_cutlass() or (not self.experts.has_fp8_qdq and self.experts.has_nvfp4): # Use padding when not using the cutlass path or when x_sf in self.experts is not None diff --git a/tensorrt_llm/_torch/models/modeling_mixtral.py b/tensorrt_llm/_torch/models/modeling_mixtral.py index 39ea39c53e2f..1e02c4fc5f9b 100644 --- a/tensorrt_llm/_torch/models/modeling_mixtral.py +++ b/tensorrt_llm/_torch/models/modeling_mixtral.py @@ -68,11 +68,10 @@ def forward( # FP4 all_gather moves this bf16 allgather in to after topk and fp4 quantization # to reduce allreduce BW if disable_fp4_allgather(): - hidden_states = allgather( - hidden_states, - self.mapping, - gather_dim=0, - all_rank_split_size=all_rank_num_tokens) + hidden_states = allgather(hidden_states, + self.mapping, + dim=0, + sizes=all_rank_num_tokens) elif not self.experts.is_cutlass() or (not self.experts.has_fp8_qdq and self.experts.has_nvfp4): # Use padding when not using the cutlass path or when x_sf in self.experts is not None diff --git a/tensorrt_llm/_torch/modules/fused_moe.py b/tensorrt_llm/_torch/modules/fused_moe.py index 38d46a68d2b2..2b9921796c90 100755 --- a/tensorrt_llm/_torch/modules/fused_moe.py +++ b/tensorrt_llm/_torch/modules/fused_moe.py @@ -733,11 +733,11 @@ def reducescatter_or_allreduce( outputs = inputs if self.parallel_size > 1 and not self.enable_alltoall: if self.use_dp: - outputs = reducescatter(inputs, - self.mapping, - scatter_dim=0, - all_rank_split_size=None if - use_dp_padding else all_rank_num_tokens) + outputs = reducescatter( + inputs, + self.mapping, + dim=0, + sizes=None if use_dp_padding else all_rank_num_tokens) elif self.reduce_results: outputs = self.all_reduce(inputs) return outputs @@ -822,17 +822,15 @@ def forward_chunk( x, token_selected_experts, token_final_scales = allgather( [x, token_selected_experts, token_final_scales], self.mapping, - gather_dim=0, - all_rank_split_size=None - if use_dp_padding else all_rank_num_tokens) + dim=0, + sizes=None if use_dp_padding else all_rank_num_tokens) else: # Fp4 gemm has extra scaling factor x, x_sf, token_selected_experts, token_final_scales = allgather( [x, x_sf, token_selected_experts, token_final_scales], self.mapping, - gather_dim=0, - all_rank_split_size=None - if use_dp_padding else all_rank_num_tokens) + dim=0, + sizes=None if use_dp_padding else all_rank_num_tokens) x_sf = reswizzle_sf(x_sf, x_row, x_col, self.scaling_vector_size) @@ -1145,9 +1143,7 @@ def alltoall_prepare_maybe_dispatch(self, all_rank_num_tokens: list, token_final_scales, (0, 0, 0, max_num_token - token_final_scales.shape[0])) gathered_token_selected_experts, gathered_token_final_scales = allgather( - [token_selected_experts, token_final_scales], - self.mapping, - gather_dim=0) + [token_selected_experts, token_final_scales], self.mapping, dim=0) gathered_token_selected_experts = torch.flatten( gathered_token_selected_experts.contiguous(), start_dim=0, diff --git a/tests/integration/test_lists/test-db/l0_dgx_h100.yml b/tests/integration/test_lists/test-db/l0_dgx_h100.yml index 3e7bb65af26c..f8ba95fb838d 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_h100.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_h100.yml @@ -51,6 +51,7 @@ l0_dgx_h100: backend: pytorch auto_trigger: deepseek tests: + - unittest/_torch/multi_gpu_modeling -k "deepseek" - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_scales_4gpus[tp4-mtp_nextn=0-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_scales_4gpus[tp4-mtp_nextn=0-fp8kv=True-attention_dp=False-cuda_graph=False-overlap_scheduler=False] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_scales_4gpus[tp4-mtp_nextn=0-fp8kv=False-attention_dp=True-cuda_graph=False-overlap_scheduler=False] diff --git a/tests/integration/test_lists/test-db/l0_dgx_h200.yml b/tests/integration/test_lists/test-db/l0_dgx_h200.yml index cc83f48679f2..f291305fa767 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_h200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_h200.yml @@ -18,4 +18,3 @@ l0_dgx_h200: # - accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_fp8_blockscale[throughput] # OOM - accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_fp8_blockscale[latency] # 1h - unittest/_torch/multi_gpu_modeling/test_llama4.py::test_llama4[pp1-ep1-enable_graph-tp8-trtllm-scout] - - unittest/_torch/multi_gpu_modeling -k "deepseek" diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index e615806f7319..83a655429094 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -17,7 +17,6 @@ l0_h100: # Only key models in H100: llama/mixtral/nemotron/deepseek - unittest/_torch -k "not (modeling or multi_gpu or auto_deploy)" - unittest/_torch -k "modeling_llama" - - unittest/_torch/multi_gpu_modeling -k "deepseek" - unittest/_torch/modeling -k "modeling_mixtral" - unittest/_torch/modeling -k "modeling_nemotron" - accuracy/test_llm_api_pytorch.py::TestGemma3_1BInstruct::test_auto_dtype diff --git a/tests/unittest/_torch/multi_gpu_modeling/test_deepseek.py b/tests/unittest/_torch/multi_gpu_modeling/test_deepseek.py index f4baeb32e7f8..6f0248b86dae 100644 --- a/tests/unittest/_torch/multi_gpu_modeling/test_deepseek.py +++ b/tests/unittest/_torch/multi_gpu_modeling/test_deepseek.py @@ -24,12 +24,7 @@ def similar(a, b, threshold=0.9): @pytest.mark.parametrize("backend", ["TRTLLM"], ids=["trtllm"]) @pytest.mark.parametrize("quant", ["bf16"]) @pytest.mark.parametrize("tp_size", [1, 4], ids=["tp1", "tp4"]) -@pytest.mark.parametrize("enable_attention_dp", [False, True], - ids=["adp_off", "adp_on"]) -@pytest.mark.parametrize("moe_max_num_tokens", [None, 64], - ids=["moe_chunk_off", "moe_chunk_on"]) -def test_deepseek_streaming(model_name, backend, quant, tp_size, - enable_attention_dp, moe_max_num_tokens): +def test_deepseek_streaming(model_name, backend, quant, tp_size): model_path = { "bf16": "bf16", "fp8": "fp8", @@ -53,6 +48,13 @@ def test_deepseek_streaming(model_name, backend, quant, tp_size, if get_total_gpu_memory(0) < 60 * 1024**3: pytest.skip(f"Not enough GPU memory to run. {get_total_gpu_memory(0)}") + if tp_size == 1: + enable_attention_dp = False + moe_max_num_tokens = None + else: + enable_attention_dp = True + moe_max_num_tokens = 64 + prompts = [ "The president of the United States is", ] * 32 From 51c0fa355af56b0fee36c525ee660a42517fe7d5 Mon Sep 17 00:00:00 2001 From: Jinyang Yuan <154768711+jinyangyuan-nvidia@users.noreply.github.com> Date: Tue, 13 May 2025 20:22:09 -0700 Subject: [PATCH 3/4] Revert changes to pass CI Signed-off-by: Jinyang Yuan <154768711+jinyangyuan-nvidia@users.noreply.github.com> --- cpp/tensorrt_llm/thop/allgatherOp.cpp | 15 ++++++++++++++- cpp/tensorrt_llm/thop/reducescatterOp.cpp | 15 ++++++++++++++- 2 files changed, 28 insertions(+), 2 deletions(-) diff --git a/cpp/tensorrt_llm/thop/allgatherOp.cpp b/cpp/tensorrt_llm/thop/allgatherOp.cpp index c2d9a5de4e09..6475527a7992 100644 --- a/cpp/tensorrt_llm/thop/allgatherOp.cpp +++ b/cpp/tensorrt_llm/thop/allgatherOp.cpp @@ -22,14 +22,17 @@ #include #include #include -#include #include #include #include #include +#if ENABLE_MULTI_DEVICE +#include +#endif // ENABLE_MULTI_DEVICE namespace torch_ext { +#if ENABLE_MULTI_DEVICE namespace { @@ -112,8 +115,11 @@ class AllgatherOp } // namespace +#endif // ENABLE_MULTI_DEVICE + torch::Tensor allgather(torch::Tensor input, torch::optional> sizes, torch::List group_) { +#if ENABLE_MULTI_DEVICE std::set group; for (int64_t rank : group_) { @@ -123,11 +129,15 @@ torch::Tensor allgather(torch::Tensor input, torch::optional allgather_list( torch::TensorList input_list, torch::optional> sizes, torch::List group_) { +#if ENABLE_MULTI_DEVICE std::set group; for (int64_t rank : group_) { @@ -137,6 +147,9 @@ std::vector allgather_list( op.initialize(); auto output_list = op.run_list(input_list, sizes); return output_list; +#else + return input_list.vec(); +#endif // ENABLE_MULTI_DEVICE } } // namespace torch_ext diff --git a/cpp/tensorrt_llm/thop/reducescatterOp.cpp b/cpp/tensorrt_llm/thop/reducescatterOp.cpp index 36a354fdb28f..535dab5751f1 100644 --- a/cpp/tensorrt_llm/thop/reducescatterOp.cpp +++ b/cpp/tensorrt_llm/thop/reducescatterOp.cpp @@ -21,8 +21,10 @@ #include #include -#include #include +#if ENABLE_MULTI_DEVICE +#include +#endif // ENABLE_MULTI_DEVICE #include #include @@ -30,6 +32,7 @@ namespace torch_ext { +#if ENABLE_MULTI_DEVICE namespace { @@ -122,9 +125,12 @@ class ReducescatterOp } // namespace +#endif // ENABLE_MULTI_DEVICE + extern torch::Tensor reducescatter( torch::Tensor input, torch::optional> sizes, torch::List group_) { +#if ENABLE_MULTI_DEVICE std::set group; for (int64_t rank : group_) { @@ -134,11 +140,15 @@ extern torch::Tensor reducescatter( op.initialize(); auto output = op.run(input, sizes); return output; +#else + return input; +#endif // ENABLE_MULTI_DEVICE } extern std::vector reducescatter_list( torch::TensorList input_list, torch::optional> sizes, torch::List group_) { +#if ENABLE_MULTI_DEVICE std::set group; for (int64_t rank : group_) { @@ -148,6 +158,9 @@ extern std::vector reducescatter_list( op.initialize(); auto output_list = op.run_list(input_list, sizes); return output_list; +#else + return input_list.vec(); +#endif // ENABLE_MULTI_DEVICE } } // namespace torch_ext From 8b225a038f2a3e6069ee55c65bc44fc3093c7f35 Mon Sep 17 00:00:00 2001 From: Jinyang Yuan <154768711+jinyangyuan-nvidia@users.noreply.github.com> Date: Thu, 15 May 2025 09:10:37 -0700 Subject: [PATCH 4/4] Fix torch.compile accuracy issue brought by PyTorch update and fix some MoE models Signed-off-by: Jinyang Yuan <154768711+jinyangyuan-nvidia@users.noreply.github.com> --- tensorrt_llm/_torch/distributed/ops.py | 67 ++++++++++++++++--- tensorrt_llm/_torch/models/modeling_llama.py | 6 +- .../_torch/models/modeling_mixtral.py | 29 +++----- .../_torch/models/modeling_qwen3_moe.py | 8 +-- .../_torch/models/modeling_qwen_moe.py | 12 ++-- tensorrt_llm/_torch/modules/fused_moe.py | 12 +++- 6 files changed, 88 insertions(+), 46 deletions(-) diff --git a/tensorrt_llm/_torch/distributed/ops.py b/tensorrt_llm/_torch/distributed/ops.py index e977202e07ee..dd62198e88eb 100644 --- a/tensorrt_llm/_torch/distributed/ops.py +++ b/tensorrt_llm/_torch/distributed/ops.py @@ -1,3 +1,4 @@ +import math import threading from typing import List, Optional, Tuple, Union @@ -35,6 +36,15 @@ def userbuffers_allreduce_finalize( return output +def get_output_info(input: torch.Tensor, dim: int) -> List[int]: + dim = dim % input.ndim + output_shape = [ + val if idx != dim else -1 for idx, val in enumerate(input.shape) + ] + numel_base = -math.prod(output_shape) + return {'output_shape': output_shape, 'numel_base': numel_base} + + def allgather( input: Union[torch.Tensor, List[torch.Tensor]], mapping: Mapping, @@ -83,12 +93,18 @@ def allgather( else: sizes = None + # Inputs are reshaped in this way to pass necessary shape information to the allgather op if isinstance(input, torch.Tensor): torch_op = torch.ops.trtllm.allgather - input = input.movedim(dim, 0).contiguous() + output_info = get_output_info(input, dim) + input = input.contiguous().view(-1, output_info['numel_base']) else: torch_op = torch.ops.trtllm.allgather_list - input = [val.movedim(dim, 0).contiguous() for val in input] + output_info = [get_output_info(val, dim) for val in input] + input = [ + val.contiguous().view(-1, val_info['numel_base']) + for val, val_info in zip(input, output_info) + ] output = torch_op( input, @@ -96,10 +112,25 @@ def allgather( mapping.tp_group, ) + def convert_output(x, x_info): + if dim == 0: + x = x.view(x_info['output_shape']) + else: + if sizes is None: + x_list = x.chunk(mapping.tp_size) + else: + x_list = x.split(sizes) + x = torch.cat([x.reshape(x_info['output_shape']) for x in x_list], + dim=dim) + return x + if isinstance(input, torch.Tensor): - output = output.movedim(0, dim).contiguous() + output = convert_output(output, output_info) else: - output = [val.movedim(0, dim).contiguous() for val in output] + output = [ + convert_output(val, val_info) + for val, val_info in zip(output, output_info) + ] return output @@ -126,12 +157,29 @@ def reducescatter( else: sizes = None + def convert_input(x, x_info): + # Inputs are reshaped in this way to pass necessary shape information to the reducescatter op + if dim == 0: + x = x.contiguous().view(-1, x_info['numel_base']) + else: + if sizes is None: + x_list = x.chunk(mapping.tp_size, dim=dim) + else: + x_list = x.split(sizes, dim=dim) + x = torch.cat([x.reshape(-1, x_info['numel_base']) for x in x_list]) + return x + if isinstance(input, torch.Tensor): torch_op = torch.ops.trtllm.reducescatter - input = input.movedim(dim, 0).contiguous() + output_info = get_output_info(input, dim) + input = convert_input(input, output_info) else: torch_op = torch.ops.trtllm.reducescatter_list - input = [val.movedim(dim, 0).contiguous() for val in input] + output_info = [get_output_info(val, dim) for val in input] + input = [ + convert_input(val, val_info) + for val, val_info in zip(input, output_info) + ] output = torch_op( input, @@ -140,9 +188,12 @@ def reducescatter( ) if isinstance(input, torch.Tensor): - output = output.movedim(0, dim).contiguous() + output = output.view(output_info['output_shape']) else: - output = [val.movedim(0, dim).contiguous() for val in output] + output = [ + val.view(val_info['output_shape']) + for val, val_info in zip(output, output_info) + ] return output diff --git a/tensorrt_llm/_torch/models/modeling_llama.py b/tensorrt_llm/_torch/models/modeling_llama.py index 1bb08f1b9a8b..ff2896d5f3fb 100644 --- a/tensorrt_llm/_torch/models/modeling_llama.py +++ b/tensorrt_llm/_torch/models/modeling_llama.py @@ -257,7 +257,10 @@ def __init__( def compute_routed_output(self, hidden_states, all_rank_num_tokens, cutlass_min_latency_mode): + use_dp_padding = False if self.enable_attention_dp and self.mapping.tp_size > 1: + # Use padding here to keep the behavior unchanged + use_dp_padding = True max_num_token_across_dp_ranks = max(all_rank_num_tokens) hidden_states = torch.nn.functional.pad( hidden_states, @@ -267,7 +270,8 @@ def compute_routed_output(self, hidden_states, all_rank_num_tokens, routed_output = self.experts(hidden_states, router_logits, cutlass_min_latency_mode, - all_rank_num_tokens=all_rank_num_tokens) + all_rank_num_tokens=all_rank_num_tokens, + use_dp_padding=use_dp_padding) return routed_output def forward( diff --git a/tensorrt_llm/_torch/models/modeling_mixtral.py b/tensorrt_llm/_torch/models/modeling_mixtral.py index 1e02c4fc5f9b..7363eaf82667 100644 --- a/tensorrt_llm/_torch/models/modeling_mixtral.py +++ b/tensorrt_llm/_torch/models/modeling_mixtral.py @@ -8,7 +8,6 @@ from ..attention_backend import AttentionMetadata from ..attention_backend.interface import PositionalEmbeddingParams, RopeParams -from ..distributed import allgather from ..models.modeling_utils import ModelConfig from ..modules.attention import Attention from ..modules.decoder_layer import DecoderLayer @@ -16,7 +15,6 @@ from ..modules.fused_moe import FusedMoE, RenormalizeMoeRoutingMethod from ..modules.linear import Linear from ..modules.rms_norm import RMSNorm -from ..utils import disable_fp4_allgather from .modeling_utils import (DecoderModel, DecoderModelForCausalLM, register_auto_model) @@ -34,7 +32,7 @@ def __init__( self.ffn_dim = config.intermediate_size self.num_experts = config.num_local_experts self.top_k = config.num_experts_per_tok - self.use_dp = model_config.mapping.enable_attention_dp + self.enable_attention_dp = model_config.mapping.enable_attention_dp # moe gate (linear layer) only runs in half/full precision for now self.gate = Linear(self.hidden_dim, @@ -55,8 +53,6 @@ def __init__( reduce_results=reduce_results, model_config=model_config) - self.mapping = model_config.mapping - def forward( self, hidden_states: torch.Tensor, @@ -64,22 +60,13 @@ def forward( ) -> torch.Tensor: all_rank_num_tokens = attn_metadata.all_rank_num_tokens use_dp_padding = False - if self.use_dp and self.mapping.tp_size > 1: - # FP4 all_gather moves this bf16 allgather in to after topk and fp4 quantization - # to reduce allreduce BW - if disable_fp4_allgather(): - hidden_states = allgather(hidden_states, - self.mapping, - dim=0, - sizes=all_rank_num_tokens) - elif not self.experts.is_cutlass() or (not self.experts.has_fp8_qdq - and self.experts.has_nvfp4): - # Use padding when not using the cutlass path or when x_sf in self.experts is not None - use_dp_padding = True - max_num_token = max(all_rank_num_tokens) - hidden_states = torch.nn.functional.pad( - hidden_states, - (0, 0, 0, max_num_token - hidden_states.shape[0])) + if self.enable_attention_dp and len(all_rank_num_tokens) > 1: + # Use padding here to keep the behavior unchanged + use_dp_padding = True + max_num_token = max(all_rank_num_tokens) + hidden_states = torch.nn.functional.pad( + hidden_states, + (0, 0, 0, max_num_token - hidden_states.shape[0])) router_logits = self.gate(hidden_states) final_hidden_states = self.experts( hidden_states, diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_moe.py b/tensorrt_llm/_torch/models/modeling_qwen3_moe.py index 9713b948a976..b0c24936688d 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3_moe.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3_moe.py @@ -68,16 +68,12 @@ def forward( hidden_states = hidden_states.view(-1, self.hidden_dim) all_rank_num_tokens = attn_metadata.all_rank_num_tokens - if self.enable_attention_dp and len(all_rank_num_tokens) > 1: - max_num_token = max(all_rank_num_tokens) - hidden_states = torch.nn.functional.pad( - hidden_states, - (0, 0, 0, max_num_token - hidden_states.shape[0])) router_logits = self.gate(hidden_states) final_hidden_states = self.experts( hidden_states, router_logits, - all_rank_num_tokens=all_rank_num_tokens) + all_rank_num_tokens=all_rank_num_tokens, + use_dp_padding=False) if not self.enable_attention_dp and self.mapping.tp_size > 1: final_hidden_states = self.allreduce( diff --git a/tensorrt_llm/_torch/models/modeling_qwen_moe.py b/tensorrt_llm/_torch/models/modeling_qwen_moe.py index 6286f6b1550d..99212d811aad 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen_moe.py +++ b/tensorrt_llm/_torch/models/modeling_qwen_moe.py @@ -82,14 +82,12 @@ def forward( hidden_states = hidden_states.view(-1, self.hidden_dim) all_rank_num_tokens = attn_metadata.all_rank_num_tokens - if self.enable_attention_dp and len(all_rank_num_tokens) > 1: - max_num_token = max(all_rank_num_tokens) - hidden_states = torch.nn.functional.pad( - hidden_states, - (0, 0, 0, max_num_token - hidden_states.shape[0])) router_logits = self.gate(hidden_states) - final_hidden_states = self.experts(hidden_states, router_logits, - all_rank_num_tokens) + final_hidden_states = self.experts( + hidden_states, + router_logits, + all_rank_num_tokens=all_rank_num_tokens, + use_dp_padding=False) shared_expert_output = self.shared_expert(hidden_states) shared_expert_output = F.sigmoid( diff --git a/tensorrt_llm/_torch/modules/fused_moe.py b/tensorrt_llm/_torch/modules/fused_moe.py index 2b9921796c90..902707bbe1ca 100755 --- a/tensorrt_llm/_torch/modules/fused_moe.py +++ b/tensorrt_llm/_torch/modules/fused_moe.py @@ -944,17 +944,22 @@ def forward_cutlass( not self.reduce_results ), "cutlass_min_latency_mode must be used with a single chunk and reduce_results must be False" + if use_dp_padding: + all_rank_num_tokens_padded = [max(all_rank_num_tokens) + ] * len(all_rank_num_tokens) + else: + all_rank_num_tokens_padded = all_rank_num_tokens if num_chunks == 1: outputs = self.forward_chunk( x, router_logits, cutlass_min_latency_mode, output_dtype, - all_rank_num_tokens=all_rank_num_tokens, + all_rank_num_tokens=all_rank_num_tokens_padded, use_dp_padding=use_dp_padding) outputs = self.reducescatter_or_allreduce( outputs, - all_rank_num_tokens=all_rank_num_tokens, + all_rank_num_tokens=all_rank_num_tokens_padded, use_dp_padding=use_dp_padding) else: @@ -967,7 +972,8 @@ def split_chunk(split_token_num: int, split_num_chunks: int): if self.use_dp: all_rank_chunk_size_list = [ - split_chunk(val, num_chunks) for val in all_rank_num_tokens + split_chunk(val, num_chunks) + for val in all_rank_num_tokens_padded ] all_rank_num_tokens_list = [[ val[idx_chunk] for val in all_rank_chunk_size_list