Profiling and misc (huggingface#10)

Dockerfile

@@ -11,6 +11,8 @@ RUN useradd -m -u $USER -s /bin/bash $USERNAME \
 
 COPY --chown=$USERNAME ./requirements.txt ./
 COPY --chown=$USERNAME transformers/ ./transformers
-RUN pip install -r requirements.txt
+
+# Stock version of pip doesn't work with editable transformers.
+RUN pip install --upgrade pip --no-cache-dir && pip install -r requirements.txt --no-cache-dir
 
 COPY --chown=$USERNAME src/ ./src

Makefile

@@ -6,69 +6,49 @@ style:
 
 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;
+install:
+	git submodule update --init
+	pip install -r requirements.txt
 
 # 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}
+	python3 src/main.py --hidden_size 2048 --n_head 16 --n_layer 24 --pipeline_class HF_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}
+	python3 src/main.py --hidden_size 2048 --n_head 16 --n_layer 24 --pipeline_class HF_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}
+	python3 src/main.py --hidden_size 2048 --n_head 16 --n_layer 24 --pipeline_class HF_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}
+	deepspeed --num_gpus 1 src/main.py --hidden_size 2048 --n_head 16 --n_layer 24 --pipeline_class DS_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}
+	python3 src/main.py --hidden_size 2048 --n_head 16 --n_layer 24 --pipeline_class HF_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}
+	python3 src/main.py --hidden_size 2048 --n_head 16 --n_layer 24 --pipeline_class HF_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}
+	python3 src/main.py --hidden_size 2048 --n_head 16 --n_layer 24 --pipeline_class HF_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}
+	deepspeed --num_gpus 1 src/main.py --hidden_size 2048 --n_head 16 --n_layer 24 --pipeline_class DS_Pipeline --model_class GPT2 --n_positions 2048 --attention_type 1 --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}
+	python3 src/main.py --hidden_size 2048 --n_head 16 --n_layer 24 --pipeline_class HF_Pipeline --model_class GPT2 --n_positions 2048 --attention_type 2 --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}
+	python3 src/main.py --hidden_size 2048 --n_head 16 --n_layer 24 --pipeline_class HF_Pipeline --model_class GPT2 --n_positions 2048 --attention_type 2 --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}
+	python3 src/main.py --hidden_size 2048 --n_head 16 --n_layer 24 --pipeline_class HF_Pipeline --model_class GPT2 --n_positions 2048 --attention_type 2 --dtype int8 --batch_size ${batch_size}
 
 # Input length experiments
 hf-1b-GPT2-mqa1-int8-input-length:
-	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} --max_input_length ${max_input_length}
+	python3 src/main.py --hidden_size 2048 --n_head 16 --n_layer 24 --pipeline_class HF_Pipeline --model_class GPT2 --n_positions 2048 --attention_type 2 --dtype int8 --batch_size ${batch_size} --max_input_length ${max_input_length}
 
 hf-1b-GPT2-mha-int8-input-length:
-	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} --max_input_length ${max_input_length}
+	python3 src/main.py --hidden_size 2048 --n_head 16 --n_layer 24 --pipeline_class HF_Pipeline --model_class GPT2 --n_positions 2048 --attention_type 1 --dtype int8 --batch_size ${batch_size} --max_input_length ${max_input_length}

requirements.txt

@@ -1,7 +1,7 @@
 accelerate==0.15.0
 bitsandbytes
 deepspeed==0.7.7
-./transformers
+-e ./transformers
 
 # TODO: Dev only
 isort>=5.5.4

src/main.py

@@ -1,19 +1,23 @@
-import pipelines
-from utils import benchmark_end_to_end, get_arg_parser, get_args, get_dummy_batch
+from typing import List, Optional
 
+from src.pipelines import get_pipeline_class
+from src.utils.arguments import parse_args
+from src.utils.benchmark import benchmark_end_to_end
+from src.utils.input import get_dummy_batch
+from src.utils.logging import configure_logging
 
-def main() -> None:
-    # deepspeed.init_distributed("nccl")
 
-    args = get_args(get_arg_parser())
+def main(argv: Optional[List[str]] = None) -> None:
+    args = parse_args(argv=argv)
 
     inputs = get_dummy_batch(args.batch_size, args.max_input_length)
 
-    generate_kwargs = dict(max_new_tokens=args.max_new_tokens, do_sample=False)
+    generate_kwargs = {"max_new_tokens": args.max_new_tokens, "do_sample": False}
 
-    pipeline_class = getattr(pipelines, args.pipeline_class)
+    pipeline_class = get_pipeline_class(args.pipeline_class)
     benchmark_end_to_end(args, pipeline_class, inputs, generate_kwargs)
 
 
 if __name__ == "__main__":
+    configure_logging()
     main()

src/pipelines/__init__.py

@@ -1,3 +1,11 @@
-from .ds_inference import DS_Inference_Pipeline
-from .hf import HF_CPU_Pipeline, HF_GPU_Pipeline
-from .pipeline import Pipeline
+def get_pipeline_class(name):
+    if name == "HF_Pipeline":
+        from src.pipelines.transformers import HF_Pipeline
+
+        return HF_Pipeline
+    elif name == "DS_Pipeline":
+        from src.pipelines.ds import DS_Pipeline
+
+        return DS_Pipeline
+    else:
+        raise NotImplementedError(f"Unsupported pipeline class: {name}")

src/pipelines/ds.py

@@ -0,0 +1,25 @@
+import os
+from argparse import Namespace
+
+import deepspeed
+import torch
+
+from src.pipelines.pipeline import Pipeline
+from src.utils.arguments import check_unused
+
+
+class DS_Pipeline(Pipeline):
+    def __init__(self, args: Namespace) -> None:
+        check_unused(args, {"device": torch.device("cuda")}, enforce=True)
+        # TODO: Works with other dtypes?
+        check_unused(args, {"dtype": torch.float16})
+        super().__init__(args)
+
+        self.model = deepspeed.init_inference(
+            self.model,
+            mp_size=int(os.getenv("WORLD_SIZE", "1")),
+            # base_dir="./",
+            dtype=args.dtype,
+            replace_with_kernel_inject=args.inject_kernel,
+            enable_cuda_graph=args.cuda_graph,
+        )