Ds-inference Int8 support through ZeroQuant technology

[RezaYazdaniAminabadi]

This PR adds the Int8 support using the ZeroQuant technology introduced here . Note that the kernels added in this PR is just doing a dequantization of the int8-weight matrices and the GeMM operation is still done in the FP16 format using FP16 tensor-core ops. However, the kernels mentioned in this paper can directly operate on the int8 data, by using the int8 tensor-cores if available (such as on A100), which can improve the GeMM throughput by 2x. There will be another PR which uses these kernels to give a much faster inference through MII system.

Here is some performance evaluation:

FP16 on 8 A100-80G:

bsz	latency	tput	tput-per-GPU	Tflops/GPU
1	60.8	16.5	2.1	0.7
8	58.7	136.2	17.0	6.0
32	64.6	495.5	61.9	21.8
64	68.4	936.2	117.0	41.2
100	78.4	1274.7	159.3	56.1

INT8 on 4 A100-80G:

bsz	latency	tput	tput-per-GPU	Tflops/GPU
1	160.5	6.2	1.6	0.5
8	165.1	48.5	12.1	4.3
32	178.7	179.1	44.8	15.8
64	190.7	335.6	83.9	29.5
100	213.0	469.4	117.4	41.3

INT8 on 8 A100-80G:

bsz	latency	tput	tput-per-GPU	Tflops/GPU
1	92.9	10.8	1.3	0.5
8	95.5	83.7	10.5	3.7
32	107.5	297.6	37.2	13.1
64	116.5	549.6	68.7	24.2
128	133.4	959.3	119.9	42.2
196	166.7	1175.5	146.9	51.7
210	170.3	1232.8	154.1	54.2
240	177.9	1349.1	168.6	59.4
256	181.8	1408.2	176.0	62.0

cc: @stas00 @jeffra

[sdpmas]

hey @RezaYazdaniAminabadi I have been waiting for this PR for a long time and recently tried out your branch for ZeroQuant (with GPT-J). I found a couple of issues: quantizer here seems to have been used without initializing. I tried using GroupQuantizer and it works but during forward pass, why is it not doing selfAttention_int8 here. I get !!!! kernel execution error. at that point.

[RezaYazdaniAminabadi]

hey @RezaYazdaniAminabadi I have been waiting for this PR for a long time and recently tried out your branch for ZeroQuant (with GPT-J). I found a couple of issues: quantizer here seems to have been used without initializing. I tried using GroupQuantizer and it works but during forward pass, why is it not doing selfAttention_int8 here. I get !!!! kernel execution error. at that point.

Hey @sdpmas
Thanks for letting me know of this issue. I am actually excited to hear someone is using it :) The thing is that we are using FP16 GeMM for now and we have a plan to release the INT8 kernels soon. I will check this and resolve the issue.
Thanks,
Reza

[RezaYazdaniAminabadi]

@sdpmas, btw, can you please paste the whole trace here to help me debug this better? Thanks

[sdpmas]

yea, sure, here's what the trace looks like:

