github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-18 16:00:49 +00:00
Code Issues Packages Projects Releases Wiki Activity

Optimized MoE layer and fixed some bugs;

Browse Source 
Decreased moe tests;

Added FFNExperts and ViTMoE model
This commit is contained in:
1SAA
2022-02-18 20:42:31 +08:00
committed by Frank Lee
parent 3dba070580
commit 219df6e685
15 changed files with 1552 additions and 203 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
colossalai/global_variables.py
View File

@@ -56,6 +56,7 @@ class MoeEnv:
self.data_parallel_size = None
self.model_parallel_size = None
self.aux_loss = None
self.enable_cuda = True
def setup(self, moe_model_size):
from .core import global_context as gpc
@@ -71,6 +72,9 @@ class MoeEnv:
def is_initialized(self):
return self.model_parallel_size is not None
def set_cuda_false(self):
self.enable_cuda = False
def reset_loss(self):
self.aux_loss = 0

2
colossalai/kernel/cuda_native/csrc/layer_norm_cuda_kernel.cu
View File

@@ -5,7 +5,7 @@
#include "ATen/ATen.h"
#include "ATen/AccumulateType.h"
#include "ATen/cuda/CUDAContext.h"
#include <THC/THCDeviceUtils.cuh>
#include "ATen/cuda/DeviceUtils.cuh"
#include <cuda.h>
#include <cuda_runtime.h>

118
colossalai/kernel/cuda_native/csrc/moe_cuda.cpp Normal file
View File

@@ -0,0 +1,118 @@
#include <torch/extension.h>
torch::Tensor moe_dispatch_cuda_forward(
int s, int ec, int h,
torch::Tensor batch_tokens,
torch::Tensor mask,
torch::Tensor dest_idx);
torch::Tensor moe_dispatch_cuda_backward(
int s, int ec, int h,
torch::Tensor expert_grad,
torch::Tensor mask,
torch::Tensor dest_idx);
torch::Tensor moe_combine_cuda_forward(
int s, int e, int c, int h,
torch::Tensor expert_tokens,
torch::Tensor logits,
torch::Tensor mask,
torch::Tensor dest_idx);
std::vector<torch::Tensor> moe_combine_cuda_backward(
int s, int e, int c, int h,
torch::Tensor tokens_grad,
torch::Tensor expert_tokens,
torch::Tensor logits,
torch::Tensor mask,
torch::Tensor dest_idx);
torch::Tensor cumsum_sub_one_in_dim0(torch::Tensor mask);
#define CHECK_CUDA(x) TORCH_CHECK(x.device().is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
torch::Tensor moe_dispatch_forward(
int s, int ec, int h,
torch::Tensor batch_tokens,
torch::Tensor mask,
torch::Tensor dest_idx) {
CHECK_INPUT(batch_tokens);
CHECK_CUDA(mask);
CHECK_CUDA(dest_idx);
return moe_dispatch_cuda_forward(
s, ec, h,
batch_tokens, mask, dest_idx);
}
torch::Tensor moe_dispatch_backward(
int s, int ec, int h,
torch::Tensor expert_grad,
torch::Tensor mask,
torch::Tensor dest_idx) {
CHECK_INPUT(expert_grad);
CHECK_CUDA(mask);
CHECK_CUDA(dest_idx);
return moe_dispatch_cuda_backward(
s, ec, h,
expert_grad, mask, dest_idx);
}
torch::Tensor moe_combine_forward(
int s, int e, int c, int h,
torch::Tensor expert_tokens,
torch::Tensor logits,
torch::Tensor mask,
torch::Tensor dest_idx) {
CHECK_INPUT(expert_tokens);
CHECK_INPUT(logits);
CHECK_CUDA(mask);
CHECK_CUDA(dest_idx);
return moe_combine_cuda_forward(
s, e, c, h,
expert_tokens, logits, mask, dest_idx);
}
std::vector<torch::Tensor> moe_combine_backward(
int s, int e, int c, int h,
torch::Tensor tokens_grad,
torch::Tensor expert_tokens,
torch::Tensor logits,
torch::Tensor mask,
torch::Tensor dest_idx) {
CHECK_INPUT(tokens_grad);
CHECK_INPUT(logits);
CHECK_CUDA(mask);
CHECK_CUDA(dest_idx);
return moe_combine_cuda_backward(
s, e, c, h,
tokens_grad, expert_tokens, logits, mask, dest_idx);
}
torch::Tensor moe_cumsum(torch::Tensor mask) {
CHECK_INPUT(mask);
return cumsum_sub_one_in_dim0(mask);
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("cumsum_sub_one", &moe_cumsum,
"Fast cumsum operation in dim0");
m.def("dispatch_forward", &moe_dispatch_forward,
"Forward operation in MoE dispatch function");
m.def("dispatch_backward", &moe_dispatch_backward,
"Backward operation in MoE dispatch function");
m.def("combine_forward", &moe_combine_forward,
"Combine operation in MoE combine function");
m.def("combine_backward", &moe_combine_backward,
"Combine operation in MoE combine function");
}

702
colossalai/kernel/cuda_native/csrc/moe_cuda_kernel.cu Normal file
View File

