About inference time

[Blank-z0]

Hi, great work! I think Zamba2 is a great hybrid SSM model for the whole community.
I'm trying to reproduce the time-to-first-token result presented in your tech report. But I found that I couldn't reproduce the TTFT of Zamba2 reported in the paper. On 1.2B and 2.7B models, I tried setting the input prompt length to 2048, the output token number to 1, and batch size is 1.
However, the TTFT of Zamba2 is higher than attention-based models, such as Phi2-2.7B, Qwen2-1.5B, Qwen2.5-3B.
For example, Zamba2-2.7B vs. Phi2-2.7B 150ms vs. 90ms, Zamba2-1.2B vs. Qwen2-1.5B 94ms vs. 81ms, I use 1 A100 (40G)
For Zamba2, I used mamba-ssm and casual-conv1d to speed up the inference, and for attention-based LLM I used flash attention 2.
The following are my machine envs

PRETTY_NAME="Ubuntu 22.04.2 LTS"
NAME="Ubuntu"
VERSION_ID="22.04"
VERSION="22.04.2 LTS (Jammy Jellyfish)"
VERSION_CODENAME=jammy
ID=ubuntu
ID_LIKE=debian
HOME_URL="https://www.ubuntu.com/"
SUPPORT_URL="https://help.ubuntu.com/"
BUG_REPORT_URL="https://bugs.launchpad.net/ubuntu/"
PRIVACY_POLICY_URL="https://www.ubuntu.com/legal/terms-and-policies/privacy-policy"
UBUNTU_CODENAME=jammy

libnccl2.18.3-1+cuda12.2

python pkgs are

