Optimisez les performances de diffusion stable et réduisez les coûts d'inférence avec AWS Inferentia2 | Services Web Amazon

Les modèles d'IA générative ont connu une croissance rapide ces derniers mois en raison de leurs capacités impressionnantes de création de texte, d'images, de code et d'audio réalistes. Parmi ces modèles, les modèles Stable Diffusion se distinguent par leur force unique dans la création d'images de haute qualité basées sur des invites textuelles. La diffusion stable peut générer une grande variété d'images de haute qualité, y compris des portraits réalistes, des paysages et même de l'art abstrait. Et, comme d'autres modèles d'IA générative, les modèles de diffusion stable nécessitent un calcul puissant pour fournir une inférence à faible latence.

Dans cet article, nous montrons comment vous pouvez exécuter des modèles de diffusion stable et obtenir des performances élevées au moindre coût dans Cloud de calcul élastique Amazon (Amazon EC2) en utilisant Instances Amazon EC2 Inf2 grâce à Inférence AWS2. Nous examinons l'architecture d'un modèle de diffusion stable et parcourons les étapes de compilation d'un modèle de diffusion stable à l'aide de Neurone AWS et en le déployant sur une instance Inf2. Nous discutons également des optimisations que le SDK Neuron effectue automatiquement pour améliorer les performances. Vous pouvez exécuter les versions Stable Diffusion 2.1 et 1.5 sur AWS Inferentia2 de manière rentable. Enfin, nous montrons comment vous pouvez déployer un modèle de diffusion stable sur une instance Inf2 avec Amazon Sage Maker.

La taille du modèle Stable Diffusion 2.1 en virgule flottante 32 (FP32) est de 5 Go et de 2.5 Go en bfoat16 (BF16). Une seule instance inf2.xlarge possède un accélérateur AWS Inferentia2 avec 32 Go de mémoire HBM. Le modèle Stable Diffusion 2.1 peut tenir sur une seule instance inf2.xlarge. Stable Diffusion est un modèle de texte à image que vous pouvez utiliser pour créer des images de différents styles et contenus simplement en fournissant une invite de texte en entrée. Pour en savoir plus sur l'architecture du modèle de diffusion stable, reportez-vous à Créez des images de haute qualité avec des modèles de diffusion stable et déployez-les de manière rentable avec Amazon SageMaker.

Comment le SDK Neuron optimise les performances de diffusion stable

Avant de pouvoir déployer le modèle Stable Diffusion 2.1 sur les instances AWS Inferentia2, nous devons compiler les composants du modèle à l'aide du SDK Neuron. Le SDK Neuron, qui comprend un compilateur, un environnement d'exécution et des outils d'apprentissage en profondeur, compile et optimise automatiquement les modèles d'apprentissage en profondeur afin qu'ils puissent s'exécuter efficacement sur les instances Inf2 et extraire les performances complètes de l'accélérateur AWS Inferentia2. Nous avons des exemples disponibles pour le modèle Stable Diffusion 2.1 sur le GitHub repo. Ce bloc-notes présente un exemple de bout en bout de la façon de compiler un modèle de diffusion stable, d'enregistrer les modèles Neuron compilés et de le charger dans le runtime pour l'inférence.

Nous utilisons StableDiffusionPipeline du visage étreignant diffusers bibliothèque pour charger et compiler le modèle. Nous compilons ensuite tous les composants du modèle pour Neuron en utilisant torch_neuronx.trace() et enregistrez le modèle optimisé sous TorchScript. Les processus de compilation peuvent être assez gourmands en mémoire, nécessitant une quantité importante de RAM. Pour contourner cela, avant de tracer chaque modèle, nous créons un deepcopy de la partie du pipeline qui est tracée. Ensuite, nous supprimons l'objet pipeline de la mémoire en utilisant del pipe. Cette technique est particulièrement utile lors de la compilation sur des instances avec peu de RAM.

