mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-18 18:52:16 +00:00
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 <petermca@microsoft.com>
This commit is contained in:
petermcaughan
GitHub
parent
000a600080
commit
e5189330d5
3 changed files with 13 additions and 51 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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())
|
||||
|
|
|
|||
Loading…
Reference in a new issue