From 5e88ce7e91244b928f01b898d2ba6e87cdbff9e1 Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Wed, 28 Aug 2024 16:09:17 +0000 Subject: [PATCH 01/89] Upgrade to commit 74e19e81e2a23809af192532b9b0e7ea202be6f2 --- .../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 +- .../run_multitask_prompt_tuning.py | 2 +- .../run_prompt_tuning_clm.py | 2 +- examples/question-answering/run_qa.py | 2 +- examples/question-answering/run_seq2seq_qa.py | 2 +- .../run_speech_recognition_ctc.py | 2 +- .../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 +- .../transformers/generation/__init__.py | 1 - .../generation/stopping_criteria.py | 12 - .../habana/transformers/generation/utils.py | 306 ++++++------------ optimum/habana/transformers/modeling_utils.py | 2 - .../models/bloom/modeling_bloom.py | 38 ++- .../models/codegen/modeling_codegen.py | 46 ++- .../models/falcon/modeling_falcon.py | 37 ++- .../models/gemma/modeling_gemma.py | 21 +- .../models/gpt_neox/modeling_gpt_neox.py | 46 ++- .../transformers/models/gptj/modeling_gptj.py | 50 ++- .../models/llama/configuration_llama.py | 1 + .../models/llama/modeling_llama.py | 19 +- .../models/llava/modeling_llava.py | 11 +- .../models/llava_next/modeling_llava_next.py | 3 + .../models/mamba/modeling_mamba.py | 15 +- .../models/mistral/modeling_mistral.py | 18 +- .../models/mixtral/modeling_mixtral.py | 17 +- .../models/persimmon/modeling_persimmon.py | 12 +- .../transformers/models/phi/modeling_phi.py | 23 +- .../models/qwen2/modeling_qwen2.py | 20 +- .../models/stablelm/modeling_stablelm.py | 12 +- .../models/starcoder2/modeling_starcoder2.py | 21 +- optimum/habana/transformers/trainer.py | 35 +- .../habana/transformers/trainer_seq2seq.py | 2 +- optimum/habana/transformers/training_args.py | 6 +- setup.py | 2 +- .../generation/test_stopping_criteria.py | 16 - 43 files changed, 452 insertions(+), 372 deletions(-) diff --git a/examples/audio-classification/run_audio_classification.py b/examples/audio-classification/run_audio_classification.py index 86dc6627dd..bb5754f6a9 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.43.0") +check_min_version("4.45.0.dev0") check_optimum_habana_min_version("1.13.0") 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 c22682203e..7a9e92a640 100644 --- a/examples/contrastive-image-text/run_bridgetower.py +++ b/examples/contrastive-image-text/run_bridgetower.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.43.0") +check_min_version("4.45.0.dev0") check_optimum_habana_min_version("1.13.0") 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 2358412de6..f55cb1b241 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.43.0") +check_min_version("4.45.0.dev0") check_optimum_habana_min_version("1.13.0") 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 7bd1d23c4d..7a91b88317 100644 --- a/examples/image-classification/run_image_classification.py +++ b/examples/image-classification/run_image_classification.py @@ -63,7 +63,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.43.0") +check_min_version("4.45.0.dev0") check_optimum_habana_min_version("1.13.0") 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 3e372d17a6..c53110a2f7 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.43.0") +check_min_version("4.45.0.dev0") check_optimum_habana_min_version("1.13.0") 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 7fb0ce8494..c790df437e 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.43.0") +check_min_version("4.45.0.dev0") check_optimum_habana_min_version("1.13.0") require_version("datasets>=2.14.0", "To fix: pip install -r examples/pytorch/language-modeling/requirements.txt") diff --git a/examples/language-modeling/run_multitask_prompt_tuning.py b/examples/language-modeling/run_multitask_prompt_tuning.py index 9f7d10655c..5ae9667be4 100644 --- a/examples/language-modeling/run_multitask_prompt_tuning.py +++ b/examples/language-modeling/run_multitask_prompt_tuning.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 risk. -check_min_version("4.38.0") +check_min_version("4.45.0.dev0") check_optimum_habana_min_version("1.10.0") 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 42798c0d5e..11b784ac64 100644 --- a/examples/language-modeling/run_prompt_tuning_clm.py +++ b/examples/language-modeling/run_prompt_tuning_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.38.0") +check_min_version("4.45.0.dev0") check_optimum_habana_min_version("1.10.0") require_version("datasets>=1.8.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 6e0c35620f..36bc131de0 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.43.0") +check_min_version("4.45.0.dev0") check_optimum_habana_min_version("1.13.0") 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 e9e789b440..9b7d862e2e 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.43.0") +check_min_version("4.45.0.dev0") check_optimum_habana_min_version("1.13.0") 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 f494d5ea29..49dd2dc2e3 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.43.0") +check_min_version("4.45.0.dev0") check_optimum_habana_min_version("1.13.0") require_version("datasets>=1.18.0", "To fix: pip install -r examples/pytorch/speech-recognition/requirements.txt") diff --git a/examples/speech-recognition/run_speech_recognition_seq2seq.py b/examples/speech-recognition/run_speech_recognition_seq2seq.py index 66ed34f476..ad29827fde 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.43.0") +check_min_version("4.45.0.dev0") check_optimum_habana_min_version("1.13.0") 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 93ebb59824..c02d485d51 100644 --- 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.37.0") +check_min_version("4.45.0.dev0") check_optimum_habana_min_version("1.10.4") # Setup logging diff --git a/examples/summarization/run_summarization.py b/examples/summarization/run_summarization.py index 122477aed4..b5548d6250 100755 --- a/examples/summarization/run_summarization.py +++ b/examples/summarization/run_summarization.py @@ -65,7 +65,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.43.0") +check_min_version("4.45.0.dev0") check_optimum_habana_min_version("1.13.0") 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 5f5cb45b1b..eb6d41ef2a 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.43.0") +check_min_version("4.45.0.dev0") check_optimum_habana_min_version("1.13.0") 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 0dec28ed39..2eec5e3151 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.43.0") +check_min_version("4.45.0.dev0") check_optimum_habana_min_version("1.13.0") require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/translation/requirements.txt") diff --git a/optimum/habana/transformers/generation/__init__.py b/optimum/habana/transformers/generation/__init__.py index 6b43ee2ae3..09f85a5451 100644 --- a/optimum/habana/transformers/generation/__init__.py +++ b/optimum/habana/transformers/generation/__init__.py @@ -3,7 +3,6 @@ from .stopping_criteria import ( gaudi_EosTokenCriteria_call, gaudi_MaxLengthCriteria_call, - gaudi_MaxNewTokensCriteria_call, gaudi_MaxTimeCriteria_call, gaudi_StoppingCriteriaList_call, ) diff --git a/optimum/habana/transformers/generation/stopping_criteria.py b/optimum/habana/transformers/generation/stopping_criteria.py index dac7aadd92..69325ab7b3 100644 --- a/optimum/habana/transformers/generation/stopping_criteria.py +++ b/optimum/habana/transformers/generation/stopping_criteria.py @@ -52,18 +52,6 @@ def gaudi_MaxLengthCriteria_call( return create_return_const_tensor(input_ids, is_done) -def gaudi_MaxNewTokensCriteria_call( - self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs -) -> Union[torch.BoolTensor, bool]: - token_idx = kwargs.get("token_idx", None) - if token_idx is not None: - assert not kwargs["needs_tensor_output"] - return token_idx >= self.max_length - else: - is_done = input_ids.shape[-1] >= self.max_length - return create_return_const_tensor(input_ids, is_done) - - def gaudi_MaxTimeCriteria_call( self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs ) -> Union[torch.BoolTensor, bool]: diff --git a/optimum/habana/transformers/generation/utils.py b/optimum/habana/transformers/generation/utils.py index d333986679..89cc340dc3 100755 --- a/optimum/habana/transformers/generation/utils.py +++ b/optimum/habana/transformers/generation/utils.py @@ -22,7 +22,7 @@ import torch import torch.distributed as dist -from transformers.cache_utils import Cache, DynamicCache, EncoderDecoderCache, QuantizedCacheConfig +from transformers.cache_utils import Cache, DynamicCache, EncoderDecoderCache from transformers.generation.beam_constraints import DisjunctiveConstraint, PhrasalConstraint from transformers.generation.beam_search import BeamScorer, BeamSearchScorer, ConstrainedBeamSearchScorer from transformers.generation.candidate_generator import ( @@ -41,8 +41,6 @@ StopStringCriteria, ) from transformers.generation.utils import ( - NEED_SETUP_CACHE_CLASSES_MAPPING, - QUANT_BACKEND_CLASSES_MAPPING, GenerateBeamDecoderOnlyOutput, GenerateBeamEncoderDecoderOutput, GenerateBeamOutput, @@ -59,7 +57,7 @@ ) from transformers.integrations.deepspeed import is_deepspeed_zero3_enabled from transformers.modeling_outputs import CausalLMOutputWithPast, Seq2SeqLMOutput -from transformers.utils import ModelOutput, is_hqq_available, is_quanto_available, is_torchdynamo_compiling +from transformers.utils import ModelOutput, is_torchdynamo_compiling from optimum.utils import logging @@ -290,6 +288,10 @@ def _expand_inputs_for_generation( Copied from Transformers: https://github.com/huggingface/transformers/blob/527ab894e59b6582578008e3b47648a65063f73d/src/transformers/generation/utils.py#L704 The tensor `token_idx` is not expanded. """ + # Do not call torch.repeat_interleave if expand_size is 1 because it clones + # the input tensor and thus requires more memory although no change is applied + if expand_size == 1: + return input_ids, model_kwargs def _expand_dict_for_generation(dict_to_expand): for key in dict_to_expand: @@ -343,7 +345,6 @@ def _update_model_kwargs_for_generation( outputs: ModelOutput, model_kwargs: Dict[str, Any], is_encoder_decoder: bool = False, - standardize_cache_format: bool = False, num_new_tokens: int = 1, ) -> Dict[str, Any]: """ @@ -355,9 +356,7 @@ def _update_model_kwargs_for_generation( model_kwargs["first_token"] = False if not model_kwargs.get("pad_done", False): # update past_key_values keeping its naming used in model code - cache_name, cache = self._extract_past_from_model_output( - outputs, standardize_cache_format=standardize_cache_format - ) + cache_name, cache = self._extract_past_from_model_output(outputs) model_kwargs[cache_name] = cache if getattr(outputs, "state", None) is not None: model_kwargs["state"] = outputs.state @@ -495,6 +494,7 @@ def _get_candidate_generator( ) -> CandidateGenerator: if generation_config.prompt_lookup_num_tokens is not None: candidate_generator = PromptLookupCandidateGenerator( + eos_token_id=generation_config._eos_token_tensor, num_output_tokens=generation_config.prompt_lookup_num_tokens, max_matching_ngram_size=generation_config.max_matching_ngram_size, max_length=generation_config.max_length, @@ -615,19 +615,18 @@ def _prepare_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, - # three conditions must be met + # the following conditions must be met # 1) the generation config must have been created from the model config (`_from_model_config` field); # 2) the generation config must have seen no modification since its creation (the hash is the same); # 3) the user must have set 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 - and self.generation_config._original_object_hash == hash(self.generation_config) - and self.config._has_non_default_generation_parameters() + and self.generation_config._from_model_config # 1) + and self.generation_config._original_object_hash == hash(self.generation_config) # 2) ): new_generation_config = GaudiGenerationConfig.from_model_config(self.config) - if new_generation_config != self.generation_config: + if new_generation_config != self.generation_config: # 3) warnings.warn( "You have modified the pretrained model configuration to control generation. This is a" " deprecated strategy to control generation and will be removed soon, in a future version." @@ -637,20 +636,12 @@ def _prepare_generation_config( self.generation_config = new_generation_config using_model_generation_config = True 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. - if is_torchdynamo_compiling(): - model_kwargs = kwargs - generate_attributes_in_kwargs = [ - key for key, value in kwargs.items() if getattr(generation_config, key, None) != value - ] - if len(generate_attributes_in_kwargs) > 0: - raise ValueError( - "`torch.compile` exception: all generation configuration attributes must be passed within a " - f"`generation_config` instance passed to `generate` (found: {generate_attributes_in_kwargs})." - ) - else: + # 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 @@ -676,6 +667,8 @@ 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 return generation_config, model_kwargs @@ -984,76 +977,11 @@ def generate( has_token_idx="token_idx" in model_kwargs, ) - use_dynamic_cache_by_default = False - if "mamba" in self.__class__.__name__.lower(): - cache_name = "cache_params" - else: - cache_name = "past_key_values" - if generation_config.cache_implementation is not None and (model_kwargs.get(cache_name) is not None): - raise ValueError( - f"Passing both `cache_implementation` (used to initialize certain caches) and `{cache_name}` (a " - "Cache object) is unsupported. Please use only one of the two." - ) - elif 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: - raise ValueError( - "This model does not support `cache_implementation='static'`. Please check the following " - "issue: https://github.com/huggingface/transformers/issues/28981" - ) - model_kwargs[cache_name] = self._get_cache( - generation_config.cache_implementation, - getattr(generation_config, "num_beams", 1) * batch_size, - generation_config.max_length, - model_kwargs, - ) - elif generation_config.cache_implementation == "quantized": - if not self._supports_quantized_cache: - raise ValueError( - "This model does not support the quantized cache. If you want your model to support quantized " - "cache, please open an issue." - ) - - cache_config = ( - generation_config.cache_config - if generation_config.cache_config is not None - else QuantizedCacheConfig() - ) - cache_class = QUANT_BACKEND_CLASSES_MAPPING[cache_config.backend] - - if cache_config.backend == "quanto" and not is_quanto_available(): - raise ImportError( - "You need to install `quanto` in order to use KV cache quantization with quanto backend. " - "Please install it via with `pip install quanto`" - ) - 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) - # Use DynamicCache() instance by default. This will avoid back and forth from legacy format that - # keeps copying the cache thus using much more memory - # elif generation_config.cache_implementation is None and self._supports_default_dynamic_cache(): - # past = model_kwargs.get(cache_name, None) - # requires_cross_attention_cache = ( - # self.config.is_encoder_decoder or model_kwargs.get("encoder_outputs") is not None - # ) - # if past is None: - # model_kwargs[cache_name] = ( - # DynamicCache() - # if not requires_cross_attention_cache - # else EncoderDecoderCache(DynamicCache(), DynamicCache()) - # ) - # use_dynamic_cache_by_default = True - # elif isinstance(past, tuple): - # model_kwargs[cache_name] = ( - # DynamicCache.from_legacy_cache(past) - # if not requires_cross_attention_cache - # else EncoderDecoderCache.from_legacy_cache(past) - # ) - # use_dynamic_cache_by_default = True + # If the model supports `num_logits_to_keep` in forward(), set it to 1 to avoid computing the whole + # logit matrix. This can save a lot of memory during the first forward pass. Note that assisted decoding + # dynamically overrides this value as it can need more than the last token logits + if self._supports_num_logits_to_keep() and "num_logits_to_keep" not in model_kwargs: + model_kwargs["num_logits_to_keep"] = 1 self._validate_generated_length( generation_config, @@ -1061,6 +989,15 @@ def generate( has_default_max_length, ) + # 7. Prepare the cache. + # - `model_kwargs` may be updated in place with a cache as defined by the parameters in `generation_config`. + # - different models have a different cache name expected by the model (default = "past_key_values") + # - `max_length`, prepared above, is used to determine the maximum cache length + # TODO (joao): remove `user_defined_cache` after v4.47 (remove default conversion to legacy format) + cache_name = "past_key_values" if "mamba" not in self.__class__.__name__.lower() else "cache_params" + user_defined_cache = model_kwargs.get(cache_name) + self._prepare_cache_for_generation(generation_config, model_kwargs, assistant_model, batch_size, device) + # determine whether introduce trim_logits feature model_kwargs["trim_logits"] = generation_config.trim_logits @@ -1101,7 +1038,7 @@ def generate( if self.config.max_position_embeddings < calculated_max_length: unwrap_deepspeed_model(self).update_sincos_cache(seq_len=calculated_max_length) - # 7. determine generation mode + # 8. determine generation mode generation_mode = generation_config.get_generation_mode(assistant_model) if generation_config.bucket_size > 0: @@ -1121,7 +1058,7 @@ def generate( "`streamer` cannot be used with beam search (yet!). Make sure that `num_beams` is set to 1." ) - if self.device.type != input_ids.device.type: + if not is_torchdynamo_compiling() and self.device.type != input_ids.device.type: warnings.warn( ( "You are calling .generate() with the `input_ids` being on a device type different" @@ -1134,7 +1071,7 @@ def generate( UserWarning, ) - # 8. prepare distribution pre_processing samplers + # 9. prepare logits processors and stopping criteria prepared_logits_processor = self._get_logits_processor( generation_config=generation_config, input_ids_seq_length=input_ids_length, @@ -1146,8 +1083,6 @@ def generate( negative_prompt_ids=negative_prompt_ids, negative_prompt_attention_mask=negative_prompt_attention_mask, ) - - # 9. prepare stopping criteria self.generation_config.generation_mode = generation_mode prepared_stopping_criteria = self._get_stopping_criteria( generation_config=generation_config, @@ -1192,22 +1127,11 @@ def generate( model_kwargs=model_kwargs, ) - # 12. prepare logits warper (if `do_sample` is `True`) - prepared_logits_warper = ( - self._get_logits_warper( - generation_config, - device=input_ids.device, - ) - if generation_config.do_sample - else None - ) - - # 13. run assisted generate + # 12. run assisted generate result = self._assisted_decoding( input_ids, candidate_generator=candidate_generator, logits_processor=prepared_logits_processor, - logits_warper=prepared_logits_warper, stopping_criteria=prepared_stopping_criteria, generation_config=generation_config, synced_gpus=synced_gpus, @@ -1225,16 +1149,10 @@ def generate( raise ValueError( f"dola decoding is not supported with stateful models, such as {self.__class__.__name__}" ) - prepared_logits_warper = ( - self._get_logits_warper(generation_config, device=input_ids.device) - if generation_config.do_sample - else None - ) result = self._dola_decoding( input_ids, dola_layers=generation_config.dola_layers, logits_processor=prepared_logits_processor, - logits_warper=prepared_logits_warper, stopping_criteria=prepared_stopping_criteria, generation_config=generation_config, synced_gpus=synced_gpus, @@ -1268,26 +1186,18 @@ def generate( ) elif generation_mode in (GenerationMode.SAMPLE, GenerationMode.GREEDY_SEARCH): - # 11. prepare logits warper - prepared_logits_warper = ( - self._get_logits_warper(generation_config, device=input_ids.device) - if generation_config.do_sample - else None + # 11. expand input_ids with `num_return_sequences` additional sequences per batch + input_ids, model_kwargs = self._expand_inputs_for_generation( + input_ids=input_ids, + expand_size=generation_config.num_return_sequences, + is_encoder_decoder=self.config.is_encoder_decoder, + **model_kwargs, ) - if generation_mode == GenerationMode.SAMPLE: - # 12. expand input_ids with `num_return_sequences` additional sequences per batch - input_ids, model_kwargs = self._expand_inputs_for_generation( - input_ids=input_ids, - expand_size=generation_config.num_return_sequences, - is_encoder_decoder=self.config.is_encoder_decoder, - **model_kwargs, - ) - # 13. run sample (it degenerates to greedy search when `generation_config.do_sample=False`) + # 12. run sample (it degenerates to greedy search when `generation_config.do_sample=False`) result = self._sample( input_ids, logits_processor=prepared_logits_processor, - logits_warper=prepared_logits_warper, stopping_criteria=prepared_stopping_criteria, generation_config=generation_config, synced_gpus=synced_gpus, @@ -1302,14 +1212,7 @@ def generate( ) elif generation_mode in (GenerationMode.BEAM_SAMPLE, GenerationMode.BEAM_SEARCH): - # 11. prepare logits warper - prepared_logits_warper = ( - self._get_logits_warper(generation_config, device=input_ids.device) - if generation_config.do_sample - else None - ) - - # 12. prepare beam search scorer + # 11. prepare beam search scorer beam_scorer = BeamSearchScorer( batch_size=batch_size, num_beams=generation_config.num_beams, @@ -1320,7 +1223,7 @@ def generate( max_length=generation_config.max_length, ) - # 13. interleave input_ids with `num_beams` additional sequences per batch + # 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, @@ -1328,12 +1231,11 @@ def generate( **model_kwargs, ) - # 14. run beam sample + # 13. run beam sample result = self._beam_search( input_ids, beam_scorer, logits_processor=prepared_logits_processor, - logits_warper=prepared_logits_warper, stopping_criteria=prepared_stopping_criteria, generation_config=generation_config, synced_gpus=synced_gpus, @@ -1455,11 +1357,34 @@ def typeerror(): **model_kwargs, ) - # Convert to legacy cache if needed - if use_dynamic_cache_by_default and generation_config.return_legacy_cache: - if isinstance(result, ModelOutput) and hasattr(result, "past_key_values"): - if isinstance(result.past_key_values, (DynamicCache, EncoderDecoderCache)): - result.past_key_values = result.past_key_values.to_legacy_cache() + # Convert to legacy cache format if requested + if ( + generation_config.return_legacy_cache is not False # Should check for `True` after v4.47 + and not is_torchdynamo_compiling() + and hasattr(result, "past_key_values") + and hasattr(result.past_key_values, "to_legacy_cache") + and result.past_key_values.to_legacy_cache is not None + ): + # handle BC (convert by default if he user hasn't passed a cache AND the cache is of the default type) + should_convert_cache = generation_config.return_legacy_cache + is_user_defined_cache = user_defined_cache is not None + is_default_cache_type = ( + type(result.past_key_values) == DynamicCache # noqa E721 + or ( + isinstance(result.past_key_values, EncoderDecoderCache) + and type(result.past_key_values.self_attention_cache) == DynamicCache # noqa E721 + and type(result.past_key_values.cross_attention_cache) == DynamicCache # noqa E721 + ) + ) + if not is_user_defined_cache and is_default_cache_type: + logger.warning_once( + "From v4.47 onwards, when a model cache is to be returned, `generate` will return a `Cache` " + "instance instead by default (as opposed to the legacy tuple of tuples format). If you want to " + "keep returning the legacy format, please set `return_legacy_cache=True`." + ) + should_convert_cache = True + if should_convert_cache: + result.past_key_values = result.past_key_values.to_legacy_cache() return result @@ -1472,7 +1397,6 @@ def _dola_decoding( generation_config: GaudiGenerationConfig, synced_gpus: bool, streamer: "BaseStreamer", - logits_warper: Optional[LogitsProcessorList], **model_kwargs, ) -> Union[GenerateNonBeamOutput, torch.LongTensor]: r""" @@ -1501,10 +1425,6 @@ def _dola_decoding( streamer (`BaseStreamer`, *optional*): Streamer object that will be used to stream the generated sequences. Generated tokens are passed through `streamer.put(token_ids)` and the streamer is responsible for any further processing. - logits_warper (`LogitsProcessorList`, *optional*): - An instance of [`LogitsProcessorList`]. List of instances of class derived from [`LogitsWarper`] used - to warp the prediction score distribution of the language modeling head applied before multinomial - sampling at each generation step. model_kwargs: Additional model specific keyword arguments will be forwarded to the `forward` function of the model. If model is an encoder-decoder model the kwargs should include `encoder_outputs`. @@ -1698,15 +1618,13 @@ def _contrastive_search( else: logit_for_next_step = torch.index_select(outputs.logits, -2, token_idx - 1).squeeze(-2) else: - # Clone is needed to avoid keeping a hanging ref to outputs.logits which may be very large for this first iteration - # (the clone itself is always small) - logit_for_next_step = outputs.logits[:, -1, :].clone() + # .float() is needed to retain precision for later logits manipulations + logit_for_next_step = outputs.logits[:, -1, :].float() model_kwargs = self._update_model_kwargs_for_generation( outputs, model_kwargs, is_encoder_decoder=self.config.is_encoder_decoder, - standardize_cache_format=True, ) if not sequential: @@ -1866,7 +1784,8 @@ def _contrastive_search( next_hidden = outputs.hidden_states[-1] full_hidden_states = outputs.hidden_states - logits = outputs.logits[:, -1, :] + # .float() is needed to retain precision for later logits manipulations + logits = outputs.logits[:, -1, :].float() context_hidden = last_hidden_states.repeat_interleave(top_k, dim=0) # compute the degeneration penalty and re-rank the candidates based on the degeneration penalty and the @@ -1916,7 +1835,7 @@ def _contrastive_search( next_past_key_values = selected_outputs["past_key_values"] else: - _, next_past_key_values = self._extract_past_from_model_output(outputs, standardize_cache_format=True) + _, next_past_key_values = self._extract_past_from_model_output(outputs) # Do it in-place layer per layer to save memory if isinstance(next_past_key_values, DynamicCache) or ( isinstance(next_past_key_values, EncoderDecoderCache) @@ -2106,7 +2025,6 @@ def _sample( generation_config: GaudiGenerationConfig, synced_gpus: bool, streamer: Optional["BaseStreamer"], - logits_warper: Optional[LogitsProcessorList], lazy_mode: Optional[bool] = False, ignore_eos: Optional[bool] = False, profiling_warmup_steps: Optional[int] = 0, @@ -2135,11 +2053,6 @@ def _sample( streamer (`BaseStreamer`, *optional*): Streamer object that will be used to stream the generated sequences. Generated tokens are passed through `streamer.put(token_ids)` and the streamer is responsible for any further processing. - logits_warper (`LogitsProcessorList`, *optional*): - An instance of [`LogitsProcessorList`]. List of instances of class derived from [`LogitsWarper`] used - to warp the prediction score distribution of the language modeling head applied before multinomial - sampling at each generation step. Only required with sampling strategies (i.e. `do_sample` is set in - `generation_config`) lazy_mode (`bool`, *optional*, defaults to `False`): Whether the run is executed in lazy mode or not (i.e. eager mode). ignore_eos (`bool`, *optional*, defaults to `False`): @@ -2169,13 +2082,9 @@ 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 - if do_sample is True and not isinstance(logits_warper, LogitsProcessorList): - raise ValueError( - "`do_sample` is set to `True`, `logits_warper` must be a `LogitsProcessorList` instance (it is " - f"{logits_warper})." - ) # init attention / hidden states / scores tuples scores = () if (return_dict_in_generate and output_scores) else None @@ -2222,7 +2131,9 @@ def _sample( time_to_first_token_done = False model_kwargs["pad_done"] = 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, cur_len=cur_len, max_length=max_length + ): if lazy_mode: self.htcore_generation.mark_step() @@ -2256,7 +2167,7 @@ def _sample( if token_idx is not None and outputs.logits.shape[-2] > 1: # case1 (w/o KV caching): outputs.logits.shape: [batch_size, max_length, vocab_size] if self.config.is_encoder_decoder: - next_token_logits = outputs.logits[:, token_idx - 1, :] + next_token_logits = outputs.logits[:, token_idx - 1, :].float() next_token_scores = logits_processor(input_ids[:, :token_idx], next_token_logits) else: if model_kwargs.get("num_virtual_tokens", 0) > 0: @@ -2270,7 +2181,8 @@ def _sample( next_token_logits = torch.index_select(outputs.logits, -2, token_idx - 1).squeeze(-2) next_token_scores = logits_processor(input_ids, next_token_logits) else: - next_token_logits = outputs.logits[:, -1, :] + # .float() is needed to retain precision for later logits manipulations + next_token_logits = outputs.logits[:, -1, :].float() if token_idx is not None and self.config.is_encoder_decoder: # case2 (with KV caching): outputs.logits.shape: [batch_size, 1, vocab_size] next_token_scores = logits_processor(input_ids[:, :token_idx], next_token_logits) @@ -2278,10 +2190,6 @@ def _sample( # case3 (default case): token_idx is None next_token_scores = logits_processor(input_ids, next_token_logits) - # pre-process distribution - if do_sample: - next_token_scores = logits_warper(input_ids, next_token_scores) - # Store scores, attentions and hidden_states when required if return_dict_in_generate: if output_scores: @@ -2305,6 +2213,7 @@ def _sample( # token selection if do_sample: probs = torch.nn.functional.softmax(next_token_scores, dim=-1) + # TODO (joao): this OP throws "skipping cudagraphs due to ['incompatible ops']", find solution next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1) else: next_tokens = torch.argmax(next_token_scores, dim=-1) @@ -2437,7 +2346,6 @@ def _beam_search( stopping_criteria: StoppingCriteriaList, generation_config: GaudiGenerationConfig, synced_gpus: bool, - logits_warper: Optional[LogitsProcessorList], lazy_mode: Optional[bool] = False, profiling_warmup_steps: Optional[int] = 0, profiling_steps: Optional[int] = 0, @@ -2465,11 +2373,6 @@ def _beam_search( The generation configuration to be used as parametrization of the decoding method. synced_gpus (`bool`): Whether to continue running the while loop until max_length (needed for ZeRO stage 3) - logits_warper (`LogitsProcessorList`, *optional*): - An instance of [`LogitsProcessorList`]. List of instances of class derived from [`LogitsWarper`] used - to warp the prediction score distribution of the language modeling head applied before multinomial - sampling at each generation step. Only required with sampling strategies (i.e. `do_sample` is set in - `generation_config`) lazy_mode (`bool`, *optional*, defaults to `False`): Whether the run is executed in lazy mode or not (i.e. eager mode). profiling_warmup_steps (`int`, *optional*, defaults to 0): @@ -2499,11 +2402,6 @@ def _beam_search( return_dict_in_generate = generation_config.return_dict_in_generate sequential = generation_config.low_memory do_sample = generation_config.do_sample - if do_sample is True and not isinstance(logits_warper, LogitsProcessorList): - raise ValueError( - "`do_sample` is set to `True`, `logits_warper` must be a `LogitsProcessorList` instance (it is " - f"{logits_warper})." - ) batch_size = len(beam_scorer._beam_hyps) num_beams = beam_scorer.num_beams @@ -2674,7 +2572,6 @@ def expand_if_needed(tensor, new_size, value, dim=-1): for model_name in [ "fsmt", "reformer", - "bloom", "ctrl", "gpt_bigcode", "transo_xl", @@ -2720,9 +2617,7 @@ def expand_if_needed(tensor, new_size, value, dim=-1): else: next_token_logits = torch.index_select(outputs.logits, -2, token_idx - 1).squeeze(-2) else: - # Clone is needed to avoid keeping a hanging ref to outputs.logits which may be very large for first iteration - # (the clone itself is always small) - next_token_logits = outputs.logits[:, -1, :].clone() + next_token_logits = outputs.logits[:, -1, :].float() next_token_scores = torch.nn.functional.log_softmax( next_token_logits, dim=-1 @@ -2732,8 +2627,6 @@ def expand_if_needed(tensor, new_size, value, dim=-1): next_token_scores_processed = logits_processor(input_ids[:, :token_idx], next_token_scores) else: next_token_scores_processed = logits_processor(input_ids, next_token_scores) - if do_sample: - next_token_scores_processed = logits_warper(input_ids, next_token_scores_processed) next_token_scores = next_token_scores_processed + beam_scores[:, None].expand_as( next_token_scores_processed ) @@ -3051,10 +2944,6 @@ def _constrained_beam_search( stopping_criteria (`StoppingCriteriaList`): An instance of [`StoppingCriteriaList`]. List of instances of class derived from [`StoppingCriteria`] used to tell if the generation loop should stop. - logits_warper (`LogitsProcessorList`): - An instance of [`LogitsProcessorList`]. List of instances of class derived from [`LogitsWarper`] used - to warp the prediction score distribution of the language modeling head applied before multinomial - sampling at each generation step. generation_config ([`GaudiGenerationConfig`]): The generation configuration to be used as parametrization of the decoding method. synced_gpus (`bool`): @@ -3168,9 +3057,7 @@ def _constrained_beam_search( else: next_token_logits = torch.index_select(outputs.logits, -2, token_idx - 1).squeeze(-2) else: - # Clone is needed to avoid keeping a hanging ref to outputs.logits which may be very large for first iteration - # (the clone itself is always small) - next_token_logits = outputs.logits[:, -1, :].clone() + next_token_logits = outputs.logits[:, -1, :].float() next_token_scores = torch.nn.functional.log_softmax( next_token_logits, dim=-1 @@ -3326,7 +3213,6 @@ def _assisted_decoding( input_ids: torch.LongTensor, candidate_generator: "GaudiCandidateGenerator", logits_processor: LogitsProcessorList, - logits_warper: LogitsProcessorList, stopping_criteria: StoppingCriteriaList, generation_config: GaudiGenerationConfig, synced_gpus: bool, @@ -3354,10 +3240,6 @@ def _assisted_decoding( 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. - logits_warper (`LogitsProcessorList`): - An instance of [`LogitsProcessorList`]. List of instances of class derived from [`LogitsWarper`] used - to warp the prediction score distribution of the language modeling head applied before multinomial - sampling at each generation step. Only used if sampling is active. stopping_criteria (`StoppingCriteriaList`): An instance of [`StoppingCriteriaList`]. List of instances of class derived from [`StoppingCriteria`] used to tell if the generation loop should stop. @@ -3388,7 +3270,7 @@ def _assisted_decoding( `model.config.is_encoder_decoder=True`. """ # init values - do_sample = logits_warper is not None + do_sample = generation_config.do_sample output_attentions = generation_config.output_attentions output_hidden_states = generation_config.output_hidden_states output_scores = generation_config.output_scores @@ -3446,9 +3328,7 @@ def _assisted_decoding( model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs) # 1. Fetch candidate sequences from a `CandidateGenerator` - 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) @@ -3494,14 +3374,12 @@ def _assisted_decoding( ) # 2.3. Process the new logits - new_logits = outputs.logits[:, -candidate_length - 1 :] # excludes the input prompt if present + # .float() is needed to retain precision for later logits manipulations + new_logits = outputs.logits[:, -candidate_length - 1 :].float() # excludes the input prompt if present next_token_logits = new_logits.clone() if len(logits_processor) > 0: for i in range(candidate_length + 1): new_logits[:, i, :] = logits_processor(candidate_input_ids[:, : cur_len + i], new_logits[:, i, :]) - if do_sample and len(logits_warper) > 0: - for i in range(candidate_length + 1): - new_logits[:, i, :] = logits_warper(candidate_input_ids[:, : cur_len + i], new_logits[:, i, :]) # 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) diff --git a/optimum/habana/transformers/modeling_utils.py b/optimum/habana/transformers/modeling_utils.py index 2b7bb32bce..8f4706c053 100644 --- a/optimum/habana/transformers/modeling_utils.py +++ b/optimum/habana/transformers/modeling_utils.py @@ -23,7 +23,6 @@ GaudiGenerationMixin, gaudi_EosTokenCriteria_call, gaudi_MaxLengthCriteria_call, - gaudi_MaxNewTokensCriteria_call, gaudi_MaxTimeCriteria_call, gaudi_StoppingCriteriaList_call, ) @@ -267,7 +266,6 @@ def adapt_transformers_to_gaudi(): transformers.generation.GenerationConfig = GaudiGenerationConfig transformers.modeling_utils.GenerationConfig = GaudiGenerationConfig transformers.generation.MaxLengthCriteria.__call__ = gaudi_MaxLengthCriteria_call - transformers.generation.MaxNewTokensCriteria.__call__ = gaudi_MaxNewTokensCriteria_call transformers.generation.MaxTimeCriteria.__call__ = gaudi_MaxTimeCriteria_call transformers.generation.EosTokenCriteria.__call__ = gaudi_EosTokenCriteria_call transformers.generation.StoppingCriteriaList.__call__ = gaudi_StoppingCriteriaList_call diff --git a/optimum/habana/transformers/models/bloom/modeling_bloom.py b/optimum/habana/transformers/models/bloom/modeling_bloom.py index df99463c15..8c1ebc54c0 100644 --- a/optimum/habana/transformers/models/bloom/modeling_bloom.py +++ b/optimum/habana/transformers/models/bloom/modeling_bloom.py @@ -23,6 +23,7 @@ import torch from torch.nn import CrossEntropyLoss from torch.nn import functional as F +from transformers.cache_utils import Cache from transformers.modeling_outputs import BaseModelOutputWithPastAndCrossAttentions, CausalLMOutputWithCrossAttentions from transformers.models.bloom.modeling_bloom import BloomForCausalLM, BloomMLP, dropout_add from transformers.utils import logging @@ -124,16 +125,17 @@ def gaudi_bloom_attention_forward( residual: torch.Tensor, alibi: torch.Tensor, attention_mask: torch.Tensor, - layer_past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + layer_past: Optional[Cache] = None, head_mask: Optional[torch.Tensor] = None, use_cache: bool = False, output_attentions: bool = False, + cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, ): + batch_size, q_length, _ = hidden_states.shape fused_qkv = self.query_key_value(hidden_states) # [batch_size, seq_length, 3 x hidden_size] - - # 3 x [batch_size, seq_length, num_heads, head_dim] - (query_layer, key_layer, value_layer) = self._split_heads(fused_qkv) + # 3 x [batch_size, num_heads, seq_length, head_dim] + query_layer, key_layer, value_layer = self._reshape(fused_qkv) batch_size, q_length, _, _ = query_layer.shape @@ -225,10 +227,11 @@ def gaudi_bloom_block_forward( hidden_states: torch.Tensor, alibi: torch.Tensor, attention_mask: torch.Tensor, - layer_past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + layer_past: Optional[Cache] = None, head_mask: Optional[torch.Tensor] = None, use_cache: bool = False, output_attentions: bool = False, + cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, ): # hidden_states: [batch_size, seq_length, hidden_size] @@ -252,6 +255,7 @@ def gaudi_bloom_block_forward( head_mask=head_mask, use_cache=use_cache, output_attentions=output_attentions, + cache_position=cache_position, token_idx=token_idx, ) @@ -326,7 +330,7 @@ def gaudi_bloom_convert_to_bloom_cache( def gaudi_bloom_model_forward( self, input_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[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, @@ -334,6 +338,7 @@ def gaudi_bloom_model_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, **deprecated_arguments, ) -> Union[Tuple[torch.Tensor, ...], BaseModelOutputWithPastAndCrossAttentions]: @@ -429,6 +434,7 @@ def gaudi_bloom_model_forward( head_mask[i], use_cache, output_attentions, + cache_position, None, ) else: @@ -440,6 +446,7 @@ def gaudi_bloom_model_forward( use_cache=use_cache, output_attentions=output_attentions, alibi=alibi, + cache_position=cache_position, token_idx=token_idx, ) @@ -477,10 +484,12 @@ def set_tp_for_inference(tp_for_inference: int): def prepare_inputs_for_generation( self, - input_ids: torch.LongTensor, - past_key_values: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - inputs_embeds: Optional[torch.Tensor] = None, + input_ids, + past_key_values=None, + attention_mask=None, + inputs_embeds=None, + cache_position=None, + use_cache=True, token_idx: Optional[torch.Tensor] = None, **kwargs, ) -> dict: @@ -499,12 +508,13 @@ def prepare_inputs_for_generation( if inputs_embeds is not None and past_key_values is None: model_inputs = {"inputs_embeds": inputs_embeds} else: - model_inputs = {"input_ids": input_ids} + model_inputs = {"input_ids": input_ids.contiguous()} # `contiguous()` needed for compilation use cases model_inputs.update( { + "cache_position": cache_position, "past_key_values": past_key_values, - "use_cache": kwargs.get("use_cache"), + "use_cache": use_cache, "attention_mask": attention_mask, "token_idx": token_idx, } @@ -514,7 +524,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[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, @@ -523,6 +533,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, token_idx: Optional[torch.Tensor] = None, **deprecated_arguments, ) -> Union[Tuple[torch.Tensor], CausalLMOutputWithCrossAttentions]: @@ -554,6 +565,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, ) hidden_states = transformer_outputs[0] diff --git a/optimum/habana/transformers/models/codegen/modeling_codegen.py b/optimum/habana/transformers/models/codegen/modeling_codegen.py index 536cb5d423..80e1ce5710 100644 --- a/optimum/habana/transformers/models/codegen/modeling_codegen.py +++ b/optimum/habana/transformers/models/codegen/modeling_codegen.py @@ -10,18 +10,20 @@ apply_rotary_pos_emb, logger, ) +from transfromers.cache_utils import Cache class GaudiCodeGenAttention(CodeGenAttention): def forward( self, hidden_states: Optional[torch.FloatTensor], - layer_past: Optional[Tuple[torch.Tensor]] = None, + layer_past: Optional[Cache] = None, attention_mask: Optional[torch.FloatTensor] = None, position_ids: Optional[torch.LongTensor] = None, head_mask: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = False, output_attentions: Optional[bool] = False, + cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, ) -> Union[ Tuple[torch.Tensor, Tuple[torch.Tensor]], @@ -106,12 +108,13 @@ def forward( def gaudi_codegen_block_forward( self, hidden_states: Optional[torch.FloatTensor], - layer_past: Optional[Tuple[torch.Tensor]] = None, + layer_past: Optional[Cache] = None, attention_mask: Optional[torch.FloatTensor] = None, position_ids: Optional[torch.LongTensor] = None, head_mask: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = False, 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, ...]]]]: """ @@ -129,6 +132,7 @@ def gaudi_codegen_block_forward( head_mask=head_mask, use_cache=use_cache, output_attentions=output_attentions, + cache_position=cache_position, token_idx=token_idx, ) attn_output = attn_outputs[0] # output_attn: a, present, (attentions) @@ -148,7 +152,7 @@ def gaudi_codegen_block_forward( def gaudi_codegen_model_forward( self, input_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[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, @@ -158,6 +162,7 @@ def gaudi_codegen_model_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, BaseModelOutputWithPast]: """ @@ -229,14 +234,16 @@ def gaudi_codegen_model_forward( if inputs_embeds is None: inputs_embeds = self.wte(input_ids) + seq_length = inputs_embeds.shape[1] + hidden_states = inputs_embeds if token_type_ids is not None: + token_type_ids = token_type_ids.view(-1, seq_length) token_type_embeds = self.wte(token_type_ids) hidden_states = hidden_states + token_type_embeds hidden_states = self.drop(hidden_states) - output_shape = input_shape + (hidden_states.size(-1),) if self.gradient_checkpointing and self.training: @@ -264,6 +271,7 @@ def gaudi_codegen_model_forward( head_mask[i], use_cache, output_attentions, + cache_position, None, ) else: @@ -275,6 +283,7 @@ def gaudi_codegen_model_forward( head_mask=head_mask[i], use_cache=use_cache, output_attentions=output_attentions, + cache_position=cache_position, token_idx=token_idx, ) @@ -314,7 +323,17 @@ class GaudiCodeGenForCausalLM(CodeGenForCausalLM): """ def prepare_inputs_for_generation( - self, input_ids, inputs_embeds=None, past_key_values=None, token_idx=None, **kwargs + self, + input_ids, + attention_mask=None, + token_type_ids=None, + position_ids=None, + past_key_values=None, + inputs_embeds=None, + cache_position=None, + use_cache=True, + token_idx=None, + **kwargs, ): token_type_ids = kwargs.get("token_type_ids", None) # Omit tokens covered by past_key_values @@ -328,9 +347,6 @@ def prepare_inputs_for_generation( if token_type_ids is not None: token_type_ids = token_type_ids[:, -1] - attention_mask = kwargs.get("attention_mask", None) - position_ids = kwargs.get("position_ids", None) - if attention_mask is not None and position_ids is None: # create position_ids on the fly for batch generation position_ids = attention_mask.long().cumsum(-1) - 1 @@ -341,17 +357,21 @@ def prepare_inputs_for_generation( else: position_ids = position_ids[:, -1] + # This `clone` call is needed to avoid recapturing cuda graphs with `torch.compile`'s `mode="reduce-overhead`, as otherwise the input `position_ids` would have various stride during the decoding. Here, simply using `.contiguous()` is not sufficient as in the batch size = 1 case, `position_ids` is already contiguous but with varying stride which retriggers a capture. + position_ids = position_ids.clone(memory_format=torch.contiguous_format) + # if `inputs_embeds` are passed, we only want to use them in the 1st generation step if inputs_embeds is not None and past_key_values is None: model_inputs = {"inputs_embeds": inputs_embeds} else: - model_inputs = {"input_ids": input_ids.contiguous()} + model_inputs = {"input_ids": input_ids.clone(memory_format=torch.contiguous_format)} model_inputs.update( { - "past_key_values": past_key_values, - "use_cache": kwargs.get("use_cache"), "position_ids": position_ids, + "cache_position": cache_position, + "past_key_values": past_key_values, + "use_cache": use_cache, "attention_mask": attention_mask, "token_type_ids": token_type_ids, "token_idx": token_idx, @@ -362,7 +382,7 @@ 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[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, @@ -373,6 +393,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, token_idx: Optional[torch.Tensor] = None, ) -> Union[Tuple, CausalLMOutputWithPast]: r""" @@ -395,6 +416,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, ) hidden_states = transformer_outputs[0] diff --git a/optimum/habana/transformers/models/falcon/modeling_falcon.py b/optimum/habana/transformers/models/falcon/modeling_falcon.py index a7a0c0e920..d600e03bfd 100644 --- a/optimum/habana/transformers/models/falcon/modeling_falcon.py +++ b/optimum/habana/transformers/models/falcon/modeling_falcon.py @@ -30,6 +30,7 @@ from torch import nn from torch.nn import CrossEntropyLoss from torch.nn import functional as F +from transformers.cache_utils import Cache from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask_for_sdpa from transformers.modeling_outputs import ( BaseModelOutputWithPastAndCrossAttentions, @@ -253,7 +254,7 @@ class GaudiFalconAttention(FalconAttention): 4. not use_flash_attention, bf16: F.scaled_dot_product_attention. Slowest option """ - def __init__(self, config: FalconConfig): + def __init__(self, config: FalconConfig, layer_idx=None): super().__init__(config) self.is_fp8 = os.getenv("QUANT_CONFIG", "") != "" @@ -337,10 +338,11 @@ def pre_attn_forward( 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[Cache] = None, head_mask: Optional[torch.Tensor] = None, use_cache: bool = False, output_attentions: bool = False, + cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, reuse_cache: Optional[bool] = False, cache_idx: int = None, @@ -597,9 +599,9 @@ class GaudiFalconDecoderLayer(FalconDecoderLayer): - add new arg flash_attention_causal_mask """ - def __init__(self, config: FalconConfig): + def __init__(self, config: FalconConfig, layer_idx=None): super().__init__(config) - self.self_attention = GaudiFalconAttention(config) + self.self_attention = GaudiFalconAttention(config, layer_idx) def allocate_kv_cache(self, batch_size, max_seq_len, inp_seq_len): self.self_attention.allocate_kv_cache(batch_size, max_seq_len, inp_seq_len) @@ -613,10 +615,11 @@ def forward( 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[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, token_idx: Optional[torch.Tensor] = None, reuse_cache: Optional[bool] = False, cache_idx: int = None, @@ -642,6 +645,7 @@ def forward( head_mask=head_mask, use_cache=use_cache, output_attentions=output_attentions, + cache_position=cache_position, token_idx=token_idx, reuse_cache=reuse_cache, cache_idx=cache_idx, @@ -699,6 +703,7 @@ def pre_attn( head_mask: Optional[torch.Tensor] = None, use_cache: bool = False, output_attentions: bool = False, + cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, reuse_cache: Optional[bool] = False, cache_idx: int = None, @@ -723,6 +728,7 @@ def pre_attn( head_mask=head_mask, use_cache=use_cache, output_attentions=output_attentions, + cache_position=cache_position, token_idx=token_idx, reuse_cache=reuse_cache, cache_idx=cache_idx, @@ -757,7 +763,7 @@ def update_sincos_cache(self, seq_len): def forward( self, input_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[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, @@ -766,6 +772,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, token_idx: Optional[torch.Tensor] = None, reuse_cache: Optional[bool] = False, cache_idx: int = None, @@ -901,6 +908,7 @@ def forward( layer_past, use_cache, output_attentions, + cache_position, None, use_flash_attention, flash_attention_recompute, @@ -916,6 +924,7 @@ def forward( use_cache=use_cache, output_attentions=output_attentions, alibi=alibi, + cache_position=cache_position, token_idx=token_idx, reuse_cache=reuse_cache, cache_idx=cache_idx, @@ -972,10 +981,12 @@ def update_sincos_cache(self, seq_len): def prepare_inputs_for_generation( self, input_ids: torch.LongTensor, - past_key_values: Optional[torch.Tensor] = None, + past_key_values: Optional[Union[Cache, torch.Tensor]] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.Tensor] = None, inputs_embeds: Optional[torch.Tensor] = None, + cache_position: Optional[torch.LongTensor] = None, + use_cache: bool = True, token_idx: Optional[torch.Tensor] = None, **kwargs, ) -> dict: @@ -1017,16 +1028,20 @@ def prepare_inputs_for_generation( else: position_ids = position_ids[:, -input_ids.shape[1] :] + # This `clone` call is needed to avoid recapturing cuda graphs with `torch.compile`'s `mode="reduce-overhead`, as otherwise the input `position_ids` would have various stride during the decoding. Here, simply using `.contiguous()` is not sufficient as in the batch size = 1 case, `position_ids` is already contiguous but with varying stride which retriggers a capture. + position_ids = position_ids.clone(memory_format=torch.contiguous_format) + if inputs_embeds is not None and past_key_values is None: model_inputs = {"inputs_embeds": inputs_embeds} else: - model_inputs = {"input_ids": input_ids} + model_inputs = {"input_ids": input_ids.clone(memory_format=torch.contiguous_format)} model_inputs.update( { "position_ids": position_ids, + "cache_position": cache_position, "past_key_values": past_key_values, - "use_cache": kwargs.get("use_cache"), + "use_cache": use_cache, "attention_mask": attention_mask, "token_idx": token_idx, "reuse_cache": reuse_cache, @@ -1041,7 +1056,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[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, @@ -1051,6 +1066,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, token_idx: Optional[torch.Tensor] = None, reuse_cache: Optional[bool] = False, trim_logits: Optional[bool] = False, @@ -1084,6 +1100,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, reuse_cache=reuse_cache, cache_idx=cache_idx, diff --git a/optimum/habana/transformers/models/gemma/modeling_gemma.py b/optimum/habana/transformers/models/gemma/modeling_gemma.py index 6c537dfa31..88793e2cc3 100644 --- a/optimum/habana/transformers/models/gemma/modeling_gemma.py +++ b/optimum/habana/transformers/models/gemma/modeling_gemma.py @@ -34,7 +34,7 @@ apply_rotary_pos_emb, repeat_kv, ) -from transformers.utils import logging +from transformers.utils import is_torchdynamo_compiling, logging from ...modeling_attn_mask_utils import ( _gaudi_prepare_4d_causal_attention_mask, @@ -331,6 +331,7 @@ def forward( output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, + num_logits_to_keep: int = 0, token_idx: Optional[torch.Tensor] = None, ) -> Union[Tuple, CausalLMOutputWithPast]: """ @@ -360,10 +361,18 @@ def forward( ) hidden_states = outputs[0] - logits = self.lm_head(hidden_states) - logits = logits.float() + if labels is None and not is_torchdynamo_compiling(): + logger.warning_once( + "Starting from v4.46, the `logits` model output will have the same type as the model (except at train time, where it will always be FP32)" + ) + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + # TODO: remove the float() operation in v4.46 + logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]).float() + 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 so that tokens < n predict n shift_logits = logits[..., :-1, :].contiguous() shift_labels = labels[..., 1:].contiguous() @@ -396,6 +405,7 @@ def prepare_inputs_for_generation( cache_position=None, position_ids=None, use_cache=True, + num_logits_to_keep=0, **kwargs, ): """ @@ -430,6 +440,8 @@ def prepare_inputs_for_generation( position_ids = torch.index_select(position_ids, 1, token_idx - 1) else: position_ids = position_ids[:, -input_ids.shape[1] :] + # This `clone` call is needed to avoid recapturing cuda graphs with `torch.compile`'s `mode="reduce-overhead`, as otherwise the input `position_ids` would have various stride during the decoding. Here, simply using `.contiguous()` is not sufficient as in the batch size = 1 case, `position_ids` is already contiguous but with varying stride which retriggers a capture. + position_ids = position_ids.clone(memory_format=torch.contiguous_format) if token_idx is None: if past_key_value := getattr(self.model.layers[0].self_attn, "past_key_value", None): @@ -442,7 +454,7 @@ def prepare_inputs_for_generation( if inputs_embeds is not None and past_key_values is None: model_inputs = {"inputs_embeds": inputs_embeds} else: - model_inputs = {"input_ids": input_ids.contiguous()} + model_inputs = {"input_ids": input_ids.clone(memory_format=torch.contiguous_format)} model_inputs.update( { @@ -451,6 +463,7 @@ def prepare_inputs_for_generation( "past_key_values": past_key_values, "use_cache": use_cache, "attention_mask": attention_mask, + "num_logits_to_keep": num_logits_to_keep, "token_idx": token_idx, } ) 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 aa6423d2b1..a759cf7787 100644 --- a/optimum/habana/transformers/models/gpt_neox/modeling_gpt_neox.py +++ b/optimum/habana/transformers/models/gpt_neox/modeling_gpt_neox.py @@ -2,6 +2,7 @@ import torch from torch.nn import CrossEntropyLoss +from transformers.cache_utils import Cache from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from transformers.models.gpt_neox.modeling_gpt_neox import ( GPTNeoXAttention, @@ -29,9 +30,11 @@ def gaudi_gpt_neox_attention_forward( attention_mask: torch.FloatTensor, position_ids: torch.LongTensor, head_mask: Optional[torch.FloatTensor] = None, - layer_past: Optional[Tuple[torch.Tensor]] = None, + layer_past: Optional[Cache] = None, use_cache: Optional[bool] = False, output_attentions: Optional[bool] = False, + padding_mask: Optional[torch.Tensor] = None, + cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, ): """ @@ -103,14 +106,14 @@ def gaudi_gpt_neox_attention_forward( class GaudiGPTNeoXLayer(GPTNeoXLayer): - def __init__(self, config): + def __init__(self, config, layer_idx): super(GPTNeoXLayer, self).__init__() self.use_parallel_residual = config.use_parallel_residual self.input_layernorm = torch.nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) self.post_attention_layernorm = torch.nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) self.post_attention_dropout = torch.nn.Dropout(config.hidden_dropout) self.post_mlp_dropout = torch.nn.Dropout(config.hidden_dropout) - self.attention = GPTNeoXAttention(config) + self.attention = GPTNeoXAttention(config, layer_idx) self.mlp = GPTNeoXMLP(config) def forward( @@ -120,8 +123,9 @@ def forward( position_ids: Optional[torch.LongTensor] = None, head_mask: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = False, - layer_past: Optional[Tuple[torch.Tensor]] = None, + layer_past: Optional[Cache] = None, output_attentions: Optional[bool] = False, + cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, ): """ @@ -137,6 +141,7 @@ def forward( head_mask=head_mask, use_cache=use_cache, output_attentions=output_attentions, + cache_position=cache_position, token_idx=token_idx, ) attn_output = attention_layer_outputs[0] # output_attn: attn_output, present, (attn_weights) @@ -173,11 +178,12 @@ 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[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, return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, ) -> Union[Tuple, BaseModelOutputWithPast]: """ @@ -260,6 +266,7 @@ def gaudi_gpt_neox_model_forward( use_cache, None, output_attentions, + cache_position, None, ) else: @@ -271,6 +278,7 @@ def gaudi_gpt_neox_model_forward( layer_past=layer_past, use_cache=use_cache, output_attentions=output_attentions, + cache_position=cache_position, token_idx=token_idx, ) hidden_states = outputs[0] @@ -322,12 +330,13 @@ def forward( position_ids: Optional[torch.LongTensor] = None, inputs_embeds: Optional[torch.FloatTensor] = None, head_mask: Optional[torch.FloatTensor] = None, - past_key_values: Optional[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, 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, CausalLMOutputWithPast]: return_dict = return_dict if return_dict is not None else self.config.use_return_dict @@ -343,6 +352,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, ) @@ -372,7 +382,16 @@ def forward( ) def prepare_inputs_for_generation( - self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, token_idx=None, **kwargs + self, + input_ids, + past_key_values=None, + attention_mask=None, + inputs_embeds=None, + cache_position=None, + position_ids=None, + use_cache=True, + token_idx=None, + **kwargs, ): input_shape = input_ids.shape @@ -392,7 +411,6 @@ def prepare_inputs_for_generation( input_ids = input_ids[:, remove_prefix_length:] - position_ids = kwargs.get("position_ids", None) if attention_mask is not None and position_ids is None: # create position_ids on the fly for batch generation position_ids = attention_mask.long().cumsum(-1) - 1 @@ -402,6 +420,8 @@ def prepare_inputs_for_generation( position_ids = torch.index_select(position_ids, 1, token_idx - 1) else: position_ids = position_ids[:, -input_ids.shape[1] :] + # This `clone` call is needed to avoid recapturing cuda graphs with `torch.compile`'s `mode="reduce-overhead`, as otherwise the input `position_ids` would have various stride during the decoding. Here, simply using `.contiguous()` is not sufficient as in the batch size = 1 case, `position_ids` is already contiguous but with varying stride which retriggers a capture. + position_ids = position_ids.clone(memory_format=torch.contiguous_format) # if model is used as a decoder in encoder-decoder model, the decoder attention mask is created on the fly if attention_mask is None: @@ -411,13 +431,15 @@ def prepare_inputs_for_generation( if inputs_embeds is not None and past_key_values is None: model_inputs = {"inputs_embeds": inputs_embeds} else: - model_inputs = {"input_ids": input_ids} + model_inputs = {"input_ids": input_ids.clone(memory_format=torch.contiguous_format)} + model_inputs.update( { - "attention_mask": attention_mask, - "past_key_values": past_key_values, "position_ids": position_ids, - "use_cache": kwargs.get("use_cache"), + "cache_position": cache_position, + "past_key_values": past_key_values, + "use_cache": use_cache, + "attention_mask": attention_mask, "token_idx": token_idx, } ) diff --git a/optimum/habana/transformers/models/gptj/modeling_gptj.py b/optimum/habana/transformers/models/gptj/modeling_gptj.py index 4793766f6e..22b2b7a989 100644 --- a/optimum/habana/transformers/models/gptj/modeling_gptj.py +++ b/optimum/habana/transformers/models/gptj/modeling_gptj.py @@ -16,6 +16,7 @@ create_sinusoidal_positions, logger, ) +from transfroemrs.cache_utils import Cache class Matmul(nn.Module): @@ -68,7 +69,7 @@ def forward(self, cur, dim, idx): class GaudiGPTJAttention(GPTJAttention): - def __init__(self, config: GPTJConfig): + def __init__(self, config: GPTJConfig, layer_idx=None): super().__init__(config) self.config = config @@ -155,12 +156,13 @@ def _attn( def forward( self, hidden_states: torch.FloatTensor, - layer_past: Optional[Tuple[torch.Tensor]] = None, + layer_past: Optional[Cache] = None, attention_mask: Optional[torch.FloatTensor] = None, position_ids: Optional[torch.LongTensor] = None, head_mask: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = False, output_attentions: Optional[bool] = False, + cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, sin: Optional[torch.Tensor] = None, cos: Optional[torch.Tensor] = None, @@ -265,11 +267,11 @@ class GaudiGPTJBlock(GPTJBlock): Inherits from GPTJBlock: https://github.com/huggingface/transformers/blob/v4.38.2/src/transformers/models/gptj/modeling_gptj.py#291 """ - def __init__(self, config: GPTJConfig): - super().__init__(config) + def __init__(self, config: GPTJConfig, layer_idx=None): + super().__init__(config, layer_idx=None) inner_dim = config.n_inner if config.n_inner is not None else 4 * config.n_embd self.ln_1 = torch.nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon) - self.attn = GaudiGPTJAttention(config) + self.attn = GaudiGPTJAttention(config, layer_idx) self.mlp = GPTJMLP(inner_dim, config) def allocate_kv_cache(self, batch_size, max_seq_len, inp_seq_len): @@ -284,12 +286,13 @@ def update_sincos_cache(self, seq_len): def forward( self, hidden_states: Optional[torch.FloatTensor], - layer_past: Optional[Tuple[torch.Tensor]] = None, + layer_past: Optional[Cache] = None, attention_mask: Optional[torch.FloatTensor] = None, position_ids: Optional[torch.LongTensor] = None, head_mask: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = False, output_attentions: Optional[bool] = False, + cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, sin: Optional[torch.Tensor] = None, cos: Optional[torch.Tensor] = None, @@ -312,6 +315,7 @@ def forward( head_mask=head_mask, use_cache=use_cache, output_attentions=output_attentions, + cache_position=cache_position, token_idx=token_idx, reuse_cache=reuse_cache, cache_idx=cache_idx, @@ -351,7 +355,7 @@ def update_sincos_cache(self, seq_len): def forward( self, input_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[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, @@ -361,6 +365,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, token_idx: Optional[torch.Tensor] = None, reuse_cache: Optional[bool] = False, cache_idx: Optional[int] = None, @@ -489,6 +494,7 @@ def forward( head_mask[i], use_cache, output_attentions, + cache_position, None, sin, cos, @@ -502,6 +508,7 @@ def forward( 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, @@ -555,11 +562,19 @@ def update_sincos_cache(self, seq_len): self.transformer.update_sincos_cache(seq_len) def prepare_inputs_for_generation( - self, input_ids, past_key_values=None, inputs_embeds=None, token_idx=None, **kwargs + self, + input_ids, + attention_mask=None, + token_type_ids=None, + position_ids=None, + past_key_values=None, + inputs_embeds=None, + cache_position=None, + use_cache=True, + token_idx=None, + **kwargs, ): reuse_cache = kwargs.get("reuse_cache") - token_type_ids = kwargs.get("token_type_ids", None) - attention_mask = kwargs.get("attention_mask", None) # Omit tokens covered by past_key_values if past_key_values: if token_idx is not None: @@ -586,8 +601,6 @@ def prepare_inputs_for_generation( input_ids = input_ids[:, :token_idx] attention_mask = attention_mask[:, :token_idx] - position_ids = kwargs.get("position_ids", None) - if attention_mask is not None and position_ids is None: # create position_ids on the fly for batch generation position_ids = attention_mask.long().cumsum(-1) - 1 @@ -597,18 +610,21 @@ def prepare_inputs_for_generation( position_ids = torch.index_select(position_ids, 1, token_idx - 1) else: position_ids = position_ids[:, -input_ids.shape[1] :] + # This `clone` call is needed to avoid recapturing cuda graphs with `torch.compile`'s `mode="reduce-overhead`, as otherwise the input `position_ids` would have various stride during the decoding. Here, simply using `.contiguous()` is not sufficient as in the batch size = 1 case, `position_ids` is already contiguous but with varying stride which retriggers a capture. + position_ids = position_ids.clone(memory_format=torch.contiguous_format) # if `inputs_embeds` are passed, we only want to use them in the 1st generation step if inputs_embeds is not None and past_key_values is None: model_inputs = {"inputs_embeds": inputs_embeds} else: - model_inputs = {"input_ids": input_ids} + model_inputs = {"input_ids": input_ids.clone(memory_format=torch.contiguous_format)} model_inputs.update( { - "past_key_values": past_key_values, - "use_cache": kwargs.get("use_cache"), "position_ids": position_ids, + "cache_position": cache_position, + "past_key_values": past_key_values, + "use_cache": use_cache, "attention_mask": attention_mask, "token_type_ids": token_type_ids, "token_idx": token_idx, @@ -622,7 +638,7 @@ 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[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, @@ -633,6 +649,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, token_idx: Optional[torch.Tensor] = None, reuse_cache: Optional[bool] = False, cache_idx: Optional[int] = None, @@ -657,6 +674,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, reuse_cache=reuse_cache, cache_idx=cache_idx, diff --git a/optimum/habana/transformers/models/llama/configuration_llama.py b/optimum/habana/transformers/models/llama/configuration_llama.py index ce754dadb5..fb159cfc48 100644 --- a/optimum/habana/transformers/models/llama/configuration_llama.py +++ b/optimum/habana/transformers/models/llama/configuration_llama.py @@ -25,6 +25,7 @@ def __init__( attention_bias=False, attention_dropout=0.0, mlp_bias=False, + head_dim=None, fused_qkv=False, parallel_strategy=None, **kwargs, diff --git a/optimum/habana/transformers/models/llama/modeling_llama.py b/optimum/habana/transformers/models/llama/modeling_llama.py index 7d41126390..eabd821278 100755 --- a/optimum/habana/transformers/models/llama/modeling_llama.py +++ b/optimum/habana/transformers/models/llama/modeling_llama.py @@ -20,6 +20,7 @@ apply_rotary_pos_emb, logger, ) +from transformers.utils import is_torchdynamo_compiling from .... import distributed from ....distributed.strategy import DistributedStrategy, NoOpStrategy @@ -1245,6 +1246,7 @@ def forward( output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, + num_logits_to_keep: int = 0, token_idx: Optional[torch.Tensor] = None, trim_logits: Optional[bool] = False, attn_softmax_bf16: Optional[bool] = False, @@ -1302,11 +1304,18 @@ def forward( logits = [F.linear(hidden_states, lm_head_slices[i]) for i in range(self.config.pretraining_tp)] logits = torch.cat(logits, dim=-1) else: - logits = self.lm_head(hidden_states) - logits = logits.float() + if labels is None and not is_torchdynamo_compiling(): + logger.warning_once( + "Starting from v4.46, the `logits` model output will have the same type as the model (except at train time, where it will always be FP32)" + ) + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + # TODO: remove the float() operation in v4.46 + logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]).float() 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 so that tokens < n predict n shift_logits = logits[..., :-1, :].contiguous() shift_labels = labels[..., 1:].contiguous() @@ -1339,6 +1348,7 @@ def prepare_inputs_for_generation( cache_position=None, position_ids=None, use_cache=True, + num_logits_to_keep=0, token_idx=None, **kwargs, ): @@ -1369,6 +1379,8 @@ def prepare_inputs_for_generation( position_ids = torch.index_select(position_ids, 1, token_idx - 1) else: position_ids = position_ids[:, -input_ids.shape[1] :] + # This `clone` call is needed to avoid recapturing cuda graphs with `torch.compile`'s `mode="reduce-overhead`, as otherwise the input `position_ids` would have various stride during the decoding. Here, simply using `.contiguous()` is not sufficient as in the batch size = 1 case, `position_ids` is already contiguous but with varying stride which retriggers a capture. + position_ids = position_ids.clone(memory_format=torch.contiguous_format) # keep cache_position implementation as None for HPU cache_position = None @@ -1377,7 +1389,7 @@ def prepare_inputs_for_generation( if inputs_embeds is not None and past_key_values is None: model_inputs = {"inputs_embeds": inputs_embeds} else: - model_inputs = {"input_ids": input_ids.contiguous()} + model_inputs = {"input_ids": input_ids.clone(memory_format=torch.contiguous_format)} model_inputs.update( { @@ -1386,6 +1398,7 @@ def prepare_inputs_for_generation( "past_key_values": past_key_values, "use_cache": use_cache, "attention_mask": attention_mask, + "num_logits_to_keep": num_logits_to_keep, "token_idx": token_idx, "trim_logits": kwargs.get("trim_logits"), "attn_softmax_bf16": kwargs.get("attn_softmax_bf16"), diff --git a/optimum/habana/transformers/models/llava/modeling_llava.py b/optimum/habana/transformers/models/llava/modeling_llava.py index 8119f442c5..4300f6c7b3 100644 --- a/optimum/habana/transformers/models/llava/modeling_llava.py +++ b/optimum/habana/transformers/models/llava/modeling_llava.py @@ -120,6 +120,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, token_idx: Optional[torch.Tensor] = None, image_offset: Optional[int] = None, tokens_pos: Optional[torch.LongTensor] = None, @@ -186,6 +187,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 + image_offset, use_flash_attention=use_flash_attention, flash_attention_recompute=flash_attention_recompute, @@ -230,7 +232,14 @@ def forward( ) def prepare_inputs_for_generation( - self, input_ids, past_key_values=None, inputs_embeds=None, pixel_values=None, attention_mask=None, **kwargs + self, + input_ids, + past_key_values=None, + inputs_embeds=None, + pixel_values=None, + attention_mask=None, + cache_position=None, + **kwargs, ): """ Inherits from LlavaForConditionalGeneration: https://github.com/huggingface/transformers/blob/v4.37.2/src/transformers/models/llava/modeling_llava.py 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 4670469e9e..dca9e8d28a 100644 --- a/optimum/habana/transformers/models/llava_next/modeling_llava_next.py +++ b/optimum/habana/transformers/models/llava_next/modeling_llava_next.py @@ -53,6 +53,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, token_idx: Optional[torch.Tensor] = None, use_flash_attention: Optional[bool] = False, flash_attention_recompute: Optional[bool] = False, @@ -84,6 +85,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 + self.image_offset, use_flash_attention=use_flash_attention, flash_attention_recompute=flash_attention_recompute, @@ -230,6 +232,7 @@ def prepare_inputs_for_generation( pixel_values=None, image_sizes=None, attention_mask=None, + cache_position=None, **kwargs, ): """ diff --git a/optimum/habana/transformers/models/mamba/modeling_mamba.py b/optimum/habana/transformers/models/mamba/modeling_mamba.py index ea7c112c7d..b9ac519318 100644 --- a/optimum/habana/transformers/models/mamba/modeling_mamba.py +++ b/optimum/habana/transformers/models/mamba/modeling_mamba.py @@ -24,10 +24,18 @@ def gaudi_MambaForCausalLM_update_model_kwargs_for_generation( and model_kwargs["cache_position"] is not None ): model_kwargs["cache_position"] = model_kwargs["cache_position"][-1:] + num_new_tokens + + if "attention_mask" in model_kwargs: + attention_mask = model_kwargs["attention_mask"] + model_kwargs["attention_mask"] = torch.cat( + [attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1 + ) + if token_idx is not None: token_idx.add_(1) if "token_idx_cpu" in model_kwargs: model_kwargs["token_idx_cpu"] += 1 + return model_kwargs @@ -38,7 +46,7 @@ def gaudi_MambaForCausalLM_prepare_inputs_for_generation( use_cache=None, cache_params: Optional[MambaCache] = None, cache_position: Optional[torch.LongTensor] = None, - attention_mask=None, + attention_mask: Optional[torch.LongTensor] = None, **kwargs, ): token_idx = kwargs.get("token_idx", None) @@ -54,6 +62,10 @@ def gaudi_MambaForCausalLM_prepare_inputs_for_generation( ) if cache_position[0] > 0: input_ids = input_ids[:, -1].unsqueeze(-1) + + if attention_mask is not None: + attention_mask = None + 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 @@ -75,6 +87,7 @@ def gaudi_MambaForCausalLM_prepare_inputs_for_generation( "cache_params": cache_params, "use_cache": use_cache, "cache_position": cache_position, + "attention_mask": attention_mask, } ) return model_inputs diff --git a/optimum/habana/transformers/models/mistral/modeling_mistral.py b/optimum/habana/transformers/models/mistral/modeling_mistral.py index 7d95e548ce..6ae4ede549 100644 --- a/optimum/habana/transformers/models/mistral/modeling_mistral.py +++ b/optimum/habana/transformers/models/mistral/modeling_mistral.py @@ -39,7 +39,7 @@ MistralRMSNorm, apply_rotary_pos_emb, ) -from transformers.utils import logging +from transformers.utils import is_torchdynamo_compiling, logging from ...modeling_attn_mask_utils import ( _gaudi_prepare_4d_causal_attention_mask, @@ -696,6 +696,7 @@ def forward( output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, + num_logits_to_keep: int = 0, token_idx: Optional[torch.Tensor] = None, reuse_cache: Optional[bool] = False, trim_logits: Optional[bool] = False, @@ -750,11 +751,18 @@ def forward( hidden_states = hidden_states.index_select(1, token_idx - 1) else: hidden_states = hidden_states[:, -1, :] - logits = self.lm_head(hidden_states) - logits = logits.float() + if labels is None and not is_torchdynamo_compiling(): + logger.warning_once( + "Starting from v4.46, the `logits` model output will have the same type as the model (except at train time, where it will always be FP32)" + ) + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + # TODO: remove the float() operation in v4.46 + logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]).float() 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 so that tokens < n predict n shift_logits = logits[..., :-1, :].contiguous() shift_labels = labels[..., 1:].contiguous() @@ -787,6 +795,7 @@ def prepare_inputs_for_generation( cache_position=None, position_ids=None, use_cache=True, + num_logits_to_keep=0, **kwargs, ): """ @@ -825,6 +834,8 @@ def prepare_inputs_for_generation( position_ids = torch.index_select(position_ids, 1, token_idx - 1) else: position_ids = position_ids[:, -input_ids.shape[1] :] + # This `clone` call is needed to avoid recapturing cuda graphs with `torch.compile`'s `mode="reduce-overhead`, as otherwise the input `position_ids` would have various stride during the decoding. Here, simply using `.contiguous()` is not sufficient as in the batch size = 1 case, `position_ids` is already contiguous but with varying stride which retriggers a capture. + position_ids = position_ids.clone(memory_format=torch.contiguous_format) # if `inputs_embeds` are passed, we only want to use them in the 1st generation step if inputs_embeds is not None and past_key_values is None: @@ -839,6 +850,7 @@ def prepare_inputs_for_generation( "past_key_values": past_key_values, "use_cache": use_cache, "attention_mask": attention_mask, + "num_logits_to_keep": num_logits_to_keep, "token_idx": token_idx, "reuse_cache": kwargs.get("reuse_cache"), "trim_logits": kwargs.get("trim_logits"), diff --git a/optimum/habana/transformers/models/mixtral/modeling_mixtral.py b/optimum/habana/transformers/models/mixtral/modeling_mixtral.py index 43dfc7e48a..a91444600f 100644 --- a/optimum/habana/transformers/models/mixtral/modeling_mixtral.py +++ b/optimum/habana/transformers/models/mixtral/modeling_mixtral.py @@ -44,7 +44,7 @@ apply_rotary_pos_emb, load_balancing_loss_func, ) -from transformers.utils import logging +from transformers.utils import is_torchdynamo_compiling, logging from ..llama.modeling_llama import ( GaudiLlamaDynamicNTKScalingRotaryEmbedding, @@ -745,6 +745,7 @@ def forward( output_router_logits: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, + num_logits_to_keep: int = 0, token_idx: Optional[torch.Tensor] = None, reuse_cache: Optional[bool] = None, flash_attention_recompute: Optional[bool] = False, @@ -780,11 +781,18 @@ def forward( ) hidden_states = outputs[0] - logits = self.lm_head(hidden_states) - logits = logits.float() + if labels is None and not is_torchdynamo_compiling(): + logger.warning_once( + "Starting from v4.46, the `logits` model output will have the same type as the model (except at train time, where it will always be FP32)" + ) + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + # TODO: remove the float() operation in v4.46 + logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]).float() 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 so that tokens < n predict n shift_logits = logits[..., :-1, :].contiguous() shift_labels = labels[..., 1:].contiguous() @@ -833,6 +841,7 @@ def prepare_inputs_for_generation( output_router_logits=False, position_ids=None, use_cache=True, + num_logits_to_keep=0, **kwargs, ): reuse_cache = kwargs.get("reuse_cache") @@ -877,6 +886,8 @@ def prepare_inputs_for_generation( "past_key_values": past_key_values, "use_cache": use_cache, "attention_mask": attention_mask, + "output_router_logits": output_router_logits, + "num_logits_to_keep": num_logits_to_keep, "token_idx": token_idx, "reuse_cache": reuse_cache, "flash_attention_recompute": kwargs.get("flash_attention_recompute"), diff --git a/optimum/habana/transformers/models/persimmon/modeling_persimmon.py b/optimum/habana/transformers/models/persimmon/modeling_persimmon.py index 4c7b24b988..8fb48d8f2c 100644 --- a/optimum/habana/transformers/models/persimmon/modeling_persimmon.py +++ b/optimum/habana/transformers/models/persimmon/modeling_persimmon.py @@ -339,6 +339,7 @@ def forward( output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, + num_logits_to_keep: int = 0, token_idx: Optional[torch.Tensor] = None, ) -> Union[Tuple, CausalLMOutputWithPast]: """ @@ -369,7 +370,8 @@ def forward( ) hidden_states = outputs[0] - logits = self.lm_head(hidden_states) + # No upscaling to float was ever done for Persimmon + logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]) loss = None if labels is not None: @@ -405,6 +407,7 @@ def prepare_inputs_for_generation( cache_position=None, position_ids=None, use_cache=True, + num_logits_to_keep=0, **kwargs, ): """ @@ -436,12 +439,16 @@ def prepare_inputs_for_generation( position_ids = torch.index_select(position_ids, 1, token_idx - 1) else: position_ids = position_ids[:, -input_ids.shape[1] :] + # This `clone` call is needed to avoid recapturing cuda graphs with `torch.compile`'s `mode="reduce-overhead`, as otherwise the input `position_ids` would have various stride during the decoding. Here, simply using `.contiguous()` is not sufficient as in the batch size = 1 case, `position_ids` is already contiguous but with varying stride which retriggers a capture. + position_ids = position_ids.clone(memory_format=torch.contiguous_format) # if `inputs_embeds` are passed, we only want to use them in the 1st generation step if inputs_embeds is not None and past_key_values is None: model_inputs = {"inputs_embeds": inputs_embeds} else: - model_inputs = {"input_ids": input_ids.contiguous()} # `contiguous()` needed for compilation use cases + model_inputs = { + "input_ids": input_ids.clone(memory_format=torch.contiguous_format) + } # `contiguous()` needed for compilation use cases model_inputs.update( { @@ -450,6 +457,7 @@ def prepare_inputs_for_generation( "past_key_values": past_key_values, "use_cache": use_cache, "attention_mask": attention_mask, + "num_logits_to_keep": num_logits_to_keep, "token_idx": token_idx, } ) diff --git a/optimum/habana/transformers/models/phi/modeling_phi.py b/optimum/habana/transformers/models/phi/modeling_phi.py index 1e21735add..81c56bec4f 100644 --- a/optimum/habana/transformers/models/phi/modeling_phi.py +++ b/optimum/habana/transformers/models/phi/modeling_phi.py @@ -35,7 +35,7 @@ PhiModel, apply_rotary_pos_emb, ) -from transformers.utils import logging +from transformers.utils import is_torchdynamo_compiling, logging from ...modeling_attn_mask_utils import ( _gaudi_prepare_4d_causal_attention_mask, @@ -532,6 +532,7 @@ def forward( output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, + num_logits_to_keep: int = 0, token_idx: Optional[torch.Tensor] = None, reuse_cache: Optional[bool] = False, trim_logits: Optional[bool] = False, @@ -575,11 +576,18 @@ def forward( hidden_states = hidden_states.index_select(1, token_idx - 1) else: hidden_states = hidden_states[:, -1, :] - logits = self.lm_head(hidden_states) - logits = logits.float() + if labels is None and not is_torchdynamo_compiling(): + logger.warning_once( + "Starting from v4.46, the `logits` model output will have the same type as the model (except at train time, where it will always be FP32)" + ) + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + # TODO: remove the float() operation in v4.46 + logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]).float() 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 so that tokens < n predict n shift_logits = logits[..., :-1, :].contiguous() shift_labels = labels[..., 1:].contiguous() @@ -612,6 +620,7 @@ def prepare_inputs_for_generation( cache_position=None, position_ids=None, use_cache=True, + num_logits_to_keep=0, token_idx=None, **kwargs, ): @@ -649,12 +658,16 @@ def prepare_inputs_for_generation( position_ids = torch.index_select(position_ids, 1, token_idx - 1) else: position_ids = position_ids[:, -input_ids.shape[1] :] + # This `clone` call is needed to avoid recapturing cuda graphs with `torch.compile`'s `mode="reduce-overhead`, as otherwise the input `position_ids` would have various stride during the decoding. Here, simply using `.contiguous()` is not sufficient as in the batch size = 1 case, `position_ids` is already contiguous but with varying stride which retriggers a capture. + position_ids = position_ids.clone(memory_format=torch.contiguous_format) # if `inputs_embeds` are passed, we only want to use them in the 1st generation step if inputs_embeds is not None and past_key_values is None: model_inputs = {"inputs_embeds": inputs_embeds} else: - model_inputs = {"input_ids": input_ids.contiguous()} # `contiguous()` needed for compilation use cases + model_inputs = { + "input_ids": input_ids.clone(memory_format=torch.contiguous_format) + } # `contiguous()` needed for compilation use cases model_inputs.update( { @@ -663,10 +676,12 @@ def prepare_inputs_for_generation( "past_key_values": past_key_values, "use_cache": use_cache, "attention_mask": attention_mask, + "num_logits_to_keep": num_logits_to_keep, "token_idx": token_idx, "reuse_cache": kwargs.get("reuse_cache"), "trim_logits": kwargs.get("trim_logits"), "cache_idx": kwargs.get("cache_idx"), } ) + return model_inputs diff --git a/optimum/habana/transformers/models/qwen2/modeling_qwen2.py b/optimum/habana/transformers/models/qwen2/modeling_qwen2.py index 0c8970dd88..53ac275bdc 100644 --- a/optimum/habana/transformers/models/qwen2/modeling_qwen2.py +++ b/optimum/habana/transformers/models/qwen2/modeling_qwen2.py @@ -35,6 +35,7 @@ apply_rotary_pos_emb, logger, ) +from transformers.utils import is_torchdynamo_compiling from ...modeling_attn_mask_utils import ( _gaudi_prepare_4d_causal_attention_mask, @@ -763,6 +764,7 @@ def forward( output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, + num_logits_to_keep: int = 0, token_idx: Optional[torch.Tensor] = None, trim_logits: Optional[bool] = False, attn_softmax_bf16: Optional[bool] = False, @@ -816,10 +818,18 @@ def forward( else: hidden_states = hidden_states[:, -1, :] - logits = self.lm_head(hidden_states).float() + if labels is None and not is_torchdynamo_compiling(): + logger.warning_once( + "Starting from v4.46, the `logits` model output will have the same type as the model (except at train time, where it will always be FP32)" + ) + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + # TODO: remove the float() operation in v4.46 + logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]).float() 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 so that tokens < n predict n shift_logits = logits[..., :-1, :].contiguous() shift_labels = labels[..., 1:].contiguous() @@ -852,6 +862,7 @@ def prepare_inputs_for_generation( cache_position=None, position_ids=None, use_cache=True, + num_logits_to_keep=0, token_idx=None, **kwargs, ): @@ -882,6 +893,8 @@ def prepare_inputs_for_generation( position_ids = torch.index_select(position_ids, 1, token_idx - 1) else: position_ids = position_ids[:, -input_ids.shape[1] :] + # This `clone` call is needed to avoid recapturing cuda graphs with `torch.compile`'s `mode="reduce-overhead`, as otherwise the input `position_ids` would have various stride during the decoding. Here, simply using `.contiguous()` is not sufficient as in the batch size = 1 case, `position_ids` is already contiguous but with varying stride which retriggers a capture. + position_ids = position_ids.clone(memory_format=torch.contiguous_format) cache_position = None @@ -889,7 +902,9 @@ def prepare_inputs_for_generation( if inputs_embeds is not None and past_key_values is None: model_inputs = {"inputs_embeds": inputs_embeds} else: - model_inputs = {"input_ids": input_ids.contiguous()} # `contiguous()` needed for compilation use cases + model_inputs = { + "input_ids": input_ids.clone(memory_format=torch.contiguous_format) + } # `contiguous()` needed for compilation use cases model_inputs.update( { @@ -898,6 +913,7 @@ def prepare_inputs_for_generation( "past_key_values": past_key_values, "use_cache": use_cache, "attention_mask": attention_mask, + "num_logits_to_keep": num_logits_to_keep, "token_idx": token_idx, "trim_logits": kwargs.get("trim_logits"), "attn_softmax_bf16": kwargs.get("attn_softmax_bf16"), diff --git a/optimum/habana/transformers/models/stablelm/modeling_stablelm.py b/optimum/habana/transformers/models/stablelm/modeling_stablelm.py index 08becc263a..aadd9469c1 100644 --- a/optimum/habana/transformers/models/stablelm/modeling_stablelm.py +++ b/optimum/habana/transformers/models/stablelm/modeling_stablelm.py @@ -370,6 +370,7 @@ def forward( output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, + num_logits_to_keep: int = 0, token_idx: Optional[torch.Tensor] = None, ) -> Union[Tuple, CausalLMOutputWithPast]: """ @@ -398,7 +399,8 @@ def forward( ) hidden_states = outputs[0] - logits = self.lm_head(hidden_states) + # No upscaling to float was ever done for StableLm + logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]) loss = None if labels is not None: @@ -434,6 +436,7 @@ def prepare_inputs_for_generation( cache_position=None, position_ids=None, use_cache=True, + num_logits_to_keep=0, **kwargs, ): """ @@ -465,12 +468,16 @@ def prepare_inputs_for_generation( position_ids = torch.index_select(position_ids, 1, token_idx - 1) else: position_ids = position_ids[:, -input_ids.shape[1] :] + # This `clone` call is needed to avoid recapturing cuda graphs with `torch.compile`'s `mode="reduce-overhead`, as otherwise the input `position_ids` would have various stride during the decoding. Here, simply using `.contiguous()` is not sufficient as in the batch size = 1 case, `position_ids` is already contiguous but with varying stride which retriggers a capture. + position_ids = position_ids.clone(memory_format=torch.contiguous_format) # if `inputs_embeds` are passed, we only want to use them in the 1st generation step if inputs_embeds is not None and past_key_values is None: model_inputs = {"inputs_embeds": inputs_embeds} else: - model_inputs = {"input_ids": input_ids.contiguous()} # `contiguous()` needed for compilation use cases + model_inputs = { + "input_ids": input_ids.clone(memory_format=torch.contiguous_format) + } # `contiguous()` needed for compilation use cases model_inputs.update( { @@ -479,6 +486,7 @@ def prepare_inputs_for_generation( "past_key_values": past_key_values, "use_cache": use_cache, "attention_mask": attention_mask, + "num_logits_to_keep": num_logits_to_keep, "token_idx": token_idx, } ) diff --git a/optimum/habana/transformers/models/starcoder2/modeling_starcoder2.py b/optimum/habana/transformers/models/starcoder2/modeling_starcoder2.py index 36d5379e4f..c62b579c2a 100644 --- a/optimum/habana/transformers/models/starcoder2/modeling_starcoder2.py +++ b/optimum/habana/transformers/models/starcoder2/modeling_starcoder2.py @@ -33,7 +33,7 @@ Starcoder2Model, apply_rotary_pos_emb, ) -from transformers.utils import logging +from transformers.utils import is_torchdynamo_compiling, logging from ...modeling_attn_mask_utils import ( _gaudi_prepare_4d_causal_attention_mask, @@ -734,6 +734,7 @@ def forward( output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, + num_logits_to_keep: int = 0, token_idx: Optional[torch.Tensor] = None, trim_logits: Optional[bool] = False, attn_softmax_bf16: Optional[bool] = False, @@ -785,10 +786,18 @@ def forward( else: hidden_states = hidden_states[:, -1, :] - logits = self.lm_head(hidden_states).float() + if labels is None and not is_torchdynamo_compiling(): + logger.warning_once( + "Starting from v4.46, the `logits` model output will have the same type as the model (except at train time, where it will always be FP32)" + ) + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + # TODO: remove the float() operation in v4.46 + logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]).float() 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 so that tokens < n predict n shift_logits = logits[..., :-1, :].contiguous() shift_labels = labels[..., 1:].contiguous() @@ -821,6 +830,7 @@ def prepare_inputs_for_generation( cache_position=None, position_ids=None, use_cache=True, + num_logits_to_keep=0, token_idx=None, **kwargs, ): @@ -849,6 +859,8 @@ def prepare_inputs_for_generation( position_ids = torch.index_select(position_ids, 1, token_idx - 1) else: position_ids = position_ids[:, -input_ids.shape[1] :] + # This `clone` call is needed to avoid recapturing cuda graphs with `torch.compile`'s `mode="reduce-overhead`, as otherwise the input `position_ids` would have various stride during the decoding. Here, simply using `.contiguous()` is not sufficient as in the batch size = 1 case, `position_ids` is already contiguous but with varying stride which retriggers a capture. + position_ids = position_ids.clone(memory_format=torch.contiguous_format) cache_position = None @@ -856,7 +868,9 @@ def prepare_inputs_for_generation( if inputs_embeds is not None and past_key_values is None: model_inputs = {"inputs_embeds": inputs_embeds} else: - model_inputs = {"input_ids": input_ids.contiguous()} # `contiguous()` needed for compilation use cases + model_inputs = { + "input_ids": input_ids.clone(memory_format=torch.contiguous_format) + } # `contiguous()` needed for compilation use cases model_inputs.update( { @@ -865,6 +879,7 @@ def prepare_inputs_for_generation( "past_key_values": past_key_values, "use_cache": use_cache, "attention_mask": attention_mask, + "num_logits_to_keep": num_logits_to_keep, "token_idx": token_idx, "trim_logits": kwargs.get("trim_logits"), "attn_softmax_bf16": kwargs.get("attn_softmax_bf16"), diff --git a/optimum/habana/transformers/trainer.py b/optimum/habana/transformers/trainer.py index 5c418e66b7..946765f9f2 100644 --- a/optimum/habana/transformers/trainer.py +++ b/optimum/habana/transformers/trainer.py @@ -807,7 +807,7 @@ def hpu_deepspeed_checkpointing(function, *checkpoint_args): self.state = TrainerState.load_from_json(os.path.join(resume_from_checkpoint, TRAINER_STATE_NAME)) self.compare_trainer_and_checkpoint_args(self.args, self.state) self._load_callback_state() - epochs_trained = self.state.global_step // num_update_steps_per_epoch + epochs_trained = int(self.state.global_step // num_update_steps_per_epoch) if not args.ignore_data_skip: steps_trained_in_current_epoch = self.state.global_step % (num_update_steps_per_epoch) steps_trained_in_current_epoch *= args.gradient_accumulation_steps @@ -1058,7 +1058,7 @@ def hpu_deepspeed_checkpointing(function, *checkpoint_args): break if step < 0: logger.warning( - "There seems to be not a single sample in your epoch_iterator, stopping training at step" + "There seems not to be a single sample in your epoch_iterator, stopping training at step" f" {self.state.global_step}! This is expected if you're using an IterableDataset and set" f" num_steps ({max_steps}) higher than the number of available samples." ) @@ -1356,8 +1356,16 @@ def _save_checkpoint(self, model, trial, metrics=None): # Save the Trainer state if self.args.should_save: - # Update the `TrainerControl` state to where we are currently - self.state.stateful_callbacks["TrainerControl"] = self.control.state() + # Update `ExportableState` callbacks and `TrainerControl` state to where we are currently + for cb in [ + cb for cb in self.callback_handler.callbacks + [self.control] if isinstance(cb, ExportableState) + ]: + cb_name = cb.__class__.__name__ + cb_state = cb.state() + if isinstance(self.state.stateful_callbacks[cb_name], list): + self.state.stateful_callbacks[cb_name].append(cb_state) + else: + self.state.stateful_callbacks[cb_name] = cb_state self.state.save_to_json(os.path.join(output_dir, TRAINER_STATE_NAME)) if self.args.push_to_hub: @@ -2429,24 +2437,21 @@ def create_accelerator_and_postprocess(self): 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 - # post accelerator creation setup - # copy of https://github.com/huggingface/transformers/blob/b71f20a7c9f3716d30f6738501559acf863e2c5c/src/transformers/trainer.py#L3991 # post accelerator creation setup if self.is_fsdp_enabled: fsdp_plugin = self.accelerator.state.fsdp_plugin fsdp_plugin.limit_all_gathers = self.args.fsdp_config.get( "limit_all_gathers", fsdp_plugin.limit_all_gathers ) - if is_accelerate_available("0.23.0"): - fsdp_plugin.activation_checkpointing = self.args.fsdp_config.get( - "activation_checkpointing", fsdp_plugin.activation_checkpointing + fsdp_plugin.activation_checkpointing = self.args.fsdp_config.get( + "activation_checkpointing", fsdp_plugin.activation_checkpointing + ) + if fsdp_plugin.activation_checkpointing and self.args.gradient_checkpointing: + raise ValueError( + "The activation_checkpointing in FSDP config and the gradient_checkpointing in training arg " + "can't be set to True simultaneously. Please use FSDP's activation_checkpointing logic " + "when using FSDP." ) - if fsdp_plugin.activation_checkpointing and self.args.gradient_checkpointing: - raise ValueError( - "The activation_checkpointing in FSDP config and the gradient_checkpointing in training arg " - "can't be set to True simultaneously. Please use FSDP's activation_checkpointing logic " - "when using FSDP." - ) if self.is_deepspeed_enabled and getattr(self.args, "hf_deepspeed_config", None) is None: self.propagate_args_to_deepspeed() diff --git a/optimum/habana/transformers/trainer_seq2seq.py b/optimum/habana/transformers/trainer_seq2seq.py index 52977e30a0..734e73c80e 100644 --- a/optimum/habana/transformers/trainer_seq2seq.py +++ b/optimum/habana/transformers/trainer_seq2seq.py @@ -85,7 +85,7 @@ def load_generation_config(gen_config_arg: Union[str, GaudiGenerationConfig]) -> Loads a `~generation.GaudiGenerationConfig` from the `GaudiSeq2SeqTrainingArguments.generation_config` arguments. Args: - gen_config_arg (`str` or [`~generation.GaudiGenerationConfig`]): + gen_config_arg (`str` or [`~generation.GaudiGenerationConfig]`): `GaudiSeq2SeqTrainingArguments.generation_config` argument. Returns: diff --git a/optimum/habana/transformers/training_args.py b/optimum/habana/transformers/training_args.py index 5a65074fc9..3a71d46506 100644 --- a/optimum/habana/transformers/training_args.py +++ b/optimum/habana/transformers/training_args.py @@ -581,8 +581,8 @@ def __post_init__(self): " during training" ) - if not isinstance(self.warmup_steps, int) or self.warmup_steps < 0 or 0 < self.warmup_steps <= 1: - raise ValueError("warmup_steps must be either 0 or > 1") + if not isinstance(self.warmup_steps, int) or self.warmup_steps < 0: + raise ValueError("warmup_steps must be of type int and must be 0 or a positive integer.") # Copy of https://github.com/huggingface/transformers/blob/b71f20a7c9f3716d30f6738501559acf863e2c5c/src/transformers/training_args.py#L1563 # except following changes, (1) Remove XLA specific code & (2) change fsdp_backward_prefetch to backward_prefetch @@ -654,7 +654,7 @@ def __post_init__(self): self.fsdp_config["xla_fsdp_grad_ckpt"] = self.fsdp_config.get("xla_fsdp_grad_ckpt", False) # accelerate integration for FSDP - if len(self.fsdp) > 0 and not self.fsdp_config["xla"]: + if len(self.fsdp) > 0: os.environ["ACCELERATE_USE_FSDP"] = "true" from accelerate.utils.constants import ( FSDP_AUTO_WRAP_POLICY, diff --git a/setup.py b/setup.py index cea680353e..3f1a8c121a 100644 --- a/setup.py +++ b/setup.py @@ -29,7 +29,7 @@ INSTALL_REQUIRES = [ - "transformers >= 4.43.0, < 4.44.0", + "transformers @ git+https://github.com/huggingface/transformers.git", "optimum", "torch", "accelerate >= 0.33.0, < 0.34.0", diff --git a/tests/transformers/tests/generation/test_stopping_criteria.py b/tests/transformers/tests/generation/test_stopping_criteria.py index 0ce7838eee..9f177f9630 100644 --- a/tests/transformers/tests/generation/test_stopping_criteria.py +++ b/tests/transformers/tests/generation/test_stopping_criteria.py @@ -27,7 +27,6 @@ from transformers.generation import ( EosTokenCriteria, MaxLengthCriteria, - MaxNewTokensCriteria, MaxTimeCriteria, StoppingCriteriaList, validate_stopping_criteria, @@ -74,21 +73,6 @@ def test_max_length_criteria(self): input_ids, scores = self._get_tensors(10) self.assertTrue(all(criteria(input_ids, scores))) - def test_max_new_tokens_criteria(self): - criteria = MaxNewTokensCriteria(start_length=5, max_new_tokens=5) - - input_ids, scores = self._get_tensors(5) - self.assertFalse(all(criteria(input_ids, scores))) - - input_ids, scores = self._get_tensors(9) - self.assertFalse(all(criteria(input_ids, scores))) - - input_ids, scores = self._get_tensors(10) - self.assertTrue(all(criteria(input_ids, scores))) - - criteria_list = StoppingCriteriaList([criteria]) - self.assertEqual(criteria_list.max_length, 10) - def test_max_time_criteria(self): input_ids, scores = self._get_tensors(5) From 8eea643c3d70f624d29785139be82184a3a1d6ad Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Tue, 3 Sep 2024 07:59:22 +0000 Subject: [PATCH 02/89] Add specific commit in setup.py --- examples/stable-diffusion/unconditional_image_generation.py | 5 ----- .../habana/transformers/generation/candidate_generator.py | 2 +- .../habana/transformers/models/codegen/modeling_codegen.py | 2 +- optimum/habana/transformers/models/gptj/modeling_gptj.py | 2 +- setup.py | 2 +- 5 files changed, 4 insertions(+), 9 deletions(-) diff --git a/examples/stable-diffusion/unconditional_image_generation.py b/examples/stable-diffusion/unconditional_image_generation.py index 06bc6504c7..9b7442358f 100644 --- a/examples/stable-diffusion/unconditional_image_generation.py +++ b/examples/stable-diffusion/unconditional_image_generation.py @@ -19,13 +19,8 @@ def check_optimum_habana_min_version(*a, **b): return () -<<<<<<< HEAD check_min_version("4.45.0.dev0") -check_optimum_habana_min_version("1.10.4") -======= -check_min_version("4.43.0") check_optimum_habana_min_version("1.14.0.dev0") ->>>>>>> main # Setup logging logging.basicConfig( diff --git a/optimum/habana/transformers/generation/candidate_generator.py b/optimum/habana/transformers/generation/candidate_generator.py index 171161074f..6688553459 100644 --- a/optimum/habana/transformers/generation/candidate_generator.py +++ b/optimum/habana/transformers/generation/candidate_generator.py @@ -8,8 +8,8 @@ if TYPE_CHECKING: + from transformers.generation.logits_process import LogitsProcessorList from transformers.modeling_utils import PreTrainedModel - from transfromers.generation.logits_process import LogitsProcessorList from .configuration_utils import GaudiGenerationConfig diff --git a/optimum/habana/transformers/models/codegen/modeling_codegen.py b/optimum/habana/transformers/models/codegen/modeling_codegen.py index 80e1ce5710..a96192db3c 100644 --- a/optimum/habana/transformers/models/codegen/modeling_codegen.py +++ b/optimum/habana/transformers/models/codegen/modeling_codegen.py @@ -3,6 +3,7 @@ import torch import torch.utils.checkpoint from torch.nn import CrossEntropyLoss +from transformers.cache_utils import Cache from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from transformers.models.codegen.modeling_codegen import ( CodeGenAttention, @@ -10,7 +11,6 @@ apply_rotary_pos_emb, logger, ) -from transfromers.cache_utils import Cache class GaudiCodeGenAttention(CodeGenAttention): diff --git a/optimum/habana/transformers/models/gptj/modeling_gptj.py b/optimum/habana/transformers/models/gptj/modeling_gptj.py index 22b2b7a989..0415769d14 100644 --- a/optimum/habana/transformers/models/gptj/modeling_gptj.py +++ b/optimum/habana/transformers/models/gptj/modeling_gptj.py @@ -4,6 +4,7 @@ import torch from torch import nn from torch.nn import CrossEntropyLoss +from transformers.cache_utils import Cache from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from transformers.models.gptj.configuration_gptj import GPTJConfig from transformers.models.gptj.modeling_gptj import ( @@ -16,7 +17,6 @@ create_sinusoidal_positions, logger, ) -from transfroemrs.cache_utils import Cache class Matmul(nn.Module): diff --git a/setup.py b/setup.py index 3f1a8c121a..9baeffde67 100644 --- a/setup.py +++ b/setup.py @@ -29,7 +29,7 @@ INSTALL_REQUIRES = [ - "transformers @ git+https://github.com/huggingface/transformers.git", + "transformers @ git+https://github.com/huggingface/transformers.git@74e19e81e2a23809af192532b9b0e7ea202be6f2", "optimum", "torch", "accelerate >= 0.33.0, < 0.34.0", From a7be363a42ce505f82b8608c0eeab69be1756b35 Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Fri, 6 Sep 2024 15:32:10 +0000 Subject: [PATCH 03/89] Upgrade to commit e48e5f1f13e05380e24f4f31f5fee07aa6f959eb --- .../habana/transformers/generation/utils.py | 16 ++- .../models/llama/modeling_llama.py | 6 +- .../models/llava/modeling_llava.py | 7 ++ .../models/llava_next/modeling_llava_next.py | 10 ++ .../models/mistral/modeling_mistral.py | 6 +- .../models/mixtral/modeling_mixtral.py | 6 +- .../models/persimmon/modeling_persimmon.py | 6 +- .../transformers/models/phi/modeling_phi.py | 6 +- .../models/qwen2/modeling_qwen2.py | 6 +- .../models/stablelm/modeling_stablelm.py | 6 +- .../models/starcoder2/modeling_starcoder2.py | 6 +- optimum/habana/transformers/trainer.py | 22 ++-- setup.py | 2 +- tests/test_trainer.py | 113 ++++++++++++------ 14 files changed, 152 insertions(+), 66 deletions(-) diff --git a/optimum/habana/transformers/generation/utils.py b/optimum/habana/transformers/generation/utils.py index 89cc340dc3..cdc5ab5318 100755 --- a/optimum/habana/transformers/generation/utils.py +++ b/optimum/habana/transformers/generation/utils.py @@ -996,7 +996,16 @@ def generate( # TODO (joao): remove `user_defined_cache` after v4.47 (remove default conversion to legacy format) cache_name = "past_key_values" if "mamba" not in self.__class__.__name__.lower() else "cache_params" user_defined_cache = model_kwargs.get(cache_name) - self._prepare_cache_for_generation(generation_config, model_kwargs, assistant_model, batch_size, device) + max_cache_length = generation_config.max_length + if ( + inputs_tensor.shape[1] != input_ids_length + and model_input_name == "inputs_embeds" + and not self.config.is_encoder_decoder + ): + max_cache_length += inputs_tensor.shape[1] + self._prepare_cache_for_generation( + generation_config, model_kwargs, assistant_model, batch_size, max_cache_length, device + ) # determine whether introduce trim_logits feature model_kwargs["trim_logits"] = generation_config.trim_logits @@ -1108,8 +1117,8 @@ def generate( raise ValueError("assisted generate is only supported for batch_size = 1") if not model_kwargs["use_cache"]: raise ValueError("assisted generate requires `use_cache=True`") - if generation_config.cache_implementation == "static": - raise ValueError("assisted generate is not supported with `static_cache`") + if generation_config.cache_implementation in ["static", "hybrid", "sliding_window"]: + raise ValueError("assisted generate is not supported with Static cache classes`") if self._is_stateful: # In assisted generation we need the ability to confirm whether the model would pick certain tokens, # which is not possible with stateful models (they can't reset to a previous subset of generated text) @@ -3329,6 +3338,7 @@ def _assisted_decoding( # 1. Fetch candidate sequences from a `CandidateGenerator` 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) diff --git a/optimum/habana/transformers/models/llama/modeling_llama.py b/optimum/habana/transformers/models/llama/modeling_llama.py index 11e269e056..212d3b1dbe 100755 --- a/optimum/habana/transformers/models/llama/modeling_llama.py +++ b/optimum/habana/transformers/models/llama/modeling_llama.py @@ -1348,7 +1348,7 @@ def prepare_inputs_for_generation( cache_position=None, position_ids=None, use_cache=True, - num_logits_to_keep=0, + num_logits_to_keep=None, token_idx=None, **kwargs, ): @@ -1391,6 +1391,9 @@ def prepare_inputs_for_generation( else: model_inputs = {"input_ids": input_ids.clone(memory_format=torch.contiguous_format)} + if num_logits_to_keep is not None: + model_inputs["num_logits_to_keep"] = num_logits_to_keep + model_inputs.update( { "position_ids": position_ids, @@ -1398,7 +1401,6 @@ def prepare_inputs_for_generation( "past_key_values": past_key_values, "use_cache": use_cache, "attention_mask": attention_mask, - "num_logits_to_keep": num_logits_to_keep, "token_idx": token_idx, "trim_logits": kwargs.get("trim_logits"), "attn_softmax_bf16": kwargs.get("attn_softmax_bf16"), diff --git a/optimum/habana/transformers/models/llava/modeling_llava.py b/optimum/habana/transformers/models/llava/modeling_llava.py index 4300f6c7b3..9e718256d0 100644 --- a/optimum/habana/transformers/models/llava/modeling_llava.py +++ b/optimum/habana/transformers/models/llava/modeling_llava.py @@ -121,6 +121,7 @@ def forward( output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, + num_logits_to_keep: int = 0, token_idx: Optional[torch.Tensor] = None, image_offset: Optional[int] = None, tokens_pos: Optional[torch.LongTensor] = None, @@ -188,6 +189,7 @@ def forward( output_hidden_states=output_hidden_states, return_dict=return_dict, cache_position=cache_position, + num_logits_to_keep=num_logits_to_keep, token_idx=token_idx + image_offset, use_flash_attention=use_flash_attention, flash_attention_recompute=flash_attention_recompute, @@ -239,6 +241,7 @@ def prepare_inputs_for_generation( pixel_values=None, attention_mask=None, cache_position=None, + num_logits_to_keep=None, **kwargs, ): """ @@ -310,6 +313,10 @@ def prepare_inputs_for_generation( model_inputs = {"input_ids": input_ids} use_flash_attention = kwargs.get("use_flash_attention", False) flash_attention_recompute = kwargs.get("flash_attention_recompute", False) + + if num_logits_to_keep is not None: + model_inputs["num_logits_to_keep"] = num_logits_to_keep + model_inputs.update( { "position_ids": position_ids, 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 dca9e8d28a..8697acfdd6 100644 --- a/optimum/habana/transformers/models/llava_next/modeling_llava_next.py +++ b/optimum/habana/transformers/models/llava_next/modeling_llava_next.py @@ -54,6 +54,7 @@ def forward( output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, + num_logits_to_keep: int = 0, token_idx: Optional[torch.Tensor] = None, use_flash_attention: Optional[bool] = False, flash_attention_recompute: Optional[bool] = False, @@ -86,6 +87,7 @@ def forward( output_hidden_states=output_hidden_states, return_dict=return_dict, cache_position=cache_position, + num_logits_to_keep=num_logits_to_keep, token_idx=token_idx + self.image_offset, use_flash_attention=use_flash_attention, flash_attention_recompute=flash_attention_recompute, @@ -144,6 +146,8 @@ def forward( output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, + cache_position=cache_position, + num_logits_to_keep=num_logits_to_keep, ) # Copied from https://github.com/huggingface/transformers/blob/v4.40.0/src/transformers/models/llava_next/modeling_llava_next.py#L356 @@ -233,6 +237,7 @@ def prepare_inputs_for_generation( image_sizes=None, attention_mask=None, cache_position=None, + num_logits_to_keep=None, **kwargs, ): """ @@ -250,6 +255,8 @@ def prepare_inputs_for_generation( pixel_values=pixel_values, image_sizes=image_sizes, attention_mask=attention_mask, + cache_position=cache_position, + num_logits_to_keep=num_logits_to_keep, **kwargs, ) else: @@ -389,6 +396,9 @@ def prepare_inputs_for_generation( else: model_inputs = {"input_ids": input_ids} + if num_logits_to_keep is not None: + model_inputs["num_logits_to_keep"] = num_logits_to_keep + model_inputs.update( { "position_ids": position_ids, diff --git a/optimum/habana/transformers/models/mistral/modeling_mistral.py b/optimum/habana/transformers/models/mistral/modeling_mistral.py index 6ae4ede549..1684b2aee1 100644 --- a/optimum/habana/transformers/models/mistral/modeling_mistral.py +++ b/optimum/habana/transformers/models/mistral/modeling_mistral.py @@ -795,7 +795,7 @@ def prepare_inputs_for_generation( cache_position=None, position_ids=None, use_cache=True, - num_logits_to_keep=0, + num_logits_to_keep=None, **kwargs, ): """ @@ -843,6 +843,9 @@ def prepare_inputs_for_generation( else: model_inputs = {"input_ids": input_ids.contiguous()} # `contiguous()` needed for compilation use cases + if num_logits_to_keep is not None: + model_inputs["num_logits_to_keep"] = num_logits_to_keep + model_inputs.update( { "position_ids": position_ids, @@ -850,7 +853,6 @@ def prepare_inputs_for_generation( "past_key_values": past_key_values, "use_cache": use_cache, "attention_mask": attention_mask, - "num_logits_to_keep": num_logits_to_keep, "token_idx": token_idx, "reuse_cache": kwargs.get("reuse_cache"), "trim_logits": kwargs.get("trim_logits"), diff --git a/optimum/habana/transformers/models/mixtral/modeling_mixtral.py b/optimum/habana/transformers/models/mixtral/modeling_mixtral.py index a91444600f..9117cdc408 100644 --- a/optimum/habana/transformers/models/mixtral/modeling_mixtral.py +++ b/optimum/habana/transformers/models/mixtral/modeling_mixtral.py @@ -841,7 +841,7 @@ def prepare_inputs_for_generation( output_router_logits=False, position_ids=None, use_cache=True, - num_logits_to_keep=0, + num_logits_to_keep=None, **kwargs, ): reuse_cache = kwargs.get("reuse_cache") @@ -879,6 +879,9 @@ def prepare_inputs_for_generation( else: model_inputs = {"input_ids": input_ids.contiguous()} # `contiguous()` needed for compilation use cases + if num_logits_to_keep is not None: + model_inputs["num_logits_to_keep"] = num_logits_to_keep + model_inputs.update( { "position_ids": position_ids, @@ -887,7 +890,6 @@ def prepare_inputs_for_generation( "use_cache": use_cache, "attention_mask": attention_mask, "output_router_logits": output_router_logits, - "num_logits_to_keep": num_logits_to_keep, "token_idx": token_idx, "reuse_cache": reuse_cache, "flash_attention_recompute": kwargs.get("flash_attention_recompute"), diff --git a/optimum/habana/transformers/models/persimmon/modeling_persimmon.py b/optimum/habana/transformers/models/persimmon/modeling_persimmon.py index 8fb48d8f2c..2b5a842285 100644 --- a/optimum/habana/transformers/models/persimmon/modeling_persimmon.py +++ b/optimum/habana/transformers/models/persimmon/modeling_persimmon.py @@ -407,7 +407,7 @@ def prepare_inputs_for_generation( cache_position=None, position_ids=None, use_cache=True, - num_logits_to_keep=0, + num_logits_to_keep=None, **kwargs, ): """ @@ -450,6 +450,9 @@ def prepare_inputs_for_generation( "input_ids": input_ids.clone(memory_format=torch.contiguous_format) } # `contiguous()` needed for compilation use cases + if num_logits_to_keep is not None: + model_inputs["num_logits_to_keep"] = num_logits_to_keep + model_inputs.update( { "position_ids": position_ids, @@ -457,7 +460,6 @@ def prepare_inputs_for_generation( "past_key_values": past_key_values, "use_cache": use_cache, "attention_mask": attention_mask, - "num_logits_to_keep": num_logits_to_keep, "token_idx": token_idx, } ) diff --git a/optimum/habana/transformers/models/phi/modeling_phi.py b/optimum/habana/transformers/models/phi/modeling_phi.py index 81c56bec4f..f8b0d14181 100644 --- a/optimum/habana/transformers/models/phi/modeling_phi.py +++ b/optimum/habana/transformers/models/phi/modeling_phi.py @@ -620,7 +620,7 @@ def prepare_inputs_for_generation( cache_position=None, position_ids=None, use_cache=True, - num_logits_to_keep=0, + num_logits_to_keep=None, token_idx=None, **kwargs, ): @@ -669,6 +669,9 @@ def prepare_inputs_for_generation( "input_ids": input_ids.clone(memory_format=torch.contiguous_format) } # `contiguous()` needed for compilation use cases + if num_logits_to_keep is not None: + model_inputs["num_logits_to_keep"] = num_logits_to_keep + model_inputs.update( { "position_ids": position_ids, @@ -676,7 +679,6 @@ def prepare_inputs_for_generation( "past_key_values": past_key_values, "use_cache": use_cache, "attention_mask": attention_mask, - "num_logits_to_keep": num_logits_to_keep, "token_idx": token_idx, "reuse_cache": kwargs.get("reuse_cache"), "trim_logits": kwargs.get("trim_logits"), diff --git a/optimum/habana/transformers/models/qwen2/modeling_qwen2.py b/optimum/habana/transformers/models/qwen2/modeling_qwen2.py index 53ac275bdc..bf0ac2689e 100644 --- a/optimum/habana/transformers/models/qwen2/modeling_qwen2.py +++ b/optimum/habana/transformers/models/qwen2/modeling_qwen2.py @@ -862,7 +862,7 @@ def prepare_inputs_for_generation( cache_position=None, position_ids=None, use_cache=True, - num_logits_to_keep=0, + num_logits_to_keep=None, token_idx=None, **kwargs, ): @@ -906,6 +906,9 @@ def prepare_inputs_for_generation( "input_ids": input_ids.clone(memory_format=torch.contiguous_format) } # `contiguous()` needed for compilation use cases + if num_logits_to_keep is not None: + model_inputs["num_logits_to_keep"] = num_logits_to_keep + model_inputs.update( { "position_ids": position_ids.contiguous(), @@ -913,7 +916,6 @@ def prepare_inputs_for_generation( "past_key_values": past_key_values, "use_cache": use_cache, "attention_mask": attention_mask, - "num_logits_to_keep": num_logits_to_keep, "token_idx": token_idx, "trim_logits": kwargs.get("trim_logits"), "attn_softmax_bf16": kwargs.get("attn_softmax_bf16"), diff --git a/optimum/habana/transformers/models/stablelm/modeling_stablelm.py b/optimum/habana/transformers/models/stablelm/modeling_stablelm.py index aadd9469c1..6777760860 100644 --- a/optimum/habana/transformers/models/stablelm/modeling_stablelm.py +++ b/optimum/habana/transformers/models/stablelm/modeling_stablelm.py @@ -436,7 +436,7 @@ def prepare_inputs_for_generation( cache_position=None, position_ids=None, use_cache=True, - num_logits_to_keep=0, + num_logits_to_keep=None, **kwargs, ): """ @@ -479,6 +479,9 @@ def prepare_inputs_for_generation( "input_ids": input_ids.clone(memory_format=torch.contiguous_format) } # `contiguous()` needed for compilation use cases + if num_logits_to_keep is not None: + model_inputs["num_logits_to_keep"] = num_logits_to_keep + model_inputs.update( { "position_ids": position_ids, @@ -486,7 +489,6 @@ def prepare_inputs_for_generation( "past_key_values": past_key_values, "use_cache": use_cache, "attention_mask": attention_mask, - "num_logits_to_keep": num_logits_to_keep, "token_idx": token_idx, } ) diff --git a/optimum/habana/transformers/models/starcoder2/modeling_starcoder2.py b/optimum/habana/transformers/models/starcoder2/modeling_starcoder2.py index c62b579c2a..5cf2653055 100644 --- a/optimum/habana/transformers/models/starcoder2/modeling_starcoder2.py +++ b/optimum/habana/transformers/models/starcoder2/modeling_starcoder2.py @@ -830,7 +830,7 @@ def prepare_inputs_for_generation( cache_position=None, position_ids=None, use_cache=True, - num_logits_to_keep=0, + num_logits_to_keep=None, token_idx=None, **kwargs, ): @@ -872,6 +872,9 @@ def prepare_inputs_for_generation( "input_ids": input_ids.clone(memory_format=torch.contiguous_format) } # `contiguous()` needed for compilation use cases + if num_logits_to_keep is not None: + model_inputs["num_logits_to_keep"] = num_logits_to_keep + model_inputs.update( { "position_ids": position_ids.contiguous(), @@ -879,7 +882,6 @@ def prepare_inputs_for_generation( "past_key_values": past_key_values, "use_cache": use_cache, "attention_mask": attention_mask, - "num_logits_to_keep": num_logits_to_keep, "token_idx": token_idx, "trim_logits": kwargs.get("trim_logits"), "attn_softmax_bf16": kwargs.get("attn_softmax_bf16"), diff --git a/optimum/habana/transformers/trainer.py b/optimum/habana/transformers/trainer.py index 946765f9f2..f06f598658 100644 --- a/optimum/habana/transformers/trainer.py +++ b/optimum/habana/transformers/trainer.py @@ -482,7 +482,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.fp16_full_eval or args.bf16_full_eval) and not args.do_train: + if args.bf16_full_eval and not args.do_train and not self.is_model_parallel: self._move_model_to_device(self.model, args.device) if "model_path" in kwargs: @@ -675,11 +675,6 @@ def _inner_training_loop( # Activate gradient checkpointing if needed if args.gradient_checkpointing: - if args.gradient_checkpointing_kwargs is None: - gradient_checkpointing_kwargs = {} - else: - gradient_checkpointing_kwargs = args.gradient_checkpointing_kwargs - import transformers.modeling_utils if args.deepspeed: @@ -703,7 +698,7 @@ def hpu_deepspeed_checkpointing(function, *checkpoint_args): torch.utils.checkpoint.checkpoint = lazy_mode_checkpointing transformers.modeling_utils.checkpoint = lazy_mode_checkpointing - self.model.gradient_checkpointing_enable(gradient_checkpointing_kwargs=gradient_checkpointing_kwargs) + self.model.gradient_checkpointing_enable(gradient_checkpointing_kwargs=args.gradient_checkpointing_kwargs) # Wrap `_gradient_checkpointing_func` in the model with `transformer_engine` `activation_checkpointing` context. if self.accelerator.state.is_fp8_enabled: @@ -2465,10 +2460,15 @@ def create_accelerator_and_postprocess(self): wrapper = "DeepSpeed" if self.is_deepspeed_enabled else "FSDP" raise ValueError(f"{wrapper} can't be used with `save_only_model` along with `load_best_model_at_end`.") - # `auto_find_batch_size` isn't yet supported with DeepSpeed/FSDP - if (self.is_deepspeed_enabled or self.is_fsdp_enabled) and self.args.auto_find_batch_size: - wrapper = "DeepSpeed" if self.is_deepspeed_enabled else "FSDP" - raise NotImplementedError(f"`{wrapper}` doesn't support `auto_find_batch_size`.") + # `auto_find_batch_size` isn't supported yet with DeepSpeed Zero-3 + if ( + self.is_deepspeed_enabled + and self.accelerator.state.deepspeed_plugin.zero_stage == 3 + and self.args.auto_find_batch_size + ): + raise ValueError( + "`auto_find_batch_size` isn't supported yet with DeepSpeed Zero-3. Please consider using Zero-2, Zero-1, or FSDP" + ) def propagate_args_to_deepspeed(self, auto_find_batch_size=False): """ diff --git a/setup.py b/setup.py index 9baeffde67..8403dbba8a 100644 --- a/setup.py +++ b/setup.py @@ -29,7 +29,7 @@ INSTALL_REQUIRES = [ - "transformers @ git+https://github.com/huggingface/transformers.git@74e19e81e2a23809af192532b9b0e7ea202be6f2", + "transformers @ git+https://github.com/huggingface/transformers.git@e48e5f1f13e05380e24f4f31f5fee07aa6f959eb", "optimum", "torch", "accelerate >= 0.33.0, < 0.34.0", diff --git a/tests/test_trainer.py b/tests/test_trainer.py index ba78bbd2cc..2fca44db9a 100644 --- a/tests/test_trainer.py +++ b/tests/test_trainer.py @@ -108,6 +108,21 @@ adapt_transformers_to_gaudi() +class MockOOMCallback(TrainerCallback): + """ + Simple callback to simulate CUDA OOM error if + the batch size is >= to `batch_size_limit`. + """ + + def __init__(self, batch_size_limit=16): + self.batch_size_limit = batch_size_limit + + def on_step_end(self, args, state, control, **kwargs): + # simulate OOM on the first step + if state.train_batch_size >= self.batch_size_limit: + raise RuntimeError("Out of memory.") + + class RegressionDataset: def __init__(self, a=2, b=3, length=64, seed=42, label_names=None): np.random.seed(seed) @@ -1855,45 +1870,73 @@ def test_resume_training_with_randomness(self): self.assertAlmostEqual(a, a1, delta=1e-5) self.assertAlmostEqual(b, b1, delta=1e-5) - def test_auto_batch_size_with_resume_from_checkpoint(self): - train_dataset = RegressionDataset(length=128) + # @require_deepspeed + # def test_auto_batch_size_with_deepspeed(self): + # train_dataset = RegressionDataset(length=128) + + # config = RegressionModelConfig(a=0, b=2) + # model = RegressionRandomPreTrainedModel(config) + + # tmp_dir = self.get_auto_remove_tmp_dir() + + # for stage in [1, 2]: + # deepspeed = { + # "zero_optimization": { + # "stage": stage, + # }, + # "train_batch_size": "auto", + # "train_micro_batch_size_per_gpu": "auto", + # } + + # args = RegressionGaudiTrainingArguments( + # tmp_dir, + # do_train=True, + # max_steps=2, + # save_strategy="no", + # per_device_train_batch_size=16, + # auto_find_batch_size=True, + # deepspeed=deepspeed, + # use_habana=True, + # use_lazy_mode=True, + # ) + # 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) - config = RegressionModelConfig(a=0, b=2) - model = RegressionRandomPreTrainedModel(config) + # def test_auto_batch_size_with_resume_from_checkpoint(self): + # train_dataset = RegressionDataset(length=128) - tmp_dir = self.get_auto_remove_tmp_dir() + # config = RegressionModelConfig(a=0, b=2) + # model = RegressionRandomPreTrainedModel(config) - class MockCudaOOMCallback(TrainerCallback): - def on_step_end(self, args, state, control, **kwargs): - # simulate OOM on the first step - if state.train_batch_size >= 16: - raise RuntimeError("CUDA out of memory.") + # tmp_dir = self.get_auto_remove_tmp_dir() - args = RegressionGaudiTrainingArguments( - tmp_dir, - do_train=True, - max_steps=2, - save_steps=1, - per_device_train_batch_size=16, - auto_find_batch_size=True, - use_habana=True, - use_lazy_mode=True, - ) - gaudi_config = get_gaudi_config() - trainer = GaudiTrainer( - model, gaudi_config, args, train_dataset=train_dataset, callbacks=[MockCudaOOMCallback()] - ) - trainer.train() - # After `auto_find_batch_size` is ran we should now be at 8 - self.assertEqual(trainer._train_batch_size, 8) - - # 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) + # args = RegressionGaudiTrainingArguments( + # tmp_dir, + # do_train=True, + # max_steps=2, + # save_steps=1, + # per_device_train_batch_size=16, + # auto_find_batch_size=True, + # use_habana=True, + # use_lazy_mode=True, + # ) + # gaudi_config = get_gaudi_config() + # trainer = GaudiTrainer( + # 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) + + # # 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) # regression for this issue: https://github.com/huggingface/transformers/issues/12970 def test_training_with_resume_from_checkpoint_false(self): From d99f18f456901b145a195878182646b76e1159cb Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Mon, 9 Sep 2024 09:54:47 +0000 Subject: [PATCH 04/89] Fix default cache --- .../habana/transformers/generation/utils.py | 118 +++++++++++++++++- optimum/habana/transformers/modeling_utils.py | 3 + 2 files changed, 119 insertions(+), 2 deletions(-) diff --git a/optimum/habana/transformers/generation/utils.py b/optimum/habana/transformers/generation/utils.py index cdc5ab5318..87304a366a 100755 --- a/optimum/habana/transformers/generation/utils.py +++ b/optimum/habana/transformers/generation/utils.py @@ -22,7 +22,7 @@ import torch import torch.distributed as dist -from transformers.cache_utils import Cache, DynamicCache, EncoderDecoderCache +from transformers.cache_utils import Cache, DynamicCache, EncoderDecoderCache, OffloadedCache, QuantizedCacheConfig from transformers.generation.beam_constraints import DisjunctiveConstraint, PhrasalConstraint from transformers.generation.beam_search import BeamScorer, BeamSearchScorer, ConstrainedBeamSearchScorer from transformers.generation.candidate_generator import ( @@ -32,6 +32,7 @@ _prepare_attention_mask, _prepare_token_type_ids, ) +from transformers.generation.configuration_utils import NEED_SETUP_CACHE_CLASSES_MAPPING, QUANT_BACKEND_CLASSES_MAPPING from transformers.generation.logits_process import LogitsProcessorList from transformers.generation.stopping_criteria import ( EosTokenCriteria, @@ -57,7 +58,7 @@ ) from transformers.integrations.deepspeed import is_deepspeed_zero3_enabled from transformers.modeling_outputs import CausalLMOutputWithPast, Seq2SeqLMOutput -from transformers.utils import ModelOutput, is_torchdynamo_compiling +from transformers.utils import ModelOutput, is_hqq_available, is_quanto_available, is_torchdynamo_compiling from optimum.utils import logging @@ -672,6 +673,119 @@ def _prepare_generation_config( return generation_config, model_kwargs + def _prepare_cache_for_generation( + self, + generation_config: GaudiGenerationConfig, + model_kwargs: Dict, + assistant_model: "PreTrainedModel", + batch_size: int, + max_cache_length: int, + device: torch.device, + ) -> bool: + """ + Copied from: https://github.com/huggingface/transformers/blob/65bb28444849976f853063edb958b3ef3dd59d12/src/transformers/generation/utils.py#L1467 + + Changes: + - change the default from DynamicCache to tuples + """ + + cache_name = "past_key_values" if "mamba" not in self.__class__.__name__.lower() else "cache_params" + requires_cross_attention_cache = ( + self.config.is_encoder_decoder or model_kwargs.get("encoder_outputs") is not None + ) + + # Quick escape route 1: if the user specifies a cache, we only need to: + # a) check for conflicting `generate` arguments + # b) convert to the new cache format (if the user passes a legacy cache and model supports it) + user_defined_cache = model_kwargs.get(cache_name) + if user_defined_cache is not None: + if generation_config.cache_implementation is not None: + raise ValueError( + f"Passing both `cache_implementation` (used to initialize certain caches) and `{cache_name}` (a " + "Cache object) is unsupported. Please use only one of the two." + ) + if isinstance(user_defined_cache, tuple) and self._supports_default_dynamic_cache(): + model_kwargs[cache_name] = ( + DynamicCache.from_legacy_cache(user_defined_cache) + if not requires_cross_attention_cache + else EncoderDecoderCache.from_legacy_cache(user_defined_cache) + ) + return + + # Quick escape route 2: if the user specifies no cache is to be used. (conflicting arguments are handled in + # `generation_config.validate()`) + if generation_config.use_cache is False: + return + + # Quick escape route 3: model that only supports legacy caches = nothing to prepare + if not self._supports_default_dynamic_cache(): + if generation_config.cache_implementation is not None: + warnings.warn( + "This model does not support `Cache` instances, it only supports the legacy cache format (tuple " + f"of tuples). `cache_implementation` (set to {generation_config.cache_implementation}) will be " + "ignored.", + UserWarning, + ) + return + + # Otherwise we NEED to prepare a cache, based on `generation_config.cache_implementation` + + # TODO(joao): support static caches in assisted generation. assisted generation needs to roll back caches, + # which is only supported in dynamic caches atm + if assistant_model is not None and generation_config.cache_implementation is not None: + logger.warning_once( + "An assistant model is provided, using a dynamic cache instead of a cache of type=" + f"'{generation_config.cache_implementation}'." + ) + generation_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: + raise ValueError( + "This model does not support `cache_implementation='static'`. Please check the following " + "issue: https://github.com/huggingface/transformers/issues/28981" + ) + model_kwargs[cache_name] = self._get_cache( + cache_implementation=generation_config.cache_implementation, + batch_size=max(generation_config.num_beams, generation_config.num_return_sequences) * batch_size, + max_cache_len=max_cache_length, + device=device, + model_kwargs=model_kwargs, + ) + elif generation_config.cache_implementation == "quantized": + if not self._supports_quantized_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." + ) + + cache_config = ( + generation_config.cache_config + if generation_config.cache_config is not None + else QuantizedCacheConfig() + ) + cache_class = QUANT_BACKEND_CLASSES_MAPPING[cache_config.backend] + + if cache_config.backend == "quanto" and not is_quanto_available(): + raise ImportError( + "You need to install `quanto` in order to use KV cache quantization with quanto backend. " + "Please install it via with `pip install quanto`" + ) + 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) + elif generation_config.cache_implementation == "offloaded": + model_kwargs[cache_name] = OffloadedCache() + + # Use tuples by default (.i.e. legacy format). + else: + return + @torch.no_grad() def generate( self, diff --git a/optimum/habana/transformers/modeling_utils.py b/optimum/habana/transformers/modeling_utils.py index 8f4706c053..b9b09ec90f 100644 --- a/optimum/habana/transformers/modeling_utils.py +++ b/optimum/habana/transformers/modeling_utils.py @@ -263,6 +263,9 @@ def adapt_transformers_to_gaudi(): transformers.generation.GenerationMixin._contrastive_search = GaudiGenerationMixin._contrastive_search transformers.generation.GenerationMixin._assisted_decoding = GaudiGenerationMixin._assisted_decoding transformers.generation.GenerationMixin._get_candidate_generator = GaudiGenerationMixin._get_candidate_generator + transformers.generation.GenerationMixin._prepare_cache_for_generation = ( + GaudiGenerationMixin._prepare_cache_for_generation + ) transformers.generation.GenerationConfig = GaudiGenerationConfig transformers.modeling_utils.GenerationConfig = GaudiGenerationConfig transformers.generation.MaxLengthCriteria.__call__ = gaudi_MaxLengthCriteria_call From 47ad03c15599f09072b767ffb850df07cc4f556a Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Tue, 24 Sep 2024 22:00:39 +0000 Subject: [PATCH 05/89] Upgrade to commit 238b13478df209ab534f2195a397dc64a3930883 --- .../habana/transformers/generation/utils.py | 41 ++++++++++++++++--- .../models/bloom/modeling_bloom.py | 14 ++++--- .../models/falcon/modeling_falcon.py | 9 ++++ .../models/llama/modeling_llama.py | 12 +++--- .../models/llava/modeling_llava.py | 1 + .../models/persimmon/modeling_persimmon.py | 10 ++--- .../transformers/models/phi/modeling_phi.py | 8 ++-- .../models/stablelm/modeling_stablelm.py | 10 ++--- optimum/habana/transformers/trainer.py | 9 ++++ setup.py | 2 +- tests/test_trainer.py | 21 +++++++++- 11 files changed, 103 insertions(+), 34 deletions(-) diff --git a/optimum/habana/transformers/generation/utils.py b/optimum/habana/transformers/generation/utils.py index 1b27515197..f33f04b8d5 100755 --- a/optimum/habana/transformers/generation/utils.py +++ b/optimum/habana/transformers/generation/utils.py @@ -35,6 +35,7 @@ from transformers.generation.configuration_utils import NEED_SETUP_CACHE_CLASSES_MAPPING, QUANT_BACKEND_CLASSES_MAPPING from transformers.generation.logits_process import LogitsProcessorList from transformers.generation.stopping_criteria import ( + ConfidenceCriteria, EosTokenCriteria, MaxLengthCriteria, MaxTimeCriteria, @@ -540,6 +541,13 @@ def _get_stopping_criteria( criteria.append(StopStringCriteria(stop_strings=generation_config.stop_strings, tokenizer=tokenizer)) if not generation_config.ignore_eos and generation_config._eos_token_tensor is not None: criteria.append(EosTokenCriteria(eos_token_id=generation_config._eos_token_tensor)) + if ( + generation_config.assistant_confidence_threshold is not None + and generation_config.assistant_confidence_threshold > 0 + ): + criteria.append( + ConfidenceCriteria(assistant_confidence_threshold=generation_config.assistant_confidence_threshold) + ) criteria = self._merge_criteria_processor_list(criteria, stopping_criteria) return criteria @@ -620,23 +628,26 @@ def _prepare_generation_config( # the following conditions must be met # 1) the generation config must have been created from the model config (`_from_model_config` field); # 2) the generation config must have seen no modification since its creation (the hash is the same); - # 3) the user must have set generation parameters in the model config. + # 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) and self.generation_config._original_object_hash == hash(self.generation_config) # 2) + and len(self.config._get_non_default_generation_parameters()) > 0 # 3) ): new_generation_config = GaudiGenerationConfig.from_model_config(self.config) - if new_generation_config != self.generation_config: # 3) + if new_generation_config != self.generation_config: # 4) warnings.warn( "You have modified the pretrained model configuration to control generation. This is a" - " deprecated strategy to control generation and will be removed soon, in a future version." + " deprecated strategy to control generation and will be removed in v5." " Please use and modify the model generation configuration (see" - " https://huggingface.co/docs/transformers/generation_strategies#default-text-generation-configuration )" + " https://huggingface.co/docs/transformers/generation_strategies#default-text-generation-configuration )", + UserWarning, ) self.generation_config = new_generation_config - using_model_generation_config = True + generation_config = self.generation_config using_model_generation_config = True @@ -973,6 +984,10 @@ def generate( model_kwargs["attention_mask"] = self._prepare_attention_mask_for_generation( inputs_tensor, generation_config._pad_token_tensor, generation_config._eos_token_tensor ) + elif kwargs_has_attention_mask: + # TODO (joao): generalize this check with other types of inputs + if model_input_name == "input_ids" and len(model_kwargs["attention_mask"].shape) > 2: + raise ValueError("`attention_mask` passed to `generate` must be 2D.") is_greedy_or_beam_and_bucket = ( not generation_config.bucket_internal @@ -1695,6 +1710,15 @@ def _contrastive_search( unfinished_sequences = torch.ones(batch_size, dtype=torch.long, device=input_ids.device) model_kwargs["cache_position"] = torch.arange(cur_len, device=input_ids.device) + # Create cosine_matrix_mask based on the attention_mask + cosine_matrix_mask = torch.ones_like(input_ids, dtype=torch.long) + if self.config.is_encoder_decoder: + if "decoder_attention_mask" in model_kwargs and model_kwargs["decoder_attention_mask"] is not None: + cosine_matrix_mask = model_kwargs["decoder_attention_mask"] + else: + cosine_matrix_mask = model_kwargs["attention_mask"] + cosine_matrix_mask = cosine_matrix_mask.repeat_interleave(top_k, dim=0) + this_peer_finished = False hb_profer = HabanaProfile( @@ -1948,7 +1972,12 @@ def _contrastive_search( # compute the degeneration penalty and re-rank the candidates based on the degeneration penalty and the # model confidence. Keeping `selected_idx` on CPU enables multi-device contrastive search and doesn't # introduce (noticeable) slowdowns on single-device runs. - selected_idx = _ranking_fast(context_hidden, next_hidden, top_k_probs, penalty_alpha, top_k) + selected_idx = _ranking_fast( + context_hidden, next_hidden, top_k_probs, cosine_matrix_mask, penalty_alpha, top_k + ) + cosine_matrix_mask = torch.cat( + [cosine_matrix_mask, cosine_matrix_mask.new_ones((cosine_matrix_mask.shape[0], 1))], dim=-1 + ) # This will be used instead of the previous inneficient torch.stack(torch.split()) augmented_idx = torch.tensor( diff --git a/optimum/habana/transformers/models/bloom/modeling_bloom.py b/optimum/habana/transformers/models/bloom/modeling_bloom.py index 4bff984d82..c06d42e34d 100644 --- a/optimum/habana/transformers/models/bloom/modeling_bloom.py +++ b/optimum/habana/transformers/models/bloom/modeling_bloom.py @@ -164,8 +164,7 @@ def gaudi_bloom_attention_forward( present = None # [batch_size * num_heads, q_length, kv_length] - # we use `torch.Tensor.baddbmm` instead of `torch.baddbmm` as the latter isn't supported by TorchScript v1.11 - matmul_result = alibi.baddbmm( + attention_scores = alibi.baddbmm( batch1=query_layer, batch2=key_layer, beta=self.beta, @@ -173,7 +172,7 @@ def gaudi_bloom_attention_forward( ) # change view to [batch_size, num_heads, q_length, kv_length] - attention_scores = matmul_result.view(batch_size, self.num_heads, q_length, kv_length) + attention_scores = attention_scores.view(batch_size, self.num_heads, q_length, -1) # cast attention scores to fp32, compute scaled softmax and cast back to initial dtype - [batch_size, num_heads, q_length, kv_length] input_dtype = attention_scores.dtype @@ -187,7 +186,7 @@ def gaudi_bloom_attention_forward( attention_probs = attention_probs * head_mask # change view [batch_size x num_heads, q_length, kv_length] - attention_probs_reshaped = attention_probs.view(batch_size * self.num_heads, q_length, kv_length) + attention_probs_reshaped = attention_probs.view(batch_size * self.num_heads, q_length, -1) # matmul: [batch_size * num_heads, q_length, head_dim] context_layer = torch.bmm(attention_probs_reshaped, value_layer) @@ -507,9 +506,12 @@ def prepare_inputs_for_generation( # if `inputs_embeds` are passed, we only want to use them in the 1st generation step if inputs_embeds is not None and past_key_values is None: - model_inputs = {"inputs_embeds": inputs_embeds} + model_inputs = {"inputs_embeds": inputs_embeds, "input_ids": None} else: - model_inputs = {"input_ids": input_ids.contiguous()} # `contiguous()` needed for compilation use cases + # This `clone` call is needed to avoid recapturing cuda graphs with `torch.compile`'s `mode="reduce-overhead`, as otherwise the + # input `position_ids` would have various stride during the decoding. Here, simply using `.contiguous()` is not sufficient as in + # the batch size = 1 case, `position_ids` is already contiguous but with varying stride which retriggers a capture. + model_inputs = {"input_ids": input_ids.clone(memory_format=torch.contiguous_format), "inputs_embeds": None} model_inputs.update( { diff --git a/optimum/habana/transformers/models/falcon/modeling_falcon.py b/optimum/habana/transformers/models/falcon/modeling_falcon.py index c066fab951..0277668422 100644 --- a/optimum/habana/transformers/models/falcon/modeling_falcon.py +++ b/optimum/habana/transformers/models/falcon/modeling_falcon.py @@ -343,6 +343,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, # will become mandatory in v4.46 token_idx: Optional[torch.Tensor] = None, reuse_cache: Optional[bool] = False, cache_idx: int = None, @@ -632,6 +633,7 @@ def forward( 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 token_idx: Optional[torch.Tensor] = None, reuse_cache: Optional[bool] = False, cache_idx: int = None, @@ -658,6 +660,7 @@ def forward( use_cache=use_cache, output_attentions=output_attentions, cache_position=cache_position, + position_embeddings=position_embeddings, token_idx=token_idx, reuse_cache=reuse_cache, cache_idx=cache_idx, @@ -716,6 +719,7 @@ def pre_attn( 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 token_idx: Optional[torch.Tensor] = None, reuse_cache: Optional[bool] = False, cache_idx: int = None, @@ -741,6 +745,7 @@ def pre_attn( use_cache=use_cache, output_attentions=output_attentions, cache_position=cache_position, + position_embeddings=position_embeddings, token_idx=token_idx, reuse_cache=reuse_cache, cache_idx=cache_idx, @@ -905,6 +910,8 @@ def forward( # head_mask has shape n_layer x batch x num_heads x N x N head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers) + position_embeddings = None + for i, (block, layer_past) in enumerate(zip(self.h, past_key_values)): if output_hidden_states: all_hidden_states = all_hidden_states + (hidden_states,) @@ -921,6 +928,7 @@ def forward( use_cache, output_attentions, cache_position, + position_embeddings, None, use_flash_attention, flash_attention_recompute, @@ -937,6 +945,7 @@ def forward( 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, diff --git a/optimum/habana/transformers/models/llama/modeling_llama.py b/optimum/habana/transformers/models/llama/modeling_llama.py index 75f3ea1bc7..f59b048684 100755 --- a/optimum/habana/transformers/models/llama/modeling_llama.py +++ b/optimum/habana/transformers/models/llama/modeling_llama.py @@ -100,7 +100,7 @@ def __init__( if config is None: logger.warning_once( "`LlamaRotaryEmbedding` can now be fully parameterized by passing the model config through the " - "`config` argument. All other arguments will be removed in v4.45" + "`config` argument. All other arguments will be removed in v4.46" ) self.rope_kwargs = { "rope_type": rope_type, @@ -186,7 +186,7 @@ def forward(self, x, seq_len=None): class GaudiLlamaLinearScalingRotaryEmbedding(GaudiLlamaRotaryEmbedding): def __init__(self, *args, **kwargs): logger.warning_once( - "`LlamaLinearScalingRotaryEmbedding` is deprecated an will be removed in v4.45. Please use " + "`LlamaLinearScalingRotaryEmbedding` is deprecated an will be removed in v4.46. Please use " "`LlamaRotaryEmbedding`, which now also does linear scaling (simply pass the model config to __init__)." ) kwargs["rope_type"] = "linear" @@ -207,7 +207,7 @@ def _set_cos_sin_cache(self, seq_len, device, dtype): class GaudiLlamaDynamicNTKScalingRotaryEmbedding(GaudiLlamaRotaryEmbedding): def __init__(self, *args, **kwargs): logger.warning_once( - "`LlamaDynamicNTKScalingRotaryEmbedding` is deprecated an will be removed in v4.45. Please use " + "`LlamaDynamicNTKScalingRotaryEmbedding` is deprecated an will be removed in v4.46. Please use " "`LlamaRotaryEmbedding`, which now also does dynamic ntk scaling (simply pass the model config to " "__init__)." ) @@ -481,7 +481,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.46 token_idx: Optional[torch.Tensor] = None, attn_softmax_bf16: Optional[bool] = False, reuse_cache: Optional[bool] = False, @@ -563,7 +563,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.45 `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) @@ -830,7 +830,7 @@ def forward( 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.46 token_idx: Optional[torch.Tensor] = None, attn_softmax_bf16: Optional[bool] = False, reuse_cache: Optional[bool] = False, diff --git a/optimum/habana/transformers/models/llava/modeling_llava.py b/optimum/habana/transformers/models/llava/modeling_llava.py index 9e718256d0..d1f72896b9 100644 --- a/optimum/habana/transformers/models/llava/modeling_llava.py +++ b/optimum/habana/transformers/models/llava/modeling_llava.py @@ -214,6 +214,7 @@ def forward( 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, ) else: diff --git a/optimum/habana/transformers/models/persimmon/modeling_persimmon.py b/optimum/habana/transformers/models/persimmon/modeling_persimmon.py index 00cba0308c..c1fb019d66 100644 --- a/optimum/habana/transformers/models/persimmon/modeling_persimmon.py +++ b/optimum/habana/transformers/models/persimmon/modeling_persimmon.py @@ -68,12 +68,12 @@ def gaudi_persimmon_attention_forward( # Partial rotary embedding query_rot, query_pass = ( - query_states[..., : self.rotary_emb.dim], - query_states[..., self.rotary_emb.dim :], + query_states[..., : self.rotary_ndims], + query_states[..., self.rotary_ndims :], ) key_rot, key_pass = ( - key_states[..., : self.rotary_emb.dim], - key_states[..., self.rotary_emb.dim :], + key_states[..., : self.rotary_ndims], + key_states[..., self.rotary_ndims :], ) # [batch_size, seq_length, num_heads, head_dim // config.partial_rotary_factor] query_rot, key_rot = apply_rotary_pos_emb(query_rot, key_rot, cos, sin, position_ids) @@ -97,7 +97,7 @@ def gaudi_persimmon_attention_forward( cache_kwargs = { "sin": sin, "cos": cos, - "partial_rotation_size": self.rotary_emb.dim, + "partial_rotation_size": self.rotary_ndims, "cache_position": cache_position, } key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) diff --git a/optimum/habana/transformers/models/phi/modeling_phi.py b/optimum/habana/transformers/models/phi/modeling_phi.py index ac12454730..53a4b1f73a 100644 --- a/optimum/habana/transformers/models/phi/modeling_phi.py +++ b/optimum/habana/transformers/models/phi/modeling_phi.py @@ -202,12 +202,12 @@ def forward( # Partial rotary embedding query_rot, query_pass = ( - query_states[..., : self.rotary_emb.dim], - query_states[..., self.rotary_emb.dim :], + query_states[..., : self.rotary_ndims], + query_states[..., self.rotary_ndims :], ) key_rot, key_pass = ( - key_states[..., : self.rotary_emb.dim], - key_states[..., self.rotary_emb.dim :], + key_states[..., : self.rotary_ndims], + key_states[..., self.rotary_ndims :], ) # [batch_size, seq_length, num_heads, head_dim // config.partial_rotary_factor] query_rot, key_rot = apply_rotary_pos_emb(query_rot, key_rot, cos, sin, position_ids) diff --git a/optimum/habana/transformers/models/stablelm/modeling_stablelm.py b/optimum/habana/transformers/models/stablelm/modeling_stablelm.py index 2cf91e1906..22eca3c9da 100644 --- a/optimum/habana/transformers/models/stablelm/modeling_stablelm.py +++ b/optimum/habana/transformers/models/stablelm/modeling_stablelm.py @@ -72,12 +72,12 @@ def gaudi_stablelm_attention_forward( # Partial rotary embedding query_rot, query_pass = ( - query_states[..., : self.rotary_emb.dim], - query_states[..., self.rotary_emb.dim :], + query_states[..., : self.rotary_ndims], + query_states[..., self.rotary_ndims :], ) key_rot, key_pass = ( - key_states[..., : self.rotary_emb.dim], - key_states[..., self.rotary_emb.dim :], + key_states[..., : self.rotary_ndims], + key_states[..., self.rotary_ndims :], ) # [batch_size, seq_length, num_heads, head_dim // config.partial_rotary_factor] query_rot, key_rot = apply_rotary_pos_emb(query_rot, key_rot, cos, sin, position_ids) @@ -101,7 +101,7 @@ def gaudi_stablelm_attention_forward( cache_kwargs = { "sin": sin, "cos": cos, - "partial_rotation_size": self.rotary_emb.dim, + "partial_rotation_size": self.rotary_ndims, "cache_position": cache_position, } key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) diff --git a/optimum/habana/transformers/trainer.py b/optimum/habana/transformers/trainer.py index fa79aa9556..e6406abe63 100644 --- a/optimum/habana/transformers/trainer.py +++ b/optimum/habana/transformers/trainer.py @@ -1036,6 +1036,8 @@ def hpu_deepspeed_checkpointing(function, *checkpoint_args, use_reentrant: Optio args.max_grad_norm, ) + self.control = self.callback_handler.on_pre_optimizer_step(args, self.state, self.control) + optimizer_was_run = True self.optimizer.step() @@ -1582,6 +1584,9 @@ def training_step(self, model: torch.nn.Module, inputs: Dict[str, Union[torch.Te `torch.Tensor`: The tensor with training loss on this batch. """ model.train() + if hasattr(self.optimizer, "train") and callable(self.optimizer.train): + self.optimizer.train() + inputs = self._prepare_inputs(inputs) with self.compute_loss_context_manager(): @@ -1819,6 +1824,8 @@ def evaluation_loop( self.deepspeed = self.model_wrapped model.eval() + if hasattr(self.optimizer, "eval") and callable(self.optimizer.eval): + self.optimizer.eval() # Do not use HPU graphs if the training is ongoing because it detaches gradients if args.use_hpu_graphs_for_inference and not self.is_in_train: @@ -2226,6 +2233,8 @@ def prediction_loop( if self.is_deepspeed_enabled: self.deepspeed = self.model_wrapped model.eval() + if hasattr(self.optimizer, "eval") and callable(self.optimizer.eval): + self.optimizer.eval() # Do not use HPU graphs if the training is ongoing because it detaches gradients if args.use_hpu_graphs_for_inference and not self.is_in_train: diff --git a/setup.py b/setup.py index 8403dbba8a..e0da79f728 100644 --- a/setup.py +++ b/setup.py @@ -29,7 +29,7 @@ INSTALL_REQUIRES = [ - "transformers @ git+https://github.com/huggingface/transformers.git@e48e5f1f13e05380e24f4f31f5fee07aa6f959eb", + "transformers @ git+https://github.com/huggingface/transformers.git@238b13478df209ab534f2195a397dc64a3930883", "optimum", "torch", "accelerate >= 0.33.0, < 0.34.0", diff --git a/tests/test_trainer.py b/tests/test_trainer.py index 2fca44db9a..eddb82b500 100644 --- a/tests/test_trainer.py +++ b/tests/test_trainer.py @@ -28,7 +28,7 @@ from typing import Dict, List, Optional, Union import numpy as np -from huggingface_hub import HfFolder, ModelCard, delete_repo, list_repo_commits, list_repo_files +from huggingface_hub import HfFolder, ModelCard, create_branch, delete_repo, list_repo_commits, list_repo_files from parameterized import parameterized from pytest import mark from requests.exceptions import HTTPError @@ -2946,6 +2946,25 @@ def test_push_to_hub_tags(self): model_card = ModelCard.load(repo_name) self.assertTrue("test-trainer-tags" in model_card.data.tags) + def test_push_to_hub_with_revision(self): + # Checks if `trainer.push_to_hub()` works correctly by adding revision + with tempfile.TemporaryDirectory() as tmp_dir: + trainer = get_regression_trainer( + output_dir=os.path.join(tmp_dir, "test-trainer-revision"), + push_to_hub=True, + hub_token=self._token, + ) + branch = "v1.0" + create_branch(repo_id=trainer.hub_model_id, branch=branch, token=self._token, exist_ok=True) + url = trainer.push_to_hub(revision=branch) + + # Extract branch from the url + re_search = re.search(r"tree/([^/]+)/", url) + self.assertIsNotNone(re_search) + + branch_name = re_search.groups()[0] + self.assertEqual(branch_name, branch) + @require_torch @require_optuna From 94c23ba8f520b2376ad8cc9b39d54cfd02413dfb Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Tue, 24 Sep 2024 22:11:10 +0000 Subject: [PATCH 06/89] Fix --- optimum/habana/transformers/trainer.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/optimum/habana/transformers/trainer.py b/optimum/habana/transformers/trainer.py index e6406abe63..aafb4e19e5 100644 --- a/optimum/habana/transformers/trainer.py +++ b/optimum/habana/transformers/trainer.py @@ -1584,8 +1584,8 @@ def training_step(self, model: torch.nn.Module, inputs: Dict[str, Union[torch.Te `torch.Tensor`: The tensor with training loss on this batch. """ model.train() - if hasattr(self.optimizer, "train") and callable(self.optimizer.train): - self.optimizer.train() + # if hasattr(self.optimizer, "train") and callable(self.optimizer.train): + # self.optimizer.train() inputs = self._prepare_inputs(inputs) @@ -1824,8 +1824,8 @@ def evaluation_loop( self.deepspeed = self.model_wrapped model.eval() - if hasattr(self.optimizer, "eval") and callable(self.optimizer.eval): - self.optimizer.eval() + # if hasattr(self.optimizer, "eval") and callable(self.optimizer.eval): + # self.optimizer.eval() # Do not use HPU graphs if the training is ongoing because it detaches gradients if args.use_hpu_graphs_for_inference and not self.is_in_train: @@ -2233,8 +2233,8 @@ def prediction_loop( if self.is_deepspeed_enabled: self.deepspeed = self.model_wrapped model.eval() - if hasattr(self.optimizer, "eval") and callable(self.optimizer.eval): - self.optimizer.eval() + # if hasattr(self.optimizer, "eval") and callable(self.optimizer.eval): + # self.optimizer.eval() # Do not use HPU graphs if the training is ongoing because it detaches gradients if args.use_hpu_graphs_for_inference and not self.is_in_train: From c19dedd2b2f730f469c09c47e60ec2441b1ca4ed Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Wed, 25 Sep 2024 18:50:40 +0000 Subject: [PATCH 07/89] Upgrade to v4.45.0 --- .../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 +- .../run_multitask_prompt_tuning.py | 2 +- .../run_prompt_tuning_clm.py | 2 +- examples/question-answering/run_qa.py | 2 +- examples/question-answering/run_seq2seq_qa.py | 2 +- .../run_speech_recognition_ctc.py | 2 +- .../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 +- .../habana/transformers/generation/utils.py | 12 ++++--- setup.py | 2 +- .../example_diff/run_audio_classification.txt | 20 +++++------ tests/example_diff/run_clip.txt | 16 ++++----- tests/example_diff/run_clm.txt | 32 ++++++++--------- tests/example_diff/run_glue.txt | 24 ++++++------- .../example_diff/run_image_classification.txt | 12 +++---- tests/example_diff/run_mlm.txt | 28 +++++++-------- tests/example_diff/run_qa.txt | 18 +++++----- tests/example_diff/run_seq2seq_qa.txt | 18 +++++----- .../run_speech_recognition_ctc.txt | 14 ++++---- .../run_speech_recognition_seq2seq.txt | 14 ++++---- tests/example_diff/run_summarization.txt | 36 +++++++++---------- tests/example_diff/run_translation.txt | 18 +++++----- 30 files changed, 150 insertions(+), 146 deletions(-) diff --git a/examples/audio-classification/run_audio_classification.py b/examples/audio-classification/run_audio_classification.py index 4feca220b5..9a23428866 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.45.0.dev0") +check_min_version("4.45.0") check_optimum_habana_min_version("1.14.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 f12bd91a5f..42b9e8a468 100644 --- a/examples/contrastive-image-text/run_bridgetower.py +++ b/examples/contrastive-image-text/run_bridgetower.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.45.0.dev0") +check_min_version("4.45.0") check_optimum_habana_min_version("1.14.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 a8621ffa1c..6a8ca235e1 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.45.0.dev0") +check_min_version("4.45.0") check_optimum_habana_min_version("1.14.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 4f0a830282..b2694665a3 100644 --- a/examples/image-classification/run_image_classification.py +++ b/examples/image-classification/run_image_classification.py @@ -63,7 +63,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.45.0.dev0") +check_min_version("4.45.0") check_optimum_habana_min_version("1.14.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 7282f179ab..5a8d25b0ed 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.45.0.dev0") +check_min_version("4.45.0") check_optimum_habana_min_version("1.14.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 30d9a9b3a7..30315bfc84 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.45.0.dev0") +check_min_version("4.45.0") check_optimum_habana_min_version("1.14.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_multitask_prompt_tuning.py b/examples/language-modeling/run_multitask_prompt_tuning.py index 84153d4a80..1d81bcc496 100644 --- a/examples/language-modeling/run_multitask_prompt_tuning.py +++ b/examples/language-modeling/run_multitask_prompt_tuning.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 risk. -check_min_version("4.45.0.dev0") +check_min_version("4.45.0") check_optimum_habana_min_version("1.14.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 03b98eccfc..e263c0c1b6 100644 --- a/examples/language-modeling/run_prompt_tuning_clm.py +++ b/examples/language-modeling/run_prompt_tuning_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.45.0.dev0") +check_min_version("4.45.0") check_optimum_habana_min_version("1.14.0.dev0") require_version("datasets>=1.8.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 1a6d4db2de..d22949c076 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.45.0.dev0") +check_min_version("4.45.0") check_optimum_habana_min_version("1.14.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 3065a98103..1f045552bd 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.45.0.dev0") +check_min_version("4.45.0") check_optimum_habana_min_version("1.14.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 26f4f8f3c7..83865556d1 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.45.0.dev0") +check_min_version("4.45.0") check_optimum_habana_min_version("1.14.0.dev0") require_version("datasets>=1.18.0", "To fix: pip install -r examples/pytorch/speech-recognition/requirements.txt") diff --git a/examples/speech-recognition/run_speech_recognition_seq2seq.py b/examples/speech-recognition/run_speech_recognition_seq2seq.py index e49410043e..ff9702e80c 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.45.0.dev0") +check_min_version("4.45.0") check_optimum_habana_min_version("1.14.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 9b7442358f..baca71b6ba 100644 --- 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.45.0.dev0") +check_min_version("4.45.0") check_optimum_habana_min_version("1.14.0.dev0") # Setup logging diff --git a/examples/summarization/run_summarization.py b/examples/summarization/run_summarization.py index 632fe5d430..8715c4e75f 100755 --- a/examples/summarization/run_summarization.py +++ b/examples/summarization/run_summarization.py @@ -65,7 +65,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.45.0.dev0") +check_min_version("4.45.0") check_optimum_habana_min_version("1.14.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 bdc227a8f3..57bf7cbb05 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.45.0.dev0") +check_min_version("4.45.0") check_optimum_habana_min_version("1.14.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 4cb7a89598..c2def132a7 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.45.0.dev0") +check_min_version("4.45.0") check_optimum_habana_min_version("1.14.0.dev0") require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/translation/requirements.txt") diff --git a/optimum/habana/transformers/generation/utils.py b/optimum/habana/transformers/generation/utils.py index f33f04b8d5..d88a45d67e 100755 --- a/optimum/habana/transformers/generation/utils.py +++ b/optimum/habana/transformers/generation/utils.py @@ -542,7 +542,8 @@ def _get_stopping_criteria( if not generation_config.ignore_eos and generation_config._eos_token_tensor is not None: criteria.append(EosTokenCriteria(eos_token_id=generation_config._eos_token_tensor)) if ( - generation_config.assistant_confidence_threshold is not None + generation_config.is_assistant + and generation_config.assistant_confidence_threshold is not None and generation_config.assistant_confidence_threshold > 0 ): criteria.append( @@ -1934,7 +1935,7 @@ def _contrastive_search( model_kwargs["past_key_values"].crop(-1) all_outputs.append(outputs) - outputs = stack_model_outputs(all_outputs) + outputs = stack_model_outputs(all_outputs, self.config.get_text_config()) else: # compute the candidate tokens by the language model and collect their hidden_states @@ -2772,13 +2773,16 @@ def expand_if_needed(tensor, new_size, value, dim=-1): ) inputs_per_sub_batches = _split_model_inputs( - model_inputs, split_size=batch_size, full_batch_size=batch_beam_size + 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) + 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( diff --git a/setup.py b/setup.py index e0da79f728..37c16d8e2f 100644 --- a/setup.py +++ b/setup.py @@ -29,7 +29,7 @@ INSTALL_REQUIRES = [ - "transformers @ git+https://github.com/huggingface/transformers.git@238b13478df209ab534f2195a397dc64a3930883", + "transformers >= 4.45.0, < 4.46.0", "optimum", "torch", "accelerate >= 0.33.0, < 0.34.0", diff --git a/tests/example_diff/run_audio_classification.txt b/tests/example_diff/run_audio_classification.txt index 5e98ce8248..238cad957b 100644 --- a/tests/example_diff/run_audio_classification.txt +++ b/tests/example_diff/run_audio_classification.txt @@ -2,7 +2,7 @@ < import warnings 28,29d26 < from datasets import DatasetDict, load_dataset -< +< 31,39c28,29 < from transformers import ( < AutoConfig, @@ -19,18 +19,18 @@ 43a34,44 > from optimum.habana import GaudiConfig, GaudiTrainer, GaudiTrainingArguments > from optimum.habana.utils import set_seed -> -> +> +> > try: > from optimum.habana.utils import check_optimum_habana_min_version > except ImportError: -> +> > def check_optimum_habana_min_version(*a, **b): > return () -> +> 47,48c48,50 < # Will error if the minimal version of Transformers is not installed. Remove at your own risks. -< check_min_version("4.45.0.dev0") +< check_min_version("4.45.0") --- > # Will error if the minimal version of Transformers and Optimum Habana are not installed. Remove at your own risks. > check_min_version("4.43.0") @@ -54,7 +54,7 @@ < "should not be used in combination with `--freeze_feature_encoder`. " < "Only make use of `--freeze_feature_encoder`." < ) -< +< 203c187 < parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)) --- @@ -66,7 +66,7 @@ > revision=model_args.model_revision, > token=model_args.token, > ) -> +> 232a224 > mixed_precision = training_args.bf16 or gaudi_config.use_torch_autocast 234,235c226,228 @@ -79,9 +79,9 @@ 304a298,300 > # Max input length > max_length = int(round(feature_extractor.sampling_rate * data_args.max_length_seconds)) -> +> 309a306 -> +> 315c312,318 < inputs = feature_extractor(subsampled_wavs, sampling_rate=feature_extractor.sampling_rate) --- diff --git a/tests/example_diff/run_clip.txt b/tests/example_diff/run_clip.txt index f57b3b3240..7cd5a15451 100644 --- a/tests/example_diff/run_clip.txt +++ b/tests/example_diff/run_clip.txt @@ -1,11 +1,11 @@ 18d17 -< +< 32a32 > import transformers 33a34 > from habana_dataloader_trainer import HabanaDataloaderTrainer 38,39d38 -< +< < import transformers 45,47d43 < Trainer, @@ -14,18 +14,18 @@ 52a49,59 > from optimum.habana import GaudiConfig, GaudiTrainer, GaudiTrainingArguments > from optimum.habana.utils import set_seed -> -> +> +> > try: > from optimum.habana.utils import check_optimum_habana_min_version > except ImportError: -> +> > def check_optimum_habana_min_version(*a, **b): > return () -> +> 56,57c63,65 < # Will error if the minimal version of Transformers is not installed. Remove at your own risks. -< check_min_version("4.45.0.dev0") +< check_min_version("4.45.0") --- > # Will error if the minimal version of Transformers and Optimum Habana are not installed. Remove at your own risks. > check_min_version("4.43.0") @@ -45,7 +45,7 @@ > revision=model_args.model_revision, > token=model_args.token, > ) -> +> 269a288 > mixed_precision = training_args.bf16 or gaudi_config.use_torch_autocast 271,272c290,292 diff --git a/tests/example_diff/run_clm.txt b/tests/example_diff/run_clm.txt index 580f3c9684..daf04e96df 100644 --- a/tests/example_diff/run_clm.txt +++ b/tests/example_diff/run_clm.txt @@ -4,14 +4,14 @@ > # Copyright 2022 The HuggingFace Inc. team. All rights reserved. 17,19c17,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 < from datasets import load_dataset -< +< 37a35 > from datasets import load_dataset 45,46d42 @@ -25,24 +25,24 @@ > from optimum.habana.utils import set_seed 57,58d52 < # Will error if the minimal version of Transformers is not installed. Remove at your own risks. -< check_min_version("4.45.0.dev0") +< check_min_version("4.45.0") 60c54,60 < require_version("datasets>=2.14.0", "To fix: pip install -r examples/pytorch/language-modeling/requirements.txt") --- > try: > from optimum.habana.utils import check_optimum_habana_min_version > except ImportError: -> +> > def check_optimum_habana_min_version(*a, **b): > return () -> +> 63a64,69 > # Will error if the minimal version of Transformers and Optimum Habana are not installed. Remove at your own risks. > check_min_version("4.43.0") > check_optimum_habana_min_version("1.14.0.dev0") -> +> > require_version("datasets>=2.14.0", "To fix: pip install -r examples/pytorch/language-modeling/requirements.txt") -> +> 79c85,86 < "The model checkpoint for weights initialization. Don't set if you want to train a model from scratch." --- @@ -65,7 +65,7 @@ 195c211,212 < streaming: bool = field(default=False, metadata={"help": "Enable streaming mode"}) --- -> +> > streaming: bool = field(default=False, metadata={"help": "Enable streaming mode."}) 221a239,241 > save_last_ckpt: bool = field( @@ -82,7 +82,7 @@ > revision=model_args.model_revision, > token=model_args.token, > ) -> +> 273a301 > mixed_precision = training_args.bf16 or gaudi_config.use_torch_autocast 275,276c303,305 @@ -95,12 +95,12 @@ 390a420 > "use_cache": False if training_args.gradient_checkpointing else model_args.use_cache, 486a517 -> +> 550a582,585 -> +> > def tensor_mapper(x): > return {i: torch.tensor(x[i], dtype=torch.int32) for i in x} -> +> 553a589,590 > if training_args.resume_from_checkpoint is not None and training_args.resume_from_checkpoint != "": > train_dataset = train_dataset.map(tensor_mapper) @@ -137,7 +137,7 @@ > ) > metrics["train_samples"] = min(max_train_samples, len(train_dataset)) 622d661 -< +< 625,626c664,669 < 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)) @@ -147,10 +147,10 @@ > 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)) -> +> 649,653d691 -< -< +< +< < def _mp_fn(index): < # For xla_spawn (TPUs) < main() diff --git a/tests/example_diff/run_glue.txt b/tests/example_diff/run_glue.txt index 26d2e245c0..f969aa8923 100644 --- a/tests/example_diff/run_glue.txt +++ b/tests/example_diff/run_glue.txt @@ -1,6 +1,6 @@ 29,30d28 < from datasets import load_dataset -< +< 31a30 > from datasets import load_dataset 40,41d38 @@ -11,27 +11,27 @@ 48a45,54 > from optimum.habana import GaudiConfig, GaudiTrainer, GaudiTrainingArguments > from optimum.habana.utils import set_seed -> -> +> +> > try: > from optimum.habana.utils import check_optimum_habana_min_version > except ImportError: -> +> > def check_optimum_habana_min_version(*a, **b): > return () 50,51c56,61 < # Will error if the minimal version of Transformers is not installed. Remove at your own risks. -< check_min_version("4.45.0.dev0") +< check_min_version("4.45.0") --- -> +> > 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.43.0") > check_optimum_habana_min_version("1.14.0.dev0") 67,68d76 < logger = logging.getLogger(__name__) -< +< 143a152,155 > problem_type: Optional[str] = field( > default="single_label_classification", @@ -53,7 +53,7 @@ > revision=model_args.model_revision, > token=model_args.token, > ) -> +> 251a275 > mixed_precision = training_args.bf16 or gaudi_config.use_torch_autocast 253,254c277,279 @@ -70,7 +70,7 @@ > 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 < trainer = Trainer( --- @@ -78,8 +78,8 @@ 529a561 > gaudi_config=gaudi_config, 629,633d660 -< -< +< +< < def _mp_fn(index): < # For xla_spawn (TPUs) < main() diff --git a/tests/example_diff/run_image_classification.txt b/tests/example_diff/run_image_classification.txt index 7a3e696fd6..d353d75f4c 100644 --- a/tests/example_diff/run_image_classification.txt +++ b/tests/example_diff/run_image_classification.txt @@ -4,7 +4,7 @@ 24a27 > import transformers 37,38d39 -< +< < import transformers 45,47d45 < Trainer, @@ -13,19 +13,19 @@ 52a51,60 > from optimum.habana import GaudiConfig, GaudiTrainer, GaudiTrainingArguments > from optimum.habana.utils import set_seed -> -> +> +> > try: > from optimum.habana.utils import check_optimum_habana_min_version > except ImportError: -> +> > def check_optimum_habana_min_version(*a, **b): > return () 54d61 < """ Fine-tuning a 🤗 Transformers model for image classification""" 58,59c65,67 < # Will error if the minimal version of Transformers is not installed. Remove at your own risks. -< check_min_version("4.45.0.dev0") +< check_min_version("4.45.0") --- > # Will error if the minimal version of Transformers and Optimum Habana are not installed. Remove at your own risks. > check_min_version("4.43.0") @@ -41,7 +41,7 @@ > revision=model_args.model_revision, > token=model_args.token, > ) -> +> 213a229 > mixed_precision = training_args.bf16 or gaudi_config.use_torch_autocast 215,216c231,233 diff --git a/tests/example_diff/run_mlm.txt b/tests/example_diff/run_mlm.txt index a3e97b56c7..d87d2fe4c0 100644 --- a/tests/example_diff/run_mlm.txt +++ b/tests/example_diff/run_mlm.txt @@ -1,13 +1,13 @@ 17,19c17,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 < from datasets import load_dataset -< +< 37a35 > from datasets import load_dataset 46,49d43 @@ -20,26 +20,26 @@ > from optimum.habana.utils import set_seed 56,57d51 < # Will error if the minimal version of Transformers is not installed. Remove at your own risks. -< check_min_version("4.45.0.dev0") +< check_min_version("4.45.0") 59c53,59 < require_version("datasets>=2.14.0", "To fix: pip install -r examples/pytorch/language-modeling/requirements.txt") --- > try: > from optimum.habana.utils import check_optimum_habana_min_version > except ImportError: -> +> > def check_optimum_habana_min_version(*a, **b): > return () -> +> 61a62,69 -> +> > # Will error if the minimal version of Transformers and Optimum Habana are not installed. Remove at your own risks. > check_min_version("4.43.0") > check_optimum_habana_min_version("1.14.0.dev0") -> +> > require_version("datasets>=2.14.0", "To fix: pip install -r examples/pytorch/language-modeling/requirements.txt") -> -> +> +> 137c145 < "choices": ["auto", "bfloat16", "float16", "float32"], --- @@ -63,7 +63,7 @@ > revision=model_args.model_revision, > token=model_args.token, > ) -> +> 284a300 > mixed_precision = training_args.bf16 or gaudi_config.use_torch_autocast 286,287c302,304 @@ -103,7 +103,7 @@ > ) > metrics["train_samples"] = min(max_train_samples, len(train_dataset)) 656d673 -< +< 659,660c676,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)) @@ -113,10 +113,10 @@ > 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)) -> +> 683,687d703 -< -< +< +< < def _mp_fn(index): < # For xla_spawn (TPUs) < main() diff --git a/tests/example_diff/run_qa.txt b/tests/example_diff/run_qa.txt index 4d289c5faa..ce15c1c30f 100644 --- a/tests/example_diff/run_qa.txt +++ b/tests/example_diff/run_qa.txt @@ -6,7 +6,7 @@ > import transformers 32,34d32 < from utils_qa import postprocess_qa_predictions -< +< < import transformers 43d40 < TrainingArguments, @@ -19,24 +19,24 @@ > from optimum.habana.utils import set_seed 52,53d50 < # Will error if the minimal version of Transformers is not installed. Remove at your own risks. -< check_min_version("4.45.0.dev0") +< check_min_version("4.45.0") 55c52,58 < require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/question-answering/requirements.txt") --- > try: > from optimum.habana.utils import check_optimum_habana_min_version > except ImportError: -> +> > def check_optimum_habana_min_version(*a, **b): > return () -> +> 58a62,67 > # Will error if the minimal version of Transformers and Optimum Habana are not installed. Remove at your own risks. > check_min_version("4.43.0") > check_optimum_habana_min_version("1.14.0.dev0") -> +> > require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/question-answering/requirements.txt") -> +> 146c155 < " batching to the maximum length in the batch (which can be faster on GPU but will be slower on TPU)." --- @@ -52,7 +52,7 @@ > revision=model_args.model_revision, > token=model_args.token, > ) -> +> 263a280 > mixed_precision = training_args.bf16 or gaudi_config.use_torch_autocast 265,266c282,284 @@ -70,8 +70,8 @@ 638a661 > gaudi_config=gaudi_config, 707,711d729 -< -< +< +< < def _mp_fn(index): < # For xla_spawn (TPUs) < main() diff --git a/tests/example_diff/run_seq2seq_qa.txt b/tests/example_diff/run_seq2seq_qa.txt index 96bcd84b82..7f1a733850 100644 --- a/tests/example_diff/run_seq2seq_qa.txt +++ b/tests/example_diff/run_seq2seq_qa.txt @@ -1,7 +1,7 @@ 29a30 > import transformers 32,33d32 -< +< < import transformers 40,41d38 < Seq2SeqTrainingArguments, @@ -11,24 +11,24 @@ > from optimum.habana.utils import set_seed 48,49d46 < # Will error if the minimal version of Transformers is not installed. Remove at your own risks. -< check_min_version("4.45.0.dev0") +< check_min_version("4.45.0") 51c48,54 < require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/question-answering/requirements.txt") --- > try: > from optimum.habana.utils import check_optimum_habana_min_version > except ImportError: -> +> > def check_optimum_habana_min_version(*a, **b): > return () -> +> 54a58,63 > # Will error if the minimal version of Transformers and Optimum Habana are not installed. Remove at your own risks. > check_min_version("4.43.0") > check_optimum_habana_min_version("1.14.0.dev0") -> +> > require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/question-answering/requirements.txt") -> +> 178c187 < " batching to the maximum length in the batch (which can be faster on GPU but will be slower on TPU)." --- @@ -44,7 +44,7 @@ > revision=model_args.model_revision, > token=model_args.token, > ) -> +> 308a325 > mixed_precision = training_args.bf16 or gaudi_config.use_torch_autocast 310,311c327,329 @@ -57,8 +57,8 @@ 661a680 > gaudi_config=gaudi_config, 735,739d753 -< -< +< +< < def _mp_fn(index): < # For xla_spawn (TPUs) < main() diff --git a/tests/example_diff/run_speech_recognition_ctc.txt b/tests/example_diff/run_speech_recognition_ctc.txt index d9bb9d115e..71f9665cfe 100644 --- a/tests/example_diff/run_speech_recognition_ctc.txt +++ b/tests/example_diff/run_speech_recognition_ctc.txt @@ -1,6 +1,6 @@ 32,33d31 < from datasets import DatasetDict, load_dataset -< +< 34a33 > from datasets import DatasetDict, load_dataset 42,43d40 @@ -13,29 +13,29 @@ > from optimum.habana.utils import set_seed 52,53d49 < # Will error if the minimal version of Transformers is not installed. Remove at your own risks. -< check_min_version("4.45.0.dev0") +< check_min_version("4.45.0") 55c51,56 < require_version("datasets>=1.18.0", "To fix: pip install -r examples/pytorch/speech-recognition/requirements.txt") --- > try: > from optimum.habana.utils import check_optimum_habana_min_version > except ImportError: -> +> > def check_optimum_habana_min_version(*a, **b): > return () 59a61,66 > # Will error if the minimal version of Transformers and Optimum Habana are not installed. Remove at your own risks. > check_min_version("4.43.0") > check_optimum_habana_min_version("1.14.0.dev0") -> +> > require_version("datasets>=1.18.0", "To fix: pip install -r examples/pytorch/speech-recognition/requirements.txt") -> +> 144c151 < "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 -< +< 400c406 < parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)) --- @@ -46,7 +46,7 @@ > cache_dir=model_args.cache_dir, > token=data_args.token, > ) -> +> 435a448 > mixed_precision = training_args.bf16 or gaudi_config.use_torch_autocast 437,438c450,452 diff --git a/tests/example_diff/run_speech_recognition_seq2seq.txt b/tests/example_diff/run_speech_recognition_seq2seq.txt index 0fce8cc3e0..40f4ab43dc 100644 --- a/tests/example_diff/run_speech_recognition_seq2seq.txt +++ b/tests/example_diff/run_speech_recognition_seq2seq.txt @@ -1,6 +1,6 @@ 31,32d30 < from datasets import DatasetDict, load_dataset -< +< 33a32 > from datasets import DatasetDict, load_dataset 41,43d39 @@ -10,17 +10,17 @@ 48a45,55 > from optimum.habana import GaudiConfig, GaudiSeq2SeqTrainer, GaudiSeq2SeqTrainingArguments > from optimum.habana.utils import set_seed -> -> +> +> > try: > from optimum.habana.utils import check_optimum_habana_min_version > except ImportError: -> +> > def check_optimum_habana_min_version(*a, **b): > return () -> +> 51c58,59 -< check_min_version("4.45.0.dev0") +< check_min_version("4.45.0") --- > check_min_version("4.43.0") > check_optimum_habana_min_version("1.14.0.dev0") @@ -49,7 +49,7 @@ > cache_dir=model_args.cache_dir, > token=model_args.token, > ) -> +> 310a334 > mixed_precision = training_args.bf16 or gaudi_config.use_torch_autocast 312,313c336,338 diff --git a/tests/example_diff/run_summarization.txt b/tests/example_diff/run_summarization.txt index aaa348da39..c9fc832ff7 100644 --- a/tests/example_diff/run_summarization.txt +++ b/tests/example_diff/run_summarization.txt @@ -8,7 +8,7 @@ > import torch > import transformers 33,34d35 -< +< < import transformers 45,47c46 < Seq2SeqTrainer, @@ -23,27 +23,27 @@ > from optimum.habana.utils import set_seed 54,55d55 < # Will error if the minimal version of Transformers is not installed. Remove at your own risks. -< check_min_version("4.45.0.dev0") +< check_min_version("4.45.0") 57c57,63 < require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/summarization/requirements.txt") --- > try: > from optimum.habana.utils import check_optimum_habana_min_version > except ImportError: -> +> > def check_optimum_habana_min_version(*a, **b): > return () -> +> 60a67,72 > # Will error if the minimal version of Transformers and Optimum Habana are not installed. Remove at your own risks. > check_min_version("4.43.0") > check_optimum_habana_min_version("1.14.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 > use_cache: bool = field( > default=True, @@ -71,7 +71,7 @@ > revision=model_args.model_revision, > token=model_args.token, > ) -> +> 347a379 > mixed_precision = training_args.bf16 or gaudi_config.use_torch_autocast 349,350c381,383 @@ -89,14 +89,14 @@ > raise ValueError( > "Training is not yet supported for BART. Eval or predict can be enabled with `--do_eval` and `--do_predict`." > ) -> +> 454c494,501 < embedding_size = model.get_input_embeddings().weight.shape[0] --- > embeddings = model.get_input_embeddings() > if is_deepspeed_zero3_enabled(): > import deepspeed -> +> > with deepspeed.zero.GatheredParameters(embeddings.weight, modifier_rank=None): > embedding_size = embeddings.weight.shape[0] > else: @@ -113,7 +113,7 @@ 575a626,665 > def preprocess_bucketing_function(examples): > # remove pairs where at least one record is None -> +> > inputs, targets = [], [] > for i in range(len(examples[text_column])): > if examples[text_column][i] and examples[summary_column][i]: @@ -121,7 +121,7 @@ > targets.append(examples[summary_column][i]) > else: > raise ValueError("Found case where either text or summary is missing.") -> +> > inputs = [prefix + inp + suffix for inp in inputs] > model_inputs = tokenizer(inputs, return_tensors="pt", padding=True) > new_model_inputs = {"input_ids": []} @@ -140,24 +140,24 @@ > model_inputs = new_model_inputs > # Tokenize targets with the `text_target` keyword argument > labels = tokenizer(text_target=targets, max_length=max_target_length, padding=padding, truncation=True) -> +> > # If we are padding here, replace all tokenizer.pad_token_id in the labels by -100 when we want to ignore > # padding in the loss. > if padding == "max_length" and data_args.ignore_pad_token_for_loss: > labels["input_ids"] = [ > [(l if l != tokenizer.pad_token_id else -100) for l in label] for label in labels["input_ids"] > ] -> +> > model_inputs["labels"] = labels["input_ids"] > return model_inputs -> +> 590a681,686 > def wrapper_preprocess_function(examples): > if model.config.is_encoder_decoder: > return preprocess_bucketing_function(examples) > else: > return preprocess_function(examples) -> +> 599c695 < preprocess_function, --- @@ -212,8 +212,8 @@ 676a780 > gaudi_config=gaudi_config, 765,769d868 -< -< +< +< < def _mp_fn(index): < # For xla_spawn (TPUs) < main() diff --git a/tests/example_diff/run_translation.txt b/tests/example_diff/run_translation.txt index 95f2749242..5d06e5c2f6 100644 --- a/tests/example_diff/run_translation.txt +++ b/tests/example_diff/run_translation.txt @@ -1,6 +1,6 @@ 30,31d29 < from datasets import load_dataset -< +< 32a31 > from datasets import load_dataset 44,45c43 @@ -15,24 +15,24 @@ > from optimum.habana.utils import set_seed 54,55d52 < # Will error if the minimal version of Transformers is not installed. Remove at your own risks. -< check_min_version("4.45.0.dev0") +< check_min_version("4.45.0") 57c54,60 < require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/translation/requirements.txt") --- > try: > from optimum.habana.utils import check_optimum_habana_min_version > except ImportError: -> +> > def check_optimum_habana_min_version(*a, **b): > return () -> +> 60a64,69 > # Will error if the minimal version of Transformers and Optimum Habana are not installed. Remove at your own risks. > check_min_version("4.43.0") > check_optimum_habana_min_version("1.14.0.dev0") -> +> > require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/translation/requirements.txt") -> +> 62c71,78 < MULTILINGUAL_TOKENIZERS = [MBartTokenizer, MBartTokenizerFast, MBart50Tokenizer, MBart50TokenizerFast, M2M100Tokenizer] --- @@ -69,7 +69,7 @@ > revision=model_args.model_revision, > token=model_args.token, > ) -> +> 296a329 > mixed_precision = training_args.bf16 or gaudi_config.use_torch_autocast 298,299c331,333 @@ -92,8 +92,8 @@ 596a632 > gaudi_config=gaudi_config, 689,693d724 -< -< +< +< < def _mp_fn(index): < # For xla_spawn (TPUs) < main() From fc399fa41c478ee335fa72ff336c572f947d5c58 Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Wed, 25 Sep 2024 19:42:38 +0000 Subject: [PATCH 08/89] Fix --- optimum/habana/transformers/modeling_utils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/optimum/habana/transformers/modeling_utils.py b/optimum/habana/transformers/modeling_utils.py index 9e4c4c648b..0aa2720d61 100644 --- a/optimum/habana/transformers/modeling_utils.py +++ b/optimum/habana/transformers/modeling_utils.py @@ -274,6 +274,7 @@ def adapt_transformers_to_gaudi(): GaudiGenerationMixin._prepare_cache_for_generation ) transformers.generation.GenerationConfig = GaudiGenerationConfig + transformers.generation.configuration_utils.GenerationConfig = GaudiGenerationConfig transformers.modeling_utils.GenerationConfig = GaudiGenerationConfig transformers.generation.MaxLengthCriteria.__call__ = gaudi_MaxLengthCriteria_call transformers.generation.MaxTimeCriteria.__call__ = gaudi_MaxTimeCriteria_call From 921615966494bff7085e093c34571150647bf939 Mon Sep 17 00:00:00 2001 From: Jimin Ha Date: Thu, 26 Sep 2024 00:51:29 -0700 Subject: [PATCH 09/89] Add bias to gptj (#1363) --- optimum/habana/transformers/models/gptj/modeling_gptj.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/optimum/habana/transformers/models/gptj/modeling_gptj.py b/optimum/habana/transformers/models/gptj/modeling_gptj.py index b7f6951427..3927e1feb9 100644 --- a/optimum/habana/transformers/models/gptj/modeling_gptj.py +++ b/optimum/habana/transformers/models/gptj/modeling_gptj.py @@ -73,6 +73,14 @@ def __init__(self, config: GPTJConfig, layer_idx=None): super().__init__(config) self.config = config + max_positions = config.max_position_embeddings + self.register_buffer( + "bias", + torch.tril(torch.ones((max_positions, max_positions), dtype=torch.bool)).view( + 1, 1, max_positions, max_positions + ), + persistent=False, + ) self.matmul_qk = Matmul() self.matmul_av = Matmul() self.k_cache = KVCache() From 679365abfcfbe91e37a16cd3473531eafd7b2771 Mon Sep 17 00:00:00 2001 From: Shiv Kaul Date: Thu, 26 Sep 2024 00:54:48 -0700 Subject: [PATCH 10/89] Switch roberta from sdpa to eager attn (#1361) --- optimum/habana/transformers/models/modeling_all_models.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/optimum/habana/transformers/models/modeling_all_models.py b/optimum/habana/transformers/models/modeling_all_models.py index c9eb95524e..90aa2d5e0f 100644 --- a/optimum/habana/transformers/models/modeling_all_models.py +++ b/optimum/habana/transformers/models/modeling_all_models.py @@ -115,7 +115,7 @@ def gaudi_conv1d_forward(self, x): @classmethod def gaudi_check_and_enable_sdpa(cls, config, hard_check_only: bool = False) -> PretrainedConfig: # This model doesn't support SDPA in Gaudi yet, fallback to original code. - MODELS_ATTN_IMPLEMENTATION_EAGER = ["bart", "gpt_bigcode", "mistral", "mixtral", "wav2vec2"] + MODELS_ATTN_IMPLEMENTATION_EAGER = ["bart", "gpt_bigcode", "mistral", "mixtral", "wav2vec2", "roberta"] if config.model_type in MODELS_ATTN_IMPLEMENTATION_EAGER: config._attn_implementation = "eager" From 1abd6ee0a244367601e0dee3718e2b30301d551d Mon Sep 17 00:00:00 2001 From: Yeonsil Yoon Date: Thu, 26 Sep 2024 00:56:04 -0700 Subject: [PATCH 11/89] Update bloom attention forward reshape follwing the transformer change (#1360) --- .../habana/transformers/models/bloom/modeling_bloom.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/optimum/habana/transformers/models/bloom/modeling_bloom.py b/optimum/habana/transformers/models/bloom/modeling_bloom.py index c06d42e34d..5b0a770451 100644 --- a/optimum/habana/transformers/models/bloom/modeling_bloom.py +++ b/optimum/habana/transformers/models/bloom/modeling_bloom.py @@ -137,11 +137,9 @@ def gaudi_bloom_attention_forward( # 3 x [batch_size, num_heads, seq_length, head_dim] query_layer, key_layer, value_layer = self._reshape(fused_qkv) - batch_size, q_length, _, _ = query_layer.shape - - query_layer = query_layer.transpose(1, 2).reshape(batch_size * self.num_heads, q_length, self.head_dim) - key_layer = key_layer.permute(0, 2, 3, 1).reshape(batch_size * self.num_heads, self.head_dim, q_length) - value_layer = value_layer.transpose(1, 2).reshape(batch_size * self.num_heads, q_length, self.head_dim) + query_layer = query_layer.reshape(batch_size * self.num_heads, -1, self.head_dim) + key_layer = key_layer.reshape(batch_size * self.num_heads, -1, self.head_dim).transpose(1, 2) + value_layer = value_layer.reshape(batch_size * self.num_heads, -1, self.head_dim) # Collapse views to improve performance on HPU query_layer = query_layer.contiguous() From 8043d2cef69edc9eae6c7282bbb7fa41f268e5b6 Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Thu, 26 Sep 2024 09:57:14 +0000 Subject: [PATCH 12/89] Workaround for Llava/Llava-next --- optimum/habana/transformers/models/llava/modeling_llava.py | 3 ++- .../transformers/models/llava_next/modeling_llava_next.py | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/optimum/habana/transformers/models/llava/modeling_llava.py b/optimum/habana/transformers/models/llava/modeling_llava.py index d1f72896b9..cccbf8ebb9 100644 --- a/optimum/habana/transformers/models/llava/modeling_llava.py +++ b/optimum/habana/transformers/models/llava/modeling_llava.py @@ -189,7 +189,8 @@ def forward( output_hidden_states=output_hidden_states, return_dict=return_dict, cache_position=cache_position, - num_logits_to_keep=num_logits_to_keep, + # TODO: from Transformers v4.45, `generate` sets `num_logits_to_keep` to 1 if not given, which we don't want here + # num_logits_to_keep=num_logits_to_keep, token_idx=token_idx + image_offset, use_flash_attention=use_flash_attention, flash_attention_recompute=flash_attention_recompute, 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 8697acfdd6..6cf728d014 100644 --- a/optimum/habana/transformers/models/llava_next/modeling_llava_next.py +++ b/optimum/habana/transformers/models/llava_next/modeling_llava_next.py @@ -87,7 +87,8 @@ def forward( output_hidden_states=output_hidden_states, return_dict=return_dict, cache_position=cache_position, - num_logits_to_keep=num_logits_to_keep, + # TODO: from Transformers v4.45, `generate` sets `num_logits_to_keep` to 1 if not given, which we don't want here + # num_logits_to_keep=num_logits_to_keep, token_idx=token_idx + self.image_offset, use_flash_attention=use_flash_attention, flash_attention_recompute=flash_attention_recompute, From 047e7ffc81f8346f85ebba66b2ba8ca1b6086c69 Mon Sep 17 00:00:00 2001 From: Harish Subramony <81822986+hsubramony@users.noreply.github.com> Date: Sat, 28 Sep 2024 04:11:36 -0700 Subject: [PATCH 13/89] Fix reshape error in mamba (#1369) --- optimum/habana/transformers/models/mamba/modeling_mamba.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/optimum/habana/transformers/models/mamba/modeling_mamba.py b/optimum/habana/transformers/models/mamba/modeling_mamba.py index 8a7af8c914..e23ce65dd8 100644 --- a/optimum/habana/transformers/models/mamba/modeling_mamba.py +++ b/optimum/habana/transformers/models/mamba/modeling_mamba.py @@ -75,6 +75,8 @@ def gaudi_MambaForCausalLM_prepare_inputs_for_generation( else: idx = token_idx + kwargs.get("inputs_embeds_offset", 0) - 1 input_ids = torch.index_select(input_ids, 1, idx) + if attention_mask is not None: + attention_mask = None else: if token_idx is not None: input_ids = torch.index_select(input_ids, 1, torch.arange(token_idx_cpu, device=input_ids.device)) From 1b8a3f7347d8d497e79b58aba5a15128816714b3 Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Tue, 1 Oct 2024 13:03:22 +0000 Subject: [PATCH 14/89] Fix contrastive search --- .../habana/transformers/generation/utils.py | 41 +++++++++++-------- 1 file changed, 25 insertions(+), 16 deletions(-) diff --git a/optimum/habana/transformers/generation/utils.py b/optimum/habana/transformers/generation/utils.py index efdd6d6126..25c454b4d5 100755 --- a/optimum/habana/transformers/generation/utils.py +++ b/optimum/habana/transformers/generation/utils.py @@ -52,7 +52,6 @@ GenerateOutput, GenerationMixin, GenerationMode, - _ranking_fast, _split_model_inputs, _split_model_outputs, stack_model_outputs, @@ -1733,15 +1732,6 @@ def _contrastive_search( unfinished_sequences = torch.ones(batch_size, dtype=torch.long, device=input_ids.device) model_kwargs["cache_position"] = torch.arange(cur_len, device=input_ids.device) - # Create cosine_matrix_mask based on the attention_mask - cosine_matrix_mask = torch.ones_like(input_ids, dtype=torch.long) - if self.config.is_encoder_decoder: - if "decoder_attention_mask" in model_kwargs and model_kwargs["decoder_attention_mask"] is not None: - cosine_matrix_mask = model_kwargs["decoder_attention_mask"] - else: - cosine_matrix_mask = model_kwargs["attention_mask"] - cosine_matrix_mask = cosine_matrix_mask.repeat_interleave(top_k, dim=0) - this_peer_finished = False hb_profer = HabanaProfile( @@ -1996,12 +1986,7 @@ def _contrastive_search( # compute the degeneration penalty and re-rank the candidates based on the degeneration penalty and the # model confidence. Keeping `selected_idx` on CPU enables multi-device contrastive search and doesn't # introduce (noticeable) slowdowns on single-device runs. - selected_idx = _ranking_fast( - context_hidden, next_hidden, top_k_probs, cosine_matrix_mask, penalty_alpha, top_k - ) - cosine_matrix_mask = torch.cat( - [cosine_matrix_mask, cosine_matrix_mask.new_ones((cosine_matrix_mask.shape[0], 1))], dim=-1 - ) + selected_idx = _ranking_fast(context_hidden, next_hidden, top_k_probs, penalty_alpha, top_k) # This will be used instead of the previous inneficient torch.stack(torch.split()) augmented_idx = torch.tensor( @@ -3810,3 +3795,27 @@ def _assisted_decoding( ) else: return input_ids + + +def _ranking_fast( + context_hidden: torch.FloatTensor, + next_hidden: torch.FloatTensor, + next_top_k_probs: torch.FloatTensor, + alpha: float, + beam_width: int, +) -> torch.FloatTensor: + """ + Reranks the top_k candidates based on a degeneration penalty (cosine similarity with previous tokens), as described + in the paper "A Contrastive Framework for Neural Text Generation". Returns the index of the best candidate for each + row in the batch. + """ + norm_context_hidden = context_hidden / context_hidden.norm(dim=2, keepdim=True) + norm_next_hidden = next_hidden / next_hidden.norm(dim=2, keepdim=True) + cosine_matrix = torch.matmul(norm_context_hidden, norm_next_hidden.transpose(1, 2)).squeeze(-1) # [B*K, S] + + degeneration_penalty, _ = torch.max(cosine_matrix, dim=-1) # [B*K] + next_top_k_probs = next_top_k_probs.view(-1) # [B*K] + contrastive_score = (1.0 - alpha) * next_top_k_probs - alpha * degeneration_penalty + contrastive_score = torch.stack(torch.split(contrastive_score, beam_width)) # [B, K] + _, selected_idx = contrastive_score.max(dim=-1) # [B] + return selected_idx From 2332afbac80215f91b50f15f2384f4acba2e8059 Mon Sep 17 00:00:00 2001 From: Vidya Galli Date: Tue, 1 Oct 2024 07:49:22 -0700 Subject: [PATCH 15/89] Fix local variable 'image_features' referenced before assignment (#1383) --- optimum/habana/transformers/models/llava/modeling_llava.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/optimum/habana/transformers/models/llava/modeling_llava.py b/optimum/habana/transformers/models/llava/modeling_llava.py index cccbf8ebb9..402a1850fe 100644 --- a/optimum/habana/transformers/models/llava/modeling_llava.py +++ b/optimum/habana/transformers/models/llava/modeling_llava.py @@ -153,7 +153,8 @@ def forward( # 1. Extra the input embeddings inputs_embeds = self.get_input_embeddings()(input_ids) - + + image_features = None # 2. Merge text and images if pixel_values is not None and input_ids.shape[1] != 1: image_outputs = self.vision_tower( From f62ecde48a94d97b45ab779faee7bd3cd4f24304 Mon Sep 17 00:00:00 2001 From: Harish Subramony <81822986+hsubramony@users.noreply.github.com> Date: Wed, 2 Oct 2024 05:08:20 -0700 Subject: [PATCH 16/89] Use model.generation_config instead of model.config (#1384) Co-authored-by: regisss <15324346+regisss@users.noreply.github.com> --- examples/speech-recognition/run_speech_recognition_seq2seq.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/examples/speech-recognition/run_speech_recognition_seq2seq.py b/examples/speech-recognition/run_speech_recognition_seq2seq.py index ff9702e80c..4dcf0b498b 100755 --- a/examples/speech-recognition/run_speech_recognition_seq2seq.py +++ b/examples/speech-recognition/run_speech_recognition_seq2seq.py @@ -580,7 +580,8 @@ def compute_metrics(pred): # save feature extractor, tokenizer and config feature_extractor.save_pretrained(training_args.output_dir) tokenizer.save_pretrained(training_args.output_dir) - 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) processor = AutoProcessor.from_pretrained(training_args.output_dir) From a8fb8ac449e848cd210aee7366dfb4eb54629bb8 Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Wed, 2 Oct 2024 12:09:52 +0000 Subject: [PATCH 17/89] Make style --- optimum/habana/transformers/models/llava/modeling_llava.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/optimum/habana/transformers/models/llava/modeling_llava.py b/optimum/habana/transformers/models/llava/modeling_llava.py index 402a1850fe..997c16d700 100644 --- a/optimum/habana/transformers/models/llava/modeling_llava.py +++ b/optimum/habana/transformers/models/llava/modeling_llava.py @@ -153,7 +153,7 @@ def forward( # 1. Extra the input embeddings inputs_embeds = self.get_input_embeddings()(input_ids) - + image_features = None # 2. Merge text and images if pixel_values is not None and input_ids.shape[1] != 1: From dd07c16c7764d570a69348d8490e870e6942e131 Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Tue, 7 Jan 2025 16:06:06 +0000 Subject: [PATCH 18/89] Upgrade to Transformers v4.47.1 --- .../run_audio_classification.py | 4 +- .../contrastive-image-text/run_bridgetower.py | 49 +- examples/contrastive-image-text/run_clip.py | 46 +- .../run_image_classification.py | 4 +- .../run_image2text_lora_finetune.py | 2 +- examples/language-modeling/run_clm.py | 4 +- examples/language-modeling/run_lora_clm.py | 2 +- examples/language-modeling/run_mlm.py | 4 +- .../run_multitask_prompt_tuning.py | 4 +- .../run_prompt_tuning_clm.py | 4 +- .../README.md | 0 .../run_example.py | 0 .../run_example_sam.py | 0 examples/question-answering/run_qa.py | 4 +- examples/question-answering/run_seq2seq_qa.py | 4 +- .../run_speech_recognition_ctc.py | 4 +- .../run_speech_recognition_seq2seq.py | 4 +- .../unconditional_image_generation.py | 2 +- examples/summarization/run_summarization.py | 4 +- examples/text-classification/run_glue.py | 4 +- examples/translation/run_translation.py | 4 +- .../habana/transformers/generation/utils.py | 204 ++++---- .../models/bloom/modeling_bloom.py | 2 +- .../transformers/models/clip/modeling_clip.py | 18 +- .../models/codegen/modeling_codegen.py | 2 +- .../models/cohere/modeling_cohere.py | 24 +- .../models/falcon/modeling_falcon.py | 10 +- .../falcon_mamba/modeling_falcon_mamba.py | 4 +- .../models/gemma/modeling_gemma.py | 28 +- .../models/gemma2/modeling_gemma2.py | 26 +- .../transformers/models/gpt2/modeling_gpt2.py | 3 +- .../models/gpt_neo/modeling_gpt_neo.py | 2 +- .../models/gpt_neox/modeling_gpt_neox.py | 2 +- .../transformers/models/gptj/modeling_gptj.py | 2 +- .../models/idefics2/modeling_idefics2.py | 8 +- .../models/llama/modeling_llama.py | 84 +--- .../models/llava_next/modeling_llava_next.py | 4 +- .../models/mistral/modeling_mistral.py | 17 +- .../models/mixtral/configuration_mixtral.py | 2 + .../models/mixtral/modeling_mixtral.py | 28 +- .../models/mllama/modeling_mllama.py | 26 +- .../transformers/models/opt/modeling_opt.py | 12 + .../models/paligemma/modeling_paligemma.py | 18 +- .../models/persimmon/modeling_persimmon.py | 2 +- .../transformers/models/phi/modeling_phi.py | 26 +- .../models/qwen2/modeling_qwen2.py | 33 +- .../models/qwen2_moe/modeling_qwen2_moe.py | 24 +- .../models/speecht5/modeling_speecht5.py | 5 +- .../models/stablelm/modeling_stablelm.py | 2 +- .../models/starcoder2/modeling_starcoder2.py | 31 +- .../transformers/models/t5/modeling_t5.py | 25 +- .../models/wav2vec2/modeling_wav2vec2.py | 7 +- optimum/habana/transformers/trainer.py | 457 ++++++++++-------- .../habana/transformers/trainer_seq2seq.py | 40 +- optimum/habana/transformers/training_args.py | 41 +- tests/test_trainer.py | 399 ++++++++++----- tests/test_trainer_seq2seq.py | 4 +- 57 files changed, 927 insertions(+), 848 deletions(-) rename examples/{object-segementation => object-segmentation}/README.md (100%) rename examples/{object-segementation => object-segmentation}/run_example.py (100%) rename examples/{object-segementation => object-segmentation}/run_example_sam.py (100%) diff --git a/examples/audio-classification/run_audio_classification.py b/examples/audio-classification/run_audio_classification.py index 6defd566d3..682615a18e 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.45.0") +check_min_version("4.47.0") check_optimum_habana_min_version("1.16.0.dev0") require_version("datasets>=1.14.0", "To fix: pip install -r examples/pytorch/audio-classification/requirements.txt") @@ -404,7 +404,7 @@ def compute_metrics(eval_pred): train_dataset=raw_datasets["train"] if training_args.do_train else None, eval_dataset=raw_datasets["eval"] if training_args.do_eval else None, compute_metrics=compute_metrics, - tokenizer=feature_extractor, + processing_class=feature_extractor, ) # Training diff --git a/examples/contrastive-image-text/run_bridgetower.py b/examples/contrastive-image-text/run_bridgetower.py index 5964b2cdcc..42ee164cdf 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.45.0") +check_min_version("4.47.0") check_optimum_habana_min_version("1.16.0.dev0") require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/contrastive-image-text/requirements.txt") @@ -153,10 +153,6 @@ class DataTrainingArguments: default=None, metadata={"help": "An optional input evaluation data file (a jsonlines file)."}, ) - test_file: Optional[str] = field( - default=None, - metadata={"help": "An optional input testing data file (a jsonlines file)."}, - ) max_seq_length: Optional[int] = field( default=128, metadata={ @@ -205,9 +201,6 @@ def __post_init__(self): if self.validation_file is not None: extension = self.validation_file.split(".")[-1] assert extension in ["csv", "json"], "`validation_file` should be a csv or a json file." - if self.test_file is not None: - extension = self.test_file.split(".")[-1] - assert extension in ["csv", "json"], "`test_file` should be a csv or a json file." dataset_name_mapping = { @@ -340,9 +333,6 @@ def main(): if data_args.validation_file is not None: data_files["validation"] = data_args.validation_file extension = data_args.validation_file.split(".")[-1] - if data_args.test_file is not None: - data_files["test"] = data_args.test_file - extension = data_args.test_file.split(".")[-1] dataset = load_dataset( extension, data_files=data_files, @@ -426,8 +416,6 @@ def _freeze_params(module): column_names = dataset["train"].column_names elif training_args.do_eval: column_names = dataset["validation"].column_names - elif training_args.do_predict: - column_names = dataset["test"].column_names else: logger.info("There is nothing to do. Please pass `do_train`, `do_eval` and/or `do_predict`.") return @@ -549,33 +537,6 @@ def transform_images(examples): # Transform images on the fly as doing it on the whole dataset takes too much time. eval_dataset.set_transform(transform_images) - if training_args.do_predict: - if "test" not in dataset: - raise ValueError("--do_predict requires a test dataset") - test_dataset = dataset["test"] - if data_args.max_eval_samples is not None: - max_eval_samples = min(len(test_dataset), data_args.max_eval_samples) - test_dataset = test_dataset.select(range(max_eval_samples)) - - test_dataset = test_dataset.map( - function=tokenize_captions, - batched=True, - num_proc=data_args.preprocessing_num_workers, - remove_columns=[col for col in column_names if col != image_column], - load_from_cache_file=not data_args.overwrite_cache, - desc="Running tokenizer on test dataset", - ) - - 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) - # 8. Initialize our trainer trainer_cls = HabanaDataloaderTrainer if data_args.mediapipe_dataloader else GaudiTrainer trainer = trainer_cls( @@ -608,13 +569,7 @@ def transform_images(examples): trainer.log_metrics("validation", metrics) trainer.save_metrics("validation", metrics) - # 11. Test - if training_args.do_predict: - metrics = trainer.evaluate(eval_dataset=test_dataset) - trainer.log_metrics("test", metrics) - trainer.save_metrics("test", metrics) - - # 12. Write Training Stats and push to hub. + # 11. Write Training Stats and push to hub. finetuned_from = model_args.model_name_or_path # If from a local directory, don't set `finetuned_from` as this is required to be a valid repo. id on the Hub. if os.path.isdir(finetuned_from): diff --git a/examples/contrastive-image-text/run_clip.py b/examples/contrastive-image-text/run_clip.py index fc3bb4886e..f7ca7f6862 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.45.0") +check_min_version("4.47.0") check_optimum_habana_min_version("1.16.0.dev0") require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/contrastive-image-text/requirements.txt") @@ -149,10 +149,6 @@ class DataTrainingArguments: default=None, metadata={"help": "An optional input evaluation data file (a jsonlines file)."}, ) - test_file: Optional[str] = field( - default=None, - metadata={"help": "An optional input testing data file (a jsonlines file)."}, - ) max_seq_length: Optional[int] = field( default=128, metadata={ @@ -201,9 +197,6 @@ def __post_init__(self): if self.validation_file is not None: extension = self.validation_file.split(".")[-1] assert extension in ["csv", "json"], "`validation_file` should be a csv or a json file." - if self.test_file is not None: - extension = self.test_file.split(".")[-1] - assert extension in ["csv", "json"], "`test_file` should be a csv or a json file." dataset_name_mapping = { @@ -335,9 +328,6 @@ def main(): if data_args.validation_file is not None: data_files["validation"] = data_args.validation_file extension = data_args.validation_file.split(".")[-1] - if data_args.test_file is not None: - data_files["test"] = data_args.test_file - extension = data_args.test_file.split(".")[-1] dataset = load_dataset( extension, data_files=data_files, @@ -407,8 +397,6 @@ def _freeze_params(module): column_names = dataset["train"].column_names elif training_args.do_eval: column_names = dataset["validation"].column_names - elif training_args.do_predict: - column_names = dataset["test"].column_names else: logger.info("There is nothing to do. Please pass `do_train`, `do_eval` and/or `do_predict`.") return @@ -523,38 +511,6 @@ def filter_corrupt_images(examples): # Transform images on the fly as doing it on the whole dataset takes too much time. eval_dataset.set_transform(transform_images) - if training_args.do_predict: - if "test" not in dataset: - raise ValueError("--do_predict requires a test dataset") - test_dataset = dataset["test"] - if data_args.max_eval_samples is not None: - max_eval_samples = min(len(test_dataset), data_args.max_eval_samples) - test_dataset = test_dataset.select(range(max_eval_samples)) - - test_dataset = test_dataset.filter( - filter_corrupt_images, batched=True, num_proc=data_args.preprocessing_num_workers - ) - test_dataset = test_dataset.map( - function=tokenize_captions, - batched=True, - num_proc=data_args.preprocessing_num_workers, - remove_columns=[col for col in column_names if col != image_column], - load_from_cache_file=not data_args.overwrite_cache, - desc="Running tokenizer on test dataset", - ) - - # Transform images on the fly as doing it on the whole dataset takes too much time. - test_dataset.set_transform(transform_images) - 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) - # 8. Initialize our trainer trainer_cls = HabanaDataloaderTrainer if data_args.mediapipe_dataloader else GaudiTrainer trainer = trainer_cls( diff --git a/examples/image-classification/run_image_classification.py b/examples/image-classification/run_image_classification.py index bc45087f9e..440cf64264 100644 --- a/examples/image-classification/run_image_classification.py +++ b/examples/image-classification/run_image_classification.py @@ -63,7 +63,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.45.0") +check_min_version("4.47.0") check_optimum_habana_min_version("1.16.0.dev0") require_version("datasets>=2.14.0", "To fix: pip install -r examples/pytorch/image-classification/requirements.txt") @@ -414,7 +414,7 @@ def val_transforms(example_batch): train_dataset=dataset["train"] if training_args.do_train else None, eval_dataset=dataset["validation"] if training_args.do_eval else None, compute_metrics=compute_metrics, - tokenizer=image_processor, + processing_class=image_processor, data_collator=collate_fn, ) diff --git a/examples/image-to-text/run_image2text_lora_finetune.py b/examples/image-to-text/run_image2text_lora_finetune.py index ded60e6d52..b2ebb9424c 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.10.0") +check_optimum_habana_min_version("1.16.0.dev0") def normalized_levenshtein(s1, s2): diff --git a/examples/language-modeling/run_clm.py b/examples/language-modeling/run_clm.py index feac065364..87b6528260 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.45.0") +check_min_version("4.47.0") check_optimum_habana_min_version("1.16.0.dev0") require_version("datasets>=2.14.0", "To fix: pip install -r examples/pytorch/language-modeling/requirements.txt") @@ -630,7 +630,7 @@ def compute_metrics(eval_preds): args=training_args, train_dataset=train_dataset if training_args.do_train else None, eval_dataset=eval_dataset if training_args.do_eval else None, - tokenizer=tokenizer, + processing_class=tokenizer, # Data collator will default to DataCollatorWithPadding, so we change it. data_collator=default_data_collator, compute_metrics=compute_metrics if training_args.do_eval else None, diff --git a/examples/language-modeling/run_lora_clm.py b/examples/language-modeling/run_lora_clm.py index 3ff7fbfd3a..df460ec2fd 100644 --- a/examples/language-modeling/run_lora_clm.py +++ b/examples/language-modeling/run_lora_clm.py @@ -942,7 +942,7 @@ def compute_metrics(eval_preds): args=training_args, train_dataset=train_dataset if training_args.do_train else None, eval_dataset=eval_dataset if training_args.do_eval else None, - tokenizer=tokenizer, + processing_class=tokenizer, data_collator=data_collator, 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, diff --git a/examples/language-modeling/run_mlm.py b/examples/language-modeling/run_mlm.py index 2de43c910b..abea9c0eb1 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.45.0") +check_min_version("4.47.0") check_optimum_habana_min_version("1.16.0.dev0") require_version("datasets>=2.14.0", "To fix: pip install -r examples/pytorch/language-modeling/requirements.txt") @@ -639,7 +639,7 @@ def compute_metrics(eval_preds): args=training_args, train_dataset=train_dataset if training_args.do_train else None, eval_dataset=eval_dataset if training_args.do_eval else None, - tokenizer=tokenizer, + processing_class=tokenizer, data_collator=data_collator, 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, diff --git a/examples/language-modeling/run_multitask_prompt_tuning.py b/examples/language-modeling/run_multitask_prompt_tuning.py index 9f955db44e..7f788fc26c 100644 --- a/examples/language-modeling/run_multitask_prompt_tuning.py +++ b/examples/language-modeling/run_multitask_prompt_tuning.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 risk. -check_min_version("4.45.0") +check_min_version("4.47.0") check_optimum_habana_min_version("1.16.0.dev0") require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/language-modeling/requirements.txt") @@ -370,7 +370,7 @@ def compute_metrics(pred): data_collator=collate_fn, train_dataset=MyDataset("train"), eval_dataset=MyDataset("val"), - tokenizer=tokenizer, + processing_class=tokenizer, compute_metrics=compute_metrics, ) diff --git a/examples/language-modeling/run_prompt_tuning_clm.py b/examples/language-modeling/run_prompt_tuning_clm.py index 44ea542d14..f08280e695 100644 --- a/examples/language-modeling/run_prompt_tuning_clm.py +++ b/examples/language-modeling/run_prompt_tuning_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.45.0") +check_min_version("4.47.0") check_optimum_habana_min_version("1.16.0.dev0") require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/language-modeling/requirements.txt") @@ -348,7 +348,7 @@ def preprocess_function(examples): data_collator=default_data_collator, train_dataset=train_dataset, eval_dataset=eval_dataset, - tokenizer=tokenizer, + processing_class=tokenizer, ) if training_args.do_train: diff --git a/examples/object-segementation/README.md b/examples/object-segmentation/README.md similarity index 100% rename from examples/object-segementation/README.md rename to examples/object-segmentation/README.md diff --git a/examples/object-segementation/run_example.py b/examples/object-segmentation/run_example.py similarity index 100% rename from examples/object-segementation/run_example.py rename to examples/object-segmentation/run_example.py diff --git a/examples/object-segementation/run_example_sam.py b/examples/object-segmentation/run_example_sam.py similarity index 100% rename from examples/object-segementation/run_example_sam.py rename to examples/object-segmentation/run_example_sam.py diff --git a/examples/question-answering/run_qa.py b/examples/question-answering/run_qa.py index 5ad77be381..5b93fa5f1b 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.45.0") +check_min_version("4.47.0") check_optimum_habana_min_version("1.16.0.dev0") require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/question-answering/requirements.txt") @@ -663,7 +663,7 @@ def compute_metrics(p: EvalPrediction): train_dataset=train_dataset if training_args.do_train else None, eval_dataset=eval_dataset if training_args.do_eval else None, eval_examples=eval_examples if training_args.do_eval else None, - tokenizer=tokenizer, + processing_class=tokenizer, data_collator=data_collator, post_process_function=post_processing_function, compute_metrics=compute_metrics, diff --git a/examples/question-answering/run_seq2seq_qa.py b/examples/question-answering/run_seq2seq_qa.py index aaadbee417..bc9d9beff4 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.45.0") +check_min_version("4.47.0") check_optimum_habana_min_version("1.16.0.dev0") require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/question-answering/requirements.txt") @@ -682,7 +682,7 @@ def post_processing_function( train_dataset=train_dataset if training_args.do_train else None, eval_dataset=eval_dataset if training_args.do_eval else None, eval_examples=eval_examples if training_args.do_eval else None, - tokenizer=tokenizer, + processing_class=tokenizer, data_collator=data_collator, compute_metrics=compute_metrics if training_args.predict_with_generate else None, post_process_function=post_processing_function, diff --git a/examples/speech-recognition/run_speech_recognition_ctc.py b/examples/speech-recognition/run_speech_recognition_ctc.py index 9d53e58519..3403d00f3c 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.45.0") +check_min_version("4.47.0") check_optimum_habana_min_version("1.16.0.dev0") require_version("datasets>=1.18.0", "To fix: pip install -r examples/pytorch/speech-recognition/requirements.txt") @@ -773,7 +773,7 @@ def compute_metrics(pred): compute_metrics=compute_metrics, train_dataset=vectorized_datasets["train"] if training_args.do_train else None, eval_dataset=vectorized_datasets["eval"] if training_args.do_eval else None, - tokenizer=processor, + processing_class=processor, preprocess_logits_for_metrics=preprocess_logits_for_metrics, ) diff --git a/examples/speech-recognition/run_speech_recognition_seq2seq.py b/examples/speech-recognition/run_speech_recognition_seq2seq.py index bb745af049..d61973f5c6 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.45.0") +check_min_version("4.47.0") check_optimum_habana_min_version("1.16.0.dev0") require_version("datasets>=1.18.0", "To fix: pip install -r examples/pytorch/speech-recognition/requirements.txt") @@ -600,7 +600,7 @@ def compute_metrics(pred): args=training_args, train_dataset=vectorized_datasets["train"] if training_args.do_train else None, eval_dataset=vectorized_datasets["eval"] if training_args.do_eval else None, - tokenizer=feature_extractor, + processing_class=feature_extractor, data_collator=data_collator, compute_metrics=compute_metrics if training_args.predict_with_generate else None, ) diff --git a/examples/stable-diffusion/unconditional_image_generation.py b/examples/stable-diffusion/unconditional_image_generation.py index bd70d0e4d6..f908c4fb9c 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.45.0") +check_min_version("4.47.0") check_optimum_habana_min_version("1.16.0.dev0") # Setup logging diff --git a/examples/summarization/run_summarization.py b/examples/summarization/run_summarization.py index 65755d24a2..dc22580f20 100755 --- a/examples/summarization/run_summarization.py +++ b/examples/summarization/run_summarization.py @@ -65,7 +65,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.45.0") +check_min_version("4.47.0") check_optimum_habana_min_version("1.16.0.dev0") require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/summarization/requirements.txt") @@ -787,7 +787,7 @@ def compute_metrics(eval_preds): args=training_args, train_dataset=train_dataset if training_args.do_train else None, eval_dataset=eval_dataset if training_args.do_eval else None, - tokenizer=tokenizer, + processing_class=tokenizer, data_collator=data_collator, compute_metrics=compute_metrics if training_args.predict_with_generate else None, ) diff --git a/examples/text-classification/run_glue.py b/examples/text-classification/run_glue.py index 68f5e9a2aa..5cfe00ff6e 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.45.0") +check_min_version("4.47.0") check_optimum_habana_min_version("1.16.0.dev0") require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/text-classification/requirements.txt") @@ -563,7 +563,7 @@ def compute_metrics(p: EvalPrediction): train_dataset=train_dataset if training_args.do_train else None, eval_dataset=eval_dataset if training_args.do_eval else None, compute_metrics=compute_metrics, - tokenizer=tokenizer, + processing_class=tokenizer, data_collator=data_collator, ) diff --git a/examples/translation/run_translation.py b/examples/translation/run_translation.py index 6f55ae1350..1a6f3379aa 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.45.0") +check_min_version("4.47.0") check_optimum_habana_min_version("1.16.0.dev0") require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/translation/requirements.txt") @@ -633,7 +633,7 @@ def compute_metrics(eval_preds): args=training_args, train_dataset=train_dataset if training_args.do_train else None, eval_dataset=eval_dataset if training_args.do_eval else None, - tokenizer=tokenizer, + processing_class=tokenizer, data_collator=data_collator, compute_metrics=compute_metrics if training_args.predict_with_generate else None, ) diff --git a/optimum/habana/transformers/generation/utils.py b/optimum/habana/transformers/generation/utils.py index 9c1b802baf..defa93c6c0 100644 --- a/optimum/habana/transformers/generation/utils.py +++ b/optimum/habana/transformers/generation/utils.py @@ -26,7 +26,9 @@ from transformers.generation.beam_constraints import DisjunctiveConstraint, PhrasalConstraint from transformers.generation.beam_search import BeamScorer, BeamSearchScorer, ConstrainedBeamSearchScorer from transformers.generation.candidate_generator import ( + AssistedCandidateGeneratorDifferentTokenizers, CandidateGenerator, + EarlyExitCandidateGenerator, PromptLookupCandidateGenerator, _crop_past_key_values, _prepare_attention_mask, @@ -57,8 +59,9 @@ stack_model_outputs, ) 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_quanto_available, is_torchdynamo_compiling +from transformers.utils import ModelOutput, is_hqq_available, is_optimum_quanto_available, is_torchdynamo_compiling from optimum.utils import logging @@ -556,15 +559,39 @@ def _get_candidate_generator( inputs_tensor: torch.Tensor, assistant_model: "PreTrainedModel", logits_processor: LogitsProcessorList, + target_tokenizer: "PreTrainedTokenizerBase", + assistant_tokenizer: "PreTrainedTokenizerBase", model_kwargs: Dict, ) -> CandidateGenerator: - if generation_config.prompt_lookup_num_tokens is not None: + different_tokenizers = all(v is not None for v in (assistant_model, target_tokenizer, assistant_tokenizer)) + + if generation_config.assistant_early_exit is not None: + candidate_generator = EarlyExitCandidateGenerator( + input_ids=input_ids, + assistant_model=self, + generation_config=generation_config, + model_kwargs=model_kwargs, + inputs_tensor=inputs_tensor, + logits_processor=logits_processor, + ) + elif generation_config.prompt_lookup_num_tokens is not None: candidate_generator = PromptLookupCandidateGenerator( eos_token_id=generation_config._eos_token_tensor, num_output_tokens=generation_config.prompt_lookup_num_tokens, max_matching_ngram_size=generation_config.max_matching_ngram_size, 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, + ) else: candidate_generator = GaudiAssistedCandidateGenerator( input_ids=input_ids, @@ -625,7 +652,7 @@ def _prepare_generated_length( inputs_tensor, has_token_idx, ): - """Prepared max and min length in generaion configs to avoid clashes between similar attributes""" + """Prepared max and min length in generation configs to avoid clashes between similar attributes""" if generation_config.max_new_tokens is not None: if not has_default_max_length and generation_config.max_length is not None: @@ -648,6 +675,12 @@ def _prepare_generated_length( and not self.config.is_encoder_decoder ): generation_config.max_length -= inputs_tensor.shape[1] + elif has_default_max_length: # by default let's always generate 20 new tokens + if generation_config.max_length == GaudiGenerationConfig().max_length: + generation_config.max_length = generation_config.max_length + input_ids_length + max_position_embeddings = getattr(self.config, "max_position_embeddings", None) + if max_position_embeddings is not None: + generation_config.max_length = min(generation_config.max_length, max_position_embeddings) # same for min length if generation_config.min_new_tokens is not None: @@ -843,10 +876,10 @@ def _prepare_cache_for_generation( ) cache_class = QUANT_BACKEND_CLASSES_MAPPING[cache_config.backend] - if cache_config.backend == "quanto" and not is_quanto_available(): + if cache_config.backend == "quanto" and not is_optimum_quanto_available(): raise ImportError( - "You need to install `quanto` in order to use KV cache quantization with quanto backend. " - "Please install it via with `pip install quanto`" + "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(): raise ImportError( @@ -930,12 +963,12 @@ def generate( for constrained generation conditioned on the prefix, as described in [Autoregressive Entity Retrieval](https://arxiv.org/abs/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` under DeepSpeed ZeRO Stage 3 multiple GPUs environment to avoid hanging if one GPU finished - generating before other GPUs. Otherwise it'll be set to `False`. + 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 + deadlocking if one GPU finishes generating before other GPUs. Otherwise, defaults to `False`. assistant_model (`PreTrainedModel`, *optional*): An assistant model that can be used to accelerate generation. The assistant model must have the exact - same tokenizer. The acceleration is achieved when forecasting candidate tokens with the assistent model + same tokenizer. The acceleration is achieved when forecasting candidate tokens with the assistant model is much faster than running generation with the model you're calling generate from. As such, the assistant model should be much smaller. streamer (`BaseStreamer`, *optional*): @@ -988,6 +1021,7 @@ def generate( # 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: raise ValueError( "`hpu_graphs` is True but `lazy_mode` is False. HPU graphs require `lazy_mode` to be set to True." @@ -995,14 +1029,11 @@ def generate( num_virtual_tokens = kwargs.pop("num_virtual_tokens", 0) generation_config, model_kwargs = self._prepare_generation_config(generation_config, **kwargs) self._validate_model_kwargs(model_kwargs.copy()) - self._validate_assistant(assistant_model) + self._validate_assistant(assistant_model, tokenizer, assistant_tokenizer) # 2. Set generation parameters if not already defined if synced_gpus is None: - if is_deepspeed_zero3_enabled() and dist.get_world_size() > 1: - synced_gpus = True - else: - synced_gpus = False + synced_gpus = (is_deepspeed_zero3_enabled() or is_fsdp_managed_module(self)) and dist.get_world_size() > 1 logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList() stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList() @@ -1039,15 +1070,13 @@ def generate( # decoder-only models with inputs_embeds forwarding must use caching (otherwise we can't detect whether we are # generating the first new token or not, and we only want to use the embeddings for the first new token) if not self.config.is_encoder_decoder and model_input_name == "inputs_embeds": - model_kwargs["use_cache"] = True - else: - model_kwargs["use_cache"] = generation_config.use_cache + generation_config.use_cache = True self.generation_config.max_length = generation_config.max_length if not kwargs_has_attention_mask and requires_attention_mask and accepts_attention_mask: model_kwargs["attention_mask"] = self._prepare_attention_mask_for_generation( - inputs_tensor, generation_config._pad_token_tensor, generation_config._eos_token_tensor + inputs_tensor, generation_config, model_kwargs ) elif kwargs_has_attention_mask: # TODO (joao): generalize this check with other types of inputs @@ -1361,6 +1390,9 @@ def generate( **kwargs, ) + # Set model_kwargs `use_cache` so we can use it later in forward runs + model_kwargs["use_cache"] = generation_config.use_cache + # In lazy mode, import Habana torch to be able to add mark_step() if lazy_mode: import habana_frameworks.torch.core as htcore @@ -1394,6 +1426,8 @@ def generate( inputs_tensor=inputs_tensor, assistant_model=assistant_model, logits_processor=logits_processor, + target_tokenizer=tokenizer, + assistant_tokenizer=assistant_tokenizer, model_kwargs=model_kwargs, ) @@ -1691,7 +1725,8 @@ def _dola_decoding( generation_config ([`~generation.GenerationConfig`]): The generation configuration to be used as parametrization of the decoding method. synced_gpus (`bool`): - Whether to continue running the while loop until max_length (needed for ZeRO stage 3) + Whether to continue running the while loop until max_length (needed to avoid deadlocking with + `FullyShardedDataParallel` and DeepSpeed ZeRO Stage 3). streamer (`BaseStreamer`, *optional*): Streamer object that will be used to stream the generated sequences. Generated tokens are passed through `streamer.put(token_ids)` and the streamer is responsible for any further processing. @@ -1748,7 +1783,8 @@ def _contrastive_search( generation_config ([`~generation.GenerationConfig`]): The generation configuration to be used as parametrization of the decoding method. synced_gpus (`bool`): - Whether to continue running the while loop until max_length (needed for ZeRO stage 3) + Whether to continue running the while loop until max_length (needed to avoid deadlocking with + `FullyShardedDataParallel` and DeepSpeed ZeRO Stage 3). streamer (`BaseStreamer`, *optional*): Streamer object that will be used to stream the generated sequences. Generated tokens are passed through `streamer.put(token_ids)` and the streamer is responsible for any further processing. @@ -1889,6 +1925,7 @@ def _contrastive_search( 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) model_kwargs = self._update_model_kwargs_for_generation( outputs, @@ -2042,7 +2079,7 @@ def _contrastive_search( output_attentions=output_attentions, ) - # This is essential to avoid having a last reference to the big past K-V and double the necesary memory + # This is essential to avoid having a last reference to the big past K-V and double the necessary memory # in the next loop del next_model_inputs @@ -2125,6 +2162,7 @@ def _contrastive_search( next_past_key_values = tuple(new_key_values) logit_for_next_step = torch.stack(torch.split(logits, top_k))[batch_indices, selected_idx, :] + logit_for_next_step = logit_for_next_step.to(input_ids.device) # Rebuilds the relevant parts of the model output for the selected token, for use in the next iteration if self.config.is_encoder_decoder: @@ -2156,8 +2194,14 @@ def _contrastive_search( ) # contrastive_search main logic end + # synced_gpus: don't waste resources running the code we don't need; kwargs must be updated before skipping + model_kwargs = self._update_model_kwargs_for_generation( + outputs, + model_kwargs, + is_encoder_decoder=self.config.is_encoder_decoder, + ) if synced_gpus and this_peer_finished: - continue # don't waste resources running the code we don't need + continue # finished sentences should have their next token be a padding token if not ignore_eos and has_eos_stopping_criteria: @@ -2173,11 +2217,6 @@ def _contrastive_search( if streamer is not None: streamer.put(next_tokens.cpu()) - model_kwargs = self._update_model_kwargs_for_generation( - outputs, - model_kwargs, - is_encoder_decoder=self.config.is_encoder_decoder, - ) # increase cur_len cur_len = cur_len + 1 @@ -2338,7 +2377,8 @@ def _sample( generation_config ([`GaudiGenerationConfig`]): The generation configuration to be used as parametrization of the decoding method. synced_gpus (`bool`): - Whether to continue running the while loop until max_length (needed for ZeRO stage 3) + Whether to continue running the while loop until max_length (needed to avoid deadlocking with + `FullyShardedDataParallel` and DeepSpeed ZeRO Stage 3). streamer (`BaseStreamer`, *optional*): Streamer object that will be used to stream the generated sequences. Generated tokens are passed through `streamer.put(token_ids)` and the streamer is responsible for any further processing. @@ -2452,14 +2492,21 @@ def _sample( **hpu_graphs_kwargs, ) + # synced_gpus: don't waste resources running the code we don't need; kwargs must be updated before skipping + model_kwargs = self._update_model_kwargs_for_generation( + outputs, + model_kwargs, + is_encoder_decoder=self.config.is_encoder_decoder, + ) if synced_gpus and this_peer_finished: - continue # don't waste resources running the code we don't need + continue token_idx = model_kwargs.get("token_idx", None) if token_idx is not None and outputs.logits.shape[-2] > 1: # case1 (w/o KV caching): outputs.logits.shape: [batch_size, max_length, vocab_size] if self.config.is_encoder_decoder: next_token_logits = outputs.logits[:, token_idx - 1, :].float() + next_token_logits = next_token_logits.to(input_ids.device) next_token_scores = logits_processor(input_ids[:, :token_idx], next_token_logits) else: if model_kwargs.get("num_virtual_tokens", 0) > 0: @@ -2471,10 +2518,12 @@ def _sample( next_token_logits = torch.index_select(outputs.logits, -2, output_idx - 1).squeeze(-2) else: next_token_logits = torch.index_select(outputs.logits, -2, token_idx - 1).squeeze(-2) + next_token_logits = next_token_logits.to(input_ids.device) next_token_scores = logits_processor(input_ids, next_token_logits) else: # .float() is needed to retain precision for later logits manipulations next_token_logits = outputs.logits[:, -1, :].float() + next_token_logits = next_token_logits.to(input_ids.device) if token_idx is not None and self.config.is_encoder_decoder: # case2 (with KV caching): outputs.logits.shape: [batch_size, 1, vocab_size] next_token_scores = logits_processor(input_ids[:, :token_idx], next_token_logits) @@ -2528,12 +2577,6 @@ def _sample( if streamer is not None: streamer.put(next_tokens.cpu()) - model_kwargs = self._update_model_kwargs_for_generation( - outputs, - model_kwargs, - is_encoder_decoder=self.config.is_encoder_decoder, - ) - cur_len = cur_len + 1 if bucket_size > 0 and bucket_internal: # Calculate slice idx for kv cache during the decode phase. @@ -2693,7 +2736,8 @@ def _beam_search( generation_config ([`GaudiGenerationConfig`]): The generation configuration to be used as parametrization of the decoding method. synced_gpus (`bool`): - Whether to continue running the while loop until max_length (needed for ZeRO stage 3) + Whether to continue running the while loop until max_length (needed to avoid deadlocking with + `FullyShardedDataParallel` and DeepSpeed ZeRO Stage 3). lazy_mode (`bool`, *optional*, defaults to `False`): Whether the run is executed in lazy mode or not (i.e. eager mode). profiling_warmup_steps (`int`, *optional*, defaults to 0): @@ -2954,9 +2998,15 @@ def expand_if_needed(tensor, new_size, value, dim=-1): **hpu_graphs_kwargs, ) + # synced_gpus: don't waste resources running the code we don't need; kwargs must be updated before skipping + model_kwargs = self._update_model_kwargs_for_generation( + outputs, + model_kwargs, + is_encoder_decoder=self.config.is_encoder_decoder, + ) if synced_gpus and this_peer_finished: cur_len = cur_len + 1 - continue # don't waste resources running the code we don't need + continue token_idx = model_kwargs.get("token_idx", None) if token_idx is not None and outputs.logits.shape[-2] > 1: @@ -2971,6 +3021,7 @@ def expand_if_needed(tensor, new_size, value, dim=-1): 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_scores = torch.nn.functional.log_softmax( next_token_logits, dim=-1 @@ -3087,12 +3138,6 @@ def expand_if_needed(tensor, new_size, value, dim=-1): else: input_ids = torch.cat([input_ids[beam_idx, :], beam_next_tokens.unsqueeze(-1)], dim=-1) - model_kwargs = self._update_model_kwargs_for_generation( - outputs, - model_kwargs, - is_encoder_decoder=self.config.is_encoder_decoder, - ) - if model_kwargs.get("past_key_values", None) is not None: if model_kwargs["reuse_cache"]: model_kwargs["past_key_values"] = unwrap_deepspeed_model(self).reorder_kv_cache(beam_idx) @@ -3276,7 +3321,8 @@ def _group_beam_search( generation_config ([`GaudiGenerationConfig`]): The generation configuration to be used as parametrization of the decoding method. synced_gpus (`bool`): - Whether to continue running the while loop until max_length (needed for ZeRO stage 3) + Whether to continue running the while loop until max_length (needed to avoid deadlocking with + `FullyShardedDataParallel` and DeepSpeed ZeRO Stage 3). lazy_mode (`bool`, *optional*, defaults to `False`): Whether the run is executed in lazy mode or not (i.e. eager mode). profiling_warmup_steps (`int`, *optional*, defaults to 0): @@ -3334,7 +3380,8 @@ def _constrained_beam_search( generation_config ([`GaudiGenerationConfig`]): The generation configuration to be used as parametrization of the decoding method. synced_gpus (`bool`): - Whether to continue running the while loop until max_length (needed for ZeRO stage 3) + Whether to continue running the while loop until max_length (needed to avoid deadlocking with + `FullyShardedDataParallel` and DeepSpeed ZeRO Stage 3). lazy_mode (`bool`, *optional*, defaults to `False`): Whether the run is executed in lazy mode or not (i.e. eager mode). profiling_warmup_steps (`int`, *optional*, defaults to 0): @@ -3433,9 +3480,15 @@ def _constrained_beam_search( **hpu_graphs_kwargs, ) + # synced_gpus: don't waste resources running the code we don't need; kwargs must be updated before skipping + model_kwargs = self._update_model_kwargs_for_generation( + outputs, + model_kwargs, + is_encoder_decoder=self.config.is_encoder_decoder, + ) if synced_gpus and this_peer_finished: cur_len = cur_len + 1 - continue # don't waste resources running the code we don't need + continue if token_idx is not None and outputs.logits.shape[-2] > 1: if model_kwargs.get("num_virtual_tokens", 0) > 0: @@ -3449,6 +3502,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_scores = torch.nn.functional.log_softmax( next_token_logits, dim=-1 @@ -3518,11 +3572,6 @@ def _constrained_beam_search( ) else: input_ids = torch.cat([input_ids[beam_idx, :], beam_next_tokens.unsqueeze(-1)], dim=-1) - model_kwargs = self._update_model_kwargs_for_generation( - outputs, - model_kwargs, - is_encoder_decoder=self.config.is_encoder_decoder, - ) # This is needed to properly delete outputs.logits which may be very large for first iteration # Otherwise a reference to outputs is kept which keeps the logits alive in the next iteration @@ -3638,7 +3687,8 @@ def _assisted_decoding( generation_config ([`~generation.GenerationConfig`]): The generation configuration to be used as parametrization of the decoding method. synced_gpus (`bool`): - Whether to continue running the while loop until max_length (needed for ZeRO stage 3) + Whether to continue running the while loop until max_length (needed to avoid deadlocking with + `FullyShardedDataParallel` and DeepSpeed ZeRO Stage 3). streamer (`BaseStreamer`, *optional*): Streamer object that will be used to stream the generated sequences. Generated tokens are passed through `streamer.put(token_ids)` and the streamer is responsible for any further processing. @@ -3689,19 +3739,10 @@ def _assisted_decoding( unfinished_sequences = torch.ones(batch_size, dtype=torch.long, device=input_ids.device) model_kwargs = self._get_initial_cache_position(input_ids, model_kwargs) - # This is needed if return_dict_in_generate is True - start_from_empty_dynamic_cache = False - past_key_values = model_kwargs.get("past_key_values", None) - if isinstance(past_key_values, DynamicCache) or ( - isinstance(past_key_values, EncoderDecoderCache) - and isinstance(past_key_values.self_attention_cache, DynamicCache) - ): - if len(past_key_values) == 0: - start_from_empty_dynamic_cache = True - hb_profer = HabanaProfile(warmup=profiling_warmup_steps, active=profiling_steps) hb_profer.start() this_peer_finished = False + is_first_iteration = True # to preserve the same API in the output as other generation methods token_idx = model_kwargs.get("token_idx", None) time_to_first_token_done = False @@ -3721,7 +3762,7 @@ def _assisted_decoding( # 1. Fetch candidate sequences from a `CandidateGenerator` 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) @@ -3769,6 +3810,7 @@ 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) next_token_logits = new_logits.clone() if len(logits_processor) > 0: for i in range(candidate_length + 1): @@ -3830,55 +3872,44 @@ def _assisted_decoding( # Store scores, attentions and hidden_states when required # Assistant: modified to append one tuple element per token, as in the other generation methods. if return_dict_in_generate: + newly_added_length = n_matches + 1 if output_scores: - scores += tuple(new_logits[:, i, :] for i in range(n_matches + 1)) + scores += tuple(new_logits[:, i, :] for i in range(newly_added_length)) if output_logits: - raw_logits += (next_token_logits,) - - if "past_key_values" not in model_kwargs or start_from_empty_dynamic_cache: - added_len = new_cur_len - # set it to false for other iterations - start_from_empty_dynamic_cache = False - else: - added_len = n_matches + 1 + raw_logits += tuple(next_token_logits[:, i, :] for i in range(newly_added_length)) + newly_added_length = new_cur_len if is_first_iteration else newly_added_length if output_attentions: if self.config.is_encoder_decoder: cross_attentions = _split_model_outputs( - cross_attentions, outputs.cross_attentions, cur_len, added_len + cross_attentions, outputs.cross_attentions, cur_len, newly_added_length ) decoder_attentions = _split_model_outputs( decoder_attentions, outputs.decoder_attentions, cur_len, - added_len, + newly_added_length, is_decoder_attention=True, ) - else: + # some (V)LLMs have hard requirement on SDPA and thus never return attn + elif outputs.attentions[0] is not None: decoder_attentions = _split_model_outputs( decoder_attentions, outputs.attentions, cur_len, - added_len, + newly_added_length, is_decoder_attention=True, ) if output_hidden_states: if self.config.is_encoder_decoder: decoder_hidden_states = _split_model_outputs( - decoder_hidden_states, outputs.decoder_hidden_states, cur_len, added_len + decoder_hidden_states, outputs.decoder_hidden_states, cur_len, newly_added_length ) else: decoder_hidden_states = _split_model_outputs( - decoder_hidden_states, outputs.hidden_states, cur_len, added_len + decoder_hidden_states, outputs.hidden_states, cur_len, newly_added_length ) - model_kwargs = self._update_model_kwargs_for_generation( - outputs, - model_kwargs, - is_encoder_decoder=self.config.is_encoder_decoder, - num_new_tokens=n_matches + 1, - ) - if ignore_eos: this_peer_finished = stopping_criteria( input_ids, @@ -3896,6 +3927,7 @@ def _assisted_decoding( eos_token_id=generation_config.eos_token_id, ) this_peer_finished = unfinished_sequences.max() == 0 + is_first_iteration = False if hb_gen_time is not None: if not time_to_first_token_done: diff --git a/optimum/habana/transformers/models/bloom/modeling_bloom.py b/optimum/habana/transformers/models/bloom/modeling_bloom.py index 5b0a770451..3edab86a60 100644 --- a/optimum/habana/transformers/models/bloom/modeling_bloom.py +++ b/optimum/habana/transformers/models/bloom/modeling_bloom.py @@ -357,7 +357,7 @@ def gaudi_bloom_model_forward( 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") + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") elif input_ids is not None: batch_size, seq_length = input_ids.shape elif inputs_embeds is not None: diff --git a/optimum/habana/transformers/models/clip/modeling_clip.py b/optimum/habana/transformers/models/clip/modeling_clip.py index 98eb7e2861..b48ba858ca 100644 --- a/optimum/habana/transformers/models/clip/modeling_clip.py +++ b/optimum/habana/transformers/models/clip/modeling_clip.py @@ -25,8 +25,12 @@ class GaudiCLIPVisionEmbeddings(CLIPVisionEmbeddings): - def forward(self, pixel_values: torch.FloatTensor) -> torch.Tensor: - batch_size = pixel_values.shape[0] + def forward(self, pixel_values: torch.FloatTensor, interpolate_pos_encoding=False) -> torch.Tensor: + batch_size, _, height, width = pixel_values.shape + if not interpolate_pos_encoding and (height != self.image_size or width != self.image_size): + raise ValueError( + f"Input image size ({height}*{width}) doesn't match model" f" ({self.image_size}*{self.image_size})." + ) target_dtype = self.patch_embedding.weight.dtype # if HQT quantization enabled, remove the explicit cast to float8 to avoid HQT casting error if "float8" in str(target_dtype) and pixel_values.device.type == "hpu": @@ -36,7 +40,10 @@ def forward(self, pixel_values: torch.FloatTensor) -> torch.Tensor: class_embeds = self.class_embedding.expand(batch_size, 1, -1) embeddings = torch.cat([class_embeds, patch_embeds], dim=1) - embeddings = embeddings + self.position_embedding(self.position_ids) + if interpolate_pos_encoding: + embeddings = embeddings + self.interpolate_pos_encoding(embeddings, height, width) + else: + embeddings = embeddings + self.position_embedding(self.position_ids) return embeddings @@ -288,6 +295,7 @@ def forward( 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]: @@ -306,7 +314,7 @@ def forward( if pixel_values is None: raise ValueError("You have to specify pixel_values") - hidden_states = self.embeddings(pixel_values) + hidden_states = self.embeddings(pixel_values, interpolate_pos_encoding=interpolate_pos_encoding) hidden_states = self.pre_layrnorm(hidden_states) encoder_outputs = self.encoder( @@ -339,6 +347,7 @@ def forward( pixel_values: Optional[torch.FloatTensor] = None, 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, @@ -356,6 +365,7 @@ def forward( 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/codegen/modeling_codegen.py b/optimum/habana/transformers/models/codegen/modeling_codegen.py index a7f15d32d4..cfe450ab6c 100644 --- a/optimum/habana/transformers/models/codegen/modeling_codegen.py +++ b/optimum/habana/transformers/models/codegen/modeling_codegen.py @@ -178,7 +178,7 @@ def gaudi_codegen_model_forward( 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") + 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() diff --git a/optimum/habana/transformers/models/cohere/modeling_cohere.py b/optimum/habana/transformers/models/cohere/modeling_cohere.py index c0785c88ed..119df106fb 100644 --- a/optimum/habana/transformers/models/cohere/modeling_cohere.py +++ b/optimum/habana/transformers/models/cohere/modeling_cohere.py @@ -3,7 +3,6 @@ import torch from torch import nn -from torch.nn import CrossEntropyLoss from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from transformers.models.cohere.modeling_cohere import ( Cache, @@ -192,9 +191,7 @@ def gaudi_cohere_model_forward( 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 cannot specify both input_ids and inputs_embeds at the same time, and must specify either one" - ) + 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`.") @@ -232,7 +229,7 @@ def gaudi_cohere_model_forward( all_self_attns = () if output_attentions else None next_decoder_cache = None - for decoder_layer in self.layers: + for decoder_layer in self.layers[: self.config.num_hidden_layers]: if output_hidden_states: all_hidden_states += (hidden_states,) @@ -310,7 +307,9 @@ def forward( output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, + num_logits_to_keep: int = 0, token_idx: Optional[torch.Tensor] = None, + **loss_kwargs, ) -> Union[Tuple, CausalLMOutputWithPast]: output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = ( @@ -334,22 +333,13 @@ def forward( ) hidden_states = outputs[0] - logits = self.lm_head(hidden_states) + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]) logits = logits * self.logit_scale - logits = logits.float() loss = None if labels is not None: - # Shift so that tokens < n predict n - shift_logits = logits[..., :-1, :].contiguous() - shift_labels = labels[..., 1:].contiguous() - # Flatten the tokens - loss_fct = CrossEntropyLoss() - shift_logits = shift_logits.view(-1, self.config.vocab_size) - shift_labels = shift_labels.view(-1) - # Enable model parallelism - shift_labels = shift_labels.to(shift_logits.device) - loss = loss_fct(shift_logits, shift_labels) + loss = self.loss_function(logits, labels, self.vocab_size, **loss_kwargs) if not return_dict: output = (logits,) + outputs[1:] diff --git a/optimum/habana/transformers/models/falcon/modeling_falcon.py b/optimum/habana/transformers/models/falcon/modeling_falcon.py index 8895f32459..3ef9edbdbb 100644 --- a/optimum/habana/transformers/models/falcon/modeling_falcon.py +++ b/optimum/habana/transformers/models/falcon/modeling_falcon.py @@ -748,7 +748,7 @@ def forward( 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") + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") elif input_ids is not None: batch_size, seq_length = input_ids.shape elif inputs_embeds is not None: @@ -1032,6 +1032,7 @@ def forward( output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, + num_logits_to_keep: int = 0, token_idx: Optional[torch.Tensor] = None, reuse_cache: Optional[bool] = False, trim_logits: Optional[bool] = False, @@ -1045,6 +1046,11 @@ def forward( Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100` are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]` + + num_logits_to_keep (`int`, *optional*): + Calculate logits for the last `num_logits_to_keep` tokens. If `0`, calculate logits for all + `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that + token can save memory, which becomes pretty significant for long sequences or large vocabulary size. """ return_dict = return_dict if return_dict is not None else self.config.use_return_dict if use_flash_attention: @@ -1082,7 +1088,7 @@ def forward( else: hidden_states = hidden_states[:, -1:, :] - lm_logits = self.lm_head(hidden_states) + lm_logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]) loss = None if labels is not None: 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 5d618fac91..dddaa5055a 100644 --- a/optimum/habana/transformers/models/falcon_mamba/modeling_falcon_mamba.py +++ b/optimum/habana/transformers/models/falcon_mamba/modeling_falcon_mamba.py @@ -53,9 +53,7 @@ def gaudi_FalconMambaModel_forward( 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): # ^ is python for xor - raise ValueError( - "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one" - ) + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") if inputs_embeds is None: inputs_embeds = self.embeddings(input_ids) diff --git a/optimum/habana/transformers/models/gemma/modeling_gemma.py b/optimum/habana/transformers/models/gemma/modeling_gemma.py index 532539065d..30b01c8aad 100755 --- a/optimum/habana/transformers/models/gemma/modeling_gemma.py +++ b/optimum/habana/transformers/models/gemma/modeling_gemma.py @@ -25,7 +25,6 @@ import torch import torch.nn.functional as F -from torch.nn import CrossEntropyLoss from transformers.cache_utils import Cache, DynamicCache from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from transformers.models.gemma.modeling_gemma import ( @@ -37,7 +36,7 @@ GemmaModel, apply_rotary_pos_emb, ) -from transformers.utils import is_torchdynamo_compiling, logging +from transformers.utils import logging from ...modeling_attn_mask_utils import ( _gaudi_prepare_4d_causal_attention_mask, @@ -611,7 +610,7 @@ def forward( self._attn_implementation = "eager" 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") + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") elif input_ids is not None: batch_size, seq_length = input_ids.shape[:2] elif inputs_embeds is not None: @@ -673,7 +672,7 @@ def forward( if lazy_mode: htcore.mark_step() - for layer_idx, decoder_layer in enumerate(self.layers): + for layer_idx, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]): if ( lazy_mode and not self.training @@ -778,6 +777,7 @@ def forward( flash_attention_recompute: Optional[bool] = False, flash_attention_causal_mask: Optional[bool] = False, attn_softmax_bf16: Optional[bool] = False, + **loss_kwargs, ) -> Union[Tuple, CausalLMOutputWithPast]: """ Inherits from GemmaForCausalLM: https://github.com/huggingface/transformers/blob/v4.38.1/src/transformers/models/gemma/modeling_gemma.py @@ -812,28 +812,12 @@ def forward( ) hidden_states = outputs[0] - if labels is None and not is_torchdynamo_compiling(): - logger.warning_once( - "Starting from v4.46, the `logits` model output will have the same type as the model (except at train time, where it will always be FP32)" - ) # Only compute necessary logits, and do not upcast them to float if we are not computing the loss - # TODO: remove the float() operation in v4.46 - logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]).float() + logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]) 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 so that tokens < n predict n - shift_logits = logits[..., :-1, :].contiguous() - shift_labels = labels[..., 1:].contiguous() - # Flatten the tokens - loss_fct = CrossEntropyLoss() - shift_logits = shift_logits.view(-1, self.config.vocab_size) - shift_labels = shift_labels.view(-1) - # Enable model parallelism - shift_labels = shift_labels.to(shift_logits.device) - loss = loss_fct(shift_logits, shift_labels) + loss = self.loss_function(logits, labels, self.vocab_size, **loss_kwargs) if not return_dict: output = (logits,) + outputs[1:] diff --git a/optimum/habana/transformers/models/gemma2/modeling_gemma2.py b/optimum/habana/transformers/models/gemma2/modeling_gemma2.py index 4196775c19..fff49d4649 100755 --- a/optimum/habana/transformers/models/gemma2/modeling_gemma2.py +++ b/optimum/habana/transformers/models/gemma2/modeling_gemma2.py @@ -20,7 +20,6 @@ import torch import torch.nn.functional as F -from torch.nn import CrossEntropyLoss from transformers.cache_utils import Cache, DynamicCache, StaticCache from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS @@ -339,11 +338,6 @@ def pre_attn_forward( - add new args use_flash_attention - add new arg flash_attention_recompute """ - if "padding_mask" in kwargs: - logger.warning_once( - "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`" - ) - bsz, q_len, _ = hidden_states.size() query_states = self.q_proj(hidden_states) @@ -688,7 +682,7 @@ def forward( self._attn_implementation = "eager" 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") + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") elif input_ids is not None: batch_size, seq_length = input_ids.shape[:2] elif inputs_embeds is not None: @@ -765,7 +759,7 @@ def forward( if lazy_mode: htcore.mark_step() - for layer_idx, decoder_layer in enumerate(self.layers): + for layer_idx, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]): if ( lazy_mode and not self.training @@ -870,6 +864,7 @@ def forward( output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, + num_logits_to_keep: int = 0, token_idx: Optional[torch.Tensor] = None, trim_logits: Optional[bool] = False, attn_softmax_bf16: Optional[bool] = False, @@ -880,6 +875,7 @@ def forward( flash_attention_fast_softmax: Optional[bool] = False, cache_idx: int = None, lazy_mode: Optional[bool] = True, + **loss_kwargs, ) -> Union[Tuple, CausalLMOutputWithPast]: """ Inherits from GemmaForCausalLM: https://github.com/huggingface/transformers/blob/v4.38.1/src/transformers/models/gemma/modeling_gemma.py @@ -924,21 +920,11 @@ def forward( else: hidden_states = hidden_states[:, -1, :] - logits = self.lm_head(hidden_states) - logits = logits.float() + logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]) loss = None if labels is not None: - # Shift so that tokens < n predict n - shift_logits = logits[..., :-1, :].contiguous() - shift_labels = labels[..., 1:].contiguous() - # Flatten the tokens - loss_fct = CrossEntropyLoss() - shift_logits = shift_logits.view(-1, self.config.vocab_size) - shift_labels = shift_labels.view(-1) - # Enable model parallelism - shift_labels = shift_labels.to(shift_logits.device) - loss = loss_fct(shift_logits, shift_labels) + loss = self.loss_function(logits, labels, self.vocab_size, **loss_kwargs) if not return_dict: output = (logits,) + outputs[1:] diff --git a/optimum/habana/transformers/models/gpt2/modeling_gpt2.py b/optimum/habana/transformers/models/gpt2/modeling_gpt2.py index 20039bb6a5..8c226a458b 100644 --- a/optimum/habana/transformers/models/gpt2/modeling_gpt2.py +++ b/optimum/habana/transformers/models/gpt2/modeling_gpt2.py @@ -383,7 +383,8 @@ def gaudi_gpt2_forward( all_self_attentions = () if output_attentions else None all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None all_hidden_states = () if output_hidden_states else None - for i, (block, layer_past) in enumerate(zip(self.h, past_key_values)): + 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) 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 76f8f0a0c0..b5ef987752 100644 --- a/optimum/habana/transformers/models/gpt_neo/modeling_gpt_neo.py +++ b/optimum/habana/transformers/models/gpt_neo/modeling_gpt_neo.py @@ -167,7 +167,7 @@ def gaudi_gpt_neo_model_forward( 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") + 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() 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 57dfca70a0..658147afbe 100644 --- a/optimum/habana/transformers/models/gpt_neox/modeling_gpt_neox.py +++ b/optimum/habana/transformers/models/gpt_neox/modeling_gpt_neox.py @@ -208,7 +208,7 @@ def gaudi_gpt_neox_model_forward( 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: - raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + 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() diff --git a/optimum/habana/transformers/models/gptj/modeling_gptj.py b/optimum/habana/transformers/models/gptj/modeling_gptj.py index 3927e1feb9..c61f496cb3 100644 --- a/optimum/habana/transformers/models/gptj/modeling_gptj.py +++ b/optimum/habana/transformers/models/gptj/modeling_gptj.py @@ -391,7 +391,7 @@ def forward( 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") + 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() diff --git a/optimum/habana/transformers/models/idefics2/modeling_idefics2.py b/optimum/habana/transformers/models/idefics2/modeling_idefics2.py index 7b92bca9c3..b9e616fe09 100644 --- a/optimum/habana/transformers/models/idefics2/modeling_idefics2.py +++ b/optimum/habana/transformers/models/idefics2/modeling_idefics2.py @@ -195,6 +195,7 @@ def forward( attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, + use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, @@ -256,6 +257,7 @@ def forward( output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, + num_logits_to_keep: int = 0, token_idx: Optional[torch.Tensor] = None, ) -> Union[Tuple, Idefics2CausalLMOutputWithPast]: """ @@ -333,15 +335,15 @@ def forward( outputs[1] = outputs[1].to_legacy_cache() if isinstance(outputs[1], Cache) else outputs[1] hidden_states = outputs[0] - logits = self.lm_head(hidden_states) - logits = logits.float() + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]) 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[..., 1:].to(logits.device) + 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: diff --git a/optimum/habana/transformers/models/llama/modeling_llama.py b/optimum/habana/transformers/models/llama/modeling_llama.py index 67f07437a1..16fc68fcc3 100755 --- a/optimum/habana/transformers/models/llama/modeling_llama.py +++ b/optimum/habana/transformers/models/llama/modeling_llama.py @@ -3,13 +3,13 @@ from typing import List, Optional, Tuple, Union import torch -import torch.nn.functional as F from torch.distributed.distributed_c10d import ProcessGroup from transformers.activations import ACT2FN from transformers.cache_utils import Cache, DynamicCache, StaticCache 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, @@ -19,7 +19,7 @@ apply_rotary_pos_emb, logger, ) -from transformers.utils import is_torchdynamo_compiling +from transformers.processing_utils import Unpack from .... import distributed from ....distributed import parallel_state @@ -246,25 +246,8 @@ def __init__(self, config): self.act_fn = ACT2FN[config.hidden_act] def pre_mlp_forward(self, x): - if self.config.pretraining_tp > 1: - slice = self.intermediate_size // self.config.pretraining_tp - gate_proj_slices = self.gate_proj.weight.split(slice, dim=0) - up_proj_slices = self.up_proj.weight.split(slice, dim=0) - down_proj_slices = self.down_proj.weight.split(slice, dim=1) - - gate_proj = torch.cat( - [F.linear(x, gate_proj_slices[i]) for i in range(self.config.pretraining_tp)], dim=-1 - ) - up_proj = torch.cat([F.linear(x, up_proj_slices[i]) for i in range(self.config.pretraining_tp)], dim=-1) - - intermediate_states = (self.act_fn(gate_proj) * up_proj).split(slice, dim=2) - down_proj = [ - F.linear(intermediate_states[i], down_proj_slices[i]) for i in range(self.config.pretraining_tp) - ] - output = sum(down_proj) - else: - input = self.act_fn(self.gate_proj(x)) * self.up_proj(x) - output = self.down_proj(input) + input = self.act_fn(self.gate_proj(x)) * self.up_proj(x) + output = self.down_proj(input) return output def mlp_all_reduce(self, x): @@ -272,8 +255,6 @@ def mlp_all_reduce(self, x): self.down_proj.all_reduce(x) def post_mlp_forward(self, x): - if self.config.pretraining_tp > 1: - return x if hasattr(self.down_proj, "post_all_reduce"): return self.down_proj.post_all_reduce(x) return x @@ -558,35 +539,16 @@ def pre_attn_forward( """ bsz, q_len, _ = hidden_states.size() - if self.config.pretraining_tp > 1: - key_value_slicing = (self.num_key_value_heads * self.head_dim) // self.config.pretraining_tp - query_slices = self.q_proj.weight.split( - (self.num_heads * self.head_dim) // self.config.pretraining_tp, dim=0 - ) - key_slices = self.get_k_proj_weight().split(key_value_slicing, dim=0) - value_slices = self.v_proj.weight.split(key_value_slicing, dim=0) - - query_states = [F.linear(hidden_states, query_slices[i]) for i in range(self.config.pretraining_tp)] - query_states = torch.cat(query_states, dim=-1) - - key_states = [F.linear(hidden_states, key_slices[i]) for i in range(self.config.pretraining_tp)] - key_states = torch.cat(key_states, dim=-1) - - value_states = [F.linear(hidden_states, value_slices[i]) for i in range(self.config.pretraining_tp)] - value_states = torch.cat(value_states, dim=-1) - + if hasattr(self.config, "fused_qkv") and self.config.fused_qkv: + qkv_states = self.qkv_proj(hidden_states) + query_states, key_states, value_states = torch.split(qkv_states, [self.dim1, self.dim2, self.dim2], dim=-1) else: - if hasattr(self.config, "fused_qkv") and self.config.fused_qkv: - qkv_states = self.qkv_proj(hidden_states) - query_states, key_states, value_states = torch.split( - qkv_states, [self.dim1, self.dim2, self.dim2], dim=-1 - ) - else: - query_states = self.q_proj(hidden_states) - key_states = self.k_proj(hidden_states) - value_states = self.v_proj(hidden_states) + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) - query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + # use -1 to infer num_heads and num_key_value_heads as they may vary if tensor parallel is used + query_states = query_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) # TODO: update when auto mp params is enabled in DeepSpeed (cf. https://github.com/HabanaAI/DeepSpeed/blob/94309c7b5dfc1a69858f5c9f25737b2f81a332a5/deepspeed/module_inject/replace_module.py#L440) key_states = key_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) @@ -1139,9 +1101,7 @@ def forward( 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 cannot specify both input_ids and inputs_embeds at the same time, and must specify either one" - ) + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") elif input_ids is not None: batch_size, seq_length = input_ids.shape[:2] elif inputs_embeds is not None: @@ -1225,7 +1185,7 @@ def forward( if lazy_mode: htcore.mark_step() - for layer_idx, decoder_layer in enumerate(self.layers): + for layer_idx, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]): if ( lazy_mode and not self.training @@ -1357,6 +1317,7 @@ def forward( cache_idx: int = None, lazy_mode: Optional[bool] = True, num_virtual_tokens: int = None, + **kwargs: Unpack[KwargsForCausalLM], ) -> Union[Tuple, CausalLMOutputWithPast]: output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = ( @@ -1390,6 +1351,7 @@ def forward( cache_idx=cache_idx, lazy_mode=lazy_mode, num_virtual_tokens=num_virtual_tokens, + **kwargs, ) hidden_states = outputs[0] _, seq_len, _ = hidden_states.shape @@ -1399,18 +1361,8 @@ def forward( else: hidden_states = hidden_states[:, -1, :] - if self.config.pretraining_tp > 1: - lm_head_slices = self.lm_head.weight.split(self.vocab_size // self.config.pretraining_tp, dim=0) - logits = [F.linear(hidden_states, lm_head_slices[i]) for i in range(self.config.pretraining_tp)] - logits = torch.cat(logits, dim=-1) - else: - if labels is None and not is_torchdynamo_compiling(): - logger.warning_once( - "Starting from v4.46, the `logits` model output will have the same type as the model (except at train time, where it will always be FP32)" - ) - # Only compute necessary logits, and do not upcast them to float if we are not computing the loss - # TODO: remove the float() operation in v4.46 - logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]).float() + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]) loss = None if labels is not 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 6cf728d014..274387d7bf 100644 --- a/optimum/habana/transformers/models/llava_next/modeling_llava_next.py +++ b/optimum/habana/transformers/models/llava_next/modeling_llava_next.py @@ -107,7 +107,9 @@ def forward( if labels is not None: # Shift so that tokens < n predict n if attention_mask is not None: - shift_attention_mask = attention_mask[..., 1:] + # 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[:, -(logits.shape[1] - 1) :].to(logits.device) shift_logits = logits[..., :-1, :][shift_attention_mask.to(logits.device) != 0].contiguous() shift_labels = labels[..., 1:][shift_attention_mask.to(labels.device) != 0].contiguous() else: diff --git a/optimum/habana/transformers/models/mistral/modeling_mistral.py b/optimum/habana/transformers/models/mistral/modeling_mistral.py index ae864b6f47..26a8567517 100644 --- a/optimum/habana/transformers/models/mistral/modeling_mistral.py +++ b/optimum/habana/transformers/models/mistral/modeling_mistral.py @@ -39,7 +39,7 @@ MistralRMSNorm, apply_rotary_pos_emb, ) -from transformers.utils import is_torchdynamo_compiling, logging +from transformers.utils import logging from ...modeling_attn_mask_utils import ( _gaudi_prepare_4d_causal_attention_mask, @@ -245,9 +245,9 @@ def forward( key_states = self.k_proj(hidden_states) value_states = self.v_proj(hidden_states) - query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) - key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) - value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + query_states = query_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) kv_seq_len = key_states.shape[-2] if past_key_value is not None: @@ -492,7 +492,7 @@ def forward( # retrieve input_ids and inputs_embeds if input_ids is not None and inputs_embeds is not None: - raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time") + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") elif input_ids is not None: batch_size, seq_length = input_ids.shape elif inputs_embeds is not None: @@ -705,13 +705,8 @@ def forward( hidden_states = hidden_states.index_select(1, token_idx - 1) else: hidden_states = hidden_states[:, -1, :] - if labels is None and not is_torchdynamo_compiling(): - logger.warning_once( - "Starting from v4.46, the `logits` model output will have the same type as the model (except at train time, where it will always be FP32)" - ) # Only compute necessary logits, and do not upcast them to float if we are not computing the loss - # TODO: remove the float() operation in v4.46 - logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]).float() + logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]) loss = None if labels is not None: diff --git a/optimum/habana/transformers/models/mixtral/configuration_mixtral.py b/optimum/habana/transformers/models/mixtral/configuration_mixtral.py index b9121cfbd4..a22f1cc947 100644 --- a/optimum/habana/transformers/models/mixtral/configuration_mixtral.py +++ b/optimum/habana/transformers/models/mixtral/configuration_mixtral.py @@ -17,6 +17,7 @@ def __init__( num_hidden_layers=32, num_attention_heads=32, num_key_value_heads=8, + head_dim=None, hidden_act="silu", max_position_embeddings=4096 * 32, initializer_range=0.02, @@ -44,6 +45,7 @@ def __init__( num_hidden_layers, num_attention_heads, num_key_value_heads, + head_dim, hidden_act, max_position_embeddings, initializer_range, diff --git a/optimum/habana/transformers/models/mixtral/modeling_mixtral.py b/optimum/habana/transformers/models/mixtral/modeling_mixtral.py index c11d7a277a..6ae2fda6d9 100644 --- a/optimum/habana/transformers/models/mixtral/modeling_mixtral.py +++ b/optimum/habana/transformers/models/mixtral/modeling_mixtral.py @@ -29,7 +29,6 @@ import torch import torch.nn.functional as F from torch import nn -from torch.nn import CrossEntropyLoss from transformers.cache_utils import Cache, DynamicCache, StaticCache from transformers.integrations.deepspeed import is_deepspeed_available from transformers.modeling_attn_mask_utils import ( @@ -45,7 +44,7 @@ apply_rotary_pos_emb, load_balancing_loss_func, ) -from transformers.utils import is_torchdynamo_compiling, logging +from transformers.utils import logging from ..llama.modeling_llama import ( GaudiLlamaDynamicNTKScalingRotaryEmbedding, @@ -347,7 +346,7 @@ def forward( attn_output = attn_output.reshape(bsz, self.num_heads, q_len, self.head_dim).contiguous() attn_output = attn_output.transpose(1, 2).contiguous() - attn_output = attn_output.reshape(bsz, q_len, self.hidden_size) + attn_output = attn_output.reshape(bsz, q_len, -1) attn_output = self.o_proj(attn_output) @@ -588,7 +587,7 @@ def forward( # retrieve input_ids and inputs_embeds if input_ids is not None and inputs_embeds is not None: - raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time") + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") elif input_ids is not None: batch_size, seq_length = input_ids.shape elif inputs_embeds is not None: @@ -770,6 +769,7 @@ def forward( reuse_cache: Optional[bool] = None, flash_attention_recompute: Optional[bool] = False, cache_idx: int = None, + **loss_kwargs, ) -> Union[Tuple, MoeCausalLMOutputWithPast]: output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_router_logits = ( @@ -801,28 +801,12 @@ def forward( ) hidden_states = outputs[0] - if labels is None and not is_torchdynamo_compiling(): - logger.warning_once( - "Starting from v4.46, the `logits` model output will have the same type as the model (except at train time, where it will always be FP32)" - ) # Only compute necessary logits, and do not upcast them to float if we are not computing the loss - # TODO: remove the float() operation in v4.46 - logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]).float() + logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]) 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 so that tokens < n predict n - shift_logits = logits[..., :-1, :].contiguous() - shift_labels = labels[..., 1:].contiguous() - # Flatten the tokens - loss_fct = CrossEntropyLoss() - shift_logits = shift_logits.view(-1, self.config.vocab_size) - shift_labels = shift_labels.view(-1) - # Enable model parallelism - shift_labels = shift_labels.to(shift_logits.device) - loss = loss_fct(shift_logits, shift_labels) + loss = self.loss_function(logits, labels, self.vocab_size, **loss_kwargs) aux_loss = None if output_router_logits: diff --git a/optimum/habana/transformers/models/mllama/modeling_mllama.py b/optimum/habana/transformers/models/mllama/modeling_mllama.py index 7e73868249..9ecbff58bd 100644 --- a/optimum/habana/transformers/models/mllama/modeling_mllama.py +++ b/optimum/habana/transformers/models/mllama/modeling_mllama.py @@ -23,7 +23,6 @@ import torch.nn.functional as F import torch.utils.checkpoint from torch import nn -from torch.nn import CrossEntropyLoss from transformers.cache_utils import Cache from transformers.modeling_attn_mask_utils import AttentionMaskConverter from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPast, CausalLMOutputWithPast @@ -41,7 +40,6 @@ MllamaVisionEncoder, MllamaVisionEncoderLayer, MllamaVisionModel, - _prepare_4d_causal_attention_mask_with_cache_position, _prepare_aspect_ratio_attention_mask, apply_rotary_pos_emb, repeat_kv, @@ -639,9 +637,7 @@ def forward( 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 cannot specify both input_ids and inputs_embeds at the same time, and must specify either one" - ) + 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( @@ -823,7 +819,7 @@ def _update_causal_mask( ) # In case the provided `attention` mask is 2D, we generate a causal mask here (4D). - causal_mask = _prepare_4d_causal_attention_mask_with_cache_position( + causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position( attention_mask, sequence_length=sequence_length, target_length=target_length, @@ -869,6 +865,7 @@ def forward( token_idx: Optional[torch.Tensor] = None, use_flash_attention: Optional[bool] = False, flash_attention_recompute: Optional[bool] = False, + **loss_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 @@ -912,18 +909,7 @@ def forward( 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 so that tokens < n predict n - shift_logits = logits[..., :-1, :].contiguous() - shift_labels = labels[..., 1:].contiguous() - # Flatten the tokens - loss_fct = CrossEntropyLoss() - shift_logits = shift_logits.view(-1, self.config.vocab_size) - shift_labels = shift_labels.view(-1) - # Enable model parallelism - shift_labels = shift_labels.to(shift_logits.device) - loss = loss_fct(shift_logits, shift_labels) + loss = self.loss_function(logits, labels, self.vocab_size, **loss_kwargs) if not return_dict: output = (logits,) + outputs[1:] @@ -981,9 +967,7 @@ def forward( 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 cannot specify both input_ids and inputs_embeds at the same time, and must specify either one" - ) + 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( diff --git a/optimum/habana/transformers/models/opt/modeling_opt.py b/optimum/habana/transformers/models/opt/modeling_opt.py index dda2a6c204..f30a1e4435 100644 --- a/optimum/habana/transformers/models/opt/modeling_opt.py +++ b/optimum/habana/transformers/models/opt/modeling_opt.py @@ -20,6 +20,7 @@ def forward( self, attention_mask: torch.LongTensor, past_key_values_length: int = 0, + position_ids: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, ): attention_mask = attention_mask.long() @@ -42,6 +43,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, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: """ @@ -170,6 +173,7 @@ def gaudi_opt_decoder_layer_forward( 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]]]: """ @@ -187,6 +191,7 @@ def gaudi_opt_decoder_layer_forward( hidden_states, self_attn_weights, present_key_value = self.self_attn( hidden_states=hidden_states, past_key_value=past_key_value, + position_ids=position_ids, attention_mask=attention_mask, layer_head_mask=layer_head_mask, output_attentions=output_attentions, @@ -242,6 +247,7 @@ def gaudi_opt_decoder_forward( output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, + position_ids: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, ) -> Union[Tuple, BaseModelOutputWithPast]: """ @@ -342,12 +348,14 @@ def gaudi_opt_decoder_forward( 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, @@ -395,6 +403,7 @@ def gaudi_opt_model_forward( output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, + position_ids: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, ) -> Union[Tuple, BaseModelOutputWithPast]: """ @@ -413,6 +422,7 @@ def gaudi_opt_model_forward( decoder_outputs = self.decoder( input_ids=input_ids, attention_mask=attention_mask, + position_ids=position_ids, head_mask=head_mask, past_key_values=past_key_values, inputs_embeds=inputs_embeds, @@ -455,6 +465,7 @@ def forward( output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, + position_ids: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, ) -> Union[Tuple, CausalLMOutputWithPast]: output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions @@ -467,6 +478,7 @@ def forward( outputs = self.model.decoder( input_ids=input_ids, attention_mask=attention_mask, + position_ids=position_ids, head_mask=head_mask, past_key_values=past_key_values, inputs_embeds=inputs_embeds, diff --git a/optimum/habana/transformers/models/paligemma/modeling_paligemma.py b/optimum/habana/transformers/models/paligemma/modeling_paligemma.py index 84d5014135..3b2487772f 100644 --- a/optimum/habana/transformers/models/paligemma/modeling_paligemma.py +++ b/optimum/habana/transformers/models/paligemma/modeling_paligemma.py @@ -57,9 +57,7 @@ def forward( """ if (input_ids is None) ^ (inputs_embeds is not None): - raise ValueError( - "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one" - ) + 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( @@ -88,10 +86,7 @@ def forward( # Merge text and images if pixel_values is not None: - image_outputs = self.vision_tower(pixel_values.to(inputs_embeds.dtype)) - selected_image_feature = image_outputs.last_hidden_state - image_features = self.multi_modal_projector(selected_image_feature) - image_features = image_features / (self.config.hidden_size**0.5) + 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) @@ -114,9 +109,8 @@ def forward( labels = torch.where(input_ids == self.pad_token_id, self.config.ignore_index, labels) causal_mask = self._update_causal_mask( - attention_mask, token_type_ids, inputs_embeds, past_key_values, cache_position, is_training + attention_mask, token_type_ids, past_key_values, cache_position, input_ids, inputs_embeds, is_training ) - outputs = self.language_model( attention_mask=causal_mask, position_ids=position_ids, @@ -133,14 +127,16 @@ def forward( ) logits = outputs.logits - logits = logits.float() 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. - shift_attention_mask = attention_mask[..., 1:] + # 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: diff --git a/optimum/habana/transformers/models/persimmon/modeling_persimmon.py b/optimum/habana/transformers/models/persimmon/modeling_persimmon.py index 98ff8d4bbf..d76c87b2f6 100644 --- a/optimum/habana/transformers/models/persimmon/modeling_persimmon.py +++ b/optimum/habana/transformers/models/persimmon/modeling_persimmon.py @@ -241,7 +241,7 @@ def gaudi_persimmon_model_forward( # retrieve input_ids and inputs_embeds if input_ids is not None and inputs_embeds is not None: - raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time") + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") elif input_ids is not None: batch_size, seq_length = input_ids.shape elif inputs_embeds is not None: diff --git a/optimum/habana/transformers/models/phi/modeling_phi.py b/optimum/habana/transformers/models/phi/modeling_phi.py index 47875afb91..ab200d2332 100644 --- a/optimum/habana/transformers/models/phi/modeling_phi.py +++ b/optimum/habana/transformers/models/phi/modeling_phi.py @@ -24,7 +24,6 @@ import torch from torch import nn -from torch.nn import CrossEntropyLoss from transformers.cache_utils import Cache, DynamicCache from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from transformers.models.phi.configuration_phi import PhiConfig @@ -35,7 +34,7 @@ PhiModel, apply_rotary_pos_emb, ) -from transformers.utils import is_torchdynamo_compiling, logging +from transformers.utils import logging from ...modeling_attn_mask_utils import ( _gaudi_prepare_4d_causal_attention_mask, @@ -351,7 +350,7 @@ def forward( # retrieve input_ids and inputs_embeds 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") + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") elif input_ids is not None: batch_size, seq_length = input_ids.shape[:2] elif inputs_embeds is not None: @@ -491,6 +490,7 @@ def forward( reuse_cache: Optional[bool] = False, trim_logits: Optional[bool] = False, cache_idx: Optional[int] = None, + **loss_kwargs, ) -> Union[Tuple, CausalLMOutputWithPast]: """ Inherits from PhiForCausalLM: https://github.com/huggingface/transformers/blob/v4.37.1/src/transformers/models/phi/modeling_phi.py @@ -530,28 +530,12 @@ def forward( hidden_states = hidden_states.index_select(1, token_idx - 1) else: hidden_states = hidden_states[:, -1, :] - if labels is None and not is_torchdynamo_compiling(): - logger.warning_once( - "Starting from v4.46, the `logits` model output will have the same type as the model (except at train time, where it will always be FP32)" - ) # Only compute necessary logits, and do not upcast them to float if we are not computing the loss - # TODO: remove the float() operation in v4.46 - logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]).float() + logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]) 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 so that tokens < n predict n - shift_logits = logits[..., :-1, :].contiguous() - shift_labels = labels[..., 1:].contiguous() - # Flatten the tokens - loss_fct = CrossEntropyLoss() - shift_logits = shift_logits.view(-1, self.config.vocab_size) - shift_labels = shift_labels.view(-1) - # Enable model parallelism - shift_labels = shift_labels.to(shift_logits.device) - loss = loss_fct(shift_logits, shift_labels) + loss = self.loss_function(logits, labels, self.vocab_size, **loss_kwargs) if not return_dict: output = (logits,) + outputs[1:] diff --git a/optimum/habana/transformers/models/qwen2/modeling_qwen2.py b/optimum/habana/transformers/models/qwen2/modeling_qwen2.py index 73803604cf..e646188e39 100644 --- a/optimum/habana/transformers/models/qwen2/modeling_qwen2.py +++ b/optimum/habana/transformers/models/qwen2/modeling_qwen2.py @@ -34,7 +34,6 @@ apply_rotary_pos_emb, logger, ) -from transformers.utils import is_torchdynamo_compiling from ...modeling_attn_mask_utils import ( _gaudi_prepare_4d_causal_attention_mask, @@ -274,9 +273,9 @@ def pre_attn_forward( key_states = self.k_proj(hidden_states) value_states = self.v_proj(hidden_states) - query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) - key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) - value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + query_states = query_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) kv_seq_len = key_states.shape[-2] if past_key_value is not None: @@ -655,7 +654,6 @@ def forward( cache_idx: int = None, lazy_mode: Optional[bool] = True, num_virtual_tokens: int = None, - **kwargs, ) -> Union[Tuple, BaseModelOutputWithPast]: output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = ( @@ -667,7 +665,7 @@ def forward( # retrieve input_ids and inputs_embeds if input_ids is not None and inputs_embeds is not None: - raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time") + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") elif input_ids is not None: batch_size, seq_length = input_ids.shape elif inputs_embeds is not None: @@ -859,7 +857,7 @@ def forward( cache_idx: int = None, lazy_mode: Optional[bool] = True, num_virtual_tokens: int = None, - **kwargs, + **loss_kwargs, ) -> Union[Tuple, CausalLMOutputWithPast]: output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = ( @@ -893,7 +891,6 @@ def forward( cache_idx=cache_idx, lazy_mode=lazy_mode, num_virtual_tokens=num_virtual_tokens, - **kwargs, ) hidden_states = outputs[0] @@ -904,28 +901,12 @@ def forward( else: hidden_states = hidden_states[:, -1, :] - if labels is None and not is_torchdynamo_compiling(): - logger.warning_once( - "Starting from v4.46, the `logits` model output will have the same type as the model (except at train time, where it will always be FP32)" - ) # Only compute necessary logits, and do not upcast them to float if we are not computing the loss - # TODO: remove the float() operation in v4.46 - logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]).float() + logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]) 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 so that tokens < n predict n - shift_logits = logits[..., :-1, :].contiguous() - shift_labels = labels[..., 1:].contiguous() - # Flatten the tokens - loss_fct = torch.nn.CrossEntropyLoss() - shift_logits = shift_logits.view(-1, self.config.vocab_size) - shift_labels = shift_labels.view(-1) - # Enable model parallelism - shift_labels = shift_labels.to(shift_logits.device) - loss = loss_fct(shift_logits, shift_labels) + loss = self.loss_function(logits, labels, self.vocab_size, **loss_kwargs) if not return_dict: output = (logits,) + outputs[1:] 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 721abfa8ff..efddd47dc5 100755 --- a/optimum/habana/transformers/models/qwen2_moe/modeling_qwen2_moe.py +++ b/optimum/habana/transformers/models/qwen2_moe/modeling_qwen2_moe.py @@ -26,7 +26,6 @@ import habana_frameworks.torch.core as htcore import torch import torch.nn.functional as F -from torch.nn import CrossEntropyLoss from transformers.cache_utils import Cache, DynamicCache, StaticCache from transformers.integrations.deepspeed import is_deepspeed_available from transformers.modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast @@ -353,7 +352,7 @@ def pre_attn_forward( key_states = self.k_proj(hidden_states) value_states = self.v_proj(hidden_states) - query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + query_states = query_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) key_states = key_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) @@ -825,7 +824,6 @@ def forward( cache_idx: int = None, lazy_mode: Optional[bool] = True, num_virtual_tokens: int = None, - **kwargs, ) -> Union[Tuple, MoeModelOutputWithPast]: """ Copied from LlamaModel.forward: https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py @@ -850,9 +848,7 @@ def forward( 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 cannot specify both input_ids and inputs_embeds at the same time, and must specify either one" - ) + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") elif input_ids is not None: batch_size, seq_length = input_ids.shape[:2] elif inputs_embeds is not None: @@ -1061,7 +1057,7 @@ def forward( cache_idx: int = None, lazy_mode: Optional[bool] = True, num_virtual_tokens: int = None, - **kwargs, + **loss_kwargs, ) -> Union[Tuple, MoeCausalLMOutputWithPast]: output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_router_logits = ( @@ -1100,7 +1096,6 @@ def forward( cache_idx=cache_idx, lazy_mode=lazy_mode, num_virtual_tokens=num_virtual_tokens, - **kwargs, ) hidden_states = outputs[0] @@ -1111,20 +1106,11 @@ def forward( else: hidden_states = hidden_states[:, -1, :] - logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]).float() + logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]) loss = None if labels is not None: - # Shift so that tokens < n predict n - shift_logits = logits[..., :-1, :].contiguous() - shift_labels = labels[..., 1:].contiguous() - # Flatten the tokens - loss_fct = CrossEntropyLoss() - shift_logits = shift_logits.view(-1, self.config.vocab_size) - shift_labels = shift_labels.view(-1) - # Enable model parallelism - shift_labels = shift_labels.to(shift_logits.device) - loss = loss_fct(shift_logits, shift_labels) + loss = self.loss_function(logits, labels, self.vocab_size, **loss_kwargs) aux_loss = None if output_router_logits: diff --git a/optimum/habana/transformers/models/speecht5/modeling_speecht5.py b/optimum/habana/transformers/models/speecht5/modeling_speecht5.py index 07c4fa8a14..ac0fb472ae 100644 --- a/optimum/habana/transformers/models/speecht5/modeling_speecht5.py +++ b/optimum/habana/transformers/models/speecht5/modeling_speecht5.py @@ -4,6 +4,7 @@ import torch.utils.checkpoint from torch import nn from transformers.integrations.deepspeed import is_deepspeed_zero3_enabled +from transformers.integrations.fsdp import is_fsdp_managed_module from transformers.modeling_attn_mask_utils import _prepare_4d_attention_mask from transformers.modeling_outputs import BaseModelOutputWithPastAndCrossAttentions from transformers.models.speecht5.modeling_speecht5 import SpeechT5EncoderWithSpeechPrenet, SpeechT5PreTrainedModel @@ -269,7 +270,7 @@ def gaudi_SpeechT5Decoder_forward( encoder_attention_mask, hidden_states.dtype, tgt_len=input_shape[-1] ) - deepspeed_zero3_is_enabled = is_deepspeed_zero3_enabled() + synced_gpus = is_deepspeed_zero3_enabled() or is_fsdp_managed_module(self) if self.gradient_checkpointing and self.training: if use_cache: @@ -302,7 +303,7 @@ def gaudi_SpeechT5Decoder_forward( if self.training: dropout_probability = torch.rand([]) skip_the_layer = dropout_probability < self.layerdrop - if skip_the_layer and not deepspeed_zero3_is_enabled: + if skip_the_layer and not synced_gpus: continue past_key_value = past_key_values[idx] if past_key_values is not None else None diff --git a/optimum/habana/transformers/models/stablelm/modeling_stablelm.py b/optimum/habana/transformers/models/stablelm/modeling_stablelm.py index 8454995ef7..f017f38b87 100644 --- a/optimum/habana/transformers/models/stablelm/modeling_stablelm.py +++ b/optimum/habana/transformers/models/stablelm/modeling_stablelm.py @@ -257,7 +257,7 @@ def gaudi_stablelm_model_forward( # retrieve input_ids and inputs_embeds if input_ids is not None and inputs_embeds is not None: - raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time") + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") elif input_ids is not None: batch_size, seq_length = input_ids.shape elif inputs_embeds is not None: diff --git a/optimum/habana/transformers/models/starcoder2/modeling_starcoder2.py b/optimum/habana/transformers/models/starcoder2/modeling_starcoder2.py index a69526a03d..a5df50b9c3 100644 --- a/optimum/habana/transformers/models/starcoder2/modeling_starcoder2.py +++ b/optimum/habana/transformers/models/starcoder2/modeling_starcoder2.py @@ -34,7 +34,7 @@ Starcoder2Model, apply_rotary_pos_emb, ) -from transformers.utils import is_torchdynamo_compiling, logging +from transformers.utils import logging from ...modeling_attn_mask_utils import ( _gaudi_prepare_4d_causal_attention_mask, @@ -207,9 +207,9 @@ def pre_attn_forward( key_states = self.k_proj(hidden_states) value_states = self.v_proj(hidden_states) - query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) - key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) - value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + query_states = query_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) kv_seq_len = key_states.shape[-2] if past_key_value is not None: @@ -525,7 +525,7 @@ def forward( # retrieve input_ids and inputs_embeds 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") + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") elif input_ids is not None: batch_size, seq_length = input_ids.shape[:2] elif inputs_embeds is not None: @@ -693,6 +693,7 @@ def forward( flash_attention_causal_mask: Optional[bool] = False, cache_idx: int = None, lazy_mode: Optional[bool] = True, + **loss_kwargs, ) -> Union[Tuple, CausalLMOutputWithPast]: output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = ( @@ -735,28 +736,12 @@ def forward( else: hidden_states = hidden_states[:, -1, :] - if labels is None and not is_torchdynamo_compiling(): - logger.warning_once( - "Starting from v4.46, the `logits` model output will have the same type as the model (except at train time, where it will always be FP32)" - ) # Only compute necessary logits, and do not upcast them to float if we are not computing the loss - # TODO: remove the float() operation in v4.46 - logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]).float() + logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]) 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 so that tokens < n predict n - shift_logits = logits[..., :-1, :].contiguous() - shift_labels = labels[..., 1:].contiguous() - # Flatten the tokens - loss_fct = torch.nn.CrossEntropyLoss() - shift_logits = shift_logits.view(-1, self.config.vocab_size) - shift_labels = shift_labels.view(-1) - # Enable model parallelism - shift_labels = shift_labels.to(shift_logits.device) - loss = loss_fct(shift_logits, shift_labels) + loss = self.loss_function(logits, labels, self.vocab_size, **loss_kwargs) if not return_dict: output = (logits,) + outputs[1:] diff --git a/optimum/habana/transformers/models/t5/modeling_t5.py b/optimum/habana/transformers/models/t5/modeling_t5.py index b7d7f9957e..bdba215617 100644 --- a/optimum/habana/transformers/models/t5/modeling_t5.py +++ b/optimum/habana/transformers/models/t5/modeling_t5.py @@ -57,6 +57,7 @@ def gaudi_T5Attention_forward( query_length=None, use_cache=False, output_attentions=False, + cache_position=None, token_idx=None, ): # Input is (batch_size, seq_length, dim) @@ -196,6 +197,7 @@ def gaudi_T5LayerSelfAttention_forward( past_key_value=None, use_cache=False, output_attentions=False, + cache_position=None, token_idx=None, ): normed_hidden_states = self.layer_norm(hidden_states) @@ -207,6 +209,7 @@ def gaudi_T5LayerSelfAttention_forward( past_key_value=past_key_value, use_cache=use_cache, output_attentions=output_attentions, + cache_position=cache_position, token_idx=token_idx, ) hidden_states = hidden_states + self.dropout(attention_output[0]) @@ -228,6 +231,7 @@ def gaudi_T5Block_forward( use_cache=False, output_attentions=False, return_dict=True, + cache_position=None, token_idx=None, ): if past_key_value is not None: @@ -255,6 +259,7 @@ def gaudi_T5Block_forward( past_key_value=self_attn_past_key_value, use_cache=use_cache, output_attentions=output_attentions, + cache_position=cache_position, token_idx=token_idx, ) hidden_states, present_key_value_state = self_attention_outputs[:2] @@ -316,6 +321,7 @@ def gaudi_T5Stack_forward( output_attentions=None, output_hidden_states=None, return_dict=None, + cache_position=None, token_idx=None, ): use_cache = use_cache if use_cache is not None else self.config.use_cache @@ -339,6 +345,13 @@ def gaudi_T5Stack_forward( err_msg_prefix = "decoder_" if self.is_decoder else "" raise ValueError(f"You have to specify either {err_msg_prefix}input_ids or {err_msg_prefix}inputs_embeds") + 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 inputs_embeds is None: if self.embed_tokens is None: raise ValueError("You have to initialize the model with valid token embeddings") @@ -378,13 +391,6 @@ def gaudi_T5Stack_forward( else: encoder_extended_attention_mask = 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 - # Prepare head mask if needed head_mask = self.get_head_mask(head_mask, self.config.num_layers) cross_attn_head_mask = self.get_head_mask(cross_attn_head_mask, self.config.num_layers) @@ -419,6 +425,7 @@ def gaudi_T5Stack_forward( use_cache, output_attentions, True, + cache_position, None, ) else: @@ -434,6 +441,8 @@ def gaudi_T5Stack_forward( 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, ) @@ -505,6 +514,7 @@ def gaudi_T5ForConditionalGeneration_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.LongTensor] = None, ) -> Union[Tuple[torch.FloatTensor], Seq2SeqLMOutput]: use_cache = use_cache if use_cache is not None else self.config.use_cache @@ -555,6 +565,7 @@ def gaudi_T5ForConditionalGeneration_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/wav2vec2/modeling_wav2vec2.py b/optimum/habana/transformers/models/wav2vec2/modeling_wav2vec2.py index 4608a56d3f..e03d9056e7 100644 --- a/optimum/habana/transformers/models/wav2vec2/modeling_wav2vec2.py +++ b/optimum/habana/transformers/models/wav2vec2/modeling_wav2vec2.py @@ -21,6 +21,7 @@ import torch from habana_frameworks.torch.hpu import get_device_name from transformers.integrations.deepspeed import is_deepspeed_zero3_enabled +from transformers.integrations.fsdp import is_fsdp_managed_module from transformers.modeling_outputs import ( BaseModelOutput, CausalLMOutput, @@ -231,7 +232,7 @@ def gaudi_wav2vec2_encoder_forward( hidden_states = self.layer_norm(hidden_states) hidden_states = self.dropout(hidden_states) - deepspeed_zero3_is_enabled = is_deepspeed_zero3_enabled() + synced_gpus = is_deepspeed_zero3_enabled() or is_fsdp_managed_module(self) for layer in self.layers: if output_hidden_states: @@ -241,8 +242,8 @@ def gaudi_wav2vec2_encoder_forward( dropout_probability = torch.rand([]) skip_the_layer = True if self.training and (dropout_probability < self.config.layerdrop) else False - if not skip_the_layer or deepspeed_zero3_is_enabled: - # under deepspeed zero3 all gpus must run in sync + 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__, diff --git a/optimum/habana/transformers/trainer.py b/optimum/habana/transformers/trainer.py index ec7d31e3a6..5e016c79c8 100644 --- a/optimum/habana/transformers/trainer.py +++ b/optimum/habana/transformers/trainer.py @@ -30,7 +30,7 @@ import warnings from collections.abc import Mapping from pathlib import Path -from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Type, Union import huggingface_hub.utils as hf_hub_utils import numpy as np @@ -39,7 +39,6 @@ from accelerate.data_loader import SeedableRandomSampler from accelerate.utils import ( DistributedDataParallelKwargs, - GradientAccumulationPlugin, load_fsdp_model, load_fsdp_optimizer, save_fsdp_model, @@ -50,15 +49,18 @@ from transformers import Trainer from transformers.data.data_collator import DataCollator from transformers.debug_utils import DebugOption, DebugUnderflowOverflow +from transformers.feature_extraction_utils import FeatureExtractionMixin +from transformers.image_processing_utils import BaseImageProcessor from transformers.integrations import hp_params from transformers.integrations.deepspeed import ( deepspeed_load_checkpoint, is_deepspeed_available, is_deepspeed_zero3_enabled, ) -from transformers.modeling_utils import PreTrainedModel, load_sharded_checkpoint +from transformers.modeling_utils import PreTrainedModel, load_sharded_checkpoint, unwrap_model +from transformers.processing_utils import ProcessorMixin from transformers.tokenization_utils_base import PreTrainedTokenizerBase -from transformers.trainer import _get_fsdp_ckpt_kwargs, _is_peft_model +from transformers.trainer import _get_fsdp_ckpt_kwargs, _is_peft_model, safe_globals from transformers.trainer_callback import ExportableState, TrainerCallback, TrainerState from transformers.trainer_pt_utils import ( DistributedTensorGatherer, @@ -79,8 +81,8 @@ EvalPrediction, HPSearchBackend, HubStrategy, - IntervalStrategy, PredictionOutput, + SaveStrategy, TrainOutput, denumpify_detensorize, enable_full_determinism, @@ -99,10 +101,12 @@ WEIGHTS_INDEX_NAME, WEIGHTS_NAME, PushInProgress, + is_accelerate_available, is_datasets_available, is_peft_available, is_safetensors_available, ) +from transformers.utils.deprecation import deprecate_kwarg from optimum.utils import logging @@ -213,6 +217,7 @@ class GaudiTrainer(Trainer): deployment on Habana's Gaudi. """ + @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, @@ -221,11 +226,15 @@ def __init__( 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, - tokenizer: Optional[PreTrainedTokenizerBase] = 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), + 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: @@ -251,7 +260,7 @@ def __init__( data_collator, train_dataset, eval_dataset, - tokenizer, + processing_class, model_init, compute_metrics, callbacks, @@ -347,7 +356,9 @@ def _get_train_sampler(self) -> Optional[torch.utils.data.Sampler]: ) else: lengths = None - model_input_name = self.tokenizer.model_input_names[0] if self.tokenizer is not None else None + model_input_name = ( + self.processing_class.model_input_names[0] if self.processing_class is not None else None + ) return LengthGroupedSampler( self.args.train_batch_size * self.args.gradient_accumulation_steps, dataset=self.train_dataset, @@ -409,6 +420,8 @@ def create_optimizer(self): "betas": (self.args.adam_beta1, self.args.adam_beta2), "eps": self.args.adam_epsilon, } + elif self.optimizer_cls_and_kwargs is not None: + optimizer_cls, optimizer_kwargs = self.optimizer_cls_and_kwargs else: optimizer_cls, optimizer_kwargs = self.get_optimizer_cls_and_kwargs(self.args, self.model) @@ -761,10 +774,17 @@ def hpu_deepspeed_checkpointing(function, *checkpoint_args, use_reentrant: Optio # FSDP-XLA, SageMaker MP/DP, DataParallel, IPEX use_accelerator_prepare = True if model is self.model else False + if use_accelerator_prepare and self.is_fsdp_enabled: + # In case of auto_find_batch_size=True + # Remove FSDP wrapping from sub-models. + self.model = unwrap_model(self.model, recursive=True) + if delay_optimizer_creation: if use_accelerator_prepare: + # configure fsdp plugin for qlora if any self._fsdp_qlora_plugin_updates() - self.model = self.accelerator.prepare(self.model) + if self.accelerator.mixed_precision != "fp8": + self.model = self.accelerator.prepare(self.model) self.create_optimizer_and_scheduler(num_training_steps=max_steps) # prepare using `accelerator` prepare @@ -930,22 +950,22 @@ def hpu_deepspeed_checkpointing(function, *checkpoint_args, use_reentrant: Optio ) hb_profiler.start() - total_batched_samples = 0 if _is_peft_model(self.model) and self.model.peft_type == PeftType.ADALORA: self.model.base_model.peft_config[self.model.trainable_adapter_name].total_step = max_steps if max_steps < self.model.base_model.peft_config[self.model.trainable_adapter_name].tfinal: self.model.base_model.peft_config[self.model.trainable_adapter_name].tfinal = 0 + for epoch in range(epochs_trained, num_train_epochs): - epoch_iterator = train_dataloader - if hasattr(epoch_iterator, "set_epoch"): - epoch_iterator.set_epoch(epoch) + epoch_dataloader = train_dataloader + if hasattr(epoch_dataloader, "set_epoch"): + epoch_dataloader.set_epoch(epoch) # Reset the past mems state at the beginning of each epoch if necessary. if args.past_index >= 0: self._past = None steps_in_epoch = ( - len(epoch_iterator) + len(epoch_dataloader) if len_dataloader is not None else args.max_steps * args.gradient_accumulation_steps ) @@ -957,147 +977,157 @@ def hpu_deepspeed_checkpointing(function, *checkpoint_args, use_reentrant: Optio rng_to_sync = False steps_skipped = 0 if steps_trained_in_current_epoch > 0: - epoch_iterator = skip_first_batches(epoch_iterator, steps_trained_in_current_epoch) + epoch_dataloader = skip_first_batches(epoch_dataloader, steps_trained_in_current_epoch) steps_skipped = steps_trained_in_current_epoch steps_trained_in_current_epoch = 0 rng_to_sync = True step = -1 - for step, inputs in enumerate(epoch_iterator): - if ( - args.throughput_warmup_steps > 0 - and (args.throughput_warmup_steps * args.gradient_accumulation_steps) - == epoch * steps_in_epoch + step - ): - start_time_after_warmup = time.time() - - total_batched_samples += 1 - - if self.args.include_num_input_tokens_seen: - main_input_name = getattr(self.model, "main_input_name", "input_ids") - if main_input_name not in inputs: - logger.warning( - "Tried to track the number of tokens seen, however the current model is " - "not configured properly to know what item is the input. To fix this, add " - "a `main_input_name` attribute to the model class you are using." - ) + epoch_iterator = iter(epoch_dataloader) + # We chunkify the epoch iterator into gradient accumulation steps `n` batches + remainder = num_examples % 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 + 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(epoch_iterator, num_batches) + for i, inputs in enumerate(batch_samples): + step += 1 + + if ( + args.throughput_warmup_steps > 0 + and (args.throughput_warmup_steps * args.gradient_accumulation_steps) + == epoch * steps_in_epoch + step + ): + start_time_after_warmup = time.time() + + do_sync_step = (step + 1) % args.gradient_accumulation_steps == 0 or (step + 1) == steps_in_epoch + # Since we perform prefetching, we need to manually set sync_gradients + if not do_sync_step: + self.accelerator.gradient_state._set_sync_gradients(False) else: - self.state.num_input_tokens_seen += ( - torch.sum( - self.accelerator.gather( - torch.tensor( - inputs[main_input_name].numel(), device=self.args.device, dtype=torch.int64 - ) - ) + self.accelerator.gradient_state._set_sync_gradients(True) + + if self.args.include_num_input_tokens_seen: + main_input_name = getattr(self.model, "main_input_name", "input_ids") + if main_input_name not in inputs: + logger.warning( + "Tried to track the number of tokens seen, however the current model is " + "not configured properly to know what item is the input. To fix this, add " + "a `main_input_name` attribute to the model class you are using." ) - .cpu() - .item() - ) - if rng_to_sync: - self._load_rng_state(resume_from_checkpoint) - rng_to_sync = False - - # Skip past any already trained steps if resuming training - if steps_trained_in_current_epoch > 0: - steps_trained_in_current_epoch -= 1 - if steps_trained_progress_bar is not None: - steps_trained_progress_bar.update(1) - if steps_trained_in_current_epoch == 0: + 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() + ) + if rng_to_sync: self._load_rng_state(resume_from_checkpoint) - continue - elif steps_trained_progress_bar is not None: - steps_trained_progress_bar.close() - steps_trained_progress_bar = None - - if step % args.gradient_accumulation_steps == 0: - self.control = self.callback_handler.on_step_begin(args, self.state, self.control) - - # attn_softmax_bf16 and use_flash_attention is enabled only for llama, qwen2, starcoder2, gemma, baichuan and chatglm - # lazy_mode for llama, qwen2, starcoder2 and mistral - if _should_update_inputs: - inputs.update(_inputs_update) - - # TODO: keep syncs for fast DDP? - with self.accelerator.accumulate(model): - tr_loss_step = self.training_step(model, inputs) - - is_last_step_and_steps_less_than_grad_acc = ( - steps_in_epoch <= args.gradient_accumulation_steps and (step + 1) == steps_in_epoch - ) - - is_optimization_step = ( - total_batched_samples % args.gradient_accumulation_steps == 0 - or - # last step in epoch but step is always smaller than gradient_accumulation_steps - is_last_step_and_steps_less_than_grad_acc - ) - - if ( - args.parallel_mode == ParallelMode.DISTRIBUTED - and args.distribution_strategy == "fast_ddp" - and is_optimization_step - ): - all_reduce_gradients( - model, use_hpu_graphs=True - ) # use HPU graphs for gradient fusion regardless of args.use_hpu_graphs_for_training setting - - if args.logging_nan_inf_filter and (torch.isnan(tr_loss_step) or torch.isinf(tr_loss_step)): - # if loss is nan or inf simply add the average of previous logged losses - tr_loss += tr_loss / (1 + self.state.global_step - self._globalstep_last_logged) - else: - if tr_loss.device != tr_loss_step.device: - raise ValueError( - f"Calculated loss must be on the original device: {tr_loss.device} but device in use is {tr_loss_step.device}" - ) - tr_loss += tr_loss_step + rng_to_sync = False + + # Skip past any already trained steps if resuming training + if steps_trained_in_current_epoch > 0: + steps_trained_in_current_epoch -= 1 + if steps_trained_progress_bar is not None: + steps_trained_progress_bar.update(1) + if steps_trained_in_current_epoch == 0: + self._load_rng_state(resume_from_checkpoint) + continue + elif steps_trained_progress_bar is not None: + steps_trained_progress_bar.close() + steps_trained_progress_bar = None + + if step % args.gradient_accumulation_steps == 0: + self.control = self.callback_handler.on_step_begin(args, self.state, self.control) + + # attn_softmax_bf16 and use_flash_attention is enabled only for llama, qwen2, starcoder2, gemma, baichuan and chatglm + # lazy_mode for llama, qwen2, starcoder2 and mistral + if _should_update_inputs: + inputs.update(_inputs_update) + + # TODO: keep syncs for fast DDP? + # We explicitly want to avoid relying on `accelerator.accumulate` for generation training + context = ( + functools.partial(self.accelerator.no_sync, model=model) + if i != len(batch_samples) - 1 + and self.accelerator.distributed_type != GaudiDistributedType.DEEPSPEED + else contextlib.nullcontext + ) + with context(): + tr_loss_step = self.training_step(model, inputs, num_items_in_batch) + + if ( + args.parallel_mode == ParallelMode.DISTRIBUTED + and args.distribution_strategy == "fast_ddp" + and do_sync_step + ): + all_reduce_gradients( + model, use_hpu_graphs=True + ) # use HPU graphs for gradient fusion regardless of args.use_hpu_graphs_for_training setting + + if args.logging_nan_inf_filter and (torch.isnan(tr_loss_step) or torch.isinf(tr_loss_step)): + # if loss is nan or inf simply add the average of previous logged losses + tr_loss += tr_loss / (1 + self.state.global_step - self._globalstep_last_logged) + else: + if tr_loss.device != tr_loss_step.device: + raise ValueError( + f"Calculated loss must be on the original device: {tr_loss.device} but device in use is {tr_loss_step.device}" + ) + tr_loss += tr_loss_step - self.current_flos += float(self.floating_point_ops(inputs)) - if args.use_lazy_mode: - self.htcore.mark_step() + self.current_flos += float(self.floating_point_ops(inputs)) + if args.use_lazy_mode: + self.htcore.mark_step() - if is_optimization_step: - # the `or` condition of `is_last_step_and_steps_less_than_grad_acc` is not covered - # in accelerate. So, explicitly enable sync gradients to True in that case. - if is_last_step_and_steps_less_than_grad_acc: + if do_sync_step: + # Since we perform prefetching, we need to manually set sync_gradients to True self.accelerator.gradient_state._set_sync_gradients(True) - # If the condition is true, we need to compute _grad_norm - if _should_compute_grad_norm: - # deepspeed does its own clipping - if self.gaudi_config.use_fused_clip_norm and args.use_habana: - # 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, - ) - - self.control = self.callback_handler.on_pre_optimizer_step(args, self.state, self.control) + # If the condition is true, we need to compute _grad_norm + if _should_compute_grad_norm: + # deepspeed does its own clipping + if self.gaudi_config.use_fused_clip_norm and args.use_habana: + # 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, + ) - self.optimizer.step() + self.control = self.callback_handler.on_pre_optimizer_step(args, self.state, self.control) - self.control = self.callback_handler.on_optimizer_step(args, self.state, self.control) + self.optimizer.step() - optimizer_was_run = not self.accelerator.optimizer_step_was_skipped - if optimizer_was_run: - # Delay optimizer scheduling until metrics are generated - if not isinstance(self.lr_scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau): - self.lr_scheduler.step() + self.control = self.callback_handler.on_optimizer_step(args, self.state, self.control) - self._zero_model_grad(model) - self.state.global_step += 1 - self.state.epoch = epoch + (step + 1 + steps_skipped) / steps_in_epoch - if args.use_lazy_mode: - self.htcore.mark_step() - self.control = self.callback_handler.on_step_end(args, self.state, self.control) + optimizer_was_run = not self.accelerator.optimizer_step_was_skipped + if optimizer_was_run: + # Delay optimizer scheduling until metrics are generated + if not isinstance(self.lr_scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau): + self.lr_scheduler.step() - self._maybe_log_save_evaluate(tr_loss, _grad_norm, model, trial, epoch, ignore_keys_for_eval) - else: - self.control = self.callback_handler.on_substep_end(args, self.state, self.control) + self._zero_model_grad(model) + self.state.global_step += 1 + self.state.epoch = epoch + (step + 1 + steps_skipped) / steps_in_epoch + if args.use_lazy_mode: + 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 + ) + else: + self.control = self.callback_handler.on_substep_end(args, self.state, self.control) - hb_profiler.step() + hb_profiler.step() + if self.control.should_epoch_stop or self.control.should_training_stop: + break + # We also need to break out of the nested loop if self.control.should_epoch_stop or self.control.should_training_stop: break if step < 0: @@ -1109,7 +1139,7 @@ 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) + self._maybe_log_save_evaluate(tr_loss, _grad_norm, model, trial, epoch, ignore_keys_for_eval, start_time) if self.control.should_training_stop: break @@ -1264,7 +1294,7 @@ 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): + def _maybe_log_save_evaluate(self, tr_loss, _grad_norm, model, trial, epoch, ignore_keys_for_eval, start_time): if self.args.adjust_throughput: save_start = time.perf_counter() @@ -1303,14 +1333,18 @@ def _maybe_log_save_evaluate(self, tr_loss, _grad_norm, model, trial, epoch, ign self._globalstep_last_logged = self.state.global_step self.store_flos() - self.log(logs) + self.log(logs, start_time) metrics = None if self.control.should_evaluate: metrics = self._evaluate(trial, ignore_keys_for_eval) + is_new_best_metric = self._determine_best_metric(metrics=metrics, trial=trial) + + if self.args.save_strategy == SaveStrategy.BEST: + self.control.should_save = is_new_best_metric if self.control.should_save: - self._save_checkpoint(model, trial, metrics=metrics) + self._save_checkpoint(model, trial) self.control = self.callback_handler.on_save(self.args, self.state, self.control) if self.args.adjust_throughput: @@ -1339,7 +1373,8 @@ def _load_rng_state(self, checkpoint): ) return - checkpoint_rng_state = torch.load(rng_file) + with safe_globals(): + 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"]) @@ -1475,18 +1510,22 @@ 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]) -> None: + def log(self, logs: Dict[str, float], start_time: Optional[float] = None) -> None: """ Log `logs` on the various objects watching training. Subclass and override this method to inject custom behavior. Args: logs (`Dict[str, float]`): The values to log. + start_time (`Optional[float]`): + The start of training. """ if self.state.epoch is not None: logs["epoch"] = self.state.epoch 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) mem_stats = get_hpu_memory_stats(self.args.device) logs.update(mem_stats) @@ -1545,7 +1584,9 @@ 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]]) -> torch.Tensor: + def training_step( + 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. @@ -1570,7 +1611,7 @@ def training_step(self, model: torch.nn.Module, inputs: Dict[str, Union[torch.Te inputs = self._prepare_inputs(inputs) with self.compute_loss_context_manager(): - loss = self.compute_loss(model, inputs) + loss = self.compute_loss(model, inputs, num_items_in_batch=num_items_in_batch) del inputs kwargs = {} @@ -1585,6 +1626,10 @@ def training_step(self, model: torch.nn.Module, inputs: Dict[str, Union[torch.Te 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 num_items_in_batch is None: + loss = loss / self.args.gradient_accumulation_steps + if _is_peft_model(self.model) and self.model.peft_type == PeftType.ADALORA: assert not ( self.accelerator.state.is_fp8_enabled and self.args.gradient_checkpointing @@ -1606,7 +1651,7 @@ def training_step(self, model: torch.nn.Module, inputs: Dict[str, Union[torch.Te self.accelerator.backward(loss, **kwargs) else: self.accelerator.backward(loss, **kwargs) - return loss.detach() / self.args.gradient_accumulation_steps + return loss.detach() def save_model(self, output_dir: Optional[str] = None, _internal_call: bool = False): """ @@ -1683,8 +1728,8 @@ def _save(self, output_dir: Optional[str] = None, state_dict=None): output_dir, state_dict=state_dict, safe_serialization=self.args.save_safetensors ) - if self.tokenizer is not None: - self.tokenizer.save_pretrained(output_dir) + if self.processing_class is not None: + self.processing_class.save_pretrained(output_dir) self.gaudi_config.save_pretrained(output_dir) @@ -1838,7 +1883,7 @@ def evaluation_loop( start_time = time.time() model = ( self.accelerator.prepare(model) - if self.is_deepspeed_enabled + if self.is_deepspeed_enabled or (self.is_fsdp_enabled and self.accelerator.mixed_precision != "fp8") else self.accelerator.prepare_model(model, evaluation_mode=True) ) self.model_preparation_time = round(time.time() - start_time, 4) @@ -1899,6 +1944,7 @@ def evaluation_loop( all_inputs = EvalLoopContainer(self.args.eval_do_concat_batches, padding_index=-100) metrics = None + eval_set_kwargs = {} # Will be useful when we have an iterable dataset so don't know its length. observed_num_examples = 0 @@ -1935,7 +1981,9 @@ def evaluation_loop( # Prediction step losses, logits, labels = self.prediction_step(model, inputs, prediction_loss_only, ignore_keys=ignore_keys) main_input_name = getattr(self.model, "main_input_name", "input_ids") - inputs_decode = self._prepare_input(inputs[main_input_name]) if args.include_inputs_for_metrics else None + inputs_decode = ( + self._prepare_input(inputs[main_input_name]) if "inputs" in args.include_for_metrics else None + ) # Update containers if losses is not None: @@ -1973,16 +2021,13 @@ def evaluation_loop( if self.args.batch_eval_metrics: if self.compute_metrics is not None and logits is not None and labels is not None: is_last_step = self.accelerator.gradient_state.end_of_dataloader - if args.include_inputs_for_metrics: - metrics = self.compute_metrics( - EvalPrediction(predictions=logits, label_ids=labels, inputs=inputs), - compute_result=is_last_step, - ) - else: - metrics = self.compute_metrics( - EvalPrediction(predictions=logits, label_ids=labels), - compute_result=is_last_step, - ) + batch_kwargs = {} + batch_kwargs["losses"] = losses if "loss" in args.include_for_metrics else None + batch_kwargs["inputs"] = inputs if "inputs" in args.include_for_metrics else None + metrics = self.compute_metrics( + EvalPrediction(predictions=logits, label_ids=labels, **batch_kwargs), + compute_result=is_last_step, + ) del losses, logits, labels, inputs @@ -2038,12 +2083,11 @@ def evaluation_loop( and all_labels is not None and not self.args.batch_eval_metrics ): - if args.include_inputs_for_metrics: - metrics = self.compute_metrics( - EvalPrediction(predictions=all_preds, label_ids=all_labels, inputs=all_inputs) - ) - else: - metrics = self.compute_metrics(EvalPrediction(predictions=all_preds, label_ids=all_labels)) + eval_set_kwargs["losses"] = all_losses if "loss" in args.include_for_metrics else None + eval_set_kwargs["inputs"] = all_inputs if "inputs" in args.include_for_metrics else None + metrics = self.compute_metrics( + EvalPrediction(predictions=all_preds, label_ids=all_labels, **eval_set_kwargs) + ) elif metrics is None: metrics = {} @@ -2182,13 +2226,13 @@ def _push_from_checkpoint(self, checkpoint_folder): for modeling_file in modeling_files: if os.path.isfile(os.path.join(checkpoint_folder, modeling_file)): shutil.copy(os.path.join(checkpoint_folder, modeling_file), os.path.join(output_dir, modeling_file)) - # Saving the tokenizer is fast and we don't know how many files it may have spawned, so we resave it to be sure. - if self.tokenizer is not None: - self.tokenizer.save_pretrained(output_dir) + # Saving the processing class is fast and we don't know how many files it may have spawned, so we resave it to be sure. + if self.processing_class is not None: + self.processing_class.save_pretrained(output_dir) # Same for the training arguments torch.save(self.args, os.path.join(output_dir, TRAINING_ARGS_NAME)) - if self.args.save_strategy == IntervalStrategy.STEPS: + if self.args.save_strategy == SaveStrategy.STEPS: commit_message = f"Training in progress, step {self.state.global_step}" else: commit_message = f"Training in progress, epoch {int(self.state.epoch)}" @@ -2254,7 +2298,7 @@ def prediction_loop( if len(self.accelerator._models) == 0 and model is self.model: model = ( self.accelerator.prepare(model) - if self.is_deepspeed_enabled + if self.is_deepspeed_enabled or self.is_fsdp_enabled else self.accelerator.prepare_model(model, evaluation_mode=True) ) @@ -2293,7 +2337,17 @@ def prediction_loop( elif args.bf16_full_eval: model = model.to(dtype=torch.bfloat16, device=args.device) - batch_size = dataloader.batch_size + batch_size = ( + dataloader.total_batch_size + if getattr(dataloader, "_is_accelerate_prepared", False) + else dataloader.batch_size + ) + + if batch_size is None: + raise ValueError( + "Batch size cannot be None. Ensure the dataloader has a valid batch_size or total_batch_size." + ) + num_examples = self.num_examples(dataloader) logger.info(f"\n***** Running {description} *****") logger.info(f" Num examples = {num_examples}") @@ -2304,6 +2358,7 @@ def prediction_loop( labels_host: Union[torch.Tensor, List[torch.Tensor]] = None inputs_host: Union[torch.Tensor, List[torch.Tensor]] = None metrics: Optional[dict] = None + eval_set_kwargs: dict = {} world_size = max(1, args.world_size) @@ -2326,7 +2381,9 @@ def prediction_loop( for step, inputs in enumerate(dataloader): loss, logits, labels = self.prediction_step(model, inputs, prediction_loss_only, ignore_keys=ignore_keys) main_input_name = getattr(self.model, "main_input_name", "input_ids") - inputs_decode = self._prepare_input(inputs[main_input_name]) if args.include_inputs_for_metrics else None + inputs_decode = ( + self._prepare_input(inputs[main_input_name]) if "inputs" in args.include_for_metrics else None + ) if loss is not None: losses = loss.repeat(batch_size) @@ -2346,16 +2403,13 @@ def prediction_loop( if self.args.batch_eval_metrics: if self.compute_metrics is not None and preds_host is not None and labels_host is not None: is_last_step = self.accelerator.gradient_state.end_of_dataloader - if args.include_inputs_for_metrics: - metrics = self.compute_metrics( - EvalPrediction(predictions=preds_host, label_ids=labels_host, inputs=inputs_host), - compute_result=is_last_step, - ) - else: - metrics = self.compute_metrics( - EvalPrediction(predictions=preds_host, label_ids=labels_host), - compute_result=is_last_step, - ) + batch_kwargs = {} + batch_kwargs["losses"] = losses_host if "loss" in args.include_for_metrics else None + batch_kwargs["inputs"] = inputs_host if "inputs" in args.include_for_metrics else None + metrics = self.compute_metrics( + EvalPrediction(predictions=preds_host, label_ids=labels_host, **batch_kwargs), + compute_result=is_last_step, + ) if self.args.batch_eval_metrics or ( args.eval_accumulation_steps is not None and (step + 1) % args.eval_accumulation_steps == 0 @@ -2398,12 +2452,9 @@ def prediction_loop( and label_ids is not None and not self.args.batch_eval_metrics ): - if args.include_inputs_for_metrics: - metrics = self.compute_metrics( - EvalPrediction(predictions=preds, label_ids=label_ids, inputs=inputs_ids) - ) - else: - metrics = self.compute_metrics(EvalPrediction(predictions=preds, label_ids=label_ids)) + eval_set_kwargs["losses"] = eval_loss if "loss" in args.include_for_metrics else None + eval_set_kwargs["inputs"] = inputs_ids if "inputs" in args.include_for_metrics else None + metrics = self.compute_metrics(EvalPrediction(predictions=preds, label_ids=label_ids, **eval_set_kwargs)) elif metrics is None: metrics = {} @@ -2421,24 +2472,21 @@ def prediction_loop( return EvalLoopOutput(predictions=preds, label_ids=label_ids, metrics=metrics, num_samples=num_examples) def create_accelerator_and_postprocess(self): + # We explicitly don't rely on the `Accelerator` to do gradient accumulation grad_acc_kwargs = {} if self.args.accelerator_config.gradient_accumulation_kwargs is not None: grad_acc_kwargs = self.args.accelerator_config.gradient_accumulation_kwargs # check if num_steps is attempted to be passed in gradient_accumulation_kwargs - if "num_steps" in grad_acc_kwargs and self.args.gradient_accumulation_steps > 1: - # raise because we do not know which setting is intended. - raise ValueError( - "The `AcceleratorConfig`'s `num_steps` is set but `gradient_accumulation_steps` is greater than 1 in the passed `TrainingArguments`" - "If using the passed `AcceleratorConfig` is desired, do not set the `TrainingArguments` `gradient_accumulation_steps`." - ) - elif "num_steps" not in grad_acc_kwargs: - # take the gradient_accumulation_steps setting from TrainingArguments. - grad_acc_kwargs["num_steps"] = self.args.gradient_accumulation_steps - - grad_acc_kwargs["sync_with_dataloader"] = False - - gradient_accumulation_plugin = GradientAccumulationPlugin(**grad_acc_kwargs) + if "num_steps" in grad_acc_kwargs: + if self.args.gradient_accumulation_steps > 1: + # raise because we do not know which setting is intended. + raise ValueError( + "The `AcceleratorConfig`'s `num_steps` is set but `gradient_accumulation_steps` is greater than 1 in the passed `TrainingArguments`" + "If using the passed `AcceleratorConfig` is desired, do not set the `TrainingArguments` `gradient_accumulation_steps`." + ) + else: + self.args.gradient_accumulation_steps = grad_acc_kwargs["num_steps"] accelerator_config = self.args.accelerator_config.to_dict() @@ -2448,6 +2496,8 @@ def create_accelerator_and_postprocess(self): even_batches=accelerator_config.pop("even_batches"), use_seedable_sampler=accelerator_config.pop("use_seedable_sampler"), ) + if is_accelerate_available("1.1.0"): + dataloader_config.data_seed = self.args.data_seed non_blocking = accelerator_config.pop("non_blocking") if non_blocking and not self.args.dataloader_pin_memory: logger.warning( @@ -2459,7 +2509,6 @@ def create_accelerator_and_postprocess(self): args = { "deepspeed_plugin": self.args.deepspeed_plugin, - "gradient_accumulation_plugin": gradient_accumulation_plugin, "distribution_strategy": self.args.distribution_strategy, "dynamic": self.args.compile_dynamic, "dataloader_config": dataloader_config, diff --git a/optimum/habana/transformers/trainer_seq2seq.py b/optimum/habana/transformers/trainer_seq2seq.py index 7a327b5a7b..0864d819b3 100644 --- a/optimum/habana/transformers/trainer_seq2seq.py +++ b/optimum/habana/transformers/trainer_seq2seq.py @@ -13,14 +13,19 @@ # See the License for the specific language governing permissions and # limitations under the License. +import contextlib import warnings from copy import deepcopy from pathlib import Path from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union import torch +from torch.distributed.fsdp import FullyShardedDataParallel from torch.utils.data import Dataset from transformers.integrations.deepspeed import is_deepspeed_zero3_enabled +from transformers.integrations.fsdp import is_fsdp_managed_module +from transformers.utils import is_datasets_available +from transformers.utils.deprecation import deprecate_kwarg from optimum.utils import logging @@ -28,9 +33,17 @@ from .trainer import GaudiTrainer +if is_datasets_available(): + import datasets + + if TYPE_CHECKING: + from torch.utils.data import IterableDataset from transformers.data.data_collator import DataCollator + from transformers.feature_extraction_utils import FeatureExtractionMixin + from transformers.image_processing_utils import BaseImageProcessor from transformers.modeling_utils import PreTrainedModel + from transformers.processing_utils import ProcessorMixin from transformers.tokenization_utils_base import PreTrainedTokenizerBase from transformers.trainer_callback import TrainerCallback from transformers.trainer_utils import EvalPrediction, PredictionOutput @@ -43,15 +56,18 @@ class GaudiSeq2SeqTrainer(GaudiTrainer): + @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, gaudi_config: "GaudiConfig" = None, args: "GaudiTrainingArguments" = None, data_collator: Optional["DataCollator"] = None, - train_dataset: Optional[Dataset] = None, + train_dataset: Optional[Union[Dataset, "IterableDataset", "datasets.Dataset"]] = None, eval_dataset: Optional[Union[Dataset, Dict[str, Dataset]]] = None, - tokenizer: Optional["PreTrainedTokenizerBase"] = None, + processing_class: Optional[ + Union["PreTrainedTokenizerBase", "BaseImageProcessor", "FeatureExtractionMixin", "ProcessorMixin"] + ] = None, model_init: Optional[Callable[[], "PreTrainedModel"]] = None, compute_metrics: Optional[Callable[["EvalPrediction"], Dict]] = None, callbacks: Optional[List["TrainerCallback"]] = None, @@ -65,7 +81,7 @@ def __init__( data_collator=data_collator, train_dataset=train_dataset, eval_dataset=eval_dataset, - tokenizer=tokenizer, + processing_class=processing_class, model_init=model_init, compute_metrics=compute_metrics, callbacks=callbacks, @@ -281,10 +297,8 @@ def prediction_step( if "max_length" in gen_kwargs and gen_kwargs["max_length"] is None: gen_kwargs.pop("max_length") - default_synced_gpus = True if is_deepspeed_zero3_enabled() else False - gen_kwargs["synced_gpus"] = ( - gen_kwargs["synced_gpus"] if gen_kwargs.get("synced_gpus") is not None else default_synced_gpus - ) + default_synced_gpus = is_deepspeed_zero3_enabled() or is_fsdp_managed_module(self.model) + gen_kwargs["synced_gpus"] = gen_kwargs.get("synced_gpus", default_synced_gpus) # pad batches to max_length on-the-fly in lazy mode gen_kwargs["lazy_mode"] = ( gen_kwargs["lazy_mode"] if gen_kwargs.get("lazy_mode") is not None else self.args.use_lazy_mode @@ -309,8 +323,18 @@ def prediction_step( generation_inputs = { k: v for k, v in inputs.items() if k not in ("decoder_input_ids", "decoder_attention_mask") } + + summon_full_params_context = ( + FullyShardedDataParallel.summon_full_params(self.model) + if isinstance(self.model, FullyShardedDataParallel) + else contextlib.nullcontext() + ) + try: - with torch.autocast(device_type="hpu", dtype=torch.bfloat16, enabled=self.use_hpu_amp): + with ( + torch.autocast(device_type="hpu", dtype=torch.bfloat16, enabled=self.use_hpu_amp), + summon_full_params_context, + ): generated_tokens = self.model.generate( **generation_inputs, generation_config=self.model.generation_config, diff --git a/optimum/habana/transformers/training_args.py b/optimum/habana/transformers/training_args.py index 4a2b12593f..56fdb1d154 100644 --- a/optimum/habana/transformers/training_args.py +++ b/optimum/habana/transformers/training_args.py @@ -26,7 +26,14 @@ 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 -from transformers.trainer_utils import EvaluationStrategy, FSDPOption, HubStrategy, IntervalStrategy, SchedulerType +from transformers.trainer_utils import ( + EvaluationStrategy, + FSDPOption, + HubStrategy, + IntervalStrategy, + SaveStrategy, + SchedulerType, +) from transformers.training_args import ( _VALID_DICT_FIELDS, OptimizerNames, @@ -409,7 +416,7 @@ def __post_init__(self): self.eval_strategy = IntervalStrategy(self.eval_strategy) self.logging_strategy = IntervalStrategy(self.logging_strategy) - self.save_strategy = IntervalStrategy(self.save_strategy) + self.save_strategy = SaveStrategy(self.save_strategy) self.hub_strategy = HubStrategy(self.hub_strategy) self.lr_scheduler_type = SchedulerType(self.lr_scheduler_type) @@ -445,7 +452,7 @@ def __post_init__(self): if self.eval_steps != int(self.eval_steps): raise ValueError(f"--eval_steps must be an integer if bigger than 1: {self.eval_steps}") self.eval_steps = int(self.eval_steps) - if self.save_strategy == IntervalStrategy.STEPS and self.save_steps > 1: + if self.save_strategy == SaveStrategy.STEPS and self.save_steps > 1: if self.save_steps != int(self.save_steps): raise ValueError(f"--save_steps must be an integer if bigger than 1: {self.save_steps}") self.save_steps = int(self.save_steps) @@ -553,6 +560,19 @@ def __post_init__(self): if self.dataloader_drop_last: self.accelerator_config.even_batches = False + # 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." + ) + self.average_tokens_across_devices = False + except ImportError as e: + logger.warning(f"Can not specify world size due to {e}. Turn average_tokens_across_devices to False.") + self.average_tokens_across_devices = False + if (self.torch_compile_mode is not None or self.torch_compile_backend is not None) and not self.torch_compile: assert get_habana_frameworks_version().minor > 12, "Torch compile is not available" self.torch_compile = True @@ -683,7 +703,7 @@ def __post_init__(self): self.fsdp_config["xla_fsdp_grad_ckpt"] = self.fsdp_config.get("xla_fsdp_grad_ckpt", False) # accelerate integration for FSDP - if len(self.fsdp) > 0: + if len(self.fsdp) > 0 and not self.fsdp_config["xla"]: os.environ["ACCELERATE_USE_FSDP"] = "true" from accelerate.utils.constants import ( FSDP_AUTO_WRAP_POLICY, @@ -825,6 +845,19 @@ def __post_init__(self): "This is not supported and we recommend you to update your version." ) + if self.data_seed is not None: + if not is_accelerate_available("1.1.0"): + raise NotImplementedError( + "data_seed requires Accelerate version `accelerate` >= 1.1.0. " + "This is not supported and we recommend you to update your version." + ) + + if self.include_inputs_for_metrics: + logger.warning( + "Using `include_inputs_for_metrics` is deprecated and will be removed in version 5 of 🤗 Transformers. Please use `include_for_metrics` list argument instead." + ) + self.include_for_metrics.append("inputs") + def __str__(self): self_as_dict = asdict(self) diff --git a/tests/test_trainer.py b/tests/test_trainer.py index eddb82b500..61e5daf198 100644 --- a/tests/test_trainer.py +++ b/tests/test_trainer.py @@ -28,12 +28,14 @@ from typing import Dict, List, Optional, Union import numpy as np -from huggingface_hub import HfFolder, ModelCard, create_branch, delete_repo, list_repo_commits, list_repo_files +from huggingface_hub import HfFolder, ModelCard, create_branch, list_repo_commits, list_repo_files from parameterized import parameterized from pytest import mark -from requests.exceptions import HTTPError from transformers import ( + AutoFeatureExtractor, + AutoImageProcessor, AutoModelForCausalLM, + AutoProcessor, AutoTokenizer, GPT2LMHeadModel, IntervalStrategy, @@ -50,6 +52,7 @@ USER, CaptureLogger, LoggingLevel, + TemporaryHubRepo, TestCasePlus, get_gpu_count, get_tests_dir, @@ -62,6 +65,7 @@ require_tensorboard, require_tokenizers, require_torch, + require_vision, ) from transformers.trainer_pt_utils import AcceleratorConfig from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR, HPSearchBackend @@ -659,7 +663,7 @@ def test_model_init(self): def test_gradient_accumulation(self): with tempfile.TemporaryDirectory() as tmpdir: - # Training with half the batch size but accumulation steps as 2 should give the same results. + # Training with half the batch size but accumulation steps as 2 should give the same training losses. trainer = get_regression_trainer( output_dir=tmpdir, gradient_accumulation_steps=2, per_device_train_batch_size=4, learning_rate=0.1 ) @@ -1051,14 +1055,18 @@ def test_multiple_peft_adapters(self): use_lazy_mode=True, ) gaudi_config = get_gaudi_config() - trainer = GaudiTrainer(tiny_model, gaudi_config, args, tokenizer=tokenizer, train_dataset=train_dataset) + trainer = GaudiTrainer( + tiny_model, gaudi_config, args, processing_class=tokenizer, train_dataset=train_dataset + ) trainer.train() parameters = dict(tiny_model.named_parameters()) state = dataclasses.asdict(trainer.state) # Reinitialize trainer - trainer = GaudiTrainer(tiny_model, gaudi_config, args, tokenizer=tokenizer, train_dataset=train_dataset) + trainer = GaudiTrainer( + tiny_model, gaudi_config, args, processing_class=tokenizer, train_dataset=train_dataset + ) checkpoint = os.path.join(tmpdir, "checkpoint-5") @@ -2455,9 +2463,6 @@ def test_accelerator_config_from_dict(self): self.assertEqual(trainer.accelerator.even_batches, False) self.assertEqual(trainer.accelerator.use_seedable_sampler, True) - if GRAD_ACCUM_KWARGS_VERSION_AVAILABLE: - self.assertEqual(trainer.accelerator.gradient_state.plugin_kwargs["sync_each_batch"], True) - def test_accelerator_config_from_yaml(self): # Checks that accelerator kwargs can be passed through # and the accelerator is initialized respectively @@ -2470,8 +2475,6 @@ def test_accelerator_config_from_yaml(self): "even_batches": False, "use_seedable_sampler": False, } - if GRAD_ACCUM_KWARGS_VERSION_AVAILABLE: - accelerator_config["gradient_accumulation_kwargs"] = {"sync_each_batch": True} json.dump(accelerator_config, f) config = RegressionModelConfig(a=1.5, b=2.5) model = RegressionPreTrainedModel(config) @@ -2486,9 +2489,6 @@ def test_accelerator_config_from_yaml(self): self.assertEqual(trainer.accelerator.even_batches, False) self.assertEqual(trainer.accelerator.use_seedable_sampler, False) - if GRAD_ACCUM_KWARGS_VERSION_AVAILABLE: - self.assertEqual(trainer.accelerator.gradient_state.plugin_kwargs["sync_each_batch"], True) - def test_accelerator_config_from_dataclass(self): # Checks that accelerator kwargs can be passed through # and the accelerator is initialized respectively @@ -2540,10 +2540,7 @@ def test_accelerate_config_from_dataclass_grad_accum(self): output_dir=tmp_dir, accelerator_config=accelerator_config, use_habana=True ) trainer = GaudiTrainer(model=model, gaudi_config=gaudi_config, args=args, eval_dataset=eval_dataset) - self.assertEqual(trainer.accelerator.gradient_state.plugin_kwargs["num_steps"], 10) - self.assertEqual(trainer.accelerator.gradient_state.plugin_kwargs["adjust_scheduler"], False) - self.assertEqual(trainer.accelerator.gradient_state.plugin_kwargs["sync_with_dataloader"], False) - self.assertEqual(trainer.accelerator.gradient_state.plugin_kwargs["sync_each_batch"], True) + self.assertEqual(trainer.args.gradient_accumulation_steps, 10) def test_accelerator_config_from_partial(self): # Checks that accelerator kwargs can be passed through @@ -2754,6 +2751,191 @@ def test_eval_use_gather_object(self): _ = trainer.evaluate() _ = trainer.predict(eval_dataset) + def test_trainer_saves_tokenizer(self): + MODEL_ID = "google-bert/bert-base-uncased" + tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=False) + + with tempfile.TemporaryDirectory() as tmp_dir: + config = RegressionModelConfig(a=1.5, b=2.5) + gaudi_config = get_gaudi_config() + trainer = GaudiTrainer( + model=RegressionPreTrainedModel(config), + args=GaudiTrainingArguments(output_dir=tmp_dir, use_habana=True, use_lazy_mode=True), + gaudi_config=gaudi_config, + processing_class=tokenizer, + ) + trainer.save_model() + + reloaded_tokenizer = AutoTokenizer.from_pretrained(tmp_dir) + + # For tokenizers, there isn't a direct to_dict method and the properties stored in the configs e.g. + # saved tokens change overtime, so we check that two tokenizers are equal by comparing their encoded outputs + test_sentence = "This is a test sentence" + self.assertListEqual( + tokenizer(test_sentence, padding="max_length").input_ids, + reloaded_tokenizer(test_sentence, padding="max_length").input_ids, + ) + + @require_vision + def test_trainer_saves_image_processor(self): + MODEL_ID = "openai/clip-vit-base-patch32" + image_processor = AutoImageProcessor.from_pretrained(MODEL_ID) + + with tempfile.TemporaryDirectory() as tmp_dir: + config = RegressionModelConfig(a=1.5, b=2.5) + gaudi_config = get_gaudi_config() + trainer = GaudiTrainer( + model=RegressionPreTrainedModel(config), + args=GaudiTrainingArguments(output_dir=tmp_dir, use_habana=True, use_lazy_mode=True), + gaudi_config=gaudi_config, + processing_class=image_processor, + ) + trainer.save_model() + reloaded_image_processor = AutoImageProcessor.from_pretrained(tmp_dir) + + self.assertDictEqual(image_processor.to_dict(), reloaded_image_processor.to_dict()) + + def test_trainer_saves_feature_extractor(self): + MODEL_ID = "facebook/wav2vec2-base-960h" + feature_extractor = AutoFeatureExtractor.from_pretrained(MODEL_ID) + + with tempfile.TemporaryDirectory() as tmp_dir: + config = RegressionModelConfig(a=1.5, b=2.5) + gaudi_config = get_gaudi_config() + trainer = GaudiTrainer( + model=RegressionPreTrainedModel(config), + args=GaudiTrainingArguments(output_dir=tmp_dir, use_habana=True, use_lazy_mode=True), + gaudi_config=gaudi_config, + processing_class=feature_extractor, + ) + trainer.save_model() + + reloaded_feature_extractor = AutoFeatureExtractor.from_pretrained(tmp_dir) + + self.assertDictEqual(feature_extractor.to_dict(), reloaded_feature_extractor.to_dict()) + + @require_vision + def test_trainer_saves_processor(self): + MODEL_ID = "openai/clip-vit-base-patch32" + image_processor = AutoImageProcessor.from_pretrained(MODEL_ID) + tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=False) + processor = AutoProcessor.from_pretrained(MODEL_ID) + + with tempfile.TemporaryDirectory() as tmp_dir: + config = RegressionModelConfig(a=1.5, b=2.5) + gaudi_config = get_gaudi_config() + trainer = GaudiTrainer( + model=RegressionPreTrainedModel(config), + args=GaudiTrainingArguments(output_dir=tmp_dir, use_habana=True, use_lazy_mode=True), + gaudi_config=gaudi_config, + processing_class=processor, + ) + trainer.save_model() + + reloaded_processor = AutoProcessor.from_pretrained(tmp_dir) + reloaded_image_processor = AutoImageProcessor.from_pretrained(tmp_dir) + reloaded_tokenizer = AutoTokenizer.from_pretrained(tmp_dir) + + self.assertDictEqual(reloaded_processor.to_dict(), processor.to_dict()) + + image_processor_dict = image_processor.to_dict() + reloaded_image_processor_dict = reloaded_image_processor.to_dict() + # When the processor is saved in the trainer, the _processor_class gets set in the reload_image_processor dict + image_processor_dict.pop("_processor_class") + reloaded_image_processor_dict.pop("_processor_class") + self.assertDictEqual(image_processor_dict, reloaded_image_processor_dict) + + # For tokenizers, there isn't a direct to_dict method and the properties stored in the configs e.g. + # saved tokens change overtime, so we check that two tokenizers are equal by comparing their encoded outputs + test_sentence = "This is a test sentence" + self.assertListEqual( + tokenizer(test_sentence, padding="max_length").input_ids, + reloaded_tokenizer(test_sentence, padding="max_length").input_ids, + ) + + def test_save_best_checkpoint(self): + freq = int(64 / self.batch_size) + total = int(self.n_epochs * 64 / self.batch_size) + + # Case 1: args.metric_for_best_model == "accuracy". + with tempfile.TemporaryDirectory() as tmpdir: + trainer = get_regression_trainer( + a=1.5, + b=2.5, + output_dir=tmpdir, + learning_rate=0.1, + eval_strategy="epoch", + save_strategy="best", + metric_for_best_model="accuracy", + compute_metrics=AlmostAccuracy(), + ) + self.assertTrue(trainer.args.metric_for_best_model == "accuracy") + + with unittest.mock.patch.object( + trainer, + "_evaluate", + side_effect=[ + {"eval_loss": 0.03, "eval_accuracy": 0.60, "epoch": 1.0}, + {"eval_loss": 0.02, "eval_accuracy": 0.65, "epoch": 2.0}, + {"eval_loss": 0.01, "eval_accuracy": 0.64, "epoch": 3.0}, + ], + ): + trainer.train() + + self.assertEqual(len(os.listdir(tmpdir)), 2) + self.check_saved_checkpoints( + output_dir=tmpdir, + freq=freq, + total=total, + ) + + # Case 2: args.metric_for_best_model == "loss". + with tempfile.TemporaryDirectory() as tmpdir: + trainer = get_regression_trainer( + a=1.5, + b=2.5, + output_dir=tmpdir, + learning_rate=0.1, + eval_strategy="epoch", + save_strategy="best", + metric_for_best_model="loss", + compute_metrics=AlmostAccuracy(), + ) + self.assertTrue(trainer.args.metric_for_best_model == "loss") + + with unittest.mock.patch.object( + trainer, + "_evaluate", + side_effect=[ + {"eval_loss": 0.03, "eval_accuracy": 0.60, "epoch": 1.0}, + {"eval_loss": 0.02, "eval_accuracy": 0.65, "epoch": 2.0}, + {"eval_loss": 0.03, "eval_accuracy": 0.66, "epoch": 3.0}, + ], + ): + trainer.train() + + self.assertEqual(len(os.listdir(tmpdir)), 2) + self.check_saved_checkpoints( + output_dir=tmpdir, + freq=freq, + total=total, + ) + + # Case 3: Metric name not provided; throw error. + with tempfile.TemporaryDirectory() as tmpdir: + with self.assertRaises(ValueError) as context: + trainer = get_regression_trainer( + a=1.5, + b=2.5, + output_dir=tmpdir, + learning_rate=0.1, + eval_strategy="epoch", + save_strategy="best", + compute_metrics=AlmostAccuracy(), + ) + + self.assertIn("`args.metric_for_best_model` must be provided", str(context.exception)) + def test_profiling(self): with tempfile.TemporaryDirectory() as tmp_dir: # 24 total steps and compilation takes place during the 1st three steps @@ -2769,64 +2951,49 @@ def setUpClass(cls): cls._token = TOKEN HfFolder.save_token(TOKEN) - @classmethod - def tearDownClass(cls): - for model in [ - "test-trainer", - "test-trainer-epoch", - "test-trainer-step", - "test-trainer-tensorboard", - "test-trainer-tags", - ]: - try: - delete_repo(token=cls._token, repo_id=model) - except HTTPError: - pass - - try: - delete_repo(token=cls._token, repo_id="valid_org/test-trainer-org") - except HTTPError: - pass - def test_push_to_hub(self): - with tempfile.TemporaryDirectory() as tmp_dir: - trainer = get_regression_trainer( - output_dir=os.path.join(tmp_dir, "test-trainer"), - push_to_hub=True, - hub_token=self._token, - ) - url = trainer.push_to_hub() + with TemporaryHubRepo(token=self._token) as tmp_repo: + output_dir_name = tmp_repo.repo_name + with tempfile.TemporaryDirectory() as tmp_dir: + trainer = get_regression_trainer( + output_dir=os.path.join(tmp_dir, output_dir_name), + push_to_hub=True, + hub_token=self._token, + ) + url = trainer.push_to_hub() # Extract repo_name from the url re_search = re.search(ENDPOINT_STAGING + r"/([^/]+/[^/]+)/", url) self.assertTrue(re_search is not None) repo_name = re_search.groups()[0] - self.assertEqual(repo_name, f"{USER}/test-trainer") + self.assertEqual(repo_name, f"{USER}/{output_dir_name}") model = RegressionPreTrainedModel.from_pretrained(repo_name) self.assertEqual(model.a.item(), trainer.model.a.item()) self.assertEqual(model.b.item(), trainer.model.b.item()) def test_push_to_hub_in_organization(self): - with tempfile.TemporaryDirectory() as tmp_dir: - trainer = get_regression_trainer(output_dir=tmp_dir) - trainer.save_model() - trainer = get_regression_trainer( - output_dir=os.path.join(tmp_dir, "test-trainer-org"), - push_to_hub=True, - hub_model_id="valid_org/test-trainer-org", - hub_token=self._token, - ) - url = trainer.push_to_hub() + with TemporaryHubRepo(namespace="valid_org", token=self._token) as tmp_repo: + with tempfile.TemporaryDirectory() as tmp_dir: + trainer = get_regression_trainer(output_dir=tmp_dir) + trainer.save_model() + output_dir_name = tmp_repo.repo_name + trainer = get_regression_trainer( + output_dir=os.path.join(tmp_dir, output_dir_name), + push_to_hub=True, + hub_model_id=f"valid_org/{output_dir_name}", + hub_token=self._token, + ) + url = trainer.push_to_hub() # Extract repo_name from the url re_search = re.search(ENDPOINT_STAGING + r"/([^/]+/[^/]+)/", url) self.assertTrue(re_search is not None) repo_name = re_search.groups()[0] - self.assertEqual(repo_name, "valid_org/test-trainer-org") + self.assertEqual(repo_name, f"valid_org/{output_dir_name}") - model = RegressionPreTrainedModel.from_pretrained("valid_org/test-trainer-org") + model = RegressionPreTrainedModel.from_pretrained(f"valid_org/{output_dir_name}") self.assertEqual(model.a.item(), trainer.model.a.item()) self.assertEqual(model.b.item(), trainer.model.b.item()) @@ -2843,19 +3010,21 @@ def get_commit_history(self, repo): return [commit.strip() for commit in commits] def test_push_to_hub_with_saves_each_epoch(self): - with tempfile.TemporaryDirectory() as tmp_dir: - with self.assertLogs(level="WARNING") as logs: - trainer = get_regression_trainer( - output_dir=os.path.join(tmp_dir, "test-trainer-epoch"), - push_to_hub=True, - hub_token=self._token, - # To avoid any flakiness if the training goes faster than the uploads. - hub_always_push=True, - save_strategy="epoch", - ) - trainer.train() + with TemporaryHubRepo(token=self._token) as tmp_repo: + with tempfile.TemporaryDirectory() as tmp_dir: + with self.assertLogs(level="WARNING") as logs: + output_dir_name = tmp_repo.repo_name + trainer = get_regression_trainer( + output_dir=os.path.join(tmp_dir, output_dir_name), + push_to_hub=True, + hub_token=self._token, + # To avoid any flakiness if the training goes faster than the uploads. + hub_always_push=True, + save_strategy="epoch", + ) + trainer.train() - commits = list_repo_commits(f"{USER}/test-trainer-epoch", token=self._token) + commits = list_repo_commits(f"{USER}/{output_dir_name}", token=self._token) commits = [c.title for c in commits] self.assertIn("initial commit", commits) self.assertIn("Training in progress, epoch 1", commits) @@ -2868,20 +3037,22 @@ def test_push_to_hub_with_saves_each_n_steps(self): if num_gpus > 2: self.skipTest(reason="More than 2 GPUs available") - with tempfile.TemporaryDirectory() as tmp_dir: - with self.assertLogs(level="WARNING") as logs: - trainer = get_regression_trainer( - output_dir=os.path.join(tmp_dir, "test-trainer-step"), - push_to_hub=True, - hub_token=self._token, - # To avoid any flakiness if the training goes faster than the uploads. - hub_always_push=True, - save_strategy="steps", - save_steps=5, - ) - trainer.train() + with TemporaryHubRepo(token=self._token) as tmp_repo: + with tempfile.TemporaryDirectory() as tmp_dir: + with self.assertLogs(level="WARNING") as logs: + output_dir_name = tmp_repo.repo_name + trainer = get_regression_trainer( + output_dir=os.path.join(tmp_dir, output_dir_name), + push_to_hub=True, + hub_token=self._token, + # To avoid any flakiness if the training goes faster than the uploads. + hub_always_push=True, + save_strategy="steps", + save_steps=5, + ) + trainer.train() - commits = list_repo_commits(f"{USER}/test-trainer-step", token=self._token) + commits = list_repo_commits(f"{USER}/{output_dir_name}", token=self._token) commits = [c.title for c in commits] self.assertIn("initial commit", commits) @@ -2901,19 +3072,21 @@ def test_push_to_hub_with_saves_each_n_steps(self): @require_tensorboard def test_push_to_hub_with_tensorboard_logs(self): - with tempfile.TemporaryDirectory() as tmp_dir: - trainer = get_regression_trainer( - output_dir=os.path.join(tmp_dir, "test-trainer-tensorboard"), - hub_token=self._token, - save_strategy="epoch", - report_to=["tensorboard"], - keep_report_to=True, - ) - trainer.train() - # Push the runs via `push_to_hub()` - trainer.push_to_hub() + with TemporaryHubRepo(token=self._token) as tmp_repo: + with tempfile.TemporaryDirectory() as tmp_dir: + output_dir_name = tmp_repo.repo_name + trainer = get_regression_trainer( + output_dir=os.path.join(tmp_dir, output_dir_name), + hub_token=self._token, + save_strategy="epoch", + report_to=["tensorboard"], + keep_report_to=True, + ) + trainer.train() + # Push the runs via `push_to_hub()` + trainer.push_to_hub() - files = list_repo_files(f"{USER}/test-trainer-tensorboard", token=self._token) + files = list_repo_files(f"{USER}/{output_dir_name}", token=self._token) found_log = False for f in files: if len(f.split("runs")) > 1 and "events.out.tfevents" in f: @@ -2925,38 +3098,42 @@ def test_push_to_hub_tags(self): # Checks if `trainer.push_to_hub()` works correctly by adding the desired # tag without having to pass `tags` in `push_to_hub` # see: - with tempfile.TemporaryDirectory() as tmp_dir: - trainer = get_regression_trainer( - output_dir=os.path.join(tmp_dir, "test-trainer-tags"), - push_to_hub=True, - hub_token=self._token, - ) + with TemporaryHubRepo(token=self._token) as tmp_repo: + with tempfile.TemporaryDirectory() as tmp_dir: + output_dir_name = tmp_repo.repo_name + trainer = get_regression_trainer( + output_dir=os.path.join(tmp_dir, output_dir_name), + push_to_hub=True, + hub_token=self._token, + ) - trainer.model.add_model_tags(["test-trainer-tags"]) + trainer.model.add_model_tags(["test-trainer-tags"]) - url = trainer.push_to_hub() + url = trainer.push_to_hub() # Extract repo_name from the url re_search = re.search(ENDPOINT_STAGING + r"/([^/]+/[^/]+)/", url) self.assertTrue(re_search is not None) repo_name = re_search.groups()[0] - self.assertEqual(repo_name, f"{USER}/test-trainer-tags") + self.assertEqual(repo_name, f"{USER}/{output_dir_name}") model_card = ModelCard.load(repo_name) self.assertTrue("test-trainer-tags" in model_card.data.tags) def test_push_to_hub_with_revision(self): # Checks if `trainer.push_to_hub()` works correctly by adding revision - with tempfile.TemporaryDirectory() as tmp_dir: - trainer = get_regression_trainer( - output_dir=os.path.join(tmp_dir, "test-trainer-revision"), - push_to_hub=True, - hub_token=self._token, - ) - branch = "v1.0" - create_branch(repo_id=trainer.hub_model_id, branch=branch, token=self._token, exist_ok=True) - url = trainer.push_to_hub(revision=branch) + with TemporaryHubRepo(token=self._token) as tmp_repo: + with tempfile.TemporaryDirectory() as tmp_dir: + output_dir_name = tmp_repo.repo_name + trainer = get_regression_trainer( + output_dir=os.path.join(tmp_dir, output_dir_name), + push_to_hub=True, + hub_token=self._token, + ) + branch = "v1.0" + create_branch(repo_id=trainer.hub_model_id, branch=branch, token=self._token, exist_ok=True) + url = trainer.push_to_hub(revision=branch) # Extract branch from the url re_search = re.search(r"tree/([^/]+)/", url) diff --git a/tests/test_trainer_seq2seq.py b/tests/test_trainer_seq2seq.py index cb1d5811aa..89905e97e8 100644 --- a/tests/test_trainer_seq2seq.py +++ b/tests/test_trainer_seq2seq.py @@ -118,7 +118,7 @@ def _compute_metrics(pred): compute_metrics=_compute_metrics, train_dataset=train_dataset, eval_dataset=val_dataset, - tokenizer=tokenizer, + processing_class=tokenizer, ) # start training @@ -153,7 +153,7 @@ def test_bad_generation_config_fail_early(self): model=model, gaudi_config=GaudiConfig(), args=training_args, - tokenizer=tokenizer, + processing_class=tokenizer, data_collator=data_collator, compute_metrics=lambda x: {"samples": x[0].shape[0]}, ) From 1924c8942f7103e2d02029bc85bc0bfd58fac499 Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Wed, 8 Jan 2025 16:07:30 +0000 Subject: [PATCH 19/89] Fix Transformers version to install --- .../models/llama/modeling_llama.py | 1 + optimum/habana/transformers/trainer.py | 2 + setup.py | 2 +- tests/test_trainer.py | 54 +++++++++---------- 4 files changed, 31 insertions(+), 28 deletions(-) diff --git a/optimum/habana/transformers/models/llama/modeling_llama.py b/optimum/habana/transformers/models/llama/modeling_llama.py index 16fc68fcc3..92c82503a0 100755 --- a/optimum/habana/transformers/models/llama/modeling_llama.py +++ b/optimum/habana/transformers/models/llama/modeling_llama.py @@ -1080,6 +1080,7 @@ def forward( cache_idx: int = None, lazy_mode: Optional[bool] = True, num_virtual_tokens: int = None, + **kwargs, ) -> Union[Tuple, BaseModelOutputWithPast]: """ Copied from LlamaModel.forward: https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py diff --git a/optimum/habana/transformers/trainer.py b/optimum/habana/transformers/trainer.py index 5e016c79c8..25f380c42b 100644 --- a/optimum/habana/transformers/trainer.py +++ b/optimum/habana/transformers/trainer.py @@ -262,9 +262,11 @@ def __init__( eval_dataset, processing_class, model_init, + compute_loss_func, compute_metrics, callbacks, optimizers, + optimizer_cls_and_kwargs, preprocess_logits_for_metrics, ) diff --git a/setup.py b/setup.py index 0bb36466ee..57d184cce2 100644 --- a/setup.py +++ b/setup.py @@ -29,7 +29,7 @@ INSTALL_REQUIRES = [ - "transformers >= 4.45.2, < 4.46.0", + "transformers >= 4.47.1, < 4.48.0", "optimum", "torch", "accelerate >= 0.33.0, < 0.34.0", diff --git a/tests/test_trainer.py b/tests/test_trainer.py index 61e5daf198..5df6fd7c2b 100644 --- a/tests/test_trainer.py +++ b/tests/test_trainer.py @@ -3024,13 +3024,13 @@ def test_push_to_hub_with_saves_each_epoch(self): ) trainer.train() - commits = list_repo_commits(f"{USER}/{output_dir_name}", token=self._token) - commits = [c.title for c in commits] - self.assertIn("initial commit", commits) - self.assertIn("Training in progress, epoch 1", commits) - self.assertIn("Training in progress, epoch 2", commits) - # Epochs 3 and 4 are not guaranteed to be present (empty commits) - self.assertTrue(any("Skipping to prevent empty commit." in record.message for record in logs.records)) + commits = list_repo_commits(f"{USER}/{output_dir_name}", token=self._token) + commits = [c.title for c in commits] + self.assertIn("initial commit", commits) + self.assertIn("Training in progress, epoch 1", commits) + self.assertIn("Training in progress, epoch 2", commits) + # Epochs 3 and 4 are not guaranteed to be present (empty commits) + self.assertTrue(any("Skipping to prevent empty commit." in record.message for record in logs.records)) def test_push_to_hub_with_saves_each_n_steps(self): num_gpus = max(1, get_gpu_count()) @@ -3052,23 +3052,23 @@ def test_push_to_hub_with_saves_each_n_steps(self): ) trainer.train() - commits = list_repo_commits(f"{USER}/{output_dir_name}", token=self._token) - commits = [c.title for c in commits] - self.assertIn("initial commit", commits) + commits = list_repo_commits(f"{USER}/{output_dir_name}", token=self._token) + commits = [c.title for c in commits] + self.assertIn("initial commit", commits) - # Some commits are skipped if nothing has changed - # We expect 1 commit per 5 epochs + 1 commit at the end - nb_empty_commits = len( - [record for record in logs.records if "Skipping to prevent empty commit." in record.message] - ) - nb_epoch_commits = len([commit for commit in commits if "Training in progress, step" in commit]) + # Some commits are skipped if nothing has changed + # We expect 1 commit per 5 epochs + 1 commit at the end + nb_empty_commits = len( + [record for record in logs.records if "Skipping to prevent empty commit." in record.message] + ) + nb_epoch_commits = len([commit for commit in commits if "Training in progress, step" in commit]) - # max_steps depend on the number of available GPUs - max_steps = math.ceil(trainer.args.num_train_epochs * len(trainer.get_train_dataloader())) - nb_expected_commits = len(range(5, max_steps, 5)) + # max_steps depend on the number of available GPUs + max_steps = math.ceil(trainer.args.num_train_epochs * len(trainer.get_train_dataloader())) + nb_expected_commits = len(range(5, max_steps, 5)) - # '>=' since final commit might be an empty commit as well (not deterministic) - self.assertGreaterEqual(nb_empty_commits + nb_epoch_commits, nb_expected_commits) + # '>=' since final commit might be an empty commit as well (not deterministic) + self.assertGreaterEqual(nb_empty_commits + nb_epoch_commits, nb_expected_commits) @require_tensorboard def test_push_to_hub_with_tensorboard_logs(self): @@ -3086,13 +3086,13 @@ def test_push_to_hub_with_tensorboard_logs(self): # Push the runs via `push_to_hub()` trainer.push_to_hub() - files = list_repo_files(f"{USER}/{output_dir_name}", token=self._token) - found_log = False - for f in files: - if len(f.split("runs")) > 1 and "events.out.tfevents" in f: - found_log = True + files = list_repo_files(f"{USER}/{output_dir_name}", token=self._token) + found_log = False + for f in files: + if len(f.split("runs")) > 1 and "events.out.tfevents" in f: + found_log = True - assert found_log is True, "No tensorboard log found in repo" + assert found_log is True, "No tensorboard log found in repo" def test_push_to_hub_tags(self): # Checks if `trainer.push_to_hub()` works correctly by adding the desired From f0926aef08f28c30c2ce3190314066ef38f1c1a5 Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Thu, 9 Jan 2025 16:42:37 +0000 Subject: [PATCH 20/89] Temporary workaround for GaudiTrainer --- examples/image-to-text/run_pipeline.py | 2 +- examples/language-modeling/run_clm.py | 2 +- .../pytorch-image-models/train_hpu_graph.py | 4 +- .../pytorch-image-models/train_hpu_lazy.py | 4 +- .../run_speech_recognition_ctc.py | 2 +- .../image_to_image_generation.py | 6 +-- .../text_to_image_generation.py | 6 +-- .../training/train_dreambooth_lora_flux.py | 2 +- .../training/train_dreambooth_lora_sdxl.py | 4 +- .../training/train_text_to_image_sdxl.py | 8 ++-- examples/summarization/run_summarization.py | 6 +-- examples/text-classification/run_glue.py | 12 +++--- examples/text-generation/run_generation.py | 18 ++++---- .../text-generation-pipeline/run_pipeline.py | 6 +-- .../run_pipeline_langchain.py | 4 +- examples/text-to-speech/run_pipeline.py | 2 +- .../visual-question-answering/run_pipeline.py | 2 +- optimum/habana/accelerate/accelerator.py | 6 +-- .../pipeline_stable_diffusion_inpaint.py | 2 +- ...eline_stable_diffusion_instruct_pix2pix.py | 2 +- .../pipeline_stable_diffusion_upscale.py | 2 +- .../pipeline_stable_diffusion_xl_inpaint.py | 2 +- optimum/habana/distributed/parallel_state.py | 8 ++-- optimum/habana/distributed/serialization.py | 6 +-- .../habana/transformers/generation/utils.py | 41 +++++++++---------- .../models/baichuan/modeling_baichuan.py | 6 +-- .../transformers/models/bart/modeling_bart.py | 3 +- .../models/chatglm/modeling_chatglm.py | 6 +-- .../transformers/models/clip/modeling_clip.py | 2 +- .../models/falcon/modeling_falcon.py | 4 +- .../models/gemma/modeling_gemma.py | 6 +-- .../models/gemma2/modeling_gemma2.py | 6 +-- .../gpt_bigcode/modeling_gpt_bigcode.py | 12 +++--- .../transformers/models/gptj/modeling_gptj.py | 6 +-- .../models/llama/modeling_llama.py | 6 +-- .../models/modeling_all_models.py | 6 +-- .../transformers/models/opt/modeling_opt.py | 3 +- .../models/qwen2_moe/modeling_qwen2_moe.py | 6 +-- .../seamless_m4t/modeling_seamless_m4t.py | 2 +- .../models/speecht5/modeling_speecht5.py | 3 +- .../transformers/models/t5/modeling_t5.py | 2 +- .../transformers/models/xglm/modeling_xglm.py | 5 +-- optimum/habana/transformers/trainer.py | 31 ++++++++++++-- optimum/habana/trl/trainer/dpo_trainer.py | 3 +- optimum/habana/trl/trainer/sft_trainer.py | 6 +-- tests/test_diffusers.py | 36 ++++++++-------- tests/test_encoder_decoder.py | 2 +- tests/test_text_generation_example.py | 6 +-- tests/test_trainer.py | 8 ++-- .../tests/models/gpt2/test_modeling_gpt2.py | 6 +-- .../models/gpt_neox/test_modeling_gpt_neox.py | 6 +-- .../tests/test_modeling_common.py | 6 +-- 52 files changed, 187 insertions(+), 166 deletions(-) diff --git a/examples/image-to-text/run_pipeline.py b/examples/image-to-text/run_pipeline.py index 44eb8d575a..b218e81daf 100644 --- a/examples/image-to-text/run_pipeline.py +++ b/examples/image-to-text/run_pipeline.py @@ -355,7 +355,7 @@ def preprocess(self, image, prompt=None, timeout=None): throughput = total_new_tokens_generated / duration logger.info(f"result = {result}") logger.info( - f"time = {(end-start) * 1000 / args.n_iterations }ms, Throughput (including tokenization) = {throughput} tokens/second" + f"time = {(end - start) * 1000 / args.n_iterations}ms, Throughput (including tokenization) = {throughput} tokens/second" ) # Store results if necessary diff --git a/examples/language-modeling/run_clm.py b/examples/language-modeling/run_clm.py index 87b6528260..8430792dff 100644 --- a/examples/language-modeling/run_clm.py +++ b/examples/language-modeling/run_clm.py @@ -472,7 +472,7 @@ def main(): else: model = AutoModelForCausalLM.from_config(config, trust_remote_code=model_args.trust_remote_code) n_params = sum({p.data_ptr(): p.numel() for p in model.parameters()}.values()) - logger.info(f"Training new model from scratch - Total size={n_params/2**20:.2f}M params") + logger.info(f"Training new model from scratch - Total size={n_params / 2**20:.2f}M params") # We resize the embeddings only when necessary to avoid index errors. If you are creating a model from scratch # on a small vocab and want a smaller embedding size, remove this test. diff --git a/examples/pytorch-image-models/train_hpu_graph.py b/examples/pytorch-image-models/train_hpu_graph.py index 0bcfbe7295..01e11f8e88 100755 --- a/examples/pytorch-image-models/train_hpu_graph.py +++ b/examples/pytorch-image-models/train_hpu_graph.py @@ -1092,7 +1092,7 @@ def main(): if utils.is_primary(args): _logger.info( - f'Scheduled epochs: {num_epochs}. LR stepped per {"epoch" if lr_scheduler.t_in_epochs else "update"}.' + f"Scheduled epochs: {num_epochs}. LR stepped per {'epoch' if lr_scheduler.t_in_epochs else 'update'}." ) results = [] @@ -1324,7 +1324,7 @@ def _backward(_loss): if utils.is_primary(args): _logger.info( f"Train: {epoch} [{update_idx:>4d}/{updates_per_epoch} " - f"({100. * (update_idx + 1) / updates_per_epoch:>3.0f}%)] " + f"({100.0 * (update_idx + 1) / updates_per_epoch:>3.0f}%)] " f"Loss: {losses_m.val:#.3g} ({losses_m.avg:#.3g}) " f"Time: {update_time_m.val:.3f}s, {update_sample_count / update_time_m.val:>7.2f}/s " f"({update_time_m.avg:.3f}s, {update_sample_count / update_time_m.avg:>7.2f}/s) " diff --git a/examples/pytorch-image-models/train_hpu_lazy.py b/examples/pytorch-image-models/train_hpu_lazy.py index bca523c9b4..f70ae7d7b6 100755 --- a/examples/pytorch-image-models/train_hpu_lazy.py +++ b/examples/pytorch-image-models/train_hpu_lazy.py @@ -1091,7 +1091,7 @@ def main(): if utils.is_primary(args): _logger.info( - f'Scheduled epochs: {num_epochs}. LR stepped per {"epoch" if lr_scheduler.t_in_epochs else "update"}.' + f"Scheduled epochs: {num_epochs}. LR stepped per {'epoch' if lr_scheduler.t_in_epochs else 'update'}." ) results = [] @@ -1325,7 +1325,7 @@ def _backward(_loss): if utils.is_primary(args): _logger.info( f"Train: {epoch} [{update_idx:>4d}/{updates_per_epoch} " - f"({100. * (update_idx + 1) / updates_per_epoch:>3.0f}%)] " + f"({100.0 * (update_idx + 1) / updates_per_epoch:>3.0f}%)] " f"Loss: {losses_m.val:#.3g} ({losses_m.avg:#.3g}) " f"Time: {update_time_m.val:.3f}s, {update_sample_count / update_time_m.val:>7.2f}/s " f"({update_time_m.avg:.3f}s, {update_sample_count / update_time_m.avg:>7.2f}/s) " diff --git a/examples/speech-recognition/run_speech_recognition_ctc.py b/examples/speech-recognition/run_speech_recognition_ctc.py index 3403d00f3c..2b0b6093c3 100644 --- a/examples/speech-recognition/run_speech_recognition_ctc.py +++ b/examples/speech-recognition/run_speech_recognition_ctc.py @@ -504,7 +504,7 @@ def main(): # E.g. characters, such as `,` and `.` do not really have an acoustic characteristic # that could be easily picked up by the model chars_to_ignore_regex = ( - f'[{"".join(data_args.chars_to_ignore).replace(" ", "")}]' 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 ) text_column_name = data_args.text_column_name diff --git a/examples/stable-diffusion/image_to_image_generation.py b/examples/stable-diffusion/image_to_image_generation.py index c76d3c0f5a..acc2536a26 100755 --- a/examples/stable-diffusion/image_to_image_generation.py +++ b/examples/stable-diffusion/image_to_image_generation.py @@ -370,12 +370,12 @@ def main(): logger.info(f"Saving images in {image_save_dir.resolve()}...") if args.ldm3d: for i, rgb in enumerate(outputs.rgb): - rgb.save(image_save_dir / f"rgb_{i+1}.png") + rgb.save(image_save_dir / f"rgb_{i + 1}.png") for i, depth in enumerate(outputs.depth): - depth.save(image_save_dir / f"depth_{i+1}.png") + depth.save(image_save_dir / f"depth_{i + 1}.png") else: for i, image in enumerate(outputs.images): - image.save(image_save_dir / f"image_{i+1}.png") + image.save(image_save_dir / f"image_{i + 1}.png") else: logger.warning("--output_type should be equal to 'pil' to save images in --image_save_dir.") diff --git a/examples/stable-diffusion/text_to_image_generation.py b/examples/stable-diffusion/text_to_image_generation.py index 8fd48c99a8..b4668e7d99 100755 --- a/examples/stable-diffusion/text_to_image_generation.py +++ b/examples/stable-diffusion/text_to_image_generation.py @@ -687,12 +687,12 @@ def main(): logger.info(f"Saving images in {image_save_dir.resolve()}...") if args.ldm3d: for i, rgb in enumerate(outputs.rgb): - rgb.save(image_save_dir / f"rgb_{i+1}.png") + rgb.save(image_save_dir / f"rgb_{i + 1}.png") for i, depth in enumerate(outputs.depth): - depth.save(image_save_dir / f"depth_{i+1}.png") + depth.save(image_save_dir / f"depth_{i + 1}.png") else: for i, image in enumerate(outputs.images): - image.save(image_save_dir / f"image_{i+1}.png") + image.save(image_save_dir / f"image_{i + 1}.png") else: logger.warning("--output_type should be equal to 'pil' to save images in --image_save_dir.") diff --git a/examples/stable-diffusion/training/train_dreambooth_lora_flux.py b/examples/stable-diffusion/training/train_dreambooth_lora_flux.py index 68b5320d19..1117d0a43f 100755 --- a/examples/stable-diffusion/training/train_dreambooth_lora_flux.py +++ b/examples/stable-diffusion/training/train_dreambooth_lora_flux.py @@ -784,7 +784,7 @@ def load_model_hook(models, input_dir): lora_state_dict = FluxPipeline.lora_state_dict(input_dir) transformer_state_dict = { - f'{k.replace("transformer.", "")}': v for k, v in lora_state_dict.items() if k.startswith("transformer.") + f"{k.replace('transformer.', '')}": v for k, v in lora_state_dict.items() if k.startswith("transformer.") } transformer_state_dict = convert_unet_state_dict_to_peft(transformer_state_dict) incompatible_keys = set_peft_model_state_dict(transformer_, transformer_state_dict, adapter_name="default") diff --git a/examples/stable-diffusion/training/train_dreambooth_lora_sdxl.py b/examples/stable-diffusion/training/train_dreambooth_lora_sdxl.py index b177cf12e6..4e96ee8e0d 100755 --- a/examples/stable-diffusion/training/train_dreambooth_lora_sdxl.py +++ b/examples/stable-diffusion/training/train_dreambooth_lora_sdxl.py @@ -94,7 +94,7 @@ def save_model_card( for i, image in enumerate(images): image.save(os.path.join(repo_folder, f"image_{i}.png")) img_str += f""" - - text: '{validation_prompt if validation_prompt else ' ' }' + - text: '{validation_prompt if validation_prompt else " "}' output: url: "image_{i}.png" @@ -1083,7 +1083,7 @@ def load_model_hook(models, input_dir): lora_state_dict, network_alphas = LoraLoaderMixin.lora_state_dict(input_dir) - unet_state_dict = {f'{k.replace("unet.", "")}': v for k, v in lora_state_dict.items() if k.startswith("unet.")} + unet_state_dict = {f"{k.replace('unet.', '')}": v for k, v in lora_state_dict.items() if k.startswith("unet.")} unet_state_dict = convert_unet_state_dict_to_peft(unet_state_dict) incompatible_keys = set_peft_model_state_dict(unet_, unet_state_dict, adapter_name="default") if incompatible_keys is not None: diff --git a/examples/stable-diffusion/training/train_text_to_image_sdxl.py b/examples/stable-diffusion/training/train_text_to_image_sdxl.py index b78c84bbe1..7bb96e51a1 100755 --- a/examples/stable-diffusion/training/train_text_to_image_sdxl.py +++ b/examples/stable-diffusion/training/train_text_to_image_sdxl.py @@ -884,9 +884,9 @@ def main(args): # download the dataset. if args.dataset_name is not None: if len(args.mediapipe) > 0: - assert ( - args.resolution == args.crop_resolution - ), f"To use hardware pipe, --resolution ({args.resolution}) must equal --crop_resolution ({args.crop_resolution})" + assert args.resolution == args.crop_resolution, ( + f"To use hardware pipe, --resolution ({args.resolution}) must equal --crop_resolution ({args.crop_resolution})" + ) if args.local_rank == 0: if not os.path.exists(args.mediapipe): os.mkdir(args.mediapipe) @@ -1532,7 +1532,7 @@ def compute_time_ids(original_size, crops_coords_top_left): image_save_dir.mkdir(parents=True, exist_ok=True) logger.info(f"Saving images in {image_save_dir.resolve()}...") for i, image in enumerate(images): - image.save(image_save_dir / f"image_{epoch}_{i+1}.png") + image.save(image_save_dir / f"image_{epoch}_{i + 1}.png") else: logger.warning("--output_type should be equal to 'pil' to save images in --image_save_dir.") diff --git a/examples/summarization/run_summarization.py b/examples/summarization/run_summarization.py index dc22580f20..97dbe32944 100755 --- a/examples/summarization/run_summarization.py +++ b/examples/summarization/run_summarization.py @@ -559,9 +559,9 @@ def main(): return if isinstance(tokenizer, tuple(MULTILINGUAL_TOKENIZERS)): - assert ( - data_args.lang is not None - ), f"{tokenizer.__class__.__name__} is a multilingual tokenizer which requires --lang argument" + assert data_args.lang is not None, ( + f"{tokenizer.__class__.__name__} is a multilingual tokenizer which requires --lang argument" + ) tokenizer.src_lang = data_args.lang tokenizer.tgt_lang = data_args.lang diff --git a/examples/text-classification/run_glue.py b/examples/text-classification/run_glue.py index 5cfe00ff6e..2e9694b404 100755 --- a/examples/text-classification/run_glue.py +++ b/examples/text-classification/run_glue.py @@ -168,9 +168,9 @@ def __post_init__(self): train_extension = self.train_file.split(".")[-1] assert train_extension in ["csv", "json"], "`train_file` should be a csv or a json file." validation_extension = self.validation_file.split(".")[-1] - assert ( - validation_extension == train_extension - ), "`validation_file` should have the same extension (csv or json) as `train_file`." + assert validation_extension == train_extension, ( + "`validation_file` should have the same extension (csv or json) as `train_file`." + ) @dataclass @@ -338,9 +338,9 @@ def main(): if data_args.test_file is not None: train_extension = data_args.train_file.split(".")[-1] test_extension = data_args.test_file.split(".")[-1] - assert ( - test_extension == train_extension - ), "`test_file` should have the same extension (csv or json) as `train_file`." + assert test_extension == train_extension, ( + "`test_file` should have the same extension (csv or json) as `train_file`." + ) data_files["test"] = data_args.test_file else: raise ValueError("Need either a GLUE task or a test file for `do_predict`.") diff --git a/examples/text-generation/run_generation.py b/examples/text-generation/run_generation.py index ef2252a989..e5df7f2c7c 100755 --- a/examples/text-generation/run_generation.py +++ b/examples/text-generation/run_generation.py @@ -526,7 +526,7 @@ def compute_valid_sequence_lengths_tensor(input_tokens): 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 @@ -541,10 +541,10 @@ def compute_valid_sequence_lengths_tensor(input_tokens): 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: @@ -559,7 +559,7 @@ def rounder(x): 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 @@ -586,12 +586,12 @@ def rounder(x): 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 {i+1}.{j+1}: {output}") + print(f"output {i + 1}.{j + 1}: {output}") all_outputs.append(output) print() @@ -747,10 +747,10 @@ def generate_dataset(batch): 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: diff --git a/examples/text-generation/text-generation-pipeline/run_pipeline.py b/examples/text-generation/text-generation-pipeline/run_pipeline.py index 43aea65cec..11e542d7a5 100644 --- a/examples/text-generation/text-generation-pipeline/run_pipeline.py +++ b/examples/text-generation/text-generation-pipeline/run_pipeline.py @@ -45,14 +45,14 @@ def main(): duration = 0 for iteration in range(args.n_iterations): - logger.info(f"Running inference iteration {iteration+1}...") + logger.info(f"Running inference iteration {iteration + 1}...") t0 = time.perf_counter() output = pipe(input_sentences) duration += time.perf_counter() - t0 for i, (input_sentence, generated_text) in enumerate(zip(input_sentences, output)): - print(f"Prompt[{iteration+1}][{i+1}]: {input_sentence}") - print(f"Generated Text[{iteration+1}][{i+1}]: {repr(generated_text)}\n") + print(f"Prompt[{iteration + 1}][{i + 1}]: {input_sentence}") + print(f"Generated Text[{iteration + 1}][{i + 1}]: {repr(generated_text)}\n") throughput = args.n_iterations * args.batch_size * args.max_new_tokens / duration print(f"Inference Duration (for {args.n_iterations} iterations): {duration} seconds") diff --git a/examples/text-generation/text-generation-pipeline/run_pipeline_langchain.py b/examples/text-generation/text-generation-pipeline/run_pipeline_langchain.py index 556494cd37..6212e808aa 100644 --- a/examples/text-generation/text-generation-pipeline/run_pipeline_langchain.py +++ b/examples/text-generation/text-generation-pipeline/run_pipeline_langchain.py @@ -87,8 +87,8 @@ def main(): duration += time.perf_counter() - t0 for i, (question, answer) in enumerate(zip(input_questions, responses)): - print(f"Question[{iteration+1}][{i+1}]: {question['question']}") - print(f"Response[{iteration+1}][{i+1}]: {answer}\n") + print(f"Question[{iteration + 1}][{i + 1}]: {question['question']}") + print(f"Response[{iteration + 1}][{i + 1}]: {answer}\n") throughput = args.n_iterations * args.batch_size * args.max_new_tokens / duration print(f"Inference Duration (for {args.n_iterations} iterations): {duration} seconds") diff --git a/examples/text-to-speech/run_pipeline.py b/examples/text-to-speech/run_pipeline.py index 1d9b53de7d..81546b0cb9 100644 --- a/examples/text-to-speech/run_pipeline.py +++ b/examples/text-to-speech/run_pipeline.py @@ -129,7 +129,7 @@ def main(): text, batch_size=args.batch_size, forward_params=forward_params, generate_kwargs=generate_kwargs ) end = time.time() - logger.info(f"speech = {speech} time = {(end-start) * 1000 / args.n_iterations }ms") + logger.info(f"speech = {speech} time = {(end - start) * 1000 / args.n_iterations}ms") sf.write("speech.wav", speech[0]["audio"].squeeze(), samplerate=speech[0]["sampling_rate"]) diff --git a/examples/visual-question-answering/run_pipeline.py b/examples/visual-question-answering/run_pipeline.py index 7b4e817bb7..82b05933bc 100644 --- a/examples/visual-question-answering/run_pipeline.py +++ b/examples/visual-question-answering/run_pipeline.py @@ -135,7 +135,7 @@ def main(): with torch.autocast(device_type="hpu", dtype=torch.bfloat16, enabled=autocast_enable): result = generator(model_input, batch_size=args.batch_size, topk=args.topk) end = time.time() - logger.info(f"result = {result}, time = {(end-start) * 1000/args.n_iterations}ms") + logger.info(f"result = {result}, time = {(end - start) * 1000 / args.n_iterations}ms") if __name__ == "__main__": diff --git a/optimum/habana/accelerate/accelerator.py b/optimum/habana/accelerate/accelerator.py index b2d93730a4..f73769692d 100644 --- a/optimum/habana/accelerate/accelerator.py +++ b/optimum/habana/accelerate/accelerator.py @@ -197,9 +197,9 @@ def __init__( if kwargs_handlers is not None: for handler in kwargs_handlers: - assert isinstance( - handler, KwargsHandler - ), f"Unsupported kwargs handler passed: {handler}, must be one that inherits `accelerate.utils.KwargsHandler`." + assert isinstance(handler, KwargsHandler), ( + f"Unsupported kwargs handler passed: {handler}, must be one that inherits `accelerate.utils.KwargsHandler`." + ) if isinstance(handler, DistributedDataParallelKwargs): if self.ddp_handler is not None: raise ValueError("You can only pass one `DistributedDataParallelKwargs` in `kwargs_handler`.") diff --git a/optimum/habana/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py b/optimum/habana/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py index 2884831732..f937423d13 100644 --- a/optimum/habana/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py +++ b/optimum/habana/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py @@ -527,7 +527,7 @@ def __call__( f"Incorrect configuration settings! The config of `pipeline.unet`: {self.unet.config} expects" f" {self.unet.config.in_channels} but received `num_channels_latents`: {num_channels_latents} +" f" `num_channels_mask`: {num_channels_mask} + `num_channels_masked_image`: {num_channels_masked_image}" - f" = {num_channels_latents+num_channels_masked_image+num_channels_mask}. Please verify the config of" + f" = {num_channels_latents + num_channels_masked_image + num_channels_mask}. Please verify the config of" " `pipeline.unet` or your `mask_image` or `image` input." ) elif num_channels_unet != 4: diff --git a/optimum/habana/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_instruct_pix2pix.py b/optimum/habana/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_instruct_pix2pix.py index 0f8eb39f92..c4b0d0e742 100644 --- a/optimum/habana/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_instruct_pix2pix.py +++ b/optimum/habana/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_instruct_pix2pix.py @@ -360,7 +360,7 @@ def __call__( f"Incorrect configuration settings! The config of `pipeline.unet`: {self.unet.config} expects" f" {self.unet.config.in_channels} but received `num_channels_latents`: {num_channels_latents} +" f" `num_channels_image`: {num_channels_image} " - f" = {num_channels_latents+num_channels_image}. Please verify the config of" + f" = {num_channels_latents + num_channels_image}. Please verify the config of" " `pipeline.unet` or your `image` input." ) diff --git a/optimum/habana/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py b/optimum/habana/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py index 58f2f977a9..136ff0dace 100644 --- a/optimum/habana/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py +++ b/optimum/habana/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py @@ -413,7 +413,7 @@ def __call__( f"Incorrect configuration settings! The config of `pipeline.unet`: {self.unet.config} expects" f" {self.unet.config.in_channels} but received `num_channels_latents`: {num_channels_latents} +" f" `num_channels_image`: {num_channels_image} " - f" = {num_channels_latents+num_channels_image}. Please verify the config of" + f" = {num_channels_latents + num_channels_image}. Please verify the config of" " `pipeline.unet` or your `image` input." ) diff --git a/optimum/habana/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py b/optimum/habana/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py index 8d94596e3b..dab18e82e2 100644 --- a/optimum/habana/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py +++ b/optimum/habana/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py @@ -639,7 +639,7 @@ def denoising_value_valid(dnv): f"Incorrect configuration settings! The config of `pipeline.unet`: {self.unet.config} expects" f" {self.unet.config.in_channels} but received `num_channels_latents`: {num_channels_latents} +" f" `num_channels_mask`: {num_channels_mask} + `num_channels_masked_image`: {num_channels_masked_image}" - f" = {num_channels_latents+num_channels_masked_image+num_channels_mask}. Please verify the config of" + f" = {num_channels_latents + num_channels_masked_image + num_channels_mask}. Please verify the config of" " `pipeline.unet` or your `mask_image` or `image` input." ) elif num_channels_unet != 4: diff --git a/optimum/habana/distributed/parallel_state.py b/optimum/habana/distributed/parallel_state.py index c370d88229..3d5c5d9a74 100644 --- a/optimum/habana/distributed/parallel_state.py +++ b/optimum/habana/distributed/parallel_state.py @@ -146,9 +146,9 @@ def initialize_model_parallel( enable_ds_sequence_parallel = sequence_parallel_size > 1 if enable_ds_sequence_parallel: - assert ( - tensor_model_parallel_size == 1 and pipeline_model_parallel_size == 1 - ), "DeepSpeed's sequence parallel does not work with tensor parallel or pipeline parallel" + assert tensor_model_parallel_size == 1 and pipeline_model_parallel_size == 1, ( + "DeepSpeed's sequence parallel does not work with tensor parallel or pipeline parallel" + ) if world_size % sequence_parallel_size != 0: raise RuntimeError( @@ -168,7 +168,7 @@ def initialize_model_parallel( if virtual_pipeline_model_parallel_size is not None: if not pipeline_model_parallel_size > 2: - raise RuntimeError("pipeline-model-parallel size should be greater than 2 with " "interleaved schedule") + raise RuntimeError("pipeline-model-parallel size should be greater than 2 with interleaved schedule") global _VIRTUAL_PIPELINE_MODEL_PARALLEL_RANK global _VIRTUAL_PIPELINE_MODEL_PARALLEL_WORLD_SIZE _VIRTUAL_PIPELINE_MODEL_PARALLEL_RANK = 0 diff --git a/optimum/habana/distributed/serialization.py b/optimum/habana/distributed/serialization.py index bf59fb2445..14842d24ca 100644 --- a/optimum/habana/distributed/serialization.py +++ b/optimum/habana/distributed/serialization.py @@ -191,9 +191,9 @@ def load_state_dict( assert len(checkpoints) > 0, f"Can't find the requested checkpoint data at {model_path}" if checkpoint_sharding is not None and checkpoint_sharding != "layer": - assert ( - world_size == len(checkpoints) - ), f"Loading a {checkpoint_sharding}-sharded checkpoint with len={len(checkpoints)} but world size is {world_size}" + assert world_size == len(checkpoints), ( + f"Loading a {checkpoint_sharding}-sharded checkpoint with len={len(checkpoints)} but world size is {world_size}" + ) checkpoints = [checkpoints[rank]] diff --git a/optimum/habana/transformers/generation/utils.py b/optimum/habana/transformers/generation/utils.py index defa93c6c0..cdd7ce8c19 100644 --- a/optimum/habana/transformers/generation/utils.py +++ b/optimum/habana/transformers/generation/utils.py @@ -1107,28 +1107,27 @@ def generate( assert generation_config.bucket_size >= 0, "please set bucket_size to use bucket_internal" assert generation_config.use_cache, "please set use_cache flag to use bucket_internal" if generation_config.reuse_cache: - assert ( - self.config.model_type - in [ - "llama", - "mistral", - "falcon", - "mixtral", - "phi", - "qwen2", - "gptj", - "starcoder2", - "qwen2_moe", - "gemma", - "gemma2", - "baichuan", - "chatglm", - ] - ), "reuse_cache only supported by llama, mistral, falcon, mixtral, phi, qwen2, qwen2_moe, gemma, gemma2, starcoder2, baichuan and chatglm at the moment" + assert self.config.model_type in [ + "llama", + "mistral", + "falcon", + "mixtral", + "phi", + "qwen2", + "gptj", + "starcoder2", + "qwen2_moe", + "gemma", + "gemma2", + "baichuan", + "chatglm", + ], ( + "reuse_cache only supported by llama, mistral, falcon, mixtral, phi, qwen2, qwen2_moe, gemma, gemma2, starcoder2, baichuan and chatglm at the moment" + ) if not generation_config.bucket_internal: - assert ( - generation_config.bucket_size <= 0 - ), "please set bucket_internal along with reuse_cache and bucket_size" + assert generation_config.bucket_size <= 0, ( + "please set bucket_internal along with reuse_cache and bucket_size" + ) else: assert generation_config.bucket_size >= 0, "please set valid bucket_size to use bucket_internal" diff --git a/optimum/habana/transformers/models/baichuan/modeling_baichuan.py b/optimum/habana/transformers/models/baichuan/modeling_baichuan.py index b733712fbb..ca9498e0f1 100644 --- a/optimum/habana/transformers/models/baichuan/modeling_baichuan.py +++ b/optimum/habana/transformers/models/baichuan/modeling_baichuan.py @@ -133,9 +133,9 @@ def allocate(self, inp_seq_len, dtype, device, shape): self.inp_seq_len = inp_seq_len self.cache = torch.zeros(shape, dtype=dtype, device=device) else: - assert ( - self.inp_seq_len == inp_seq_len - ), f"inp_seq_len must be the same. self.inp_seq_len:{self.inp_seq_len} inp_seq_len:{inp_seq_len}" + assert self.inp_seq_len == inp_seq_len, ( + f"inp_seq_len must be the same. self.inp_seq_len:{self.inp_seq_len} inp_seq_len:{inp_seq_len}" + ) self.cache.fill_(0) def update(self, prev, cur, dim, idx, inp_seq_len): diff --git a/optimum/habana/transformers/models/bart/modeling_bart.py b/optimum/habana/transformers/models/bart/modeling_bart.py index 3e5f822cb1..2fdfbcc6d0 100644 --- a/optimum/habana/transformers/models/bart/modeling_bart.py +++ b/optimum/habana/transformers/models/bart/modeling_bart.py @@ -158,8 +158,7 @@ def gaudi_BartAttention_forward( if layer_head_mask is not None: if layer_head_mask.size() != (self.num_heads,): raise ValueError( - f"Head mask for a single layer should be of size {(self.num_heads,)}, but is" - f" {layer_head_mask.size()}" + 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) diff --git a/optimum/habana/transformers/models/chatglm/modeling_chatglm.py b/optimum/habana/transformers/models/chatglm/modeling_chatglm.py index 01c508aa5d..3afa86c4a9 100644 --- a/optimum/habana/transformers/models/chatglm/modeling_chatglm.py +++ b/optimum/habana/transformers/models/chatglm/modeling_chatglm.py @@ -148,9 +148,9 @@ def allocate(self, inp_seq_len, dtype, device, shape): # self.cache = torch.zeros(shape, dtype=dtype, device=device) self.cache = torch.zeros(shape, dtype=torch.bfloat16, device=device) else: - assert ( - self.inp_seq_len == inp_seq_len - ), f"inp_seq_len must be the same. self.inp_seq_len:{self.inp_seq_len} inp_seq_len:{inp_seq_len}" + assert self.inp_seq_len == inp_seq_len, ( + f"inp_seq_len must be the same. self.inp_seq_len:{self.inp_seq_len} inp_seq_len:{inp_seq_len}" + ) self.cache.fill_(0) def update(self, prev, cur, dim, idx, inp_seq_len): diff --git a/optimum/habana/transformers/models/clip/modeling_clip.py b/optimum/habana/transformers/models/clip/modeling_clip.py index b48ba858ca..310bdef1fa 100644 --- a/optimum/habana/transformers/models/clip/modeling_clip.py +++ b/optimum/habana/transformers/models/clip/modeling_clip.py @@ -29,7 +29,7 @@ def forward(self, pixel_values: torch.FloatTensor, interpolate_pos_encoding=Fals batch_size, _, height, width = pixel_values.shape if not interpolate_pos_encoding and (height != self.image_size or width != self.image_size): raise ValueError( - f"Input image size ({height}*{width}) doesn't match model" f" ({self.image_size}*{self.image_size})." + f"Input image size ({height}*{width}) doesn't match model ({self.image_size}*{self.image_size})." ) target_dtype = self.patch_embedding.weight.dtype # if HQT quantization enabled, remove the explicit cast to float8 to avoid HQT casting error diff --git a/optimum/habana/transformers/models/falcon/modeling_falcon.py b/optimum/habana/transformers/models/falcon/modeling_falcon.py index 3ef9edbdbb..92e42deb33 100644 --- a/optimum/habana/transformers/models/falcon/modeling_falcon.py +++ b/optimum/habana/transformers/models/falcon/modeling_falcon.py @@ -1054,7 +1054,9 @@ def forward( """ return_dict = return_dict if return_dict is not None else self.config.use_return_dict if use_flash_attention: - assert FusedSDPA, "`use_flash_attention` is True, but cannot find FusedSDPA. Please import it as `from habana_frameworks.torch.hpex.kernels import FusedSDPA` or set use_flash_attention to False (at the expense of a possible performance degradation)." + assert FusedSDPA, ( + "`use_flash_attention` is True, but cannot find FusedSDPA. Please import it as `from habana_frameworks.torch.hpex.kernels import FusedSDPA` or set use_flash_attention to False (at the expense of a possible performance degradation)." + ) if flash_attention_recompute: assert use_flash_attention, "flash_attention_recompute is set, but use_flash_attention is not" if flash_attention_causal_mask: diff --git a/optimum/habana/transformers/models/gemma/modeling_gemma.py b/optimum/habana/transformers/models/gemma/modeling_gemma.py index 30b01c8aad..8e34b12b7f 100755 --- a/optimum/habana/transformers/models/gemma/modeling_gemma.py +++ b/optimum/habana/transformers/models/gemma/modeling_gemma.py @@ -131,9 +131,9 @@ def allocate(self, inp_seq_len, dtype, device, shape): self.inp_seq_len = inp_seq_len self.cache = torch.zeros(shape, dtype=dtype, device=device) else: - assert ( - self.inp_seq_len == inp_seq_len - ), f"inp_seq_len must be the same. self.inp_seq_len:{self.inp_seq_len} inp_seq_len:{inp_seq_len}" + assert self.inp_seq_len == inp_seq_len, ( + f"inp_seq_len must be the same. self.inp_seq_len:{self.inp_seq_len} inp_seq_len:{inp_seq_len}" + ) self.cache.fill_(0) def update(self, prev, cur, dim, idx, inp_seq_len): diff --git a/optimum/habana/transformers/models/gemma2/modeling_gemma2.py b/optimum/habana/transformers/models/gemma2/modeling_gemma2.py index fff49d4649..5927b04285 100755 --- a/optimum/habana/transformers/models/gemma2/modeling_gemma2.py +++ b/optimum/habana/transformers/models/gemma2/modeling_gemma2.py @@ -213,9 +213,9 @@ def allocate(self, inp_seq_len, dtype, device, shape): self.inp_seq_len = inp_seq_len self.cache = torch.zeros(shape, dtype=dtype, device=device) else: - assert ( - self.inp_seq_len == inp_seq_len - ), f"inp_seq_len must be the same. self.inp_seq_len:{self.inp_seq_len} inp_seq_len:{inp_seq_len}" + assert self.inp_seq_len == inp_seq_len, ( + f"inp_seq_len must be the same. self.inp_seq_len:{self.inp_seq_len} inp_seq_len:{inp_seq_len}" + ) self.cache.fill_(0) def update(self, prev, cur, dim, idx, inp_seq_len): 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 9f451256c9..f01255624f 100644 --- a/optimum/habana/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py +++ b/optimum/habana/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py @@ -306,9 +306,9 @@ def forward( - optimize KV cache """ if use_flash_attention: - assert ( - self.fused_scaled_dot_product_attention is not None - ), "Can't load HPU fused scaled dot-product attention kernel. Please retry without flash attention" + assert self.fused_scaled_dot_product_attention is not None, ( + "Can't load HPU fused scaled dot-product attention kernel. Please retry without flash attention" + ) if encoder_hidden_states is not None: if not hasattr(self, "q_attn") or not self.is_cross_attention: @@ -353,9 +353,9 @@ def forward( present = torch.cat((key, value), dim=-1) if use_cache else None else: assert token_idx is not None, "Invalid parameters: token_idx is None at decode stage with bucket_internal" - assert ( - layer_past is not None - ), "Invalid parameters: layer_past is None at decode stage with bucket_internal" + assert layer_past is not None, ( + "Invalid parameters: layer_past is None at decode stage with bucket_internal" + ) past_key, past_value = layer_past.split((self.head_dim, self.head_dim), dim=-1) key = past_key.index_copy_(1, token_idx - 1, key) diff --git a/optimum/habana/transformers/models/gptj/modeling_gptj.py b/optimum/habana/transformers/models/gptj/modeling_gptj.py index c61f496cb3..d4da76d6f2 100644 --- a/optimum/habana/transformers/models/gptj/modeling_gptj.py +++ b/optimum/habana/transformers/models/gptj/modeling_gptj.py @@ -38,9 +38,9 @@ def allocate(self, inp_seq_len, dtype, device, shape): self.inp_seq_len = inp_seq_len self.cache = torch.zeros(shape, dtype=dtype, device=device) else: - assert ( - self.inp_seq_len == inp_seq_len - ), f"inp_seq_len must be the same. self.inp_seq_len:{self.inp_seq_len} inp_seq_len:{inp_seq_len}" + assert self.inp_seq_len == inp_seq_len, ( + f"inp_seq_len must be the same. self.inp_seq_len:{self.inp_seq_len} inp_seq_len:{inp_seq_len}" + ) self.cache.fill_(0) def update(self, prev, cur, dim, idx, inp_seq_len): diff --git a/optimum/habana/transformers/models/llama/modeling_llama.py b/optimum/habana/transformers/models/llama/modeling_llama.py index 92c82503a0..da26c16567 100755 --- a/optimum/habana/transformers/models/llama/modeling_llama.py +++ b/optimum/habana/transformers/models/llama/modeling_llama.py @@ -378,9 +378,9 @@ def allocate(self, inp_seq_len, dtype, device, shape): self.inp_seq_len = inp_seq_len self.cache = torch.zeros(shape, dtype=dtype, device=device) else: - assert ( - self.inp_seq_len == inp_seq_len - ), f"inp_seq_len must be the same. self.inp_seq_len:{self.inp_seq_len} inp_seq_len:{inp_seq_len}" + assert self.inp_seq_len == inp_seq_len, ( + f"inp_seq_len must be the same. self.inp_seq_len:{self.inp_seq_len} inp_seq_len:{inp_seq_len}" + ) self.cache.fill_(0) @staticmethod diff --git a/optimum/habana/transformers/models/modeling_all_models.py b/optimum/habana/transformers/models/modeling_all_models.py index 5a78359e3a..3f9304db74 100644 --- a/optimum/habana/transformers/models/modeling_all_models.py +++ b/optimum/habana/transformers/models/modeling_all_models.py @@ -48,9 +48,9 @@ def allocate(self, inp_seq_len, dtype, device, shape): self.inp_seq_len = inp_seq_len self.cache = torch.zeros(shape, dtype=dtype, device=device) else: - assert ( - self.inp_seq_len == inp_seq_len - ), f"inp_seq_len must be the same. self.inp_seq_len:{self.inp_seq_len} inp_seq_len:{inp_seq_len}" + assert self.inp_seq_len == inp_seq_len, ( + f"inp_seq_len must be the same. self.inp_seq_len:{self.inp_seq_len} inp_seq_len:{inp_seq_len}" + ) self.cache.fill_(0) @staticmethod diff --git a/optimum/habana/transformers/models/opt/modeling_opt.py b/optimum/habana/transformers/models/opt/modeling_opt.py index f30a1e4435..179495d776 100644 --- a/optimum/habana/transformers/models/opt/modeling_opt.py +++ b/optimum/habana/transformers/models/opt/modeling_opt.py @@ -127,8 +127,7 @@ def gaudi_opt_attention_forward( if layer_head_mask is not None: if layer_head_mask.size() != (self.num_heads,): raise ValueError( - f"Head mask for a single layer should be of size {(self.num_heads,)}, but is" - f" {layer_head_mask.size()}" + 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) 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 efddd47dc5..0dc677d9bd 100755 --- a/optimum/habana/transformers/models/qwen2_moe/modeling_qwen2_moe.py +++ b/optimum/habana/transformers/models/qwen2_moe/modeling_qwen2_moe.py @@ -188,9 +188,9 @@ def allocate(self, inp_seq_len, dtype, device, shape): self.inp_seq_len = inp_seq_len self.cache = torch.zeros(shape, dtype=dtype, device=device) else: - assert ( - self.inp_seq_len == inp_seq_len - ), f"inp_seq_len must be the same. self.inp_seq_len:{self.inp_seq_len} inp_seq_len:{inp_seq_len}" + assert self.inp_seq_len == inp_seq_len, ( + f"inp_seq_len must be the same. self.inp_seq_len:{self.inp_seq_len} inp_seq_len:{inp_seq_len}" + ) self.cache.fill_(0) @staticmethod 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 53cea37255..061aebb3c6 100644 --- a/optimum/habana/transformers/models/seamless_m4t/modeling_seamless_m4t.py +++ b/optimum/habana/transformers/models/seamless_m4t/modeling_seamless_m4t.py @@ -732,7 +732,7 @@ def gaudi_SeamlessM4TForTextToSpeech_generate( elif tgt_lang not in lang_code_to_id: raise ValueError( f"""`tgt_lang={tgt_lang}` is not supported by this model. - Please specify a `tgt_lang` in {','.join(lang_code_to_id.keys())}. Note that SeamlessM4T supports + Please specify a `tgt_lang` in {",".join(lang_code_to_id.keys())}. Note that SeamlessM4T supports more languages for text translation than for speech synthesis.""" ) if kwargs.get("hpu_graphs", True): diff --git a/optimum/habana/transformers/models/speecht5/modeling_speecht5.py b/optimum/habana/transformers/models/speecht5/modeling_speecht5.py index ac0fb472ae..25f47176ed 100644 --- a/optimum/habana/transformers/models/speecht5/modeling_speecht5.py +++ b/optimum/habana/transformers/models/speecht5/modeling_speecht5.py @@ -115,8 +115,7 @@ def gaudi_SpeechT5Attention_forward( if layer_head_mask is not None: if layer_head_mask.size() != (self.num_heads,): raise ValueError( - f"Head mask for a single layer should be of size {(self.num_heads,)}, but is" - f" {layer_head_mask.size()}" + 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) diff --git a/optimum/habana/transformers/models/t5/modeling_t5.py b/optimum/habana/transformers/models/t5/modeling_t5.py index bdba215617..15e7a4e92b 100644 --- a/optimum/habana/transformers/models/t5/modeling_t5.py +++ b/optimum/habana/transformers/models/t5/modeling_t5.py @@ -70,7 +70,7 @@ def gaudi_T5Attention_forward( if past_key_value is not None: if len(past_key_value) != 2: raise ValueError( - f"past_key_value should have 2 past states: keys and values. Got { len(past_key_value)} past states" + f"past_key_value should have 2 past states: keys and values. Got {len(past_key_value)} past states" ) if token_idx is None: real_seq_length += past_key_value[0].shape[2] if query_length is None else query_length diff --git a/optimum/habana/transformers/models/xglm/modeling_xglm.py b/optimum/habana/transformers/models/xglm/modeling_xglm.py index ef5a16801a..f69eb3b990 100644 --- a/optimum/habana/transformers/models/xglm/modeling_xglm.py +++ b/optimum/habana/transformers/models/xglm/modeling_xglm.py @@ -109,8 +109,7 @@ def gaudi_xglm_attention_forward( if layer_head_mask is not None: if layer_head_mask.size() != (self.num_heads,): raise ValueError( - f"Head mask for a single layer should be of size {(self.num_heads,)}, but is" - f" {layer_head_mask.size()}" + 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) @@ -300,7 +299,7 @@ def gaudi_xglm_model_forward( 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 = True` is incompatible with gradient checkpointing`. Setting `use_cache = False`..." ) use_cache = False diff --git a/optimum/habana/transformers/trainer.py b/optimum/habana/transformers/trainer.py index 25f380c42b..44690f4b6a 100644 --- a/optimum/habana/transformers/trainer.py +++ b/optimum/habana/transformers/trainer.py @@ -1633,9 +1633,9 @@ def training_step( loss = loss / self.args.gradient_accumulation_steps if _is_peft_model(self.model) and self.model.peft_type == PeftType.ADALORA: - assert not ( - self.accelerator.state.is_fp8_enabled and self.args.gradient_checkpointing - ), "FP8 precision with gradient_checkpointing is currently not supported with PeftType.ADALORA" + assert not (self.accelerator.state.is_fp8_enabled and self.args.gradient_checkpointing), ( + "FP8 precision with gradient_checkpointing is currently not supported with PeftType.ADALORA" + ) if self.is_deepspeed_enabled and not is_deepspeed_zero3_enabled(): self.accelerator.deepspeed_engine_wrapped.engine.backward(loss) self.model.base_model.update_and_allocate(self.state.global_step) @@ -2596,3 +2596,28 @@ def _zero_model_grad(self, model): except TypeError: model.zero_grad() model._zero_grad_kwargs = {} + + def get_batch_samples(self, epoch_iterator, num_batches): + batch_samples = [] + num_items_in_batch = None + for _ in range(num_batches): + try: + batch_samples += [next(epoch_iterator)] + except StopIteration: + break + + # TODO: execute get_batch_samples outside of the training loop (before training) and uncomment the following lines + # if len(batch_samples) > 0 and "labels" in batch_samples[0]: + # # For now we don't support object detection + # try: + # num_items_in_batch = sum([(batch["labels"].ne(-100)).sum() for batch in batch_samples]) + # except (TypeError, AttributeError): + # pass + + # if self.args.average_tokens_across_devices: + # num_items_in_batch = self.accelerator.gather(num_items_in_batch).sum().item() + + # if torch.is_tensor(num_items_in_batch): + # num_items_in_batch = num_items_in_batch.item() + + return batch_samples, num_items_in_batch diff --git a/optimum/habana/trl/trainer/dpo_trainer.py b/optimum/habana/trl/trainer/dpo_trainer.py index bd07a981bb..84c48f1782 100644 --- a/optimum/habana/trl/trainer/dpo_trainer.py +++ b/optimum/habana/trl/trainer/dpo_trainer.py @@ -167,8 +167,7 @@ def __init__( if isinstance(ref_model, str): warnings.warn( - "You passed a ref model_id to the DPOTrainer. This will automatically create an " - "`AutoModelForCausalLM`" + "You passed a ref model_id to the DPOTrainer. This will automatically create an `AutoModelForCausalLM`" ) ref_model = AutoModelForCausalLM.from_pretrained(ref_model, **ref_model_init_kwargs) diff --git a/optimum/habana/trl/trainer/sft_trainer.py b/optimum/habana/trl/trainer/sft_trainer.py index 04e648a161..6fb6365655 100644 --- a/optimum/habana/trl/trainer/sft_trainer.py +++ b/optimum/habana/trl/trainer/sft_trainer.py @@ -133,9 +133,9 @@ def __init__( - num_buckets: Number of buckets. > 0 means apply bucketing, <= 0 means no bucketing """ if num_buckets > 0: - assert ( - data_collator is None - ), "For bucketing (num_buckets > 0), we only support data_collator=None (later it becomes DataCollatorForLanguageModeling)" + assert data_collator is None, ( + "For bucketing (num_buckets > 0), we only support data_collator=None (later it becomes DataCollatorForLanguageModeling)" + ) if args is None: output_dir = "tmp_trainer" warnings.warn(f"No `SFTConfig` passed, using `output_dir={output_dir}`.") diff --git a/tests/test_diffusers.py b/tests/test_diffusers.py index 03663b7fc8..b26878551a 100755 --- a/tests/test_diffusers.py +++ b/tests/test_diffusers.py @@ -1616,15 +1616,15 @@ def test_fused_qkv_projections(self): image = pipe(**inputs).images image_slice_disabled = image[0, -3:, -3:, -1] - assert np.allclose( - original_image_slice, image_slice_fused, atol=1e-3, rtol=1e-3 - ), "Fusion of QKV projections shouldn't affect the outputs." - assert np.allclose( - image_slice_fused, image_slice_disabled, atol=1e-3, rtol=1e-3 - ), "Outputs, with QKV projection fusion enabled, shouldn't change when fused QKV projections are disabled." - assert np.allclose( - original_image_slice, image_slice_disabled, atol=1e-2, rtol=1e-2 - ), "Original outputs should match when fused QKV projections are disabled." + assert np.allclose(original_image_slice, image_slice_fused, atol=1e-3, rtol=1e-3), ( + "Fusion of QKV projections shouldn't affect the outputs." + ) + assert np.allclose(image_slice_fused, image_slice_disabled, atol=1e-3, rtol=1e-3), ( + "Outputs, with QKV projection fusion enabled, shouldn't change when fused QKV projections are disabled." + ) + assert np.allclose(original_image_slice, image_slice_disabled, atol=1e-2, rtol=1e-2), ( + "Original outputs should match when fused QKV projections are disabled." + ) class GaudiStableDiffusionControlNetPipelineTester(TestCase): @@ -2536,7 +2536,7 @@ def test_train_controlnet(self): cmd_line = f""" python3 - {path_to_script.parent.parent.parent / 'gaudi_spawn.py'} + {path_to_script.parent.parent.parent / "gaudi_spawn.py"} --use_mpi --world_size 8 {path_to_script} @@ -2624,7 +2624,7 @@ def _test_dreambooth(self, extra_config, train_text_encoder=False): python3 {path_to_script} --pretrained_model_name_or_path hf-internal-testing/tiny-stable-diffusion-pipe - --instance_data_dir {Path(os.path.dirname(__file__))/'resource/img'} + --instance_data_dir {Path(os.path.dirname(__file__)) / "resource/img"} --resolution 64 --train_batch_size 1 --gradient_accumulation_steps 1 @@ -2720,7 +2720,7 @@ def _test_dreambooth_lora_sdxl(self, train_text_encoder=False): python3 {path_to_script} --pretrained_model_name_or_path hf-internal-testing/tiny-stable-diffusion-xl-pipe - --instance_data_dir {Path(os.path.dirname(__file__))/'resource/img'} + --instance_data_dir {Path(os.path.dirname(__file__)) / "resource/img"} --resolution 64 --train_batch_size 1 --gradient_accumulation_steps 1 @@ -5939,9 +5939,9 @@ def new_step(self, *args, **kwargs): inputs_1 = {**inputs, **{"denoising_end": split_1, "output_type": "latent"}} latents = pipe_1(**inputs_1).images[0] - assert ( - expected_steps_1 == done_steps - ), f"Failure with {scheduler_cls.__name__} and {num_steps} and {split_1} and {split_2}" + assert expected_steps_1 == done_steps, ( + f"Failure with {scheduler_cls.__name__} and {num_steps} and {split_1} and {split_2}" + ) inputs_2 = { **inputs, @@ -5955,9 +5955,9 @@ def new_step(self, *args, **kwargs): pipe_3(**inputs_3).images[0] assert expected_steps_3 == done_steps[len(expected_steps_1) + len(expected_steps_2) :] - assert ( - expected_steps == done_steps - ), f"Failure with {scheduler_cls.__name__} and {num_steps} and {split_1} and {split_2}" + assert expected_steps == done_steps, ( + f"Failure with {scheduler_cls.__name__} and {num_steps} and {split_1} and {split_2}" + ) for steps in [7, 11, 20]: for split_1, split_2 in zip([0.19, 0.32], [0.81, 0.68]): diff --git a/tests/test_encoder_decoder.py b/tests/test_encoder_decoder.py index 20d808b69f..723739eb5b 100644 --- a/tests/test_encoder_decoder.py +++ b/tests/test_encoder_decoder.py @@ -189,7 +189,7 @@ def _test_text_translation( "--do_predict", "--source_lang en", "--target_lang ro", - '--source_prefix "translate English to Romanian: "' "--dataset_name wmt16", + '--source_prefix "translate English to Romanian: "--dataset_name wmt16', "--dataset_config_name ro-en", f"--per_device_eval_batch_size {batch_size}", f"--generation_num_beams {num_beams}", diff --git a/tests/test_text_generation_example.py b/tests/test_text_generation_example.py index ec1cc67475..912cbefae8 100644 --- a/tests/test_text_generation_example.py +++ b/tests/test_text_generation_example.py @@ -369,9 +369,9 @@ def _test_text_generation( # Verify output for 1 HPU, BF16 if check_output: - assert ( - model_name in MODEL_OUTPUTS - ), f"Failed functional testing, missing expected output in MODEL_OUTPUTS for model {model_name}" + assert model_name in MODEL_OUTPUTS, ( + f"Failed functional testing, missing expected output in MODEL_OUTPUTS for model {model_name}" + ) expected_output = MODEL_OUTPUTS[model_name] assert results["output"][0][0] == expected_output diff --git a/tests/test_trainer.py b/tests/test_trainer.py index 5df6fd7c2b..92118a5b55 100644 --- a/tests/test_trainer.py +++ b/tests/test_trainer.py @@ -544,7 +544,7 @@ def convert_to_sharded_checkpoint(self, folder, save_safe=True, load_safe=True): keys = list(state_dict.keys()) shard_files = [ - shard_name.replace(f".{extension}", f"-{idx+1:05d}-of-{len(keys):05d}.{extension}") + shard_name.replace(f".{extension}", f"-{idx + 1:05d}-of-{len(keys):05d}.{extension}") for idx in range(len(keys)) ] index = {"metadata": {}, "weight_map": {key: shard_files[i] for i, key in enumerate(keys)}} @@ -1706,9 +1706,9 @@ def test_load_best_model_with_save(self): ) trainer.train() # Check that we have the last known step: - assert os.path.exists( - os.path.join(tmpdir, f"checkpoint-{trainer.state.max_steps}") - ), f"Could not find checkpoint-{trainer.state.max_steps}" + assert os.path.exists(os.path.join(tmpdir, f"checkpoint-{trainer.state.max_steps}")), ( + f"Could not find checkpoint-{trainer.state.max_steps}" + ) # And then check the last step assert os.path.exists(os.path.join(tmpdir, "checkpoint-9")), "Could not find checkpoint-9" diff --git a/tests/transformers/tests/models/gpt2/test_modeling_gpt2.py b/tests/transformers/tests/models/gpt2/test_modeling_gpt2.py index eae4e5571a..b479f2b237 100644 --- a/tests/transformers/tests/models/gpt2/test_modeling_gpt2.py +++ b/tests/transformers/tests/models/gpt2/test_modeling_gpt2.py @@ -392,9 +392,9 @@ def create_and_check_cached_forward_with_and_without_attention_mask(self, config model.eval() # We want this for SDPA, eager works with a `None` attention mask - assert ( - model.config._attn_implementation == "sdpa" - ), "This test assumes the model to have the SDPA implementation for its attention calculations." + assert model.config._attn_implementation == "sdpa", ( + "This test assumes the model to have the SDPA implementation for its attention calculations." + ) # Prepare cache and non_cache input, needs a full attention mask cached_len = input_ids.shape[-1] // 2 diff --git a/tests/transformers/tests/models/gpt_neox/test_modeling_gpt_neox.py b/tests/transformers/tests/models/gpt_neox/test_modeling_gpt_neox.py index 14561c2080..5026ff87d8 100644 --- a/tests/transformers/tests/models/gpt_neox/test_modeling_gpt_neox.py +++ b/tests/transformers/tests/models/gpt_neox/test_modeling_gpt_neox.py @@ -213,9 +213,9 @@ def create_and_check_cached_forward_with_and_without_attention_mask(self, config model.to(torch_device) model.eval() # We want this for SDPA, eager works with a `None` attention mask - assert ( - model.config._attn_implementation == "sdpa" - ), "This test assumes the model to have the SDPA implementation for its attention calculations." + assert model.config._attn_implementation == "sdpa", ( + "This test assumes the model to have the SDPA implementation for its attention calculations." + ) # Prepare cache and non_cache input, needs a full attention mask cached_len = input_ids.shape[-1] // 2 input_mask = torch.ones(size=input_ids.size()).to(torch_device) diff --git a/tests/transformers/tests/test_modeling_common.py b/tests/transformers/tests/test_modeling_common.py index e08860278b..55c7aa8dae 100755 --- a/tests/transformers/tests/test_modeling_common.py +++ b/tests/transformers/tests/test_modeling_common.py @@ -2261,9 +2261,9 @@ def test_model_is_small(self): for model_class in self.all_model_classes: model = model_class(config) num_params = model.num_parameters() - assert ( - num_params < 1000000 - ), f"{model_class} is too big for the common tests ({num_params})! It should have 1M max." + assert num_params < 1000000, ( + f"{model_class} is too big for the common tests ({num_params})! It should have 1M max." + ) global_rng = random.Random() From e50e1792e327386f897a88d21369a4e48623a346 Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Fri, 10 Jan 2025 09:13:08 +0000 Subject: [PATCH 21/89] Fixes for text generation --- .../habana/transformers/generation/utils.py | 39 ++--- optimum/habana/transformers/modeling_utils.py | 4 +- .../habana/transformers/models/__init__.py | 2 +- .../models/gemma/modeling_gemma.py | 2 +- .../models/gpt_neox/modeling_gpt_neox.py | 64 +++++++- .../transformers/models/opt/__init__.py | 2 +- .../transformers/models/opt/modeling_opt.py | 146 +++++++++++------- 7 files changed, 176 insertions(+), 83 deletions(-) diff --git a/optimum/habana/transformers/generation/utils.py b/optimum/habana/transformers/generation/utils.py index cdd7ce8c19..3486463480 100644 --- a/optimum/habana/transformers/generation/utils.py +++ b/optimum/habana/transformers/generation/utils.py @@ -2491,12 +2491,7 @@ def _sample( **hpu_graphs_kwargs, ) - # synced_gpus: don't waste resources running the code we don't need; kwargs must be updated before skipping - model_kwargs = self._update_model_kwargs_for_generation( - outputs, - model_kwargs, - is_encoder_decoder=self.config.is_encoder_decoder, - ) + # synced_gpus: don't waste resources running the code we don't need if synced_gpus and this_peer_finished: continue @@ -2576,6 +2571,12 @@ def _sample( if streamer is not None: streamer.put(next_tokens.cpu()) + model_kwargs = self._update_model_kwargs_for_generation( + outputs, + model_kwargs, + is_encoder_decoder=self.config.is_encoder_decoder, + ) + cur_len = cur_len + 1 if bucket_size > 0 and bucket_internal: # Calculate slice idx for kv cache during the decode phase. @@ -2997,12 +2998,7 @@ def expand_if_needed(tensor, new_size, value, dim=-1): **hpu_graphs_kwargs, ) - # synced_gpus: don't waste resources running the code we don't need; kwargs must be updated before skipping - model_kwargs = self._update_model_kwargs_for_generation( - outputs, - model_kwargs, - is_encoder_decoder=self.config.is_encoder_decoder, - ) + # synced_gpus: don't waste resources running the code we don't need if synced_gpus and this_peer_finished: cur_len = cur_len + 1 continue @@ -3137,6 +3133,12 @@ def expand_if_needed(tensor, new_size, value, dim=-1): else: input_ids = torch.cat([input_ids[beam_idx, :], beam_next_tokens.unsqueeze(-1)], dim=-1) + model_kwargs = self._update_model_kwargs_for_generation( + outputs, + model_kwargs, + is_encoder_decoder=self.config.is_encoder_decoder, + ) + if model_kwargs.get("past_key_values", None) is not None: if model_kwargs["reuse_cache"]: model_kwargs["past_key_values"] = unwrap_deepspeed_model(self).reorder_kv_cache(beam_idx) @@ -3479,12 +3481,7 @@ def _constrained_beam_search( **hpu_graphs_kwargs, ) - # synced_gpus: don't waste resources running the code we don't need; kwargs must be updated before skipping - model_kwargs = self._update_model_kwargs_for_generation( - outputs, - model_kwargs, - is_encoder_decoder=self.config.is_encoder_decoder, - ) + # synced_gpus: don't waste resources running the code we don't need if synced_gpus and this_peer_finished: cur_len = cur_len + 1 continue @@ -3572,6 +3569,12 @@ def _constrained_beam_search( else: input_ids = torch.cat([input_ids[beam_idx, :], beam_next_tokens.unsqueeze(-1)], dim=-1) + model_kwargs = self._update_model_kwargs_for_generation( + outputs, + model_kwargs, + is_encoder_decoder=self.config.is_encoder_decoder, + ) + # This is needed to properly delete outputs.logits which may be very large for first iteration # Otherwise a reference to outputs is kept which keeps the logits alive in the next iteration # IMPORTANT: Note that this should appear BEFORE the call to _reorder_cache() to save the maximum memory diff --git a/optimum/habana/transformers/modeling_utils.py b/optimum/habana/transformers/modeling_utils.py index 8fe0ba7b99..1a71465d9d 100644 --- a/optimum/habana/transformers/modeling_utils.py +++ b/optimum/habana/transformers/modeling_utils.py @@ -119,6 +119,7 @@ GaudiMptBlock, GaudiMptForCausalLM, GaudiMptModel, + GaudiOPTDecoderLayer, GaudiOPTForCausalLM, GaudiOPTLearnedPositionalEmbedding, GaudiPaliGemmaForConditionalGeneration, @@ -218,7 +219,6 @@ gaudi_mixtral_rmsnorm_forward, gaudi_opt_attention_forward, gaudi_opt_decoder_forward, - gaudi_opt_decoder_layer_forward, gaudi_opt_model_forward, gaudi_owlvitclasspredictionhead_forward, gaudi_persimmon_model_forward, @@ -407,7 +407,7 @@ def adapt_transformers_to_gaudi(): transformers.models.opt.modeling_opt.OPTDecoder.forward = gaudi_opt_decoder_forward transformers.models.opt.modeling_opt.OPTForCausalLM = GaudiOPTForCausalLM transformers.models.opt.modeling_opt.OPTModel.forward = gaudi_opt_model_forward - transformers.models.opt.modeling_opt.OPTDecoderLayer.forward = gaudi_opt_decoder_layer_forward + transformers.models.opt.modeling_opt.OPTDecoderLayer = GaudiOPTDecoderLayer transformers.models.opt.modeling_opt.OPTLearnedPositionalEmbedding = GaudiOPTLearnedPositionalEmbedding # Optimization for GPTJ on Gaudi diff --git a/optimum/habana/transformers/models/__init__.py b/optimum/habana/transformers/models/__init__.py index 13b84d48b1..7c81b01c8c 100644 --- a/optimum/habana/transformers/models/__init__.py +++ b/optimum/habana/transformers/models/__init__.py @@ -210,11 +210,11 @@ GaudiMptModel, ) from .opt import ( + GaudiOPTDecoderLayer, GaudiOPTForCausalLM, GaudiOPTLearnedPositionalEmbedding, gaudi_opt_attention_forward, gaudi_opt_decoder_forward, - gaudi_opt_decoder_layer_forward, gaudi_opt_model_forward, ) from .owlvit import gaudi_owlvitclasspredictionhead_forward diff --git a/optimum/habana/transformers/models/gemma/modeling_gemma.py b/optimum/habana/transformers/models/gemma/modeling_gemma.py index 8e34b12b7f..ee4ff65a9e 100755 --- a/optimum/habana/transformers/models/gemma/modeling_gemma.py +++ b/optimum/habana/transformers/models/gemma/modeling_gemma.py @@ -813,7 +813,7 @@ def forward( hidden_states = outputs[0] # Only compute necessary logits, and do not upcast them to float if we are not computing the loss - logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]) + logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]).float() loss = None if labels 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 658147afbe..96a955974c 100644 --- a/optimum/habana/transformers/models/gpt_neox/modeling_gpt_neox.py +++ b/optimum/habana/transformers/models/gpt_neox/modeling_gpt_neox.py @@ -28,6 +28,56 @@ from ..modeling_all_models import apply_customized_rope_module +def gaudi_eager_attention_forward( + query, key, value, attention_mask, head_mask, norm_factor, attention_dropout, training, **_kwargs +): + """ + Copied from: https://github.com/huggingface/transformers/blob/v4.47.1/src/transformers/models/gpt_neox/modeling_gpt_neox.py#L98 + Changes: + - transposition at the end is commented + """ + # q, k, v: [bs, num_attention_heads, seq_len, attn_head_size] + batch_size, num_attention_heads, query_length, attn_head_size = query.size() + key_length = key.size(-2) + + query = query.view(batch_size * num_attention_heads, query_length, attn_head_size) + key = key.view(batch_size * num_attention_heads, key_length, attn_head_size) + attn_scores = torch.zeros( + batch_size * num_attention_heads, + query_length, + key_length, + dtype=query.dtype, + device=key.device, + ) + attn_scores = torch.baddbmm( + attn_scores, + query, + key.transpose(1, 2), + beta=1.0, + alpha=norm_factor, + ) + attn_scores = attn_scores.view(batch_size, num_attention_heads, query_length, key_length) + + if attention_mask is not None: # no matter the length, we just slice it + causal_mask = attention_mask[:, :, :, : key.shape[-2]] + attn_scores = attn_scores + causal_mask + + attn_weights = torch.nn.functional.softmax(attn_scores, dim=-1) + attn_weights = attn_weights.to(value.dtype) + + # Mask heads if we want to + if head_mask is not None: + attn_weights = attn_weights * head_mask + + attn_weights = torch.nn.functional.dropout(attn_weights, p=attention_dropout, training=training) + attn_output = torch.matmul(attn_weights, value) + + # # Reshape outputs + # attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + class GaudiGPTNeoXAttention(GPTNeoXAttention): def __init__(self, config: GPTNeoXConfig, layer_idx=None): super().__init__(config, layer_idx) @@ -52,6 +102,7 @@ def forward( - add new args token_idx - optimize KV cache """ + bsz, seq_len, _ = hidden_states.shape has_layer_past = layer_past is not None # Compute QKV @@ -101,9 +152,18 @@ def forward( present = (key, value) if use_cache else None # Compute attention - attn_output, attn_weights = self._attn(query, key, value, attention_mask, head_mask) + attn_output, attn_weights = gaudi_eager_attention_forward( + query, + key, + value, + attention_mask=attention_mask, + head_mask=head_mask, + norm_factor=self.norm_factor, + attention_dropout=self.config.attention_dropout, + training=self.training, + ) - # Reshape outputs + # Reshape outputs and final projection attn_output = self._merge_heads(attn_output, self.num_attention_heads, self.head_size) attn_output = self.dense(attn_output) diff --git a/optimum/habana/transformers/models/opt/__init__.py b/optimum/habana/transformers/models/opt/__init__.py index 9ea5a435ee..aeaa92cfd3 100644 --- a/optimum/habana/transformers/models/opt/__init__.py +++ b/optimum/habana/transformers/models/opt/__init__.py @@ -1,8 +1,8 @@ from .modeling_opt import ( + GaudiOPTDecoderLayer, GaudiOPTForCausalLM, GaudiOPTLearnedPositionalEmbedding, gaudi_opt_attention_forward, gaudi_opt_decoder_forward, - gaudi_opt_decoder_layer_forward, gaudi_opt_model_forward, ) diff --git a/optimum/habana/transformers/models/opt/modeling_opt.py b/optimum/habana/transformers/models/opt/modeling_opt.py index 179495d776..3a7c99d96e 100644 --- a/optimum/habana/transformers/models/opt/modeling_opt.py +++ b/optimum/habana/transformers/models/opt/modeling_opt.py @@ -2,8 +2,15 @@ import torch from torch.nn import CrossEntropyLoss +from transformers.activations import ACT2FN from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast -from transformers.models.opt.modeling_opt import OPTForCausalLM, OPTLearnedPositionalEmbedding, logger +from transformers.models.opt.configuration_opt import OPTConfig +from transformers.models.opt.modeling_opt import ( + OPT_ATTENTION_CLASSES, + OPTForCausalLM, + OPTLearnedPositionalEmbedding, + logger, +) from ...modeling_attn_mask_utils import _gaudi_prepare_4d_causal_attention_mask @@ -164,75 +171,98 @@ def gaudi_opt_attention_forward( return attn_output, attn_weights_reshaped, past_key_value -def gaudi_opt_decoder_layer_forward( - self, - 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, - 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]]]: - """ - Copied from OPTDecoderLayer.forward: https://github.com/huggingface/transformers/blob/main/src/transformers/models/opt/modeling_opt.py - The only differences are: - - add new args token_idx - """ - residual = hidden_states +class GaudiOPTDecoderLayer(torch.nn.Module): + def __init__(self, config: OPTConfig): + """ + Attention implementation is set to "eager" (default in Transformers is "sdpa"). + """ + super().__init__() + self.embed_dim = config.hidden_size - # 125m, 1.7B, ..., 175B applies layer norm BEFORE attention - if self.do_layer_norm_before: - hidden_states = self.self_attn_layer_norm(hidden_states) + self.self_attn = OPT_ATTENTION_CLASSES["eager"](config=config, is_decoder=True) - # Self Attention - hidden_states, self_attn_weights, present_key_value = self.self_attn( - hidden_states=hidden_states, - past_key_value=past_key_value, - position_ids=position_ids, - attention_mask=attention_mask, - layer_head_mask=layer_head_mask, - output_attentions=output_attentions, - token_idx=token_idx, - ) - hidden_states = torch.nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) - hidden_states = residual + hidden_states + self.do_layer_norm_before = config.do_layer_norm_before + self.dropout = config.dropout + self.activation_fn = ACT2FN[config.activation_function] - # 350m applies layer norm AFTER attention - if not self.do_layer_norm_before: - hidden_states = self.self_attn_layer_norm(hidden_states) + self.self_attn_layer_norm = torch.nn.LayerNorm( + self.embed_dim, elementwise_affine=config.layer_norm_elementwise_affine + ) + self.fc1 = torch.nn.Linear(self.embed_dim, config.ffn_dim, bias=config.enable_bias) + self.fc2 = torch.nn.Linear(config.ffn_dim, self.embed_dim, bias=config.enable_bias) + self.final_layer_norm = torch.nn.LayerNorm( + self.embed_dim, elementwise_affine=config.layer_norm_elementwise_affine + ) - # Fully Connected - hidden_states_shape = hidden_states.shape - hidden_states = hidden_states.reshape(-1, hidden_states.size(-1)) - residual = hidden_states + def forward( + self, + 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, + 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]]]: + """ + Copied from OPTDecoderLayer.forward: https://github.com/huggingface/transformers/blob/main/src/transformers/models/opt/modeling_opt.py + The only differences are: + - add new args token_idx + """ + residual = hidden_states + + # 125m, 1.7B, ..., 175B applies layer norm BEFORE attention + if self.do_layer_norm_before: + hidden_states = self.self_attn_layer_norm(hidden_states) + + # Self Attention + hidden_states, self_attn_weights, present_key_value = self.self_attn( + hidden_states=hidden_states, + past_key_value=past_key_value, + position_ids=position_ids, + attention_mask=attention_mask, + layer_head_mask=layer_head_mask, + output_attentions=output_attentions, + token_idx=token_idx, + ) + hidden_states = torch.nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) + hidden_states = residual + hidden_states - # 125m, 1.7B, ..., 175B applies layer norm BEFORE attention - if self.do_layer_norm_before: - hidden_states = self.final_layer_norm(hidden_states) + # 350m applies layer norm AFTER attention + if not self.do_layer_norm_before: + hidden_states = self.self_attn_layer_norm(hidden_states) - hidden_states = self.fc1(hidden_states) - hidden_states = self.activation_fn(hidden_states) + # Fully Connected + hidden_states_shape = hidden_states.shape + hidden_states = hidden_states.reshape(-1, hidden_states.size(-1)) + residual = hidden_states - hidden_states = self.fc2(hidden_states) - hidden_states = torch.nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) + # 125m, 1.7B, ..., 175B applies layer norm BEFORE attention + if self.do_layer_norm_before: + hidden_states = self.final_layer_norm(hidden_states) - hidden_states = (residual + hidden_states).view(hidden_states_shape) + hidden_states = self.fc1(hidden_states) + hidden_states = self.activation_fn(hidden_states) - # 350m applies layer norm AFTER attention - if not self.do_layer_norm_before: - hidden_states = self.final_layer_norm(hidden_states) + hidden_states = self.fc2(hidden_states) + hidden_states = torch.nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) - outputs = (hidden_states,) + hidden_states = (residual + hidden_states).view(hidden_states_shape) - if output_attentions: - outputs += (self_attn_weights,) + # 350m applies layer norm AFTER attention + if not self.do_layer_norm_before: + hidden_states = self.final_layer_norm(hidden_states) - if use_cache: - outputs += (present_key_value,) + outputs = (hidden_states,) + + if output_attentions: + outputs += (self_attn_weights,) + + if use_cache: + outputs += (present_key_value,) - return outputs + return outputs def gaudi_opt_decoder_forward( @@ -298,7 +328,7 @@ def gaudi_opt_decoder_forward( attention_mask, input_shape, inputs_embeds, past_key_values_length ) - pos_embeds = self.embed_positions(attention_mask, past_key_values_length, token_idx) + pos_embeds = self.embed_positions(attention_mask, past_key_values_length, position_ids, token_idx) if self.project_in is not None: inputs_embeds = self.project_in(inputs_embeds) From c804270d206a45eb54cb9db58a40a505672a51ee Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Fri, 10 Jan 2025 13:44:18 +0000 Subject: [PATCH 22/89] Set eager attention for distilbert, gpt_neox --- .../transformers/models/modeling_all_models.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/optimum/habana/transformers/models/modeling_all_models.py b/optimum/habana/transformers/models/modeling_all_models.py index 3f9304db74..e52d40e206 100644 --- a/optimum/habana/transformers/models/modeling_all_models.py +++ b/optimum/habana/transformers/models/modeling_all_models.py @@ -199,7 +199,17 @@ def gaudi_conv1d_forward(self, x): @classmethod def gaudi_check_and_enable_sdpa(cls, config, hard_check_only: bool = False) -> PretrainedConfig: # This model doesn't support SDPA in Gaudi yet, fallback to original code. - MODELS_ATTN_IMPLEMENTATION_EAGER = ["albert", "bart", "gpt_bigcode", "mistral", "mixtral", "wav2vec2", "roberta"] + MODELS_ATTN_IMPLEMENTATION_EAGER = [ + "albert", + "bart", + "gpt_bigcode", + "mistral", + "mixtral", + "wav2vec2", + "roberta", + "distilbert", + "gpt_neox", + ] if config.model_type in MODELS_ATTN_IMPLEMENTATION_EAGER: config._attn_implementation = "eager" From 0000de522a506701753ec0d868bfe2387585c08c Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Wed, 15 Jan 2025 09:41:46 +0000 Subject: [PATCH 23/89] Upgrade to Transformers v4.48 --- .../run_audio_classification.py | 2 +- .../contrastive-image-text/run_bridgetower.py | 2 +- examples/contrastive-image-text/run_clip.py | 2 +- .../run_image_classification.py | 56 +++-- examples/language-modeling/run_clm.py | 2 +- examples/language-modeling/run_mlm.py | 2 +- .../run_multitask_prompt_tuning.py | 2 +- .../run_prompt_tuning_clm.py | 2 +- examples/question-answering/run_qa.py | 2 +- examples/question-answering/run_seq2seq_qa.py | 2 +- .../run_speech_recognition_ctc.py | 2 +- .../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 +- .../habana/transformers/generation/utils.py | 4 +- .../transformers/modeling_attn_mask_utils.py | 6 + optimum/habana/transformers/modeling_utils.py | 10 +- .../habana/transformers/models/__init__.py | 4 +- .../transformers/models/cohere/__init__.py | 2 +- .../models/cohere/modeling_cohere.py | 177 +++++++------- .../models/falcon/modeling_falcon.py | 4 +- .../models/gemma/modeling_gemma.py | 114 ++++----- .../models/gemma2/modeling_gemma2.py | 110 +++++---- .../transformers/models/gpt2/modeling_gpt2.py | 88 +++---- .../models/gpt_neox/modeling_gpt_neox.py | 1 + .../transformers/models/llama/__init__.py | 2 - .../models/llama/modeling_llama.py | 219 +++++++----------- .../models/minicpm/modeling_minicpm.py | 5 +- .../models/mistral/modeling_mistral.py | 162 ++++++------- .../models/mixtral/modeling_mixtral.py | 146 ++++++------ .../models/mllama/modeling_mllama.py | 2 +- .../models/paligemma/modeling_paligemma.py | 2 +- .../models/persimmon/modeling_persimmon.py | 3 + .../transformers/models/phi/modeling_phi.py | 127 +++++----- .../models/qwen2/modeling_qwen2.py | 145 +++++++----- .../models/qwen2_moe/modeling_qwen2_moe.py | 4 + .../models/stablelm/modeling_stablelm.py | 2 + .../models/starcoder2/modeling_starcoder2.py | 180 +++++++------- optimum/habana/transformers/trainer.py | 31 ++- .../habana/transformers/trainer_seq2seq.py | 8 +- optimum/habana/transformers/training_args.py | 4 +- .../transformers/training_args_seq2seq.py | 5 - setup.py | 2 +- tests/test_trainer.py | 12 +- 46 files changed, 809 insertions(+), 858 deletions(-) diff --git a/examples/audio-classification/run_audio_classification.py b/examples/audio-classification/run_audio_classification.py index 682615a18e..95057317a6 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.47.0") +check_min_version("4.48.0") check_optimum_habana_min_version("1.16.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 42ee164cdf..84876cf906 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.47.0") +check_min_version("4.48.0") check_optimum_habana_min_version("1.16.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 f7ca7f6862..bd3d52bd1d 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.47.0") +check_min_version("4.48.0") check_optimum_habana_min_version("1.16.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 440cf64264..902eefebdd 100644 --- a/examples/image-classification/run_image_classification.py +++ b/examples/image-classification/run_image_classification.py @@ -43,6 +43,7 @@ AutoImageProcessor, AutoModelForImageClassification, HfArgumentParser, + TimmWrapperImageProcessor, ) from transformers.trainer_utils import get_last_checkpoint from transformers.utils import check_min_version, send_example_telemetry @@ -63,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.47.0") +check_min_version("4.48.0") check_optimum_habana_min_version("1.16.0.dev0") require_version("datasets>=2.14.0", "To fix: pip install -r examples/pytorch/image-classification/requirements.txt") @@ -346,31 +347,36 @@ def compute_metrics(p): ) # Define torchvision transforms to be applied to each image. - if "shortest_edge" in image_processor.size: - size = image_processor.size["shortest_edge"] + if isinstance(image_processor, TimmWrapperImageProcessor): + _train_transforms = image_processor.train_transforms + _val_transforms = image_processor.val_transforms else: - size = (image_processor.size["height"], image_processor.size["width"]) - normalize = ( - Normalize(mean=image_processor.image_mean, std=image_processor.image_std) - if hasattr(image_processor, "image_mean") and hasattr(image_processor, "image_std") - else Lambda(lambda x: x) - ) - _train_transforms = Compose( - [ - RandomResizedCrop(size), - RandomHorizontalFlip(), - ToTensor(), - normalize, - ] - ) - _val_transforms = Compose( - [ - Resize(size), - CenterCrop(size), - ToTensor(), - normalize, - ] - ) + if "shortest_edge" in image_processor.size: + size = image_processor.size["shortest_edge"] + else: + size = (image_processor.size["height"], image_processor.size["width"]) + + # Create normalization transform + if hasattr(image_processor, "image_mean") and hasattr(image_processor, "image_std"): + normalize = Normalize(mean=image_processor.image_mean, std=image_processor.image_std) + else: + normalize = Lambda(lambda x: x) + _train_transforms = Compose( + [ + RandomResizedCrop(size), + RandomHorizontalFlip(), + ToTensor(), + normalize, + ] + ) + _val_transforms = Compose( + [ + Resize(size), + CenterCrop(size), + ToTensor(), + normalize, + ] + ) def train_transforms(example_batch): """Apply _train_transforms across a batch.""" diff --git a/examples/language-modeling/run_clm.py b/examples/language-modeling/run_clm.py index 8430792dff..1b4b806004 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.47.0") +check_min_version("4.48.0") check_optimum_habana_min_version("1.16.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 abea9c0eb1..1d95c44ee7 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.47.0") +check_min_version("4.48.0") check_optimum_habana_min_version("1.16.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_multitask_prompt_tuning.py b/examples/language-modeling/run_multitask_prompt_tuning.py index 7f788fc26c..7030a26a3b 100644 --- a/examples/language-modeling/run_multitask_prompt_tuning.py +++ b/examples/language-modeling/run_multitask_prompt_tuning.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 risk. -check_min_version("4.47.0") +check_min_version("4.48.0") check_optimum_habana_min_version("1.16.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 f08280e695..4d7b958ae4 100644 --- a/examples/language-modeling/run_prompt_tuning_clm.py +++ b/examples/language-modeling/run_prompt_tuning_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.47.0") +check_min_version("4.48.0") check_optimum_habana_min_version("1.16.0.dev0") require_version("datasets>=1.8.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 5b93fa5f1b..261daaec4a 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.47.0") +check_min_version("4.48.0") check_optimum_habana_min_version("1.16.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 bc9d9beff4..b6c297f0f4 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.47.0") +check_min_version("4.48.0") check_optimum_habana_min_version("1.16.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 2b0b6093c3..197e74720e 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.47.0") +check_min_version("4.48.0") check_optimum_habana_min_version("1.16.0.dev0") require_version("datasets>=1.18.0", "To fix: pip install -r examples/pytorch/speech-recognition/requirements.txt") diff --git a/examples/speech-recognition/run_speech_recognition_seq2seq.py b/examples/speech-recognition/run_speech_recognition_seq2seq.py index d61973f5c6..5fe794e173 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.47.0") +check_min_version("4.48.0") check_optimum_habana_min_version("1.16.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 f908c4fb9c..4484fee11e 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.47.0") +check_min_version("4.48.0") check_optimum_habana_min_version("1.16.0.dev0") # Setup logging diff --git a/examples/summarization/run_summarization.py b/examples/summarization/run_summarization.py index 97dbe32944..f288bb063a 100755 --- a/examples/summarization/run_summarization.py +++ b/examples/summarization/run_summarization.py @@ -65,7 +65,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.47.0") +check_min_version("4.48.0") check_optimum_habana_min_version("1.16.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 2e9694b404..65c19a0bf2 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.47.0") +check_min_version("4.48.0") check_optimum_habana_min_version("1.16.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 1a6f3379aa..5c6e7f4bfd 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.47.0") +check_min_version("4.48.0") check_optimum_habana_min_version("1.16.0.dev0") require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/translation/requirements.txt") diff --git a/optimum/habana/transformers/generation/utils.py b/optimum/habana/transformers/generation/utils.py index 3486463480..f55ff55220 100644 --- a/optimum/habana/transformers/generation/utils.py +++ b/optimum/habana/transformers/generation/utils.py @@ -672,6 +672,7 @@ def _prepare_generated_length( elif ( model_input_name == "inputs_embeds" and input_ids_length != inputs_tensor.shape[1] + and input_ids_length != 0 and not self.config.is_encoder_decoder ): generation_config.max_length -= inputs_tensor.shape[1] @@ -3762,9 +3763,10 @@ def _assisted_decoding( model_kwargs["lazy_mode"] = lazy_mode model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs) - # 1. Fetch candidate sequences from a `CandidateGenerator` + # 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) diff --git a/optimum/habana/transformers/modeling_attn_mask_utils.py b/optimum/habana/transformers/modeling_attn_mask_utils.py index 4d2b928620..eb1ba79ed4 100755 --- a/optimum/habana/transformers/modeling_attn_mask_utils.py +++ b/optimum/habana/transformers/modeling_attn_mask_utils.py @@ -16,6 +16,7 @@ import torch from transformers.modeling_attn_mask_utils import AttentionMaskConverter +from transformers.utils.import_utils import is_torchdynamo_compiling @dataclass @@ -57,6 +58,11 @@ def _make_causal_mask( col_indices = torch.arange(mask.size(1), device=mask.device) context_mask = (col_indices <= row_indices + diagonal).bool().expand_as(mask) # Expand to match mask shape + # Recent changes in PyTorch prevent mutations on tensors converted with aten::_to_copy + # See https://github.com/pytorch/pytorch/issues/127571 + if is_torchdynamo_compiling(): + mask = mask.clone() + mask.masked_fill_(context_mask, torch.finfo(dtype).min) return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length) diff --git a/optimum/habana/transformers/modeling_utils.py b/optimum/habana/transformers/modeling_utils.py index 1a71465d9d..c86a245bf7 100644 --- a/optimum/habana/transformers/modeling_utils.py +++ b/optimum/habana/transformers/modeling_utils.py @@ -51,6 +51,7 @@ GaudiCLIPVisionTransformer, GaudiCodeGenAttention, GaudiCodeGenForCausalLM, + GaudiCohereAttention, GaudiCohereDecoderLayer, GaudiCohereForCausalLM, GaudiFalconAttention, @@ -88,9 +89,7 @@ GaudiIdefics2VisionEmbeddings, GaudiLlamaAttention, GaudiLlamaDecoderLayer, - GaudiLlamaDynamicNTKScalingRotaryEmbedding, GaudiLlamaForCausalLM, - GaudiLlamaLinearScalingRotaryEmbedding, GaudiLlamaMLP, GaudiLlamaModel, GaudiLlamaRotaryEmbedding, @@ -189,7 +188,6 @@ gaudi_check_and_enable_sdpa, gaudi_codegen_block_forward, gaudi_codegen_model_forward, - gaudi_cohere_attention_forward, gaudi_cohere_model_forward, gaudi_conv1d_forward, gaudi_DetrConvModel_forward, @@ -445,10 +443,6 @@ def adapt_transformers_to_gaudi(): transformers.models.llama.modeling_llama.LlamaMLP = GaudiLlamaMLP transformers.models.llama.modeling_llama.LlamaDecoderLayer = GaudiLlamaDecoderLayer transformers.models.llama.modeling_llama.LlamaRotaryEmbedding = GaudiLlamaRotaryEmbedding - transformers.models.llama.modeling_llama.LlamaLinearScalingRotaryEmbedding = GaudiLlamaLinearScalingRotaryEmbedding - transformers.models.llama.modeling_llama.LlamaDynamicNTKScalingRotaryEmbedding = ( - GaudiLlamaDynamicNTKScalingRotaryEmbedding - ) transformers.models.llama.modeling_llama.LlamaRMSNorm.forward = gaudi_llama_rmsnorm_forward transformers.models.llama.configuration_llama.LlamaConfig = LlamaConfig @@ -706,7 +700,7 @@ def adapt_transformers_to_gaudi(): transformers.models.cohere.modeling_cohere.CohereDecoderLayer = GaudiCohereDecoderLayer transformers.models.cohere.modeling_cohere.CohereForCausalLM = GaudiCohereForCausalLM transformers.models.cohere.modeling_cohere.CohereModel.forward = gaudi_cohere_model_forward - transformers.models.cohere.modeling_cohere.CohereAttention.forward = gaudi_cohere_attention_forward + transformers.models.cohere.modeling_cohere.CohereAttention = GaudiCohereAttention # Optimization for xglm on Gaudi transformers.models.xglm.modeling_xglm.XGLMForCausalLM = GaudiXGLMForCausalLM diff --git a/optimum/habana/transformers/models/__init__.py b/optimum/habana/transformers/models/__init__.py index 7c81b01c8c..1b4af85036 100644 --- a/optimum/habana/transformers/models/__init__.py +++ b/optimum/habana/transformers/models/__init__.py @@ -57,9 +57,9 @@ gaudi_codegen_model_forward, ) from .cohere import ( + GaudiCohereAttention, GaudiCohereDecoderLayer, GaudiCohereForCausalLM, - gaudi_cohere_attention_forward, gaudi_cohere_model_forward, ) from .decilm import ( @@ -146,9 +146,7 @@ from .llama import ( GaudiLlamaAttention, GaudiLlamaDecoderLayer, - GaudiLlamaDynamicNTKScalingRotaryEmbedding, GaudiLlamaForCausalLM, - GaudiLlamaLinearScalingRotaryEmbedding, GaudiLlamaMLP, GaudiLlamaModel, GaudiLlamaRotaryEmbedding, diff --git a/optimum/habana/transformers/models/cohere/__init__.py b/optimum/habana/transformers/models/cohere/__init__.py index ec3a43831c..94e2ddb055 100644 --- a/optimum/habana/transformers/models/cohere/__init__.py +++ b/optimum/habana/transformers/models/cohere/__init__.py @@ -1,6 +1,6 @@ from .modeling_cohere import ( + GaudiCohereAttention, GaudiCohereDecoderLayer, GaudiCohereForCausalLM, - gaudi_cohere_attention_forward, gaudi_cohere_model_forward, ) diff --git a/optimum/habana/transformers/models/cohere/modeling_cohere.py b/optimum/habana/transformers/models/cohere/modeling_cohere.py index 119df106fb..119989988b 100644 --- a/optimum/habana/transformers/models/cohere/modeling_cohere.py +++ b/optimum/habana/transformers/models/cohere/modeling_cohere.py @@ -1,8 +1,6 @@ -import math from typing import List, Optional, Tuple, Union import torch -from torch import nn from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from transformers.models.cohere.modeling_cohere import ( Cache, @@ -10,120 +8,104 @@ CohereConfig, CohereDecoderLayer, CohereForCausalLM, - CohereLayerNorm, - CohereMLP, + CohereRotaryEmbedding, DynamicCache, + KwargsForCausalLM, StaticCache, apply_rotary_pos_emb, + eager_attention_forward, logger, - repeat_kv, ) +from transformers.processing_utils import Unpack from ...modeling_attn_mask_utils import _gaudi_prepare_4d_causal_attention_mask -def gaudi_cohere_attention_forward( - self, - hidden_states: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, - past_key_value: Optional[Cache] = None, - output_attentions: bool = False, - use_cache: bool = False, - cache_position: Optional[torch.LongTensor] = None, - token_idx: Optional[torch.Tensor] = None, -) -> 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: - - add new args token_idx - - optimize KV cache - """ - bsz, q_len, _ = hidden_states.size() - - query_states = self.q_proj(hidden_states) - key_states = self.k_proj(hidden_states) - value_states = self.v_proj(hidden_states) - - query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim) - key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim) - if self.use_qk_norm: - query_states = self.q_norm(query_states) - key_states = self.k_norm(key_states) - - query_states = query_states.transpose(1, 2) - key_states = key_states.transpose(1, 2) - value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) - - cos, sin = self.rotary_emb(value_states, position_ids) - query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) - - if past_key_value is not None: - # sin and cos are specific to RoPE models; position_ids needed for the static cache - cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} - if token_idx is not None: - if len(past_key_value.key_cache) <= self.layer_idx: - past_key_value.key_cache.append(key_states) - past_key_value.value_cache.append(value_states) - else: - past_key_value.key_cache[self.layer_idx].index_copy_(2, token_idx - 1, key_states) - past_key_value.value_cache[self.layer_idx].index_copy_(2, token_idx - 1, value_states) - key_states = past_key_value.key_cache[self.layer_idx] - value_states = past_key_value.value_cache[self.layer_idx] - else: - key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) - - key_states = repeat_kv(key_states, self.num_key_value_groups) - value_states = repeat_kv(value_states, self.num_key_value_groups) - - attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim) +class GaudiCohereAttention(CohereAttention): + def __init__(self, config: CohereConfig, layer_idx: Optional[int] = None): + super().__init__(config, layer_idx) - 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 + self.rotary_emb = CohereRotaryEmbedding(config=config) - # upcast attention to fp32 - attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) - attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training) - 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()}" + def forward( + self, + hidden_states: 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]]]: + """ + Copied from CohereAttention.forward: https://github.com/huggingface/transformers/blob/main/src/transformers/models/cohere/modeling_cohere.py + The only differences are: + - add new args token_idx + - optimize KV cache + """ + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + query_states = self.q_proj(hidden_states).view(hidden_shape) + key_states = self.k_proj(hidden_states).view(hidden_shape) + value_states = self.v_proj(hidden_states).view(hidden_shape) + + if self.use_qk_norm: # main diff from Llama + query_states = self.q_norm(query_states) + key_states = self.k_norm(key_states) + + query_states = query_states.transpose(1, 2) + key_states = key_states.transpose(1, 2) + value_states = value_states.transpose(1, 2) + + cos, sin = self.rotary_emb(value_states, kwargs["position_ids"]) + # print("SHAPEEEEEEEEEEEE", cos.shape, sin.shape, query_states.shape, key_states.shape) + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + if past_key_value is not None: + # sin and cos are specific to RoPE models; position_ids needed for the static cache + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} + if token_idx is not None: + if len(past_key_value.key_cache) <= self.layer_idx: + past_key_value.key_cache.append(key_states) + past_key_value.value_cache.append(value_states) + else: + past_key_value.key_cache[self.layer_idx].index_copy_(2, token_idx - 1, key_states) + past_key_value.value_cache[self.layer_idx].index_copy_(2, token_idx - 1, value_states) + key_states = past_key_value.key_cache[self.layer_idx] + value_states = past_key_value.value_cache[self.layer_idx] + else: + key_states, value_states = past_key_value.update( + key_states, value_states, self.layer_idx, cache_kwargs + ) + + 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.attention_dropout, + scaling=self.scaling, ) - attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) - attn_output = attn_output.reshape(bsz, q_len, self.hidden_size) - - attn_output = self.o_proj(attn_output) - - if not output_attentions: - attn_weights = None - - return attn_output, attn_weights, past_key_value + return attn_output, attn_weights, past_key_value class GaudiCohereDecoderLayer(CohereDecoderLayer): - def __init__(self, config: CohereConfig, layer_idx: int): - super(CohereDecoderLayer, self).__init__() - self.hidden_size = config.hidden_size - - self.self_attn = CohereAttention(config=config, layer_idx=layer_idx) - - self.mlp = CohereMLP(config) - self.input_layernorm = CohereLayerNorm(hidden_size=(config.hidden_size), eps=config.layer_norm_eps) - def forward( self, 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[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 token_idx: Optional[torch.Tensor] = None, ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: """ @@ -144,6 +126,7 @@ def forward( output_attentions=output_attentions, use_cache=use_cache, cache_position=cache_position, + position_embeddings=position_embeddings, token_idx=token_idx, ) @@ -154,10 +137,8 @@ def forward( 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,) @@ -169,7 +150,7 @@ def gaudi_cohere_model_forward( input_ids: 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[Cache] = None, inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, @@ -299,7 +280,7 @@ def forward( input_ids: 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[Union[Cache, List[torch.FloatTensor]]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, @@ -309,7 +290,7 @@ def forward( cache_position: Optional[torch.LongTensor] = None, num_logits_to_keep: int = 0, token_idx: Optional[torch.Tensor] = None, - **loss_kwargs, + **kwargs: Unpack[KwargsForCausalLM], ) -> Union[Tuple, CausalLMOutputWithPast]: output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = ( @@ -335,11 +316,11 @@ def forward( hidden_states = outputs[0] # Only compute necessary logits, and do not upcast them to float if we are not computing the loss logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]) - logits = logits * self.logit_scale + logits = logits * self.logit_scale # main diff from Llama loss = None if labels is not None: - loss = self.loss_function(logits, labels, self.vocab_size, **loss_kwargs) + loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs) if not return_dict: output = (logits,) + outputs[1:] diff --git a/optimum/habana/transformers/models/falcon/modeling_falcon.py b/optimum/habana/transformers/models/falcon/modeling_falcon.py index 92e42deb33..ddc52a4a74 100644 --- a/optimum/habana/transformers/models/falcon/modeling_falcon.py +++ b/optimum/habana/transformers/models/falcon/modeling_falcon.py @@ -282,7 +282,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, # will become mandatory in v4.46 + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, token_idx: Optional[torch.Tensor] = None, reuse_cache: Optional[bool] = False, cache_idx: int = None, @@ -576,7 +576,7 @@ def forward( 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, token_idx: Optional[torch.Tensor] = None, reuse_cache: Optional[bool] = False, cache_idx: int = None, diff --git a/optimum/habana/transformers/models/gemma/modeling_gemma.py b/optimum/habana/transformers/models/gemma/modeling_gemma.py index ee4ff65a9e..7ec22d6c12 100755 --- a/optimum/habana/transformers/models/gemma/modeling_gemma.py +++ b/optimum/habana/transformers/models/gemma/modeling_gemma.py @@ -34,8 +34,10 @@ GemmaForCausalLM, GemmaMLP, GemmaModel, + KwargsForCausalLM, apply_rotary_pos_emb, ) +from transformers.processing_utils import Unpack from transformers.utils import logging from ...modeling_attn_mask_utils import ( @@ -161,6 +163,37 @@ def forward(self, cur, dim, idx): return self.update(self.cache, cur, dim, idx, self.inp_seq_len) +def 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, +): + query_states, key_states, value_states, attention_mask = gaudi_gemma_repeat_kv( + query, key, value, attention_mask, module.num_key_value_groups + ) + + attn_weights = module.matmul_qk(query, key_states.transpose(2, 3)) * scaling + if attention_mask is not None: + causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] + attn_weights = attn_weights + causal_mask + + if kwargs["attn_softmax_bf16"]: + attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1, dtype=query_states.dtype) + else: + # upcast attention to fp32 + 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=dropout, training=module.training) + attn_output = module.matmul_av(attn_weights, value_states) + # attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + class GaudiGemmaAttention(GemmaAttention): def __init__(self, config: GemmaConfig, layer_idx: Optional[int] = None): super().__init__(config, layer_idx) @@ -170,7 +203,6 @@ def __init__(self, config: GemmaConfig, layer_idx: Optional[int] = None): self.k_cache = KVCache() self.v_cache = KVCache() self.inp_seq_len = -1 - self.norm_factor = 1.0 / math.sqrt(self.head_dim) self.block_size = 4096 self.rotary_emb = GaudiRotaryEmbedding(config=self.config) @@ -238,10 +270,9 @@ def gaudi_flash_attn_v1( def pre_attn_forward( self, hidden_states: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, + position_embeddings: Tuple[torch.Tensor, torch.Tensor], + attention_mask: Optional[torch.Tensor], past_key_value: Optional[Cache] = None, - output_attentions: bool = False, use_cache: bool = False, cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, @@ -262,20 +293,13 @@ def pre_attn_forward( - add new args use_flash_attention - add new arg flash_attention_recompute """ - if "padding_mask" in kwargs: - logger.warning_once( - "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`" - ) + input_shape = hidden_states.shape[:-1] + q_len = input_shape[1] + hidden_shape = (*input_shape, -1, self.head_dim) - bsz, q_len, _ = hidden_states.size() - - query_states = self.q_proj(hidden_states) - key_states = self.k_proj(hidden_states) - value_states = self.v_proj(hidden_states) - - query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) - key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) - value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) + key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) + value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) kv_seq_len = key_states.shape[-2] if past_key_value is not None: @@ -291,7 +315,9 @@ def pre_attn_forward( kv_seq_len = past_key_value[0].shape[-2] cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len) - query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos[position_ids], sin[position_ids]) + query_states, key_states = apply_rotary_pos_emb( + query_states, key_states, cos[kwargs["position_ids"]], sin[kwargs["position_ids"]] + ) if use_cache: # reuse k, v, self_attention @@ -321,6 +347,7 @@ def pre_attn_forward( past_key_value = None if use_flash_attention and FusedSDPA: + attn_weights = None if q_len == 1: # next token use_recompute = True if os.getenv("QUANT_CONFIG", "") else False @@ -359,43 +386,22 @@ def pre_attn_forward( ) else: - query_states, key_states, value_states, attention_mask = gaudi_gemma_repeat_kv( - query_states, key_states, value_states, attention_mask, self.num_key_value_groups - ) - - attn_weights = self.matmul_qk(query_states, key_states.transpose(-2, -1)) * self.norm_factor - - if attention_mask is not None: # no matter the length, we just slice it - causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] - if cache_position is not None: - causal_mask = attention_mask[:, :, cache_position, : key_states.shape[-2]] - attn_weights = attn_weights + causal_mask - - if attn_softmax_bf16: - attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1, dtype=query_states.dtype) - else: - # upcast attention to fp32 - 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.attention_dropout, training=self.training) - attn_output = self.matmul_av(attn_weights, value_states) - attn_output = attn_output.reshape(bsz, -1, q_len, self.head_dim) - - 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()}" + kwargs["attn_softmax_bf16"] = attn_softmax_bf16 + 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.attention_dropout, + scaling=self.scaling, + **kwargs, ) attn_output = attn_output.transpose(1, 2).contiguous() - - attn_output = attn_output.reshape(bsz, q_len, -1) + attn_output = attn_output.reshape(*input_shape, -1).contiguous() attn_output = self.o_proj(attn_output) - if not output_attentions: - attn_weights = None - return attn_output, attn_weights, past_key_value def attention_all_reduce(self, attn_output): @@ -448,6 +454,7 @@ def pre_attn( 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 token_idx: Optional[torch.Tensor] = None, attn_softmax_bf16: Optional[bool] = False, reuse_cache: Optional[bool] = False, @@ -465,6 +472,7 @@ def pre_attn( output_attentions=output_attentions, use_cache=use_cache, cache_position=cache_position, + position_embeddings=position_embeddings, token_idx=token_idx, attn_softmax_bf16=attn_softmax_bf16, reuse_cache=reuse_cache, @@ -526,10 +534,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,) @@ -777,7 +783,7 @@ def forward( flash_attention_recompute: Optional[bool] = False, flash_attention_causal_mask: Optional[bool] = False, attn_softmax_bf16: Optional[bool] = False, - **loss_kwargs, + **kwargs: Unpack[KwargsForCausalLM], ) -> Union[Tuple, CausalLMOutputWithPast]: """ Inherits from GemmaForCausalLM: https://github.com/huggingface/transformers/blob/v4.38.1/src/transformers/models/gemma/modeling_gemma.py @@ -817,7 +823,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=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs) if not return_dict: output = (logits,) + outputs[1:] diff --git a/optimum/habana/transformers/models/gemma2/modeling_gemma2.py b/optimum/habana/transformers/models/gemma2/modeling_gemma2.py index 5927b04285..9cd07b560d 100755 --- a/optimum/habana/transformers/models/gemma2/modeling_gemma2.py +++ b/optimum/habana/transformers/models/gemma2/modeling_gemma2.py @@ -99,7 +99,7 @@ def __init__( self.original_max_seq_len = max_position_embeddings else: # BC: "rope_type" was originally "type" - if config.rope_scaling is not None: + if hasattr(config, "rope_scaling") and config.rope_scaling is not None: self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type")) else: self.rope_type = "default" @@ -243,16 +243,56 @@ def forward(self, cur, dim, idx): return self.update(self.cache, cur, dim, idx, self.inp_seq_len) +def gaudi_eager_attention_forward( + module: torch.nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: Optional[torch.Tensor], + dropout: float = 0.0, + scaling: Optional[float] = None, + softcap: Optional[float] = None, + **kwargs, +) -> Tuple[torch.Tensor, torch.Tensor]: + if scaling is None: + scaling = module.head_dim**-0.5 + + query_states, key_states, value_states, attention_mask = gaudi_gemma2_repeat_kv( + query, key, value, attention_mask, module.num_key_value_groups + ) + + attn_weights = module.matmul_qk(query_states, key_states.transpose(2, 3)) * scaling + + if softcap is not None: + attn_weights = attn_weights / softcap + attn_weights = torch.tanh(attn_weights) + attn_weights = attn_weights * softcap + 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 = torch.nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = torch.nn.functional.dropout(attn_weights, p=dropout, training=module.training) + attn_output = module.matmul_av(attn_weights, value_states) + return attn_output, attn_weights + + class GaudiGemma2Attention(Gemma2Attention): def __init__(self, config: Gemma2Config, layer_idx: Optional[int] = None): super().__init__(config, layer_idx) + self.rotary_emb = GaudiGemma2RotaryEmbedding( + self.head_dim, + max_position_embeddings=config.max_position_embeddings, + base=config.rope_theta, + ) + self.matmul_qk = Matmul() self.matmul_av = Matmul() self.k_cache = KVCache() self.v_cache = KVCache() self.inp_seq_len = -1 - self.norm_factor = 1.0 / math.sqrt(self.head_dim) self.block_size = 4096 def allocate_kv_cache(self, batch_size, max_seq_len, inp_seq_len): @@ -313,10 +353,9 @@ def gaudi_flash_attn_v1(self, query_layer, key_layer, value_layer, attention_mas def pre_attn_forward( self, hidden_states: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, + position_embeddings: Tuple[torch.Tensor, torch.Tensor], + attention_mask: Optional[torch.Tensor], past_key_value: Optional[Cache] = None, - output_attentions: bool = False, use_cache: bool = False, cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, @@ -338,15 +377,13 @@ def pre_attn_forward( - add new args use_flash_attention - add new arg flash_attention_recompute """ - bsz, q_len, _ = hidden_states.size() + input_shape = hidden_states.shape[:-1] + q_len = input_shape[1] + hidden_shape = (*input_shape, -1, self.head_dim) - query_states = self.q_proj(hidden_states) - key_states = self.k_proj(hidden_states) - value_states = self.v_proj(hidden_states) - - query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) - key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) - value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) + key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) + value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) kv_seq_len = key_states.shape[-2] if past_key_value is not None: @@ -362,7 +399,7 @@ def pre_attn_forward( kv_seq_len = past_key_value[0].shape[-2] cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len) - query_states, key_states = apply_customized_rope(query_states, key_states, cos, sin, position_ids) + query_states, key_states = apply_customized_rope(query_states, key_states, cos, sin, kwargs["position_ids"]) if use_cache: # reuse k, v, self_attention @@ -392,6 +429,7 @@ def pre_attn_forward( past_key_value = None if use_flash_attention and FusedSDPA: + attn_weights = None import habana_frameworks.torch.hpu as ht softmax_mode = "fast" if flash_attention_fast_softmax else "None" @@ -421,40 +459,24 @@ def pre_attn_forward( ) else: - query_states, key_states, value_states, attention_mask = gaudi_gemma2_repeat_kv( - query_states, key_states, value_states, attention_mask, self.num_key_value_groups - ) - - attn_weights = self.matmul_qk(query_states, key_states.transpose(-2, -1)) * self.norm_factor - - if attention_mask is not None: # no matter the length, we just slice it - causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] - if cache_position is not None: - causal_mask = attention_mask[:, :, cache_position, : key_states.shape[-2]] - attn_weights = attn_weights + causal_mask - - # upcast attention to fp32 - 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.attention_dropout, training=self.training) - attn_output = self.matmul_av(attn_weights, value_states) - attn_output = attn_output.reshape(bsz, -1, q_len, self.head_dim) - - 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_weights = gaudi_eager_attention_forward( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=self.attention_dropout if self.training else 0.0, + scaling=self.scaling, + sliding_window=self.sliding_window, + softcap=self.attn_logit_softcapping, + **kwargs, ) attn_output = attn_output.transpose(1, 2).contiguous() - attn_output = attn_output.reshape(bsz, q_len, -1) + attn_output = attn_output.reshape(*input_shape, -1).contiguous() attn_output = self.o_proj(attn_output) - if not output_attentions: - attn_weights = None - if not reuse_cache and token_idx is not None and cache_idx is not None and q_len == 1: # Return only past key value shapes and not the tensors during decode phase (q len is 1) # to avoid making past key values as persistent output tensors of HPU graphs. @@ -506,6 +528,7 @@ def update_sincos_cache(self, seq_len): def pre_attn( self, hidden_states: 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, @@ -525,6 +548,7 @@ def pre_attn( hidden_states, attn_weights, present_key_value = self.self_attn.pre_attn_forward( hidden_states=hidden_states, + position_embeddings=position_embeddings, attention_mask=attention_mask, position_ids=position_ids, past_key_value=past_key_value, @@ -545,6 +569,7 @@ def pre_attn( def forward( self, hidden_states: torch.Tensor, + 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, @@ -569,6 +594,7 @@ def forward( hidden_states, self_attn_weights, present_key_value = self.pre_attn( hidden_states, + position_embeddings, attention_mask, position_ids, past_key_value, diff --git a/optimum/habana/transformers/models/gpt2/modeling_gpt2.py b/optimum/habana/transformers/models/gpt2/modeling_gpt2.py index 8c226a458b..546ee7ef47 100644 --- a/optimum/habana/transformers/models/gpt2/modeling_gpt2.py +++ b/optimum/habana/transformers/models/gpt2/modeling_gpt2.py @@ -9,6 +9,7 @@ GPT2DoubleHeadsModel, GPT2DoubleHeadsModelOutput, GPT2LMHeadModel, + eager_attention_forward, logger, ) @@ -20,48 +21,6 @@ class GaudiGPT2Attention(GPT2Attention): - optimize KV cache """ - def _attn(self, query, key, value, attention_mask=None, head_mask=None): - key = key.contiguous() - value = value.contiguous() - attn_weights = torch.matmul(query, key.transpose(-1, -2)) - - if self.scale_attn_weights: - attn_weights = attn_weights / torch.full( - [], value.size(-1) ** 0.5, dtype=attn_weights.dtype, device=attn_weights.device - ) - - # Layer-wise attention scaling - if self.scale_attn_by_inverse_layer_idx: - attn_weights = attn_weights / float(self.layer_idx + 1) - - if not self.is_cross_attention: - # if only "normal" attention layer implements causal mask - query_length, key_length = query.size(-2), key.size(-2) - causal_mask = self.bias[:, :, key_length - query_length : key_length, :key_length] - 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.full([], mask_value, dtype=attn_weights.dtype, device=attn_weights.device) - attn_weights = torch.where(causal_mask, attn_weights.to(attn_weights.dtype), mask_value) - - if attention_mask is not None: - # Apply the attention mask - attn_weights = attn_weights + attention_mask - - attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1) - - # Downcast (if necessary) back to V's dtype (if in mixed-precision) -- No-Op otherwise - attn_weights = attn_weights.type(value.dtype) - attn_weights = self.attn_dropout(attn_weights) - - # Mask heads if we want to - if head_mask is not None: - attn_weights = attn_weights * head_mask - - attn_output = torch.matmul(attn_weights, value) - - return attn_output, attn_weights - def _upcast_and_reordered_attn(self, query, key, value, attention_mask=None, head_mask=None): key = key.contiguous() value = value.contiguous() @@ -133,38 +92,51 @@ def forward( "Please make sure to instantiate class with `GPT2Attention(..., is_cross_attention=True)`." ) - query = self.q_attn(hidden_states) - key, value = self.c_attn(encoder_hidden_states).split(self.split_size, dim=2) + query_states = self.q_attn(hidden_states) + key_states, value_states = self.c_attn(encoder_hidden_states).split(self.split_size, dim=2) attention_mask = encoder_attention_mask else: - query, key, value = self.c_attn(hidden_states).split(self.split_size, dim=2) + query_states, key_states, value_states = self.c_attn(hidden_states).split(self.split_size, dim=2) + + shape_q = (*query_states.shape[:-1], -1, self.head_dim) + shape_kv = (*key_states.shape[:-1], -1, self.head_dim) - query = self._split_heads(query, self.num_heads, self.head_dim).contiguous() - key = self._split_heads(key, self.num_heads, self.head_dim).contiguous() - value = self._split_heads(value, self.num_heads, self.head_dim).contiguous() + query_states = query_states.view(shape_q).transpose(1, 2).contiguous() + key_states = key_states.view(shape_kv).transpose(1, 2).contiguous() + value_states = value_states.view(shape_kv).transpose(1, 2).contiguous() if layer_past is not None: past_key, past_value = layer_past if token_idx is not None: - past_key.index_copy_(2, token_idx - 1, key) - past_value.index_copy_(2, token_idx - 1, value) - key = past_key - value = past_value + past_key.index_copy_(2, token_idx - 1, key_states) + past_value.index_copy_(2, token_idx - 1, value_states) + key_states = past_key + value_states = past_value else: - key = torch.cat((past_key, key), dim=-2) - value = torch.cat((past_value, value), dim=-2) + key_states = torch.cat((past_key, key_states), dim=-2) + value_states = torch.cat((past_value, value_states), dim=-2) if use_cache is True: - present = (key, value) + present = (key_states, value_states) else: present = None if self.reorder_and_upcast_attn: - attn_output, attn_weights = self._upcast_and_reordered_attn(query, key, value, attention_mask, head_mask) + attn_output, attn_weights = self._upcast_and_reordered_attn( + query_states, key_states, value_states, attention_mask, head_mask + ) else: - attn_output, attn_weights = self._attn(query, key, value, attention_mask, head_mask) + attn_output, attn_weights = eager_attention_forward( + self, + query_states, + key_states, + value_states, + attention_mask, + head_mask=head_mask, + dropout=self.attn_dropout.p if self.training else 0.0, + ) - attn_output = self._merge_heads(attn_output, self.num_heads, self.head_dim) + attn_output = attn_output.reshape(*attn_output.shape[:-2], -1).contiguous() attn_output = self.c_proj(attn_output) attn_output = self.resid_dropout(attn_output) 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 96a955974c..4f4a152c67 100644 --- a/optimum/habana/transformers/models/gpt_neox/modeling_gpt_neox.py +++ b/optimum/habana/transformers/models/gpt_neox/modeling_gpt_neox.py @@ -94,6 +94,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 token_idx: Optional[torch.Tensor] = None, ): """ diff --git a/optimum/habana/transformers/models/llama/__init__.py b/optimum/habana/transformers/models/llama/__init__.py index 0a8758d894..ae6a8ecaa7 100644 --- a/optimum/habana/transformers/models/llama/__init__.py +++ b/optimum/habana/transformers/models/llama/__init__.py @@ -2,9 +2,7 @@ from .modeling_llama import ( GaudiLlamaAttention, GaudiLlamaDecoderLayer, - GaudiLlamaDynamicNTKScalingRotaryEmbedding, GaudiLlamaForCausalLM, - GaudiLlamaLinearScalingRotaryEmbedding, GaudiLlamaMLP, GaudiLlamaModel, GaudiLlamaRotaryEmbedding, diff --git a/optimum/habana/transformers/models/llama/modeling_llama.py b/optimum/habana/transformers/models/llama/modeling_llama.py index da26c16567..eb4e32d53f 100755 --- a/optimum/habana/transformers/models/llama/modeling_llama.py +++ b/optimum/habana/transformers/models/llama/modeling_llama.py @@ -1,5 +1,4 @@ import copy -import math from typing import List, Optional, Tuple, Union import torch @@ -84,48 +83,30 @@ def gaudi_llama_rmsnorm_forward(self, hidden_states): class GaudiLlamaRotaryEmbedding(torch.nn.Module): - def __init__( - self, - dim=None, - max_position_embeddings=2048, - base=10000, - device=None, - scaling_factor=1.0, - rope_type="default", - config: Optional[LlamaConfig] = None, - ): + def __init__(self, config: LlamaConfig, device=None): super().__init__() - # TODO (joao): remove the `if` below, only used for BC - self.rope_kwargs = {} - if config is None: - logger.warning_once( - "`LlamaRotaryEmbedding` can now be fully parameterized by passing the model config through the " - "`config` argument. All other arguments will be removed in v4.46" - ) - self.rope_kwargs = { - "rope_type": rope_type, - "factor": scaling_factor, - "dim": dim, - "base": base, - "max_position_embeddings": max_position_embeddings, - } - self.rope_type = rope_type - self.max_seq_len_cached = max_position_embeddings - self.original_max_seq_len = max_position_embeddings + # BC: "rope_type" was originally "type" + if hasattr(config, "rope_scaling") and config.rope_scaling is not None: + self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type")) else: - # BC: "rope_type" was originally "type" - if config.rope_scaling is not None: - self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type")) - else: - self.rope_type = "default" - self.max_seq_len_cached = config.max_position_embeddings - self.original_max_seq_len = config.max_position_embeddings + self.rope_type = "default" + self.max_seq_len_cached = config.max_position_embeddings + self.original_max_seq_len = config.max_position_embeddings + + if self.rope_type == "linear": + self.scaling_factor = config.rope_scaling["factor"] + elif self.rope_type == "dynamic": + self.scaling_factor = config.rope_scaling["factor"] + self.base = config.rope_theta + partial_rotary_factor = config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0 + head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + self.dim = int(head_dim * partial_rotary_factor) self.config = config self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] - inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device, **self.rope_kwargs) + inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device) self.register_buffer("inv_freq", inv_freq, persistent=False) self.original_inv_freq = self.inv_freq @@ -136,8 +117,19 @@ def __init__( def _set_cos_sin_cache(self, seq_len, device, dtype): self.max_seq_len_cached = seq_len + + if self.rope_type == "dynamic" and seq_len > self.config.max_position_embeddings: + base = self.base * ( + (self.scaling_factor * seq_len / self.config.max_position_embeddings) - (self.scaling_factor - 1) + ) ** (self.dim / (self.dim - 2)) + inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim)) + self.register_buffer("inv_freq", inv_freq, persistent=False) + t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype) + if self.rope_type == "linear": + t = t / self.scaling_factor + freqs = torch.outer(t, self.inv_freq) # Different from paper, but it uses a different permutation in order to obtain the same calculation emb = torch.cat((freqs, freqs), dim=-1) @@ -152,9 +144,7 @@ def _dynamic_frequency_update(self, seq_len, device): """ # seq_len = torch.max(position_ids) + 1 if seq_len > self.max_seq_len_cached: # growth - inv_freq, self.attention_scaling = self.rope_init_fn( - self.config, device, seq_len=seq_len, **self.rope_kwargs - ) + inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device, seq_len=seq_len) self.register_buffer("inv_freq", inv_freq, persistent=False) # TODO joao: may break with compilation self.max_seq_len_cached = seq_len @@ -184,56 +174,6 @@ def forward(self, x, seq_len=None): ) -class GaudiLlamaLinearScalingRotaryEmbedding(GaudiLlamaRotaryEmbedding): - def __init__(self, *args, **kwargs): - logger.warning_once( - "`LlamaLinearScalingRotaryEmbedding` is deprecated an will be removed in v4.46. Please use " - "`LlamaRotaryEmbedding`, which now also does linear scaling (simply pass the model config to __init__)." - ) - kwargs["rope_type"] = "linear" - super().__init__(*args, **kwargs) - - def _set_cos_sin_cache(self, seq_len, device, dtype): - self.max_seq_len_cached = seq_len - t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype) - t = t / self.scaling_factor - - freqs = torch.outer(t, self.inv_freq) - # Different from paper, but it uses a different permutation in order to obtain the same calculation - emb = torch.cat((freqs, freqs), dim=-1) - self.register_buffer("_cos_cached", emb.cos().to(dtype), persistent=False) - self.register_buffer("_sin_cached", emb.sin().to(dtype), persistent=False) - - -class GaudiLlamaDynamicNTKScalingRotaryEmbedding(GaudiLlamaRotaryEmbedding): - def __init__(self, *args, **kwargs): - logger.warning_once( - "`LlamaDynamicNTKScalingRotaryEmbedding` is deprecated an will be removed in v4.46. Please use " - "`LlamaRotaryEmbedding`, which now also does dynamic ntk scaling (simply pass the model config to " - "__init__)." - ) - kwargs["rope_type"] = "dynamic" - super().__init__(*args, **kwargs) - - def _set_cos_sin_cache(self, seq_len, device, dtype): - self.max_seq_len_cached = seq_len - - if seq_len > self.max_position_embeddings: - base = self.base * ( - (self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1) - ) ** (self.dim / (self.dim - 2)) - inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim)) - self.register_buffer("inv_freq", inv_freq, persistent=False) - - t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype) - - freqs = torch.outer(t, self.inv_freq) - # Different from paper, but it uses a different permutation in order to obtain the same calculation - emb = torch.cat((freqs, freqs), dim=-1) - self.register_buffer("_cos_cached", emb.cos().to(dtype), persistent=False) - self.register_buffer("_sin_cached", emb.sin().to(dtype), persistent=False) - - class GaudiLlamaMLP(LlamaMLP): def __init__(self, config): super(LlamaMLP, self).__init__() @@ -415,6 +355,36 @@ def GaudiDistributedAttention(fused_scaled_dot_product_attention, fused_scaled_d return fused_scaled_dot_product_attention +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, + attn_softmax_bf16: bool = False, +): + query_states, key_states, value_states, attention_mask = gaudi_llama_repeat_kv( + query, key, value, attention_mask, module.num_key_value_groups + ) + + attn_weights = module.matmul_qk(query_states, key_states.transpose(-2, -1)) * scaling + if attention_mask is not None: + causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] + attn_weights = attn_weights + causal_mask + + if attn_softmax_bf16: + attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1, dtype=query_states.dtype) + else: + # upcast attention to fp32 + 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=dropout, training=module.training) + attn_output = module.matmul_av(attn_weights, value_states) + + return attn_output, attn_weights + + class GaudiLlamaAttention(LlamaAttention): def __init__(self, config: LlamaConfig, layer_idx: Optional[int] = None): super().__init__(config, layer_idx) @@ -424,6 +394,9 @@ def __init__(self, config: LlamaConfig, layer_idx: Optional[int] = None): self.k_cache = KVCache() self.v_cache = KVCache() + self.rotary_emb = GaudiLlamaRotaryEmbedding(config=config) + self.num_key_value_heads = config.num_key_value_heads + if hasattr(config, "fused_qkv") and config.fused_qkv: self.num_heads = config.num_attention_heads self.head_dim = config.hidden_size // self.num_heads @@ -438,11 +411,10 @@ def __init__(self, config: LlamaConfig, layer_idx: Optional[int] = None): self.k_proj = None self.v_proj = None self.inp_seq_len = -1 - self.norm_factor = 1.0 / math.sqrt(self.head_dim) self.fused_scaled_dot_product_attention = ( ModuleFusedSDPA( FusedSDPA, - scale=self.norm_factor, + scale=self.scaling, attention_dropout=self.attention_dropout, enable_recompute=False, flash_attention_fp8=getattr(config, "flash_attention_fp8", False), @@ -505,13 +477,12 @@ def reorder_kv_cache(self, beam_idx: torch.LongTensor): def pre_attn_forward( self, hidden_states: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, + position_embeddings: Tuple[torch.Tensor, torch.Tensor], + attention_mask: Optional[torch.Tensor], past_key_value: Optional[Cache] = None, - 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.46 + position_ids: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, attn_softmax_bf16: Optional[bool] = False, reuse_cache: Optional[bool] = False, @@ -537,7 +508,9 @@ def pre_attn_forward( - add new arg flash_attention_fast_softmax - add new arg num_virtual_tokens """ - bsz, q_len, _ = hidden_states.size() + input_shape = hidden_states.shape[:-1] + q_len = input_shape[1] + hidden_shape = (*input_shape, -1, self.head_dim) if hasattr(self.config, "fused_qkv") and self.config.fused_qkv: qkv_states = self.qkv_proj(hidden_states) @@ -548,10 +521,10 @@ def pre_attn_forward( value_states = self.v_proj(hidden_states) # use -1 to infer num_heads and num_key_value_heads as they may vary if tensor parallel is used - query_states = query_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + query_states = query_states.view(hidden_shape).transpose(1, 2) # TODO: update when auto mp params is enabled in DeepSpeed (cf. https://github.com/HabanaAI/DeepSpeed/blob/94309c7b5dfc1a69858f5c9f25737b2f81a332a5/deepspeed/module_inject/replace_module.py#L440) - key_states = key_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) - value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + key_states = key_states.view(hidden_shape).transpose(1, 2) + value_states = value_states.view(hidden_shape).transpose(1, 2) kv_seq_len = key_states.shape[-2] if past_key_value is not None: @@ -649,6 +622,7 @@ def pre_attn_forward( self.fused_scaled_dot_product_attention, self.fused_scaled_dot_product_attention_distributed ) if use_flash_attention and FusedSDPA is not None: + attn_weights = None if q_len == 1: # next token attn_output = fused_scaled_dot_product_attention( @@ -698,44 +672,21 @@ def pre_attn_forward( ) else: - query_states, key_states, value_states, attention_mask = gaudi_llama_repeat_kv( - query_states, key_states, value_states, attention_mask, self.num_key_value_groups - ) - - attn_weights = self.matmul_qk(query_states, key_states.transpose(-2, -1)) * self.norm_factor - - if attention_mask is not None: # no matter the length, we just slice it - causal_mask = attention_mask - if cache_position is not None: - causal_mask = attention_mask[:, :, cache_position, : key_states.shape[-2]] - attn_weights = attn_weights + causal_mask - - if attn_softmax_bf16: - attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1, dtype=query_states.dtype) - else: - # upcast attention to fp32 - 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.attention_dropout, training=self.training) - attn_output = self.matmul_av(attn_weights, value_states) - attn_output = attn_output.reshape(bsz, -1, q_len, self.head_dim) - - 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_weights = gaudi_eager_attention_forward( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + attn_softmax_bf16=attn_softmax_bf16, ) attn_output = attn_output.transpose(1, 2).contiguous() - - attn_output = attn_output.reshape(bsz, q_len, -1) - + attn_output = attn_output.reshape(*input_shape, -1).contiguous() attn_output = self.o_proj(attn_output) - if not output_attentions: - attn_weights = None - if not reuse_cache and token_idx is not None and cache_idx is not None and q_len == 1: # Return only past key value shapes and not the tensors during decode phase (q len is 1) # to avoid making past key values as persistent output tensors of HPU graphs. @@ -791,7 +742,6 @@ def __init__( self.o_proj = torch.nn.Linear( self.config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias ) - self.norm_factor = 1.0 / math.sqrt(self.head_dim) self.setup_tp(rank, world_size) def colwise_param_names(self) -> List[str]: @@ -883,7 +833,7 @@ def forward( 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.46 + 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, @@ -936,7 +886,6 @@ def forward( hidden_states = self.post_mlp(hidden_states, residual) outputs = (hidden_states,) - if output_attentions: outputs += (self_attn_weights,) if use_cache: diff --git a/optimum/habana/transformers/models/minicpm/modeling_minicpm.py b/optimum/habana/transformers/models/minicpm/modeling_minicpm.py index 9e7656de22..1f2e4a7ff3 100644 --- a/optimum/habana/transformers/models/minicpm/modeling_minicpm.py +++ b/optimum/habana/transformers/models/minicpm/modeling_minicpm.py @@ -45,7 +45,7 @@ CausalLMOutputWithPast, ) from transformers.modeling_utils import PreTrainedModel -from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS, is_torch_greater_or_equal_than_1_13 +from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS from transformers.utils import ( add_start_docstrings, add_start_docstrings_to_model_forward, @@ -67,9 +67,6 @@ # This makes `_prepare_4d_causal_attention_mask` a leaf function in the FX graph. # It means that the function will not be traced through and simply appear as a node in the graph. if is_torch_fx_available(): - if not is_torch_greater_or_equal_than_1_13: - import torch.fx - _prepare_4d_causal_attention_mask = torch.fx.wrap(_prepare_4d_causal_attention_mask) diff --git a/optimum/habana/transformers/models/mistral/modeling_mistral.py b/optimum/habana/transformers/models/mistral/modeling_mistral.py index 26a8567517..2c5b28b307 100644 --- a/optimum/habana/transformers/models/mistral/modeling_mistral.py +++ b/optimum/habana/transformers/models/mistral/modeling_mistral.py @@ -19,18 +19,16 @@ # limitations under the License. """PyTorch Mistral model.""" -import math import os from typing import List, Optional, Tuple, Union import habana_frameworks.torch.core as htcore import torch -from torch import nn -from torch.nn import CrossEntropyLoss from transformers.cache_utils import Cache, DynamicCache 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,16 +37,13 @@ MistralRMSNorm, apply_rotary_pos_emb, ) +from transformers.processing_utils import Unpack from transformers.utils import logging from ...modeling_attn_mask_utils import ( _gaudi_prepare_4d_causal_attention_mask, ) -from ..llama.modeling_llama import ( - GaudiLlamaDynamicNTKScalingRotaryEmbedding, - GaudiLlamaLinearScalingRotaryEmbedding, - GaudiLlamaRotaryEmbedding, -) +from ..llama.modeling_llama import GaudiLlamaRotaryEmbedding from ..modeling_all_models import KVCache, Matmul, apply_customized_rope_module @@ -141,10 +136,42 @@ def gaudi_mistral_rmsnorm_forward(self, hidden_states): return self.weight * hidden_states.to(input_dtype) +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, + attn_softmax_bf16: bool = False, + **kwargs, +): + bsz, q_len = kwargs["input_shape"] + query_states, key_states, value_states, attention_mask = gaudi_mistral_repeat_kv( + query, key, value, attention_mask, module.num_key_value_groups + ) + + attn_weights = module.matmul_qk(query_states, key_states.transpose(-2, -1)) * scaling + if attention_mask is not None: + causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] + attn_weights = attn_weights + causal_mask + + if attn_softmax_bf16: + attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1, dtype=query_states.dtype) + else: + # upcast attention to fp32 + 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=dropout, training=module.training) + attn_output = module.matmul_av(attn_weights, value_states) + attn_output = attn_output.reshape(bsz, -1, q_len, module.head_dim) + + return attn_output, attn_weights + + class GaudiMistralAttention(MistralAttention): def __init__(self, config: MistralConfig, layer_idx: Optional[int] = None): super().__init__(config, layer_idx) - config.rope_scaling = config.rope_scaling if hasattr(config, "rope_scaling") else None self.config = config self.k_cache = KVCache() self.v_cache = KVCache() @@ -152,38 +179,8 @@ def __init__(self, config: MistralConfig, layer_idx: Optional[int] = None): self.matmul_av = Matmul() self.fused_scaled_dot_product_attention = ModuleFusedSDPA(FusedSDPA) if FusedSDPA else None self.inp_seq_len = -1 - self._init_rope() - self.norm_factor = 1.0 / math.sqrt(self.head_dim) - - def _init_rope(self): - """ - Copied from: https://github.com/huggingface/transformers/blob/v4.38.2/src/transformers/models/llama/modeling_llama.py#L294 - """ - if self.config.rope_scaling is None: - self.rotary_emb = GaudiLlamaRotaryEmbedding( - self.head_dim, - max_position_embeddings=self.max_position_embeddings, - base=self.rope_theta, - ) - else: - scaling_type = self.config.rope_scaling["type"] - scaling_factor = self.config.rope_scaling["factor"] - if scaling_type == "linear": - self.rotary_emb = GaudiLlamaLinearScalingRotaryEmbedding( - self.head_dim, - max_position_embeddings=self.max_position_embeddings, - scaling_factor=scaling_factor, - base=self.rope_theta, - ) - elif scaling_type == "dynamic": - self.rotary_emb = GaudiLlamaDynamicNTKScalingRotaryEmbedding( - self.head_dim, - max_position_embeddings=self.max_position_embeddings, - scaling_factor=scaling_factor, - base=self.rope_theta, - ) - else: - raise ValueError(f"Unknown RoPE scaling type {scaling_type}") + self.rotary_emb = GaudiLlamaRotaryEmbedding(config=config) + self.num_key_value_heads = config.num_key_value_heads def allocate_kv_cache(self, batch_size, max_seq_len, inp_seq_len): cache_shape = (batch_size, self.num_key_value_heads, max_seq_len, self.head_dim) @@ -217,10 +214,9 @@ def reorder_kv_cache(self, beam_idx: torch.LongTensor): def forward( self, hidden_states: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, + position_embeddings: Tuple[torch.Tensor, torch.Tensor], + attention_mask: Optional[torch.Tensor], past_key_value: Optional[Cache] = None, - output_attentions: bool = False, use_cache: bool = False, cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, @@ -239,15 +235,13 @@ def forward( - add new args reuse_cache - add new args cache_idx """ - bsz, q_len, _ = hidden_states.size() - - query_states = self.q_proj(hidden_states) - key_states = self.k_proj(hidden_states) - value_states = self.v_proj(hidden_states) + input_shape = hidden_states.shape[:-1] + q_len = input_shape[1] + hidden_shape = (*input_shape, -1, self.head_dim) - query_states = query_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) - key_states = key_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) - value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) + key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) + value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) kv_seq_len = key_states.shape[-2] if past_key_value is not None: @@ -268,7 +262,7 @@ def forward( kv_seq_len += kv_shape cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len) query_states, key_states = apply_customized_rope( - query_states, key_states, cos, sin, position_ids, self.training + query_states, key_states, cos, sin, kwargs["position_ids"], self.training ) if use_cache: @@ -301,6 +295,7 @@ def forward( import habana_frameworks.torch.hpu as ht if FusedSDPA and use_flash_attention: + attn_weights = None if q_len == 1: # next token use_recompute = True if os.getenv("QUANT_CONFIG", "") else False @@ -323,39 +318,23 @@ def forward( ) else: - # repeat k/v heads if n_kv_heads < n_heads - query_states, key_states, value_states, attention_mask = gaudi_mistral_repeat_kv( - query_states, key_states, value_states, attention_mask, self.num_key_value_groups - ) - attn_weights = self.matmul_qk(query_states, key_states.transpose(-2, -1)) * self.norm_factor - - 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 - - if attn_softmax_bf16: - attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=query_states.dtype) - else: - # upcast attention to fp32 - attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) - attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training) - attn_output = self.matmul_av(attn_weights, value_states) - attn_output = attn_output.reshape(bsz, -1, q_len, self.head_dim) - - 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_weights = gaudi_eager_attention_forward( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + sliding_window=getattr(self.config, "sliding_window", None), # main diff with Llama + attn_softmax_bf16=attn_softmax_bf16, + input_shape=input_shape, ) attn_output = attn_output.transpose(1, 2).contiguous() - - attn_output = attn_output.view(bsz, q_len, -1) + attn_output = attn_output.reshape(*input_shape, -1).contiguous() attn_output = self.o_proj(attn_output) - if not output_attentions: - attn_weights = None - return attn_output, attn_weights, past_key_value @@ -388,6 +367,7 @@ def forward( 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 token_idx: Optional[torch.Tensor] = None, reuse_cache: Optional[bool] = False, cache_idx: Optional[int] = None, @@ -415,6 +395,7 @@ def forward( output_attentions=output_attentions, use_cache=use_cache, cache_position=cache_position, + position_embeddings=position_embeddings, token_idx=token_idx, reuse_cache=reuse_cache, cache_idx=cache_idx, @@ -490,7 +471,6 @@ def forward( 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") elif input_ids is not None: @@ -502,7 +482,7 @@ def forward( 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=True` is incompatible with gradient checkpointing. Setting `use_cache=False`." ) use_cache = False @@ -549,7 +529,7 @@ def forward( if lazy_mode: htcore.mark_step() - for layer_idx, decoder_layer in enumerate(self.layers): + for layer_idx, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]): if layer_idx == len(self.layers) // 2 or ( lazy_mode and not self.training @@ -660,6 +640,7 @@ def forward( use_flash_attention: Optional[bool] = False, flash_attention_recompute: Optional[bool] = False, flash_attention_causal_mask: Optional[bool] = False, + **kwargs: Unpack[KwargsForCausalLM], ) -> Union[Tuple, CausalLMOutputWithPast]: """ Inherits from MistralForCausalLM: https://github.com/huggingface/transformers/blob/v4.34.1/src/transformers/models/mistral/modeling_mistral.py @@ -710,18 +691,7 @@ def forward( 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 so that tokens < n predict n - shift_logits = logits[..., :-1, :].contiguous() - shift_labels = labels[..., 1:].contiguous() - # Flatten the tokens - loss_fct = CrossEntropyLoss() - shift_logits = shift_logits.view(-1, self.config.vocab_size) - shift_labels = shift_labels.view(-1) - # Enable model parallelism - shift_labels = shift_labels.to(shift_logits.device) - loss = loss_fct(shift_logits, shift_labels) + loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs) if not return_dict: output = (logits,) + outputs[1:] diff --git a/optimum/habana/transformers/models/mixtral/modeling_mixtral.py b/optimum/habana/transformers/models/mixtral/modeling_mixtral.py index 6ae2fda6d9..e009d7f8a9 100644 --- a/optimum/habana/transformers/models/mixtral/modeling_mixtral.py +++ b/optimum/habana/transformers/models/mixtral/modeling_mixtral.py @@ -28,15 +28,20 @@ import habana_frameworks.torch.core as htcore 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.modeling_attn_mask_utils import ( _prepare_4d_causal_attention_mask, _prepare_4d_causal_attention_mask_for_sdpa, ) -from transformers.modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast +from transformers.modeling_outputs import ( + BaseModelOutputWithPast, + CausalLMOutputWithPast, + MoeCausalLMOutputWithPast, + MoeModelOutputWithPast, +) from transformers.models.mixtral.modeling_mixtral import ( + KwargsForCausalLM, MixtralAttention, MixtralDecoderLayer, MixtralForCausalLM, @@ -44,13 +49,10 @@ apply_rotary_pos_emb, load_balancing_loss_func, ) +from transformers.processing_utils import Unpack from transformers.utils import logging -from ..llama.modeling_llama import ( - GaudiLlamaDynamicNTKScalingRotaryEmbedding, - GaudiLlamaLinearScalingRotaryEmbedding, - GaudiLlamaRotaryEmbedding, -) +from ..llama.modeling_llama import GaudiLlamaRotaryEmbedding from ..modeling_all_models import KVCache, apply_customized_rope_module from .configuration_mixtral import MixtralConfig @@ -173,48 +175,45 @@ def forward(q, k, v, mask, causal, q_block_size): return attn_output +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, +): + bsz, q_len = kwargs["input_shape"] + query_states, key_states, value_states, attention_mask = gaudi_mixtral_repeat_kv( + query, key, value, attention_mask, module.num_key_value_groups + ) + + attn_weights = torch.matmul(query_states, key_states.transpose(-2, -1)) * scaling + if attention_mask is not None: + causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] + attn_weights = attn_weights + causal_mask + + attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = torch.nn.functional.dropout(attn_weights, p=dropout, training=module.training) + attn_output = torch.matmul(attn_weights, value_states) + + attn_output = attn_output.reshape(bsz, -1, q_len, module.head_dim).contiguous() + + return attn_output, attn_weights + + class GaudiMixtralAttention(MixtralAttention): def __init__(self, config: MixtralConfig, layer_idx: Optional[int] = None): super().__init__(config, layer_idx) - config.rope_scaling = config.rope_scaling if hasattr(config, "rope_scaling") else None self.config = config - self._init_rope() self.k_cache = KVCache() self.v_cache = KVCache() self.inp_seq_len = -1 - self.norm_factor = 1.0 / math.sqrt(self.head_dim) + self.rotary_emb = GaudiLlamaRotaryEmbedding(config=config) self.block_size = 1024 - def _init_rope(self): - """ - Copied from: https://github.com/huggingface/transformers/blob/v4.38.2/src/transformers/models/llama/modeling_llama.py#L294 - """ - if self.config.rope_scaling is None: - self.rotary_emb = GaudiLlamaRotaryEmbedding( - self.head_dim, - max_position_embeddings=self.max_position_embeddings, - base=self.rope_theta, - ) - else: - scaling_type = self.config.rope_scaling["type"] - scaling_factor = self.config.rope_scaling["factor"] - if scaling_type == "linear": - self.rotary_emb = GaudiLlamaLinearScalingRotaryEmbedding( - self.head_dim, - max_position_embeddings=self.max_position_embeddings, - scaling_factor=scaling_factor, - base=self.rope_theta, - ) - elif scaling_type == "dynamic": - self.rotary_emb = GaudiLlamaDynamicNTKScalingRotaryEmbedding( - self.head_dim, - max_position_embeddings=self.max_position_embeddings, - scaling_factor=scaling_factor, - base=self.rope_theta, - ) - else: - raise ValueError(f"Unknown RoPE scaling type {scaling_type}") - def allocate_kv_cache(self, batch_size, max_seq_len, inp_seq_len): cache_shape = (batch_size, self.num_key_value_heads, max_seq_len, self.head_dim) device = self.k_proj.weight.device @@ -225,16 +224,16 @@ def allocate_kv_cache(self, batch_size, max_seq_len, inp_seq_len): def forward( self, hidden_states: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, + position_embeddings: Tuple[torch.Tensor, torch.Tensor], + attention_mask: Optional[torch.Tensor], past_key_value: Optional[Cache] = None, - output_attentions: bool = False, use_cache: bool = False, 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, + **kwargs, ) -> 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 @@ -245,15 +244,13 @@ def forward( - add new args flash_attention_recompute - add new args cache_idx """ - bsz, q_len, _ = hidden_states.size() - - query_states = self.q_proj(hidden_states) - key_states = self.k_proj(hidden_states) - value_states = self.v_proj(hidden_states) + input_shape = hidden_states.shape[:-1] + q_len = input_shape[1] + hidden_shape = (*input_shape, -1, self.head_dim) - query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) - key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) - value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) + key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) + value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) kv_seq_len = key_states.shape[-2] if past_key_value is not None: @@ -273,9 +270,10 @@ def forward( kv_seq_len = past_key_value[0][-2] else: kv_seq_len = past_key_value[0].shape[-2] + cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len) query_states, key_states = apply_customized_rope( - query_states, key_states, cos, sin, position_ids, self.training + query_states, key_states, cos, sin, kwargs["position_ids"], self.training ) if use_cache: @@ -305,6 +303,7 @@ def forward( past_key_value = None if FusedSDPA: + attn_weights = None if query_states.dtype != key_states.dtype: key_states = key_states.type(query_states.dtype) value_states = value_states.type(query_states.dtype) @@ -328,31 +327,22 @@ def forward( query_states, key_states, value_states, attention_mask, 0.0, False, None ) else: - query_states, key_states, value_states, attention_mask = gaudi_mixtral_repeat_kv( - query_states, key_states, value_states, attention_mask, self.num_key_value_groups + attn_output, attn_weights = gaudi_eager_attention_forward( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + sliding_window=getattr(self.config, "sliding_window", None), # main diff with Llama + input_shape=input_shape, ) - attn_weights = torch.matmul(query_states, key_states.transpose(-2, -1)) * self.norm_factor - - if attention_mask is not None: - attention_mask = attention_mask.unsqueeze(2) - attn_weights = attn_weights + attention_mask - - # upcast attention to fp32 - attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) - attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training) - attn_output = torch.matmul(attn_weights, value_states) - - attn_output = attn_output.reshape(bsz, self.num_heads, q_len, self.head_dim).contiguous() - attn_output = attn_output.transpose(1, 2).contiguous() - attn_output = attn_output.reshape(bsz, q_len, -1) - + attn_output = attn_output.reshape(*input_shape, -1).contiguous() attn_output = self.o_proj(attn_output) - if not output_attentions or FusedSDPA: - attn_weights = None - return attn_output, attn_weights, past_key_value @@ -371,6 +361,8 @@ def gaudi_mixtral_block_sparse_moe_forward(self, hidden_states: torch.Tensor) -> - optimize expert forward, remove dynamic control and dynamic shape """ batch_size, sequence_length, hidden_dim = hidden_states.shape + if self.training and self.jitter_noise > 0: + hidden_states *= torch.empty_like(hidden_states).uniform_(1.0 - self.jitter_noise, 1.0 + self.jitter_noise) hidden_states = hidden_states.view(-1, hidden_dim) # router_logits: (batch * sequence_length, n_experts) router_logits = self.gate(hidden_states) @@ -486,6 +478,7 @@ def forward( 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 token_idx: Optional[torch.Tensor] = None, reuse_cache: Optional[bool] = False, flash_attention_recompute: Optional[bool] = False, @@ -507,6 +500,7 @@ def forward( # Self Attention hidden_states, self_attn_weights, present_key_value = self.self_attn( hidden_states=hidden_states, + position_embeddings=position_embeddings, attention_mask=attention_mask, position_ids=position_ids, past_key_value=past_key_value, @@ -565,7 +559,7 @@ def forward( reuse_cache: Optional[bool] = False, flash_attention_recompute: Optional[bool] = False, cache_idx: int = None, - ) -> Union[Tuple, MoeModelOutputWithPast]: + ) -> Union[Tuple, BaseModelOutputWithPast]: """ 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: @@ -769,8 +763,8 @@ def forward( reuse_cache: Optional[bool] = None, flash_attention_recompute: Optional[bool] = False, cache_idx: int = None, - **loss_kwargs, - ) -> Union[Tuple, MoeCausalLMOutputWithPast]: + **kwargs: Unpack[KwargsForCausalLM], + ) -> Union[Tuple, CausalLMOutputWithPast]: 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 @@ -806,7 +800,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: diff --git a/optimum/habana/transformers/models/mllama/modeling_mllama.py b/optimum/habana/transformers/models/mllama/modeling_mllama.py index 9ecbff58bd..6d2d2a08fb 100644 --- a/optimum/habana/transformers/models/mllama/modeling_mllama.py +++ b/optimum/habana/transformers/models/mllama/modeling_mllama.py @@ -785,7 +785,7 @@ def _update_causal_mask( - add support if past_key_value is not Cache """ if self.config._attn_implementation == "flash_attention_2": - if attention_mask is not None and 0.0 in attention_mask: + if attention_mask is not None and (attention_mask == 0.0).any(): return attention_mask return None diff --git a/optimum/habana/transformers/models/paligemma/modeling_paligemma.py b/optimum/habana/transformers/models/paligemma/modeling_paligemma.py index 3b2487772f..1d2db48d41 100644 --- a/optimum/habana/transformers/models/paligemma/modeling_paligemma.py +++ b/optimum/habana/transformers/models/paligemma/modeling_paligemma.py @@ -103,7 +103,7 @@ def forward( # mask out pad-token-ids in labels for BC if labels is not None and self.pad_token_id in labels: logger.warning_once( - "`labels` contains `pad_token_id` which will be masked with `config.ignore_index`. ", + "`labels` contains `pad_token_id` which will be masked with `config.ignore_index`. " "You have to mask out `pad_token_id` when preparing `labels`, this behavior will be removed in v.4.46.", ) labels = torch.where(input_ids == self.pad_token_id, self.config.ignore_index, labels) diff --git a/optimum/habana/transformers/models/persimmon/modeling_persimmon.py b/optimum/habana/transformers/models/persimmon/modeling_persimmon.py index d76c87b2f6..3e56f3c9e2 100644 --- a/optimum/habana/transformers/models/persimmon/modeling_persimmon.py +++ b/optimum/habana/transformers/models/persimmon/modeling_persimmon.py @@ -40,6 +40,7 @@ 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 token_idx: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: """ @@ -167,6 +168,7 @@ def forward( 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 token_idx: Optional[torch.Tensor] = None, ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: """ @@ -187,6 +189,7 @@ def forward( output_attentions=output_attentions, use_cache=use_cache, cache_position=cache_position, + position_embeddings=position_embeddings, token_idx=token_idx, ) hidden_states = residual + hidden_states diff --git a/optimum/habana/transformers/models/phi/modeling_phi.py b/optimum/habana/transformers/models/phi/modeling_phi.py index ab200d2332..c86e7563ac 100644 --- a/optimum/habana/transformers/models/phi/modeling_phi.py +++ b/optimum/habana/transformers/models/phi/modeling_phi.py @@ -19,21 +19,21 @@ # limitations under the License. """PyTorch Phi model.""" -import math from typing import List, Optional, Tuple, Union import torch -from torch import nn from transformers.cache_utils import Cache, DynamicCache 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, PhiModel, apply_rotary_pos_emb, ) +from transformers.processing_utils import Unpack from transformers.utils import logging from ...modeling_attn_mask_utils import ( @@ -80,6 +80,34 @@ def gaudi_phi_repeat_kv( return query_states, key_states, value_states, attention_mask +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, +): + bsz, q_len = kwargs["input_shape"] + query_states, key_states, value_states, attention_mask = gaudi_phi_repeat_kv( + query, key, value, attention_mask, module.num_key_value_groups + ) + + attn_weights = module.matmul_qk(query_states, key_states.transpose(2, 3)) * scaling + if attention_mask is not None: + causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] + attn_weights = attn_weights + causal_mask + + attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = torch.nn.functional.dropout(attn_weights, p=dropout, training=module.training) + attn_output = module.matmul_av(attn_weights, value_states) + attn_output = attn_output.reshape(bsz, -1, q_len, module.head_dim) + + return attn_output, attn_weights + + class GaudiPhiAttention(PhiAttention): def __init__(self, config: PhiConfig, layer_idx: Optional[int] = None): super().__init__(config, layer_idx) @@ -100,10 +128,9 @@ def allocate_kv_cache(self, batch_size, max_seq_len, inp_seq_len): def forward( self, hidden_states: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, + position_embeddings: Tuple[torch.Tensor, torch.Tensor], + attention_mask: Optional[torch.Tensor], past_key_value: Optional[Cache] = None, - output_attentions: bool = False, use_cache: bool = False, cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, @@ -119,20 +146,18 @@ def forward( - add new args reuse_cache - add new args cache_idx """ - bsz, q_len, _ = hidden_states.size() + input_shape = hidden_states.shape[:-1] + q_len = input_shape[1] + hidden_shape = (*input_shape, -1, self.head_dim) - query_states = self.q_proj(hidden_states) - key_states = self.k_proj(hidden_states) - value_states = self.v_proj(hidden_states) + query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) + key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) + value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) if self.qk_layernorm: query_states = self.q_layernorm(query_states) key_states = self.k_layernorm(key_states) - query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) - key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) - value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) - kv_seq_len = key_states.shape[-2] if past_key_value is not None: if self.layer_idx is None: @@ -163,7 +188,9 @@ def forward( key_states[..., self.rotary_ndims :], ) # [batch_size, seq_length, num_heads, head_dim // config.partial_rotary_factor] - query_rot, key_rot = apply_rotary_pos_emb(query_rot, key_rot, cos[position_ids], sin[position_ids]) + query_rot, key_rot = apply_rotary_pos_emb( + query_rot, key_rot, cos[kwargs["position_ids"]], sin[kwargs["position_ids"]] + ) # [batch_size, seq_length, num_heads, head_dim] query_states = torch.cat((query_rot, query_pass), dim=-1) @@ -196,54 +223,21 @@ def forward( else: past_key_value = None - query_states, key_states, value_states, attention_mask = gaudi_phi_repeat_kv( - query_states, key_states, value_states, attention_mask, self.num_key_value_groups + attn_output, attn_weights = gaudi_eager_attention_forward( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + input_shape=input_shape, ) - # Queries and keys upcast to fp32 is required by Phi-2 to avoid overflow - attn_weights = self.matmul_qk( - query_states.to(torch.float32), key_states.to(torch.float32).transpose(2, 3) - ) / math.sqrt(self.head_dim) - - if attn_weights.size() not in [ - (bsz, self.num_heads, q_len, kv_seq_len), - (bsz, self.num_key_value_heads, self.num_key_value_groups, q_len, kv_seq_len), - ]: - raise ValueError( - f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)} or" - f" {(bsz, self.num_key_value_heads, self.num_key_value_groups, q_len, kv_seq_len)}, but is" - f" {attn_weights.size()}" - ) - - if attention_mask is not None: - if attention_mask.size() not in [(bsz, 1, q_len, kv_seq_len), (bsz, 1, 1, q_len, kv_seq_len)]: - raise ValueError( - f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)} or {(bsz, 1, 1, q_len, kv_seq_len)}," - f" but is {attention_mask.size()}" - ) - attn_weights = attn_weights + attention_mask - - # upcast attention to fp32 - attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(value_states.dtype) - attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training) - - attn_output = self.matmul_av(attn_weights, value_states) - attn_output = attn_output.reshape(bsz, -1, q_len, self.head_dim) - - 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_output.reshape(*input_shape, -1).contiguous() attn_output = self.dense(attn_output) - if not output_attentions: - attn_weights = None - return attn_output, attn_weights, past_key_value @@ -263,10 +257,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, output_attentions: Optional[bool] = False, use_cache: Optional[bool] = False, - past_key_value: Optional[Tuple[torch.Tensor]] = None, cache_position: Optional[torch.LongTensor] = None, + 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, @@ -293,6 +288,7 @@ def forward( output_attentions=output_attentions, use_cache=use_cache, cache_position=cache_position, + position_embeddings=position_embeddings, token_idx=token_idx, reuse_cache=reuse_cache, cache_idx=cache_idx, @@ -360,7 +356,7 @@ def forward( 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=True` is incompatible with gradient checkpointing. Setting `use_cache=False`." ) use_cache = False @@ -400,7 +396,7 @@ def forward( attention_mask, (batch_size, seq_length), inputs_embeds, past_seen_tokens ) - inputs_embeds = self.embed_dropout(inputs_embeds) + inputs_embeds = self.embed_dropout(inputs_embeds) # diff with Llama hidden_states = inputs_embeds # decoder layers @@ -408,7 +404,7 @@ def forward( all_self_attns = () if output_attentions else None next_decoder_cache = () if not use_new_cache else None - for layer_idx, decoder_layer in enumerate(self.layers): + for layer_idx, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]): if output_hidden_states: all_hidden_states += (hidden_states,) @@ -418,9 +414,9 @@ def forward( hidden_states, attention_mask, position_ids, + None if past_key_values is None else past_key_values[layer_idx], output_attentions, use_cache, - None if past_key_values is None else past_key_values[layer_idx], cache_position, None, ) @@ -446,7 +442,7 @@ def forward( if output_attentions: all_self_attns += (layer_outputs[1],) - hidden_states = self.final_layernorm(hidden_states) + hidden_states = self.final_layernorm(hidden_states) # diff with Llama # add hidden states from the last decoder layer if output_hidden_states: @@ -490,7 +486,7 @@ def forward( reuse_cache: Optional[bool] = False, trim_logits: Optional[bool] = False, cache_idx: Optional[int] = None, - **loss_kwargs, + **kwargs: Unpack[KwargsForCausalLM], ) -> Union[Tuple, CausalLMOutputWithPast]: """ Inherits from PhiForCausalLM: https://github.com/huggingface/transformers/blob/v4.37.1/src/transformers/models/phi/modeling_phi.py @@ -499,7 +495,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 @@ -535,7 +530,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=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs) if not return_dict: output = (logits,) + outputs[1:] diff --git a/optimum/habana/transformers/models/qwen2/modeling_qwen2.py b/optimum/habana/transformers/models/qwen2/modeling_qwen2.py index e646188e39..5573aa19a6 100644 --- a/optimum/habana/transformers/models/qwen2/modeling_qwen2.py +++ b/optimum/habana/transformers/models/qwen2/modeling_qwen2.py @@ -16,15 +16,14 @@ # Copyright (C) 2022-2024 Habana Labs, Ltd. an Intel Company ############################################################################### -import math from typing import List, Optional, Tuple, Union import torch -import torch.nn as nn from transformers.cache_utils import Cache, DynamicCache, StaticCache from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from transformers.models.qwen2.configuration_qwen2 import Qwen2Config from transformers.models.qwen2.modeling_qwen2 import ( + KwargsForCausalLM, Qwen2Attention, Qwen2DecoderLayer, Qwen2ForCausalLM, @@ -34,6 +33,7 @@ apply_rotary_pos_emb, logger, ) +from transformers.processing_utils import Unpack from ...modeling_attn_mask_utils import ( _gaudi_prepare_4d_causal_attention_mask, @@ -166,6 +166,41 @@ def forward( ) +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, + attn_softmax_bf16: bool = False, + **kwargs, +): + bsz, q_len = kwargs["input_shape"] + query_states, key_states, value_states, attention_mask = gaudi_qwen2_repeat_kv( + query, key, value, attention_mask, module.num_key_value_groups + ) + + query_states = query_states * scaling + attn_weights = module.matmul_qk(query_states, key_states.transpose(-2, -1)).float() + htcore.mark_step() + if attention_mask is not None: + causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] + attn_weights = attn_weights + causal_mask + + if attn_softmax_bf16: + attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1, dtype=query_states.dtype) + else: + # upcast attention to fp32 + 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=dropout, training=module.training) + attn_output = module.matmul_av(attn_weights, value_states) + attn_output = attn_output.reshape(bsz, -1, q_len, module.head_dim) + + return attn_output, attn_weights + + class GaudiQwen2Attention(Qwen2Attention): def __init__(self, config: Qwen2Config, layer_idx: Optional[int] = None): super().__init__(config, layer_idx) @@ -176,14 +211,13 @@ def __init__(self, config: Qwen2Config, layer_idx: Optional[int] = None): self.v_cache = KVCache() self.inp_seq_len = -1 - self.norm_factor = 1.0 / math.sqrt(self.head_dim) self.rotary_emb = GaudiRotaryEmbedding(config=self.config) self.fused_scaled_dot_product_attention = ( ModuleFusedSDPA( FusedSDPA, - scale=self.norm_factor, + scale=self.scaling, attention_dropout=self.attention_dropout, enable_recompute=False, flash_attention_fp8=getattr(config, "flash_attention_fp8", False), @@ -237,10 +271,9 @@ def reorder_kv_cache(self, beam_idx: torch.LongTensor): def pre_attn_forward( self, hidden_states: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, + position_embeddings: Tuple[torch.Tensor, torch.Tensor], + attention_mask: Optional[torch.Tensor], past_key_value: Optional[Cache] = None, - output_attentions: bool = False, use_cache: bool = False, cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, @@ -267,15 +300,13 @@ def pre_attn_forward( - add new arg flash_attention_fast_softmax - add new arg num_virtual_tokens """ - bsz, q_len, _ = hidden_states.size() + input_shape = hidden_states.shape[:-1] + q_len = input_shape[1] + hidden_shape = (*input_shape, -1, self.head_dim) - query_states = self.q_proj(hidden_states) - key_states = self.k_proj(hidden_states) - value_states = self.v_proj(hidden_states) - - query_states = query_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) - key_states = key_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) - value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) + key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) + value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) kv_seq_len = key_states.shape[-2] if past_key_value is not None: @@ -295,7 +326,7 @@ def pre_attn_forward( cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len) query_states, key_states = apply_customized_rope( - query_states, key_states, cos, sin, position_ids, self.training + query_states, key_states, cos, sin, kwargs["position_ids"], self.training ) if use_cache: @@ -343,7 +374,16 @@ def pre_attn_forward( else: past_key_value = None + 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 if q_len == 1: # next token attn_output = self.fused_scaled_dot_product_attention( @@ -392,46 +432,23 @@ def pre_attn_forward( ) else: - query_states, key_states, value_states, attention_mask = gaudi_qwen2_repeat_kv( - query_states, key_states, value_states, attention_mask, self.num_key_value_groups - ) - - query_states = query_states * self.norm_factor - attn_weights = self.matmul_qk(query_states, key_states.transpose(-2, -1)).float() - htcore.mark_step() - - if attention_mask is not None: # no matter the length, we just slice it - causal_mask = attention_mask - if cache_position is not None: - causal_mask = attention_mask[:, :, cache_position, : key_states.shape[-2]] - attn_weights = attn_weights + causal_mask.float() - - if attn_softmax_bf16: - attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1, dtype=query_states.dtype) - else: - # upcast attention to fp32 - 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.attention_dropout, training=self.training) - attn_output = self.matmul_av(attn_weights, value_states) - attn_output = attn_output.reshape(bsz, -1, q_len, self.head_dim) - - 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_weights = gaudi_eager_attention_forward( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + sliding_window=sliding_window, # main diff with Llama + attn_softmax_bf16=attn_softmax_bf16, + input_shape=input_shape, ) attn_output = attn_output.transpose(1, 2).contiguous() - - attn_output = attn_output.reshape(bsz, q_len, -1) - + attn_output = attn_output.reshape(*input_shape, -1).contiguous() attn_output = self.o_proj(attn_output) - if not output_attentions: - attn_weights = None - if not reuse_cache and token_idx is not None and cache_idx is not None and q_len == 1: # Return only past key value shapes and not the tensors during decode phase (q len is 1) # to avoid making past key values as persistent output tensors of HPU graphs. @@ -478,6 +495,7 @@ def forward( 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 token_idx: Optional[torch.Tensor] = None, attn_softmax_bf16: Optional[bool] = False, reuse_cache: Optional[bool] = False, @@ -500,6 +518,7 @@ def forward( output_attentions=output_attentions, use_cache=use_cache, cache_position=cache_position, + position_embeddings=position_embeddings, token_idx=token_idx, attn_softmax_bf16=attn_softmax_bf16, reuse_cache=reuse_cache, @@ -536,6 +555,7 @@ def pre_attn( 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, token_idx: Optional[torch.Tensor] = None, attn_softmax_bf16: Optional[bool] = False, reuse_cache: Optional[bool] = False, @@ -557,6 +577,7 @@ def pre_attn( output_attentions=output_attentions, use_cache=use_cache, cache_position=cache_position, + position_embeddings=position_embeddings, token_idx=token_idx, attn_softmax_bf16=attn_softmax_bf16, reuse_cache=reuse_cache, @@ -609,11 +630,10 @@ def __init__(self, config: Qwen2Config): self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size - self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) - self.layers = nn.ModuleList( + self.embed_tokens = torch.nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = torch.nn.ModuleList( [GaudiQwen2DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] ) - self._attn_implementation = config._attn_implementation self.norm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.gradient_checkpointing = False @@ -673,12 +693,11 @@ def forward( else: raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds") - 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 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) @@ -740,7 +759,7 @@ def forward( if lazy_mode: htcore.mark_step() - for layer_idx, decoder_layer in enumerate(self.layers): + for layer_idx, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]): if ( lazy_mode and not self.training @@ -857,7 +876,7 @@ def forward( cache_idx: int = None, lazy_mode: Optional[bool] = True, num_virtual_tokens: int = None, - **loss_kwargs, + **kwargs: Unpack[KwargsForCausalLM], ) -> Union[Tuple, CausalLMOutputWithPast]: output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = ( @@ -906,7 +925,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=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs) if not return_dict: output = (logits,) + outputs[1:] 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 0dc677d9bd..861a30dff4 100755 --- a/optimum/habana/transformers/models/qwen2_moe/modeling_qwen2_moe.py +++ b/optimum/habana/transformers/models/qwen2_moe/modeling_qwen2_moe.py @@ -321,6 +321,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 token_idx: Optional[torch.Tensor] = None, attn_softmax_bf16: Optional[bool] = False, reuse_cache: Optional[bool] = False, @@ -635,6 +636,7 @@ def forward( 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 token_idx: Optional[torch.Tensor] = None, attn_softmax_bf16: Optional[bool] = False, reuse_cache: Optional[bool] = False, @@ -707,6 +709,7 @@ def pre_attn( 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 token_idx: Optional[torch.Tensor] = None, attn_softmax_bf16: Optional[bool] = False, reuse_cache: Optional[bool] = False, @@ -728,6 +731,7 @@ def pre_attn( output_attentions=output_attentions, use_cache=use_cache, cache_position=cache_position, + position_embeddings=position_embeddings, token_idx=token_idx, attn_softmax_bf16=attn_softmax_bf16, reuse_cache=reuse_cache, diff --git a/optimum/habana/transformers/models/stablelm/modeling_stablelm.py b/optimum/habana/transformers/models/stablelm/modeling_stablelm.py index f017f38b87..97a78077d7 100644 --- a/optimum/habana/transformers/models/stablelm/modeling_stablelm.py +++ b/optimum/habana/transformers/models/stablelm/modeling_stablelm.py @@ -41,6 +41,7 @@ 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 token_idx: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: """ @@ -177,6 +178,7 @@ def forward( 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 token_idx: Optional[torch.Tensor] = None, ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: """ diff --git a/optimum/habana/transformers/models/starcoder2/modeling_starcoder2.py b/optimum/habana/transformers/models/starcoder2/modeling_starcoder2.py index a5df50b9c3..42f4cd5e9a 100644 --- a/optimum/habana/transformers/models/starcoder2/modeling_starcoder2.py +++ b/optimum/habana/transformers/models/starcoder2/modeling_starcoder2.py @@ -27,6 +27,7 @@ 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, @@ -34,6 +35,7 @@ Starcoder2Model, apply_rotary_pos_emb, ) +from transformers.processing_utils import Unpack from transformers.utils import logging from ...modeling_attn_mask_utils import ( @@ -106,6 +108,39 @@ def gaudi_starcoder2_repeat_kv( return query_states, key_states, value_states, attention_mask +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, + attn_softmax_bf16: bool = False, + **kwargs, +): + bsz, q_len = kwargs["input_shape"] + query_states, key_states, value_states, attention_mask = gaudi_starcoder2_repeat_kv( + query, key, value, attention_mask, module.num_key_value_groups + ) + + attn_weights = module.matmul_qk(query_states, key_states.transpose(-2, -1)) * scaling + if attention_mask is not None: + causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] + attn_weights = attn_weights + causal_mask + + if attn_softmax_bf16: + attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1, dtype=query_states.dtype) + else: + # upcast attention to fp32 + 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=dropout, training=module.training) + attn_output = module.matmul_av(attn_weights, value_states) + attn_output = attn_output.reshape(bsz, -1, q_len, module.head_dim) + + return attn_output, attn_weights + + class GaudiStarcoder2Attention(Starcoder2Attention): def __init__(self, config: Starcoder2Config, layer_idx: Optional[int] = None): super().__init__(config, layer_idx) @@ -115,7 +150,6 @@ def __init__(self, config: Starcoder2Config, layer_idx: Optional[int] = None): self.k_cache = KVCache() self.v_cache = KVCache() self.inp_seq_len = -1 - self.norm_factor = 1.0 / math.sqrt(self.head_dim) self.block_size = 4096 self.rotary_emb = GaudiRotaryEmbedding(config=self.config) @@ -177,10 +211,9 @@ def gaudi_flash_attn_v1(self, query_layer, key_layer, value_layer, attention_mas def pre_attn_forward( self, hidden_states: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, + position_embeddings: Tuple[torch.Tensor, torch.Tensor], + attention_mask: Optional[torch.Tensor], past_key_value: Optional[Cache] = None, - output_attentions: bool = False, use_cache: bool = False, cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, @@ -201,15 +234,13 @@ def pre_attn_forward( - add new args use_flash_attention - add new arg flash_attention_recompute """ - bsz, q_len, _ = hidden_states.size() - - query_states = self.q_proj(hidden_states) - key_states = self.k_proj(hidden_states) - value_states = self.v_proj(hidden_states) + input_shape = hidden_states.shape[:-1] + q_len = input_shape[1] + hidden_shape = (*input_shape, -1, self.head_dim) - query_states = query_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) - key_states = key_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) - value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) + key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) + value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) kv_seq_len = key_states.shape[-2] if past_key_value is not None: @@ -226,7 +257,7 @@ def pre_attn_forward( cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len) query_states, key_states = apply_customized_rope( - query_states, key_states, cos, sin, position_ids, self.training + query_states, key_states, cos, sin, kwargs["position_ids"], self.training ) if use_cache: @@ -257,6 +288,7 @@ def pre_attn_forward( past_key_value = None if use_flash_attention and FusedSDPA: + attn_weights = None import habana_frameworks.torch.hpu as ht if q_len == 1: @@ -285,42 +317,22 @@ def pre_attn_forward( ) else: - query_states, key_states, value_states, attention_mask = gaudi_starcoder2_repeat_kv( - query_states, key_states, value_states, attention_mask, self.num_key_value_groups - ) - - attn_weights = self.matmul_qk(query_states, key_states.transpose(-2, -1)) * self.norm_factor - - 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 - - if attn_softmax_bf16: - attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1, dtype=query_states.dtype) - else: - # upcast attention to fp32 - 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.attention_dropout, training=self.training) - attn_output = self.matmul_av(attn_weights, value_states) - attn_output = attn_output.reshape(bsz, -1, q_len, self.head_dim) - - 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_weights = gaudi_eager_attention_forward( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + sliding_window=getattr(self.config, "sliding_window", None), # diff with Llama + **kwargs, ) attn_output = attn_output.transpose(1, 2).contiguous() - - attn_output = attn_output.reshape(bsz, q_len, self.hidden_size) - + attn_output = attn_output.reshape(*input_shape, -1).contiguous() attn_output = self.o_proj(attn_output) - if not output_attentions: - attn_weights = None - return attn_output, attn_weights, past_key_value def attention_all_reduce(self, attn_output): @@ -363,6 +375,7 @@ def forward( 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 token_idx: Optional[torch.Tensor] = None, attn_softmax_bf16: Optional[bool] = False, reuse_cache: Optional[bool] = False, @@ -380,6 +393,7 @@ def forward( output_attentions=output_attentions, use_cache=use_cache, cache_position=cache_position, + position_embeddings=position_embeddings, token_idx=token_idx, attn_softmax_bf16=attn_softmax_bf16, reuse_cache=reuse_cache, @@ -412,6 +426,7 @@ def pre_attn( 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 token_idx: Optional[torch.Tensor] = None, attn_softmax_bf16: Optional[bool] = False, reuse_cache: Optional[bool] = False, @@ -423,10 +438,9 @@ def pre_attn( hidden_states = self.input_layernorm(hidden_states) hidden_states, attn_weights, present_key_value = self.self_attn.pre_attn_forward( hidden_states, + position_embeddings, attention_mask, - position_ids, past_key_value, - output_attentions, use_cache, cache_position, token_idx, @@ -436,6 +450,7 @@ def pre_attn( flash_attention_recompute, flash_attention_causal_mask, cache_idx=cache_idx, + position_ids=position_ids, ) return hidden_states, attn_weights, present_key_value @@ -477,7 +492,6 @@ def __init__(self, config: Starcoder2Config): self.layers = torch.nn.ModuleList( [GaudiStarcoder2DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] ) - self._attn_implementation = "eager" self.norm = torch.nn.LayerNorm(config.hidden_size, eps=config.norm_epsilon) self.gradient_checkpointing = False # Initialize weights and apply final processing @@ -533,12 +547,11 @@ def forward( else: raise ValueError("You have to specify either input_ids or inputs_embeds") - 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 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) @@ -572,8 +585,11 @@ def forward( inputs_embeds, past_seen_tokens, ) - # embed positions + hidden_states = inputs_embeds + hidden_states = torch.nn.functional.dropout( + hidden_states, p=self.embedding_dropout, training=self.training + ) # main diff with Llama # decoder layers all_hidden_states = () if output_hidden_states else None @@ -583,44 +599,26 @@ def forward( if lazy_mode: htcore.mark_step() - for layer_idx, decoder_layer in enumerate(self.layers): + 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], + 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] @@ -693,7 +691,7 @@ def forward( flash_attention_causal_mask: Optional[bool] = False, cache_idx: int = None, lazy_mode: Optional[bool] = True, - **loss_kwargs, + **kwargs: Unpack[KwargsForCausalLM], ) -> Union[Tuple, CausalLMOutputWithPast]: output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = ( @@ -741,7 +739,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=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs) if not return_dict: output = (logits,) + outputs[1:] diff --git a/optimum/habana/transformers/trainer.py b/optimum/habana/transformers/trainer.py index 44690f4b6a..654e727599 100644 --- a/optimum/habana/transformers/trainer.py +++ b/optimum/habana/transformers/trainer.py @@ -1239,6 +1239,7 @@ 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 @@ -1254,7 +1255,22 @@ def _load_best_model(self): active_adapter = model.active_adapter if os.path.exists(best_adapter_model_path) or os.path.exists(best_safe_adapter_model_path): - model.load_adapter(self.state.best_model_checkpoint, active_adapter) + try: + model.load_adapter(self.state.best_model_checkpoint, active_adapter) + except RuntimeError as exc: + if model.peft_config[active_adapter].is_prompt_learning: + # for context: https://github.com/huggingface/peft/issues/2256 + msg = ( + "When using prompt learning PEFT methods such as " + f"{model.peft_config[active_adapter].peft_type.value}, setting " + "load_best_model_at_end=True can lead to errors, it is recommended " + "to set this to False and to load the model manually from the checkpoint " + "directory using PeftModel.from_pretrained(base_model, ) after training " + "has finished." + ) + raise RuntimeError(msg) from exc + else: + raise # Load_adapter has no return value present, modify it when appropriate. from torch.nn.modules.module import _IncompatibleKeys @@ -1277,7 +1293,7 @@ def _load_best_model(self): state_dict = torch.load( best_model_path, map_location="cpu", - weights_only=True, + **weights_only_kwarg, ) # If the model is on the GPU, it still works! @@ -1613,7 +1629,10 @@ def training_step( inputs = self._prepare_inputs(inputs) with self.compute_loss_context_manager(): - loss = self.compute_loss(model, inputs, num_items_in_batch=num_items_in_batch) + if self.model_accepts_loss_kwargs: + loss = self.compute_loss(model, inputs) + else: + loss = self.compute_loss(model, inputs, num_items_in_batch=num_items_in_batch) del inputs kwargs = {} @@ -2607,6 +2626,10 @@ def get_batch_samples(self, epoch_iterator, num_batches): break # TODO: execute get_batch_samples outside of the training loop (before training) and uncomment the following lines + # Keep default behavior the same + # if not self.model_accepts_loss_kwargs: + # return batch_samples, None + # if len(batch_samples) > 0 and "labels" in batch_samples[0]: # # For now we don't support object detection # try: @@ -2614,7 +2637,7 @@ def get_batch_samples(self, epoch_iterator, num_batches): # except (TypeError, AttributeError): # pass - # if self.args.average_tokens_across_devices: + # if self.args.average_tokens_across_devices and num_items_in_batch is not None: # num_items_in_batch = self.accelerator.gather(num_items_in_batch).sum().item() # if torch.is_tensor(num_items_in_batch): diff --git a/optimum/habana/transformers/trainer_seq2seq.py b/optimum/habana/transformers/trainer_seq2seq.py index 0864d819b3..65880ac4a9 100644 --- a/optimum/habana/transformers/trainer_seq2seq.py +++ b/optimum/habana/transformers/trainer_seq2seq.py @@ -69,6 +69,7 @@ def __init__( 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), @@ -83,6 +84,7 @@ def __init__( eval_dataset=eval_dataset, processing_class=processing_class, model_init=model_init, + compute_loss_func=compute_loss_func, compute_metrics=compute_metrics, callbacks=callbacks, optimizers=optimizers, @@ -401,10 +403,12 @@ def prediction_step( return loss, generated_tokens, labels def _pad_tensors_to_max_len(self, tensor, max_length): - if self.tokenizer is not None and hasattr(self.tokenizer, "pad_token_id"): + if self.processing_class is not None and hasattr(self.processing_class, "pad_token_id"): # If PAD token is not defined at least EOS token has to be defined pad_token_id = ( - self.tokenizer.pad_token_id if self.tokenizer.pad_token_id is not None else self.tokenizer.eos_token_id + self.processing_class.pad_token_id + if self.processing_class.pad_token_id is not None + else self.processing_class.eos_token_id ) else: if self.model.config.pad_token_id is not None: diff --git a/optimum/habana/transformers/training_args.py b/optimum/habana/transformers/training_args.py index 56fdb1d154..1e631e593f 100644 --- a/optimum/habana/transformers/training_args.py +++ b/optimum/habana/transformers/training_args.py @@ -458,7 +458,7 @@ def __post_init__(self): self.save_steps = int(self.save_steps) # Sanity checks for load_best_model_at_end: we require save and eval strategies to be compatible. - if self.load_best_model_at_end: + if self.load_best_model_at_end and self.save_strategy != SaveStrategy.BEST: if self.eval_strategy != self.save_strategy: raise ValueError( "--load_best_model_at_end requires the save and eval strategy to match, but found\n- Evaluation " @@ -897,7 +897,7 @@ def _setup_devices(self) -> "torch.device": if not is_accelerate_available(): raise ImportError( f"Using the `Trainer` with `PyTorch` requires `accelerate>={ACCELERATE_MIN_VERSION}`: " - "Please run `pip install transformers[torch]` or `pip install accelerate -U`" + 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} diff --git a/optimum/habana/transformers/training_args_seq2seq.py b/optimum/habana/transformers/training_args_seq2seq.py index 82e02bb491..58269c5862 100644 --- a/optimum/habana/transformers/training_args_seq2seq.py +++ b/optimum/habana/transformers/training_args_seq2seq.py @@ -33,11 +33,6 @@ class GaudiSeq2SeqTrainingArguments(GaudiTrainingArguments): to enable deployment on Habana's Gaudi. Args: - sortish_sampler (`bool`, *optional*, defaults to `False`): - Whether to use a *sortish sampler* or not. Only possible if the underlying datasets are *Seq2SeqDataset* - for now but will become generally available in the near future. - It sorts the inputs according to lengths in order to minimize the padding size, with a bit of randomness - for the training set. predict_with_generate (`bool`, *optional*, defaults to `False`): Whether to use generate to calculate generative metrics (ROUGE, BLEU). generation_max_length (`int`, *optional*): diff --git a/setup.py b/setup.py index 57d184cce2..4043511fc1 100644 --- a/setup.py +++ b/setup.py @@ -29,7 +29,7 @@ INSTALL_REQUIRES = [ - "transformers >= 4.47.1, < 4.48.0", + "transformers >= 4.48.0, < 4.49.0", "optimum", "torch", "accelerate >= 0.33.0, < 0.34.0", diff --git a/tests/test_trainer.py b/tests/test_trainer.py index 92118a5b55..f3bc5b2d65 100644 --- a/tests/test_trainer.py +++ b/tests/test_trainer.py @@ -414,7 +414,15 @@ def get_gaudi_config(gaudi_config_name_or_path: Optional[Union[str, Path]] = Non return GaudiConfig.from_pretrained(gaudi_config_name_or_path) def get_regression_trainer( - a=0, b=0, double_output=False, train_len=64, eval_len=64, pretrained=True, keep_report_to=False, **kwargs + a=0, + b=0, + double_output=False, + train_len=64, + eval_len=64, + pretrained=True, + keep_report_to=False, + output_dir=None, + **kwargs, ): label_names = kwargs.get("label_names", None) gradient_checkpointing = kwargs.get("gradient_checkpointing", False) @@ -442,8 +450,8 @@ def get_regression_trainer( compute_metrics = kwargs.pop("compute_metrics", None) data_collator = kwargs.pop("data_collator", None) optimizers = kwargs.pop("optimizers", (None, None)) - output_dir = kwargs.pop("output_dir", "./regression") preprocess_logits_for_metrics = kwargs.pop("preprocess_logits_for_metrics", None) + assert output_dir is not None, "output_dir should be specified for testing" args = RegressionGaudiTrainingArguments( output_dir, use_habana=True, use_lazy_mode=True, a=a, b=b, keep_report_to=keep_report_to, **kwargs From 2a3affa969b4cb70ec471e3453a27a0e37607790 Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Wed, 15 Jan 2025 10:46:57 +0000 Subject: [PATCH 24/89] Small fixes --- .../models/cohere/modeling_cohere.py | 1 - .../transformers/models/gemma/modeling_gemma.py | 17 ++++++++++------- .../models/gemma2/modeling_gemma2.py | 9 ++++++--- .../transformers/models/llama/modeling_llama.py | 4 ++++ 4 files changed, 20 insertions(+), 11 deletions(-) diff --git a/optimum/habana/transformers/models/cohere/modeling_cohere.py b/optimum/habana/transformers/models/cohere/modeling_cohere.py index 119989988b..495ae2f9f0 100644 --- a/optimum/habana/transformers/models/cohere/modeling_cohere.py +++ b/optimum/habana/transformers/models/cohere/modeling_cohere.py @@ -59,7 +59,6 @@ def forward( value_states = value_states.transpose(1, 2) cos, sin = self.rotary_emb(value_states, kwargs["position_ids"]) - # print("SHAPEEEEEEEEEEEE", cos.shape, sin.shape, query_states.shape, key_states.shape) query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) if past_key_value is not None: diff --git a/optimum/habana/transformers/models/gemma/modeling_gemma.py b/optimum/habana/transformers/models/gemma/modeling_gemma.py index 7ec22d6c12..a4de41d29a 100755 --- a/optimum/habana/transformers/models/gemma/modeling_gemma.py +++ b/optimum/habana/transformers/models/gemma/modeling_gemma.py @@ -163,7 +163,7 @@ def forward(self, cur, dim, idx): return self.update(self.cache, cur, dim, idx, self.inp_seq_len) -def eager_attention_forward( +def gaudi_eager_attention_forward( module: torch.nn.Module, query: torch.Tensor, key: torch.Tensor, @@ -171,25 +171,28 @@ def eager_attention_forward( attention_mask: Optional[torch.Tensor], scaling: float, dropout: float = 0.0, + attn_softmax_bf16: bool = False, **kwargs, ): + bsz, q_len = kwargs["input_shape"] query_states, key_states, value_states, attention_mask = gaudi_gemma_repeat_kv( query, key, value, attention_mask, module.num_key_value_groups ) - attn_weights = module.matmul_qk(query, key_states.transpose(2, 3)) * scaling + attn_weights = module.matmul_qk(query_states, key_states.transpose(-2, -1)) * scaling if attention_mask is not None: causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] attn_weights = attn_weights + causal_mask - if kwargs["attn_softmax_bf16"]: + if attn_softmax_bf16: attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1, dtype=query_states.dtype) else: # upcast attention to fp32 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=dropout, training=module.training) attn_output = module.matmul_av(attn_weights, value_states) - # attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.reshape(bsz, -1, q_len, module.head_dim) return attn_output, attn_weights @@ -386,8 +389,7 @@ def pre_attn_forward( ) else: - kwargs["attn_softmax_bf16"] = attn_softmax_bf16 - attn_output, attn_weights = eager_attention_forward( + attn_output, attn_weights = gaudi_eager_attention_forward( self, query_states, key_states, @@ -395,7 +397,8 @@ def pre_attn_forward( attention_mask, dropout=0.0 if not self.training else self.attention_dropout, scaling=self.scaling, - **kwargs, + attn_softmax_bf16=attn_softmax_bf16, + input_shape=input_shape, ) attn_output = attn_output.transpose(1, 2).contiguous() diff --git a/optimum/habana/transformers/models/gemma2/modeling_gemma2.py b/optimum/habana/transformers/models/gemma2/modeling_gemma2.py index 9cd07b560d..5eb5baf632 100755 --- a/optimum/habana/transformers/models/gemma2/modeling_gemma2.py +++ b/optimum/habana/transformers/models/gemma2/modeling_gemma2.py @@ -254,6 +254,8 @@ def gaudi_eager_attention_forward( softcap: Optional[float] = None, **kwargs, ) -> Tuple[torch.Tensor, torch.Tensor]: + bsz, q_len = kwargs["input_shape"] + if scaling is None: scaling = module.head_dim**-0.5 @@ -261,7 +263,7 @@ def gaudi_eager_attention_forward( query, key, value, attention_mask, module.num_key_value_groups ) - attn_weights = module.matmul_qk(query_states, key_states.transpose(2, 3)) * scaling + attn_weights = module.matmul_qk(query_states, key_states.transpose(-2, -1)) * scaling if softcap is not None: attn_weights = attn_weights / softcap @@ -275,6 +277,8 @@ def gaudi_eager_attention_forward( attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) attn_weights = torch.nn.functional.dropout(attn_weights, p=dropout, training=module.training) attn_output = module.matmul_av(attn_weights, value_states) + attn_output = attn_output.reshape(bsz, -1, q_len, module.head_dim) + return attn_output, attn_weights @@ -469,11 +473,10 @@ def pre_attn_forward( scaling=self.scaling, sliding_window=self.sliding_window, softcap=self.attn_logit_softcapping, - **kwargs, + input_shape=input_shape, ) attn_output = attn_output.transpose(1, 2).contiguous() - attn_output = attn_output.reshape(*input_shape, -1).contiguous() attn_output = self.o_proj(attn_output) diff --git a/optimum/habana/transformers/models/llama/modeling_llama.py b/optimum/habana/transformers/models/llama/modeling_llama.py index eb4e32d53f..021b5e42a1 100755 --- a/optimum/habana/transformers/models/llama/modeling_llama.py +++ b/optimum/habana/transformers/models/llama/modeling_llama.py @@ -364,7 +364,9 @@ def gaudi_eager_attention_forward( scaling: float, dropout: float = 0.0, attn_softmax_bf16: bool = False, + **kwargs, ): + bsz, q_len = kwargs["input_shape"] query_states, key_states, value_states, attention_mask = gaudi_llama_repeat_kv( query, key, value, attention_mask, module.num_key_value_groups ) @@ -381,6 +383,7 @@ def gaudi_eager_attention_forward( 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=dropout, training=module.training) attn_output = module.matmul_av(attn_weights, value_states) + attn_output = attn_output.reshape(bsz, -1, q_len, module.head_dim) return attn_output, attn_weights @@ -681,6 +684,7 @@ def pre_attn_forward( dropout=0.0 if not self.training else self.attention_dropout, scaling=self.scaling, attn_softmax_bf16=attn_softmax_bf16, + input_shape=input_shape, ) attn_output = attn_output.transpose(1, 2).contiguous() From 064f4c1e596b33d1437a3edffe6e3433abd38ed5 Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Wed, 15 Jan 2025 16:02:13 +0000 Subject: [PATCH 25/89] Fix integration tests --- tests/test_trainer.py | 38 +++++++++++++++++++++++++++----------- 1 file changed, 27 insertions(+), 11 deletions(-) diff --git a/tests/test_trainer.py b/tests/test_trainer.py index f3bc5b2d65..bca097be1f 100644 --- a/tests/test_trainer.py +++ b/tests/test_trainer.py @@ -1847,17 +1847,18 @@ def test_can_resume_training(self): # Now check failures # 1. fail to find a bogus checkpoint - trainer = get_regression_trainer() - with self.assertRaises(Exception) as context: - trainer.train(resume_from_checkpoint=f"{checkpoint}-bogus") - self.assertTrue("Can't find a valid checkpoint at" in str(context.exception)) + with tempfile.TemporaryDirectory() as tmpdir: + trainer = get_regression_trainer(output_dir=tmpdir) + with self.assertRaises(Exception) as context: + trainer.train(resume_from_checkpoint=f"{checkpoint}-bogus") + self.assertTrue("Can't find a valid checkpoint at" in str(context.exception)) # 2. fail to find any checkpoint - due a fresh output_dir - output_dir2 = self.get_auto_remove_tmp_dir() - trainer = get_regression_trainer(output_dir=output_dir2) - with self.assertRaises(Exception) as context: - trainer.train(resume_from_checkpoint=True) - self.assertTrue("No valid checkpoint found in output directory" in str(context.exception)) + with tempfile.TemporaryDirectory() as tmpdir: + trainer = get_regression_trainer(output_dir=tmpdir) + with self.assertRaises(Exception) as context: + trainer.train(resume_from_checkpoint=True) + self.assertTrue("No valid checkpoint found in output directory" in str(context.exception)) def test_resume_training_with_randomness(self): train_dataset = RegressionDataset(length=128) @@ -2929,7 +2930,9 @@ def test_save_best_checkpoint(self): total=total, ) - # Case 3: Metric name not provided; throw error. + def test_metric_for_best_model_behavior(self): + # Case 1: Metric name not provided when `save_strategy == "best"`. + # Should raise ValueError. with tempfile.TemporaryDirectory() as tmpdir: with self.assertRaises(ValueError) as context: trainer = get_regression_trainer( @@ -2941,9 +2944,22 @@ def test_save_best_checkpoint(self): save_strategy="best", compute_metrics=AlmostAccuracy(), ) - self.assertIn("`args.metric_for_best_model` must be provided", str(context.exception)) + # Case 2: Metric name not provided when `load_best_model_at_end == True`. + # `metric_for_best_model` should be set to `"loss"` by default. + with tempfile.TemporaryDirectory() as tmpdir: + trainer = get_regression_trainer( + a=1.5, + b=2.5, + output_dir=tmpdir, + learning_rate=0.1, + eval_strategy="steps", + save_strategy="steps", + load_best_model_at_end=True, + ) + self.assertTrue(trainer.args.metric_for_best_model == "loss") + def test_profiling(self): with tempfile.TemporaryDirectory() as tmp_dir: # 24 total steps and compilation takes place during the 1st three steps From 21714f7306d59403d966cef223e945f0e8c9368b Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Thu, 16 Jan 2025 10:27:28 +0000 Subject: [PATCH 26/89] Fixes for text-generation --- .../models/decilm/modeling_decilm.py | 48 +++++++++++++++++++ .../models/mixtral/modeling_mixtral.py | 1 + .../transformers/models/phi/modeling_phi.py | 1 + .../models/qwen2/modeling_qwen2.py | 2 + tests/test_text_generation_example.py | 8 ++-- 5 files changed, 56 insertions(+), 4 deletions(-) diff --git a/optimum/habana/transformers/models/decilm/modeling_decilm.py b/optimum/habana/transformers/models/decilm/modeling_decilm.py index 03651cf985..6618911530 100644 --- a/optimum/habana/transformers/models/decilm/modeling_decilm.py +++ b/optimum/habana/transformers/models/decilm/modeling_decilm.py @@ -179,6 +179,54 @@ def __init__(self, config: DeciLMConfig, layer_idx: int): self.input_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + def forward( + self, + hidden_states: torch.Tensor, + 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, + token_idx: Optional[torch.Tensor] = None, + **kwargs, + ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + residual = hidden_states + + hidden_states = self.input_layernorm(hidden_states) + + # Self Attention + hidden_states, self_attn_weights, 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, + **kwargs, + ) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + outputs = (hidden_states,) + + if output_attentions: + outputs += (self_attn_weights,) + + if use_cache: + outputs += (present_key_value,) + + return outputs + @add_start_docstrings( "The bare DeciLM Model outputting raw hidden-states without any specific head on top.", diff --git a/optimum/habana/transformers/models/mixtral/modeling_mixtral.py b/optimum/habana/transformers/models/mixtral/modeling_mixtral.py index e009d7f8a9..b921fbf3ea 100644 --- a/optimum/habana/transformers/models/mixtral/modeling_mixtral.py +++ b/optimum/habana/transformers/models/mixtral/modeling_mixtral.py @@ -213,6 +213,7 @@ def __init__(self, config: MixtralConfig, layer_idx: Optional[int] = None): self.inp_seq_len = -1 self.rotary_emb = GaudiLlamaRotaryEmbedding(config=config) self.block_size = 1024 + self.num_key_value_heads = config.num_key_value_heads def allocate_kv_cache(self, batch_size, max_seq_len, inp_seq_len): cache_shape = (batch_size, self.num_key_value_heads, max_seq_len, self.head_dim) diff --git a/optimum/habana/transformers/models/phi/modeling_phi.py b/optimum/habana/transformers/models/phi/modeling_phi.py index c86e7563ac..e7bd7b3b52 100644 --- a/optimum/habana/transformers/models/phi/modeling_phi.py +++ b/optimum/habana/transformers/models/phi/modeling_phi.py @@ -117,6 +117,7 @@ def __init__(self, config: PhiConfig, layer_idx: Optional[int] = None): self.v_cache = KVCache() self.inp_seq_len = -1 self.rotary_emb = GaudiRotaryEmbedding(config=self.config) + self.num_key_value_heads = config.num_key_value_heads def allocate_kv_cache(self, batch_size, max_seq_len, inp_seq_len): cache_shape = (batch_size, self.num_key_value_heads, max_seq_len, self.head_dim) diff --git a/optimum/habana/transformers/models/qwen2/modeling_qwen2.py b/optimum/habana/transformers/models/qwen2/modeling_qwen2.py index 5573aa19a6..e8536662ae 100644 --- a/optimum/habana/transformers/models/qwen2/modeling_qwen2.py +++ b/optimum/habana/transformers/models/qwen2/modeling_qwen2.py @@ -226,6 +226,8 @@ def __init__(self, config: Qwen2Config, layer_idx: Optional[int] = None): else None ) + self.num_key_value_heads = config.num_key_value_heads + def get_k_proj_weight(self): """4bit quantization in GPTQ replaces the k_proj.weight with qweight.""" if hasattr(self.k_proj, "qweight"): diff --git a/tests/test_text_generation_example.py b/tests/test_text_generation_example.py index 912cbefae8..7f1f2543f1 100644 --- a/tests/test_text_generation_example.py +++ b/tests/test_text_generation_example.py @@ -108,10 +108,10 @@ "bigcode/starcoder": 'def print_hello_world():\n print("Hello World")\n\ndef print_hello_world_twice():\n print_hello_world()\n print_hello_world()\n\ndef print_hello_world_thrice():\n print_hello_world()\n print_hello_world()\n print_hello_world()\n\ndef print_hello_world_four_times():\n print_hello_world()\n print_hello_world()\n print_hello_world()\n ', "bigcode/starcoder2-3b": 'def print_hello_world():\n print("Hello World")\n\ndef print_hello_world_with_name(name):\n print("Hello World, " + name)\n\ndef print_hello_world_with_name_and_age(name, age):\n print("Hello World, " + name + ", " + str(age))\n\ndef print_hello_world_with_name_and_age_and_gender(name, age, gender):\n print("Hello', "google/gemma-7b": "DeepSpeed is a machine learning framework that enables training of large-scale models on commodity hardware. It is designed to be a drop-in replacement for PyTorch, and it is compatible with the existing PyTorch ecosystem. DeepSpeed is designed to be easy to use, and it provides a number of features that make it easy to train large-scale models. DeepSpeed is designed to be scalable, and it can be used to train models on a single machine or on a cluster of machines. DeepSpeed is designed to be efficient,", - "google/gemma-2-9b": "DeepSpeed is a machine learning framework that enables training of large-scale deep learning models on a single GPU or across multiple GPUs. It is designed to be easy to use and highly scalable, making it a powerful tool for researchers and practitioners working with large-scale deep learning models.\n\nDeepSpeed is built on top of PyTorch, a popular deep learning framework, and provides a set of tools and libraries that make it easy to train large-scale models. It includes features such as zero-shot inference, which allows models to be", - "meta-llama/Llama-2-7b-hf": "DeepSpeed is a machine learning framework for deep learning. It is designed to be fast and efficient, while also being easy to use. DeepSpeed is based on the TensorFlow framework, and it uses the TensorFlow library to perform computations.\nDeepSpeed is a deep learning framework that is designed to be fast and efficient. It is based on the TensorFlow library and uses the TensorFlow library to perform computations. DeepSpeed is designed to be easy to use and to provide a high level of flex", - "mistralai/Mistral-7B-v0.1": "DeepSpeed is a machine learning framework that accelerates training of large models on a single machine or distributed systems. It is designed to be compatible with PyTorch and TensorFlow, and can be used to train models on a single machine or on a distributed system.\n\nDeepSpeed is a machine learning framework that accelerates training of large models on a single machine or distributed systems. It is designed to be compatible with PyTorch and TensorFlow, and can be used to train models on a single machine or on a distributed system", - "mistralai/Mixtral-8x7B-v0.1": "DeepSpeed is a machine learning framework that enables training of large models on a single machine with a single GPU. It is designed to be easy to use and efficient, and it can be used to train models on a variety of tasks.\n\n## Introduction\n\nDeepSpeed is a machine learning framework that enables training of large models on a single machine with a single GPU. It is designed to be easy to use and efficient, and it can be used to train models on a variety of tasks.\n\n## What is DeepSpeed", + "google/gemma-2-9b": "DeepSpeed is a machine learning framework that enables training of large-scale deep learning models on a single GPU or across multiple GPUs. It is designed to be easy to use and highly scalable, making it a popular choice for training large-scale models such as GPT-3 and BERT.\n\nDeepSpeed is built on top of PyTorch, a popular deep learning framework, and provides a set of tools and libraries that make it easy to train large-scale models. It includes features such as zero-shot learning, which allows models to", + "meta-llama/Llama-2-7b-hf": "DeepSpeed is a machine learning framework for deep learning. It is designed to be fast and efficient, while also being easy to use. DeepSpeed is based on the TensorFlow framework, and it uses the TensorFlow library to perform computations.\nDeepSpeed is a deep learning framework that is designed to be fast and efficient. It is based on the TensorFlow library and uses the TensorFlow library to perform computations. DeepSpeed is designed to be easy to use and to provide a high level of performance", + "mistralai/Mistral-7B-v0.1": "DeepSpeed is a machine learning framework that accelerates training of large models on a single machine or distributed systems. It is designed to be easy to use and flexible, allowing users to quickly train models on a variety of hardware platforms.\n\nDeepSpeed is a machine learning framework that accelerates training of large models on a single machine or distributed systems. It is designed to be easy to use and flexible, allowing users to quickly train models on a variety of hardware platforms.\n\nDeepSpeed is a machine learning framework that accelerates training", + "mistralai/Mixtral-8x7B-v0.1": "DeepSpeed is a machine learning framework that enables training of large models on a single machine with a single GPU. It is designed to be easy to use and efficient, and it can be used to train models on a variety of tasks.\n\n## What is DeepSpeed?\n\nDeepSpeed is a machine learning framework that enables training of large models on a single machine with a single GPU. It is designed to be easy to use and efficient, and it can be used to train models on a variety of tasks.\n\n##", "Qwen/Qwen2-7B": "DeepSpeed is a machine learning framework that provides a unified interface for training deep learning models. It is designed to be easy to use and to provide high performance. DeepSpeed is built on top of PyTorch and TensorFlow, and it supports a wide range of models, including transformers, convolutional neural networks, and recurrent neural networks.\nDeepSpeed is a machine learning framework that provides a unified interface for training deep learning models. It is designed to be easy to use and to provide high performance. DeepSpeed is built on top of Py", } else: From 1cfd53be563d28a278052e13f101299c676c416a Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Thu, 16 Jan 2025 17:55:35 +0000 Subject: [PATCH 27/89] Fixes --- .../transformers/models/llama/modeling_llama.py | 3 +++ .../transformers/models/mllama/modeling_mllama.py | 11 +++-------- optimum/habana/transformers/trainer.py | 9 ++++++--- optimum/habana/trl/trainer/dpo_trainer.py | 13 +++++++++++++ optimum/habana/trl/trainer/reward_trainer.py | 2 +- 5 files changed, 26 insertions(+), 12 deletions(-) diff --git a/optimum/habana/transformers/models/llama/modeling_llama.py b/optimum/habana/transformers/models/llama/modeling_llama.py index 021b5e42a1..d8d9344684 100755 --- a/optimum/habana/transformers/models/llama/modeling_llama.py +++ b/optimum/habana/transformers/models/llama/modeling_llama.py @@ -399,6 +399,9 @@ def __init__(self, config: LlamaConfig, layer_idx: Optional[int] = None): self.rotary_emb = GaudiLlamaRotaryEmbedding(config=config) self.num_key_value_heads = config.num_key_value_heads + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = getattr(config, "head_dim", self.hidden_size // self.num_heads) if hasattr(config, "fused_qkv") and config.fused_qkv: self.num_heads = config.num_attention_heads diff --git a/optimum/habana/transformers/models/mllama/modeling_mllama.py b/optimum/habana/transformers/models/mllama/modeling_mllama.py index 6d2d2a08fb..88168c9e94 100644 --- a/optimum/habana/transformers/models/mllama/modeling_mllama.py +++ b/optimum/habana/transformers/models/mllama/modeling_mllama.py @@ -928,7 +928,7 @@ class GaudiMllamaForConditionalGeneration(MllamaForConditionalGeneration): def __init__(self, config: MllamaConfig): # sdpa is better for vision model in HPU config._attn_implementation = "sdpa" - super(GaudiMllamaForConditionalGeneration, self).__init__(config) + super().__init__(config) def forward( self, @@ -1260,13 +1260,8 @@ 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] - intermediate_hidden_states = [ - hidden_state - for idx, hidden_state in enumerate(all_intermediate_hidden_states) - if idx in self.intermediate_layers_indices - ] - intermediate_hidden_states = torch.stack(intermediate_hidden_states, dim=-1) + all_intermediate_hidden_states = [output[1][i] for i in self.intermediate_layers_indices] + intermediate_hidden_states = torch.stack(all_intermediate_hidden_states, dim=-1) """ intermediate_hidden_states = torch.stack(all_intermediate_hidden_states, dim=-1) diff --git a/optimum/habana/transformers/trainer.py b/optimum/habana/transformers/trainer.py index 654e727599..290af48444 100644 --- a/optimum/habana/transformers/trainer.py +++ b/optimum/habana/transformers/trainer.py @@ -995,7 +995,7 @@ def hpu_deepspeed_checkpointing(function, *checkpoint_args, use_reentrant: Optio 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(epoch_iterator, num_batches) + batch_samples, num_items_in_batch = self.get_batch_samples_transformers(epoch_iterator, num_batches) for i, inputs in enumerate(batch_samples): step += 1 @@ -1351,7 +1351,7 @@ def _maybe_log_save_evaluate(self, tr_loss, _grad_norm, model, trial, epoch, ign self._globalstep_last_logged = self.state.global_step self.store_flos() - self.log(logs, start_time) + self.log(logs, start_time=start_time) metrics = None if self.control.should_evaluate: @@ -2616,7 +2616,10 @@ def _zero_model_grad(self, model): model.zero_grad() model._zero_grad_kwargs = {} - def get_batch_samples(self, epoch_iterator, num_batches): + def get_batch_samples_transformers(self, epoch_iterator, num_batches): + """ + Added "_transformers" at the end of the method name to avoid a wrong call to a similarly named method in TRL trainers. + """ batch_samples = [] num_items_in_batch = None for _ in range(num_batches): diff --git a/optimum/habana/trl/trainer/dpo_trainer.py b/optimum/habana/trl/trainer/dpo_trainer.py index 84c48f1782..3af14d6555 100644 --- a/optimum/habana/trl/trainer/dpo_trainer.py +++ b/optimum/habana/trl/trainer/dpo_trainer.py @@ -668,3 +668,16 @@ def cross_entropy_loss(logits, labels): return (chosen_logps, rejected_logps, chosen_logits, rejected_logits, nll_loss, outputs.aux_loss) return (chosen_logps, rejected_logps, chosen_logits, rejected_logits, nll_loss) + + def log(self, logs: Dict[str, float], **kwargs) -> None: + """ + Changes: + - add `**kwargs` to the method arguments to make sure it's compatible with Transformers + """ + # logs either has 'loss' or 'eval_loss' + train_eval = "train" if "loss" in logs else "eval" + # Add averaged stored metrics to logs + for key, metrics in self._stored_metrics[train_eval].items(): + logs[key] = torch.tensor(metrics).mean().item() + del self._stored_metrics[train_eval] + return super().log(logs) diff --git a/optimum/habana/trl/trainer/reward_trainer.py b/optimum/habana/trl/trainer/reward_trainer.py index bbb0c761fe..cd551ef60c 100644 --- a/optimum/habana/trl/trainer/reward_trainer.py +++ b/optimum/habana/trl/trainer/reward_trainer.py @@ -28,7 +28,7 @@ class GaudiRewardTrainer(GaudiTrainer): Copied from https://github.com/huggingface/trl/blob/v0.7.6/examples/research_projects/stack_llama/scripts/reward_modeling.py#L266 """ - def compute_loss(self, model, inputs, return_outputs=False): + def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None): rewards_j = model(input_ids=inputs["input_ids_j"], attention_mask=inputs["attention_mask_j"])[0] rewards_k = model(input_ids=inputs["input_ids_k"], attention_mask=inputs["attention_mask_k"])[0] loss = -nn.functional.logsigmoid(rewards_j - rewards_k).mean() From 573cc574669e70ff54d5a61c4632e5c6fb7949ec Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Thu, 16 Jan 2025 18:18:26 +0000 Subject: [PATCH 28/89] Style --- .../models/starcoder2/modeling_starcoder2.py | 50 +++++++++++++------ tests/test_encoder_decoder.py | 2 +- 2 files changed, 35 insertions(+), 17 deletions(-) diff --git a/optimum/habana/transformers/models/starcoder2/modeling_starcoder2.py b/optimum/habana/transformers/models/starcoder2/modeling_starcoder2.py index 42f4cd5e9a..ecc6dce685 100644 --- a/optimum/habana/transformers/models/starcoder2/modeling_starcoder2.py +++ b/optimum/habana/transformers/models/starcoder2/modeling_starcoder2.py @@ -603,22 +603,40 @@ def forward( if output_hidden_states: all_hidden_states += (hidden_states,) - 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, - ) + 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] diff --git a/tests/test_encoder_decoder.py b/tests/test_encoder_decoder.py index 723739eb5b..25e7f69b01 100644 --- a/tests/test_encoder_decoder.py +++ b/tests/test_encoder_decoder.py @@ -189,7 +189,7 @@ def _test_text_translation( "--do_predict", "--source_lang en", "--target_lang ro", - '--source_prefix "translate English to Romanian: "--dataset_name wmt16', + '--source_prefix "translate English to Romanian: "--dataset_name wmt16', # noqa "--dataset_config_name ro-en", f"--per_device_eval_batch_size {batch_size}", f"--generation_num_beams {num_beams}", From a7bc5171e2daadd3188a82248746abd7c4ef8ff7 Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Thu, 16 Jan 2025 18:28:08 +0000 Subject: [PATCH 29/89] Again --- tests/test_encoder_decoder.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_encoder_decoder.py b/tests/test_encoder_decoder.py index 25e7f69b01..78ffc437a4 100644 --- a/tests/test_encoder_decoder.py +++ b/tests/test_encoder_decoder.py @@ -189,7 +189,8 @@ def _test_text_translation( "--do_predict", "--source_lang en", "--target_lang ro", - '--source_prefix "translate English to Romanian: "--dataset_name wmt16', # noqa + '--source_prefix "translate English to Romanian: "', + "--dataset_name wmt16", "--dataset_config_name ro-en", f"--per_device_eval_batch_size {batch_size}", f"--generation_num_beams {num_beams}", From f69e957d802f428afb5f79f4ada761292acee8b3 Mon Sep 17 00:00:00 2001 From: Vidya Galli Date: Tue, 28 Jan 2025 14:11:08 -0800 Subject: [PATCH 30/89] Fix for image2text lora llama test (#1731) --- tests/baselines/Llama_3_2_11B_Vision_Instruct.json | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/baselines/Llama_3_2_11B_Vision_Instruct.json b/tests/baselines/Llama_3_2_11B_Vision_Instruct.json index c2a58cc25c..fd90ab97f0 100644 --- a/tests/baselines/Llama_3_2_11B_Vision_Instruct.json +++ b/tests/baselines/Llama_3_2_11B_Vision_Instruct.json @@ -1,13 +1,13 @@ { "gaudi2": { "image2text_lora_finetune": { - "num_train_epochs": 2, + "num_train_epochs": 1, "eval_batch_size": 4, "distribution": { "multi_card": { "learning_rate": 5e-5, "train_batch_size": 2, - "train_runtime": 470, + "train_runtime": 350, "train_samples_per_second": 20.48, "eval_accuracy": 0.6, "extra_arguments": [ From 265e6a1e887adee55885561186b5de868494c2d4 Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Fri, 17 Jan 2025 08:51:03 +0000 Subject: [PATCH 31/89] Cherry-pick https://github.com/huggingface/transformers/pull/35651 --- optimum/habana/transformers/trainer.py | 11 ++--------- 1 file changed, 2 insertions(+), 9 deletions(-) diff --git a/optimum/habana/transformers/trainer.py b/optimum/habana/transformers/trainer.py index 290af48444..1e03283e45 100644 --- a/optimum/habana/transformers/trainer.py +++ b/optimum/habana/transformers/trainer.py @@ -1629,10 +1629,7 @@ def training_step( inputs = self._prepare_inputs(inputs) with self.compute_loss_context_manager(): - if self.model_accepts_loss_kwargs: - loss = self.compute_loss(model, inputs) - else: - loss = self.compute_loss(model, inputs, num_items_in_batch=num_items_in_batch) + loss = self.compute_loss(model, inputs, num_items_in_batch=num_items_in_batch) del inputs kwargs = {} @@ -1648,7 +1645,7 @@ def training_step( self.htcore.mark_step() # Finally we need to normalize the loss for reporting - if num_items_in_batch is None: + if not self.model_accepts_loss_kwargs and self.compute_loss_func is None: loss = loss / self.args.gradient_accumulation_steps if _is_peft_model(self.model) and self.model.peft_type == PeftType.ADALORA: @@ -2629,10 +2626,6 @@ def get_batch_samples_transformers(self, epoch_iterator, num_batches): break # TODO: execute get_batch_samples outside of the training loop (before training) and uncomment the following lines - # Keep default behavior the same - # if not self.model_accepts_loss_kwargs: - # return batch_samples, None - # if len(batch_samples) > 0 and "labels" in batch_samples[0]: # # For now we don't support object detection # try: From 32478f5d091f1d19a316a102d00f1cd5ef55bfd8 Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Fri, 31 Jan 2025 09:45:55 +0000 Subject: [PATCH 32/89] Upgrade to Transformers v4.48.2 --- .../habana/transformers/models/gemma2/modeling_gemma2.py | 9 +++++++++ setup.py | 2 +- 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/optimum/habana/transformers/models/gemma2/modeling_gemma2.py b/optimum/habana/transformers/models/gemma2/modeling_gemma2.py index 5eb5baf632..505c8c3ac3 100755 --- a/optimum/habana/transformers/models/gemma2/modeling_gemma2.py +++ b/optimum/habana/transformers/models/gemma2/modeling_gemma2.py @@ -546,6 +546,7 @@ def pre_attn( flash_attention_causal_mask: Optional[bool] = False, flash_attention_fast_softmax: Optional[bool] = False, cache_idx: int = None, + **kwargs, ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: hidden_states = self.input_layernorm(hidden_states) @@ -566,6 +567,7 @@ def pre_attn( flash_attention_causal_mask=flash_attention_causal_mask, flash_attention_fast_softmax=flash_attention_fast_softmax, cache_idx=cache_idx, + **kwargs, ) return hidden_states, attn_weights, present_key_value @@ -579,6 +581,7 @@ def forward( output_attentions: Optional[bool] = False, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, + last_cache_position: int = 0, token_idx: Optional[torch.Tensor] = None, attn_softmax_bf16: Optional[bool] = False, reuse_cache: Optional[bool] = False, @@ -587,6 +590,7 @@ def forward( flash_attention_causal_mask: Optional[bool] = False, flash_attention_fast_softmax: Optional[bool] = False, cache_idx: int = None, + **kwargs, ) -> 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 @@ -612,6 +616,7 @@ def forward( flash_attention_causal_mask=flash_attention_causal_mask, flash_attention_fast_softmax=flash_attention_fast_softmax, cache_idx=cache_idx, + **kwargs, ) self.self_attn.attention_all_reduce(hidden_states) @@ -685,6 +690,7 @@ def forward( 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, attn_softmax_bf16: Optional[bool] = False, reuse_cache: Optional[bool] = False, @@ -809,6 +815,7 @@ def forward( output_attentions, use_cache, cache_position, + last_cache_position, None, attn_softmax_bf16, False, @@ -827,6 +834,7 @@ def forward( 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, @@ -938,6 +946,7 @@ def forward( flash_attention_fast_softmax=flash_attention_fast_softmax, cache_idx=cache_idx, lazy_mode=lazy_mode, + **loss_kwargs, ) hidden_states = outputs[0] diff --git a/setup.py b/setup.py index 510a0c3658..c472e03326 100644 --- a/setup.py +++ b/setup.py @@ -29,7 +29,7 @@ INSTALL_REQUIRES = [ - "transformers >= 4.48.0, < 4.49.0", + "transformers >= 4.48.2, < 4.49.0", "optimum", "torch", "accelerate >= 0.33.0, < 0.34.0", From 1b79cf3525c57ee2054d89ac3023faa25ea28830 Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Fri, 31 Jan 2025 13:13:54 +0000 Subject: [PATCH 33/89] Fix deprecated imports following merged changes for DETR and Qwen2-VL --- optimum/habana/transformers/modeling_utils.py | 18 +- .../habana/transformers/models/__init__.py | 3 +- .../transformers/models/detr/__init__.py | 1 - .../transformers/models/detr/modeling_detr.py | 6 +- .../transformers/models/qwen2_vl/__init__.py | 2 +- .../models/qwen2_vl/modeling_qwen2_vl.py | 172 ++++++++++-------- 6 files changed, 113 insertions(+), 89 deletions(-) diff --git a/optimum/habana/transformers/modeling_utils.py b/optimum/habana/transformers/modeling_utils.py index 582c9ce67a..3dbafa4a5a 100644 --- a/optimum/habana/transformers/modeling_utils.py +++ b/optimum/habana/transformers/modeling_utils.py @@ -139,7 +139,6 @@ GaudiQwen2MoeForCausalLM, GaudiQwen2MoeMLP, GaudiQwen2MoeModel, - GaudiQwen2VisionSdpaAttention, GaudiQwen2VisionTransformerPretrainedModel, GaudiQwen2VLDecoderLayer, GaudiQwen2VLForConditionalGeneration, @@ -153,6 +152,7 @@ GaudiStarcoder2DecoderLayer, GaudiStarcoder2ForCausalLM, GaudiStarcoder2Model, + GaudiVisionSdpaAttention, GaudiWav2Vec2SdpaAttention, GaudiWhisperDecoder, GaudiWhisperDecoderLayer, @@ -200,7 +200,6 @@ gaudi_DetrConvModel_forward, gaudi_DetrHungarianMatcher_forward, gaudi_DetrLoss_forward, - gaudi_DetrLoss_get_targets_without_no_objects, gaudi_DetrLoss_loss_boxes, gaudi_DetrLoss_loss_cardinality, gaudi_DetrLoss_loss_labels, @@ -651,7 +650,7 @@ def adapt_transformers_to_gaudi(): ) # Optimization for qwen2-vl Gaudi - transformers.models.qwen2_vl.modeling_qwen2_vl.VisionSdpaAttention = GaudiQwen2VisionSdpaAttention + transformers.models.qwen2_vl.modeling_qwen2_vl.VisionSdpaAttention = GaudiVisionSdpaAttention transformers.models.qwen2_vl.modeling_qwen2_vl.Qwen2VLVisionBlock = GaudiQwen2VLVisionBlock transformers.models.qwen2_vl.modeling_qwen2_vl.Qwen2VisionTransformerPretrainedModel = ( GaudiQwen2VisionTransformerPretrainedModel @@ -755,11 +754,8 @@ def adapt_transformers_to_gaudi(): # Optimization for DETR model on Gaudi transformers.models.detr.modeling_detr.DetrConvModel.forward = gaudi_DetrConvModel_forward - transformers.models.detr.modeling_detr.DetrHungarianMatcher.forward = gaudi_DetrHungarianMatcher_forward - transformers.models.detr.modeling_detr.DetrLoss.get_targets_without_no_objects = ( - gaudi_DetrLoss_get_targets_without_no_objects - ) - transformers.models.detr.modeling_detr.DetrLoss.loss_labels = gaudi_DetrLoss_loss_labels - transformers.models.detr.modeling_detr.DetrLoss.loss_cardinality = gaudi_DetrLoss_loss_cardinality - transformers.models.detr.modeling_detr.DetrLoss.loss_boxes = gaudi_DetrLoss_loss_boxes - transformers.models.detr.modeling_detr.DetrLoss.forward = gaudi_DetrLoss_forward + transformers.loss.loss_for_object_detection.HungarianMatcher.forward = gaudi_DetrHungarianMatcher_forward + transformers.loss.loss_for_object_detection.ImageLoss.loss_labels = gaudi_DetrLoss_loss_labels + transformers.loss.loss_for_object_detection.ImageLoss.loss_cardinality = gaudi_DetrLoss_loss_cardinality + transformers.loss.loss_for_object_detection.ImageLoss.loss_boxes = gaudi_DetrLoss_loss_boxes + transformers.loss.loss_for_object_detection.ImageLoss.forward = gaudi_DetrLoss_forward diff --git a/optimum/habana/transformers/models/__init__.py b/optimum/habana/transformers/models/__init__.py index 4b41626c7f..6a0e72e9c1 100644 --- a/optimum/habana/transformers/models/__init__.py +++ b/optimum/habana/transformers/models/__init__.py @@ -75,7 +75,6 @@ gaudi_DetrConvModel_forward, gaudi_DetrHungarianMatcher_forward, gaudi_DetrLoss_forward, - gaudi_DetrLoss_get_targets_without_no_objects, gaudi_DetrLoss_loss_boxes, gaudi_DetrLoss_loss_cardinality, gaudi_DetrLoss_loss_labels, @@ -255,13 +254,13 @@ gaudi_qwen2moe_rmsnorm_forward, ) from .qwen2_vl import ( - GaudiQwen2VisionSdpaAttention, GaudiQwen2VisionTransformerPretrainedModel, GaudiQwen2VLDecoderLayer, GaudiQwen2VLForConditionalGeneration, GaudiQwen2VLModel, GaudiQwen2VLSdpaAttention, GaudiQwen2VLVisionBlock, + GaudiVisionSdpaAttention, ) from .seamless_m4t import ( gaudi_SeamlessM4TAttention_forward, diff --git a/optimum/habana/transformers/models/detr/__init__.py b/optimum/habana/transformers/models/detr/__init__.py index cc6452cf40..d31f2ae55b 100644 --- a/optimum/habana/transformers/models/detr/__init__.py +++ b/optimum/habana/transformers/models/detr/__init__.py @@ -2,7 +2,6 @@ gaudi_DetrConvModel_forward, gaudi_DetrHungarianMatcher_forward, gaudi_DetrLoss_forward, - gaudi_DetrLoss_get_targets_without_no_objects, gaudi_DetrLoss_loss_boxes, gaudi_DetrLoss_loss_cardinality, gaudi_DetrLoss_loss_labels, diff --git a/optimum/habana/transformers/models/detr/modeling_detr.py b/optimum/habana/transformers/models/detr/modeling_detr.py index e23699fbf3..75d6789e49 100644 --- a/optimum/habana/transformers/models/detr/modeling_detr.py +++ b/optimum/habana/transformers/models/detr/modeling_detr.py @@ -1,7 +1,7 @@ import torch from scipy.optimize import linear_sum_assignment from torch import nn -from transformers.models.detr.modeling_detr import center_to_corners_format, generalized_box_iou +from transformers.loss.loss_deformable_detr import center_to_corners_format, generalized_box_iou from transformers.utils import is_accelerate_available @@ -138,6 +138,7 @@ def gaudi_DetrLoss_loss_boxes(self, outputs, targets, indices, num_boxes): losses = {} losses["loss_bbox"] = loss_bbox.sum() / num_boxes + loss_giou = 1 - torch.diag( generalized_box_iou(center_to_corners_format(source_boxes), center_to_corners_format(target_boxes)) ) @@ -153,7 +154,6 @@ def gaudi_DetrLoss_loss_cardinality(self, outputs, targets, indices, num_boxes): """ logits = outputs["logits"] target_lengths = torch.as_tensor([len(v) for v in targets], device="cpu") - # Count the number of predictions that are NOT "no-object" (which is the last class) card_pred = (logits.argmax(-1) != logits.shape[-1] - 1).sum(1) card_err = nn.functional.l1_loss(card_pred.to("cpu").float(), target_lengths.float()) @@ -175,7 +175,7 @@ def gaudi_DetrLoss_forward(self, outputs, targets): # Retrieve the matching between the outputs of the last layer and the targets device = outputs["logits"].device - target_copy = self.get_targets_without_no_objects(targets) + target_copy = self.gaudi_DetrLoss_get_targets_without_no_objects(targets) indices = self.matcher(outputs_without_aux, target_copy) # Compute the average number of target boxes across all nodes, for normalization purposes diff --git a/optimum/habana/transformers/models/qwen2_vl/__init__.py b/optimum/habana/transformers/models/qwen2_vl/__init__.py index 72a587c799..1a22399f10 100644 --- a/optimum/habana/transformers/models/qwen2_vl/__init__.py +++ b/optimum/habana/transformers/models/qwen2_vl/__init__.py @@ -1,9 +1,9 @@ from .modeling_qwen2_vl import ( - GaudiQwen2VisionSdpaAttention, GaudiQwen2VisionTransformerPretrainedModel, GaudiQwen2VLDecoderLayer, GaudiQwen2VLForConditionalGeneration, GaudiQwen2VLModel, GaudiQwen2VLSdpaAttention, GaudiQwen2VLVisionBlock, + GaudiVisionSdpaAttention, ) 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 d2f0706dd6..79d11e9cff 100644 --- a/optimum/habana/transformers/models/qwen2_vl/modeling_qwen2_vl.py +++ b/optimum/habana/transformers/models/qwen2_vl/modeling_qwen2_vl.py @@ -33,7 +33,6 @@ Qwen2VLSdpaAttention, Qwen2VLVisionBlock, VisionSdpaAttention, - _prepare_4d_causal_attention_mask_with_cache_position, apply_multimodal_rotary_pos_emb, apply_rotary_pos_emb_vision, repeat_kv, @@ -60,7 +59,7 @@ 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 GaudiQwen2VisionSdpaAttention(VisionSdpaAttention): +class GaudiVisionSdpaAttention(VisionSdpaAttention): def __init__(self, dim: int, num_heads: int = 16) -> None: super().__init__(dim, num_heads) self.fused_scaled_dot_product_attention = ModuleFusedSDPA(FusedSDPA) if FusedSDPA else None @@ -107,7 +106,7 @@ class GaudiQwen2VLVisionBlock(Qwen2VLVisionBlock): def __init__(self, config, attn_implementation: str = "sdpa") -> None: super().__init__(config, attn_implementation) - self.attn = GaudiQwen2VisionSdpaAttention(config.embed_dim, num_heads=config.num_heads) + self.attn = GaudiVisionSdpaAttention(config.embed_dim, num_heads=config.num_heads) def forward( self, @@ -131,38 +130,6 @@ def forward( return hidden_states -# from: https://github.com/huggingface/transformers/blob/v4.45.2/src/transformers/models/qwen2_vl/modeling_qwen2_vl.py#L1058 -class GaudiQwen2VisionTransformerPretrainedModel(Qwen2VisionTransformerPretrainedModel): - def forward( - self, - hidden_states: torch.Tensor, - grid_thw: torch.Tensor, - use_flash_attention: Optional[bool] = False, - ) -> torch.Tensor: - """ - Copied from https://github.com/huggingface/transformers/blob/53fad641cfdb5105e2470bcf3ef17ea8e25cc300/src/transformers/models/qwen2_vl/modeling_qwen2_vl.py#L1118 - The only differences are: - - add new args use_flash_attention - """ - hidden_states = self.patch_embed(hidden_states) - rotary_pos_emb = self.rot_pos_emb(grid_thw) - - cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum( - dim=0, dtype=torch.int32 - ) - cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0) - - for blk in self.blocks: - hidden_states = blk( - hidden_states, - cu_seqlens=cu_seqlens, - rotary_pos_emb=rotary_pos_emb, - use_flash_attention=use_flash_attention, - ) - - return self.merger(hidden_states) - - # from: https://github.com/huggingface/transformers/blob/v4.45.2/src/transformers/models/qwen2_vl/modeling_qwen2_vl.py#L821 class GaudiQwen2VLSdpaAttention(Qwen2VLSdpaAttention): """ @@ -186,7 +153,7 @@ 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, # will become mandatory in v4.46 + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC use_flash_attention: Optional[bool] = False, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: """ @@ -209,16 +176,18 @@ def forward( output_attentions=output_attentions, use_cache=use_cache, cache_position=cache_position, + position_embeddings=position_embeddings, ) bsz, q_len, _ = hidden_states.size() + query_states = self.q_proj(hidden_states) key_states = self.k_proj(hidden_states) value_states = self.v_proj(hidden_states) - query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) - key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) - value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + query_states = query_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) kv_seq_len = key_states.shape[-2] if past_key_value is not None: @@ -304,7 +273,7 @@ def forward( 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.46 + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC **kwargs, ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: """ @@ -370,6 +339,43 @@ def forward( return outputs +# from: https://github.com/huggingface/transformers/blob/v4.45.2/src/transformers/models/qwen2_vl/modeling_qwen2_vl.py#L1058 +class GaudiQwen2VisionTransformerPretrainedModel(Qwen2VisionTransformerPretrainedModel): + def forward( + self, + hidden_states: torch.Tensor, + grid_thw: torch.Tensor, + use_flash_attention: Optional[bool] = False, + ) -> torch.Tensor: + """ + Copied from https://github.com/huggingface/transformers/blob/53fad641cfdb5105e2470bcf3ef17ea8e25cc300/src/transformers/models/qwen2_vl/modeling_qwen2_vl.py#L1118 + The only differences are: + - add new args use_flash_attention + """ + hidden_states = self.patch_embed(hidden_states) + rotary_pos_emb = self.rot_pos_emb(grid_thw) + + cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum( + dim=0, dtype=torch.int32 + ) + cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0) + + for blk in self.blocks: + if self.gradient_checkpointing and self.training: + hidden_states = self._gradient_checkpointing_func( + blk.__call__, hidden_states, cu_seqlens, rotary_pos_emb, use_flash_attention + ) + else: + hidden_states = blk( + hidden_states, + cu_seqlens=cu_seqlens, + rotary_pos_emb=rotary_pos_emb, + use_flash_attention=use_flash_attention, + ) + + return self.merger(hidden_states) + + # from: https://github.com/huggingface/transformers/blob/v4.45.2/src/transformers/models/qwen2_vl/modeling_qwen2_vl.py#L1137 class GaudiQwen2VLModel(Qwen2VLModel): def forward( @@ -401,9 +407,7 @@ def forward( 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 cannot specify both input_ids and inputs_embeds at the same time, and must specify either one" - ) + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") if self.gradient_checkpointing and self.training: if use_cache: @@ -514,6 +518,7 @@ def forward( image_grid_thw: Optional[torch.LongTensor] = None, video_grid_thw: Optional[torch.LongTensor] = None, rope_deltas: Optional[torch.LongTensor] = None, + cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, use_flash_attention: Optional[bool] = False, ) -> Union[Tuple, Qwen2VLCausalLMOutputWithPast]: @@ -589,20 +594,61 @@ def forward( image_embeds = self.visual( pixel_values, grid_thw=image_grid_thw, use_flash_attention=use_flash_attention ) - image_mask = (input_ids == self.config.image_token_id).unsqueeze(-1).expand_as(inputs_embeds) + n_image_tokens = (input_ids == self.config.image_token_id).sum().item() + n_image_features = image_embeds.shape[0] + if n_image_tokens != n_image_features: + raise ValueError( + f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}" + ) + image_mask = ( + (input_ids == self.config.image_token_id) + .unsqueeze(-1) + .expand_as(inputs_embeds) + .to(inputs_embeds.device) + ) image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype) inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) if pixel_values_videos is not None: pixel_values_videos = pixel_values_videos.type(self.visual.get_dtype()) video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw) - video_mask = (input_ids == self.config.video_token_id).unsqueeze(-1).expand_as(inputs_embeds) + n_video_tokens = (input_ids == self.config.video_token_id).sum().item() + n_video_features = video_embeds.shape[0] + if n_video_tokens != n_video_features: + raise ValueError( + f"Video features and video tokens do not match: tokens: {n_video_tokens}, features {n_video_features}" + ) + video_mask = ( + (input_ids == self.config.video_token_id) + .unsqueeze(-1) + .expand_as(inputs_embeds) + .to(inputs_embeds.device) + ) video_embeds = video_embeds.to(inputs_embeds.device, inputs_embeds.dtype) inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) if attention_mask is not None: attention_mask = attention_mask.to(inputs_embeds.device) + # if we get 4D attention mask we cannot calculate rope deltas anymore. TODO @raushan fixme + if position_ids is None and (attention_mask is None or attention_mask.ndim == 2): + # calculate RoPE index once per generation in the pre-fill stage only + if (cache_position is not None and cache_position[0] == 0) or self.rope_deltas is None: + position_ids, rope_deltas = self.get_rope_index( + input_ids, image_grid_thw, video_grid_thw, attention_mask + ) + self.rope_deltas = rope_deltas + # then use the prev pre-calculated rope-deltas to get the correct position ids + else: + batch_size, seq_length, _ = inputs_embeds.shape + delta = cache_position[0] + self.rope_deltas if cache_position is not None else 0 + position_ids = torch.arange(seq_length, device=inputs_embeds.device) + position_ids = position_ids.view(1, -1).expand(batch_size, -1) + if cache_position is not None: # otherwise `deltas` is an int `0` + delta = delta.repeat_interleave(batch_size // delta.shape[0], dim=0) + position_ids = position_ids.add(delta) + position_ids = position_ids.unsqueeze(0).expand(3, -1, -1) + outputs = self.model( input_ids=None, position_ids=position_ids, @@ -613,15 +659,17 @@ def forward( output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, + cache_position=cache_position, use_flash_attention=use_flash_attention, ) hidden_states = outputs[0] logits = self.lm_head(hidden_states) - logits = logits.float() 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 so that tokens < n predict n shift_logits = logits[..., :-1, :].contiguous() shift_labels = labels[..., 1:].contiguous() @@ -643,7 +691,7 @@ def forward( past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, - rope_deltas=rope_deltas, + rope_deltas=self.rope_deltas, ) def prepare_inputs_for_generation( @@ -688,22 +736,6 @@ def prepare_inputs_for_generation( 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] - rope_deltas = kwargs.get("rope_deltas", None) - if attention_mask is not None and position_ids is None: - if cache_position is None or (cache_position is not None and cache_position[0] == 0): - position_ids, rope_deltas = self.get_rope_index( - input_ids, image_grid_thw, video_grid_thw, attention_mask - ) - else: - batch_size, seq_length = input_ids.shape - delta = ( - cache_position[0] + rope_deltas if cache_position is not None and rope_deltas is not None else 0 - ) - position_ids = torch.arange(seq_length, device=input_ids.device) - position_ids = position_ids.view(1, -1).expand(batch_size, -1) - position_ids = position_ids.add(delta) - position_ids = position_ids.unsqueeze(0).expand(3, -1, -1) - if cache_position[0] != 0: pixel_values = None pixel_values_videos = None @@ -722,18 +754,16 @@ def prepare_inputs_for_generation( batch_size, sequence_length = input_ids.shape device = input_ids.device - dtype = self.lm_head.weight.dtype - min_dtype = torch.finfo(dtype).min - - attention_mask = _prepare_4d_causal_attention_mask_with_cache_position( + attention_mask = self.model._prepare_4d_causal_attention_mask_with_cache_position( attention_mask, sequence_length=sequence_length, - target_length=past_key_values.get_max_length(), - dtype=dtype, + target_length=past_key_values.get_max_cache_shape(), + dtype=self.lm_head.weight.dtype, device=device, - min_dtype=min_dtype, cache_position=cache_position, batch_size=batch_size, + config=self.config, + past_key_values=past_key_values, ) model_inputs.update( @@ -746,7 +776,7 @@ def prepare_inputs_for_generation( "pixel_values_videos": pixel_values_videos, "image_grid_thw": image_grid_thw, "video_grid_thw": video_grid_thw, - "rope_deltas": rope_deltas, + "cache_position": cache_position, "token_idx": token_idx, "use_flash_attention": use_flash_attention, } From c1f30d80b16e714dda24a34ac42d4b5efc9ff526 Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Fri, 31 Jan 2025 15:22:26 +0000 Subject: [PATCH 34/89] Workaround for textual inversion --- .../training/textual_inversion.py | 3 ++ .../training/textual_inversion_sdxl.py | 3 ++ optimum/habana/transformers/modeling_utils.py | 6 ++++ .../modeling_utils_transformers.py | 35 +++++++++++++++++++ 4 files changed, 47 insertions(+) create mode 100644 optimum/habana/transformers/modeling_utils_transformers.py diff --git a/examples/stable-diffusion/training/textual_inversion.py b/examples/stable-diffusion/training/textual_inversion.py index 2f465699b3..2dc0d9d41d 100755 --- a/examples/stable-diffusion/training/textual_inversion.py +++ b/examples/stable-diffusion/training/textual_inversion.py @@ -53,6 +53,7 @@ from optimum.habana import GaudiConfig from optimum.habana.accelerate import GaudiAccelerator from optimum.habana.diffusers import GaudiDDIMScheduler, GaudiStableDiffusionPipeline +from optimum.habana.transformers.modeling_utils import adapt_transformers_to_gaudi from optimum.habana.utils import set_seed @@ -677,6 +678,8 @@ def main(): placeholder_token_ids = tokenizer.convert_tokens_to_ids(placeholder_tokens) # Resize the token embeddings as we are adding new special tokens to the tokenizer + # TODO: remove the call to `adapt_transformers_to_gaudi` once torch.linalg.eigvals is supported on HPU + adapt_transformers_to_gaudi() text_encoder.resize_token_embeddings(len(tokenizer)) # Initialise the newly added placeholder token with the embeddings of the initializer token diff --git a/examples/stable-diffusion/training/textual_inversion_sdxl.py b/examples/stable-diffusion/training/textual_inversion_sdxl.py index 3ab6c57602..da382fbf30 100755 --- a/examples/stable-diffusion/training/textual_inversion_sdxl.py +++ b/examples/stable-diffusion/training/textual_inversion_sdxl.py @@ -52,6 +52,7 @@ from optimum.habana.diffusers import ( GaudiStableDiffusionXLPipeline, ) +from optimum.habana.transformers.modeling_utils import adapt_transformers_to_gaudi from optimum.habana.utils import set_seed @@ -678,6 +679,8 @@ def main(): placeholder_token_ids_2 = tokenizer_2.convert_tokens_to_ids(placeholder_tokens) # Resize the token embeddings as we are adding new special tokens to the tokenizer + # TODO: remove the call to `adapt_transformers_to_gaudi` once torch.linalg.eigvals is supported on HPU + adapt_transformers_to_gaudi() text_encoder_1.resize_token_embeddings(len(tokenizer_1)) text_encoder_2.resize_token_embeddings(len(tokenizer_2)) diff --git a/optimum/habana/transformers/modeling_utils.py b/optimum/habana/transformers/modeling_utils.py index 3dbafa4a5a..cbfdf29743 100644 --- a/optimum/habana/transformers/modeling_utils.py +++ b/optimum/habana/transformers/modeling_utils.py @@ -26,6 +26,7 @@ gaudi_MaxTimeCriteria_call, gaudi_StoppingCriteriaList_call, ) +from .modeling_utils_transformers import _gaudi_init_added_embeddings_weights_with_mean from .models import ( GAUDI_WHISPER_ATTENTION_CLASSES, BaichuanConfig, @@ -759,3 +760,8 @@ def adapt_transformers_to_gaudi(): transformers.loss.loss_for_object_detection.ImageLoss.loss_cardinality = gaudi_DetrLoss_loss_cardinality transformers.loss.loss_for_object_detection.ImageLoss.loss_boxes = gaudi_DetrLoss_loss_boxes transformers.loss.loss_for_object_detection.ImageLoss.forward = gaudi_DetrLoss_forward + + # Workaround for textual inversion + transformers.modeling_utils.PreTrainedModel._init_added_embeddings_weights_with_mean = ( + _gaudi_init_added_embeddings_weights_with_mean + ) diff --git a/optimum/habana/transformers/modeling_utils_transformers.py b/optimum/habana/transformers/modeling_utils_transformers.py new file mode 100644 index 0000000000..d2f1a49d97 --- /dev/null +++ b/optimum/habana/transformers/modeling_utils_transformers.py @@ -0,0 +1,35 @@ +import torch + + +def _gaudi_init_added_embeddings_weights_with_mean( + self, old_embeddings, new_embeddings, old_embedding_dim, old_num_tokens, added_num_tokens +): + """ + Copied from: https://github.com/huggingface/transformers/blob/v4.48.2/src/transformers/modeling_utils.py#L2406 + Changes: + - torch.linalg.eigvals is not supported on HPU so run it on CPU + """ + old_embeddings_weight = old_embeddings.weight.data.to(torch.float32) + mean_embeddings = torch.mean(old_embeddings_weight, axis=0) + old_centered_embeddings = old_embeddings_weight - mean_embeddings + covariance = old_centered_embeddings.T @ old_centered_embeddings / old_num_tokens + + # Check if the covariance is positive definite. + # TODO: do not move `covariance` to the host once torch.linalg.eigvals is supported on HPU + eigenvalues = torch.linalg.eigvals(covariance.to("cpu")) + is_covariance_psd = bool( + (covariance == covariance.T).all() and not torch.is_complex(eigenvalues) and (eigenvalues > 0).all() + ) + if is_covariance_psd: + # If covariances is positive definite, a distribution can be created. and we can sample new weights from it. + distribution = torch.distributions.multivariate_normal.MultivariateNormal( + mean_embeddings, covariance_matrix=1e-9 * covariance + ) + new_embeddings.weight.data[-1 * added_num_tokens :, :] = distribution.sample( + sample_shape=(added_num_tokens,) + ).to(old_embeddings.weight.dtype) + else: + # Otherwise, just initialize with the mean. because distribtion will not be created. + new_embeddings.weight.data[-1 * added_num_tokens :, :] = ( + mean_embeddings[None, :].repeat(added_num_tokens, 1).to(old_embeddings.weight.dtype) + ) From 7eadac6eea439a1bc99ad331f7d7959fc188331e Mon Sep 17 00:00:00 2001 From: Iman Gohari Date: Fri, 31 Jan 2025 14:04:05 -0700 Subject: [PATCH 35/89] Fixes for v4.48 pytest (#1699) --- .../tests/models/falcon/test_modeling_falcon.py | 8 ++++---- .../tests/models/gpt_neox/test_modeling_gpt_neox.py | 12 +++++------- .../tests/models/gptj/test_modeling_gptj.py | 9 --------- 3 files changed, 9 insertions(+), 20 deletions(-) diff --git a/tests/transformers/tests/models/falcon/test_modeling_falcon.py b/tests/transformers/tests/models/falcon/test_modeling_falcon.py index 6d44b2c98b..660f900fea 100644 --- a/tests/transformers/tests/models/falcon/test_modeling_falcon.py +++ b/tests/transformers/tests/models/falcon/test_modeling_falcon.py @@ -52,8 +52,6 @@ FalconModel, ) from transformers.models.falcon.modeling_falcon import ( - FalconDynamicNTKScalingRotaryEmbedding, - FalconLinearScalingRotaryEmbedding, FalconRotaryEmbedding, ) @@ -456,11 +454,12 @@ def test_model_rope_scaling(self): torch.testing.assert_close(original_sin_short, original_sin_long[:short_input_length, :]) # Sanity check linear RoPE scaling # New position "x" should match original position with index "x/scaling_factor" - linear_scaling_rope = FalconLinearScalingRotaryEmbedding( + linear_scaling_rope = FalconRotaryEmbedding( head_dim, max_position_embeddings=config.max_position_embeddings, base=config.rope_theta, scaling_factor=scaling_factor, + rope_type="linear", ).to(torch_device) linear_cos_short, linear_sin_short = linear_scaling_rope(x, short_input_length) linear_cos_long, linear_sin_long = linear_scaling_rope(x, long_input_length) @@ -473,11 +472,12 @@ def test_model_rope_scaling(self): # Sanity check Dynamic NTK RoPE scaling # Scaling should only be observed after a long input is fed. We can observe that the frequencies increase # with scaling_factor (or that `inv_freq` decreases) - ntk_scaling_rope = FalconDynamicNTKScalingRotaryEmbedding( + ntk_scaling_rope = FalconRotaryEmbedding( head_dim, max_position_embeddings=config.max_position_embeddings, base=config.rope_theta, scaling_factor=scaling_factor, + rope_type="dynamic", ).to(torch_device) ntk_cos_short, ntk_sin_short = ntk_scaling_rope(x, short_input_length) ntk_cos_long, ntk_sin_long = ntk_scaling_rope(x, long_input_length) diff --git a/tests/transformers/tests/models/gpt_neox/test_modeling_gpt_neox.py b/tests/transformers/tests/models/gpt_neox/test_modeling_gpt_neox.py index 5026ff87d8..eb5ef0893c 100644 --- a/tests/transformers/tests/models/gpt_neox/test_modeling_gpt_neox.py +++ b/tests/transformers/tests/models/gpt_neox/test_modeling_gpt_neox.py @@ -38,11 +38,7 @@ GPTNeoXForTokenClassification, GPTNeoXModel, ) - from transformers.models.gpt_neox.modeling_gpt_neox import ( - GPTNeoXDynamicNTKScalingRotaryEmbedding, - GPTNeoXLinearScalingRotaryEmbedding, - GPTNeoXRotaryEmbedding, - ) + from transformers.models.gpt_neox.modeling_gpt_neox import GPTNeoXRotaryEmbedding class GPTNeoXModelTester: @@ -371,11 +367,12 @@ def test_model_rope_scaling(self): torch.testing.assert_close(original_sin_short, original_sin_long[:short_input_length, :]) # Sanity check linear RoPE scaling # New position "x" should match original position with index "x/scaling_factor" - linear_scaling_rope = GPTNeoXLinearScalingRotaryEmbedding( + linear_scaling_rope = GPTNeoXRotaryEmbedding( head_dim, max_position_embeddings=config.max_position_embeddings, base=config.rotary_emb_base, scaling_factor=scaling_factor, + rope_type="linear", ).to(torch_device) linear_cos_short, linear_sin_short = linear_scaling_rope(x, short_input_length) linear_cos_long, linear_sin_long = linear_scaling_rope(x, long_input_length) @@ -388,11 +385,12 @@ def test_model_rope_scaling(self): # Sanity check Dynamic NTK RoPE scaling # Scaling should only be observed after a long input is fed. We can observe that the frequencies increase # with scaling_factor (or that `inv_freq` decreases) - ntk_scaling_rope = GPTNeoXDynamicNTKScalingRotaryEmbedding( + ntk_scaling_rope = GPTNeoXRotaryEmbedding( head_dim, max_position_embeddings=config.max_position_embeddings, base=config.rotary_emb_base, scaling_factor=scaling_factor, + rope_type="dynamic", ).to(torch_device) ntk_cos_short, ntk_sin_short = ntk_scaling_rope(x, short_input_length) ntk_cos_long, ntk_sin_long = ntk_scaling_rope(x, long_input_length) diff --git a/tests/transformers/tests/models/gptj/test_modeling_gptj.py b/tests/transformers/tests/models/gptj/test_modeling_gptj.py index bc988d958a..f4c8ad29b6 100644 --- a/tests/transformers/tests/models/gptj/test_modeling_gptj.py +++ b/tests/transformers/tests/models/gptj/test_modeling_gptj.py @@ -43,9 +43,6 @@ GPTJForSequenceClassification, GPTJModel, ) - from transformers.pytorch_utils import is_torch_greater_or_equal_than_1_12 -else: - is_torch_greater_or_equal_than_1_12 = False class GPTJModelTester: @@ -393,16 +390,10 @@ class GPTJModelTest(ModelTesterMixin, GenerationTesterMixin, unittest.TestCase): test_model_parallel = False test_head_masking = False - @unittest.skipIf( - not is_torch_greater_or_equal_than_1_12, reason="PR #22069 made changes that require torch v1.12+." - ) @pytest.mark.skip("Skipped for Gaudi") def test_torch_fx(self): super().test_torch_fx() - @unittest.skipIf( - not is_torch_greater_or_equal_than_1_12, reason="PR #22069 made changes that require torch v1.12+." - ) @pytest.mark.skip("Skipped for Gaudi") def test_torch_fx_output_loss(self): super().test_torch_fx_output_loss() From 5cee21807ed44f9f570aff36c53e1b4ccff8c30b Mon Sep 17 00:00:00 2001 From: Iman Gohari Date: Fri, 31 Jan 2025 15:45:28 -0700 Subject: [PATCH 36/89] fea(): Applied changes in HF #35235 (#1738) --- .../models/falcon/test_modeling_falcon.py | 27 ++++--------------- .../models/gpt_neox/test_modeling_gpt_neox.py | 27 ++++--------------- 2 files changed, 10 insertions(+), 44 deletions(-) diff --git a/tests/transformers/tests/models/falcon/test_modeling_falcon.py b/tests/transformers/tests/models/falcon/test_modeling_falcon.py index 660f900fea..51f7e34cc9 100644 --- a/tests/transformers/tests/models/falcon/test_modeling_falcon.py +++ b/tests/transformers/tests/models/falcon/test_modeling_falcon.py @@ -434,33 +434,21 @@ def test_model_rope_scaling_from_config(self, scaling_type): def test_model_rope_scaling(self): config, _ = self.model_tester.prepare_config_and_inputs_for_common() - hidden_size = config.hidden_size - num_heads = config.num_attention_heads - head_dim = hidden_size // num_heads scaling_factor = 10 short_input_length = 10 long_input_length = int(config.max_position_embeddings * 1.5) # Inputs x = torch.randn(1, dtype=torch.float32, device=torch_device) # used exlusively to get the dtype and the device # Sanity check original RoPE - original_rope = FalconRotaryEmbedding( - head_dim, - max_position_embeddings=config.max_position_embeddings, - base=config.rope_theta, - ).to(torch_device) + original_rope = FalconRotaryEmbedding(config).to(torch_device) original_cos_short, original_sin_short = original_rope(x, short_input_length) original_cos_long, original_sin_long = original_rope(x, long_input_length) torch.testing.assert_close(original_cos_short, original_cos_long[:short_input_length, :]) torch.testing.assert_close(original_sin_short, original_sin_long[:short_input_length, :]) # Sanity check linear RoPE scaling # New position "x" should match original position with index "x/scaling_factor" - linear_scaling_rope = FalconRotaryEmbedding( - head_dim, - max_position_embeddings=config.max_position_embeddings, - base=config.rope_theta, - scaling_factor=scaling_factor, - rope_type="linear", - ).to(torch_device) + config.rope_scaling = {"type": "linear", "factor": scaling_factor} + linear_scaling_rope = FalconRotaryEmbedding(config).to(torch_device) linear_cos_short, linear_sin_short = linear_scaling_rope(x, short_input_length) linear_cos_long, linear_sin_long = linear_scaling_rope(x, long_input_length) torch.testing.assert_close(linear_cos_short, linear_cos_long[:short_input_length, :]) @@ -472,13 +460,8 @@ def test_model_rope_scaling(self): # Sanity check Dynamic NTK RoPE scaling # Scaling should only be observed after a long input is fed. We can observe that the frequencies increase # with scaling_factor (or that `inv_freq` decreases) - ntk_scaling_rope = FalconRotaryEmbedding( - head_dim, - max_position_embeddings=config.max_position_embeddings, - base=config.rope_theta, - scaling_factor=scaling_factor, - rope_type="dynamic", - ).to(torch_device) + config.rope_scaling = {"type": "dynamic", "factor": scaling_factor} + ntk_scaling_rope = FalconRotaryEmbedding(config).to(torch_device) ntk_cos_short, ntk_sin_short = ntk_scaling_rope(x, short_input_length) ntk_cos_long, ntk_sin_long = ntk_scaling_rope(x, long_input_length) torch.testing.assert_close(ntk_cos_short, original_cos_short) diff --git a/tests/transformers/tests/models/gpt_neox/test_modeling_gpt_neox.py b/tests/transformers/tests/models/gpt_neox/test_modeling_gpt_neox.py index eb5ef0893c..905b9474dc 100644 --- a/tests/transformers/tests/models/gpt_neox/test_modeling_gpt_neox.py +++ b/tests/transformers/tests/models/gpt_neox/test_modeling_gpt_neox.py @@ -347,33 +347,21 @@ def test_model_rope_scaling_from_config(self, scaling_type): # Copied from tests.models.falcon.test_modeling_falcon.FalconModelTest.test_model_rope_scaling with Falcon->GPTNeoX, rope_theta->rotary_emb_base def test_model_rope_scaling(self): config, _ = self.model_tester.prepare_config_and_inputs_for_common() - hidden_size = config.hidden_size - num_heads = config.num_attention_heads - head_dim = hidden_size // num_heads scaling_factor = 10 short_input_length = 10 long_input_length = int(config.max_position_embeddings * 1.5) # Inputs x = torch.randn(1, dtype=torch.float32, device=torch_device) # used exlusively to get the dtype and the device # Sanity check original RoPE - original_rope = GPTNeoXRotaryEmbedding( - head_dim, - max_position_embeddings=config.max_position_embeddings, - base=config.rotary_emb_base, - ).to(torch_device) + original_rope = GPTNeoXRotaryEmbedding(config).to(torch_device) original_cos_short, original_sin_short = original_rope(x, short_input_length) original_cos_long, original_sin_long = original_rope(x, long_input_length) torch.testing.assert_close(original_cos_short, original_cos_long[:short_input_length, :]) torch.testing.assert_close(original_sin_short, original_sin_long[:short_input_length, :]) # Sanity check linear RoPE scaling # New position "x" should match original position with index "x/scaling_factor" - linear_scaling_rope = GPTNeoXRotaryEmbedding( - head_dim, - max_position_embeddings=config.max_position_embeddings, - base=config.rotary_emb_base, - scaling_factor=scaling_factor, - rope_type="linear", - ).to(torch_device) + config.rope_scaling = {"type": "linear", "factor": scaling_factor} + linear_scaling_rope = GPTNeoXRotaryEmbedding(config).to(torch_device) linear_cos_short, linear_sin_short = linear_scaling_rope(x, short_input_length) linear_cos_long, linear_sin_long = linear_scaling_rope(x, long_input_length) torch.testing.assert_close(linear_cos_short, linear_cos_long[:short_input_length, :]) @@ -385,13 +373,8 @@ def test_model_rope_scaling(self): # Sanity check Dynamic NTK RoPE scaling # Scaling should only be observed after a long input is fed. We can observe that the frequencies increase # with scaling_factor (or that `inv_freq` decreases) - ntk_scaling_rope = GPTNeoXRotaryEmbedding( - head_dim, - max_position_embeddings=config.max_position_embeddings, - base=config.rotary_emb_base, - scaling_factor=scaling_factor, - rope_type="dynamic", - ).to(torch_device) + config.rope_scaling = {"type": "dynamic", "factor": scaling_factor} + ntk_scaling_rope = GPTNeoXRotaryEmbedding(config).to(torch_device) ntk_cos_short, ntk_sin_short = ntk_scaling_rope(x, short_input_length) ntk_cos_long, ntk_sin_long = ntk_scaling_rope(x, long_input_length) torch.testing.assert_close(ntk_cos_short, original_cos_short) From 17943de147497ac973983276fb63de3d430c98e9 Mon Sep 17 00:00:00 2001 From: Bhargav Date: Wed, 5 Feb 2025 16:19:57 +0530 Subject: [PATCH 37/89] Removing HL_DS_DISTRIBUTED_ATTENTION_SEQ_DIM as it's not needed from SynapseAI 1.20 (#1726) --- examples/language-modeling/README.md | 1 - 1 file changed, 1 deletion(-) diff --git a/examples/language-modeling/README.md b/examples/language-modeling/README.md index 693dd49241..06f09dbda9 100644 --- a/examples/language-modeling/README.md +++ b/examples/language-modeling/README.md @@ -965,7 +965,6 @@ We have added support for [Deepspeed Ulysses](https://github.com/microsoft/DeepS > This feature is still in beta version and may not work out of the box for all transformer model architectures and configurations. ```bash -HL_DS_DISTRIBUTED_ATTENTION_SEQ_DIM=1 \ python3 ../gaudi_spawn.py \ --world_size 8 --use_deepspeed run_lora_clm.py \ --model_name_or_path meta-llama/Llama-3.1-8B \ From d2148196a9eb42dc9105889e66dd369bf5b8951f Mon Sep 17 00:00:00 2001 From: Chetan Kumar Verma <39086835+ckvermaAI@users.noreply.github.com> Date: Wed, 5 Feb 2025 16:21:39 +0530 Subject: [PATCH 38/89] Update DS config to align with recommended settings (#1730) --- examples/language-modeling/llama3_ds_zero1_config.json | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/examples/language-modeling/llama3_ds_zero1_config.json b/examples/language-modeling/llama3_ds_zero1_config.json index b04ef0f0a9..50a1f46b7d 100755 --- a/examples/language-modeling/llama3_ds_zero1_config.json +++ b/examples/language-modeling/llama3_ds_zero1_config.json @@ -8,6 +8,13 @@ }, "gradient_clipping": 1.0, "zero_optimization": { - "stage": 1 + "stage": 1, + "contiguous_gradients": false + }, + "timers": { + "throughput": { + "enabled": true, + "synchronized": false + } } } From 6a520fff5b8169dcbe6923c03881219fa1ac68e1 Mon Sep 17 00:00:00 2001 From: Sheng Yang Date: Wed, 5 Feb 2025 19:10:36 +0800 Subject: [PATCH 39/89] Fix graph breaks in Mixtral (#65) (#1705) --- .../models/mixtral/modeling_mixtral.py | 40 +++++++++---------- 1 file changed, 19 insertions(+), 21 deletions(-) diff --git a/optimum/habana/transformers/models/mixtral/modeling_mixtral.py b/optimum/habana/transformers/models/mixtral/modeling_mixtral.py index 97e9a8026f..d7548a3cfd 100644 --- a/optimum/habana/transformers/models/mixtral/modeling_mixtral.py +++ b/optimum/habana/transformers/models/mixtral/modeling_mixtral.py @@ -20,7 +20,6 @@ """PyTorch Mixtral model.""" -import contextlib import math import os from typing import List, Optional, Tuple, Union @@ -76,18 +75,12 @@ print("Not using HPU fused kernel for apply_rotary_pos_emb") FusedRoPE = None -try: - from habana_frameworks.torch.hpu import sdp_kernel - - SDPContext = True -except ImportError: - SDPContext = False - +deepspeed_available = is_deepspeed_available() logger = logging.get_logger(__name__) def apply_customized_rope(q, k, cos, sin, position_ids, training=True): - if q.device.type == "hpu" and FusedRoPE: + if q.device.type == "hpu" and FusedRoPE is not None: return apply_customized_rope_module(q, k, cos, sin, position_ids, training) else: return apply_rotary_pos_emb(q, k, cos, sin, position_ids) @@ -99,7 +92,7 @@ def gaudi_mixtral_rmsnorm_forward(self, hidden_states): The only differences are: - override RMSNorm with Habana fused RMSNorm """ - if hidden_states.device.type == "hpu" and FusedRMSNorm: + if hidden_states.device.type == "hpu" and FusedRMSNorm is not None: # mixed dtypes are not good for FusedRMSNorm, both inputs need to have same dtype if hidden_states.dtype != self.weight.dtype: orig_dtype = hidden_states.dtype @@ -307,7 +300,7 @@ def forward( else: past_key_value = None - if FusedSDPA: + if FusedSDPA is not None: if query_states.dtype != key_states.dtype: key_states = key_states.type(query_states.dtype) value_states = value_states.type(query_states.dtype) @@ -324,12 +317,17 @@ def forward( ) htcore.mark_step() else: - with ( - sdp_kernel(enable_recompute=flash_attention_recompute) if SDPContext else contextlib.nullcontext() - ): - attn_output = FusedSDPA.apply( - query_states, key_states, value_states, attention_mask, 0.0, False, None - ) + attn_output = FusedSDPA.apply( + query_states, + key_states, + value_states, + attention_mask, + 0.0, + False, + None, + "None", + flash_attention_recompute, + ) else: query_states, key_states, value_states, attention_mask = gaudi_mixtral_repeat_kv( query_states, key_states, value_states, attention_mask, self.num_key_value_groups @@ -353,7 +351,7 @@ def forward( attn_output = self.o_proj(attn_output) - if not output_attentions or FusedSDPA: + if not output_attentions or FusedSDPA is not None: attn_weights = None return attn_output, attn_weights, past_key_value @@ -379,7 +377,7 @@ def gaudi_mixtral_block_sparse_moe_forward(self, hidden_states: torch.Tensor) -> # router_logits: (batch * sequence_length, n_experts) router_logits = self.gate(hidden_states) - if is_deepspeed_available() and (not self.training): + if deepspeed_available and (not self.training): from deepspeed import comm as dist if dist.is_initialized(): @@ -427,7 +425,7 @@ def gaudi_mixtral_block_dynamic_moe_forward(self, hidden_states: torch.Tensor) - # router_logits: (batch * sequence_length, n_experts) router_logits = self.gate(hidden_states) - if is_deepspeed_available() and (not self.training): + if deepspeed_available and (not self.training): from deepspeed import comm as dist if dist.is_initialized(): @@ -453,7 +451,7 @@ def gaudi_mixtral_block_dynamic_moe_forward(self, hidden_states: torch.Tensor) - experts_min=0, experts_max=7, ) - if is_deepspeed_available() and (not self.training): + if deepspeed_available and (not self.training): from deepspeed import comm as dist if dist.is_initialized(): From bedc041f9d4ce1e249be17912422ca111dd85ac2 Mon Sep 17 00:00:00 2001 From: Bhargav Date: Thu, 6 Feb 2025 15:25:34 +0530 Subject: [PATCH 40/89] Add batch dim idx to support latest deepspeed DistributedAttention (#1725) --- .../models/llama/modeling_llama.py | 80 +++++++++++++++++-- 1 file changed, 73 insertions(+), 7 deletions(-) diff --git a/optimum/habana/transformers/models/llama/modeling_llama.py b/optimum/habana/transformers/models/llama/modeling_llama.py index 18867ff8a4..6ab636f565 100755 --- a/optimum/habana/transformers/models/llama/modeling_llama.py +++ b/optimum/habana/transformers/models/llama/modeling_llama.py @@ -430,7 +430,68 @@ def forward(self, cur, dim, idx): return self.update(self.cache, cur, dim, idx, self.inp_seq_len) -def GaudiDistributedAttention(fused_scaled_dot_product_attention, fused_scaled_dot_product_attention_distributed): +class GaudiDistributedAttention(torch.nn.Module): + def __init__( + self, hpu_module_fsdpa: ModuleFusedSDPA, scale, attention_dropout, enable_recompute, flash_attention_fp8 + ): + super().__init__() + self._hpu_module_fsdpa = hpu_module_fsdpa + if parallel_state.sequence_parallel_is_initialized() and parallel_state.get_sequence_parallel_world_size() > 1: + from deepspeed.sequence.layer import DistributedAttention + + self._hpu_module_fsdpa_distributed = DistributedAttention( + self._hpu_module_fsdpa, parallel_state.get_sequence_parallel_group(), 1, 2 + ) + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_mask: torch.Tensor, + dropout_p: float, + is_casual, + scale, + softmax_mode, + recompute_mode, + valid_sequence_lengths, + padding_side="left", + ): + if parallel_state.sequence_parallel_is_initialized() and parallel_state.get_sequence_parallel_world_size() > 1: + return self._hpu_module_fsdpa_distributed( + query, + key, + value, + 0, # As the shape for inputs is [B, N, S, H] + None, + attn_mask, + dropout_p, + is_casual, + scale, + softmax_mode, + recompute_mode, + valid_sequence_lengths, + padding_side, + ) + else: + return self._hpu_module_fsdpa( + query, + key, + value, + attn_mask, + dropout_p, + is_casual, + scale, + softmax_mode, + recompute_mode, + valid_sequence_lengths, + padding_side, + ) + + +def get_gaudi_distributed_attention( + fused_scaled_dot_product_attention, fused_scaled_dot_product_attention_distributed +): if parallel_state.sequence_parallel_is_initialized() and parallel_state.get_sequence_parallel_world_size() > 1: return fused_scaled_dot_product_attention_distributed else: @@ -472,14 +533,19 @@ def __init__(self, config: LlamaConfig, layer_idx: Optional[int] = None): if FusedSDPA else None ) - # https://github.com/microsoft/DeepSpeed/issues/4359 # for all2all comm, Distributed Attention cares about sequence (s) and number of heads (h) dimensions. In HPU, they are at 1 and 2 indices self.fused_scaled_dot_product_attention_distributed = None if parallel_state.sequence_parallel_is_initialized() and parallel_state.get_sequence_parallel_world_size() > 1: - from deepspeed.sequence.layer import DistributedAttention - - self.fused_scaled_dot_product_attention_distributed = DistributedAttention( - self.fused_scaled_dot_product_attention, parallel_state.get_sequence_parallel_group(), 1, 2 + self.fused_scaled_dot_product_attention_distributed = ( + GaudiDistributedAttention( + self.fused_scaled_dot_product_attention, + scale=self.norm_factor, + attention_dropout=self.attention_dropout, + enable_recompute=False, + flash_attention_fp8=getattr(config, "flash_attention_fp8", False), + ) + if FusedSDPA + else None ) def get_k_proj_weight(self): @@ -696,7 +762,7 @@ def pre_attn_forward( kv_seq_len = key_states.shape[-2] else: past_key_value = None - fused_scaled_dot_product_attention = GaudiDistributedAttention( + fused_scaled_dot_product_attention = get_gaudi_distributed_attention( self.fused_scaled_dot_product_attention, self.fused_scaled_dot_product_attention_distributed ) if use_flash_attention and FusedSDPA is not None: From ce57e4042c6a0722db029a0de947d667e766f0cd Mon Sep 17 00:00:00 2001 From: Yaser Afshar Date: Fri, 7 Feb 2025 06:48:59 -0800 Subject: [PATCH 41/89] Add _prepare_inputs_for_generation (#1743) --- .../habana/transformers/generation/utils.py | 163 +++++++++++++++++- optimum/habana/transformers/modeling_utils.py | 3 + 2 files changed, 165 insertions(+), 1 deletion(-) diff --git a/optimum/habana/transformers/generation/utils.py b/optimum/habana/transformers/generation/utils.py index 454ab7da07..3825b33c55 100755 --- a/optimum/habana/transformers/generation/utils.py +++ b/optimum/habana/transformers/generation/utils.py @@ -22,7 +22,14 @@ import torch import torch.distributed as dist -from transformers.cache_utils import Cache, DynamicCache, EncoderDecoderCache, OffloadedCache, QuantizedCacheConfig +from transformers.cache_utils import ( + Cache, + DynamicCache, + EncoderDecoderCache, + OffloadedCache, + QuantizedCacheConfig, + StaticCache, +) from transformers.generation.beam_constraints import DisjunctiveConstraint, PhrasalConstraint from transformers.generation.beam_search import BeamScorer, BeamSearchScorer, ConstrainedBeamSearchScorer from transformers.generation.candidate_generator import ( @@ -170,6 +177,160 @@ class GaudiGenerationMixin(GenerationMixin): sizes allows to make the most of lazy mode and HPU graphs. """ + def _prepare_inputs_for_generation( + self, + input_ids: torch.LongTensor, + past_key_values: Optional[Cache] = None, + attention_mask: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + cache_position: Optional[torch.LongTensor] = None, + **kwargs, + ): + """ + Prepare the model inputs for generation. In includes operations like computing the 4D attention mask or + slicing inputs given the existing cache. + + See the forward pass in the model documentation for expected arguments (different models might have different + requirements for e.g. `past_key_values`). This function should work as is for most LLMs. + + Copied from https://github.com/huggingface/transformers/blob/v4.48.2/src/transformers/generation/utils.py#L349 + Extended with custom modifications to remove keys not used in the forward method. + """ + + # 1. Handle BC: + model_inputs = {} + # - some models don't have `Cache` support (which implies they don't expect `cache_position` in `forward`) + if self._supports_cache_class: + model_inputs["cache_position"] = cache_position + # - `cache_position` was not a mandatory input in `prepare_inputs_for_generation` for those models, and this + # function may be called outside of `generate`. Handle most use cases by creating `cache_position` on the fly + # (this alternative is not as robust as calling `generate` and letting it create `cache_position`) + elif cache_position is None: + past_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0 + 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) + if past_key_values is not None: + model_inputs["past_key_values"] = past_key_values + if ( + 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] + + # 3. Prepare base model inputs + input_ids_key = "decoder_input_ids" if self.config.is_encoder_decoder else "input_ids" + # if `inputs_embeds` are passed, we only want to use them in the 1st generation step + if not self.config.is_encoder_decoder: + if inputs_embeds is not None and cache_position[0] == 0: + model_inputs[input_ids_key] = None + model_inputs["inputs_embeds"] = inputs_embeds + else: + # `clone` calls in this function ensure a consistent stride. See #32227 + model_inputs[input_ids_key] = input_ids.clone(memory_format=torch.contiguous_format) + model_inputs["inputs_embeds"] = None + else: + model_inputs[input_ids_key] = input_ids.clone(memory_format=torch.contiguous_format) + + # 4. Create missing `position_ids` on the fly + if ( + attention_mask is not None + and kwargs.get("position_ids") is None + and "position_ids" in set(inspect.signature(self.forward).parameters.keys()) + ): + position_ids = attention_mask.long().cumsum(-1) - 1 + position_ids.masked_fill_(attention_mask == 0, 1) + kwargs["position_ids"] = position_ids # placed in kwargs for further processing (see below) + + # 5. Slice model inputs if it's an input that should have the same length as `input_ids` + for model_input_name in ["position_ids", "token_type_ids"]: + model_input = kwargs.get(model_input_name) + if model_input is not None: + if past_key_values is not None: + current_input_length = ( + model_inputs["inputs_embeds"].shape[1] + if model_inputs["inputs_embeds"] is not None + else model_inputs[input_ids_key].shape[1] + ) + model_input = model_input[:, -current_input_length:] + 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: + 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 + + # 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: + causal_mask_creation_function = getattr( + base_model, "_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." + ) + else: + attention_mask = causal_mask_creation_function( + attention_mask, + 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, + past_key_values=past_key_values, + ) + if attention_mask is not None: + model_inputs["attention_mask"] = 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: + model_inputs[key] = value + + # 8. Remove unexpected `generate` inputs (TODO @joao: fix trainer and examples) + model_inputs.pop("labels", None) + + # 9. Custom logic to remove unused keys + forward_call = self._slow_forward if torch._C._get_tracing_state() else self.forward + forward_call_signature = inspect.signature(forward_call) + forward_call_has_kwargs = False + for param in forward_call_signature.parameters.values(): + if param.kind == param.VAR_KEYWORD: + forward_call_has_kwargs = True + break + + if not forward_call_has_kwargs: + forward_call_keys = set(forward_call_signature.parameters.keys()) + model_inputs_keys = list(model_inputs.keys()) + for key in model_inputs_keys: + if key not in forward_call_keys: + del model_inputs[key] + + return model_inputs + def _get_hpu_graphs_kwargs(self, model_kwargs): hpu_graphs_kwargs = {} if model_kwargs["limit_hpu_graphs"]: diff --git a/optimum/habana/transformers/modeling_utils.py b/optimum/habana/transformers/modeling_utils.py index fb7bb8e372..c33a58007e 100644 --- a/optimum/habana/transformers/modeling_utils.py +++ b/optimum/habana/transformers/modeling_utils.py @@ -312,6 +312,9 @@ def adapt_transformers_to_gaudi(): transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2SdpaAttention = GaudiWav2Vec2SdpaAttention # Generation is modified to run faster in lazy mode + transformers.generation.GenerationMixin.prepare_inputs_for_generation = ( + GaudiGenerationMixin._prepare_inputs_for_generation + ) transformers.generation.GenerationMixin.generate = GaudiGenerationMixin.generate transformers.generation.GenerationMixin._update_model_kwargs_for_generation = ( GaudiGenerationMixin._update_model_kwargs_for_generation From be34027b211fc4c7f8156e1c0a1ae38a3a6fea3a Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Fri, 7 Feb 2025 18:30:30 +0000 Subject: [PATCH 42/89] Upgrade to v4.48.3 --- .../habana/transformers/modeling_rope_utils.py | 3 +++ .../models/bloom/modeling_bloom.py | 15 +++++++-------- .../models/codegen/modeling_codegen.py | 14 ++++++++------ .../models/falcon/modeling_falcon.py | 15 ++++++--------- .../models/gemma/modeling_gemma.py | 1 + .../models/gemma2/modeling_gemma2.py | 3 +++ .../transformers/models/gpt2/modeling_gpt2.py | 14 +++++++------- .../models/gpt_bigcode/modeling_gpt_bigcode.py | 14 +++++++------- .../models/gpt_neo/modeling_gpt_neo.py | 14 ++++++++------ .../models/gpt_neox/modeling_gpt_neox.py | 1 + .../transformers/models/gptj/modeling_gptj.py | 13 +++++++------ .../models/llama/modeling_llama.py | 3 +++ .../transformers/models/mpt/modeling_mpt.py | 15 +++++++-------- .../transformers/models/opt/modeling_opt.py | 14 +++++++------- .../models/paligemma/modeling_paligemma.py | 5 +++-- .../models/persimmon/modeling_persimmon.py | 18 +++++++----------- .../models/stablelm/modeling_stablelm.py | 18 +++++++----------- .../transformers/models/xglm/modeling_xglm.py | 16 ++++++++-------- 18 files changed, 100 insertions(+), 96 deletions(-) diff --git a/optimum/habana/transformers/modeling_rope_utils.py b/optimum/habana/transformers/modeling_rope_utils.py index 639219c9ab..0a05e51a2f 100644 --- a/optimum/habana/transformers/modeling_rope_utils.py +++ b/optimum/habana/transformers/modeling_rope_utils.py @@ -88,6 +88,9 @@ def _dynamic_frequency_update(self, seq_len, device): self.max_seq_len_cached = seq_len if seq_len < self.original_max_seq_len and self.max_seq_len_cached > self.original_max_seq_len: # reset + # This .to() is needed if the model has been moved to a device after being initialized (because + # the buffer is automatically moved, but not the original copy) + self.original_inv_freq = self.original_inv_freq.to(device) self.register_buffer("inv_freq", self.original_inv_freq, persistent=False) self.max_seq_len_cached = self.original_max_seq_len diff --git a/optimum/habana/transformers/models/bloom/modeling_bloom.py b/optimum/habana/transformers/models/bloom/modeling_bloom.py index 3edab86a60..f36c9cd578 100644 --- a/optimum/habana/transformers/models/bloom/modeling_bloom.py +++ b/optimum/habana/transformers/models/bloom/modeling_bloom.py @@ -21,7 +21,6 @@ from typing import Optional, Tuple, Union import torch -from torch.nn import CrossEntropyLoss from torch.nn import functional as F from transformers.cache_utils import Cache from transformers.modeling_outputs import BaseModelOutputWithPastAndCrossAttentions, CausalLMOutputWithCrossAttentions @@ -544,6 +543,8 @@ def forward( `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100` are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]` """ + # Bloom has deprecated kwargs, so we need to pop num_items_in_batch explicitly + num_items_in_batch = deprecated_arguments.pop("num_items_in_batch", None) 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` warnings.warn( @@ -577,14 +578,12 @@ def forward( if labels is not None: # move labels to correct device to enable model parallelism labels = labels.to(lm_logits.device) - # Shift so that tokens < n predict n - shift_logits = lm_logits[..., :-1, :].contiguous() - shift_labels = labels[..., 1:].contiguous() - batch_size, seq_length, vocab_size = shift_logits.shape # Flatten the tokens - loss_fct = CrossEntropyLoss() - loss = loss_fct( - shift_logits.view(batch_size * seq_length, vocab_size), shift_labels.view(batch_size * seq_length) + loss = self.loss_function( + lm_logits, + labels, + vocab_size=self.config.vocab_size, + num_items_in_batch=num_items_in_batch, ) if not return_dict: diff --git a/optimum/habana/transformers/models/codegen/modeling_codegen.py b/optimum/habana/transformers/models/codegen/modeling_codegen.py index cfe450ab6c..963cead407 100644 --- a/optimum/habana/transformers/models/codegen/modeling_codegen.py +++ b/optimum/habana/transformers/models/codegen/modeling_codegen.py @@ -2,7 +2,6 @@ import torch import torch.utils.checkpoint -from torch.nn import CrossEntropyLoss from transformers.cache_utils import Cache from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from transformers.models.codegen.modeling_codegen import ( @@ -164,6 +163,7 @@ def gaudi_codegen_model_forward( return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, + **kwargs, # NOOP kwargs, for now ) -> Union[Tuple, BaseModelOutputWithPast]: """ Copied from CodeGenBlock.forward: https://github.com/huggingface/transformers/blob/main/src/transformers/models/codegen/modeling_codegen.py @@ -397,6 +397,7 @@ def forward( return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, + **kwargs, ) -> Union[Tuple, CausalLMOutputWithPast]: r""" labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): @@ -432,12 +433,13 @@ def forward( if labels is not None: # move labels to correct device to enable model parallelism labels = labels.to(lm_logits.device) - # Shift so that tokens < n predict n - shift_logits = lm_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)) + loss = self.loss_function( + lm_logits, + labels, + vocab_size=self.config.vocab_size, + **kwargs, + ) loss = loss.to(hidden_states.dtype) diff --git a/optimum/habana/transformers/models/falcon/modeling_falcon.py b/optimum/habana/transformers/models/falcon/modeling_falcon.py index ddc52a4a74..508fab27af 100644 --- a/optimum/habana/transformers/models/falcon/modeling_falcon.py +++ b/optimum/habana/transformers/models/falcon/modeling_falcon.py @@ -27,7 +27,6 @@ import habana_frameworks.torch.core as htcore from torch import nn -from torch.nn import CrossEntropyLoss from torch.nn import functional as F from transformers.cache_utils import Cache from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask_for_sdpa @@ -1040,6 +1039,7 @@ def forward( use_flash_attention: Optional[bool] = False, flash_attention_recompute: Optional[bool] = False, flash_attention_causal_mask: Optional[bool] = False, + **kwargs, ) -> Union[Tuple[torch.Tensor], CausalLMOutputWithCrossAttentions]: r""" labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): @@ -1094,14 +1094,11 @@ def forward( loss = None if labels is not None: - # Shift so that tokens < n predict n - shift_logits = lm_logits[..., :-1, :].contiguous() - shift_labels = labels[..., 1:].contiguous() - batch_size, seq_length, vocab_size = shift_logits.shape - # Flatten the tokens - loss_fct = CrossEntropyLoss() - loss = loss_fct( - shift_logits.view(batch_size * seq_length, vocab_size), shift_labels.view(batch_size * seq_length) + loss = self.loss_function( + lm_logits, + labels, + vocab_size=self.config.vocab_size, + **kwargs, ) if not return_dict: diff --git a/optimum/habana/transformers/models/gemma/modeling_gemma.py b/optimum/habana/transformers/models/gemma/modeling_gemma.py index a4de41d29a..d2d4209d0e 100755 --- a/optimum/habana/transformers/models/gemma/modeling_gemma.py +++ b/optimum/habana/transformers/models/gemma/modeling_gemma.py @@ -603,6 +603,7 @@ def forward( flash_attention_causal_mask: Optional[bool] = False, cache_idx: int = None, lazy_mode: Optional[bool] = True, + **kwargs, # NOOP kwarg for now ) -> Union[Tuple, BaseModelOutputWithPast]: """ Copied from GemmaModel.forward: https://github.com/huggingface/transformers/blob/v4.38.1/src/transformers/models/gemma/modeling_gemma.py diff --git a/optimum/habana/transformers/models/gemma2/modeling_gemma2.py b/optimum/habana/transformers/models/gemma2/modeling_gemma2.py index 505c8c3ac3..7178d8f970 100755 --- a/optimum/habana/transformers/models/gemma2/modeling_gemma2.py +++ b/optimum/habana/transformers/models/gemma2/modeling_gemma2.py @@ -143,6 +143,9 @@ def _dynamic_frequency_update(self, seq_len, device): self.max_seq_len_cached = seq_len if seq_len < self.original_max_seq_len and self.max_seq_len_cached > self.original_max_seq_len: # reset + # This .to() is needed if the model has been moved to a device after being initialized (because + # the buffer is automatically moved, but not the original copy) + self.original_inv_freq = self.original_inv_freq.to(device) self.register_buffer("inv_freq", self.original_inv_freq, persistent=False) self.max_seq_len_cached = self.original_max_seq_len diff --git a/optimum/habana/transformers/models/gpt2/modeling_gpt2.py b/optimum/habana/transformers/models/gpt2/modeling_gpt2.py index 546ee7ef47..e42a8308fa 100644 --- a/optimum/habana/transformers/models/gpt2/modeling_gpt2.py +++ b/optimum/habana/transformers/models/gpt2/modeling_gpt2.py @@ -516,6 +516,7 @@ def forward( output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, token_idx: Optional[torch.Tensor] = None, + **kwargs, ) -> Union[Tuple, CausalLMOutputWithCrossAttentions]: return_dict = return_dict if return_dict is not None else self.config.use_return_dict @@ -546,14 +547,13 @@ def forward( loss = None if labels is not None: - # move labels to correct device to enable model parallelism - labels = labels.to(lm_logits.device) - # Shift so that tokens < n predict n - shift_logits = lm_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)) + loss = self.loss_function( + lm_logits, + labels, + vocab_size=self.config.vocab_size, + **kwargs, + ) if not return_dict: output = (lm_logits,) + transformer_outputs[1:] 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 f01255624f..608c272135 100644 --- a/optimum/habana/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py +++ b/optimum/habana/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py @@ -22,7 +22,6 @@ import torch import torch.nn.functional as F import torch.utils.checkpoint -from torch.nn import CrossEntropyLoss from transformers.modeling_outputs import BaseModelOutputWithPastAndCrossAttentions, CausalLMOutputWithCrossAttentions from transformers.models.gpt_bigcode.modeling_gpt_bigcode import ( GPTBigCodeAttention, @@ -806,6 +805,7 @@ def forward( flash_attention_fast_softmax: Optional[bool] = False, flash_attention_causal_mask: Optional[bool] = False, cache_idx: Optional[int] = None, + **kwargs, ) -> Union[Tuple, CausalLMOutputWithCrossAttentions]: r""" labels (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): @@ -842,12 +842,12 @@ def forward( loss = None if labels is not None: - # Shift so that tokens < n predict n - shift_logits = lm_logits[..., :-1, :].contiguous() - shift_labels = labels[..., 1:].contiguous().to(shift_logits.device) - # Flatten the tokens - loss_fct = CrossEntropyLoss() - loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)) + loss = self.loss_function( + lm_logits, + labels, + vocab_size=self.config.vocab_size, + **kwargs, + ) if not return_dict: output = (lm_logits,) + transformer_outputs[1:] 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 b5ef987752..1cb65bffd0 100644 --- a/optimum/habana/transformers/models/gpt_neo/modeling_gpt_neo.py +++ b/optimum/habana/transformers/models/gpt_neo/modeling_gpt_neo.py @@ -1,7 +1,6 @@ from typing import Optional, Tuple, Union import torch -from torch.nn import CrossEntropyLoss from transformers.modeling_outputs import ( BaseModelOutputWithPast, BaseModelOutputWithPastAndCrossAttentions, @@ -305,7 +304,9 @@ def 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, + **kwargs, ) -> Union[Tuple[torch.Tensor], CausalLMOutputWithCrossAttentions]: r""" labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): @@ -341,12 +342,13 @@ def forward( # https://github.com/EleutherAI/gpt-neo/blob/89ce74164da2fb16179106f54e2269b5da8db333/models/gpt2/gpt2.py#L179 lm_logits = lm_logits.to(torch.float32) - # Shift so that tokens < n predict n - shift_logits = lm_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)) + loss = self.loss_function( + lm_logits, + labels, + vocab_size=self.config.vocab_size, + **kwargs, + ) lm_logits = lm_logits.to(hidden_states.dtype) loss = loss.to(hidden_states.dtype) 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 4f4a152c67..dd41d7b557 100644 --- a/optimum/habana/transformers/models/gpt_neox/modeling_gpt_neox.py +++ b/optimum/habana/transformers/models/gpt_neox/modeling_gpt_neox.py @@ -408,6 +408,7 @@ def forward( return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, + **kwargs, # Unused for now, mostly for the loss correction ) -> Union[Tuple, CausalLMOutputWithPast]: return_dict = return_dict if return_dict is not None else self.config.use_return_dict diff --git a/optimum/habana/transformers/models/gptj/modeling_gptj.py b/optimum/habana/transformers/models/gptj/modeling_gptj.py index d4da76d6f2..a719dc645a 100644 --- a/optimum/habana/transformers/models/gptj/modeling_gptj.py +++ b/optimum/habana/transformers/models/gptj/modeling_gptj.py @@ -3,7 +3,6 @@ import habana_frameworks.torch.core as htcore import torch from torch import nn -from torch.nn import CrossEntropyLoss from transformers.cache_utils import Cache from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from transformers.models.gptj.configuration_gptj import GPTJConfig @@ -662,6 +661,7 @@ def forward( token_idx: Optional[torch.Tensor] = None, reuse_cache: Optional[bool] = False, cache_idx: Optional[int] = None, + **kwargs, ) -> Union[Tuple, CausalLMOutputWithPast]: r""" labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): @@ -699,12 +699,13 @@ def forward( if labels is not None: # move labels to correct device to enable model parallelism labels = labels.to(lm_logits.device) - # Shift so that tokens < n predict n - shift_logits = lm_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)) + loss = self.loss_function( + lm_logits, + labels, + vocab_size=self.config.vocab_size, + **kwargs, + ) loss = loss.to(hidden_states.dtype) diff --git a/optimum/habana/transformers/models/llama/modeling_llama.py b/optimum/habana/transformers/models/llama/modeling_llama.py index ce795c0cd8..e10d9e683e 100755 --- a/optimum/habana/transformers/models/llama/modeling_llama.py +++ b/optimum/habana/transformers/models/llama/modeling_llama.py @@ -149,6 +149,9 @@ def _dynamic_frequency_update(self, seq_len, device): self.max_seq_len_cached = seq_len if seq_len < self.original_max_seq_len and self.max_seq_len_cached > self.original_max_seq_len: # reset + # This .to() is needed if the model has been moved to a device after being initialized (because + # the buffer is automatically moved, but not the original copy) + self.original_inv_freq = self.original_inv_freq.to(device) self.register_buffer("inv_freq", self.original_inv_freq, persistent=False) self.max_seq_len_cached = self.original_max_seq_len diff --git a/optimum/habana/transformers/models/mpt/modeling_mpt.py b/optimum/habana/transformers/models/mpt/modeling_mpt.py index 309e0d7acc..7219ac0f29 100755 --- a/optimum/habana/transformers/models/mpt/modeling_mpt.py +++ b/optimum/habana/transformers/models/mpt/modeling_mpt.py @@ -19,7 +19,6 @@ import torch from torch import nn -from torch.nn import CrossEntropyLoss from transformers.modeling_outputs import BaseModelOutputWithPastAndCrossAttentions, CausalLMOutputWithCrossAttentions from transformers.models.mpt.modeling_mpt import ( MptAttention, @@ -244,6 +243,7 @@ def forward( 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]: """ Copied from MptModel.forward: https://github.com/huggingface/transformers/blob/v4.32.0/src/transformers/models/mpt/modeling_mpt.py @@ -444,6 +444,7 @@ def forward( use_flash_attention: Optional[bool] = False, flash_attention_recompute: Optional[bool] = False, cache_idx: Optional[torch.Tensor] = None, + **kwargs, ) -> Union[Tuple[torch.Tensor], CausalLMOutputWithCrossAttentions]: """ Inherits from MptForCausalLM: https://github.com/huggingface/transformers/blob/v4.32.0/src/transformers/models/mpt/modeling_mpt.py @@ -477,14 +478,12 @@ def forward( if labels is not None: # move labels to correct device to enable model parallelism labels = labels.to(lm_logits.device) - # Shift so that tokens < n predict n - shift_logits = lm_logits[..., :-1, :].contiguous() - shift_labels = labels[..., 1:].contiguous() - batch_size, seq_length, vocab_size = shift_logits.shape # Flatten the tokens - loss_fct = CrossEntropyLoss() - loss = loss_fct( - shift_logits.view(batch_size * seq_length, vocab_size), shift_labels.view(batch_size * seq_length) + loss = self.loss_function( + lm_logits, + labels, + vocab_size=self.config.vocab_size, + **kwargs, ) if not return_dict: diff --git a/optimum/habana/transformers/models/opt/modeling_opt.py b/optimum/habana/transformers/models/opt/modeling_opt.py index 3a7c99d96e..0d7afa4de8 100644 --- a/optimum/habana/transformers/models/opt/modeling_opt.py +++ b/optimum/habana/transformers/models/opt/modeling_opt.py @@ -1,7 +1,6 @@ from typing import List, Optional, Tuple, Union import torch -from torch.nn import CrossEntropyLoss from transformers.activations import ACT2FN from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from transformers.models.opt.configuration_opt import OPTConfig @@ -496,6 +495,7 @@ def forward( return_dict: Optional[bool] = None, position_ids: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, + **kwargs, ) -> Union[Tuple, CausalLMOutputWithPast]: output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = ( @@ -524,12 +524,12 @@ def forward( if labels is not None: # move labels to correct device to enable model parallelism labels = labels.to(logits.device) - # Shift so that tokens < n predict n - shift_logits = logits[..., :-1, :].contiguous() - shift_labels = labels[..., 1:].contiguous() - # Flatten the tokens - loss_fct = CrossEntropyLoss() - loss = loss_fct(shift_logits.view(-1, self.config.vocab_size), shift_labels.view(-1)) + loss = self.loss_function( + logits, + labels, + vocab_size=self.config.vocab_size, + **kwargs, + ) if not return_dict: output = (logits,) + outputs[1:] diff --git a/optimum/habana/transformers/models/paligemma/modeling_paligemma.py b/optimum/habana/transformers/models/paligemma/modeling_paligemma.py index 1d2db48d41..6f2a2817d0 100644 --- a/optimum/habana/transformers/models/paligemma/modeling_paligemma.py +++ b/optimum/habana/transformers/models/paligemma/modeling_paligemma.py @@ -48,7 +48,7 @@ def forward( return_dict: Optional[bool] = None, num_logits_to_keep: int = 0, token_idx: Optional[torch.Tensor] = None, - **kwargs, + **lm_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 @@ -109,7 +109,7 @@ def forward( labels = torch.where(input_ids == self.pad_token_id, self.config.ignore_index, labels) causal_mask = self._update_causal_mask( - attention_mask, token_type_ids, past_key_values, cache_position, input_ids, inputs_embeds, is_training + attention_mask, token_type_ids, past_key_values, cache_position, inputs_embeds, is_training ) outputs = self.language_model( attention_mask=causal_mask, @@ -124,6 +124,7 @@ def forward( # TODO: from Transformers v4.45, `generate` sets `num_logits_to_keep` to 1 if not given, which we don't want here # num_logits_to_keep=num_logits_to_keep, token_idx=token_idx, + **lm_kwargs, ) logits = outputs.logits diff --git a/optimum/habana/transformers/models/persimmon/modeling_persimmon.py b/optimum/habana/transformers/models/persimmon/modeling_persimmon.py index 3e56f3c9e2..62fbe16f3c 100644 --- a/optimum/habana/transformers/models/persimmon/modeling_persimmon.py +++ b/optimum/habana/transformers/models/persimmon/modeling_persimmon.py @@ -3,7 +3,6 @@ import torch from torch import nn -from torch.nn import CrossEntropyLoss from transformers.cache_utils import Cache, DynamicCache from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from transformers.models.persimmon.configuration_persimmon import PersimmonConfig @@ -365,6 +364,7 @@ def forward( cache_position: Optional[torch.LongTensor] = None, num_logits_to_keep: int = 0, token_idx: Optional[torch.Tensor] = None, + **kwargs, ) -> Union[Tuple, CausalLMOutputWithPast]: """ Inherits from PersimmonForCausalLM: https://github.com/huggingface/transformers/blob/v4.38.2/src/transformers/models/persimmon/modeling_persimmon.py @@ -399,16 +399,12 @@ def forward( loss = None if labels is not None: - # Shift so that tokens < n predict n - shift_logits = logits[..., :-1, :].contiguous() - shift_labels = labels[..., 1:].contiguous() - # Flatten the tokens - loss_fct = CrossEntropyLoss() - shift_logits = shift_logits.view(-1, self.config.vocab_size) - shift_labels = shift_labels.view(-1) - # Enable model parallelism - shift_labels = shift_labels.to(shift_logits.device) - loss = loss_fct(shift_logits, shift_labels) + loss = self.loss_function( + logits, + labels, + vocab_size=self.config.vocab_size, + **kwargs, + ) if not return_dict: output = (logits,) + outputs[1:] diff --git a/optimum/habana/transformers/models/stablelm/modeling_stablelm.py b/optimum/habana/transformers/models/stablelm/modeling_stablelm.py index 97a78077d7..7457b8f886 100644 --- a/optimum/habana/transformers/models/stablelm/modeling_stablelm.py +++ b/optimum/habana/transformers/models/stablelm/modeling_stablelm.py @@ -3,7 +3,6 @@ import torch from torch import nn -from torch.nn import CrossEntropyLoss from transformers.cache_utils import Cache, DynamicCache from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from transformers.models.stablelm.configuration_stablelm import StableLmConfig @@ -384,6 +383,7 @@ def forward( cache_position: Optional[torch.LongTensor] = None, num_logits_to_keep: int = 0, token_idx: Optional[torch.Tensor] = None, + **kwargs, ) -> Union[Tuple, CausalLMOutputWithPast]: """ Inherits from StableLmForCausalLM: https://github.com/huggingface/transformers/blob/v4.38.2/src/transformers/models/stablelm/modeling_stablelm.py @@ -416,16 +416,12 @@ def forward( loss = None if labels is not None: - # Shift so that tokens < n predict n - shift_logits = logits[..., :-1, :].contiguous() - shift_labels = labels[..., 1:].contiguous() - # Flatten the tokens - loss_fct = CrossEntropyLoss() - shift_logits = shift_logits.view(-1, self.config.vocab_size) - shift_labels = shift_labels.view(-1) - # Enable model parallelism - shift_labels = shift_labels.to(shift_logits.device) - loss = loss_fct(shift_logits, shift_labels) + loss = self.loss_function( + logits, + labels, + vocab_size=self.config.vocab_size, + **kwargs, + ) if not return_dict: output = (logits,) + outputs[1:] diff --git a/optimum/habana/transformers/models/xglm/modeling_xglm.py b/optimum/habana/transformers/models/xglm/modeling_xglm.py index f69eb3b990..289e0eb55f 100644 --- a/optimum/habana/transformers/models/xglm/modeling_xglm.py +++ b/optimum/habana/transformers/models/xglm/modeling_xglm.py @@ -2,7 +2,6 @@ import torch from torch import nn -from torch.nn import CrossEntropyLoss from transformers.modeling_outputs import BaseModelOutputWithPastAndCrossAttentions, CausalLMOutputWithCrossAttentions from transformers.models.xglm.modeling_xglm import XGLMForCausalLM from transformers.utils import logging @@ -405,6 +404,7 @@ def forward( output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, token_idx: Optional[torch.Tensor] = None, + **kwargs, ) -> Union[Tuple[torch.Tensor], CausalLMOutputWithCrossAttentions]: """ Inherits from XGLMForCausalLM: https://github.com/huggingface/transformers/blob/v4.44.1/src/transformers/models/xglm/modeling_xglm.py @@ -440,13 +440,13 @@ def forward( loss = None if labels is not None: - # shift labels and add a pad token to the end - shift_labels = labels.new_zeros(labels.shape) - shift_labels[:, :-1] = labels[:, 1:].clone() - shift_labels[:, -1] = self.config.pad_token_id - - loss_fct = CrossEntropyLoss() - loss = loss_fct(logits.view(-1, self.config.vocab_size), shift_labels.view(-1)) + loss = self.loss_function( + logits, + labels, + vocab_size=self.config.vocab_size, + pad_token_id=self.config.pad_token_id, + **kwargs, + ) if not return_dict: output = (logits,) + outputs[1:] From bd9a60ecc7b0ea43ebdda38170d040811183c9af Mon Sep 17 00:00:00 2001 From: Sun Choi Date: Fri, 7 Feb 2025 13:05:39 -0800 Subject: [PATCH 43/89] Fix the issue with --load_quantized_model_with_autoawq (#1747) Co-authored-by: regisss <15324346+regisss@users.noreply.github.com> --- Makefile | 2 +- examples/text-generation/README.md | 2 +- examples/text-generation/requirements_awq.txt | 3 +++ optimum/habana/transformers/integrations/awq.py | 2 +- 4 files changed, 6 insertions(+), 3 deletions(-) create mode 100644 examples/text-generation/requirements_awq.txt diff --git a/Makefile b/Makefile index 80fb7b8c62..24ef8476ab 100644 --- a/Makefile +++ b/Makefile @@ -107,7 +107,7 @@ slow_tests_diffusers: test_installs # Run text-generation non-regression tests slow_tests_text_generation_example: test_installs - python -m pip install triton==3.1.0 autoawq + python -m pip install -r examples/text-generation/requirements_awq.txt BUILD_CUDA_EXT=0 python -m pip install -vvv --no-build-isolation git+https://github.com/HabanaAI/AutoGPTQ.git python -m pip install git+https://github.com/HabanaAI/DeepSpeed.git@1.19.0 python -m pytest tests/test_text_generation_example.py tests/test_encoder_decoder.py -v -s --token $(TOKEN) diff --git a/examples/text-generation/README.md b/examples/text-generation/README.md index a884877bee..699cb55dca 100755 --- a/examples/text-generation/README.md +++ b/examples/text-generation/README.md @@ -734,7 +734,7 @@ Currently, this support is limited to UINT4 inference of pre-quantized models on Please run the following command to install AutoAWQ: ```bash -pip install triton==3.1.0 autoawq +pip install -r requirements_awq.txt ``` You can run a *UINT4 weight quantized* model using AutoAWQ by including the argument `--load_quantized_model_with_autoawq`. diff --git a/examples/text-generation/requirements_awq.txt b/examples/text-generation/requirements_awq.txt new file mode 100644 index 0000000000..dff2632403 --- /dev/null +++ b/examples/text-generation/requirements_awq.txt @@ -0,0 +1,3 @@ +triton==3.1.0 +autoawq +transformers>=4.48.2,<4.49.0 diff --git a/optimum/habana/transformers/integrations/awq.py b/optimum/habana/transformers/integrations/awq.py index 7ad1cd454c..a816ddbb1d 100644 --- a/optimum/habana/transformers/integrations/awq.py +++ b/optimum/habana/transformers/integrations/awq.py @@ -168,7 +168,7 @@ def post_init_awq_gemm_hpu_modules(model): return model -def gaudi_awq_quantizer_process_model_after_weight_loading(self, model): +def gaudi_awq_quantizer_process_model_after_weight_loading(self, model, **kwargs): if self.quantization_config.version == GaudiAWQLinearVersion.HPU: model = post_init_awq_gemm_hpu_modules(model) else: From 2f665e8f0a7c150aff23d4bdaa711aa4cc6bff40 Mon Sep 17 00:00:00 2001 From: "Wang, Yi" Date: Wed, 12 Feb 2025 17:22:53 +0800 Subject: [PATCH 44/89] Fix dpo crash in transformers 4.48 (#1750) Signed-off-by: Wang, Yi A --- optimum/habana/trl/trainer/dpo_trainer.py | 39 ++++++++++++++++++++++- 1 file changed, 38 insertions(+), 1 deletion(-) diff --git a/optimum/habana/trl/trainer/dpo_trainer.py b/optimum/habana/trl/trainer/dpo_trainer.py index 3af14d6555..d57a032983 100644 --- a/optimum/habana/trl/trainer/dpo_trainer.py +++ b/optimum/habana/trl/trainer/dpo_trainer.py @@ -15,7 +15,9 @@ import inspect import warnings from collections import defaultdict -from typing import Callable, Dict, List, Literal, Optional, Tuple, Union +from contextlib import nullcontext +from functools import partial +from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union import torch import torch.nn as nn @@ -681,3 +683,38 @@ def log(self, logs: Dict[str, float], **kwargs) -> None: logs[key] = torch.tensor(metrics).mean().item() del self._stored_metrics[train_eval] return super().log(logs) + + def compute_loss( + self, + model: Union[PreTrainedModel, nn.Module], + inputs: Dict[str, Union[torch.Tensor, Any]], + return_outputs=False, + num_items_in_batch=None, + ) -> Union[torch.Tensor, Tuple[torch.Tensor, Dict[str, torch.Tensor]]]: + """ + Copied from DPOTrainer.compute_loss: https://github.com/huggingface/trl/blob/v0.9.6/trl/trainer/dpo_trainer.py#L1393 + - add num_items_in_batch to work with transformers 4.48 + - use hpu autocast + """ + if not self.use_dpo_data_collator: + warnings.warn( + "compute_loss is only implemented for DPODataCollatorWithPadding, and you passed a datacollator that is different than " + "DPODataCollatorWithPadding - you might see unexpected behavior. Alternatively, you can implement your own prediction_step method if you are using a custom data collator" + ) + compute_loss_context_manager = ( + partial(torch.autocast, device_type="hpu", dtype=torch.bfloat16) + if self._peft_has_been_casted_to_bf16 + else nullcontext + ) + + with compute_loss_context_manager(): + loss, metrics = self.get_batch_loss_metrics(model, inputs, train_eval="train") + + # Make sure to move the loss to the device the original accumulating loss is at back in the `Trainer` class: + loss = loss.to(self.args.device) + # force log the metrics + self.store_metrics(metrics, train_eval="train") + + if return_outputs: + return (loss, metrics) + return loss From 595b816504d1a62eff6a73ee3d86906e9fe51b23 Mon Sep 17 00:00:00 2001 From: Sun Choi Date: Wed, 12 Feb 2025 01:35:46 -0800 Subject: [PATCH 45/89] Fix for Falcon image-to-text crash (#1760) --- examples/image-to-text/run_pipeline.py | 1 + 1 file changed, 1 insertion(+) diff --git a/examples/image-to-text/run_pipeline.py b/examples/image-to-text/run_pipeline.py index de849e3469..fe59fbcd5c 100644 --- a/examples/image-to-text/run_pipeline.py +++ b/examples/image-to-text/run_pipeline.py @@ -328,6 +328,7 @@ def main(): if "falcon-11B-vlm" in args.model_name_or_path: # WA falcon vlm issue that image_token_id == embed size. generator.model.resize_token_embeddings(generator.tokenizer.vocab_size + 1) + processor.patch_size = config.vision_config.patch_size generate_kwargs = { "lazy_mode": True, "hpu_graphs": args.use_hpu_graphs, From f3729a438fc83fba15a8261f83a3f2e55af883f3 Mon Sep 17 00:00:00 2001 From: Akihiro Takahashi Date: Wed, 12 Feb 2025 02:24:22 -0800 Subject: [PATCH 46/89] Fix llama attr (#1771) --- optimum/habana/transformers/models/llama/modeling_llama.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/optimum/habana/transformers/models/llama/modeling_llama.py b/optimum/habana/transformers/models/llama/modeling_llama.py index a9ffc4ed73..66610b9fc6 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 +import math from typing import List, Optional, Tuple, Union import torch @@ -485,10 +486,11 @@ def __init__(self, config: LlamaConfig, layer_idx: Optional[int] = None): self.k_proj = None self.v_proj = None self.inp_seq_len = -1 + self.norm_factor = 1.0 / math.sqrt(self.head_dim) self.fused_scaled_dot_product_attention = ( ModuleFusedSDPA( FusedSDPA, - scale=self.scaling, + scale=self.norm_factor, attention_dropout=self.attention_dropout, enable_recompute=False, flash_attention_fp8=getattr(config, "flash_attention_fp8", False), From bcb0778e78e290c8d5f4f5258083d42dd3cb23ce Mon Sep 17 00:00:00 2001 From: Akihiro Takahashi Date: Wed, 12 Feb 2025 14:58:08 -0800 Subject: [PATCH 47/89] Update llama scaling (#1775) Use super class self.scaling for scale. --- optimum/habana/transformers/models/llama/modeling_llama.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/optimum/habana/transformers/models/llama/modeling_llama.py b/optimum/habana/transformers/models/llama/modeling_llama.py index 66610b9fc6..43eb83f2ed 100755 --- a/optimum/habana/transformers/models/llama/modeling_llama.py +++ b/optimum/habana/transformers/models/llama/modeling_llama.py @@ -1,5 +1,4 @@ import copy -import math from typing import List, Optional, Tuple, Union import torch @@ -486,11 +485,10 @@ def __init__(self, config: LlamaConfig, layer_idx: Optional[int] = None): self.k_proj = None self.v_proj = None self.inp_seq_len = -1 - self.norm_factor = 1.0 / math.sqrt(self.head_dim) self.fused_scaled_dot_product_attention = ( ModuleFusedSDPA( FusedSDPA, - scale=self.norm_factor, + scale=self.scaling, attention_dropout=self.attention_dropout, enable_recompute=False, flash_attention_fp8=getattr(config, "flash_attention_fp8", False), @@ -504,7 +502,7 @@ def __init__(self, config: LlamaConfig, layer_idx: Optional[int] = None): self.fused_scaled_dot_product_attention_distributed = ( GaudiDistributedAttention( self.fused_scaled_dot_product_attention, - scale=self.norm_factor, + scale=self.scaling, attention_dropout=self.attention_dropout, enable_recompute=False, flash_attention_fp8=getattr(config, "flash_attention_fp8", False), From d053218304dda4f4d569a905915f23208b9f0847 Mon Sep 17 00:00:00 2001 From: Edward Mascarenhas Date: Fri, 14 Feb 2025 10:23:25 -0800 Subject: [PATCH 48/89] Fix loss calculation (Workaround), final fix TBD (#1784) --- optimum/habana/transformers/trainer.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/optimum/habana/transformers/trainer.py b/optimum/habana/transformers/trainer.py index a51bb49a89..6f186e521c 100644 --- a/optimum/habana/transformers/trainer.py +++ b/optimum/habana/transformers/trainer.py @@ -1645,7 +1645,9 @@ def training_step( 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: + if (not self.model_accepts_loss_kwargs and self.compute_loss_func is None) or (num_items_in_batch is None): + # TODO refer to todo in function get_batch_samples_transformers - + # temporary fix to calculate loss correctly loss = loss / self.args.gradient_accumulation_steps if _is_peft_model(self.model) and self.model.peft_type == PeftType.ADALORA: From 8b006c4a73ff041f750f2d1fb398b50bae34c76f Mon Sep 17 00:00:00 2001 From: Libin Tang Date: Tue, 18 Feb 2025 09:32:51 -0800 Subject: [PATCH 49/89] Simplify text-gen readme (#1780) Co-authored-by: Sayantan Sarkar --- examples/text-generation/README.md | 144 ++--------------------------- 1 file changed, 7 insertions(+), 137 deletions(-) diff --git a/examples/text-generation/README.md b/examples/text-generation/README.md index 78dcd44c30..4732ca1877 100755 --- a/examples/text-generation/README.md +++ b/examples/text-generation/README.md @@ -145,48 +145,7 @@ python ../gaudi_spawn.py --use_deepspeed --world_size 8 run_generation.py \ --sdp_on_bf16 ``` -You can also run Llama2-70B on Gaudi2 with all optimizations enabled using the following command: -```bash -python ../gaudi_spawn.py --use_deepspeed --world_size 8 run_generation.py \ ---model_name_or_path meta-llama/Llama-2-70b-hf \ ---max_new_tokens 4096 \ ---bf16 \ ---use_hpu_graphs \ ---use_kv_cache \ ---batch_size 180 \ ---attn_softmax_bf16 \ ---limit_hpu_graphs \ ---reuse_cache \ ---trim_logits \ ---sdp_on_bf16 -``` -To run Falcon-7B inference, use the following command: -```bash -python run_generation.py \ - --model_name_or_path tiiuae/falcon-7b \ - --bf16 \ - --use_hpu_graphs \ - --use_kv_cache \ - --batch_size 1 \ - --max_new_tokens 128 \ - --do_sample \ - --sdp_on_bf16 -``` - -To run Falcon-40B inference on 8 Gaudi2 cards, use the following command: -```bash -python ../gaudi_spawn.py --use_deepspeed --world_size 8 run_generation.py \ ---model_name_or_path tiiuae/falcon-40b \ ---max_new_tokens 2048 \ ---bf16 \ ---use_hpu_graphs \ ---use_kv_cache \ ---batch_size 1 \ ---do_sample \ ---use_flash_attention \ ---flash_attention_causal_mask -``` To run Llama3-405B inference on 8 Gaudi3 cards use the following command: ```bash @@ -260,29 +219,6 @@ python run_generation.py \ --sdp_on_bf16 ``` -### Using growing bucket optimization - -With `--bucket_size`, instead of padding up the kv-cache up to full size before starting, we grow the cache/input in multiples of `bucket_size`. This helps increase throughput and also reduce number of compilations if the dataset has varying prompt lengths. - -> For now, it is available only for greedy and beam search generation, and cannot be used with `--reuse_cache`. - -Here is an example: -```bash -python run_generation.py \ ---model_name_or_path path_to_model \ ---use_hpu_graphs \ ---use_kv_cache \ ---bf16 \ ---max_new_tokens 200 \ ---batch_size=2 \ ---bucket_size 50 -``` - -`--bucket_size` option is especially useful when processing an input stream with varying lengths, that is when you have something like `--dataset_name squad --column_name context --max_input_tokens -1`. `--max_input_tokens -1` specifies no truncation of input prompt in the dataset. - -Another way to simulate dynamic input is to use `--simulate_dyn_prompt`. For example `--simulate_dyn_prompt 25 35 45` will extend or crop the default prompt (or the prompt passed in using `--prompt`) to sizes 25, 35, and 45, and throughput will be measured for these 3 lengths. If `--simulate_dyn_prompt` is used, the min and max input lengths from it are computed to perform warmup as well. One final optimization that can be used in case of dynamic inputs is `--reduce_recompile`. Thus the suggested configuration to simulate dynamicity after warmup is to use all three arguments: `--simulate_dyn_prompt 25 35 45 --reduce_recompile --bucket_size 30` - -While `--bucket_size` works for any model without model file changes, an even more optimized version of bucketing is supported for certain models like Llama. This can be enabled by setting `--bucket_internal` flag (along with `--bucket_size` to specify the bucket size) ### Using Beam Search @@ -353,66 +289,11 @@ PT_ENABLE_INT64_SUPPORT=1 PT_HPU_LAZY_MODE=0 python ../gaudi_spawn.py --world_s ### Running with FP8 -Llama2-70b, Llama2-7b, Llama3-70b, Llama3-8b, Mixtral-8x7B, Falcon-7B, Falcon-40B, Falcon-180B, phi-2 and Llama3-405B in FP8 are enabled using the [Intel Neural Compressor (INC)](https://docs.habana.ai/en/latest/PyTorch/Inference_on_PyTorch/Inference_Using_FP8.html), which provides model measurement and quantization capabilities in PyTorch. From synapse 1.17 / optimum-habana 1.13 release, INC is used by default for measuring and quantization. Habana Quantization Toolkit (HQT), which was used earlier, will be removed in future releases. To use HQT, disable INC by setting the following environment variable: `USE_INC=0`. +Llama2-70b, Llama2-7b, Llama3-70b, Llama3-8b, Mixtral-8x7B, Falcon-180B and Llama3-405B in FP8 are enabled using the [Intel Neural Compressor (INC)](https://docs.habana.ai/en/latest/PyTorch/Inference_on_PyTorch/Inference_Using_FP8.html), which provides model measurement and quantization capabilities in PyTorch. From synapse 1.17 / optimum-habana 1.13 release, INC is used by default for measuring and quantization. Habana Quantization Toolkit (HQT), which was used earlier, will be removed in future releases. To use HQT, disable INC by setting the following environment variable: `USE_INC=0`. More information on enabling fp8 in SynapseAI is available here: https://docs.habana.ai/en/latest/PyTorch/Inference_on_PyTorch/Inference_Using_FP8.html -Here is an example to measure the tensor quantization statistics on LLama2-70b: -```bash -QUANT_CONFIG=./quantization_config/maxabs_measure.json python ../gaudi_spawn.py \ ---use_deepspeed --world_size 8 run_lm_eval.py \ --o acc_70b_bs1_measure.txt \ ---model_name_or_path meta-llama/Llama-2-70b-hf \ ---attn_softmax_bf16 \ ---use_hpu_graphs \ ---trim_logits \ ---use_kv_cache \ ---bucket_size=128 \ ---bucket_internal \ ---use_flash_attention \ ---flash_attention_recompute \ ---bf16 \ ---batch_size 1 -``` - -Here is an example to quantize the model based on previous measurements for LLama2-70b: -```bash -QUANT_CONFIG=./quantization_config/maxabs_quant.json python ../gaudi_spawn.py \ ---use_deepspeed --world_size 8 run_lm_eval.py \ --o acc_70b_bs1_quant.txt \ ---model_name_or_path meta-llama/Llama-2-70b-hf \ ---attn_softmax_bf16 \ ---use_hpu_graphs \ ---trim_logits \ ---use_kv_cache \ ---bucket_size=128 \ ---bucket_internal \ ---use_flash_attention \ ---flash_attention_recompute \ ---bf16 \ ---batch_size 1 -``` - -Alternatively, here is another example to quantize the model based on previous measurements for LLama2-70b: -```bash -QUANT_CONFIG=./quantization_config/maxabs_quant.json python ../gaudi_spawn.py \ ---use_deepspeed --world_size 8 run_generation.py \ ---model_name_or_path meta-llama/Llama-2-70b-hf \ ---attn_softmax_bf16 \ ---use_hpu_graphs \ ---trim_logits \ ---use_kv_cache \ ---reuse_cache \ ---use_flash_attention \ ---flash_attention_recompute \ ---bf16 \ ---batch_size 350 \ ---max_new_tokens 2048 \ ---max_input_tokens 2048 \ ---limit_hpu_graphs -``` - Here is an example to measure the tensor quantization statistics on Mixtral-8x7B with 1 card: ```bash QUANT_CONFIG=./quantization_config/maxabs_measure.json python run_generation.py \ @@ -514,12 +395,12 @@ QUANT_CONFIG=./quantization_config/maxabs_quant.json python ../gaudi_spawn.py \ --flash_attention_causal_mask ``` -Here is an example to measure the tensor quantization statistics on phi-2 with 1 card: +Here is an example to measure the tensor quantization statistics on Llama3-8b with 1 card: ```bash QUANT_CONFIG=./quantization_config/maxabs_measure.json python run_lm_eval.py \ --o acc_phi-2_bs1_measure.txt \ ---model_name_or_path microsoft/phi-2 \ +-o acc_Llama3-8b_bs1_measure.txt \ +--model_name_or_path meta-llama/Meta-Llama-3-8B \ --use_hpu_graphs \ --use_kv_cache \ --max_new_tokens 100 \ @@ -529,10 +410,10 @@ QUANT_CONFIG=./quantization_config/maxabs_measure.json python run_lm_eval.py \ --bf16 ``` -Here is an example to quantize the model based on previous measurements for phi-2 with 1 card: +Here is an example to quantize the model based on previous measurements for Llama3-8b with 1 card: ```bash -QUANT_CONFIG=./quantization_config/maxabs_quant_phi.json python run_generation.py \ ---model_name_or_path microsoft/phi-2 \ +QUANT_CONFIG=./quantization_config/maxabs_quant.json python run_generation.py \ +--model_name_or_path meta-llama/Meta-Llama-3-8B \ --use_hpu_graphs \ --use_kv_cache \ --max_new_tokens 100 \ @@ -844,17 +725,6 @@ python run_lm_eval.py \ -o eval.json ``` -Evaluate Llama 70B on 8 Gaudi2 cards on task WinoGrande, using the BF16 data type: -``` -deepspeed --num_gpus 8 run_lm_eval.py \ ---model_name_or_path meta-llama/Llama-2-70b-hf \ ---use_hpu_graphs \ ---use_kv_cache \ ---bf16 \ ---batch_size=1 \ ---tasks winogrande \ --o eval.json -``` ## Text-Generation Pipeline From 6772b4fa91188d7366f24deb4c36951e2a887cc3 Mon Sep 17 00:00:00 2001 From: Iman Gohari Date: Thu, 20 Feb 2025 09:37:55 -0700 Subject: [PATCH 50/89] Diffusers: Simplified the README files. Updated CI tests. (#1718) Signed-off-by: Daniel Socek Co-authored-by: Daniel Socek Co-authored-by: regisss <15324346+regisss@users.noreply.github.com> --- README.md | 8 +- examples/stable-diffusion/README.md | 557 ++----------------- examples/stable-diffusion/training/README.md | 377 +------------ tests/test_diffusers.py | 16 + 4 files changed, 94 insertions(+), 864 deletions(-) diff --git a/README.md b/README.md index e8fdf07116..92c002e043 100644 --- a/README.md +++ b/README.md @@ -291,11 +291,11 @@ The following model architectures, tasks and device distributions have been vali | Architecture | Training | Inference | Tasks | |:--------------------|:--------:|:---------:|:------| -| Stable Diffusion | :heavy_check_mark: | :heavy_check_mark: |
  • [text-to-image generation](https://github.com/huggingface/optimum-habana/tree/main/examples/stable-diffusion#stable-diffusion)
  • [image-to-image generation](https://github.com/huggingface/optimum-habana/tree/main/examples/stable-diffusion#stable-diffusion-based-image-to-image)
    • | +| Stable Diffusion | :heavy_check_mark: | :heavy_check_mark: |
  • [text-to-image generation](https://github.com/huggingface/optimum-habana/tree/main/examples/stable-diffusion#text-to-image-generation)
  • [image-to-image generation](https://github.com/huggingface/optimum-habana/tree/main/examples/stable-diffusion#image-to-image-generation)
    • | | Stable Diffusion XL | :heavy_check_mark: | :heavy_check_mark: |
  • [text-to-image generation](https://github.com/huggingface/optimum-habana/tree/main/examples/stable-diffusion#stable-diffusion-xl-sdxl)
  • [image-to-image generation](https://github.com/huggingface/optimum-habana/tree/main/examples/stable-diffusion#stable-diffusion-xl-refiner)
    • | -| Stable Diffusion Depth2img | |
  • Single card
    • |
  • [depth-to-image generation](https://github.com/huggingface/optimum-habana/tree/main/examples/stable-diffusion#depth-to-image-generation)
    • | -| Stable Diffusion 3 | |
  • Single card
    • |
  • [text-to-image generation](https://github.com/huggingface/optimum-habana/tree/main/examples/stable-diffusion#stable-diffusion-3-sd3)
    • | -| LDM3D | |
  • Single card
    • |
  • [text-to-image generation](https://github.com/huggingface/optimum-habana/tree/main/examples/stable-diffusion#latent-diffusion-model-for-3d-ldm3d)
    • | +| Stable Diffusion Depth2img | |
  • Single card
    • |
  • [depth-to-image generation](https://github.com/huggingface/optimum-habana/tree/main/examples/stable-diffusion)
    • | +| Stable Diffusion 3 | |
  • Single card
    • |
  • [text-to-image generation](https://github.com/huggingface/optimum-habana/tree/main/examples/stable-diffusion#stable-diffusion-3-and-35-sd3)
    • | +| LDM3D | |
  • Single card
    • |
  • [text-to-image generation](https://github.com/huggingface/optimum-habana/tree/main/examples/stable-diffusion#text-to-image-generation)
    • | | FLUX.1 |
  • LoRA
    • |
  • Single card
    • |
  • [text-to-image generation](https://github.com/huggingface/optimum-habana/tree/main/examples/stable-diffusion#flux1)
  • [image-to-image generation](https://github.com/huggingface/optimum-habana/tree/main/examples/stable-diffusion#flux1-image-to-image)
    • | | Text to Video | |
  • Single card
    • |
  • [text-to-video generation](https://github.com/huggingface/optimum-habana/tree/main/examples/stable-diffusion#text-to-video-generation)
    • | | Image to Video | |
  • Single card
    • |
  • [image-to-video generation](https://github.com/huggingface/optimum-habana/tree/main/examples/stable-diffusion#image-to-video-generation)
    • | diff --git a/examples/stable-diffusion/README.md b/examples/stable-diffusion/README.md index 9919780543..71f887ab7d 100644 --- a/examples/stable-diffusion/README.md +++ b/examples/stable-diffusion/README.md @@ -30,143 +30,8 @@ pip install -r requirements.txt ## Text-to-Image Generation -### Stable Diffusion - -Here's how to generate images using the Stable Diffusion 1.4 model with a single prompt: - -```bash -python text_to_image_generation.py \ - --model_name_or_path CompVis/stable-diffusion-v1-4 \ - --prompts "An image of a squirrel in Picasso style" \ - --num_images_per_prompt 28 \ - --batch_size 7 \ - --image_save_dir /tmp/stable_diffusion_images \ - --use_habana \ - --use_hpu_graphs \ - --gaudi_config Habana/stable-diffusion \ - --sdp_on_bf16 \ - --bf16 -``` - -> [!NOTE] -> HPU graphs are recommended when generating images by batches to get the fastest possible generations. -> The first batch of images entails a performance penalty. All subsequent batches will be generated much faster. -> You can enable this mode with `--use_hpu_graphs`. - -To generate images with multiple prompts, simply include two prompts in your input as shown below: - -```bash -python text_to_image_generation.py \ - --model_name_or_path CompVis/stable-diffusion-v1-4 \ - --prompts "An image of a squirrel in Picasso style" "A shiny flying horse taking off" \ - --num_images_per_prompt 32 \ - --batch_size 8 \ - --image_save_dir /tmp/stable_diffusion_images \ - --use_habana \ - --use_hpu_graphs \ - --gaudi_config Habana/stable-diffusion \ - --sdp_on_bf16 \ - --bf16 -``` - -Distributed inference with multiple HPUs is also supported. Below is an example demonstrating how to generate images with two prompts on two HPUs: - -```bash -python ../gaudi_spawn.py \ - --world_size 2 text_to_image_generation.py \ - --model_name_or_path CompVis/stable-diffusion-v1-4 \ - --prompts "An image of a squirrel in Picasso style" "A shiny flying horse taking off" \ - --num_images_per_prompt 20 \ - --batch_size 4 \ - --image_save_dir /tmp/stable_diffusion_images \ - --use_habana \ - --use_hpu_graphs \ - --gaudi_config Habana/stable-diffusion \ - --sdp_on_bf16 \ - --bf16 \ - --distributed -``` - -> [!NOTE] -> HPU graphs are recommended when generating images by batches to get the fastest possible generations. -> The first batch of images entails a performance penalty. All subsequent batches will be generated much faster. -> You can enable this mode with `--use_hpu_graphs`. - -You can run other older Stable Diffusion models in a similar manner. For example, to generate images with Stable Diffusion 1.5, use the option: -`--model_name_or_path stable-diffusion-v1-5/stable-diffusion-v1-5`. Examples showcasing Stable Diffusion 2 are provided next. - -### Stable Diffusion 2 - -[Stable Diffusion 2](https://huggingface.co/docs/diffusers/main/en/api/pipelines/stable_diffusion_2) can also be used -to generate images with this script. Here is an example demonstrating image generation with a single prompt: - -```bash -python text_to_image_generation.py \ - --model_name_or_path stabilityai/stable-diffusion-2-1 \ - --prompts "An image of a squirrel in Picasso style" \ - --num_images_per_prompt 28 \ - --batch_size 7 \ - --height 768 \ - --width 768 \ - --image_save_dir /tmp/stable_diffusion_images \ - --use_habana \ - --use_hpu_graphs \ - --gaudi_config Habana/stable-diffusion-2 \ - --sdp_on_bf16 \ - --bf16 -``` - -> [!NOTE] -> There are two different checkpoints for Stable Diffusion 2: -> - use [stabilityai/stable-diffusion-2-1](https://huggingface.co/stabilityai/stable-diffusion-2-1) for generating 768x768 images -> - use [stabilityai/stable-diffusion-2-1-base](https://huggingface.co/stabilityai/stable-diffusion-2-1-base) for generating 512x512 images - -### Latent Diffusion Model for 3D (LDM3D) - -[LDM3D](https://arxiv.org/abs/2305.10853) generates both image and depth map data from a given text prompt, allowing users -to generate RGBD images from text prompts. - -[Original checkpoint](https://huggingface.co/Intel/ldm3d) and [latest checkpoint](https://huggingface.co/Intel/ldm3d-4c) -are open source. A [demo](https://huggingface.co/spaces/Intel/ldm3d) is also available. Here is how to run this model: - -```bash -python text_to_image_generation.py \ - --model_name_or_path "Intel/ldm3d-4c" \ - --prompts "An image of a squirrel in Picasso style" \ - --num_images_per_prompt 28 \ - --batch_size 7 \ - --height 768 \ - --width 768 \ - --image_save_dir /tmp/stable_diffusion_images \ - --use_habana \ - --use_hpu_graphs \ - --gaudi_config Habana/stable-diffusion-2 \ - --ldm3d -``` - -Here is how to generate images and depth maps with two prompts on two HPUs: - -```bash -python ../gaudi_spawn.py --world_size 2 text_to_image_generation.py \ - --model_name_or_path "Intel/ldm3d-4c" \ - --prompts "An image of a squirrel in Picasso style" "A shiny flying horse taking off" \ - --num_images_per_prompt 10 \ - --batch_size 2 \ - --height 768 \ - --width 768 \ - --image_save_dir /tmp/stable_diffusion_images \ - --use_habana \ - --use_hpu_graphs \ - --gaudi_config Habana/stable-diffusion-2 \ - --ldm3d \ - --distributed -``` - -> [!NOTE] -> There are three different checkpoints for LDM3D: -> - use [original checkpoint](https://huggingface.co/Intel/ldm3d) to generate outputs from the paper -> - use [the latest checkpoint](https://huggingface.co/Intel/ldm3d-4c) for generating improved results -> - use [the pano checkpoint](https://huggingface.co/Intel/ldm3d-pano) to generate panoramic view +Optimum for Intel Gaudi supports state-of-the-art diffusion-based text-to-image generation models, including SDXL, SD3/3.5, and FLUX. We provide +brief inference examples for these models. For running legacy Stable Diffusion (SD) models, please refer to [this](README_legacy.md) document. ### Stable Diffusion XL (SDXL) @@ -196,113 +61,27 @@ python text_to_image_generation.py \ > The first batch of images entails a performance penalty. All subsequent batches will be generated much faster. > You can enable this mode with `--use_hpu_graphs`. -SDXL integrates a second text encoder (OpenCLIP ViT-bigG/14), alongside the original Stable Diffusion text encoder. This addition significantly increases the number of parameters, enabling more detailed and descriptive prompts. Below is an example of how to generate images using multiple prompts for both `prompt` (primary text encoder) and `prompt_2` (secondary text encoder), along with their respective negative prompts: - -```bash -python text_to_image_generation.py \ - --model_name_or_path stabilityai/stable-diffusion-xl-base-1.0 \ - --prompts "Sailing ship painting by Van Gogh" "A shiny flying horse taking off" \ - --prompts_2 "Red tone" "Blue tone" \ - --negative_prompts "Low quality" "Sketch" \ - --negative_prompts_2 "Clouds" "Clouds" \ - --num_images_per_prompt 32 \ - --batch_size 8 \ - --image_save_dir /tmp/stable_diffusion_xl_images \ - --scheduler euler_discrete \ - --use_habana \ - --use_hpu_graphs \ - --gaudi_config Habana/stable-diffusion \ - --sdp_on_bf16 \ - --bf16 -``` - -SDXL also supports distributed inferencing with Intel Gaudi accelerators. Below is an example of generating SDXL images in a distributed manner using two prompts on two HPUs: - -```bash -python ../gaudi_spawn.py --world_size 2 text_to_image_generation.py \ - --model_name_or_path stabilityai/stable-diffusion-xl-base-1.0 \ - --prompts "Sailing ship painting by Van Gogh" "A shiny flying horse taking off" \ - --prompts_2 "Red tone" "Blue tone" \ - --negative_prompts "Low quality" "Sketch" \ - --negative_prompts_2 "Clouds" "Clouds" \ - --num_images_per_prompt 32 \ - --batch_size 8 \ - --image_save_dir /tmp/stable_diffusion_xl_images \ - --scheduler euler_discrete \ - --use_habana \ - --use_hpu_graphs \ - --gaudi_config Habana/stable-diffusion \ - --sdp_on_bf16 \ - --bf16 \ - --distributed -``` - -The performance-optimized SDXL pipeline can be enabled using the `--optimize` option. This option utilizes a more aggressively optimized attention mechanism for enhanced performance. Additionally, it supports running -inference in mixed FP8 precision. - -Here is how to generate SDXL images with optimized pipeline in FP8 precision: -```bash -QUANT_CONFIG=quantization/stable-diffusion-xl/quantize_config.json \ -python text_to_image_generation.py \ - --model_name_or_path stabilityai/stable-diffusion-xl-base-1.0 \ - --prompts "Sailing ship painting by Van Gogh" \ - --num_images_per_prompt 28 \ - --batch_size 7 \ - --image_save_dir /tmp/stable_diffusion_xl_images \ - --scheduler euler_discrete \ - --use_habana \ - --use_hpu_graphs \ - --gaudi_config Habana/stable-diffusion \ - --sdp_on_bf16 \ - --bf16 \ - --optimize -``` - -### SDXL-Turbo - -The knowledge distillation technique can be used to train a distilled version of SDXL, allowing for high-quality -image generation with fewer inference steps. SDXL-Turbo is a distilled version of Stable Diffusion XL 1.0, -optimized for real-time synthesis. - -Here is how to generate images with multiple prompts: - -```bash -python text_to_image_generation.py \ - --model_name_or_path stabilityai/sdxl-turbo \ - --prompts "Sailing ship painting by Van Gogh" "A shiny flying horse taking off" \ - --num_images_per_prompt 32 \ - --batch_size 8 \ - --image_save_dir /tmp/stable_diffusion_xl_turbo_images \ - --scheduler euler_ancestral_discrete \ - --use_habana \ - --use_hpu_graphs \ - --gaudi_config Habana/stable-diffusion \ - --sdp_on_bf16 \ - --bf16 \ - --num_inference_steps 1 \ - --guidance_scale 1.000001 \ - --timestep_spacing trailing -``` - > [!WARNING] > There is a regression with `--guidance_scale 0.0` in current release which will be addressed in later releases. > Setting `--guidance_scale` to a value larger than 1 resolves the regression. -### Stable Diffusion 3 (SD3) +To input multiple prompts, pass prompt strings separated by spaces. SDXL improves text-to-image generation by combining +OpenCLIP ViT-bigG/14 with the original Stable Diffusion text encoder, thus allowing for more descriptive prompts. +You can pass single or multiple prompts for both `prompt` and `prompt_2` (2nd text encoder), as well as their negative prompts. -Stable Diffusion 3 was introduced by Stability AI [here](https://stability.ai/news/stable-diffusion-3). -It uses Diffusion Transformer instead of UNet for denoising, which yields improved image quality. +Additionally, you can run inference on multiple HPUs by replacing `python text_to_image_generation.py` +with `python ../gaudi_spawn.py --world_size text_to_image_generation.py` and adding option `--distributed`. -Before running SD3 pipeline, you need to: +A version of the SDXL pipeline optimized for FP8 on Intel Gaudi is also available. Set +`QUANT_CONFIG=quantization/stable-diffusion-xl/quantize_config.json` enviromement variable and use option `--optimize` +to run FP8-optimized SDXL pipeline. -1. Agree to the Terms and Conditions for using SD3 model at [HuggingFace model page](https://huggingface.co/stabilityai/stable-diffusion-3-medium) -2. Authenticate with HuggingFace using your HF Token. For authentication, run: +To run SDXL-Turbo, the distilled version of SDXL, use `--model_name_or_path stabilityai/sdxl-turbo` in the input. -```bash -huggingface-cli login -``` +### Stable Diffusion 3 and 3.5 (SD3) -Here is how to generate SD3 images with a single prompt: +Stable Diffusion 3 was introduced by Stability AI [here](https://stability.ai/news/stable-diffusion-3). +It uses Diffusion Transformer instead of UNet for denoising, which yields improved image quality. ```bash python text_to_image_generation.py \ @@ -320,79 +99,20 @@ python text_to_image_generation.py \ --bf16 ``` -This model can also be quantized with some ops running in FP8 precision. +> [!NOTE] +> The access to SD3 requires agreeing to its terms and conditions at [HuggingFace model page](https://huggingface.co/stabilityai/stable-diffusion-3-medium), +> and then authenticating using your HF token via `huggingface-cli login`. -Before quantization, run stats collection using measure mode: +This model can also be quantized with some ops running in FP8 precision. Before quantization, run stats collection using measure mode by setting +runtime variable `QUANT_CONFIG=quantization/stable-diffusion-3/measure_config.json` and `--quant_mode measure`. After stats collection, you can run +SD3 in quantization mode by setting runtime variable `QUANT_CONFIG=quantization/stable-diffusion-3/quantize_config.json` and `--quant_mode quantize`. -```bash -QUANT_CONFIG=quantization/stable-diffusion-3/measure_config.json \ -python text_to_image_generation.py \ - --model_name_or_path stabilityai/stable-diffusion-3-medium-diffusers \ - --prompts "Sailing ship painting by Van Gogh" \ - --num_images_per_prompt 10 \ - --batch_size 1 \ - --num_inference_steps 28 \ - --image_save_dir /tmp/stable_diffusion_3_images \ - --scheduler default \ - --use_habana \ - --use_hpu_graphs \ - --gaudi_config Habana/stable-diffusion \ - --sdp_on_bf16 \ - --bf16 \ - --quant_mode measure -``` +To run Stable Diffusion 3.5 Large, use `--model_name_or_path stabilityai/stable-diffusion-3.5-large` in the input. -After stats collection, here is how to run SD3 in quantization mode: - -```bash -QUANT_CONFIG=quantization/stable-diffusion-3/quantize_config.json \ -python text_to_image_generation.py \ - --model_name_or_path stabilityai/stable-diffusion-3-medium-diffusers \ - --prompts "Sailing ship painting by Van Gogh" \ - --num_images_per_prompt 10 \ - --batch_size 1 \ - --num_inference_steps 28 \ - --image_save_dir /tmp/stable_diffusion_3_images \ - --scheduler default \ - --use_habana \ - --use_hpu_graphs \ - --gaudi_config Habana/stable-diffusion \ - --sdp_on_bf16 \ - --bf16 \ - --quant_mode quantize -``` - -### FLUX.1 +### FLUX FLUX.1 was introduced by Black Forest Labs [here](https://blackforestlabs.ai/announcing-black-forest-labs/). -Here is how to run FLUX.1-schnell model (distilled fast version of FLUX.1): - -```bash -python text_to_image_generation.py \ - --model_name_or_path black-forest-labs/FLUX.1-schnell \ - --prompts "A cat holding a sign that says hello world" \ - --num_images_per_prompt 10 \ - --batch_size 1 \ - --num_inference_steps 4 \ - --image_save_dir /tmp/flux_1_images \ - --scheduler flow_match_euler_discrete \ - --use_habana \ - --use_hpu_graphs \ - --gaudi_config Habana/stable-diffusion \ - --sdp_on_bf16 \ - --bf16 -``` - -Before running FLUX.1-dev model, you need to: - -1. Agree to the Terms and Conditions for using FLUX.1-dev model at [HuggingFace model page](https://huggingface.co/black-forest-labs/FLUX.1-dev) -2. Authenticate with HuggingFace using your HF Token. For authentication, run: - -```bash -huggingface-cli login -``` - Here is how to run FLUX.1-dev model: ```bash @@ -411,59 +131,28 @@ python text_to_image_generation.py \ --bf16 ``` -This model can also be quantized with some ops running in FP8 precision. - -Before quantization, run stats collection using measure mode: - -```bash -QUANT_CONFIG=quantization/flux/measure_config.json \ -python text_to_image_generation.py \ - --model_name_or_path black-forest-labs/FLUX.1-dev \ - --prompts "A cat holding a sign that says hello world" \ - --num_images_per_prompt 10 \ - --batch_size 1 \ - --num_inference_steps 30 \ - --image_save_dir /tmp/flux_1_images \ - --scheduler flow_match_euler_discrete \ - --use_habana \ - --use_hpu_graphs \ - --gaudi_config Habana/stable-diffusion \ - --sdp_on_bf16 \ - --bf16 \ - --quant_mode measure -``` +> [!NOTE] +> The access to FLUX.1-dev model requires agreeing to its terms and conditions at [HuggingFace model page](https://huggingface.co/black-forest-labs/FLUX.1-dev), +> and then authenticating using your HF token via `huggingface-cli login`. -After stats collection, here is how to run FLUX.1-dev in quantization mode: +This model can also be quantized with some ops running in FP8 precision. Before quantization, run stats collection using measure mode by setting +runtime variable `QUANT_CONFIG=quantization/flux/measure_config.json` and `--quant_mode measure`. After stats collection, you can run +FLUX in quantization mode by setting runtime variable `QUANT_CONFIG=quantization/flux/quantize_config.json` and `--quant_mode quantize`. -```bash -QUANT_CONFIG=quantization/flux/quantize_config.json \ -python text_to_image_generation.py \ - --model_name_or_path black-forest-labs/FLUX.1-dev \ - --prompts "A cat holding a sign that says hello world" \ - --num_images_per_prompt 10 \ - --batch_size 1 \ - --num_inference_steps 30 \ - --image_save_dir /tmp/flux_1_images \ - --scheduler flow_match_euler_discrete \ - --use_habana \ - --use_hpu_graphs \ - --gaudi_config Habana/stable-diffusion \ - --sdp_on_bf16 \ - --bf16 \ - --quant_mode quantize -``` +To run with FLUX.1-schnell model, a distilled version of FLUX.1 (which is not gated), use `--model_name_or_path black-forest-labs/FLUX.1-schnell`. ## ControlNet - ControlNet was introduced in [Adding Conditional Control to Text-to-Image Diffusion Models](https://huggingface.co/papers/2302.05543) -by Lvmin Zhang and Maneesh Agrawala, enables conditioning the Stable Diffusion model with an additional input image. This allows for precise control over the composition of generated images using various features such as edges, pose, depth, and more. +by Lvmin Zhang and Maneesh Agrawala, enables conditioning the Stable Diffusion model with an additional input image. +This allows for precise control over the composition of generated images using various features such as edges, +pose, depth, and more. Here is how to generate images conditioned by Canny edge model: ```bash python text_to_image_generation.py \ - --model_name_or_path CompVis/stable-diffusion-v1-4 \ + --model_name_or_path stable-diffusion-v1-5/stable-diffusion-v1-5 \ --controlnet_model_name_or_path lllyasviel/sd-controlnet-canny \ --prompts "futuristic-looking woman" \ --control_image https://hf.co/datasets/huggingface/documentation-images/resolve/main/diffusers/input_image_vermeer.png \ @@ -477,65 +166,29 @@ python text_to_image_generation.py \ --bf16 ``` -The ControlNet example can be run with multiple prompts by supplying more than one prompt in the input. -Additionally, it supports distributed execution. Below is an example of generating images conditioned by the Canny edge model using two prompts on two HPUs: - -```bash -python ../gaudi_spawn.py --world_size 2 text_to_image_generation.py \ - --model_name_or_path CompVis/stable-diffusion-v1-4 \ - --controlnet_model_name_or_path lllyasviel/sd-controlnet-canny \ - --prompts "futuristic-looking woman" "a rusty robot" \ - --control_image https://hf.co/datasets/huggingface/documentation-images/resolve/main/diffusers/input_image_vermeer.png \ - --num_images_per_prompt 16 \ - --batch_size 4 \ - --image_save_dir /tmp/controlnet_images \ - --use_habana \ - --use_hpu_graphs \ - --gaudi_config Habana/stable-diffusion \ - --sdp_on_bf16 \ - --bf16 \ - --distributed -``` +You can run inference on multiple HPUs by replacing `python text_to_image_generation.py` with +`python ../gaudi_spawn.py --world_size text_to_image_generation.py` and adding option `--distributed`. -These ControlNet examples will preprocess the input image to derive Canny edges. Alternatively, you can use `--control_preprocessing_type none` to supply a preprocessed control image directly, enabling many additional use cases. +This ControlNet example will preprocess the input image to derive Canny edges. Alternatively, you can use `--control_preprocessing_type none` +to supply a preprocessed control image directly, enabling many additional use cases. ## Inpainting Inpainting replaces or edits specific areas of an image. For more details, please refer to [Hugging Face Diffusers doc](https://huggingface.co/docs/diffusers/en/using-diffusers/inpaint). -### Stable Diffusion Inpainting - -```bash -python text_to_image_generation.py \ - --model_name_or_path stabilityai/stable-diffusion-2-inpainting \ - --base_image https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint.png \ - --mask_image https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint_mask.png \ - --prompts "concept art digital painting of an elven castle, inspired by lord of the rings, highly detailed, 8k" \ - --seed 0 \ - --num_images_per_prompt 12 \ - --batch_size 4 \ - --image_save_dir /tmp/inpaiting_images \ - --use_habana \ - --use_hpu_graphs \ - --gaudi_config Habana/stable-diffusion \ - --sdp_on_bf16 \ - --bf16 -``` - ### Stable Diffusion XL Inpainting ```bash python text_to_image_generation.py \ - --model_name_or_path diffusers/stable-diffusion-xl-1.0-inpainting-0.1 \ + --model_name_or_path diffusers/stable-diffusion-xl-1.0-inpainting-0.1 \ --base_image https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint.png \ --mask_image https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint_mask.png \ --prompts "concept art digital painting of an elven castle, inspired by lord of the rings, highly detailed, 8k" \ --seed 0 \ - --scheduler euler_discrete \ --num_images_per_prompt 12 \ --batch_size 4 \ - --image_save_dir /tmp/xl_inpaiting_images \ + --image_save_dir /tmp/inpaiting_images \ --use_habana \ --use_hpu_graphs \ --gaudi_config Habana/stable-diffusion \ @@ -548,25 +201,6 @@ python text_to_image_generation.py \ This section provides examples of additional inference techniques based on Stable Diffusion. For more details, please refer to [Hugging Face Diffusers documentation](https://huggingface.co/docs/diffusers/main/en/using-diffusers/overview_techniques). -### Unconditional Image Generation - -Here is how to perform unconditional image generation on Intel Gaudi. For more details, please refer to the -[Unconditional Image Generation](https://huggingface.co/docs/diffusers/using-diffusers/unconditional_image_generation) -section in the Hugging Face documentation. - -```bash -python unconditional_image_generation.py \ - --model_name_or_path "google/ddpm-ema-celebahq-256" \ - --batch_size 16 \ - --use_habana \ - --use_gaudi_ddim_scheduler \ - --use_hpu_graphs \ - --sdp_on_bf16 \ - --bf16 \ - --save_outputs \ - --output_dir "/tmp/" -``` - ### Controlling Brightness Here is an example of how to control brightness. For more information, please refer to the @@ -597,12 +231,12 @@ section in the Hugging Face documentation. ```bash python text_to_image_generation.py \ - --model_name_or_path CompVis/stable-diffusion-v1-4 \ - --prompts "a red cat playing with a ball+++" "a red cat playing with a ball---" \ + --model_name_or_path stabilityai/stable-diffusion-xl-base-1.0 \ + --prompts "a red cat--- playing with a ball+++" "a red cat+++ playing with a ball---" \ --num_images_per_prompt 4 \ --batch_size 4 \ --use_habana --use_hpu_graphs \ - --image_save_dir /tmp/stable_diffusion_images_compel \ + --image_save_dir /tmp/stable_diffusion_xl_images_compel \ --seed 33 \ --sdp_on_bf16 \ --bf16 \ @@ -618,12 +252,12 @@ section in the Hugging Face documentation. ```bash python text_to_image_generation.py \ - --model_name_or_path CompVis/stable-diffusion-v1-4 \ + --model_name_or_path stabilityai/stable-diffusion-xl-base-1.0 \ --prompts "A squirrel eating a burger" \ --num_images_per_prompt 4 \ --batch_size 4 \ --use_habana \ - --image_save_dir /tmp/stable_diffusion_images_freeu \ + --image_save_dir /tmp/stable_diffusion_xl_images_freeu \ --seed 33 \ --use_freeu \ --sdp_on_bf16 \ @@ -634,34 +268,7 @@ python text_to_image_generation.py \ Images can also be generated using initial input images to guide the diffusion-based image generation process. -### Stable Diffusion-based Image-to-Image - -Here is how to generate images using a single prompt and an input image with the `timbrooks/instruct-pix2pix` model, which is based on Stable Diffusion: - -```bash -python image_to_image_generation.py \ - --model_name_or_path "timbrooks/instruct-pix2pix" \ - --src_image_path "https://raw.githubusercontent.com/timothybrooks/instruct-pix2pix/main/imgs/example.jpg" \ - --prompts "turn him into cyborg" \ - --num_images_per_prompt 20 \ - --batch_size 4 \ - --guidance_scale 7.5 \ - --image_guidance_scale 1 \ - --num_inference_steps 10 \ - --image_save_dir /tmp/stable_diffusion_images \ - --use_habana \ - --use_hpu_graphs \ - --gaudi_config Habana/stable-diffusion \ - --sdp_on_bf16 \ - --bf16 -``` - -> [!NOTE] -> HPU graphs are recommended when generating images by batches to get the fastest possible generations. -> The first batch of images entails a performance penalty. All subsequent batches will be generated much faster. -> You can enable this mode with `--use_hpu_graphs`. - -### Stable Diffusion XL Refiner +### Stable Diffusion XL Image-to-Image Here is how to refine SDXL images using a single image and prompt: @@ -682,7 +289,7 @@ python image_to_image_generation.py \ --bf16 ``` -### FLUX.1 Image-to-Image +### FLUX Image-to-Image Here is how to generate a FLUX.1 image using a single input image and prompt: @@ -704,41 +311,6 @@ python image_to_image_generation.py \ --bf16 ``` -### Stable Diffusion Image Variations - -Here is how to generate image variations of a single image (without any input prompts): - -```bash -python image_to_image_generation.py \ - --model_name_or_path "lambdalabs/sd-image-variations-diffusers" \ - --src_image_path "https://github.com/SHI-Labs/Versatile-Diffusion/blob/master/assets/demo/reg_example/ghibli.jpg?raw=true" \ - --num_images_per_prompt 20 \ - --batch_size 4 \ - --image_save_dir /tmp/stable_diffusion_images \ - --guidance_scale 3 \ - --use_habana \ - --use_hpu_graphs \ - --gaudi_config Habana/stable-diffusion \ - --sdp_on_bf16 \ - --bf16 -``` - -### Depth to Image Generation - -Here is an example of performing depth-guided image generation: - -```bash -python depth_to_image_generation.py \ - --model_name_or_path "stabilityai/stable-diffusion-2-depth" \ - --prompts "two tigers" \ - --base_image "http://images.cocodataset.org/val2017/000000039769.jpg" \ - --image_save_dir /tmp/stable_diffusion_images \ - --use_habana \ - --use_hpu_graphs \ - --sdp_on_bf16 \ - --bf16 -``` - ## Text-to-Video Generation This section demonstrates how to use the `GaudiTextToVideoSDPipeline` for text-to-video generation tasks on HPUs. @@ -758,12 +330,10 @@ python text_to_video_generation.py \ Stable Video Diffusion (SVD) was unveiled in [Stable Video Diffusion Announcement](https://stability.ai/news/stable-video-diffusion-open-ai-video-model) by the Stability AI team. Stable Video Diffusion XT version (SVD-XT) is tuned to generate 25 frames of video from a single image. -## Image-to-video Generation +## Image-to-Video Generation Script `image_to_video_generation.py` showcases how to perform image-to-video generation using Stable Video Diffusion on Intel Gaudi. -### Single Image Prompt - Here is how to generate video with one image prompt: ```bash @@ -782,35 +352,10 @@ python image_to_video_generation.py \ ``` > [!NOTE] -> For improved performance of the image-to-video pipeline on Gaudi, it is recommended to configure the environment -> by setting PT_HPU_MAX_COMPOUND_OP_SIZE to 1. - -### Multiple Image Prompts - -Here is how to generate videos with several image prompts: +> For improved performance of the image-to-video pipeline on Gaudi, it is recommended to set the following env variable: `PT_HPU_MAX_COMPOUND_OP_SIZE=1`. -```bash -PT_HPU_MAX_COMPOUND_OP_SIZE=1 \ -python image_to_video_generation.py \ - --model_name_or_path "stabilityai/stable-video-diffusion-img2vid-xt" \ - --image_path \ - "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/svd/rocket.png" \ - "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png" \ - "https://hf.co/datasets/huggingface/documentation-images/resolve/main/diffusers/input_image_vermeer.png" \ - "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/hf-logo.png" \ - --num_videos_per_prompt 1 \ - --video_save_dir /tmp/stable_video_diffusion_xt \ - --save_frames_as_images \ - --use_habana \ - --use_hpu_graphs \ - --gaudi_config Habana/stable-diffusion \ - --sdp_on_bf16 \ - --bf16 -``` - -> [!NOTE] -> For improved performance of the image-to-video pipeline on Gaudi, it is recommended to configure the environment -> by setting PT_HPU_MAX_COMPOUND_OP_SIZE to 1. +You can pass multiple image prompts strings separated via space, i.e. +`--image_path "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/svd/rocket.png" "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png"`. ### Image-to-Video ControlNet @@ -876,7 +421,7 @@ python image_to_video_generation.py \ # Important Notes for Gaudi3 Users -- **Batch Size Limitation**: Due to a known issue, batch sizes for some Stable Diffusion models need to be reduced. + - **Batch Size Limitation**: Due to a known issue, batch sizes for some Stable Diffusion models need to be reduced. This issue is expected to be resolved in a future release. - **Image-to-Video ControlNet**: The Image-to-Video ControlNet command is currently not supported on Gaudi3. diff --git a/examples/stable-diffusion/training/README.md b/examples/stable-diffusion/training/README.md index 4ea85c9e36..4c1add8b76 100644 --- a/examples/stable-diffusion/training/README.md +++ b/examples/stable-diffusion/training/README.md @@ -18,91 +18,6 @@ limitations under the License. This directory contains scripts that showcase how to perform training/fine-tuning of Stable Diffusion models on Habana Gaudi. -## Textual Inversion - -[Textual Inversion](https://arxiv.org/abs/2208.01618) is a method to personalize text2image models like Stable Diffusion on your own images using just 3-5 examples. - -The `textual_inversion.py` script shows how to implement the training procedure on Habana Gaudi. - -In the examples below, we will use a set of cat images from the following dataset: -[https://huggingface.co/datasets/diffusers/cat_toy_example](https://huggingface.co/datasets/diffusers/cat_toy_example) - -To download this and other example training datasets locally, run: -```bash -python download_train_datasets.py -``` - -Now we can launch the training using: - -```bash -python textual_inversion.py \ - --pretrained_model_name_or_path CompVis/stable-diffusion-v1-4 \ - --train_data_dir ./cat \ - --learnable_property object \ - --placeholder_token "" \ - --initializer_token toy \ - --resolution 512 \ - --train_batch_size 4 \ - --max_train_steps 3000 \ - --learning_rate 5.0e-04 \ - --scale_lr \ - --lr_scheduler constant \ - --lr_warmup_steps 0 \ - --output_dir /tmp/textual_inversion_cat \ - --save_as_full_pipeline \ - --gaudi_config_name Habana/stable-diffusion \ - --throughput_warmup_steps 3 -``` - -> [!NOTE] -> Change `--resolution` to 768 if you are using the [stable-diffusion-2](https://huggingface.co/stabilityai/stable-diffusion-2) 768x768 model. - -> [!NOTE] -> As described in [the official paper](https://arxiv.org/abs/2208.01618), only one embedding vector is used for the placeholder token, *e.g.* `""`. -> However, one can also add multiple embedding vectors for the placeholder token to increase the number of fine-tuneable parameters. -> This can help the model to learn more complex details. To use multiple embedding vectors, you can define `--num_vectors` to a number larger than one, -> *e.g.*: `--num_vectors 5`. The saved textual inversion vectors will then be larger in size compared to the default case. - -Once you have trained a model as described above, inference can be done using `GaudiStableDiffusionPipeline`. -Please make sure to include the `placeholder_token` in your prompt so that textual inversion guided inference can take effect. - -You can use `text_to_image_generation.py` sample to run inference with the fine-tuned model: - -```bash -python ../text_to_image_generation.py \ - --model_name_or_path /tmp/textual_inversion_cat \ - --prompts "A backpack" \ - --num_images_per_prompt 5 \ - --batch_size 1 \ - --image_save_dir /tmp/textual_inversion_cat_images \ - --use_habana \ - --use_hpu_graphs \ - --gaudi_config Habana/stable-diffusion \ - --sdp_on_bf16 \ - --bf16 -``` - -Alternatively, you can run inference with the fine-tuned model using a simple Python script like this: - -```python -from optimum.habana.diffusers import GaudiStableDiffusionPipeline -import torch - -model_id = "/tmp/textual_inversion_cat" -pipe = GaudiStableDiffusionPipeline.from_pretrained( - model_id, - torch_dtype=torch.bfloat16, - use_habana=True, - use_hpu_graphs=True, - gaudi_config="Habana/stable-diffusion", - sdp_on_bf16=True, -) - -prompt = "A backpack" -image = pipe(prompt, num_inference_steps=50, guidance_scale=7.5).images[0] -image.save(f"cat-backpack.png") -``` - ## Textual Inversion XL The `textual_inversion_sdxl.py` script shows how to implement textual inversion fine-tuning on Gaudi for XL diffusion models @@ -155,36 +70,6 @@ python ../text_to_image_generation.py \ --bf16 ``` -Alternatively, you can run inference with the fine-tuned model using a simple standalone Python script. -The following script can be used to run inference using the fine-tuned model with both text encoders, -separately and in combination: - -```python -from optimum.habana.diffusers import GaudiStableDiffusionXLPipeline -import torch - -model_id = "/tmp/textual_inversion_cat_sdxl" -pipe = GaudiStableDiffusionXLPipeline.from_pretrained( - model_id, - torch_dtype=torch.bfloat16, - use_habana=True, - use_hpu_graphs=True, - gaudi_config="Habana/stable-diffusion", - sdp_on_bf16=True, -) - -prompt = "A backpack" -image = pipe(prompt, num_inference_steps=50, guidance_scale=7.5).images[0] -image.save(f"cat-backpack.png") - -image = pipe(prompt="", prompt_2=prompt, num_inference_steps=50, guidance_scale=7.5).images[0] -image.save(f"cat-backpack_p2.png") - -prompt_2 = "A colored backpack" -image = pipe(prompt=prompt, prompt_2=prompt_2, num_inference_steps=50, guidance_scale=7.5).images[0] -image.save(f"cat-backpack_p1and2.png") -``` - ## ControlNet Training ControlNet was introduced in [Adding Conditional Control to Text-to-Image Diffusion Models ](https://huggingface.co/papers/2302.05543) @@ -200,8 +85,8 @@ Then proceed to training with command: ```bash python train_controlnet.py \ - --pretrained_model_name_or_path=CompVis/stable-diffusion-v1-4\ - --output_dir=/tmp/stable_diffusion1_4 \ + --pretrained_model_name_or_path=stabilityai/stable-diffusion-2-1 \ + --output_dir=/tmp/stable_diffusion2_1 \ --dataset_name=fusing/fill50k \ --resolution=512 \ --learning_rate=1e-5 \ @@ -212,28 +97,12 @@ python train_controlnet.py \ --use_hpu_graphs \ --sdp_on_bf16 \ --bf16 \ + --max_train_steps 2500 \ --trust_remote_code ``` -### Multi-Card Training - -You can run these fine-tuning scripts in a distributed fashion as follows: -```bash -python ../../gaudi_spawn.py --use_mpi --world_size 8 train_controlnet.py \ - --pretrained_model_name_or_path CompVis/stable-diffusion-v1-4 \ - --output_dir=/tmp/stable_diffusion1_4 \ - --dataset_name=fusing/fill50k \ - --resolution=512 \ - --learning_rate=1e-5 \ - --validation_image "./cnet/conditioning_image_1.png" "./cnet/conditioning_image_2.png" \ - --validation_prompt "red circle with blue background" "cyan circle with brown floral background" \ - --train_batch_size=4 \ - --throughput_warmup_steps 3 \ - --use_hpu_graphs \ - --sdp_on_bf16 \ - --bf16 \ - --trust_remote_code -``` +You can run inference on multiple HPUs by replacing `python train_controlnet.py` +with `python ../gaudi_spawn.py --world_size train_controlnet.py`. ### Inference @@ -241,8 +110,8 @@ After training completes, you can use `text_to_image_generation.py` sample to ru ```bash python ../text_to_image_generation.py \ - --model_name_or_path CompVis/stable-diffusion-v1-4 \ - --controlnet_model_name_or_path /tmp/stable_diffusion1_4 \ + --model_name_or_path stabilityai/stable-diffusion-2-1 \ + --controlnet_model_name_or_path /tmp/stable_diffusion2_1 \ --prompts "pale golden rod circle with old lace background" \ --control_image "./cnet/conditioning_image_1.png" \ --num_images_per_prompt 5 \ @@ -254,43 +123,6 @@ python ../text_to_image_generation.py \ --sdp_on_bf16 \ --bf16 ``` - -Alternatively, you can run inference using a simple standalone Python script, as shown below: - -```python -from diffusers import ControlNetModel, UniPCMultistepScheduler -from diffusers.utils import load_image -import torch -from optimum.habana.diffusers import GaudiStableDiffusionControlNetPipeline - -base_model_path = "CompVis/stable-diffusion-v1-4" -controlnet_path = "/tmp/stable_diffusion1_4" - -controlnet = ControlNetModel.from_pretrained(controlnet_path, torch_dtype=torch.bfloat16) -pipe = GaudiStableDiffusionControlNetPipeline.from_pretrained( - base_model_path, - controlnet=controlnet, - torch_dtype=torch.bfloat16, - use_habana=True, - use_hpu_graphs=True, - gaudi_config="Habana/stable-diffusion", - sdp_on_bf16=True, -) - -# speed up diffusion process with faster scheduler and memory optimization -pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config) - -control_image = load_image("./cnet/conditioning_image_1.png") -prompt = "pale golden rod circle with old lace background" - -# generate image -generator = torch.manual_seed(0) -image = pipe( - prompt, num_inference_steps=20, generator=generator, image=control_image -).images[0] -image.save("./output.png") -``` - ## Fine-Tuning for Stable Diffusion XL The `train_text_to_image_sdxl.py` script shows how to implement the fine-tuning of Stable Diffusion XL models on Gaudi. @@ -336,76 +168,11 @@ python train_text_to_image_sdxl.py \ --adjust_throughput ``` -### Multi-Card Training +> [!WARNING] +> There is a known issue that in the first 2 steps, graph compilation takes longer than 10 seconds. This will be fixed in a future release. -To train Stable Diffusion XL on a multi-card Gaudi system, use: -```bash -PT_HPU_RECIPE_CACHE_CONFIG=/tmp/stdxl_recipe_cache,True,1024 \ -python ../../gaudi_spawn.py --world_size 8 --use_mpi train_text_to_image_sdxl.py \ - --pretrained_model_name_or_path stabilityai/stable-diffusion-xl-base-1.0 \ - --pretrained_vae_model_name_or_path madebyollin/sdxl-vae-fp16-fix \ - --dataset_name lambdalabs/naruto-blip-captions \ - --resolution 512 \ - --crop_resolution 512 \ - --center_crop \ - --random_flip \ - --proportion_empty_prompts=0.2 \ - --train_batch_size 16 \ - --max_train_steps 336 \ - --learning_rate 1e-05 \ - --max_grad_norm 1 \ - --lr_scheduler constant \ - --lr_warmup_steps 0 \ - --output_dir sdxl_model_output \ - --gaudi_config_name Habana/stable-diffusion \ - --throughput_warmup_steps 3 \ - --dataloader_num_workers 8 \ - --sdp_on_bf16 \ - --bf16 \ - --use_hpu_graphs_for_training \ - --use_hpu_graphs_for_inference \ - --validation_prompt="a cute naruto creature" \ - --validation_epochs 48 \ - --checkpointing_steps 336 \ - --mediapipe dataset_sdxl_mediapipe \ - --adjust_throughput -``` - -### Single Card Training on Gaudi1 - -To train Stable Diffusion XL on a single Gaudi1 card, use: -```bash -python train_text_to_image_sdxl.py \ - --pretrained_model_name_or_path stabilityai/stable-diffusion-xl-base-1.0 \ - --pretrained_vae_model_name_or_path madebyollin/sdxl-vae-fp16-fix \ - --dataset_name lambdalabs/naruto-blip-captions \ - --resolution 256 \ - --center_crop \ - --random_flip \ - --proportion_empty_prompts=0.2 \ - --train_batch_size 1 \ - --gradient_accumulation_steps 4 \ - --max_train_steps 3000 \ - --learning_rate 1e-05 \ - --max_grad_norm 1 \ - --lr_scheduler constant \ - --lr_warmup_steps 0 \ - --output_dir sdxl_model_output \ - --gaudi_config_name Habana/stable-diffusion \ - --throughput_warmup_steps 3 \ - --use_hpu_graphs_for_training \ - --use_hpu_graphs_for_inference \ - --checkpointing_steps 3000 \ - --sdp_on_bf16 \ - --bf16 -``` - -> [!NOTE] -> There is a known issue that in the first 2 steps, graph compilation takes longer than 10 seconds. -> This will be fixed in a future release. - -> [!NOTE] -> `--mediapipe` only works on Gaudi2. +You can run inference on multiple HPUs by replacing `python train_text_to_image_sdxl.py` +with `PT_HPU_RECIPE_CACHE_CONFIG=/tmp/stdxl_recipe_cache,True,1024 python ../gaudi_spawn.py --world_size train_text_to_image_sdxl.py`. ### Inference @@ -445,7 +212,7 @@ python download_train_datasets.py To launch the multi-card Stable Diffusion training, use: ```bash python ../../gaudi_spawn.py --world_size 8 --use_mpi train_dreambooth.py \ - --pretrained_model_name_or_path="CompVis/stable-diffusion-v1-4" \ + --pretrained_model_name_or_path="stabilityai/stable-diffusion-2-1" \ --instance_data_dir="dog" \ --output_dir="dog_sd" \ --class_data_dir="path-to-class-images" \ @@ -482,7 +249,7 @@ UNet or text encoder. To run the multi-card training, use: ```bash python ../../gaudi_spawn.py --world_size 8 --use_mpi train_dreambooth.py \ - --pretrained_model_name_or_path="CompVis/stable-diffusion-v1-4" \ + --pretrained_model_name_or_path="stabilityai/stable-diffusion-2-1" \ --instance_data_dir="dog" \ --output_dir="dog_sd" \ --class_data_dir="path-to-class-images" \ @@ -513,9 +280,9 @@ Similar command could be applied with `loha`, `lokr`, or `oft` adapters. You could check each adapter's specific arguments with `--help`, for example: ```bash -python3 train_dreambooth.py oft --help +python train_dreambooth.py oft --help ``` -> [!NOTE] +> [!WARNING] > Currently, the `oft` adapter is not supported in HPU graph mode, as it triggers `torch.inverse`, > causing a CPU fallback that is incompatible with HPU graph capturing. @@ -523,7 +290,7 @@ After training completes, you can use `text_to_image_generation.py` sample for i ```bash python ../text_to_image_generation.py \ - --model_name_or_path CompVis/stable-diffusion-v1-4 \ + --model_name_or_path stabilityai/stable-diffusion-2-1 \ --unet_adapter_name_or_path dog_sd/unet \ --prompts "a sks dog" \ --num_images_per_prompt 5 \ @@ -564,31 +331,11 @@ python train_dreambooth_lora_sdxl.py \ --gaudi_config_name Habana/stable-diffusion ``` -To launch Stable Diffusion XL LoRA training on a multi-card Gaudi system, use:" -```bash -python ../../gaudi_spawn.py --world_size 8 --use_mpi train_dreambooth_lora_sdxl.py \ - --pretrained_model_name_or_path="stabilityai/stable-diffusion-xl-base-1.0" \ - --instance_data_dir="dog" \ - --pretrained_vae_model_name_or_path="madebyollin/sdxl-vae-fp16-fix" \ - --output_dir="lora-trained-xl" \ - --mixed_precision="bf16" \ - --instance_prompt="a photo of sks dog" \ - --resolution=1024 \ - --train_batch_size=1 \ - --gradient_accumulation_steps=4 \ - --learning_rate=1e-4 \ - --lr_scheduler="constant" \ - --lr_warmup_steps=0 \ - --max_train_steps=500 \ - --validation_prompt="A photo of sks dog in a bucket" \ - --validation_epochs=25 \ - --seed=0 \ - --use_hpu_graphs_for_inference \ - --use_hpu_graphs_for_training \ - --gaudi_config_name Habana/stable-diffusion -``` > [!NOTE] -> To use DeepSpeed instead of MPI, replace `--use_mpi` with `--deepspeed` in the previous example +> To use DeepSpeed instead of MPI, replace `--use_mpi` with `--deepspeed` in the previous example. + +You can run inference on multiple HPUs by replacing `python train_dreambooth_lora_sdxl.py` +with `python ../gaudi_spawn.py --world_size train_dreambooth_lora_sdxl.py`. After training is completed, you can directly use `text_to_image_generation.py` sample for inference, as shown below: ```bash @@ -606,34 +353,6 @@ python ../text_to_image_generation.py \ --bf16 ``` -Alternatively, you can run inference with a simple Python script such as this: -```python -import torch -from optimum.habana import GaudiConfig -from optimum.habana.diffusers import GaudiStableDiffusionXLPipeline - -pipe = GaudiStableDiffusionXLPipeline.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", - torch_dtype=torch.bfloat16, - use_hpu_graphs=True, - use_habana=True, - gaudi_config="Habana/stable-diffusion", - sdp_on_bf16=True, -) -pipe.load_lora_weights("lora-trained-xl") - -prompt = "A photo of sks dog in a bucket" -image = pipe( - prompt, - height=1024, - width=1024, - guidance_scale=3.5, - num_inference_steps=30, - max_sequence_length=512, -).images[0] -image.save("sdxl-lora.png") -``` - ### DreamBooth LoRA Fine-Tuning with FLUX.1-dev We can use the same `dog` dataset for the following examples. @@ -665,35 +384,12 @@ python train_dreambooth_lora_flux.py \ --gaudi_config_name="Habana/stable-diffusion" ``` -To launch FLUX.1-dev LoRA training on a multi-card Gaudi system, use:" -```bash -python ../../gaudi_spawn.py --world_size 8 --use_mpi train_dreambooth_lora_flux.py \ - --pretrained_model_name_or_path="black-forest-labs/FLUX.1-dev" \ - --dataset="dog" \ - --prompt="a photo of sks dog" \ - --output_dir="dog_lora_flux" \ - --mixed_precision="bf16" \ - --weighting_scheme="none" \ - --resolution=1024 \ - --train_batch_size=1 \ - --learning_rate=1e-4 \ - --guidance_scale=1 \ - --report_to="tensorboard" \ - --gradient_accumulation_steps=4 \ - --gradient_checkpointing \ - --lr_scheduler="constant" \ - --lr_warmup_steps=0 \ - --cache_latents \ - --rank=4 \ - --max_train_steps=500 \ - --seed="0" \ - --use_hpu_graphs_for_inference \ - --use_hpu_graphs_for_training \ - --gaudi_config_name="Habana/stable-diffusion" -``` > [!NOTE] > To use DeepSpeed instead of MPI, replace `--use_mpi` with `--use_deepspeed` in the previous example +You can run inference on multiple HPUs by replacing `python train_dreambooth_lora_flux.py` +with `python ../gaudi_spawn.py --world_size train_dreambooth_lora_flux.py`. + After training completes, you could directly use `text_to_image_generation.py` sample for inference as follows: ```bash python ../text_to_image_generation.py \ @@ -709,30 +405,3 @@ python ../text_to_image_generation.py \ --sdp_on_bf16 \ --bf16 ``` - -Alternatively, you can run inference on Gaudi system with a simple Python script like this: -```python -import torch -from optimum.habana import GaudiConfig -from optimum.habana.diffusers import GaudiFluxPipeline - -pipe = GaudiFluxPipeline.from_pretrained( - "black-forest-labs/FLUX.1-dev", - torch_dtype=torch.bfloat16, - use_hpu_graphs=True, - use_habana=True, - gaudi_config="Habana/stable-diffusion", - sdp_on_bf16=True, -) -pipe.load_lora_weights("dog_lora_flux") - -prompt = "A photo of sks dog in a bucket" -image = pipe( - prompt, - height=1024, - width=1024, - guidance_scale=3.5, - num_inference_steps=30, -).images[0] -image.save("flux-dev.png") -``` diff --git a/tests/test_diffusers.py b/tests/test_diffusers.py index 396dc8f35e..557d047d88 100644 --- a/tests/test_diffusers.py +++ b/tests/test_diffusers.py @@ -208,6 +208,14 @@ def check_8xhpu(test_case): return pytest.mark.skipif(skip, reason="test requires 8xHPU multi-card system")(test_case) +def legacy(test_case): + """ + Decorator used to skip tests for legacy models + """ + skip = os.environ.get("RUN_DIFFUSERS_LEGACY", "0") != "1" + return pytest.mark.skipif(skip, reason="This test is for old/legacy model. Skipped starting 1.16.0.")(test_case) + + class GaudiPipelineUtilsTester(TestCase): """ Tests the features added on top of diffusers/pipeline_utils.py. @@ -627,6 +635,7 @@ def test_stable_diffusion_hpu_graphs(self): self.assertEqual(images[-1].shape, (64, 64, 3)) @slow + @legacy def test_no_throughput_regression_bf16(self): prompts = [ "An image of a squirrel in Picasso style", @@ -677,6 +686,7 @@ def test_no_throughput_regression_bf16(self): @custom_bf16_ops @slow + @legacy def test_no_throughput_regression_autocast(self): prompts = [ "An image of a squirrel in Picasso style", @@ -710,6 +720,7 @@ def test_no_throughput_regression_autocast(self): @custom_bf16_ops @slow + @legacy def test_no_generation_regression_ldm3d(self): prompts = [ "An image of a squirrel in Picasso style", @@ -800,6 +811,7 @@ def test_no_generation_regression_upscale(self): @slow @check_8xhpu + @legacy def test_sd_textual_inversion(self): path_to_script = ( Path(os.path.dirname(__file__)).parent @@ -2470,6 +2482,7 @@ def test_depth2img_pipeline_hpu_graphs(self): assert images[0].shape == (32, 32, 3) @slow + @legacy def test_depth2img_pipeline(self): gaudi_config = GaudiConfig(use_torch_autocast=True) model_name = "stabilityai/stable-diffusion-2-depth" @@ -2610,6 +2623,7 @@ def test_script_train_controlnet(self): @slow @check_8xhpu + @legacy def test_train_controlnet(self): with tempfile.TemporaryDirectory() as tmpdir: path_to_script = ( @@ -5025,6 +5039,7 @@ def test_inference_batch_single_identical(self): super().test_inference_batch_single_identical(expected_max_diff=3e-3) @slow + @legacy def test_stable_diffusion_inpaint_no_throughput_regression(self): """Test that stable diffusion inpainting no throughput regression autocast""" @@ -5814,6 +5829,7 @@ def test_ddpmpipline_hpu_graphs(self): self.assertEqual(np.array(images[-1]).shape, (256, 256, 3)) @slow + @legacy def test_no_throughput_regression_bf16(self): batch_size = 16 # use batch size 16 as the baseline model_name = "google/ddpm-ema-celebahq-256" From fe65b051b360ad381e42b7f86682f2e2af0d21fc Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Fri, 21 Feb 2025 17:45:35 +0000 Subject: [PATCH 51/89] Switch version number --- .github/workflows/fast_tests.yml | 8 ++--- .github/workflows/slow_tests.yml | 36 +++++++++---------- .github/workflows/slow_tests_gaudi2.yml | 36 +++++++++---------- Makefile | 4 +-- README.md | 4 +-- docs/Dockerfile | 2 +- docs/source/installation.mdx | 2 +- docs/source/quickstart.mdx | 14 ++++---- docs/source/usage_guides/deepspeed.mdx | 4 +-- examples/gaudi_spawn.py | 2 +- examples/kubernetes/Dockerfile | 6 ++-- examples/kubernetes/README.md | 6 ++-- examples/kubernetes/README.md.gotmpl | 6 ++-- examples/kubernetes/docker-compose.yaml | 18 +++++----- examples/multi-node-training/EFA/Dockerfile | 4 +-- .../multi-node-training/GaudiNIC/Dockerfile | 4 +-- examples/speech-recognition/README.md | 2 +- examples/text-generation/README.md | 18 +++++----- .../text-generation-pipeline/README.md | 2 +- notebooks/AI_HW_Summit_2022.ipynb | 2 +- optimum/habana/accelerate/accelerator.py | 2 +- optimum/habana/accelerate/state.py | 2 +- optimum/habana/utils.py | 2 +- 23 files changed, 93 insertions(+), 93 deletions(-) diff --git a/.github/workflows/fast_tests.yml b/.github/workflows/fast_tests.yml index cdd7d1dbf5..5a1e982926 100644 --- a/.github/workflows/fast_tests.yml +++ b/.github/workflows/fast_tests.yml @@ -21,7 +21,7 @@ jobs: uses: actions/checkout@v2 - name: Pull image run: | - docker pull vault.habana.ai/gaudi-docker/1.19.0/ubuntu22.04/habanalabs/pytorch-installer-2.5.1:latest + docker pull vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest - name: Run tests run: | docker run \ @@ -36,7 +36,7 @@ jobs: --cap-add=sys_nice \ --net=host \ --ipc=host \ - vault.habana.ai/gaudi-docker/1.19.0/ubuntu22.04/habanalabs/pytorch-installer-2.5.1:latest \ + vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest \ /bin/bash tests/ci/fast_tests.sh diffusers: name: Run tests for optimum.habana.diffusers @@ -46,7 +46,7 @@ jobs: uses: actions/checkout@v2 - name: Pull image run: | - docker pull vault.habana.ai/gaudi-docker/1.19.0/ubuntu22.04/habanalabs/pytorch-installer-2.5.1:latest + docker pull vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest - name: Run tests run: | docker run \ @@ -61,5 +61,5 @@ jobs: --cap-add=sys_nice \ --net=host \ --ipc=host \ - vault.habana.ai/gaudi-docker/1.19.0/ubuntu22.04/habanalabs/pytorch-installer-2.5.1:latest \ + vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest \ /bin/bash tests/ci/fast_tests_diffusers.sh diff --git a/.github/workflows/slow_tests.yml b/.github/workflows/slow_tests.yml index d0fcb85051..e7fb736923 100644 --- a/.github/workflows/slow_tests.yml +++ b/.github/workflows/slow_tests.yml @@ -19,7 +19,7 @@ jobs: uses: actions/checkout@v2 - name: Pull image run: | - docker pull vault.habana.ai/gaudi-docker/1.19.0/ubuntu22.04/habanalabs/pytorch-installer-2.5.1:latest + docker pull vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest - name: Run tests run: | docker run \ @@ -31,7 +31,7 @@ jobs: --cap-add=sys_nice \ --net=host \ --ipc=host \ - vault.habana.ai/gaudi-docker/1.19.0/ubuntu22.04/habanalabs/pytorch-installer-2.5.1:latest \ + vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest \ /bin/bash tests/ci/example_diff_tests.sh stable-diffusion: name: Test Stable Diffusion @@ -45,7 +45,7 @@ jobs: uses: actions/checkout@v2 - name: Pull image run: | - docker pull vault.habana.ai/gaudi-docker/1.19.0/ubuntu22.04/habanalabs/pytorch-installer-2.5.1:latest + docker pull vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest - name: Run tests run: | docker run \ @@ -57,7 +57,7 @@ jobs: --cap-add=sys_nice \ --net=host \ --ipc=host \ - vault.habana.ai/gaudi-docker/1.19.0/ubuntu22.04/habanalabs/pytorch-installer-2.5.1:latest \ + vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest \ /bin/bash tests/ci/slow_tests_diffusers.sh deepspeed: name: Test DeepSpeed models @@ -72,7 +72,7 @@ jobs: uses: actions/checkout@v2 - name: Pull image run: | - docker pull vault.habana.ai/gaudi-docker/1.19.0/ubuntu22.04/habanalabs/pytorch-installer-2.5.1:latest + docker pull vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest - name: Run tests run: | docker run \ @@ -84,7 +84,7 @@ jobs: --cap-add=sys_nice \ --net=host \ --ipc=host \ - vault.habana.ai/gaudi-docker/1.19.0/ubuntu22.04/habanalabs/pytorch-installer-2.5.1:latest \ + vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest \ /bin/bash tests/ci/slow_tests_deepspeed.sh multi-card: name: Test multi-card models @@ -99,7 +99,7 @@ jobs: uses: actions/checkout@v2 - name: Pull image run: | - docker pull vault.habana.ai/gaudi-docker/1.19.0/ubuntu22.04/habanalabs/pytorch-installer-2.5.1:latest + docker pull vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest - name: Run tests run: | docker run \ @@ -111,7 +111,7 @@ jobs: --cap-add=sys_nice \ --net=host \ --ipc=host \ - vault.habana.ai/gaudi-docker/1.19.0/ubuntu22.04/habanalabs/pytorch-installer-2.5.1:latest \ + vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest \ /bin/bash tests/ci/slow_tests_8x.sh single-card: name: Test single-card models @@ -127,7 +127,7 @@ jobs: uses: actions/checkout@v2 - name: Pull image run: | - docker pull vault.habana.ai/gaudi-docker/1.19.0/ubuntu22.04/habanalabs/pytorch-installer-2.5.1:latest + docker pull vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest - name: Run tests run: | docker run \ @@ -139,7 +139,7 @@ jobs: --cap-add=sys_nice \ --net=host \ --ipc=host \ - vault.habana.ai/gaudi-docker/1.19.0/ubuntu22.04/habanalabs/pytorch-installer-2.5.1:latest \ + vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest \ /bin/bash tests/ci/slow_tests_1x.sh albert-xxl-single-card: name: Test single-card ALBERT XXL @@ -158,7 +158,7 @@ jobs: - name: Pull image if: github.event.schedule == '0 21 * * 6' run: | - docker pull vault.habana.ai/gaudi-docker/1.19.0/ubuntu22.04/habanalabs/pytorch-installer-2.5.1:latest + docker pull vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest - name: Run test if: github.event.schedule == '0 21 * * 6' run: | @@ -171,7 +171,7 @@ jobs: --cap-add=sys_nice \ --net=host \ --ipc=host \ - vault.habana.ai/gaudi-docker/1.19.0/ubuntu22.04/habanalabs/pytorch-installer-2.5.1:latest \ + vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest \ /bin/bash tests/ci/albert_xxl_1x.sh - name: Warning if: github.event.schedule != '0 21 * * 6' @@ -192,7 +192,7 @@ jobs: uses: actions/checkout@v2 - name: Pull image run: | - docker pull vault.habana.ai/gaudi-docker/1.19.0/ubuntu22.04/habanalabs/pytorch-installer-2.5.1:latest + docker pull vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest - name: Run tests run: | docker run \ @@ -204,7 +204,7 @@ jobs: --cap-add=sys_nice \ --net=host \ --ipc=host \ - vault.habana.ai/gaudi-docker/1.19.0/ubuntu22.04/habanalabs/pytorch-installer-2.5.1:latest \ + vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest \ make slow_tests_text_generation_example TOKEN=${{ secrets.TEXT_GENERATION_CI_HUB_TOKEN }} trl: name: Test TRL integration @@ -223,7 +223,7 @@ jobs: uses: actions/checkout@v2 - name: Pull image run: | - docker pull vault.habana.ai/gaudi-docker/1.19.0/ubuntu22.04/habanalabs/pytorch-installer-2.5.1:latest + docker pull vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest - name: Run tests run: | docker run \ @@ -235,7 +235,7 @@ jobs: --cap-add=sys_nice \ --net=host \ --ipc=host \ - vault.habana.ai/gaudi-docker/1.19.0/ubuntu22.04/habanalabs/pytorch-installer-2.5.1:latest \ + vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest \ /bin/bash tests/ci/slow_tests_trl.sh sentence-transformers: name: Test Sentence Transformers integration @@ -263,7 +263,7 @@ jobs: path: sentence-transformers - name: Pull image run: | - docker pull vault.habana.ai/gaudi-docker/1.19.0/ubuntu22.04/habanalabs/pytorch-installer-2.5.1:latest + docker pull vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest - name: Run tests run: | docker run \ @@ -275,5 +275,5 @@ jobs: --cap-add=sys_nice \ --net=host \ --ipc=host \ - vault.habana.ai/gaudi-docker/1.19.0/ubuntu22.04/habanalabs/pytorch-installer-2.5.1:latest \ + vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest \ /bin/bash optimum-habana/tests/ci/sentence_transformers.sh diff --git a/.github/workflows/slow_tests_gaudi2.yml b/.github/workflows/slow_tests_gaudi2.yml index 86b50d6e2c..c5b7dbbb2c 100644 --- a/.github/workflows/slow_tests_gaudi2.yml +++ b/.github/workflows/slow_tests_gaudi2.yml @@ -17,7 +17,7 @@ jobs: uses: actions/checkout@v2 - name: Pull image run: | - docker pull vault.habana.ai/gaudi-docker/1.19.0/ubuntu22.04/habanalabs/pytorch-installer-2.5.1:latest + docker pull vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest - name: Run tests run: | docker run \ @@ -30,7 +30,7 @@ jobs: --cap-add=sys_nice \ --net=host \ --ipc=host \ - vault.habana.ai/gaudi-docker/1.19.0/ubuntu22.04/habanalabs/pytorch-installer-2.5.1:latest \ + vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest \ /bin/bash tests/ci/example_diff_tests.sh stable-diffusion: name: Test Stable Diffusion @@ -43,7 +43,7 @@ jobs: uses: actions/checkout@v2 - name: Pull image run: | - docker pull vault.habana.ai/gaudi-docker/1.19.0/ubuntu22.04/habanalabs/pytorch-installer-2.5.1:latest + docker pull vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest - name: Run tests run: | docker run \ @@ -59,7 +59,7 @@ jobs: --cap-add=sys_nice \ --net=host \ --ipc=host \ - vault.habana.ai/gaudi-docker/1.19.0/ubuntu22.04/habanalabs/pytorch-installer-2.5.1:latest \ + vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest \ /bin/bash tests/ci/slow_tests_diffusers.sh ${{ secrets.TEXT_GENERATION_CI_HUB_TOKEN }} deepspeed: name: Test DeepSpeed models @@ -72,7 +72,7 @@ jobs: uses: actions/checkout@v2 - name: Pull image run: | - docker pull vault.habana.ai/gaudi-docker/1.19.0/ubuntu22.04/habanalabs/pytorch-installer-2.5.1:latest + docker pull vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest - name: Run tests run: | docker run \ @@ -88,7 +88,7 @@ jobs: --cap-add=sys_nice \ --net=host \ --ipc=host \ - vault.habana.ai/gaudi-docker/1.19.0/ubuntu22.04/habanalabs/pytorch-installer-2.5.1:latest \ + vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest \ /bin/bash tests/ci/slow_tests_deepspeed.sh ${{ secrets.TEXT_GENERATION_CI_HUB_TOKEN }} fsdp: name: Test FSDP models @@ -101,7 +101,7 @@ jobs: uses: actions/checkout@v2 - name: Pull image run: | - docker pull vault.habana.ai/gaudi-docker/1.19.0/ubuntu22.04/habanalabs/pytorch-installer-2.5.1:latest + docker pull vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest - name: Run tests run: | docker run \ @@ -117,7 +117,7 @@ jobs: --cap-add=sys_nice \ --net=host \ --ipc=host \ - vault.habana.ai/gaudi-docker/1.19.0/ubuntu22.04/habanalabs/pytorch-installer-2.5.1:latest \ + vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest \ make slow_tests_fsdp TOKEN=${{ secrets.TEXT_GENERATION_CI_HUB_TOKEN }} multi-card: name: Test multi-card models @@ -130,7 +130,7 @@ jobs: uses: actions/checkout@v2 - name: Pull image run: | - docker pull vault.habana.ai/gaudi-docker/1.19.0/ubuntu22.04/habanalabs/pytorch-installer-2.5.1:latest + docker pull vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest - name: Run tests run: | docker run \ @@ -146,7 +146,7 @@ jobs: --cap-add=sys_nice \ --net=host \ --ipc=host \ - vault.habana.ai/gaudi-docker/1.19.0/ubuntu22.04/habanalabs/pytorch-installer-2.5.1:latest \ + vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest \ /bin/bash tests/ci/slow_tests_8x.sh ${{ secrets.TEXT_GENERATION_CI_HUB_TOKEN }} single-card: name: Test single-card models @@ -160,7 +160,7 @@ jobs: uses: actions/checkout@v2 - name: Pull image run: | - docker pull vault.habana.ai/gaudi-docker/1.19.0/ubuntu22.04/habanalabs/pytorch-installer-2.5.1:latest + docker pull vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest - name: Run tests run: | docker run \ @@ -177,7 +177,7 @@ jobs: --cap-add=sys_nice \ --net=host \ --ipc=host \ - vault.habana.ai/gaudi-docker/1.19.0/ubuntu22.04/habanalabs/pytorch-installer-2.5.1:latest \ + vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest \ /bin/bash tests/ci/slow_tests_1x.sh text-generation: name: Test text-generation example @@ -192,7 +192,7 @@ jobs: uses: actions/checkout@v2 - name: Pull image run: | - docker pull vault.habana.ai/gaudi-docker/1.19.0/ubuntu22.04/habanalabs/pytorch-installer-2.5.1:latest + docker pull vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest - name: Run tests run: | docker run \ @@ -208,7 +208,7 @@ jobs: --cap-add=sys_nice \ --net=host \ --ipc=host \ - vault.habana.ai/gaudi-docker/1.19.0/ubuntu22.04/habanalabs/pytorch-installer-2.5.1:latest \ + vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest \ make slow_tests_text_generation_example TOKEN=${{ secrets.TEXT_GENERATION_CI_HUB_TOKEN }} trl: name: Test TRL integration @@ -221,7 +221,7 @@ jobs: uses: actions/checkout@v2 - name: Pull image run: | - docker pull vault.habana.ai/gaudi-docker/1.19.0/ubuntu22.04/habanalabs/pytorch-installer-2.5.1:latest + docker pull vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest - name: Run tests run: | docker run \ @@ -237,7 +237,7 @@ jobs: --cap-add=sys_nice \ --net=host \ --ipc=host \ - vault.habana.ai/gaudi-docker/1.19.0/ubuntu22.04/habanalabs/pytorch-installer-2.5.1:latest \ + vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest \ /bin/bash tests/ci/slow_tests_trl.sh sentence-transformers: name: Test Sentence Transformers integration @@ -258,7 +258,7 @@ jobs: path: sentence-transformers - name: Pull image run: | - docker pull vault.habana.ai/gaudi-docker/1.19.0/ubuntu22.04/habanalabs/pytorch-installer-2.5.1:latest + docker pull vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest - name: Run tests run: | docker run \ @@ -274,5 +274,5 @@ jobs: --cap-add=sys_nice \ --net=host \ --ipc=host \ - vault.habana.ai/gaudi-docker/1.19.0/ubuntu22.04/habanalabs/pytorch-installer-2.5.1:latest \ + vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest \ /bin/bash optimum-habana/tests/ci/sentence_transformers.sh diff --git a/Makefile b/Makefile index 2b74ed3fe1..8065ba4b69 100644 --- a/Makefile +++ b/Makefile @@ -93,7 +93,7 @@ slow_tests_8x: test_installs # Run DeepSpeed non-regression tests slow_tests_deepspeed: test_installs - python -m pip install git+https://github.com/HabanaAI/DeepSpeed.git@1.19.0 + python -m pip install git+https://github.com/HabanaAI/DeepSpeed.git@1.20.0 python -m pytest tests/test_examples.py -v -s -k "deepspeed" slow_tests_diffusers: test_installs @@ -109,7 +109,7 @@ slow_tests_diffusers: test_installs slow_tests_text_generation_example: test_installs python -m pip install -r examples/text-generation/requirements_awq.txt BUILD_CUDA_EXT=0 python -m pip install -vvv --no-build-isolation git+https://github.com/HabanaAI/AutoGPTQ.git - python -m pip install git+https://github.com/HabanaAI/DeepSpeed.git@1.19.0 + python -m pip install git+https://github.com/HabanaAI/DeepSpeed.git@1.20.0 python -m pytest tests/test_text_generation_example.py tests/test_encoder_decoder.py -v -s --token $(TOKEN) # Run image-to-text non-regression tests diff --git a/README.md b/README.md index dc4636d308..751b698dd4 100644 --- a/README.md +++ b/README.md @@ -48,7 +48,7 @@ Please refer to the Intel Gaudi AI Accelerator official [installation guide](htt > Tests should be run in a Docker container based on Intel Gaudi's official images. Instructions to > obtain the latest containers from the Intel Gaudi Vault are available > [here](https://docs.habana.ai/en/latest/Installation_Guide/Additional_Installation/Docker_Installation.html#use-intel-gaudi-containers). -> The current Optimum for Intel Gaudi has been validated with Intel Gaudi v1.19 stack. +> The current Optimum for Intel Gaudi has been validated with Intel Gaudi v1.20 stack. ## Install the library and get example scripts @@ -95,7 +95,7 @@ git clone -b transformers_future https://github.com/huggingface/optimum-habana To use DeepSpeed on HPUs, you also need to run the following command: ```bash -pip install git+https://github.com/HabanaAI/DeepSpeed.git@1.19.0 +pip install git+https://github.com/HabanaAI/DeepSpeed.git@1.20.0 ``` To install the requirements for every example: diff --git a/docs/Dockerfile b/docs/Dockerfile index 060b7413dc..ead30b7412 100644 --- a/docs/Dockerfile +++ b/docs/Dockerfile @@ -1,4 +1,4 @@ -FROM vault.habana.ai/gaudi-docker/1.19.0/ubuntu22.04/habanalabs/pytorch-installer-2.5.1:latest +FROM vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest ARG commit_sha ARG clone_url diff --git a/docs/source/installation.mdx b/docs/source/installation.mdx index fa54c4446e..6b39fa1084 100644 --- a/docs/source/installation.mdx +++ b/docs/source/installation.mdx @@ -24,7 +24,7 @@ python -m pip install --upgrade-strategy eager optimum[habana] To use Microsoft® DeepSpeed with Intel Gaudi devices, you also need to run the following command: ```bash -python -m pip install git+https://github.com/HabanaAI/DeepSpeed.git@1.19.0 +python -m pip install git+https://github.com/HabanaAI/DeepSpeed.git@1.20.0 ``` To ensure that you are installing the correct Intel Gaudi Software, please run the `hl-smi` command to confirm the software version diff --git a/docs/source/quickstart.mdx b/docs/source/quickstart.mdx index c882de2629..57d0bf90cb 100644 --- a/docs/source/quickstart.mdx +++ b/docs/source/quickstart.mdx @@ -32,12 +32,12 @@ platform for deep learning and follow the steps to start and connect to the node ## Docker Setup Now that you have access to the node, you will use the latest Intel Gaudi AI Accelerator docker image by executing the docker run command which will -automatically download and run the docker. At the time of writing this guide, latest Gaudi docker version was 1.19.0: +automatically download and run the docker. At the time of writing this guide, latest Gaudi docker version was 1.20.0: ```bash -release=1.19.0 +release=1.20.0 os=ubuntu22.04 -torch=2.5.1 +torch=2.6.0 docker_image=vault.habana.ai/gaudi-docker/$release/$os/habanalabs/pytorch-installer-$torch:latest ``` @@ -65,11 +65,11 @@ docker run -itd \ ## Optimum for Intel Gaudi Setup Check latest release of Optimum for Intel Gaudi [here](https://github.com/huggingface/optimum-habana/releases). -At the time of writing this guide, latest Optimum for Intel Gaudi release version was v1.15.0, which is paired with Intel Gaudi Software release -version 1.19.0. Install Optimum for Intel Gaudi as follows: +At the time of writing this guide, latest Optimum for Intel Gaudi release version was v1.16.0, which is paired with Intel Gaudi Software release +version 1.20.0. Install Optimum for Intel Gaudi as follows: ```bash -git clone -b v1.15.0 https://github.com/huggingface/optimum-habana +git clone -b v1.16.0 https://github.com/huggingface/optimum-habana pip install ./optimum-habana ``` @@ -115,7 +115,7 @@ Microsoft® DeepSpeed. Gaudi-specific fork of the library is maintained by Intel To install the library compatible with the same Gaudi software release stack, use: ```bash -pip install git+https://github.com/HabanaAI/DeepSpeed.git@1.19.0 +pip install git+https://github.com/HabanaAI/DeepSpeed.git@1.20.0 ``` With DeepSpeed successfully installed we can now run a distributed GPT-2 inference on an 8 HPU system as follows: diff --git a/docs/source/usage_guides/deepspeed.mdx b/docs/source/usage_guides/deepspeed.mdx index f6617e92ce..6fc34f2261 100644 --- a/docs/source/usage_guides/deepspeed.mdx +++ b/docs/source/usage_guides/deepspeed.mdx @@ -32,7 +32,7 @@ You can find more information about DeepSpeed Gaudi integration [here](https://d To use DeepSpeed on Gaudi, you need to install Optimum for Intel Gaudi and [DeepSpeed fork for Intel Gaudi](https://github.com/HabanaAI/DeepSpeed) with: ```bash pip install optimum[habana] -pip install git+https://github.com/HabanaAI/DeepSpeed.git@1.19.0 +pip install git+https://github.com/HabanaAI/DeepSpeed.git@1.20.0 ``` @@ -79,7 +79,7 @@ It is strongly advised to read [this section](https://huggingface.co/docs/transf -Other examples of configurations for HPUs are proposed [here](https://github.com/HabanaAI/Model-References/tree/1.19.0/PyTorch/nlp/DeepSpeedExamples/deepspeed-bert/scripts) by Intel. +Other examples of configurations for HPUs are proposed [here](https://github.com/HabanaAI/Model-References/tree/1.20.0/PyTorch/nlp/DeepSpeedExamples/deepspeed-bert/scripts) by Intel. The [Transformers documentation](https://huggingface.co/docs/transformers/main_classes/deepspeed#configuration) explains how to write a configuration from scratch very well. A more complete description of all configuration possibilities is available [here](https://www.deepspeed.ai/docs/config-json/). diff --git a/examples/gaudi_spawn.py b/examples/gaudi_spawn.py index f282809a31..6817ca0565 100644 --- a/examples/gaudi_spawn.py +++ b/examples/gaudi_spawn.py @@ -84,7 +84,7 @@ def main(): if not is_deepspeed_available(): raise ImportError( "--use_deepspeed requires deepspeed: `pip install" - " git+https://github.com/HabanaAI/DeepSpeed.git@1.19.0`." + " git+https://github.com/HabanaAI/DeepSpeed.git@1.20.0`." ) # Patch sys.argv diff --git a/examples/kubernetes/Dockerfile b/examples/kubernetes/Dockerfile index 7ebfd93894..2264dfca57 100644 --- a/examples/kubernetes/Dockerfile +++ b/examples/kubernetes/Dockerfile @@ -1,7 +1,7 @@ -ARG GAUDI_SW_VER=1.19.0 +ARG GAUDI_SW_VER=1.20.0 ARG OS=ubuntu22.04 -ARG TORCH_VER=2.5.1 -ARG OPTIMUM_HABANA_VER=1.15.0 +ARG TORCH_VER=2.6.0 +ARG OPTIMUM_HABANA_VER=1.16.0 FROM vault.habana.ai/gaudi-docker/${GAUDI_SW_VER}/${OS}/habanalabs/pytorch-installer-${TORCH_VER}:latest AS optimum-habana diff --git a/examples/kubernetes/README.md b/examples/kubernetes/README.md index 06f4f01d09..fe65d41482 100644 --- a/examples/kubernetes/README.md +++ b/examples/kubernetes/README.md @@ -43,12 +43,12 @@ Use the the following commands to build the containers: ```bash # Specify the Gaudi SW version, OS, and PyTorch version which will be used for the base container -export GAUDI_SW_VER=1.19.0 +export GAUDI_SW_VER=1.20.0 export OS=ubuntu22.04 -export TORCH_VER=2.5.1 +export TORCH_VER=2.6.0 # Specify the version of optimum-habana to install in the container -export OPTIMUM_HABANA_VER=1.15.0 +export OPTIMUM_HABANA_VER=1.16.0 git clone https://github.com/huggingface/optimum-habana.git diff --git a/examples/kubernetes/README.md.gotmpl b/examples/kubernetes/README.md.gotmpl index 431f8ad611..48f0af8259 100644 --- a/examples/kubernetes/README.md.gotmpl +++ b/examples/kubernetes/README.md.gotmpl @@ -43,12 +43,12 @@ Use the the following commands to build the containers: ```bash # Specify the Gaudi SW version, OS, and PyTorch version which will be used for the base container -export GAUDI_SW_VER=1.19.0 +export GAUDI_SW_VER=1.20.0 export OS=ubuntu22.04 -export TORCH_VER=2.5.1 +export TORCH_VER=2.6.0 # Specify the version of optimum-habana to install in the container -export OPTIMUM_HABANA_VER=1.15.0 +export OPTIMUM_HABANA_VER=1.16.0 git clone https://github.com/huggingface/optimum-habana.git diff --git a/examples/kubernetes/docker-compose.yaml b/examples/kubernetes/docker-compose.yaml index 6bdea75bbd..4ab69f1021 100644 --- a/examples/kubernetes/docker-compose.yaml +++ b/examples/kubernetes/docker-compose.yaml @@ -5,30 +5,30 @@ services: http_proxy: ${http_proxy:-""} https_proxy: ${https_proxy:-""} no_proxy: ${no_proxy:-""} - GAUDI_SW_VER: ${GAUDI_SW_VER:-1.19.0} + GAUDI_SW_VER: ${GAUDI_SW_VER:-1.20.0} OS: ${OS:-ubuntu22.04} - OPTIMUM_HABANA_VER: ${OPTIMUM_HABANA_VER:-1.15.0} - TORCH_VER: ${TORCH_VER:-2.5.1} + OPTIMUM_HABANA_VER: ${OPTIMUM_HABANA_VER:-1.16.0} + TORCH_VER: ${TORCH_VER:-2.6.0} REGISTRY: ${REGISTRY} REPO: ${REPO} context: . labels: - org.opencontainers.base.name: "vault.habana.ai/gaudi-docker/${GAUDI_SW_VER:-1.19.0}/${OS:-ubuntu22.04}/habanalabs/pytorch-installer-${TORCH_VER:-2.5.1}:latest" + org.opencontainers.base.name: "vault.habana.ai/gaudi-docker/${GAUDI_SW_VER:-1.20.0}/${OS:-ubuntu22.04}/habanalabs/pytorch-installer-${TORCH_VER:-2.6.0}:latest" org.opencontainers.image.title: "Optimum for Intel® Gaudi® Accelerators" - org.opencontainers.image.version: gaudi-${GAUDI_SW_VER:-1.19.0}-optimum-habana-${OPTIMUM_HABANA_VER:-1.15.0} + org.opencontainers.image.version: gaudi-${GAUDI_SW_VER:-1.20.0}-optimum-habana-${OPTIMUM_HABANA_VER:-1.16.0} command: > sh -c "python -c 'from optimum import habana; print(\"optimum-habana:\", habana.__version__)'" - image: ${REGISTRY}/${REPO}:b-${GITHUB_RUN_NUMBER:-0}-gaudi-${GAUDI_SW_VER:-1.19.0}-optimum-habana-${OPTIMUM_HABANA_VER:-1.15.0} + image: ${REGISTRY}/${REPO}:b-${GITHUB_RUN_NUMBER:-0}-gaudi-${GAUDI_SW_VER:-1.20.0}-optimum-habana-${OPTIMUM_HABANA_VER:-1.16.0} pull_policy: always optimum-habana-examples: build: labels: - org.opencontainers.base.name: "${REGISTRY}/${REPO}:gaudi-${GAUDI_SW_VER:-1.19.0}-optimum-habana-${OPTIMUM_HABANA_VER:-1.15.0}" + org.opencontainers.base.name: "${REGISTRY}/${REPO}:gaudi-${GAUDI_SW_VER:-1.20.0}-optimum-habana-${OPTIMUM_HABANA_VER:-1.16.0}" org.opencontainers.image.title: "Optimum for Intel® Gaudi® Accelerators Examples" - org.opencontainers.image.version: gaudi-${GAUDI_SW_VER:-1.19.0}-optimum-habana-examples-${OPTIMUM_HABANA_VER:-1.15.0} + org.opencontainers.image.version: gaudi-${GAUDI_SW_VER:-1.20.0}-optimum-habana-examples-${OPTIMUM_HABANA_VER:-1.16.0} target: optimum-habana-examples command: > sh -c "python -c 'from optimum import habana; print(\"optimum-habana:\", habana.__version__)'" extends: optimum-habana - image: ${REGISTRY}/${REPO}:b-${GITHUB_RUN_NUMBER:-0}-gaudi-${GAUDI_SW_VER:-1.19.0}-optimum-habana-examples-${OPTIMUM_HABANA_VER:-1.15.0} + image: ${REGISTRY}/${REPO}:b-${GITHUB_RUN_NUMBER:-0}-gaudi-${GAUDI_SW_VER:-1.20.0}-optimum-habana-examples-${OPTIMUM_HABANA_VER:-1.16.0} diff --git a/examples/multi-node-training/EFA/Dockerfile b/examples/multi-node-training/EFA/Dockerfile index bc6f827164..8b83af7d9d 100644 --- a/examples/multi-node-training/EFA/Dockerfile +++ b/examples/multi-node-training/EFA/Dockerfile @@ -1,4 +1,4 @@ -FROM vault.habana.ai/gaudi-docker/1.19.0/ubuntu22.04/habanalabs/pytorch-installer-2.5.1:latest +FROM vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest # Installs pdsh and upgrade pip RUN apt-get update && apt-get install -y pdsh && \ @@ -19,7 +19,7 @@ RUN sed -i 's/#Port 22/Port 3022/g' /etc/ssh/sshd_config && \ # Installs Optimum Habana and Habana's fork of DeepSpeed RUN pip install optimum[habana] && \ - pip install git+https://github.com/HabanaAI/DeepSpeed.git@1.19.0 + pip install git+https://github.com/HabanaAI/DeepSpeed.git@1.20.0 CMD ssh-keygen -t rsa -b 4096 -N '' -f ~/.ssh/id_rsa && \ chmod 600 ~/.ssh/id_rsa && \ diff --git a/examples/multi-node-training/GaudiNIC/Dockerfile b/examples/multi-node-training/GaudiNIC/Dockerfile index 5375a6fcc7..09a98e6bb9 100644 --- a/examples/multi-node-training/GaudiNIC/Dockerfile +++ b/examples/multi-node-training/GaudiNIC/Dockerfile @@ -1,4 +1,4 @@ -FROM vault.habana.ai/gaudi-docker/1.19.0/ubuntu22.04/habanalabs/pytorch-installer-2.5.1:latest +FROM vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest # Installs pdsh and upgrade pip RUN apt-get update && apt-get install -y pdsh && \ @@ -13,7 +13,7 @@ RUN sed -i 's/#Port 22/Port 3022/g' /etc/ssh/sshd_config && \ # Installs Optimum Habana and Habana's fork of DeepSpeed RUN pip install optimum[habana] && \ - pip install git+https://github.com/HabanaAI/DeepSpeed.git@1.19.0 + pip install git+https://github.com/HabanaAI/DeepSpeed.git@1.20.0 CMD ssh-keygen -t rsa -b 4096 -N '' -f ~/.ssh/id_rsa && \ chmod 600 ~/.ssh/id_rsa && \ diff --git a/examples/speech-recognition/README.md b/examples/speech-recognition/README.md index 1f0f8fbe38..d51d990db7 100644 --- a/examples/speech-recognition/README.md +++ b/examples/speech-recognition/README.md @@ -145,7 +145,7 @@ On 8 HPUs, this script should run in *ca.* 49 minutes and yield a CTC loss of ** > You need to install DeepSpeed with: > ```bash -> pip install git+https://github.com/HabanaAI/DeepSpeed.git@1.19.0 +> pip install git+https://github.com/HabanaAI/DeepSpeed.git@1.20.0 > ``` DeepSpeed can be used with almost the same command as for a multi-card run: diff --git a/examples/text-generation/README.md b/examples/text-generation/README.md index 5adf348217..daaa44aac6 100755 --- a/examples/text-generation/README.md +++ b/examples/text-generation/README.md @@ -33,7 +33,7 @@ pip install -r requirements_lm_eval.txt Then, if you plan to use [DeepSpeed-inference](https://docs.habana.ai/en/latest/PyTorch/DeepSpeed/Inference_Using_DeepSpeed.html) (e.g. to use BLOOM/BLOOMZ), you should install DeepSpeed as follows: ```bash -pip install git+https://github.com/HabanaAI/DeepSpeed.git@1.19.0 +pip install git+https://github.com/HabanaAI/DeepSpeed.git@1.20.0 ``` @@ -204,14 +204,14 @@ python ../gaudi_spawn.py --use_deepspeed --world_size 8 run_generation.py \ To run Deepseek-R1-BF16 inference on 16 Gaudi3 cards (2 nodes) use the following command. Ensure you replace the hostfile parameter with the appropriate file. Sample hostfile reference [here](https://github.com/huggingface/optimum-habana/blob/main/examples/multi-node-training/hostfile) ```bash -python3 ../gaudi_spawn.py --hostfile= --use_deepspeed \ ---world_size 16 ./run_generation.py \ ---model_name_or_path opensourcerelease/DeepSeek-R1-bf16 \ ---bf16 \ +python3 ../gaudi_spawn.py --hostfile= --use_deepspeed \ +--world_size 16 ./run_generation.py \ +--model_name_or_path opensourcerelease/DeepSeek-R1-bf16 \ +--bf16 \ --trim_logits \ ---batch_size 1 \ ---use_hpu_graphs \ ---use_kv_cache \ +--batch_size 1 \ +--use_hpu_graphs \ +--use_kv_cache \ --parallel_strategy "ep" \ --prompt "DeepSpeed is a machine learning framework" ``` @@ -637,7 +637,7 @@ python run_generation.py \ ### Saving FP8 Checkpoints in Hugging Face format After quantizing the model, we can save it to a local path. -> [!NOTE] +> [!NOTE] > Before executing the command below, please refer to the [Running with FP8](#running-with-fp8) section to measure the model quantization statistics. Here is an example of how to quantize and save the LLama3.1-70B model on two cards: diff --git a/examples/text-generation/text-generation-pipeline/README.md b/examples/text-generation/text-generation-pipeline/README.md index 2aa036ec3a..ec28462501 100644 --- a/examples/text-generation/text-generation-pipeline/README.md +++ b/examples/text-generation/text-generation-pipeline/README.md @@ -22,7 +22,7 @@ The text-generation pipeline can be used to perform text-generation by providing If you plan to use [DeepSpeed-inference](https://docs.habana.ai/en/latest/PyTorch/DeepSpeed/Inference_Using_DeepSpeed.html), you should install DeepSpeed as follows: ```bash -pip install git+https://github.com/HabanaAI/DeepSpeed.git@1.19.0 +pip install git+https://github.com/HabanaAI/DeepSpeed.git@1.20.0 ``` If you would like to use the pipeline with LangChain classes, you can install LangChain as follows: diff --git a/notebooks/AI_HW_Summit_2022.ipynb b/notebooks/AI_HW_Summit_2022.ipynb index 4ebb252cf3..0b0f34c8f2 100644 --- a/notebooks/AI_HW_Summit_2022.ipynb +++ b/notebooks/AI_HW_Summit_2022.ipynb @@ -262,7 +262,7 @@ "metadata": {}, "outputs": [], "source": [ - "!pip install git+https://github.com/HabanaAI/DeepSpeed.git@1.19.0" + "!pip install git+https://github.com/HabanaAI/DeepSpeed.git@1.20.0" ] }, { diff --git a/optimum/habana/accelerate/accelerator.py b/optimum/habana/accelerate/accelerator.py index 8566c9a7e5..de027eff8e 100644 --- a/optimum/habana/accelerate/accelerator.py +++ b/optimum/habana/accelerate/accelerator.py @@ -157,7 +157,7 @@ def __init__( if deepspeed_plugin: if not is_deepspeed_available(): raise ImportError( - "DeepSpeed is not installed => run `pip install git+https://github.com/HabanaAI/DeepSpeed.git@1.19.0`." + "DeepSpeed is not installed => run `pip install git+https://github.com/HabanaAI/DeepSpeed.git@1.20.0`." ) mixed_precision = ( diff --git a/optimum/habana/accelerate/state.py b/optimum/habana/accelerate/state.py index c5d241e384..a50d5039fe 100644 --- a/optimum/habana/accelerate/state.py +++ b/optimum/habana/accelerate/state.py @@ -57,7 +57,7 @@ def __init__(self, cpu: bool = False, **kwargs): if not is_deepspeed_available(): raise ImportError( "DeepSpeed is not available, install it with: `pip install" - " git+https://github.com/HabanaAI/DeepSpeed.git@1.19.0`." + " git+https://github.com/HabanaAI/DeepSpeed.git@1.20.0`." ) self.distributed_type = GaudiDistributedType.DEEPSPEED import deepspeed diff --git a/optimum/habana/utils.py b/optimum/habana/utils.py index 244b52e203..65354380d4 100755 --- a/optimum/habana/utils.py +++ b/optimum/habana/utils.py @@ -31,7 +31,7 @@ logger = logging.get_logger(__name__) -CURRENTLY_VALIDATED_SYNAPSE_VERSION = version.parse("1.19.0") +CURRENTLY_VALIDATED_SYNAPSE_VERSION = version.parse("1.20.0") def to_device_dtype(my_input: Any, target_device: torch.device = None, target_dtype: torch.dtype = None): From ffda2a03ef6fa73c232040c66bbbdfef71f779eb Mon Sep 17 00:00:00 2001 From: Silvia Colabrese Date: Fri, 28 Feb 2025 10:48:15 +0100 Subject: [PATCH 52/89] Temporary WA for get_type error (#1806) Co-authored-by: Yaser Afshar --- examples/text-generation/run_lm_eval.py | 30 +++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/examples/text-generation/run_lm_eval.py b/examples/text-generation/run_lm_eval.py index 4e536dc757..7f0797489f 100644 --- a/examples/text-generation/run_lm_eval.py +++ b/examples/text-generation/run_lm_eval.py @@ -214,6 +214,36 @@ def _model_call(self, inps: torch.Tensor) -> torch.Tensor: logits = logits.to(torch.float32) return logits + def get_model_info(self) -> dict: + """ + Patched method to get Hugging Face model information for experiment reproducibility. + source: https://github.com/EleutherAI/lm-evaluation-harness/blob/v0.4.7/lm_eval/models/huggingface.py/#L1375 + Remove from SynapseAI 1.21 + """ + + def get_model_num_params(model) -> int: + if hasattr(model, "num_parameters"): + return model.num_parameters() + elif hasattr(model, "parameters"): + return sum(p.numel() for p in model.parameters()) + else: + return -1 + + def get_model_dtype(model) -> str: + if hasattr(model, "dtype"): + return model.dtype + elif hasattr(model, "parameters"): + return next(model.parameters()).dtype + else: + return "" + + model_info = { + "model_num_parameters": get_model_num_params(self._model), + "model_dtype": get_model_dtype(self._model), + "model_revision": self.revision, + } + return model_info + def main() -> None: # Modified based on cli_evaluate function in https://github.com/EleutherAI/lm-evaluation-harness/blob/v0.4.7/lm_eval/__main__.py/#L268 From 167a218f18f84ba53a3ce606cb67fca36cb3dc28 Mon Sep 17 00:00:00 2001 From: Yaser Afshar Date: Wed, 5 Mar 2025 15:20:45 -0800 Subject: [PATCH 53/89] Loss Computation for Compatibility with Transformers 4.48.3 (#1794) --- optimum/habana/distributed/contextparallel.py | 52 ++++- .../models/llama/modeling_llama.py | 29 +-- optimum/habana/transformers/trainer.py | 188 ++++++++++++------ 3 files changed, 183 insertions(+), 86 deletions(-) diff --git a/optimum/habana/distributed/contextparallel.py b/optimum/habana/distributed/contextparallel.py index 2020b6a84e..66d10d0f72 100644 --- a/optimum/habana/distributed/contextparallel.py +++ b/optimum/habana/distributed/contextparallel.py @@ -7,13 +7,26 @@ ) -# Gather losses across context parallel group -class _ContextParallelLoss(torch.autograd.Function): +class ContextParallelLossFunction(torch.autograd.Function): + """ + Gather losses across context parallel group. + + This custom autograd function is designed to handle the distribution of loss computation + across multiple parallel contexts in a distributed training setup. It ensures that the loss + is gathered from all devices involved in the parallel context, allowing for consistent and + accurate computation of the overall loss. + + The forward method gathers the loss from all ranks in the context parallel group, while the + backward method ensures that gradients are correctly synchronized across the different parallel + contexts. + """ + @staticmethod def forward(ctx, loss): ctx.seqlen = loss.size(0) * get_sequence_parallel_world_size() - + # Create a tensor to gather all losses from context parallel group loss_all = torch.empty(ctx.seqlen, dtype=loss.dtype, device=loss.device) + # Gather losses from all ranks in the group torch.distributed.all_gather_into_tensor(loss_all, loss, group=get_sequence_parallel_group()) return loss_all @@ -21,10 +34,37 @@ def forward(ctx, loss): def backward(ctx, grad_output): step_seqlen = ctx.seqlen // get_sequence_parallel_world_size() sp_rank = get_sequence_parallel_rank() + # Extract the relevant part of the gradient for this rank grad_output_part = grad_output[step_seqlen * sp_rank : step_seqlen * (sp_rank + 1)] - return grad_output_part, None -def _get_loss_from_context_parallel(vocab_parallel_loss): - return _ContextParallelLoss.apply(vocab_parallel_loss) +def fixed_cross_entropy(source, target, num_items_in_batch: int = None, ignore_index: int = -100, **kwargs): + loss_all = torch.nn.functional.cross_entropy(source, target, ignore_index=ignore_index, reduction="none") + # Apply context parallel loss + loss_all = ContextParallelLossFunction.apply(loss_all) + if num_items_in_batch is None: + loss = torch.mean(loss_all) + else: + loss = torch.sum(loss_all) / num_items_in_batch + return loss + + +def ForCausalLMContextParallelLoss( + logits, labels, vocab_size: int, num_items_in_batch: int = None, ignore_index: int = -100, **kwargs +): + # Upcast to float if we need to compute the loss to avoid potential precision issues + logits = logits.float() + labels = labels.to(logits.device) + # Shift so that tokens < n predict n + shift_logits = logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + + # Flatten the tokens + shift_logits = shift_logits.view(-1, vocab_size) + shift_labels = shift_labels.view(-1) + # Enable model parallelism + shift_labels = shift_labels.to(shift_logits.device) + + loss = fixed_cross_entropy(shift_logits, shift_labels, num_items_in_batch, ignore_index, **kwargs) + return loss diff --git a/optimum/habana/transformers/models/llama/modeling_llama.py b/optimum/habana/transformers/models/llama/modeling_llama.py index 242b3b8fa3..3b40bb6ce9 100755 --- a/optimum/habana/transformers/models/llama/modeling_llama.py +++ b/optimum/habana/transformers/models/llama/modeling_llama.py @@ -1418,6 +1418,10 @@ class GaudiLlamaForCausalLM(LlamaForCausalLM): def __init__(self, config, parallel_strategy: DistributedStrategy = NoOpStrategy): config.parallel_strategy = parallel_strategy super().__init__(config) + if parallel_state.sequence_parallel_is_initialized() and parallel_state.get_sequence_parallel_world_size() > 1: + from ....distributed.contextparallel import ForCausalLMContextParallelLoss + + self._loss_function = ForCausalLMContextParallelLoss def allocate_kv_cache(self, batch_size, max_seq_len, inp_seq_len): self.model.allocate_kv_cache(batch_size, max_seq_len, inp_seq_len) @@ -1506,30 +1510,7 @@ def forward( 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 so that tokens < n predict n - shift_logits = logits[..., :-1, :].contiguous() - shift_labels = labels[..., 1:].contiguous() - # Flatten the tokens - loss_fct = torch.nn.CrossEntropyLoss() - shift_logits = shift_logits.view(-1, self.config.vocab_size) - shift_labels = shift_labels.view(-1) - # Enable model parallelism - shift_labels = shift_labels.to(shift_logits.device) - # Collect losses from context parallel group - # Each rank in group calculates loss on partial outputs - if ( - parallel_state.sequence_parallel_is_initialized() - and parallel_state.get_sequence_parallel_world_size() > 1 - ): - from ....distributed.contextparallel import _get_loss_from_context_parallel - - loss_fct = torch.nn.CrossEntropyLoss(reduction="none") - loss_all = _get_loss_from_context_parallel(loss_fct(shift_logits, shift_labels)) - loss = torch.mean(loss_all) - else: - loss = loss_fct(shift_logits, shift_labels) + loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs) if not return_dict: output = (logits,) + outputs[1:] diff --git a/optimum/habana/transformers/trainer.py b/optimum/habana/transformers/trainer.py index 6f186e521c..62761944a9 100644 --- a/optimum/habana/transformers/trainer.py +++ b/optimum/habana/transformers/trainer.py @@ -207,8 +207,8 @@ def _get_input_update_settings(model, lazy_mode: Optional[bool] = None) -> Tuple TRAINER_STATE_NAME = "trainer_state.json" OPTIMIZER_NAME = "optimizer.pt" OPTIMIZER_NAME_BIN = "optimizer.bin" -SCHEDULER_NAME = "scheduler.pt" SCALER_NAME = "scaler.pt" +SCHEDULER_NAME = "scheduler.pt" class GaudiTrainer(Trainer): @@ -450,6 +450,9 @@ def _tune_save_checkpoint(self, checkpoint_dir: str): output_dir = os.path.join(checkpoint_dir, f"{PREFIX_CHECKPOINT_DIR}-{self.state.global_step}") self.save_model(output_dir, _internal_call=True) if self.args.should_save: + # TODO + # Update the `TrainerControl` state to where we are currently + # self.state.stateful_callbacks["TrainerControl"] = self.control.state() self.state.save_to_json(os.path.join(output_dir, TRAINER_STATE_NAME)) torch.save(self.optimizer.state_dict(), os.path.join(output_dir, OPTIMIZER_NAME)) torch.save(self.lr_scheduler.state_dict(), os.path.join(output_dir, SCHEDULER_NAME)) @@ -467,13 +470,22 @@ def _wrap_model(self, model, training=True, dataloader=None): if self.args.parallel_mode == ParallelMode.DISTRIBUTED and self.args.distribution_strategy == "ddp": kwargs = {} - kwargs["find_unused_parameters"] = self.args.ddp_find_unused_parameters - if self.args.ddp_find_unused_parameters and self.args.gradient_checkpointing: - logger.warning( - "ddp_find_unused_parameters and gradient_checkpointing are both True, which may lead to an error:" - " https://github.com/huggingface/transformers/pull/4659#issuecomment-643356021" - ) - kwargs["bucket_cap_mb"] = self.args.ddp_bucket_cap_mb + if self.args.ddp_find_unused_parameters is not None: + kwargs["find_unused_parameters"] = self.args.ddp_find_unused_parameters + if self.args.ddp_find_unused_parameters and self.args.gradient_checkpointing: + logger.warning( + "ddp_find_unused_parameters and gradient_checkpointing are both True, which may lead to an error:" + " https://github.com/huggingface/transformers/pull/4659#issuecomment-643356021" + ) + elif isinstance(model, PreTrainedModel): + # find_unused_parameters breaks checkpointing as per + # https://github.com/huggingface/transformers/pull/4659#issuecomment-643356021 + kwargs["find_unused_parameters"] = not model.is_gradient_checkpointing + else: + kwargs["find_unused_parameters"] = True + + if self.args.ddp_bucket_cap_mb is not None: + kwargs["bucket_cap_mb"] = self.args.ddp_bucket_cap_mb if self.args.use_habana: kwargs["gradient_as_bucket_view"] = True @@ -499,6 +511,7 @@ def train( ): """ Main training entry point. + Args: resume_from_checkpoint (`str` or `bool`, *optional*): If a `str`, local path to a saved checkpoint as saved by a previous instance of [`Trainer`]. If a @@ -541,7 +554,7 @@ def train( FutureWarning, ) if len(kwargs) > 0: - raise TypeError(f"train() received got unexpected keyword arguments: {', '.join(list(kwargs.keys()))}.") + raise TypeError(f"train() got unexpected keyword arguments: {', '.join(list(kwargs.keys()))}.") # This might change the seed so needs to run first. self._hp_search_setup(trial) self._train_batch_size = self.args.train_batch_size @@ -826,18 +839,15 @@ def hpu_deepspeed_checkpointing(function, *checkpoint_args, use_reentrant: Optio # Check if saved optimizer or scheduler states exist self._load_optimizer_and_scheduler(resume_from_checkpoint) - if self.gaudi_config.use_fused_clip_norm: + if self.gaudi_config.use_fused_clip_norm and self.args.use_habana: try: from habana_frameworks.torch.hpex.normalization import FusedClipNorm except ImportError as error: - error.msg = ( - f"Could not import 'FusedClipNorm' from 'habana_frameworks.torch.hpex.normalization'. {error.msg}." - ) + error.msg = f"Could not import habana_frameworks.torch.hpex.normalization. {error.msg}." raise error - self.FusedNorm = FusedClipNorm( - model.parameters(), - args.max_grad_norm, - ) + self.FusedNorm = FusedClipNorm(model.parameters(), args.max_grad_norm) + else: + self.FusedNorm = None # important: at this point: # self.model is the Transformers Model @@ -924,9 +934,10 @@ def hpu_deepspeed_checkpointing(function, *checkpoint_args, use_reentrant: Optio self._total_loss_scalar = 0.0 self._globalstep_last_logged = self.state.global_step self._zero_model_grad(model) - _grad_norm: Optional[float] = None - _should_compute_grad_norm: bool = not self.accelerator.distributed_type == GaudiDistributedType.DEEPSPEED and ( - # Gradient clipping + grad_norm: Optional[float] = None + + # Gradient clipping + _should_compute_grad_norm: bool = self.accelerator.distributed_type != GaudiDistributedType.DEEPSPEED and ( args.max_grad_norm is not None and args.max_grad_norm > 0 ) @@ -944,6 +955,16 @@ def hpu_deepspeed_checkpointing(function, *checkpoint_args, use_reentrant: Optio else: self.log_evaluate_save_time = None + # Calculate the number of items in each batch for all epochs + num_items_in_batches = self.get_num_items_in_batches( + args, + epochs_trained, + num_train_epochs, + train_dataloader, + len_dataloader, + num_examples, + ) + hb_profiler = HabanaProfile( warmup=self.args.profiling_warmup_steps, active=self.args.profiling_steps, @@ -992,10 +1013,13 @@ def hpu_deepspeed_checkpointing(function, *checkpoint_args, use_reentrant: Optio 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 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) + batch_samples = self.get_iterator_batch_samples(epoch_iterator, num_batches) + num_items_in_batch = num_items_in_batches[epoch][update_step] for i, inputs in enumerate(batch_samples): step += 1 @@ -1008,10 +1032,7 @@ def hpu_deepspeed_checkpointing(function, *checkpoint_args, use_reentrant: Optio do_sync_step = (step + 1) % args.gradient_accumulation_steps == 0 or (step + 1) == steps_in_epoch # Since we perform prefetching, we need to manually set sync_gradients - if not do_sync_step: - self.accelerator.gradient_state._set_sync_gradients(False) - else: - self.accelerator.gradient_state._set_sync_gradients(True) + self.accelerator.gradient_state._set_sync_gradients(do_sync_step) if self.args.include_num_input_tokens_seen: main_input_name = getattr(self.model, "main_input_name", "input_ids") @@ -1073,15 +1094,16 @@ def hpu_deepspeed_checkpointing(function, *checkpoint_args, use_reentrant: Optio if args.logging_nan_inf_filter and (torch.isnan(tr_loss_step) or torch.isinf(tr_loss_step)): # if loss is nan or inf simply add the average of previous logged losses - tr_loss += tr_loss / (1 + self.state.global_step - self._globalstep_last_logged) + tr_loss = tr_loss + tr_loss / (1 + self.state.global_step - self._globalstep_last_logged) else: if tr_loss.device != tr_loss_step.device: raise ValueError( f"Calculated loss must be on the original device: {tr_loss.device} but device in use is {tr_loss_step.device}" ) - tr_loss += tr_loss_step + tr_loss = tr_loss + tr_loss_step self.current_flos += float(self.floating_point_ops(inputs)) + if args.use_lazy_mode: self.htcore.mark_step() @@ -1089,15 +1111,15 @@ def hpu_deepspeed_checkpointing(function, *checkpoint_args, use_reentrant: Optio # Since we perform prefetching, we need to manually set sync_gradients to True self.accelerator.gradient_state._set_sync_gradients(True) - # If the condition is true, we need to compute _grad_norm + # If the condition is true, we need to compute grad_norm, deepspeed does its own clipping if _should_compute_grad_norm: - # deepspeed does its own clipping - if self.gaudi_config.use_fused_clip_norm and args.use_habana: + # Gradient clipping + if self.FusedNorm is not None: # TODO: to merge self.accelerator.clip_grad_norm_ when HMP is removed - _grad_norm = self.FusedNorm.clip_norm(model.parameters()) + grad_norm = self.FusedNorm.clip_norm(model.parameters()) else: # Revert to normal clipping otherwise - _grad_norm = self.accelerator.clip_grad_norm_( + grad_norm = self.accelerator.clip_grad_norm_( model.parameters(), args.max_grad_norm, ) @@ -1121,7 +1143,7 @@ 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 ) else: self.control = self.callback_handler.on_substep_end(args, self.state, self.control) @@ -1141,7 +1163,7 @@ 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) if self.control.should_training_stop: break @@ -1297,6 +1319,8 @@ def _load_best_model(self): ) # If the model is on the GPU, it still works! + # workaround for FSDP bug https://github.com/pytorch/pytorch/issues/82963 + # which takes *args instead of **kwargs load_result = model.load_state_dict(state_dict, False) if has_been_loaded: @@ -1324,6 +1348,7 @@ def _maybe_log_save_evaluate(self, tr_loss, _grad_norm, model, trial, epoch, ign # reset tr_loss to zero tr_loss -= tr_loss + logs["loss"] = round(tr_loss_scalar / (self.state.global_step - self._globalstep_last_logged), 4) # This grad_norm block was outside of _maybe_log_save_evaluate method causing perf degradation. @@ -1351,7 +1376,7 @@ def _maybe_log_save_evaluate(self, tr_loss, _grad_norm, model, trial, epoch, ign self._globalstep_last_logged = self.state.global_step self.store_flos() - self.log(logs, start_time=start_time) + self.log(logs, start_time) metrics = None if self.control.should_evaluate: @@ -1531,7 +1556,9 @@ def _load_optimizer_and_scheduler(self, checkpoint): def log(self, logs: Dict[str, float], start_time: Optional[float] = None) -> None: """ Log `logs` on the various objects watching training. + Subclass and override this method to inject custom behavior. + Args: logs (`Dict[str, float]`): The values to log. @@ -1586,7 +1613,9 @@ def _prepare_input(self, data: Union[torch.Tensor, Any]) -> Union[torch.Tensor, def autocast_smart_context_manager(self, cache_enabled: Optional[bool] = True): """ A helper wrapper that creates an appropriate context manager for `autocast` while feeding it the desired - arguments, depending on the situation. Modified by Habana to enable using `autocast` on Gaudi devices. + arguments, depending on the situation. + + Modified by Habana to enable using `autocast` on Gaudi devices. """ if self.use_cpu_amp: ctx_manager = torch.autocast(device_type="cpu", dtype=torch.bfloat16, cache_enabled=cache_enabled) @@ -1623,6 +1652,7 @@ def training_step( `torch.Tensor`: The tensor with training loss on this batch. """ model.train() + # TODO # if hasattr(self.optimizer, "train") and callable(self.optimizer.train): # self.optimizer.train() @@ -1645,8 +1675,7 @@ def training_step( 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) or (num_items_in_batch is None): - # TODO refer to todo in function get_batch_samples_transformers - + 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 @@ -2300,6 +2329,7 @@ def prediction_loop( ) -> EvalLoopOutput: """ Prediction/evaluation loop, shared by `Trainer.evaluate()` and `Trainer.predict()`. + Works both with or without labels. """ args = self.args @@ -2334,6 +2364,7 @@ def prediction_loop( self.deepspeed = self.model_wrapped model.eval() + # TODO # if hasattr(self.optimizer, "eval") and callable(self.optimizer.eval): # self.optimizer.eval() @@ -2518,6 +2549,7 @@ def create_accelerator_and_postprocess(self): ) if is_accelerate_available("1.1.0"): dataloader_config.data_seed = self.args.data_seed + non_blocking = accelerator_config.pop("non_blocking") if non_blocking and not self.args.dataloader_pin_memory: logger.warning( @@ -2616,30 +2648,74 @@ def _zero_model_grad(self, model): model.zero_grad() model._zero_grad_kwargs = {} - def get_batch_samples_transformers(self, epoch_iterator, num_batches): + def get_num_items_in_batches( + self, args, epochs_trained, num_train_epochs, train_dataloader, len_dataloader, num_examples + ): """ - Added "_transformers" at the end of the method name to avoid a wrong call to a similarly named method in TRL trainers. + Calculate the number of items in each batch for all epochs during training. """ - batch_samples = [] - num_items_in_batch = None - for _ in range(num_batches): + steps_in_epoch = ( + len_dataloader if len_dataloader is not None else args.max_steps * args.gradient_accumulation_steps + ) + + remainder = num_examples % args.gradient_accumulation_steps + if remainder == 0: + remainder = args.gradient_accumulation_steps + + total_updates = steps_in_epoch // args.gradient_accumulation_steps + 1 + if args.gradient_accumulation_steps == 1: + total_updates -= 1 + + num_items_in_batches = [] + for epoch in range(epochs_trained, num_train_epochs): + epoch_dataloader = train_dataloader + if hasattr(epoch_dataloader, "set_epoch"): + epoch_dataloader.set_epoch(epoch) + + epoch_iterator = iter(epoch_dataloader) try: - batch_samples += [next(epoch_iterator)] + first_batch = next(epoch_iterator) except StopIteration: break + # Check if the batch contains "labels" (once per epoch) + if "labels" not in first_batch: + num_items_in_batches.append([None] * total_updates) + continue + + device = first_batch["labels"].device - # TODO: execute get_batch_samples outside of the training loop (before training) and uncomment the following lines - # if len(batch_samples) > 0 and "labels" in batch_samples[0]: - # # For now we don't support object detection - # try: - # num_items_in_batch = sum([(batch["labels"].ne(-100)).sum() for batch in batch_samples]) - # except (TypeError, AttributeError): - # pass + # Reset the iterator + epoch_iterator = iter(epoch_dataloader) + + num_items_in_batches.append([]) + for update_step in range(total_updates): + num_batches = args.gradient_accumulation_steps if update_step != (total_updates - 1) else remainder + + num_items_in_batch = 0 + for _ in range(num_batches): + try: + batch = next(epoch_iterator) + num_items_in_batch += (batch["labels"].ne(-100)).sum().item() + except StopIteration: + break + + if self.args.average_tokens_across_devices and num_items_in_batch > 0: + num_items_in_batch = torch.tensor(num_items_in_batch, device=device) + num_items_in_batch = self.accelerator.gather(num_items_in_batch).sum().item() + + # Set to None if no items in batch + if num_items_in_batch == 0: + num_items_in_batch = None - # if self.args.average_tokens_across_devices and num_items_in_batch is not None: - # num_items_in_batch = self.accelerator.gather(num_items_in_batch).sum().item() + num_items_in_batches[epoch].append(num_items_in_batch) - # if torch.is_tensor(num_items_in_batch): - # num_items_in_batch = num_items_in_batch.item() + return num_items_in_batches - return batch_samples, num_items_in_batch + def get_iterator_batch_samples(self, epoch_iterator, num_batches): + batch_samples = [] + for _ in range(num_batches): + try: + batch_samples += [next(epoch_iterator)] + except StopIteration: + break + return batch_samples From 379524c1c77a66cdfbd0d68606aa8a5def15bda0 Mon Sep 17 00:00:00 2001 From: Mieszko Dziadowiec Date: Thu, 6 Mar 2025 00:29:33 +0100 Subject: [PATCH 54/89] Move model to device before wrapping with FSDP (#1801) --- optimum/habana/accelerate/accelerator.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/optimum/habana/accelerate/accelerator.py b/optimum/habana/accelerate/accelerator.py index 8566c9a7e5..7b397822a1 100644 --- a/optimum/habana/accelerate/accelerator.py +++ b/optimum/habana/accelerate/accelerator.py @@ -476,6 +476,9 @@ def prepare_model(self, model: torch.nn.Module, device_placement: bool = None, e "limit_all_gathers": fsdp_plugin.limit_all_gathers, "device_id": torch.device("hpu", torch.hpu.current_device()), } + # There's issue with moving view tensors to device within FSDP class [See: https://github.com/pytorch/pytorch/issues/147321] + # Due to above issue, view tensor's may lead to silent incorrent behavior, while pretending to be view they're really not + model = model.to(kwargs["device_id"]) model = FSDP(model, **kwargs) if fsdp_plugin.activation_checkpointing: from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import ( From 46bad3ba41c183d74c1891bdf6183c7dbd569602 Mon Sep 17 00:00:00 2001 From: Dmitry Date: Wed, 5 Mar 2025 23:59:03 +0100 Subject: [PATCH 55/89] v1.16 Llama3-405B text-generation. Added DEEPSPEED_USE_HABANA_FRAMEWORKS_DETERMINISTIC_API flag (#1812) --- examples/text-generation/README.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/examples/text-generation/README.md b/examples/text-generation/README.md index 7fa3e5ca70..c2966329f9 100755 --- a/examples/text-generation/README.md +++ b/examples/text-generation/README.md @@ -190,6 +190,7 @@ python ../gaudi_spawn.py --use_deepspeed --world_size 8 run_generation.py \ To run Llama3-405B inference on 8 Gaudi3 cards use the following command: ```bash +DEEPSPEED_USE_HABANA_FRAMEWORKS_DETERMINISTIC_API=true \ ENABLE_LB_BUNDLE_ALL_COMPUTE_MME=0 ENABLE_EXPERIMENTAL_FLAGS=1 \ python ../gaudi_spawn.py --use_deepspeed --world_size 8 run_generation.py \ --model_name_or_path meta-llama/Llama-3.1-405B-Instruct \ @@ -496,6 +497,7 @@ QUANT_CONFIG=./quantization_config/maxabs_quant.json python ../gaudi_spawn.py \ Here is an example to measure the tensor quantization statistics on Llama3-405B with 8 cards: > Please note that Llama3-405B requires minimum 16 cards Gaudi2 and 8 cards Gaudi3. ```bash +DEEPSPEED_USE_HABANA_FRAMEWORKS_DETERMINISTIC_API=true \ QUANT_CONFIG=./quantization_config/maxabs_measure_include_outputs.json python ../gaudi_spawn.py \ --use_deepspeed --world_size 8 run_lm_eval.py \ -o acc_llama3_405b_bs1_quant.txt \ @@ -514,6 +516,7 @@ QUANT_CONFIG=./quantization_config/maxabs_measure_include_outputs.json python .. Here is an example to quantize the model based on previous measurements for Llama3-405B with 8 cards: > Please note that Llama3-405B requires minimum 16 cards Gaudi2 and 8 cards Gaudi3. ```bash +DEEPSPEED_USE_HABANA_FRAMEWORKS_DETERMINISTIC_API=true \ QUANT_CONFIG=./quantization_config/maxabs_quant.json python ../gaudi_spawn.py \ --use_deepspeed --world_size 8 run_generation.py \ --model_name_or_path meta-llama/Llama-3.1-405B-Instruct \ From 00782271fbf1e360b6c7a302cf0099e2c0cf0314 Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Thu, 6 Mar 2025 09:22:11 +0000 Subject: [PATCH 56/89] Make style --- optimum/habana/transformers/models/mixtral/modeling_mixtral.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/optimum/habana/transformers/models/mixtral/modeling_mixtral.py b/optimum/habana/transformers/models/mixtral/modeling_mixtral.py index 8873bc8402..d84b44dbab 100644 --- a/optimum/habana/transformers/models/mixtral/modeling_mixtral.py +++ b/optimum/habana/transformers/models/mixtral/modeling_mixtral.py @@ -144,7 +144,8 @@ def gaudi_mixtral_repeat_kv( class GaudiMixtralSparseMoeBlock(MixtralSparseMoeBlock): def forward(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: - batch_size, sequence_length, hidden_dim = hidden_states.shape + original_shape = hidden_states.shape + hidden_dim = original_shape[2] if self.training and self.jitter_noise > 0: hidden_states *= torch.empty_like(hidden_states).uniform_(1.0 - self.jitter_noise, 1.0 + self.jitter_noise) hidden_states = hidden_states.view(-1, hidden_dim) From 81f33ed4171dea9f4d5db870ffee9a80169a477a Mon Sep 17 00:00:00 2001 From: Urszula Golowicz Date: Thu, 6 Mar 2025 17:49:30 +0100 Subject: [PATCH 57/89] Revert placing llama on cpu (#1827) --- examples/text-generation/README.md | 3 --- .../quantization_config/unit_scale_quant.json | 7 +------ examples/text-generation/utils.py | 8 +------- 3 files changed, 2 insertions(+), 16 deletions(-) diff --git a/examples/text-generation/README.md b/examples/text-generation/README.md index c2966329f9..7fa3e5ca70 100755 --- a/examples/text-generation/README.md +++ b/examples/text-generation/README.md @@ -190,7 +190,6 @@ python ../gaudi_spawn.py --use_deepspeed --world_size 8 run_generation.py \ To run Llama3-405B inference on 8 Gaudi3 cards use the following command: ```bash -DEEPSPEED_USE_HABANA_FRAMEWORKS_DETERMINISTIC_API=true \ ENABLE_LB_BUNDLE_ALL_COMPUTE_MME=0 ENABLE_EXPERIMENTAL_FLAGS=1 \ python ../gaudi_spawn.py --use_deepspeed --world_size 8 run_generation.py \ --model_name_or_path meta-llama/Llama-3.1-405B-Instruct \ @@ -497,7 +496,6 @@ QUANT_CONFIG=./quantization_config/maxabs_quant.json python ../gaudi_spawn.py \ Here is an example to measure the tensor quantization statistics on Llama3-405B with 8 cards: > Please note that Llama3-405B requires minimum 16 cards Gaudi2 and 8 cards Gaudi3. ```bash -DEEPSPEED_USE_HABANA_FRAMEWORKS_DETERMINISTIC_API=true \ QUANT_CONFIG=./quantization_config/maxabs_measure_include_outputs.json python ../gaudi_spawn.py \ --use_deepspeed --world_size 8 run_lm_eval.py \ -o acc_llama3_405b_bs1_quant.txt \ @@ -516,7 +514,6 @@ QUANT_CONFIG=./quantization_config/maxabs_measure_include_outputs.json python .. Here is an example to quantize the model based on previous measurements for Llama3-405B with 8 cards: > Please note that Llama3-405B requires minimum 16 cards Gaudi2 and 8 cards Gaudi3. ```bash -DEEPSPEED_USE_HABANA_FRAMEWORKS_DETERMINISTIC_API=true \ QUANT_CONFIG=./quantization_config/maxabs_quant.json python ../gaudi_spawn.py \ --use_deepspeed --world_size 8 run_generation.py \ --model_name_or_path meta-llama/Llama-3.1-405B-Instruct \ diff --git a/examples/text-generation/quantization_config/unit_scale_quant.json b/examples/text-generation/quantization_config/unit_scale_quant.json index 20783ea3f1..216cf27e68 100644 --- a/examples/text-generation/quantization_config/unit_scale_quant.json +++ b/examples/text-generation/quantization_config/unit_scale_quant.json @@ -3,10 +3,5 @@ "mode": "QUANTIZE", "observer": "maxabs", "scale_method": "unit_scale", - "whitelist": {"types": [], "names": []}, - "blacklist": {"types": [], "names": []}, - "quantize_weight": false, - "dump_stats_path": "./results/hk", - "ignore_modules_wo_measures": "True", - "dump_stats_xlsx_path": "./run_outputs/fp8stats.xlsx" + "dump_stats_path": "./hqt_output/measure" } diff --git a/examples/text-generation/utils.py b/examples/text-generation/utils.py index 61270ca218..63a1a32fb7 100644 --- a/examples/text-generation/utils.py +++ b/examples/text-generation/utils.py @@ -439,12 +439,7 @@ def setup_distributed_model(args, model_dtype, model_kwargs, logger): logger.info("DeepSpeed is enabled.") deepspeed.init_distributed(dist_backend="hccl") config = AutoConfig.from_pretrained(args.model_name_or_path, torch_dtype=model_dtype, **model_kwargs) - - keep_module_on_host = False - if "Llama-3.1-405B" in args.model_name_or_path: - keep_module_on_host = True - - load_to_meta = False if keep_module_on_host else model_on_meta(config) + load_to_meta = model_on_meta(config) if args.assistant_model is None: assistant_model = None @@ -499,7 +494,6 @@ def setup_distributed_model(args, model_dtype, model_kwargs, logger): # Initialize the model ds_inference_kwargs = {"dtype": model_dtype} - ds_inference_kwargs["keep_module_on_host"] = keep_module_on_host ds_inference_kwargs["tensor_parallel"] = {"tp_size": args.world_size} ds_inference_kwargs["enable_cuda_graph"] = args.use_hpu_graphs ds_inference_kwargs["injection_policy"] = get_ds_injection_policy(config) From 195fdf8485da1d8c59c67953f24101f0e85f162b Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Tue, 11 Mar 2025 10:16:09 +0000 Subject: [PATCH 58/89] Fix contrastive search --- optimum/habana/transformers/generation/utils.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/optimum/habana/transformers/generation/utils.py b/optimum/habana/transformers/generation/utils.py index 37b8a9f41a..d53ac286fb 100755 --- a/optimum/habana/transformers/generation/utils.py +++ b/optimum/habana/transformers/generation/utils.py @@ -2367,12 +2367,6 @@ def _contrastive_search( ) # contrastive_search main logic end - # synced_gpus: don't waste resources running the code we don't need; kwargs must be updated before skipping - model_kwargs = self._update_model_kwargs_for_generation( - outputs, - model_kwargs, - is_encoder_decoder=self.config.is_encoder_decoder, - ) if synced_gpus and this_peer_finished: continue @@ -2390,6 +2384,11 @@ def _contrastive_search( if streamer is not None: streamer.put(next_tokens.cpu()) + model_kwargs = self._update_model_kwargs_for_generation( + outputs, + model_kwargs, + is_encoder_decoder=self.config.is_encoder_decoder, + ) # increase cur_len cur_len = cur_len + 1 From a22b82171e940642af9ea6d8979a56977e224f1d Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Fri, 14 Mar 2025 02:55:28 -0600 Subject: [PATCH 59/89] Upgrade to Transformers v4.49 (#1810) Signed-off-by: Daniel Socek Signed-off-by: U. Artie Eoff Co-authored-by: Libin Tang Co-authored-by: Daniel Socek Co-authored-by: Nikolay Protasov Co-authored-by: Harish Subramony <81822986+hsubramony@users.noreply.github.com> Co-authored-by: U. Artie Eoff Co-authored-by: Yeonsil Yoon Co-authored-by: Luca Calabria Co-authored-by: Shiv Kaul Co-authored-by: Iman Gohari Co-authored-by: Harshvardhan Chauhan --- .../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 +- .../run_image2text_lora_finetune.py | 11 +- examples/language-modeling/run_clm.py | 2 +- examples/language-modeling/run_mlm.py | 2 +- .../run_multitask_prompt_tuning.py | 2 +- .../run_prompt_tuning_clm.py | 2 +- examples/question-answering/run_qa.py | 2 +- examples/question-answering/run_seq2seq_qa.py | 2 +- .../run_speech_recognition_ctc.py | 2 +- .../run_speech_recognition_seq2seq.py | 2 +- examples/summarization/run_summarization.py | 2 +- examples/text-classification/run_glue.py | 2 +- examples/text-generation/requirements_awq.txt | 2 +- examples/translation/run_translation.py | 2 +- .../diffusers/pipelines/pipeline_utils.py | 15 +- .../pipeline_stable_diffusion_xl_mlperf.py | 21 + .../habana/transformers/generation/utils.py | 93 +- optimum/habana/transformers/modeling_utils.py | 6 - .../modeling_utils_transformers.py | 35 - .../models/cohere/modeling_cohere.py | 5 +- .../models/falcon/modeling_falcon.py | 5 +- .../models/gemma/modeling_gemma.py | 5 +- .../models/gemma2/modeling_gemma2.py | 9 +- .../transformers/models/gpt2/modeling_gpt2.py | 2 +- .../gpt_bigcode/modeling_gpt_bigcode.py | 2 +- .../models/gpt_neox/modeling_gpt_neox.py | 55 +- .../models/idefics2/modeling_idefics2.py | 7 +- .../models/llama/modeling_llama.py | 7 +- .../models/llava/modeling_llava.py | 26 +- .../models/llava_next/modeling_llava_next.py | 19 +- .../models/mistral/modeling_mistral.py | 5 +- .../models/mixtral/modeling_mixtral.py | 5 +- .../models/mllama/modeling_mllama.py | 23 +- .../transformers/models/opt/modeling_opt.py | 34 +- .../models/paligemma/modeling_paligemma.py | 8 +- .../models/persimmon/modeling_persimmon.py | 5 +- .../transformers/models/phi/modeling_phi.py | 5 +- .../models/qwen2/modeling_qwen2.py | 5 +- .../models/qwen2_moe/modeling_qwen2_moe.py | 6 +- .../models/qwen2_vl/modeling_qwen2_vl.py | 57 +- .../models/stablelm/modeling_stablelm.py | 5 +- .../models/starcoder2/modeling_starcoder2.py | 5 +- .../video_llava/modeling_video_llava.py | 68 +- .../models/whisper/modeling_whisper.py | 2 +- .../transformers/models/xglm/modeling_xglm.py | 2 +- optimum/habana/transformers/trainer.py | 118 +- optimum/habana/transformers/training_args.py | 8 + setup.py | 2 +- .../fixture/tests/test_diffusers.json | 4 +- .../fixture/tests/test_encoder_decoder.json | 12 +- .../fixture/tests/test_examples.json | 302 ++--- .../fixture/tests/test_fsdp_examples.json | 4 +- .../tests/test_image_to_text_example.json | 24 +- .../tests/test_object_segmentation.json | 2 +- .../fixture/tests/test_openclip_vqa.json | 4 +- .../tests/test_sentence_transformers.json | 26 +- .../tests/test_text_generation_example.json | 138 +-- .../fixture/tests/test_video_llava.json | 2 +- tests/test_trainer.py | 1091 +++++++++++------ .../models/gpt_neox/test_modeling_gpt_neox.py | 3 + .../tests/test_modeling_common.py | 6 +- 64 files changed, 1341 insertions(+), 995 deletions(-) delete mode 100644 optimum/habana/transformers/modeling_utils_transformers.py diff --git a/examples/audio-classification/run_audio_classification.py b/examples/audio-classification/run_audio_classification.py index c1049e3e8e..bdae71bb99 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.48.0") +check_min_version("4.49.0") check_optimum_habana_min_version("1.17.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 d18bef8bcf..01ffac6fe9 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.48.0") +check_min_version("4.49.0") check_optimum_habana_min_version("1.17.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 43fb51457a..741408e238 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.48.0") +check_min_version("4.49.0") check_optimum_habana_min_version("1.17.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 cb88496a77..5958b9f9de 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.48.0") +check_min_version("4.49.0") check_optimum_habana_min_version("1.17.0.dev0") require_version("datasets>=2.14.0", "To fix: pip install -r examples/pytorch/image-classification/requirements.txt") diff --git a/examples/image-to-text/run_image2text_lora_finetune.py b/examples/image-to-text/run_image2text_lora_finetune.py index 0fc1911d7c..7201cd12ec 100644 --- a/examples/image-to-text/run_image2text_lora_finetune.py +++ b/examples/image-to-text/run_image2text_lora_finetune.py @@ -382,8 +382,8 @@ def eval(processor, model, dataset, batch_size, use_lazy_mode, use_hpu_graphs, m images, texts, return_tensors="pt", - padding="max_length", - truncation=True, + padding=True, + truncation=False, max_length=max_seq_length, padding_side="left", ) @@ -611,15 +611,12 @@ def main(): text = processor.apply_chat_template(messages, add_generation_prompt=True) if config.model_type == "llava": - # don't expand image_token_id - setattr(processor, "patch_size", None) - setattr(processor, "vision_feature_select_strategy", None) inputs = processor( [image], [text.strip()], return_tensors="pt", - padding="max_length", - truncation=True, + padding=True, + truncation=False, max_length=data_args.max_seq_length, padding_side="left", ) diff --git a/examples/language-modeling/run_clm.py b/examples/language-modeling/run_clm.py index f39438d6ff..e7fd5d3d83 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.48.0") +check_min_version("4.49.0") check_optimum_habana_min_version("1.17.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 32e2b93987..087e020439 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.48.0") +check_min_version("4.49.0") check_optimum_habana_min_version("1.17.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_multitask_prompt_tuning.py b/examples/language-modeling/run_multitask_prompt_tuning.py index 73e226956c..ef757cb763 100644 --- a/examples/language-modeling/run_multitask_prompt_tuning.py +++ b/examples/language-modeling/run_multitask_prompt_tuning.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 risk. -check_min_version("4.48.0") +check_min_version("4.49.0") check_optimum_habana_min_version("1.17.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 2ddf3d59a0..fd541b872f 100644 --- a/examples/language-modeling/run_prompt_tuning_clm.py +++ b/examples/language-modeling/run_prompt_tuning_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.48.0") +check_min_version("4.49.0") check_optimum_habana_min_version("1.17.0.dev0") require_version("datasets>=1.8.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 f37072fce9..ff23237c5b 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.48.0") +check_min_version("4.49.0") check_optimum_habana_min_version("1.17.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 db9cdd9f39..d9d15e76af 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.48.0") +check_min_version("4.49.0") check_optimum_habana_min_version("1.17.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 676b7e14dc..8279ff7a5d 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.48.0") +check_min_version("4.49.0") check_optimum_habana_min_version("1.17.0.dev0") require_version("datasets>=1.18.0", "To fix: pip install -r examples/pytorch/speech-recognition/requirements.txt") diff --git a/examples/speech-recognition/run_speech_recognition_seq2seq.py b/examples/speech-recognition/run_speech_recognition_seq2seq.py index 1dbe973d10..dbdd000851 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.48.0") +check_min_version("4.49.0") check_optimum_habana_min_version("1.17.0.dev0") require_version("datasets>=1.18.0", "To fix: pip install -r examples/pytorch/speech-recognition/requirements.txt") diff --git a/examples/summarization/run_summarization.py b/examples/summarization/run_summarization.py index 977e1a3644..8b64e8bdd8 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.48.0") +check_min_version("4.49.0") check_optimum_habana_min_version("1.17.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 140e1511dd..c74b08e207 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.48.0") +check_min_version("4.49.0") check_optimum_habana_min_version("1.17.0.dev0") require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/text-classification/requirements.txt") diff --git a/examples/text-generation/requirements_awq.txt b/examples/text-generation/requirements_awq.txt index dff2632403..812d48b233 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.48.2,<=4.49.0 diff --git a/examples/translation/run_translation.py b/examples/translation/run_translation.py index 86c05b7ef8..ef833a173d 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.48.0") +check_min_version("4.49.0") check_optimum_habana_min_version("1.17.0.dev0") require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/translation/requirements.txt") diff --git a/optimum/habana/diffusers/pipelines/pipeline_utils.py b/optimum/habana/diffusers/pipelines/pipeline_utils.py index 5215fd6603..24f436e3f2 100644 --- a/optimum/habana/diffusers/pipelines/pipeline_utils.py +++ b/optimum/habana/diffusers/pipelines/pipeline_utils.py @@ -381,10 +381,23 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P # Import htcore here to support model quantization import habana_frameworks.torch.core as htcore # noqa: F401 - return super().from_pretrained( + # Normally we just need to return super().from_pretrained. However this is a + # workaround for Transformers 4.49.0 issue (sub_model torch_dtype option ignored). + # Note this issue is already fixed in 4.50.0dev working branch.. + model = super().from_pretrained( pretrained_model_name_or_path, **kwargs, ) + if bf16_full_eval: + # Get the component names + component_names = [name for name in model.__dict__ if not name.startswith("_")] + # Iterate through the component names and fix dtype + for name in component_names: + component = getattr(model, name, None) + if component is not None and hasattr(component, "dtype"): + component.to(torch.bfloat16) + + return model @classmethod def save_lora_weights( diff --git a/optimum/habana/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_mlperf.py b/optimum/habana/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_mlperf.py index 3cca208954..e6f6517de0 100644 --- a/optimum/habana/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_mlperf.py +++ b/optimum/habana/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_mlperf.py @@ -260,6 +260,27 @@ def run_unet( return latents + # Normally we do not wrap from_pretrained. However this is a + # workaround for Transformers 4.49.0 issue (sub_model torch_dtype option ignored). + # Note this issue is already fixed in 4.50.0dev working branch.. + @classmethod + def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs): + bf16_full_eval = kwargs.get("torch_dtype", None) == torch.bfloat16 + model = super().from_pretrained( + pretrained_model_name_or_path, + **kwargs, + ) + if bf16_full_eval: + # Get the component names + component_names = [name for name in model.__dict__ if not name.startswith("_")] + # Iterate through the component names and fix dtype + for name in component_names: + component = getattr(model, name, None) + if component is not None and hasattr(component, "dtype"): + component.to(torch.bfloat16) + + return model + @classmethod def _split_inputs_into_batches( cls, diff --git a/optimum/habana/transformers/generation/utils.py b/optimum/habana/transformers/generation/utils.py index d53ac286fb..e8488abc69 100755 --- a/optimum/habana/transformers/generation/utils.py +++ b/optimum/habana/transformers/generation/utils.py @@ -52,6 +52,7 @@ StopStringCriteria, ) from transformers.generation.utils import ( + ALL_CACHE_NAMES, GenerateBeamDecoderOnlyOutput, GenerateBeamEncoderDecoderOutput, GenerateBeamOutput, @@ -217,9 +218,13 @@ def _prepare_inputs_for_generation( # 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 ( + 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 ): @@ -229,9 +234,9 @@ def _prepare_inputs_for_generation( # 3. Prepare base model inputs input_ids_key = "decoder_input_ids" if self.config.is_encoder_decoder else "input_ids" - # if `inputs_embeds` are passed, we only want to use them in the 1st generation step + # if `inputs_embeds` are passed, we only want to use them in the 1st generation step for every prompt. if not self.config.is_encoder_decoder: - if inputs_embeds is not None and cache_position[0] == 0: + if inputs_embeds is not None and len(cache_position) == inputs_embeds.shape[1]: model_inputs[input_ids_key] = None model_inputs["inputs_embeds"] = inputs_embeds else: @@ -242,23 +247,28 @@ 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 + attention_mask = ( + kwargs.pop("decoder_attention_mask", None) if self.config.is_encoder_decoder else attention_mask + ) + attention_mask_key = "decoder_attention_mask" if self.config.is_encoder_decoder else "attention_mask" + position_ids_key = "decoder_position_ids" if self.config.is_encoder_decoder else "position_ids" if ( attention_mask is not None - and kwargs.get("position_ids") is None - and "position_ids" in set(inspect.signature(self.forward).parameters.keys()) + and kwargs.get(position_ids_key) is None + and position_ids_key in set(inspect.signature(self.forward).parameters.keys()) ): position_ids = attention_mask.long().cumsum(-1) - 1 position_ids.masked_fill_(attention_mask == 0, 1) - kwargs["position_ids"] = position_ids # placed in kwargs for further processing (see below) + kwargs[position_ids_key] = position_ids # placed in kwargs for further processing (see below) # 5. Slice model inputs if it's an input that should have the same length as `input_ids` - for model_input_name in ["position_ids", "token_type_ids"]: + for model_input_name in ["position_ids", "token_type_ids", "decoder_position_ids"]: model_input = kwargs.get(model_input_name) if model_input is not None: if past_key_values is not None: current_input_length = ( model_inputs["inputs_embeds"].shape[1] - if model_inputs["inputs_embeds"] is not None + if model_inputs.get("inputs_embeds") is not None else model_inputs[input_ids_key].shape[1] ) model_input = model_input[:, -current_input_length:] @@ -305,7 +315,7 @@ def _prepare_inputs_for_generation( past_key_values=past_key_values, ) if attention_mask is not None: - model_inputs["attention_mask"] = attention_mask + model_inputs[attention_mask_key] = attention_mask # 7. Forward ALL kwargs that are uninitialized (e.g. `use_cache`). for key, value in kwargs.items(): @@ -571,10 +581,15 @@ def _update_model_kwargs_for_generation( model_kwargs["first_token"] = False if not model_kwargs.get("pad_done", False): # update past_key_values keeping its naming used in model code - cache_name, cache = self._extract_past_from_model_output(outputs) - model_kwargs[cache_name] = cache - if getattr(outputs, "state", None) is not None: - model_kwargs["state"] = outputs.state + for possible_cache_name in ALL_CACHE_NAMES: + if possible_cache_name in outputs: + # TODO (joao): remove output/input mismatch when these old models (xlnet, reformer) are deprecated + if possible_cache_name in ("past_buckets_states", "mems"): + cache_name = "past_key_values" + else: + cache_name = possible_cache_name + model_kwargs[cache_name] = getattr(outputs, possible_cache_name) + break # update token_type_ids with last value if "token_type_ids" in model_kwargs: @@ -836,7 +851,6 @@ def _prepare_generated_length( elif ( model_input_name == "inputs_embeds" and input_ids_length != inputs_tensor.shape[1] - and input_ids_length != 0 and not self.config.is_encoder_decoder ): generation_config.max_length -= inputs_tensor.shape[1] @@ -1415,13 +1429,13 @@ def generate( has_token_idx="token_idx" in model_kwargs, ) - # If the model supports `num_logits_to_keep` in forward(), set it to 1 to avoid computing the whole + # If the model supports `logits_to_keep` in forward(), set it to 1 to avoid computing the whole # logit matrix. This can save a lot of memory during the first forward pass. Note that assisted decoding # dynamically overrides this value as it can need more than the last token logits # # Use trim_logits in HPU to save memory (in replacement of the num_logits_to_keep) - # if self._supports_num_logits_to_keep() and "num_logits_to_keep" not in model_kwargs: - # model_kwargs["num_logits_to_keep"] = 1 + # if self._supports_logits_to_keep() and "logits_to_keep" not in model_kwargs: + # model_kwargs["logits_to_keep"] = 1 self._validate_generated_length( generation_config, @@ -1433,10 +1447,7 @@ def generate( # - `model_kwargs` may be updated in place with a cache as defined by the parameters in `generation_config`. # - different models have a different cache name expected by the model (default = "past_key_values") # - `max_length`, prepared above, is used to determine the maximum cache length - # TODO (joao): remove `user_defined_cache` after v4.47 (remove default conversion to legacy format) - cache_name = "past_key_values" if "mamba" not in self.__class__.__name__.lower() else "cache_params" - user_defined_cache = model_kwargs.get(cache_name) - max_cache_length = generation_config.max_length + max_cache_length = generation_config.max_length - 1 if ( inputs_tensor.shape[1] != input_ids_length and model_input_name == "inputs_embeds" @@ -1836,32 +1847,12 @@ def typeerror(): # Convert to legacy cache format if requested if ( - generation_config.return_legacy_cache is not False # Should check for `True` after v4.47 + generation_config.return_legacy_cache is True and not is_torchdynamo_compiling() and hasattr(result, "past_key_values") - and hasattr(result.past_key_values, "to_legacy_cache") - and result.past_key_values.to_legacy_cache is not None + and getattr(result.past_key_values, "to_legacy_cache") is not None ): - # handle BC (convert by default if he user hasn't passed a cache AND the cache is of the default type) - should_convert_cache = generation_config.return_legacy_cache - is_user_defined_cache = user_defined_cache is not None - is_default_cache_type = ( - type(result.past_key_values) == DynamicCache # noqa E721 - or ( - isinstance(result.past_key_values, EncoderDecoderCache) - and type(result.past_key_values.self_attention_cache) == DynamicCache # noqa E721 - and type(result.past_key_values.cross_attention_cache) == DynamicCache # noqa E721 - ) - ) - if not is_user_defined_cache and is_default_cache_type: - logger.warning_once( - "From v4.47 onwards, when a model cache is to be returned, `generate` will return a `Cache` " - "instance instead by default (as opposed to the legacy tuple of tuples format). If you want to " - "keep returning the legacy format, please set `return_legacy_cache=True`." - ) - should_convert_cache = True - if should_convert_cache: - result.past_key_values = result.past_key_values.to_legacy_cache() + result.past_key_values = result.past_key_values.to_legacy_cache() return result @@ -2108,8 +2099,12 @@ def _contrastive_search( if not sequential: # Expands model inputs top_k times, for batched forward passes (akin to beam search). + # input_ids is required for expanding visual inputs in qwen2vl _, model_kwargs = self._expand_inputs_for_generation( - expand_size=top_k, is_encoder_decoder=self.config.is_encoder_decoder, **model_kwargs + input_ids=input_ids, + expand_size=top_k, + is_encoder_decoder=self.config.is_encoder_decoder, + **model_kwargs, ) past_key_values = model_kwargs.get("past_key_values") @@ -2316,7 +2311,9 @@ def _contrastive_search( next_past_key_values = selected_outputs["past_key_values"] else: - _, next_past_key_values = self._extract_past_from_model_output(outputs) + next_past_key_values = None + for possible_cache_name in ALL_CACHE_NAMES: + next_past_key_values = next_past_key_values or getattr(outputs, possible_cache_name, None) # Do it in-place layer per layer to save memory if isinstance(next_past_key_values, DynamicCache) or ( isinstance(next_past_key_values, EncoderDecoderCache) @@ -3976,8 +3973,8 @@ def _assisted_decoding( ) model_inputs = self.prepare_inputs_for_generation(candidate_input_ids, **candidate_kwargs) - if "num_logits_to_keep" in model_inputs: - model_inputs["num_logits_to_keep"] = candidate_length + 1 + if "logits_to_keep" in model_inputs: + model_inputs["logits_to_keep"] = candidate_length + 1 hpu_graphs_kwargs = self._get_hpu_graphs_kwargs(model_kwargs) diff --git a/optimum/habana/transformers/modeling_utils.py b/optimum/habana/transformers/modeling_utils.py index b6776eaa5d..60d9e56484 100644 --- a/optimum/habana/transformers/modeling_utils.py +++ b/optimum/habana/transformers/modeling_utils.py @@ -41,7 +41,6 @@ gaudi_awq_quantizer_process_model_before_weight_loading, gaudi_awq_quantizer_validate_environment, ) -from .modeling_utils_transformers import _gaudi_init_added_embeddings_weights_with_mean from .models import ( GAUDI_WHISPER_ATTENTION_CLASSES, BaichuanConfig, @@ -804,8 +803,3 @@ def adapt_transformers_to_gaudi(): transformers.loss.loss_for_object_detection.ImageLoss.loss_cardinality = gaudi_DetrLoss_loss_cardinality transformers.loss.loss_for_object_detection.ImageLoss.loss_boxes = gaudi_DetrLoss_loss_boxes transformers.loss.loss_for_object_detection.ImageLoss.forward = gaudi_DetrLoss_forward - - # Workaround for textual inversion - transformers.modeling_utils.PreTrainedModel._init_added_embeddings_weights_with_mean = ( - _gaudi_init_added_embeddings_weights_with_mean - ) diff --git a/optimum/habana/transformers/modeling_utils_transformers.py b/optimum/habana/transformers/modeling_utils_transformers.py deleted file mode 100644 index d2f1a49d97..0000000000 --- a/optimum/habana/transformers/modeling_utils_transformers.py +++ /dev/null @@ -1,35 +0,0 @@ -import torch - - -def _gaudi_init_added_embeddings_weights_with_mean( - self, old_embeddings, new_embeddings, old_embedding_dim, old_num_tokens, added_num_tokens -): - """ - Copied from: https://github.com/huggingface/transformers/blob/v4.48.2/src/transformers/modeling_utils.py#L2406 - Changes: - - torch.linalg.eigvals is not supported on HPU so run it on CPU - """ - old_embeddings_weight = old_embeddings.weight.data.to(torch.float32) - mean_embeddings = torch.mean(old_embeddings_weight, axis=0) - old_centered_embeddings = old_embeddings_weight - mean_embeddings - covariance = old_centered_embeddings.T @ old_centered_embeddings / old_num_tokens - - # Check if the covariance is positive definite. - # TODO: do not move `covariance` to the host once torch.linalg.eigvals is supported on HPU - eigenvalues = torch.linalg.eigvals(covariance.to("cpu")) - is_covariance_psd = bool( - (covariance == covariance.T).all() and not torch.is_complex(eigenvalues) and (eigenvalues > 0).all() - ) - if is_covariance_psd: - # If covariances is positive definite, a distribution can be created. and we can sample new weights from it. - distribution = torch.distributions.multivariate_normal.MultivariateNormal( - mean_embeddings, covariance_matrix=1e-9 * covariance - ) - new_embeddings.weight.data[-1 * added_num_tokens :, :] = distribution.sample( - sample_shape=(added_num_tokens,) - ).to(old_embeddings.weight.dtype) - else: - # Otherwise, just initialize with the mean. because distribtion will not be created. - new_embeddings.weight.data[-1 * added_num_tokens :, :] = ( - mean_embeddings[None, :].repeat(added_num_tokens, 1).to(old_embeddings.weight.dtype) - ) diff --git a/optimum/habana/transformers/models/cohere/modeling_cohere.py b/optimum/habana/transformers/models/cohere/modeling_cohere.py index 495ae2f9f0..e5ce7c1081 100644 --- a/optimum/habana/transformers/models/cohere/modeling_cohere.py +++ b/optimum/habana/transformers/models/cohere/modeling_cohere.py @@ -287,7 +287,7 @@ def forward( output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, - num_logits_to_keep: int = 0, + logits_to_keep: Union[int, torch.Tensor] = 0, token_idx: Optional[torch.Tensor] = None, **kwargs: Unpack[KwargsForCausalLM], ) -> Union[Tuple, CausalLMOutputWithPast]: @@ -314,7 +314,8 @@ def forward( hidden_states = outputs[0] # Only compute necessary logits, and do not upcast them to float if we are not computing the loss - logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]) + 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, :]) logits = logits * self.logit_scale # main diff from Llama loss = None diff --git a/optimum/habana/transformers/models/falcon/modeling_falcon.py b/optimum/habana/transformers/models/falcon/modeling_falcon.py index 508fab27af..4c1d2b1a42 100644 --- a/optimum/habana/transformers/models/falcon/modeling_falcon.py +++ b/optimum/habana/transformers/models/falcon/modeling_falcon.py @@ -1031,7 +1031,7 @@ def forward( output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, - num_logits_to_keep: int = 0, + logits_to_keep: Union[int, torch.Tensor] = 0, token_idx: Optional[torch.Tensor] = None, reuse_cache: Optional[bool] = False, trim_logits: Optional[bool] = False, @@ -1090,7 +1090,8 @@ def forward( else: hidden_states = hidden_states[:, -1:, :] - lm_logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]) + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + lm_logits = self.lm_head(hidden_states[:, slice_indices, :]) loss = None if labels is not None: diff --git a/optimum/habana/transformers/models/gemma/modeling_gemma.py b/optimum/habana/transformers/models/gemma/modeling_gemma.py index d2d4209d0e..eb2ba9b89d 100755 --- a/optimum/habana/transformers/models/gemma/modeling_gemma.py +++ b/optimum/habana/transformers/models/gemma/modeling_gemma.py @@ -781,7 +781,7 @@ def forward( output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, - num_logits_to_keep: int = 0, + 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, @@ -823,7 +823,8 @@ def forward( hidden_states = outputs[0] # Only compute necessary logits, and do not upcast them to float if we are not computing the loss - logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]).float() + 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() loss = None if labels is not None: diff --git a/optimum/habana/transformers/models/gemma2/modeling_gemma2.py b/optimum/habana/transformers/models/gemma2/modeling_gemma2.py index 172a5f218d..cb4c6ab65f 100755 --- a/optimum/habana/transformers/models/gemma2/modeling_gemma2.py +++ b/optimum/habana/transformers/models/gemma2/modeling_gemma2.py @@ -899,7 +899,7 @@ def forward( output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, - num_logits_to_keep: int = 0, + logits_to_keep: Union[int, torch.Tensor] = 0, token_idx: Optional[torch.Tensor] = None, trim_logits: Optional[bool] = False, attn_softmax_bf16: Optional[bool] = False, @@ -956,7 +956,12 @@ def forward( else: hidden_states = hidden_states[:, -1, :] - logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]) + 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() + if self.config.final_logit_softcapping is not None: + logits = logits / self.config.final_logit_softcapping + logits = torch.tanh(logits) + logits = logits * self.config.final_logit_softcapping loss = None if labels is not None: diff --git a/optimum/habana/transformers/models/gpt2/modeling_gpt2.py b/optimum/habana/transformers/models/gpt2/modeling_gpt2.py index e42a8308fa..301f9b6633 100644 --- a/optimum/habana/transformers/models/gpt2/modeling_gpt2.py +++ b/optimum/habana/transformers/models/gpt2/modeling_gpt2.py @@ -297,7 +297,7 @@ def gaudi_gpt2_forward( if inputs_embeds is None: inputs_embeds = self.wte(input_ids) position_embeds = self.wpe(position_ids) - hidden_states = inputs_embeds + position_embeds + hidden_states = inputs_embeds + position_embeds.to(inputs_embeds.device) # GPT2Attention mask. attention_mask = attention_mask.view(batch_size, -1) if attention_mask is not None else None 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 608c272135..ffc27dc931 100644 --- a/optimum/habana/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py +++ b/optimum/habana/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py @@ -619,7 +619,7 @@ def gaudi_gpt_bigcode_model_forward( if inputs_embeds is None: inputs_embeds = self.wte(input_ids) position_embeds = self.wpe(position_ids) - hidden_states = inputs_embeds + position_embeds + hidden_states = inputs_embeds + position_embeds.to(inputs_embeds.device) if token_type_ids is not None: token_type_embeds = self.wte(token_type_ids) 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 dd41d7b557..30b8ee79ee 100644 --- a/optimum/habana/transformers/models/gpt_neox/modeling_gpt_neox.py +++ b/optimum/habana/transformers/models/gpt_neox/modeling_gpt_neox.py @@ -1,7 +1,6 @@ from typing import Optional, Tuple, Union import torch -from torch.nn import CrossEntropyLoss from transformers.cache_utils import Cache from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from transformers.models.gpt_neox.configuration_gpt_neox import GPTNeoXConfig @@ -11,9 +10,11 @@ GPTNeoXLayer, GPTNeoXMLP, GPTNeoXModel, + KwargsForCausalLM, 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 @@ -82,6 +83,7 @@ class GaudiGPTNeoXAttention(GPTNeoXAttention): def __init__(self, config: GPTNeoXConfig, layer_idx=None): super().__init__(config, layer_idx) self.rotary_emb = GaudiRotaryEmbedding(config=self.config) + self.num_attention_heads = config.num_attention_heads def forward( self, @@ -159,7 +161,7 @@ def forward( value, attention_mask=attention_mask, head_mask=head_mask, - norm_factor=self.norm_factor, + norm_factor=self.scaling, attention_dropout=self.config.attention_dropout, training=self.training, ) @@ -174,6 +176,18 @@ def forward( return outputs + @classmethod + def _merge_heads(cls, tensor, num_attention_heads, attn_head_size): + """ + Merges attn_head_size dim and num_attn_heads dim into hidden dim + """ + # tensor [bs, num_attention_heads, seq_len, attn_head_size] + tensor = tensor.permute(0, 2, 1, 3).contiguous() + # -> [bs, seq_len, num_attention_heads, attn_head_size] + tensor = tensor.view(tensor.size(0), tensor.size(1), num_attention_heads * attn_head_size) + # -> [bs, seq_len, hidden_size] + return tensor + class GaudiGPTNeoXLayer(GPTNeoXLayer): def __init__(self, config, layer_idx): @@ -375,7 +389,7 @@ def gaudi_gpt_neox_model_forward( class GaudiGPTNeoXForCausalLM(GPTNeoXForCausalLM): """ - Inherits from GPTNeoXForCausalLM: https://github.com/huggingface/transformers/blob/main/src/transformers/models/opt_neox/modeling_gpt_neox.py + Inherits from GPTNeoXForCausalLM: https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt_neox/modeling_gpt_neox.py The only differences are: - add new args token_idx - add token_idx into model_inputs @@ -408,7 +422,8 @@ def forward( return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, - **kwargs, # Unused for now, mostly for the loss correction + 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 @@ -425,28 +440,25 @@ def forward( return_dict=return_dict, cache_position=cache_position, token_idx=token_idx, + **kwargs, ) hidden_states = outputs[0] - lm_logits = self.embed_out(hidden_states) + # 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, :]) - lm_loss = None + loss = None if labels is not None: - # move labels to correct device to enable model parallelism - labels = labels.to(lm_logits.device) - # we are doing next-token prediction; shift prediction scores and input ids by one - shift_logits = lm_logits[:, :-1, :].contiguous() - labels = labels[:, 1:].contiguous() - loss_fct = CrossEntropyLoss() - lm_loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), labels.view(-1)) + loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs) if not return_dict: - output = (lm_logits,) + outputs[1:] - return ((lm_loss,) + output) if lm_loss is not None else output + output = (logits,) + outputs[1:] + return ((loss,) + output) if loss is not None else output return CausalLMOutputWithPast( - loss=lm_loss, - logits=lm_logits, + loss=loss, + logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, @@ -518,6 +530,15 @@ def prepare_inputs_for_generation( return model_inputs + def _reorder_cache(self, past_key_values, beam_idx): + reordered_past = () + for layer_past in past_key_values: + reordered_past += ( + tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past[:2]) + + layer_past[2:], + ) + return reordered_past + def apply_customized_rope(q, k, cos, sin, position_ids, training=True): if q.device.type == "hpu" and FusedRoPE is not None: diff --git a/optimum/habana/transformers/models/idefics2/modeling_idefics2.py b/optimum/habana/transformers/models/idefics2/modeling_idefics2.py index b9e616fe09..cf4dd06452 100644 --- a/optimum/habana/transformers/models/idefics2/modeling_idefics2.py +++ b/optimum/habana/transformers/models/idefics2/modeling_idefics2.py @@ -237,7 +237,7 @@ def inputs_merger( 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 + new_inputs_embeds[special_image_token_mask] = reshaped_image_hidden_states.to(new_inputs_embeds.device) return new_inputs_embeds @@ -257,7 +257,7 @@ def forward( output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, - num_logits_to_keep: int = 0, + logits_to_keep: Union[int, torch.Tensor] = 0, token_idx: Optional[torch.Tensor] = None, ) -> Union[Tuple, Idefics2CausalLMOutputWithPast]: """ @@ -336,7 +336,8 @@ def forward( hidden_states = outputs[0] # Only compute necessary logits, and do not upcast them to float if we are not computing the loss - logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]) + 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: diff --git a/optimum/habana/transformers/models/llama/modeling_llama.py b/optimum/habana/transformers/models/llama/modeling_llama.py index 3b40bb6ce9..3bb0589e6b 100755 --- a/optimum/habana/transformers/models/llama/modeling_llama.py +++ b/optimum/habana/transformers/models/llama/modeling_llama.py @@ -1146,7 +1146,7 @@ def __init__(self, config: LlamaConfig): layers = [] for layer_idx in range(config.num_hidden_layers): layer = GaudiLlamaDecoderLayer(config, layer_idx) - if config.parallel_strategy is not None: + if hasattr(config, "paralle_strategy") and config.parallel_strategy is not None: layer = config.parallel_strategy.distribute_layer(layer, layer_idx) layers.append(layer) self.layers = torch.nn.ModuleList(layers) @@ -1445,7 +1445,7 @@ def forward( output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, - num_logits_to_keep: int = 0, + logits_to_keep: Union[int, torch.Tensor] = 0, token_idx: Optional[torch.Tensor] = None, trim_logits: Optional[bool] = False, attn_softmax_bf16: Optional[bool] = False, @@ -1506,7 +1506,8 @@ def forward( hidden_states = hidden_states[:, -1, :] # Only compute necessary logits, and do not upcast them to float if we are not computing the loss - logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]) + 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() loss = None if labels is not None: diff --git a/optimum/habana/transformers/models/llava/modeling_llava.py b/optimum/habana/transformers/models/llava/modeling_llava.py index f8fba446e6..474bd41fc3 100644 --- a/optimum/habana/transformers/models/llava/modeling_llava.py +++ b/optimum/habana/transformers/models/llava/modeling_llava.py @@ -114,7 +114,7 @@ def forward( position_ids: Optional[torch.LongTensor] = 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, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, @@ -122,12 +122,14 @@ def forward( output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, - num_logits_to_keep: int = 0, + logits_to_keep: Union[int, torch.Tensor] = 0, + image_sizes: torch.Tensor = None, token_idx: Optional[torch.Tensor] = None, image_offset: Optional[int] = None, tokens_pos: Optional[torch.LongTensor] = None, use_flash_attention: Optional[bool] = False, flash_attention_recompute: Optional[bool] = False, + **lm_kwargs, ) -> Union[Tuple, LlavaCausalLMOutputWithPast]: """ Inherits from LlavaForConditionalGeneration: https://github.com/huggingface/transformers/blob/v4.45.2/src/transformers/models/llava/modeling_llava.py#L362 @@ -152,9 +154,7 @@ def forward( ) if (input_ids is None) ^ (inputs_embeds is not None): - raise ValueError( - "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one" - ) + 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( @@ -199,10 +199,11 @@ def forward( output_hidden_states=output_hidden_states, return_dict=return_dict, cache_position=cache_position, - num_logits_to_keep=num_logits_to_keep, + logits_to_keep=logits_to_keep, token_idx=token_idx + image_offset, use_flash_attention=use_flash_attention, flash_attention_recompute=flash_attention_recompute, + **lm_kwargs, ) if input_ids.shape[1] != 1 and pixel_values is not None and tokens_pos is not None: @@ -238,9 +239,10 @@ def forward( output_hidden_states=output_hidden_states, return_dict=return_dict, cache_position=cache_position, - num_logits_to_keep=num_logits_to_keep, + logits_to_keep=logits_to_keep, use_flash_attention=use_flash_attention, flash_attention_recompute=flash_attention_recompute, + **lm_kwargs, ) logits = outputs[0] @@ -249,7 +251,9 @@ def forward( if labels is not None: # Shift so that tokens < n predict n if attention_mask is not None: - shift_attention_mask = attention_mask[..., 1:] + # 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[:, -(logits.shape[1] - 1) :].to(logits.device) shift_logits = logits[..., :-1, :][shift_attention_mask.to(logits.device) != 0].contiguous() shift_labels = labels[..., 1:][shift_attention_mask.to(labels.device) != 0].contiguous() else: @@ -282,7 +286,7 @@ def prepare_inputs_for_generation( pixel_values=None, attention_mask=None, cache_position=None, - num_logits_to_keep=None, + logits_to_keep=None, **kwargs, ): """ @@ -358,8 +362,8 @@ def prepare_inputs_for_generation( use_flash_attention = kwargs.get("use_flash_attention", False) flash_attention_recompute = kwargs.get("flash_attention_recompute", False) - if num_logits_to_keep is not None: - model_inputs["num_logits_to_keep"] = num_logits_to_keep + if logits_to_keep is not None: + model_inputs["logits_to_keep"] = logits_to_keep model_inputs.update( { 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 72c2b0a01b..a61ef20599 100644 --- a/optimum/habana/transformers/models/llava_next/modeling_llava_next.py +++ b/optimum/habana/transformers/models/llava_next/modeling_llava_next.py @@ -46,7 +46,7 @@ def forward( position_ids: Optional[torch.LongTensor] = 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, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, @@ -54,10 +54,11 @@ def forward( output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, - num_logits_to_keep: int = 0, + 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]: """ Inherits from LlavaForConditionalGeneration: https://github.com/huggingface/transformers/blob/v4.40.0/src/transformers/models/llava_next/modeling_llava_next.py#L433 @@ -88,10 +89,11 @@ def forward( return_dict=return_dict, 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 - # num_logits_to_keep=num_logits_to_keep, + # 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, ) if inputs_embeds.shape[1] != 1 and pixel_values is not None and self.text_tokens_pos is not None: @@ -150,7 +152,8 @@ def forward( output_hidden_states=output_hidden_states, return_dict=return_dict, cache_position=cache_position, - num_logits_to_keep=num_logits_to_keep, + logits_to_keep=logits_to_keep, + **lm_kwargs, ) # Copied from https://github.com/huggingface/transformers/blob/v4.40.0/src/transformers/models/llava_next/modeling_llava_next.py#L356 @@ -240,7 +243,7 @@ def prepare_inputs_for_generation( image_sizes=None, attention_mask=None, cache_position=None, - num_logits_to_keep=None, + logits_to_keep=None, **kwargs, ): """ @@ -259,7 +262,7 @@ def prepare_inputs_for_generation( image_sizes=image_sizes, attention_mask=attention_mask, cache_position=cache_position, - num_logits_to_keep=num_logits_to_keep, + logits_to_keep=logits_to_keep, **kwargs, ) else: @@ -418,8 +421,8 @@ def prepare_inputs_for_generation( else: model_inputs = {"input_ids": input_ids} - if num_logits_to_keep is not None: - model_inputs["num_logits_to_keep"] = num_logits_to_keep + if logits_to_keep is not None: + model_inputs["logits_to_keep"] = logits_to_keep model_inputs.update( { diff --git a/optimum/habana/transformers/models/mistral/modeling_mistral.py b/optimum/habana/transformers/models/mistral/modeling_mistral.py index 2c5b28b307..38e3a4d3f4 100644 --- a/optimum/habana/transformers/models/mistral/modeling_mistral.py +++ b/optimum/habana/transformers/models/mistral/modeling_mistral.py @@ -630,7 +630,7 @@ def forward( output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, - num_logits_to_keep: int = 0, + logits_to_keep: Union[int, torch.Tensor] = 0, token_idx: Optional[torch.Tensor] = None, reuse_cache: Optional[bool] = False, trim_logits: Optional[bool] = False, @@ -687,7 +687,8 @@ def forward( else: hidden_states = hidden_states[:, -1, :] # Only compute necessary logits, and do not upcast them to float if we are not computing the loss - logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]) + 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() loss = None if labels is not None: diff --git a/optimum/habana/transformers/models/mixtral/modeling_mixtral.py b/optimum/habana/transformers/models/mixtral/modeling_mixtral.py index d84b44dbab..2c9e6ba2f1 100644 --- a/optimum/habana/transformers/models/mixtral/modeling_mixtral.py +++ b/optimum/habana/transformers/models/mixtral/modeling_mixtral.py @@ -751,7 +751,7 @@ def forward( output_router_logits: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, - num_logits_to_keep: int = 0, + logits_to_keep: Union[int, torch.Tensor] = 0, token_idx: Optional[torch.Tensor] = None, reuse_cache: Optional[bool] = None, flash_attention_recompute: Optional[bool] = False, @@ -789,7 +789,8 @@ def forward( hidden_states = outputs[0] # Only compute necessary logits, and do not upcast them to float if we are not computing the loss - logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]) + 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() loss = None if labels is not None: diff --git a/optimum/habana/transformers/models/mllama/modeling_mllama.py b/optimum/habana/transformers/models/mllama/modeling_mllama.py index b1d5286469..450cfc9523 100644 --- a/optimum/habana/transformers/models/mllama/modeling_mllama.py +++ b/optimum/habana/transformers/models/mllama/modeling_mllama.py @@ -122,7 +122,7 @@ def forward( attention_mask: Optional[torch.Tensor] = None, output_attentions: bool = None, use_flash_attention: Optional[bool] = False, - ) -> 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 @@ -865,7 +865,7 @@ def forward( output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, - num_logits_to_keep: int = 0, + 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, @@ -906,8 +906,9 @@ def forward( hidden_states = outputs[0] - if token_idx is None and num_logits_to_keep != 0: - logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]).float() + if token_idx is None and logits_to_keep != 0: + 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() else: logits = self.lm_head(hidden_states).float() @@ -952,7 +953,7 @@ def forward( output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, - num_logits_to_keep: int = 0, + 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, @@ -1037,7 +1038,7 @@ def forward( output_attentions=output_attentions, return_dict=return_dict, cache_position=cache_position, - num_logits_to_keep=num_logits_to_keep, + logits_to_keep=logits_to_keep, token_idx=token_idx, use_flash_attention=use_flash_attention, flash_attention_recompute=flash_attention_recompute, @@ -1058,7 +1059,7 @@ def prepare_inputs_for_generation( past_key_values=None, use_cache=False, cache_position=None, - num_logits_to_keep=None, + logits_to_keep=None, **kwargs, ): """ @@ -1105,8 +1106,8 @@ def prepare_inputs_for_generation( # The clone here is for the same reason as for `position_ids`. model_inputs = {"input_ids": input_ids.clone(memory_format=torch.contiguous_format), "inputs_embeds": None} - if num_logits_to_keep is not None: - model_inputs["num_logits_to_keep"] = num_logits_to_keep + if logits_to_keep is not None: + model_inputs["logits_to_keep"] = logits_to_keep # keep cache_position implementation as None for HPU cache_position = None @@ -1193,7 +1194,9 @@ def forward( aspect_ratio_ids = aspect_ratio_ids.reshape(batch_size * num_concurrent_media, -1) # Patch embedding - patch_embeds = self.patch_embedding(pixel_values.to(self.dtype).to(self.device)) + target_dtype = self.patch_embedding.weight.dtype + target_device = self.patch_embedding.weight.device + patch_embeds = self.patch_embedding(pixel_values.to(target_device, target_dtype)) hidden_state = patch_embeds.flatten(2).transpose(1, 2) # Tile embeddings diff --git a/optimum/habana/transformers/models/opt/modeling_opt.py b/optimum/habana/transformers/models/opt/modeling_opt.py index 0d7afa4de8..2b0fa0c99b 100644 --- a/optimum/habana/transformers/models/opt/modeling_opt.py +++ b/optimum/habana/transformers/models/opt/modeling_opt.py @@ -41,6 +41,10 @@ def forward( return torch.nn.Embedding.forward(self, token_idx + self.offset) +def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int) -> torch.Tensor: + return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous() + + def gaudi_opt_attention_forward( self, hidden_states: torch.Tensor, @@ -74,12 +78,12 @@ def gaudi_opt_attention_forward( value_states = past_key_value[1] elif is_cross_attention: # cross_attentions - key_states = self._shape(self.k_proj(key_value_states), -1, bsz) - value_states = self._shape(self.v_proj(key_value_states), -1, bsz) + key_states = _shape(self, self.k_proj(key_value_states), -1, bsz) + value_states = _shape(self, self.v_proj(key_value_states), -1, bsz) elif past_key_value is not None: # reuse k, v, self_attention - key_states = self._shape(self.k_proj(hidden_states), -1, bsz) - value_states = self._shape(self.v_proj(hidden_states), -1, bsz) + key_states = _shape(self, self.k_proj(hidden_states), -1, bsz) + value_states = _shape(self, self.v_proj(hidden_states), -1, bsz) if token_idx is not None: past_key_value[0].index_copy_(2, token_idx - 1, key_states) past_key_value[1].index_copy_(2, token_idx - 1, value_states) @@ -90,21 +94,13 @@ def gaudi_opt_attention_forward( value_states = torch.cat([past_key_value[1], value_states], dim=2) else: # self_attention - key_states = self._shape(self.k_proj(hidden_states), -1, bsz) - 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. - # 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 - # 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` - past_key_value = (key_states, value_states) + key_states = _shape(self, self.k_proj(hidden_states), -1, bsz) + value_states = _shape(self, self.v_proj(hidden_states), -1, bsz) + + past_key_value = (key_states, value_states) proj_shape = (bsz * self.num_heads, -1, self.head_dim) - query_states = self._shape(query_states, tgt_len, bsz).view(*proj_shape) + query_states = _shape(self, query_states, tgt_len, bsz).view(*proj_shape) key_states = key_states.view(*proj_shape) value_states = value_states.view(*proj_shape) @@ -171,14 +167,14 @@ def gaudi_opt_attention_forward( class GaudiOPTDecoderLayer(torch.nn.Module): - def __init__(self, config: OPTConfig): + def __init__(self, config: OPTConfig, layer_idx: int = None): """ Attention implementation is set to "eager" (default in Transformers is "sdpa"). """ super().__init__() self.embed_dim = config.hidden_size - self.self_attn = OPT_ATTENTION_CLASSES["eager"](config=config, is_decoder=True) + self.self_attn = OPT_ATTENTION_CLASSES["eager"](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/paligemma/modeling_paligemma.py b/optimum/habana/transformers/models/paligemma/modeling_paligemma.py index 6f2a2817d0..ade847111e 100644 --- a/optimum/habana/transformers/models/paligemma/modeling_paligemma.py +++ b/optimum/habana/transformers/models/paligemma/modeling_paligemma.py @@ -24,7 +24,7 @@ PaliGemmaCausalLMOutputWithPast, PaliGemmaForConditionalGeneration, ) -from transformers.utils import logging +from transformers.utils import is_torchdynamo_compiling, logging logger = logging.get_logger(__name__) @@ -46,7 +46,7 @@ def forward( output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, - num_logits_to_keep: int = 0, + logits_to_keep: Union[int, torch.Tensor] = 0, token_idx: Optional[torch.Tensor] = None, **lm_kwargs, ) -> Union[Tuple, PaliGemmaCausalLMOutputWithPast]: @@ -90,7 +90,7 @@ def forward( 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 inputs_embeds[special_image_mask].numel() != image_features.numel(): + 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) raise ValueError( f"Number of images does not match number of special image tokens in the input text. " @@ -122,7 +122,7 @@ def forward( return_dict=return_dict, 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 - # num_logits_to_keep=num_logits_to_keep, + # logits_to_keep=logits_to_keep, token_idx=token_idx, **lm_kwargs, ) diff --git a/optimum/habana/transformers/models/persimmon/modeling_persimmon.py b/optimum/habana/transformers/models/persimmon/modeling_persimmon.py index 62fbe16f3c..1c02f414e0 100644 --- a/optimum/habana/transformers/models/persimmon/modeling_persimmon.py +++ b/optimum/habana/transformers/models/persimmon/modeling_persimmon.py @@ -362,7 +362,7 @@ def forward( output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, - num_logits_to_keep: int = 0, + logits_to_keep: Union[int, torch.Tensor] = 0, token_idx: Optional[torch.Tensor] = None, **kwargs, ) -> Union[Tuple, CausalLMOutputWithPast]: @@ -395,7 +395,8 @@ def forward( hidden_states = outputs[0] # No upscaling to float was ever done for Persimmon - logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]) + 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: diff --git a/optimum/habana/transformers/models/phi/modeling_phi.py b/optimum/habana/transformers/models/phi/modeling_phi.py index e7bd7b3b52..b72258aef7 100644 --- a/optimum/habana/transformers/models/phi/modeling_phi.py +++ b/optimum/habana/transformers/models/phi/modeling_phi.py @@ -482,7 +482,7 @@ def forward( output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, - num_logits_to_keep: int = 0, + logits_to_keep: Union[int, torch.Tensor] = 0, token_idx: Optional[torch.Tensor] = None, reuse_cache: Optional[bool] = False, trim_logits: Optional[bool] = False, @@ -527,7 +527,8 @@ def forward( else: hidden_states = hidden_states[:, -1, :] # Only compute necessary logits, and do not upcast them to float if we are not computing the loss - logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]) + 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: diff --git a/optimum/habana/transformers/models/qwen2/modeling_qwen2.py b/optimum/habana/transformers/models/qwen2/modeling_qwen2.py index e8536662ae..6956d6e4a6 100644 --- a/optimum/habana/transformers/models/qwen2/modeling_qwen2.py +++ b/optimum/habana/transformers/models/qwen2/modeling_qwen2.py @@ -865,7 +865,7 @@ def forward( output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, - num_logits_to_keep: int = 0, + logits_to_keep: Union[int, torch.Tensor] = 0, token_idx: Optional[torch.Tensor] = None, trim_logits: Optional[bool] = False, attn_softmax_bf16: Optional[bool] = False, @@ -923,7 +923,8 @@ def forward( hidden_states = hidden_states[:, -1, :] # Only compute necessary logits, and do not upcast them to float if we are not computing the loss - logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]) + 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() loss = None if labels is not None: 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 861a30dff4..3b7077aca9 100755 --- a/optimum/habana/transformers/models/qwen2_moe/modeling_qwen2_moe.py +++ b/optimum/habana/transformers/models/qwen2_moe/modeling_qwen2_moe.py @@ -1048,7 +1048,7 @@ def forward( output_router_logits: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, - num_logits_to_keep: int = 0, + logits_to_keep: Union[int, torch.Tensor] = 0, token_idx: Optional[torch.Tensor] = None, trim_logits: Optional[bool] = False, reuse_cache: Optional[bool] = None, @@ -1110,7 +1110,9 @@ def forward( else: hidden_states = hidden_states[:, -1, :] - logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]) + # 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: 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 79d11e9cff..007bf91ac8 100644 --- a/optimum/habana/transformers/models/qwen2_vl/modeling_qwen2_vl.py +++ b/optimum/habana/transformers/models/qwen2_vl/modeling_qwen2_vl.py @@ -37,7 +37,7 @@ apply_rotary_pos_emb_vision, repeat_kv, ) -from transformers.utils import logging +from transformers.utils import is_torchdynamo_compiling, logging try: @@ -68,7 +68,8 @@ def forward( self, hidden_states: torch.Tensor, cu_seqlens: torch.Tensor, - rotary_pos_emb: torch.Tensor = None, + rotary_pos_emb: Optional[torch.Tensor] = None, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, use_flash_attention: Optional[bool] = False, ) -> torch.Tensor: """ @@ -79,8 +80,19 @@ def forward( """ seq_length = hidden_states.shape[0] q, k, v = self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0) - q = apply_rotary_pos_emb_vision(q.unsqueeze(0), rotary_pos_emb).squeeze(0) - k = apply_rotary_pos_emb_vision(k.unsqueeze(0), rotary_pos_emb).squeeze(0) + if position_embeddings is None: + logger.warning_once( + "The attention layers in this model are transitioning from computing the RoPE embeddings internally " + "through `rotary_pos_emb` (2D tensor of RoPE theta values), to using externally computed " + "`position_embeddings` (Tuple of tensors, containing cos and sin). In v4.54 `rotary_pos_emb` will be " + "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() + else: + cos, sin = position_embeddings + q, k = apply_rotary_pos_emb_vision(q, k, cos, sin) attention_mask = torch.zeros([1, seq_length, seq_length], device=q.device, dtype=torch.bool) for i in range(1, len(cu_seqlens)): @@ -110,9 +122,10 @@ def __init__(self, config, attn_implementation: str = "sdpa") -> None: def forward( self, - hidden_states, - cu_seqlens, - rotary_pos_emb, + hidden_states: torch.Tensor, + cu_seqlens: torch.Tensor, + rotary_pos_emb: Optional[torch.Tensor] = None, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, use_flash_attention: Optional[bool] = False, ) -> torch.Tensor: """ @@ -124,6 +137,7 @@ def forward( self.norm1(hidden_states), cu_seqlens=cu_seqlens, rotary_pos_emb=rotary_pos_emb, + position_embeddings=position_embeddings, use_flash_attention=use_flash_attention, ) hidden_states = hidden_states + self.mlp(self.norm2(hidden_states)) @@ -354,6 +368,8 @@ def forward( """ hidden_states = self.patch_embed(hidden_states) rotary_pos_emb = self.rot_pos_emb(grid_thw) + emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1) + position_embeddings = (emb.cos(), emb.sin()) cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum( dim=0, dtype=torch.int32 @@ -363,13 +379,13 @@ def forward( for blk in self.blocks: if self.gradient_checkpointing and self.training: hidden_states = self._gradient_checkpointing_func( - blk.__call__, hidden_states, cu_seqlens, rotary_pos_emb, use_flash_attention + blk.__call__, hidden_states, cu_seqlens, None, position_embeddings, use_flash_attention ) else: hidden_states = blk( hidden_states, cu_seqlens=cu_seqlens, - rotary_pos_emb=rotary_pos_emb, + position_embeddings=position_embeddings, use_flash_attention=use_flash_attention, ) @@ -501,6 +517,9 @@ def forward( # from: https://github.com/huggingface/transformers/blob/v4.45.2/src/transformers/models/qwen2_vl/modeling_qwen2_vl.py#L1420 class GaudiQwen2VLForConditionalGeneration(Qwen2VLForConditionalGeneration): + # todo: change when the following gets fixed https://github.com/huggingface/transformers/blame/66f29aaaf55c8fe0c3dbcd24beede2ca4effac56/src/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py#L390C5-L390C27 + _supports_static_cache = True + def forward( self, input_ids: torch.LongTensor = None, @@ -633,7 +652,11 @@ def forward( # if we get 4D attention mask we cannot calculate rope deltas anymore. TODO @raushan fixme if position_ids is None and (attention_mask is None or attention_mask.ndim == 2): # calculate RoPE index once per generation in the pre-fill stage only - if (cache_position is not None and cache_position[0] == 0) or self.rope_deltas is None: + if ( + (cache_position is not None and cache_position[0] == 0) + or self.rope_deltas is None + or (past_key_values is None or past_key_values.get_seq_length() == 0) + ): position_ids, rope_deltas = self.get_rope_index( input_ids, image_grid_thw, video_grid_thw, attention_mask ) @@ -646,6 +669,7 @@ def forward( position_ids = position_ids.view(1, -1).expand(batch_size, -1) if cache_position is not None: # otherwise `deltas` is an int `0` delta = delta.repeat_interleave(batch_size // delta.shape[0], dim=0) + delta = delta.to(position_ids.device) position_ids = position_ids.add(delta) position_ids = position_ids.unsqueeze(0).expand(3, -1, -1) @@ -730,8 +754,17 @@ def prepare_inputs_for_generation( # 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) + # Exception 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: - if inputs_embeds is not None: # Exception 1 + 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] @@ -741,7 +774,7 @@ def prepare_inputs_for_generation( pixel_values_videos = None # if `inputs_embeds` are passed, we only want to use them in the 1st generation step - if inputs_embeds is not None and cache_position[0] == 0: + if inputs_embeds is not None and len(cache_position) == inputs_embeds.shape[1]: model_inputs = {"inputs_embeds": inputs_embeds, "input_ids": None} else: model_inputs = {"input_ids": input_ids, "inputs_embeds": None} diff --git a/optimum/habana/transformers/models/stablelm/modeling_stablelm.py b/optimum/habana/transformers/models/stablelm/modeling_stablelm.py index 7457b8f886..b86fdad63a 100644 --- a/optimum/habana/transformers/models/stablelm/modeling_stablelm.py +++ b/optimum/habana/transformers/models/stablelm/modeling_stablelm.py @@ -381,7 +381,7 @@ def forward( output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, - num_logits_to_keep: int = 0, + logits_to_keep: Union[int, torch.Tensor] = 0, token_idx: Optional[torch.Tensor] = None, **kwargs, ) -> Union[Tuple, CausalLMOutputWithPast]: @@ -412,7 +412,8 @@ def forward( hidden_states = outputs[0] # No upscaling to float was ever done for StableLm - logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]) + 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: diff --git a/optimum/habana/transformers/models/starcoder2/modeling_starcoder2.py b/optimum/habana/transformers/models/starcoder2/modeling_starcoder2.py index ecc6dce685..24fb2a4e17 100644 --- a/optimum/habana/transformers/models/starcoder2/modeling_starcoder2.py +++ b/optimum/habana/transformers/models/starcoder2/modeling_starcoder2.py @@ -699,7 +699,7 @@ def forward( output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, - num_logits_to_keep: int = 0, + logits_to_keep: Union[int, torch.Tensor] = 0, token_idx: Optional[torch.Tensor] = None, trim_logits: Optional[bool] = False, attn_softmax_bf16: Optional[bool] = False, @@ -753,7 +753,8 @@ def forward( hidden_states = hidden_states[:, -1, :] # Only compute necessary logits, and do not upcast them to float if we are not computing the loss - logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]) + 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: 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 2ba890c8d5..209045e0b5 100644 --- a/optimum/habana/transformers/models/video_llava/modeling_video_llava.py +++ b/optimum/habana/transformers/models/video_llava/modeling_video_llava.py @@ -18,6 +18,7 @@ import torch from torch import nn +from transformers.modeling_outputs import BaseModelOutputWithPooling from transformers.models.video_llava.modeling_video_llava import ( VideoLlavaCausalLMOutputWithPast, VideoLlavaConfig, @@ -123,6 +124,42 @@ def _merge_input_ids_with_visual_features( return final_embedding, final_attention_mask, final_labels, position_ids, final_input_ids + def _get_vision_features( + self, + pixel_values_images: Optional[torch.FloatTensor] = None, + pixel_values_videos: Optional[torch.FloatTensor] = None, + vision_feature_layer: Optional[int] = None, + vision_feature_select_strategy: Optional[str] = None, + ) -> 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`") + + # videos do not need to select features and it's always "full" (as it is done in the orig implementation) + if pixel_values_videos is not None: + batch_size_vid, num_frames, channels, height, width = pixel_values_videos.shape + + pixel_values = pixel_values_videos.reshape(batch_size_vid * num_frames, channels, height, width) + video_outputs = self.video_tower(pixel_values, output_hidden_states=True) + video_outputs = video_outputs.hidden_states[vision_feature_layer].squeeze(1) + else: + video_outputs = None + num_frames = 0 + + if pixel_values_images is not None: + image_outputs = self.image_tower(pixel_values_images, output_hidden_states=True) + image_outputs = image_outputs.hidden_states[vision_feature_layer].squeeze(1) + + if vision_feature_select_strategy == "default": + image_outputs = image_outputs[:, 1:] + elif vision_feature_select_strategy == "full": + image_outputs = image_outputs + else: + raise ValueError(f"Unexpected select feature strategy: {self.config.vision_feature_select_strategy}") + else: + image_outputs = None + + return image_outputs, video_outputs, num_frames + def forward( self, input_ids: torch.LongTensor = None, @@ -132,7 +169,7 @@ def forward( position_ids: Optional[torch.LongTensor] = 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, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, @@ -140,7 +177,7 @@ def forward( output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, - num_logits_to_keep: int = 0, + logits_to_keep: Union[int, torch.Tensor] = 0, token_idx: Optional[torch.Tensor] = None, **kwargs, ) -> Union[Tuple, VideoLlavaCausalLMOutputWithPast]: @@ -161,6 +198,7 @@ def forward( 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.language_model( attention_mask=attention_mask, position_ids=position_ids, @@ -171,19 +209,9 @@ def forward( output_hidden_states=output_hidden_states, return_dict=return_dict, cache_position=cache_position, - num_logits_to_keep=0, + logits_to_keep=0, token_idx=token_idx, - trim_logits=kwargs.get("trim_logits"), - attn_softmax_bf16=kwargs.get("attn_softmax_bf16"), - reuse_cache=kwargs.get("reuse_cache"), - use_flash_attention=kwargs.get("use_flash_attention"), - flash_attention_recompute=kwargs.get("flash_attention_recompute"), - flash_attention_causal_mask=kwargs.get("flash_attention_causal_mask"), - flash_attention_fast_softmax=kwargs.get("flash_attention_fast_softmax"), - valid_sequence_lengths=kwargs.get("valid_sequence_lengths"), - cache_idx=kwargs.get("cache_idx"), - lazy_mode=kwargs.get("lazy_mode"), - num_virtual_tokens=kwargs.get("num_virtual_tokens"), + **kwargs, ) logits = outputs[0] @@ -194,7 +222,9 @@ def forward( if labels is not None: # Shift so that tokens < n predict n if attention_mask is not None: - shift_attention_mask = attention_mask[..., 1:] + # 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[:, -(logits.shape[1] - 1) :].to(logits.device) shift_logits = logits[..., :-1, :][shift_attention_mask.to(logits.device) != 0].contiguous() shift_labels = labels[..., 1:][shift_attention_mask.to(labels.device) != 0].contiguous() else: @@ -229,7 +259,7 @@ def prepare_inputs_for_generation( pixel_values_videos=None, attention_mask=None, cache_position=None, - num_logits_to_keep=None, + logits_to_keep=None, **kwargs, ): token_idx = kwargs.get("token_idx", None) @@ -242,7 +272,7 @@ def prepare_inputs_for_generation( pixel_values_videos=pixel_values_videos, attention_mask=attention_mask, cache_position=cache_position, - num_logits_to_keep=num_logits_to_keep, + logits_to_keep=logits_to_keep, **kwargs, ) # Else, we need to update token_idx when merging features from videos/images with input embeddings @@ -277,7 +307,7 @@ def prepare_inputs_for_generation( inputs_embeds=inputs_embeds, attention_mask=attention_mask, cache_position=cache_position, - num_logits_to_keep=num_logits_to_keep, + logits_to_keep=logits_to_keep, **kwargs, ) position_ids = model_inputs["position_ids"] @@ -401,7 +431,7 @@ def prepare_inputs_for_generation( "inputs_embeds": inputs_embeds, } ) - if legacy_processing or cache_position[0] == 0: + if legacy_processing or (cache_position is not None and cache_position[0]) == 0: # If we're in cached decoding stage, pixel values should be None because input ids do not contain special image token anymore # Otherwise we need pixel values to be passed to model model_inputs["pixel_values_images"] = pixel_values_images diff --git a/optimum/habana/transformers/models/whisper/modeling_whisper.py b/optimum/habana/transformers/models/whisper/modeling_whisper.py index e5bf8b458a..fc86606eb6 100644 --- a/optimum/habana/transformers/models/whisper/modeling_whisper.py +++ b/optimum/habana/transformers/models/whisper/modeling_whisper.py @@ -300,7 +300,7 @@ def forward( if token_idx is not None: position_ids = (token_idx - 1).unsqueeze(0) else: - position_ids = cache_position.unsqueeze(0) + position_ids = cache_position.unsqueeze(0).repeat(input_shape[0], 1) # embed positions if input_ids is not None: positions = self.embed_positions( diff --git a/optimum/habana/transformers/models/xglm/modeling_xglm.py b/optimum/habana/transformers/models/xglm/modeling_xglm.py index 289e0eb55f..daf85e5e73 100644 --- a/optimum/habana/transformers/models/xglm/modeling_xglm.py +++ b/optimum/habana/transformers/models/xglm/modeling_xglm.py @@ -292,7 +292,7 @@ def gaudi_xglm_model_forward( encoder_attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1] ) - hidden_states = inputs_embeds + self.embed_positions(position_ids, past_key_values_length) + hidden_states = inputs_embeds + self.embed_positions(position_ids, past_key_values_length).to(inputs_embeds.device) hidden_states = nn.functional.dropout(hidden_states, p=float(self.dropout), training=self.training) if self.gradient_checkpointing and self.training: diff --git a/optimum/habana/transformers/trainer.py b/optimum/habana/transformers/trainer.py index 62761944a9..1931081bee 100644 --- a/optimum/habana/transformers/trainer.py +++ b/optimum/habana/transformers/trainer.py @@ -25,7 +25,6 @@ import os import random import shutil -import sys import time import warnings from collections.abc import Mapping @@ -51,7 +50,6 @@ from transformers.debug_utils import DebugOption, DebugUnderflowOverflow from transformers.feature_extraction_utils import FeatureExtractionMixin from transformers.image_processing_utils import BaseImageProcessor -from transformers.integrations import hp_params from transformers.integrations.deepspeed import ( deepspeed_load_checkpoint, is_deepspeed_available, @@ -79,7 +77,6 @@ PREFIX_CHECKPOINT_DIR, EvalLoopOutput, EvalPrediction, - HPSearchBackend, HubStrategy, PredictionOutput, SaveStrategy, @@ -650,51 +647,30 @@ def _inner_training_loop( # 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 + ( + num_train_epochs, + num_update_steps_per_epoch, + num_examples, + num_train_samples, + epoch_based, + len_dataloader, + max_steps, + ) = self.set_initial_training_values(args, train_dataloader, total_train_batch_size) if ( self.accelerator.mpu.sequence_parallel_is_initialized() and self.accelerator.mpu.get_sequence_parallel_world_size() > 1 ): total_train_batch_size = total_train_batch_size / self.accelerator.mpu.get_sequence_parallel_world_size() - len_dataloader = None num_train_tokens = None - if has_length(train_dataloader): - len_dataloader = len(train_dataloader) - num_update_steps_per_epoch = len_dataloader // args.gradient_accumulation_steps - num_update_steps_per_epoch = max(num_update_steps_per_epoch, 1) - num_examples = self.num_examples(train_dataloader) - if args.max_steps > 0: - max_steps = args.max_steps - num_train_epochs = args.max_steps // num_update_steps_per_epoch + int( - args.max_steps % num_update_steps_per_epoch > 0 - ) - # May be slightly incorrect if the last batch in the training dataloader has a smaller size but it's - # the best we can do. - num_train_samples = args.max_steps * total_train_batch_size - if args.include_tokens_per_second: - num_train_tokens = ( - self.num_tokens(train_dataloader, args.max_steps) * args.gradient_accumulation_steps - ) + if self.args.include_tokens_per_second: + num_train_tokens = self.num_tokens(train_dataloader, None if epoch_based else max_steps) + # If going by epochs, multiply tokens linearly + if len_dataloader is not None and epoch_based: + num_train_tokens *= args.num_train_epochs + # Otherwise since its steps, we just multiply by grad accum else: - max_steps = math.ceil(args.num_train_epochs * num_update_steps_per_epoch) - num_train_epochs = math.ceil(args.num_train_epochs) - num_train_samples = self.num_examples(train_dataloader) * args.num_train_epochs - if args.include_tokens_per_second: - num_train_tokens = self.num_tokens(train_dataloader) * args.num_train_epochs - elif args.max_steps > 0: # Rely on max_steps when dataloader does not have a working size - max_steps = args.max_steps - # Setting a very large number of epochs so we go as many times as necessary over the iterator. - num_train_epochs = sys.maxsize - num_update_steps_per_epoch = max_steps - num_examples = total_train_batch_size * args.max_steps - num_train_samples = args.max_steps * total_train_batch_size - if args.include_tokens_per_second: - num_train_tokens = self.num_tokens(train_dataloader, args.max_steps) * args.gradient_accumulation_steps - else: - raise ValueError( - "args.max_steps must be set to a positive value if dataloader does not have a length, was" - f" {args.max_steps}" - ) + num_train_tokens *= args.gradient_accumulation_steps if DebugOption.UNDERFLOW_OVERFLOW in self.args.debug: debug_overflow = DebugUnderflowOverflow(self.model) # noqa @@ -721,21 +697,7 @@ def _inner_training_loop( self.state.train_batch_size = self._train_batch_size # Compute absolute values for logging, eval, and save if given as ratio - if args.logging_steps is not None: - if args.logging_steps < 1: - self.state.logging_steps = math.ceil(max_steps * args.logging_steps) - else: - self.state.logging_steps = args.logging_steps - if args.eval_steps is not None: - if args.eval_steps < 1: - self.state.eval_steps = math.ceil(max_steps * args.eval_steps) - else: - self.state.eval_steps = args.eval_steps - if args.save_steps is not None: - if args.save_steps < 1: - self.state.save_steps = math.ceil(max_steps * args.save_steps) - else: - self.state.save_steps = args.save_steps + self.state.compute_steps(args, max_steps) # Activate gradient checkpointing if needed if args.gradient_checkpointing: @@ -838,6 +800,7 @@ def hpu_deepspeed_checkpointing(function, *checkpoint_args, use_reentrant: Optio # Check if saved optimizer or scheduler states exist self._load_optimizer_and_scheduler(resume_from_checkpoint) + self._load_scaler(resume_from_checkpoint) if self.gaudi_config.use_fused_clip_norm and self.args.use_habana: try: @@ -908,25 +871,7 @@ def hpu_deepspeed_checkpointing(function, *checkpoint_args, use_reentrant: Optio torch.distributed.broadcast(param.data, src=0) # Update the references - self.callback_handler.model = self.model - self.callback_handler.optimizer = self.optimizer - self.callback_handler.lr_scheduler = self.lr_scheduler - self.callback_handler.train_dataloader = train_dataloader - if self.hp_name is not None and self._trial is not None: - # use self._trial because the SigOpt/Optuna hpo only call `_hp_search_setup(trial)` instead of passing trial - # parameter to Train when using DDP. - self.state.trial_name = self.hp_name(self._trial) - if trial is not None: - assignments = trial.assignments if self.hp_search_backend == HPSearchBackend.SIGOPT else trial - self.state.trial_params = hp_params(assignments) - else: - self.state.trial_params = None - # This should be the same if the state has been saved but in case the training arguments changed, it's safer - # to set this after the load. - self.state.max_steps = max_steps - self.state.num_train_epochs = num_train_epochs - self.state.is_local_process_zero = self.is_local_process_zero() - self.state.is_world_process_zero = self.is_world_process_zero() + self.state.init_training_references(self, train_dataloader, 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) @@ -1130,8 +1075,7 @@ def hpu_deepspeed_checkpointing(function, *checkpoint_args, use_reentrant: Optio self.control = self.callback_handler.on_optimizer_step(args, self.state, self.control) - optimizer_was_run = not self.accelerator.optimizer_step_was_skipped - if optimizer_was_run: + 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): self.lr_scheduler.step() @@ -1679,6 +1623,11 @@ def training_step( # temporary fix to calculate loss correctly loss = loss / self.args.gradient_accumulation_steps + # Turning off loss scaling w.r.t. gradient accumulation when DeepSpeed is enabled + # https://github.com/huggingface/transformers/pull/35808 + if self.accelerator.distributed_type == GaudiDistributedType.DEEPSPEED: + kwargs["scale_wrt_gas"] = False + if _is_peft_model(self.model) and self.model.peft_type == PeftType.ADALORA: assert not (self.accelerator.state.is_fp8_enabled and self.args.gradient_checkpointing), ( "FP8 precision with gradient_checkpointing is currently not supported with PeftType.ADALORA" @@ -2197,7 +2146,7 @@ def prediction_step( inputs = self._prepare_inputs(inputs) if ignore_keys is None: if hasattr(self.model, "config"): - ignore_keys = getattr(self.model.config, "keys_to_ignore_at_inference", []) + ignore_keys = getattr(self.model.config, "keys_to_ignore_at_inference", ["past_key_values"]) else: ignore_keys = [] @@ -2541,11 +2490,10 @@ def create_accelerator_and_postprocess(self): accelerator_config = self.args.accelerator_config.to_dict() + # Extract dataloader config params from accelerator config + dataloader_params = ["split_batches", "dispatch_batches", "even_batches", "use_seedable_sampler"] dataloader_config = DataLoaderConfiguration( - split_batches=accelerator_config.pop("split_batches"), - dispatch_batches=accelerator_config.pop("dispatch_batches"), - even_batches=accelerator_config.pop("even_batches"), - use_seedable_sampler=accelerator_config.pop("use_seedable_sampler"), + **{param: accelerator_config.pop(param) for param in dataloader_params} ) if is_accelerate_available("1.1.0"): dataloader_config.data_seed = self.args.data_seed @@ -2584,12 +2532,8 @@ def create_accelerator_and_postprocess(self): # post accelerator creation setup if self.is_fsdp_enabled: fsdp_plugin = self.accelerator.state.fsdp_plugin - fsdp_plugin.limit_all_gathers = self.args.fsdp_config.get( - "limit_all_gathers", fsdp_plugin.limit_all_gathers - ) - fsdp_plugin.activation_checkpointing = self.args.fsdp_config.get( - "activation_checkpointing", fsdp_plugin.activation_checkpointing - ) + for param in ["limit_all_gathers", "activation_checkpointing"]: + setattr(fsdp_plugin, param, self.args.fsdp_config.get(param, getattr(fsdp_plugin, param))) if fsdp_plugin.activation_checkpointing and self.args.gradient_checkpointing: raise ValueError( "The activation_checkpointing in FSDP config and the gradient_checkpointing in training arg " diff --git a/optimum/habana/transformers/training_args.py b/optimum/habana/transformers/training_args.py index b4d87b275c..f9c8ba5467 100644 --- a/optimum/habana/transformers/training_args.py +++ b/optimum/habana/transformers/training_args.py @@ -409,6 +409,14 @@ def __post_init__(self): if self.throughput_warmup_steps < 0: raise ValueError("--throughput_warmup_steps must be positive.") + # Set default output_dir if not provided + if self.output_dir is None: + self.output_dir = "trainer_output" + logger.info( + "No output directory specified, defaulting to 'trainer_output'. " + "To change this behavior, specify --output_dir when creating TrainingArguments." + ) + # Parse in args that could be `dict` sent in from the CLI as a string for field in _VALID_DICT_FIELDS: passed_value = getattr(self, field) diff --git a/setup.py b/setup.py index c472e03326..b8a0774b07 100644 --- a/setup.py +++ b/setup.py @@ -29,7 +29,7 @@ INSTALL_REQUIRES = [ - "transformers >= 4.48.2, < 4.49.0", + "transformers >= 4.49.0, < 4.50.0", "optimum", "torch", "accelerate >= 0.33.0, < 0.34.0", diff --git a/tests/baselines/fixture/tests/test_diffusers.json b/tests/baselines/fixture/tests/test_diffusers.json index cde044dfd2..cdbb43a232 100644 --- a/tests/baselines/fixture/tests/test_diffusers.json +++ b/tests/baselines/fixture/tests/test_diffusers.json @@ -7,7 +7,7 @@ "throughput": 0.145 }, "gaudi3": { - "throughput": 0.145 + "throughput": 0.221 } }, "tests/test_diffusers.py::GaudiFluxImg2ImgPipelineTester::test_flux_img2img_inference": { @@ -64,7 +64,7 @@ "throughput": 1.086 }, "gaudi3": { - "throughput": 1.086 + "throughput": 2.168 } }, "tests/test_diffusers.py::GaudiStableDiffusionPipelineTester::test_sd_textual_inversion": { diff --git a/tests/baselines/fixture/tests/test_encoder_decoder.json b/tests/baselines/fixture/tests/test_encoder_decoder.json index 5f73275d56..670e29464c 100644 --- a/tests/baselines/fixture/tests/test_encoder_decoder.json +++ b/tests/baselines/fixture/tests/test_encoder_decoder.json @@ -9,8 +9,8 @@ "predict_samples_per_second": 4.339 }, "gaudi3": { - "predict_rougeLsum": 28.9801, - "predict_samples_per_second": 4.339 + "predict_rougeLsum": 15.618, + "predict_samples_per_second": 1.091 } }, "tests/test_encoder_decoder.py::TestEncoderDecoderModels::test_text_summarization_bf16[t5-3b-Habana/t5-2-1]": { @@ -23,8 +23,8 @@ "predict_samples_per_second": 3.848 }, "gaudi3": { - "predict_rougeLsum": 21.8877, - "predict_samples_per_second": 3.848 + "predict_rougeLsum": 21.7057, + "predict_samples_per_second": 5.032 } }, "tests/test_encoder_decoder.py::TestEncoderDecoderModels::test_text_translation_bf16[t5-small-Habana/t5-2-1]": { @@ -37,8 +37,8 @@ "predict_samples_per_second": 11.648 }, "gaudi3": { - "predict_bleu": 11.7277, - "predict_samples_per_second": 11.648 + "predict_bleu": 11.7168, + "predict_samples_per_second": 18.174 } } } \ No newline at end of file diff --git a/tests/baselines/fixture/tests/test_examples.json b/tests/baselines/fixture/tests/test_examples.json index 37845920e4..e281343d76 100644 --- a/tests/baselines/fixture/tests/test_examples.json +++ b/tests/baselines/fixture/tests/test_examples.json @@ -6,9 +6,9 @@ "train_samples_per_second": 14.06 }, "gaudi3": { - "perplexity": 26.39, - "train_runtime": 356.07, - "train_samples_per_second": 14.06 + "perplexity": 26.271165167474585, + "train_runtime": 218.4737, + "train_samples_per_second": 23.781 } }, "tests/test_examples.py::CausalLanguageModelingLORAExampleTester::test_run_lora_clm_llama-7b_single_card": { @@ -23,9 +23,9 @@ "train_samples_per_second": 18.428 }, "gaudi3": { - "perplexity": 3.8436, - "train_runtime": 113.9713, - "train_samples_per_second": 18.428 + "perplexity": 3.843924462719278, + "train_runtime": 148.7151, + "train_samples_per_second": 32.357 } }, "tests/test_examples.py::DeepSpeedTextClassificationExampleTester::test_run_glue_LlamaGuard-7b_deepspeed": { @@ -35,9 +35,9 @@ "train_samples_per_second": 342.169 }, "gaudi3": { - "eval_f1": 0.8873483535528596, - "train_runtime": 62.4539, - "train_samples_per_second": 342.169 + "eval_f1": 0.8809523809523809, + "train_runtime": 232.6707, + "train_samples_per_second": 560.75 } }, "tests/test_examples.py::DeepspeedCausalLanguageModelingExampleTester::test_run_clm_CodeLlama-13b-Instruct-hf_deepspeed": { @@ -47,9 +47,9 @@ "train_samples_per_second": 18.789 }, "gaudi3": { - "perplexity": 6.877496628184696, - "train_runtime": 542.2985, - "train_samples_per_second": 18.789 + "perplexity": 6.877100646486551, + "train_runtime": 477.7145, + "train_samples_per_second": 29.814 } }, "tests/test_examples.py::DeepspeedCausalLanguageModelingExampleTester::test_run_clm_bloom-7b1_deepspeed": { @@ -65,9 +65,9 @@ "train_samples_per_second": 18.216 }, "gaudi3": { - "perplexity": 16.51629, - "train_runtime": 445, - "train_samples_per_second": 18.216 + "perplexity": 16.260238201071928, + "train_runtime": 243.1757, + "train_samples_per_second": 34.196 } }, "tests/test_examples.py::DeepspeedCausalLanguageModelingExampleTester::test_run_clm_gemma-2b-it_deepspeed": { @@ -77,9 +77,9 @@ "train_samples_per_second": 81.097 }, "gaudi3": { - "perplexity": 924.062, - "train_runtime": 75.518, - "train_samples_per_second": 81.097 + "perplexity": 980.9833890324784, + "train_runtime": 51.73, + "train_samples_per_second": 142.775 } }, "tests/test_examples.py::DeepspeedCausalLanguageModelingExampleTester::test_run_clm_gpt-neox-20b_deepspeed": { @@ -89,9 +89,9 @@ "train_samples_per_second": 7.328 }, "gaudi3": { - "perplexity": 8.169664686471043, - "train_runtime": 445, - "train_samples_per_second": 7.328 + "perplexity": 7.827201417363628, + "train_runtime": 445.3031, + "train_samples_per_second": 11.704 } }, "tests/test_examples.py::DeepspeedCausalLanguageModelingExampleTester::test_run_clm_gpt2-xl_deepspeed": { @@ -106,9 +106,9 @@ "train_samples_per_second": 95.539 }, "gaudi3": { - "perplexity": 13.237754028004865, - "train_runtime": 206.5775, - "train_samples_per_second": 95.539 + "perplexity": 13.155277331993139, + "train_runtime": 159.357, + "train_samples_per_second": 150.538 } }, "tests/test_examples.py::DeepspeedSFTExampleTester::test_sft_Qwen2-72B_deepspeed": { @@ -118,9 +118,9 @@ "train_samples_per_second": 7.554 }, "gaudi3": { - "perplexity": 3.7020898897918824, - "train_runtime": 918.8018, - "train_samples_per_second": 7.554 + "perplexity": 3.728595328528421, + "train_runtime": 440.2459, + "train_samples_per_second": 19.627 } }, "tests/test_examples.py::DeepspeedSummarizationExampleTester::test_run_summarization_flan-t5-xxl_deepspeed": { @@ -130,9 +130,9 @@ "train_samples_per_second": 28.387 }, "gaudi3": { - "eval_rougeLsum": 29.308, - "train_runtime": 155.86, - "train_samples_per_second": 28.387 + "eval_rougeLsum": 28.0738, + "train_runtime": 118.419, + "train_samples_per_second": 52.048 } }, "tests/test_examples.py::EagerModeCausalLanguageModelingExampleTester::test_run_clm_gemma-2b-it_single_card": { @@ -142,9 +142,9 @@ "train_samples_per_second": 8.597 }, "gaudi3": { - "perplexity": 26.69, - "train_runtime": 560.8188, - "train_samples_per_second": 8.597 + "perplexity": 26.299428898047232, + "train_runtime": 318.8908, + "train_samples_per_second": 15.166 } }, "tests/test_examples.py::ImageClassificationExampleTester::test_run_image_classification_swin-base-patch4-window7-224-in22k_single_card": { @@ -159,9 +159,9 @@ "train_samples_per_second": 826.766 }, "gaudi3": { - "eval_accuracy": 0.9850666666666666, - "train_runtime": 77.8934, - "train_samples_per_second": 826.766 + "eval_accuracy": 0.9849333333333333, + "train_runtime": 73.8308, + "train_samples_per_second": 1155.964 } }, "tests/test_examples.py::ImageClassificationExampleTester::test_run_image_classification_vit-base-patch16-224-in21k_single_card": { @@ -177,8 +177,8 @@ }, "gaudi3": { "eval_accuracy": 0.9690666666666666, - "train_runtime": 54.9734, - "train_samples_per_second": 870.272 + "train_runtime": 47.9419, + "train_samples_per_second": 1164.009 } }, "tests/test_examples.py::MultiCardAudioClassificationExampleTester::test_run_audio_classification_ast-finetuned-speech-commands-v2_multi_card": { @@ -189,10 +189,10 @@ "train_samples_per_second": 1955.74 }, "gaudi3": { - "eval_accuracy": 0.1871, - "eval_samples_per_second": 2301.088, - "train_runtime": 139.9477, - "train_samples_per_second": 1955.74 + "eval_accuracy": 0.19650135869565216, + "eval_samples_per_second": 3352.901, + "train_runtime": 106.5372, + "train_samples_per_second": 2676.242 } }, "tests/test_examples.py::MultiCardAudioClassificationExampleTester::test_run_audio_classification_wav2vec2-base_multi_card": { @@ -209,10 +209,10 @@ "train_samples_per_second": 2975.844 }, "gaudi3": { - "eval_accuracy": 0.7228, - "eval_samples_per_second": 3640.021, - "train_runtime": 63.4079, - "train_samples_per_second": 2975.844 + "eval_accuracy": 0.7352241847826086, + "eval_samples_per_second": 2059.992, + "train_runtime": 57.0028, + "train_samples_per_second": 4213.033 } }, "tests/test_examples.py::MultiCardBridgetowerExampleTester::test_run_bridgetower_bridgetower-large-itm-mlm-itc_multi_card": { @@ -221,8 +221,8 @@ "train_samples_per_second": 904.93 }, "gaudi3": { - "train_runtime": 224.42, - "train_samples_per_second": 904.93 + "train_runtime": 342.4851, + "train_samples_per_second": 1009.467 } }, "tests/test_examples.py::MultiCardCausalLanguageModelingAdaloraExampleTester::test_run_lora_clm_llama-7b_multi_card": { @@ -232,9 +232,9 @@ "train_samples_per_second": 107 }, "gaudi3": { - "perplexity": 2.59, - "train_runtime": 459, - "train_samples_per_second": 107 + "perplexity": 2.592915682175543, + "train_runtime": 818.9693, + "train_samples_per_second": 85.059 } }, "tests/test_examples.py::MultiCardCausalLanguageModelingExampleTester::test_run_clm_gemma-2b-it_multi_card": { @@ -244,9 +244,9 @@ "train_samples_per_second": 94.524 }, "gaudi3": { - "perplexity": 954.5995, - "train_runtime": 82.6617, - "train_samples_per_second": 94.524 + "perplexity": 902.0585179806482, + "train_runtime": 66.2529, + "train_samples_per_second": 159.47 } }, "tests/test_examples.py::MultiCardCausalLanguageModelingIA3ExampleTester::test_run_lora_clm_llama-7b_multi_card": { @@ -256,9 +256,9 @@ "train_samples_per_second": 161 }, "gaudi3": { - "perplexity": 3.3, - "train_runtime": 262.8, - "train_samples_per_second": 161 + "perplexity": 3.291398111098924, + "train_runtime": 390.7556, + "train_samples_per_second": 256.027 } }, "tests/test_examples.py::MultiCardCausalLanguageModelingLORAExampleTester2::test_run_lora_clm_falcon-40b_multi_card": { @@ -268,9 +268,9 @@ "train_samples_per_second": 15.0 }, "gaudi3": { - "perplexity": 1.6, - "train_runtime": 710, - "train_samples_per_second": 15.0 + "perplexity": 1.588740773299791, + "train_runtime": 408.8298, + "train_samples_per_second": 33.87 } }, "tests/test_examples.py::MultiCardCausalLanguageModelingLORAExampleTester2::test_run_lora_clm_llama-7b_multi_card": { @@ -280,9 +280,9 @@ "train_samples_per_second": 148.093 }, "gaudi3": { - "perplexity": 2.3665, - "train_runtime": 294.5707, - "train_samples_per_second": 148.093 + "perplexity": 1.570946503005108, + "train_runtime": 342.6741, + "train_samples_per_second": 267.801 } }, "tests/test_examples.py::MultiCardCausalLanguageModelingLORAExampleTester::test_run_lora_clm_falcon-40b_multi_card": { @@ -292,9 +292,9 @@ "train_samples_per_second": 15.0 }, "gaudi3": { - "perplexity": 4.0, - "train_runtime": 550, - "train_samples_per_second": 15.0 + "perplexity": 3.694849124063941, + "train_runtime": 320.063, + "train_samples_per_second": 35.863 } }, "tests/test_examples.py::MultiCardCausalLanguageModelingLORAExampleTester::test_run_lora_clm_llama-7b_multi_card": { @@ -309,9 +309,9 @@ "train_samples_per_second": 148.093 }, "gaudi3": { - "perplexity": 2.3665, - "train_runtime": 294.5707, - "train_samples_per_second": 148.093 + "perplexity": 2.3665888138128466, + "train_runtime": 394.5646, + "train_samples_per_second": 238.486 } }, "tests/test_examples.py::MultiCardCausalLanguageModelingLORAFSDPCompileExampleTester::test_run_lora_clm_llama-7b_multi_card": { @@ -321,9 +321,9 @@ "train_samples_per_second": 93.5 }, "gaudi3": { - "perplexity": 2.4259, - "train_runtime": 186.2483, - "train_samples_per_second": 93.5 + "perplexity": 2.42632366178759, + "train_runtime": 98.5791, + "train_samples_per_second": 126.028 } }, "tests/test_examples.py::MultiCardCausalLanguageModelingLlamaAdapterExampleTester::test_run_lora_clm_llama-7b_multi_card": { @@ -333,9 +333,9 @@ "train_samples_per_second": 294 }, "gaudi3": { - "perplexity": 5.575, - "train_runtime": 131.7, - "train_samples_per_second": 294 + "perplexity": 5.575957971980852, + "train_runtime": 227.3213, + "train_samples_per_second": 504.974 } }, "tests/test_examples.py::MultiCardCausalLanguageModelingLnExampleTester::test_run_lora_clm_llama-7b_multi_card": { @@ -345,9 +345,9 @@ "train_samples_per_second": 165 }, "gaudi3": { - "perplexity": 2.83, - "train_runtime": 249, - "train_samples_per_second": 165 + "perplexity": 2.842264808115683, + "train_runtime": 332.9477, + "train_samples_per_second": 267.004 } }, "tests/test_examples.py::MultiCardCausalLanguageModelingLoRACPExampleTester::test_run_lora_clm_llama-7b_deepspeed": { @@ -357,9 +357,9 @@ "train_samples_per_second": 34.41 }, "gaudi3": { - "perplexity": 2.8889, - "train_runtime": 147.3597, - "train_samples_per_second": 34.41 + "perplexity": 2.8421374130082477, + "train_runtime": 219.1417, + "train_samples_per_second": 55.554 } }, "tests/test_examples.py::MultiCardCausalLanguageModelingLoRAFP8ExampleTester::test_run_lora_clm_llama-7b_multi_card": { @@ -369,9 +369,9 @@ "train_samples_per_second": 232.439 }, "gaudi3": { - "perplexity": 2.3692, - "train_runtime": 411.9935, - "train_samples_per_second": 232.439 + "perplexity": 2.3750491436810424, + "train_runtime": 547.5649, + "train_samples_per_second": 323.175 } }, "tests/test_examples.py::MultiCardCausalLanguageModelingPTuningExampleTester::test_run_prompt_tuning_clm_llama-7b_multi_card": { @@ -381,9 +381,9 @@ "train_samples_per_second": 63.161 }, "gaudi3": { - "perplexity": 1.047, - "train_runtime": 18.7, - "train_samples_per_second": 63.161 + "perplexity": 1.0262332298756216, + "train_runtime": 16.2913, + "train_samples_per_second": 78.376 } }, "tests/test_examples.py::MultiCardCausalLanguageModelingPrefixTuningExampleTester::test_run_prompt_tuning_clm_llama-7b_multi_card": { @@ -393,9 +393,9 @@ "train_samples_per_second": 63.249 }, "gaudi3": { - "perplexity": 1.172, - "train_runtime": 16.1, - "train_samples_per_second": 63.249 + "perplexity": 1.1720024747280242, + "train_runtime": 15.1138, + "train_samples_per_second": 67.894 } }, "tests/test_examples.py::MultiCardCausalLanguageModelingPromptTuningExampleTester::test_run_prompt_tuning_clm_llama-7b_multi_card": { @@ -405,9 +405,9 @@ "train_samples_per_second": 63.161 }, "gaudi3": { - "perplexity": 1.224, - "train_runtime": 16.5, - "train_samples_per_second": 63.161 + "perplexity": 1.2158095633720596, + "train_runtime": 14.0663, + "train_samples_per_second": 75.406 } }, "tests/test_examples.py::MultiCardCausalLanguageModelingVeraExampleTester::test_run_lora_clm_llama-7b_multi_card": { @@ -417,9 +417,9 @@ "train_samples_per_second": 127.305 }, "gaudi3": { - "perplexity": 9.064502567217577, - "train_runtime": 312.9258, - "train_samples_per_second": 127.305 + "perplexity": 8.65669958765362, + "train_runtime": 261.8749, + "train_samples_per_second": 199.0 } }, "tests/test_examples.py::MultiCardDPOExampleTester::test_dpo_llama-7b_multi_card": { @@ -428,8 +428,8 @@ "train_samples_per_second": 13.499 }, "gaudi3": { - "train_runtime": 234.6471, - "train_samples_per_second": 13.499 + "train_runtime": 194.4848, + "train_samples_per_second": 16.454 } }, "tests/test_examples.py::MultiCardImageClassificationExampleTester::test_run_image_classification_swin-base-patch4-window7-224-in22k_multi_card": { @@ -444,9 +444,9 @@ "train_samples_per_second": 6202.525 }, "gaudi3": { - "eval_accuracy": 0.9821, - "train_runtime": 62.9986, - "train_samples_per_second": 6202.525 + "eval_accuracy": 0.9817333333333333, + "train_runtime": 74.7483, + "train_samples_per_second": 8253.709 } }, "tests/test_examples.py::MultiCardImageClassificationExampleTester::test_run_image_classification_vit-base-patch16-224-in21k_multi_card": { @@ -461,9 +461,9 @@ "train_samples_per_second": 6718.643 }, "gaudi3": { - "eval_accuracy": 0.9679, - "train_runtime": 23.99, - "train_samples_per_second": 6718.643 + "eval_accuracy": 0.9677333333333333, + "train_runtime": 33.4011, + "train_samples_per_second": 6636.054 } }, "tests/test_examples.py::MultiCardImageToTextModelingLoRAExampleTester::test_run_image2text_lora_finetune_Llama-3.2-11B-Vision-Instruct_multi_card": { @@ -473,9 +473,9 @@ "train_samples_per_second": 20.48 }, "gaudi3": { - "eval_accuracy": 0.6, - "train_runtime": 350, - "train_samples_per_second": 20.48 + "eval_accuracy": 0.9044574025188373, + "train_runtime": 397.9607, + "train_samples_per_second": 39.088 } }, "tests/test_examples.py::MultiCardImageToTextModelingLoRAExampleTester::test_run_image2text_lora_finetune_idefics2-8b_multi_card": { @@ -485,9 +485,9 @@ "train_samples_per_second": 11.8 }, "gaudi3": { - "eval_accuracy": 0.6, - "train_runtime": 286, - "train_samples_per_second": 11.8 + "eval_accuracy": 0.6910165783279163, + "train_runtime": 273.7778, + "train_samples_per_second": 17.93 } }, "tests/test_examples.py::MultiCardImageToTextModelingLoRAExampleTester::test_run_image2text_lora_finetune_llava-1.5-7b-hf_multi_card": { @@ -497,9 +497,9 @@ "train_samples_per_second": 25.146 }, "gaudi3": { - "eval_accuracy": 0.2122, - "train_runtime": 118.5782, - "train_samples_per_second": 25.146 + "eval_accuracy": 0.20785648331296863, + "train_runtime": 184.9003, + "train_samples_per_second": 27.828 } }, "tests/test_examples.py::MultiCardMaskedLanguageModelingExampleTester::test_run_mlm_roberta-large_multi_card": { @@ -514,9 +514,9 @@ "train_samples_per_second": 1056.875 }, "gaudi3": { - "perplexity": 2.829522488584474, - "train_runtime": 22.7101, - "train_samples_per_second": 1056.875 + "perplexity": 2.8534683742096933, + "train_runtime": 53.0805, + "train_samples_per_second": 1335.957 } }, "tests/test_examples.py::MultiCardPPOExampleTester::test_ppo_llama-7b_multi_card": { @@ -525,8 +525,8 @@ "train_samples_per_second": 0.5 }, "gaudi3": { - "train_runtime": 62, - "train_samples_per_second": 0.5 + "train_runtime": 40.73775029182434, + "train_samples_per_second": 0.7855122035647137 } }, "tests/test_examples.py::MultiCardProteinFoldingClassificationTester::test_run_sequence_classification_protst-esm1b-for-sequential-classification_multi_card": { @@ -536,9 +536,9 @@ "train_samples_per_second": 768.648 }, "gaudi3": { - "eval_accuracy": 0.5436668594563332, - "train_runtime": 38.9504, - "train_samples_per_second": 768.648 + "eval_accuracy": 0.5442452284557547, + "train_runtime": 40.0248, + "train_samples_per_second": 1564.079 } }, "tests/test_examples.py::MultiCardQuestionAnsweringExampleTester::test_run_qa_roberta-large_multi_card": { @@ -553,9 +553,9 @@ "train_samples_per_second": 2138.366 }, "gaudi3": { - "eval_f1": 94.09, - "train_runtime": 79.333, - "train_samples_per_second": 2138.366 + "eval_f1": 94.33668918864852, + "train_runtime": 153.0279, + "train_samples_per_second": 3146.332 } }, "tests/test_examples.py::MultiCardRewardExampleTester::test_reward_modeling_llama-7b_multi_card": { @@ -564,8 +564,8 @@ "train_samples_per_second": 1.6 }, "gaudi3": { - "train_runtime": 250, - "train_samples_per_second": 1.6 + "train_runtime": 135.1176, + "train_samples_per_second": 3.027 } }, "tests/test_examples.py::MultiCardSFTChatExampleTester::test_sft_Qwen2-7B_multi_card": { @@ -574,8 +574,8 @@ "train_samples_per_second": 7.342 }, "gaudi3": { - "train_runtime": 423.995, - "train_samples_per_second": 7.342 + "train_runtime": 587.8481, + "train_samples_per_second": 13.968 } }, "tests/test_examples.py::MultiCardSFTChatPeftExampleTester::test_sft_Qwen2-7B_multi_card": { @@ -584,8 +584,8 @@ "train_samples_per_second": 120 }, "gaudi3": { - "train_runtime": 410, - "train_samples_per_second": 120 + "train_runtime": 364.7036, + "train_samples_per_second": 193.023 } }, "tests/test_examples.py::MultiCardSFTExampleTester::test_sft_llama-7b_multi_card": { @@ -594,8 +594,8 @@ "train_samples_per_second": 51.54 }, "gaudi3": { - "train_runtime": 206, - "train_samples_per_second": 51.54 + "train_runtime": 316.0836, + "train_samples_per_second": 86.193 } }, "tests/test_examples.py::MultiCardSeq2SeqSpeechRecognitionExampleTester::test_run_speech_recognition_seq2seq_whisper-small_multi_card": { @@ -612,10 +612,10 @@ "train_samples_per_second": 218.0 }, "gaudi3": { - "eval_samples_per_second": 31.0, - "eval_wer": 0.4693843594009983, - "train_runtime": 380.0, - "train_samples_per_second": 218.0 + "eval_samples_per_second": 64.339, + "eval_wer": 0.38905990016638936, + "train_runtime": 290.6815, + "train_samples_per_second": 463.628 } }, "tests/test_examples.py::MultiCardSpeechRecognitionExampleTester::test_run_speech_recognition_ctc_wav2vec2-large-lv60_multi_card": { @@ -632,10 +632,10 @@ "train_samples_per_second": 225.572 }, "gaudi3": { - "eval_samples_per_second": 196.665, - "eval_wer": 0.1109, - "train_runtime": 308.8036, - "train_samples_per_second": 225.572 + "eval_samples_per_second": 491.004, + "eval_wer": 0.06197937326457755, + "train_runtime": 255.782, + "train_samples_per_second": 292.161 } }, "tests/test_examples.py::MultiCardTextClassificationExampleTester::test_run_glue_bert-large-uncased-whole-word-masking_multi_card": { @@ -650,9 +650,9 @@ "train_samples_per_second": 2845.068 }, "gaudi3": { - "eval_f1": 0.8452579034941764, - "train_runtime": 31.445, - "train_samples_per_second": 2845.068 + "eval_f1": 0.89198606271777, + "train_runtime": 61.3444, + "train_samples_per_second": 1826.566 } }, "tests/test_examples.py::MultiCardVisionLanguageExampleTester::test_run_clip_clip-roberta_multi_card": { @@ -665,8 +665,8 @@ "train_samples_per_second": 14124 }, "gaudi3": { - "train_runtime": 59.5, - "train_samples_per_second": 14124 + "train_runtime": 64.3878, + "train_samples_per_second": 19625.412 } }, "tests/test_examples.py::QuestionAnsweringExampleTester::test_run_qa_roberta-large_single_card": { @@ -681,9 +681,9 @@ "train_samples_per_second": 266.47 }, "gaudi3": { - "eval_f1": 94.5886, - "train_runtime": 361.4789, - "train_samples_per_second": 266.47 + "eval_f1": 94.36192902198283, + "train_runtime": 260.988, + "train_samples_per_second": 423.007 } }, "tests/test_examples.py::TextClassificationExampleTester::test_run_glue_bert-large-uncased-whole-word-masking_single_card": { @@ -698,9 +698,9 @@ "train_samples_per_second": 1100.598 }, "gaudi3": { - "eval_f1": 0.867, - "train_runtime": 33.2909, - "train_samples_per_second": 1100.598 + "eval_f1": 0.8826446280991735, + "train_runtime": 74.0631, + "train_samples_per_second": 1652.436 } } } \ No newline at end of file diff --git a/tests/baselines/fixture/tests/test_fsdp_examples.json b/tests/baselines/fixture/tests/test_fsdp_examples.json index b9e17c7354..67e7d56879 100644 --- a/tests/baselines/fixture/tests/test_fsdp_examples.json +++ b/tests/baselines/fixture/tests/test_fsdp_examples.json @@ -15,8 +15,8 @@ "train_samples_per_second": 85.016 }, "gaudi3": { - "train_loss": 0.9093, - "train_samples_per_second": 85.016 + "train_loss": 0.9092939383912795, + "train_samples_per_second": 119.866 } } } \ No newline at end of file diff --git a/tests/baselines/fixture/tests/test_image_to_text_example.json b/tests/baselines/fixture/tests/test_image_to_text_example.json index e95c6d88d8..58dbd84613 100644 --- a/tests/baselines/fixture/tests/test_image_to_text_example.json +++ b/tests/baselines/fixture/tests/test_image_to_text_example.json @@ -4,7 +4,7 @@ "throughput": 21.89944593215077 }, "gaudi3": { - "throughput": 21.89944593215077 + "throughput": 55.82131026867695 } }, "tests/test_image_to_text_example.py::test_image_to_text_bf16[Qwen/Qwen2-VL-2B-Instruct-1]": { @@ -12,7 +12,7 @@ "throughput": 28.755882208438422 }, "gaudi3": { - "throughput": 28.755882208438422 + "throughput": 85.53160250422563 } }, "tests/test_image_to_text_example.py::test_image_to_text_bf16[Qwen/Qwen2-VL-7B-Instruct-1]": { @@ -20,7 +20,7 @@ "throughput": 19.32562189532818 }, "gaudi3": { - "throughput": 19.32562189532818 + "throughput": 17.216165111759725 } }, "tests/test_image_to_text_example.py::test_image_to_text_bf16[google/paligemma-3b-mix-224-1]": { @@ -28,7 +28,7 @@ "throughput": 132.8949150246155 }, "gaudi3": { - "throughput": 132.8949150246155 + "throughput": 215.66261236773295 } }, "tests/test_image_to_text_example.py::test_image_to_text_bf16[llava-hf/llava-1.5-13b-hf-1]": { @@ -61,7 +61,7 @@ "throughput": 33.17984878151546 }, "gaudi3": { - "throughput": 33.17984878151546 + "throughput": 72.22445594285129 } }, "tests/test_image_to_text_example.py::test_image_to_text_bf16[llava-hf/llava-v1.6-vicuna-13b-hf-1]": { @@ -72,7 +72,7 @@ "throughput": 23.527610042925 }, "gaudi3": { - "throughput": 23.527610042925 + "throughput": 45.50628237484548 } }, "tests/test_image_to_text_example.py::test_image_to_text_bf16[llava-hf/llava-v1.6-vicuna-7b-hf-1]": { @@ -80,7 +80,7 @@ "throughput": 35.00608681379742 }, "gaudi3": { - "throughput": 35.00608681379742 + "throughput": 73.24265508277661 } }, "tests/test_image_to_text_example.py::test_image_to_text_bf16[meta-llama/Llama-3.2-11B-Vision-Instruct-1]": { @@ -88,7 +88,7 @@ "throughput": 18.974541922240313 }, "gaudi3": { - "throughput": 18.974541922240313 + "throughput": 60.21756704358577 } }, "tests/test_image_to_text_example.py::test_image_to_text_bf16[tiiuae/falcon-11B-vlm-1]": { @@ -96,7 +96,7 @@ "throughput": 23.69260849957278 }, "gaudi3": { - "throughput": 23.69260849957278 + "throughput": 42.77946694511338 } }, "tests/test_image_to_text_example.py::test_image_to_text_fp8[llava-hf/llava-1.5-13b-hf-1]": { @@ -120,7 +120,7 @@ "throughput": 45.011551008367086 }, "gaudi3": { - "throughput": 45.011551008367086 + "throughput": 85.4014722462956 } }, "tests/test_image_to_text_example.py::test_image_to_text_fp8[llava-hf/llava-v1.6-vicuna-13b-hf-1]": { @@ -128,7 +128,7 @@ "throughput": 30.9535718774675 }, "gaudi3": { - "throughput": 30.9535718774675 + "throughput": 56.447951664149116 } }, "tests/test_image_to_text_example.py::test_image_to_text_fp8[llava-hf/llava-v1.6-vicuna-7b-hf-1]": { @@ -136,7 +136,7 @@ "throughput": 45.18544502949674 }, "gaudi3": { - "throughput": 45.18544502949674 + "throughput": 83.9326869276268 } } } \ No newline at end of file diff --git a/tests/baselines/fixture/tests/test_object_segmentation.json b/tests/baselines/fixture/tests/test_object_segmentation.json index 65ae50ea0f..d70be2c64b 100644 --- a/tests/baselines/fixture/tests/test_object_segmentation.json +++ b/tests/baselines/fixture/tests/test_object_segmentation.json @@ -4,7 +4,7 @@ "latency": 5.3107380867004395 }, "gaudi3": { - "latency": 5.3107380867004395 + "latency": 3.9719343185424805 } } } \ No newline at end of file diff --git a/tests/baselines/fixture/tests/test_openclip_vqa.json b/tests/baselines/fixture/tests/test_openclip_vqa.json index 2daee462ac..bb47580588 100644 --- a/tests/baselines/fixture/tests/test_openclip_vqa.json +++ b/tests/baselines/fixture/tests/test_openclip_vqa.json @@ -7,7 +7,7 @@ "throughput": 1472 }, "gaudi3": { - "throughput": 1472 + "throughput": 1289.3560859645047 } }, "tests/test_openclip_vqa.py::test_openclip_vqa_bf16[microsoft/BiomedCLIP-PubMedBERT_256-vit_base_patch16_224]": { @@ -18,7 +18,7 @@ "throughput": 1816 }, "gaudi3": { - "throughput": 1816 + "throughput": 1876.4408565804385 } } } \ No newline at end of file diff --git a/tests/baselines/fixture/tests/test_sentence_transformers.json b/tests/baselines/fixture/tests/test_sentence_transformers.json index dfa5753e50..36b07cd3ea 100644 --- a/tests/baselines/fixture/tests/test_sentence_transformers.json +++ b/tests/baselines/fixture/tests/test_sentence_transformers.json @@ -7,7 +7,7 @@ "measured_throughput": 3614.2610109716247 }, "gaudi3": { - "measured_throughput": 3614.2610109716247 + "measured_throughput": 5674.813347163265 } }, "tests/test_sentence_transformers.py::test_compute_embeddings_throughput[sentence-transformers/all-MiniLM-L6-v2]": { @@ -18,7 +18,7 @@ "measured_throughput": 2615.6975354038477 }, "gaudi3": { - "measured_throughput": 2615.6975354038477 + "measured_throughput": 6489.3086857211365 } }, "tests/test_sentence_transformers.py::test_compute_embeddings_throughput[sentence-transformers/all-distilroberta-v1]": { @@ -29,7 +29,7 @@ "measured_throughput": 958.5097903298335 }, "gaudi3": { - "measured_throughput": 958.5097903298335 + "measured_throughput": 6105.954239105652 } }, "tests/test_sentence_transformers.py::test_compute_embeddings_throughput[sentence-transformers/all-mpnet-base-v2]": { @@ -40,7 +40,7 @@ "measured_throughput": 762.5595168883357 }, "gaudi3": { - "measured_throughput": 762.5595168883357 + "measured_throughput": 5025.5970390534085 } }, "tests/test_sentence_transformers.py::test_compute_embeddings_throughput[sentence-transformers/distiluse-base-multilingual-cased-v1]": { @@ -51,7 +51,7 @@ "measured_throughput": 3487.3319366004903 }, "gaudi3": { - "measured_throughput": 3487.3319366004903 + "measured_throughput": 5908.987916285729 } }, "tests/test_sentence_transformers.py::test_compute_embeddings_throughput[sentence-transformers/distiluse-base-multilingual-cased-v2]": { @@ -62,7 +62,7 @@ "measured_throughput": 3807.2486282025716 }, "gaudi3": { - "measured_throughput": 3807.2486282025716 + "measured_throughput": 5995.942563633102 } }, "tests/test_sentence_transformers.py::test_compute_embeddings_throughput[sentence-transformers/multi-qa-MiniLM-L6-cos-v1]": { @@ -73,7 +73,7 @@ "measured_throughput": 1208.3672807492396 }, "gaudi3": { - "measured_throughput": 1208.3672807492396 + "measured_throughput": 6369.4219807072195 } }, "tests/test_sentence_transformers.py::test_compute_embeddings_throughput[sentence-transformers/multi-qa-distilbert-cos-v1]": { @@ -84,7 +84,7 @@ "measured_throughput": 944.6166139694299 }, "gaudi3": { - "measured_throughput": 944.6166139694299 + "measured_throughput": 6167.298763111252 } }, "tests/test_sentence_transformers.py::test_compute_embeddings_throughput[sentence-transformers/multi-qa-mpnet-base-dot-v1]": { @@ -95,7 +95,7 @@ "measured_throughput": 545.3360251829846 }, "gaudi3": { - "measured_throughput": 545.3360251829846 + "measured_throughput": 5011.953212884994 } }, "tests/test_sentence_transformers.py::test_compute_embeddings_throughput[sentence-transformers/paraphrase-MiniLM-L3-v2]": { @@ -106,7 +106,7 @@ "measured_throughput": 5734.318427972881 }, "gaudi3": { - "measured_throughput": 5734.318427972881 + "measured_throughput": 7073.782785445982 } }, "tests/test_sentence_transformers.py::test_compute_embeddings_throughput[sentence-transformers/paraphrase-albert-small-v2]": { @@ -117,7 +117,7 @@ "measured_throughput": 3896.1911011860166 }, "gaudi3": { - "measured_throughput": 3896.1911011860166 + "measured_throughput": 6136.85257090509 } }, "tests/test_sentence_transformers.py::test_compute_embeddings_throughput[sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2]": { @@ -128,7 +128,7 @@ "measured_throughput": 3558.0778715789693 }, "gaudi3": { - "measured_throughput": 3558.0778715789693 + "measured_throughput": 5650.834160594289 } }, "tests/test_sentence_transformers.py::test_compute_embeddings_throughput[sentence-transformers/paraphrase-multilingual-mpnet-base-v2]": { @@ -139,7 +139,7 @@ "measured_throughput": 2392.1654748794062 }, "gaudi3": { - "measured_throughput": 2392.1654748794062 + "measured_throughput": 4906.993110085868 } } } \ 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 7679a8171b..b0c1f40f81 100644 --- a/tests/baselines/fixture/tests/test_text_generation_example.json +++ b/tests/baselines/fixture/tests/test_text_generation_example.json @@ -4,7 +4,7 @@ "throughput": 456.7 }, "gaudi3": { - "throughput": 456.7 + "throughput": 828.916211466145 } }, "tests/test_text_generation_example.py::test_text_generation_beam_search[Qwen/Qwen2-7b-Instruct-1-True]": { @@ -12,7 +12,7 @@ "throughput": 91.24938949709826 }, "gaudi3": { - "throughput": 91.24938949709826 + "throughput": 98.57537548249874 } }, "tests/test_text_generation_example.py::test_text_generation_bf16_1x[CohereForAI/c4ai-command-r-v01-1-False-False]": { @@ -33,10 +33,10 @@ }, "tests/test_text_generation_example.py::test_text_generation_bf16_1x[EleutherAI/gpt-j-6b-1-False-False]": { "gaudi2": { - "throughput": 160.5823842101192 + "throughput": 143.64228300147943 }, "gaudi3": { - "throughput": 160.5823842101192 + "throughput": 165.9126964936202 } }, "tests/test_text_generation_example.py::test_text_generation_bf16_1x[EleutherAI/gpt-j-6b-1-True-False]": { @@ -57,7 +57,7 @@ "throughput": 50.67672679310354 }, "gaudi3": { - "throughput": 50.67672679310354 + "throughput": 61.74067195778036 } }, "tests/test_text_generation_example.py::test_text_generation_bf16_1x[Qwen/Qwen1.5-7B-1-False-False]": { @@ -78,7 +78,7 @@ "throughput": 44.25834541569395 }, "gaudi3": { - "throughput": 44.25834541569395 + "throughput": 179.15343204459856 } }, "tests/test_text_generation_example.py::test_text_generation_bf16_1x[Qwen/Qwen2-7B-256-False-True]": { @@ -87,8 +87,8 @@ "throughput": 8870.945160540245 }, "gaudi3": { - "output": "DeepSpeed is a machine learning framework that provides a unified interface for training deep learning models. It is designed to be easy to use and to provide high performance. DeepSpeed is built on top of PyTorch and TensorFlow, and it supports a wide range of models, including transformers, convolutional neural networks, and recurrent neural networks.\nDeepSpeed is a machine learning framework that provides a unified interface for training deep learning models. It is designed to be easy to use and to provide high performance. DeepSpeed is built on top of Py", - "throughput": 8870.945160540245 + "output": "DeepSpeed is a machine learning framework that provides a unified interface for training deep learning models. It is designed to be easy to use and to provide high performance. DeepSpeed is built on top of PyTorch and TensorFlow, and it supports both CPU and GPU training. It also provides a number of features that are not available in other frameworks, such as automatic mixed precision training and distributed training.\nDeepSpeed is a machine learning framework that provides a unified interface for training deep learning models. It is designed to be easy to use and to", + "throughput": 14633.079557607358 } }, "tests/test_text_generation_example.py::test_text_generation_bf16_1x[Qwen/Qwen2.5-7B-4-False-False]": { @@ -96,7 +96,7 @@ "throughput": 490 }, "gaudi3": { - "throughput": 490 + "throughput": 633.0694674407139 } }, "tests/test_text_generation_example.py::test_text_generation_bf16_1x[Salesforce/codegen2-1B-1-False-False]": { @@ -107,7 +107,7 @@ "throughput": 446.4029486883532 }, "gaudi3": { - "throughput": 446.4029486883532 + "throughput": 405.96090453183643 } }, "tests/test_text_generation_example.py::test_text_generation_bf16_1x[THUDM/chatglm2-6b-1-True-False]": { @@ -115,7 +115,7 @@ "throughput": 150 }, "gaudi3": { - "throughput": 150 + "throughput": 169.28444068272802 } }, "tests/test_text_generation_example.py::test_text_generation_bf16_1x[THUDM/chatglm3-6b-1-True-False]": { @@ -123,7 +123,7 @@ "throughput": 150 }, "gaudi3": { - "throughput": 150 + "throughput": 168.9312894863455 } }, "tests/test_text_generation_example.py::test_text_generation_bf16_1x[adept/persimmon-8b-base-1-False-False]": { @@ -136,7 +136,7 @@ "throughput": 366.73968820698406 }, "gaudi3": { - "throughput": 366.73968820698406 + "throughput": 359.5154721132213 } }, "tests/test_text_generation_example.py::test_text_generation_bf16_1x[baichuan-inc/Baichuan2-13B-Chat-1-False-False]": { @@ -144,7 +144,7 @@ "throughput": 66 }, "gaudi3": { - "throughput": 66 + "throughput": 83.1114363254922 } }, "tests/test_text_generation_example.py::test_text_generation_bf16_1x[baichuan-inc/Baichuan2-7B-Chat-1-True-False]": { @@ -152,7 +152,7 @@ "throughput": 108 }, "gaudi3": { - "throughput": 108 + "throughput": 129.18924637215144 } }, "tests/test_text_generation_example.py::test_text_generation_bf16_1x[bigcode/starcoder-1-False-False]": { @@ -167,7 +167,7 @@ }, "gaudi3": { "output": "def print_hello_world():\n print(\"Hello World\")\n\ndef print_hello_world_twice():\n print_hello_world()\n print_hello_world()\n\ndef print_hello_world_thrice():\n print_hello_world()\n print_hello_world()\n print_hello_world()\n\ndef print_hello_world_four_times():\n print_hello_world()\n print_hello_world()\n print_hello_world()\n ", - "throughput": 6846.575763562658 + "throughput": 14438.542540850205 } }, "tests/test_text_generation_example.py::test_text_generation_bf16_1x[bigcode/starcoder2-3b-1-False-False]": { @@ -182,7 +182,7 @@ }, "gaudi3": { "output": "def print_hello_world():\n print(\"Hello World\")\n\ndef print_hello_world_with_name(name):\n print(\"Hello World, \" + name)\n\ndef print_hello_world_with_name_and_age(name, age):\n print(\"Hello World, \" + name + \", \" + str(age))\n\ndef print_hello_world_with_name_and_age_and_gender(name, age, gender):\n print(\"Hello", - "throughput": 261.07213776344133 + "throughput": 279.92066126452653 } }, "tests/test_text_generation_example.py::test_text_generation_bf16_1x[bigscience/bloomz-7b1-1-False-False]": { @@ -193,7 +193,7 @@ "throughput": 130.0472971205316 }, "gaudi3": { - "throughput": 130.0472971205316 + "throughput": 155.29323724597498 } }, "tests/test_text_generation_example.py::test_text_generation_bf16_1x[codellama/CodeLlama-34b-hf-1-True-False]": { @@ -201,7 +201,7 @@ "throughput": 32.644 }, "gaudi3": { - "throughput": 32.644 + "throughput": 42.94755856794396 } }, "tests/test_text_generation_example.py::test_text_generation_bf16_1x[deepseek-ai/DeepSeek-V2-Lite-1-False-False]": { @@ -209,7 +209,7 @@ "throughput": 35 }, "gaudi3": { - "throughput": 35 + "throughput": 149.2189570033595 } }, "tests/test_text_generation_example.py::test_text_generation_bf16_1x[facebook/xglm-1.7B-1-False-False]": { @@ -226,8 +226,8 @@ "throughput": 36.578709544111 }, "gaudi3": { - "output": "DeepSpeed is a machine learning framework that enables you to train models with trillions of parameters and beyond, using model parallelism to partition large models over multiple GPUs.\n\nThe following is a brief introduction to the DeepSpeed model parallel training.\n\n

    1. Introduction

    \n\nThe DeepSpeed model parallel training is a simple and effective way to train large models. It is a framework that enables you to train models with trillions of parameters and beyond.\n\nDeepSpeed is a distributed deep learning optimization toolkit that makes it easy and efficient", - "throughput": 36.578709544111 + "output": "DeepSpeed is a machine learning framework that enables you to train large models on a single GPU. It is a framework that is used to train large models on a single GPU.\n\nThe main idea is to use a large amount of memory to fit the model on a single GPU.\n\nThe main idea of \u200b\u200bthe algorithm is to use the gradient of the loss function to update the model parameters.\n\nThe main idea of \u200b\u200bthe algorithm is to use the gradient of the loss function to update the model parameters.\n\nThe main idea of", + "throughput": 46.04685368495098 } }, "tests/test_text_generation_example.py::test_text_generation_bf16_1x[google/gemma-2-9b-1-False-True]": { @@ -237,7 +237,7 @@ }, "gaudi3": { "output": "DeepSpeed is a machine learning framework that enables training of large-scale deep learning models on a single GPU or across multiple GPUs. It is designed to be easy to use and highly scalable, making it a powerful tool for researchers and practitioners working with large-scale deep learning models.\n\nDeepSpeed is built on top of PyTorch, a popular deep learning framework, and provides a set of tools and libraries that make it easy to train large-scale models. It includes features such as zero-shot inference, which allows models to be", - "throughput": 92.302359446567 + "throughput": 111.60209707224463 } }, "tests/test_text_generation_example.py::test_text_generation_bf16_1x[google/gemma-7b-1-False-False]": { @@ -252,7 +252,7 @@ }, "gaudi3": { "output": "DeepSpeed is a machine learning framework that enables training of large-scale models on commodity hardware. It is designed to be a drop-in replacement for PyTorch, and it is compatible with the existing PyTorch ecosystem. DeepSpeed is designed to be easy to use, and it provides a number of features that make it easy to train large-scale models. DeepSpeed is designed to be scalable, and it can be used to train models on a single machine or on a cluster of machines. DeepSpeed is designed to be efficient,", - "throughput": 109.70751574382221 + "throughput": 135.97272017864475 } }, "tests/test_text_generation_example.py::test_text_generation_bf16_1x[gpt2-xl-1-False-False]": { @@ -263,7 +263,7 @@ "throughput": 281.8734689674413 }, "gaudi3": { - "throughput": 281.8734689674413 + "throughput": 286.8456278152758 } }, "tests/test_text_generation_example.py::test_text_generation_bf16_1x[meta-llama/Llama-2-7b-hf-1-True-False]": { @@ -273,12 +273,12 @@ }, "tests/test_text_generation_example.py::test_text_generation_bf16_1x[meta-llama/Llama-2-7b-hf-1-True-True]": { "gaudi2": { - "output": "DeepSpeed is a machine learning framework for deep learning. It is designed to be fast and efficient, while also being easy to use. DeepSpeed is based on the TensorFlow framework, and it uses the TensorFlow library to perform computations.\nDeepSpeed is a deep learning framework that is designed to be fast and efficient. It is based on the TensorFlow library and uses the TensorFlow library to perform computations. DeepSpeed is designed to be easy to use and to provide a high level of performance", + "output": "DeepSpeed is a machine learning framework for deep learning. It is designed to be fast and efficient, while also being easy to use. DeepSpeed is based on the TensorFlow framework, and it uses the TensorFlow library to perform computations.\nDeepSpeed is a deep learning framework that is designed to be fast and efficient. It is based on the TensorFlow library and uses the TensorFlow library to perform computations. DeepSpeed is designed to be easy to use and to provide a high level of flex", "throughput": 141.25776956002076 }, "gaudi3": { "output": "DeepSpeed is a machine learning framework for deep learning. It is designed to be fast and efficient, while also being easy to use. DeepSpeed is based on the TensorFlow framework, and it uses the TensorFlow library to perform computations.\nDeepSpeed is a deep learning framework that is designed to be fast and efficient. It is based on the TensorFlow library and uses the TensorFlow library to perform computations. DeepSpeed is designed to be easy to use and to provide a high level of flex", - "throughput": 141.25776956002076 + "throughput": 173.7868608608374 } }, "tests/test_text_generation_example.py::test_text_generation_bf16_1x[meta-llama/Llama-2-7b-hf-512-False-False]": { @@ -286,7 +286,7 @@ "throughput": 8711 }, "gaudi3": { - "throughput": 8711 + "throughput": 15150.480373545233 } }, "tests/test_text_generation_example.py::test_text_generation_bf16_1x[meta-llama/Llama-2-7b-hf-512-True-False]": { @@ -294,7 +294,7 @@ "throughput": 12808 }, "gaudi3": { - "throughput": 12808 + "throughput": 23362.95410956595 } }, "tests/test_text_generation_example.py::test_text_generation_bf16_1x[meta-llama/Meta-Llama-3-8B-1-True-False]": { @@ -302,7 +302,7 @@ "throughput": 129 }, "gaudi3": { - "throughput": 129 + "throughput": 162.03504027530752 } }, "tests/test_text_generation_example.py::test_text_generation_bf16_1x[microsoft/phi-2-1-False-False]": { @@ -313,7 +313,7 @@ "throughput": 224.72307766211117 }, "gaudi3": { - "throughput": 224.72307766211117 + "throughput": 236.53539137265457 } }, "tests/test_text_generation_example.py::test_text_generation_bf16_1x[mistralai/Mistral-7B-v0.1-1-True-False]": { @@ -323,22 +323,22 @@ }, "tests/test_text_generation_example.py::test_text_generation_bf16_1x[mistralai/Mistral-7B-v0.1-1-True-True]": { "gaudi2": { - "output": "DeepSpeed is a machine learning framework that accelerates training of large models on a single machine or distributed systems. It is designed to be easy to use and flexible, allowing users to quickly train models on a variety of hardware platforms.\n\nDeepSpeed is a machine learning framework that accelerates training of large models on a single machine or distributed systems. It is designed to be easy to use and flexible, allowing users to quickly train models on a variety of hardware platforms.\n\nDeepSpeed is a machine learning framework that accelerates training", - "throughput": 130.2172236767782 + "output": "DeepSpeed is a machine learning framework that accelerates training of large models on a single machine or distributed systems. It is designed to be compatible with PyTorch and TensorFlow, and can be used to train models on a single machine or on a distributed system.\n\nDeepSpeed is a machine learning framework that accelerates training of large models on a single machine or distributed systems. It is designed to be compatible with PyTorch and TensorFlow, and can be used to train models on a single machine or on a distributed system", + "throughput": 134.94827207337997 }, "gaudi3": { "output": "DeepSpeed is a machine learning framework that accelerates training of large models on a single machine or distributed systems. It is designed to be compatible with PyTorch and TensorFlow, and can be used to train models on a single machine or on a distributed system.\n\nDeepSpeed is a machine learning framework that accelerates training of large models on a single machine or distributed systems. It is designed to be compatible with PyTorch and TensorFlow, and can be used to train models on a single machine or on a distributed system", - "throughput": 130.2172236767782 + "throughput": 160.48685620965531 } }, "tests/test_text_generation_example.py::test_text_generation_bf16_1x[mistralai/Mixtral-8x7B-v0.1-1-False-True]": { "gaudi2": { - "output": "DeepSpeed is a machine learning framework that enables training of large models on a single machine with a single GPU. It is designed to be easy to use and efficient, and it can be used to train models on a variety of tasks.\n\n## What is DeepSpeed?\n\nDeepSpeed is a machine learning framework that enables training of large models on a single machine with a single GPU. It is designed to be easy to use and efficient, and it can be used to train models on a variety of tasks.\n\n##", - "throughput": 23.7931001677926 + "output": "DeepSpeed is a machine learning framework that enables training of large models on a single machine with a single GPU. It is designed to be easy to use and efficient, and it can be used to train models on a variety of tasks.\n\n## Introduction\n\nDeepSpeed is a machine learning framework that enables training of large models on a single machine with a single GPU. It is designed to be easy to use and efficient, and it can be used to train models on a variety of tasks.\n\n## What is DeepSpeed", + "throughput": 71.29570003665306 }, "gaudi3": { - "output": "DeepSpeed is a machine learning framework that enables training of large models on a single machine with a single GPU. It is designed to be easy to use and efficient, and it can be used to train models on a variety of tasks.\n\n## Introduction\n\nDeepSpeed is a machine learning framework that enables training of large models on a single machine with a single GPU. It is designed to be easy to use and efficient, and it can be used to train models on a variety of tasks.\n\n## What is DeepSpeed", - "throughput": 23.7931001677926 + "output": "DeepSpeed is a machine learning framework that enables training of large models on a single machine with multiple GPUs. It is designed to be easy to use and efficient, and it supports a wide range of models and tasks.\n\nDeepSpeed is a deep learning framework that enables training of large models on a single machine with multiple GPUs. It is designed to be easy to use and efficient, and it supports a wide range of models and tasks.\n\nDeepSpeed is a deep learning framework that enables training of large models on a", + "throughput": 81.6817273229847 } }, "tests/test_text_generation_example.py::test_text_generation_bf16_1x[mosaicml/mpt-30b-1-False-False]": { @@ -346,7 +346,7 @@ "throughput": 36.06464336116623 }, "gaudi3": { - "throughput": 36.06464336116623 + "throughput": 42.05243284402848 } }, "tests/test_text_generation_example.py::test_text_generation_bf16_1x[mosaicml/mpt-7b-1-False-False]": { @@ -359,7 +359,7 @@ "throughput": 65.116 }, "gaudi3": { - "throughput": 65.116 + "throughput": 67.06139602530865 } }, "tests/test_text_generation_example.py::test_text_generation_bf16_1x[stabilityai/stablelm-2-12b-1-False-False]": { @@ -375,10 +375,10 @@ }, "tests/test_text_generation_example.py::test_text_generation_bf16_1x[state-spaces/mamba-130m-hf-1536-False-False]": { "gaudi2": { - "throughput": 5385.511100161605 + "throughput": 3100.9825044466907 }, "gaudi3": { - "throughput": 5385.511100161605 + "throughput": 1948.1615848330302 } }, "tests/test_text_generation_example.py::test_text_generation_bf16_1x[state-spaces/mamba-130m-hf-224-False-False]": { @@ -391,7 +391,7 @@ "throughput": 25.202450111088346 }, "gaudi3": { - "throughput": 25.202450111088346 + "throughput": 34.03571811480758 } }, "tests/test_text_generation_example.py::test_text_generation_bf16_1x[tiiuae/falcon-7b-1-True-False]": { @@ -404,7 +404,7 @@ "throughput": 47.1464839567739 }, "gaudi3": { - "throughput": 47.1464839567739 + "throughput": 45.90538768350833 } }, "tests/test_text_generation_example.py::test_text_generation_contrastive_search[gpt2-xl-1-False]": { @@ -415,7 +415,7 @@ "throughput": 51.61471298016438 }, "gaudi3": { - "throughput": 51.61471298016438 + "throughput": 69.74689153288725 } }, "tests/test_text_generation_example.py::test_text_generation_deepspeed[Qwen/Qwen2.5-72B-2-1]": { @@ -423,7 +423,7 @@ "throughput": 26 }, "gaudi3": { - "throughput": 26 + "throughput": 32.54000413829271 } }, "tests/test_text_generation_example.py::test_text_generation_deepspeed[bigscience/bloomz-7b1-8-1]": { @@ -436,7 +436,7 @@ "throughput": 36.77314954096159 }, "gaudi3": { - "throughput": 36.77314954096159 + "throughput": 42.964481338739304 } }, "tests/test_text_generation_example.py::test_text_generation_deepspeed[facebook/opt-66b-2-1]": { @@ -444,7 +444,7 @@ "throughput": 28.48069266504111 }, "gaudi3": { - "throughput": 28.48069266504111 + "throughput": 36.79515723258173 } }, "tests/test_text_generation_example.py::test_text_generation_deepspeed[google/gemma-2-27b-8-1]": { @@ -452,7 +452,7 @@ "throughput": 87.578709544111 }, "gaudi3": { - "throughput": 87.578709544111 + "throughput": 107.59395201764178 } }, "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": 110.12610917383735 + "throughput": 123.69992293361813 } }, "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": 64 + "throughput": 75.6224035651044 } }, "tests/test_text_generation_example.py::test_text_generation_distributed_tp[meta-llama/Llama-2-7b-hf]": { @@ -484,7 +484,7 @@ "throughput": 1345.2369318328463 }, "gaudi3": { - "throughput": 1345.2369318328463 + "throughput": 4660.026752215663 } }, "tests/test_text_generation_example.py::test_text_generation_fp8[meta-llama/Llama-2-70b-hf-4-207-False-2048-128]": { @@ -492,7 +492,7 @@ "throughput": 568.5 }, "gaudi3": { - "throughput": 568.5 + "throughput": 918.3333993444961 } }, "tests/test_text_generation_example.py::test_text_generation_fp8[meta-llama/Llama-2-70b-hf-4-3042-False-128-128]": { @@ -500,7 +500,7 @@ "throughput": 5374.6 }, "gaudi3": { - "throughput": 5374.6 + "throughput": 9105.741034094377 } }, "tests/test_text_generation_example.py::test_text_generation_fp8[meta-llama/Llama-2-70b-hf-4-750-False-128-2048]": { @@ -508,7 +508,7 @@ "throughput": 7422.4 }, "gaudi3": { - "throughput": 7422.4 + "throughput": 12966.32808044709 } }, "tests/test_text_generation_example.py::test_text_generation_fp8[meta-llama/Llama-2-70b-hf-8-172-False-2048-2048]": { @@ -516,7 +516,7 @@ "throughput": 4656.2 }, "gaudi3": { - "throughput": 4656.2 + "throughput": 6968.716105590979 } }, "tests/test_text_generation_example.py::test_text_generation_fp8[meta-llama/Llama-2-7b-hf-1-1230-False-128-128]": { @@ -524,7 +524,7 @@ "throughput": 13152.7 }, "gaudi3": { - "throughput": 13152.7 + "throughput": 19132.3193582529 } }, "tests/test_text_generation_example.py::test_text_generation_fp8[meta-llama/Llama-2-7b-hf-1-163-False-128-2048]": { @@ -532,7 +532,7 @@ "throughput": 4774.7 }, "gaudi3": { - "throughput": 4774.7 + "throughput": 7240.988993899055 } }, "tests/test_text_generation_example.py::test_text_generation_fp8[meta-llama/Llama-2-7b-hf-1-81-False-2048-2048]": { @@ -540,7 +540,7 @@ "throughput": 1942.9 }, "gaudi3": { - "throughput": 1942.9 + "throughput": 2868.2782272085133 } }, "tests/test_text_generation_example.py::test_text_generation_fp8[meta-llama/Llama-2-7b-hf-1-94-False-2048-128]": { @@ -548,7 +548,7 @@ "throughput": 1293.3 }, "gaudi3": { - "throughput": 1293.3 + "throughput": 1852.6696711170073 } }, "tests/test_text_generation_example.py::test_text_generation_fp8[microsoft/phi-2-1-1-True-128-128]": { @@ -556,7 +556,7 @@ "throughput": 254.08932787178165 }, "gaudi3": { - "throughput": 254.08932787178165 + "throughput": 298.62002948546194 } }, "tests/test_text_generation_example.py::test_text_generation_fp8[mistralai/Mistral-7B-Instruct-v0.2-1-120-True-128-2048]": { @@ -580,7 +580,7 @@ "throughput": 3393.149396451692 }, "gaudi3": { - "throughput": 3393.149396451692 + "throughput": 4877.759076826148 } }, "tests/test_text_generation_example.py::test_text_generation_fp8[mistralai/Mistral-7B-Instruct-v0.2-1-896-True-128-128]": { @@ -588,7 +588,7 @@ "throughput": 17068.965283763682 }, "gaudi3": { - "throughput": 17068.965283763682 + "throughput": 25100.757003294264 } }, "tests/test_text_generation_example.py::test_text_generation_fp8[mistralai/Mixtral-8x7B-v0.1-1-1-True-128-128]": { @@ -596,7 +596,7 @@ "throughput": 40.94 }, "gaudi3": { - "throughput": 40.94 + "throughput": 114.8447433058542 } }, "tests/test_text_generation_example.py::test_text_generation_fp8[mistralai/Mixtral-8x7B-v0.1-2-48-True-2048-2048]": { @@ -604,7 +604,7 @@ "throughput": 1147.5 }, "gaudi3": { - "throughput": 1147.5 + "throughput": 2632.4017718271375 } }, "tests/test_text_generation_example.py::test_text_generation_fp8[mistralai/Mixtral-8x7B-v0.1-2-768-True-128-128]": { @@ -636,7 +636,7 @@ "throughput": 2506.68 }, "gaudi3": { - "throughput": 2506.68 + "throughput": 3716.3864966397186 } }, "tests/test_text_generation_example.py::test_text_generation_gptq[TheBloke/Llama-2-7b-Chat-GPTQ-1-10-False-128-2048]": { @@ -644,7 +644,7 @@ "throughput": 456.7 }, "gaudi3": { - "throughput": 456.7 + "throughput": 828.9133748373866 } }, "tests/test_text_generation_example.py::test_text_generation_torch_compile[meta-llama/Llama-2-7b-hf]": { @@ -652,7 +652,7 @@ "throughput": 102.27823420713148 }, "gaudi3": { - "throughput": 102.27823420713148 + "throughput": 170.08149766812704 } }, "tests/test_text_generation_example.py::test_text_generation_torch_compile_distributed[meta-llama/Llama-2-7b-hf]": { @@ -660,7 +660,7 @@ "throughput": 39.72973199515235 }, "gaudi3": { - "throughput": 39.72973199515235 + "throughput": 182.2741046353745 } } } \ No newline at end of file diff --git a/tests/baselines/fixture/tests/test_video_llava.json b/tests/baselines/fixture/tests/test_video_llava.json index 90146af1f5..f2c67def28 100644 --- a/tests/baselines/fixture/tests/test_video_llava.json +++ b/tests/baselines/fixture/tests/test_video_llava.json @@ -7,7 +7,7 @@ "throughput": 27.72902536827787 }, "gaudi3": { - "throughput": 27.72902536827787 + "throughput": 41.32754713852968 } } } \ No newline at end of file diff --git a/tests/test_trainer.py b/tests/test_trainer.py index bca097be1f..a8d55f341a 100644 --- a/tests/test_trainer.py +++ b/tests/test_trainer.py @@ -37,6 +37,7 @@ AutoModelForCausalLM, AutoProcessor, AutoTokenizer, + DataCollatorForLanguageModeling, GPT2LMHeadModel, IntervalStrategy, LineByLineTextDataset, @@ -82,6 +83,7 @@ from optimum.habana import GaudiConfig, GaudiTrainingArguments from optimum.habana.accelerate import GaudiAccelerator, GaudiAcceleratorState +from optimum.habana.utils import set_seed from optimum.utils import logging @@ -112,6 +114,19 @@ adapt_transformers_to_gaudi() +class StoreLossCallback(TrainerCallback): + """ + Simple callback to store the loss. + """ + + def __init__(self): + self.losses = [] + + def on_log(self, args, state, control, logs=None, **kwargs): + if "loss" in logs: + self.losses.append(logs["loss"]) + + class MockOOMCallback(TrainerCallback): """ Simple callback to simulate CUDA OOM error if @@ -127,6 +142,26 @@ def on_step_end(self, args, state, control, **kwargs): raise RuntimeError("Out of memory.") +def ForCausalLMLoss(logits, labels, vocab_size, num_items_in_batch, disable_num_items_in_batch=False): + # Upcast to float if we need to compute the loss to avoid potential precision issues + logits = logits.float() + # Shift so that tokens < n predict n + shift_logits = logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + + # Flatten the tokens + shift_logits = shift_logits.view(-1, vocab_size) + shift_labels = shift_labels.view(-1) + # Enable model parallelism + shift_labels = shift_labels.to(shift_logits.device) + if num_items_in_batch is None or disable_num_items_in_batch: + loss = nn.functional.cross_entropy(shift_logits, shift_labels, ignore_index=-100, reduction="mean") + else: + loss = nn.functional.cross_entropy(shift_logits, shift_labels, ignore_index=-100, reduction="sum") + loss = loss / num_items_in_batch + return loss + + class RegressionDataset: def __init__(self, a=2, b=3, length=64, seed=42, label_names=None): np.random.seed(seed) @@ -470,14 +505,44 @@ def get_regression_trainer( preprocess_logits_for_metrics=preprocess_logits_for_metrics, ) + def get_language_model_trainer(**kwargs): + import datasets + + dataset = datasets.load_dataset("fka/awesome-chatgpt-prompts") + model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2") + tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2") + tokenizer.pad_token = tokenizer.eos_token + + def _tokenize_function(examples): + model_inputs = tokenizer(examples["prompt"], padding="max_length", truncation=True) + model_inputs["labels"] = np.array(model_inputs["input_ids"]).astype(np.int64) + return model_inputs + + tokenized_datasets = dataset.map(_tokenize_function, batched=True) + training_args = GaudiTrainingArguments(use_habana=True, use_lazy_mode=True, **kwargs) + gaudi_config = get_gaudi_config() + + trainer = GaudiTrainer( + model=model, + gaudi_config=gaudi_config, + args=training_args, + train_dataset=tokenized_datasets["train"], + ) + + return trainer + class GaudiTrainerIntegrationCommon: - def check_saved_checkpoints(self, output_dir, freq, total, is_pretrained=True, safe_weights=True): + def check_saved_checkpoints( + self, output_dir, freq, total, is_pretrained=True, safe_weights=True, use_scaler=False + ): weights_file = WEIGHTS_NAME if not safe_weights else SAFE_WEIGHTS_NAME file_list = [weights_file, "training_args.bin", "optimizer.pt", "scheduler.pt", "trainer_state.json"] if is_pretrained: file_list.append("config.json") file_list.append("gaudi_config.json") + if use_scaler: + file_list.append("scaler.pt") for step in range(freq, total, freq): checkpoint = os.path.join(output_dir, f"checkpoint-{step}") self.assertTrue(os.path.isdir(checkpoint)) @@ -505,8 +570,8 @@ def check_best_model_has_been_loaded( state_dict = safetensors.torch.load_file(os.path.join(checkpoint, SAFE_WEIGHTS_NAME)) best_model.load_state_dict(state_dict) best_model.to(trainer.args.device) - self.assertTrue(torch.allclose(best_model.a, trainer.model.a)) - self.assertTrue(torch.allclose(best_model.b, trainer.model.b)) + torch.testing.assert_close(best_model.a, trainer.model.a) + torch.testing.assert_close(best_model.b, trainer.model.b) metrics = trainer.evaluate() self.assertEqual(metrics[metric], best_value) @@ -594,8 +659,8 @@ def check_trained_model(self, model, alternate_seed=False, bf16=False): # 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: - self.assertTrue(torch.allclose(model.a, a)) - self.assertTrue(torch.allclose(model.b, b)) + torch.testing.assert_close(model.a, a) + torch.testing.assert_close(model.b, b) 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)) @@ -669,6 +734,226 @@ def test_model_init(self): trainer.train() self.check_trained_model(trainer.model, alternate_seed=True) + def test_gradient_accumulation_loss_alignment_with_model_loss(self): + set_seed(42) + import datasets + + model_name = "nickypro/tinyllama-15M" + dataset_name = "wikitext" + dataset_config = "wikitext-2-raw-v1" + dataset = datasets.load_dataset(dataset_name, dataset_config, split="train[:40]") + tokenizer = AutoTokenizer.from_pretrained(model_name) + + tokenizer.pad_token = tokenizer.eos_token + + def tokenize_function(examples): + return tokenizer(examples["text"], max_length=16, padding="max_length", truncation=True) + + tokenized_dataset = dataset.map(tokenize_function, batched=True, remove_columns=dataset.column_names) + + 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, + "max_steps": 5, + "learning_rate": 3e-4, + "disable_tqdm": True, + } + + with tempfile.TemporaryDirectory() as tmp_dir: + args = GaudiTrainingArguments( + tmp_dir, + use_habana=True, + use_lazy_mode=True, + **args_kwargs, + ) + gaudi_config = get_gaudi_config() + trainer = GaudiTrainer( + model, + gaudi_config, + args, + train_dataset=tokenized_dataset, + callbacks=[base_loss_callback], + data_collator=data_collator, + ) + 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, + gradient_accumulation_steps=2, + per_device_train_batch_size=4, + use_habana=True, + use_lazy_mode=True, + ) + set_seed(42) + model = AutoModelForCausalLM.from_pretrained(model_name) + trainer = GaudiTrainer( + model, + gaudi_config, + args, + train_dataset=tokenized_dataset, + callbacks=[grad_accum_loss_callback], + data_collator=data_collator, + ) + trainer.train() + + set_seed(42) + model.load_state_dict(state_dict) + broken_loss_callback = StoreLossCallback() + trainer = GaudiTrainer( + model, + gaudi_config, + args, + train_dataset=tokenized_dataset, + callbacks=[broken_loss_callback], + data_collator=data_collator, + ) + # disable model_accepts_loss_kwargs + 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") + + # 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.2, f"Difference {max(diff_broken)} is not greater than 1.2") + + 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") + + def test_gradient_accumulation_loss_alignment_with_loss_func(self): + set_seed(42) + import datasets + + model_name = "roneneldan/TinyStories-33M" + dataset_name = "wikitext" + dataset_config = "wikitext-2-raw-v1" + dataset = datasets.load_dataset(dataset_name, dataset_config, split="train[:40]") + tokenizer = AutoTokenizer.from_pretrained(model_name) + + tokenizer.pad_token = tokenizer.eos_token + + def tokenize_function(examples): + return tokenizer(examples["text"], max_length=16, padding="max_length", truncation=True) + + tokenized_dataset = dataset.map(tokenize_function, batched=True) + + tokenizer.pad_token = tokenizer.eos_token + data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False) + + model = AutoModelForCausalLM.from_pretrained(model_name) + + def compute_loss(logits, labels, vocab_size, num_items_in_batch, disable_num_items_in_batch=False): + return ForCausalLMLoss( + logits["logits"], labels, vocab_size, num_items_in_batch, disable_num_items_in_batch + ) + + loss_fn = partial(compute_loss, vocab_size=model.config.vocab_size, disable_num_items_in_batch=False) + + base_loss_callback = StoreLossCallback() + + args_kwargs = { + "report_to": "none", + "logging_steps": 1, + "max_steps": 5, + "learning_rate": 3e-4, + "disable_tqdm": True, + } + + with tempfile.TemporaryDirectory() as tmp_dir: + args = GaudiTrainingArguments( + tmp_dir, + use_habana=True, + use_lazy_mode=True, + **args_kwargs, + ) + gaudi_config = get_gaudi_config() + trainer = GaudiTrainer( + model, + gaudi_config, + args, + train_dataset=tokenized_dataset, + callbacks=[base_loss_callback], + compute_loss_func=loss_fn, + data_collator=data_collator, + ) + trainer.train() + + grad_accum_loss_callback = StoreLossCallback() + with tempfile.TemporaryDirectory() as tmp_dir: + args = GaudiTrainingArguments( + tmp_dir, + **args_kwargs, + gradient_accumulation_steps=2, + per_device_train_batch_size=4, + use_habana=True, + use_lazy_mode=True, + ) + set_seed(42) + model = AutoModelForCausalLM.from_pretrained(model_name) + trainer = GaudiTrainer( + model, + gaudi_config, + args, + train_dataset=tokenized_dataset, + callbacks=[grad_accum_loss_callback], + compute_loss_func=loss_fn, + data_collator=data_collator, + ) + trainer.train() + + set_seed(42) + model = AutoModelForCausalLM.from_pretrained(model_name) + broken_loss_callback = StoreLossCallback() + loss_fn = partial(compute_loss, vocab_size=model.config.vocab_size, disable_num_items_in_batch=True) + trainer = GaudiTrainer( + model, + gaudi_config, + args, + train_dataset=tokenized_dataset, + callbacks=[broken_loss_callback], + compute_loss_func=loss_fn, + data_collator=data_collator, + ) + 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") + + # max diff broken should be very off + self.assertGreater(max(diff_broken), 3, f"Difference {max(diff_broken)} is not greater than 3") + def test_gradient_accumulation(self): with tempfile.TemporaryDirectory() as tmpdir: # Training with half the batch size but accumulation steps as 2 should give the same training losses. @@ -969,57 +1254,79 @@ def test_trainer_works_with_dict(self): eval_dataset = RegressionDataset() model = RegressionDictModel() gaudi_config = get_gaudi_config() - with tempfile.TemporaryDirectory() as tmpdir: - args = GaudiTrainingArguments(tmpdir, use_habana=True, use_lazy_mode=True, report_to="none") - trainer = GaudiTrainer(model, gaudi_config, args, train_dataset=train_dataset, eval_dataset=eval_dataset) - trainer.train() - _ = trainer.evaluate() - _ = trainer.predict(eval_dataset) + args = GaudiTrainingArguments( + self.get_auto_remove_tmp_dir(), use_habana=True, use_lazy_mode=True, report_to="none" + ) + trainer = GaudiTrainer(model, gaudi_config, args, train_dataset=train_dataset, eval_dataset=eval_dataset) + trainer.train() + _ = trainer.evaluate() + _ = trainer.predict(eval_dataset) def test_evaluation_with_keys_to_drop(self): - with tempfile.TemporaryDirectory() as tmpdir: - config = GPT2Config(vocab_size=100, n_positions=128, n_embd=32, n_layer=3, n_head=4) - tiny_gpt2 = GaudiGPT2LMHeadModel(config) - x = torch.randint(0, 100, (128,)) - eval_dataset = RepeatDataset(x) - args = GaudiTrainingArguments(tmpdir, use_habana=True, use_lazy_mode=True, report_to="none") - gaudi_config = get_gaudi_config() - trainer = GaudiTrainer(tiny_gpt2, gaudi_config, args, eval_dataset=eval_dataset) - # By default the past_key_values are removed - result = trainer.predict(eval_dataset) - self.assertTrue(isinstance(result.predictions, np.ndarray)) - # We can still get them by setting ignore_keys to [] - result = trainer.predict(eval_dataset, ignore_keys=[]) - self.assertTrue(isinstance(result.predictions, tuple)) - self.assertEqual(len(result.predictions), 2) + config = GPT2Config(vocab_size=100, n_positions=128, n_embd=32, n_layer=3, n_head=4) + tiny_gpt2 = GaudiGPT2LMHeadModel(config) + x = torch.randint(0, 100, (128,)) + eval_dataset = RepeatDataset(x) + args = GaudiTrainingArguments( + self.get_auto_remove_tmp_dir(), use_habana=True, use_lazy_mode=True, report_to="none" + ) + gaudi_config = get_gaudi_config() + trainer = GaudiTrainer(tiny_gpt2, gaudi_config, args, eval_dataset=eval_dataset) + # By default the past_key_values are removed + result = trainer.predict(eval_dataset) + self.assertTrue(isinstance(result.predictions, np.ndarray)) + # We can still get them by setting ignore_keys to [] + result = trainer.predict(eval_dataset, ignore_keys=[]) + self.assertTrue(isinstance(result.predictions, tuple)) + self.assertEqual(len(result.predictions), 2) def test_training_arguments_are_left_untouched(self): - with tempfile.TemporaryDirectory() as tmpdir: - trainer = get_regression_trainer(output_dir=tmpdir) - trainer.train() - args = GaudiTrainingArguments(tmpdir, use_habana=True, use_lazy_mode=True, report_to=[]) - dict1, dict2 = args.to_dict(), trainer.args.to_dict() - for key in dict1.keys(): - # Logging dir can be slightly different as they default to something with the time. - if key != "logging_dir": - self.assertEqual(dict1[key], dict2[key]) + tmp_dir = self.get_auto_remove_tmp_dir() + trainer = get_regression_trainer(output_dir=tmp_dir) + 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(): + # Logging dir can be slightly different as they default to something with the time. + if key != "logging_dir": + self.assertEqual(dict1[key], dict2[key]) def test_number_of_steps_in_training(self): - with tempfile.TemporaryDirectory() as tmpdir: - # Regular training has n_epochs * len(train_dl) steps - trainer = get_regression_trainer(output_dir=tmpdir, learning_rate=0.1) - train_output = trainer.train() - self.assertEqual(train_output.global_step, self.n_epochs * 64 / self.batch_size) + # Regular training has n_epochs * len(train_dl) steps + tmp_dir = self.get_auto_remove_tmp_dir() + trainer = get_regression_trainer(output_dir=tmp_dir, learning_rate=0.1) + train_output = trainer.train() + self.assertEqual(train_output.global_step, self.n_epochs * 64 / self.batch_size) - # Check passing num_train_epochs works (and a float version too): - trainer = get_regression_trainer(output_dir=tmpdir, learning_rate=0.1, num_train_epochs=1.5) - train_output = trainer.train() - self.assertEqual(train_output.global_step, int(1.5 * 64 / self.batch_size)) + # Check passing num_train_epochs works (and a float version too): + trainer = get_regression_trainer(output_dir=tmp_dir, learning_rate=0.1, num_train_epochs=1.5) + train_output = trainer.train() + self.assertEqual(train_output.global_step, int(1.5 * 64 / self.batch_size)) - # If we pass a max_steps, num_train_epochs is ignored - trainer = get_regression_trainer(output_dir=tmpdir, learning_rate=0.1, max_steps=10) - train_output = trainer.train() - self.assertEqual(train_output.global_step, 10) + # If we pass a max_steps, num_train_epochs is ignored + trainer = get_regression_trainer(output_dir=tmp_dir, learning_rate=0.1, max_steps=10) + train_output = trainer.train() + self.assertEqual(train_output.global_step, 10) + + # TODO: enable this test when torch.compile becomes the default on Gaudi + # def test_torch_compile_loss_func_compatibility(self): + # config = LlamaConfig(vocab_size=100, hidden_size=32, num_hidden_layers=3, num_attention_heads=4) + # tiny_llama = LlamaForCausalLM(config) + + # x = torch.randint(0, 100, (128,)) + # train_dataset = RepeatDataset(x) + + # args = GaudiTrainingArguments( + # self.get_auto_remove_tmp_dir(), + # per_device_train_batch_size=2, + # torch_compile=True, + # max_steps=1, # compile happens on the first step + # use_habana=True, + # use_lazy_mode=True, + # ) + # gaudi_config = get_gaudi_config() + # trainer = GaudiTrainer(model=tiny_llama, gaudi_config=gaudi_config, args=args, train_dataset=train_dataset) # noqa + # trainer.train() @require_peft def test_multiple_peft_adapters(self): @@ -1051,38 +1358,34 @@ def test_multiple_peft_adapters(self): tokenizer.pad_token = tokenizer.eos_token - with tempfile.TemporaryDirectory() as tmpdir: - args = GaudiTrainingArguments( - tmpdir, - per_device_train_batch_size=1, - learning_rate=1e-9, - save_steps=5, - logging_steps=5, - max_steps=10, - use_habana=True, - use_lazy_mode=True, - ) - gaudi_config = get_gaudi_config() - trainer = GaudiTrainer( - tiny_model, gaudi_config, args, processing_class=tokenizer, train_dataset=train_dataset - ) + tmp_dir = self.get_auto_remove_tmp_dir() + args = GaudiTrainingArguments( + tmp_dir, + per_device_train_batch_size=1, + learning_rate=1e-9, + save_steps=5, + logging_steps=5, + max_steps=10, + use_habana=True, + use_lazy_mode=True, + ) + gaudi_config = get_gaudi_config() + trainer = GaudiTrainer(tiny_model, gaudi_config, args, processing_class=tokenizer, train_dataset=train_dataset) - trainer.train() - parameters = dict(tiny_model.named_parameters()) - state = dataclasses.asdict(trainer.state) + trainer.train() + parameters = dict(tiny_model.named_parameters()) + state = dataclasses.asdict(trainer.state) - # Reinitialize trainer - trainer = GaudiTrainer( - tiny_model, gaudi_config, args, processing_class=tokenizer, train_dataset=train_dataset - ) + # Reinitialize trainer + trainer = GaudiTrainer(tiny_model, gaudi_config, args, processing_class=tokenizer, train_dataset=train_dataset) - checkpoint = os.path.join(tmpdir, "checkpoint-5") + checkpoint = os.path.join(tmp_dir, "checkpoint-5") - trainer.train(resume_from_checkpoint=checkpoint) - parameters1 = dict(tiny_model.named_parameters()) - state1 = dataclasses.asdict(trainer.state) - self.assertEqual(parameters, parameters1) - self.check_trainer_state_are_the_same(state, state1) + trainer.train(resume_from_checkpoint=checkpoint) + parameters1 = dict(tiny_model.named_parameters()) + state1 = dataclasses.asdict(trainer.state) + self.assertEqual(parameters, parameters1) + self.check_trainer_state_are_the_same(state, state1) # TODO: investigate why this test fails # def test_neftune(self): @@ -1136,201 +1439,195 @@ def test_logging_inf_nan_filter(self): x = torch.randint(0, 100, (128,)) train_dataset = RepeatDataset(x) - with tempfile.TemporaryDirectory() as tmpdir: - # GaudiTrainer without inf/nan filter - gaudi_config = get_gaudi_config() - args = GaudiTrainingArguments( - tmpdir, - learning_rate=1e9, - logging_steps=5, - logging_nan_inf_filter=False, - use_habana=True, - use_lazy_mode=True, - report_to="none", - ) - trainer = GaudiTrainer(tiny_gpt2, gaudi_config, args, train_dataset=train_dataset) - trainer.train() - log_history_no_filter = trainer.state.log_history - - # GaudiTrainer with inf/nan filter - args = GaudiTrainingArguments( - tmpdir, - learning_rate=1e9, - logging_steps=5, - logging_nan_inf_filter=True, - use_habana=True, - use_lazy_mode=True, - report_to="none", - ) - trainer = GaudiTrainer(tiny_gpt2, gaudi_config, args, train_dataset=train_dataset) - trainer.train() - log_history_filter = trainer.state.log_history + # GaudiTrainer without inf/nan filter + gaudi_config = get_gaudi_config() + args = GaudiTrainingArguments( + self.get_auto_remove_tmp_dir(), + learning_rate=1e9, + logging_steps=5, + logging_nan_inf_filter=False, + use_habana=True, + use_lazy_mode=True, + report_to="none", + ) + trainer = GaudiTrainer(tiny_gpt2, gaudi_config, args, train_dataset=train_dataset) + trainer.train() + log_history_no_filter = trainer.state.log_history + + # GaudiTrainer with inf/nan filter + args = GaudiTrainingArguments( + self.get_auto_remove_tmp_dir(), + learning_rate=1e9, + logging_steps=5, + logging_nan_inf_filter=True, + use_habana=True, + use_lazy_mode=True, + report_to="none", + ) + trainer = GaudiTrainer(tiny_gpt2, gaudi_config, args, train_dataset=train_dataset) + trainer.train() + log_history_filter = trainer.state.log_history - def is_any_loss_nan_or_inf(log_history): - losses = [l["loss"] for l in log_history[:-1]] - return any(math.isnan(x) for x in losses) or any(math.isinf(x) for x in losses) + def is_any_loss_nan_or_inf(log_history): + losses = [l["loss"] for l in log_history[:-1]] + return any(math.isnan(x) for x in losses) or any(math.isinf(x) for x in losses) - self.assertTrue(is_any_loss_nan_or_inf(log_history_no_filter)) - self.assertFalse(is_any_loss_nan_or_inf(log_history_filter)) + self.assertTrue(is_any_loss_nan_or_inf(log_history_no_filter)) + self.assertFalse(is_any_loss_nan_or_inf(log_history_filter)) def test_train_and_eval_dataloaders(self): - with tempfile.TemporaryDirectory() as tmpdir: - trainer = get_regression_trainer(output_dir=tmpdir, learning_rate=0.1, per_device_train_batch_size=16) - self.assertEqual(trainer.get_train_dataloader().total_batch_size, 16) - trainer = get_regression_trainer(output_dir=tmpdir, learning_rate=0.1, per_device_eval_batch_size=16) - self.assertEqual(trainer.get_eval_dataloader().total_batch_size, 16) - - # Check drop_last works - trainer = get_regression_trainer( - output_dir=tmpdir, - train_len=66, - eval_len=74, - learning_rate=0.1, - per_device_train_batch_size=16, - per_device_eval_batch_size=32, - ) - self.assertEqual(len(trainer.get_train_dataloader()), 66 // (16) + 1) - self.assertEqual(len(trainer.get_eval_dataloader()), 74 // (32) + 1) - - trainer = get_regression_trainer( - output_dir=tmpdir, - train_len=66, - eval_len=74, - learning_rate=0.1, - per_device_train_batch_size=16, - per_device_eval_batch_size=32, - dataloader_drop_last=True, - ) - self.assertEqual(len(trainer.get_train_dataloader()), 66 // (16)) - self.assertEqual(len(trainer.get_eval_dataloader()), 74 // (32)) + tmp_dir = self.get_auto_remove_tmp_dir() + trainer = get_regression_trainer(output_dir=tmp_dir, learning_rate=0.1, per_device_train_batch_size=16) + self.assertEqual(trainer.get_train_dataloader().total_batch_size, 16) + trainer = get_regression_trainer(output_dir=tmp_dir, learning_rate=0.1, per_device_eval_batch_size=16) + self.assertEqual(trainer.get_eval_dataloader().total_batch_size, 16) + + # Check drop_last works + trainer = get_regression_trainer( + output_dir=tmp_dir, + train_len=66, + eval_len=74, + learning_rate=0.1, + per_device_train_batch_size=16, + per_device_eval_batch_size=32, + ) + self.assertEqual(len(trainer.get_train_dataloader()), 66 // (16) + 1) + self.assertEqual(len(trainer.get_eval_dataloader()), 74 // (32) + 1) + + trainer = get_regression_trainer( + output_dir=tmp_dir, + train_len=66, + eval_len=74, + learning_rate=0.1, + per_device_train_batch_size=16, + per_device_eval_batch_size=32, + dataloader_drop_last=True, + ) + self.assertEqual(len(trainer.get_train_dataloader()), 66 // (16)) + self.assertEqual(len(trainer.get_eval_dataloader()), 74 // (32)) - # Check passing a new dataset for evaluation works - new_eval_dataset = RegressionDataset(length=128) - self.assertEqual(len(trainer.get_eval_dataloader(new_eval_dataset)), 128 // (32)) + # Check passing a new dataset for evaluation works + new_eval_dataset = RegressionDataset(length=128) + self.assertEqual(len(trainer.get_eval_dataloader(new_eval_dataset)), 128 // (32)) # tests that we do not require dataloader to have a .dataset attribute def test_dataloader_without_dataset(self): train_dataset = RegressionDataset(length=128) - with tempfile.TemporaryDirectory() as tmp_dir: - args = GaudiTrainingArguments(output_dir=tmp_dir, use_habana=True, use_lazy_mode=True, report_to="none") - trainer = CustomDataloaderTrainer( - model=RegressionModel(), - gaudi_config=get_gaudi_config(), - args=args, - train_dataset=train_dataset, - eval_dataset=train_dataset, - ) - trainer.train() - trainer.evaluate() + args = GaudiTrainingArguments( + output_dir=self.get_auto_remove_tmp_dir(), use_habana=True, use_lazy_mode=True, report_to="none" + ) + trainer = CustomDataloaderTrainer( + model=RegressionModel(), + gaudi_config=get_gaudi_config(), + args=args, + train_dataset=train_dataset, + eval_dataset=train_dataset, + ) + trainer.train() + trainer.evaluate() def test_get_eval_dataloader_without_persistent_workers(self): - with tempfile.TemporaryDirectory() as tmpdir: - train_dataset = RegressionDataset() - config = GPT2Config(vocab_size=100, n_positions=128, n_embd=32, n_layer=3, n_head=4) - tiny_gpt2 = GPT2LMHeadModel(config) - args = GaudiTrainingArguments( - tmpdir, - report_to="none", - dataloader_persistent_workers=False, - use_habana=True, - use_lazy_mode=True, - ) - - # Single evaluation dataset - eval_dataset = RegressionDataset() - gaudi_config = get_gaudi_config() - trainer = GaudiTrainer( - tiny_gpt2, gaudi_config, args, train_dataset=train_dataset, eval_dataset=eval_dataset - ) - # Mocking the prepare method to avoid the dataloader changing with each call to get_eval_dataloader - trainer.accelerator.prepare = lambda x: x - - default_dataloader = trainer.get_eval_dataloader() - dataloader_with_dataset = trainer.get_eval_dataloader(eval_dataset) - - self.assertEqual(default_dataloader.dataset, eval_dataset) - self.assertEqual(dataloader_with_dataset.dataset, eval_dataset) - self.assertNotEqual(default_dataloader, dataloader_with_dataset) + train_dataset = RegressionDataset() + config = GPT2Config(vocab_size=100, n_positions=128, n_embd=32, n_layer=3, n_head=4) + tiny_gpt2 = GPT2LMHeadModel(config) + args = GaudiTrainingArguments( + self.get_auto_remove_tmp_dir(), + report_to="none", + dataloader_persistent_workers=False, + use_habana=True, + use_lazy_mode=True, + ) - # Multiple evaluation datasets - first_dataset = RegressionDataset() - second_dataset = RegressionDataset() - trainer = GaudiTrainer( - tiny_gpt2, - gaudi_config, - args, - train_dataset=train_dataset, - eval_dataset={"first": first_dataset, "second": second_dataset}, - ) - # Mocking the prepare method to avoid the dataloader changing with each call to get_eval_dataloader - trainer.accelerator.prepare = lambda x: x + # Single evaluation dataset + eval_dataset = RegressionDataset() + gaudi_config = get_gaudi_config() + trainer = GaudiTrainer(tiny_gpt2, gaudi_config, args, train_dataset=train_dataset, eval_dataset=eval_dataset) + # Mocking the prepare method to avoid the dataloader changing with each call to get_eval_dataloader + trainer.accelerator.prepare = lambda x: x + + default_dataloader = trainer.get_eval_dataloader() + dataloader_with_dataset = trainer.get_eval_dataloader(eval_dataset) + + self.assertEqual(default_dataloader.dataset, eval_dataset) + self.assertEqual(dataloader_with_dataset.dataset, eval_dataset) + self.assertNotEqual(default_dataloader, dataloader_with_dataset) + + # Multiple evaluation datasets + first_dataset = RegressionDataset() + second_dataset = RegressionDataset() + trainer = GaudiTrainer( + tiny_gpt2, + gaudi_config, + args, + train_dataset=train_dataset, + eval_dataset={"first": first_dataset, "second": second_dataset}, + ) + # Mocking the prepare method to avoid the dataloader changing with each call to get_eval_dataloader + trainer.accelerator.prepare = lambda x: x - first_dataloader = trainer.get_eval_dataloader("first") - first_dataloader_repeated = trainer.get_eval_dataloader("first") - second_dataloader = trainer.get_eval_dataloader("second") - second_dataloader_repeated = trainer.get_eval_dataloader("second") + first_dataloader = trainer.get_eval_dataloader("first") + first_dataloader_repeated = trainer.get_eval_dataloader("first") + second_dataloader = trainer.get_eval_dataloader("second") + second_dataloader_repeated = trainer.get_eval_dataloader("second") - self.assertEqual(first_dataset, first_dataloader.dataset) - self.assertEqual(first_dataloader.dataset, first_dataloader_repeated.dataset) - self.assertEqual(second_dataset, second_dataloader.dataset) - self.assertEqual(second_dataloader.dataset, second_dataloader_repeated.dataset) - self.assertNotEqual(first_dataloader, first_dataloader_repeated) - self.assertNotEqual(second_dataloader, second_dataloader_repeated) + self.assertEqual(first_dataset, first_dataloader.dataset) + self.assertEqual(first_dataloader.dataset, first_dataloader_repeated.dataset) + self.assertEqual(second_dataset, second_dataloader.dataset) + self.assertEqual(second_dataloader.dataset, second_dataloader_repeated.dataset) + self.assertNotEqual(first_dataloader, first_dataloader_repeated) + self.assertNotEqual(second_dataloader, second_dataloader_repeated) def test_get_eval_dataloader_with_persistent_workers(self): - with tempfile.TemporaryDirectory() as tmpdir: - train_dataset = RegressionDataset() - config = GPT2Config(vocab_size=100, n_positions=128, n_embd=32, n_layer=3, n_head=4) - tiny_gpt2 = GPT2LMHeadModel(config) - args = GaudiTrainingArguments( - tmpdir, - report_to="none", - dataloader_persistent_workers=True, - dataloader_num_workers=2, - use_habana=True, - use_lazy_mode=True, - ) - - # Single evaluation dataset - eval_dataset = RegressionDataset() - gaudi_config = get_gaudi_config() - trainer = GaudiTrainer( - tiny_gpt2, gaudi_config, args, train_dataset=train_dataset, eval_dataset=eval_dataset - ) - # Mocking the prepare method to avoid the dataloader changing with each call to get_eval_dataloader - trainer.accelerator.prepare = lambda x: x - - default_dataloader = trainer.get_eval_dataloader() - dataloader_with_dataset = trainer.get_eval_dataloader(eval_dataset) - - self.assertEqual(default_dataloader.dataset, eval_dataset) - self.assertEqual(dataloader_with_dataset.dataset, eval_dataset) - self.assertEqual(default_dataloader, dataloader_with_dataset) + train_dataset = RegressionDataset() + config = GPT2Config(vocab_size=100, n_positions=128, n_embd=32, n_layer=3, n_head=4) + tiny_gpt2 = GPT2LMHeadModel(config) + args = GaudiTrainingArguments( + self.get_auto_remove_tmp_dir(), + report_to="none", + dataloader_persistent_workers=True, + dataloader_num_workers=2, + use_habana=True, + use_lazy_mode=True, + ) - # Multiple evaluation datasets - first_dataset = RegressionDataset() - second_dataset = RegressionDataset() - trainer = GaudiTrainer( - tiny_gpt2, - gaudi_config, - args, - train_dataset=train_dataset, - eval_dataset={"first": first_dataset, "second": second_dataset}, - ) - # Mocking the prepare method to avoid the dataloader changing with each call to get_eval_dataloader - trainer.accelerator.prepare = lambda x: x + # Single evaluation dataset + eval_dataset = RegressionDataset() + gaudi_config = get_gaudi_config() + trainer = GaudiTrainer(tiny_gpt2, gaudi_config, args, train_dataset=train_dataset, eval_dataset=eval_dataset) + # Mocking the prepare method to avoid the dataloader changing with each call to get_eval_dataloader + trainer.accelerator.prepare = lambda x: x + + default_dataloader = trainer.get_eval_dataloader() + dataloader_with_dataset = trainer.get_eval_dataloader(eval_dataset) + + self.assertEqual(default_dataloader.dataset, eval_dataset) + self.assertEqual(dataloader_with_dataset.dataset, eval_dataset) + self.assertEqual(default_dataloader, dataloader_with_dataset) + + # Multiple evaluation datasets + first_dataset = RegressionDataset() + second_dataset = RegressionDataset() + trainer = GaudiTrainer( + tiny_gpt2, + gaudi_config, + args, + train_dataset=train_dataset, + eval_dataset={"first": first_dataset, "second": second_dataset}, + ) + # Mocking the prepare method to avoid the dataloader changing with each call to get_eval_dataloader + trainer.accelerator.prepare = lambda x: x - first_dataloader = trainer.get_eval_dataloader("first") - first_dataloader_repeated = trainer.get_eval_dataloader("first") - second_dataloader = trainer.get_eval_dataloader("second") - second_dataloader_repeated = trainer.get_eval_dataloader("second") + first_dataloader = trainer.get_eval_dataloader("first") + first_dataloader_repeated = trainer.get_eval_dataloader("first") + second_dataloader = trainer.get_eval_dataloader("second") + second_dataloader_repeated = trainer.get_eval_dataloader("second") - self.assertEqual(first_dataset, first_dataloader.dataset) - self.assertEqual(first_dataloader.dataset, first_dataloader_repeated.dataset) - self.assertEqual(second_dataset, second_dataloader.dataset) - self.assertEqual(second_dataloader.dataset, second_dataloader_repeated.dataset) - self.assertEqual(first_dataloader, first_dataloader_repeated) - self.assertEqual(second_dataloader, second_dataloader_repeated) + self.assertEqual(first_dataset, first_dataloader.dataset) + self.assertEqual(first_dataloader.dataset, first_dataloader_repeated.dataset) + self.assertEqual(second_dataset, second_dataloader.dataset) + self.assertEqual(second_dataloader.dataset, second_dataloader_repeated.dataset) + self.assertEqual(first_dataloader, first_dataloader_repeated) + self.assertEqual(second_dataloader, second_dataloader_repeated) def test_data_is_not_parallelized_when_model_is_parallel(self): model = RegressionModel() @@ -1672,193 +1969,191 @@ def test_log_level(self): self.assertNotIn(log_info_string, cl.out) def test_save_checkpoints(self): - with tempfile.TemporaryDirectory() as tmpdir: - trainer = get_regression_trainer(output_dir=tmpdir, save_steps=5) - trainer.train() - self.check_saved_checkpoints(tmpdir, 5, int(self.n_epochs * 64 / self.batch_size)) + tmp_dir = self.get_auto_remove_tmp_dir() + trainer = get_regression_trainer(output_dir=tmp_dir, save_steps=5) + trainer.train() + self.check_saved_checkpoints(tmp_dir, 5, int(self.n_epochs * 64 / self.batch_size)) # With a regular model that is not a PreTrainedModel - with tempfile.TemporaryDirectory() as tmpdir: - trainer = get_regression_trainer(output_dir=tmpdir, save_steps=5, pretrained=False) - trainer.train() - self.check_saved_checkpoints(tmpdir, 5, int(self.n_epochs * 64 / self.batch_size), False) + tmp_dir = self.get_auto_remove_tmp_dir() + trainer = get_regression_trainer(output_dir=tmp_dir, save_steps=5, pretrained=False) + trainer.train() + self.check_saved_checkpoints(tmp_dir, 5, int(self.n_epochs * 64 / self.batch_size), False) @require_safetensors def test_safe_checkpoints(self): for save_safetensors in [True, False]: - with tempfile.TemporaryDirectory() as tmpdir: - trainer = get_regression_trainer(output_dir=tmpdir, save_steps=5, save_safetensors=save_safetensors) - trainer.train() - self.check_saved_checkpoints( - tmpdir, 5, int(self.n_epochs * 64 / self.batch_size), safe_weights=save_safetensors - ) + tmp_dir = self.get_auto_remove_tmp_dir() + trainer = get_regression_trainer(output_dir=tmp_dir, save_steps=5, save_safetensors=save_safetensors) + trainer.train() + self.check_saved_checkpoints( + tmp_dir, 5, int(self.n_epochs * 64 / self.batch_size), safe_weights=save_safetensors + ) # With a regular model that is not a PreTrainedModel - with tempfile.TemporaryDirectory() as tmpdir: - trainer = get_regression_trainer( - output_dir=tmpdir, save_steps=5, pretrained=False, save_safetensors=save_safetensors - ) - trainer.train() - self.check_saved_checkpoints( - tmpdir, 5, int(self.n_epochs * 64 / self.batch_size), False, safe_weights=save_safetensors - ) - - def test_load_best_model_with_save(self): - with tempfile.TemporaryDirectory() as tmpdir: + tmp_dir = self.get_auto_remove_tmp_dir() trainer = get_regression_trainer( - output_dir=tmpdir, - save_steps=5, - evaluation_strategy="steps", - eval_steps=5, - max_steps=9, + output_dir=tmp_dir, save_steps=5, pretrained=False, save_safetensors=save_safetensors ) trainer.train() - # Check that we have the last known step: - assert os.path.exists(os.path.join(tmpdir, f"checkpoint-{trainer.state.max_steps}")), ( - f"Could not find checkpoint-{trainer.state.max_steps}" + self.check_saved_checkpoints( + tmp_dir, 5, int(self.n_epochs * 64 / self.batch_size), False, safe_weights=save_safetensors ) - # And then check the last step - assert os.path.exists(os.path.join(tmpdir, "checkpoint-9")), "Could not find checkpoint-9" + + 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_steps=5, + max_steps=9, + ) + trainer.train() + # Check that we have the last known step: + assert os.path.exists(os.path.join(tmp_dir, f"checkpoint-{trainer.state.max_steps}")), ( + f"Could not find checkpoint-{trainer.state.max_steps}" + ) + # And then check the last step + assert os.path.exists(os.path.join(tmp_dir, "checkpoint-9")), "Could not find checkpoint-9" # Now test that using a limit works # Should result in: # - save at step 5 (but is deleted) # - save at step 10 (loaded in at the end when `load_best_model=True`) # - save at step 11 - with tempfile.TemporaryDirectory() as tmpdir: - trainer = get_regression_trainer( - output_dir=tmpdir, - save_steps=5, - evaluation_strategy="steps", - eval_steps=5, - load_best_model_at_end=True, - save_total_limit=2, - max_steps=11, - ) - trainer.train() - # Check that we have the last known step: - assert os.path.exists(os.path.join(tmpdir, "checkpoint-11")), "Could not find checkpoint-11" - # And then check the last multiple - assert os.path.exists(os.path.join(tmpdir, "checkpoint-10")), "Could not find checkpoint-10" - # Finally check that we don't have an old one - assert not os.path.exists(os.path.join(tmpdir, "checkpoint-5")), "Found checkpoint-5, limit not respected" - - # Finally check that the right model was loaded in, checkpoint-10 - # this goes by the last `eval` step check to do so, so it won't be - # the last model *saved* - model_state = trainer.model.state_dict() - final_model_weights = safetensors.torch.load_file( - os.path.join(tmpdir, "checkpoint-10", "model.safetensors") - ) - for k, v in model_state.items(): - assert torch.allclose(v, final_model_weights[k]), f"{k} is not the same" + tmp_dir = self.get_auto_remove_tmp_dir() + trainer = get_regression_trainer( + output_dir=tmp_dir, + save_steps=5, + evaluation_strategy="steps", + eval_steps=5, + load_best_model_at_end=True, + save_total_limit=2, + max_steps=11, + ) + trainer.train() + # Check that we have the last known step: + assert os.path.exists(os.path.join(tmp_dir, "checkpoint-11")), "Could not find checkpoint-11" + # And then check the last multiple + assert os.path.exists(os.path.join(tmp_dir, "checkpoint-10")), "Could not find checkpoint-10" + # Finally check that we don't have an old one + assert not os.path.exists(os.path.join(tmp_dir, "checkpoint-5")), "Found checkpoint-5, limit not respected" + + # Finally check that the right model was loaded in, checkpoint-10 + # this goes by the last `eval` step check to do so, so it won't be + # the last model *saved* + model_state = trainer.model.state_dict() + final_model_weights = safetensors.torch.load_file(os.path.join(tmp_dir, "checkpoint-10", "model.safetensors")) + for k, v in model_state.items(): + assert torch.allclose(v, final_model_weights[k]), f"{k} is not the same" def test_can_resume_training(self): - with tempfile.TemporaryDirectory() as tmpdir: - kwargs = { - "output_dir": tmpdir, - "train_len": 128, - "save_steps": 5, - "learning_rate": 0.1, - "logging_steps": 5, - } - trainer = get_regression_trainer(**kwargs) - # Disable FusedClipNorm because it makes the test fail - trainer.gaudi_config.use_fused_clip_norm = False - trainer.train() - (a, b) = trainer.model.a.item(), trainer.model.b.item() - state = dataclasses.asdict(trainer.state) + tmp_dir = self.get_auto_remove_tmp_dir() + kwargs = { + "output_dir": tmp_dir, + "train_len": 128, + "save_steps": 5, + "learning_rate": 0.1, + "logging_steps": 5, + } + trainer = get_regression_trainer(**kwargs) + # Disable FusedClipNorm because it makes the test fail + trainer.gaudi_config.use_fused_clip_norm = False + trainer.train() + (a, b) = trainer.model.a.item(), trainer.model.b.item() + state = dataclasses.asdict(trainer.state) - checkpoint = os.path.join(tmpdir, "checkpoint-5") + checkpoint = os.path.join(tmp_dir, "checkpoint-5") - # Reinitialize trainer - trainer = get_regression_trainer(**kwargs) - # Disable FusedClipNorm because it makes the test fail - trainer.gaudi_config.use_fused_clip_norm = False + # Reinitialize trainer + trainer = get_regression_trainer(**kwargs) + # Disable FusedClipNorm because it makes the test fail + trainer.gaudi_config.use_fused_clip_norm = False - trainer.train(resume_from_checkpoint=checkpoint) - (a1, b1) = trainer.model.a.item(), trainer.model.b.item() - state1 = dataclasses.asdict(trainer.state) - self.assertEqual(a, a1) - self.assertEqual(b, b1) - self.check_trainer_state_are_the_same(state, state1) + trainer.train(resume_from_checkpoint=checkpoint) + (a1, b1) = trainer.model.a.item(), trainer.model.b.item() + state1 = dataclasses.asdict(trainer.state) + self.assertEqual(a, a1) + self.assertEqual(b, b1) + self.check_trainer_state_are_the_same(state, state1) - # Now check with a later checkpoint that it also works when we span over one epoch - checkpoint = os.path.join(tmpdir, "checkpoint-15") + # Now check with a later checkpoint that it also works when we span over one epoch + checkpoint = os.path.join(tmp_dir, "checkpoint-15") - # Reinitialize trainer and load model - trainer = get_regression_trainer(**kwargs) - # Disable FusedClipNorm because it makes the test fail - trainer.gaudi_config.use_fused_clip_norm = False + # Reinitialize trainer and load model + trainer = get_regression_trainer(**kwargs) + # Disable FusedClipNorm because it makes the test fail + trainer.gaudi_config.use_fused_clip_norm = False - trainer.train(resume_from_checkpoint=checkpoint) - (a1, b1) = trainer.model.a.item(), trainer.model.b.item() - state1 = dataclasses.asdict(trainer.state) - self.assertEqual(a, a1) - self.assertEqual(b, b1) - self.check_trainer_state_are_the_same(state, state1) + trainer.train(resume_from_checkpoint=checkpoint) + (a1, b1) = trainer.model.a.item(), trainer.model.b.item() + state1 = dataclasses.asdict(trainer.state) + self.assertEqual(a, a1) + self.assertEqual(b, b1) + self.check_trainer_state_are_the_same(state, state1) # With a regular model that is not a PreTrainedModel - with tempfile.TemporaryDirectory() as tmpdir: - kwargs = { - "output_dir": tmpdir, - "train_len": 128, - "save_steps": 5, - "learning_rate": 0.1, - "pretrained": False, - } + tmp_dir = self.get_auto_remove_tmp_dir() + kwargs = { + "output_dir": tmp_dir, + "train_len": 128, + "save_steps": 5, + "learning_rate": 0.1, + "pretrained": False, + } - trainer = get_regression_trainer(**kwargs) - # Disable FusedClipNorm because it makes the test fail - trainer.gaudi_config.use_fused_clip_norm = False - trainer.train() - (a, b) = trainer.model.a.item(), trainer.model.b.item() - state = dataclasses.asdict(trainer.state) + trainer = get_regression_trainer(**kwargs) + # Disable FusedClipNorm because it makes the test fail + trainer.gaudi_config.use_fused_clip_norm = False + trainer.train() + (a, b) = trainer.model.a.item(), trainer.model.b.item() + state = dataclasses.asdict(trainer.state) - checkpoint = os.path.join(tmpdir, "checkpoint-5") + checkpoint = os.path.join(tmp_dir, "checkpoint-5") - # Reinitialize trainer and load model - trainer = get_regression_trainer(**kwargs) - # Disable FusedClipNorm because it makes the test fail - trainer.gaudi_config.use_fused_clip_norm = False + # Reinitialize trainer and load model + trainer = get_regression_trainer(**kwargs) + # Disable FusedClipNorm because it makes the test fail + trainer.gaudi_config.use_fused_clip_norm = False - trainer.train(resume_from_checkpoint=checkpoint) - (a1, b1) = trainer.model.a.item(), trainer.model.b.item() - state1 = dataclasses.asdict(trainer.state) - self.assertEqual(a, a1) - self.assertEqual(b, b1) - self.check_trainer_state_are_the_same(state, state1) + trainer.train(resume_from_checkpoint=checkpoint) + (a1, b1) = trainer.model.a.item(), trainer.model.b.item() + state1 = dataclasses.asdict(trainer.state) + self.assertEqual(a, a1) + self.assertEqual(b, b1) + self.check_trainer_state_are_the_same(state, state1) - # Now check with a later checkpoint that it also works when we span over one epoch - checkpoint = os.path.join(tmpdir, "checkpoint-15") + # Now check with a later checkpoint that it also works when we span over one epoch + checkpoint = os.path.join(tmp_dir, "checkpoint-15") - # Reinitialize trainer and load model - trainer = get_regression_trainer(**kwargs) - # Disable FusedClipNorm because it makes the test fail - trainer.gaudi_config.use_fused_clip_norm = False + # Reinitialize trainer and load model + trainer = get_regression_trainer(**kwargs) + # Disable FusedClipNorm because it makes the test fail + trainer.gaudi_config.use_fused_clip_norm = False - trainer.train(resume_from_checkpoint=checkpoint) - (a1, b1) = trainer.model.a.item(), trainer.model.b.item() - state1 = dataclasses.asdict(trainer.state) - self.assertEqual(a, a1) - self.assertEqual(b, b1) - self.check_trainer_state_are_the_same(state, state1) + trainer.train(resume_from_checkpoint=checkpoint) + (a1, b1) = trainer.model.a.item(), trainer.model.b.item() + state1 = dataclasses.asdict(trainer.state) + self.assertEqual(a, a1) + self.assertEqual(b, b1) + self.check_trainer_state_are_the_same(state, state1) # Now check failures # 1. fail to find a bogus checkpoint - with tempfile.TemporaryDirectory() as tmpdir: - trainer = get_regression_trainer(output_dir=tmpdir) - with self.assertRaises(Exception) as context: - trainer.train(resume_from_checkpoint=f"{checkpoint}-bogus") - self.assertTrue("Can't find a valid checkpoint at" in str(context.exception)) + tmp_dir = self.get_auto_remove_tmp_dir() + trainer = get_regression_trainer(output_dir=tmp_dir) + with self.assertRaises(Exception) as context: + trainer.train(resume_from_checkpoint=f"{checkpoint}-bogus") + self.assertTrue("Can't find a valid checkpoint at" in str(context.exception)) # 2. fail to find any checkpoint - due a fresh output_dir - with tempfile.TemporaryDirectory() as tmpdir: - trainer = get_regression_trainer(output_dir=tmpdir) - with self.assertRaises(Exception) as context: - trainer.train(resume_from_checkpoint=True) - self.assertTrue("No valid checkpoint found in output directory" in str(context.exception)) + tmp_dir = self.get_auto_remove_tmp_dir() + trainer = get_regression_trainer(output_dir=tmp_dir) + with self.assertRaises(Exception) as context: + trainer.train(resume_from_checkpoint=True) + self.assertTrue("No valid checkpoint found in output directory" in str(context.exception)) def test_resume_training_with_randomness(self): train_dataset = RegressionDataset(length=128) diff --git a/tests/transformers/tests/models/gpt_neox/test_modeling_gpt_neox.py b/tests/transformers/tests/models/gpt_neox/test_modeling_gpt_neox.py index 905b9474dc..aac27b4ec3 100644 --- a/tests/transformers/tests/models/gpt_neox/test_modeling_gpt_neox.py +++ b/tests/transformers/tests/models/gpt_neox/test_modeling_gpt_neox.py @@ -209,6 +209,9 @@ def create_and_check_cached_forward_with_and_without_attention_mask(self, config model.to(torch_device) model.eval() # We want this for SDPA, eager works with a `None` attention mask + # TODO: Starting v4.49, gpt_neox _attn_implementation is set to eager: https://github.com/huggingface/optimum-habana/blob/transformers_4_49/optimum/habana/transformers/models/modeling_all_models.py + # here we manually set it back to sdpa for testing + model.config._attn_implementation = "sdpa" assert model.config._attn_implementation == "sdpa", ( "This test assumes the model to have the SDPA implementation for its attention calculations." ) diff --git a/tests/transformers/tests/test_modeling_common.py b/tests/transformers/tests/test_modeling_common.py index 55c7aa8dae..e61d5b75c8 100755 --- a/tests/transformers/tests/test_modeling_common.py +++ b/tests/transformers/tests/test_modeling_common.py @@ -1637,7 +1637,7 @@ def test_load_save_without_tied_weights(self): def test_tied_weights_keys(self): config, _ = self.model_tester.prepare_config_and_inputs_for_common() - config.tie_word_embeddings = True + config.get_text_config().tie_word_embeddings = True for model_class in self.all_model_classes: model_tied = model_class(config) @@ -1651,8 +1651,8 @@ def test_tied_weights_keys(self): tied_weight_keys = model_tied._tied_weights_keys if model_tied._tied_weights_keys is not None else [] # Detect we get a hit for each key for key in tied_weight_keys: - if not any(re.search(key, p) for group in tied_params for p in group): - raise ValueError(f"{key} is not a tied weight key for {model_class}.") + is_tied_key = any(re.search(key, p) for group in tied_params for p in group) + self.assertTrue(is_tied_key, f"{key} is not a tied weight key for {model_class}.") # Removed tied weights found from tied params -> there should only be one left after for key in tied_weight_keys: From dd42c9243b2264b3c94c54b20fab28d52753666b Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Fri, 14 Mar 2025 14:19:16 +0000 Subject: [PATCH 60/89] Fix `get_num_items_in_batches` for iterable datasets and when resuming training --- optimum/habana/transformers/trainer.py | 22 ++++++++++++++++++++-- 1 file changed, 20 insertions(+), 2 deletions(-) diff --git a/optimum/habana/transformers/trainer.py b/optimum/habana/transformers/trainer.py index 1931081bee..2f10a95191 100644 --- a/optimum/habana/transformers/trainer.py +++ b/optimum/habana/transformers/trainer.py @@ -908,6 +908,7 @@ def hpu_deepspeed_checkpointing(function, *checkpoint_args, use_reentrant: Optio train_dataloader, len_dataloader, num_examples, + steps_trained_in_current_epoch, ) hb_profiler = HabanaProfile( @@ -2593,7 +2594,14 @@ def _zero_model_grad(self, model): model._zero_grad_kwargs = {} def get_num_items_in_batches( - self, args, epochs_trained, num_train_epochs, train_dataloader, len_dataloader, num_examples + self, + args, + epochs_trained, + num_train_epochs, + train_dataloader, + len_dataloader, + num_examples, + steps_trained_in_current_epoch, ): """ Calculate the number of items in each batch for all epochs during training. @@ -2609,10 +2617,15 @@ def get_num_items_in_batches( total_updates = steps_in_epoch // args.gradient_accumulation_steps + 1 if args.gradient_accumulation_steps == 1: total_updates -= 1 + global_step = 0 num_items_in_batches = [] for epoch in range(epochs_trained, num_train_epochs): - epoch_dataloader = train_dataloader + if epoch == epochs_trained and steps_trained_in_current_epoch > 0: + epoch_dataloader = skip_first_batches(train_dataloader, steps_trained_in_current_epoch) + else: + epoch_dataloader = train_dataloader + if hasattr(epoch_dataloader, "set_epoch"): epoch_dataloader.set_epoch(epoch) @@ -2652,6 +2665,11 @@ def get_num_items_in_batches( num_items_in_batch = None num_items_in_batches[epoch].append(num_items_in_batch) + global_step += 1 + + # For iterable datasets, don't do more than max_steps steps + if len_dataloader is None and global_step >= args.max_steps: + break return num_items_in_batches From 69f7e6d1aaa3851ed67bc6ac9f41198359735e1a Mon Sep 17 00:00:00 2001 From: Shifani Rajabose Date: Fri, 14 Mar 2025 14:42:30 -0400 Subject: [PATCH 61/89] Fixes pytest runtime error - Incompatible input shapes, broadcast not possible (#1796) --- optimum/habana/transformers/models/gpt2/modeling_gpt2.py | 1 + 1 file changed, 1 insertion(+) diff --git a/optimum/habana/transformers/models/gpt2/modeling_gpt2.py b/optimum/habana/transformers/models/gpt2/modeling_gpt2.py index 301f9b6633..ed2f0d0134 100644 --- a/optimum/habana/transformers/models/gpt2/modeling_gpt2.py +++ b/optimum/habana/transformers/models/gpt2/modeling_gpt2.py @@ -70,6 +70,7 @@ def _upcast_and_reordered_attn(self, query, key, value, attention_mask=None, hea attn_weights = attn_weights * head_mask attn_output = torch.matmul(attn_weights, value) + attn_output = attn_output.transpose(1, 2) return attn_output, attn_weights From d0d017257be9da3fb952a9e0808f71fae647cc82 Mon Sep 17 00:00:00 2001 From: Dmitry Date: Fri, 14 Mar 2025 19:13:40 +0100 Subject: [PATCH 62/89] Fix for AutoModelForCausalLM.from_pretrained() (#1844) --- examples/language-modeling/run_clm.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/examples/language-modeling/run_clm.py b/examples/language-modeling/run_clm.py index e7fd5d3d83..42341b6c80 100644 --- a/examples/language-modeling/run_clm.py +++ b/examples/language-modeling/run_clm.py @@ -459,7 +459,7 @@ def main(): # Note that chatglm2/3 has float16 dtype from config.json, and on Gaudi we need to use bfloat16. if config.model_type == "chatglm": - config.dtype = "torch.bfloat16" + config.torch_dtype = torch.bfloat16 tokenizer_kwargs = { "cache_dir": model_args.cache_dir, @@ -484,6 +484,11 @@ def main(): if model_args.torch_dtype in ["auto", None] else getattr(torch, model_args.torch_dtype) ) + # 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) + model = AutoModelForCausalLM.from_pretrained( model_args.model_name_or_path, from_tf=bool(".ckpt" in model_args.model_name_or_path), From adbaa23e209d84b3e2b439db204ff78f1ff683fd Mon Sep 17 00:00:00 2001 From: Mounika Mandava Date: Fri, 14 Mar 2025 11:47:14 -0700 Subject: [PATCH 63/89] Fix unexpected 'num_items_in_batch' argument in GPT-NeoX forward (#1850) Co-authored-by: regisss <15324346+regisss@users.noreply.github.com> --- optimum/habana/transformers/models/gpt_neox/modeling_gpt_neox.py | 1 + 1 file changed, 1 insertion(+) 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 30b8ee79ee..bca96fb5c9 100644 --- a/optimum/habana/transformers/models/gpt_neox/modeling_gpt_neox.py +++ b/optimum/habana/transformers/models/gpt_neox/modeling_gpt_neox.py @@ -269,6 +269,7 @@ def gaudi_gpt_neox_model_forward( return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, token_idx: Optional[torch.Tensor] = None, + **kwargs, ) -> Union[Tuple, BaseModelOutputWithPast]: """ Copied from GPTNeoxModel.forward: https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt_neox/modeling_gpt_neox.py From e802f5f3e12afad4aad5cbd1427e77bd354597ea Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Fri, 14 Mar 2025 18:52:19 +0000 Subject: [PATCH 64/89] Make style --- examples/language-modeling/run_clm.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/language-modeling/run_clm.py b/examples/language-modeling/run_clm.py index 42341b6c80..10c69ae51c 100644 --- a/examples/language-modeling/run_clm.py +++ b/examples/language-modeling/run_clm.py @@ -487,7 +487,7 @@ def main(): # 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) + torch_dtype = getattr(config, "torch_dtype", None) model = AutoModelForCausalLM.from_pretrained( model_args.model_name_or_path, From f461199d5397e6dc71492942ae305f2050ecba1b Mon Sep 17 00:00:00 2001 From: Silvia Colabrese Date: Mon, 17 Mar 2025 15:42:00 +0100 Subject: [PATCH 65/89] Fix for `GaudiLlamaAttention` object has no attribute 'max_position_embeddings' (#1854) --- optimum/habana/transformers/models/llama/modeling_llama.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/optimum/habana/transformers/models/llama/modeling_llama.py b/optimum/habana/transformers/models/llama/modeling_llama.py index 3bb0589e6b..a5f98ad644 100755 --- a/optimum/habana/transformers/models/llama/modeling_llama.py +++ b/optimum/habana/transformers/models/llama/modeling_llama.py @@ -532,11 +532,11 @@ def allocate_kv_cache(self, batch_size, max_seq_len, inp_seq_len): self.v_cache.allocate(inp_seq_len, dtype, device, cache_shape) def update_sincos_cache(self, seq_len): - # Call rotary emb forward() to update cos/sin cache when infering more than self.max_position_embeddings + # Call rotary emb forward() to update cos/sin cache when infering more than self.rotary_emb.original_max_seq_len # This helps in avoiding creation of these caches during actual model forward pass and # reduce memory consumption and improve performance. - if seq_len > self.max_position_embeddings: - self.max_position_embeddings = seq_len + if seq_len > self.rotary_emb.original_max_seq_len: + self.rotary_emb.original_max_seq_len = seq_len _, _ = self.rotary_emb(self.get_k_proj_weight(), seq_len=seq_len) def reorder(self, tensor, beam_idx, dim_a, dim_b): From 9cf57be9cff53c6972a5ba154d91b4c2856bbc28 Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Mon, 17 Mar 2025 15:02:05 +0000 Subject: [PATCH 66/89] Fix error with TRL examples --- optimum/habana/transformers/trainer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/optimum/habana/transformers/trainer.py b/optimum/habana/transformers/trainer.py index 2f10a95191..2857bfe792 100644 --- a/optimum/habana/transformers/trainer.py +++ b/optimum/habana/transformers/trainer.py @@ -1321,7 +1321,7 @@ def _maybe_log_save_evaluate(self, tr_loss, _grad_norm, model, trial, epoch, ign self._globalstep_last_logged = self.state.global_step self.store_flos() - self.log(logs, start_time) + self.log(logs, start_time=start_time) metrics = None if self.control.should_evaluate: From dbd987bcaa35cf7683b64a2d73b2abada221ff55 Mon Sep 17 00:00:00 2001 From: Harshvardhan Chauhan Date: Wed, 19 Mar 2025 01:54:25 -0700 Subject: [PATCH 67/89] Adjust precision of eval_accuracy to avoid random failure in pytest for lora finetune Llava image-to-text (#1855) --- tests/baselines/fixture/tests/test_examples.json | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/baselines/fixture/tests/test_examples.json b/tests/baselines/fixture/tests/test_examples.json index e281343d76..df8e0d7e73 100644 --- a/tests/baselines/fixture/tests/test_examples.json +++ b/tests/baselines/fixture/tests/test_examples.json @@ -492,7 +492,7 @@ }, "tests/test_examples.py::MultiCardImageToTextModelingLoRAExampleTester::test_run_image2text_lora_finetune_llava-1.5-7b-hf_multi_card": { "gaudi2": { - "eval_accuracy": 0.2122, + "eval_accuracy": 0.21, "train_runtime": 118.5782, "train_samples_per_second": 25.146 }, @@ -703,4 +703,4 @@ "train_samples_per_second": 1652.436 } } -} \ No newline at end of file +} From 78e50b951b0d3ae2596387e3a72d41cb90527da0 Mon Sep 17 00:00:00 2001 From: Harish Subramony <81822986+hsubramony@users.noreply.github.com> Date: Wed, 19 Mar 2025 01:56:12 -0700 Subject: [PATCH 68/89] Missing num_key_value_heads attribute in GaudiGemmaAttention (#1861) --- optimum/habana/transformers/models/gemma/modeling_gemma.py | 1 + 1 file changed, 1 insertion(+) diff --git a/optimum/habana/transformers/models/gemma/modeling_gemma.py b/optimum/habana/transformers/models/gemma/modeling_gemma.py index eb2ba9b89d..d0908301c9 100755 --- a/optimum/habana/transformers/models/gemma/modeling_gemma.py +++ b/optimum/habana/transformers/models/gemma/modeling_gemma.py @@ -207,6 +207,7 @@ def __init__(self, config: GemmaConfig, layer_idx: Optional[int] = None): self.v_cache = KVCache() self.inp_seq_len = -1 self.block_size = 4096 + self.num_key_value_heads = config.num_key_value_heads self.rotary_emb = GaudiRotaryEmbedding(config=self.config) self.fused_scaled_dot_product_attention = ModuleFusedSDPA(FusedSDPA) if FusedSDPA else None From bff38033a0e2ae50a188ed53a626301c2d4b7580 Mon Sep 17 00:00:00 2001 From: ZhengHongming888 Date: Wed, 19 Mar 2025 01:57:44 -0700 Subject: [PATCH 69/89] Update Sentence Transformer CI/Ref (#1862) --- .../tests/test_sentence_transformers.json | 10 +++++----- tests/test_sentence_transformers.py | 16 ++++++++++------ 2 files changed, 15 insertions(+), 11 deletions(-) diff --git a/tests/baselines/fixture/tests/test_sentence_transformers.json b/tests/baselines/fixture/tests/test_sentence_transformers.json index 36b07cd3ea..8badb310a7 100644 --- a/tests/baselines/fixture/tests/test_sentence_transformers.json +++ b/tests/baselines/fixture/tests/test_sentence_transformers.json @@ -40,7 +40,7 @@ "measured_throughput": 762.5595168883357 }, "gaudi3": { - "measured_throughput": 5025.5970390534085 + "measured_throughput": 4922.539053408532 } }, "tests/test_sentence_transformers.py::test_compute_embeddings_throughput[sentence-transformers/distiluse-base-multilingual-cased-v1]": { @@ -62,7 +62,7 @@ "measured_throughput": 3807.2486282025716 }, "gaudi3": { - "measured_throughput": 5995.942563633102 + "measured_throughput": 5905.9363310232243 } }, "tests/test_sentence_transformers.py::test_compute_embeddings_throughput[sentence-transformers/multi-qa-MiniLM-L6-cos-v1]": { @@ -84,7 +84,7 @@ "measured_throughput": 944.6166139694299 }, "gaudi3": { - "measured_throughput": 6167.298763111252 + "measured_throughput": 6044.311125223232 } }, "tests/test_sentence_transformers.py::test_compute_embeddings_throughput[sentence-transformers/multi-qa-mpnet-base-dot-v1]": { @@ -95,7 +95,7 @@ "measured_throughput": 545.3360251829846 }, "gaudi3": { - "measured_throughput": 5011.953212884994 + "measured_throughput": 4962.288434499423 } }, "tests/test_sentence_transformers.py::test_compute_embeddings_throughput[sentence-transformers/paraphrase-MiniLM-L3-v2]": { @@ -142,4 +142,4 @@ "measured_throughput": 4906.993110085868 } } -} \ No newline at end of file +} diff --git a/tests/test_sentence_transformers.py b/tests/test_sentence_transformers.py index a8ddcbb78a..559e14fa4b 100644 --- a/tests/test_sentence_transformers.py +++ b/tests/test_sentence_transformers.py @@ -50,12 +50,16 @@ def _test_sentence_transformers( sentences = list(sentences) - for i in range(2): - start_time = time.perf_counter() - _ = model.encode(sentences, batch_size=32) - end_time = time.perf_counter() - diff_time = end_time - start_time - measured_throughput = len(sentences) / diff_time + measured_throughput0 =[] + for j in range(10): + for i in range(2): + start_time = time.perf_counter() + _ = model.encode(sentences, batch_size=32) + end_time = time.perf_counter() + diff_time = end_time - start_time + measured_throughput0.append(len(sentences) / diff_time) + measured_throughput0.sort() + measured_throughput = sum(measured_throughput0[2:8])/6 # Only assert the last measured throughtput as the first iteration is used as a warmup baseline.assertRef( From 5d2fbde1ac793b6f6a856c16bb7c15ac33c0aae7 Mon Sep 17 00:00:00 2001 From: Harish Subramony <81822986+hsubramony@users.noreply.github.com> Date: Thu, 20 Mar 2025 02:18:34 -0700 Subject: [PATCH 70/89] Fix typo in modeling llama (#1864) --- examples/text-generation/README.md | 2 +- .../transformers/models/llama/modeling_llama.py | 2 +- .../tests/test_text_generation_example.json | 2 +- tests/test_text_generation_example.py | 15 +++++++++++++-- 4 files changed, 16 insertions(+), 5 deletions(-) diff --git a/examples/text-generation/README.md b/examples/text-generation/README.md index e89774686b..2f778c1792 100755 --- a/examples/text-generation/README.md +++ b/examples/text-generation/README.md @@ -297,7 +297,7 @@ PT_ENABLE_INT64_SUPPORT=1 PT_HPU_LAZY_MODE=0 python ../gaudi_spawn.py --world_s --bucket_size=128 \ --use_flash_attention \ --flash_attention_recompute \ ---batch_size 246 \ +--batch_size 220 \ --max_input_tokens 2048 \ --max_new_tokens 2048 \ --torch_compile \ diff --git a/optimum/habana/transformers/models/llama/modeling_llama.py b/optimum/habana/transformers/models/llama/modeling_llama.py index a5f98ad644..ca469dfd3a 100755 --- a/optimum/habana/transformers/models/llama/modeling_llama.py +++ b/optimum/habana/transformers/models/llama/modeling_llama.py @@ -1146,7 +1146,7 @@ def __init__(self, config: LlamaConfig): layers = [] for layer_idx in range(config.num_hidden_layers): layer = GaudiLlamaDecoderLayer(config, layer_idx) - if hasattr(config, "paralle_strategy") and config.parallel_strategy is not None: + if hasattr(config, "parallel_strategy") and config.parallel_strategy is not None: layer = config.parallel_strategy.distribute_layer(layer, layer_idx) layers.append(layer) self.layers = torch.nn.ModuleList(layers) diff --git a/tests/baselines/fixture/tests/test_text_generation_example.json b/tests/baselines/fixture/tests/test_text_generation_example.json index b0c1f40f81..44a3181beb 100644 --- a/tests/baselines/fixture/tests/test_text_generation_example.json +++ b/tests/baselines/fixture/tests/test_text_generation_example.json @@ -484,7 +484,7 @@ "throughput": 1345.2369318328463 }, "gaudi3": { - "throughput": 4660.026752215663 + "throughput": 5057.520303949097 } }, "tests/test_text_generation_example.py::test_text_generation_fp8[meta-llama/Llama-2-70b-hf-4-207-False-2048-128]": { diff --git a/tests/test_text_generation_example.py b/tests/test_text_generation_example.py index 1b50cf9c7f..df6226cd8b 100644 --- a/tests/test_text_generation_example.py +++ b/tests/test_text_generation_example.py @@ -304,6 +304,10 @@ def _test_text_generation( command += [ f"--parallel_strategy={parallel_strategy}", ] + if "llama-2-7b-hf" in model_name.lower() and torch_compile and parallel_strategy == "tp": + command.insert(-2, "--bucket_size 128") + command.insert(-2, "--bucket_internal") + command.insert(-2, "--max_input_tokens 2048") with TemporaryDirectory() as tmp_dir: command.append(f"--output_dir {tmp_dir}") @@ -510,12 +514,19 @@ def test_text_generation_torch_compile_distributed(model_name: str, baseline, to @pytest.mark.parametrize("model_name", MODELS_TO_TEST["distributed_tp"]) def test_text_generation_distributed_tp(model_name: str, baseline, token): world_size = 8 + batch_size=64 + max_input_tokens=128 + if "llama-2-7b-hf" in model_name.lower(): + #match the params from examples/readme + batch_size=220 + max_input_tokens=2048 + _test_text_generation( model_name, baseline, token, - batch_size=64, - max_input_tokens=128, + batch_size=batch_size, + max_input_tokens=max_input_tokens, world_size=world_size, torch_compile=True, parallel_strategy="tp", From 0ec8b04c939c7febebf6bbe23f6785eb55576ab3 Mon Sep 17 00:00:00 2001 From: Iman Gohari Date: Thu, 20 Mar 2025 02:19:10 -0700 Subject: [PATCH 71/89] fea(): Added the updated skip list for mistral/mixtral tests (#1863) --- .../models/mistral/test_modeling_mistral.py | 16 ++++++++++++++++ .../models/mixtral/test_modeling_mixtral.py | 12 ++++++++++++ 2 files changed, 28 insertions(+) diff --git a/tests/transformers/tests/models/mistral/test_modeling_mistral.py b/tests/transformers/tests/models/mistral/test_modeling_mistral.py index 962eea1b0e..7caaefff10 100644 --- a/tests/transformers/tests/models/mistral/test_modeling_mistral.py +++ b/tests/transformers/tests/models/mistral/test_modeling_mistral.py @@ -297,6 +297,22 @@ class MistralModelTest(ModelTesterMixin, GenerationTesterMixin, unittest.TestCas test_headmasking = False test_pruning = False + @unittest.skip(reason="This test is not supported for Mistral") + def test_assisted_decoding_sample(self): + pass + + @unittest.skip(reason="This test is not supported for Mistral") + def test_generate_from_inputs_embeds_decoder_only(self): + pass + + @unittest.skip(reason="This test is not supported for Mistral") + def test_greedy_generate(self): + pass + + @unittest.skip(reason="This test is not supported for Mistral") + def test_sample_generate(self): + pass + @unittest.skip(reason="This test is not supported for Mistral") def test_beam_search_generate(self): pass diff --git a/tests/transformers/tests/models/mixtral/test_modeling_mixtral.py b/tests/transformers/tests/models/mixtral/test_modeling_mixtral.py index 1b2230aaf2..a08be61136 100644 --- a/tests/transformers/tests/models/mixtral/test_modeling_mixtral.py +++ b/tests/transformers/tests/models/mixtral/test_modeling_mixtral.py @@ -298,6 +298,18 @@ class MixtralModelTest(ModelTesterMixin, GenerationTesterMixin, unittest.TestCas test_headmasking = False test_pruning = False + @unittest.skip(reason="This test is not supported for Mixtral") + def test_assisted_decoding_sample(self): + pass + + @unittest.skip(reason="This test is not supported for Mixtral") + def test_generate_from_inputs_embeds_decoder_only(self): + pass + + @unittest.skip(reason="This test is not supported for Mixtral") + def test_sample_generate(self): + pass + @unittest.skip(reason="This test is not supported for Mixtral") def test_beam_search_generate(self): pass From 639f96d714b2e62577ef5a47fb8dad8614551599 Mon Sep 17 00:00:00 2001 From: Daniel Socek Date: Fri, 21 Mar 2025 04:44:18 -0400 Subject: [PATCH 72/89] Fix llama internal bucketing issue (#1871) Signed-off-by: Daniel Socek --- optimum/habana/transformers/models/llama/modeling_llama.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/optimum/habana/transformers/models/llama/modeling_llama.py b/optimum/habana/transformers/models/llama/modeling_llama.py index ca469dfd3a..b43eff206d 100755 --- a/optimum/habana/transformers/models/llama/modeling_llama.py +++ b/optimum/habana/transformers/models/llama/modeling_llama.py @@ -1600,6 +1600,10 @@ def prepare_inputs_for_generation( if num_logits_to_keep is not None: model_inputs["num_logits_to_keep"] = num_logits_to_keep + if bucket_internal and reuse_cache is not True: + # update input with kv cache len to capture padding changes during internal bucketing without cache reuse + model_inputs["kv_cache_len"] = kwargs.get("kv_cache_len") + model_inputs.update( { "position_ids": position_ids, From f3124e7af2fd299943f5a13b8307c38efec1ff76 Mon Sep 17 00:00:00 2001 From: Shifani Rajabose Date: Fri, 21 Mar 2025 04:45:04 -0400 Subject: [PATCH 73/89] Fix regression for test_run_image2text_lora_finetune_idefics2-8b_multi_card (#1870) --- 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 df8e0d7e73..20a26bd895 100644 --- a/tests/baselines/fixture/tests/test_examples.json +++ b/tests/baselines/fixture/tests/test_examples.json @@ -485,7 +485,7 @@ "train_samples_per_second": 11.8 }, "gaudi3": { - "eval_accuracy": 0.6910165783279163, + "eval_accuracy": 0.6, "train_runtime": 273.7778, "train_samples_per_second": 17.93 } From bbada81b45f7878e0bb9c5c07af9bae86d7c096e Mon Sep 17 00:00:00 2001 From: Silvia Colabrese Date: Mon, 24 Mar 2025 18:32:19 +0100 Subject: [PATCH 74/89] Revert "Move model to device before wrapping with FSDP (#1801)" (#1865) --- optimum/habana/accelerate/accelerator.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/optimum/habana/accelerate/accelerator.py b/optimum/habana/accelerate/accelerator.py index 73d42d2dca..de027eff8e 100644 --- a/optimum/habana/accelerate/accelerator.py +++ b/optimum/habana/accelerate/accelerator.py @@ -476,9 +476,6 @@ def prepare_model(self, model: torch.nn.Module, device_placement: bool = None, e "limit_all_gathers": fsdp_plugin.limit_all_gathers, "device_id": torch.device("hpu", torch.hpu.current_device()), } - # There's issue with moving view tensors to device within FSDP class [See: https://github.com/pytorch/pytorch/issues/147321] - # Due to above issue, view tensor's may lead to silent incorrent behavior, while pretending to be view they're really not - model = model.to(kwargs["device_id"]) model = FSDP(model, **kwargs) if fsdp_plugin.activation_checkpointing: from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import ( From 677acbe9bf68eae7d731da1dd9fdfd018ae0abc8 Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Thu, 27 Mar 2025 00:10:49 +0000 Subject: [PATCH 75/89] Fix Qwen2 gradient checkpointing --- optimum/habana/transformers/models/qwen2/modeling_qwen2.py | 1 + 1 file changed, 1 insertion(+) diff --git a/optimum/habana/transformers/models/qwen2/modeling_qwen2.py b/optimum/habana/transformers/models/qwen2/modeling_qwen2.py index 6956d6e4a6..cfffa62c2a 100644 --- a/optimum/habana/transformers/models/qwen2/modeling_qwen2.py +++ b/optimum/habana/transformers/models/qwen2/modeling_qwen2.py @@ -783,6 +783,7 @@ def forward( use_cache, cache_position, None, + None, attn_softmax_bf16, False, use_flash_attention, From bb0e1faf8763a54f07f8998888268d135df0b3c1 Mon Sep 17 00:00:00 2001 From: Silvia Colabrese Date: Tue, 1 Apr 2025 09:13:09 +0200 Subject: [PATCH 76/89] Llama2 70 b OOM fix (#1902) --- optimum/habana/transformers/models/llama/modeling_llama.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/optimum/habana/transformers/models/llama/modeling_llama.py b/optimum/habana/transformers/models/llama/modeling_llama.py index b43eff206d..61a8031bfd 100755 --- a/optimum/habana/transformers/models/llama/modeling_llama.py +++ b/optimum/habana/transformers/models/llama/modeling_llama.py @@ -559,6 +559,7 @@ def pre_attn_forward( position_embeddings: Tuple[torch.Tensor, torch.Tensor], attention_mask: Optional[torch.Tensor], past_key_value: Optional[Cache] = None, + output_attentions: bool = False, use_cache: bool = False, cache_position: Optional[torch.LongTensor] = None, position_ids: Optional[torch.LongTensor] = None, @@ -779,6 +780,8 @@ def pre_attn_forward( attn_output = attn_output.transpose(1, 2).contiguous() attn_output = attn_output.reshape(*input_shape, -1).contiguous() attn_output = self.o_proj(attn_output) + if not output_attentions: + attn_weights = None if not reuse_cache and token_idx is not None and cache_idx is not None and q_len == 1: # Return only past key value shapes and not the tensors during decode phase (q len is 1) From 5301225bbaca44f3e22832559c54f7a7df74cf77 Mon Sep 17 00:00:00 2001 From: Libin Tang Date: Wed, 2 Apr 2025 02:23:19 -0700 Subject: [PATCH 77/89] Fix the flant5-xl training issue (#1905) Co-authored-by: Yaser Afshar --- optimum/habana/transformers/models/t5/modeling_t5.py | 1 + 1 file changed, 1 insertion(+) diff --git a/optimum/habana/transformers/models/t5/modeling_t5.py b/optimum/habana/transformers/models/t5/modeling_t5.py index 2f70e8c85f..ff96861250 100644 --- a/optimum/habana/transformers/models/t5/modeling_t5.py +++ b/optimum/habana/transformers/models/t5/modeling_t5.py @@ -516,6 +516,7 @@ def gaudi_T5ForConditionalGeneration_forward( 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 From c2851e1127d0e3de803cc37c137979183917efa0 Mon Sep 17 00:00:00 2001 From: Daniel Socek Date: Thu, 3 Apr 2025 12:22:00 -0400 Subject: [PATCH 78/89] Improve diffusers training download script (#1906) Signed-off-by: Daniel Socek --- examples/stable-diffusion/training/README.md | 8 ++++++++ .../stable-diffusion/training/download_train_datasets.py | 4 ++-- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/examples/stable-diffusion/training/README.md b/examples/stable-diffusion/training/README.md index 4c1add8b76..df98e8aacc 100644 --- a/examples/stable-diffusion/training/README.md +++ b/examples/stable-diffusion/training/README.md @@ -23,6 +23,14 @@ This directory contains scripts that showcase how to perform training/fine-tunin The `textual_inversion_sdxl.py` script shows how to implement textual inversion fine-tuning on Gaudi for XL diffusion models such as `stabilityai/stable-diffusion-xl-base-1.0` or `cagliostrolab/animagine-xl-3.1` for example. +For this example we will use a set of cat toy images from the following dataset: +[https://huggingface.co/datasets/diffusers/cat_toy_example](https://huggingface.co/datasets/diffusers/cat_toy_example). + +To download this and other example training datasets locally, run: +```bash +python download_train_datasets.py +``` + Assuming the afforemenioned cat toy dataset has been obtained, we can launch textual inversion XL training using: ```bash diff --git a/examples/stable-diffusion/training/download_train_datasets.py b/examples/stable-diffusion/training/download_train_datasets.py index 6ff500c9ef..5df883121d 100755 --- a/examples/stable-diffusion/training/download_train_datasets.py +++ b/examples/stable-diffusion/training/download_train_datasets.py @@ -45,8 +45,8 @@ repo_type="dataset", local_dir=local_dir, ) -shutil.move(file_path1, local_dir) -shutil.move(file_path2, local_dir) +shutil.copy2(file_path1, local_dir) +shutil.copy2(file_path2, local_dir) cache_dir = Path(local_dir, ".cache") if cache_dir.is_dir(): shutil.rmtree(cache_dir) From 4e12497e1e31a9e6bf0c5d4d4f1fc6a895e04c40 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil <57442720+IlyasMoutawwakil@users.noreply.github.com> Date: Fri, 4 Apr 2025 01:25:07 +0200 Subject: [PATCH 79/89] Count `num_items_in_batch` only when necessary (#1869) Co-authored-by: regisss <15324346+regisss@users.noreply.github.com> --- optimum/habana/transformers/trainer.py | 126 +++-------- .../fixture/tests/test_examples.json | 198 +++++++++--------- tests/configs/examples/llama_7b.json | 5 +- tests/test_sentence_transformers.py | 4 +- tests/test_text_generation_example.py | 10 +- 5 files changed, 141 insertions(+), 202 deletions(-) diff --git a/optimum/habana/transformers/trainer.py b/optimum/habana/transformers/trainer.py index 2857bfe792..510ad8a5a0 100644 --- a/optimum/habana/transformers/trainer.py +++ b/optimum/habana/transformers/trainer.py @@ -900,17 +900,6 @@ def hpu_deepspeed_checkpointing(function, *checkpoint_args, use_reentrant: Optio else: self.log_evaluate_save_time = None - # Calculate the number of items in each batch for all epochs - num_items_in_batches = self.get_num_items_in_batches( - args, - epochs_trained, - num_train_epochs, - train_dataloader, - len_dataloader, - num_examples, - steps_trained_in_current_epoch, - ) - hb_profiler = HabanaProfile( warmup=self.args.profiling_warmup_steps, active=self.args.profiling_steps, @@ -964,8 +953,9 @@ def hpu_deepspeed_checkpointing(function, *checkpoint_args, use_reentrant: Optio for _ in range(total_updates): update_step += 1 num_batches = args.gradient_accumulation_steps if update_step != (total_updates - 1) else remainder - batch_samples = self.get_iterator_batch_samples(epoch_iterator, num_batches) - num_items_in_batch = num_items_in_batches[epoch][update_step] + batch_samples, num_items_in_batch = self.get_batch_samples_transformers( + epoch_iterator, num_batches, args.device + ) for i, inputs in enumerate(batch_samples): step += 1 @@ -2593,91 +2583,41 @@ def _zero_model_grad(self, model): model.zero_grad() model._zero_grad_kwargs = {} - def get_num_items_in_batches( - self, - args, - epochs_trained, - num_train_epochs, - train_dataloader, - len_dataloader, - num_examples, - steps_trained_in_current_epoch, - ): - """ - Calculate the number of items in each batch for all epochs during training. - """ - steps_in_epoch = ( - len_dataloader if len_dataloader is not None else args.max_steps * args.gradient_accumulation_steps - ) - - remainder = num_examples % args.gradient_accumulation_steps - if remainder == 0: - remainder = args.gradient_accumulation_steps - - total_updates = steps_in_epoch // args.gradient_accumulation_steps + 1 - if args.gradient_accumulation_steps == 1: - total_updates -= 1 - global_step = 0 - - num_items_in_batches = [] - for epoch in range(epochs_trained, num_train_epochs): - if epoch == epochs_trained and steps_trained_in_current_epoch > 0: - epoch_dataloader = skip_first_batches(train_dataloader, steps_trained_in_current_epoch) - else: - epoch_dataloader = train_dataloader - - if hasattr(epoch_dataloader, "set_epoch"): - epoch_dataloader.set_epoch(epoch) + def get_batch_samples_transformers(self, epoch_iterator, num_batches, device): + batch_samples = [] + num_items_in_batch = None - epoch_iterator = iter(epoch_dataloader) + for _ in range(num_batches): try: - first_batch = next(epoch_iterator) + batch_samples.append(next(epoch_iterator)) except StopIteration: break - # Check if the batch contains "labels" (once per epoch) - if "labels" not in first_batch: - num_items_in_batches.append([None] * total_updates) - continue - - device = first_batch["labels"].device - - # Reset the iterator - epoch_iterator = iter(epoch_dataloader) - - num_items_in_batches.append([]) - for update_step in range(total_updates): - num_batches = args.gradient_accumulation_steps if update_step != (total_updates - 1) else remainder - - num_items_in_batch = 0 - for _ in range(num_batches): - try: - batch = next(epoch_iterator) - num_items_in_batch += (batch["labels"].ne(-100)).sum().item() - except StopIteration: - break - - if self.args.average_tokens_across_devices and num_items_in_batch > 0: - num_items_in_batch = torch.tensor(num_items_in_batch, device=device) - num_items_in_batch = self.accelerator.gather(num_items_in_batch).sum().item() - # Set to None if no items in batch - if num_items_in_batch == 0: - num_items_in_batch = None - - num_items_in_batches[epoch].append(num_items_in_batch) - global_step += 1 + count_num_items_in_batch = ( + len(batch_samples) > 0 + and "labels" in batch_samples[0] + and ( + # num_items_in_batch is passed to model forward + # https://github.com/huggingface/transformers/blob/v4.49.0/src/transformers/trainer.py#L3757 + self.model_accepts_loss_kwargs + # num_items_in_batch is passed to compute_loss_func + # https://github.com/huggingface/transformers/blob/v4.49.0/src/transformers/trainer.py#L3773 + or self.compute_loss_func is not None + # num_items_in_batch is also verified if (self.model_accepts_loss_kwargs or self.compute_loss_func) + # https://github.com/huggingface/transformers/blob/v4.49.0/src/transformers/trainer.py#L3790 + ) + ) - # For iterable datasets, don't do more than max_steps steps - if len_dataloader is None and global_step >= args.max_steps: - break + if count_num_items_in_batch: + # For now we don't support object detection + try: + num_items_in_batch = torch.cat(batch_samples).ne(-100).sum() + except (TypeError, AttributeError, RuntimeError): + pass - return num_items_in_batches + if num_items_in_batch is not None: + if self.args.average_tokens_across_devices: + num_items_in_batch = self.accelerator.gather(num_items_in_batch).sum() + num_items_in_batch = num_items_in_batch.to(device) - def get_iterator_batch_samples(self, epoch_iterator, num_batches): - batch_samples = [] - for _ in range(num_batches): - try: - batch_samples += [next(epoch_iterator)] - except StopIteration: - break - return batch_samples + return batch_samples, num_items_in_batch diff --git a/tests/baselines/fixture/tests/test_examples.json b/tests/baselines/fixture/tests/test_examples.json index 20a26bd895..ab8358df87 100644 --- a/tests/baselines/fixture/tests/test_examples.json +++ b/tests/baselines/fixture/tests/test_examples.json @@ -25,13 +25,13 @@ "gaudi3": { "perplexity": 3.843924462719278, "train_runtime": 148.7151, - "train_samples_per_second": 32.357 + "train_samples_per_second": 29.824 } }, "tests/test_examples.py::DeepSpeedTextClassificationExampleTester::test_run_glue_LlamaGuard-7b_deepspeed": { "gaudi2": { - "eval_f1": 0.8873483535528596, - "train_runtime": 62.4539, + "eval_f1": 0.8811, + "train_runtime": 68.1838, "train_samples_per_second": 342.169 }, "gaudi3": { @@ -48,7 +48,7 @@ }, "gaudi3": { "perplexity": 6.877100646486551, - "train_runtime": 477.7145, + "train_runtime": 519.1738, "train_samples_per_second": 29.814 } }, @@ -66,8 +66,8 @@ }, "gaudi3": { "perplexity": 16.260238201071928, - "train_runtime": 243.1757, - "train_samples_per_second": 34.196 + "train_runtime": 375.5908, + "train_samples_per_second": 21.888 } }, "tests/test_examples.py::DeepspeedCausalLanguageModelingExampleTester::test_run_clm_gemma-2b-it_deepspeed": { @@ -78,8 +78,8 @@ }, "gaudi3": { "perplexity": 980.9833890324784, - "train_runtime": 51.73, - "train_samples_per_second": 142.775 + "train_runtime": 57.7676, + "train_samples_per_second": 135.39 } }, "tests/test_examples.py::DeepspeedCausalLanguageModelingExampleTester::test_run_clm_gpt-neox-20b_deepspeed": { @@ -101,13 +101,13 @@ "train_samples_per_second": 16.464 }, "gaudi2": { - "perplexity": 13.237754028004865, - "train_runtime": 206.5775, - "train_samples_per_second": 95.539 + "perplexity": 13.1786, + "train_runtime": 191.6778, + "train_samples_per_second": 96.789 }, "gaudi3": { "perplexity": 13.155277331993139, - "train_runtime": 159.357, + "train_runtime": 175.0512, "train_samples_per_second": 150.538 } }, @@ -119,20 +119,20 @@ }, "gaudi3": { "perplexity": 3.728595328528421, - "train_runtime": 440.2459, - "train_samples_per_second": 19.627 + "train_runtime": 712.4347, + "train_samples_per_second": 13.065 } }, "tests/test_examples.py::DeepspeedSummarizationExampleTester::test_run_summarization_flan-t5-xxl_deepspeed": { "gaudi2": { - "eval_rougeLsum": 29.308, - "train_runtime": 155.86, - "train_samples_per_second": 28.387 + "eval_rougeLsum": 27.9095, + "train_runtime": 141.557, + "train_samples_per_second": 32.239 }, "gaudi3": { "eval_rougeLsum": 28.0738, - "train_runtime": 118.419, - "train_samples_per_second": 52.048 + "train_runtime": 137.2073, + "train_samples_per_second": 44.762 } }, "tests/test_examples.py::EagerModeCausalLanguageModelingExampleTester::test_run_clm_gemma-2b-it_single_card": { @@ -154,14 +154,14 @@ "train_samples_per_second": 212.722 }, "gaudi2": { - "eval_accuracy": 0.9850666666666666, - "train_runtime": 77.8934, - "train_samples_per_second": 826.766 + "eval_accuracy": 0.9871, + "train_runtime": 74.7251, + "train_samples_per_second": 831.235 }, "gaudi3": { "eval_accuracy": 0.9849333333333333, - "train_runtime": 73.8308, - "train_samples_per_second": 1155.964 + "train_runtime": 86.4819, + "train_samples_per_second": 1043.6 } }, "tests/test_examples.py::ImageClassificationExampleTester::test_run_image_classification_vit-base-patch16-224-in21k_single_card": { @@ -171,14 +171,14 @@ "train_samples_per_second": 359.584 }, "gaudi2": { - "eval_accuracy": 0.9690666666666666, - "train_runtime": 54.9734, - "train_samples_per_second": 870.272 + "eval_accuracy": 0.9696, + "train_runtime": 54.0092, + "train_samples_per_second": 886.623 }, "gaudi3": { "eval_accuracy": 0.9690666666666666, - "train_runtime": 47.9419, - "train_samples_per_second": 1164.009 + "train_runtime": 51.3575, + "train_samples_per_second": 1100.39 } }, "tests/test_examples.py::MultiCardAudioClassificationExampleTester::test_run_audio_classification_ast-finetuned-speech-commands-v2_multi_card": { @@ -190,9 +190,9 @@ }, "gaudi3": { "eval_accuracy": 0.19650135869565216, - "eval_samples_per_second": 3352.901, - "train_runtime": 106.5372, - "train_samples_per_second": 2676.242 + "eval_samples_per_second": 2442.221, + "train_runtime": 118.0048, + "train_samples_per_second": 2462.94 } }, "tests/test_examples.py::MultiCardAudioClassificationExampleTester::test_run_audio_classification_wav2vec2-base_multi_card": { @@ -211,8 +211,8 @@ "gaudi3": { "eval_accuracy": 0.7352241847826086, "eval_samples_per_second": 2059.992, - "train_runtime": 57.0028, - "train_samples_per_second": 4213.033 + "train_runtime": 71.6177, + "train_samples_per_second": 3977.496 } }, "tests/test_examples.py::MultiCardBridgetowerExampleTester::test_run_bridgetower_bridgetower-large-itm-mlm-itc_multi_card": { @@ -221,8 +221,8 @@ "train_samples_per_second": 904.93 }, "gaudi3": { - "train_runtime": 342.4851, - "train_samples_per_second": 1009.467 + "train_runtime": 496.3196, + "train_samples_per_second": 783.481 } }, "tests/test_examples.py::MultiCardCausalLanguageModelingAdaloraExampleTester::test_run_lora_clm_llama-7b_multi_card": { @@ -233,8 +233,8 @@ }, "gaudi3": { "perplexity": 2.592915682175543, - "train_runtime": 818.9693, - "train_samples_per_second": 85.059 + "train_runtime": 1026.9536, + "train_samples_per_second": 71.308 } }, "tests/test_examples.py::MultiCardCausalLanguageModelingExampleTester::test_run_clm_gemma-2b-it_multi_card": { @@ -257,8 +257,8 @@ }, "gaudi3": { "perplexity": 3.291398111098924, - "train_runtime": 390.7556, - "train_samples_per_second": 256.027 + "train_runtime": 485.2917, + "train_samples_per_second": 233.759 } }, "tests/test_examples.py::MultiCardCausalLanguageModelingLORAExampleTester2::test_run_lora_clm_falcon-40b_multi_card": { @@ -269,8 +269,8 @@ }, "gaudi3": { "perplexity": 1.588740773299791, - "train_runtime": 408.8298, - "train_samples_per_second": 33.87 + "train_runtime": 574.862, + "train_samples_per_second": 24.276 } }, "tests/test_examples.py::MultiCardCausalLanguageModelingLORAExampleTester2::test_run_lora_clm_llama-7b_multi_card": { @@ -293,8 +293,8 @@ }, "gaudi3": { "perplexity": 3.694849124063941, - "train_runtime": 320.063, - "train_samples_per_second": 35.863 + "train_runtime": 463.3435, + "train_samples_per_second": 24.637 } }, "tests/test_examples.py::MultiCardCausalLanguageModelingLORAExampleTester::test_run_lora_clm_llama-7b_multi_card": { @@ -310,8 +310,8 @@ }, "gaudi3": { "perplexity": 2.3665888138128466, - "train_runtime": 394.5646, - "train_samples_per_second": 238.486 + "train_runtime": 442.8761, + "train_samples_per_second": 219.733 } }, "tests/test_examples.py::MultiCardCausalLanguageModelingLORAFSDPCompileExampleTester::test_run_lora_clm_llama-7b_multi_card": { @@ -322,8 +322,8 @@ }, "gaudi3": { "perplexity": 2.42632366178759, - "train_runtime": 98.5791, - "train_samples_per_second": 126.028 + "train_runtime": 166.8724, + "train_samples_per_second": 73.249 } }, "tests/test_examples.py::MultiCardCausalLanguageModelingLlamaAdapterExampleTester::test_run_lora_clm_llama-7b_multi_card": { @@ -334,8 +334,8 @@ }, "gaudi3": { "perplexity": 5.575957971980852, - "train_runtime": 227.3213, - "train_samples_per_second": 504.974 + "train_runtime": 268.6561, + "train_samples_per_second": 456.126 } }, "tests/test_examples.py::MultiCardCausalLanguageModelingLnExampleTester::test_run_lora_clm_llama-7b_multi_card": { @@ -346,8 +346,8 @@ }, "gaudi3": { "perplexity": 2.842264808115683, - "train_runtime": 332.9477, - "train_samples_per_second": 267.004 + "train_runtime": 407.6245, + "train_samples_per_second": 240.47 } }, "tests/test_examples.py::MultiCardCausalLanguageModelingLoRACPExampleTester::test_run_lora_clm_llama-7b_deepspeed": { @@ -358,8 +358,8 @@ }, "gaudi3": { "perplexity": 2.8421374130082477, - "train_runtime": 219.1417, - "train_samples_per_second": 55.554 + "train_runtime": 268.0088, + "train_samples_per_second": 52.488 } }, "tests/test_examples.py::MultiCardCausalLanguageModelingLoRAFP8ExampleTester::test_run_lora_clm_llama-7b_multi_card": { @@ -370,7 +370,7 @@ }, "gaudi3": { "perplexity": 2.3750491436810424, - "train_runtime": 547.5649, + "train_runtime": 670.2323, "train_samples_per_second": 323.175 } }, @@ -382,8 +382,8 @@ }, "gaudi3": { "perplexity": 1.0262332298756216, - "train_runtime": 16.2913, - "train_samples_per_second": 78.376 + "train_runtime": 20.7457, + "train_samples_per_second": 60.479 } }, "tests/test_examples.py::MultiCardCausalLanguageModelingPrefixTuningExampleTester::test_run_prompt_tuning_clm_llama-7b_multi_card": { @@ -394,8 +394,8 @@ }, "gaudi3": { "perplexity": 1.1720024747280242, - "train_runtime": 15.1138, - "train_samples_per_second": 67.894 + "train_runtime": 17.2662, + "train_samples_per_second": 61.545 } }, "tests/test_examples.py::MultiCardCausalLanguageModelingPromptTuningExampleTester::test_run_prompt_tuning_clm_llama-7b_multi_card": { @@ -406,20 +406,20 @@ }, "gaudi3": { "perplexity": 1.2158095633720596, - "train_runtime": 14.0663, - "train_samples_per_second": 75.406 + "train_runtime": 16.7016, + "train_samples_per_second": 61.478 } }, "tests/test_examples.py::MultiCardCausalLanguageModelingVeraExampleTester::test_run_lora_clm_llama-7b_multi_card": { "gaudi2": { - "perplexity": 9.064502567217577, - "train_runtime": 312.9258, - "train_samples_per_second": 127.305 + "perplexity": 2.6174, + "train_runtime": 312.3, + "train_samples_per_second": 125.798 }, "gaudi3": { - "perplexity": 8.65669958765362, - "train_runtime": 261.8749, - "train_samples_per_second": 199.0 + "perplexity": 2.613, + "train_runtime": 321.3809, + "train_samples_per_second": 168.296 } }, "tests/test_examples.py::MultiCardDPOExampleTester::test_dpo_llama-7b_multi_card": { @@ -428,8 +428,8 @@ "train_samples_per_second": 13.499 }, "gaudi3": { - "train_runtime": 194.4848, - "train_samples_per_second": 16.454 + "train_runtime": 211.3438, + "train_samples_per_second": 15.141 } }, "tests/test_examples.py::MultiCardImageClassificationExampleTester::test_run_image_classification_swin-base-patch4-window7-224-in22k_multi_card": { @@ -445,8 +445,8 @@ }, "gaudi3": { "eval_accuracy": 0.9817333333333333, - "train_runtime": 74.7483, - "train_samples_per_second": 8253.709 + "train_runtime": 106.7478, + "train_samples_per_second": 7042.508 } }, "tests/test_examples.py::MultiCardImageClassificationExampleTester::test_run_image_classification_vit-base-patch16-224-in21k_multi_card": { @@ -462,7 +462,7 @@ }, "gaudi3": { "eval_accuracy": 0.9677333333333333, - "train_runtime": 33.4011, + "train_runtime": 38.7232, "train_samples_per_second": 6636.054 } }, @@ -474,8 +474,8 @@ }, "gaudi3": { "eval_accuracy": 0.9044574025188373, - "train_runtime": 397.9607, - "train_samples_per_second": 39.088 + "train_runtime": 583.0969, + "train_samples_per_second": 35.056 } }, "tests/test_examples.py::MultiCardImageToTextModelingLoRAExampleTester::test_run_image2text_lora_finetune_idefics2-8b_multi_card": { @@ -515,8 +515,8 @@ }, "gaudi3": { "perplexity": 2.8534683742096933, - "train_runtime": 53.0805, - "train_samples_per_second": 1335.957 + "train_runtime": 109.6454, + "train_samples_per_second": 994.854 } }, "tests/test_examples.py::MultiCardPPOExampleTester::test_ppo_llama-7b_multi_card": { @@ -525,8 +525,8 @@ "train_samples_per_second": 0.5 }, "gaudi3": { - "train_runtime": 40.73775029182434, - "train_samples_per_second": 0.7855122035647137 + "train_runtime": 47.45752716064453, + "train_samples_per_second": 0.6742871345082827 } }, "tests/test_examples.py::MultiCardProteinFoldingClassificationTester::test_run_sequence_classification_protst-esm1b-for-sequential-classification_multi_card": { @@ -537,7 +537,7 @@ }, "gaudi3": { "eval_accuracy": 0.5442452284557547, - "train_runtime": 40.0248, + "train_runtime": 50.2717, "train_samples_per_second": 1564.079 } }, @@ -554,7 +554,7 @@ }, "gaudi3": { "eval_f1": 94.33668918864852, - "train_runtime": 153.0279, + "train_runtime": 184.2499, "train_samples_per_second": 3146.332 } }, @@ -564,8 +564,8 @@ "train_samples_per_second": 1.6 }, "gaudi3": { - "train_runtime": 135.1176, - "train_samples_per_second": 3.027 + "train_runtime": 169.1724, + "train_samples_per_second": 2.418 } }, "tests/test_examples.py::MultiCardSFTChatExampleTester::test_sft_Qwen2-7B_multi_card": { @@ -574,7 +574,7 @@ "train_samples_per_second": 7.342 }, "gaudi3": { - "train_runtime": 587.8481, + "train_runtime": 693.6076, "train_samples_per_second": 13.968 } }, @@ -584,7 +584,7 @@ "train_samples_per_second": 120 }, "gaudi3": { - "train_runtime": 364.7036, + "train_runtime": 383.0226, "train_samples_per_second": 193.023 } }, @@ -594,7 +594,7 @@ "train_samples_per_second": 51.54 }, "gaudi3": { - "train_runtime": 316.0836, + "train_runtime": 347.6082, "train_samples_per_second": 86.193 } }, @@ -612,10 +612,10 @@ "train_samples_per_second": 218.0 }, "gaudi3": { - "eval_samples_per_second": 64.339, - "eval_wer": 0.38905990016638936, - "train_runtime": 290.6815, - "train_samples_per_second": 463.628 + "eval_samples_per_second": 48.493, + "eval_wer": 0.4296589018302829, + "train_runtime": 321.6139, + "train_samples_per_second": 504.087 } }, "tests/test_examples.py::MultiCardSpeechRecognitionExampleTester::test_run_speech_recognition_ctc_wav2vec2-large-lv60_multi_card": { @@ -634,8 +634,8 @@ "gaudi3": { "eval_samples_per_second": 491.004, "eval_wer": 0.06197937326457755, - "train_runtime": 255.782, - "train_samples_per_second": 292.161 + "train_runtime": 235.8412, + "train_samples_per_second": 392.016 } }, "tests/test_examples.py::MultiCardTextClassificationExampleTester::test_run_glue_bert-large-uncased-whole-word-masking_multi_card": { @@ -650,8 +650,8 @@ "train_samples_per_second": 2845.068 }, "gaudi3": { - "eval_f1": 0.89198606271777, - "train_runtime": 61.3444, + "eval_f1": 0.8809523809523809, + "train_runtime": 88.3481, "train_samples_per_second": 1826.566 } }, @@ -665,8 +665,8 @@ "train_samples_per_second": 14124 }, "gaudi3": { - "train_runtime": 64.3878, - "train_samples_per_second": 19625.412 + "train_runtime": 71.4956, + "train_samples_per_second": 20525.583 } }, "tests/test_examples.py::QuestionAnsweringExampleTester::test_run_qa_roberta-large_single_card": { @@ -698,9 +698,9 @@ "train_samples_per_second": 1100.598 }, "gaudi3": { - "eval_f1": 0.8826446280991735, - "train_runtime": 74.0631, - "train_samples_per_second": 1652.436 + "eval_f1": 0.8826, + "train_runtime": 78.5458, + "train_samples_per_second": 1674.509 } } -} +} \ No newline at end of file diff --git a/tests/configs/examples/llama_7b.json b/tests/configs/examples/llama_7b.json index f021337f67..18f52ef431 100644 --- a/tests/configs/examples/llama_7b.json +++ b/tests/configs/examples/llama_7b.json @@ -584,7 +584,7 @@ "eval_batch_size": 4, "distribution": { "multi_card": { - "learning_rate": 0.01, + "learning_rate": 0.03, "train_batch_size": 8, "metrics": [ "perplexity", @@ -593,7 +593,6 @@ ], "extra_arguments": [ "--bf16", - "--gradient_accumulation_steps 1", "--eval_strategy no", "--save_strategy no", "--warmup_ratio 0.03", @@ -1202,7 +1201,7 @@ "eval_batch_size": 4, "distribution": { "multi_card": { - "learning_rate": 0.01, + "learning_rate": 0.03, "train_batch_size": 8, "metrics": [ "perplexity", diff --git a/tests/test_sentence_transformers.py b/tests/test_sentence_transformers.py index 559e14fa4b..0952082fe6 100644 --- a/tests/test_sentence_transformers.py +++ b/tests/test_sentence_transformers.py @@ -50,7 +50,7 @@ def _test_sentence_transformers( sentences = list(sentences) - measured_throughput0 =[] + measured_throughput0 = [] for j in range(10): for i in range(2): start_time = time.perf_counter() @@ -59,7 +59,7 @@ def _test_sentence_transformers( diff_time = end_time - start_time measured_throughput0.append(len(sentences) / diff_time) measured_throughput0.sort() - measured_throughput = sum(measured_throughput0[2:8])/6 + measured_throughput = sum(measured_throughput0[2:8]) / 6 # Only assert the last measured throughtput as the first iteration is used as a warmup baseline.assertRef( diff --git a/tests/test_text_generation_example.py b/tests/test_text_generation_example.py index dafd16bfb8..498aaf0d9a 100644 --- a/tests/test_text_generation_example.py +++ b/tests/test_text_generation_example.py @@ -514,12 +514,12 @@ def test_text_generation_torch_compile_distributed(model_name: str, baseline, to @pytest.mark.parametrize("model_name", MODELS_TO_TEST["distributed_tp"]) def test_text_generation_distributed_tp(model_name: str, baseline, token): world_size = 8 - batch_size=64 - max_input_tokens=128 + batch_size = 64 + max_input_tokens = 128 if "llama-2-7b-hf" in model_name.lower(): - #match the params from examples/readme - batch_size=220 - max_input_tokens=2048 + # match the params from examples/readme + batch_size = 220 + max_input_tokens = 2048 _test_text_generation( model_name, From 3e831a1bbda09a21bb9d708999788be2dc7c400e Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Fri, 4 Apr 2025 07:47:12 +0000 Subject: [PATCH 80/89] Fix issue when computing `num_items_in_batch` --- optimum/habana/transformers/trainer.py | 2 +- tests/baselines/fixture/tests/test_examples.json | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/optimum/habana/transformers/trainer.py b/optimum/habana/transformers/trainer.py index 510ad8a5a0..fad51b10a0 100644 --- a/optimum/habana/transformers/trainer.py +++ b/optimum/habana/transformers/trainer.py @@ -2611,7 +2611,7 @@ def get_batch_samples_transformers(self, epoch_iterator, num_batches, device): if count_num_items_in_batch: # For now we don't support object detection try: - num_items_in_batch = torch.cat(batch_samples).ne(-100).sum() + num_items_in_batch = torch.cat([batch["labels"] for batch in batch_samples]).ne(-100).sum() except (TypeError, AttributeError, RuntimeError): pass diff --git a/tests/baselines/fixture/tests/test_examples.json b/tests/baselines/fixture/tests/test_examples.json index ab8358df87..44f239d60d 100644 --- a/tests/baselines/fixture/tests/test_examples.json +++ b/tests/baselines/fixture/tests/test_examples.json @@ -106,9 +106,9 @@ "train_samples_per_second": 96.789 }, "gaudi3": { - "perplexity": 13.155277331993139, - "train_runtime": 175.0512, - "train_samples_per_second": 150.538 + "perplexity": 13.143, + "train_runtime": 198.1031, + "train_samples_per_second": 147.801 } }, "tests/test_examples.py::DeepspeedSFTExampleTester::test_sft_Qwen2-72B_deepspeed": { From 6aaef3a53f94a9c248033948604ab34d70fc0222 Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Mon, 7 Apr 2025 14:57:53 +0000 Subject: [PATCH 81/89] Recover throughput for Llava LoRA fine-tuning test --- optimum/habana/transformers/models/llava/modeling_llava.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/optimum/habana/transformers/models/llava/modeling_llava.py b/optimum/habana/transformers/models/llava/modeling_llava.py index 474bd41fc3..a6096644a8 100644 --- a/optimum/habana/transformers/models/llava/modeling_llava.py +++ b/optimum/habana/transformers/models/llava/modeling_llava.py @@ -251,9 +251,7 @@ def forward( if labels is not None: # Shift so that tokens < n predict n 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[:, -(logits.shape[1] - 1) :].to(logits.device) + shift_attention_mask = attention_mask[..., 1:] shift_logits = logits[..., :-1, :][shift_attention_mask.to(logits.device) != 0].contiguous() shift_labels = labels[..., 1:][shift_attention_mask.to(labels.device) != 0].contiguous() else: From 7ae3c3a8d5eb4dc7c37bcc9815fda6f58fe9473a Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Wed, 9 Apr 2025 08:34:59 +0000 Subject: [PATCH 82/89] Update Gaudi3 CI baselines --- tests/baselines/fixture/tests/test_examples.json | 6 +++--- .../baselines/fixture/tests/test_image_to_text_example.json | 2 +- tests/configs/examples/llama_7b.json | 3 ++- 3 files changed, 6 insertions(+), 5 deletions(-) diff --git a/tests/baselines/fixture/tests/test_examples.json b/tests/baselines/fixture/tests/test_examples.json index 44f239d60d..4c5e3bb7e0 100644 --- a/tests/baselines/fixture/tests/test_examples.json +++ b/tests/baselines/fixture/tests/test_examples.json @@ -232,7 +232,7 @@ "train_samples_per_second": 107 }, "gaudi3": { - "perplexity": 2.592915682175543, + "perplexity": 2.631517742463695, "train_runtime": 1026.9536, "train_samples_per_second": 71.308 } @@ -281,8 +281,8 @@ }, "gaudi3": { "perplexity": 1.570946503005108, - "train_runtime": 342.6741, - "train_samples_per_second": 267.801 + "train_runtime": 418.2808, + "train_samples_per_second": 238.235 } }, "tests/test_examples.py::MultiCardCausalLanguageModelingLORAExampleTester::test_run_lora_clm_falcon-40b_multi_card": { diff --git a/tests/baselines/fixture/tests/test_image_to_text_example.json b/tests/baselines/fixture/tests/test_image_to_text_example.json index 58dbd84613..0330059e50 100644 --- a/tests/baselines/fixture/tests/test_image_to_text_example.json +++ b/tests/baselines/fixture/tests/test_image_to_text_example.json @@ -12,7 +12,7 @@ "throughput": 28.755882208438422 }, "gaudi3": { - "throughput": 85.53160250422563 + "throughput": 75.94218824701237 } }, "tests/test_image_to_text_example.py::test_image_to_text_bf16[Qwen/Qwen2-VL-7B-Instruct-1]": { diff --git a/tests/configs/examples/llama_7b.json b/tests/configs/examples/llama_7b.json index 18f52ef431..ad566eb9bf 100644 --- a/tests/configs/examples/llama_7b.json +++ b/tests/configs/examples/llama_7b.json @@ -788,7 +788,8 @@ "--adam_epsilon 1e-08", "--ddp_bucket_cap_mb 50", "--validation_split_percentage 10", - "--attn_softmax_bf16" + "--attn_softmax_bf16", + "--dataloader_drop_last True" ] } } From 9025f8cc74bde36c2959da8b4bd93c40bc83af64 Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Thu, 10 Apr 2025 08:23:35 +0000 Subject: [PATCH 83/89] Update Gaudi2 CI baseline --- 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 4c5e3bb7e0..8a2723ec7f 100644 --- a/tests/baselines/fixture/tests/test_examples.json +++ b/tests/baselines/fixture/tests/test_examples.json @@ -102,7 +102,7 @@ }, "gaudi2": { "perplexity": 13.1786, - "train_runtime": 191.6778, + "train_runtime": 224.3465, "train_samples_per_second": 96.789 }, "gaudi3": { From a781bd9e5f3c4134bab12744b40ff56a50743ba0 Mon Sep 17 00:00:00 2001 From: Harish Subramony <81822986+hsubramony@users.noreply.github.com> Date: Fri, 11 Apr 2025 00:44:56 -0700 Subject: [PATCH 84/89] Update Habana_Validated_Models.md (#1922) --- tests/Habana_Validated_Models.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/Habana_Validated_Models.md b/tests/Habana_Validated_Models.md index 77f8ec54be..0c63ca6ae8 100644 --- a/tests/Habana_Validated_Models.md +++ b/tests/Habana_Validated_Models.md @@ -109,7 +109,9 @@ |tests/test_text_generation_example.py|baichuan-inc/Baichuan2-7B-Chat|mixed-precision(bf16)|single_card| |tests/test_text_generation_example.py|baichuan-inc/Baichuan2-13B-Chat|mixed-precision(bf16)|single_card| |tests/test_text_generation_example.py|deepseek-ai/DeepSeek-V2-Lite|mixed-precision(bf16)|single_card| +|tests/test_text_generation_example.py|THUDM/chatglm2-6b|mixed-precision(bf16)|single_card| |tests/test_text_generation_example.py|THUDM/chatglm3-6b|mixed-precision(bf16)|single_card| +|tests/test_text_generation_example.py|Qwen/Qwen2.5-7B|mixed-precision(bf16)|single_card| |tests/test_text_generation_example.py|tiiuae/falcon-180B|fp8|multi_card| |tests/test_text_generation_example.py|meta-llama/Llama-2-7b-hf|fp8|single_card| |tests/test_text_generation_example.py|meta-llama/Llama-2-70b-hf|fp8|multi_card| @@ -123,6 +125,7 @@ |tests/test_text_generation_example.py|facebook/opt-66b|mixed-precision(bf16)|deepspeed| |tests/test_text_generation_example.py|google/gemma-2-9b|mixed-precision(bf16)|deepspeed| |tests/test_text_generation_example.py|google/gemma-2-27b|mixed-precision(bf16)|deepspeed| +|tests/test_text_generation_example.py|[Qwen/Qwen2.5-72B|mixed-precision(bf16)|deepspeed| |tests/test_text_generation_example.py|meta-llama/Llama-2-7b-hf|mixed-precision(bf16)|torch.compile| |tests/test_text_generation_example.py|meta-llama/Llama-2-7b-hf|mixed-precision(bf16)|distributed_tp| |tests/test_text_generation_example.py|gpt2-xl|mixed-precision(bf16)|contrastive_search| @@ -134,3 +137,4 @@ |tests/test_sentence_transformers.py|sentence-transformers/all-mpnet-base-v2|mixed-precision(bf16)|single_card| |tests/test_custom_file_input.py|bigcode/starcoder|mixed-precision(bf16)|single_card| |tests/test_diffusers.py|CompVis/stable-diffusion-v1-4|mixed-precision(bf16)|single_card| +|tests/test_video_llava.py|LanguageBind/Video-LLaVA-7B-hf|mixed-precision(bf16)|single_card| From 01a18eb41646787e6e6c825c6c2af5823562cebc Mon Sep 17 00:00:00 2001 From: Daniel Socek Date: Fri, 11 Apr 2025 03:45:49 -0400 Subject: [PATCH 85/89] Update FP8 diffusers doc (#1910) Signed-off-by: Daniel Socek --- examples/stable-diffusion/README.md | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/examples/stable-diffusion/README.md b/examples/stable-diffusion/README.md index 71f887ab7d..e0398c84d4 100644 --- a/examples/stable-diffusion/README.md +++ b/examples/stable-diffusion/README.md @@ -103,10 +103,14 @@ python text_to_image_generation.py \ > The access to SD3 requires agreeing to its terms and conditions at [HuggingFace model page](https://huggingface.co/stabilityai/stable-diffusion-3-medium), > and then authenticating using your HF token via `huggingface-cli login`. -This model can also be quantized with some ops running in FP8 precision. Before quantization, run stats collection using measure mode by setting +This model can also be quantized with some ops running in FP8 precision. Before quantization, run stats collection once using measure mode by setting runtime variable `QUANT_CONFIG=quantization/stable-diffusion-3/measure_config.json` and `--quant_mode measure`. After stats collection, you can run SD3 in quantization mode by setting runtime variable `QUANT_CONFIG=quantization/stable-diffusion-3/quantize_config.json` and `--quant_mode quantize`. +> [!NOTE] +> If you are running SD3 Gaudi pipeline as a service, run quantization mode only once and pipeline in memory will be quantized to use FP8 precision. +> Running quantization mode multiple times on the same pipeline object may cause errors. + To run Stable Diffusion 3.5 Large, use `--model_name_or_path stabilityai/stable-diffusion-3.5-large` in the input. ### FLUX @@ -135,10 +139,14 @@ python text_to_image_generation.py \ > The access to FLUX.1-dev model requires agreeing to its terms and conditions at [HuggingFace model page](https://huggingface.co/black-forest-labs/FLUX.1-dev), > and then authenticating using your HF token via `huggingface-cli login`. -This model can also be quantized with some ops running in FP8 precision. Before quantization, run stats collection using measure mode by setting +This model can also be quantized with some ops running in FP8 precision. Before quantization, run stats collection once using measure mode by setting runtime variable `QUANT_CONFIG=quantization/flux/measure_config.json` and `--quant_mode measure`. After stats collection, you can run FLUX in quantization mode by setting runtime variable `QUANT_CONFIG=quantization/flux/quantize_config.json` and `--quant_mode quantize`. +> [!NOTE] +> If you are running Flux Gaudi pipeline as a service, run quantization mode only once and pipeline in memory will be quantized to use FP8 precision. +> Running quantization mode multiple times on the same pipeline object may cause errors. + To run with FLUX.1-schnell model, a distilled version of FLUX.1 (which is not gated), use `--model_name_or_path black-forest-labs/FLUX.1-schnell`. ## ControlNet From 3d093b2b923898e5fbee7efd5971e8eec9eda566 Mon Sep 17 00:00:00 2001 From: Silvia Colabrese Date: Fri, 11 Apr 2025 09:48:54 +0200 Subject: [PATCH 86/89] Move model to device before wrapping with fsdp (#1909) --- examples/question-answering/run_qa.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/question-answering/run_qa.py b/examples/question-answering/run_qa.py index ff23237c5b..49facbd401 100644 --- a/examples/question-answering/run_qa.py +++ b/examples/question-answering/run_qa.py @@ -376,7 +376,7 @@ def main(): token=model_args.token, trust_remote_code=model_args.trust_remote_code, ) - + model = model.to("hpu") # Tokenizer check: this script requires a fast tokenizer. if not isinstance(tokenizer, PreTrainedTokenizerFast): raise ValueError( From 9e24cb219c764659c1a018f80b9dc80b0aa9e8e9 Mon Sep 17 00:00:00 2001 From: Harshvardhan Chauhan Date: Fri, 11 Apr 2025 01:09:13 -0700 Subject: [PATCH 87/89] Update train_runtime reference number for clm gpt2_xl to fix pytest failure (#1878) --- 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 8a2723ec7f..3bc15272af 100644 --- a/tests/baselines/fixture/tests/test_examples.json +++ b/tests/baselines/fixture/tests/test_examples.json @@ -102,7 +102,7 @@ }, "gaudi2": { "perplexity": 13.1786, - "train_runtime": 224.3465, + "train_runtime": 241.093, "train_samples_per_second": 96.789 }, "gaudi3": { From cb346352592b379522fbf6dfa950eecb5b5136cc Mon Sep 17 00:00:00 2001 From: regisss <15324346+regisss@users.noreply.github.com> Date: Fri, 11 Apr 2025 13:01:41 +0000 Subject: [PATCH 88/89] Update Gaudi3 CI baselines --- tests/baselines/fixture/tests/test_image_to_text_example.json | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/baselines/fixture/tests/test_image_to_text_example.json b/tests/baselines/fixture/tests/test_image_to_text_example.json index 0330059e50..34e0900d32 100644 --- a/tests/baselines/fixture/tests/test_image_to_text_example.json +++ b/tests/baselines/fixture/tests/test_image_to_text_example.json @@ -12,7 +12,7 @@ "throughput": 28.755882208438422 }, "gaudi3": { - "throughput": 75.94218824701237 + "throughput": 85.45308676024212 } }, "tests/test_image_to_text_example.py::test_image_to_text_bf16[Qwen/Qwen2-VL-7B-Instruct-1]": { @@ -20,7 +20,7 @@ "throughput": 19.32562189532818 }, "gaudi3": { - "throughput": 17.216165111759725 + "throughput": 61.22236102130072 } }, "tests/test_image_to_text_example.py::test_image_to_text_bf16[google/paligemma-3b-mix-224-1]": { From db78ffb8e8d5077a72f9fe93632d9c1f5bf0ff39 Mon Sep 17 00:00:00 2001 From: Nikolay Protasov Date: Fri, 11 Apr 2025 19:07:50 +0200 Subject: [PATCH 89/89] Fix gemma2b-it perplexity (#1929) --- tests/baselines/fixture/tests/test_examples.json | 12 ++++++------ tests/configs/examples/gemma_2b_it.json | 8 ++++---- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/tests/baselines/fixture/tests/test_examples.json b/tests/baselines/fixture/tests/test_examples.json index 3bc15272af..2efcc37194 100644 --- a/tests/baselines/fixture/tests/test_examples.json +++ b/tests/baselines/fixture/tests/test_examples.json @@ -72,12 +72,12 @@ }, "tests/test_examples.py::DeepspeedCausalLanguageModelingExampleTester::test_run_clm_gemma-2b-it_deepspeed": { "gaudi2": { - "perplexity": 924.062, - "train_runtime": 75.518, - "train_samples_per_second": 81.097 + "perplexity": 12.8, + "train_runtime": 80.3429, + "train_samples_per_second": 76.06 }, "gaudi3": { - "perplexity": 980.9833890324784, + "perplexity": 12.8, "train_runtime": 57.7676, "train_samples_per_second": 135.39 } @@ -239,12 +239,12 @@ }, "tests/test_examples.py::MultiCardCausalLanguageModelingExampleTester::test_run_clm_gemma-2b-it_multi_card": { "gaudi2": { - "perplexity": 954.5995, + "perplexity": 12.8, "train_runtime": 82.6617, "train_samples_per_second": 94.524 }, "gaudi3": { - "perplexity": 902.0585179806482, + "perplexity": 12.8, "train_runtime": 66.2529, "train_samples_per_second": 159.47 } diff --git a/tests/configs/examples/gemma_2b_it.json b/tests/configs/examples/gemma_2b_it.json index 6ecab478ad..5e612bec6d 100644 --- a/tests/configs/examples/gemma_2b_it.json +++ b/tests/configs/examples/gemma_2b_it.json @@ -18,7 +18,7 @@ ] }, "multi_card": { - "learning_rate": 0.0008, + "learning_rate": 0.0002, "train_batch_size": 4, "metrics": [ "perplexity", @@ -31,7 +31,7 @@ ] }, "deepspeed": { - "learning_rate": 0.0008, + "learning_rate": 0.0002, "train_batch_size": 4, "metrics": [ "perplexity", @@ -66,7 +66,7 @@ ] }, "multi_card": { - "learning_rate": 0.0008, + "learning_rate": 0.0002, "train_batch_size": 4, "metrics": [ "perplexity", @@ -79,7 +79,7 @@ ] }, "deepspeed": { - "learning_rate": 0.0008, + "learning_rate": 0.0002, "train_batch_size": 4, "metrics": [ "perplexity",