From e6980cb6a1b604efb0f97000e518aef7a8d4977c Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Fri, 18 Apr 2025 08:05:11 +0000 Subject: [PATCH 01/22] Update `setup.py` and examples to Transformers v4.51 --- .../run_audio_classification.py | 2 +- .../contrastive-image-text/run_bridgetower.py | 2 +- examples/contrastive-image-text/run_clip.py | 2 +- .../run_image_classification.py | 2 +- examples/language-modeling/run_clm.py | 2 +- examples/language-modeling/run_mlm.py | 2 +- examples/question-answering/run_qa.py | 2 +- examples/question-answering/run_seq2seq_qa.py | 2 +- .../run_speech_recognition_ctc.py | 4 +- .../run_speech_recognition_seq2seq.py | 2 +- .../unconditional_image_generation.py | 2 +- examples/summarization/run_summarization.py | 2 +- examples/text-classification/run_glue.py | 2 +- examples/translation/run_translation.py | 2 +- setup.py | 2 +- .../example_diff/run_audio_classification.txt | 79 ++-- tests/example_diff/run_clip.txt | 52 +-- tests/example_diff/run_clm.txt | 128 ++++-- tests/example_diff/run_generation.txt | 380 +++++++++++------- tests/example_diff/run_glue.txt | 48 ++- .../example_diff/run_image_classification.txt | 38 +- tests/example_diff/run_mlm.txt | 58 ++- tests/example_diff/run_qa.txt | 51 +-- tests/example_diff/run_seq2seq_qa.txt | 48 ++- .../run_speech_recognition_ctc.txt | 108 +++-- .../run_speech_recognition_seq2seq.txt | 64 +-- tests/example_diff/run_summarization.txt | 114 +++--- tests/example_diff/run_translation.txt | 52 ++- 28 files changed, 752 insertions(+), 500 deletions(-) diff --git a/examples/audio-classification/run_audio_classification.py b/examples/audio-classification/run_audio_classification.py index d3b553a503..073b8ad577 100644 --- a/examples/audio-classification/run_audio_classification.py +++ b/examples/audio-classification/run_audio_classification.py @@ -46,7 +46,7 @@ def check_optimum_habana_min_version(*a, **b): logger = logging.getLogger(__name__) # Will error if the minimal version of Transformers and Optimum Habana are not installed. Remove at your own risks. -check_min_version("4.49.0") +check_min_version("4.51.0") check_optimum_habana_min_version("1.18.0.dev0") require_version("datasets>=1.14.0", "To fix: pip install -r examples/pytorch/audio-classification/requirements.txt") diff --git a/examples/contrastive-image-text/run_bridgetower.py b/examples/contrastive-image-text/run_bridgetower.py index a884343fab..b1534d8e82 100644 --- a/examples/contrastive-image-text/run_bridgetower.py +++ b/examples/contrastive-image-text/run_bridgetower.py @@ -58,7 +58,7 @@ def check_optimum_habana_min_version(*a, **b): logger = logging.getLogger(__name__) # Will error if the minimal version of Transformers and Optimum Habana are not installed. Remove at your own risks. -check_min_version("4.49.0") +check_min_version("4.51.0") check_optimum_habana_min_version("1.18.0.dev0") require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/contrastive-image-text/requirements.txt") diff --git a/examples/contrastive-image-text/run_clip.py b/examples/contrastive-image-text/run_clip.py index ddd4b2c4fa..f5b657d6e3 100644 --- a/examples/contrastive-image-text/run_clip.py +++ b/examples/contrastive-image-text/run_clip.py @@ -61,7 +61,7 @@ def check_optimum_habana_min_version(*a, **b): logger = logging.getLogger(__name__) # Will error if the minimal version of Transformers and Optimum Habana are not installed. Remove at your own risks. -check_min_version("4.49.0") +check_min_version("4.51.0") check_optimum_habana_min_version("1.18.0.dev0") require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/contrastive-image-text/requirements.txt") diff --git a/examples/image-classification/run_image_classification.py b/examples/image-classification/run_image_classification.py index d8e1506014..a82428eb94 100644 --- a/examples/image-classification/run_image_classification.py +++ b/examples/image-classification/run_image_classification.py @@ -64,7 +64,7 @@ def check_optimum_habana_min_version(*a, **b): logger = logging.getLogger(__name__) # Will error if the minimal version of Transformers and Optimum Habana are not installed. Remove at your own risks. -check_min_version("4.49.0") +check_min_version("4.51.0") check_optimum_habana_min_version("1.18.0.dev0") require_version("datasets>=2.14.0", "To fix: pip install -r examples/pytorch/image-classification/requirements.txt") diff --git a/examples/language-modeling/run_clm.py b/examples/language-modeling/run_clm.py index ef5e3d6188..934a74292f 100644 --- a/examples/language-modeling/run_clm.py +++ b/examples/language-modeling/run_clm.py @@ -62,7 +62,7 @@ def check_optimum_habana_min_version(*a, **b): logger = logging.getLogger(__name__) # Will error if the minimal version of Transformers and Optimum Habana are not installed. Remove at your own risks. -check_min_version("4.49.0") +check_min_version("4.51.0") check_optimum_habana_min_version("1.18.0.dev0") require_version("datasets>=2.14.0", "To fix: pip install -r examples/pytorch/language-modeling/requirements.txt") diff --git a/examples/language-modeling/run_mlm.py b/examples/language-modeling/run_mlm.py index 0d05e7a35e..98741f2b4b 100644 --- a/examples/language-modeling/run_mlm.py +++ b/examples/language-modeling/run_mlm.py @@ -61,7 +61,7 @@ def check_optimum_habana_min_version(*a, **b): logger = logging.getLogger(__name__) # Will error if the minimal version of Transformers and Optimum Habana are not installed. Remove at your own risks. -check_min_version("4.49.0") +check_min_version("4.51.0") check_optimum_habana_min_version("1.18.0.dev0") require_version("datasets>=2.14.0", "To fix: pip install -r examples/pytorch/language-modeling/requirements.txt") diff --git a/examples/question-answering/run_qa.py b/examples/question-answering/run_qa.py index db5dab8768..c4b3224269 100644 --- a/examples/question-answering/run_qa.py +++ b/examples/question-answering/run_qa.py @@ -60,7 +60,7 @@ def check_optimum_habana_min_version(*a, **b): logger = logging.getLogger(__name__) # Will error if the minimal version of Transformers and Optimum Habana are not installed. Remove at your own risks. -check_min_version("4.49.0") +check_min_version("4.51.0") check_optimum_habana_min_version("1.18.0.dev0") require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/question-answering/requirements.txt") diff --git a/examples/question-answering/run_seq2seq_qa.py b/examples/question-answering/run_seq2seq_qa.py index 8899f11c17..374ec915ca 100644 --- a/examples/question-answering/run_seq2seq_qa.py +++ b/examples/question-answering/run_seq2seq_qa.py @@ -56,7 +56,7 @@ def check_optimum_habana_min_version(*a, **b): logger = logging.getLogger(__name__) # Will error if the minimal version of Transformers and Optimum Habana are not installed. Remove at your own risks. -check_min_version("4.49.0") +check_min_version("4.51.0") check_optimum_habana_min_version("1.18.0.dev0") require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/question-answering/requirements.txt") diff --git a/examples/speech-recognition/run_speech_recognition_ctc.py b/examples/speech-recognition/run_speech_recognition_ctc.py index 74c0f814bd..0828d8cb4f 100644 --- a/examples/speech-recognition/run_speech_recognition_ctc.py +++ b/examples/speech-recognition/run_speech_recognition_ctc.py @@ -59,7 +59,7 @@ def check_optimum_habana_min_version(*a, **b): logger = logging.getLogger(__name__) # Will error if the minimal version of Transformers and Optimum Habana are not installed. Remove at your own risks. -check_min_version("4.49.0") +check_min_version("4.51.0") check_optimum_habana_min_version("1.18.0.dev0") require_version("datasets>=1.18.0", "To fix: pip install -r examples/pytorch/speech-recognition/requirements.txt") @@ -325,7 +325,7 @@ class DataCollatorCTCWithPadding: Data collator that will dynamically pad the inputs received. Args: processor (:class:`~transformers.AutoProcessor`) - The processor used for proccessing the data. + The processor used for processing the data. padding (:obj:`bool`, :obj:`str` or :class:`~transformers.tokenization_utils_base.PaddingStrategy`, `optional`, defaults to :obj:`True`): Select a strategy to pad the returned sequences (according to the model's padding side and padding index) among: diff --git a/examples/speech-recognition/run_speech_recognition_seq2seq.py b/examples/speech-recognition/run_speech_recognition_seq2seq.py index 61516cce4c..f52bd73887 100755 --- a/examples/speech-recognition/run_speech_recognition_seq2seq.py +++ b/examples/speech-recognition/run_speech_recognition_seq2seq.py @@ -55,7 +55,7 @@ def check_optimum_habana_min_version(*a, **b): # Will error if the minimal version of Transformers is not installed. Remove at your own risks. -check_min_version("4.49.0") +check_min_version("4.51.0") check_optimum_habana_min_version("1.18.0.dev0") require_version("datasets>=1.18.0", "To fix: pip install -r examples/pytorch/speech-recognition/requirements.txt") diff --git a/examples/stable-diffusion/unconditional_image_generation.py b/examples/stable-diffusion/unconditional_image_generation.py index 8504a2fcb9..979f60b838 100755 --- a/examples/stable-diffusion/unconditional_image_generation.py +++ b/examples/stable-diffusion/unconditional_image_generation.py @@ -19,7 +19,7 @@ def check_optimum_habana_min_version(*a, **b): return () -check_min_version("4.49.0") +check_min_version("4.51.0") check_optimum_habana_min_version("1.18.0.dev0") # Setup logging diff --git a/examples/summarization/run_summarization.py b/examples/summarization/run_summarization.py index 6cad764d56..87e5faa9d1 100755 --- a/examples/summarization/run_summarization.py +++ b/examples/summarization/run_summarization.py @@ -64,7 +64,7 @@ def check_optimum_habana_min_version(*a, **b): logger = logging.getLogger(__name__) # Will error if the minimal version of Transformers and Optimum Habana are not installed. Remove at your own risks. -check_min_version("4.49.0") +check_min_version("4.51.0") check_optimum_habana_min_version("1.18.0.dev0") require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/summarization/requirements.txt") diff --git a/examples/text-classification/run_glue.py b/examples/text-classification/run_glue.py index 76a19e84c2..b4bbc78364 100755 --- a/examples/text-classification/run_glue.py +++ b/examples/text-classification/run_glue.py @@ -57,7 +57,7 @@ def check_optimum_habana_min_version(*a, **b): logger = logging.getLogger(__name__) # Will error if the minimal version of Transformers and Optimum Habana are not installed. Remove at your own risks. -check_min_version("4.49.0") +check_min_version("4.51.0") check_optimum_habana_min_version("1.18.0.dev0") require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/text-classification/requirements.txt") diff --git a/examples/translation/run_translation.py b/examples/translation/run_translation.py index cddcd4e059..c1d8a07d1d 100644 --- a/examples/translation/run_translation.py +++ b/examples/translation/run_translation.py @@ -62,7 +62,7 @@ def check_optimum_habana_min_version(*a, **b): logger = logging.getLogger(__name__) # Will error if the minimal version of Transformers and Optimum Habana are not installed. Remove at your own risks. -check_min_version("4.49.0") +check_min_version("4.51.0") check_optimum_habana_min_version("1.18.0.dev0") require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/translation/requirements.txt") diff --git a/setup.py b/setup.py index ee22ddfae6..7b3fe3b5dc 100644 --- a/setup.py +++ b/setup.py @@ -29,7 +29,7 @@ INSTALL_REQUIRES = [ - "transformers >= 4.49.0, < 4.50.0", + "transformers[hf_xet] >= 4.51.0, < 4.52.0", "optimum", "torch", "accelerate >= 1.5.0", diff --git a/tests/example_diff/run_audio_classification.txt b/tests/example_diff/run_audio_classification.txt index ea9a7e34ef..b0055ed170 100644 --- a/tests/example_diff/run_audio_classification.txt +++ b/tests/example_diff/run_audio_classification.txt @@ -1,9 +1,11 @@ -20d19 +1a2 +> # coding=utf-8 +19d19 < import warnings -28,29d26 +27,28d26 < from datasets import DatasetDict, load_dataset < -31,39c28,29 +30,38c28,29 < from transformers import ( < AutoConfig, < AutoFeatureExtractor, @@ -16,7 +18,7 @@ --- > from datasets import DatasetDict, load_dataset > from transformers import AutoConfig, AutoFeatureExtractor, AutoModelForAudioClassification, HfArgumentParser -43a34,44 +42a34,44 > from optimum.habana import GaudiConfig, GaudiTrainer, GaudiTrainingArguments > from optimum.habana.utils import set_seed > @@ -28,19 +30,36 @@ > def check_optimum_habana_min_version(*a, **b): > return () > -47,48c48,50 +46,47c48,50 < # Will error if the minimal version of Transformers is not installed. Remove at your own risks. -< check_min_version("4.46.0.dev0") +< check_min_version("4.52.0.dev0") --- > # Will error if the minimal version of Transformers and Optimum Habana are not installed. Remove at your own risks. -> check_min_version("4.45.0") -> check_optimum_habana_min_version("1.14.0.dev0") -174,176d175 +> check_min_version("4.51.0") +> check_optimum_habana_min_version("1.18.0.dev0") +173,175d175 < freeze_feature_extractor: Optional[bool] = field( < default=None, metadata={"help": "Whether to freeze the feature extractor layers of the model."} < ) -182,196d180 -< def __post_init__(self): +179a180,196 +> use_flash_attention: bool = field( +> default=False, metadata={"help": "Whether to use Habana flash attention for fine-tuning"} +> ) +> flash_attention_recompute: bool = field( +> default=False, +> metadata={ +> "help": "Whether to enable recompute in Habana flash attention for fine-tuning." +> " It is applicable only when use_flash_attention is True." +> }, +> ) +> flash_attention_fast_softmax: bool = field( +> default=False, +> metadata={ +> "help": "Whether to use fast softmax for Habana flash attention." +> " It is applicable only when use_flash_attention is True." +> }, +> ) +182,194c199,206 < if not self.freeze_feature_extractor and self.freeze_feature_encoder: < warnings.warn( < "The argument `--freeze_feature_extractor` is deprecated and " @@ -54,12 +73,20 @@ < "should not be used in combination with `--freeze_feature_encoder`. " < "Only make use of `--freeze_feature_encoder`." < ) -< -203c187 +--- +> if self.use_flash_attention: +> os.environ["USE_FLASH_ATTENTION"] = "1" +> if self.flash_attention_recompute: +> assert self.use_flash_attention, "flash_attention_recompute is set, but use_flash_attention is not" +> os.environ["FLASH_ATTENTION_RECOMPUTE"] = "1" +> if self.flash_attention_fast_softmax: +> assert self.use_flash_attention, "flash_attention_fast_softmax is set, but use_flash_attention is not" +> os.environ["FLASH_ATTENTION_FAST_SOFTMAX"] = "1" +202c214 < parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)) --- > parser = HfArgumentParser((ModelArguments, DataTrainingArguments, GaudiTrainingArguments)) -231a216,222 +230a243,249 > gaudi_config = GaudiConfig.from_pretrained( > training_args.gaudi_config_name, > cache_dir=model_args.cache_dir, @@ -67,22 +94,22 @@ > token=model_args.token, > ) > -232a224 +231a251 > mixed_precision = training_args.bf16 or gaudi_config.use_torch_autocast -234,235c226,228 +233,234c253,255 < f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}, " < + f"distributed training: {training_args.parallel_mode.value == 'distributed'}, 16-bits training: {training_args.fp16}" --- > f"Process rank: {training_args.local_rank}, device: {training_args.device}, " > + f"distributed training: {training_args.parallel_mode.value == 'distributed'}, " > + f"mixed-precision training: {mixed_precision}" -304a298,300 +303a325,327 > # Max input length > max_length = int(round(feature_extractor.sampling_rate * data_args.max_length_seconds)) > -309a306 +308a333 > -315c312,318 +314c339,345 < inputs = feature_extractor(subsampled_wavs, sampling_rate=feature_extractor.sampling_rate) --- > inputs = feature_extractor( @@ -92,7 +119,7 @@ > padding="max_length", > truncation=True, > ) -324c327,333 +323c354,360 < inputs = feature_extractor(wavs, sampling_rate=feature_extractor.sampling_rate) --- > inputs = feature_extractor( @@ -102,19 +129,17 @@ > padding="max_length", > truncation=True, > ) -370,371c379,380 +356a394 +> attn_implementation=training_args.attn_implementation, +369,370c407,408 < # freeze the convolutional waveform encoder < if model_args.freeze_feature_encoder: --- > # freeze the convolutional waveform encoder if supported by model > if hasattr(model, "freeze_feature_encoder") and model_args.freeze_feature_encoder: -391c400 +390c428 < trainer = Trainer( --- > trainer = GaudiTrainer( -392a402 +391a430 > gaudi_config=gaudi_config, -397c407 -< processing_class=feature_extractor, ---- -> tokenizer=feature_extractor, diff --git a/tests/example_diff/run_clip.txt b/tests/example_diff/run_clip.txt index ce6b6fd9ba..be2fe6baca 100644 --- a/tests/example_diff/run_clip.txt +++ b/tests/example_diff/run_clip.txt @@ -1,17 +1,19 @@ -18d17 +1a2 +> # coding=utf-8 +17d17 < -32a32 +31a32 > import transformers -33a34 +32a34 > from habana_dataloader_trainer import HabanaDataloaderTrainer -38,39d38 +37,38d38 < < import transformers -45,47d43 +44,46d43 < Trainer, < TrainingArguments, < set_seed, -52a49,59 +51a49,59 > from optimum.habana import GaudiConfig, GaudiTrainer, GaudiTrainingArguments > from optimum.habana.utils import set_seed > @@ -23,22 +25,22 @@ > def check_optimum_habana_min_version(*a, **b): > return () > -56,57c63,65 +55,56c63,65 < # Will error if the minimal version of Transformers is not installed. Remove at your own risks. -< check_min_version("4.46.0.dev0") +< check_min_version("4.52.0.dev0") --- > # Will error if the minimal version of Transformers and Optimum Habana are not installed. Remove at your own risks. -> check_min_version("4.45.0") -> check_optimum_habana_min_version("1.14.0.dev0") -181a190,192 +> check_min_version("4.51.0") +> check_optimum_habana_min_version("1.18.0.dev0") +176a186,188 > mediapipe_dataloader: bool = field( > default=False, metadata={"help": "Turn on MediaPipe hardware-based accelerated data loading."} > ) -240c251 +232c244 < parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)) --- > parser = HfArgumentParser((ModelArguments, DataTrainingArguments, GaudiTrainingArguments)) -268a280,286 +260a273,279 > gaudi_config = GaudiConfig.from_pretrained( > training_args.gaudi_config_name, > cache_dir=model_args.cache_dir, @@ -46,18 +48,18 @@ > token=model_args.token, > ) > -269a288 +261a281 > mixed_precision = training_args.bf16 or gaudi_config.use_torch_autocast -271,272c290,292 +263,264c283,285 < f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}, " < + f"distributed training: {training_args.parallel_mode.value == 'distributed'}, 16-bits training: {training_args.fp16}" --- > f"Process rank: {training_args.local_rank}, device: {training_args.device}, " > + f"distributed training: {training_args.parallel_mode.value == 'distributed'}, " > + f"mixed-precision training: {mixed_precision}" -420d439 +407d427 < image_transformations = torch.jit.script(image_transformations) -467,468c486,494 +454,455c474,482 < # Transform images on the fly as doing it on the whole dataset takes too much time. < train_dataset.set_transform(transform_images) --- @@ -70,7 +72,7 @@ > else: > # Transform images on the fly as doing it on the whole dataset takes too much time. > train_dataset.set_transform(transform_images) -490,491c516,524 +477,478c504,512 < # Transform images on the fly as doing it on the whole dataset takes too much time. < eval_dataset.set_transform(transform_images) --- @@ -83,20 +85,10 @@ > else: > # Transform images on the fly as doing it on the whole dataset takes too much time. > eval_dataset.set_transform(transform_images) -514a548,556 -> if data_args.mediapipe_dataloader: -> test_dataset.image_mean = image_processor.image_mean -> test_dataset.image_std = image_processor.image_std -> test_dataset.text_max_length = data_args.max_seq_length -> test_dataset.image_resize = config.vision_config.image_size -> test_dataset.transform_func = transform_images -> else: -> # Transform images on the fly as doing it on the whole dataset takes too much time. -> test_dataset.set_transform(transform_images) -517c559,560 +481c515,516 < trainer = Trainer( --- > trainer_cls = HabanaDataloaderTrainer if data_args.mediapipe_dataloader else GaudiTrainer > trainer = trainer_cls( -518a562 +482a518 > gaudi_config=gaudi_config, diff --git a/tests/example_diff/run_clm.txt b/tests/example_diff/run_clm.txt index 392843941e..8f50b571ba 100644 --- a/tests/example_diff/run_clm.txt +++ b/tests/example_diff/run_clm.txt @@ -1,32 +1,33 @@ -3c3 +2c2,3 < # Copyright 2020 The HuggingFace Inc. team. All rights reserved. --- +> # coding=utf-8 > # Copyright 2022 The HuggingFace Inc. team. All rights reserved. -17,19c17,18 +16,18c17,18 < Fine-tuning the library models for causal language modeling (GPT, GPT-2, CTRL, ...) on a text file or a dataset. < < Here is the full list of checkpoints on the hub that can be fine-tuned by this script: --- > Training the library models for causal language modeling (GPT, GPT-2, CTRL, ...) on a text file or a dataset. > Here is the full list of checkpoints on the hub that can be trained by this script: -35,36d33 +34,35d33 < from datasets import load_dataset < -37a35 +36a35 > from datasets import load_dataset -45,46d42 +44,45d42 < Trainer, < TrainingArguments, -48,49d43 +47,48d43 < is_torch_xla_available, < set_seed, -55a50,51 +54a50,51 > from optimum.habana import GaudiConfig, GaudiTrainer, GaudiTrainingArguments > from optimum.habana.utils import set_seed -57,58d52 +56,57d52 < # Will error if the minimal version of Transformers is not installed. Remove at your own risks. -< check_min_version("4.46.0.dev0") -60c54,60 +< check_min_version("4.52.0.dev0") +59c54,60 < require_version("datasets>=2.14.0", "To fix: pip install -r examples/pytorch/language-modeling/requirements.txt") --- > try: @@ -36,19 +37,19 @@ > def check_optimum_habana_min_version(*a, **b): > return () > -63a64,69 +62a64,69 > # Will error if the minimal version of Transformers and Optimum Habana are not installed. Remove at your own risks. -> check_min_version("4.45.0") -> check_optimum_habana_min_version("1.14.0.dev0") +> check_min_version("4.51.0") +> check_optimum_habana_min_version("1.18.0.dev0") > > require_version("datasets>=2.14.0", "To fix: pip install -r examples/pytorch/language-modeling/requirements.txt") > -79c85,86 +78c85,86 < "The model checkpoint for weights initialization. Don't set if you want to train a model from scratch." --- > "The model checkpoint for weights initialization. Don't set it if you want to train a model from" > " scratch." -142a150,158 +141a150,184 > use_cache: bool = field( > default=True, > metadata={ @@ -58,24 +59,50 @@ > ) > }, > ) -148c164 +> attn_softmax_bf16: bool = field( +> default=False, +> metadata={"help": ("Whether to run attention softmax layer in bf16 precision for fine-tuning.")}, +> ) +> use_flash_attention: bool = field( +> default=False, +> metadata={"help": ("Whether to use Habana flash attention for fine-tuning.")}, +> ) +> flash_attention_recompute: bool = field( +> default=False, +> metadata={ +> "help": ( +> "Whether to enable recompute in Habana flash attention for fine-tuning." +> " It is applicable only when use_flash_attention is True." +> ) +> }, +> ) +> flash_attention_causal_mask: bool = field( +> default=False, +> metadata={ +> "help": ( +> "Whether to enable causal mask in Habana flash attention for fine-tuning." +> " It is applicable only when use_flash_attention is True." +> ) +> }, +> ) +147c190 < "set True will benefit LLM loading time and RAM consumption." --- > "Setting it to True will benefit LLM loading time and RAM consumption." -195c211,212 +194c237,238 < streaming: bool = field(default=False, metadata={"help": "Enable streaming mode"}) --- > > streaming: bool = field(default=False, metadata={"help": "Enable streaming mode."}) -221a239,241 +220a265,267 > save_last_ckpt: bool = field( > default=True, metadata={"help": "Whether to save checkpoint at the end of the training."} > ) -243c263 +242c289 < parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)) --- > parser = HfArgumentParser((ModelArguments, DataTrainingArguments, GaudiTrainingArguments)) -272a293,299 +271a319,325 > gaudi_config = GaudiConfig.from_pretrained( > training_args.gaudi_config_name, > cache_dir=model_args.cache_dir, @@ -83,38 +110,63 @@ > token=model_args.token, > ) > -273a301 +272a327 > mixed_precision = training_args.bf16 or gaudi_config.use_torch_autocast -275,276c303,305 +274,275c329,331 < f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}, " < + f"distributed training: {training_args.parallel_mode.value == 'distributed'}, 16-bits training: {training_args.fp16}" --- > f"Process rank: {training_args.local_rank}, device: {training_args.device}, " > + f"distributed training: {training_args.parallel_mode.value == 'distributed'}, " > + f"mixed-precision training: {mixed_precision}" -390a420 +389a446 > "use_cache": False if training_args.gradient_checkpointing else model_args.use_cache, -486a517 +402a460,463 +> # Note that chatglm2/3 has float16 dtype from config.json, and on Gaudi we need to use bfloat16. +> if config.model_type == "chatglm": +> config.torch_dtype = torch.bfloat16 +> +425a487,491 +> # workaraund for https://github.com/huggingface/transformers/issues/36258 +> # TODO: remove after fix is avalible in a release version of `transformers`` +> if torch_dtype is None: +> torch_dtype = getattr(config, "torch_dtype", None) > -550a582,585 +444,446c510,522 +< embedding_size = model.get_input_embeddings().weight.shape[0] +< if len(tokenizer) > embedding_size: +< model.resize_token_embeddings(len(tokenizer)) +--- +> # We need to skip this test for baichuan and chatglm pretrain +> if config.model_type not in ("baichuan", "chatglm"): +> embedding_size = model.get_input_embeddings().weight.shape[0] +> if len(tokenizer) > embedding_size: +> model.resize_token_embeddings(len(tokenizer)) +> +> # We need to add these fused kernels config +> if model_args.attn_softmax_bf16: +> model.generation_config.attn_softmax_bf16 = True +> if model_args.use_flash_attention: +> model.generation_config.use_flash_attention = True +> model.generation_config.flash_attention_recompute = model_args.flash_attention_recompute +> model.generation_config.flash_attention_causal_mask = model_args.flash_attention_causal_mask +485a562 +> +549a627,630 > > def tensor_mapper(x): > return {i: torch.tensor(x[i], dtype=torch.int32) for i in x} > -553a589,590 +552a634,635 > if training_args.resume_from_checkpoint is not None and training_args.resume_from_checkpoint != "": > train_dataset = train_dataset.map(tensor_mapper) -584c621 +583c666 < trainer = Trainer( --- > trainer = GaudiTrainer( -585a623 +584a668 > gaudi_config=gaudi_config, -589c627 -< processing_class=tokenizer, ---- -> tokenizer=tokenizer, -592,595c630,631 +591,594c675,676 < compute_metrics=compute_metrics if training_args.do_eval and not is_torch_xla_available() else None, < preprocess_logits_for_metrics=preprocess_logits_for_metrics < if training_args.do_eval and not is_torch_xla_available() @@ -122,12 +174,12 @@ --- > compute_metrics=compute_metrics if training_args.do_eval else None, > preprocess_logits_for_metrics=preprocess_logits_for_metrics if training_args.do_eval else None, -606c642,643 +605c687,688 < trainer.save_model() # Saves the tokenizer too for easy upload --- > if data_args.save_last_ckpt: > trainer.save_model() # Saves the tokenizer too for easy upload -610,613c647,653 +609,612c692,698 < max_train_samples = ( < data_args.max_train_samples if data_args.max_train_samples is not None else len(train_dataset) < ) @@ -140,9 +192,9 @@ > data_args.max_train_samples if data_args.max_train_samples is not None else len(train_dataset) > ) > metrics["train_samples"] = min(max_train_samples, len(train_dataset)) -622d661 +621d706 < -625,626c664,669 +624,625c709,714 < max_eval_samples = data_args.max_eval_samples if data_args.max_eval_samples is not None else len(eval_dataset) < metrics["eval_samples"] = min(max_eval_samples, len(eval_dataset)) --- @@ -152,7 +204,7 @@ > ) > metrics["eval_samples"] = min(max_eval_samples, len(eval_dataset)) > -649,653d691 +648,652d736 < < < def _mp_fn(index): diff --git a/tests/example_diff/run_generation.txt b/tests/example_diff/run_generation.txt index e7ad521192..31caa6fa2a 100644 --- a/tests/example_diff/run_generation.txt +++ b/tests/example_diff/run_generation.txt @@ -1,27 +1,25 @@ -17c17,19 +1a2 +> # coding=utf-8 +16c17,19 < """Conditional text generation with the auto-regressive models of the library (GPT/GPT-2/CTRL/Transformer-XL/XLNet)""" --- > """ > Conditional text generation on Habana Gaudi/Gaudi2. > """ -20c22 +19c22 < import inspect --- > import json -22c24,28 -< from typing import Tuple ---- +20a24,28 > import math > import os > import time > from itertools import cycle > from pathlib import Path -25,26c31 +23,46c31,38 < from accelerate import PartialState < from accelerate.utils import set_seed ---- -> from utils import adjust_batch, count_hpu_graphs, finalize_quantization, initialize_model -28,50c33 +< < from transformers import ( < AutoTokenizer, < BloomForCausalLM, @@ -43,11 +41,21 @@ < XLMWithLMHeadModel, < XLNetLMHeadModel, < XLNetTokenizer, -< ) +--- +> from transformers import BatchEncoding +> from utils import ( +> SetTrueOrFalseOrNone, +> adjust_batch, +> count_hpu_graphs, +> finalize_quantization, +> initialize_model, +> save_model, +48c40,41 < from transformers.modeling_outputs import CausalLMOutputWithPast --- +> > from optimum.habana.utils import get_hpu_memory_stats -60,216d42 +58,280d50 < MAX_LENGTH = int(10000) # Hardcoded max length to avoid infinite loop < < MODEL_CLASSES = { @@ -205,7 +213,7 @@ < for _ in range(num_block_layers) < ) < return past_key_values -218,285c44,46 +< < < def prepare_jit_inputs(inputs, model, tokenizer): < batch_size = len(inputs) @@ -262,8 +270,8 @@ < ) < < def _reorder_cache( -< self, past_key_values: Tuple[Tuple[torch.Tensor]], beam_idx: torch.Tensor -< ) -> Tuple[Tuple[torch.Tensor]]: +< self, past_key_values: tuple[tuple[torch.Tensor]], beam_idx: torch.Tensor +< ) -> tuple[tuple[torch.Tensor]]: < """ < This function is used to re-order the `past_key_values` cache if [`~PretrainedModel.beam_search`] or < [`~PretrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct @@ -271,22 +279,22 @@ < """ < return self._default._reorder_cache(past_key_values, beam_idx) < -< +282,283c52,54 < def main(): < parser = argparse.ArgumentParser() --- > def setup_parser(parser): > # Arguments management > parser.add_argument("--device", "-d", type=str, choices=["hpu"], help="Device to run", default="hpu") -287c48 +285c56 < "--model_type", --- > "--model_name_or_path", -291c52 +289c60 < help="Model type selected in the list: " + ", ".join(MODEL_CLASSES.keys()), --- > help="Path to pre-trained model (on the HF Hub or locally).", -294c55,83 +292c63,91 < "--model_name_or_path", --- > "--bf16", @@ -318,7 +326,7 @@ > ) > parser.add_argument( > "--dataset_name", -297,298c86,97 +295,296c94,117 < required=True, < help="Path to pre-trained model or shortcut name selected in the list: " + ", ".join(MODEL_CLASSES.keys()), --- @@ -334,15 +342,8 @@ > "--do_sample", > action="store_true", > help="Whether to use sampling for generation.", -300,304d98 -< -< parser.add_argument("--prompt", type=str, default="") -< parser.add_argument("--length", type=int, default=20) -< parser.add_argument("--stop_token", type=str, default=None, help="Token at which text generation is stopped") -< -306c100,113 -< "--temperature", ---- +> ) +> parser.add_argument( > "--num_beams", > default=1, > type=int, @@ -353,16 +354,23 @@ > default=None, > type=int, > help="Size of candidate set used for re-ranking in contrastive search. top_k > 1 enables contrastive search.", -> ) -> parser.add_argument( +298,302d118 +< +< parser.add_argument("--prompt", type=str, default="") +< parser.add_argument("--length", type=int, default=20) +< parser.add_argument("--stop_token", type=str, default=None, help="Token at which text generation is stopped") +< +304c120,121 +< "--temperature", +--- > "--penalty_alpha", > default=None, -308,309c115 +306,307c123 < default=1.0, < help="temperature of 1.0 has no effect, lower tend toward greedy sampling", --- > help="Degeneration penalty for contrastive search. penalty_alpha > 0 enables contrastive search.", -312c118,249 +310c126,146 < "--repetition_penalty", type=float, default=1.0, help="primarily useful for CTRL model; in that case, use 1.2" --- > "--trim_logits", @@ -386,13 +394,28 @@ > default=0, > type=int, > help="Number of steps to capture for profiling.", -> ) -> parser.add_argument( +312,319d147 +< parser.add_argument("--k", type=int, default=0) +< parser.add_argument("--p", type=float, default=0.9) +< +< parser.add_argument("--prefix", type=str, default="", help="Text added prior to input.") +< parser.add_argument("--padding_text", type=str, default="", help="Deprecated, the use of `--prefix` is preferred.") +< parser.add_argument("--xlm_language", type=str, default="", help="Optional language when used with the XLM model.") +< +< parser.add_argument("--seed", type=int, default=42, help="random seed for initialization") +321c149 +< "--use_cpu", +--- > "--profiling_record_shapes", -> action="store_true", +323c151 +< help="Whether or not to use cpu. If set to False, we will use gpu/npu or mps device if available", +--- > help="Record shapes when enabling profiling.", -> ) -> parser.add_argument( +325d152 +< parser.add_argument("--num_return_sequences", type=int, default=1, help="The number of samples to generate.") +327c154,201 +< "--fp16", +--- > "--prompt", > default=None, > type=str, @@ -441,7 +464,9 @@ > ) > parser.add_argument( > "--attn_softmax_bf16", -> action="store_true", +329c203,349 +< help="Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit", +--- > help="Whether to run attention softmax layer in lower precision provided that the model supports it and " > "is also running in lower precision.", > ) @@ -476,6 +501,11 @@ > help="Skip HPU Graph usage for first token to save memory", > ) > parser.add_argument( +> "--clear_hpu_graphs_cache", +> action="store_true", +> help="Clear HPU graphs cache", +> ) +> parser.add_argument( > "--show_graphs_count", > action="store_true", > help="Show statistics of HPU graph compilation.", @@ -497,18 +527,13 @@ > "--reduce_recompile", > action="store_true", > help="Preprocess on cpu, and some other optimizations. Useful to prevent recompilations when using dynamic prompts (simulate_dyn_prompt)", -314,319d250 -< parser.add_argument("--k", type=int, default=0) -< parser.add_argument("--p", type=float, default=0.9) -< -< parser.add_argument("--prefix", type=str, default="", help="Text added prior to input.") -< parser.add_argument("--padding_text", type=str, default="", help="Deprecated, the use of `--prefix` is preferred.") -< parser.add_argument("--xlm_language", type=str, default="", help="Optional language when used with the XLM model.") -321d251 -< parser.add_argument("--seed", type=int, default=42, help="random seed for initialization") -323c253,263 -< "--use_cpu", ---- +> ) +> parser.add_argument( +> "--use_chat_template", +> action="store_true", +> help="Wraps the prompt(s) in a chat template of `{ user: }`", +> ) +> parser.add_argument( > "--use_flash_attention", > action="store_true", > help="Whether to enable Habana Flash Attention, provided that the model supports it.", @@ -520,19 +545,14 @@ > ) > parser.add_argument( > "--flash_attention_causal_mask", -325c265 -< help="Whether or not to use cpu. If set to False, " "we will use gpu/npu or mps device if available", ---- +> action="store_true", > help="Whether to enable Habana Flash Attention in causal mode on first token generation.", -327d266 -< parser.add_argument("--num_return_sequences", type=int, default=1, help="The number of samples to generate.") -329c268 -< "--fp16", ---- +> ) +> parser.add_argument( > "--flash_attention_fast_softmax", -331c270,294 -< help="Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit", ---- +> nargs="?", +> const=None, +> action=SetTrueOrFalseOrNone, > help="Whether to enable Habana Flash Attention in fast softmax mode.", > ) > parser.add_argument( @@ -558,13 +578,6 @@ > "--csp", > type=str, > help="Path to serialize const params. Const params will be held on disk memory instead of being allocated on host memory.", -333c296,323 -< parser.add_argument("--jit", action="store_true", help="Whether or not to use jit trace to accelerate inference") ---- -> parser.add_argument( -> "--disk_offload", -> action="store_true", -> help="Whether to enable device map auto. In case no space left on cpu, weights will be offloaded to disk.", > ) > parser.add_argument( > "--trust_remote_code", @@ -572,41 +585,83 @@ > help="Whether to trust the execution of code from datasets/models defined on the Hub. This option should only be set to `True` for repositories you trust and in which you have read the code, as it will execute code present on the Hub on your local machine.", > ) > parser.add_argument( -> "--load_quantized_model", -> action="store_true", -> help="Whether to load model from hugging face checkpoint.", -> ) -> parser.add_argument( > "--parallel_strategy", > type=str, -> choices=["tp", "none"], # Add other strategies as needed +> choices=["tp", "ep", "none"], # Add other strategies as needed > default="none", -> help="Run multi card with the specified parallel strategy. Choices are 'tp' for Tensor Parallel Strategy or 'none'.", +> help="Run multi card with the specified parallel strategy. Choices are 'tp' for Tensor Parallel Strategy or 'ep' for Expert Parallel Strategy or 'none'.", > ) > parser.add_argument( > "--input_embeds", > action="store_true", > help="Whether to enable inputs_embeds or not.", > ) -> -336,337c326,327 +> parser.add_argument( +> "--run_partial_dataset", +> action="store_true", +> help="Run the inference with dataset for specified --n_iterations(default:5)", +> ) +> parser.add_argument( +> "--sdp_on_bf16", action="store_true", help="Allow pyTorch to use reduced precision in the SDPA math backend" +> ) +> parser.add_argument( +> "--save_quantized_model_with_inc", +> action="store_true", +> help="Save quantized Huggingface checkpoint using INC.", +> ) +> parser.add_argument( +> "--saved_model_path", +> type=str, +> default="inc_quantized_model", +> help="A path to save quantized checkpoint.", +331,335d350 +< parser.add_argument("--jit", action="store_true", help="Whether or not to use jit trace to accelerate inference") +< args = parser.parse_args() +< < # Initialize the distributed state. < distributed_state = PartialState(cpu=args.use_cpu) ---- -> if args.torch_compile: -> args.use_hpu_graphs = False -339c329,330 +337c352,384 < logger.warning(f"device: {distributed_state.device}, 16-bits inference: {args.fp16}") --- -> if not args.use_hpu_graphs: -> args.limit_hpu_graphs = False -341,342c332,333 +> quant_parser_group = parser.add_mutually_exclusive_group() +> quant_parser_group.add_argument( +> "--load_quantized_model_with_autogptq", +> action="store_true", +> help="Load an AutoGPTQ quantized checkpoint using AutoGPTQ.", +> ) +> quant_parser_group.add_argument( +> "--load_quantized_model_with_autoawq", +> action="store_true", +> help="Load an AutoAWQ quantized checkpoint using AutoAWQ.", +> ) +> quant_parser_group.add_argument( +> "--disk_offload", +> action="store_true", +> help="Whether to enable device map auto. In case no space left on cpu, weights will be offloaded to disk.", +> ) +> quant_parser_group.add_argument( +> "--load_quantized_model_with_inc", +> action="store_true", +> help="Load a quantized Huggingface checkpoint using INC.", +> ) +> quant_parser_group.add_argument( +> "--local_quantized_inc_model_path", +> type=str, +> default=None, +> help="Path to neural-compressor quantized model, if set, the checkpoint will be loaded.", +> ) +> parser.add_argument( +> "--attn_batch_split", +> default=1, +> type=int, +> help="Specify the batch size split for attention and mlp layers. 1 for no split. This is enabled only for prompt.", +> ) +339,340c386 < if args.seed is not None: < set_seed(args.seed) --- -> if args.use_flash_attention and not args.flash_attention_fast_softmax: -> args.flash_attention_fast_softmax = True -344,371c335,340 +> args = parser.parse_args() +342,369c388,389 < # Initialize the model and tokenizer < try: < args.model_type = args.model_type.lower() @@ -636,26 +691,48 @@ < prepare_input = PREPROCESSING_FUNCTIONS.get(args.model_type) < preprocessed_prompt_text = prepare_input(args, model, tokenizer, prompt_text) --- -> args.quant_config = os.getenv("QUANT_CONFIG", "") -> if args.quant_config == "" and args.disk_offload: -> logger.warning( -> "`--disk_offload` was tested only with fp8, it may not work with full precision. If error raises try to remove the --disk_offload flag." -> ) -> return args -373,376d341 +> if args.torch_compile: +> args.use_hpu_graphs = False +371,374c391,392 < if model.__class__.__name__ in ["TransfoXLLMHeadModel"]: < tokenizer_kwargs = {"add_space_before_punct_symbol": True} < else: < tokenizer_kwargs = {} -378,384c343,349 +--- +> if not args.use_hpu_graphs: +> args.limit_hpu_graphs = False +376,377c394,396 < encoded_prompt = tokenizer.encode( < preprocessed_prompt_text, add_special_tokens=False, return_tensors="pt", **tokenizer_kwargs -< ) -< else: +--- +> if args.use_flash_attention and args.flash_attention_fast_softmax is None: +> logger.warning( +> "`--flash_attention_fast_softmax` was not set; defaulting to True due to `--use_flash_attention` being enabled." +378a398 +> args.flash_attention_fast_softmax = True +380,382c400 < prefix = args.prefix if args.prefix else args.padding_text < encoded_prompt = tokenizer.encode(prefix + prompt_text, add_special_tokens=False, return_tensors="pt") < encoded_prompt = encoded_prompt.to(distributed_state.device) --- +> args.flash_attention_fast_softmax = False +384,385c402,668 +< if encoded_prompt.size()[-1] == 0: +< input_ids = None +--- +> args.quant_config = os.getenv("QUANT_CONFIG", "") +> if args.quant_config and args.load_quantized_model_with_autogptq: +> raise RuntimeError("Setting both quant_config and load_quantized_model_with_autogptq is unsupported. ") +> if args.quant_config and args.load_quantized_model_with_autoawq: +> raise RuntimeError("Setting both quant_config and load_quantized_model_with_autoawq is unsupported. ") +> +> if args.quant_config == "" and args.disk_offload: +> logger.warning( +> "`--disk_offload` was tested only with fp8, it may not work with full precision. If error raises try to remove the --disk_offload flag." +> ) +> return args +> +> > def prepare_generation_embedding(model, model_name, input_tokens): > batch_size = input_tokens["input_ids"].size(0) > @@ -663,10 +740,7 @@ > > if inputs_embeds.size(0) != batch_size: > inputs_embeds = inputs_embeds.expand(batch_size, -1, -1) -386,387c351,570 -< if encoded_prompt.size()[-1] == 0: -< input_ids = None ---- +> > attention_mask = input_tokens["attention_mask"] > return {"inputs_embeds": inputs_embeds, "attention_mask": attention_mask} > @@ -682,6 +756,9 @@ > > import habana_frameworks.torch.hpu as torch_hpu > +> if args.sdp_on_bf16: +> torch._C._set_math_sdp_allow_fp16_bf16_reduction(True) +> > if args.dataset_name is None: > # Benchmark over the prompts below > if args.prompt: @@ -750,13 +827,31 @@ > encode_t0 = time.perf_counter() > # Tokenization > if args.max_input_tokens > 0: +> if hasattr(model.config, "type_vocab_size") and model.config.type_vocab_size > 0: +> return_token_type_ids = True +> else: +> return_token_type_ids = False +> > input_tokens = tokenizer.batch_encode_plus( > input_sentences, > return_tensors="pt", > padding="max_length", > max_length=args.max_input_tokens, > truncation=True, +> return_token_type_ids=return_token_type_ids, > ) +> +> def compute_valid_sequence_lengths_tensor(input_tokens): +> attn_mask = input_tokens["attention_mask"] +> return torch.sum(attn_mask, dim=1, dtype=torch.int32) +> +> valid_sequence_lengths = compute_valid_sequence_lengths_tensor(input_tokens).to(args.device) +> generation_config.valid_sequence_lengths = valid_sequence_lengths +> elif args.use_chat_template: +> input_messages = [{"role": "user", "content": sentence} for sentence in input_sentences] +> input_ids = tokenizer.apply_chat_template(input_messages, return_tensors="pt", padding=True) +> attention_mask = torch.ones_like(input_ids) +> input_tokens = BatchEncoding({"input_ids": input_ids, "attention_mask": attention_mask}) > else: > input_tokens = tokenizer.batch_encode_plus(input_sentences, return_tensors="pt", padding=True) > encode_duration = time.perf_counter() - encode_t0 @@ -795,7 +890,7 @@ > profiling_record_shapes=args.profiling_record_shapes, > ).cpu() > first_token_time = iteration_times[0] + encode_duration -> logger.info(f"Time to first token = {first_token_time*1000}ms") +> logger.info(f"Time to first token = {first_token_time * 1000}ms") > return tokenizer.batch_decode(outputs, skip_special_tokens=True) > > from optimum.habana.utils import HabanaProfile @@ -810,10 +905,10 @@ > if dyn_prompt_lens is None or len(set(dyn_prompt_lens)) == 1: > for i in range(args.warmup): > if dyn_prompt_lens is None: -> print(f"Warming up iteration {i+1}/{args.warmup}", flush=True) +> print(f"Warming up iteration {i + 1}/{args.warmup}", flush=True) > generate(None, args.reduce_recompile) > else: -> print(f"Warming up for shape {dyn_prompt_lens[0]} iteration {i+1}/{args.warmup}", flush=True) +> print(f"Warming up for shape {dyn_prompt_lens[0]} iteration {i + 1}/{args.warmup}", flush=True) > generate(dyn_prompt_lens[0], args.reduce_recompile) > else: > if args.bucket_size > 0: @@ -828,7 +923,7 @@ > for i in range(args.warmup): > lst = list(range(min_prompt_len, max_sentence_len + 1, args.bucket_size)) > for sz in lst: -> print(f"Warming up for shape {sz - 1} iteration {i+1}/{args.warmup}", flush=True) +> print(f"Warming up for shape {sz - 1} iteration {i + 1}/{args.warmup}", flush=True) > generate(sz - 1, args.reduce_recompile) > torch_hpu.synchronize() > compilation_duration = time.perf_counter() - t0 @@ -852,12 +947,16 @@ > > print() > print("Input/outputs:") +> all_inputs = [] +> all_outputs = [] > for i, input_sentence in enumerate(zip(input_sentences)): -> print(f"input {i+1}: {input_sentence}") +> print(f"input {i + 1}: {input_sentence}") +> all_inputs.append(input_sentence) > for j, output in enumerate( > zip(generated[args.num_return_sequences * i : args.num_return_sequences * (i + 1)]) > ): -> print(f"output {j+1}: {output}") +> print(f"output {i + 1}.{j + 1}: {output}") +> all_outputs.append(output) > print() > > # Store results if necessary @@ -867,7 +966,8 @@ > > results = { > "throughput": throughput, -> "output": output, +> "input": all_inputs, +> "output": all_outputs, > } > with (output_dir / "results.json").open("w", encoding="utf-8") as f: > json.dump(results, f, ensure_ascii=False, indent=4) @@ -887,7 +987,7 @@ > print(f"Graph compilation duration = {compilation_duration} seconds") > print(separator) > print() -389c572,589 +387c670,687 < input_ids = encoded_prompt --- > # Downloading and loading a dataset from the hub. @@ -896,7 +996,7 @@ > > assert not args.simulate_dyn_prompt, "Both dataset_name and simulate_dyn_prompt are set" > -> raw_dataset = load_dataset(args.dataset_name) +> raw_dataset = load_dataset(args.dataset_name, trust_remote_code=args.trust_remote_code) > if "test" in raw_dataset: > split = "test" > elif "validation" in raw_dataset: @@ -908,7 +1008,7 @@ > .shuffle() > .select(range(args.dataset_max_samples if args.dataset_max_samples > 0 else (raw_dataset[split]).num_rows)) > ) -391,397c591,598 +389,395c689,696 < if args.jit: < jit_input_texts = ["enable jit"] < jit_inputs = prepare_jit_inputs(jit_input_texts, model, tokenizer) @@ -925,7 +1025,7 @@ > logger.info( > f"No column name was given so automatically choosing '{column_name}' for prompts. If you would like to use another column of the dataset, you can set the argument `--column_name`." > ) -399,437c600,620 +397,435c698,718 < sig = inspect.signature(model.__call__) < jit_inputs = tuple(jit_inputs[key] for key in sig.parameters if jit_inputs.get(key, None) is not None) < traced_model = torch.jit.trace(model, jit_inputs, strict=False) @@ -987,7 +1087,7 @@ > preprocess_function, > batched=True, > desc="Running tokenizer on dataset", -438a622,704 +436a720,743 > # After tokenization, we can remove the column of interest > raw_dataset = raw_dataset.remove_columns([column_name]) > raw_dataset.set_format(type="torch") @@ -1012,7 +1112,10 @@ > collate_fn = None > > dataloader = DataLoader(raw_dataset, batch_size=args.batch_size, collate_fn=collate_fn) -> +438,439c745,822 +< generated_sequences.append(total_sequence) +< print(total_sequence) +--- > def generate_dataset(batch): > prompt = tokenizer.batch_decode(batch["input_ids"], skip_special_tokens=True) > # Move inputs to target device(s) @@ -1033,22 +1136,21 @@ > return prompt, outputs > > # warmup -> if prompt_length > 0: -> from optimum.habana.utils import HabanaProfile +> from optimum.habana.utils import HabanaProfile > -> # compilation stage disable profiling -> HabanaProfile.disable() -> # Compilation -> logger.info("Graph compilation...") -> t0 = time.perf_counter() -> for i, batch in enumerate(dataloader): -> generate_dataset(batch) -> # The first three iterations take longer because of graph compilation -> if (i + 1) == 3: -> break -> torch_hpu.synchronize() -> compilation_duration = time.perf_counter() - t0 -> HabanaProfile.enable() +> # compilation stage disable profiling +> HabanaProfile.disable() +> # Compilation +> logger.info("Graph compilation...") +> t0 = time.perf_counter() +> for i, batch in enumerate(dataloader): +> generate_dataset(batch) +> # The first three iterations take longer because of graph compilation +> if (i + 1) == 3: +> break +> torch_hpu.synchronize() +> compilation_duration = time.perf_counter() - t0 +> HabanaProfile.enable() > > total_new_tokens_generated = 0 > duration = 0 @@ -1061,20 +1163,19 @@ > duration += time.perf_counter() - t0 > total_new_tokens_generated += args.batch_size * args.max_new_tokens > print(separator) -> print(f"Batch n°{i+1}") -> print(f"Input: {prompt[:args.batch_size]}") +> print(f"Batch n°{i + 1}") +> print(f"Input: {prompt[: args.batch_size]}") > print( -> f"Output: {tokenizer.batch_decode(outputs, skip_special_tokens=True)[:args.batch_size*args.num_return_sequences]}" +> f"Output: {tokenizer.batch_decode(outputs, skip_special_tokens=True)[: args.batch_size * args.num_return_sequences]}" > ) > print(separator) +> if args.run_partial_dataset and args.n_iterations == i + 1: +> break > t_end = time.time() > > throughput = total_new_tokens_generated / duration > # Print Stats -440,441c706,722 -< generated_sequences.append(total_sequence) -< print(total_sequence) ---- +> > stats = f"Throughput (including tokenization) = {throughput} tokens/second" > separator = "-" * len(stats) > print() @@ -1085,14 +1186,15 @@ > mem = get_hpu_memory_stats() > for k, v in mem.items(): > print("{:35} = {} GB".format(k[:-5].replace("_", " ").capitalize(), v)) -> if prompt_length > 0: -> print(f"Graph compilation duration = {compilation_duration} seconds") +> print(f"Graph compilation duration = {compilation_duration} seconds") > print(separator) > if args.quant_config: > finalize_quantization(model) +> if args.save_quantized_model_with_inc: +> save_model(model, tokenizer, args.saved_model_path) > if args.const_serialization_path and os.path.isdir(args.const_serialization_path): > import shutil -443c724 +441c824 < return generated_sequences --- > shutil.rmtree(args.const_serialization_path) diff --git a/tests/example_diff/run_glue.txt b/tests/example_diff/run_glue.txt index 7594d07d88..a332e54bac 100644 --- a/tests/example_diff/run_glue.txt +++ b/tests/example_diff/run_glue.txt @@ -1,14 +1,16 @@ -29,30d28 +1a2 +> # coding=utf-8 +28,29d28 < from datasets import load_dataset < -31a30 +30a30 > from datasets import load_dataset -40,41d38 +39,40d38 < Trainer, < TrainingArguments, -43d39 +42d39 < set_seed, -48a45,54 +47a45,54 > from optimum.habana import GaudiConfig, GaudiTrainer, GaudiTrainingArguments > from optimum.habana.utils import set_seed > @@ -19,34 +21,34 @@ > > def check_optimum_habana_min_version(*a, **b): > return () -50,51c56,61 +49,50c56,61 < # Will error if the minimal version of Transformers is not installed. Remove at your own risks. -< check_min_version("4.46.0.dev0") +< check_min_version("4.52.0.dev0") --- > > logger = logging.getLogger(__name__) > > # Will error if the minimal version of Transformers and Optimum Habana are not installed. Remove at your own risks. -> check_min_version("4.45.0") -> check_optimum_habana_min_version("1.14.0.dev0") -67,68d76 +> check_min_version("4.51.0") +> check_optimum_habana_min_version("1.18.0.dev0") +66,67d76 < logger = logging.getLogger(__name__) < -143a152,155 +142a152,155 > problem_type: Optional[str] = field( > default="single_label_classification", > metadata={"help": "Problem type, such as single_label_classification or multi_label_classification"}, > ) -213a226,229 +212a226,229 > add_pad_token: bool = field( > default=False, > metadata={"help": "Will add `pad_token` to tokenizer and model's config as `eos_token` if it's not defined."}, > ) -221c237 +220c237 < parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)) --- > parser = HfArgumentParser((ModelArguments, DataTrainingArguments, GaudiTrainingArguments)) -250a267,273 +249a267,273 > gaudi_config = GaudiConfig.from_pretrained( > training_args.gaudi_config_name, > cache_dir=model_args.cache_dir, @@ -54,34 +56,30 @@ > token=model_args.token, > ) > -251a275 +250a275 > mixed_precision = training_args.bf16 or gaudi_config.use_torch_autocast -253,254c277,279 +252,253c277,279 < f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}, " < + f"distributed training: {training_args.parallel_mode.value == 'distributed'}, 16-bits training: {training_args.fp16}" --- > f"Process rank: {training_args.local_rank}, device: {training_args.device}, " > + f"distributed training: {training_args.parallel_mode.value == 'distributed'}, " > + f"mixed-precision training: {mixed_precision}" -376a402 +375a402 > problem_type=data_args.problem_type, -417a444,448 +416a444,448 > if model_args.add_pad_token: > if not model.config.pad_token_id and not tokenizer.pad_token: > tokenizer.pad_token = tokenizer.eos_token > model.config.pad_token_id = tokenizer.eos_token_id > -528c559 +527c559 < trainer = Trainer( --- > trainer = GaudiTrainer( -529a561 +528a561 > gaudi_config=gaudi_config, -534c566 -< processing_class=tokenizer, ---- -> tokenizer=tokenizer, -629,633d660 +628,632d660 < < < def _mp_fn(index): diff --git a/tests/example_diff/run_image_classification.txt b/tests/example_diff/run_image_classification.txt index e13bf04a37..e42e130c2b 100644 --- a/tests/example_diff/run_image_classification.txt +++ b/tests/example_diff/run_image_classification.txt @@ -1,16 +1,18 @@ -14a15,16 +1a2 +> # coding=utf-8 +13a15,16 > # limitations under the License. > """Fine-tuning a 🤗 Transformers model for image classification""" -24a27 +23a27 > import transformers -37,38d39 +36,37d39 < < import transformers -45,47d45 +45,47d46 < Trainer, < TrainingArguments, < set_seed, -52a51,60 +52a52,61 > from optimum.habana import GaudiConfig, GaudiTrainer, GaudiTrainingArguments > from optimum.habana.utils import set_seed > @@ -21,20 +23,20 @@ > > def check_optimum_habana_min_version(*a, **b): > return () -54d61 +54d62 < """ Fine-tuning a 🤗 Transformers model for image classification""" -58,59c65,67 +58,59c66,68 < # Will error if the minimal version of Transformers is not installed. Remove at your own risks. -< check_min_version("4.46.0.dev0") +< check_min_version("4.52.0.dev0") --- > # Will error if the minimal version of Transformers and Optimum Habana are not installed. Remove at your own risks. -> check_min_version("4.45.0") -> check_optimum_habana_min_version("1.14.0.dev0") -184c192 +> check_min_version("4.51.0") +> check_optimum_habana_min_version("1.18.0.dev0") +184c193 < parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)) --- > parser = HfArgumentParser((ModelArguments, DataTrainingArguments, GaudiTrainingArguments)) -212a221,227 +212a222,228 > gaudi_config = GaudiConfig.from_pretrained( > training_args.gaudi_config_name, > cache_dir=model_args.cache_dir, @@ -42,22 +44,18 @@ > token=model_args.token, > ) > -213a229 +213a230 > mixed_precision = training_args.bf16 or gaudi_config.use_torch_autocast -215,216c231,233 +215,216c232,234 < f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}, " < + f"distributed training: {training_args.parallel_mode.value == 'distributed'}, 16-bits training: {training_args.fp16}" --- > f"Process rank: {training_args.local_rank}, device: {training_args.device}, " > + f"distributed training: {training_args.parallel_mode.value == 'distributed'}, " > + f"mixed-precision training: {mixed_precision}" -393c410 +398c416 < trainer = Trainer( --- > trainer = GaudiTrainer( -394a412 +399a418 > gaudi_config=gaudi_config, -399c417 -< processing_class=image_processor, ---- -> tokenizer=image_processor, diff --git a/tests/example_diff/run_mlm.txt b/tests/example_diff/run_mlm.txt index 6774aac565..5ae5abda17 100644 --- a/tests/example_diff/run_mlm.txt +++ b/tests/example_diff/run_mlm.txt @@ -1,27 +1,29 @@ -17,19c17,18 +1a2 +> # coding=utf-8 +16,18c17,18 < Fine-tuning the library models for masked language modeling (BERT, ALBERT, RoBERTa...) on a text file or a dataset. < < Here is the full list of checkpoints on the hub that can be fine-tuned by this script: --- > Training the library models for masked language modeling (BERT, ALBERT, RoBERTa...) on a text file or a dataset. > Here is the full list of checkpoints on the hub that can be trained by this script: -35,36d33 +34,35d33 < from datasets import load_dataset < -37a35 +36a35 > from datasets import load_dataset -46,49d43 +45,48d43 < Trainer, < TrainingArguments, < is_torch_xla_available, < set_seed, -54a49,50 +53a49,50 > from optimum.habana import GaudiConfig, GaudiTrainer, GaudiTrainingArguments > from optimum.habana.utils import set_seed -56,57d51 +55,56d51 < # Will error if the minimal version of Transformers is not installed. Remove at your own risks. -< check_min_version("4.46.0.dev0") -59c53,59 +< check_min_version("4.52.0.dev0") +58c53,59 < require_version("datasets>=2.14.0", "To fix: pip install -r examples/pytorch/language-modeling/requirements.txt") --- > try: @@ -31,32 +33,32 @@ > def check_optimum_habana_min_version(*a, **b): > return () > -61a62,69 +60a62,69 > > # Will error if the minimal version of Transformers and Optimum Habana are not installed. Remove at your own risks. -> check_min_version("4.45.0") -> check_optimum_habana_min_version("1.14.0.dev0") +> check_min_version("4.51.0") +> check_optimum_habana_min_version("1.18.0.dev0") > > require_version("datasets>=2.14.0", "To fix: pip install -r examples/pytorch/language-modeling/requirements.txt") > > -137c145 +136c145 < "choices": ["auto", "bfloat16", "float16", "float32"], --- > "choices": ["auto", "bfloat16", "float32"], -145c153 +144c153 < "set True will benefit LLM loading time and RAM consumption." --- > "Setting it to True will benefit LLM loading time and RAM consumption." -230c238 +229c238 < streaming: bool = field(default=False, metadata={"help": "Enable streaming mode"}) --- > streaming: bool = field(default=False, metadata={"help": "Enable streaming mode."}) -254c262 +253c262 < parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)) --- > parser = HfArgumentParser((ModelArguments, DataTrainingArguments, GaudiTrainingArguments)) -283a292,298 +282a292,298 > gaudi_config = GaudiConfig.from_pretrained( > training_args.gaudi_config_name, > cache_dir=model_args.cache_dir, @@ -64,28 +66,24 @@ > token=model_args.token, > ) > -284a300 +283a300 > mixed_precision = training_args.bf16 or gaudi_config.use_torch_autocast -286,287c302,304 +285,286c302,304 < f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}, " < + f"distributed training: {training_args.parallel_mode.value == 'distributed'}, 16-bits training: {training_args.fp16}" --- > f"Process rank: {training_args.local_rank}, device: {training_args.device}, " > + f"distributed training: {training_args.parallel_mode.value == 'distributed'}, " > + f"mixed-precision training: {mixed_precision}" -289d305 +288d305 < # Set the verbosity to info of the Transformers logger (on main process only): -620c636 +619c636 < trainer = Trainer( --- > trainer = GaudiTrainer( -621a638 +620a638 > gaudi_config=gaudi_config, -625c642 -< processing_class=tokenizer, ---- -> tokenizer=tokenizer, -627,630c644,645 +626,629c644,645 < compute_metrics=compute_metrics if training_args.do_eval and not is_torch_xla_available() else None, < preprocess_logits_for_metrics=preprocess_logits_for_metrics < if training_args.do_eval and not is_torch_xla_available() @@ -93,7 +91,7 @@ --- > compute_metrics=compute_metrics if training_args.do_eval else None, > preprocess_logits_for_metrics=preprocess_logits_for_metrics if training_args.do_eval else None, -644,647c659,665 +643,646c659,665 < max_train_samples = ( < data_args.max_train_samples if data_args.max_train_samples is not None else len(train_dataset) < ) @@ -106,9 +104,9 @@ > data_args.max_train_samples if data_args.max_train_samples is not None else len(train_dataset) > ) > metrics["train_samples"] = min(max_train_samples, len(train_dataset)) -656d673 +655d673 < -659,660c676,681 +658,659c676,681 < max_eval_samples = data_args.max_eval_samples if data_args.max_eval_samples is not None else len(eval_dataset) < metrics["eval_samples"] = min(max_eval_samples, len(eval_dataset)) --- @@ -118,7 +116,7 @@ > ) > metrics["eval_samples"] = min(max_eval_samples, len(eval_dataset)) > -683,687d703 +682,686d703 < < < def _mp_fn(index): diff --git a/tests/example_diff/run_qa.txt b/tests/example_diff/run_qa.txt index d5f93c68de..9293643ea3 100644 --- a/tests/example_diff/run_qa.txt +++ b/tests/example_diff/run_qa.txt @@ -1,26 +1,27 @@ -3c3 +2c2,3 < # Copyright 2020 The HuggingFace Team All rights reserved. --- +> # coding=utf-8 > # Copyright 2022 The HuggingFace Team All rights reserved. -29a30 +28a30 > import transformers -32,34d32 +31,33d32 < from utils_qa import postprocess_qa_predictions < < import transformers -43d40 +42d40 < TrainingArguments, -45d41 +44d41 < set_seed, -49a46 +48a46 > from utils_qa import postprocess_qa_predictions -50a48,49 +49a48,49 > from optimum.habana import GaudiConfig, GaudiTrainingArguments > from optimum.habana.utils import set_seed -52,53d50 +51,52d50 < # Will error if the minimal version of Transformers is not installed. Remove at your own risks. -< check_min_version("4.46.0.dev0") -55c52,58 +< check_min_version("4.52.0.dev0") +54c52,58 < require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/question-answering/requirements.txt") --- > try: @@ -30,22 +31,22 @@ > def check_optimum_habana_min_version(*a, **b): > return () > -58a62,67 +57a62,67 > # Will error if the minimal version of Transformers and Optimum Habana are not installed. Remove at your own risks. -> check_min_version("4.45.0") -> check_optimum_habana_min_version("1.14.0.dev0") +> check_min_version("4.51.0") +> check_optimum_habana_min_version("1.18.0.dev0") > > require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/question-answering/requirements.txt") > -146c155 +145c155 < " batching to the maximum length in the batch (which can be faster on GPU but will be slower on TPU)." --- > " batching to the maximum length in the batch (which can be faster on GPU but will be slower on HPU)." -233c242 +232c242 < parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)) --- > parser = HfArgumentParser((ModelArguments, DataTrainingArguments, GaudiTrainingArguments)) -262a272,278 +261a272,278 > gaudi_config = GaudiConfig.from_pretrained( > training_args.gaudi_config_name, > cache_dir=model_args.cache_dir, @@ -53,27 +54,27 @@ > token=model_args.token, > ) > -263a280 +262a280 > mixed_precision = training_args.bf16 or gaudi_config.use_torch_autocast -265,266c282,284 +264,265c282,284 < f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}, " < + f"distributed training: {training_args.parallel_mode.value == 'distributed'}, 16-bits training: {training_args.fp16}" --- > f"Process rank: {training_args.local_rank}, device: {training_args.device}, " > + f"distributed training: {training_args.parallel_mode.value == 'distributed'}, " > + f"mixed-precision training: {mixed_precision}" -347a366,369 +346a366,369 > if config.model_type == "llama": > if tokenizer.pad_token is None: > tokenizer.add_special_tokens({"pad_token": "[PAD]"}) > tokenizer.cls_token = tokenizer.bos_token -638a661 -> gaudi_config=gaudi_config, -643c666 -< processing_class=tokenizer, +356c379 +< --- -> tokenizer=tokenizer, -707,711d729 +> model = model.to("hpu") +637a661 +> gaudi_config=gaudi_config, +706,710d729 < < < def _mp_fn(index): diff --git a/tests/example_diff/run_seq2seq_qa.txt b/tests/example_diff/run_seq2seq_qa.txt index 3f99d2477b..73e5f5777c 100644 --- a/tests/example_diff/run_seq2seq_qa.txt +++ b/tests/example_diff/run_seq2seq_qa.txt @@ -1,18 +1,24 @@ -29a30 +1a2 +> # coding=utf-8 +24c25 +< from typing import Optional +--- +> from typing import List, Optional, Tuple +28a30 > import transformers -32,33d32 +31,32d32 < < import transformers -40,41d38 +39,40d38 < Seq2SeqTrainingArguments, < set_seed, -46a44,45 +45a44,45 > from optimum.habana import GaudiConfig, GaudiSeq2SeqTrainingArguments > from optimum.habana.utils import set_seed -48,49d46 +47,48d46 < # Will error if the minimal version of Transformers is not installed. Remove at your own risks. -< check_min_version("4.46.0.dev0") -51c48,54 +< check_min_version("4.52.0.dev0") +50c48,54 < require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/question-answering/requirements.txt") --- > try: @@ -22,22 +28,22 @@ > def check_optimum_habana_min_version(*a, **b): > return () > -54a58,63 +53a58,63 > # Will error if the minimal version of Transformers and Optimum Habana are not installed. Remove at your own risks. -> check_min_version("4.45.0") -> check_optimum_habana_min_version("1.14.0.dev0") +> check_min_version("4.51.0") +> check_optimum_habana_min_version("1.18.0.dev0") > > require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/question-answering/requirements.txt") > -178c187 +177c187 < " batching to the maximum length in the batch (which can be faster on GPU but will be slower on TPU)." --- > " batching to the maximum length in the batch (which can be faster on GPU but will be slower on HPU)." -278c287 +277c287 < parser = HfArgumentParser((ModelArguments, DataTrainingArguments, Seq2SeqTrainingArguments)) --- > parser = HfArgumentParser((ModelArguments, DataTrainingArguments, GaudiSeq2SeqTrainingArguments)) -307a317,323 +306a317,323 > gaudi_config = GaudiConfig.from_pretrained( > training_args.gaudi_config_name, > cache_dir=model_args.cache_dir, @@ -45,22 +51,22 @@ > token=model_args.token, > ) > -308a325 +307a325 > mixed_precision = training_args.bf16 or gaudi_config.use_torch_autocast -310,311c327,329 +309,310c327,329 < f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}, " < + f"distributed training: {training_args.parallel_mode.value == 'distributed'}, 16-bits training: {training_args.fp16}" --- > f"Process rank: {training_args.local_rank}, device: {training_args.device}, " > + f"distributed training: {training_args.parallel_mode.value == 'distributed'}, " > + f"mixed-precision training: {mixed_precision}" -661a680 -> gaudi_config=gaudi_config, -666c685 -< processing_class=tokenizer, +471c490 +< ) -> tuple[list[str], list[str]]: --- -> tokenizer=tokenizer, -735,739d753 +> ) -> Tuple[List[str], List[str]]: +660a680 +> gaudi_config=gaudi_config, +734,738d753 < < < def _mp_fn(index): diff --git a/tests/example_diff/run_speech_recognition_ctc.txt b/tests/example_diff/run_speech_recognition_ctc.txt index b78cfbbebd..bb62540d08 100644 --- a/tests/example_diff/run_speech_recognition_ctc.txt +++ b/tests/example_diff/run_speech_recognition_ctc.txt @@ -1,20 +1,26 @@ -32,33d31 +1a2 +> # coding=utf-8 +26c27 +< from typing import Optional, Union +--- +> from typing import Dict, List, Optional, Union +31,32d31 < from datasets import DatasetDict, load_dataset < -34a33 +33a33 > from datasets import DatasetDict, load_dataset -42,43d40 +41,42d40 < Trainer, < TrainingArguments, -45d41 +44d41 < set_seed, -50a47,48 +49a47,48 > from optimum.habana import GaudiConfig, GaudiTrainer, GaudiTrainingArguments > from optimum.habana.utils import set_seed -52,53d49 +51,52d49 < # Will error if the minimal version of Transformers is not installed. Remove at your own risks. -< check_min_version("4.46.0.dev0") -55c51,56 +< check_min_version("4.52.0.dev0") +54c51,56 < require_version("datasets>=1.18.0", "To fix: pip install -r examples/pytorch/speech-recognition/requirements.txt") --- > try: @@ -23,40 +29,80 @@ > > def check_optimum_habana_min_version(*a, **b): > return () -59a61,66 +58a61,66 > # Will error if the minimal version of Transformers and Optimum Habana are not installed. Remove at your own risks. -> check_min_version("4.45.0") -> check_optimum_habana_min_version("1.14.0.dev0") +> check_min_version("4.51.0") +> check_optimum_habana_min_version("1.18.0.dev0") > > require_version("datasets>=1.18.0", "To fix: pip install -r examples/pytorch/speech-recognition/requirements.txt") > -144c151 +143c151 < "help": "Whether a convolutional attention network should be stacked on top of the Wav2Vec2Bert Encoder. Can be very" --- > "help": "Whether a convolutional attention network should be stacked on top of the Wav2Vec2Bert Encoder. Can be very " -154d160 +146a155,181 +> use_flash_attention: bool = field( +> default=False, metadata={"help": "Whether to use Habana flash attention for fine-tuning"} +> ) +> flash_attention_recompute: bool = field( +> default=False, +> metadata={ +> "help": "Whether to enable recompute in Habana flash attention for fine-tuning." +> " It is applicable only when use_flash_attention is True." +> }, +> ) +> flash_attention_fast_softmax: bool = field( +> default=False, +> metadata={ +> "help": "Whether to use fast softmax for Habana flash attention." +> " It is applicable only when use_flash_attention is True." +> }, +> ) +> +> def __post_init__(self): +> if self.use_flash_attention: +> os.environ["USE_FLASH_ATTENTION"] = "1" +> if self.flash_attention_recompute: +> assert self.use_flash_attention, "flash_attention_recompute is set, but use_flash_attention is not" +> os.environ["FLASH_ATTENTION_RECOMPUTE"] = "1" +> if self.flash_attention_fast_softmax: +> assert self.use_flash_attention, "flash_attention_fast_softmax is set, but use_flash_attention is not" +> os.environ["FLASH_ATTENTION_FAST_SOFTMAX"] = "1" +153d187 < -400c406 +213c247 +< chars_to_ignore: Optional[list[str]] = list_field( +--- +> chars_to_ignore: Optional[List[str]] = list_field( +217c251 +< eval_metrics: list[str] = list_field( +--- +> eval_metrics: List[str] = list_field( +320c354 +< def __call__(self, features: list[dict[str, Union[list[int], torch.Tensor]]]) -> dict[str, torch.Tensor]: +--- +> def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]: +399c433 < parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)) --- > parser = HfArgumentParser((ModelArguments, DataTrainingArguments, GaudiTrainingArguments)) -434a441,446 +433a468,473 > gaudi_config = GaudiConfig.from_pretrained( > training_args.gaudi_config_name, > cache_dir=model_args.cache_dir, > token=data_args.token, > ) > -435a448 +434a475 > mixed_precision = training_args.bf16 or gaudi_config.use_torch_autocast -437,438c450,452 +436,437c477,479 < f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}, " < f"distributed training: {training_args.parallel_mode.value == 'distributed'}, 16-bits training: {training_args.fp16}" --- > f"Process rank: {training_args.local_rank}, device: {training_args.device}, " > + f"distributed training: {training_args.parallel_mode.value == 'distributed'}, " > + f"mixed-precision training: {mixed_precision}" -451,458c465,471 +450,457c492,498 < if training_args.do_train: < raw_datasets["train"] = load_dataset( < data_args.dataset_name, @@ -73,7 +119,7 @@ > token=data_args.token, > trust_remote_code=data_args.trust_remote_code, > ) -460,465c473,478 +459,464c500,505 < if data_args.audio_column_name not in raw_datasets["train"].column_names: < raise ValueError( < f"--audio_column_name '{data_args.audio_column_name}' not found in dataset '{data_args.dataset_name}'." @@ -87,7 +133,7 @@ > " Make sure to set `--audio_column_name` to the correct audio column - one of" > f" {', '.join(raw_datasets['train'].column_names)}." > ) -467,472c480,485 +466,471c507,512 < if data_args.text_column_name not in raw_datasets["train"].column_names: < raise ValueError( < f"--text_column_name {data_args.text_column_name} not found in dataset '{data_args.dataset_name}'. " @@ -101,36 +147,34 @@ > "Make sure to set `--text_column_name` to the correct text column - one of " > f"{', '.join(raw_datasets['train'].column_names)}." > ) -474,475c487,488 +473,474c514,515 < if data_args.max_train_samples is not None: < raw_datasets["train"] = raw_datasets["train"].select(range(data_args.max_train_samples)) --- > if data_args.max_train_samples is not None: > raw_datasets["train"] = raw_datasets["train"].select(range(data_args.max_train_samples)) -494c507 -< f'[{"".join(data_args.chars_to_ignore)}]' if data_args.chars_to_ignore is not None else None +493c534 +< f"[{''.join(data_args.chars_to_ignore)}]" if data_args.chars_to_ignore is not None else None --- -> f'[{"".join(data_args.chars_to_ignore).replace(" ", "")}]' if data_args.chars_to_ignore is not None else None -633a647,651 +> f"[{''.join(data_args.chars_to_ignore).replace(' ', '')}]" if data_args.chars_to_ignore is not None else None +523a565 +> attn_implementation=training_args.attn_implementation, +632a675,679 > raise RuntimeError( > f"The dataset sampling rate ({dataset_sampling_rate}) is different from the feature extractor one" > f" ({feature_extractor.sampling_rate}).Data resampling should be done. The Datasets library does not" > " support it on HPUs yet." > ) -743c761,764 +742c789,792 < processor=processor, feature_extractor_input_name=feature_extractor_input_name --- > processor=processor, > feature_extractor_input_name=feature_extractor_input_name, > pad_to_multiple_of=int(max_input_length), > pad_to_multiple_of_labels=500, -747c768 +746c796 < trainer = Trainer( --- > trainer = GaudiTrainer( -748a770 +747a798 > gaudi_config=gaudi_config, -754c776 -< processing_class=processor, ---- -> tokenizer=processor, diff --git a/tests/example_diff/run_speech_recognition_seq2seq.txt b/tests/example_diff/run_speech_recognition_seq2seq.txt index 188eeb1a6b..cd35ec0718 100644 --- a/tests/example_diff/run_speech_recognition_seq2seq.txt +++ b/tests/example_diff/run_speech_recognition_seq2seq.txt @@ -1,13 +1,19 @@ -31,32d30 +1a2 +> # coding=utf-8 +25c26 +< from typing import Any, Optional, Union +--- +> from typing import Any, Dict, List, Optional, Union +30,31d30 < from datasets import DatasetDict, load_dataset < -33a32 +32a32 > from datasets import DatasetDict, load_dataset -41,43d39 +40,42d39 < Seq2SeqTrainer, < Seq2SeqTrainingArguments, < set_seed, -48a45,55 +47a45,55 > from optimum.habana import GaudiConfig, GaudiSeq2SeqTrainer, GaudiSeq2SeqTrainingArguments > from optimum.habana.utils import set_seed > @@ -19,19 +25,31 @@ > def check_optimum_habana_min_version(*a, **b): > return () > -51c58,59 -< check_min_version("4.46.0.dev0") +50c58,59 +< check_min_version("4.52.0.dev0") +--- +> check_min_version("4.51.0") +> check_optimum_habana_min_version("1.18.0.dev0") +112c121 +< forced_decoder_ids: list[list[int]] = field( +--- +> forced_decoder_ids: List[List[int]] = field( +116c125 +< suppress_tokens: list[int] = field( --- -> check_min_version("4.45.0") -> check_optimum_habana_min_version("1.14.0.dev0") -230a239,242 +> suppress_tokens: List[int] = field( +229a239,242 > label_features_max_length: int = field( > default=None, > metadata={"help": "Max length for padding label features."}, > ) -248a261 +247a261 > label_features_max_length: int -262c275,279 +249c263 +< def __call__(self, features: list[dict[str, Union[list[int], torch.Tensor]]]) -> dict[str, torch.Tensor]: +--- +> def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]: +261c275,279 < labels_batch = self.processor.tokenizer.pad(label_features, return_tensors="pt") --- > kwargs = {} @@ -39,47 +57,43 @@ > kwargs["padding"] = "max_length" > kwargs["max_length"] = self.label_features_max_length > labels_batch = self.processor.tokenizer.pad(label_features, return_tensors="pt", **kwargs) -282c299 +281c299 < parser = HfArgumentParser((ModelArguments, DataTrainingArguments, Seq2SeqTrainingArguments)) --- > parser = HfArgumentParser((ModelArguments, DataTrainingArguments, GaudiSeq2SeqTrainingArguments)) -309a327,332 +308a327,332 > gaudi_config = GaudiConfig.from_pretrained( > training_args.gaudi_config_name, > cache_dir=model_args.cache_dir, > token=model_args.token, > ) > -310a334 +309a334 > mixed_precision = training_args.bf16 or gaudi_config.use_torch_autocast -312,313c336,338 +311,312c336,338 < f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}, " < f"distributed training: {training_args.parallel_mode.value == 'distributed'}, 16-bits training: {training_args.fp16}" --- > f"Process rank: {training_args.local_rank}, device: {training_args.device}, " > + f"distributed training: {training_args.parallel_mode.value == 'distributed'}, " > + f"mixed-precision training: {mixed_precision}" -444d468 +443d468 < model.generation_config.forced_decoder_ids = model_args.forced_decoder_ids -458a483,486 +457a483,486 > logger.warning( > f"The dataset sampling rate ({dataset_sampling_rate}) is different from the feature extractor one" > f" ({feature_extractor.sampling_rate}).Data resampling should be done." > ) -555c583,584 +554c583,584 < config.save_pretrained(training_args.output_dir) --- > # TODO: uncomment the line below when this is fixed in Transformers > # config.save_pretrained(training_args.output_dir) -563a593 +562a593 > label_features_max_length=data_args.label_features_max_length, -567c597 +566c597 < trainer = Seq2SeqTrainer( --- > trainer = GaudiSeq2SeqTrainer( -568a599 +567a599 > gaudi_config=gaudi_config, -572c603 -< processing_class=feature_extractor, ---- -> tokenizer=feature_extractor, diff --git a/tests/example_diff/run_summarization.txt b/tests/example_diff/run_summarization.txt index 2e84d598e1..caf7d53b6a 100644 --- a/tests/example_diff/run_summarization.txt +++ b/tests/example_diff/run_summarization.txt @@ -1,30 +1,33 @@ -3c3 +2c2,3 < # Copyright 2021 The HuggingFace Team. All rights reserved. --- +> # coding=utf-8 > # Copyright 2022 The HuggingFace Team. All rights reserved. -20a21 +19a21 > import copy -30a32,33 -> import torch -> import transformers -33,34d35 +30,32c32 +< from datasets import load_dataset +< from filelock import FileLock < -< import transformers -45,47c46 +--- +> import torch +33a34 +> from datasets import load_dataset +44,46c45 < Seq2SeqTrainer, < Seq2SeqTrainingArguments, < set_seed, --- > default_data_collator, -48a48 +47a47 > from transformers.integrations.deepspeed import is_deepspeed_zero3_enabled -52a53,54 +51a52,53 > from optimum.habana import GaudiConfig, GaudiSeq2SeqTrainer, GaudiSeq2SeqTrainingArguments > from optimum.habana.utils import set_seed -54,55d55 +53,54d54 < # Will error if the minimal version of Transformers is not installed. Remove at your own risks. -< check_min_version("4.46.0.dev0") -57c57,63 +< check_min_version("4.52.0.dev0") +56c56,62 < require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/summarization/requirements.txt") --- > try: @@ -34,17 +37,24 @@ > def check_optimum_habana_min_version(*a, **b): > return () > -60a67,72 +60,68c66,71 +< try: +< nltk.data.find("tokenizers/punkt") +< except (LookupError, OSError): +< if is_offline_mode(): +< raise LookupError( +< "Offline mode: run this script without TRANSFORMERS_OFFLINE first to download nltk data files" +< ) +< with FileLock(".lock") as lock: +< nltk.download("punkt", quiet=True) +--- > # Will error if the minimal version of Transformers and Optimum Habana are not installed. Remove at your own risks. -> check_min_version("4.45.0") -> check_optimum_habana_min_version("1.14.0.dev0") +> check_min_version("4.51.0") +> check_optimum_habana_min_version("1.18.0.dev0") > > require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/summarization/requirements.txt") > -70a83,84 -> nltk.download("punkt_tab") # Needed for version 3.8.2 -> -129a144,152 +128a132,140 > use_cache: bool = field( > default=True, > metadata={ @@ -54,17 +64,29 @@ > ) > }, > ) -213c236 +212c224 < "efficient on GPU but very bad for TPU." --- > "efficient on GPU but very bad for HPU in lazy mode." -261a285 +260a273 > source_suffix: Optional[str] = field(default="", metadata={"help": "A suffix to add after every source text."}) -317c341 +316c329 < parser = HfArgumentParser((ModelArguments, DataTrainingArguments, Seq2SeqTrainingArguments)) --- > parser = HfArgumentParser((ModelArguments, DataTrainingArguments, GaudiSeq2SeqTrainingArguments)) -346a371,377 +345a359,383 +> with training_args.main_process_first(): +> try: +> nltk.data.find("tokenizers/punkt") +> nltk.data.find("punkt_tab") # Needed for version 3.8.2 +> except (LookupError, OSError): +> if is_offline_mode(): +> raise LookupError( +> "Offline mode: run this script without TRANSFORMERS_OFFLINE first to download nltk data files" +> ) +> nltk.download("punkt", quiet=True) +> nltk.download("punkt_tab", quiet=True) +> > gaudi_config = GaudiConfig.from_pretrained( > training_args.gaudi_config_name, > cache_dir=model_args.cache_dir, @@ -72,25 +94,31 @@ > token=model_args.token, > ) > -347a379 +> if training_args.do_train and training_args.use_compiled_autograd: +> from habana_frameworks.torch.dynamo.compile_backend.experimental import enable_compiled_autograd +> +> enable_compiled_autograd() +> torch._C._set_autograd_fallback_mode("nothing") +> +346a385 > mixed_precision = training_args.bf16 or gaudi_config.use_torch_autocast -349,350c381,383 +348,349c387,389 < f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}, " < + f"distributed training: {training_args.parallel_mode.value == 'distributed'}, 16-bits training: {training_args.fp16}" --- > f"Process rank: {training_args.local_rank}, device: {training_args.device}, " > + f"distributed training: {training_args.parallel_mode.value == 'distributed'}, " > + f"mixed-precision training: {mixed_precision}" -432a466 +431a472 > use_cache=False if training_args.gradient_checkpointing else model_args.use_cache, -451a486,491 +450a492,497 > is_bart = model.config.model_type == "bart" > if is_bart and training_args.do_train: > raise ValueError( > "Training is not yet supported for BART. Eval or predict can be enabled with `--do_eval` and `--do_predict`." > ) > -454c494,501 +453c500,507 < embedding_size = model.get_input_embeddings().weight.shape[0] --- > embeddings = model.get_input_embeddings() @@ -101,16 +129,16 @@ > embedding_size = embeddings.weight.shape[0] > else: > embedding_size = embeddings.weight.shape[0] -487a535 +486a541 > suffix = data_args.source_suffix if data_args.source_suffix is not None else "" -558a607,608 +557a613,614 > else: > raise ValueError("Found case where either text or summary is missing.") -560c610 +559c616 < inputs = [prefix + inp for inp in inputs] --- > inputs = [prefix + inp + suffix for inp in inputs] -575a626,665 +574a632,671 > def preprocess_bucketing_function(examples): > # remove pairs where at least one record is None > @@ -151,22 +179,22 @@ > model_inputs["labels"] = labels["input_ids"] > return model_inputs > -590a681,686 +589a687,692 > def wrapper_preprocess_function(examples): > if model.config.is_encoder_decoder: > return preprocess_bucketing_function(examples) > else: > return preprocess_function(examples) > -599c695 +598c701 < preprocess_function, --- > wrapper_preprocess_function, -615c711 +614c717 < preprocess_function, --- > wrapper_preprocess_function, -625,630c721,729 +624,629c727,735 < data_collator = DataCollatorForSeq2Seq( < tokenizer, < model=model, @@ -183,7 +211,7 @@ > label_pad_token_id=label_pad_token_id, > pad_to_multiple_of=8 if training_args.fp16 else None, > ) -665,672c764,775 +664,671c770,781 < training_args.generation_max_length = ( < training_args.generation_max_length < if training_args.generation_max_length is not None @@ -205,17 +233,13 @@ > training_args.generation_config.num_beams = data_args.num_beams > elif training_args.generation_num_beams is not None: > training_args.generation_config.num_beams = training_args.generation_num_beams -675c778 +674c784 < trainer = Seq2SeqTrainer( --- > trainer = GaudiSeq2SeqTrainer( -676a780 +675a786 > gaudi_config=gaudi_config, -680c784 -< processing_class=tokenizer, ---- -> tokenizer=tokenizer, -765,769d868 +764,768d874 < < < def _mp_fn(index): diff --git a/tests/example_diff/run_translation.txt b/tests/example_diff/run_translation.txt index ad61fd0b21..1f7c845aa4 100644 --- a/tests/example_diff/run_translation.txt +++ b/tests/example_diff/run_translation.txt @@ -1,22 +1,24 @@ -30,31d29 +1a2 +> # coding=utf-8 +29,30d29 < from datasets import load_dataset < -32a31 +31a31 > from datasets import load_dataset -44,45c43 +43,44c43 < Seq2SeqTrainer, < Seq2SeqTrainingArguments, --- > NllbTokenizerFast, -47d44 +46d44 < set_seed, -52a50,51 +51a50,51 > from optimum.habana import GaudiConfig, GaudiSeq2SeqTrainer, GaudiSeq2SeqTrainingArguments > from optimum.habana.utils import set_seed -54,55d52 +53,54d52 < # Will error if the minimal version of Transformers is not installed. Remove at your own risks. -< check_min_version("4.46.0.dev0") -57c54,60 +< check_min_version("4.52.0.dev0") +56c54,60 < require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/translation/requirements.txt") --- > try: @@ -26,14 +28,14 @@ > def check_optimum_habana_min_version(*a, **b): > return () > -60a64,69 +59a64,69 > # Will error if the minimal version of Transformers and Optimum Habana are not installed. Remove at your own risks. -> check_min_version("4.45.0") -> check_optimum_habana_min_version("1.14.0.dev0") +> check_min_version("4.51.0") +> check_optimum_habana_min_version("1.18.0.dev0") > > require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/translation/requirements.txt") > -62c71,78 +61c71,78 < MULTILINGUAL_TOKENIZERS = [MBartTokenizer, MBartTokenizerFast, MBart50Tokenizer, MBart50TokenizerFast, M2M100Tokenizer] --- > MULTILINGUAL_TOKENIZERS = [ @@ -44,7 +46,7 @@ > M2M100Tokenizer, > NllbTokenizerFast, > ] -110a127,135 +109a127,135 > use_cache: bool = field( > default=True, > metadata={ @@ -54,15 +56,15 @@ > ) > }, > ) -181c206 +180c206 < "efficient on GPU but very bad for TPU." --- > "efficient on GPU but very bad for HPU in lazy mode." -266c291 +265c291 < parser = HfArgumentParser((ModelArguments, DataTrainingArguments, Seq2SeqTrainingArguments)) --- > parser = HfArgumentParser((ModelArguments, DataTrainingArguments, GaudiSeq2SeqTrainingArguments)) -295a321,327 +294a321,327 > gaudi_config = GaudiConfig.from_pretrained( > training_args.gaudi_config_name, > cache_dir=model_args.cache_dir, @@ -70,32 +72,28 @@ > token=model_args.token, > ) > -296a329 +295a329 > mixed_precision = training_args.bf16 or gaudi_config.use_torch_autocast -298,299c331,333 +297,298c331,333 < f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}, " < + f"distributed training: {training_args.parallel_mode.value == 'distributed'}, 16-bits training: {training_args.fp16}" --- > f"Process rank: {training_args.local_rank}, device: {training_args.device}, " > + f"distributed training: {training_args.parallel_mode.value == 'distributed'}, " > + f"mixed-precision training: {mixed_precision}" -385a420 +384a420 > use_cache=False if training_args.gradient_checkpointing else model_args.use_cache, -457c492 +456c492 < # Check the whether the source target length fits in the model, if it has absolute positional embeddings --- > # Check whether the source target length fits in the model, if it has absolute positional embeddings -595c630 +594c630 < trainer = Seq2SeqTrainer( --- > trainer = GaudiSeq2SeqTrainer( -596a632 +595a632 > gaudi_config=gaudi_config, -600c636 -< processing_class=tokenizer, ---- -> tokenizer=tokenizer, -689,693d724 +688,692d724 < < < def _mp_fn(index): From 9786545a711e0c8f05a3d098aa464447f238419c Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Fri, 18 Apr 2025 13:34:47 +0000 Subject: [PATCH 02/22] Upgrade `generation/utils.py` to v4.51.3 (except `_beam_search`, will be done later) --- .../habana/transformers/generation/utils.py | 297 +++++++++++------- 1 file changed, 190 insertions(+), 107 deletions(-) diff --git a/optimum/habana/transformers/generation/utils.py b/optimum/habana/transformers/generation/utils.py index baa48485ae..25cb4926af 100755 --- a/optimum/habana/transformers/generation/utils.py +++ b/optimum/habana/transformers/generation/utils.py @@ -34,9 +34,11 @@ from transformers.generation.beam_search import BeamScorer, BeamSearchScorer, ConstrainedBeamSearchScorer from transformers.generation.candidate_generator import ( AssistedCandidateGeneratorDifferentTokenizers, + AssistantVocabTranslatorCache, CandidateGenerator, EarlyExitCandidateGenerator, PromptLookupCandidateGenerator, + UniversalSpeculativeDecodingGenerator, _crop_past_key_values, _prepare_attention_mask, _prepare_token_type_ids, @@ -213,24 +215,11 @@ def _prepare_inputs_for_generation( cache_position = torch.arange(past_length, input_ids.shape[1], dtype=torch.long, device=input_ids.device) # 2. Generic cache-dependent input preparation - # If we have cache: let's slice `input_ids` through `cache_position`, to keep only the unprocessed tokens - # Exception 1: when passing input_embeds, input_ids may be missing entries - # Exception 2: some generation methods do special slicing of input_ids, so we don't need to do it here - # Exception 3: with synced GPUs cache_position may go out of bounds, but we only want dummy token in that case. - # (we can't check exception 3 while compiling) - # Excpetion 4: If input_embeds are passed then slice it through `cache_position`, to keep only the unprocessed tokens and - # generate the first token for each sequence. Later use the generated Input ids for continuation. if past_key_values is not None: model_inputs["past_key_values"] = past_key_values - if inputs_embeds is not None and input_ids.shape[1] == 0: # Exception 4 - inputs_embeds = inputs_embeds[:, -cache_position.shape[0] :] - elif ( - inputs_embeds is not None # Exception 1 - or (is_torchdynamo_compiling() or cache_position[-1] >= input_ids.shape[1]) # Exception 3 - ): - input_ids = input_ids[:, -cache_position.shape[0] :] - elif input_ids.shape[1] != cache_position.shape[0]: # Default case (the "else", a no op, is Exception 2) - input_ids = input_ids[:, cache_position] + inputs_embeds, input_ids = self._cache_dependant_input_preparation( + input_ids, inputs_embeds, cache_position + ) # 3. Prepare base model inputs input_ids_key = "decoder_input_ids" if self.config.is_encoder_decoder else "input_ids" @@ -247,6 +236,7 @@ def _prepare_inputs_for_generation( model_inputs[input_ids_key] = input_ids.clone(memory_format=torch.contiguous_format) # 4. Create missing `position_ids` on the fly + encoder_attention_mask = attention_mask if self.config.is_encoder_decoder else None attention_mask = ( kwargs.pop("decoder_attention_mask", None) if self.config.is_encoder_decoder else attention_mask ) @@ -317,6 +307,9 @@ def _prepare_inputs_for_generation( if attention_mask is not None: model_inputs[attention_mask_key] = attention_mask + if encoder_attention_mask is not None: + model_inputs["attention_mask"] = encoder_attention_mask + # 7. Forward ALL kwargs that are uninitialized (e.g. `use_cache`). for key, value in kwargs.items(): if key not in model_inputs: @@ -369,7 +362,7 @@ def _prepare_decoder_input_ids_for_generation( model_input_name: str, model_kwargs: Dict[str, torch.Tensor], decoder_start_token_id: torch.Tensor, - device: torch.device = None, + device: Optional[torch.device] = None, max_new_tokens: int = None, pad_token_id: int = None, ) -> Tuple[torch.LongTensor, Dict[str, torch.Tensor]]: @@ -761,16 +754,36 @@ def _get_candidate_generator( max_length=generation_config.max_length, ) elif different_tokenizers: - candidate_generator = AssistedCandidateGeneratorDifferentTokenizers( - input_ids=input_ids, - assistant_model=assistant_model, - generation_config=generation_config, - model_kwargs=model_kwargs, - inputs_tensor=inputs_tensor, - logits_processor=logits_processor, - target_tokenizer=target_tokenizer, - assistant_tokenizer=assistant_tokenizer, - ) + if generation_config.do_sample is True: + atm_translator = AssistantVocabTranslatorCache.get_translator( + target_tokenizer, assistant_tokenizer, self.config.vocab_size, assistant_model.device + ) + candidate_generator = UniversalSpeculativeDecodingGenerator( + input_ids=input_ids, + assistant_model=assistant_model, + generation_config=generation_config, + model_kwargs=model_kwargs, + inputs_tensor=inputs_tensor, + logits_processor=logits_processor, + target_tokenizer=target_tokenizer, + assistant_tokenizer=assistant_tokenizer, + atm_translator=atm_translator, + ) + elif generation_config.do_sample is False: + candidate_generator = AssistedCandidateGeneratorDifferentTokenizers( + input_ids=input_ids, + assistant_model=assistant_model, + generation_config=generation_config, + model_kwargs=model_kwargs, + inputs_tensor=inputs_tensor, + logits_processor=logits_processor, + target_tokenizer=target_tokenizer, + assistant_tokenizer=assistant_tokenizer, + ) + else: + raise ValueError( + f"Invalid value for `do_sample`: expected a boolean, got {type(generation_config.do_sample).__name__}" + ) else: candidate_generator = GaudiAssistedCandidateGenerator( input_ids=input_ids, @@ -885,7 +898,7 @@ def _prepare_generated_length( return generation_config def _prepare_generation_config( - self, generation_config: Optional[GaudiGenerationConfig], **kwargs: Dict + self, generation_config: Optional[GaudiGenerationConfig], use_model_defaults: Optional[bool] = None, **kwargs: Dict ) -> Tuple[GaudiGenerationConfig, Dict]: """ Copied from https://github.com/huggingface/transformers/blob/v4.40.2/src/transformers/generation/utils.py#L1230 @@ -893,11 +906,10 @@ def _prepare_generation_config( - add management of `static_shapes` and `ignore_eos` in the generation config - workaround for `token_type_ids` for Falcon """ - # TODO joao: when we can detect `fullgraph=True` in `torch.compile` (https://github.com/pytorch/pytorch/pull/120400) - # replace `is_torchdynamo_compiling` by the corresponding check. As it is, we are being too restrictive with - # the parameterization in `fullgraph=False` so as to enable `fullgraph=True`. + # parameterization priority: + # kwargs > non-global default values in `generation_config` > `model.generation_config` > GenerationConfig() + # TODO (joao): per-model generation config classes. - # priority: `generation_config` argument > `model.generation_config` (the default generation config) using_model_generation_config = False if generation_config is None: # legacy: users may modify the model configuration to control generation. To trigger this legacy behavior, @@ -906,10 +918,8 @@ def _prepare_generation_config( # 2) the generation config must have seen no modification since its creation (the hash is the same); # 3) there are non-default generation parameters in the model config. # 4) the user must have set new generation parameters in the model config. - # NOTE: `torch.compile` can't compile `hash`, this legacy support is disabled with compilation. if ( - not is_torchdynamo_compiling() - and self.generation_config._from_model_config # 1) + self.generation_config._from_model_config # 1) and self.generation_config._original_object_hash == hash(self.generation_config) # 2) and len(self.config._get_non_default_generation_parameters()) > 0 # 3) ): @@ -927,27 +937,34 @@ def _prepare_generation_config( generation_config = self.generation_config using_model_generation_config = True - # `torch.compile` can't compile `copy.deepcopy`, arguments in `kwargs` that are part of `generation_config` - # will mutate the object with `.update`. As such, passing these arguments through `kwargs` is disabled -- an - # exception will be raised in `_validate_model_kwargs` - if not is_torchdynamo_compiling(): - generation_config = copy.deepcopy(generation_config) - if generation_config.static_shapes is None: - generation_config.static_shapes = self.config.model_type in MODELS_OPTIMIZED_WITH_STATIC_SHAPES - if self.config.model_type == "vision-encoder-decoder": - generation_config.static_shapes = ( - self.config.decoder.model_type in MODELS_OPTIMIZED_WITH_STATIC_SHAPES + # `torch.export.export` usually raises an exception if it is called + # with ``strict=True``. deepcopy can only be processed if ``strict=False``. + generation_config = copy.deepcopy(generation_config) + + if not using_model_generation_config: + # If `generation_config` is provided: + # - `use_model_defaults`: let's fallback ALL default values to the model's generation config + # - otherwise: legacy behavior, let's just make sure we have the tokens defined + model_base_version = version.parse(version.parse(self.generation_config.transformers_version).base_version) + if use_model_defaults is True or ( + use_model_defaults is None and model_base_version >= version.parse("4.50.0") + ): + modified_values = {} + default_generation_config = GaudiGenerationConfig() + for key, default_value in default_generation_config.__dict__.items(): + if key.startswith("_") or key == "transformers_version": # metadata + continue + custom_gen_config_value = getattr(generation_config, key) + model_gen_config_value = getattr(self.generation_config, key) + if custom_gen_config_value == default_value and model_gen_config_value != default_value: + modified_values[key] = model_gen_config_value + setattr(generation_config, key, model_gen_config_value) + if len(modified_values) > 0: + logger.warning_once( + f"`generation_config` default values have been modified to match model-specific defaults: " + f"{modified_values}. If this is not desired, please set these values explicitly." ) - self.generation_config.static_shapes = generation_config.static_shapes - if generation_config.ignore_eos is None: - generation_config.ignore_eos = kwargs.get("ignore_eos", kwargs.get("lazy_mode", None)) - self.generation_config.ignore_eos = generation_config.ignore_eos - model_kwargs = generation_config.update(**kwargs) # All unused kwargs must be model kwargs - if self.config.model_type == "falcon" and "token_type_ids" in kwargs.keys(): - for key in ["token_type_ids"]: - model_kwargs.pop(key, None) - # If `generation_config` is provided, let's fallback ALL special tokens to the default values for the model - if not using_model_generation_config: + else: if generation_config.bos_token_id is None: generation_config.bos_token_id = self.generation_config.bos_token_id if generation_config.eos_token_id is None: @@ -956,8 +973,24 @@ def _prepare_generation_config( generation_config.pad_token_id = self.generation_config.pad_token_id if generation_config.decoder_start_token_id is None: generation_config.decoder_start_token_id = self.generation_config.decoder_start_token_id - else: - model_kwargs = kwargs + + if generation_config.static_shapes is None: + generation_config.static_shapes = self.config.model_type in MODELS_OPTIMIZED_WITH_STATIC_SHAPES + if self.config.model_type == "vision-encoder-decoder": + generation_config.static_shapes = ( + self.config.decoder.model_type in MODELS_OPTIMIZED_WITH_STATIC_SHAPES + ) + self.generation_config.static_shapes = generation_config.static_shapes + if generation_config.ignore_eos is None: + generation_config.ignore_eos = kwargs.get("ignore_eos", kwargs.get("lazy_mode", None)) + self.generation_config.ignore_eos = generation_config.ignore_eos + + # Finally, apply any passed kwargs + model_kwargs = generation_config.update(**kwargs) + + if self.config.model_type == "falcon" and "token_type_ids" in kwargs.keys(): + for key in ["token_type_ids"]: + model_kwargs.pop(key, None) return generation_config, model_kwargs @@ -1027,6 +1060,9 @@ def _prepare_cache_for_generation( ) generation_config.cache_implementation = None + generation_config.cache_implementation = generation_config.cache_implementation or getattr( + self.config.get_text_config(), "cache_implementation", None + ) if generation_config.cache_implementation is not None: if generation_config.cache_implementation in NEED_SETUP_CACHE_CLASSES_MAPPING: if generation_config.cache_implementation == "static" and not self._supports_static_cache: @@ -1069,6 +1105,8 @@ def _prepare_cache_for_generation( model_kwargs[cache_name] = cache_class(cache_config) elif generation_config.cache_implementation == "offloaded": model_kwargs[cache_name] = OffloadedCache() + elif generation_config.cache_implementation == "dynamic": + model_kwargs[cache_name] = DynamicCache() # Use tuples by default (.i.e. legacy format). else: @@ -1087,6 +1125,7 @@ def generate( streamer: Optional["BaseStreamer"] = None, negative_prompt_ids: Optional[torch.Tensor] = None, negative_prompt_attention_mask: Optional[torch.Tensor] = None, + use_model_defaults: Optional[bool] = None, lazy_mode: Optional[bool] = False, hpu_graphs: Optional[bool] = False, profiling_warmup_steps: Optional[int] = 0, @@ -1158,6 +1197,11 @@ def generate( size. This is an experimental feature, subject to breaking API changes in future versions. negative_prompt_attention_mask (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Attention_mask for `negative_prompt_ids`. + use_model_defaults (`bool`, *optional*): + When it is `True`, unset parameters in `generation_config` will be set to the model-specific default + generation configuration (`model.generation_config`), as opposed to the global defaults + (`GenerationConfig()`). If unset, models saved starting from `v4.50` will consider this flag to be + `True`. lazy_mode (`bool`, *optional*, defaults to `False`): Whether the run is executed in lazy mode or not (i.e. eager mode). hpu_graphs (`bool`, *optional*, defaults to `False`): @@ -1206,7 +1250,9 @@ def generate( "`hpu_graphs` is True but `lazy_mode` is False. HPU graphs require `lazy_mode` to be set to True." ) num_virtual_tokens = kwargs.pop("num_virtual_tokens", 0) - generation_config, model_kwargs = self._prepare_generation_config(generation_config, **kwargs) + generation_config, model_kwargs = self._prepare_generation_config( + generation_config, use_model_defaults, **kwargs + ) self._validate_model_kwargs(model_kwargs.copy()) self._validate_assistant(assistant_model, tokenizer, assistant_tokenizer) @@ -1231,7 +1277,7 @@ def generate( self._prepare_special_tokens(generation_config, kwargs_has_attention_mask, device=device) # decoder-only models must use left-padding for batched generation. - if not self.config.is_encoder_decoder and not is_torchdynamo_compiling(): + if not self.config.is_encoder_decoder: # If `input_ids` was given, check if the last id in any sequence is `pad_token_id` # Note: If using, `inputs_embeds` this check does not work, because we want to be more hands-off. if ( @@ -1544,7 +1590,7 @@ def generate( "`streamer` cannot be used with beam search (yet!). Make sure that `num_beams` is set to 1." ) - if not is_torchdynamo_compiling() and self.device.type != input_ids.device.type: + if self.device.type != input_ids.device.type: warnings.warn( ( "You are calling .generate() with the `input_ids` being on a device type different" @@ -1703,18 +1749,7 @@ def generate( ) elif generation_mode in (GenerationMode.BEAM_SAMPLE, GenerationMode.BEAM_SEARCH): - # 11. prepare beam search scorer - beam_scorer = BeamSearchScorer( - batch_size=batch_size, - num_beams=generation_config.num_beams, - device=inputs_tensor.device, - length_penalty=generation_config.length_penalty, - do_early_stopping=generation_config.early_stopping, - num_beam_hyps_to_keep=generation_config.num_return_sequences, - max_length=generation_config.max_length, - ) - - # 12. interleave input_ids with `num_beams` additional sequences per batch + # 11. interleave input_ids with `num_beams` additional sequences per batch input_ids, model_kwargs = self._expand_inputs_for_generation( input_ids=input_ids, expand_size=generation_config.num_beams, @@ -1722,10 +1757,9 @@ def generate( **model_kwargs, ) - # 13. run beam sample + # 12. run beam sample result = self._beam_search( input_ids, - beam_scorer, logits_processor=prepared_logits_processor, stopping_criteria=prepared_stopping_criteria, generation_config=generation_config, @@ -1851,7 +1885,6 @@ def typeerror(): # Convert to legacy cache format if requested if ( generation_config.return_legacy_cache is True - and not is_torchdynamo_compiling() and hasattr(result, "past_key_values") and getattr(result.past_key_values, "to_legacy_cache") is not None ): @@ -2090,9 +2123,11 @@ def _contrastive_search( else: logit_for_next_step = torch.index_select(outputs.logits, -2, token_idx - 1).squeeze(-2) else: - # .float() is needed to retain precision for later logits manipulations - logit_for_next_step = outputs.logits[:, -1, :].float() - logit_for_next_step = logit_for_next_step.to(input_ids.device) + logit_for_next_step = outputs.logits[:, -1, :] + # torch.float32 is needed to retain precision for later logits manipulations + logit_for_next_step = logit_for_next_step.to( + copy=True, dtype=torch.float32, device=input_ids.device + ) model_kwargs = self._update_model_kwargs_for_generation( outputs, @@ -2583,7 +2618,6 @@ def _sample( output_scores = generation_config.output_scores output_logits = generation_config.output_logits return_dict_in_generate = generation_config.return_dict_in_generate - max_length = generation_config.max_length has_eos_stopping_criteria = any(hasattr(criteria, "eos_token_id") for criteria in stopping_criteria) do_sample = generation_config.do_sample @@ -2636,7 +2670,7 @@ def _sample( model_kwargs["mqa_model"] = False model_kwargs["lazy_mode"] = lazy_mode while self._has_unfinished_sequences( - this_peer_finished, synced_gpus, device=input_ids.device, cur_len=cur_len, max_length=max_length + this_peer_finished, synced_gpus, device=input_ids.device ): if lazy_mode: self.htcore_generation.mark_step() @@ -2886,7 +2920,6 @@ def _sample( def _beam_search( self, input_ids: torch.LongTensor, - beam_scorer: BeamScorer, logits_processor: LogitsProcessorList, stopping_criteria: StoppingCriteriaList, generation_config: GaudiGenerationConfig, @@ -2902,12 +2935,15 @@ def _beam_search( Generates sequences of token ids for models with a language modeling head using **beam search decoding** and can be used for text-decoder, text-to-text, speech-to-text, and vision-to-text models. + If it's the first time you're diving into Beam Search, we recommend you read the following blog post: + https://huggingface.co/blog/how-to-generate (especially the beam search section). + + You can recompute the sequence scores from the individual scores using the `compute_transition_scores` function + (https://huggingface.co/docs/transformers/main_classes/text_generation#transformers.GenerationMixin.compute_transition_scores) + Parameters: - input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + input_ids (`torch.LongTensor` of shape `(batch_size*num_beams, sequence_length)`): The sequence used as a prompt for the generation. - beam_scorer (`BeamScorer`): - An derived instance of [`BeamScorer`] that defines how beam hypotheses are constructed, stored and - sorted during generation. For more information, the documentation of [`BeamScorer`] should be read. logits_processor (`LogitsProcessorList`): An instance of [`LogitsProcessorList`]. List of instances of class derived from [`LogitsProcessor`] used to modify the prediction scores of the language modeling head applied at each generation step. @@ -2938,7 +2974,8 @@ def _beam_search( `return_dict_in_generate=True` or a [`transformers.generation.GenerateBeamEncoderDecoderOutput`] if `model.config.is_encoder_decoder=True`. """ - # init values + + # 1. init beam_search values pad_token_id = generation_config._pad_token_tensor eos_token_id = generation_config._eos_token_tensor output_attentions = generation_config.output_attentions @@ -2946,13 +2983,35 @@ def _beam_search( output_scores = generation_config.output_scores output_logits = generation_config.output_logits return_dict_in_generate = generation_config.return_dict_in_generate - sequential = generation_config.low_memory do_sample = generation_config.do_sample + early_stopping = generation_config.early_stopping + length_penalty = generation_config.length_penalty + max_length = generation_config.max_length + num_beams = generation_config.num_beams + num_return_sequences = generation_config.num_return_sequences + + batch_size_unflattened, cur_len = input_ids.shape + batch_size = batch_size_unflattened // num_beams + # TODO (joao): standardize special cases + if self.__class__.__name__ == "MoshiDepthDecoder": + vocab_size = self.config.audio_vocab_size + elif self.__class__.__name__ == "ImageGPTForCausalImageModeling": + vocab_size = self.get_output_embeddings().out_features + else: + vocab_size = self.config.get_text_config().vocab_size + decoder_prompt_len = cur_len + this_peer_finished = False - batch_size = len(beam_scorer._beam_hyps) - num_beams = beam_scorer.num_beams - - batch_beam_size, cur_len = input_ids.shape + # At each beam search step, we want to keep top K [K = (number of EOS tokens + 1) * `num_beams`] candidates + # with the highest log-probabilities, or sample K continuations without replacement. We gather the top K + # (as opposed to `num_beams`, or any number lower than K) so that we have at least `num_beams` sequences + # non-finished to continue the live beam search, in case the top `num_beams` all select an EOS token. + n_eos_tokens = eos_token_id.shape[0] if eos_token_id is not None else 0 + beams_to_keep = max(2, 1 + n_eos_tokens) * num_beams + top_num_beam_mask = torch.cat( + (torch.ones((num_beams), dtype=torch.bool), torch.zeros((beams_to_keep - num_beams), dtype=torch.bool)), + dim=0, + ).to(input_ids.device) if "inputs_embeds" in model_kwargs: cur_len = model_kwargs["inputs_embeds"].shape[1] token_idx = model_kwargs.get("token_idx", None) @@ -2962,17 +3021,19 @@ def _beam_search( model_kwargs["cache_position"] = torch.arange(cur_len, device=input_ids.device) - if num_beams * batch_size != batch_beam_size: + # (joao) feature lost in the refactor. Probably won't implement, hurts readbility with minimal gains (there + # are newer low-memory alternatives like the offloaded cache) + sequential = generation_config.low_memory + if sequential: raise ValueError( - f"Batch dimension of `input_ids` should be {num_beams * batch_size}, but is {batch_beam_size}." + "`low_memory=True` is not supported after the beam search refactor. Please check the discussion in " + "#35802 *after the PR got merged*, and add a comment there if your questions are not yet answered." ) - # init attention / hidden states / scores tuples - scores = () if (return_dict_in_generate and output_scores) else None + # 2. init output tuples + all_scores = () if (return_dict_in_generate and output_scores) else None raw_logits = () if (return_dict_in_generate and output_logits) else None - beam_indices = ( - tuple(() for _ in range(batch_beam_size)) if (return_dict_in_generate and output_scores) else None - ) + beam_indices = () if (return_dict_in_generate and output_logits) else None decoder_attentions = () if (return_dict_in_generate and output_attentions) else None cross_attentions = () if (return_dict_in_generate and output_attentions) else None decoder_hidden_states = () if (return_dict_in_generate and output_hidden_states) else None @@ -2984,11 +3045,33 @@ def _beam_search( model_kwargs["encoder_outputs"].get("hidden_states") if output_hidden_states else None ) + # 3. init running tensors and static-shaped placeholders + + # per batch, beam-item holding current token in loop and completed sequences + output_fill_value = pad_token_id or eos_token_id[0] if eos_token_id is not None else -1 + running_sequences = torch.full( + (batch_size, num_beams, max_length), + fill_value=output_fill_value, + dtype=torch.int64, + device=input_ids.device, + ) + running_sequences[:, :, :cur_len] = self._unflatten_beam_dim(input_ids, batch_size, num_beams) + sequences = running_sequences.detach().clone() + + # per batch, beam-item score, logprobs # initialise score of first beam with 0 and the rest with -1e9. This makes sure that only tokens # of the first beam are considered to avoid sampling the exact same tokens across all beams. - beam_scores = torch.zeros((batch_size, num_beams), dtype=torch.float, device=input_ids.device) - beam_scores[:, 1:] = -1e9 - beam_scores = beam_scores.view((batch_size * num_beams,)) + running_beam_scores = torch.zeros((batch_size, num_beams), dtype=torch.float, device=input_ids.device) + running_beam_scores[:, 1:] = -1e9 + beam_scores = torch.full((batch_size, num_beams), fill_value=-1e9, dtype=torch.float, device=input_ids.device) + + # per batch, beam-item state bit indicating if sentence has finished. + is_sent_finished = torch.zeros((batch_size, num_beams), dtype=torch.bool, device=input_ids.device) + + # per batch, beam-item state bit indicating if there are valid continuations. + next_token_hits_stopping_criteria = torch.zeros( + (batch_size, num_beams), dtype=torch.bool, device=input_ids.device + ) # Beam token selection: pick 1 + eos_token_id.shape[0] next tokens for each beam so we have at least 1 # non eos token per beam. @@ -3489,7 +3572,7 @@ def _group_beam_search( decoding** and can be used for text-decoder, text-to-text, speech-to-text, and vision-to-text models. Parameters: - input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + input_ids (`torch.LongTensor` of shape `(batch_size*num_beams, sequence_length)`): The sequence used as a prompt for the generation. beam_scorer (`BeamScorer`): An derived instance of [`BeamScorer`] that defines how beam hypotheses are constructed, stored and @@ -3547,7 +3630,7 @@ def _constrained_beam_search( decoding** and can be used for text-decoder, text-to-text, speech-to-text, and vision-to-text models. Parameters: - input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + input_ids (`torch.LongTensor` of shape `(batch_size*num_beams, sequence_length)`): The sequence used as a prompt for the generation. constrained_beam_scorer (`ConstrainedBeamSearchScorer`): A derived instance of [`BeamScorer`] that defines how beam hypotheses are constructed, stored and @@ -3679,7 +3762,7 @@ def _constrained_beam_search( next_token_logits = torch.index_select(outputs.logits, -2, token_idx - 1).squeeze(-2) else: next_token_logits = outputs.logits[:, -1, :].float() - next_token_logits = next_token_logits.to(input_ids.device) + next_token_logits = next_token_logits.to(copy=True, dtype=torch.float32, device=input_ids.device) next_token_scores = torch.nn.functional.log_softmax( next_token_logits, dim=-1 @@ -3945,7 +4028,6 @@ def _assisted_decoding( # 1. Fetch candidate sequences from a `CandidateGenerator` and move to the correct device candidate_input_ids, candidate_logits = candidate_generator.get_candidates(input_ids[:, :cur_len]) - candidate_input_ids = candidate_input_ids.to(self.device) if candidate_logits is not None: candidate_logits = candidate_logits.to(self.device) @@ -3993,8 +4075,9 @@ def _assisted_decoding( # 2.3. Process the new logits # .float() is needed to retain precision for later logits manipulations - new_logits = outputs.logits[:, -candidate_length - 1 :].float() # excludes the input prompt if present - new_logits = new_logits.to(input_ids.device) + new_logits = outputs.logits[:, -candidate_length - 1 :].to( + dtype=torch.float32, device=input_ids.device + ) # excludes the input prompt if present next_token_logits = new_logits.clone() if len(logits_processor) > 0: for i in range(candidate_length + 1): From 0cd573e58a065ad12333b58e80c04158ec1e923d Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Mon, 21 Apr 2025 18:54:42 +0000 Subject: [PATCH 03/22] Upgrade to Transformers v4.51 --- .../habana/transformers/generation/utils.py | 7 +- .../transformers/models/bart/modeling_bart.py | 8 +- .../transformers/models/blip/modeling_blip.py | 6 +- .../transformers/models/clip/modeling_clip.py | 31 +- .../models/cohere/modeling_cohere.py | 27 +- .../models/gemma/modeling_gemma.py | 24 +- .../models/gemma2/modeling_gemma2.py | 27 +- .../transformers/models/gpt2/modeling_gpt2.py | 2 +- .../models/gpt_neox/modeling_gpt_neox.py | 18 +- .../models/idefics2/modeling_idefics2.py | 8 +- .../models/llama/modeling_llama.py | 26 +- .../models/llava/modeling_llava.py | 6 +- .../models/llava_next/modeling_llava_next.py | 4 +- .../models/mistral/modeling_mistral.py | 28 +- .../models/mixtral/modeling_mixtral.py | 38 +- .../models/mllama/modeling_mllama.py | 12 +- .../transformers/models/opt/modeling_opt.py | 8 +- .../models/paligemma/modeling_paligemma.py | 44 +- .../models/persimmon/modeling_persimmon.py | 24 +- .../transformers/models/phi/modeling_phi.py | 28 +- .../models/qwen2/modeling_qwen2.py | 28 +- .../models/qwen2_moe/modeling_qwen2_moe.py | 31 +- .../models/qwen2_vl/modeling_qwen2_vl.py | 15 +- .../seamless_m4t/modeling_seamless_m4t.py | 18 +- .../models/stablelm/modeling_stablelm.py | 24 +- .../models/starcoder2/modeling_starcoder2.py | 24 +- .../video_llava/modeling_video_llava.py | 6 +- .../transformers/models/vit/modeling_vit.py | 65 ++- .../models/wav2vec2/modeling_wav2vec2.py | 2 +- optimum/habana/transformers/trainer.py | 145 ++++--- .../habana/transformers/trainer_seq2seq.py | 26 +- optimum/habana/transformers/training_args.py | 40 +- tests/test_trainer.py | 410 +++++++++++++----- tests/test_trainer_seq2seq.py | 2 +- 34 files changed, 646 insertions(+), 566 deletions(-) diff --git a/optimum/habana/transformers/generation/utils.py b/optimum/habana/transformers/generation/utils.py index 25cb4926af..268edd593a 100755 --- a/optimum/habana/transformers/generation/utils.py +++ b/optimum/habana/transformers/generation/utils.py @@ -22,6 +22,7 @@ import torch import torch.distributed as dist +from packaging import version from transformers.cache_utils import ( Cache, DynamicCache, @@ -33,8 +34,8 @@ from transformers.generation.beam_constraints import DisjunctiveConstraint, PhrasalConstraint from transformers.generation.beam_search import BeamScorer, BeamSearchScorer, ConstrainedBeamSearchScorer from transformers.generation.candidate_generator import ( - AssistedCandidateGeneratorDifferentTokenizers, AssistantVocabTranslatorCache, + AssistedCandidateGeneratorDifferentTokenizers, CandidateGenerator, EarlyExitCandidateGenerator, PromptLookupCandidateGenerator, @@ -71,7 +72,7 @@ from transformers.integrations.deepspeed import is_deepspeed_zero3_enabled from transformers.integrations.fsdp import is_fsdp_managed_module from transformers.modeling_outputs import CausalLMOutputWithPast, Seq2SeqLMOutput -from transformers.utils import ModelOutput, is_hqq_available, is_optimum_quanto_available, is_torchdynamo_compiling +from transformers.utils import ModelOutput, is_hqq_available, is_optimum_quanto_available from optimum.utils import logging @@ -2888,7 +2889,7 @@ def _sample( + start_token_idx ) # Create a mask for positions greater than the first eos_token_id - mask = torch.arange(max_length).expand(batch_size, max_length) > eos_positions.unsqueeze(1) + mask = torch.arange(generation_config.max_length).expand(batch_size, generation_config.max_length) > eos_positions.unsqueeze(1) # Apply the mask to set positions greater than the first eos_token_id to pad_token_id input_ids[mask] = pad_token_id diff --git a/optimum/habana/transformers/models/bart/modeling_bart.py b/optimum/habana/transformers/models/bart/modeling_bart.py index 2fdfbcc6d0..16d1bbd0e7 100644 --- a/optimum/habana/transformers/models/bart/modeling_bart.py +++ b/optimum/habana/transformers/models/bart/modeling_bart.py @@ -314,7 +314,7 @@ def gaudi_BartDecoderLayer_forward( def gaudi_BartEncoder_forward( self, - input_ids: torch.LongTensor = None, + input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, head_mask: Optional[torch.Tensor] = None, inputs_embeds: Optional[torch.FloatTensor] = None, @@ -421,7 +421,7 @@ def gaudi_BartEncoder_forward( def gaudi_BartDecoder_forward( self, - input_ids: torch.LongTensor = None, + input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, encoder_hidden_states: Optional[torch.FloatTensor] = None, encoder_attention_mask: Optional[torch.LongTensor] = None, @@ -601,7 +601,7 @@ def gaudi_BartDecoder_forward( def gaudi_BartModel_forward( self, - input_ids: torch.LongTensor = None, + input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, decoder_input_ids: Optional[torch.LongTensor] = None, decoder_attention_mask: Optional[torch.LongTensor] = None, @@ -690,7 +690,7 @@ def gaudi_BartModel_forward( def gaudi_BartForConditionalGeneration_forward( self, - input_ids: torch.LongTensor = None, + input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, decoder_input_ids: Optional[torch.LongTensor] = None, decoder_attention_mask: Optional[torch.LongTensor] = None, diff --git a/optimum/habana/transformers/models/blip/modeling_blip.py b/optimum/habana/transformers/models/blip/modeling_blip.py index 512664210a..7421d7c4a5 100644 --- a/optimum/habana/transformers/models/blip/modeling_blip.py +++ b/optimum/habana/transformers/models/blip/modeling_blip.py @@ -37,7 +37,7 @@ def gaudi_BlipForConditionalGeneration_generate( image_embeds = vision_outputs[0] - image_attention_mask = torch.ones(image_embeds.size()[:-1], dtype=torch.long).to(image_embeds.device) + image_attention_mask = torch.ones(image_embeds.size()[:-1], dtype=torch.long, device=image_embeds.device) if isinstance(input_ids, list): input_ids = torch.LongTensor(input_ids) @@ -96,7 +96,7 @@ def gaudi_BlipForQuestionAnswering_generate( image_embeds = vision_outputs[0] - image_attention_mask = torch.ones(image_embeds.size()[:-1], dtype=torch.long).to(image_embeds.device) + image_attention_mask = torch.ones(image_embeds.size()[:-1], dtype=torch.long, device=image_embeds.device) if isinstance(input_ids, list): input_ids = torch.LongTensor(input_ids) @@ -111,7 +111,7 @@ def gaudi_BlipForQuestionAnswering_generate( question_embeds = question_outputs[0] - question_attention_mask = torch.ones(question_embeds.size()[:-1], dtype=torch.long).to(question_embeds.device) + question_attention_mask = torch.ones(question_embeds.size()[:-1], dtype=torch.long, device=question_embeds.device) bos_ids = torch.full( (question_embeds.size(0), 1), diff --git a/optimum/habana/transformers/models/clip/modeling_clip.py b/optimum/habana/transformers/models/clip/modeling_clip.py index 310bdef1fa..2f24e76848 100644 --- a/optimum/habana/transformers/models/clip/modeling_clip.py +++ b/optimum/habana/transformers/models/clip/modeling_clip.py @@ -1,4 +1,4 @@ -from typing import Optional, Tuple, Union +from typing import Optional, Tuple import torch from torch import nn @@ -146,7 +146,7 @@ def forward( attn_weights = self.softmax(attn_weights, dim=-1) if output_attentions: - # this operation is a bit akward, but it's required to + # this operation is a bit awkward, but it's required to # make sure that attn_weights keeps its gradient. # In order to do so, attn_weights have to reshaped # twice and have to be reused in the following @@ -232,10 +232,9 @@ def forward( causal_attention_mask: Optional[torch.Tensor] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, use_flash_attention: Optional[bool] = False, flash_attention_recompute: Optional[bool] = False, - ) -> Union[Tuple, BaseModelOutput]: + ) -> BaseModelOutput: """ Copied from CLIPEncoder.forward: https://github.com/huggingface/transformers/blob/ab0f050b42d903f34d6eb97f3f8c0c07f0517ad2/src/transformers/models/clip/modeling_clip.py The only differences are: @@ -246,7 +245,6 @@ def forward( 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 encoder_states = () if output_hidden_states else None all_attentions = () if output_attentions else None @@ -281,10 +279,10 @@ def forward( 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 + last_hidden_state=hidden_states, + hidden_states=encoder_states, + attentions=all_attentions, ) @@ -294,11 +292,10 @@ def forward( pixel_values: Optional[torch.FloatTensor] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, interpolate_pos_encoding: Optional[bool] = False, use_flash_attention: Optional[bool] = False, flash_attention_recompute: Optional[bool] = False, - ) -> Union[Tuple, BaseModelOutputWithPooling]: + ) -> BaseModelOutputWithPooling: """ Copied from CLIPVisionTransformer.forward: https://github.com/huggingface/transformers/blob/ab0f050b42d903f34d6eb97f3f8c0c07f0517ad2/src/transformers/models/clip/modeling_clip.py The only differences are: @@ -309,7 +306,6 @@ def forward( 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 if pixel_values is None: raise ValueError("You have to specify pixel_values") @@ -317,22 +313,18 @@ def forward( hidden_states = self.embeddings(pixel_values, interpolate_pos_encoding=interpolate_pos_encoding) hidden_states = self.pre_layrnorm(hidden_states) - encoder_outputs = self.encoder( + encoder_outputs: BaseModelOutput = self.encoder( inputs_embeds=hidden_states, output_attentions=output_attentions, output_hidden_states=output_hidden_states, - return_dict=return_dict, use_flash_attention=use_flash_attention, flash_attention_recompute=flash_attention_recompute, ) - last_hidden_state = encoder_outputs[0] + last_hidden_state = encoder_outputs.last_hidden_state pooled_output = last_hidden_state[:, 0, :] pooled_output = self.post_layernorm(pooled_output) - if not return_dict: - return (last_hidden_state, pooled_output) + encoder_outputs[1:] - return BaseModelOutputWithPooling( last_hidden_state=last_hidden_state, pooler_output=pooled_output, @@ -348,23 +340,20 @@ def forward( output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, interpolate_pos_encoding: bool = False, - return_dict: Optional[bool] = None, use_flash_attention: Optional[bool] = False, flash_attention_recompute: Optional[bool] = False, - ) -> Union[Tuple, BaseModelOutputWithPooling]: + ) -> BaseModelOutputWithPooling: """ Copied from CLIPVisionModel.forward: https://github.com/huggingface/transformers/blob/ab0f050b42d903f34d6eb97f3f8c0c07f0517ad2/src/transformers/models/clip/modeling_clip.py The only differences are: - add new args use_flash_attention - add new args flash_attention_recompute """ - return_dict = return_dict if return_dict is not None else self.config.use_return_dict return self.vision_model( pixel_values=pixel_values, output_attentions=output_attentions, output_hidden_states=output_hidden_states, - return_dict=return_dict, interpolate_pos_encoding=interpolate_pos_encoding, use_flash_attention=use_flash_attention, flash_attention_recompute=flash_attention_recompute, diff --git a/optimum/habana/transformers/models/cohere/modeling_cohere.py b/optimum/habana/transformers/models/cohere/modeling_cohere.py index e5ce7c1081..0890367808 100644 --- a/optimum/habana/transformers/models/cohere/modeling_cohere.py +++ b/optimum/habana/transformers/models/cohere/modeling_cohere.py @@ -1,3 +1,4 @@ +from functools import partial from typing import List, Optional, Tuple, Union import torch @@ -146,7 +147,7 @@ def forward( def gaudi_cohere_model_forward( self, - input_ids: torch.LongTensor = None, + input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[Cache] = None, @@ -154,10 +155,10 @@ def gaudi_cohere_model_forward( use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, -) -> Union[Tuple, BaseModelOutputWithPast]: + **kwargs, +) -> BaseModelOutputWithPast: """ Copied from CohereModel.forward: https://github.com/huggingface/transformers/blob/main/src/transformers/models/cohere/modeling_cohere.py The only differences are: @@ -168,7 +169,6 @@ def gaudi_cohere_model_forward( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) use_cache = use_cache if use_cache is not None else self.config.use_cache - return_dict = return_dict if return_dict is not None else self.config.use_return_dict if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") @@ -215,7 +215,7 @@ def gaudi_cohere_model_forward( if self.gradient_checkpointing and self.training: layer_outputs = self._gradient_checkpointing_func( - decoder_layer.__call__, + partial(decoder_layer.__call__, **kwargs), hidden_states, causal_mask, position_ids, @@ -254,8 +254,6 @@ def gaudi_cohere_model_forward( if return_legacy_cache: next_cache = next_cache.to_legacy_cache() - if not return_dict: - return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None) return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=next_cache, @@ -276,7 +274,7 @@ class GaudiCohereForCausalLM(CohereForCausalLM): def forward( self, - input_ids: torch.LongTensor = None, + input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, @@ -285,20 +283,18 @@ def forward( use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, logits_to_keep: Union[int, torch.Tensor] = 0, token_idx: Optional[torch.Tensor] = None, **kwargs: Unpack[KwargsForCausalLM], - ) -> Union[Tuple, CausalLMOutputWithPast]: + ) -> CausalLMOutputWithPast: 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 # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) - outputs = self.model( + outputs: BaseModelOutputWithPast = self.model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, @@ -307,12 +303,11 @@ def forward( use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, - return_dict=return_dict, cache_position=cache_position, token_idx=token_idx, ) - hidden_states = outputs[0] + hidden_states = outputs.last_hidden_state # Only compute necessary logits, and do not upcast them to float if we are not computing the loss slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep logits = self.lm_head(hidden_states[:, slice_indices, :]) @@ -322,10 +317,6 @@ def forward( if labels is not None: loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs) - if not return_dict: - output = (logits,) + outputs[1:] - return (loss,) + output if loss is not None else output - return CausalLMOutputWithPast( loss=loss, logits=logits, diff --git a/optimum/habana/transformers/models/gemma/modeling_gemma.py b/optimum/habana/transformers/models/gemma/modeling_gemma.py index d0908301c9..425db073e5 100755 --- a/optimum/habana/transformers/models/gemma/modeling_gemma.py +++ b/optimum/habana/transformers/models/gemma/modeling_gemma.py @@ -586,7 +586,7 @@ def update_sincos_cache(self, seq_len): def forward( self, - input_ids: torch.LongTensor = None, + input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, @@ -594,7 +594,6 @@ def forward( use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, attn_softmax_bf16: Optional[bool] = False, @@ -605,7 +604,7 @@ def forward( cache_idx: int = None, lazy_mode: Optional[bool] = True, **kwargs, # NOOP kwarg for now - ) -> Union[Tuple, BaseModelOutputWithPast]: + ) -> BaseModelOutputWithPast: """ Copied from GemmaModel.forward: https://github.com/huggingface/transformers/blob/v4.38.1/src/transformers/models/gemma/modeling_gemma.py The only differences are: @@ -616,7 +615,6 @@ def forward( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) use_cache = use_cache if use_cache is not None else self.config.use_cache - return_dict = return_dict if return_dict is not None else self.config.use_return_dict self._attn_implementation = "eager" @@ -748,8 +746,7 @@ def forward( next_cache = ( next_decoder_cache.to_legacy_cache() if isinstance(next_decoder_cache, Cache) else next_decoder_cache ) - if not return_dict: - return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None) + return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=next_cache, @@ -770,7 +767,7 @@ def update_sincos_cache(self, seq_len): def forward( self, - input_ids: torch.LongTensor = None, + input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, @@ -780,7 +777,6 @@ def forward( reuse_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, logits_to_keep: Union[int, torch.Tensor] = 0, token_idx: Optional[torch.Tensor] = None, @@ -789,7 +785,7 @@ def forward( flash_attention_causal_mask: Optional[bool] = False, attn_softmax_bf16: Optional[bool] = False, **kwargs: Unpack[KwargsForCausalLM], - ) -> Union[Tuple, CausalLMOutputWithPast]: + ) -> CausalLMOutputWithPast: """ Inherits from GemmaForCausalLM: https://github.com/huggingface/transformers/blob/v4.38.1/src/transformers/models/gemma/modeling_gemma.py The only differences are: @@ -800,10 +796,9 @@ def forward( 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 # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) - outputs = self.model( + outputs: BaseModelOutputWithPast = self.model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, @@ -813,7 +808,6 @@ def forward( reuse_cache=reuse_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, - return_dict=return_dict, cache_position=cache_position, token_idx=token_idx, use_flash_attention=use_flash_attention, @@ -822,7 +816,7 @@ def forward( attn_softmax_bf16=attn_softmax_bf16, ) - hidden_states = outputs[0] + hidden_states = outputs.last_hidden_state # Only compute necessary logits, and do not upcast them to float if we are not computing the loss slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep logits = self.lm_head(hidden_states[:, slice_indices, :]).float() @@ -831,10 +825,6 @@ def forward( if labels is not None: loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs) - if not return_dict: - output = (logits,) + outputs[1:] - return (loss,) + output if loss is not None else output - return CausalLMOutputWithPast( loss=loss, logits=logits, diff --git a/optimum/habana/transformers/models/gemma2/modeling_gemma2.py b/optimum/habana/transformers/models/gemma2/modeling_gemma2.py index cb4c6ab65f..595f5fc7d0 100755 --- a/optimum/habana/transformers/models/gemma2/modeling_gemma2.py +++ b/optimum/habana/transformers/models/gemma2/modeling_gemma2.py @@ -16,6 +16,7 @@ """PyTorch Gemma2 model.""" import math +from functools import partial from typing import List, Optional, Tuple, Union import torch @@ -682,7 +683,7 @@ def update_sincos_cache(self, seq_len): def forward( self, - input_ids: torch.LongTensor = None, + input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, @@ -690,7 +691,6 @@ def forward( use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, last_cache_position: Optional[int] = None, token_idx: Optional[torch.Tensor] = None, @@ -702,7 +702,8 @@ def forward( flash_attention_fast_softmax: Optional[bool] = False, cache_idx: int = None, lazy_mode: Optional[bool] = True, - ) -> Union[Tuple, BaseModelOutputWithPast]: + **kwargs, + ) -> BaseModelOutputWithPast: """ Copied from GemmaModel.forward: https://github.com/huggingface/transformers/blob/v4.38.1/src/transformers/models/gemma/modeling_gemma.py The only differences are: @@ -714,7 +715,6 @@ def forward( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) use_cache = use_cache if use_cache is not None else self.config.use_cache - return_dict = return_dict if return_dict is not None else self.config.use_return_dict self._attn_implementation = "eager" @@ -809,7 +809,7 @@ def forward( if self.gradient_checkpointing and self.training: layer_outputs = self._gradient_checkpointing_func( - decoder_layer.__call__, + partial(decoder_layer.__call__, **kwargs), hidden_states, causal_mask, position_ids, @@ -866,8 +866,6 @@ def forward( next_cache = ( next_decoder_cache.to_legacy_cache() if isinstance(next_decoder_cache, Cache) else next_decoder_cache ) - if not return_dict: - return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None) return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=next_cache, @@ -888,7 +886,7 @@ def update_sincos_cache(self, seq_len): def forward( self, - input_ids: torch.LongTensor = None, + input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, @@ -897,7 +895,6 @@ def forward( use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, logits_to_keep: Union[int, torch.Tensor] = 0, token_idx: Optional[torch.Tensor] = None, @@ -911,7 +908,7 @@ def forward( cache_idx: int = None, lazy_mode: Optional[bool] = True, **loss_kwargs, - ) -> Union[Tuple, CausalLMOutputWithPast]: + ) -> CausalLMOutputWithPast: """ Inherits from GemmaForCausalLM: https://github.com/huggingface/transformers/blob/v4.38.1/src/transformers/models/gemma/modeling_gemma.py The only differences are: @@ -921,10 +918,9 @@ def forward( 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 # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) - outputs = self.model( + outputs: BaseModelOutputWithPast = self.model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, @@ -933,7 +929,6 @@ def forward( use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, - return_dict=return_dict, cache_position=cache_position, token_idx=token_idx, attn_softmax_bf16=attn_softmax_bf16, @@ -947,7 +942,7 @@ def forward( **loss_kwargs, ) - hidden_states = outputs[0] + hidden_states = outputs.last_hidden_state _, seq_len, _ = hidden_states.shape if seq_len > 1 and trim_logits and not self.training: @@ -967,10 +962,6 @@ def forward( if labels is not None: loss = self.loss_function(logits, labels, self.vocab_size, **loss_kwargs) - if not return_dict: - output = (logits,) + outputs[1:] - return (loss,) + output if loss is not None else output - return CausalLMOutputWithPast( loss=loss, logits=logits, diff --git a/optimum/habana/transformers/models/gpt2/modeling_gpt2.py b/optimum/habana/transformers/models/gpt2/modeling_gpt2.py index ed2f0d0134..ab7460c098 100644 --- a/optimum/habana/transformers/models/gpt2/modeling_gpt2.py +++ b/optimum/habana/transformers/models/gpt2/modeling_gpt2.py @@ -50,7 +50,7 @@ def _upcast_and_reordered_attn(self, query, key, value, attention_mask=None, hea mask_value = torch.finfo(attn_weights.dtype).min # Need to be a tensor, otherwise we get error: `RuntimeError: expected scalar type float but found double`. # Need to be on the same device, otherwise `RuntimeError: ..., x and y to be on the same device` - mask_value = torch.tensor(mask_value, dtype=attn_weights.dtype).to(attn_weights.device) + mask_value = torch.tensor(mask_value, dtype=attn_weights.dtype, device=attn_weights.device) attn_weights = torch.where(causal_mask, attn_weights, mask_value) if attention_mask is not None: diff --git a/optimum/habana/transformers/models/gpt_neox/modeling_gpt_neox.py b/optimum/habana/transformers/models/gpt_neox/modeling_gpt_neox.py index bca96fb5c9..7a93ed8c07 100644 --- a/optimum/habana/transformers/models/gpt_neox/modeling_gpt_neox.py +++ b/optimum/habana/transformers/models/gpt_neox/modeling_gpt_neox.py @@ -266,11 +266,10 @@ def gaudi_gpt_neox_model_forward( use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, **kwargs, -) -> Union[Tuple, BaseModelOutputWithPast]: +) -> BaseModelOutputWithPast: """ Copied from GPTNeoxModel.forward: https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt_neox/modeling_gpt_neox.py The only differences are: @@ -280,7 +279,6 @@ def gaudi_gpt_neox_model_forward( 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 use_cache = use_cache if use_cache is not None else self.config.use_cache if input_ids is not None and inputs_embeds is not None: @@ -377,9 +375,6 @@ def gaudi_gpt_neox_model_forward( if output_hidden_states: all_hidden_states = all_hidden_states + (hidden_states,) - if not return_dict: - return tuple(v for v in [hidden_states, presents, all_hidden_states, all_attentions] if v is not None) - return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=presents, @@ -420,15 +415,13 @@ def forward( use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, logits_to_keep: Union[int, torch.Tensor] = 0, **kwargs: Unpack[KwargsForCausalLM], ) -> Union[Tuple, CausalLMOutputWithPast]: - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - outputs = self.gpt_neox( + outputs: BaseModelOutputWithPast = self.gpt_neox( input_ids, attention_mask=attention_mask, position_ids=position_ids, @@ -438,13 +431,12 @@ def forward( use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, - return_dict=return_dict, cache_position=cache_position, token_idx=token_idx, **kwargs, ) - hidden_states = outputs[0] + hidden_states = outputs.last_hidden_state # Only compute necessary logits, and do not upcast them to float if we are not computing the loss slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep logits = self.embed_out(hidden_states[:, slice_indices, :]) @@ -453,10 +445,6 @@ def forward( if labels is not None: loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs) - if not return_dict: - output = (logits,) + outputs[1:] - return ((loss,) + output) if loss is not None else output - return CausalLMOutputWithPast( loss=loss, logits=logits, diff --git a/optimum/habana/transformers/models/idefics2/modeling_idefics2.py b/optimum/habana/transformers/models/idefics2/modeling_idefics2.py index cf4dd06452..efc74746c6 100644 --- a/optimum/habana/transformers/models/idefics2/modeling_idefics2.py +++ b/optimum/habana/transformers/models/idefics2/modeling_idefics2.py @@ -74,7 +74,7 @@ def forward(self, pixel_values: torch.FloatTensor, patch_attention_mask: torch.B class GaudiIdefics2Model(Idefics2Model): def forward( self, - input_ids: torch.LongTensor = None, + input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, @@ -85,6 +85,7 @@ def forward( use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, return_dict: Optional[bool] = None, ) -> Union[Tuple, Idefics2BaseModelOutputWithPast]: """ @@ -198,6 +199,7 @@ def forward( use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, + cache_position=cache_position, return_dict=return_dict, ) @@ -244,7 +246,7 @@ def inputs_merger( class GaudiIdefics2ForConditionalGeneration(Idefics2ForConditionalGeneration): def forward( self, - input_ids: torch.LongTensor = None, + input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, @@ -257,6 +259,7 @@ def forward( output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, logits_to_keep: Union[int, torch.Tensor] = 0, token_idx: Optional[torch.Tensor] = None, ) -> Union[Tuple, Idefics2CausalLMOutputWithPast]: @@ -380,6 +383,7 @@ def forward( use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, + cache_position=cache_position, return_dict=return_dict, ) diff --git a/optimum/habana/transformers/models/llama/modeling_llama.py b/optimum/habana/transformers/models/llama/modeling_llama.py index d467a17850..ea1e9df1ec 100755 --- a/optimum/habana/transformers/models/llama/modeling_llama.py +++ b/optimum/habana/transformers/models/llama/modeling_llama.py @@ -1,4 +1,5 @@ import copy +from functools import partial from typing import List, Optional, Tuple, Union import torch @@ -1175,7 +1176,7 @@ def update_sincos_cache(self, seq_len): def forward( self, - input_ids: torch.LongTensor = None, + input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, @@ -1183,7 +1184,6 @@ def forward( use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, attn_softmax_bf16: Optional[bool] = False, @@ -1198,7 +1198,7 @@ def forward( num_virtual_tokens: int = None, attn_batch_split: int = 1, **kwargs, - ) -> Union[Tuple, BaseModelOutputWithPast]: + ) -> BaseModelOutputWithPast: """ Copied from LlamaModel.forward: https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py The only differences are: @@ -1216,7 +1216,6 @@ def forward( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) use_cache = use_cache if use_cache is not None else self.config.use_cache - return_dict = return_dict if return_dict is not None else self.config.use_return_dict if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") @@ -1328,7 +1327,7 @@ def forward( if self.gradient_checkpointing and self.training: layer_outputs = self._gradient_checkpointing_func( - decoder_layer.__call__, + partial(decoder_layer.__call__, **kwargs), hidden_states, causal_mask, position_ids, @@ -1402,8 +1401,6 @@ def forward( if not use_new_cache and isinstance(next_cache, Cache): next_cache = next_cache.to_legacy_cache() - if not return_dict: - return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None) return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=next_cache, @@ -1443,7 +1440,7 @@ def update_sincos_cache(self, seq_len): def forward( self, - input_ids: torch.LongTensor = None, + input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, @@ -1452,7 +1449,6 @@ def forward( use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, logits_to_keep: Union[int, torch.Tensor] = 0, token_idx: Optional[torch.Tensor] = None, @@ -1469,18 +1465,17 @@ def forward( num_virtual_tokens: int = None, attn_batch_split: int = 1, **kwargs: Unpack[KwargsForCausalLM], - ) -> Union[Tuple, CausalLMOutputWithPast]: + ) -> CausalLMOutputWithPast: 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 if self.generation_config.use_fused_rope is False: global has_fused_rope has_fused_rope = False # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) - outputs = self.model( + outputs: BaseModelOutputWithPast = self.model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, @@ -1489,7 +1484,6 @@ def forward( use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, - return_dict=return_dict, cache_position=cache_position, token_idx=token_idx, attn_softmax_bf16=attn_softmax_bf16, @@ -1506,7 +1500,7 @@ def forward( **kwargs, ) - hidden_states = outputs[0] + hidden_states = outputs.last_hidden_state _, seq_len, _ = hidden_states.shape if seq_len > 1 and trim_logits and not self.training: if token_idx is not None: @@ -1522,10 +1516,6 @@ def forward( if labels is not None: loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs) - if not return_dict: - output = (logits,) + outputs[1:] - return (loss,) + output if loss is not None else output - return CausalLMOutputWithPast( loss=loss, logits=logits, diff --git a/optimum/habana/transformers/models/llava/modeling_llava.py b/optimum/habana/transformers/models/llava/modeling_llava.py index a6096644a8..4db9a1b36f 100644 --- a/optimum/habana/transformers/models/llava/modeling_llava.py +++ b/optimum/habana/transformers/models/llava/modeling_llava.py @@ -108,8 +108,8 @@ def _merge_input_ids_with_image_features(image_features, inputs_embeds, input_id class GaudiLlavaForConditionalGeneration(LlavaForConditionalGeneration): def forward( self, - input_ids: torch.LongTensor = None, - pixel_values: torch.FloatTensor = None, + input_ids: Optional[torch.LongTensor] = None, + pixel_values: Optional[torch.FloatTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, @@ -123,7 +123,7 @@ def forward( return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, logits_to_keep: Union[int, torch.Tensor] = 0, - image_sizes: torch.Tensor = None, + image_sizes: Optional[torch.Tensor] = None, token_idx: Optional[torch.Tensor] = None, image_offset: Optional[int] = None, tokens_pos: Optional[torch.LongTensor] = None, diff --git a/optimum/habana/transformers/models/llava_next/modeling_llava_next.py b/optimum/habana/transformers/models/llava_next/modeling_llava_next.py index a61ef20599..6be69d6bb1 100644 --- a/optimum/habana/transformers/models/llava_next/modeling_llava_next.py +++ b/optimum/habana/transformers/models/llava_next/modeling_llava_next.py @@ -39,8 +39,8 @@ class GaudiLlavaNextForConditionalGeneration(LlavaNextForConditionalGeneration): def forward( self, - input_ids: torch.LongTensor = None, - pixel_values: torch.FloatTensor = None, + input_ids: Optional[torch.LongTensor] = None, + pixel_values: Optional[torch.FloatTensor] = None, image_sizes: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, diff --git a/optimum/habana/transformers/models/mistral/modeling_mistral.py b/optimum/habana/transformers/models/mistral/modeling_mistral.py index 38e3a4d3f4..4c2f469f60 100644 --- a/optimum/habana/transformers/models/mistral/modeling_mistral.py +++ b/optimum/habana/transformers/models/mistral/modeling_mistral.py @@ -20,6 +20,7 @@ """PyTorch Mistral model.""" import os +from functools import partial from typing import List, Optional, Tuple, Union import habana_frameworks.torch.core as htcore @@ -438,7 +439,7 @@ def update_sincos_cache(self, seq_len): def forward( self, - input_ids: torch.LongTensor = None, + input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, @@ -446,7 +447,6 @@ def forward( use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, reuse_cache: Optional[bool] = False, @@ -456,7 +456,8 @@ def forward( use_flash_attention: Optional[bool] = False, flash_attention_recompute: Optional[bool] = False, flash_attention_causal_mask: Optional[bool] = False, - ) -> Union[Tuple, BaseModelOutputWithPast]: + **kwargs, + ) -> BaseModelOutputWithPast: """ Copied from MistralModel.forward: https://github.com/huggingface/transformers/blob/v4.34.1/src/transformers/models/mistral/modeling_mistral.py The only differences are: @@ -469,8 +470,6 @@ def forward( ) use_cache = use_cache if use_cache is not None else self.config.use_cache - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - if input_ids is not None and inputs_embeds is not None: raise ValueError("You must specify exactly one of input_ids or inputs_embeds") elif input_ids is not None: @@ -541,7 +540,7 @@ def forward( if self.gradient_checkpointing and self.training: layer_outputs = self._gradient_checkpointing_func( - decoder_layer.__call__, + partial(decoder_layer.__call__, **kwargs), hidden_states, causal_mask, position_ids, @@ -597,8 +596,6 @@ def forward( else (next_decoder_cache.to_legacy_cache() if return_legacy_cache else next_decoder_cache) ) - if not return_dict: - return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None) return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=next_cache, @@ -619,7 +616,7 @@ def update_sincos_cache(self, seq_len): def forward( self, - input_ids: torch.LongTensor = None, + input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, @@ -628,7 +625,6 @@ def forward( use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, logits_to_keep: Union[int, torch.Tensor] = 0, token_idx: Optional[torch.Tensor] = None, @@ -641,7 +637,7 @@ def forward( flash_attention_recompute: Optional[bool] = False, flash_attention_causal_mask: Optional[bool] = False, **kwargs: Unpack[KwargsForCausalLM], - ) -> Union[Tuple, CausalLMOutputWithPast]: + ) -> CausalLMOutputWithPast: """ Inherits from MistralForCausalLM: https://github.com/huggingface/transformers/blob/v4.34.1/src/transformers/models/mistral/modeling_mistral.py The only differences are: @@ -652,14 +648,13 @@ def forward( 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 if self.generation_config.use_fused_rope is False: global has_fused_rope has_fused_rope = False # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) - outputs = self.model( + outputs: BaseModelOutputWithPast = self.model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, @@ -668,7 +663,6 @@ def forward( use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, - return_dict=return_dict, cache_position=cache_position, token_idx=token_idx, reuse_cache=reuse_cache, @@ -679,7 +673,7 @@ def forward( flash_attention_recompute=flash_attention_recompute, flash_attention_causal_mask=flash_attention_causal_mask, ) - hidden_states = outputs[0] + hidden_states = outputs.last_hidden_state _, seq_len, _ = hidden_states.shape if seq_len > 1 and trim_logits and not self.training: if token_idx is not None: @@ -694,10 +688,6 @@ def forward( if labels is not None: loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs) - if not return_dict: - output = (logits,) + outputs[1:] - return (loss,) + output if loss is not None else output - return CausalLMOutputWithPast( loss=loss, logits=logits, diff --git a/optimum/habana/transformers/models/mixtral/modeling_mixtral.py b/optimum/habana/transformers/models/mixtral/modeling_mixtral.py index 2c9e6ba2f1..6bcf268c3d 100644 --- a/optimum/habana/transformers/models/mixtral/modeling_mixtral.py +++ b/optimum/habana/transformers/models/mixtral/modeling_mixtral.py @@ -21,6 +21,7 @@ """PyTorch Mixtral model.""" import math +from functools import partial from typing import List, Optional, Tuple, Union import habana_frameworks.torch.core as htcore @@ -33,8 +34,6 @@ _prepare_4d_causal_attention_mask_for_sdpa, ) from transformers.modeling_outputs import ( - BaseModelOutputWithPast, - CausalLMOutputWithPast, MoeCausalLMOutputWithPast, MoeModelOutputWithPast, ) @@ -537,7 +536,7 @@ def allocate_kv_cache(self, batch_size, max_seq_len, inp_seq_len): def forward( self, - input_ids: torch.LongTensor = None, + input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, @@ -546,13 +545,13 @@ def forward( output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, output_router_logits: Optional[bool] = None, - return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, reuse_cache: Optional[bool] = False, flash_attention_recompute: Optional[bool] = False, cache_idx: int = None, - ) -> Union[Tuple, BaseModelOutputWithPast]: + **kwargs, + ) -> MoeModelOutputWithPast: """ Copied from MixtralModel.forward: https://github.com/huggingface/transformers/blob/v4.37.0/src/transformers/models/mixtral/modeling_mixtral.py#L1069 The only differences are: @@ -570,8 +569,6 @@ def forward( ) use_cache = use_cache if use_cache is not None else self.config.use_cache - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - # retrieve input_ids and inputs_embeds if input_ids is not None and inputs_embeds is not None: raise ValueError("You must specify exactly one of input_ids or inputs_embeds") @@ -657,7 +654,7 @@ def forward( if self.gradient_checkpointing and self.training: layer_outputs = self._gradient_checkpointing_func( - decoder_layer.__call__, + partial(decoder_layer.__call__, **kwargs), hidden_states, attention_mask, position_ids, @@ -708,12 +705,6 @@ def forward( next_decoder_cache.to_legacy_cache() if isinstance(next_decoder_cache, Cache) else next_decoder_cache ) - if not return_dict: - return tuple( - v - for v in [hidden_states, next_cache, all_hidden_states, all_self_attns, all_router_logits] - if v is not None - ) return MoeModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=next_cache, @@ -739,7 +730,7 @@ def allocate_kv_cache(self, batch_size, max_seq_len, inp_seq_len): def forward( self, - input_ids: torch.LongTensor = None, + input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, @@ -749,7 +740,6 @@ def forward( output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, output_router_logits: Optional[bool] = None, - return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, logits_to_keep: Union[int, torch.Tensor] = 0, token_idx: Optional[torch.Tensor] = None, @@ -757,7 +747,7 @@ def forward( flash_attention_recompute: Optional[bool] = False, cache_idx: int = None, **kwargs: Unpack[KwargsForCausalLM], - ) -> Union[Tuple, CausalLMOutputWithPast]: + ) -> MoeCausalLMOutputWithPast: output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_router_logits = ( output_router_logits if output_router_logits is not None else self.config.output_router_logits @@ -766,10 +756,9 @@ def forward( 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 # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) - outputs = self.model( + outputs: MoeModelOutputWithPast = self.model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, @@ -779,7 +768,6 @@ def forward( output_attentions=output_attentions, output_hidden_states=output_hidden_states, output_router_logits=output_router_logits, - return_dict=return_dict, cache_position=cache_position, token_idx=token_idx, reuse_cache=reuse_cache, @@ -787,7 +775,7 @@ def forward( cache_idx=cache_idx, ) - hidden_states = outputs[0] + hidden_states = outputs.last_hidden_state # Only compute necessary logits, and do not upcast them to float if we are not computing the loss slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep logits = self.lm_head(hidden_states[:, slice_indices, :]).float() @@ -799,7 +787,7 @@ def forward( aux_loss = None if output_router_logits: aux_loss = load_balancing_loss_func( - outputs.router_logits if return_dict else outputs[-1], + outputs.router_logits, self.num_experts, self.num_experts_per_tok, attention_mask, @@ -807,12 +795,6 @@ def forward( if labels is not None: loss += self.router_aux_loss_coef * aux_loss.to(loss.device) # make sure to reside in the same device - if not return_dict: - output = (logits,) + outputs[1:] - if output_router_logits: - output = (aux_loss,) + output - return (loss,) + output if loss is not None else output - return MoeCausalLMOutputWithPast( loss=loss, aux_loss=aux_loss, diff --git a/optimum/habana/transformers/models/mllama/modeling_mllama.py b/optimum/habana/transformers/models/mllama/modeling_mllama.py index 665c40d5e5..7c2481431b 100644 --- a/optimum/habana/transformers/models/mllama/modeling_mllama.py +++ b/optimum/habana/transformers/models/mllama/modeling_mllama.py @@ -158,7 +158,7 @@ def forward( self, hidden_state: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, - output_attentions: bool = None, + output_attentions: Optional[bool] = None, use_flash_attention: Optional[bool] = False, ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented. @@ -214,7 +214,7 @@ def forward( self, hidden_state: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, - output_attentions: bool = None, + output_attentions: Optional[bool] = None, use_flash_attention: Optional[bool] = False, ): """ @@ -317,7 +317,7 @@ def forward( past_key_value: Optional[Cache] = None, attention_mask: Optional[torch.Tensor] = None, output_attentions: bool = False, - use_cache: bool = None, + use_cache: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, use_flash_attention: Optional[bool] = False, @@ -890,7 +890,7 @@ def _update_causal_mask( class GaudiMllamaForCausalLM(MllamaForCausalLM): def forward( self, - input_ids: torch.LongTensor = None, + input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, cross_attention_states: Optional[torch.LongTensor] = None, @@ -1001,7 +1001,7 @@ def forward( use_flash_attention: Optional[bool] = False, flash_attention_recompute: Optional[bool] = False, logits_bf16: Optional[bool] = False, - **kwargs, + **loss_kwargs, ) -> Union[Tuple, CausalLMOutputWithPast]: """ Copied from MllamaForConditionalGeneration::forward: https://github.com/huggingface/transformers/blob/v4.45.2/src/transformers/models/mllama/modeling_mllama.py#L2077 @@ -1077,7 +1077,6 @@ def forward( past_key_values=past_key_values, use_cache=use_cache, inputs_embeds=inputs_embeds, - labels=labels, output_hidden_states=output_hidden_states, output_attentions=output_attentions, return_dict=return_dict, @@ -1087,6 +1086,7 @@ def forward( use_flash_attention=use_flash_attention, flash_attention_recompute=flash_attention_recompute, logits_bf16=logits_bf16, + **loss_kwargs, ) return outputs diff --git a/optimum/habana/transformers/models/opt/modeling_opt.py b/optimum/habana/transformers/models/opt/modeling_opt.py index 2b0fa0c99b..54e91c7486 100644 --- a/optimum/habana/transformers/models/opt/modeling_opt.py +++ b/optimum/habana/transformers/models/opt/modeling_opt.py @@ -167,7 +167,7 @@ def gaudi_opt_attention_forward( class GaudiOPTDecoderLayer(torch.nn.Module): - def __init__(self, config: OPTConfig, layer_idx: int = None): + def __init__(self, config: OPTConfig, layer_idx: Optional[int] = None): """ Attention implementation is set to "eager" (default in Transformers is "sdpa"). """ @@ -262,7 +262,7 @@ def forward( def gaudi_opt_decoder_forward( self, - input_ids: torch.LongTensor = None, + input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, head_mask: Optional[torch.Tensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, @@ -418,7 +418,7 @@ def gaudi_opt_decoder_forward( def gaudi_opt_model_forward( self, - input_ids: torch.LongTensor = None, + input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, head_mask: Optional[torch.Tensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, @@ -479,7 +479,7 @@ class GaudiOPTForCausalLM(OPTForCausalLM): def forward( self, - input_ids: torch.LongTensor = None, + input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, head_mask: Optional[torch.Tensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, diff --git a/optimum/habana/transformers/models/paligemma/modeling_paligemma.py b/optimum/habana/transformers/models/paligemma/modeling_paligemma.py index ade847111e..31b7e35be0 100644 --- a/optimum/habana/transformers/models/paligemma/modeling_paligemma.py +++ b/optimum/habana/transformers/models/paligemma/modeling_paligemma.py @@ -20,6 +20,7 @@ import torch.utils.checkpoint from torch import nn from transformers.cache_utils import Cache +from transformers.modeling_outputs import CausalLMOutputWithPast from transformers.models.paligemma.modeling_paligemma import ( PaliGemmaCausalLMOutputWithPast, PaliGemmaForConditionalGeneration, @@ -33,8 +34,8 @@ class GaudiPaliGemmaForConditionalGeneration(PaliGemmaForConditionalGeneration): def forward( self, - input_ids: torch.LongTensor = None, - pixel_values: torch.FloatTensor = None, + input_ids: Optional[torch.LongTensor] = None, + pixel_values: Optional[torch.FloatTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[Union[List[torch.FloatTensor], Cache]] = None, @@ -59,11 +60,6 @@ def forward( if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") - if pixel_values is not None and inputs_embeds is not None: - raise ValueError( - "You cannot specify both pixel_values and inputs_embeds at the same time, and must specify either one" - ) - 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 @@ -72,8 +68,16 @@ def forward( is_training = token_type_ids is not None and labels is not None + # Replace image id woth PAD if the image token if OOV, to avoid index-errors + if input_ids is not None and self.config.image_token_index >= self.vocab_size: + special_image_mask = input_ids == self.config.image_token_index + llm_input_ids = input_ids.clone() + llm_input_ids[special_image_mask] = 0 + else: + llm_input_ids = input_ids + if inputs_embeds is None: - inputs_embeds = self.get_input_embeddings()(input_ids) + inputs_embeds = self.get_input_embeddings()(llm_input_ids) if cache_position is None: past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 @@ -88,10 +92,16 @@ def forward( if pixel_values is not None: image_features = self.get_image_features(pixel_values) - special_image_mask = (input_ids == self.config.image_token_index).unsqueeze(-1) - special_image_mask = special_image_mask.expand_as(inputs_embeds).to(inputs_embeds.device) + if input_ids is None: + special_image_mask = inputs_embeds == self.get_input_embeddings()( + torch.tensor(self.config.image_token_index, dtype=torch.long, device=inputs_embeds.device) + ) + else: + special_image_mask = (input_ids == self.config.image_token_index).unsqueeze(-1) + special_image_mask = special_image_mask.expand_as(inputs_embeds).to(inputs_embeds.device) + if not is_torchdynamo_compiling() and inputs_embeds[special_image_mask].numel() != image_features.numel(): - image_tokens_in_text = torch.sum(input_ids == self.config.image_token_index) + image_tokens_in_text = (special_image_mask).sum(dim=1).sum(dim=0)[0] raise ValueError( f"Number of images does not match number of special image tokens in the input text. " f"Got {image_tokens_in_text} image tokens in the text but {image_features.shape[0] * image_features.shape[1]} " @@ -111,7 +121,7 @@ def forward( causal_mask = self._update_causal_mask( attention_mask, token_type_ids, past_key_values, cache_position, inputs_embeds, is_training ) - outputs = self.language_model( + outputs: CausalLMOutputWithPast = self.language_model( attention_mask=causal_mask, position_ids=position_ids, past_key_values=past_key_values, @@ -119,7 +129,7 @@ def forward( use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, - return_dict=return_dict, + return_dict=True, cache_position=cache_position, # TODO: from Transformers v4.45, `generate` sets `num_logits_to_keep` to 1 if not given, which we don't want here # logits_to_keep=logits_to_keep, @@ -127,7 +137,7 @@ def forward( **lm_kwargs, ) - logits = outputs.logits + logits = outputs[0] loss = None if labels is not None: # Upcast to float if we need to compute the loss to avoid potential precision issues @@ -149,11 +159,8 @@ def forward( flat_logits = shift_logits.view(-1, self.config.text_config.vocab_size) flat_labels = shift_labels.view(-1).to(shift_logits.device) loss = loss_fct(flat_logits, flat_labels) - if not return_dict: - output = (logits,) + outputs[1:] - return (loss,) + output if loss is not None else output - return PaliGemmaCausalLMOutputWithPast( + output = PaliGemmaCausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=outputs.past_key_values, @@ -161,3 +168,4 @@ def forward( attentions=outputs.attentions, image_hidden_states=image_features if pixel_values is not None else None, ) + return output if return_dict else output.to_tuple() diff --git a/optimum/habana/transformers/models/persimmon/modeling_persimmon.py b/optimum/habana/transformers/models/persimmon/modeling_persimmon.py index 1c02f414e0..408600280d 100644 --- a/optimum/habana/transformers/models/persimmon/modeling_persimmon.py +++ b/optimum/habana/transformers/models/persimmon/modeling_persimmon.py @@ -214,7 +214,7 @@ def forward( def gaudi_persimmon_model_forward( self, - input_ids: torch.LongTensor = None, + input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, @@ -222,10 +222,9 @@ def gaudi_persimmon_model_forward( use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, -) -> Union[Tuple, BaseModelOutputWithPast]: +) -> BaseModelOutputWithPast: """ Copied from PersimmonModel.forward: https://github.com/huggingface/transformers/blob/v4.38.2/src/transformers/models/persimmon/modeling_persimmon.py The only differences are: @@ -239,8 +238,6 @@ def gaudi_persimmon_model_forward( ) use_cache = use_cache if use_cache is not None else self.config.use_cache - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - # retrieve input_ids and inputs_embeds if input_ids is not None and inputs_embeds is not None: raise ValueError("You must specify exactly one of input_ids or inputs_embeds") @@ -338,8 +335,6 @@ def gaudi_persimmon_model_forward( if use_cache: next_cache = next_decoder_cache.to_legacy_cache() if use_legacy_cache else next_decoder_cache - if not return_dict: - return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None) return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=next_cache, @@ -351,7 +346,7 @@ def gaudi_persimmon_model_forward( class GaudiPersimmonForCausalLM(PersimmonForCausalLM): def forward( self, - input_ids: torch.LongTensor = None, + input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, @@ -360,12 +355,11 @@ def forward( use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, logits_to_keep: Union[int, torch.Tensor] = 0, token_idx: Optional[torch.Tensor] = None, **kwargs, - ) -> Union[Tuple, CausalLMOutputWithPast]: + ) -> CausalLMOutputWithPast: """ Inherits from PersimmonForCausalLM: https://github.com/huggingface/transformers/blob/v4.38.2/src/transformers/models/persimmon/modeling_persimmon.py The only differences are: @@ -376,10 +370,9 @@ def forward( 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 # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) - outputs = self.model( + outputs: BaseModelOutputWithPast = self.model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, @@ -388,12 +381,11 @@ def forward( use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, - return_dict=return_dict, cache_position=cache_position, token_idx=token_idx, ) - hidden_states = outputs[0] + hidden_states = outputs.last_hidden_state # No upscaling to float was ever done for Persimmon slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep logits = self.lm_head(hidden_states[:, slice_indices, :]) @@ -407,10 +399,6 @@ def forward( **kwargs, ) - if not return_dict: - output = (logits,) + outputs[1:] - return (loss,) + output if loss is not None else output - return CausalLMOutputWithPast( loss=loss, logits=logits, diff --git a/optimum/habana/transformers/models/phi/modeling_phi.py b/optimum/habana/transformers/models/phi/modeling_phi.py index b72258aef7..b50c3547ac 100644 --- a/optimum/habana/transformers/models/phi/modeling_phi.py +++ b/optimum/habana/transformers/models/phi/modeling_phi.py @@ -19,6 +19,7 @@ # limitations under the License. """PyTorch Phi model.""" +from functools import partial from typing import List, Optional, Tuple, Union import torch @@ -316,7 +317,7 @@ def allocate_kv_cache(self, batch_size, max_seq_len, inp_seq_len): def forward( self, - input_ids: torch.LongTensor = None, + input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, @@ -324,12 +325,12 @@ def forward( use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, reuse_cache: Optional[bool] = False, cache_idx: Optional[int] = None, - ) -> Union[Tuple, BaseModelOutputWithPast]: + **kwargs, + ) -> BaseModelOutputWithPast: """ Copied from PhiModel.forward: https://github.com/huggingface/transformers/blob/v4.37.1/src/transformers/models/phi/modeling_phi.py The only differences are: @@ -343,8 +344,6 @@ def forward( ) use_cache = use_cache if use_cache is not None else self.config.use_cache - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - # retrieve input_ids and inputs_embeds if input_ids is not None and inputs_embeds is not None: raise ValueError("You must specify exactly one of input_ids or inputs_embeds") @@ -411,7 +410,7 @@ def forward( if self.gradient_checkpointing and self.training: layer_outputs = self._gradient_checkpointing_func( - decoder_layer.__call__, + partial(decoder_layer.__call__, **kwargs), hidden_states, attention_mask, position_ids, @@ -455,8 +454,6 @@ def forward( next_decoder_cache.to_legacy_cache() if isinstance(next_decoder_cache, Cache) else next_decoder_cache ) - if not return_dict: - return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None) return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=next_cache, @@ -471,7 +468,7 @@ def allocate_kv_cache(self, batch_size, max_seq_len, inp_seq_len): def forward( self, - input_ids: torch.LongTensor = None, + input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, @@ -480,7 +477,6 @@ def forward( use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, logits_to_keep: Union[int, torch.Tensor] = 0, token_idx: Optional[torch.Tensor] = None, @@ -488,7 +484,7 @@ def forward( trim_logits: Optional[bool] = False, cache_idx: Optional[int] = None, **kwargs: Unpack[KwargsForCausalLM], - ) -> Union[Tuple, CausalLMOutputWithPast]: + ) -> CausalLMOutputWithPast: """ Inherits from PhiForCausalLM: https://github.com/huggingface/transformers/blob/v4.37.1/src/transformers/models/phi/modeling_phi.py The only differences are: @@ -500,10 +496,9 @@ def forward( 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 # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) - outputs = self.model( + outputs: BaseModelOutputWithPast = self.model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, @@ -512,14 +507,13 @@ def forward( use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, - return_dict=return_dict, cache_position=cache_position, token_idx=token_idx, reuse_cache=reuse_cache, cache_idx=cache_idx, ) - hidden_states = outputs[0] + hidden_states = outputs.last_hidden_state _, seq_len, _ = hidden_states.shape if seq_len > 1 and trim_logits and not self.training: if token_idx is not None: @@ -534,10 +528,6 @@ def forward( if labels is not None: loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs) - if not return_dict: - output = (logits,) + outputs[1:] - return (loss,) + output if loss is not None else output - return CausalLMOutputWithPast( loss=loss, logits=logits, diff --git a/optimum/habana/transformers/models/qwen2/modeling_qwen2.py b/optimum/habana/transformers/models/qwen2/modeling_qwen2.py index c5b2dcecb5..46f29202e2 100644 --- a/optimum/habana/transformers/models/qwen2/modeling_qwen2.py +++ b/optimum/habana/transformers/models/qwen2/modeling_qwen2.py @@ -16,6 +16,7 @@ # Copyright (C) 2022-2024 Habana Labs, Ltd. an Intel Company ############################################################################### +from functools import partial from typing import List, Optional, Tuple, Union import torch @@ -755,7 +756,7 @@ def update_sincos_cache(self, seq_len): def forward( self, - input_ids: torch.LongTensor = None, + input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, @@ -763,7 +764,6 @@ def forward( use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, attn_softmax_bf16: Optional[bool] = False, @@ -776,15 +776,14 @@ def forward( cache_idx: int = None, lazy_mode: Optional[bool] = True, num_virtual_tokens: int = None, - ) -> Union[Tuple, BaseModelOutputWithPast]: + **kwargs, + ) -> BaseModelOutputWithPast: 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 ) use_cache = use_cache if use_cache is not None else self.config.use_cache - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - # retrieve input_ids and inputs_embeds if input_ids is not None and inputs_embeds is not None: raise ValueError("You must specify exactly one of input_ids or inputs_embeds") @@ -874,7 +873,7 @@ def forward( if self.gradient_checkpointing and self.training: layer_outputs = self._gradient_checkpointing_func( - decoder_layer.__call__, + partial(decoder_layer.__call__, **kwargs), hidden_states, causal_mask, position_ids, @@ -933,8 +932,6 @@ def forward( next_cache = ( next_decoder_cache.to_legacy_cache() if isinstance(next_decoder_cache, Cache) else next_decoder_cache ) - if not return_dict: - return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None) return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=next_cache, @@ -955,7 +952,7 @@ def update_sincos_cache(self, seq_len): def forward( self, - input_ids: torch.LongTensor = None, + input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, @@ -964,7 +961,6 @@ def forward( use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, logits_to_keep: Union[int, torch.Tensor] = 0, token_idx: Optional[torch.Tensor] = None, @@ -980,18 +976,17 @@ def forward( lazy_mode: Optional[bool] = True, num_virtual_tokens: int = None, **kwargs: Unpack[KwargsForCausalLM], - ) -> Union[Tuple, CausalLMOutputWithPast]: + ) -> CausalLMOutputWithPast: 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 if self.generation_config.use_fused_rope is False: global has_fused_rope has_fused_rope = False # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) - outputs = self.model( + outputs: BaseModelOutputWithPast = self.model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, @@ -1000,7 +995,6 @@ def forward( use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, - return_dict=return_dict, cache_position=cache_position, token_idx=token_idx, attn_softmax_bf16=attn_softmax_bf16, @@ -1015,7 +1009,7 @@ def forward( num_virtual_tokens=num_virtual_tokens, ) - hidden_states = outputs[0] + hidden_states = outputs.last_hidden_state _, seq_len, _ = hidden_states.shape if seq_len > 1 and trim_logits and not self.training: if token_idx is not None: @@ -1031,10 +1025,6 @@ def forward( if labels is not None: loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs) - if not return_dict: - output = (logits,) + outputs[1:] - return (loss,) + output if loss is not None else output - return CausalLMOutputWithPast( loss=loss, logits=logits, diff --git a/optimum/habana/transformers/models/qwen2_moe/modeling_qwen2_moe.py b/optimum/habana/transformers/models/qwen2_moe/modeling_qwen2_moe.py index 3b7077aca9..b7c2c556a3 100755 --- a/optimum/habana/transformers/models/qwen2_moe/modeling_qwen2_moe.py +++ b/optimum/habana/transformers/models/qwen2_moe/modeling_qwen2_moe.py @@ -806,7 +806,7 @@ def update_sincos_cache(self, seq_len): def forward( self, - input_ids: torch.LongTensor = None, + input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, @@ -815,7 +815,6 @@ def forward( output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, output_router_logits: Optional[bool] = None, - return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, attn_softmax_bf16: Optional[bool] = False, @@ -828,7 +827,7 @@ def forward( cache_idx: int = None, lazy_mode: Optional[bool] = True, num_virtual_tokens: int = None, - ) -> Union[Tuple, MoeModelOutputWithPast]: + ) -> MoeModelOutputWithPast: """ Copied from LlamaModel.forward: https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py The only differences are: @@ -849,7 +848,6 @@ def forward( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) use_cache = use_cache if use_cache is not None else self.config.use_cache - return_dict = return_dict if return_dict is not None else self.config.use_return_dict if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") @@ -998,12 +996,6 @@ def forward( next_cache = ( next_decoder_cache.to_legacy_cache() if isinstance(next_decoder_cache, Cache) else next_decoder_cache ) - if not return_dict: - return tuple( - v - for v in [hidden_states, next_cache, all_hidden_states, all_self_attns, all_router_logits] - if v is not None - ) return MoeModelOutputWithPast( last_hidden_state=hidden_states, @@ -1036,7 +1028,7 @@ def update_sincos_cache(self, seq_len): def forward( self, - input_ids: torch.LongTensor = None, + input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, @@ -1046,7 +1038,6 @@ def forward( output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, output_router_logits: Optional[bool] = None, - return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, logits_to_keep: Union[int, torch.Tensor] = 0, token_idx: Optional[torch.Tensor] = None, @@ -1062,7 +1053,7 @@ def forward( lazy_mode: Optional[bool] = True, num_virtual_tokens: int = None, **loss_kwargs, - ) -> Union[Tuple, MoeCausalLMOutputWithPast]: + ) -> MoeCausalLMOutputWithPast: output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_router_logits = ( output_router_logits if output_router_logits is not None else self.config.output_router_logits @@ -1071,13 +1062,12 @@ def forward( 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 if self.generation_config.use_fused_rope is False: global has_fused_rope has_fused_rope = False # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) - outputs = self.model( + outputs: MoeModelOutputWithPast = self.model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, @@ -1087,7 +1077,6 @@ def forward( output_attentions=output_attentions, output_hidden_states=output_hidden_states, output_router_logits=output_router_logits, - return_dict=return_dict, cache_position=cache_position, token_idx=token_idx, attn_softmax_bf16=attn_softmax_bf16, @@ -1102,7 +1091,7 @@ def forward( num_virtual_tokens=num_virtual_tokens, ) - hidden_states = outputs[0] + hidden_states = outputs.last_hidden_state _, seq_len, _ = hidden_states.shape if seq_len > 1 and trim_logits and not self.training: if token_idx is not None: @@ -1121,7 +1110,7 @@ def forward( aux_loss = None if output_router_logits: aux_loss = load_balancing_loss_func( - outputs.router_logits if return_dict else outputs[-1], + outputs.router_logits, self.num_experts, self.num_experts_per_tok, attention_mask, @@ -1129,12 +1118,6 @@ def forward( if labels is not None: loss += self.router_aux_loss_coef * aux_loss.to(loss.device) # make sure to reside in the same device - if not return_dict: - output = (logits,) + outputs[1:] - if output_router_logits: - output = (aux_loss,) + output - return (loss,) + output if loss is not None else output - return MoeCausalLMOutputWithPast( loss=loss, aux_loss=aux_loss, diff --git a/optimum/habana/transformers/models/qwen2_vl/modeling_qwen2_vl.py b/optimum/habana/transformers/models/qwen2_vl/modeling_qwen2_vl.py index 79a9744c4a..a729b8f28b 100644 --- a/optimum/habana/transformers/models/qwen2_vl/modeling_qwen2_vl.py +++ b/optimum/habana/transformers/models/qwen2_vl/modeling_qwen2_vl.py @@ -88,8 +88,8 @@ def forward( "removed and `position_embeddings` will be mandatory." ) emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1) - cos = emb.cos().float() - sin = emb.sin().float() + cos = emb.cos() + sin = emb.sin() else: cos, sin = position_embeddings q, k = apply_rotary_pos_emb_vision(q, k, cos, sin) @@ -102,11 +102,11 @@ def forward( v = v.transpose(0, 1) if FusedSDPA is not None and use_flash_attention: - attn_output = self.fused_scaled_dot_product_attention(q, k, v, attention_mask, 0.0, False, None, "None") + attn_output = self.fused_scaled_dot_product_attention(q.unsqueeze(0), k.unsqueeze(0), v.unsqueeze(0), attention_mask, 0.0, False, None, "None") else: - attn_output = F.scaled_dot_product_attention(q, k, v, attention_mask, dropout_p=0.0) + attn_output = F.scaled_dot_product_attention(q.unsqueeze(0), k.unsqueeze(0), v.unsqueeze(0), attention_mask, dropout_p=0.0) - attn_output = attn_output.transpose(0, 1) + attn_output = attn_output.squeeze(0).transpose(0, 1) attn_output = attn_output.reshape(seq_length, -1) attn_output = self.proj(attn_output) del attention_mask @@ -396,7 +396,7 @@ def forward( class GaudiQwen2VLModel(Qwen2VLModel): def forward( self, - input_ids: torch.LongTensor = None, + input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, @@ -522,7 +522,7 @@ class GaudiQwen2VLForConditionalGeneration(Qwen2VLForConditionalGeneration): def forward( self, - input_ids: torch.LongTensor = None, + input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, @@ -549,7 +549,6 @@ def forward( - add Gaudi Example """ r""" - Args: labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored diff --git a/optimum/habana/transformers/models/seamless_m4t/modeling_seamless_m4t.py b/optimum/habana/transformers/models/seamless_m4t/modeling_seamless_m4t.py index 061aebb3c6..7c962c6def 100644 --- a/optimum/habana/transformers/models/seamless_m4t/modeling_seamless_m4t.py +++ b/optimum/habana/transformers/models/seamless_m4t/modeling_seamless_m4t.py @@ -226,7 +226,7 @@ def gaudi_SeamlessM4TDecoderLayer_forward( def gaudi_SeamlessM4TDecoder_forward( self, - input_ids: torch.LongTensor = None, + input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, encoder_hidden_states: Optional[torch.FloatTensor] = None, encoder_attention_mask: Optional[torch.LongTensor] = None, @@ -446,7 +446,7 @@ def gaudi_SeamlessM4TTextToUnitModel_forward( def gaudi_SeamlessM4TTextToUnitForConditionalGeneration_forward( self, - input_ids: torch.LongTensor = None, + input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, decoder_input_ids: Optional[torch.LongTensor] = None, decoder_attention_mask: Optional[torch.LongTensor] = None, @@ -599,7 +599,7 @@ def _transpose_conv_out_length(input_length, kernel_size, stride, pad, dilation= def gaudi_SeamlessM4TForTextToSpeech_forward( self, - input_ids: torch.LongTensor = None, + input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, decoder_input_ids: Optional[torch.LongTensor] = None, decoder_attention_mask: Optional[torch.LongTensor] = None, @@ -754,7 +754,7 @@ def gaudi_SeamlessM4TForTextToSpeech_generate( # overwrite text_decoder_input_ids if tgt_lang is passed. The latter gets priority over decoder_input_ids. text_tgt_lang_id = self.generation_config.text_decoder_lang_to_code_id.get(tgt_lang) - text_decoder_input_ids = torch.tensor([[text_tgt_lang_id]] * batch_size).to(self.device) + text_decoder_input_ids = torch.tensor([[text_tgt_lang_id]] * batch_size, device=self.device) kwargs_text["decoder_input_ids"] = text_decoder_input_ids @@ -775,7 +775,7 @@ def gaudi_SeamlessM4TForTextToSpeech_generate( idx_most_probable_sequences_per_batch = text_generation_output.sequences_scores.view(batch_size, -1) idx_most_probable_sequences_per_batch = idx_most_probable_sequences_per_batch.argmax(-1) idx_most_probable_sequences_per_batch = ( - idx_most_probable_sequences_per_batch + torch.arange(batch_size).to(self.device) * num_return_sequences + idx_most_probable_sequences_per_batch + torch.arange(batch_size, device=self.device) * num_return_sequences ) sequences = sequences[idx_most_probable_sequences_per_batch] @@ -796,8 +796,8 @@ def gaudi_SeamlessM4TForTextToSpeech_generate( # Compute t2u decoder_input_ids t2u_decoder_input_ids = kwargs_speech.get("decoder_input_ids") t2u_tgt_lang_id = self.generation_config.t2u_lang_code_to_id.get(tgt_lang) - t2u_decoder_input_ids = torch.tensor([[self.config.t2u_eos_token_id, t2u_tgt_lang_id]] * batch_size).to( - self.device + t2u_decoder_input_ids = torch.tensor( + [[self.config.t2u_eos_token_id, t2u_tgt_lang_id]] * batch_size, device=self.device ) kwargs_speech["decoder_input_ids"] = t2u_decoder_input_ids @@ -815,9 +815,9 @@ def gaudi_SeamlessM4TForTextToSpeech_generate( unit_ids = torch.where(unit_ids == self.config.t2u_pad_token_id, unit_ids, unit_ids - self.config.vocoder_offset) vocoder_tgt_lang_id = self.generation_config.vocoder_lang_code_to_id.get(tgt_lang) - vocoder_tgt_lang_id = torch.tensor([[vocoder_tgt_lang_id]] * len(unit_ids)).to(self.device) + vocoder_tgt_lang_id = torch.tensor([[vocoder_tgt_lang_id]] * len(unit_ids), device=self.device) - spkr_id = torch.tensor([[spkr_id]] * len(unit_ids)).to(self.device) + spkr_id = torch.tensor([[spkr_id]] * len(unit_ids), device=self.device) waveform, waveform_lengths = self.vocoder(input_ids=unit_ids, spkr_id=spkr_id, lang_id=vocoder_tgt_lang_id) diff --git a/optimum/habana/transformers/models/stablelm/modeling_stablelm.py b/optimum/habana/transformers/models/stablelm/modeling_stablelm.py index b86fdad63a..92fdbd9bbc 100644 --- a/optimum/habana/transformers/models/stablelm/modeling_stablelm.py +++ b/optimum/habana/transformers/models/stablelm/modeling_stablelm.py @@ -230,7 +230,7 @@ def forward( def gaudi_stablelm_model_forward( self, - input_ids: torch.LongTensor = None, + input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, @@ -238,10 +238,9 @@ def gaudi_stablelm_model_forward( use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, -) -> Union[Tuple, BaseModelOutputWithPast]: +) -> BaseModelOutputWithPast: """ Copied from StableLmModel.forward: https://github.com/huggingface/transformers/blob/v4.38.2/src/transformers/models/stablelm/modeling_stablelm.py The only differences are: @@ -254,8 +253,6 @@ def gaudi_stablelm_model_forward( ) use_cache = use_cache if use_cache is not None else self.config.use_cache - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - # retrieve input_ids and inputs_embeds if input_ids is not None and inputs_embeds is not None: raise ValueError("You must specify exactly one of input_ids or inputs_embeds") @@ -357,8 +354,6 @@ def gaudi_stablelm_model_forward( if use_cache: next_cache = next_decoder_cache.to_legacy_cache() if use_legacy_cache else next_decoder_cache - if not return_dict: - return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None) return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=next_cache, @@ -370,7 +365,7 @@ def gaudi_stablelm_model_forward( class GaudiStableLmForCausalLM(StableLmForCausalLM): def forward( self, - input_ids: torch.LongTensor = None, + input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, @@ -379,12 +374,11 @@ def forward( use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, logits_to_keep: Union[int, torch.Tensor] = 0, token_idx: Optional[torch.Tensor] = None, **kwargs, - ) -> Union[Tuple, CausalLMOutputWithPast]: + ) -> CausalLMOutputWithPast: """ Inherits from StableLmForCausalLM: https://github.com/huggingface/transformers/blob/v4.38.2/src/transformers/models/stablelm/modeling_stablelm.py The only differences are: @@ -394,9 +388,8 @@ def forward( 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 - outputs = self.model( + outputs: BaseModelOutputWithPast = self.model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, @@ -405,12 +398,11 @@ def forward( use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, - return_dict=return_dict, cache_position=cache_position, token_idx=token_idx, ) - hidden_states = outputs[0] + hidden_states = outputs.last_hidden_state # No upscaling to float was ever done for StableLm slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep logits = self.lm_head(hidden_states[:, slice_indices, :]) @@ -424,10 +416,6 @@ def forward( **kwargs, ) - if not return_dict: - output = (logits,) + outputs[1:] - return (loss,) + output if loss is not None else output - return CausalLMOutputWithPast( loss=loss, logits=logits, diff --git a/optimum/habana/transformers/models/starcoder2/modeling_starcoder2.py b/optimum/habana/transformers/models/starcoder2/modeling_starcoder2.py index 24fb2a4e17..8fa5b42fe4 100644 --- a/optimum/habana/transformers/models/starcoder2/modeling_starcoder2.py +++ b/optimum/habana/transformers/models/starcoder2/modeling_starcoder2.py @@ -510,7 +510,7 @@ def update_sincos_cache(self, seq_len): def forward( self, - input_ids: torch.LongTensor = None, + input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, @@ -518,7 +518,6 @@ def forward( use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, attn_softmax_bf16: Optional[bool] = False, @@ -528,15 +527,13 @@ def forward( flash_attention_causal_mask: Optional[bool] = False, cache_idx: int = None, lazy_mode: Optional[bool] = True, - ) -> Union[Tuple, BaseModelOutputWithPast]: + ) -> BaseModelOutputWithPast: 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 ) use_cache = use_cache if use_cache is not None else self.config.use_cache - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - # retrieve input_ids and inputs_embeds if input_ids is not None and inputs_embeds is not None: raise ValueError("You must specify exactly one of input_ids or inputs_embeds") @@ -657,8 +654,6 @@ def forward( next_cache = ( next_decoder_cache.to_legacy_cache() if isinstance(next_decoder_cache, Cache) else next_decoder_cache ) - if not return_dict: - return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None) return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=next_cache, @@ -688,7 +683,7 @@ def update_sincos_cache(self, seq_len): def forward( self, - input_ids: torch.LongTensor = None, + input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, @@ -697,7 +692,6 @@ def forward( use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, logits_to_keep: Union[int, torch.Tensor] = 0, token_idx: Optional[torch.Tensor] = None, @@ -710,12 +704,11 @@ def forward( cache_idx: int = None, lazy_mode: Optional[bool] = True, **kwargs: Unpack[KwargsForCausalLM], - ) -> Union[Tuple, CausalLMOutputWithPast]: + ) -> CausalLMOutputWithPast: 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 if not hasattr(self.config, "_attn_implementation"): setattr(self.config, "_attn_implementation", "eager") @@ -723,7 +716,7 @@ def forward( self.config._attn_implementation = "eager" # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) - outputs = self.model( + outputs: BaseModelOutputWithPast = self.model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, @@ -732,7 +725,6 @@ def forward( use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, - return_dict=return_dict, cache_position=cache_position, token_idx=token_idx, attn_softmax_bf16=attn_softmax_bf16, @@ -744,7 +736,7 @@ def forward( lazy_mode=lazy_mode, ) - hidden_states = outputs[0] + hidden_states = outputs.last_hidden_state _, seq_len, _ = hidden_states.shape if seq_len > 1 and trim_logits and not self.training: if token_idx is not None: @@ -760,10 +752,6 @@ def forward( if labels is not None: loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs) - if not return_dict: - output = (logits,) + outputs[1:] - return (loss,) + output if loss is not None else output - return CausalLMOutputWithPast( loss=loss, logits=logits, diff --git a/optimum/habana/transformers/models/video_llava/modeling_video_llava.py b/optimum/habana/transformers/models/video_llava/modeling_video_llava.py index 209045e0b5..6670b23375 100644 --- a/optimum/habana/transformers/models/video_llava/modeling_video_llava.py +++ b/optimum/habana/transformers/models/video_llava/modeling_video_llava.py @@ -162,9 +162,9 @@ def _get_vision_features( def forward( self, - input_ids: torch.LongTensor = None, - pixel_values_images: torch.FloatTensor = None, - pixel_values_videos: torch.FloatTensor = None, + input_ids: Optional[torch.LongTensor] = None, + pixel_values_images: Optional[torch.FloatTensor] = None, + pixel_values_videos: Optional[torch.FloatTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, diff --git a/optimum/habana/transformers/models/vit/modeling_vit.py b/optimum/habana/transformers/models/vit/modeling_vit.py index 4fd5990e14..3988f6ddf8 100644 --- a/optimum/habana/transformers/models/vit/modeling_vit.py +++ b/optimum/habana/transformers/models/vit/modeling_vit.py @@ -20,6 +20,37 @@ import torch +def gaudi_eager_attention_forward( + module: torch.nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: Optional[torch.Tensor], + scaling: float, + dropout: float = 0.0, + **kwargs, +): + # Take the dot product between "query" and "key" to get the raw attention scores. + # The div has been put inside the matmul because it achieves better performance on HPU. + attn_weights = torch.matmul(query, key.transpose(-1, -2) / math.sqrt(module.attention_head_size)) + + # Normalize the attention scores to probabilities. + attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + + # This is actually dropping out entire tokens to attend to, which might + # seem a bit unusual, but is taken from the original Transformer paper. + attn_weights = torch.nn.functional.dropout(attn_weights, p=dropout, training=module.training) + + # Mask heads if we want to + if attention_mask is not None: + attn_weights = attn_weights * attention_mask + + attn_output = torch.matmul(attn_weights, value) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + def gaudi_vit_self_attention_forward( self, hidden_states, head_mask: Optional[torch.Tensor] = None, output_attentions: bool = False ) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor]]: @@ -28,33 +59,23 @@ def gaudi_vit_self_attention_forward( the division is performed before the matmul for computing attention scores. This gives better performance on HPU. """ - - mixed_query_layer = self.query(hidden_states) - key_layer = self.transpose_for_scores(self.key(hidden_states)) value_layer = self.transpose_for_scores(self.value(hidden_states)) - query_layer = self.transpose_for_scores(mixed_query_layer) - - # Take the dot product between "query" and "key" to get the raw attention scores. - # The div has been put inside the matmul because it achieves better performance on HPU. - attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2) / math.sqrt(self.attention_head_size)) - - # Normalize the attention scores to probabilities. - attention_probs = torch.nn.functional.softmax(attention_scores, dim=-1) - - # This is actually dropping out entire tokens to attend to, which might - # seem a bit unusual, but is taken from the original Transformer paper. - attention_probs = self.dropout(attention_probs) - - # Mask heads if we want to - if head_mask is not None: - attention_probs = attention_probs * head_mask + query_layer = self.transpose_for_scores(self.query(hidden_states)) - context_layer = torch.matmul(attention_probs, value_layer) + context_layer, attention_probs = gaudi_eager_attention_forward( + self, + query_layer, + key_layer, + value_layer, + head_mask, + is_causal=self.is_causal, + scaling=self.scaling, + dropout=0.0 if not self.training else self.dropout_prob, + ) - context_layer = context_layer.permute(0, 2, 1, 3).contiguous() new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,) - context_layer = context_layer.view(new_context_layer_shape) + context_layer = context_layer.reshape(new_context_layer_shape) outputs = (context_layer, attention_probs) if output_attentions else (context_layer,) diff --git a/optimum/habana/transformers/models/wav2vec2/modeling_wav2vec2.py b/optimum/habana/transformers/models/wav2vec2/modeling_wav2vec2.py index e03d9056e7..f58cba1baf 100644 --- a/optimum/habana/transformers/models/wav2vec2/modeling_wav2vec2.py +++ b/optimum/habana/transformers/models/wav2vec2/modeling_wav2vec2.py @@ -89,7 +89,7 @@ def compute_num_masked_span(input_length): # compute number of masked spans in batch input_lengths = ( - attention_mask.sum(-1).detach().tolist() + attention_mask.detach().sum(-1).tolist() if attention_mask is not None else [sequence_length for _ in range(batch_size)] ) diff --git a/optimum/habana/transformers/trainer.py b/optimum/habana/transformers/trainer.py index 06d35f2e2b..3d7aa3ffda 100644 --- a/optimum/habana/transformers/trainer.py +++ b/optimum/habana/transformers/trainer.py @@ -29,7 +29,7 @@ import warnings from collections.abc import Mapping from pathlib import Path -from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Type, Union +from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Tuple, Union import huggingface_hub.utils as hf_hub_utils import numpy as np @@ -38,6 +38,7 @@ from accelerate.data_loader import SeedableRandomSampler from accelerate.utils import ( DistributedDataParallelKwargs, + TorchTensorParallelPlugin, load_fsdp_model, load_fsdp_optimizer, save_fsdp_model, @@ -217,21 +218,21 @@ class GaudiTrainer(Trainer): @deprecate_kwarg("tokenizer", new_name="processing_class", version="5.0.0", raise_if_both_names=True) def __init__( self, - model: Union[PreTrainedModel, torch.nn.Module] = None, + model: Union[PreTrainedModel, torch.nn.Module, None] = None, gaudi_config: GaudiConfig = None, args: TrainingArguments = None, data_collator: Optional[DataCollator] = None, train_dataset: Optional[Union[Dataset, IterableDataset, "datasets.Dataset"]] = None, - eval_dataset: Optional[Union[Dataset, Dict[str, Dataset], "datasets.Dataset"]] = None, + eval_dataset: Optional[Union[Dataset, dict[str, Dataset], "datasets.Dataset"]] = None, processing_class: Optional[ Union[PreTrainedTokenizerBase, BaseImageProcessor, FeatureExtractionMixin, ProcessorMixin] ] = None, model_init: Optional[Callable[[], PreTrainedModel]] = None, compute_loss_func: Optional[Callable] = None, - compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None, - callbacks: Optional[List[TrainerCallback]] = None, - optimizers: Tuple[Optional[torch.optim.Optimizer], Optional[torch.optim.lr_scheduler.LambdaLR]] = (None, None), - optimizer_cls_and_kwargs: Optional[Tuple[Type[torch.optim.Optimizer], Dict[str, Any]]] = None, + compute_metrics: Optional[Callable[[EvalPrediction], dict]] = None, + callbacks: Optional[list[TrainerCallback]] = None, + optimizers: tuple[Optional[torch.optim.Optimizer], Optional[torch.optim.lr_scheduler.LambdaLR]] = (None, None), + optimizer_cls_and_kwargs: Optional[tuple[type[torch.optim.Optimizer], dict[str, Any]]] = None, preprocess_logits_for_metrics: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None, ): if args is None: @@ -502,8 +503,8 @@ def _wrap_model(self, model, training=True, dataloader=None): def train( self, resume_from_checkpoint: Optional[Union[str, bool]] = None, - trial: Union["optuna.Trial", Dict[str, Any]] = None, - ignore_keys_for_eval: Optional[List[str]] = None, + trial: Union["optuna.Trial", dict[str, Any], None] = None, + ignore_keys_for_eval: Optional[list[str]] = None, **kwargs, ): """ @@ -538,7 +539,7 @@ def train( # do_train is not a reliable argument, as it might not be set and .train() still called, so # the following is a workaround: - if args.bf16_full_eval and not args.do_train and not self.is_model_parallel: + if args.bf16_full_eval and not args.do_train and not self.is_model_parallel and self.model_init is None: self._move_model_to_device(self.model, args.device) if "model_path" in kwargs: @@ -677,6 +678,11 @@ def _inner_training_loop( delay_optimizer_creation = self.is_fsdp_enabled + # Can't delay optimizer creation when using FSDP2: https://github.com/huggingface/accelerate/blob/3f636d626063ffcf9a337c7d3624d61b7d187d59/src/accelerate/accelerator.py#L1404 + is_fsdp2 = self.is_fsdp_enabled and (getattr(self.accelerator.state.fsdp_plugin, "fsdp_version", 1) == 2) + if is_fsdp2: + delay_optimizer_creation = False + # We need to reset the scheduler, as its parameters may be different on subsequent calls if self._created_lr_scheduler: self.lr_scheduler = None @@ -871,10 +877,14 @@ def hpu_deepspeed_checkpointing(function, *checkpoint_args, use_reentrant: Optio torch.distributed.broadcast(param.data, src=0) # Update the references - self.state.init_training_references(self, train_dataloader, max_steps, num_train_epochs, trial) + for attr in ("model", "optimizer", "lr_scheduler"): + setattr(self.callback_handler, attr, getattr(self, attr)) + self.callback_handler.train_dataloader = train_dataloader + + self.state.init_training_references(self, max_steps, num_train_epochs, trial) # tr_loss is a tensor to avoid synchronization of TPUs through .item() - tr_loss = torch.tensor(0.0).to(args.device) + tr_loss = torch.tensor(0.0, device=args.device) # _total_loss_scalar is updated every time .item() has to be called on tr_loss and stores the sum of all losses self._total_loss_scalar = 0.0 self._globalstep_last_logged = self.state.global_step @@ -886,6 +896,8 @@ def hpu_deepspeed_checkpointing(function, *checkpoint_args, use_reentrant: Optio args.max_grad_norm is not None and args.max_grad_norm > 0 ) + learning_rate = None + # attn_softmax_bf16 and use_flash_attention are enabled only for llama, qwen2, starcoder2, gemma and baichuan # lazy_mode for llama, qwen2, starcoder2 and mistral _should_update_inputs, _inputs_update = _get_input_update_settings(self.model, lazy_mode=args.use_lazy_mode) @@ -981,9 +993,7 @@ def hpu_deepspeed_checkpointing(function, *checkpoint_args, use_reentrant: Optio else: input_tokens = inputs[main_input_name].numel() input_tokens = torch.tensor(input_tokens, device=self.args.device, dtype=torch.int64) - self.state.num_input_tokens_seen += ( - self.accelerator.gather(input_tokens).sum().cpu().item() - ) + self.state.num_input_tokens_seen += self.accelerator.gather(input_tokens).sum().item() if rng_to_sync: self._load_rng_state(resume_from_checkpoint) rng_to_sync = False @@ -1066,6 +1076,9 @@ def hpu_deepspeed_checkpointing(function, *checkpoint_args, use_reentrant: Optio self.control = self.callback_handler.on_optimizer_step(args, self.state, self.control) + # get leaning rate before update + learning_rate = self._get_learning_rate() + if not self.accelerator.optimizer_step_was_skipped: # Delay optimizer scheduling until metrics are generated if not isinstance(self.lr_scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau): @@ -1078,7 +1091,14 @@ def hpu_deepspeed_checkpointing(function, *checkpoint_args, use_reentrant: Optio self.htcore.mark_step() self.control = self.callback_handler.on_step_end(args, self.state, self.control) self._maybe_log_save_evaluate( - tr_loss, grad_norm, model, trial, epoch, ignore_keys_for_eval, start_time + tr_loss, + grad_norm, + model, + trial, + epoch, + ignore_keys_for_eval, + start_time, + learning_rate=learning_rate, ) else: self.control = self.callback_handler.on_substep_end(args, self.state, self.control) @@ -1098,7 +1118,9 @@ def hpu_deepspeed_checkpointing(function, *checkpoint_args, use_reentrant: Optio self.control.should_training_stop = True self.control = self.callback_handler.on_epoch_end(args, self.state, self.control) - self._maybe_log_save_evaluate(tr_loss, grad_norm, model, trial, epoch, ignore_keys_for_eval, start_time) + self._maybe_log_save_evaluate( + tr_loss, grad_norm, model, trial, epoch, ignore_keys_for_eval, start_time, learning_rate=learning_rate + ) if self.control.should_training_stop: break @@ -1174,7 +1196,6 @@ def _load_best_model(self): best_safe_adapter_model_path = os.path.join(self.state.best_model_checkpoint, ADAPTER_SAFE_WEIGHTS_NAME) model = self.model - if self.is_deepspeed_enabled: deepspeed_load_checkpoint( self.model_wrapped, @@ -1196,7 +1217,6 @@ def _load_best_model(self): or os.path.exists(best_safe_adapter_model_path) ): has_been_loaded = True - weights_only_kwarg = {"weights_only": True} if _is_peft_model(model): # If train a model using PEFT & LoRA, assume that adapter have been saved properly. # TODO: in the future support only specific min PEFT versions @@ -1247,11 +1267,7 @@ def _load_best_model(self): if self.args.save_safetensors and os.path.isfile(best_safe_model_path): state_dict = safetensors.torch.load_file(best_safe_model_path, device="cpu") else: - state_dict = torch.load( - best_model_path, - map_location="cpu", - **weights_only_kwarg, - ) + state_dict = torch.load(best_model_path, map_location="cpu", weights_only=True) # If the model is on the GPU, it still works! # workaround for FSDP bug https://github.com/pytorch/pytorch/issues/82963 @@ -1271,12 +1287,14 @@ def _load_best_model(self): "on multiple nodes, you should activate `--save_on_each_node`." ) - def _maybe_log_save_evaluate(self, tr_loss, _grad_norm, model, trial, epoch, ignore_keys_for_eval, start_time): + def _maybe_log_save_evaluate( + self, tr_loss, _grad_norm, model, trial, epoch, ignore_keys_for_eval, start_time, learning_rate=None + ): if self.args.adjust_throughput: save_start = time.perf_counter() if self.control.should_log and self.state.global_step > self._globalstep_last_logged: - logs: Dict[str, float] = {} + logs: dict[str, float] = {} # all_gather + mean() to get average loss over all processes tr_loss_scalar = self._nested_gather(tr_loss).mean().item() @@ -1305,7 +1323,10 @@ def _maybe_log_save_evaluate(self, tr_loss, _grad_norm, model, trial, epoch, ign if grad_norm is not None: logs["grad_norm"] = grad_norm - logs["learning_rate"] = self._get_learning_rate() + if learning_rate is not None: + logs["learning_rate"] = learning_rate + else: + logs["learning_rate"] = self._get_learning_rate() self._total_loss_scalar += tr_loss_scalar self._globalstep_last_logged = self.state.global_step @@ -1352,7 +1373,7 @@ def _load_rng_state(self, checkpoint): return with safe_globals(): - checkpoint_rng_state = torch.load(rng_file) + checkpoint_rng_state = torch.load(rng_file, weights_only=True) random.setstate(checkpoint_rng_state["python"]) np.random.set_state(checkpoint_rng_state["numpy"]) torch.random.set_rng_state(checkpoint_rng_state["cpu"]) @@ -1442,7 +1463,9 @@ def _load_optimizer_and_scheduler(self, checkpoint): # deepspeed loads optimizer/lr_scheduler together with the model in deepspeed_init if not isinstance(self.lr_scheduler, DeepSpeedSchedulerWrapper): with warnings.catch_warnings(record=True) as caught_warnings: - self.lr_scheduler.load_state_dict(torch.load(os.path.join(checkpoint, SCHEDULER_NAME))) + self.lr_scheduler.load_state_dict( + torch.load(os.path.join(checkpoint, SCHEDULER_NAME), weights_only=True) + ) reissue_pt_warnings(caught_warnings) return @@ -1475,12 +1498,12 @@ def _load_optimizer_and_scheduler(self, checkpoint): ) else: self.optimizer.load_state_dict( - torch.load(os.path.join(checkpoint, OPTIMIZER_NAME), map_location=map_location) + torch.load(os.path.join(checkpoint, OPTIMIZER_NAME), map_location=map_location, weights_only=True) ) with warnings.catch_warnings(record=True) as caught_warnings: self.lr_scheduler.load_state_dict( - torch.load(os.path.join(checkpoint, SCHEDULER_NAME), map_location=map_location) + torch.load(os.path.join(checkpoint, SCHEDULER_NAME), map_location=map_location, weights_only=True) ) reissue_pt_warnings(caught_warnings) @@ -1488,7 +1511,7 @@ def _load_optimizer_and_scheduler(self, checkpoint): if self.args.use_habana: to_device_dtype(self.optimizer.state.values(), target_device=torch.device("hpu")) - def log(self, logs: Dict[str, float], start_time: Optional[float] = None) -> None: + def log(self, logs: dict[str, float], start_time: Optional[float] = None) -> None: """ Log `logs` on the various objects watching training. @@ -1567,7 +1590,7 @@ def autocast_smart_context_manager(self, cache_enabled: Optional[bool] = True): return ctx_manager def training_step( - self, model: torch.nn.Module, inputs: Dict[str, Union[torch.Tensor, Any]], num_items_in_batch=None + self, model: torch.nn.Module, inputs: dict[str, Union[torch.Tensor, Any]], num_items_in_batch=None ) -> torch.Tensor: """ Perform a training step on a batch of inputs. @@ -1719,6 +1742,13 @@ def _save(self, output_dir: Optional[str] = None, state_dict=None): if self.processing_class is not None: self.processing_class.save_pretrained(output_dir) + elif ( + self.data_collator is not None + and hasattr(self.data_collator, "tokenizer") + and self.data_collator.tokenizer is not None + ): + logger.info("Saving Trainer.data_collator.tokenizer by default as Trainer.processing_class is `None`") + self.data_collator.tokenizer.save_pretrained(output_dir) self.gaudi_config.save_pretrained(output_dir) @@ -1727,10 +1757,10 @@ def _save(self, output_dir: Optional[str] = None, state_dict=None): def evaluate( self, - eval_dataset: Optional[Union[Dataset, Dict[str, Dataset]]] = None, - ignore_keys: Optional[List[str]] = None, + eval_dataset: Optional[Union[Dataset, dict[str, Dataset]]] = None, + ignore_keys: Optional[list[str]] = None, metric_key_prefix: str = "eval", - ) -> Dict[str, float]: + ) -> dict[str, float]: """ From https://github.com/huggingface/transformers/blob/v4.38.2/src/transformers/trainer.py#L3162 with the following modification 1. use throughput_warmup_steps in evaluation throughput calculation @@ -1802,7 +1832,7 @@ def evaluate( return output.metrics def predict( - self, test_dataset: Dataset, ignore_keys: Optional[List[str]] = None, metric_key_prefix: str = "test" + self, test_dataset: Dataset, ignore_keys: Optional[list[str]] = None, metric_key_prefix: str = "test" ) -> PredictionOutput: """ From https://github.com/huggingface/transformers/blob/v4.45.2/src/transformers/trainer.py#L3904 with the following modification @@ -1851,7 +1881,7 @@ def evaluation_loop( dataloader: DataLoader, description: str, prediction_loss_only: Optional[bool] = None, - ignore_keys: Optional[List[str]] = None, + ignore_keys: Optional[list[str]] = None, metric_key_prefix: str = "eval", ) -> EvalLoopOutput: """ @@ -1976,11 +2006,11 @@ def evaluation_loop( # Update containers if losses is not None: - losses = self.gather_function((losses.repeat(batch_size))) + losses = self.gather_function(losses.repeat(batch_size)) all_losses.add(losses) if inputs_decode is not None: inputs_decode = self.accelerator.pad_across_processes(inputs_decode, dim=1, pad_index=-100) - inputs_decode = self.gather_function((inputs_decode)) + inputs_decode = self.gather_function(inputs_decode) if not self.args.batch_eval_metrics or description == "Prediction": all_inputs.add(inputs_decode) if labels is not None: @@ -1995,13 +2025,13 @@ def evaluation_loop( logits = self.accelerator.pad_across_processes(logits, dim=1, pad_index=-100) if self.preprocess_logits_for_metrics is not None: logits = self.preprocess_logits_for_metrics(logits, labels) - logits = self.gather_function((logits)) + logits = self.gather_function(logits) if not self.args.batch_eval_metrics or description == "Prediction": all_preds.add(logits) if labels is not None: if self.args.context_parallel_size != 1: labels = labels.clone() - labels = self.gather_function((labels)) + labels = self.gather_function(labels) if not self.args.batch_eval_metrics or description == "Prediction": all_labels.add(labels) @@ -2102,10 +2132,10 @@ def evaluation_loop( def prediction_step( self, model: torch.nn.Module, - inputs: Dict[str, Union[torch.Tensor, Any]], + inputs: dict[str, Union[torch.Tensor, Any]], prediction_loss_only: bool, - ignore_keys: Optional[List[str]] = None, - ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]: + ignore_keys: Optional[list[str]] = None, + ) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]: """ Perform an evaluation step on `model` using `inputs`. Subclass and override to inject custom behavior. @@ -2154,7 +2184,7 @@ def prediction_step( if has_labels or loss_without_labels: with self.compute_loss_context_manager(): loss, outputs = self.compute_loss(model, inputs, return_outputs=True) - loss = loss.mean().detach() + loss = loss.detach().mean() if isinstance(outputs, dict): logits = tuple(v for k, v in outputs.items() if k not in ignore_keys + ["loss"]) @@ -2264,7 +2294,7 @@ def prediction_loop( dataloader: DataLoader, description: str, prediction_loss_only: Optional[bool] = None, - ignore_keys: Optional[List[str]] = None, + ignore_keys: Optional[list[str]] = None, metric_key_prefix: str = "eval", ) -> EvalLoopOutput: """ @@ -2344,10 +2374,10 @@ def prediction_loop( logger.info(f" Num examples = {num_examples}") logger.info(f" Batch size = {batch_size}") - losses_host: torch.Tensor = None - preds_host: Union[torch.Tensor, List[torch.Tensor]] = None - labels_host: Union[torch.Tensor, List[torch.Tensor]] = None - inputs_host: Union[torch.Tensor, List[torch.Tensor]] = None + losses_host: Optional[torch.Tensor] = None + preds_host: Union[torch.Tensor, list[torch.Tensor], None] = None + labels_host: Union[torch.Tensor, list[torch.Tensor], None] = None + inputs_host: Union[torch.Tensor, list[torch.Tensor], None] = None metrics: Optional[dict] = None eval_set_kwargs: dict = {} @@ -2506,6 +2536,11 @@ def create_accelerator_and_postprocess(self): "use_regional_compilation": self.args.use_regional_compilation, "compiled_autograd_enable": self.args.use_compiled_autograd, } + # tp is initialized at Accelerator init phase so + # args should be prepared here + if self.args.tp_size > 1: + self.is_tp_enabled = True + args["torch_tp_plugin"] = TorchTensorParallelPlugin(tp_size=self.args.tp_size) # create accelerator object self.accelerator = GaudiAccelerator(**args) @@ -2520,7 +2555,7 @@ def create_accelerator_and_postprocess(self): # deepspeed and accelerate flags covering both trainer args and accelerate launcher self.is_deepspeed_enabled = getattr(self.accelerator.state, "deepspeed_plugin", None) is not None self.is_fsdp_enabled = getattr(self.accelerator.state, "fsdp_plugin", None) is not None - + self.is_tp_enabled = getattr(self.accelerator.state, "torch_tp_plugin", None) is not None # post accelerator creation setup if self.is_fsdp_enabled: fsdp_plugin = self.accelerator.state.fsdp_plugin @@ -2554,6 +2589,12 @@ def create_accelerator_and_postprocess(self): raise ValueError( "`auto_find_batch_size` isn't supported yet with DeepSpeed Zero-3. Please consider using Zero-2, Zero-1, or FSDP" ) + if ( + self.args.save_only_model + and self.is_fsdp_enabled + and "SHARDED_STATE_DICT" in str(self.accelerator.state.fsdp_plugin.state_dict_type) + ): + raise ValueError("save_only_model option is not compatible with FSDP state dict type 'SHARDED_STATE_DICT'") def propagate_args_to_deepspeed(self, auto_find_batch_size=False): """ diff --git a/optimum/habana/transformers/trainer_seq2seq.py b/optimum/habana/transformers/trainer_seq2seq.py index 65880ac4a9..91e677b5d3 100644 --- a/optimum/habana/transformers/trainer_seq2seq.py +++ b/optimum/habana/transformers/trainer_seq2seq.py @@ -17,7 +17,7 @@ import warnings from copy import deepcopy from pathlib import Path -from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union +from typing import TYPE_CHECKING, Any, Callable, Optional, Union import torch from torch.distributed.fsdp import FullyShardedDataParallel @@ -64,15 +64,15 @@ def __init__( args: "GaudiTrainingArguments" = None, data_collator: Optional["DataCollator"] = None, train_dataset: Optional[Union[Dataset, "IterableDataset", "datasets.Dataset"]] = None, - eval_dataset: Optional[Union[Dataset, Dict[str, Dataset]]] = None, + eval_dataset: Optional[Union[Dataset, dict[str, Dataset]]] = None, processing_class: Optional[ Union["PreTrainedTokenizerBase", "BaseImageProcessor", "FeatureExtractionMixin", "ProcessorMixin"] ] = None, model_init: Optional[Callable[[], "PreTrainedModel"]] = None, compute_loss_func: Optional[Callable] = None, - compute_metrics: Optional[Callable[["EvalPrediction"], Dict]] = None, - callbacks: Optional[List["TrainerCallback"]] = None, - optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None), + compute_metrics: Optional[Callable[["EvalPrediction"], dict]] = None, + callbacks: Optional[list["TrainerCallback"]] = None, + optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None), preprocess_logits_for_metrics: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None, ): super().__init__( @@ -149,10 +149,10 @@ def load_generation_config(gen_config_arg: Union[str, GaudiGenerationConfig]) -> def evaluate( self, eval_dataset: Optional[Dataset] = None, - ignore_keys: Optional[List[str]] = None, + ignore_keys: Optional[list[str]] = None, metric_key_prefix: str = "eval", **gen_kwargs, - ) -> Dict[str, float]: + ) -> dict[str, float]: """ Run evaluation and returns metrics. The calling script will be responsible for providing a method to compute metrics, as they are task-dependent @@ -201,7 +201,7 @@ def evaluate( def predict( self, test_dataset: Dataset, - ignore_keys: Optional[List[str]] = None, + ignore_keys: Optional[list[str]] = None, metric_key_prefix: str = "test", **gen_kwargs, ) -> "PredictionOutput": @@ -258,11 +258,11 @@ def predict( def prediction_step( self, model: torch.nn.Module, - inputs: Dict[str, Union[torch.Tensor, Any]], + inputs: dict[str, Union[torch.Tensor, Any]], prediction_loss_only: bool, - ignore_keys: Optional[List[str]] = None, + ignore_keys: Optional[list[str]] = None, **gen_kwargs, - ) -> Tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]: + ) -> tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]: """ Perform an evaluation step on `model` using `inputs`. Subclass and override to inject custom behavior. @@ -373,9 +373,9 @@ def prediction_step( with self.compute_loss_context_manager(): outputs = model(**inputs) if self.label_smoother is not None: - loss = self.label_smoother(outputs, inputs["labels"]).mean().detach() + loss = self.label_smoother(outputs, inputs["labels"]).detach().mean() else: - loss = (outputs["loss"] if isinstance(outputs, dict) else outputs[0]).mean().detach() + loss = (outputs["loss"] if isinstance(outputs, dict) else outputs[0]).detach().mean() else: loss = None except RuntimeError as error: diff --git a/optimum/habana/transformers/training_args.py b/optimum/habana/transformers/training_args.py index 07ea0c30c5..318eed6891 100644 --- a/optimum/habana/transformers/training_args.py +++ b/optimum/habana/transformers/training_args.py @@ -13,7 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import io import json import os import warnings @@ -37,7 +36,6 @@ SchedulerType, ) from transformers.training_args import ( - _VALID_DICT_FIELDS, OptimizerNames, ParallelMode, TrainingArguments, @@ -428,9 +426,9 @@ def __post_init__(self): ) # Parse in args that could be `dict` sent in from the CLI as a string - for field in _VALID_DICT_FIELDS: + for field in self._VALID_DICT_FIELDS: passed_value = getattr(self, field) - # We only want to do this if the str starts with a bracket to indiciate a `dict` + # We only want to do this if the str starts with a bracket to indicate a `dict` # else its likely a filename if supported if isinstance(passed_value, str) and passed_value.startswith("{"): loaded_dict = json.loads(passed_value) @@ -451,13 +449,6 @@ def __post_init__(self): if self.disable_tqdm is None: self.disable_tqdm = logger.getEffectiveLevel() > logging.WARN - if self.evaluation_strategy is not None: - warnings.warn( - "`evaluation_strategy` is deprecated and will be removed in version 4.46 of 🤗 Transformers. Use `eval_strategy` instead", - FutureWarning, - ) - self.eval_strategy = self.evaluation_strategy - if isinstance(self.eval_strategy, EvaluationStrategy): warnings.warn( "using `EvaluationStrategy` for `eval_strategy` is deprecated and will be removed in version 5" @@ -477,7 +468,7 @@ def __post_init__(self): self.do_eval = True if self.torch_empty_cache_steps is not None: - if not (isinstance(self.torch_empty_cache_steps, int) or self.torch_empty_cache_steps > 0): + if not (isinstance(self.torch_empty_cache_steps, int) and self.torch_empty_cache_steps > 0): raise ValueError( f"`torch_empty_cache_steps` must be an integer bigger than 0, got {self.torch_empty_cache_steps}." ) @@ -579,7 +570,7 @@ def __post_init__(self): # We need to setup the accelerator config here *before* the first call to `self.device` if is_accelerate_available(): - if not isinstance(self.accelerator_config, (AcceleratorConfig)): + if not isinstance(self.accelerator_config, AcceleratorConfig): if self.accelerator_config is None: self.accelerator_config = AcceleratorConfig() elif isinstance(self.accelerator_config, dict): @@ -594,22 +585,6 @@ def __post_init__(self): else: self.accelerator_config = AcceleratorConfig.from_json_file(self.accelerator_config) - if self.dispatch_batches is not None: - warnings.warn( - "Using `--dispatch_batches` is deprecated and will be removed in version 4.41 of 🤗 Transformers. Use" - " `--accelerator_config {'dispatch_batches':VALUE} instead", - FutureWarning, - ) - self.accelerator_config.dispatch_batches = self.dispatch_batches - - if self.split_batches is not None: - warnings.warn( - "Using `--split_batches` is deprecated and will be removed in version 4.41 of 🤗 Transformers. Use" - " `--accelerator_config {'split_batches':VALUE} instead", - FutureWarning, - ) - self.accelerator_config.split_batches = self.split_batches - if self.dataloader_drop_last: self.accelerator_config.even_batches = False @@ -633,7 +608,7 @@ def __post_init__(self): # Note: PT_HPU_LAZY_MODE=0 needs to be set before library is loaded, # setting it here would be too late - hence assertion. if self.torch_compile and self.torch_compile_backend is None: - self.torch_compile_backend = "inductor" + self.torch_compile_backend = "hpu_backend" # accelerate integration for torch compile if self.torch_compile: @@ -715,7 +690,7 @@ def __post_init__(self): if isinstance(self.fsdp_config, str): if len(self.fsdp) == 0: warnings.warn("`--fsdp_config` is useful only when `--fsdp` is specified.") - with io.open(self.fsdp_config, "r", encoding="utf-8") as f: + with open(self.fsdp_config, encoding="utf-8") as f: self.fsdp_config = json.load(f) for k in list(self.fsdp_config.keys()): if k.startswith("fsdp_"): @@ -755,6 +730,9 @@ def __post_init__(self): self.fsdp_config["xla_fsdp_v2"] = self.fsdp_config.get("xla_fsdp_v2", False) self.fsdp_config["xla_fsdp_grad_ckpt"] = self.fsdp_config.get("xla_fsdp_grad_ckpt", False) + if self.tp_size > 1: + os.environ["ACCELERATE_USE_TP"] = "true" + os.environ["TP_SIZE"] = str(self.tp_size) # accelerate integration for FSDP if len(self.fsdp) > 0 and not self.fsdp_config["xla"]: os.environ["ACCELERATE_USE_FSDP"] = "true" diff --git a/tests/test_trainer.py b/tests/test_trainer.py index a5a6df83dd..d32e2209a0 100644 --- a/tests/test_trainer.py +++ b/tests/test_trainer.py @@ -25,7 +25,7 @@ from functools import partial from itertools import product from pathlib import Path -from typing import Dict, List, Optional, Union +from typing import Any, Optional, Union import numpy as np from accelerate.state import AcceleratorState @@ -44,6 +44,7 @@ LineByLineTextDataset, PretrainedConfig, TrainerCallback, + default_data_collator, get_polynomial_decay_schedule_with_warmup, is_torch_available, ) @@ -56,7 +57,9 @@ LoggingLevel, TemporaryHubRepo, TestCasePlus, + evaluate_side_effect_factory, get_gpu_count, + get_steps_per_epoch, get_tests_dir, is_staging_test, require_accelerate, @@ -566,7 +569,7 @@ def check_best_model_has_been_loaded( else: best_model = RegressionModel() if not safe_weights: - state_dict = torch.load(os.path.join(checkpoint, WEIGHTS_NAME)) + state_dict = torch.load(os.path.join(checkpoint, WEIGHTS_NAME), weights_only=True) else: state_dict = safetensors.torch.load_file(os.path.join(checkpoint, SAFE_WEIGHTS_NAME)) best_model.load_state_dict(state_dict) @@ -656,12 +659,12 @@ def setUp(self): trainer.train() self.alternate_trained_model = (trainer.model.a, trainer.model.b) - def check_trained_model(self, model, alternate_seed=False, bf16=False): + def check_trained_model(self, model, alternate_seed=False, bf16=False, **kwargs): # Checks a training seeded with learning_rate = 0.1 (a, b) = self.alternate_trained_model if alternate_seed else self.default_trained_model if not bf16: - torch.testing.assert_close(model.a, a) - torch.testing.assert_close(model.b, b) + torch.testing.assert_close(model.a, a, **kwargs) + torch.testing.assert_close(model.b, b, **kwargs) else: self.assertTrue(torch.allclose(model.a, a, atol=1e-03, rtol=0)) self.assertTrue(torch.allclose(model.b, b, atol=1e-03, rtol=0)) @@ -754,11 +757,6 @@ def tokenize_function(examples): data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False) - model = AutoModelForCausalLM.from_pretrained(model_name) - state_dict = model.state_dict() - - base_loss_callback = StoreLossCallback() - args_kwargs = { "report_to": "none", "logging_steps": 1, @@ -774,6 +772,10 @@ def tokenize_function(examples): use_lazy_mode=True, **args_kwargs, ) + # train with base loss + set_seed(42) + model = AutoModelForCausalLM.from_pretrained(model_name) + base_loss_callback = StoreLossCallback() gaudi_config = get_gaudi_config() trainer = GaudiTrainer( model, @@ -786,8 +788,6 @@ def tokenize_function(examples): assert trainer.model_accepts_loss_kwargs trainer.train() - grad_accum_loss_callback = StoreLossCallback() - with tempfile.TemporaryDirectory() as tmp_dir: args = GaudiTrainingArguments( tmp_dir, **args_kwargs, @@ -796,8 +796,11 @@ def tokenize_function(examples): use_habana=True, use_lazy_mode=True, ) + + # train with gradient accumulation set_seed(42) model = AutoModelForCausalLM.from_pretrained(model_name) + grad_accum_loss_callback = StoreLossCallback() trainer = GaudiTrainer( model, gaudi_config, @@ -806,10 +809,12 @@ def tokenize_function(examples): callbacks=[grad_accum_loss_callback], data_collator=data_collator, ) + assert trainer.model_accepts_loss_kwargs trainer.train() + # train with broken loss set_seed(42) - model.load_state_dict(state_dict) + model = AutoModelForCausalLM.from_pretrained(model_name) broken_loss_callback = StoreLossCallback() trainer = GaudiTrainer( model, @@ -819,31 +824,30 @@ def tokenize_function(examples): callbacks=[broken_loss_callback], data_collator=data_collator, ) - # disable model_accepts_loss_kwargs + # disable model_accepts_loss_kwargs so that "num_items_in_batch" is not passed to the model trainer.model_accepts_loss_kwargs = False trainer.train() - # Calculate the difference between the base loss and the grad_accum loss - diff_truth = [ - abs(base - grad) for base, grad in zip(base_loss_callback.losses, grad_accum_loss_callback.losses) - ] - diff_broken = [ - abs(base - grad) for base, grad in zip(base_loss_callback.losses, broken_loss_callback.losses) - ] - - # all diff truth should be quite close - self.assertLess(max(diff_truth), 0.01, f"Difference {max(diff_truth)} is not within 0.01") + # Calculate the difference between the base loss and the grad_accum loss + diff_truth = [ + abs(base - grad) for base, grad in zip(base_loss_callback.losses, grad_accum_loss_callback.losses) + ] + diff_broken = [ + abs(base - grad) for base, grad in zip(base_loss_callback.losses, broken_loss_callback.losses) + ] - # max diff broken should be very off - # updated target value compared original implementation https://github.com/huggingface/transformers/blob/v4.49.0/tests/trainer/test_trainer.py#L888 - self.assertGreater(max(diff_broken), 1.0, f"Difference {max(diff_broken)} is not greater than 1.0") + # all diff truth should be quite close + self.assertLess(max(diff_truth), 0.01, f"Difference {max(diff_truth)} is not within 0.01") + # max diff broken should be very off ("very off" is arbitrary, but as long as it's bigger than 0.1, it's fine) + # updated target value compared original implementation https://github.com/huggingface/transformers/blob/v4.49.0/tests/trainer/test_trainer.py#L888 + self.assertGreater(max(diff_broken), 1.0, f"Difference {max(diff_broken)} is not greater than 1.0") - loss_base = sum(base_loss_callback.losses) - loss_broken = sum(broken_loss_callback.losses) + loss_base = sum(base_loss_callback.losses) + loss_broken = sum(broken_loss_callback.losses) - # mean/sum loss should not vary too much. - relative_diff = abs(loss_base - loss_broken) / max(loss_base, loss_broken) - self.assertLess(relative_diff, 0.2, f"Relative difference {relative_diff} is not within 0.2") + # mean/sum loss should not vary too much. + relative_diff = abs(loss_base - loss_broken) / max(loss_base, loss_broken) + self.assertLess(relative_diff, 0.2, f"Relative difference {relative_diff} is not within 0.2") def test_gradient_accumulation_loss_alignment_with_loss_func(self): set_seed(42) @@ -2001,12 +2005,30 @@ def test_safe_checkpoints(self): tmp_dir, 5, int(self.n_epochs * 64 / self.batch_size), False, safe_weights=save_safetensors ) + def test_save_collator_tokenizer_by_default(self): + class FakeCollator: + def __init__(self): + self.tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased") + self.tokenizer.add_tokens(["", ""]) + + def __call__(self, features: list[Any], return_tensors="pt") -> dict[str, Any]: + return default_data_collator(features, return_tensors) + + data_collator = FakeCollator() + tmp_dir = self.get_auto_remove_tmp_dir() + trainer = get_regression_trainer( + output_dir=tmp_dir, save_steps=5, save_safetensors=True, data_collator=data_collator + ) + trainer.train() + loaded_tokenizer = AutoTokenizer.from_pretrained(os.path.join(tmp_dir, os.listdir(tmp_dir)[0])) + assert len(loaded_tokenizer) == len(trainer.data_collator.tokenizer), "Failed to load updated tokenizer" + def test_load_best_model_with_save(self): tmp_dir = self.get_auto_remove_tmp_dir() trainer = get_regression_trainer( output_dir=tmp_dir, save_steps=5, - evaluation_strategy="steps", + eval_strategy="steps", eval_steps=5, max_steps=9, ) @@ -2027,7 +2049,7 @@ def test_load_best_model_with_save(self): trainer = get_regression_trainer( output_dir=tmp_dir, save_steps=5, - evaluation_strategy="steps", + eval_strategy="steps", eval_steps=5, load_best_model_at_end=True, save_total_limit=2, @@ -2746,7 +2768,7 @@ def test_accelerator_config_from_dict(self): model = RegressionPreTrainedModel(config) eval_dataset = SampleIterableDataset() - accelerator_config = { + accelerator_config: dict[str, Any] = { "split_batches": True, "dispatch_batches": True, "even_batches": False, @@ -2870,61 +2892,6 @@ def test_accelerator_config_from_partial(self): self.assertEqual(trainer.accelerator.even_batches, True) self.assertEqual(trainer.accelerator.use_seedable_sampler, True) - def test_accelerator_config_from_dict_with_deprecated_args(self): - # Checks that accelerator kwargs can be passed through - # and the accelerator is initialized respectively - # and maintains the deprecated args if passed in - with tempfile.TemporaryDirectory() as tmp_dir: - config = RegressionModelConfig(a=1.5, b=2.5) - model = RegressionPreTrainedModel(config) - eval_dataset = SampleIterableDataset() - - # Leaves all options as something *not* basic - with self.assertWarns(FutureWarning) as cm: - gaudi_config = get_gaudi_config() - args = RegressionGaudiTrainingArguments( - output_dir=tmp_dir, - accelerator_config={ - "split_batches": True, - }, - dispatch_batches=False, - use_habana=True, - ) - self.assertIn("dispatch_batches", str(cm.warnings[0].message)) - trainer = GaudiTrainer(model=model, gaudi_config=gaudi_config, args=args, eval_dataset=eval_dataset) - self.assertEqual(trainer.accelerator.dispatch_batches, False) - self.assertEqual(trainer.accelerator.split_batches, True) - with self.assertWarns(FutureWarning) as cm: - args = RegressionGaudiTrainingArguments( - output_dir=tmp_dir, - accelerator_config={ - "even_batches": False, - }, - split_batches=True, - use_habana=True, - ) - self.assertIn("split_batches", str(cm.warnings[0].message)) - trainer = GaudiTrainer(model=model, gaudi_config=gaudi_config, args=args, eval_dataset=eval_dataset) - self.assertEqual(trainer.accelerator.split_batches, True) - self.assertEqual(trainer.accelerator.even_batches, False) - self.assertEqual(trainer.accelerator.dispatch_batches, None) - - def test_accelerator_config_only_deprecated_args(self): - with tempfile.TemporaryDirectory() as tmp_dir: - with self.assertWarns(FutureWarning) as cm: - gaudi_config = get_gaudi_config() - args = RegressionGaudiTrainingArguments( - output_dir=tmp_dir, - split_batches=True, - use_habana=True, - ) - self.assertIn("split_batches", str(cm.warnings[0].message)) - config = RegressionModelConfig(a=1.5, b=2.5) - model = RegressionPreTrainedModel(config) - eval_dataset = SampleIterableDataset() - trainer = GaudiTrainer(model=model, gaudi_config=gaudi_config, args=args, eval_dataset=eval_dataset) - self.assertEqual(trainer.accelerator.split_batches, True) - def test_accelerator_custom_state(self): AcceleratorState._reset_state(reset_partial_state=True) with tempfile.TemporaryDirectory() as tmp_dir: @@ -3256,6 +3223,191 @@ def test_metric_for_best_model_behavior(self): ) self.assertTrue(trainer.args.metric_for_best_model == "loss") + def test_best_model_checkpoint_behavior(self): + # Case 1. Never evaluated, save_total_limit > 1 and save_steps == 1. + # Both best_metric and best_model_checkpoint should be None. + with tempfile.TemporaryDirectory() as tmpdir: + trainer = get_regression_trainer( + output_dir=tmpdir, + eval_strategy="steps", + save_strategy="steps", + save_steps=1, + metric_for_best_model="accuracy", + greater_is_better=True, + ) + trainer.train() + + assert trainer.state.best_metric is None + assert trainer.state.best_model_checkpoint is None + assert len(os.listdir(tmpdir)) == trainer.state.global_step + + # Case 2. Never evaluated and save_total_limit == 1. + # Both best_metric and best_model_checkpoint should be None. + # Only the last checkpoint should remain. + with tempfile.TemporaryDirectory() as tmpdir: + trainer = get_regression_trainer( + output_dir=tmpdir, + eval_strategy="steps", + save_strategy="steps", + save_steps=1, + metric_for_best_model="accuracy", + greater_is_better=True, + save_total_limit=1, + ) + trainer.train() + + num_steps = trainer.state.global_step + + assert trainer.state.best_metric is None + assert trainer.state.best_model_checkpoint is None + assert len(os.listdir(tmpdir)) == 1 + + ckpt = os.path.join(tmpdir, f"{PREFIX_CHECKPOINT_DIR}-{num_steps}") + assert os.path.isdir(ckpt) + assert os.listdir(tmpdir)[0] == f"{PREFIX_CHECKPOINT_DIR}-{num_steps}" + + # Case 3. eval_strategy == save_strategy. + # best_model_checkpoint should be at epoch 1. + with tempfile.TemporaryDirectory() as tmpdir: + trainer = get_regression_trainer( + output_dir=tmpdir, + eval_strategy="epoch", + save_strategy="epoch", + metric_for_best_model="accuracy", + compute_metrics=AlmostAccuracy(), + greater_is_better=True, + load_best_model_at_end=False, + ) + + with unittest.mock.patch.object( + trainer, + "_evaluate", + side_effect=evaluate_side_effect_factory( + [ + {"eval_accuracy": 0.59}, + {"eval_accuracy": 0.57}, + {"eval_accuracy": 0.55}, + ] + ), + ): + trainer.train() + + steps_per_epoch = get_steps_per_epoch(trainer) + + assert trainer.state.best_metric == 0.59 + assert trainer.state.best_global_step == steps_per_epoch + + best_ckpt = os.path.join(tmpdir, f"{PREFIX_CHECKPOINT_DIR}-{trainer.state.best_global_step}") + assert trainer.state.best_model_checkpoint == best_ckpt + + assert len(os.listdir(tmpdir)) == trainer.state.num_train_epochs + + # Case 4. eval_strategy != save_strategy. + with tempfile.TemporaryDirectory() as tmpdir: + trainer = get_regression_trainer( + output_dir=tmpdir, + eval_strategy="epoch", + save_strategy="steps", + save_steps=1, + metric_for_best_model="accuracy", + compute_metrics=AlmostAccuracy(), + greater_is_better=True, + load_best_model_at_end=False, + ) + + with unittest.mock.patch.object( + trainer, + "_evaluate", + side_effect=evaluate_side_effect_factory( + [ + {"eval_accuracy": 0.59}, + {"eval_accuracy": 0.57}, + {"eval_accuracy": 0.55}, + ] + ), + ): + trainer.train() + + steps_per_epoch = get_steps_per_epoch(trainer) + + assert trainer.state.best_metric == 0.59 + assert trainer.state.best_global_step == steps_per_epoch + + best_ckpt = os.path.join(tmpdir, f"{PREFIX_CHECKPOINT_DIR}-{trainer.state.best_global_step}") + assert trainer.state.best_model_checkpoint == best_ckpt + + assert len(os.listdir(tmpdir)) == trainer.state.global_step + + # Case 5. Multiple checkpoints, save_total_limit == 1. + # Best metric is found at step 1 and that checkpoint should be saved. + with tempfile.TemporaryDirectory() as tmpdir: + trainer = get_regression_trainer( + output_dir=tmpdir, + eval_strategy="steps", + eval_steps=1, + save_strategy="steps", + save_steps=1, + metric_for_best_model="accuracy", + compute_metrics=AlmostAccuracy(), + greater_is_better=True, + save_total_limit=1, + ) + + with unittest.mock.patch.object( + trainer, + "_evaluate", + side_effect=evaluate_side_effect_factory( + [ + {"eval_accuracy": 0.90}, + {"eval_accuracy": 0.80}, + {"eval_accuracy": 0.70}, + ] + ), + ): + trainer.train() + + assert trainer.state.best_metric == 0.90 + assert trainer.state.best_global_step == 1 + + best_ckpt = os.path.join(tmpdir, f"{PREFIX_CHECKPOINT_DIR}-{trainer.state.best_global_step}") + assert trainer.state.best_model_checkpoint == best_ckpt + + assert len(os.listdir(tmpdir)) == 1 + + # Case 6. Saving happens more often and eval/save mismatch. + # `best_model_checkpoint` should be None due to a step mismatch. + with tempfile.TemporaryDirectory() as tmpdir: + trainer = get_regression_trainer( + output_dir=tmpdir, + eval_strategy="steps", + eval_steps=3, + save_strategy="steps", + save_steps=2, + metric_for_best_model="accuracy", + compute_metrics=AlmostAccuracy(), + greater_is_better=True, + ) + + with unittest.mock.patch.object( + trainer, + "_evaluate", + side_effect=evaluate_side_effect_factory( + [ + {"eval_accuracy": 0.90}, + {"eval_accuracy": 0.80}, + {"eval_accuracy": 0.70}, + ] + ), + ): + trainer.train() + + assert trainer.state.best_metric == 0.90 + assert trainer.state.best_global_step == 3 + + assert trainer.state.best_model_checkpoint is None + + assert len(os.listdir(tmpdir)) == trainer.state.global_step // 2 + def test_profiling(self): with tempfile.TemporaryDirectory() as tmp_dir: # 24 total steps and compilation takes place during the 1st three steps @@ -3538,7 +3690,7 @@ def model_init(trial): def hp_name(trial): return MyTrialShortNamer.shortname(trial.params) - def compute_objective(metrics: Dict[str, float]) -> List[float]: + def compute_objective(metrics: dict[str, float]) -> list[float]: return metrics["eval_loss"], metrics["eval_accuracy"] with tempfile.TemporaryDirectory() as tmp_dir: @@ -3565,6 +3717,38 @@ def compute_objective(metrics: Dict[str, float]) -> List[float]: ) +@require_torch +@require_optuna +class TrainerHyperParameterOptunaIntegrationTestWithFullEval(unittest.TestCase): + def test_hyperparameter_search(self): + def hp_space(trial): + return {} + + def model_init(trial): + if trial is not None: + a = trial.suggest_int("a", -4, 4) + b = trial.suggest_int("b", -4, 4) + else: + a = 0 + b = 0 + config = RegressionModelConfig(a=a, b=b, double_output=False) + + return RegressionPreTrainedModel(config) + + with tempfile.TemporaryDirectory() as tmp_dir: + trainer = get_regression_trainer( + output_dir=tmp_dir, + disable_tqdm=True, + model_init=model_init, + fp16_full_eval=True, + ) + trainer.hyperparameter_search( + direction="minimize", + hp_space=hp_space, + n_trials=2, + ) + + # TODO: crashes because `TypeError: cannot pickle 'PyCapsule' object` # @require_torch # @require_ray @@ -3691,16 +3875,6 @@ def compute_objective(metrics: Dict[str, float]) -> List[float]: } optim_test_params = [ - ( - OptimizerNames.ADAMW_HF, - transformers.optimization.AdamW, - default_adam_kwargs, - ), - ( - OptimizerNames.ADAMW_HF.value, - transformers.optimization.AdamW, - default_adam_kwargs, - ), ( OptimizerNames.ADAMW_TORCH, torch.optim.AdamW, @@ -3752,9 +3926,6 @@ def test_optim_supported(self, name: str, expected_cls, mandatory_kwargs): # self.batch_size = args.train_batch_size # def test_hyperparameter_search(self): -# class MyTrialShortNamer(TrialShortNamer): -# DEFAULTS = {"a": 0, "b": 0} - # def hp_space(trial): # return { # "method": "random", @@ -3776,9 +3947,6 @@ def test_optim_supported(self, name: str, expected_cls, mandatory_kwargs): # return RegressionPreTrainedModel(model_config) -# def hp_name(params): -# return MyTrialShortNamer.shortname(params) - # with tempfile.TemporaryDirectory() as tmp_dir: # trainer = get_regression_trainer( # output_dir=tmp_dir, @@ -3793,9 +3961,31 @@ def test_optim_supported(self, name: str, expected_cls, mandatory_kwargs): # run_name="test", # model_init=model_init, # ) -# trainer.hyperparameter_search( -# direction="minimize", hp_space=hp_space, hp_name=hp_name, backend="wandb", n_trials=4, anonymous="must" -# ) +# sweep_kwargs = { +# "direction": "minimize", +# "hp_space": hp_space, +# "backend": "wandb", +# "n_trials": 4, +# } +# best_run = trainer.hyperparameter_search(**sweep_kwargs) + +# self.assertIsNotNone(best_run.run_id) +# self.assertIsNotNone(best_run.run_summary) +# hp_keys = set(best_run.hyperparameters.keys()) +# self.assertSetEqual(hp_keys, {"a", "b", "assignments", "metric"}) + +# # pretend restarting the process purged the environ +# import os + +# del os.environ["WANDB_ENTITY"] +# del os.environ["WANDB_PROJECT"] +# sweep_kwargs["sweep_id"] = best_run.run_summary +# updated_best_run = trainer.hyperparameter_search(**sweep_kwargs) + +# self.assertIsNotNone(updated_best_run.run_id) +# self.assertEqual(updated_best_run.run_summary, best_run.run_summary) +# updated_hp_keys = set(updated_best_run.hyperparameters.keys()) +# self.assertSetEqual(updated_hp_keys, {"a", "b", "assignments", "metric"}) class HyperParameterSearchBackendsTest(unittest.TestCase): diff --git a/tests/test_trainer_seq2seq.py b/tests/test_trainer_seq2seq.py index 89905e97e8..f5866043a8 100644 --- a/tests/test_trainer_seq2seq.py +++ b/tests/test_trainer_seq2seq.py @@ -132,7 +132,7 @@ def _compute_metrics(pred): @require_torch def test_bad_generation_config_fail_early(self): - # Tests that a bad geneartion config causes the trainer to fail early + # Tests that a bad generation config causes the trainer to fail early model = AutoModelForSeq2SeqLM.from_pretrained("google-t5/t5-small") tokenizer = AutoTokenizer.from_pretrained("google-t5/t5-small") data_collator = DataCollatorForSeq2Seq(tokenizer, model=model, return_tensors="pt", padding="longest") From 34f32a9b956ceb8ce9ec472e9c25707acb7a0978 Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Mon, 21 Apr 2025 20:00:37 +0000 Subject: [PATCH 04/22] Fixes --- optimum/habana/transformers/modeling_utils.py | 6 ++++-- optimum/habana/transformers/trainer.py | 2 +- tests/test_trainer.py | 1 - 3 files changed, 5 insertions(+), 4 deletions(-) diff --git a/optimum/habana/transformers/modeling_utils.py b/optimum/habana/transformers/modeling_utils.py index 0d490a8e49..c4601d278b 100644 --- a/optimum/habana/transformers/modeling_utils.py +++ b/optimum/habana/transformers/modeling_utils.py @@ -743,8 +743,10 @@ def adapt_transformers_to_gaudi(): transformers.AutoConfig.register("deepseek_v2", DeepseekV2Config) transformers.AutoModelForCausalLM.register(DeepseekV2Config, DeepseekV2ForCausalLM) transformers.AutoTokenizer.register(DeepseekV2Config, fast_tokenizer_class=DeepseekTokenizerFast) - transformers.AutoConfig.register("deepseek_v3", DeepseekV3Config) - transformers.AutoModelForCausalLM.register(DeepseekV3Config, DeepseekV3ForCausalLM) + # transformers.AutoConfig.register("deepseek_v3", DeepseekV3Config) + # transformers.AutoModelForCausalLM.register(DeepseekV3Config, DeepseekV3ForCausalLM) + transformers.models.deepseek_v3.configuration_deepseek_v3.DeepseekV3Config = DeepseekV3Config + transformers.models.deepseek_v3.modeling_deepseek_v3.DeepseekV3ForCausalLM = DeepseekV3ForCausalLM # Optimization for cohere on Gaudi transformers.models.cohere.modeling_cohere.CohereDecoderLayer = GaudiCohereDecoderLayer diff --git a/optimum/habana/transformers/trainer.py b/optimum/habana/transformers/trainer.py index 3d7aa3ffda..30f30d866b 100644 --- a/optimum/habana/transformers/trainer.py +++ b/optimum/habana/transformers/trainer.py @@ -1373,7 +1373,7 @@ def _load_rng_state(self, checkpoint): return with safe_globals(): - checkpoint_rng_state = torch.load(rng_file, weights_only=True) + checkpoint_rng_state = torch.load(rng_file) random.setstate(checkpoint_rng_state["python"]) np.random.set_state(checkpoint_rng_state["numpy"]) torch.random.set_rng_state(checkpoint_rng_state["cpu"]) diff --git a/tests/test_trainer.py b/tests/test_trainer.py index d32e2209a0..7707ce0ff1 100644 --- a/tests/test_trainer.py +++ b/tests/test_trainer.py @@ -3740,7 +3740,6 @@ def model_init(trial): output_dir=tmp_dir, disable_tqdm=True, model_init=model_init, - fp16_full_eval=True, ) trainer.hyperparameter_search( direction="minimize", From 25b053c7bed718aee4bb1b149d9e16e1b477666b Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Tue, 22 Apr 2025 14:09:02 +0000 Subject: [PATCH 05/22] Fix AWQ requirements --- examples/text-generation/requirements_awq.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/text-generation/requirements_awq.txt b/examples/text-generation/requirements_awq.txt index 812d48b233..ef52a8b105 100644 --- a/examples/text-generation/requirements_awq.txt +++ b/examples/text-generation/requirements_awq.txt @@ -1,3 +1,3 @@ triton==3.1.0 autoawq -transformers>=4.48.2,<=4.49.0 +transformers>=4.51.0,<4.52.0 From 071fb639ce1fc14efdbaa1d6c1364a14e12f2a89 Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Wed, 30 Apr 2025 08:13:39 +0000 Subject: [PATCH 06/22] Beam search --- .../habana/transformers/generation/utils.py | 130 ++++++------------ 1 file changed, 41 insertions(+), 89 deletions(-) diff --git a/optimum/habana/transformers/generation/utils.py b/optimum/habana/transformers/generation/utils.py index ded5602844..bc3be24025 100755 --- a/optimum/habana/transformers/generation/utils.py +++ b/optimum/habana/transformers/generation/utils.py @@ -1750,7 +1750,18 @@ def generate( ) elif generation_mode in (GenerationMode.BEAM_SAMPLE, GenerationMode.BEAM_SEARCH): - # 11. interleave input_ids with `num_beams` additional sequences per batch + # 11. prepare beam search scorer + beam_scorer = BeamSearchScorer( + batch_size=batch_size, + num_beams=generation_config.num_beams, + device=inputs_tensor.device, + length_penalty=generation_config.length_penalty, + do_early_stopping=generation_config.early_stopping, + num_beam_hyps_to_keep=generation_config.num_return_sequences, + max_length=generation_config.max_length, + ) + + # 12. interleave input_ids with `num_beams` additional sequences per batch input_ids, model_kwargs = self._expand_inputs_for_generation( input_ids=input_ids, expand_size=generation_config.num_beams, @@ -1758,9 +1769,10 @@ def generate( **model_kwargs, ) - # 12. run beam sample + # 13. run beam sample result = self._beam_search( input_ids, + beam_scorer, logits_processor=prepared_logits_processor, stopping_criteria=prepared_stopping_criteria, generation_config=generation_config, @@ -2921,6 +2933,7 @@ def _sample( def _beam_search( self, input_ids: torch.LongTensor, + beam_scorer: BeamScorer, logits_processor: LogitsProcessorList, stopping_criteria: StoppingCriteriaList, generation_config: GaudiGenerationConfig, @@ -2943,8 +2956,11 @@ def _beam_search( (https://huggingface.co/docs/transformers/main_classes/text_generation#transformers.GenerationMixin.compute_transition_scores) Parameters: - input_ids (`torch.LongTensor` of shape `(batch_size*num_beams, sequence_length)`): + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): The sequence used as a prompt for the generation. + beam_scorer (`BeamScorer`): + An derived instance of [`BeamScorer`] that defines how beam hypotheses are constructed, stored and + sorted during generation. For more information, the documentation of [`BeamScorer`] should be read. logits_processor (`LogitsProcessorList`): An instance of [`LogitsProcessorList`]. List of instances of class derived from [`LogitsProcessor`] used to modify the prediction scores of the language modeling head applied at each generation step. @@ -2991,28 +3007,10 @@ def _beam_search( num_beams = generation_config.num_beams num_return_sequences = generation_config.num_return_sequences - batch_size_unflattened, cur_len = input_ids.shape - batch_size = batch_size_unflattened // num_beams - # TODO (joao): standardize special cases - if self.__class__.__name__ == "MoshiDepthDecoder": - vocab_size = self.config.audio_vocab_size - elif self.__class__.__name__ == "ImageGPTForCausalImageModeling": - vocab_size = self.get_output_embeddings().out_features - else: - vocab_size = self.config.get_text_config().vocab_size - decoder_prompt_len = cur_len - this_peer_finished = False + batch_size = len(beam_scorer._beam_hyps) + num_beams = beam_scorer.num_beams - # At each beam search step, we want to keep top K [K = (number of EOS tokens + 1) * `num_beams`] candidates - # with the highest log-probabilities, or sample K continuations without replacement. We gather the top K - # (as opposed to `num_beams`, or any number lower than K) so that we have at least `num_beams` sequences - # non-finished to continue the live beam search, in case the top `num_beams` all select an EOS token. - n_eos_tokens = eos_token_id.shape[0] if eos_token_id is not None else 0 - beams_to_keep = max(2, 1 + n_eos_tokens) * num_beams - top_num_beam_mask = torch.cat( - (torch.ones((num_beams), dtype=torch.bool), torch.zeros((beams_to_keep - num_beams), dtype=torch.bool)), - dim=0, - ).to(input_ids.device) + batch_beam_size, cur_len = input_ids.shape if "inputs_embeds" in model_kwargs: cur_len = model_kwargs["inputs_embeds"].shape[1] token_idx = model_kwargs.get("token_idx", None) @@ -3022,6 +3020,11 @@ def _beam_search( model_kwargs["cache_position"] = torch.arange(cur_len, device=input_ids.device) + if num_beams * batch_size != batch_beam_size: + raise ValueError( + f"Batch dimension of `input_ids` should be {num_beams * batch_size}, but is {batch_beam_size}." + ) + # (joao) feature lost in the refactor. Probably won't implement, hurts readbility with minimal gains (there # are newer low-memory alternatives like the offloaded cache) sequential = generation_config.low_memory @@ -3032,9 +3035,11 @@ def _beam_search( ) # 2. init output tuples - all_scores = () if (return_dict_in_generate and output_scores) else None + scores = () if (return_dict_in_generate and output_scores) else None raw_logits = () if (return_dict_in_generate and output_logits) else None - beam_indices = () if (return_dict_in_generate and output_logits) else None + beam_indices = ( + tuple(() for _ in range(batch_beam_size)) if (return_dict_in_generate and output_scores) else None + ) decoder_attentions = () if (return_dict_in_generate and output_attentions) else None cross_attentions = () if (return_dict_in_generate and output_attentions) else None decoder_hidden_states = () if (return_dict_in_generate and output_hidden_states) else None @@ -3046,33 +3051,11 @@ def _beam_search( model_kwargs["encoder_outputs"].get("hidden_states") if output_hidden_states else None ) - # 3. init running tensors and static-shaped placeholders - - # per batch, beam-item holding current token in loop and completed sequences - output_fill_value = pad_token_id or eos_token_id[0] if eos_token_id is not None else -1 - running_sequences = torch.full( - (batch_size, num_beams, max_length), - fill_value=output_fill_value, - dtype=torch.int64, - device=input_ids.device, - ) - running_sequences[:, :, :cur_len] = self._unflatten_beam_dim(input_ids, batch_size, num_beams) - sequences = running_sequences.detach().clone() - - # per batch, beam-item score, logprobs # initialise score of first beam with 0 and the rest with -1e9. This makes sure that only tokens # of the first beam are considered to avoid sampling the exact same tokens across all beams. - running_beam_scores = torch.zeros((batch_size, num_beams), dtype=torch.float, device=input_ids.device) - running_beam_scores[:, 1:] = -1e9 - beam_scores = torch.full((batch_size, num_beams), fill_value=-1e9, dtype=torch.float, device=input_ids.device) - - # per batch, beam-item state bit indicating if sentence has finished. - is_sent_finished = torch.zeros((batch_size, num_beams), dtype=torch.bool, device=input_ids.device) - - # per batch, beam-item state bit indicating if there are valid continuations. - next_token_hits_stopping_criteria = torch.zeros( - (batch_size, num_beams), dtype=torch.bool, device=input_ids.device - ) + beam_scores = torch.zeros((batch_size, num_beams), dtype=torch.float, device=input_ids.device) + beam_scores[:, 1:] = -1e9 + beam_scores = beam_scores.view((batch_size * num_beams,)) # Beam token selection: pick 1 + eos_token_id.shape[0] next tokens for each beam so we have at least 1 # non eos token per beam. @@ -3206,6 +3189,7 @@ def expand_if_needed(tensor, new_size, value, dim=-1): initial_ids = input_ids[::num_beams, 0:cur_len] time_to_first_token_done = False + # 4. run the generation loop while self._has_unfinished_sequences(this_peer_finished, synced_gpus, device=input_ids.device): if lazy_mode: self.htcore_generation.mark_step() @@ -3224,44 +3208,12 @@ def expand_if_needed(tensor, new_size, value, dim=-1): model_inputs.update({"output_attentions": output_attentions} if output_attentions else {}) model_inputs.update({"output_hidden_states": output_hidden_states} if output_hidden_states else {}) - # if sequential is True, split the input to batches of batch_size and run sequentially - if sequential: - if any( - model_name in self.__class__.__name__.lower() - for model_name in [ - "fsmt", - "reformer", - "ctrl", - "gpt_bigcode", - "transo_xl", - "xlnet", - "cpm", - "jamba", - ] - ): - raise RuntimeError( - f"Currently generation for {self.__class__.__name__} is not supported " - f"for `low_memory beam_search`. Please open an issue on GitHub if you need this feature." - ) - - inputs_per_sub_batches = _split_model_inputs( - model_inputs, - split_size=batch_size, - full_batch_size=batch_beam_size, - config=self.config.get_text_config(), - ) - outputs_per_sub_batch = [ - self(**inputs_per_sub_batch, return_dict=True) for inputs_per_sub_batch in inputs_per_sub_batches - ] - - outputs = stack_model_outputs(outputs_per_sub_batch, self.config.get_text_config()) - else: - hpu_graphs_kwargs = self._get_hpu_graphs_kwargs(model_kwargs) - outputs = self( - **model_inputs, - return_dict=True, - **hpu_graphs_kwargs, - ) + hpu_graphs_kwargs = self._get_hpu_graphs_kwargs(model_kwargs) + outputs = self( + **model_inputs, + return_dict=True, + **hpu_graphs_kwargs, + ) # synced_gpus: don't waste resources running the code we don't need if synced_gpus and this_peer_finished: From b4e3ed35d3434e885e0d579f001253b2db7de0d2 Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Wed, 30 Apr 2025 08:18:15 +0000 Subject: [PATCH 07/22] Make style --- .../habana/transformers/generation/utils.py | 36 ++++++++----------- .../models/gpt_neox/modeling_gpt_neox.py | 1 - .../models/qwen2_vl/modeling_qwen2_vl.py | 8 +++-- tests/test_trainer.py | 4 +-- 4 files changed, 22 insertions(+), 27 deletions(-) diff --git a/optimum/habana/transformers/generation/utils.py b/optimum/habana/transformers/generation/utils.py index bc3be24025..2dcd78ad9d 100755 --- a/optimum/habana/transformers/generation/utils.py +++ b/optimum/habana/transformers/generation/utils.py @@ -65,7 +65,6 @@ GenerateOutput, GenerationMixin, GenerationMode, - _split_model_inputs, _split_model_outputs, stack_model_outputs, ) @@ -899,7 +898,10 @@ def _prepare_generated_length( return generation_config def _prepare_generation_config( - self, generation_config: Optional[GaudiGenerationConfig], use_model_defaults: Optional[bool] = None, **kwargs: Dict + self, + generation_config: Optional[GaudiGenerationConfig], + use_model_defaults: Optional[bool] = None, + **kwargs: Dict, ) -> Tuple[GaudiGenerationConfig, Dict]: """ Copied from https://github.com/huggingface/transformers/blob/v4.40.2/src/transformers/generation/utils.py#L1230 @@ -978,9 +980,7 @@ def _prepare_generation_config( if generation_config.static_shapes is None: generation_config.static_shapes = self.config.model_type in MODELS_OPTIMIZED_WITH_STATIC_SHAPES if self.config.model_type == "vision-encoder-decoder": - generation_config.static_shapes = ( - self.config.decoder.model_type in MODELS_OPTIMIZED_WITH_STATIC_SHAPES - ) + generation_config.static_shapes = self.config.decoder.model_type in MODELS_OPTIMIZED_WITH_STATIC_SHAPES self.generation_config.static_shapes = generation_config.static_shapes if generation_config.ignore_eos is None: generation_config.ignore_eos = kwargs.get("ignore_eos", kwargs.get("lazy_mode", None)) @@ -2138,9 +2138,7 @@ def _contrastive_search( else: logit_for_next_step = outputs.logits[:, -1, :] # torch.float32 is needed to retain precision for later logits manipulations - logit_for_next_step = logit_for_next_step.to( - copy=True, dtype=torch.float32, device=input_ids.device - ) + logit_for_next_step = logit_for_next_step.to(copy=True, dtype=torch.float32, device=input_ids.device) model_kwargs = self._update_model_kwargs_for_generation( outputs, @@ -2682,9 +2680,7 @@ def _sample( model_kwargs["pad_done"] = False model_kwargs["mqa_model"] = False model_kwargs["lazy_mode"] = lazy_mode - while self._has_unfinished_sequences( - this_peer_finished, synced_gpus, device=input_ids.device - ): + while self._has_unfinished_sequences(this_peer_finished, synced_gpus, device=input_ids.device): if lazy_mode: self.htcore_generation.mark_step() @@ -2901,7 +2897,9 @@ def _sample( + start_token_idx ) # Create a mask for positions greater than the first eos_token_id - mask = torch.arange(generation_config.max_length).expand(batch_size, generation_config.max_length) > eos_positions.unsqueeze(1) + mask = torch.arange(generation_config.max_length).expand( + batch_size, generation_config.max_length + ) > eos_positions.unsqueeze(1) # Apply the mask to set positions greater than the first eos_token_id to pad_token_id input_ids[mask] = pad_token_id @@ -3003,9 +3001,9 @@ def _beam_search( do_sample = generation_config.do_sample early_stopping = generation_config.early_stopping length_penalty = generation_config.length_penalty - max_length = generation_config.max_length + # max_length = generation_config.max_length num_beams = generation_config.num_beams - num_return_sequences = generation_config.num_return_sequences + # num_return_sequences = generation_config.num_return_sequences batch_size = len(beam_scorer._beam_hyps) num_beams = beam_scorer.num_beams @@ -3302,7 +3300,7 @@ def expand_if_needed(tensor, new_size, value, dim=-1): beam_trace_tokens.index_copy_(0, beam_trace_idx, beam_tokens.unsqueeze(0)) beam_trace_idx.add_(1) - if self.generation_config.early_stopping: + if early_stopping: num_eos_tokens.add_(beam_tokens[0:num_beams].eq(self.config.eos_token_id).sum()) beam_scores.add_(torch.where(beam_tokens.eq(self.config.eos_token_id), float("-inf"), 0.0)) @@ -3385,11 +3383,7 @@ def expand_if_needed(tensor, new_size, value, dim=-1): is_min_length_reached = ( self.generation_config.min_length and cur_len >= self.generation_config.min_length ) - if ( - self.generation_config.early_stopping - and is_min_length_reached - and num_eos_tokens >= num_beams_tensor - ): + if early_stopping and is_min_length_reached and num_eos_tokens >= num_beams_tensor: break elif get_final_stopping_criteria(stopping_criteria(input_ids, scores, token_idx=cur_len)): break @@ -3458,7 +3452,7 @@ def move(obj, device): sequence_outputs = {} sequence_outputs["sequences"] = finalize_beams( - initial_ids.cpu(), move(beam_trace, "cpu"), self.config, self.generation_config.length_penalty + initial_ids.cpu(), move(beam_trace, "cpu"), self.config, length_penalty ) else: sequence_outputs = beam_scorer.finalize( diff --git a/optimum/habana/transformers/models/gpt_neox/modeling_gpt_neox.py b/optimum/habana/transformers/models/gpt_neox/modeling_gpt_neox.py index 7a93ed8c07..fcaedb1066 100644 --- a/optimum/habana/transformers/models/gpt_neox/modeling_gpt_neox.py +++ b/optimum/habana/transformers/models/gpt_neox/modeling_gpt_neox.py @@ -420,7 +420,6 @@ def forward( logits_to_keep: Union[int, torch.Tensor] = 0, **kwargs: Unpack[KwargsForCausalLM], ) -> Union[Tuple, CausalLMOutputWithPast]: - outputs: BaseModelOutputWithPast = self.gpt_neox( input_ids, attention_mask=attention_mask, diff --git a/optimum/habana/transformers/models/qwen2_vl/modeling_qwen2_vl.py b/optimum/habana/transformers/models/qwen2_vl/modeling_qwen2_vl.py index a729b8f28b..57bc4e3294 100644 --- a/optimum/habana/transformers/models/qwen2_vl/modeling_qwen2_vl.py +++ b/optimum/habana/transformers/models/qwen2_vl/modeling_qwen2_vl.py @@ -102,9 +102,13 @@ def forward( v = v.transpose(0, 1) if FusedSDPA is not None and use_flash_attention: - attn_output = self.fused_scaled_dot_product_attention(q.unsqueeze(0), k.unsqueeze(0), v.unsqueeze(0), attention_mask, 0.0, False, None, "None") + attn_output = self.fused_scaled_dot_product_attention( + q.unsqueeze(0), k.unsqueeze(0), v.unsqueeze(0), attention_mask, 0.0, False, None, "None" + ) else: - attn_output = F.scaled_dot_product_attention(q.unsqueeze(0), k.unsqueeze(0), v.unsqueeze(0), attention_mask, dropout_p=0.0) + attn_output = F.scaled_dot_product_attention( + q.unsqueeze(0), k.unsqueeze(0), v.unsqueeze(0), attention_mask, dropout_p=0.0 + ) attn_output = attn_output.squeeze(0).transpose(0, 1) attn_output = attn_output.reshape(seq_length, -1) diff --git a/tests/test_trainer.py b/tests/test_trainer.py index 7707ce0ff1..4ad2d23bf5 100644 --- a/tests/test_trainer.py +++ b/tests/test_trainer.py @@ -832,9 +832,7 @@ def tokenize_function(examples): diff_truth = [ abs(base - grad) for base, grad in zip(base_loss_callback.losses, grad_accum_loss_callback.losses) ] - diff_broken = [ - abs(base - grad) for base, grad in zip(base_loss_callback.losses, broken_loss_callback.losses) - ] + diff_broken = [abs(base - grad) for base, grad in zip(base_loss_callback.losses, broken_loss_callback.losses)] # all diff truth should be quite close self.assertLess(max(diff_truth), 0.01, f"Difference {max(diff_truth)} is not within 0.01") From 831bfae2c649c97a4806b63df89aafbf6110c544 Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Wed, 30 Apr 2025 09:38:56 +0000 Subject: [PATCH 08/22] Text-generation fix --- examples/text-generation/run_generation.py | 4 +- examples/text-generation/utils.py | 49 +++++++++---------- optimum/habana/checkpoint_utils.py | 3 +- .../habana/transformers/generation/utils.py | 6 +-- .../fixture/tests/test_encoder_decoder.json | 4 +- .../tests/test_text_generation_example.json | 16 +++--- 6 files changed, 41 insertions(+), 41 deletions(-) diff --git a/examples/text-generation/run_generation.py b/examples/text-generation/run_generation.py index eb96346735..d868bc572c 100755 --- a/examples/text-generation/run_generation.py +++ b/examples/text-generation/run_generation.py @@ -37,8 +37,6 @@ save_model, ) -from optimum.habana.utils import HabanaGenerationTime, get_hpu_memory_stats - logging.basicConfig( format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", @@ -442,6 +440,8 @@ def main(): if args.sdp_on_bf16: torch._C._set_math_sdp_allow_fp16_bf16_reduction(True) + from optimum.habana.utils import HabanaGenerationTime, get_hpu_memory_stats + if args.dataset_name is None: # Benchmark over the prompts below if args.prompt: diff --git a/examples/text-generation/utils.py b/examples/text-generation/utils.py index 43ea9cf6d5..e99544aa27 100644 --- a/examples/text-generation/utils.py +++ b/examples/text-generation/utils.py @@ -29,21 +29,6 @@ from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer from transformers.utils import check_min_version -from optimum.habana.checkpoint_utils import ( - get_ds_injection_policy, - get_repo_root, - model_is_optimized, - model_on_meta, - write_checkpoints_json, -) -from optimum.habana.utils import ( - HabanaGenerationTime, - check_habana_frameworks_version, - check_optimum_habana_min_version, - get_habana_frameworks_version, - set_seed, -) - def adjust_batch(batch, size): curr_size = batch["input_ids"].shape[1] @@ -105,6 +90,8 @@ def setup_distributed(args): def setup_inference(args, model): import habana_frameworks.torch.core as htcore + from optimum.habana.utils import get_habana_frameworks_version + habana_version = get_habana_frameworks_version() print("Initializing inference mode") @@ -130,9 +117,6 @@ def setup_const_serialization(const_serialization_path): def setup_env(args): - # Will error if the minimal version of Transformers is not installed. Remove at your own risks. - check_min_version("4.45.0") - check_optimum_habana_min_version("1.18.0.dev0") # TODO: SW-167588 - WA for memory issue in hqt prep_model os.environ.setdefault("EXPERIMENTAL_WEIGHT_SHARING", "FALSE") @@ -149,6 +133,13 @@ def setup_env(args): # we can call HPU graphs clear_inputs(). os.environ.setdefault("PT_HPUGRAPH_DISABLE_TENSOR_CACHE", "1") + # Will error if the minimal version of Transformers is not installed. Remove at your own risks. + check_min_version("4.51.0") + + from optimum.habana.utils import check_optimum_habana_min_version + + check_optimum_habana_min_version("1.18.0.dev0") + # Tweak generation so that it runs faster on Gaudi from optimum.habana.transformers.modeling_utils import adapt_transformers_to_gaudi @@ -308,11 +299,8 @@ def setup_model(args, model_dtype, model_kwargs, logger): from optimum.habana.transformers.trainer import _is_peft_model - if check_habana_frameworks_version("1.13.0") and model.config.model_type == "falcon": - model = wrap_in_hpu_graph(model, hash_with_views=False) - else: - max_graphs = getattr(args, "max_graphs", None) - model = wrap_in_hpu_graph(model, max_graphs=max_graphs) + max_graphs = getattr(args, "max_graphs", None) + model = wrap_in_hpu_graph(model, max_graphs=max_graphs) if args.assistant_model is not None: assistant_model = wrap_in_hpu_graph(assistant_model) if _is_peft_model(model): @@ -437,6 +425,8 @@ def setup_distributed_model_ep(args, model_dtype, model_kwargs, logger): def setup_distributed_model(args, model_dtype, model_kwargs, logger): import deepspeed + from optimum.habana.checkpoint_utils import get_ds_injection_policy, model_on_meta, write_checkpoints_json + # List of model types that need max position embeddings capped at 8192 MODELS_WITH_POS_EMBEDDING_LIMIT = ["llama"] @@ -653,6 +643,8 @@ def setup_generation_config(args, model, assistant_model, tokenizer): if args.force_words is not None: force_words_ids = [tokenizer.encode(force_word, add_special_tokens=False) for force_word in args.force_words] + from optimum.habana.checkpoint_utils import model_is_optimized + is_optimized = model_is_optimized(model.config) # Generation configuration @@ -706,8 +698,6 @@ def exclude_hpu_graph_configs(args): def initialize_model(args, logger): - timer = HabanaGenerationTime() - timer.start() setup_distributed(args) if not args.world_size > 0 and args.attn_batch_split > 1: logger.warning("Disabling attention batch splitting as it's unnecessary for single-card execution") @@ -717,10 +707,19 @@ def initialize_model(args, logger): override_prints(args.global_rank == 0 or args.verbose_workers, logger) setup_env(args) setup_device(args) + + from optimum.habana.utils import HabanaGenerationTime, set_seed + + timer = HabanaGenerationTime() + timer.start() set_seed(args.seed) + + from optimum.habana.checkpoint_utils import get_repo_root + cache_dir = get_repo_root(args.model_name_or_path, local_rank=args.local_rank, token=args.token) if args.assistant_model is not None: get_repo_root(args.assistant_model, local_rank=args.local_rank, token=args.token) + use_deepspeed = args.world_size > 0 if use_deepspeed or args.bf16: model_dtype = torch.bfloat16 diff --git a/optimum/habana/checkpoint_utils.py b/optimum/habana/checkpoint_utils.py index 6a6001d5e0..f9c7859f45 100644 --- a/optimum/habana/checkpoint_utils.py +++ b/optimum/habana/checkpoint_utils.py @@ -4,7 +4,6 @@ import torch from huggingface_hub import list_repo_files, snapshot_download -from transformers import modeling_utils from transformers.utils import is_offline_mode @@ -64,6 +63,8 @@ def get_checkpoint_files(model_name_or_path, local_rank, token=None): # Extensions: .bin | .safetensors | .pt # Creates a list of paths from all downloaded files in cache dir + from transformers import modeling_utils + if any(file.suffix == ".bin" for file in Path(cached_repo_dir).rglob("*")): (name, ext) = os.path.splitext(modeling_utils.WEIGHTS_NAME) elif any(file.suffix == ".safetensors" for file in Path(cached_repo_dir).rglob("*")): diff --git a/optimum/habana/transformers/generation/utils.py b/optimum/habana/transformers/generation/utils.py index 2dcd78ad9d..e08f7c304e 100755 --- a/optimum/habana/transformers/generation/utils.py +++ b/optimum/habana/transformers/generation/utils.py @@ -1061,9 +1061,9 @@ def _prepare_cache_for_generation( ) generation_config.cache_implementation = None - generation_config.cache_implementation = generation_config.cache_implementation or getattr( - self.config.get_text_config(), "cache_implementation", None - ) + # generation_config.cache_implementation = generation_config.cache_implementation or getattr( + # self.config.get_text_config(), "cache_implementation", None + # ) if generation_config.cache_implementation is not None: if generation_config.cache_implementation in NEED_SETUP_CACHE_CLASSES_MAPPING: if generation_config.cache_implementation == "static" and not self._supports_static_cache: diff --git a/tests/baselines/fixture/tests/test_encoder_decoder.json b/tests/baselines/fixture/tests/test_encoder_decoder.json index 670e29464c..5959b46026 100644 --- a/tests/baselines/fixture/tests/test_encoder_decoder.json +++ b/tests/baselines/fixture/tests/test_encoder_decoder.json @@ -10,7 +10,7 @@ }, "gaudi3": { "predict_rougeLsum": 15.618, - "predict_samples_per_second": 1.091 + "predict_samples_per_second": 0.894 } }, "tests/test_encoder_decoder.py::TestEncoderDecoderModels::test_text_summarization_bf16[t5-3b-Habana/t5-2-1]": { @@ -38,7 +38,7 @@ }, "gaudi3": { "predict_bleu": 11.7168, - "predict_samples_per_second": 18.174 + "predict_samples_per_second": 10.944 } } } \ No newline at end of file diff --git a/tests/baselines/fixture/tests/test_text_generation_example.json b/tests/baselines/fixture/tests/test_text_generation_example.json index 8b8c2d15ef..376b87db6b 100644 --- a/tests/baselines/fixture/tests/test_text_generation_example.json +++ b/tests/baselines/fixture/tests/test_text_generation_example.json @@ -78,7 +78,7 @@ "throughput": 44.25834541569395 }, "gaudi3": { - "throughput": 179.15343204459856 + "throughput": 143.5977987468177 } }, "tests/test_text_generation_example.py::test_text_generation_bf16_1x[Qwen/Qwen2-7B-256-False-True]": { @@ -209,7 +209,7 @@ "throughput": 35 }, "gaudi3": { - "throughput": 149.2189570033595 + "throughput": 95.405132122058 } }, "tests/test_text_generation_example.py::test_text_generation_bf16_1x[facebook/xglm-1.7B-1-False-False]": { @@ -415,7 +415,7 @@ "throughput": 51.61471298016438 }, "gaudi3": { - "throughput": 69.74689153288725 + "throughput": 48.4989896475782 } }, "tests/test_text_generation_example.py::test_text_generation_deepspeed[Qwen/Qwen2.5-72B-2-1]": { @@ -436,7 +436,7 @@ "throughput": 36.77314954096159 }, "gaudi3": { - "throughput": 42.964481338739304 + "throughput": 34.82088387988216 } }, "tests/test_text_generation_example.py::test_text_generation_deepspeed[facebook/opt-66b-2-1]": { @@ -452,7 +452,7 @@ "throughput": 87.578709544111 }, "gaudi3": { - "throughput": 107.59395201764178 + "throughput": 60.55676566393658 } }, "tests/test_text_generation_example.py::test_text_generation_deepspeed[google/gemma-2-9b-8-1]": { @@ -460,7 +460,7 @@ "throughput": 110.12610917383735 }, "gaudi3": { - "throughput": 123.69992293361813 + "throughput": 68.9805898434787 } }, "tests/test_text_generation_example.py::test_text_generation_deepspeed[meta-llama/Llama-2-70b-hf-8-1]": { @@ -476,7 +476,7 @@ "throughput": 64 }, "gaudi3": { - "throughput": 75.6224035651044 + "throughput": 40.96632613314894 } }, "tests/test_text_generation_example.py::test_text_generation_distributed_tp[meta-llama/Llama-2-7b-hf]": { @@ -516,7 +516,7 @@ "throughput": 4656.2 }, "gaudi3": { - "throughput": 6968.716105590979 + "throughput": 6059.088893259565 } }, "tests/test_text_generation_example.py::test_text_generation_fp8[meta-llama/Llama-2-7b-hf-1-1230-False-128-128]": { From 0947f5b8f20724177196a08bfe83c180013a7a06 Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Mon, 5 May 2025 12:54:59 +0000 Subject: [PATCH 09/22] 1x tests fix --- tests/baselines/fixture/tests/test_examples.json | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/baselines/fixture/tests/test_examples.json b/tests/baselines/fixture/tests/test_examples.json index 2efcc37194..810390ccc8 100644 --- a/tests/baselines/fixture/tests/test_examples.json +++ b/tests/baselines/fixture/tests/test_examples.json @@ -161,7 +161,7 @@ "gaudi3": { "eval_accuracy": 0.9849333333333333, "train_runtime": 86.4819, - "train_samples_per_second": 1043.6 + "train_samples_per_second": 949.371 } }, "tests/test_examples.py::ImageClassificationExampleTester::test_run_image_classification_vit-base-patch16-224-in21k_single_card": { @@ -682,8 +682,8 @@ }, "gaudi3": { "eval_f1": 94.36192902198283, - "train_runtime": 260.988, - "train_samples_per_second": 423.007 + "train_runtime": 276.0212, + "train_samples_per_second": 400.151 } }, "tests/test_examples.py::TextClassificationExampleTester::test_run_glue_bert-large-uncased-whole-word-masking_single_card": { @@ -699,7 +699,7 @@ }, "gaudi3": { "eval_f1": 0.8826, - "train_runtime": 78.5458, + "train_runtime": 85.085, "train_samples_per_second": 1674.509 } } From 8b2e896e28ee587c1edeb08a5e2748ef7da600ef Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Thu, 8 May 2025 08:21:40 +0000 Subject: [PATCH 10/22] Fix 8x tests --- tests/baselines/fixture/tests/test_examples.json | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/baselines/fixture/tests/test_examples.json b/tests/baselines/fixture/tests/test_examples.json index 810390ccc8..d4098b1030 100644 --- a/tests/baselines/fixture/tests/test_examples.json +++ b/tests/baselines/fixture/tests/test_examples.json @@ -189,7 +189,7 @@ "train_samples_per_second": 1955.74 }, "gaudi3": { - "eval_accuracy": 0.19650135869565216, + "eval_accuracy": 0.18546195652173914, "eval_samples_per_second": 2442.221, "train_runtime": 118.0048, "train_samples_per_second": 2462.94 @@ -209,7 +209,7 @@ "train_samples_per_second": 2975.844 }, "gaudi3": { - "eval_accuracy": 0.7352241847826086, + "eval_accuracy": 0.7202785326086957, "eval_samples_per_second": 2059.992, "train_runtime": 71.6177, "train_samples_per_second": 3977.496 @@ -446,7 +446,7 @@ "gaudi3": { "eval_accuracy": 0.9817333333333333, "train_runtime": 106.7478, - "train_samples_per_second": 7042.508 + "train_samples_per_second": 5241.263 } }, "tests/test_examples.py::MultiCardImageClassificationExampleTester::test_run_image_classification_vit-base-patch16-224-in21k_multi_card": { @@ -614,8 +614,8 @@ "gaudi3": { "eval_samples_per_second": 48.493, "eval_wer": 0.4296589018302829, - "train_runtime": 321.6139, - "train_samples_per_second": 504.087 + "train_runtime": 351.7181, + "train_samples_per_second": 417.561 } }, "tests/test_examples.py::MultiCardSpeechRecognitionExampleTester::test_run_speech_recognition_ctc_wav2vec2-large-lv60_multi_card": { @@ -632,7 +632,7 @@ "train_samples_per_second": 225.572 }, "gaudi3": { - "eval_samples_per_second": 491.004, + "eval_samples_per_second": 369.034, "eval_wer": 0.06197937326457755, "train_runtime": 235.8412, "train_samples_per_second": 392.016 From 762c08fe351b1b18bf0eb09f9acaca7409aa4ba3 Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Thu, 22 May 2025 13:44:14 +0000 Subject: [PATCH 11/22] Fix 8x tests --- examples/trl/reward_modeling.py | 2 +- optimum/habana/transformers/models/clip/modeling_clip.py | 1 + tests/baselines/fixture/tests/test_examples.json | 6 +++--- 3 files changed, 5 insertions(+), 4 deletions(-) diff --git a/examples/trl/reward_modeling.py b/examples/trl/reward_modeling.py index ec81d2eca8..122f7820c7 100644 --- a/examples/trl/reward_modeling.py +++ b/examples/trl/reward_modeling.py @@ -78,7 +78,7 @@ class ScriptArguments: metadata={"help": "Enables gradient checkpointing."}, ) optim: Optional[str] = field( - default="adamw_hf", + default="adamw_torch", metadata={"help": "The optimizer to use."}, ) lr_scheduler_type: Optional[str] = field( diff --git a/optimum/habana/transformers/models/clip/modeling_clip.py b/optimum/habana/transformers/models/clip/modeling_clip.py index 2f24e76848..274a8f12bd 100644 --- a/optimum/habana/transformers/models/clip/modeling_clip.py +++ b/optimum/habana/transformers/models/clip/modeling_clip.py @@ -342,6 +342,7 @@ def forward( interpolate_pos_encoding: bool = False, use_flash_attention: Optional[bool] = False, flash_attention_recompute: Optional[bool] = False, + return_dict: Optional[bool] = None, ) -> BaseModelOutputWithPooling: """ Copied from CLIPVisionModel.forward: https://github.com/huggingface/transformers/blob/ab0f050b42d903f34d6eb97f3f8c0c07f0517ad2/src/transformers/models/clip/modeling_clip.py diff --git a/tests/baselines/fixture/tests/test_examples.json b/tests/baselines/fixture/tests/test_examples.json index d4098b1030..8d502e379a 100644 --- a/tests/baselines/fixture/tests/test_examples.json +++ b/tests/baselines/fixture/tests/test_examples.json @@ -405,7 +405,7 @@ "train_samples_per_second": 63.161 }, "gaudi3": { - "perplexity": 1.2158095633720596, + "perplexity": 1.2817338202117678, "train_runtime": 16.7016, "train_samples_per_second": 61.478 } @@ -564,8 +564,8 @@ "train_samples_per_second": 1.6 }, "gaudi3": { - "train_runtime": 169.1724, - "train_samples_per_second": 2.418 + "train_runtime": 196.1679, + "train_samples_per_second": 2.085 } }, "tests/test_examples.py::MultiCardSFTChatExampleTester::test_sft_Qwen2-7B_multi_card": { From 1c91d3f597510d1a5d96cd9ce80f368464972585 Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Thu, 29 May 2025 08:59:16 +0000 Subject: [PATCH 12/22] Fix mllama test --- optimum/habana/transformers/models/mllama/modeling_mllama.py | 2 ++ tests/baselines/fixture/tests/test_examples.json | 4 ++-- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/optimum/habana/transformers/models/mllama/modeling_mllama.py b/optimum/habana/transformers/models/mllama/modeling_mllama.py index 154807bd89..223bf26a22 100644 --- a/optimum/habana/transformers/models/mllama/modeling_mllama.py +++ b/optimum/habana/transformers/models/mllama/modeling_mllama.py @@ -1084,6 +1084,7 @@ def forward( else: cross_attention_mask = cross_attention_mask[:, :, -1:] full_text_row_masked_out_mask = full_text_row_masked_out_mask[:, :, -1:] + outputs = self.language_model( input_ids=input_ids, attention_mask=attention_mask, @@ -1094,6 +1095,7 @@ def forward( past_key_values=past_key_values, use_cache=use_cache, inputs_embeds=inputs_embeds, + labels=labels, output_hidden_states=output_hidden_states, output_attentions=output_attentions, return_dict=return_dict, diff --git a/tests/baselines/fixture/tests/test_examples.json b/tests/baselines/fixture/tests/test_examples.json index 8d502e379a..2749b22e6b 100644 --- a/tests/baselines/fixture/tests/test_examples.json +++ b/tests/baselines/fixture/tests/test_examples.json @@ -417,7 +417,7 @@ "train_samples_per_second": 125.798 }, "gaudi3": { - "perplexity": 2.613, + "perplexity": 2.6391937381661483, "train_runtime": 321.3809, "train_samples_per_second": 168.296 } @@ -498,7 +498,7 @@ }, "gaudi3": { "eval_accuracy": 0.20785648331296863, - "train_runtime": 184.9003, + "train_runtime": 257.3229, "train_samples_per_second": 27.828 } }, From 31bf07ef37a7b18243819327b4eaa8e441b714e4 Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Fri, 30 May 2025 16:51:06 +0000 Subject: [PATCH 13/22] Fix CI tests --- tests/baselines/fixture/tests/test_examples.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/baselines/fixture/tests/test_examples.json b/tests/baselines/fixture/tests/test_examples.json index 2749b22e6b..d739aa211d 100644 --- a/tests/baselines/fixture/tests/test_examples.json +++ b/tests/baselines/fixture/tests/test_examples.json @@ -209,7 +209,7 @@ "train_samples_per_second": 2975.844 }, "gaudi3": { - "eval_accuracy": 0.7202785326086957, + "eval_accuracy": 0.7048233695652174, "eval_samples_per_second": 2059.992, "train_runtime": 71.6177, "train_samples_per_second": 3977.496 From 3515a12cf6af02ddf8ca5d25a2b1fe205c0b4656 Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Mon, 11 Aug 2025 17:14:40 +0000 Subject: [PATCH 14/22] Update examples --- examples/audio-classification/README.md | 4 +- .../audio-classification/requirements.txt | 2 +- .../run_audio_classification.py | 7 +- .../contrastive-image-text/run_bridgetower.py | 7 +- examples/contrastive-image-text/run_clip.py | 7 +- examples/image-classification/README.md | 6 +- .../run_image_classification.py | 9 +- .../run_image2text_lora_finetune.py | 2 +- examples/language-modeling/run_clm.py | 173 ++++++++++++------ examples/language-modeling/run_lora_clm.py | 2 +- examples/language-modeling/run_mlm.py | 23 +-- .../run_multitask_prompt_tuning.py | 4 +- .../run_prompt_tuning_clm.py | 4 +- examples/protein-folding/run_esmfold.py | 2 +- .../run_sequence_classification.py | 2 +- .../protein-folding/run_zero_shot_eval.py | 2 +- examples/question-answering/run_qa.py | 11 +- examples/question-answering/run_seq2seq_qa.py | 11 +- examples/question-answering/trainer_qa.py | 1 - .../question-answering/trainer_seq2seq_qa.py | 7 +- examples/question-answering/utils_qa.py | 15 +- examples/speech-recognition/requirements.txt | 2 +- .../run_speech_recognition_ctc.py | 15 +- .../run_speech_recognition_seq2seq.py | 15 +- .../depth_to_image_generation.py | 2 +- .../image_to_image_generation.py | 2 +- .../image_to_video_generation.py | 2 +- .../text_to_image_generation.py | 2 +- .../text_to_video_generation.py | 2 +- .../training/train_controlnet.py | 2 +- .../unconditional_image_generation.py | 4 +- examples/summarization/run_summarization.py | 7 +- examples/text-classification/run_glue.py | 29 ++- examples/translation/run_translation.py | 7 +- 34 files changed, 221 insertions(+), 171 deletions(-) diff --git a/examples/audio-classification/README.md b/examples/audio-classification/README.md index ffc38e6709..224224edc0 100644 --- a/examples/audio-classification/README.md +++ b/examples/audio-classification/README.md @@ -150,10 +150,10 @@ PT_HPU_LAZY_MODE=1 python run_audio_classification.py \ $ apt install git-lfs ``` -2. Log in with your HuggingFace account credentials using `huggingface-cli` +2. Log in with your HuggingFace account credentials using `hf` ```bash -$ huggingface-cli login +$ hf auth login # ...follow the prompts ``` diff --git a/examples/audio-classification/requirements.txt b/examples/audio-classification/requirements.txt index bae36f7451..e440989242 100644 --- a/examples/audio-classification/requirements.txt +++ b/examples/audio-classification/requirements.txt @@ -1,4 +1,4 @@ -datasets>=1.14.0 +datasets[audio]>=1.14.0 evaluate numba==0.60.0 librosa diff --git a/examples/audio-classification/run_audio_classification.py b/examples/audio-classification/run_audio_classification.py index 073b8ad577..a60110b141 100644 --- a/examples/audio-classification/run_audio_classification.py +++ b/examples/audio-classification/run_audio_classification.py @@ -1,5 +1,4 @@ #!/usr/bin/env python -# coding=utf-8 # Copyright 2021 The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -46,8 +45,8 @@ def check_optimum_habana_min_version(*a, **b): logger = logging.getLogger(__name__) # Will error if the minimal version of Transformers and Optimum Habana are not installed. Remove at your own risks. -check_min_version("4.51.0") -check_optimum_habana_min_version("1.18.0.dev0") +check_min_version("4.55.0") +check_optimum_habana_min_version("1.19.0.dev0") require_version("datasets>=1.14.0", "To fix: pip install -r examples/pytorch/audio-classification/requirements.txt") @@ -159,7 +158,7 @@ class ModelArguments: metadata={ "help": ( "The token to use as HTTP bearer authorization for remote files. If not specified, will use the token " - "generated when running `huggingface-cli login` (stored in `~/.huggingface`)." + "generated when running `hf auth login` (stored in `~/.huggingface`)." ) }, ) diff --git a/examples/contrastive-image-text/run_bridgetower.py b/examples/contrastive-image-text/run_bridgetower.py index 0b98b7a0b2..0fc04ae95e 100644 --- a/examples/contrastive-image-text/run_bridgetower.py +++ b/examples/contrastive-image-text/run_bridgetower.py @@ -1,5 +1,4 @@ #!/usr/bin/env python -# coding=utf-8 # Copyright 2023 The HuggingFace Team All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -58,8 +57,8 @@ def check_optimum_habana_min_version(*a, **b): logger = logging.getLogger(__name__) # Will error if the minimal version of Transformers and Optimum Habana are not installed. Remove at your own risks. -check_min_version("4.51.0") -check_optimum_habana_min_version("1.18.0.dev0") +check_min_version("4.55.0") +check_optimum_habana_min_version("1.19.0.dev0") require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/contrastive-image-text/requirements.txt") @@ -96,7 +95,7 @@ class ModelArguments: metadata={ "help": ( "The token to use as HTTP bearer authorization for remote files. If not specified, will use the token " - "generated when running `huggingface-cli login` (stored in `~/.huggingface`)." + "generated when running `hf auth login` (stored in `~/.huggingface`)." ) }, ) diff --git a/examples/contrastive-image-text/run_clip.py b/examples/contrastive-image-text/run_clip.py index 2e928ec3d6..2d16de2d04 100644 --- a/examples/contrastive-image-text/run_clip.py +++ b/examples/contrastive-image-text/run_clip.py @@ -1,5 +1,4 @@ #!/usr/bin/env python -# coding=utf-8 # Copyright 2022 The HuggingFace Team All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -61,8 +60,8 @@ def check_optimum_habana_min_version(*a, **b): logger = logging.getLogger(__name__) # Will error if the minimal version of Transformers and Optimum Habana are not installed. Remove at your own risks. -check_min_version("4.51.0") -check_optimum_habana_min_version("1.18.0.dev0") +check_min_version("4.55.0") +check_optimum_habana_min_version("1.19.0.dev0") require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/contrastive-image-text/requirements.txt") @@ -99,7 +98,7 @@ class ModelArguments: metadata={ "help": ( "The token to use as HTTP bearer authorization for remote files. If not specified, will use the token " - "generated when running `huggingface-cli login` (stored in `~/.huggingface`)." + "generated when running `hf auth login` (stored in `~/.huggingface`)." ) }, ) diff --git a/examples/image-classification/README.md b/examples/image-classification/README.md index 7a0f7fef38..4253210f08 100644 --- a/examples/image-classification/README.md +++ b/examples/image-classification/README.md @@ -143,7 +143,7 @@ dataset = load_dataset("imagefolder", data_files={"train": ["path/to/file1", "pa Next, push it to the hub! ```python -# assuming you have ran the huggingface-cli login command in a terminal +# assuming you have ran the hf auth login command in a terminal dataset.push_to_hub("name_of_your_dataset") # if you want to push to a private repo, simply pass private=True: @@ -166,10 +166,10 @@ $ git config --global user.email "you@example.com" $ git config --global user.name "Your Name" ``` -2. Log in with your HuggingFace account credentials using `huggingface-cli`: +2. Log in with your HuggingFace account credentials using `hf`: ```bash -$ huggingface-cli login +$ hf auth login # ...follow the prompts ``` diff --git a/examples/image-classification/run_image_classification.py b/examples/image-classification/run_image_classification.py index a82428eb94..39ea1141d4 100644 --- a/examples/image-classification/run_image_classification.py +++ b/examples/image-classification/run_image_classification.py @@ -1,5 +1,4 @@ #!/usr/bin/env python -# coding=utf-8 # Copyright 2021 The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -64,8 +63,8 @@ def check_optimum_habana_min_version(*a, **b): logger = logging.getLogger(__name__) # Will error if the minimal version of Transformers and Optimum Habana are not installed. Remove at your own risks. -check_min_version("4.51.0") -check_optimum_habana_min_version("1.18.0.dev0") +check_min_version("4.55.0") +check_optimum_habana_min_version("1.19.0.dev0") require_version("datasets>=2.14.0", "To fix: pip install -r examples/pytorch/image-classification/requirements.txt") @@ -165,7 +164,7 @@ class ModelArguments: metadata={ "help": ( "The token to use as HTTP bearer authorization for remote files. If not specified, will use the token " - "generated when running `huggingface-cli login` (stored in `~/.huggingface`)." + "generated when running `hf auth login` (stored in `~/.huggingface`)." ) }, ) @@ -294,7 +293,7 @@ def collate_fn(examples): return {"pixel_values": pixel_values, "labels": labels} # If we don't have a validation split, split off a percentage of train as validation. - data_args.train_val_split = None if "validation" in dataset.keys() else data_args.train_val_split + data_args.train_val_split = None if "validation" in dataset else data_args.train_val_split if isinstance(data_args.train_val_split, float) and data_args.train_val_split > 0.0: split = dataset["train"].train_test_split(data_args.train_val_split) dataset["train"] = split["train"] diff --git a/examples/image-to-text/run_image2text_lora_finetune.py b/examples/image-to-text/run_image2text_lora_finetune.py index 927d58749f..95307f229a 100644 --- a/examples/image-to-text/run_image2text_lora_finetune.py +++ b/examples/image-to-text/run_image2text_lora_finetune.py @@ -55,7 +55,7 @@ def check_optimum_habana_min_version(*a, **b): logger = logging.getLogger(__name__) # Will error if the minimal version of Optimum Habana is not installed. Remove at your own risks. -check_optimum_habana_min_version("1.18.0.dev0") +check_optimum_habana_min_version("1.19.0.dev0") def normalized_levenshtein(s1, s2): diff --git a/examples/language-modeling/run_clm.py b/examples/language-modeling/run_clm.py index 1c90b93ed3..fb4e6ae555 100644 --- a/examples/language-modeling/run_clm.py +++ b/examples/language-modeling/run_clm.py @@ -1,5 +1,4 @@ #!/usr/bin/env python -# coding=utf-8 # Copyright 2022 The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -32,7 +31,7 @@ import evaluate import torch import transformers -from datasets import load_dataset +from datasets import IterableDataset, IterableDatasetDict, load_dataset from transformers import ( CONFIG_MAPPING, MODEL_FOR_CAUSAL_LM_MAPPING, @@ -62,8 +61,8 @@ def check_optimum_habana_min_version(*a, **b): logger = logging.getLogger(__name__) # Will error if the minimal version of Transformers and Optimum Habana are not installed. Remove at your own risks. -check_min_version("4.51.0") -check_optimum_habana_min_version("1.18.0.dev0") +check_min_version("4.55.0") +check_optimum_habana_min_version("1.19.0.dev0") require_version("datasets>=2.14.0", "To fix: pip install -r examples/pytorch/language-modeling/requirements.txt") @@ -123,7 +122,7 @@ class ModelArguments: metadata={ "help": ( "The token to use as HTTP bearer authorization for remote files. If not specified, will use the token " - "generated when running `huggingface-cli login` (stored in `~/.huggingface`)." + "generated when running `hf auth login` (stored in `~/.huggingface`)." ) }, ) @@ -182,15 +181,6 @@ class ModelArguments: ) }, ) - low_cpu_mem_usage: bool = field( - default=False, - metadata={ - "help": ( - "It is an option to create the model as an empty shell, then only materialize its parameters when the pretrained weights are loaded. " - "Setting it to True will benefit LLM loading time and RAM consumption." - ) - }, - ) def __post_init__(self): if self.config_overrides is not None and (self.config_name is not None or self.model_name_or_path is not None): @@ -281,6 +271,45 @@ def __post_init__(self): assert extension in ["csv", "json", "txt"], "`validation_file` should be a csv, a json or a txt file." +def split_streaming_dataset( + full_streaming_dataset, + validation_percentage: int = 5, +) -> IterableDatasetDict: + """ + Splits a streaming dataset into + training and validation IterableDatasets, and supports methods like .map(), .filter(), + .take() and properties like .features on the resulting streams. + + Args: + full_streaming_dataset (Dataset): The name of the dataset to load (e.g., "HuggingFaceFW/fineweb"). + validation_percentage (int): The proportion of the dataset to be used for validation split. + + Returns: + IterableDatasetDict: An IterableDatasetDict containing two IterableDataset objects: (train_stream, validation_stream). + """ + if not (0 < validation_percentage < 100): + raise ValueError( + f"validation_percentage must be between 0 and 100 (exclusive). Passed: {validation_percentage}" + ) + + def split_generator(is_train: bool): + for i, example in enumerate(full_streaming_dataset): + if is_train: + if i % 100 > validation_percentage: + yield example + else: + if i % 100 < validation_percentage: + yield example + + features = full_streaming_dataset.features + train_stream = IterableDataset.from_generator(split_generator, gen_kwargs={"is_train": True}, features=features) + validation_stream = IterableDataset.from_generator( + split_generator, gen_kwargs={"is_train": False}, features=features + ) + + return IterableDatasetDict({"train": train_stream, "validation": validation_stream}) + + def main(): # See all possible arguments in src/transformers/training_args.py # or by passing the --help flag to this script. @@ -369,25 +398,37 @@ def main(): streaming=data_args.streaming, trust_remote_code=model_args.trust_remote_code, ) - if "validation" not in raw_datasets.keys(): - raw_datasets["validation"] = load_dataset( - data_args.dataset_name, - data_args.dataset_config_name, - split=f"train[:{data_args.validation_split_percentage}%]", - cache_dir=model_args.cache_dir, - token=model_args.token, - streaming=data_args.streaming, - trust_remote_code=model_args.trust_remote_code, - ) - raw_datasets["train"] = load_dataset( - data_args.dataset_name, - data_args.dataset_config_name, - split=f"train[{data_args.validation_split_percentage}%:]", - cache_dir=model_args.cache_dir, - token=model_args.token, - streaming=data_args.streaming, - trust_remote_code=model_args.trust_remote_code, - ) + if "validation" not in raw_datasets: + if data_args.streaming: + dataset_stream = load_dataset( + data_args.dataset_name, + data_args.dataset_config_name, + split="train", + cache_dir=model_args.cache_dir, + token=model_args.token, + streaming=data_args.streaming, + trust_remote_code=model_args.trust_remote_code, + ) + raw_datasets = split_streaming_dataset(dataset_stream, data_args.validation_split_percentage) + else: + raw_datasets["validation"] = load_dataset( + data_args.dataset_name, + data_args.dataset_config_name, + split=f"train[:{data_args.validation_split_percentage}%]", + cache_dir=model_args.cache_dir, + token=model_args.token, + streaming=data_args.streaming, + trust_remote_code=model_args.trust_remote_code, + ) + raw_datasets["train"] = load_dataset( + data_args.dataset_name, + data_args.dataset_config_name, + split=f"train[{data_args.validation_split_percentage}%:]", + cache_dir=model_args.cache_dir, + token=model_args.token, + streaming=data_args.streaming, + trust_remote_code=model_args.trust_remote_code, + ) else: data_files = {} dataset_args = {} @@ -411,23 +452,34 @@ def main(): **dataset_args, ) # If no validation data is there, validation_split_percentage will be used to divide the dataset. - if "validation" not in raw_datasets.keys(): - raw_datasets["validation"] = load_dataset( - extension, - data_files=data_files, - split=f"train[:{data_args.validation_split_percentage}%]", - cache_dir=model_args.cache_dir, - token=model_args.token, - **dataset_args, - ) - raw_datasets["train"] = load_dataset( - extension, - data_files=data_files, - split=f"train[{data_args.validation_split_percentage}%:]", - cache_dir=model_args.cache_dir, - token=model_args.token, - **dataset_args, - ) + if "validation" not in raw_datasets: + if data_args.streaming: + dataset_stream = load_dataset( + extension, + data_files=data_files, + split="train", + cache_dir=model_args.cache_dir, + token=model_args.token, + **dataset_args, + ) + raw_datasets = split_streaming_dataset(dataset_stream, data_args.validation_split_percentage) + else: + raw_datasets["validation"] = load_dataset( + extension, + data_files=data_files, + split=f"train[:{data_args.validation_split_percentage}%]", + cache_dir=model_args.cache_dir, + token=model_args.token, + **dataset_args, + ) + raw_datasets["train"] = load_dataset( + extension, + data_files=data_files, + split=f"train[{data_args.validation_split_percentage}%:]", + cache_dir=model_args.cache_dir, + token=model_args.token, + **dataset_args, + ) # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at # https://huggingface.co/docs/datasets/loading_datasets. @@ -498,7 +550,6 @@ def main(): token=model_args.token, trust_remote_code=model_args.trust_remote_code, torch_dtype=torch_dtype, - low_cpu_mem_usage=model_args.low_cpu_mem_usage, ) else: model = AutoModelForCausalLM.from_config(config, trust_remote_code=model_args.trust_remote_code) @@ -588,7 +639,7 @@ def tokenize_function(examples): # Main data processing function that will concatenate all texts from our dataset and generate chunks of block_size. def group_texts(examples): # Concatenate all texts. - concatenated_examples = {k: list(chain(*examples[k])) for k in examples.keys()} + concatenated_examples = {k: list(chain(*examples[k])) for k in examples} total_length = len(concatenated_examples[list(examples.keys())[0]]) # We drop the small remainder, and if the total_length < block_size we exclude this batch and return an empty dict. # We could add padding if the model supported it instead of this drop, you can customize this part to your needs. @@ -634,16 +685,22 @@ def tensor_mapper(x): if training_args.resume_from_checkpoint is not None and training_args.resume_from_checkpoint != "": train_dataset = train_dataset.map(tensor_mapper) if data_args.max_train_samples is not None: - max_train_samples = min(len(train_dataset), data_args.max_train_samples) - train_dataset = train_dataset.select(range(max_train_samples)) + if data_args.streaming: + train_dataset = train_dataset.take(data_args.max_train_samples) + else: + max_train_samples = min(len(train_dataset), data_args.max_train_samples) + train_dataset = train_dataset.select(range(max_train_samples)) if training_args.do_eval: if "validation" not in tokenized_datasets: raise ValueError("--do_eval requires a validation dataset") eval_dataset = lm_datasets["validation"] if data_args.max_eval_samples is not None: - max_eval_samples = min(len(eval_dataset), data_args.max_eval_samples) - eval_dataset = eval_dataset.select(range(max_eval_samples)) + if data_args.streaming: + eval_dataset = eval_dataset.take(data_args.max_eval_samples) + else: + max_eval_samples = min(len(eval_dataset), data_args.max_eval_samples) + eval_dataset = eval_dataset.select(range(max_eval_samples)) def preprocess_logits_for_metrics(logits, labels): if isinstance(logits, tuple): @@ -706,7 +763,9 @@ def compute_metrics(eval_preds): logger.info("*** Evaluate ***") metrics = trainer.evaluate() - if not data_args.streaming: + if data_args.streaming: + metrics["eval_samples"] = max_eval_samples + else: max_eval_samples = ( data_args.max_eval_samples if data_args.max_eval_samples is not None else len(eval_dataset) ) diff --git a/examples/language-modeling/run_lora_clm.py b/examples/language-modeling/run_lora_clm.py index 1d4d328139..b22eabba44 100644 --- a/examples/language-modeling/run_lora_clm.py +++ b/examples/language-modeling/run_lora_clm.py @@ -70,7 +70,7 @@ def check_optimum_habana_min_version(*a, **b): logger = logging.getLogger(__name__) # Will error if the minimal version of Optimum Habana is not installed. Remove at your own risks. -check_optimum_habana_min_version("1.18.0.dev0") +check_optimum_habana_min_version("1.19.0.dev0") @dataclass diff --git a/examples/language-modeling/run_mlm.py b/examples/language-modeling/run_mlm.py index 98741f2b4b..e065228381 100644 --- a/examples/language-modeling/run_mlm.py +++ b/examples/language-modeling/run_mlm.py @@ -1,5 +1,4 @@ #!/usr/bin/env python -# coding=utf-8 # Copyright 2020 The HuggingFace Team All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -61,8 +60,8 @@ def check_optimum_habana_min_version(*a, **b): logger = logging.getLogger(__name__) # Will error if the minimal version of Transformers and Optimum Habana are not installed. Remove at your own risks. -check_min_version("4.51.0") -check_optimum_habana_min_version("1.18.0.dev0") +check_min_version("4.55.0") +check_optimum_habana_min_version("1.19.0.dev0") require_version("datasets>=2.14.0", "To fix: pip install -r examples/pytorch/language-modeling/requirements.txt") @@ -121,7 +120,7 @@ class ModelArguments: metadata={ "help": ( "The token to use as HTTP bearer authorization for remote files. If not specified, will use the token " - "generated when running `huggingface-cli login` (stored in `~/.huggingface`)." + "generated when running `hf auth login` (stored in `~/.huggingface`)." ) }, ) @@ -145,15 +144,6 @@ class ModelArguments: "choices": ["auto", "bfloat16", "float32"], }, ) - low_cpu_mem_usage: bool = field( - default=False, - metadata={ - "help": ( - "It is an option to create the model as an empty shell, then only materialize its parameters when the pretrained weights are loaded. " - "Setting it to True will benefit LLM loading time and RAM consumption." - ) - }, - ) def __post_init__(self): if self.config_overrides is not None and (self.config_name is not None or self.model_name_or_path is not None): @@ -342,7 +332,7 @@ def main(): streaming=data_args.streaming, trust_remote_code=model_args.trust_remote_code, ) - if "validation" not in raw_datasets.keys(): + if "validation" not in raw_datasets: raw_datasets["validation"] = load_dataset( data_args.dataset_name, data_args.dataset_config_name, @@ -379,7 +369,7 @@ def main(): ) # If no validation data is there, validation_split_percentage will be used to divide the dataset. - if "validation" not in raw_datasets.keys(): + if "validation" not in raw_datasets: raw_datasets["validation"] = load_dataset( extension, data_files=data_files, @@ -453,7 +443,6 @@ def main(): token=model_args.token, trust_remote_code=model_args.trust_remote_code, torch_dtype=torch_dtype, - low_cpu_mem_usage=model_args.low_cpu_mem_usage, ) else: logger.info("Training new model from scratch") @@ -553,7 +542,7 @@ def tokenize_function(examples): # max_seq_length. def group_texts(examples): # Concatenate all texts. - concatenated_examples = {k: list(chain(*examples[k])) for k in examples.keys()} + concatenated_examples = {k: list(chain(*examples[k])) for k in examples} total_length = len(concatenated_examples[list(examples.keys())[0]]) # We drop the small remainder, and if the total_length < max_seq_length we exclude this batch and return an empty dict. # We could add padding if the model supported it instead of this drop, you can customize this part to your needs. diff --git a/examples/language-modeling/run_multitask_prompt_tuning.py b/examples/language-modeling/run_multitask_prompt_tuning.py index 1cd743a874..1c2c70266d 100644 --- a/examples/language-modeling/run_multitask_prompt_tuning.py +++ b/examples/language-modeling/run_multitask_prompt_tuning.py @@ -60,8 +60,8 @@ def check_optimum_habana_min_version(*a, **b): logger = logging.getLogger(__name__) # Will error if the minimal version of Transformers and Optimum Habana are not installed. Remove at your own risk. -check_min_version("4.49.0") -check_optimum_habana_min_version("1.18.0.dev0") +check_min_version("4.55.0") +check_optimum_habana_min_version("1.19.0.dev0") require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/language-modeling/requirements.txt") diff --git a/examples/language-modeling/run_prompt_tuning_clm.py b/examples/language-modeling/run_prompt_tuning_clm.py index 1a35196445..17f3cd0077 100644 --- a/examples/language-modeling/run_prompt_tuning_clm.py +++ b/examples/language-modeling/run_prompt_tuning_clm.py @@ -62,8 +62,8 @@ def check_optimum_habana_min_version(*a, **b): logger = logging.getLogger(__name__) # Will error if the minimal version of Transformers and Optimum Habana are not installed. Remove at your own risks. -check_min_version("4.49.0") -check_optimum_habana_min_version("1.18.0.dev0") +check_min_version("4.55.0") +check_optimum_habana_min_version("1.19.0.dev0") require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/language-modeling/requirements.txt") diff --git a/examples/protein-folding/run_esmfold.py b/examples/protein-folding/run_esmfold.py index 2e2003ab1b..94036a423e 100644 --- a/examples/protein-folding/run_esmfold.py +++ b/examples/protein-folding/run_esmfold.py @@ -40,7 +40,7 @@ def check_optimum_habana_min_version(*a, **b): # Will error if the minimal version of Optimum Habana is not installed. Remove at your own risks. -check_optimum_habana_min_version("1.18.0.dev0") +check_optimum_habana_min_version("1.19.0.dev0") def convert_outputs_to_pdb(outputs): diff --git a/examples/protein-folding/run_sequence_classification.py b/examples/protein-folding/run_sequence_classification.py index 6c69e2f62e..a2aed8bc6a 100644 --- a/examples/protein-folding/run_sequence_classification.py +++ b/examples/protein-folding/run_sequence_classification.py @@ -41,7 +41,7 @@ def check_optimum_habana_min_version(*a, **b): # Will error if the minimal version of Optimum Habana is not installed. Remove at your own risks. -check_optimum_habana_min_version("1.18.0.dev0") +check_optimum_habana_min_version("1.19.0.dev0") logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) diff --git a/examples/protein-folding/run_zero_shot_eval.py b/examples/protein-folding/run_zero_shot_eval.py index 83aa1c0ce0..3d3a4edadf 100644 --- a/examples/protein-folding/run_zero_shot_eval.py +++ b/examples/protein-folding/run_zero_shot_eval.py @@ -36,7 +36,7 @@ def check_optimum_habana_min_version(*a, **b): # Will error if the minimal version of Optimum Habana is not installed. Remove at your own risks. -check_optimum_habana_min_version("1.18.0.dev0") +check_optimum_habana_min_version("1.19.0.dev0") logging.basicConfig( diff --git a/examples/question-answering/run_qa.py b/examples/question-answering/run_qa.py index 064717d80f..7cc38496d4 100644 --- a/examples/question-answering/run_qa.py +++ b/examples/question-answering/run_qa.py @@ -1,5 +1,4 @@ #!/usr/bin/env python -# coding=utf-8 # Copyright 2022 The HuggingFace Team All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -59,8 +58,8 @@ def check_optimum_habana_min_version(*a, **b): logger = logging.getLogger(__name__) # Will error if the minimal version of Transformers and Optimum Habana are not installed. Remove at your own risks. -check_min_version("4.51.0") -check_optimum_habana_min_version("1.18.0.dev0") +check_min_version("4.55.0") +check_optimum_habana_min_version("1.19.0.dev0") require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/question-answering/requirements.txt") @@ -93,7 +92,7 @@ class ModelArguments: metadata={ "help": ( "The token to use as HTTP bearer authorization for remote files. If not specified, will use the token " - "generated when running `huggingface-cli login` (stored in `~/.huggingface`)." + "generated when running `hf auth login` (stored in `~/.huggingface`)." ) }, ) @@ -378,8 +377,8 @@ def main(): model = model.to("hpu") # Tokenizer check: this script requires a fast tokenizer. if not isinstance(tokenizer, PreTrainedTokenizerFast): - raise ValueError( - "This example script only works for models that have a fast tokenizer. Checkout the big table of models at" + raise TypeError( + "This example script only works for models that have a fast tokenizer. Check out the big table of models at" " https://huggingface.co/transformers/index.html#supported-frameworks to find the model types that meet" " this requirement" ) diff --git a/examples/question-answering/run_seq2seq_qa.py b/examples/question-answering/run_seq2seq_qa.py index 374ec915ca..8d85cae733 100644 --- a/examples/question-answering/run_seq2seq_qa.py +++ b/examples/question-answering/run_seq2seq_qa.py @@ -1,5 +1,4 @@ #!/usr/bin/env python -# coding=utf-8 # Copyright 2021 The HuggingFace Team All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -22,7 +21,7 @@ import os import sys from dataclasses import dataclass, field -from typing import List, Optional, Tuple +from typing import Optional import datasets import evaluate @@ -56,8 +55,8 @@ def check_optimum_habana_min_version(*a, **b): logger = logging.getLogger(__name__) # Will error if the minimal version of Transformers and Optimum Habana are not installed. Remove at your own risks. -check_min_version("4.51.0") -check_optimum_habana_min_version("1.18.0.dev0") +check_min_version("4.55.0") +check_optimum_habana_min_version("1.19.0.dev0") require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/question-answering/requirements.txt") @@ -94,7 +93,7 @@ class ModelArguments: metadata={ "help": ( "The token to use as HTTP bearer authorization for remote files. If not specified, will use the token " - "generated when running `huggingface-cli login` (stored in `~/.huggingface`)." + "generated when running `hf auth login` (stored in `~/.huggingface`)." ) }, ) @@ -487,7 +486,7 @@ def preprocess_squad_batch( question_column: str, context_column: str, answer_column: str, - ) -> Tuple[List[str], List[str]]: + ) -> tuple[list[str], list[str]]: questions = examples[question_column] contexts = examples[context_column] answers = examples[answer_column] diff --git a/examples/question-answering/trainer_qa.py b/examples/question-answering/trainer_qa.py index 3f684502a2..dd57d8f731 100644 --- a/examples/question-answering/trainer_qa.py +++ b/examples/question-answering/trainer_qa.py @@ -1,4 +1,3 @@ -# coding=utf-8 # Copyright 2022 The HuggingFace Team All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/examples/question-answering/trainer_seq2seq_qa.py b/examples/question-answering/trainer_seq2seq_qa.py index 7a8a31c420..f79606b9e7 100644 --- a/examples/question-answering/trainer_seq2seq_qa.py +++ b/examples/question-answering/trainer_seq2seq_qa.py @@ -1,4 +1,3 @@ -# coding=utf-8 # Copyright 2021 The HuggingFace Team All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -18,7 +17,7 @@ import math import time -from typing import Dict, List, Optional +from typing import Optional from torch.utils.data import Dataset from transformers.trainer_utils import PredictionOutput, speed_metrics @@ -37,10 +36,10 @@ def evaluate( self, eval_dataset: Optional[Dataset] = None, eval_examples=None, - ignore_keys: Optional[List[str]] = None, + ignore_keys: Optional[list[str]] = None, metric_key_prefix: str = "eval", **gen_kwargs, - ) -> Dict[str, float]: + ) -> dict[str, float]: gen_kwargs = gen_kwargs.copy() # Use legacy argument setting if a) the option is not explicitly passed; and b) the argument is set in the diff --git a/examples/question-answering/utils_qa.py b/examples/question-answering/utils_qa.py index c596bf98ef..eb07ee6df7 100644 --- a/examples/question-answering/utils_qa.py +++ b/examples/question-answering/utils_qa.py @@ -1,4 +1,3 @@ -# coding=utf-8 # Copyright 2022 The HuggingFace Team All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -20,7 +19,7 @@ import json import logging import os -from typing import Optional, Tuple +from typing import Optional import numpy as np from tqdm.auto import tqdm @@ -32,7 +31,7 @@ def postprocess_qa_predictions( examples, features, - predictions: Tuple[np.ndarray, np.ndarray], + predictions: tuple[np.ndarray, np.ndarray], version_2_with_negative: bool = False, n_best_size: int = 20, max_answer_length: int = 30, @@ -47,7 +46,7 @@ def postprocess_qa_predictions( Args: examples: The non-preprocessed dataset (see the main script for more information). features: The processed dataset (see the main script for more information). - predictions (:obj:`Tuple[np.ndarray, np.ndarray]`): + predictions (:obj:`tuple[np.ndarray, np.ndarray]`): The predictions of the model: two arrays containing the start logits and the end logits respectively. Its first dimension must match the number of elements of :obj:`features`. version_2_with_negative (:obj:`bool`, `optional`, defaults to :obj:`False`): @@ -221,7 +220,7 @@ def postprocess_qa_predictions( # If we have an output_dir, let's save all those dicts. if output_dir is not None: if not os.path.isdir(output_dir): - raise EnvironmentError(f"{output_dir} is not a directory.") + raise OSError(f"{output_dir} is not a directory.") prediction_file = os.path.join( output_dir, "predictions.json" if prefix is None else f"{prefix}_predictions.json" @@ -251,7 +250,7 @@ def postprocess_qa_predictions( def postprocess_qa_predictions_with_beam_search( examples, features, - predictions: Tuple[np.ndarray, np.ndarray], + predictions: tuple[np.ndarray, np.ndarray], version_2_with_negative: bool = False, n_best_size: int = 20, max_answer_length: int = 30, @@ -268,7 +267,7 @@ def postprocess_qa_predictions_with_beam_search( Args: examples: The non-preprocessed dataset (see the main script for more information). features: The processed dataset (see the main script for more information). - predictions (:obj:`Tuple[np.ndarray, np.ndarray]`): + predictions (:obj:`tuple[np.ndarray, np.ndarray]`): The predictions of the model: two arrays containing the start logits and the end logits respectively. Its first dimension must match the number of elements of :obj:`features`. version_2_with_negative (:obj:`bool`, `optional`, defaults to :obj:`False`): @@ -414,7 +413,7 @@ def postprocess_qa_predictions_with_beam_search( # If we have an output_dir, let's save all those dicts. if output_dir is not None: if not os.path.isdir(output_dir): - raise EnvironmentError(f"{output_dir} is not a directory.") + raise OSError(f"{output_dir} is not a directory.") prediction_file = os.path.join( output_dir, "predictions.json" if prefix is None else f"{prefix}_predictions.json" diff --git a/examples/speech-recognition/requirements.txt b/examples/speech-recognition/requirements.txt index 67aeeaaa30..5b66d1c038 100644 --- a/examples/speech-recognition/requirements.txt +++ b/examples/speech-recognition/requirements.txt @@ -1,4 +1,4 @@ -datasets >= 1.18.0 +datasets[audio] >= 1.18.0 numba==0.60.0 librosa jiwer diff --git a/examples/speech-recognition/run_speech_recognition_ctc.py b/examples/speech-recognition/run_speech_recognition_ctc.py index 3e40517af3..0d6ecfdff4 100644 --- a/examples/speech-recognition/run_speech_recognition_ctc.py +++ b/examples/speech-recognition/run_speech_recognition_ctc.py @@ -1,5 +1,4 @@ #!/usr/bin/env python -# coding=utf-8 # Copyright 2021 The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -23,7 +22,7 @@ import re import sys from dataclasses import dataclass, field -from typing import Dict, List, Optional, Union +from typing import Optional, Union import datasets import evaluate @@ -58,8 +57,8 @@ def check_optimum_habana_min_version(*a, **b): logger = logging.getLogger(__name__) # Will error if the minimal version of Transformers and Optimum Habana are not installed. Remove at your own risks. -check_min_version("4.51.0") -check_optimum_habana_min_version("1.18.0.dev0") +check_min_version("4.55.0") +check_optimum_habana_min_version("1.19.0.dev0") require_version("datasets>=1.18.0", "To fix: pip install -r examples/pytorch/speech-recognition/requirements.txt") @@ -243,11 +242,11 @@ class DataTrainingArguments: ) }, ) - chars_to_ignore: Optional[List[str]] = list_field( + chars_to_ignore: Optional[list[str]] = list_field( default=None, metadata={"help": "A list of characters to remove from the transcripts."}, ) - eval_metrics: List[str] = list_field( + eval_metrics: list[str] = list_field( default=["wer"], metadata={"help": "A list of metrics the model should be evaluated on. E.g. `'wer cer'`"}, ) @@ -279,7 +278,7 @@ class DataTrainingArguments: metadata={ "help": ( "The token to use as HTTP bearer authorization for remote files. If not specified, will use the token " - "generated when running `huggingface-cli login` (stored in `~/.huggingface`)." + "generated when running `hf auth login` (stored in `~/.huggingface`)." ) }, ) @@ -350,7 +349,7 @@ class DataCollatorCTCWithPadding: pad_to_multiple_of_labels: Optional[int] = None feature_extractor_input_name: Optional[str] = "input_values" - def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]: + def __call__(self, features: list[dict[str, Union[list[int], torch.Tensor]]]) -> dict[str, torch.Tensor]: # split inputs and labels since they have to be of different lengths and need # different padding methods input_features = [ diff --git a/examples/speech-recognition/run_speech_recognition_seq2seq.py b/examples/speech-recognition/run_speech_recognition_seq2seq.py index f52bd73887..de8ef98350 100755 --- a/examples/speech-recognition/run_speech_recognition_seq2seq.py +++ b/examples/speech-recognition/run_speech_recognition_seq2seq.py @@ -1,5 +1,4 @@ #!/usr/bin/env python -# coding=utf-8 # Copyright 2021 The HuggingFace Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -23,7 +22,7 @@ import os import sys from dataclasses import dataclass, field -from typing import Any, Dict, List, Optional, Union +from typing import Any, Optional, Union import datasets import evaluate @@ -55,8 +54,8 @@ def check_optimum_habana_min_version(*a, **b): # Will error if the minimal version of Transformers is not installed. Remove at your own risks. -check_min_version("4.51.0") -check_optimum_habana_min_version("1.18.0.dev0") +check_min_version("4.55.0") +check_optimum_habana_min_version("1.19.0.dev0") require_version("datasets>=1.18.0", "To fix: pip install -r examples/pytorch/speech-recognition/requirements.txt") @@ -98,7 +97,7 @@ class ModelArguments: metadata={ "help": ( "The token to use as HTTP bearer authorization for remote files. If not specified, will use the token " - "generated when running `huggingface-cli login` (stored in `~/.huggingface`)." + "generated when running `hf auth login` (stored in `~/.huggingface`)." ) }, ) @@ -118,11 +117,11 @@ class ModelArguments: freeze_encoder: bool = field( default=False, metadata={"help": "Whether to freeze the entire encoder of the seq2seq model."} ) - forced_decoder_ids: List[List[int]] = field( + forced_decoder_ids: list[list[int]] = field( default=None, metadata={"help": "Deprecated. Please use the `language` and `task` arguments instead."}, ) - suppress_tokens: List[int] = field( + suppress_tokens: list[int] = field( default=None, metadata={ "help": ( @@ -260,7 +259,7 @@ class DataCollatorSpeechSeq2SeqWithPadding: forward_attention_mask: bool label_features_max_length: int - def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]: + def __call__(self, features: list[dict[str, Union[list[int], torch.Tensor]]]) -> dict[str, torch.Tensor]: # split inputs and labels since they have to be of different lengths and need # different padding methods model_input_name = self.processor.model_input_names[0] diff --git a/examples/stable-diffusion/depth_to_image_generation.py b/examples/stable-diffusion/depth_to_image_generation.py index fcd89257d1..0443bc5080 100755 --- a/examples/stable-diffusion/depth_to_image_generation.py +++ b/examples/stable-diffusion/depth_to_image_generation.py @@ -41,7 +41,7 @@ def check_optimum_habana_min_version(*a, **b): # Will error if the minimal version of Optimum Habana is not installed. Remove at your own risks. -check_optimum_habana_min_version("1.18.0.dev0") +check_optimum_habana_min_version("1.19.0.dev0") logger = logging.getLogger(__name__) diff --git a/examples/stable-diffusion/image_to_image_generation.py b/examples/stable-diffusion/image_to_image_generation.py index 9542931b11..f55c01a6f9 100755 --- a/examples/stable-diffusion/image_to_image_generation.py +++ b/examples/stable-diffusion/image_to_image_generation.py @@ -41,7 +41,7 @@ def check_optimum_habana_min_version(*a, **b): # Will error if the minimal version of Optimum Habana is not installed. Remove at your own risks. -check_optimum_habana_min_version("1.18.0.dev0") +check_optimum_habana_min_version("1.19.0.dev0") logger = logging.getLogger(__name__) diff --git a/examples/stable-diffusion/image_to_video_generation.py b/examples/stable-diffusion/image_to_video_generation.py index 3aacdb51a1..c2be57980a 100755 --- a/examples/stable-diffusion/image_to_video_generation.py +++ b/examples/stable-diffusion/image_to_video_generation.py @@ -38,7 +38,7 @@ def check_optimum_habana_min_version(*a, **b): # Will error if the minimal version of Optimum Habana is not installed. Remove at your own risks. -check_optimum_habana_min_version("1.18.0.dev0") +check_optimum_habana_min_version("1.19.0.dev0") logger = logging.getLogger(__name__) diff --git a/examples/stable-diffusion/text_to_image_generation.py b/examples/stable-diffusion/text_to_image_generation.py index e3b0beed48..aac565dcd5 100755 --- a/examples/stable-diffusion/text_to_image_generation.py +++ b/examples/stable-diffusion/text_to_image_generation.py @@ -42,7 +42,7 @@ def check_optimum_habana_min_version(*a, **b): # Will error if the minimal version of Optimum Habana is not installed. Remove at your own risks. -check_optimum_habana_min_version("1.18.0.dev0") +check_optimum_habana_min_version("1.19.0.dev0") logger = logging.getLogger(__name__) diff --git a/examples/stable-diffusion/text_to_video_generation.py b/examples/stable-diffusion/text_to_video_generation.py index 144727cbc1..5ab6bf8697 100755 --- a/examples/stable-diffusion/text_to_video_generation.py +++ b/examples/stable-diffusion/text_to_video_generation.py @@ -37,7 +37,7 @@ def check_optimum_habana_min_version(*a, **b): # Will error if the minimal version of Optimum Habana is not installed. Remove at your own risks. -check_optimum_habana_min_version("1.18.0.dev0") +check_optimum_habana_min_version("1.19.0.dev0") logger = logging.getLogger(__name__) diff --git a/examples/stable-diffusion/training/train_controlnet.py b/examples/stable-diffusion/training/train_controlnet.py index d6c1a391e9..cb199802bf 100755 --- a/examples/stable-diffusion/training/train_controlnet.py +++ b/examples/stable-diffusion/training/train_controlnet.py @@ -67,7 +67,7 @@ def check_optimum_habana_min_version(*a, **b): # Will error if the minimal version of Optimum Habana is not installed. Remove at your own risks. -check_optimum_habana_min_version("1.18.0.dev0") +check_optimum_habana_min_version("1.19.0.dev0") if is_wandb_available(): import wandb diff --git a/examples/stable-diffusion/unconditional_image_generation.py b/examples/stable-diffusion/unconditional_image_generation.py index 979f60b838..1fd9026e1d 100755 --- a/examples/stable-diffusion/unconditional_image_generation.py +++ b/examples/stable-diffusion/unconditional_image_generation.py @@ -19,8 +19,8 @@ def check_optimum_habana_min_version(*a, **b): return () -check_min_version("4.51.0") -check_optimum_habana_min_version("1.18.0.dev0") +check_min_version("4.55.0") +check_optimum_habana_min_version("1.19.0.dev0") # Setup logging logging.basicConfig( diff --git a/examples/summarization/run_summarization.py b/examples/summarization/run_summarization.py index 87e5faa9d1..7a36a18e3e 100755 --- a/examples/summarization/run_summarization.py +++ b/examples/summarization/run_summarization.py @@ -1,5 +1,4 @@ #!/usr/bin/env python -# coding=utf-8 # Copyright 2022 The HuggingFace Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -64,8 +63,8 @@ def check_optimum_habana_min_version(*a, **b): logger = logging.getLogger(__name__) # Will error if the minimal version of Transformers and Optimum Habana are not installed. Remove at your own risks. -check_min_version("4.51.0") -check_optimum_habana_min_version("1.18.0.dev0") +check_min_version("4.55.0") +check_optimum_habana_min_version("1.19.0.dev0") require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/summarization/requirements.txt") @@ -106,7 +105,7 @@ class ModelArguments: metadata={ "help": ( "The token to use as HTTP bearer authorization for remote files. If not specified, will use the token " - "generated when running `huggingface-cli login` (stored in `~/.huggingface`)." + "generated when running `hf auth login` (stored in `~/.huggingface`)." ) }, ) diff --git a/examples/text-classification/run_glue.py b/examples/text-classification/run_glue.py index be36e601ad..12509510ca 100755 --- a/examples/text-classification/run_glue.py +++ b/examples/text-classification/run_glue.py @@ -1,5 +1,4 @@ #!/usr/bin/env python -# coding=utf-8 # Copyright 2020 The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -20,6 +19,7 @@ import os import random import sys +from collections import Counter from dataclasses import dataclass, field from typing import Optional @@ -57,8 +57,8 @@ def check_optimum_habana_min_version(*a, **b): logger = logging.getLogger(__name__) # Will error if the minimal version of Transformers and Optimum Habana are not installed. Remove at your own risks. -check_min_version("4.51.0") -check_optimum_habana_min_version("1.18.0.dev0") +check_min_version("4.55.0") +check_optimum_habana_min_version("1.19.0.dev0") require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/text-classification/requirements.txt") @@ -158,7 +158,7 @@ class DataTrainingArguments: def __post_init__(self): if self.task_name is not None: self.task_name = self.task_name.lower() - if self.task_name not in task_to_keys.keys(): + if self.task_name not in task_to_keys: raise ValueError("Unknown task, you should pick one in " + ",".join(task_to_keys.keys())) elif self.dataset_name is not None: pass @@ -205,7 +205,7 @@ class ModelArguments: metadata={ "help": ( "The token to use as HTTP bearer authorization for remote files. If not specified, will use the token " - "generated when running `huggingface-cli login` (stored in `~/.huggingface`)." + "generated when running `hf auth login` (stored in `~/.huggingface`)." ) }, ) @@ -345,7 +345,7 @@ def main(): else: raise ValueError("Need either a GLUE task or a test file for `do_predict`.") - for key in data_files.keys(): + for key in data_files: logger.info(f"load a local file for {key}: {data_files[key]}") if data_args.train_file.endswith(".csv"): @@ -501,6 +501,14 @@ def preprocess_function(examples): load_from_cache_file=not data_args.overwrite_cache, desc="Running tokenizer on dataset", ) + + def print_class_distribution(dataset, split_name): + label_counts = Counter(dataset["label"]) + total = sum(label_counts.values()) + logger.info(f"Class distribution in {split_name} set:") + for label, count in label_counts.items(): + logger.info(f" Label {label}: {count} ({count / total:.2%})") + if training_args.do_train: if "train" not in raw_datasets: raise ValueError("--do_train requires a train dataset") @@ -508,6 +516,7 @@ def preprocess_function(examples): if data_args.max_train_samples is not None: max_train_samples = min(len(train_dataset), data_args.max_train_samples) train_dataset = train_dataset.select(range(max_train_samples)) + print_class_distribution(train_dataset, "train") if training_args.do_eval: if "validation" not in raw_datasets and "validation_matched" not in raw_datasets: @@ -516,6 +525,7 @@ def preprocess_function(examples): if data_args.max_eval_samples is not None: max_eval_samples = min(len(eval_dataset), data_args.max_eval_samples) eval_dataset = eval_dataset.select(range(max_eval_samples)) + print_class_distribution(eval_dataset, "validation") if training_args.do_predict or data_args.task_name is not None or data_args.test_file is not None: if "test" not in raw_datasets and "test_matched" not in raw_datasets: @@ -524,6 +534,7 @@ def preprocess_function(examples): if data_args.max_predict_samples is not None: max_predict_samples = min(len(predict_dataset), data_args.max_predict_samples) predict_dataset = predict_dataset.select(range(max_predict_samples)) + print_class_distribution(predict_dataset, "test") # Log a few random samples from the training set: if training_args.do_train: @@ -542,8 +553,12 @@ def preprocess_function(examples): # predictions and label_ids field) and has to return a dictionary string to float. def compute_metrics(p: EvalPrediction): preds = p.predictions[0] if isinstance(p.predictions, tuple) else p.predictions + labels = p.label_ids + if not training_args.eval_do_concat_batches: + preds = np.concatenate(preds, axis=0) + labels = np.concatenate(p.label_ids, axis=0) preds = np.squeeze(preds) if is_regression else np.argmax(preds, axis=1) - result = metric.compute(predictions=preds, references=p.label_ids) + result = metric.compute(predictions=preds, references=labels) if len(result) > 1: result["combined_score"] = np.mean(list(result.values())).item() return result diff --git a/examples/translation/run_translation.py b/examples/translation/run_translation.py index c1d8a07d1d..4bf587ad38 100644 --- a/examples/translation/run_translation.py +++ b/examples/translation/run_translation.py @@ -1,5 +1,4 @@ #!/usr/bin/env python -# coding=utf-8 # Copyright The HuggingFace Team and The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -62,8 +61,8 @@ def check_optimum_habana_min_version(*a, **b): logger = logging.getLogger(__name__) # Will error if the minimal version of Transformers and Optimum Habana are not installed. Remove at your own risks. -check_min_version("4.51.0") -check_optimum_habana_min_version("1.18.0.dev0") +check_min_version("4.55.0") +check_optimum_habana_min_version("1.19.0.dev0") require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/translation/requirements.txt") @@ -110,7 +109,7 @@ class ModelArguments: metadata={ "help": ( "The token to use as HTTP bearer authorization for remote files. If not specified, will use the token " - "generated when running `huggingface-cli login` (stored in `~/.huggingface`)." + "generated when running `hf auth login` (stored in `~/.huggingface`)." ) }, ) From d4960e774aba7b524718ec706bf8d6090d32188f Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Mon, 11 Aug 2025 21:11:50 +0000 Subject: [PATCH 15/22] generation + integrations + loss folders --- .../generation/candidate_generator.py | 4 +- .../generation/stopping_criteria.py | 2 +- .../habana/transformers/generation/utils.py | 231 ++++++++++++------ .../transformers/integrations/deepspeed.py | 10 + .../habana/transformers/loss/loss_rt_detr.py | 4 +- 5 files changed, 167 insertions(+), 84 deletions(-) diff --git a/optimum/habana/transformers/generation/candidate_generator.py b/optimum/habana/transformers/generation/candidate_generator.py index 6688553459..699e1fbcb3 100644 --- a/optimum/habana/transformers/generation/candidate_generator.py +++ b/optimum/habana/transformers/generation/candidate_generator.py @@ -1,5 +1,5 @@ import inspect -from typing import TYPE_CHECKING, Dict, Optional +from typing import TYPE_CHECKING, Optional import torch from transformers.generation.candidate_generator import ( @@ -20,7 +20,7 @@ def __init__( input_ids: torch.LongTensor, assistant_model: "PreTrainedModel", generation_config: "GaudiGenerationConfig", - model_kwargs: Dict, + model_kwargs: dict, inputs_tensor: Optional[torch.Tensor] = None, logits_processor: "LogitsProcessorList" = None, ): diff --git a/optimum/habana/transformers/generation/stopping_criteria.py b/optimum/habana/transformers/generation/stopping_criteria.py index 844ffa50f2..c9dc16f2b2 100644 --- a/optimum/habana/transformers/generation/stopping_criteria.py +++ b/optimum/habana/transformers/generation/stopping_criteria.py @@ -43,7 +43,7 @@ def gaudi_MaxLengthCriteria_call( else: return create_return_const_tensor(input_ids, token_idx >= self.max_length) else: - cur_len = input_ids.shape[-1] + cur_len = input_ids.shape[1] is_done = cur_len >= self.max_length if self.max_position_embeddings is not None and not is_done and cur_len >= self.max_position_embeddings: logger.warning_once( diff --git a/optimum/habana/transformers/generation/utils.py b/optimum/habana/transformers/generation/utils.py index a8b1858e99..5697a01e49 100755 --- a/optimum/habana/transformers/generation/utils.py +++ b/optimum/habana/transformers/generation/utils.py @@ -17,7 +17,7 @@ import copy import inspect import math -from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union +from typing import TYPE_CHECKING, Any, Callable, Optional, Union import torch import torch.distributed as dist @@ -27,8 +27,6 @@ DynamicCache, EncoderDecoderCache, OffloadedCache, - QuantizedCacheConfig, - StaticCache, ) from transformers.generation.beam_constraints import DisjunctiveConstraint, PhrasalConstraint from transformers.generation.beam_search import BeamScorer, BeamSearchScorer, ConstrainedBeamSearchScorer @@ -39,7 +37,6 @@ EarlyExitCandidateGenerator, PromptLookupCandidateGenerator, UniversalSpeculativeDecodingGenerator, - _crop_past_key_values, _prepare_attention_mask, _prepare_token_type_ids, ) @@ -69,6 +66,7 @@ ) from transformers.integrations.deepspeed import is_deepspeed_zero3_enabled from transformers.integrations.fsdp import is_fsdp_managed_module +from transformers.masking_utils import create_masks_for_generate from transformers.modeling_outputs import CausalLMOutputWithPast, Seq2SeqLMOutput from transformers.utils import ModelOutput, is_hqq_available, is_optimum_quanto_available @@ -267,32 +265,46 @@ def _prepare_inputs_for_generation( model_input = model_input.clone(memory_format=torch.contiguous_format) model_inputs[model_input_name] = model_input - # 6. Create 4D attention mask is we are using a `StaticCache` (important for performant compiled forward pass) - if isinstance(past_key_values, StaticCache) and attention_mask.ndim == 2: - if model_inputs["inputs_embeds"] is not None: + # 6. Create 4D attention mask is we are using a compilable cache (important for performant compiled forward + # pass) + if ( + isinstance(past_key_values, Cache) + and past_key_values.is_compileable + and attention_mask is not None + and attention_mask.ndim == 2 + ): + if not self.config.is_encoder_decoder and model_inputs["inputs_embeds"] is not None: batch_size, sequence_length, _ = model_inputs["inputs_embeds"].shape - device = model_inputs["inputs_embeds"].device else: - batch_size, sequence_length = model_inputs[input_ids_key].shape - device = model_inputs[input_ids_key].device + batch_size, sequence_length = model_inputs[input_ids_key].shape[:2] # Create the causal mask with fixed shape in advance, to reduce recompilations. If the function to create - # the 4D causal mask exists, it should be present in the base model (XXXModel class). - base_model = getattr(self, self.base_model_prefix, None) - if base_model is None: - causal_mask_creation_function = getattr( - self, "_prepare_4d_causal_attention_mask_with_cache_position", None - ) - else: + # the 4D causal mask exists, it should be present in the base model (XXXModel class) or in its decoder. + base_model = getattr(self, self.base_model_prefix, self) + decoder = base_model.get_decoder() if hasattr(base_model, "get_decoder") else None + causal_mask_creation_function = getattr( + base_model, "_prepare_4d_causal_attention_mask_with_cache_position", None + ) + if causal_mask_creation_function is None and decoder is not None: # it may be in the decoder causal_mask_creation_function = getattr( - base_model, "_prepare_4d_causal_attention_mask_with_cache_position", None + decoder, "_prepare_4d_causal_attention_mask_with_cache_position", None ) - if causal_mask_creation_function is None: - logger.warning_once( - f"{self.__class__.__name__} has no `_prepare_4d_causal_attention_mask_with_cache_position` method " - "defined in its base modeling class. Compiled forward passes will be sub-optimal. If you're " - "writing code, see Llama for an example implementation. If you're a user, please report this " - "issue on GitHub." + + # If it's not defined, it means the model uses the new general mask API + if causal_mask_creation_function is None: # can't be found + token_type_ids = model_inputs.get("token_type_ids", None) + position_ids = model_inputs.get(position_ids_key, None) + # Some models may overwrite the general one + causal_mask_creation_function = getattr(self, "create_masks_for_generate", create_masks_for_generate) + attention_mask = causal_mask_creation_function( + config=self.config, + # we only need batch size, seq_length and dtype here - we don't care about the values of the embeddings + input_embeds=torch.empty((batch_size, sequence_length), dtype=self.dtype), + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, + position_ids=position_ids, + token_type_ids=token_type_ids, ) else: attention_mask = causal_mask_creation_function( @@ -300,7 +312,6 @@ def _prepare_inputs_for_generation( sequence_length=sequence_length, target_length=past_key_values.get_max_cache_shape(), dtype=self.dtype, - device=device, cache_position=cache_position, batch_size=batch_size, config=self.config, @@ -362,12 +373,12 @@ def _prepare_decoder_input_ids_for_generation( self, batch_size: int, model_input_name: str, - model_kwargs: Dict[str, torch.Tensor], + model_kwargs: dict[str, torch.Tensor], decoder_start_token_id: torch.Tensor, device: Optional[torch.device] = None, max_new_tokens: int = None, pad_token_id: int = None, - ) -> Tuple[torch.LongTensor, Dict[str, torch.Tensor]]: + ) -> tuple[torch.LongTensor, dict[str, torch.Tensor]]: """Prepares `decoder_input_ids` for generation with encoder-decoder models""" # 1. Check whether the user has defined `decoder_input_ids` manually. To facilitate in terms of input naming, # we also allow the user to pass it under `input_ids`, if the encoder does not use it as the main input. @@ -472,7 +483,7 @@ def _expand_inputs_for_generation( is_encoder_decoder: bool = False, input_ids: Optional[torch.LongTensor] = None, **model_kwargs, - ) -> Tuple[torch.LongTensor, Dict[str, Any]]: + ) -> tuple[torch.LongTensor, dict[str, Any]]: """ Expands tensors from [batch_size, ...] to [batch_size * expand_size, ...]. @@ -563,10 +574,10 @@ def _remove_past_key_values(self, model_kwargs): def _update_model_kwargs_for_generation( self, outputs: ModelOutput, - model_kwargs: Dict[str, Any], + model_kwargs: dict[str, Any], is_encoder_decoder: bool = False, num_new_tokens: int = 1, - ) -> Dict[str, Any]: + ) -> dict[str, Any]: """ Copied from Transformers: https://github.com/huggingface/transformers/blob/527ab894e59b6582578008e3b47648a65063f73d/src/transformers/generation/utils.py#L745 @@ -735,7 +746,7 @@ def _get_candidate_generator( logits_processor: LogitsProcessorList, target_tokenizer: "PreTrainedTokenizerBase", assistant_tokenizer: "PreTrainedTokenizerBase", - model_kwargs: Dict, + model_kwargs: dict, ) -> CandidateGenerator: different_tokenizers = all(v is not None for v in (assistant_model, target_tokenizer, assistant_tokenizer)) @@ -758,8 +769,14 @@ def _get_candidate_generator( elif different_tokenizers: if generation_config.do_sample is True: atm_translator = AssistantVocabTranslatorCache.get_translator( - target_tokenizer, assistant_tokenizer, self.config.vocab_size, assistant_model.device + target_tokenizer, + assistant_tokenizer, + self.config.get_text_config().vocab_size, + assistant_model=assistant_model, + assistant_prune_lm_head=True, # prune LM head of assistant model ) + # Since we prune the LM head, we cannot use the repetition penalty on the assistant model due to mismatches between token ids and logits index + assistant_model.generation_config.repetition_penalty = None candidate_generator = UniversalSpeculativeDecodingGenerator( input_ids=input_ids, assistant_model=assistant_model, @@ -903,8 +920,8 @@ def _prepare_generation_config( self, generation_config: Optional[GaudiGenerationConfig], use_model_defaults: Optional[bool] = None, - **kwargs: Dict, - ) -> Tuple[GaudiGenerationConfig, Dict]: + **kwargs: dict, + ) -> tuple[GaudiGenerationConfig, dict]: """ Copied from https://github.com/huggingface/transformers/blob/v4.40.2/src/transformers/generation/utils.py#L1230 Differences: @@ -955,16 +972,25 @@ def _prepare_generation_config( use_model_defaults is None and model_base_version >= version.parse("4.50.0") ): modified_values = {} - default_generation_config = GaudiGenerationConfig() - for key, default_value in default_generation_config.__dict__.items(): + global_default_generation_config = GaudiGenerationConfig() + model_generation_config = self.generation_config + # we iterate over the model's generation config: it may hold custom keys, which we'll want to copy + for key, model_gen_config_value in model_generation_config.__dict__.items(): if key.startswith("_") or key == "transformers_version": # metadata continue - custom_gen_config_value = getattr(generation_config, key) - model_gen_config_value = getattr(self.generation_config, key) - if custom_gen_config_value == default_value and model_gen_config_value != default_value: + global_default_value = getattr(global_default_generation_config, key, None) + custom_gen_config_value = getattr(generation_config, key, None) + if ( + custom_gen_config_value == global_default_value + and model_gen_config_value != global_default_value + ): modified_values[key] = model_gen_config_value setattr(generation_config, key, model_gen_config_value) - if len(modified_values) > 0: + # edge case: we may set `temperature=0.0` and `do_sample=False`, but the model defaults to + # `do_sample=True` + if generation_config.temperature == 0.0: + generation_config.do_sample = False + if use_model_defaults is None and len(modified_values) > 0: logger.warning_once( f"`generation_config` default values have been modified to match model-specific defaults: " f"{modified_values}. If this is not desired, please set these values explicitly." @@ -1000,7 +1026,7 @@ def _prepare_generation_config( def _prepare_cache_for_generation( self, generation_config: GaudiGenerationConfig, - model_kwargs: Dict, + model_kwargs: dict, assistant_model: "PreTrainedModel", batch_size: int, max_cache_length: int, @@ -1013,7 +1039,9 @@ def _prepare_cache_for_generation( - change the default from DynamicCache to tuples """ - cache_name = "past_key_values" if "mamba" not in self.__class__.__name__.lower() else "cache_params" + is_hybrid_cache = any(class_name in self.__class__.__name__.lower() for class_name in ["mamba", "falconh1"]) + cache_name = "past_key_values" if not is_hybrid_cache else "cache_params" + requires_cross_attention_cache = ( self.config.is_encoder_decoder or model_kwargs.get("encoder_outputs") is not None ) @@ -1041,7 +1069,7 @@ def _prepare_cache_for_generation( if generation_config.use_cache is False: return - # Quick escape route 3: model that only supports legacy caches = nothing to prepare + # Quick escape route 3: model that only supports legacy caches or models that supply it in `prepare_inputs_for_generation` (mamba, zamba, ...) if not self._supports_default_dynamic_cache(): if generation_config.cache_implementation is not None: warn0( @@ -1064,11 +1092,11 @@ def _prepare_cache_for_generation( generation_config.cache_implementation = None # generation_config.cache_implementation = generation_config.cache_implementation or getattr( - # self.config.get_text_config(), "cache_implementation", None + # self.config.get_text_config(decoder=True), "cache_implementation", None # ) if generation_config.cache_implementation is not None: if generation_config.cache_implementation in NEED_SETUP_CACHE_CLASSES_MAPPING: - if generation_config.cache_implementation == "static" and not self._supports_static_cache: + if generation_config.cache_implementation == "static" and not self._can_compile_fullgraph: raise ValueError( "This model does not support `cache_implementation='static'`. Please check the following " "issue: https://github.com/huggingface/transformers/issues/28981" @@ -1081,7 +1109,7 @@ def _prepare_cache_for_generation( model_kwargs=model_kwargs, ) elif generation_config.cache_implementation == "quantized": - if not self._supports_quantized_cache: + if self.config.is_encoder_decoder or not self._supports_default_dynamic_cache(): raise ValueError( "This model does not support the quantized cache. If you want your model to support quantized " "cache, please open an issue and tag @zucchini-nlp." @@ -1090,22 +1118,22 @@ def _prepare_cache_for_generation( cache_config = ( generation_config.cache_config if generation_config.cache_config is not None - else QuantizedCacheConfig() + else {"backend": "quanto"} ) - cache_class = QUANT_BACKEND_CLASSES_MAPPING[cache_config.backend] + cache_class = QUANT_BACKEND_CLASSES_MAPPING[cache_config["backend"]] - if cache_config.backend == "quanto" and not is_optimum_quanto_available(): + if cache_config["backend"] == "quanto" and not is_optimum_quanto_available(): raise ImportError( "You need to install optimum-quanto in order to use KV cache quantization with optimum-quanto backend. " "Please install it via with `pip install optimum-quanto`" ) - elif cache_config.backend == "HQQ" and not is_hqq_available(): + elif cache_config["backend"] == "HQQ" and not is_hqq_available(): raise ImportError( "You need to install `HQQ` in order to use KV cache quantization with HQQ backend. " "Please install it via with `pip install hqq`" ) - model_kwargs[cache_name] = cache_class(cache_config) + model_kwargs[cache_name] = cache_class(**cache_config) elif generation_config.cache_implementation == "offloaded": model_kwargs[cache_name] = OffloadedCache() elif generation_config.cache_implementation == "dynamic": @@ -1122,16 +1150,17 @@ def generate( generation_config: Optional[GaudiGenerationConfig] = None, logits_processor: Optional[LogitsProcessorList] = None, stopping_criteria: Optional[StoppingCriteriaList] = None, - prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], List[int]]] = None, + prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], list[int]]] = None, synced_gpus: Optional[bool] = None, assistant_model: Optional["PreTrainedModel"] = None, streamer: Optional["BaseStreamer"] = None, negative_prompt_ids: Optional[torch.Tensor] = None, negative_prompt_attention_mask: Optional[torch.Tensor] = None, use_model_defaults: Optional[bool] = None, + custom_generate: Optional[str] = None, lazy_mode: Optional[bool] = False, hpu_graphs: Optional[bool] = False, - iteration_times: Optional[List[float]] = None, + iteration_times: Optional[list[float]] = None, profiler: Optional[HabanaProfile] = None, **kwargs, ) -> Union[GenerateOutput, torch.LongTensor]: @@ -1174,13 +1203,13 @@ def generate( generation config an error is thrown. If your stopping criteria depends on the `scores` input, make sure you pass `return_dict_in_generate=True, output_scores=True` to `generate`. This feature is intended for advanced users. - prefix_allowed_tokens_fn (`Callable[[int, torch.Tensor], List[int]]`, *optional*): + prefix_allowed_tokens_fn (`Callable[[int, torch.Tensor], list[int]]`, *optional*): If provided, this function constraints the beam search to allowed tokens only at each step. If not provided no constraint is applied. This function takes 2 arguments: the batch ID `batch_id` and `input_ids`. It has to return a list with the allowed tokens for the next generation step conditioned on the batch ID `batch_id` and the previously generated tokens `inputs_ids`. This argument is useful for constrained generation conditioned on the prefix, as described in [Autoregressive Entity - Retrieval](https://arxiv.org/abs/2010.00904). + Retrieval](https://huggingface.co/papers/2010.00904). synced_gpus (`bool`, *optional*): Whether to continue running the while loop until max_length. Unless overridden, this flag will be set to `True` if using `FullyShardedDataParallel` or DeepSpeed ZeRO Stage 3 with multiple GPUs to avoid @@ -1203,13 +1232,18 @@ def generate( generation configuration (`model.generation_config`), as opposed to the global defaults (`GenerationConfig()`). If unset, models saved starting from `v4.50` will consider this flag to be `True`. + custom_generate (`str`, *optional*): + A string containing the name of a huggingface.co repository. If provided, the custom `generate` + function defined in that reposity's `custom_generate/generate.py` file will be executed instead of the + standard `generate` method. Note that the logic is for generation is entirely defined in that + repository, and the return type may be different from the standard `generate` method. lazy_mode (`bool`, *optional*, defaults to `False`): Whether the run is executed in lazy mode or not (i.e. eager mode). hpu_graphs (`bool`, *optional*, defaults to `False`): Whether to use HPU graphs for inference. profiler (`HabanaProfile`, *optional*, defaults to None): HabanaProfile object to use for profiling. - kwargs (`Dict[str, Any]`, *optional*): + kwargs (`dict[str, Any]`, *optional*): Ad hoc parametrization of `generation_config` and/or additional model-specific kwargs that will be forwarded to the `forward` function of the model. If the model is an encoder-decoder model, encoder specific kwargs should not be prefixed and decoder specific kwargs should be prefixed with *decoder_*. @@ -1238,8 +1272,28 @@ def generate( else: synced_gpus = False + # 0. If requested, load an arbitrary generation recipe from the Hub and run it instead + trust_remote_code = kwargs.pop("trust_remote_code", None) + if custom_generate is not None: + # Get all `generate` arguments in a single variable. Custom functions are responsible for handling them: + # they receive the same inputs as `generate`, with `model` instead of `self` and excluding the arguments to + # trigger the custom generation. They can access to methods from `GenerationMixin` through `model`. + global_keys_to_exclude = { + "self", + "kwargs", + "global_keys_to_exclude", + "trust_remote_code", + "custom_generate", + } + generate_arguments = {key: value for key, value in locals().items() if key not in global_keys_to_exclude} + generate_arguments.update(kwargs) + + custom_generate_function = self.load_custom_generate( + custom_generate, trust_remote_code=trust_remote_code, **kwargs + ) + return custom_generate_function(model=self, **generate_arguments) + # 1. Handle `generation_config` and kwargs that might update it, and validate the `.generate()` call - self._validate_model_class() tokenizer = kwargs.pop("tokenizer", None) # Pull this out first, we only use it for stopping criteria assistant_tokenizer = kwargs.pop("assistant_tokenizer", None) # only used for assisted generation if hpu_graphs and not lazy_mode: @@ -1681,6 +1735,11 @@ def generate( **model_kwargs, ) elif generation_mode == GenerationMode.DOLA_GENERATION: + if not trust_remote_code: + logger.warning_once( + "DoLa Decoding is scheduled to be moved to a `custom_generate` repository in v4.55.0. " + "To prevent loss of backward compatibility, add `trust_remote_code=True` to your `generate` call." + ) if self._is_stateful: # DoLa decoding was not designed for stateful models, and would require some changes raise ValueError( @@ -1698,6 +1757,11 @@ def generate( ) elif generation_mode == GenerationMode.CONTRASTIVE_SEARCH: + if not trust_remote_code: + logger.warning_once( + "Contrastive Search is scheduled to be moved to a `custom_generate` repository in v4.55.0. " + "To prevent loss of backward compatibility, add `trust_remote_code=True` to your `generate` call." + ) if not model_kwargs["use_cache"]: raise ValueError("Contrastive search requires `use_cache=True`") if self._is_stateful: @@ -1779,6 +1843,10 @@ def generate( ) elif generation_mode == GenerationMode.GROUP_BEAM_SEARCH: + logger.warning_once( + "Group Beam Search is scheduled to be moved to a `custom_generate` repository in v4.55.0. " + "To prevent loss of backward compatibility, add `trust_remote_code=True` to your `generate` call." + ) # 11. prepare beam search scorer beam_scorer = BeamSearchScorer( batch_size=batch_size, @@ -1812,6 +1880,10 @@ def generate( ) elif generation_mode == GenerationMode.CONSTRAINED_BEAM_SEARCH: + logger.warning_once( + "Constrained Beam Search is scheduled to be moved to a `custom_generate` repository in v4.55.0. " + "To prevent loss of backward compatibility, add `trust_remote_code=True` to your `generate` call." + ) final_constraints = [] if generation_config.constraints is not None: final_constraints = generation_config.constraints @@ -1820,7 +1892,7 @@ def generate( def typeerror(): raise ValueError( - "`force_words_ids` has to either be a `List[List[List[int]]]` or `List[List[int]]` " + "`force_words_ids` has to either be a `list[list[list[int]]]` or `list[list[int]]` " f"of positive integers, but is {generation_config.force_words_ids}." ) @@ -1897,7 +1969,7 @@ def typeerror(): def _dola_decoding( self, input_ids: torch.LongTensor, - dola_layers: Union[str, List[int]], + dola_layers: Union[str, list[int]], logits_processor: LogitsProcessorList, stopping_criteria: StoppingCriteriaList, generation_config: GaudiGenerationConfig, @@ -1909,12 +1981,12 @@ def _dola_decoding( Generates sequences of token ids for models with a language modeling head using **dola decoding** and can be used for decoder-only text models. The method is based on the paper "DoLa: Decoding by Contrasting Layers Improves Factuality in Large Language - Models" (https://arxiv.org/abs/2309.03883) in ICLR 2024. + Models" (https://huggingface.co/papers/2309.03883) in ICLR 2024. Parameters: input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): The sequence used as a prompt for the generation. - dola_layers (`Union[str, List[int]]`): + dola_layers (`Union[str, list[int]]`): The candidate layers used in contrasting layers of DoLa. It can be either 1) 'low' or 'high', which means the lower part or higher part of the model layers, respectively, or 2) a list of layer indices to be used for candidate layers. The 0-th layer is the word embedding layer of the model. @@ -2033,7 +2105,7 @@ def _contrastive_search( ) # keep track of which sequences are already finished - batch_size, cur_len = input_ids.shape + batch_size, cur_len = input_ids.shape[:2] if not ignore_eos: unfinished_sequences = torch.ones(batch_size, dtype=torch.long, device=input_ids.device) @@ -2627,11 +2699,11 @@ def _sample( ) # keep track of which sequences are already finished - batch_size, cur_len = input_ids.shape + batch_size, cur_len = input_ids.shape[:2] this_peer_finished = False if not ignore_eos: unfinished_sequences = torch.ones(batch_size, dtype=torch.long, device=input_ids.device) - model_kwargs = self._get_initial_cache_position(input_ids, model_kwargs) + model_kwargs = self._get_initial_cache_position(cur_len, input_ids.device, model_kwargs) bucket_size = model_kwargs.get("bucket_size", -1) prev_idx = -1 # avoiding calculate cache_idx when its value is not changing @@ -2983,7 +3055,7 @@ def _beam_search( batch_size = len(beam_scorer._beam_hyps) num_beams = beam_scorer.num_beams - batch_beam_size, cur_len = input_ids.shape + batch_beam_size, cur_len = input_ids.shape[:2] if "inputs_embeds" in model_kwargs: cur_len = model_kwargs["inputs_embeds"].shape[1] token_idx = model_kwargs.get("token_idx", None) @@ -2998,7 +3070,7 @@ def _beam_search( f"Batch dimension of `input_ids` should be {num_beams * batch_size}, but is {batch_beam_size}." ) - # (joao) feature lost in the refactor. Probably won't implement, hurts readbility with minimal gains (there + # (joao) feature lost in the refactor. Probably won't implement, hurts readability with minimal gains (there # are newer low-memory alternatives like the offloaded cache) sequential = generation_config.low_memory if sequential: @@ -3592,7 +3664,7 @@ def _constrained_beam_search( batch_size = len(constrained_beam_scorer._beam_hyps) num_beams = constrained_beam_scorer.num_beams - batch_beam_size, cur_len = input_ids.shape + batch_beam_size, cur_len = input_ids.shape[:2] token_idx = model_kwargs.get("token_idx", None) if token_idx is not None: @@ -3635,7 +3707,7 @@ def _constrained_beam_search( if token_idx is not None: decoder_prompt_len = cur_len else: - decoder_prompt_len = input_ids.shape[-1] + decoder_prompt_len = input_ids.shape[1] if profiler is not None: profiler.start() @@ -3757,13 +3829,15 @@ def _constrained_beam_search( # (that way the memory peak does not include outputs.logits) del outputs + # NOTE: we need to check if `self._reorder_cache` exists for special models like RAG, RecurrentGemma etc. if model_kwargs.get("past_key_values", None) is not None: - model_kwargs["past_key_values"] = self._temporary_reorder_cache( - model_kwargs["past_key_values"], beam_idx - ) + if hasattr(self, "_reorder_cache"): + model_kwargs["past_key_values"] = self._reorder_cache(model_kwargs["past_key_values"], beam_idx) + else: + model_kwargs["past_key_values"].reorder_cache(beam_idx) if return_dict_in_generate and output_scores: - beam_indices = tuple((beam_indices[beam_idx[i]] + (beam_idx[i],) for i in range(len(beam_indices)))) + beam_indices = tuple(beam_indices[beam_idx[i]] + (beam_idx[i],) for i in range(len(beam_indices))) # increase cur_len cur_len = cur_len + 1 @@ -3909,10 +3983,10 @@ def _assisted_decoding( ) # keep track of which sequences are already finished - batch_size, cur_len = input_ids.shape + batch_size, cur_len = input_ids.shape[:2] if not ignore_eos: unfinished_sequences = torch.ones(batch_size, dtype=torch.long, device=input_ids.device) - model_kwargs = self._get_initial_cache_position(input_ids, model_kwargs) + model_kwargs = self._get_initial_cache_position(cur_len, input_ids.device, model_kwargs) if profiler is not None: profiler.start() @@ -3930,7 +4004,7 @@ def _assisted_decoding( # Update cur_len in case of static shapes cur_len = (token_idx + model_kwargs.get("inputs_embeds_offset", 0)).item() else: - cur_len = input_ids.shape[-1] + cur_len = input_ids.shape[1] # prepare model inputs model_kwargs["lazy_mode"] = lazy_mode @@ -3995,7 +4069,7 @@ def _assisted_decoding( # 3. Select the accepted tokens. There are two possible cases: # Case 1: `do_sample=True` and we have logits for the candidates (originally from speculative decoding) - # 👉 Apply algorithm 1 from the speculative decoding paper (https://arxiv.org/pdf/2211.17192.pdf). + # 👉 Apply algorithm 1 from the speculative decoding paper (https://huggingface.co/papers/2211.17192). if do_sample and candidate_logits is not None: from transformers.generation.utils import _speculative_sampling @@ -4037,11 +4111,10 @@ def _assisted_decoding( input_ids = torch.cat((input_ids, valid_tokens), dim=-1) if streamer is not None: streamer.put(valid_tokens.cpu()) - new_cur_len = input_ids.shape[-1] + new_cur_len = input_ids.shape[1] # 4.2. Discard past key values relative to unused assistant tokens - new_cache_size = new_cur_len - 1 - outputs.past_key_values = _crop_past_key_values(self, outputs.past_key_values, new_cache_size) + outputs.past_key_values.crop(new_cur_len - 1) # 5. Update the candidate generation strategy if needed candidate_generator.update_candidate_strategy(input_ids, new_logits, n_matches) diff --git a/optimum/habana/transformers/integrations/deepspeed.py b/optimum/habana/transformers/integrations/deepspeed.py index 716d7cd141..c6f65aa89c 100644 --- a/optimum/habana/transformers/integrations/deepspeed.py +++ b/optimum/habana/transformers/integrations/deepspeed.py @@ -136,6 +136,16 @@ def deepspeed_init(trainer, num_training_steps, inference=False): model_parameters = None else: trainer.optimizer = None # important for when deepspeed_init is used as re-init + deepspeed_tp_size = hf_deepspeed_config.config.get("tensor_parallel", {}).get("autotp_size", 1) + if deepspeed_tp_size > 1: + import deepspeed + + model = deepspeed.tp_model_init( + model=model, + tp_size=deepspeed_tp_size, + dtype=hf_deepspeed_config.dtype(), + config=hf_deepspeed_config.config, + ) model_parameters = list(filter(lambda p: p.requires_grad, model.parameters())) optimizer, lr_scheduler = deepspeed_optim_sched( trainer, hf_deepspeed_config, args, num_training_steps, model_parameters diff --git a/optimum/habana/transformers/loss/loss_rt_detr.py b/optimum/habana/transformers/loss/loss_rt_detr.py index 50cdc3a3b7..b119cd8de6 100644 --- a/optimum/habana/transformers/loss/loss_rt_detr.py +++ b/optimum/habana/transformers/loss/loss_rt_detr.py @@ -54,7 +54,7 @@ def gaudi_RTDetrHungarianMatcher_forward(self, outputs, targets): target_bbox = torch.cat([v["boxes"] for v in targets]) # Compute the classification cost. Contrary to the loss, we don't use the NLL, # but approximate it in 1 - proba[target class]. - # The 1 is a constant that doesn't change the matching, it can be ommitted. + # The 1 is a constant that doesn't change the matching, it can be omitted. if self.use_focal_loss: out_prob = F.sigmoid(outputs["logits"].flatten(0, 1)) out_prob = out_prob[:, target_ids] @@ -67,7 +67,7 @@ def gaudi_RTDetrHungarianMatcher_forward(self, outputs, targets): # Compute the L1 cost between boxes bbox_cost = torch.cdist(out_bbox, target_bbox, p=1) - # Compute the giou cost betwen boxes + # Compute the giou cost between boxes giou_cost = -generalized_box_iou(center_to_corners_format(out_bbox), center_to_corners_format(target_bbox)) # Compute the final cost matrix cost_matrix = self.bbox_cost * bbox_cost + self.class_cost * class_cost + self.giou_cost * giou_cost From 91e29c424246e85b6c2ef2489911256a91bec6f5 Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Wed, 13 Aug 2025 16:44:14 +0000 Subject: [PATCH 16/22] models --- .../models/albert/modeling_albert.py | 4 +- .../transformers/models/bart/modeling_bart.py | 14 +- .../models/blip/modeling_blip_text.py | 11 +- .../models/bloom/modeling_bloom.py | 97 +++------ .../models/codegen/modeling_codegen.py | 70 +++---- .../models/cohere/modeling_cohere.py | 107 +++------- .../models/esm/modeling_esmfold.py | 1 + .../models/falcon/modeling_falcon.py | 77 +++---- .../falcon_mamba/modeling_falcon_mamba.py | 74 +++---- .../models/gemma/modeling_gemma.py | 122 +++-------- .../models/gemma2/modeling_gemma2.py | 90 ++++----- .../models/gpt_neo/modeling_gpt_neo.py | 105 ++++------ .../models/gpt_neox/modeling_gpt_neox.py | 46 ++--- .../transformers/models/gptj/modeling_gptj.py | 66 +++--- .../idefics2/image_processing_idefics2.py | 4 +- .../models/idefics2/modeling_idefics2.py | 87 ++++---- .../models/llama/modeling_llama.py | 190 ++++++------------ .../models/llava/modeling_llava.py | 19 +- .../models/llava_next/modeling_llava_next.py | 27 +-- .../modeling_llava_onevision.py | 11 +- .../models/mamba/modeling_mamba.py | 6 +- .../models/mistral/modeling_mistral.py | 119 +++-------- .../models/mixtral/modeling_mixtral.py | 126 +++--------- .../models/mllama/modeling_mllama.py | 38 ++-- .../models/modeling_all_models.py | 4 +- .../transformers/models/mpt/modeling_mpt.py | 76 +++---- .../transformers/models/opt/modeling_opt.py | 22 +- .../models/owlvit/modeling_owlvit.py | 4 +- .../transformers/models/phi/modeling_phi.py | 70 +++---- .../models/qwen2/modeling_qwen2.py | 163 ++++++--------- .../models/qwen2_moe/modeling_qwen2_moe.py | 94 ++++----- .../models/qwen3/modeling_qwen3.py | 165 ++++++--------- .../models/qwen3_moe/modeling_qwen3_moe.py | 166 +++++---------- .../seamless_m4t/modeling_seamless_m4t.py | 87 ++++---- .../models/siglip/modeling_siglip.py | 36 ++-- .../models/speecht5/modeling_speecht5.py | 64 +++--- .../models/stablelm/modeling_stablelm.py | 50 ++--- .../models/starcoder2/modeling_starcoder2.py | 128 ++++-------- .../transformers/models/t5/modeling_t5.py | 59 ++---- .../video_llava/modeling_video_llava.py | 16 +- .../video_llava/processing_video_llava.py | 18 +- .../transformers/models/vit/modeling_vit.py | 28 ++- .../models/wav2vec2/modeling_wav2vec2.py | 46 ++--- .../transformers/models/xglm/modeling_xglm.py | 68 +++---- .../xlm_roberta/modeling_xlm_roberta.py | 21 +- 45 files changed, 1038 insertions(+), 1858 deletions(-) diff --git a/optimum/habana/transformers/models/albert/modeling_albert.py b/optimum/habana/transformers/models/albert/modeling_albert.py index 6ac9b80073..7cae39c145 100644 --- a/optimum/habana/transformers/models/albert/modeling_albert.py +++ b/optimum/habana/transformers/models/albert/modeling_albert.py @@ -14,7 +14,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Optional, Tuple, Union +from typing import Optional, Union import torch from transformers.modeling_outputs import BaseModelOutputWithPooling @@ -31,7 +31,7 @@ def gaudi_albert_forward( output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, -) -> Union[BaseModelOutputWithPooling, Tuple]: +) -> Union[BaseModelOutputWithPooling, tuple]: """ Same as https://github.com/huggingface/transformers/blob/a9eee2ffecc874df7dd635b2c6abb246fdb318cc/src/transformers/models/albert/modeling_albert.py#L689 except that mixed precision is disabled for computing: diff --git a/optimum/habana/transformers/models/bart/modeling_bart.py b/optimum/habana/transformers/models/bart/modeling_bart.py index 16d1bbd0e7..6dbb19b0e5 100644 --- a/optimum/habana/transformers/models/bart/modeling_bart.py +++ b/optimum/habana/transformers/models/bart/modeling_bart.py @@ -58,7 +58,7 @@ def __init__(self, num_embeddings: int, embedding_dim: int): self.offset = 2 super().__init__(num_embeddings + self.offset, embedding_dim) - def forward(self, input_ids: torch.Tensor, past_key_values_length: torch.Tensor = torch.tensor(0)): + def forward(self, input_ids: torch.Tensor, past_key_values_length: torch.Tensor = torch.tensor(0), position_ids: torch.Tensor = None): """`input_ids' shape is expected to be [bsz x seqlen].""" bsz, seq_len = input_ids.shape[:2] @@ -72,19 +72,25 @@ def gaudi_BartAttention_forward( self, hidden_states: torch.Tensor, key_value_states: Optional[torch.Tensor] = None, - past_key_value: Optional[Tuple[torch.Tensor]] = None, + past_key_value: Optional[tuple[torch.Tensor]] = None, attention_mask: Optional[torch.Tensor] = None, layer_head_mask: Optional[torch.Tensor] = None, output_attentions: bool = False, + cache_position: Optional[torch.Tensor] = None, token_idx: Optional[torch.Tensor] = None, -) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + # TODO: we need a refactor so that the different attention modules can get their specific kwargs + # ATM, we have mixed things encoder, decoder, and encoder-decoder attn + **kwargs, +) -> 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() + # determine input shapes + bsz, tgt_len = hidden_states.shape[:-1] + src_len = key_value_states.shape[1] if is_cross_attention else tgt_len # get query proj query_states = self.q_proj(hidden_states) * self.scaling diff --git a/optimum/habana/transformers/models/blip/modeling_blip_text.py b/optimum/habana/transformers/models/blip/modeling_blip_text.py index ef64437659..960a4db17d 100644 --- a/optimum/habana/transformers/models/blip/modeling_blip_text.py +++ b/optimum/habana/transformers/models/blip/modeling_blip_text.py @@ -25,19 +25,26 @@ def gaudi_BlipTextSelfAttention_forward( encoder_attention_mask: Optional[torch.FloatTensor] = None, past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, output_attentions: Optional[bool] = False, + cache_position: Optional[torch.Tensor] = None, token_idx: Optional[torch.Tensor] = None, -) -> Tuple[torch.Tensor]: +) -> tuple[torch.Tensor]: """ Copied from BlipTextSelfAttention.forward: https://github.com/huggingface/transformers/blob/v4.37.2/src/transformers/models/blip/modeling_blip_text.py#L143 The only differences are: - add token_idx """ - mixed_query_layer = self.query(hidden_states) + batch_size, seq_length, _ = hidden_states.shape + query_layer = ( + self.query(hidden_states) + .view(batch_size, -1, self.num_attention_heads, self.attention_head_size) + .transpose(1, 2) + ) # If this is instantiated as a cross-attention module, the keys # and values come from an encoder; the attention mask needs to be # such that the encoder's padding tokens are not attended to. is_cross_attention = encoder_hidden_states is not None + attention_mask = encoder_attention_mask if is_cross_attention else attention_mask if is_cross_attention: key_layer = self.transpose_for_scores(self.key(encoder_hidden_states)) diff --git a/optimum/habana/transformers/models/bloom/modeling_bloom.py b/optimum/habana/transformers/models/bloom/modeling_bloom.py index 85802156bf..51a877a13a 100644 --- a/optimum/habana/transformers/models/bloom/modeling_bloom.py +++ b/optimum/habana/transformers/models/bloom/modeling_bloom.py @@ -17,7 +17,7 @@ ############################################################################### import math import os -from typing import Optional, Tuple, Union +from typing import Optional, Union import torch from torch.nn import functional as F @@ -37,7 +37,7 @@ def gaudi_bloom_build_alibi_tensor( attention_mask: torch.Tensor, num_heads: int, dtype: torch.dtype, training: bool ) -> torch.Tensor: """ - Link to paper: https://arxiv.org/abs/2108.12409 Alibi tensor is not causal as the original paper mentions, it + Link to paper: https://huggingface.co/papers/2108.12409 Alibi tensor is not causal as the original paper mentions, it relies on a translation invariance of softmax for quick implementation: with l being a tensor, and a fixed value `softmax(l+a) = softmax(l)`. Based on https://github.com/ofirpress/attention_with_linear_biases/blob/a35aaca144e0eb6b789dfcb46784c4b8e31b7983/fairseq/models/transformer.py#L742 @@ -204,10 +204,7 @@ def gaudi_bloom_attention_forward( output_tensor = self.dense(context_layer) output_tensor = dropout_add(output_tensor, residual, self.hidden_dropout, self.training) - outputs = (output_tensor, present) - if output_attentions: - outputs += (attention_probs,) return outputs @@ -242,7 +239,7 @@ def gaudi_bloom_block_forward( residual = hidden_states # Self attention. - attn_outputs = self.self_attention( + attention_output, attn_weights = self.self_attention( layernorm_output, residual, layer_past=layer_past, @@ -255,10 +252,6 @@ def gaudi_bloom_block_forward( token_idx=token_idx, ) - attention_output = attn_outputs[0] - - outputs = attn_outputs[1:] - layernorm_output = self.post_attention_layernorm(attention_output) # Get residual @@ -270,17 +263,12 @@ def gaudi_bloom_block_forward( # MLP. output = self.mlp(layernorm_output, residual) - if use_cache: - outputs = (output,) + outputs - else: - outputs = (output,) + outputs[1:] - - return outputs # hidden_states, present, attentions + return output, attn_weights # hidden_states, present, attentions def gaudi_bloom_convert_to_standard_cache( - self, past_key_value: Tuple[Tuple[torch.Tensor, torch.Tensor]], batch_size: int, training: bool -) -> Tuple[Tuple[torch.Tensor, torch.Tensor]]: + self, past_key_value: tuple[tuple[torch.Tensor, torch.Tensor]], batch_size: int, training: bool +) -> tuple[tuple[torch.Tensor, torch.Tensor]]: """ Standardizes the format of the cache so as to match most implementations, i.e. to tuple(tuple([batch_size, num_heads, ...])) @@ -305,8 +293,8 @@ def gaudi_bloom_convert_to_standard_cache( def gaudi_bloom_convert_to_bloom_cache( - self, past_key_value: Tuple[Tuple[torch.Tensor, torch.Tensor]] -) -> Tuple[Tuple[torch.Tensor, torch.Tensor]]: + self, past_key_value: tuple[tuple[torch.Tensor, torch.Tensor]] +) -> tuple[tuple[torch.Tensor, torch.Tensor]]: """ Converts the cache to the format expected by Bloom, i.e. to tuple(tuple([batch_size * num_heads, ...])) """ @@ -326,7 +314,7 @@ def gaudi_bloom_convert_to_bloom_cache( def gaudi_bloom_model_forward( self, input_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[Union[Cache, Tuple[Tuple[torch.Tensor, torch.Tensor], ...]]] = None, + past_key_values: Optional[Union[Cache, tuple[tuple[torch.Tensor, torch.Tensor], ...]]] = None, attention_mask: Optional[torch.Tensor] = None, head_mask: Optional[torch.LongTensor] = None, inputs_embeds: Optional[torch.LongTensor] = None, @@ -337,7 +325,7 @@ def gaudi_bloom_model_forward( cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, **deprecated_arguments, -) -> Union[Tuple[torch.Tensor, ...], BaseModelOutputWithPastAndCrossAttentions]: +) -> Union[tuple[torch.Tensor, ...], BaseModelOutputWithPastAndCrossAttentions]: if deprecated_arguments.pop("position_ids", False) is not False: # `position_ids` could have been `torch.Tensor` or `None` so defaulting pop to `False` allows to detect if users were passing explicitly `None` warn0( @@ -421,31 +409,17 @@ def gaudi_bloom_model_forward( if output_hidden_states: all_hidden_states = all_hidden_states + (hidden_states,) - if self.gradient_checkpointing and self.training: - outputs = self._gradient_checkpointing_func( - block.__call__, - hidden_states, - alibi, - causal_mask, - layer_past, - head_mask[i], - use_cache, - output_attentions, - cache_position, - None, - ) - else: - outputs = block( - hidden_states, - layer_past=layer_past, - attention_mask=causal_mask, - head_mask=head_mask[i], - use_cache=use_cache, - output_attentions=output_attentions, - alibi=alibi, - cache_position=cache_position, - token_idx=token_idx, - ) + outputs = block( + hidden_states, + layer_past=layer_past, + attention_mask=causal_mask, + head_mask=head_mask[i], + use_cache=use_cache, + output_attentions=output_attentions, + alibi=alibi, + cache_position=cache_position, + token_idx=token_idx, + ) hidden_states = outputs[0] if use_cache is True: @@ -525,7 +499,7 @@ def prepare_inputs_for_generation( def forward( self, input_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[Union[Cache, Tuple[Tuple[torch.Tensor, torch.Tensor], ...]]] = None, + past_key_values: Optional[Union[Cache, tuple[tuple[torch.Tensor, torch.Tensor], ...]]] = None, attention_mask: Optional[torch.Tensor] = None, head_mask: Optional[torch.Tensor] = None, inputs_embeds: Optional[torch.Tensor] = None, @@ -537,7 +511,7 @@ def forward( cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, **deprecated_arguments, - ) -> Union[Tuple[torch.Tensor], CausalLMOutputWithCrossAttentions]: + ) -> Union[tuple[torch.Tensor], CausalLMOutputWithCrossAttentions]: r""" labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set @@ -599,28 +573,3 @@ def forward( hidden_states=transformer_outputs.hidden_states, attentions=transformer_outputs.attentions, ) - - def _reorder_cache( - self, past: Tuple[Tuple[torch.Tensor, torch.Tensor], ...], beam_idx: torch.LongTensor - ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], ...]: - """ - This function is used to re-order the `past_key_values` cache if [`~PreTrainedModel.beam_search`] or - [`~PreTrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct - beam_idx at every generation step. - - Output shares the same memory storage as `past`. - """ - standardized_past = self._convert_to_standard_cache(past, batch_size=len(beam_idx), training=self.training) - - # Get a copy of `beam_idx` on all the devices where we need those indices. - device_to_beam_idx = { - past_state.device: beam_idx.to(past_state.device) for layer_past in past for past_state in layer_past - } - reordered_past = tuple( - ( - layer_past[0].index_select(0, device_to_beam_idx[layer_past[0].device]), - layer_past[1].index_select(0, device_to_beam_idx[layer_past[0].device]), - ) - for layer_past in standardized_past - ) - return self._convert_to_bloom_cache(reordered_past) diff --git a/optimum/habana/transformers/models/codegen/modeling_codegen.py b/optimum/habana/transformers/models/codegen/modeling_codegen.py index 963cead407..0e5fbb4a27 100644 --- a/optimum/habana/transformers/models/codegen/modeling_codegen.py +++ b/optimum/habana/transformers/models/codegen/modeling_codegen.py @@ -1,4 +1,4 @@ -from typing import Optional, Tuple, Union +from typing import Optional, Union import torch import torch.utils.checkpoint @@ -25,8 +25,8 @@ def forward( cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, ) -> Union[ - Tuple[torch.Tensor, Tuple[torch.Tensor]], - Optional[Tuple[torch.Tensor, Tuple[torch.Tensor], Tuple[torch.Tensor, ...]]], + tuple[torch.Tensor, tuple[torch.Tensor]], + Optional[tuple[torch.Tensor, tuple[torch.Tensor], tuple[torch.Tensor, ...]]], ]: """ Copied from CodeGenAttention.forward: https://github.com/huggingface/transformers/blob/main/src/transformers/models/codegen/modeling_codegen.py @@ -98,10 +98,8 @@ def forward( attn_output = self.resid_dropout(attn_output) outputs = (attn_output, present) - if output_attentions: - outputs += (attn_weights,) - return outputs # a, present, (attentions) + return outputs, attn_weights def gaudi_codegen_block_forward( @@ -115,7 +113,7 @@ def gaudi_codegen_block_forward( output_attentions: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, -) -> Union[Tuple[torch.Tensor], Optional[Tuple[torch.Tensor, Tuple[torch.FloatTensor, ...]]]]: +) -> Union[tuple[torch.Tensor], Optional[tuple[torch.Tensor, tuple[torch.FloatTensor, ...]]]]: """ Copied from CodeGenBlock.forward: https://github.com/huggingface/transformers/blob/main/src/transformers/models/codegen/modeling_codegen.py The only differences are: @@ -123,7 +121,7 @@ def gaudi_codegen_block_forward( """ residual = hidden_states hidden_states = self.ln_1(hidden_states) - attn_outputs = self.attn( + attn_outputs, attn_weights = self.attn( hidden_states=hidden_states, layer_past=layer_past, attention_mask=attention_mask, @@ -134,24 +132,16 @@ def gaudi_codegen_block_forward( cache_position=cache_position, token_idx=token_idx, ) - attn_output = attn_outputs[0] # output_attn: a, present, (attentions) - outputs = attn_outputs[1:] - feed_forward_hidden_states = self.mlp(hidden_states) - hidden_states = attn_output + feed_forward_hidden_states + residual - - if use_cache: - outputs = (hidden_states,) + outputs - else: - outputs = (hidden_states,) + outputs[1:] + hidden_states = attn_outputs + feed_forward_hidden_states + residual - return outputs # hidden_states, present, (attentions) + return hidden_states, attn_weights def gaudi_codegen_model_forward( self, input_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[Union[Cache, Tuple[Tuple[torch.Tensor]]]] = None, + past_key_values: Optional[Union[Cache, tuple[tuple[torch.Tensor]]]] = None, attention_mask: Optional[torch.FloatTensor] = None, token_type_ids: Optional[torch.LongTensor] = None, position_ids: Optional[torch.LongTensor] = None, @@ -164,7 +154,7 @@ def gaudi_codegen_model_forward( cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, **kwargs, # NOOP kwargs, for now -) -> Union[Tuple, BaseModelOutputWithPast]: +) -> Union[tuple, BaseModelOutputWithPast]: """ Copied from CodeGenBlock.forward: https://github.com/huggingface/transformers/blob/main/src/transformers/models/codegen/modeling_codegen.py The only differences are: @@ -261,31 +251,17 @@ def gaudi_codegen_model_forward( if output_hidden_states: all_hidden_states = all_hidden_states + (hidden_states,) - if self.gradient_checkpointing and self.training: - outputs = self._gradient_checkpointing_func( - block.__call__, - hidden_states, - None, - attention_mask, - position_ids, - head_mask[i], - use_cache, - output_attentions, - cache_position, - None, - ) - else: - outputs = block( - hidden_states=hidden_states, - layer_past=layer_past, - attention_mask=attention_mask, - position_ids=position_ids, - head_mask=head_mask[i], - use_cache=use_cache, - output_attentions=output_attentions, - cache_position=cache_position, - token_idx=token_idx, - ) + outputs = block( + hidden_states, + layer_past=layer_past, + attention_mask=attention_mask, + position_ids=position_ids, + head_mask=head_mask[i], + use_cache=use_cache, + output_attentions=output_attentions, + cache_position=cache_position, + token_idx=token_idx, + ) hidden_states = outputs[0] if use_cache is True: @@ -384,7 +360,7 @@ def prepare_inputs_for_generation( def forward( self, input_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[Union[Cache, Tuple[Tuple[torch.Tensor]]]] = None, + past_key_values: Optional[Union[Cache, tuple[tuple[torch.Tensor]]]] = None, attention_mask: Optional[torch.FloatTensor] = None, token_type_ids: Optional[torch.LongTensor] = None, position_ids: Optional[torch.LongTensor] = None, @@ -398,7 +374,7 @@ def forward( cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, **kwargs, - ) -> Union[Tuple, CausalLMOutputWithPast]: + ) -> Union[tuple, CausalLMOutputWithPast]: r""" labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set diff --git a/optimum/habana/transformers/models/cohere/modeling_cohere.py b/optimum/habana/transformers/models/cohere/modeling_cohere.py index 0890367808..e2cad79930 100644 --- a/optimum/habana/transformers/models/cohere/modeling_cohere.py +++ b/optimum/habana/transformers/models/cohere/modeling_cohere.py @@ -1,7 +1,7 @@ -from functools import partial -from typing import List, Optional, Tuple, Union +from typing import Optional, Union import torch +from transformers.masking_utils import create_causal_mask from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from transformers.models.cohere.modeling_cohere import ( Cache, @@ -11,13 +11,12 @@ CohereForCausalLM, CohereRotaryEmbedding, DynamicCache, - KwargsForCausalLM, StaticCache, apply_rotary_pos_emb, eager_attention_forward, - logger, ) from transformers.processing_utils import Unpack +from transformers.utils import TransformersKwargs from ...modeling_attn_mask_utils import _gaudi_prepare_4d_causal_attention_mask @@ -31,13 +30,13 @@ def __init__(self, config: CohereConfig, layer_idx: Optional[int] = None): def forward( self, hidden_states: torch.Tensor, - position_embeddings: Tuple[torch.Tensor, torch.Tensor], + position_embeddings: tuple[torch.Tensor, torch.Tensor], attention_mask: Optional[torch.Tensor], past_key_value: Optional[Cache] = None, cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, **kwargs, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]: """ Copied from CohereAttention.forward: https://github.com/huggingface/transformers/blob/main/src/transformers/models/cohere/modeling_cohere.py The only differences are: @@ -102,43 +101,34 @@ def forward( attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Cache] = None, - output_attentions: Optional[bool] = False, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, - position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC + position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC token_idx: Optional[torch.Tensor] = None, - ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]: """ Copied from CohereDecoderLayer.forward: https://github.com/huggingface/transformers/blob/main/src/transformers/models/cohere/modeling_cohere.py The only differences are: - add new args token_idx """ residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - # Self Attention - hidden_states_attention, self_attn_weights, present_key_value = self.self_attn( + hidden_states_attention, _, present_key_value = self.self_attn( hidden_states=hidden_states, attention_mask=attention_mask, position_ids=position_ids, past_key_value=past_key_value, - output_attentions=output_attentions, use_cache=use_cache, cache_position=cache_position, position_embeddings=position_embeddings, token_idx=token_idx, ) - # Fully Connected hidden_states_mlp = self.mlp(hidden_states) - - # Add everything together hidden_states = residual + hidden_states_attention + hidden_states_mlp outputs = (hidden_states,) - if output_attentions: - outputs += (self_attn_weights,) if use_cache: outputs += (present_key_value,) @@ -152,10 +142,8 @@ def gaudi_cohere_model_forward( position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[Cache] = None, inputs_embeds: Optional[torch.FloatTensor] = None, - use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, token_idx: Optional[torch.Tensor] = None, **kwargs, ) -> BaseModelOutputWithPast: @@ -164,21 +152,13 @@ def gaudi_cohere_model_forward( The only differences are: - add new args token_idx """ - 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 - ) use_cache = use_cache if use_cache is not None else self.config.use_cache if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") - if self.gradient_checkpointing and self.training and use_cache: - logger.warning_once("`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`.") - use_cache = False - if inputs_embeds is None: - inputs_embeds = self.embed_tokens(input_ids) + inputs_embeds: torch.Tensor = self.embed_tokens(input_ids) past_seen_tokens = 0 return_legacy_cache = False @@ -190,66 +170,43 @@ def gaudi_cohere_model_forward( if cache_position is None: past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 - cache_position = torch.arange( + cache_position: torch.Tensor = torch.arange( past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device ) if position_ids is None: position_ids = cache_position.unsqueeze(0) - causal_mask = self._update_causal_mask( - attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions + causal_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, + position_ids=position_ids, ) - # embed positions hidden_states = inputs_embeds - # decoder layers - all_hidden_states = () if output_hidden_states else None - all_self_attns = () if output_attentions else None next_decoder_cache = None for decoder_layer in self.layers[: self.config.num_hidden_layers]: - if output_hidden_states: - all_hidden_states += (hidden_states,) - - if self.gradient_checkpointing and self.training: - layer_outputs = self._gradient_checkpointing_func( - partial(decoder_layer.__call__, **kwargs), - hidden_states, - causal_mask, - position_ids, - past_key_values, - output_attentions, - use_cache, - cache_position, - ) - else: - layer_outputs = decoder_layer( - hidden_states, - attention_mask=causal_mask, - position_ids=position_ids, - past_key_value=past_key_values, - output_attentions=output_attentions, - use_cache=use_cache, - cache_position=cache_position, - token_idx=token_idx, - ) - - hidden_states = layer_outputs[0] + hidden_states = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_ids=position_ids, + past_key_value=past_key_values, + use_cache=use_cache, + cache_position=cache_position, + token_idx=token_idx, + **kwargs, + ) if use_cache: - next_decoder_cache = layer_outputs[2 if output_attentions else 1] - - if output_attentions: - all_self_attns += (layer_outputs[1],) + next_decoder_cache = hidden_states[1] hidden_states = self.norm(hidden_states) - # add hidden states from the last decoder layer - if output_hidden_states: - all_hidden_states += (hidden_states,) - next_cache = next_decoder_cache if use_cache else None if return_legacy_cache: next_cache = next_cache.to_legacy_cache() @@ -257,8 +214,6 @@ def gaudi_cohere_model_forward( return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=next_cache, - hidden_states=all_hidden_states, - attentions=all_self_attns, ) @@ -277,7 +232,7 @@ def forward( input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, + past_key_values: Optional[Union[Cache, list[torch.FloatTensor]]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, @@ -286,7 +241,7 @@ def forward( cache_position: Optional[torch.LongTensor] = None, logits_to_keep: Union[int, torch.Tensor] = 0, token_idx: Optional[torch.Tensor] = None, - **kwargs: Unpack[KwargsForCausalLM], + **kwargs: Unpack[TransformersKwargs], ) -> CausalLMOutputWithPast: output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = ( diff --git a/optimum/habana/transformers/models/esm/modeling_esmfold.py b/optimum/habana/transformers/models/esm/modeling_esmfold.py index 88b6dac0d2..42b4cecac5 100644 --- a/optimum/habana/transformers/models/esm/modeling_esmfold.py +++ b/optimum/habana/transformers/models/esm/modeling_esmfold.py @@ -108,6 +108,7 @@ def gaudi_esm_for_protein_folding_forward( position_ids: Optional[torch.Tensor] = None, masking_pattern: Optional[torch.Tensor] = None, num_recycles: Optional[int] = None, + output_hidden_states: Optional[bool] = False, ) -> EsmForProteinFoldingOutput: r""" Returns: diff --git a/optimum/habana/transformers/models/falcon/modeling_falcon.py b/optimum/habana/transformers/models/falcon/modeling_falcon.py index a59f3426ec..c773295df3 100644 --- a/optimum/habana/transformers/models/falcon/modeling_falcon.py +++ b/optimum/habana/transformers/models/falcon/modeling_falcon.py @@ -1,7 +1,7 @@ import contextlib import math import os -from typing import Optional, Tuple, Union +from typing import Optional, Union import torch @@ -210,7 +210,7 @@ def __init__(self, config: FalconConfig, layer_idx=None): def _split_heads( self, fused_qkv: torch.Tensor, broadcast: Optional[bool] = True - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: if self.new_decoder_architecture: batch, seq_len, _ = fused_qkv.shape @@ -281,7 +281,7 @@ def pre_attn_forward( use_cache: bool = False, output_attentions: bool = False, cache_position: Optional[torch.LongTensor] = None, - position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, token_idx: Optional[torch.Tensor] = None, reuse_cache: Optional[bool] = False, cache_idx: int = None, @@ -570,12 +570,12 @@ def forward( alibi: Optional[torch.Tensor], attention_mask: torch.Tensor, position_ids: Optional[torch.LongTensor] = None, - layer_past: Optional[Union[Cache, Tuple[torch.Tensor, torch.Tensor]]] = None, + layer_past: Optional[Union[Cache, tuple[torch.Tensor, torch.Tensor]]] = None, head_mask: Optional[torch.Tensor] = None, use_cache: bool = False, output_attentions: bool = False, cache_position: Optional[torch.LongTensor] = None, - position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, token_idx: Optional[torch.Tensor] = None, reuse_cache: Optional[bool] = False, cache_idx: int = None, @@ -656,12 +656,12 @@ def pre_attn( alibi: Optional[torch.Tensor], attention_mask: torch.Tensor, position_ids: Optional[torch.LongTensor] = None, - layer_past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + layer_past: Optional[tuple[torch.Tensor, torch.Tensor]] = None, head_mask: Optional[torch.Tensor] = None, use_cache: bool = False, output_attentions: bool = False, cache_position: Optional[torch.LongTensor] = None, - position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # will become mandatory in v4.46 + position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, # will become mandatory in v4.46 token_idx: Optional[torch.Tensor] = None, reuse_cache: Optional[bool] = False, cache_idx: int = None, @@ -722,7 +722,7 @@ def update_sincos_cache(self, seq_len): def forward( self, input_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[Union[Cache, Tuple[Tuple[torch.Tensor, torch.Tensor], ...]]] = None, + past_key_values: Optional[Union[Cache, tuple[tuple[torch.Tensor, torch.Tensor], ...]]] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, head_mask: Optional[torch.LongTensor] = None, @@ -738,7 +738,7 @@ def forward( use_flash_attention: Optional[bool] = False, flash_attention_recompute: Optional[bool] = False, flash_attention_causal_mask: Optional[bool] = False, - ) -> Union[Tuple[torch.Tensor, ...], BaseModelOutputWithPastAndCrossAttentions]: + ) -> Union[tuple[torch.Tensor, ...], BaseModelOutputWithPastAndCrossAttentions]: 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 @@ -858,43 +858,24 @@ def forward( if output_hidden_states: all_hidden_states = all_hidden_states + (hidden_states,) - if self.gradient_checkpointing and self.training: - outputs = self._gradient_checkpointing_func( - block.__call__, - hidden_states, - alibi, - attention_mask, - position_ids, - head_mask[i], - layer_past, - use_cache, - output_attentions, - cache_position, - position_embeddings, - None, - use_flash_attention, - flash_attention_recompute, - flash_attention_causal_mask, - ) - else: - outputs = block( - hidden_states, - layer_past=layer_past, - attention_mask=attention_mask, - position_ids=position_ids, - head_mask=head_mask[i], - use_cache=use_cache, - output_attentions=output_attentions, - alibi=alibi, - cache_position=cache_position, - position_embeddings=position_embeddings, - token_idx=token_idx, - reuse_cache=reuse_cache, - cache_idx=cache_idx, - use_flash_attention=use_flash_attention, - flash_attention_recompute=flash_attention_recompute, - flash_attention_causal_mask=flash_attention_causal_mask, - ) + outputs = block( + hidden_states, + layer_past=layer_past, + attention_mask=attention_mask, + position_ids=position_ids, + head_mask=head_mask[i], + use_cache=use_cache, + output_attentions=output_attentions, + alibi=alibi, + cache_position=cache_position, + position_embeddings=position_embeddings, + token_idx=token_idx, + reuse_cache=reuse_cache, + cache_idx=cache_idx, + use_flash_attention=use_flash_attention, + flash_attention_recompute=flash_attention_recompute, + flash_attention_causal_mask=flash_attention_causal_mask, + ) hidden_states = outputs[0] if use_cache is True: @@ -1016,7 +997,7 @@ def prepare_inputs_for_generation( def forward( self, input_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[Union[Cache, Tuple[Tuple[torch.Tensor, torch.Tensor], ...]]] = None, + past_key_values: Optional[Union[Cache, tuple[tuple[torch.Tensor, torch.Tensor], ...]]] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, head_mask: Optional[torch.Tensor] = None, @@ -1036,7 +1017,7 @@ def forward( flash_attention_recompute: Optional[bool] = False, flash_attention_causal_mask: Optional[bool] = False, **kwargs, - ) -> Union[Tuple[torch.Tensor], CausalLMOutputWithCrossAttentions]: + ) -> Union[tuple[torch.Tensor], CausalLMOutputWithCrossAttentions]: r""" labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set diff --git a/optimum/habana/transformers/models/falcon_mamba/modeling_falcon_mamba.py b/optimum/habana/transformers/models/falcon_mamba/modeling_falcon_mamba.py index dddaa5055a..98212a3ccb 100644 --- a/optimum/habana/transformers/models/falcon_mamba/modeling_falcon_mamba.py +++ b/optimum/habana/transformers/models/falcon_mamba/modeling_falcon_mamba.py @@ -14,12 +14,11 @@ # limitations under the License. """PyTorch FALCONMAMBA model.""" -from typing import Optional, Tuple, Union +from typing import Optional, Union import habana_frameworks.torch.core as htcore import torch -from transformers.cache_utils import MambaCache -from transformers.models.falcon_mamba.modeling_falcon_mamba import FalconMambaOutput +from transformers.models.falcon_mamba.modeling_falcon_mamba import FalconMambaCache, FalconMambaOutput from transformers.utils import ( logging, ) @@ -38,14 +37,14 @@ def gaudi_FalconMambaModel_forward( self, input_ids: Optional[torch.LongTensor] = None, inputs_embeds: Optional[torch.LongTensor] = None, - cache_params: Optional[MambaCache] = None, + cache_params: Optional[FalconMambaCache] = None, use_cache: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.LongTensor] = None, lazy_mode: Optional[bool] = True, -) -> Union[Tuple, FalconMambaOutput]: +) -> Union[tuple, FalconMambaOutput]: output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) @@ -63,7 +62,7 @@ def gaudi_FalconMambaModel_forward( if use_cache: if cache_params is None: - cache_params = MambaCache( + cache_params = FalconMambaCache( self.config, inputs_embeds.size(0), device=inputs_embeds.device, dtype=inputs_embeds.dtype ) cache_position = torch.arange(0, self.config.conv_kernel, device=inputs_embeds.device) @@ -78,22 +77,18 @@ def gaudi_FalconMambaModel_forward( ) else: cache_params = None + hidden_states = inputs_embeds all_hidden_states = () if output_hidden_states else None for mixer_block in self.layers: if lazy_mode: htcore.mark_step() - if self.gradient_checkpointing and self.training: - hidden_states = self._gradient_checkpointing_func( - mixer_block.__call__, hidden_states, cache_params, cache_position, attention_mask - ) - else: - hidden_states = mixer_block( - hidden_states, - cache_params=cache_params, - cache_position=cache_position, - attention_mask=attention_mask, - ) + hidden_states = mixer_block( + hidden_states, + cache_params=cache_params, + cache_position=cache_position, + attention_mask=attention_mask, + ) if output_hidden_states: all_hidden_states = all_hidden_states + (hidden_states,) @@ -123,38 +118,33 @@ def gaudi_FalconMambaForCausalLM_prepare_inputs_for_generation( input_ids, inputs_embeds=None, use_cache=None, - cache_params: Optional[MambaCache] = None, + cache_params: Optional[FalconMambaCache] = None, cache_position: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.LongTensor] = None, **kwargs, ): - if use_cache: - # `cache_position` should have been initialized in `generate` - if cache_position is None: - raise ValueError( - "`cache_position` should not be None as it should have been initialized in " - "`model.generate`, you are responsible for passing in a valid `cache_position` if " - "you are calling `prepare_inputs_for_generation` directly with `use_cache=True`" - ) - if cache_position[0] > 0: - # input_ids = input_ids[:, -1].unsqueeze(-1) - idx = torch.tensor([input_ids.size(1) - 1], device=input_ids.device) - input_ids = torch.index_select(input_ids, 1, idx) - - if attention_mask is not None: - attention_mask = None - + model_inputs = {"input_ids": input_ids.contiguous()} + if use_cache and cache_params is None: + # we initialize the `cache_position` to full size of `conv_states` at prefill stage + # considering padding will be applied when input length is shorter, and truncation + # will be applied when it is longer, so it will be equivalent to always have it match + # the length of `cache_params.conv_states`, which is `config.conv_kernel` + cache_position = torch.arange(0, self.backbone.config.conv_kernel, device=input_ids.device) + if inputs_embeds is not None: + model_inputs = {"inputs_embeds": inputs_embeds} + max_batch_size = inputs_embeds.size(0) else: - # we initialize the `cache_position` to full size of `conv_states` at prefill stage - # considering padding will be applied when input length is shorter, and truncation - # will be applied when it is longer, so it will be equivalent to always have it match - # the length of `cache_params.conv_states`, which is `config.conv_kernel` - cache_position = torch.arange(0, self.config.conv_kernel, device=input_ids.device) + max_batch_size = input_ids.size(0) + cache_params = FalconMambaCache(self.backbone.config, max_batch_size, device=self.device, dtype=self.dtype) - if inputs_embeds is not None and cache_params is None: + if use_cache and cache_position[0] > 0: + # input_ids = input_ids[:, -1].unsqueeze(-1) + idx = torch.tensor([input_ids.size(1) - 1], device=input_ids.device) + model_inputs["input_ids"] = torch.index_select(input_ids, 1, idx).contiguous() + attention_mask = None + + if not use_cache and inputs_embeds is not None: model_inputs = {"inputs_embeds": inputs_embeds} - else: - model_inputs = {"input_ids": input_ids.contiguous()} model_inputs.update( { diff --git a/optimum/habana/transformers/models/gemma/modeling_gemma.py b/optimum/habana/transformers/models/gemma/modeling_gemma.py index 425db073e5..f02ae65365 100755 --- a/optimum/habana/transformers/models/gemma/modeling_gemma.py +++ b/optimum/habana/transformers/models/gemma/modeling_gemma.py @@ -21,7 +21,7 @@ import math import os -from typing import List, Optional, Tuple, Union +from typing import Optional, Union import torch import torch.nn.functional as F @@ -34,11 +34,10 @@ GemmaForCausalLM, GemmaMLP, GemmaModel, - KwargsForCausalLM, apply_rotary_pos_emb, ) from transformers.processing_utils import Unpack -from transformers.utils import logging +from transformers.utils import TransformersKwargs, logging from ...modeling_attn_mask_utils import ( _gaudi_prepare_4d_causal_attention_mask, @@ -274,7 +273,7 @@ def gaudi_flash_attn_v1( def pre_attn_forward( self, hidden_states: torch.Tensor, - position_embeddings: Tuple[torch.Tensor, torch.Tensor], + position_embeddings: tuple[torch.Tensor, torch.Tensor], attention_mask: Optional[torch.Tensor], past_key_value: Optional[Cache] = None, use_cache: bool = False, @@ -287,7 +286,7 @@ def pre_attn_forward( flash_attention_causal_mask: Optional[bool] = False, cache_idx: Optional[int] = None, **kwargs, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]: """ The only differences are: - add new args token_idx @@ -454,11 +453,10 @@ def pre_attn( hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_value: Optional[Tuple[torch.Tensor]] = None, - output_attentions: Optional[bool] = False, + past_key_value: Optional[tuple[torch.Tensor]] = None, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, - position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC + position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC token_idx: Optional[torch.Tensor] = None, attn_softmax_bf16: Optional[bool] = False, reuse_cache: Optional[bool] = False, @@ -466,14 +464,13 @@ def pre_attn( flash_attention_recompute: Optional[bool] = False, flash_attention_causal_mask: Optional[bool] = False, cache_idx: Optional[int] = None, - ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]: hidden_states = self.input_layernorm(hidden_states) hidden_states, attn_weights, present_key_value = self.self_attn.pre_attn_forward( hidden_states=hidden_states, attention_mask=attention_mask, position_ids=position_ids, past_key_value=past_key_value, - output_attentions=output_attentions, use_cache=use_cache, cache_position=cache_position, position_embeddings=position_embeddings, @@ -493,7 +490,6 @@ def forward( attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Cache] = None, - output_attentions: Optional[bool] = False, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, @@ -503,7 +499,7 @@ def forward( flash_attention_recompute: Optional[bool] = False, flash_attention_causal_mask: Optional[bool] = False, cache_idx: Optional[int] = None, - ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]: """ Copied from GemmaDecoderLayer.forward: https://github.com/huggingface/transformers/blob/v4.38.1/src/transformers/models/gemma/modeling_gemma.py The only differences are: @@ -517,7 +513,6 @@ def forward( attention_mask=attention_mask, position_ids=position_ids, past_key_value=past_key_value, - output_attentions=output_attentions, use_cache=use_cache, cache_position=cache_position, token_idx=token_idx, @@ -538,8 +533,6 @@ def forward( hidden_states = self.post_mlp(hidden_states, residual) outputs = (hidden_states,) - if output_attentions: - outputs += (self_attn_weights,) if use_cache: outputs += (present_key_value,) @@ -589,11 +582,9 @@ def forward( input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, + past_key_values: Optional[list[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, attn_softmax_bf16: Optional[bool] = False, @@ -603,17 +594,13 @@ def forward( flash_attention_causal_mask: Optional[bool] = False, cache_idx: int = None, lazy_mode: Optional[bool] = True, - **kwargs, # NOOP kwarg for now + **kwargs, ) -> BaseModelOutputWithPast: """ Copied from GemmaModel.forward: https://github.com/huggingface/transformers/blob/v4.38.1/src/transformers/models/gemma/modeling_gemma.py The only differences are: - add new args token_idx """ - 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 - ) use_cache = use_cache if use_cache is not None else self.config.use_cache self._attn_implementation = "eager" @@ -627,12 +614,6 @@ def forward( else: raise ValueError("You have to specify either input_ids or inputs_embeds") - if self.gradient_checkpointing and self.training and use_cache: - logger.warning_once( - "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`." - ) - use_cache = False - if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) @@ -673,9 +654,6 @@ def forward( normalizer = torch.tensor(self.config.hidden_size**0.5, dtype=hidden_states.dtype, device=inputs_embeds.device) hidden_states = hidden_states * normalizer - # decoder layers - all_hidden_states = () if output_hidden_states else None - all_self_attns = () if output_attentions else None next_decoder_cache = () if not use_new_cache else None if lazy_mode: @@ -689,58 +667,28 @@ def forward( ): htcore.mark_step() - if output_hidden_states: - all_hidden_states += (hidden_states,) - - if self.gradient_checkpointing and self.training: - layer_outputs = self._gradient_checkpointing_func( - decoder_layer.__call__, - hidden_states, - attention_mask, - position_ids, - past_key_values, - output_attentions, - use_cache, - cache_position, - None, - attn_softmax_bf16, - False, - use_flash_attention, - flash_attention_recompute, - flash_attention_causal_mask, - ) - else: - layer_outputs = decoder_layer( - hidden_states, - attention_mask=attention_mask, - position_ids=position_ids, - past_key_value=None if past_key_values is None else past_key_values[layer_idx], - output_attentions=output_attentions, - use_cache=use_cache, - cache_position=cache_position, - token_idx=token_idx, - attn_softmax_bf16=attn_softmax_bf16, - reuse_cache=reuse_cache, - use_flash_attention=use_flash_attention, - flash_attention_recompute=flash_attention_recompute, - flash_attention_causal_mask=flash_attention_causal_mask, - cache_idx=cache_idx, - ) - - hidden_states = layer_outputs[0] + hidden_states = decoder_layer( + hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=None if past_key_values is None else past_key_values[layer_idx], + use_cache=use_cache, + cache_position=cache_position, + token_idx=token_idx, + attn_softmax_bf16=attn_softmax_bf16, + reuse_cache=reuse_cache, + use_flash_attention=use_flash_attention, + flash_attention_recompute=flash_attention_recompute, + flash_attention_causal_mask=flash_attention_causal_mask, + cache_idx=cache_idx, + **kwargs, + ) if use_cache: - next_decoder_cache += (layer_outputs[2 if output_attentions else 1],) - - if output_attentions: - all_self_attns += (layer_outputs[1],) + next_decoder_cache += (hidden_states[1],) hidden_states = self.norm(hidden_states) - # add hidden states from the last decoder layer - if output_hidden_states: - all_hidden_states += (hidden_states,) - next_cache = None if use_cache: next_cache = ( @@ -750,8 +698,6 @@ def forward( return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=next_cache, - hidden_states=all_hidden_states, - attentions=all_self_attns, ) @@ -770,13 +716,11 @@ def forward( input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, + past_key_values: Optional[Union[Cache, list[torch.FloatTensor]]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, reuse_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, logits_to_keep: Union[int, torch.Tensor] = 0, token_idx: Optional[torch.Tensor] = None, @@ -784,7 +728,7 @@ def forward( flash_attention_recompute: Optional[bool] = False, flash_attention_causal_mask: Optional[bool] = False, attn_softmax_bf16: Optional[bool] = False, - **kwargs: Unpack[KwargsForCausalLM], + **kwargs: Unpack[TransformersKwargs], ) -> CausalLMOutputWithPast: """ Inherits from GemmaForCausalLM: https://github.com/huggingface/transformers/blob/v4.38.1/src/transformers/models/gemma/modeling_gemma.py @@ -792,12 +736,6 @@ def forward( - add new args token_idx - add new args attn_softmax_bf16 """ - 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 - ) - - # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) outputs: BaseModelOutputWithPast = self.model( input_ids=input_ids, attention_mask=attention_mask, @@ -806,8 +744,6 @@ def forward( inputs_embeds=inputs_embeds, use_cache=use_cache, reuse_cache=reuse_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, cache_position=cache_position, token_idx=token_idx, use_flash_attention=use_flash_attention, diff --git a/optimum/habana/transformers/models/gemma2/modeling_gemma2.py b/optimum/habana/transformers/models/gemma2/modeling_gemma2.py index 1143ae95f6..583f888c68 100755 --- a/optimum/habana/transformers/models/gemma2/modeling_gemma2.py +++ b/optimum/habana/transformers/models/gemma2/modeling_gemma2.py @@ -16,8 +16,7 @@ """PyTorch Gemma2 model.""" import math -from functools import partial -from typing import List, Optional, Tuple, Union +from typing import Optional, Union import torch import torch.nn.functional as F @@ -255,7 +254,7 @@ def gaudi_eager_attention_forward( scaling: Optional[float] = None, softcap: Optional[float] = None, **kwargs, -) -> Tuple[torch.Tensor, torch.Tensor]: +) -> tuple[torch.Tensor, torch.Tensor]: bsz, q_len = kwargs["input_shape"] if scaling is None: @@ -359,7 +358,7 @@ def gaudi_flash_attn_v1(self, query_layer, key_layer, value_layer, attention_mas def pre_attn_forward( self, hidden_states: torch.Tensor, - position_embeddings: Tuple[torch.Tensor, torch.Tensor], + position_embeddings: tuple[torch.Tensor, torch.Tensor], attention_mask: Optional[torch.Tensor], past_key_value: Optional[Cache] = None, use_cache: bool = False, @@ -373,7 +372,7 @@ def pre_attn_forward( flash_attention_fast_softmax: Optional[bool] = False, cache_idx: int = None, **kwargs, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]: """ The only differences are: - add new args token_idx @@ -533,10 +532,10 @@ def update_sincos_cache(self, seq_len): def pre_attn( self, hidden_states: torch.Tensor, - position_embeddings: Tuple[torch.Tensor, torch.Tensor], + position_embeddings: tuple[torch.Tensor, torch.Tensor], attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_value: Optional[Tuple[torch.Tensor]] = None, + past_key_value: Optional[tuple[torch.Tensor]] = None, output_attentions: Optional[bool] = False, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, @@ -549,7 +548,7 @@ def pre_attn( flash_attention_fast_softmax: Optional[bool] = False, cache_idx: int = None, **kwargs, - ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]: hidden_states = self.input_layernorm(hidden_states) hidden_states, attn_weights, present_key_value = self.self_attn.pre_attn_forward( @@ -576,7 +575,7 @@ def pre_attn( def forward( self, hidden_states: torch.Tensor, - position_embeddings: Tuple[torch.Tensor, torch.Tensor] = None, + position_embeddings: tuple[torch.Tensor, torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Cache] = None, @@ -593,7 +592,7 @@ def forward( flash_attention_fast_softmax: Optional[bool] = False, cache_idx: int = None, **kwargs, - ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]: """ Copied from GemmaDecoderLayer.forward: https://github.com/huggingface/transformers/blob/v4.38.1/src/transformers/models/gemma/modeling_gemma.py The only differences are: @@ -687,7 +686,7 @@ def forward( input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, + past_key_values: Optional[Union[Cache, list[torch.FloatTensor]]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, @@ -811,47 +810,26 @@ def forward( if output_hidden_states: all_hidden_states += (hidden_states,) - if self.gradient_checkpointing and self.training: - layer_outputs = self._gradient_checkpointing_func( - partial(decoder_layer.__call__, **kwargs), - hidden_states, - position_embeddings, - causal_mask, - position_ids, - past_key_values, - output_attentions, - use_cache, - cache_position, - last_cache_position, - None, - attn_softmax_bf16, - False, - use_flash_attention, - flash_attention_recompute, - flash_attention_causal_mask, - flash_attention_fast_softmax, - None, - ) - else: - layer_outputs = decoder_layer( - hidden_states, - position_embeddings=position_embeddings, - attention_mask=causal_mask, - position_ids=position_ids, - past_key_value=None if past_key_values is None else past_key_values[layer_idx], - output_attentions=output_attentions, - use_cache=use_cache, - cache_position=cache_position, - last_cache_position=last_cache_position, - token_idx=token_idx, - attn_softmax_bf16=attn_softmax_bf16, - reuse_cache=reuse_cache, - use_flash_attention=use_flash_attention, - flash_attention_recompute=flash_attention_recompute, - flash_attention_causal_mask=flash_attention_causal_mask, - flash_attention_fast_softmax=flash_attention_fast_softmax, - cache_idx=cache_idx, - ) + layer_outputs = decoder_layer( + hidden_states, + position_embeddings=position_embeddings, + attention_mask=causal_mask, + position_ids=position_ids, + past_key_value=None if past_key_values is None else past_key_values[layer_idx], + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + last_cache_position=last_cache_position, + token_idx=token_idx, + attn_softmax_bf16=attn_softmax_bf16, + reuse_cache=reuse_cache, + use_flash_attention=use_flash_attention, + flash_attention_recompute=flash_attention_recompute, + flash_attention_causal_mask=flash_attention_causal_mask, + flash_attention_fast_softmax=flash_attention_fast_softmax, + cache_idx=cache_idx, + **kwargs, + ) hidden_states = layer_outputs[0] @@ -895,7 +873,7 @@ def forward( input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, + past_key_values: Optional[Union[Cache, list[torch.FloatTensor]]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, @@ -913,7 +891,7 @@ def forward( flash_attention_fast_softmax: Optional[bool] = False, cache_idx: int = None, lazy_mode: Optional[bool] = True, - **loss_kwargs, + **kwargs, ) -> CausalLMOutputWithPast: """ Inherits from GemmaForCausalLM: https://github.com/huggingface/transformers/blob/v4.38.1/src/transformers/models/gemma/modeling_gemma.py @@ -945,7 +923,7 @@ def forward( flash_attention_fast_softmax=flash_attention_fast_softmax, cache_idx=cache_idx, lazy_mode=lazy_mode, - **loss_kwargs, + **kwargs, ) hidden_states = outputs.last_hidden_state @@ -966,7 +944,7 @@ def forward( loss = None if labels is not None: - loss = self.loss_function(logits, labels, self.vocab_size, **loss_kwargs) + loss = self.loss_function(logits, labels, self.vocab_size, **kwargs) return CausalLMOutputWithPast( loss=loss, diff --git a/optimum/habana/transformers/models/gpt_neo/modeling_gpt_neo.py b/optimum/habana/transformers/models/gpt_neo/modeling_gpt_neo.py index 1cb65bffd0..fb2c317752 100644 --- a/optimum/habana/transformers/models/gpt_neo/modeling_gpt_neo.py +++ b/optimum/habana/transformers/models/gpt_neo/modeling_gpt_neo.py @@ -1,4 +1,4 @@ -from typing import Optional, Tuple, Union +from typing import Optional, Union import torch from transformers.modeling_outputs import ( @@ -12,32 +12,6 @@ from ...modeling_attn_mask_utils import _gaudi_prepare_4d_causal_attention_mask -def gaudi_gpt_neo_attention_forward( - self, - hidden_states, - attention_mask=None, - layer_past=None, - head_mask=None, - use_cache=False, - output_attentions=False, - token_idx=None, -): - """ - Copied from GPTNeoAttention.forward: https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt_neo/modeling_gpt_neo.py - The only differences are: - - add new args token_idx - """ - return self.attention( - hidden_states, - attention_mask=attention_mask, - layer_past=layer_past, - head_mask=head_mask, - use_cache=use_cache, - output_attentions=output_attentions, - token_idx=token_idx, - ) - - def gaudi_gpt_neo_selfattention_forward( self, hidden_states, @@ -86,12 +60,36 @@ def gaudi_gpt_neo_selfattention_forward( attn_output = self.resid_dropout(attn_output) outputs = (attn_output, present) - if output_attentions: - outputs += (attn_weights,) return outputs # a, present, (attentions) +def gaudi_gpt_neo_attention_forward( + self, + hidden_states, + attention_mask=None, + layer_past=None, + head_mask=None, + use_cache=False, + output_attentions=False, + token_idx=None, +): + """ + Copied from GPTNeoAttention.forward: https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt_neo/modeling_gpt_neo.py + The only differences are: + - add new args token_idx + """ + return self.attention( + hidden_states, + attention_mask=attention_mask, + layer_past=layer_past, + head_mask=head_mask, + use_cache=use_cache, + output_attentions=output_attentions, + token_idx=token_idx, + ) + + def gaudi_gpt_neo_block_forward( self, hidden_states, @@ -109,7 +107,7 @@ def gaudi_gpt_neo_block_forward( """ residual = hidden_states hidden_states = self.ln_1(hidden_states) - attn_outputs = self.attn( + attn_output, attn_weights = self.attn( hidden_states, layer_past=layer_past, attention_mask=attention_mask, @@ -118,8 +116,7 @@ def gaudi_gpt_neo_block_forward( output_attentions=output_attentions, token_idx=token_idx, ) - attn_output = attn_outputs[0] # output_attn: a, present, (attentions) - outputs = attn_outputs[1:] + # residual connection hidden_states = attn_output + residual @@ -129,18 +126,13 @@ def gaudi_gpt_neo_block_forward( # residual connection hidden_states = residual + feed_forward_hidden_states - if use_cache: - outputs = (hidden_states,) + outputs - else: - outputs = (hidden_states,) + outputs[1:] - - return outputs # hidden_states, present, (attentions, cross_attentions) + return hidden_states, attn_weights def gaudi_gpt_neo_model_forward( self, input_ids: Optional[torch.Tensor] = None, - past_key_values: Optional[Tuple[torch.FloatTensor]] = None, + past_key_values: Optional[tuple[torch.FloatTensor]] = None, attention_mask: Optional[torch.Tensor] = None, token_type_ids: Optional[torch.Tensor] = None, position_ids: Optional[torch.Tensor] = None, @@ -151,7 +143,7 @@ def gaudi_gpt_neo_model_forward( output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, token_idx: Optional[torch.Tensor] = None, -) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPastAndCrossAttentions]: +) -> Union[tuple[torch.Tensor], BaseModelOutputWithPastAndCrossAttentions]: """ Copied from GPTNeoModel.forward: https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt_neo/modeling_gpt_neo.py The only differences are: @@ -234,26 +226,15 @@ def gaudi_gpt_neo_model_forward( if output_hidden_states: all_hidden_states = all_hidden_states + (hidden_states,) - if self.gradient_checkpointing and self.training: - outputs = self._gradient_checkpointing_func( - block.__call__, - hidden_states, - None, - attention_mask, - head_mask[i], - use_cache, - output_attentions, - ) - else: - outputs = block( - hidden_states, - layer_past=layer_past, - attention_mask=attention_mask, - head_mask=head_mask[i], - use_cache=use_cache, - output_attentions=output_attentions, - token_idx=token_idx, - ) + outputs = block( + hidden_states, + layer_past=layer_past, + attention_mask=attention_mask, + head_mask=head_mask[i], + use_cache=use_cache, + output_attentions=output_attentions, + token_idx=token_idx, + ) hidden_states = outputs[0] if use_cache is True: @@ -293,7 +274,7 @@ class GaudiGPTNeoForCausalLM(GPTNeoForCausalLM): def forward( self, input_ids: Optional[torch.Tensor] = None, - past_key_values: Optional[Tuple[torch.FloatTensor]] = None, + past_key_values: Optional[tuple[torch.FloatTensor]] = None, attention_mask: Optional[torch.Tensor] = None, token_type_ids: Optional[torch.Tensor] = None, position_ids: Optional[torch.Tensor] = None, @@ -307,7 +288,7 @@ def forward( cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, **kwargs, - ) -> Union[Tuple[torch.Tensor], CausalLMOutputWithCrossAttentions]: + ) -> Union[tuple[torch.Tensor], CausalLMOutputWithCrossAttentions]: r""" labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set diff --git a/optimum/habana/transformers/models/gpt_neox/modeling_gpt_neox.py b/optimum/habana/transformers/models/gpt_neox/modeling_gpt_neox.py index fcaedb1066..5de06556a0 100644 --- a/optimum/habana/transformers/models/gpt_neox/modeling_gpt_neox.py +++ b/optimum/habana/transformers/models/gpt_neox/modeling_gpt_neox.py @@ -1,4 +1,4 @@ -from typing import Optional, Tuple, Union +from typing import Optional, Union import torch from transformers.cache_utils import Cache @@ -96,7 +96,7 @@ def forward( output_attentions: Optional[bool] = False, padding_mask: Optional[torch.Tensor] = None, cache_position: Optional[torch.LongTensor] = None, - position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC + position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC token_idx: Optional[torch.Tensor] = None, ): """ @@ -262,7 +262,7 @@ def gaudi_gpt_neox_model_forward( position_ids: Optional[torch.LongTensor] = None, head_mask: Optional[torch.FloatTensor] = None, inputs_embeds: Optional[torch.FloatTensor] = None, - past_key_values: Optional[Union[Cache, Tuple[Tuple[torch.FloatTensor]]]] = None, + past_key_values: Optional[Union[Cache, tuple[tuple[torch.FloatTensor]]]] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, @@ -339,31 +339,17 @@ def gaudi_gpt_neox_model_forward( if output_hidden_states: all_hidden_states = all_hidden_states + (hidden_states,) - if self.gradient_checkpointing and self.training: - outputs = self._gradient_checkpointing_func( - layer.__call__, - hidden_states, - attention_mask, - position_ids, - head_mask[i], - use_cache, - None, - output_attentions, - cache_position, - None, - ) - else: - outputs = layer( - hidden_states, - attention_mask=attention_mask, - position_ids=position_ids, - head_mask=head_mask[i], - layer_past=layer_past, - use_cache=use_cache, - output_attentions=output_attentions, - cache_position=cache_position, - token_idx=token_idx, - ) + outputs = layer( + hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + head_mask=head_mask[i], + layer_past=layer_past, + use_cache=use_cache, + output_attentions=output_attentions, + cache_position=cache_position, + token_idx=token_idx, + ) hidden_states = outputs[0] if use_cache is True: presents = presents + (outputs[1],) @@ -410,7 +396,7 @@ def forward( position_ids: Optional[torch.LongTensor] = None, inputs_embeds: Optional[torch.FloatTensor] = None, head_mask: Optional[torch.FloatTensor] = None, - past_key_values: Optional[Union[Cache, Tuple[Tuple[torch.FloatTensor]]]] = None, + past_key_values: Optional[Union[Cache, tuple[tuple[torch.FloatTensor]]]] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, @@ -419,7 +405,7 @@ def forward( token_idx: Optional[torch.Tensor] = None, logits_to_keep: Union[int, torch.Tensor] = 0, **kwargs: Unpack[KwargsForCausalLM], - ) -> Union[Tuple, CausalLMOutputWithPast]: + ) -> Union[tuple, CausalLMOutputWithPast]: outputs: BaseModelOutputWithPast = self.gpt_neox( input_ids, attention_mask=attention_mask, diff --git a/optimum/habana/transformers/models/gptj/modeling_gptj.py b/optimum/habana/transformers/models/gptj/modeling_gptj.py index a719dc645a..d0eee3e34b 100644 --- a/optimum/habana/transformers/models/gptj/modeling_gptj.py +++ b/optimum/habana/transformers/models/gptj/modeling_gptj.py @@ -1,4 +1,4 @@ -from typing import Optional, Tuple, Union +from typing import Optional, Union import habana_frameworks.torch.core as htcore import torch @@ -176,8 +176,8 @@ def forward( reuse_cache: Optional[bool] = False, cache_idx: Optional[int] = None, ) -> Union[ - Tuple[torch.Tensor, Tuple[torch.Tensor]], - Optional[Tuple[torch.Tensor, Tuple[torch.Tensor], Tuple[torch.Tensor, ...]]], + tuple[torch.Tensor, tuple[torch.Tensor]], + Optional[tuple[torch.Tensor, tuple[torch.Tensor], tuple[torch.Tensor, ...]]], ]: """ Copied from GPTJAttention.forward: https://github.com/huggingface/transformers/blob/main/src/transformers/models/gptj/modeling_gptj.py#194 @@ -263,10 +263,8 @@ def forward( attn_output = self.resid_dropout(attn_output) outputs = (attn_output, present) - if output_attentions: - outputs += (attn_weights,) - return outputs # a, present, (attentions) + return outputs class GaudiGPTJBlock(GPTJBlock): @@ -305,7 +303,7 @@ def forward( cos: Optional[torch.Tensor] = None, reuse_cache: Optional[bool] = False, cache_idx: Optional[int] = None, - ) -> Union[Tuple[torch.Tensor], Optional[Tuple[torch.Tensor, Tuple[torch.FloatTensor, ...]]]]: + ) -> Union[tuple[torch.Tensor], Optional[tuple[torch.Tensor, tuple[torch.FloatTensor, ...]]]]: """ Copied from GPTJBlock.forward: https://github.com/huggingface/transformers/blob/main/src/transformers/models/gptj/modeling_gptj.py The only differences are: @@ -362,7 +360,7 @@ def update_sincos_cache(self, seq_len): def forward( self, input_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[Union[Cache, Tuple[Tuple[torch.Tensor]]]] = None, + past_key_values: Optional[Union[Cache, tuple[tuple[torch.Tensor]]]] = None, attention_mask: Optional[torch.FloatTensor] = None, token_type_ids: Optional[torch.LongTensor] = None, position_ids: Optional[torch.LongTensor] = None, @@ -376,7 +374,7 @@ def forward( token_idx: Optional[torch.Tensor] = None, reuse_cache: Optional[bool] = False, cache_idx: Optional[int] = None, - ) -> Union[Tuple, BaseModelOutputWithPast]: + ) -> Union[tuple, BaseModelOutputWithPast]: """ Copied from https://github.com/huggingface/transformers/blob/v4.38.2/src/transformers/models/gptj/modeling_gptj.py#L554 The only differences are: @@ -491,37 +489,21 @@ def forward( if output_hidden_states: all_hidden_states = all_hidden_states + (hidden_states,) - if self.gradient_checkpointing and self.training: - outputs = self._gradient_checkpointing_func( - block.__call__, - hidden_states, - None, - attention_mask, - position_ids, - head_mask[i], - use_cache, - output_attentions, - cache_position, - None, - sin, - cos, - ) - else: - outputs = block( - hidden_states=hidden_states, - layer_past=layer_past, - attention_mask=attention_mask, - position_ids=position_ids, - head_mask=head_mask[i], - use_cache=use_cache, - output_attentions=output_attentions, - cache_position=cache_position, - token_idx=token_idx, - reuse_cache=reuse_cache, - cache_idx=cache_idx, - sin=sin, - cos=cos, - ) + outputs = block( + hidden_states, + layer_past=layer_past, + attention_mask=attention_mask, + position_ids=position_ids, + head_mask=head_mask[i], + use_cache=use_cache, + output_attentions=output_attentions, + cache_position=cache_position, + token_idx=token_idx, + reuse_cache=reuse_cache, + cache_idx=cache_idx, + sin=sin, + cos=cos, + ) hidden_states = outputs[0] if use_cache is True: @@ -646,7 +628,7 @@ def prepare_inputs_for_generation( def forward( self, input_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[Union[Cache, Tuple[Tuple[torch.Tensor]]]] = None, + past_key_values: Optional[Union[Cache, tuple[tuple[torch.Tensor]]]] = None, attention_mask: Optional[torch.FloatTensor] = None, token_type_ids: Optional[torch.LongTensor] = None, position_ids: Optional[torch.LongTensor] = None, @@ -662,7 +644,7 @@ def forward( reuse_cache: Optional[bool] = False, cache_idx: Optional[int] = None, **kwargs, - ) -> Union[Tuple, CausalLMOutputWithPast]: + ) -> Union[tuple, CausalLMOutputWithPast]: r""" labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set diff --git a/optimum/habana/transformers/models/idefics2/image_processing_idefics2.py b/optimum/habana/transformers/models/idefics2/image_processing_idefics2.py index 3bf45a36de..c2b4247928 100644 --- a/optimum/habana/transformers/models/idefics2/image_processing_idefics2.py +++ b/optimum/habana/transformers/models/idefics2/image_processing_idefics2.py @@ -14,7 +14,7 @@ # limitations under the License. -from typing import Iterable, List, Optional, Union +from typing import Iterable, Optional, Union import numpy as np from transformers.image_processing_utils import BatchFeature @@ -30,7 +30,7 @@ class Gaudi2Idefics2ImageProcessor(Idefics2ImageProcessor): def pad( self, - images: List[np.ndarray], + images: list[np.ndarray], constant_values: Union[float, Iterable[float]] = 0, return_pixel_mask: bool = True, return_tensors: Optional[Union[str, TensorType]] = None, diff --git a/optimum/habana/transformers/models/idefics2/modeling_idefics2.py b/optimum/habana/transformers/models/idefics2/modeling_idefics2.py index efc74746c6..1d11134774 100644 --- a/optimum/habana/transformers/models/idefics2/modeling_idefics2.py +++ b/optimum/habana/transformers/models/idefics2/modeling_idefics2.py @@ -14,11 +14,10 @@ # limitations under the License. """PyTorch Idefics2 model.""" -from typing import List, Optional, Tuple, Union +from typing import Optional, Union import torch import torch.utils.checkpoint -from torch.nn import CrossEntropyLoss from transformers.cache_utils import Cache, DynamicCache from transformers.models.idefics2.modeling_idefics2 import ( Idefics2BaseModelOutputWithPast, @@ -59,8 +58,11 @@ def forward(self, pixel_values: torch.FloatTensor, patch_attention_mask: torch.B nb_patches_h = int(p_attn_mask[:, 0].sum()) nb_patches_w = int(p_attn_mask[0].sum()) - fractional_coords_h = torch.arange(0, 1 - 1e-6, 1 / nb_patches_h) - fractional_coords_w = torch.arange(0, 1 - 1e-6, 1 / nb_patches_w) + h_indices = torch.arange(nb_patches_h, device=pixel_values.device, dtype=pixel_values.dtype) + w_indices = torch.arange(nb_patches_w, device=pixel_values.device, dtype=pixel_values.dtype) + + fractional_coords_h = h_indices / nb_patches_h * (1 - 1e-6) + fractional_coords_w = w_indices / nb_patches_w * (1 - 1e-6) bucket_coords_h = torch.bucketize(fractional_coords_h, boundaries, right=True) bucket_coords_w = torch.bucketize(fractional_coords_w, boundaries, right=True) @@ -72,12 +74,36 @@ def forward(self, pixel_values: torch.FloatTensor, patch_attention_mask: torch.B class GaudiIdefics2Model(Idefics2Model): + def inputs_merger( + self, + input_ids: torch.LongTensor, + inputs_embeds: Optional[torch.Tensor], + image_hidden_states: Optional[torch.Tensor], + ): + """ + Inherits from Idefics2Model::inputs_merger https://github.com/huggingface/transformers/blob/v4.43.4/src/transformers/models/idefics2/modeling_idefics2.py#L1268 + The only differences are: + - replace `==` with torch.where to fix the issue in hpu graph + """ + if input_ids is None: + special_image_mask = torch.where(inputs_embeds == self.get_input_embeddings()( + torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device) + )) + special_image_mask = special_image_mask.all(-1) + else: + special_image_mask = torch.where(input_ids == self.config.image_token_id) + + special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device) + image_hidden_states = image_hidden_states.to(inputs_embeds.device, inputs_embeds.dtype) + inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_hidden_states) + return inputs_embeds + def forward( self, input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, + past_key_values: Optional[list[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, pixel_values: Optional[torch.FloatTensor] = None, pixel_attention_mask: Optional[torch.BoolTensor] = None, @@ -87,7 +113,8 @@ def forward( output_hidden_states: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, return_dict: Optional[bool] = None, - ) -> Union[Tuple, Idefics2BaseModelOutputWithPast]: + **kwargs, + ) -> Union[tuple, Idefics2BaseModelOutputWithPast]: """ Inherits from Idefics2Model::forward https://github.com/huggingface/transformers/blob/v4.43.4/src/transformers/models/idefics2/modeling_idefics2.py#L1303 The only differences are: @@ -182,7 +209,7 @@ def forward( elif image_hidden_states is not None: image_hidden_states = image_hidden_states.to(dtype=self.dtype, device=input_ids.device) - if past_seen_tokens == 0 and inputs_embeds is not None and image_hidden_states is not None: + if image_hidden_states is not None: # When we generate, we don't want to replace the potential image_token_id that we generated by images # that simply don't exist inputs_embeds = self.inputs_merger( @@ -200,7 +227,8 @@ def forward( output_attentions=output_attentions, output_hidden_states=output_hidden_states, cache_position=cache_position, - return_dict=return_dict, + return_dict=True, + **kwargs, ) if return_legacy_cache and use_cache: @@ -224,24 +252,6 @@ def forward( image_hidden_states=image_hidden_states, ) - def inputs_merger( - self, - input_ids: torch.LongTensor, - inputs_embeds: Optional[torch.Tensor], - image_hidden_states: Optional[torch.Tensor], - ): - """ - Inherits from Idefics2Model::inputs_merger https://github.com/huggingface/transformers/blob/v4.43.4/src/transformers/models/idefics2/modeling_idefics2.py#L1268 - The only differences are: - - replace `==` with torch.where to fix the issue in hpu graph - """ - num_images, _, vision_hidden_size = image_hidden_states.shape - special_image_token_mask = torch.where(input_ids == self.image_token_id) - new_inputs_embeds = inputs_embeds.clone() - reshaped_image_hidden_states = image_hidden_states.view(-1, vision_hidden_size) - new_inputs_embeds[special_image_token_mask] = reshaped_image_hidden_states.to(new_inputs_embeds.device) - return new_inputs_embeds - class GaudiIdefics2ForConditionalGeneration(Idefics2ForConditionalGeneration): def forward( @@ -249,7 +259,7 @@ def forward( input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, + past_key_values: Optional[list[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, pixel_values: Optional[torch.FloatTensor] = None, pixel_attention_mask: Optional[torch.BoolTensor] = None, @@ -262,7 +272,8 @@ def forward( cache_position: Optional[torch.LongTensor] = None, logits_to_keep: Union[int, torch.Tensor] = 0, token_idx: Optional[torch.Tensor] = None, - ) -> Union[Tuple, Idefics2CausalLMOutputWithPast]: + **kwargs, + ) -> Union[tuple, Idefics2CausalLMOutputWithPast]: """ Inherits from Idefics2ForConditionalGeneration::forward https://github.com/huggingface/transformers/blob/v4.43.4/src/transformers/models/idefics2/modeling_idefics2.py#L1505 The only differences are: @@ -344,22 +355,9 @@ def forward( loss = None if labels is not None: - labels = labels.to(logits.device) - # Shift so that tokens < n predict n - if attention_mask is not None: - shift_attention_mask = attention_mask[:, -(logits.shape[1] - 1) :].to(logits.device) - shift_logits = logits[..., :-1, :][shift_attention_mask != 0].contiguous() - shift_labels = labels[..., 1:][shift_attention_mask != 0].contiguous() - else: - shift_logits = logits[..., :-1, :].contiguous() - shift_labels = labels[..., 1:].contiguous() - # Flatten the tokens - loss_fct = CrossEntropyLoss() - loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)) - - if not return_dict: - output = (logits,) + outputs[1:] - return (loss,) + output if loss is not None else output + loss = self.loss_function( + logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size, **kwargs + ) return Idefics2CausalLMOutputWithPast( loss=loss, @@ -385,6 +383,7 @@ def forward( output_hidden_states=output_hidden_states, cache_position=cache_position, return_dict=return_dict, + **kwargs, ) def prepare_inputs_for_generation( diff --git a/optimum/habana/transformers/models/llama/modeling_llama.py b/optimum/habana/transformers/models/llama/modeling_llama.py index 75d8097cc8..0562f9630f 100755 --- a/optimum/habana/transformers/models/llama/modeling_llama.py +++ b/optimum/habana/transformers/models/llama/modeling_llama.py @@ -1,6 +1,5 @@ import copy -from functools import partial -from typing import List, Optional, Tuple, Union +from typing import Optional, Union import torch from torch.distributed.distributed_c10d import ProcessGroup @@ -9,7 +8,6 @@ from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS from transformers.models.llama.modeling_llama import ( - KwargsForCausalLM, LlamaAttention, LlamaDecoderLayer, LlamaForCausalLM, @@ -17,9 +15,9 @@ LlamaModel, LlamaRMSNorm, apply_rotary_pos_emb, - logger, ) from transformers.processing_utils import Unpack +from transformers.utils import TransformersKwargs from .... import distributed from ....distributed import parallel_state @@ -86,7 +84,7 @@ def __init__(self, config: LlamaConfig, device=None): super().__init__() # BC: "rope_type" was originally "type" - if hasattr(config, "rope_scaling") and config.rope_scaling is not None: + if hasattr(config, "rope_scaling") and isinstance(config.rope_scaling, dict): self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type")) else: self.rope_type = "default" @@ -219,10 +217,10 @@ def __init__( GaudiLlamaMLP.__init__(self, self.config) self.setup_tp(rank, world_size) - def colwise_param_names(self) -> List[str]: + def colwise_param_names(self) -> list[str]: return ["up_proj", "gate_proj"] - def rowwise_param_names(self) -> List[str]: + def rowwise_param_names(self) -> list[str]: return ["down_proj"] @staticmethod @@ -559,7 +557,7 @@ def reorder_kv_cache(self, beam_idx: torch.LongTensor): def pre_attn_forward( self, hidden_states: torch.Tensor, - position_embeddings: Tuple[torch.Tensor, torch.Tensor], + position_embeddings: tuple[torch.Tensor, torch.Tensor], attention_mask: Optional[torch.Tensor], past_key_value: Optional[Cache] = None, output_attentions: bool = False, @@ -577,7 +575,7 @@ def pre_attn_forward( cache_idx: int = None, num_virtual_tokens: int = None, **kwargs, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]: """ Copied from LlamaAttention.forward: https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py The only differences are: @@ -631,7 +629,7 @@ def pre_attn_forward( # logger.warning_once( # "The attention layers in this model are transitioning from computing the RoPE embeddings internally " # "through `position_ids` (2D tensor with the indexes of the tokens), to using externally computed " - # "`position_embeddings` (Tuple of tensors, containing cos and sin). In v4.46 `position_ids` will be " + # "`position_embeddings` (tuple of tensors, containing cos and sin). In v4.46 `position_ids` will be " # "removed and `position_embeddings` will be mandatory." # ) # cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len) @@ -843,14 +841,14 @@ def __init__( ) self.setup_tp(rank, world_size) - def colwise_param_names(self) -> List[str]: + def colwise_param_names(self) -> list[str]: colwise_weights = ["q_proj"] if self.pre_tp_kvheads != 1: colwise_weights.append("k_proj") colwise_weights.append("v_proj") return colwise_weights - def rowwise_param_names(self) -> List[str]: + def rowwise_param_names(self) -> list[str]: return ["o_proj"] @staticmethod @@ -867,7 +865,7 @@ def pre_attn_forward( output_attentions: bool = False, use_cache: bool = False, cache_position: Optional[torch.LongTensor] = None, - position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # will become mandatory in v4.45 + position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, # will become mandatory in v4.45 token_idx: Optional[torch.Tensor] = None, attn_softmax_bf16: Optional[bool] = False, reuse_cache: Optional[bool] = False, @@ -877,7 +875,7 @@ def pre_attn_forward( flash_attention_fast_softmax: Optional[bool] = False, cache_idx: int = None, **kwargs, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]: hidden_states, attn_weights, present_key_value = GaudiLlamaAttention.pre_attn_forward( self, hidden_states=hidden_states, @@ -928,10 +926,9 @@ def forward( attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Cache] = None, - output_attentions: Optional[bool] = False, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, - position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC + position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC token_idx: Optional[torch.Tensor] = None, attn_softmax_bf16: Optional[bool] = False, reuse_cache: Optional[bool] = False, @@ -945,7 +942,7 @@ def forward( attn_batch_split: int = 1, prev_layer_residual: Optional[torch.Tensor] = None, **kwargs, - ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]: """ Copied from LlamaDecoderLayer.forward: https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py The only differences are: @@ -986,7 +983,6 @@ def forward( attention_mask=sub_attention_mask[i], position_ids=sub_position_ids[i], past_key_value=past_key_value, - output_attentions=output_attentions, use_cache=use_cache, cache_position=cache_position, position_embeddings=position_embeddings, @@ -1003,8 +999,6 @@ def forward( **kwargs, ) self.self_attn.attention_all_reduce(split_hidden_states[i]) - if output_attentions: - split_attn_weights.append(self_attn_weights) if use_cache: split_present_key_values.append(present_key_value) @@ -1030,7 +1024,6 @@ def forward( attention_mask=attention_mask, position_ids=position_ids, past_key_value=past_key_value, - output_attentions=output_attentions, use_cache=use_cache, cache_position=cache_position, position_embeddings=position_embeddings, @@ -1052,8 +1045,6 @@ def forward( hidden_states = self.post_mlp(hidden_states, residual) outputs = (hidden_states,) - if output_attentions: - outputs += (self_attn_weights,) if use_cache: outputs += (present_key_value,) # Store the residual splits to add them in the beginning of the next layer @@ -1067,11 +1058,11 @@ def pre_attn( hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_value: Optional[Tuple[torch.Tensor]] = None, + past_key_value: Optional[tuple[torch.Tensor]] = None, output_attentions: Optional[bool] = False, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, - position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # will become mandatory in v4.45 + position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, # will become mandatory in v4.45 token_idx: Optional[torch.Tensor] = None, attn_softmax_bf16: Optional[bool] = False, reuse_cache: Optional[bool] = False, @@ -1082,7 +1073,7 @@ def pre_attn( valid_sequence_lengths: Optional[torch.Tensor] = None, cache_idx: int = None, num_virtual_tokens: int = None, - ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]: hidden_states = self.input_layernorm(hidden_states) hidden_states, attn_weights, present_key_value = self.self_attn.pre_attn_forward( hidden_states=hidden_states, @@ -1181,12 +1172,10 @@ def forward( input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, + past_key_values: Optional[Union[Cache, list[torch.FloatTensor]]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, - use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, token_idx: Optional[torch.Tensor] = None, attn_softmax_bf16: Optional[bool] = False, reuse_cache: Optional[bool] = False, @@ -1213,10 +1202,6 @@ def forward( - add new arg flash_attention_fast_softmax - add new arg lazy_mode """ - 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 - ) use_cache = use_cache if use_cache is not None else self.config.use_cache if (input_ids is None) ^ (inputs_embeds is not None): @@ -1235,14 +1220,8 @@ def forward( global has_fused_rms_norm has_fused_rms_norm = False - if self.gradient_checkpointing and self.training and use_cache: - logger.warning_once( - "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`." - ) - use_cache = False - if inputs_embeds is None: - inputs_embeds = self.embed_tokens(input_ids) + inputs_embeds: torch.Tensor = self.embed_tokens(input_ids) ignore_cache_position = True # Ignoring cache position for HPU use_new_cache = False # Ignoring new Cache path for HPU @@ -1290,15 +1269,9 @@ def forward( else: causal_mask = self._update_causal_mask(attention_mask, inputs_embeds, cache_position, past_seen_tokens) - # embed positions hidden_states = inputs_embeds - - # create position embeddings to be shared across the decoder layers position_embeddings = None # self.rotary_emb(hidden_states, position_ids) - # decoder layers - all_hidden_states = () if output_hidden_states else None - all_self_attns = () if output_attentions else None next_decoder_cache = () if not use_new_cache else None if lazy_mode: @@ -1324,81 +1297,47 @@ def forward( ): htcore.mark_step() - if output_hidden_states: - all_hidden_states += (hidden_states,) - - if self.gradient_checkpointing and self.training: - layer_outputs = self._gradient_checkpointing_func( - partial(decoder_layer.__call__, **kwargs), - hidden_states, - causal_mask, - position_ids, - past_key_values, - output_attentions, - use_cache, - cache_position, - position_embeddings, - None, - attn_softmax_bf16, - False, - use_flash_attention, - flash_attention_recompute, - flash_attention_causal_mask, - flash_attention_fast_softmax, - valid_sequence_lengths, - None, - ) - hidden_states = layer_outputs[0] + # Calling the layer with positional arguments + # This is a workaround for an issue with DeepSpeed where + # it cannot handle keyword arguments and throws a RuntimError + use_prev_layer_residual = attn_batch_split > 1 and past_key_values is None + layer_prev_layer_residual = prev_layer_residual if use_prev_layer_residual else None + layer_hidden_states = hidden_states_split if split_prompt else hidden_states + past_key_value = None if past_key_values is None else past_key_values[layer_idx] + layer_outputs = decoder_layer( + layer_hidden_states, + causal_mask, + position_ids, + past_key_value, + use_cache, + cache_position, + position_embeddings, + token_idx, + attn_softmax_bf16, + reuse_cache, + use_flash_attention, + flash_attention_recompute, + flash_attention_causal_mask, + flash_attention_fast_softmax, + valid_sequence_lengths, + cache_idx, + num_virtual_tokens, + attn_batch_split, + layer_prev_layer_residual, + ) + if use_prev_layer_residual: + index = 1 + int(use_cache) + prev_layer_residual = layer_outputs[index] + if split_prompt: + hidden_states_split = layer_outputs[0] else: - # Calling the layer with positional arguments - # This is a workaround for an issue with DeepSpeed where - # it cannot handle keyword arguments and throws a RuntimError - use_prev_layer_residual = attn_batch_split > 1 and past_key_values is None - layer_prev_layer_residual = prev_layer_residual if use_prev_layer_residual else None - layer_hidden_states = hidden_states_split if split_prompt else hidden_states - past_key_value = None if past_key_values is None else past_key_values[layer_idx] - layer_outputs = decoder_layer( - layer_hidden_states, - causal_mask, - position_ids, - past_key_value, - output_attentions, - use_cache, - cache_position, - position_embeddings, - token_idx, - attn_softmax_bf16, - reuse_cache, - use_flash_attention, - flash_attention_recompute, - flash_attention_causal_mask, - flash_attention_fast_softmax, - valid_sequence_lengths, - cache_idx, - num_virtual_tokens, - attn_batch_split, - layer_prev_layer_residual, - ) - if use_prev_layer_residual: - index = 1 + int(use_cache) + int(output_attentions) - prev_layer_residual = layer_outputs[index] - if split_prompt: - hidden_states_split = layer_outputs[0] - else: - hidden_states = layer_outputs[0] + hidden_states = layer_outputs[0] if use_cache: - next_decoder_cache += (layer_outputs[2 if output_attentions else 1],) - - if output_attentions: - all_self_attns += (layer_outputs[1],) + next_decoder_cache += (layer_outputs[1],) hidden_states = self.norm(hidden_states) - # add hidden states from the last decoder layer - if output_hidden_states: - all_hidden_states += (hidden_states,) - next_cache = next_decoder_cache if use_cache else None if not use_new_cache and isinstance(next_cache, Cache): next_cache = next_cache.to_legacy_cache() @@ -1406,8 +1345,6 @@ def forward( return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=next_cache, - hidden_states=all_hidden_states, - attentions=all_self_attns, ) @@ -1445,12 +1382,10 @@ def forward( input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, + past_key_values: Optional[Union[Cache, list[torch.FloatTensor]]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, logits_to_keep: Union[int, torch.Tensor] = 0, token_idx: Optional[torch.Tensor] = None, @@ -1466,17 +1401,12 @@ def forward( lazy_mode: Optional[bool] = True, num_virtual_tokens: int = None, attn_batch_split: int = 1, - **kwargs: Unpack[KwargsForCausalLM], + **kwargs: Unpack[TransformersKwargs], ) -> CausalLMOutputWithPast: - 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 - ) if self.generation_config.use_fused_rope is False: global has_fused_rope has_fused_rope = False - # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) outputs: BaseModelOutputWithPast = self.model( input_ids=input_ids, attention_mask=attention_mask, @@ -1484,8 +1414,6 @@ def forward( past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, cache_position=cache_position, token_idx=token_idx, attn_softmax_bf16=attn_softmax_bf16, @@ -1528,8 +1456,8 @@ def forward( @staticmethod def _reorder_cache( - past: Tuple[Tuple[torch.Tensor, torch.Tensor], ...], beam_idx: torch.LongTensor - ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], ...]: + past: tuple[tuple[torch.Tensor, torch.Tensor], ...], beam_idx: torch.LongTensor + ) -> tuple[tuple[torch.Tensor, torch.Tensor], ...]: """ This function is used to re-order the `past_key_values` cache if [`~PreTrainedModel.beam_search`] or [`~PreTrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct diff --git a/optimum/habana/transformers/models/llava/modeling_llava.py b/optimum/habana/transformers/models/llava/modeling_llava.py index 42851c9890..5c14251b1e 100644 --- a/optimum/habana/transformers/models/llava/modeling_llava.py +++ b/optimum/habana/transformers/models/llava/modeling_llava.py @@ -19,7 +19,7 @@ # limitations under the License. """PyTorch Llava model.""" -from typing import List, Optional, Tuple, Union +from typing import Optional, Union import torch import torch.nn as nn @@ -108,13 +108,13 @@ def _merge_input_ids_with_image_features(image_features, inputs_embeds, input_id class GaudiLlavaForConditionalGeneration(LlavaForConditionalGeneration): def forward( self, - input_ids: Optional[torch.LongTensor] = None, - pixel_values: Optional[torch.FloatTensor] = None, + input_ids: torch.LongTensor = None, + pixel_values: torch.FloatTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, + past_key_values: Optional[list[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, - vision_feature_layer: Optional[Union[int, List[int]]] = None, + vision_feature_layer: Optional[Union[int, list[int]]] = None, vision_feature_select_strategy: Optional[str] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, @@ -129,8 +129,8 @@ def forward( tokens_pos: Optional[torch.LongTensor] = None, use_flash_attention: Optional[bool] = False, flash_attention_recompute: Optional[bool] = False, - **lm_kwargs, - ) -> Union[Tuple, LlavaCausalLMOutputWithPast]: + **kwargs, + ) -> Union[tuple, LlavaCausalLMOutputWithPast]: """ Inherits from LlavaForConditionalGeneration: https://github.com/huggingface/transformers/blob/v4.45.2/src/transformers/models/llava/modeling_llava.py#L362 The only differences are: @@ -138,7 +138,6 @@ def forward( - add new args image_offset - add new args tokens_pos """ - 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 @@ -203,7 +202,7 @@ def forward( token_idx=token_idx + image_offset, use_flash_attention=use_flash_attention, flash_attention_recompute=flash_attention_recompute, - **lm_kwargs, + **kwargs, ) if input_ids.shape[1] != 1 and pixel_values is not None and tokens_pos is not None: @@ -242,7 +241,7 @@ def forward( logits_to_keep=logits_to_keep, use_flash_attention=use_flash_attention, flash_attention_recompute=flash_attention_recompute, - **lm_kwargs, + **kwargs, ) logits = outputs[0] diff --git a/optimum/habana/transformers/models/llava_next/modeling_llava_next.py b/optimum/habana/transformers/models/llava_next/modeling_llava_next.py index 410bfd9b2e..b7d46600d1 100644 --- a/optimum/habana/transformers/models/llava_next/modeling_llava_next.py +++ b/optimum/habana/transformers/models/llava_next/modeling_llava_next.py @@ -19,7 +19,7 @@ # limitations under the License. """PyTorch Llava-NeXT model.""" -from typing import List, Optional, Tuple, Union +from typing import Optional, Union import torch import torch.utils.checkpoint @@ -39,27 +39,26 @@ class GaudiLlavaNextForConditionalGeneration(LlavaNextForConditionalGeneration): def forward( self, - input_ids: Optional[torch.LongTensor] = None, - pixel_values: Optional[torch.FloatTensor] = None, + input_ids: torch.LongTensor = None, + pixel_values: torch.FloatTensor = None, image_sizes: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, + past_key_values: Optional[list[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, - vision_feature_layer: Optional[Union[int, List[int]]] = None, + vision_feature_layer: Optional[Union[int, list[int]]] = None, vision_feature_select_strategy: Optional[str] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, logits_to_keep: Union[int, torch.Tensor] = 0, token_idx: Optional[torch.Tensor] = None, use_flash_attention: Optional[bool] = False, flash_attention_recompute: Optional[bool] = False, - **lm_kwargs, - ) -> Union[Tuple, LlavaNextCausalLMOutputWithPast]: + **kwargs, + ) -> Union[tuple, LlavaNextCausalLMOutputWithPast]: """ Inherits from LlavaForConditionalGeneration: https://github.com/huggingface/transformers/blob/v4.40.0/src/transformers/models/llava_next/modeling_llava_next.py#L433 The only differences are: @@ -74,7 +73,6 @@ def forward( 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 if inputs_embeds is None: inputs_embeds = self.get_input_embeddings()(input_ids) @@ -86,14 +84,14 @@ def forward( use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, - return_dict=return_dict, + return_dict=True, cache_position=cache_position, # TODO: from Transformers v4.45, `generate` sets `num_logits_to_keep` to 1 if not given, which we don't want here # logits_to_keep=logits_to_keep, token_idx=token_idx + self.image_offset, use_flash_attention=use_flash_attention, flash_attention_recompute=flash_attention_recompute, - **lm_kwargs, + **kwargs, ) if inputs_embeds.shape[1] != 1 and pixel_values is not None and self.text_tokens_pos is not None: @@ -123,10 +121,6 @@ def forward( shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1).to(shift_logits.device) ) - if not return_dict: - output = (logits,) + outputs[1:] - return (loss,) + output if loss is not None else output - return LlavaNextCausalLMOutputWithPast( loss=loss, logits=logits, @@ -150,10 +144,9 @@ def forward( use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, - return_dict=return_dict, cache_position=cache_position, logits_to_keep=logits_to_keep, - **lm_kwargs, + **kwargs, ) # Copied from https://github.com/huggingface/transformers/blob/v4.40.0/src/transformers/models/llava_next/modeling_llava_next.py#L356 diff --git a/optimum/habana/transformers/models/llava_onevision/modeling_llava_onevision.py b/optimum/habana/transformers/models/llava_onevision/modeling_llava_onevision.py index 50986e0ec2..83c608b686 100644 --- a/optimum/habana/transformers/models/llava_onevision/modeling_llava_onevision.py +++ b/optimum/habana/transformers/models/llava_onevision/modeling_llava_onevision.py @@ -15,7 +15,7 @@ """PyTorch Llava-Onevision model.""" # Portions copied from optimum.habana.transformers.models.llava_next and transformers.models.llava_onevision -from typing import List, Optional, Tuple, Union +from typing import Optional, Union import torch import torch.utils.checkpoint @@ -41,11 +41,12 @@ def forward( image_sizes_videos: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, + past_key_values: Optional[list[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, - vision_feature_layer: Optional[int] = None, + vision_feature_layer: Optional[Union[int, list[int]]] = None, vision_feature_select_strategy: Optional[str] = None, vision_aspect_ratio: Optional[str] = None, + batch_num_images: Optional[torch.LongTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, @@ -56,7 +57,8 @@ def forward( token_idx: Optional[torch.Tensor] = None, use_flash_attention: Optional[bool] = False, flash_attention_recompute: Optional[bool] = False, - ) -> Union[Tuple, LlavaOnevisionCausalLMOutputWithPast]: + **kwargs, + ) -> Union[tuple, LlavaOnevisionCausalLMOutputWithPast]: """ Inherits from LlavaForConditionalGeneration: https://github.com/huggingface/transformers/blob/v4.40.0/src/transformers/models/llava_next/modeling_llava_next.py#L433 The only differences are: @@ -151,6 +153,7 @@ def forward( vision_feature_layer=vision_feature_layer, vision_feature_select_strategy=vision_feature_select_strategy, vision_aspect_ratio=vision_aspect_ratio, + batch_num_images=batch_num_images, labels=labels, use_cache=use_cache, output_attentions=output_attentions, diff --git a/optimum/habana/transformers/models/mamba/modeling_mamba.py b/optimum/habana/transformers/models/mamba/modeling_mamba.py index e23ce65dd8..db07b00abd 100644 --- a/optimum/habana/transformers/models/mamba/modeling_mamba.py +++ b/optimum/habana/transformers/models/mamba/modeling_mamba.py @@ -1,4 +1,4 @@ -from typing import Any, Dict, Optional +from typing import Any, Optional import torch from transformers.models.mamba.modeling_mamba import ( @@ -14,8 +14,8 @@ def gaudi_MambaForCausalLM_update_model_kwargs_for_generation( - self, outputs: ModelOutput, model_kwargs: Dict[str, Any], num_new_tokens: int = 1, **kwargs -) -> Dict[str, Any]: + self, outputs: ModelOutput, model_kwargs: dict[str, Any], num_new_tokens: int = 1, **kwargs +) -> dict[str, Any]: model_kwargs["cache_params"] = outputs.get("cache_params", None) token_idx = model_kwargs.get("token_idx", None) if ( diff --git a/optimum/habana/transformers/models/mistral/modeling_mistral.py b/optimum/habana/transformers/models/mistral/modeling_mistral.py index 4c2f469f60..8d8401b3bc 100644 --- a/optimum/habana/transformers/models/mistral/modeling_mistral.py +++ b/optimum/habana/transformers/models/mistral/modeling_mistral.py @@ -20,8 +20,7 @@ """PyTorch Mistral model.""" import os -from functools import partial -from typing import List, Optional, Tuple, Union +from typing import Optional, Union import habana_frameworks.torch.core as htcore import torch @@ -29,7 +28,6 @@ from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from transformers.models.mistral.configuration_mistral import MistralConfig from transformers.models.mistral.modeling_mistral import ( - KwargsForCausalLM, MistralAttention, MistralDecoderLayer, MistralForCausalLM, @@ -39,7 +37,7 @@ apply_rotary_pos_emb, ) from transformers.processing_utils import Unpack -from transformers.utils import logging +from transformers.utils import TransformersKwargs, logging from ...modeling_attn_mask_utils import ( _gaudi_prepare_4d_causal_attention_mask, @@ -215,7 +213,7 @@ def reorder_kv_cache(self, beam_idx: torch.LongTensor): def forward( self, hidden_states: torch.Tensor, - position_embeddings: Tuple[torch.Tensor, torch.Tensor], + position_embeddings: tuple[torch.Tensor, torch.Tensor], attention_mask: Optional[torch.Tensor], past_key_value: Optional[Cache] = None, use_cache: bool = False, @@ -228,7 +226,7 @@ def forward( flash_attention_recompute: Optional[bool] = False, flash_attention_causal_mask: Optional[bool] = False, **kwargs, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]: """ Copied from MistralAttention.forward: https://github.com/huggingface/transformers/blob/v4.34.1/src/transformers/models/mistral/modeling_mistral.py The only differences are: @@ -365,10 +363,9 @@ def forward( attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Cache] = None, - output_attentions: Optional[bool] = False, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, - position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC + position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC token_idx: Optional[torch.Tensor] = None, reuse_cache: Optional[bool] = False, cache_idx: Optional[int] = None, @@ -377,7 +374,7 @@ def forward( flash_attention_recompute: Optional[bool] = False, flash_attention_causal_mask: Optional[bool] = False, **kwargs, - ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]: """ Copied from MistralDecoderLayer.forward: https://github.com/huggingface/transformers/blob/v4.34.1/src/transformers/models/mistral/modeling_mistral.py The only differences are: @@ -393,7 +390,6 @@ def forward( attention_mask=attention_mask, position_ids=position_ids, past_key_value=past_key_value, - output_attentions=output_attentions, use_cache=use_cache, cache_position=cache_position, position_embeddings=position_embeddings, @@ -416,9 +412,6 @@ def forward( outputs = (hidden_states,) - if output_attentions: - outputs += (self_attn_weights,) - if use_cache: outputs += (present_key_value,) @@ -442,11 +435,9 @@ def forward( input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, + past_key_values: Optional[Union[Cache, list[torch.FloatTensor]]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, reuse_cache: Optional[bool] = False, @@ -464,10 +455,6 @@ def forward( - add new args token_idx - add new arg lazy_mode """ - 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 - ) use_cache = use_cache if use_cache is not None else self.config.use_cache if input_ids is not None and inputs_embeds is not None: @@ -479,12 +466,6 @@ def forward( else: raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds") - if self.gradient_checkpointing and self.training and use_cache: - logger.warning_once( - "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`." - ) - use_cache = False - past_key_values_length = 0 use_new_cache = False return_legacy_cache = False @@ -520,9 +501,6 @@ def forward( hidden_states = inputs_embeds - # decoder layers - all_hidden_states = () if output_hidden_states else None - all_self_attns = () if output_attentions else None next_decoder_cache = () if not use_new_cache else None if lazy_mode: @@ -535,59 +513,31 @@ def forward( and (torch.distributed.is_initialized() is False or torch.distributed.get_world_size() == 1) ): htcore.mark_step() - if output_hidden_states: - all_hidden_states += (hidden_states,) - - if self.gradient_checkpointing and self.training: - layer_outputs = self._gradient_checkpointing_func( - partial(decoder_layer.__call__, **kwargs), - hidden_states, - causal_mask, - position_ids, - None if past_key_values is None else past_key_values[layer_idx], - output_attentions, - use_cache, - cache_position, - None, - False, - cache_idx, - attn_softmax_bf16, - use_flash_attention, - flash_attention_recompute, - flash_attention_causal_mask, - ) - else: - layer_outputs = decoder_layer( - hidden_states, - attention_mask=causal_mask, - position_ids=position_ids, - past_key_value=None if past_key_values is None else past_key_values[layer_idx], - output_attentions=output_attentions, - use_cache=use_cache, - cache_position=cache_position, - token_idx=token_idx, - reuse_cache=reuse_cache, - cache_idx=cache_idx, - attn_softmax_bf16=attn_softmax_bf16, - use_flash_attention=use_flash_attention, - flash_attention_recompute=flash_attention_recompute, - flash_attention_causal_mask=flash_attention_causal_mask, - ) + + layer_outputs = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_ids=position_ids, + past_key_value=None if past_key_values is None else past_key_values[layer_idx], + use_cache=use_cache, + cache_position=cache_position, + token_idx=token_idx, + reuse_cache=reuse_cache, + cache_idx=cache_idx, + attn_softmax_bf16=attn_softmax_bf16, + use_flash_attention=use_flash_attention, + flash_attention_recompute=flash_attention_recompute, + flash_attention_causal_mask=flash_attention_causal_mask, + **kwargs, + ) hidden_states = layer_outputs[0] if use_cache: - next_decoder_cache += (layer_outputs[2 if output_attentions else 1],) - - if output_attentions: - all_self_attns += (layer_outputs[1],) + next_decoder_cache += (layer_outputs[1],) hidden_states = self.norm(hidden_states) - # add hidden states from the last decoder layer - if output_hidden_states: - all_hidden_states += (hidden_states,) - next_cache = None if use_cache: next_cache = ( @@ -599,8 +549,6 @@ def forward( return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=next_cache, - hidden_states=all_hidden_states, - attentions=all_self_attns, ) @@ -619,12 +567,10 @@ def forward( input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, + past_key_values: Optional[Union[Cache, list[torch.FloatTensor]]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, logits_to_keep: Union[int, torch.Tensor] = 0, token_idx: Optional[torch.Tensor] = None, @@ -636,24 +582,17 @@ def forward( use_flash_attention: Optional[bool] = False, flash_attention_recompute: Optional[bool] = False, flash_attention_causal_mask: Optional[bool] = False, - **kwargs: Unpack[KwargsForCausalLM], + **kwargs: Unpack[TransformersKwargs], ) -> CausalLMOutputWithPast: """ Inherits from MistralForCausalLM: https://github.com/huggingface/transformers/blob/v4.34.1/src/transformers/models/mistral/modeling_mistral.py The only differences are: - add new args token_idx """ - - 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 - ) - if self.generation_config.use_fused_rope is False: global has_fused_rope has_fused_rope = False - # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) outputs: BaseModelOutputWithPast = self.model( input_ids=input_ids, attention_mask=attention_mask, @@ -661,8 +600,6 @@ def forward( past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, cache_position=cache_position, token_idx=token_idx, reuse_cache=reuse_cache, @@ -672,7 +609,9 @@ def forward( use_flash_attention=use_flash_attention, flash_attention_recompute=flash_attention_recompute, flash_attention_causal_mask=flash_attention_causal_mask, + **kwargs, ) + hidden_states = outputs.last_hidden_state _, seq_len, _ = hidden_states.shape if seq_len > 1 and trim_logits and not self.training: diff --git a/optimum/habana/transformers/models/mixtral/modeling_mixtral.py b/optimum/habana/transformers/models/mixtral/modeling_mixtral.py index f31e7e4760..dab8e3f932 100644 --- a/optimum/habana/transformers/models/mixtral/modeling_mixtral.py +++ b/optimum/habana/transformers/models/mixtral/modeling_mixtral.py @@ -21,8 +21,7 @@ """PyTorch Mixtral model.""" import math -from functools import partial -from typing import List, Optional, Tuple, Union +from typing import Optional, Union import habana_frameworks.torch.core as htcore import torch @@ -40,7 +39,6 @@ MoeModelOutputWithPast, ) from transformers.models.mixtral.modeling_mixtral import ( - KwargsForCausalLM, MixtralAttention, MixtralBlockSparseTop2MLP, MixtralDecoderLayer, @@ -50,7 +48,7 @@ load_balancing_loss_func, ) from transformers.processing_utils import Unpack -from transformers.utils import logging +from transformers.utils import TransformersKwargs, logging from ....distributed.tensorparallel import _all_reduce from ..llama.modeling_llama import GaudiLlamaRotaryEmbedding @@ -218,7 +216,7 @@ def __init__(self, config): # Jitter parameters self.jitter_noise = config.router_jitter_noise - def forward(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: original_shape = hidden_states.shape hidden_dim = original_shape[2] if self.training and self.jitter_noise > 0: @@ -389,7 +387,7 @@ def allocate_kv_cache(self, batch_size, max_seq_len, inp_seq_len): def forward( self, hidden_states: torch.Tensor, - position_embeddings: Tuple[torch.Tensor, torch.Tensor], + position_embeddings: tuple[torch.Tensor, torch.Tensor], attention_mask: Optional[torch.Tensor], past_key_value: Optional[Cache] = None, use_cache: bool = False, @@ -399,7 +397,7 @@ def forward( flash_attention_recompute: Optional[bool] = False, cache_idx: int = None, **kwargs, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]: """ Copied from MixtralAttention.forward: https://github.com/huggingface/transformers/blob/v4.37.0/src/transformers/models/mixtral/modeling_mixtral.py The only differences are: @@ -519,7 +517,7 @@ def forward( def calculate_routing_tensors( score: torch.Tensor, topk: int, hidden_states_dtype: torch.dtype -) -> Tuple[torch.Tensor, torch.Tensor]: +) -> tuple[torch.Tensor, torch.Tensor]: """Based on https://github.com/huggingface/transformers/blob/main/src/transformers/models/mixtral/modeling_mixtral.py#L641""" routing_weights = F.softmax(score, dim=1, dtype=torch.float32) routing_weights, selected_experts = torch.topk(routing_weights, topk, dim=-1) @@ -541,18 +539,16 @@ def forward( hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_value: Optional[Tuple[torch.Tensor]] = None, - output_attentions: Optional[bool] = False, - output_router_logits: Optional[bool] = False, + past_key_value: Optional[tuple[torch.Tensor]] = None, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, - position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC + position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC token_idx: Optional[torch.Tensor] = None, reuse_cache: Optional[bool] = False, flash_attention_recompute: Optional[bool] = False, cache_idx: int = None, **kwargs, - ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]: """ Copied from MixtralDecoderLayer.forward: https://github.com/huggingface/transformers/blob/v4.37.0/src/transformers/models/mixtral/modeling_mixtral.py The only differences are: @@ -572,7 +568,6 @@ def forward( attention_mask=attention_mask, position_ids=position_ids, past_key_value=past_key_value, - output_attentions=output_attentions, use_cache=use_cache, cache_position=cache_position, token_idx=token_idx, @@ -590,15 +585,9 @@ def forward( outputs = (hidden_states,) - if output_attentions: - outputs += (self_attn_weights,) - if use_cache: outputs += (present_key_value,) - if output_router_logits: - outputs += (router_logits,) - return outputs @@ -612,12 +601,9 @@ def forward( input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, + past_key_values: Optional[list[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - output_router_logits: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, reuse_cache: Optional[bool] = False, @@ -633,13 +619,6 @@ def forward( - add new args flash_attention_recompute - add new args cache_idx """ - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_router_logits = ( - output_router_logits if output_router_logits is not None else self.config.output_router_logits - ) - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) use_cache = use_cache if use_cache is not None else self.config.use_cache # retrieve input_ids and inputs_embeds @@ -655,13 +634,6 @@ def forward( past_key_values_length = 0 use_new_cache = False # Ignoring new Cache path for HPU - if self.gradient_checkpointing and self.training: - if use_cache: - logger.warning_once( - "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." - ) - use_cache = False - if past_key_values is not None and use_cache: if reuse_cache: past_key_values_length = past_key_values[0][0][2] @@ -694,7 +666,7 @@ def forward( if self.config._attn_implementation == "flash_attention_2": # 2d mask is passed through the layers attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None - elif self.config._attn_implementation == "sdpa" and not output_attentions: + elif self.config._attn_implementation == "sdpa": # output_attentions=True can not be supported when using SDPA, and we fall back on # the manual implementation that requires a 4D causal mask in all cases. attention_mask = _prepare_4d_causal_attention_mask_for_sdpa( @@ -716,62 +688,32 @@ def forward( hidden_states = inputs_embeds # decoder layers - all_hidden_states = () if output_hidden_states else None - all_self_attns = () if output_attentions else None - all_router_logits = () if output_router_logits else None next_decoder_cache = () if not use_new_cache else None - for layer_idx, decoder_layer in enumerate(self.layers): - if output_hidden_states: - all_hidden_states += (hidden_states,) - - if self.gradient_checkpointing and self.training: - layer_outputs = self._gradient_checkpointing_func( - partial(decoder_layer.__call__, **kwargs), - hidden_states, - attention_mask, - position_ids, - past_key_values, - output_attentions, - output_router_logits, - use_cache, - cache_position, - ) - else: - layer_outputs = decoder_layer( - hidden_states, - attention_mask=attention_mask, - position_ids=position_ids, - past_key_value=None if past_key_values is None else past_key_values[layer_idx], - output_attentions=output_attentions, - output_router_logits=output_router_logits, - use_cache=use_cache, - cache_position=cache_position, - token_idx=token_idx, - reuse_cache=reuse_cache, - flash_attention_recompute=flash_attention_recompute, - cache_idx=cache_idx, - ) + for layer_idx, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]): + layer_outputs = decoder_layer( + hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=None if past_key_values is None else past_key_values[layer_idx], + use_cache=use_cache, + cache_position=cache_position, + token_idx=token_idx, + reuse_cache=reuse_cache, + flash_attention_recompute=flash_attention_recompute, + cache_idx=cache_idx, + **kwargs, + ) hidden_states = layer_outputs[0] if use_cache: - next_decoder_cache += (layer_outputs[2 if output_attentions else 1],) - - if output_attentions: - all_self_attns += (layer_outputs[1],) - - if output_router_logits: - all_router_logits += (layer_outputs[-1],) + next_decoder_cache += (layer_outputs[1],) htcore.mark_step() hidden_states = self.norm(hidden_states) - # add hidden states from the last decoder layer - if output_hidden_states: - all_hidden_states += (hidden_states,) - next_cache = None if use_cache: next_cache = ( @@ -781,9 +723,6 @@ def forward( return MoeModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=next_cache, - hidden_states=all_hidden_states, - attentions=all_self_attns, - router_logits=all_router_logits, ) @@ -806,12 +745,10 @@ def forward( input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, + past_key_values: Optional[list[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, output_router_logits: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, logits_to_keep: Union[int, torch.Tensor] = 0, @@ -819,17 +756,12 @@ def forward( reuse_cache: Optional[bool] = None, flash_attention_recompute: Optional[bool] = False, cache_idx: int = None, - **kwargs: Unpack[KwargsForCausalLM], + **kwargs: Unpack[TransformersKwargs], ) -> MoeCausalLMOutputWithPast: - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_router_logits = ( output_router_logits if output_router_logits is not None else self.config.output_router_logits ) - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) outputs: MoeModelOutputWithPast = self.model( input_ids=input_ids, @@ -838,8 +770,6 @@ def forward( past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, output_router_logits=output_router_logits, cache_position=cache_position, token_idx=token_idx, diff --git a/optimum/habana/transformers/models/mllama/modeling_mllama.py b/optimum/habana/transformers/models/mllama/modeling_mllama.py index 8171a373ae..f0950db0d8 100644 --- a/optimum/habana/transformers/models/mllama/modeling_mllama.py +++ b/optimum/habana/transformers/models/mllama/modeling_mllama.py @@ -17,7 +17,7 @@ import math import os -from typing import List, Optional, Tuple, Union +from typing import Optional, Union import habana_frameworks.torch.core as htcore import torch @@ -114,7 +114,7 @@ def _prepare_cross_attention_mask( num_vision_tokens: int, dtype: str, token_idx: Optional[torch.Tensor] = None, -) -> Tuple[torch.Tensor, torch.Tensor]: +) -> tuple[torch.Tensor, torch.Tensor]: """ Copied from _prepare_cross_attention_mask: https://github.com/huggingface/transformers/blob/v4.45.2/src/transformers/models/mllama/modeling_mllama.py#L99 The only differences are: @@ -168,7 +168,7 @@ def forward( attention_mask: Optional[torch.Tensor] = None, output_attentions: Optional[bool] = None, use_flash_attention: Optional[bool] = False, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented. """ Copied from MllamaVisionSdpaAttention::forward:https://github.com/huggingface/transformers/blob/v4.45.2/src/transformers/models/mllama/modeling_mllama.py#L283 @@ -265,7 +265,7 @@ def forward( output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, use_flash_attention: Optional[bool] = False, - ) -> Union[Tuple, BaseModelOutput]: + ) -> Union[tuple, BaseModelOutput]: """ Copied from MllamaVisionEncoder::forward:https://github.com/huggingface/transformers/blob/v4.45.2/src/transformers/models/mllama/modeling_mllama.py#L394 The only differences are: @@ -330,7 +330,7 @@ def forward( token_idx: Optional[torch.Tensor] = None, use_flash_attention: Optional[bool] = False, flash_attention_recompute: Optional[bool] = False, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]: """ Copied from MllamaTextCrossAttention::forward: https://github.com/huggingface/transformers/blob/v4.45.2/src/transformers/models/mllama/modeling_mllama.py#L512 The only differences are: @@ -531,17 +531,17 @@ def forward( cross_attention_states: Optional[torch.Tensor] = None, cross_attention_mask: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, - full_text_row_masked_out_mask: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + full_text_row_masked_out_mask: Optional[tuple[torch.Tensor, torch.Tensor]] = None, position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Cache] = None, output_attentions: Optional[bool] = False, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, - position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # will become mandatory in v4.45 + position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, # will become mandatory in v4.45 token_idx: Optional[torch.Tensor] = None, use_flash_attention: Optional[bool] = False, flash_attention_recompute: Optional[bool] = False, - ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]: """ Copied from MllamaSelfAttentionDecoderLayer::forward: https://github.com/huggingface/transformers/blob/v4.45.2/src/transformers/models/mllama/modeling_mllama.py#L904 The only differences are: @@ -596,7 +596,7 @@ def forward( cross_attention_states: torch.Tensor, cross_attention_mask: torch.Tensor, attention_mask: torch.Tensor, - full_text_row_masked_out_mask: Tuple[torch.Tensor, torch.Tensor], + full_text_row_masked_out_mask: tuple[torch.Tensor, torch.Tensor], position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Cache] = None, output_attentions: Optional[bool] = False, @@ -606,7 +606,7 @@ def forward( token_idx: Optional[torch.Tensor] = None, use_flash_attention: Optional[bool] = False, flash_attention_recompute: Optional[bool] = False, - ) -> Tuple[torch.Tensor]: + ) -> tuple[torch.Tensor]: """ Copied from MllamaCrossAttentionDecoderLayer::forward: https://github.com/huggingface/transformers/blob/v4.45.2/src/transformers/models/mllama/modeling_mllama.py#L989 The only differences are: @@ -657,8 +657,8 @@ def forward( position_ids: Optional[torch.LongTensor] = None, cross_attention_states: Optional[torch.FloatTensor] = None, cross_attention_mask: Optional[torch.Tensor] = None, - full_text_row_masked_out_mask: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, - past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, + full_text_row_masked_out_mask: Optional[tuple[torch.Tensor, torch.Tensor]] = None, + past_key_values: Optional[Union[Cache, list[torch.FloatTensor]]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, @@ -668,7 +668,7 @@ def forward( token_idx: Optional[torch.Tensor] = None, use_flash_attention: Optional[bool] = False, flash_attention_recompute: Optional[bool] = False, - ) -> Union[Tuple, BaseModelOutputWithPast]: + ) -> Union[tuple, BaseModelOutputWithPast]: """ Copied from MllamaTextModel::forward: https://github.com/huggingface/transformers/blob/v4.45.2/src/transformers/models/mllama/modeling_mllama.py#L1617 The only differences are: @@ -903,8 +903,8 @@ def forward( position_ids: Optional[torch.LongTensor] = None, cross_attention_states: Optional[torch.LongTensor] = None, cross_attention_mask: Optional[torch.LongTensor] = None, - full_text_row_masked_out_mask: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, - past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, + full_text_row_masked_out_mask: Optional[tuple[torch.Tensor, torch.Tensor]] = None, + past_key_values: Optional[Union[Cache, list[torch.FloatTensor]]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, @@ -919,7 +919,7 @@ def forward( flash_attention_recompute: Optional[bool] = False, logits_bf16: Optional[bool] = False, **loss_kwargs, - ) -> Union[Tuple, CausalLMOutputWithPast]: + ) -> Union[tuple, CausalLMOutputWithPast]: """ Copied from MllamaForCausalLM::forward: https://github.com/huggingface/transformers/blob/v4.45.2/src/transformers/models/mllama/modeling_mllama.py#L1871 The only differences are: @@ -1003,7 +1003,7 @@ def forward( cross_attention_mask: Optional[torch.Tensor] = None, cross_attention_states: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, + past_key_values: Optional[list[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, @@ -1018,7 +1018,7 @@ def forward( flash_attention_recompute: Optional[bool] = False, logits_bf16: Optional[bool] = False, **loss_kwargs, - ) -> Union[Tuple, CausalLMOutputWithPast]: + ) -> Union[tuple, CausalLMOutputWithPast]: """ Copied from MllamaForConditionalGeneration::forward: https://github.com/huggingface/transformers/blob/v4.45.2/src/transformers/models/mllama/modeling_mllama.py#L2077 The only differences are: @@ -1242,7 +1242,7 @@ def forward( output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, use_flash_attention: Optional[bool] = False, - ) -> Union[BaseModelOutput, Tuple[torch.Tensor, ...]]: + ) -> Union[BaseModelOutput, tuple[torch.Tensor, ...]]: """ Copied from MllamaVisionModel::forward: https://github.com/huggingface/transformers/blob/v4.45.2/src/transformers/models/mllama/modeling_mllama.py#L1425 The only differences are: diff --git a/optimum/habana/transformers/models/modeling_all_models.py b/optimum/habana/transformers/models/modeling_all_models.py index 290c200618..9b1af1b141 100644 --- a/optimum/habana/transformers/models/modeling_all_models.py +++ b/optimum/habana/transformers/models/modeling_all_models.py @@ -14,8 +14,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Tuple - import torch from transformers.modeling_utils import ModuleUtilsMixin, PretrainedConfig from transformers.utils.import_utils import is_torch_sdpa_available @@ -135,7 +133,7 @@ def gaudi_invert_attention_mask(self, encoder_attention_mask: torch.Tensor) -> t def gaudi_get_extended_attention_mask( - self, attention_mask: torch.Tensor, input_shape: Tuple[int], device: torch.device = None, dtype: torch.float = None + self, attention_mask: torch.Tensor, input_shape: tuple[int], device: torch.device = None, dtype: torch.float = None ) -> torch.Tensor: """ Same as https://github.com/huggingface/transformers/blob/a9eee2ffecc874df7dd635b2c6abb246fdb318cc/src/transformers/modeling_utils.py#L692 diff --git a/optimum/habana/transformers/models/mpt/modeling_mpt.py b/optimum/habana/transformers/models/mpt/modeling_mpt.py index a6dcae591e..6dc17dc726 100755 --- a/optimum/habana/transformers/models/mpt/modeling_mpt.py +++ b/optimum/habana/transformers/models/mpt/modeling_mpt.py @@ -15,7 +15,7 @@ ############################################################################### # Copyright (C) 2022-2023 Habana Labs, Ltd. an Intel Company ############################################################################### -from typing import Optional, Tuple, Union +from typing import Optional, Union import torch from torch import nn @@ -43,15 +43,16 @@ class GaudiMptAttention(MptAttention): - def __init__(self, config: MptConfig): - super().__init__(config) + def __init__(self, config: MptConfig, layer_idx: Optional[int] = None): + super().__init__(config, layer_idx) def forward( self, hidden_states: torch.Tensor, position_bias: torch.Tensor, - past_key_value: Optional[Tuple[torch.Tensor]] = None, + past_key_value: Optional[tuple[torch.Tensor]] = None, attention_mask: Optional[torch.Tensor] = None, + cache_position: Optional[torch.Tensor] = None, token_idx: Optional[torch.Tensor] = None, use_flash_attention: Optional[bool] = False, flash_attention_recompute: Optional[bool] = False, @@ -155,18 +156,19 @@ def forward( class GaudiMptBlock(MptBlock): - def __init__(self, config: MptConfig): - super().__init__(config) - self.attn = GaudiMptAttention(config) + def __init__(self, config: MptConfig, layer_idx: Optional[int] = None): + super().__init__(config, layer_idx) + self.attn = GaudiMptAttention(config, layer_idx) def forward( self, hidden_states: torch.Tensor, position_bias: torch.Tensor, attention_mask: torch.Tensor, - layer_past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + layer_past: Optional[tuple[torch.Tensor, torch.Tensor]] = None, use_cache: bool = False, output_attentions: bool = False, + cache_position: Optional[torch.Tensor] = None, token_idx: Optional[torch.Tensor] = None, use_flash_attention: Optional[bool] = False, flash_attention_recompute: Optional[bool] = False, @@ -192,6 +194,7 @@ def forward( position_bias=position_bias, attention_mask=attention_mask, past_key_value=layer_past, + cache_position=cache_position, token_idx=token_idx, use_flash_attention=use_flash_attention, flash_attention_recompute=flash_attention_recompute, @@ -212,9 +215,6 @@ def forward( if use_cache: outputs += (past_key_value,) - if output_attentions: - outputs += (attn_weights,) - return outputs # hidden_states, present, attentions @@ -222,19 +222,20 @@ class GaudiMptModel(MptModel): def forward( self, input_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None, + past_key_values: Optional[tuple[tuple[torch.Tensor, torch.Tensor], ...]] = None, attention_mask: Optional[torch.Tensor] = None, inputs_embeds: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, + cache_position: Optional[torch.Tensor] = None, token_idx: Optional[torch.Tensor] = None, use_flash_attention: Optional[bool] = False, flash_attention_recompute: Optional[bool] = False, cache_idx: Optional[torch.Tensor] = None, **kwargs, # NOOP kwargs, for now - ) -> Union[Tuple[torch.Tensor, ...], BaseModelOutputWithPastAndCrossAttentions]: + ) -> Union[tuple[torch.Tensor, ...], BaseModelOutputWithPastAndCrossAttentions]: """ Copied from MptModel.forward: https://github.com/huggingface/transformers/blob/v4.32.0/src/transformers/models/mpt/modeling_mpt.py The only differences are: @@ -271,13 +272,6 @@ def forward( all_self_attentions = () if output_attentions else None all_hidden_states = () if output_hidden_states else None - if self.gradient_checkpointing and self.training: - if use_cache: - logger.warning_once( - "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." - ) - use_cache = False - # Compute alibi tensor: check build_alibi_tensor documentation seq_length_with_past = seq_length past_key_values_length = 0 @@ -300,30 +294,18 @@ def forward( if output_hidden_states: all_hidden_states = all_hidden_states + (hidden_states,) - if self.gradient_checkpointing and self.training: - outputs = self._gradient_checkpointing_func( - block.__call__, - hidden_states, - alibi, - causal_mask, - layer_past, - use_cache, - output_attentions, - None, - ) - else: - outputs = block( - hidden_states, - layer_past=layer_past, - attention_mask=causal_mask, - use_cache=use_cache, - output_attentions=output_attentions, - position_bias=alibi, - token_idx=token_idx, - use_flash_attention=use_flash_attention, - flash_attention_recompute=flash_attention_recompute, - cache_idx=cache_idx, - ) + outputs = block( + hidden_states, + layer_past=layer_past, + attention_mask=causal_mask, + use_cache=use_cache, + output_attentions=output_attentions, + position_bias=alibi, + token_idx=token_idx, + use_flash_attention=use_flash_attention, + flash_attention_recompute=flash_attention_recompute, + cache_idx=cache_idx, + ) hidden_states = outputs[0] if use_cache is True: @@ -422,7 +404,7 @@ def prepare_inputs_for_generation( def forward( self, input_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None, + past_key_values: Optional[tuple[tuple[torch.Tensor, torch.Tensor], ...]] = None, attention_mask: Optional[torch.Tensor] = None, inputs_embeds: Optional[torch.Tensor] = None, labels: Optional[torch.Tensor] = None, @@ -430,12 +412,13 @@ def forward( output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, + cache_position: Optional[torch.Tensor] = None, token_idx: Optional[torch.Tensor] = None, use_flash_attention: Optional[bool] = False, flash_attention_recompute: Optional[bool] = False, cache_idx: Optional[torch.Tensor] = None, **kwargs, - ) -> Union[Tuple[torch.Tensor], CausalLMOutputWithCrossAttentions]: + ) -> Union[tuple[torch.Tensor], CausalLMOutputWithCrossAttentions]: """ Inherits from MptForCausalLM: https://github.com/huggingface/transformers/blob/v4.32.0/src/transformers/models/mpt/modeling_mpt.py The only differences are: @@ -455,6 +438,7 @@ def forward( output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, + cache_position=cache_position, token_idx=token_idx, use_flash_attention=use_flash_attention, flash_attention_recompute=flash_attention_recompute, diff --git a/optimum/habana/transformers/models/opt/modeling_opt.py b/optimum/habana/transformers/models/opt/modeling_opt.py index 54e91c7486..bd11f5ca56 100644 --- a/optimum/habana/transformers/models/opt/modeling_opt.py +++ b/optimum/habana/transformers/models/opt/modeling_opt.py @@ -1,4 +1,4 @@ -from typing import List, Optional, Tuple, Union +from typing import Optional, Union import torch from transformers.activations import ACT2FN @@ -49,14 +49,14 @@ def gaudi_opt_attention_forward( self, hidden_states: torch.Tensor, key_value_states: Optional[torch.Tensor] = None, - past_key_value: Optional[Tuple[torch.Tensor]] = None, + past_key_value: Optional[tuple[torch.Tensor]] = None, attention_mask: Optional[torch.Tensor] = None, layer_head_mask: Optional[torch.Tensor] = None, output_attentions: bool = False, # isn't needed in normal attention, but needed in flash attention so to keep the signature same position_ids: Optional[torch.Tensor] = None, token_idx: Optional[torch.Tensor] = None, -) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: +) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]: """ Copied from OPTAttention.forward: https://github.com/huggingface/transformers/blob/main/src/transformers/models/opt/modeling_opt.py The only differences are: @@ -194,12 +194,12 @@ def forward( hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, layer_head_mask: Optional[torch.Tensor] = None, - past_key_value: Optional[Tuple[torch.Tensor]] = None, + past_key_value: Optional[tuple[torch.Tensor]] = None, output_attentions: Optional[bool] = False, use_cache: Optional[bool] = False, position_ids: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, - ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]: """ Copied from OPTDecoderLayer.forward: https://github.com/huggingface/transformers/blob/main/src/transformers/models/opt/modeling_opt.py The only differences are: @@ -265,7 +265,7 @@ def gaudi_opt_decoder_forward( input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, head_mask: Optional[torch.Tensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, + past_key_values: Optional[list[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, @@ -273,7 +273,7 @@ def gaudi_opt_decoder_forward( return_dict: Optional[bool] = None, position_ids: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, -) -> Union[Tuple, BaseModelOutputWithPast]: +) -> Union[tuple, BaseModelOutputWithPast]: """ Copied from OPTDecoder.forward: https://github.com/huggingface/transformers/blob/main/src/transformers/models/opt/modeling_opt.py The only differences are: @@ -421,7 +421,7 @@ def gaudi_opt_model_forward( input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, head_mask: Optional[torch.Tensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, + past_key_values: Optional[list[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, @@ -429,7 +429,7 @@ def gaudi_opt_model_forward( return_dict: Optional[bool] = None, position_ids: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, -) -> Union[Tuple, BaseModelOutputWithPast]: +) -> Union[tuple, BaseModelOutputWithPast]: """ Copied from OPTModel.forward: https://github.com/huggingface/transformers/blob/main/src/transformers/models/opt/modeling_opt.py The only differences are: @@ -482,7 +482,7 @@ def forward( input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, head_mask: Optional[torch.Tensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, + past_key_values: Optional[list[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, @@ -492,7 +492,7 @@ def forward( position_ids: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, **kwargs, - ) -> Union[Tuple, CausalLMOutputWithPast]: + ) -> Union[tuple, CausalLMOutputWithPast]: 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 diff --git a/optimum/habana/transformers/models/owlvit/modeling_owlvit.py b/optimum/habana/transformers/models/owlvit/modeling_owlvit.py index 4ba4cd2ad1..9d58b810ea 100644 --- a/optimum/habana/transformers/models/owlvit/modeling_owlvit.py +++ b/optimum/habana/transformers/models/owlvit/modeling_owlvit.py @@ -1,4 +1,4 @@ -from typing import Optional, Tuple +from typing import Optional import torch @@ -8,7 +8,7 @@ def gaudi_owlvitclasspredictionhead_forward( image_embeds: torch.FloatTensor, query_embeds: Optional[torch.FloatTensor], query_mask: Optional[torch.Tensor], -) -> Tuple[torch.FloatTensor]: +) -> tuple[torch.FloatTensor]: """ Copied from modeling_owlvit: https://github.com/huggingface/transformers/blob/v4.37.2/src/transformers/models/owlvit/modeling_owlvit.py#L1233 The only modification is: diff --git a/optimum/habana/transformers/models/phi/modeling_phi.py b/optimum/habana/transformers/models/phi/modeling_phi.py index 8d0837e30e..119ba8a632 100644 --- a/optimum/habana/transformers/models/phi/modeling_phi.py +++ b/optimum/habana/transformers/models/phi/modeling_phi.py @@ -19,15 +19,13 @@ # limitations under the License. """PyTorch Phi model.""" -from functools import partial -from typing import List, Optional, Tuple, Union +from typing import Optional, Union import torch from transformers.cache_utils import Cache from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from transformers.models.phi.configuration_phi import PhiConfig from transformers.models.phi.modeling_phi import ( - KwargsForCausalLM, PhiAttention, PhiForCausalLM, PhiMLP, @@ -35,7 +33,7 @@ apply_rotary_pos_emb, ) from transformers.processing_utils import Unpack -from transformers.utils import logging +from transformers.utils import TransformersKwargs, logging from ...modeling_attn_mask_utils import ( _gaudi_prepare_4d_causal_attention_mask, @@ -130,7 +128,7 @@ def allocate_kv_cache(self, batch_size, max_seq_len, inp_seq_len): def forward( self, hidden_states: torch.Tensor, - position_embeddings: Tuple[torch.Tensor, torch.Tensor], + position_embeddings: tuple[torch.Tensor, torch.Tensor], attention_mask: Optional[torch.Tensor], past_key_value: Optional[Cache] = None, use_cache: bool = False, @@ -139,7 +137,7 @@ def forward( reuse_cache: Optional[bool] = False, cache_idx: Optional[int] = None, **kwargs, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]: """ Copied from PhiAttention.forward: https://github.com/huggingface/transformers/blob/v4.37.1/src/transformers/models/phi/modeling_phi.py The only differences are: @@ -259,16 +257,16 @@ def forward( hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_value: Optional[Tuple[torch.Tensor]] = None, + past_key_value: Optional[tuple[torch.Tensor]] = None, output_attentions: Optional[bool] = False, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, - position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC + position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC token_idx: Optional[torch.Tensor] = None, reuse_cache: Optional[bool] = False, cache_idx: Optional[int] = None, **kwargs, - ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]: """ Copied from PhiDecoderLayer.forward: https://github.com/huggingface/transformers/blob/v4.37.1/src/transformers/models/phi/modeling_phi.py The only differences are: @@ -320,7 +318,7 @@ def forward( input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, + past_key_values: Optional[list[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, @@ -400,31 +398,19 @@ def forward( if output_hidden_states: all_hidden_states += (hidden_states,) - if self.gradient_checkpointing and self.training: - layer_outputs = self._gradient_checkpointing_func( - partial(decoder_layer.__call__, **kwargs), - hidden_states, - attention_mask, - position_ids, - None if past_key_values is None else past_key_values[layer_idx], - output_attentions, - use_cache, - cache_position, - None, - ) - else: - layer_outputs = decoder_layer( - hidden_states, - attention_mask=attention_mask, - position_ids=position_ids, - past_key_value=None if past_key_values is None else past_key_values[layer_idx], - output_attentions=output_attentions, - use_cache=use_cache, - cache_position=cache_position, - token_idx=token_idx, - reuse_cache=reuse_cache, - cache_idx=cache_idx, - ) + layer_outputs = decoder_layer( + hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=None if past_key_values is None else past_key_values[layer_idx], + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + token_idx=token_idx, + reuse_cache=reuse_cache, + cache_idx=cache_idx, + **kwargs, + ) hidden_states = layer_outputs[0] @@ -463,19 +449,17 @@ def forward( input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, + past_key_values: Optional[list[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, logits_to_keep: Union[int, torch.Tensor] = 0, token_idx: Optional[torch.Tensor] = None, reuse_cache: Optional[bool] = False, trim_logits: Optional[bool] = False, cache_idx: Optional[int] = None, - **kwargs: Unpack[KwargsForCausalLM], + **kwargs: Unpack[TransformersKwargs], ) -> CausalLMOutputWithPast: """ Inherits from PhiForCausalLM: https://github.com/huggingface/transformers/blob/v4.37.1/src/transformers/models/phi/modeling_phi.py @@ -484,12 +468,6 @@ def forward( - add new args reuse_cache - add new args cache_idx """ - 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 - ) - - # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) outputs: BaseModelOutputWithPast = self.model( input_ids=input_ids, attention_mask=attention_mask, @@ -497,8 +475,6 @@ def forward( past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, cache_position=cache_position, token_idx=token_idx, reuse_cache=reuse_cache, diff --git a/optimum/habana/transformers/models/qwen2/modeling_qwen2.py b/optimum/habana/transformers/models/qwen2/modeling_qwen2.py index 6bc79b46f0..3bad0f8b11 100644 --- a/optimum/habana/transformers/models/qwen2/modeling_qwen2.py +++ b/optimum/habana/transformers/models/qwen2/modeling_qwen2.py @@ -16,11 +16,11 @@ # Copyright (C) 2022-2024 Habana Labs, Ltd. an Intel Company ############################################################################### -from functools import partial -from typing import List, Optional, Tuple, Union +from typing import Optional, Union import torch from transformers.cache_utils import Cache, DynamicCache, StaticCache +from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask from transformers.modeling_outputs import ( BaseModelOutputWithPast, CausalLMOutputWithPast, @@ -29,7 +29,6 @@ ) from transformers.models.qwen2.configuration_qwen2 import Qwen2Config from transformers.models.qwen2.modeling_qwen2 import ( - KwargsForCausalLM, Qwen2Attention, Qwen2DecoderLayer, Qwen2ForCausalLM, @@ -42,6 +41,7 @@ logger, ) from transformers.processing_utils import Unpack +from transformers.utils import TransformersKwargs from ....distributed import parallel_state from ...modeling_attn_mask_utils import ( @@ -365,7 +365,7 @@ def reorder_kv_cache(self, beam_idx: torch.LongTensor): def pre_attn_forward( self, hidden_states: torch.Tensor, - position_embeddings: Tuple[torch.Tensor, torch.Tensor], + position_embeddings: tuple[torch.Tensor, torch.Tensor], attention_mask: Optional[torch.Tensor], past_key_value: Optional[Cache] = None, use_cache: bool = False, @@ -381,7 +381,7 @@ def pre_attn_forward( cache_idx: int = None, num_virtual_tokens: int = None, **kwargs, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]: """ The only differences are: - add new args token_idx @@ -485,13 +485,6 @@ def pre_attn_forward( fused_scaled_dot_product_attention = get_gaudi_distributed_attention( self.fused_scaled_dot_product_attention, self.fused_scaled_dot_product_attention_distributed ) - sliding_window = None - if ( - self.config.use_sliding_window - and getattr(self.config, "sliding_window", None) is not None - and self.layer_idx >= self.config.max_window_layers - ): - sliding_window = self.config.sliding_window if use_flash_attention and FusedSDPA is not None: attn_weights = None @@ -552,9 +545,10 @@ def pre_attn_forward( attention_mask, dropout=0.0 if not self.training else self.attention_dropout, scaling=self.scaling, - sliding_window=sliding_window, # main diff with Llama + sliding_window=self.sliding_window, # main diff with Llama attn_softmax_bf16=attn_softmax_bf16, input_shape=input_shape, + **kwargs, ) attn_output = attn_output.transpose(1, 2).contiguous() @@ -603,11 +597,11 @@ def forward( hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_value: Optional[Tuple[torch.Tensor]] = None, + past_key_value: Optional[tuple[torch.Tensor]] = None, output_attentions: Optional[bool] = False, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, - position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC + position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC token_idx: Optional[torch.Tensor] = None, attn_softmax_bf16: Optional[bool] = False, reuse_cache: Optional[bool] = False, @@ -619,7 +613,7 @@ def forward( cache_idx: int = None, num_virtual_tokens: int = None, **kwargs, - ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]: residual = hidden_states hidden_states, self_attn_weights, present_key_value = self.pre_attn( @@ -663,11 +657,11 @@ def pre_attn( hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_value: Optional[Tuple[torch.Tensor]] = None, + past_key_value: Optional[tuple[torch.Tensor]] = None, output_attentions: Optional[bool] = False, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, - position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, token_idx: Optional[torch.Tensor] = None, attn_softmax_bf16: Optional[bool] = False, reuse_cache: Optional[bool] = False, @@ -679,7 +673,7 @@ def pre_attn( cache_idx: int = None, num_virtual_tokens: int = None, **kwargs, - ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]: hidden_states = self.input_layernorm(hidden_states) hidden_states, attn_weights, present_key_value = self.self_attn.pre_attn_forward( hidden_states=hidden_states, @@ -747,8 +741,9 @@ def __init__(self, config: Qwen2Config): [GaudiQwen2DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] ) self.norm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.gradient_checkpointing = False + self.has_sliding_layers = "sliding_attention" in self.config.layer_types + # Initialize weights and apply final processing self.post_init() @@ -768,11 +763,9 @@ def forward( input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, + past_key_values: Optional[list[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, attn_softmax_bf16: Optional[bool] = False, @@ -787,10 +780,6 @@ def forward( num_virtual_tokens: int = None, **kwargs, ) -> BaseModelOutputWithPast: - 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 - ) use_cache = use_cache if use_cache is not None else self.config.use_cache # retrieve input_ids and inputs_embeds @@ -803,12 +792,6 @@ def forward( else: raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds") - if self.gradient_checkpointing and self.training and use_cache: - logger.warning_once( - "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`." - ) - use_cache = False - if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) @@ -856,14 +839,29 @@ def forward( past_seen_tokens, ) else: - causal_mask = self._update_causal_mask(attention_mask, inputs_embeds, cache_position, past_seen_tokens) + # It may already have been prepared by e.g. `generate` + if not isinstance(causal_mask_mapping := attention_mask, dict): + # Prepare mask arguments + mask_kwargs = { + "config": self.config, + "input_embeds": inputs_embeds, + "attention_mask": attention_mask, + "cache_position": cache_position, + "past_key_values": past_seen_tokens, + "position_ids": position_ids, + } + # Create the masks + causal_mask_mapping = { + "full_attention": create_causal_mask(**mask_kwargs), + } + # The sliding window alternating layers are not always activated depending on the config + if self.has_sliding_layers: + causal_mask_mapping["sliding_attention"] = create_sliding_window_causal_mask(**mask_kwargs) + causal_mask = causal_mask_mapping # embed positions hidden_states = inputs_embeds - # decoder layers - all_hidden_states = () if output_hidden_states else None - all_self_attns = () if output_attentions else None next_decoder_cache = () if not use_new_cache else None if lazy_mode: @@ -877,65 +875,33 @@ def forward( ): htcore.mark_step() - if output_hidden_states: - all_hidden_states += (hidden_states,) - - if self.gradient_checkpointing and self.training: - layer_outputs = self._gradient_checkpointing_func( - partial(decoder_layer.__call__, **kwargs), - hidden_states, - causal_mask, - position_ids, - past_key_values, - output_attentions, - use_cache, - cache_position, - None, - None, - attn_softmax_bf16, - False, - use_flash_attention, - flash_attention_recompute, - flash_attention_causal_mask, - flash_attention_fast_softmax, - valid_sequence_lengths, - None, - ) - else: - layer_outputs = decoder_layer( - hidden_states, - attention_mask=causal_mask, - position_ids=position_ids, - past_key_value=None if past_key_values is None else past_key_values[layer_idx], - output_attentions=output_attentions, - use_cache=use_cache, - cache_position=cache_position, - token_idx=token_idx, - attn_softmax_bf16=attn_softmax_bf16, - reuse_cache=reuse_cache, - use_flash_attention=use_flash_attention, - flash_attention_recompute=flash_attention_recompute, - flash_attention_causal_mask=flash_attention_causal_mask, - flash_attention_fast_softmax=flash_attention_fast_softmax, - valid_sequence_lengths=valid_sequence_lengths, - cache_idx=cache_idx, - num_virtual_tokens=num_virtual_tokens, - ) + layer_outputs = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_ids=position_ids, + past_key_value=None if past_key_values is None else past_key_values[layer_idx], + use_cache=use_cache, + cache_position=cache_position, + token_idx=token_idx, + attn_softmax_bf16=attn_softmax_bf16, + reuse_cache=reuse_cache, + use_flash_attention=use_flash_attention, + flash_attention_recompute=flash_attention_recompute, + flash_attention_causal_mask=flash_attention_causal_mask, + flash_attention_fast_softmax=flash_attention_fast_softmax, + valid_sequence_lengths=valid_sequence_lengths, + cache_idx=cache_idx, + num_virtual_tokens=num_virtual_tokens, + **kwargs, + ) hidden_states = layer_outputs[0] if use_cache: - next_decoder_cache += (layer_outputs[2 if output_attentions else 1],) - - if output_attentions: - all_self_attns += (layer_outputs[1],) + next_decoder_cache += (layer_outputs[1],) hidden_states = self.norm(hidden_states) - # add hidden states from the last decoder layer - if output_hidden_states: - all_hidden_states += (hidden_states,) - next_cache = None if use_cache: next_cache = ( @@ -944,8 +910,6 @@ def forward( return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=next_cache, - hidden_states=all_hidden_states, - attentions=all_self_attns, ) @@ -964,12 +928,10 @@ def forward( input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, + past_key_values: Optional[list[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, logits_to_keep: Union[int, torch.Tensor] = 0, token_idx: Optional[torch.Tensor] = None, @@ -984,17 +946,12 @@ def forward( cache_idx: int = None, lazy_mode: Optional[bool] = True, num_virtual_tokens: int = None, - **kwargs: Unpack[KwargsForCausalLM], + **kwargs: Unpack[TransformersKwargs], ) -> CausalLMOutputWithPast: - 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 - ) if self.generation_config.use_fused_rope is False: global has_fused_rope has_fused_rope = False - # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) outputs: BaseModelOutputWithPast = self.model( input_ids=input_ids, attention_mask=attention_mask, @@ -1002,8 +959,6 @@ def forward( past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, cache_position=cache_position, token_idx=token_idx, attn_softmax_bf16=attn_softmax_bf16, @@ -1044,8 +999,8 @@ def forward( @staticmethod def _reorder_cache( - past: Tuple[Tuple[torch.Tensor, torch.Tensor], ...], beam_idx: torch.LongTensor - ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], ...]: + past: tuple[tuple[torch.Tensor, torch.Tensor], ...], beam_idx: torch.LongTensor + ) -> tuple[tuple[torch.Tensor, torch.Tensor], ...]: """ This function is used to re-order the `past_key_values` cache if [`~PreTrainedModel.beam_search`] or [`~PreTrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct diff --git a/optimum/habana/transformers/models/qwen2_moe/modeling_qwen2_moe.py b/optimum/habana/transformers/models/qwen2_moe/modeling_qwen2_moe.py index 6fadcfe279..07238749f0 100755 --- a/optimum/habana/transformers/models/qwen2_moe/modeling_qwen2_moe.py +++ b/optimum/habana/transformers/models/qwen2_moe/modeling_qwen2_moe.py @@ -20,7 +20,7 @@ """PyTorch Qwen2MoE model.""" import math -from typing import List, Optional, Tuple, Union +from typing import Optional, Union import habana_frameworks.torch.core as htcore import torch @@ -322,7 +322,7 @@ def pre_attn_forward( output_attentions: bool = False, use_cache: bool = False, cache_position: Optional[torch.LongTensor] = None, - position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC + position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC token_idx: Optional[torch.Tensor] = None, attn_softmax_bf16: Optional[bool] = False, reuse_cache: Optional[bool] = False, @@ -334,7 +334,7 @@ def pre_attn_forward( cache_idx: int = None, num_virtual_tokens: int = None, **kwargs, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]: """ Copied from LlamaAttention.forward: https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py The only differences are: @@ -527,7 +527,7 @@ def post_attn_forward(self, attn_output): return attn_output -def gaudi_qwen2moe_block_sparse_moe_forward(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: +def gaudi_qwen2moe_block_sparse_moe_forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """ - optimize expert forward, remove dynamic control and dynamic shape """ @@ -633,12 +633,12 @@ def forward( hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_value: Optional[Tuple[torch.Tensor]] = None, + past_key_value: Optional[tuple[torch.Tensor]] = None, output_attentions: Optional[bool] = False, output_router_logits: Optional[bool] = False, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, - position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC + position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC token_idx: Optional[torch.Tensor] = None, attn_softmax_bf16: Optional[bool] = False, reuse_cache: Optional[bool] = False, @@ -650,7 +650,7 @@ def forward( cache_idx: int = None, num_virtual_tokens: int = None, **kwargs, - ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]: """ Copied from LlamaDecoderLayer.forward: https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py The only differences are: @@ -708,11 +708,11 @@ def pre_attn( hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_value: Optional[Tuple[torch.Tensor]] = None, + past_key_value: Optional[tuple[torch.Tensor]] = None, output_attentions: Optional[bool] = False, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, - position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC + position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC token_idx: Optional[torch.Tensor] = None, attn_softmax_bf16: Optional[bool] = False, reuse_cache: Optional[bool] = False, @@ -724,7 +724,7 @@ def pre_attn( cache_idx: int = None, num_virtual_tokens: int = None, **kwargs, - ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]: hidden_states = self.input_layernorm(hidden_states) hidden_states, attn_weights, present_key_value = self.self_attn.pre_attn_forward( hidden_states=hidden_states, @@ -812,7 +812,7 @@ def forward( input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, + past_key_values: Optional[list[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, @@ -935,48 +935,26 @@ def forward( if output_hidden_states: all_hidden_states += (hidden_states,) - if self.gradient_checkpointing and self.training: - layer_outputs = self._gradient_checkpointing_func( - decoder_layer.__call__, - hidden_states, - causal_mask, - position_ids, - past_key_values, - output_attentions, - output_router_logits, - use_cache, - cache_position, - None, - attn_softmax_bf16, - False, - use_flash_attention, - flash_attention_recompute, - flash_attention_causal_mask, - flash_attention_fast_softmax, - valid_sequence_lengths, - None, - ) - else: - layer_outputs = decoder_layer( - hidden_states, - attention_mask=causal_mask, - position_ids=position_ids, - past_key_value=None if past_key_values is None else past_key_values[layer_idx], - output_attentions=output_attentions, - output_router_logits=output_router_logits, - use_cache=use_cache, - cache_position=cache_position, - token_idx=token_idx, - attn_softmax_bf16=attn_softmax_bf16, - reuse_cache=reuse_cache, - use_flash_attention=use_flash_attention, - flash_attention_recompute=flash_attention_recompute, - flash_attention_causal_mask=flash_attention_causal_mask, - flash_attention_fast_softmax=flash_attention_fast_softmax, - valid_sequence_lengths=valid_sequence_lengths, - cache_idx=cache_idx, - num_virtual_tokens=num_virtual_tokens, - ) + layer_outputs = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_ids=position_ids, + past_key_value=None if past_key_values is None else past_key_values[layer_idx], + output_attentions=output_attentions, + output_router_logits=output_router_logits, + use_cache=use_cache, + cache_position=cache_position, + token_idx=token_idx, + attn_softmax_bf16=attn_softmax_bf16, + reuse_cache=reuse_cache, + use_flash_attention=use_flash_attention, + flash_attention_recompute=flash_attention_recompute, + flash_attention_causal_mask=flash_attention_causal_mask, + flash_attention_fast_softmax=flash_attention_fast_softmax, + valid_sequence_lengths=valid_sequence_lengths, + cache_idx=cache_idx, + num_virtual_tokens=num_virtual_tokens, + ) hidden_states = layer_outputs[0] @@ -1034,7 +1012,7 @@ def forward( input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, + past_key_values: Optional[list[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, @@ -1055,7 +1033,7 @@ def forward( cache_idx: int = None, lazy_mode: Optional[bool] = True, num_virtual_tokens: int = None, - **loss_kwargs, + **kwargs, ) -> MoeCausalLMOutputWithPast: output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_router_logits = ( @@ -1108,7 +1086,7 @@ def forward( loss = None if labels is not None: - loss = self.loss_function(logits, labels, self.vocab_size, **loss_kwargs) + loss = self.loss_function(logits, labels, self.vocab_size, **kwargs) aux_loss = None if output_router_logits: @@ -1133,8 +1111,8 @@ def forward( @staticmethod def _reorder_cache( - past: Tuple[Tuple[torch.Tensor, torch.Tensor], ...], beam_idx: torch.LongTensor - ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], ...]: + past: tuple[tuple[torch.Tensor, torch.Tensor], ...], beam_idx: torch.LongTensor + ) -> tuple[tuple[torch.Tensor, torch.Tensor], ...]: """ This function is used to re-order the `past_key_values` cache if [`~PreTrainedModel.beam_search`] or [`~PreTrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct diff --git a/optimum/habana/transformers/models/qwen3/modeling_qwen3.py b/optimum/habana/transformers/models/qwen3/modeling_qwen3.py index 1021fc17dc..db35ff73a6 100644 --- a/optimum/habana/transformers/models/qwen3/modeling_qwen3.py +++ b/optimum/habana/transformers/models/qwen3/modeling_qwen3.py @@ -16,15 +16,14 @@ # Copyright (C) 2022-2024 Habana Labs, Ltd. an Intel Company ############################################################################### -from functools import partial -from typing import List, Optional, Tuple, Union +from typing import Optional, Union import torch from transformers.cache_utils import Cache, DynamicCache, StaticCache +from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from transformers.models.qwen3.configuration_qwen3 import Qwen3Config from transformers.models.qwen3.modeling_qwen3 import ( - KwargsForCausalLM, Qwen3Attention, Qwen3DecoderLayer, Qwen3ForCausalLM, @@ -32,9 +31,9 @@ Qwen3Model, Qwen3RMSNorm, apply_rotary_pos_emb, - logger, ) from transformers.processing_utils import Unpack +from transformers.utils import TransformersKwargs from ....distributed import parallel_state from ...modeling_attn_mask_utils import ( @@ -358,7 +357,7 @@ def reorder_kv_cache(self, beam_idx: torch.LongTensor): def pre_attn_forward( self, hidden_states: torch.Tensor, - position_embeddings: Tuple[torch.Tensor, torch.Tensor], + position_embeddings: tuple[torch.Tensor, torch.Tensor], attention_mask: Optional[torch.Tensor], past_key_value: Optional[Cache] = None, use_cache: bool = False, @@ -374,7 +373,7 @@ def pre_attn_forward( cache_idx: int = None, num_virtual_tokens: int = None, **kwargs, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]: """ The only differences are: - add new args token_idx @@ -478,13 +477,6 @@ def pre_attn_forward( fused_scaled_dot_product_attention = get_gaudi_distributed_attention( self.fused_scaled_dot_product_attention, self.fused_scaled_dot_product_attention_distributed ) - sliding_window = None - if ( - self.config.use_sliding_window - and getattr(self.config, "sliding_window", None) is not None - and self.layer_idx >= self.config.max_window_layers - ): - sliding_window = self.config.sliding_window if use_flash_attention and FusedSDPA is not None: attn_weights = None @@ -544,7 +536,7 @@ def pre_attn_forward( attention_mask, dropout=0.0 if not self.training else self.attention_dropout, scaling=self.scaling, - sliding_window=sliding_window, # main diff with Llama + sliding_window=self.sliding_window, # main diff with Llama attn_softmax_bf16=attn_softmax_bf16, input_shape=input_shape, ) @@ -595,11 +587,11 @@ def forward( hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_value: Optional[Tuple[torch.Tensor]] = None, + past_key_value: Optional[tuple[torch.Tensor]] = None, output_attentions: Optional[bool] = False, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, - position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC + position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC token_idx: Optional[torch.Tensor] = None, attn_softmax_bf16: Optional[bool] = False, reuse_cache: Optional[bool] = False, @@ -611,7 +603,7 @@ def forward( cache_idx: int = None, num_virtual_tokens: int = None, **kwargs, - ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]: residual = hidden_states hidden_states, self_attn_weights, present_key_value = self.pre_attn( @@ -619,7 +611,6 @@ def forward( attention_mask=attention_mask, position_ids=position_ids, past_key_value=past_key_value, - output_attentions=output_attentions, use_cache=use_cache, cache_position=cache_position, position_embeddings=position_embeddings, @@ -643,8 +634,6 @@ def forward( hidden_states = self.post_mlp(hidden_states, residual) outputs = (hidden_states,) - if output_attentions: - outputs += (self_attn_weights,) if use_cache: outputs += (present_key_value,) @@ -655,11 +644,11 @@ def pre_attn( hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_value: Optional[Tuple[torch.Tensor]] = None, + past_key_value: Optional[tuple[torch.Tensor]] = None, output_attentions: Optional[bool] = False, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, - position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, token_idx: Optional[torch.Tensor] = None, attn_softmax_bf16: Optional[bool] = False, reuse_cache: Optional[bool] = False, @@ -671,7 +660,7 @@ def pre_attn( cache_idx: int = None, num_virtual_tokens: int = None, **kwargs, - ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]: hidden_states = self.input_layernorm(hidden_states) hidden_states, attn_weights, present_key_value = self.self_attn.pre_attn_forward( hidden_states=hidden_states, @@ -739,8 +728,9 @@ def __init__(self, config: Qwen3Config): [GaudiQwen3DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] ) self.norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.gradient_checkpointing = False + self.has_sliding_layers = "sliding_attention" in self.config.layer_types + # Initialize weights and apply final processing self.post_init() @@ -760,11 +750,9 @@ def forward( input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, + past_key_values: Optional[list[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, attn_softmax_bf16: Optional[bool] = False, @@ -779,10 +767,6 @@ def forward( num_virtual_tokens: int = None, **kwargs, ) -> BaseModelOutputWithPast: - 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 - ) use_cache = use_cache if use_cache is not None else self.config.use_cache # retrieve input_ids and inputs_embeds @@ -795,12 +779,6 @@ def forward( else: raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds") - if self.gradient_checkpointing and self.training and use_cache: - logger.warning_once( - "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`." - ) - use_cache = False - if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) @@ -849,13 +827,29 @@ def forward( ) else: causal_mask = self._update_causal_mask(attention_mask, inputs_embeds, cache_position, past_seen_tokens) + # It may already have been prepared by e.g. `generate` + if not isinstance(causal_mask_mapping := attention_mask, dict): + # Prepare mask arguments + mask_kwargs = { + "config": self.config, + "input_embeds": inputs_embeds, + "attention_mask": attention_mask, + "cache_position": cache_position, + "past_key_values": past_seen_tokens, + "position_ids": position_ids, + } + # Create the masks + causal_mask_mapping = { + "full_attention": create_causal_mask(**mask_kwargs), + } + # The sliding window alternating layers are not always activated depending on the config + if self.has_sliding_layers: + causal_mask_mapping["sliding_attention"] = create_sliding_window_causal_mask(**mask_kwargs) + causal_mask = causal_mask_mapping # embed positions hidden_states = inputs_embeds - # decoder layers - all_hidden_states = () if output_hidden_states else None - all_self_attns = () if output_attentions else None next_decoder_cache = () if not use_new_cache else None if lazy_mode: @@ -869,65 +863,33 @@ def forward( ): htcore.mark_step() - if output_hidden_states: - all_hidden_states += (hidden_states,) - - if self.gradient_checkpointing and self.training: - layer_outputs = self._gradient_checkpointing_func( - partial(decoder_layer.__call__, **kwargs), - hidden_states, - causal_mask, - position_ids, - past_key_values, - output_attentions, - use_cache, - cache_position, - None, - None, - attn_softmax_bf16, - False, - use_flash_attention, - flash_attention_recompute, - flash_attention_causal_mask, - flash_attention_fast_softmax, - valid_sequence_lengths, - None, - ) - else: - layer_outputs = decoder_layer( - hidden_states, - attention_mask=causal_mask, - position_ids=position_ids, - past_key_value=None if past_key_values is None else past_key_values[layer_idx], - output_attentions=output_attentions, - use_cache=use_cache, - cache_position=cache_position, - token_idx=token_idx, - attn_softmax_bf16=attn_softmax_bf16, - reuse_cache=reuse_cache, - use_flash_attention=use_flash_attention, - flash_attention_recompute=flash_attention_recompute, - flash_attention_causal_mask=flash_attention_causal_mask, - flash_attention_fast_softmax=flash_attention_fast_softmax, - valid_sequence_lengths=valid_sequence_lengths, - cache_idx=cache_idx, - num_virtual_tokens=num_virtual_tokens, - ) + layer_outputs = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_ids=position_ids, + past_key_value=None if past_key_values is None else past_key_values[layer_idx], + use_cache=use_cache, + cache_position=cache_position, + token_idx=token_idx, + attn_softmax_bf16=attn_softmax_bf16, + reuse_cache=reuse_cache, + use_flash_attention=use_flash_attention, + flash_attention_recompute=flash_attention_recompute, + flash_attention_causal_mask=flash_attention_causal_mask, + flash_attention_fast_softmax=flash_attention_fast_softmax, + valid_sequence_lengths=valid_sequence_lengths, + cache_idx=cache_idx, + num_virtual_tokens=num_virtual_tokens, + **kwargs, + ) hidden_states = layer_outputs[0] if use_cache: - next_decoder_cache += (layer_outputs[2 if output_attentions else 1],) - - if output_attentions: - all_self_attns += (layer_outputs[1],) + next_decoder_cache += (layer_outputs[1],) hidden_states = self.norm(hidden_states) - # add hidden states from the last decoder layer - if output_hidden_states: - all_hidden_states += (hidden_states,) - next_cache = None if use_cache: next_cache = ( @@ -936,8 +898,6 @@ def forward( return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=next_cache, - hidden_states=all_hidden_states, - attentions=all_self_attns, ) @@ -956,12 +916,10 @@ def forward( input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, + past_key_values: Optional[list[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, logits_to_keep: Union[int, torch.Tensor] = 0, token_idx: Optional[torch.Tensor] = None, @@ -976,17 +934,12 @@ def forward( cache_idx: int = None, lazy_mode: Optional[bool] = True, num_virtual_tokens: int = None, - **kwargs: Unpack[KwargsForCausalLM], + **kwargs: Unpack[TransformersKwargs], ) -> CausalLMOutputWithPast: - 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 - ) if self.generation_config.use_fused_rope is False: global has_fused_rope has_fused_rope = False - # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) outputs: BaseModelOutputWithPast = self.model( input_ids=input_ids, attention_mask=attention_mask, @@ -994,8 +947,6 @@ def forward( past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, cache_position=cache_position, token_idx=token_idx, attn_softmax_bf16=attn_softmax_bf16, @@ -1036,8 +987,8 @@ def forward( @staticmethod def _reorder_cache( - past: Tuple[Tuple[torch.Tensor, torch.Tensor], ...], beam_idx: torch.LongTensor - ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], ...]: + past: tuple[tuple[torch.Tensor, torch.Tensor], ...], beam_idx: torch.LongTensor + ) -> tuple[tuple[torch.Tensor, torch.Tensor], ...]: """ This function is used to re-order the `past_key_values` cache if [`~PreTrainedModel.beam_search`] or [`~PreTrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct diff --git a/optimum/habana/transformers/models/qwen3_moe/modeling_qwen3_moe.py b/optimum/habana/transformers/models/qwen3_moe/modeling_qwen3_moe.py index e11da95567..8dee09080f 100644 --- a/optimum/habana/transformers/models/qwen3_moe/modeling_qwen3_moe.py +++ b/optimum/habana/transformers/models/qwen3_moe/modeling_qwen3_moe.py @@ -19,13 +19,14 @@ # limitations under the License. """PyTorch Qwen3MoE model.""" -from typing import List, Optional, Tuple, Union +from typing import Optional, Union import torch import torch.nn.functional as F from torch import nn from transformers.cache_utils import Cache, DynamicCache, StaticCache from transformers.integrations.deepspeed import is_deepspeed_available +from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask from transformers.modeling_outputs import ( MoeCausalLMOutputWithPast, MoeModelOutputWithPast, @@ -41,7 +42,6 @@ Qwen3MoeSparseMoeBlock, apply_rotary_pos_emb, load_balancing_loss_func, - logger, ) from ....distributed import parallel_state @@ -374,7 +374,7 @@ def reorder_kv_cache(self, beam_idx: torch.LongTensor): def pre_attn_forward( self, hidden_states: torch.Tensor, - position_embeddings: Tuple[torch.Tensor, torch.Tensor], + position_embeddings: tuple[torch.Tensor, torch.Tensor], attention_mask: Optional[torch.Tensor], past_key_value: Optional[Cache] = None, use_cache: bool = False, @@ -390,7 +390,7 @@ def pre_attn_forward( cache_idx: int = None, num_virtual_tokens: int = None, **kwargs, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]: """ The only differences are: - add new args token_idx @@ -494,13 +494,6 @@ def pre_attn_forward( fused_scaled_dot_product_attention = get_gaudi_distributed_attention( self.fused_scaled_dot_product_attention, self.fused_scaled_dot_product_attention_distributed ) - sliding_window = None - if ( - self.config.use_sliding_window - and getattr(self.config, "sliding_window", None) is not None - and self.layer_idx >= self.config.max_window_layers - ): - sliding_window = self.config.sliding_window if use_flash_attention and FusedSDPA is not None: attn_weights = None @@ -560,7 +553,7 @@ def pre_attn_forward( attention_mask, dropout=0.0 if not self.training else self.attention_dropout, scaling=self.scaling, - sliding_window=sliding_window, # main diff with Llama + sliding_window=self.sliding_window, # main diff with Llama attn_softmax_bf16=attn_softmax_bf16, input_shape=input_shape, ) @@ -600,7 +593,7 @@ def __init__(self, config: Qwen3MoeConfig): [Qwen3MoeMLP(config, intermediate_size=config.moe_intermediate_size) for _ in range(self.num_experts)] ) - def forward(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """ - optimize expert forward, remove dynamic control and dynamic shape """ @@ -676,12 +669,10 @@ def forward( hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_value: Optional[Tuple[torch.Tensor]] = None, - output_attentions: Optional[bool] = False, - output_router_logits: Optional[bool] = False, + past_key_value: Optional[tuple[torch.Tensor]] = None, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, - position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC + position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC token_idx: Optional[torch.Tensor] = None, attn_softmax_bf16: Optional[bool] = False, reuse_cache: Optional[bool] = False, @@ -693,7 +684,7 @@ def forward( cache_idx: int = None, num_virtual_tokens: int = None, **kwargs, - ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]: if "padding_mask" in kwargs: warn0( "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`", @@ -705,7 +696,6 @@ def forward( attention_mask=attention_mask, position_ids=position_ids, past_key_value=past_key_value, - output_attentions=output_attentions, use_cache=use_cache, cache_position=cache_position, position_embeddings=position_embeddings, @@ -727,12 +717,8 @@ def forward( hidden_states = self.post_mlp(hidden_states, residual) outputs = (hidden_states,) - if output_attentions: - outputs += (self_attn_weights,) if use_cache: outputs += (present_key_value,) - if output_router_logits: - outputs += (router_logits,) return outputs @@ -741,11 +727,11 @@ def pre_attn( hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_value: Optional[Tuple[torch.Tensor]] = None, + past_key_value: Optional[tuple[torch.Tensor]] = None, output_attentions: Optional[bool] = False, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, - position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, token_idx: Optional[torch.Tensor] = None, attn_softmax_bf16: Optional[bool] = False, reuse_cache: Optional[bool] = False, @@ -757,7 +743,7 @@ def pre_attn( cache_idx: int = None, num_virtual_tokens: int = None, **kwargs, - ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]: hidden_states = self.input_layernorm(hidden_states) hidden_states, attn_weights, present_key_value = self.self_attn.pre_attn_forward( hidden_states=hidden_states, @@ -845,12 +831,9 @@ def forward( input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, + past_key_values: Optional[list[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - output_router_logits: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, attn_softmax_bf16: Optional[bool] = False, @@ -863,6 +846,7 @@ def forward( cache_idx: int = None, lazy_mode: Optional[bool] = True, num_virtual_tokens: int = None, + **kwargs, ) -> MoeModelOutputWithPast: """ Copied from LlamaModel.forward: https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py @@ -876,13 +860,6 @@ def forward( - add new arg flash_attention_fast_softmax - add new arg lazy_mode """ - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_router_logits = ( - output_router_logits if output_router_logits is not None else self.config.output_router_logits - ) - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) use_cache = use_cache if use_cache is not None else self.config.use_cache if (input_ids is None) ^ (inputs_embeds is not None): @@ -894,12 +871,6 @@ def forward( else: raise ValueError("You have to specify either input_ids or inputs_embeds") - if self.gradient_checkpointing and self.training and use_cache: - logger.warning_once( - "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`." - ) - use_cache = False - if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) @@ -950,83 +921,51 @@ def forward( past_seen_tokens, ) else: - causal_mask = self._update_causal_mask(attention_mask, inputs_embeds, cache_position, past_seen_tokens) + mask_function = create_causal_mask if self.config.sliding_window is None else create_sliding_window_causal_mask + causal_mask = mask_function( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_seen_tokens, + position_ids=position_ids, + ) # embed positions hidden_states = inputs_embeds - # decoder layers - all_hidden_states = () if output_hidden_states else None - all_self_attns = () if output_attentions else None - all_router_logits = () if output_router_logits else None next_decoder_cache = () if not use_new_cache else None if lazy_mode: htcore.mark_step() for layer_idx, decoder_layer in enumerate(self.layers): - if output_hidden_states: - all_hidden_states += (hidden_states,) - - if self.gradient_checkpointing and self.training: - layer_outputs = self._gradient_checkpointing_func( - decoder_layer.__call__, - hidden_states, - causal_mask, - position_ids, - past_key_values, - output_attentions, - output_router_logits, - use_cache, - cache_position, - None, - attn_softmax_bf16, - False, - use_flash_attention, - flash_attention_recompute, - flash_attention_causal_mask, - flash_attention_fast_softmax, - valid_sequence_lengths, - None, - ) - else: - layer_outputs = decoder_layer( - hidden_states, - attention_mask=causal_mask, - position_ids=position_ids, - past_key_value=None if past_key_values is None else past_key_values[layer_idx], - output_attentions=output_attentions, - output_router_logits=output_router_logits, - use_cache=use_cache, - cache_position=cache_position, - token_idx=token_idx, - attn_softmax_bf16=attn_softmax_bf16, - reuse_cache=reuse_cache, - use_flash_attention=use_flash_attention, - flash_attention_recompute=flash_attention_recompute, - flash_attention_causal_mask=flash_attention_causal_mask, - flash_attention_fast_softmax=flash_attention_fast_softmax, - valid_sequence_lengths=valid_sequence_lengths, - cache_idx=cache_idx, - num_virtual_tokens=num_virtual_tokens, - ) + layer_outputs = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_ids=position_ids, + past_key_value=None if past_key_values is None else past_key_values[layer_idx], + use_cache=use_cache, + cache_position=cache_position, + token_idx=token_idx, + attn_softmax_bf16=attn_softmax_bf16, + reuse_cache=reuse_cache, + use_flash_attention=use_flash_attention, + flash_attention_recompute=flash_attention_recompute, + flash_attention_causal_mask=flash_attention_causal_mask, + flash_attention_fast_softmax=flash_attention_fast_softmax, + valid_sequence_lengths=valid_sequence_lengths, + cache_idx=cache_idx, + num_virtual_tokens=num_virtual_tokens, + ) hidden_states = layer_outputs[0] if use_cache: - next_decoder_cache += (layer_outputs[2 if output_attentions else 1],) - if output_attentions: - all_self_attns += (layer_outputs[1],) - - if output_router_logits: - all_router_logits += (layer_outputs[-1],) + next_decoder_cache += (layer_outputs[1],) hidden_states = self.norm(hidden_states) - # add hidden states from the last decoder layer - if output_hidden_states: - all_hidden_states += (hidden_states,) - next_cache = None if use_cache: next_cache = ( @@ -1036,9 +975,6 @@ def forward( return MoeModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=next_cache, - hidden_states=all_hidden_states, - attentions=all_self_attns, - router_logits=all_router_logits, ) @@ -1067,12 +1003,10 @@ def forward( input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, + past_key_values: Optional[list[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, output_router_logits: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, logits_to_keep: Union[int, torch.Tensor] = 0, @@ -1088,16 +1022,12 @@ def forward( cache_idx: int = None, lazy_mode: Optional[bool] = True, num_virtual_tokens: int = None, - **loss_kwargs, + **kwargs, ) -> MoeCausalLMOutputWithPast: - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_router_logits = ( output_router_logits if output_router_logits is not None else self.config.output_router_logits ) - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) if self.generation_config.use_fused_rope is False: global has_fused_rope has_fused_rope = False @@ -1110,8 +1040,6 @@ def forward( past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, output_router_logits=output_router_logits, cache_position=cache_position, token_idx=token_idx, @@ -1141,7 +1069,7 @@ def forward( loss = None if labels is not None: - loss = self.loss_function(logits, labels, self.vocab_size, **loss_kwargs) + loss = self.loss_function(logits, labels, self.vocab_size, **kwargs) aux_loss = None if output_router_logits: @@ -1166,8 +1094,8 @@ def forward( @staticmethod def _reorder_cache( - past: Tuple[Tuple[torch.Tensor, torch.Tensor], ...], beam_idx: torch.LongTensor - ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], ...]: + past: tuple[tuple[torch.Tensor, torch.Tensor], ...], beam_idx: torch.LongTensor + ) -> tuple[tuple[torch.Tensor, torch.Tensor], ...]: """ This function is used to re-order the `past_key_values` cache if [`~PreTrainedModel.beam_search`] or [`~PreTrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct diff --git a/optimum/habana/transformers/models/seamless_m4t/modeling_seamless_m4t.py b/optimum/habana/transformers/models/seamless_m4t/modeling_seamless_m4t.py index 7c962c6def..5d22b6e485 100644 --- a/optimum/habana/transformers/models/seamless_m4t/modeling_seamless_m4t.py +++ b/optimum/habana/transformers/models/seamless_m4t/modeling_seamless_m4t.py @@ -1,4 +1,4 @@ -from typing import Optional, Tuple, Union +from typing import Optional, Union import torch import torch.utils.checkpoint @@ -32,11 +32,12 @@ def gaudi_SeamlessM4TAttention_forward( self, hidden_states: torch.Tensor, encoder_hidden_states: Optional[torch.Tensor] = None, - past_key_value: Optional[Tuple[torch.Tensor]] = None, + past_key_value: Optional[tuple[torch.Tensor]] = None, attention_mask: Optional[torch.Tensor] = None, output_attentions: bool = False, + cache_position: Optional[torch.Tensor] = None, token_idx: Optional[torch.Tensor] = None, -) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: +) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]: """ Copied from SeamlessM4TAttention.forward: https://github.com/huggingface/transformers/blob/v4.38.2/src/transformers/models/seamless_m4t/modeling_seamless_m4t.py The only differences are: @@ -85,10 +86,10 @@ def gaudi_SeamlessM4TAttention_forward( value_states = self._shape(self.v_proj(hidden_states), -1, bsz) if self.is_decoder: - # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states. + # if cross_attention save tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states. # Further calls to cross_attention layer can then reuse all cross-attention # key/value_states (first "if" case) - # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of + # if uni-directional self-attention (decoder) save tuple(torch.Tensor, torch.Tensor) of # all previous decoder key/value_states. Further calls to uni-directional self-attention # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case) # if encoder bi-directional self-attention `past_key_value` is always `None` @@ -156,9 +157,10 @@ def gaudi_SeamlessM4TDecoderLayer_forward( attention_mask: Optional[torch.Tensor] = None, encoder_hidden_states: Optional[torch.Tensor] = None, encoder_attention_mask: Optional[torch.Tensor] = None, - past_key_value: Optional[Tuple[torch.Tensor]] = None, + past_key_value: Optional[tuple[torch.Tensor]] = None, output_attentions: Optional[bool] = False, use_cache: Optional[bool] = True, + cache_position: Optional[torch.Tensor] = None, token_idx: Optional[torch.Tensor] = None, ) -> torch.Tensor: """ @@ -178,6 +180,7 @@ def gaudi_SeamlessM4TDecoderLayer_forward( past_key_value=self_attn_past_key_value, attention_mask=attention_mask, output_attentions=output_attentions, + cache_position=cache_position, token_idx=token_idx, ) hidden_states = self.attn_dropout(hidden_states) @@ -199,6 +202,7 @@ def gaudi_SeamlessM4TDecoderLayer_forward( past_key_value=cross_attn_past_key_value, attention_mask=encoder_attention_mask, output_attentions=output_attentions, + cache_position=cache_position, ) hidden_states = self.attn_dropout(hidden_states) hidden_states = residual + hidden_states @@ -230,14 +234,15 @@ def gaudi_SeamlessM4TDecoder_forward( attention_mask: Optional[torch.Tensor] = None, encoder_hidden_states: Optional[torch.FloatTensor] = None, encoder_attention_mask: Optional[torch.LongTensor] = None, - past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, + past_key_values: Optional[tuple[tuple[torch.FloatTensor]]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, + cache_position: Optional[torch.Tensor] = None, token_idx: Optional[torch.Tensor] = None, -) -> Union[Tuple, BaseModelOutputWithPastAndCrossAttentions]: +) -> Union[tuple, BaseModelOutputWithPastAndCrossAttentions]: """ Copied from SeamlessM4TDecoder.forward: https://github.com/huggingface/transformers/blob/v4.38.2/src/transformers/models/seamless_m4t/modeling_seamless_m4t.py The only differences are: @@ -290,13 +295,6 @@ def gaudi_SeamlessM4TDecoder_forward( hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) - if self.gradient_checkpointing and self.training: - if use_cache: - logger.warning_once( - "`use_cache=True` is incompatible with gradient checkpointing`. Setting `use_cache=False`..." - ) - use_cache = False - # decoder layers all_hidden_states = () if output_hidden_states else None all_self_attns = () if output_attentions else None @@ -304,7 +302,7 @@ def gaudi_SeamlessM4TDecoder_forward( next_decoder_cache = () if use_cache else None for idx, decoder_layer in enumerate(self.layers): - # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description) + # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description) if output_hidden_states: all_hidden_states += (hidden_states,) if self.training: @@ -314,28 +312,17 @@ def gaudi_SeamlessM4TDecoder_forward( past_key_value = past_key_values[idx] if past_key_values is not None else None - if self.gradient_checkpointing and self.training: - layer_outputs = self._gradient_checkpointing_func( - decoder_layer.__call__, - hidden_states, - attention_mask, - encoder_hidden_states, - encoder_attention_mask, - None, - output_attentions, - use_cache, - ) - else: - layer_outputs = decoder_layer( - hidden_states, - attention_mask=attention_mask, - encoder_hidden_states=encoder_hidden_states, - encoder_attention_mask=encoder_attention_mask, - past_key_value=past_key_value, - output_attentions=output_attentions, - use_cache=use_cache, - token_idx=token_idx, - ) + layer_outputs = decoder_layer( + hidden_states, + attention_mask=attention_mask, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + token_idx=token_idx, + ) hidden_states = layer_outputs[0] if use_cache: @@ -375,16 +362,17 @@ def gaudi_SeamlessM4TTextToUnitModel_forward( attention_mask: Optional[torch.Tensor] = None, decoder_input_ids: Optional[torch.LongTensor] = None, decoder_attention_mask: Optional[torch.LongTensor] = None, - encoder_outputs: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, - past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, + encoder_outputs: Optional[tuple[tuple[torch.FloatTensor]]] = None, + past_key_values: Optional[tuple[tuple[torch.FloatTensor]]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, decoder_inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, + cache_position: Optional[torch.Tensor] = None, token_idx: Optional[torch.Tensor] = None, -) -> Union[Tuple[torch.Tensor], Seq2SeqModelOutput]: +) -> Union[tuple[torch.Tensor], Seq2SeqModelOutput]: """ Copied from SeamlessM4TTextToUnitModel.forward: https://github.com/huggingface/transformers/blob/v4.38.2/src/transformers/models/seamless_m4t/modeling_seamless_m4t.py The only differences are: @@ -426,6 +414,7 @@ def gaudi_SeamlessM4TTextToUnitModel_forward( output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, + cache_position=cache_position, token_idx=token_idx, ) @@ -450,8 +439,8 @@ def gaudi_SeamlessM4TTextToUnitForConditionalGeneration_forward( attention_mask: Optional[torch.Tensor] = None, decoder_input_ids: Optional[torch.LongTensor] = None, decoder_attention_mask: Optional[torch.LongTensor] = None, - encoder_outputs: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, - past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, + encoder_outputs: Optional[tuple[tuple[torch.FloatTensor]]] = None, + past_key_values: Optional[tuple[tuple[torch.FloatTensor]]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, decoder_inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, @@ -459,8 +448,9 @@ def gaudi_SeamlessM4TTextToUnitForConditionalGeneration_forward( output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, + cache_position: Optional[torch.Tensor] = None, token_idx: Optional[torch.Tensor] = None, -) -> Union[Seq2SeqLMOutput, Tuple[torch.FloatTensor]]: +) -> Union[Seq2SeqLMOutput, tuple[torch.FloatTensor]]: """ Copied from SeamlessM4TTextToUnitForConditionalGeneration.forward: https://github.com/huggingface/transformers/blob/v4.38.2/src/transformers/models/seamless_m4t/modeling_seamless_m4t.py The only differences are: @@ -490,6 +480,7 @@ def gaudi_SeamlessM4TTextToUnitForConditionalGeneration_forward( output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, + cache_position=cache_position, token_idx=token_idx, ) lm_logits = self.lm_head(outputs[0]) @@ -603,8 +594,8 @@ def gaudi_SeamlessM4TForTextToSpeech_forward( attention_mask: Optional[torch.Tensor] = None, decoder_input_ids: Optional[torch.LongTensor] = None, decoder_attention_mask: Optional[torch.LongTensor] = None, - encoder_outputs: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, - past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, + encoder_outputs: Optional[tuple[tuple[torch.FloatTensor]]] = None, + past_key_values: Optional[tuple[tuple[torch.FloatTensor]]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, decoder_inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, @@ -612,8 +603,9 @@ def gaudi_SeamlessM4TForTextToSpeech_forward( output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, + cache_position: Optional[torch.Tensor] = None, token_idx: Optional[torch.Tensor] = None, -) -> Union[Seq2SeqLMOutput, Tuple[torch.FloatTensor]]: +) -> Union[Seq2SeqLMOutput, tuple[torch.FloatTensor]]: """ Copied from SeamlessM4TForTextToSpeech.forward: https://github.com/huggingface/transformers/blob/v4.38.2/src/transformers/models/seamless_m4t/modeling_seamless_m4t.py The only differences are: @@ -672,6 +664,7 @@ def gaudi_SeamlessM4TForTextToSpeech_forward( output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, + cache_position=cache_position, token_idx=token_idx, ) diff --git a/optimum/habana/transformers/models/siglip/modeling_siglip.py b/optimum/habana/transformers/models/siglip/modeling_siglip.py index a90735b083..8ccd64bdb3 100644 --- a/optimum/habana/transformers/models/siglip/modeling_siglip.py +++ b/optimum/habana/transformers/models/siglip/modeling_siglip.py @@ -1,4 +1,4 @@ -from typing import Optional, Tuple, Union +from typing import Optional, Union import torch from torch import nn @@ -74,7 +74,8 @@ def forward( use_flash_attention: Optional[bool] = False, flash_attention_recompute: Optional[bool] = False, flash_attention_fast_softmax: Optional[bool] = False, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + **kwargs, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: """ Copied from CLIPAttention.forward: https://github.com/huggingface/transformers/blob/ab0f050b42d903f34d6eb97f3f8c0c07f0517ad2/src/transformers/models/clip/modeling_clip.py The only differences are: @@ -163,7 +164,7 @@ def forward( output_attentions: Optional[bool] = False, use_flash_attention: Optional[bool] = False, flash_attention_recompute: Optional[bool] = False, - ) -> Tuple[torch.FloatTensor]: + ) -> tuple[torch.FloatTensor]: """ Args: hidden_states (`torch.FloatTensor`): @@ -209,7 +210,7 @@ def forward( return_dict: Optional[bool] = None, use_flash_attention: Optional[bool] = False, flash_attention_recompute: Optional[bool] = False, - ) -> Union[Tuple, BaseModelOutput]: + ) -> Union[tuple, BaseModelOutput]: r""" Args: inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): @@ -245,21 +246,14 @@ def forward( for encoder_layer in self.layers: if output_hidden_states: encoder_states = encoder_states + (hidden_states,) - if self.gradient_checkpointing and self.training: - layer_outputs = self._gradient_checkpointing_func( - encoder_layer.__call__, - hidden_states, - attention_mask, - output_attentions, - ) - else: - layer_outputs = encoder_layer( - hidden_states, - attention_mask, - output_attentions=output_attentions, - use_flash_attention=use_flash_attention, - flash_attention_recompute=flash_attention_recompute, - ) + + layer_outputs = encoder_layer( + hidden_states, + attention_mask, + output_attentions=output_attentions, + use_flash_attention=use_flash_attention, + flash_attention_recompute=flash_attention_recompute, + ) hidden_states = layer_outputs[0] @@ -286,7 +280,7 @@ def forward( interpolate_pos_encoding: Optional[bool] = False, use_flash_attention: Optional[bool] = False, flash_attention_recompute: Optional[bool] = False, - ) -> Union[Tuple, BaseModelOutputWithPooling]: + ) -> Union[tuple, BaseModelOutputWithPooling]: r""" Returns: @@ -332,7 +326,7 @@ def forward( return_dict: Optional[bool] = None, use_flash_attention: Optional[bool] = False, flash_attention_recompute: Optional[bool] = False, - ) -> Union[Tuple, BaseModelOutputWithPooling]: + ) -> Union[tuple, BaseModelOutputWithPooling]: """ Copied from CLIPVisionModel.forward: https://github.com/huggingface/transformers/blob/ab0f050b42d903f34d6eb97f3f8c0c07f0517ad2/src/transformers/models/clip/modeling_clip.py The only differences are: diff --git a/optimum/habana/transformers/models/speecht5/modeling_speecht5.py b/optimum/habana/transformers/models/speecht5/modeling_speecht5.py index 25f47176ed..9af6cdac26 100644 --- a/optimum/habana/transformers/models/speecht5/modeling_speecht5.py +++ b/optimum/habana/transformers/models/speecht5/modeling_speecht5.py @@ -1,4 +1,4 @@ -from typing import List, Optional, Tuple, Union +from typing import Optional, Union import torch import torch.utils.checkpoint @@ -22,13 +22,14 @@ def gaudi_SpeechT5Attention_forward( self, hidden_states: torch.Tensor, key_value_states: Optional[torch.Tensor] = None, - past_key_value: Optional[Tuple[torch.Tensor]] = None, + past_key_value: Optional[tuple[torch.Tensor]] = None, attention_mask: Optional[torch.Tensor] = None, layer_head_mask: Optional[torch.Tensor] = None, position_bias: Optional[torch.Tensor] = None, output_attentions: bool = False, + cache_position: Optional[torch.Tensor] = None, token_idx: Optional[torch.Tensor] = None, -) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: +) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]: """ Copied from SpeechT5Attention.forward: https://github.com/huggingface/transformers/blob/v4.37.2/src/transformers/models/speecht5/modeling_speecht5.py The only differences are: @@ -70,10 +71,10 @@ def gaudi_SpeechT5Attention_forward( value_states = self._shape(self.v_proj(hidden_states), -1, bsz) if self.is_decoder: - # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states. + # if cross_attention save tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states. # Further calls to cross_attention layer can then reuse all cross-attention # key/value_states (first "if" case) - # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of + # if uni-directional self-attention (decoder) save tuple(torch.Tensor, torch.Tensor) of # all previous decoder key/value_states. Further calls to uni-directional self-attention # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case) # if encoder bi-directional self-attention `past_key_value` is always `None` @@ -160,9 +161,10 @@ def gaudi_SpeechT5DecoderLayer_forward( encoder_attention_mask: Optional[torch.Tensor] = None, layer_head_mask: Optional[torch.Tensor] = None, cross_attn_layer_head_mask: Optional[torch.Tensor] = None, - past_key_value: Optional[Tuple[torch.Tensor]] = None, + past_key_value: Optional[tuple[torch.Tensor]] = None, output_attentions: Optional[bool] = False, use_cache: Optional[bool] = True, + cache_position: Optional[torch.Tensor] = None, token_idx: Optional[torch.Tensor] = None, ): """ @@ -182,6 +184,7 @@ def gaudi_SpeechT5DecoderLayer_forward( attention_mask=attention_mask, layer_head_mask=layer_head_mask, output_attentions=output_attentions, + cache_position=cache_position, token_idx=token_idx, ) hidden_states = self.dropout(hidden_states) @@ -203,6 +206,7 @@ def gaudi_SpeechT5DecoderLayer_forward( layer_head_mask=cross_attn_layer_head_mask, past_key_value=cross_attn_past_key_value, output_attentions=output_attentions, + cache_position=cache_position, ) hidden_states = self.dropout(hidden_states) hidden_states = residual + hidden_states @@ -234,13 +238,14 @@ def gaudi_SpeechT5Decoder_forward( encoder_attention_mask: Optional[torch.LongTensor] = None, head_mask: Optional[torch.Tensor] = None, cross_attn_head_mask: Optional[torch.Tensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, + past_key_values: Optional[list[torch.FloatTensor]] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, + cache_position: Optional[torch.Tensor] = None, token_idx: Optional[torch.Tensor] = None, -) -> Union[Tuple, BaseModelOutputWithPastAndCrossAttentions]: +) -> Union[tuple, BaseModelOutputWithPastAndCrossAttentions]: """ Copied from SpeechT5Decoder.forward: https://github.com/huggingface/transformers/blob/v4.37.2/src/transformers/models/speecht5/modeling_speecht5.py The only differences are: @@ -297,7 +302,7 @@ def gaudi_SpeechT5Decoder_forward( if output_hidden_states: all_hidden_states = all_hidden_states + (hidden_states,) - # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description) + # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description) skip_the_layer = False if self.training: dropout_probability = torch.rand([]) @@ -307,32 +312,19 @@ def gaudi_SpeechT5Decoder_forward( past_key_value = past_key_values[idx] if past_key_values is not None else None - if self.gradient_checkpointing and self.training: - layer_outputs = self._gradient_checkpointing_func( - decoder_layer.__call__, - hidden_states, - attention_mask, - encoder_hidden_states, - encoder_attention_mask, - head_mask[idx] if head_mask is not None else None, - cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None, - None, - output_attentions, - use_cache, - ) - else: - layer_outputs = decoder_layer( - hidden_states, - attention_mask=attention_mask, - encoder_hidden_states=encoder_hidden_states, - encoder_attention_mask=encoder_attention_mask, - layer_head_mask=(head_mask[idx] if head_mask is not None else None), - cross_attn_layer_head_mask=(cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None), - past_key_value=past_key_value, - output_attentions=output_attentions, - use_cache=use_cache, - token_idx=token_idx, - ) + layer_outputs = decoder_layer( + hidden_states, + attention_mask=attention_mask, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + layer_head_mask=(head_mask[idx] if head_mask is not None else None), + cross_attn_layer_head_mask=(cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None), + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + token_idx=token_idx, + ) hidden_states = layer_outputs[0] if use_cache: @@ -375,7 +367,7 @@ def gaudi_generate_speech( vocoder: Optional[nn.Module] = None, output_cross_attentions: bool = False, return_output_lengths: bool = False, -) -> Union[torch.FloatTensor, Tuple[torch.FloatTensor, torch.FloatTensor]]: +) -> Union[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor]]: """ Copied from _generate_speech: https://github.com/huggingface/transformers/blob/v4.37.2/src/transformers/models/speecht5/modeling_speecht5.py The only differences are: diff --git a/optimum/habana/transformers/models/stablelm/modeling_stablelm.py b/optimum/habana/transformers/models/stablelm/modeling_stablelm.py index 92fdbd9bbc..c659b24dd9 100644 --- a/optimum/habana/transformers/models/stablelm/modeling_stablelm.py +++ b/optimum/habana/transformers/models/stablelm/modeling_stablelm.py @@ -1,5 +1,5 @@ import math -from typing import List, Optional, Tuple, Union +from typing import Optional, Union import torch from torch import nn @@ -40,9 +40,9 @@ def forward( output_attentions: bool = False, use_cache: bool = False, cache_position: Optional[torch.LongTensor] = None, - position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC + position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC token_idx: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]: """ Copied from StableLmAttention.forward: https://github.com/huggingface/transformers/blob/v4.38.2/src/transformers/models/stablelm/modeling_stablelm.py The only differences are: @@ -173,13 +173,13 @@ def forward( hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_value: Optional[Tuple[torch.Tensor]] = None, + past_key_value: Optional[tuple[torch.Tensor]] = None, output_attentions: Optional[bool] = False, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, - position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC + position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC token_idx: Optional[torch.Tensor] = None, - ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]: """ Copied from StableLmDecoderLayer.forward: https://github.com/huggingface/transformers/blob/v4.38.2/src/transformers/models/stablelm/modeling_stablelm.py The only differences are: @@ -233,7 +233,7 @@ def gaudi_stablelm_model_forward( input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, + past_key_values: Optional[list[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, @@ -312,29 +312,16 @@ def gaudi_stablelm_model_forward( if output_hidden_states: all_hidden_states += (hidden_states,) - if self.gradient_checkpointing and self.training: - layer_outputs = self._gradient_checkpointing_func( - decoder_layer.__call__, - hidden_states, - attention_mask, - position_ids, - past_key_values, - output_attentions, - use_cache, - cache_position, - None, - ) - else: - layer_outputs = decoder_layer( - hidden_states, - attention_mask=attention_mask, - position_ids=position_ids, - past_key_value=past_key_values, - output_attentions=output_attentions, - use_cache=use_cache, - cache_position=cache_position, - token_idx=token_idx, - ) + layer_outputs = decoder_layer( + hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + token_idx=token_idx, + ) hidden_states = layer_outputs[0] @@ -368,7 +355,7 @@ def forward( input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, + past_key_values: Optional[list[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, @@ -389,6 +376,7 @@ def forward( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) outputs: BaseModelOutputWithPast = self.model( input_ids=input_ids, attention_mask=attention_mask, diff --git a/optimum/habana/transformers/models/starcoder2/modeling_starcoder2.py b/optimum/habana/transformers/models/starcoder2/modeling_starcoder2.py index c0435babdf..c5a37cd568 100644 --- a/optimum/habana/transformers/models/starcoder2/modeling_starcoder2.py +++ b/optimum/habana/transformers/models/starcoder2/modeling_starcoder2.py @@ -18,7 +18,7 @@ import math import os -from typing import List, Optional, Tuple, Union +from typing import Optional, Union import torch import torch.nn.functional as F @@ -27,7 +27,6 @@ from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from transformers.models.starcoder2.configuration_starcoder2 import Starcoder2Config from transformers.models.starcoder2.modeling_starcoder2 import ( - KwargsForCausalLM, Starcoder2Attention, Starcoder2DecoderLayer, Starcoder2ForCausalLM, @@ -36,7 +35,7 @@ apply_rotary_pos_emb, ) from transformers.processing_utils import Unpack -from transformers.utils import logging +from transformers.utils import TransformersKwargs, logging from ...modeling_attn_mask_utils import ( _gaudi_prepare_4d_causal_attention_mask, @@ -211,7 +210,7 @@ def gaudi_flash_attn_v1(self, query_layer, key_layer, value_layer, attention_mas def pre_attn_forward( self, hidden_states: torch.Tensor, - position_embeddings: Tuple[torch.Tensor, torch.Tensor], + position_embeddings: tuple[torch.Tensor, torch.Tensor], attention_mask: Optional[torch.Tensor], past_key_value: Optional[Cache] = None, use_cache: bool = False, @@ -224,7 +223,7 @@ def pre_attn_forward( flash_attention_causal_mask: Optional[bool] = False, cache_idx: int = None, **kwargs, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]: """ The only differences are: - add new args token_idx @@ -371,11 +370,10 @@ def forward( hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_value: Optional[Tuple[torch.Tensor]] = None, - output_attentions: Optional[bool] = False, + past_key_value: Optional[tuple[torch.Tensor]] = None, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, - position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC + position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC token_idx: Optional[torch.Tensor] = None, attn_softmax_bf16: Optional[bool] = False, reuse_cache: Optional[bool] = False, @@ -383,14 +381,14 @@ def forward( flash_attention_recompute: Optional[bool] = False, flash_attention_causal_mask: Optional[bool] = False, cache_idx: int = None, - ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + **kwargs, + ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]: residual = hidden_states hidden_states, self_attn_weights, present_key_value = self.pre_attn( hidden_states=hidden_states, attention_mask=attention_mask, position_ids=position_ids, past_key_value=past_key_value, - output_attentions=output_attentions, use_cache=use_cache, cache_position=cache_position, position_embeddings=position_embeddings, @@ -401,6 +399,7 @@ def forward( flash_attention_recompute=flash_attention_recompute, flash_attention_causal_mask=flash_attention_causal_mask, cache_idx=cache_idx, + **kwargs, ) self.self_attn.attention_all_reduce(hidden_states) hidden_states, residual = self.post_attn_pre_mlp(hidden_states, residual) @@ -409,9 +408,6 @@ def forward( outputs = (hidden_states,) - if output_attentions: - outputs += (self_attn_weights,) - if use_cache: outputs += (present_key_value,) @@ -422,11 +418,10 @@ def pre_attn( hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_value: Optional[Tuple[torch.Tensor]] = None, - output_attentions: Optional[bool] = False, + past_key_value: Optional[tuple[torch.Tensor]] = None, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, - position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC + position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC token_idx: Optional[torch.Tensor] = None, attn_softmax_bf16: Optional[bool] = False, reuse_cache: Optional[bool] = False, @@ -434,7 +429,8 @@ def pre_attn( flash_attention_recompute: Optional[bool] = False, flash_attention_causal_mask: Optional[bool] = False, cache_idx: int = None, - ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + **kwargs, + ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]: hidden_states = self.input_layernorm(hidden_states) hidden_states, attn_weights, present_key_value = self.self_attn.pre_attn_forward( hidden_states, @@ -513,11 +509,9 @@ def forward( input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, + past_key_values: Optional[list[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, attn_softmax_bf16: Optional[bool] = False, @@ -527,11 +521,8 @@ def forward( flash_attention_causal_mask: Optional[bool] = False, cache_idx: int = None, lazy_mode: Optional[bool] = True, + **kwargs, ) -> BaseModelOutputWithPast: - 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 - ) use_cache = use_cache if use_cache is not None else self.config.use_cache # retrieve input_ids and inputs_embeds @@ -544,12 +535,6 @@ def forward( else: raise ValueError("You have to specify either input_ids or inputs_embeds") - if self.gradient_checkpointing and self.training and use_cache: - logger.warning_once( - "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`." - ) - use_cache = False - if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) @@ -580,77 +565,45 @@ def forward( hidden_states, p=self.embedding_dropout, training=self.training ) # main diff with Llama - # decoder layers - all_hidden_states = () if output_hidden_states else None - all_self_attns = () if output_attentions else None next_decoder_cache = () if lazy_mode: htcore.mark_step() for layer_idx, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]): - if output_hidden_states: - all_hidden_states += (hidden_states,) - - if self.gradient_checkpointing and self.training: - layer_outputs = self._gradient_checkpointing_func( - decoder_layer.__call__, - hidden_states, - attention_mask, - position_ids, - past_key_values, - output_attentions, - use_cache, - cache_position, - None, - attn_softmax_bf16, - False, - use_flash_attention, - flash_attention_recompute, - flash_attention_causal_mask, - ) - else: - layer_outputs = decoder_layer( - hidden_states, - attention_mask=attention_mask, - position_ids=position_ids, - past_key_value=None if past_key_values is None else past_key_values[layer_idx], - output_attentions=output_attentions, - use_cache=use_cache, - cache_position=cache_position, - token_idx=token_idx, - attn_softmax_bf16=attn_softmax_bf16, - reuse_cache=reuse_cache, - use_flash_attention=use_flash_attention, - flash_attention_recompute=flash_attention_recompute, - flash_attention_causal_mask=flash_attention_causal_mask, - cache_idx=cache_idx, - ) + layer_outputs = decoder_layer( + hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=None if past_key_values is None else past_key_values[layer_idx], + use_cache=use_cache, + cache_position=cache_position, + token_idx=token_idx, + attn_softmax_bf16=attn_softmax_bf16, + reuse_cache=reuse_cache, + use_flash_attention=use_flash_attention, + flash_attention_recompute=flash_attention_recompute, + flash_attention_causal_mask=flash_attention_causal_mask, + cache_idx=cache_idx, + **kwargs, + ) hidden_states = layer_outputs[0] if use_cache: - next_decoder_cache += (layer_outputs[2 if output_attentions else 1],) - - if output_attentions: - all_self_attns += (layer_outputs[1],) + next_decoder_cache += (layer_outputs[1],) hidden_states = self.norm(hidden_states) - # add hidden states from the last decoder layer - if output_hidden_states: - all_hidden_states += (hidden_states,) - next_cache = None if use_cache: next_cache = ( next_decoder_cache.to_legacy_cache() if isinstance(next_decoder_cache, Cache) else next_decoder_cache ) + return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=next_cache, - hidden_states=all_hidden_states, - attentions=all_self_attns, ) @@ -678,12 +631,10 @@ def forward( input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, + past_key_values: Optional[list[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, logits_to_keep: Union[int, torch.Tensor] = 0, token_idx: Optional[torch.Tensor] = None, @@ -695,19 +646,13 @@ def forward( flash_attention_causal_mask: Optional[bool] = False, cache_idx: int = None, lazy_mode: Optional[bool] = True, - **kwargs: Unpack[KwargsForCausalLM], + **kwargs: Unpack[TransformersKwargs], ) -> CausalLMOutputWithPast: - 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 - ) - if not hasattr(self.config, "_attn_implementation"): setattr(self.config, "_attn_implementation", "eager") else: self.config._attn_implementation = "eager" - # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) outputs: BaseModelOutputWithPast = self.model( input_ids=input_ids, attention_mask=attention_mask, @@ -715,8 +660,6 @@ def forward( past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, cache_position=cache_position, token_idx=token_idx, attn_softmax_bf16=attn_softmax_bf16, @@ -726,6 +669,7 @@ def forward( flash_attention_causal_mask=flash_attention_causal_mask, cache_idx=cache_idx, lazy_mode=lazy_mode, + **kwargs, ) hidden_states = outputs.last_hidden_state diff --git a/optimum/habana/transformers/models/t5/modeling_t5.py b/optimum/habana/transformers/models/t5/modeling_t5.py index 69ebbf8e27..dc53afe679 100644 --- a/optimum/habana/transformers/models/t5/modeling_t5.py +++ b/optimum/habana/transformers/models/t5/modeling_t5.py @@ -1,4 +1,4 @@ -from typing import Optional, Tuple, Union +from typing import Optional, Union import habana_frameworks.torch.core as htcore import torch @@ -411,41 +411,22 @@ def gaudi_T5Stack_forward( if output_hidden_states: all_hidden_states = all_hidden_states + (hidden_states,) - if self.gradient_checkpointing and self.training: - layer_outputs = self._gradient_checkpointing_func( - layer_module.__call__, - hidden_states, - extended_attention_mask, - position_bias, - encoder_hidden_states, - encoder_extended_attention_mask, - encoder_decoder_position_bias, - layer_head_mask, - cross_attn_layer_head_mask, - None, # past_key_value is always None with gradient checkpointing - use_cache, - output_attentions, - True, - cache_position, - None, - ) - else: - layer_outputs = layer_module( - hidden_states, - attention_mask=extended_attention_mask, - position_bias=position_bias, - encoder_hidden_states=encoder_hidden_states, - encoder_attention_mask=encoder_extended_attention_mask, - encoder_decoder_position_bias=encoder_decoder_position_bias, - layer_head_mask=layer_head_mask, - cross_attn_layer_head_mask=cross_attn_layer_head_mask, - past_key_value=past_key_value, - use_cache=use_cache, - output_attentions=output_attentions, - return_dict=return_dict, - cache_position=cache_position, - token_idx=token_idx, - ) + layer_outputs = layer_module( + hidden_states, + attention_mask=extended_attention_mask, + position_bias=position_bias, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_extended_attention_mask, + encoder_decoder_position_bias=encoder_decoder_position_bias, + layer_head_mask=layer_head_mask, + cross_attn_layer_head_mask=cross_attn_layer_head_mask, + past_key_value=past_key_value, + use_cache=use_cache, + output_attentions=output_attentions, + return_dict=return_dict, + cache_position=cache_position, + token_idx=token_idx, + ) # layer_outputs is a tuple with: # hidden-states, key-value-states, (self-attention position bias), (self-attention weights), (cross-attention position bias), (cross-attention weights) @@ -506,8 +487,8 @@ def gaudi_T5ForConditionalGeneration_forward( head_mask: Optional[torch.FloatTensor] = None, decoder_head_mask: Optional[torch.FloatTensor] = None, cross_attn_head_mask: Optional[torch.Tensor] = None, - encoder_outputs: Optional[Tuple[Tuple[torch.Tensor]]] = None, - past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None, + encoder_outputs: Optional[tuple[tuple[torch.Tensor]]] = None, + past_key_values: Optional[tuple[tuple[torch.Tensor]]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, decoder_inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, @@ -518,7 +499,7 @@ def gaudi_T5ForConditionalGeneration_forward( cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.LongTensor] = None, **kwargs, -) -> Union[Tuple[torch.FloatTensor], Seq2SeqLMOutput]: +) -> Union[tuple[torch.FloatTensor], Seq2SeqLMOutput]: use_cache = use_cache if use_cache is not None else self.config.use_cache return_dict = return_dict if return_dict is not None else self.config.use_return_dict diff --git a/optimum/habana/transformers/models/video_llava/modeling_video_llava.py b/optimum/habana/transformers/models/video_llava/modeling_video_llava.py index 6670b23375..d5d8db565f 100644 --- a/optimum/habana/transformers/models/video_llava/modeling_video_llava.py +++ b/optimum/habana/transformers/models/video_llava/modeling_video_llava.py @@ -14,7 +14,7 @@ # limitations under the License. """PyTorch VideoLlava model.""" -from typing import List, Optional, Tuple, Union +from typing import Optional, Union import torch from torch import nn @@ -130,7 +130,7 @@ def _get_vision_features( pixel_values_videos: Optional[torch.FloatTensor] = None, vision_feature_layer: Optional[int] = None, vision_feature_select_strategy: Optional[str] = None, - ) -> Union[Tuple, BaseModelOutputWithPooling]: + ) -> Union[tuple, BaseModelOutputWithPooling]: if pixel_values_images is None and pixel_values_videos is None: raise ValueError("You have to specify `pixel_values_images` or `pixel_values_videos`") @@ -162,14 +162,14 @@ def _get_vision_features( def forward( self, - input_ids: Optional[torch.LongTensor] = None, - pixel_values_images: Optional[torch.FloatTensor] = None, - pixel_values_videos: Optional[torch.FloatTensor] = None, + input_ids: torch.LongTensor = None, + pixel_values_images: torch.FloatTensor = None, + pixel_values_videos: torch.FloatTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, + past_key_values: Optional[list[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, - vision_feature_layer: Optional[Union[int, List[int]]] = None, + vision_feature_layer: Optional[Union[int, list[int]]] = None, vision_feature_select_strategy: Optional[str] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, @@ -180,7 +180,7 @@ def forward( logits_to_keep: Union[int, torch.Tensor] = 0, token_idx: Optional[torch.Tensor] = None, **kwargs, - ) -> Union[Tuple, VideoLlavaCausalLMOutputWithPast]: + ) -> Union[tuple, VideoLlavaCausalLMOutputWithPast]: r""" Copied from VideoLlavaForConditionalGeneration.forward: https://github.com/huggingface/transformers/blob/v4.45.2/src/transformers/models/video_llava/modeling_video_llava.py The only differences are: diff --git a/optimum/habana/transformers/models/video_llava/processing_video_llava.py b/optimum/habana/transformers/models/video_llava/processing_video_llava.py index 9ab480220c..72ae665cb1 100644 --- a/optimum/habana/transformers/models/video_llava/processing_video_llava.py +++ b/optimum/habana/transformers/models/video_llava/processing_video_llava.py @@ -1,4 +1,4 @@ -from typing import List, Optional, Union +from typing import Optional, Union from transformers.image_processing_utils import BatchFeature from transformers.image_utils import ImageInput, get_image_size, to_numpy_array @@ -21,6 +21,7 @@ class GaudiVideoLlavaProcessor(VideoLlavaProcessor): def __init__( self, image_processor=None, + video_processor=None, tokenizer=None, patch_size=None, vision_feature_select_strategy=None, @@ -33,7 +34,7 @@ def __init__( Copied from https://github.com/huggingface/transformers/blob/v4.45.2/src/transformers/models/video_llava/processing_video_llava.py#L61 except move super().__init__ to the top of the function so that it will not change the default value for self.patch_size and self.vision_feature_select_strategy """ - super().__init__(image_processor, tokenizer, chat_template=chat_template) + super().__init__(image_processor, video_processor, tokenizer, chat_template=chat_template) self.patch_size = patch_size self.vision_feature_select_strategy = vision_feature_select_strategy self.image_token = image_token @@ -41,7 +42,7 @@ def __init__( def __call__( self, - text: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]] = None, + text: Union[TextInput, PreTokenizedInput, list[TextInput], list[PreTokenizedInput]] = None, images: ImageInput = None, videos: ImageInput = None, padding: Union[bool, str, PaddingStrategy] = False, @@ -53,15 +54,16 @@ def __call__( Copied from https://github.com/huggingface/transformers/blob/v4.45.2/src/transformers/models/video_llava/processing_video_llava.py#L78 Here we use the processor logit in transformers 4.45.2 to avoid the performance regression for op `masked_scatter` in transformers 4.49 on Gaudi3 """ - data = {} - if images is not None or videos is not None: - encoded_images = self.image_processor(images=images, videos=videos, return_tensors=return_tensors) - data.update(encoded_images) if isinstance(text, str): text = [text] elif not isinstance(text, list) and not isinstance(text[0], str): - raise ValueError("Invalid input text. Please provide a string, or a list of strings") + raise TypeError("Invalid input text. Please provide a string, or a list of strings") + + data = {} + if images is not None or videos is not None: + encoded_images = self.image_processor(images=images, videos=videos, return_tensors=return_tensors) + data.update(encoded_images) prompt_strings = text if encoded_images is not None and (self.patch_size is None or self.vision_feature_select_strategy is None): diff --git a/optimum/habana/transformers/models/vit/modeling_vit.py b/optimum/habana/transformers/models/vit/modeling_vit.py index 3988f6ddf8..09ccd1db5b 100644 --- a/optimum/habana/transformers/models/vit/modeling_vit.py +++ b/optimum/habana/transformers/models/vit/modeling_vit.py @@ -15,7 +15,7 @@ # limitations under the License. import math -from typing import Optional, Tuple, Union +from typing import Optional, Union import torch @@ -52,16 +52,32 @@ def gaudi_eager_attention_forward( def gaudi_vit_self_attention_forward( - self, hidden_states, head_mask: Optional[torch.Tensor] = None, output_attentions: bool = False -) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor]]: + self, + hidden_states, + head_mask: Optional[torch.Tensor] = None, + output_attentions: bool = False, +) -> Union[tuple[torch.Tensor, torch.Tensor], tuple[torch.Tensor]]: """ Same method as transformers.models.vit.modeling_vit.ViTSelfAttention.forward with a small tweak: the division is performed before the matmul for computing attention scores. This gives better performance on HPU. """ - key_layer = self.transpose_for_scores(self.key(hidden_states)) - value_layer = self.transpose_for_scores(self.value(hidden_states)) - query_layer = self.transpose_for_scores(self.query(hidden_states)) + batch_size, seq_length, _ = hidden_states.shape + key_layer = ( + self.key(hidden_states) + .view(batch_size, -1, self.num_attention_heads, self.attention_head_size) + .transpose(1, 2) + ) + value_layer = ( + self.value(hidden_states) + .view(batch_size, -1, self.num_attention_heads, self.attention_head_size) + .transpose(1, 2) + ) + query_layer = ( + self.query(hidden_states) + .view(batch_size, -1, self.num_attention_heads, self.attention_head_size) + .transpose(1, 2) + ) context_layer, attention_probs = gaudi_eager_attention_forward( self, diff --git a/optimum/habana/transformers/models/wav2vec2/modeling_wav2vec2.py b/optimum/habana/transformers/models/wav2vec2/modeling_wav2vec2.py index f58cba1baf..58da36dde2 100644 --- a/optimum/habana/transformers/models/wav2vec2/modeling_wav2vec2.py +++ b/optimum/habana/transformers/models/wav2vec2/modeling_wav2vec2.py @@ -16,7 +16,7 @@ import os import random -from typing import Optional, Tuple, Union +from typing import Optional, Union import torch from habana_frameworks.torch.hpu import get_device_name @@ -47,7 +47,7 @@ def _gaudi_wav2vec2_compute_mask_indices( - shape: Tuple[int, int], + shape: tuple[int, int], mask_prob: float, mask_length: int, attention_mask: Optional[torch.LongTensor] = None, @@ -160,7 +160,7 @@ def compute_num_masked_span(input_length): def _gaudi_wav2vec2_sample_negative_indices( - features_shape: Tuple, num_negatives: int, mask_time_indices: Optional[torch.Tensor] = None + features_shape: tuple, num_negatives: int, mask_time_indices: Optional[torch.Tensor] = None ): """ Copied from Transformers: https://github.com/huggingface/transformers/blob/bd469c40659ce76c81f69c7726759d249b4aef49/src/transformers/models/wav2vec2/modeling_wav2vec2.py#L254 @@ -220,12 +220,10 @@ def gaudi_wav2vec2_encoder_forward( expand_attention_mask = attention_mask.unsqueeze(-1).repeat(1, 1, hidden_states.shape[2]) hidden_states[~expand_attention_mask] = 0 - # extend attention_mask - attention_mask = 1.0 - attention_mask[:, None, None, :].to(dtype=hidden_states.dtype) - attention_mask = attention_mask * torch.finfo(hidden_states.dtype).min - attention_mask = attention_mask.expand( - attention_mask.shape[0], 1, attention_mask.shape[-1], attention_mask.shape[-1] - ) + attention_mask = self._update_full_mask( + attention_mask, + hidden_states, + ) position_embeddings = self.pos_conv_embed(hidden_states) hidden_states = hidden_states + position_embeddings @@ -238,23 +236,15 @@ def gaudi_wav2vec2_encoder_forward( if output_hidden_states: all_hidden_states = all_hidden_states + (hidden_states,) - # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description) + # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description) dropout_probability = torch.rand([]) - skip_the_layer = True if self.training and (dropout_probability < self.config.layerdrop) else False + skip_the_layer = self.training and (dropout_probability < self.config.layerdrop) if not skip_the_layer or synced_gpus: # under fsdp or deepspeed zero3 all gpus must run in sync - if self.gradient_checkpointing and self.training: - layer_outputs = self._gradient_checkpointing_func( - layer.__call__, - hidden_states, - attention_mask, - output_attentions, - ) - else: - layer_outputs = layer( - hidden_states, attention_mask=attention_mask, output_attentions=output_attentions - ) + layer_outputs = layer( + hidden_states, attention_mask=attention_mask, output_attentions=output_attentions + ) hidden_states = layer_outputs[0] if skip_the_layer: @@ -283,7 +273,7 @@ def gaudi_wav2vec2_forward( output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, -) -> Union[Tuple, Wav2Vec2BaseModelOutput]: +) -> Union[tuple, Wav2Vec2BaseModelOutput]: """ Copied from Transformers: https://github.com/huggingface/transformers/blob/bd469c40659ce76c81f69c7726759d249b4aef49/src/transformers/models/wav2vec2/modeling_wav2vec2.py#L1282 The only difference is that a clone of `hidden_states` is given to _mask_hidden_states to avoid an error. @@ -393,7 +383,7 @@ def gaudi_wav2vec2forctc_forward( output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, labels: Optional[torch.Tensor] = None, -) -> Union[Tuple, CausalLMOutput]: +) -> Union[tuple, CausalLMOutput]: """ copied from Transformers https://github.com/huggingface/transformers/blob/e770f0316d2a9b787c9d1440f204fcb65e176682/src/transformers/models/wav2vec2/modeling_wav2vec2.py#L1950 only differences are (1) attention_mask tensor generation using ones_like is done on HPU, (2) masked_select is not applied on labels to compute flattened_targets to avoid @@ -506,11 +496,11 @@ def forward( self, hidden_states: torch.Tensor, key_value_states: Optional[torch.Tensor] = None, - past_key_value: Optional[Tuple[torch.Tensor]] = None, + past_key_value: Optional[tuple[torch.Tensor]] = None, attention_mask: Optional[torch.Tensor] = None, layer_head_mask: Optional[torch.Tensor] = None, output_attentions: bool = False, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]: """ Copied from Wav2Vec2SdpaAttention.forward: https://github.com/huggingface/transformers/blob/main/src/transformers/models/wav2vec2/modeling_wav2vec2.py The only difference is If the `USE_FLASH_ATTENTION` switch is enabled, then use the HPU's fused SDPA; otherwise, use PyTorch's native SDPA @@ -567,10 +557,10 @@ def forward( value_states = self._shape(self.v_proj(hidden_states), -1, bsz) if self.is_decoder: - # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states. + # if cross_attention save tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states. # Further calls to cross_attention layer can then reuse all cross-attention # key/value_states (first "if" case) - # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of + # if uni-directional self-attention (decoder) save tuple(torch.Tensor, torch.Tensor) of # all previous decoder key/value_states. Further calls to uni-directional self-attention # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case) # if encoder bi-directional self-attention `past_key_value` is always `None` diff --git a/optimum/habana/transformers/models/xglm/modeling_xglm.py b/optimum/habana/transformers/models/xglm/modeling_xglm.py index daf85e5e73..8d37da54a8 100644 --- a/optimum/habana/transformers/models/xglm/modeling_xglm.py +++ b/optimum/habana/transformers/models/xglm/modeling_xglm.py @@ -1,4 +1,4 @@ -from typing import List, Optional, Tuple, Union +from typing import Optional, Union import torch from torch import nn @@ -18,12 +18,13 @@ def gaudi_xglm_attention_forward( self, hidden_states: torch.Tensor, key_value_states: Optional[torch.Tensor] = None, - past_key_value: Optional[Tuple[torch.Tensor]] = None, + past_key_value: Optional[tuple[torch.Tensor]] = None, attention_mask: Optional[torch.Tensor] = None, layer_head_mask: Optional[torch.Tensor] = None, output_attentions: bool = False, + cache_position: Optional[torch.Tensor] = None, token_idx: Optional[torch.Tensor] = None, -) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: +) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]: """ Copied from XGLMAttention.forward: https://github.com/huggingface/transformers/blob/v4.44.1/src/transformers/models/xglm/modeling_xglm.py The only differences are: @@ -65,10 +66,10 @@ def gaudi_xglm_attention_forward( value_states = self._shape(self.v_proj(hidden_states), -1, bsz) if self.is_decoder: - # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states. + # if cross_attention save tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states. # Further calls to cross_attention layer can then reuse all cross-attention # key/value_states (first "if" case) - # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of + # if uni-directional self-attention (decoder) save tuple(torch.Tensor, torch.Tensor) of # all previous decoder key/value_states. Further calls to uni-directional self-attention # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case) # if encoder bi-directional self-attention `past_key_value` is always `None` @@ -153,9 +154,10 @@ def gaudi_xglm_decoder_layer_forward( encoder_attention_mask: Optional[torch.Tensor] = None, layer_head_mask: Optional[torch.Tensor] = None, cross_attn_layer_head_mask: Optional[torch.Tensor] = None, - past_key_value: Optional[Tuple[torch.Tensor]] = None, + past_key_value: Optional[tuple[torch.Tensor]] = None, output_attentions: Optional[bool] = False, use_cache: Optional[bool] = True, + cache_position: Optional[torch.Tensor] = None, token_idx: Optional[torch.Tensor] = None, ) -> torch.Tensor: """ @@ -176,6 +178,7 @@ def gaudi_xglm_decoder_layer_forward( attention_mask=attention_mask, layer_head_mask=layer_head_mask, output_attentions=output_attentions, + cache_position=cache_position, token_idx=token_idx, ) hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) @@ -197,6 +200,7 @@ def gaudi_xglm_decoder_layer_forward( layer_head_mask=cross_attn_layer_head_mask, past_key_value=cross_attn_past_key_value, output_attentions=output_attentions, + cache_position=cache_position, ) hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) hidden_states = residual + hidden_states @@ -233,14 +237,15 @@ def gaudi_xglm_model_forward( encoder_attention_mask: Optional[torch.Tensor] = None, head_mask: Optional[torch.Tensor] = None, cross_attn_head_mask: Optional[torch.Tensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, + past_key_values: Optional[list[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.Tensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, + cache_position: Optional[torch.Tensor] = None, token_idx: Optional[torch.Tensor] = None, -) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPastAndCrossAttentions]: +) -> Union[tuple[torch.Tensor], BaseModelOutputWithPastAndCrossAttentions]: """ Copied from XGLMModel.forward: https://github.com/huggingface/transformers/blob/v4.44.1/src/transformers/models/xglm/modeling_xglm.py The only differences are: @@ -317,7 +322,7 @@ def gaudi_xglm_model_forward( f" {head_mask.size()[0]}." ) for idx, decoder_layer in enumerate(self.layers): - # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description) + # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description) if output_hidden_states: all_hidden_states += (hidden_states,) if self.training: @@ -327,32 +332,19 @@ def gaudi_xglm_model_forward( past_key_value = past_key_values[idx] if past_key_values is not None else None - if self.gradient_checkpointing and self.training: - layer_outputs = self._gradient_checkpointing_func( - decoder_layer.__call__, - hidden_states, - attention_mask, - encoder_hidden_states, - encoder_attention_mask, - head_mask[idx] if head_mask is not None else None, - cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None, - None, - output_attentions, - use_cache, - ) - else: - layer_outputs = decoder_layer( - hidden_states, - attention_mask=attention_mask, - encoder_hidden_states=encoder_hidden_states, - encoder_attention_mask=encoder_attention_mask, - layer_head_mask=(head_mask[idx] if head_mask is not None else None), - cross_attn_layer_head_mask=(cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None), - past_key_value=past_key_value, - output_attentions=output_attentions, - use_cache=use_cache, - token_idx=token_idx, - ) + layer_outputs = decoder_layer( + hidden_states, + attention_mask, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + layer_head_mask=(head_mask[idx] if head_mask is not None else None), + cross_attn_layer_head_mask=(cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None), + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + token_idx=token_idx, + ) hidden_states = layer_outputs[0] if use_cache: @@ -396,16 +388,17 @@ def forward( encoder_attention_mask: Optional[torch.Tensor] = None, head_mask: Optional[torch.Tensor] = None, cross_attn_head_mask: Optional[torch.Tensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, + past_key_values: Optional[list[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.Tensor] = None, labels: Optional[torch.Tensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, + cache_position: Optional[torch.Tensor] = None, token_idx: Optional[torch.Tensor] = None, **kwargs, - ) -> Union[Tuple[torch.Tensor], CausalLMOutputWithCrossAttentions]: + ) -> Union[tuple[torch.Tensor], CausalLMOutputWithCrossAttentions]: """ Inherits from XGLMForCausalLM: https://github.com/huggingface/transformers/blob/v4.44.1/src/transformers/models/xglm/modeling_xglm.py The only differences are: @@ -433,6 +426,7 @@ def forward( output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, + cache_position=cache_position, token_idx=token_idx, ) diff --git a/optimum/habana/transformers/models/xlm_roberta/modeling_xlm_roberta.py b/optimum/habana/transformers/models/xlm_roberta/modeling_xlm_roberta.py index 2d6d5532d5..bc1956f396 100644 --- a/optimum/habana/transformers/models/xlm_roberta/modeling_xlm_roberta.py +++ b/optimum/habana/transformers/models/xlm_roberta/modeling_xlm_roberta.py @@ -15,7 +15,7 @@ # limitations under the License. """PyTorch XLM-RoBERTa model.""" -from typing import Optional, Tuple +from typing import Optional import torch import torch.utils.checkpoint @@ -33,10 +33,10 @@ def gaudi_XLMRoberta_Sdpa_SelfAttention_forward( attention_mask: Optional[torch.Tensor] = None, head_mask: Optional[torch.FloatTensor] = None, encoder_hidden_states: Optional[torch.FloatTensor] = None, - encoder_attention_mask: Optional[torch.FloatTensor] = None, - past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, + past_key_value: Optional[tuple[tuple[torch.FloatTensor]]] = None, output_attentions: Optional[bool] = False, -) -> Tuple[torch.Tensor]: + cache_position: Optional[torch.Tensor] = None, +) -> tuple[torch.Tensor]: r""" Copied from https://github.com/huggingface/transformers/blob/v4.46.3/src/transformers/models/xlm_roberta/modeling_xlm_roberta.py#L295 Changes: @@ -56,21 +56,20 @@ def gaudi_XLMRoberta_Sdpa_SelfAttention_forward( attention_mask, head_mask, encoder_hidden_states, - encoder_attention_mask, past_key_value, output_attentions, + cache_position, ) bsz, tgt_len, _ = hidden_states.size() - query_layer = self.transpose_for_scores(self.query(hidden_states)) + query_layer = ( + self.query(hidden_states).view(bsz, -1, self.num_attention_heads, self.attention_head_size).transpose(1, 2) + ) - # If this is instantiated as a cross-attention module, the keys and values come from an encoder; the attention - # mask needs to be such that the encoder's padding tokens are not attended to. is_cross_attention = encoder_hidden_states is not None current_states = encoder_hidden_states if is_cross_attention else hidden_states - attention_mask = encoder_attention_mask if is_cross_attention else attention_mask # Check `seq_length` of `past_key_value` == `len(current_states)` to support prefix tuning if is_cross_attention and past_key_value and past_key_value[0].shape[2] == current_states.shape[1]: @@ -85,9 +84,7 @@ def gaudi_XLMRoberta_Sdpa_SelfAttention_forward( if self.is_decoder: past_key_value = (key_layer, value_layer) - is_causal = ( - True if self.is_decoder and not is_cross_attention and attention_mask is None and tgt_len > 1 else False - ) + is_causal = self.is_decoder and not is_cross_attention and attention_mask is None and tgt_len > 1 attn_output = FusedSDPA.apply( query_layer, key_layer, value_layer, attention_mask, 0.0, is_causal, None, "fast", False From 5b39327082e9b4a3889802a8fa83b4360f76d241 Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Wed, 13 Aug 2025 20:55:05 +0000 Subject: [PATCH 17/22] trainer + tests --- .../habana/transformers/generation/utils.py | 32 +- optimum/habana/transformers/modeling_utils.py | 28 +- .../habana/transformers/models/__init__.py | 4 +- .../models/cohere/modeling_cohere.py | 2 +- .../models/decilm/configuration_decilm.py | 80 ++++- .../models/decilm/modeling_decilm.py | 11 - .../models/gpt_neox/modeling_gpt_neox.py | 3 +- .../transformers/models/opt/modeling_opt.py | 4 +- .../models/qwen2/modeling_qwen2.py | 6 +- .../models/qwen2_vl/modeling_qwen2_vl.py | 12 +- .../habana/transformers/models/t5/__init__.py | 3 +- .../transformers/models/t5/modeling_t5.py | 317 ++++++++++-------- .../transformers/models/whisper/__init__.py | 1 - .../models/whisper/modeling_whisper.py | 12 +- optimum/habana/transformers/trainer.py | 152 +++++---- .../habana/transformers/trainer_seq2seq.py | 12 +- optimum/habana/transformers/training_args.py | 52 ++- setup.py | 2 +- tests/test_trainer.py | 98 +++++- tests/test_trainer_distributed.py | 3 +- tests/test_trainer_seq2seq.py | 3 +- 21 files changed, 519 insertions(+), 318 deletions(-) diff --git a/optimum/habana/transformers/generation/utils.py b/optimum/habana/transformers/generation/utils.py index 5697a01e49..894e91bb99 100755 --- a/optimum/habana/transformers/generation/utils.py +++ b/optimum/habana/transformers/generation/utils.py @@ -174,6 +174,34 @@ def get_final_stopping_criteria(x): raise TypeError(f"The stopping criteria should be either a boolean or a torch.tensor but got {type(x)}.") +# Auxiliary functions for beam search +def _temporary_reorder_cache(self, past_key_values, beam_idx): + """ + Temporary function to handle the different types of cache reordering processes while we roll out `Cache`. + + TODO: standardize cache formats and make all models compatible with `Cache`. It would remove the need + for this function, with `Cache.reorder_cache` being the sole remaining code path + """ + model_class = self.__class__.__name__.lower() + # Exception 1: code path for models using the legacy cache format + if isinstance(past_key_values, (tuple, list)): + past_key_values = self._reorder_cache(past_key_values, beam_idx) + # Exception 2: models with different cache formats. These are limited to `DynamicCache` until their + # cache format is standardized, to avoid adding complexity to the codebase. + elif "gptbigcode" in model_class: + if not isinstance(past_key_values, (DynamicCache, EncoderDecoderCache)): + raise ValueError( + f"Using an unsupported cache format with {model_class}. Currently, it only supports the " + "legacy tuple format or `DynamicCache`" + ) + past_key_values = self._reorder_cache(past_key_values, beam_idx) + past_key_values = DynamicCache.from_legacy_cache(past_key_values) + # Standard code path: use the `Cache.reorder_cache` + else: + past_key_values.reorder_cache(beam_idx) + return past_key_values + + class GaudiGenerationMixin(GenerationMixin): """ This class enables to perform fast generation in lazy mode and with HPU graphs. @@ -3404,8 +3432,8 @@ def expand_if_needed(tensor, new_size, value, dim=-1): if model_kwargs["reuse_cache"]: model_kwargs["past_key_values"] = unwrap_deepspeed_model(self).reorder_kv_cache(beam_idx) else: - model_kwargs["past_key_values"] = self._temporary_reorder_cache( - model_kwargs["past_key_values"], beam_idx + model_kwargs["past_key_values"] = _temporary_reorder_cache( + self, model_kwargs["past_key_values"], beam_idx ) if return_dict_in_generate and output_scores: diff --git a/optimum/habana/transformers/modeling_utils.py b/optimum/habana/transformers/modeling_utils.py index 66802f7135..29c7864a3d 100644 --- a/optimum/habana/transformers/modeling_utils.py +++ b/optimum/habana/transformers/modeling_utils.py @@ -42,7 +42,6 @@ ) from .loss import gaudi_RTDetrHungarianMatcher_forward from .models import ( - GAUDI_WHISPER_ATTENTION_CLASSES, BaichuanConfig, BaichuanForCausalLM, BaichuanTokenizer, @@ -298,8 +297,7 @@ gaudi_t5_layernorm_forward, gaudi_T5Attention_forward, gaudi_T5Block_forward, - gaudi_T5ForConditionalGeneration_forward, - gaudi_T5ForConditionalGeneration_prepare_inputs_for_generation, + GaudiT5ForConditionalGeneration, gaudi_T5LayerSelfAttention_forward, gaudi_T5Stack_forward, gaudi_table_transformer_conv_encoder_forward, @@ -350,10 +348,7 @@ def adapt_transformers_to_gaudi(): transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2Encoder.forward = gaudi_wav2vec2_encoder_forward transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2ForCTC.forward = gaudi_wav2vec2forctc_forward transformers.models.wav2vec2.modeling_wav2vec2.TDNNLayer.forward = gaudi_wav2vec2_tdnnlayer_forward - transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2SdpaAttention = GaudiWav2Vec2SdpaAttention - transformers.models.wav2vec2.modeling_wav2vec2.WAV2VEC2_ATTENTION_CLASSES.update( - {"sdpa": GaudiWav2Vec2SdpaAttention} - ) + transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2Attention = GaudiWav2Vec2SdpaAttention # Generation is modified to run faster in lazy mode transformers.generation.GenerationMixin.prepare_inputs_for_generation = ( @@ -489,9 +484,6 @@ def adapt_transformers_to_gaudi(): transformers.models.gpt_bigcode.modeling_gpt_bigcode.GPTBigCodeForCausalLM = GaudiGPTBigCodeForCausalLM transformers.models.gpt_bigcode.modeling_gpt_bigcode.GPTBigCodeBlock.forward = gaudi_gpt_bigcode_block_forward transformers.models.gpt_bigcode.modeling_gpt_bigcode.GPTBigCodeModel.forward = gaudi_gpt_bigcode_model_forward - transformers.models.gpt_bigcode.modeling_gpt_bigcode.GPTBIGCODE_ATTENTION_CLASSES.update( - {"eager": GaudiGPTBigCodeAttention} - ) # Optimization for gpt-neo generation on Gaudi transformers.models.gpt_neo.modeling_gpt_neo.GPTNeoForCausalLM = GaudiGPTNeoForCausalLM @@ -566,10 +558,7 @@ def adapt_transformers_to_gaudi(): transformers.models.t5.modeling_t5.T5LayerNorm.forward = gaudi_t5_layernorm_forward transformers.models.t5.modeling_t5.T5Stack.forward = gaudi_T5Stack_forward transformers.models.t5.modeling_t5.T5LayerSelfAttention.forward = gaudi_T5LayerSelfAttention_forward - transformers.models.t5.modeling_t5.T5ForConditionalGeneration.forward = gaudi_T5ForConditionalGeneration_forward - transformers.models.t5.modeling_t5.T5ForConditionalGeneration.prepare_inputs_for_generation = ( - gaudi_T5ForConditionalGeneration_prepare_inputs_for_generation - ) + transformers.models.t5.modeling_t5.T5ForConditionalGeneration = GaudiT5ForConditionalGeneration transformers.models.t5.modeling_t5.T5Attention.forward = gaudi_T5Attention_forward transformers.models.t5.modeling_t5.T5Block.forward = gaudi_T5Block_forward @@ -784,7 +773,6 @@ def adapt_transformers_to_gaudi(): transformers.models.whisper.modeling_whisper.WhisperDecoder = GaudiWhisperDecoder transformers.models.whisper.modeling_whisper.WhisperModel = GaudiWhisperModel transformers.models.whisper.modeling_whisper.WhisperForConditionalGeneration = GaudiWhisperForConditionalGeneration - transformers.models.whisper.modeling_whisper.WHISPER_ATTENTION_CLASSES = GAUDI_WHISPER_ATTENTION_CLASSES # Optimization for mllama on Gaudi transformers.models.mllama.modeling_mllama.MllamaSelfAttentionDecoderLayer = GaudiMllamaSelfAttentionDecoderLayer @@ -799,13 +787,13 @@ def adapt_transformers_to_gaudi(): transformers.models.mllama.modeling_mllama.MllamaVisionEncoderLayer = GaudiMllamaVisionEncoderLayer transformers.models.mllama.modeling_mllama.MllamaVisionSdpaAttention = GaudiMllamaVisionSdpaAttention - transformers.AutoConfig.register("deci", DeciLMConfig) - transformers.AutoModelForCausalLM.register(DeciLMConfig, DeciLMForCausalLM) + # transformers.AutoConfig.register("deci", DeciLMConfig) + # transformers.AutoModelForCausalLM.register(DeciLMConfig, DeciLMForCausalLM) # Optimization for deepseek on Gaudi - transformers.AutoConfig.register("deepseek_v2", DeepseekV2Config) - transformers.AutoModelForCausalLM.register(DeepseekV2Config, DeepseekV2ForCausalLM) - transformers.AutoTokenizer.register(DeepseekV2Config, fast_tokenizer_class=DeepseekTokenizerFast) + # transformers.AutoConfig.register("deepseek_v2", DeepseekV2Config) + # transformers.AutoModelForCausalLM.register(DeepseekV2Config, DeepseekV2ForCausalLM) + # transformers.AutoTokenizer.register(DeepseekV2Config, fast_tokenizer_class=DeepseekTokenizerFast) # transformers.AutoConfig.register("deepseek_v3", DeepseekV3Config) # transformers.AutoModelForCausalLM.register(DeepseekV3Config, DeepseekV3ForCausalLM) transformers.models.deepseek_v3.configuration_deepseek_v3.DeepseekV3Config = DeepseekV3Config diff --git a/optimum/habana/transformers/models/__init__.py b/optimum/habana/transformers/models/__init__.py index 81ebae6e51..d6c8d62f40 100644 --- a/optimum/habana/transformers/models/__init__.py +++ b/optimum/habana/transformers/models/__init__.py @@ -335,8 +335,7 @@ gaudi_t5_layernorm_forward, gaudi_T5Attention_forward, gaudi_T5Block_forward, - gaudi_T5ForConditionalGeneration_forward, - gaudi_T5ForConditionalGeneration_prepare_inputs_for_generation, + GaudiT5ForConditionalGeneration, gaudi_T5LayerSelfAttention_forward, gaudi_T5Stack_forward, ) @@ -358,7 +357,6 @@ gaudi_wav2vec2forctc_forward, ) from .whisper import ( - GAUDI_WHISPER_ATTENTION_CLASSES, GaudiWhisperDecoder, GaudiWhisperDecoderLayer, GaudiWhisperForConditionalGeneration, diff --git a/optimum/habana/transformers/models/cohere/modeling_cohere.py b/optimum/habana/transformers/models/cohere/modeling_cohere.py index e2cad79930..602e209c90 100644 --- a/optimum/habana/transformers/models/cohere/modeling_cohere.py +++ b/optimum/habana/transformers/models/cohere/modeling_cohere.py @@ -1,6 +1,7 @@ from typing import Optional, Union import torch +from transformers.cache_utils import StaticCache from transformers.masking_utils import create_causal_mask from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from transformers.models.cohere.modeling_cohere import ( @@ -11,7 +12,6 @@ CohereForCausalLM, CohereRotaryEmbedding, DynamicCache, - StaticCache, apply_rotary_pos_emb, eager_attention_forward, ) diff --git a/optimum/habana/transformers/models/decilm/configuration_decilm.py b/optimum/habana/transformers/models/decilm/configuration_decilm.py index 35bbdad6f8..71f6ee02cc 100644 --- a/optimum/habana/transformers/models/decilm/configuration_decilm.py +++ b/optimum/habana/transformers/models/decilm/configuration_decilm.py @@ -3,10 +3,10 @@ https://huggingface.co/Deci/DeciLM-7B/blob/main/configuration_decilm.py """ -from transformers.models.llama.configuration_llama import LlamaConfig +from transformers.configuration_utils import PretrainedConfig -class DeciLMConfig(LlamaConfig): +class DeciLMConfig(PretrainedConfig): r""" Args: num_key_value_heads_per_layer (`List[int]`): @@ -15,10 +15,84 @@ class DeciLMConfig(LlamaConfig): model_type = "deci" + keys_to_ignore_at_inference = ["past_key_values"] + # Default tensor parallel plan for base model `LlamaModel` + base_model_tp_plan = { + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise", + } + base_model_pp_plan = { + "embed_tokens": (["input_ids"], ["inputs_embeds"]), + "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), + "norm": (["hidden_states"], ["hidden_states"]), + } + def __init__( self, + vocab_size=32000, + hidden_size=4096, + intermediate_size=11008, + num_hidden_layers=32, + num_attention_heads=32, + num_key_value_heads=None, + hidden_act="silu", + max_position_embeddings=2048, + initializer_range=0.02, + rms_norm_eps=1e-6, + use_cache=True, + pad_token_id=None, + bos_token_id=1, + eos_token_id=2, + pretraining_tp=1, + tie_word_embeddings=False, + rope_theta=10000.0, + rope_scaling=None, + attention_bias=False, + attention_dropout=0.0, + mlp_bias=False, + head_dim=None, num_key_value_heads_per_layer: list = None, **kwargs, ): + self.vocab_size = vocab_size + self.max_position_embeddings = max_position_embeddings + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + + # for backward compatibility + if num_key_value_heads is None: + num_key_value_heads = num_attention_heads + + self.num_key_value_heads = num_key_value_heads + self.hidden_act = hidden_act + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.pretraining_tp = pretraining_tp + self.use_cache = use_cache + self.rope_theta = rope_theta + self.rope_scaling = rope_scaling + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.mlp_bias = mlp_bias + self.head_dim = head_dim if head_dim is not None else self.hidden_size // self.num_attention_heads self.num_key_value_heads_per_layer = num_key_value_heads_per_layer - super().__init__(**kwargs) + # Validate the correctness of rotary position embeddings parameters + # BC: if there is a 'type' field, copy it it to 'rope_type'. + if self.rope_scaling is not None and "type" in self.rope_scaling: + self.rope_scaling["rope_type"] = self.rope_scaling["type"] + rope_config_validation(self) + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) diff --git a/optimum/habana/transformers/models/decilm/modeling_decilm.py b/optimum/habana/transformers/models/decilm/modeling_decilm.py index 6618911530..8a734404c3 100644 --- a/optimum/habana/transformers/models/decilm/modeling_decilm.py +++ b/optimum/habana/transformers/models/decilm/modeling_decilm.py @@ -14,8 +14,6 @@ from transformers.modeling_outputs import CausalLMOutputWithPast from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES from transformers.models.llama.modeling_llama import ( - LLAMA_INPUTS_DOCSTRING, - LLAMA_START_DOCSTRING, BaseModelOutputWithPast, LlamaAttention, LlamaDecoderLayer, @@ -228,20 +226,12 @@ def forward( return outputs -@add_start_docstrings( - "The bare DeciLM Model outputting raw hidden-states without any specific head on top.", - LLAMA_START_DOCSTRING, -) class DeciLMPreTrainedModel(LlamaPreTrainedModel): config_class = DeciLMConfig _no_split_modules = ["DeciLMDecoderLayer"] _keys_to_ignore_on_load_missing = ["self_attn.rotary_emb.inv_freq"] -@add_start_docstrings( - "The bare DeciLM Model outputting raw hidden-states without any specific head on top.", - LLAMA_START_DOCSTRING, -) class DeciLMModel(LlamaModel, DeciLMPreTrainedModel): """ Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`DeciLMDecoderLayer`] @@ -265,7 +255,6 @@ def __init__(self, config: DeciLMConfig): # Initialize weights and apply final processing self.post_init() - @add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING) def forward( self, input_ids: torch.LongTensor = None, diff --git a/optimum/habana/transformers/models/gpt_neox/modeling_gpt_neox.py b/optimum/habana/transformers/models/gpt_neox/modeling_gpt_neox.py index 5de06556a0..c10875e0c1 100644 --- a/optimum/habana/transformers/models/gpt_neox/modeling_gpt_neox.py +++ b/optimum/habana/transformers/models/gpt_neox/modeling_gpt_neox.py @@ -10,7 +10,6 @@ GPTNeoXLayer, GPTNeoXMLP, GPTNeoXModel, - KwargsForCausalLM, apply_rotary_pos_emb, logger, ) @@ -404,7 +403,7 @@ def forward( cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, logits_to_keep: Union[int, torch.Tensor] = 0, - **kwargs: Unpack[KwargsForCausalLM], + **kwargs, ) -> Union[tuple, CausalLMOutputWithPast]: outputs: BaseModelOutputWithPast = self.gpt_neox( input_ids, diff --git a/optimum/habana/transformers/models/opt/modeling_opt.py b/optimum/habana/transformers/models/opt/modeling_opt.py index bd11f5ca56..408953efca 100644 --- a/optimum/habana/transformers/models/opt/modeling_opt.py +++ b/optimum/habana/transformers/models/opt/modeling_opt.py @@ -5,7 +5,7 @@ from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from transformers.models.opt.configuration_opt import OPTConfig from transformers.models.opt.modeling_opt import ( - OPT_ATTENTION_CLASSES, + OPTAttention, OPTForCausalLM, OPTLearnedPositionalEmbedding, logger, @@ -174,7 +174,7 @@ def __init__(self, config: OPTConfig, layer_idx: Optional[int] = None): super().__init__() self.embed_dim = config.hidden_size - self.self_attn = OPT_ATTENTION_CLASSES["eager"](config=config, layer_idx=layer_idx) + self.self_attn = OPTAttention(config=config, layer_idx=layer_idx) self.do_layer_norm_before = config.do_layer_norm_before self.dropout = config.dropout diff --git a/optimum/habana/transformers/models/qwen2/modeling_qwen2.py b/optimum/habana/transformers/models/qwen2/modeling_qwen2.py index 3bad0f8b11..5057b14665 100644 --- a/optimum/habana/transformers/models/qwen2/modeling_qwen2.py +++ b/optimum/habana/transformers/models/qwen2/modeling_qwen2.py @@ -38,10 +38,9 @@ Qwen2Model, Qwen2RMSNorm, apply_rotary_pos_emb, - logger, ) from transformers.processing_utils import Unpack -from transformers.utils import TransformersKwargs +from transformers.utils import TransformersKwargs, logging from ....distributed import parallel_state from ...modeling_attn_mask_utils import ( @@ -76,6 +75,9 @@ import habana_frameworks.torch.core as htcore +logger = logging.get_logger(__name__) + + def gaudi_qwen2_rmsnorm_forward(self, hidden_states): if hidden_states.device.type == "hpu" and has_fused_rms_norm: # mixed dtypes are not good for FusedRMSNorm, both inputs need to have same dtype diff --git a/optimum/habana/transformers/models/qwen2_vl/modeling_qwen2_vl.py b/optimum/habana/transformers/models/qwen2_vl/modeling_qwen2_vl.py index cdb1259db1..f8c63fe518 100644 --- a/optimum/habana/transformers/models/qwen2_vl/modeling_qwen2_vl.py +++ b/optimum/habana/transformers/models/qwen2_vl/modeling_qwen2_vl.py @@ -25,14 +25,13 @@ ) from transformers.models.qwen2_vl.modeling_qwen2_vl import ( Qwen2VisionTransformerPretrainedModel, + Qwen2VLAttention, Qwen2VLCausalLMOutputWithPast, Qwen2VLConfig, Qwen2VLDecoderLayer, Qwen2VLForConditionalGeneration, Qwen2VLModel, - Qwen2VLSdpaAttention, Qwen2VLVisionBlock, - VisionSdpaAttention, apply_multimodal_rotary_pos_emb, apply_rotary_pos_emb_vision, repeat_kv, @@ -59,9 +58,12 @@ def forward(self, query, key, value, attn_mask, dropout_p, is_casual, scale, sof # from: https://github.com/huggingface/transformers/blob/v4.45.2/src/transformers/models/qwen2_vl/modeling_qwen2_vl.py#L383 -class GaudiVisionSdpaAttention(VisionSdpaAttention): +class GaudiVisionSdpaAttention(torch.nn.Module): def __init__(self, dim: int, num_heads: int = 16) -> None: - super().__init__(dim, num_heads) + super().__init__() + self.num_heads = num_heads + self.qkv = torch.nn.Linear(dim, dim * 3, bias=True) + self.proj = torch.nn.Linear(dim, dim) self.fused_scaled_dot_product_attention = ModuleFusedSDPA(FusedSDPA) if FusedSDPA else None def forward( @@ -149,7 +151,7 @@ def forward( # from: https://github.com/huggingface/transformers/blob/v4.45.2/src/transformers/models/qwen2_vl/modeling_qwen2_vl.py#L821 -class GaudiQwen2VLSdpaAttention(Qwen2VLSdpaAttention): +class GaudiQwen2VLSdpaAttention(Qwen2VLAttention): """ Qwen2 attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from `Qwen2Attention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to diff --git a/optimum/habana/transformers/models/t5/__init__.py b/optimum/habana/transformers/models/t5/__init__.py index e92116128d..a6379c20fa 100644 --- a/optimum/habana/transformers/models/t5/__init__.py +++ b/optimum/habana/transformers/models/t5/__init__.py @@ -2,8 +2,7 @@ gaudi_t5_layernorm_forward, gaudi_T5Attention_forward, gaudi_T5Block_forward, - gaudi_T5ForConditionalGeneration_forward, - gaudi_T5ForConditionalGeneration_prepare_inputs_for_generation, + GaudiT5ForConditionalGeneration, gaudi_T5LayerSelfAttention_forward, gaudi_T5Stack_forward, ) diff --git a/optimum/habana/transformers/models/t5/modeling_t5.py b/optimum/habana/transformers/models/t5/modeling_t5.py index dc53afe679..3c92a336dd 100644 --- a/optimum/habana/transformers/models/t5/modeling_t5.py +++ b/optimum/habana/transformers/models/t5/modeling_t5.py @@ -9,7 +9,7 @@ BaseModelOutputWithPastAndCrossAttentions, Seq2SeqLMOutput, ) -from transformers.models.t5.modeling_t5 import __HEAD_MASK_WARNING_MSG +from transformers.models.t5.modeling_t5 import __HEAD_MASK_WARNING_MSG, T5ForConditionalGeneration from transformers.utils import ( logging, ) @@ -478,154 +478,185 @@ def gaudi_T5Stack_forward( ) -def gaudi_T5ForConditionalGeneration_forward( - self, - input_ids: Optional[torch.LongTensor] = None, - attention_mask: Optional[torch.FloatTensor] = None, - decoder_input_ids: Optional[torch.LongTensor] = None, - decoder_attention_mask: Optional[torch.BoolTensor] = None, - head_mask: Optional[torch.FloatTensor] = None, - decoder_head_mask: Optional[torch.FloatTensor] = None, - cross_attn_head_mask: Optional[torch.Tensor] = None, - encoder_outputs: Optional[tuple[tuple[torch.Tensor]]] = None, - past_key_values: Optional[tuple[tuple[torch.Tensor]]] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - decoder_inputs_embeds: Optional[torch.FloatTensor] = None, - labels: Optional[torch.LongTensor] = None, - use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - cache_position: Optional[torch.LongTensor] = None, - token_idx: Optional[torch.LongTensor] = None, - **kwargs, -) -> Union[tuple[torch.FloatTensor], Seq2SeqLMOutput]: - use_cache = use_cache if use_cache is not None else self.config.use_cache - return_dict = return_dict if return_dict is not None else self.config.use_return_dict +class GaudiT5ForConditionalGeneration(T5ForConditionalGeneration): + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.FloatTensor] = None, + decoder_input_ids: Optional[torch.LongTensor] = None, + decoder_attention_mask: Optional[torch.BoolTensor] = None, + head_mask: Optional[torch.FloatTensor] = None, + decoder_head_mask: Optional[torch.FloatTensor] = None, + cross_attn_head_mask: Optional[torch.Tensor] = None, + encoder_outputs: Optional[tuple[tuple[torch.Tensor]]] = None, + past_key_values: Optional[tuple[tuple[torch.Tensor]]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + decoder_inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + token_idx: Optional[torch.LongTensor] = None, + **kwargs, + ) -> Union[tuple[torch.FloatTensor], Seq2SeqLMOutput]: + use_cache = use_cache if use_cache is not None else self.config.use_cache + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # FutureWarning: head_mask was separated into two input args - head_mask, decoder_head_mask + if head_mask is not None and decoder_head_mask is None: + if self.config.num_layers == self.config.num_decoder_layers: + warn0(__HEAD_MASK_WARNING_MSG, category=FutureWarning, state=self.accelerate.state) + decoder_head_mask = head_mask + + # Encode if needed (training, first prediction pass) + if encoder_outputs is None: + # Convert encoder inputs in embeddings if needed + encoder_outputs = self.encoder( + input_ids=input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + head_mask=head_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + elif return_dict and not isinstance(encoder_outputs, BaseModelOutput): + encoder_outputs = BaseModelOutput( + last_hidden_state=encoder_outputs[0], + hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None, + attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None, + ) - # FutureWarning: head_mask was separated into two input args - head_mask, decoder_head_mask - if head_mask is not None and decoder_head_mask is None: - if self.config.num_layers == self.config.num_decoder_layers: - warn0(__HEAD_MASK_WARNING_MSG, category=FutureWarning, state=self.accelerate.state) - decoder_head_mask = head_mask - - # Encode if needed (training, first prediction pass) - if encoder_outputs is None: - # Convert encoder inputs in embeddings if needed - encoder_outputs = self.encoder( - input_ids=input_ids, - attention_mask=attention_mask, - inputs_embeds=inputs_embeds, - head_mask=head_mask, + hidden_states = encoder_outputs[0] + + if labels is not None and decoder_input_ids is None and decoder_inputs_embeds is None: + # get decoder inputs from shifting lm labels to the right + decoder_input_ids = self._shift_right(labels) + + # Decode + decoder_outputs = self.decoder( + input_ids=decoder_input_ids, + attention_mask=decoder_attention_mask, + inputs_embeds=decoder_inputs_embeds, + past_key_values=past_key_values, + encoder_hidden_states=hidden_states, + encoder_attention_mask=attention_mask, + head_mask=decoder_head_mask, + cross_attn_head_mask=cross_attn_head_mask, + use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, + cache_position=cache_position, + token_idx=token_idx, ) - elif return_dict and not isinstance(encoder_outputs, BaseModelOutput): - encoder_outputs = BaseModelOutput( - last_hidden_state=encoder_outputs[0], - hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None, - attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None, - ) - - hidden_states = encoder_outputs[0] - - if labels is not None and decoder_input_ids is None and decoder_inputs_embeds is None: - # get decoder inputs from shifting lm labels to the right - decoder_input_ids = self._shift_right(labels) - - # Decode - decoder_outputs = self.decoder( - input_ids=decoder_input_ids, - attention_mask=decoder_attention_mask, - inputs_embeds=decoder_inputs_embeds, - past_key_values=past_key_values, - encoder_hidden_states=hidden_states, - encoder_attention_mask=attention_mask, - head_mask=decoder_head_mask, - cross_attn_head_mask=cross_attn_head_mask, - use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - cache_position=cache_position, - token_idx=token_idx, - ) - - sequence_output = decoder_outputs[0] - - if self.config.tie_word_embeddings: - # Rescale output before projecting on vocab - # See https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/transformer/transformer.py#L586 - sequence_output = sequence_output * (self.model_dim**-0.5) - lm_logits = self.lm_head(sequence_output) - - loss = None - if labels is not None: - loss_fct = CrossEntropyLoss(ignore_index=-100) - # move labels to correct device to enable PP - labels = labels.to(lm_logits.device) - loss = loss_fct(lm_logits.view(-1, lm_logits.size(-1)), labels.view(-1)) - # TODO(thom): Add z_loss https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/layers.py#L666 - - if not return_dict: - output = (lm_logits,) + decoder_outputs[1:] + encoder_outputs - return ((loss,) + output) if loss is not None else output - - return Seq2SeqLMOutput( - loss=loss, - logits=lm_logits, - past_key_values=decoder_outputs.past_key_values, - decoder_hidden_states=decoder_outputs.hidden_states, - decoder_attentions=decoder_outputs.attentions, - cross_attentions=decoder_outputs.cross_attentions, - encoder_last_hidden_state=encoder_outputs.last_hidden_state, - encoder_hidden_states=encoder_outputs.hidden_states, - encoder_attentions=encoder_outputs.attentions, - ) + sequence_output = decoder_outputs[0] + + if self.config.tie_word_embeddings: + # Rescale output before projecting on vocab + # See https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/transformer/transformer.py#L586 + sequence_output = sequence_output * (self.model_dim**-0.5) + + lm_logits = self.lm_head(sequence_output) + + loss = None + if labels is not None: + loss_fct = CrossEntropyLoss(ignore_index=-100) + # move labels to correct device to enable PP + labels = labels.to(lm_logits.device) + loss = loss_fct(lm_logits.view(-1, lm_logits.size(-1)), labels.view(-1)) + # TODO(thom): Add z_loss https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/layers.py#L666 + + if not return_dict: + output = (lm_logits,) + decoder_outputs[1:] + encoder_outputs + return ((loss,) + output) if loss is not None else output + + return Seq2SeqLMOutput( + loss=loss, + logits=lm_logits, + past_key_values=decoder_outputs.past_key_values, + decoder_hidden_states=decoder_outputs.hidden_states, + decoder_attentions=decoder_outputs.attentions, + cross_attentions=decoder_outputs.cross_attentions, + encoder_last_hidden_state=encoder_outputs.last_hidden_state, + encoder_hidden_states=encoder_outputs.hidden_states, + encoder_attentions=encoder_outputs.attentions, + ) -def gaudi_T5ForConditionalGeneration_prepare_inputs_for_generation( - self, - input_ids, - past_key_values=None, - attention_mask=None, - head_mask=None, - decoder_head_mask=None, - decoder_attention_mask=None, - cross_attn_head_mask=None, - use_cache=None, - encoder_outputs=None, - token_idx=None, - **kwargs, -): - # cut decoder_input_ids if past_key_values is used - if past_key_values is not None: - if token_idx is not None: - idx = token_idx + kwargs.get("inputs_embeds_offset", 0) - 1 - input_ids = torch.index_select(input_ids, 1, idx) - else: - past_length = past_key_values[0][0].shape[2] - - # Some generation methods already pass only the last input ID - if input_ids.shape[1] > past_length: - remove_prefix_length = past_length + def prepare_inputs_for_generation( + self, + input_ids, + past_key_values=None, + attention_mask=None, + head_mask=None, + decoder_head_mask=None, + decoder_attention_mask=None, + cross_attn_head_mask=None, + use_cache=None, + encoder_outputs=None, + token_idx=None, + **kwargs, + ): + # cut decoder_input_ids if past_key_values is used + if past_key_values is not None: + if token_idx is not None: + idx = token_idx + kwargs.get("inputs_embeds_offset", 0) - 1 + input_ids = torch.index_select(input_ids, 1, idx) else: - # Default to old behavior: keep only final ID - remove_prefix_length = input_ids.shape[1] - 1 - - input_ids = input_ids[:, remove_prefix_length:] - - return { - "decoder_input_ids": input_ids, - "past_key_values": past_key_values, - "encoder_outputs": encoder_outputs, - "attention_mask": attention_mask, - "decoder_attention_mask": decoder_attention_mask, - "head_mask": head_mask, - "decoder_head_mask": decoder_head_mask, - "cross_attn_head_mask": cross_attn_head_mask, - "use_cache": use_cache, - "token_idx": token_idx, - } + past_length = past_key_values[0][0].shape[2] + + # Some generation methods already pass only the last input ID + if input_ids.shape[1] > past_length: + remove_prefix_length = past_length + else: + # Default to old behavior: keep only final ID + remove_prefix_length = input_ids.shape[1] - 1 + + input_ids = input_ids[:, remove_prefix_length:] + + return { + "decoder_input_ids": input_ids, + "past_key_values": past_key_values, + "encoder_outputs": encoder_outputs, + "attention_mask": attention_mask, + "decoder_attention_mask": decoder_attention_mask, + "head_mask": head_mask, + "decoder_head_mask": decoder_head_mask, + "cross_attn_head_mask": cross_attn_head_mask, + "use_cache": use_cache, + "token_idx": token_idx, + } + + def _reorder_cache(self, past_key_values, beam_idx): + # if decoder past is not included in output + # speedy decoding is disabled and no need to reorder + if past_key_values is None: + logger.warning("You might want to consider setting `use_cache=True` to speed up decoding") + return past_key_values + + reordered_decoder_past = () + for layer_past_states in past_key_values: + # get the correct batch idx from layer past batch dim + # batch dim of `past` is at 2nd position + reordered_layer_past_states = () + for layer_past_state in layer_past_states: + # need to set correct `past` for each of the four key / value states + reordered_layer_past_states = reordered_layer_past_states + ( + layer_past_state.index_select(0, beam_idx.to(layer_past_state.device)), + ) + + if reordered_layer_past_states[0].shape != layer_past_states[0].shape: + raise ValueError( + f"reordered_layer_past_states[0] shape {reordered_layer_past_states[0].shape} and layer_past_states[0] shape {layer_past_states[0].shape} mismatched" + ) + if len(reordered_layer_past_states) != len(layer_past_states): + raise ValueError( + f"length of reordered_layer_past_states {len(reordered_layer_past_states)} and length of layer_past_states {len(layer_past_states)} mismatched" + ) + + reordered_decoder_past = reordered_decoder_past + (reordered_layer_past_states,) + return reordered_decoder_past diff --git a/optimum/habana/transformers/models/whisper/__init__.py b/optimum/habana/transformers/models/whisper/__init__.py index 1530983341..e43cb709ad 100644 --- a/optimum/habana/transformers/models/whisper/__init__.py +++ b/optimum/habana/transformers/models/whisper/__init__.py @@ -1,5 +1,4 @@ from .modeling_whisper import ( - GAUDI_WHISPER_ATTENTION_CLASSES, GaudiWhisperDecoder, GaudiWhisperDecoderLayer, GaudiWhisperForConditionalGeneration, diff --git a/optimum/habana/transformers/models/whisper/modeling_whisper.py b/optimum/habana/transformers/models/whisper/modeling_whisper.py index fc86606eb6..0b586aa78e 100644 --- a/optimum/habana/transformers/models/whisper/modeling_whisper.py +++ b/optimum/habana/transformers/models/whisper/modeling_whisper.py @@ -15,10 +15,8 @@ WhisperAttention, WhisperDecoder, WhisperDecoderLayer, - WhisperFlashAttention2, WhisperForConditionalGeneration, WhisperModel, - WhisperSdpaAttention, shift_tokens_right, ) from transformers.utils import logging @@ -27,7 +25,7 @@ logger = logging.get_logger(__name__) -class GaudiWhisperSdpaAttention(WhisperSdpaAttention): +class GaudiWhisperSdpaAttention(WhisperAttention): def forward( self, hidden_states: torch.Tensor, @@ -143,14 +141,6 @@ def forward( return attn_output, None, past_key_value -# Now only support the default GaudiWhisperSdpaAttention -GAUDI_WHISPER_ATTENTION_CLASSES = { - "eager": WhisperAttention, - "flash_attention_2": WhisperFlashAttention2, - "sdpa": GaudiWhisperSdpaAttention, -} - - class GaudiWhisperDecoderLayer(WhisperDecoderLayer): def forward( self, diff --git a/optimum/habana/transformers/trainer.py b/optimum/habana/transformers/trainer.py index 35ac284b7f..768b3e45ce 100644 --- a/optimum/habana/transformers/trainer.py +++ b/optimum/habana/transformers/trainer.py @@ -27,9 +27,9 @@ import shutil import time import warnings -from collections.abc import Mapping +from collections.abc import Iterator, Mapping from pathlib import Path -from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Tuple, Union +from typing import TYPE_CHECKING, Any, Callable, Optional, Union import huggingface_hub.utils as hf_hub_utils import numpy as np @@ -99,6 +99,7 @@ WEIGHTS_INDEX_NAME, WEIGHTS_NAME, PushInProgress, + check_torch_load_is_safe, is_accelerate_available, is_datasets_available, is_peft_available, @@ -141,7 +142,7 @@ from accelerate.utils import DataLoaderConfiguration, is_torch_version -def _get_input_update_settings(model, lazy_mode: Optional[bool] = None) -> Tuple[bool, Dict]: +def _get_input_update_settings(model, lazy_mode: Optional[bool] = None) -> tuple[bool, dict]: """ Determines whether the input settings need to be updated. @@ -156,10 +157,10 @@ def _get_input_update_settings(model, lazy_mode: Optional[bool] = None) -> Tuple lazy_mode[Optional[bool]]: Whether to use lazy mode for the model (defaults to `None`) Returns: - Tuple[bool, Dict]: A flag indicating whether the input settings should be updated. + tuple[bool, dict]: A flag indicating whether the input settings should be updated. A dictionary containing the specific input settings that need to be updated, if any """ - inputs_update: Dict = {} + inputs_update: dict = {} should_update_inputs = (getattr(model, "generation_config", None) is not None) and ( model.config.model_type in ("llama", "qwen2", "starcoder2", "gemma", "baichuan", "chatglm", "deepseek_v2") @@ -338,22 +339,34 @@ def __init__( "ignore", message="User provided device_type of 'cuda', but CUDA is not available. Disabling" ) + @property + def tokenizer(self) -> Optional[PreTrainedTokenizerBase]: + # Removed warning about usage of Trainer.tokenizer + return self.processing_class + + @tokenizer.setter + def tokenizer(self, processing_class) -> None: + # Removed warning about usage of Trainer.tokenizer + self.processing_class = processing_class + def _move_model_to_device(self, model, device): model = model.to(device) # Moving a model to HPU disconnects the tied weights, so we have to retie them. if self.args.use_habana and hasattr(model, "tie_weights"): model.tie_weights() - def _get_train_sampler(self) -> Optional[torch.utils.data.Sampler]: - if self.train_dataset is None or not has_length(self.train_dataset): + def _get_train_sampler(self, train_dataset: Optional[Dataset] = None) -> Optional[torch.utils.data.Sampler]: + if train_dataset is None: + train_dataset = self.train_dataset + if train_dataset is None or not has_length(train_dataset): return None # Build the sampler. if self.args.group_by_length: - if is_datasets_available() and isinstance(self.train_dataset, datasets.Dataset): + if is_datasets_available() and isinstance(train_dataset, datasets.Dataset): lengths = ( - self.train_dataset[self.args.length_column_name] - if self.args.length_column_name in self.train_dataset.column_names + train_dataset[self.args.length_column_name] + if self.args.length_column_name in train_dataset.column_names else None ) else: @@ -363,13 +376,13 @@ def _get_train_sampler(self) -> Optional[torch.utils.data.Sampler]: ) return LengthGroupedSampler( self.args.train_batch_size * self.args.gradient_accumulation_steps, - dataset=self.train_dataset, + dataset=train_dataset, lengths=lengths, model_input_name=model_input_name, ) else: - num_samples = len(self.train_dataset) + num_samples = len(train_dataset) if ( not self.args.dataloader_drop_last and num_samples % self.args.per_device_train_batch_size != 0 @@ -379,7 +392,7 @@ def _get_train_sampler(self) -> Optional[torch.utils.data.Sampler]: num_samples += ( self.args.per_device_train_batch_size - num_samples % self.args.per_device_train_batch_size ) - return RandomSampler(self.train_dataset, num_samples=num_samples) + return RandomSampler(train_dataset, num_samples=num_samples) def create_optimizer(self): """ @@ -517,12 +530,12 @@ def train( If a `str`, local path to a saved checkpoint as saved by a previous instance of [`Trainer`]. If a `bool` and equals `True`, load the last checkpoint in *args.output_dir* as saved by a previous instance of [`Trainer`]. If present, training will resume from the model/optimizer/scheduler states loaded here. - trial (`optuna.Trial` or `Dict[str, Any]`, *optional*): + trial (`optuna.Trial` or `dict[str, Any]`, *optional*): The trial run or the hyperparameter dictionary for hyperparameter search. ignore_keys_for_eval (`List[str]`, *optional*) A list of keys in the output of your model (if it is a dictionary) that should be ignored when gathering predictions for evaluation during the training. - kwargs (`Dict[str, Any]`, *optional*): + kwargs (`dict[str, Any]`, *optional*): Additional keyword arguments used to hide deprecated arguments """ if resume_from_checkpoint is False: @@ -650,7 +663,7 @@ def _inner_training_loop( # number of training epochs: num_train_epochs # number of training steps per epoch: num_update_steps_per_epoch # total number of training steps to execute: max_steps - total_train_batch_size = self._train_batch_size * args.gradient_accumulation_steps * args.world_size + total_train_batch_size = self.get_total_train_batch_size(args) ( num_train_epochs, num_update_steps_per_epoch, @@ -758,7 +771,7 @@ def hpu_deepspeed_checkpointing(function, *checkpoint_args, use_reentrant: Optio # as the model is wrapped, don't use `accelerator.prepare` # this is for unhandled cases such as # FSDP-XLA, SageMaker MP/DP, DataParallel, IPEX - use_accelerator_prepare = True if model is self.model else False + use_accelerator_prepare = model is self.model if use_accelerator_prepare and self.is_fsdp_enabled: # In case of auto_find_batch_size=True @@ -777,7 +790,11 @@ def hpu_deepspeed_checkpointing(function, *checkpoint_args, use_reentrant: Optio if use_accelerator_prepare: self.model.train() if hasattr(self.lr_scheduler, "step"): - model, self.optimizer = self.accelerator.prepare(self.model, self.optimizer) + # We should avoid accelerate preparing the model in TP case since we dont need it as it is handled by transformers from_pretrained and also it goes into DDP based preparation. + if self.is_tp_enabled: + self.optimizer = self.accelerator.prepare(self.optimizer) + else: + model, self.optimizer = self.accelerator.prepare(self.model, self.optimizer) else: # to handle cases wherein we pass "DummyScheduler" such as when it is specified in DeepSpeed config. model, self.optimizer, self.lr_scheduler = self.accelerator.prepare( @@ -959,19 +976,22 @@ def hpu_deepspeed_checkpointing(function, *checkpoint_args, use_reentrant: Optio step = -1 epoch_iterator = iter(epoch_dataloader) # We chunkify the epoch iterator into gradient accumulation steps `n` batches - remainder = num_examples % args.gradient_accumulation_steps + remainder = steps_in_epoch % args.gradient_accumulation_steps if remainder == 0: remainder = args.gradient_accumulation_steps update_step = -1 - total_updates = steps_in_epoch // args.gradient_accumulation_steps + 1 - if args.gradient_accumulation_steps == 1: - total_updates -= 1 + total_updates = steps_in_epoch // args.gradient_accumulation_steps + int( + remainder < args.gradient_accumulation_steps + ) for _ in range(total_updates): update_step += 1 num_batches = args.gradient_accumulation_steps if update_step != (total_updates - 1) else remainder batch_samples, num_items_in_batch = self.get_batch_samples_transformers( epoch_iterator, num_batches, args.device ) + # Store the number of batches for current gradient accumulation + # This is used to correctly scale the loss when the last accumulation step has fewer batches + self.current_gradient_accumulation_steps = len(batch_samples) for i, inputs in enumerate(batch_samples): step += 1 if self.args.compile_from_sec_iteration and is_torch_version(">=", "2.6.0"): @@ -1070,15 +1090,27 @@ def hpu_deepspeed_checkpointing(function, *checkpoint_args, use_reentrant: Optio # TODO: to merge self.accelerator.clip_grad_norm_ when HMP is removed grad_norm = self.FusedNorm.clip_norm(model.parameters()) else: - # Revert to normal clipping otherwise - grad_norm = self.accelerator.clip_grad_norm_( - model.parameters(), - args.max_grad_norm, - ) + grad_norm_context = contextlib.nullcontext + if self.is_tp_enabled: + from torch.distributed._tensor.experimental import implicit_replication + + grad_norm_context = implicit_replication + with grad_norm_context(): + grad_norm = self.accelerator.clip_grad_norm_( + model.parameters(), + args.max_grad_norm, + ) self.control = self.callback_handler.on_pre_optimizer_step(args, self.state, self.control) - self.optimizer.step() + context = contextlib.nullcontext + if self.is_tp_enabled: + from torch.distributed._tensor.experimental import implicit_replication + + context = implicit_replication + + with context(): + self.optimizer.step() self.control = self.callback_handler.on_optimizer_step(args, self.state, self.control) @@ -1273,6 +1305,7 @@ def _load_best_model(self): if self.args.save_safetensors and os.path.isfile(best_safe_model_path): state_dict = safetensors.torch.load_file(best_safe_model_path, device="cpu") else: + check_torch_load_is_safe() state_dict = torch.load(best_model_path, map_location="cpu", weights_only=True) # If the model is on the GPU, it still works! @@ -1478,6 +1511,7 @@ def _load_optimizer_and_scheduler(self, checkpoint): # deepspeed loads optimizer/lr_scheduler together with the model in deepspeed_init if not isinstance(self.lr_scheduler, DeepSpeedSchedulerWrapper): with warnings.catch_warnings(record=True) as caught_warnings: + check_torch_load_is_safe() self.lr_scheduler.load_state_dict( torch.load(os.path.join(checkpoint, SCHEDULER_NAME), weights_only=True) ) @@ -1512,11 +1546,13 @@ def _load_optimizer_and_scheduler(self, checkpoint): **_get_fsdp_ckpt_kwargs(), ) else: + check_torch_load_is_safe() self.optimizer.load_state_dict( torch.load(os.path.join(checkpoint, OPTIMIZER_NAME), map_location=map_location, weights_only=True) ) with warnings.catch_warnings(record=True) as caught_warnings: + check_torch_load_is_safe() self.lr_scheduler.load_state_dict( torch.load(os.path.join(checkpoint, SCHEDULER_NAME), map_location=map_location, weights_only=True) ) @@ -1533,7 +1569,7 @@ def log(self, logs: dict[str, float], start_time: Optional[float] = None) -> Non Subclass and override this method to inject custom behavior. Args: - logs (`Dict[str, float]`): + logs (`dict[str, float]`): The values to log. start_time (`Optional[float]`): The start of training. @@ -1543,7 +1579,7 @@ def log(self, logs: dict[str, float], start_time: Optional[float] = None) -> Non if self.args.include_num_input_tokens_seen: logs["num_input_tokens_seen"] = self.state.num_input_tokens_seen if start_time is not None: - speed_metrics("train", start_time, num_tokens=self.state.num_input_tokens_seen) + logs.update(speed_metrics("train", start_time, num_tokens=self.state.num_input_tokens_seen)) mem_stats = get_hpu_memory_stats(self.args.device) logs.update(mem_stats) @@ -1605,7 +1641,10 @@ def autocast_smart_context_manager(self, cache_enabled: Optional[bool] = True): return ctx_manager def training_step( - self, model: torch.nn.Module, inputs: dict[str, Union[torch.Tensor, Any]], num_items_in_batch=None + self, + model: torch.nn.Module, + inputs: dict[str, Union[torch.Tensor, Any]], + num_items_in_batch: Optional[torch.Tensor] = None, ) -> torch.Tensor: """ Perform a training step on a batch of inputs. @@ -1615,7 +1654,7 @@ def training_step( Args: model (`torch.nn.Module`): The model to train. - inputs (`Dict[str, Union[torch.Tensor, Any]]`): + inputs (`dict[str, Union[torch.Tensor, Any]]`): The inputs and targets of the model. The dictionary will be unpacked before being fed to the model. Most models expect the targets under the @@ -1647,10 +1686,10 @@ def training_step( if self.args.use_lazy_mode and self.args.pipelining_fwd_bwd: self.htcore.mark_step() - # Finally we need to normalize the loss for reporting - if not self.model_accepts_loss_kwargs and self.compute_loss_func is None: - # temporary fix to calculate loss correctly - loss = loss / self.args.gradient_accumulation_steps + # Finally we need to normalize the loss for reporting if GA loss bug is not fixed during compute loss + if (not self.model_accepts_loss_kwargs or num_items_in_batch is None) and self.compute_loss_func is None: + # If the model does not accept loss kwargs, we need to normalize the loss by the number of gradient accumulation steps + loss = loss / self.current_gradient_accumulation_steps # Turning off loss scaling w.r.t. gradient accumulation when DeepSpeed is enabled # https://github.com/huggingface/transformers/pull/35808 @@ -1688,7 +1727,14 @@ def save_model(self, output_dir: Optional[str] = None, _internal_call: bool = Fa if output_dir is None: output_dir = self.args.output_dir - if self.is_fsdp_enabled: + # We are in N-D parallelism if we have parallelism_config set, so we check accelerate if we're on a to_save rank + if getattr(self.accelerator, "parallelism_config", None) is not None: + if self.accelerator.should_save_model: + self._save(output_dir) + # If we drop to here, we're in 1D parallelism, so all ranks need to go to `save_pretrained` + elif (tp_size := getattr(self.model, "_tp_size", 0)) is not None and tp_size > 1: + self._save(output_dir) + elif self.is_fsdp_enabled: if "FULL_STATE_DICT" in str(self.accelerator.state.fsdp_plugin.state_dict_type): state_dict = self.accelerator.get_state_dict(self.model) if self.args.should_save: @@ -1719,7 +1765,7 @@ def save_model(self, output_dir: Optional[str] = None, _internal_call: bool = Fa # Push to the Hub when `save_model` is called by the user. if self.args.push_to_hub and not _internal_call: - self.push_to_hub(commit_message="Model save") + self.push_to_hub(commit_message="Model save", revision=self.args.hub_revision) def _save(self, output_dir: Optional[str] = None, state_dict=None): # If we are executing this function, we are the process zero, so we don't check for that. @@ -2171,7 +2217,7 @@ def prediction_step( Args: model (`torch.nn.Module`): The model to evaluate. - inputs (`Dict[str, Union[torch.Tensor, Any]]`): + inputs (`dict[str, Union[torch.Tensor, Any]]`): The inputs and targets of the model. The dictionary will be unpacked before being fed to the model. Most models expect the targets under the argument `labels`. Check your model's documentation for all accepted arguments. @@ -2181,17 +2227,17 @@ def prediction_step( A list of keys in the output of your model (if it is a dictionary) that should be ignored when gathering predictions. Return: - Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]: A tuple with the loss, + tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]: A tuple with the loss, logits and labels (each being optional). """ has_labels = False if len(self.label_names) == 0 else all(inputs.get(k) is not None for k in self.label_names) # For CLIP-like models capable of returning loss values. # If `return_loss` is not specified or being `None` in `inputs`, we check if the default value of `return_loss` # is `True` in `model.forward`. - return_loss = inputs.get("return_loss", None) + return_loss = inputs.get("return_loss") if return_loss is None: return_loss = self.can_return_loss - loss_without_labels = True if len(self.label_names) == 0 and return_loss else False + loss_without_labels = len(self.label_names) == 0 and return_loss inputs = self._prepare_inputs(inputs) if ignore_keys is None: @@ -2292,6 +2338,7 @@ def _push_from_checkpoint(self, checkpoint_folder): token=self.args.hub_token, run_as_future=True, ignore_patterns=["_*", f"{PREFIX_CHECKPOINT_DIR}-*"], + revision=self.args.hub_revision, ) push_jobs = [model_push_job] @@ -2307,6 +2354,7 @@ def _push_from_checkpoint(self, checkpoint_folder): commit_message=commit_message + ", checkpoint", token=self.args.hub_token, run_as_future=True, + revision=self.args.hub_revision, ) push_jobs.append(checkpoint_push) @@ -2567,16 +2615,16 @@ def create_accelerator_and_postprocess(self): } # tp is initialized at Accelerator init phase so # args should be prepared here - if self.args.tp_size > 1: + if hasattr(self.model, "tp_size") and self.model.tp_size is not None and self.model.tp_size > 1: self.is_tp_enabled = True - args["torch_tp_plugin"] = TorchTensorParallelPlugin(tp_size=self.args.tp_size) + args["torch_tp_plugin"] = TorchTensorParallelPlugin(tp_size=self.model.tp_size) # create accelerator object self.accelerator = GaudiAccelerator(**args) # some Trainer classes need to use `gather` instead of `gather_for_metrics`, thus we store a flag self.gather_function = self.accelerator.gather_for_metrics - if "use_gather_object" in inspect.signature(self.gather_function).parameters.keys(): + if "use_gather_object" in inspect.signature(self.gather_function).parameters: self.gather_function = functools.partial( self.gather_function, use_gather_object=self.args.eval_use_gather_object ) @@ -2654,7 +2702,9 @@ def _zero_model_grad(self, model): model.zero_grad() model._zero_grad_kwargs = {} - def get_batch_samples_transformers(self, epoch_iterator, num_batches, device): + def get_batch_samples_transformers( + self, epoch_iterator: Iterator, num_batches: int, device: torch.device + ) -> tuple[list, Optional[torch.Tensor]]: batch_samples = [] num_items_in_batch = None @@ -2692,13 +2742,3 @@ def get_batch_samples_transformers(self, epoch_iterator, num_batches, device): num_items_in_batch = num_items_in_batch.to(device) return batch_samples, num_items_in_batch - - @property - def tokenizer(self) -> Optional[PreTrainedTokenizerBase]: - # Removed warning about usage of Trainer.tokenizer - return self.processing_class - - @tokenizer.setter - def tokenizer(self, processing_class) -> None: - # Removed warning about usage of Trainer.tokenizer - self.processing_class = processing_class diff --git a/optimum/habana/transformers/trainer_seq2seq.py b/optimum/habana/transformers/trainer_seq2seq.py index 91e677b5d3..da7315181f 100644 --- a/optimum/habana/transformers/trainer_seq2seq.py +++ b/optimum/habana/transformers/trainer_seq2seq.py @@ -14,7 +14,6 @@ # limitations under the License. import contextlib -import warnings from copy import deepcopy from pathlib import Path from typing import TYPE_CHECKING, Any, Callable, Optional, Union @@ -135,15 +134,10 @@ def load_generation_config(gen_config_arg: Union[str, GaudiGenerationConfig]) -> # Strict validation to fail early. `GenerationConfig.save_pretrained()`, run at the end of training, throws # an exception if there are warnings at validation time. try: - with warnings.catch_warnings(record=True) as caught_warnings: - gen_config.validate() - if len(caught_warnings) > 0: - raise ValueError(str([w.message for w in caught_warnings])) + gen_config.validate(strict=True) except ValueError as exc: - raise ValueError( - "The loaded generation config instance is invalid -- `GenerationConfig.validate()` throws warnings " - "and/or exceptions. Fix these issues to train your model.\n\nThrown during validation:\n" + str(exc) - ) + raise ValueError(str(exc) + "\n\nFix these issues to train your model.") + return gen_config def evaluate( diff --git a/optimum/habana/transformers/training_args.py b/optimum/habana/transformers/training_args.py index b27b20dce6..49c57d5b21 100644 --- a/optimum/habana/transformers/training_args.py +++ b/optimum/habana/transformers/training_args.py @@ -18,12 +18,11 @@ from dataclasses import asdict, dataclass, field from datetime import timedelta from pathlib import Path -from typing import Optional, Union +from typing import Any, Optional, Union import torch.distributed as dist from accelerate import DistributedType, PartialState from accelerate.state import AcceleratorState -from packaging import version from transformers.debug_utils import DebugOption from transformers.file_utils import cached_property, is_torch_available, requires_backends from transformers.trainer_pt_utils import AcceleratorConfig @@ -544,10 +543,11 @@ def __post_init__(self): "--load_best_model_at_end requires the saving steps to be a multiple of the evaluation " f"steps, but found {self.save_steps}, which is not a multiple of {self.eval_steps}." ) - raise ValueError( - "--load_best_model_at_end requires the saving steps to be a round multiple of the evaluation " - f"steps, but found {self.save_steps}, which is not a round multiple of {self.eval_steps}." - ) + else: + raise ValueError( + "--load_best_model_at_end requires the saving steps to be a round multiple of the evaluation " + f"steps, but found {self.save_steps}, which is not a round multiple of {self.eval_steps}." + ) safetensors_available = is_safetensors_available() if self.save_safetensors and not safetensors_available: @@ -565,9 +565,7 @@ def __post_init__(self): ) and self.metric_for_best_model is None: self.metric_for_best_model = "loss" if self.greater_is_better is None and self.metric_for_best_model is not None: - self.greater_is_better = not (self.metric_for_best_model.endswith("loss")) - if self.run_name is None: - self.run_name = self.output_dir + self.greater_is_better = not self.metric_for_best_model.endswith("loss") if self.lr_scheduler_type == SchedulerType.REDUCE_ON_PLATEAU: if self.eval_strategy == IntervalStrategy.NO: @@ -585,9 +583,6 @@ def __post_init__(self): FutureWarning, ) self.optim = OptimizerNames.ADAFACTOR - if self.optim == OptimizerNames.ADAMW_TORCH_FUSED and is_torch_available(): - if version.parse(version.parse(torch.__version__).base_version) < version.parse("2.0.0"): - raise ValueError("--optim adamw_torch_fused requires PyTorch 2.0 or higher") # We need to setup the accelerator config here *before* the first call to `self.device` if is_accelerate_available(): @@ -605,17 +600,24 @@ def __post_init__(self): ) else: self.accelerator_config = AcceleratorConfig.from_json_file(self.accelerator_config) - + if self.accelerator_config.split_batches: + logger.info( + "Using `split_batches=True` in `accelerator_config` will override the `per_device_train_batch_size` " + "Batches will be split across all processes equally when using `split_batches=True`." + ) if self.dataloader_drop_last: self.accelerator_config.even_batches = False + # This call to self.device is necessary to call _setup_devices so that + # torch.distributed is initialized + device_is_hpu = self.device.type == "hpu" + # Disable average tokens when using single device if self.average_tokens_across_devices: try: if self.world_size == 1: - logger.warning( - "average_tokens_across_devices is set to True but it is invalid when world size is" - "1. Turn it to False automatically." + logger.info( + "average_tokens_across_devices is True but world size is 1. Setting it to False automatically." ) self.average_tokens_across_devices = False except ImportError as e: @@ -715,10 +717,12 @@ def __post_init__(self): warn0("`--fsdp_config` is useful only when `--fsdp` is specified.") with open(self.fsdp_config, encoding="utf-8") as f: self.fsdp_config = json.load(f) - for k in list(self.fsdp_config.keys()): - if k.startswith("fsdp_"): - v = self.fsdp_config.pop(k) - self.fsdp_config[k[5:]] = v + + if self.fsdp_config is not None and isinstance(self.fsdp_config, dict): + for k in list(self.fsdp_config.keys()): + if k.startswith("fsdp_"): + v = self.fsdp_config.pop(k) + self.fsdp_config[k[5:]] = v if self.fsdp_min_num_params > 0: warn0("using `--fsdp_min_num_params` is deprecated. Use fsdp_config instead ", FutureWarning) @@ -754,9 +758,6 @@ def __post_init__(self): self.fsdp_config["xla_fsdp_v2"] = self.fsdp_config.get("xla_fsdp_v2", False) self.fsdp_config["xla_fsdp_grad_ckpt"] = self.fsdp_config.get("xla_fsdp_grad_ckpt", False) - if self.tp_size > 1: - os.environ["ACCELERATE_USE_TP"] = "true" - os.environ["TP_SIZE"] = str(self.tp_size) # accelerate integration for FSDP if len(self.fsdp) > 0 and not self.fsdp_config["xla"]: os.environ["ACCELERATE_USE_FSDP"] = "true" @@ -809,9 +810,6 @@ def __post_init__(self): elif self.debug is None: self.debug = [] - # This call to self.device is necessary to call _setup_devices so that - # torch.distributed is initialized - device_is_hpu = self.device.type == "hpu" self.deepspeed_plugin = None if self.deepspeed: if not device_is_hpu: @@ -964,7 +962,7 @@ def _setup_devices(self) -> "torch.device": f"Please run `pip install transformers[torch]` or `pip install accelerate -U`" ) # We delay the init of `PartialState` to the end for clarity - accelerator_state_kwargs = {"enabled": True, "use_configured_state": False} + accelerator_state_kwargs: dict[str, Any] = {"enabled": True, "use_configured_state": False} if isinstance(self.accelerator_config, AcceleratorConfig): accelerator_state_kwargs["use_configured_state"] = self.accelerator_config.pop( "use_configured_state", False diff --git a/setup.py b/setup.py index 9c414eff33..cd897feee0 100644 --- a/setup.py +++ b/setup.py @@ -29,7 +29,7 @@ INSTALL_REQUIRES = [ - "transformers >= 4.51.0, < 4.52.0", + "transformers >= 4.55.0, < 4.56.0", "optimum ~= 1.25", "torch", "accelerate >= 1.7.0", diff --git a/tests/test_trainer.py b/tests/test_trainer.py index 4ad2d23bf5..f88e133fcf 100644 --- a/tests/test_trainer.py +++ b/tests/test_trainer.py @@ -1,4 +1,3 @@ -# coding=utf-8 # Copyright 2022 the HuggingFace Inc. team. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -80,6 +79,7 @@ SAFE_WEIGHTS_NAME, WEIGHTS_INDEX_NAME, WEIGHTS_NAME, + check_torch_load_is_safe, is_accelerate_available, is_safetensors_available, ) @@ -463,7 +463,7 @@ def get_regression_trainer( output_dir=None, **kwargs, ): - label_names = kwargs.get("label_names", None) + label_names = kwargs.get("label_names") gradient_checkpointing = kwargs.get("gradient_checkpointing", False) train_dataset = RegressionDataset(length=train_len, label_names=label_names) eval_dataset = RegressionDataset(length=eval_len, label_names=label_names) @@ -569,6 +569,7 @@ def check_best_model_has_been_loaded( else: best_model = RegressionModel() if not safe_weights: + check_torch_load_is_safe() state_dict = torch.load(os.path.join(checkpoint, WEIGHTS_NAME), weights_only=True) else: state_dict = safetensors.torch.load_file(os.path.join(checkpoint, SAFE_WEIGHTS_NAME)) @@ -580,6 +581,11 @@ def check_best_model_has_been_loaded( metrics = trainer.evaluate() self.assertEqual(metrics[metric], best_value) + def remove_nan_logs(self, log): + for key in list(log.keys()): + if log[key] != log[key]: # Check if the value is NaN + del log[key] + def check_trainer_state_are_the_same(self, trainer_state, trainer_state1): # We'll pop things so operate on copies. state = trainer_state.copy() @@ -593,6 +599,10 @@ def check_trainer_state_are_the_same(self, trainer_state, trainer_state1): for key in skip_log_keys: _ = log.pop(key, None) _ = log1.pop(key, None) + + self.remove_nan_logs(log) + self.remove_nan_logs(log1) + self.assertEqual(log, log1) def convert_to_sharded_checkpoint(self, folder, save_safe=True, load_safe=True): @@ -601,6 +611,7 @@ def convert_to_sharded_checkpoint(self, folder, save_safe=True, load_safe=True): loader = safetensors.torch.load_file weights_file = os.path.join(folder, SAFE_WEIGHTS_NAME) else: + check_torch_load_is_safe() loader = torch.load weights_file = os.path.join(folder, WEIGHTS_NAME) @@ -1088,6 +1099,37 @@ def test_cosine_with_min_lr_scheduler(self): trainer.lr_scheduler.step() self.assertEqual(trainer.lr_scheduler.get_last_lr()[0], 1e-5) + def test_cosine_with_min_lr_schedule_with_warmup_lr_rate(self): + train_dataset = RegressionDataset() + model = RegressionModel() + num_steps, num_warmup_steps = 10, 2 + extra_kwargs = {"min_lr": 1e-5} # Non-default arguments + with tempfile.TemporaryDirectory() as tmpdir: + args = GaudiTrainingArguments( + tmpdir, + lr_scheduler_type="cosine_warmup_with_min_lr", + lr_scheduler_kwargs=extra_kwargs, + learning_rate=0.2, + warmup_steps=num_warmup_steps, + report_to="none", + use_habana=True, + use_lazy_mode=True, + ) + trainer = GaudiTrainer(model, gaudi_config=get_gaudi_config(), args=args, train_dataset=train_dataset) + trainer.create_optimizer_and_scheduler(num_training_steps=num_steps) + + # Checking that the scheduler was created + self.assertIsNotNone(trainer.lr_scheduler) + + # Check the last learning rate + step_lrs = [] + for _ in range(num_steps): + step_lrs.append(trainer.optimizer.param_groups[0]["lr"]) + trainer.lr_scheduler.step() + self.assertEqual(step_lrs[0], 0.1) + self.assertEqual(step_lrs[1], 0.2) + self.assertEqual(step_lrs[-1], 1e-05) + def test_reduce_lr_on_plateau_args(self): # test passed arguments for a custom ReduceLROnPlateau scheduler train_dataset = RegressionDataset(length=64) @@ -1289,7 +1331,7 @@ def test_training_arguments_are_left_untouched(self): trainer.train() args = GaudiTrainingArguments(tmp_dir, use_habana=True, use_lazy_mode=True, report_to=[]) dict1, dict2 = args.to_dict(), trainer.args.to_dict() - for key in dict1.keys(): + for key in dict1: # Logging dir can be slightly different as they default to something with the time. if key != "logging_dir": self.assertEqual(dict1[key], dict2[key]) @@ -1324,6 +1366,7 @@ def test_number_of_steps_in_training(self): # per_device_train_batch_size=2, # torch_compile=True, # max_steps=1, # compile happens on the first step + # report_to="none", # use_habana=True, # use_lazy_mode=True, # ) @@ -2235,7 +2278,7 @@ def test_resume_training_with_randomness(self): # gaudi_config = get_gaudi_config() # trainer = GaudiTrainer(model, gaudi_config, args, train_dataset=train_dataset, callbacks=[MockOOMCallback()]) # trainer.train() - # self.assertEqual(trainer._train_batch_size, 8) + # self.assertEqual(trainer._train_batch_size, 14) # def test_auto_batch_size_with_resume_from_checkpoint(self): # train_dataset = RegressionDataset(length=128) @@ -2260,16 +2303,16 @@ def test_resume_training_with_randomness(self): # model, gaudi_config, args, train_dataset=train_dataset, callbacks=[MockOOMCallback()] # ) # trainer.train() - # # After `auto_find_batch_size` is ran we should now be at 8 - # self.assertEqual(trainer._train_batch_size, 8) + # # After `auto_find_batch_size` is ran we should now be at 16*0.9=14 + # self.assertEqual(trainer._train_batch_size, 14) # # We can then make a new Trainer # trainer = GaudiTrainer(model, gaudi_config, args, train_dataset=train_dataset) # # Check we are at 16 to start # self.assertEqual(trainer._train_batch_size, 16 * max(trainer.args.n_gpu, 1)) # trainer.train(resume_from_checkpoint=True) - # # We should be back to 8 again, picking up based upon the last ran Trainer - # self.assertEqual(trainer._train_batch_size, 8) + # # We should be back to 14 again, picking up based upon the last ran Trainer + # self.assertEqual(trainer._train_batch_size, 14) # regression for this issue: https://github.com/huggingface/transformers/issues/12970 def test_training_with_resume_from_checkpoint_false(self): @@ -2609,6 +2652,36 @@ def test_num_train_epochs_in_training(self): train_output = trainer.train() self.assertEqual(train_output.global_step, int(self.n_epochs)) + def test_num_batches_in_training_with_gradient_accumulation(self): + with tempfile.TemporaryDirectory() as tmp_dir: + for num_train_epochs in [1, 2]: + for train_len in [123, 120]: + trainer = get_regression_trainer( + train_len=train_len, + per_device_train_batch_size=4, + gradient_accumulation_steps=5, + num_train_epochs=num_train_epochs, + output_dir=tmp_dir, + ) + + total_batch_samples = [] + + def wrap_get_batch_samples(fn): + def wrapped_fn(epoch_iterator, num_batches, device): + self.assertGreater(num_batches, 0) + batch_samples, num_items_in_batch = fn(epoch_iterator, num_batches, device) + self.assertEqual(len(batch_samples), num_batches) + total_batch_samples.append(num_batches) + return batch_samples, num_items_in_batch + + return wrapped_fn + + trainer.get_batch_samples = wrap_get_batch_samples(trainer.get_batch_samples) + + trainer.train() + + self.assertEqual(len(trainer.get_train_dataloader()) * num_train_epochs, sum(total_batch_samples)) + def test_early_stopping_callback(self): # early stopping stops training before num_training_epochs with tempfile.TemporaryDirectory() as tmp_dir: @@ -3056,7 +3129,7 @@ def test_trainer_saves_image_processor(self): gaudi_config = get_gaudi_config() trainer = GaudiTrainer( model=RegressionPreTrainedModel(config), - args=GaudiTrainingArguments(output_dir=tmp_dir, use_habana=True, use_lazy_mode=True), + args=GaudiTrainingArguments(output_dir=tmp_dir, report_to="none", use_habana=True, use_lazy_mode=True), gaudi_config=gaudi_config, processing_class=image_processor, ) @@ -3074,7 +3147,7 @@ def test_trainer_saves_feature_extractor(self): gaudi_config = get_gaudi_config() trainer = GaudiTrainer( model=RegressionPreTrainedModel(config), - args=GaudiTrainingArguments(output_dir=tmp_dir, use_habana=True, use_lazy_mode=True), + args=GaudiTrainingArguments(output_dir=tmp_dir, report_to="none", use_habana=True, use_lazy_mode=True), gaudi_config=gaudi_config, processing_class=feature_extractor, ) @@ -3096,7 +3169,7 @@ def test_trainer_saves_processor(self): gaudi_config = get_gaudi_config() trainer = GaudiTrainer( model=RegressionPreTrainedModel(config), - args=GaudiTrainingArguments(output_dir=tmp_dir, use_habana=True, use_lazy_mode=True), + args=GaudiTrainingArguments(output_dir=tmp_dir, report_to="none", use_habana=True, use_lazy_mode=True), gaudi_config=gaudi_config, processing_class=processor, ) @@ -3470,8 +3543,7 @@ def test_push_to_hub_in_organization(self): def get_commit_history(self, repo): commit_logs = subprocess.run( "git log".split(), - stderr=subprocess.PIPE, - stdout=subprocess.PIPE, + capture_output=True, check=True, encoding="utf-8", cwd=repo, diff --git a/tests/test_trainer_distributed.py b/tests/test_trainer_distributed.py index b16460e14e..2738a21752 100644 --- a/tests/test_trainer_distributed.py +++ b/tests/test_trainer_distributed.py @@ -14,7 +14,6 @@ # limitations under the License. from pathlib import Path -from typing import Dict from transformers import EvalPrediction, HfArgumentParser, is_torch_available from transformers.testing_utils import TestCasePlus @@ -135,7 +134,7 @@ def test_gaudi_trainer_distributed_hpu_graphs(self): for dataset_length in [101, 40, 7]: dataset = DummyDataset(dataset_length) - def compute_metrics(p: EvalPrediction) -> Dict: + def compute_metrics(p: EvalPrediction) -> dict: sequential = list(range(len(dataset))) success = p.predictions.tolist() == sequential and p.label_ids.tolist() == sequential if not success and training_args.local_rank == 0: diff --git a/tests/test_trainer_seq2seq.py b/tests/test_trainer_seq2seq.py index f5866043a8..e8c4203f0e 100644 --- a/tests/test_trainer_seq2seq.py +++ b/tests/test_trainer_seq2seq.py @@ -1,4 +1,3 @@ -# coding=utf-8 # Copyright 2022 the HuggingFace Inc. team. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -157,4 +156,4 @@ def test_bad_generation_config_fail_early(self): data_collator=data_collator, compute_metrics=lambda x: {"samples": x[0].shape[0]}, ) - self.assertIn("The loaded generation config instance is invalid", str(exc.exception)) + self.assertIn("Fix these issues to train your model", str(exc.exception)) From 5450d7ab50139c01864740a8ca9ff19bd3a5169c Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Mon, 18 Aug 2025 09:31:35 +0000 Subject: [PATCH 18/22] Fix tests + comment remote models (to fix later) --- optimum/habana/transformers/modeling_utils.py | 7 +------ optimum/habana/transformers/models/__init__.py | 2 +- .../transformers/models/decilm/modeling_decilm.py | 2 +- .../models/gpt_neox/modeling_gpt_neox.py | 1 - optimum/habana/transformers/models/t5/__init__.py | 2 +- tests/test_trainer.py | 12 ++++++------ 6 files changed, 10 insertions(+), 16 deletions(-) diff --git a/optimum/habana/transformers/modeling_utils.py b/optimum/habana/transformers/modeling_utils.py index 29c7864a3d..1a2aa61366 100644 --- a/optimum/habana/transformers/modeling_utils.py +++ b/optimum/habana/transformers/modeling_utils.py @@ -50,11 +50,6 @@ ChatGLMForConditionalGeneration, ChatGLMForSequenceClassification, ChatGLMTokenizer, - DeciLMConfig, - DeciLMForCausalLM, - DeepseekTokenizerFast, - DeepseekV2Config, - DeepseekV2ForCausalLM, DeepseekV3Config, DeepseekV3ForCausalLM, Gaudi2Idefics2ImageProcessor, @@ -190,6 +185,7 @@ GaudiStarcoder2DecoderLayer, GaudiStarcoder2ForCausalLM, GaudiStarcoder2Model, + GaudiT5ForConditionalGeneration, GaudiVideoLlavaForConditionalGeneration, GaudiVideoLlavaProcessor, GaudiVisionSdpaAttention, @@ -297,7 +293,6 @@ gaudi_t5_layernorm_forward, gaudi_T5Attention_forward, gaudi_T5Block_forward, - GaudiT5ForConditionalGeneration, gaudi_T5LayerSelfAttention_forward, gaudi_T5Stack_forward, gaudi_table_transformer_conv_encoder_forward, diff --git a/optimum/habana/transformers/models/__init__.py b/optimum/habana/transformers/models/__init__.py index d6c8d62f40..4d93c789c6 100644 --- a/optimum/habana/transformers/models/__init__.py +++ b/optimum/habana/transformers/models/__init__.py @@ -332,10 +332,10 @@ GaudiStarcoder2Model, ) from .t5 import ( + GaudiT5ForConditionalGeneration, gaudi_t5_layernorm_forward, gaudi_T5Attention_forward, gaudi_T5Block_forward, - GaudiT5ForConditionalGeneration, gaudi_T5LayerSelfAttention_forward, gaudi_T5Stack_forward, ) diff --git a/optimum/habana/transformers/models/decilm/modeling_decilm.py b/optimum/habana/transformers/models/decilm/modeling_decilm.py index 8a734404c3..fdab2dc475 100644 --- a/optimum/habana/transformers/models/decilm/modeling_decilm.py +++ b/optimum/habana/transformers/models/decilm/modeling_decilm.py @@ -25,7 +25,7 @@ apply_rotary_pos_emb, repeat_kv, ) -from transformers.utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging +from transformers.utils import logging from ...modeling_attn_mask_utils import _gaudi_prepare_4d_causal_attention_mask from ..llama.modeling_llama import GaudiLlamaRotaryEmbedding diff --git a/optimum/habana/transformers/models/gpt_neox/modeling_gpt_neox.py b/optimum/habana/transformers/models/gpt_neox/modeling_gpt_neox.py index c10875e0c1..ff92dd3d29 100644 --- a/optimum/habana/transformers/models/gpt_neox/modeling_gpt_neox.py +++ b/optimum/habana/transformers/models/gpt_neox/modeling_gpt_neox.py @@ -13,7 +13,6 @@ apply_rotary_pos_emb, logger, ) -from transformers.processing_utils import Unpack from ...modeling_attn_mask_utils import _gaudi_prepare_4d_causal_attention_mask from ...modeling_rope_utils import GaudiRotaryEmbedding diff --git a/optimum/habana/transformers/models/t5/__init__.py b/optimum/habana/transformers/models/t5/__init__.py index a6379c20fa..549d55352c 100644 --- a/optimum/habana/transformers/models/t5/__init__.py +++ b/optimum/habana/transformers/models/t5/__init__.py @@ -1,8 +1,8 @@ from .modeling_t5 import ( + GaudiT5ForConditionalGeneration, gaudi_t5_layernorm_forward, gaudi_T5Attention_forward, gaudi_T5Block_forward, - GaudiT5ForConditionalGeneration, gaudi_T5LayerSelfAttention_forward, gaudi_T5Stack_forward, ) diff --git a/tests/test_trainer.py b/tests/test_trainer.py index f88e133fcf..2c64d0ac31 100644 --- a/tests/test_trainer.py +++ b/tests/test_trainer.py @@ -783,7 +783,7 @@ def tokenize_function(examples): use_lazy_mode=True, **args_kwargs, ) - # train with base loss + # train with base loss (per_device_train_batch_size is equal to 8 by default) set_seed(42) model = AutoModelForCausalLM.from_pretrained(model_name) base_loss_callback = StoreLossCallback() @@ -799,6 +799,7 @@ def tokenize_function(examples): assert trainer.model_accepts_loss_kwargs trainer.train() + # train with gradient accumulation args = GaudiTrainingArguments( tmp_dir, **args_kwargs, @@ -808,7 +809,6 @@ def tokenize_function(examples): use_lazy_mode=True, ) - # train with gradient accumulation set_seed(42) model = AutoModelForCausalLM.from_pretrained(model_name) grad_accum_loss_callback = StoreLossCallback() @@ -848,15 +848,15 @@ def tokenize_function(examples): # all diff truth should be quite close self.assertLess(max(diff_truth), 0.01, f"Difference {max(diff_truth)} is not within 0.01") # max diff broken should be very off ("very off" is arbitrary, but as long as it's bigger than 0.1, it's fine) - # updated target value compared original implementation https://github.com/huggingface/transformers/blob/v4.49.0/tests/trainer/test_trainer.py#L888 - self.assertGreater(max(diff_broken), 1.0, f"Difference {max(diff_broken)} is not greater than 1.0") + self.assertGreater(max(diff_broken), 0.7, f"Difference {max(diff_broken)} is not greater than 0.7") loss_base = sum(base_loss_callback.losses) loss_broken = sum(broken_loss_callback.losses) # mean/sum loss should not vary too much. relative_diff = abs(loss_base - loss_broken) / max(loss_base, loss_broken) - self.assertLess(relative_diff, 0.2, f"Relative difference {relative_diff} is not within 0.2") + # updated target value compared to original implementation: https://github.com/huggingface/transformers/blob/v4.55.0/tests/trainer/test_trainer.py#L922 + self.assertLess(relative_diff, 0.6, f"Relative difference {relative_diff} is not within 0.6") def test_gradient_accumulation_loss_alignment_with_loss_func(self): set_seed(42) @@ -2676,7 +2676,7 @@ def wrapped_fn(epoch_iterator, num_batches, device): return wrapped_fn - trainer.get_batch_samples = wrap_get_batch_samples(trainer.get_batch_samples) + trainer.get_batch_samples_transformers = wrap_get_batch_samples(trainer.get_batch_samples_transformers) trainer.train() From f8aa93a209b0b6cc722791679cc35e5120351055 Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Mon, 18 Aug 2025 12:28:52 +0000 Subject: [PATCH 19/22] Fix decilm --- optimum/habana/transformers/modeling_utils.py | 6 +- .../models/decilm/configuration_decilm.py | 80 +------------------ .../models/decilm/modeling_decilm.py | 2 + 3 files changed, 9 insertions(+), 79 deletions(-) diff --git a/optimum/habana/transformers/modeling_utils.py b/optimum/habana/transformers/modeling_utils.py index 1a2aa61366..890014f60c 100644 --- a/optimum/habana/transformers/modeling_utils.py +++ b/optimum/habana/transformers/modeling_utils.py @@ -50,6 +50,8 @@ ChatGLMForConditionalGeneration, ChatGLMForSequenceClassification, ChatGLMTokenizer, + DeciLMConfig, + DeciLMForCausalLM, DeepseekV3Config, DeepseekV3ForCausalLM, Gaudi2Idefics2ImageProcessor, @@ -782,8 +784,8 @@ def adapt_transformers_to_gaudi(): transformers.models.mllama.modeling_mllama.MllamaVisionEncoderLayer = GaudiMllamaVisionEncoderLayer transformers.models.mllama.modeling_mllama.MllamaVisionSdpaAttention = GaudiMllamaVisionSdpaAttention - # transformers.AutoConfig.register("deci", DeciLMConfig) - # transformers.AutoModelForCausalLM.register(DeciLMConfig, DeciLMForCausalLM) + transformers.AutoConfig.register("deci", DeciLMConfig) + transformers.AutoModelForCausalLM.register(DeciLMConfig, DeciLMForCausalLM) # Optimization for deepseek on Gaudi # transformers.AutoConfig.register("deepseek_v2", DeepseekV2Config) diff --git a/optimum/habana/transformers/models/decilm/configuration_decilm.py b/optimum/habana/transformers/models/decilm/configuration_decilm.py index 71f6ee02cc..35bbdad6f8 100644 --- a/optimum/habana/transformers/models/decilm/configuration_decilm.py +++ b/optimum/habana/transformers/models/decilm/configuration_decilm.py @@ -3,10 +3,10 @@ https://huggingface.co/Deci/DeciLM-7B/blob/main/configuration_decilm.py """ -from transformers.configuration_utils import PretrainedConfig +from transformers.models.llama.configuration_llama import LlamaConfig -class DeciLMConfig(PretrainedConfig): +class DeciLMConfig(LlamaConfig): r""" Args: num_key_value_heads_per_layer (`List[int]`): @@ -15,84 +15,10 @@ class DeciLMConfig(PretrainedConfig): model_type = "deci" - keys_to_ignore_at_inference = ["past_key_values"] - # Default tensor parallel plan for base model `LlamaModel` - base_model_tp_plan = { - "layers.*.self_attn.q_proj": "colwise", - "layers.*.self_attn.k_proj": "colwise", - "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", - } - base_model_pp_plan = { - "embed_tokens": (["input_ids"], ["inputs_embeds"]), - "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), - "norm": (["hidden_states"], ["hidden_states"]), - } - def __init__( self, - vocab_size=32000, - hidden_size=4096, - intermediate_size=11008, - num_hidden_layers=32, - num_attention_heads=32, - num_key_value_heads=None, - hidden_act="silu", - max_position_embeddings=2048, - initializer_range=0.02, - rms_norm_eps=1e-6, - use_cache=True, - pad_token_id=None, - bos_token_id=1, - eos_token_id=2, - pretraining_tp=1, - tie_word_embeddings=False, - rope_theta=10000.0, - rope_scaling=None, - attention_bias=False, - attention_dropout=0.0, - mlp_bias=False, - head_dim=None, num_key_value_heads_per_layer: list = None, **kwargs, ): - self.vocab_size = vocab_size - self.max_position_embeddings = max_position_embeddings - self.hidden_size = hidden_size - self.intermediate_size = intermediate_size - self.num_hidden_layers = num_hidden_layers - self.num_attention_heads = num_attention_heads - - # for backward compatibility - if num_key_value_heads is None: - num_key_value_heads = num_attention_heads - - self.num_key_value_heads = num_key_value_heads - self.hidden_act = hidden_act - self.initializer_range = initializer_range - self.rms_norm_eps = rms_norm_eps - self.pretraining_tp = pretraining_tp - self.use_cache = use_cache - self.rope_theta = rope_theta - self.rope_scaling = rope_scaling - self.attention_bias = attention_bias - self.attention_dropout = attention_dropout - self.mlp_bias = mlp_bias - self.head_dim = head_dim if head_dim is not None else self.hidden_size // self.num_attention_heads self.num_key_value_heads_per_layer = num_key_value_heads_per_layer - # Validate the correctness of rotary position embeddings parameters - # BC: if there is a 'type' field, copy it it to 'rope_type'. - if self.rope_scaling is not None and "type" in self.rope_scaling: - self.rope_scaling["rope_type"] = self.rope_scaling["type"] - rope_config_validation(self) - - super().__init__( - pad_token_id=pad_token_id, - bos_token_id=bos_token_id, - eos_token_id=eos_token_id, - tie_word_embeddings=tie_word_embeddings, - **kwargs, - ) + super().__init__(**kwargs) diff --git a/optimum/habana/transformers/models/decilm/modeling_decilm.py b/optimum/habana/transformers/models/decilm/modeling_decilm.py index fdab2dc475..bb3bb8b00c 100644 --- a/optimum/habana/transformers/models/decilm/modeling_decilm.py +++ b/optimum/habana/transformers/models/decilm/modeling_decilm.py @@ -381,6 +381,8 @@ def forward( class DeciLMForCausalLM(LlamaForCausalLM, DeciLMPreTrainedModel): + config_class = DeciLMConfig + def __init__(self, config): DeciLMPreTrainedModel.__init__(self, config) self.model = DeciLMModel(config) From f1de53764854653e333a45f4779f19a32edfb658 Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Mon, 18 Aug 2025 13:17:12 +0000 Subject: [PATCH 20/22] Fix deepseek_v2 --- optimum/habana/transformers/modeling_utils.py | 7 +--- .../habana/transformers/models/__init__.py | 1 - .../models/deepseek_v2/__init__.py | 1 - .../deepseek_v2/configuration_deepseek_v2.py | 2 - .../deepseek_v2/tokenization_deepseek_v2.py | 38 ------------------- 5 files changed, 2 insertions(+), 47 deletions(-) delete mode 100644 optimum/habana/transformers/models/deepseek_v2/tokenization_deepseek_v2.py diff --git a/optimum/habana/transformers/modeling_utils.py b/optimum/habana/transformers/modeling_utils.py index 890014f60c..4a040c85ab 100644 --- a/optimum/habana/transformers/modeling_utils.py +++ b/optimum/habana/transformers/modeling_utils.py @@ -310,6 +310,7 @@ gaudi_xglm_model_forward, gaudi_XLMRoberta_Sdpa_SelfAttention_forward, ) +from .models.deepseek_v2.modeling_deepseek_v2 import DeepseekV2ForCausalLM as GaudiDeepseekV2ForCausalLM def adapt_transformers_to_gaudi(): @@ -788,11 +789,7 @@ def adapt_transformers_to_gaudi(): transformers.AutoModelForCausalLM.register(DeciLMConfig, DeciLMForCausalLM) # Optimization for deepseek on Gaudi - # transformers.AutoConfig.register("deepseek_v2", DeepseekV2Config) - # transformers.AutoModelForCausalLM.register(DeepseekV2Config, DeepseekV2ForCausalLM) - # transformers.AutoTokenizer.register(DeepseekV2Config, fast_tokenizer_class=DeepseekTokenizerFast) - # transformers.AutoConfig.register("deepseek_v3", DeepseekV3Config) - # transformers.AutoModelForCausalLM.register(DeepseekV3Config, DeepseekV3ForCausalLM) + transformers.models.deepseek_v2.modeling_deepseek_v2.DeepseekV2ForCausalLM = GaudiDeepseekV2ForCausalLM transformers.models.deepseek_v3.configuration_deepseek_v3.DeepseekV3Config = DeepseekV3Config transformers.models.deepseek_v3.modeling_deepseek_v3.DeepseekV3ForCausalLM = DeepseekV3ForCausalLM diff --git a/optimum/habana/transformers/models/__init__.py b/optimum/habana/transformers/models/__init__.py index 4d93c789c6..8efbc3bfa5 100644 --- a/optimum/habana/transformers/models/__init__.py +++ b/optimum/habana/transformers/models/__init__.py @@ -70,7 +70,6 @@ DeciLMForCausalLM, ) from .deepseek_v2 import ( - DeepseekTokenizerFast, DeepseekV2Config, DeepseekV2ForCausalLM, ) diff --git a/optimum/habana/transformers/models/deepseek_v2/__init__.py b/optimum/habana/transformers/models/deepseek_v2/__init__.py index 66d303b31f..8e5d7b89aa 100644 --- a/optimum/habana/transformers/models/deepseek_v2/__init__.py +++ b/optimum/habana/transformers/models/deepseek_v2/__init__.py @@ -1,3 +1,2 @@ from .configuration_deepseek_v2 import DeepseekV2Config from .modeling_deepseek_v2 import DeepseekV2ForCausalLM -from .tokenization_deepseek_v2 import DeepseekTokenizerFast diff --git a/optimum/habana/transformers/models/deepseek_v2/configuration_deepseek_v2.py b/optimum/habana/transformers/models/deepseek_v2/configuration_deepseek_v2.py index 6f9d03ddcd..b35a76963c 100644 --- a/optimum/habana/transformers/models/deepseek_v2/configuration_deepseek_v2.py +++ b/optimum/habana/transformers/models/deepseek_v2/configuration_deepseek_v2.py @@ -7,8 +7,6 @@ logger = logging.get_logger(__name__) -DEEPSEEK_PRETRAINED_CONFIG_ARCHIVE_MAP = {} - class DeepseekV2Config(PretrainedConfig): r""" diff --git a/optimum/habana/transformers/models/deepseek_v2/tokenization_deepseek_v2.py b/optimum/habana/transformers/models/deepseek_v2/tokenization_deepseek_v2.py deleted file mode 100644 index 12499a38b5..0000000000 --- a/optimum/habana/transformers/models/deepseek_v2/tokenization_deepseek_v2.py +++ /dev/null @@ -1,38 +0,0 @@ -""" -Copied from https://huggingface.co/deepseek-ai/DeepSeek-V2-Lite/resolve/main/tokenization_deepseek_fast.py -""" - -from typing import List, Optional, Union - -from transformers.models.llama import LlamaTokenizerFast - - -class DeepseekTokenizerFast(LlamaTokenizerFast): - def convert_ids_to_tokens( - self, ids: Union[int, List[int]], skip_special_tokens: bool = False - ) -> Union[str, List[str]]: - """ - Converts a single index or a sequence of indices in a token or a sequence of tokens, using the vocabulary and - added tokens. - Args: - ids (`int` or `List[int]`): - The token id (or token ids) to convert to tokens. - skip_special_tokens (`bool`, *optional*, defaults to `False`): - Whether or not to remove special tokens in the decoding. - Returns: - `str` or `List[str]`: The decoded token(s). - """ - if isinstance(ids, int): - return self._convert_id_to_token(ids) - tokens = [] - for index in ids: - index = int(index) - if skip_special_tokens and index in self.all_special_ids: - continue - token = self._tokenizer.id_to_token(index) - tokens.append(token if token is not None else "") - return tokens - - def _convert_id_to_token(self, index: int) -> Optional[str]: - token = self._tokenizer.id_to_token(int(index)) - return token if token is not None else "" From 4d0d67c678ad94f58e9606f461916bca29314a6a Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Tue, 19 Aug 2025 08:22:14 +0000 Subject: [PATCH 21/22] finish model updates --- optimum/habana/transformers/modeling_utils.py | 2 +- .../transformers/models/bart/modeling_bart.py | 214 ++++++-------- .../transformers/models/bert/modeling_bert.py | 80 ++++-- .../models/blip/modeling_blip_text.py | 131 +++++---- .../transformers/models/clip/modeling_clip.py | 136 ++++----- .../transformers/models/gpt2/modeling_gpt2.py | 101 +++---- .../gpt_bigcode/modeling_gpt_bigcode.py | 82 +++--- .../models/idefics2/modeling_idefics2.py | 9 +- .../models/mllama/modeling_mllama.py | 261 ++++-------------- .../transformers/models/opt/modeling_opt.py | 125 +++------ .../models/paligemma/modeling_paligemma.py | 52 ++-- .../models/persimmon/modeling_persimmon.py | 94 +++---- .../models/qwen3_moe/modeling_qwen3_moe.py | 4 +- .../transformers/models/t5/modeling_t5.py | 1 - .../models/wav2vec2/modeling_wav2vec2.py | 4 +- .../models/whisper/modeling_whisper.py | 81 +++--- optimum/habana/transformers/trainer.py | 2 +- tests/test_trainer.py | 4 +- 18 files changed, 523 insertions(+), 860 deletions(-) diff --git a/optimum/habana/transformers/modeling_utils.py b/optimum/habana/transformers/modeling_utils.py index 4a040c85ab..b5e57094a3 100644 --- a/optimum/habana/transformers/modeling_utils.py +++ b/optimum/habana/transformers/modeling_utils.py @@ -766,7 +766,7 @@ def adapt_transformers_to_gaudi(): transformers.models.video_llava.processing_video_llava.VideoLlavaProcessor = GaudiVideoLlavaProcessor # Optimization for Whisper on Gaudi - transformers.models.whisper.modeling_whisper.WhisperSdpaAttention = GaudiWhisperSdpaAttention + transformers.models.whisper.modeling_whisper.WhisperAttention = GaudiWhisperSdpaAttention transformers.models.whisper.modeling_whisper.WhisperDecoderLayer = GaudiWhisperDecoderLayer transformers.models.whisper.modeling_whisper.WhisperDecoder = GaudiWhisperDecoder transformers.models.whisper.modeling_whisper.WhisperModel = GaudiWhisperModel diff --git a/optimum/habana/transformers/models/bart/modeling_bart.py b/optimum/habana/transformers/models/bart/modeling_bart.py index 6dbb19b0e5..0482e80a23 100644 --- a/optimum/habana/transformers/models/bart/modeling_bart.py +++ b/optimum/habana/transformers/models/bart/modeling_bart.py @@ -15,7 +15,7 @@ # limitations under the License. """PyTorch BART model.""" -from typing import List, Optional, Tuple, Union +from typing import Optional, Union import torch import torch.utils.checkpoint @@ -32,7 +32,7 @@ Seq2SeqLMOutput, Seq2SeqModelOutput, ) -from transformers.models.bart.modeling_bart import shift_tokens_right +from transformers.models.bart.modeling_bart import eager_attention_forward, shift_tokens_right from transformers.utils import logging from ...modeling_attn_mask_utils import ( @@ -58,7 +58,12 @@ def __init__(self, num_embeddings: int, embedding_dim: int): self.offset = 2 super().__init__(num_embeddings + self.offset, embedding_dim) - def forward(self, input_ids: torch.Tensor, past_key_values_length: torch.Tensor = torch.tensor(0), position_ids: torch.Tensor = None): + def forward( + self, + input_ids: torch.Tensor, + past_key_values_length: torch.Tensor = torch.tensor(0), + position_ids: torch.Tensor = None, + ): """`input_ids' shape is expected to be [bsz x seqlen].""" bsz, seq_len = input_ids.shape[:2] @@ -90,7 +95,6 @@ def gaudi_BartAttention_forward( # determine input shapes bsz, tgt_len = hidden_states.shape[:-1] - src_len = key_value_states.shape[1] if is_cross_attention else tgt_len # get query proj query_states = self.q_proj(hidden_states) * self.scaling @@ -128,10 +132,10 @@ def gaudi_BartAttention_forward( value_states = self._shape(self.v_proj(hidden_states), -1, bsz) if self.is_decoder: - # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states. + # if cross_attention save tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states. # Further calls to cross_attention layer can then reuse all cross-attention # key/value_states (first "if" case) - # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of + # if uni-directional self-attention (decoder) save tuple(torch.Tensor, torch.Tensor) of # all previous decoder key/value_states. Further calls to uni-directional self-attention # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case) # if encoder bi-directional self-attention `past_key_value` is always `None` @@ -142,63 +146,23 @@ def gaudi_BartAttention_forward( key_states = key_states.reshape(*proj_shape) value_states = value_states.reshape(*proj_shape) - src_len = key_states.size(1) - attn_weights = torch.bmm(query_states, key_states.transpose(1, 2)) - - if attn_weights.size() != (bsz * self.num_heads, tgt_len, src_len): - raise ValueError( - f"Attention weights should be of size {(bsz * self.num_heads, tgt_len, src_len)}, but is" - f" {attn_weights.size()}" - ) - - if attention_mask is not None: - if attention_mask.size() != (bsz, 1, tgt_len, src_len): - raise ValueError( - f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is {attention_mask.size()}" - ) - attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + attention_mask - attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len) - - 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 {layer_head_mask.size()}" - ) - attn_weights = layer_head_mask.view(1, -1, 1, 1) * attn_weights.view(bsz, self.num_heads, tgt_len, src_len) - attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len) - - if output_attentions: - # this operation is a bit awkward, but it's required to - # make sure that attn_weights keeps its gradient. - # In order to do so, attn_weights have to be reshaped - # twice and have to be reused in the following - attn_weights_reshaped = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) - attn_weights = attn_weights_reshaped.view(bsz * self.num_heads, tgt_len, src_len) - else: - attn_weights_reshaped = None - - attn_probs = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training) - - attn_output = torch.bmm(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.view(bsz, self.num_heads, tgt_len, self.head_dim) - 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, attn_weights = eager_attention_forward( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.dropout, + scaling=self.scaling, + output_attentions=output_attentions, + head_mask=layer_head_mask, + **kwargs, + ) + attn_output = attn_output.reshape(bsz, tgt_len, -1).contiguous() attn_output = self.out_proj(attn_output) - return attn_output, attn_weights_reshaped, past_key_value + return attn_output, attn_weights, past_key_value def gaudi_BartEncoderLayer_forward( @@ -208,7 +172,7 @@ def gaudi_BartEncoderLayer_forward( layer_head_mask: torch.FloatTensor, output_attentions: Optional[bool] = False, token_idx: Optional[torch.Tensor] = None, -) -> Tuple[torch.FloatTensor, Optional[torch.FloatTensor]]: +) -> tuple[torch.FloatTensor, Optional[torch.FloatTensor]]: residual = hidden_states hidden_states, attn_weights, _ = self.self_attn( hidden_states=hidden_states, @@ -245,11 +209,12 @@ def gaudi_BartDecoderLayer_forward( encoder_attention_mask: Optional[torch.Tensor] = None, layer_head_mask: Optional[torch.Tensor] = None, cross_attn_layer_head_mask: Optional[torch.Tensor] = None, - past_key_value: Optional[Tuple[torch.Tensor]] = None, + past_key_value: Optional[tuple[torch.Tensor]] = None, output_attentions: Optional[bool] = False, use_cache: Optional[bool] = True, + cache_position: Optional[torch.Tensor] = None, token_idx: Optional[torch.Tensor] = None, -) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: +) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]: residual = hidden_states # Self Attention @@ -262,6 +227,7 @@ def gaudi_BartDecoderLayer_forward( attention_mask=attention_mask, layer_head_mask=layer_head_mask, output_attentions=output_attentions, + cache_position=cache_position, token_idx=token_idx, ) hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) @@ -289,6 +255,7 @@ def gaudi_BartDecoderLayer_forward( layer_head_mask=cross_attn_layer_head_mask, past_key_value=cross_attn_past_key_value, output_attentions=output_attentions, + cache_position=cache_position, token_idx=token_idx, ) hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) @@ -328,7 +295,7 @@ def gaudi_BartEncoder_forward( output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, token_idx: Optional[torch.Tensor] = None, -) -> Union[Tuple, BaseModelOutput]: +) -> Union[tuple, BaseModelOutput]: 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 @@ -359,16 +326,10 @@ def gaudi_BartEncoder_forward( hidden_states = self.layernorm_embedding(hidden_states) hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) - # expand attention_mask - if attention_mask is not None: - if self._use_sdpa and head_mask is None and not output_attentions: - # output_attentions=True & head_mask can not be supported when using SDPA, fall back to - # the manual implementation that requires a 4D causal mask in all cases. - # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] - attention_mask = _prepare_4d_attention_mask_for_sdpa(attention_mask, inputs_embeds.dtype) - else: - # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] - attention_mask = _prepare_4d_attention_mask(attention_mask, inputs_embeds.dtype) + attention_mask = self._update_full_mask( + attention_mask, + inputs_embeds, + ) encoder_states = () if output_hidden_states else None all_attentions = () if output_attentions else None @@ -384,7 +345,7 @@ def gaudi_BartEncoder_forward( 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) + # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description) to_drop = False if self.training: dropout_probability = torch.rand([]) @@ -394,23 +355,13 @@ def gaudi_BartEncoder_forward( 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, - None, - ) - 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, - token_idx=token_idx, - ) + 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, + token_idx=token_idx, + ) hidden_states = layer_outputs[0] @@ -433,14 +384,15 @@ def gaudi_BartDecoder_forward( encoder_attention_mask: Optional[torch.LongTensor] = None, head_mask: Optional[torch.Tensor] = None, cross_attn_head_mask: Optional[torch.Tensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, + past_key_values: Optional[list[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, -) -> Union[Tuple, BaseModelOutputWithPastAndCrossAttentions]: +) -> Union[tuple, BaseModelOutputWithPastAndCrossAttentions]: 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 @@ -448,8 +400,15 @@ def gaudi_BartDecoder_forward( use_cache = use_cache if use_cache is not None else self.config.use_cache return_dict = return_dict if return_dict is not None else self.config.use_return_dict + if self.gradient_checkpointing and self.training: + if use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." + ) + use_cache = False + # retrieve input_ids and inputs_embeds - if input_ids is not None and inputs_embeds is not None: + if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time") elif input_ids is not None: input = input_ids @@ -514,13 +473,6 @@ def gaudi_BartDecoder_forward( hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) - if self.gradient_checkpointing and self.training: - if use_cache: - logger.warning_once( - "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." - ) - use_cache = False - # decoder layers all_hidden_states = () if output_hidden_states else None all_self_attns = () if output_attentions else None @@ -537,7 +489,7 @@ def gaudi_BartDecoder_forward( ) for idx, decoder_layer in enumerate(self.layers): - # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description) + # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description) if output_hidden_states: all_hidden_states += (hidden_states,) if self.training: @@ -547,33 +499,19 @@ def gaudi_BartDecoder_forward( past_key_value = past_key_values[idx] if past_key_values is not None else None - if self.gradient_checkpointing and self.training: - layer_outputs = self._gradient_checkpointing_func( - decoder_layer.__call__, - hidden_states, - attention_mask, - encoder_hidden_states, - encoder_attention_mask, - head_mask[idx] if head_mask is not None else None, - cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None, - None, - output_attentions, - use_cache, - None, - ) - else: - layer_outputs = decoder_layer( - hidden_states, - attention_mask=attention_mask, - encoder_hidden_states=encoder_hidden_states, - encoder_attention_mask=encoder_attention_mask, - layer_head_mask=(head_mask[idx] if head_mask is not None else None), - cross_attn_layer_head_mask=(cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None), - past_key_value=past_key_value, - output_attentions=output_attentions, - use_cache=use_cache, - token_idx=token_idx, - ) + layer_outputs = decoder_layer( + hidden_states, + attention_mask, + encoder_hidden_states, # as a positional argument for gradient checkpointing + encoder_attention_mask=encoder_attention_mask, + layer_head_mask=(head_mask[idx] if head_mask is not None else None), + cross_attn_layer_head_mask=(cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None), + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + token_idx=token_idx, + ) hidden_states = layer_outputs[0] if use_cache: @@ -614,16 +552,17 @@ def gaudi_BartModel_forward( head_mask: Optional[torch.Tensor] = None, decoder_head_mask: Optional[torch.Tensor] = None, cross_attn_head_mask: Optional[torch.Tensor] = None, - encoder_outputs: Optional[List[torch.FloatTensor]] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, + encoder_outputs: Optional[list[torch.FloatTensor]] = None, + past_key_values: Optional[list[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, decoder_inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, -) -> Union[Tuple, Seq2SeqModelOutput]: +) -> Union[tuple, Seq2SeqModelOutput]: # different to other models, Bart automatically creates decoder_input_ids from # input_ids if no decoder_input_ids are provided if decoder_input_ids is None and decoder_inputs_embeds is None: @@ -676,6 +615,7 @@ def gaudi_BartModel_forward( output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, + cache_position=cache_position, token_idx=token_idx, ) @@ -703,8 +643,8 @@ def gaudi_BartForConditionalGeneration_forward( head_mask: Optional[torch.Tensor] = None, decoder_head_mask: Optional[torch.Tensor] = None, cross_attn_head_mask: Optional[torch.Tensor] = None, - encoder_outputs: Optional[List[torch.FloatTensor]] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, + encoder_outputs: Optional[list[torch.FloatTensor]] = None, + past_key_values: Optional[list[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, decoder_inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, @@ -712,8 +652,9 @@ def gaudi_BartForConditionalGeneration_forward( output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, -) -> Union[Tuple, Seq2SeqLMOutput]: +) -> Union[tuple, Seq2SeqLMOutput]: return_dict = return_dict if return_dict is not None else self.config.use_return_dict if labels is not None: @@ -741,6 +682,7 @@ def gaudi_BartForConditionalGeneration_forward( output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, + cache_position=cache_position, token_idx=token_idx, ) diff --git a/optimum/habana/transformers/models/bert/modeling_bert.py b/optimum/habana/transformers/models/bert/modeling_bert.py index 02f18ece02..c67a448e49 100644 --- a/optimum/habana/transformers/models/bert/modeling_bert.py +++ b/optimum/habana/transformers/models/bert/modeling_bert.py @@ -1,8 +1,9 @@ -from typing import List, Optional, Tuple, Union +from typing import Optional, Union import torch import torch.utils.checkpoint from habana_frameworks.torch.hpex.kernels import FusedSDPA +from transformers.cache_utils import Cache, EncoderDecoderCache from transformers.models.bert.modeling_bert import BaseModelOutputWithPoolingAndCrossAttentions, BertSdpaSelfAttention from optimum.utils import logging @@ -21,12 +22,13 @@ def gaudi_BertModel_forward( inputs_embeds: Optional[torch.Tensor] = None, encoder_hidden_states: Optional[torch.Tensor] = None, encoder_attention_mask: Optional[torch.Tensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, + past_key_values: Optional[list[torch.FloatTensor]] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, -) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPoolingAndCrossAttentions]: + cache_position: Optional[torch.Tensor] = None, +) -> Union[tuple[torch.Tensor], BaseModelOutputWithPoolingAndCrossAttentions]: r""" Copied from https://github.com/huggingface/transformers/blob/15c74a28294fe9082b81b24efe58df16fed79a9e/src/transformers/models/bert/modeling_bert.py Changes: @@ -56,8 +58,13 @@ def gaudi_BertModel_forward( batch_size, seq_length = input_shape device = input_ids.device if input_ids is not None else inputs_embeds.device - # past_key_values_length - past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0 + past_key_values_length = 0 + if past_key_values is not None: + past_key_values_length = ( + past_key_values[0][0].shape[-2] + if not isinstance(past_key_values, Cache) + else past_key_values.get_seq_length() + ) if token_type_ids is None: if hasattr(self.embeddings, "token_type_ids"): @@ -111,6 +118,7 @@ def gaudi_BertModel_forward( output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, + cache_position=cache_position, ) sequence_output = encoder_outputs[0] pooled_output = self.pooler(sequence_output) if self.pooler is not None else None @@ -135,9 +143,10 @@ def gaudi_Bert_Sdpa_SelfAttention_forward( head_mask: Optional[torch.FloatTensor] = None, encoder_hidden_states: Optional[torch.FloatTensor] = None, encoder_attention_mask: Optional[torch.FloatTensor] = None, - past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, + past_key_value: Optional[tuple[tuple[torch.FloatTensor]]] = None, output_attentions: Optional[bool] = False, -) -> Tuple[torch.Tensor]: + cache_position: Optional[torch.Tensor] = None, +) -> tuple[torch.Tensor]: r""" Copied from https://github.com/huggingface/transformers/blob/v4.46.3/src/transformers/models/bert/modeling_bert.py Changes: @@ -160,32 +169,54 @@ def gaudi_Bert_Sdpa_SelfAttention_forward( encoder_attention_mask, past_key_value, output_attentions, + cache_position, ) bsz, tgt_len, _ = hidden_states.size() - query_layer = self.transpose_for_scores(self.query(hidden_states)) + query_layer = ( + self.query(hidden_states).view(bsz, -1, self.num_attention_heads, self.attention_head_size).transpose(1, 2) + ) is_cross_attention = encoder_hidden_states is not None - current_states = encoder_hidden_states if is_cross_attention else hidden_states - attention_mask = encoder_attention_mask if is_cross_attention else attention_mask + if past_key_value is not None: + if isinstance(past_key_value, EncoderDecoderCache): + 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 + curr_past_key_value = past_key_value.cross_attention_cache + else: + curr_past_key_value = past_key_value.self_attention_cache + else: + curr_past_key_value = past_key_value - if is_cross_attention and past_key_value and past_key_value[0].shape[2] == current_states.shape[1]: - key_layer, value_layer = past_key_value + current_states = encoder_hidden_states if is_cross_attention else hidden_states + if is_cross_attention and past_key_value is not None and is_updated: + # reuse k,v, cross_attentions + key_layer = curr_past_key_value.layers[self.layer_idx].keys + value_layer = curr_past_key_value.layers[self.layer_idx].values else: - key_layer = self.transpose_for_scores(self.key(current_states)) - value_layer = self.transpose_for_scores(self.value(current_states)) - if past_key_value is not None and not is_cross_attention: - key_layer = torch.cat([past_key_value[0], key_layer], dim=2) - value_layer = torch.cat([past_key_value[1], value_layer], dim=2) + key_layer = ( + self.key(current_states).view(bsz, -1, self.num_attention_heads, self.attention_head_size).transpose(1, 2) + ) + value_layer = ( + self.value(current_states) + .view(bsz, -1, self.num_attention_heads, self.attention_head_size) + .transpose(1, 2) + ) - if self.is_decoder: - past_key_value = (key_layer, value_layer) + if past_key_value is not None: + # save all key/value_layer to cache to be re-used for fast auto-regressive generation + cache_position = cache_position if not is_cross_attention else None + key_layer, value_layer = curr_past_key_value.update( + key_layer, value_layer, self.layer_idx, {"cache_position": cache_position} + ) + # set flag that curr layer for cross-attn is already updated so we can re-use in subsequent calls + if is_cross_attention: + past_key_value.is_updated[self.layer_idx] = True - is_causal = ( - True if self.is_decoder and not is_cross_attention and attention_mask is None and tgt_len > 1 else False - ) + is_causal = self.is_decoder and not is_cross_attention and attention_mask is None and tgt_len > 1 attention_mask = attention_mask.to(query_layer.dtype) softmax_algo = "None" @@ -198,7 +229,4 @@ def gaudi_Bert_Sdpa_SelfAttention_forward( attn_output = attn_output.transpose(1, 2) attn_output = attn_output.reshape(bsz, tgt_len, self.all_head_size) - outputs = (attn_output,) - if self.is_decoder: - outputs = outputs + (past_key_value,) - return outputs + return attn_output, None diff --git a/optimum/habana/transformers/models/blip/modeling_blip_text.py b/optimum/habana/transformers/models/blip/modeling_blip_text.py index 960a4db17d..7c22969a93 100644 --- a/optimum/habana/transformers/models/blip/modeling_blip_text.py +++ b/optimum/habana/transformers/models/blip/modeling_blip_text.py @@ -1,5 +1,5 @@ import math -from typing import List, Optional, Tuple, Union +from typing import Optional, Union import torch from torch import nn @@ -23,7 +23,7 @@ def gaudi_BlipTextSelfAttention_forward( head_mask: Optional[torch.FloatTensor] = None, encoder_hidden_states: Optional[torch.FloatTensor] = None, encoder_attention_mask: Optional[torch.FloatTensor] = None, - past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, + past_key_value: Optional[tuple[tuple[torch.FloatTensor]]] = None, output_attentions: Optional[bool] = False, cache_position: Optional[torch.Tensor] = None, token_idx: Optional[torch.Tensor] = None, @@ -47,12 +47,28 @@ def gaudi_BlipTextSelfAttention_forward( attention_mask = encoder_attention_mask if is_cross_attention else attention_mask if is_cross_attention: - key_layer = self.transpose_for_scores(self.key(encoder_hidden_states)) - value_layer = self.transpose_for_scores(self.value(encoder_hidden_states)) + key_layer = ( + self.key(encoder_hidden_states) + .view(batch_size, -1, self.num_attention_heads, self.attention_head_size) + .transpose(1, 2) + ) + value_layer = ( + self.value(encoder_hidden_states) + .view(batch_size, -1, self.num_attention_heads, self.attention_head_size) + .transpose(1, 2) + ) attention_mask = encoder_attention_mask elif past_key_value is not None: - key_layer = self.transpose_for_scores(self.key(hidden_states)) - value_layer = self.transpose_for_scores(self.value(hidden_states)) + key_layer = ( + self.key(hidden_states) + .view(batch_size, -1, self.num_attention_heads, self.attention_head_size) + .transpose(1, 2) + ) + value_layer = ( + self.value(hidden_states) + .view(batch_size, -1, self.num_attention_heads, self.attention_head_size) + .transpose(1, 2) + ) if token_idx is not None: past_key_value[0].index_copy_(2, token_idx - 1, key_layer) past_key_value[1].index_copy_(2, token_idx - 1, value_layer) @@ -62,10 +78,16 @@ def gaudi_BlipTextSelfAttention_forward( key_layer = torch.cat([past_key_value[0], key_layer], dim=2) value_layer = torch.cat([past_key_value[1], value_layer], dim=2) else: - key_layer = self.transpose_for_scores(self.key(hidden_states)) - value_layer = self.transpose_for_scores(self.value(hidden_states)) - - query_layer = self.transpose_for_scores(mixed_query_layer) + key_layer = ( + self.key(hidden_states) + .view(batch_size, -1, self.num_attention_heads, self.attention_head_size) + .transpose(1, 2) + ) + value_layer = ( + self.value(hidden_states) + .view(batch_size, -1, self.num_attention_heads, self.attention_head_size) + .transpose(1, 2) + ) past_key_value = (key_layer, value_layer) @@ -122,11 +144,11 @@ def gaudi_BlipTextAttention_forward( attention_mask: Optional[torch.FloatTensor] = None, head_mask: Optional[torch.FloatTensor] = None, encoder_hidden_states: Optional[torch.FloatTensor] = None, - encoder_attention_mask: Optional[torch.FloatTensor] = None, - past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, + past_key_value: Optional[tuple[tuple[torch.FloatTensor]]] = None, output_attentions: Optional[bool] = False, + cache_position: Optional[torch.Tensor] = None, token_idx: Optional[torch.Tensor] = None, -) -> Tuple[torch.Tensor]: +) -> tuple[torch.Tensor]: """ Copied from BlipTextAttention.forward: https://github.com/huggingface/transformers/blob/v4.37.2/src/transformers/models/blip/modeling_blip_text.py#L265 The only differences are: @@ -134,12 +156,12 @@ def gaudi_BlipTextAttention_forward( """ self_outputs = self.self( hidden_states, - attention_mask, - head_mask, - encoder_hidden_states, - encoder_attention_mask, - past_key_value, - output_attentions, + attention_mask=attention_mask, + head_mask=head_mask, + encoder_hidden_states=encoder_hidden_states, + past_key_value=past_key_value, + output_attentions=output_attentions, + cache_position=cache_position, token_idx=token_idx, ) attention_output = self.output(self_outputs[0], hidden_states) @@ -154,10 +176,11 @@ def gaudi_BlipTextLayer_forward( head_mask: Optional[torch.FloatTensor] = None, encoder_hidden_states: Optional[torch.FloatTensor] = None, encoder_attention_mask: Optional[torch.FloatTensor] = None, - past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, + past_key_value: Optional[tuple[tuple[torch.FloatTensor]]] = None, output_attentions: Optional[bool] = False, + cache_position: Optional[torch.Tensor] = None, token_idx: Optional[torch.Tensor] = None, -) -> Tuple[torch.Tensor]: +) -> tuple[torch.Tensor]: """ Copied from BlipTextLayer.forward: https://github.com/huggingface/transformers/blob/v4.37.2/src/transformers/models/blip/modeling_blip_text.py#L333 The only differences are: @@ -166,10 +189,11 @@ def gaudi_BlipTextLayer_forward( self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None self_attention_outputs = self.attention( hidden_states, - attention_mask, - head_mask, + attention_mask=attention_mask, + head_mask=head_mask, output_attentions=output_attentions, past_key_value=self_attn_past_key_value, + cache_position=cache_position, token_idx=token_idx, ) attention_output = self_attention_outputs[0] @@ -180,11 +204,12 @@ def gaudi_BlipTextLayer_forward( if encoder_hidden_states is not None: cross_attention_outputs = self.crossattention( attention_output, - attention_mask, - head_mask, - encoder_hidden_states, - encoder_attention_mask, + attention_mask=encoder_attention_mask, + head_mask=head_mask, + encoder_hidden_states=encoder_hidden_states, + past_key_value=past_key_value, output_attentions=output_attentions, + cache_position=cache_position, ) attention_output = cross_attention_outputs[0] outputs = outputs + cross_attention_outputs[1:-1] # add cross attentions if we output attention weights @@ -205,13 +230,14 @@ def gaudi_BlipTextEncoder_forward( head_mask: Optional[torch.FloatTensor] = None, encoder_hidden_states: Optional[torch.FloatTensor] = None, encoder_attention_mask: Optional[torch.FloatTensor] = None, - past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, + past_key_values: Optional[tuple[tuple[torch.FloatTensor]]] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = False, output_hidden_states: Optional[bool] = False, return_dict: Optional[bool] = True, + cache_position: Optional[torch.Tensor] = None, token_idx: Optional[torch.Tensor] = None, -) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPastAndCrossAttentions]: +) -> Union[tuple[torch.Tensor], BaseModelOutputWithPastAndCrossAttentions]: """ Copied from BlipTextEncoder.forward: https://github.com/huggingface/transformers/blob/v4.37.2/src/transformers/models/blip/modeling_blip_text.py#L391 The only differences are: @@ -237,28 +263,17 @@ def gaudi_BlipTextEncoder_forward( layer_head_mask = head_mask[i] if head_mask is not None else None past_key_value = past_key_values[i] if past_key_values is not None else None - if self.gradient_checkpointing and self.training: - layer_outputs = self._gradient_checkpointing_func( - layer_module.__call__, - hidden_states, - attention_mask, - layer_head_mask, - encoder_hidden_states, - encoder_attention_mask, - past_key_value, - output_attentions, - ) - else: - layer_outputs = layer_module( - hidden_states, - attention_mask, - layer_head_mask, - encoder_hidden_states, - encoder_attention_mask, - past_key_value, - output_attentions, - token_idx=token_idx, - ) + layer_outputs = layer_module( + hidden_states, + attention_mask, + layer_head_mask, + encoder_hidden_states, + encoder_attention_mask, + past_key_value, + output_attentions, + cache_position, + token_idx=token_idx, + ) hidden_states = layer_outputs[0] if use_cache: @@ -301,14 +316,15 @@ def gaudi_BlipTextModel_forward( encoder_embeds: Optional[torch.Tensor] = None, encoder_hidden_states: Optional[torch.Tensor] = None, encoder_attention_mask: Optional[torch.Tensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, + past_key_values: Optional[list[torch.FloatTensor]] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, is_decoder: Optional[bool] = False, + cache_position: Optional[torch.Tensor] = None, token_idx: Optional[torch.Tensor] = None, -) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPoolingAndCrossAttentions]: +) -> Union[tuple[torch.Tensor], BaseModelOutputWithPoolingAndCrossAttentions]: """ Copied from BlipTextModel.forward: https://github.com/huggingface/transformers/blob/v4.37.2/src/transformers/models/blip/modeling_blip_text.py#L666 The only differences are: @@ -347,7 +363,7 @@ def gaudi_BlipTextModel_forward( past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0 if attention_mask is None: - attention_mask = torch.ones(((batch_size, seq_length + past_key_values_length))).to(device) + attention_mask = torch.ones((batch_size, seq_length + past_key_values_length)).to(device) # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length] # ourselves in which case we just need to make it broadcastable to all heads. @@ -402,6 +418,7 @@ def gaudi_BlipTextModel_forward( output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, + cache_position=cache_position, token_idx=token_idx, ) sequence_output = encoder_outputs[0] @@ -430,7 +447,7 @@ def gaudi_BlipTextLMHead_forward( encoder_hidden_states: Optional[torch.Tensor] = None, encoder_attention_mask: Optional[torch.Tensor] = None, labels: Optional[torch.Tensor] = None, - past_key_values: Optional[List[torch.Tensor]] = None, + past_key_values: Optional[list[torch.Tensor]] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, @@ -438,8 +455,9 @@ def gaudi_BlipTextLMHead_forward( return_logits: Optional[bool] = False, is_decoder: Optional[bool] = True, reduction: Optional[str] = "mean", + cache_position: Optional[torch.Tensor] = None, token_idx: Optional[torch.Tensor] = None, -) -> Union[Tuple[torch.Tensor], CausalLMOutputWithCrossAttentions]: +) -> Union[tuple[torch.Tensor], CausalLMOutputWithCrossAttentions]: """ Copied from BlipTextLMHeadModel.forward: https://github.com/huggingface/transformers/blob/v4.37.2/src/transformers/models/blip/modeling_blip_text.py#L820 The only differences are: @@ -463,6 +481,7 @@ def gaudi_BlipTextLMHead_forward( output_hidden_states=output_hidden_states, return_dict=return_dict, is_decoder=is_decoder, + cache_position=cache_position, token_idx=token_idx, ) diff --git a/optimum/habana/transformers/models/clip/modeling_clip.py b/optimum/habana/transformers/models/clip/modeling_clip.py index 274a8f12bd..e0177d18f8 100644 --- a/optimum/habana/transformers/models/clip/modeling_clip.py +++ b/optimum/habana/transformers/models/clip/modeling_clip.py @@ -1,9 +1,9 @@ -from typing import Optional, Tuple +from typing import Optional, Union import torch from torch import nn from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling -from transformers.models.clip.configuration_clip import CLIPConfig +from transformers.models.clip.configuration_clip import CLIPTextConfig, CLIPVisionConfig from transformers.models.clip.modeling_clip import ( CLIPMLP, CLIPAttention, @@ -61,11 +61,11 @@ def __init__(self): super().__init__() def forward(self, x, dim=None, invAttnHead=None): - return torch.nn.functional.softmax(x, dim) + return torch.nn.functional.softmax(x, dim, dtype=torch.float32) class GaudiCLIPAttention(CLIPAttention): - def __init__(self, config): + def __init__(self, config: Union[CLIPVisionConfig, CLIPTextConfig]): super().__init__(config=config) self.fused_scaled_dot_product_attention = ModuleFusedSDPA(FusedSDPA) if FusedSDPA else None self.bmm1 = Matmul() @@ -81,7 +81,7 @@ def forward( use_flash_attention: Optional[bool] = False, flash_attention_recompute: Optional[bool] = False, flash_attention_fast_softmax: Optional[bool] = False, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: """ Copied from CLIPAttention.forward: https://github.com/huggingface/transformers/blob/ab0f050b42d903f34d6eb97f3f8c0c07f0517ad2/src/transformers/models/clip/modeling_clip.py The only differences are: @@ -89,28 +89,34 @@ def forward( - add new args flash_attention_recompute - add new args flash_attention_fast_softmax """ - bsz, tgt_len, _ = hidden_states.size() - attn_weights_reshaped = None - # get query proj - query_states = self.q_proj(hidden_states) * self.scale - key_states = self._shape(self.k_proj(hidden_states), -1, bsz) - value_states = self._shape(self.v_proj(hidden_states), -1, bsz) - - proj_shape = (bsz * self.num_heads, -1, self.head_dim) - query_states = self._shape(query_states, tgt_len, bsz).view(*proj_shape) - key_states = key_states.view(*proj_shape) - value_states = value_states.view(*proj_shape) - - src_len = key_states.size(1) + batch_size, seq_length, embed_dim = hidden_states.shape + + queries = self.q_proj(hidden_states) + keys = self.k_proj(hidden_states) + values = self.v_proj(hidden_states) + + queries = queries.view(batch_size, seq_length, -1, self.head_dim).transpose(1, 2) + keys = keys.view(batch_size, seq_length, -1, self.head_dim).transpose(1, 2) + values = values.view(batch_size, seq_length, -1, self.head_dim).transpose(1, 2) + # CLIP text model uses both `causal_attention_mask` and `attention_mask` + # in case FA2 kernel is called, `is_causal` should be inferred from `causal_attention_mask` + if self.config._attn_implementation == "flash_attention_2": + self.is_causal = causal_attention_mask is not None + else: + if attention_mask is not None and causal_attention_mask is not None: + attention_mask = attention_mask + causal_attention_mask + elif causal_attention_mask is not None: + attention_mask = causal_attention_mask + if FusedSDPA and use_flash_attention: import habana_frameworks.torch.hpu as ht softmax_mode = "fast" if flash_attention_fast_softmax else "None" with ht.sdp_kernel(enable_recompute=flash_attention_recompute): attn_output = self.fused_scaled_dot_product_attention( - query_states, - key_states, - value_states, + queries, + keys, + values, attention_mask, self.dropout, False, @@ -118,64 +124,27 @@ def forward( softmax_mode, ) else: - attn_weights = self.bmm1(query_states, key_states.transpose(1, 2)) - if attn_weights.size() != (bsz * self.num_heads, tgt_len, src_len): - raise ValueError( - f"Attention weights should be of size {(bsz * self.num_heads, tgt_len, src_len)}, but is" - f" {attn_weights.size()}" - ) - - # apply the causal_attention_mask first - if causal_attention_mask is not None: - if causal_attention_mask.size() != (bsz, 1, tgt_len, src_len): - raise ValueError( - f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is" - f" {causal_attention_mask.size()}" - ) - attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + causal_attention_mask - attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len) - + attn_weights = self.bmm1(queries, keys.transpose(1, 2)) * self.scale if attention_mask is not None: - if attention_mask.size() != (bsz, 1, tgt_len, src_len): - raise ValueError( - f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is {attention_mask.size()}" - ) - attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + attention_mask - attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len) + attn_weights = attn_weights + attention_mask + attn_weights = self.softmax(attn_weights, dim=-1).to(queries.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training) - attn_weights = self.softmax(attn_weights, dim=-1) + attn_output = self.bmm2(attn_weights, values) - if output_attentions: - # this operation is a bit awkward, but it's required to - # make sure that attn_weights keeps its gradient. - # In order to do so, attn_weights have to reshaped - # twice and have to be reused in the following - attn_weights_reshaped = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) - attn_weights = attn_weights_reshaped.view(bsz * self.num_heads, tgt_len, src_len) - else: - attn_weights_reshaped = None - - attn_probs = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training) - - attn_output = self.bmm2(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.view(bsz, self.num_heads, tgt_len, self.head_dim) - attn_output = attn_output.transpose(1, 2) - attn_output = attn_output.reshape(bsz, tgt_len, -1) + # attn_output = attn_output.view(bsz, self.num_heads, tgt_len, self.head_dim) + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.reshape(batch_size, seq_length, embed_dim).contiguous() attn_output = self.out_proj(attn_output) - return attn_output, attn_weights_reshaped + if not output_attentions: + attn_weights = None + return attn_output, attn_weights class GaudiCLIPEncoderLayer(CLIPEncoderLayer): - def __init__(self, config: CLIPConfig): + def __init__(self, config: Union[CLIPVisionConfig, CLIPTextConfig]): super(CLIPEncoderLayer, self).__init__() self.embed_dim = config.hidden_size self.self_attn = GaudiCLIPAttention(config) @@ -191,7 +160,7 @@ def forward( output_attentions: Optional[bool] = False, use_flash_attention: Optional[bool] = False, flash_attention_recompute: Optional[bool] = False, - ) -> Tuple[torch.FloatTensor]: + ) -> tuple[torch.FloatTensor]: """ Copied from CLIPEncoderLayer.forward: https://github.com/huggingface/transformers/blob/ab0f050b42d903f34d6eb97f3f8c0c07f0517ad2/src/transformers/models/clip/modeling_clip.py The only differences are: @@ -253,23 +222,14 @@ def forward( for idx, encoder_layer in enumerate(self.layers): if output_hidden_states: encoder_states = encoder_states + (hidden_states,) - if self.gradient_checkpointing and self.training: - layer_outputs = self._gradient_checkpointing_func( - encoder_layer.__call__, - hidden_states, - attention_mask, - causal_attention_mask, - output_attentions, - ) - else: - layer_outputs = encoder_layer( - hidden_states, - attention_mask, - causal_attention_mask, - output_attentions=output_attentions, - use_flash_attention=use_flash_attention, - flash_attention_recompute=flash_attention_recompute, - ) + layer_outputs = encoder_layer( + hidden_states, + attention_mask, + causal_attention_mask, + output_attentions=output_attentions, + use_flash_attention=use_flash_attention, + flash_attention_recompute=flash_attention_recompute, + ) hidden_states = layer_outputs[0] diff --git a/optimum/habana/transformers/models/gpt2/modeling_gpt2.py b/optimum/habana/transformers/models/gpt2/modeling_gpt2.py index ab7460c098..45d4fa59da 100644 --- a/optimum/habana/transformers/models/gpt2/modeling_gpt2.py +++ b/optimum/habana/transformers/models/gpt2/modeling_gpt2.py @@ -1,4 +1,4 @@ -from typing import Optional, Tuple, Union +from typing import Optional, Union import torch from torch.nn import CrossEntropyLoss @@ -76,8 +76,9 @@ def _upcast_and_reordered_attn(self, query, key, value, attention_mask=None, hea def forward( self, - hidden_states: Optional[Tuple[torch.FloatTensor]], - layer_past: Optional[Tuple[torch.Tensor]] = None, + hidden_states: Optional[tuple[torch.FloatTensor]], + layer_past: Optional[tuple[torch.Tensor]] = None, + cache_position: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.FloatTensor] = None, head_mask: Optional[torch.FloatTensor] = None, encoder_hidden_states: Optional[torch.Tensor] = None, @@ -85,7 +86,8 @@ def forward( use_cache: Optional[bool] = False, output_attentions: Optional[bool] = False, token_idx: Optional[torch.Tensor] = None, - ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]], ...]: + **kwargs, + ) -> tuple[Union[torch.Tensor, tuple[torch.Tensor]], ...]: if encoder_hidden_states is not None: if not hasattr(self, "q_attn"): raise ValueError( @@ -167,8 +169,9 @@ def __init__(self, config, layer_idx=None): def forward( self, - hidden_states: Optional[Tuple[torch.FloatTensor]], - layer_past: Optional[Tuple[torch.Tensor]] = None, + hidden_states: Optional[tuple[torch.FloatTensor]], + layer_past: Optional[tuple[torch.Tensor]] = None, + cache_position: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.FloatTensor] = None, head_mask: Optional[torch.FloatTensor] = None, encoder_hidden_states: Optional[torch.Tensor] = None, @@ -176,7 +179,8 @@ def forward( use_cache: Optional[bool] = False, output_attentions: Optional[bool] = False, token_idx: Optional[torch.Tensor] = None, - ) -> Union[Tuple[torch.Tensor], Optional[Tuple[torch.Tensor, Tuple[torch.FloatTensor, ...]]]]: + **kwargs, + ) -> Union[tuple[torch.Tensor], Optional[tuple[torch.Tensor, tuple[torch.FloatTensor, ...]]]]: """ Copied from GPT2Block.forward: https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt2/modeling_gpt2.py The only differences are: @@ -194,6 +198,7 @@ def forward( use_cache=use_cache, output_attentions=output_attentions, token_idx=token_idx, + **kwargs, ) attn_output = attn_outputs[0] # output_attn: a, present, (attentions) outputs = attn_outputs[1:] @@ -240,7 +245,8 @@ def forward( def gaudi_gpt2_forward( self, input_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None, + past_key_values: Optional[tuple[tuple[torch.Tensor]]] = None, + cache_position: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.FloatTensor] = None, token_type_ids: Optional[torch.LongTensor] = None, position_ids: Optional[torch.LongTensor] = None, @@ -253,7 +259,8 @@ def gaudi_gpt2_forward( output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, token_idx: Optional[torch.Tensor] = None, -) -> Union[Tuple, BaseModelOutputWithPastAndCrossAttentions]: + **kwargs, +) -> Union[tuple, BaseModelOutputWithPastAndCrossAttentions]: """ Copied from GPT2Model.forward: https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt2/modeling_gpt2.py The only differences are: @@ -358,45 +365,22 @@ def gaudi_gpt2_forward( all_hidden_states = () if output_hidden_states else None for i in range(len(self.h)): block, layer_past = self.h[i], past_key_values[i] - # Model parallel - if self.model_parallel: - torch.cuda.set_device(hidden_states.device) - # Ensure layer_past is on same device as hidden_states (might not be correct) - if layer_past is not None: - layer_past = tuple(past_state.to(hidden_states.device) for past_state in layer_past) - # Ensure that attention_mask is always on the same device as hidden_states - if attention_mask is not None: - attention_mask = attention_mask.to(hidden_states.device) - if isinstance(head_mask, torch.Tensor): - head_mask = head_mask.to(hidden_states.device) if output_hidden_states: all_hidden_states = all_hidden_states + (hidden_states,) - if self.gradient_checkpointing and self.training: - outputs = self._gradient_checkpointing_func( - block.__call__, - hidden_states, - None, - attention_mask, - head_mask[i], - encoder_hidden_states, - encoder_attention_mask, - use_cache, - output_attentions, - None, - ) - else: - outputs = block( - hidden_states, - layer_past=layer_past, - attention_mask=attention_mask, - head_mask=head_mask[i], - encoder_hidden_states=encoder_hidden_states, - encoder_attention_mask=encoder_attention_mask, - use_cache=use_cache, - output_attentions=output_attentions, - token_idx=token_idx, - ) + outputs = block( + hidden_states, + layer_past=layer_past, + cache_position=cache_position, + attention_mask=attention_mask, + head_mask=head_mask[i], + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + use_cache=use_cache, + output_attentions=output_attentions, + token_idx=token_idx, + **kwargs, + ) hidden_states = outputs[0] if use_cache is True: @@ -503,7 +487,8 @@ def prepare_inputs_for_generation( def forward( self, input_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None, + past_key_values: Optional[tuple[tuple[torch.Tensor]]] = None, + cache_position: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.FloatTensor] = None, token_type_ids: Optional[torch.LongTensor] = None, position_ids: Optional[torch.LongTensor] = None, @@ -516,15 +501,17 @@ def forward( output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, + logits_to_keep: Union[int, torch.Tensor] = 0, token_idx: Optional[torch.Tensor] = None, **kwargs, - ) -> Union[Tuple, CausalLMOutputWithCrossAttentions]: + ) -> Union[tuple, CausalLMOutputWithCrossAttentions]: return_dict = return_dict if return_dict is not None else self.config.use_return_dict transformer_outputs = self.transformer( input_ids, past_key_values=past_key_values, attention_mask=attention_mask, + cache_position=cache_position, token_type_ids=token_type_ids, position_ids=position_ids, head_mask=head_mask, @@ -539,30 +526,26 @@ def forward( ) hidden_states = transformer_outputs[0] - # Set device for model parallelism - if self.model_parallel: - torch.cuda.set_device(self.transformer.first_device) - hidden_states = hidden_states.to(self.lm_head.weight.device) - - lm_logits = self.lm_head(hidden_states) + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + logits = self.lm_head(hidden_states[:, slice_indices, :]) loss = None if labels is not None: # Flatten the tokens loss = self.loss_function( - lm_logits, + logits, labels, vocab_size=self.config.vocab_size, **kwargs, ) if not return_dict: - output = (lm_logits,) + transformer_outputs[1:] + output = (logits,) + transformer_outputs[1:] return ((loss,) + output) if loss is not None else output return CausalLMOutputWithCrossAttentions( loss=loss, - logits=lm_logits, + logits=logits, past_key_values=transformer_outputs.past_key_values, hidden_states=transformer_outputs.hidden_states, attentions=transformer_outputs.attentions, @@ -640,7 +623,8 @@ def prepare_inputs_for_generation( def forward( self, input_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None, + past_key_values: Optional[tuple[tuple[torch.Tensor]]] = None, + cache_position: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.FloatTensor] = None, token_type_ids: Optional[torch.LongTensor] = None, position_ids: Optional[torch.LongTensor] = None, @@ -655,12 +639,13 @@ def forward( return_dict: Optional[bool] = None, token_idx: Optional[torch.Tensor] = None, **kwargs, - ) -> Union[Tuple, GPT2DoubleHeadsModelOutput]: + ) -> Union[tuple, GPT2DoubleHeadsModelOutput]: return_dict = return_dict if return_dict is not None else self.config.use_return_dict transformer_outputs = self.transformer( input_ids, past_key_values=past_key_values, + cache_position=cache_position, attention_mask=attention_mask, token_type_ids=token_type_ids, position_ids=position_ids, diff --git a/optimum/habana/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py b/optimum/habana/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py index ffc27dc931..c7de607cad 100644 --- a/optimum/habana/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py +++ b/optimum/habana/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py @@ -17,7 +17,7 @@ ############################################################################### import math -from typing import List, Optional, Tuple, Union +from typing import Optional, Union import torch import torch.nn.functional as F @@ -288,15 +288,17 @@ def forward( encoder_attention_mask: Optional[torch.Tensor] = None, use_cache: Optional[bool] = False, output_attentions: Optional[bool] = False, + cache_position: Optional[torch.Tensor] = None, token_idx: Optional[torch.Tensor] = None, use_flash_attention: Optional[bool] = False, flash_attention_recompute: Optional[bool] = False, flash_attention_fast_softmax: Optional[bool] = False, flash_attention_causal_mask: Optional[bool] = False, cache_idx: Optional[int] = None, + **kwargs, ) -> Union[ - Tuple[torch.Tensor, Optional[torch.Tensor]], - Tuple[torch.Tensor, Optional[torch.Tensor], Tuple[torch.Tensor, ...]], + tuple[torch.Tensor, Optional[torch.Tensor]], + tuple[torch.Tensor, Optional[torch.Tensor], tuple[torch.Tensor, ...]], ]: """ Copied from GPTBigCodeAttention.forward: https://github.com/huggingface/transformers/blob/v4.40-release/src/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py @@ -403,7 +405,7 @@ def forward( def gaudi_gpt_bigcode_block_forward( self, - hidden_states: Optional[Tuple[torch.Tensor]], + hidden_states: Optional[tuple[torch.Tensor]], layer_past: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, head_mask: Optional[torch.Tensor] = None, @@ -411,6 +413,7 @@ def gaudi_gpt_bigcode_block_forward( encoder_attention_mask: Optional[torch.Tensor] = None, use_cache: Optional[bool] = False, output_attentions: Optional[bool] = False, + cache_position: Optional[torch.Tensor] = None, token_idx: Optional[torch.Tensor] = None, use_flash_attention: Optional[bool] = False, flash_attention_recompute: Optional[bool] = False, @@ -418,7 +421,7 @@ def gaudi_gpt_bigcode_block_forward( flash_attention_causal_mask: Optional[bool] = False, cache_idx: Optional[int] = None, **kwargs, -) -> Union[Tuple[torch.Tensor], Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: +) -> Union[tuple[torch.Tensor], tuple[torch.Tensor, torch.Tensor], tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: """ Copied from GPTBigCodeBlock.forward: https://github.com/huggingface/transformers/blob/v4.40-release/src/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py The only differences are: @@ -433,12 +436,14 @@ def gaudi_gpt_bigcode_block_forward( head_mask=head_mask, use_cache=use_cache, output_attentions=output_attentions, + cache_position=cache_position, token_idx=token_idx, use_flash_attention=use_flash_attention, flash_attention_recompute=flash_attention_recompute, flash_attention_fast_softmax=flash_attention_fast_softmax, flash_attention_causal_mask=flash_attention_causal_mask, cache_idx=cache_idx, + **kwargs, ) attn_output = attn_outputs[0] # output_attn: a, present, (attentions) outputs = attn_outputs[1:] @@ -461,6 +466,8 @@ def gaudi_gpt_bigcode_block_forward( encoder_hidden_states=encoder_hidden_states, encoder_attention_mask=encoder_attention_mask, output_attentions=output_attentions, + cache_position=cache_position, + **kwargs, ) attn_output = cross_attn_outputs[0] # residual connection @@ -484,7 +491,7 @@ def gaudi_gpt_bigcode_block_forward( def gaudi_gpt_bigcode_model_forward( self, input_ids: Optional[torch.Tensor] = None, - past_key_values: Optional[List[torch.Tensor]] = None, + past_key_values: Optional[list[torch.Tensor]] = None, attention_mask: Optional[torch.Tensor] = None, token_type_ids: Optional[torch.Tensor] = None, position_ids: Optional[torch.Tensor] = None, @@ -496,13 +503,15 @@ def gaudi_gpt_bigcode_model_forward( output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, + cache_position: Optional[torch.Tensor] = None, token_idx: Optional[torch.Tensor] = None, use_flash_attention: Optional[bool] = False, flash_attention_recompute: Optional[bool] = False, flash_attention_fast_softmax: Optional[bool] = False, flash_attention_causal_mask: Optional[bool] = False, cache_idx: Optional[int] = None, -) -> Union[Tuple, BaseModelOutputWithPastAndCrossAttentions]: + **kwargs, +) -> Union[tuple, BaseModelOutputWithPastAndCrossAttentions]: """ Copied from GPTBigCodeModel.forward: https://github.com/huggingface/transformers/blob/v4.40-release/src/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py The only differences are: @@ -520,10 +529,9 @@ def gaudi_gpt_bigcode_model_forward( use_cache = use_cache if use_cache is not None else self.config.use_cache return_dict = return_dict if return_dict is not None else self.config.use_return_dict - if input_ids is not None and inputs_embeds is not None: - raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") elif input_ids is not None: - self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask) input_shape = input_ids.size() input_ids = input_ids.view(-1, input_shape[-1]) batch_size = input_ids.shape[0] @@ -637,36 +645,24 @@ def gaudi_gpt_bigcode_model_forward( if output_hidden_states: all_hidden_states = all_hidden_states + (hidden_states,) - if self.gradient_checkpointing and self.training: - outputs = self._gradient_checkpointing_func( - block.__call__, - hidden_states, - None, - attention_mask, - head_mask[i], - encoder_hidden_states, - encoder_attention_mask, - use_cache, - output_attentions, - None, - ) - else: - outputs = block( - hidden_states, - layer_past=layer_past, - attention_mask=attention_mask, - head_mask=head_mask[i], - encoder_hidden_states=encoder_hidden_states, - encoder_attention_mask=encoder_attention_mask, - use_cache=use_cache, - output_attentions=output_attentions, - token_idx=token_idx, - use_flash_attention=use_flash_attention, - flash_attention_recompute=flash_attention_recompute, - flash_attention_fast_softmax=flash_attention_fast_softmax, - flash_attention_causal_mask=flash_attention_causal_mask, - cache_idx=cache_idx, - ) + outputs = block( + hidden_states, + layer_past=layer_past, + attention_mask=attention_mask, + head_mask=head_mask[i], + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + use_cache=use_cache, + output_attentions=output_attentions, + cache_position=cache_position, + token_idx=token_idx, + use_flash_attention=use_flash_attention, + flash_attention_recompute=flash_attention_recompute, + flash_attention_fast_softmax=flash_attention_fast_softmax, + flash_attention_causal_mask=flash_attention_causal_mask, + cache_idx=cache_idx, + **kwargs, + ) hidden_states = outputs[0] if use_cache: @@ -786,7 +782,7 @@ def prepare_inputs_for_generation( def forward( self, input_ids: Optional[torch.Tensor] = None, - past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None, + past_key_values: Optional[tuple[tuple[torch.Tensor]]] = None, attention_mask: Optional[torch.Tensor] = None, token_type_ids: Optional[torch.Tensor] = None, position_ids: Optional[torch.Tensor] = None, @@ -799,6 +795,7 @@ def forward( output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, + cache_position: Optional[torch.Tensor] = None, token_idx: Optional[torch.Tensor] = None, use_flash_attention: Optional[bool] = False, flash_attention_recompute: Optional[bool] = False, @@ -806,7 +803,7 @@ def forward( flash_attention_causal_mask: Optional[bool] = False, cache_idx: Optional[int] = None, **kwargs, - ) -> Union[Tuple, CausalLMOutputWithCrossAttentions]: + ) -> Union[tuple, CausalLMOutputWithCrossAttentions]: r""" labels (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set @@ -829,6 +826,7 @@ def forward( output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, + cache_position=cache_position, token_idx=token_idx, use_flash_attention=use_flash_attention, flash_attention_recompute=flash_attention_recompute, diff --git a/optimum/habana/transformers/models/idefics2/modeling_idefics2.py b/optimum/habana/transformers/models/idefics2/modeling_idefics2.py index 1d11134774..bef37c7470 100644 --- a/optimum/habana/transformers/models/idefics2/modeling_idefics2.py +++ b/optimum/habana/transformers/models/idefics2/modeling_idefics2.py @@ -86,9 +86,12 @@ def inputs_merger( - replace `==` with torch.where to fix the issue in hpu graph """ if input_ids is None: - special_image_mask = torch.where(inputs_embeds == self.get_input_embeddings()( - torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device) - )) + special_image_mask = torch.where( + inputs_embeds + == self.get_input_embeddings()( + torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device) + ) + ) special_image_mask = special_image_mask.all(-1) else: special_image_mask = torch.where(input_ids == self.config.image_token_id) diff --git a/optimum/habana/transformers/models/mllama/modeling_mllama.py b/optimum/habana/transformers/models/mllama/modeling_mllama.py index f0950db0d8..7f303f965a 100644 --- a/optimum/habana/transformers/models/mllama/modeling_mllama.py +++ b/optimum/habana/transformers/models/mllama/modeling_mllama.py @@ -222,7 +222,6 @@ def forward( self, hidden_state: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, - output_attentions: Optional[bool] = None, use_flash_attention: Optional[bool] = False, ): """ @@ -248,12 +247,7 @@ def forward( hidden_state = self.gate_ffn.tanh() * hidden_state hidden_state = residual + hidden_state - outputs = (hidden_state,) - - if output_attentions: - outputs += (attn_weights,) - - return outputs + return hidden_state class GaudiMllamaVisionEncoder(MllamaVisionEncoder): @@ -261,9 +255,6 @@ def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, use_flash_attention: Optional[bool] = False, ) -> Union[tuple, BaseModelOutput]: """ @@ -271,46 +262,16 @@ def forward( The only differences are: - add use_flash_attention """ - 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 - - encoder_states = () if output_hidden_states else None - all_attentions = () if output_attentions else None - for encoder_layer in self.layers: - if output_hidden_states: - encoder_states = encoder_states + (hidden_states,) - if self.gradient_checkpointing and self.training: - layer_outputs = self._gradient_checkpointing_func( - encoder_layer.__call__, - hidden_states, - attention_mask, - output_attentions, - ) - else: - layer_outputs = encoder_layer( - hidden_state=hidden_states, - attention_mask=attention_mask, - output_attentions=output_attentions, - use_flash_attention=use_flash_attention, - ) + hidden_states = encoder_layer( + hidden_state=hidden_states, + attention_mask=attention_mask, + use_flash_attention=use_flash_attention, + ) - if output_attentions: - all_attentions = all_attentions + (layer_outputs[1],) htcore.mark_step() - hidden_states = layer_outputs[0] - - 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 - ) + return BaseModelOutput(last_hidden_state=hidden_states) class GaudiMllamaTextCrossAttention(MllamaTextCrossAttention): @@ -324,12 +285,12 @@ def forward( cross_attention_states: Optional[torch.Tensor] = None, past_key_value: Optional[Cache] = None, attention_mask: Optional[torch.Tensor] = None, - output_attentions: bool = False, use_cache: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, use_flash_attention: Optional[bool] = False, flash_attention_recompute: Optional[bool] = False, + **kwargs, ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]: """ Copied from MllamaTextCrossAttention::forward: https://github.com/huggingface/transformers/blob/v4.45.2/src/transformers/models/mllama/modeling_mllama.py#L512 @@ -350,9 +311,7 @@ def forward( value_states = self.v_proj(cross_attention_states) key_states = key_states.view(bsz, -1, self.num_key_value_heads, self.head_dim).transpose(1, 2) value_states = value_states.view(bsz, -1, self.num_key_value_heads, self.head_dim).transpose(1, 2) - if not (FusedSDPA and use_flash_attention): - key_states = repeat_kv(key_states, self.num_key_value_groups) - value_states = repeat_kv(value_states, self.num_key_value_groups) + key_states = self.k_norm(key_states) if past_key_value is not None: # if we have a new image + new tokens, we only computed key_states on that new image @@ -376,8 +335,8 @@ def forward( key_states, value_states = (past_key_value[0], past_key_value[1]) elif cache_position is not None and cache_position[0] != 0: key_states, value_states = ( - past_key_value.key_cache[self.layer_idx], - past_key_value.value_cache[self.layer_idx], + past_key_value.layers[self.layer_idx].keys, + past_key_value.layers[self.layer_idx].values, ) else: raise ValueError( @@ -413,9 +372,6 @@ def forward( attn_output = attn_output.reshape(bsz, q_len, -1) attn_output = self.o_proj(attn_output) - if not output_attentions: - attn_weights = None - return attn_output, attn_weights, past_key_value @@ -429,7 +385,6 @@ def forward( hidden_states: torch.Tensor, attention_mask: torch.Tensor, position_embeddings: torch.Tensor, - output_attentions: bool = False, use_cache: bool = False, past_key_value=None, cache_position=None, @@ -512,9 +467,6 @@ def forward( attn_output = self.o_proj(attn_output) - if not output_attentions: - attn_weights = None - return attn_output, attn_weights, past_key_value @@ -534,13 +486,13 @@ def forward( full_text_row_masked_out_mask: Optional[tuple[torch.Tensor, torch.Tensor]] = None, position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Cache] = None, - output_attentions: Optional[bool] = False, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, # will become mandatory in v4.45 token_idx: Optional[torch.Tensor] = None, use_flash_attention: Optional[bool] = False, flash_attention_recompute: Optional[bool] = False, + **kwargs, ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]: """ Copied from MllamaSelfAttentionDecoderLayer::forward: https://github.com/huggingface/transformers/blob/v4.45.2/src/transformers/models/mllama/modeling_mllama.py#L904 @@ -558,13 +510,13 @@ def forward( attention_mask=attention_mask, position_ids=position_ids, past_key_value=past_key_value, - output_attentions=output_attentions, use_cache=use_cache, cache_position=cache_position, position_embeddings=position_embeddings, token_idx=token_idx, use_flash_attention=use_flash_attention, flash_attention_recompute=flash_attention_recompute, + **kwargs, ) hidden_states = residual + hidden_states @@ -576,9 +528,6 @@ def forward( outputs = (hidden_states,) - if output_attentions: - outputs += (self_attn_weights,) - if use_cache: outputs += (present_key_value,) @@ -599,13 +548,13 @@ def forward( full_text_row_masked_out_mask: tuple[torch.Tensor, torch.Tensor], position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Cache] = None, - output_attentions: Optional[bool] = False, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, position_embeddings: Optional[torch.Tensor] = None, token_idx: Optional[torch.Tensor] = None, use_flash_attention: Optional[bool] = False, flash_attention_recompute: Optional[bool] = False, + **kwargs, ) -> tuple[torch.Tensor]: """ Copied from MllamaCrossAttentionDecoderLayer::forward: https://github.com/huggingface/transformers/blob/v4.45.2/src/transformers/models/mllama/modeling_mllama.py#L989 @@ -622,12 +571,12 @@ def forward( attention_mask=cross_attention_mask, cross_attention_states=cross_attention_states, past_key_value=past_key_value, - output_attentions=output_attentions, cache_position=cache_position, use_cache=use_cache, token_idx=token_idx, use_flash_attention=use_flash_attention, flash_attention_recompute=flash_attention_recompute, + **kwargs, ) hidden_states = residual + self.cross_attn_attn_gate.tanh() * hidden_states @@ -640,9 +589,6 @@ def forward( outputs = (hidden_states,) - if output_attentions: - outputs += (attn_weights,) - if use_cache: outputs += (past_key_value,) @@ -661,13 +607,11 @@ def forward( past_key_values: Optional[Union[Cache, list[torch.FloatTensor]]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, use_flash_attention: Optional[bool] = False, flash_attention_recompute: Optional[bool] = False, + **kwargs, ) -> Union[tuple, BaseModelOutputWithPast]: """ Copied from MllamaTextModel::forward: https://github.com/huggingface/transformers/blob/v4.45.2/src/transformers/models/mllama/modeling_mllama.py#L1617 @@ -676,22 +620,11 @@ def forward( - add support if past_key_value is not Cache - add use_flash_attention and flash_attention_recompute """ - 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 - ) use_cache = use_cache if use_cache is not None else self.config.use_cache - return_dict = return_dict if return_dict is not None else self.config.use_return_dict if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") - if self.gradient_checkpointing and self.training and use_cache: - logger.warning_once( - "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`." - ) - use_cache = False - if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) @@ -708,9 +641,7 @@ def forward( ) if position_ids is None: position_ids = cache_position.unsqueeze(0) - causal_mask = self._update_causal_mask( - attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions - ) + causal_mask = self._update_causal_mask(attention_mask, inputs_embeds, cache_position, past_key_values) else: if position_ids is None: position_ids = torch.arange( @@ -732,8 +663,6 @@ def forward( position_embeddings = self.rotary_emb(hidden_states, position_ids) # decoder layers - all_hidden_states = () if output_hidden_states else None - all_self_attns = () if output_attentions else None next_decoder_cache = None if isinstance(past_key_values, Cache) else () for idx, decoder_layer in enumerate(self.layers): @@ -741,8 +670,6 @@ def forward( not torch.distributed.is_initialized() or torch.distributed.get_world_size() == 1 ): htcore.mark_step() - if output_hidden_states: - all_hidden_states += (hidden_states,) # For text-only path we should skip cross attention layers. # Let's check if the layer is cross attention layer and if we have cross attention states @@ -757,69 +684,42 @@ def forward( if is_cross_attention_layer and cross_attention_states is None and is_cross_attention_cache_empty: continue - if self.gradient_checkpointing and self.training: - layer_outputs = self._gradient_checkpointing_func( - decoder_layer.__call__, - hidden_states, - cross_attention_states, - cross_attention_mask, - causal_mask, - full_text_row_masked_out_mask, - position_ids, - past_key_values, - output_attentions, - use_cache, - cache_position, - position_embeddings, - ) + if isinstance(past_key_values, Cache): + past_key_value = past_key_values else: - if isinstance(past_key_values, Cache): - past_key_value = past_key_values - else: - past_key_value = None if past_key_values is None else past_key_values[idx] - layer_outputs = decoder_layer( - hidden_states, - cross_attention_states=cross_attention_states, - cross_attention_mask=cross_attention_mask, - attention_mask=causal_mask, - full_text_row_masked_out_mask=full_text_row_masked_out_mask, - position_ids=position_ids, - past_key_value=past_key_value, - output_attentions=output_attentions, - use_cache=use_cache, - cache_position=cache_position, - position_embeddings=position_embeddings, - token_idx=token_idx, - use_flash_attention=use_flash_attention, - flash_attention_recompute=flash_attention_recompute, - ) + past_key_value = None if past_key_values is None else past_key_values[idx] + layer_outputs = decoder_layer( + hidden_states, + cross_attention_states=cross_attention_states, + cross_attention_mask=cross_attention_mask, + attention_mask=causal_mask, + full_text_row_masked_out_mask=full_text_row_masked_out_mask, + position_ids=position_ids, + past_key_value=past_key_value, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + token_idx=token_idx, + use_flash_attention=use_flash_attention, + flash_attention_recompute=flash_attention_recompute, + **kwargs, + ) hidden_states = layer_outputs[0] if use_cache: if isinstance(past_key_values, Cache): - next_decoder_cache = layer_outputs[2 if output_attentions else 1] + next_decoder_cache = layer_outputs[1] else: - next_decoder_cache += (layer_outputs[2 if output_attentions else 1],) - - if output_attentions: - all_self_attns += (layer_outputs[1],) + next_decoder_cache += (layer_outputs[1],) hidden_states = self.norm(hidden_states) - # add hidden states from the last decoder layer - if output_hidden_states: - all_hidden_states += (hidden_states,) - next_cache = next_decoder_cache if use_cache else None - if not return_dict: - return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None) return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=next_cache, - hidden_states=all_hidden_states, - attentions=all_self_attns, ) def _update_causal_mask( @@ -828,7 +728,7 @@ def _update_causal_mask( input_tensor: torch.Tensor, cache_position: torch.Tensor, past_key_values: Cache, - output_attentions: bool, + output_attentions: bool = False, ): """ Copied from MllamaTextModel::_update_causal_mask: https://github.com/huggingface/transformers/blob/v4.45.2/src/transformers/models/mllama/modeling_mllama.py#L1768 @@ -908,9 +808,6 @@ def forward( inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, logits_to_keep: Union[int, torch.Tensor] = 0, token_idx: Optional[torch.Tensor] = None, @@ -918,7 +815,7 @@ def forward( use_flash_attention: Optional[bool] = False, flash_attention_recompute: Optional[bool] = False, logits_bf16: Optional[bool] = False, - **loss_kwargs, + **kwargs, ) -> Union[tuple, CausalLMOutputWithPast]: """ Copied from MllamaForCausalLM::forward: https://github.com/huggingface/transformers/blob/v4.45.2/src/transformers/models/mllama/modeling_mllama.py#L1871 @@ -927,12 +824,6 @@ def forward( - add logits handle if token_idx is not None - add use_flash_attention and flash_attention_recompute """ - 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 - # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) outputs = self.model( input_ids=input_ids, @@ -944,16 +835,14 @@ def forward( past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, cache_position=cache_position, token_idx=token_idx, use_flash_attention=use_flash_attention, flash_attention_recompute=flash_attention_recompute, + **kwargs, ) - hidden_states = outputs[0] + hidden_states = outputs.last_hidden_state _, seq_len, _ = hidden_states.shape if seq_len > 1 and trim_logits and not self.training: if token_idx is not None: @@ -972,11 +861,7 @@ def forward( loss = None if labels is not None: - loss = self.loss_function(logits, labels, self.vocab_size, **loss_kwargs) - - if not return_dict: - output = (logits,) + outputs[1:] - return (loss,) + output if loss is not None else output + loss = self.loss_function(logits, labels, self.vocab_size, **kwargs) return CausalLMOutputWithPast( loss=loss, @@ -1007,9 +892,6 @@ def forward( inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, logits_to_keep: Union[int, torch.Tensor] = 0, token_idx: Optional[torch.Tensor] = None, @@ -1017,7 +899,7 @@ def forward( use_flash_attention: Optional[bool] = False, flash_attention_recompute: Optional[bool] = False, logits_bf16: Optional[bool] = False, - **loss_kwargs, + **kwargs, ) -> Union[tuple, CausalLMOutputWithPast]: """ Copied from MllamaForConditionalGeneration::forward: https://github.com/huggingface/transformers/blob/v4.45.2/src/transformers/models/mllama/modeling_mllama.py#L2077 @@ -1025,12 +907,6 @@ def forward( - add token_idx input - add use_flash_attention and flash_attention_recompute """ - 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 - if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") @@ -1050,9 +926,6 @@ def forward( pixel_values=pixel_values, aspect_ratio_ids=aspect_ratio_ids, aspect_ratio_mask=aspect_ratio_mask, - output_hidden_states=output_hidden_states, - output_attentions=output_attentions, - return_dict=return_dict, use_flash_attention=use_flash_attention, ) cross_attention_states = vision_outputs[0] @@ -1095,9 +968,6 @@ def forward( use_cache=use_cache, inputs_embeds=inputs_embeds, labels=labels, - output_hidden_states=output_hidden_states, - output_attentions=output_attentions, - return_dict=return_dict, cache_position=cache_position, logits_to_keep=logits_to_keep, token_idx=token_idx, @@ -1105,7 +975,7 @@ def forward( use_flash_attention=use_flash_attention, flash_attention_recompute=flash_attention_recompute, logits_bf16=logits_bf16, - **loss_kwargs, + **kwargs, ) return outputs @@ -1238,22 +1108,14 @@ def forward( pixel_values: torch.Tensor, aspect_ratio_ids: torch.Tensor, aspect_ratio_mask: torch.Tensor, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, use_flash_attention: Optional[bool] = False, + **kwargs, ) -> Union[BaseModelOutput, tuple[torch.Tensor, ...]]: """ Copied from MllamaVisionModel::forward: https://github.com/huggingface/transformers/blob/v4.45.2/src/transformers/models/mllama/modeling_mllama.py#L1425 The only differences are: - optimize perf of stage "Collect intermediate layer outputs from encoder output" """ - 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 - batch_size, num_concurrent_media, num_tiles, num_channels, height, width = pixel_values.shape pixel_values = pixel_values.reshape(batch_size * num_concurrent_media * num_tiles, num_channels, height, width) @@ -1303,11 +1165,9 @@ def forward( output = self.transformer( hidden_state, attention_mask=attention_mask, - output_hidden_states=True, - output_attentions=output_attentions, use_flash_attention=use_flash_attention, ) - hidden_state = output[0] + hidden_state = output.last_hidden_state hidden_state = self.layernorm_post(hidden_state) @@ -1322,11 +1182,9 @@ def forward( global_output = self.global_transformer( hidden_state, attention_mask=attention_mask, - output_hidden_states=output_hidden_states, - output_attentions=output_attentions, use_flash_attention=use_flash_attention, ) - hidden_state = global_output[0] + hidden_state = global_output.last_hidden_state # Remove padding form hidden state hidden_state = hidden_state.reshape( @@ -1336,7 +1194,7 @@ def forward( hidden_state = hidden_state.reshape(batch_size, num_concurrent_media, num_tiles, num_patches, dim) # Collect intermediate layer outputs from encoder output - all_intermediate_hidden_states = [output[1][i] for i in self.intermediate_layers_indices] + all_intermediate_hidden_states = [output.last_hidden_state for _ in self.intermediate_layers_indices] intermediate_hidden_states = torch.stack(all_intermediate_hidden_states, dim=-1) """ @@ -1356,23 +1214,4 @@ def forward( # Concatenate final hidden state and intermediate hidden states hidden_state = torch.cat([hidden_state, intermediate_hidden_states], dim=-1) - if output_hidden_states: - hidden_states = tuple(all_intermediate_hidden_states) + tuple(global_output[1]) - else: - hidden_states = None - - if output_attentions: - # global transformer in contrast to `self.transformer` doesn't always return hidden states so we might go index out-of-range - global_attn = tuple(global_output[2]) if output_hidden_states else tuple(global_output[1]) - attentions = tuple(output[2]) + global_attn - else: - attentions = None - - if not return_dict: - return tuple(v for v in [hidden_state, hidden_states, attentions] if v is not None) - - return BaseModelOutput( - last_hidden_state=hidden_state, - hidden_states=hidden_states, - attentions=attentions, - ) + return BaseModelOutput(last_hidden_state=hidden_state) diff --git a/optimum/habana/transformers/models/opt/modeling_opt.py b/optimum/habana/transformers/models/opt/modeling_opt.py index 408953efca..045a10f44a 100644 --- a/optimum/habana/transformers/models/opt/modeling_opt.py +++ b/optimum/habana/transformers/models/opt/modeling_opt.py @@ -53,9 +53,8 @@ def gaudi_opt_attention_forward( attention_mask: Optional[torch.Tensor] = None, layer_head_mask: Optional[torch.Tensor] = None, output_attentions: bool = False, - # isn't needed in normal attention, but needed in flash attention so to keep the signature same - position_ids: Optional[torch.Tensor] = None, token_idx: Optional[torch.Tensor] = None, + **kwargs, ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]: """ Copied from OPTAttention.forward: https://github.com/huggingface/transformers/blob/main/src/transformers/models/opt/modeling_opt.py @@ -69,7 +68,11 @@ def gaudi_opt_attention_forward( bsz, tgt_len, _ = hidden_states.size() - # get query proj + # Scaling is susceptible to floating point arithmetics' inprecisions + # which can lead to different results (this is dependent from model + # to model, e.g. whisper is one such case). We therefore keep the + # original order of scaling to follow the original implementation + # and enforce no scaling (1.0) in the attention call below. query_states = self.q_proj(hidden_states) * self.scaling # get key, value proj if is_cross_attention and past_key_value is not None: @@ -104,66 +107,21 @@ def gaudi_opt_attention_forward( key_states = key_states.view(*proj_shape) value_states = value_states.view(*proj_shape) - src_len = key_states.size(1) - attn_weights = torch.bmm(query_states, key_states.transpose(1, 2)) - - if attn_weights.size() != (bsz * self.num_heads, tgt_len, src_len): - raise ValueError( - f"Attention weights should be of size {(bsz * self.num_heads, tgt_len, src_len)}, but is" - f" {attn_weights.size()}" - ) - + attn_weights = torch.bmm(query_states, key_states.transpose(1, 2)) * self.scaling if attention_mask is not None: - if attention_mask.size() != (bsz, 1, tgt_len, src_len): - raise ValueError( - f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is {attention_mask.size()}" - ) - attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + attention_mask - attn_weights = torch.max( - attn_weights, torch.tensor(torch.finfo(attn_weights.dtype).min, device=attn_weights.device) - ) - attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len) - - attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1) + attn_weights = attn_weights + attention_mask - 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 {layer_head_mask.size()}" - ) - attn_weights = layer_head_mask.view(1, -1, 1, 1) * attn_weights.view(bsz, self.num_heads, tgt_len, src_len) - attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len) - - if output_attentions: - # this operation is a bit awkward, but it's required to - # make sure that attn_weights keeps its gradient. - # In order to do so, attn_weights have to be reshaped - # twice and have to be reused in the following - attn_weights_reshaped = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) - attn_weights = attn_weights_reshaped.view(bsz * self.num_heads, tgt_len, src_len) - else: - attn_weights_reshaped = None - - attn_probs = torch.nn.functional.dropout(attn_weights, p=self.dropout, training=self.training) - - attn_output = torch.bmm(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_weights = torch.nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) + attn_weights = torch.nn.functional.dropout(attn_weights, p=self.dropout, training=self.training) + attn_output = torch.bmm(attn_weights, value_states) attn_output = attn_output.view(bsz, self.num_heads, tgt_len, self.head_dim) - 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 aross GPUs when using tensor-parallelism. - attn_output = attn_output.reshape(bsz, tgt_len, self.embed_dim) + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.reshape(bsz, tgt_len, -1).contiguous() attn_output = self.out_proj(attn_output) - return attn_output, attn_weights_reshaped, past_key_value + return attn_output, attn_weights, past_key_value class GaudiOPTDecoderLayer(torch.nn.Module): @@ -199,6 +157,7 @@ def forward( use_cache: Optional[bool] = False, position_ids: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, + **kwargs, ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]: """ Copied from OPTDecoderLayer.forward: https://github.com/huggingface/transformers/blob/main/src/transformers/models/opt/modeling_opt.py @@ -220,6 +179,7 @@ def forward( layer_head_mask=layer_head_mask, output_attentions=output_attentions, token_idx=token_idx, + **kwargs, ) hidden_states = torch.nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) hidden_states = residual + hidden_states @@ -273,6 +233,7 @@ def gaudi_opt_decoder_forward( return_dict: Optional[bool] = None, position_ids: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, + **kwargs, ) -> Union[tuple, BaseModelOutputWithPast]: """ Copied from OPTDecoder.forward: https://github.com/huggingface/transformers/blob/main/src/transformers/models/opt/modeling_opt.py @@ -352,7 +313,7 @@ def gaudi_opt_decoder_forward( ) for idx, decoder_layer in enumerate(self.layers): - # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description) + # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description) if output_hidden_states: all_hidden_states += (hidden_states,) @@ -363,29 +324,17 @@ def gaudi_opt_decoder_forward( past_key_value = past_key_values[idx] if past_key_values is not None else None - if self.gradient_checkpointing and self.training: - layer_outputs = self._gradient_checkpointing_func( - decoder_layer.__call__, - hidden_states, - causal_attention_mask, - head_mask[idx] if head_mask is not None else None, - None, - output_attentions, - use_cache, - position_ids, - None, - ) - else: - layer_outputs = decoder_layer( - hidden_states, - attention_mask=causal_attention_mask, - position_ids=position_ids, - layer_head_mask=(head_mask[idx] if head_mask is not None else None), - past_key_value=past_key_value, - output_attentions=output_attentions, - use_cache=use_cache, - token_idx=token_idx, - ) + layer_outputs = decoder_layer( + hidden_states, + attention_mask=causal_attention_mask, + position_ids=position_ids, + layer_head_mask=(head_mask[idx] if head_mask is not None else None), + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + token_idx=token_idx, + **kwargs, + ) hidden_states = layer_outputs[0] @@ -406,8 +355,6 @@ def gaudi_opt_decoder_forward( all_hidden_states += (hidden_states,) next_cache = next_decoder_cache if use_cache else None - if not return_dict: - return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None) return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=next_cache, @@ -429,6 +376,7 @@ def gaudi_opt_model_forward( return_dict: Optional[bool] = None, position_ids: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, + **kwargs, ) -> Union[tuple, BaseModelOutputWithPast]: """ Copied from OPTModel.forward: https://github.com/huggingface/transformers/blob/main/src/transformers/models/opt/modeling_opt.py @@ -453,13 +401,11 @@ def gaudi_opt_model_forward( use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, - return_dict=return_dict, + return_dict=True, token_idx=token_idx, + **kwargs, ) - if not return_dict: - return decoder_outputs - return BaseModelOutputWithPast( last_hidden_state=decoder_outputs.last_hidden_state, past_key_values=decoder_outputs.past_key_values, @@ -510,8 +456,9 @@ def forward( use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, - return_dict=return_dict, + return_dict=True, token_idx=token_idx, + **kwargs, ) logits = self.lm_head(outputs[0]).contiguous() @@ -527,10 +474,6 @@ def forward( **kwargs, ) - if not return_dict: - output = (logits,) + outputs[1:] - return (loss,) + output if loss is not None else output - return CausalLMOutputWithPast( loss=loss, logits=logits, diff --git a/optimum/habana/transformers/models/paligemma/modeling_paligemma.py b/optimum/habana/transformers/models/paligemma/modeling_paligemma.py index 31b7e35be0..187460dc5d 100644 --- a/optimum/habana/transformers/models/paligemma/modeling_paligemma.py +++ b/optimum/habana/transformers/models/paligemma/modeling_paligemma.py @@ -14,11 +14,10 @@ # limitations under the License. """PyTorch Paligemma model.""" -from typing import List, Optional, Tuple, Union +from typing import Optional, Union import torch import torch.utils.checkpoint -from torch import nn from transformers.cache_utils import Cache from transformers.modeling_outputs import CausalLMOutputWithPast from transformers.models.paligemma.modeling_paligemma import ( @@ -34,11 +33,11 @@ class GaudiPaliGemmaForConditionalGeneration(PaliGemmaForConditionalGeneration): def forward( self, - input_ids: Optional[torch.LongTensor] = None, - pixel_values: Optional[torch.FloatTensor] = None, + input_ids: torch.LongTensor = None, + pixel_values: torch.FloatTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[Union[List[torch.FloatTensor], Cache]] = None, + past_key_values: Optional[Union[list[torch.FloatTensor], Cache]] = None, token_type_ids: Optional[torch.LongTensor] = None, cache_position: Optional[torch.LongTensor] = None, inputs_embeds: Optional[torch.FloatTensor] = None, @@ -49,17 +48,13 @@ def forward( return_dict: Optional[bool] = None, logits_to_keep: Union[int, torch.Tensor] = 0, token_idx: Optional[torch.Tensor] = None, - **lm_kwargs, - ) -> Union[Tuple, PaliGemmaCausalLMOutputWithPast]: + **kwargs, + ) -> Union[tuple, PaliGemmaCausalLMOutputWithPast]: """ Inherits from PaliGemmaForConditionalGeneration::forward https://github.com/huggingface/transformers/blob/v4.45.1/src/transformers/models/paligemma/modeling_paligemma.py#L402 The only differences are: - add new args token_idx """ - - if (input_ids is None) ^ (inputs_embeds is not None): - raise ValueError("You must specify exactly one of input_ids or inputs_embeds") - 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 @@ -134,38 +129,25 @@ def forward( # TODO: from Transformers v4.45, `generate` sets `num_logits_to_keep` to 1 if not given, which we don't want here # logits_to_keep=logits_to_keep, token_idx=token_idx, - **lm_kwargs, + **kwargs, ) - logits = outputs[0] + hidden_states = outputs[0] + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + logits = self.lm_head(hidden_states[:, slice_indices, :]) + loss = None if labels is not None: - # Upcast to float if we need to compute the loss to avoid potential precision issues - logits = logits.float() - shift_logits = logits[..., :-1, :] - shift_labels = labels[..., 1:] - if attention_mask is not None: - # we use the input attention mask to shift the logits and labels, because it is 2D. - # we also crop attn mask in case it is longer, which happens in PrefixTuning with peft - shift_attention_mask = attention_mask[:, -shift_logits.shape[1] :].to(logits.device) - shift_logits = shift_logits[shift_attention_mask.to(logits.device) != 0].contiguous() - shift_labels = shift_labels[shift_attention_mask.to(shift_labels.device) != 0].contiguous() - else: - shift_logits = shift_logits.contiguous() - shift_labels = shift_labels.contiguous() - # Flatten the tokens - loss_fct = nn.CrossEntropyLoss() - - flat_logits = shift_logits.view(-1, self.config.text_config.vocab_size) - flat_labels = shift_labels.view(-1).to(shift_logits.device) - loss = loss_fct(flat_logits, flat_labels) + loss = self.loss_function( + logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size, **kwargs + ) - output = PaliGemmaCausalLMOutputWithPast( + return PaliGemmaCausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, - image_hidden_states=image_features if pixel_values is not None else None, + image_hidden_states=outputs.image_hidden_states, ) - return output if return_dict else output.to_tuple() diff --git a/optimum/habana/transformers/models/persimmon/modeling_persimmon.py b/optimum/habana/transformers/models/persimmon/modeling_persimmon.py index 408600280d..4209d0c94c 100644 --- a/optimum/habana/transformers/models/persimmon/modeling_persimmon.py +++ b/optimum/habana/transformers/models/persimmon/modeling_persimmon.py @@ -1,8 +1,6 @@ -import math -from typing import List, Optional, Tuple, Union +from typing import Optional, Union import torch -from torch import nn from transformers.cache_utils import Cache, DynamicCache from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from transformers.models.persimmon.configuration_persimmon import PersimmonConfig @@ -11,6 +9,7 @@ PersimmonDecoderLayer, PersimmonForCausalLM, apply_rotary_pos_emb, + eager_attention_forward, ) from transformers.utils import logging @@ -39,9 +38,10 @@ def forward( output_attentions: bool = False, use_cache: bool = False, cache_position: Optional[torch.LongTensor] = None, - position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC + position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC token_idx: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + **kwargs, + ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]: """ Copied from PersimmonAttention.forward: https://github.com/huggingface/transformers/blob/v4.38.2/src/transformers/models/persimmon/modeling_persimmon.py The only differences are: @@ -118,33 +118,18 @@ def forward( key_states, value_states, self.layer_idx, cache_kwargs ) - attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim) - - if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len): - raise ValueError( - f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is" - f" {attn_weights.size()}" - ) - - 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 - - # upcast attention to fp32 - attn_weights = nn.functional.softmax(attn_weights, dtype=torch.float32, dim=-1).to(query_states.dtype) - attn_weights = self.attention_dropout(attn_weights) - - attn_output = torch.matmul(attn_weights, value_states) - - if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim): - raise ValueError( - f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is" - f" {attn_output.size()}" - ) - - attn_output = attn_output.transpose(1, 2).contiguous() - attn_output = attn_output.reshape(bsz, q_len, self.hidden_size) + attn_output, attn_weights = eager_attention_forward( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.config.attention_dropout, + scaling=self.scaling, + **kwargs, + ) + attn_output = attn_output.reshape(bsz, q_len, -1) attn_output = self.dense(attn_output) if not output_attentions: @@ -163,13 +148,14 @@ def forward( hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_value: Optional[Tuple[torch.Tensor]] = None, + past_key_value: Optional[tuple[torch.Tensor]] = None, output_attentions: Optional[bool] = False, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, - position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC + position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC token_idx: Optional[torch.Tensor] = None, - ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + **kwargs, + ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]: """ Copied from PersimmonDecoderLayer.forward: https://github.com/huggingface/transformers/blob/v4.38.2/src/transformers/models/persimmon/modeling_persimmon.py The only differences are: @@ -190,6 +176,7 @@ def forward( cache_position=cache_position, position_embeddings=position_embeddings, token_idx=token_idx, + **kwargs, ) hidden_states = residual + hidden_states @@ -217,13 +204,14 @@ def gaudi_persimmon_model_forward( input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, + past_key_values: Optional[list[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, + **kwargs, ) -> BaseModelOutputWithPast: """ Copied from PersimmonModel.forward: https://github.com/huggingface/transformers/blob/v4.38.2/src/transformers/models/persimmon/modeling_persimmon.py @@ -294,28 +282,17 @@ def gaudi_persimmon_model_forward( if output_hidden_states: all_hidden_states += (hidden_states,) - if self.gradient_checkpointing and self.training: - layer_outputs = self._gradient_checkpointing_func( - decoder_layer.__call__, - hidden_states, - attention_mask, - position_ids, - past_key_values, - output_attentions, - use_cache, - cache_position, - ) - else: - layer_outputs = decoder_layer( - hidden_states, - attention_mask=attention_mask, - position_ids=position_ids, - past_key_value=past_key_values, - output_attentions=output_attentions, - use_cache=use_cache, - cache_position=cache_position, - token_idx=token_idx, - ) + layer_outputs = decoder_layer( + hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + token_idx=token_idx, + **kwargs, + ) hidden_states = layer_outputs[0] @@ -349,7 +326,7 @@ def forward( input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, + past_key_values: Optional[list[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, @@ -383,6 +360,7 @@ def forward( output_hidden_states=output_hidden_states, cache_position=cache_position, token_idx=token_idx, + **kwargs, ) hidden_states = outputs.last_hidden_state diff --git a/optimum/habana/transformers/models/qwen3_moe/modeling_qwen3_moe.py b/optimum/habana/transformers/models/qwen3_moe/modeling_qwen3_moe.py index 8dee09080f..55290530d1 100644 --- a/optimum/habana/transformers/models/qwen3_moe/modeling_qwen3_moe.py +++ b/optimum/habana/transformers/models/qwen3_moe/modeling_qwen3_moe.py @@ -921,7 +921,9 @@ def forward( past_seen_tokens, ) else: - mask_function = create_causal_mask if self.config.sliding_window is None else create_sliding_window_causal_mask + mask_function = ( + create_causal_mask if self.config.sliding_window is None else create_sliding_window_causal_mask + ) causal_mask = mask_function( config=self.config, input_embeds=inputs_embeds, diff --git a/optimum/habana/transformers/models/t5/modeling_t5.py b/optimum/habana/transformers/models/t5/modeling_t5.py index 3c92a336dd..1e05eeec04 100644 --- a/optimum/habana/transformers/models/t5/modeling_t5.py +++ b/optimum/habana/transformers/models/t5/modeling_t5.py @@ -586,7 +586,6 @@ def forward( encoder_attentions=encoder_outputs.attentions, ) - def prepare_inputs_for_generation( self, input_ids, diff --git a/optimum/habana/transformers/models/wav2vec2/modeling_wav2vec2.py b/optimum/habana/transformers/models/wav2vec2/modeling_wav2vec2.py index 58da36dde2..e6d62d466d 100644 --- a/optimum/habana/transformers/models/wav2vec2/modeling_wav2vec2.py +++ b/optimum/habana/transformers/models/wav2vec2/modeling_wav2vec2.py @@ -242,9 +242,7 @@ def gaudi_wav2vec2_encoder_forward( skip_the_layer = self.training and (dropout_probability < self.config.layerdrop) if not skip_the_layer or synced_gpus: # under fsdp or deepspeed zero3 all gpus must run in sync - layer_outputs = layer( - hidden_states, attention_mask=attention_mask, output_attentions=output_attentions - ) + layer_outputs = layer(hidden_states, attention_mask=attention_mask, output_attentions=output_attentions) hidden_states = layer_outputs[0] if skip_the_layer: diff --git a/optimum/habana/transformers/models/whisper/modeling_whisper.py b/optimum/habana/transformers/models/whisper/modeling_whisper.py index 0b586aa78e..848466fde8 100644 --- a/optimum/habana/transformers/models/whisper/modeling_whisper.py +++ b/optimum/habana/transformers/models/whisper/modeling_whisper.py @@ -1,10 +1,11 @@ -from typing import Optional, Tuple, Union +from typing import Optional, Union import torch import torch.utils.checkpoint from torch import nn from torch.nn import CrossEntropyLoss from transformers.cache_utils import Cache, DynamicCache, EncoderDecoderCache +from transformers.masking_utils import create_causal_mask from transformers.modeling_outputs import ( BaseModelOutput, BaseModelOutputWithPastAndCrossAttentions, @@ -36,7 +37,7 @@ def forward( output_attentions: bool = False, cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]: """ Inherits from WhisperDecoderLayer: https://github.com/huggingface/transformers/blob/v4.44.0/src/transformers/models/whisper/modeling_whisper.py The only differences are: @@ -304,12 +305,13 @@ def forward( hidden_states = inputs_embeds + positions.to(inputs_embeds.device) hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) - causal_mask = self._update_causal_mask( - attention_mask, - inputs_embeds, - cache_position, - past_key_values.self_attention_cache if past_key_values is not None else None, - output_attentions, + causal_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, + position_ids=position_ids, ) if self.gradient_checkpointing and self.training: @@ -331,7 +333,7 @@ def forward( f" {head_mask.size()[0]}." ) for idx, decoder_layer in enumerate(self.layers): - # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description) + # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description) if output_hidden_states: all_hidden_states += (hidden_states,) if self.training: @@ -339,35 +341,18 @@ def forward( if dropout_probability < self.layerdrop: continue - if self.gradient_checkpointing and self.training: - layer_outputs = self._gradient_checkpointing_func( - decoder_layer.__call__, - hidden_states, - causal_mask, - encoder_hidden_states, - None, # encoder attention mask - head_mask[idx] if head_mask is not None else None, - cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None, - None, # past_key_value - output_attentions, - use_cache, - cache_position, - ) - else: - layer_outputs = decoder_layer( - hidden_states, - attention_mask=causal_mask, - encoder_hidden_states=encoder_hidden_states, - layer_head_mask=(head_mask[idx] if head_mask is not None else None), - cross_attn_layer_head_mask=( - cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None - ), - past_key_value=past_key_values if use_cache else None, - output_attentions=output_attentions, - use_cache=use_cache, - cache_position=cache_position, - token_idx=token_idx, - ) + layer_outputs = decoder_layer( + hidden_states, + attention_mask=causal_mask, + encoder_hidden_states=encoder_hidden_states, + layer_head_mask=(head_mask[idx] if head_mask is not None else None), + cross_attn_layer_head_mask=(cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None), + past_key_value=past_key_values if use_cache else None, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + token_idx=token_idx, + ) hidden_states = layer_outputs[0] if output_attentions: @@ -411,17 +396,17 @@ def forward( head_mask: Optional[torch.Tensor] = None, decoder_head_mask: Optional[torch.Tensor] = None, cross_attn_head_mask: Optional[torch.Tensor] = None, - encoder_outputs: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, - past_key_values: Optional[Union[EncoderDecoderCache, Tuple[torch.FloatTensor]]] = None, - decoder_inputs_embeds: Optional[Tuple[torch.FloatTensor]] = None, - decoder_position_ids: Optional[Tuple[torch.LongTensor]] = None, + encoder_outputs: Optional[tuple[tuple[torch.FloatTensor]]] = None, + past_key_values: Optional[Union[EncoderDecoderCache, tuple[torch.FloatTensor]]] = None, + decoder_inputs_embeds: Optional[tuple[torch.FloatTensor]] = None, + decoder_position_ids: Optional[tuple[torch.LongTensor]] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, - ) -> Union[Tuple[torch.Tensor], Seq2SeqModelOutput]: + ) -> Union[tuple[torch.Tensor], Seq2SeqModelOutput]: """ Inherits from WhisperModel: https://github.com/huggingface/transformers/blob/v4.44.0/src/transformers/models/whisper/modeling_whisper.py The only differences are: @@ -495,10 +480,10 @@ def forward( head_mask: Optional[torch.Tensor] = None, decoder_head_mask: Optional[torch.Tensor] = None, cross_attn_head_mask: Optional[torch.Tensor] = None, - encoder_outputs: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, - past_key_values: Optional[Union[EncoderDecoderCache, Tuple[torch.FloatTensor]]] = None, - decoder_inputs_embeds: Optional[Tuple[torch.FloatTensor]] = None, - decoder_position_ids: Optional[Tuple[torch.LongTensor]] = None, + encoder_outputs: Optional[tuple[tuple[torch.FloatTensor]]] = None, + past_key_values: Optional[Union[EncoderDecoderCache, tuple[torch.FloatTensor]]] = None, + decoder_inputs_embeds: Optional[tuple[torch.FloatTensor]] = None, + decoder_position_ids: Optional[tuple[torch.LongTensor]] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, @@ -506,7 +491,7 @@ def forward( return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, - ) -> Union[Tuple[torch.Tensor], Seq2SeqLMOutput]: + ) -> Union[tuple[torch.Tensor], Seq2SeqLMOutput]: """ Inherits from WhisperForConditionalGeneration: https://github.com/huggingface/transformers/blob/v4.44.0/src/transformers/models/whisper/modeling_whisper.py The only differences are: diff --git a/optimum/habana/transformers/trainer.py b/optimum/habana/transformers/trainer.py index 768b3e45ce..33889f2ced 100644 --- a/optimum/habana/transformers/trainer.py +++ b/optimum/habana/transformers/trainer.py @@ -1686,7 +1686,7 @@ def training_step( if self.args.use_lazy_mode and self.args.pipelining_fwd_bwd: self.htcore.mark_step() - # Finally we need to normalize the loss for reporting if GA loss bug is not fixed during compute loss + # Finally we need to normalize the loss for reporting if GA loss bug is not fixed during compute loss if (not self.model_accepts_loss_kwargs or num_items_in_batch is None) and self.compute_loss_func is None: # If the model does not accept loss kwargs, we need to normalize the loss by the number of gradient accumulation steps loss = loss / self.current_gradient_accumulation_steps diff --git a/tests/test_trainer.py b/tests/test_trainer.py index 2c64d0ac31..c09cfbcc2d 100644 --- a/tests/test_trainer.py +++ b/tests/test_trainer.py @@ -2676,7 +2676,9 @@ def wrapped_fn(epoch_iterator, num_batches, device): return wrapped_fn - trainer.get_batch_samples_transformers = wrap_get_batch_samples(trainer.get_batch_samples_transformers) + trainer.get_batch_samples_transformers = wrap_get_batch_samples( + trainer.get_batch_samples_transformers + ) trainer.train() From 41fc61caffd7784bdb958253e712c7d353eb03ab Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Tue, 19 Aug 2025 08:43:52 +0000 Subject: [PATCH 22/22] Upgrade to v4.55.2 --- .../models/idefics2/modeling_idefics2.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/optimum/habana/transformers/models/idefics2/modeling_idefics2.py b/optimum/habana/transformers/models/idefics2/modeling_idefics2.py index bef37c7470..040e62e4bd 100644 --- a/optimum/habana/transformers/models/idefics2/modeling_idefics2.py +++ b/optimum/habana/transformers/models/idefics2/modeling_idefics2.py @@ -47,19 +47,19 @@ def forward(self, pixel_values: torch.FloatTensor, patch_attention_mask: torch.B embeddings = patch_embeds.flatten(2).transpose(1, 2) max_nb_patches_h, max_nb_patches_w = max_im_h // self.patch_size, max_im_w // self.patch_size - boundaries = torch.arange(1 / self.num_patches_per_side, 1.0, 1 / self.num_patches_per_side) + boundaries = torch.arange( + 1 / self.num_patches_per_side, 1.0, 1 / self.num_patches_per_side, device=pixel_values.device + ) position_ids = torch.full( - size=(batch_size, max_nb_patches_h * max_nb_patches_w), - fill_value=0, - device=self.position_embedding.weight.device, + size=(batch_size, max_nb_patches_h * max_nb_patches_w), fill_value=0, device=pixel_values.device ) for batch_idx, p_attn_mask in enumerate(patch_attention_mask): nb_patches_h = int(p_attn_mask[:, 0].sum()) nb_patches_w = int(p_attn_mask[0].sum()) - h_indices = torch.arange(nb_patches_h, device=pixel_values.device, dtype=pixel_values.dtype) - w_indices = torch.arange(nb_patches_w, device=pixel_values.device, dtype=pixel_values.dtype) + h_indices = torch.arange(nb_patches_h, device=position_ids.device, dtype=position_ids.dtype) + w_indices = torch.arange(nb_patches_w, device=position_ids.device, dtype=position_ids.dtype) fractional_coords_h = h_indices / nb_patches_h * (1 - 1e-6) fractional_coords_w = w_indices / nb_patches_w * (1 - 1e-6) @@ -69,6 +69,7 @@ def forward(self, pixel_values: torch.FloatTensor, patch_attention_mask: torch.B pos_ids = (bucket_coords_h[:, None] * self.num_patches_per_side + bucket_coords_w).flatten() position_ids[batch_idx][p_attn_mask.view(-1)] = pos_ids.to(self.position_embedding.weight.device) + embeddings = embeddings + self.position_embedding(position_ids) return embeddings