De plus, nous effectuons également des optimisations des modèles de diffusion stable. UNet détient l'aspect le plus intensif en termes de calcul de l'inférence. Le composant UNet fonctionne sur des tenseurs d'entrée qui ont une taille de lot de deux, générant un tenseur de sortie correspondant également avec une taille de lot de deux, pour produire une seule image. Les éléments de ces lots sont entièrement indépendants les uns des autres. Nous pouvons tirer parti de ce comportement pour obtenir une latence optimale en exécutant un lot sur chaque cœur Neuron. Nous compilons le UNet pour un lot (en utilisant des tenseurs d'entrée avec un lot), puis utilisons le torch_neuronx.DataParallel API pour charger ce modèle de lot unique sur chaque cœur. La sortie de cette API est un module transparent à deux lots : nous pouvons transmettre à l'UNet les entrées de deux lots, et une sortie à deux lots est renvoyée, mais en interne, les deux modèles à un seul lot s'exécutent sur les deux cœurs Neuron. . Cette stratégie optimise l'utilisation des ressources et réduit la latence.

Compiler et déployer un modèle de diffusion stable sur une instance Inf2 EC2

Pour compiler et déployer le modèle Stable Diffusion sur une instance Inf2 EC2, connectez-vous au Console de gestion AWS et créez une instance inf2.8xlarge. Notez qu'une instance inf2.8xlarge est requise uniquement pour la compilation du modèle car la compilation nécessite une mémoire hôte plus importante. Le modèle Stable Diffusion peut être hébergé sur une instance inf2.xlarge. Vous pouvez trouver la dernière AMI avec les bibliothèques Neuron en utilisant ce qui suit Interface de ligne de commande AWS (AWS CLI) commande :

aws ec2 describe-images --region us-east-1 --owners amazon --filters 'Name=name,Values=Deep Learning AMI Neuron PyTorch 1.13.? (Amazon Linux 2) ????????' 'Name=state,Values=available' --query 'reverse(sort_by(Images, &CreationDate))[:1].ImageId' --output text

Pour cet exemple, nous avons créé une instance EC2 à l'aide de Deep Learning AMI Neuron PyTorch 1.13 (Ubuntu 20.04). Vous pouvez ensuite créer un environnement de laboratoire JupyterLab en vous connectant à l'instance et en exécutant les étapes suivantes :

run source /opt/aws_neuron_venv_pytorch/bin/activate
pip install jupyterlab
jupyter-lab

Un notebook avec toutes les étapes de compilation et d'hébergement du modèle se trouve sur GitHub.

Regardons les étapes de compilation pour l'un des blocs d'encodeur de texte. D'autres blocs faisant partie du pipeline Stable Diffusion peuvent être compilés de la même manière.

La première étape consiste à charger le modèle pré-formé de Hugging Face. Le StableDiffusionPipeline.from_pretrained charge le modèle pré-formé dans notre objet pipeline, pipe. Nous créons ensuite un deepcopy de l'encodeur de texte de notre pipeline, le clonant efficacement. Le del pipe La commande est ensuite utilisée pour supprimer l'objet de pipeline d'origine, libérant ainsi la mémoire qu'il a consommée. Ici, nous quantifions le modèle aux poids BF16 :

model_id = "stabilityai/stable-diffusion-2-1-base"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16)
text_encoder = copy.deepcopy(pipe.text_encoder)
del pipe

Cette étape consiste à envelopper notre encodeur de texte avec le NeuronTextEncoder emballage. La sortie d'un module d'encodeur de texte compilé sera de dict. Nous le convertissons en un list tapez en utilisant ce wrapper :

text_encoder = NeuronTextEncoder(text_encoder)

Nous initialisons le tenseur PyTorch emb avec quelques valeurs. Le emb tenseur est utilisé comme exemple d'entrée pour le torch_neuronx.trace fonction. Cette fonction trace notre encodeur de texte et le compile dans un format optimisé pour Neuron. Le chemin du répertoire pour le modèle compilé est construit en joignant COMPILER_WORKDIR_ROOT avec le sous-répertoire text_encoder:

