Massimizza le prestazioni stabili di Diffusion e riduci i costi di inferenza con AWS Inferentia2
Massimizza le prestazioni di Diffusion e riduci i costi con AWS Inferentia2.
I modelli di intelligenza artificiale generativa stanno vivendo una rapida crescita negli ultimi mesi grazie alle loro impressionanti capacità nel creare testo, immagini, codice e audio realistici. Tra questi modelli, i modelli di Diffusione Stabile si distinguono per la loro unica capacità di creare immagini di alta qualità basate su prompt di testo. La Diffusione Stabile può generare una vasta gamma di immagini di alta qualità, tra cui ritratti realistici, paesaggi e persino arte astratta. E, come altri modelli di intelligenza artificiale generativa, i modelli di Diffusione Stabile richiedono un potente calcolo per fornire inferenze a bassa latenza.
In questo post, mostriamo come eseguire modelli di Diffusione Stabile e ottenere prestazioni elevate al costo più basso possibile in Amazon Elastic Compute Cloud (Amazon EC2) utilizzando le istanze Amazon EC2 Inf2 con tecnologia AWS Inferentia2. Esaminiamo l’architettura di un modello di Diffusione Stabile e seguiamo i passaggi per compilare un modello di Diffusione Stabile utilizzando AWS Neuron e distribuirlo su un’istanza Inf2. Discutiamo anche le ottimizzazioni che il Neuron SDK apporta automaticamente per migliorare le prestazioni. È possibile eseguire sia le versioni Stable Diffusion 2.1 che 1.5 su AWS Inferentia2 in modo economico. Infine, mostriamo come è possibile distribuire un modello di Diffusione Stabile su un’istanza Inf2 con Amazon SageMaker.
La dimensione del modello Stable Diffusion 2.1 in floating point 32 (FP32) è di 5 GB e di 2,5 GB in bfoat16 (BF16). Un’istanza inf2.xlarge singola dispone di un acceleratore AWS Inferentia2 con 32 GB di memoria HBM. Il modello Stable Diffusion 2.1 può adattarsi a una singola istanza inf2.xlarge. Stable Diffusion è un modello di testo-immagine che è possibile utilizzare per creare immagini di diversi stili e contenuti semplicemente fornendo un prompt di testo in input. Per saperne di più sull’architettura del modello di Diffusione Stabile, fare riferimento alla guida “Creare immagini di alta qualità con modelli di Diffusione Stabile e distribuirle in modo economico con Amazon SageMaker”.
Come il Neuron SDK ottimizza le prestazioni della Diffusione Stabile
Prima di poter distribuire il modello Stable Diffusion 2.1 sulle istanze AWS Inferentia2, è necessario compilare i componenti del modello utilizzando il Neuron SDK. Il Neuron SDK, che include un compilatore, un runtime e strumenti di deep learning, compila e ottimizza automaticamente i modelli di deep learning in modo che possano essere eseguiti efficientemente su istanze Inf2 ed estrarre le prestazioni complete dell’acceleratore AWS Inferentia2. Abbiamo esempi disponibili per il modello Stable Diffusion 2.1 nel repository GitHub. Questo notebook presenta un esempio completo di come compilare un modello di Diffusione Stabile, salvare i modelli Neuron compilati e caricarli nel runtime per l’inferenza.
- Politica sull’IA @🤗 Considerazioni sull’apprendimento automatico aperto nella legge sull’IA dell’UE
- Google a ICML 2023
- La funzione di rilevamento dell’usurpazione protegge i marchi e le personalità dai falsi account sui social media
Utilizziamo StableDiffusionPipeline dalla libreria diffusers di Hugging Face per caricare e compilare il modello. Compiliamo quindi tutti i componenti del modello per Neuron utilizzando torch_neuronx.trace() e salviamo il modello ottimizzato come TorchScript. I processi di compilazione possono richiedere molta memoria, richiedendo una quantità significativa di RAM. Per aggirare questo problema, prima di tracciare ogni modello, creiamo una deepcopy della parte della pipeline che viene tracciata. Successivamente, eliminiamo l’oggetto pipeline dalla memoria utilizzando del pipe. Questa tecnica è particolarmente utile quando si compila su istanze con RAM limitata.
Inoltre, eseguiamo anche ottimizzazioni sui modelli di Diffusione Stabile. UNet rappresenta l’aspetto più computazionalmente intensivo della fase di inferenza. Il componente UNet opera su tensori di input che hanno una dimensione di batch di due, generando un tensore di output corrispondente sempre con una dimensione di batch di due, per produrre un singolo’immagine. Gli elementi all’interno di questi batch sono completamente indipendenti tra loro. Possiamo sfruttare questo comportamento per ottenere una latenza ottimale eseguendo un batch su ogni core di Neuron. Compiliamo il UNet per un batch (utilizzando tensori di input con un batch singolo), quindi utilizziamo l’API torch_neuronx.DataParallel per caricare questo modello a batch singolo su ogni core. L’output di questa API è un modulo a due batch senza soluzione di continuità: possiamo passare al UNet gli input di due batch e viene restituito un output a due batch, ma internamente i due modelli a batch singolo vengono eseguiti sui due core di Neuron. Questa strategia ottimizza l’utilizzo delle risorse e riduce la latenza.
Compilare e distribuire un modello di Diffusione Stabile su un’istanza EC2 Inf2
Per compilare e distribuire il modello di Diffusione Stabile su un’istanza EC2 Inf2, accedi alla Console di Gestione AWS e crea un’istanza inf2.8xlarge. Nota che un’istanza inf2.8xlarge è necessaria solo per la compilazione del modello perché la compilazione richiede una memoria host superiore. Il modello di Diffusione Stabile può essere ospitato su un’istanza inf2.xlarge. Puoi trovare l’ultima AMI con le librerie Neuron utilizzando il seguente comando AWS Command Line Interface (AWS CLI):
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 textPer questo esempio, abbiamo creato un’istanza EC2 utilizzando l’AMI Deep Learning Neuron PyTorch 1.13 (Ubuntu 20.04). È quindi possibile creare un ambiente di laboratorio JupyterLab collegandosi all’istanza ed eseguendo i seguenti passaggi:
eseguire source /opt/aws_neuron_venv_pytorch/bin/activate
pip install jupyterlab
jupyter-labUn notebook con tutti i passaggi per la compilazione e l’hosting del modello si trova su GitHub.
Esaminiamo i passaggi di compilazione per uno dei blocchi encoder di testo. Altri blocchi che fanno parte del pipeline Stable Diffusion possono essere compilati in modo simile.
Il primo passaggio consiste nel caricare il modello pre-addestrato da Hugging Face. Il metodo StableDiffusionPipeline.from_pretrained carica il modello pre-addestrato nel nostro oggetto pipeline, pipe. Creiamo quindi una deepcopy del blocco encoder di testo dalla nostra pipeline, clonandolo effettivamente. Il comando del pipe viene quindi utilizzato per eliminare l’oggetto pipeline originale, liberando la memoria che è stata consumata da esso. Qui, stiamo quantizzando il modello con pesi 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 pipeQuesto passaggio coinvolge l’incapsulamento del nostro blocco encoder di testo con il wrapper NeuronTextEncoder. L’output di un modulo di encoder di testo compilato sarà di tipo dict. Lo convertiamo in un tipo list utilizzando questo wrapper:
text_encoder = NeuronTextEncoder(text_encoder)Inizializziamo il tensore PyTorch emb con alcuni valori. Il tensore emb viene utilizzato come input di esempio per la funzione torch_neuronx.trace. Questa funzione traccia il nostro blocco encoder di testo e lo compila in un formato ottimizzato per Neuron. Il percorso della directory per il modello compilato è costruito unendo COMPILER_WORKDIR_ROOT con la sottodirectory 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’encoder di testo compilato viene salvato utilizzando torch.jit.save. Viene memorizzato con il nome di file model.pt nella directory text_encoder dello spazio di lavoro del compilatore:
text_encoder_filename = os.path.join(COMPILER_WORKDIR_ROOT, 'text_encoder/model.pt')
torch.jit.save(text_encoder_neuron, text_encoder_filename)Il notebook include passaggi simili per compilare altri componenti del modello: UNet, decoder VAE e post_quant_conv VAE. Dopo aver compilato tutti i modelli, è possibile caricare ed eseguire il modello seguendo questi passaggi:
- Definire i percorsi per i modelli compilati.
- Caricare un modello pre-addestrato
StableDiffusionPipeline, con la sua configurazione specificata per utilizzare il tipo di dati bfloat16. - Caricare il modello UNet su due core di Neuron utilizzando l’API
torch_neuronx.DataParallel. Ciò consente di eseguire l’elaborazione parallela dei dati, il che può migliorare significativamente le prestazioni del modello. - Caricare le parti rimanenti del modello (
text_encoder,decoderepost_quant_conv) su un singolo core di Neuron.
Puoi quindi eseguire la pipeline fornendo il testo di input come prompt. Di seguito sono riportate alcune immagini generate dal modello per i prompt:
- Ritratto di renaud sechan, penna e inchiostro, disegni a linee intricate, di craig mullins, ruan jia, kentaro miura, greg rutkowski, loundraw
- Ritratto di vecchio minatore di carbone del XIX secolo, bellissimo dipinto, con un volto altamente dettagliato dipinto da Greg Rutkowski
- Un castello nel mezzo di una foresta
Ospita Stable Diffusion 2.1 su AWS Inferentia2 e SageMaker
Ospitare modelli di Stable Diffusion con SageMaker richiede anche la compilazione con il Neuron SDK. È possibile completare la compilazione in anticipo o durante l’esecuzione utilizzando i container di Large Model Inference (LMI). La compilazione in anticipo consente tempi di caricamento del modello più rapidi ed è l’opzione preferita.
I container SageMaker LMI offrono due modalità di distribuzione del modello:
- Un’opzione senza codice in cui forniamo solo un file
serving.propertiescon le configurazioni necessarie - Porta il tuo script di inferenza
Esamineremo entrambe le soluzioni, le configurazioni e lo script di inferenza (model.py). In questo post, dimostriamo il deployment utilizzando un modello pre-compilato memorizzato in un bucket di Amazon Simple Storage Service (Amazon S3). Puoi utilizzare questo modello pre-compilato per i tuoi deployment.
Configura il modello con uno script fornito
In questa sezione, mostriamo come configurare il container LMI per ospitare i modelli di Stable Diffusion. Il notebook SD2.1 è disponibile su GitHub. Il primo passo è creare il pacchetto di configurazione del modello secondo la seguente struttura di directory. Il nostro obiettivo è utilizzare le configurazioni del modello minime necessarie per ospitare il modello. La struttura di directory necessaria è la seguente:
<directory-radice-configurazione> /
├── serving.properties
│
└── model.py [OPZIONALE]Successivamente, creiamo il file serving.properties con i seguenti parametri:
%%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=bf16I parametri specificano quanto segue:
- option.model_id – I container LMI utilizzano s5cmd per caricare il modello dalla posizione S3 e quindi è necessario specificare la posizione dei nostri pesi compilati.
- option.entryPoint – Per utilizzare gli handler integrati, specificare la classe transformers-neuronx. Se si dispone di uno script di inferenza personalizzato, è necessario fornirlo invece.
- option.dtype – Specifica il caricamento dei pesi in una dimensione specifica. Per questo post, utilizziamo BF16, che riduce ulteriormente i requisiti di memoria e riduce la latenza.
- option.tensor_parallel_degree – Questo parametro specifica il numero di acceleratori che utilizziamo per questo modello. L’acceleratore del chip AWS Inferentia2 ha due core Neuron e quindi specificando un valore di 2 significa che utilizziamo un acceleratore (due core). Ciò significa che ora possiamo creare più worker per aumentare il throughput del punto di accesso.
- option.engine – Viene impostato su Python per indicare che non utilizzeremo altri compilatori come DeepSpeed o Faster Transformer per questo hosting.
Porta il tuo script
Se si desidera utilizzare il proprio script di inferenza personalizzato, è necessario rimuovere l’opzione option.entryPoint da serving.properties. In tal caso, il container LMI cercherà un file model.py nella stessa posizione di serving.properties e lo utilizzerà per eseguire l’inferenza.
Crea il tuo script di inferenza personalizzato (model.py)
Creare il proprio script di inferenza è relativamente semplice utilizzando il contenitore LMI. Il contenitore richiede che il file model.py abbia un’implementazione del seguente metodo:
def handle(inputs: Input) che restituisce un oggetto di tipo OutputsEsaminiamo alcune delle aree critiche del notebook allegato, che illustra la funzione di script personalizzato.
Sostituisci il modulo cross_attention con la versione ottimizzata:
# Sostituisci il modulo di cross-attention originale con il modulo di cross-attention personalizzato per prestazioni migliori
CrossAttention.get_attention_scores = get_attention_scores
Carica i pesi compilati per i seguenti
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')Questi sono i nomi dei file dei pesi compilati che abbiamo utilizzato durante la creazione delle compilazioni. Sentiti libero di modificare i nomi dei file, ma assicurati che i nomi dei file dei pesi corrispondano a quanto specificato qui.
Successivamente, dobbiamo caricarli utilizzando il Neuron SDK e impostarli nei pesi effettivi del modello. Durante il caricamento dei pesi ottimizzati UNet, nota che stiamo specificando anche il numero di core Neuron su cui caricarli. Qui, carichiamo su un singolo acceleratore con due core:
# Carica il UNet compilato su due core Neuron.
pipe.unet = NeuronUNet(UNetWrap(pipe.unet))
logging.info(f"Caricamento modello: unet: creato")
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)
# Carica altri modelli compilati su un singolo core Neuron.
# - carica gli encoder
pipe.text_encoder = NeuronTextEncoder(pipe.text_encoder)
clip_compiled = torch.jit.load(text_encoder_filename)
pipe.text_encoder.neuron_text_encoder = clip_compiled
#- carica i decoder
pipe.vae.decoder = torch.jit.load(decoder_filename)
pipe.vae.post_quant_conv = torch.jit.load(post_quant_conv_filename)Eseguire l’inferenza con un prompt invoca l’oggetto pipe per generare un’immagine.
Crea il punto di accesso SageMaker
Utilizziamo le API Boto3 per creare un punto di accesso SageMaker. Completa i seguenti passaggi:
- Crea il pacchetto tarball solo con i file di servizio e, facoltativamente,
model.py, e caricalo su Amazon S3. - Crea il modello utilizzando il contenitore dell’immagine e il pacchetto tarball caricato in precedenza.
- Crea la configurazione del punto di accesso utilizzando i seguenti parametri chiave:
- Utilizza un’istanza
ml.inf2.xlarge. - Imposta
ContainerStartupHealthCheckTimeoutInSecondssu 240 per garantire che il controllo di integrità inizi dopo il dispiegamento del modello. - Imposta
VolumeInGBsu un valore maggiore in modo che possa essere utilizzato per caricare i pesi del modello di dimensioni 32 GB.
- Utilizza un’istanza
Crea un modello SageMaker
Dopo aver creato il file model.tar.gz e caricato su Amazon S3, è necessario creare un modello SageMaker. Utilizziamo il contenitore LMI e l’artefatto del modello dal passaggio precedente per creare il modello SageMaker. SageMaker ci consente di personalizzare e iniettare varie variabili d’ambiente. Per questo flusso di lavoro, possiamo lasciare tutto come predefinito. Vedi il seguente codice:
inference_image_uri = (
f"763104351884.dkr.ecr.{region}.amazonaws.com/djl-inference:0 djl-serving-inf2"
)Crea l’oggetto del modello, che essenzialmente crea un contenitore di lockdown che viene caricato sull’istanza e utilizzato per l’inferenza:
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},
)Crea un endpoint SageMaker
In questa demo utilizziamo un’istanza ml.inf2.xlarge. Dobbiamo impostare i parametri VolumeSizeInGB per fornire lo spazio su disco necessario per caricare il modello e i pesi. Questo parametro è applicabile alle istanze che supportano l’attacco del volume Amazon Elastic Block Store (Amazon EBS). Possiamo lasciare il timeout per il download del modello e il controllo di avvio del contenitore a un valore più alto, che darà il tempo adeguato al contenitore per recuperare i pesi da Amazon S3 e caricarli negli acceleratori AWS Inferentia2. Per ulteriori dettagli, fare riferimento a CreateEndpointConfig.
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
},
],
)Infine, creiamo un endpoint SageMaker:
create_endpoint_response = boto3_sm_client.create_endpoint(
EndpointName=f"{endpoint_name}", EndpointConfigName=endpoint_config_name
)Invoca il modello endpoint
Questo è un modello generativo, quindi passiamo il prompt che il modello utilizza per generare l’immagine. Il payload è di tipo JSON:
response_model = boto3_sm_run_client.invoke_endpoint(
EndpointName=endpoint_name,
Body=json.dumps(
{
"prompt": "Paesaggio di montagna",
"parameters": {} #
}
),
ContentType="application/json",
)Benchmarking del modello Stable Diffusion su Inf2
Abbiamo eseguito alcuni test per valutare le prestazioni del modello Stable Diffusion con il tipo di dati BF 16 su Inf2, e siamo in grado di ottenere valori di latenza che si avvicinano o superano alcuni degli altri acceleratori per Stable Diffusion. Questo, unito al costo inferiore dei chip AWS Inferentia2, rende questa una proposta estremamente vantaggiosa.
I seguenti numeri sono del modello Stable Diffusion distribuito su un’istanza inf2.xl. Per ulteriori informazioni sui costi, fare riferimento alle istanze Amazon EC2 Inf2.
| Modello Stable Diffusion 1.5 Stable Diffusion 1.5 Stable Diffusion 1.5 Stable Diffusion 1.5 | Risoluzione 512×512 768×768 512×512 768×768 | Tipo di dati bf16 bf16 bf16 bf16 | Iterazioni 50 50 30 30 | Latency P95 (ms) 2.427,4 8.235,9 1.456,5 4.941,6 | Costo Inf2.xl su richiesta per ora $0,76 $0,76 $0,76 $0,76 | Inf2.xl (Costo per immagine) $0,0005125 $0,0017387 $0,0003075 $0,0010432 |
| Stable Diffusion 2.1 Stable Diffusion 2.1 Stable Diffusion 2.1 Stable Diffusion 2.1 | 512×512 768×768 512×512 768×768 | bf16 bf16 bf16 bf16 | 50 50 30 30 | 1.976,9 6.836,3 1.186,2 4.101,8 | $0,76 $0,76 $0,76 $0,76 | $0,0004174 $0,0014432 $0,0002504 $0,0008659 |
Conclusion
In questo post, abbiamo approfondito la compilazione, l’ottimizzazione e la distribuzione del modello Stable Diffusion 2.1 utilizzando istanze Inf2. Abbiamo anche dimostrato la distribuzione dei modelli Stable Diffusion utilizzando SageMaker. Le istanze Inf2 offrono anche ottime prestazioni a un ottimo rapporto qualità-prezzo per Stable Diffusion 1.5. Per saperne di più su perché le istanze Inf2 sono ideali per l’IA generativa e i modelli di linguaggio di grandi dimensioni, consulta Amazon EC2 Inf2 Instances for Low-Cost, High-Performance Generative AI Inference are Now Generally Available. Per i dettagli sulle prestazioni, consulta Inf2 Performance. Dai un’occhiata agli esempi aggiuntivi nel repository di GitHub.
Un ringraziamento speciale a Matthew Mcclain, Beni Hegedus, Kamran Khan, Shruti Koparkar e Qing Lan per la revisione e i preziosi contributi forniti.
