diff --git a/aiter/dist/communication_op.py b/aiter/dist/communication_op.py index 258ccc8175a..f88716c8278 100644 --- a/aiter/dist/communication_op.py +++ b/aiter/dist/communication_op.py @@ -30,6 +30,12 @@ def tensor_model_parallel_all_reduce( return get_tp_group().all_reduce(input_, open_fp8_quant) +def tensor_model_parallel_fused_allreduce_rmsnorm( + input_: torch.Tensor, weight_: torch.Tensor, eps: float +) -> torch.Tensor: + return get_tp_group().fused_allreduce_rmsnorm(input_, weight_, eps) + + def tensor_model_parallel_custom_all_gather(input_: torch.Tensor) -> torch.Tensor: return get_tp_group().custom_all_gather(input_) diff --git a/aiter/dist/device_communicators/communicator_cuda.py b/aiter/dist/device_communicators/communicator_cuda.py index 97c3eefb712..339337cc657 100644 --- a/aiter/dist/device_communicators/communicator_cuda.py +++ b/aiter/dist/device_communicators/communicator_cuda.py @@ -158,6 +158,38 @@ def all_reduce(self, input_, ca_fp8_quant: bool = False) -> torch.Tensor: torch.distributed.all_reduce(out, group=self.device_group) return out + def fused_allreduce_rmsnorm(self, input_, weight_, eps) -> torch.Tensor: + n = input_.shape[-1] + can_use_fuse_ar_rms = ( + n <= 16384 and input_.numel() * input_.element_size() < 8 * 1024 * 8192 + ) + ca_comm = self.ca_comm + if ( + ca_comm is not None + and not ca_comm.disabled + and ca_comm.should_custom_ar(input_) + and can_use_fuse_ar_rms + ): + out = ca_comm.custom_fused_ar_rms(input_, weight_, eps) + assert out is not None + return out + # call split kernel + ar_out = all_reduce(input_) + out = torch.empty_like(ar_out) + residual_out = torch.empty_like(ar_out) + from aiter import rmsnorm2d_fwd_with_add + + rmsnorm2d_fwd_with_add( + out, + ar_out, + input_, + residual_out, + weight_, + eps, + 0, + ) + return out + def reduce_scatter(self, input_: torch.Tensor, dim: int = -1): world_size = self.world_size pynccl_comm = self.pynccl_comm diff --git a/aiter/dist/device_communicators/custom_all_reduce.py b/aiter/dist/device_communicators/custom_all_reduce.py index 12e6ee9a563..7a788b90464 100644 --- a/aiter/dist/device_communicators/custom_all_reduce.py +++ b/aiter/dist/device_communicators/custom_all_reduce.py @@ -336,6 +336,41 @@ def custom_all_gather(self, inp: torch.Tensor) -> Optional[torch.Tensor]: else: return self.all_gather_unreg(inp) + def fused_ar_rms( + self, + inp: torch.Tensor, + *, + out: Optional[torch.Tensor] = None, + w: torch.Tensor, + eps: float, + registered: bool = False, + ): + if out is None: + out = torch.empty_like(inp) + ops.fused_allreduce_rmsnorm( + self._ptr, + inp, + out, + w, + eps, + None if registered else self.buffer, + ) + return out + + def custom_fused_ar_rms( + self, input: torch.Tensor, weight: torch.Tensor, eps: float + ) -> Optional[torch.Tensor]: + # when custom allreduce is disabled, this will be None + if self.disabled or not self.should_custom_ar(input): + return None + if self._IS_CAPTURING: + if torch.cuda.is_current_stream_capturing(): + return self.fused_ar_rms(input, w=weight, eps=eps, registered=True) + else: + return torch.empty_like(input) + else: + return self.fused_ar_rms(input, w=weight, eps=eps, registered=False) + def close(self): if not self.disabled and self._ptr: ops.dispose(self._ptr) diff --git a/aiter/dist/parallel_state.py b/aiter/dist/parallel_state.py index 0f336ba1f1d..ec438224c08 100644 --- a/aiter/dist/parallel_state.py +++ b/aiter/dist/parallel_state.py @@ -119,6 +119,23 @@ def all_reduce_( return group._all_reduce_out_place(tensor, ca_fp8_quant) +def fused_allreduce_rmsnorm_fake( + inp: torch.Tensor, w: torch.Tensor, eps: float, group_name: str +) -> torch.Tensor: + return torch.empty_like(inp) + + +@torch_compile_guard(gen_fake=fused_allreduce_rmsnorm_fake) +def fused_allreduce_rmsnorm_( + inp: torch.Tensor, w: torch.Tensor, eps: float, group_name: str +) -> torch.Tensor: + assert group_name in _groups, f"Group {group_name} is not found." + group = _groups[group_name]() + if group is None: + raise ValueError(f"Group {group_name} is destroyed.") + return group._fused_allreduce_rmsnorm_out_place(inp, w, eps) + + if supports_custom_op(): # @torch.library.custom_op("aiter::outplace_all_gather", mutates_args=[]) @@ -329,6 +346,20 @@ def _all_reduce_out_place( raise ValueError("No device communicator found") return self.device_communicator.all_reduce(input_, ca_fp8_quant) + def fused_allreduce_rmsnorm( + self, input_: torch.Tensor, weight_: torch.Tensor, eps: float + ) -> torch.Tensor: + return fused_allreduce_rmsnorm_( + input_, weight_, eps, group_name=self.unique_name + ) + + def _fused_allreduce_rmsnorm_out_place( + self, input_: torch.Tensor, weight_: torch.Tensor, eps: float + ) -> torch.Tensor: + if self.device_communicator is None: + raise ValueError("No device communicator found") + return self.device_communicator.fused_allreduce_rmsnorm(input_, weight_, eps) + def _all_gather_out_place(self, input_: torch.Tensor) -> torch.Tensor: ca_comm = self.device_communicator.ca_comm assert ca_comm is not None diff --git a/aiter/ops/custom_all_reduce.py b/aiter/ops/custom_all_reduce.py index 53bd1d46da2..be62dae8d34 100644 --- a/aiter/ops/custom_all_reduce.py +++ b/aiter/ops/custom_all_reduce.py @@ -41,6 +41,17 @@ def all_gather_unreg( ) -> None: ... +@compile_ops("module_custom_all_reduce") +def fused_allreduce_rmsnorm( + _fa: int, + inp: torch.Tensor, + out: torch.Tensor, + w: torch.Tensor, + eps: float, + reg_buffer: Optional[torch.Tensor] = None, +) -> None: ... + + def all_reduce_asm_fake_tensor( inp: torch.Tensor, ca: int, diff --git a/csrc/include/custom_all_reduce.cuh b/csrc/include/custom_all_reduce.cuh index 86388d775a3..a03ebe99f0d 100644 --- a/csrc/include/custom_all_reduce.cuh +++ b/csrc/include/custom_all_reduce.cuh @@ -873,6 +873,293 @@ namespace aiter } } + // fused allreduce rmsnorm first step + template + __global__ void __launch_bounds__(512, 1) reduce_scatter_cross_device_store( + RankData* _dp, + RankSignals sg, + Signal* self_sg, + int rank, + int size + ) + { + constexpr int pack_size = packed_t::P::size; + constexpr int tnum_gpu = THREAD_NUM / ngpus; + using P = typename packed_t::P; + using A = typename packed_t::A; + __shared__ T tmp_smem[tnum_gpu * ngpus * pack_size]; + int warp_id = threadIdx.x / tnum_gpu; + int lane_id = threadIdx.x % tnum_gpu; + const P* ptrs[ngpus]; + P* tmps[ngpus]; +#pragma unroll + for (int i = 0; i < ngpus; ++i) + { + ptrs[i] = (const P*)_dp->ptrs[i]; + tmps[i] = get_tmp_buf

(sg.signals[i]); + } + start_sync(sg, self_sg, rank); + + // the case of fused_allreduce_rmsnorm does not need thread level boundary check + int part = size / (pack_size * tnum_gpu) / ngpus; + for (int bid = blockIdx.x; bid < part; bid += gridDim.x) + { + // cross device read by all warp + P input_reg = ptrs[warp_id][(rank * part + bid) * tnum_gpu + lane_id]; + *(reinterpret_cast(&tmp_smem[0]) + threadIdx.x) = input_reg; + __syncthreads(); + // calculate and save in first warp + if (warp_id == 0) + { + A add_reg; +#pragma unroll + for (int i = 0; i < pack_size; ++i) + { + add_reg.data[i] = ck_tile::type_convert(tmp_smem[pack_size * threadIdx.x + i]); + } +#pragma unroll + for (int i = 1; i < ngpus; ++i) + { +#pragma unroll + for (int j = 0; j < pack_size; ++j) + { + add_reg.data[j] += ck_tile::type_convert(tmp_smem[i * pack_size * tnum_gpu + pack_size * threadIdx.x + j]); + } + } + *(reinterpret_cast(&tmp_smem[0]) + lane_id) = add_reg; + } + __syncthreads(); + + // cross device store + P rslt; +#pragma unroll + for (int i = 0; i < pack_size; ++i) + { + float res_x = ck_tile::type_convert(input_reg.data[i]); + float sum_x = *(reinterpret_cast(&tmp_smem[0]) + lane_id * pack_size + i); + rslt.data[i] = ck_tile::type_convert(res_x + sum_x); + } + tmps[warp_id][(rank * part + bid) * tnum_gpu + lane_id] = rslt; + } + } + + template + DINLINE void smemReduceSum(float* smem_addr) + { + // a warp executes the same instruction +#pragma unroll + for (int stride = reduce_range / 2; stride > 32; stride >>= 1) + { + if (threadIdx.x < stride) + { + smem_addr[threadIdx.x] += smem_addr[threadIdx.x + stride]; + } + __syncthreads(); + } + volatile float* v_smem = &smem_addr[0]; + if (threadIdx.x < 32) + { + v_smem[threadIdx.x] += v_smem[threadIdx.x + 32]; + v_smem[threadIdx.x] += v_smem[threadIdx.x + 16]; + v_smem[threadIdx.x] += v_smem[threadIdx.x + 8]; + v_smem[threadIdx.x] += v_smem[threadIdx.x + 4]; + v_smem[threadIdx.x] += v_smem[threadIdx.x + 2]; + v_smem[threadIdx.x] += v_smem[threadIdx.x + 1]; + } + __syncthreads(); + } + + /* + * input case n dim should be divided by 4096 with dtype bf16 + * and should be divided by 2048 with dtype fp32 + * */ + template + __global__ void __launch_bounds__(tnum, 1) local_device_load_rmsnorm_naive( + RankSignals sg, + T* __restrict__ results, + T* __restrict__ weight, + float eps, + int rank, + int m, + int n + ) + { + constexpr int pack_size = packed_t::P::size; + using P = typename packed_t::P; + using A = typename packed_t::A; + __shared__ float smem[tnum]; + P* tmps = get_tmp_buf