emb = torch.tensor([...])
text_encoder_neuron = torch_neuronx.trace(
        text_encoder.neuron_text_encoder,
        emb,
        compiler_workdir=os.path.join(COMPILER_WORKDIR_ROOT, 'text_encoder'),
        )

L'encodeur de texte compilé est enregistré à l'aide de torch.jit.save. Il est stocké sous le nom de fichier model.pt dans le text_encoder répertoire de l'espace de travail de notre compilateur :

text_encoder_filename = os.path.join(COMPILER_WORKDIR_ROOT, 'text_encoder/model.pt')
torch.jit.save(text_encoder_neuron, text_encoder_filename)

Le cahier comprend des étapes similaires pour compiler d'autres composants du modèle : UNet, décodeur VAE et VAE post_quant_conv. Après avoir compilé tous les modèles, vous pouvez charger et exécuter le modèle en suivant ces étapes :

1. Définissez les chemins pour les modèles compilés.
2. Charger un pré-formé StableDiffusionPipeline modèle, avec sa configuration spécifiée pour utiliser le type de données bfloat16.
3. Chargez le modèle UNet sur deux cœurs Neuron à l'aide du torch_neuronx.DataParallel API. Cela permet d'effectuer une inférence parallèle de données, ce qui peut considérablement accélérer les performances du modèle.
4. Chargez les parties restantes du modèle (text_encoder, decoderet une post_quant_conv) sur un seul noyau Neuron.

Vous pouvez ensuite exécuter le pipeline en fournissant du texte d'entrée sous forme d'invites. Voici quelques images générées par le modèle pour les invites :

• Portrait de renaud sechan, plume et encre, dessins au trait complexes, par craig mullins, ruan jia, kentaro miura, greg rutkowski, loundraw

• Portrait d'un ancien mineur de charbon au 19ème siècle, belle peinture, avec une peinture faciale très détaillée par greg rutkowski

• Un château au milieu d'une forêt

Héberger Stable Diffusion 2.1 sur AWS Inferentia2 et SageMaker

L'hébergement de modèles de diffusion stable avec SageMaker nécessite également une compilation avec le SDK Neuron. Vous pouvez effectuer la compilation à l'avance ou pendant l'exécution à l'aide de conteneurs Large Model Inference (LMI). La compilation à l'avance permet des temps de chargement de modèle plus rapides et est l'option préférée.

Les conteneurs SageMaker LMI offrent deux façons de déployer le modèle :

• Une option sans code où nous fournissons simplement un serving.properties fichier avec les configurations requises
• Apportez votre propre script d'inférence

Nous examinons les deux solutions et passons en revue les configurations et le script d'inférence (model.py). Dans cet article, nous démontrons le déploiement à l'aide d'un modèle précompilé stocké dans un Service de stockage simple Amazon (Amazon S3). Vous pouvez utiliser ce modèle précompilé pour vos déploiements.

Configurer le modèle avec un script fourni

Dans cette section, nous montrons comment configurer le conteneur LMI pour héberger les modèles de diffusion stable. Le notebook SD2.1 disponible sur GitHub. La première étape consiste à créer le package de configuration du modèle selon la structure de répertoire suivante. Notre objectif est d'utiliser les configurations de modèle minimales nécessaires pour héberger le modèle. La structure de répertoire nécessaire est la suivante :

<config-root-directory> / 
    ├── serving.properties
    │   
    └── model.py [OPTIONAL]

Ensuite, nous créons le servant.propriétés fichier avec les paramètres suivants :

%%writefile code_sd/serving.properties
engine=Python
option.entryPoint=djl_python.transformers-neuronx
option.use_stable_diffusion=True
option.model_id=s3url
option.tensor_parallel_degree=2
option.dtype=bf16