@@ -0,0 +1,702 @@
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_fp16.h>
#include <cub/cub.cuh>
#include "block_reduce.h"
template<typename T, int block_size, int pack_size>
__device__ void moe_dpch_one_fwd(T *src_row, T *dst_row, const int cols) {
assert(cols % pack_size == 0);
const int bpack_size = block_size * pack_size;
typedef cub::BlockLoad<T, block_size, pack_size,
cub::BLOCK_LOAD_VECTORIZE> BlockLoad;
__shared__ typename BlockLoad::TempStorage ts_load;
typedef cub::BlockStore<T, block_size, pack_size,
cub::BLOCK_STORE_VECTORIZE> BlockStore;
__shared__ typename BlockStore::TempStorage ts_store;
int tps = threadIdx.x * pack_size; T pack[pack_size];
for (int idx = 0; idx + tps < cols; idx += bpack_size) {
BlockLoad(ts_load).Load(src_row + idx, pack);
BlockStore(ts_store).Store(dst_row + idx, pack);
}
}
template<typename T, int block_size, int pack_size>
__device__ void moe_dpch_one_bwd(T *src_row, T *dst_row, const int cols) {
assert(cols % pack_size == 0);
const int bpack_size = block_size * pack_size;
typedef cub::BlockLoad<T, block_size, pack_size,
cub::BLOCK_LOAD_VECTORIZE> BlockLoad;
__shared__ typename BlockLoad::TempStorage ts_load;
typedef cub::BlockStore<T, block_size, pack_size,
cub::BLOCK_STORE_VECTORIZE> BlockStore;
__shared__ typename BlockStore::TempStorage ts_store;
int tps = threadIdx.x * pack_size; T pack[pack_size];
for (int idx = 0; idx + tps < cols; idx += bpack_size) {
BlockLoad(ts_load).Load(dst_row + idx, pack);
BlockStore(ts_store).Store(src_row + idx, pack);
}
}
template<typename T, int block_size, int pack_size>
__device__ void moe_dpch_two_fwd(T *src_row, T *dst_row1, T *dst_row2, const int cols) {
assert(cols % pack_size == 0);
const int bpack_size = block_size * pack_size;
typedef cub::BlockLoad<T, block_size, pack_size,
cub::BLOCK_LOAD_VECTORIZE> BlockLoad;
__shared__ typename BlockLoad::TempStorage ts_load;
typedef cub::BlockStore<T, block_size, pack_size,
cub::BLOCK_STORE_VECTORIZE> BlockStore;
__shared__ typename BlockStore::TempStorage ts_store;
int tps = threadIdx.x * pack_size; T pack[pack_size];
for (int idx = 0; idx + tps < cols; idx += bpack_size) {
BlockLoad(ts_load).Load(src_row + idx, pack);
BlockStore(ts_store).Store(dst_row1 + idx, pack);
BlockStore(ts_store).Store(dst_row2 + idx, pack);
}
}
template<typename T, int block_size, int pack_size>
__device__ void moe_dpch_two_bwd(T *src_row, T *dst_row1, T *dst_row2, const int cols) {
assert(cols % pack_size == 0);
const int bpack_size = block_size * pack_size;
typedef cub::BlockLoad<T, block_size, pack_size,
cub::BLOCK_LOAD_VECTORIZE> BlockLoad;
__shared__ typename BlockLoad::TempStorage ts_load;
typedef cub::BlockStore<T, block_size, pack_size,
cub::BLOCK_STORE_VECTORIZE> BlockStore;
__shared__ typename BlockStore::TempStorage ts_store;
int tps = threadIdx.x * pack_size;
T pack1[pack_size], pack2[pack_size];
for (int idx = 0; idx + tps < cols; idx += bpack_size) {
BlockLoad(ts_load).Load(dst_row1 + idx, pack1);
BlockLoad(ts_load).Load(dst_row2 + idx, pack2);
#pragma unroll
for (int i = 0; i < pack_size; ++i) {
pack1[i] += pack2[i];
}
BlockStore(ts_store).Store(src_row + idx, pack1);
}
}
template<typename T, int block_size, int pack_size>
__device__ void moe_cb_one_fwd(
T *src_row, T *dst_row,
const T weight, const int cols) {
assert(cols % pack_size == 0);
const int bpack_size = block_size * pack_size;
typedef cub::BlockLoad<T, block_size, pack_size,
cub::BLOCK_LOAD_VECTORIZE> BlockLoad;
__shared__ typename BlockLoad::TempStorage ts_load;
typedef cub::BlockStore<T, block_size, pack_size,
cub::BLOCK_STORE_VECTORIZE> BlockStore;
__shared__ typename BlockStore::TempStorage ts_store;
int tps = threadIdx.x * pack_size; T pack[pack_size];
for (int idx = 0; idx + tps < cols; idx += bpack_size) {
BlockLoad(ts_load).Load(src_row + idx, pack);
#pragma unroll
for (int i = 0; i < pack_size; ++i) {
pack[i] *= weight;
}
BlockStore(ts_store).Store(dst_row + idx, pack);
}
}
template<typename T, int block_size, int pack_size>
__device__ void moe_cb_one_bwd(
T *src_row, T *dst_row, T *tks_row, T *weight_grad,
const T weight, const int cols) {
assert(cols % pack_size == 0);
const int bpack_size = block_size * pack_size;
typedef cub::BlockLoad<T, block_size, pack_size,
cub::BLOCK_LOAD_VECTORIZE> BlockLoad;
__shared__ typename BlockLoad::TempStorage ts_load;
typedef cub::BlockStore<T, block_size, pack_size,
cub::BLOCK_STORE_VECTORIZE> BlockStore;
__shared__ typename BlockStore::TempStorage ts_store;
int tps = threadIdx.x * pack_size;
T grad[pack_size], tokens[pack_size];
float thread_sum = 0;
for (int idx = 0; idx + tps < cols; idx += bpack_size) {
BlockLoad(ts_load).Load(dst_row + idx, grad);
BlockLoad(ts_load).Load(tks_row + idx, tokens);
#pragma unroll
for (int i = 0; i < pack_size; ++i) {
thread_sum += grad[i] * tokens[i];
grad[i] *= weight;
}
BlockStore(ts_store).Store(src_row + idx, grad);
}
blockReduce<ReduceType::kSum, 1>(&thread_sum);
if (threadIdx.x == 0)
*weight_grad = static_cast<T>(thread_sum);
}
template<typename T, int block_size, int pack_size>
__device__ void moe_cb_two_fwd(
T *src_row1, T *src_row2, T *dst_row,
const T weight1, const T weight2, const int cols) {
assert(cols % pack_size == 0);
const int bpack_size = block_size * pack_size;
typedef cub::BlockLoad<T, block_size, pack_size,
cub::BLOCK_LOAD_VECTORIZE> BlockLoad;
__shared__ typename BlockLoad::TempStorage ts_load;
typedef cub::BlockStore<T, block_size, pack_size,
cub::BLOCK_STORE_VECTORIZE> BlockStore;
__shared__ typename BlockStore::TempStorage ts_store;
int tps = threadIdx.x * pack_size;
T pack1[pack_size], pack2[pack_size];
for (int idx = 0; idx + tps < cols; idx += bpack_size) {
BlockLoad(ts_load).Load(src_row1 + idx, pack1);
BlockLoad(ts_load).Load(src_row2 + idx, pack2);
#pragma unroll
for (int i = 0; i < pack_size; ++i) {
pack1[i] = pack1[i] * weight1 + pack2[i] * weight2;
}
BlockStore(ts_store).Store(dst_row + idx, pack1);
}
}
template<typename T, int block_size, int pack_size>
__device__ void moe_cb_two_bwd(
T *src_row1, T *src_row2, T *dst_row,
T *tks_row1, T *tks_row2, T *weight_grad1, T *weight_grad2,
const T weight1, const T weight2, const int cols) {
assert(cols % pack_size == 0);
const int bpack_size = block_size * pack_size;
typedef cub::BlockLoad<T, block_size, pack_size,
cub::BLOCK_LOAD_VECTORIZE> BlockLoad;
__shared__ typename BlockLoad::TempStorage ts_load;
typedef cub::BlockStore<T, block_size, pack_size,
cub::BLOCK_STORE_VECTORIZE> BlockStore;
__shared__ typename BlockStore::TempStorage ts_store;
int tps = threadIdx.x * pack_size;
T grad[pack_size], tokens1[pack_size], tokens2[pack_size],
sgrad1[pack_size], sgrad2[pack_size];
float thread_sum[2] = {0, 0};
for (int idx = 0; idx + tps < cols; idx += bpack_size) {
BlockLoad(ts_load).Load(dst_row + idx, grad);
BlockLoad(ts_load).Load(tks_row1 + idx, tokens1);
BlockLoad(ts_load).Load(tks_row2 + idx, tokens2);
#pragma unroll
for (int i = 0; i < pack_size; ++i) {
thread_sum[0] += grad[i] * tokens1[i];
thread_sum[1] += grad[i] * tokens2[i];
sgrad1[i] = weight1 * grad[i];
sgrad2[i] = weight2 * grad[i];
}
BlockStore(ts_store).Store(src_row1 + idx, sgrad1);
BlockStore(ts_store).Store(src_row2 + idx, sgrad2);
}
blockReduce<ReduceType::kSum, 2>(thread_sum);
if (threadIdx.x == 0)
*weight_grad1 = static_cast<T>(thread_sum[0]);
else if (threadIdx.x == 1)
*weight_grad2 = static_cast<T>(thread_sum[1]);
}
// DISPATCH KERNELS --------------------------------
template<typename T, int block_size, int pack_size>
__device__ void moe_dpch_fwd_selector(
T *src_row, T *dst_row1, T *dst_row2, const int cols,
const int indicator1, const int indicator2) {
if (indicator1 != 0 && indicator2 != 0)
moe_dpch_two_fwd<T, block_size, pack_size>(
src_row, dst_row1, dst_row2, cols);
else if (indicator1 != 0)
moe_dpch_one_fwd<T, block_size, pack_size>(
src_row, dst_row1, cols);
else if (indicator2 != 0)
moe_dpch_one_fwd<T, block_size, pack_size>(
src_row, dst_row2, cols);
else
return;
}
template<typename T, int block_size, int pack_size>
__device__ void moe_dpch_bwd_selector(
T *src_row, T *dst_row1, T *dst_row2, const int cols,
const int indicator1, const int indicator2) {
if (indicator1 != 0 && indicator2 != 0)
moe_dpch_two_bwd<T, block_size, pack_size>(
src_row, dst_row1, dst_row2, cols);
else if (indicator1 != 0)
moe_dpch_one_bwd<T, block_size, pack_size>(
src_row, dst_row1, cols);
else if (indicator2 != 0)
moe_dpch_one_bwd<T, block_size, pack_size>(
src_row, dst_row2, cols);
else
return;
}
template<typename T, int block_size, int pack_size>
__global__ void moe_dpch_fwd_kernel(
T *batch_tokens, T *expert_input,
int *mask1, int *mask2,
int *dest1, int *dest2, const int h) {
int row = blockIdx.x;
int indicator2 = mask2 == nullptr ? 0 : mask2[row];
moe_dpch_fwd_selector<T, block_size, pack_size>(
batch_tokens + (row * h),
expert_input + (dest1[row] * h), expert_input + (dest2[row] * h),
h, mask1[row], indicator2);
}
template<typename T, int block_size, int pack_size>
__global__ void moe_dpch_bwd_kernel(
T *tokens_grad, T *expert_grad,
int *mask1, int *mask2,
int *dest1, int *dest2, const int h) {
int row = blockIdx.x;
int indicator2 = mask2 == nullptr ? 0 : mask2[row];
moe_dpch_bwd_selector<T, block_size, pack_size>(
tokens_grad + (row * h),
expert_grad + (dest1[row] * h), expert_grad + (dest2[row] * h),
h, mask1[row], indicator2);
}
// COMBINE KERNELS --------------------------------
template<typename T, int block_size, int pack_size>
__device__ void moe_cb_fwd_selector(
T *src_row1, T *src_row2, T *dst_row, const int cols,
const T weight1, const T weight2,
const int indicator1, const int indicator2) {
if (indicator1 != 0 && indicator2 != 0)
moe_cb_two_fwd<T, block_size, pack_size>(
src_row1, src_row2, dst_row, weight1, weight2, cols);
else if (indicator1 != 0)
moe_cb_one_fwd<T, block_size, pack_size>(
src_row1, dst_row, weight1, cols);
else if (indicator2 != 0)
moe_cb_one_fwd<T, block_size, pack_size>(
src_row2, dst_row, weight2, cols);
else
return;
}
template<typename T, int block_size, int pack_size>
__device__ void moe_cb_bwd_selector(
T *src_row1, T *src_row2, T *dst_row, const int cols,
T *tks_row1, T *tks_row2, T *wt_grad1, T *wt_grad2,
const T weight1, const T weight2,
const int indicator1, const int indicator2) {
if (indicator1 != 0 && indicator2 != 0)
moe_cb_two_bwd<T, block_size, pack_size>(
src_row1, src_row2, dst_row,
tks_row1, tks_row2, wt_grad1, wt_grad2,
weight1, weight2, cols);
else if (indicator1 != 0)
moe_cb_one_bwd<T, block_size, pack_size>(
src_row1, dst_row, tks_row1, wt_grad1, weight1, cols);
else if (indicator2 != 0)
moe_cb_one_bwd<T, block_size, pack_size>(
src_row2, dst_row, tks_row2, wt_grad2, weight2, cols);
else
return;
}
template<typename T, int block_size, int pack_size>
__global__ void moe_cb_fwd_kernel(
T *expert_tokens, T *combine_tokens, T *logits,
int *mask1, int *mask2,
int *dest1, int *dest2,
const int e, const int c, const int h) {
int row = blockIdx.x, eid1 = dest1[row] / c, eid2 = dest2[row] / c;
int indicator2 = mask2 == nullptr ? 0 : mask2[row];
T *row_log = logits + (row * e);
moe_cb_fwd_selector<T, block_size, pack_size>(
expert_tokens + (dest1[row] * h), expert_tokens + (dest2[row] * h),
combine_tokens + (row * h), h,
row_log[eid1], row_log[eid2],
mask1[row], indicator2);
}
template<typename T, int block_size, int pack_size>
__global__ void moe_cb_bwd_kernel(
T *tokens_grad, T *expert_grad, T *tks,
T *logits, T *logits_grad,
int *mask1, int *mask2,
int *dest1, int *dest2,
const int e, const int c, const int h) {
int row = blockIdx.x, eid1 = dest1[row] / c, eid2 = dest2[row] / c;
int indicator2 = mask2 == nullptr ? 0 : mask2[row];
T *row_log = logits + (row * e), *row_grad = logits_grad + (row * e);
moe_cb_bwd_selector<T, block_size, pack_size>(
expert_grad + (dest1[row] * h), expert_grad + (dest2[row] * h),
tokens_grad + (row * h), h,
tks + (dest1[row] * h), tks + (dest2[row] * h),
row_grad + eid1, row_grad + eid2,
row_log[eid1], row_log[eid2],
mask1[row], indicator2);
}
//CUMSUM KERNEL --------------------------------
template<int block_size, int pack_size>
__global__ void cumsum_kernel(
int *inputs, int *outputs,
const int s, const int e) {
assert(s % pack_size == 0);
constexpr int bpack_size = block_size * pack_size;
int tid = threadIdx.x, bid = blockIdx.x,
tps = tid * pack_size, last_sum = -1;
__shared__ int temp[block_size + 1]; int pack[pack_size];
for (int idx = 0; idx < s; idx += bpack_size) {
int offset = 1;
if (idx + tps < s) {
temp[tid] = inputs[tps * e + bid];
#pragma unroll
for (int i = 1; i < pack_size; ++i) {
pack[i] = inputs[(tps + i) * e + bid];
}
#pragma unroll
for (int i = 1; i < pack_size; ++i) {
temp[tid] += pack[i];
}
}
for (int i = block_size >> 1; i > 0; i >>= 1) {
__syncthreads();
if (tid < i) {
int j = offset * (2 * tid + 1) - 1;
temp[j + offset] += temp[j];
}
offset <<= 1;
}
if (tid == 0) {
temp[block_size] = temp[block_size - 1];
temp[block_size - 1] = 0;
}
for (int i = 1; i < block_size; i <<= 1) {
offset >>= 1;
__syncthreads();
if (tid < i) {
int j = offset * (2 * tid + 1) - 1,
k = j + offset, ts = temp[j];
temp[j] = temp[k];
temp[k] += ts;
}
}
__syncthreads();
if (tid == 0)
temp[0] = temp[block_size];
__syncthreads();
if (idx + tps < s) {
temp[tid + 1] += last_sum;
#pragma unroll
for (int i = pack_size - 1; i > 0; --i) {
outputs[(tps + i) * e + bid] = temp[tid + 1];
temp[tid + 1] -= pack[i];
}
outputs[tps * e + bid] = temp[tid + 1];
}
__syncthreads();
last_sum += temp[0];
inputs += bpack_size * e;
outputs += bpack_size * e;
}
}
//LAUNCH FUNCTIONS --------------------------------
template<typename T>
void moe_dpch_fwd_launch(
T *batch_tokens, T *expert_input,
int *mask1, int *mask2,
int *dest1, int *dest2,
const int s, const int h) {
if (h < 256)
moe_dpch_fwd_kernel<T, 32, 4><<<s, 32>>>(batch_tokens, expert_input, mask1, mask2, dest1, dest2, h);
else if (h < 512)
moe_dpch_fwd_kernel<T, 32, 8><<<s, 32>>>(batch_tokens, expert_input, mask1, mask2, dest1, dest2, h);
else if (h < 1024)
moe_dpch_fwd_kernel<T, 32, 16><<<s, 32>>>(batch_tokens, expert_input, mask1, mask2, dest1, dest2, h);
else if (h < 2048)
moe_dpch_fwd_kernel<T, 64, 16><<<s, 64>>>(batch_tokens, expert_input, mask1, mask2, dest1, dest2, h);
else
moe_dpch_fwd_kernel<T, 128, 16><<<s, 128>>>(batch_tokens, expert_input, mask1, mask2, dest1, dest2, h);
}
template<typename T>
void moe_dpch_bwd_launch(
T *tokens_grad, T *expert_grad,
int *mask1, int *mask2,
int *dest1, int *dest2,
const int s, const int h) {
if (h < 256)
moe_dpch_bwd_kernel<T, 32, 4><<<s, 32>>>(tokens_grad, expert_grad, mask1, mask2, dest1, dest2, h);
else if (h < 512)
moe_dpch_bwd_kernel<T, 32, 8><<<s, 32>>>(tokens_grad, expert_grad, mask1, mask2, dest1, dest2, h);
else if (h < 1024)
moe_dpch_bwd_kernel<T, 32, 16><<<s, 32>>>(tokens_grad, expert_grad, mask1, mask2, dest1, dest2, h);
else if (h < 2048)
moe_dpch_bwd_kernel<T, 64, 16><<<s, 64>>>(tokens_grad, expert_grad, mask1, mask2, dest1, dest2, h);
else
moe_dpch_bwd_kernel<T, 128, 16><<<s, 128>>>(tokens_grad, expert_grad, mask1, mask2, dest1, dest2, h);
}
template<typename T>
void moe_cb_fwd_launch(
T *expert_tokens, T *combine_tokens, T *logits,
int *mask1, int *mask2,
int *dest1, int *dest2,
const int s, const int e, const int c, const int h) {
if (h < 256)
moe_cb_fwd_kernel<T, 32, 4><<<s, 32>>>
(expert_tokens, combine_tokens, logits, mask1, mask2, dest1, dest2, e, c, h);
else if (h < 512)
moe_cb_fwd_kernel<T, 32, 8><<<s, 32>>>
(expert_tokens, combine_tokens, logits, mask1, mask2, dest1, dest2, e, c, h);
else if (h < 1024)
moe_cb_fwd_kernel<T, 32, 16><<<s, 32>>>
(expert_tokens, combine_tokens, logits, mask1, mask2, dest1, dest2, e, c, h);
else if (h < 2048)
moe_cb_fwd_kernel<T, 64, 16><<<s, 64>>>
(expert_tokens, combine_tokens, logits, mask1, mask2, dest1, dest2, e, c, h);
else
moe_cb_fwd_kernel<T, 128, 16><<<s, 128>>>
(expert_tokens, combine_tokens, logits, mask1, mask2, dest1, dest2, e, c, h);
}
template<typename T>
void moe_cb_bwd_launch(
T *tokens_grad, T *expert_grad, T *tks,
T *logits, T *logits_grad,
int *mask1, int *mask2,
int *dest1, int *dest2,
const int s, const int e, const int c, const int h) {
if (h < 256)
moe_cb_bwd_kernel<T, 32, 4><<<s, 32>>>
(tokens_grad, expert_grad, tks, logits, logits_grad, mask1, mask2, dest1, dest2, e, c, h);
else // if (h < 512)
moe_cb_bwd_kernel<T, 64, 4><<<s, 64>>>
(tokens_grad, expert_grad, tks, logits, logits_grad, mask1, mask2, dest1, dest2, e, c, h);
// else if (h < 1024)
// moe_cb_bwd_kernel<T, 128, 4><<<s, 128>>>
// (tokens_grad, expert_grad, tks, logits, logits_grad, mask1, mask2, dest1, dest2, e, c, h);
// else
// moe_cb_bwd_kernel<T, 256, 4><<<s, 256>>>
// (tokens_grad, expert_grad, tks, logits, logits_grad, mask1, mask2, dest1, dest2, e, c, h);
}
void cumsum_launch(
int *inputs, int *outputs,
const int s, const int e) {
if (s <= 256)
cumsum_kernel<256, 1><<<e, 256>>>(inputs, outputs, s, e);
else if (s <= 512)
cumsum_kernel<512, 1><<<e, 512>>>(inputs, outputs, s, e);
else if (s <= 1024)
cumsum_kernel<1024, 1><<<e, 1024>>>(inputs, outputs, s, e);
else if (s <= 2048)
cumsum_kernel<1024, 2><<<e, 1024>>>(inputs, outputs, s, e);
else
cumsum_kernel<1024, 4><<<e, 1024>>>(inputs, outputs, s, e);
}
// API FUNCTIONS --------------------------------
#define DISPATCH_FLOAT_AND_HALF(TYPE, NAME, ...) \
switch (TYPE) \
{ \
case at::ScalarType::Float: \
{ \
using scalar_t = float; \
__VA_ARGS__; \
break; \
} \
case at::ScalarType::Half: \
{ \
using scalar_t = at::Half; \
__VA_ARGS__; \
break; \
} \
default: \
AT_ERROR(#NAME, " not implemented yet for specific data type.");\
}
torch::Tensor moe_dispatch_cuda_forward(
int s, int ec, int h,
torch::Tensor batch_tokens,
torch::Tensor mask,
torch::Tensor dest_idx) {
assert(h % 16 == 0);
auto res = torch::zeros({ec, h},
torch::dtype(batch_tokens.dtype()).device(batch_tokens.device()));
auto k = mask.size(0);
DISPATCH_FLOAT_AND_HALF(
batch_tokens.scalar_type(), "moe dispatch forward",
moe_dpch_fwd_launch<scalar_t>(
batch_tokens.data<scalar_t>(), res.data<scalar_t>(),
mask[0].data<int>(), k == 1 ? nullptr : mask[1].data<int>(),
dest_idx[0].data<int>(), k == 1 ? dest_idx[0].data<int>() : dest_idx[1].data<int>(),
s, h)
);
return res;
}
torch::Tensor moe_dispatch_cuda_backward(
int s, int ec, int h,
torch::Tensor expert_grad,
torch::Tensor mask,
torch::Tensor dest_idx) {
assert(h % 16 == 0);
auto res = torch::zeros({s, h},
torch::dtype(expert_grad.dtype()).device(expert_grad.device()));
auto k = mask.size(0);
DISPATCH_FLOAT_AND_HALF(
expert_grad.scalar_type(), "moe dispatch backward",
moe_dpch_bwd_launch<scalar_t>(
res.data<scalar_t>(), expert_grad.data<scalar_t>(),
mask[0].data<int>(), k == 1 ? nullptr : mask[1].data<int>(),
dest_idx[0].data<int>(), k == 1 ? dest_idx[0].data<int>() : dest_idx[1].data<int>(),
s, h)
);
return res;
}
torch::Tensor moe_combine_cuda_forward(
int s, int e, int c, int h,
torch::Tensor expert_tokens,
torch::Tensor logits,
torch::Tensor mask,
torch::Tensor dest_idx) {
assert(h % 16 == 0);
assert(expert_tokens.dtype() == logits.dtype());
auto res = torch::zeros({s, h},
torch::dtype(expert_tokens.dtype()).device(expert_tokens.device()));
auto k = mask.size(0);
DISPATCH_FLOAT_AND_HALF(
expert_tokens.scalar_type(), "moe combine forward",
moe_cb_fwd_launch<scalar_t>(
expert_tokens.data<scalar_t>(), res.data<scalar_t>(), logits.data<scalar_t>(),
mask[0].data<int>(), k == 1 ? nullptr : mask[1].data<int>(),
dest_idx[0].data<int>(), k == 1 ? dest_idx[0].data<int>() : dest_idx[1].data<int>(),
s, e, c, h)
);
return res;
}
std::vector<torch::Tensor> moe_combine_cuda_backward(
int s, int e, int c, int h,
torch::Tensor tokens_grad,
torch::Tensor expert_tokens,
torch::Tensor logits,
torch::Tensor mask,
torch::Tensor dest_idx) {
assert(h % 16 == 0);
assert(tokens_grad.dtype() == expert_tokens.dtype());
assert(expert_tokens.dtype() == logits.dtype());
auto egrad = torch::zeros({e * c, h},
torch::dtype(tokens_grad.dtype()).device(tokens_grad.device())),
wgrad = torch::zeros({s, e}, torch::dtype(logits.dtype()).device(logits.device()));
auto k = mask.size(0);
DISPATCH_FLOAT_AND_HALF(
tokens_grad.scalar_type(), "moe combine backward",
moe_cb_bwd_launch<scalar_t>(
tokens_grad.data<scalar_t>(), egrad.data<scalar_t>(), expert_tokens.data<scalar_t>(),
logits.data<scalar_t>(), wgrad.data<scalar_t>(),
mask[0].data<int>(), k == 1 ? nullptr : mask[1].data<int>(),
dest_idx[0].data<int>(), k == 1 ? dest_idx[0].data<int>() : dest_idx[1].data<int>(),
s, e, c, h)
);
return {egrad, wgrad};
}
torch::Tensor cumsum_sub_one_in_dim0(torch::Tensor mask) {
assert(mask.dim() == 2);
assert(mask.dtype() == torch::kInt32);
const int s = mask.size(0), e = mask.size(1);
auto res = torch::empty({s, e}, torch::dtype(torch::kInt32).device(mask.device()));
cumsum_launch(mask.data<int>(), res.data<int>(), s, e);
return res;
}

11
colossalai/nn/layer/moe/__init__.py
View File

@@ -1,8 +1,5 @@
from ._operation import AllToAll
from .layers import Experts, MoeLayer, \
NormalNoiseGenerator, Top1Router, Top2Router
from .experts import Experts, FFNExperts
from .layers import MoeLayer, Top1Router, Top2Router
from .utils import NormalNoiseGenerator
__all__ = [
'AllToAll', 'Experts', 'Top1Router', 'Top2Router',
'MoeLayer', 'NormalNoiseGenerator'
]
__all__ = ['Experts', 'FFNExperts', 'Top1Router', 'Top2Router', 'MoeLayer', 'NormalNoiseGenerator']

89
colossalai/nn/layer/moe/_operation.py
View File

@@ -6,16 +6,26 @@ from colossalai.context import ParallelMode
from colossalai.core import global_context as gpc
from typing import Any, Tuple
U_CUDA_MODE = False
try:
import colossal_moe_cuda
U_CUDA_MODE = True
except ImportError:
print("If you want to activate cuda mode for MoE, please install with cuda_ext!")
class AllToAll(torch.autograd.Function):
"""Dispatches input tensor [e, c, h] to all experts by all_to_all_single
operation in torch.distributed.
"""
@staticmethod
def forward(ctx: Any,
inputs: Tensor,
parallel_mode: ParallelMode) -> Tensor:
ctx.parallel_mode = parallel_mode
if ctx is not None:
ctx.parallel_mode = parallel_mode
if not inputs.is_contiguous():
inputs = inputs.contiguous()
@@ -26,4 +36,79 @@ class AllToAll(torch.autograd.Function):
@staticmethod
def backward(ctx: Any, *grad_outputs: Tensor) -> Tuple[Tensor, None]:
return AllToAll.apply(*grad_outputs, ctx.parallel_mode), None
return AllToAll.forward(None, *grad_outputs, ctx.parallel_mode), None
class MoeDispatch(torch.autograd.Function):
@staticmethod
def forward(ctx, tokens, mask, dest_idx, ec):
s = tokens.size(0)
h = tokens.size(1)
expert_input = colossal_moe_cuda.dispatch_forward(s, ec, h, tokens, mask, dest_idx)
ctx.save_for_backward(mask, dest_idx)
ctx.s = s
ctx.h = h
ctx.ec = ec
return expert_input
@staticmethod
def backward(ctx, output_grad):
mask, dest_idx = ctx.saved_tensors
d_tokens = colossal_moe_cuda.dispatch_backward(
ctx.s, ctx.ec, ctx.h, output_grad, mask, dest_idx)
return d_tokens, None, None, None
class MoeCombine(torch.autograd.Function):
@staticmethod
def forward(ctx, expert_tokens, logits, mask, dest_idx, ec):
assert logits.dtype == torch.float32
s = logits.size(0)
e = logits.size(1)
c = ec // e
h = expert_tokens.size(-1)
fp16_flag = (expert_tokens.dtype == torch.float16)
cb_input = expert_tokens.to(torch.float32) if fp16_flag else expert_tokens
ctokens = colossal_moe_cuda.combine_forward(s, e, c, h,
cb_input, logits,
mask, dest_idx)
output = ctokens.to(torch.float16) if fp16_flag else ctokens
ctx.save_for_backward(expert_tokens, logits, mask, dest_idx)
ctx.s = s
ctx.e = e
ctx.c = c
ctx.h = h
ctx.fp16_flag = fp16_flag
return output
@staticmethod
def backward(ctx, tokens_grad):
expert_tokens, logits, mask, dest_idx = ctx.saved_tensors
cb_grad = tokens_grad.to(torch.float32) if tokens_grad.dtype is torch.float16 \
else tokens_grad
cb_input = expert_tokens.to(torch.float32) if ctx.fp16_flag else expert_tokens
d_expert, d_logits = colossal_moe_cuda.combine_backward(
ctx.s, ctx.e, ctx.c, ctx.h,
cb_grad, cb_input, logits, mask, dest_idx)
d_expert = d_expert.to(torch.float16) if ctx.fp16_flag else d_expert
return d_expert, d_logits, None, None, None
def moe_cumsum(inputs: Tensor):
dim0 = inputs.size(0)
flag = (dim0 <= 1024) or (dim0 <= 2048 and dim0 % 2 == 0) or (dim0 % 4 == 0)
if flag and U_CUDA_MODE:
return colossal_moe_cuda.cumsum_sub_one(inputs)
else:
return torch.cumsum(inputs, dim=0) - 1

96
colossalai/nn/layer/moe/experts.py Normal file
View File

@@ -0,0 +1,96 @@
import math
import torch
import torch.nn as nn
from colossalai.global_variables import moe_env
from colossalai.context import ParallelMode, seed
from colossalai.utils import get_current_device
class Experts(nn.Module):
"""A wrapper class to create experts. It will create E experts across the
moe model parallel group, where E is the number of experts. Every expert
is a instence of the class, 'expert' in initialization parameters.
:param expert: The class of all experts
:param num_experts: The number of experts
:param expert_args: Args used to initialize experts
:type num_experts: int
"""
def __init__(self, expert, num_experts, **expert_args):
super().__init__()
assert num_experts % moe_env.model_parallel_size == 0, \
"The number of experts should be divied by moe model size"
num_local_experts = num_experts // moe_env.model_parallel_size
with seed(ParallelMode.MOE_MODEL):
self.experts = nn.ModuleList([expert(**expert_args) for _ in range(num_local_experts)])
self.num_local_experts = num_local_experts
for exp in self.experts:
for param in exp.parameters():
param.__setattr__('moe_param', True)
def forward(self, inputs):
expert_input = torch.chunk(inputs, self.num_local_experts, dim=1)
expert_output = []
for i in range(self.num_local_experts):
expert_output.append(self.experts[i](expert_input[i]))
output = torch.cat(expert_output, dim=1).contiguous()
return output
class FFNExperts(nn.Module):
def __init__(self, num_experts: int, d_model: int, d_ff: int, activation=None, drop_rate: float = 0):
super().__init__()
assert num_experts % moe_env.model_parallel_size == 0, \
"The number of experts should be divied by moe model size"
num_local_experts = num_experts // moe_env.model_parallel_size
self.w1 = nn.Parameter(torch.empty(num_local_experts, d_model, d_ff, device=get_current_device()))
self.b1 = nn.Parameter(torch.empty(num_local_experts, 1, d_ff, device=get_current_device()))
self.w2 = nn.Parameter(torch.empty(num_local_experts, d_ff, d_model, device=get_current_device()))
self.b2 = nn.Parameter(torch.empty(num_local_experts, 1, d_model, device=get_current_device()))
s1 = math.sqrt(0.1 / d_model)
s2 = math.sqrt(0.1 / d_ff)
nn.init.trunc_normal_(self.w1, std=s1)
nn.init.trunc_normal_(self.b1, std=s1)
nn.init.trunc_normal_(self.w2, std=s2)
nn.init.trunc_normal_(self.b2, std=s2)
self.act = nn.GELU() if activation is None else activation
self.drop = nn.Dropout(p=drop_rate)
for param in self.parameters():
param.__setattr__('moe_param', True)
def forward(self, inputs): # x [g, el, c, h]
el = inputs.size(1)
h = inputs.size(-1)
inputs = inputs.transpose(0, 1)
inshape = inputs.shape
inputs = inputs.reshape(el, -1, h)
out_ff = torch.baddbmm(self.b1, inputs, self.w1)
out_act = self.act(out_ff)
with seed(ParallelMode.TENSOR):
inter = self.drop(out_act)
out_model = torch.baddbmm(self.b2, inter, self.w2)
with seed(ParallelMode.TENSOR):
outputs = self.drop(out_model) # outputs [el, gc, h]
outputs = outputs.reshape(inshape)
outputs = outputs.transpose(0, 1).contiguous()
return outputs

300
colossalai/nn/layer/moe/layers.py
View File

@@ -3,70 +3,13 @@ import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.cuda.amp import autocast
import torch.distributed as dist
from colossalai.core import global_context as gpc
from colossalai.global_variables import moe_env
from colossalai.context import ParallelMode, seed
from colossalai.context import ParallelMode
from colossalai.utils import get_current_device
from ._operation import AllToAll
class NormalNoiseGenerator:
"""Generates a random noisy mask for logtis tensor.
All noise is generated from a normal distribution (0, 1 / E^2), where
E = the number of experts.
:param num_experts: The number of experts
:type num_experts: int
"""
def __init__(self, num_experts: int):
self.normal = torch.distributions.normal.Normal(
loc=torch.tensor(0.0, device=get_current_device()),
scale=torch.tensor(1.0 / num_experts ** 2, device=get_current_device())
).rsample
def __call__(self, inputs: torch.Tensor):
noisy = self.normal(inputs.shape)
return inputs + noisy
class Experts(nn.Module):
"""A wrapper class to create experts. It will create E experts across the
moe model parallel group, where E is the number of experts. Every expert
is a instence of the class, 'expert' in initialization parameters.
:param expert: The class of all experts
:param num_experts: The number of experts
:param expert_args: Args used to initialize experts
:type num_experts: int
"""
def __init__(self, expert, num_experts, **expert_args):
super().__init__()
assert num_experts % moe_env.model_parallel_size == 0, \
"The number of experts should be divied by moe model size"
num_local_experts = num_experts // moe_env.model_parallel_size
with seed(ParallelMode.MOE_MODEL):
self.experts = nn.ModuleList([
expert(**expert_args) for _ in range(num_local_experts)])
self.num_local_experts = num_local_experts
for exp in self.experts:
for param in exp.parameters():
param.__setattr__('moe_param', 1)
def forward(self, inputs):
expert_input = torch.chunk(inputs, self.num_local_experts, dim=0)
expert_output = []
for i in range(self.num_local_experts):
expert_output.append(self.experts[i](expert_input[i]))
output = torch.cat(expert_output, dim=0)
return output
from ._operation import U_CUDA_MODE, AllToAll, MoeDispatch, MoeCombine, moe_cumsum
from .utils import autocast_softmax
class Top1Router(nn.Module):
@@ -83,63 +26,79 @@ class Top1Router(nn.Module):
:type noisy_func: Callable, optional
"""
def __init__(self,
capacity_factor: float,
min_capacity: int,
noisy_func=None):
def __init__(self, capacity_factor: float, min_capacity: int = 0, select_policy: str = "first", noisy_func=None):
super().__init__()
self.capacity_factor = capacity_factor
self.min_capacity = min_capacity
self.select_policy = select_policy
self.noisy_func = noisy_func
self.uniform = torch.distributions.uniform.Uniform(
low=torch.tensor(0.0, device=get_current_device()),
high=torch.tensor(1.0, device=get_current_device())).rsample
def get_capacity(self, logits_shape):
capacity = math.ceil(self.capacity_factor *
logits_shape[0] / logits_shape[1])
if capacity < self.min_capacity:
capacity = self.min_capacity
assert select_policy in {"first", "random"}
if select_policy == "random":
self.uniform = torch.distributions.uniform.Uniform(low=torch.tensor(0.0, device=get_current_device()),
high=torch.tensor(1.0,
device=get_current_device())).rsample
def get_capacity(
self,
logits_shape,
):
capacity = math.floor(self.capacity_factor * logits_shape[-2] / logits_shape[-1])
capacity += capacity % 2
capacity = max(capacity, self.min_capacity)
assert capacity > 0
return capacity
def forward(self, inputs):
def forward(self, inputs: torch.Tensor, cuda_mode: bool = False):
if self.noisy_func is not None:
inputs_noisy = self.noisy_func(inputs)
else:
inputs_noisy = inputs
logits = F.softmax(inputs, dim=1)
num_experts = logits.shape[1]
logits = autocast_softmax(inputs, dim=-1)
num_experts = logits.size(-1)
capacity = self.get_capacity(logits.shape)
expert_idx = torch.argmax(inputs_noisy, dim=1)
expert_mask = F.one_hot(expert_idx, num_classes=num_experts)
expert_mask_f = expert_mask.float()
top1_idx = torch.argmax(inputs_noisy, dim=-1)
mask = F.one_hot(top1_idx, num_classes=num_experts).to(torch.int32)
exp_counts = torch.sum(expert_mask, dim=0).detach().to('cpu')
if self.training:
me = torch.mean(logits, dim=0)
ce = torch.mean(mask.float(), dim=0)
l_aux = num_experts * torch.sum(me * ce)
moe_env.add_loss(l_aux)
else:
max_num = torch.max(torch.sum(mask, dim=0))
dist.all_reduce(max_num, op=dist.ReduceOp.MAX, group=gpc.get_group(ParallelMode.MOE_MODEL))
capacity = max_num.item()
me = torch.mean(logits, dim=0)
ce = torch.mean(expert_mask_f, dim=0)
l_aux = torch.sum(me * ce) * num_experts
moe_env.add_loss(l_aux)
if not self.training:
ranks = moe_cumsum(mask)
elif self.select_policy == "random":
rand_mask = mask * self.uniform(mask.shape)
_, dispatch_idx = torch.topk(rand_mask, k=capacity, dim=0)
mask = mask * torch.zeros_like(mask).scatter_(0, dispatch_idx, 1)
ranks = moe_cumsum(mask)
elif self.select_policy == "first":
ranks = moe_cumsum(mask)
mask = mask * torch.lt(ranks, capacity)
else:
raise NotImplementedError("Not support such select policy yet.")
rand_mask = expert_mask * self.uniform(logits.shape)
_, dispatch_idx = torch.topk(rand_mask, k=capacity, dim=0)
ranks = torch.sum(mask * ranks, dim=-1)
dispatch_mask = \
expert_mask * torch.zeros_like(expert_mask).scatter_(0, dispatch_idx, 1)
locations = torch.cumsum(dispatch_mask, dim=0) - 1
locations = torch.sum(dispatch_mask * locations, dim=1)
locations = F.one_hot(locations, num_classes=capacity)
logits = logits * dispatch_mask
combine_weights = logits.unsqueeze(2) * locations.unsqueeze(1)
sec_mask = combine_weights.bool()
return combine_weights, sec_mask, exp_counts
if cuda_mode:
mask = torch.sum(mask, dim=-1)
mask = torch.stack([mask], dim=0).to(torch.int32)
dest_idx = torch.stack([top1_idx * capacity + ranks], dim=0).to(torch.int32)
return logits, mask, dest_idx, num_experts * capacity
else:
ranks = F.one_hot(ranks, num_classes=capacity)
weight = mask * logits.type_as(inputs)
combine_weights = weight.unsqueeze(2) * ranks.unsqueeze(1)
sec_mask = combine_weights.bool()
return combine_weights, sec_mask
class Top2Router(nn.Module):
@@ -159,53 +118,67 @@ class Top2Router(nn.Module):
self.noisy_func = noisy_func
def get_capacity(self, logits_shape):
capacity = math.ceil(2 * self.capacity_factor *
logits_shape[0] / logits_shape[1])
capacity = math.floor(2 * self.capacity_factor * logits_shape[-2] / logits_shape[-1])
capacity += capacity % 2
assert capacity > 0
return capacity
def forward(self, inputs):
def forward(self, inputs: torch.Tensor, cuda_mode: bool = False):
# inputs: [s, h]
if self.noisy_func is not None:
inputs = self.noisy_func(inputs)
logits = F.softmax(inputs, dim=-1)
logits = autocast_softmax(inputs, dim=-1) # logits: [s, e]
num_experts = logits.size(-1)
capacity = self.get_capacity(logits.shape)
_, expert_idx = torch.topk(logits, k=2, dim=-1, largest=True, sorted=True)
top1_idx = expert_idx[:, 0]
top2_idx = expert_idx[:, 1]
top1_idx = torch.argmax(logits, dim=-1)
mask1 = F.one_hot(top1_idx, num_classes=num_experts).to(torch.int32)
logits_except1 = logits.masked_fill(mask1.bool(), float("-inf"))
top2_idx = torch.argmax(logits_except1, dim=-1)
mask2 = F.one_hot(top2_idx, num_classes=num_experts).to(torch.int32)
mask1 = F.one_hot(top1_idx, num_classes=num_experts)
mask2 = F.one_hot(top2_idx, num_classes=num_experts)
cmask = (mask1 + mask2) # loss: [s, e]
if self.training:
me = torch.mean(logits, dim=0)
ce = torch.mean(cmask.float(), dim=0)
l_aux = num_experts * torch.sum(me * ce) / 2.0
moe_env.add_loss(l_aux)
else:
max_num = torch.max(torch.sum(cmask, dim=0))
dist.all_reduce(max_num, op=dist.ReduceOp.MAX, group=gpc.get_group(ParallelMode.MOE_MODEL))
capacity = max_num.item()
loss_mask = (mask1 + mask2)
exp_counts = torch.sum(loss_mask, dim=0).detach().to('cpu')
me = torch.mean(logits, dim=0)
ce = torch.mean(loss_mask.float(), dim=0)
l_aux = num_experts * torch.sum(me * ce) / 2.0
moe_env.add_loss(l_aux)
rank1 = moe_cumsum(mask1) # rank1: [s, e]
rank2 = moe_cumsum(mask2)
rank2 += torch.sum(mask1, dim=-2, keepdim=True)
locations1 = torch.cumsum(mask1, dim=0) - 1
locations2 = torch.cumsum(mask2, dim=0) - 1
locations2 += torch.sum(mask1, dim=0, keepdim=True)
mask1 *= torch.lt(rank1, capacity)
mask2 *= torch.lt(rank2, capacity)
mask1 *= torch.lt(locations1, capacity)
mask2 *= torch.lt(locations2, capacity)
rank1 = torch.sum(mask1 * rank1, dim=-1)
rank2 = torch.sum(mask2 * rank2, dim=-1)
weight1 = mask1 * logits
weight2 = mask2 * logits
if cuda_mode:
mask1 = torch.sum(mask1, dim=-1)
mask2 = torch.sum(mask2, dim=-1)
locations1 = torch.sum(mask1 * locations1, dim=1)
locations2 = torch.sum(mask2 * locations2, dim=1)
locations1_sc = F.one_hot(locations1, num_classes=capacity)
locations2_sc = F.one_hot(locations2, num_classes=capacity)
mask = torch.stack([mask1, mask2], dim=0).to(torch.int32)
dest_idx = torch.stack([top1_idx * capacity + rank1, top2_idx * capacity + rank2], dim=0).to(torch.int32)
combine_weights1 = weight1.unsqueeze(2) * locations1_sc.unsqueeze(1)
combine_weights2 = weight2.unsqueeze(2) * locations2_sc.unsqueeze(1)
combine_weights = combine_weights1 + combine_weights2
sec_mask = combine_weights.bool()
return logits, mask, dest_idx, num_experts * capacity
else:
weight1 = mask1 * logits.type_as(inputs)
weight2 = mask2 * logits.type_as(inputs)
rank1_sc = F.one_hot(rank1, num_classes=capacity)
rank2_sc = F.one_hot(rank2, num_classes=capacity)
return combine_weights, sec_mask, exp_counts
cb_weight1 = weight1.unsqueeze(2) * rank1_sc.unsqueeze(1)
cb_weight2 = weight2.unsqueeze(2) * rank2_sc.unsqueeze(1)
cb_weight = cb_weight1 + cb_weight2
sec_mask = cb_weight.bool()
return cb_weight, sec_mask
class MoeLayer(nn.Module):
@@ -225,52 +198,47 @@ class MoeLayer(nn.Module):
:type experts: nn.Module
"""
def __init__(self,
dim_model: int,
num_experts: int,
router: nn.Module,
experts: nn.Module):
def __init__(self, dim_model: int, num_experts: int, router: nn.Module, experts: nn.Module):
super().__init__()
self.d_model = dim_model
self.num_experts = num_experts
self.gate = nn.Linear(dim_model, num_experts, device=get_current_device())
self.gate = nn.Linear(dim_model, num_experts, bias=False, device=get_current_device())
self.router = router
self.experts = experts
self.cuda_mode = True if U_CUDA_MODE and moe_env.enable_cuda else False
def _router_part(self, tokens: torch.Tensor):
gate_output = self.gate(tokens)
return self.router(gate_output)
def expert_part(self, expert_input: torch.Tensor):
expert_input = AllToAll.apply(expert_input, ParallelMode.MOE_MODEL)
def router_part(self, tokens: torch.Tensor):
autocast_context = torch.is_autocast_enabled()
if not autocast_context:
return self._router_part(tokens)
else:
with autocast(enabled=False):
if tokens.dtype == torch.float16:
input_tokens = tokens.float()
else:
input_tokens = tokens
return self._router_part(input_tokens)
input_shape = expert_input.shape
expert_input = expert_input.reshape(moe_env.model_parallel_size,
self.num_experts // moe_env.model_parallel_size, -1, self.d_model)
expert_output = self.experts(expert_input)
expert_output = expert_output.reshape(input_shape)
expert_output = AllToAll.apply(expert_output, ParallelMode.MOE_MODEL)
return expert_output
def forward(self, inputs: torch.Tensor) -> torch.Tensor:
tokens = inputs.reshape(-1, self.d_model)
gate_output = self.gate(tokens)
router_res = self.router(gate_output, self.cuda_mode)
combine_weights, sec_mask, exp_counts = self.router_part(tokens)
if self.cuda_mode:
logits, mask, dest_idx, ec = router_res
expert_input = MoeDispatch.apply(tokens, mask, dest_idx, ec)
expert_output = self.expert_part(expert_input)
ret = MoeCombine.apply(expert_output, logits, mask, dest_idx, ec)
else:
combine_weights, sec_mask = router_res
sec_mask_f = sec_mask.type_as(inputs)
expert_input = torch.matmul(sec_mask_f.permute(1, 2, 0), tokens)
expert_output = self.expert_part(expert_input)
combine_weights = combine_weights.view(combine_weights.shape[0], -1)
expert_output = expert_output.view(-1, expert_output.shape[-1])
ret = torch.matmul(combine_weights, expert_output)
sec_mask_f = sec_mask.type_as(inputs)
dispatch_data = torch.matmul(sec_mask_f.permute(1, 2, 0), tokens)
dispatch_data = AllToAll.apply(dispatch_data, ParallelMode.MOE_MODEL)
expert_output = self.experts(dispatch_data)
expert_output = AllToAll.apply(expert_output, ParallelMode.MOE_MODEL)
combine_weights = combine_weights.view(combine_weights.shape[0], -1)
expert_output = expert_output.view(-1, expert_output.shape[-1])
ret = torch.matmul(combine_weights, expert_output)
ret = ret.reshape(inputs.shape)
return ret

32
colossalai/nn/layer/moe/utils.py Normal file
View File

@@ -0,0 +1,32 @@
import torch
import torch.nn.functional as F
from colossalai.utils import get_current_device
class NormalNoiseGenerator:
"""Generates a random noisy mask for logtis tensor.
All noise is generated from a normal distribution (0, 1 / E^2), where
E = the number of experts.
:param num_experts: The number of experts
:type num_experts: int
"""
def __init__(self, num_experts: int):
self.normal = torch.distributions.normal.Normal(
loc=torch.tensor(0.0, device=get_current_device()),
scale=torch.tensor(1.0 / num_experts ** 2, device=get_current_device())
).rsample
def __call__(self, inputs: torch.Tensor):
noisy = self.normal(inputs.shape)
return inputs + noisy
def autocast_softmax(inputs: torch.Tensor, dim: int):
assert inputs.dtype in {torch.float16, torch.float32}
fp16_flag = (inputs.dtype == torch.float16)
sm_input = inputs.to(torch.float32) if fp16_flag else inputs
sm_output = F.softmax(sm_input, dim)
return sm_output