Skip to content

Add Inference support for running the BigScience-BLOOM Architecture #2083

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jeffra merged 27 commits into from
Jul 18, 2022
Merged

Add Inference support for running the BigScience-BLOOM Architecture #2083

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
27 commits
Select commit Hold shift + click to select a range
7cf7c89
Add Inference support for running the BigScience-BLOOM Architecture
Jul 9, 2022
6d33b82
Merge branch 'master' into ds-inference/bloom-support
jeffra Jul 11, 2022
6cc2340
formatting
jeffra Jul 11, 2022
2a69756
add the checkpoint loading at the same of kernel-injection
Jul 12, 2022
e449ac7
Merge branch 'ds-inference/bloom-support' of github.com:microsoft/Dee…
Jul 12, 2022
6e304aa
releasing checkpoint CPU-memory after loading it
Jul 12, 2022
16dd1dc
some fixes plus formatting
Jul 12, 2022
e0dd488
fix layer_past; this caused issues when running inference on several …
Jul 13, 2022
02f9945
Add support for multi-batch inference
Jul 13, 2022
1cef202
fix the padding issue for large bach inference
Jul 14, 2022
644fea4
fixing some bug in softmax kernel for batch_size>1
Jul 14, 2022
adb0b97
align alibi&mask addition with HF new changes
Jul 15, 2022
bd3c0a0
revert back some changes and support for very large batch size
Jul 15, 2022
1f92e55
reduce the max_token_length for now
Jul 15, 2022
5d1f351
fix mask-adding
Jul 15, 2022
3c12b89
fix the large-batch inference for MP > 1
Jul 16, 2022
a5bdd58
Merge branch 'master' into ds-inference/bloom-support
RezaYazdaniAminabadi Jul 18, 2022
aa5e01f
Ds inference/bloom support meta (#2104)
jeffra Jul 18, 2022
b6503ed
fix the Bert and GPT-J unit tests
Jul 18, 2022
a9459d6
fix for OneDevice
jeffra Jul 18, 2022
2cd301e
Merge branch 'master' into ds-inference/bloom-support
jeffra Jul 18, 2022
72aba56
added bloom inference tests
mrwyattii Jul 18, 2022
332f69d
fixing the masking stride for the GPT models
Jul 18, 2022
0ddf41c
Merge branch 'ds-inference/bloom-support' of github.com:microsoft/Dee…
Jul 18, 2022
d38464e
revert back some changes on replace_module
Jul 18, 2022
ac2d092
fix fp32 softmax
Jul 18, 2022
1e3ea74
allocate tensors initially on cpu at inference-api
Jul 18, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 12 additions & 12 deletions csrc/transformer/inference/csrc/gelu.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -174,7 +174,7 @@ __global__ void fused_bias_residual(float* input,
float* attnbias,
int total_count,
int intermediate_size,
int mp_size,
float mp_size,
bool preln)
{
float4* input_cast = reinterpret_cast<float4*>(input);
Expand All @@ -191,10 +191,10 @@ __global__ void fused_bias_residual(float* input,
float4 bias_data = bias_cast[offset % intermediate_size];
float4 attn_bias = attnbias_cast[offset % intermediate_size];
if (preln) {
data.x = (data.x + res_vec.x) * mp_size + (out.x + bias_data.x + attn_bias.x);
data.y = (data.y + res_vec.y) * mp_size + (out.y + bias_data.y + attn_bias.y);
data.z = (data.z + res_vec.z) * mp_size + (out.z + bias_data.z + attn_bias.z);
data.w = (data.w + res_vec.w) * mp_size + (out.w + bias_data.w + attn_bias.w);
data.x = (data.x + res_vec.x + bias_data.x + attn_bias.x) * mp_size + (out.x);
data.y = (data.y + res_vec.y + bias_data.y + attn_bias.y) * mp_size + (out.y);
data.z = (data.z + res_vec.z + bias_data.z + attn_bias.z) * mp_size + (out.z);
data.w = (data.w + res_vec.w + bias_data.w + attn_bias.w) * mp_size + (out.w);
} else {
data.x = data.x + out.x + bias_data.x;
data.y = data.y + out.y + bias_data.y;
Expand All @@ -212,7 +212,7 @@ __global__ void fused_bias_residual(__half* input,
__half* attn_bias,
int total_count,
int intermediate_size,
int mp_size,
float mp_size,
Comment thread
RezaYazdaniAminabadi marked this conversation as resolved.
Outdated
bool preln)
{
#ifdef HALF_PRECISION_AVAILABLE
Expand Down Expand Up @@ -257,13 +257,13 @@ __global__ void fused_bias_residual(__half* input,

if (preln) {
low_data.x =
(low_data.x + low_res.x) * mp_size + (low_out.x + (low_bias.x + attn_low_bias.x));
(low_data.x + low_res.x + (low_bias.x + attn_low_bias.x)) * mp_size + low_out.x;
low_data.y =
(low_data.y + low_res.y) * mp_size + (low_out.y + (low_bias.y + attn_low_bias.y));
high_data.x = (high_data.x + high_res.x) * mp_size +
(high_out.x + (high_bias.x + attn_high_bias.x));
high_data.y = (high_data.y + high_res.y) * mp_size +
(high_out.y + (high_bias.y + attn_high_bias.y));
(low_data.y + low_res.y + (low_bias.y + attn_low_bias.y)) * mp_size + low_out.y;
high_data.x = (high_data.x + high_res.x + (high_bias.x + attn_high_bias.x)) * mp_size +
high_out.x;
high_data.y = (high_data.y + high_res.y + (high_bias.y + attn_high_bias.y)) * mp_size +
high_out.y;
} else {
low_data.x = (low_data.x + low_out.x + low_bias.x);
low_data.y = (low_data.y + low_out.y + low_bias.y);
Expand Down
3 changes: 2 additions & 1 deletion csrc/transformer/inference/csrc/normalize.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -88,6 +88,7 @@ __global__ void fused_bias_residual_layer_norm(__half* output,
int row_stride)
{
#ifdef HALF_PRECISION_AVAILABLE

int iteration_stride = blockDim.x;
int iterations = row_stride / iteration_stride;

Expand All @@ -99,7 +100,7 @@ __global__ void fused_bias_residual_layer_norm(__half* output,
int gid = id >> 5;
int warp_num = iteration_stride >> 5;

__half2 inp_reg[NORM_REG];
__half2 inp_reg[8];
Comment thread
RezaYazdaniAminabadi marked this conversation as resolved.
Outdated

const __half2* vals_cast = reinterpret_cast<const __half2*>(vals);
__half2* out_cast = reinterpret_cast<__half2*>(output);
Expand Down
27 changes: 21 additions & 6 deletions csrc/transformer/inference/csrc/pt_binding.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,8 @@ at::Tensor ds_softmax(at::Tensor& attn_scores,
bool recompute,
bool local_attention,
int window_size,
bool async_op)
bool async_op,
float layer_scale)
{
auto attn_scores_c = attn_scores.contiguous();
int bsz = attn_scores_c.size(0);
Expand All @@ -32,6 +33,8 @@ at::Tensor ds_softmax(at::Tensor& attn_scores,

launch_attn_softmax_v2((T*)attn_scores_c.data_ptr(),
(attn_mask.sizes().size() > 1 ? (T*)attn_mask.data_ptr() : nullptr),
(T*)attn_mask.data_ptr(),
layer_scale,
triangular,
recompute,
local_attention,
Expand Down Expand Up @@ -145,7 +148,7 @@ void attention_unfused(at::Tensor& prev_key_cont,
CUBLAS_GEMM_DEFAULT_TENSOR_OP);
#endif
attn_score = ds_softmax<T>(
attn_score, attn_mask, triangular, recompute, local_attention, window_size, false);
attn_score, attn_mask, triangular, recompute, local_attention, window_size, false, 1.0);
alpha = 1.0;
cublas_strided_batched_gemm(Context::Instance().GetCublasHandle(),
k,
Expand Down Expand Up @@ -225,6 +228,8 @@ std::vector<at::Tensor> ds_softmax_context1(at::Tensor& query,
template <typename T>
void ds_softmax_internal(T* attn_scores,
at::Tensor& attn_mask,
at::Tensor& alibi,
float& layer_scale,
bool triangular,
bool recompute,
bool local_attention,
Expand All @@ -236,6 +241,8 @@ void ds_softmax_internal(T* attn_scores,
{
launch_attn_softmax_v2((T*)attn_scores,
(attn_mask.sizes().size() > 1 ? (T*)attn_mask.data_ptr() : nullptr),
(alibi.sizes().size() > 1 ? (T*)alibi.data_ptr() : nullptr),
layer_scale,
triangular,
recompute,
local_attention,
Expand Down Expand Up @@ -263,9 +270,12 @@ void attention_unfused(T* prev_key_cont,
bool triangular,
bool recompute,
bool local_attention,
int window_size)
int window_size,
at::Tensor& alibi,
int layer_id)
{
float alpha = norm_factor * norm_factor;
float layer_scale = std::max(1, layer_id);
float alpha = norm_factor * norm_factor / layer_scale;
float gemm_beta = 0.0;
T* workspace = (T*)output + bsz * seq_len * heads * k;

Expand All @@ -292,6 +302,8 @@ void attention_unfused(T* prev_key_cont,
#endif
ds_softmax_internal<T>(workspace,
attn_mask,
alibi,
layer_scale,
triangular,
recompute,
local_attention,
Expand Down Expand Up @@ -336,7 +348,8 @@ std::vector<at::Tensor> ds_softmax_context(at::Tensor& query_key_value,
int window_size,
bool no_masking,
unsigned layer_id,
unsigned num_layers)
unsigned num_layers,
at::Tensor& alibi)
{
unsigned bsz = query_key_value.size(0);
unsigned seq_len = query_key_value.size(1);
Expand Down Expand Up @@ -410,7 +423,9 @@ std::vector<at::Tensor> ds_softmax_context(at::Tensor& query_key_value,
(triangular && is_prompt),
is_prompt,
local_attention,
window_size);
window_size,
alibi,
layer_id);
launch_transform4d_0213<T>((T*)output.data_ptr(),
temp_buf,
bsz,
Expand Down
36 changes: 34 additions & 2 deletions csrc/transformer/inference/csrc/softmax.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,8 @@ namespace cg = cooperative_groups;

__global__ void attn_softmax_v2(__half* vals,
__half* mask,
__half* alibi,
float layer_scale,
bool triangular,
bool recompute,
bool local_attention,
Expand Down Expand Up @@ -64,7 +66,7 @@ __global__ void attn_softmax_v2(__half* vals,
if (iter_offset < total_count) {
vals += (iter_offset * sequence_length);

int mask_offset = (iter_offset / (heads * num_seq)) * (sequence_length);
int mask_offset = (iter_offset / (num_seq)) * (sequence_length);
int seq_id = iter_offset % num_seq;
int seq_id4 = seq_id >> 2;

Expand All @@ -76,7 +78,7 @@ __global__ void attn_softmax_v2(__half* vals,
(local_attention && real_seq_id >= window_size) ? real_seq_id - window_size : -1;

float max_val = minus_infinity;

// if (lane == 0) printf("%d, %d: %d \n", wid, blockIdx.x, mask_offset);
for (int i = 0; i < iterations; i++) {
int data_id = i * (reduceWidth << 2) + (seq_lane << 2);
if ((!triangular || ((data_id >> 2) <= seq_id4)) && (data_id >> 2) >= window_stride4 &&
Expand All @@ -102,6 +104,16 @@ __global__ void attn_softmax_v2(__half* vals,
high_data[i].x += __half2float(mask[data_id + mask_offset + 2]);
high_data[i].y += __half2float(mask[data_id + mask_offset + 3]);
}
if (alibi) {
low_data[i].x = low_data[i].x * layer_scale +
__half2float(alibi[data_id + mask_offset]);
low_data[i].y = low_data[i].y * layer_scale +
__half2float(alibi[data_id + mask_offset + 1]);
high_data[i].x = high_data[i].x * layer_scale +
__half2float(alibi[data_id + mask_offset + 2]);
high_data[i].y = high_data[i].y * layer_scale +
__half2float(alibi[data_id + mask_offset + 3]);
}
} else {
low_data[i].x = data_id > window_stride ? __half2float(vals[data_id])
: minus_infinity;
Expand All @@ -123,6 +135,16 @@ __global__ void attn_softmax_v2(__half* vals,
if ((data_id + 2) < sequence_length)
high_data[i].x += __half2float(mask[data_id + mask_offset + 2]);
}
if (alibi) {
low_data[i].x = low_data[i].x * layer_scale +
__half2float(alibi[data_id + mask_offset]);
if ((data_id + 1) < sequence_length)
low_data[i].y = low_data[i].y * layer_scale +
__half2float(alibi[data_id + mask_offset + 1]);
if ((data_id + 2) < sequence_length)
high_data[i].x = high_data[i].x * layer_scale +
__half2float(alibi[data_id + mask_offset + 2]);
}
}
// if(lane == 0) printf("%f , %d, %d \n", low_data[i].x, data_id, seq_id);
max_val = (low_data[i].x > max_val ? low_data[i].x : max_val);
Expand Down Expand Up @@ -204,6 +226,8 @@ __global__ void attn_softmax_v2(__half* vals,

__global__ void attn_softmax_v2(float* vals,
float* attn_mask,
float* alibi,
float layer_scale,
bool triangular,
bool recompute,
bool local_attention,
Expand Down Expand Up @@ -371,6 +395,8 @@ __global__ void attn_softmax_v2(float* vals,
template <typename T>
void launch_attn_softmax_v2(T* vals,
T* mask,
T* alibi,
float layer_scale,
bool triangular,
bool recompute,
bool local_attention,
Expand All @@ -393,6 +419,8 @@ void launch_attn_softmax_v2(T* vals,
attn_softmax_v2<<<grid_dim, block_dim, 0, stream>>>(
vals,
mask,
alibi,
layer_scale,
triangular,
recompute,
local_attention,
Expand All @@ -410,6 +438,8 @@ void launch_attn_softmax_v2(T* vals,

template void launch_attn_softmax_v2(float* vals,
float* mask,
float* alibi,
float layer_scale,
bool triangular,
bool recompute,
bool local_attention,
Expand All @@ -422,6 +452,8 @@ template void launch_attn_softmax_v2(float* vals,
cudaStream_t stream);
template void launch_attn_softmax_v2(__half* vals,
__half* mask,
__half* alibi,
float layer_scale,
bool triangular,
bool recompute,
bool local_attention,
Expand Down
1 change: 1 addition & 0 deletions csrc/transformer/inference/includes/context.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -89,6 +89,7 @@ class Context {
cudaMalloc(&_workspace, size);
}

if (!_workspace) { throw std::runtime_error("Workspace is null."); }
_workSpaceSize = size;
}

Expand Down
4 changes: 3 additions & 1 deletion csrc/transformer/inference/includes/custom_cuda_layers.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@
#include <cassert>
#include <iostream>

#define MAX_OUT_TOKES 1024
#define MAX_OUT_TOKES 128
#define MAX_WARP_NUM 32
#define WARP_SIZE 32

Expand All @@ -28,6 +28,8 @@
template <typename T>
void launch_attn_softmax_v2(T* vals,
T* mask,
T* alibi,
float layer_scale,
bool triangular,
bool recompute,
bool local_attention,
Expand Down
Loading