From 72a3bde3305175ca28d0e5158282c9742da7a2d8 Mon Sep 17 00:00:00 2001 From: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com> Date: Fri, 17 May 2024 23:23:57 -0700 Subject: [PATCH] Add GQA on CPU in LLaMA scripts (#20720) ### Description This PR adds support for adding GroupQueryAttention (GQA) in models that are running on CPU. ### Motivation and Context Previously, the LLaMA scripts supported creating models that have GQA for CUDA only. With the recently added support for [GQA on CPU](https://github.com/microsoft/onnxruntime/pull/20299), models where `num_attention_heads != num_key_value_heads` can now use the GQA op and [run much faster on CPU](https://github.com/microsoft/onnxruntime/pull/20598). --- .../tools/transformers/models/llama/README.md | 16 +++---- .../transformers/models/llama/benchmark.py | 10 ++++- .../models/llama/benchmark_e2e.py | 4 ++ .../models/llama/convert_to_onnx.py | 45 ++++++++++--------- .../transformers/models/llama/llama_inputs.py | 6 +-- .../transformers/models/llama/llama_parity.py | 4 +- .../transformers/models/llama/llama_torch.py | 9 ++-- 7 files changed, 54 insertions(+), 40 deletions(-) diff --git a/onnxruntime/python/tools/transformers/models/llama/README.md b/onnxruntime/python/tools/transformers/models/llama/README.md index 04671e47c0..6fba98c14e 100644 --- a/onnxruntime/python/tools/transformers/models/llama/README.md +++ b/onnxruntime/python/tools/transformers/models/llama/README.md @@ -99,7 +99,7 @@ $ python3 -m models.llama.convert_to_onnx -m meta-llama/Llama-2-7b-hf --input ./ $ python3 -m onnxruntime.transformers.models.llama.convert_to_onnx -m meta-llama/Llama-2-7b-hf --input ./Llama-2-7b-hf --output ./llama2-7b ``` -Export for FP32 CUDA +Export for FP32 CUDA (with MultiHeadAttention) ``` # From source: $ python3 -m models.llama.convert_to_onnx -m meta-llama/Llama-2-7b-hf --output llama2-7b-fp32-gpu --precision fp32 --execution_provider cuda @@ -108,13 +108,13 @@ $ python3 -m models.llama.convert_to_onnx -m meta-llama/Llama-2-7b-hf --output l $ python3 -m onnxruntime.transformers.models.llama.convert_to_onnx -m meta-llama/Llama-2-7b-hf --output llama2-7b-fp32-gpu --precision fp32 --execution_provider cuda ``` -Export for FP32 CPU +Export for FP32 CPU (with GroupQueryAttention) ``` # From source: -$ python3 -m models.llama.convert_to_onnx -m meta-llama/Llama-2-7b-hf --output llama2-7b-fp32-cpu --precision fp32 --execution_provider cpu +$ python3 -m models.llama.convert_to_onnx -m meta-llama/Llama-2-7b-hf --output llama2-7b-fp32-cpu --precision fp32 --execution_provider cpu --use_gqa # From wheel: -$ python3 -m onnxruntime.transformers.models.llama.convert_to_onnx -m meta-llama/Llama-2-7b-hf --output llama2-7b-fp32-cpu --precision fp32 --execution_provider cpu +$ python3 -m onnxruntime.transformers.models.llama.convert_to_onnx -m meta-llama/Llama-2-7b-hf --output llama2-7b-fp32-cpu --precision fp32 --execution_provider cpu --use_gqa ``` Export for FP16 CUDA (with MultiHeadAttention) @@ -135,7 +135,7 @@ $ python3 -m models.llama.convert_to_onnx -m meta-llama/Llama-2-7b-hf --output l $ python3 -m onnxruntime.transformers.models.llama.convert_to_onnx -m meta-llama/Llama-2-7b-hf --output llama2-7b-fp16 --precision fp16 --execution_provider cuda --use_gqa ``` -Note: GroupQueryAttention currently works with the FP16 CUDA and INT4 CUDA models, and it can provide faster inference than MultiHeadAttention, especially for large sequence lengths (e.g. 1024 or larger). For the best performance, you should pre-allocate the KV cache buffers to have size `(batch_size, num_heads, max_sequence_length, head_size)` so that the past KV and present KV caches share the same memory. You also need to bind them with ONNX Runtime's [IO binding](https://onnxruntime.ai/docs/api/python/api_summary.html#iobinding). +Note: GroupQueryAttention can provide faster inference than MultiHeadAttention, especially for large sequence lengths (e.g. 1024 or larger). For the best performance, you should pre-allocate the KV cache buffers to have size `(batch_size, num_heads, max_sequence_length, head_size)` so that the past KV and present KV caches share the same memory. You also need to bind them with ONNX Runtime's [IO binding](https://onnxruntime.ai/docs/api/python/api_summary.html#iobinding). Here is an example of how you can bind directly to `torch.tensor` objects: ``` @@ -212,15 +212,13 @@ $ python3 -m models.llama.convert_to_onnx -m meta-llama/Llama-2-7b-hf --output l $ python3 -m onnxruntime.transformers.models.llama.convert_to_onnx -m meta-llama/Llama-2-7b-hf --output llama2-7b-int4-gpu --precision int4 --quantization_method blockwise --execution_provider cuda --use_gqa ``` -Note: See the FP16 CUDA notes about GroupQueryAttention. The `--use_gqa` flag is optional. - Export for INT4 CPU ``` # From source: -$ python3 -m models.llama.convert_to_onnx -m meta-llama/Llama-2-7b-hf --output llama2-7b-int4-cpu --precision int4 --quantization_method blockwise --execution_provider cpu +$ python3 -m models.llama.convert_to_onnx -m meta-llama/Llama-2-7b-hf --output llama2-7b-int4-cpu --precision int4 --quantization_method blockwise --execution_provider cpu --use_gqa # From wheel: -$ python3 -m onnxruntime.transformers.models.llama.convert_to_onnx -m meta-llama/Llama-2-7b-hf --output llama2-7b-int4-cpu --precision int4 --quantization_method blockwise --execution_provider cpu +$ python3 -m onnxruntime.transformers.models.llama.convert_to_onnx -m meta-llama/Llama-2-7b-hf --output llama2-7b-int4-cpu --precision int4 --quantization_method blockwise --execution_provider cpu --use_gqa ``` Export LLaMA-2 70B sharded model into 4 partitions diff --git a/onnxruntime/python/tools/transformers/models/llama/benchmark.py b/onnxruntime/python/tools/transformers/models/llama/benchmark.py index b9d6b30baa..d05de369b3 100644 --- a/onnxruntime/python/tools/transformers/models/llama/benchmark.py +++ b/onnxruntime/python/tools/transformers/models/llama/benchmark.py @@ -198,6 +198,7 @@ def get_model(args: argparse.Namespace): source, torch_dtype=torch.float16 if args.use_fp16 else torch.float32, use_auth_token=args.auth, + trust_remote_code=args.auth, use_cache=True, cache_dir=args.cache_dir, ).to(args.target_device) @@ -240,6 +241,7 @@ def get_model(args: argparse.Namespace): decoder_file_name=decoder_file_name, decoder_with_past_file_name=decoder_with_past_file_name, use_auth_token=args.auth, + trust_remote_code=args.auth, use_io_binding=True, # Large perf gain even for cpu due to avoiding output copy. use_merged=(True if decoder_file_name == "model.onnx" else None), provider=provider, @@ -658,8 +660,12 @@ def main(): args.rank = rank args.world_size = world_size - tokenizer = AutoTokenizer.from_pretrained(args.model_name, cache_dir=args.cache_dir) - config = AutoConfig.from_pretrained(args.model_name, cache_dir=args.cache_dir) + tokenizer = AutoTokenizer.from_pretrained( + args.model_name, cache_dir=args.cache_dir, use_auth_token=args.auth, trust_remote_code=args.auth + ) + config = AutoConfig.from_pretrained( + args.model_name, cache_dir=args.cache_dir, use_auth_token=args.auth, trust_remote_code=args.auth + ) target_device = f"cuda:{args.rank}" if args.device != "cpu" else args.device use_fp16 = args.precision == "fp16" diff --git a/onnxruntime/python/tools/transformers/models/llama/benchmark_e2e.py b/onnxruntime/python/tools/transformers/models/llama/benchmark_e2e.py index b69bd22974..c7d4dbfa59 100644 --- a/onnxruntime/python/tools/transformers/models/llama/benchmark_e2e.py +++ b/onnxruntime/python/tools/transformers/models/llama/benchmark_e2e.py @@ -69,6 +69,7 @@ def get_model(args: argparse.Namespace): cache_dir=args.cache_dir, torch_dtype=args.torch_dtype, use_auth_token=args.auth, + trust_remote_code=args.auth, use_cache=True, attn_implementation="flash_attention_2", quantization_config=bnb_config, @@ -80,6 +81,7 @@ def get_model(args: argparse.Namespace): cache_dir=args.cache_dir, torch_dtype=args.torch_dtype, use_auth_token=args.auth, + trust_remote_code=args.auth, use_cache=True, attn_implementation=("flash_attention_2" if args.device == "cuda" else "sdpa"), ).to(args.target_device) @@ -338,11 +340,13 @@ def main(): args.hf_dir_path if args.hf_dir_path != "" else args.model_name, cache_dir=args.cache_dir, use_auth_token=args.auth, + trust_remote_code=args.auth, ) tokenizer = AutoTokenizer.from_pretrained( args.hf_dir_path if args.hf_dir_path != "" else args.model_name, cache_dir=args.cache_dir, use_auth_token=args.auth, + trust_remote_code=args.auth, ) model = get_model(args) diff --git a/onnxruntime/python/tools/transformers/models/llama/convert_to_onnx.py b/onnxruntime/python/tools/transformers/models/llama/convert_to_onnx.py index 9990c1d006..8a33544654 100644 --- a/onnxruntime/python/tools/transformers/models/llama/convert_to_onnx.py +++ b/onnxruntime/python/tools/transformers/models/llama/convert_to_onnx.py @@ -263,7 +263,7 @@ def run_torchscript_separate_export( device, batch_size, sequence_length, - use_fp16=args.precision == Precision.FLOAT16, + use_fp16=False, world_size=world_size, ) input_names = [ @@ -343,7 +343,7 @@ def run_torchscript_merged_export( sequence_length, past_sequence_length, max_seq_len=max_sequence_length, - use_fp16=args.precision == Precision.FLOAT16, + use_fp16=False, world_size=world_size, ) input_names = [ @@ -400,7 +400,15 @@ def run_torchscript_merged_export( # Optimize the model as FP32 -def optimize_export(config: AutoConfig, input_path: str, output_path: str, remove_model: bool = True): +def optimize_export( + args: argparse.Namespace, + config: AutoConfig, + input_path: str, + output_path: str, + remove_model: bool = True, + world_size: int = 1, + window_size: int = -1, +): from fusion_options import FusionOptions optimization_options = FusionOptions("gpt2") @@ -414,6 +422,8 @@ def optimize_export(config: AutoConfig, input_path: str, output_path: str, remov optimization_options=optimization_options, only_onnxruntime=False, ) + if args.use_gqa: + model_opt = use_group_query_attention(config, model_opt, world_size, window_size) model_opt.save_model_to_file(output_path, use_external_data_format=True) # Run symbolic shape inference on optimized model to avoid shape errors during runtime @@ -445,9 +455,7 @@ def optimize_export(config: AutoConfig, input_path: str, output_path: str, remov remove_existing_model(input_path) -def convert_to_float16( - args: argparse.Namespace, config: AutoConfig, old_paths: list[str], rank: int = 0, world_size: int = 1 -): +def convert_to_float16(args: argparse.Namespace, old_paths: list[str], rank: int = 0): decoder_model_fp16_path = os.path.join(args.output, f"rank_{rank}_{args.model_name}_decoder_model_fp16.onnx") decoder_with_past_model_fp16_path = os.path.join( args.output, f"rank_{rank}_{args.model_name}_decoder_with_past_model_fp16.onnx" @@ -462,8 +470,6 @@ def convert_to_float16( if os.path.exists(fp32_path): model = OnnxModel(onnx.load_model(fp32_path, load_external_data=True)) model.convert_float_to_float16(keep_io_types=False) - if args.use_gqa: - model = use_group_query_attention(config, model, world_size) model.save_model_to_file(fp16_path, use_external_data_format=True) del model logger.info(f"The ONNX model at {fp32_path} has been converted to float16 and saved at {fp16_path}!") @@ -473,12 +479,12 @@ def convert_to_float16( return new_paths -def use_group_query_attention(config: AutoConfig, fp16_model_opt: OnnxModel, world_size: int = 1, window_size: int = 0): +def use_group_query_attention(config: AutoConfig, model_opt: OnnxModel, world_size: int = 1, window_size: int = -1): # Replace MultiHeadAttention with GroupQueryAttention - fp16_model_opt = replace_mha_with_gqa(fp16_model_opt, "attention_mask", config.num_key_value_heads, world_size) - fp16_model_opt.prune_graph() - fp16_model_opt.update_graph(allow_remove_graph_inputs=True) - return fp16_model_opt + model_opt = replace_mha_with_gqa(model_opt, "attention_mask", config.num_key_value_heads, world_size, window_size) + model_opt.prune_graph() + model_opt.update_graph(allow_remove_graph_inputs=True) + return model_opt def smooth_quant( @@ -577,13 +583,12 @@ def remove_existing_files(output_path: str): def optimize_optimum(config: AutoConfig, args: argparse.Namespace): tmp_file = os.path.join(args.output, args.model_name + ".tmp.onnx") output_file = os.path.join(args.output, args.model_name + ".onnx") - optimize_export(config, args.input, tmp_file, remove_model=False) + window_size = -1 if not hasattr(config, "sliding_window") else config.sliding_window + optimize_export(args, config, args.input, tmp_file, remove_model=False, window_size=window_size) logger.info(f"Model successfully optimized to {tmp_file}") opt_model = OnnxModel(onnx.load_model(tmp_file, load_external_data=True)) if args.precision == Precision.FLOAT16: opt_model.convert_float_to_float16(keep_io_types=False) - window_size = 0 if not hasattr(config, "sliding_window") else config.sliding_window - opt_model = use_group_query_attention(config, opt_model, args.world_size, window_size) logger.info("Model successfully fused and quantized to FP16!") opt_model.save_model_to_file(output_file, use_external_data_format=True) logger.info(f"Output model successfully saved to {output_file}") @@ -829,7 +834,7 @@ def main(): location = args.original_model_name if use_auth_token else args.input if args.optimize_optimum: - config = AutoConfig.from_pretrained(args.original_model_name) + config = AutoConfig.from_pretrained(args.original_model_name, cache_dir=args.cache_dir) optimize_optimum(config, args) return @@ -901,7 +906,7 @@ def main(): logger.info("Optimizing models...") for orig_path, opt_path in zip(old_paths, new_paths): if os.path.exists(orig_path): - optimize_export(l_config, input_path=orig_path, output_path=opt_path) + optimize_export(args, l_config, input_path=orig_path, output_path=opt_path, world_size=world_size) # Re-assign default FP32 model paths as their optimized versions decoder_model_fp32_path = decoder_model_fp32_opt_path @@ -915,7 +920,7 @@ def main(): # Change precision of exported models from FP32 if args.precision == Precision.FLOAT16: - new_paths = convert_to_float16(args, l_config, old_paths, rank, world_size) + new_paths = convert_to_float16(args, old_paths, rank) elif args.precision == Precision.INT8: decoder_model_int8_path = os.path.join( @@ -971,7 +976,7 @@ def main(): elif args.precision == Precision.INT4: if args.execution_provider != "cpu": - old_paths = convert_to_float16(args, l_config, old_paths, rank, world_size) + old_paths = convert_to_float16(args, old_paths, rank) decoder_model_int4_path = os.path.join( args.output, f"rank_{rank}_{args.model_name}_decoder_model_int4.onnx" diff --git a/onnxruntime/python/tools/transformers/models/llama/llama_inputs.py b/onnxruntime/python/tools/transformers/models/llama/llama_inputs.py index 7b3caf0b70..39f0588436 100644 --- a/onnxruntime/python/tools/transformers/models/llama/llama_inputs.py +++ b/onnxruntime/python/tools/transformers/models/llama/llama_inputs.py @@ -255,8 +255,6 @@ def convert_inputs_for_ort( use_buffer_share: bool = False, past_seq_len: int = 0, max_seq_len: int = 2048, - device: str = "", - device_id: int = -1, ): ort_inputs = {} for k, v in pt_inputs.items(): @@ -268,7 +266,7 @@ def convert_inputs_for_ort( ort_inputs[k] = v.detach().cpu().numpy() # Reshape KV caches if using past-present-share-buffer - if use_buffer_share and device != "" and device != "cpu" and device_id > -1: + if use_buffer_share: ort_inputs = enable_past_present_share_buffer(ort_inputs, past_seq_len, max_seq_len) return ort_inputs @@ -420,7 +418,7 @@ def get_initial_inputs_and_outputs( use_buffer_share: bool, engine: str, ): - tokenizer.pad_token = "[PAD]" + tokenizer.pad_token = tokenizer.eos_token encodings_dict = tokenizer.batch_encode_plus(prompt, padding=True) torch_dtype = torch.float16 if use_fp16 else torch.float32 diff --git a/onnxruntime/python/tools/transformers/models/llama/llama_parity.py b/onnxruntime/python/tools/transformers/models/llama/llama_parity.py index c3e754e3df..eab55154b5 100644 --- a/onnxruntime/python/tools/transformers/models/llama/llama_parity.py +++ b/onnxruntime/python/tools/transformers/models/llama/llama_parity.py @@ -20,6 +20,7 @@ from llama_inputs import ( get_merged_sample_with_past_kv_inputs, get_sample_inputs, get_sample_with_past_kv_inputs, + verify_ort_inputs, ) from llama_torch import setup_torch_model from transformers import AutoConfig @@ -112,8 +113,6 @@ def verify_parity( use_buffer_share=args.use_buffer_share, past_seq_len=past_sequence_length, max_seq_len=max_sequence_length, - device=args.execution_provider, - device_id=int(args.rank), ) ep = f"{args.execution_provider.upper()}ExecutionProvider" @@ -124,6 +123,7 @@ def verify_parity( sess_options=ort.SessionOptions(), providers=[ep], ) + inputs = verify_ort_inputs(ort_model, inputs) # Add IO bindings for non-CPU execution providers if args.execution_provider != "cpu": diff --git a/onnxruntime/python/tools/transformers/models/llama/llama_torch.py b/onnxruntime/python/tools/transformers/models/llama/llama_torch.py index d570e2d7ee..643b21ce61 100644 --- a/onnxruntime/python/tools/transformers/models/llama/llama_torch.py +++ b/onnxruntime/python/tools/transformers/models/llama/llama_torch.py @@ -13,7 +13,7 @@ from transformers import AutoConfig, AutoModelForCausalLM logger = logging.getLogger("") -def setup_torch_model(args, location, use_auth_token, torch_dtype=torch.float32, device=None): +def setup_torch_model(args, location, auth, torch_dtype=torch.float32, device=None): world_size = get_size() logger.info(f"world_size: {world_size}") rank = get_rank() @@ -24,12 +24,15 @@ def setup_torch_model(args, location, use_auth_token, torch_dtype=torch.float32, for i in range(world_size): if i == rank % (world_size): - l_config = AutoConfig.from_pretrained(location, use_auth_token=use_auth_token, cache_dir=args.cache_dir) + l_config = AutoConfig.from_pretrained( + location, use_auth_token=auth, cache_dir=args.cache_dir, trust_remote_code=auth + ) l_config.use_cache = True l_config._attn_implementation = "eager" # "eager" uses LlamaAttention for attention layer llama = AutoModelForCausalLM.from_pretrained( location, - use_auth_token=use_auth_token, + use_auth_token=auth, + trust_remote_code=auth, config=l_config, torch_dtype=torch_dtype, cache_dir=args.cache_dir,