Package                   Version              Editable project location
------------------------- -------------------- ----------------------------------------------------------
absl-py                   1.4.0
accelerate                0.33.0
accelerator               2024.3.8.dev1
aiohttp                   3.8.4
aiosignal                 1.3.1
annotated-types           0.7.0
apex                      0.1
argon2-cffi               21.3.0
argon2-cffi-bindings      21.2.0
asttokens                 2.2.1
astunparse                1.6.3
async-timeout             4.0.2
attrs                     23.1.0
audioread                 3.0.0
backcall                  0.2.0
beautifulsoup4            4.12.2
bitsandbytes              0.43.0
bleach                    6.0.0
blis                      0.7.10
bottle                    0.12.25
cachetools                5.3.1
catalogue                 2.0.9
causal-conv1d             1.3.0.post1
certifi                   2023.7.22
cffi                      1.15.1
charset-normalizer        3.2.0
click                     8.1.5
cloudpickle               2.2.1
cmake                     3.27.1
comm                      0.1.4
confection                0.1.1
contourpy                 1.1.0
cubinlinker               0.3.0+2.g7c3675e
cuda-python               12.1.0rc5+1.g994d8d0
cudf                      23.6.0
cugraph                   23.6.0
cugraph-dgl               23.6.0
cugraph-service-client    23.6.0
cugraph-service-server    23.6.0
cuml                      23.6.0
cupy-cuda12x              12.1.0
cycler                    0.11.0
cymem                     2.0.7
Cython                    3.0.0
dask                      2023.3.2
dask-cuda                 23.6.0
dask-cudf                 23.6.0
datasets                  2.18.0
debugpy                   1.8.7
decorator                 5.1.1
deepspeed                 0.14.4
defusedxml                0.7.1
dill                      0.3.8
distributed               2023.3.2.1
dm-tree                   0.1.8
docker-pycreds            0.4.0
einops                    0.6.1
exceptiongroup            1.1.2
execnet                   2.0.2
executing                 1.2.0
expecttest                0.1.3
fastjsonschema            2.18.0
fastrlock                 0.8.1
fbgemm-gpu                0.6.0
filelock                  3.12.2
flash-attn                2.6.3
fonttools                 4.42.0
frozenlist                1.4.0
fsspec                    2023.6.0
gast                      0.5.4
gin-config                0.5.0
gitdb                     4.0.11
GitPython                 3.1.43
google-auth               2.22.0
google-auth-oauthlib      0.4.6
graphsurgeon              0.4.6
grpcio                    1.56.2
hjson                     3.1.0
huggingface-hub           0.26.5
hypothesis                5.35.1
idna                      3.4
importlib-metadata        6.8.0
iniconfig                 2.0.0
intel-openmp              2021.4.0
iopath                    0.1.10
ipykernel                 6.25.0
ipython                   8.14.0
ipython-genutils          0.2.0
jedi                      0.19.0
Jinja2                    3.1.2
joblib                    1.3.1
json5                     0.9.14
jsonlines                 4.0.0
jsonschema                4.18.6
jsonschema-specifications 2023.7.1
jupyter_client            8.3.0
jupyter_core              5.3.1
jupyter-tensorboard       0.2.0
jupyterlab                2.3.2
jupyterlab-pygments       0.2.2
jupyterlab-server         1.2.0
jupytext                  1.15.0
kiwisolver                1.4.4
langcodes                 3.3.0
librosa                   0.9.2
llvmlite                  0.40.1
locket                    1.0.0
mamba-ssm                 2.1.0
Markdown                  3.4.4
markdown-it-py            3.0.0
MarkupSafe                2.1.3
matplotlib                3.7.2
matplotlib-inline         0.1.6
mdit-py-plugins           0.4.0
mdurl                     0.1.2
mistune                   3.0.1
mkl                       2021.1.1
mkl-devel                 2021.1.1
mkl-include               2021.1.1
mock                      5.1.0
modelscope                1.22.0
mpmath                    1.3.0
msgpack                   1.0.5
multidict                 6.0.4
multiprocess              0.70.16
murmurhash                1.0.9
nbclient                  0.8.0
nbconvert                 7.7.3
nbformat                  5.9.2
nest-asyncio              1.5.7
networkx                  2.6.3
ninja                     1.11.1
notebook                  6.4.10
numba                     0.57.1+1.gc785c8f1f
numpy                     1.22.2
nvidia-cublas-cu12        12.1.3.1
nvidia-cuda-cupti-cu12    12.1.105
nvidia-cuda-nvrtc-cu12    12.1.105
nvidia-cuda-runtime-cu12  12.1.105
nvidia-cudnn-cu12         9.1.0.70
nvidia-cufft-cu12         11.0.2.54
nvidia-curand-cu12        10.3.2.106
nvidia-cusolver-cu12      11.4.5.107
nvidia-cusparse-cu12      12.1.0.106
nvidia-dali-cuda120       1.28.0
nvidia-ml-py              12.560.30
nvidia-nccl-cu12          2.20.5
nvidia-nvjitlink-cu12     12.6.68
nvidia-nvtx-cu12          12.1.105
nvidia-pyindex            1.0.9
nvtx                      0.2.5
oauthlib                  3.2.2
onnx                      1.14.0
opencv                    4.7.0
packaging                 23.1
pandas                    1.5.2
pandocfilters             1.5.0
parso                     0.8.3
partd                     1.4.0
pathy                     0.10.2
peft                      0.11.0
pexpect                   4.8.0
pickleshare               0.7.5
Pillow                    9.2.0
pip                       23.2.1
platformdirs              3.10.0
pluggy                    1.2.0
ply                       3.11
polygraphy                0.47.1
pooch                     1.7.0
portalocker               2.10.1
preshed                   3.0.8
prettytable               3.8.0
prometheus-client         0.17.1
prompt-toolkit            3.0.39
protobuf                  4.21.12
psutil                    5.9.4
ptxcompiler               0.8.1+1.g4a94326
ptyprocess                0.7.0
pure-eval                 0.2.2
py-cpuinfo                9.0.0
pyarrow                   17.0.0
pyarrow-hotfix            0.6
pyasn1                    0.5.0
pyasn1-modules            0.3.0
pybind11                  2.11.1
pycocotools               2.0+nv0.7.3
pycparser                 2.21
pydantic                  2.9.2
pydantic_core             2.23.4
Pygments                  2.16.1
pylibcugraph              23.6.0
pylibcugraphops           23.6.0
pylibraft                 23.6.0
Pympler                   1.1
pynvml                    11.4.1
pyparsing                 3.0.9
pytest                    7.4.0
pytest-flakefinder        1.1.0
pytest-rerunfailures      12.0
pytest-shard              0.1.2
pytest-xdist              3.3.1
python-dateutil           2.8.2
python-hostlist           1.23.0
pytorch-quantization      2.1.2
pytz                      2023.3
PyYAML                    6.0.1
pyzmq                     25.1.0
raft-dask                 23.6.0
referencing               0.30.2
regex                     2023.6.3
requests                  2.31.0
requests-oauthlib         1.3.1
resampy                   0.4.2
rmm                       23.6.0
rpds-py                   0.9.2
rsa                       4.9
safetensors               0.4.5
scikit-learn              1.2.0
scipy                     1.11.1
Send2Trash                1.8.2
sentencepiece             0.1.99
sentry-sdk                2.17.0
setproctitle              1.3.3
setuptools                68.0.0
six                       1.16.0
smart-open                6.3.0
smmap                     5.0.1
sortedcontainers          2.4.0
soundfile                 0.12.1
soupsieve                 2.4.1
spacy                     3.6.0
spacy-legacy              3.0.12
spacy-loggers             1.0.4
sphinx-glpi-theme         0.3
srsly                     2.4.7
stack-data                0.6.2
sympy                     1.12
tabulate                  0.9.0
tbb                       2021.10.0
tblib                     2.0.0
tensorboard               2.9.0
tensorboard-data-server   0.6.1
tensorboard-plugin-wit    1.8.1
tensorrt                  8.6.1
terminado                 0.17.1
thinc                     8.1.10
threadpoolctl             3.2.0
thriftpy2                 0.4.16
tinycss2                  1.2.1
tokenizers                0.19.1
toml                      0.10.2
tomli                     2.0.1
toolz                     0.12.0
torch                     2.4.0
torchdata                 0.7.1
torchvision               0.19.0
tornado                   6.3.2
tqdm                      4.65.0
traitlets                 5.9.0
transformer_engine_cu12   1.12.0
transformers              4.43.0.dev0          /codes/transformers_zamba2
treelite                  3.2.0
treelite-runtime          3.2.0
triton                    3.0.0
typer                     0.9.0
types-dataclasses         0.6.6
typing_extensions         4.12.2
ucx-py                    0.32.0
uff                       0.6.9
urllib3                   1.26.16
waitress                  3.0.2
wandb                     0.18.5
wasabi                    1.1.2
wcwidth                   0.2.6
webencodings              0.5.1
Werkzeug                  2.3.6
wheel                     0.41.1
xdoctest                  1.0.2
xgboost                   1.7.5
xxhash                    3.5.0
yarl                      1.9.2
zict                      3.0.0
zipp                      3.16.2

