From 470b42558e181b218139c93a6661619552592134 Mon Sep 17 00:00:00 2001 From: Kyle Sayers Date: Sun, 19 Jan 2025 00:23:42 +0000 Subject: [PATCH 01/21] WIP Signed-off-by: Kyle Sayers --- examples/multimodal_audio/whisper_example.py | 316 ++++++++++++++++++ .../modifiers/utils/pytorch_helpers.py | 4 +- .../transformers/finetune/data/base.py | 2 +- .../transformers/utils/data_collator.py | 7 + 4 files changed, 327 insertions(+), 2 deletions(-) create mode 100644 examples/multimodal_audio/whisper_example.py diff --git a/examples/multimodal_audio/whisper_example.py b/examples/multimodal_audio/whisper_example.py new file mode 100644 index 000000000..ca705fc89 --- /dev/null +++ b/examples/multimodal_audio/whisper_example.py @@ -0,0 +1,316 @@ +import torch +from datasets import load_dataset +from transformers import WhisperForConditionalGeneration, WhisperProcessor + +from llmcompressor.modifiers.quantization import GPTQModifier +from llmcompressor.transformers import oneshot +from llmcompressor.transformers.utils.data_collator import whisper_data_collator + +# Select model and load it. +MODEL_ID = "openai/whisper-tiny" + +model = WhisperForConditionalGeneration.from_pretrained( + MODEL_ID, + device_map="auto", + torch_dtype="auto", +) +model.config.forced_decoder_ids = None +processor = WhisperProcessor.from_pretrained(MODEL_ID) + +# Select calibration dataset. +DATASET_ID = "hf-internal-testing/librispeech_asr_dummy" +DATASET_SPLIT = f"validation[:512]" + +# Select number of samples. 512 samples is a good place to start. +# Increasing the number of samples can improve accuracy. +NUM_CALIBRATION_SAMPLES = 512 +MAX_SEQUENCE_LENGTH = 2048 + +# Load dataset and preprocess. +ds = load_dataset(DATASET_ID, "clean", split=DATASET_SPLIT) + + +def preprocess(example): + return { + "array": example["audio"]["array"], + "sampling_rate": example["audio"]["sampling_rate"], + } + + +ds = ds.map(preprocess, remove_columns=ds.column_names) + + +# Tokenize inputs. +def tokenize(sample): + generation_config = None + return_token_timestamps = None + logprob_threshold = None + return_timestamps = None + language = None + task = None + is_multilingual = None + + input_features = None, + generation_config = None, + logits_processor = None, + stopping_criteria = None, + prefix_allowed_tokens_fn = None, + synced_gpus = False, + return_timestamps = None, + task = None, + language = None, + is_multilingual = None, + prompt_ids = None, + prompt_condition_type = None, # first-segment, all-segments + condition_on_prev_tokens = None, + temperature = None, + compression_ratio_threshold = None, + logprob_threshold = None, + no_speech_threshold = None, + num_segment_frames = None, + attention_mask = None, + time_precision = 0.02, + time_precision_features = 0.01, + return_token_timestamps = None, + return_segments = False, + return_dict_in_generate = None, + + + input_features = processor( + sample["array"], + sampling_rate=sample["sampling_rate"], + ).input_features + + # 1. prepare generation config + generation_config, kwargs = model._prepare_generation_config(generation_config, **kwargs) + + # 2. set global generate variables + input_stride = model.model.encoder.conv1.stride[0] * model.model.encoder.conv2.stride[0] + num_segment_frames = input_stride * model.config.max_source_positions + batch_size, total_input_frames = model._retrieve_total_input_frames( + input_features=input_features, input_stride=input_stride, kwargs=kwargs + ) + is_shortform = total_input_frames <= num_segment_frames + + # 3. Make sure generation config is correctly set + # Make sure the generation config is correctly set depending on whether timestamps are to be returned or not + return_dict_in_generate = model._set_return_outputs( + return_dict_in_generate=return_dict_in_generate, + return_token_timestamps=return_token_timestamps, + logprob_threshold=logprob_threshold, + generation_config=generation_config, + ) + timestamp_begin = model._set_return_timestamps( + return_timestamps=return_timestamps, is_shortform=is_shortform, generation_config=generation_config + ) + model._set_language_and_task( + language=language, task=task, is_multilingual=is_multilingual, generation_config=generation_config + ) + model._set_num_frames( + return_token_timestamps=return_token_timestamps, generation_config=generation_config, kwargs=kwargs + ) + model._set_thresholds_and_condition( + generation_config=generation_config, + logprob_threshold=logprob_threshold, + compression_ratio_threshold=compression_ratio_threshold, + no_speech_threshold=no_speech_threshold, + condition_on_prev_tokens=condition_on_prev_tokens, + ) + model._set_prompt_condition_type( + generation_config=generation_config, + prompt_condition_type=prompt_condition_type, + ) + + # pass self.config for backward compatibility + init_tokens = model._retrieve_init_tokens( + input_features, + batch_size=batch_size, + generation_config=generation_config, + config=model.config, + num_segment_frames=num_segment_frames, + kwargs=kwargs, + ) + # passing `decoder_input_ids` is deprecated - the only exception is for assisted generation + # where the input ids are handled explicitly by the generate method + model._check_decoder_input_ids(kwargs=kwargs) + + # 3. Retrieve logits processors + device = kwargs["encoder_outputs"][0].device if "encoder_outputs" in kwargs else input_features.device + begin_index = init_tokens.shape[1] + logits_processor = model._retrieve_logit_processors( + generation_config=generation_config, + logits_processor=logits_processor, + begin_index=begin_index, # begin index is index of first generated decoder token + num_beams=kwargs.get("num_beams", 1), + device=device, + ) + + # 4 Set and retrieve global generation variables + model._set_condition_on_prev_tokens( + condition_on_prev_tokens=condition_on_prev_tokens, generation_config=generation_config + ) + + temperatures = [temperature] if not isinstance(temperature, (list, tuple)) else temperature + temperature = temperatures[0] + + max_frames, seek = model._retrieve_max_frames_and_seek( + batch_size=batch_size, + attention_mask=attention_mask, + total_input_frames=total_input_frames, + is_shortform=is_shortform, + ) + + # 5 Prepare running variables, list for generation + num_return_sequences = generation_config.num_return_sequences + ( + batch_idx_map, + cur_bsz, + input_features, + seek, + max_frames, + init_tokens, + do_condition_on_prev_tokens, + ) = model._expand_variables_for_generation( + input_features=input_features, + seek=seek, + max_frames=max_frames, + init_tokens=init_tokens, + batch_size=batch_size, + condition_on_prev_tokens=condition_on_prev_tokens, + generation_config=generation_config, + ) + + current_segments = model._prepare_segments( + prompt_ids=prompt_ids, + batch_size=cur_bsz, + generation_config=generation_config, + ) + + # 6 Transcribe audio until we reach the end of all input audios + while (seek < max_frames).any(): + # 6.1 NOTE: When in longform transcription mode and batch size > 1 we need to dynamically reduce the batch size during the loop + # in case one audio finished earlier than another one. Thus, we need to keep a table of "previous-index-2-current-index" in order + # to know which original audio is being decoded + # Set updated index map, duration of previously decoded chunks and number of max frames of current decoding chunk + input_features, cur_bsz, batch_idx_map = model._maybe_reduce_batch( + input_features=input_features, + seek=seek, + max_frames=max_frames, + cur_bsz=cur_bsz, + batch_idx_map=batch_idx_map, + ) + time_offset = ( + seek.to(torch.float32 if device.type == "mps" else torch.float64) * time_precision / input_stride + ) + seek_num_frames = (max_frames - seek).clamp(max=num_segment_frames) + + # 6.2 cut out next 30s segment from input features + segment_input = model._get_input_segment( + input_features=input_features, + seek=seek, + seek_num_frames=seek_num_frames, + num_segment_frames=num_segment_frames, + cur_bsz=cur_bsz, + batch_idx_map=batch_idx_map, + ) + + # 6.3 prepare decoder input ids + suppress_tokens = _get_attr_from_logit_processors( + logits_processor, SuppressTokensLogitsProcessor, "suppress_tokens" + ) + + decoder_input_ids, kwargs = model._prepare_decoder_input_ids( + cur_bsz=cur_bsz, + init_tokens=init_tokens, + current_segments=current_segments, + batch_idx_map=batch_idx_map, + do_condition_on_prev_tokens=do_condition_on_prev_tokens, + prompt_ids=prompt_ids, + generation_config=generation_config, + config=model.config, + device=init_tokens.device, + suppress_tokens=suppress_tokens, + timestamp_begin=timestamp_begin, + kwargs=kwargs, + ) + + # 6.4 set max new tokens or max length + model._set_max_new_tokens_and_length( + config=model.config, + decoder_input_ids=decoder_input_ids, + generation_config=generation_config, + ) + + # 6.5 Set current `begin_index` for all logit processors + if logits_processor is not None: + for proc in logits_processor: + if hasattr(proc, "set_begin_index"): + proc.set_begin_index(decoder_input_ids.shape[-1]) + + # 6.6 Run generate with fallback + ( + seek_sequences, + seek_outputs, + should_skip, + do_condition_on_prev_tokens, + model_output_type, + ) = model.generate_with_fallback( + segment_input=segment_input, + decoder_input_ids=decoder_input_ids, + cur_bsz=cur_bsz, + batch_idx_map=batch_idx_map, + seek=seek, + num_segment_frames=num_segment_frames, + max_frames=max_frames, + temperatures=temperatures, + generation_config=generation_config, + logits_processor=logits_processor, + stopping_criteria=stopping_criteria, + prefix_allowed_tokens_fn=prefix_allowed_tokens_fn, + synced_gpus=synced_gpus, + return_token_timestamps=return_token_timestamps, + do_condition_on_prev_tokens=do_condition_on_prev_tokens, + is_shortform=is_shortform, + batch_size=batch_size, + attention_mask=attention_mask, + kwargs=kwargs, + ) + + return segment_input["input_features"] + + +ds = ds.map(tokenize, remove_columns=ds.column_names) + +# Configure the quantization algorithm to run. +# * quantize the weights to 4 bit with GPTQ with a group size 128 +breakpoint() +sample_input = next(iter(ds)) +output = model(**sample_input) + + +recipe = GPTQModifier(targets="Linear", scheme="W4A16", ignore=["lm_head"]) + +# Apply algorithms. +oneshot( + model=model, + dataset=ds, + recipe=recipe, + max_seq_length=MAX_SEQUENCE_LENGTH, + num_calibration_samples=NUM_CALIBRATION_SAMPLES, + data_collator=whisper_data_collator, +) +breakpoint() + +# Confirm generations of the quantized model look sane. +print("\n\n") +print("========== SAMPLE GENERATION ==============") +sample_input = next(iter(ds)) +output = model.generate(sample_input) +print(processor.batch_decode(output, skip_special_tokens=True)) +#[' Mr. Quilter is the apostle of the middle classes and we are glad to welcome his gospel.'] +print("==========================================\n\n") + +# Save to disk compressed. +SAVE_DIR = MODEL_ID.split("/")[1] + "-W4A16-G128" +model.save_pretrained(SAVE_DIR, save_compressed=True) +processor.save_pretrained(SAVE_DIR) \ No newline at end of file diff --git a/src/llmcompressor/modifiers/utils/pytorch_helpers.py b/src/llmcompressor/modifiers/utils/pytorch_helpers.py index c9869f267..444a0bac2 100644 --- a/src/llmcompressor/modifiers/utils/pytorch_helpers.py +++ b/src/llmcompressor/modifiers/utils/pytorch_helpers.py @@ -41,7 +41,9 @@ def apply_pad_mask_to_batch(batch: Dict[str, torch.Tensor]) -> Dict[str, torch.T :param batch: batch to apply padding to if it exists :return: batch with padding zeroed out in the input_ids """ - batch["input_ids"] = batch["input_ids"] * batch["attention_mask"] + print(batch.keys()) + if "input_ids" in batch and "attention_mask" in batch: + batch["input_ids"] = batch["input_ids"] * batch["attention_mask"] return batch diff --git a/src/llmcompressor/transformers/finetune/data/base.py b/src/llmcompressor/transformers/finetune/data/base.py index 81a3fc95f..8a4c150cd 100644 --- a/src/llmcompressor/transformers/finetune/data/base.py +++ b/src/llmcompressor/transformers/finetune/data/base.py @@ -105,7 +105,7 @@ def __call__(self, add_labels: bool = True) -> DatasetType: dataset = self.rename_columns(dataset) logger.debug(f"Dataset after column renaming: {get_columns(dataset)}") - if "input_ids" not in get_columns(dataset): + if "input_ids" not in get_columns(dataset) and "input_features" not in get_columns(dataset): # tokenize/ process dataset = self.filter_tokenizer_args(dataset) logger.debug(f"Tokenizer args after filtering: {get_columns(dataset)}") diff --git a/src/llmcompressor/transformers/utils/data_collator.py b/src/llmcompressor/transformers/utils/data_collator.py index b2dc7c651..b4e13e60e 100644 --- a/src/llmcompressor/transformers/utils/data_collator.py +++ b/src/llmcompressor/transformers/utils/data_collator.py @@ -46,3 +46,10 @@ def qwen2_vl_data_collator(batch): "pixel_values": torch.tensor(batch[0]["pixel_values"]), "image_grid_thw": torch.tensor(batch[0]["image_grid_thw"]), } + + +def whisper_data_collator(batch): + assert len(batch) == 1 + return { + "input_features": torch.LongTensor(batch[0]["input_features"]), + } From 5276c9f3a0ca7776bacee18cc40f34e7e224adf2 Mon Sep 17 00:00:00 2001 From: Kyle Sayers Date: Sun, 19 Jan 2025 01:12:58 +0000 Subject: [PATCH 02/21] WIP: traceable, sample generation WIP Signed-off-by: Kyle Sayers --- examples/multimodal_audio/whisper_example.py | 277 +++--------------- .../transformers/tracing/__init__.py | 4 + .../transformers/tracing/whisper.py | 151 ++++++++++ .../transformers/utils/data_collator.py | 3 +- 4 files changed, 196 insertions(+), 239 deletions(-) create mode 100644 src/llmcompressor/transformers/tracing/whisper.py diff --git a/examples/multimodal_audio/whisper_example.py b/examples/multimodal_audio/whisper_example.py index ca705fc89..228f7718f 100644 --- a/examples/multimodal_audio/whisper_example.py +++ b/examples/multimodal_audio/whisper_example.py @@ -1,15 +1,16 @@ import torch from datasets import load_dataset -from transformers import WhisperForConditionalGeneration, WhisperProcessor +from transformers import WhisperProcessor from llmcompressor.modifiers.quantization import GPTQModifier from llmcompressor.transformers import oneshot from llmcompressor.transformers.utils.data_collator import whisper_data_collator +from llmcompressor.transformers.tracing import TraceableWhisperForConditionalGeneration # Select model and load it. MODEL_ID = "openai/whisper-tiny" -model = WhisperForConditionalGeneration.from_pretrained( +model = TraceableWhisperForConditionalGeneration.from_pretrained( MODEL_ID, device_map="auto", torch_dtype="auto", @@ -19,11 +20,11 @@ # Select calibration dataset. DATASET_ID = "hf-internal-testing/librispeech_asr_dummy" -DATASET_SPLIT = f"validation[:512]" +DATASET_SPLIT = f"validation[:1]" # Select number of samples. 512 samples is a good place to start. # Increasing the number of samples can improve accuracy. -NUM_CALIBRATION_SAMPLES = 512 +NUM_CALIBRATION_SAMPLES = 1 # 512 MAX_SEQUENCE_LENGTH = 2048 # Load dataset and preprocess. @@ -39,253 +40,52 @@ def preprocess(example): ds = ds.map(preprocess, remove_columns=ds.column_names) +r""" +Returns: -# Tokenize inputs. -def tokenize(sample): - generation_config = None - return_token_timestamps = None - logprob_threshold = None - return_timestamps = None - language = None - task = None - is_multilingual = None +Example: + ```python + >>> import torch + >>> from transformers import AutoFeatureExtractor, WhisperModel + >>> from datasets import load_dataset - input_features = None, - generation_config = None, - logits_processor = None, - stopping_criteria = None, - prefix_allowed_tokens_fn = None, - synced_gpus = False, - return_timestamps = None, - task = None, - language = None, - is_multilingual = None, - prompt_ids = None, - prompt_condition_type = None, # first-segment, all-segments - condition_on_prev_tokens = None, - temperature = None, - compression_ratio_threshold = None, - logprob_threshold = None, - no_speech_threshold = None, - num_segment_frames = None, - attention_mask = None, - time_precision = 0.02, - time_precision_features = 0.01, - return_token_timestamps = None, - return_segments = False, - return_dict_in_generate = None, + >>> model = WhisperModel.from_pretrained("openai/whisper-base") + >>> feature_extractor = AutoFeatureExtractor.from_pretrained("openai/whisper-base") + >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation") + >>> inputs = feature_extractor(ds[0]["audio"]["array"], return_tensors="pt") + >>> input_features = inputs.input_features + >>> decoder_input_ids = torch.tensor([[1, 1]]) * model.config.decoder_start_token_id + >>> last_hidden_state = model(input_features, decoder_input_ids=decoder_input_ids).last_hidden_state + >>> list(last_hidden_state.shape) + [1, 2, 512] + ``` +""" +# Tokenize inputs. +def tokenize(sample): + batch_size = 1 input_features = processor( sample["array"], sampling_rate=sample["sampling_rate"], + return_tensors="pt", ).input_features - # 1. prepare generation config - generation_config, kwargs = model._prepare_generation_config(generation_config, **kwargs) - - # 2. set global generate variables - input_stride = model.model.encoder.conv1.stride[0] * model.model.encoder.conv2.stride[0] - num_segment_frames = input_stride * model.config.max_source_positions - batch_size, total_input_frames = model._retrieve_total_input_frames( - input_features=input_features, input_stride=input_stride, kwargs=kwargs - ) - is_shortform = total_input_frames <= num_segment_frames - - # 3. Make sure generation config is correctly set - # Make sure the generation config is correctly set depending on whether timestamps are to be returned or not - return_dict_in_generate = model._set_return_outputs( - return_dict_in_generate=return_dict_in_generate, - return_token_timestamps=return_token_timestamps, - logprob_threshold=logprob_threshold, - generation_config=generation_config, - ) - timestamp_begin = model._set_return_timestamps( - return_timestamps=return_timestamps, is_shortform=is_shortform, generation_config=generation_config - ) - model._set_language_and_task( - language=language, task=task, is_multilingual=is_multilingual, generation_config=generation_config - ) - model._set_num_frames( - return_token_timestamps=return_token_timestamps, generation_config=generation_config, kwargs=kwargs - ) - model._set_thresholds_and_condition( - generation_config=generation_config, - logprob_threshold=logprob_threshold, - compression_ratio_threshold=compression_ratio_threshold, - no_speech_threshold=no_speech_threshold, - condition_on_prev_tokens=condition_on_prev_tokens, - ) - model._set_prompt_condition_type( - generation_config=generation_config, - prompt_condition_type=prompt_condition_type, - ) - - # pass self.config for backward compatibility - init_tokens = model._retrieve_init_tokens( - input_features, - batch_size=batch_size, - generation_config=generation_config, - config=model.config, - num_segment_frames=num_segment_frames, - kwargs=kwargs, - ) - # passing `decoder_input_ids` is deprecated - the only exception is for assisted generation - # where the input ids are handled explicitly by the generate method - model._check_decoder_input_ids(kwargs=kwargs) - - # 3. Retrieve logits processors - device = kwargs["encoder_outputs"][0].device if "encoder_outputs" in kwargs else input_features.device - begin_index = init_tokens.shape[1] - logits_processor = model._retrieve_logit_processors( - generation_config=generation_config, - logits_processor=logits_processor, - begin_index=begin_index, # begin index is index of first generated decoder token - num_beams=kwargs.get("num_beams", 1), - device=device, - ) - - # 4 Set and retrieve global generation variables - model._set_condition_on_prev_tokens( - condition_on_prev_tokens=condition_on_prev_tokens, generation_config=generation_config - ) - - temperatures = [temperature] if not isinstance(temperature, (list, tuple)) else temperature - temperature = temperatures[0] - - max_frames, seek = model._retrieve_max_frames_and_seek( - batch_size=batch_size, - attention_mask=attention_mask, - total_input_frames=total_input_frames, - is_shortform=is_shortform, - ) - - # 5 Prepare running variables, list for generation - num_return_sequences = generation_config.num_return_sequences - ( - batch_idx_map, - cur_bsz, - input_features, - seek, - max_frames, - init_tokens, - do_condition_on_prev_tokens, - ) = model._expand_variables_for_generation( - input_features=input_features, - seek=seek, - max_frames=max_frames, - init_tokens=init_tokens, - batch_size=batch_size, - condition_on_prev_tokens=condition_on_prev_tokens, - generation_config=generation_config, - ) - - current_segments = model._prepare_segments( - prompt_ids=prompt_ids, - batch_size=cur_bsz, - generation_config=generation_config, - ) + decoder_input_ids = torch.ones((batch_size, 1), dtype=torch.long) * model.config.decoder_start_token_id - # 6 Transcribe audio until we reach the end of all input audios - while (seek < max_frames).any(): - # 6.1 NOTE: When in longform transcription mode and batch size > 1 we need to dynamically reduce the batch size during the loop - # in case one audio finished earlier than another one. Thus, we need to keep a table of "previous-index-2-current-index" in order - # to know which original audio is being decoded - # Set updated index map, duration of previously decoded chunks and number of max frames of current decoding chunk - input_features, cur_bsz, batch_idx_map = model._maybe_reduce_batch( - input_features=input_features, - seek=seek, - max_frames=max_frames, - cur_bsz=cur_bsz, - batch_idx_map=batch_idx_map, - ) - time_offset = ( - seek.to(torch.float32 if device.type == "mps" else torch.float64) * time_precision / input_stride - ) - seek_num_frames = (max_frames - seek).clamp(max=num_segment_frames) - - # 6.2 cut out next 30s segment from input features - segment_input = model._get_input_segment( - input_features=input_features, - seek=seek, - seek_num_frames=seek_num_frames, - num_segment_frames=num_segment_frames, - cur_bsz=cur_bsz, - batch_idx_map=batch_idx_map, - ) - - # 6.3 prepare decoder input ids - suppress_tokens = _get_attr_from_logit_processors( - logits_processor, SuppressTokensLogitsProcessor, "suppress_tokens" - ) - - decoder_input_ids, kwargs = model._prepare_decoder_input_ids( - cur_bsz=cur_bsz, - init_tokens=init_tokens, - current_segments=current_segments, - batch_idx_map=batch_idx_map, - do_condition_on_prev_tokens=do_condition_on_prev_tokens, - prompt_ids=prompt_ids, - generation_config=generation_config, - config=model.config, - device=init_tokens.device, - suppress_tokens=suppress_tokens, - timestamp_begin=timestamp_begin, - kwargs=kwargs, - ) - - # 6.4 set max new tokens or max length - model._set_max_new_tokens_and_length( - config=model.config, - decoder_input_ids=decoder_input_ids, - generation_config=generation_config, - ) - - # 6.5 Set current `begin_index` for all logit processors - if logits_processor is not None: - for proc in logits_processor: - if hasattr(proc, "set_begin_index"): - proc.set_begin_index(decoder_input_ids.shape[-1]) - - # 6.6 Run generate with fallback - ( - seek_sequences, - seek_outputs, - should_skip, - do_condition_on_prev_tokens, - model_output_type, - ) = model.generate_with_fallback( - segment_input=segment_input, - decoder_input_ids=decoder_input_ids, - cur_bsz=cur_bsz, - batch_idx_map=batch_idx_map, - seek=seek, - num_segment_frames=num_segment_frames, - max_frames=max_frames, - temperatures=temperatures, - generation_config=generation_config, - logits_processor=logits_processor, - stopping_criteria=stopping_criteria, - prefix_allowed_tokens_fn=prefix_allowed_tokens_fn, - synced_gpus=synced_gpus, - return_token_timestamps=return_token_timestamps, - do_condition_on_prev_tokens=do_condition_on_prev_tokens, - is_shortform=is_shortform, - batch_size=batch_size, - attention_mask=attention_mask, - kwargs=kwargs, - ) - - return segment_input["input_features"] + return { + "input_features": input_features, + "decoder_input_ids": decoder_input_ids + } ds = ds.map(tokenize, remove_columns=ds.column_names) # Configure the quantization algorithm to run. # * quantize the weights to 4 bit with GPTQ with a group size 128 -breakpoint() -sample_input = next(iter(ds)) -output = model(**sample_input) +#breakpoint() +#sample_input = next(iter(ds)) +#output = model(**sample_input) recipe = GPTQModifier(targets="Linear", scheme="W4A16", ignore=["lm_head"]) @@ -299,13 +99,14 @@ def tokenize(sample): num_calibration_samples=NUM_CALIBRATION_SAMPLES, data_collator=whisper_data_collator, ) -breakpoint() # Confirm generations of the quantized model look sane. print("\n\n") print("========== SAMPLE GENERATION ==============") -sample_input = next(iter(ds)) -output = model.generate(sample_input) +sample_input = whisper_data_collator([next(iter(ds))]).to(model.device) +sample_input = {k: v.to("cuda:0") for k, v in sample_input.items()} +output = model.generate(**sample_input) +breakpoint() print(processor.batch_decode(output, skip_special_tokens=True)) #[' Mr. Quilter is the apostle of the middle classes and we are glad to welcome his gospel.'] print("==========================================\n\n") diff --git a/src/llmcompressor/transformers/tracing/__init__.py b/src/llmcompressor/transformers/tracing/__init__.py index 4baa5864d..b9202d95b 100644 --- a/src/llmcompressor/transformers/tracing/__init__.py +++ b/src/llmcompressor/transformers/tracing/__init__.py @@ -5,9 +5,13 @@ from .mllama import ( MllamaForConditionalGeneration as TraceableMllamaForConditionalGeneration, ) +from .whisper import ( + WhisperForConditionalGeneration as TraceableWhisperForConditionalGeneration +) __all__ = [ "TraceableLlavaForConditionalGeneration", "TraceableMllamaForConditionalGeneration", "TraceableMistralForCausalLM", + "TraceableWhisperForConditionalGeneration", ] diff --git a/src/llmcompressor/transformers/tracing/whisper.py b/src/llmcompressor/transformers/tracing/whisper.py new file mode 100644 index 000000000..b8241f764 --- /dev/null +++ b/src/llmcompressor/transformers/tracing/whisper.py @@ -0,0 +1,151 @@ +# coding=utf-8 +# Copyright 2022 The OpenAI Authors and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# vllm-project: no copyright +"""PyTorch Whisper model.""" + +import torch +from torch import nn + +from transformers import WhisperConfig +from transformers.models.whisper.modeling_whisper import ( + WhisperEncoder, + WhisperDecoder, + WhisperModel, + WhisperForConditionalGeneration, + WhisperForAudioClassification, +) +from transformers.modeling_outputs import BaseModelOutput + + +class WhisperEncoder(WhisperEncoder): + def forward( + self, + input_features, + attention_mask=None, + head_mask=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + ): + + expected_seq_length = self.config.max_source_positions * self.conv1.stride[0] * self.conv2.stride[0] + # TRACING: assume preprocessing is correct + # if input_features.shape[-1] != expected_seq_length: + if False: + raise ValueError( + f"Whisper expects the mel input features to be of length {expected_seq_length}, but found {input_features.shape[-1]}. Make sure to pad the input mel features to {expected_seq_length}." + ) + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + inputs_embeds = nn.functional.gelu(self.conv1(input_features)) + inputs_embeds = nn.functional.gelu(self.conv2(inputs_embeds)) + + inputs_embeds = inputs_embeds.permute(0, 2, 1) + embed_pos = self.embed_positions.weight + + hidden_states = inputs_embeds + embed_pos + hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) + + encoder_states = () if output_hidden_states else None + all_attentions = () if output_attentions else None + + # check if head_mask has a correct number of layers specified if desired + if head_mask is not None: + assert head_mask.size()[0] == ( + len(self.layers) + ), f"The head_mask should be specified for {len(self.layers)} layers, but it is for {head_mask.size()[0]}." + + for idx, encoder_layer in enumerate(self.layers): + if output_hidden_states: + encoder_states = encoder_states + (hidden_states,) + # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description) + to_drop = False + if self.training: + dropout_probability = torch.rand([]) + if dropout_probability < self.layerdrop: # skip the layer + to_drop = True + + if to_drop: + layer_outputs = (None, None) + else: + if self.gradient_checkpointing and self.training: + layer_outputs = self._gradient_checkpointing_func( + encoder_layer.__call__, + hidden_states, + None, + (head_mask[idx] if head_mask is not None else None), + output_attentions, + ) + else: + layer_outputs = encoder_layer( + hidden_states, + None, + layer_head_mask=(head_mask[idx] if head_mask is not None else None), + output_attentions=output_attentions, + ) + + hidden_states = layer_outputs[0] + + if output_attentions: + all_attentions = all_attentions + (layer_outputs[1],) + + hidden_states = self.layer_norm(hidden_states) + if output_hidden_states: + encoder_states = encoder_states + (hidden_states,) + + if not return_dict: + return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None) + return BaseModelOutput( + last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions + ) + + +class WhisperModel(WhisperModel): + def __init__(self, config: WhisperConfig): + super().__init__(config) + + self.encoder = WhisperEncoder(config) + self.decoder = WhisperDecoder(config) + # Initialize weights and apply final processing + self.post_init() + +class WhisperForConditionalGeneration(WhisperForConditionalGeneration): + def __init__(self, config: WhisperConfig): + super().__init__(config) + self.model = WhisperModel(config) + self.proj_out = nn.Linear(config.d_model, config.vocab_size, bias=False) + self.max_target_positions = config.max_target_positions + + # Initialize weights and apply final processing + self.post_init() + + +class WhisperForAudioClassification(WhisperForAudioClassification): + def __init__(self, config): + super().__init__(config) + + self.encoder = WhisperEncoder(config) + num_layers = config.num_hidden_layers + 1 # transformer layers + input embeddings + if config.use_weighted_layer_sum: + self.layer_weights = nn.Parameter(torch.ones(num_layers) / num_layers) + self.projector = nn.Linear(config.hidden_size, config.classifier_proj_size) + self.classifier = nn.Linear(config.classifier_proj_size, config.num_labels) + + # Initialize weights and apply final processing + self.post_init() \ No newline at end of file diff --git a/src/llmcompressor/transformers/utils/data_collator.py b/src/llmcompressor/transformers/utils/data_collator.py index b4e13e60e..111e12b07 100644 --- a/src/llmcompressor/transformers/utils/data_collator.py +++ b/src/llmcompressor/transformers/utils/data_collator.py @@ -51,5 +51,6 @@ def qwen2_vl_data_collator(batch): def whisper_data_collator(batch): assert len(batch) == 1 return { - "input_features": torch.LongTensor(batch[0]["input_features"]), + "input_features": torch.tensor(batch[0]["input_features"]), + "decoder_input_ids": torch.tensor(batch[0]["decoder_input_ids"]), } From 898f86ec4f9cd20fd2271e916c53b8644d2555f3 Mon Sep 17 00:00:00 2001 From: Kyle Sayers Date: Sun, 19 Jan 2025 01:32:54 +0000 Subject: [PATCH 03/21] WIP: working, need to change ds split Signed-off-by: Kyle Sayers --- examples/multimodal_audio/whisper_example.py | 63 +++++++------------ .../transformers/finetune/data/base.py | 4 +- .../transformers/tracing/whisper.py | 1 + 3 files changed, 27 insertions(+), 41 deletions(-) diff --git a/examples/multimodal_audio/whisper_example.py b/examples/multimodal_audio/whisper_example.py index 228f7718f..dd3bf2fb1 100644 --- a/examples/multimodal_audio/whisper_example.py +++ b/examples/multimodal_audio/whisper_example.py @@ -1,11 +1,10 @@ -import torch from datasets import load_dataset from transformers import WhisperProcessor from llmcompressor.modifiers.quantization import GPTQModifier from llmcompressor.transformers import oneshot -from llmcompressor.transformers.utils.data_collator import whisper_data_collator from llmcompressor.transformers.tracing import TraceableWhisperForConditionalGeneration +from llmcompressor.transformers.utils.data_collator import whisper_data_collator # Select model and load it. MODEL_ID = "openai/whisper-tiny" @@ -20,11 +19,11 @@ # Select calibration dataset. DATASET_ID = "hf-internal-testing/librispeech_asr_dummy" -DATASET_SPLIT = f"validation[:1]" +DATASET_SPLIT = "validation[:512]" # Select number of samples. 512 samples is a good place to start. # Increasing the number of samples can improve accuracy. -NUM_CALIBRATION_SAMPLES = 1 # 512 +NUM_CALIBRATION_SAMPLES = 512 MAX_SEQUENCE_LENGTH = 2048 # Load dataset and preprocess. @@ -40,27 +39,6 @@ def preprocess(example): ds = ds.map(preprocess, remove_columns=ds.column_names) -r""" -Returns: - -Example: - ```python - >>> import torch - >>> from transformers import AutoFeatureExtractor, WhisperModel - >>> from datasets import load_dataset - - >>> model = WhisperModel.from_pretrained("openai/whisper-base") - >>> feature_extractor = AutoFeatureExtractor.from_pretrained("openai/whisper-base") - >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation") - >>> inputs = feature_extractor(ds[0]["audio"]["array"], return_tensors="pt") - >>> input_features = inputs.input_features - >>> decoder_input_ids = torch.tensor([[1, 1]]) * model.config.decoder_start_token_id - >>> last_hidden_state = model(input_features, decoder_input_ids=decoder_input_ids).last_hidden_state - >>> list(last_hidden_state.shape) - [1, 2, 512] - ``` -""" - # Tokenize inputs. def tokenize(sample): @@ -71,23 +49,29 @@ def tokenize(sample): return_tensors="pt", ).input_features - decoder_input_ids = torch.ones((batch_size, 1), dtype=torch.long) * model.config.decoder_start_token_id + generation_config, _kwargs = model._prepare_generation_config(None) - return { - "input_features": input_features, - "decoder_input_ids": decoder_input_ids - } + input_stride = ( + model.model.encoder.conv1.stride[0] * model.model.encoder.conv2.stride[0] + ) + num_segment_frames = input_stride * model.config.max_source_positions + + decoder_input_ids = model._retrieve_init_tokens( + input_features, + batch_size=batch_size, + generation_config=generation_config, + config=model.config, + num_segment_frames=num_segment_frames, + kwargs={}, + ) + + return {"input_features": input_features, "decoder_input_ids": decoder_input_ids} ds = ds.map(tokenize, remove_columns=ds.column_names) # Configure the quantization algorithm to run. # * quantize the weights to 4 bit with GPTQ with a group size 128 -#breakpoint() -#sample_input = next(iter(ds)) -#output = model(**sample_input) - - recipe = GPTQModifier(targets="Linear", scheme="W4A16", ignore=["lm_head"]) # Apply algorithms. @@ -103,15 +87,14 @@ def tokenize(sample): # Confirm generations of the quantized model look sane. print("\n\n") print("========== SAMPLE GENERATION ==============") -sample_input = whisper_data_collator([next(iter(ds))]).to(model.device) +sample_input = whisper_data_collator([next(iter(ds))]) sample_input = {k: v.to("cuda:0") for k, v in sample_input.items()} -output = model.generate(**sample_input) -breakpoint() +output = model.generate(**sample_input, language="en") print(processor.batch_decode(output, skip_special_tokens=True)) -#[' Mr. Quilter is the apostle of the middle classes and we are glad to welcome his gospel.'] +# Mr. Quilter is the apostle of the middle classes and we are glad to welcome his gospel print("==========================================\n\n") # Save to disk compressed. SAVE_DIR = MODEL_ID.split("/")[1] + "-W4A16-G128" model.save_pretrained(SAVE_DIR, save_compressed=True) -processor.save_pretrained(SAVE_DIR) \ No newline at end of file +processor.save_pretrained(SAVE_DIR) diff --git a/src/llmcompressor/transformers/finetune/data/base.py b/src/llmcompressor/transformers/finetune/data/base.py index 8a4c150cd..a5ec88570 100644 --- a/src/llmcompressor/transformers/finetune/data/base.py +++ b/src/llmcompressor/transformers/finetune/data/base.py @@ -105,7 +105,9 @@ def __call__(self, add_labels: bool = True) -> DatasetType: dataset = self.rename_columns(dataset) logger.debug(f"Dataset after column renaming: {get_columns(dataset)}") - if "input_ids" not in get_columns(dataset) and "input_features" not in get_columns(dataset): + if "input_ids" not in get_columns( + dataset + ) and "input_features" not in get_columns(dataset): # tokenize/ process dataset = self.filter_tokenizer_args(dataset) logger.debug(f"Tokenizer args after filtering: {get_columns(dataset)}") diff --git a/src/llmcompressor/transformers/tracing/whisper.py b/src/llmcompressor/transformers/tracing/whisper.py index b8241f764..6e245760c 100644 --- a/src/llmcompressor/transformers/tracing/whisper.py +++ b/src/llmcompressor/transformers/tracing/whisper.py @@ -1,3 +1,4 @@ +# flake8: noqa # coding=utf-8 # Copyright 2022 The OpenAI Authors and The HuggingFace Inc. team. All rights reserved. # From 8ca9b6d718826fd57a08050b338e8f46a905e264 Mon Sep 17 00:00:00 2001 From: Kyle Sayers Date: Sun, 19 Jan 2025 01:35:52 +0000 Subject: [PATCH 04/21] readme todo Signed-off-by: Kyle Sayers --- examples/multimodal_audio/README.md | 1 + 1 file changed, 1 insertion(+) create mode 100644 examples/multimodal_audio/README.md diff --git a/examples/multimodal_audio/README.md b/examples/multimodal_audio/README.md new file mode 100644 index 000000000..30404ce4c --- /dev/null +++ b/examples/multimodal_audio/README.md @@ -0,0 +1 @@ +TODO \ No newline at end of file From 98aca1686ded0947364e381f654930ec7f58b7c3 Mon Sep 17 00:00:00 2001 From: Kyle Sayers Date: Sun, 19 Jan 2025 05:46:39 +0000 Subject: [PATCH 05/21] split to peoples_speech dataset Signed-off-by: Kyle Sayers --- .../multimodal_audio/qwen2_audio_example.py | 118 ++++++++++++++++++ examples/multimodal_audio/whisper_example.py | 14 ++- 2 files changed, 128 insertions(+), 4 deletions(-) create mode 100644 examples/multimodal_audio/qwen2_audio_example.py diff --git a/examples/multimodal_audio/qwen2_audio_example.py b/examples/multimodal_audio/qwen2_audio_example.py new file mode 100644 index 000000000..a3a1d6b13 --- /dev/null +++ b/examples/multimodal_audio/qwen2_audio_example.py @@ -0,0 +1,118 @@ +from datasets import load_dataset +from transformers import Qwen2AudioForConditionalGeneration, AutoProcessor + +from llmcompressor.modifiers.quantization import GPTQModifier +from llmcompressor.transformers import oneshot +from llmcompressor.transformers.tracing import TraceableWhisperForConditionalGeneration +from llmcompressor.transformers.utils.data_collator import whisper_data_collator + +# Select model and load it. +MODEL_ID = "Qwen/Qwen2-Audio-7B-Instruct" + +model = Qwen2AudioForConditionalGeneration.from_pretrained( + MODEL_ID, + device_map="auto", + torch_dtype="auto", +) +processor = AutoProcessor.from_pretrained(MODEL_ID) + +# Select calibration dataset. +DATASET_ID = "MLCommons/peoples_speech" +DATASET_SUBSET = "test" +DATASET_SPLIT = "test" + +# Select number of samples. 512 samples is a good place to start. +# Increasing the number of samples can improve accuracy. +NUM_CALIBRATION_SAMPLES = 512 +MAX_SEQUENCE_LENGTH = 2048 + +# Load dataset and preprocess. +ds = load_dataset( + DATASET_ID, + DATASET_SUBSET, + split=f"{DATASET_SPLIT}[:{NUM_CALIBRATION_SAMPLES}]", + trust_remote_code=True, +) + + +def preprocess(example): + conversation = [ + {"role": "user", "content": [ + {"type": "audio", "audio_url": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2-Audio/audio/guess_age_gender.wav"}, + ]}, + {"role": "assistant", "content": "Yes, the speaker is female and in her twenties."}, + {"role": "user", "content": [ + {"type": "audio", "audio_url": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2-Audio/audio/translate_to_chinese.wav"}, + ]}, + ] + return { + "text": processor.apply_chat_template(conversation, add_generation_prompt=True, tokenize=False) + } + return { + "array": example["audio"]["array"], + "sampling_rate": example["audio"]["sampling_rate"], + } + + +ds = ds.map(preprocess, remove_columns=ds.column_names) + + +# Tokenize inputs. +def tokenize(sample): + batch_size = 1 + input_features = processor( + sample["array"], + sampling_rate=sample["sampling_rate"], + return_tensors="pt", + ).input_features + + generation_config, _kwargs = model._prepare_generation_config(None) + + input_stride = ( + model.model.encoder.conv1.stride[0] * model.model.encoder.conv2.stride[0] + ) + num_segment_frames = input_stride * model.config.max_source_positions + + decoder_input_ids = model._retrieve_init_tokens( + input_features, + batch_size=batch_size, + generation_config=generation_config, + config=model.config, + num_segment_frames=num_segment_frames, + kwargs={}, + ) + + return {"input_features": input_features, "decoder_input_ids": decoder_input_ids} + + +ds = ds.map(tokenize, remove_columns=ds.column_names) + +# Configure the quantization algorithm to run. +# * quantize the weights to 4 bit with GPTQ with a group size 128 +recipe = GPTQModifier(targets="Linear", scheme="W4A16", ignore=["lm_head"]) + +# Apply algorithms. +oneshot( + model=model, + dataset=ds, + recipe=recipe, + max_seq_length=MAX_SEQUENCE_LENGTH, + num_calibration_samples=NUM_CALIBRATION_SAMPLES, + data_collator=whisper_data_collator, +) + +# Confirm generations of the quantized model look sane. +print("\n\n") +print("========== SAMPLE GENERATION ==============") +sample_input = whisper_data_collator([next(iter(ds))]) +sample_input = {k: v.to("cuda:0") for k, v in sample_input.items()} +output = model.generate(**sample_input, language="en") +print(processor.batch_decode(output, skip_special_tokens=True)[0]) +print("==========================================\n\n") +# If you are interested in doing something to your house, go to the green building +# adviser, look it up and see what the experts are talking about + +# Save to disk compressed. +SAVE_DIR = MODEL_ID.split("/")[1] + "-W4A16-G128" +model.save_pretrained(SAVE_DIR, save_compressed=True) +processor.save_pretrained(SAVE_DIR) diff --git a/examples/multimodal_audio/whisper_example.py b/examples/multimodal_audio/whisper_example.py index dd3bf2fb1..b87185252 100644 --- a/examples/multimodal_audio/whisper_example.py +++ b/examples/multimodal_audio/whisper_example.py @@ -18,8 +18,9 @@ processor = WhisperProcessor.from_pretrained(MODEL_ID) # Select calibration dataset. -DATASET_ID = "hf-internal-testing/librispeech_asr_dummy" -DATASET_SPLIT = "validation[:512]" +DATASET_ID = "MLCommons/peoples_speech" +DATASET_SUBSET = "test" +DATASET_SPLIT = "test" # Select number of samples. 512 samples is a good place to start. # Increasing the number of samples can improve accuracy. @@ -27,7 +28,12 @@ MAX_SEQUENCE_LENGTH = 2048 # Load dataset and preprocess. -ds = load_dataset(DATASET_ID, "clean", split=DATASET_SPLIT) +ds = load_dataset( + DATASET_ID, + DATASET_SUBSET, + split=f"{DATASET_SPLIT}[:{NUM_CALIBRATION_SAMPLES}]", + trust_remote_code=True, +) def preprocess(example): @@ -91,8 +97,8 @@ def tokenize(sample): sample_input = {k: v.to("cuda:0") for k, v in sample_input.items()} output = model.generate(**sample_input, language="en") print(processor.batch_decode(output, skip_special_tokens=True)) -# Mr. Quilter is the apostle of the middle classes and we are glad to welcome his gospel print("==========================================\n\n") +# The track appears on the compilation album "Kraftworks" # Save to disk compressed. SAVE_DIR = MODEL_ID.split("/")[1] + "-W4A16-G128" From 7067c3f6d087012c020ecd42a2301b11eb4585c2 Mon Sep 17 00:00:00 2001 From: Kyle Sayers Date: Sun, 19 Jan 2025 06:09:24 +0000 Subject: [PATCH 06/21] use cleanup example, add todo check Signed-off-by: Kyle Sayers --- examples/multimodal_audio/whisper_example.py | 14 ++++++-------- .../transformers/finetune/data/base.py | 1 + 2 files changed, 7 insertions(+), 8 deletions(-) diff --git a/examples/multimodal_audio/whisper_example.py b/examples/multimodal_audio/whisper_example.py index b87185252..ad7830d6f 100644 --- a/examples/multimodal_audio/whisper_example.py +++ b/examples/multimodal_audio/whisper_example.py @@ -7,7 +7,7 @@ from llmcompressor.transformers.utils.data_collator import whisper_data_collator # Select model and load it. -MODEL_ID = "openai/whisper-tiny" +MODEL_ID = "openai/whisper-base" model = TraceableWhisperForConditionalGeneration.from_pretrained( MODEL_ID, @@ -46,25 +46,23 @@ def preprocess(example): ds = ds.map(preprocess, remove_columns=ds.column_names) -# Tokenize inputs. -def tokenize(sample): - batch_size = 1 +# Process inputs. +def process(sample): input_features = processor( sample["array"], sampling_rate=sample["sampling_rate"], return_tensors="pt", ).input_features + # decoder_input_ids define the task context generation_config, _kwargs = model._prepare_generation_config(None) - input_stride = ( model.model.encoder.conv1.stride[0] * model.model.encoder.conv2.stride[0] ) num_segment_frames = input_stride * model.config.max_source_positions - decoder_input_ids = model._retrieve_init_tokens( input_features, - batch_size=batch_size, + batch_size=1, generation_config=generation_config, config=model.config, num_segment_frames=num_segment_frames, @@ -74,7 +72,7 @@ def tokenize(sample): return {"input_features": input_features, "decoder_input_ids": decoder_input_ids} -ds = ds.map(tokenize, remove_columns=ds.column_names) +ds = ds.map(process, remove_columns=ds.column_names) # Configure the quantization algorithm to run. # * quantize the weights to 4 bit with GPTQ with a group size 128 diff --git a/src/llmcompressor/transformers/finetune/data/base.py b/src/llmcompressor/transformers/finetune/data/base.py index a5ec88570..10e014949 100644 --- a/src/llmcompressor/transformers/finetune/data/base.py +++ b/src/llmcompressor/transformers/finetune/data/base.py @@ -105,6 +105,7 @@ def __call__(self, add_labels: bool = True) -> DatasetType: dataset = self.rename_columns(dataset) logger.debug(f"Dataset after column renaming: {get_columns(dataset)}") + # TODO: investigate processor.model_input_names if "input_ids" not in get_columns( dataset ) and "input_features" not in get_columns(dataset): From 0848fb64649741162b9ad0ffb6adc2c4283dae82 Mon Sep 17 00:00:00 2001 From: Kyle Sayers Date: Sun, 19 Jan 2025 06:29:27 +0000 Subject: [PATCH 07/21] qwen2, need to add traceability Signed-off-by: Kyle Sayers --- .../multimodal_audio/qwen2_audio_example.py | 67 +- examples/multimodal_audio/whisper_example.py | 2 +- .../transformers/tracing/qwen2_audio.py | 1368 +++++++++++++++++ .../transformers/utils/data_collator.py | 10 + 4 files changed, 1401 insertions(+), 46 deletions(-) create mode 100644 src/llmcompressor/transformers/tracing/qwen2_audio.py diff --git a/examples/multimodal_audio/qwen2_audio_example.py b/examples/multimodal_audio/qwen2_audio_example.py index a3a1d6b13..2f0f1cc35 100644 --- a/examples/multimodal_audio/qwen2_audio_example.py +++ b/examples/multimodal_audio/qwen2_audio_example.py @@ -1,10 +1,9 @@ from datasets import load_dataset -from transformers import Qwen2AudioForConditionalGeneration, AutoProcessor +from transformers import AutoProcessor, Qwen2AudioForConditionalGeneration from llmcompressor.modifiers.quantization import GPTQModifier from llmcompressor.transformers import oneshot -from llmcompressor.transformers.tracing import TraceableWhisperForConditionalGeneration -from llmcompressor.transformers.utils.data_collator import whisper_data_collator +from llmcompressor.transformers.utils.data_collator import qwen2_audio_data_collator # Select model and load it. MODEL_ID = "Qwen/Qwen2-Audio-7B-Instruct" @@ -36,20 +35,21 @@ def preprocess(example): - conversation = [ - {"role": "user", "content": [ - {"type": "audio", "audio_url": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2-Audio/audio/guess_age_gender.wav"}, - ]}, - {"role": "assistant", "content": "Yes, the speaker is female and in her twenties."}, - {"role": "user", "content": [ - {"type": "audio", "audio_url": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2-Audio/audio/translate_to_chinese.wav"}, - ]}, + messages = [ + { + "role": "user", + "content": [ + {"type": "audio", "audio": None}, + {"type": "text", "text": "What does the person say?"}, + ], + }, ] + return { - "text": processor.apply_chat_template(conversation, add_generation_prompt=True, tokenize=False) - } - return { - "array": example["audio"]["array"], + "text": processor.apply_chat_template( + messages, add_generation_prompt=True, tokenize=False + ), + "audios": [example["audio"]["array"]], "sampling_rate": example["audio"]["sampling_rate"], } @@ -59,30 +59,7 @@ def preprocess(example): # Tokenize inputs. def tokenize(sample): - batch_size = 1 - input_features = processor( - sample["array"], - sampling_rate=sample["sampling_rate"], - return_tensors="pt", - ).input_features - - generation_config, _kwargs = model._prepare_generation_config(None) - - input_stride = ( - model.model.encoder.conv1.stride[0] * model.model.encoder.conv2.stride[0] - ) - num_segment_frames = input_stride * model.config.max_source_positions - - decoder_input_ids = model._retrieve_init_tokens( - input_features, - batch_size=batch_size, - generation_config=generation_config, - config=model.config, - num_segment_frames=num_segment_frames, - kwargs={}, - ) - - return {"input_features": input_features, "decoder_input_ids": decoder_input_ids} + return processor(**sample, return_tensors="pt") ds = ds.map(tokenize, remove_columns=ds.column_names) @@ -98,19 +75,19 @@ def tokenize(sample): recipe=recipe, max_seq_length=MAX_SEQUENCE_LENGTH, num_calibration_samples=NUM_CALIBRATION_SAMPLES, - data_collator=whisper_data_collator, + data_collator=qwen2_audio_data_collator, ) # Confirm generations of the quantized model look sane. print("\n\n") print("========== SAMPLE GENERATION ==============") -sample_input = whisper_data_collator([next(iter(ds))]) -sample_input = {k: v.to("cuda:0") for k, v in sample_input.items()} -output = model.generate(**sample_input, language="en") +breakpoint() +sample_input = qwen2_audio_data_collator([next(iter(ds))]) +sample_input = {k: v.to(model.device) for k, v in sample_input.items()} +output = model.generate(**sample_input) print(processor.batch_decode(output, skip_special_tokens=True)[0]) print("==========================================\n\n") -# If you are interested in doing something to your house, go to the green building -# adviser, look it up and see what the experts are talking about +# that's where you have a lot of windows in the # Save to disk compressed. SAVE_DIR = MODEL_ID.split("/")[1] + "-W4A16-G128" diff --git a/examples/multimodal_audio/whisper_example.py b/examples/multimodal_audio/whisper_example.py index ad7830d6f..55c9a2fab 100644 --- a/examples/multimodal_audio/whisper_example.py +++ b/examples/multimodal_audio/whisper_example.py @@ -92,7 +92,7 @@ def process(sample): print("\n\n") print("========== SAMPLE GENERATION ==============") sample_input = whisper_data_collator([next(iter(ds))]) -sample_input = {k: v.to("cuda:0") for k, v in sample_input.items()} +sample_input = {k: v.to(model.device) for k, v in sample_input.items()} output = model.generate(**sample_input, language="en") print(processor.batch_decode(output, skip_special_tokens=True)) print("==========================================\n\n") diff --git a/src/llmcompressor/transformers/tracing/qwen2_audio.py b/src/llmcompressor/transformers/tracing/qwen2_audio.py new file mode 100644 index 000000000..42fca41b9 --- /dev/null +++ b/src/llmcompressor/transformers/tracing/qwen2_audio.py @@ -0,0 +1,1368 @@ +# flake8: noqa +# coding=utf-8 +# Copyright 2024 the HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# vllm-project: no copyright +"""PyTorch Qwen2Audio model.""" + +import math +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Tuple, Union + +import torch +import torch.utils.checkpoint +from torch import nn + +from ...activations import ACT2FN +from ...cache_utils import Cache, EncoderDecoderCache, StaticCache +from ...generation import GenerationMixin +from ...modeling_outputs import BaseModelOutput, ModelOutput +from ...modeling_utils import PreTrainedModel +from ...utils import ( + add_start_docstrings, + add_start_docstrings_to_model_forward, + is_flash_attn_2_available, + is_flash_attn_greater_or_equal_2_10, + logging, + replace_return_docstrings, +) +from ..auto import AutoModel, AutoModelForCausalLM +from .configuration_qwen2_audio import Qwen2AudioConfig, Qwen2AudioEncoderConfig + + +if is_flash_attn_2_available(): + from ...modeling_flash_attention_utils import _flash_attention_forward + + +logger = logging.get_logger(__name__) + +_CONFIG_FOR_DOC = "Qwen2AudioConfig" + + +@dataclass +class Qwen2AudioCausalLMOutputWithPast(ModelOutput): + """ + Base class for Qwen2Audio causal language model (or autoregressive) outputs. + + Args: + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Language modeling loss (for next-token prediction). + logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`): + Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). + past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape + `(batch_size, num_heads, sequence_length, embed_size_per_head)`) + + Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see + `past_key_values` input) to speed up sequential decoding. + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + attention_mask (`torch.FloatTensor`, *optional*): + Attentions mask, used to update attention mask and position_ids. + """ + + loss: Optional[torch.FloatTensor] = None + logits: torch.FloatTensor = None + past_key_values: Optional[List[torch.FloatTensor]] = None + hidden_states: Optional[Tuple[torch.FloatTensor]] = None + attentions: Optional[Tuple[torch.FloatTensor]] = None + attention_mask: Optional[torch.FloatTensor] = None + + +# Copied from transformers.models.whisper.modeling_whisper.WhisperAttention with Whisper->Qwen2Audio +class Qwen2AudioAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__( + self, + embed_dim: int, + num_heads: int, + dropout: float = 0.0, + is_decoder: bool = False, + bias: bool = True, + is_causal: bool = False, + layer_idx: Optional[int] = None, + config: Optional[Qwen2AudioConfig] = None, + ): + super().__init__() + self.embed_dim = embed_dim + self.num_heads = num_heads + self.dropout = dropout + self.head_dim = embed_dim // num_heads + self.config = config + + if (self.head_dim * num_heads) != self.embed_dim: + raise ValueError( + f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}" + f" and `num_heads`: {num_heads})." + ) + self.scaling = self.head_dim**-0.5 + self.is_decoder = is_decoder + self.is_causal = is_causal + + if layer_idx is None and is_decoder: + logger.warning_once( + f"Instantiating a decoder {self.__class__.__name__} without passing `layer_idx` is not recommended and " + "will to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` " + "when creating this class." + ) + self.layer_idx = layer_idx + + self.k_proj = nn.Linear(embed_dim, embed_dim, bias=False) + self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias) + self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias) + self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias) + + # Copied from transformers.models.bart.modeling_bart.BartAttention._shape with BART->whisper + def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int): + return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous() + + def forward( + self, + hidden_states: torch.Tensor, + key_value_states: Optional[torch.Tensor] = None, + past_key_value: Optional[EncoderDecoderCache] = None, + attention_mask: Optional[torch.Tensor] = None, + layer_head_mask: Optional[torch.Tensor] = None, + output_attentions: bool = False, + cache_position: Optional[torch.LongTensor] = None, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + """Input shape: Batch x Time x Channel""" + + # if key_value_states are provided this layer is used as a cross-attention layer + # for the decoder + is_cross_attention = key_value_states is not None + bsz, tgt_len, _ = hidden_states.size() + + # get query proj + query_states = self._shape(self.q_proj(hidden_states) * self.scaling, tgt_len, bsz) + + if past_key_value is not None: + is_updated = past_key_value.is_updated.get(self.layer_idx) + if is_cross_attention: + # after the first generated id, we can subsequently re-use all key/value_states from cache + past_key_value.is_updated[self.layer_idx] = True + past_key_value = past_key_value.cross_attention_cache + else: + past_key_value = past_key_value.self_attention_cache + + # use key_value_states if cross attention + current_states = key_value_states if key_value_states is not None else hidden_states + if is_cross_attention and past_key_value and is_updated: + # reuse k,v, cross_attentions + key_states = past_key_value.key_cache[self.layer_idx] + value_states = past_key_value.value_cache[self.layer_idx] + else: + key_states = self._shape(self.k_proj(current_states), -1, bsz) + value_states = self._shape(self.v_proj(current_states), -1, bsz) + if past_key_value is not None: + # save all key/value_states to cache to be re-used for fast auto-regressive generation + cache_position = cache_position if not is_cross_attention else None + key_states, value_states = past_key_value.update( + key_states, value_states, self.layer_idx, {"cache_position": cache_position} + ) + + attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) + + if attention_mask is not None: # no matter the length, we just slice it + causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] + attn_weights = attn_weights + causal_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1) + + if layer_head_mask is not None: + if layer_head_mask.size() != (self.num_heads,): + raise ValueError( + f"Head mask for a single layer should be of size {(self.num_heads,)}, but is" + f" {layer_head_mask.size()}" + ) + attn_weights = layer_head_mask.view(1, -1, 1, 1) * attn_weights + + attn_probs = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training) + attn_output = torch.matmul(attn_probs, value_states) + + if attn_output.size() != (bsz, self.num_heads, tgt_len, self.head_dim): + raise ValueError( + f"`attn_output` should be of size {(bsz, self.num_heads, tgt_len, self.head_dim)}, but is" + f" {attn_output.size()}" + ) + + attn_output = attn_output.transpose(1, 2) + # Use the `embed_dim` from the config (stored in the class) rather than `hidden_state` because `attn_output` can be + # partitioned across GPUs when using tensor-parallelism. + attn_output = attn_output.reshape(bsz, tgt_len, self.embed_dim) + + attn_output = self.out_proj(attn_output) + + return attn_output, attn_weights, past_key_value + + +# Copied from transformers.models.whisper.modeling_whisper.WhisperFlashAttention2 with Whisper->Qwen2Audio +class Qwen2AudioFlashAttention2(Qwen2AudioAttention): + """ + Qwen2Audio flash attention module. This module inherits from `Qwen2AudioAttention` as the weights of the module stays + untouched. The only required change would be on the forward pass where it needs to correctly call the public API of + flash attention and deal with padding tokens in case the input contains any of them. + """ + + # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2.__init__ + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1. + # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0. + # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left). + self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10() + + def forward( + self, + hidden_states: torch.Tensor, + key_value_states: Optional[torch.Tensor] = None, + past_key_value: Optional[EncoderDecoderCache] = None, + attention_mask: Optional[torch.Tensor] = None, + layer_head_mask: Optional[torch.Tensor] = None, + output_attentions: bool = False, + cache_position: Optional[torch.LongTensor] = None, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + if isinstance(past_key_value, StaticCache): + raise ValueError( + "The `static` cache implementation is not compatible with `attn_implementation='flash_attention_2'`. " + "Use `attn_implementation='sdpa'` in the meantime, and open an issue at https://github.com/huggingface/transformers" + ) + # Qwen2AudioFlashAttention2 attention does not support output_attentions + if output_attentions: + raise ValueError("Qwen2AudioFlashAttention2 attention does not support output_attentions") + + # if key_value_states are provided this layer is used as a cross-attention layer + # for the decoder + is_cross_attention = key_value_states is not None + bsz, tgt_len, _ = hidden_states.size() + + # get query proj + query_states = torch.reshape(self.q_proj(hidden_states), (bsz, tgt_len, self.num_heads, self.head_dim)) + + if past_key_value is not None: + is_updated = past_key_value.is_updated.get(self.layer_idx) + if is_cross_attention: + # after the first generated id, we can subsequently re-use all key/value_states from cache + past_key_value.is_updated[self.layer_idx] = True + past_key_value = past_key_value.cross_attention_cache + else: + past_key_value = past_key_value.self_attention_cache + + # use key_value_states if cross attention + current_states = key_value_states if key_value_states is not None else hidden_states + if is_cross_attention and past_key_value and is_updated: + # reuse k,v, cross_attentions + key_states = past_key_value.key_cache[self.layer_idx] + value_states = past_key_value.value_cache[self.layer_idx] + else: + key_states = self._shape(self.k_proj(current_states), -1, bsz) + value_states = self._shape(self.v_proj(current_states), -1, bsz) + if past_key_value is not None: + # save all key/value_states to cache to be re-used for fast auto-regressive generation + cache_position = cache_position if not is_cross_attention else None + key_states, value_states = past_key_value.update( + key_states, value_states, self.layer_idx, {"cache_position": cache_position} + ) + + # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim] + # We would need to refactor the KV cache to be able to avoid many of these transpose/reshape/view. + key_states = key_states.transpose(1, 2) + value_states = value_states.transpose(1, 2) + + causal_mask = attention_mask + if attention_mask is not None: # no matter the length, we just slice it + causal_mask = attention_mask[:, : key_states.shape[-2]] + + # In PEFT, usually we cast the layer norms in float32 for training stability reasons + # therefore the input hidden states gets silently casted in float32. Hence, we need + # cast them back in the correct dtype just to be sure everything works as expected. + # This might slowdown training & inference so it is recommended to not cast the LayerNorms + # in fp32. (LlamaRMSNorm handles it correctly) + + input_dtype = query_states.dtype + if input_dtype == torch.float32: + if torch.is_autocast_enabled(): + target_dtype = torch.get_autocast_gpu_dtype() + # Handle the case where the model is quantized + elif hasattr(self.config, "_pre_quantization_dtype"): + target_dtype = self.config._pre_quantization_dtype + else: + target_dtype = self.q_proj.weight.dtype + + logger.warning_once( + f"The input hidden states seems to be silently casted in float32, this might be related to" + f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in" + f" {target_dtype}." + ) + + query_states = query_states.to(target_dtype) + key_states = key_states.to(target_dtype) + value_states = value_states.to(target_dtype) + + attn_output = _flash_attention_forward( + query_states, + key_states, + value_states, + causal_mask, + tgt_len, + dropout=self.dropout if self.training else 0.0, + is_causal=self.is_causal, + use_top_left_mask=self._flash_attn_uses_top_left_mask, + ) + + attn_output = attn_output.reshape(bsz, tgt_len, -1) + attn_output = self.out_proj(attn_output) + + if not output_attentions: + attn_weights = None + + return attn_output, attn_weights, past_key_value + + +# Copied from transformers.models.whisper.modeling_whisper.WhisperSdpaAttention with Whisper->Qwen2Audio +class Qwen2AudioSdpaAttention(Qwen2AudioAttention): + def forward( + self, + hidden_states: torch.Tensor, + key_value_states: Optional[torch.Tensor] = None, + past_key_value: Optional[EncoderDecoderCache] = None, + attention_mask: Optional[torch.Tensor] = None, + layer_head_mask: Optional[torch.Tensor] = None, + output_attentions: bool = False, + cache_position: Optional[torch.LongTensor] = None, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + """Input shape: Batch x Time x Channel""" + if output_attentions or layer_head_mask is not None: + # TODO: Improve this warning with e.g. `model.config._attn_implementation = "manual"` once this is implemented. + logger.warning_once( + "Qwen2AudioModel is using Qwen2AudioSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True` or `layer_head_mask` not None. Falling back to the manual attention" + ' implementation, but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.' + ) + return super().forward( + hidden_states, + key_value_states=key_value_states, + past_key_value=past_key_value, + attention_mask=attention_mask, + layer_head_mask=layer_head_mask, + output_attentions=output_attentions, + cache_position=cache_position, + ) + + # if key_value_states are provided this layer is used as a cross-attention layer + # for the decoder + is_cross_attention = key_value_states is not None + bsz, tgt_len, _ = hidden_states.size() + + # get query proj + query_states = self._shape(self.q_proj(hidden_states), tgt_len, bsz) + + if past_key_value is not None: + is_updated = past_key_value.is_updated.get(self.layer_idx) + if is_cross_attention: + # after the first generated id, we can subsequently re-use all key/value_states from cache + past_key_value.is_updated[self.layer_idx] = True + past_key_value = past_key_value.cross_attention_cache + else: + past_key_value = past_key_value.self_attention_cache + + # use key_value_states if cross attention + current_states = key_value_states if key_value_states is not None else hidden_states + if is_cross_attention and past_key_value and is_updated: + # reuse k,v, cross_attentions + key_states = past_key_value.key_cache[self.layer_idx] + value_states = past_key_value.value_cache[self.layer_idx] + else: + key_states = self._shape(self.k_proj(current_states), -1, bsz) + value_states = self._shape(self.v_proj(current_states), -1, bsz) + if past_key_value is not None: + # save all key/value_states to cache to be re-used for fast auto-regressive generation + cache_position = cache_position if not is_cross_attention else None + key_states, value_states = past_key_value.update( + key_states, value_states, self.layer_idx, {"cache_position": cache_position} + ) + + causal_mask = attention_mask + if attention_mask is not None: # no matter the length, we just slice it + causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] + + # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an inline conditional assignment + # in SDPA to support both torch.compile's dynamic shapes and full graph options. An inline conditional prevents dynamic shapes from compiling. + # The tgt_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case tgt_len == 1. + is_causal = True if self.is_causal and causal_mask is None and tgt_len > 1 else False + + # NOTE: SDPA with memory-efficient backend is currently (torch==2.1.2) bugged when using non-contiguous inputs and a custom attn_mask, + # but we are fine here as `_shape` do call `.contiguous()`. Reference: https://github.com/pytorch/pytorch/issues/112577 + attn_output = torch.nn.functional.scaled_dot_product_attention( + query_states, + key_states, + value_states, + attn_mask=causal_mask, + dropout_p=self.dropout if self.training else 0.0, + is_causal=is_causal, + ) + + if attn_output.size() != (bsz, self.num_heads, tgt_len, self.head_dim): + raise ValueError( + f"`attn_output` should be of size {(bsz, self.num_heads, tgt_len, self.head_dim)}, but is" + f" {attn_output.size()}" + ) + + attn_output = attn_output.transpose(1, 2) + + # Use the `embed_dim` from the config (stored in the class) rather than `hidden_state` because `attn_output` can be + # partitioned across GPUs when using tensor-parallelism. + attn_output = attn_output.reshape(bsz, tgt_len, self.embed_dim) + + attn_output = self.out_proj(attn_output) + + return attn_output, None, past_key_value + + +QWEN2AUDIO_ATTENTION_CLASSES = { + "eager": Qwen2AudioAttention, + "flash_attention_2": Qwen2AudioFlashAttention2, + "sdpa": Qwen2AudioSdpaAttention, +} + + +# Copied from transformers.models.whisper.modeling_whisper.WhisperEncoderLayer with Whisper->Qwen2Audio, WHISPER->QWEN2AUDIO +class Qwen2AudioEncoderLayer(nn.Module): + def __init__(self, config: Qwen2AudioConfig): + super().__init__() + self.embed_dim = config.d_model + + self.self_attn = QWEN2AUDIO_ATTENTION_CLASSES[config._attn_implementation]( + embed_dim=self.embed_dim, + num_heads=config.encoder_attention_heads, + dropout=config.attention_dropout, + config=config, + ) + self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim) + self.dropout = config.dropout + self.activation_fn = ACT2FN[config.activation_function] + self.activation_dropout = config.activation_dropout + self.fc1 = nn.Linear(self.embed_dim, config.encoder_ffn_dim) + self.fc2 = nn.Linear(config.encoder_ffn_dim, self.embed_dim) + self.final_layer_norm = nn.LayerNorm(self.embed_dim) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor, + layer_head_mask: torch.Tensor, + output_attentions: bool = False, + ) -> torch.Tensor: + """ + Args: + hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` + attention_mask (`torch.FloatTensor`): attention mask of size + `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. + layer_head_mask (`torch.FloatTensor`): mask for attention heads in a given layer of size + `(encoder_attention_heads,)`. + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + """ + residual = hidden_states + hidden_states = self.self_attn_layer_norm(hidden_states) + hidden_states, attn_weights, _ = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + layer_head_mask=layer_head_mask, + output_attentions=output_attentions, + ) + hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.final_layer_norm(hidden_states) + hidden_states = self.activation_fn(self.fc1(hidden_states)) + hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training) + hidden_states = self.fc2(hidden_states) + hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) + hidden_states = residual + hidden_states + + if hidden_states.dtype == torch.float16 and ( + torch.isinf(hidden_states).any() or torch.isnan(hidden_states).any() + ): + clamp_value = torch.finfo(hidden_states.dtype).max - 1000 + hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) + + outputs = (hidden_states,) + + if output_attentions: + outputs += (attn_weights,) + + return outputs + + +QWEN2AUDIO_START_DOCSTRING = r""" + This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the + library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads + etc.) + + This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. + Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage + and behavior. + + Parameters: + config ([`Qwen2AudioConfig`]): + Model configuration class with all the parameters of the model. Initializing with a config file does not + load the weights associated with the model, only the configuration. Check out the + [`~PreTrainedModel.from_pretrained`] method to load the model weights. +""" + + +@add_start_docstrings( + "The bare Qwen2Audio Model outputting raw hidden-states without any specific head on top.", + QWEN2AUDIO_START_DOCSTRING, +) +class Qwen2AudioPreTrainedModel(PreTrainedModel): + config_class = Qwen2AudioConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["Qwen2AudioAttention"] + _skip_keys_device_placement = "past_key_values" + _supports_flash_attn_2 = True + _supports_sdpa = True + + def _init_weights(self, module): + # important: this ported version of Qwen2Audio isn't meant for training from scratch - only + # inference and fine-tuning - so the proper init weights code has been removed + std = self.config.init_std if hasattr(self.config, "init_std") else self.config.audio_config.init_std + + if isinstance(module, (nn.Linear, nn.Conv1d)): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + + +QWEN2AUDIOENCODER_START_DOCSTRING = r""" + This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the + library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads + etc.) + + This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. + Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage + and behavior. + + Parameters: + config ([`Qwen2AudioEncoderConfig`]): + Model configuration class with all the parameters of the model. Initializing with a config file does not + load the weights associated with the model, only the configuration. Check out the + [`~PreTrainedModel.from_pretrained`] method to load the model weights. +""" + + +@add_start_docstrings( + """The audio model from Qwen2Audio without any head or projection on top.""", + QWEN2AUDIOENCODER_START_DOCSTRING, +) +# Copied from transformers.models.whisper.modeling_whisper.WhisperEncoder with Whisper->Qwen2Audio +class Qwen2AudioEncoder(Qwen2AudioPreTrainedModel): + """ + Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a + [`Qwen2AudioEncoderLayer`]. + + Args: + config: Qwen2AudioEncoderConfig + """ + + # Ignore copy + config_class = Qwen2AudioEncoderConfig + main_input_name = "input_features" + _no_split_modules = ["Qwen2AudioEncoderLayer"] + + def __init__(self, config: Qwen2AudioEncoderConfig): + super().__init__(config) + self.dropout = config.dropout + self.layerdrop = config.encoder_layerdrop + + embed_dim = config.d_model + self.num_mel_bins = config.num_mel_bins + self.padding_idx = config.pad_token_id + self.max_source_positions = config.max_source_positions + self.embed_scale = math.sqrt(embed_dim) if config.scale_embedding else 1.0 + + self.conv1 = nn.Conv1d(self.num_mel_bins, embed_dim, kernel_size=3, padding=1) + self.conv2 = nn.Conv1d(embed_dim, embed_dim, kernel_size=3, stride=2, padding=1) + + self.embed_positions = nn.Embedding(self.max_source_positions, embed_dim) + self.embed_positions.requires_grad_(False) + + self.layers = nn.ModuleList([Qwen2AudioEncoderLayer(config) for _ in range(config.encoder_layers)]) + self.layer_norm = nn.LayerNorm(config.d_model) + # Ignore copy + self.avg_pooler = nn.AvgPool1d(2, stride=2) + + self.gradient_checkpointing = False + # Initialize weights and apply final processing + self.post_init() + + def _freeze_parameters(self): + for param in self.parameters(): + param.requires_grad = False + self._requires_grad = False + + def get_input_embeddings(self) -> nn.Module: + return self.conv1 + + def set_input_embeddings(self, value: nn.Module): + self.conv1 = value + + def forward( + self, + input_features, + attention_mask=None, + head_mask=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + ): + r""" + Args: + input_features (`torch.LongTensor` of shape `(batch_size, feature_size, sequence_length)`): + Float values of mel features extracted from the raw speech waveform. Raw speech waveform can be + obtained by loading a `.flac` or `.wav` audio file into an array of type `List[float]` or a + `numpy.ndarray`, *e.g.* via the soundfile library (`pip install soundfile`). To prepare the array into + `input_features`, the [`AutoFeatureExtractor`] should be used for extracting the mel features, padding + and conversion into a tensor of type `torch.FloatTensor`. See [`~WhisperFeatureExtractor.__call__`] + attention_mask (`torch.Tensor`)`, *optional*): + Qwen2Audio does not support masking of the `input_features`, this argument is preserved for compatibility, + but it is not used. By default the silence in the input log mel spectrogram are ignored. + head_mask (`torch.Tensor` of shape `(encoder_layers, encoder_attention_heads)`, *optional*): + Mask to nullify selected heads of the attention modules. Mask values selected in `[0, 1]`: + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors + for more detail. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. + """ + + expected_seq_length = self.config.max_source_positions * self.conv1.stride[0] * self.conv2.stride[0] + if input_features.shape[-1] != expected_seq_length: + raise ValueError( + f"Qwen2Audio expects the mel input features to be of length {expected_seq_length}, but found {input_features.shape[-1]}. Make sure to pad the input mel features to {expected_seq_length}." + ) + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # Ignore copy + input_features = input_features.to(dtype=self.conv1.weight.dtype, device=self.conv1.weight.device) + + inputs_embeds = nn.functional.gelu(self.conv1(input_features)) + inputs_embeds = nn.functional.gelu(self.conv2(inputs_embeds)) + + inputs_embeds = inputs_embeds.permute(0, 2, 1) + embed_pos = self.embed_positions.weight + + hidden_states = inputs_embeds + embed_pos + hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) + + encoder_states = () if output_hidden_states else None + all_attentions = () if output_attentions else None + + # check if head_mask has a correct number of layers specified if desired + if head_mask is not None: + assert head_mask.size()[0] == ( + len(self.layers) + ), f"The head_mask should be specified for {len(self.layers)} layers, but it is for {head_mask.size()[0]}." + + for idx, encoder_layer in enumerate(self.layers): + if output_hidden_states: + encoder_states = encoder_states + (hidden_states,) + # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description) + to_drop = False + if self.training: + dropout_probability = torch.rand([]) + if dropout_probability < self.layerdrop: # skip the layer + to_drop = True + + # Ignore copy + if to_drop: + layer_outputs = (None, None) + else: + if self.gradient_checkpointing and self.training: + layer_outputs = self._gradient_checkpointing_func( + encoder_layer.__call__, + hidden_states, + attention_mask, + (head_mask[idx] if head_mask is not None else None), + output_attentions, + ) + else: + layer_outputs = encoder_layer( + hidden_states, + attention_mask, + layer_head_mask=(head_mask[idx] if head_mask is not None else None), + output_attentions=output_attentions, + ) + + hidden_states = layer_outputs[0] + + if output_attentions: + all_attentions = all_attentions + (layer_outputs[1],) + + # Ignore copy + hidden_states = hidden_states.permute(0, 2, 1) + hidden_states = self.avg_pooler(hidden_states) + hidden_states = hidden_states.permute(0, 2, 1) + + hidden_states = self.layer_norm(hidden_states) + if output_hidden_states: + encoder_states = encoder_states + (hidden_states,) + + if not return_dict: + return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None) + return BaseModelOutput( + last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions + ) + + # Ignore copy + def _get_feat_extract_output_lengths(self, input_lengths: torch.LongTensor): + """ + Computes the output length of the convolutional layers and the output length of the audio encoder + """ + input_lengths = (input_lengths - 1) // 2 + 1 + output_lengths = (input_lengths - 2) // 2 + 1 + return input_lengths, output_lengths + + +class Qwen2AudioMultiModalProjector(nn.Module): + def __init__(self, config: Qwen2AudioConfig): + super().__init__() + self.linear = nn.Linear(config.audio_config.d_model, config.text_config.hidden_size, bias=True) + + def forward(self, audio_features): + hidden_states = self.linear(audio_features) + return hidden_states + + +QWEN2AUDIO_INPUTS_DOCSTRING = r""" + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide + it. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + input_features (`torch.FloatTensor` of shape `(batch_size, feature_size, feature_sequence_length)`): + Float values mel features extracted from the raw speech waveform. Raw speech waveform can be obtained by + loading a `.flac` or `.wav` audio file into an array of type `List[float]` or a `numpy.ndarray`, *e.g.* via + the soundfile library (`pip install soundfile`). To prepare the array into `input_features`, the + [`AutoFeatureExtractor`] should be used for extracting the mel features, padding and conversion into a + tensor of type `torch.FloatTensor`. See [`~WhisperFeatureExtractor.__call__`] + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see + `past_key_values`). + + If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`] + and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more + information on the default strategy. + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + feature_attention_mask (`torch.Tensor` of shape `(batch_size, feature_sequence_length)`): + Mask to avoid performing attention on padding feature indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, + config.n_positions - 1]`. [What are position IDs?](../glossary#position-ids) + past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape + `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of shape + `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`. + + Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention + blocks) that can be used (see `past_key_values` input) to speed up sequential decoding. + + If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that + don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all + `decoder_input_ids` of shape `(batch_size, sequence_length)`. + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This + is useful if you want more control over how to convert `input_ids` indices into associated vectors than the + model's internal embedding lookup matrix. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see + `past_key_values`). + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned + tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for + more detail. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. +""" + + +@add_start_docstrings( + """The QWEN2AUDIO model which consists of a audio backbone and a language model.""", + QWEN2AUDIO_START_DOCSTRING, +) +class Qwen2AudioForConditionalGeneration(Qwen2AudioPreTrainedModel, GenerationMixin): + def __init__(self, config: Qwen2AudioConfig): + super().__init__(config) + self.audio_tower = AutoModel.from_config(config.audio_config) + + self.multi_modal_projector = Qwen2AudioMultiModalProjector(config) + self.vocab_size = config.text_config.vocab_size + self.language_model = AutoModelForCausalLM.from_config(config.text_config) + self.pad_token_id = self.config.pad_token_id if self.config.pad_token_id is not None else -1 + self._padding_side = "left" # set it to left by default, user can use setter to change padding_sides + self.post_init() + + @property + def padding_side(self): + return self._padding_side + + @padding_side.setter + def padding_side(self, padding_side: str): + if padding_side not in ["left", "right"]: + raise ValueError(f"{padding_side} is not `left` or `right`.") + self._padding_side = padding_side + + # Copied from transformers.models.llava.modeling_llava.LlavaForConditionalGeneration.get_input_embeddings + def get_input_embeddings(self): + return self.language_model.get_input_embeddings() + + # Copied from transformers.models.llava.modeling_llava.LlavaForConditionalGeneration.set_input_embeddings + def set_input_embeddings(self, value): + self.language_model.set_input_embeddings(value) + + # Copied from transformers.models.llava.modeling_llava.LlavaForConditionalGeneration.get_output_embeddings + def get_output_embeddings(self): + return self.language_model.get_output_embeddings() + + # Copied from transformers.models.llava.modeling_llava.LlavaForConditionalGeneration.set_output_embeddings + def set_output_embeddings(self, new_embeddings): + self.language_model.set_output_embeddings(new_embeddings) + + # Copied from transformers.models.llava.modeling_llava.LlavaForConditionalGeneration.set_decoder + def set_decoder(self, decoder): + self.language_model.set_decoder(decoder) + + # Copied from transformers.models.llava.modeling_llava.LlavaForConditionalGeneration.get_decoder + def get_decoder(self): + return self.language_model.get_decoder() + + # Copied from transformers.models.llava.modeling_llava.LlavaForConditionalGeneration.tie_weights + def tie_weights(self): + return self.language_model.tie_weights() + + # Copied from transformers.models.llava.modeling_llava.LlavaForConditionalGeneration.resize_token_embeddings + def resize_token_embeddings(self, new_num_tokens: Optional[int] = None, pad_to_multiple_of=None) -> nn.Embedding: + model_embeds = self.language_model.resize_token_embeddings(new_num_tokens, pad_to_multiple_of) + # update vocab size + self.config.text_config.vocab_size = model_embeds.num_embeddings + self.vocab_size = model_embeds.num_embeddings + return model_embeds + + def _merge_input_ids_with_audio_features( + self, audio_features, num_audio_tokens, inputs_embeds, input_ids, attention_mask, labels + ): + """ + Merge input_ids with with audio features into final embeddings + + Args: + audio_features (`torch.Tensor` of shape `(num_audios, max_audio_tokens, embed_dim)`): + All audio vectors of all audios in the batch + num_audio_tokens (`torch.LongTensor` of shape `(num_audios)`): + The length of audio embeddings of each audio as stacked in `audio_features` + inputs_embeds (`torch.Tensor` of shape `(batch_size, sequence_length, embed_dim)`): + Token embeddings before merging with audio embeddings + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Input_ids of tokens, possibly filled with audio token + attention_mask (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Mask to avoid performing attention on padding token indices. + labels (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*) + labels need to be recalculated to support training (if provided) + Returns: + final_embedding, final_attention_mask, final_labels, position_ids, final_input_ids + + Explanation: + each audio has variable length embeddings, with length specified by num_audio_tokens + audio_features is concatenation of all audio embed vectors + task: fill each <|AUDIO|> with the correct number of audio embeddings + Example: + X (5 tokens), Y (3 tokens), Z (8 tokens) + X, Y are in the same sequence (in-context learning) + if right padding + input_ids: [ + a b c d e f X g h i j k Y l m + o p q r Z s t u v _ _ _ _ _ _ + ] + input_ids should be: [ + a b c d e f X X X X X g h i j k Y Y Y l m + o p q r Z Z Z Z Z Z Z Z s t u v _ _ _ _ _ + ] + labels should be: [ + a b c d e f _ _ _ _ _ g h i j k _ _ _ l m + o p q r _ _ _ _ _ _ _ _ s t u v _ _ _ _ _ + ] + elif left padding + input_ids: [ + a b c d e f X g h i j k Y l m + _ _ _ _ _ _ o p q r Z s t u v + ] + input_ids should be: [ + a b c d e f X X X X X g h i j k Y Y Y l m + _ _ _ _ _ o p q r Z Z Z Z Z Z Z Z s t u v + ] + labels should be: [ + a b c d e f _ _ _ _ _ g h i j k _ _ _ l m + _ _ _ _ _ o p q r _ _ _ _ _ _ _ _ s t u v + ] + Edge cases: + * If tokens are same but audio token sizes are different, then cannot infer left or right padding + ```python + url1 = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2-Audio/audio/glass-breaking-151256.mp3" + audio1, _ = librosa.load(BytesIO(urlopen(url1).read()), sr=processor.feature_extractor.sampling_rate) + url2 = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2-Audio/audio/f2641_0_throatclearing.wav" + audio2, _ = librosa.load(BytesIO(urlopen(url2).read()), sr=processor.feature_extractor.sampling_rate) + prompts = [ + "[INST] <|AUDIO|>\nWhat is that in this audio? [/INST]", + "[INST] <|AUDIO|>\nWhat is that in this audio? [/INST]", + ] + inputs = processor(text=prompts, audios=[audio1, audio2], return_tensors='pt', padding=True).to("cuda") + audio1 has 101 tokens, while audio2 has 72 tokens + ``` + + input_ids: [ + a b c d X g h + i j Y k l m n + ] + where X is 3 tokens while Y is 5, this mean after merge + if left-padding (batched generation) + input_ids should be: [ + _ _ a b c d X X X g h + i j Y Y Y Y Y k l m n + ] + elif (right padding) (training) + input_ids should be: [ + a b c d X X X g h _ _ + i j Y Y Y Y Y k l m n + ] + """ + num_audios, max_audio_tokens, embed_dim = audio_features.shape + audio_features_mask = torch.arange(max_audio_tokens).expand(num_audios, max_audio_tokens).to( + num_audio_tokens.device + ) < num_audio_tokens.unsqueeze(1) + masked_audio_features = audio_features[audio_features_mask].view(-1, embed_dim) + batch_size, sequence_length = input_ids.shape + _left_padding = torch.any(attention_mask[:, 0] == 0) + _right_padding = torch.any(attention_mask[:, -1] == 0) + + left_padding = True + if batch_size > 1: + if _left_padding and not _right_padding: + left_padding = True + elif not _left_padding and _right_padding: + left_padding = False + elif not _left_padding and not _right_padding: + # both side is 1, so cannot tell + left_padding = self.padding_side == "left" + else: + # invalid attention_mask + raise ValueError(f"both side of attention_mask has zero, invalid. {attention_mask}") + + # 1. Create a mask to know where special audio tokens are + special_audio_token_mask = input_ids == self.config.audio_token_index + num_special_audio_tokens = torch.sum(special_audio_token_mask, dim=-1) + + # In case the Audio model or the Language model has been offloaded to CPU, we need to manually + # set the corresponding tensors into their correct target device. + target_device = inputs_embeds.device + attention_mask = attention_mask.to(target_device) + input_ids = input_ids.to(target_device) + num_audio_tokens = num_audio_tokens.to(target_device) + batch_indices, non_audio_indices = torch.where( + (input_ids != self.config.audio_token_index) & (attention_mask == 1) + ) + + # 2. Compute the positions where text should be written + # Calculate new positions for text tokens in merged audio-text sequence. + # `special_audio_token_mask` identifies audio tokens. Each audio token will be replaced by `audio_feat_lengths - 1` text tokens. + # `torch.cumsum` computes how each audio token shifts subsequent text token positions. + token_placeholder_num = torch.zeros_like(input_ids) + token_placeholder_num[special_audio_token_mask] = num_audio_tokens.long() - 1 + token_placeholder_num = token_placeholder_num + 1 + new_token_positions = torch.cumsum(token_placeholder_num, -1) - 1 + max_token_num = token_placeholder_num.sum(-1).max() + nb_audio_pad = max_token_num - 1 - new_token_positions[:, -1] + if left_padding: + new_token_positions += nb_audio_pad[:, None] # offset for left padding + text_to_overwrite = new_token_positions[batch_indices, non_audio_indices] + batch_indices, non_audio_indices, text_to_overwrite = ( + batch_indices.to(target_device), + non_audio_indices.to(target_device), + text_to_overwrite.to(target_device), + ) + + # 3. Create the full embedding, already padded to the maximum position + final_embedding = torch.zeros( + batch_size, max_token_num, embed_dim, dtype=inputs_embeds.dtype, device=inputs_embeds.device + ) + final_attention_mask = torch.zeros( + batch_size, max_token_num, dtype=attention_mask.dtype, device=inputs_embeds.device + ) + final_input_ids = torch.full( + (batch_size, max_token_num), self.pad_token_id, dtype=input_ids.dtype, device=inputs_embeds.device + ) + + # 4. Fill the embeddings based on the mask. If we have ["hey" "