(sg.signals[rank]); + + for (int bid = blockIdx.x; bid < m; bid += gridDim.x) + { + float square_sum = 0.0f; + P rmsnorm_inp[n_loop]; + P w_arr[n_loop]; +#pragma unroll + for (int n_iter = 0; n_iter < n_loop; ++n_iter) + { + int read_idx = bid * n_loop * blockDim.x + n_iter * blockDim.x + threadIdx.x; + rmsnorm_inp[n_iter] = tmps[read_idx]; + w_arr[n_iter] = *(reinterpret_cast(weight) + n_iter * blockDim.x + threadIdx.x); + A reduce_pack; +#pragma unroll + for (int i = 0; i < pack_size; ++i) + { + float ar_elem = ck_tile::type_convert(rmsnorm_inp[n_iter].data[i]); + reduce_pack.data[i] = ar_elem * ar_elem; + } + square_sum += packReduce(reduce_pack); + } + smem[threadIdx.x] = square_sum; + __syncthreads(); + smemReduceSum(&smem[0]); + square_sum = smem[0]; + float denom = rsqrtf(square_sum / n + eps); +#pragma unroll + for (int n_iter = 0; n_iter < n_loop; ++n_iter) + { + P rmsnorm_rslt; +#pragma unroll + for (int i = 0; i < pack_size; ++i) + { + float x_f32 = ck_tile::type_convert(rmsnorm_inp[n_iter].data[i]); + float w_f32 = ck_tile::type_convert(w_arr[n_iter].data[i]); + rmsnorm_rslt.data[i] = ck_tile::type_convert(x_f32 * w_f32 * denom); + } + int write_idx = bid * n_loop * blockDim.x + n_iter * blockDim.x + threadIdx.x; + *(reinterpret_cast(results) + write_idx) = rmsnorm_rslt; + } + } + } + + /* + * block size can be 256 and 512 + * corresponding 2048 and 4096 elem per block + * */ + template + __global__ void __launch_bounds__(tnum, 1) local_device_load_rmsnorm( + RankSignals sg, + T* __restrict__ results, + T* __restrict__ weight, + float eps, + int rank, + int m, + int n + ) + { + constexpr int pack_size = packed_t::P::size; + using P = typename packed_t::P; + using A = typename packed_t::A; + __shared__ float smem[tnum]; + P* tmps = get_tmp_buf

(sg.signals[rank]); + + for (int bid = blockIdx.x; bid < m; bid += gridDim.x) + { + float square_sum = 0.0f; + P rmsnorm_inp[n_loop]; + P w_arr[n_loop]; +#pragma unroll + for (int n_iter = 0; n_iter < n_loop; ++n_iter) + { + if (n_iter * tnum + threadIdx.x < (n / pack_size)) + { + int read_idx = bid * (n / pack_size) + n_iter * tnum + threadIdx.x; + rmsnorm_inp[n_iter] = tmps[read_idx]; + w_arr[n_iter] = *(reinterpret_cast(weight) + n_iter * tnum + threadIdx.x); + A reduce_pack; +#pragma unroll + for (int i = 0; i < pack_size; ++i) + { + float ar_elem = ck_tile::type_convert(rmsnorm_inp[n_iter].data[i]); + reduce_pack.data[i] = ar_elem * ar_elem; + } + square_sum += packReduce(reduce_pack); + } + } + smem[threadIdx.x] = square_sum; + __syncthreads(); + smemReduceSum(&smem[0]); + square_sum = smem[0]; + float denom = rsqrtf(square_sum / n + eps); +#pragma unroll + for (int n_iter = 0; n_iter < n_loop; ++n_iter) + { + if (n_iter * tnum + threadIdx.x < (n / pack_size)) + { + P rmsnorm_rslt; +#pragma unroll + for (int i = 0; i < pack_size; ++i) + { + float x_f32 = ck_tile::type_convert(rmsnorm_inp[n_iter].data[i]); + float w_f32 = ck_tile::type_convert(w_arr[n_iter].data[i]); + rmsnorm_rslt.data[i] = ck_tile::type_convert(x_f32 * w_f32 * denom); + } + int write_idx = bid * (n / pack_size) + n_iter * tnum + threadIdx.x; + *(reinterpret_cast(results) + write_idx) = rmsnorm_rslt; + } + } + } + } + + template + __global__ void __launch_bounds__(256, 1) local_device_load_rmsnorm_512n( + RankSignals sg, + T* __restrict__ results, + T* __restrict__ weight, + float eps, + int rank, + int m, + int n + ) + { + constexpr int pack_size = packed_t::P::size; + using P = typename packed_t::P; + using A = typename packed_t::A; + P* tmps = get_tmp_buf

(sg.signals[rank]); + int warp_id = threadIdx.x / 64; + int lane_id = threadIdx.x % 64; + int warp_num = blockDim.x / 64; + + for (int bid = blockIdx.x * warp_num + warp_id; bid < m; bid += gridDim.x * warp_num) + { + float square_sum = 0.0f; + P rmsnorm_inp[n_loop]; + P w_arr[n_loop]; +#pragma unroll + for (int n_iter = 0; n_iter < n_loop; ++n_iter) + { + int read_idx = bid * 64 * n_loop + n_iter * 64 + lane_id; + rmsnorm_inp[n_iter] = tmps[read_idx]; + w_arr[n_iter] = *(reinterpret_cast(weight) + n_iter * 64 + lane_id); + A reduce_pack; +#pragma unroll + for (int i = 0; i < pack_size; ++i) + { + float ar_elem = ck_tile::type_convert(rmsnorm_inp[n_iter].data[i]); + reduce_pack.data[i] = ar_elem * ar_elem; + } + float tmp_sum = packReduce(reduce_pack); + square_sum += tmp_sum; + } + square_sum = warpReduce(square_sum); + float denom = rsqrtf(square_sum / n + eps); +#pragma unroll + for (int n_iter = 0; n_iter < n_loop; ++n_iter) + { + P rmsnorm_rslt; +#pragma unroll + for (int i = 0; i < pack_size; ++i) + { + float x_f32 = ck_tile::type_convert(rmsnorm_inp[n_iter].data[i]); + float w_f32 = ck_tile::type_convert(w_arr[n_iter].data[i]); + rmsnorm_rslt.data[i] = ck_tile::type_convert(x_f32 * w_f32 * denom); + } + int write_idx = bid * 64 * n_loop + n_iter * 64 + lane_id; + *(reinterpret_cast(results) + write_idx) = rmsnorm_rslt; + } + } + } + using IPC_KEY = std::array; static_assert(sizeof(IPC_KEY) == sizeof(hipIpcMemHandle_t)); static_assert(alignof(IPC_KEY) == alignof(hipIpcMemHandle_t)); @@ -1293,6 +1580,146 @@ namespace aiter } } + template + void dispatchFusedAllReduceRMSNorm(hipStream_t stream, T* input, T* output, T* weight, float eps, int m, int n) + { + auto d = packed_t::P::size; + int size = m * n; + if (size % d != 0) + { + throw std::runtime_error( + "custom allreduce currently requires input length to be multiple " + "of " + + std::to_string(d)); + } + RankData* ptrs = get_buffer_RD(stream, input); + hipDevice_t dev; + hipDeviceProp_t dev_prop; + hipGetDevice(&dev); + hipGetDeviceProperties(&dev_prop, dev); + uint32_t num_cu = dev_prop.multiProcessorCount; + + // step 1, run reduce-scatter + allgather cross device save + dim3 block(512); + int block_num = ((size / world_size_) + 512 - 1) / 512; + dim3 grid(std::min(block_num, 80)); + switch (world_size_) + { + case 8: + reduce_scatter_cross_device_store<<>>(ptrs, sg_, self_sg_, rank_, size); + break; + case 4: + reduce_scatter_cross_device_store<<>>(ptrs, sg_, self_sg_, rank_, size); + break; + case 2: + reduce_scatter_cross_device_store<<>>(ptrs, sg_, self_sg_, rank_, size); + break; + default: + printf("fused allreduce rmsnorm world size error\n"); + } + + // step 2, run allgather local device load + rmsnorm + int n_bytes = n * sizeof(T); + auto setGrid = [&](int naive_grid_size, const void* kernel_ptr) + { + int occupancy; + hipOccupancyMaxActiveBlocksPerMultiprocessor(&occupancy, kernel_ptr, block.x, 0); + grid.x = naive_grid_size < num_cu * occupancy ? naive_grid_size : num_cu * occupancy; + }; + +#define launch_fused_allreduce_rmsnorm(template_kernel) \ + do \ + { \ + auto kernel_ptr = reinterpret_cast(template_kernel); \ + setGrid(naive_grid_size, kernel_ptr); \ + template_kernel<<>>(sg_, output, weight, eps, rank_, m, n); \ + } while (0) + + if (n_bytes % 1024 == 0) + { + if (8192 <= n_bytes && n_bytes <= 32768) + { + int naive_grid_size = m; + int n_loop = n_bytes / 8192; // 1, 2, 3, 4 + if (n_bytes % 8192 == 0) + { + switch (n_loop) + { + case 1: + launch_fused_allreduce_rmsnorm((local_device_load_rmsnorm_naive)); + break; + case 2: + launch_fused_allreduce_rmsnorm((local_device_load_rmsnorm_naive)); + break; + case 3: + launch_fused_allreduce_rmsnorm((local_device_load_rmsnorm_naive)); + break; + case 4: + launch_fused_allreduce_rmsnorm((local_device_load_rmsnorm_naive)); + break; + } + } + else + { + n_loop += 1; + switch (n_loop) + { + case 2: + launch_fused_allreduce_rmsnorm((local_device_load_rmsnorm)); + break; + case 3: + launch_fused_allreduce_rmsnorm((local_device_load_rmsnorm)); + break; + case 4: + launch_fused_allreduce_rmsnorm((local_device_load_rmsnorm)); + break; + } + } + } + else if (4096 <= n_bytes && n_bytes < 8192) + { + block.x = 256; + int naive_grid_size = m; + if (n_bytes == 4096) + { + // naive n_loop = 1 + launch_fused_allreduce_rmsnorm((local_device_load_rmsnorm_naive)); + } + else + { + // n_loop = 2 + launch_fused_allreduce_rmsnorm((local_device_load_rmsnorm)); + } + } + else if (1024 <= n_bytes && n_bytes < 4096) + { + block.x = 256; + int naive_grid_size = (m + 3) / 4; + int n_loop = n_bytes / 1024; + switch (n_loop) + { + case 1: + launch_fused_allreduce_rmsnorm((local_device_load_rmsnorm_512n)); + break; + case 2: + launch_fused_allreduce_rmsnorm((local_device_load_rmsnorm_512n)); + break; + case 3: + launch_fused_allreduce_rmsnorm((local_device_load_rmsnorm_512n)); + break; + } + } + else + { + printf("fused allreduce rmsnorm shape size error\n"); + } + } + else + { + printf("fused allreduce rmsnorm shape error\n"); + } + } + ~CustomAllreduce() { for (auto [_, ptr] : ipc_handles_) diff --git a/csrc/include/custom_all_reduce.h b/csrc/include/custom_all_reduce.h index 72f368b7542..d8a370e514d 100644 --- a/csrc/include/custom_all_reduce.h +++ b/csrc/include/custom_all_reduce.h @@ -32,12 +32,18 @@ void all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out, bool open_fp8_quant, - std::optional& reg_buffer); + std::optional reg_buffer); void all_gather_reg(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out); void all_gather_unreg(fptr_t _fa, torch::Tensor& inp, torch::Tensor& reg_buffer, torch::Tensor& out); +void fused_allreduce_rmsnorm(fptr_t _fa, + torch::Tensor& inp, + torch::Tensor& out, + torch::Tensor& w, + float eps, + std::optional reg_buffer); void dispose(fptr_t _fa); int64_t meta_size(); diff --git a/csrc/include/rocm_ops.hpp b/csrc/include/rocm_ops.hpp index e8dec972f7c..cf7834a20b0 100644 --- a/csrc/include/rocm_ops.hpp +++ b/csrc/include/rocm_ops.hpp @@ -294,6 +294,14 @@ py::arg("out"), \ py::arg("open_fp8_quant"), \ py::arg("reg_buffer") = std::nullopt); \ + m.def("fused_allreduce_rmsnorm", \ + &aiter::fused_allreduce_rmsnorm, \ + py::arg("_fa"), \ + py::arg("inp"), \ + py::arg("out"), \ + py::arg("w"), \ + py::arg("eps"), \ + py::arg("reg_buffer") = std::nullopt); \ m.def("all_reduce_asm_", &all_reduce_asm, ""); \ m.def("all_reduce_rmsnorm_", &all_reduce_rmsnorm, "all_reduce_rmsnorm"); \ m.def("all_reduce_rmsnorm_quant_", &all_reduce_rmsnorm_quant, "all_reduce_rmsnorm_quant"); \ diff --git a/csrc/kernels/custom_all_reduce.cu b/csrc/kernels/custom_all_reduce.cu index 2e25b40f23d..7845fae0861 100644 --- a/csrc/kernels/custom_all_reduce.cu +++ b/csrc/kernels/custom_all_reduce.cu @@ -133,7 +133,7 @@ void all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out, bool open_fp8_quant, - std::optional& reg_buffer) + std::optional reg_buffer) { const at::hip::OptionalHIPGuardMasqueradingAsCUDA device_guard(device_of(inp)); auto stream = c10::hip::getCurrentHIPStreamMasqueradingAsCUDA().stream(); @@ -216,6 +216,76 @@ void all_gather_unreg(fptr_t _fa, torch::Tensor& inp, torch::Tensor& reg_buffer, _all_gather(_fa, reg_buffer, out, inp.numel(), stream); } +void _fused_allreduce_rmsnorm( + fptr_t _fa, torch::Tensor& inp, torch::Tensor& out, torch::Tensor& w, int eps, int m, int n, hipStream_t stream) +{ + auto fa = reinterpret_cast(_fa); + TORCH_CHECK(_is_weak_contiguous(out)); + switch(out.scalar_type()) + { + case at::ScalarType::Float: { + fa->dispatchFusedAllReduceRMSNorm(stream, + reinterpret_cast(inp.data_ptr()), + reinterpret_cast(out.data_ptr()), + reinterpret_cast(w.data_ptr()), + eps, m, n); + break; + } + case at::ScalarType::Half: { + fa->dispatchFusedAllReduceRMSNorm(stream, + reinterpret_cast(inp.data_ptr()), + reinterpret_cast(out.data_ptr()), + reinterpret_cast(w.data_ptr()), + eps, m, n); + break; + } +#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) + case at::ScalarType::BFloat16: { + fa->dispatchFusedAllReduceRMSNorm<__hip_bfloat16>(stream, + reinterpret_cast<__hip_bfloat16*>(inp.data_ptr()), + reinterpret_cast<__hip_bfloat16*>(out.data_ptr()), + reinterpret_cast<__hip_bfloat16*>(w.data_ptr()), + eps, m, n); + break; + } +#endif + default: + throw std::runtime_error("custom allreduce only supports float32, float16 and bfloat16"); + } +} + +void fused_allreduce_rmsnorm(fptr_t _fa, + torch::Tensor& inp, + torch::Tensor& out, + torch::Tensor& w, + float eps, + std::optional reg_buffer) +{ + const at::hip::OptionalHIPGuardMasqueradingAsCUDA device_guard(device_of(inp)); + auto stream = c10::hip::getCurrentHIPStreamMasqueradingAsCUDA().stream(); + TORCH_CHECK_EQ(inp.scalar_type(), out.scalar_type()); + TORCH_CHECK_EQ(inp.numel(), out.numel()); + int n = w.numel(); + int m = inp.numel() / n; + + if(reg_buffer.has_value()) + { + auto input_size = inp.numel() * inp.element_size(); + TORCH_CHECK(input_size <= reg_buffer.value().numel() * reg_buffer.value().element_size(), + "registered buffer is too small to contain the input"); + HIP_CALL(hipMemcpyAsync(reg_buffer.value().data_ptr(), + inp.data_ptr(), + input_size, + hipMemcpyDeviceToDevice, + stream)); + _fused_allreduce_rmsnorm(_fa, reg_buffer.value(), out, w, eps, m, n, stream); + } + else + { + _fused_allreduce_rmsnorm(_fa, inp, out, w, eps, m, n, stream); + } +} + void dispose(fptr_t _fa) { auto fa = reinterpret_cast(_fa); diff --git a/op_tests/multigpu_tests/test_fused_ar_rms.py b/op_tests/multigpu_tests/test_fused_ar_rms.py new file mode 100644 index 00000000000..33e0e16b934 --- /dev/null +++ b/op_tests/multigpu_tests/test_fused_ar_rms.py @@ -0,0 +1,533 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. + +import os +import aiter +import torch +import torch.nn.functional as F +import torch.distributed as dist +import argparse +import itertools +from aiter import dtypes + +from aiter.dist.parallel_state import ( + ensure_model_parallel_initialized, + init_distributed_environment, + set_custom_all_reduce, + get_tp_group, + graph_capture, + destroy_model_parallel, + destroy_distributed_environment, +) +from aiter.dist.utils import get_open_port, get_distributed_init_method, get_ip +from aiter.dist.communication_op import ( + tensor_model_parallel_all_reduce, + tensor_model_parallel_fused_allreduce_rmsnorm, +) +from aiter.test_common import ( + checkAllclose, + perftest, + benchmark, +) +from multiprocessing import set_start_method, Pool, freeze_support +import logging + +logger = logging.getLogger("aiter") + +set_start_method("spawn", force=True) + + +def fused_ar_rmsnorm(tp_size, pp_size, rankID, x, weight, eps, withGraph=False): + device = torch.device(f"cuda:{rankID}") + torch.cuda.set_device(device) + # init + logger.info(f"RANK: {rankID} {tp_size} init_process_group...") + set_custom_all_reduce(True) + init_distributed_environment( + world_size=tp_size, + rank=rankID, + distributed_init_method=get_distributed_init_method(get_ip(), get_open_port()), + ) + ensure_model_parallel_initialized(tp_size, pp_size) + x = x.to(device) + weight = weight.to(device) + # dist.barrier(device_ids=[i for i in range(tp_size)]) + + # warmup and align all gpu + group = get_tp_group().device_group + dist.all_reduce(torch.zeros(1).cuda(), group=group) + torch.cuda.synchronize() + + if withGraph: + graph = torch.cuda.CUDAGraph() + with graph_capture() as gc: + with torch.cuda.graph(graph, stream=gc.stream): + out = tensor_model_parallel_fused_allreduce_rmsnorm(x, weight, eps) + out.fill_(0) + + @perftest() + def run_ca(): + graph.replay() + + _, us = run_ca() + out = (out, us) + else: + + @perftest() + def run_ca(x): + return tensor_model_parallel_fused_allreduce_rmsnorm(x, weight, eps) + + out = run_ca(x) + + # destroy + if dist.is_initialized(): + destroy_model_parallel() + destroy_distributed_environment() + torch.cuda.empty_cache() + return out + + +def get_acc_value_with_cudagraph(tp_size, pp_size, rankID, x, weight, eps, loop_time=1): + device = torch.device(f"cuda:{rankID}") + torch.cuda.set_device(device) + # init + logger.info(f"RANK: {rankID} {tp_size} init_process_group...") + set_custom_all_reduce(True) + init_distributed_environment( + world_size=tp_size, + rank=rankID, + distributed_init_method=get_distributed_init_method(get_ip(), get_open_port()), + ) + ensure_model_parallel_initialized(tp_size, pp_size) + x = x.to(device) + weight = weight.to(device) + # dist.barrier(device_ids=[i for i in range(tp_size)]) + + # warmup and align all gpu + group = get_tp_group().device_group + dist.all_reduce(torch.zeros(1).cuda(), group=group) + torch.cuda.synchronize() + + # out = torch.empty_like(x) + graph = torch.cuda.CUDAGraph() + with graph_capture() as gc: + with torch.cuda.graph(graph, stream=gc.stream): + # out = torch.empty_like(x) + out = tensor_model_parallel_fused_allreduce_rmsnorm(x, weight, eps) + out.fill_(0) + + def run_ca(): + graph.replay() + rslt = out.clone() + out.fill_(0) + return rslt + + for i in range(loop_time): + out = run_ca() + + # destroy + if dist.is_initialized(): + destroy_model_parallel() + destroy_distributed_environment() + torch.cuda.empty_cache() + return out + + +def get_acc_value_only(tp_size, pp_size, rankID, x, weight, eps, loop_time=1): + device = torch.device(f"cuda:{rankID}") + torch.cuda.set_device(device) + # init + logger.info(f"RANK: {rankID} {tp_size} init_process_group...") + set_custom_all_reduce(True) + init_distributed_environment( + world_size=tp_size, + rank=rankID, + distributed_init_method=get_distributed_init_method(get_ip(), get_open_port()), + ) + ensure_model_parallel_initialized(tp_size, pp_size) + x = x.to(device) + weight = weight.to(device) + # dist.barrier(device_ids=[i for i in range(tp_size)]) + + # warmup and align all gpu + group = get_tp_group().device_group + torch.cuda.synchronize() + + for i in range(loop_time): + out = tensor_model_parallel_fused_allreduce_rmsnorm(x, weight, eps) + + # destroy + if dist.is_initialized(): + destroy_model_parallel() + destroy_distributed_environment() + torch.cuda.empty_cache() + return out + + +def split_ar_rmsnorm(tp_size, pp_size, rankID, x, weight, eps, withGraph=False): + device = torch.device(f"cuda:{rankID}") + torch.cuda.set_device(device) + # init + logger.info(f"RANK: {rankID} {tp_size} init_process_group...") + set_custom_all_reduce(True) + init_distributed_environment( + world_size=tp_size, + rank=rankID, + distributed_init_method=get_distributed_init_method(get_ip(), get_open_port()), + ) + ensure_model_parallel_initialized(tp_size, pp_size) + x = x.to(device) + weight = weight.to(device) + # dist.barrier(device_ids=[i for i in range(tp_size)]) + + # warmup and align all gpu + group = get_tp_group().device_group + dist.all_reduce(torch.zeros(1).cuda(), group=group) + torch.cuda.synchronize() + + if withGraph: + graph = torch.cuda.CUDAGraph() + with graph_capture() as gc: + with torch.cuda.graph(graph, stream=gc.stream): + ar_out = tensor_model_parallel_all_reduce(x) + # out = aiter.rms_norm(ar_out, weight, eps, 0) + out = torch.empty_like(ar_out) + residual_out = torch.empty_like(ar_out) + aiter.rmsnorm2d_fwd_with_add( + out, + ar_out, + x, + residual_out, + weight, + eps, + 0, + ) + out.fill_(0) + + @perftest() + def run_ca(): + graph.replay() + + _, us = run_ca() + out = (out, us) + else: + + @perftest() + def run_ca(x): + ar_out = tensor_model_parallel_all_reduce(x) + out = torch.empty_like(ar_out) + residual_out = torch.empty_like(ar_out) + aiter.rmsnorm2d_fwd_with_add( + out, + ar_out, + x, + residual_out, + weight, + eps, + 0, + ) + return out + + out = run_ca(x) + + # destroy + if dist.is_initialized(): + destroy_model_parallel() + destroy_distributed_environment() + torch.cuda.empty_cache() + return out + + +def run_cu(input, weight, eps, device_id): + device = f"cuda:{device_id}" + input = input.to(device) + weight = weight.to(device) + + @perftest() + def compute(): + output = torch.empty_like(input) + aiter.rms_norm_cu(output, input, weight, eps) + + return compute() + + +@benchmark() +def test_split_ar_rmsnorm(tp_size, pp_size, shape, dtype, withGraph=False): + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = "49373" + pool = Pool(processes=tp_size) + ref = torch.zeros(shape, dtype=dtype) + rets = [] + cpu_rslt = [] + weight_list = [] + res_inp = [] + # print(type(shape[0]), shape[1], ref.device) + m = shape[0] + n = shape[1] + eps = 1e-6 + for i in range(tp_size): + x = torch.randn(shape, dtype=dtype) + res_inp.append(x) + ref += x + weight = torch.randn((n,), dtype=dtype) + weight_list.append(weight) + rets.append( + pool.apply_async( + split_ar_rmsnorm, args=(tp_size, pp_size, i, x, weight, eps, withGraph) + ) + # pool.apply_async(run_cu, args=(x, weight, eps, i)) + ) + pool.close() + pool.join() + for i in range(tp_size): + host_rslt = F.rms_norm( + input=(ref + res_inp[i]), + normalized_shape=(ref.shape[-1],), + weight=weight_list[i], + eps=eps, + ) + cpu_rslt.append(host_rslt) + rets = [el.get() for el in rets] + for out, us in rets: + msg = f"test_split_ar_rmsnorm: {shape=} {dtype=} {withGraph=} {us:>8.2f}" + # print(cpu_rslt[out.device.index]) + checkAllclose(cpu_rslt[out.device.index], out.to(ref), msg=msg) + + +@benchmark() +def test_fused_ar_rmsnorm(tp_size, pp_size, shape, dtype, withGraph=False): + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = "49373" + pool = Pool(processes=tp_size) + ref = torch.zeros(shape, dtype=dtype) + rets = [] + cpu_rslt = [] + weight_list = [] + res_inp = [] + # print(type(shape[0]), shape[1], ref.device) + m = shape[0] + n = shape[1] + eps = 1e-6 + for i in range(tp_size): + x = torch.randn(shape, dtype=dtype) + # x = torch.ones(shape, dtype=dtype) + res_inp.append(x) + # print(f"device {i}, x[0][0] = {x[0][0]}") + ref += x + weight = torch.randn((n,), dtype=dtype) + weight_list.append(weight) + rets.append( + pool.apply_async( + fused_ar_rmsnorm, args=(tp_size, pp_size, i, x, weight, eps, withGraph) + ) + # pool.apply_async(run_cu, args=(x, weight, eps, i)) + ) + pool.close() + pool.join() + print(f"rslt[0][0] = {ref[0][0]}") + + for i in range(tp_size): + host_rslt = F.rms_norm( + input=(ref + res_inp[i]), + normalized_shape=(ref.shape[-1],), + weight=weight_list[i], + eps=eps, + ) + # host_rslt = ref + res_inp[i] + cpu_rslt.append(host_rslt) + + rets = [el.get() for el in rets] + for out, us in rets: + msg = f"test_fused_ar_rmsnorm: {shape=} {dtype=} {withGraph=} {us:>8.2f}" + # print(cpu_rslt[out.device.index]) + checkAllclose(cpu_rslt[out.device.index], out.to(ref), msg=msg) + # checkAllclose(ref, out.to(ref), msg=msg) + + +def acc_test(tp_size, pp_size, shape, dtype): + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = "49373" + pool = Pool(processes=tp_size) + ref = torch.zeros(shape, dtype=dtype) + rets = [] + cpu_rslt = [] + weight_list = [] + # print(type(shape[0]), shape[1], ref.device) + m = shape[0] + n = shape[1] + eps = 1e-6 + for i in range(tp_size): + x = torch.randn(shape, dtype=dtype) + ref += x + weight = torch.randn((n,), dtype=dtype) + weight_list.append(weight) + rets.append( + pool.apply_async( + get_acc_value_only, args=(tp_size, pp_size, i, x, weight, eps, 1) + ) + ) + pool.close() + pool.join() + + ar_rslt = [] + for i, ret in enumerate(rets): + rslt = ret.get() + ar_rslt.append(rslt) + for i in ar_rslt: + checkAllclose(ref, i.to(ref)) + + +def acc_test_cudagraph_on(tp_size, pp_size, shape, dtype, loop_time=1): + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = "49373" + pool = Pool(processes=tp_size) + ref = torch.zeros(shape, dtype=dtype) + rets = [] + cpu_rslt = [] + weight_list = [] + # print(type(shape[0]), shape[1], ref.device) + m = shape[0] + n = shape[1] + eps = 1e-6 + for i in range(tp_size): + x = torch.randn(shape, dtype=dtype) + ref += x + weight = torch.randn((n,), dtype=dtype) + weight_list.append(weight) + rets.append( + pool.apply_async( + get_acc_value_with_cudagraph, + args=(tp_size, pp_size, i, x, weight, eps, loop_time), + ) + ) + pool.close() + pool.join() + + ar_rslt = [] + for i, ret in enumerate(rets): + rslt = ret.get() + ar_rslt.append(rslt) + for i in ar_rslt: + checkAllclose(ref, i.to(ref)) + + +# def acc_test(tp_size, pp_size, shape, dtype): +# os.environ["MASTER_ADDR"] = "127.0.0.1" +# os.environ["MASTER_PORT"] = "49373" +# pool = Pool(processes=tp_size) +# ref = torch.zeros(shape, dtype=dtype) +# rets = [] +# cpu_rslt = [] +# weight_list = [] +# # print(type(shape[0]), shape[1], ref.device) +# m = shape[0] +# n = shape[1] +# eps = 1e-6 +# for i in range(tp_size): +# x = torch.randn(shape, dtype=dtype) +# print(f"device {i}, x[0][0] = {x[0][0]}") +# ref += x +# weight = torch.randn((n,), dtype=dtype) +# weight_list.append(weight) +# rets.append( +# pool.apply_async(get_acc_value_only, args=(tp_size, pp_size, i, x, weight, eps)) +# ) +# pool.close() +# pool.join() +# for i in range(tp_size): +# host_rslt = F.rms_norm( +# input=ref, normalized_shape=(ref.shape[-1],), weight=weight_list[i], eps=eps +# ) +# cpu_rslt.append(host_rslt) +# +# ar_rslt = [] +# for i, ret in enumerate(rets): +# rslt = ret.get() +# ar_rslt.append(rslt) +# for i in range(len(ar_rslt)): +# checkAllclose(cpu_rslt[i], ar_rslt[i].to(ref)) + +l_dtype = ["bf16"] +l_shape = [ + # (4096, 2048) + (64, 7168) + # (64, 512 * 99) + # (16, 512) +] +l_tp = [8] +l_pp = [1] +l_graph = [True, False] + +parser = argparse.ArgumentParser(description="config input of test") +parser.add_argument( + "-d", + "--dtype", + type=str, + choices=l_dtype, + nargs="?", + const=None, + default=None, + help="data type", +) +parser.add_argument( + "-s", + "--shape", + type=dtypes.str2tuple, + nargs="?", + const=None, + default=None, + help="shape. e.g. -s 128,8192", +) + +parser.add_argument( + "-t", + "--tp", + type=int, + nargs="?", + const=None, + default=None, + help="tp num. e.g. -t 8", +) + +parser.add_argument( + "-p", + "--pp", + type=int, + nargs="?", + const=None, + default=None, + help="tp num. e.g. -p 1", +) + +parser.add_argument( + "-g", + "--graphon", + type=int, + nargs="?", + const=None, + default=None, + help="open cudagraph. e.g. -g 1", +) + + +if __name__ == "__main__": + freeze_support() + args = parser.parse_args() + if args.dtype is None: + l_dtype = [dtypes.d_dtypes[key] for key in l_dtype] + else: + l_dtype = [dtypes.d_dtypes[args.dtype]] + if args.shape is not None: + l_shape = [args.shape] + if args.tp is not None: + l_tp = [args.tp] + if args.pp is not None: + l_pp = [args.pp] + if args.graphon is not None: + print(args.graphon) + l_graph = [args.graphon] + for dtype, shape, tp, pp, graph_on in itertools.product( + l_dtype, l_shape, l_tp, l_pp, l_graph + ): + test_split_ar_rmsnorm(tp, pp, shape, dtype, withGraph=graph_on) + test_fused_ar_rmsnorm(tp, pp, shape, dtype, withGraph=graph_on)