My mamba-ssm and causal-conv1d use the same version as in setup.py. Below is my code to reproduce TTFT:

# Modified from https://github.com/state-spaces/mamba/blob/main/benchmarks/benchmark_generation_mamba_simple.py
import argparse
import time
import json
import torch
import torch.nn.functional as F
import sys
from einops import rearrange
from transformers import AutoTokenizer, AutoModelForCausalLM

parser = argparse.ArgumentParser(description="Generation benchmarking")
parser.add_argument("--model-name", type=str, default="Qwen2-1.5B")
parser.add_argument("--prompt", type=str, default=None)
parser.add_argument("--promptlen", type=int, default=2048)
parser.add_argument("--genlen", type=int, default=1)
parser.add_argument("--temperature", type=float, default=1.0)
parser.add_argument("--topk", type=int, default=0)
parser.add_argument("--topp", type=float, default=1.0)
parser.add_argument("--minp", type=float, default=0.0)
parser.add_argument("--repetition-penalty", type=float, default=1.0)
parser.add_argument("--batch", type=int, default=1)
# parser.add_argument("--output_to_files", type=bool, default=True)
args = parser.parse_args()

repeats = 3
device = "cuda"
dtype = torch.bfloat16
root = '../../model/'
cache_dir = './'

print(f"Loading model {args.model_name}")