***************** Creating model in RANK (0) with WORLD_SIZE = 1 *****************
[2022-08-19 16:22:31,093] [INFO] [logging.py:68:log_dist] [Rank -1] DeepSpeed info: version=0.7.1+84e0d03b, git-hash=84e0d03b, git-branch=ds-inference/ZeroQuant-Int8
[2022-08-19 16:22:31,094] [INFO] [logging.py:68:log_dist] [Rank -1] quantize_bits = 8 mlp_extra_grouping = False, quantize_groups = 1
Using /home/samip_dahal/.cache/torch_extensions/py37_cu113 as PyTorch extensions root...
Detected CUDA files, patching ldflags
Emitting ninja build file /home/samip_dahal/.cache/torch_extensions/py37_cu113/transformer_inference/build.ninja...
Building extension module transformer_inference...
Allowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N)
ninja: no work to do.
Loading extension module transformer_inference...
Time to load transformer_inference op: 0.22836589813232422 seconds
[2022-08-19 16:22:35,624] [INFO] [logging.py:68:log_dist] [Rank -1] DeepSpeed-Inference config: {'layer_id': 0, 'hidden_size': 2560, 'intermediate_size': 10240, 'heads': 32, 'num_hidden_layers': -1, 'fp16': True, 'pre_layer_norm': True, 'local_rank': -1, 'stochastic_mode': False, 'epsilon': 1e-05, 'mp_size': 1, 'q_int8': True, 'scale_attention': True, 'triangular_masking': True, 'local_attention': False, 'window_size': 1, 'rotary_dim': 64, 'rotate_half': False, 'rotate_every_two': True, 'return_tuple': True, 'mlp_after_attn': False, 'mlp_act_func_type': <ActivationFuncType.GELU: 1>, 'specialized_mode': False, 'training_mp_size': 1, 'bigscience_bloom': False}
Setting `pad_token_id` to `eos_token_id`:2 for open-end generation.
!!!! kernel execution error. (m: 2560, n: 22, k: 10240, error: 13) 
Traceback (most recent call last):
  File "deep/infer.py", line 53, in <module>
    output=model.generate(**inp_tokens,max_length=150,min_length=100)
  File "/home/samip_dahal/.deep/lib/python3.7/site-packages/torch/autograd/grad_mode.py", line 28, in decorate_context
    return func(*args, **kwargs)
  File "/home/samip_dahal/.deep/lib/python3.7/site-packages/transformers/generation_utils.py", line 1303, in generate
    **model_kwargs,
  File "/home/samip_dahal/.deep/lib/python3.7/site-packages/transformers/generation_utils.py", line 1693, in greedy_search
    output_hidden_states=output_hidden_states,
  File "/home/samip_dahal/.deep/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1120, in _call_impl
    result = forward_call(*input, **kwargs)
  File "/home/samip_dahal/.deep/lib/python3.7/site-packages/transformers/models/gptj/modeling_gptj.py", line 832, in forward
    return_dict=return_dict,
  File "/home/samip_dahal/.deep/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/samip_dahal/.deep/lib/python3.7/site-packages/transformers/models/gptj/modeling_gptj.py", line 682, in forward
    output_attentions=output_attentions,
  File "/home/samip_dahal/.deep/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/samip_dahal/deep/DeepSpeed/deepspeed/ops/transformer/inference/transformer_inference.py", line 872, in forward
    output = output.to(input_type)
RuntimeError: CUDA error: misaligned address
CUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.

and my code looks something like this:

model=deepspeed.init_inference(model,
                                 mp_size=world_size,
                                 dtype=torch.int8,
                                 replace_method='auto',
                                 replace_with_kernel_inject=True,
                                 quantization_setting=1)
        
model.cuda().to(f'cuda:{local_rank}')
inp_tokens = tokenizer("#function to determine if a number is prime and do a lot of things related to fibonacci\nimport", return_tensors="pt")
for token in inp_tokens:
    if torch.is_tensor(inp_tokens[token]):
        inp_tokens[token] = inp_tokens[token].to(f'cuda:{local_rank}')
output=model.generate(**inp_tokens,max_length=150,min_length=100)

[sdpmas]

@RezaYazdaniAminabadi any estimates on when you will release int8 kernels? that would be really helpful. I've been trying to speed up generation for GPT-J and ZeroQuant seems to be the way to go!

[mayank31398]

@pai4451 this has been fixed.
Please reinstall deepspeed from master branch
This was a bug in loading checkpoints

[pai4451]

@pai4451 this has been fixed.
Please reinstall deepspeed from master branch
This was a bug in loading checkpoints

@mayank31398 I did install the latest master branch of DeepSpeed (you can see my output of pip freeze). But I initialize DeepSpeed by specifying the keyword save_mp_checkpoint_path for the first time, and for the next time and so on I load the checkpoints with my saved mp checkpoint path, and then I got these repetitive output. Am I doing wrong in these steps? Should I always load from the downloaded checkpoint path from the hub?

[pai4451]

@mayank31398

After upgrade to transformers==4.21.2 the repetitive output seems to be solved for this input (DeepSpeed is a machine learning framework). However, if I change to other input the output of the model is still very repetitive. For example, this is another input I tried