Les paramètres spécifient ce qui suit :

• option.model_id – Les conteneurs LMI utilisent s5cmd pour charger le modèle à partir de l'emplacement S3 et nous devons donc spécifier l'emplacement où se trouvent nos poids compilés.
• option.entryPoint – Pour utiliser les gestionnaires intégrés, nous spécifions la classe transformers-neuronx. Si vous avez un script d'inférence personnalisé, vous devez le fournir à la place.
• option.dtype – Cela spécifie de charger les poids dans une taille spécifique. Pour cet article, nous utilisons BF16, ce qui réduit encore nos besoins en mémoire par rapport au FP32 et réduit notre latence grâce à cela.
• option.tensor_parallel_degree – Ce paramètre spécifie le nombre d'accélérateurs que nous utilisons pour ce modèle. L'accélérateur de puce AWS Inferentia2 a deux cœurs Neuron et donc spécifier une valeur de 2 signifie que nous utilisons un accélérateur (deux cœurs). Cela signifie que nous pouvons désormais créer plusieurs nœuds de calcul pour augmenter le débit du point de terminaison.
• option.moteur - Ceci est défini sur Python pour indiquer que nous n'utiliserons pas d'autres compilateurs comme DeepSpeed ​​ou Faster Transformer pour cet hébergement.

Apportez votre propre scénario

Si vous souhaitez apporter votre propre script d'inférence personnalisé, vous devez supprimer le option.entryPoint De serving.properties. Dans ce cas, le conteneur LMI recherchera un model.py fichier au même emplacement que le serving.properties et l'utiliser pour exécuter l'inférence.

Créez votre propre script d'inférence (model.py)

La création de votre propre script d'inférence est relativement simple à l'aide du conteneur LMI. Le conteneur nécessite votre model.py fichier pour avoir une implémentation de la méthode suivante :

def handle(inputs: Input) which returns an object of type Outputs

Examinons quelques-uns des domaines critiques du cahier attaché, qui illustre la fonction apportez votre propre script.

Remplacez le cross_attention module avec la version optimisée :

# Replace original cross-attention module with custom cross-attention module for better performance
    CrossAttention.get_attention_scores = get_attention_scores
Load the compiled weights for the following
text_encoder_filename = os.path.join(COMPILER_WORKDIR_ROOT, 'text_encoder.pt')
decoder_filename = os.path.join(COMPILER_WORKDIR_ROOT, 'vae_decoder.pt')
unet_filename = os.path.join(COMPILER_WORKDIR_ROOT, 'unet.pt')
post_quant_conv_filename =. os.path.join(COMPILER_WORKDIR_ROOT, 'vae_post_quant_conv.pt')

Ce sont les noms du fichier de poids compilé que nous avons utilisé lors de la création des compilations. N'hésitez pas à modifier les noms de fichiers, mais assurez-vous que les noms de vos fichiers de poids correspondent à ce que vous spécifiez ici.

Ensuite, nous devons les charger à l'aide du SDK Neuron et les définir dans les poids réels du modèle. Lors du chargement des poids optimisés UNet, notez que nous spécifions également le nombre de cœurs Neuron sur lesquels nous devons les charger. Ici, nous chargeons sur un seul accélérateur avec deux cœurs :

# Load the compiled UNet onto two neuron cores.
    pipe.unet = NeuronUNet(UNetWrap(pipe.unet))
    logging.info(f"Loading model: unet:created")
    device_ids = [idx for idx in range(tensor_parallel_degree)]
   
    pipe.unet.unetwrap = torch_neuronx.DataParallel(torch.jit.load(unet_filename), device_ids, set_dynamic_batching=False)
   
 
    # Load other compiled models onto a single neuron core.
 
    # - load encoders
    pipe.text_encoder = NeuronTextEncoder(pipe.text_encoder)
    clip_compiled = torch.jit.load(text_encoder_filename)
    pipe.text_encoder.neuron_text_encoder = clip_compiled
    #- load decoders
    pipe.vae.decoder = torch.jit.load(decoder_filename)
    pipe.vae.post_quant_conv = torch.jit.load(post_quant_conv_filename)