tokenizer = AutoTokenizer.from_pretrained(root+args.model_name)
model = AutoModelForCausalLM.from_pretrained(root+args.model_name, device_map=device, torch_dtype=dtype, attn_implementation="flash_attention_2")
model.eval()
print(f"Number of parameters: {sum(p.numel() for p in model.parameters() if p.requires_grad)}")

torch.random.manual_seed(0)
if args.prompt is None:
    input_ids = torch.randint(1, 1000, (args.batch, args.promptlen), dtype=torch.long, device="cuda")
    attn_mask = torch.ones_like(input_ids, dtype=torch.long, device="cuda")
else:
    tokens = tokenizer(args.prompt, return_tensors="pt")
    input_ids = tokens.input_ids.to(device=device)
    attn_mask = tokens.attention_mask.to(device=device)
max_length = input_ids.shape[1] + args.genlen
print(input_ids.shape)
fn = lambda: model.generate(
    input_ids=input_ids,
    attention_mask=attn_mask,
    max_new_tokens=args.genlen,
    return_dict_in_generate=True,
    pad_token_id=tokenizer.eos_token_id,
    do_sample=False,
    temperature=args.temperature,
    top_k=args.topk,
    top_p=args.topp,
    repetition_penalty=args.repetition_penalty,
)
try:
    out = fn()
except RuntimeError as e:
    if "out of memory" in str(e).lower():
        with open('ouput.txt', mode='a', encoding='utf-8') as f:
            print(f"Prompt length: {len(input_ids[0])}, generation length: {len(out.sequences[0]) - len(input_ids[0])}\n", file=f)
            print(f"{args.model_name} OOM!\n\n", file=f)
        print(f"{args.model_name} OOM!")
        sys.exit(1)

if args.prompt is not None:
    print(tokenizer.batch_decode(out.sequences.tolist()))

torch.cuda.synchronize()
start = time.time()
for _ in range(repeats):
    try:
        fn()
    except RuntimeError as e:
        if "out of memory" in str(e).lower():
            with open(f'{args.model_name}-speed.txt', mode='a', encoding='utf-8') as f:
                print(f"Prompt length: {len(input_ids[0])}, generation length: {len(out.sequences[0]) - len(input_ids[0])}", file=f)
                print(f"{args.model_name} OOM!\n\n", file=f)
            print(f"{args.model_name} OOM!")
            sys.exit(1)
               
torch.cuda.synchronize()
end=time.time()
print(f"Batch size: {args.batch}, Prompt length: {len(input_ids[0])}, generation length: {len(out.sequences[0]) - len(input_ids[0])}")
print(f"{args.model_name} prompt processing + decoding time: {(end - start) / repeats * 1000:.0f}ms\n\n")
with open(f'{args.model_name}-speed.txt', mode='a', encoding='utf-8') as f:
    print(f"Batch size: {args.batch}, Prompt length: {len(input_ids[0])}, generation length: {len(out.sequences[0]) - len(input_ids[0])}", file=f)
    print(f"{args.model_name} prompt processing + decoding time: {(end - start) / repeats * 1000:.0f}ms\n\n", file=f)

I would like to know if my reproduction process is correct and how you calculated the TTFT results shown in the tech report.
Looking forward to your reply, thanks!