Add MQA

Makefile

@@ -0,0 +1,61 @@
+batch_size := 1
+
+install-mqa-transformers:
+git clone https://github.com/bigcode-project/transformers.git; \
+cd transformers; \
+git checkout mayank/multi_query; \
+pip install .; \
+cd ..; \
+rm -rf transformers;
+
+# BLOOM AliBi
+hf-1b-bloom-fp32:
+python src/main.py --hidden_size 2048 --n_head 16 --n_layer 24 --pipeline_class HF_GPU_Pipeline --model_class BLOOM --dtype float32 --batch_size ${batch_size}
+
+hf-1b-bloom-bf16:
+python src/main.py --hidden_size 2048 --n_head 16 --n_layer 24 --pipeline_class HF_GPU_Pipeline --model_class BLOOM --dtype bfloat16 --batch_size ${batch_size}
+
+hf-1b-bloom-int8:
+python src/main.py --hidden_size 2048 --n_head 16 --n_layer 24 --pipeline_class HF_GPU_Pipeline --model_class BLOOM --dtype int8 --batch_size ${batch_size}
+
+ds-inference-1b-bloom-fp16:
+deepspeed --num_gpus 1 src/main.py --hidden_size 2048 --n_head 16 --n_layer 24 --pipeline_class DS_Inference_Pipeline --model_class BLOOM --batch_size ${batch_size}
+
+# GPT2 MHA
+hf-1b-GPT2-mha-fp32:
+python src/main.py --hidden_size 2048 --n_head 16 --n_layer 24 --pipeline_class HF_GPU_Pipeline --model_class GPT2 --n_positions 2048 --attention_type 1 --dtype float32 --batch_size ${batch_size}
+
+hf-1b-GPT2-mha-bf16:
+python src/main.py --hidden_size 2048 --n_head 16 --n_layer 24 --pipeline_class HF_GPU_Pipeline --model_class GPT2 --n_positions 2048 --attention_type 1 --dtype bfloat16 --batch_size ${batch_size}
+
+hf-1b-GPT2-mha-int8:
+python src/main.py --hidden_size 2048 --n_head 16 --n_layer 24 --pipeline_class HF_GPU_Pipeline --model_class GPT2 --n_positions 2048 --attention_type 1 --dtype int8 --batch_size ${batch_size}
+
+ds-inference-1b-GPT2-mha-fp16:
+deepspeed --num_gpus 1 src/main.py --hidden_size 2048 --n_head 16 --n_layer 24 --pipeline_class DS_Inference_Pipeline --model_class GPT2 --n_positions 2048 --attention_type 1 --batch_size ${batch_size}
+
+# GPT2 MQA
+hf-1b-GPT2-mqa-fp32:
+python src/main.py --hidden_size 2048 --n_head 16 --n_layer 24 --pipeline_class HF_GPU_Pipeline --model_class GPT2 --n_positions 2048 --attention_type 2 --dtype float32 --batch_size ${batch_size}
+
+hf-1b-GPT2-mqa-bf16:
+python src/main.py --hidden_size 2048 --n_head 16 --n_layer 24 --pipeline_class HF_GPU_Pipeline --model_class GPT2 --n_positions 2048 --attention_type 2 --dtype bfloat16 --batch_size ${batch_size}
+
+hf-1b-GPT2-mqa-int8:
+python src/main.py --hidden_size 2048 --n_head 16 --n_layer 24 --pipeline_class HF_GPU_Pipeline --model_class GPT2 --n_positions 2048 --attention_type 2 --dtype int8 --batch_size ${batch_size}
+
+ds-inference-1b-GPT2-mqa-fp16:
+deepspeed --num_gpus 1 src/main.py --hidden_size 2048 --n_head 16 --n_layer 24 --pipeline_class DS_Inference_Pipeline --model_class GPT2 --n_positions 2048 --attention_type 2 --batch_size ${batch_size}
+
+# GPT2 MQA1
+hf-1b-GPT2-mqa1-fp32:
+python src/main.py --hidden_size 2048 --n_head 16 --n_layer 24 --pipeline_class HF_GPU_Pipeline --model_class GPT2 --n_positions 2048 --attention_type 3 --dtype float32 --batch_size ${batch_size}
+
+hf-1b-GPT2-mqa1-bf16:
+python src/main.py --hidden_size 2048 --n_head 16 --n_layer 24 --pipeline_class HF_GPU_Pipeline --model_class GPT2 --n_positions 2048 --attention_type 3 --dtype bfloat16 --batch_size ${batch_size}
+
+hf-1b-GPT2-mqa1-int8:
+python src/main.py --hidden_size 2048 --n_head 16 --n_layer 24 --pipeline_class HF_GPU_Pipeline --model_class GPT2 --n_positions 2048 --attention_type 3 --dtype int8 --batch_size ${batch_size}
+
+ds-inference-1b-GPT2-mqa1-fp16:
+deepspeed --num_gpus 1 src/main.py --hidden_size 2048 --n_head 16 --n_layer 24 --pipeline_class DS_Inference_Pipeline --model_class GPT2 --n_positions 2048 --attention_type 3 --batch_size ${batch_size}