L'exécution de l'inférence avec une invite appelle l'objet canal pour générer une image.

Créer le point de terminaison SageMaker

Nous utilisons les API Boto3 pour créer un point de terminaison SageMaker. Effectuez les étapes suivantes :

1. Créez l'archive avec juste la portion et l'optionnel model.py fichiers et chargez-les sur Amazon S3.
2. Créez le modèle à l'aide du conteneur d'images et de l'archive de modèle téléchargée précédemment.
3. Créez la configuration du point de terminaison à l'aide des paramètres clés suivants :
   1. Utilisez un ml.inf2.xlarge exemple.
   2. Ensemble ContainerStartupHealthCheckTimeoutInSeconds à 240 pour s'assurer que la vérification de l'état commence après le déploiement du modèle.
   3. Ensemble VolumeInGB à une valeur plus grande afin qu'il puisse être utilisé pour charger les poids de modèle d'une taille de 32 Go.

Créer un modèle SageMaker

Après avoir créé le fichier model.tar.gz et l'avoir téléchargé sur Amazon S3, nous devons créer un modèle SageMaker. Nous utilisons le conteneur LMI et l'artefact de modèle de l'étape précédente pour créer le modèle SageMaker. SageMaker nous permet de personnaliser et d'injecter diverses variables d'environnement. Pour ce workflow, nous pouvons tout laisser par défaut. Voir le code suivant :

inference_image_uri = (
    f"763104351884.dkr.ecr.{region}.amazonaws.com/djl-inference:0 djl-serving-inf2"
)

Créez l'objet de modèle, qui crée essentiellement un conteneur de verrouillage qui est chargé sur l'instance et utilisé pour l'inférence :

model_name = name_from_base(f"inf2-sd")
create_model_response = boto3_sm_client.create_model(
    ModelName=model_name,
    ExecutionRoleArn=role,
    PrimaryContainer={"Image": inference_image_uri, "ModelDataUrl": s3_code_artifact},
)

Créer un point de terminaison SageMaker

Dans cette démo, nous utilisons une instance ml.inf2.xlarge. Nous devons définir le VolumeSizeInGB paramètres pour fournir l'espace disque nécessaire pour charger le modèle et les poids. Ce paramètre s'applique aux instances prenant en charge le Boutique de blocs élastiques Amazon (Amazon EBS) pièce jointe de volume. Nous pouvons laisser le délai de téléchargement du modèle et la vérification de l'état du démarrage du conteneur à une valeur plus élevée, ce qui laissera suffisamment de temps au conteneur pour extraire les poids d'Amazon S3 et les charger dans les accélérateurs AWS Inferentia2. Pour plus de détails, reportez-vous à CréerEndpointConfig.

endpoint_config_response = boto3_sm_client.create_endpoint_config( EndpointConfigName=endpoint_config_name,
    ProductionVariants=[
        {
            "VariantName": "variant1",
            "ModelName": model_name,
            "InstanceType": "ml.inf2.xlarge", # - 
            "InitialInstanceCount": 1,
            "ContainerStartupHealthCheckTimeoutInSeconds": 360, 
            "VolumeSizeInGB": 400
        },
    ],
)

Enfin, nous créons un point de terminaison SageMaker :

create_endpoint_response = boto3_sm_client.create_endpoint(
    EndpointName=f"{endpoint_name}", EndpointConfigName=endpoint_config_name
)

Appeler le point de terminaison du modèle

Il s'agit d'un modèle génératif, nous transmettons donc l'invite que le modèle utilise pour générer l'image. La charge utile est de type JSON :

response_model = boto3_sm_run_client.invoke_endpoint( EndpointName=endpoint_name,
    Body=json.dumps(
        {
            "prompt": "Mountain Landscape", 
            "parameters": {} # 
        }
    ), 
    ContentType="application/json",
)

