Skip to content

Extend scratch buffer for long prompts #2212

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
Sep 23, 2022
Merged

Extend scratch buffer for long prompts #2212

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
27 commits
Select commit Hold shift + click to select a range
4844966
Extend scratch buffer for long prompts
cmikeh2 Aug 12, 2022
be4714d
Merge branch 'master' into cholmes/fix-long-seq-len-inference
cmikeh2 Aug 12, 2022
c6411d1
Fetch correct tail buffer for batched inputs.
cmikeh2 Aug 12, 2022
c074ed3
Style change
cmikeh2 Aug 12, 2022
3135774
Merge branch 'cholmes/fix-long-seq-len-inference' of https://github.c…
cmikeh2 Aug 12, 2022
777a36e
Merge branch 'master' into cholmes/fix-long-seq-len-inference
RezaYazdaniAminabadi Aug 13, 2022
ec6b1ad
Fix variable rename
cmikeh2 Aug 16, 2022
d897f98
Merge branch 'cholmes/fix-long-seq-len-inference' of https://github.c…
cmikeh2 Aug 16, 2022
6da9234
Merge branch 'master' into cholmes/fix-long-seq-len-inference
cmikeh2 Aug 16, 2022
606d344
Reduce maximum sequence length
cmikeh2 Aug 16, 2022
5269ba1
Merge branch 'master' into cholmes/fix-long-seq-len-inference
cmikeh2 Aug 23, 2022
89f2ded
Merge branch 'master' into cholmes/fix-long-seq-len-inference
RezaYazdaniAminabadi Sep 1, 2022
8e29808
Merge branch 'master' into cholmes/fix-long-seq-len-inference
cmikeh2 Sep 9, 2022
c824330
Add debug print
cmikeh2 Sep 9, 2022
d9eb076
Merge branch 'master' into cholmes/fix-long-seq-len-inference
cmikeh2 Sep 9, 2022
aafba00
Multi-batch inference fix
cmikeh2 Sep 10, 2022
4abd455
add batch-size at the tranform launch for the half-precision implemen…
Sep 11, 2022
603cc5b
Merge branch 'master' into cholmes/fix-long-seq-len-inference
RezaYazdaniAminabadi Sep 13, 2022
508712a
Merge branch 'master' into cholmes/fix-long-seq-len-inference
RezaYazdaniAminabadi Sep 19, 2022
51a6371
no need to throw error when there is no mask passed
Sep 22, 2022
9effa9e
Merge branch 'cholmes/fix-long-seq-len-inference' of github.com:micro…
Sep 22, 2022
c9652ec
Merge branch 'master' into cholmes/fix-long-seq-len-inference
RezaYazdaniAminabadi Sep 22, 2022
d8f5203
Increasing the token-length based on available memory for GPT models …
RezaYazdaniAminabadi Sep 22, 2022
b64c117
Merge branch 'master' into cholmes/fix-long-seq-len-inference
RezaYazdaniAminabadi Sep 22, 2022
48a8b96
fix bert issue & remove some dead code
Sep 22, 2022
c1d83f9
fix formating
Sep 22, 2022
a5f4d31
Merge branch 'master' into cholmes/fix-long-seq-len-inference
jeffra Sep 22, 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
41 changes: 32 additions & 9 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 @@ -59,13 +59,17 @@ at::Tensor ds_softmax(at::Tensor& attn_scores,

template <typename T>
void allocate_workspace(size_t hidden_dim,
size_t prompt_len,
size_t max_seq_len,
size_t batch_size,
unsigned num_layers,
size_t head_size = 128)
{
size_t _workSpaceSize = 16 * (hidden_dim * batch_size * max_seq_len) +
(num_layers * batch_size * max_seq_len * hidden_dim * 2); // KV-cache
size_t prompt_output = batch_size * head_size * prompt_len * prompt_len;
_workSpaceSize = _workSpaceSize + prompt_output;

Context::Instance().GenWorkSpace(_workSpaceSize * sizeof(T));
}

Expand All @@ -81,10 +85,13 @@ at::Tensor einsum_sec_sm_ecm(at::Tensor& Q, at::Tensor& W)
float alpha = 1;
float gemm_beta = 0.0;

if (!workspace) {
allocate_workspace<T>(W.size(1), MAX_OUT_TOKES, Q.size(0), 1);
/*
// Reallocate memory if we received a new prompt
if (!workspace || input.size(1) != 1) {
allocate_workspace<T>(W.size(1), MAX_OUT_TOKES, Q.size(0), 1, head_size);
workspace = (T*)Context::Instance().GetWorkSpace();
}
*/

auto O = at::from_blob(workspace, {Q.size(1), Q.size(2), W.size(1)}, options);
unsigned m = W.size(1);
Expand Down Expand Up @@ -305,7 +312,15 @@ void attention_unfused(T* prev_key_cont,
float layer_scale = alibi.sizes().size() > 1 ? std::max(1, layer_id) : 1.0;
float alpha = norm_factor * norm_factor / layer_scale;
float gemm_beta = 0.0;
T* workspace = (T*)output + bsz * seq_len * heads * k;
T* workspace;
if (seq_len == 1) {
workspace = (T*)output + bsz * seq_len * heads * k;
} else {
// If we are doing the prompt, switch to the tail workspace
T* scratch = (T*)Context::Instance().GetWorkSpace();
workspace = scratch + (Context::Instance().get_workspace_size() / sizeof(T)) -
heads * seq_len * seq_len;

@RezaYazdaniAminabadi RezaYazdaniAminabadi Aug 12, 2022

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@cmikeh2, do we want to add bsz here too, since we might get a batched input?

@cmikeh2 cmikeh2 Aug 12, 2022

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, this is primed to cause a segfault on any batched input.

}

cublasSetStream(Context::Instance().GetCublasHandle(), Context::Instance().GetCurrentStream());
cublas_strided_batched_gemm(Context::Instance().GetCublasHandle(),
Expand Down Expand Up @@ -591,14 +606,19 @@ std::vector<at::Tensor> ds_qkv_gemm(at::Tensor& input,
at::Tensor& beta,
const float epsilon,
bool add_bias,
unsigned num_layers)
unsigned num_layers,
bool bloom_seq_len,

@RezaYazdaniAminabadi RezaYazdaniAminabadi Aug 12, 2022

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

maybe change this to is_bloom?

unsigned head_size)
{
int bsz = input.size(0) * input.size(1);
T* workspace = (T*)Context::Instance().GetWorkSpace();
if (!workspace) {
// Reallocate memory if we receive a new prompt
if (!workspace || input.size(1) != 1) {
cublasSetStream(Context::Instance().GetCublasHandle(),
Context::Instance().GetCurrentStream());
allocate_workspace<T>(input.size(2), MAX_OUT_TOKES, input.size(0), num_layers);
const int max_seq_len = (bloom_seq_len) ? 128 : MAX_OUT_TOKES;
allocate_workspace<T>(
input.size(2), input.size(1), max_seq_len, input.size(0), num_layers, head_size);
workspace = (T*)Context::Instance().GetWorkSpace();
}
auto options = at::TensorOptions()
Expand Down Expand Up @@ -699,7 +719,8 @@ template <typename T>
at::Tensor ds_linear_layer(at::Tensor& input,
at::Tensor& weight,
at::Tensor& bias,
unsigned num_layers)
unsigned num_layers,
unsigned head_size)
{
auto input_cont = input.contiguous();
auto options = at::TensorOptions()
Expand All @@ -710,10 +731,12 @@ at::Tensor ds_linear_layer(at::Tensor& input,

int bsz = input.size(0) * input.size(1);
T* workspace = (T*)Context::Instance().GetWorkSpace();
if (!workspace) {
// Reallocate memory if we received a new prompt
if (!workspace || input.size(1) != 1) {
cublasSetStream(Context::Instance().GetCublasHandle(),
Context::Instance().GetCurrentStream());
allocate_workspace<T>(input.size(2), MAX_OUT_TOKES, input.size(0), num_layers);
allocate_workspace<T>(
input.size(2), input.size(1), MAX_OUT_TOKES, input.size(0), num_layers, head_size);
workspace = (T*)Context::Instance().GetWorkSpace();
}
auto output = at::from_blob(workspace, {input.size(0), input.size(1), weight.size(1)}, options);
Expand Down
2 changes: 1 addition & 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 128
#define MAX_OUT_TOKES 1024
#define MAX_WARP_NUM 32
#define WARP_SIZE 32

Expand Down
9 changes: 7 additions & 2 deletions deepspeed/ops/transformer/inference/transformer_inference.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -401,14 +401,17 @@ def compute_attention(qkv_out, input_mask):
def selfAttention_fp():
vector_matmul_func = inference_cuda_module.vector_matmul_fp16 if config.fp16 else \
inference_cuda_module.vector_matmul_fp32

head_size = (attn_qkvw.shape[-1] // 3 // num_attention_heads_per_partition)
if not config.pre_layer_norm:
linear_func = inference_cuda_module.linear_layer_fp16 if config.fp16 else \
inference_cuda_module.linear_layer_fp32

qkv_out = linear_func(input,
attn_qkvw,
attn_qkvb,
DeepSpeedTransformerInference.layer_id)
DeepSpeedTransformerInference.layer_id,
head_size)
else:
qkv_func = inference_cuda_module.qkv_gemm_fp16 if config.fp16 else \
inference_cuda_module.qkv_gemm_fp32
Expand All @@ -421,7 +424,9 @@ def selfAttention_fp():
config.epsilon,
(attn_qkvb is not None),
1 if config.bigscience_bloom else
DeepSpeedTransformerInference.layer_id)
DeepSpeedTransformerInference.layer_id,
config.bigscience_bloom,
head_size)
context_layer, key_layer, value_layer = compute_attention(qkv_out[0] if isinstance(qkv_out, list) else qkv_out, input_mask)
output = vector_matmul_func(context_layer, attn_ow, False)

Expand Down