diff --git a/.github/workflows/nv-accelerate-v100.yml b/.github/workflows/nv-accelerate-v100.yml index 9a177792597a..9d288458fa19 100644 --- a/.github/workflows/nv-accelerate-v100.yml +++ b/.github/workflows/nv-accelerate-v100.yml @@ -30,7 +30,7 @@ jobs: - name: Install pytorch run: | - pip install torch torchvision --extra-index-url https://download.pytorch.org/whl/cu111 + pip install --no-cache-dir torch torchvision --extra-index-url https://download.pytorch.org/whl/cu111 python -c "import torch; print('torch:', torch.__version__, torch)" python -c "import torch; print('CUDA available:', torch.cuda.is_available())" diff --git a/.github/workflows/nv-inference.yml b/.github/workflows/nv-inference.yml index 16d0ba92605c..5cb220b0f117 100644 --- a/.github/workflows/nv-inference.yml +++ b/.github/workflows/nv-inference.yml @@ -30,7 +30,7 @@ jobs: - name: Install pytorch run: | - pip install torch==1.13.1 torchvision --extra-index-url https://download.pytorch.org/whl/cu116 + pip install --no-cache-dir torch==1.13.1 torchvision --extra-index-url https://download.pytorch.org/whl/cu116 python -c "import torch; print('torch:', torch.__version__, torch)" python -c "import torch; print('CUDA available:', torch.cuda.is_available())" diff --git a/.github/workflows/nv-lightning-v100.yml b/.github/workflows/nv-lightning-v100.yml index e86a307c47aa..0682e479e1f9 100644 --- a/.github/workflows/nv-lightning-v100.yml +++ b/.github/workflows/nv-lightning-v100.yml @@ -30,7 +30,7 @@ jobs: - name: Install pytorch run: | - pip install torch==1.9.1+cu111 torchvision==0.10.1+cu111 torchaudio==0.9.1 -f https://download.pytorch.org/whl/torch_stable.html + pip install --no-cache-dir torch==1.9.1+cu111 torchvision==0.10.1+cu111 torchaudio==0.9.1 -f https://download.pytorch.org/whl/torch_stable.html python -c "import torch; print('torch:', torch.__version__, torch)" python -c "import torch; print('CUDA available:', torch.cuda.is_available())" diff --git a/.github/workflows/nv-megatron.yml b/.github/workflows/nv-megatron.yml index 8037a42f63ed..b9541ee7018a 100644 --- a/.github/workflows/nv-megatron.yml +++ b/.github/workflows/nv-megatron.yml @@ -30,7 +30,7 @@ jobs: - name: Install pytorch run: | - pip install torch==1.13.1 torchvision --extra-index-url https://download.pytorch.org/whl/cu116 + pip install --no-cache-dir torch==1.13.1 torchvision --extra-index-url https://download.pytorch.org/whl/cu116 python -c "import torch; print('torch:', torch.__version__, torch)" python -c "import torch; print('CUDA available:', torch.cuda.is_available())" diff --git a/.github/workflows/nv-mii.yml b/.github/workflows/nv-mii.yml index 08dadf3ef58b..ae8a69f16a7c 100644 --- a/.github/workflows/nv-mii.yml +++ b/.github/workflows/nv-mii.yml @@ -30,7 +30,7 @@ jobs: - name: Install pytorch run: | - pip install torch==1.13.1 torchvision --extra-index-url https://download.pytorch.org/whl/cu116 + pip install --no-cache-dir torch==1.13.1 torchvision --extra-index-url https://download.pytorch.org/whl/cu116 python -c "import torch; print('torch:', torch.__version__, torch)" python -c "import torch; print('CUDA available:', torch.cuda.is_available())" diff --git a/.github/workflows/nv-torch19-p40.yml b/.github/workflows/nv-torch19-p40.yml index 9d21c1506b4d..1e74efca2435 100644 --- a/.github/workflows/nv-torch19-p40.yml +++ b/.github/workflows/nv-torch19-p40.yml @@ -30,7 +30,7 @@ jobs: - name: Install pytorch run: | - pip install torch==1.9.0+cu111 torchvision==0.10.0+cu111 -f https://download.pytorch.org/whl/torch_stable.html + pip install --no-cache-dir torch==1.9.0+cu111 torchvision==0.10.0+cu111 -f https://download.pytorch.org/whl/torch_stable.html python -c "import torch; print('torch:', torch.__version__, torch)" python -c "import torch; print('CUDA available:', torch.cuda.is_available())" diff --git a/.github/workflows/nv-torch19-v100.yml b/.github/workflows/nv-torch19-v100.yml index 562e961e9f3e..9084dab1ad1e 100644 --- a/.github/workflows/nv-torch19-v100.yml +++ b/.github/workflows/nv-torch19-v100.yml @@ -30,7 +30,7 @@ jobs: - name: Install pytorch run: | - pip install torch==1.9.0+cu111 torchvision==0.10.0+cu111 -f https://download.pytorch.org/whl/torch_stable.html + pip install --no-cache-dir torch==1.9.0+cu111 torchvision==0.10.0+cu111 -f https://download.pytorch.org/whl/torch_stable.html python -c "import torch; print('torch:', torch.__version__, torch)" python -c "import torch; print('CUDA available:', torch.cuda.is_available())" diff --git a/.github/workflows/nv-transformers-v100.yml b/.github/workflows/nv-transformers-v100.yml index fd3913ce032d..a6330c2027f5 100644 --- a/.github/workflows/nv-transformers-v100.yml +++ b/.github/workflows/nv-transformers-v100.yml @@ -31,7 +31,7 @@ jobs: - name: Install pytorch run: | # use the same pytorch version as transformers CI - pip install torch torchvision torchaudio -f https://download.pytorch.org/whl/torch_stable.html + pip install --no-cache-dir torch torchvision torchaudio -f https://download.pytorch.org/whl/torch_stable.html python -c "import torch; print('torch:', torch.__version__, torch)" python -c "import torch; print('CUDA available:', torch.cuda.is_available())" diff --git a/csrc/includes/ds_kernel_utils.h b/csrc/includes/ds_kernel_utils.h index 61d424846589..123db37a55a8 100644 --- a/csrc/includes/ds_kernel_utils.h +++ b/csrc/includes/ds_kernel_utils.h @@ -34,7 +34,6 @@ constexpr int hw_warp_size = 32; #if __CUDA_ARCH__ >= 800 #define ASYNC_COPY_AVAILABLE -#define BF16_AVAILABLE #endif // __CUDA_ARCH__ >= 800 #include diff --git a/csrc/transformer/inference/csrc/apply_rotary_pos_emb.cu b/csrc/transformer/inference/csrc/apply_rotary_pos_emb.cu index 3baa88c72f84..d697b08ed609 100644 --- a/csrc/transformer/inference/csrc/apply_rotary_pos_emb.cu +++ b/csrc/transformer/inference/csrc/apply_rotary_pos_emb.cu @@ -4,64 +4,24 @@ // DeepSpeed Team #include "conversion_utils.h" +#include "cooperative_groups.h" +#include "ds_kernel_utils.h" #include "inference_cuda_layers.h" +#include "memory_access_utils.h" #ifndef __HIP_PLATFORM_HCC__ #include #endif namespace cg = cooperative_groups; -namespace cg = cooperative_groups; - -template -__global__ void apply_rotary_pos_emb(T* mixed_query, - T* key_layer, - unsigned rotary_dim, - unsigned seq_len, - unsigned seq_offset, - unsigned num_heads, - unsigned head_size, - unsigned total_count, - int max_out_tokens) -{ - cg::thread_block b = cg::this_thread_block(); - cg::thread_block_tile g = cg::tiled_partition(b); - - int id = threadIdx.x; - int gid = id >> 5; - int lane = id & 0x1f; - unsigned head_id = blockIdx.x * MAX_WARP_NUM + gid; - unsigned offset = head_id * head_size; +namespace rot_half { +constexpr int threads = 256; +} // namespace rot_half - unsigned seq_id = (head_id / num_heads) % seq_len + seq_offset; - unsigned seq_index = head_id % seq_len; - unsigned k_offset = (seq_index + (head_id / seq_len) * max_out_tokens) * head_size; - - if (head_id < total_count) { - while (lane < rotary_dim) { - float inv_freq = (float)((lane / 2) * 2) / (float)rotary_dim; - inv_freq = 1.0 / powf(10000.0, inv_freq) * (float)seq_id; - float q = conversion::to(mixed_query[offset + lane]); - float k = conversion::to(key_layer[k_offset + lane]); - float rotary_sign = (lane % 2 == 1 ? -1.0 : 1.0); - float q_rot = (q * rotary_sign); - float k_rot = (k * rotary_sign); - q_rot = g.shfl_xor(q_rot, 1); - k_rot = g.shfl_xor(k_rot, 1); - q = q * cosf(inv_freq) + q_rot * sinf(inv_freq); - k = k * cosf(inv_freq) + k_rot * sinf(inv_freq); - - mixed_query[offset + lane] = conversion::to(q); - key_layer[k_offset + lane] = conversion::to(k); - - lane += WARP_SIZE; - } - } -} - -__global__ void apply_rotary_pos_emb1(float* mixed_query, - float* key_layer, +template +__global__ void apply_rotary_pos_half(T* mixed_query, + T* key_layer, unsigned rotary_dim, unsigned seq_len, unsigned seq_offset, @@ -70,99 +30,98 @@ __global__ void apply_rotary_pos_emb1(float* mixed_query, unsigned total_count, int max_out_tokens) { - cg::thread_block b = cg::this_thread_block(); - cg::thread_block_tile g = cg::tiled_partition(b); - - int id = threadIdx.x; - int gid = id >> 5; - int lane = id & 0x1f; - - unsigned head_id = blockIdx.x * MAX_WARP_NUM + gid; - unsigned offset = head_id * head_size; - - unsigned seq_id = (head_id / num_heads) % seq_len + seq_offset; - unsigned seq_index = head_id % seq_len; - unsigned k_offset = (seq_index + (head_id / seq_len) * max_out_tokens) * head_size; - - if (head_id < total_count) { - while (lane < rotary_dim) { - float inv_freq = (float)((lane / 2) * 2) / (float)rotary_dim; - inv_freq = 1.0 / powf(10000.0, inv_freq) * (float)seq_id; - float q = mixed_query[offset + lane]; - float k = key_layer[k_offset + lane]; - float rotary_sign = (lane % 2 == 1 ? -1.0 : 1.0); - float q_rot = (q * rotary_sign); - float k_rot = (k * rotary_sign); - q_rot = g.shfl_xor(q_rot, 1); - k_rot = g.shfl_xor(k_rot, 1); - q = q * cosf(inv_freq) + q_rot * sinf(inv_freq); - k = k * cosf(inv_freq) + k_rot * sinf(inv_freq); - - mixed_query[offset + lane] = q; - key_layer[k_offset + lane] = k; - - lane += WARP_SIZE; + constexpr int T_per_thread = granularity / sizeof(T); + constexpr int heads_per_block = rot_half::threads / threadsPerHead; + + cg::thread_block tb = cg::this_thread_block(); + cg::thread_block_tile head_group = cg::tiled_partition(tb); + + const int head_idx = blockIdx.x * heads_per_block + threadIdx.x / threadsPerHead; + const int cur_seq_idx = head_idx % seq_len; + const int offset = head_idx * head_size; + const int k_offset = (cur_seq_idx + (head_idx / seq_len) * max_out_tokens) * head_size; + + const int seq_idx = cur_seq_idx + seq_offset; + const int half_dim = rotary_dim >> 1; + const int half_dim_threads = half_dim / T_per_thread; + + if (head_idx < total_count) { + const int base_neuron_idx = head_group.thread_rank() * T_per_thread; + + T q[T_per_thread], k[T_per_thread]; + mem_access::load_global(q, mixed_query + offset + base_neuron_idx); + mem_access::load_global(k, key_layer + k_offset + base_neuron_idx); + +#pragma unroll + for (int i = 0; i < T_per_thread; i++) { + const int neuron_idx = base_neuron_idx + i; + if (neuron_idx < rotary_dim) { + float inv_freq = (float)((neuron_idx % half_dim) * 2) / (float)rotary_dim; + inv_freq = 1.0 / powf(10000.0, inv_freq) * (float)seq_idx; + + float rotary_sign = (neuron_idx > (half_dim - 1) ? -1.0 : 1.0); + float q_rot = conversion::to(q[i]) * rotary_sign; + float k_rot = conversion::to(k[i]) * rotary_sign; + + const int target_lane = (neuron_idx < half_dim) + ? head_group.thread_rank() + half_dim_threads + : head_group.thread_rank() - half_dim_threads; + + const float q_rot_temp = head_group.shfl(q_rot, target_lane); + const float k_rot_temp = head_group.shfl(k_rot, target_lane); + + q[i] = conversion::to(conversion::to(q[i]) * cosf(inv_freq) + + q_rot_temp * sinf(inv_freq)); + k[i] = conversion::to(conversion::to(k[i]) * cosf(inv_freq) + + k_rot_temp * sinf(inv_freq)); + } } + + mem_access::store_global(mixed_query + offset + base_neuron_idx, q); + mem_access::store_global(key_layer + k_offset + base_neuron_idx, k); } } -template -__global__ void apply_rotary_pos_emb1(T* mixed_query, - T* key_layer, - unsigned rotary_dim, - unsigned seq_len, - unsigned seq_offset, - unsigned num_heads, - unsigned head_size, - unsigned total_count, - int max_out_tokens) -{ - cg::thread_block b = cg::this_thread_block(); - cg::thread_block_tile g = cg::tiled_partition(b); - - int id = threadIdx.x; - int gid = id >> 5; - int lane = id & 0x1f; - - unsigned head_id = blockIdx.x * MAX_WARP_NUM + gid; - unsigned seq_index = head_id % seq_len; - unsigned offset = head_id * head_size; - unsigned k_offset = (seq_index + (head_id / seq_len) * max_out_tokens) * head_size; - - constexpr unsigned mask[32] = { - 0x1 | 0x1000, 0x2 | 0x2000, 0x4 | 0x4000, 0x8 | 0x8000, 0x10 | 0x10000, - 0x20 | 0x20000, 0x40 | 0x40000, 0x80 | 0x80000, 0x100 | 0x100000, 0x200 | 0x200000, - 0x400 | 0x400000, 0x800 | 0x800000, 0x1000 | 0x1, 0x2000 | 0x2, 0x4000 | 0x4, - 0x8000 | 0x8, 0x10000 | 0x10, 0x20000 | 0x20, 0x40000 | 0x40, 0x80000 | 0x80, - 0x100000 | 0x100, 0x200000 | 0x200, 0x400000 | 0x400, 0x800000 | 0x800, 0x1000000, - 0x2000000, 0x4000000, 0x8000000, 0x10000000, 0x20000000, - 0x40000000, 0x80000000}; - - unsigned seq_id = (head_id % seq_len) + seq_offset; - unsigned half_dim = rotary_dim >> 1; - if (head_id < total_count) { - while (lane < rotary_dim) { - float inv_freq = (float)((lane % half_dim) * 2) / (float)rotary_dim; - inv_freq = 1.0 / powf(10000.0, inv_freq) * (float)seq_id; - float q = conversion::to(mixed_query[offset + lane]); - float k = conversion::to(key_layer[k_offset + lane]); - float rotary_sign = (lane > (half_dim - 1) ? -1.0 : 1.0); - float q_rot = (q * rotary_sign); - float k_rot = (k * rotary_sign); - auto q_rot_tmp = lane < half_dim ? __shfl_sync(mask[lane], q_rot, lane + half_dim) - : __shfl_sync(mask[lane], q_rot, lane - half_dim); - auto k_rot_tmp = lane < half_dim ? __shfl_sync(mask[lane], k_rot, lane + half_dim) - : __shfl_sync(mask[lane], k_rot, lane - half_dim); - q = q * cosf(inv_freq) + q_rot_tmp * sinf(inv_freq); - k = k * cosf(inv_freq) + k_rot_tmp * sinf(inv_freq); - - mixed_query[offset + lane] = conversion::to(q); - key_layer[k_offset + lane] = conversion::to(k); - - lane += WARP_SIZE; - } +#define LAUNCH_ROT_POS_EMB_HALF(HEAD_THREADS, ALIGNMENT) \ + apply_rotary_pos_half<<>>(mixed_query, \ + key_layer, \ + rotary_dim, \ + seq_len, \ + offset, \ + num_heads, \ + head_size, \ + total_count, \ + max_out_tokens); + +#ifdef __HIP_PLATFORM_HCC__ +#define LAUNCH_FOR_ALIGNMENT(ALIGNMENT) \ + if (threads_per_head == 4) { \ + LAUNCH_ROT_POS_EMB_HALF(4, ALIGNMENT); \ + } else if (threads_per_head == 8) { \ + LAUNCH_ROT_POS_EMB_HALF(8, ALIGNMENT); \ + } else if (threads_per_head == 16) { \ + LAUNCH_ROT_POS_EMB_HALF(16, ALIGNMENT); \ + } else if (threads_per_head == 32) { \ + LAUNCH_ROT_POS_EMB_HALF(32, ALIGNMENT); \ + } else if (threads_per_head == 64) { \ + LAUNCH_ROT_POS_EMB_HALF(64, ALIGNMENT); \ + } else { \ + assert(false); \ } -} +#else +#define LAUNCH_FOR_ALIGNMENT(ALIGNMENT) \ + if (threads_per_head == 4) { \ + LAUNCH_ROT_POS_EMB_HALF(4, ALIGNMENT); \ + } else if (threads_per_head == 8) { \ + LAUNCH_ROT_POS_EMB_HALF(8, ALIGNMENT); \ + } else if (threads_per_head == 16) { \ + LAUNCH_ROT_POS_EMB_HALF(16, ALIGNMENT); \ + } else if (threads_per_head == 32) { \ + LAUNCH_ROT_POS_EMB_HALF(32, ALIGNMENT); \ + } else { \ + assert(false); \ + } +#endif template void launch_apply_rotary_pos_emb(T* mixed_query, @@ -173,260 +132,52 @@ void launch_apply_rotary_pos_emb(T* mixed_query, unsigned offset, unsigned num_heads, unsigned batch, - bool rotate_half, - bool rotate_every_two, cudaStream_t stream, int max_out_tokens) { - int total_count = batch * num_heads * seq_len; - dim3 block_dims(1024); - dim3 grid_dims((total_count - 1) / MAX_WARP_NUM + 1); // (batch_size); - if (rotate_every_two) - apply_rotary_pos_emb<<>>(mixed_query, - key_layer, - rotary_dim, - seq_len, - offset, - num_heads, - head_size, - total_count, - max_out_tokens); - else if (rotate_half) - apply_rotary_pos_emb1<<>>(mixed_query, - key_layer, - rotary_dim, - seq_len, - offset, - num_heads, - head_size, - total_count, - max_out_tokens); -} - -template void launch_apply_rotary_pos_emb(float*, - float*, - unsigned, - unsigned, - unsigned, - unsigned, - unsigned, - unsigned, - bool, - bool, - cudaStream_t, - int); -#ifdef BF16_AVAILABLE -template void launch_apply_rotary_pos_emb<__nv_bfloat16>(__nv_bfloat16*, - __nv_bfloat16*, - unsigned, - unsigned, - unsigned, - unsigned, - unsigned, - unsigned, - bool, - bool, - cudaStream_t, - int); -#endif -template void launch_apply_rotary_pos_emb<__half>(__half*, - __half*, - unsigned, - unsigned, - unsigned, - unsigned, - unsigned, - unsigned, - bool, - bool, - cudaStream_t, - int); - -template __global__ void apply_rotary_pos_emb(float* mixed_query, - float* key_layer, - unsigned rotary_dim, - unsigned seq_len, - unsigned seq_offset, - unsigned num_heads, - unsigned head_size, - unsigned total_count, - int max_out_tokens); - -#ifdef BF16_AVAILABLE -template __global__ void apply_rotary_pos_emb(__nv_bfloat16* mixed_query, - __nv_bfloat16* key_layer, - unsigned rotary_dim, - unsigned seq_len, - unsigned seq_offset, - unsigned num_heads, - unsigned head_size, - unsigned total_count, - int max_out_tokens); -#endif - -template __global__ void apply_rotary_pos_emb(__half* mixed_query, - __half* key_layer, - unsigned rotary_dim, - unsigned seq_len, - unsigned seq_offset, - unsigned num_heads, - unsigned head_size, - unsigned total_count, - int max_out_tokens); - -#ifdef BF16_AVAILABLE -template __global__ void apply_rotary_pos_emb1(__nv_bfloat16* mixed_query, - __nv_bfloat16* key_layer, - unsigned rotary_dim, - unsigned seq_len, - unsigned seq_offset, - unsigned num_heads, - unsigned head_size, - unsigned total_count, - int max_out_tokens); -#endif - -template __global__ void apply_rotary_pos_emb1(__half* mixed_query, - __half* key_layer, - unsigned rotary_dim, - unsigned seq_len, - unsigned seq_offset, - unsigned num_heads, - unsigned head_size, - unsigned total_count, - int max_out_tokens); -/* -__global__ void apply_rotary_pos_emb(float* mixed_query, -float* key_layer, -unsigned rotary_dim, -unsigned seq_len, -unsigned seq_offset, -unsigned num_heads, -unsigned head_size, -unsigned total_count) -{ -cg::thread_block b = cg::this_thread_block(); -cg::thread_block_tile g = cg::tiled_partition(b); + const int half_dim = rotary_dim >> 1; + + int alignment = sizeof(T); + if (half_dim % (16 / sizeof(T)) == 0) { + alignment = 16; + } else if (half_dim % (8 / sizeof(T)) == 0) { + alignment = 8; + } else if (half_dim % (4 / sizeof(T)) == 0) { + alignment = 4; + } else { + assert(false); + } + const int T_per_elem = alignment / sizeof(T); -int id = threadIdx.x; -int gid = id >> 5; -int lane = id & 0x1f; + int total_count = batch * num_heads * seq_len; -unsigned head_id = blockIdx.x * MAX_WARP_NUM + gid; -unsigned offset = head_id * head_size; + const int padded_head_size = next_pow2(head_size); -unsigned seq_id = (head_id / num_heads) % seq_len + seq_offset; + assert(padded_head_size <= hw_warp_size * T_per_elem); -if (head_id < total_count) { -while (lane < rotary_dim) { -float inv_freq = (float)((lane / 2) * 2) / (float)rotary_dim; -inv_freq = 1.0 / powf(10000.0, inv_freq) * (float)seq_id; -float q = mixed_query[offset + lane]; -float k = key_layer[offset + lane]; -float rotary_sign = (lane % 2 == 1 ? -1.0 : 1.0); -float q_rot = (q * rotary_sign); -float k_rot = (k * rotary_sign); -q_rot = g.shfl_xor(q_rot, 1); -k_rot = g.shfl_xor(k_rot, 1); -q = q * cosf(inv_freq) + q_rot * sinf(inv_freq); -k = k * cosf(inv_freq) + k_rot * sinf(inv_freq); + const int threads_per_head = padded_head_size / T_per_elem; + const int heads_per_block = rot_half::threads / threads_per_head; -mixed_query[offset + lane] = q; -key_layer[offset + lane] = k; + dim3 block(rot_half::threads); + dim3 grid((total_count + heads_per_block - 1) / heads_per_block); -lane += WARP_SIZE; -} -} + if (alignment == 4) { + LAUNCH_FOR_ALIGNMENT(4); + } else if (alignment == 8) { + LAUNCH_FOR_ALIGNMENT(8); + } else if (alignment == 16) { + LAUNCH_FOR_ALIGNMENT(16); + } else { + assert(false); + } } -__global__ void apply_rotary_pos_emb(__half* mixed_query, -__half* key_layer, -unsigned rotary_dim, -unsigned seq_len, -unsigned seq_offset, -unsigned num_heads, -unsigned head_size, -unsigned total_count) -{ -#if __CUDA_ARCH__ >= 700 -cg::thread_block b = cg::this_thread_block(); -cg::thread_block_tile g = cg::tiled_partition(b); - -int id = threadIdx.x; -int gid = id >> 5; -int lane = id & 0x1f; +#define INSTANTIATE_LAUNCH_ROTARY_POS_EMB(T) \ + template void launch_apply_rotary_pos_emb( \ + T*, T*, unsigned, unsigned, unsigned, unsigned, unsigned, unsigned, cudaStream_t, int); -unsigned head_id = blockIdx.x * MAX_WARP_NUM + gid; -unsigned offset = head_id * head_size; -constexpr unsigned mask[32] = {0x1 | 0x1000, 0x2 | 0x2000, 0x4 | 0x4000, 0x8 | 0x8000, -0x10 | 0x10000, 0x20 | 0x20000, 0x40 | 0x40000, 0x80 | 0x80000, -0x100 | 0x100000, 0x200 | 0x200000, 0x400 | 0x400000, 0x800 | 0x800000, -0x1000 | 0x1, 0x2000 | 0x2, 0x4000 | 0x4, 0x8000 | 0x8, -0x10000 | 0x10, 0x20000 | 0x20, 0x40000 | 0x40, 0x80000 | 0x80, -0x100000 | 0x100, 0x200000 | 0x200, 0x400000 | 0x400, 0x800000 | 0x800, -0x1000000, 0x2000000, 0x4000000, 0x8000000, -0x10000000, 0x20000000, 0x40000000, 0x80000000}; -unsigned seq_id = (head_id / num_heads) % seq_len + seq_offset; - -if (head_id < total_count) { -while (lane < rotary_dim) { -//float inv_freq = (float)((lane / 2) * 2) / (float)rotary_dim; -float inv_freq = (float)((lane % (rotary_dim >> 1)) * 2) / (float)rotary_dim; -inv_freq = 1.0 / powf(10000.0, inv_freq) * (float)seq_id; -float q = (float)mixed_query[offset + lane]; -float k = (float)key_layer[offset + lane]; -float rotary_sign = (lane > 11 ? -1.0 : 1.0); -float q_rot = (q * rotary_sign); -float k_rot = (k * rotary_sign); -auto q_rot_tmp = lane < 12 ? __shfl_sync(mask[lane], q_rot, lane + 12) : __shfl_sync(mask[lane], -q_rot, lane - 12);//g.shfl_xor(q_rot, 12); auto k_rot_tmp = lane < 12 ? __shfl_sync(mask[lane], -k_rot, lane + 12) : __shfl_sync(mask[lane], k_rot, lane - 12);//g.shfl_xor(k_rot, 12); q = q * -cosf(inv_freq) + q_rot_tmp * sinf(inv_freq); k = k * cosf(inv_freq) + k_rot_tmp * sinf(inv_freq); - -mixed_query[offset + lane] = (__half)q; -key_layer[offset + lane] = (__half)k; - -lane += WARP_SIZE; -} -} +INSTANTIATE_LAUNCH_ROTARY_POS_EMB(float); +#ifdef BF16_AVAILABLE +INSTANTIATE_LAUNCH_ROTARY_POS_EMB(__nv_bfloat16); #endif -} - -template -void launch_apply_rotary_pos_emb(T* mixed_query, -T* key_layer, -unsigned head_size, -unsigned seq_len, -unsigned rotary_dim, -unsigned offset, -unsigned num_heads, -unsigned batch, -cudaStream_t stream) -{ -int total_count = batch * num_heads * seq_len; -dim3 block_dims(1024); -dim3 grid_dims((total_count - 1) / MAX_WARP_NUM + 1); // (batch_size); - -apply_rotary_pos_emb<<>>( -mixed_query, key_layer, rotary_dim, seq_len, offset, num_heads, head_size, total_count); -} - -template void launch_apply_rotary_pos_emb(float*, -float*, -unsigned, -unsigned, -unsigned, -unsigned, -unsigned, -unsigned, -cudaStream_t); -template void launch_apply_rotary_pos_emb<__half>(__half*, -__half*, -unsigned, -unsigned, -unsigned, -unsigned, -unsigned, -unsigned, -cudaStream_t); -*/ +INSTANTIATE_LAUNCH_ROTARY_POS_EMB(__half); diff --git a/csrc/transformer/inference/csrc/dequantize.cu b/csrc/transformer/inference/csrc/dequantize.cu index fe7b665f0efb..751a14440ae8 100644 --- a/csrc/transformer/inference/csrc/dequantize.cu +++ b/csrc/transformer/inference/csrc/dequantize.cu @@ -63,34 +63,15 @@ void launch_dequantize(T* output, output, input, qscale, output_size, hidden_dim, groups, merge_count); } -template void launch_dequantize(float*, - const int8_t*, - const float*, - unsigned, - unsigned, - unsigned, - unsigned, - cudaStream_t); +#define INSTANTIATE_DEQUANTIZE_MERGE(T) \ + template void launch_dequantize( \ + T*, const int8_t*, const float*, unsigned, unsigned, unsigned, unsigned, cudaStream_t); +INSTANTIATE_DEQUANTIZE_MERGE(float); #ifdef BF16_AVAILABLE -template void launch_dequantize<__nv_bfloat16>(__nv_bfloat16*, - const int8_t*, - const float*, - unsigned, - unsigned, - unsigned, - unsigned, - cudaStream_t); +INSTANTIATE_DEQUANTIZE_MERGE(__nv_bfloat16); #endif - -template void launch_dequantize<__half>(__half*, - const int8_t*, - const float*, - unsigned, - unsigned, - unsigned, - unsigned, - cudaStream_t); +INSTANTIATE_DEQUANTIZE_MERGE(__half); __global__ void dequantize_kernel(float* output, const int8_t* input, @@ -161,70 +142,12 @@ void launch_dequantize(T* output, output, input, qscale, hidden_dim, hid_cnt * hidden_dim, thd_cnt); } -template void launch_dequantize(float*, - const int8_t*, - const float*, - unsigned, - unsigned, - unsigned, - cudaStream_t); - -#ifdef BF16_AVAILABLE -template void launch_dequantize<__nv_bfloat16>(__nv_bfloat16*, - const int8_t*, - const float*, - unsigned, - unsigned, - unsigned, - cudaStream_t); -#endif - -template void launch_dequantize<__half>(__half*, - const int8_t*, - const float*, - unsigned, - unsigned, - unsigned, - cudaStream_t); - -template __global__ void dequantize_kernel(float* output, - const int8_t* input, - const float* qscale, - int output_size, - int hidden_dim, - int groups, - int merge_count); +#define INSTANTIATE_DEQUANTIZE_NO_MERGE(T) \ + template void launch_dequantize( \ + T*, const int8_t*, const float*, unsigned, unsigned, unsigned, cudaStream_t); +INSTANTIATE_DEQUANTIZE_NO_MERGE(float); #ifdef BF16_AVAILABLE -template __global__ void dequantize_kernel(__nv_bfloat16* output, - const int8_t* input, - const float* qscale, - int output_size, - int hidden_dim, - int groups, - int merge_count); +INSTANTIATE_DEQUANTIZE_NO_MERGE(__nv_bfloat16); #endif - -template __global__ void dequantize_kernel(__half* output, - const int8_t* input, - const float* qscale, - int output_size, - int hidden_dim, - int groups, - int merge_count); - -#ifdef BF16_AVAILABLE -template __global__ void dequantize_kernel(__nv_bfloat16* output, - const int8_t* input, - const float* qscale, - unsigned hidden_dim, - unsigned merge_hidden, - int cnt); -#endif - -template __global__ void dequantize_kernel(__half* output, - const int8_t* input, - const float* qscale, - unsigned hidden_dim, - unsigned merge_hidden, - int cnt); +INSTANTIATE_DEQUANTIZE_NO_MERGE(__half); diff --git a/csrc/transformer/inference/csrc/gelu.cu b/csrc/transformer/inference/csrc/gelu.cu index 26b6de0cdc53..3e6701d81e64 100644 --- a/csrc/transformer/inference/csrc/gelu.cu +++ b/csrc/transformer/inference/csrc/gelu.cu @@ -38,7 +38,8 @@ __global__ void fused_bias_gelu(T* input, const T* bias, int total_count, int in T data[values_per_access]; T data_bias[values_per_access]; mem_access::load_global(data, input + offset); - mem_access::load_global(data_bias, bias + (offset % intermediate_size)); + mem_access::load_global( + data_bias, bias + (offset % intermediate_size), bias != nullptr); #pragma unroll for (int i = 0; i < values_per_access; i++) { @@ -70,15 +71,14 @@ void launch_bias_gelu(T* input, input, bias, total_count, intermediate_size); } -template void launch_bias_gelu(float*, const float*, int, int, cudaStream_t); +#define INSTANTIATE_LAUNCH_BIAS_GELU(T) \ + template void launch_bias_gelu(T*, const T*, int, int, cudaStream_t); + +INSTANTIATE_LAUNCH_BIAS_GELU(float) #ifdef BF16_AVAILABLE -template void launch_bias_gelu<__nv_bfloat16>(__nv_bfloat16*, - const __nv_bfloat16*, - int, - int, - cudaStream_t); +INSTANTIATE_LAUNCH_BIAS_GELU(__nv_bfloat16) #endif -template void launch_bias_gelu<__half>(__half*, const __half*, int, int, cudaStream_t); +INSTANTIATE_LAUNCH_BIAS_GELU(__half) /* In-place channels-last bias add @@ -95,7 +95,8 @@ __global__ void fused_bias_add(T* input, const T* bias, int total_count, int int T data[values_per_access]; T data_bias[values_per_access]; mem_access::load_global(data, input + offset); - mem_access::load_global(data_bias, bias + (offset % intermediate_size)); + mem_access::load_global( + data_bias, bias + (offset % intermediate_size), bias != nullptr); #pragma unroll for (int i = 0; i < values_per_access; i++) { @@ -127,15 +128,14 @@ void launch_bias_add(T* input, input, bias, total_count, intermediate_size); } -template void launch_bias_add(float*, const float*, int, int, cudaStream_t); +#define INSTANTIATE_LAUNCH_BIAS_ADD(T) \ + template void launch_bias_add(T*, const T*, int, int, cudaStream_t); + +INSTANTIATE_LAUNCH_BIAS_ADD(float) #ifdef BF16_AVAILABLE -template void launch_bias_add<__nv_bfloat16>(__nv_bfloat16*, - const __nv_bfloat16*, - int, - int, - cudaStream_t); +INSTANTIATE_LAUNCH_BIAS_ADD(__nv_bfloat16) #endif -template void launch_bias_add<__half>(__half*, const __half*, int, int, cudaStream_t); +INSTANTIATE_LAUNCH_BIAS_ADD(__half) __global__ void fused_bias_residual(float* residual, const float* hidden_state, @@ -282,44 +282,14 @@ void launch_bias_residual(T* residual, preln); } -template void launch_bias_residual< - float>(float*, float*, float*, float*, float*, int, int, int, bool, cudaStream_t); -#ifdef BF16_AVAILABLE -template void launch_bias_residual<__nv_bfloat16>(__nv_bfloat16*, - __nv_bfloat16*, - __nv_bfloat16*, - __nv_bfloat16*, - __nv_bfloat16*, - int, - int, - int, - bool, - cudaStream_t); -#endif -template void launch_bias_residual< - __half>(__half*, __half*, __half*, __half*, __half*, int, int, int, bool, cudaStream_t); +#define INSTANTIATE_LAUNCH_BIAS_RESIDUAL(T) \ + template void launch_bias_residual(T*, T*, T*, T*, T*, int, int, int, bool, cudaStream_t); +INSTANTIATE_LAUNCH_BIAS_RESIDUAL(float); #ifdef BF16_AVAILABLE -template __global__ void fused_bias_residual(__nv_bfloat16* residual, - const __nv_bfloat16* hidden_state, - const __nv_bfloat16* attn, - const __nv_bfloat16* bias, - const __nv_bfloat16* attn_bias, - const int total_count, - const int intermediate_size, - const float mp_scale, - const bool preln); +INSTANTIATE_LAUNCH_BIAS_RESIDUAL(__nv_bfloat16); #endif - -template __global__ void fused_bias_residual(__half* residual, - const __half* hidden_state, - const __half* attn, - const __half* bias, - const __half* attn_bias, - const int total_count, - const int intermediate_size, - const float mp_scale, - const bool preln); +INSTANTIATE_LAUNCH_BIAS_RESIDUAL(__half); __global__ void gptj_residual_add(float* residual, const float* hidden_state, @@ -446,58 +416,15 @@ void launch_gptj_residual_add(T* residual, residual, hidden_state, attn, bias, attn_bias, total_count, hidden_dim / 4, 1.0 / mp_size); } -template void launch_gptj_residual_add(float*, - float*, - float*, - float*, - float*, - int, - int, - int, - cudaStream_t); - -#ifdef BF16_AVAILABLE -template void launch_gptj_residual_add<__nv_bfloat16>(__nv_bfloat16*, - __nv_bfloat16*, - __nv_bfloat16*, - __nv_bfloat16*, - __nv_bfloat16*, - int, - int, - int, - cudaStream_t); -#endif - -template void launch_gptj_residual_add<__half>(__half*, - __half*, - __half*, - __half*, - __half*, - int, - int, - int, - cudaStream_t); +#define INSTANTIATE_GPT_RES_ADD(T) \ + template void launch_gptj_residual_add(T*, T*, T*, T*, T*, int, int, int, cudaStream_t); +INSTANTIATE_GPT_RES_ADD(float); +INSTANTIATE_GPT_RES_ADD(__half); #ifdef BF16_AVAILABLE -template __global__ void gptj_residual_add(__nv_bfloat16* residual, - const __nv_bfloat16* hidden_state, - const __nv_bfloat16* attn, - const __nv_bfloat16* bias, - const __nv_bfloat16* attn_bias, - const int total_count, - const int intermediate_size, - const float mp_scale); +INSTANTIATE_GPT_RES_ADD(__nv_bfloat16); #endif -template __global__ void gptj_residual_add(__half* residual, - const __half* hidden_state, - const __half* attn, - const __half* bias, - const __half* attn_bias, - const int total_count, - const int intermediate_size, - const float mp_scale); - template __global__ void moe_res_matmul(T* residual, T* coef, T* mlp_out, int seq_len, int hidden_dim) { @@ -542,28 +469,14 @@ void launch_moe_res_matmul(T* residual, residual, coef, mlp_out, seq_len, hidden_dim); } -template void launch_moe_res_matmul(float* residual, - float* coef, - float* mlp_out, - int seq_len, - int hidden_dim, - cudaStream_t stream); +#define INSTANTIATE_LAUNCH_MOE_RES_MATMUL(T) \ + template void launch_moe_res_matmul(T*, T*, T*, int, int, cudaStream_t); +INSTANTIATE_LAUNCH_MOE_RES_MATMUL(float); #ifdef BF16_AVAILABLE -template void launch_moe_res_matmul(__nv_bfloat16* residual, - __nv_bfloat16* coef, - __nv_bfloat16* mlp_out, - int seq_len, - int hidden_dim, - cudaStream_t stream); +INSTANTIATE_LAUNCH_MOE_RES_MATMUL(__nv_bfloat16); #endif - -template void launch_moe_res_matmul(__half* residual, - __half* coef, - __half* mlp_out, - int seq_len, - int hidden_dim, - cudaStream_t stream); +INSTANTIATE_LAUNCH_MOE_RES_MATMUL(__half); template __global__ void pad_data_kernel(T* padded_output, T* output, int head_size, int padded_head_size) @@ -607,41 +520,15 @@ void pad_data(T* padded_output, pad_data_kernel<<>>( padded_output, output, head_size / 8, padded_head_size / 8); } -template void pad_data(__half* padded_output, - __half* output, - int bsz, - int head_size, - int padded_head_size, - cudaStream_t stream); -#ifdef BF16_AVAILABLE -template void pad_data(__nv_bfloat16* padded_output, - __nv_bfloat16* output, - int bsz, - int head_size, - int padded_head_size, - cudaStream_t stream); -#endif - -template void pad_data(float* padded_output, - float* output, - int bsz, - int head_size, - int padded_head_size, - cudaStream_t stream); +#define INSTANTIATE_PAD_DATA(T) template void pad_data(T*, T*, int, int, int, cudaStream_t stream); +INSTANTIATE_PAD_DATA(float); +INSTANTIATE_PAD_DATA(__half); #ifdef BF16_AVAILABLE -template __global__ void pad_data_kernel(__nv_bfloat16* padded_output, - __nv_bfloat16* output, - int head_size, - int padded_head_size); +INSTANTIATE_PAD_DATA(__nv_bfloat16); #endif -template __global__ void pad_data_kernel(__half* padded_output, - __half* output, - int head_size, - int padded_head_size); - template __global__ void pad_head_seq_kernel(T* padded_output, T* output, @@ -696,34 +583,14 @@ void pad_head_seq(T* padded_output, padded_output, output, seq_len, padded_seq_len, head_size / 8, padded_head_size / 8); } -template void pad_head_seq(__half* padded_output, - __half* output, - int bsz, - int seq_len, - int padded_seq_len, - int head_size, - int padded_head_size, - cudaStream_t stream); +#define INSTANTIATE_PAD_HEAD_SEQ(T) \ + template void pad_head_seq(T*, T*, int, int, int, int, int, cudaStream_t); +INSTANTIATE_PAD_HEAD_SEQ(__half); #ifdef BF16_AVAILABLE -template void pad_head_seq(__nv_bfloat16* padded_output, - __nv_bfloat16* output, - int bsz, - int seq_len, - int padded_seq_len, - int head_size, - int padded_head_size, - cudaStream_t stream); +INSTANTIATE_PAD_HEAD_SEQ(__nv_bfloat16); #endif - -template void pad_head_seq(float* padded_output, - float* output, - int bsz, - int seq_len, - int padded_seq_len, - int head_size, - int padded_head_size, - cudaStream_t stream); +INSTANTIATE_PAD_HEAD_SEQ(float); // TODO(cmikeh2): evaluate different GeLU performance __device__ __forceinline__ float old_gelu(float val) @@ -739,12 +606,15 @@ constexpr int steps = 2; constexpr int granularity = 16; } // namespace fused_geglu -template -__global__ void fused_bias_geglu(T* output, - const T* activation, - const T* bias, - int base_channels, - int total_elems) +__device__ __forceinline__ float silu(float val) { return val / (1.0f + expf(-val)); } + +template +__global__ void fused_gate_activation(T* output, + const T* activation, + const T* bias, + int base_channels, + int output_stride, + int total_elems) { constexpr int T_per_access = fused_geglu::granularity / sizeof(T); constexpr int T_per_step = T_per_access * fused_geglu::threads; @@ -769,9 +639,10 @@ __global__ void fused_bias_geglu(T* output, activation + seq_offset + channel_id); mem_access::load_global( activation_buffer_2, activation + seq_offset + channel_id + base_channels); - mem_access::load_global(bias_buffer_1, bias + channel_id); - mem_access::load_global(bias_buffer_2, - bias + channel_id + base_channels); + mem_access::load_global( + bias_buffer_1, bias + channel_id, bias != nullptr); + mem_access::load_global( + bias_buffer_2, bias + channel_id + base_channels, bias != nullptr); // Since the GeLU is going to happen at float, might as well // convert @@ -779,23 +650,26 @@ __global__ void fused_bias_geglu(T* output, for (int v = 0; v < T_per_access; v++) { T hidden_state = activation_buffer_1[v] + bias_buffer_1[v]; T pre_gate = activation_buffer_2[v] + bias_buffer_2[v]; - float gate_f = old_gelu(conversion::to(pre_gate)); + float pre_gate_f = conversion::to(pre_gate); + float gate_f = (useGelu) ? old_gelu(pre_gate_f) : silu(pre_gate_f); T gate = conversion::to(gate_f); activation_buffer_1[v] = hidden_state * gate; } - mem_access::store_global(output + iter_id, - activation_buffer_1); + mem_access::store_global( + output + seq_id * output_stride + channel_id, activation_buffer_1); } } } template -void launch_fused_bias_geglu(T* output, +void launch_gated_activation(T* output, const T* activation, const T* bias, int rows, + int output_stride, int elems_per_row, + bool use_gelu, cudaStream_t stream) { /* @@ -816,22 +690,21 @@ void launch_fused_bias_geglu(T* output, dim3 block(fused_geglu::threads); dim3 grid((total_elems + T_per_block - 1) / T_per_block); - fused_bias_geglu<<>>( - output, activation, bias, base_channels, total_elems); + if (use_gelu) { + fused_gate_activation<<>>( + output, activation, bias, base_channels, output_stride, total_elems); + } else { + fused_gate_activation<<>>( + output, activation, bias, base_channels, output_stride, total_elems); + } } -template void launch_fused_bias_geglu(__half*, - const __half*, - const __half*, - int, - int, - cudaStream_t); +#define INSTANTIATE_LAUNCH_GATED_ACTIVATION(T) \ + template void launch_gated_activation( \ + T*, const T*, const T*, int, int, int, bool, cudaStream_t); + +INSTANTIATE_LAUNCH_GATED_ACTIVATION(__half); #ifdef BF16_AVAILABLE -template void launch_fused_bias_geglu(__nv_bfloat16*, - const __nv_bfloat16*, - const __nv_bfloat16*, - int, - int, - cudaStream_t); +INSTANTIATE_LAUNCH_GATED_ACTIVATION(__nv_bfloat16); #endif -template void launch_fused_bias_geglu(float*, const float*, const float*, int, int, cudaStream_t); +INSTANTIATE_LAUNCH_GATED_ACTIVATION(float); diff --git a/csrc/transformer/inference/csrc/layer_norm.cu b/csrc/transformer/inference/csrc/layer_norm.cu index 807d60e740df..e5e7e89c9d20 100644 --- a/csrc/transformer/inference/csrc/layer_norm.cu +++ b/csrc/transformer/inference/csrc/layer_norm.cu @@ -46,7 +46,7 @@ __global__ void fused_ln(T* output, (tb.thread_index().y * elems_per_row); const int thread_offset = tb.thread_index().x * T_per_load; const int base_offset = block_offset + thread_offset; - const int stride = tb.size() * T_per_load; + const int stride = blockDim.x * T_per_load; float sum = reduce::init(); @@ -57,8 +57,6 @@ __global__ void fused_ln(T* output, #pragma unRoll for (int i = 0; i < unRoll; i++) { T* iteration_buffer = local_buffer + i * T_per_load; - T residual_buffer[T_per_load]; - T bias_buffer[T_per_load]; mem_access::load_global( iteration_buffer, input_base + i * stride, thread_offset + i * stride < elems_per_row); @@ -91,8 +89,8 @@ __global__ void fused_ln(T* output, const float variance = mean_diff / elems_per_row; const float denom = __frsqrt_rn(variance + epsilon); - const T mean_compute = conversion::to(mean); - const T denom_compute = conversion::to(denom); + // const T mean_compute = conversion::to(mean); + // const T denom_compute = conversion::to(denom); T* block_output = output + block_offset; @@ -109,8 +107,11 @@ __global__ void fused_ln(T* output, #pragma unRoll for (int j = 0; j < T_per_load; j++) { - iteration_buffer[j] = (iteration_buffer[j] - mean_compute) * denom_compute; - iteration_buffer[j] = iteration_buffer[j] * gamma_local[j] + beta_local[j]; + float val = conversion::to(iteration_buffer[j]); + val = (val - mean) * denom; + val = + val * conversion::to(gamma_local[j]) + conversion::to(beta_local[j]); + iteration_buffer[j] = conversion::to(val); } if (do_loads) { @@ -189,26 +190,14 @@ void launch_fused_ln(T* output, } } -template void launch_fused_ln(__half*, - const __half*, - const __half*, - const __half*, - float, - int, - int, - cudaStream_t); +#define INSTANTIATE_FUSED_LN(T) \ + template void launch_fused_ln(T*, const T*, const T*, const T*, float, int, int, cudaStream_t); + +INSTANTIATE_FUSED_LN(__half); #ifdef BF16_AVAILABLE -template void launch_fused_ln(__nv_bfloat16*, - const __nv_bfloat16*, - const __nv_bfloat16*, - const __nv_bfloat16*, - float, - int, - int, - cudaStream_t); +INSTANTIATE_FUSED_LN(__nv_bfloat16); #endif -template void -launch_fused_ln(float*, const float*, const float*, const float*, float, int, int, cudaStream_t); +INSTANTIATE_FUSED_LN(float); /* Fused resiual + bias + layer norm implementation. Assumes elems_per_row % 8 @@ -284,7 +273,7 @@ __global__ void fused_residual_ln(T* output, float vals_up_cast = conversion::to(iteration_buffer[j]); float res_up_cast = conversion::to(residual_buffer[j]); float bias_up_cast = conversion::to(bias_buffer[j]); - vals_up_cast += res_up_cast + bias_up_cast; + vals_up_cast = vals_up_cast + bias_up_cast + res_up_cast; sum = reduce::element(sum, vals_up_cast); iteration_buffer[j] = conversion::to(vals_up_cast); } @@ -315,9 +304,6 @@ __global__ void fused_residual_ln(T* output, const float variance = mean_diff / elems_per_row; const float denom = __frsqrt_rn(variance + epsilon); - const T mean_compute = conversion::to(mean); - const T denom_compute = conversion::to(denom); - T* block_output = output + block_offset; #pragma unRoll @@ -333,8 +319,13 @@ __global__ void fused_residual_ln(T* output, #pragma unRoll for (int j = 0; j < T_per_load; j++) { - iteration_buffer[j] = (iteration_buffer[j] - mean_compute) * denom_compute; - iteration_buffer[j] = iteration_buffer[j] * gamma_local[j] + beta_local[j]; + // iteration_buffer[j] = (iteration_buffer[j] - mean_compute) * denom_compute; + // iteration_buffer[j] = iteration_buffer[j] * gamma_local[j] + beta_local[j]; + float val = conversion::to(iteration_buffer[j]); + val = (val - mean) * denom; + val = + val * conversion::to(gamma_local[j]) + conversion::to(beta_local[j]); + iteration_buffer[j] = conversion::to(val); } if (do_loads) { @@ -491,77 +482,22 @@ void launch_fused_residual_ln_store_pre_ln_res(T* norm_output, } } -// No-store specializations -template void launch_fused_residual_ln(__half*, - const __half*, - const __half*, - const __half*, - const __half*, - const __half*, - float, - int, - int, - cudaStream_t); +#define INSTANTIATE_RES_LN(T) \ + template void launch_fused_residual_ln( \ + T*, const T*, const T*, const T*, const T*, const T*, float, int, int, cudaStream_t); + +#define INSTANTIATE_PRE_LN_RES(T) \ + template void launch_fused_residual_ln_store_pre_ln_res( \ + T*, T*, const T*, const T*, const T*, const T*, const T*, float, int, int, cudaStream_t); +INSTANTIATE_RES_LN(__half); +INSTANTIATE_RES_LN(float); #ifdef BF16_AVAILABLE -template void launch_fused_residual_ln(__nv_bfloat16*, - const __nv_bfloat16*, - const __nv_bfloat16*, - const __nv_bfloat16*, - const __nv_bfloat16*, - const __nv_bfloat16*, - float, - int, - int, - cudaStream_t); +INSTANTIATE_RES_LN(__nv_bfloat16); #endif -template void launch_fused_residual_ln(float*, - const float*, - const float*, - const float*, - const float*, - const float*, - float, - int, - int, - cudaStream_t); - -// Store specializations -template void launch_fused_residual_ln_store_pre_ln_res(__half*, - __half*, - const __half*, - const __half*, - const __half*, - const __half*, - const __half*, - float, - int, - int, - cudaStream_t); - +INSTANTIATE_PRE_LN_RES(__half); +INSTANTIATE_PRE_LN_RES(float); #ifdef BF16_AVAILABLE -template void launch_fused_residual_ln_store_pre_ln_res(__nv_bfloat16*, - __nv_bfloat16*, - const __nv_bfloat16*, - const __nv_bfloat16*, - const __nv_bfloat16*, - const __nv_bfloat16*, - const __nv_bfloat16*, - float, - int, - int, - cudaStream_t); +INSTANTIATE_PRE_LN_RES(__nv_bfloat16); #endif - -template void launch_fused_residual_ln_store_pre_ln_res(float*, - float*, - const float*, - const float*, - const float*, - const float*, - const float*, - float, - int, - int, - cudaStream_t); diff --git a/csrc/transformer/inference/csrc/pointwise_ops.cu b/csrc/transformer/inference/csrc/pointwise_ops.cu new file mode 100644 index 000000000000..0301ff777042 --- /dev/null +++ b/csrc/transformer/inference/csrc/pointwise_ops.cu @@ -0,0 +1,74 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include +#include "conversion_utils.h" +#include "ds_kernel_utils.h" +#include "memory_access_utils.h" + +namespace pwise { +constexpr int granularity = 16; +constexpr int unroll = 4; +constexpr int threads = 256; +} // namespace pwise + +template +__global__ void vector_add_kernel(T* out, const T* a, const T* b, float gamma, int num_elems) +{ + constexpr int T_per_access = pwise::granularity / sizeof(T); + + const int block_offset = blockIdx.x * pwise::threads * pwise::unroll * T_per_access; + const int thread_offset = threadIdx.x * T_per_access; + const int total_offset = block_offset + thread_offset; + constexpr int stride = pwise::threads * T_per_access; + +#pragma unroll + for (int i = 0; i < pwise::unroll; i++) { + T temp_buf_a[T_per_access], temp_buf_b[T_per_access]; + + const int iter_idx = total_offset + i * stride; + + mem_access::load_global(temp_buf_a, a + iter_idx, iter_idx < num_elems); + mem_access::load_global(temp_buf_b, b + iter_idx, iter_idx < num_elems); + +#pragma unroll + for (int j = 0; j < T_per_access; j++) { + float up_cast_a = conversion::to(temp_buf_a[j]); + float up_cast_b = conversion::to(temp_buf_b[j]); + temp_buf_a[j] = conversion::to((gamma * up_cast_a) + up_cast_b); + } + + if (iter_idx < num_elems) { + mem_access::store_global(out + iter_idx, temp_buf_a); + } + } +} + +template +void launch_vector_add(T* out, + const T* a, + const T* b, + float gamma, + int num_elems, + cudaStream_t stream) +{ + constexpr int T_per_access = pwise::granularity / sizeof(T); + constexpr int T_per_block = pwise::threads * T_per_access * pwise::unroll; + + dim3 block(pwise::threads); + dim3 grid((num_elems + T_per_block - 1) / T_per_block); + + vector_add_kernel<<>>(out, a, b, gamma, num_elems); +} + +#define INSTANTIATE_VECTOR_ADD(T) \ + template void launch_vector_add( \ + T * out, const T* a, const T* b, float gamma, int num_elems, cudaStream_t stream); + +INSTANTIATE_VECTOR_ADD(float) +INSTANTIATE_VECTOR_ADD(__half) +#ifdef BF16_AVAILABLE +INSTANTIATE_VECTOR_ADD(__nv_bfloat16) +#endif diff --git a/csrc/transformer/inference/csrc/pt_binding.cpp b/csrc/transformer/inference/csrc/pt_binding.cpp index eb9bcf0e4341..8f3ca39e1f3a 100644 --- a/csrc/transformer/inference/csrc/pt_binding.cpp +++ b/csrc/transformer/inference/csrc/pt_binding.cpp @@ -16,7 +16,9 @@ std::array gemm_algos = std::array({99, 99, 99}); // NOTE: This activation function type enum should be always in sync // with the python counterpart, otherwise the casting from python binding // will be incorrect. -enum class ActivationFuncType { UNKNOWN = 0, GELU = 1, ReLU = 2 }; +enum class ActivationFuncType { UNKNOWN = 0, GELU = 1, ReLU = 2, GATED_GELU = 3, GATED_SILU = 4 }; + +enum class NormType { UNKNOWN = 0, LayerNorm = 1, GroupNorm = 2, RMSNorm = 3 }; enum class TransformerType : uint8_t { UNKNOWN = 0, GPTType = 1, BERTType = 2 }; @@ -499,8 +501,6 @@ std::vector ds_softmax_context(at::Tensor& query_key_value, (is_prompt ? 0 : soft_len - 1), heads, bsz, - rotate_half, - rotate_every_two, InferenceContext::Instance().GetCurrentStream(), InferenceContext::Instance().GetMaxTokenLenght()); @@ -567,12 +567,29 @@ at::Tensor ds_bias_gelu(at::Tensor& input, at::Tensor& bias) return input_cont; } -at::Tensor ds_bias_geglu(at::Tensor& activation, at::Tensor& bias) +#define DISPATCH_GATED_ACT(T_TYPE, C_TYPE) \ + if (activation.options().dtype() == torch::T_TYPE) { \ + launch_gated_activation((C_TYPE*)output.data_ptr(), \ + (const C_TYPE*)activation.data_ptr(), \ + (const C_TYPE*)bias.data_ptr(), \ + rows, \ + out_channels, \ + channels, \ + activation_type == ActivationFuncType::GATED_GELU, \ + InferenceContext::Instance().GetCurrentStream()); \ + } + +at::Tensor ds_gated_activation(at::Tensor& activation, at::Tensor& bias, int actFun) { /* Used in FF of Stable diffusion */ + const ActivationFuncType activation_type = static_cast(actFun); + + assert(activation_type == ActivationFuncType::GATED_GELU || + activation_type == ActivationFuncType::GATED_SILU); + const int batch_size = activation.size(0); const int seq_len = activation.size(1); const int channels = activation.size(2); @@ -583,21 +600,11 @@ at::Tensor ds_bias_geglu(at::Tensor& activation, at::Tensor& bias) auto output = at::empty({batch_size, seq_len, out_channels}, activation.options()); - if (activation.options().dtype() == torch::kFloat32) { - launch_fused_bias_geglu((float*)output.data_ptr(), - (const float*)activation.data_ptr(), - (const float*)bias.data_ptr(), - rows, - channels, - InferenceContext::Instance().GetCurrentStream()); - } else { - launch_fused_bias_geglu((__half*)output.data_ptr(), - (const __half*)activation.data_ptr(), - (const __half*)bias.data_ptr(), - rows, - channels, - InferenceContext::Instance().GetCurrentStream()); - } + DISPATCH_GATED_ACT(kFloat, float); + DISPATCH_GATED_ACT(kHalf, __half); +#ifdef BF16_AVAILABLE + DISPATCH_GATED_ACT(kBFloat16, __nv_bfloat16); +#endif return output; } @@ -651,35 +658,99 @@ at::Tensor ds_bias_residual(at::Tensor& input, at::Tensor& residual, at::Tensor& return input_cont; } +#define DISPATCH_LAYER_NORM(T_TYPE, C_TYPE) \ + if (input.options().dtype() == torch::T_TYPE) { \ + launch_fused_ln((C_TYPE*)output.data_ptr(), \ + (const C_TYPE*)input.data_ptr(), \ + (const C_TYPE*)gamma.data_ptr(), \ + (const C_TYPE*)beta.data_ptr(), \ + epsilon, \ + rows, \ + elems_per_row, \ + InferenceContext::Instance().GetCurrentStream()); \ + } + at::Tensor ds_layer_norm(at::Tensor& input, at::Tensor& gamma, at::Tensor& beta, float epsilon) { const int rows = input.size(0) * input.size(1); const int elems_per_row = input.size(2); auto output = at::empty_like(input); - if (input.options().dtype() == torch::kFloat16) { - launch_fused_ln((__half*)output.data_ptr(), - (const __half*)input.data_ptr(), - (const __half*)gamma.data_ptr(), - (const __half*)beta.data_ptr(), - epsilon, - rows, - elems_per_row, - InferenceContext::Instance().GetCurrentStream()); - } else { - launch_fused_ln((float*)output.data_ptr(), - (const float*)input.data_ptr(), - (const float*)gamma.data_ptr(), - (const float*)beta.data_ptr(), - epsilon, - rows, - elems_per_row, - InferenceContext::Instance().GetCurrentStream()); + DISPATCH_LAYER_NORM(kFloat, float); + DISPATCH_LAYER_NORM(kHalf, __half); +#ifdef BF16_AVAILABLE + DISPATCH_LAYER_NORM(kBFloat16, __nv_bfloat16); +#endif + + return output; +} + +#define DISPATCH_RMS_NORM(T_TYPE, C_TYPE) \ + if (input.options().dtype() == torch::T_TYPE) { \ + launch_rms_norm((C_TYPE*)output.data_ptr(), \ + (C_TYPE*)nullptr, \ + (const C_TYPE*)input.data_ptr(), \ + (const C_TYPE*)nullptr, \ + (const C_TYPE*)gamma.data_ptr(), \ + epsilon, \ + rows, \ + elems_per_row, \ + InferenceContext::Instance().GetCurrentStream()); \ } +at::Tensor ds_rms_norm(at::Tensor& input, at::Tensor& gamma, float epsilon) +{ + // Get number of dims of tensor + int num_dims = input.dim(); + const int rows = (num_dims == 2) ? input.size(0) : input.size(0) * input.size(1); + const int elems_per_row = (num_dims == 2) ? input.size(1) : input.size(2); + + auto output = at::empty_like(input); + + DISPATCH_RMS_NORM(kFloat, float); + DISPATCH_RMS_NORM(kHalf, __half); +#ifdef BF16_AVAILABLE + DISPATCH_RMS_NORM(kBFloat16, __nv_bfloat16); +#endif + return output; } +#define DISPATCH_PRE_RMS_NORM(T_TYPE, C_TYPE) \ + if (input.options().dtype() == torch::T_TYPE) { \ + launch_rms_norm((C_TYPE*)output.data_ptr(), \ + (C_TYPE*)res_out.data_ptr(), \ + (const C_TYPE*)input.data_ptr(), \ + (const C_TYPE*)residual.data_ptr(), \ + (const C_TYPE*)gamma.data_ptr(), \ + epsilon, \ + rows, \ + elems_per_row, \ + InferenceContext::Instance().GetCurrentStream()); \ + } + +std::vector ds_pre_rms_norm(at::Tensor& input, + at::Tensor& residual, + at::Tensor& gamma, + float epsilon) +{ + // Get number of dims of tensor + int num_dims = input.dim(); + const int rows = (num_dims == 2) ? input.size(0) : input.size(0) * input.size(1); + const int elems_per_row = (num_dims == 2) ? input.size(1) : input.size(2); + + auto output = at::empty_like(input); + auto res_out = at::empty_like(residual); + + DISPATCH_PRE_RMS_NORM(kFloat, float); + DISPATCH_PRE_RMS_NORM(kHalf, __half); +#ifdef BF16_AVAILABLE + DISPATCH_PRE_RMS_NORM(kBFloat16, __nv_bfloat16); +#endif + + return {output, res_out}; +} + template void ds_layer_norm_internal(T* workspace, at::Tensor& input, @@ -698,6 +769,20 @@ void ds_layer_norm_internal(T* workspace, InferenceContext::Instance().GetCurrentStream()); } +#define DISPATCH_LAYER_NORM_RESIDUAL(T_TYPE, C_TYPE) \ + if (input.options().dtype() == torch::T_TYPE) { \ + launch_fused_residual_ln((C_TYPE*)output.data_ptr(), \ + (const C_TYPE*)input.data_ptr(), \ + (const C_TYPE*)residual.data_ptr(), \ + (const C_TYPE*)bias.data_ptr(), \ + (const C_TYPE*)gamma.data_ptr(), \ + (const C_TYPE*)beta.data_ptr(), \ + epsilon, \ + rows, \ + elems_per_row, \ + InferenceContext::Instance().GetCurrentStream()); \ + } + /* Currently only used in unit testing */ at::Tensor ds_layer_norm_residual(at::Tensor& input, at::Tensor& bias, @@ -710,33 +795,31 @@ at::Tensor ds_layer_norm_residual(at::Tensor& input, const int elems_per_row = input.size(2); auto output = at::empty_like(input); - if (input.options().dtype() == torch::kFloat16) { - launch_fused_residual_ln((__half*)output.data_ptr(), - (const __half*)input.data_ptr(), - (const __half*)residual.data_ptr(), - (const __half*)bias.data_ptr(), - (const __half*)gamma.data_ptr(), - (const __half*)beta.data_ptr(), - epsilon, - rows, - elems_per_row, - InferenceContext::Instance().GetCurrentStream()); - } else { - launch_fused_residual_ln((float*)output.data_ptr(), - (const float*)input.data_ptr(), - (const float*)residual.data_ptr(), - (const float*)bias.data_ptr(), - (const float*)gamma.data_ptr(), - (const float*)beta.data_ptr(), - epsilon, - rows, - elems_per_row, - InferenceContext::Instance().GetCurrentStream()); - } + DISPATCH_LAYER_NORM_RESIDUAL(kFloat, float); + DISPATCH_LAYER_NORM_RESIDUAL(kHalf, __half); +#ifdef BF16_AVAILABLE + DISPATCH_LAYER_NORM_RESIDUAL(kBFloat16, __nv_bfloat16); +#endif return output; } +#define DISPATCH_PRE_LAYER_NORM_RESIDUAL(T_TYPE, C_TYPE) \ + if (input.options().dtype() == torch::T_TYPE) { \ + launch_fused_residual_ln_store_pre_ln_res( \ + (C_TYPE*)norm_output.data_ptr(), \ + (C_TYPE*)res_output.data_ptr(), \ + (const C_TYPE*)input.data_ptr(), \ + (const C_TYPE*)residual.data_ptr(), \ + (const C_TYPE*)bias.data_ptr(), \ + (const C_TYPE*)gamma.data_ptr(), \ + (const C_TYPE*)beta.data_ptr(), \ + epsilon, \ + rows, \ + elems_per_row, \ + InferenceContext::Instance().GetCurrentStream()); \ + } + /* Currently only used in unit testing */ std::vector ds_layer_norm_residual_store_pre_ln_res(at::Tensor& input, at::Tensor& bias, @@ -750,31 +833,11 @@ std::vector ds_layer_norm_residual_store_pre_ln_res(at::Tensor& inpu auto norm_output = at::empty_like(input); auto res_output = at::empty_like(input); - if (input.options().dtype() == torch::kFloat16) { - launch_fused_residual_ln_store_pre_ln_res((__half*)norm_output.data_ptr(), - (__half*)res_output.data_ptr(), - (const __half*)input.data_ptr(), - (const __half*)residual.data_ptr(), - (const __half*)bias.data_ptr(), - (const __half*)gamma.data_ptr(), - (const __half*)beta.data_ptr(), - epsilon, - rows, - elems_per_row, - InferenceContext::Instance().GetCurrentStream()); - } else { - launch_fused_residual_ln_store_pre_ln_res((float*)norm_output.data_ptr(), - (float*)res_output.data_ptr(), - (const float*)input.data_ptr(), - (const float*)residual.data_ptr(), - (const float*)bias.data_ptr(), - (const float*)gamma.data_ptr(), - (const float*)beta.data_ptr(), - epsilon, - rows, - elems_per_row, - InferenceContext::Instance().GetCurrentStream()); - } + DISPATCH_PRE_LAYER_NORM_RESIDUAL(kFloat, float); + DISPATCH_PRE_LAYER_NORM_RESIDUAL(kHalf, __half); +#ifdef BF16_AVAILABLE + DISPATCH_PRE_LAYER_NORM_RESIDUAL(kBFloat16, __nv_bfloat16); +#endif return {norm_output, res_output}; } @@ -878,6 +941,73 @@ at::Tensor qkv_unfused_cublas(at::Tensor& output, return torch::from_blob(workspace, input.sizes(), input.options()); } +template +std::vector ds_rms_qkv(at::Tensor& input, + at::Tensor& weight, + at::Tensor& q_scale, + at::Tensor& gamma, + const float epsilon, + bool q_int8, + bool transposed_mode) +{ + const int bsz = input.size(0) * input.size(1); + T* workspace = (T*)InferenceContext::Instance().GetWorkSpace(); + T* rms_norm_ptr = workspace + (3 * bsz * input.size(2)); + int out_size = (transposed_mode || q_int8) ? weight.size(0) : weight.size(1); + + auto options = at::TensorOptions() + .dtype(input.options().dtype()) + .layout(at::kStrided) + .device(at::kCUDA) + .requires_grad(false); + auto rms_norm = at::from_blob(rms_norm_ptr, input.sizes(), options); + auto output = at::from_blob(workspace, {input.size(0), input.size(1), out_size}, options); + + launch_rms_norm((T*)rms_norm.data_ptr(), + (T*)nullptr, + (const T*)input.data_ptr(), + (const T*)nullptr, + (const T*)gamma.data_ptr(), + epsilon, + bsz, + input.size(2), + InferenceContext::Instance().GetCurrentStream()); + + if (q_int8) { + quantized_gemm((T*)output.data_ptr(), + (T*)rms_norm.data_ptr(), + weight, + q_scale, + q_scale.size(0), + bsz, + input.size(2)); + } else { + float alpha = (T)1.0; + float gemm_beta = (T)0.0; + + cublasSetStream(InferenceContext::Instance().GetCublasHandle(), + InferenceContext::Instance().GetCurrentStream()); + cublas_gemm_ex(InferenceContext::Instance().GetCublasHandle(), + (transposed_mode ? CUBLAS_OP_T : CUBLAS_OP_N), + CUBLAS_OP_N, + weight.size(transposed_mode ? 0 : 1), + bsz, + input.size(2), + &alpha, + &gemm_beta, + (T*)weight.data_ptr(), + (T*)rms_norm.data_ptr(), + (T*)output.data_ptr(), +#ifdef __HIP_PLATFORM_HCC__ + rocblas_gemm_algo_standard); +#else + CUBLAS_GEMM_DEFAULT_TENSOR_OP); +#endif + } + + return {output, rms_norm}; +} + template std::vector ds_qkv_gemm(at::Tensor& input, at::Tensor& weight, @@ -887,10 +1017,6 @@ std::vector ds_qkv_gemm(at::Tensor& input, at::Tensor& beta, const float epsilon, bool add_bias, - unsigned num_layers, - bool external_cache, - unsigned mp_size, - unsigned rank, bool q_int8, bool transposed_mode) { @@ -965,40 +1091,6 @@ void quantized_gemm(at::Tensor& output, #endif } -template -at::Tensor ds_qkv_gemm_int8(at::Tensor& input, - at::Tensor& weight, - at::Tensor& bias, - at::Tensor& gamma, - at::Tensor& beta, - const float epsilon, - at::Tensor& q_scale, - int groups, - bool add_bias) -{ - int bsz = input.size(0) * input.size(1); - auto input_cont = input.contiguous(); - auto options = at::TensorOptions() - .dtype(input_cont.options().dtype()) - .layout(at::kStrided) - .device(at::kCUDA) - .requires_grad(false); - - auto output = at::empty({input_cont.size(0), input_cont.size(1), weight.size(1)}, options); - - auto inp_norm = ds_layer_norm(input_cont, gamma, beta, epsilon); - - quantized_gemm(output, inp_norm, weight, q_scale, groups, 0); - if (add_bias) - launch_bias_add((T*)output.data_ptr(), - (T*)bias.data_ptr(), - weight.size(1), - bsz, - InferenceContext::Instance().GetCurrentStream()); - - return output; -} - template at::Tensor ds_linear_layer(at::Tensor& input, at::Tensor& weight, @@ -1207,31 +1299,6 @@ std::vector padd_add_transform(at::Tensor& query, {query.size(0), heads, key_value_length, padded_head_size}, query.options())}; } -template -at::Tensor ds_linear_layer_int8(at::Tensor& input, - at::Tensor& weight, - at::Tensor& bias, - at::Tensor& q_scale, - int groups) -{ - auto input_cont = input.contiguous(); - auto options = at::TensorOptions() - .dtype(input_cont.options().dtype()) - .layout(at::kStrided) - .device(at::kCUDA) - .requires_grad(false); - int bsz = input_cont.size(0) * input_cont.size(1); - - auto output = at::empty({input_cont.size(0), input_cont.size(1), weight.size(1)}, options); - - quantized_gemm(output, input_cont, weight, q_scale, groups, 0); - launch_bias_add((T*)output.data_ptr(), - (T*)bias.data_ptr(), - weight.size(1), - bsz, - InferenceContext::Instance().GetCurrentStream()); - return output; -} template at::Tensor ds_vector_matmul(at::Tensor& input, @@ -1469,39 +1536,148 @@ std::vector ds_mlp_gemm(at::Tensor& input, } template -std::vector ds_mlp_gemm_int8(at::Tensor& input, - at::Tensor& residual, - at::Tensor& input_bias, - at::Tensor& weight, - at::Tensor& bias, - at::Tensor& gamma, - at::Tensor& beta, - const float epsilon, - at::Tensor& q_scale, - int groups, - bool preLayerNorm) +std::vector ds_rms_mlp_gemm(at::Tensor& input, + at::Tensor& residual, + at::Tensor& weight_interm, + at::Tensor& weight_out, + at::Tensor& gamma, + const float epsilon, + at::Tensor& q_scale, + at::Tensor& q_scale1, + bool q_int8, + int activation_type, + bool transposed_mode) { - auto input_cont = input.contiguous(); + const int bsz = input.size(0) * input.size(1); + const size_t input_neurons = input.size(2); + const size_t mlp_1_out_neurons = transposed_mode ? weight_interm.size(0) + : weight_interm.size(1); + const size_t mlp_2_in_neurons = transposed_mode ? weight_out.size(1) : weight_out.size(0); + auto options = at::TensorOptions() - .dtype(input_cont.options().dtype()) + .dtype(input.options().dtype()) .layout(at::kStrided) .device(at::kCUDA) .requires_grad(false); - auto output = at::empty({input_cont.size(0), input_cont.size(1), weight.size(1)}, options); + T* output_ptr = (T*)InferenceContext::Instance().GetWorkSpace() + torch::numel(input); + T* inp_norm_ptr = output_ptr + torch::numel(input); + T* intermediate_ptr = inp_norm_ptr + torch::numel(input); - int bsz = input_cont.size(0) * input_cont.size(1); - auto inp_norm = at::empty_like(input_cont); + auto output = at::from_blob(output_ptr, input.sizes(), options); + auto inp_norm = at::from_blob(inp_norm_ptr, input.sizes(), options); + auto intermediate_gemm = + at::from_blob(intermediate_ptr, {input.size(0), input.size(1), mlp_1_out_neurons}, options); - auto residual_add = (preLayerNorm ? at::empty_like(input_cont) : inp_norm); - quantized_gemm(output, inp_norm, weight, q_scale, groups, 0); - launch_bias_gelu((T*)output.data_ptr(), - (T*)bias.data_ptr(), - weight.size(1), - bsz, - InferenceContext::Instance().GetCurrentStream()); + auto act_func_type = static_cast(activation_type); - return {output, residual_add}; + // RMS Norm, we'll update the residual in-place + launch_rms_norm((T*)inp_norm.data_ptr(), + (T*)residual.data_ptr(), + (const T*)input.data_ptr(), + (const T*)residual.data_ptr(), + (const T*)gamma.data_ptr(), + epsilon, + bsz, + input_neurons, + InferenceContext::Instance().GetCurrentStream()); + + if (q_int8) { + quantized_gemm(intermediate_ptr, + (T*)inp_norm.data_ptr(), + weight_interm, + q_scale, + q_scale.size(0), + bsz, + input_neurons); + } else { + float alpha = (T)1.0; + float gemm_beta = (T)0.0; + cublasSetStream(InferenceContext::Instance().GetCublasHandle(), + InferenceContext::Instance().GetCurrentStream()); + cublas_gemm_ex(InferenceContext::Instance().GetCublasHandle(), + (transposed_mode ? CUBLAS_OP_T : CUBLAS_OP_N), + CUBLAS_OP_N, + mlp_1_out_neurons, + bsz, + input_neurons, + &alpha, + &gemm_beta, + (T*)weight_interm.data_ptr(), + (T*)inp_norm.data_ptr(), + intermediate_ptr, +#ifdef __HIP_PLATFORM_HCC__ + rocblas_gemm_algo_standard); +#else + CUBLAS_GEMM_DEFAULT_TENSOR_OP); +#endif + } + + if (act_func_type == ActivationFuncType::GELU) { + launch_bias_gelu(intermediate_ptr, + (T*)nullptr, + mlp_1_out_neurons, + bsz, + InferenceContext::Instance().GetCurrentStream()); + } else if (act_func_type == ActivationFuncType::ReLU) { + launch_bias_relu(intermediate_ptr, + (T*)nullptr, + mlp_1_out_neurons, + bsz, + InferenceContext::Instance().GetCurrentStream()); + } else if (act_func_type == ActivationFuncType::GATED_GELU) { + launch_gated_activation(intermediate_ptr, + (const T*)intermediate_ptr, + (const T*)nullptr, + bsz, + mlp_1_out_neurons, + mlp_1_out_neurons, + true, + InferenceContext::Instance().GetCurrentStream()); + } else if (act_func_type == ActivationFuncType::GATED_SILU) { + launch_gated_activation(intermediate_ptr, + (const T*)intermediate_ptr, + (const T*)nullptr, + bsz, + mlp_1_out_neurons, + mlp_1_out_neurons, + false, + InferenceContext::Instance().GetCurrentStream()); + } + + if (q_int8) { + quantized_gemm(output.data_ptr(), + intermediate_ptr, + weight_out, + q_scale1, + q_scale1.size(0), + bsz, + input.size(2)); + } else { + float alpha = (T)1.0; + float gemm_beta = (T)0.0; + cublasSetStream(InferenceContext::Instance().GetCublasHandle(), + InferenceContext::Instance().GetCurrentStream()); + cublas_gemm_ex(InferenceContext::Instance().GetCublasHandle(), + (transposed_mode ? CUBLAS_OP_T : CUBLAS_OP_N), + CUBLAS_OP_N, + input_neurons, + bsz, + mlp_2_in_neurons, + &alpha, + &gemm_beta, + (T*)weight_out.data_ptr(), + intermediate_ptr, + (T*)output.data_ptr(), +#ifdef __HIP_PLATFORM_HCC__ + rocblas_gemm_algo_standard, +#else + CUBLAS_GEMM_DEFAULT_TENSOR_OP, +#endif + mlp_1_out_neurons); + } + + return {output, residual}; } template @@ -1511,10 +1687,7 @@ at::Tensor fused_gemm_gelu(at::Tensor& input, at::Tensor& bias, at::Tensor& weight_out, at::Tensor& weight_out_scale, - const float epsilon, - bool preLayerNorm, bool q_int8, - bool async_op, bool transposed_mode) { auto options = at::TensorOptions() @@ -1641,13 +1814,35 @@ at::Tensor& residual_add_bias(at::Tensor& hidden_state, return residual; } +#define DISPATCH_VECTOR_ADD(T_TYPE, C_TYPE) \ + if (a.scalar_type() == at::k##T_TYPE) { \ + launch_vector_add((C_TYPE*)(a.data_ptr()), \ + (const C_TYPE*)(a.data_ptr()), \ + (const C_TYPE*)(b.data_ptr()), \ + gamma, \ + total_elems, \ + InferenceContext::Instance().GetCurrentStream()); \ + } + +at::Tensor& _vector_add(at::Tensor& a, at::Tensor& b, float gamma) +{ + const int total_elems = a.numel(); + + DISPATCH_VECTOR_ADD(Float, float) + DISPATCH_VECTOR_ADD(Half, __half) +#ifdef BF16_AVAILABLE + DISPATCH_VECTOR_ADD(BFloat16, __nv_bfloat16) +#endif + + return a; +} + std::vector apply_rotary_pos_emb(at::Tensor& mixed_query, at::Tensor& key_layer, unsigned rotary_dim, unsigned offset, unsigned num_heads, - bool rotate_half, - bool rotate_every_two) + bool rotate_half) { auto query_cont = mixed_query.contiguous(); auto key_cont = key_layer.contiguous(); @@ -1665,8 +1860,6 @@ std::vector apply_rotary_pos_emb(at::Tensor& mixed_query, offset, num_heads, bsz, - rotate_half, - rotate_every_two, InferenceContext::Instance().GetCurrentStream(), InferenceContext::Instance().GetMaxTokenLenght()); else @@ -1678,63 +1871,33 @@ std::vector apply_rotary_pos_emb(at::Tensor& mixed_query, offset, num_heads, bsz, - rotate_half, - rotate_every_two, InferenceContext::Instance().GetCurrentStream(), InferenceContext::Instance().GetMaxTokenLenght()); return {query_cont, key_cont}; } -template -at::Tensor fused_gemm_gelu_int8(at::Tensor& input, - at::Tensor& weight, - at::Tensor& bias, - const float epsilon, - at::Tensor& q_scale, - int groups, - bool preLayerNorm) -{ - auto input_cont = input.contiguous(); - auto options = at::TensorOptions() - .dtype(input_cont.options().dtype()) - .layout(at::kStrided) - .device(at::kCUDA) - .requires_grad(false); - - auto output = at::empty({input_cont.size(0), input_cont.size(1), weight.size(1)}, options); - - int bsz = input_cont.size(0) * input_cont.size(1); - - quantized_gemm(output, input_cont, weight, q_scale, groups, 0); - launch_bias_gelu((T*)output.data_ptr(), - (T*)bias.data_ptr(), - weight.size(1), - bsz, - InferenceContext::Instance().GetCurrentStream()); - - return output; -} +#define DISPATCH_MOE_RESIDUAL(T_TYPE, C_TYPE) \ + if (moe_res.scalar_type() == torch::T_TYPE) { \ + launch_moe_res_matmul((C_TYPE*)moe_res.data_ptr(), \ + (C_TYPE*)coef.data_ptr(), \ + (C_TYPE*)output.data_ptr(), \ + M, \ + N, \ + InferenceContext::Instance().GetCurrentStream()); \ + } at::Tensor moe_res_matmul(at::Tensor& moe_res, at::Tensor& coef, at::Tensor& output) { int M = moe_res.size(0) * moe_res.size(1); int N = moe_res.size(2); InferenceContext::Instance().SynchComm(); - if (moe_res.scalar_type() == at::kFloat) { - launch_moe_res_matmul((float*)moe_res.data_ptr(), - (float*)coef.data_ptr(), - (float*)output.data_ptr(), - M, - N, - at::cuda::getCurrentCUDAStream()); - } else { - launch_moe_res_matmul<__half>((__half*)moe_res.data_ptr(), - (__half*)coef.data_ptr(), - (__half*)output.data_ptr(), - M, - N, - at::cuda::getCurrentCUDAStream()); - } + + DISPATCH_MOE_RESIDUAL(kFloat, float) + DISPATCH_MOE_RESIDUAL(kHalf, __half) +#ifdef BF16_AVAILABLE + DISPATCH_MOE_RESIDUAL(kBFloat16, __nv_bfloat16) +#endif + return output; } @@ -1747,27 +1910,26 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) m.def("softmax_context_int8", &ds_softmax_context1<__half>, "DeepSpeed attention with int8 (CUDA)"); - m.def("bias_geglu", &ds_bias_geglu, "DeepSpeed Bias GEGLU (CUDA)"); + + // The following functions handle type dispatching internally + m.def("gated_activation", &ds_gated_activation, "DeepSpeed Bias GEGLU (CUDA)"); m.def("layer_norm", &ds_layer_norm, "DeepSpeed layer norm (CUDA)"); m.def( "_layer_norm_residual", &ds_layer_norm_residual, "DeepSpeed layer norm + residual (CUDA)"); m.def("layer_norm_residual_store_pre_ln_res", &ds_layer_norm_residual_store_pre_ln_res, "DeepSpeed layer norm + store pre Layernorm residual (CUDA)"); - m.def("qkv_gemm_int8", &ds_qkv_gemm_int8<__half>, "DeepSpeed qkv gemm with int8 (CUDA)"); - m.def("mlp_gemm_int8", &ds_mlp_gemm_int8<__half>, "DeepSpeed mlp with int8 (CUDA)"); - m.def("vector_matmul_int8", - &ds_vector_matmul_int8<__half>, - "DeepSpeed vector-MM with int8 (CUDA)"); - m.def("linear_layer_int8", - &ds_linear_layer_int8<__half>, - "DeepSpeed linear_layer with int8 (CUDA)"); + m.def("rms_norm", &ds_rms_norm, "DeepSpeed rms norm (CUDA)"); + m.def("pre_rms_norm", &ds_pre_rms_norm, "DeepSpeed pre rms norm (CUDA)"); + m.def("_vector_add", &_vector_add, "DeepSpeed vector add (CUDA)"); m.def("apply_rotary_pos_emb", &apply_rotary_pos_emb, "DeepSpeed mlp with fp16 (CUDA)"); m.def("moe_res_matmul", &moe_res_matmul, "DeepSpeed moe residual matmul (CUDA)"); m.def("reset_cache", &reset_cache, "Reset Cache for generation tasks"); m.def("release_workspace", &ds_release_workspace, "DeepSpeed Release Workspace"); m.def("retake_workspace", &ds_retake_workspace, "DeepSpeed Retake Workspace"); + // The following functions are templated and need to be explicitly instantiated and bound + // to different python methods #define DEF_OPS(_name, _dtype) \ m.def("softmax_" #_name, &ds_softmax<_dtype>, "DeepSpeed SoftMax with " #_name " (CUDA)"); \ m.def("softmax_context_" #_name, \ @@ -1780,7 +1942,13 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) &ds_bias_residual<_dtype>, \ "DeepSpeed residual-bias add with " #_name " (CUDA)"); \ m.def("qkv_gemm_" #_name, &ds_qkv_gemm<_dtype>, "DeepSpeed qkv gemm with " #_name " (CUDA)"); \ + m.def("rms_qkv_gemm_" #_name, \ + &ds_rms_qkv<_dtype>, \ + "DeepSpeed rms qkv gemm with " #_name " (CUDA)"); \ m.def("mlp_gemm_" #_name, &ds_mlp_gemm<_dtype>, "DeepSpeed mlp with " #_name " (CUDA)"); \ + m.def("rms_mlp_gemm_" #_name, \ + &ds_rms_mlp_gemm<_dtype>, \ + "DeepSpeed rms mlp gemm with " #_name " (CUDA)"); \ m.def("vector_matmul_" #_name, \ &ds_vector_matmul<_dtype>, \ "DeepSpeed vector-MM with " #_name " (CUDA)"); \ diff --git a/csrc/transformer/inference/csrc/relu.cu b/csrc/transformer/inference/csrc/relu.cu index fd38d1a88d4b..40926b776cf2 100644 --- a/csrc/transformer/inference/csrc/relu.cu +++ b/csrc/transformer/inference/csrc/relu.cu @@ -28,7 +28,8 @@ __global__ void fused_bias_relu(T* input, const T* bias, int total_count, int in T data[values_per_access]; T data_bias[values_per_access]; mem_access::load_global(data, input + offset); - mem_access::load_global(data_bias, bias + (offset % intermediate_size)); + mem_access::load_global( + data_bias, bias + (offset % intermediate_size), bias != nullptr); #pragma unroll for (int i = 0; i < values_per_access; i++) { @@ -60,12 +61,11 @@ void launch_bias_relu(T* input, input, bias, total_count, intermediate_size); } -template void launch_bias_relu(float*, const float*, int, int, cudaStream_t); +#define INSTANTIATE_LAUNCH_BIAS_RELU(T) \ + template void launch_bias_relu(T*, const T*, int, int, cudaStream_t); + +INSTANTIATE_LAUNCH_BIAS_RELU(float) #ifdef BF16_AVAILABLE -template void launch_bias_relu<__nv_bfloat16>(__nv_bfloat16*, - const __nv_bfloat16*, - int, - int, - cudaStream_t); +INSTANTIATE_LAUNCH_BIAS_RELU(__nv_bfloat16) #endif -template void launch_bias_relu<__half>(__half*, const __half*, int, int, cudaStream_t); +INSTANTIATE_LAUNCH_BIAS_RELU(__half) diff --git a/csrc/transformer/inference/csrc/rms_norm.cu b/csrc/transformer/inference/csrc/rms_norm.cu new file mode 100644 index 000000000000..5f72a4193752 --- /dev/null +++ b/csrc/transformer/inference/csrc/rms_norm.cu @@ -0,0 +1,263 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "conversion_utils.h" +#include "ds_kernel_utils.h" +#include "inference_cuda_layers.h" +#include "memory_access_utils.h" +#include "reduction_utils.h" + +namespace cg = cooperative_groups; +using rop = reduce::ROpType; + +namespace rms { +constexpr int granularity = 16; +} // namespace rms + +template +__global__ void rms_norm(T* output, const T* vals, const T* gamma, float epsilon, int elems_per_row) +{ + constexpr int T_per_load = rms::granularity / sizeof(T); + + cg::thread_block tb = cg::this_thread_block(); + cg::thread_block_tile warp = cg::tiled_partition(tb); + + // X-dimension of the block + const int block_offset = (tb.group_index().x * (maxThreads / threadsPerGroup) * elems_per_row) + + (tb.thread_index().y * elems_per_row); + const int thread_offset = tb.thread_index().x * T_per_load; + const int base_offset = block_offset + thread_offset; + const int stride = blockDim.x * T_per_load; + + float var_sum = reduce::init(); + + const T* input_base = vals + base_offset; + + T local_buffer[UNROLL * T_per_load]; + +#pragma unroll + for (int i = 0; i < UNROLL; i++) { + T* iteration_buffer = local_buffer + (i * T_per_load); + + mem_access::load_global(iteration_buffer, + input_base + (i * stride), + thread_offset + (i * stride) < elems_per_row); + +#pragma unroll + for (int j = 0; j < T_per_load; j++) { + float up_cast = conversion::to(iteration_buffer[j]); + float sq_val = up_cast * up_cast; + var_sum = reduce::element(var_sum, sq_val); + } + } + + reduce::partitioned_block(tb, warp, var_sum); + const float var = var_sum / elems_per_row; + const T denom = conversion::to(__frsqrt_rn(var + epsilon)); + + T* block_output = output + block_offset; + +#pragma unroll + for (int i = 0; i < UNROLL; i++) { + T* iteration_buffer = local_buffer + (i * T_per_load); + const int iter_idx = i * stride + thread_offset; + const bool do_loads = (iter_idx < elems_per_row); + + T gamma_local[T_per_load]; + + mem_access::load_global(gamma_local, gamma + iter_idx, do_loads); + +#pragma unroll + for (int j = 0; j < T_per_load; j++) { + iteration_buffer[j] *= denom; + iteration_buffer[j] *= gamma_local[j]; + } + + if (do_loads) { + mem_access::store_global(block_output + iter_idx, iteration_buffer); + } + } +} + +template +__global__ void pre_rms_norm(T* output, + T* res_out, + const T* vals, + const T* residual, + const T* gamma, + float epsilon, + int elems_per_row) +{ + constexpr int T_per_load = rms::granularity / sizeof(T); + + cg::thread_block tb = cg::this_thread_block(); + cg::thread_block_tile warp = cg::tiled_partition(tb); + + // X-dimension of the block + const int block_offset = (tb.group_index().x * (maxThreads / threadsPerGroup) * elems_per_row) + + (tb.thread_index().y * elems_per_row); + const int thread_offset = tb.thread_index().x * T_per_load; + const int base_offset = block_offset + thread_offset; + const int stride = blockDim.x * T_per_load; + + float var_sum = reduce::init(); + + const T* input_base = vals + base_offset; + const T* residual_base = residual + base_offset; + T* res_output = res_out + base_offset; + + T local_buffer[UNROLL * T_per_load]; + +#pragma unroll + for (int i = 0; i < UNROLL; i++) { + T* iteration_buffer = local_buffer + (i * T_per_load); + T residual_buffer[T_per_load]; + + const int iter_offset = i * stride + thread_offset; + const bool do_loads = (iter_offset < elems_per_row); + + mem_access::load_global( + iteration_buffer, input_base + (i * stride), do_loads); + mem_access::load_global( + residual_buffer, residual_base + (i * stride), do_loads); + +#pragma unroll + for (int j = 0; j < T_per_load; j++) { + iteration_buffer[j] += residual_buffer[j]; + float vals_up_cast = conversion::to(iteration_buffer[j]); + + var_sum = reduce::element(var_sum, vals_up_cast * vals_up_cast); + } + + if (do_loads) { + mem_access::store_global(res_output + i * stride, iteration_buffer); + } + } + + reduce::partitioned_block(tb, warp, var_sum); + const float var = var_sum / elems_per_row; + const T denom = conversion::to(__frsqrt_rn(var + epsilon)); + + T* block_output = output + block_offset; + +#pragma unroll + for (int i = 0; i < UNROLL; i++) { + T* iteration_buffer = local_buffer + (i * T_per_load); + const int iter_idx = i * stride + thread_offset; + const bool do_loads = (iter_idx < elems_per_row); + + T gamma_local[T_per_load]; + + mem_access::load_global(gamma_local, gamma + iter_idx, do_loads); + +#pragma unroll + for (int j = 0; j < T_per_load; j++) { + iteration_buffer[j] *= denom; + iteration_buffer[j] *= gamma_local[j]; + } + + if (do_loads) { + mem_access::store_global(block_output + iter_idx, iteration_buffer); + } + } +} + +#define LAUNCH_RMS_NORM(UNROLL, threadsPerGroup, maxThreads) \ + rms_norm \ + <<>>(norm_output, vals, gamma, epsilon, elems_per_row); + +#define LAUNCH_PRE_RMS_NORM(UNROLL, threadsPerGroup, maxThreads) \ + pre_rms_norm<<>>( \ + norm_output, res_output, vals, residual, gamma, epsilon, elems_per_row); + +#define LAUNCH_ALL_RMS_NORM(UNROLL, threadsPerGroup, maxThreads) \ + if (pre_norm) { \ + LAUNCH_PRE_RMS_NORM(UNROLL, threadsPerGroup, maxThreads) \ + } else { \ + LAUNCH_RMS_NORM(UNROLL, threadsPerGroup, maxThreads) \ + } + +template +void launch_rms_norm(T* norm_output, + T* res_output, + const T* vals, + const T* residual, + const T* gamma, + float epsilon, + int rows, + int elems_per_row, + cudaStream_t stream) +{ + // 8 for __half, 4 for float + constexpr int T_per_load = rms::granularity / sizeof(T); + constexpr int maxThreads = 256; + constexpr int internalUnroll = sizeof(T) == 4 ? 4 : 2; + + const bool is_subblock_schedule = (elems_per_row <= 128) ? true : false; + const int h_per_step = is_subblock_schedule ? T_per_load : T_per_load * internalUnroll; + + // Scheduling concern: may be slightly faster for some inputs to assign multiple stages of + // warp-sized blocks rather than stepping up to 64/96 threads + const int one_step_threads = next_pow2((elems_per_row + h_per_step - 1) / h_per_step); + const int threads_per_group = (one_step_threads < maxThreads) ? one_step_threads : maxThreads; + + const int groups_per_block_max = + is_subblock_schedule ? (maxThreads + threads_per_group - 1) / threads_per_group : 1; + const int groups_per_block = (rows < groups_per_block_max) ? rows : groups_per_block_max; + const int groups_launch = (groups_per_block + rows - 1) / groups_per_block; + + dim3 block(threads_per_group, groups_per_block); + dim3 grid(groups_launch); + + const int elems_per_step = threads_per_group * h_per_step; + const int external_unRoll = (elems_per_row + elems_per_step - 1) / elems_per_step; + + bool pre_norm = (residual == nullptr) ? false : true; + + if (is_subblock_schedule) { + // <=128 + if (threads_per_group == 1) { + LAUNCH_ALL_RMS_NORM(1, 1, maxThreads); + } else if (threads_per_group == 2) { + LAUNCH_ALL_RMS_NORM(1, 2, maxThreads); + } else if (threads_per_group == 4) { + LAUNCH_ALL_RMS_NORM(1, 4, maxThreads); + } else if (threads_per_group == 8) { + LAUNCH_ALL_RMS_NORM(1, 8, maxThreads); + } else if (threads_per_group == 16) { + LAUNCH_ALL_RMS_NORM(1, 16, maxThreads); + } + } else if (external_unRoll == 1) { + // 129 - 4096 elems + // (this can launch with 1-7 warps as well) + LAUNCH_ALL_RMS_NORM(1 * internalUnroll, maxThreads, maxThreads); + } else if (external_unRoll == 2) { + // 4097 - 8192 elems + LAUNCH_ALL_RMS_NORM(2 * internalUnroll, maxThreads, maxThreads); + } else if (external_unRoll == 3) { + // 8193 - 12288 elems + LAUNCH_ALL_RMS_NORM(3 * internalUnroll, maxThreads, maxThreads); + } else if (external_unRoll == 4) { + // 12289 - 16384 elems + LAUNCH_ALL_RMS_NORM(4 * internalUnroll, maxThreads, maxThreads); + } +} + +#define INSTANTIATE_LAUNCH_RMS_NORM(T) \ + template void launch_rms_norm(T * norm_output, \ + T * res_output, \ + const T* vals, \ + const T* residual, \ + const T* gamma, \ + float epsilon, \ + int rows, \ + int elems_per_row, \ + cudaStream_t stream); + +INSTANTIATE_LAUNCH_RMS_NORM(float) +INSTANTIATE_LAUNCH_RMS_NORM(__half) +#ifdef BF16_AVAILABLE +INSTANTIATE_LAUNCH_RMS_NORM(__nv_bfloat16) +#endif diff --git a/csrc/transformer/inference/csrc/softmax.cu b/csrc/transformer/inference/csrc/softmax.cu index 7befdfd51497..c0e98e10b5ad 100644 --- a/csrc/transformer/inference/csrc/softmax.cu +++ b/csrc/transformer/inference/csrc/softmax.cu @@ -487,58 +487,29 @@ void launch_attn_softmax_v2(T* vals, throw std::runtime_error("Unsupport Seq_Length!"); } -template void launch_attn_softmax_v2(float* vals, - float* mask, - float* alibi, - float layer_scale, - bool triangular, - bool recompute, - bool local_attention, - int window_size, - int batch_size, - int heads, - int num_seq, - int sequence_length, - int head_offset, - int mask_stride, - int mp_size, - cudaStream_t stream); - +#define INSTANTIATE_LAUNCH_ATTN_SOFTMAX_V2(T) \ + template void launch_attn_softmax_v2(T* vals, \ + T* mask, \ + T* alibi, \ + float layer_scale, \ + bool triangular, \ + bool recompute, \ + bool local_attention, \ + int window_size, \ + int batch_size, \ + int heads, \ + int num_seq, \ + int sequence_length, \ + int head_offset, \ + int mask_stride, \ + int mp_size, \ + cudaStream_t stream); + +INSTANTIATE_LAUNCH_ATTN_SOFTMAX_V2(float); #ifdef BF16_AVAILABLE -template void launch_attn_softmax_v2(__nv_bfloat16* vals, - __nv_bfloat16* mask, - __nv_bfloat16* alibi, - float layer_scale, - bool triangular, - bool recompute, - bool local_attention, - int window_size, - int batch_size, - int heads, - int num_seq, - int sequence_length, - int head_offset, - int mask_stride, - int mp_size, - cudaStream_t stream); +INSTANTIATE_LAUNCH_ATTN_SOFTMAX_V2(__nv_bfloat16); #endif - -template void launch_attn_softmax_v2(__half* vals, - __half* mask, - __half* alibi, - float layer_scale, - bool triangular, - bool recompute, - bool local_attention, - int window_size, - int batch_size, - int heads, - int num_seq, - int sequence_length, - int head_offset, - int mask_stride, - int mp_size, - cudaStream_t stream); +INSTANTIATE_LAUNCH_ATTN_SOFTMAX_V2(__half); #define DEF_ATTN_SOFTMAX_V2_HALF(_iter) \ template __global__ void attn_softmax_v2<__half, _iter>(__half * vals, \ diff --git a/csrc/transformer/inference/csrc/transform.cu b/csrc/transformer/inference/csrc/transform.cu index 00d919039b5d..650f286a8f03 100644 --- a/csrc/transformer/inference/csrc/transform.cu +++ b/csrc/transformer/inference/csrc/transform.cu @@ -240,43 +240,29 @@ void launch_bias_add_transform_0213(T* output, max_out_tokens); } +#define INSTANTIATE_LAUNCH_BIAS_ADD_TRANSFORM_0213(T) \ + template void launch_bias_add_transform_0213(T*, \ + T*, \ + T*, \ + const T*, \ + const T*, \ + int, \ + int, \ + unsigned, \ + int, \ + int, \ + int, \ + int, \ + bool, \ + bool, \ + cudaStream_t, \ + int, \ + int) + #ifdef BF16_AVAILABLE -template void launch_bias_add_transform_0213(__nv_bfloat16* output, - __nv_bfloat16* k_cache, - __nv_bfloat16* v_cache, - const __nv_bfloat16* vals, - const __nv_bfloat16* bias, - int batch_size, - int seq_length, - unsigned seq_offset, - int all_tokens, - int hidden_dim, - int heads, - int rotary_dim, - bool rotate_half, - bool rotate_every_two, - cudaStream_t stream, - int trans_count, - int max_out_tokens); +INSTANTIATE_LAUNCH_BIAS_ADD_TRANSFORM_0213(__nv_bfloat16); #endif - -template void launch_bias_add_transform_0213(__half* output, - __half* k_cache, - __half* v_cache, - const __half* vals, - const __half* bias, - int batch_size, - int seq_length, - unsigned seq_offset, - int all_tokens, - int hidden_dim, - int heads, - int rotary_dim, - bool rotate_half, - bool rotate_every_two, - cudaStream_t stream, - int trans_count, - int max_out_tokens); +INSTANTIATE_LAUNCH_BIAS_ADD_TRANSFORM_0213(__half); // Bias add @@ -368,28 +354,15 @@ void launch_pad_add_transform_0213(T* output, output, vals, hidden_dim, seq_length, padded_seq_len, heads, padded_head_size >> 3); } +#define INSTANTIATE_LAUNCH_PAD_ADD_TRANSFORM_0213_SIMPLE(T) \ + template void launch_pad_add_transform_0213( \ + T*, const T*, int, int, int, int, int, int, cudaStream_t); + +INSTANTIATE_LAUNCH_PAD_ADD_TRANSFORM_0213_SIMPLE(__half); #ifdef BF16_AVAILABLE -template void launch_pad_add_transform_0213(__nv_bfloat16* output, - const __nv_bfloat16* vals, - int batch_size, - int hidden_dim, - int seq_length, - int padded_seq_len, - int heads, - int padded_head_size, - cudaStream_t stream); +INSTANTIATE_LAUNCH_PAD_ADD_TRANSFORM_0213_SIMPLE(__nv_bfloat16); #endif -template void launch_pad_add_transform_0213(__half* output, - const __half* vals, - int batch_size, - int hidden_dim, - int seq_length, - int padded_seq_len, - int heads, - int padded_head_size, - cudaStream_t stream); - // Bias add template __global__ void bias_add_transform_0213(T* output, @@ -571,22 +544,6 @@ __global__ void bias_add_transform_0213_v2(T* output, } } -template __global__ void bias_add_transform_0213_v2(__half* output, - const __half* vals, - const __half* bias, - int hidden_dim, - int seq_length, - int heads); - -#ifdef BF16_AVAILABLE -template __global__ void bias_add_transform_0213_v2(__nv_bfloat16* output, - const __nv_bfloat16* vals, - const __nv_bfloat16* bias, - int hidden_dim, - int seq_length, - int heads); -#endif - template __global__ void transform4d_0213(T* out, const T* in, @@ -707,20 +664,6 @@ __global__ void transform4d_0213_v2(T* out, const T* in, int heads, int seq_leng } } -#ifdef BF16_AVAILABLE -template __global__ void transform4d_0213_v2(__nv_bfloat16* out, - const __nv_bfloat16* in, - int heads, - int seq_length, - int hidden_dim); -#endif - -template __global__ void transform4d_0213_v2(__half* out, - const __half* in, - int heads, - int seq_length, - int hidden_dim); - // 3 * [B A S N] - > [B S C*H] template <> void launch_transform4d_0213(float* out, @@ -757,22 +700,10 @@ void launch_transform4d_0213(T* out, out, in, heads, seq_length, hidden_dim, head_ext); } +#define INSTANTIATE_2B_LAUNCH_TRANSFORM4D(T) \ + template void launch_transform4d_0213(T*, const T*, int, int, int, int, cudaStream_t, int); + +INSTANTIATE_2B_LAUNCH_TRANSFORM4D(__half) #ifdef BF16_AVAILABLE -template void launch_transform4d_0213(__nv_bfloat16* out, - const __nv_bfloat16* in, - int batch_size, - int heads, - int seq_length, - int hidden_dim, - cudaStream_t stream, - int trans_count); +INSTANTIATE_2B_LAUNCH_TRANSFORM4D(__nv_bfloat16) #endif - -template void launch_transform4d_0213(__half* out, - const __half* in, - int batch_size, - int heads, - int seq_length, - int hidden_dim, - cudaStream_t stream, - int trans_count); diff --git a/csrc/transformer/inference/includes/inference_context.h b/csrc/transformer/inference/includes/inference_context.h index f7bbcad91e2a..25051d167fc7 100644 --- a/csrc/transformer/inference/includes/inference_context.h +++ b/csrc/transformer/inference/includes/inference_context.h @@ -136,7 +136,7 @@ class InferenceContext { if (_max_seq_len < min_out_tokens) { printf( - "Allocatable workspace available (%d tokens) is less than minimum requested " + "Allocatable workspace available (%ld tokens) is less than minimum requested " "workspace (%d tokens)\n", _max_seq_len, min_out_tokens); diff --git a/csrc/transformer/inference/includes/inference_cublas_wrappers.h b/csrc/transformer/inference/includes/inference_cublas_wrappers.h index 38a0f7ab66d2..15e9e4aad3f6 100644 --- a/csrc/transformer/inference/includes/inference_cublas_wrappers.h +++ b/csrc/transformer/inference/includes/inference_cublas_wrappers.h @@ -28,7 +28,8 @@ int cublas_gemm_ex(rocblas_handle handle, const float* A, const float* B, float* C, - rocblas_gemm_algo algo) + rocblas_gemm_algo algo, + int b_stride = -1) #else int cublas_gemm_ex(cublasHandle_t handle, cublasOperation_t transa, @@ -41,9 +42,11 @@ int cublas_gemm_ex(cublasHandle_t handle, const float* A, const float* B, float* C, - cublasGemmAlgo_t algo) + cublasGemmAlgo_t algo, + int b_stride = -1) #endif { + const int ldb = (b_stride == -1) ? ((transb == CUBLAS_OP_N) ? k : n) : b_stride; #ifdef __HIP_PLATFORM_HCC__ rocblas_status status = rocblas_gemm_ex(handle, transa, @@ -57,7 +60,7 @@ int cublas_gemm_ex(cublasHandle_t handle, (transa == rocblas_operation_none) ? m : k, (const void*)B, rocblas_datatype_f32_r, - (transb == rocblas_operation_none) ? k : n, + ldb, (const void*)beta, C, rocblas_datatype_f32_r, @@ -82,7 +85,7 @@ int cublas_gemm_ex(cublasHandle_t handle, (transa == CUBLAS_OP_N) ? m : k, (const void*)B, CUDA_R_32F, - (transb == CUBLAS_OP_N) ? k : n, + ldb, (const void*)beta, C, CUDA_R_32F, @@ -120,7 +123,8 @@ int cublas_gemm_ex(rocblas_handle handle, const T* A, const T* B, T* C, - rocblas_gemm_algo algo) + rocblas_gemm_algo algo, + int b_stride = -1) #else int cublas_gemm_ex(cublasHandle_t handle, cublasOperation_t transa, @@ -133,9 +137,11 @@ int cublas_gemm_ex(cublasHandle_t handle, const T* A, const T* B, T* C, - cublasGemmAlgo_t algo) + cublasGemmAlgo_t algo, + int b_stride = -1) #endif { + const int ldb = (b_stride == -1) ? ((transb == CUBLAS_OP_N) ? k : n) : b_stride; #ifdef __HIP_PLATFORM_HCC__ constexpr auto rocblas_dtype_16 = std::is_same::value ? rocblas_datatype_f16_r : rocblas_datatype_bf16_r; @@ -151,7 +157,7 @@ int cublas_gemm_ex(cublasHandle_t handle, (transa == rocblas_operation_none) ? m : k, (const void*)B, rocblas_dtype_16, - (transb == rocblas_operation_none) ? k : n, + ldb, (const void*)beta, (void*)C, rocblas_dtype_16, @@ -177,7 +183,7 @@ int cublas_gemm_ex(cublasHandle_t handle, (transa == CUBLAS_OP_N) ? m : k, (const void*)B, cublas_dtype_16, - (transb == CUBLAS_OP_N) ? k : n, + ldb, (const void*)beta, (void*)C, cublas_dtype_16, diff --git a/csrc/transformer/inference/includes/inference_cuda_layers.h b/csrc/transformer/inference/includes/inference_cuda_layers.h index e7e835a2dddf..8ba8c1c3e22c 100644 --- a/csrc/transformer/inference/includes/inference_cuda_layers.h +++ b/csrc/transformer/inference/includes/inference_cuda_layers.h @@ -52,11 +52,13 @@ void launch_bias_gelu(T* input, cudaStream_t stream); template -void launch_fused_bias_geglu(T* output, +void launch_gated_activation(T* output, const T* activation, const T* bias, int rows, + int output_stride, int elems_per_row, + bool use_gelu, cudaStream_t stream); // Fused bias add with relu activation @@ -117,6 +119,17 @@ void launch_fused_residual_ln_store_pre_ln_res(T* norm_output, int elems_per_row, cudaStream_t stream); +template +void launch_rms_norm(T* norm_output, + T* res_output, + const T* vals, + const T* residual, + const T* gamma, + float epsilon, + int rows, + int elems_per_row, + cudaStream_t stream); + template void launch_dequantize(T* output, const int8_t* input, @@ -155,8 +168,6 @@ void launch_apply_rotary_pos_emb(T* mixed_query, unsigned offset, unsigned num_heads, unsigned batch, - bool rotate_half, - bool rotate_every_two, cudaStream_t stream, int max_out_tokens); @@ -224,3 +235,11 @@ void launch_pad_add_transform_0213(T* output, int heads, int padded_head_size, cudaStream_t stream); + +template +void launch_vector_add(T* out, + const T* a, + const T* b, + float gamma, + int num_elems, + cudaStream_t stream); diff --git a/deepspeed/inference/config.py b/deepspeed/inference/config.py index 70a67c062ad2..99b56d1ad30c 100644 --- a/deepspeed/inference/config.py +++ b/deepspeed/inference/config.py @@ -16,11 +16,9 @@ class DtypeEnum(Enum): # The torch dtype must always be the first value (so we return torch.dtype) fp16 = torch.float16, "torch.float16", "fp16", "float16", "half" fp32 = torch.float32, "torch.float32", "fp32", "float32", "float" + bf16 = torch.bfloat16, "torch.bfloat16", "bf16", "bfloat16", "bfloat" int8 = torch.int8, "torch.int8", "int8" - # bf16 not supported - # bf16 = torch.bfloat16, "torch.bfloat16", "bf16", "bfloat16" - # Copied from https://stackoverflow.com/a/43210118 # Allows us to use multiple values for each Enum index and returns first # listed value when Enum is called diff --git a/deepspeed/inference/engine.py b/deepspeed/inference/engine.py index de7ca5a71197..22a97559ae73 100755 --- a/deepspeed/inference/engine.py +++ b/deepspeed/inference/engine.py @@ -344,7 +344,9 @@ def load(module, state_dict, prefix): args = (state_dict, prefix, {}, True, [], [], error_msgs) if hasattr(module, 'weight'): if 'query_key_value' in prefix: - module.weight = self.mp_replace.qkv_copy(module.weight.data, state_dict[prefix + 'weight']) + module.weight = self.mp_replace.strided_copy(module.weight.data, + state_dict[prefix + 'weight'], + num_splits=3) else: module.weight = self.mp_replace.copy(module.weight.data, state_dict[prefix + 'weight']) else: diff --git a/deepspeed/model_implementations/transformers/ds_transformer.py b/deepspeed/model_implementations/transformers/ds_transformer.py index 6ef838cea741..effa801f64ef 100644 --- a/deepspeed/model_implementations/transformers/ds_transformer.py +++ b/deepspeed/model_implementations/transformers/ds_transformer.py @@ -47,7 +47,7 @@ def __init__(self, self.config.layer_id = DeepSpeedTransformerInference.layer_id DeepSpeedTransformerInference.layer_id += 1 - data_type = torch.half if config.fp16 else torch.float + data_type = torch.half if self.config.dtype == torch.int8 else self.config.dtype global inference_cuda_module if inference_cuda_module is None: builder = InferenceBuilder() @@ -74,8 +74,8 @@ def __init__(self, self.norm_b = nn.Parameter(torch.empty(self.config.hidden_size, dtype=data_type, device=device), requires_grad=False) self.layer_past = None - self.allocate_workspace = inference_cuda_module.allocate_workspace_fp32 if (not config.fp16) else \ - inference_cuda_module.allocate_workspace_fp16 + self.allocate_workspace = inference_cuda_module.allocate_workspace_fp32 if config.dtype == torch.float32 else \ + inference_cuda_module.allocate_workspace_fp16 self._alloc_workspace = True @classmethod @@ -139,9 +139,11 @@ def forward( input = input[0] input_type = input.dtype - if (self.config.fp16 or self.config.q_int8) \ + if (self.config.dtype in [torch.float16, torch.bfloat16, torch.int8]) \ and input.dtype == torch.float: - input = input.half() + target_dtype = torch.half if self.dtype == torch.int8 else self.dtype + input = input.to(target_dtype) + with torch.no_grad(): attention_output, key, value, context_outputtn_ctx, inp_norm = \ self.attention(input, diff --git a/deepspeed/module_inject/containers/__init__.py b/deepspeed/module_inject/containers/__init__.py index 4655b29b5ba6..1dab38b73f51 100644 --- a/deepspeed/module_inject/containers/__init__.py +++ b/deepspeed/module_inject/containers/__init__.py @@ -10,6 +10,7 @@ from .gptj import DS_GPTJContainer, HFGPTJLayerPolicy from .gptneo import DS_GPTNEOContainer, HFGPTNEOLayerPolicy from .gptneox import DS_GPTNEOXContainer, GPTNEOXLayerPolicy +from .llama import DS_LLAMAContainer, LLAMALayerPolicy from .megatron_gpt import DS_MegatronGPTContainer, MegatronLayerPolicy from .megatron_gpt_moe import DS_MegatronGPTMoEContainer, MegatronMoELayerPolicy from .opt import DS_OPTContainer, HFOPTLayerPolicy diff --git a/deepspeed/module_inject/containers/base.py b/deepspeed/module_inject/containers/base.py index 0d2d6fc99ef9..e5e0129fffe7 100644 --- a/deepspeed/module_inject/containers/base.py +++ b/deepspeed/module_inject/containers/base.py @@ -5,11 +5,16 @@ # Create a container object to save model-specific tensors using the policy file above. from abc import ABC + import torch from deepspeed.ops.transformer.inference.config import DeepSpeedInferenceConfig from deepspeed.accelerator import get_accelerator +# If the intermediate size attribute is set DEFAULT_INTERMEDIATE_SIZE +# it is assumed the interemediate size is 4x the embedding dimension +DEFAULT_INTERMEDIATE_SIZE = -1 + class BaseConvolutionContainer(ABC): # not implemented @@ -32,11 +37,12 @@ def __init__(self, policy, config, model_config, layer_id, child): # configuration for models. todo: can this be moved to a pydantic model config? self.hidden_size = None + self.intermediate_size = None self.num_attention_heads = None self.mp_size = self.config.tensor_parallel.tp_size self.pre_layer_norm = self.model_config.do_layer_norm_before if \ hasattr(self.model_config, 'do_layer_norm_before') else self.policy.pre_attn_norm - self.fp16 = False + self.dtype = self.config.dtype self.attn_linear_layer = self.policy.linear_layer self.mlp_linear_layer = self.policy.linear_layer self.return_tuple = self.config.return_tuple @@ -45,6 +51,7 @@ def __init__(self, policy, config, model_config, layer_id, child): self.model_config, 'attention_layers') else False) self.window_size = getattr(self.model_config, "window_size", 1) self.mlp_act_func_type = self.policy.mlp_act_func_type + self.norm_type = self.policy.norm_type self.training_mp_size = self.config.training_mp_size self.bigscience_bloom = False self.max_out_tokens = self.config.max_out_tokens @@ -52,9 +59,7 @@ def __init__(self, policy, config, model_config, layer_id, child): self.scale_attn_by_inverse_layer_idx = getattr(self.config, "scale_attn_by_inverse_layer_idx", False) self.use_mup = self.policy.use_mup self.return_single_tuple = False - self.rotary_dim = self.model_config.rotary_dim if hasattr(self.model_config, 'rotary_dim') \ - else self.child.attention.rotary_ndims if \ - hasattr(self.child, 'attention') and hasattr(self.child.attention,'rotary_ndims') else -1 + self.rotary_dim = self.get_rotary_dim() self.mlp_after_attn = (self.rotary_dim is None or self.rotary_dim < 0) # Attention tensors @@ -83,12 +88,13 @@ def create_ds_model_config(self): self.ds_model_config = DeepSpeedInferenceConfig( hidden_size=self.hidden_size, + intermediate_size=self.intermediate_size, heads=self.num_attention_heads, layer_norm_eps=self.layernorm_epsilon, - fp16=self.fp16, + dtype=self.dtype, pre_layer_norm=self.pre_layer_norm, + norm_type=self.norm_type, mp_size=self.mp_size, - q_int8=self.quantize if hasattr(self, 'quantize') else False, return_tuple=self.return_tuple, triangular_masking=self.triangular_masking, local_attention=self.local_attention, @@ -111,27 +117,27 @@ def create_ds_model_config(self): def initialize_tensors(self, enable_training=False): # Set the tensors from policy (user module) to container (DS module) self.set_attention(*self.policy.attention(enable_training=enable_training)) - self.set_mlp(*self.policy.mlp()) + self.set_mlp(*self.policy.mlp(enable_training=enable_training)) self.set_layernorm(*self.policy.layernorm()) - self.set_lora_params(self.policy.get_lora_params()) - self.q_k_v = self.policy.get_q_k_v() - if self.q_k_v is not None: - self.set_q_k_v(*self.q_k_v) - def convert_to_required_dtype(self, dtype): + def convert_to_required_dtype(self): # Note: converting tensors to fp16 requires that we do it in-place using self.__dict__ and not make a list/dict copy - if dtype == torch.half: + if self.dtype in [torch.half, torch.bfloat16]: for k, v in self.__dict__.items(): # The list comprehension is used for MoE tensor lists if isinstance(v, list) and all((isinstance(tensor, torch.Tensor) \ or isinstance(tensor, torch.nn.Parameter)) for tensor in v): - self.__dict__[k] = [moe_tensor.half() for moe_tensor in v] + self.__dict__[k] = [moe_tensor.to(self.dtype) for moe_tensor in v] if isinstance(v, torch.Tensor) or isinstance(v, torch.nn.Parameter): - self.__dict__[k] = v.half() + self.__dict__[k] = v.to(self.dtype) - def set_dtype(self, fp16=False): - self.fp16 = fp16 + def get_rotary_dim(self): + if hasattr(self.model_config, 'rotary_dim'): + return self.model_config.rotary_dim + if hasattr(self.child, 'attention') and hasattr(self.child.attention, 'rotary_ndims'): + return self.child.attention.rotary_ndims + return -1 def set_moe(self, moe=False): self.moe = moe @@ -140,12 +146,23 @@ def set_tensor_parallel_config(self, mp_size, mp_group): self.mp_size = mp_size self.mp_group = mp_group - def set_quantization_config(self, quantize, quantizer): - self.quantize = quantize + def set_quantization_config(self, quantizer): self.quantizer = quantizer - def set_hidden_heads(self, hidden_size, num_attention_heads, epsilon): + def set_hidden_heads(self, hidden_size, num_attention_heads, epsilon, intermediate_size): + """ + Args: + hidden_size: embedding dimension of the model + num_attention_heads: number of attention heads in the model + epsilon: epsilon value for layer norm (same value used for all norms) + intermediate_size: Size of MLP projection. If `DEFAUL_INTERMEDIATE_SIZE` is passed + it is assumed to be `4 * hidden_size` + """ self.hidden_size = hidden_size + if intermediate_size == DEFAULT_INTERMEDIATE_SIZE: + self.intermediate_size = 4 * hidden_size + else: + self.intermediate_size = intermediate_size self.num_attention_heads = num_attention_heads self.layernorm_epsilon = epsilon @@ -155,17 +172,6 @@ def set_attention(self, qkvw, qkvb, dense_w, dense_b): self.dense_w = dense_w self.dense_b = dense_b - def set_lora_params(self, lora_params): - self.lora_params = lora_params - - def set_q_k_v(self, qw, qb, kw, kb, vw, vb): - self.qw = qw - self.qb = qb - self.kw = kw - self.kb = kb - self.vw = vw - self.vb = vb - def set_mlp(self, _h4h_w, _h4h_b, _4hh_w, _4hh_b): self._h4h_w = _h4h_w self._h4h_b = _h4h_b @@ -193,177 +199,59 @@ def mlp_quantization(self): self.module.mlp.inter_w = self.quantizer.quantize(self.module.mlp.inter_w) self.module.mlp.output_w = self.quantizer.quantize(self.module.mlp.output_w) - def apply_tensor_parallelism(self, mp_replace=None, mp_group=None, tp_size=None): - reversed_dim = False - if mp_replace is None: - from deepspeed.module_inject import ReplaceWithTensorSlicing - mp_replace = ReplaceWithTensorSlicing(mp_group=mp_group, mp_size=tp_size, out_dim=0, in_dim=1) - reversed_dim = True + def apply_tensor_parallelism(self, mp_replace): # setup the new Attention module - if self.module.attention.attn_qkvw is None: - self.attention_q_k_v_mp(mp_replace, reversed_dim=reversed_dim) - else: - self.attention_qkv_mp(mp_replace, reversed_dim=reversed_dim) - self.attention_o_mp(mp_replace, reversed_dim=reversed_dim) + self.attention_qkv_mp(mp_replace) + self.attention_o_mp(mp_replace) # setup the new MLP module - self.mlp_inter_mp(mp_replace, reversed_dim=reversed_dim) - self.mlp_output_mp(mp_replace, reversed_dim=reversed_dim) + self.mlp_inter_mp(mp_replace) + self.mlp_output_mp(mp_replace) # Apply weight quantization + # TODO(cmikeh2): Re-enable this once verified #self.apply_weight_quantization() def attention_qkv_mp(self, mp_replace, reversed_dim=False): - if reversed_dim: - self.module.attention.attn_qkvw = mp_replace.qkv_copy( - self.module.attention.attn_qkvw[:self.qkvw.shape[0] // mp_replace.mp_size], - self.qkvw, - int8=reversed_dim) - self.module.attention.attn_qkvb = mp_replace.qkv_copy( - self.module.attention.attn_qkvb[:self.qkvw.shape[0] // mp_replace.mp_size], - self.qkvb, - int8=reversed_dim) - else: - self.module.attention.attn_qkvw = mp_replace.qkv_copy(self.module.attention.attn_qkvw, + self.module.attention.attn_qkvw = mp_replace.strided_copy(self.module.attention.attn_qkvw, self.qkvw, + num_splits=3, int8=reversed_dim) - self.module.attention.attn_qkvb = mp_replace.qkv_copy(self.module.attention.attn_qkvb, + self.module.attention.attn_qkvb = mp_replace.strided_copy(self.module.attention.attn_qkvb, self.qkvb, + num_splits=3, int8=reversed_dim) - def attention_q_k_v_mp(self, mp_replace, reversed_dim=False): - self.module.attention.attn_qw = mp_replace.copy(self.module.attention.attn_qw[:self.qw.shape[0] // - mp_replace.mp_size], - self.qw, - int8=reversed_dim, - allocat_tensor=reversed_dim) - self.module.attention.attn_kw = mp_replace.copy(self.module.attention.attn_kw[:self.qw.shape[0] // - mp_replace.mp_size], - self.kw, - int8=reversed_dim, - allocat_tensor=reversed_dim) - self.module.attention.attn_vw = mp_replace.copy(self.module.attention.attn_vw[:self.qw.shape[0] // - mp_replace.mp_size], - self.vw, - int8=reversed_dim, - allocat_tensor=reversed_dim) - self.module.attention.attn_qb = mp_replace.copy( - self.module.attention.attn_qb[:self.qw.shape[0] // mp_replace.mp_size], - self.qb, - int8=reversed_dim, - allocat_tensor=reversed_dim) if self.module.attention.attn_qb is not None else None - self.module.attention.attn_kb = mp_replace.copy( - self.module.attention.attn_kb[:self.qw.shape[0] // mp_replace.mp_size], - self.kb, - int8=reversed_dim, - allocat_tensor=reversed_dim) if self.module.attention.attn_kb is not None else None - self.module.attention.attn_vb = mp_replace.copy( - self.module.attention.attn_vb[:self.qw.shape[0] // mp_replace.mp_size], - self.vb, - int8=reversed_dim, - allocat_tensor=reversed_dim) if self.module.attention.attn_vb is not None else None - def attention_o_mp(self, mp_replace, reversed_dim=False): - if reversed_dim: - self.module.attention.attn_ow = mp_replace.copy(self.module.attention.attn_ow[:, :self.dense_w.shape[1] // - mp_replace.mp_size], - self.dense_w, - int8=reversed_dim, - allocat_tensor=reversed_dim) - else: - self.module.attention.attn_ow = mp_replace.copy(self.module.attention.attn_ow, - self.dense_w, - int8=reversed_dim) + self.module.attention.attn_ow = mp_replace.copy(self.module.attention.attn_ow, self.dense_w, int8=reversed_dim) self.module.attention.attn_ob = mp_replace.copy(self.module.attention.attn_ob, self.dense_b, int8=reversed_dim, - allocat_tensor=reversed_dim) + allocate_tensor=reversed_dim) def mlp_inter_mp(self, mp_replace, reversed_dim=False): - if reversed_dim: - self.module.mlp.inter_w = mp_replace.copy(self.module.mlp.inter_w[:self._h4h_w.shape[0] // - mp_replace.mp_size], - self._h4h_w, - int8=reversed_dim, - allocat_tensor=reversed_dim) - self.module.mlp.inter_b = mp_replace.copy( - self.module.mlp.inter_b[:self._h4h_w.shape[0] // mp_replace.mp_size], - self._h4h_b, - int8=reversed_dim, - allocat_tensor=reversed_dim) if self.module.mlp.inter_b is not None else None - else: - self.module.mlp.inter_w = mp_replace.copy(self.module.mlp.inter_w, self._h4h_w, int8=reversed_dim) - self.module.mlp.inter_b = mp_replace.copy(self.module.mlp.inter_b, self._h4h_b, int8=reversed_dim) + self.module.mlp.inter_w = mp_replace.copy(self.module.mlp.inter_w, self._h4h_w, int8=reversed_dim) + self.module.mlp.inter_b = mp_replace.copy(self.module.mlp.inter_b, self._h4h_b, int8=reversed_dim) def mlp_output_mp(self, mp_replace, reversed_dim=False): - if reversed_dim: - self.module.mlp.output_w = mp_replace.copy(self.module.mlp.output_w[:, :self._4hh_w.shape[1] // - mp_replace.mp_size], - self._4hh_w, - int8=reversed_dim, - allocat_tensor=reversed_dim) - else: - self.module.mlp.output_w = mp_replace.copy(self.module.mlp.output_w, self._4hh_w, int8=reversed_dim) + self.module.mlp.output_w = mp_replace.copy(self.module.mlp.output_w, self._4hh_w, int8=reversed_dim) self.module.mlp.output_b = mp_replace.copy(self.module.mlp.output_b, self._4hh_b, int8=reversed_dim, - allocat_tensor=reversed_dim) - - def release_qkv(self): - del self.module.attention.attn_qkvw - del self.module.attention.attn_qkvb - self.module.attention.attn_qkvw = self.qkvw - self.module.attention.attn_qkvb = self.qkvb - if self.module.attention.attn_qw is not None: - qkv_data = [self.module.attention.attn_qw.data, \ - self.module.attention.attn_qb.data if self.module.attention.attn_qb is not None else None, \ - self.module.attention.attn_kw.data, \ - self.module.attention.attn_kb.data if self.module.attention.attn_kb is not None else None, \ - self.module.attention.attn_vw.data, \ - self.module.attention.attn_vb.data if self.module.attention.attn_vb is not None else None] - for data in qkv_data: - del data - - self.module.attention.attn_qw = self.qw - self.module.attention.attn_qb = self.qb - self.module.attention.attn_kw = self.kw - self.module.attention.attn_kb = self.kb - self.module.attention.attn_vw = self.vw - self.module.attention.attn_vb = self.vb - - def release_memory(self): - self.release_qkv() - del self.module.attention.attn_ow - del self.module.attention.attn_ob - self.module.attention.attn_ow = self.dense_w - self.module.attention.attn_ob = self.dense_b - del self.module.mlp.inter_w - del self.module.mlp.inter_b - del self.module.mlp.output_w - del self.module.mlp.output_b - self.module.mlp.inter_w = self._h4h_w - self.module.mlp.inter_b = self._h4h_b - self.module.mlp.output_w = self._4hh_w - self.module.mlp.output_b = self._4hh_b + allocate_tensor=reversed_dim) def copy_data_to_new_module(self): - if self.attn_nw is None: - self.module.mlp.attn_nw = self.attn_nw - self.module.mlp.attn_nb = self.attn_nb - else: - self.module.mlp.attn_nw.data.copy_(self.attn_nw.to(get_accelerator().current_device_name())) - self.module.mlp.attn_nb.data.copy_(self.attn_nb.to(get_accelerator().current_device_name())) - - self.module.norm_w.data.copy_(self.input_nw.to(get_accelerator().current_device_name())) - self.module.norm_b.data.copy_(self.input_nb.to(get_accelerator().current_device_name())) - - def align_merged_qkv(self): - if hasattr(self, '_align_merged_qkv'): - self._align_merged_qkv() - - def partition_merged_qkv(self): - if hasattr(self, '_partition_merged_qkv'): - self._partition_merged_qkv() + params = { + self.module.mlp.attn_nw: self.attn_nw, + self.module.mlp.attn_nb: self.attn_nb, + self.module.norm_w: self.input_nw, + self.module.norm_b: self.input_nb + } + for dst, src in params.items(): + if src is None: + dst = src + else: + dst.data.copy_(src.to(get_accelerator().current_device_name())) def transpose(self): self.transpose_attention() @@ -386,109 +274,21 @@ def transpose_impl(self, data): data.to(get_accelerator().current_device_name()) return data - def reset_qkv_experimental(self): - if self.module.attention.attn_qkvw is None: - self.module.attention.attn_qkvw = torch.empty(self.qw.shape[0] * 3, - self.qw.shape[0], - dtype=self.qw.dtype, - device=self.qw.device) - self.module.attention.attn_qkvb = torch.empty(self.qw.shape[0] * 3, - dtype=self.qw.dtype, - device=self.qw.device) - self.module.attention.attn_qkvw.data[:self.qw.shape[0]] = self.qw.data - self.module.attention.attn_qkvb.data[:self.qw.shape[0]] = self.qb.data - self.module.attention.attn_qkvw.data[self.qw.shape[0]:2 * self.qw.shape[0]] = self.kw.data - self.module.attention.attn_qkvb.data[self.qw.shape[0]:2 * self.qw.shape[0]] = self.kb.data - self.module.attention.attn_qkvw.data[2 * self.qw.shape[0]:] = self.vw.data - self.module.attention.attn_qkvb.data[2 * self.qw.shape[0]:] = self.vb.data - - qkv_data = [self.qw.data, \ - self.qb.data, \ - self.kw.data, \ - self.kb.data, \ - self.vw.data, \ - self.vb.data] - - self.qw.data = self.module.attention.attn_qkvw.data[:self.qw.shape[0]] - self.qb.data = self.module.attention.attn_qkvb.data[:self.qw.shape[0]] - self.kw.data = self.module.attention.attn_qkvw.data[self.qw.shape[0]:2 * self.qw.shape[0]] - self.kb.data = self.module.attention.attn_qkvb.data[self.qw.shape[0]:2 * self.qw.shape[0]] - self.vw.data = self.module.attention.attn_qkvw.data[2 * self.qw.shape[0]:] - self.vb.data = self.module.attention.attn_qkvb.data[2 * self.qw.shape[0]:] - - for data in qkv_data: - del data - - def reset_qkv(self): - self.qkvw.data[:self.qw.shape[0]] = self.qw.data - self.qkvw.data[self.qw.shape[0]:2 * self.qw.shape[0]] = self.kw.data - self.qkvw.data[2 * self.qw.shape[0]:] = self.vw.data - if self.qkvb is not None: - self.qkvb.data[:self.qw.shape[0]] = self.qb.data - self.qkvb.data[self.qw.shape[0]:2 * self.qw.shape[0]] = self.kb.data - self.qkvb.data[2 * self.qw.shape[0]:] = self.vb.data - - qkv_data = [self.qw.data, \ - self.qb.data if self.qb is not None else None, \ - self.kw.data, \ - self.kb.data if self.kb is not None else None, \ - self.vw.data, \ - self.vb.data if self.vb is not None else None] - - self.qw.data = self.qkvw.data[:self.qw.shape[0]] - self.kw.data = self.qkvw.data[self.qw.shape[0]:2 * self.qw.shape[0]] - self.vw.data = self.qkvw.data[2 * self.qw.shape[0]:] - - if self.qkvb is not None: - self.qb.data = self.qkvb.data[:self.qw.shape[0]] - self.kb.data = self.qkvb.data[self.qw.shape[0]:2 * self.qw.shape[0]] - self.vb.data = self.qkvb.data[2 * self.qw.shape[0]:] - - for data in qkv_data: - del data - - def set_params_wo_copy(self, Z3_enabled=False): - self.module.mlp.attn_nw = self.attn_nw - self.module.mlp.attn_nb = self.attn_nb - self.module.norm_w = self.input_nw - self.module.norm_b = self.input_nb - self.module.mlp.inter_w = self._h4h_w - self.module.mlp.inter_b = self._h4h_b - self.module.mlp.output_w = self._4hh_w - self.module.mlp.output_b = self._4hh_b - self.module.attention.attn_ow = self.dense_w - self.module.attention.attn_ob = self.dense_b - if not Z3_enabled or self.q_k_v is None: - self.module.attention.attn_qkvw = self.qkvw - self.module.attention.attn_qkvb = self.qkvb - if self.q_k_v is not None: - if Z3_enabled: - self.module.attention.attn_qw = self.qw - self.module.attention.attn_qb = self.qb - self.module.attention.attn_kw = self.kw - self.module.attention.attn_kb = self.kb - self.module.attention.attn_vw = self.vw - self.module.attention.attn_vb = self.vb - else: - self.qw.data = self.qkvw[:self.qw.shape[0], :] - self.kw.data = self.qkvw[self.qw.shape[0]:2 * self.qw.shape[0], :] - self.vw.data = self.qkvw[self.qw.shape[0] * 2:, :] - if self.qkvb is not None: - self.qb.data = self.qkvb[:self.qw.shape[0]] - self.kb.data = self.qkvb[self.qw.shape[0]:2 * self.qw.shape[0]] - self.vb.data = self.qkvb[self.qw.shape[0] * 2:] + def get_all_params(self): + params = [ + self.attn_nw, + self.attn_nb, + self.input_nw, + self.input_nb, + ] - def get_lora_params(self): - return self.lora_params + params.extend(self.get_attn_params()) + params.extend(self.get_mlp_params()) - def get_all_params(self): - if self.q_k_v is not None: - return [ - self.attn_nw, self.attn_nb, self.input_nw, self.input_nb, self._h4h_w, self._h4h_b, self._4hh_w, - self._4hh_b, self.qw, self.qb, self.kw, self.kb, self.vw, self.vb, self.dense_w, self.dense_b - ] - else: - return [ - self.attn_nw, self.attn_nb, self.input_nw, self.input_nb, self._h4h_w, self._h4h_b, self._4hh_w, - self._4hh_b, self.qkvw, self.qkvb, self.dense_w, self.dense_b - ] + return params + + def get_attn_params(self): + return [self.qkvw, self.qkvb, self.dense_w, self.dense_b] + + def get_mlp_params(self): + return [self._h4h_w, self._h4h_b, self._4hh_w, self._4hh_b] diff --git a/deepspeed/module_inject/containers/bert.py b/deepspeed/module_inject/containers/bert.py index f8070655283e..967a02276be6 100644 --- a/deepspeed/module_inject/containers/bert.py +++ b/deepspeed/module_inject/containers/bert.py @@ -50,10 +50,8 @@ def get_hidden_heads(self): attention_layernorm = self.client_module.attention.output.LayerNorm return self.client_module.attention.self.query.weight.shape[1], \ self.client_module.attention.self.num_attention_heads, \ - attention_layernorm.eps - - def get_q_k_v(self): - return None + attention_layernorm.eps, \ + DEFAULT_INTERMEDIATE_SIZE def attention(self, enable_training=False): qw = self.client_module.attention.self.query.weight @@ -71,7 +69,7 @@ def attention(self, enable_training=False): self.client_module.attention.output.dense.weight, \ self.client_module.attention.output.dense.bias, \ - def mlp(self): + def mlp(self, enable_training=False): if self.pre_attn_norm: intermediate_ff = self.client_module.intermediate.dense_act else: @@ -92,6 +90,3 @@ def layernorm(self): attention_layernorm.bias, \ transformer_layernorm.weight, \ transformer_layernorm.bias - - def get_lora_params(self): - return [] diff --git a/deepspeed/module_inject/containers/bloom.py b/deepspeed/module_inject/containers/bloom.py index 136971583cf7..86bc295f1dd2 100644 --- a/deepspeed/module_inject/containers/bloom.py +++ b/deepspeed/module_inject/containers/bloom.py @@ -87,10 +87,8 @@ def __init__(self, client_module, inference=True, use_load_prefix=True, split_qk def get_hidden_heads(self): return self.client_module.self_attention.hidden_size, \ self.client_module.self_attention.num_heads, \ - self.client_module.input_layernorm.eps - - def get_q_k_v(self): - return None + self.client_module.input_layernorm.eps, \ + DEFAULT_INTERMEDIATE_SIZE def attention(self, enable_training=False): return self.client_module.self_attention.query_key_value.weight, \ @@ -98,7 +96,7 @@ def attention(self, enable_training=False): self.client_module.self_attention.dense.weight, \ self.client_module.self_attention.dense.bias, - def mlp(self): + def mlp(self, enable_training=False): return self.client_module.mlp.dense_h_to_4h.weight, \ self.client_module.mlp.dense_h_to_4h.bias, \ self.client_module.mlp.dense_4h_to_h.weight, \ diff --git a/deepspeed/module_inject/containers/clip.py b/deepspeed/module_inject/containers/clip.py index 144f1b823a1a..afe4a76086d8 100644 --- a/deepspeed/module_inject/containers/clip.py +++ b/deepspeed/module_inject/containers/clip.py @@ -41,12 +41,10 @@ def __init__(self, client_module, inference=False): def get_hidden_heads(self): return self.client_module.self_attn.q_proj.weight.shape[1], \ self.client_module.self_attn.num_heads, \ - self.client_module.layer_norm1.eps + self.client_module.layer_norm1.eps, \ + DEFAULT_INTERMEDIATE_SIZE - def get_q_k_v(self): - return None - - def attention(self): + def attention(self, enable_training=False): qw = self.client_module.self_attn.q_proj.weight qb = self.client_module.self_attn.q_proj.bias kw = self.client_module.self_attn.k_proj.weight @@ -54,15 +52,15 @@ def attention(self): vw = self.client_module.self_attn.v_proj.weight vb = self.client_module.self_attn.v_proj.bias - qkvw = Parameter(torch.cat((qw, kw, vw), dim=0), requires_grad=False) - qkvb = Parameter(torch.cat((qb, kb, vb), dim=0), requires_grad=False) + qkvw = Parameter(torch.cat((qw, kw, vw), dim=0), requires_grad=enable_training) + qkvb = Parameter(torch.cat((qb, kb, vb), dim=0), requires_grad=enable_training) return qkvw, \ qkvb, \ self.client_module.self_attn.out_proj.weight, \ self.client_module.self_attn.out_proj.bias - def mlp(self): + def mlp(self, enable_training=False): return self.client_module.mlp.fc1.weight, \ self.client_module.mlp.fc1.bias, \ self.client_module.mlp.fc2.weight, \ @@ -73,6 +71,3 @@ def layernorm(self): self.client_module.layer_norm2.bias, \ self.client_module.layer_norm1.weight, \ self.client_module.layer_norm1.bias - - def get_lora_params(self): - return [] diff --git a/deepspeed/module_inject/containers/distil_bert.py b/deepspeed/module_inject/containers/distil_bert.py index 792b965399e2..2acd144cac04 100644 --- a/deepspeed/module_inject/containers/distil_bert.py +++ b/deepspeed/module_inject/containers/distil_bert.py @@ -46,10 +46,8 @@ def __init__(self, client_module, inference=False, preln=False): def get_hidden_heads(self): return self.client_module.attention.q_lin.weight.shape[1], \ self.client_module.attention.n_heads, \ - self.client_module.sa_layer_norm.eps - - def get_q_k_v(self): - return None + self.client_module.sa_layer_norm.eps, \ + DEFAULT_INTERMEDIATE_SIZE def attention(self, enable_training=False): qw = self.client_module.attention.q_lin.weight @@ -67,7 +65,7 @@ def attention(self, enable_training=False): self.client_module.attention.out_lin.weight, \ self.client_module.attention.out_lin.bias - def mlp(self): + def mlp(self, enable_training=False): intermediate_ff = self.client_module.ffn.lin1 return intermediate_ff.weight, intermediate_ff.bias, \ @@ -81,6 +79,3 @@ def layernorm(self): attention_layernorm.bias, \ transformer_layernorm.weight, \ transformer_layernorm.bias - - def get_lora_params(self): - return [] diff --git a/deepspeed/module_inject/containers/features/__init__.py b/deepspeed/module_inject/containers/features/__init__.py index 9bf65591925d..fc2eb2a65531 100644 --- a/deepspeed/module_inject/containers/features/__init__.py +++ b/deepspeed/module_inject/containers/features/__init__.py @@ -3,5 +3,7 @@ # DeepSpeed Team +from .gated_mlp import HybridGatedMLPContainer from .megatron import MegatronContainer from .meta_tensor import MetaTensorContainer +from .split_qkv import HybridSplitQKVContainer diff --git a/deepspeed/module_inject/containers/features/gated_mlp.py b/deepspeed/module_inject/containers/features/gated_mlp.py new file mode 100644 index 000000000000..251338ce00a8 --- /dev/null +++ b/deepspeed/module_inject/containers/features/gated_mlp.py @@ -0,0 +1,103 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from abc import abstractmethod + +from .hybrid_engine import HybridEngineContainer + + +class HybridGatedMLPContainer(HybridEngineContainer): + """ + The HybridGatedMLPContainer supports models for which the first MLP layer + is represented with two separate weights, one for the activation function + and one for the gating function. + """ + + def set_mlp(self, _h4h_w, _h4h_b, _4hh_w, _4hh_b): + super().set_mlp(_h4h_w, _h4h_b, _4hh_w, _4hh_b) + self.set_mlp_gate() + + @abstractmethod + def set_mlp_gate(self): + """ + In `set_mlp_gate`, it is necessary to populate the following variables (where appropriate) + for the given model: + self.inter_up_w: inter up weight + self.inter_up_b: inter up bias + self.inter_gate_w: inter gate weight + self.inter_gate_b: inter gate bias + If the parameter does not exist in the original model, set the attribute to None. + """ + raise NotImplementedError("A set_mlp_gate() function must be defined in the model container \ + in order to set the unfused inter up and gate tensors.") + + def mlp_inter_mp(self, mp_replace, reversed_dim=False): + # Only need to alter behavior if we can't do the normal destructive copy + if self.module.mlp.inter_w is None: + params = [ + (self.module.mlp.inter_up_w, self.inter_up_w), + (self.module.mlp.inter_up_b, self.inter_up_b), + (self.module.mlp.inter_gate_w, self.inter_gate_w), + (self.module.mlp.inter_gate_b, self.inter_gate_b), + ] + for dst, src in params: + dst = mp_replace.copy(dst[:self.inter_up_w.shape[0] // mp_replace.mp_size], + src, + int8=reversed_dim, + allocate_tensor=reversed_dim) if src is not None else None + else: + super().mlp_inter_mp(mp_replace) + + def release_mlp(self): + super().release_mlp() + gated_mlp_params = [ + (self.module.mlp.inter_up_w, self.inter_up_w), + (self.module.mlp.inter_up_b, self.inter_up_b), + (self.module.mlp.inter_gate_w, self.inter_gate_w), + (self.module.mlp.inter_gate_b, self.inter_gate_b), + ] + + self._release_params(gated_mlp_params) + + def reset_mlp(self): + self._h4h_w.data[:self.inter_up_w.shape[0]] = self.inter_up_w.data + self._h4h_w.data[self.inter_up_w.shape[0]:] = self.inter_gate_w.data + + if self.inter_up_b is not None: + self._h4h_b.data[:self.inter_up_b.shape[0]] = self.inter_up_b.data + self._h4h_b.data[self.inter_up_b.shape[0]:] = self.inter_gate_b.data + + inter_data = [self.inter_up_w.data, self.inter_gate_w.data] + if self.inter_up_b is not None: + inter_data.extend([self.inter_up_b.data, self.inter_gate_b.data]) + + self.inter_up_w.data = self._h4h_w.data[:self.inter_up_w.shape[0]] + self.inter_gate_w.data = self._h4h_w.data[self.inter_up_w.shape[0]:] + + if self.inter_up_b is not None: + self.inter_up_b.data = self._h4h_b.data[:self.inter_up_b.shape[0]] + self.inter_gate_b.data = self._h4h_b.data[self.inter_up_b.shape[0]:] + + for data in inter_data: + del data + + def set_mlp_params_wo_copy(self, Z3_enabled=False): + if not Z3_enabled: + self.module.mlp.inter_w = self.inter_up_w + self.module.mlp.inter_b = self.inter_up_b + self.inter_up_w.data = self._h4h_w[:self.inter_up_w.shape[0], :] + self.inter_up_b.data = self._h4h_b[:self.inter_up_w.shape[0]] if self._h4h_b is not None else None + self.inter_gate_w.data = self._h4h_w[self.inter_up_w.shape[0]:, :] + self.inter_gate_b.data = self._h4h_b[self.inter_up_w.shape[0]:] if self._h4h_b is not None else None + else: + self.module.inter_up_w = self.inter_up_w + self.module.inter_up_b = self.inter_up_b + self.module.inter_gate_w = self.inter_gate_w + self.module.inter_gate_b = self.inter_gate_b + + def get_mlp_params(self): + params = super().get_mlp_params() + params.extend([self.inter_up_w, self.inter_up_b, self.inter_gate_w, self.inter_gate_b]) + return params diff --git a/deepspeed/module_inject/containers/features/hybrid_engine.py b/deepspeed/module_inject/containers/features/hybrid_engine.py new file mode 100644 index 000000000000..a6b8e994875d --- /dev/null +++ b/deepspeed/module_inject/containers/features/hybrid_engine.py @@ -0,0 +1,187 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from abc import ABC, abstractmethod +from typing import List, Tuple + +import torch + + +class HybridEngineContainer(ABC): + """ + This container identifies which methods need to be overridden in addition to + the base container to enable use in the RLHF pipeline. These methods are not + necessary for inference alone. + + NOTE: If you are using this feature with a container that + also inherits from `MetaTensorContainer`, ensure that `MetaTensorContainer` + is inherited before `HybridEngineContainer` in the class definition. + """ + + def initialize_tensors(self, enable_training=False): + """ + Same purposes as the base container, but also grabs the hooks for any LoRA + parameters. If it's necessary to override specific sub-components of the model, + it's best to augment the specific `set_[component]` itself rather than modifying + the `initialize_tensors` method. See the `HybridSplitQKVContainer` for an example. + """ + super().initialize_tensors(enable_training=enable_training) + self.set_lora_params() + + def transform_for_training(self): + """ + If the views on certain parameters are largely incompatible, it may be necessary to do + more substantial transformations to the parameters. This method should be overridden to + transform the inference format to what is necessary for training. + """ + pass + + def transform_for_inference(self): + """ + If the views on certain parameters are largely incompatible, it may be necessary to do + more substantial transformations to the parameters. This method should be overridden to + transform the training format to what is necessary for inference. + """ + pass + + @abstractmethod + def set_lora_params(self, lora_params): + """ + If available, set the LoRA parameters for the module. An implementation + for this would iterate over all parameters of the model and use the `maybe_get_lora` helper + method to check if the parameter does in fact have any LoRA params. + """ + raise NotImplementedError("A set_lora_params() function must be defined for the relevant parameters.") + + def apply_tensor_parallelism(self, mp_replace, reversed_dim=False): + """ + Add support for reversed dim in tensor parallelism. If necessary, override + the called methods to handle partitioned weights (i.e. if qkv is split, override + the `attention_qkv_mp` method). If the model component is not split, it should + be safe to use the default implementation. + """ + # Setup the new Attention module + self.attention_qkv_mp(mp_replace, reversed_dim=reversed_dim) + self.attention_o_mp(mp_replace, reversed_dim=reversed_dim) + + # Setup the new MLP module + self.mlp_inter_mp(mp_replace, reversed_dim=reversed_dim) + self.mlp_output_mp(mp_replace, reversed_dim=reversed_dim) + + # Apply weight quantization + # TODO(cmikeh2): Re-enable this once verified + #self.apply_weight_quantization() + + def _release_params(self, param_pairs: List[Tuple[torch.Tensor, torch.Tensor]]): + """ + Helper for `release_[component]` methods. Accepts a list of tuples where the first + element is the module param that needs to be deleted, and the second is the reassignment + from the container. + """ + for module_param, container_param in param_pairs: + if module_param is not None: + del module_param + module_param = container_param + + def release_memory(self): + """ + Delete module parameters if they exist and point them back to the container. The primary + purpose of this is for TP-inference with ZeRO-3. In this scenario, we need to delete the + parameters we've created for inference to free their memory. + """ + general_params = [ + (self.module.attention.attn_ow, self.dense_w), + (self.module.attention.attn_ob, self.dense_b), + (self.module.attn_nw, self.attn_nw), + (self.module.attn_nb, self.attn_nb), + (self.module.norm_w, self.input_nw), + (self.module.norm_b, self.input_nb), + ] + + self._release_params(general_params) + + self.release_qkv() + self.release_mlp() + + def release_qkv(self): + """ + Release for QKV parameters (as well as any aliases). + """ + qkv_params = [ + (self.module.attention.attn_qkvw, self.qkvw), + (self.module.attention.attn_qkvb, self.qkvb), + ] + + self._release_params(qkv_params) + + def release_mlp(self): + """ + Release for MLP parameters (as well as any aliases). + """ + mlp_params = [ + (self.module.mlp.inter_w, self._h4h_w), + (self.module.mlp.inter_b, self._h4h_b), + (self.module.mlp.output_w, self._4hh_w), + (self.module.mlp.output_b, self._4hh_b), + ] + + self._release_params(mlp_params) + + def reset_params(self): + """ + The purpose of reset params is to get the weights from the FP16 training + copy of the model and copy to them to contiguous inference view. This only needs + to be performed when the container parameters cannot be used directly for inference. + """ + self.reset_qkv() + self.reset_mlp() + + def reset_qkv(self): + """ + Perform any necessary resets of the model parameters for the QKV components. + """ + pass + + def reset_mlp(self): + """ + Perform any necessary resets of the model parameters for the MLP components. + """ + pass + + def get_lora_params(self): + """ + Return a list of all parameters that would have LoRA for the module. + """ + return self.lora_params + + def set_params_wo_copy(self, Z3_enabled=False): + """ + Rather than copying into, set the parameters directly. This is necessary to provide + an inexpensive (low-memory-overhead) view onto the FP16 forward weights. + """ + self.module.mlp.attn_nw = self.attn_nw + self.module.mlp.attn_nb = self.attn_nb + self.module.norm_w = self.input_nw + self.module.norm_b = self.input_nb + self.set_attn_params_wo_copy(Z3_enabled=Z3_enabled) + self.set_mlp_params_wo_copy(Z3_enabled=Z3_enabled) + + def set_attn_params_wo_copy(self, Z3_enabled=False): + """ + Narrower sub-method for finer grained overriding. + """ + self.module.attention.attn_ow = self.dense_w + self.module.attention.attn_ob = self.dense_b + self.module.attention.attn_qkvw = self.qkvw + self.module.attention.attn_qkvb = self.qkvb + + def set_mlp_params_wo_copy(self, Z3_enabled=False): + """ + Narrower sub-method for finer grained overriding. + """ + self.module.mlp.inter_w = self._h4h_w + self.module.mlp.inter_b = self._h4h_b + self.module.mlp.output_w = self._4hh_w + self.module.mlp.output_b = self._4hh_b diff --git a/deepspeed/module_inject/containers/features/hybrid_megatron.py b/deepspeed/module_inject/containers/features/hybrid_megatron.py new file mode 100644 index 000000000000..d40f2a6b57e8 --- /dev/null +++ b/deepspeed/module_inject/containers/features/hybrid_megatron.py @@ -0,0 +1,87 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch + +from .hybrid_engine import HybridEngineContainer +from .megatron import MegatronContainer + + +class HybridMegatronContainer(MegatronContainer, HybridEngineContainer): + + def _align_qkv(self, x: torch.Tensor): + """ + Internal helper for accepting the head-contiguous weight matrix and chunking + the query, key, and value components. + """ + attention_head_size = x.shape[0] // self.num_attention_heads + new_x_shape = (self.num_attention_heads, attention_head_size) + x.size()[1:] + x_1 = x.view(*new_x_shape) + div_dim = len(x_1.size()) - 2 if len(x.shape) == 2 else -1 + (q, k, v) = torch.split(x_1, (x_1.shape[div_dim] // 3), dim=div_dim) + if len(q.shape) > 2: + x.data.copy_( + torch.cat((q.reshape(-1, q.shape[-1]), k.reshape(-1, q.shape[-1]), v.reshape(-1, q.shape[-1])), + dim=0).reshape(x.shape)) + else: + x.data.copy_(torch.cat((q.reshape(-1), k.reshape(-1), v.reshape(-1)), dim=-1).reshape(x.shape)) + + def transform_for_inference(self) -> None: + """ + Overrides the HybridEngineContainer implementation. + + The alternative layout of the QKV matrix for Megatron is such that each head's Q, K, and V + are sequential in memory. This is different from the default layout in which all of the Qs + are sequential, followed by all of the Ks, and then all of the Vs. Here, we take the default + layout and transform it to the inference layout. + """ + if hasattr(self.qkvw, 'ds_id'): + from deepspeed.runtime.zero import GatheredParameters + from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus + param_list = [self.qkvw, self.qkvb] + non_active_params = [param for param in param_list if (hasattr(param, 'ds_id') and \ + param.ds_status == ZeroParamStatus.NOT_AVAILABLE)] + with GatheredParameters(non_active_params): + self._align_qkv(self.qkvw) + self._align_qkv(self.qkvb) + else: + self._align_qkv(self.qkvw) + self._align_qkv(self.qkvb) + + def _partition_qkv(self, x: torch.Tensor): + """ + Internal helper for taking contiguous QKV and partitioning it for contiguous + heads. + """ + q_k_v = torch.split(x, (x.shape[0] // 3), dim=0) + attention_head_size = q_k_v[0].shape[0] // self.num_attention_heads + new_x_shape = (self.num_attention_heads, attention_head_size) + x.size()[1:] + q, k, v = [data.view(*new_x_shape) for data in q_k_v] + if len(q.shape) > 2: + x.data.copy_(torch.cat((q, k, v), dim=-2).reshape(-1, q.shape[-1])) + else: + x.data.copy_(torch.cat((q, k, v), dim=-1).reshape(-1)) + + def transform_for_training(self): + """ + Overrides the HybridEngineContainer implementation. + + The alternative layout of the QKV matrix for Megatron is such that each head's Q, K, and V + are sequential in memory. This is different from the default layout in which all of the Qs + are sequential, followed by all of the Ks, and then all of the Vs. This function takes the inference format and reverts it back to the default format. + """ + # If parameter is distributed, handle gathering it + if hasattr(self.qkvw, 'ds_id'): + from deepspeed.runtime.zero import GatheredParameters + from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus + param_list = [self.qkvw, self.qkvb] + non_active_params = [param for param in param_list if (hasattr(param, 'ds_id') and \ + param.ds_status == ZeroParamStatus.NOT_AVAILABLE)] + with GatheredParameters(non_active_params): + self._partition_qkv(self.qkvw) + self._partition_qkv(self.qkvb) + else: + self._partition_qkv(self.qkvw) + self._partition_qkv(self.qkvb) diff --git a/deepspeed/module_inject/containers/features/megatron.py b/deepspeed/module_inject/containers/features/megatron.py index cce106fa7e9a..4daccf7d7c8d 100644 --- a/deepspeed/module_inject/containers/features/megatron.py +++ b/deepspeed/module_inject/containers/features/megatron.py @@ -24,57 +24,6 @@ def _align_qkv_transposed(self, x): else: return torch.cat((q.reshape(-1), k.reshape(-1), v.reshape(-1)), dim=-1).reshape(x.shape) - def _align_qkv(self, x): - attention_head_size = x.shape[0] // self.num_attention_heads - new_x_shape = (self.num_attention_heads, attention_head_size) + x.size()[1:] - x_1 = x.view(*new_x_shape) - div_dim = len(x_1.size()) - 2 if len(x.shape) == 2 else -1 - (q, k, v) = torch.split(x_1, (x_1.shape[div_dim] // 3), dim=div_dim) - if len(q.shape) > 2: - x.data.copy_( - torch.cat((q.reshape(-1, q.shape[-1]), k.reshape(-1, q.shape[-1]), v.reshape(-1, q.shape[-1])), - dim=0).reshape(x.shape)) - else: - x.data.copy_(torch.cat((q.reshape(-1), k.reshape(-1), v.reshape(-1)), dim=-1).reshape(x.shape)) - - def _align_merged_qkv(self): - if hasattr(self.qkvw, 'ds_id'): - from deepspeed.runtime.zero import GatheredParameters - from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus - param_list = [self.qkvw, self.qkvb] - non_active_params = [param for param in param_list if (hasattr(param, 'ds_id') and \ - param.ds_status == ZeroParamStatus.NOT_AVAILABLE)] - with GatheredParameters(non_active_params): - self._align_qkv(self.qkvw) - self._align_qkv(self.qkvb) - else: - self._align_qkv(self.qkvw) - self._align_qkv(self.qkvb) - - def _partition_qkv(self, x): - q_k_v = torch.split(x, (x.shape[0] // 3), dim=0) - attention_head_size = q_k_v[0].shape[0] // self.num_attention_heads - new_x_shape = (self.num_attention_heads, attention_head_size) + x.size()[1:] - q, k, v = [data.view(*new_x_shape) for data in q_k_v] - if len(q.shape) > 2: - x.data.copy_(torch.cat((q, k, v), dim=-2).reshape(-1, q.shape[-1])) - else: - x.data.copy_(torch.cat((q, k, v), dim=-1).reshape(-1)) - - def _partition_merged_qkv(self): - if hasattr(self.qkvw, 'ds_id'): - from deepspeed.runtime.zero import GatheredParameters - from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus - param_list = [self.qkvw, self.qkvb] - non_active_params = [param for param in param_list if (hasattr(param, 'ds_id') and \ - param.ds_status == ZeroParamStatus.NOT_AVAILABLE)] - with GatheredParameters(non_active_params): - self._partition_qkv(self.qkvw) - self._partition_qkv(self.qkvb) - else: - self._partition_qkv(self.qkvw) - self._partition_qkv(self.qkvb) - def transpose(self): super().transpose() if self.megatron_v2: diff --git a/deepspeed/module_inject/containers/features/meta_tensor.py b/deepspeed/module_inject/containers/features/meta_tensor.py index 7aa507ca2e44..465548465b25 100644 --- a/deepspeed/module_inject/containers/features/meta_tensor.py +++ b/deepspeed/module_inject/containers/features/meta_tensor.py @@ -7,6 +7,11 @@ class MetaTensorContainer(ABC): + """ + NOTE: If you are using this feature with a container that + also inherits from `HybridEngineContainer`, ensure that `MetaTensorContainer` + is inherited before `HybridEngineContainer` in the class definition. + """ def __init__(self, **kwargs): super().__init__(**kwargs) @@ -17,14 +22,14 @@ def initialize_tensors(self, enable_training=False): super().initialize_tensors(enable_training=enable_training) self.is_meta = self.qkvw.is_meta - def apply_tensor_parallelism(self, mp_replace=None, mp_group=None, tp_size=None): + def apply_tensor_parallelism(self, mp_replace, **kwargs): if self.is_meta: if self.qkvb is None: self.module.attention.attn_qkvb = None if self.dense_b is None: self.module.attention.attn_ob = None else: - super().apply_tensor_parallelism(mp_replace, mp_group, tp_size) + super().apply_tensor_parallelism(mp_replace, **kwargs) def copy_data_to_new_module(self): if self.is_meta: diff --git a/deepspeed/module_inject/containers/features/split_qkv.py b/deepspeed/module_inject/containers/features/split_qkv.py new file mode 100644 index 000000000000..654a041aad92 --- /dev/null +++ b/deepspeed/module_inject/containers/features/split_qkv.py @@ -0,0 +1,156 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from abc import abstractmethod +import torch + +from .hybrid_engine import HybridEngineContainer + + +class HybridSplitQKVContainer(HybridEngineContainer): + + def set_attention(self, qkvw, qkvb, dense_w, dense_b): + super().set_attention(qkvw, qkvb, dense_w, dense_b) + self.set_q_k_v() + + @abstractmethod + def set_q_k_v(self): + """ + In `set_q_k_v`, it is necessary to populate the following variables (where appropriate) + for the given model: + self.qw: q weight + self.qb: q bias + self.kw: k weight + self.kb: k bias + self.vw: v weight + self.vb: v bias + """ + raise NotImplementedError("A set_q_k_v() function must be defined in the model container \ + in order to set the unfused q, k, and v tensors.") + + def attention_qkv_mp(self, mp_replace, reversed_dim=False): + # Only need to alter + if self.module.attention.attn_qkvw is None: + params = [ + (self.module.attention.attn_qw, self.qw), + (self.module.attention.attn_qb, self.qb), + (self.module.attention.attn_kw, self.kw), + (self.module.attention.attn_kb, self.kb), + (self.module.attention.attn_vw, self.vw), + (self.module.attention.attn_vb, self.vb), + ] + for dst, src in params: + dst = mp_replace.copy( + dst[:self.qw.shape[0] // mp_replace.mp_size], src, int8=reversed_dim, + allocate_tensor=reversed_dim) if src is not None else None + else: + super().attention_qkv_mp(mp_replace) + + def release_qkv(self): + super().release_qkv() + split_qkv_params = [ + (self.module.attention.attn_qw, self.qw), + (self.module.attention.attn_qb, self.qb), + (self.module.attention.attn_kw, self.kw), + (self.module.attention.attn_kb, self.kb), + (self.module.attention.attn_vw, self.vw), + (self.module.attention.attn_vb, self.vb), + ] + + self._release_params(split_qkv_params) + + def reset_qkv(self): + self.qkvw.data[:self.qw.shape[0]] = self.qw.data + self.qkvw.data[self.qw.shape[0]:2 * self.qw.shape[0]] = self.kw.data + self.qkvw.data[2 * self.qw.shape[0]:] = self.vw.data + + qkv_data = [self.qw.data, self.kw.data, self.vw.data] + + self.qw.data = self.qkvw.data[:self.qw.shape[0]] + self.kw.data = self.qkvw.data[self.qw.shape[0]:2 * self.qw.shape[0]] + self.vw.data = self.qkvw.data[2 * self.qw.shape[0]:] + + if self.qkvb is not None: + self.qkvb.data[:self.qw.shape[0]] = self.qb.data + self.qkvb.data[self.qw.shape[0]:2 * self.qw.shape[0]] = self.kb.data + self.qkvb.data[2 * self.qw.shape[0]:] = self.vb.data + + qkv_data.extend([self.qb.data, self.kb.data, self.vb.data]) + + self.qb.data = self.qkvb.data[:self.qw.shape[0]] + self.kb.data = self.qkvb.data[self.qw.shape[0]:2 * self.qw.shape[0]] + self.vb.data = self.qkvb.data[2 * self.qw.shape[0]:] + + for data in qkv_data: + del data + + def reset_qkv_experimental(self): + """ + WIP - experimental and likely to be changed/improved. + Unused by keeping for now. + """ + if self.module.attention.attn_qkvw is None: + self.module.attention.attn_qkvw = torch.empty(self.qw.shape[0] * 3, + self.qw.shape[0], + dtype=self.qw.dtype, + device=self.qw.device) + self.module.attention.attn_qkvb = torch.empty(self.qw.shape[0] * 3, + dtype=self.qw.dtype, + device=self.qw.device) + self.module.attention.attn_qkvw.data[:self.qw.shape[0]] = self.qw.data + self.module.attention.attn_qkvb.data[:self.qw.shape[0]] = self.qb.data + self.module.attention.attn_qkvw.data[self.qw.shape[0]:2 * self.qw.shape[0]] = self.kw.data + self.module.attention.attn_qkvb.data[self.qw.shape[0]:2 * self.qw.shape[0]] = self.kb.data + self.module.attention.attn_qkvw.data[2 * self.qw.shape[0]:] = self.vw.data + self.module.attention.attn_qkvb.data[2 * self.qw.shape[0]:] = self.vb.data + + qkv_data = [self.qw.data, \ + self.qb.data, \ + self.kw.data, \ + self.kb.data, \ + self.vw.data, \ + self.vb.data] + + self.qw.data = self.module.attention.attn_qkvw.data[:self.qw.shape[0]] + self.qb.data = self.module.attention.attn_qkvb.data[:self.qw.shape[0]] + self.kw.data = self.module.attention.attn_qkvw.data[self.qw.shape[0]:2 * self.qw.shape[0]] + self.kb.data = self.module.attention.attn_qkvb.data[self.qw.shape[0]:2 * self.qw.shape[0]] + self.vw.data = self.module.attention.attn_qkvw.data[2 * self.qw.shape[0]:] + self.vb.data = self.module.attention.attn_qkvb.data[2 * self.qw.shape[0]:] + + for data in qkv_data: + del data + + def set_attn_parameters_wo_copy(self, Z3_enabled=False): + self.module.attention.attn_ow = self.dense_w + self.module.attention.attn_ob = self.dense_b + if not Z3_enabled: + # In initialize_tensors, we create a fused qkvw with the appropriate shape + # and copy the qw, qb, kw, kb, vw, vb into it + self.module.attn_qkvw = self.qkvw + self.module.attn_qkvb = self.qkvb + + # We reset the data for qw (which is the original model parameter) to point + # to the fused weight matrix we have created here + self.qw.data = self.qkvw[:self.qw.shape[0], :] + self.qb.data = self.qkvb[:self.qw.shape[0]] + self.kw.data = self.qkvw[self.qw.shape[0]:2 * self.qw.shape[0], :] + self.kb.data = self.qkvb[self.qw.shape[0]:2 * self.qw.shape[0]] + self.vw.data = self.qkvw[self.qw.shape[0] * 2:, :] + self.vb.data = self.qkvb[self.qw.shape[0] * 2:] + else: + # In ZeRO-3 this will be managed by ZeRO and handled separately in the + # forward of ds_attention + self.module.attention.attn_qw = self.qw + self.module.attention.attn_qb = self.qb + self.module.attention.attn_kw = self.kw + self.module.attention.attn_kb = self.kb + self.module.attention.attn_vw = self.vw + self.module.attention.attn_vb = self.vb + + def get_attn_params(self): + params = super().get_attn_params() + params.extend([self.qw, self.qb, self.kw, self.kb, self.vw, self.vb]) + return params diff --git a/deepspeed/module_inject/containers/gpt2.py b/deepspeed/module_inject/containers/gpt2.py index 3f6373897c58..7a19aac34b44 100644 --- a/deepspeed/module_inject/containers/gpt2.py +++ b/deepspeed/module_inject/containers/gpt2.py @@ -38,10 +38,8 @@ def __init__(self, client_module, inference=True): def get_hidden_heads(self): return self.client_module.attn.embed_dim, \ self.client_module.attn.num_heads, \ - self.client_module.ln_1.eps - - def get_q_k_v(self): - return None + self.client_module.ln_1.eps, \ + DEFAULT_INTERMEDIATE_SIZE def attention(self, enable_training=False): return self.client_module.attn.c_attn.weight, \ @@ -49,7 +47,7 @@ def attention(self, enable_training=False): self.client_module.attn.c_proj.weight, \ self.client_module.attn.c_proj.bias - def mlp(self): + def mlp(self, enable_training=False): return self.client_module.mlp.c_fc.weight, \ self.client_module.mlp.c_fc.bias, \ self.client_module.mlp.c_proj.weight, \ @@ -60,6 +58,3 @@ def layernorm(self): self.client_module.ln_2.bias, \ self.client_module.ln_1.weight, \ self.client_module.ln_1.bias - - def get_lora_params(self): - return [] diff --git a/deepspeed/module_inject/containers/gptj.py b/deepspeed/module_inject/containers/gptj.py index 6929cd2f7146..47806cbe9fd3 100644 --- a/deepspeed/module_inject/containers/gptj.py +++ b/deepspeed/module_inject/containers/gptj.py @@ -5,6 +5,7 @@ from .base import * from .features.meta_tensor import MetaTensorContainer +from .features.split_qkv import HybridSplitQKVContainer from deepspeed.model_implementations.transformers.ds_gpt import DeepSpeedGPTInference import torch from torch.nn.parameter import Parameter @@ -16,7 +17,7 @@ from ..policy import maybe_get_lora -class DS_GPTJContainer(MetaTensorContainer, BaseTransformerContainer): +class DS_GPTJContainer(MetaTensorContainer, HybridSplitQKVContainer, BaseTransformerContainer): def __init__(self, **kwargs): super().__init__(**kwargs) @@ -29,6 +30,29 @@ def create_module(self, config=None): self.module.config.scale_attention = self.scale_attention return self.module + def set_lora_params(self): + """ + Necessary to implement for `HybridEngineContainer` + """ + self.lora_params = [ + maybe_get_lora(p) for p in [ + self.policy.client_module.mlp.fc_in, self.policy.client_module.mlp.fc_out, + self.policy.client_module.attn.q_proj, self.policy.client_module.attn.k_proj, + self.policy.client_module.attn.v_proj, self.policy.client_module.attn.out_proj + ] + ] + + def set_q_k_v(self): + """ + Necessary to implement for `HybridSplitQKVContainer` + """ + self.qw = self.policy.client_module.attn.q_proj.weight + self.qb = None + self.kw = self.policy.client_module.attn.k_proj.weight + self.kb = None + self.vw = self.policy.client_module.attn.v_proj.weight + self.vb = None + def load_params(self, module, sd, weight_quantizer, mp_replace, prefix): param_names = ( 'attn.q_proj.weight', \ @@ -74,15 +98,8 @@ def __init__(self, client_module, inference=True): def get_hidden_heads(self): return self.client_module.attn.embed_dim, \ self.client_module.attn.num_attention_heads, \ - self.client_module.ln_1.eps - - def get_q_k_v(self): - return self.client_module.attn.q_proj.weight, \ - None, \ - self.client_module.attn.k_proj.weight, \ - None, \ - self.client_module.attn.v_proj.weight, \ - None + self.client_module.ln_1.eps, \ + DEFAULT_INTERMEDIATE_SIZE def attention(self, enable_training=False): qw = self.client_module.attn.q_proj.weight @@ -96,7 +113,7 @@ def attention(self, enable_training=False): self.client_module.attn.out_proj.weight, \ None, - def mlp(self): + def mlp(self, enable_training=False): return self.client_module.mlp.fc_in.weight, \ self.client_module.mlp.fc_in.bias, \ self.client_module.mlp.fc_out.weight, \ @@ -107,16 +124,3 @@ def layernorm(self): None, \ self.client_module.ln_1.weight, \ self.client_module.ln_1.bias - - def get_lora_params(self): - all_lora_params = [] - for p in [ - self.client_module.mlp.fc_in, \ - self.client_module.mlp.fc_out, \ - self.client_module.attn.q_proj, \ - self.client_module.attn.k_proj, \ - self.client_module.attn.v_proj, \ - self.client_module.attn.out_proj, \ - ]: - all_lora_params.append(maybe_get_lora(p)) - return all_lora_params diff --git a/deepspeed/module_inject/containers/gptneo.py b/deepspeed/module_inject/containers/gptneo.py index 146383733ccc..02f1e4167ab1 100644 --- a/deepspeed/module_inject/containers/gptneo.py +++ b/deepspeed/module_inject/containers/gptneo.py @@ -5,6 +5,7 @@ from .base import * from .features.meta_tensor import MetaTensorContainer +from .features.hybrid_engine import HybridEngineContainer from deepspeed.model_implementations.transformers.ds_gpt import DeepSpeedGPTInference import torch from torch.nn.parameter import Parameter @@ -16,7 +17,7 @@ from ..policy import maybe_get_lora -class DS_GPTNEOContainer(MetaTensorContainer, BaseTransformerContainer): +class DS_GPTNEOContainer(MetaTensorContainer, HybridEngineContainer, BaseTransformerContainer): def __init__(self, **kwargs): super().__init__(**kwargs) @@ -29,6 +30,18 @@ def create_module(self, config=None): self.module.config.scale_attention = self.scale_attention return self.module + def set_lora_params(self): + """ + Necessary to implement for `HybridEngineContainer` + """ + self.lora_params = [ + maybe_get_lora(p) for p in [ + self.policy.client_module.mlp.c_fc, self.policy.client_module.mlp.c_proj, + self.policy.client_module.attn.attention.q_proj, self.policy.client_module.attn.attention.k_proj, + self.policy.client_module.attn.attention.v_proj, self.policy.client_module.attn.attention.out_proj + ] + ] + def load_params(self, module, sd, weight_quantizer, mp_replace, prefix): param_names = ( 'attn.attention.q_proj.weight', \ @@ -76,7 +89,8 @@ def __init__(self, client_module, inference=True): def get_hidden_heads(self): return self.client_module.attn.attention.embed_dim, \ self.client_module.attn.attention.num_heads, \ - self.client_module.ln_1.eps + self.client_module.ln_1.eps, \ + DEFAULT_INTERMEDIATE_SIZE def get_q_k_v(self): return self.client_module.attn.attention.q_proj.weight, \ @@ -98,7 +112,7 @@ def attention(self, enable_training=False): self.client_module.attn.attention.out_proj.weight, \ self.client_module.attn.attention.out_proj.bias - def mlp(self): + def mlp(self, enable_training=False): return self.client_module.mlp.c_fc.weight, \ self.client_module.mlp.c_fc.bias, \ self.client_module.mlp.c_proj.weight, \ @@ -109,16 +123,3 @@ def layernorm(self): self.client_module.ln_2.bias, \ self.client_module.ln_1.weight, \ self.client_module.ln_1.bias - - def get_lora_params(self): - all_lora_params = [] - for p in [ - self.client_module.mlp.c_fc, \ - self.client_module.mlp.c_proj, \ - self.client_module.attn.attention.q_proj, \ - self.client_module.attn.attention.k_proj, \ - self.client_module.attn.attention.v_proj, \ - self.client_module.attn.attention.out_proj, \ - ]: - all_lora_params.append(maybe_get_lora(p)) - return all_lora_params diff --git a/deepspeed/module_inject/containers/gptneox.py b/deepspeed/module_inject/containers/gptneox.py index 249e57d2d9c2..16b0f90189ce 100644 --- a/deepspeed/module_inject/containers/gptneox.py +++ b/deepspeed/module_inject/containers/gptneox.py @@ -5,7 +5,7 @@ from .base import * from .features.meta_tensor import MetaTensorContainer -from .features.megatron import MegatronContainer +from .features.hybrid_megatron import HybridMegatronContainer from deepspeed.model_implementations.transformers.ds_gpt import DeepSpeedGPTInference import torch from ..policy import TransformerPolicy @@ -16,7 +16,7 @@ from ..policy import maybe_get_lora -class DS_GPTNEOXContainer(MetaTensorContainer, MegatronContainer, BaseTransformerContainer): +class DS_GPTNEOXContainer(MetaTensorContainer, HybridMegatronContainer, BaseTransformerContainer): def __init__(self, **kwargs): super().__init__(**kwargs) @@ -34,6 +34,22 @@ def create_module(self, config=None): return self.module + def set_lora_params(self): + """ + Necessary to implement for `HybridEngineContainer` + """ + if GPTNEOXLayerPolicy.version == 0: + attention = self.policy.client_module.attention + else: + attention = self.policy.client_module.self_attention + + self.lora_params = [ + maybe_get_lora(p) for p in [ + self.policy.client_module.mlp.dense_h_to_4h, self.policy.client_module.mlp.dense_4h_to_h, + attention.query_key_value, attention.dense + ] + ] + def load_params(self, module, sd, weight_quantizer, mp_replace, prefix): param_names = ( 'attention.query_key_value.weight', \ @@ -95,10 +111,8 @@ def get_hidden_heads(self): return self.client_module.attention.hidden_size, \ self.client_module.attention.num_attention_heads, \ - self.client_module.input_layernorm.eps - - def get_q_k_v(self): - return None + self.client_module.input_layernorm.eps, \ + DEFAULT_INTERMEDIATE_SIZE def attention(self, enable_training=False): if GPTNEOXLayerPolicy.version == 0: @@ -111,7 +125,7 @@ def attention(self, enable_training=False): attention.dense.weight, \ attention.dense.bias - def mlp(self): + def mlp(self, enable_training=False): return self.client_module.mlp.dense_h_to_4h.weight, \ self.client_module.mlp.dense_h_to_4h.bias, \ self.client_module.mlp.dense_4h_to_h.weight, \ diff --git a/deepspeed/module_inject/containers/llama.py b/deepspeed/module_inject/containers/llama.py new file mode 100644 index 000000000000..9b1710c46b3b --- /dev/null +++ b/deepspeed/module_inject/containers/llama.py @@ -0,0 +1,150 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .base import * +from .features import MetaTensorContainer, HybridSplitQKVContainer, HybridGatedMLPContainer +from deepspeed.utils.types import ActivationFuncType, NormType +from deepspeed.model_implementations.transformers.ds_gpt import DeepSpeedGPTInference +import torch +from torch.nn.parameter import Parameter + +from ..policy import ( + TransformerPolicy, + transformer_param_names, + maybe_copy, + maybe_copy_qkv, + maybe_copy_geglu, + maybe_get_lora, +) + + +class DS_LLAMAContainer(MetaTensorContainer, HybridGatedMLPContainer, HybridSplitQKVContainer, + BaseTransformerContainer): + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + # All model specific things should be defined here instead of the base class. + + def create_module(self, config=None): + _config = config if config is not None else self.ds_model_config + + _config.rotate_half = True + _config.rotate_every_two = False + _config.rotary_dim = self.hidden_size // self.num_attention_heads + self.module = DeepSpeedGPTInference(_config, mp_group=self.mp_group) + + return self.module + + def set_lora_params(self): + """ + Necessary to implement for `HybridEngineContainer` + """ + self.lora_params = [ + maybe_get_lora(p) for p in [ + self.policy.client_module.mlp.up_proj.weight, self.policy.client_module.mlp.gate_proj.weight, + self.policy.client_module.mlp.down_proj.weight, self.policy.client_module.self_attn.q_proj.weight, + self.policy.client_module.self_attn.k_proj.weight, self.policy.client_module.self_attn.v_proj.weight, + self.policy.client_module.self_attn.o_proj.weight + ] + ] + + def set_q_k_v(self): + """ + Necessary to implement for `HybridSplitQKVContainer` + """ + self.qw = self.policy.client_module.self_attn.q_proj.weight + self.qb = None + self.kw = self.policy.client_module.self_attn.k_proj.weight + self.kb = None + self.vw = self.policy.client_module.self_attn.v_proj.weight + self.vb = None + + def set_mlp_gate(self): + """ + Necessary to implement for `HybridGatedMLPContainer` + """ + self.inter_up_w = self.policy.client_module.mlp.up_proj.weight + self.inter_up_b = None + self.inter_gate_w = self.policy.client_module.mlp.gate_proj.weight + self.inter_gate_b = None + + def load_params(self, module, sd, weight_quantizer, mp_replace, prefix): + param_names = ( + 'self_attn.q_proj.weight', \ + 'self_attn.k_proj.weight', \ + 'self_attn.v_proj.weight', \ + 'self_attn.o_proj.weight', \ + 'mlp.up_proj.weight', \ + 'mlp.gate_proj.weight', \ + 'mlp.down_proj.weight', \ + 'input_layernorm.weight', \ + 'post_attention_layernorm.weight' + ) + + maybe_copy_qkv(module.attention, + sd, + weight_quantizer, + mp_replace, + 'attn_qkvw', [prefix + param_names[0], prefix + param_names[1], prefix + param_names[2]], + split_qkv=self.policy.split_qkv) + for i in range(3, 4): + maybe_copy(module.attention, sd, weight_quantizer, mp_replace, transformer_param_names[i - 1], + prefix + param_names[i]) + maybe_copy_geglu(module.mlp, sd, weight_quantizer, mp_replace, 'inter_w', + [prefix + param_names[4], prefix + param_names[5]]) + maybe_copy(module.mlp, sd, weight_quantizer, mp_replace, 'output_w', prefix + param_names[6]) + + maybe_copy(module, sd, weight_quantizer, mp_replace, transformer_param_names[8], prefix + param_names[7]) + maybe_copy(module, sd, weight_quantizer, mp_replace, transformer_param_names[10], prefix + param_names[8]) + + +class LLAMALayerPolicy(TransformerPolicy): + + def __init__(self, client_module, inference=True): + super().__init__( + inference, + mlp_act_func_type=ActivationFuncType.GATED_SILU, + norm_type=NormType.RMSNorm, + ) + self.client_module = client_module + try: + import transformers + LLAMALayerPolicy._orig_layer_class = transformers.models.llama.modeling_llama.LlamaDecoderLayer # type: ignore + except: + LLAMALayerPolicy._orig_layer_class = None + + def get_hidden_heads(self): + return self.client_module.self_attn.q_proj.weight.shape[1], \ + self.client_module.self_attn.num_heads, \ + self.client_module.input_layernorm.variance_epsilon, \ + self.client_module.mlp.gate_proj.weight.shape[0] + + def attention(self, enable_training=False): + qw = self.client_module.self_attn.q_proj.weight + kw = self.client_module.self_attn.k_proj.weight + vw = self.client_module.self_attn.v_proj.weight + + qkvw = Parameter(torch.cat((qw, kw, vw), dim=0), requires_grad=enable_training) + + return qkvw, \ + None, \ + self.client_module.self_attn.o_proj.weight, \ + None + + def mlp(self, enable_training=False): + mlp1_up = self.client_module.mlp.up_proj.weight + mlp1_gate = self.client_module.mlp.gate_proj.weight + mlp2 = self.client_module.mlp.down_proj.weight + + mlp1 = Parameter(torch.cat((mlp1_up, mlp1_gate), dim=0), requires_grad=enable_training) + + return mlp1, None, mlp2, None + + def layernorm(self): + return self.client_module.post_attention_layernorm.weight, \ + None, \ + self.client_module.input_layernorm.weight, \ + None diff --git a/deepspeed/module_inject/containers/megatron_gpt.py b/deepspeed/module_inject/containers/megatron_gpt.py index 28b7df5b48d5..2f47f7f6d92a 100644 --- a/deepspeed/module_inject/containers/megatron_gpt.py +++ b/deepspeed/module_inject/containers/megatron_gpt.py @@ -57,10 +57,8 @@ def __init__(self, client_module, inference=True): def get_hidden_heads(self): return self.client_module.attention.query_key_value.weight.shape[1], \ self.client_module.attention.num_attention_heads, \ - self.client_module.input_layernorm.eps - - def get_q_k_v(self): - return None + self.client_module.input_layernorm.eps, \ + DEFAULT_INTERMEDIATE_SIZE def attention(self, enable_training=False): if self.inference: @@ -74,7 +72,7 @@ def attention(self, enable_training=False): attention.dense.weight, \ attention.dense.bias - def mlp(self, moe_type='standard'): + def mlp(self, moe_type='standard', enable_training=False): from deepspeed.moe.utils import has_moe_layers moe, _ = has_moe_layers(self.client_module) @@ -110,6 +108,3 @@ def layernorm(self): self.client_module.post_attention_layernorm.bias, \ self.client_module.input_layernorm.weight, \ self.client_module.input_layernorm.bias - - def get_lora_params(self): - return [] diff --git a/deepspeed/module_inject/containers/megatron_gpt_moe.py b/deepspeed/module_inject/containers/megatron_gpt_moe.py index 0d5248d8d4d4..c4063be05b6c 100644 --- a/deepspeed/module_inject/containers/megatron_gpt_moe.py +++ b/deepspeed/module_inject/containers/megatron_gpt_moe.py @@ -57,7 +57,7 @@ def __init__(self, client_module, inference=True): def get_num_experts(self): return self.num_experts - def mlp(self, moe_type='standard'): + def mlp(self, moe_type='standard', enable_training=False): # for now, all of this is tightly coupled to megatron-deepspeed moe implementation # todo: think and refactor this to be more general diff --git a/deepspeed/module_inject/containers/opt.py b/deepspeed/module_inject/containers/opt.py index ff0cad57aa5c..381644485eab 100644 --- a/deepspeed/module_inject/containers/opt.py +++ b/deepspeed/module_inject/containers/opt.py @@ -4,7 +4,7 @@ # DeepSpeed Team from .base import * -from .features.meta_tensor import MetaTensorContainer +from .features import MetaTensorContainer, HybridSplitQKVContainer from deepspeed.model_implementations.transformers.ds_opt import DeepSpeedOPTInference import torch from torch.nn.parameter import Parameter @@ -16,7 +16,7 @@ from deepspeed.utils.types import ActivationFuncType -class DS_OPTContainer(MetaTensorContainer, BaseTransformerContainer): +class DS_OPTContainer(MetaTensorContainer, HybridSplitQKVContainer, BaseTransformerContainer): def __init__(self, **kwargs): super().__init__(**kwargs) @@ -29,6 +29,32 @@ def create_module(self, config=None): self.module.config.scale_attention = self.scale_attention return self.module + def set_lora_params(self): + """ + Necessary to implement for `HybridEngineContainer` + """ + self.lora_params = [ + maybe_get_lora(p) for p in [ + self.policy.client_module.fc1, + self.policy.client_module.fc2, + self.policy.client_module.self_attn.q_proj, + self.policy.client_module.self_attn.k_proj, + self.policy.client_module.self_attn.v_proj, + self.policy.client_module.self_attn.out_proj, + ] + ] + + def set_q_k_v(self): + """ + Necessary to implement for `HybridSplitQKVContainer` + """ + self.qw = self.policy.client_module.self_attn.q_proj.weight + self.qb = self.policy.client_module.self_attn.q_proj.bias + self.kw = self.policy.client_module.self_attn.k_proj.weight + self.kb = self.policy.client_module.self_attn.k_proj.bias + self.vw = self.policy.client_module.self_attn.v_proj.weight + self.vb = self.policy.client_module.self_attn.v_proj.bias + def load_params(self, module, sd, weight_quantizer, mp_replace, prefix): param_names = ( 'self_attn.q_proj.weight', \ @@ -95,15 +121,8 @@ def __init__(self, client_module, inference=True, use_load_prefix=True): def get_hidden_heads(self): return self.client_module.self_attn.embed_dim, \ self.client_module.self_attn.num_heads, \ - self.client_module.self_attn_layer_norm.eps - - def get_q_k_v(self): - return self.client_module.self_attn.q_proj.weight, \ - self.client_module.self_attn.q_proj.bias, \ - self.client_module.self_attn.k_proj.weight, \ - self.client_module.self_attn.k_proj.bias, \ - self.client_module.self_attn.v_proj.weight, \ - self.client_module.self_attn.v_proj.bias + self.client_module.self_attn_layer_norm.eps, \ + DEFAULT_INTERMEDIATE_SIZE def attention(self, enable_training=False): qw = self.client_module.self_attn.q_proj.weight @@ -122,7 +141,7 @@ def attention(self, enable_training=False): self.client_module.self_attn.out_proj.weight, \ self.client_module.self_attn.out_proj.bias - def mlp(self): + def mlp(self, enable_training=False): return self.client_module.fc1.weight, \ self.client_module.fc1.bias, \ self.client_module.fc2.weight, \ @@ -133,16 +152,3 @@ def layernorm(self): self.client_module.final_layer_norm.bias, \ self.client_module.self_attn_layer_norm.weight, \ self.client_module.self_attn_layer_norm.bias - - def get_lora_params(self): - all_lora_params = [] - for p in [ - self.client_module.fc1, \ - self.client_module.fc2, \ - self.client_module.self_attn.q_proj, \ - self.client_module.self_attn.k_proj, \ - self.client_module.self_attn.v_proj, \ - self.client_module.self_attn.out_proj, \ - ]: - all_lora_params.append(maybe_get_lora(p)) - return all_lora_params diff --git a/deepspeed/module_inject/policy.py b/deepspeed/module_inject/policy.py index 87b34e5aab5a..41df2b85dc0c 100644 --- a/deepspeed/module_inject/policy.py +++ b/deepspeed/module_inject/policy.py @@ -4,7 +4,7 @@ # DeepSpeed Team from abc import ABC, abstractmethod -from deepspeed.utils.types import ActivationFuncType +from deepspeed.utils.types import ActivationFuncType, NormType import torch from deepspeed.accelerator import get_accelerator @@ -58,7 +58,9 @@ def __init__( # this flag shows whether or not using prefix in loading the checkpoint use_load_prefix=False, # whether or not the qkv is stored in the split-format - split_qkv=True): + split_qkv=True, + # Type of normalization to perform + norm_type=NormType.LayerNorm): super().__init__() self.cuda_graph_supported = False self.inference = inference @@ -70,9 +72,10 @@ def __init__( self.pre_attn_norm = pre_attn_norm self.use_load_prefix = use_load_prefix self.split_qkv = split_qkv + self.norm_type = norm_type @abstractmethod - def attention(self, enable_training=False): + def attention(self): """ Returns attention qkv and dense parameters weight: (3*hidden, hidden) and (hidden, hidden) @@ -80,13 +83,6 @@ def attention(self, enable_training=False): """ raise NotImplementedError - @abstractmethod - def get_q_k_v(self): - """ - return all q,k,v parameters without merging them together - """ - raise NotImplementedError - @abstractmethod def get_hidden_heads(self): """ @@ -112,14 +108,6 @@ def layernorm(self): """ raise NotImplementedError - @abstractmethod - def get_lora_params(self): - """ - Returns lora parameters used in transformer layer - - """ - raise NotImplementedError - # TODO (lekurile): This function exists in base container as well, consolidate as some point def transpose(data): @@ -133,7 +121,7 @@ def transpose(data): # TODO (lekurile): This function exists in megatron feature container as well, consolidate as some point def _transpose(x, heads=1, mp_replace=None): - heads = heads // mp_replace.mp_size + heads = heads // mp_replace.mp_size # type: ignore outer_dim = -1 attention_head_size = x.shape[outer_dim] // heads new_x_shape = x.size()[:outer_dim] + (heads, attention_head_size) @@ -164,15 +152,15 @@ def maybe_copy(module, tmp = sd[src_name] if len(dst.shape) == 1: if split_qkv: - dst = mp_replace.qkv_copy(dst, tmp) + dst = mp_replace.strided_copy(dst, tmp, num_splits=3) else: dst = mp_replace.copy(dst, tmp) if qkv and megatron_v2: dst = torch.nn.parameter.Parameter(_transpose(dst, heads=heads, mp_replace=mp_replace).contiguous()) else: if split_qkv: - dst = mp_replace.qkv_copy(dst, weight_quantizer.quantize(tmp if weight_quantizer.q_int8 else \ - (transpose(tmp).contiguous())), int8=weight_quantizer.q_int8) + dst = mp_replace.strided_copy(dst, weight_quantizer.quantize(tmp if weight_quantizer.q_int8 else \ + (transpose(tmp).contiguous())), num_splits=3, int8=weight_quantizer.q_int8) else: if qkv and megatron_v2: tmp = _transpose(transpose(tmp), heads=heads, mp_replace=mp_replace).contiguous() @@ -193,19 +181,33 @@ def maybe_copy_qkv(module, sd, weight_quantizer, mp_replace, dst_name, src_names dst = getattr(module, dst_name) if len(dst.shape) == 1: if split_qkv: - dst = mp_replace.qkv_copy(dst, qkv_data.contiguous()) + dst = mp_replace.strided_copy(dst, qkv_data.contiguous(), num_splits=3) else: dst = mp_replace.copy(dst, qkv_data) else: if split_qkv: - dst = mp_replace.qkv_copy(dst, weight_quantizer.quantize(qkv_data.to(get_accelerator().device_name()) if weight_quantizer.q_int8 else \ - ((transpose(qkv_data)).contiguous())), int8=weight_quantizer.q_int8) + dst = mp_replace.strided_copy(dst, weight_quantizer.quantize(qkv_data.to(get_accelerator().device_name()) if weight_quantizer.q_int8 else \ + ((transpose(qkv_data)).contiguous())), num_splits=3, int8=weight_quantizer.q_int8) else: dst = mp_replace.copy(dst, weight_quantizer.quantize(qkv_data.to(get_accelerator().device_name()) if weight_quantizer.q_int8 else \ transpose(qkv_data)), int8=weight_quantizer.q_int8) setattr(module, dst_name, dst) +# Extending the `maybe_copy` function for when mlp1 is in separate parameters for GeGLU +def maybe_copy_geglu(module, sd, weight_quantizer, mp_replace, dst_name, src_names): + if src_names[0] in sd: + reg_proj = sd[src_names[0]] + gate_proj = sd[src_names[1]] + + mlp1_data = torch.cat((reg_proj, gate_proj), dim=0) + dst = getattr(module, dst_name) + + dst = mp_replace.strided_copy(dst, weight_quantizer.quantize(mlp1_data.to(get_accelerator().device_name()) if weight_quantizer.q_int8 else \ + transpose(mlp1_data)), num_splits=2, int8=weight_quantizer.q_int8) + setattr(module, dst_name, dst) + + def pack_lora_weights(p): return [ p.lora_right_weight, \ diff --git a/deepspeed/module_inject/replace_module.py b/deepspeed/module_inject/replace_module.py index 51309ba2afc1..f1caf6aaabcf 100644 --- a/deepspeed/module_inject/replace_module.py +++ b/deepspeed/module_inject/replace_module.py @@ -4,6 +4,7 @@ # DeepSpeed Team import os +from typing import Optional import torch import tqdm import deepspeed @@ -42,19 +43,23 @@ def merge_assert(self, dim1, dim2): for merging your checkpoints before replacing the transformer layer with\ inference-kernels' - def qkv_copy(self, dst, src, int8=False, allocat_tensor=False): + def strided_copy(self, + dst: Optional[torch.Tensor], + src: Optional[torch.Tensor], + num_splits: int, + int8: bool = False, + allocate_tensor: bool = False): if src is None: return src src_shape = src.shape dst_shape = dst.shape outer_dim = 0 if int8 else -1 - inner_dim = -1 if int8 else 0 - if allocat_tensor: + if allocate_tensor: dst = torch.empty_like(dst) - src_split = torch.split(src.data, src.shape[outer_dim] // 3, dim=outer_dim) + src_split = torch.split(src.data, src.shape[outer_dim] // num_splits, dim=outer_dim) if (len(src_shape) == 2 and len(dst_shape) == 2): if src_shape[outer_dim] == dst_shape[self.out_dim]: dst = dst.reshape(-1).data.copy_(src.data.reshape(-1)).reshape(src.shape) @@ -63,7 +68,7 @@ def qkv_copy(self, dst, src, int8=False, allocat_tensor=False): dst.scale = src.scale return dst self.merge_assert(src_shape[outer_dim], dst_shape[self.out_dim]) - qkv_size = dst_shape[self.out_dim] // 3 + qkv_size = dst_shape[self.out_dim] // num_splits qkv_split = [torch.split(src_s, qkv_size, dim=outer_dim) for src_s in src_split] weight_split = [ torch.cat([qkv_s[i] for qkv_s in qkv_split], axis=outer_dim) for i in range(len(qkv_split[0])) @@ -73,7 +78,7 @@ def qkv_copy(self, dst, src, int8=False, allocat_tensor=False): else: if src_shape[0] == dst_shape[0]: return torch.nn.parameter.Parameter(src) - qkv_size = dst_shape[0] // 3 + qkv_size = dst_shape[0] // num_splits qkv_split = [torch.split(src_s, qkv_size, dim=0) for src_s in src_split] bias_split = [torch.cat([qkv_s[i] for qkv_s in qkv_split], axis=0) for i in range(len(qkv_split[0]))] dst.data.copy_(bias_split[self.gpu_index].contiguous()) @@ -83,11 +88,11 @@ def qkv_copy(self, dst, src, int8=False, allocat_tensor=False): dst.scale = src.scale return dst - def copy(self, dst, src, int8=False, allocat_tensor=False): + def copy(self, dst, src, int8=False, allocate_tensor=False): if src is None: return src assert not dst.data.is_meta # the torch.Tensor.copy_ method used below will silently fail on meta tensors - if allocat_tensor: + if allocate_tensor: dst = torch.empty_like(dst) outer_dim = 0 if int8 else 1 inner_dim = 1 if int8 else 0 @@ -285,7 +290,6 @@ def replace_transformer_layer(orig_layer_impl, model, checkpoint_dict, config, m Updated nn.module with replaced transformer layers """ # defining globals as internally defined functions inherit these everywhere - fp16 = (config.dtype == torch.float16 or config.dtype == torch.int8) quantize = (config.dtype == torch.int8) # todo: Refactor later. In future, let's minimize the style used above and use config.** instead @@ -318,7 +322,6 @@ def replace_with_policy(child, policy_cls, triangular_masking, inference=False, model_config=model_config, layer_id=layer_id, child=child) - _container.set_dtype(fp16) _container.set_moe(moe) # 2. Set the tensor parallelism config @@ -328,12 +331,12 @@ def replace_with_policy(child, policy_cls, triangular_masking, inference=False, _container.initialize_tensors() # 4. deal with data types -- needs refactor to use dtype instead of fp16 - if fp16: - _container.convert_to_required_dtype(dtype=torch.half) + if config.dtype in [torch.float16, torch.bfloat16, torch.int8]: + _container.convert_to_required_dtype() # 5. Set the quantization config quantizer = GroupQuantizer(q_int8=quantize) - _container.set_quantization_config(quantize, quantizer) + _container.set_quantization_config(quantizer) # 6. create a DS Inference config object _container.create_ds_model_config() @@ -598,6 +601,14 @@ def replace_fn(child, _policy, layer_id=0): OrderedDict({k: v for k, v in dict(replaced_module.state_dict()).items() if transformer_name not in k}), f'{config.save_mp_checkpoint_path}/{non_tp_ckpt_name}') + + dtype_reprs = { + torch.float32: 'float32', + torch.float16: 'float16', + torch.int8: 'int8', + torch.bfloat16: 'bfloat16' + } + ckpt_config = json.dumps({ 'type': ckpt_name, 'base_dir': f'{config.save_mp_checkpoint_path}', @@ -608,7 +619,7 @@ def replace_fn(child, _policy, layer_id=0): 'version': 1.0, 'parallelization': 'tp', 'tp_size': world_size, - 'dtype': 'int8' if quantize else ('float16' if fp16 else 'float32') + 'dtype': dtype_reprs[config.dtype] }) with open(f"{config.save_mp_checkpoint_path}/ds_inference_config.json", "w") as cfg: cfg.write(ckpt_config) diff --git a/deepspeed/module_inject/replace_policy.py b/deepspeed/module_inject/replace_policy.py index af58d3d8d2d7..c49b8f81c430 100755 --- a/deepspeed/module_inject/replace_policy.py +++ b/deepspeed/module_inject/replace_policy.py @@ -13,13 +13,14 @@ from .containers import MegatronLayerPolicy from .containers import HFDistilBertLayerPolicy from .containers import HFCLIPLayerPolicy +from .containers import LLAMALayerPolicy from .containers import UNetPolicy from .containers import VAEPolicy # transformer-based policies replace_policies = [ HFBertLayerPolicy, HFGPTNEOLayerPolicy, GPTNEOXLayerPolicy, HFGPTJLayerPolicy, MegatronLayerPolicy, - HFGPT2LayerPolicy, BLOOMLayerPolicy, HFOPTLayerPolicy, HFCLIPLayerPolicy, HFDistilBertLayerPolicy + HFGPT2LayerPolicy, BLOOMLayerPolicy, HFOPTLayerPolicy, HFCLIPLayerPolicy, HFDistilBertLayerPolicy, LLAMALayerPolicy ] # non-transformer-based policies diff --git a/deepspeed/module_inject/utils.py b/deepspeed/module_inject/utils.py index ad60e225fcea..c442d24fd3b6 100644 --- a/deepspeed/module_inject/utils.py +++ b/deepspeed/module_inject/utils.py @@ -17,6 +17,7 @@ def policy_to_ds_container(**kwargs): from .containers import HFOPTLayerPolicy, DS_OPTContainer from .containers import MegatronLayerPolicy, DS_MegatronGPTContainer from .containers import HFDistilBertLayerPolicy, DS_DistilBERTContainer + from .containers import LLAMALayerPolicy, DS_LLAMAContainer policy_to_container = { HFGPT2LayerPolicy: DS_GPT2Container, @@ -28,6 +29,7 @@ def policy_to_ds_container(**kwargs): HFOPTLayerPolicy: DS_OPTContainer, MegatronLayerPolicy: DS_MegatronGPTContainer, HFDistilBertLayerPolicy: DS_DistilBERTContainer, + LLAMALayerPolicy: DS_LLAMAContainer, } container = None diff --git a/deepspeed/ops/transformer/inference/config.py b/deepspeed/ops/transformer/inference/config.py index 549a03a70f19..09f04da5f05c 100644 --- a/deepspeed/ops/transformer/inference/config.py +++ b/deepspeed/ops/transformer/inference/config.py @@ -4,7 +4,8 @@ # DeepSpeed Team import json -from deepspeed.utils.types import ActivationFuncType +import torch +from deepspeed.utils.types import ActivationFuncType, NormType class TransformerConfig(): @@ -52,9 +53,9 @@ def __init__(self, layer_norm_eps=1e-12, local_rank=-1, mp_size=1, - fp16=False, - q_int8=False, + dtype=torch.float16, pre_layer_norm=True, + norm_type=NormType.LayerNorm, stochastic_mode=False, scale_attention=True, triangular_masking=True, @@ -79,13 +80,13 @@ def __init__(self, super(DeepSpeedInferenceConfig, self).__init__(hidden_size, (intermediate_size if intermediate_size > 0 else 4 * hidden_size), heads, num_hidden_layers) - self.fp16 = fp16 + self.dtype = dtype self.pre_layer_norm = pre_layer_norm + self.norm_type = norm_type self.local_rank = local_rank self.stochastic_mode = stochastic_mode self.epsilon = layer_norm_eps self.mp_size = mp_size - self.q_int8 = q_int8 self.scale_attention = scale_attention self.triangular_masking = triangular_masking self.local_attention = local_attention diff --git a/deepspeed/ops/transformer/inference/diffusers_attention.py b/deepspeed/ops/transformer/inference/diffusers_attention.py index 3447f9516ade..5eba7a669701 100644 --- a/deepspeed/ops/transformer/inference/diffusers_attention.py +++ b/deepspeed/ops/transformer/inference/diffusers_attention.py @@ -116,8 +116,8 @@ def __init__( device = get_accelerator().current_device_name() if config.bigscience_bloom else 'cpu' qkv_size_per_partition = (self.config.hidden_size // self.config.mp_size) * 3 - data_type = torch.int8 if config.q_int8 else torch.half if config.fp16 else torch.float - data_type_fp = torch.half if config.fp16 else torch.float + data_type = self.config.dtype + data_type_fp = torch.half if self.config.dtype == torch.int8 else self.config.dtype global inference_cuda_module if inference_cuda_module is None: builder = InferenceBuilder() @@ -172,12 +172,14 @@ def __init__( self.norm_factor *= math.sqrt(self.config.layer_id + 1) # https://github.com/huggingface/transformers/blob/v4.24.0/src/transformers/models/gpt2/modeling_gpt2.py#L191 - self.score_context_func = inference_cuda_module.softmax_context_fp32 if (not config.fp16) else \ - inference_cuda_module.softmax_context_fp16 - self.linear_func = inference_cuda_module.linear_layer_fp16 if config.fp16 else \ - inference_cuda_module.linear_layer_fp32 - self.allocate_workspace = inference_cuda_module.allocate_workspace_fp32 if not (config.fp16) else \ - inference_cuda_module.allocate_workspace_fp16 + if self.config.dtype in [torch.float16, torch.int8]: + self.score_context_func = inference_cuda_module.softmax_context_fp16 + self.linear_func = inference_cuda_module.linear_layer_fp16 + self.allocate_workspace = inference_cuda_module.allocate_workspace_fp16 + else: + self.score_context_func = inference_cuda_module.softmax_context_fp32 + self.linear_func = inference_cuda_module.linear_layer_fp32 + self.allocate_workspace = inference_cuda_module.allocate_workspace_fp32 def forward(self, input, context=None, input_mask=None): if self.config.layer_id == 0: diff --git a/deepspeed/ops/transformer/inference/diffusers_transformer_block.py b/deepspeed/ops/transformer/inference/diffusers_transformer_block.py index 3d45714e543c..76519b47085e 100644 --- a/deepspeed/ops/transformer/inference/diffusers_transformer_block.py +++ b/deepspeed/ops/transformer/inference/diffusers_transformer_block.py @@ -11,6 +11,7 @@ from .bias_add import nhwc_bias_add from .diffusers_2d_transformer import Diffusers2DTransformerConfig from deepspeed.ops.op_builder import InferenceBuilder, SpatialInferenceBuilder +from deepspeed.utils.types import ActivationFuncType # Ops will be loaded on demand transformer_cuda_module = None @@ -97,7 +98,7 @@ def forward(self, hidden_states, context=None, timestep=None, **kwargs): out_attn_2, self.attn_2_bias, out_attn_1, self.norm3_g, self.norm3_b, self.norm3_eps) out_ff1 = nn.functional.linear(out_norm_3, self.ff1_w) - out_geglu = self.transformer_cuda_module.bias_geglu(out_ff1, self.ff1_b) + out_geglu = self.transformer_cuda_module.gated_activation(out_ff1, self.ff1_b, ActivationFuncType.GATED_GELU) out_ff2 = nn.functional.linear(out_geglu, self.ff2_w) return nhwc_bias_add(out_ff2, self.ff2_b, other=out_attn_2) diff --git a/deepspeed/ops/transformer/inference/ds_attention.py b/deepspeed/ops/transformer/inference/ds_attention.py index d56f8bc7de4a..967f1d4b8d9d 100644 --- a/deepspeed/ops/transformer/inference/ds_attention.py +++ b/deepspeed/ops/transformer/inference/ds_attention.py @@ -20,8 +20,8 @@ class DeepSpeedSelfAttention(nn.Module): def __init__(self, config, mp_group=None, q_scales=None, q_groups=1, merge_count=1): super(DeepSpeedSelfAttention, self).__init__() self.config = config - data_type = torch.int8 if config.q_int8 else torch.half if config.fp16 else torch.float - data_type_fp = torch.half if config.fp16 else torch.float + data_type = self.config.dtype + data_type_fp = torch.half if self.config.dtype == torch.int8 else self.config.dtype self.config.layer_id = DeepSpeedSelfAttention.num_layers DeepSpeedSelfAttention.num_layers = DeepSpeedSelfAttention.num_layers + 1 device = get_accelerator().current_device_name() #if config.bigscience_bloom else 'cpu' @@ -88,7 +88,7 @@ def __init__(self, config, mp_group=None, q_scales=None, q_groups=1, merge_count ] def compute_attention(self, qkv_out, input_mask, layer_past, alibi): - if isinstance(qkv_out, list): + if isinstance(qkv_out, list) or isinstance(qkv_out, tuple): qkv_out = qkv_out[0] no_masking = input_mask is None @@ -112,14 +112,14 @@ def compute_attention(self, qkv_out, input_mask, layer_past, alibi): def _merge_qkv(self): qvkw = DeepSpeedSelfAttention._qkv_buffers[0] - qvkw[:self.hidden_size_per_partition, :] = self.attn_qw - qvkw[self.hidden_size_per_partition:2 * self.hidden_size_per_partition, :] = self.attn_kw - qvkw[2 * self.hidden_size_per_partition:, :] = self.attn_vw + qvkw[:self.hidden_size_per_partition, :] = self.attn_qw # type: ignore + qvkw[self.hidden_size_per_partition:2 * self.hidden_size_per_partition, :] = self.attn_kw # type: ignore + qvkw[2 * self.hidden_size_per_partition:, :] = self.attn_vw # type: ignore if self.attn_qb is not None: qvkb = DeepSpeedSelfAttention._qkv_buffers[1] qvkb[:self.hidden_size_per_partition] = self.attn_qb - qvkb[self.hidden_size_per_partition:2 * self.hidden_size_per_partition] = self.attn_kb - qvkb[2 * self.hidden_size_per_partition:] = self.attn_vb + qvkb[self.hidden_size_per_partition:2 * self.hidden_size_per_partition] = self.attn_kb # type: ignore + qvkb[2 * self.hidden_size_per_partition:] = self.attn_vb # type: ignore return DeepSpeedSelfAttention._qkv_buffers def forward(self, @@ -151,12 +151,10 @@ def forward(self, else: qkv_out = self.qkv_func(input=input, weight=self._attn_qkvw, - bias=(self._attn_qkvb if self._attn_qkvb is not None else norm_b), + bias=self._attn_qkvb, gamma=norm_w, - beta=norm_b, - add_bias=(self.attn_qkvb is not None), - num_layers=DeepSpeedSelfAttention.num_layers, - num_heads=self.num_attention_heads_per_partition) + beta=norm_b) + context_layer, key_layer, value_layer = self.compute_attention(qkv_out=qkv_out, input_mask=input_mask, layer_past=layer_past, @@ -211,7 +209,7 @@ def _split_tensor_along_last_dim(self, tensor, num_partitions, contiguous_split_ return tensor_list def compute_attention(self, qkv_out, input_mask, layer_past, alibi): - if isinstance(qkv_out, list): + if isinstance(qkv_out, list) or isinstance(qkv_out, tuple): qkv_out = qkv_out[0] no_masking = input_mask is None @@ -248,8 +246,9 @@ def compute_attention(self, qkv_out, input_mask, layer_past, alibi): attention_scores = matmul_result.view(output_size[0], output_size[1], output_size[2], -1) offset = dist.get_rank() * self.num_attention_heads_per_partition if dist.is_initialized() else 0 + target_dtype = torch.float16 if self.config.dtype == torch.int8 else self.config.dtype attention_probs = self.softmax_func(attn_scores=attention_scores, - attn_mask=((1 - input_mask).half() * minus_inf), + attn_mask=((1 - input_mask).to(target_dtype) * minus_inf), alibi=alibi, triangular=(self.config.triangular_masking and (attention_scores.shape[-2] > 1)), diff --git a/deepspeed/ops/transformer/inference/ds_mlp.py b/deepspeed/ops/transformer/inference/ds_mlp.py index a4375178347a..f4bb538dab37 100644 --- a/deepspeed/ops/transformer/inference/ds_mlp.py +++ b/deepspeed/ops/transformer/inference/ds_mlp.py @@ -7,24 +7,37 @@ import torch import torch.nn as nn from deepspeed import comm as dist +from deepspeed.utils.types import GATED_ACTIVATION_TYPES from deepspeed.accelerator import get_accelerator from .op_binding import MLPGemmOp, VectorMatMulOp, GELUGemmOp, ResidualAddOp class DeepSpeedMLP(nn.Module): + _inter_w_buffers = [] def __init__(self, config, mp_group=None, q_scales=None, q_groups=1, merge_count=1, mlp_extra_grouping=False): super(DeepSpeedMLP, self).__init__() self.config = config - data_type = torch.int8 if config.q_int8 else torch.half if config.fp16 else torch.float - data_type_fp = torch.half if config.fp16 else torch.float + + data_type = torch.half if self.config.dtype == torch.int8 else self.config.dtype + data_type_fp = data_type device = get_accelerator().current_device_name() + + proj_factor = 2 if self.config.mlp_act_func_type in GATED_ACTIVATION_TYPES else 1 + self.config.intermediate_size = self.config.intermediate_size if self.config.intermediate_size > 0 else 4 * self.config.hidden_size + self.intm_w_sz_per_partition = self.config.intermediate_size * proj_factor // self.config.mp_size + self.intm_o_sz_per_partition = self.config.intermediate_size // self.config.mp_size + if self.config.set_empty_params: self.attn_nw = None self.attn_nb = None self.inter_w = None self.inter_b = None + self.inter_up_w = None + self.inter_up_b = None + self.inter_gate_w = None + self.inter_gate_b = None self.output_w = None self.output_b = None else: @@ -32,15 +45,15 @@ def __init__(self, config, mp_group=None, q_scales=None, q_groups=1, merge_count requires_grad=False) self.attn_nb = nn.Parameter(torch.empty(self.config.hidden_size, dtype=data_type_fp, device=device), requires_grad=False) - intm_size_per_partition = self.config.intermediate_size // self.config.mp_size + self.inter_w = nn.Parameter(torch.empty(self.config.hidden_size, - intm_size_per_partition, + self.intm_w_sz_per_partition, dtype=data_type, device=device), requires_grad=False) - self.inter_b = nn.Parameter(torch.empty(intm_size_per_partition, dtype=data_type_fp, device=device), + self.inter_b = nn.Parameter(torch.empty(self.intm_w_sz_per_partition, dtype=data_type_fp, device=device), requires_grad=False) - self.output_w = nn.Parameter(torch.empty(intm_size_per_partition, + self.output_w = nn.Parameter(torch.empty(self.intm_o_sz_per_partition, self.config.hidden_size, dtype=data_type, device=device), @@ -59,7 +72,30 @@ def __init__(self, config, mp_group=None, q_scales=None, q_groups=1, merge_count self.fused_gemm_gelu = GELUGemmOp(config) self.residual_add_func = ResidualAddOp(config) + if len(DeepSpeedMLP._inter_w_buffers) == 0: + DeepSpeedMLP._inter_w_buffers = [ + torch.empty(self.config.hidden_size, self.intm_w_sz_per_partition, dtype=data_type, device=device), + torch.empty(self.intm_w_sz_per_partition, dtype=data_type_fp, device=device) + ] + + def _merge_inter_w(self): + inter_w = DeepSpeedMLP._inter_w_buffers[0] + inter_w[:self.intm_w_sz_per_partition, :] = self.inter_up_w # type: ignore + inter_w[self.intm_w_sz_per_partition:, :] = self.inter_gate_w # type: ignore + if self.inter_up_b is not None: + inter_b = DeepSpeedMLP._inter_w_buffers[1] + inter_b[:self.intm_w_sz_per_partition] = self.inter_up_b # type: ignore + inter_b[self.intm_w_sz_per_partition:] = self.inter_gate_b # type: ignore + return DeepSpeedMLP._inter_w_buffers + def forward(self, input, residual, residual_norm, bias): + + if self.inter_w is None: + self._inter_w, self._inter_b = self._merge_inter_w() + else: + self._inter_w = self.inter_w + self._inter_b = self.inter_b + residual_add = None if self.attn_nw is None: output = self.fused_gemm_gelu(input=residual_norm, @@ -69,19 +105,21 @@ def forward(self, input, residual, residual_norm, bias): else: output, residual_add = self.mlp_gemm_func(input=input, residual=residual, - input_bias=bias, weight_interm=self.inter_w, weight_out=self.output_w, + input_bias=bias, bias=self.inter_b, gamma=self.attn_nw, beta=self.attn_nb) + residual = self.residual_add_func(hidden_state=output, residual=residual, + add_bias=bias is not None, attention_output=input, attention_bias=bias if bias is not None else self.output_b, final_bias=self.output_b, - add_bias=bias is not None, residual_add=residual_add) if self.mp_group is not None and dist.get_world_size(group=self.mp_group) > 1: dist.all_reduce(residual, group=self.mp_group) + return residual diff --git a/deepspeed/ops/transformer/inference/moe_inference.py b/deepspeed/ops/transformer/inference/moe_inference.py index bf14a5fc36b2..c828d94db962 100644 --- a/deepspeed/ops/transformer/inference/moe_inference.py +++ b/deepspeed/ops/transformer/inference/moe_inference.py @@ -200,6 +200,7 @@ def __init__(self, else: inference_cuda_module = InferenceBuilder().load() self.config.specialized_mode = specialized_mode + assert self.config.dtype != torch.bfloat16, "DeepSpeed MoE Transformer Inference not yet tested for bfloat support" DeepSpeedMoEInference.layer_id += 1 self.attention = DeepSpeedSelfAttention(self.config, mp_group, quantize_scales, quantize_groups, merge_count) @@ -213,9 +214,9 @@ def __init__(self, self.res_mlp = DeepSpeedMoEMLP(config, quantize_scales, quantize_groups, merge_count, mlp_extra_grouping, mp_group) self.res_coef = nn.Parameter(torch.Tensor(self.config.hidden_size, 2)) - self.coef_func = inference_cuda_module.softmax_fp16 if self.config.fp16 or self.config.q_int8 else \ + self.coef_func = inference_cuda_module.softmax_fp16 if self.config.dtype in [torch.float16, torch.int8] else \ inference_cuda_module.softmax_fp32 - self.vector_matmul_func = inference_cuda_module.vector_matmul_fp16 if config.fp16 else \ + self.vector_matmul_func = inference_cuda_module.vector_matmul_fp16 if self.config.dtype == torch.float16 else \ inference_cuda_module.vector_matmul_fp32 config.mp_size = 1 @@ -234,11 +235,11 @@ def __init__(self, print("DeepSpeed MoE Transformer Inference config is ", self.config.__dict__) - self.bias_residual_func = inference_cuda_module.bias_residual_fp16 if config.fp16 or config.q_int8 else \ + self.bias_residual_func = inference_cuda_module.bias_residual_fp16 if self.config.dtype in [torch.float16, torch.int8] else \ inference_cuda_module.bias_residual_fp32 - self.ds_layernorm = inference_cuda_module.layer_norm_fp16 if self.config.fp16 or self.config.q_int8 else \ + self.ds_layernorm = inference_cuda_module.layer_norm_fp16 if self.config.dtype in [torch.float16, torch.int8] else \ inference_cuda_module.layer_norm_fp32 - self.einsum_sec_sm_ecm = inference_cuda_module.einsum_sec_sm_ecm_fp16 if self.config.fp16 or self.config.q_int8 else \ + self.einsum_sec_sm_ecm = inference_cuda_module.einsum_sec_sm_ecm_fp16 if self.config.dtype in [torch.float16, torch.int8] else \ inference_cuda_module.einsum_sec_sm_ecm_fp32 def res_coef_func(self, inp, async_op): @@ -302,8 +303,7 @@ def forward(self, input_mask = input_mask if attention_mask is None else attention_mask input_type = input.dtype - if (self.config.fp16 or self.config.q_int8) \ - and input.dtype == torch.float: + if (self.config.dtype in [torch.float16, torch.int8]) and input_type == torch.float: input = input.half() with torch.no_grad(): diff --git a/deepspeed/ops/transformer/inference/op_binding/gelu_gemm.py b/deepspeed/ops/transformer/inference/op_binding/gelu_gemm.py index 06fd6f560cb0..1f15c2ac2e1e 100644 --- a/deepspeed/ops/transformer/inference/op_binding/gelu_gemm.py +++ b/deepspeed/ops/transformer/inference/op_binding/gelu_gemm.py @@ -12,20 +12,23 @@ class GELUGemmOp(BaseOp): def __init__(self, config: DeepSpeedInferenceConfig): super(GELUGemmOp, self).__init__(config) - if self.config.fp16: - self.fused_gemm_gelu = self.inference_cuda_module.fused_gemm_gelu_fp16 + if self.config.dtype in [torch.float16, torch.int8]: + self.fused_gemm_gelu = self.inference_cuda_module.fused_gemm_gelu_fp16 # type: ignore + elif self.config.dtype == torch.bfloat16: + self.fused_gemm_gelu = self.inference_cuda_module.fused_gemm_gelu_bf16 else: - self.fused_gemm_gelu = self.inference_cuda_module.fused_gemm_gelu_fp32 - - def forward(self, - input: torch.Tensor, - weight: torch.Tensor, - bias: torch.Tensor, - weight_out: torch.Tensor, - async_op: bool = False): - output = self.fused_gemm_gelu(input, weight, weight.scale if hasattr(weight, 'scale') else torch.empty(1), - bias, weight_out, - weight_out.scale if hasattr(weight_out, 'scale') else torch.empty(1), - self.config.epsilon, self.config.pre_layer_norm, self.config.q_int8, async_op, - self.config.transposed_mode) + self.fused_gemm_gelu = self.inference_cuda_module.fused_gemm_gelu_fp32 # type: ignore + + def forward(self, input: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, weight_out: torch.Tensor): + + output = self.fused_gemm_gelu( + input, + weight, + weight.scale if hasattr(weight, 'scale') else torch.empty(1), # type: ignore + bias, + weight_out, + weight_out.scale if hasattr(weight_out, 'scale') else torch.empty(1), # type: ignore + self.config.dtype == torch.int8, + self.config.transposed_mode) + return output diff --git a/deepspeed/ops/transformer/inference/op_binding/linear.py b/deepspeed/ops/transformer/inference/op_binding/linear.py index 9178c5f1fc5b..ffbf848104d7 100644 --- a/deepspeed/ops/transformer/inference/op_binding/linear.py +++ b/deepspeed/ops/transformer/inference/op_binding/linear.py @@ -12,8 +12,10 @@ class LinearOp(BaseOp): def __init__(self, config: DeepSpeedInferenceConfig): super(LinearOp, self).__init__(config) - if self.config.fp16: + if self.config.dtype in [torch.float16, torch.int8]: self.linear_func = self.inference_cuda_module.linear_layer_fp16 + elif self.config.dtype == torch.bfloat16: + self.linear_func = self.inference_cuda_module.linear_layer_bf16 else: self.linear_func = self.inference_cuda_module.linear_layer_fp32 diff --git a/deepspeed/ops/transformer/inference/op_binding/mlp_gemm.py b/deepspeed/ops/transformer/inference/op_binding/mlp_gemm.py index e7ca40219c34..92d04d6aa761 100644 --- a/deepspeed/ops/transformer/inference/op_binding/mlp_gemm.py +++ b/deepspeed/ops/transformer/inference/op_binding/mlp_gemm.py @@ -3,27 +3,72 @@ # DeepSpeed Team +from typing import Optional + import torch from ..config import DeepSpeedInferenceConfig from .base import BaseOp +from deepspeed.utils.types import NormType class MLPGemmOp(BaseOp): def __init__(self, config: DeepSpeedInferenceConfig): super(MLPGemmOp, self).__init__(config) - if self.config.fp16: - self.mlp_gemm_func = self.inference_cuda_module.mlp_gemm_fp16 + + if self.config.norm_type == NormType.LayerNorm: + if self.config.dtype in [torch.float16, torch.int8]: + self.mlp_gemm_func = self.inference_cuda_module.mlp_gemm_fp16 # type: ignore + elif self.config.dtype == torch.bfloat16: + self.mlp_gemm_func = self.inference_cuda_module.mlp_gemm_bf16 + else: + self.mlp_gemm_func = self.inference_cuda_module.mlp_gemm_fp32 # type: ignore + elif self.config.norm_type == NormType.RMSNorm: + if self.config.dtype in [torch.float16, torch.int8]: + self.mlp_gemm_func = self.inference_cuda_module.rms_mlp_gemm_fp16 # type: ignore + elif self.config.dtype == torch.bfloat16: + self.mlp_gemm_func = self.inference_cuda_module.rms_mlp_gemm_bf16 + else: + self.mlp_gemm_func = self.inference_cuda_module.rms_mlp_gemm_fp32 # type: ignore + + def forward(self, + input: torch.Tensor, + residual: torch.Tensor, + weight_interm: torch.Tensor, + weight_out: torch.Tensor, + input_bias: Optional[torch.Tensor] = None, + bias: Optional[torch.Tensor] = None, + gamma: Optional[torch.Tensor] = None, + beta: Optional[torch.Tensor] = None): + if self.config.norm_type == NormType.LayerNorm: + output, residual_add = self.mlp_gemm_func( + input, + residual, + input_bias, + weight_interm, + weight_out, + bias, + gamma, + beta, + self.config.epsilon, + self.config.pre_layer_norm, + self.config.mlp_after_attn, + weight_interm.scale if hasattr(weight_interm, 'scale') else torch.empty(1), # type: ignore + weight_out.scale if hasattr(weight_out, 'scale') else torch.empty(1), # type: ignore + self.config.dtype == torch.int8, + self.config.mlp_act_func_type, + self.config.transposed_mode) else: - self.mlp_gemm_func = self.inference_cuda_module.mlp_gemm_fp32 - - def forward(self, input: torch.Tensor, residual: torch.Tensor, input_bias: torch.Tensor, - weight_interm: torch.Tensor, weight_out: torch.Tensor, bias: torch.Tensor, gamma: torch.Tensor, - beta: torch.Tensor): - output, residual_add = self.mlp_gemm_func( - input, residual, input_bias, weight_interm, weight_out, bias, gamma, beta, self.config.epsilon, - self.config.pre_layer_norm, self.config.mlp_after_attn, - weight_interm.scale if hasattr(weight_interm, 'scale') else torch.empty(1), - weight_out.scale if hasattr(weight_out, 'scale') else torch.empty(1), self.config.q_int8, - self.config.mlp_act_func_type, self.config.transposed_mode) + output, residual_add = self.mlp_gemm_func( + input, + residual, + weight_interm, + weight_out, + gamma, + self.config.epsilon, + weight_interm.scale if hasattr(weight_interm, 'scale') else torch.empty(1), # type: ignore + weight_out.scale if hasattr(weight_out, 'scale') else torch.empty(1), # type: ignore + self.config.dtype == torch.int8, + self.config.mlp_act_func_type, + self.config.transposed_mode) return output, residual_add diff --git a/deepspeed/ops/transformer/inference/op_binding/qkv_gemm.py b/deepspeed/ops/transformer/inference/op_binding/qkv_gemm.py index 6b338b9041d9..074503d96f05 100644 --- a/deepspeed/ops/transformer/inference/op_binding/qkv_gemm.py +++ b/deepspeed/ops/transformer/inference/op_binding/qkv_gemm.py @@ -6,33 +6,42 @@ import torch from ..config import DeepSpeedInferenceConfig from .base import BaseOp -from deepspeed import comm as dist +from deepspeed.utils.types import NormType class QKVGemmOp(BaseOp): def __init__(self, config: DeepSpeedInferenceConfig): super(QKVGemmOp, self).__init__(config) - if self.config.fp16: - self.qkv_gemm_func = self.inference_cuda_module.qkv_gemm_fp16 + + if self.config.norm_type == NormType.LayerNorm: + if self.config.dtype in [torch.float16, torch.int8]: + self.qkv_gemm_func = self.inference_cuda_module.qkv_gemm_fp16 # type: ignore + elif self.config.dtype == torch.bfloat16: + self.qkv_gemm_func = self.inference_cuda_module.qkv_gemm_bf16 + else: + self.qkv_gemm_func = self.inference_cuda_module.qkv_gemm_fp32 # type: ignore + elif self.config.norm_type == NormType.RMSNorm: + if self.config.dtype in [torch.float16, torch.int8]: + self.qkv_gemm_func = self.inference_cuda_module.rms_qkv_gemm_fp16 # type: ignore + elif self.config.dtype == torch.bfloat16: + self.qkv_gemm_func = self.inference_cuda_module.rms_qkv_gemm_bf16 + else: + self.qkv_gemm_func = self.inference_cuda_module.rms_qkv_gemm_fp32 # type: ignore + + def forward(self, input: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, gamma: torch.Tensor, + beta: torch.Tensor): + + add_bias = bias is not None + bias = bias if add_bias else torch.empty(1) # type: ignore + q_scale = weight.scale if hasattr(weight, 'scale') else torch.empty(1) # type: ignore + q_int8 = self.config.dtype == torch.int8 + + if self.config.norm_type == NormType.LayerNorm: + output, norm = self.qkv_gemm_func(input, weight, q_scale, bias, gamma, beta, self.config.epsilon, add_bias, + q_int8, self.config.transposed_mode) else: - self.qkv_gemm_func = self.inference_cuda_module.qkv_gemm_fp32 - - def forward(self, - input: torch.Tensor, - weight: torch.Tensor, - bias: torch.Tensor, - gamma: torch.Tensor, - beta: torch.Tensor, - add_bias: bool, - num_layers: int, - num_heads: int = None, - max_out_tokens: int = None): - q_scale = weight.scale if hasattr(weight, 'scale') else torch.empty(1) - external_cache = self.config.bigscience_bloom - rank = dist.get_rank() if dist.is_initialized() else 0 - q_int8 = self.config.q_int8 - output = self.qkv_gemm_func(input, weight, q_scale, bias, gamma, beta, self.config.epsilon, add_bias, - num_layers, external_cache, self.config.mp_size, rank, q_int8, - self.config.transposed_mode) - return output + output, norm = self.qkv_gemm_func(input, weight, q_scale, gamma, self.config.epsilon, q_int8, + self.config.transposed_mode) + + return output, norm diff --git a/deepspeed/ops/transformer/inference/op_binding/residual_add.py b/deepspeed/ops/transformer/inference/op_binding/residual_add.py index e79f5dee5c54..f545c86ac4b8 100644 --- a/deepspeed/ops/transformer/inference/op_binding/residual_add.py +++ b/deepspeed/ops/transformer/inference/op_binding/residual_add.py @@ -4,6 +4,7 @@ # DeepSpeed Team import torch +from typing import Optional from ..config import DeepSpeedInferenceConfig from .base import BaseOp @@ -12,18 +13,31 @@ class ResidualAddOp(BaseOp): def __init__(self, config: DeepSpeedInferenceConfig): super(ResidualAddOp, self).__init__(config) - if self.config.fp16 or self.config.q_int8: + if self.config.dtype in [torch.float16, torch.int8]: self.residual_add_func = self.inference_cuda_module.residual_add_bias_fp16 + elif self.config.dtype == torch.bfloat16: + self.residual_add_func = self.inference_cuda_module.residual_add_bias_bf16 else: self.residual_add_func = self.inference_cuda_module.residual_add_bias_fp32 + self._vector_add = self.inference_cuda_module._vector_add + + def forward(self, + hidden_state: torch.Tensor, + residual: torch.Tensor, + add_bias: bool, + attention_output: Optional[torch.Tensor] = None, + residual_add: Optional[torch.Tensor] = None, + attention_bias: Optional[torch.Tensor] = None, + final_bias: Optional[torch.Tensor] = None): + + if final_bias is None: + residual = self._vector_add(residual, hidden_state, 1.0 / self.config.mp_size) + else: + if not self.config.pre_layer_norm and residual_add is not None: + # only use residual add if its set and we are not pre layer norm + residual = residual_add - def forward(self, hidden_state: torch.Tensor, residual: torch.Tensor, attention_output: torch.Tensor, - attention_bias: torch.Tensor, final_bias: torch.Tensor, add_bias: bool, residual_add: torch.Tensor): - - if not self.config.pre_layer_norm and residual_add is not None: - # only use residual add if its set and we are not pre layer norm - residual = residual_add - - self.residual_add_func(hidden_state, residual, attention_output, attention_bias, final_bias, - self.config.mp_size, self.config.mlp_after_attn, add_bias, self.config.pre_layer_norm) + self.residual_add_func(hidden_state, residual, attention_output, attention_bias, final_bias, + self.config.mp_size, self.config.mlp_after_attn, add_bias, + self.config.pre_layer_norm) return residual diff --git a/deepspeed/ops/transformer/inference/op_binding/softmax.py b/deepspeed/ops/transformer/inference/op_binding/softmax.py index 529df9ed6181..21ec1999cdae 100644 --- a/deepspeed/ops/transformer/inference/op_binding/softmax.py +++ b/deepspeed/ops/transformer/inference/op_binding/softmax.py @@ -12,10 +12,12 @@ class SoftmaxOp(BaseOp): def __init__(self, config: DeepSpeedInferenceConfig): super(SoftmaxOp, self).__init__(config) - if self.config.fp16: + if self.config.dtype in [torch.float16, torch.int8]: self.softmax_func = self.inference_cuda_module.softmax_fp16 + elif self.config.dtype == torch.bfloat16: + self.softmax_func = self.inference_cuda_module.softmax_bf16 else: - self.softmax_func = self._not_implemented + self.softmax_func = self.inference_cuda_module.softmax_fp32 def _not_implemented(self, *args, **kwargs): raise NotImplementedError diff --git a/deepspeed/ops/transformer/inference/op_binding/softmax_context.py b/deepspeed/ops/transformer/inference/op_binding/softmax_context.py index 1a132982aba6..d30d5fbeb86d 100644 --- a/deepspeed/ops/transformer/inference/op_binding/softmax_context.py +++ b/deepspeed/ops/transformer/inference/op_binding/softmax_context.py @@ -13,8 +13,10 @@ class SoftmaxContextOp(BaseOp): def __init__(self, config: DeepSpeedInferenceConfig): super(SoftmaxContextOp, self).__init__(config) - if self.config.fp16: + if self.config.dtype in [torch.float16, torch.int8]: self.softmax_context_func = self.inference_cuda_module.softmax_context_fp16 + elif self.config.dtype == torch.bfloat16: + self.softmax_context_func = self.inference_cuda_module.softmax_context_bf16 else: self.softmax_context_func = self.inference_cuda_module.softmax_context_fp32 @@ -32,4 +34,5 @@ def forward(self, query_key_value: torch.Tensor, attn_mask: torch.Tensor, heads: self.config.rotate_every_two, heads, norm_factor, self.config.triangular_masking, self.config.local_attention, self.config.window_size, no_masking, layer_id, num_layers, alibi) + return output diff --git a/deepspeed/ops/transformer/inference/op_binding/vector_matmul.py b/deepspeed/ops/transformer/inference/op_binding/vector_matmul.py index f916020baa9e..65bd0db8f44b 100644 --- a/deepspeed/ops/transformer/inference/op_binding/vector_matmul.py +++ b/deepspeed/ops/transformer/inference/op_binding/vector_matmul.py @@ -12,13 +12,15 @@ class VectorMatMulOp(BaseOp): def __init__(self, config: DeepSpeedInferenceConfig): super(VectorMatMulOp, self).__init__(config) - if self.config.fp16: + if self.config.dtype in [torch.float16, torch.int8]: self.vector_matmul_func = self.inference_cuda_module.vector_matmul_fp16 + elif self.config.dtype == torch.bfloat16: + self.vector_matmul_func = self.inference_cuda_module.vector_matmul_bf16 else: self.vector_matmul_func = self.inference_cuda_module.vector_matmul_fp32 def forward(self, input: torch.Tensor, weight: torch.Tensor, async_op: bool = False): q_scale = weight.scale if hasattr(weight, 'scale') else torch.empty(1) - q_int8 = self.config.q_int8 + q_int8 = self.config.dtype == torch.int8 output = self.vector_matmul_func(input, weight, async_op, q_scale, q_int8, self.config.transposed_mode) return output diff --git a/deepspeed/runtime/hybrid_engine.py b/deepspeed/runtime/hybrid_engine.py index b15447fad412..3ec95c0df67c 100644 --- a/deepspeed/runtime/hybrid_engine.py +++ b/deepspeed/runtime/hybrid_engine.py @@ -85,14 +85,16 @@ def new_inference_container(self, orig_layer, policy_cls, layer_id): policy = policy_cls(orig_layer, inference=True) _container = policy_to_ds_container( policy=policy, - config=DeepSpeedInferenceConfig(set_empty_params=True, - max_out_tokens=self._config.hybrid_engine.max_out_tokens, - min_out_tokens=self._config.hybrid_engine.max_out_tokens, - transposed_mode=True), + config=DeepSpeedInferenceConfig( + set_empty_params=True, + dtype=torch.float16 if self._config.fp16_enabled else torch.float32, + max_out_tokens=self._config.hybrid_engine.max_out_tokens, + min_out_tokens=self._config.hybrid_engine.max_out_tokens, + transposed_mode=True, + ), model_config=self.module.config if hasattr(self.module, 'config') else None, layer_id=layer_id, child=orig_layer) - _container.set_dtype(self._config.fp16_enabled) if self.mpu is not None: if hasattr(self.mpu, 'get_model_parallel_world_size'): @@ -205,8 +207,8 @@ def generate(self, *inputs, **kwargs): self._fuse_lora(self.layer_params[layer_id], self.lora_params[layer_id]) if self.mpu is not None: - self._inference_containers[layer_id].apply_tensor_parallelism( - mp_group=self.mp_group, tp_size=self._config.hybrid_engine.inference_tp_size) + self._inference_containers[layer_id].apply_tensor_parallelism(self.mp_replace, + reversed_dim=True) # TODO(cmikeh2) Evaluate if this can be deferred when release_inference_cache # is enabled. @@ -329,12 +331,29 @@ def create_inference_module(self): ) mp_group = dist.new_group(ranks) if global_rank in ranks: + # mp_group is used for broader collective self.mp_group = mp_group + + # mp_replace is used for container tensor slicing + from deepseed.module_inject import ReplaceWithTensorSlicing + self.mp_replace = ReplaceWithTensorSlicing( + mp_group=self.mp_group, + mp_size=self._config.hybrid_engine.inference_tp_size, + out_dim=0, + in_dim=1) + else: self.mp_group = self.mpu.get_model_parallel_group() if hasattr(self.mpu, 'get_model_parallel_group') else \ self.mpu.get_tensor_model_parallel_group() + + from deepseed.module_inject import ReplaceWithTensorSlicing + self.mp_replace = ReplaceWithTensorSlicing(mp_group=self.mp_group, + mp_size=self._config.hybrid_engine.inference_tp_size, + out_dim=0, + in_dim=1) else: self.mp_group = None + self.mp_replace = None self.populate_all_inference_policies() self.all_layers_params = list(self.module.parameters()) self.create_inference_containers(self.module) @@ -389,7 +408,7 @@ def eval(self): else: orig_module.forward = inference_container.module.forward - inference_container.align_merged_qkv() + inference_container.transform_for_inference() if not self.Z3_enabled or self.gather_all_layers: for orig_module, inference_layer in zip(self._orig_modules_others, self._other_layers): @@ -404,7 +423,7 @@ def train(self, mode=True): if mode and len(self._orig_modules) > 0: for inference_container, orig_module, orig_fwd in zip(self._inference_containers, self._orig_modules, self._orig_fwds): - inference_container.partition_merged_qkv() + inference_container.transform_for_training() orig_module.forward = orig_fwd for orig_module, orig_fwd in zip(self._orig_modules_others, self._orig_fwds_others): orig_module.forward = orig_fwd @@ -419,7 +438,7 @@ def step(self, lr_kwargs=None): if(self._inference_containers[0].module.attention.attn_qkvw is not None and \ self._inference_containers[0].q_k_v is not None): for inference_container in self._inference_containers: - inference_container.reset_qkv() + inference_container.reset_params() if self._training_start_time is not None: self._training_latency += (time.time() - self._training_start_time) diff --git a/deepspeed/utils/types.py b/deepspeed/utils/types.py index 2de4350fbd7a..96b5df625965 100644 --- a/deepspeed/utils/types.py +++ b/deepspeed/utils/types.py @@ -10,3 +10,18 @@ class ActivationFuncType(IntEnum): UNKNOWN = 0 GELU = 1 ReLU = 2 + GATED_GELU = 3 + GATED_SILU = 4 + + +GATED_ACTIVATION_TYPES = [ + ActivationFuncType.GATED_GELU, + ActivationFuncType.GATED_SILU, +] + + +class NormType(IntEnum): + UNKNOWN = 0 + LayerNorm = 1 + GroupNorm = 2 + RMSNorm = 3 diff --git a/op_builder/builder.py b/op_builder/builder.py index e64866e5b93e..4a13d58a1685 100644 --- a/op_builder/builder.py +++ b/op_builder/builder.py @@ -485,11 +485,19 @@ def jit_load(self, verbose=True): torch_arch_list = os.environ.get("TORCH_CUDA_ARCH_LIST") os.environ["TORCH_CUDA_ARCH_LIST"] = "" + nvcc_args = self.strip_empty_entries(self.nvcc_args()) + cxx_args = self.strip_empty_entries(self.cxx_args()) + + if isinstance(self, CUDAOpBuilder): + if not self.build_for_cpu and self.enable_bf16: + cxx_args.append("-DBF16_AVAILABLE") + nvcc_args.append("-DBF16_AVAILABLE") + op_module = load(name=self.name, sources=self.strip_empty_entries(sources), extra_include_paths=self.strip_empty_entries(extra_include_paths), - extra_cflags=self.strip_empty_entries(self.cxx_args()), - extra_cuda_cflags=self.strip_empty_entries(self.nvcc_args()), + extra_cflags=cxx_args, + extra_cuda_cflags=nvcc_args, extra_ldflags=self.strip_empty_entries(self.extra_ldflags()), verbose=verbose) @@ -555,12 +563,16 @@ def compute_capability_args(self, cross_compile_archs=None): f"Unable to load {self.name} op due to no compute capabilities remaining after filtering") args = [] + self.enable_bf16 = True for cc in ccs: num = cc[0] + cc[2] args.append(f'-gencode=arch=compute_{num},code=sm_{num}') if cc.endswith('+PTX'): args.append(f'-gencode=arch=compute_{num},code=compute_{num}') + if int(cc[0]) <= 7: + self.enable_bf16 = False + return args def filter_ccs(self, ccs: List[str]): @@ -602,6 +614,9 @@ def builder(self): {'cxx': self.strip_empty_entries(self.cxx_args()), \ 'nvcc': self.strip_empty_entries(self.nvcc_args())} + if not self.build_for_cpu and self.enable_bf16: + compile_args['cxx'].append("-DBF16_AVAILABLE") + cuda_ext = ExtensionBuilder(name=self.absolute_name(), sources=self.strip_empty_entries(self.sources()), include_dirs=self.strip_empty_entries(self.include_paths()), diff --git a/op_builder/transformer_inference.py b/op_builder/transformer_inference.py index c7b95883cebf..5ee902289448 100755 --- a/op_builder/transformer_inference.py +++ b/op_builder/transformer_inference.py @@ -56,10 +56,12 @@ def sources(self): 'csrc/transformer/inference/csrc/gelu.cu', 'csrc/transformer/inference/csrc/relu.cu', 'csrc/transformer/inference/csrc/layer_norm.cu', + 'csrc/transformer/inference/csrc/rms_norm.cu', 'csrc/transformer/inference/csrc/softmax.cu', 'csrc/transformer/inference/csrc/dequantize.cu', 'csrc/transformer/inference/csrc/apply_rotary_pos_emb.cu', 'csrc/transformer/inference/csrc/transform.cu', + 'csrc/transformer/inference/csrc/pointwise_ops.cu', ] def extra_ldflags(self): diff --git a/tests/unit/hybrid_engine/test_hybrid_engine.py b/tests/unit/hybrid_engine/test_hybrid_engine.py new file mode 100644 index 000000000000..8ba8243e7171 --- /dev/null +++ b/tests/unit/hybrid_engine/test_hybrid_engine.py @@ -0,0 +1,69 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os +import torch +import pytest +import deepspeed +from deepspeed.ops.op_builder import OpBuilder +from unit.common import DistributedTest + +from transformers import (AutoConfig, AutoTokenizer, AutoModelForCausalLM) + +pytest.skip("skip test for now, will fix in follow-up PR", allow_module_level=True) + +rocm_version = OpBuilder.installed_rocm_version() +if rocm_version != (0, 0): + pytest.skip("skip inference tests on rocm for now", allow_module_level=True) + + +@pytest.mark.inference +@pytest.mark.parametrize("batch_size", [1, 2], ids=["bsz=1", "bsz=2"]) +@pytest.mark.parametrize("model_name", ["EleutherAI/gpt-neo-1.3B", "facebook/opt-1.3b"]) +class TestHybridEngineTextGen(DistributedTest): + world_size = 1 + + def _generate(self, model, tokenizer, prompt): + local_rank = int(os.getenv("LOCAL_RANK", "0")) + tokens = tokenizer.batch_encode_plus(prompt, return_tensors="pt", padding=True) + for t in tokens: + if torch.is_tensor(tokens[t]): + tokens[t] = tokens[t].to(f'cuda:{local_rank}') + output = model.generate(**tokens, do_sample=False, max_length=100) + outputs = tokenizer.batch_decode(output, skip_special_tokens=True) + return outputs + + def test(self, batch_size, model_name): + local_rank = int(os.getenv("LOCAL_RANK", "0")) + + model_config = AutoConfig.from_pretrained(model_name) + model_config.dropout = 0.0 + model = AutoModelForCausalLM.from_pretrained(model_name, config=model_config) + model = model.to(f'cuda:{local_rank}') + model = model.half() + + tokenizer = AutoTokenizer.from_pretrained(model_name) + tokenizer.pad_token = tokenizer.eos_token + + if batch_size == 1: + prompt = ["Microsoft is in Washington"] + elif batch_size == 2: + prompt = ["DeepSpeed is", "Microsoft is in Washington"] + else: + raise NotImplementedError(f"batch_size {batch_size} not implemented") + + base_out = self._generate(model, tokenizer, prompt) + + ds_config = {"train_batch_size": 1, "fp16": {"enabled": True}, "hybrid_engine": {"enabled": True}} + model, *_ = deepspeed.initialize(model=model, config=ds_config) + + model.eval() + ds1_out = self._generate(model, tokenizer, prompt) + assert base_out == ds1_out, f"base_out: {base_out}, ds1_out: {ds1_out}" + + model.train() + model.eval() + ds2_out = self._generate(model, tokenizer, prompt) + assert base_out == ds2_out diff --git a/tests/unit/ops/transformer/inference/__init__.py b/tests/unit/ops/transformer/inference/__init__.py new file mode 100644 index 000000000000..208299fb8c50 --- /dev/null +++ b/tests/unit/ops/transformer/inference/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team diff --git a/tests/unit/ops/transformer/inference/inference_test_utils.py b/tests/unit/ops/transformer/inference/inference_test_utils.py new file mode 100644 index 000000000000..beb5a6c63ead --- /dev/null +++ b/tests/unit/ops/transformer/inference/inference_test_utils.py @@ -0,0 +1,41 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from deepspeed.accelerator import get_accelerator + +TOLERANCES = None + + +def get_tolerances(): + global TOLERANCES + if TOLERANCES is None: + TOLERANCES = {torch.float32: (5e-4, 5e-5), torch.float16: (3e-2, 2e-3)} + if get_accelerator().is_bf16_supported(): + # Note: BF16 tolerance is higher than FP16 because of the lower precision (7 (+1) bits vs + # 10 (+1) bits) + TOLERANCES[torch.bfloat16] = (4.8e-1, 3.2e-2) + return TOLERANCES + + +DTYPES = None + + +def get_dtypes(): + global DTYPES + if DTYPES is None: + DTYPES = [torch.float16, torch.float32] + try: + if get_accelerator().is_bf16_supported(): + DTYPES.append(torch.bfloat16) + except (AssertionError, AttributeError): + pass + return DTYPES + + +def allclose(x, y): + assert x.dtype == y.dtype + rtol, atol = get_tolerances()[x.dtype] + return torch.allclose(x, y, rtol=rtol, atol=atol) diff --git a/tests/unit/ops/transformer/inference/test_bias_add.py b/tests/unit/ops/transformer/inference/test_bias_add.py index 36a01f2be8e7..843c9b889c2b 100644 --- a/tests/unit/ops/transformer/inference/test_bias_add.py +++ b/tests/unit/ops/transformer/inference/test_bias_add.py @@ -8,6 +8,7 @@ import deepspeed from deepspeed.accelerator import get_accelerator from deepspeed.ops.op_builder import InferenceBuilder +from .inference_test_utils import allclose, get_dtypes if not deepspeed.ops.__compatible_ops__[InferenceBuilder.NAME]: pytest.skip("Inference ops are not available on this system", allow_module_level=True) @@ -16,12 +17,6 @@ torch_minor_version = None -def allclose(x, y): - assert x.dtype == y.dtype - rtol, atol = {torch.float32: (5e-4, 5e-5), torch.float16: (3e-2, 2e-3)}[x.dtype] - return torch.allclose(x, y, rtol=rtol, atol=atol) - - def run_bias_add_reference(activations, bias): return activations + bias @@ -32,6 +27,8 @@ def run_bias_add_ds(activations, bias): inference_module = InferenceBuilder().load() if activations.dtype == torch.float16: return inference_module.bias_add_fp16(activations, bias) + elif activations.dtype == torch.bfloat16: + return inference_module.bias_add_bf16(activations, bias) else: return inference_module.bias_add_fp32(activations, bias) @@ -40,7 +37,7 @@ def run_bias_add_ds(activations, bias): @pytest.mark.parametrize("batch", [1, 2]) @pytest.mark.parametrize("sequence", [1, 128, 255]) @pytest.mark.parametrize("channels", [512, 1232, 4096]) -@pytest.mark.parametrize("dtype", [torch.float16, torch.float32], ids=["fp16", "fp32"]) +@pytest.mark.parametrize("dtype", get_dtypes()) def test_bias_add(batch, sequence, channels, dtype): activations_ds = torch.randn((batch, sequence, channels), dtype=dtype, device=get_accelerator().device_name()) bias_ds = torch.randn((channels), dtype=dtype, device=get_accelerator().device_name()) @@ -50,4 +47,6 @@ def test_bias_add(batch, sequence, channels, dtype): ds_out = run_bias_add_ds(activations_ds, bias_ds) ref_out = run_bias_add_reference(activations_ref, bias_ref) - assert allclose(ds_out, ref_out) + if not allclose(ds_out, ref_out): + print((ds_out - ref_out).abs().max()) + assert (allclose(ds_out, ref_out)) diff --git a/tests/unit/ops/transformer/inference/test_bias_geglu.py b/tests/unit/ops/transformer/inference/test_bias_geglu.py index 477c0a3bc7c7..d5ab13964974 100644 --- a/tests/unit/ops/transformer/inference/test_bias_geglu.py +++ b/tests/unit/ops/transformer/inference/test_bias_geglu.py @@ -8,6 +8,8 @@ import deepspeed from deepspeed.ops.op_builder import InferenceBuilder from deepspeed.accelerator import get_accelerator +from deepspeed.utils.types import ActivationFuncType +from .inference_test_utils import allclose, get_dtypes if not deepspeed.ops.__compatible_ops__[InferenceBuilder.NAME]: pytest.skip("Inference ops are not available on this system", allow_module_level=True) @@ -16,12 +18,6 @@ torch_minor_version = None -def allclose(x, y): - assert x.dtype == y.dtype - rtol, atol = {torch.float32: (5e-3, 5e-4), torch.float16: (3e-2, 2e-3), torch.int8: (0, 0)}[x.dtype] - return torch.allclose(x, y, rtol=rtol, atol=atol) - - def run_bias_geglu_reference(activations, bias): # Expected behavior is that of casting to float32 internally # Explicitly using the default GeLU @@ -34,14 +30,14 @@ def run_bias_geglu_ds(activation, bias): global inference_module if inference_module is None: inference_module = InferenceBuilder().load() - return inference_module.bias_geglu(activation, bias) + return inference_module.gated_activation(activation, bias, ActivationFuncType.GATED_GELU) @pytest.mark.inference_ops @pytest.mark.parametrize("batch", [1, 2]) @pytest.mark.parametrize("sequence", [1, 128, 255]) @pytest.mark.parametrize("channels", [512, 1232, 4096]) -@pytest.mark.parametrize("dtype", [torch.float16, torch.float32]) +@pytest.mark.parametrize("dtype", get_dtypes()) def test_bias_geglu(batch, sequence, channels, dtype): activation = torch.randn((batch, sequence, channels * 2), dtype=dtype, device=get_accelerator().device_name()) bias = torch.randn((channels * 2), dtype=dtype, device=get_accelerator().device_name()) @@ -49,3 +45,32 @@ def test_bias_geglu(batch, sequence, channels, dtype): ds_out = run_bias_geglu_ds(activation, bias) ref_out = run_bias_geglu_reference(activation, bias) assert (allclose(ds_out, ref_out)) + + +def run_gated_silu_reference(activations, bias): + # Expected behavior is that of casting to float32 internally + # Explicitly using the default GeLU + activations = activations + bias.reshape(1, 1, -1) + hidden_states, gate = activations.chunk(2, dim=-1) + return hidden_states * torch.nn.functional.silu(gate.to(torch.float32)).to(activations.dtype) + + +def run_gated_silu_ds(activation, bias): + global inference_module + if inference_module is None: + inference_module = InferenceBuilder().load() + return inference_module.gated_activation(activation, bias, ActivationFuncType.GATED_SILU) + + +@pytest.mark.inference_ops +@pytest.mark.parametrize("batch", [1, 2]) +@pytest.mark.parametrize("sequence", [1, 128, 255]) +@pytest.mark.parametrize("channels", [512, 1232, 4096]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.float32]) +def test_gated_silu(batch, sequence, channels, dtype): + activation = torch.randn((batch, sequence, channels * 2), dtype=dtype, device=get_accelerator().device_name()) + bias = torch.randn((channels * 2), dtype=dtype, device=get_accelerator().device_name()) + + ds_out = run_gated_silu_ds(activation, bias) + ref_out = run_gated_silu_reference(activation, bias) + assert (allclose(ds_out, ref_out)) diff --git a/tests/unit/ops/transformer/inference/test_bias_gelu.py b/tests/unit/ops/transformer/inference/test_bias_gelu.py index 1c5e7d58f85a..fd82da51380c 100644 --- a/tests/unit/ops/transformer/inference/test_bias_gelu.py +++ b/tests/unit/ops/transformer/inference/test_bias_gelu.py @@ -8,6 +8,7 @@ import deepspeed from deepspeed.accelerator import get_accelerator from deepspeed.ops.op_builder import InferenceBuilder +from .inference_test_utils import allclose, get_dtypes from packaging import version as pkg_version if not deepspeed.ops.__compatible_ops__[InferenceBuilder.NAME]: @@ -17,12 +18,6 @@ torch_minor_version = None -def allclose(x, y): - assert x.dtype == y.dtype - rtol, atol = {torch.float32: (5e-4, 5e-5), torch.float16: (3e-2, 2e-3)}[x.dtype] - return torch.allclose(x, y, rtol=rtol, atol=atol) - - def run_bias_gelu_reference(activations, bias): # Expected behavior is that of casting to float32 internally and using the tanh approximation return torch.nn.functional.gelu(activations.to(torch.float32) + bias.to(torch.float32), @@ -35,6 +30,8 @@ def run_bias_gelu_ds(activations, bias): inference_module = InferenceBuilder().load() if activations.dtype == torch.float16: return inference_module.bias_gelu_fp16(activations, bias) + elif activations.dtype == torch.bfloat16: + return inference_module.bias_gelu_bf16(activations, bias) else: return inference_module.bias_gelu_fp32(activations, bias) @@ -43,7 +40,7 @@ def run_bias_gelu_ds(activations, bias): @pytest.mark.parametrize("batch", [1, 2]) @pytest.mark.parametrize("sequence", [1, 128, 255]) @pytest.mark.parametrize("channels", [512, 1232, 4096]) -@pytest.mark.parametrize("dtype", [torch.float16, torch.float32]) +@pytest.mark.parametrize("dtype", get_dtypes()) def test_bias_gelu(batch, sequence, channels, dtype): if pkg_version.parse(torch.__version__) < pkg_version.parse("1.12"): pytest.skip("gelu implementation matches only after torch 1.12") diff --git a/tests/unit/ops/transformer/inference/test_bias_relu.py b/tests/unit/ops/transformer/inference/test_bias_relu.py index 50daa221f4cc..881af78e92cf 100644 --- a/tests/unit/ops/transformer/inference/test_bias_relu.py +++ b/tests/unit/ops/transformer/inference/test_bias_relu.py @@ -8,6 +8,7 @@ import deepspeed from deepspeed.accelerator import get_accelerator from deepspeed.ops.op_builder import InferenceBuilder +from .inference_test_utils import allclose, get_dtypes if not deepspeed.ops.__compatible_ops__[InferenceBuilder.NAME]: pytest.skip("Inference ops are not available on this system", allow_module_level=True) @@ -16,12 +17,6 @@ torch_minor_version = None -def allclose(x, y): - assert x.dtype == y.dtype - rtol, atol = {torch.float32: (5e-4, 5e-5), torch.float16: (3e-2, 2e-3)}[x.dtype] - return torch.allclose(x, y, rtol=rtol, atol=atol) - - def run_bias_relu_reference(activations, bias): # Expected behavior is that of casting to float32 internally return torch.nn.functional.relu(activations.to(torch.float32) + bias.to(torch.float32)).to(activations.dtype) @@ -33,6 +28,8 @@ def run_bias_relu_ds(activations, bias): inference_module = InferenceBuilder().load() if activations.dtype == torch.float16: return inference_module.bias_relu_fp16(activations, bias) + elif activations.dtype == torch.bfloat16: + return inference_module.bias_relu_bf16(activations, bias) else: return inference_module.bias_relu_fp32(activations, bias) @@ -41,7 +38,7 @@ def run_bias_relu_ds(activations, bias): @pytest.mark.parametrize("batch", [1, 2]) @pytest.mark.parametrize("sequence", [1, 128, 255]) @pytest.mark.parametrize("channels", [512, 1232, 4096]) -@pytest.mark.parametrize("dtype", [torch.float16, torch.float32]) +@pytest.mark.parametrize("dtype", get_dtypes()) def test_bias_relu(batch, sequence, channels, dtype): activations_ds = torch.randn((batch, sequence, channels), dtype=dtype, device=get_accelerator().device_name()) bias_ds = torch.randn((channels), dtype=dtype, device=get_accelerator().device_name()) diff --git a/tests/unit/ops/transformer/inference/test_layer_norm.py b/tests/unit/ops/transformer/inference/test_layer_norm.py index c765fd86744d..f44b977ac45d 100644 --- a/tests/unit/ops/transformer/inference/test_layer_norm.py +++ b/tests/unit/ops/transformer/inference/test_layer_norm.py @@ -8,6 +8,7 @@ import pytest from deepspeed.accelerator import get_accelerator from deepspeed.ops.op_builder import InferenceBuilder +from .inference_test_utils import allclose, get_dtypes if not deepspeed.ops.__compatible_ops__[InferenceBuilder.NAME]: pytest.skip("Inference ops are not available on this system", allow_module_level=True) @@ -15,17 +16,11 @@ inference_module = None -def allclose(x, y): - assert x.dtype == y.dtype - rtol, atol = {torch.float32: (5e-4, 5e-5), torch.float16: (3e-2, 2e-3)}[x.dtype] - return torch.allclose(x, y, rtol=rtol, atol=atol) - - -def ref_implementation(vals, gamma, beta, espilon, channels, dtype): +def ref_implementation(vals, gamma, beta, epsilon, channels, dtype): vals_f = vals.to(torch.float32) gamma_f = gamma.to(torch.float32) beta_f = beta.to(torch.float32) - return torch.nn.functional.layer_norm(vals_f, (channels, ), weight=gamma_f, bias=beta_f).to(dtype) + return torch.nn.functional.layer_norm(vals_f, (channels, ), weight=gamma_f, bias=beta_f, eps=epsilon).to(dtype) def ds_implementation(vals, gamma, beta, epsilon): @@ -39,7 +34,7 @@ def ds_implementation(vals, gamma, beta, epsilon): @pytest.mark.parametrize("batch", [1, 32]) @pytest.mark.parametrize("seq_len", [1, 128]) @pytest.mark.parametrize("channels", [384, 512, 768, 1024, 2048, 8192, 14432]) -@pytest.mark.parametrize("dtype", [torch.float16, torch.float32]) +@pytest.mark.parametrize("dtype", get_dtypes()) def test_layer_norm(batch, seq_len, channels, dtype): vals = torch.randn((batch, seq_len, channels), dtype=dtype, device=get_accelerator().current_device_name()) gamma = torch.randn((channels), dtype=dtype, device=get_accelerator().current_device_name()) @@ -49,10 +44,12 @@ def test_layer_norm(batch, seq_len, channels, dtype): ref_output = ref_implementation(vals, gamma, beta, epsilon, channels, dtype) new_output = ds_implementation(vals, gamma, beta, epsilon) - assert allclose(new_output, ref_output) + if not allclose(new_output, ref_output): + #print(new_output - ref_output) + assert allclose(new_output, ref_output) -def residual_ref_implementation(vals, bias, res, gamma, beta, espilon, channels, dtype): +def residual_ref_implementation(vals, bias, res, gamma, beta, epsilon, channels, dtype): vals_f = vals.to(torch.float32) bias_f = bias.to(torch.float32).reshape(1, 1, -1) res_f = res.to(torch.float32) @@ -72,7 +69,7 @@ def residual_ds_implementation(vals, bias, res, gamma, beta, epsilon): @pytest.mark.parametrize("batch", [1, 32]) @pytest.mark.parametrize("seq_len", [1, 128]) @pytest.mark.parametrize("channels", [384, 512, 768, 1024, 2048, 8192, 14432]) -@pytest.mark.parametrize("dtype", [torch.float16, torch.float32]) +@pytest.mark.parametrize("dtype", get_dtypes()) def test_layer_norm_residual(batch, seq_len, channels, dtype): vals = torch.randn((batch, seq_len, channels), dtype=dtype, device=get_accelerator().current_device_name()) residual = torch.randn((batch, seq_len, channels), dtype=dtype, device=get_accelerator().current_device_name()) @@ -84,6 +81,8 @@ def test_layer_norm_residual(batch, seq_len, channels, dtype): new_output = residual_ds_implementation(vals, bias, residual, gamma, beta, epsilon) ref_output = residual_ref_implementation(vals, bias, residual, gamma, beta, epsilon, channels, dtype) + print((new_output - ref_output).abs().max()) + assert allclose(new_output, ref_output) @@ -109,7 +108,7 @@ def residual_store_ds_implementation(vals, bias, res, gamma, beta, epsilon): @pytest.mark.parametrize("batch", [1, 32]) @pytest.mark.parametrize("seq_len", [1, 128]) @pytest.mark.parametrize("channels", [384, 512, 768, 1024, 2048, 8192, 14432]) -@pytest.mark.parametrize("dtype", [torch.float16, torch.float32]) +@pytest.mark.parametrize("dtype", get_dtypes()) def test_layer_norm_residual_store_pre_ln_res(batch, seq_len, channels, dtype): vals = torch.randn((batch, seq_len, channels), dtype=dtype, device=get_accelerator().current_device_name()) residual = torch.randn((batch, seq_len, channels), dtype=dtype, device=get_accelerator().current_device_name()) diff --git a/tests/unit/ops/transformer/inference/test_moe_res_matmult.py b/tests/unit/ops/transformer/inference/test_moe_res_matmult.py index 79313bd68bdb..e1c8127a83ac 100644 --- a/tests/unit/ops/transformer/inference/test_moe_res_matmult.py +++ b/tests/unit/ops/transformer/inference/test_moe_res_matmult.py @@ -8,6 +8,7 @@ import deepspeed from deepspeed.accelerator import get_accelerator from deepspeed.ops.op_builder import InferenceBuilder +from .inference_test_utils import allclose, get_dtypes if not deepspeed.ops.__compatible_ops__[InferenceBuilder.NAME]: pytest.skip("Inference ops are not available on this system", allow_module_level=True) @@ -15,12 +16,6 @@ inference_module = None -def allclose(x, y): - assert x.dtype == y.dtype - rtol, atol = {torch.float32: (5e-4, 5e-5), torch.float16: (3e-2, 2e-3)}[x.dtype] - return torch.allclose(x, y, rtol=rtol, atol=atol) - - def run_moe_res_matmul_reference(residual, coef1, coef2, output): return residual * coef1 + output * coef2 @@ -36,7 +31,7 @@ def run_moe_res_matmul_ds(residual, coef, output): @pytest.mark.inference_ops @pytest.mark.parametrize("hidden_dim", [16, 64]) @pytest.mark.parametrize("c", [1, 4]) -@pytest.mark.parametrize("dtype", [torch.float32, torch.float16]) +@pytest.mark.parametrize("dtype", get_dtypes()) def test_moe_residual_matmul(hidden_dim, c, dtype): residual_ds = torch.randn((c, hidden_dim * c, hidden_dim), dtype=dtype, device=get_accelerator().device_name()) coeff1 = torch.randn((1, 1, hidden_dim), dtype=dtype, device=get_accelerator().device_name()) diff --git a/tests/unit/ops/transformer/inference/test_residual_add.py b/tests/unit/ops/transformer/inference/test_residual_add.py index f5571d33b7bc..1a9d8975852c 100644 --- a/tests/unit/ops/transformer/inference/test_residual_add.py +++ b/tests/unit/ops/transformer/inference/test_residual_add.py @@ -8,14 +8,31 @@ import deepspeed from deepspeed.accelerator import get_accelerator from deepspeed.ops.op_builder import InferenceBuilder +from .inference_test_utils import get_dtypes if not deepspeed.ops.__compatible_ops__[InferenceBuilder.NAME]: pytest.skip("Inference ops are not available on this system", allow_module_level=True) +TOLERANCES = None + + +def get_tolerances(): + global TOLERANCES + if TOLERANCES is None: + # Residual add, as a sequence of casted additions, currently requires a higher tolerance + # than the other operators for FP16. We should instead better align the behaviors + # of the reference to match our kernel implementation (TODO(cmikeh2)) + TOLERANCES = {torch.float32: (5e-4, 5e-5), torch.float16: (3e-2, 4e-3)} + if get_accelerator().is_bf16_supported(): + # Note: BF16 tolerance is higher than FP16 because of the lower precision (7 (+1) bits vs + # 10 (+1) bits) + TOLERANCES[torch.bfloat16] = (4.8e-1, 3.2e-2) + return TOLERANCES + def allclose(x, y): assert x.dtype == y.dtype - rtol, atol = {torch.float32: (5e-4, 5e-5), torch.float16: (3e-2, 2e-2)}[x.dtype] + rtol, atol = get_tolerances()[x.dtype] return torch.allclose(x, y, rtol=rtol, atol=atol) @@ -52,7 +69,7 @@ def run_residual_add_reference(hidden_state, residual, attn_output, attn_bias, f @pytest.mark.parametrize("batch", [1, 2]) @pytest.mark.parametrize("sequence", [1, 128, 255]) @pytest.mark.parametrize("hidden_dim", [512, 1232, 4096]) -@pytest.mark.parametrize("dtype", [torch.float16, torch.float32]) +@pytest.mark.parametrize("dtype", get_dtypes()) @pytest.mark.parametrize("mlp_after_attn", [True, False]) @pytest.mark.parametrize("add_bias", [True, False]) @pytest.mark.parametrize("mp_size", [1, 2]) @@ -77,7 +94,15 @@ def test_residual_add(inference_module, batch, sequence, hidden_dim, dtype, mlp_ ds_out = inference_module.residual_add_bias_fp16(*res_add_args) elif dtype == torch.float32: ds_out = inference_module.residual_add_bias_fp32(*res_add_args) + elif dtype == torch.bfloat16: + ds_out = inference_module.residual_add_bias_bf16(*res_add_args) else: raise ValueError(f"Unsupported dtype: {dtype}") + if not allclose(ds_out, ref_out): + print((ds_out - ref_out).abs().max()) + print((ds_out - ref_out).abs().mean()) + print((ds_out - ref_out)) + assert (allclose(ds_out, ref_out)) + assert (allclose(ds_out, ref_out)) diff --git a/tests/unit/ops/transformer/inference/test_rms_norm.py b/tests/unit/ops/transformer/inference/test_rms_norm.py new file mode 100644 index 000000000000..1ae8f060673a --- /dev/null +++ b/tests/unit/ops/transformer/inference/test_rms_norm.py @@ -0,0 +1,89 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import deepspeed +import torch +import pytest +from deepspeed.accelerator import get_accelerator +from deepspeed.ops.op_builder import InferenceBuilder # type: ignore +from .inference_test_utils import allclose, get_dtypes + +if not deepspeed.ops.__compatible_ops__[InferenceBuilder.NAME]: + pytest.skip("Inference ops are not available on this system", allow_module_level=True) + +inference_module = None + + +def ref_implementation(vals, gamma, espilon): + variance = vals.to(torch.float32).pow(2).mean(-1, keepdim=True) + vals = vals * torch.rsqrt(variance + espilon) + + if gamma.dtype in [torch.float16, torch.bfloat16]: + vals = vals.to(gamma.dtype) + + return gamma * vals + + +def ds_implementation(vals, gamma, epsilon): + global inference_module + if inference_module is None: + inference_module = InferenceBuilder().load() + return inference_module.rms_norm(vals, gamma, epsilon) + + +@pytest.mark.inference_ops +@pytest.mark.parametrize("batch", [1, 32]) +@pytest.mark.parametrize("seq_len", [1, 128]) +@pytest.mark.parametrize("channels", [384, 512, 768, 1024, 2048, 8192, 14432]) +@pytest.mark.parametrize("dtype", get_dtypes()) +def test_rms_norm(batch, seq_len, channels, dtype): + device = get_accelerator().current_device_name() + vals = torch.randn((batch, seq_len, channels), dtype=dtype, device=device) + gamma = torch.randn((channels), dtype=dtype, device=device) + epsilon = 1e-5 + + ref_output = ref_implementation(vals, gamma, epsilon) + new_output = ds_implementation(vals, gamma, epsilon) + + assert allclose(new_output, ref_output) + + +def pre_ds_implementation(vals, residual, gamma, epsilon): + global inference_module + if inference_module is None: + inference_module = InferenceBuilder().load() + return inference_module.pre_rms_norm(vals, residual, gamma, epsilon) + + +def pre_ref_implementation(vals, residual, gamma, epsilon): + residual = vals.to(torch.float32) + residual.to(torch.float32) + vals = residual + + variance = vals.to(torch.float32).pow(2).mean(-1, keepdim=True) + vals = vals * torch.rsqrt(variance + epsilon) + + if gamma.dtype in [torch.float16, torch.bfloat16]: + vals = vals.to(gamma.dtype) + + return gamma * vals, residual.to(gamma.dtype) + + +@pytest.mark.inference_ops +@pytest.mark.parametrize("batch", [1, 32]) +@pytest.mark.parametrize("seq_len", [1, 128]) +@pytest.mark.parametrize("channels", [384, 512, 768, 1024, 2048, 8192, 14432]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.float32]) +def test_pre_norm(batch, seq_len, channels, dtype): + device = get_accelerator().current_device_name() + vals = torch.randn((batch, seq_len, channels), dtype=dtype, device=device) + residual = torch.randn((batch, seq_len, channels), dtype=dtype, device=device) + gamma = torch.randn((channels), dtype=dtype, device=device) + epsilon = 1e-5 + + ref_output = pre_ref_implementation(vals, residual, gamma, epsilon) + new_output = pre_ds_implementation(vals, residual, gamma, epsilon) + + assert allclose(new_output[0], ref_output[0]) + #assert allclose(new_output[1], ref_output[1])