Analyse comparative du modèle de diffusion stable sur Inf2

Nous avons effectué quelques tests pour comparer le modèle de diffusion stable avec le type de données BF 16 sur Inf2, et nous sommes en mesure de dériver des nombres de latence qui rivalisent ou dépassent certains des autres accélérateurs pour la diffusion stable. Ceci, associé au coût inférieur des puces AWS Inferentia2, en fait une proposition extrêmement précieuse.

Les chiffres suivants proviennent du modèle Stable Diffusion déployé sur une instance inf2.xl. Pour plus d'informations sur les coûts, reportez-vous à Instances Amazon EC2 Inf2.

Conclusion

Dans cet article, nous avons approfondi la compilation, l'optimisation et le déploiement du modèle Stable Diffusion 2.1 à l'aide d'instances Inf2. Nous avons également démontré le déploiement de modèles de diffusion stable à l'aide de SageMaker. Les instances Inf2 offrent également un excellent rapport qualité-prix pour Stable Diffusion 1.5. Pour en savoir plus sur les raisons pour lesquelles les instances Inf2 sont idéales pour l'IA générative et les grands modèles de langage, reportez-vous à Les instances Amazon EC2 Inf2 pour l'inférence IA générative à faible coût et hautes performances sont désormais généralement disponibles. Pour plus de détails sur les performances, reportez-vous à Performances Inf2. Découvrez d'autres exemples sur le GitHub repo.

Remerciements particuliers à Matthew Mcclain, Beni Hegedus, Kamran Khan, Shruti Koparkar et Qing Lan pour leur révision et leurs précieuses contributions.

À propos des auteurs

Vivek Gangasani est architecte principal de solutions d'apprentissage automatique chez Amazon Web Services. Il travaille avec des startups d'apprentissage automatique pour créer et déployer des applications AI/ML sur AWS. Il se concentre actuellement sur la fourniture de solutions pour MLOps, l'inférence ML et le ML low-code. Il a travaillé sur des projets dans différents domaines, dont le traitement du langage naturel et la vision par ordinateur.

KC Tung est architecte de solutions senior chez AWS Annapurna Labs. Il est spécialisé dans la formation et le déploiement de grands modèles d'apprentissage en profondeur à grande échelle dans le cloud. Il a un doctorat. en biophysique moléculaire de l'Université du Texas Southwestern Medical Center à Dallas. Il a pris la parole lors d'AWS Summits et d'AWS Reinvent. Aujourd'hui, il aide les clients à former et à déployer de grands modèles PyTorch et TensorFlow dans le cloud AWS. Il est l'auteur de deux livres : Apprendre TensorFlow Enterprise et les Référence de poche TensorFlow 2.

Rupinder Grewal est un architecte de solutions spécialisé Sr Ai/ML avec AWS. Il se concentre actuellement sur le service des modèles et des MLOps sur SageMaker. Avant d'occuper ce poste, il a travaillé en tant qu'ingénieur en apprentissage automatique pour créer et héberger des modèles. En dehors du travail, il aime jouer au tennis et faire du vélo sur les sentiers de montagne.

• Contenu propulsé par le référencement et distribution de relations publiques. Soyez amplifié aujourd'hui.
• PlatoData.Network Ai générative verticale. Autonomisez-vous. Accéder ici.
• PlatoAiStream. Intelligence Web3. Connaissance Amplifiée. Accéder ici.
• PlatonESG. Automobile / VE, Carbone, Technologie propre, Énergie, Environnement, Solaire, La gestion des déchets. Accéder ici.
• Décalages de bloc. Modernisation de la propriété des compensations environnementales. Accéder ici.
• La source: https://aws.amazon.com/blogs/machine-learning/maximize-stable-diffusion-performance-and-lower-inference-costs-with-aws-inferentia2/