From e5189330d591d48d34cb08aa6f144141419217f2 Mon Sep 17 00:00:00 2001 From: petermcaughan Date: Fri, 12 May 2023 11:23:07 -0700 Subject: [PATCH] Address OOM Issue when exporting Whisper (#15880) ### Description Remove attention_mask from unnecessary code paths in the whisper export process. ### Motivation and Context Current export script frequently hits OOM error when export whisper-large. Memory profiling shows that this is a result of generating dummy inputs for the `encoder_attention_mask` input for a model pass during exporting - in whisper-large, this dummy tensor can be around 20GB in size. `encoder_attention_mask` is ultimately a dummy input - it's just there to satisfy certain BeamSearch requirements. Thus, we're currently creating a 20GB tensor and passing it to the model, which then discards the input anyways. By removing the code path to generate a dummy encoder_mask tensor, we can reduce the memory requirements to export whisper substantially, while keeping the BeamSearch checks satisfied. --------- Co-authored-by: Peter McAughan --- .../models/whisper/whisper_decoder.py | 31 ++----------------- .../models/whisper/whisper_encoder.py | 9 ++---- .../whisper/whisper_encoder_decoder_init.py | 24 +++++--------- 3 files changed, 13 insertions(+), 51 deletions(-) diff --git a/onnxruntime/python/tools/transformers/models/whisper/whisper_decoder.py b/onnxruntime/python/tools/transformers/models/whisper/whisper_decoder.py index 6247b68c8c..21073a6a2c 100644 --- a/onnxruntime/python/tools/transformers/models/whisper/whisper_decoder.py +++ b/onnxruntime/python/tools/transformers/models/whisper/whisper_decoder.py @@ -15,7 +15,6 @@ import numpy import onnx import torch from transformers import WhisperConfig, file_utils -from whisper_encoder import WhisperEncoderInputs from onnxruntime import InferenceSession @@ -51,7 +50,6 @@ class WhisperDecoderInit(torch.nn.Module): def forward( self, decoder_input_ids: torch.Tensor, - encoder_attention_mask: torch.Tensor, encoder_hidden_states: torch.FloatTensor, ): encoder_outputs = file_utils.ModelOutput() @@ -63,7 +61,6 @@ class WhisperDecoderInit(torch.nn.Module): None, encoder_outputs=encoder_outputs, decoder_input_ids=decoder_input_ids, - head_mask=encoder_attention_mask, past_key_values=None, use_cache=True, return_dict=True, @@ -81,7 +78,7 @@ class WhisperDecoder(torch.nn.Module): self.lm_head = lm_head self.config = config - def forward(self, decoder_input_ids, encoder_attention_mask, *past): + def forward(self, decoder_input_ids, *past): encoder_outputs = file_utils.ModelOutput() dummy_encoder_hidden_states = torch.randn((decoder_input_ids.shape[0], 3000, int(self.config.d_model))) encoder_outputs["last_hidden_state"] = dummy_encoder_hidden_states @@ -96,7 +93,6 @@ class WhisperDecoder(torch.nn.Module): None, encoder_outputs=encoder_outputs, decoder_input_ids=decoder_input_ids, - # decoder_attention_mask=encoder_attention_mask, past_key_values=past_key_values, use_cache=True, return_dict=True, @@ -110,11 +106,9 @@ class WhisperDecoderInputs: def __init__( self, decoder_input_ids, - encoder_attention_mask=None, past_key_values=None, ): self.decoder_input_ids: torch.LongTensor = decoder_input_ids - self.encoder_attention_mask: torch.LongTensor = encoder_attention_mask self.past_key_values: Union[List[torch.FloatTensor], List[torch.HalfTensor], None] = past_key_values @staticmethod @@ -158,14 +152,6 @@ class WhisperDecoderInputs: device=device, ) - encoder_inputs = WhisperEncoderInputs.create_dummy( - batch_size, - encode_sequence_length, - vocab_size, - device, - use_int32_inputs=use_int32_inputs, - ) - float_type = torch.float16 if float16 else torch.float32 if past_decode_sequence_length > 0: @@ -191,16 +177,10 @@ class WhisperDecoderInputs: else: past = None - encoder_attention_mask = torch.zeros( - (encoder_inputs.input_ids.shape[0], 1, encoder_inputs.input_ids.shape[1], encoder_inputs.input_ids.shape[1]) - ).type(torch.int8) - return WhisperDecoderInputs(decoder_input_ids, encoder_attention_mask, past) + return WhisperDecoderInputs(decoder_input_ids, past) def to_list(self) -> List: - input_list = [ - self.decoder_input_ids, - self.encoder_attention_mask, - ] + input_list = [self.decoder_input_ids] if self.past_key_values: input_list.extend(self.past_key_values) return input_list @@ -209,7 +189,6 @@ class WhisperDecoderInputs: past = [p.to(dtype=torch.float32) for p in self.past_key_values] if self.past_key_values else None return WhisperDecoderInputs( self.decoder_input_ids.clone(), - self.encoder_attention_mask.clone(), past, ) @@ -259,7 +238,6 @@ class WhisperDecoderHelper: # Shape of input tensors (sequence_length==1): # input_ids: (batch_size, sequence_length) - # encoder_attention_mask: (batch_size, encode_sequence_length) # past_self_*: (batch_size, num_heads, past_decode_sequence_length, head_size) # past_cross_*: (batch_size, num_heads, encode_sequence_length, head_size) @@ -269,12 +247,10 @@ class WhisperDecoderHelper: # past_cross_*: (batch_size, num_heads, encode_sequence_length, head_size) input_names = ["input_ids"] - input_names.append("encoder_attention_mask") input_names.extend(input_past_names) dynamic_axes = { "input_ids": {0: "batch_size"}, - "encoder_attention_mask": {0: "batch_size", 1: "encode_sequence_length"}, "encoder_hidden_states": {0: "batch_size", 1: "encode_sequence_length / 2"}, "logits": {0: "batch_size", 1: "sequence_length"}, } @@ -335,7 +311,6 @@ class WhisperDecoderHelper: ort_inputs = { "input_ids": numpy.ascontiguousarray(inputs.decoder_input_ids.cpu().numpy()), - "encoder_attention_mask": numpy.ascontiguousarray(inputs.encoder_attention_mask.cpu().numpy()), } if inputs.past_key_values: diff --git a/onnxruntime/python/tools/transformers/models/whisper/whisper_encoder.py b/onnxruntime/python/tools/transformers/models/whisper/whisper_encoder.py index fcf3ce0904..41cadcc2b0 100644 --- a/onnxruntime/python/tools/transformers/models/whisper/whisper_encoder.py +++ b/onnxruntime/python/tools/transformers/models/whisper/whisper_encoder.py @@ -33,12 +33,12 @@ class WhisperEncoder(torch.nn.Module): self.encoder = encoder self.config = config - def forward(self, input_features, attention_mask): + def forward(self, input_features): return self.encoder.model.encoder(input_features)[0] class WhisperEncoderInputs: - def __init__(self, input_features, attention_mask): + def __init__(self, input_features): self.input_ids: torch.LongTensor = input_features # HF Whisper model doesn't support Attention Mask functionality @@ -57,14 +57,12 @@ class WhisperEncoderInputs: Returns: WhisperEncoderInputs: dummy inputs for encoder """ - dtype = torch.float32 input_features = torch.randn( size=(batch_size, feature_size, sequence_length), device=device, ) - attention_mask = torch.ones([batch_size, feature_size, sequence_length], dtype=dtype, device=device) - return WhisperEncoderInputs(input_features, attention_mask) + return WhisperEncoderInputs(input_features) def to_list(self) -> List: if self.input_features is None: @@ -134,7 +132,6 @@ class WhisperEncoderHelper: """Run inference of ONNX model.""" ort_inputs = { "input_ids": numpy.ascontiguousarray(inputs.input_ids.cpu().numpy()), - "attention_mask": numpy.ascontiguousarray(inputs.attention_mask.cpu().numpy()), } return ort_session.run(None, ort_inputs) diff --git a/onnxruntime/python/tools/transformers/models/whisper/whisper_encoder_decoder_init.py b/onnxruntime/python/tools/transformers/models/whisper/whisper_encoder_decoder_init.py index ce8b72d89a..ae9d931dc9 100644 --- a/onnxruntime/python/tools/transformers/models/whisper/whisper_encoder_decoder_init.py +++ b/onnxruntime/python/tools/transformers/models/whisper/whisper_encoder_decoder_init.py @@ -47,21 +47,19 @@ class WhisperEncoderDecoderInit(torch.nn.Module): def forward( self, encoder_input_ids: torch.Tensor, - encoder_attention_mask: torch.Tensor, decoder_input_ids: torch.Tensor = None, ): - encoder_hidden_states: torch.FloatTensor = self.whisper_encoder(encoder_input_ids, None) + encoder_hidden_states: torch.FloatTensor = self.whisper_encoder(encoder_input_ids) # Decoder out: (logits, past_key_values, encoder_hidden_state) - decinit_out = self.whisper_decoder_init(decoder_input_ids, encoder_attention_mask, encoder_hidden_states) + decinit_out = self.whisper_decoder_init(decoder_input_ids, encoder_hidden_states) present_self, present_cross = PastKeyValuesHelper.group_by_self_and_cross(decinit_out[1]) present = present_self + present_cross return decinit_out[0], encoder_hidden_states, present class WhisperEncoderDecoderInitInputs: - def __init__(self, encoder_input_ids, encoder_attention_mask, decoder_input_ids=None): + def __init__(self, encoder_input_ids, decoder_input_ids=None): self.encoder_input_ids: torch.LongTensor = encoder_input_ids - self.encoder_attention_mask: torch.LongTensor = encoder_attention_mask self.decoder_input_ids: torch.LongTensor = decoder_input_ids @staticmethod @@ -81,17 +79,14 @@ class WhisperEncoderDecoderInitInputs: use_int32_inputs=use_int32_inputs, ) decoder_input_ids = None - encoder_attention_mask = torch.zeros( - (encoder_inputs.input_ids.shape[0], 1, encoder_inputs.input_ids.shape[1], encoder_inputs.input_ids.shape[1]) - ).type(torch.int8) if use_decoder_input_ids: dtype = torch.int32 if use_int32_inputs else torch.int64 decoder_input_ids = torch.ones((batch_size, 1), dtype=dtype, device=device) * config.decoder_start_token_id - return WhisperEncoderDecoderInitInputs(encoder_inputs.input_ids, encoder_attention_mask, decoder_input_ids) + return WhisperEncoderDecoderInitInputs(encoder_inputs.input_ids, decoder_input_ids) def to_list(self) -> List: - input_list = [self.encoder_input_ids, self.encoder_attention_mask] + input_list = [self.encoder_input_ids] if self.decoder_input_ids is not None: input_list.append(self.decoder_input_ids) return input_list @@ -129,17 +124,14 @@ class WhisperEncoderDecoderInitHelper: ) input_list = inputs.to_list() - out = model(inputs.encoder_input_ids, inputs.encoder_attention_mask, inputs.decoder_input_ids) + out = model(inputs.encoder_input_ids, inputs.decoder_input_ids) present = out[2] - # pdb.set_trace() present_names = PastKeyValuesHelper.get_input_names(present, encoder=True) - # present_names = PastKeyValuesHelper.get_past_names(model.config.num_layers, present=True) output_names = ["logits", "encoder_hidden_states", *present_names] # Shape of input tensors (sequence_length==1): # input_ids: (batch_size, sequence_length) - # encoder_attention_mask: (batch_size, encode_sequence_length) # encoder_hidden_states: (batch_size, encode_sequence_length, hidden_size) # past_self_*: (batch_size, num_heads, past_decode_sequence_length, head_size) # past_cross_*: (batch_size, num_heads, encode_sequence_length, head_size) @@ -149,7 +141,7 @@ class WhisperEncoderDecoderInitHelper: # past_self_*: (batch_size, num_heads, past_decode_sequence_length + sequence_length, head_size) # past_cross_*: (batch_size, num_heads, encode_sequence_length, head_size) - input_names = ["encoder_input_ids", "encoder_attention_mask"] + input_names = ["encoder_input_ids"] # ONNX exporter might mark dimension like 'Transposepresent_value_self_1_dim_2' in shape inference. # We use a workaround here: first use dim_param "1" for sequence_length, and later change to dim_value. @@ -159,7 +151,6 @@ class WhisperEncoderDecoderInitHelper: head_size = str(model.config.d_model // model.config.encoder_attention_heads) dynamic_axes = { "encoder_input_ids": {0: "batch_size", 1: "encode_sequence_length"}, - "encoder_attention_mask": {0: "batch_size", 1: "encode_sequence_length"}, "encoder_hidden_states": { 0: "batch_size", 1: "encode_sequence_length", @@ -240,7 +231,6 @@ class WhisperEncoderDecoderInitHelper: ort_inputs = { "encoder_input_ids": numpy.ascontiguousarray(inputs.encoder_input_ids.cpu().numpy()), - "encoder_attention_mask": numpy.ascontiguousarray(inputs.encoder_attention_mask.cpu().numpy()), } if inputs.decoder_input_ids is not None: ort_inputs["decoder_input_ids"] = numpy.ascontiguousarray(inputs.decoder_input_ids.cpu().numpy())