README.md

@@ -1 +1,61 @@
-# bigcode-inference-benchmark
+# bigcode-inference-benchmark
+A100 80GB
+
+BLOOM\
+```python
+hidden_size = 2048
+n_head = 16
+n_layer = 24
+```
+
+Throughput (tokens/sec | msec/token)
+| batch_size | HF (fp32) | HF (bf16) | HF (int8) | DS-inference (fp16) |
+|:----------:|:---------------:|:---------------:|:---------------:|:-------------------:|
+| 1 | 51.59 \| 19.38 | 47.46 \| 21.07 | 16.53 \| 60.49 | 61.61 \| 16.23 |
+| 2 | 103.92 \| 9.62 | 96.88 \| 10.32 | 33.79 \| 29.60 | 121.55 \| 8.23 |
+| 4 | 211.96 \| 4.72 | 193.72 \| 5.16 | 67.38 \| 14.84 | 240.06 \| 4.17 |
+| 8 | 411.79 \| 2.43 | 370.67 \| 2.70 | 134.34 \| 7.44 | 492.42 \| 2.03 |
+| 16 | 804.55 \| 1.24 | 781.29 \| 1.28 | 275.69 \| 3.63 | 970.59 \| 1.03 |
+| 32 | 1574.68 \| 0.64 | 1539.19 \| 0.65 | 537.14 \| 1.86 | 1999.04 \| 0.50 |
+| 64 | 2712.46 \| 0.37 | 3038.01 \| 0.33 | 1070.50 \| 0.93 | 3971.09 \| 0.25 |
+| 128 | 2974.36 \| 0.34 | 5795.97 \| 0.17 | 2055.34 \| 0.49 | 7514.59 \| 0.13 |
+| 256 | 3695.44 \| 0.27 | 8216.27 \| 0.12 | 3523.77 \| 0.28 | 10226.50 \| 0.10 |
+| 384 | 3591.13 \| 0.28 | 9328.18 \| 0.11 | 4585.33 \| 0.22 | 11094.27 \| 0.09 |
+| 512 | 3708.54 \| 0.27 | 9446.34 \| 0.11 | 5416.48 \| 0.18 | 11390.85 \| 0.09 |
+| 640 | 3859.43 \| 0.26 | 9572.53 \| 0.10 | 6113.65 \| 0.16 | 11625.71 \| 0.09 |
+| 768 | 3804.82 \| 0.26 | 9464.75 \| 0.11 | 6582.52 \| 0.15 | 11814.31 \| 0.08 |
+| 896 | 3652.42 \| 0.27 | 9482.11 \| 0.11 | 7111.08 \| 0.14 | 11744.38 \| 0.09 |
+| 1024 | oom | 9710.46 \| 0.10 | 7486.36 \| 0.13 | 11534.95 \| 0.09 |
+| 1152 | oom | 9712.39 \| 0.10 | 7544.99 \| 0.13 | oom |
+| 1280 | oom | 9667.19 \| 0.10 | 7858.91 \| 0.13 | oom |
+| 1408 | oom | 9771.91 \| 0.10 | 8116.30 \| 0.12 | oom |
+| 1536 | oom | 9744.56 \| 0.10 | 8201.28 \| 0.12 | oom |
+| 1664 | oom | 9719.82 \| 0.10 | 8227.56 \| 0.12 | oom |
+| 1792 | oom | 9690.61 \| 0.10 | 8344.36 \| 0.12 | oom |
+| 1920 | oom | oom | oom | oom |
+
+Latency (sec)
+| batch_size | HF (fp32) | HF (bf16) | HF (int8) | DS-inference (fp16) |
+|:----------:|:---------:|:---------:|:---------:|:-------------------:|
+| 1 | 1.94 | 2.11 | 6.05 | 1.62 |
+| 2 | 1.92 | 2.06 | 5.92 | 1.65 |
+| 4 | 1.89 | 2.06 | 5.94 | 1.67 |
+| 8 | 1.94 | 2.16 | 5.96 | 1.62 |
+| 16 | 1.99 | 2.05 | 5.80 | 1.65 |
+| 32 | 2.03 | 2.08 | 5.96 | 1.60 |
+| 64 | 2.36 | 2.11 | 5.98 | 1.61 |
+| 128 | 4.30 | 2.21 | 6.23 | 1.70 |
+| 256 | 6.93 | 3.12 | 7.26 | 2.50 |
+| 384 | 10.69 | 4.12 | 8.37 | 3.46 |
+| 512 | 14.82 | 5.42 | 9.45 | 4.49 |
+| 640 | 19.85 | 6.69 | 10.47 | 5.51 |
+| 768 | 20.18 | 8.11 | 11.67 | 6.50 |
+| 896 | 24.53 | 9.45 | 12.60 | 7.63 |
+| 1024 | oom | 10.55 | 13.68 | 8.88 |
+| 1152 | oom | 11.86 | 15.27 | oom |
+| 1280 | oom | 13.24 | 16.29 | oom |
+| 1408 | oom | 14.41 | 17.35 | oom |
+| 1536 | oom | 15.76 | 18.73 | oom |
+| 1664 | oom | 17.12 | 20.22 | oom |
+| 1792 | oom | 18.49 | 21.48 | oom |
+| 1920 | oom | oom | oom | oom |