in=He has a
out=He has a lot of money.
He has a lot of money.
He has a lot of money.
He has a lot of money.
He has a lot of money.
He has a lot of money.
He has a lot of money.
…(repeated)

I can mitigate the repetitiveness by increasing the value of repetition_penalty in model.generate(), but I’m still wondering is the repetitiveness from my ZeroQuant BLOOM normal? Did you try out other inputs as well and conclude that such issue has been fixed?

[RezaYazdaniAminabadi]

generate_kwargs = {'min_length': 100, 'max_new_tokens': 100, 'do_sample': False} batch_size = 1
in = DeepSpeed is a machine learning framework out = DeepSpeed is a machine learning frameworkRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTCRTC This is the output for me for int8 model
@RezaYazdaniAminabadi This is what is generated for me. :) Sounds about right for a 176 B model

Hi @mayank31398 can you show me how you run this? Thanks, Reza

Hi @RezaYazdaniAminabadi,

I am also getting the repetitive output from the ZeroQuant Int8 version of bloom.

in=DeepSpeed is a machine learning framework
out=DeeSpeed is a machine learning framework for deep deep deep deep deep deep deep deep deep deep deep…(repeated)

My generate_kwargs are max_new_tokens=100, do_sample=False. I installed the latest DeepSpeed from master branch and the versions are

deepspeed==0.7.3+afdc7287
transformers==4.20.1
accelerate==0.12.0

The code I use is from bloom-ds-inference-tp-sharded.txt. Do you know how to solve it or how it come from?

Hi @pai4451
I unfortunately don't see this issue neither with this input nor others. Can you please tell me how I can reproduce this?
Thanks,
Reza

[RezaYazdaniAminabadi]

I have tested this with "He has a" input text and it seems the repetition is not related to float16 vsInt8 or DeepSpeed vs HF. This is an issue with the greedy mode of text-generation for this model. So, if you want to generate good quality text, you can pass do_sample as True for now.

Here is the text I see when using fp16 version of checkpoint:

in=He has a                                                                                                                                                                                                                                                 
out=He has a very good point.                                                                                                 
I mean, you know, I don't know what I was thinking.                                                                                                                                                                                                         I mean, I just...                                                                                                                                                                                                                                           I just...                                                                                                                                                                                                                                                   
I just...                                                                                                                     
I just...                                                                                                                     
I just...                                                                                                                     
I just...                                                                                                                     
I just...                                                                                                                     
I just...                                                                                                                     
I just...                                                                                                                     
I just...                                                                                                                                                                                                                                                   
I just...                                                                                                                                                                                                                                                   
I just...                                                                                                                                                                                                                                                   
I just...                                                                                                                                                                                                                                                   
I just...                                                                                                                                                                                                                                                   
I just...                                                                                                                     
I just...                                                                                                                                                                                                                                                   
I just...                                                                                                                                                                                                                                                   
I just...                                                                                                                                                                                                                                                   
I just...                                                                                                                                                                                                                                                   
I just...                                                                                                                                                                                                                                                   
I just...                                                                                                                     
I just...                                                                                                                                                                                                                                                   
I just...                                                                                                                                                                                                                                                   
I just...                                                                                                                     
I just...                                                                                                                     
I just.

By passing do_sample as True, you resolve the repetition issue:

in=He has a
out=He has a big head, but...
- But there is something different.
- Yes.
- He can be a good guy.
- Yes.
You just need to find out who he is, what he likes...
So that I find out who he is, and what he likes?
I just got back to town.
I have things to take care of... that haven't been taken of for a while.
You mean your daughter.
Diana...
Don't try to read too much in what I say.
I'm not asking

Here is the INT8 version of the text:

in=He has a
out=He has a bad reputation now as a corrupt cop.I don't know if you know this.You don't know anything about my brother.He killed three of them and a cop, and he's got an attitude, just like you.
I'm on his side.
He wants to end the business.
So now, what about my brother?
Don't worry about that.
We want to know the truth, no matter who it is.
Are you telling me it's my brother?
Of course it's him, who else?
We know that

[pai4451]

