From c8ab68af6c3c7866ad1cef6326649897568167fc Mon Sep 17 00:00:00 2001 From: Yi Dong <43824965+yidong72@users.noreply.github.com> Date: Fri, 5 Aug 2022 12:46:05 -0400 Subject: [PATCH 01/25] Add a generic fused softmax (#1440) * new kernel Signed-off-by: Yi Dong * added the unit tests Signed-off-by: Yi Dong * clean up unittest Signed-off-by: Yi Dong * use float Signed-off-by: Yi Dong * more clean up Signed-off-by: Yi Dong * remove the long seq test case --- apex/transformer/functional/fused_softmax.py | 54 +++ .../generic_scaled_masked_softmax.cpp | 82 ++++ csrc/megatron/generic_scaled_masked_softmax.h | 384 ++++++++++++++++++ .../generic_scaled_masked_softmax_cuda.cu | 116 ++++++ csrc/megatron/scaled_masked_softmax_cuda.cu | 4 +- ...scaled_upper_triang_masked_softmax_cuda.cu | 4 +- setup.py | 9 +- .../L0/run_transformer/test_fused_softmax.py | 72 ++++ 8 files changed, 722 insertions(+), 3 deletions(-) create mode 100644 csrc/megatron/generic_scaled_masked_softmax.cpp create mode 100644 csrc/megatron/generic_scaled_masked_softmax.h create mode 100644 csrc/megatron/generic_scaled_masked_softmax_cuda.cu diff --git a/apex/transformer/functional/fused_softmax.py b/apex/transformer/functional/fused_softmax.py index 8ceaffef9..06a45c461 100644 --- a/apex/transformer/functional/fused_softmax.py +++ b/apex/transformer/functional/fused_softmax.py @@ -98,6 +98,33 @@ def scaled_masked_softmax(inputs, mask, scale): return ScaledMaskedSoftmax.apply(*args) +class GenericScaledMaskedSoftmax(torch.autograd.Function): + @staticmethod + def forward(ctx, inputs, mask, scale): + import generic_scaled_masked_softmax_cuda + + scale_t = torch.tensor([scale]) + softmax_results = generic_scaled_masked_softmax_cuda.forward(inputs, mask, scale_t[0]) + ctx.save_for_backward(softmax_results, scale_t) + return softmax_results + + @staticmethod + def backward(ctx, output_grads): + import generic_scaled_masked_softmax_cuda_new + + softmax_results, scale_t = ctx.saved_tensors + + input_grads = generic_scaled_masked_softmax_cuda.backward(output_grads, softmax_results, scale_t[0]) + return input_grads, None, None + + +def generic_scaled_masked_softmax(inputs, mask, scale): + # input is 4D tensor (b, np, sq, sk) + args = _cast_if_autocast_enabled(inputs, mask, scale) + with torch.cuda.amp.autocast(enabled=False): + return GenericScaledMaskedSoftmax.apply(*args) + + class FusedScaleMaskSoftmax(torch.nn.Module): """ fused operation: scaling + mask + softmax @@ -209,3 +236,30 @@ def get_batch_per_block(sq, sk, b, np): import scaled_masked_softmax_cuda return scaled_masked_softmax_cuda.get_batch_per_block(sq, sk, b, np) + +class GenericFusedScaleMaskSoftmax(FusedScaleMaskSoftmax): + """ + Generic version of FusedSacleMaskSoftmax. + It removes the seq-len limitations and has slight performance degragation compared with FusedScaleMaskSoftmax + + fused operation: scaling + mask + softmax + + Arguments: + input_in_fp16: flag to indicate if input in fp16 data format. + input_in_bf16: flag to indicate if input in bf16 data format. + scaled_masked_softmax_fusion: flag to indicate user want to use softmax fusion + mask_func: mask function to be applied. + softmax_in_fp32: if true, softmax in performed at fp32 precision. + scale: scaling factor used in input tensor scaling. + """ + + def __init__( + self, input_in_fp16, input_in_bf16, scaled_masked_softmax_fusion, mask_func, softmax_in_fp32, scale, + ): + super().__init__(input_in_fp16, input_in_bf16, AttnMaskType.padding, scaled_masked_softmax_fusion, mask_func, softmax_in_fp32, scale) + self.scaled_masked_softmax_fusion = generic_scaled_masked_softmax + + def is_kernel_available(self, mask, b, np, sq, sk): + if self.scaled_masked_softmax_fusion and 0 < sk: # user want to fuse # sk must be 1 ~ + return True + return False \ No newline at end of file diff --git a/csrc/megatron/generic_scaled_masked_softmax.cpp b/csrc/megatron/generic_scaled_masked_softmax.cpp new file mode 100644 index 000000000..4bd9f0e7f --- /dev/null +++ b/csrc/megatron/generic_scaled_masked_softmax.cpp @@ -0,0 +1,82 @@ +/* coding=utf-8 + * Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include + +namespace multihead_attn +{ + namespace fused_softmax + { + namespace generic_scaled_masked_softmax + { + + torch::Tensor fwd_cuda( + torch::Tensor const &input, + torch::Tensor const &mask, + float scale_factor); + + torch::Tensor bwd_cuda( + torch::Tensor const &output_grads, + torch::Tensor const &softmax_results, + float scale_factor); + + torch::Tensor fwd( + torch::Tensor const &input, + torch::Tensor const &mask, + float scale_factor) + { + AT_ASSERTM(input.dim() == 4, "expected 4D tensor"); + AT_ASSERTM((input.scalar_type() == at::ScalarType::Half) || + (input.scalar_type() == at::ScalarType::BFloat16), + "Only fp16 and bf16 are supported"); + AT_ASSERTM(mask.dim() == 4, "expected 4D tensor"); + + return fwd_cuda(input, mask, scale_factor); + } + + torch::Tensor bwd( + torch::Tensor const &output_grads, + torch::Tensor const &softmax_results, + float scale_factor) + { + + AT_ASSERTM(output_grads.dim() == 4, "expected 3D tensor"); + AT_ASSERTM(softmax_results.dim() == 4, "expected 3D tensor"); + + AT_ASSERTM((output_grads.scalar_type() == at::ScalarType::Half) || + (output_grads.scalar_type() == at::ScalarType::BFloat16), + "Only fp16 and bf16 are supported"); + AT_ASSERTM((softmax_results.scalar_type() == at::ScalarType::Half) || + (softmax_results.scalar_type() == at::ScalarType::BFloat16), + "Only fp16 and bf16 are supported"); + + return bwd_cuda(output_grads, softmax_results, scale_factor); + } + + } // end namespace generic_scaled_masked_softmax + } // end namespace fused_softmax +} // end namespace multihead_attn + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("forward", + &multihead_attn::fused_softmax::generic_scaled_masked_softmax::fwd, + "Self Multihead Attention scaled, time masked softmax -- Forward."); + + m.def("backward", + &multihead_attn::fused_softmax::generic_scaled_masked_softmax::bwd, + "Self Multihead Attention scaled, time masked softmax -- Backward."); +} diff --git a/csrc/megatron/generic_scaled_masked_softmax.h b/csrc/megatron/generic_scaled_masked_softmax.h new file mode 100644 index 000000000..555ff462f --- /dev/null +++ b/csrc/megatron/generic_scaled_masked_softmax.h @@ -0,0 +1,384 @@ +/* coding=utf-8 + * Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +namespace { + +template +struct Add { + __device__ __forceinline__ T operator()(T a, T b) const { + return a + b; + } +}; + +template +struct Max { + __device__ __forceinline__ T operator()(T a, T b) const { + return a < b ? b : a; + } +}; + +template +__device__ __forceinline__ T WARP_SHFL_DOWN_NATIVE(T value, int laneMask, int width = warpSize, unsigned int mask = 0xffffffff) +{ +#if CUDA_VERSION >= 9000 + return __shfl_down_sync(mask, value, laneMask, width); +#else + return __shfl_down(value, laneMask, width); +#endif +} + +template class ReduceOp> +__device__ __forceinline__ acc_t warp_reduce_new(acc_t val) { + ReduceOp r; + #pragma unroll + for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) + { + val = r(val, WARP_SHFL_DOWN_NATIVE(val, offset, WARP_SIZE)); + } + return val; +} + + +template +__global__ void scaled_masked_softmax_warp_backward_new( + output_t *gradInput, //[batches, attn_heads, q_len, k_len] + input_t *grad, + const input_t *output, //[batches, attn_heads, q_len, k_len] + acc_t scale, + int element_count) +{ + int threads_per_block = blockDim.x; + //the first element_count*2 elements are used for cache, the last 128 is used for reduction + extern __shared__ acc_t shared_data[]; + input_t *local_data = (input_t *)shared_data; + input_t *output_data = &local_data[element_count]; + // maximum shared cached 128, enough for 4096 elements reduction into 4096/32= 128 elements + acc_t *shared = (acc_t *)(&(local_data[element_count*2])); + + int num_reductions = (element_count - 1) / threads_per_block + 1; + + int offset = blockIdx.x * element_count; + + int local_idx = threadIdx.x; + int lane = threadIdx.x % C10_WARP_SIZE; + int wid = threadIdx.x / C10_WARP_SIZE; + int warps_per_thread_block = threads_per_block / C10_WARP_SIZE; + + // load the data to local data + acc_t val = 0.0; + for (int i = 0; i < num_reductions; i++){ + if (i*threads_per_block + local_idx < element_count){ + val = output[offset + i*threads_per_block + local_idx]; + output_data[i*threads_per_block + local_idx] = val; + local_data[i*threads_per_block + local_idx] = val * grad[offset + i*threads_per_block + local_idx]; + } + __syncthreads(); + } + + // find the sum + for (int i = local_idx; i < (element_count - 1) / C10_WARP_SIZE + 1; i += threads_per_block){ + shared[i] = 0.0; + } + __syncthreads(); + + #pragma unroll + for (int i = 0; i < num_reductions; i++){ + if (i*threads_per_block + local_idx < element_count){ + val = local_data[i*threads_per_block + local_idx]; + } + else{ + val = 0.0; + } + __syncthreads(); + val = warp_reduce_new(val); + if (lane==0 && wid + warps_per_thread_block * i < (element_count - 1) / C10_WARP_SIZE + 1) { + shared[wid + warps_per_thread_block*i] = val; + } + __syncthreads(); + } + + // final shared reduction + + int shared_mem_len = (element_count - 1) / C10_WARP_SIZE + 1; + int num_warps = (shared_mem_len - 1) / C10_WARP_SIZE + 1; + while ( shared_mem_len > 1 ){ + #pragma unroll + for (int i = 0; i < num_reductions; i++){ + if (i*threads_per_block + local_idx < shared_mem_len){ + val = shared[i*threads_per_block + local_idx]; + } + else{ + val = 0.0; + } + __syncthreads(); + val = warp_reduce_new(val); + if (lane==0) { + shared[wid + warps_per_thread_block * i] = val; + } + __syncthreads(); + } + shared_mem_len = num_warps; + num_warps = (shared_mem_len - 1) / C10_WARP_SIZE + 1; + } + val = shared[0]; + #pragma unroll + for (int i = local_idx; i < element_count; i += threads_per_block){ + gradInput[offset + i] = (output_t)(scale*(local_data[i] - output_data[i]*val)); + } +} + +} // end of anonymous namespace + +template +void dispatch_scaled_masked_softmax_backward_new( + output_t *grad_input, + input_t *grad, + const input_t *output, + const acc_t scale, + int query_seq_len, + int key_seq_len, + int batches, + int attn_heads) +{ + if (key_seq_len == 0) + { + return; + } + else + { + int batch_count = batches * attn_heads * query_seq_len; + // use 128 threads per block to maximimize gpu utilization + constexpr int threads_per_block = 128; + int num_warps = (key_seq_len - 1) / C10_WARP_SIZE + 1; + dim3 blocks(batch_count, 1, 1); + dim3 threads(threads_per_block, 1, 1); + + scaled_masked_softmax_warp_backward_new + <<>>(grad_input, grad, output, scale, key_seq_len); + } +} + +/* + * Extended softmax (from native aten pytorch) with following additional features + * 1) input scaling + * 2) Explicit masking + */ +template +__global__ void scaled_masked_softmax_warp_forward_new( + output_t *dst, + const input_t *src, + const uint8_t *mask, + const acc_t scale, + int query_len, // query_len + int attn_heads, + int element_count, // key_len + int pad_batches) // mask batch size +{ + // min threawds_per_block has to be bigger than 128 + int threads_per_block = blockDim.x; + // the first element_count is used for cache, the last 128 is used for reduction + extern __shared__ acc_t local_data[]; + // maximum shared cached 128, enough for 4096 elements reduction into 4096/32= 128 elements + acc_t *shared = &(local_data[element_count]); + // number of 1024 threads reductions + int num_reductions = (element_count - 1) / threads_per_block + 1; + + int offset = blockIdx.x * element_count; + int mask_offset; + int query_id = blockIdx.x % query_len; + if (pad_batches == 1){ + // broadcaste the mask tensor + mask_offset = query_id * element_count; + } + else{ + int mask_batch_id = blockIdx.x / attn_heads / query_len; + mask_offset = (mask_batch_id * query_len + query_id) * element_count; + } + + int local_idx = threadIdx.x; + int lane = threadIdx.x % C10_WARP_SIZE; + int wid = threadIdx.x / C10_WARP_SIZE; + int warps_per_thread_block = threads_per_block / C10_WARP_SIZE; + + // load the data to local data + for (int i = local_idx; i < element_count; i += threads_per_block) + { + // TODO, use the copy vector method + if (mask[mask_offset + i] == 1) + { + local_data[i] = -10000.0; + } + else + { + local_data[i] = src[offset + i] * scale; + } + } + + // first find the max value + for (int i = local_idx; i < (element_count - 1) / C10_WARP_SIZE + 1; i += threads_per_block){ + shared[i] = -10000.0; + } + __syncthreads(); + acc_t val = -10000.0; + #pragma unroll + for (int i = 0; i < num_reductions; i++){ + if (i*threads_per_block + local_idx < element_count){ + val = local_data[i*threads_per_block + local_idx]; + } + else{ + val = -10000.0; + } + __syncthreads(); + val = warp_reduce_new(val); + + if (lane==0 && wid + warps_per_thread_block * i < (element_count - 1) / C10_WARP_SIZE + 1) { + shared[wid + warps_per_thread_block*i] = val; + } + __syncthreads(); + } + + // final shared reduction + int shared_mem_len = (element_count - 1) / C10_WARP_SIZE + 1; + int num_warps = (shared_mem_len - 1) / C10_WARP_SIZE + 1; + while ( shared_mem_len > 1 ){ + #pragma unroll + for (int i = 0; i < num_reductions; i++){ + if (i*threads_per_block + local_idx < shared_mem_len){ + val = shared[i*threads_per_block + local_idx]; + } + else{ + val = -10000.0; + } + __syncthreads(); + val = warp_reduce_new(val); + if (lane==0) { + shared[wid + warps_per_thread_block * i] = val; + } + __syncthreads(); + } + shared_mem_len = num_warps; + num_warps = (shared_mem_len - 1) / C10_WARP_SIZE + 1; + } + + acc_t reduced_val = shared[0]; + if (reduced_val < -10000.0 + 0.1){ + // if everything is masked, pay attention to nothing + #pragma unroll + for (int i = local_idx; i < element_count; i += threads_per_block){ + dst[offset + i] = 0.0; + } + return; + } + + // update the values + #pragma unroll + for (int i = local_idx; i < element_count; i += threads_per_block){ + local_data[i] = std::exp(local_data[i] - reduced_val); + } + + // find the sum + for (int i = local_idx; i < (element_count - 1) / C10_WARP_SIZE + 1; i += threads_per_block){ + shared[i] = 0.0; + } + __syncthreads(); + + #pragma unroll + for (int i = 0; i < num_reductions; i++){ + if (i*threads_per_block + local_idx < element_count){ + val = local_data[i*threads_per_block + local_idx]; + } + else{ + val = 0.0; + } + __syncthreads(); + + val = warp_reduce_new(val); + if (lane==0 && wid + warps_per_thread_block * i < (element_count - 1) / C10_WARP_SIZE + 1) { + shared[wid + warps_per_thread_block*i] = val; + } + __syncthreads(); + } + + shared_mem_len = (element_count - 1) / C10_WARP_SIZE + 1; + num_warps = (shared_mem_len - 1) / C10_WARP_SIZE + 1; + while ( shared_mem_len > 1 ){ + #pragma unroll + for (int i = 0; i < num_reductions; i++){ + if (i*threads_per_block + local_idx < shared_mem_len){ + val = shared[i*threads_per_block + local_idx]; + } + else{ + val = 0.0; + } + __syncthreads(); + val = warp_reduce_new(val); + if (lane==0) { + shared[wid + warps_per_thread_block * i] = val; + } + __syncthreads(); + } + shared_mem_len = num_warps; + num_warps = (shared_mem_len - 1) / C10_WARP_SIZE + 1; + } + + reduced_val = shared[0]; + + #pragma unroll + for (int i = local_idx; i < element_count; i += threads_per_block){ + dst[offset + i] = local_data[i] / reduced_val; + } +} + + +template +void dispatch_scaled_masked_softmax_forward_new( + output_t *dst, + const input_t *src, + const uint8_t *mask, + const input_t scale, + int query_seq_len, + int key_seq_len, + int batches, + int attn_heads, + int pad_batches) +{ + if (key_seq_len == 0) { + return; + } else { + int batch_count = batches * attn_heads * query_seq_len; + + // use 128 threads per block to maximimize gpu utilization + constexpr int threads_per_block = 128; + + // calculate the needed shared memory + int num_warps = (key_seq_len - 1) / C10_WARP_SIZE + 1; + + dim3 blocks(batch_count, 1, 1); + dim3 threads(threads_per_block, 1, 1); + scaled_masked_softmax_warp_forward_new + <<>>(dst, src, mask, scale, query_seq_len, attn_heads, key_seq_len, pad_batches); + } +} diff --git a/csrc/megatron/generic_scaled_masked_softmax_cuda.cu b/csrc/megatron/generic_scaled_masked_softmax_cuda.cu new file mode 100644 index 000000000..5af9aa9cc --- /dev/null +++ b/csrc/megatron/generic_scaled_masked_softmax_cuda.cu @@ -0,0 +1,116 @@ +/* coding=utf-8 + * Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include +#ifndef __HIP_PLATFORM_HCC__ +#include +#endif +#include +#include +#include "generic_scaled_masked_softmax.h" +#include "type_shim.h" + +namespace multihead_attn { +namespace fused_softmax { +namespace generic_scaled_masked_softmax { + +torch::Tensor fwd_cuda( + torch::Tensor const& input, + torch::Tensor const& mask, + float scale_factor) +{ + // input is a 4d tensor with dimensions [batches, attn_heads, seq_len, seq_len] + const int batches = input.size(0); + const int pad_batches = mask.size(0); + const int attn_heads = input.size(1); + const int query_seq_len = input.size(2); + const int key_seq_len = input.size(3); + TORCH_INTERNAL_ASSERT(pad_batches == 1 || pad_batches == batches); + TORCH_INTERNAL_ASSERT(mask.size(1) == 1); + TORCH_INTERNAL_ASSERT(mask.size(2) == query_seq_len); + TORCH_INTERNAL_ASSERT(mask.size(3) == key_seq_len); + + // Output + auto act_options = input.options().requires_grad(false); + torch::Tensor softmax_results = + torch::empty({batches, attn_heads, query_seq_len, key_seq_len}, act_options); + + // Softmax Intermediate Result Ptr + void* input_ptr = static_cast(input.data_ptr()); + void* mask_ptr = static_cast(mask.data_ptr()); + void* softmax_results_ptr = static_cast(softmax_results.data_ptr()); + + DISPATCH_HALF_AND_BFLOAT( + input.scalar_type(), + "dispatch_scaled_masked_softmax_forward", + dispatch_scaled_masked_softmax_forward_new( + reinterpret_cast(softmax_results_ptr), + reinterpret_cast(input_ptr), + reinterpret_cast(mask_ptr), + scale_factor, + query_seq_len, + key_seq_len, + batches, + attn_heads, + pad_batches); + ); + return softmax_results; +} + +torch::Tensor bwd_cuda( + torch::Tensor const& output_grads_, + torch::Tensor const& softmax_results_, + float scale_factor) { + + auto output_grads = output_grads_.contiguous(); + auto softmax_results = softmax_results_.contiguous(); + + //output grads is a 4d tensor with dimensions [batches, attn_heads, seq_len, seq_len] + const int batches = output_grads.size(0); + const int attn_heads = output_grads.size(1); + const int query_seq_len = output_grads.size(2); + const int key_seq_len = output_grads.size(3); + + auto act_options = output_grads.options(); + torch::Tensor input_grad = + torch::empty({batches, attn_heads, query_seq_len, key_seq_len}, act_options); + + void* output_grads_ptr = static_cast(output_grads.data_ptr()); + + //Softmax Grad + DISPATCH_HALF_AND_BFLOAT( + output_grads_.scalar_type(), + "dispatch_scaled_masked_softmax_backward", + dispatch_scaled_masked_softmax_backward_new( + reinterpret_cast(static_cast(input_grad.data_ptr())), + reinterpret_cast(output_grads_ptr), + reinterpret_cast(softmax_results.data_ptr()), + scale_factor, + query_seq_len, + key_seq_len, + batches, + attn_heads); + ); + + //backward pass is completely in-place + return input_grad; +} +} +} +} diff --git a/csrc/megatron/scaled_masked_softmax_cuda.cu b/csrc/megatron/scaled_masked_softmax_cuda.cu index 60966706b..ac534a072 100644 --- a/csrc/megatron/scaled_masked_softmax_cuda.cu +++ b/csrc/megatron/scaled_masked_softmax_cuda.cu @@ -18,7 +18,9 @@ #include #include #include -//#include +#ifndef __HIP_PLATFORM_HCC__ +#include +#endif #include #include #include "scaled_masked_softmax.h" diff --git a/csrc/megatron/scaled_upper_triang_masked_softmax_cuda.cu b/csrc/megatron/scaled_upper_triang_masked_softmax_cuda.cu index df022cbbf..c4cbbf17c 100644 --- a/csrc/megatron/scaled_upper_triang_masked_softmax_cuda.cu +++ b/csrc/megatron/scaled_upper_triang_masked_softmax_cuda.cu @@ -18,7 +18,9 @@ #include #include #include -//#include +#ifndef __HIP_PLATFORM_HCC__ +#include +#endif #include #include #include "scaled_upper_triang_masked_softmax.h" diff --git a/setup.py b/setup.py index d86a601a4..2f033c0d8 100644 --- a/setup.py +++ b/setup.py @@ -292,6 +292,13 @@ def check_if_rocm_pytorch(): include_dirs=[os.path.join(this_dir, 'csrc')], extra_compile_args={'cxx': ['-O3'] + version_dependent_macros, 'nvcc':nvcc_args_transformer if not IS_ROCM_PYTORCH else hipcc_args_transformer})) + ext_modules.append( + CUDAExtension(name='generic_scaled_masked_softmax_cuda', + sources=['csrc/megatron/generic_scaled_masked_softmax.cpp', + 'csrc/megatron/generic_scaled_masked_softmax_cuda.cu'], + include_dirs=[os.path.join(this_dir, "csrc")], + extra_compile_args={'cxx': ['-O3'] + version_dependent_macros, + 'nvcc':nvcc_args_transformer if not IS_ROCM_PYTORCH else hipcc_args_transformer})) ext_modules.append( CUDAExtension(name='scaled_masked_softmax_cuda', sources=['csrc/megatron/scaled_masked_softmax.cpp', @@ -555,7 +562,7 @@ def check_if_rocm_pytorch(): if "--transducer" in sys.argv or "--cuda_ext" in sys.argv: if "--transducer" in sys.argv: sys.argv.remove("--transducer") - + if not IS_ROCM_PYTORCH: raise_if_cuda_home_none("--transducer") diff --git a/tests/L0/run_transformer/test_fused_softmax.py b/tests/L0/run_transformer/test_fused_softmax.py index 278df69f7..721dd291c 100644 --- a/tests/L0/run_transformer/test_fused_softmax.py +++ b/tests/L0/run_transformer/test_fused_softmax.py @@ -14,6 +14,14 @@ def attention_mask_func(attention_scores, attention_mask): return attention_scores.masked_fill(attention_mask, -10000.0) +def forward_torch_softmax(input, mask, scale): + input = input * scale + mask_output = attention_mask_func(input, mask) if mask is not None else input + probs = torch.nn.Softmax(dim=-1)(mask_output) + all_k_masked = mask.all(axis=-1) + zero_attention_mask = (1.0 - all_k_masked.float())[:, :, :, None] + probs = probs * zero_attention_mask + return probs autocast_dtypes = ( (torch.half, torch.bfloat16) if torch.cuda.is_bf16_supported() else (torch.half,) @@ -216,5 +224,69 @@ def test_autocast_fused_upper_triangle_mask_softmax(self): actual.backward(g0) expected.backward(g1) + +class TestGenericFusedSoftmaxKernel(common_utils.TestCase): + + def _setup_qk(self): + qlen = [1, 2, 1234, 2322, 2348] + klen = [1, 2, 3, 4, 5, 8, 10, 11, 13, 128, 256, 1200, 1234, 2048, 3123, 4096, 4128, 7234, 8192] + return itertools.product(qlen, klen) + + def _setup_batch_attn_scalar(self): + batch = 2 + attn = 16 + scale_t = 1.0 + return batch, attn, scale_t + + def test_forward(self): + import generic_scaled_masked_softmax_cuda + batch, attn, scale_t = self._setup_batch_attn_scalar() + for qlen, klen in self._setup_qk(): + inputs = torch.normal(0, 2, (batch, attn, qlen, klen), dtype=torch.float16, device='cuda:0') + masks = torch.randint(0, 2, (batch, 1, qlen, klen), dtype=torch.bool, device='cuda:0') + softmax_results = generic_scaled_masked_softmax_cuda.forward(inputs, masks, scale_t) + softmax_results_torch = forward_torch_softmax(inputs, masks, scale_t) + self.assertEqual(softmax_results_torch.type(torch.float16), softmax_results, atol=1e-3, rtol=1e-3) + + def test_backward(self): + import generic_scaled_masked_softmax_cuda + batch, attn, scale_t = self._setup_batch_attn_scalar() + for qlen, klen in self._setup_qk(): + inputs = torch.normal(0, 2, (batch, attn, qlen, klen), dtype=torch.float16, device='cuda:0') + backward = torch.rand_like(inputs, dtype=torch.float16, device='cuda:0') + masks = torch.randint(0, 2, (batch, 1, qlen, klen), dtype=torch.bool, device='cuda:0') + softmax_results = generic_scaled_masked_softmax_cuda.forward(inputs, masks, scale_t) + back_grad = generic_scaled_masked_softmax_cuda.backward(backward, softmax_results, scale_t) + inputs.requires_grad = True + softmax_results_torch = forward_torch_softmax(inputs, masks, scale_t) + softmax_results_torch.backward(backward) + self.assertEqual(back_grad, inputs.grad, atol=1e-3, rtol=1e-3) + + def test_allmasked(self): + import generic_scaled_masked_softmax_cuda + batch, attn, scale_t = self._setup_batch_attn_scalar() + for qlen, klen in self._setup_qk(): + inputs = torch.normal(0, 2, (batch, attn, qlen, klen), dtype=torch.float16, device='cuda:0') + masks = torch.ones((batch, 1, qlen, klen), dtype=torch.bool, device='cuda:0') + softmax_results = generic_scaled_masked_softmax_cuda.forward(inputs, masks, scale_t) + softmax_results_torch = forward_torch_softmax(inputs, masks, scale_t) + self.assertEqual(softmax_results_torch.type(torch.float16), softmax_results, atol=1e-3, rtol=1e-3) + + def test_allmask_backward(self): + import generic_scaled_masked_softmax_cuda + batch, attn, scale_t = self._setup_batch_attn_scalar() + for qlen, klen in self._setup_qk(): + inputs = torch.normal(0, 2, (batch, attn, qlen, klen), dtype=torch.float16, device='cuda:0') + backward = torch.rand_like(inputs, dtype=torch.float16, device='cuda:0') + masks = torch.ones((batch, 1, qlen, klen), dtype=torch.bool, device='cuda:0') + softmax_results = generic_scaled_masked_softmax_cuda.forward(inputs, masks, scale_t) + back_grad = generic_scaled_masked_softmax_cuda.backward(backward, softmax_results, scale_t) + + inputs.requires_grad = True + softmax_results_torch = forward_torch_softmax(inputs, masks, scale_t) + softmax_results_torch.backward(backward) + self.assertEqual(back_grad, inputs.grad, atol=1e-3, rtol=1e-3) + + if __name__ == "__main__": common_utils.run_tests() From c1174a82f01b90eab0a15b142c4212566ba7564d Mon Sep 17 00:00:00 2001 From: Masaki Kozuki Date: Tue, 9 Aug 2022 10:52:00 -0700 Subject: [PATCH 02/25] [transformer] unittest: less mem consumption for generic softmax (#1448) * less mem consumption by fused generic softmax tests ran with RTX 3070 Ti Signed-off-by: Masaki Kozuki * Deduplicate qlen of 1234 --- .../L0/run_transformer/test_fused_softmax.py | 94 +++++++++---------- 1 file changed, 43 insertions(+), 51 deletions(-) diff --git a/tests/L0/run_transformer/test_fused_softmax.py b/tests/L0/run_transformer/test_fused_softmax.py index 721dd291c..a613e51f2 100644 --- a/tests/L0/run_transformer/test_fused_softmax.py +++ b/tests/L0/run_transformer/test_fused_softmax.py @@ -227,65 +227,57 @@ def test_autocast_fused_upper_triangle_mask_softmax(self): class TestGenericFusedSoftmaxKernel(common_utils.TestCase): - def _setup_qk(self): - qlen = [1, 2, 1234, 2322, 2348] - klen = [1, 2, 3, 4, 5, 8, 10, 11, 13, 128, 256, 1200, 1234, 2048, 3123, 4096, 4128, 7234, 8192] - return itertools.product(qlen, klen) - - def _setup_batch_attn_scalar(self): - batch = 2 - attn = 16 - scale_t = 1.0 - return batch, attn, scale_t - - def test_forward(self): + def setUp(self): + super().setUp() + self.batch = 2 + self.attn = 16 + self.scale_t = 1.0 + self.dtype = torch.float16 + self.device = torch.cuda.current_device() + self.thresh = {"atol": 1e-3, "rtol": 1e-3} + + qlen = [1, 2] + klen = [1, 2, 3, 4, 5, 8, 10, 11, 13, 128, 256, 1200, 1234] + available_cuda_mem = torch.cuda.memory.mem_get_info(self.device)[0] / (1024 ** 3) + if available_cuda_mem > 40: + qlen.extend([1234, 2322, 2348]) + klen.extend([2048, 3123, 4096, 4128, 7234, 8192]) + + self.q_k_lens = itertools.product(qlen, klen) + + def test_forward(self, allmasked: bool=False): import generic_scaled_masked_softmax_cuda - batch, attn, scale_t = self._setup_batch_attn_scalar() - for qlen, klen in self._setup_qk(): - inputs = torch.normal(0, 2, (batch, attn, qlen, klen), dtype=torch.float16, device='cuda:0') - masks = torch.randint(0, 2, (batch, 1, qlen, klen), dtype=torch.bool, device='cuda:0') - softmax_results = generic_scaled_masked_softmax_cuda.forward(inputs, masks, scale_t) - softmax_results_torch = forward_torch_softmax(inputs, masks, scale_t) - self.assertEqual(softmax_results_torch.type(torch.float16), softmax_results, atol=1e-3, rtol=1e-3) - - def test_backward(self): + for qlen, klen in self.q_k_lens: + inputs = torch.normal(0, 2, (self.batch, self.attn, qlen, klen), dtype=self.dtype, device=self.device) + masks = ( + torch.randint(0, 2, (self.batch, 1, qlen, klen), dtype=torch.bool, device=self.device) + if not allmasked else torch.ones((self.batch, 1, qlen, klen), dtype=torch.bool, device=self.device) + ) + softmax_results = generic_scaled_masked_softmax_cuda.forward(inputs, masks, self.scale_t) + softmax_results_torch = forward_torch_softmax(inputs, masks, self.scale_t) + self.assertEqual(softmax_results_torch.to(self.dtype), softmax_results, **self.thresh) + + def test_backward(self, allmasked: bool=False): import generic_scaled_masked_softmax_cuda - batch, attn, scale_t = self._setup_batch_attn_scalar() - for qlen, klen in self._setup_qk(): - inputs = torch.normal(0, 2, (batch, attn, qlen, klen), dtype=torch.float16, device='cuda:0') - backward = torch.rand_like(inputs, dtype=torch.float16, device='cuda:0') - masks = torch.randint(0, 2, (batch, 1, qlen, klen), dtype=torch.bool, device='cuda:0') - softmax_results = generic_scaled_masked_softmax_cuda.forward(inputs, masks, scale_t) - back_grad = generic_scaled_masked_softmax_cuda.backward(backward, softmax_results, scale_t) + for qlen, klen in self.q_k_lens: + inputs = torch.normal(0, 2, (self.batch, self.attn, qlen, klen), dtype=self.dtype, device=self.device) + backward = torch.rand_like(inputs, dtype=torch.float16, device=self.device) + masks = ( + torch.randint(0, 2, (self.batch, 1, qlen, klen), dtype=torch.bool, device=self.device) + if not allmasked else torch.ones((self.batch, 1, qlen, klen), dtype=torch.bool, device=self.device) + ) + softmax_results = generic_scaled_masked_softmax_cuda.forward(inputs, masks, self.scale_t) + back_grad = generic_scaled_masked_softmax_cuda.backward(backward, softmax_results, self.scale_t) inputs.requires_grad = True - softmax_results_torch = forward_torch_softmax(inputs, masks, scale_t) + softmax_results_torch = forward_torch_softmax(inputs, masks, self.scale_t) softmax_results_torch.backward(backward) - self.assertEqual(back_grad, inputs.grad, atol=1e-3, rtol=1e-3) + self.assertEqual(back_grad, inputs.grad, **self.thresh) def test_allmasked(self): - import generic_scaled_masked_softmax_cuda - batch, attn, scale_t = self._setup_batch_attn_scalar() - for qlen, klen in self._setup_qk(): - inputs = torch.normal(0, 2, (batch, attn, qlen, klen), dtype=torch.float16, device='cuda:0') - masks = torch.ones((batch, 1, qlen, klen), dtype=torch.bool, device='cuda:0') - softmax_results = generic_scaled_masked_softmax_cuda.forward(inputs, masks, scale_t) - softmax_results_torch = forward_torch_softmax(inputs, masks, scale_t) - self.assertEqual(softmax_results_torch.type(torch.float16), softmax_results, atol=1e-3, rtol=1e-3) + self.test_forward(True) def test_allmask_backward(self): - import generic_scaled_masked_softmax_cuda - batch, attn, scale_t = self._setup_batch_attn_scalar() - for qlen, klen in self._setup_qk(): - inputs = torch.normal(0, 2, (batch, attn, qlen, klen), dtype=torch.float16, device='cuda:0') - backward = torch.rand_like(inputs, dtype=torch.float16, device='cuda:0') - masks = torch.ones((batch, 1, qlen, klen), dtype=torch.bool, device='cuda:0') - softmax_results = generic_scaled_masked_softmax_cuda.forward(inputs, masks, scale_t) - back_grad = generic_scaled_masked_softmax_cuda.backward(backward, softmax_results, scale_t) - - inputs.requires_grad = True - softmax_results_torch = forward_torch_softmax(inputs, masks, scale_t) - softmax_results_torch.backward(backward) - self.assertEqual(back_grad, inputs.grad, atol=1e-3, rtol=1e-3) + self.test_backward(True) if __name__ == "__main__": From 13c393e2bc12a530b2a6348eab6a099fd5b53c2f Mon Sep 17 00:00:00 2001 From: Masaki Kozuki Date: Mon, 15 Aug 2022 12:02:33 -0700 Subject: [PATCH 03/25] Use `xmlrunner.XMLTestRunner` accordingly in `tests/L0/run_test.py` (#1451) * Use xmlrunner.XMLTestRunner accordingly TODO: - [x] Remove `subTest` because it's not compatible with the current way of running L0 tests Signed-off-by: Masaki Kozuki * use `torch.testing` more to enable xmlrunner Signed-off-by: Masaki Kozuki * Remove `subTest` for xmlrunner Signed-off-by: Masaki Kozuki * removing subTest Signed-off-by: Masaki Kozuki * not depend on an env var Signed-off-by: Masaki Kozuki * fix syntax errors * open with `"wb"` * xml file per dir Signed-off-by: Masaki Kozuki * remove comment-out Signed-off-by: Masaki Kozuki * Refactor `TestTransformer`: define member methods (#5) * setUpClass to define `test_` methods Signed-off-by: Masaki Kozuki * manually define Signed-off-by: Masaki Kozuki Signed-off-by: Masaki Kozuki * add a missing test Signed-off-by: Masaki Kozuki * remove print Signed-off-by: Masaki Kozuki * remove ext Signed-off-by: Masaki Kozuki Signed-off-by: Masaki Kozuki --- .../test_fused_layer_norm.py | 44 +- tests/L0/run_mlp/test_mlp.py | 58 +- tests/L0/run_test.py | 29 +- .../L0/run_transformer/test_batch_sampler.py | 37 +- .../L0/run_transformer/test_cross_entropy.py | 45 +- .../L0/run_transformer/test_fused_softmax.py | 263 +++--- tests/L0/run_transformer/test_layers.py | 750 +++++++++--------- tests/L0/run_transformer/test_mapping.py | 96 ++- tests/L0/run_transformer/test_microbatches.py | 58 +- .../L0/run_transformer/test_parallel_state.py | 148 ++-- tests/L0/run_transformer/test_random.py | 131 ++- .../test_transformer_module.py | 115 ++- .../run_transformer/test_transformer_utils.py | 34 +- 13 files changed, 933 insertions(+), 875 deletions(-) diff --git a/tests/L0/run_fused_layer_norm/test_fused_layer_norm.py b/tests/L0/run_fused_layer_norm/test_fused_layer_norm.py index 18219522c..a3826253a 100644 --- a/tests/L0/run_fused_layer_norm/test_fused_layer_norm.py +++ b/tests/L0/run_fused_layer_norm/test_fused_layer_norm.py @@ -3,8 +3,12 @@ import torch -import apex from apex.testing.common_utils import skipFlakyTest +from apex.normalization import FusedLayerNorm +from apex.normalization import FusedRMSNorm +from apex.normalization import MixedFusedLayerNorm +from apex.normalization import MixedFusedRMSNorm + class TestFusedLayerNorm(unittest.TestCase): dtype = torch.float @@ -18,15 +22,15 @@ class TestFusedLayerNorm(unittest.TestCase): def setUp(self): # bias and weight are set to 0 and 1 respectively, so no need to copy parameters from cpu module to the gpu one if not self.mixed_fused: - self.module_cpu_ = apex.normalization.FusedLayerNorm( + self.module_cpu_ = FusedLayerNorm( normalized_shape=self.normalized_shape, elementwise_affine=self.elementwise_affine).cpu() - self.module_cuda_ = apex.normalization.FusedLayerNorm( + self.module_cuda_ = FusedLayerNorm( normalized_shape=self.normalized_shape, elementwise_affine=self.elementwise_affine).to(device="cuda", dtype=self.dtype) else: assert self.elementwise_affine - self.module_cpu_ = apex.normalization.MixedFusedLayerNorm( + self.module_cpu_ = MixedFusedLayerNorm( normalized_shape=self.normalized_shape).cpu() - self.module_cuda_ = apex.normalization.MixedFusedLayerNorm( + self.module_cuda_ = MixedFusedLayerNorm( normalized_shape=self.normalized_shape).to(device="cuda", dtype=self.dtype) @@ -65,8 +69,7 @@ def _check_same_output(self, batch_size, contiguous): def _test_same_output(self, batch_size): for contiguous in (True, False): - with self.subTest(contiguous=contiguous): - self._check_same_output(batch_size, contiguous) + self._check_same_output(batch_size, contiguous) def test_layer_norm(self): self._test_same_output(16) @@ -87,15 +90,15 @@ class TestFusedRMSNorm(unittest.TestCase): def setUp(self): # bias and weight are set to 0 and 1 respectively, so no need to copy parameters from cpu module to the gpu one if not self.mixed_fused: - self.module_cpu_ = apex.normalization.FusedRMSNorm( + self.module_cpu_ = FusedRMSNorm( normalized_shape=self.normalized_shape, elementwise_affine=self.elementwise_affine).cpu() - self.module_cuda_ = apex.normalization.FusedRMSNorm( + self.module_cuda_ = FusedRMSNorm( normalized_shape=self.normalized_shape, elementwise_affine=self.elementwise_affine).to(device="cuda", dtype=self.dtype) else: assert self.elementwise_affine - self.module_cpu_ = apex.normalization.MixedFusedRMSNorm( + self.module_cpu_ = MixedFusedRMSNorm( normalized_shape=self.normalized_shape).cpu() - self.module_cuda_ = apex.normalization.MixedFusedRMSNorm( + self.module_cuda_ = MixedFusedRMSNorm( normalized_shape=self.normalized_shape).to(device="cuda", dtype=self.dtype) def _check_same_output(self, batch_size, contiguous): @@ -136,8 +139,7 @@ def _check_same_output(self, batch_size, contiguous): def _test_same_output(self, batch_size): for contiguous in (True, False): - with self.subTest(contiguous=contiguous): - self._check_same_output(batch_size, contiguous) + self._check_same_output(batch_size, contiguous) def test_layer_norm(self): self._test_same_output(16) @@ -206,17 +208,17 @@ def _prep_layers(normalized_shape, elementwise_affine, dtype): native = torch.nn.LayerNorm( normalized_shape=normalized_shape, elementwise_affine=elementwise_affine ).to(device="cuda", dtype=dtype) - fused = apex.normalization.FusedLayerNorm( + fused = FusedLayerNorm( normalized_shape=normalized_shape, elementwise_affine=elementwise_affine ).cuda() return native, fused def _prep_rms_layers(normalized_shape, elementwise_affine, dtype): - native = apex.normalization.FusedRMSNorm( + native = FusedRMSNorm( normalized_shape=normalized_shape, elementwise_affine=elementwise_affine ) - fused = apex.normalization.FusedRMSNorm( + fused = FusedRMSNorm( normalized_shape=normalized_shape, elementwise_affine=elementwise_affine ).cuda() return native, fused @@ -261,8 +263,7 @@ def _run_test(self, dtype, elementwise_affine): def test_autocast(self): for (dtype, elementwise_affine) in itertools.product(autocast_dtypes, (True, False)): - with self.subTest(f"{dtype}-{elementwise_affine}"): - self._run_test(dtype, elementwise_affine) + self._run_test(dtype, elementwise_affine) @unittest.skip("Skipped on ROCm5.2 due to the failure of reproducing the issue locally. (Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!) Please refer to https://github.com/ROCmSoftwarePlatform/apex/pull/78") class TestAutocastFusedRMSNorm(unittest.TestCase): @@ -294,5 +295,8 @@ def _run_test(self, dtype, elementwise_affine): def test_autocast(self): for (dtype, elementwise_affine) in itertools.product(autocast_dtypes, (True, False)): - with self.subTest(f"{dtype}-{elementwise_affine}"): - self._run_test(dtype, elementwise_affine) + self._run_test(dtype, elementwise_affine) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/L0/run_mlp/test_mlp.py b/tests/L0/run_mlp/test_mlp.py index 615dec95c..b28b7e791 100644 --- a/tests/L0/run_mlp/test_mlp.py +++ b/tests/L0/run_mlp/test_mlp.py @@ -1,10 +1,13 @@ """Tests for c++ MLP""" -import unittest +from itertools import product from time import time import numpy as np import torch from torch import nn +from torch.testing._internal import common_utils +from torch.testing._internal.common_device_type import instantiate_device_type_tests +from torch.testing._internal.common_device_type import onlyCUDA from apex.mlp import MLP from apex.testing.common_utils import skipFlakyTest @@ -25,10 +28,51 @@ def test_numeric(self): mlp_layers = [] for i in range(mlp.num_layers): linear = nn.Linear(mlp_sizes[i], mlp_sizes[i + 1]) - mlp.weights[i].data.copy_(linear.weight) - mlp.biases[i].data.copy_(linear.bias) + with torch.no_grad(): + mlp.weights[i].copy_(linear.weight) + mlp.biases[i].copy_(linear.bias) mlp_layers.append(linear) - mlp_layers.append(nn.ReLU(inplace=True)) + mlp_layers.append(nn.ReLU()) + + ref_mlp = nn.Sequential(*mlp_layers).cuda() + + test_input = ( + torch.empty(batch_size, mlp_sizes[0], device="cuda") + .uniform_(-1.0, 1.0) + .requires_grad_() + ) + ref_input = test_input.clone().detach().requires_grad_() + mlp_out = mlp(test_input) + ref_out = ref_mlp(ref_input) + self.assertEqual(mlp_out, ref_out) + + # Use mean value as scalar loss. Multiply 10 to make it big enough not zero out + mlp_out.mean().mul(10.0).backward() + ref_out.mean().mul(10.0).backward() + self.assertEqual(test_input.grad, ref_input.grad) + self.assertEqual(mlp.biases[0].grad, ref_mlp[0].bias.grad) + + @common_utils.parametrize( + "use_activation,bias", + list(product(("none", "relu", "sigmoid"), (True, False))), + ) + def test_mlp(self, use_activation: str, bias: bool): + # for use_activation in ["none", "relu", "sigmoid"]: + msg = f"activation: {use_activation}, bias: {bias}" + mlp = MLP(mlp_sizes, bias=bias, activation=use_activation).cuda() + + mlp_layers = [] + for i in range(mlp.num_layers): + linear = nn.Linear(mlp_sizes[i], mlp_sizes[i + 1], bias=bias) + with torch.no_grad(): + mlp.weights[i].copy_(linear.weight) + if bias: + mlp.biases[i].copy_(linear.bias) + mlp_layers.append(linear) + if use_activation == "relu": + mlp_layers.append(nn.ReLU()) + if use_activation == "sigmoid": + mlp_layers.append(nn.Sigmoid()) ref_mlp = nn.Sequential(*mlp_layers).cuda() @@ -141,8 +185,9 @@ def test_no_grad(self): mlp_layers = [] for i in range(mlp.num_layers): linear = nn.Linear(mlp_sizes[i], mlp_sizes[i + 1]) - mlp.weights[i].data.copy_(linear.weight) - mlp.biases[i].data.copy_(linear.bias) + with torch.no_grad(): + mlp.weights[i].copy_(linear.weight) + mlp.biases[i].copy_(linear.bias) mlp_layers.append(linear) mlp_layers.append(nn.ReLU(inplace=True)) @@ -220,3 +265,4 @@ def test_performance_half(self): if __name__ == '__main__': unittest.main() + diff --git a/tests/L0/run_test.py b/tests/L0/run_test.py index e87a1e8e9..4b8f5c0ad 100644 --- a/tests/L0/run_test.py +++ b/tests/L0/run_test.py @@ -2,7 +2,9 @@ How to run this script? -1. Run all the tests: `python /path/to/apex/tests/L0/run_test.py` +1. Run all the tests: `python /path/to/apex/tests/L0/run_test.py` If you want an xml report, + pass `--xml-report`, i.e. `python /path/to/apex/tests/L0/run_test.py --xml-report` and + the file is created in `/path/to/apex/tests/L0`. 2. Run one of the tests (e.g. fused layer norm): `python /path/to/apex/tests/L0/run_test.py --include run_fused_layer_norm` 3. Run two or more of the tests (e.g. optimizers and fused layer norm): @@ -46,16 +48,35 @@ def parse_args(): default=DEFAULT_TEST_DIRS, help="select a set of tests to run (defaults to ALL tests).", ) + parser.add_argument( + "--xml-report", + action="store_true", + help="pass this argument to get a junit xml report. (requires `xmlrunner`)", + ) args, _ = parser.parse_known_args() return args -def main(args): - runner = unittest.TextTestRunner(verbosity=2) +def main(args: argparse.Namespace) -> None: + test_runner_kwargs = {"verbosity": 2} + Runner = unittest.TextTestRunner + if args.xml_report: + import xmlrunner + from datetime import date # NOQA + Runner = xmlrunner.XMLTestRunner + errcode = 0 for test_dir in args.include: + if args.xml_report: + this_dir = os.path.abspath(os.path.dirname(__file__)) + xml_output = os.path.join( + this_dir, + f"""TEST_{test_dir}_{date.today().strftime("%y%m%d")}""", + ) + test_runner_kwargs["output"] = xml_output + + runner = Runner(**test_runner_kwargs) test_dir = os.path.join(TEST_ROOT, test_dir) - print(test_dir) suite = unittest.TestLoader().discover(test_dir) print("\nExecuting tests from " + test_dir) diff --git a/tests/L0/run_transformer/test_batch_sampler.py b/tests/L0/run_transformer/test_batch_sampler.py index 52175d53a..f0fb01407 100644 --- a/tests/L0/run_transformer/test_batch_sampler.py +++ b/tests/L0/run_transformer/test_batch_sampler.py @@ -1,10 +1,6 @@ -from itertools import product - import torch from torch.testing._internal import common_utils from torch.utils.data import Dataset -from torch.utils.data import RandomSampler -from torch.utils.data import BatchSampler from torch.utils.data import DataLoader from apex.transformer.pipeline_parallel.utils import _split_batch_into_microbatch as split_batch_into_microbatch @@ -85,23 +81,22 @@ def test_batch_sampler_behavior(self): dataset = MyIterableDataset(0, 100) for num_workers in (1, 2, 4): - with self.subTest(f"{num_workers}"): - torch.manual_seed(42) - loader = DataLoader(dataset, batch_sampler=MegatronPretrainingRandomSampler(100, 0, 4, 0, 1), num_workers=num_workers) - samples = [] - for i, batch in enumerate(loader): - samples.append(batch) - if i == 2 - 1: - break - - torch.manual_seed(42) - loader = DataLoader(dataset, batch_sampler=MegatronPretrainingRandomSampler(100, 0, 2, 0, 1), num_workers=num_workers) - samples2 = [] - for i, batch in enumerate(loader): - samples2.append(batch) - if i == 4 - 1: - break - self.assertEqual(torch.cat(samples), torch.cat(samples2)) + torch.manual_seed(42) + loader = DataLoader(dataset, batch_sampler=MegatronPretrainingRandomSampler(100, 0, 4, 0, 1), num_workers=num_workers) + samples = [] + for i, batch in enumerate(loader): + samples.append(batch) + if i == 2 - 1: + break + + torch.manual_seed(42) + loader = DataLoader(dataset, batch_sampler=MegatronPretrainingRandomSampler(100, 0, 2, 0, 1), num_workers=num_workers) + samples2 = [] + for i, batch in enumerate(loader): + samples2.append(batch) + if i == 4 - 1: + break + self.assertEqual(torch.cat(samples), torch.cat(samples2), msg=f"num_workers={num_workers}") def test_split_batch(self): diff --git a/tests/L0/run_transformer/test_cross_entropy.py b/tests/L0/run_transformer/test_cross_entropy.py index 1f5162876..657ecbf6d 100644 --- a/tests/L0/run_transformer/test_cross_entropy.py +++ b/tests/L0/run_transformer/test_cross_entropy.py @@ -63,27 +63,30 @@ def test_cross_entropy(self): for tensor_model_parallel_world_size in range(1, self.world_size + 1): if self.world_size % tensor_model_parallel_world_size: continue - with self.subTest( - tensor_model_parallel_world_size=tensor_model_parallel_world_size - ): - parallel_state.initialize_model_parallel( - tensor_model_parallel_size_=tensor_model_parallel_world_size, - ) - vocab_size = vocab_size_per_partition * tensor_model_parallel_world_size - loss_torch, grad_torch = torch_cross_entropy( - batch_size, sequence_length, vocab_size, logits_scale, seed - ) - ( - loss_tensor_parallel, - grad_tensor_parallel, - ) = tensor_sharded_cross_entropy( - batch_size, sequence_length, vocab_size, logits_scale, seed - ) - - self.assertEqual(loss_torch, loss_tensor_parallel) - self.assertEqual(grad_torch, grad_tensor_parallel) - - parallel_state.destroy_model_parallel() + parallel_state.initialize_model_parallel( + tensor_model_parallel_size_=tensor_model_parallel_world_size, + ) + vocab_size = vocab_size_per_partition * tensor_model_parallel_world_size + loss_torch, grad_torch = torch_cross_entropy( + batch_size, sequence_length, vocab_size, logits_scale, seed + ) + ( + loss_tensor_parallel, + grad_tensor_parallel, + ) = tensor_sharded_cross_entropy( + batch_size, sequence_length, vocab_size, logits_scale, seed + ) + + self.assertEqual( + loss_torch, loss_tensor_parallel, + msg=f"tensor_model_parallel_size: {tensor_model_parallel_world_size}", + ) + self.assertEqual( + grad_torch, grad_tensor_parallel, + msg=f"tensor_model_parallel_size: {tensor_model_parallel_world_size}", + ) + + parallel_state.destroy_model_parallel() class NcclVocabParallelCrossEntropyTest(VocabParallelCrossEntropyTestBase, NcclDistributedTestBase): pass diff --git a/tests/L0/run_transformer/test_fused_softmax.py b/tests/L0/run_transformer/test_fused_softmax.py index a613e51f2..015f997db 100644 --- a/tests/L0/run_transformer/test_fused_softmax.py +++ b/tests/L0/run_transformer/test_fused_softmax.py @@ -65,75 +65,75 @@ def test_fused_scale_mask_softmax(self): for (dtype, scale, softmax_in_fp32, shape) in itertools.product( (torch.half, torch.bfloat16), (None, 2.0), (False, True), ((4, 12, 24, 24), (32, 12, 4, 214)) ): - with self.subTest(f"{dtype}-{scale}-{softmax_in_fp32}"): - input_in_fp16 = dtype == torch.half - input_in_bf16 = dtype == torch.bfloat16 - if not (scale is None or softmax_in_fp32): - with self.assertRaises(RuntimeError): - self._setup_fused_softmax( - input_in_fp16, - input_in_bf16, - scale, - softmax_in_fp32, - AttnMaskType.padding, - ) - return - fused_fn, torch_fn = self._setup_fused_softmax( - input_in_fp16, - input_in_bf16, - scale, - softmax_in_fp32, - AttnMaskType.padding, - ) + msg = f"{dtype}-{scale}-{softmax_in_fp32}" + input_in_fp16 = dtype == torch.half + input_in_bf16 = dtype == torch.bfloat16 + if not (scale is None or softmax_in_fp32): + with self.assertRaises(RuntimeError, msg=msg): + self._setup_fused_softmax( + input_in_fp16, + input_in_bf16, + scale, + softmax_in_fp32, + AttnMaskType.padding, + ) + return + fused_fn, torch_fn = self._setup_fused_softmax( + input_in_fp16, + input_in_bf16, + scale, + softmax_in_fp32, + AttnMaskType.padding, + ) - attention_scores_0 = ( - torch.randn(shape) - .to(device="cuda", dtype=dtype) - .requires_grad_(True) - ) - with torch.no_grad(): - attention_scores_1 = attention_scores_0.clone().requires_grad_(True) - mask_shape = (shape[0],) + (1,) + shape[2:] - mask = torch.randint(0, 2, mask_shape, device="cuda").bool() - expected = fused_fn(attention_scores_0, mask) - actual = torch_fn(attention_scores_1, mask) - self.assertEqual(actual, expected) - - g0 = torch.rand_like(actual) - with torch.no_grad(): - g1 = g0.clone() - expected.backward(g0) - actual.backward(g1) + attention_scores_0 = ( + torch.randn(shape) + .to(device="cuda", dtype=dtype) + .requires_grad_(True) + ) + with torch.no_grad(): + attention_scores_1 = attention_scores_0.clone().requires_grad_(True) + mask_shape = (shape[0],) + (1,) + shape[2:] + mask = torch.randint(0, 2, mask_shape, device="cuda").bool() + expected = fused_fn(attention_scores_0, mask) + actual = torch_fn(attention_scores_1, mask) + self.assertEqual(actual, expected, msg=msg) + + g0 = torch.rand_like(actual) + with torch.no_grad(): + g1 = g0.clone() + expected.backward(g0) + actual.backward(g1) def test_autocast_fused_scale_mask_softmax(self): for dtype in autocast_dtypes: - with self.subTest(f"{dtype}"): - input_in_fp16 = dtype == torch.half - input_in_bf16 = dtype == torch.bfloat16 - fused_fn, torch_fn = self._setup_fused_softmax( - input_in_fp16, input_in_bf16, attn_mask_type=AttnMaskType.padding - ) + msg = f"dtype: {dtype}" + input_in_fp16 = dtype == torch.half + input_in_bf16 = dtype == torch.bfloat16 + fused_fn, torch_fn = self._setup_fused_softmax( + input_in_fp16, input_in_bf16, attn_mask_type=AttnMaskType.padding + ) - attention_scores_0 = ( - torch.randn((4, 12, 24, 24)).cuda().requires_grad_(True) + attention_scores_0 = ( + torch.randn((4, 12, 24, 24)).cuda().requires_grad_(True) + ) + with torch.no_grad(): + attention_scores_1 = ( + attention_scores_0.clone().to(dtype).requires_grad_(True) ) - with torch.no_grad(): - attention_scores_1 = ( - attention_scores_0.clone().to(dtype).requires_grad_(True) - ) - mask = torch.randint(0, 2, (4, 1, 24, 24)).bool().cuda() + mask = torch.randint(0, 2, (4, 1, 24, 24)).bool().cuda() - expected = torch_fn(attention_scores_1, mask) - with torch.cuda.amp.autocast(dtype=dtype): - actual = fused_fn(attention_scores_0, mask) - self.assertEqual(actual.dtype, dtype) - self.assertEqual(actual, expected) + expected = torch_fn(attention_scores_1, mask) + with torch.cuda.amp.autocast(dtype=dtype): + actual = fused_fn(attention_scores_0, mask) + self.assertEqual(actual.dtype, dtype, msg=msg) + self.assertEqual(actual, expected, msg=msg) - g0 = torch.rand_like(actual) - with torch.no_grad(): - g1 = g0.clone() - expected.backward(g0) - actual.backward(g1) + g0 = torch.rand_like(actual) + with torch.no_grad(): + g1 = g0.clone() + expected.backward(g0) + actual.backward(g1) def test_fused_upper_triangle_mask_softmax(self): """ @@ -146,83 +146,83 @@ def test_fused_upper_triangle_mask_softmax(self): for (dtype, scale, softmax_in_fp32) in itertools.product( (torch.half, torch.bfloat16), (None, 2.0), (False, True), ): - with self.subTest(f"{dtype}-{scale}-{softmax_in_fp32}"): - input_in_fp16 = dtype == torch.half - input_in_bf16 = dtype == torch.bfloat16 - if not (scale is None or softmax_in_fp32): - with self.assertRaises(RuntimeError): - self._setup_fused_softmax( - input_in_fp16, - input_in_bf16, - scale, - softmax_in_fp32, - AttnMaskType.causal, - ) - return - fused_fn, torch_fn = self._setup_fused_softmax( - input_in_fp16, - input_in_bf16, - scale, - softmax_in_fp32, - AttnMaskType.causal, - ) + msg = f"{dtype}-{scale}-{softmax_in_fp32}" + input_in_fp16 = dtype == torch.half + input_in_bf16 = dtype == torch.bfloat16 + if not (scale is None or softmax_in_fp32): + with self.assertRaises(RuntimeError, msg=msg): + self._setup_fused_softmax( + input_in_fp16, + input_in_bf16, + scale, + softmax_in_fp32, + AttnMaskType.causal, + ) + return + fused_fn, torch_fn = self._setup_fused_softmax( + input_in_fp16, + input_in_bf16, + scale, + softmax_in_fp32, + AttnMaskType.causal, + ) - attn_weights_0 = ( - torch.randn((4, 12, 24, 24)) - .to(device="cuda", dtype=dtype) - .requires_grad_(True) - ) - with torch.no_grad(): - attn_weights_1 = attn_weights_0.clone().requires_grad_(True) - total_mask = ( - ~(torch.tril(torch.randn((24, 24), device="cuda")).bool()) - .unsqueeze(0) - .unsqueeze(0) - ) - total_mask = total_mask.repeat((4, 1, 1, 1)) - expected = fused_fn(attn_weights_0, total_mask) - actual = torch_fn(attn_weights_1, total_mask) - self.assertEqual(actual, expected) + attn_weights_0 = ( + torch.randn((4, 12, 24, 24)) + .to(device="cuda", dtype=dtype) + .requires_grad_(True) + ) + with torch.no_grad(): + attn_weights_1 = attn_weights_0.clone().requires_grad_(True) + total_mask = ( + ~(torch.tril(torch.randn((24, 24), device="cuda")).bool()) + .unsqueeze(0) + .unsqueeze(0) + ) + total_mask = total_mask.repeat((4, 1, 1, 1)) + expected = fused_fn(attn_weights_0, total_mask) + actual = torch_fn(attn_weights_1, total_mask) + self.assertEqual(actual, expected, msg=msg) - g0 = torch.randn_like(actual) - with torch.no_grad(): - g1 = g0.clone() - actual.backward(g0) - expected.backward(g1) + g0 = torch.randn_like(actual) + with torch.no_grad(): + g1 = g0.clone() + actual.backward(g0) + expected.backward(g1) def test_autocast_fused_upper_triangle_mask_softmax(self): for dtype in autocast_dtypes: - with self.subTest(f"{dtype}"): - input_in_fp16 = dtype == torch.half - input_in_bf16 = dtype == torch.bfloat16 - fused_fn, torch_fn = self._setup_fused_softmax( - input_in_fp16, input_in_bf16, attn_mask_type=AttnMaskType.causal - ) + msg = f"dtype: {dtype}" + input_in_fp16 = dtype == torch.half + input_in_bf16 = dtype == torch.bfloat16 + fused_fn, torch_fn = self._setup_fused_softmax( + input_in_fp16, input_in_bf16, attn_mask_type=AttnMaskType.causal + ) - attn_weights_0 = ( - torch.randn((4, 12, 24, 24)).cuda().requires_grad_(True) - ) - with torch.no_grad(): - attn_weights_1 = ( - attn_weights_0.clone().to(dtype).requires_grad_(True) - ) - total_mask = ( - ~(torch.tril(torch.randn((24, 24), device="cuda")).bool()) - .unsqueeze(0) - .unsqueeze(0) + attn_weights_0 = ( + torch.randn((4, 12, 24, 24)).cuda().requires_grad_(True) + ) + with torch.no_grad(): + attn_weights_1 = ( + attn_weights_0.clone().to(dtype).requires_grad_(True) ) + total_mask = ( + ~(torch.tril(torch.randn((24, 24), device="cuda")).bool()) + .unsqueeze(0) + .unsqueeze(0) + ) - with torch.cuda.amp.autocast(dtype=dtype): - actual = fused_fn(attn_weights_0, total_mask) - self.assertEqual(actual.dtype, dtype) - expected = torch_fn(attn_weights_1, total_mask) - self.assertEqual(actual, expected) + with torch.cuda.amp.autocast(dtype=dtype): + actual = fused_fn(attn_weights_0, total_mask) + self.assertEqual(actual.dtype, dtype, msg=msg) + expected = torch_fn(attn_weights_1, total_mask) + self.assertEqual(actual, expected, msg=msg) - g0 = torch.randn_like(actual) - with torch.no_grad(): - g1 = g0.clone() - actual.backward(g0) - expected.backward(g1) + g0 = torch.randn_like(actual) + with torch.no_grad(): + g1 = g0.clone() + actual.backward(g0) + expected.backward(g1) class TestGenericFusedSoftmaxKernel(common_utils.TestCase): @@ -255,7 +255,8 @@ def test_forward(self, allmasked: bool=False): ) softmax_results = generic_scaled_masked_softmax_cuda.forward(inputs, masks, self.scale_t) softmax_results_torch = forward_torch_softmax(inputs, masks, self.scale_t) - self.assertEqual(softmax_results_torch.to(self.dtype), softmax_results, **self.thresh) + self.assertEqual( + softmax_results_torch.to(self.dtype), softmax_results, **self.thresh, msg=f"(q, k) = ({qlen, klen})") def test_backward(self, allmasked: bool=False): import generic_scaled_masked_softmax_cuda @@ -271,7 +272,7 @@ def test_backward(self, allmasked: bool=False): inputs.requires_grad = True softmax_results_torch = forward_torch_softmax(inputs, masks, self.scale_t) softmax_results_torch.backward(backward) - self.assertEqual(back_grad, inputs.grad, **self.thresh) + self.assertEqual(back_grad, inputs.grad, **self.thresh, msg=f"(q, k) = ({qlen, klen})") def test_allmasked(self): self.test_forward(True) diff --git a/tests/L0/run_transformer/test_layers.py b/tests/L0/run_transformer/test_layers.py index b3b2eb2fc..80cdff0f2 100644 --- a/tests/L0/run_transformer/test_layers.py +++ b/tests/L0/run_transformer/test_layers.py @@ -47,43 +47,43 @@ def test_all_gather_parity(self) -> None: for tensor_model_parallel_world_size in range(1, self.world_size + 1): if self.world_size % tensor_model_parallel_world_size: continue - with self.subTest(tensor_model_parallel_world_size=tensor_model_parallel_world_size): - parallel_state.initialize_model_parallel( - tensor_model_parallel_size_=tensor_model_parallel_world_size, - ) - tensor_model_parallel_rank = parallel_state.get_tensor_model_parallel_rank() - cur_tensor_model_device = torch.device(f"cuda:{tensor_model_parallel_rank}") - with torch.no_grad(): - tensor = tensor_model_parallel_rank * torch.ones( - self.tensor_shape, dtype=torch.float32, device=cur_tensor_model_device) - numel = tensor.numel() - numel_gathered = tensor_model_parallel_world_size * numel - gathered = torch.empty( - torch.Size((numel_gathered,)), - device=cur_tensor_model_device, - dtype=torch.float32, - requires_grad=False, - ) - chunks = [ - gathered[i * numel : (i + 1) * numel] - for i in range(tensor_model_parallel_world_size) - ] - all_gather(chunks, tensor, group=parallel_state.get_tensor_model_parallel_group()) - - gathered_for_base = torch.empty( - torch.Size((numel_gathered,)), - device=cur_tensor_model_device, - dtype=torch.float32, - requires_grad=False, - ) - _all_gather_base( - gathered_for_base, - tensor, - group=parallel_state.get_tensor_model_parallel_group(), - ) - - self.assertEqual(gathered, gathered_for_base) - parallel_state.destroy_model_parallel() + parallel_state.initialize_model_parallel( + tensor_model_parallel_size_=tensor_model_parallel_world_size, + ) + tensor_model_parallel_rank = parallel_state.get_tensor_model_parallel_rank() + cur_tensor_model_device = torch.device(f"cuda:{tensor_model_parallel_rank}") + with torch.no_grad(): + tensor = tensor_model_parallel_rank * torch.ones( + self.tensor_shape, dtype=torch.float32, device=cur_tensor_model_device) + numel = tensor.numel() + numel_gathered = tensor_model_parallel_world_size * numel + gathered = torch.empty( + torch.Size((numel_gathered,)), + device=cur_tensor_model_device, + dtype=torch.float32, + requires_grad=False, + ) + chunks = [ + gathered[i * numel : (i + 1) * numel] + for i in range(tensor_model_parallel_world_size) + ] + all_gather(chunks, tensor, group=parallel_state.get_tensor_model_parallel_group()) + + gathered_for_base = torch.empty( + torch.Size((numel_gathered,)), + device=cur_tensor_model_device, + dtype=torch.float32, + requires_grad=False, + ) + _all_gather_base( + gathered_for_base, + tensor, + group=parallel_state.get_tensor_model_parallel_group(), + ) + + msg = f"tensor_model_parallel_world_size: {tensor_model_parallel_world_size}" + self.assertEqual(gathered, gathered_for_base, msg=msg) + parallel_state.destroy_model_parallel() @torch.no_grad() @unittest.skipIf(torch.cuda.device_count() < 2, "Requires >=2 GPUs") @@ -95,113 +95,111 @@ def test_reduce_scatter_parity(self) -> None: for tensor_model_parallel_world_size in range(2, self.world_size + 1): if self.world_size % tensor_model_parallel_world_size: continue - with self.subTest(tensor_model_parallel_world_size=tensor_model_parallel_world_size): - parallel_state.initialize_model_parallel( - tensor_model_parallel_size_=tensor_model_parallel_world_size, - ) - tensor_model_parallel_rank = parallel_state.get_tensor_model_parallel_rank() - cur_tensor_model_device = torch.device(f"cuda:{tensor_model_parallel_rank}") - with torch.no_grad(): - input = torch.cat([ - i * torch.ones(self.tensor_shape, dtype=torch.float32, device=cur_tensor_model_device) - for i in range(tensor_model_parallel_world_size) - ]) - input_list = [t.clone() for t in input.chunk(tensor_model_parallel_world_size)] - output = torch.empty( - self.tensor_shape, - device=cur_tensor_model_device, - dtype=torch.float32, - requires_grad=False, - ) - reduce_scatter( - output, input_list, - group=parallel_state.get_tensor_model_parallel_group(), - ) - - output_for_base = torch.empty( - self.tensor_shape, - device=cur_tensor_model_device, - dtype=torch.float32, - requires_grad=False, - ) - _reduce_scatter_base( - output_for_base, - input, - group=parallel_state.get_tensor_model_parallel_group(), - ) - - self.assertEqual(output, output_for_base) - self.assertEqual(input, torch.cat(input_list)) - parallel_state.destroy_model_parallel() + parallel_state.initialize_model_parallel( + tensor_model_parallel_size_=tensor_model_parallel_world_size, + ) + tensor_model_parallel_rank = parallel_state.get_tensor_model_parallel_rank() + cur_tensor_model_device = torch.device(f"cuda:{tensor_model_parallel_rank}") + with torch.no_grad(): + input = torch.cat([ + i * torch.ones(self.tensor_shape, dtype=torch.float32, device=cur_tensor_model_device) + for i in range(tensor_model_parallel_world_size) + ]) + input_list = [t.clone() for t in input.chunk(tensor_model_parallel_world_size)] + output = torch.empty( + self.tensor_shape, + device=cur_tensor_model_device, + dtype=torch.float32, + requires_grad=False, + ) + reduce_scatter( + output, input_list, + group=parallel_state.get_tensor_model_parallel_group(), + ) + + output_for_base = torch.empty( + self.tensor_shape, + device=cur_tensor_model_device, + dtype=torch.float32, + requires_grad=False, + ) + _reduce_scatter_base( + output_for_base, + input, + group=parallel_state.get_tensor_model_parallel_group(), + ) + + msg = f"tensor_model_parallel_world_size: {tensor_model_parallel_world_size}" + self.assertEqual(output, output_for_base, msg=msg) + self.assertEqual(input, torch.cat(input_list), msg=msg) + parallel_state.destroy_model_parallel() def test_parallel_embedding(self) -> None: for tensor_model_parallel_world_size in range(1, self.world_size + 1): if self.world_size % tensor_model_parallel_world_size: continue - with self.subTest( - tensor_model_parallel_world_size=tensor_model_parallel_world_size - ): - parallel_state.initialize_model_parallel( - tensor_model_parallel_size_=tensor_model_parallel_world_size, - ) - set_random_seed(self.SEED + 1) - input_tensor = torch.randint( - 0, - self.VOCAB_SIZE, - ( - self.BATCH_SIZE, - self.SEQUENCE_LENGTH, - ), - device="cuda", - ) - loss_weight = torch.randn( - ( - self.BATCH_SIZE, - self.SEQUENCE_LENGTH, - self.HIDDEN_SIZE, - ), - device="cuda", - ) - - set_random_seed(self.SEED) - embedding_torch = nn.Embedding( - self.VOCAB_SIZE, + parallel_state.initialize_model_parallel( + tensor_model_parallel_size_=tensor_model_parallel_world_size, + ) + set_random_seed(self.SEED + 1) + input_tensor = torch.randint( + 0, + self.VOCAB_SIZE, + ( + self.BATCH_SIZE, + self.SEQUENCE_LENGTH, + ), + device="cuda", + ) + loss_weight = torch.randn( + ( + self.BATCH_SIZE, + self.SEQUENCE_LENGTH, self.HIDDEN_SIZE, - ).cuda() - output_torch = embedding_torch(input_tensor) - loss_torch = torch.mul(output_torch, loss_weight).sum() - loss_torch.backward() - - # N.B.(mkozuki): With affine weight initialization on GPU, - # it's super difficult to keep the consistency with nn.Embedding. - # Thus, turning on `use_cpu_initialization`. - set_random_seed(self.SEED) - embedding_vocab_parallel = layers.VocabParallelEmbedding( - self.VOCAB_SIZE, - self.HIDDEN_SIZE, - init_method=nn.init.normal_, - use_cpu_initialization=True, - ).cuda() - output_vocab_parallel = embedding_vocab_parallel(input_tensor) - loss_vocab_parallel = torch.mul( - output_vocab_parallel, loss_weight - ).sum() - loss_vocab_parallel.backward() - - self.assertEqual(output_torch, output_vocab_parallel) - self.assertEqual(loss_torch, loss_vocab_parallel) - - splitted_weight_torch = torch.split( - embedding_torch.weight.grad, - self.VOCAB_SIZE - // tensor_model_parallel_world_size, - 0, - )[parallel_state.get_tensor_model_parallel_rank()] - self.assertEqual( - splitted_weight_torch, embedding_vocab_parallel.weight.grad - ) - - parallel_state.destroy_model_parallel() + ), + device="cuda", + ) + + set_random_seed(self.SEED) + embedding_torch = nn.Embedding( + self.VOCAB_SIZE, + self.HIDDEN_SIZE, + ).cuda() + output_torch = embedding_torch(input_tensor) + loss_torch = torch.mul(output_torch, loss_weight).sum() + loss_torch.backward() + + # N.B.(mkozuki): With affine weight initialization on GPU, + # it's super difficult to keep the consistency with nn.Embedding. + # Thus, turning on `use_cpu_initialization`. + set_random_seed(self.SEED) + embedding_vocab_parallel = layers.VocabParallelEmbedding( + self.VOCAB_SIZE, + self.HIDDEN_SIZE, + init_method=nn.init.normal_, + use_cpu_initialization=True, + ).cuda() + output_vocab_parallel = embedding_vocab_parallel(input_tensor) + loss_vocab_parallel = torch.mul( + output_vocab_parallel, loss_weight + ).sum() + loss_vocab_parallel.backward() + + msg = f"tensor_model_parallel_world_size: {tensor_model_parallel_world_size}" + self.assertEqual(output_torch, output_vocab_parallel, msg=msg) + self.assertEqual(loss_torch, loss_vocab_parallel, msg=msg) + + splitted_weight_torch = torch.split( + embedding_torch.weight.grad, + self.VOCAB_SIZE + // tensor_model_parallel_world_size, + 0, + )[parallel_state.get_tensor_model_parallel_rank()] + self.assertEqual( + splitted_weight_torch, embedding_vocab_parallel.weight.grad, msg=msg, + ) + + parallel_state.destroy_model_parallel() def _affine_weight_init_test_impl( self, init_device: str, is_column_parallel: bool @@ -210,56 +208,55 @@ def _affine_weight_init_test_impl( for tensor_model_parallel_world_size in range(1, self.world_size + 1): if self.world_size % tensor_model_parallel_world_size: continue - with self.subTest( - tensor_model_parallel_world_size=tensor_model_parallel_world_size - ): - parallel_state.initialize_model_parallel( - tensor_model_parallel_size_=tensor_model_parallel_world_size + parallel_state.initialize_model_parallel( + tensor_model_parallel_size_=tensor_model_parallel_world_size + ) + input_size: int = self.INPUT_SIZE_COEFF * tensor_model_parallel_world_size + output_size: int = self.OUTPUT_SIZE_COEFF * tensor_model_parallel_world_size + + weight_shape = ( + (self.OUTPUT_SIZE_COEFF, input_size) + if is_column_parallel + else (output_size, self.INPUT_SIZE_COEFF) + ) + weight = torch.empty(weight_shape) + set_random_seed(self.SEED) + + sharding_dim_size = ( + self.OUTPUT_SIZE_COEFF + if is_column_parallel + else self.INPUT_SIZE_COEFF + ) + + if init_device == "cpu": + layers._initialize_affine_weight_cpu( + weight, + output_size, + input_size, + sharding_dim_size, + dim, + nn.init.normal_, + params_dtype=torch.float32, ) - input_size: int = self.INPUT_SIZE_COEFF * tensor_model_parallel_world_size - output_size: int = self.OUTPUT_SIZE_COEFF * tensor_model_parallel_world_size - - weight_shape = ( - (self.OUTPUT_SIZE_COEFF, input_size) - if is_column_parallel - else (output_size, self.INPUT_SIZE_COEFF) + else: + layers._initialize_affine_weight_gpu( + weight, torch.nn.init.normal_, dim ) - weight = torch.empty(weight_shape) - set_random_seed(self.SEED) + # Target + set_random_seed(self.SEED) + if init_device == "cpu": + main_weight = torch.empty(output_size, input_size) + nn.init.normal_(main_weight) + curr_weight = torch.split(main_weight, sharding_dim_size, dim=dim)[ + parallel_state.get_tensor_model_parallel_rank() + ] + else: + curr_weight = torch.empty(*weight_shape) + nn.init.normal_(curr_weight) - sharding_dim_size = ( - self.OUTPUT_SIZE_COEFF - if is_column_parallel - else self.INPUT_SIZE_COEFF - ) - - if init_device == "cpu": - layers._initialize_affine_weight_cpu( - weight, - output_size, - input_size, - sharding_dim_size, - dim, - nn.init.normal_, - params_dtype=torch.float32, - ) - else: - layers._initialize_affine_weight_gpu( - weight, torch.nn.init.normal_, dim - ) - # Target - set_random_seed(self.SEED) - if init_device == "cpu": - main_weight = torch.empty(output_size, input_size) - nn.init.normal_(main_weight) - curr_weight = torch.split(main_weight, sharding_dim_size, dim=dim)[ - parallel_state.get_tensor_model_parallel_rank() - ] - else: - curr_weight = torch.empty(*weight_shape) - nn.init.normal_(curr_weight) - self.assertEqual(curr_weight, weight) - parallel_state.destroy_model_parallel() + self.assertEqual( + curr_weight, weight, msg=f"tensor_model_parallel_world_size: {tensor_model_parallel_world_size}") + parallel_state.destroy_model_parallel() def test_affine_weight_init_column_parallel_cpu(self) -> None: self._affine_weight_init_test_impl(init_device="cpu", is_column_parallel=True) @@ -304,103 +301,105 @@ def _row_parallel_linear_test_impl( ): if self.world_size % tensor_model_parallel_world_size: continue - with self.subTest( - tensor_model_parallel_world_size=tensor_model_parallel_world_size, - ): - parallel_state.initialize_model_parallel( - tensor_model_parallel_size_=tensor_model_parallel_world_size, - ) - set_random_seed(self.SEED) + parallel_state.initialize_model_parallel( + tensor_model_parallel_size_=tensor_model_parallel_world_size, + ) + set_random_seed(self.SEED) + + linear = layers.RowParallelLinear( + self.HIDDEN_SIZE, + self.HIDDEN_SIZE, + keep_master_weight_for_test=True, + params_dtype=torch.float32, + use_cpu_initialization=True, + gradient_accumulation_fusion=gradient_accumulation_fusion, + accumulation_in_fp16=accumulation_in_fp16, + sequence_parallel_enabled=sequence_parallel_enabled, + # n.b.(mkozuki): RowParallelLinear is constructed with `input_is_parallel=True` + # by default, e.g. https://github.com/NVIDIA/NeMo/blob/782b4e1652aaa43c8be390d9\ + # db0dc89544afa080/nemo/collections/nlp/modules/common/megatron/transformer.py#L204 + input_is_parallel=True, + ).cuda() + if accumulation_in_fp16: + linear = linear.half() + # Simulate the situation where fusion of weight grad calculation and gradient accumulation is enabled. + if gradient_accumulation_fusion: + with torch.no_grad(): + linear.weight.main_grad = torch.zeros_like(linear.weight) - linear = layers.RowParallelLinear( - self.HIDDEN_SIZE, - self.HIDDEN_SIZE, - keep_master_weight_for_test=True, - params_dtype=torch.float32, - use_cpu_initialization=True, - gradient_accumulation_fusion=gradient_accumulation_fusion, - accumulation_in_fp16=accumulation_in_fp16, - sequence_parallel_enabled=sequence_parallel_enabled, - # n.b.(mkozuki): RowParallelLinear is constructed with `input_is_parallel=True` - # by default, e.g. https://github.com/NVIDIA/NeMo/blob/782b4e1652aaa43c8be390d9\ - # db0dc89544afa080/nemo/collections/nlp/modules/common/megatron/transformer.py#L204 - input_is_parallel=True, - ).cuda() - if accumulation_in_fp16: - linear = linear.half() - # Simulate the situation where fusion of weight grad calculation and gradient accumulation is enabled. - if gradient_accumulation_fusion: - with torch.no_grad(): - linear.weight.main_grad = torch.zeros_like(linear.weight) + msg = f"tensor_model_parallel_world_size: {tensor_model_parallel_world_size}" - with torch.no_grad(): - orig_input_tensor = torch.randn(tensor_shape, requires_grad=True, device="cuda") - orig_loss_weight = torch.randn(tensor_shape, device="cuda") - input_tensor = orig_input_tensor.chunk( + with torch.no_grad(): + orig_input_tensor = torch.randn(tensor_shape, requires_grad=True, device="cuda") + orig_loss_weight = torch.randn(tensor_shape, device="cuda") + input_tensor = orig_input_tensor.chunk( + chunks=tensor_model_parallel_world_size, + dim=2, + )[parallel_state.get_tensor_model_parallel_rank()].contiguous() + if sequence_parallel_enabled: + loss_weight = orig_loss_weight.chunk( chunks=tensor_model_parallel_world_size, - dim=2, - )[parallel_state.get_tensor_model_parallel_rank()].contiguous() - if sequence_parallel_enabled: - loss_weight = orig_loss_weight.chunk( - chunks=tensor_model_parallel_world_size, - dim=0, - )[parallel_state.get_tensor_model_parallel_rank()] - else: - loss_weight = orig_loss_weight - if accumulation_in_fp16: - orig_input_tensor = orig_input_tensor.half() - input_tensor = input_tensor.half() - loss_weight = loss_weight.half() - input_tensor.requires_grad_() - output, _ = linear(input_tensor) - loss = torch.mul(output, loss_weight).sum() - loss.backward() - self.assertIsNotNone(input_tensor.grad) + dim=0, + )[parallel_state.get_tensor_model_parallel_rank()] + else: + loss_weight = orig_loss_weight + if accumulation_in_fp16: + orig_input_tensor = orig_input_tensor.half() + input_tensor = input_tensor.half() + loss_weight = loss_weight.half() + input_tensor.requires_grad_() + output, _ = linear(input_tensor) + loss = torch.mul(output, loss_weight).sum() + loss.backward() + self.assertIsNotNone(input_tensor.grad, msg=msg) + + ref_linear = nn.Linear( + in_features=self.HIDDEN_SIZE, + out_features=self.HIDDEN_SIZE, + bias=False, + device="cuda", + ) + with torch.no_grad(): + dldy = orig_loss_weight.clone() + x = orig_input_tensor.clone() + ref_linear.weight.copy_(linear.master_weight) + if accumulation_in_fp16: + ref_linear = ref_linear.half() + x.requires_grad_() + expected_output = ref_linear(x) + expected_loss = torch.mul(expected_output, dldy).sum() + expected_loss.backward() - ref_linear = nn.Linear( - in_features=self.HIDDEN_SIZE, - out_features=self.HIDDEN_SIZE, - bias=False, - device="cuda", - ) - with torch.no_grad(): - dldy = orig_loss_weight.clone() - x = orig_input_tensor.clone() - ref_linear.weight.copy_(linear.master_weight) - if accumulation_in_fp16: - ref_linear = ref_linear.half() - x.requires_grad_() - expected_output = ref_linear(x) - expected_loss = torch.mul(expected_output, dldy).sum() - expected_loss.backward() - - if not accumulation_in_fp16: - if sequence_parallel_enabled: - self.assertEqual( - x=output, - y=expected_output.chunk( - chunks=tensor_model_parallel_world_size, - dim=0, - )[parallel_state.get_tensor_model_parallel_rank()], - ) - else: - self.assertEqual( - x=output, - y=expected_output, - ) - - grad_attr_name = "main_grad" if gradient_accumulation_fusion else "grad" - # NOTE(mkozuki): Numerical errors seems to be enlarged by tensor model parallel. - if tensor_model_parallel_world_size == 1: + if not accumulation_in_fp16: + if sequence_parallel_enabled: self.assertEqual( - x=getattr(linear.weight, grad_attr_name), - y=ref_linear.weight.grad.chunk( + x=output, + y=expected_output.chunk( chunks=tensor_model_parallel_world_size, dim=0, )[parallel_state.get_tensor_model_parallel_rank()], + msg=msg, + ) + else: + self.assertEqual( + x=output, + y=expected_output, + msg=msg, ) - parallel_state.destroy_model_parallel() + grad_attr_name = "main_grad" if gradient_accumulation_fusion else "grad" + # NOTE(mkozuki): Numerical errors seems to be enlarged by tensor model parallel. + if tensor_model_parallel_world_size == 1: + self.assertEqual( + x=getattr(linear.weight, grad_attr_name), + y=ref_linear.weight.grad.chunk( + chunks=tensor_model_parallel_world_size, + dim=0, + )[parallel_state.get_tensor_model_parallel_rank()], + msg=msg, + ) + + parallel_state.destroy_model_parallel() def test_column_parallel_linear(self): self._column_parallel_linear_test_impl(False, False, False, False) @@ -438,112 +437,115 @@ def _column_parallel_linear_test_impl( if async_tensor_model_parallel_allreduce and sequence_parallel_enabled: if tensor_model_parallel_world_size == 1: continue - with self.subTest(tensor_model_parallel_world_size=tensor_model_parallel_world_size): - if self.world_size % tensor_model_parallel_world_size: - continue - parallel_state.initialize_model_parallel( - tensor_model_parallel_size_=tensor_model_parallel_world_size, - ) - - input_tensor_shape = self.tensor_shape - expected_output_shape = self.tensor_shape - # When sequence parallel, `gather_output` is disabled, i.e., - # output of matmul isn't gathered in dimension of feature/hidden (last dim). - if sequence_parallel_enabled: - expected_output_shape[-1] //= tensor_model_parallel_world_size - - # tensor's shape is [sequence length, batch size, hidden size] - set_random_seed(self.SEED) - linear = layers.ColumnParallelLinear( - self.HIDDEN_SIZE, - self.HIDDEN_SIZE, + if self.world_size % tensor_model_parallel_world_size: + continue + msg = f"tensor_model_parallel_world_size: {tensor_model_parallel_world_size}" + parallel_state.initialize_model_parallel( + tensor_model_parallel_size_=tensor_model_parallel_world_size, + ) + + input_tensor_shape = self.tensor_shape + expected_output_shape = self.tensor_shape + # When sequence parallel, `gather_output` is disabled, i.e., + # output of matmul isn't gathered in dimension of feature/hidden (last dim). + if sequence_parallel_enabled: + expected_output_shape[-1] //= tensor_model_parallel_world_size + + # tensor's shape is [sequence length, batch size, hidden size] + set_random_seed(self.SEED) + linear = layers.ColumnParallelLinear( + self.HIDDEN_SIZE, + self.HIDDEN_SIZE, + bias=False, + keep_master_weight_for_test=True, + params_dtype=torch.float32, + use_cpu_initialization=True, + gather_output=not sequence_parallel_enabled, + no_async_tensor_model_parallel_allreduce=not async_tensor_model_parallel_allreduce, + gradient_accumulation_fusion=gradient_accumulation_fusion, + accumulation_in_fp16=accumulation_in_fp16, + sequence_parallel_enabled=sequence_parallel_enabled, + ).cuda() + if accumulation_in_fp16: + linear = linear.half() + + # Simulate the situation where fusion of weight grad calculation and gradient accumulation happens. + if gradient_accumulation_fusion: + with torch.no_grad(): + linear.weight.main_grad = torch.zeros_like(linear.weight) + + orig_input_tensor = torch.randn(input_tensor_shape, device="cuda", requires_grad=True) + if accumulation_in_fp16: + orig_input_tensor = orig_input_tensor.half() + if sequence_parallel_enabled: + input_tensor = list( + orig_input_tensor.chunk(tensor_model_parallel_world_size, dim=0) + )[parallel_state.get_tensor_model_parallel_rank()] + else: + input_tensor = orig_input_tensor + output, _ = linear(input_tensor) + # The order of dimension is expected to be (sequence, batch, hidden) + self.assertEqual(output.shape, expected_output_shape, msg=msg) + + orig_loss_weight = torch.randn(input_tensor_shape, device="cuda") + if accumulation_in_fp16: + orig_loss_weight = orig_loss_weight.half() + if sequence_parallel_enabled: + loss_weight = orig_loss_weight.chunk( + tensor_model_parallel_world_size, dim=2, + )[parallel_state.get_tensor_model_parallel_rank()] + else: + loss_weight = orig_loss_weight + loss = torch.mul(output, loss_weight).sum() + loss.backward() + + with torch.no_grad(): + dldy = orig_loss_weight.clone() + x = orig_input_tensor.clone() + ref_linear = nn.Linear( + in_features=self.HIDDEN_SIZE, + out_features=self.HIDDEN_SIZE, bias=False, - keep_master_weight_for_test=True, - params_dtype=torch.float32, - use_cpu_initialization=True, - gather_output=not sequence_parallel_enabled, - no_async_tensor_model_parallel_allreduce=not async_tensor_model_parallel_allreduce, - gradient_accumulation_fusion=gradient_accumulation_fusion, - accumulation_in_fp16=accumulation_in_fp16, - sequence_parallel_enabled=sequence_parallel_enabled, - ).cuda() - if accumulation_in_fp16: - linear = linear.half() - - # Simulate the situation where fusion of weight grad calculation and gradient accumulation happens. - if gradient_accumulation_fusion: - with torch.no_grad(): - linear.weight.main_grad = torch.zeros_like(linear.weight) - - orig_input_tensor = torch.randn(input_tensor_shape, device="cuda", requires_grad=True) - if accumulation_in_fp16: - orig_input_tensor = orig_input_tensor.half() - if sequence_parallel_enabled: - input_tensor = list( - orig_input_tensor.chunk(tensor_model_parallel_world_size, dim=0) - )[parallel_state.get_tensor_model_parallel_rank()] - else: - input_tensor = orig_input_tensor - output, _ = linear(input_tensor) - # The order of dimension is expected to be (sequence, batch, hidden) - self.assertEqual(output.shape, expected_output_shape) - - orig_loss_weight = torch.randn(input_tensor_shape, device="cuda") + device="cuda", + ) if accumulation_in_fp16: - orig_loss_weight = orig_loss_weight.half() - if sequence_parallel_enabled: - loss_weight = orig_loss_weight.chunk( - tensor_model_parallel_world_size, dim=2, - )[parallel_state.get_tensor_model_parallel_rank()] - else: - loss_weight = orig_loss_weight - loss = torch.mul(output, loss_weight).sum() - loss.backward() - - with torch.no_grad(): - dldy = orig_loss_weight.clone() - x = orig_input_tensor.clone() - ref_linear = nn.Linear( - in_features=self.HIDDEN_SIZE, - out_features=self.HIDDEN_SIZE, - bias=False, - device="cuda", - ) - if accumulation_in_fp16: - ref_linear = ref_linear.half() - # NOTE(mkozuki): `master_weight` is available because `keep_master_weight_for_test` is set. - ref_linear.weight.copy_(linear.master_weight) - x.requires_grad_() - expected_output = ref_linear(x) - if sequence_parallel_enabled: - chunk = expected_output.chunk( - tensor_model_parallel_world_size, - dim=2, - )[parallel_state.get_tensor_model_parallel_rank()] - self.assertEqual( - x=output, - y=chunk, - ) - else: - self.assertEqual( - x=output, - y=expected_output, - ) + ref_linear = ref_linear.half() + # NOTE(mkozuki): `master_weight` is available because `keep_master_weight_for_test` is set. + ref_linear.weight.copy_(linear.master_weight) + x.requires_grad_() + expected_output = ref_linear(x) + if sequence_parallel_enabled: + chunk = expected_output.chunk( + tensor_model_parallel_world_size, + dim=2, + )[parallel_state.get_tensor_model_parallel_rank()] + self.assertEqual( + x=output, + y=chunk, + msg=msg, + ) + else: + self.assertEqual( + x=output, + y=expected_output, + msg=msg, + ) - expected_loss = torch.mul(expected_output, dldy).sum() - expected_loss.backward() - grad_attr_name = "main_grad" if gradient_accumulation_fusion else "grad" - # NOTE(mkozuki): Numerical errors seems to be enlarged by tensor model parallel. - if tensor_model_parallel_world_size == 1: - self.assertEqual( - x=getattr(linear.weight, grad_attr_name), - y=ref_linear.weight.grad.chunk( - chunks=tensor_model_parallel_world_size, - dim=0, - )[parallel_state.get_tensor_model_parallel_rank()], - ) + expected_loss = torch.mul(expected_output, dldy).sum() + expected_loss.backward() + grad_attr_name = "main_grad" if gradient_accumulation_fusion else "grad" + # NOTE(mkozuki): Numerical errors seems to be enlarged by tensor model parallel. + if tensor_model_parallel_world_size == 1: + self.assertEqual( + x=getattr(linear.weight, grad_attr_name), + y=ref_linear.weight.grad.chunk( + chunks=tensor_model_parallel_world_size, + dim=0, + )[parallel_state.get_tensor_model_parallel_rank()], + msg=msg, + ) - parallel_state.destroy_model_parallel() + parallel_state.destroy_model_parallel() class NcclTensorParallelLayerTest(TensorParallelLayerTestBase, NcclDistributedTestBase): diff --git a/tests/L0/run_transformer/test_mapping.py b/tests/L0/run_transformer/test_mapping.py index 9ebda6dc6..18fa07636 100644 --- a/tests/L0/run_transformer/test_mapping.py +++ b/tests/L0/run_transformer/test_mapping.py @@ -18,67 +18,65 @@ def test_reduce(self): for tensor_model_paralell_world_size in range(1, self.world_size + 1): if self.world_size % tensor_model_paralell_world_size > 0: continue - with self.subTest( - tensor_model_paralell_world_size=tensor_model_paralell_world_size - ): - parallel_state.initialize_model_parallel( - tensor_model_parallel_size_=tensor_model_paralell_world_size - ) - t = torch.full((10, 10, 10, 10), 50, device=f"cuda:{self.rank}") - expected = torch.full( - (10, 10, 10, 10), - 50 * tensor_model_paralell_world_size, - device=f"cuda:{self.rank}", - ) - self.assertTrue(torch.equal(mappings._reduce(t), expected)) - parallel_state.destroy_model_parallel() + parallel_state.initialize_model_parallel( + tensor_model_parallel_size_=tensor_model_paralell_world_size + ) + t = torch.full((10, 10, 10, 10), 50, device=f"cuda:{self.rank}") + expected = torch.full( + (10, 10, 10, 10), + 50 * tensor_model_paralell_world_size, + device=f"cuda:{self.rank}", + ) + self.assertTrue( + torch.equal(mappings._reduce(t), expected), + msg=f"tensor_model_paralell_world_size: {tensor_model_paralell_world_size}", + ) + parallel_state.destroy_model_parallel() def test_split(self): for tensor_model_paralell_world_size in range(1, self.world_size + 1): if self.world_size % tensor_model_paralell_world_size > 0: continue - with self.subTest( - tensor_model_paralell_world_size=tensor_model_paralell_world_size - ): - parallel_state.initialize_model_parallel( - tensor_model_parallel_size_=tensor_model_paralell_world_size - ) + parallel_state.initialize_model_parallel( + tensor_model_parallel_size_=tensor_model_paralell_world_size + ) - tensors = [ - torch.randn(10, 1) - for rank in range(tensor_model_paralell_world_size) - ] - x = torch.cat(tensors, 1) - out = mappings._split_along_last_dim(x) - self.assertTrue( - torch.equal( - out, tensors[parallel_state.get_tensor_model_parallel_rank()] - ) - ) - parallel_state.destroy_model_parallel() + tensors = [ + torch.randn(10, 1) + for _ in range(tensor_model_paralell_world_size) + ] + x = torch.cat(tensors, 1) + out = mappings._split_along_last_dim(x) + self.assertTrue( + torch.equal( + out, tensors[parallel_state.get_tensor_model_parallel_rank()] + ), + msg=f"tensor_model_paralell_world_size: {tensor_model_paralell_world_size}" + ) + parallel_state.destroy_model_parallel() def test_gather(self): for tensor_model_paralell_world_size in range(1, self.world_size + 1): if self.world_size % tensor_model_paralell_world_size > 0: continue - with self.subTest( - tensor_model_paralell_world_size=tensor_model_paralell_world_size - ): - parallel_state.initialize_model_parallel( - tensor_model_parallel_size_=tensor_model_paralell_world_size - ) - device = f"cuda:{self.rank}" - gathered = mappings._gather_along_last_dim( - torch.tensor( - [parallel_state.get_tensor_model_parallel_rank()], device=device - ) - ) - expected = torch.tensor( - [rank for rank in range(tensor_model_paralell_world_size)], - device=device, + parallel_state.initialize_model_parallel( + tensor_model_parallel_size_=tensor_model_paralell_world_size + ) + device = f"cuda:{self.rank}" + gathered = mappings._gather_along_last_dim( + torch.tensor( + [parallel_state.get_tensor_model_parallel_rank()], device=device ) - self.assertTrue(torch.equal(gathered, expected)) - parallel_state.destroy_model_parallel() + ) + expected = torch.tensor( + [rank for rank in range(tensor_model_paralell_world_size)], + device=device, + ) + self.assertTrue( + torch.equal(gathered, expected), + msg=f"tensor_model_paralell_world_size: {tensor_model_paralell_world_size}", + ) + parallel_state.destroy_model_parallel() class NcclMappingTest(MappingTestBase, NcclDistributedTestBase): pass diff --git a/tests/L0/run_transformer/test_microbatches.py b/tests/L0/run_transformer/test_microbatches.py index 0d4b50972..13e64881a 100644 --- a/tests/L0/run_transformer/test_microbatches.py +++ b/tests/L0/run_transformer/test_microbatches.py @@ -39,36 +39,36 @@ def _test(self, rampup_batch_size: Optional[List[int]]) -> None: continue if self.world_size % data_parallel_size != 0: continue - with self.subTest(data_parallel_size=data_parallel_size): - parallel_state.initialize_model_parallel( - tensor_model_parallel_size_=self.world_size // data_parallel_size, - pipeline_model_parallel_size_=1, - ) - self.assertEqual(data_parallel_size, parallel_state.get_data_parallel_world_size()) - - _reconfigure_microbatch_calculator( - self.rank, - rampup_batch_size, - self.GLOBAL_BATCH_SIZE, - self.MICRO_BATCH_SIZE, - data_parallel_size, - ) - - self.assertEqual(get_micro_batch_size(), expected_micro_batch_size) - self.assertEqual(get_num_microbatches(), expected_global_batch_size / expected_micro_batch_size / data_parallel_size) - current_global_batch_size = get_current_global_batch_size() - self.assertEqual(current_global_batch_size, expected_global_batch_size) - - # Make sure `global_batch_size` equals to the final global batch size after - # certain number of updates. - if rampup_batch_size: - update_num_microbatches(current_global_batch_size) - for i in range(100): - current_global_batch_size = get_current_global_batch_size() - update_num_microbatches(current_global_batch_size) + msg = f"data_parallel_size: {data_parallel_size}" + parallel_state.initialize_model_parallel( + tensor_model_parallel_size_=self.world_size // data_parallel_size, + pipeline_model_parallel_size_=1, + ) + self.assertEqual(data_parallel_size, parallel_state.get_data_parallel_world_size(), msg=msg) + + _reconfigure_microbatch_calculator( + self.rank, + rampup_batch_size, + self.GLOBAL_BATCH_SIZE, + self.MICRO_BATCH_SIZE, + data_parallel_size, + ) + + self.assertEqual(get_micro_batch_size(), expected_micro_batch_size, msg=msg) + self.assertEqual(get_num_microbatches(), expected_global_batch_size / expected_micro_batch_size / data_parallel_size, msg=msg) + current_global_batch_size = get_current_global_batch_size() + self.assertEqual(current_global_batch_size, expected_global_batch_size, msg=msg) + + # Make sure `global_batch_size` equals to the final global batch size after + # certain number of updates. + if rampup_batch_size: + update_num_microbatches(current_global_batch_size) + for i in range(100): current_global_batch_size = get_current_global_batch_size() - self.assertEqual(get_current_global_batch_size(), self.GLOBAL_BATCH_SIZE) - parallel_state.destroy_model_parallel() + update_num_microbatches(current_global_batch_size) + current_global_batch_size = get_current_global_batch_size() + self.assertEqual(get_current_global_batch_size(), self.GLOBAL_BATCH_SIZE, msg=msg) + parallel_state.destroy_model_parallel() def test_constant_microbatch_calculator(self): self._test(rampup_batch_size=None) diff --git a/tests/L0/run_transformer/test_parallel_state.py b/tests/L0/run_transformer/test_parallel_state.py index 0314895ea..a86c2fb36 100644 --- a/tests/L0/run_transformer/test_parallel_state.py +++ b/tests/L0/run_transformer/test_parallel_state.py @@ -28,42 +28,43 @@ def test_initialize_model_parallel(self) -> None: self.assertFalse(parallel_state.model_parallel_is_initialized()) for tensor_model_parallel_world_size in range(1, self.world_size + 1): - with self.subTest( - tensor_model_parallel_world_size=tensor_model_parallel_world_size - ): - if self.world_size % tensor_model_parallel_world_size: - continue - - pipeline_model_parallel_world_size = ( - self.world_size // tensor_model_parallel_world_size - ) - - parallel_state.initialize_model_parallel( - tensor_model_parallel_size_=tensor_model_parallel_world_size, - pipeline_model_parallel_size_=pipeline_model_parallel_world_size, - ) - self.assertEqual( - tensor_model_parallel_world_size, - parallel_state.get_tensor_model_parallel_world_size(), - ) - expected_tensor_model_parallel_rank = calc_expected_tensor_model_paralell_rank( - self.rank, tensor_model_parallel_world_size - ) - self.assertEqual( - expected_tensor_model_parallel_rank, - parallel_state.get_tensor_model_parallel_rank(), - ) - - expected_tensor_model_parallel_src_rank = ( - self.rank // tensor_model_parallel_world_size - ) * tensor_model_parallel_world_size - self.assertEqual( - expected_tensor_model_parallel_src_rank, - parallel_state.get_tensor_model_parallel_src_rank(), - ) - - parallel_state.destroy_model_parallel() - self.assertFalse(parallel_state.model_parallel_is_initialized()) + msg = f"tensor_model_parallel_world_siz: {tensor_model_parallel_world_size}" + if self.world_size % tensor_model_parallel_world_size: + continue + + pipeline_model_parallel_world_size = ( + self.world_size // tensor_model_parallel_world_size + ) + + parallel_state.initialize_model_parallel( + tensor_model_parallel_size_=tensor_model_parallel_world_size, + pipeline_model_parallel_size_=pipeline_model_parallel_world_size, + ) + self.assertEqual( + tensor_model_parallel_world_size, + parallel_state.get_tensor_model_parallel_world_size(), + msg=msg, + ) + expected_tensor_model_parallel_rank = calc_expected_tensor_model_paralell_rank( + self.rank, tensor_model_parallel_world_size + ) + self.assertEqual( + expected_tensor_model_parallel_rank, + parallel_state.get_tensor_model_parallel_rank(), + msg=msg, + ) + + expected_tensor_model_parallel_src_rank = ( + self.rank // tensor_model_parallel_world_size + ) * tensor_model_parallel_world_size + self.assertEqual( + expected_tensor_model_parallel_src_rank, + parallel_state.get_tensor_model_parallel_src_rank(), + msg=msg, + ) + + parallel_state.destroy_model_parallel() + self.assertFalse(parallel_state.model_parallel_is_initialized(), msg=msg) def test_initialize_model_parallel_with_virtual_and_split(self) -> None: if self.world_size < 4: @@ -138,43 +139,44 @@ def test_initialize_model_parallel_decoder_only(self) -> None: self.assertFalse(parallel_state.model_parallel_is_initialized()) for tensor_model_parallel_world_size in range(1, self.world_size + 1): - with self.subTest( - tensor_model_parallel_world_size=tensor_model_parallel_world_size - ): - if self.world_size % tensor_model_parallel_world_size: - continue - - pipeline_model_parallel_world_size = ( - self.world_size // tensor_model_parallel_world_size - ) - - parallel_state.initialize_model_parallel( - tensor_model_parallel_size_=tensor_model_parallel_world_size, - pipeline_model_parallel_size_=pipeline_model_parallel_world_size, - pipeline_model_parallel_split_rank_=0, - ) - self.assertEqual( - tensor_model_parallel_world_size, - parallel_state.get_tensor_model_parallel_world_size(), - ) - expected_tensor_model_parallel_rank = calc_expected_tensor_model_paralell_rank( - self.rank, tensor_model_parallel_world_size - ) - self.assertEqual( - expected_tensor_model_parallel_rank, - parallel_state.get_tensor_model_parallel_rank(), - ) - - expected_tensor_model_parallel_src_rank = ( - self.rank // tensor_model_parallel_world_size - ) * tensor_model_parallel_world_size - self.assertEqual( - expected_tensor_model_parallel_src_rank, - parallel_state.get_tensor_model_parallel_src_rank(), - ) - - parallel_state.destroy_model_parallel() - self.assertFalse(parallel_state.model_parallel_is_initialized()) + msg = f"tensor_model_parallel_world_size: {tensor_model_parallel_world_size}" + if self.world_size % tensor_model_parallel_world_size: + continue + + pipeline_model_parallel_world_size = ( + self.world_size // tensor_model_parallel_world_size + ) + + parallel_state.initialize_model_parallel( + tensor_model_parallel_size_=tensor_model_parallel_world_size, + pipeline_model_parallel_size_=pipeline_model_parallel_world_size, + pipeline_model_parallel_split_rank_=0, + ) + self.assertEqual( + tensor_model_parallel_world_size, + parallel_state.get_tensor_model_parallel_world_size(), + msg=msg, + ) + expected_tensor_model_parallel_rank = calc_expected_tensor_model_paralell_rank( + self.rank, tensor_model_parallel_world_size + ) + self.assertEqual( + expected_tensor_model_parallel_rank, + parallel_state.get_tensor_model_parallel_rank(), + msg=msg, + ) + + expected_tensor_model_parallel_src_rank = ( + self.rank // tensor_model_parallel_world_size + ) * tensor_model_parallel_world_size + self.assertEqual( + expected_tensor_model_parallel_src_rank, + parallel_state.get_tensor_model_parallel_src_rank(), + msg=msg, + ) + + parallel_state.destroy_model_parallel() + self.assertFalse(parallel_state.model_parallel_is_initialized(), msg=msg) class NcclParallelStateTest(ParallelStateTestBase, NcclDistributedTestBase): pass diff --git a/tests/L0/run_transformer/test_random.py b/tests/L0/run_transformer/test_random.py index 6060f9ed9..bd737ee5e 100644 --- a/tests/L0/run_transformer/test_random.py +++ b/tests/L0/run_transformer/test_random.py @@ -18,98 +18,95 @@ def test_set_cuda_rng_state(self): for tensor_model_parallel_world_size in range(1, self.world_size + 1): if self.world_size % tensor_model_parallel_world_size: continue - with self.subTest( - tensor_model_parallel_world_size=tensor_model_parallel_world_size - ): - parallel_state.initialize_model_parallel( - tensor_model_parallel_size_=tensor_model_parallel_world_size - ) + msg = f"tensor_model_parallel_world_size: {tensor_model_parallel_world_size}" + parallel_state.initialize_model_parallel( + tensor_model_parallel_size_=tensor_model_parallel_world_size + ) - size, seed = 123, 1234 - torch.cuda.manual_seed(seed) - tensor = torch.cuda.FloatTensor(size) + size, seed = 123, 1234 + torch.cuda.manual_seed(seed) + tensor = torch.cuda.FloatTensor(size) - rng_state = torch.cuda.get_rng_state() - rng_state_clone = rng_state.clone() + rng_state = torch.cuda.get_rng_state() + rng_state_clone = rng_state.clone() - for _ in range(5): - torch.randn(size, out=tensor) - result_1 = tensor.clone() + for _ in range(5): + torch.randn(size, out=tensor) + result_1 = tensor.clone() - self.assertEqual(rng_state.sub(rng_state_clone).max(), 0) - self.assertGreater( - torch.cuda.get_rng_state().sub(rng_state_clone).max(), 0 - ) + self.assertEqual(rng_state.sub(rng_state_clone).max(), 0, msg=msg) + self.assertGreater( + torch.cuda.get_rng_state().sub(rng_state_clone).max(), 0, + msg=msg, + ) - new_rng_state = torch.cuda.get_rng_state() - self.assertGreater(new_rng_state.sub(rng_state).max(), 0) + new_rng_state = torch.cuda.get_rng_state() + self.assertGreater(new_rng_state.sub(rng_state).max(), 0, msg=msg) - tensor_parallel.random._set_cuda_rng_state(rng_state) - for _ in range(5): - torch.randn(size, out=tensor) - tensor_parallel.random._set_cuda_rng_state(rng_state) - for _ in range(5): - torch.randn(size, out=tensor) - result_2 = tensor.clone() + tensor_parallel.random._set_cuda_rng_state(rng_state) + for _ in range(5): + torch.randn(size, out=tensor) + tensor_parallel.random._set_cuda_rng_state(rng_state) + for _ in range(5): + torch.randn(size, out=tensor) + result_2 = tensor.clone() - self.assertEqual(result_2, result_1) + self.assertEqual(result_2, result_1, msg=msg) - self.assertEqual(rng_state.sub(rng_state_clone).max(), 0) + self.assertEqual(rng_state.sub(rng_state_clone).max(), 0, msg=msg) - parallel_state.destroy_model_parallel() + parallel_state.destroy_model_parallel() def test_cuda_rng_tracker(self): for tensor_model_parallel_world_size in range(1, self.world_size + 1): if self.world_size % tensor_model_parallel_world_size: continue - with self.subTest( - tensor_model_parallel_world_size=tensor_model_parallel_world_size - ): - parallel_state.initialize_model_parallel( - tensor_model_parallel_size_=tensor_model_parallel_world_size - ) + msg = f"tensor_model_parallel_world_size: {tensor_model_parallel_world_size}" + parallel_state.initialize_model_parallel( + tensor_model_parallel_size_=tensor_model_parallel_world_size + ) - seed_1, seed_2, size = 1234, 4321, [12, 21] - tensor = torch.cuda.FloatTensor(size) + seed_1, seed_2, size = 1234, 4321, [12, 21] + tensor = torch.cuda.FloatTensor(size) - torch.cuda.manual_seed(seed_1) - torch.randn(size, out=tensor) - target_11 = tensor.clone() - torch.randn(size, out=tensor) - target_12 = tensor.clone() + torch.cuda.manual_seed(seed_1) + torch.randn(size, out=tensor) + target_11 = tensor.clone() + torch.randn(size, out=tensor) + target_12 = tensor.clone() - torch.cuda.manual_seed(seed_2) - torch.randn(size, out=tensor) - targt_21 = tensor.clone() - torch.randn(size, out=tensor) - target_22 = tensor.clone() + torch.cuda.manual_seed(seed_2) + torch.randn(size, out=tensor) + targt_21 = tensor.clone() + torch.randn(size, out=tensor) + target_22 = tensor.clone() - torch.cuda.manual_seed(seed_1) - tensor_parallel.random.get_cuda_rng_tracker().add("test", seed_2) + torch.cuda.manual_seed(seed_1) + tensor_parallel.random.get_cuda_rng_tracker().add("test", seed_2) + torch.randn(size, out=tensor) + result_11 = tensor.clone() + + with tensor_parallel.random.get_cuda_rng_tracker().fork("test"): torch.randn(size, out=tensor) - result_11 = tensor.clone() + result_21 = tensor.clone() - with tensor_parallel.random.get_cuda_rng_tracker().fork("test"): - torch.randn(size, out=tensor) - result_21 = tensor.clone() + torch.randn(size, out=tensor) + result_12 = tensor.clone() + with tensor_parallel.random.get_cuda_rng_tracker().fork("test"): torch.randn(size, out=tensor) - result_12 = tensor.clone() - - with tensor_parallel.random.get_cuda_rng_tracker().fork("test"): - torch.randn(size, out=tensor) - result_22 = tensor.clone() + result_22 = tensor.clone() - self.assertEqual(target_11, result_11) - self.assertEqual(target_12, result_12) - self.assertEqual(targt_21, result_21) - self.assertEqual(target_22, result_22) - self.assertNotEqual(result_11, result_21) - self.assertNotEqual(result_21, result_22) + self.assertEqual(target_11, result_11, msg=msg) + self.assertEqual(target_12, result_12, msg=msg) + self.assertEqual(targt_21, result_21, msg=msg) + self.assertEqual(target_22, result_22, msg=msg) + self.assertNotEqual(result_11, result_21, msg=msg) + self.assertNotEqual(result_21, result_22, msg=msg) - tensor_parallel.random.get_cuda_rng_tracker().reset() - parallel_state.destroy_model_parallel() + tensor_parallel.random.get_cuda_rng_tracker().reset() + parallel_state.destroy_model_parallel() class NcclTransformerRandomTest(TransformerRandomTestBase, NcclDistributedTestBase): pass diff --git a/tests/L0/run_transformer/test_transformer_module.py b/tests/L0/run_transformer/test_transformer_module.py index 77ce67edb..7be1e5654 100644 --- a/tests/L0/run_transformer/test_transformer_module.py +++ b/tests/L0/run_transformer/test_transformer_module.py @@ -31,76 +31,65 @@ def get_launch_option(test_filename) -> Tuple[bool, str]: return should_skip, "" -def run_transformer_tests(): +def get_test_command(test_file: str) -> str: + python_executable_path = sys.executable + is_denied = False + should_skip, launch_option = get_launch_option(test_file) + if should_skip: + return "" + test_run_cmd = ( + f"{python_executable_path} {launch_option} {test_file} " + "--micro-batch-size 2 --num-layers 16 --hidden-size 256 --num-attention-heads 8 --max-position-embeddings " + "512 --seq-length 512 --global-batch-size 128" + ) + if "bert" in test_file or "gpt" in test_file: + import torch + + num_devices = torch.cuda.device_count() + if "bert" in test_file: + # "bert" uses the interleaving. + tensor_model_parallel_size = 2 if num_devices % 2 == 0 and num_devices > 4 else 1 + if "gpt" in test_file: + # "gpt" uses the non-interleaving. + tensor_model_parallel_size = 2 if num_devices % 2 == 0 and num_devices >= 4 else 1 + pipeline_model_parallel_size = num_devices // tensor_model_parallel_size + test_run_cmd += f" --pipeline-model-parallel-size {pipeline_model_parallel_size} --tensor-model-parallel-size {tensor_model_parallel_size}" + + if "bert" in test_file: + test_run_cmd += f" --bert-no-binary-head" + else: + test_run_cmd += f" --use-cpu-initialization" + return test_run_cmd + + +def _get_test_file(key): python_executable_path = sys.executable directory = os.path.dirname(__file__) - files = [ + test_file = [ os.path.join(directory, f) for f in os.listdir(directory) - if f.startswith("run_") and os.path.isfile(os.path.join(directory, f)) - ] - print("#######################################################") - print(f"# Python executable path: {python_executable_path}") - print(f"# {len(files)} tests: {files}") - print("#######################################################") - errors = [] - for i, test_file in enumerate(files, 1): - is_denied = False - should_skip, launch_option = get_launch_option(test_file) - if should_skip: - print( - f"### {i} / {len(files)}: {test_file} skipped. Requires multiple GPUs." - ) - continue - test_run_cmd = ( - f"{python_executable_path} {launch_option} {test_file} " - "--micro-batch-size 2 --num-layers 16 --hidden-size 256 --num-attention-heads 8 --max-position-embeddings " - "512 --seq-length 512 --global-batch-size 128" - ) - if "bert" in test_file or "gpt" in test_file: - import torch - - num_devices = torch.cuda.device_count() - if "bert" in test_file: - # "bert" uses the interleaving. - tensor_model_parallel_size = 2 if num_devices % 2 == 0 and num_devices > 4 else 1 - if "gpt" in test_file: - # "gpt" uses the non-interleaving. - tensor_model_parallel_size = 2 if num_devices % 2 == 0 and num_devices >= 4 else 1 - pipeline_model_parallel_size = num_devices // tensor_model_parallel_size - test_run_cmd += f" --pipeline-model-parallel-size {pipeline_model_parallel_size} --tensor-model-parallel-size {tensor_model_parallel_size}" - - if "bert" in test_file: - test_run_cmd += f" --bert-no-binary-head" - else: - test_run_cmd += f" --use-cpu-initialization" - print(f"### {i} / {len(files)}: cmd: {test_run_cmd}") - try: - output = ( - subprocess.check_output(test_run_cmd, shell=True) - .decode(sys.stdout.encoding) - .strip() - ) - except Exception as e: - errors.append((test_file, str(e))) - else: - if ">> passed the test :-)" not in output: - errors.append((test_file, output)) - else: - if not errors: - print("### PASSED") - else: - print("### FAILED") - short_msg = f"{len(errors)} out of {len(files)} tests failed" - print(short_msg) - for (filename, log) in errors: - print(f"File: {filename}\nLog: {log}") - raise RuntimeError(short_msg) + if f.startswith("run_") and os.path.isfile(os.path.join(directory, f)) and key in f + ][0] + return test_file class TestTransformer(unittest.TestCase): - def test_transformer(self): - run_transformer_tests() + def _test_impl(self, key: str): + test_file = _get_test_file(key) + command = get_test_command(test_file) + if command: + subprocess.run(command, shell=True, check=True) + else: + self.skipTest("Appropriate command is not generated") + + def test_standalone_bert(self): + self._test_impl("bert") + + def test_standalone_gpt(self): + self._test_impl("gpt") + + def test_dynamic_batch_size(self): + self._test_impl("dynamic_batchsize") if __name__ == "__main__": diff --git a/tests/L0/run_transformer/test_transformer_utils.py b/tests/L0/run_transformer/test_transformer_utils.py index d5d16080d..4164eca87 100644 --- a/tests/L0/run_transformer/test_transformer_utils.py +++ b/tests/L0/run_transformer/test_transformer_utils.py @@ -17,23 +17,23 @@ def test_split_tensor_along_last_dim(self): for tensor_model_paralell_world_size in range(1, self.world_size + 1): if self.world_size % tensor_model_paralell_world_size > 0: continue - with self.subTest( - tensor_model_paralell_world_size=tensor_model_paralell_world_size - ): - parallel_state.initialize_model_parallel( - tensor_model_parallel_size_=tensor_model_paralell_world_size - ) - - device = "cpu" - input_tensor = torch.randn((100, 100, 100), device=device) - splits = utils.split_tensor_along_last_dim(input_tensor, 10) - last_dim_shapes = torch.tensor( - [int(split.size()[-1]) for split in splits] - ) - - self.assertTrue(torch.equal(last_dim_shapes, torch.full((10,), 10),)) - - parallel_state.destroy_model_parallel() + parallel_state.initialize_model_parallel( + tensor_model_parallel_size_=tensor_model_paralell_world_size + ) + + device = "cpu" + input_tensor = torch.randn((100, 100, 100), device=device) + splits = utils.split_tensor_along_last_dim(input_tensor, 10) + last_dim_shapes = torch.tensor( + [int(split.size()[-1]) for split in splits] + ) + + self.assertTrue( + torch.equal(last_dim_shapes, torch.full((10,), 10),), + msg=f"tensor_model_paralell_world_size: {tensor_model_paralell_world_size}", + ) + + parallel_state.destroy_model_parallel() if __name__ == "__main__": From 14ce2592c77821f619831d30c00d1ebf3bbc2ef9 Mon Sep 17 00:00:00 2001 From: Masaki Kozuki Date: Tue, 9 Aug 2022 15:48:57 -0700 Subject: [PATCH 04/25] Update mlp_cuda test (#1425) to use `torch.testing.assert_close` instead of `numpy.testing.assert_allclose`. The former uses a bit looser threshold values. Signed-off-by: Masaki Kozuki --- apex/mlp/mlp.py | 1 + tests/L0/run_mlp/test_mlp.py | 214 +++++++++++++++++------------------ 2 files changed, 104 insertions(+), 111 deletions(-) diff --git a/apex/mlp/mlp.py b/apex/mlp/mlp.py index bae38f3f8..be9ce4e5a 100644 --- a/apex/mlp/mlp.py +++ b/apex/mlp/mlp.py @@ -21,6 +21,7 @@ def backward(ctx, grad_o): del ctx.outputs return (None, None, *grads) +# TODO(crcrpar): Should make this compatible with torch.cuda.amp mlp_function = amp.half_function(MlpFunction.apply) class MLP(torch.nn.Module): diff --git a/tests/L0/run_mlp/test_mlp.py b/tests/L0/run_mlp/test_mlp.py index b28b7e791..0b51e58e4 100644 --- a/tests/L0/run_mlp/test_mlp.py +++ b/tests/L0/run_mlp/test_mlp.py @@ -1,7 +1,6 @@ """Tests for c++ MLP""" from itertools import product from time import time -import numpy as np import torch from torch import nn @@ -12,12 +11,14 @@ from apex.mlp import MLP from apex.testing.common_utils import skipFlakyTest + batch_size = 1024 mlp_sizes = [480, 1024, 1024, 512, 256, 1] num_iters = 10 -class TestMLP(unittest.TestCase): +# note(crcrpar): On Ampere, this test should be run without TF32 enabled. +class TestMLP(unittest.TestCase): def test_creation(self): MLP(mlp_sizes) @@ -76,107 +77,91 @@ def test_mlp(self, use_activation: str, bias: bool): ref_mlp = nn.Sequential(*mlp_layers).cuda() - test_input = torch.empty(batch_size, mlp_sizes[0], device="cuda").uniform_(-1., 1.).requires_grad_() + test_input = ( + torch.empty(batch_size, mlp_sizes[0], device="cuda") + .uniform_(-1.0, 1.0) + .requires_grad_() + ) ref_input = test_input.clone().detach().requires_grad_() mlp_out = mlp(test_input) ref_out = ref_mlp(ref_input) - np.testing.assert_allclose( - mlp_out.detach().cpu().numpy(), - ref_out.detach().cpu().numpy(), - atol=1e-7, rtol=1e-5) + torch.testing.assert_close(mlp_out, ref_out) # Use mean value as scalar loss. Multiply 10 to make it big enough not zero out - mlp_out.mean().mul(10.).backward() - ref_out.mean().mul(10.).backward() - np.testing.assert_allclose( - test_input.grad.detach().cpu().numpy(), - ref_input.grad.detach().cpu().numpy(), - atol=0, rtol=1e-5) - np.testing.assert_allclose( - mlp.biases[0].grad.detach().cpu().numpy(), - ref_mlp[0].bias.grad.detach().cpu().numpy(), - atol=1e-7, rtol=1e-5) + mlp_out.mean().mul(10.0).backward() + ref_out.mean().mul(10.0).backward() + torch.testing.assert_close(test_input.grad, ref_input.grad) + torch.testing.assert_close(mlp.biases[0].grad, ref_mlp[0].bias.grad) @skipFlakyTest def test_no_bias(self): - for use_activation in ['none', 'relu', 'sigmoid']: - mlp = MLP(mlp_sizes, bias=False, activation=use_activation).cuda() - - mlp_layers = [] - for i in range(mlp.num_layers): - linear = nn.Linear(mlp_sizes[i], mlp_sizes[i + 1], bias=False) - mlp.weights[i].data.copy_(linear.weight) - mlp_layers.append(linear) - if use_activation == 'relu': - mlp_layers.append(nn.ReLU(inplace=True)) - if use_activation == 'sigmoid': - mlp_layers.append(nn.Sigmoid()) - - ref_mlp = nn.Sequential(*mlp_layers).cuda() - - test_input = torch.empty(batch_size, mlp_sizes[0], device="cuda").uniform_(-1., 1.).requires_grad_() - ref_input = test_input.clone().detach().requires_grad_() - mlp_out = mlp(test_input) - ref_out = ref_mlp(ref_input) - np.testing.assert_allclose( - mlp_out.detach().cpu().numpy(), - ref_out.detach().cpu().numpy(), - atol=1e-7, rtol=1e-5) - - # Use mean value as scalar loss. Multiply 10 to make it big enough not zero out - mlp_out.mean().mul(10.).backward() - ref_out.mean().mul(10.).backward() - np.testing.assert_allclose( - test_input.grad.detach().cpu().numpy(), - ref_input.grad.detach().cpu().numpy(), - atol=0, rtol=100) - np.testing.assert_allclose( - mlp.weights[0].grad.detach().cpu().numpy(), - ref_mlp[0].weight.grad.detach().cpu().numpy(), - atol=1e-7, rtol=100) + for use_activation in ["none", "relu", "sigmoid"]: + with self.subTest(use_activation=use_activation): + mlp = MLP(mlp_sizes, bias=False, activation=use_activation).cuda() + + mlp_layers = [] + for i in range(mlp.num_layers): + linear = nn.Linear(mlp_sizes[i], mlp_sizes[i + 1], bias=False) + mlp.weights[i].data.copy_(linear.weight) + mlp_layers.append(linear) + if use_activation == "relu": + mlp_layers.append(nn.ReLU(inplace=True)) + if use_activation == "sigmoid": + mlp_layers.append(nn.Sigmoid()) + + ref_mlp = nn.Sequential(*mlp_layers).cuda() + + test_input = ( + torch.empty(batch_size, mlp_sizes[0], device="cuda") + .uniform_(-1.0, 1.0) + .requires_grad_() + ) + ref_input = test_input.clone().detach().requires_grad_() + mlp_out = mlp(test_input) + ref_out = ref_mlp(ref_input) + torch.testing.assert_close(mlp_out, ref_out) + + # Use mean value as scalar loss. Multiply 10 to make it big enough not zero out + mlp_out.mean().mul(10.0).backward() + ref_out.mean().mul(10.0).backward() + torch.testing.assert_close(test_input.grad, ref_input.grad) + torch.testing.assert_close(mlp.weights[0].grad, ref_mlp[0].weight.grad) @skipFlakyTest def test_with_bias(self): - for use_activation in ['none', 'relu', 'sigmoid']: - mlp = MLP(mlp_sizes, bias=True, activation=use_activation).cuda() - - mlp_layers = [] - for i in range(mlp.num_layers): - linear = nn.Linear(mlp_sizes[i], mlp_sizes[i + 1], bias=True) - mlp.weights[i].data.copy_(linear.weight) - mlp.biases[i].data.copy_(linear.bias) - mlp_layers.append(linear) - if use_activation == 'relu': - mlp_layers.append(nn.ReLU(inplace=True)) - if use_activation == 'sigmoid': - mlp_layers.append(nn.Sigmoid()) - - ref_mlp = nn.Sequential(*mlp_layers).cuda() - - test_input = torch.empty(batch_size, mlp_sizes[0], device="cuda").uniform_(-1., 1.).requires_grad_() - ref_input = test_input.clone().detach().requires_grad_() - mlp_out = mlp(test_input) - ref_out = ref_mlp(ref_input) - np.testing.assert_allclose( - mlp_out.detach().cpu().numpy(), - ref_out.detach().cpu().numpy(), - atol=1e-7, rtol=1e-5) - - # Use mean value as scalar loss. Multiply 10 to make it big enough not zero out - mlp_out.mean().mul(10.).backward() - ref_out.mean().mul(10.).backward() - np.testing.assert_allclose( - test_input.grad.detach().cpu().numpy(), - ref_input.grad.detach().cpu().numpy(), - atol=0, rtol=1) - np.testing.assert_allclose( - mlp.weights[0].grad.detach().cpu().numpy(), - ref_mlp[0].weight.grad.detach().cpu().numpy(), - atol=1e-7, rtol=1) - np.testing.assert_allclose( - mlp.biases[0].grad.detach().cpu().numpy(), - ref_mlp[0].bias.grad.detach().cpu().numpy(), - atol=1e-7, rtol=1e-5) + for use_activation in ["none", "relu", "sigmoid"]: + with self.subTest(use_activation=use_activation): + mlp = MLP(mlp_sizes, bias=True, activation=use_activation).cuda() + + mlp_layers = [] + for i in range(mlp.num_layers): + linear = nn.Linear(mlp_sizes[i], mlp_sizes[i + 1], bias=True) + mlp.weights[i].data.copy_(linear.weight) + mlp.biases[i].data.copy_(linear.bias) + mlp_layers.append(linear) + if use_activation == "relu": + mlp_layers.append(nn.ReLU(inplace=True)) + if use_activation == "sigmoid": + mlp_layers.append(nn.Sigmoid()) + + ref_mlp = nn.Sequential(*mlp_layers).cuda() + + test_input = ( + torch.empty(batch_size, mlp_sizes[0], device="cuda") + .uniform_(-1.0, 1.0) + .requires_grad_() + ) + ref_input = test_input.clone().detach().requires_grad_() + mlp_out = mlp(test_input) + ref_out = ref_mlp(ref_input) + torch.testing.assert_close(mlp_out, ref_out) + + # Use mean value as scalar loss. Multiply 10 to make it big enough not zero out + mlp_out.mean().mul(10.0).backward() + ref_out.mean().mul(10.0).backward() + torch.testing.assert_close(test_input.grad, ref_input.grad) + torch.testing.assert_close(mlp.weights[0].grad, ref_mlp[0].weight.grad) + torch.testing.assert_close(mlp.biases[0].grad, ref_mlp[0].bias.grad) @skipFlakyTest def test_no_grad(self): @@ -193,22 +178,15 @@ def test_no_grad(self): ref_mlp = nn.Sequential(*mlp_layers).cuda() - test_input = torch.empty(batch_size, mlp_sizes[0], device="cuda").uniform_(-1., 1.) + test_input = torch.empty(batch_size, mlp_sizes[0], device="cuda").uniform_(-1.0, 1.0) ref_input = test_input.clone().detach() mlp_out = mlp(test_input) ref_out = ref_mlp(ref_input) - np.testing.assert_allclose( - mlp_out.detach().cpu().numpy(), - ref_out.detach().cpu().numpy(), - atol=1e-7, rtol=1e-5) + torch.testing.assert_close(mlp_out, ref_out) # Use mean value as scalar loss. Multiply 10 to make it big enough not zero out - mlp_out.mean().mul(10.).backward() - ref_out.mean().mul(10.).backward() - np.testing.assert_allclose( - mlp.weights[0].grad.detach().cpu().numpy(), - ref_mlp[0].weight.grad.detach().cpu().numpy(), - atol=1e-7, rtol=1e-5) + ref_out.mean().mul(10.0).backward() + torch.testing.assert_close(mlp.weights[0].grad, ref_mlp[0].weight.grad) def test_performance_half(self): mlp = MLP(mlp_sizes).cuda().half() @@ -223,10 +201,16 @@ def test_performance_half(self): ref_mlp = nn.Sequential(*mlp_layers).cuda().half() - test_input = torch.empty( - batch_size, mlp_sizes[0], device="cuda", dtype=torch.half).fill_(10.).requires_grad_() - ref_input = torch.empty( - batch_size, mlp_sizes[0], device="cuda", dtype=torch.half).fill_(10.).requires_grad_() + test_input = ( + torch.empty(batch_size, mlp_sizes[0], device="cuda", dtype=torch.half) + .fill_(10.0) + .requires_grad_() + ) + ref_input = ( + torch.empty(batch_size, mlp_sizes[0], device="cuda", dtype=torch.half) + .fill_(10.0) + .requires_grad_() + ) # Warm up GPU for _ in range(100): @@ -249,7 +233,8 @@ def test_performance_half(self): ref_loss.backward() torch.cuda.synchronize() stop_time = time() - print(F"\nPytorch MLP time {(stop_time - start_time) * 1000. / num_iters:.4f} ms") + ref_time = (stop_time - start_time) * 1000.0 / num_iters + print(f"\nPytorch MLP time {ref_time:.4f} ms") torch.cuda.synchronize() start_time = time() @@ -260,9 +245,16 @@ def test_performance_half(self): test_loss.backward() torch.cuda.synchronize() stop_time = time() - print(F"C++ MLP time {(stop_time - start_time) * 1000. / num_iters:.4f} ms") - #torch.cuda.profiler.stop() + actual_time = (stop_time - start_time) * 1000.0 / num_iters + print(f"C++ MLP time {actual_time:.4f} ms") + torch.cuda.profiler.stop() + self.assertLessEqual( + actual_time, + ref_time, + msg=f"Custom extension took {actual_time:.4f} while PyTorch took {ref_time:.4f}", + ) + -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() From 7c3f9d13786fa38f9c7577e8647d079557d0775d Mon Sep 17 00:00:00 2001 From: Masaki Kozuki Date: Mon, 15 Aug 2022 12:02:33 -0700 Subject: [PATCH 05/25] Use `xmlrunner.XMLTestRunner` accordingly in `tests/L0/run_test.py` (#1451) * Use xmlrunner.XMLTestRunner accordingly TODO: - [x] Remove `subTest` because it's not compatible with the current way of running L0 tests Signed-off-by: Masaki Kozuki * use `torch.testing` more to enable xmlrunner Signed-off-by: Masaki Kozuki * Remove `subTest` for xmlrunner Signed-off-by: Masaki Kozuki * removing subTest Signed-off-by: Masaki Kozuki * not depend on an env var Signed-off-by: Masaki Kozuki * fix syntax errors * open with `"wb"` * xml file per dir Signed-off-by: Masaki Kozuki * remove comment-out Signed-off-by: Masaki Kozuki * Refactor `TestTransformer`: define member methods (#5) * setUpClass to define `test_` methods Signed-off-by: Masaki Kozuki * manually define Signed-off-by: Masaki Kozuki Signed-off-by: Masaki Kozuki * add a missing test Signed-off-by: Masaki Kozuki * remove print Signed-off-by: Masaki Kozuki * remove ext Signed-off-by: Masaki Kozuki Signed-off-by: Masaki Kozuki --- tests/L0/run_mlp/test_mlp.py | 87 +++++------------------------------- 1 file changed, 10 insertions(+), 77 deletions(-) diff --git a/tests/L0/run_mlp/test_mlp.py b/tests/L0/run_mlp/test_mlp.py index 0b51e58e4..a5cfb0c3b 100644 --- a/tests/L0/run_mlp/test_mlp.py +++ b/tests/L0/run_mlp/test_mlp.py @@ -18,7 +18,7 @@ # note(crcrpar): On Ampere, this test should be run without TF32 enabled. -class TestMLP(unittest.TestCase): +class TestMLP(common_utils.TestCase): def test_creation(self): MLP(mlp_sizes) @@ -85,83 +85,13 @@ def test_mlp(self, use_activation: str, bias: bool): ref_input = test_input.clone().detach().requires_grad_() mlp_out = mlp(test_input) ref_out = ref_mlp(ref_input) - torch.testing.assert_close(mlp_out, ref_out) + self.assertEqual(mlp_out, ref_out, msg=msg) # Use mean value as scalar loss. Multiply 10 to make it big enough not zero out mlp_out.mean().mul(10.0).backward() ref_out.mean().mul(10.0).backward() - torch.testing.assert_close(test_input.grad, ref_input.grad) - torch.testing.assert_close(mlp.biases[0].grad, ref_mlp[0].bias.grad) - - @skipFlakyTest - def test_no_bias(self): - for use_activation in ["none", "relu", "sigmoid"]: - with self.subTest(use_activation=use_activation): - mlp = MLP(mlp_sizes, bias=False, activation=use_activation).cuda() - - mlp_layers = [] - for i in range(mlp.num_layers): - linear = nn.Linear(mlp_sizes[i], mlp_sizes[i + 1], bias=False) - mlp.weights[i].data.copy_(linear.weight) - mlp_layers.append(linear) - if use_activation == "relu": - mlp_layers.append(nn.ReLU(inplace=True)) - if use_activation == "sigmoid": - mlp_layers.append(nn.Sigmoid()) - - ref_mlp = nn.Sequential(*mlp_layers).cuda() - - test_input = ( - torch.empty(batch_size, mlp_sizes[0], device="cuda") - .uniform_(-1.0, 1.0) - .requires_grad_() - ) - ref_input = test_input.clone().detach().requires_grad_() - mlp_out = mlp(test_input) - ref_out = ref_mlp(ref_input) - torch.testing.assert_close(mlp_out, ref_out) - - # Use mean value as scalar loss. Multiply 10 to make it big enough not zero out - mlp_out.mean().mul(10.0).backward() - ref_out.mean().mul(10.0).backward() - torch.testing.assert_close(test_input.grad, ref_input.grad) - torch.testing.assert_close(mlp.weights[0].grad, ref_mlp[0].weight.grad) - - @skipFlakyTest - def test_with_bias(self): - for use_activation in ["none", "relu", "sigmoid"]: - with self.subTest(use_activation=use_activation): - mlp = MLP(mlp_sizes, bias=True, activation=use_activation).cuda() - - mlp_layers = [] - for i in range(mlp.num_layers): - linear = nn.Linear(mlp_sizes[i], mlp_sizes[i + 1], bias=True) - mlp.weights[i].data.copy_(linear.weight) - mlp.biases[i].data.copy_(linear.bias) - mlp_layers.append(linear) - if use_activation == "relu": - mlp_layers.append(nn.ReLU(inplace=True)) - if use_activation == "sigmoid": - mlp_layers.append(nn.Sigmoid()) - - ref_mlp = nn.Sequential(*mlp_layers).cuda() - - test_input = ( - torch.empty(batch_size, mlp_sizes[0], device="cuda") - .uniform_(-1.0, 1.0) - .requires_grad_() - ) - ref_input = test_input.clone().detach().requires_grad_() - mlp_out = mlp(test_input) - ref_out = ref_mlp(ref_input) - torch.testing.assert_close(mlp_out, ref_out) - - # Use mean value as scalar loss. Multiply 10 to make it big enough not zero out - mlp_out.mean().mul(10.0).backward() - ref_out.mean().mul(10.0).backward() - torch.testing.assert_close(test_input.grad, ref_input.grad) - torch.testing.assert_close(mlp.weights[0].grad, ref_mlp[0].weight.grad) - torch.testing.assert_close(mlp.biases[0].grad, ref_mlp[0].bias.grad) + self.assertEqual(test_input.grad, ref_input.grad, msg=msg) + self.assertEqual(mlp.weights[0].grad, ref_mlp[0].weight.grad, msg=msg) @skipFlakyTest def test_no_grad(self): @@ -182,11 +112,11 @@ def test_no_grad(self): ref_input = test_input.clone().detach() mlp_out = mlp(test_input) ref_out = ref_mlp(ref_input) - torch.testing.assert_close(mlp_out, ref_out) + self.assertEqual(mlp_out, ref_out) # Use mean value as scalar loss. Multiply 10 to make it big enough not zero out ref_out.mean().mul(10.0).backward() - torch.testing.assert_close(mlp.weights[0].grad, ref_mlp[0].weight.grad) + self.assertEqual(mlp.weights[0].grad, ref_mlp[0].weight.grad) def test_performance_half(self): mlp = MLP(mlp_sizes).cuda().half() @@ -255,6 +185,9 @@ def test_performance_half(self): ) +instantiate_device_type_tests(TestMLP, globals(), only_for=("cuda",)) + + if __name__ == "__main__": - unittest.main() + common_utils.run_tests() From 519a038b286b8f598ffbbca846f3687cc0392695 Mon Sep 17 00:00:00 2001 From: Masaki Kozuki Date: Thu, 8 Sep 2022 17:37:14 -0700 Subject: [PATCH 06/25] Update `apex.mlp` to use fp16 in `autocast` (#1477) * apex.amp migration to torch.cuda.amp Signed-off-by: Masaki Kozuki * add autocast tests Signed-off-by: Masaki Kozuki * split with and without autocast Signed-off-by: Masaki Kozuki Signed-off-by: Masaki Kozuki --- apex/_autocast_utils.py | 3 +++ apex/mlp/mlp.py | 20 ++++++++++------ tests/L0/run_mlp/test_mlp.py | 46 ++++++++++++++++++++++++------------ 3 files changed, 47 insertions(+), 22 deletions(-) diff --git a/apex/_autocast_utils.py b/apex/_autocast_utils.py index e86c6c6a5..3a92a83f3 100644 --- a/apex/_autocast_utils.py +++ b/apex/_autocast_utils.py @@ -3,6 +3,9 @@ import torch +__all__ = ["_cast_if_autocast_enabled"] + + def _get_autocast_dtypes() -> Sequence[torch.dtype]: if torch.cuda.is_bf16_supported(): return [torch.half, torch.bfloat16] diff --git a/apex/mlp/mlp.py b/apex/mlp/mlp.py index be9ce4e5a..be0901eb8 100644 --- a/apex/mlp/mlp.py +++ b/apex/mlp/mlp.py @@ -1,9 +1,12 @@ from copy import copy import math + import torch from torch import nn + +from apex._autocast_utils import _cast_if_autocast_enabled import mlp_cuda -from .. import amp + class MlpFunction(torch.autograd.Function): @staticmethod @@ -21,8 +24,11 @@ def backward(ctx, grad_o): del ctx.outputs return (None, None, *grads) -# TODO(crcrpar): Should make this compatible with torch.cuda.amp -mlp_function = amp.half_function(MlpFunction.apply) + +def mlp_function(bias, activation, *args): + autocast_args = _cast_if_autocast_enabled(bias, activation, *args) + return MlpFunction.apply(*autocast_args) + class MLP(torch.nn.Module): """Launch MLP in C++ @@ -33,16 +39,16 @@ class MLP(torch.nn.Module): relu (bool): Default True """ def __init__(self, mlp_sizes, bias=True, activation='relu'): - super(MLP, self).__init__() + super().__init__() self.num_layers = len(mlp_sizes) - 1 self.mlp_sizes = copy(mlp_sizes) self.bias = 1 if bias else 0 - if activation is 'none': + if activation == 'none': self.activation = 0 - elif activation is 'relu': + elif activation == 'relu': self.activation = 1 - elif activation is 'sigmoid': + elif activation == 'sigmoid': self.activation = 2 else: raise TypeError("activation must be relu or none.") diff --git a/tests/L0/run_mlp/test_mlp.py b/tests/L0/run_mlp/test_mlp.py index a5cfb0c3b..3b261766d 100644 --- a/tests/L0/run_mlp/test_mlp.py +++ b/tests/L0/run_mlp/test_mlp.py @@ -53,13 +53,7 @@ def test_numeric(self): self.assertEqual(test_input.grad, ref_input.grad) self.assertEqual(mlp.biases[0].grad, ref_mlp[0].bias.grad) - @common_utils.parametrize( - "use_activation,bias", - list(product(("none", "relu", "sigmoid"), (True, False))), - ) - def test_mlp(self, use_activation: str, bias: bool): - # for use_activation in ["none", "relu", "sigmoid"]: - msg = f"activation: {use_activation}, bias: {bias}" + def _test_mlp_impl(self, use_activation: str, bias: bool, enable_autocast: bool): mlp = MLP(mlp_sizes, bias=bias, activation=use_activation).cuda() mlp_layers = [] @@ -83,15 +77,37 @@ def test_mlp(self, use_activation: str, bias: bool): .requires_grad_() ) ref_input = test_input.clone().detach().requires_grad_() - mlp_out = mlp(test_input) - ref_out = ref_mlp(ref_input) - self.assertEqual(mlp_out, ref_out, msg=msg) - # Use mean value as scalar loss. Multiply 10 to make it big enough not zero out - mlp_out.mean().mul(10.0).backward() - ref_out.mean().mul(10.0).backward() - self.assertEqual(test_input.grad, ref_input.grad, msg=msg) - self.assertEqual(mlp.weights[0].grad, ref_mlp[0].weight.grad, msg=msg) + with torch.cuda.amp.autocast_mode.autocast(enabled=enable_autocast): + mlp_out = mlp(test_input) + mlp_loss = mlp_out.mean().mul(10.0) + # Use mean value as scalar loss. Multiply 10 to make it big enough not zero out + ref_out = ref_mlp(ref_input) + ref_loss = ref_out.mean().mul(10.0) + + mlp_loss.backward() + ref_loss.backward() + if enable_autocast: + self.assertEqual(mlp_out.dtype, torch.float16) + self.assertEqual(ref_out.dtype, torch.float16) + else: + self.assertEqual(mlp_out, ref_out) + self.assertEqual(test_input.grad, ref_input.grad) + self.assertEqual(mlp.weights[0].grad, ref_mlp[0].weight.grad) + + @common_utils.parametrize( + "use_activation,bias", + list(product(("none", "relu", "sigmoid"), (True, False))), + ) + def test_mlp(self, use_activation: str, bias: bool): + self._test_mlp_impl(use_activation, bias, enable_autocast=False) + + @common_utils.parametrize( + "use_activation,bias", + list(product(("none", "relu", "sigmoid"), (True, False))), + ) + def test_mlp_autocast_fp16(self, use_activation: str, bias: bool): + self._test_mlp_impl(use_activation, bias, enable_autocast=True) @skipFlakyTest def test_no_grad(self): From f4c4b86045ec66f760ce2c57a11340fca0a5ce48 Mon Sep 17 00:00:00 2001 From: hubertlu-tw Date: Wed, 28 Dec 2022 01:41:39 +0000 Subject: [PATCH 07/25] Skip flaky test for ROCm --- tests/L0/run_mlp/test_mlp.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/L0/run_mlp/test_mlp.py b/tests/L0/run_mlp/test_mlp.py index 3b261766d..dcc97aefb 100644 --- a/tests/L0/run_mlp/test_mlp.py +++ b/tests/L0/run_mlp/test_mlp.py @@ -134,6 +134,7 @@ def test_no_grad(self): ref_out.mean().mul(10.0).backward() self.assertEqual(mlp.weights[0].grad, ref_mlp[0].weight.grad) + @skipFlakyTest def test_performance_half(self): mlp = MLP(mlp_sizes).cuda().half() From bb7b64db53c3398ad19cc3839ad525e92ceab557 Mon Sep 17 00:00:00 2001 From: Sandeep Subramanian Date: Tue, 23 Aug 2022 17:36:34 -0700 Subject: [PATCH 08/25] Label smoothing in vocab parallel cross entropy (#1457) * Label smoothing in vocab parallel cross entropy Signed-off-by: MaximumEntropy * Fix context saving Signed-off-by: MaximumEntropy * Remove .item() calls Signed-off-by: MaximumEntropy * Update tests Signed-off-by: MaximumEntropy Signed-off-by: MaximumEntropy --- .../tensor_parallel/cross_entropy.py | 41 ++++++++++++++++--- .../L0/run_transformer/test_cross_entropy.py | 8 ++-- 2 files changed, 40 insertions(+), 9 deletions(-) diff --git a/apex/transformer/tensor_parallel/cross_entropy.py b/apex/transformer/tensor_parallel/cross_entropy.py index 3918645c7..b14d7c19c 100644 --- a/apex/transformer/tensor_parallel/cross_entropy.py +++ b/apex/transformer/tensor_parallel/cross_entropy.py @@ -22,7 +22,7 @@ class _VocabParallelCrossEntropy(torch.autograd.Function): @staticmethod - def forward(ctx, vocab_parallel_logits, target): + def forward(ctx, vocab_parallel_logits, target, label_smoothing=0.0): # Maximum value along vocab dimension across all GPUs. logits_max = torch.max(vocab_parallel_logits, dim=-1)[0] @@ -72,6 +72,27 @@ def forward(ctx, vocab_parallel_logits, target): # Store softmax, target-mask and masked-target for backward pass. exp_logits.div_(sum_exp_logits.unsqueeze(dim=-1)) + + vocab_size = exp_logits.size(-1) + if label_smoothing > 0: + """ + We'd like to assign 1 / (K - 1) probability mass to every index that is not the ground truth. + = (1 - alpha) * y_gt + alpha * mean(y_{i for i != gt}) + = (1 - alpha) * y_gt + (alpha / (K - 1)) * \sum_{i != gt} y_i + = ((K - 1) * (1 - alpha) / (K - 1)) * y_gt + (alpha / (K - 1)) * \sum_{i != gt} y_i + = (K * (1 - alpha) - 1) / (K - 1)) * y_gt + (alpha / (K - 1)) * \sum_{i} y_i + = (1 - (alpha * K) / (K - 1)) * y_gt + ( (alpha * K) / (K - 1) ) * \sum_{i} y_i / K + From: https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/common/losses/smoothed_cross_entropy.py + """ + assert 1.0 > label_smoothing > 0.0 + smoothing = label_smoothing * vocab_size / (vocab_size - 1) + + # Exp logits at this point are normalized probabilities. So we can just take the log to get log-probs. + log_probs = torch.log(exp_logits) + mean_log_probs = log_probs.mean(dim=-1) + loss = (1.0 - smoothing) * loss - smoothing * mean_log_probs + + ctx.label_smoothing, ctx.vocab_size = label_smoothing, vocab_size ctx.save_for_backward(exp_logits, target_mask, masked_target_1d) return loss @@ -81,6 +102,7 @@ def backward(ctx, grad_output): # Retreive tensors from the forward path. softmax, target_mask, masked_target_1d = ctx.saved_tensors + label_smoothing, vocab_size = ctx.label_smoothing, ctx.vocab_size # All the inputs have softmax as thier gradient. grad_input = softmax @@ -90,14 +112,23 @@ def backward(ctx, grad_output): # Add the gradient from matching classes. arange_1d = torch.arange(start=0, end=grad_2d.size()[0], device=grad_2d.device) - grad_2d[arange_1d, masked_target_1d] -= 1.0 - target_mask.view(-1).float() + + softmax_update = 1.0 - target_mask.view(-1).float() + + if label_smoothing > 0: + smoothing = label_smoothing * vocab_size / (vocab_size - 1) + grad_2d[arange_1d, masked_target_1d] -= (1.0 - smoothing) * softmax_update + average_grad = 1 / vocab_size + grad_2d[arange_1d, :] -= smoothing * average_grad + else: + grad_2d[arange_1d, masked_target_1d] -= softmax_update # Finally elementwise multiplication with the output gradients. grad_input.mul_(grad_output.unsqueeze(dim=-1)) - return grad_input, None + return grad_input, None, None -def vocab_parallel_cross_entropy(vocab_parallel_logits, target): +def vocab_parallel_cross_entropy(vocab_parallel_logits, target, label_smoothing=0.0): """Helper function for the cross entropy.""" - return _VocabParallelCrossEntropy.apply(vocab_parallel_logits, target) + return _VocabParallelCrossEntropy.apply(vocab_parallel_logits, target, label_smoothing) diff --git a/tests/L0/run_transformer/test_cross_entropy.py b/tests/L0/run_transformer/test_cross_entropy.py index 657ecbf6d..d1e1f9b1e 100644 --- a/tests/L0/run_transformer/test_cross_entropy.py +++ b/tests/L0/run_transformer/test_cross_entropy.py @@ -18,7 +18,7 @@ def torch_cross_entropy( - batch_size: int, seq_length: int, vocab_size: int, logits_scale: float, seed: int, + batch_size: int, seq_length: int, vocab_size: int, logits_scale: float, seed: int, label_smoothing: float = 0.0 ) -> Tuple[torch.Tensor, torch.Tensor]: set_random_seed(seed) identity = IdentityLayer( @@ -28,7 +28,7 @@ def torch_cross_entropy( target = torch.cuda.LongTensor(size=(batch_size, seq_length)).random_(0, vocab_size) loss = ( F.cross_entropy( - logits.view(-1, logits.size()[-1]), target.view(-1), reduction="none" + logits.view(-1, logits.size()[-1]), target.view(-1), reduction="none", label_smoothing=label_smoothing ) .view_as(target) .mean() @@ -38,7 +38,7 @@ def torch_cross_entropy( def tensor_sharded_cross_entropy( - batch_size, seq_length, vocab_size, logits_scale, seed + batch_size, seq_length, vocab_size, logits_scale, seed, label_smoothing=0.0 ): set_random_seed(seed) identity = IdentityLayer( @@ -48,7 +48,7 @@ def tensor_sharded_cross_entropy( logits_parallel = tensor_parallel.scatter_to_tensor_model_parallel_region(logits) target = torch.cuda.LongTensor(size=(batch_size, seq_length)).random_(0, vocab_size) logits_parallel_ = logits_parallel.clone().detach() - loss = cross_entropy.vocab_parallel_cross_entropy(logits_parallel, target).mean() + loss = cross_entropy.vocab_parallel_cross_entropy(logits_parallel, target, label_smoothing=label_smoothing).mean() loss.backward() # check for mutation assert torch.equal(logits_parallel_, logits_parallel) From 3225191b17511a6214abc4e161a8bb734fb59a2e Mon Sep 17 00:00:00 2001 From: Masaki Kozuki Date: Fri, 16 Sep 2022 15:50:12 -0700 Subject: [PATCH 09/25] introducing `APEX_RUN_WITH_SLOW_TESTS` env var (#1489) Signed-off-by: Masaki Kozuki --- tests/L0/run_transformer/test_transformer_module.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/L0/run_transformer/test_transformer_module.py b/tests/L0/run_transformer/test_transformer_module.py index 7be1e5654..57f454a67 100644 --- a/tests/L0/run_transformer/test_transformer_module.py +++ b/tests/L0/run_transformer/test_transformer_module.py @@ -5,6 +5,8 @@ import unittest +# TODO(crcrpar): should move this to `apex._testing` or whatever in the near future. +RUN_SLOW_TESTS = os.getenv('APEX_RUN_WITH_SLOW_TESTS', '0') == '1' SEVERALGPU_TEST = [ "bert_minimal_test", "gpt_minimal_test", @@ -33,7 +35,6 @@ def get_launch_option(test_filename) -> Tuple[bool, str]: def get_test_command(test_file: str) -> str: python_executable_path = sys.executable - is_denied = False should_skip, launch_option = get_launch_option(test_file) if should_skip: return "" @@ -63,7 +64,6 @@ def get_test_command(test_file: str) -> str: def _get_test_file(key): - python_executable_path = sys.executable directory = os.path.dirname(__file__) test_file = [ os.path.join(directory, f) @@ -73,6 +73,7 @@ def _get_test_file(key): return test_file +@unittest.skipUnless(RUN_SLOW_TESTS, "this test takes long.") class TestTransformer(unittest.TestCase): def _test_impl(self, key: str): test_file = _get_test_file(key) From 7c3cae347162581abf033cb89f697060489cd70a Mon Sep 17 00:00:00 2001 From: Tim Moon <4406448+timmoon10@users.noreply.github.com> Date: Mon, 19 Sep 2022 23:25:13 -0700 Subject: [PATCH 10/25] [contrib][DistributedFusedAdam] Support overlapped grad sync with Megatron pipeline parallelism (#1475) * Refactor how dist Adam handles overlapped grad sync Each grad bucket independently keeps track of grads that have been generated. Add helper function to create callback functions. Change default param arg in grad norm functions to None. Perform communication for checkpointing in main stream to avoid memory pool overheads. * Support Megatron pipeline parallelism with async grad reduction Enables async grad reduction in first pipeline stage during last backward pass, and disables async grad reduction in all other pipeline stages. * Review suggestions from crcrpar Add unit test for pipeline parallelism with custom sync context. Style tweaks. * Use unittest assert functions in pipeline parallelism test Review suggestion from crcrpar --- .../optimizers/distributed_fused_adam.py | 236 ++++++++++-------- .../schedules/fwd_bwd_no_pipelining.py | 14 +- ...fwd_bwd_pipelining_without_interleaving.py | 28 +++ .../test_pipeline_parallel_fwd_bwd.py | 107 ++++++++ 4 files changed, 264 insertions(+), 121 deletions(-) diff --git a/apex/contrib/optimizers/distributed_fused_adam.py b/apex/contrib/optimizers/distributed_fused_adam.py index 550068022..117d7294b 100644 --- a/apex/contrib/optimizers/distributed_fused_adam.py +++ b/apex/contrib/optimizers/distributed_fused_adam.py @@ -170,6 +170,8 @@ def __init__(self): self.status = DistributedFusedAdam.GradientStatus.READY # Request object for asynchronous communication self.sync_request = None + # Params that have generated grads + self.grads_generated = set() def sync_wait(self): """Wait for asynchronous communication to finish""" @@ -298,7 +300,6 @@ def __init__(self, # Objects for gradient synchronization self._grads_buckets = collections.defaultdict(self.GradientBucket) - self._grads_generated = set() self._pipeline_streams = [torch.cuda.Stream() for _ in range(self.pipeline_size)] # Divide gradients by factor before optimizer step. Used for @@ -322,16 +323,39 @@ def __init__(self, # Attach hooks for gradient synchronization self._register_post_backward_hooks() + # Allocate contiguous gradient buffer if needed + if self.contiguous_grad_buffer: + self._init_grad_buffer() + + def _make_post_backward_hook(self, param, param_group_id, param_id): + """Create callback function to call after param generates grad + + Lazily initialize parameter and try launching grad sync. + + """ + def hook(*unused): + with self._lock: + need_to_initialize = 'fragments' not in self.state[param] + if need_to_initialize: + self._init_param_state(param, param_group_id, param_id) + if self.greedy_grad_copy: + self._grad_copy(param) + if self.overlap_grad_sync: + self._try_start_bucket_grad_sync( + params=[param], + ignore_last_bucket=need_to_initialize, + ) + return hook + def _register_post_backward_hooks(self): """Attach hooks for gradient synchronization - Optimizer state for parameters are initialized lazily as they - are encountered in the backward pass. + Also synchronizes param values between processes and counts + number of parameters being optimized. """ self._num_grads = 0 - grad_buffer_size = 0 - self._lock = threading.Lock() + self._lock = threading.Lock() # Not sure if needed self._grad_accs = [] for param_group_id, group in enumerate(self.param_groups): for param_id, param in enumerate(group['params']): @@ -343,40 +367,34 @@ def _register_post_backward_hooks(self): if param.requires_grad: self._num_grads += 1 - # Callback after gradient is generated - def wrapper(p, p_group_id, p_id): - p_tmp = p.expand_as(p) - grad_acc = p_tmp.grad_fn.next_functions[0][0] - def reduction_hook(*unused): - with self._lock: - if 'fragments' not in self.state[p]: - self._init_param_state(p, p_group_id, p_id) - if self.greedy_grad_copy: - self._grad_copy(p) - if self.overlap_grad_sync: - self._try_start_bucket_grad_sync( - params=[p], - ignore_last_bucket=True, - ) - grad_acc.register_hook(reduction_hook) - self._grad_accs.append(grad_acc) - wrapper(param, param_group_id, param_id) - - # Gradient size, with padding for alignment + # Register callback for after grad is generated + param_tmp = param.expand_as(param) + grad_acc = param_tmp.grad_fn.next_functions[0][0] + hook = self._make_post_backward_hook( + param, + param_group_id, + param_id, + ) + grad_acc.register_hook(hook) + self._grad_accs.append(grad_acc) + + def _init_grad_buffer(self): + """Allocate contiguous buffer for grad buckets""" + grad_buffer_size = 0 + for group in self.param_groups: + for param in group['params']: + if param.requires_grad: grad_size = _round_to_multiple(param.numel(), self.alignment) grad_buffer_size += grad_size - - # Allocate contiguous gradient buffer if needed - if self.contiguous_grad_buffer: - grad_buffer_size = _round_to_multiple( - grad_buffer_size, - self.bucket_size, - ) - self._grad_buffer = torch.zeros( - [grad_buffer_size], - dtype=self.dtype, - device=self.device, - ) + grad_buffer_size = _round_to_multiple( + grad_buffer_size, + self.bucket_size, + ) + self._grad_buffer = torch.zeros( + [grad_buffer_size], + dtype=self.dtype, + device=self.device, + ) def init_params(self, params=None): """Initialize optimizer state for parameters @@ -516,8 +534,7 @@ def zero_grad(self, set_to_none=True): param.grad.zero_() # Reset other state - self._grads_generated = set() - self._inv_grad_scale = torch.full([1], 1.0, dtype=self.dtype, device=self.device) + self._grad_scale = torch.full([], 1.0, dtype=self.dtype, device=self.device) self._grad_norm = None def _grad_copy(self, param): @@ -606,13 +623,10 @@ def _force_bucket_grad_sync(self): device=self.device, ) - # Reset set of generated gradients - self._grads_generated = set() - def _try_start_bucket_grad_sync( self, params=[], - ignore_last_bucket=True, + ignore_last_bucket=False, ): """Launches gradient synchronization if enough buckets are ready @@ -632,38 +646,28 @@ def _try_start_bucket_grad_sync( # Register params that have generated grads for param in params: - self._grads_generated.add(param) for fragment in self.state[param]['fragments']: bucket_id = fragment.bucket_id + bucket = self._grads_buckets[bucket_id] bucket_fragments = self.state['buckets'][bucket_id].fragments - is_filled = True - for other_fragment in reversed(bucket_fragments): - param_group_id = other_fragment.param_group_id - param_id = other_fragment.param_id - other_param = self.param_groups[param_group_id]['params'][param_id] - if other_param not in self._grads_generated: - is_filled = False - break - if is_filled: - bucket = self._grads_buckets[bucket_id] + bucket.grads_generated.add(param) + if len(bucket.grads_generated) == len(bucket_fragments): bucket.status = self.GradientStatus.FULLY_FILLED # Launch reductions if enough buckets are ready - if len(self._grads_generated) == self._num_grads: - self._force_bucket_grad_sync() - else: - filled_buckets = [] - for bucket_id, bucket in sorted(self._grads_buckets.items()): - if ignore_last_bucket and bucket_id == len(self.state['buckets'])-1: - continue - if bucket.status == self.GradientStatus.FULLY_FILLED: - filled_buckets.append(bucket) - pipeline_size = _round_to_multiple( - len(filled_buckets), - self.pipeline_size, - ) - if pipeline_size > 0: - self._start_bucket_grad_sync(filled_buckets[:pipeline_size]) + filled_buckets = [] + for bucket_id, bucket in sorted(self._grads_buckets.items()): + if ignore_last_bucket and bucket_id == len(self.state['buckets'])-1: + continue + if bucket.status == self.GradientStatus.FULLY_FILLED: + filled_buckets.append(bucket) + pipeline_size = _round_to_multiple( + len(filled_buckets), + self.pipeline_size, + round_up=False, + ) + if pipeline_size > 0: + self._start_bucket_grad_sync(filled_buckets[:pipeline_size]) def _start_bucket_grad_sync(self, buckets): """Synchronize gradient buckets @@ -692,13 +696,22 @@ def _start_bucket_grad_sync(self, buckets): stream.wait_stream(main_stream) for i, bucket in enumerate(buckets): bucket.status = self.GradientStatus.SYNCING - stream = self._pipeline_streams[i % self.pipeline_size] - with torch.cuda.stream(stream): - - # Reduce-scatter over distributed process group - bucket.sync_wait() - if self.distributed_size == 1: - bucket.sync_grads_shard = bucket.grads_bucket + bucket.grads_generated.clear() + bucket.sync_wait() + if self.distributed_size == 1: + bucket.sync_grads_shard = bucket.grads_bucket + else: + bucket.sync_grads_shard = torch.zeros( + [self.shard_size], + dtype=self.grad_sync_dtype, + device=self.device, + ) + grads_bucket_shards = [ + bucket.grads_bucket[i*self.shard_size:(i+1)*self.shard_size] + for i in range(self.distributed_size) + ] + if self._reduce_scatter_no_copy: + no_copy_kwarg = { 'no_copy': True } else: with torch.cuda.stream(main_stream): bucket.sync_grads_shard = torch.zeros( @@ -805,7 +818,7 @@ def grad_sync(self): ) self._force_bucket_grad_sync() - def _local_grad_norm(self, parameters=[], norm_type=2.0): + def _local_grad_norm(self, parameters=None, norm_type=2.0): """Local contribution to parameter gradient norm Returns square of 2-norm. Other norms are not yet supported. @@ -821,7 +834,7 @@ def _local_grad_norm(self, parameters=[], norm_type=2.0): # Make sure that gradients have been reduced self.grad_sync() - if not parameters or len(parameters) == self._num_grads: + if parameters is None or len(parameters) == self._num_grads: # Compute norm of all local gradients dummy_overflow_buf = torch.zeros([1], dtype=torch.int32, device='cuda') grad_norm_sq = multi_tensor_applier( @@ -852,7 +865,7 @@ def _local_grad_norm(self, parameters=[], norm_type=2.0): return grad_norm_sq.detach().view([]) - def grad_norm(self, parameters=[], norm_type=2.0, force=False): + def grad_norm(self, parameters=None, norm_type=2.0, force=False): """Gradient norm of parameters in optimizer The norm is computed over all gradients together, as if they @@ -886,7 +899,7 @@ def grad_norm(self, parameters=[], norm_type=2.0, force=False): self._grad_norm = grad_norm_sq.sqrt() return self._grad_norm.detach() - def clip_grad_norm(self, max_norm, parameters=[], norm_type=2.0): + def clip_grad_norm(self, max_norm, parameters=None, norm_type=2.0): """Clips gradient norm of parameters in optimizer The norm is computed over all gradients together, as if they @@ -1180,28 +1193,25 @@ def state_dict(self, gather_on_root=True): # Split data into chunks and gather on root rank # Note: Assuming we are using the NCCL backend, communication # must happen on the GPU. We split the data into fixed-size - # chunks so that the GPU memory usage is limited to - # (chunk_size * distributed_size) bytes. + # chunks to limit GPU memory usage. # TODO: Avoid chunking with direct communication between CPUs main_stream = torch.cuda.current_stream() for stream in self._pipeline_streams: stream.wait_stream(main_stream) for stream_id, offset in enumerate(range(0, max_state_size, chunk_size)): stream_id %= self.pipeline_size - - # Buffers for chunk - if self.distributed_rank == 0: - gathered_chunks = [ - gathered_chunks_buffers[stream_id][i*chunk_size:(i+1)*chunk_size] - for i in range(self.distributed_size) - ] - else: - chunk = chunk_buffers[stream_id] - - # Perform communication on parallel stream stream = self._pipeline_streams[stream_id] with torch.cuda.stream(stream): + # Buffers for chunk + if self.distributed_rank == 0: + gathered_chunks = [ + gathered_chunks_buffers[stream_id][i*chunk_size:(i+1)*chunk_size] + for i in range(self.distributed_size) + ] + else: + chunk = chunk_buffers[stream_id] + # Copy to GPU if self.distributed_rank != 0 and offset < local_state_size: local_chunk_size = min(chunk_size, local_state_size-offset) @@ -1216,24 +1226,30 @@ def state_dict(self, gather_on_root=True): ) # Gather on root - if self.distributed_rank == 0: - if self._gather_no_copy: - no_copy_kwarg = { 'no_copy': True } + # Note: Call in main stream to avoid memory pool + # overheads from internal memory allocations in + # gather. + main_stream.wait_stream(stream) + with torch.cuda.stream(main_stream): + if self.distributed_rank == 0: + if self._gather_no_copy: + no_copy_kwarg = { 'no_copy': True } + else: + no_copy_kwarg = {} + torch.distributed.gather( + gathered_chunks[0], + gathered_chunks, + dst=self._process_group_ranks[0], + group=self.process_group, + **no_copy_kwarg, + ) else: - no_copy_kwarg = {} - torch.distributed.gather( - gathered_chunks[0], - gathered_chunks, - dst=self._process_group_ranks[0], - group=self.process_group, - **no_copy_kwarg, - ) - else: - torch.distributed.gather( - chunk, - dst=self._process_group_ranks[0], - group=self.process_group, - ) + torch.distributed.gather( + chunk, + dst=self._process_group_ranks[0], + group=self.process_group, + ) + stream.wait_stream(main_stream) # Copy back to CPU if self.distributed_rank == 0: diff --git a/apex/transformer/pipeline_parallel/schedules/fwd_bwd_no_pipelining.py b/apex/transformer/pipeline_parallel/schedules/fwd_bwd_no_pipelining.py index 5500085a4..6bd3ff139 100644 --- a/apex/transformer/pipeline_parallel/schedules/fwd_bwd_no_pipelining.py +++ b/apex/transformer/pipeline_parallel/schedules/fwd_bwd_no_pipelining.py @@ -1,4 +1,4 @@ -from contextlib import contextmanager +import contextlib from typing import List, Union, Optional import torch @@ -20,14 +20,6 @@ _logger = get_transformer_logger(__name__) -@contextmanager -def placeholder_handler(): - try: - yield - finally: - pass - - def forward_backward_no_pipelining( forward_step_func: FwdStepFunc, batch: Batch, @@ -59,7 +51,7 @@ def forward_backward_no_pipelining( disable_autocast: Turn off `enabled` flag of `torch.cuda.amp.autocast` if :obj:`True`. Should be used when your forward and loss computation is in the autocast context to avoid unnecesarily nest autocast context. - custom_sync_context_handler: + custom_sync_context_handler: Context manager to disable asynchronous gradient reductions. **kwargs: Added to handle `tensor_shape` which has no effect on this function. Returns: @@ -77,7 +69,7 @@ def forward_backward_no_pipelining( elif isinstance(model, torch.nn.parallel.distributed.DistributedDataParallel): context_handler = model.no_sync else: - context_handler = placeholder_handler + context_handler = contextlib.nullcontext losses_reduced = [] input_tensor, output_tensor_grad = None, None diff --git a/apex/transformer/pipeline_parallel/schedules/fwd_bwd_pipelining_without_interleaving.py b/apex/transformer/pipeline_parallel/schedules/fwd_bwd_pipelining_without_interleaving.py index 5dc2933ba..30d903b5a 100644 --- a/apex/transformer/pipeline_parallel/schedules/fwd_bwd_pipelining_without_interleaving.py +++ b/apex/transformer/pipeline_parallel/schedules/fwd_bwd_pipelining_without_interleaving.py @@ -1,3 +1,4 @@ +import contextlib from typing import Union, List, Optional, Sequence import warnings @@ -239,6 +240,7 @@ def forward_backward_pipelining_without_interleaving( deallocate_pipeline_outputs: bool = False, async_comm: bool = False, sequence_parallel_enabled: bool = False, + custom_sync_context_handler = None, **kwargs, ) -> List[Union[torch.Tensor, Sequence[torch.Tensor]]]: """Run non-interleaved 1F1B schedule, with communication between pipeline stages. @@ -269,9 +271,18 @@ def forward_backward_pipelining_without_interleaving( sequence_parallel_enabled: Set to :obj:`True` for this function to handle sequence length. When :obj:`True`, the sequence length on each tensor model parallel rank is updated to :math:`original\_sequence\_length / tensor\_model\_parallel\_world\_size`. + custom_sync_context_handler: Context manager to disable + asynchronous gradient reductions. In the first pipeline + stage, asynchronous gradient reductions are enabled in the + backward pass of the last microbatch. In other pipeline + stages, asynchronous gradient reductions are disabled and + gradients must be manually synchronized by the caller (in + practice the runtime is covered up by the bubble + overhead). Returns: a list of loss `torch.Tensor`s if the last stage, empty list otherwise. + """ # timers = get_timers() @@ -287,6 +298,15 @@ def forward_backward_pipelining_without_interleaving( raise RuntimeError(msg) model: torch.nn.Module = model[0] + # Disable async grad reductions + if custom_sync_context_handler is not None: + context_handler = custom_sync_context_handler + else: + context_handler = contextlib.nullcontext + context = context_handler() + context.__enter__() + context_is_active = True + # Compute number of warmup microbatches. num_microbatches: int = get_num_microbatches() num_warmup_microbatches: int = ( @@ -456,6 +476,9 @@ def forward_backward_pipelining_without_interleaving( _logger.info("Cooldown phase") if not forward_only: for i in range(num_warmup_microbatches): + if i == num_warmup_microbatches-1 and rank == 0: + context.__exit__(None, None, None) + context_is_active = False _logger.debug(f"cooldown iter: {i} / {num_warmup_microbatches}") input_tensor = input_tensors.pop(0) output_tensor = output_tensors.pop(0) @@ -486,4 +509,9 @@ def forward_backward_pipelining_without_interleaving( sequence_parallel_enabled=sequence_parallel_enabled, ) + # Make sure to exit context handler for async grad reductions + if context_is_active: + context.__exit__(None, None, None) + context_is_active = False + return losses_reduced diff --git a/tests/L0/run_transformer/test_pipeline_parallel_fwd_bwd.py b/tests/L0/run_transformer/test_pipeline_parallel_fwd_bwd.py index a409c40f2..11b59d1d3 100644 --- a/tests/L0/run_transformer/test_pipeline_parallel_fwd_bwd.py +++ b/tests/L0/run_transformer/test_pipeline_parallel_fwd_bwd.py @@ -1,3 +1,4 @@ +import contextlib import logging import itertools import re @@ -442,6 +443,112 @@ def test_pipelining_without_interleaving_sequence_parallel_encoder_or_decoder(se def test_pipelining_without_interleaving_sequence_parallel_encoder_or_decoder_half(self) -> None: self._forward_backward_test_impl(forward_only=False, sequence_parallel_enabled=True, model_type=ModelType.encoder_or_decoder, dtype=torch.half) +@unittest.skipIf(torch.cuda.device_count() < 4 or torch.cuda.device_count() % 2 != 0, "Requires >= 4 GPUs") +class NcclPipelineParallelWithCustomSyncContextHandler(NcclDistributedTestBase): + + GLOBAL_BATCH_SIZE = 2 + MICRO_BATCH_SIZE = 1 + HIDDEN_SIZE = 1 + + @property + def world_size(self) -> int: + return min(torch.cuda.device_count(), 8) + + def test_pipelining_without_interleaving_with_custom_sync_context_handler(self) -> None: + + # Parallel configuration + world_size = torch.cuda.device_count() + tensor_model_parallel_world_size = 1 + data_parallel_size = 2 + pipeline_model_parallel_world_size = world_size // data_parallel_size + + # Initialize pipeline parallelism + parallel_state.initialize_model_parallel( + tensor_model_parallel_size_=tensor_model_parallel_world_size, + pipeline_model_parallel_size_=pipeline_model_parallel_world_size, + virtual_pipeline_model_parallel_size_=None, + ) + pp_utils._reconfigure_microbatch_calculator( + rank=parallel_state.get_tensor_model_parallel_rank(), + rampup_batch_size=None, + global_batch_size=self.GLOBAL_BATCH_SIZE, + micro_batch_size=self.MICRO_BATCH_SIZE, + data_parallel_size=parallel_state.get_data_parallel_world_size(), + ) + pp_utils.update_num_microbatches(0) + + # Construct synthetic data + dtype = get_dtype_for_comparison() + hidden_size = self.HIDDEN_SIZE + microbatch_size = self.MICRO_BATCH_SIZE + global_batch_shape = ( + self.GLOBAL_BATCH_SIZE + // parallel_state.get_data_parallel_world_size(), + hidden_size, + hidden_size, + ) + batch = None + if parallel_state.is_pipeline_first_stage(): + batch = (torch.ones(global_batch_shape, dtype=dtype).cuda(), ) + + # Construct model + model = build_model( + testing_utils.model_provider_func, + wrap_with_ddp=True, + virtual_pipeline_model_parallel_size=None, + hidden_size=hidden_size, + )[0] + model = model.to(dtype) + model.module.apply(get_init_weights_func(0)) + + # Construct custom context + has_entered_context = False + has_exited_context = False + no_grad_at_context_exit = False + @contextlib.contextmanager + def custom_context(): + try: + nonlocal has_entered_context + has_entered_context = True + yield + finally: + nonlocal has_exited_context, no_grad_at_context_exit + has_exited_context = True + no_grad_at_context_exit = all( + p.grad is None + for p in model.parameters() + ) + + # Training step with pipeline parallelism + loss = forward_backward_pipelining_without_interleaving( + testing_utils.fwd_step_func, + batch, + model, + forward_only=False, + tensor_shape=(microbatch_size, hidden_size, hidden_size), + dtype=dtype, + async_comm=False, + grad_scaler=None, + deallocate_pipeline_outputs=False, + sequence_parallel_enabled=False, + custom_sync_context_handler=custom_context, + ) + torch.cuda.synchronize() + + # Check context behavior + self.assertTrue(has_entered_context, 'Has not entered custom sync context') + self.assertTrue(has_exited_context, 'Has not exited custom sync context') + self.assertEqual( + no_grad_at_context_exit, + parallel_state.is_pipeline_first_stage(), + 'Expected to exit custom sync context ' + 'before backward pass in first pipeline stage ' + 'and after backward pass in other pipeline stages' + ) + + # Clean up + parallel_state.destroy_model_parallel() + if __name__ == "__main__": common_utils.run_tests() From 23ef6ff64445e804cb0c85c70a49d77627cf0db2 Mon Sep 17 00:00:00 2001 From: Masaki Kozuki Date: Tue, 11 Oct 2022 17:31:14 -0700 Subject: [PATCH 11/25] [transformer] Allow for skipping stream sync (#1505) * Optionally disable stream synchronization after batched p2p communication * Add test cases with `sync_batch_comm=False` only when https://github.com/pytorch/pytorch/pull/82450 is included in pytorch. Signed-off-by: Masaki Kozuki * utilize existing test methods Signed-off-by: Masaki Kozuki * consistent naming Signed-off-by: Masaki Kozuki Co-authored-by: Aidyn-A * silly boy, to skip the sync, set False Signed-off-by: Masaki Kozuki * cosmetic Signed-off-by: Masaki Kozuki * Test with async_pipelinign w/o sync after batch_isend_irecv Signed-off-by: Masaki Kozuki * again, set sync_batch_comm to False Signed-off-by: Masaki Kozuki Co-authored-by: Aidyn-A * Remove `torch.testing._internal.common_cuda` Signed-off-by: Masaki Kozuki Signed-off-by: Masaki Kozuki Co-authored-by: Sangkug Lym Co-authored-by: Aidyn-A --- .../pipeline_parallel/p2p_communication.py | 26 ++++- .../fwd_bwd_pipelining_with_interleaving.py | 9 ++ ...fwd_bwd_pipelining_without_interleaving.py | 31 +++++- .../test_pipeline_parallel_fwd_bwd.py | 101 +++++++++++++----- 4 files changed, 138 insertions(+), 29 deletions(-) diff --git a/apex/transformer/pipeline_parallel/p2p_communication.py b/apex/transformer/pipeline_parallel/p2p_communication.py index 6c4b0d93d..aa8470c5f 100644 --- a/apex/transformer/pipeline_parallel/p2p_communication.py +++ b/apex/transformer/pipeline_parallel/p2p_communication.py @@ -128,6 +128,7 @@ def _communicate( fp32_residual_connection: bool = False, async_comm: bool = False, sequence_parallel_enabled: bool = False, + sync_batch_comm: bool = True, ) -> Tuple[Union[torch.Tensor, FutureTensor, None], Union[torch.Tensor, FutureTensor, None]]: """Base function for communication of tensors between stages. @@ -159,6 +160,8 @@ def _communicate( sequence_parallel_enabled: Set to :obj:`True` if sequence parallel is enabled. This argument is here for consistency with Megatron-LM. This argument has an effect on the communication optimization, not on tensor_shape update. + sync_batch_comm: If :obj:`False`, disable cuda synchronization after the batched communication. + To disable, https://github.com/pytorch/pytorch/pull/82450 would be required. Returns: tuple containing @@ -267,8 +270,9 @@ def tensor_recv_next_wait(): torch.cuda.synchronize() tensor_recv_next_waitfunc = tensor_recv_next_wait else: - # To protect against race condition when using batch_isend_irecv(). - torch.cuda.synchronize() + if sync_batch_comm: + # To protect against race condition when using batch_isend_irecv(). + torch.cuda.synchronize() # If using scatter-gather optimization, gather smaller chunks. if scatter_gather_optimization_doable: @@ -325,6 +329,7 @@ def recv_forward( dtype: Optional[torch.dtype] = None, async_comm: bool = False, sequence_parallel_enabled: bool = False, + sync_batch_comm: bool = True, timers: _Timers = None, ) -> Union[torch.Tensor, FutureTensor, None]: """Receive tensor from previous rank in pipeline (forward receive).""" @@ -342,6 +347,7 @@ def recv_forward( dtype_=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) # if timers is not None: # timers("forward-recv").stop() @@ -354,6 +360,7 @@ def recv_backward( dtype: Optional[torch.dtype] = None, async_comm: bool = False, sequence_parallel_enabled: bool = False, + sync_batch_comm: bool = True, timers: _Timers = None, ) -> Union[torch.Tensor, FutureTensor, None]: """Receive tensor from next rank in pipeline (backward receive).""" @@ -370,6 +377,7 @@ def recv_backward( dtype_=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) # if timers is not None: # timers("backward-recv").stop() @@ -384,6 +392,7 @@ def send_forward( dtype: Optional[torch.dtype] = None, async_comm: bool = False, sequence_parallel_enabled: bool = False, + sync_batch_comm: bool = True, timers: _Timers = None, ) -> None: """Send tensor to next rank in pipeline (forward send).""" @@ -401,6 +410,7 @@ def send_forward( dtype_=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) # if timers is not None: # timers("forward-send").stop() @@ -413,6 +423,7 @@ def send_backward( dtype: Optional[torch.dtype] = None, async_comm: bool = False, sequence_parallel_enabled: bool = False, + sync_batch_comm: bool = True, timers: _Timers = None, ) -> None: """Send tensor to previous rank in pipeline (backward send).""" @@ -429,6 +440,7 @@ def send_backward( dtype_=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) # if timers is not None: # timers("backward-send").stop() @@ -441,6 +453,7 @@ def send_forward_recv_backward( dtype: Optional[torch.dtype] = None, async_comm: bool = False, sequence_parallel_enabled: bool = False, + sync_batch_comm: bool = True, timers: _Timers = None, ) -> Union[torch.Tensor, FutureTensor, None]: """Batched send and recv with next rank in pipeline.""" @@ -457,6 +470,7 @@ def send_forward_recv_backward( dtype_=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) # if timers is not None: # timers("forward-send-backward-recv").stop() @@ -470,6 +484,7 @@ def send_backward_recv_forward( dtype: Optional[torch.dtype] = None, async_comm: bool = False, sequence_parallel_enabled: bool = False, + sync_batch_comm: bool = True, timers: _Timers = None, ) -> Union[torch.Tensor, FutureTensor, None]: """Batched send and recv with previous rank in pipeline.""" @@ -486,6 +501,7 @@ def send_backward_recv_forward( dtype_=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) # if timers is not None: # timers("backward-send-forward-recv").stop() @@ -500,6 +516,7 @@ def send_forward_recv_forward( dtype: Optional[torch.dtype] = None, async_comm: bool = False, sequence_parallel_enabled: bool = False, + sync_batch_comm: bool = True, timers: _Timers = None, ) -> Union[torch.Tensor, FutureTensor]: """Batched recv from previous rank and send to next rank in pipeline.""" @@ -514,6 +531,7 @@ def send_forward_recv_forward( dtype_=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) # if timers is not None: # timers("forward-send-forward-recv").stop() @@ -528,6 +546,7 @@ def send_backward_recv_backward( dtype: Optional[torch.dtype] = None, async_comm: bool = False, sequence_parallel_enabled: bool = False, + sync_batch_comm: bool = True, timers: _Timers = None, ) -> Union[torch.Tensor, FutureTensor]: """Batched recv from next rank and send to previous rank in pipeline.""" @@ -542,6 +561,7 @@ def send_backward_recv_backward( dtype_=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) # if timers is not None: # timers("backward-send-backward-recv").stop() @@ -558,6 +578,7 @@ def send_forward_backward_recv_forward_backward( dtype: Optional[torch.dtype] = None, async_comm: bool = False, sequence_parallel_enabled: bool = False, + sync_batch_comm: bool = True, timers: _Timers = None, ) -> Tuple[Union[torch.Tensor, FutureTensor], Union[torch.Tensor, FutureTensor]]: """Batched send and recv with previous and next ranks in pipeline.""" @@ -572,6 +593,7 @@ def send_forward_backward_recv_forward_backward( dtype_=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) # if timers is not None: # timers("forward-backward-send-forward-backward-recv").stop() diff --git a/apex/transformer/pipeline_parallel/schedules/fwd_bwd_pipelining_with_interleaving.py b/apex/transformer/pipeline_parallel/schedules/fwd_bwd_pipelining_with_interleaving.py index 17ad8334d..0d94c0b84 100644 --- a/apex/transformer/pipeline_parallel/schedules/fwd_bwd_pipelining_with_interleaving.py +++ b/apex/transformer/pipeline_parallel/schedules/fwd_bwd_pipelining_with_interleaving.py @@ -36,6 +36,7 @@ def _forward_backward_pipelining_with_interleaving( deallocate_pipeline_outputs: bool = False, async_comm: bool = False, sequence_parallel_enabled: bool = False, + sync_batch_comm: bool = True, **kwargs, ) -> List[Union[torch.Tensor, Sequence[torch.Tensor]]]: """Run interleaved 1F1B schedule with communication between pipeline stages as needed. @@ -70,6 +71,8 @@ def _forward_backward_pipelining_with_interleaving( sequence_parallel_enabled: Set to :obj:`True` for this function to handle sequence length. When :obj:`True`, the sequence length on each tensor model parallel rank is updated to :math:`original\_sequence\_length / tensor\_model\_parallel\_world\_size`. + sync_batch_comm: If :obj:`False`, disable cuda synchronization after the batched communication. + To disable, https://github.com/pytorch/pytorch/pull/82450 would be required. Returns: a list of loss `torch.Tensor`s if the last stage, empty list otherwise. @@ -221,6 +224,7 @@ def backward_step_helper(microbatch_id: int) -> torch.Tensor: dtype=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) ) _logger.info("Warmup phase") @@ -269,6 +273,7 @@ def backward_step_helper(microbatch_id: int) -> torch.Tensor: dtype=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) output_tensor_grads[num_model_chunks - 1].append(output_tensor_grad) else: @@ -280,6 +285,7 @@ def backward_step_helper(microbatch_id: int) -> torch.Tensor: dtype=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) input_tensors[next_forward_model_chunk_id].append(input_tensor) free_output_tensor(output_tensor, deallocate_pipeline_outputs) @@ -365,6 +371,7 @@ def backward_step_helper(microbatch_id: int) -> torch.Tensor: dtype=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) free_output_tensor(output_tensor, deallocate_pipeline_outputs) @@ -387,6 +394,7 @@ def backward_step_helper(microbatch_id: int) -> torch.Tensor: dtype=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) ) for k in range(num_microbatches_remaining, num_microbatches): @@ -409,6 +417,7 @@ def backward_step_helper(microbatch_id: int) -> torch.Tensor: dtype=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) ) diff --git a/apex/transformer/pipeline_parallel/schedules/fwd_bwd_pipelining_without_interleaving.py b/apex/transformer/pipeline_parallel/schedules/fwd_bwd_pipelining_without_interleaving.py index 30d903b5a..18dc4c8e9 100644 --- a/apex/transformer/pipeline_parallel/schedules/fwd_bwd_pipelining_without_interleaving.py +++ b/apex/transformer/pipeline_parallel/schedules/fwd_bwd_pipelining_without_interleaving.py @@ -90,6 +90,7 @@ def recv_forward( dtype: Optional[torch.dtype] = None, async_comm: bool = False, sequence_parallel_enabled: bool = False, + sync_batch_comm: bool = True, ) -> List[Union[None, torch.Tensor, FutureTensor]]: input_tensors = [] for tensor_shape in tensor_shapes: @@ -102,6 +103,7 @@ def recv_forward( dtype=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) ) return input_tensors @@ -113,6 +115,7 @@ def recv_backward( dtype: Optional[torch.dtype] = None, async_comm: bool = False, sequence_parallel_enabled: bool = False, + sync_batch_comm: bool = True, ) -> List[Union[None, torch.Tensor, FutureTensor]]: output_tensor_grads = [] for tensor_shape in tensor_shapes: @@ -125,6 +128,7 @@ def recv_backward( dtype=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) ) return output_tensor_grads @@ -137,6 +141,7 @@ def send_forward( dtype: Optional[torch.dtype] = None, async_comm: bool = False, sequence_parallel_enabled: bool = False, + sync_batch_comm: bool = True, ) -> None: if not isinstance(output_tensors, list): output_tensors = [output_tensors] @@ -149,6 +154,7 @@ def send_forward( dtype=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) @@ -159,6 +165,7 @@ def send_backward( dtype: Optional[torch.dtype] = None, async_comm: bool = False, sequence_parallel_enabled: bool = False, + sync_batch_comm: bool = True, ) -> None: if not isinstance(input_tensor_grads, list): input_tensor_grads = [input_tensor_grads] @@ -171,6 +178,7 @@ def send_backward( dtype=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) @@ -181,6 +189,7 @@ def send_forward_recv_backward( dtype: Optional[torch.dtype] = None, async_comm: bool = False, sequence_parallel_enabled: bool = False, + sync_batch_comm: bool = True, ) -> List[Union[None, torch.Tensor, FutureTensor]]: if not isinstance(output_tensors, list): output_tensors = [output_tensors] @@ -195,6 +204,7 @@ def send_forward_recv_backward( dtype=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) output_tensor_grads.append(output_tensor_grad) return output_tensor_grads @@ -207,6 +217,7 @@ def send_backward_recv_forward( dtype: Optional[torch.dtype] = None, async_comm: bool = False, sequence_parallel_enabled: bool = False, + sync_batch_comm: bool = True, ) -> List[Union[None, torch.Tensor, FutureTensor]]: if not isinstance(input_tensor_grads, list): input_tensor_grads = [input_tensor_grads] @@ -221,6 +232,7 @@ def send_backward_recv_forward( dtype=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) input_tensors.append(input_tensor) return input_tensors @@ -241,6 +253,7 @@ def forward_backward_pipelining_without_interleaving( async_comm: bool = False, sequence_parallel_enabled: bool = False, custom_sync_context_handler = None, + sync_batch_comm: bool = True, **kwargs, ) -> List[Union[torch.Tensor, Sequence[torch.Tensor]]]: """Run non-interleaved 1F1B schedule, with communication between pipeline stages. @@ -279,6 +292,8 @@ def forward_backward_pipelining_without_interleaving( gradients must be manually synchronized by the caller (in practice the runtime is covered up by the bubble overhead). + sync_batch_comm: If :obj:`False`, disable cuda synchronization after the batched communication. + To disable, https://github.com/pytorch/pytorch/pull/82450 would be required. Returns: a list of loss `torch.Tensor`s if the last stage, empty list otherwise. @@ -354,6 +369,7 @@ def forward_backward_pipelining_without_interleaving( dtype=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) cur_microbatch: Optional[torch.Tensor] = get_kth_microbatch(batch, i) output_tensor = forward_step( @@ -372,6 +388,7 @@ def forward_backward_pipelining_without_interleaving( dtype=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) if not forward_only: @@ -384,7 +401,12 @@ def forward_backward_pipelining_without_interleaving( # receive this tensor here. if num_microbatches_remaining > 0: _logger.debug("recv_forward before steady state start") - input_tensor: List[Union[None, torch.Tensor, FutureTensor]] = recv_forward(tensor_shapes=recv_tensor_shapes, dtype=dtype, async_comm=async_comm) + input_tensor: List[Union[None, torch.Tensor, FutureTensor]] = recv_forward( + tensor_shapes=recv_tensor_shapes, + dtype=dtype, + async_comm=async_comm, + sync_batch_comm=sync_batch_comm, + ) ################################################################################################################### # Run 1F1B in steady state. @@ -412,6 +434,7 @@ def forward_backward_pipelining_without_interleaving( dtype=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) if not last_iteration: @@ -421,6 +444,7 @@ def forward_backward_pipelining_without_interleaving( dtype=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) else: @@ -431,6 +455,7 @@ def forward_backward_pipelining_without_interleaving( dtype=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) # Add input_tensor and output_tensor to end of list. @@ -460,6 +485,7 @@ def forward_backward_pipelining_without_interleaving( dtype=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) else: _logger.debug("send bwd and receive fwd") @@ -469,6 +495,7 @@ def forward_backward_pipelining_without_interleaving( dtype=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) ################################################################################################################### # Run cooldown backward passes. @@ -489,6 +516,7 @@ def forward_backward_pipelining_without_interleaving( dtype=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) input_tensor_grad = backward_step( @@ -507,6 +535,7 @@ def forward_backward_pipelining_without_interleaving( dtype=dtype, async_comm=async_comm, sequence_parallel_enabled=sequence_parallel_enabled, + sync_batch_comm=sync_batch_comm, ) # Make sure to exit context handler for async grad reductions diff --git a/tests/L0/run_transformer/test_pipeline_parallel_fwd_bwd.py b/tests/L0/run_transformer/test_pipeline_parallel_fwd_bwd.py index 11b59d1d3..788c7d946 100644 --- a/tests/L0/run_transformer/test_pipeline_parallel_fwd_bwd.py +++ b/tests/L0/run_transformer/test_pipeline_parallel_fwd_bwd.py @@ -1,13 +1,14 @@ import contextlib import logging import itertools +import os +from packaging.version import parse, Version import re from typing import Optional, Tuple, List import unittest import torch from torch.testing._internal import common_utils -from torch.testing._internal import common_cuda from apex._autocast_utils import _get_autocast_dtypes from apex.transformer import parallel_state @@ -39,6 +40,16 @@ weight_coeff = 1024 +# Guard for https://github.com/pytorch/pytorch/pull/82450 +CAN_SKIP_SYNC_AFTER_BATCH_ISEND_IRECV = False +ngc_container_2209, pytorch_113 = Version("22.09"), Version("1.13") +if parse(os.getenv("NVIDIA_PYTORCH_VERSION", "22.08")) >= ngc_container_2209: + CAN_SKIP_SYNC_AFTER_BATCH_ISEND_IRECV = True +elif parse(torch.__version__) >= pytorch_113: + CAN_SKIP_SYNC_AFTER_BATCH_ISEND_IRECV = True +else: + CAN_SKIP_SYNC_AFTER_BATCH_ISEND_IRECV = False + def get_init_weights_func(offset: int = 0): @torch.no_grad() @@ -114,6 +125,7 @@ def _forward_backward_test_impl( *, default_backend: Optional[str] = None, p2p_backend: Optional[str] = None, + sync_batch_comm: bool = True, ) -> None: if fwd_bwd_func == _forward_backward_pipelining_with_interleaving: self.assertIsNotNone(virtual_pipeline_model_parallel_size) @@ -167,7 +179,6 @@ def _forward_backward_test_impl( hidden_size=self.HIDDEN_SIZE, ) - offset = pipeline_model_parallel_world_size if virtual_pipeline_model_parallel_size is not None else 0 for idx, model_module in enumerate(model): model_module = model_module.to(dtype) @@ -193,6 +204,7 @@ def _forward_backward_test_impl( async_comm=async_comm, grad_scaler=grad_scaler, deallocate_pipeline_output=deallocate_pipeline_outputs, + sync_batch_comm=sync_batch_comm, ) if dtype == get_dtype_for_comparison(): @@ -236,48 +248,54 @@ def test_learning_no_pipelining(self): def test_inference_no_pipelining(self): self._forward_backward_test_impl(True, forward_backward_no_pipelining, 1, None) - def test_learning_pipelining_without_interleaving(self): + def test_learning_pipelining_without_interleaving(self, sync_batch_comm: bool = True): self._forward_backward_test_impl( - False, forward_backward_pipelining_without_interleaving, None, None + False, forward_backward_pipelining_without_interleaving, None, None, sync_batch_comm=sync_batch_comm, ) - def test_inference_pipelining_without_interleaving(self): + def test_inference_pipelining_without_interleaving(self, sync_batch_comm: bool = True): self._forward_backward_test_impl( - True, forward_backward_pipelining_without_interleaving, None, None + True, forward_backward_pipelining_without_interleaving, None, None, sync_batch_comm=sync_batch_comm, ) - def test_learning_async_pipelining_without_interleaving(self): + def test_learning_async_pipelining_without_interleaving(self, sync_batch_comm: bool = True): self._forward_backward_test_impl( - False, forward_backward_pipelining_without_interleaving, None, None, async_comm=True + False, forward_backward_pipelining_without_interleaving, None, None, async_comm=True, + sync_batch_comm=sync_batch_comm, ) - def test_inference_async_pipelining_without_interleaving(self): + def test_inference_async_pipelining_without_interleaving(self, sync_batch_comm: bool = True): self._forward_backward_test_impl( - True, forward_backward_pipelining_without_interleaving, None, None, async_comm=True + True, forward_backward_pipelining_without_interleaving, None, None, async_comm=True, + sync_batch_comm=sync_batch_comm, ) @unittest.skipUnless(_get_default_world_sizes_model_parallel_world_size()[-1] > 2, "Interleaved schedule requires pipeline_model_parallel_world_size > 2") - def test_learning_pipelining_with_interleaving(self): + def test_learning_pipelining_with_interleaving(self, sync_batch_comm: bool = True): self._forward_backward_test_impl( - False, _forward_backward_pipelining_with_interleaving, None, virtual_pipeline_model_parallel_size=2 + False, _forward_backward_pipelining_with_interleaving, None, virtual_pipeline_model_parallel_size=2, + sync_batch_comm=sync_batch_comm, ) @unittest.skipUnless(_get_default_world_sizes_model_parallel_world_size()[-1] > 2, "Interleaved schedule requires pipeline_model_parallel_world_size > 2") - def test_inference_pipelining_with_interleaving(self): + def test_inference_pipelining_with_interleaving(self, sync_batch_comm: bool = True): self._forward_backward_test_impl( - True, _forward_backward_pipelining_with_interleaving, None, virtual_pipeline_model_parallel_size=2 + True, _forward_backward_pipelining_with_interleaving, None, virtual_pipeline_model_parallel_size=2, + sync_batch_comm=sync_batch_comm, ) @unittest.skipUnless(_get_default_world_sizes_model_parallel_world_size()[-1] > 2, "Interleaved schedule requires pipeline_model_parallel_world_size > 2") - def test_learning_async_pipelining_with_interleaving(self): + def test_learning_async_pipelining_with_interleaving(self, sync_batch_comm: bool = True): self._forward_backward_test_impl( - False, _forward_backward_pipelining_with_interleaving, None, virtual_pipeline_model_parallel_size=2, async_comm=True + False, _forward_backward_pipelining_with_interleaving, None, virtual_pipeline_model_parallel_size=2, async_comm=True, + sync_batch_comm=sync_batch_comm, ) @unittest.skipUnless(_get_default_world_sizes_model_parallel_world_size()[-1] > 2, "Interleaved schedule requires pipeline_model_parallel_world_size > 2") - def test_inference_async_pipelining_with_interleaving(self): + def test_inference_async_pipelining_with_interleaving(self, sync_batch_comm: bool = True): self._forward_backward_test_impl( - True, _forward_backward_pipelining_with_interleaving, None, virtual_pipeline_model_parallel_size=2, async_comm=True + True, _forward_backward_pipelining_with_interleaving, None, virtual_pipeline_model_parallel_size=2, async_comm=True, + sync_batch_comm=sync_batch_comm, ) @@ -310,6 +328,38 @@ def test_learning_pipelining_without_interleaving_ucc_for_p2p(self): def test_inference_pipelining_without_interleaving_ucc_for_p2p(self): self._test_hybrid_backends(True) + @unittest.skipUnless(CAN_SKIP_SYNC_AFTER_BATCH_ISEND_IRECV, "Requires https://github.com/pytorch/pytorch/pull/82450") + def test_learning_pipelining_without_interleaving_skyp_sync_after_batch_isend_irecv(self): + self.test_learning_pipelining_without_interleaving(sync_batch_comm=False) + + @unittest.skipUnless(CAN_SKIP_SYNC_AFTER_BATCH_ISEND_IRECV, "Requires https://github.com/pytorch/pytorch/pull/82450") + def test_inference_pipelining_without_interleaving_skip_sync_after_batch_isend_irecv(self): + self.test_inference_pipelining_without_interleaving(sync_batch_comm=False) + + @unittest.skipUnless(CAN_SKIP_SYNC_AFTER_BATCH_ISEND_IRECV, "Requires https://github.com/pytorch/pytorch/pull/82450") + def test_learning_async_pipelining_without_interleaving_skip_sync_after_batch_isend_irecv(self): + self.test_learning_async_pipelining_without_interleaving(sync_batch_comm=False) + + @unittest.skipUnless(CAN_SKIP_SYNC_AFTER_BATCH_ISEND_IRECV, "Requires https://github.com/pytorch/pytorch/pull/82450") + def test_inference_async_pipelining_without_interleaving_skip_sync_after_batch_isend_irecv(self): + self.test_inference_async_pipelining_without_interleaving(sync_batch_comm=False) + + @unittest.skipUnless(CAN_SKIP_SYNC_AFTER_BATCH_ISEND_IRECV, "Requires https://github.com/pytorch/pytorch/pull/82450") + def test_learning_pipelining_with_interleaving_skip_sync_after_batch_isend_irecv(self): + self.test_learning_pipelining_with_interleaving(sync_batch_comm=False) + + @unittest.skipUnless(CAN_SKIP_SYNC_AFTER_BATCH_ISEND_IRECV, "Requires https://github.com/pytorch/pytorch/pull/82450") + def test_inference_pipelining_with_interleaving_skip_sync_after_batch_isend_irecv(self): + self.test_inference_pipelining_with_interleaving(sync_batch_comm=False) + + @unittest.skipUnless(CAN_SKIP_SYNC_AFTER_BATCH_ISEND_IRECV, "Requires https://github.com/pytorch/pytorch/pull/82450") + def test_learning_async_pipelining_with_interleaving_skip_sync_after_batch_isend_irecv(self): + self.test_learning_async_pipelining_with_interleaving(sync_batch_comm=False) + + @unittest.skipUnless(CAN_SKIP_SYNC_AFTER_BATCH_ISEND_IRECV, "Requires https://github.com/pytorch/pytorch/pull/82450") + def test_inference_async_pipelining_with_interleaving_skip_sync_after_batch_isend_irecv(self): + self.test_inference_async_pipelining_with_interleaving(sync_batch_comm=False) + # n.b.(mkozuki): pipeline parallel w/o interleaving with UCX_TLS=tcp,sm fails. class UccPipelineParallelForwardBackwardTest(UccDistributedTestBase, PipelineParallelForwardBackwardTestBase): @@ -328,25 +378,22 @@ def world_size(self) -> int: @unittest.skipIf(torch.cuda.device_count() < 4, "Requires >= 4 GPUs") class NcclPipelineParallelWithToyParallelMLP(NcclDistributedTestBase): - GLOBAL_BATCH_SIZE = 16 - MICRO_BATCH_SIZE = 2 - HIDDEN_SIZE = 64 + GLOBAL_BATCH_SIZE: int = 16 + MICRO_BATCH_SIZE: int = 2 + HIDDEN_SIZE: int = 64 # TODO(mkozuki): Change `DECODER_SEQUENCE_LENGTH` to a value different from `ENCODER_SEQUENCE_LENGTH`. # To test forward_backward_pipelining_without_interleaving with `model_type=ModelType.encoder_and_decoder`, # `decoder_seq_length` is necessary and ideally should be different from `encoder_sequence_length` # but my laziness let me use the same value. # Note that you may have to either update `MyModel` def or define another `MyModel`. # to support different `DECODER_SEQUENCE_LENGTH`. - ENCODER_SEQUENCE_LENGTH = 32 - DECODER_SEQUENCE_LENGTH = 32 + ENCODER_SEQUENCE_LENGTH: int = 32 + DECODER_SEQUENCE_LENGTH: int = 32 @property def world_size(self) -> int: return min(torch.cuda.device_count(), 8) - # TODO(mkozuki): Add cases of async_comm=True - # TODO(mkozuki): Add loss check. - # TODO(mkozuki): Call `build_model` with `model_type`. # TODO(mkozuki): Set `tensor_model_parallel>1` for encoder_and_decoder as well if there's enough GPUs # in order to let `sequence_parallel_enabled` have an effect on tensor shape logic. def _forward_backward_test_impl( @@ -381,6 +428,7 @@ def _forward_backward_test_impl( micro_batch_size=self.MICRO_BATCH_SIZE, data_parallel_size=parallel_state.get_data_parallel_world_size(), ) + # TODO(mkozuki): Call `build_model` with `model_type`. model = build_model( testing_utils.mlp_provider_func, wrap_with_ddp=False, @@ -443,6 +491,7 @@ def test_pipelining_without_interleaving_sequence_parallel_encoder_or_decoder(se def test_pipelining_without_interleaving_sequence_parallel_encoder_or_decoder_half(self) -> None: self._forward_backward_test_impl(forward_only=False, sequence_parallel_enabled=True, model_type=ModelType.encoder_or_decoder, dtype=torch.half) + @unittest.skipIf(torch.cuda.device_count() < 4 or torch.cuda.device_count() % 2 != 0, "Requires >= 4 GPUs") class NcclPipelineParallelWithCustomSyncContextHandler(NcclDistributedTestBase): From 127842e122e2153aa09bc5fa8683408a66e7a4de Mon Sep 17 00:00:00 2001 From: hubertlu-tw Date: Thu, 29 Dec 2022 19:18:08 +0000 Subject: [PATCH 12/25] Skip test_grad_scaler in test_dist_adam.py --- apex/contrib/test/optimizers/test_dist_adam.py | 2 ++ apex/contrib/test/run_rocm_extensions.py | 2 +- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/apex/contrib/test/optimizers/test_dist_adam.py b/apex/contrib/test/optimizers/test_dist_adam.py index bd23ce2ae..6040e8eb4 100644 --- a/apex/contrib/test/optimizers/test_dist_adam.py +++ b/apex/contrib/test/optimizers/test_dist_adam.py @@ -6,6 +6,7 @@ from torch.testing._internal import common_utils from apex.contrib.optimizers.distributed_fused_adam import DistributedFusedAdam from apex.transformer.testing.distributed_test_base import NcclDistributedTestBase +from apex.testing.common_utils import skipIfRocm class SimpleModel(torch.nn.Module): @@ -250,6 +251,7 @@ def test_clip_grad_norm(self): dist_model.parameters()): torch.testing.assert_close(dist_param, ref_param) + @skipIfRocm def test_grad_scaler(self): torch.manual_seed(self.seed + self.rank) diff --git a/apex/contrib/test/run_rocm_extensions.py b/apex/contrib/test/run_rocm_extensions.py index c7801988b..0e8b2a69b 100644 --- a/apex/contrib/test/run_rocm_extensions.py +++ b/apex/contrib/test/run_rocm_extensions.py @@ -2,7 +2,7 @@ import sys -test_dirs = ["groupbn", "fused_dense", "layer_norm", "multihead_attn", "transducer", "focal_loss", "index_mul_2d", "."] # "." for test_label_smoothing.py +test_dirs = ["groupbn", "fused_dense", "layer_norm", "multihead_attn", "transducer", "focal_loss", "index_mul_2d", "optimizers", "."] # "." for test_label_smoothing.py ROCM_BLACKLIST = [ "layer_norm" ] From 589a5a6c05cb4d44af03c5756f43a2d1a232f452 Mon Sep 17 00:00:00 2001 From: hubertlu-tw Date: Thu, 29 Dec 2022 19:22:59 +0000 Subject: [PATCH 13/25] Add run_rocm_extensions.sh for skipFlakyTest and skipIfRocm --- apex/contrib/test/run_rocm_extensions.sh | 2 ++ 1 file changed, 2 insertions(+) create mode 100644 apex/contrib/test/run_rocm_extensions.sh diff --git a/apex/contrib/test/run_rocm_extensions.sh b/apex/contrib/test/run_rocm_extensions.sh new file mode 100644 index 000000000..302b42189 --- /dev/null +++ b/apex/contrib/test/run_rocm_extensions.sh @@ -0,0 +1,2 @@ +#!/bin/bash +APEX_TEST_WITH_ROCM=1 APEX_SKIP_FLAKY_TEST=1 python3 run_rocm_extensions.py From 4ca3c93ba9984d5a3683b2569585400504a73453 Mon Sep 17 00:00:00 2001 From: Aidyn-A <31858918+Aidyn-A@users.noreply.github.com> Date: Wed, 12 Oct 2022 17:41:22 -0700 Subject: [PATCH 14/25] add tearDown (#1508) --- apex/transformer/testing/distributed_test_base.py | 1 + tests/L0/run_transformer/test_batch_sampler.py | 4 ++++ tests/L0/run_transformer/test_fused_softmax.py | 8 ++++++++ 3 files changed, 13 insertions(+) diff --git a/apex/transformer/testing/distributed_test_base.py b/apex/transformer/testing/distributed_test_base.py index 7a8168759..e0823e80c 100644 --- a/apex/transformer/testing/distributed_test_base.py +++ b/apex/transformer/testing/distributed_test_base.py @@ -38,6 +38,7 @@ def setUp(self) -> None: self._spawn_processes() def tearDown(self) -> None: + torch.cuda.empty_cache() super().tearDown() @property diff --git a/tests/L0/run_transformer/test_batch_sampler.py b/tests/L0/run_transformer/test_batch_sampler.py index f0fb01407..99b0f1aeb 100644 --- a/tests/L0/run_transformer/test_batch_sampler.py +++ b/tests/L0/run_transformer/test_batch_sampler.py @@ -77,6 +77,10 @@ def __iter__(self): # Samples 8 tensors in total. # First sample 4 tensors twice, then sample 2 tensors fourth. class TestBatchSamplerBehavior(common_utils.TestCase): + def tearDown(self) -> None: + torch.cuda.empty_cache() + super().tearDown() + def test_batch_sampler_behavior(self): dataset = MyIterableDataset(0, 100) diff --git a/tests/L0/run_transformer/test_fused_softmax.py b/tests/L0/run_transformer/test_fused_softmax.py index 015f997db..7eb452564 100644 --- a/tests/L0/run_transformer/test_fused_softmax.py +++ b/tests/L0/run_transformer/test_fused_softmax.py @@ -57,6 +57,10 @@ def _setup_fused_softmax( ) return fused_fn, torch_fn + def tearDown(self) -> None: + torch.cuda.empty_cache() + super().tearDown() + def test_fused_scale_mask_softmax(self): """ attention_scores.shape = [4, 12, 24, 24] @@ -245,6 +249,10 @@ def setUp(self): self.q_k_lens = itertools.product(qlen, klen) + def tearDown(self) -> None: + torch.cuda.empty_cache() + super().tearDown() + def test_forward(self, allmasked: bool=False): import generic_scaled_masked_softmax_cuda for qlen, klen in self.q_k_lens: From 0b55a2fbc0f630575a115958ddbcd85740d6b134 Mon Sep 17 00:00:00 2001 From: Tim Moon <4406448+timmoon10@users.noreply.github.com> Date: Wed, 26 Oct 2022 18:07:02 -0700 Subject: [PATCH 15/25] Support overlapped grad sync with Megatron interleaved pipeline parallelism (#1514) * Add option to use no_sync context with interleaved pipeline parallelism * Add unit test for no_sync context with interleaved pipeline parallelism * Debug no_sync context support in interleaved pipeline parallelism --- .../fwd_bwd_pipelining_with_interleaving.py | 73 +++++++- ...fwd_bwd_pipelining_without_interleaving.py | 61 ++++--- .../test_pipeline_parallel_fwd_bwd.py | 160 +++++++++++++++--- 3 files changed, 240 insertions(+), 54 deletions(-) diff --git a/apex/transformer/pipeline_parallel/schedules/fwd_bwd_pipelining_with_interleaving.py b/apex/transformer/pipeline_parallel/schedules/fwd_bwd_pipelining_with_interleaving.py index 0d94c0b84..748b23b4a 100644 --- a/apex/transformer/pipeline_parallel/schedules/fwd_bwd_pipelining_with_interleaving.py +++ b/apex/transformer/pipeline_parallel/schedules/fwd_bwd_pipelining_with_interleaving.py @@ -1,4 +1,5 @@ -from typing import List, Union, Optional, Sequence +import contextlib +from typing import Any, List, Optional, Sequence, Union import warnings import torch @@ -36,6 +37,7 @@ def _forward_backward_pipelining_with_interleaving( deallocate_pipeline_outputs: bool = False, async_comm: bool = False, sequence_parallel_enabled: bool = False, + custom_sync_context_handler: Optional[Any] = None, sync_batch_comm: bool = True, **kwargs, ) -> List[Union[torch.Tensor, Sequence[torch.Tensor]]]: @@ -71,11 +73,18 @@ def _forward_backward_pipelining_with_interleaving( sequence_parallel_enabled: Set to :obj:`True` for this function to handle sequence length. When :obj:`True`, the sequence length on each tensor model parallel rank is updated to :math:`original\_sequence\_length / tensor\_model\_parallel\_world\_size`. + custom_sync_context_handler: Does nothing if ``None`` (default + value). Otherwise, this is treated as a function to + construct a context manager to disable asynchronous + gradient reductions. Asynchronous gradient reductions are + only enabled in the final backward pass of each model + chunk. sync_batch_comm: If :obj:`False`, disable cuda synchronization after the batched communication. To disable, https://github.com/pytorch/pytorch/pull/82450 would be required. Returns: a list of loss `torch.Tensor`s if the last stage, empty list otherwise. + """ if not isinstance(model, list): raise RuntimeError("`model` must be a list of `nn.Module`'s'") @@ -86,6 +95,26 @@ def _forward_backward_pipelining_with_interleaving( "This option is not recommended." ) + # Construct helper functions for async grad reductions + if custom_sync_context_handler is not None: + sync_context_handler = custom_sync_context_handler + else: + sync_context_handler = contextlib.nullcontext + sync_context = None + def disable_grad_sync(): + """Disable asynchronous grad reductions""" + nonlocal sync_context + if sync_context is None: + sync_context = sync_context_handler() + sync_context.__enter__() + def enable_grad_sync(): + """Enable asynchronous grad reductions""" + nonlocal sync_context + if sync_context is not None: + sync_context.__exit__(None, None, None) + sync_context = None + disable_grad_sync() + # mypy will blame the following if statement if sequence_parallel_enabled: seq_length, batch_size, hidden = tensor_shape @@ -145,17 +174,41 @@ def _forward_backward_pipelining_with_interleaving( # Helper function definitions. ################################################################################################################### def get_model_chunk_id(microbatch_id: int, forward: bool) -> int: - """Helper function to get the model chunk ID given the iteration number.""" - pipeline_parallel_size = parallel_state.get_pipeline_model_parallel_world_size() - microbatch_id_in_group = microbatch_id % ( - pipeline_parallel_size * num_model_chunks - ) + """Helper function to get the model chunk ID given the iteration number. + + Each model chunk processes pipeline_parallel_size microbatches + at a time. We assume that the number of microbatches is a + multiple of pipeline_parallel_size*num_model_chunks. + """ + microbatch_group_size = pipeline_parallel_size * num_model_chunks + microbatch_id_in_group = microbatch_id % microbatch_group_size model_chunk_id = microbatch_id_in_group // pipeline_parallel_size if not forward: model_chunk_id = num_model_chunks - model_chunk_id - 1 return model_chunk_id - def forward_step_helper(microbatch_id: int, curr_iters: List[int]) -> torch.Tensor: + def is_last_microbatch_for_model_chunk(microbatch_id: int) -> bool: + """Helper function to check if an iteration is the last for a model + chunk. + + Each model chunk processes pipeline_parallel_size microbatches + at a time. We assume that the number of microbatches is a + multiple of pipeline_parallel_size*num_model_chunks. + """ + microbatch_group_size = pipeline_parallel_size * num_model_chunks + num_microbatch_groups = num_microbatches // microbatch_group_size + microbatch_group_id = microbatch_id // microbatch_group_size + microbatch_id_in_group = microbatch_id % microbatch_group_size + if microbatch_group_id == num_microbatch_groups - 1: + return microbatch_id_in_group % pipeline_parallel_size == pipeline_parallel_size - 1 + else: + return False + + def forward_step_helper( + microbatch_id: int, + curr_iters: List[int], + checkpoint_activations_micro_batch: Optional[bool] = None, + ) -> torch.Tensor: """Helper method to run forward step with model split into chunks (run set_virtual_pipeline_model_parallel_rank() before calling forward_step()). @@ -197,6 +250,8 @@ def backward_step_helper(microbatch_id: int) -> torch.Tensor: model_type = get_model_type(model[model_chunk_id]) parallel_state.set_virtual_pipeline_model_parallel_rank(model_chunk_id) + if is_last_microbatch_for_model_chunk(microbatch_id): + enable_grad_sync() if parallel_state.is_pipeline_last_stage(): if len(output_tensor_grads[model_chunk_id]) == 0: output_tensor_grads[model_chunk_id].append(None) @@ -211,6 +266,7 @@ def backward_step_helper(microbatch_id: int) -> torch.Tensor: grad_scaler=grad_scaler, deallocate_pipeline_outputs=deallocate_pipeline_outputs, ) + disable_grad_sync() return input_tensor_grad @@ -421,4 +477,7 @@ def backward_step_helper(microbatch_id: int) -> torch.Tensor: ) ) + # Make sure to exit context handler for async grad reductions + enable_grad_sync() + return losses_reduced diff --git a/apex/transformer/pipeline_parallel/schedules/fwd_bwd_pipelining_without_interleaving.py b/apex/transformer/pipeline_parallel/schedules/fwd_bwd_pipelining_without_interleaving.py index 18dc4c8e9..500a1d4e1 100644 --- a/apex/transformer/pipeline_parallel/schedules/fwd_bwd_pipelining_without_interleaving.py +++ b/apex/transformer/pipeline_parallel/schedules/fwd_bwd_pipelining_without_interleaving.py @@ -1,5 +1,5 @@ import contextlib -from typing import Union, List, Optional, Sequence +from typing import Any, List, Optional, Sequence, Union import warnings import torch @@ -252,7 +252,8 @@ def forward_backward_pipelining_without_interleaving( deallocate_pipeline_outputs: bool = False, async_comm: bool = False, sequence_parallel_enabled: bool = False, - custom_sync_context_handler = None, + custom_sync_context_handler: Optional[Any] = None, + custom_grad_sync_func: Optional[Any] = None, sync_batch_comm: bool = True, **kwargs, ) -> List[Union[torch.Tensor, Sequence[torch.Tensor]]]: @@ -284,14 +285,15 @@ def forward_backward_pipelining_without_interleaving( sequence_parallel_enabled: Set to :obj:`True` for this function to handle sequence length. When :obj:`True`, the sequence length on each tensor model parallel rank is updated to :math:`original\_sequence\_length / tensor\_model\_parallel\_world\_size`. - custom_sync_context_handler: Context manager to disable - asynchronous gradient reductions. In the first pipeline - stage, asynchronous gradient reductions are enabled in the - backward pass of the last microbatch. In other pipeline - stages, asynchronous gradient reductions are disabled and - gradients must be manually synchronized by the caller (in - practice the runtime is covered up by the bubble - overhead). + custom_sync_context_handler: Does nothing if ``None`` (default + value). Otherwise, a function to construct a context + manager that disable asynchronous gradient reductions. + Asynchronous gradient reductions are only enabled in the + first pipeline stage, during the last backward pass. + custom_grad_sync_func: Does nothing if ``None`` (default + value). Otherwise, a function to perform gradient + reductions. This is called in all pipeline stages except + the first, during the bubble overhead. sync_batch_comm: If :obj:`False`, disable cuda synchronization after the batched communication. To disable, https://github.com/pytorch/pytorch/pull/82450 would be required. @@ -315,12 +317,23 @@ def forward_backward_pipelining_without_interleaving( # Disable async grad reductions if custom_sync_context_handler is not None: - context_handler = custom_sync_context_handler + sync_context_handler = custom_sync_context_handler else: - context_handler = contextlib.nullcontext - context = context_handler() - context.__enter__() - context_is_active = True + sync_context_handler = contextlib.nullcontext + sync_context = None + def disable_grad_sync(): + """Disable asynchronous grad reductions""" + nonlocal sync_context + if sync_context is None: + sync_context = sync_context_handler() + sync_context.__enter__() + def enable_grad_sync(): + """Enable asynchronous grad reductions""" + nonlocal sync_context + if sync_context is not None: + sync_context.__exit__(None, None, None) + sync_context = None + disable_grad_sync() # Compute number of warmup microbatches. num_microbatches: int = get_num_microbatches() @@ -503,10 +516,13 @@ def forward_backward_pipelining_without_interleaving( _logger.info("Cooldown phase") if not forward_only: for i in range(num_warmup_microbatches): - if i == num_warmup_microbatches-1 and rank == 0: - context.__exit__(None, None, None) - context_is_active = False _logger.debug(f"cooldown iter: {i} / {num_warmup_microbatches}") + + if i == num_warmup_microbatches-1 and rank == 0: + # Async grad reduction in first pipeline stage, during + # last backward pass + enable_grad_sync() + input_tensor = input_tensors.pop(0) output_tensor = output_tensors.pop(0) @@ -538,9 +554,10 @@ def forward_backward_pipelining_without_interleaving( sync_batch_comm=sync_batch_comm, ) - # Make sure to exit context handler for async grad reductions - if context_is_active: - context.__exit__(None, None, None) - context_is_active = False + # Grad reduction in all pipeline stages except the first, during + # the bubble overhead + enable_grad_sync() + if rank != 0 and custom_grad_sync_func is not None: + custom_grad_sync_func() return losses_reduced diff --git a/tests/L0/run_transformer/test_pipeline_parallel_fwd_bwd.py b/tests/L0/run_transformer/test_pipeline_parallel_fwd_bwd.py index 788c7d946..191f09ff0 100644 --- a/tests/L0/run_transformer/test_pipeline_parallel_fwd_bwd.py +++ b/tests/L0/run_transformer/test_pipeline_parallel_fwd_bwd.py @@ -492,10 +492,9 @@ def test_pipelining_without_interleaving_sequence_parallel_encoder_or_decoder_ha self._forward_backward_test_impl(forward_only=False, sequence_parallel_enabled=True, model_type=ModelType.encoder_or_decoder, dtype=torch.half) -@unittest.skipIf(torch.cuda.device_count() < 4 or torch.cuda.device_count() % 2 != 0, "Requires >= 4 GPUs") class NcclPipelineParallelWithCustomSyncContextHandler(NcclDistributedTestBase): - GLOBAL_BATCH_SIZE = 2 + GLOBAL_BATCH_SIZE = 32 MICRO_BATCH_SIZE = 1 HIDDEN_SIZE = 1 @@ -503,19 +502,19 @@ class NcclPipelineParallelWithCustomSyncContextHandler(NcclDistributedTestBase): def world_size(self) -> int: return min(torch.cuda.device_count(), 8) + @unittest.skipIf(torch.cuda.device_count() < 2 or torch.cuda.device_count() % 2 != 0, "Requires >= 2 GPUs") def test_pipelining_without_interleaving_with_custom_sync_context_handler(self) -> None: # Parallel configuration world_size = torch.cuda.device_count() tensor_model_parallel_world_size = 1 - data_parallel_size = 2 + data_parallel_size = 2 if world_size > 2 else 1 pipeline_model_parallel_world_size = world_size // data_parallel_size # Initialize pipeline parallelism parallel_state.initialize_model_parallel( tensor_model_parallel_size_=tensor_model_parallel_world_size, pipeline_model_parallel_size_=pipeline_model_parallel_world_size, - virtual_pipeline_model_parallel_size_=None, ) pp_utils._reconfigure_microbatch_calculator( rank=parallel_state.get_tensor_model_parallel_rank(), @@ -544,29 +543,29 @@ def test_pipelining_without_interleaving_with_custom_sync_context_handler(self) model = build_model( testing_utils.model_provider_func, wrap_with_ddp=True, - virtual_pipeline_model_parallel_size=None, hidden_size=hidden_size, )[0] model = model.to(dtype) model.module.apply(get_init_weights_func(0)) - # Construct custom context - has_entered_context = False - has_exited_context = False - no_grad_at_context_exit = False + # Construct context that destroys all grads on exit + has_entered_grad_sync_context = False + has_exited_grad_sync_context = False + has_called_grad_sync_func = False @contextlib.contextmanager - def custom_context(): + def custom_grad_sync_context(): try: - nonlocal has_entered_context - has_entered_context = True + nonlocal has_entered_grad_sync_context + has_entered_grad_sync_context = True yield finally: - nonlocal has_exited_context, no_grad_at_context_exit - has_exited_context = True - no_grad_at_context_exit = all( - p.grad is None - for p in model.parameters() - ) + nonlocal has_exited_grad_sync_context + has_exited_grad_sync_context = True + for param in model.parameters(): + param.grad = None + def custom_grad_sync_func(): + nonlocal has_called_grad_sync_func + has_called_grad_sync_func = True # Training step with pipeline parallelism loss = forward_backward_pipelining_without_interleaving( @@ -580,19 +579,130 @@ def custom_context(): grad_scaler=None, deallocate_pipeline_outputs=False, sequence_parallel_enabled=False, - custom_sync_context_handler=custom_context, + custom_sync_context_handler=custom_grad_sync_context, + custom_grad_sync_func=custom_grad_sync_func, ) torch.cuda.synchronize() + # Check if model has initialized gradients + has_any_grads = any(param.grad is not None for param in model.parameters()) + has_all_grads = all(param.grad is not None for param in model.parameters()) + # Check context behavior - self.assertTrue(has_entered_context, 'Has not entered custom sync context') - self.assertTrue(has_exited_context, 'Has not exited custom sync context') + self.assertTrue(has_entered_grad_sync_context, 'Has not entered custom sync context') + self.assertTrue(has_exited_grad_sync_context, 'Has not exited custom sync context') + self.assertEqual( + has_any_grads, + has_all_grads, + 'Expected gradients to all be uninitialized or all be initialized', + ) self.assertEqual( - no_grad_at_context_exit, + has_all_grads, parallel_state.is_pipeline_first_stage(), - 'Expected to exit custom sync context ' - 'before backward pass in first pipeline stage ' - 'and after backward pass in other pipeline stages' + 'Expected gradients to be initialized only in first pipeline stage', + ) + + # Clean up + parallel_state.destroy_model_parallel() + + @unittest.skipIf(torch.cuda.device_count() < 4 or torch.cuda.device_count() % 2 != 0, "Requires >= 4 GPUs") + def test_pipelining_with_interleaving_with_custom_sync_context_handler(self) -> None: + + # Parallel configuration + world_size = torch.cuda.device_count() + tensor_model_parallel_world_size = 1 + data_parallel_size = 2 if world_size > 4 else 1 + pipeline_model_parallel_world_size = world_size // data_parallel_size + virtual_pipeline_model_parallel_size = 2 + + # Initialize pipeline parallelism + parallel_state.initialize_model_parallel( + tensor_model_parallel_size_=tensor_model_parallel_world_size, + pipeline_model_parallel_size_=pipeline_model_parallel_world_size, + virtual_pipeline_model_parallel_size_=virtual_pipeline_model_parallel_size, + ) + pp_utils._reconfigure_microbatch_calculator( + rank=parallel_state.get_tensor_model_parallel_rank(), + rampup_batch_size=None, + global_batch_size=self.GLOBAL_BATCH_SIZE, + micro_batch_size=self.MICRO_BATCH_SIZE, + data_parallel_size=parallel_state.get_data_parallel_world_size(), + ) + pp_utils.update_num_microbatches(0) + + # Construct synthetic data + dtype = get_dtype_for_comparison() + hidden_size = self.HIDDEN_SIZE + microbatch_size = self.MICRO_BATCH_SIZE + global_batch_shape = ( + self.GLOBAL_BATCH_SIZE + // parallel_state.get_data_parallel_world_size(), + hidden_size, + hidden_size, + ) + batch = None + if parallel_state.is_pipeline_first_stage(): + batch = (torch.ones(global_batch_shape, dtype=dtype).cuda(), ) + + # Construct model + model = build_model( + testing_utils.model_provider_func, + wrap_with_ddp=True, + virtual_pipeline_model_parallel_size=virtual_pipeline_model_parallel_size, + hidden_size=hidden_size, + ) + for module in model: + module.to(dtype) + module.module.apply(get_init_weights_func(0)) + + # Construct context that keeps track whenever entered/exited + grad_sync_context_enter_count = 0 + grad_sync_context_exit_count = 0 + @contextlib.contextmanager + def custom_grad_sync_context(): + try: + nonlocal grad_sync_context_enter_count + grad_sync_context_enter_count += 1 + yield + finally: + nonlocal grad_sync_context_exit_count + grad_sync_context_exit_count += 1 + for module in model: + for param in module.parameters(): + param.grad = None + + # Training step with pipeline parallelism + loss = _forward_backward_pipelining_with_interleaving( + testing_utils.fwd_step_func, + batch, + model, + forward_only=False, + tensor_shape=(microbatch_size, hidden_size, hidden_size), + dtype=dtype, + async_comm=False, + grad_scaler=None, + deallocate_pipeline_outputs=False, + sequence_parallel_enabled=False, + custom_sync_context_handler=custom_grad_sync_context, + ) + torch.cuda.synchronize() + + # Check context behavior + self.assertTrue( + grad_sync_context_enter_count > 0, + 'Has not entered custom sync context', + ) + self.assertEqual( + grad_sync_context_enter_count, + grad_sync_context_exit_count, + 'Has not entered and exited custom sync context ' + 'the same number of times', + ) + self.assertEqual( + grad_sync_context_exit_count, + virtual_pipeline_model_parallel_size + 1, + 'Expected to exit custom sync context once per model chunk ' + 'and once at the function end', ) # Clean up From 28de67e57db711b688ab735dde29fdd321d6beea Mon Sep 17 00:00:00 2001 From: Fuzzkatt Date: Thu, 27 Oct 2022 11:39:02 -0700 Subject: [PATCH 16/25] update run_transformer tests to default to using pytorch native UCC instead of torch_ucc (#1495) * update HAS_TORCH_UCC to TORCH_UCC * add comments for failing tests * move HAS_UCC to _ucc_utils.py * whitespace * small changes * newline * updated list of failing tests * update failing tests list --- apex/transformer/_ucc_util.py | 9 +++++++++ apex/transformer/parallel_state.py | 13 +++---------- apex/transformer/testing/commons.py | 5 +++-- apex/transformer/testing/distributed_test_base.py | 13 ++++--------- tests/L0/run_transformer/run_bert_minimal_test.py | 12 +++--------- tests/L0/run_transformer/run_gpt_minimal_test.py | 10 ++-------- tests/L0/run_transformer/test_layers.py | 1 + .../test_pipeline_parallel_fwd_bwd.py | 9 ++++++--- 8 files changed, 31 insertions(+), 41 deletions(-) create mode 100644 apex/transformer/_ucc_util.py diff --git a/apex/transformer/_ucc_util.py b/apex/transformer/_ucc_util.py new file mode 100644 index 000000000..c070c84d1 --- /dev/null +++ b/apex/transformer/_ucc_util.py @@ -0,0 +1,9 @@ +from torch import distributed as dist + +HAS_UCC = dist.is_ucc_available() +if not HAS_UCC: + try: + import torch_ucc + HAS_UCC = True + except ImportError: + HAS_UCC = False diff --git a/apex/transformer/parallel_state.py b/apex/transformer/parallel_state.py index a8d16bfd3..f5821c2c2 100644 --- a/apex/transformer/parallel_state.py +++ b/apex/transformer/parallel_state.py @@ -21,6 +21,7 @@ import torch from apex.transformer.log_util import get_transformer_logger +from apex.transformer._ucc_util import HAS_UCC _logger = get_transformer_logger(__name__) @@ -126,7 +127,8 @@ def initialize_model_parallel( assert default_backend is None or default_backend in ("nccl", "ucc") assert p2p_backend is None or p2p_backend in ("nccl", "ucc") if "ucc" in (default_backend, p2p_backend): - check_torch_ucc_availability() + if not HAS_UCC: + raise ImportError("UCC backend requires pytorch source build with UCC installed and enabled") warnings.warn("`ucc` backend support is experimental", ExperimentalWarning) if default_backend == "ucc": warnings.warn("The UCC's functionality as `default_backend` is not well verified", ExperimentalWarning) @@ -671,12 +673,3 @@ def destroy_model_parallel(): # Used to warn when the UCC is specified. class ExperimentalWarning(Warning): pass - - -def check_torch_ucc_availability() -> None: - try: - import torch_ucc # NOQA - except ImportError: - raise ImportError( - "UCC backend requires [torch_ucc](https://github.com/facebookresearch/torch_ucc) but not found" - ) diff --git a/apex/transformer/testing/commons.py b/apex/transformer/testing/commons.py index 226e449db..83f894cb5 100644 --- a/apex/transformer/testing/commons.py +++ b/apex/transformer/testing/commons.py @@ -35,7 +35,7 @@ Batch, ) from apex.transformer.testing import global_vars - +from apex.transformer._ucc_util import HAS_UCC TEST_SUCCESS_MESSAGE = ">> passed the test :-)" @@ -257,7 +257,8 @@ def initialize_distributed(backend="nccl"): if backend not in ("nccl", "ucc"): raise RuntimeError(f"Currently only nccl & ucc are supported but {backend}") if backend == "ucc": - import torch_ucc # NOQA + if not HAS_UCC: + raise ImportError("UCC backend requires pytorch source build with UCC installed and enabled") args = global_vars.get_args() local_rank = args.local_rank diff --git a/apex/transformer/testing/distributed_test_base.py b/apex/transformer/testing/distributed_test_base.py index e0823e80c..bcfbfc032 100644 --- a/apex/transformer/testing/distributed_test_base.py +++ b/apex/transformer/testing/distributed_test_base.py @@ -9,12 +9,7 @@ from torch.testing._internal import common_utils from torch.testing._internal import common_distributed -HAS_TORCH_UCC = None -try: - import torch_ucc - HAS_TORCH_UCC = True -except ImportError: - HAS_TORCH_UCC = False +from apex.transformer._ucc_util import HAS_UCC # NOTE(mkozuki): Version guard for ucc. ref: https://github.com/openucx/ucc/issues/496 _TORCH_UCC_COMPAT_NVIDIA_DRIVER_VERSION = Version("470.42.01") @@ -89,16 +84,16 @@ class NcclDistributedTestBase(DistributedTestBase): DISTRIBUTED_BACKEND = "nccl" - @unittest.skipUnless( - HAS_TORCH_UCC, - "Requires [`torch_ucc`](https://github.com/facebookresearch/torch_ucc)", + HAS_UCC, + "Requires either torch ucc or pytorch build from source with native ucc installed and enabled", ) @unittest.skipUnless( HAS_TORCH_UCC_COMPAT_NVIDIA_DRIVER, f"`torch_ucc` requires NVIDIA driver >= {_TORCH_UCC_COMPAT_NVIDIA_DRIVER_VERSION} but {_driver_version} found. " "See https://github.com/openucx/ucc/issues/496", ) + class UccDistributedTestBase(DistributedTestBase): DISTRIBUTED_BACKEND = "ucc" diff --git a/tests/L0/run_transformer/run_bert_minimal_test.py b/tests/L0/run_transformer/run_bert_minimal_test.py index 639c31e77..eb3dac3c5 100644 --- a/tests/L0/run_transformer/run_bert_minimal_test.py +++ b/tests/L0/run_transformer/run_bert_minimal_test.py @@ -1,13 +1,7 @@ import random import torch -try: - import torch_ucc -except ImportError: - HAS_TORCH_UCC = False -else: - HAS_TORCH_UCC = True - print("Use UCC as backend of Pipeline Parallel ProcessGroups") +from apex.transformer._ucc_util import HAS_UCC from apex.transformer.enums import ModelType from apex.transformer import tensor_parallel from apex.transformer import parallel_state @@ -187,7 +181,7 @@ def train( init = True try: virtual_pipeline_model_parallel_sizes = (None, 2,) - if HAS_TORCH_UCC: + if HAS_UCC: # Deliberately skipping test with interleaved schedule for BERT model. # It deadlocks on hybrid UCC/NCCL backend. virtual_pipeline_model_parallel_sizes = (None,) @@ -217,7 +211,7 @@ def train( args.pipeline_model_parallel_size, virtual_pipeline_model_parallel_size, default_backend="nccl", - p2p_backend="ucc" if HAS_TORCH_UCC else "nccl", + p2p_backend="ucc" if HAS_UCC else "nccl", ) pipeline_model_parallel_size = ( parallel_state.get_pipeline_model_parallel_world_size() diff --git a/tests/L0/run_transformer/run_gpt_minimal_test.py b/tests/L0/run_transformer/run_gpt_minimal_test.py index b3674c898..f5bf83422 100644 --- a/tests/L0/run_transformer/run_gpt_minimal_test.py +++ b/tests/L0/run_transformer/run_gpt_minimal_test.py @@ -3,14 +3,8 @@ import time import torch -try: - import torch_ucc -except ImportError: - HAS_TORCH_UCC = False -else: - HAS_TORCH_UCC = True - print("Use UCC as backend of Pipeline Parallel ProcessGroups") +from apex.transformer._ucc_util import HAS_UCC from apex.transformer import parallel_state from apex.transformer.enums import ModelType from apex.transformer.tensor_parallel import model_parallel_cuda_manual_seed @@ -198,7 +192,7 @@ def train(model, optim, pipeline_model_parallel_size, async_comm): tensor_model_parallel_size_=args.tensor_model_parallel_size, pipeline_model_parallel_size_=args.pipeline_model_parallel_size, default_backend="nccl", - p2p_backend="ucc" if HAS_TORCH_UCC else "nccl", + p2p_backend="ucc" if HAS_UCC else "nccl", ) pipeline_model_parallel_size = ( diff --git a/tests/L0/run_transformer/test_layers.py b/tests/L0/run_transformer/test_layers.py index 80cdff0f2..e2df2acfc 100644 --- a/tests/L0/run_transformer/test_layers.py +++ b/tests/L0/run_transformer/test_layers.py @@ -279,6 +279,7 @@ def test_row_parallel_linear_gradient_accumulation_fusion(self) -> None: def test_row_parallel_linear_gradient_accumulation_fusion_in_fp16(self) -> None: self._row_parallel_linear_test_impl(True, True, False) + # fails on native ucc and torch ucc: ucc does not support reduce scatter @unittest.skipIf(torch.cuda.device_count() < 2, "Sequence Parallel requires >=2 GPUs") def test_row_parallel_linear_sequence_parallel(self) -> None: self._row_parallel_linear_test_impl(False, False, True) diff --git a/tests/L0/run_transformer/test_pipeline_parallel_fwd_bwd.py b/tests/L0/run_transformer/test_pipeline_parallel_fwd_bwd.py index 191f09ff0..81c48e206 100644 --- a/tests/L0/run_transformer/test_pipeline_parallel_fwd_bwd.py +++ b/tests/L0/run_transformer/test_pipeline_parallel_fwd_bwd.py @@ -30,10 +30,9 @@ ) from apex.transformer.testing.distributed_test_base import NcclDistributedTestBase from apex.transformer.testing.distributed_test_base import UccDistributedTestBase -from apex.transformer.testing.distributed_test_base import HAS_TORCH_UCC from apex.transformer.testing.distributed_test_base import HAS_TORCH_UCC_COMPAT_NVIDIA_DRIVER from apex.transformer.testing import commons as testing_utils - +from apex.transformer._ucc_util import HAS_UCC logging.getLogger("torch").setLevel(logging.WARNING) logging.getLogger("apex").setLevel(logging.WARNING) @@ -270,6 +269,7 @@ def test_inference_async_pipelining_without_interleaving(self, sync_batch_comm: sync_batch_comm=sync_batch_comm, ) + # fails on native ucc: times out @unittest.skipUnless(_get_default_world_sizes_model_parallel_world_size()[-1] > 2, "Interleaved schedule requires pipeline_model_parallel_world_size > 2") def test_learning_pipelining_with_interleaving(self, sync_batch_comm: bool = True): self._forward_backward_test_impl( @@ -277,6 +277,7 @@ def test_learning_pipelining_with_interleaving(self, sync_batch_comm: bool = Tru sync_batch_comm=sync_batch_comm, ) + # fails on native ucc: times out @unittest.skipUnless(_get_default_world_sizes_model_parallel_world_size()[-1] > 2, "Interleaved schedule requires pipeline_model_parallel_world_size > 2") def test_inference_pipelining_with_interleaving(self, sync_batch_comm: bool = True): self._forward_backward_test_impl( @@ -284,6 +285,7 @@ def test_inference_pipelining_with_interleaving(self, sync_batch_comm: bool = Tr sync_batch_comm=sync_batch_comm, ) + # fails on native ucc: times out @unittest.skipUnless(_get_default_world_sizes_model_parallel_world_size()[-1] > 2, "Interleaved schedule requires pipeline_model_parallel_world_size > 2") def test_learning_async_pipelining_with_interleaving(self, sync_batch_comm: bool = True): self._forward_backward_test_impl( @@ -291,6 +293,7 @@ def test_learning_async_pipelining_with_interleaving(self, sync_batch_comm: bool sync_batch_comm=sync_batch_comm, ) + # fails on native ucc: times out @unittest.skipUnless(_get_default_world_sizes_model_parallel_world_size()[-1] > 2, "Interleaved schedule requires pipeline_model_parallel_world_size > 2") def test_inference_async_pipelining_with_interleaving(self, sync_batch_comm: bool = True): self._forward_backward_test_impl( @@ -313,7 +316,7 @@ def _run_hybrid_distributed_backend(self, forward_only: bool) -> None: @unittest.skipUnless(HAS_TORCH_UCC_COMPAT_NVIDIA_DRIVER, "Needs driver >= 470.42.01") def _test_hybrid_backends(self, forward_only: bool) -> None: - if HAS_TORCH_UCC: + if HAS_UCC: self._run_hybrid_distributed_backend(forward_only) else: with self.assertRaisesRegex( From 4a268de2e316961b36de1f650cab31c1efecdbc3 Mon Sep 17 00:00:00 2001 From: Masaki Kozuki Date: Thu, 27 Oct 2022 15:12:15 -0700 Subject: [PATCH 17/25] update exception message (#1524) Signed-off-by: Masaki Kozuki Signed-off-by: Masaki Kozuki --- tests/L0/run_transformer/test_pipeline_parallel_fwd_bwd.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/L0/run_transformer/test_pipeline_parallel_fwd_bwd.py b/tests/L0/run_transformer/test_pipeline_parallel_fwd_bwd.py index 81c48e206..a1193f1c6 100644 --- a/tests/L0/run_transformer/test_pipeline_parallel_fwd_bwd.py +++ b/tests/L0/run_transformer/test_pipeline_parallel_fwd_bwd.py @@ -321,7 +321,7 @@ def _test_hybrid_backends(self, forward_only: bool) -> None: else: with self.assertRaisesRegex( ImportError, - re.escape("UCC backend requires [torch_ucc](https://github.com/facebookresearch/torch_ucc) but not found"), + re.escape("UCC backend requires pytorch source build with UCC installed and enabled"), ): self._run_hybrid_distributed_backend(forward_only) From df833cb226e98f5e400397d7add233897c67bd19 Mon Sep 17 00:00:00 2001 From: yaoyu-33 <54727607+yaoyu-33@users.noreply.github.com> Date: Tue, 22 Nov 2022 11:32:48 -0800 Subject: [PATCH 18/25] Update megatron fused softmax follow megatron-lm (#1539) * Update megatron fused softmax follow megatron-lm Signed-off-by: Yu Yao * Add mask=None support in scaled_masked_softmax Signed-off-by: Yu Yao * Update setup.py for scaled_softmax_cuda Signed-off-by: Yu Yao * Add tests for fused_scale_softmax (mask=None) Signed-off-by: Yu Yao * Assert grad equal in fused softmax test Signed-off-by: Yu Yao * Revert "Assert grad equal in fused softmax test" Signed-off-by: Yu Yao Signed-off-by: Yu Yao Co-authored-by: Yu Yao --- apex/transformer/functional/fused_softmax.py | 48 +++- csrc/megatron/scaled_masked_softmax.h | 216 +++++++++++++++++- csrc/megatron/scaled_masked_softmax_cuda.cu | 2 +- csrc/megatron/scaled_softmax.cpp | 75 ++++++ csrc/megatron/scaled_softmax_cuda.cu | 104 +++++++++ setup.py | 33 +++ .../L0/run_transformer/test_fused_softmax.py | 78 +++++++ 7 files changed, 547 insertions(+), 9 deletions(-) create mode 100644 csrc/megatron/scaled_softmax.cpp create mode 100644 csrc/megatron/scaled_softmax_cuda.cu diff --git a/apex/transformer/functional/fused_softmax.py b/apex/transformer/functional/fused_softmax.py index 06a45c461..b0bfebb93 100644 --- a/apex/transformer/functional/fused_softmax.py +++ b/apex/transformer/functional/fused_softmax.py @@ -93,9 +93,14 @@ def backward(ctx, output_grads): def scaled_masked_softmax(inputs, mask, scale): # input is 4D tensor (b, np, sq, sk) - args = _cast_if_autocast_enabled(inputs, mask, scale) - with torch.cuda.amp.autocast(enabled=False): - return ScaledMaskedSoftmax.apply(*args) + if mask is not None: + args = _cast_if_autocast_enabled(inputs, mask, scale) + with torch.cuda.amp.autocast(enabled=False): + return ScaledMaskedSoftmax.apply(*args) + else: + args = _cast_if_autocast_enabled(inputs, scale) + with torch.cuda.amp.autocast(enabled=False): + return ScaledSoftmax.apply(*args) class GenericScaledMaskedSoftmax(torch.autograd.Function): @@ -125,6 +130,37 @@ def generic_scaled_masked_softmax(inputs, mask, scale): return GenericScaledMaskedSoftmax.apply(*args) +class ScaledSoftmax(torch.autograd.Function): + """ + Fused operation which performs following two operations in sequence + 1. Scale the tensor. + 2. Perform softmax. + """ + + @staticmethod + def forward(ctx, inputs, scale): + import scaled_softmax_cuda + + scale_t = torch.tensor([scale]) + + softmax_results = scaled_softmax_cuda.forward( + inputs, scale_t[0] + ) + ctx.save_for_backward(softmax_results, scale_t) + return softmax_results + + @staticmethod + def backward(ctx, output_grads): + import scaled_softmax_cuda + + softmax_results, scale_t = ctx.saved_tensors + + input_grads = scaled_softmax_cuda.backward( + output_grads, softmax_results, scale_t[0] + ) + return input_grads, None, None + + class FusedScaleMaskSoftmax(torch.nn.Module): """ fused operation: scaling + mask + softmax @@ -191,14 +227,14 @@ def is_kernel_available(self, mask, b, np, sq, sk): and self.input_in_float16 # input must be fp16 and ( self.attn_mask_type == AttnMaskType.causal - or (self.attn_mask_type == AttnMaskType.padding and mask is not None) + or self.attn_mask_type == AttnMaskType.padding ) - and 16 < sk <= 2048 # sk must be 16 ~ 2048 + and 16 < sk <= 4096 # sk must be 16 ~ 4096 and sq % 4 == 0 # sq must be divisor of 4 and sk % 4 == 0 # sk must be divisor of 4 and attn_batches % 4 == 0 # np * b must be divisor of 4 ): - if 0 <= sk <= 2048: + if 0 <= sk <= 4096: batch_per_block = self.get_batch_per_block(sq, sk, b, np) if self.attn_mask_type == AttnMaskType.causal: diff --git a/csrc/megatron/scaled_masked_softmax.h b/csrc/megatron/scaled_masked_softmax.h index 78a29cf3b..ad35b3baa 100644 --- a/csrc/megatron/scaled_masked_softmax.h +++ b/csrc/megatron/scaled_masked_softmax.h @@ -90,6 +90,117 @@ __device__ __forceinline__ void warp_reduce(acc_t* sum) { } } + +/* + * Extended softmax (from native aten pytorch) with following additional features + * 1) input scaling + */ +template +__global__ void scaled_softmax_warp_forward( + output_t *dst, + const input_t *src, + const acc_t scale, + int micro_batch_size, + int element_count) +{ + // WARP_SIZE and WARP_BATCH must match the return values batches_per_warp and + // warp_size of method warp_softmax_forward_kernel. + constexpr int next_power_of_two = 1 << log2_elements; + constexpr int WARP_SIZE = (next_power_of_two < C10_WARP_SIZE) ? next_power_of_two : C10_WARP_SIZE; + constexpr int WARP_ITERATIONS = next_power_of_two / WARP_SIZE; + constexpr int WARP_BATCH = (next_power_of_two <= 128) ? 2 : 1; + constexpr int ELEMENTS_PER_LDG_STG = (WARP_ITERATIONS < 4) ? 1 : 4; + + // blockDim/threadIdx = (WARP_SIZE, WARPS_PER_BLOCK, ) + // gridDim/blockIdx = (seq_len, attn_heads, batches) + int first_batch = (blockDim.y * (blockIdx.x + gridDim.x * (blockIdx.y + gridDim.y * blockIdx.z))+ threadIdx.y) * WARP_BATCH; + + // micro_batch_size might not be a multiple of WARP_BATCH. Check how + // many batches have to computed within this WARP. + int local_batches = micro_batch_size - first_batch; + if (local_batches > WARP_BATCH) + local_batches = WARP_BATCH; + + // there might be multiple batches per warp. compute the index within the batch + int local_idx = threadIdx.x; + + src += first_batch * element_count + ELEMENTS_PER_LDG_STG * local_idx; + dst += first_batch * element_count + ELEMENTS_PER_LDG_STG * local_idx; + + // load data from global memory + acc_t elements[WARP_BATCH][WARP_ITERATIONS]; + input_t temp_data[ELEMENTS_PER_LDG_STG]; + #pragma unroll + for (int i = 0; i < WARP_BATCH; ++i) { + int batch_element_count = (i >= local_batches) ? 0 : element_count; + + #pragma unroll + for (int it = 0; it < WARP_ITERATIONS; it+=ELEMENTS_PER_LDG_STG) { + int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE; + + if (element_index < batch_element_count) { + int itr_idx = i*element_count+it*WARP_SIZE; + copy_vector(temp_data, src + itr_idx); + + #pragma unroll + for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) { + elements[i][it + element] = (acc_t)temp_data[element] * scale; + } + } else { + #pragma unroll + for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) { + elements[i][it + element] = -std::numeric_limits::infinity(); + } + } + } + } + + // compute max_value + acc_t max_value[WARP_BATCH]; + #pragma unroll + for (int i = 0; i < WARP_BATCH; ++i) { + max_value[i] = elements[i][0]; + #pragma unroll + for (int it = 1; it < WARP_ITERATIONS; ++it) { + max_value[i] = (max_value[i] > elements[i][it]) ? max_value[i] : elements[i][it]; + } + } + warp_reduce(max_value); + + acc_t sum[WARP_BATCH] { 0.0f }; + #pragma unroll + for (int i = 0; i < WARP_BATCH; ++i) { + #pragma unroll + for (int it = 0; it < WARP_ITERATIONS; ++it) { + elements[i][it] = std::exp((elements[i][it] - max_value[i])); + sum[i] += elements[i][it]; + } + } + warp_reduce(sum); + + // store result + output_t out[ELEMENTS_PER_LDG_STG]; + #pragma unroll + for (int i = 0; i < WARP_BATCH; ++i) { + if (i >= local_batches) + break; + #pragma unroll + for (int it = 0; it < WARP_ITERATIONS; it+=ELEMENTS_PER_LDG_STG) { + int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE; + if (element_index < element_count) { + #pragma unroll + for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) { + out[element] = elements[i][it + element] / sum[i]; + } + copy_vector(dst + i * element_count + it * WARP_SIZE, out); + } else { + break; + } + } + } +} + + /* * Extended softmax (from native aten pytorch) with following additional features * 1) input scaling @@ -326,6 +437,98 @@ int get_batch_per_block(int query_seq_len, int key_seq_len, int batches, int att return batches_per_block; } +template +void dispatch_scaled_softmax_forward( + output_t *dst, + const input_t *src, + const input_t scale, + int query_seq_len, + int key_seq_len, + int batches, + int attn_heads) +{ + TORCH_INTERNAL_ASSERT(key_seq_len >= 0 && key_seq_len <= 4096 ); + if (key_seq_len == 0) { + return; + } else { + int log2_elements = log2_ceil(key_seq_len); + const int next_power_of_two = 1 << log2_elements; + int batch_count = batches * attn_heads * query_seq_len; + + // This value must match the WARP_SIZE constexpr value computed inside softmax_warp_forward. + int warp_size = (next_power_of_two < C10_WARP_SIZE) ? next_power_of_two : C10_WARP_SIZE; + + // This value must match the WARP_BATCH constexpr value computed inside softmax_warp_forward. + int batches_per_warp = (next_power_of_two <= 128) ? 2 : 1; + + // use 128 threads per block to maximimize gpu utilization + constexpr int threads_per_block = 128; + + int warps_per_block = (threads_per_block / warp_size); + int batches_per_block = warps_per_block * batches_per_warp; + TORCH_INTERNAL_ASSERT(query_seq_len%batches_per_block == 0); + dim3 blocks(query_seq_len/batches_per_block, attn_heads, batches); + dim3 threads(warp_size, warps_per_block, 1); + // Launch code would be more elegant if C++ supported FOR CONSTEXPR + switch (log2_elements) { + case 0: // 1 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 1: // 2 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 2: // 4 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 3: // 8 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 4: // 16 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 5: // 32 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 6: // 64 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 7: // 128 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 8: // 256 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 9: // 512 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 10: // 1024 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 11: // 2048 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + case 12: // 4096 + scaled_softmax_warp_forward + <<>>(dst, src, scale, batch_count, key_seq_len); + break; + default: + break; + } + } +} + template void dispatch_scaled_masked_softmax_forward( output_t *dst, @@ -338,7 +541,7 @@ void dispatch_scaled_masked_softmax_forward( int attn_heads, int pad_batches) { - TORCH_INTERNAL_ASSERT(key_seq_len >= 0 && key_seq_len <= 2048 ); + TORCH_INTERNAL_ASSERT(key_seq_len >= 0 && key_seq_len <= 4096 ); if (key_seq_len == 0) { return; } else { @@ -410,6 +613,10 @@ void dispatch_scaled_masked_softmax_forward( scaled_masked_softmax_warp_forward <<>>(dst, src, mask, scale, batch_count, key_seq_len, pad_batches); break; + case 12: // 4096 + scaled_masked_softmax_warp_forward + <<>>(dst, src, mask, scale, batch_count, key_seq_len, pad_batches); + break; default: break; } @@ -427,7 +634,7 @@ void dispatch_scaled_masked_softmax_backward( int batches, int attn_heads) { - TORCH_INTERNAL_ASSERT( key_seq_len >= 0 && key_seq_len <= 2048 ); + TORCH_INTERNAL_ASSERT( key_seq_len >= 0 && key_seq_len <= 4096 ); if (key_seq_len == 0) { return; } else { @@ -498,6 +705,11 @@ void dispatch_scaled_masked_softmax_backward( scaled_masked_softmax_warp_backward <<>>(grad_input, grad, output, scale, batch_count, key_seq_len); break; + case 12: // 4096 + scaled_masked_softmax_warp_backward + <<>>(grad_input, grad, output, scale, batch_count, key_seq_len); + break; + default: break; } diff --git a/csrc/megatron/scaled_masked_softmax_cuda.cu b/csrc/megatron/scaled_masked_softmax_cuda.cu index ac534a072..de911a446 100644 --- a/csrc/megatron/scaled_masked_softmax_cuda.cu +++ b/csrc/megatron/scaled_masked_softmax_cuda.cu @@ -46,7 +46,7 @@ torch::Tensor fwd_cuda( const int attn_heads = input.size(1); const int query_seq_len = input.size(2); const int key_seq_len = input.size(3); - TORCH_INTERNAL_ASSERT(key_seq_len <= 2048); + TORCH_INTERNAL_ASSERT(key_seq_len <= 4096); TORCH_INTERNAL_ASSERT(query_seq_len > 1); TORCH_INTERNAL_ASSERT(pad_batches == 1 || pad_batches == batches); TORCH_INTERNAL_ASSERT(mask.size(1) == 1); diff --git a/csrc/megatron/scaled_softmax.cpp b/csrc/megatron/scaled_softmax.cpp new file mode 100644 index 000000000..9ad989085 --- /dev/null +++ b/csrc/megatron/scaled_softmax.cpp @@ -0,0 +1,75 @@ +/* coding=utf-8 + * Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include + +namespace multihead_attn { +namespace fused_softmax { +namespace scaled_softmax { + +torch::Tensor fwd_cuda( + torch::Tensor const& input, + float scale_factor); + +torch::Tensor bwd_cuda( + torch::Tensor const& output_grads, + torch::Tensor const& softmax_results, + float scale_factor); + +torch::Tensor fwd( + torch::Tensor const& input, + float scale_factor) { + AT_ASSERTM(input.dim() == 4, "expected 4D tensor"); + AT_ASSERTM((input.scalar_type() == at::ScalarType::Half) || + (input.scalar_type() == at::ScalarType::BFloat16), + "Only fp16 and bf16 are supported"); + + return fwd_cuda(input, scale_factor); +} + +torch::Tensor bwd( + torch::Tensor const& output_grads, + torch::Tensor const& softmax_results, + float scale_factor) { + + AT_ASSERTM(output_grads.dim() == 4, "expected 3D tensor"); + AT_ASSERTM(softmax_results.dim() == 4, "expected 3D tensor"); + + AT_ASSERTM((output_grads.scalar_type() == at::ScalarType::Half) || + (output_grads.scalar_type() == at::ScalarType::BFloat16), + "Only fp16 and bf16 are supported"); + AT_ASSERTM((softmax_results.scalar_type() == at::ScalarType::Half) || + (softmax_results.scalar_type() == at::ScalarType::BFloat16), + "Only fp16 and bf16 are supported"); + + return bwd_cuda(output_grads, softmax_results, scale_factor); +} + +} // end namespace scaled_softmax +} // end namespace fused_softmax +} // end namespace multihead_attn + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("forward", + &multihead_attn::fused_softmax::scaled_softmax::fwd, + "Self Multihead Attention scaled, softmax -- Forward."); + m.def("backward", + &multihead_attn::fused_softmax::scaled_softmax::bwd, + "Self Multihead Attention scaled, softmax -- Backward."); +} + diff --git a/csrc/megatron/scaled_softmax_cuda.cu b/csrc/megatron/scaled_softmax_cuda.cu new file mode 100644 index 000000000..dc52a7e18 --- /dev/null +++ b/csrc/megatron/scaled_softmax_cuda.cu @@ -0,0 +1,104 @@ +/* coding=utf-8 + * Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include +#include +#include +#include +#include "scaled_masked_softmax.h" +#include "type_shim.h" + +namespace multihead_attn { +namespace fused_softmax { +namespace scaled_softmax { + +torch::Tensor fwd_cuda( + torch::Tensor const& input, + float scale_factor) +{ + // input is a 4d tensor with dimensions [batches, attn_heads, seq_len, seq_len] + const int batches = input.size(0); + const int attn_heads = input.size(1); + const int query_seq_len = input.size(2); + const int key_seq_len = input.size(3); + TORCH_INTERNAL_ASSERT(key_seq_len <= 4096); + TORCH_INTERNAL_ASSERT(query_seq_len > 1); + + // Output + auto act_options = input.options().requires_grad(false); + torch::Tensor softmax_results = + torch::empty({batches, attn_heads, query_seq_len, key_seq_len}, act_options); + + // Softmax Intermediate Result Ptr + void* input_ptr = static_cast(input.data_ptr()); + void* softmax_results_ptr = static_cast(softmax_results.data_ptr()); + + DISPATCH_HALF_AND_BFLOAT( + input.scalar_type(), + "dispatch_scaled_softmax_forward", + dispatch_scaled_softmax_forward( + reinterpret_cast(softmax_results_ptr), + reinterpret_cast(input_ptr), + scale_factor, + query_seq_len, + key_seq_len, + batches, + attn_heads); + ); + return softmax_results; +} + +torch::Tensor bwd_cuda( + torch::Tensor const& output_grads_, + torch::Tensor const& softmax_results_, + float scale_factor) { + + auto output_grads = output_grads_.contiguous(); + auto softmax_results = softmax_results_.contiguous(); + + //output grads is a 4d tensor with dimensions [batches, attn_heads, seq_len, seq_len] + const int batches = output_grads.size(0); + const int attn_heads = output_grads.size(1); + const int query_seq_len = output_grads.size(2); + const int key_seq_len = output_grads.size(3); + + void* output_grads_ptr = static_cast(output_grads.data_ptr()); + + //Softmax Grad + DISPATCH_HALF_AND_BFLOAT( + output_grads_.scalar_type(), + "dispatch_scaled_masked_softmax_backward", + dispatch_scaled_masked_softmax_backward( + reinterpret_cast(output_grads_ptr), + reinterpret_cast(output_grads_ptr), + reinterpret_cast(softmax_results.data_ptr()), + scale_factor, + query_seq_len, + key_seq_len, + batches, + attn_heads); + ); + + //backward pass is completely in-place + return output_grads; +} +} +} +} + diff --git a/setup.py b/setup.py index 2f033c0d8..5de5acb28 100644 --- a/setup.py +++ b/setup.py @@ -307,6 +307,39 @@ def check_if_rocm_pytorch(): os.path.join(this_dir, 'csrc/megatron')], extra_compile_args={'cxx': ['-O3'] + version_dependent_macros, 'nvcc':nvcc_args_transformer if not IS_ROCM_PYTORCH else hipcc_args_transformer})) + ext_modules.append( + CUDAExtension(name='scaled_softmax_cuda', + sources=['csrc/megatron/scaled_softmax.cpp', + 'csrc/megatron/scaled_softmax_cuda.cu'], + include_dirs=[os.path.join(this_dir, 'csrc'), + os.path.join(this_dir, 'csrc/megatron')], + extra_compile_args={'cxx': ['-O3'] + version_dependent_macros, + 'nvcc':nvcc_args_transformer if not IS_ROCM_PYTORCH else hipcc_args_transformer})) + # TODO (hubertlu): support fused_weight_gradient_mlp_cuda for ROCm + if not IS_ROCM_PYTORCH: + if bare_metal_version >= Version("11.0"): + cc_flag = [] + cc_flag.append("-gencode") + cc_flag.append("arch=compute_70,code=sm_70") + cc_flag.append("-gencode") + cc_flag.append("arch=compute_80,code=sm_80") + if bare_metal_version >= Version("11.1"): + cc_flag.append("-gencode") + cc_flag.append("arch=compute_86,code=sm_86") + if bare_metal_version >= Version("11.8"): + cc_flag.append("-gencode") + cc_flag.append("arch=compute_90,code=sm_90") + + ext_modules.append( + CUDAExtension(name='fused_weight_gradient_mlp_cuda', + sources=['csrc/megatron/fused_weight_gradient_dense.cpp', + 'csrc/megatron/fused_weight_gradient_dense_cuda.cu', + 'csrc/megatron/fused_weight_gradient_dense_16bit_prec_cuda.cu'], + include_dirs=[os.path.join(this_dir, 'csrc'), + os.path.join(this_dir, 'csrc/megatron')], + extra_compile_args={'cxx': ['-O3'] + version_dependent_macros, + 'nvcc':append_nvcc_threads(nvcc_args_transformer + ['--use_fast_math'] + cc_flag) + if not IS_ROCM_PYTORCH else hipcc_args_transformer})) if "--bnp" in sys.argv or "--cuda_ext" in sys.argv: diff --git a/tests/L0/run_transformer/test_fused_softmax.py b/tests/L0/run_transformer/test_fused_softmax.py index 7eb452564..c49967073 100644 --- a/tests/L0/run_transformer/test_fused_softmax.py +++ b/tests/L0/run_transformer/test_fused_softmax.py @@ -139,6 +139,84 @@ def test_autocast_fused_scale_mask_softmax(self): expected.backward(g0) actual.backward(g1) + def test_fused_scale_softmax(self): + """ + attention_scores.shape = [4, 12, 24, 24] + mask = None + """ + for (dtype, scale, softmax_in_fp32, shape) in itertools.product( + (torch.half, torch.bfloat16), (None, 2.0), (False, True), ((4, 12, 24, 24), (32, 12, 4, 214)) + ): + msg = f"{dtype}-{scale}-{softmax_in_fp32}" + input_in_fp16 = dtype == torch.half + input_in_bf16 = dtype == torch.bfloat16 + if not (scale is None or softmax_in_fp32): + with self.assertRaises(RuntimeError, msg=msg): + self._setup_fused_softmax( + input_in_fp16, + input_in_bf16, + scale, + softmax_in_fp32, + AttnMaskType.padding, + ) + return + fused_fn, torch_fn = self._setup_fused_softmax( + input_in_fp16, + input_in_bf16, + scale, + softmax_in_fp32, + AttnMaskType.padding, + ) + + attention_scores_0 = ( + torch.randn(shape) + .to(device="cuda", dtype=dtype) + .requires_grad_(True) + ) + with torch.no_grad(): + attention_scores_1 = attention_scores_0.clone().requires_grad_(True) + mask = None + + expected = fused_fn(attention_scores_0, mask) + actual = torch_fn(attention_scores_1, mask) + self.assertEqual(actual, expected, msg=msg) + + g0 = torch.rand_like(actual) + with torch.no_grad(): + g1 = g0.clone() + expected.backward(g0) + actual.backward(g1) + + def test_autocast_fused_scale_softmax(self): + for dtype in autocast_dtypes: + msg = f"dtype: {dtype}" + input_in_fp16 = dtype == torch.half + input_in_bf16 = dtype == torch.bfloat16 + fused_fn, torch_fn = self._setup_fused_softmax( + input_in_fp16, input_in_bf16, attn_mask_type=AttnMaskType.padding + ) + + attention_scores_0 = ( + torch.randn((4, 12, 24, 24)).cuda().requires_grad_(True) + ) + with torch.no_grad(): + attention_scores_1 = ( + attention_scores_0.clone().to(dtype).requires_grad_(True) + ) + mask = None + + expected = torch_fn(attention_scores_1, mask) + with torch.cuda.amp.autocast(dtype=dtype): + actual = fused_fn(attention_scores_0, mask) + self.assertEqual(actual.dtype, dtype, msg=msg) + self.assertEqual(actual, expected, msg=msg) + + g0 = torch.rand_like(actual) + with torch.no_grad(): + g1 = g0.clone() + expected.backward(g0) + actual.backward(g1) + def test_fused_upper_triangle_mask_softmax(self): """ attn_weights.shape: [4, 12, 24, 24] From 5e5331a18eaf85a860b36d434de6c31055e79765 Mon Sep 17 00:00:00 2001 From: hubertlu-tw Date: Thu, 29 Dec 2022 20:34:17 +0000 Subject: [PATCH 19/25] Resolve filename collision issue in compilation on ROCm (Ref: #77) --- csrc/megatron/scaled_softmax.cpp | 1 - csrc/megatron/scaled_softmax_cuda.cu | 2 +- 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/csrc/megatron/scaled_softmax.cpp b/csrc/megatron/scaled_softmax.cpp index 9ad989085..a25cfa328 100644 --- a/csrc/megatron/scaled_softmax.cpp +++ b/csrc/megatron/scaled_softmax.cpp @@ -14,7 +14,6 @@ * limitations under the License. */ -#include #include #include diff --git a/csrc/megatron/scaled_softmax_cuda.cu b/csrc/megatron/scaled_softmax_cuda.cu index dc52a7e18..7dcbca68e 100644 --- a/csrc/megatron/scaled_softmax_cuda.cu +++ b/csrc/megatron/scaled_softmax_cuda.cu @@ -18,7 +18,7 @@ #include #include #include -#include +//#include #include #include #include "scaled_masked_softmax.h" From f02cc057d1fc0842890ed11abaf8c3bc23f63c3e Mon Sep 17 00:00:00 2001 From: Fuzzkatt Date: Thu, 8 Dec 2022 01:31:26 -0800 Subject: [PATCH 20/25] Refactor run_transformer bert minimal and gpt minimal tests (#1540) * working test_bert_minimal.py * remove some debugging statements * working test_gpt_minimal.py * test_dynamic_batchsize.py having issues with torch.backends.cudnn.allow_tf32 * working test_dynamic_batchsize.py * refactor test_bert_minimal.py, need to investigate rng of MANUAL_SEED for nccl only pipeline with virtual_pipeline_model_parallel_size = 2 * add test_bert_minimal_alt.py for visibility * update test_gpt_minimal.py * lint * update loss cutoff for bert test * split with / without interleaving tests for bert * use skipTest * remove ONCE * add ignore_unknown_args=True * remove old testing files * add num_devices logic to override_args --- apex/transformer/testing/arguments.py | 8 +- apex/transformer/testing/global_vars.py | 6 +- .../run_transformer/run_bert_minimal_test.py | 254 ------------------ .../run_transformer/run_gpt_minimal_test.py | 217 --------------- tests/L0/run_transformer/test_bert_minimal.py | 250 +++++++++++++++++ ...size_test.py => test_dynamic_batchsize.py} | 169 +++++++----- tests/L0/run_transformer/test_gpt_minimal.py | 232 ++++++++++++++++ .../test_transformer_module.py | 97 ------- 8 files changed, 588 insertions(+), 645 deletions(-) delete mode 100644 tests/L0/run_transformer/run_bert_minimal_test.py delete mode 100644 tests/L0/run_transformer/run_gpt_minimal_test.py create mode 100644 tests/L0/run_transformer/test_bert_minimal.py rename tests/L0/run_transformer/{run_dynamic_batchsize_test.py => test_dynamic_batchsize.py} (54%) create mode 100644 tests/L0/run_transformer/test_gpt_minimal.py delete mode 100644 tests/L0/run_transformer/test_transformer_module.py diff --git a/apex/transformer/testing/arguments.py b/apex/transformer/testing/arguments.py index f39c288ce..5c74054ca 100644 --- a/apex/transformer/testing/arguments.py +++ b/apex/transformer/testing/arguments.py @@ -20,7 +20,7 @@ import torch -def parse_args(extra_args_provider=None, defaults={}, +def parse_args(extra_args_provider=None, defaults={}, override_args={}, ignore_unknown_args=False): """Parse all arguments.""" parser = argparse.ArgumentParser(description='Megatron-LM Arguments', @@ -59,16 +59,22 @@ def parse_args(extra_args_provider=None, defaults={}, # Distributed args. args.rank = int(os.getenv('RANK', '0')) args.world_size = int(os.getenv("WORLD_SIZE", '1')) + + for key in override_args: + setattr(args, key, override_args[key]) + # Tensor model parallel size. args.tensor_model_parallel_size = min( args.tensor_model_parallel_size, args.world_size) assert args.world_size % args.tensor_model_parallel_size == 0, 'world size'\ ' ({}) is not divisible by tensor model parallel size ({})'.format( args.world_size, args.tensor_model_parallel_size) + # Pipeline model parallel size. args.pipeline_model_parallel_size = min( args.pipeline_model_parallel_size, (args.world_size // args.tensor_model_parallel_size)) + args.transformer_pipeline_model_parallel_size = ( args.pipeline_model_parallel_size - 1 if args.standalone_embedding_stage else diff --git a/apex/transformer/testing/global_vars.py b/apex/transformer/testing/global_vars.py index 6b8537456..fa68f2c14 100644 --- a/apex/transformer/testing/global_vars.py +++ b/apex/transformer/testing/global_vars.py @@ -84,11 +84,12 @@ def get_timers(): return _GLOBAL_TIMERS -def set_global_variables(extra_args_provider=None, args_defaults={}, +def set_global_variables(extra_args_provider=None, args_defaults={}, override_args={}, ignore_unknown_args=False): """Set args, tokenizer, tensorboard-writer, adlr-autoresume, and timers.""" args = _parse_args(extra_args_provider=extra_args_provider, defaults=args_defaults, + override_args=override_args, ignore_unknown_args=ignore_unknown_args) # _build_num_microbatches_calculator(args) # if args.vocab_file: @@ -98,13 +99,14 @@ def set_global_variables(extra_args_provider=None, args_defaults={}, _set_timers() -def _parse_args(extra_args_provider=None, defaults={}, +def _parse_args(extra_args_provider=None, defaults={}, override_args={}, ignore_unknown_args=False): """Parse entire arguments.""" global _GLOBAL_ARGS _ensure_var_is_not_initialized(_GLOBAL_ARGS, 'args') _GLOBAL_ARGS = parse_args(extra_args_provider=extra_args_provider, defaults=defaults, + override_args=override_args, ignore_unknown_args=ignore_unknown_args) return _GLOBAL_ARGS diff --git a/tests/L0/run_transformer/run_bert_minimal_test.py b/tests/L0/run_transformer/run_bert_minimal_test.py deleted file mode 100644 index eb3dac3c5..000000000 --- a/tests/L0/run_transformer/run_bert_minimal_test.py +++ /dev/null @@ -1,254 +0,0 @@ -import random -import torch - -from apex.transformer._ucc_util import HAS_UCC -from apex.transformer.enums import ModelType -from apex.transformer import tensor_parallel -from apex.transformer import parallel_state -from apex.transformer.log_util import set_logging_level -from apex.transformer.tensor_parallel import vocab_parallel_cross_entropy -from apex.transformer.pipeline_parallel.utils import setup_microbatch_calculator -from apex.transformer.pipeline_parallel.utils import unwrap_model -from apex.transformer.pipeline_parallel.utils import ( - average_losses_across_data_parallel_group, -) -from apex.transformer.pipeline_parallel.schedules import get_forward_backward_func -from apex.transformer.pipeline_parallel.schedules.common import build_model -from apex.transformer.pipeline_parallel.schedules.common import ( - _get_params_for_weight_decay_optimization, -) - -from apex.transformer.testing.standalone_bert import bert_model_provider -from apex.transformer.testing import global_vars -from apex.transformer.testing.commons import TEST_SUCCESS_MESSAGE -from apex.transformer.testing.commons import initialize_distributed -from apex.transformer.testing.commons import print_separator - -import warnings - - -class DebugWarning(Warning): - pass - - -set_logging_level("WARNING") -mode = None -MANUAL_SEED = 42 -inds = None -masks = None -data_idx = 0 -MASK_PROB = 0.1 -EASY_MODE = False -EASY_MODE_SIZ = 32 -ONCE = False - - -def download_fancy_data(): - # import requests - # response = requests.get('https://internet.com/book.txt') - # text = ' '.join(response.text.split()) - text = """ - An original sentence not subject to any license restrictions, copyright, or royalty payments. Nothing to see here. Commercial or non-commercial use. Research or non-research purposes. The quick brown fox jumps over the lazy dog. Lorem ipsum. - """ - text = text * 1024 - encoded = text.encode("ascii", "replace") - ints = [int(encoded[i]) for i in range(len(encoded))] - return torch.tensor(ints) - - -# build a batch given sequence_len and batch size -def generate_fancy_data_labels(sequence_len, batch_size): - global data_idx - global inds - global masks - global MANUAL_SEED - temps = [] - for i in range(batch_size): - if inds is None or data_idx >= len(inds): - # hack as use of RNG will fall out of sync due to pipelines being different - torch.manual_seed(MANUAL_SEED) - inds = torch.randperm(effective_length, device="cuda") - masks = ( - torch.rand( - len(inds) // batch_size + 1, batch_size, sequence_len, device="cuda" - ) - >= MASK_PROB - ).long() - MANUAL_SEED += 1 - print("new epoch", len(inds)) - data_idx = 0 - print("my start", inds[0:5]) - print("masks_checksum:", torch.sum(masks)) - if EASY_MODE: - data_idx_ = data_idx % EASY_MODE_SIZ - else: - data_idx_ = data_idx - offset = inds[data_idx_] # * SEQUENCE_LEN - data_idx += 1 - - curr = fancy_data[offset : offset + sequence_len].clone().detach() - temps.append(curr) - temp = torch.stack(temps, dim=0).cuda() - mask = masks[data_idx // batch_size] - mask_not = torch.logical_not(mask).long() - data = mask * temp + mask_not * 124 - label = temp - if parallel_state.get_tensor_model_parallel_rank() == 0: - data_dict = {"text": data, "label": label, "mask_not": mask_not} - else: - data_dict = None - keys = ["text", "label", "mask_not"] - dtype = torch.int64 - broadcasted_data = tensor_parallel.broadcast_data(keys, data_dict, torch.long) - return ( - broadcasted_data["text"].long(), - broadcasted_data["label"].long(), - broadcasted_data["mask_not"], - ) - - -easy_data = None - - -def fwd_step_func(batch, model): - data, label, loss_mask = batch - y = model(data, torch.ones_like(data), lm_labels=label) - - def loss_func(output_tensor): - global ONCE - output_tensor, _ = output_tensor - lm_loss_ = output_tensor.float() - lm_loss = torch.sum(lm_loss_.view(-1) * loss_mask.reshape(-1)) / loss_mask.sum() - averaged_loss = average_losses_across_data_parallel_group([lm_loss]) - if data_idx >= 1536: - assert averaged_loss < 4.8 - if not ONCE: - print("LOSS OK") - ONCE = True - return lm_loss, {"avg": averaged_loss} - - return y, loss_func - - -def train( - model, optim, virtual_pipeline_model_parallel_size, pipeline_model_parallel_size, async_comm -): - sequence_len = global_vars.get_args().seq_length - micro_batch_size = global_vars.get_args().micro_batch_size - hidden_size = global_vars.get_args().hidden_size - forward_backward_func = get_forward_backward_func( - virtual_pipeline_model_parallel_size, pipeline_model_parallel_size - ) - tensor_shape = (args.seq_length, args.micro_batch_size, args.hidden_size) - for _ in range(16): - batch = generate_fancy_data_labels(sequence_len, batch_size) - optim.zero_grad() - forward_backward_func( - fwd_step_func, - batch, - model, - forward_only=False, - tensor_shape=tensor_shape, - async_comm=async_comm, - sequence_parallel_enabled=global_vars.get_args().sequence_parallel, - ) - # All-reduce layernorm parameters across model parallel nodes - # when sequence parallelism is used - if parallel_state.get_tensor_model_parallel_world_size() > 1 and global_vars.get_args().sequence_parallel: - for model_module in model: - unwrapped_model = unwrap_model(model_module) - for param in unwrapped_model.parameters(): - if getattr(param, 'sequence_parallel_enabled', False): - grad = param.grad - torch.distributed.all_reduce(grad, group=parallel_state.get_tensor_model_parallel_group()) - - optim.step() - - -if __name__ == "__main__": - global fancy_data - global effective_length - - global_vars.set_global_variables() - - fancy_data = download_fancy_data() - effective_length = fancy_data.size(0) // global_vars.get_args().seq_length - effective_length = fancy_data.size(0) - global_vars.get_args().seq_length - - initialize_distributed("nccl") - world_size = torch.distributed.get_world_size() - failure = None - init = True - try: - virtual_pipeline_model_parallel_sizes = (None, 2,) - if HAS_UCC: - # Deliberately skipping test with interleaved schedule for BERT model. - # It deadlocks on hybrid UCC/NCCL backend. - virtual_pipeline_model_parallel_sizes = (None,) - for virtual_pipeline_model_parallel_size in virtual_pipeline_model_parallel_sizes: - args = global_vars.get_args() - async_comm = not args.sequence_parallel and virtual_pipeline_model_parallel_size is None - data_idx = 0 - ONCE = False - if init: - init = False - args = global_vars.get_args() - args.padded_vocab_size = 128 # needed in standalone gpt - args.model_type = ModelType.encoder_or_decoder - batch_size = args.global_batch_size - micro_batch_size = args.micro_batch_size - setup_microbatch_calculator( - args.rank, - args.rampup_batch_size, - args.global_batch_size, - args.micro_batch_size, - args.data_parallel_size, - ) - else: - parallel_state.destroy_model_parallel() - parallel_state.initialize_model_parallel( - args.tensor_model_parallel_size, - args.pipeline_model_parallel_size, - virtual_pipeline_model_parallel_size, - default_backend="nccl", - p2p_backend="ucc" if HAS_UCC else "nccl", - ) - pipeline_model_parallel_size = ( - parallel_state.get_pipeline_model_parallel_world_size() - ) - - tensor_parallel.random.model_parallel_cuda_manual_seed(0) - model = build_model( - bert_model_provider, - wrap_with_ddp=parallel_state.get_data_parallel_world_size() > 1, - virtual_pipeline_model_parallel_size=virtual_pipeline_model_parallel_size, - cpu_offload=args.cpu_offload, - ) - assert isinstance(model, list) - assert len(model) == ( - 1 - if virtual_pipeline_model_parallel_size is None - else virtual_pipeline_model_parallel_size - ) - _param_groups = _get_params_for_weight_decay_optimization(model) - optim = torch.optim.Adam(_param_groups) - print(effective_length) - print(fancy_data.size(0)) - train( - model, - optim, - virtual_pipeline_model_parallel_size, - args.pipeline_model_parallel_size, - async_comm, - ) - except Exception as e: - failure = str(e) - finally: - parallel_state.destroy_model_parallel() - if failure is not None: - warnings.warn( - f"Minimal BERT Pipeline Parallel Failed with: {failure}", DebugWarning - ) - print(f"Minimal BERT Pipeline Parallel Failed with: {failure}") - torch.distributed.barrier() - print(TEST_SUCCESS_MESSAGE) diff --git a/tests/L0/run_transformer/run_gpt_minimal_test.py b/tests/L0/run_transformer/run_gpt_minimal_test.py deleted file mode 100644 index f5bf83422..000000000 --- a/tests/L0/run_transformer/run_gpt_minimal_test.py +++ /dev/null @@ -1,217 +0,0 @@ -from functools import partial -from typing import List -import time - -import torch - -from apex.transformer._ucc_util import HAS_UCC -from apex.transformer import parallel_state -from apex.transformer.enums import ModelType -from apex.transformer.tensor_parallel import model_parallel_cuda_manual_seed -from apex.transformer.pipeline_parallel.utils import setup_microbatch_calculator -from apex.transformer.pipeline_parallel.utils import unwrap_model -from apex.transformer.pipeline_parallel.utils import ( - average_losses_across_data_parallel_group, -) -from apex.transformer.pipeline_parallel.utils import get_ltor_masks_and_position_ids -from apex.transformer.pipeline_parallel.schedules.common import build_model -from apex.transformer.pipeline_parallel.schedules.common import ( - _get_params_for_weight_decay_optimization, -) -from apex.transformer.pipeline_parallel.schedules.fwd_bwd_pipelining_without_interleaving import ( - forward_backward_pipelining_without_interleaving, -) -from apex.transformer.testing.standalone_gpt import gpt_model_provider -from apex.transformer.testing import global_vars -from apex.transformer.testing.commons import TEST_SUCCESS_MESSAGE -from apex.transformer.testing.commons import initialize_distributed - -MANUAL_SEED = 42 -inds = None -data_idx = 0 -N_VOCAB = 128 - - -def download_fancy_data(): - # import requests - # response = requests.get('https://internet.com/book.txt') - # text = ' '.join(response.text.split()) - text = """ - An original sentence not subject to any license restrictions, copyright, or royalty payments. Nothing to see here. Commercial or non-commercial use. Research or non-research purposes. The quick brown fox jumps over the lazy dog. Lorem ipsum. - """ - text = text * 1024 - encoded = text.encode("ascii", "replace") - ints = [int(encoded[i]) for i in range(len(encoded))] - return torch.tensor(ints) - - -# build a batch given sequence_len and batch size -def generate_fancy_data_labels(sequence_len, batch_size): - global data_idx - global inds - global MANUAL_SEED - temps = list() - for i in range(batch_size): - if inds is None or data_idx >= len(inds): - # hack as use of RNG will fall out of sync due to pipelines being different - model_parallel_cuda_manual_seed(MANUAL_SEED) - inds = torch.randperm(effective_length, device="cuda") - MANUAL_SEED += 1 - data_idx = 0 - data_idx_ = data_idx - offset = inds[data_idx_] - data_idx += 1 - curr = fancy_data[offset : offset + sequence_len + 1].clone().detach() - temps.append(curr) - temp = torch.stack(temps, dim=0).cuda() - return temp - - -easy_data = None - - -def get_batch(int_tensors: List[torch.Tensor]): - data = int_tensors[0] - # Unpack. - tokens_ = data.long() - labels = tokens_[:, 1:].contiguous() - tokens = tokens_[:, :-1].contiguous() - # Get the masks and position ids. - attention_mask, loss_mask, position_ids = get_ltor_masks_and_position_ids( - tokens, - N_VOCAB, # tokenizer.eod, - False, # args.reset_position_ids, - False, # args.reset_attention_mask, - False, # args.eod_mask_loss, - ) - return tokens, labels, loss_mask, attention_mask, position_ids - - -# Ref: https://github.com/NVIDIA/Megatron-LM/blob/b31e1296354e979722627a6c4dedafe19b51fa97/pretrain_gpt.py#L75 -def loss_func(loss_mask, output_tensor): - losses = output_tensor.float() - loss_mask = loss_mask.view(-1).float() - loss = torch.sum(losses.view(-1) * loss_mask) / loss_mask.sum() - - # Reduce loss for logging. - averaged_loss = average_losses_across_data_parallel_group([loss]) - - return loss, {"lm loss": averaged_loss[0]} - - -# Ref: https://github.com/NVIDIA/Megatron-LM/blob/b31e1296354e979722627a6c4dedafe19b51fa97/pretrain_gpt.py#L86 -def fwd_step_func(batch, model): - """Forward step.""" - tokens, labels, loss_mask, attention_mask, position_ids = get_batch(batch) - output_tensor = model(tokens, position_ids, attention_mask, labels=labels) - return output_tensor, partial(loss_func, loss_mask) - - -def train(model, optim, pipeline_model_parallel_size, async_comm): - sequence_len = global_vars.get_args().seq_length - micro_batch_size = global_vars.get_args().micro_batch_size - hidden_size = global_vars.get_args().hidden_size - fwd_bwd_func = forward_backward_pipelining_without_interleaving - - tensor_shape = (args.seq_length, args.micro_batch_size, args.hidden_size) - runtime = 0 - # training loop - for i in range(3): - since = time.time() - if torch.distributed.get_rank() == 0: - print("begin iter", i) - batch = [ - generate_fancy_data_labels(args.seq_length, args.global_batch_size) - for _ in range(pipeline_model_parallel_size) - ] - if torch.distributed.get_rank() == 0: - print("finished making batch...") - optim.zero_grad() - fwd_bwd_func( - fwd_step_func, - batch, - model, - forward_only=False, - tensor_shape=tensor_shape, - async_comm=async_comm, - sequence_parallel_enabled=args.sequence_parallel, - ) - if torch.distributed.get_rank() == 0: - print("finished forward step") - # All-reduce layernorm parameters across model parallel nodes - # when sequence parallelism is used - if parallel_state.get_tensor_model_parallel_world_size() > 1 and global_vars.get_args().sequence_parallel: - for model_module in model: - unwrapped_model = unwrap_model(model_module) - for param in unwrapped_model.parameters(): - if getattr(param, 'sequence_parallel_enabled', False): - grad = param.grad - torch.distributed.all_reduce(grad, group=parallel_state.get_tensor_model_parallel_group()) - optim.step() - if torch.distributed.get_rank() == 0: - print("finished iter", i) - runtime += time.time() - since - return runtime / 3.0 - - -if __name__ == "__main__": - init = True - global_vars.set_global_variables() - for async_comm in (False,) if global_vars.get_args().sequence_parallel else (False, True): - global fancy_data - global effective_length - - if init: - init = False - - fancy_data = download_fancy_data() - args = global_vars.get_args() - args.model_type = ModelType.encoder_or_decoder - effective_length = fancy_data.size(0) // args.seq_length - effective_length = fancy_data.size(0) - args.seq_length - - initialize_distributed("nccl") - world_size = torch.distributed.get_world_size() - - failure = None - args.padded_vocab_size = 128 - batch_size = args.global_batch_size - micro_batch_size = args.micro_batch_size - setup_microbatch_calculator( - args.rank, - args.rampup_batch_size, - args.global_batch_size, - args.micro_batch_size, - args.data_parallel_size, # args.data_parallel_size, - ) - world_size = torch.distributed.get_world_size() - - print(args.tensor_model_parallel_size, "MODEL PARALLEL SIZE") - - parallel_state.initialize_model_parallel( - tensor_model_parallel_size_=args.tensor_model_parallel_size, - pipeline_model_parallel_size_=args.pipeline_model_parallel_size, - default_backend="nccl", - p2p_backend="ucc" if HAS_UCC else "nccl", - ) - - pipeline_model_parallel_size = ( - parallel_state.get_pipeline_model_parallel_world_size() - ) - model_parallel_cuda_manual_seed(0) - model = build_model( - gpt_model_provider, - wrap_with_ddp=parallel_state.get_data_parallel_world_size() > 1, - virtual_pipeline_model_parallel_size=None, - cpu_offload=args.cpu_offload, - ) - assert isinstance(model, list), model - _param_groups = _get_params_for_weight_decay_optimization(model) - optim = torch.optim.Adam(_param_groups) - runtime = train(model, optim, args.pipeline_model_parallel_size, async_comm) - - parallel_state.destroy_model_parallel() - torch.distributed.barrier() - if torch.distributed.get_rank() == 0: - print(TEST_SUCCESS_MESSAGE) - print("Average Iteration Time:", runtime) diff --git a/tests/L0/run_transformer/test_bert_minimal.py b/tests/L0/run_transformer/test_bert_minimal.py new file mode 100644 index 000000000..887f7e259 --- /dev/null +++ b/tests/L0/run_transformer/test_bert_minimal.py @@ -0,0 +1,250 @@ +import torch +import unittest +from apex.transformer.testing import global_vars +from apex.transformer.testing.standalone_bert import bert_model_provider +from apex.transformer.pipeline_parallel.schedules.common import ( + _get_params_for_weight_decay_optimization, build_model +) +from apex.transformer.pipeline_parallel.schedules import get_forward_backward_func +from apex.transformer.pipeline_parallel.utils import ( + average_losses_across_data_parallel_group, unwrap_model, setup_microbatch_calculator +) +from apex.transformer.log_util import set_logging_level +from apex.transformer import tensor_parallel, parallel_state +from apex.transformer.enums import ModelType +from apex.transformer._ucc_util import HAS_UCC +from apex.transformer.testing.distributed_test_base import UccDistributedTestBase, NcclDistributedTestBase +import logging + +from torch.testing._internal import common_utils + +logging.getLogger("torch").setLevel(logging.WARNING) + + +logging.getLogger("apex").setLevel(logging.WARNING) + + +set_logging_level("WARNING") + + +class BertTestBase: + + def _download_fancy_data(self): + text = """ + An original sentence not subject to any license restrictions, copyright, or royalty payments. Nothing to see here. Commercial or non-commercial use. Research or non-research purposes. The quick brown fox jumps over the lazy dog. Lorem ipsum. + """ + text = text * 1024 + encoded = text.encode("ascii", "replace") + ints = [int(encoded[i]) for i in range(len(encoded))] + return torch.tensor(ints) + + # build a batch given sequence_len and batch size + def _generate_fancy_data_labels(self, sequence_len, batch_size): + temps = [] + for i in range(batch_size): + if self.inds is None or self.data_idx >= len(self.inds): + # hack as use of RNG will fall out of sync due to pipelines being different + torch.manual_seed(self.MANUAL_SEED) + self.inds = torch.randperm( + self.effective_length, device="cuda") + self.masks = ( + torch.rand( + len(self.inds) // batch_size + 1, batch_size, sequence_len, device="cuda" + ) + >= self.MASK_PROB + ).long() + self.MANUAL_SEED += 1 + self.data_idx = 0 + if self.rank == 0: + print("new epoch", len(self.inds)) + print("my start", self.inds[0:5]) + print("masks_checksum:", torch.sum(self.masks)) + if self.EASY_MODE: + data_idx_ = self.data_idx % self.EASY_MODE_SIZ + else: + data_idx_ = self.data_idx + offset = self.inds[data_idx_] # * SEQUENCE_LEN + self.data_idx += 1 + + curr = self.fancy_data[offset: offset + + sequence_len].clone().detach() + temps.append(curr) + temp = torch.stack(temps, dim=0).cuda() + mask = self.masks[self.data_idx // batch_size] + mask_not = torch.logical_not(mask).long() + data = mask * temp + mask_not * 124 + label = temp + if parallel_state.get_tensor_model_parallel_rank() == 0: + data_dict = {"text": data, "label": label, "mask_not": mask_not} + else: + data_dict = None + keys = ["text", "label", "mask_not"] + broadcasted_data = tensor_parallel.broadcast_data( + keys, data_dict, torch.long) + return ( + broadcasted_data["text"].long(), + broadcasted_data["label"].long(), + broadcasted_data["mask_not"], + ) + + def _fwd_step_func(self, batch, model): + data, label, loss_mask = batch + y = model(data, torch.ones_like(data), lm_labels=label) + + def loss_func(output_tensor): + output_tensor, _ = output_tensor + lm_loss_ = output_tensor.float() + lm_loss = torch.sum(lm_loss_.view(-1) * + loss_mask.reshape(-1)) / loss_mask.sum() + averaged_loss = average_losses_across_data_parallel_group([ + lm_loss]) + if self.data_idx >= 1536: + # NOTE (patwang): Loss cutoff might be excessively high but roughly one in five + # unlucky random seeds do cause loss to spike to just under 8.0 + self.assertLess(averaged_loss, 8.0) + return lm_loss, {"avg": averaged_loss} + + return y, loss_func + + def _train( + self, model, optim, virtual_pipeline_model_parallel_size, pipeline_model_parallel_size, async_comm + ): + args = global_vars.get_args() + sequence_len = args.seq_length + micro_batch_size = args.micro_batch_size + hidden_size = args.hidden_size + global_batch_size = args.global_batch_size + forward_backward_func = get_forward_backward_func( + virtual_pipeline_model_parallel_size, pipeline_model_parallel_size + ) + tensor_shape = (sequence_len, micro_batch_size, hidden_size) + for _ in range(16): + batch = self._generate_fancy_data_labels( + sequence_len, global_batch_size) + optim.zero_grad() + forward_backward_func( + self._fwd_step_func, + batch, + model, + forward_only=False, + tensor_shape=tensor_shape, + async_comm=async_comm, + sequence_parallel_enabled=args.sequence_parallel, + ) + # All-reduce layernorm parameters across model parallel nodes + # when sequence parallelism is used + if parallel_state.get_tensor_model_parallel_world_size() > 1 and args.sequence_parallel: + for model_module in model: + unwrapped_model = unwrap_model(model_module) + for param in unwrapped_model.parameters(): + if getattr(param, 'sequence_parallel_enabled', False): + grad = param.grad + torch.distributed.all_reduce( + grad, group=parallel_state.get_tensor_model_parallel_group()) + + optim.step() + + @unittest.skipUnless(torch.cuda.device_count() > 2, "requires at least 3 gpus") + def test_bert_without_interleaving(self): + self._test_bert(virtual_pipeline_model_parallel_size=None) + + @unittest.skipUnless(torch.cuda.device_count() > 2, "requires at least 3 gpus") + def test_bert_with_interleaving(self): + if self.DISTRIBUTED_BACKEND == 'ucc': + self.skipTest('skip interleaving with ucc') + self._test_bert(virtual_pipeline_model_parallel_size=2) + + def _test_bert(self, virtual_pipeline_model_parallel_size): + + self.MANUAL_SEED = 42 + self.inds = None + self.masks = None + self.data_idx = 0 + self.MASK_PROB = 0.1 + self.EASY_MODE = False + self.EASY_MODE_SIZ = 32 + + num_devices = torch.cuda.device_count() + tensor_model_parallel_size = 2 if num_devices % 2 == 0 and num_devices > 4 else 1 + pipeline_model_parallel_size = num_devices // tensor_model_parallel_size + + override_args = { + "micro_batch_size": 2, + "num_layers": 16, + "hidden_size": 256, + "num_attention_heads": 8, + "max_position_embeddings": 512, + "seq_length": 512, + "global_batch_size": 128, + "pipeline_model_parallel_size": pipeline_model_parallel_size, + "tensor_model_parallel_size": tensor_model_parallel_size, + "bert_binary_head": False, + "world_size": self.world_size, + "rank": self.rank, + } + + global_vars.set_global_variables(override_args=override_args, ignore_unknown_args=True) + args = global_vars.get_args() + + self.fancy_data = self._download_fancy_data() + self.effective_length = self.fancy_data.size(0) // args.seq_length + self.effective_length = self.fancy_data.size(0) - args.seq_length + + if self.rank == 0: + print( + f'testing backend: {self.DISTRIBUTED_BACKEND} with virtual_pipeline_model_parallel_size: {virtual_pipeline_model_parallel_size}') + async_comm = not args.sequence_parallel and virtual_pipeline_model_parallel_size is None + self.data_idx = 0 + args.padded_vocab_size = 128 # needed in standalone gpt + args.model_type = ModelType.encoder_or_decoder + setup_microbatch_calculator( + args.rank, + args.rampup_batch_size, + args.global_batch_size, + args.micro_batch_size, + args.data_parallel_size, + ) + parallel_state.initialize_model_parallel( + args.tensor_model_parallel_size, + args.pipeline_model_parallel_size, + virtual_pipeline_model_parallel_size, + default_backend="nccl", + p2p_backend=self.DISTRIBUTED_BACKEND, + ) + + tensor_parallel.random.model_parallel_cuda_manual_seed(0) + model = build_model( + bert_model_provider, + wrap_with_ddp=parallel_state.get_data_parallel_world_size() > 1, + virtual_pipeline_model_parallel_size=virtual_pipeline_model_parallel_size, + cpu_offload=args.cpu_offload, + ) + assert isinstance(model, list) + assert len(model) == ( + 1 + if virtual_pipeline_model_parallel_size is None + else virtual_pipeline_model_parallel_size + ) + _param_groups = _get_params_for_weight_decay_optimization(model) + optim = torch.optim.Adam(_param_groups) + self._train( + model, + optim, + virtual_pipeline_model_parallel_size, + args.pipeline_model_parallel_size, + async_comm, + ) + torch.distributed.barrier() + + +class NcclBertTest(BertTestBase, NcclDistributedTestBase): + pass + + +@unittest.skipUnless(HAS_UCC, "requires pytorch to be built with native ucc") +class UccBertTest(BertTestBase, UccDistributedTestBase): + pass + + +if __name__ == "__main__": + common_utils.run_tests() diff --git a/tests/L0/run_transformer/run_dynamic_batchsize_test.py b/tests/L0/run_transformer/test_dynamic_batchsize.py similarity index 54% rename from tests/L0/run_transformer/run_dynamic_batchsize_test.py rename to tests/L0/run_transformer/test_dynamic_batchsize.py index b2c020a41..327aa4d83 100644 --- a/tests/L0/run_transformer/run_dynamic_batchsize_test.py +++ b/tests/L0/run_transformer/test_dynamic_batchsize.py @@ -1,45 +1,38 @@ from typing import Tuple, List import torch +import unittest from apex.transformer import parallel_state from apex.transformer.pipeline_parallel.utils import get_num_microbatches from apex.transformer.pipeline_parallel.schedules.common import ( - _get_params_for_weight_decay_optimization, + _get_params_for_weight_decay_optimization, build_model ) -from apex.transformer.pipeline_parallel.schedules.common import build_model from apex.transformer.pipeline_parallel.schedules.fwd_bwd_pipelining_with_interleaving import ( _forward_backward_pipelining_with_interleaving, ) -from apex.transformer.pipeline_parallel.utils import setup_microbatch_calculator -from apex.transformer.pipeline_parallel.utils import _reconfigure_microbatch_calculator -from apex.transformer.pipeline_parallel.utils import update_num_microbatches +from apex.transformer.pipeline_parallel.utils import ( + setup_microbatch_calculator, _reconfigure_microbatch_calculator, update_num_microbatches +) from apex.transformer.testing import global_vars -from apex.transformer.testing.commons import TEST_SUCCESS_MESSAGE -from apex.transformer.testing.commons import initialize_distributed -from apex.transformer.testing.commons import print_separator -from apex.transformer.testing.commons import fwd_step_func -from apex.transformer.log_util import get_transformer_logger, set_logging_level -from apex.transformer.testing.commons import model_provider_func -from apex.transformer._data import MegatronPretrainingRandomSampler -from apex.transformer._data import MegatronPretrainingSampler +from apex.transformer.testing.commons import ( + print_separator, fwd_step_func, model_provider_func +) +from apex.transformer.log_util import get_transformer_logger +from apex.transformer._data import MegatronPretrainingRandomSampler, MegatronPretrainingSampler +from apex.transformer.testing.distributed_test_base import NcclDistributedTestBase +from torch.testing._internal import common_utils # note(mkozuki): To see warmup, steady, cooldown iterations, uncomment the line below # set_logging_level("INFO") _logger = get_transformer_logger("pipeline_parallel_test") # note(mkozuki): To see if local batch size increases, uncomment the line below # _logger.setLevel("INFO") -global_vars.set_global_variables( - args_defaults={"global_batch_size": 512, "rampup_batch_size": [64, 64, 1000],}, - ignore_unknown_args=True, -) -RAMPUP_BATCH_SIZE = [] NUM_ITERATIONS = 20 NUM_SAMPLES = 16384 // 2 -batch_size, micro_batch_size = None, None HIDDEN_SIZE = 16 @@ -88,9 +81,10 @@ def run_interleaved_with_dynamic_batch_size( ) assert isinstance(model, list) assert len(model) == virtual_pipeline_model_parallel_size - optimizer = torch.optim.Adam(_get_params_for_weight_decay_optimization(model)) + optimizer = torch.optim.Adam( + _get_params_for_weight_decay_optimization(model)) - initial_local_minibatch_size = get_num_microbatches() * micro_batch_size + initial_local_minibatch_size = get_num_microbatches() * args.micro_batch_size dataset = Dataset(NUM_SAMPLES) data_loader = torch.utils.data.DataLoader( dataset, @@ -110,11 +104,11 @@ def get_num_samples(batch): assert isinstance(batch, (list, tuple)) return [get_num_samples(b) for b in batch] - tensor_shape = [micro_batch_size, HIDDEN_SIZE, HIDDEN_SIZE] + tensor_shape = [args.micro_batch_size, HIDDEN_SIZE, HIDDEN_SIZE] consumed_samples = 0 for i in range(NUM_ITERATIONS): update_num_microbatches(consumed_samples, consistency_check=False) - local_batch_size = get_num_microbatches() * micro_batch_size + local_batch_size = get_num_microbatches() * args.micro_batch_size data_iter._index_sampler.local_minibatch_size = local_batch_size local_mini_batch = next(data_iter) @@ -134,7 +128,7 @@ def get_num_samples(batch): consumed_samples += ( parallel_state.get_data_parallel_world_size() * get_num_microbatches() - * micro_batch_size + * args.micro_batch_size ) if not forward_only: @@ -146,57 +140,84 @@ def get_num_samples(batch): optimizer.zero_grad(set_to_none=True) torch.distributed.barrier() - if torch.distributed.get_rank() == 0: - print(TEST_SUCCESS_MESSAGE) + + +class DynamicBatchsizeTestBase: + @unittest.skipUnless(torch.cuda.device_count() > 2, "requires at least 3 gpus") + def test_dynamic_batchsize(self): + + n_tests = 0 + failures = [] + + override_args = { + "micro_batch_size": 2, + "num_layers": 16, + "hidden_size": 256, + "num_attention_heads": 8, + "max_position_embeddings": 512, + "seq_length": 512, + "global_batch_size": 128, + "use_cpu_initialization": True, + "world_size": self.world_size, + "rank": self.rank, + } + + global_vars.set_global_variables( + args_defaults={"global_batch_size": 512, + "rampup_batch_size": [64, 64, 1000], }, + ignore_unknown_args=True, + override_args=override_args, + ) + + args = global_vars.get_args() + + setup_microbatch_calculator( + args.rank, + args.rampup_batch_size, + args.global_batch_size, + args.micro_batch_size, + 1, # args.data_parallel_size, + ) + for BatchSamplerCls in ( + MegatronPretrainingSampler, + MegatronPretrainingRandomSampler, + ): + for forward_only in (False, True): + n_tests += 1 + pipeline_model_parallel_size = self.world_size + try: + run_interleaved_with_dynamic_batch_size( + pipeline_model_parallel_size, forward_only, BatchSamplerCls, + ) + except Exception as e: + msg = ( + f"\tforward_only: {forward_only}\n" + f"pipeline rank: {parallel_state.get_pipeline_model_parallel_rank()}, " + f"virtual pipeline rank: {parallel_state.get_virtual_pipeline_model_parallel_rank()}\n" + f"{str(e)}" + ) + raise RuntimeError(msg) + finally: + parallel_state.destroy_model_parallel() + print_separator("TEST RESULT") + if failures: + torch.distributed.barrier() + if torch.distributed.get_rank() == 0: + print("\n".join(failures)) + msg = f"{len(failures)} / {n_tests} cases failed" + raise RuntimeError(msg) + else: + torch.distributed.barrier() + if torch.distributed.get_rank() == 0: + print("### PASS!") + + +class NcclDynamicBatchsizeTest(DynamicBatchsizeTestBase, NcclDistributedTestBase): + pass + +# TODO: (Fuzzkatt) UCC still doesn't work with fwd_bwd_pipelining_with_interleaving if __name__ == "__main__": torch.backends.cuda.matmul.allow_tf32 = False - torch.backends.cudnn.allow_tf32 = False - n_tests = 0 - failures = [] - - initialize_distributed() - world_size = torch.distributed.get_world_size() - args = global_vars.get_args() - batch_size = args.global_batch_size - micro_batch_size = args.micro_batch_size - setup_microbatch_calculator( - args.rank, - args.rampup_batch_size, - args.global_batch_size, - args.micro_batch_size, - 1, # args.data_parallel_size, - ) - for BatchSamplerCls in ( - MegatronPretrainingSampler, - MegatronPretrainingRandomSampler, - ): - for forward_only in (False, True): - n_tests += 1 - pipeline_model_parallel_size = world_size - try: - run_interleaved_with_dynamic_batch_size( - pipeline_model_parallel_size, forward_only, BatchSamplerCls, - ) - except Exception as e: - msg = ( - f"\tforward_only: {forward_only}\n" - f"pipeline rank: {parallel_state.get_pipeline_model_parallel_rank()}, " - f"virtual pipeline rank: {parallel_state.get_virtual_pipeline_model_parallel_rank()}\n" - f"{str(e)}" - ) - raise RuntimeError(msg) - finally: - parallel_state.destroy_model_parallel() - print_separator("TEST RESULT") - if failures: - torch.distributed.barrier() - if torch.distributed.get_rank() == 0: - print("\n".join(failures)) - msg = f"{len(failures)} / {n_tests} cases failed" - raise RuntimeError(msg) - else: - torch.distributed.barrier() - if torch.distributed.get_rank() == 0: - print("### PASS!") + common_utils.run_tests() diff --git a/tests/L0/run_transformer/test_gpt_minimal.py b/tests/L0/run_transformer/test_gpt_minimal.py new file mode 100644 index 000000000..4f9b5e5eb --- /dev/null +++ b/tests/L0/run_transformer/test_gpt_minimal.py @@ -0,0 +1,232 @@ +from functools import partial +from typing import List +import time + +import torch + +import unittest + +from apex.transformer._ucc_util import HAS_UCC +from apex.transformer import parallel_state +from apex.transformer.enums import ModelType +from apex.transformer.tensor_parallel import model_parallel_cuda_manual_seed +from apex.transformer.pipeline_parallel.utils import ( + average_losses_across_data_parallel_group, unwrap_model, setup_microbatch_calculator, + get_ltor_masks_and_position_ids +) +from apex.transformer.pipeline_parallel.schedules.common import ( + _get_params_for_weight_decay_optimization, build_model +) +from apex.transformer.pipeline_parallel.schedules.fwd_bwd_pipelining_without_interleaving import ( + forward_backward_pipelining_without_interleaving, +) +from apex.transformer.testing.standalone_gpt import gpt_model_provider +from apex.transformer.testing import global_vars + +from apex.transformer.testing.distributed_test_base import UccDistributedTestBase, NcclDistributedTestBase + +from torch.testing._internal import common_utils +from torch.testing._internal.common_device_type import instantiate_device_type_tests + + +class GptTestBase: + + def _download_fancy_data(self): + text = """ + An original sentence not subject to any license restrictions, copyright, or royalty payments. Nothing to see here. Commercial or non-commercial use. Research or non-research purposes. The quick brown fox jumps over the lazy dog. Lorem ipsum. + """ + text = text * 1024 + encoded = text.encode("ascii", "replace") + ints = [int(encoded[i]) for i in range(len(encoded))] + return torch.tensor(ints) + + # build a batch given sequence_len and batch size + def _generate_fancy_data_labels(self, sequence_len, batch_size): + temps = list() + for i in range(batch_size): + if self.inds is None or self.data_idx >= len(self.inds): + # hack as use of RNG will fall out of sync due to pipelines being different + model_parallel_cuda_manual_seed(self.MANUAL_SEED) + self.inds = torch.randperm(effective_length, device="cuda") + self.MANUAL_SEED += 1 + self.data_idx = 0 + data_idx_ = self.data_idx + offset = self.inds[data_idx_] + self.data_idx += 1 + curr = fancy_data[offset: offset + + sequence_len + 1].clone().detach() + temps.append(curr) + temp = torch.stack(temps, dim=0).cuda() + return temp + + def _get_batch(self, int_tensors: List[torch.Tensor]): + data = int_tensors[0] + # Unpack. + tokens_ = data.long() + labels = tokens_[:, 1:].contiguous() + tokens = tokens_[:, :-1].contiguous() + # Get the masks and position ids. + attention_mask, loss_mask, position_ids = get_ltor_masks_and_position_ids( + tokens, + self.N_VOCAB, # tokenizer.eod, + False, # args.reset_position_ids, + False, # args.reset_attention_mask, + False, # args.eod_mask_loss, + ) + return tokens, labels, loss_mask, attention_mask, position_ids + + # Ref: https://github.com/NVIDIA/Megatron-LM/blob/b31e1296354e979722627a6c4dedafe19b51fa97/pretrain_gpt.py#L75 + def _loss_func(self, loss_mask, output_tensor): + losses = output_tensor.float() + loss_mask = loss_mask.view(-1).float() + loss = torch.sum(losses.view(-1) * loss_mask) / loss_mask.sum() + + # Reduce loss for logging. + averaged_loss = average_losses_across_data_parallel_group([loss]) + + return loss, {"lm loss": averaged_loss[0]} + + # Ref: https://github.com/NVIDIA/Megatron-LM/blob/b31e1296354e979722627a6c4dedafe19b51fa97/pretrain_gpt.py#L86 + def _fwd_step_func(self, batch, model): + """Forward step.""" + tokens, labels, loss_mask, attention_mask, position_ids = self._get_batch( + batch) + output_tensor = model(tokens, position_ids, + attention_mask, labels=labels) + return output_tensor, partial(self._loss_func, loss_mask) + + def _train(self, model, optim, pipeline_model_parallel_size, async_comm): + args = global_vars.get_args() + fwd_bwd_func = forward_backward_pipelining_without_interleaving + + tensor_shape = (args.seq_length, args.micro_batch_size, + args.hidden_size) + runtime = 0 + # training loop + for i in range(3): + since = time.time() + if torch.distributed.get_rank() == 0: + print("begin iter", i) + batch = [ + self._generate_fancy_data_labels( + args.seq_length, args.global_batch_size) + for _ in range(pipeline_model_parallel_size) + ] + if torch.distributed.get_rank() == 0: + print("finished making batch...") + optim.zero_grad() + fwd_bwd_func( + self._fwd_step_func, + batch, + model, + forward_only=False, + tensor_shape=tensor_shape, + async_comm=async_comm, + sequence_parallel_enabled=args.sequence_parallel, + ) + if torch.distributed.get_rank() == 0: + print("finished forward step") + # All-reduce layernorm parameters across model parallel nodes + # when sequence parallelism is used + if parallel_state.get_tensor_model_parallel_world_size() > 1 and global_vars.get_args().sequence_parallel: + for model_module in model: + unwrapped_model = unwrap_model(model_module) + for param in unwrapped_model.parameters(): + if getattr(param, 'sequence_parallel_enabled', False): + grad = param.grad + torch.distributed.all_reduce( + grad, group=parallel_state.get_tensor_model_parallel_group()) + optim.step() + if torch.distributed.get_rank() == 0: + print("finished iter", i) + runtime += time.time() - since + return runtime / 3.0 + + @unittest.skipUnless(torch.cuda.device_count() > 2, "requires at least 3 gpus") + def test_gpt(self): + self.MANUAL_SEED = 42 + self.inds = None + self.data_idx = 0 + self.N_VOCAB = 128 + init = True + + num_devices = torch.cuda.device_count() + tensor_model_parallel_size = 2 if num_devices % 2 == 0 and num_devices >= 4 else 1 + pipeline_model_parallel_size = num_devices // tensor_model_parallel_size + + override_args = { + "micro_batch_size": 2, + "num_layers": 16, + "hidden_size": 256, + "num_attention_heads": 8, + "max_position_embeddings": 512, + "seq_length": 512, + "global_batch_size": 128, + "pipeline_model_parallel_size": pipeline_model_parallel_size, + "tensor_model_parallel_size": tensor_model_parallel_size, + "world_size": self.world_size, + "rank": self.rank, + } + + global_vars.set_global_variables(override_args=override_args, ignore_unknown_args=True) + args = global_vars.get_args() + + for async_comm in (False,) if args.sequence_parallel else (False, True): + global fancy_data + global effective_length + + if init: + init = False + + fancy_data = self._download_fancy_data() + args = global_vars.get_args() + args.model_type = ModelType.encoder_or_decoder + effective_length = fancy_data.size(0) // args.seq_length + effective_length = fancy_data.size(0) - args.seq_length + + args.padded_vocab_size = 128 + setup_microbatch_calculator( + args.rank, + args.rampup_batch_size, + args.global_batch_size, + args.micro_batch_size, + args.data_parallel_size, + ) + + print(args.tensor_model_parallel_size, "MODEL PARALLEL SIZE") + + parallel_state.initialize_model_parallel( + tensor_model_parallel_size_=args.tensor_model_parallel_size, + pipeline_model_parallel_size_=args.pipeline_model_parallel_size, + default_backend="nccl", + p2p_backend=self.DISTRIBUTED_BACKEND, + ) + + model_parallel_cuda_manual_seed(0) + model = build_model( + gpt_model_provider, + wrap_with_ddp=parallel_state.get_data_parallel_world_size() > 1, + virtual_pipeline_model_parallel_size=None, + cpu_offload=args.cpu_offload, + ) + assert isinstance(model, list), model + _param_groups = _get_params_for_weight_decay_optimization(model) + optim = torch.optim.Adam(_param_groups) + runtime = self._train( + model, optim, args.pipeline_model_parallel_size, async_comm) + + parallel_state.destroy_model_parallel() + torch.distributed.barrier() + + +class NcclGptTest(GptTestBase, NcclDistributedTestBase): + pass + + +@unittest.skipUnless(HAS_UCC, "requires pytorch to be built with native ucc") +class UccGptTest(GptTestBase, UccDistributedTestBase): + pass + + +if __name__ == "__main__": + common_utils.run_tests() diff --git a/tests/L0/run_transformer/test_transformer_module.py b/tests/L0/run_transformer/test_transformer_module.py deleted file mode 100644 index 57f454a67..000000000 --- a/tests/L0/run_transformer/test_transformer_module.py +++ /dev/null @@ -1,97 +0,0 @@ -from typing import Tuple -import os -import subprocess -import sys -import unittest - - -# TODO(crcrpar): should move this to `apex._testing` or whatever in the near future. -RUN_SLOW_TESTS = os.getenv('APEX_RUN_WITH_SLOW_TESTS', '0') == '1' -SEVERALGPU_TEST = [ - "bert_minimal_test", - "gpt_minimal_test", - "dynamic_batchsize_test", -] - - -def get_multigpu_launch_option(min_gpu): - should_skip = False - import torch - - num_devices = torch.cuda.device_count() - if num_devices < min_gpu: - should_skip = True - distributed_run_options = f"-m torch.distributed.run --nproc_per_node={num_devices}" - return should_skip, distributed_run_options - - -def get_launch_option(test_filename) -> Tuple[bool, str]: - should_skip = False - for severalgpu_test in SEVERALGPU_TEST: - if severalgpu_test in test_filename: - return get_multigpu_launch_option(3) - return should_skip, "" - - -def get_test_command(test_file: str) -> str: - python_executable_path = sys.executable - should_skip, launch_option = get_launch_option(test_file) - if should_skip: - return "" - test_run_cmd = ( - f"{python_executable_path} {launch_option} {test_file} " - "--micro-batch-size 2 --num-layers 16 --hidden-size 256 --num-attention-heads 8 --max-position-embeddings " - "512 --seq-length 512 --global-batch-size 128" - ) - if "bert" in test_file or "gpt" in test_file: - import torch - - num_devices = torch.cuda.device_count() - if "bert" in test_file: - # "bert" uses the interleaving. - tensor_model_parallel_size = 2 if num_devices % 2 == 0 and num_devices > 4 else 1 - if "gpt" in test_file: - # "gpt" uses the non-interleaving. - tensor_model_parallel_size = 2 if num_devices % 2 == 0 and num_devices >= 4 else 1 - pipeline_model_parallel_size = num_devices // tensor_model_parallel_size - test_run_cmd += f" --pipeline-model-parallel-size {pipeline_model_parallel_size} --tensor-model-parallel-size {tensor_model_parallel_size}" - - if "bert" in test_file: - test_run_cmd += f" --bert-no-binary-head" - else: - test_run_cmd += f" --use-cpu-initialization" - return test_run_cmd - - -def _get_test_file(key): - directory = os.path.dirname(__file__) - test_file = [ - os.path.join(directory, f) - for f in os.listdir(directory) - if f.startswith("run_") and os.path.isfile(os.path.join(directory, f)) and key in f - ][0] - return test_file - - -@unittest.skipUnless(RUN_SLOW_TESTS, "this test takes long.") -class TestTransformer(unittest.TestCase): - def _test_impl(self, key: str): - test_file = _get_test_file(key) - command = get_test_command(test_file) - if command: - subprocess.run(command, shell=True, check=True) - else: - self.skipTest("Appropriate command is not generated") - - def test_standalone_bert(self): - self._test_impl("bert") - - def test_standalone_gpt(self): - self._test_impl("gpt") - - def test_dynamic_batch_size(self): - self._test_impl("dynamic_batchsize") - - -if __name__ == "__main__": - unittest.main() From 80a49541aea588a808af8ba373d34f0466088e07 Mon Sep 17 00:00:00 2001 From: Masaki Kozuki Date: Thu, 27 Oct 2022 14:26:50 -0700 Subject: [PATCH 21/25] check `is_ucc_available` (#1523) Signed-off-by: Masaki Kozuki --- apex/transformer/_ucc_util.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/apex/transformer/_ucc_util.py b/apex/transformer/_ucc_util.py index c070c84d1..a286e2ab1 100644 --- a/apex/transformer/_ucc_util.py +++ b/apex/transformer/_ucc_util.py @@ -1,6 +1,6 @@ from torch import distributed as dist -HAS_UCC = dist.is_ucc_available() +HAS_UCC = hasattr(dist, "is_ucc_available") and dist.is_ucc_available() if not HAS_UCC: try: import torch_ucc From 666e769b58b75f14190ed9ff8190d78f8af89131 Mon Sep 17 00:00:00 2001 From: hubertlu-tw Date: Fri, 30 Dec 2022 00:42:42 +0000 Subject: [PATCH 22/25] Support fused_weight_gradient_mlp_cuda for ROCm --- ...d_weight_gradient_dense_16bit_prec_cuda.cu | 58 +++++++++++++ .../fused_weight_gradient_dense_cuda.cu | 87 +++++++++++++++++++ setup.py | 21 +++-- 3 files changed, 155 insertions(+), 11 deletions(-) diff --git a/csrc/megatron/fused_weight_gradient_dense_16bit_prec_cuda.cu b/csrc/megatron/fused_weight_gradient_dense_16bit_prec_cuda.cu index 60d1e8d1f..878a08075 100644 --- a/csrc/megatron/fused_weight_gradient_dense_16bit_prec_cuda.cu +++ b/csrc/megatron/fused_weight_gradient_dense_16bit_prec_cuda.cu @@ -30,6 +30,34 @@ void gemmex_wrapper_fp16( const float* beta, at::BFloat16* C, int ldc) { +#ifdef __HIP_PLATFORM_HCC__ + rocblas_int flags = 0; + TORCH_CUDABLAS_CHECK(rocblas_gemm_ex( + handle, + transa, + transb, + m, + n, + k, + alpha, + A, + rocblas_datatype_f16_r, + lda, + B, + rocblas_datatype_f16_r, + ldb, + beta, + C, + rocblas_datatype_f16_r, + ldc, + C, + rocblas_datatype_f16_r, + ldc, + rocblas_datatype_f32_r, + rocblas_gemm_algo_standard /*algo*/, + 0 /*solution_index*/, + flags)); +#else TORCH_CUDABLAS_CHECK(cublasGemmEx( handle, transa, @@ -50,6 +78,7 @@ void gemmex_wrapper_fp16( ldc, CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP)); +#endif } // FP16 inputs and FP16 accumulation @@ -68,6 +97,34 @@ void gemmex_wrapper_fp16( const float* beta, at::Half* C, int ldc) { +#ifdef __HIP_PLATFORM_HCC__ + rocblas_int flags = 0; + TORCH_CUDABLAS_CHECK(rocblas_gemm_ex( + handle, + transa, + transb, + m, + n, + k, + alpha, + A, + rocblas_datatype_f16_r, + lda, + B, + rocblas_datatype_f16_r, + ldb, + beta, + C, + rocblas_datatype_f16_r, + ldc, + C, + rocblas_datatype_f16_r, + ldc, + rocblas_datatype_f32_r, + rocblas_gemm_algo_standard /*algo*/, + 0 /*solution_index*/, + flags)); +#else TORCH_CUDABLAS_CHECK(cublasGemmEx( handle, transa, @@ -88,6 +145,7 @@ void gemmex_wrapper_fp16( ldc, CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP)); +#endif } template diff --git a/csrc/megatron/fused_weight_gradient_dense_cuda.cu b/csrc/megatron/fused_weight_gradient_dense_cuda.cu index dfaa1345d..5d0b08876 100644 --- a/csrc/megatron/fused_weight_gradient_dense_cuda.cu +++ b/csrc/megatron/fused_weight_gradient_dense_cuda.cu @@ -30,6 +30,34 @@ void gemmex_wrapper( const float* beta, float* C, int ldc) { +#ifdef __HIP_PLATFORM_HCC__ + rocblas_int flags = 0; + TORCH_CUDABLAS_CHECK(rocblas_gemm_ex( + handle, + transa, + transb, + m, + n, + k, + alpha, + A, + rocblas_datatype_f16_r, + lda, + B, + rocblas_datatype_f16_r, + ldb, + beta, + C, + rocblas_datatype_f32_r, + ldc, + C, + rocblas_datatype_f32_r, + ldc, + rocblas_datatype_f32_r, + rocblas_gemm_algo_standard, + 0 /*solution_index*/, + flags)); +#else TORCH_CUDABLAS_CHECK(cublasGemmEx( handle, transa, @@ -50,6 +78,7 @@ void gemmex_wrapper( ldc, CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP)); +#endif } // FP16 Tensor core wrapper around cublas GEMMEx @@ -68,6 +97,34 @@ void gemmex_wrapper( const float* beta, float* C, int ldc) { +#ifdef __HIP_PLATFORM_HCC__ + rocblas_int flags = 0; + TORCH_CUDABLAS_CHECK(rocblas_gemm_ex( + handle, + transa, + transb, + m, + n, + k, + alpha, + A, + rocblas_datatype_f16_r, + lda, + B, + rocblas_datatype_f16_r, + ldb, + beta, + C, + rocblas_datatype_f32_r, + ldc, + C, + rocblas_datatype_f32_r, + ldc, + rocblas_datatype_f32_r, + rocblas_gemm_algo_standard, + 0 /*solution_index*/, + flags)); +#else TORCH_CUDABLAS_CHECK(cublasGemmEx( handle, transa, @@ -88,6 +145,7 @@ void gemmex_wrapper( ldc, CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP)); +#endif } // FP32 wrapper around cublas GEMMEx @@ -106,6 +164,34 @@ void gemmex_wrapper( const float *beta, float *C, int ldc) { +#ifdef __HIP_PLATFORM_HCC__ + rocblas_int flags = 0; + TORCH_CUDABLAS_CHECK(rocblas_gemm_ex( + handle, + transa, + transb, + m, + n, + k, + alpha, + A, + rocblas_datatype_f32_r, + lda, + B, + rocblas_datatype_f32_r, + ldb, + beta, + C, + rocblas_datatype_f32_r, + ldc, + C, + rocblas_datatype_f32_r, + ldc, + rocblas_datatype_f32_r, + rocblas_gemm_algo_standard, + 0 /*solution_index*/, + flags)); +#else TORCH_CUDABLAS_CHECK(cublasGemmEx( handle, transa, @@ -126,6 +212,7 @@ void gemmex_wrapper( ldc, CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP)); +#endif } template diff --git a/setup.py b/setup.py index 5de5acb28..9a4e7b0b2 100644 --- a/setup.py +++ b/setup.py @@ -315,7 +315,6 @@ def check_if_rocm_pytorch(): os.path.join(this_dir, 'csrc/megatron')], extra_compile_args={'cxx': ['-O3'] + version_dependent_macros, 'nvcc':nvcc_args_transformer if not IS_ROCM_PYTORCH else hipcc_args_transformer})) - # TODO (hubertlu): support fused_weight_gradient_mlp_cuda for ROCm if not IS_ROCM_PYTORCH: if bare_metal_version >= Version("11.0"): cc_flag = [] @@ -330,16 +329,16 @@ def check_if_rocm_pytorch(): cc_flag.append("-gencode") cc_flag.append("arch=compute_90,code=sm_90") - ext_modules.append( - CUDAExtension(name='fused_weight_gradient_mlp_cuda', - sources=['csrc/megatron/fused_weight_gradient_dense.cpp', - 'csrc/megatron/fused_weight_gradient_dense_cuda.cu', - 'csrc/megatron/fused_weight_gradient_dense_16bit_prec_cuda.cu'], - include_dirs=[os.path.join(this_dir, 'csrc'), - os.path.join(this_dir, 'csrc/megatron')], - extra_compile_args={'cxx': ['-O3'] + version_dependent_macros, - 'nvcc':append_nvcc_threads(nvcc_args_transformer + ['--use_fast_math'] + cc_flag) - if not IS_ROCM_PYTORCH else hipcc_args_transformer})) + ext_modules.append( + CUDAExtension(name='fused_weight_gradient_mlp_cuda', + sources=['csrc/megatron/fused_weight_gradient_dense.cpp', + 'csrc/megatron/fused_weight_gradient_dense_cuda.cu', + 'csrc/megatron/fused_weight_gradient_dense_16bit_prec_cuda.cu'], + include_dirs=[os.path.join(this_dir, 'csrc'), + os.path.join(this_dir, 'csrc/megatron')], + extra_compile_args={'cxx': ['-O3'] + version_dependent_macros, + 'nvcc':append_nvcc_threads(nvcc_args_transformer + ['--use_fast_math'] + cc_flag) + if not IS_ROCM_PYTORCH else hipcc_args_transformer})) if "--bnp" in sys.argv or "--cuda_ext" in sys.argv: From 52f036d729aa4c20bec8977b3a9467712a55fd70 Mon Sep 17 00:00:00 2001 From: hubertlu-tw Date: Fri, 30 Dec 2022 01:43:50 +0000 Subject: [PATCH 23/25] Fix bugs in run_rocm_extensions.py --- apex/contrib/test/run_rocm_extensions.py | 74 +++++++++++++++++++----- 1 file changed, 60 insertions(+), 14 deletions(-) diff --git a/apex/contrib/test/run_rocm_extensions.py b/apex/contrib/test/run_rocm_extensions.py index 0e8b2a69b..0693cec0d 100644 --- a/apex/contrib/test/run_rocm_extensions.py +++ b/apex/contrib/test/run_rocm_extensions.py @@ -1,26 +1,72 @@ +import argparse +import os import unittest import sys -test_dirs = ["groupbn", "fused_dense", "layer_norm", "multihead_attn", "transducer", "focal_loss", "index_mul_2d", "optimizers", "."] # "." for test_label_smoothing.py -ROCM_BLACKLIST = [ - "layer_norm" +#test_dirs = ["fused_dense", "layer_norm", "multihead_attn", "transducer", "focal_loss", "index_mul_2d", "optimizers", ".", "groupbn"] # "." for test_label_smoothing.py +#ROCM_BLACKLIST = [ +# "layer_norm" +#] + +TEST_ROOT = os.path.dirname(os.path.abspath(__file__)) +TEST_DIRS = [ + "fused_dense", + "layer_norm", # not fully supported on ROCm + "multihead_attn", + "transducer", + "focal_loss", + "index_mul_2d", + "optimizers", + "xentropy", + "clip_grad", + "groupbn", +] + +DEFAULT_TEST_DIRS = [ + "fused_dense", + "multihead_attn", + "transducer", + "focal_loss", + "index_mul_2d", + "optimizers", + "xentropy", + "clip_grad", + "groupbn", ] -runner = unittest.TextTestRunner(verbosity=2) +def parse_args(): + parser = argparse.ArgumentParser( + description="Extension test runner", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument( + "--include", + nargs="+", + choices=TEST_DIRS, + default=DEFAULT_TEST_DIRS, + help="select a set of tests to run (defaults to ALL tests).", + ) + args, _ = parser.parse_known_args() + return args -errcode = 0 +def main(args): + runner = unittest.TextTestRunner(verbosity=2) + errcode = 0 + for test_dir in args.include: + test_dir = os.path.join(TEST_ROOT, test_dir) + print(test_dir) + suite = unittest.TestLoader().discover(test_dir) -for test_dir in test_dirs: - if test_dir in ROCM_BLACKLIST: - continue - suite = unittest.TestLoader().discover(test_dir) + print("\nExecuting tests from " + test_dir) - print("\nExecuting tests from " + test_dir) + result = runner.run(suite) - result = runner.run(suite) + if not result.wasSuccessful(): + errcode = 1 - if not result.wasSuccessful(): - errcode = 1 + sys.exit(errcode) -sys.exit(errcode) +if __name__ == '__main__': + args = parse_args() + main(args) From 483fb507b7301da789e932d4452bea0634b09598 Mon Sep 17 00:00:00 2001 From: hubertlu-tw Date: Fri, 30 Dec 2022 01:44:53 +0000 Subject: [PATCH 24/25] Move test_label_smoothing.py to xentropy folder --- apex/contrib/test/{ => xentropy}/test_label_smoothing.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename apex/contrib/test/{ => xentropy}/test_label_smoothing.py (100%) diff --git a/apex/contrib/test/test_label_smoothing.py b/apex/contrib/test/xentropy/test_label_smoothing.py similarity index 100% rename from apex/contrib/test/test_label_smoothing.py rename to apex/contrib/test/xentropy/test_label_smoothing.py From 72f978c2190f049bb0a2aaf170bb1efc2bc52798 Mon Sep 17 00:00:00 2001 From: hubertlu-tw Date: Fri, 30 Dec 2022 01:49:30 +0000 Subject: [PATCH 25/25] Add some test folders for those ROCm does not support --- apex/contrib/test/run_rocm_extensions.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/apex/contrib/test/run_rocm_extensions.py b/apex/contrib/test/run_rocm_extensions.py index 0693cec0d..ae5e7d01f 100644 --- a/apex/contrib/test/run_rocm_extensions.py +++ b/apex/contrib/test/run_rocm_extensions.py @@ -13,6 +13,10 @@ TEST_DIRS = [ "fused_dense", "layer_norm", # not fully supported on ROCm + "conv_bias_relu",# not fully supported on ROCm + "fmha", # not fully supported on ROCm + #"cudnn_gbn", # not fully supported on ROCm + #"bottleneck", # not fully supported on ROCm "multihead_attn", "transducer", "focal_loss",