# coding=utf-8 # Copyright 2023 The HuggingFace Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Model specific Neuron configurations.""" import copy from functools import partial from typing import TYPE_CHECKING, Dict, List import torch from optimum.exporters.tasks import TasksManager from optimum.utils import ( DummyInputGenerator, DummySeq2SeqDecoderTextInputGenerator, DummyTextInputGenerator, DummyTimestepInputGenerator, DummyVisionInputGenerator, NormalizedConfig, NormalizedConfigManager, NormalizedSeq2SeqConfig, NormalizedTextAndVisionConfig, NormalizedTextConfig, NormalizedVisionConfig, is_diffusers_available, ) from ...neuron.utils import ( ASTDummyAudioInputGenerator, DummyBeamValuesGenerator, DummyControNetInputGenerator, DummyIPAdapterInputGenerator, DummyMaskedPosGenerator, WhisperDummyTextInputGenerator, is_neuronx_distributed_available, saved_model_in_temporary_directory, ) from .config import ( AudioNeuronConfig, TextAndVisionNeuronConfig, TextEncoderNeuronConfig, TextSeq2SeqNeuronConfig, VisionNeuronConfig, ) from .model_wrappers import ( CLIPVisionWithProjectionNeuronWrapper, ControlNetNeuronWrapper, NoCacheModelWrapper, PixartTransformerNeuronWrapper, SentenceTransformersCLIPNeuronWrapper, SentenceTransformersTransformerNeuronWrapper, T5DecoderWrapper, T5EncoderForSeq2SeqLMWrapper, T5EncoderWrapper, UnetNeuronWrapper, WhisperDecoderWrapper, WhisperEncoderWrapper, ) if is_neuronx_distributed_available(): import neuronx_distributed if TYPE_CHECKING: if is_diffusers_available(): from diffusers.models.vae import Decoder as VaeDecoder COMMON_TEXT_TASKS = [ "feature-extraction", "fill-mask", "multiple-choice", "question-answering", "text-classification", "token-classification", ] register_in_tasks_manager = TasksManager.create_register("neuron") @register_in_tasks_manager("bert", *COMMON_TEXT_TASKS) class BertNeuronConfig(TextEncoderNeuronConfig): NORMALIZED_CONFIG_CLASS = NormalizedConfigManager.get_normalized_config_class("bert") ATOL_FOR_VALIDATION = 1e-3 @property def inputs(self) -> List[str]: return ["input_ids", "attention_mask", "token_type_ids"] @register_in_tasks_manager("albert", *COMMON_TEXT_TASKS) class AlbertNeuronConfig(BertNeuronConfig): pass @register_in_tasks_manager("convbert", *COMMON_TEXT_TASKS) class ConvBertNeuronConfig(BertNeuronConfig): ATOL_FOR_VALIDATION = 1e-1 # TODO: why accuracy more off than other arch @property def outputs(self) -> List[str]: if self.task == "feature-extraction": return ["last_hidden_state"] return self._TASK_TO_COMMON_OUTPUTS[self.task] @register_in_tasks_manager("electra", *COMMON_TEXT_TASKS) class ElectraNeuronConfig(BertNeuronConfig): @property def outputs(self) -> List[str]: if self.task == "feature-extraction": return ["last_hidden_state"] return self._TASK_TO_COMMON_OUTPUTS[self.task] @register_in_tasks_manager("esm", *["feature-extraction", "fill-mask", "text-classification", "token-classification"]) class EsmNeuronConfig(TextEncoderNeuronConfig): NORMALIZED_CONFIG_CLASS = NormalizedConfigManager.get_normalized_config_class("bert") ATOL_FOR_VALIDATION = 1e-3 @property def inputs(self) -> List[str]: return ["input_ids", "attention_mask"] @register_in_tasks_manager("flaubert", *COMMON_TEXT_TASKS) class FlaubertNeuronConfig(ElectraNeuronConfig): pass @register_in_tasks_manager("mobilebert", *COMMON_TEXT_TASKS) class MobileBertNeuronConfig(BertNeuronConfig): pass @register_in_tasks_manager( "modernbert", *["feature-extraction", "fill-mask", "text-classification", "token-classification"] ) class ModernBertNeuronConfig(BertNeuronConfig): @property def inputs(self) -> List[str]: return ["input_ids", "attention_mask"] @property def outputs(self) -> List[str]: if self.task == "feature-extraction": return ["last_hidden_state"] return self._TASK_TO_COMMON_OUTPUTS[self.task] @register_in_tasks_manager("phi", *["feature-extraction", "text-classification", "token-classification"]) class PhiNeuronConfig(ElectraNeuronConfig): CUSTOM_MODEL_WRAPPER = NoCacheModelWrapper @property def inputs(self) -> List[str]: return ["input_ids", "attention_mask"] @register_in_tasks_manager("roformer", *COMMON_TEXT_TASKS) class RoFormerNeuronConfig(ElectraNeuronConfig): pass @register_in_tasks_manager("xlm", *COMMON_TEXT_TASKS) class XLMNeuronConfig(ElectraNeuronConfig): pass @register_in_tasks_manager("distilbert", *COMMON_TEXT_TASKS) class DistilBertNeuronConfig(BertNeuronConfig): ATOL_FOR_VALIDATION = 1e-3 @property def inputs(self) -> List[str]: return ["input_ids", "attention_mask"] @property def outputs(self) -> List[str]: if self.task == "feature-extraction": return ["last_hidden_state"] return self._TASK_TO_COMMON_OUTPUTS[self.task] @register_in_tasks_manager("camembert", *COMMON_TEXT_TASKS) class CamembertNeuronConfig(BertNeuronConfig): ATOL_FOR_VALIDATION = 1e-3 @property def inputs(self) -> List[str]: return ["input_ids", "attention_mask"] @register_in_tasks_manager("mpnet", *COMMON_TEXT_TASKS) class MPNetNeuronConfig(CamembertNeuronConfig): pass @register_in_tasks_manager("roberta", *COMMON_TEXT_TASKS) class RobertaNeuronConfig(CamembertNeuronConfig): pass @register_in_tasks_manager("xlm-roberta", *COMMON_TEXT_TASKS) class XLMRobertaNeuronConfig(CamembertNeuronConfig): pass # https://github.com/aws-neuron/aws-neuron-sdk/issues/642 # Failed only for INF1: 'XSoftmax' @register_in_tasks_manager("deberta", *([task for task in COMMON_TEXT_TASKS if task != "multiple-choice"])) class DebertaNeuronConfig(ElectraNeuronConfig): @property def inputs(self) -> List[str]: common_inputs = super().inputs if self._config.type_vocab_size == 0: # We remove token type ids. common_inputs.pop(-1) return common_inputs # https://github.com/aws-neuron/aws-neuron-sdk/issues/642 # Failed only for INF1: 'XSoftmax' @register_in_tasks_manager("deberta-v2", *([task for task in COMMON_TEXT_TASKS if task != "multiple-choice"])) class DebertaV2NeuronConfig(ElectraNeuronConfig): pass @register_in_tasks_manager( "transformer", *["feature-extraction", "sentence-similarity"], library_name="sentence_transformers" ) class SentenceTransformersTransformerNeuronConfig(TextEncoderNeuronConfig): NORMALIZED_CONFIG_CLASS = NormalizedTextConfig CUSTOM_MODEL_WRAPPER = SentenceTransformersTransformerNeuronWrapper ATOL_FOR_VALIDATION = 1e-3 @property def inputs(self) -> List[str]: return ["input_ids", "attention_mask"] @property def outputs(self) -> List[str]: return ["token_embeddings", "sentence_embedding"] class CLIPNormalizedConfig(NormalizedTextAndVisionConfig): TEXT_CONFIG = "text_config" VISION_CONFIG = "vision_config" @register_in_tasks_manager("clip-vision-with-projection", *["feature-extraction"], library_name="diffusers") class CLIPVisionWithProjectionNeuronConfig(VisionNeuronConfig): MODEL_TYPE = "clip-vision-with-projection" NORMALIZED_CONFIG_CLASS = NormalizedVisionConfig CUSTOM_MODEL_WRAPPER = CLIPVisionWithProjectionNeuronWrapper @property def inputs(self) -> List[str]: return ["pixel_values"] @property def outputs(self) -> List[str]: common_outputs = ["image_embeds", "last_hidden_state"] if self.output_hidden_states: common_outputs.append("hidden_states") return common_outputs @register_in_tasks_manager("clip", *["feature-extraction", "zero-shot-image-classification", "image-classification"]) class CLIPNeuronConfig(TextAndVisionNeuronConfig): NORMALIZED_CONFIG_CLASS = CLIPNormalizedConfig INPUT_ARGS = ("text_batch_size", "image_batch_size", "sequence_length", "num_channels", "width", "height") @property def inputs(self) -> List[str]: if self.task == "image-classification": return ["pixel_values"] else: return ["input_ids", "pixel_values", "attention_mask"] @property def outputs(self) -> List[str]: if self.task == "image-classification": return ["logits"] else: return [ "logits_per_image", "logits_per_text", "text_embeds", "image_embeds", "text_model_output", "vision_model_output", ] def _create_dummy_input_generator_classes(self, **kwargs) -> List["DummyInputGenerator"]: for name, axis_dim in self._axes.items(): self._axes[name] = kwargs.pop(name, axis_dim) self._validate_mandatory_axes() other_axes = copy.deepcopy(self._axes) text_batch_size = other_axes.pop("text_batch_size") images_batch_size = other_axes.pop("image_batch_size") return [ DummyTextInputGenerator(self.task, self._normalized_config, batch_size=text_batch_size, **other_axes), DummyVisionInputGenerator(self.task, self._normalized_config, batch_size=images_batch_size, **other_axes), ] @register_in_tasks_manager("clip-text-with-projection", *["feature-extraction"], library_name="diffusers") class CLIPTextWithProjectionNeuronConfig(TextEncoderNeuronConfig): MODEL_TYPE = "clip-text-with-projection" ATOL_FOR_VALIDATION = 1e-3 NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( vocab_size="vocab_size", sequence_length="max_position_embeddings", num_layers="num_hidden_layers", allow_new=True, ) @property def inputs(self) -> List[str]: return ["input_ids"] @property def outputs(self) -> List[str]: common_outputs = ["text_embeds", "last_hidden_state"] if self._normalized_config.output_hidden_states: common_outputs.append("hidden_states") return common_outputs @register_in_tasks_manager("clip-text-model", *["feature-extraction"], library_name="diffusers") class CLIPTextNeuronConfig(CLIPTextWithProjectionNeuronConfig): MODEL_TYPE = "clip-text-model" @property def outputs(self) -> List[str]: common_outputs = ["last_hidden_state", "pooler_output"] if self._normalized_config.output_hidden_states: common_outputs.append("hidden_states") return common_outputs # TODO: We should decouple clip text and vision, this would need fix on Optimum main. For the current workaround # users can pass dummy text inputs when encoding image, vice versa. @register_in_tasks_manager( "clip", *["feature-extraction", "sentence-similarity"], library_name="sentence_transformers" ) class SentenceTransformersCLIPNeuronConfig(CLIPNeuronConfig): CUSTOM_MODEL_WRAPPER = SentenceTransformersCLIPNeuronWrapper ATOL_FOR_VALIDATION = 1e-3 @property def outputs(self) -> List[str]: return ["text_embeds", "image_embeds"] @register_in_tasks_manager("vit", *["feature-extraction", "image-classification"]) class ViTNeuronConfig(VisionNeuronConfig): ATOL_FOR_VALIDATION = 1e-3 NORMALIZED_CONFIG_CLASS = NormalizedVisionConfig DUMMY_INPUT_GENERATOR_CLASSES = (DummyVisionInputGenerator, DummyMaskedPosGenerator) INPUT_ARGS = ("batch_size",) # `num_channels` and `image_size` are inferred from the config @property def inputs(self) -> List[str]: common_inputs = ["pixel_values"] if self.task == "masked-im": common_inputs.append("bool_masked_pos") return common_inputs @register_in_tasks_manager("beit", *["feature-extraction", "image-classification"]) class BeitNeuronConfig(ViTNeuronConfig): pass @register_in_tasks_manager("convnext", *["feature-extraction", "image-classification"]) class ConvNextNeuronConfig(ViTNeuronConfig): pass @register_in_tasks_manager("convnextv2", *["feature-extraction", "image-classification"]) class ConvNextV2NeuronConfig(ViTNeuronConfig): pass @register_in_tasks_manager("cvt", *["feature-extraction", "image-classification"]) class CvTNeuronConfig(ViTNeuronConfig): MODEL_TYPE = "cvt" @property def outputs(self) -> List[str]: common_outputs = super().outputs if self.task == "feature-extraction": return ["last_hidden_state", "cls_token_value"] else: return common_outputs @register_in_tasks_manager("deit", *["feature-extraction", "image-classification"]) class DeiTNeuronConfig(ViTNeuronConfig): pass @register_in_tasks_manager("donut-swin", *["feature-extraction"]) class DonutSwinNeuronConfig(ViTNeuronConfig): pass @register_in_tasks_manager("dpt", *["feature-extraction"]) class DptNeuronConfig(ViTNeuronConfig): pass @register_in_tasks_manager("levit", *["feature-extraction", "image-classification"]) class LevitNeuronConfig(ViTNeuronConfig): MODEL_TYPE = "levit" pass @register_in_tasks_manager( "mobilenet-v2", *["feature-extraction", "image-classification", "semantic-segmentation", "image-segmentation"] ) class MobileNetV2NeuronConfig(ViTNeuronConfig): MODEL_TYPE = "mobilenet-v2" pass @register_in_tasks_manager( "mobilevit", *["feature-extraction", "image-classification", "semantic-segmentation", "image-segmentation"] ) class MobileViTNeuronConfig(ViTNeuronConfig): MODEL_TYPE = "mobilevit" pass @register_in_tasks_manager("swin", *["feature-extraction", "image-classification"]) class SwinNeuronConfig(ViTNeuronConfig): pass @register_in_tasks_manager("yolos", *["feature-extraction", "object-detection"]) class YolosTNeuronConfig(ViTNeuronConfig): @property def outputs(self) -> List[str]: common_outputs = super().outputs if self.task == "object-detection": common_outputs.append("last_hidden_state") return common_outputs @register_in_tasks_manager( "wav2vec2", *[ "feature-extraction", "automatic-speech-recognition", "audio-classification", "audio-frame-classification", "audio-xvector", ], ) class Wav2Vec2NeuronConfig(AudioNeuronConfig): NORMALIZED_CONFIG_CLASS = NormalizedConfig MODEL_TYPE = "wav2vec2" @property def inputs(self) -> List[str]: return ["input_values"] @property def outputs(self) -> List[str]: common_outputs = super().outputs if self.task == "feature-extraction": common_outputs = ["last_hidden_state", "extract_features"] if self.task == "audio-xvector": common_outputs.append("embeddings") return common_outputs @register_in_tasks_manager( "audio-spectrogram-transformer", *[ "feature-extraction", "audio-classification", ], ) class ASTNeuronConfig(AudioNeuronConfig): NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( num_mel_bins="num_mel_bins", max_length="max_length", allow_new=True ) DUMMY_INPUT_GENERATOR_CLASSES = (ASTDummyAudioInputGenerator,) @property def inputs(self) -> List[str]: return ["input_values"] @register_in_tasks_manager( "hubert", *[ "feature-extraction", "automatic-speech-recognition", "audio-classification", ], ) class HubertNeuronConfig(Wav2Vec2NeuronConfig): MODEL_TYPE = "hubert" @property def outputs(self) -> List[str]: common_outputs = super().outputs if self.task == "feature-extraction": common_outputs = ["last_hidden_state"] return common_outputs # TODO: compilation failed due to a bug in xla: https://github.com/pytorch/xla/issues/6398. # @register_in_tasks_manager( # "sew", # *[ # "feature-extraction", # "automatic-speech-recognition", # "audio-classification", # ], # ) # class SEWNeuronConfig(Wav2Vec2NeuronConfig): # pass # TODO: compilation failed due to a bug in xla: https://github.com/pytorch/xla/issues/6398. # @register_in_tasks_manager( # "sew-d", # *[ # "feature-extraction", # "automatic-speech-recognition", # "audio-classification", # ], # ) # class SEWDNeuronConfig(Wav2Vec2NeuronConfig): # pass @register_in_tasks_manager( "unispeech", *[ "feature-extraction", "automatic-speech-recognition", "audio-classification", ], ) class UniSpeechNeuronConfig(Wav2Vec2NeuronConfig): MODEL_TYPE = "unispeech" pass @register_in_tasks_manager( "unispeech-sat", *[ "feature-extraction", "automatic-speech-recognition", "audio-classification", "audio-frame-classification", "audio-xvector", ], ) class UniSpeechSATNeuronConfig(Wav2Vec2NeuronConfig): MODEL_TYPE = "unispeech-sat" pass # TODO: compilation failed due to a bug in xla: https://github.com/pytorch/xla/issues/6398. # @register_in_tasks_manager( # "wav2vec2-bert", # *[ # "feature-extraction", # "automatic-speech-recognition", # "audio-classification", # "audio-frame-classification", # "audio-xvector", # ], # ) # class Wav2Vec2BertNeuronConfig(Wav2Vec2NeuronConfig): # pass # TODO: compilation failed due to a bug in xla: https://github.com/pytorch/xla/issues/6398. # @register_in_tasks_manager( # "wav2vec2-conformer", # *[ # "feature-extraction", # "automatic-speech-recognition", # "audio-classification", # "audio-frame-classification", # "audio-xvector", # ], # ) # class Wav2Vec2ConformerNeuronConfig(Wav2Vec2NeuronConfig): # pass @register_in_tasks_manager( "wavlm", *[ "feature-extraction", "automatic-speech-recognition", "audio-classification", "audio-frame-classification", "audio-xvector", ], ) class WavLMNeuronConfig(Wav2Vec2NeuronConfig): MODEL_TYPE = "wavlm" pass @register_in_tasks_manager("unet", *["semantic-segmentation"], library_name="diffusers") class UNetNeuronConfig(VisionNeuronConfig): ATOL_FOR_VALIDATION = 1e-3 INPUT_ARGS = ("batch_size", "sequence_length", "num_channels", "width", "height", "vae_scale_factor") MODEL_TYPE = "unet" CUSTOM_MODEL_WRAPPER = UnetNeuronWrapper NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( height="height", width="width", num_channels="in_channels", hidden_size="cross_attention_dim", vocab_size="norm_num_groups", allow_new=True, ) DUMMY_INPUT_GENERATOR_CLASSES = ( DummyVisionInputGenerator, DummyTimestepInputGenerator, DummySeq2SeqDecoderTextInputGenerator, DummyControNetInputGenerator, DummyIPAdapterInputGenerator, ) @property def inputs(self) -> List[str]: common_inputs = ["sample", "timestep", "encoder_hidden_states"] # TODO : add text_image, image and image_embeds if getattr(self._normalized_config, "addition_embed_type", None) == "text_time": common_inputs.append("text_embeds") common_inputs.append("time_ids") if getattr(self._normalized_config, "time_cond_proj_dim", None) is not None: common_inputs.append("timestep_cond") if self.with_controlnet: # outputs of controlnet common_inputs += ["down_block_additional_residuals", "mid_block_additional_residual"] if self.with_ip_adapter: # add output of image encoder if self.image_encoder_output_hidden_states: common_inputs += ["image_enc_hidden_states"] else: common_inputs += ["image_embeds"] return common_inputs @property def outputs(self) -> List[str]: return ["sample"] def generate_dummy_inputs(self, return_tuple: bool = False, **kwargs): dummy_inputs = super().generate_dummy_inputs(**kwargs) dummy_inputs["timestep"] = dummy_inputs["timestep"].float() dummy_inputs["encoder_hidden_states"] = dummy_inputs["encoder_hidden_states"][0] # break down down_block_additional_residuals num_down_block_outputs = len(self._normalized_config.down_block_types) * ( self._normalized_config.layers_per_block + 1 ) down_block_additional_residuals = dummy_inputs.pop("down_block_additional_residuals", None) if down_block_additional_residuals: for idx in range(num_down_block_outputs): dummy_inputs[f"down_block_additional_residuals_{idx}"] = down_block_additional_residuals[idx] if getattr(self._normalized_config, "addition_embed_type", None) == "text_time": dummy_inputs["added_cond_kwargs"] = { "text_embeds": dummy_inputs.pop("text_embeds"), "time_ids": dummy_inputs.pop("time_ids"), } if return_tuple is True: return tuple(dummy_inputs.values()) else: return dummy_inputs @property def is_sdxl(self) -> bool: return self._is_sdxl @is_sdxl.setter def is_sdxl(self, is_sdxl: bool): self._is_sdxl = is_sdxl @property def with_controlnet(self) -> bool: return self._with_controlnet @with_controlnet.setter def with_controlnet(self, with_controlnet: bool): self._with_controlnet = with_controlnet @property def with_ip_adapter(self) -> bool: return self._with_ip_adapter @with_ip_adapter.setter def with_ip_adapter(self, with_ip_adapter: bool): self._with_ip_adapter = with_ip_adapter if with_ip_adapter: self.mandatory_axes += ("image_encoder_shapes",) setattr(self, "image_encoder_shapes", self.input_shapes["image_encoder_shapes"]) @register_in_tasks_manager("pixart-transformer-2d", *["semantic-segmentation"], library_name="diffusers") class PixartTransformerNeuronConfig(VisionNeuronConfig): ATOL_FOR_VALIDATION = 1e-3 INPUT_ARGS = ( "batch_size", "sequence_length", "num_channels", "width", "height", "vae_scale_factor", "encoder_hidden_size", ) MODEL_TYPE = "pixart-transformer-2d" CUSTOM_MODEL_WRAPPER = PixartTransformerNeuronWrapper NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( height="height", width="width", num_channels="in_channels", hidden_size="cross_attention_dim", vocab_size="norm_num_groups", allow_new=True, ) DUMMY_INPUT_GENERATOR_CLASSES = ( DummyVisionInputGenerator, DummyControNetInputGenerator, DummyTextInputGenerator, DummySeq2SeqDecoderTextInputGenerator, ) @property def inputs(self) -> List[str]: common_inputs = ["sample", "encoder_hidden_states", "timestep", "encoder_attention_mask"] return common_inputs @property def outputs(self) -> List[str]: return ["out_hidden_states"] @register_in_tasks_manager("controlnet", *["semantic-segmentation"], library_name="diffusers") class ControlNetNeuronConfig(VisionNeuronConfig): ATOL_FOR_VALIDATION = 1e-3 INPUT_ARGS = ( "batch_size", "sequence_length", "num_channels", "height", "width", "vae_scale_factor", "encoder_hidden_size", ) MODEL_TYPE = "controlnet" CUSTOM_MODEL_WRAPPER = ControlNetNeuronWrapper NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( height="height", width="width", num_channels="in_channels", hidden_size="cross_attention_dim", vocab_size="norm_num_groups", allow_new=True, ) DUMMY_INPUT_GENERATOR_CLASSES = ( DummyVisionInputGenerator, DummyControNetInputGenerator, # Instead of `encoder_hidden_states` generated by `DummySeq2SeqDecoderTextInputGenerator` DummyTimestepInputGenerator, DummySeq2SeqDecoderTextInputGenerator, ) @property def inputs(self) -> List[str]: common_inputs = ["sample", "timestep", "encoder_hidden_states", "controlnet_cond", "conditioning_scale"] if getattr(self._normalized_config, "addition_embed_type", None) == "text_time": common_inputs.append("text_embeds") common_inputs.append("time_ids") return common_inputs @property def outputs(self) -> List[str]: return ["down_block_res_samples", "mid_block_res_sample"] @register_in_tasks_manager("vae-encoder", *["semantic-segmentation"], library_name="diffusers") class VaeEncoderNeuronConfig(VisionNeuronConfig): ATOL_FOR_VALIDATION = 1e-3 MODEL_TYPE = "vae-encoder" NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( num_channels="in_channels", allow_new=True, ) @property def inputs(self) -> List[str]: return ["sample"] @property def outputs(self) -> List[str]: return ["latent_parameters"] def generate_dummy_inputs(self, return_tuple: bool = False, **kwargs): dummy_inputs = super().generate_dummy_inputs(**kwargs) if return_tuple is True: return tuple(dummy_inputs.values()) else: return dummy_inputs @register_in_tasks_manager("vae-decoder", *["semantic-segmentation"], library_name="diffusers") class VaeDecoderNeuronConfig(VisionNeuronConfig): ATOL_FOR_VALIDATION = 1e-3 MODEL_TYPE = "vae-decoder" NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( num_channels="latent_channels", allow_new=True, ) @property def inputs(self) -> List[str]: return ["latent_sample"] @property def outputs(self) -> List[str]: return ["sample"] def patch_model_for_export( self, model: "VaeDecoder", dummy_inputs: Dict[str, torch.Tensor], **kwargs, ): return super().patch_model_for_export(model=model, dummy_inputs=dummy_inputs, forward_with_tuple=True) class T5EncoderBaseNeuronConfig(TextSeq2SeqNeuronConfig): ATOL_FOR_VALIDATION = 1e-3 NORMALIZED_CONFIG_CLASS = NormalizedSeq2SeqConfig.with_args( hidden_size="d_model", num_attention_heads="num_heads", encoder_num_layers="num_layers", decoder_num_layers="num_decoder_layers", key_value_dim="d_kv", allow_new=True, ) @property def inputs(self) -> List[str]: return ["input_ids", "attention_mask"] @register_in_tasks_manager("t5", *["feature-extraction"], library_name="diffusers") class T5EncoderForDiffusersNeuronConfig(T5EncoderBaseNeuronConfig): CUSTOM_MODEL_WRAPPER = T5EncoderWrapper INPUT_ARGS = ("batch_size", "sequence_length") @property def outputs(self) -> List[str]: return ["last_hidden_state"] @property def is_encoder_decoder(self) -> bool: return True def patch_model_for_export(self, model_or_path, **input_shapes): return self.CUSTOM_MODEL_WRAPPER(model_or_path, **input_shapes) @register_in_tasks_manager("t5-encoder", *["text2text-generation"]) class T5EncoderForTransformersNeuronConfig(T5EncoderBaseNeuronConfig): CUSTOM_MODEL_WRAPPER = T5EncoderForSeq2SeqLMWrapper INPUT_ARGS = ("batch_size", "sequence_length", "num_beams") MODEL_TYPE = "t5-encoder" @property def outputs(self) -> List[str]: common_outputs = ( [f"present.{idx}.self.key" for idx in range(self._config.num_decoder_layers)] + [f"present.{idx}.self.value" for idx in range(self._config.num_decoder_layers)] + [f"present.{idx}.cross.key" for idx in range(self._config.num_decoder_layers)] + [f"present.{idx}.cross.value" for idx in range(self._config.num_decoder_layers)] ) return common_outputs @property def is_encoder_decoder(self) -> bool: return True def patch_model_for_export(self, model_or_path, device="xla", **kwargs): num_beams = kwargs.pop("num_beams", 1) sequence_length = kwargs.pop("sequence_length", None) batch_size = kwargs.pop("batch_size", None) if self.tensor_parallel_size > 1: # `torch.nn.modules` objects not eligible for pickling, the model needs to be loaded within the func. return partial( self.get_parallel_callable, model_or_path, sequence_length, batch_size, num_beams, device, self.tensor_parallel_size, ) else: return self.CUSTOM_MODEL_WRAPPER( model_or_path, sequence_length=sequence_length, batch_size=batch_size, num_beams=num_beams, device=device, tensor_parallel_size=self.tensor_parallel_size, ) def get_parallel_callable( self, model_name_or_path, sequence_length, batch_size, num_beams, device, tensor_parallel_size ): """Unlike `torch_neuronx.trace`, `parallel_model_trace` requires a function returning a model object and a dictionary of states.""" model = TasksManager.get_model_from_task( model_name_or_path=model_name_or_path, task=self.task, framework="pt", library_name="transformers", ) # TODO: add extra args, eg. revision, trust_remote_code, etc. model.config.use_cache = True with saved_model_in_temporary_directory(model) as ckpt_path: # Plug in parallel layers from optimum.neuron.models.inference.t5.modeling_t5 import parallelize parallel_model = parallelize(model) # Load the weights into the parallel layers neuronx_distributed.parallel_layers.load(ckpt_path, parallel_model, sharded=False) encoder = self.CUSTOM_MODEL_WRAPPER( parallel_model, sequence_length=sequence_length, batch_size=batch_size, num_beams=num_beams, device=device, tensor_parallel_size=tensor_parallel_size, ) encoder.eval() aliases = self.generate_io_aliases(encoder) return encoder, aliases def generate_io_aliases(self, encoder=None): aliases = {} if self.tensor_parallel_size > 1: for i in range(len(encoder.past_key_values_sa)): aliases[encoder.past_key_values_sa[i]] = i for i in range(len(encoder.past_key_values_ca)): aliases[encoder.past_key_values_ca[i]] = len(encoder.past_key_values_sa) + i return aliases @register_in_tasks_manager("t5-decoder", "text2text-generation") class T5DecoderNeuronConfig(TextSeq2SeqNeuronConfig): ATOL_FOR_VALIDATION = 1e-3 DUMMY_INPUT_GENERATOR_CLASSES = TextSeq2SeqNeuronConfig.DUMMY_INPUT_GENERATOR_CLASSES + (DummyBeamValuesGenerator,) INPUT_ARGS = ("batch_size", "sequence_length", "num_beams") MODEL_TYPE = "t5-decoder" CUSTOM_MODEL_WRAPPER = T5DecoderWrapper NORMALIZED_CONFIG_CLASS = NormalizedSeq2SeqConfig @property def inputs(self) -> List[str]: common_inputs = [ "decoder_input_ids", "decoder_attention_mask", "encoder_hidden_states", "attention_mask", # TODO: replace with `encoder_attention_mask` after optimum 1.14 release "beam_idx", "beam_scores", ] return common_inputs @property def outputs(self) -> List[str]: beam_outputs = ["next_token_scores", "next_tokens", "next_indices"] if self.num_beams > 1 else ["next_tokens"] common_outputs = ( beam_outputs + [f"past.{idx}.self.key" for idx in range(self._config.num_decoder_layers)] + [f"past.{idx}.self.value" for idx in range(self._config.num_decoder_layers)] + [f"past.{idx}.cross.key" for idx in range(self._config.num_decoder_layers)] + [f"past.{idx}.cross.value" for idx in range(self._config.num_decoder_layers)] ) if self.output_hidden_states: # Flatten hidden states of all layers common_outputs += [ f"decoder_hidden_state.{idx}" for idx in range(self._config.num_decoder_layers + 1) ] # +1 for the embedding layer if self.output_attentions: # Flatten attentions tensors of all attention layers common_outputs += [f"decoder_attention.{idx}" for idx in range(self._config.num_decoder_layers)] if getattr(self._config, "is_encoder_decoder", False) is True: common_outputs += [f"cross_attention.{idx}" for idx in range(self._config.num_decoder_layers)] return common_outputs @property def is_encoder_decoder(self) -> bool: return True def generate_dummy_inputs(self, **kwargs): batch_size = kwargs.pop("batch_size") * kwargs.get("num_beams") dummy_inputs = super().generate_dummy_inputs(batch_size=batch_size, **kwargs) dummy_inputs["decoder_input_ids"] = dummy_inputs["decoder_input_ids"][:, :1] # sequence_length = 1 dummy_inputs["encoder_hidden_states"] = dummy_inputs["encoder_hidden_states"][0] return dummy_inputs def _create_dummy_input_generator_classes(self, **kwargs) -> List["DummyInputGenerator"]: dummy_inputs_generators = super()._create_dummy_input_generator_classes(**kwargs) dummy_beam_values_generator = self.DUMMY_INPUT_GENERATOR_CLASSES[-1]( self.task, self._normalized_config, num_beams=kwargs.pop("num_beams", 1), **kwargs, ) dummy_inputs_generators.append(dummy_beam_values_generator) return dummy_inputs_generators def patch_model_for_export(self, model, device="xla", **kwargs): batch_size = kwargs.pop("batch_size", 1) sequence_length = kwargs.pop("sequence_length", 1) num_beams = kwargs.pop("num_beams", 1) trace_args = { "model": model, "batch_size": batch_size, "sequence_length": sequence_length, "num_beams": num_beams, "output_hidden_states": self.output_hidden_states, "output_attentions": self.output_attentions, "device": device, "tensor_parallel_size": self.tensor_parallel_size, } if self.tensor_parallel_size > 1: return partial( self.get_parallel_callable, model, batch_size, sequence_length, num_beams, self.output_hidden_states, self.output_attentions, device, self.tensor_parallel_size, ) else: return self.CUSTOM_MODEL_WRAPPER(**trace_args) def get_parallel_callable( self, model_name_or_path, batch_size, sequence_length, num_beams, output_hidden_states, output_attentions, device, tensor_parallel_size, ): """Unlike `torch_neuronx.trace`, `parallel_model_trace` requires a function returning a model object and a dictionary of states.""" model = TasksManager.get_model_from_task( model_name_or_path=model_name_or_path, task=self.task, framework="pt", library_name="transformers", ) # TODO: add extra args, eg. revision, trust_remote_code, etc. model.config.use_cache = True with saved_model_in_temporary_directory(model) as ckpt_path: # Plug in parallel layers from optimum.neuron.models.inference.t5.modeling_t5 import parallelize parallel_model = parallelize(model) # Load the weights into the parallel layers neuronx_distributed.parallel_layers.load(ckpt_path, parallel_model, sharded=False) decoder = self.CUSTOM_MODEL_WRAPPER( parallel_model, batch_size=batch_size, sequence_length=sequence_length, num_beams=num_beams, output_hidden_states=output_hidden_states, output_attentions=output_attentions, device=device, tensor_parallel_size=tensor_parallel_size, ) decoder.eval() aliases = self.generate_io_aliases(decoder) return decoder, aliases def generate_io_aliases(self, decoder): num_outputs_from_trace = 3 if decoder.num_beams > 1 else 1 aliases = {} for i in range(len(decoder.past_key_values_sa)): aliases[decoder.past_key_values_sa[i]] = i + num_outputs_from_trace for i in range(len(decoder.past_key_values_ca)): aliases[decoder.past_key_values_ca[i]] = len(decoder.past_key_values_sa) + i + num_outputs_from_trace return aliases @register_in_tasks_manager("whisper-encoder", *["automatic-speech-recognition"]) class WhisperEncoderNeuronConfig(AudioNeuronConfig): ATOL_FOR_VALIDATION = 1e-3 MODEL_TYPE = "whisper-encoder" CUSTOM_MODEL_WRAPPER = WhisperEncoderWrapper INPUT_ARGS = ("batch_size", "sequence_length") DUMMY_INPUT_GENERATOR_CLASSES = AudioNeuronConfig.DUMMY_INPUT_GENERATOR_CLASSES + (WhisperDummyTextInputGenerator,) NORMALIZED_CONFIG_CLASS = NormalizedSeq2SeqConfig.with_args( encoder_num_layers="encoder_layers", decoder_num_layers="decoder_layers", feature_size="num_mel_bins", allow_new=True, ) @property def inputs(self) -> List[str]: return ["input_features", "decoder_input_ids"] @property def outputs(self) -> List[str]: return ["lm_logits", "encoder_last_hidden_state"] @property def is_encoder_decoder(self) -> bool: return True def generate_dummy_inputs(self, return_tuple: bool = False, **kwargs): kwargs["sequence_length"] = 1 # only `decoder_start_token_id` return super().generate_dummy_inputs(return_tuple=return_tuple, **kwargs) def patch_model_for_export(self, model_or_path, **input_shapes): return self.CUSTOM_MODEL_WRAPPER(model_or_path, **input_shapes) @register_in_tasks_manager("whisper-decoder", *["automatic-speech-recognition"]) class WhisperDecoderNeuronConfig(AudioNeuronConfig): ATOL_FOR_VALIDATION = 1e-3 MODEL_TYPE = "whisper-decoder" DUMMY_INPUT_GENERATOR_CLASSES = (WhisperDummyTextInputGenerator,) INPUT_ARGS = ("batch_size", "sequence_length") CUSTOM_MODEL_WRAPPER = WhisperDecoderWrapper NORMALIZED_CONFIG_CLASS = NormalizedSeq2SeqConfig.with_args( encoder_num_layers="encoder_layers", decoder_num_layers="decoder_layers", feature_size="num_mel_bins", hidden_size="d_model", allow_new=True, ) @property def inputs(self) -> List[str]: return ["decoder_input_ids", "encoder_hidden_states"] @property def outputs(self) -> List[str]: return ["lm_logits"] @property def is_encoder_decoder(self) -> bool: return True def patch_model_for_export(self, model_or_path, **input_shapes): return self.CUSTOM_MODEL_WRAPPER(model_or_path, **input_shapes)