mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-05 04:17:53 +00:00
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).
This commit is contained in:
kunal-vaishnavi
GitHub
parent
bd7a0fb377
commit
72a3bde330
7 changed files with 54 additions and 40 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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":
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
Loading…
Reference in a new issue