Thanks @RezaYazdaniAminabadi for your explanation. I will keep monitoring the output from FP16 and Int8. I guess both of them will return repetitive result in some cases. Thanks a lot!

[xk503775229]

Is there any guide to running inference on compressed models(especially ZeroQuant)?
Any help would be appreciated.

[stas00]

Please see the demo scripts for BLOOM Inference here:
https://github.com/bigscience-workshop/Megatron-DeepSpeed/tree/main/scripts/bloom-inference-scripts#deepspeed-inference
scroll down to --dtype int8

[wanghaoshuang]

@RezaYazdaniAminabadi
Is there will be another PR for fusing the dequantization with GeMM schedule and fusing the token-wise activation quantization with gelu?

[zcrypt0]

When trying to generate int8 shards (from bigscience/bloom) on a 7xA6000 node I get the following error:

RuntimeError: The size of tensor a (6144) must match the size of tensor b (14336) at non-singleton dimension 1

Full Traceback

Traceback (most recent call last):
  File "t.py", line 402, in 
    model = deepspeed.init_inference(model,
  File "/home/ubuntu/venv/lib/python3.8/site-packages/deepspeed/__init__.py", line 305, in init_inference
    engine = InferenceEngine(model,
  File "/home/ubuntu/venv/lib/python3.8/site-packages/deepspeed/inference/engine.py", line 145, in __init__
    self._apply_injection_policy(
  File "/home/ubuntu/venv/lib/python3.8/site-packages/deepspeed/inference/engine.py", line 342, in _apply_injection_policy
    replace_transformer_layer(client_module,
  File "/home/ubuntu/venv/lib/python3.8/site-packages/deepspeed/module_inject/replace_module.py", line 865, in replace_transformer_layer
    load_model_with_checkpoint(
  File "/home/ubuntu/venv/lib/python3.8/site-packages/deepspeed/module_inject/load_checkpoint.py", line 229, in load_model_with_checkpoint
    load_module_recursive(r_module)
  File "/home/ubuntu/venv/lib/python3.8/site-packages/deepspeed/module_inject/load_checkpoint.py", line 224, in load_module_recursive
    load_module_recursive(
  File "/home/ubuntu/venv/lib/python3.8/site-packages/deepspeed/module_inject/load_checkpoint.py", line 224, in load_module_recursive
    load_module_recursive(
  File "/home/ubuntu/venv/lib/python3.8/site-packages/deepspeed/module_inject/load_checkpoint.py", line 222, in load_module_recursive
    layer_policies[child.__class__](child, prefix + name + '.')
  File "/home/ubuntu/venv/lib/python3.8/site-packages/deepspeed/module_inject/load_checkpoint.py", line 142, in load_transformer_layer
    module.attention.attn_qkvw = mp_replace.copy(module.attention.attn_qkvw,
  File "/home/ubuntu/venv/lib/python3.8/site-packages/deepspeed/module_inject/replace_module.py", line 109, in copy
    dst.data.copy_(weight_split.contiguous())
RuntimeError: The size of tensor a (6144) must match the size of tensor b (14336) at non-singleton dimension 1

My code looks like this (infer_dtype==int8):

model = deepspeed.init_inference(model,
                                 mp_size=world_size,
                                 dtype=getattr(torch, infer_dtype),
                                 checkpoint=init_inference_checkpoints_json,
                                 save_mp_checkpoint_path=save_mp_path,
                                 **kwargs,
                                 )

Is it enough to simply change the dtype parameter on init_inference() or is there some special set of parameters to generate quantized params?

EDIT: I am synced to the tip of master

/cc @stas00 @RezaYazdaniAminabadi

[mayank31398]

@zcrypt0 you can't quantize just be specifying int8 as dtype
You can read ZeroQuant for details.
It might be compute heavy, not sure how long it takes.
For BLOOM-176B, Microsoft has released quantized weights

Not sure how these weights are quantized. But ZeroQuant also does layer distillation for quantization (we don't know what data Microsoft used for this). -> is this documentation available @RezaYazdaniAminabadi ?

This is an example to quantize gpt2

[zcrypt0]

Much appreciated @mayank31398

I was hoping it was just magic, but I will read the paper.

When I try to use the quantized weights (microsoft/bloom-deepspeed-inference-int8) on a 7xA6000 node, I get a size mismatch error. I suspect it's because the weights are sharded for 8 gpus. Any insight on this?

The error (note this doesn't occur on an 8xA6000 node):
RuntimeError: The size of tensor a (6144) must match the size of tensor b (4608) at non-singleton dimension 0

Full Traceback

Traceback (most recent call last):
  File "t.py", line 422, in 
    model = deepspeed.init_inference(model,
  File "/home/ubuntu/venv/lib/python3.8/site-packages/deepspeed/__init__.py", line 305, in init_inference
    engine = InferenceEngine(model,
  File "/home/ubuntu/venv/lib/python3.8/site-packages/deepspeed/inference/engine.py", line 145, in __init__
    self._apply_injection_policy(
  File "/home/ubuntu/venv/lib/python3.8/site-packages/deepspeed/inference/engine.py", line 342, in _apply_injection_policy
    replace_transformer_layer(client_module,
  File "/home/ubuntu/venv/lib/python3.8/site-packages/deepspeed/module_inject/replace_module.py", line 895, in replace_transformer_layer
    load_model_with_checkpoint(replaced_module,
  File "/home/ubuntu/venv/lib/python3.8/site-packages/deepspeed/module_inject/load_checkpoint.py", line 229, in load_model_with_checkpoint
    load_module_recursive(r_module)
  File "/home/ubuntu/venv/lib/python3.8/site-packages/deepspeed/module_inject/load_checkpoint.py", line 224, in load_module_recursive
    load_module_recursive(
  File "/home/ubuntu/venv/lib/python3.8/site-packages/deepspeed/module_inject/load_checkpoint.py", line 224, in load_module_recursive
    load_module_recursive(
  File "/home/ubuntu/venv/lib/python3.8/site-packages/deepspeed/module_inject/load_checkpoint.py", line 222, in load_module_recursive
    layer_policies[child.__class__](child, prefix + name + '.')
  File "/home/ubuntu/venv/lib/python3.8/site-packages/deepspeed/module_inject/load_checkpoint.py", line 138, in load_transformer_layer
    load_parameters(child, prefix + n + '.')
  File "/home/ubuntu/venv/lib/python3.8/site-packages/deepspeed/module_inject/load_checkpoint.py", line 128, in load_parameters
    p.data.copy_(bias_split)
RuntimeError: The size of tensor a (6144) must match the size of tensor b (4608) at non-singleton dimension 0

[mayank31398]

I haven't tried on 7 GPUs. I can give it a shot.

[JingfengYang]

I tried to run 8-bit quantized inference of BLOOM-176B on 8 40G A100 GPUs, but encountered the error: "AttributeError: 'GroupQuantizer' object has no attribute 'num_groups' ". I think it's because'GroupQuantizer' doe not initialize 'num_groups' attribute. Could you please help fix it?

[jeffra]

@JingfengYang, I've reproduced this and will update this thread once we have a fix. Sorry for the inconvenience and thanks for reporting this to us as well.

[JingfengYang]

Thanks! Also, when I'm using a prior committed version of this repo, there is not such error but a new error "AttributeError: 'Parameter' object has no attribute 'scale'" occurs. FYI, I'm using this script: https://github.com/huggingface/transformers-bloom-inference/blob/main/bloom-inference-scripts/bloom-ds-inference.py recommended by Huggingface official website to run 8-bit quantized inference of BLOOM-176B on 8 40G A100 GPUs.

[jeffra]

@JingfengYang, I think we've come up with a fix for the num_groups issue. I've pushed a PR #2645 that should fix this. I need to consult w. @RezaYazdaniAminabadi after the winter break to ensure I am not missing anything here, but feel free to give it a try on your side.

[liangxiaoyun]

Is this problem solved? I also encountered the same problem here.

hey @RezaYazdaniAminabadi I have been waiting for this PR for a long time and recently tried out your branch for ZeroQuant (with GPT-J). I found a couple of issues: quantizer here seems to have been used without initializing. I tried using GroupQuantizer and it works but during forward pass, why is it not doing selfAttention_int8 here. I get !!!! kernel execution error. at that point.

[Tracin]

@RezaYazdaniAminabadi After some simple modification of code, I ran my model in INT8 with low cuda memory usage, very appreciated for this. I wonder when will the full ZeroQuant, I mean INT8 calculation of Gemm will be released?

[kiucho]

Is there any guide I can use ZeroQuant? And is next example just about quantizing(ZeroQuant) gpt2? Then how can I inference to see how much latency is improved?
https://github.com/microsoft/DeepSpeedExamples/tree/master/compression/gpt2

[kiucho]

Sorry, I leave this comment cuz there' no specific guide for inferencing ZeroQuant model.

https://github.com/microsoft/DeepSpeedExamples/blob/master/compression/gpt2/bash_script/run_zero_quant.sh
I have quantized(zeroquant) gpt-j 6b with above script.
At the end of run_clm_no_trainer.py (https://github.com/microsoft/DeepSpeedExamples/blob/master/compression/gpt2/run_clm_no_trainer.py)
I appended next code to benchmark the quantized one refering gpt-bench.py(https://github.com/microsoft/DeepSpeedExamples/blob/master/benchmarks/inference/gpt-bench.py)

Setting dtype=torch.float16 works, but setting dtype=torch.int8 gives an error, and it seems that the Int8 kernel is not yet supported.
What I want to know is if there's something wrong with the code I added, and if the int8 kernel is supported but I'm not utilizing it properly.

Thanks.

    import time
    from transformers import pipeline
    from deepspeed.accelerator import get_accelerator

    def print_latency(latency_set, title, warmup=3):
        # trim warmup queries
        latency_set = list(latency_set)
        latency_set = latency_set[warmup:]
        count = len(latency_set)
        if count > 0:
            latency_set.sort()
            n50 = (count - 1) * 0.5 + 1
            n90 = (count - 1) * 0.9 + 1
            n95 = (count - 1) * 0.95 + 1
            n99 = (count - 1) * 0.99 + 1
            n999 = (count - 1) * 0.999 + 1

            avg = sum(latency_set) / count
            p50 = latency_set[int(n50) - 1]
            p90 = latency_set[int(n90) - 1]
            p95 = latency_set[int(n95) - 1]
            p99 = latency_set[int(n99) - 1]
            p999 = latency_set[int(n999) - 1]

            print(f"====== latency stats {title} ======")
            print("\tAvg Latency: {0:8.2f} ms".format(avg * 1000))
            print("\tP50 Latency: {0:8.2f} ms".format(p50 * 1000))
            print("\tP90 Latency: {0:8.2f} ms".format(p90 * 1000))
            print("\tP95 Latency: {0:8.2f} ms".format(p95 * 1000))
            print("\tP99 Latency: {0:8.2f} ms".format(p99 * 1000))
            print("\t999 Latency: {0:8.2f} ms".format(p999 * 1000))

    deepspeed.init_distributed()
    dtype = torch.float16

    pipe = pipeline("text-generation", model=model, framework="pt", tokenizer=tokenizer)

    if True: # using deepspeed
        pipe.model = deepspeed.init_inference(
            pipe.model,
            dtype=dtype,
            tensor_parallel={"tp_size": 1},
            replace_with_kernel_inject=True,
        )
        pipe.model.profile_model_time()
    
    responses = []
    times = []
    mtimes = []
    for i in range(30):
        get_accelerator().synchronize()
        start = time.time()
        r = pipe("DeepSpeed is", do_sample=False, max_new_tokens=50)
        get_accelerator().synchronize()
        end = time.time()
        responses.append(r)
        times.append(end - start)  # / (args.max_tokens - 3))
        if True: # using deepspeed
            mtimes.append(sum(pipe.model.model_times()))
    if args.local_rank == 0:
        print_latency(times, "(e2e) latency")
        if True: # using deepspeed
            print_latency(mtimes, "(model-only) latency")
        print_latency(map(lambda t: t / (50 - 3), times), "(e2e) per token latency")
        print(f"RESPONSE 0:")
        print("-" * 30)
        print(responses[0][0]["generated_text"])
        print("-" * 30)