run.sh

@@ -0,0 +1,14 @@
+export CUDA_VISIBLE_DEVICES=0
+
+rm -rf ./tmp
+
+for bs in {1,2,4,8,16,32,64}
+do
+ make $1 batch_size=$bs
+done
+
+for i in {1..20}
+do
+ bs=$(($i*128))
+ make $1 batch_size=$bs
+done

src/pipelines/ds_inference.py

@@ -3,7 +3,6 @@
 
 import deepspeed
 import torch
-from transformers import BloomForCausalLM
 
 from .pipeline import Pipeline
 
@@ -16,7 +15,7 @@ def __init__(self, args: Namespace) -> None:
 
  # with deepspeed.OnDevice(dtype=torch.bfloat16, device="meta"):
  # model = BloomForCausalLM._from_config(config, torch_dtype=torch.bfloat16)
- self.model = BloomForCausalLM._from_config(self.config, torch_dtype=torch.bfloat16)
+ self.model = self.model_class.from_pretrained("tmp", torch_dtype=torch.bfloat16)
  self.model.eval()
 
  # checkpoints_json = os.path.join(args.model_name, "checkpoints.json")

src/pipelines/hf.py

@@ -1,7 +1,6 @@
 from argparse import Namespace
 
 import torch
-from transformers import BloomForCausalLM
 
 from .pipeline import Pipeline
 
@@ -11,13 +10,15 @@ def __init__(self, args: Namespace, device: str = "cpu") -> None:
  super().__init__(args)
 
  model_kwargs = {}
+
  if args.dtype == torch.int8:
  model_kwargs["load_in_8bit"] = True
+ model_kwargs["device_map"] = "auto"
  else:
  model_kwargs["torch_dtype"] = args.dtype
 
  self.input_device = device
- self.model = BloomForCausalLM._from_config(self.config, **model_kwargs).to(self.input_device)
+ self.model = self.model_class.from_pretrained("tmp", **model_kwargs).to(self.input_device)
  self.model.eval()
 
 

src/pipelines/pipeline.py

@@ -1,46 +1,14 @@
+import os
 from argparse import Namespace
-from typing import List, Tuple
+from typing import List, Tuple, Union
 
 import torch
-from transformers import AutoTokenizer, BloomConfig
+from transformers import AutoTokenizer, BloomConfig, BloomForCausalLM, GPT2Config, GPT2LMHeadModel
 
 
 class Pipeline:
  def __init__(self, args: Namespace) -> None:
- self.config = BloomConfig.from_dict(
- {
- "apply_residual_connection_post_layernorm": False,
- "architectures": ["BloomModel"],
- "attention_dropout": 0.0,
- "attention_softmax_in_fp32": True,
- "bias_dropout_fusion": True,
- "bos_token_id": 1,
- "eos_token_id": 2,
- "hidden_dropout": 0.0,
- "hidden_size": args.hidden_size,
- "initializer_range": 0.02,
- "layer_norm_epsilon": 1e-05,
- "masked_softmax_fusion": True,
- "model_type": "bloom",
- "n_head": args.n_head,
- "n_inner": None,
- "n_layer": args.n_layer,
- "offset_alibi": 100,
- "pad_token_id": 3,
- "pretraining_tp": 1,
- "skip_bias_add": True,
- "skip_bias_add_qkv": False,
- "slow_but_exact": False,
- "transformers_version": "4.22.2",
- "unk_token_id": 0,
- "use_cache": True,
- "vocab_size": 250880,
- }
- )
-
- # hardcoded for now to bigscience/bloom
- self.tokenizer = AutoTokenizer.from_pretrained("bigscience/bloom")
-
+ self.config, self.tokenizer, self.model_class = get_config_tokenizer_model_class(args)
  self.model = None
  self.input_device = None
 
@@ -69,3 +37,51 @@ def get_num_parameters(self) -> int:
  for i in self.model.parameters():
  param_count += i.numel()
  return param_count
+
+
+def get_config_tokenizer_model_class(args: Namespace) -> Union[BloomConfig, GPT2Config]:
+ tokenizer = AutoTokenizer.from_pretrained("gpt2")
+ tokenizer.add_special_tokens({"pad_token": "[PAD]"})
+
+ if args.model_class.lower() == "bloom":
+ config = BloomConfig(
+ attention_softmax_in_fp32=True,
+ hidden_size=args.hidden_size,
+ n_head=args.n_head,
+ n_layer=args.n_layer,
+ vocab_size=len(tokenizer),
+ bos_token_id=tokenizer.bos_token_id,
+ eos_token_id=tokenizer.eos_token_id,
+ use_cache=True,
+ )
+ model_class = BloomForCausalLM
+ elif args.model_class.lower() == "gpt2":
+ config = GPT2Config(
+ n_embd=args.hidden_size,
+ n_head=args.n_head,
+ n_layer=args.n_layer,
+ n_positions=args.n_positions,
+ bos_token_id=tokenizer.bos_token_id,
+ eos_token_id=tokenizer.eos_token_id,
+ attention_type=get_attention_type(args.attention_type),
+ print_details=False,
+ vocab_size=len(tokenizer),
+ use_cache=True,
+ )
+ model_class = GPT2LMHeadModel
+
+ if not os.path.exists("tmp"):
+ model_class._from_config(config).save_pretrained("tmp")
+
+ return config, tokenizer, model_class
+
+
+def get_attention_type(attention_type: int):
+ from transformers.models.gpt2.modeling_gpt2 import AttentionType
+
+ if attention_type == 1:
+ return AttentionType.MULTI_HEAD
+ elif attention_type == 2:
+ return AttentionType.MULTI_QUERY
+ elif attention_type == 3:
+ return AttentionType.MULTI_QUERY_1

src/utils/arguments.py

@@ -6,14 +6,18 @@
 def get_arg_parser() -> ArgumentParser:
  parser = ArgumentParser()
  parser.add_argument("--pipeline_class", default="HF_GPU_Pipeline", type=str)
+ parser.add_argument("--model_class", default="GPT2", type=str)
  parser.add_argument("--batch_size", default=1, type=int)
  parser.add_argument("--dtype", default="bfloat16", type=str)
  parser.add_argument("--max_new_tokens", default=100, type=int)
  parser.add_argument("--local_rank", type=int)
  parser.add_argument("--hidden_size", type=int)
+ parser.add_argument("--attention_type", type=int)
+ parser.add_argument("--n_positions", type=int)
  parser.add_argument("--n_head", type=int)
  parser.add_argument("--n_layer", type=int)
  parser.add_argument("--benchmark_cycles", type=int, default=5)
+ parser.add_argument("--clear_every_run", action="store_true")
  return parser
 
 

src/utils/benchmark.py

@@ -53,9 +53,10 @@ def benchmark_end_to_end(args: Namespace, pipeline_class: Pipeline, text: List[s
  if args.benchmark_cycles > 0:
  print_rank_n(f"*** Running benchmark")
 
- torch.cuda.empty_cache()
- gc.collect()
- torch.cuda.synchronize()
+ if args.clear_every_run:
+ torch.cuda.empty_cache()
+ gc.collect()
+ torch.cuda.synchronize()
 
  # benchmark
  total_new_tokens_generated, benchmark_time = run_and_log_time(