Maximize o desempenho do Stable Diffusion e reduza os custos de inferência com o AWS Inferentia2 | Amazon Web Services

Os modelos generativos de IA têm experimentado um rápido crescimento nos últimos meses devido às suas impressionantes capacidades na criação de texto, imagens, código e áudio realistas. Entre esses modelos, os modelos Stable Diffusion se destacam por sua força única na criação de imagens de alta qualidade com base em instruções de texto. O Stable Diffusion pode gerar uma ampla variedade de imagens de alta qualidade, incluindo retratos realistas, paisagens e até arte abstrata. E, como outros modelos generativos de IA, os modelos de difusão estável requerem computação poderosa para fornecer inferência de baixa latência.

Neste post, mostramos como você pode executar modelos de difusão estável e obter alto desempenho com o menor custo em Amazon Elastic Compute Nuvem (Amazon EC2) usando Instâncias do Amazon EC2 Inf2 alimentado por AWS Inferencia2. Observamos a arquitetura de um modelo de difusão estável e percorremos as etapas de compilação de um modelo de difusão estável usando Neurônio AWS e implantá-lo em uma instância Inf2. Também discutimos as otimizações que o Neuron SDK faz automaticamente para melhorar o desempenho. Você pode executar as versões Stable Diffusion 2.1 e 1.5 no AWS Inferentia2 de maneira econômica. Por fim, mostramos como você pode implantar um modelo de difusão estável em uma instância Inf2 com Amazon Sage Maker.

O tamanho do modelo Stable Diffusion 2.1 em ponto flutuante 32 (FP32) é de 5 GB e 2.5 GB em bfoat16 (BF16). Uma única instância inf2.xlarge possui um acelerador AWS Inferentia2 com 32 GB de memória HBM. O modelo Stable Diffusion 2.1 pode caber em uma única instância inf2.xlarge. Stable Diffusion é um modelo de texto para imagem que você pode usar para criar imagens de diferentes estilos e conteúdos simplesmente fornecendo um prompt de texto como entrada. Para saber mais sobre a arquitetura do modelo de difusão estável, consulte Crie imagens de alta qualidade com modelos Stable Diffusion e implante-as de maneira econômica com o Amazon SageMaker.

Como o Neuron SDK otimiza o desempenho do Stable Diffusion

Antes de podermos implantar o modelo Stable Diffusion 2.1 em instâncias AWS Inferentia2, precisamos compilar os componentes do modelo usando o SDK de neurônios. O Neuron SDK, que inclui um compilador de aprendizado profundo, tempo de execução e ferramentas, compila e otimiza automaticamente modelos de aprendizado profundo para que possam ser executados com eficiência em instâncias Inf2 e extrair desempenho total do acelerador AWS Inferentia2. Temos exemplos disponíveis para o modelo Stable Diffusion 2.1 no GitHub repo. Este notebook apresenta um exemplo completo de como compilar um modelo de difusão estável, salvar os modelos Neuron compilados e carregá-los no tempo de execução para inferência.

Usamos StableDiffusionPipeline do rosto que abraça diffusers biblioteca para carregar e compilar o modelo. Em seguida, compilamos todos os componentes do modelo para Neuron usando torch_neuronx.trace() e salve o modelo otimizado como TorchScript. Os processos de compilação podem consumir bastante memória, exigindo uma quantidade significativa de RAM. Para contornar isso, antes de traçar cada modelo, criamos um deepcopy da parte do pipeline que está sendo rastreada. Em seguida, excluímos o objeto pipeline da memória usando del pipe. Esta técnica é particularmente útil ao compilar em instâncias com pouca RAM.

Além disso, também realizamos otimizações nos modelos de Difusão Estável. UNet contém o aspecto computacionalmente mais intensivo da inferência. O componente UNet opera em tensores de entrada que possuem tamanho de lote dois, gerando um tensor de saída correspondente também com tamanho de lote dois, para produzir uma única imagem. Os elementos desses lotes são totalmente independentes uns dos outros. Podemos aproveitar esse comportamento para obter a latência ideal executando um lote em cada núcleo do Neuron. Compilamos o UNet para um lote (usando tensores de entrada com um lote) e, em seguida, usamos o torch_neuronx.DataParallel API para carregar este modelo de lote único em cada núcleo. A saída desta API é um módulo contínuo de dois lotes: podemos passar para a UNet as entradas de dois lotes e uma saída de dois lotes é retornada, mas internamente, os dois modelos de lote único estão sendo executados nos dois núcleos do Neuron . Essa estratégia otimiza a utilização de recursos e reduz a latência.

Compilar e implantar um modelo de difusão estável em uma instância Inf2 EC2

Para compilar e implantar o modelo Stable Diffusion em uma instância Inf2 EC2, assine o Console de gerenciamento da AWS e crie uma instância inf2.8xlarge. Observe que uma instância inf2.8xlarge é necessária apenas para a compilação do modelo porque a compilação requer uma memória de host maior. O modelo Stable Diffusion pode ser hospedado em uma instância inf2.xlarge. Você pode encontrar a AMI mais recente com bibliotecas Neuron usando o seguinte Interface de linha de comando da AWS Comando (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 text

Para este exemplo, criamos uma instância EC2 usando Deep Learning AMI Neuron PyTorch 1.13 (Ubuntu 20.04). Você pode então criar um ambiente de laboratório JupyterLab conectando-se à instância e executando as seguintes etapas:

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

Um notebook com todos os passos para compilação e hospedagem do modelo está localizado em GitHub.

Vejamos as etapas de compilação de um dos blocos do codificador de texto. Outros blocos que fazem parte do pipeline Stable Diffusion podem ser compilados de forma semelhante.

A primeira etapa é carregar o modelo pré-treinado do Hugging Face. O StableDiffusionPipeline.from_pretrained método carrega o modelo pré-treinado em nosso objeto de pipeline, pipe. Criamos então um deepcopy do codificador de texto do nosso pipeline, clonando-o efetivamente. O del pipe O comando é então usado para excluir o objeto do pipeline original, liberando a memória que foi consumida por ele. Aqui, estamos quantizando o modelo para pesos 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

Esta etapa envolve envolver nosso codificador de texto com o NeuronTextEncoder embrulho. A saída de um módulo codificador de texto compilado será de dict. Nós o convertemos em um list digite usando este wrapper:

text_encoder = NeuronTextEncoder(text_encoder)

Inicializamos o tensor PyTorch emb com alguns valores. O emb tensor é usado como exemplo de entrada para o torch_neuronx.trace função. Esta função rastreia nosso codificador de texto e o compila em um formato otimizado para Neuron. O caminho do diretório para o modelo compilado é construído juntando-se COMPILER_WORKDIR_ROOT com o subdiretório 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'),
        )

O codificador de texto compilado é salvo usando torch.jit.save. É armazenado sob o nome de arquivo model.pt no text_encoder diretório do espaço de trabalho do nosso compilador:

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

A caderno inclui etapas semelhantes para compilar outros componentes do modelo: UNet, decodificador VAE e VAE post_quant_conv. Depois de compilar todos os modelos, você poderá carregar e executar o modelo seguindo estas etapas:

1. Defina os caminhos para os modelos compilados.
2. Carregue um pré-treinado StableDiffusionPipeline model, com sua configuração especificada para usar o tipo de dados bfloat16.
3. Carregue o modelo UNet em dois núcleos Neuron usando o torch_neuronx.DataParallel API. Isso permite a realização de inferência paralela de dados, o que pode acelerar significativamente o desempenho do modelo.
4. Carregue as partes restantes do modelo (text_encoder, decoder e post_quant_conv) em um único núcleo do Neuron.

Você pode então executar o pipeline fornecendo texto de entrada como prompts. A seguir estão algumas imagens geradas pelo modelo para os prompts:

• Retrato de renaud sechan, caneta e tinta, desenhos intricados, de craig mullins, ruan jia, kentaro miura, greg rutkowski, loundraw

• Retrato de um antigo mineiro de carvão do século XIX, bela pintura, com pintura facial altamente detalhada de Greg Rutkowski

• Um castelo no meio de uma floresta

Host Stable Diffusion 2.1 no AWS Inferentia2 e SageMaker

Hospedar modelos de difusão estável com SageMaker também requer compilação com o Neuron SDK. Você pode concluir a compilação antecipadamente ou durante o tempo de execução usando contêineres Large Model Inference (LMI). A compilação antecipada permite tempos de carregamento do modelo mais rápidos e é a opção preferida.

Os contêineres SageMaker LMI fornecem duas maneiras de implantar o modelo:

• Uma opção sem código onde apenas fornecemos um serving.properties arquivo com as configurações necessárias
• Traga seu próprio script de inferência

Analisamos ambas as soluções e examinamos as configurações e o script de inferência (model.py). Neste post, demonstramos a implantação usando um modelo pré-compilado armazenado em um Serviço de armazenamento simples da Amazon (Amazon S3) balde. Você pode usar esse modelo pré-compilado para suas implantações.

Configure o modelo com um script fornecido

Nesta seção, mostramos como configurar o contêiner LMI para hospedar os modelos de difusão estável. O notebook SD2.1 disponível em GitHub. A primeira etapa é criar o pacote de configuração do modelo de acordo com a estrutura de diretório a seguir. Nosso objetivo é usar as configurações mínimas de modelo necessárias para hospedar o modelo. A estrutura de diretório necessária é a seguinte:

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

Em seguida, criamos o servindo.propriedades arquivo com os seguintes parâmetros:

%%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

Os parâmetros especificam o seguinte:

• opção.model_id – Os contêineres LMI usam s5cmd para carregar o modelo do local S3 e, portanto, precisamos especificar o local onde estão nossos pesos compilados.
• opção.entryPoint – Para usar os manipuladores integrados, especificamos a classe transformers-neuronx. Se você tiver um script de inferência personalizado, será necessário fornecê-lo.
• opção.dtype – Especifica o carregamento dos pesos em um tamanho específico. Para este post, usamos o BF16, que reduz ainda mais nossos requisitos de memória em relação ao FP32 e diminui nossa latência devido a isso.
• opção.tensor_parallel_degree – Este parâmetro especifica o número de aceleradores que usamos para este modelo. O acelerador de chip AWS Inferentia2 tem dois núcleos Neuron e, portanto, especificar um valor 2 significa que usamos um acelerador (dois núcleos). Isso significa que agora podemos criar vários trabalhadores para aumentar o rendimento do endpoint.
• opção.motor – Isso é definido como Python para indicar que não usaremos outros compiladores como DeepSpeed ​​​​ou Faster Transformer para esta hospedagem.

Traga seu próprio roteiro

Se quiser trazer seu próprio script de inferência personalizado, você precisará remover o option.entryPoint da serving.properties. O contêiner LMI, nesse caso, procurará um model.py arquivo no mesmo local que o serving.properties e use isso para executar a inferência.

Crie seu próprio script de inferência (model.py)

Criar seu próprio script de inferência é relativamente simples usando o contêiner LMI. O contêiner requer seu model.py arquivo para ter uma implementação do seguinte método:

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

Vamos examinar algumas das áreas críticas do caderno anexado, que demonstra a função trazer seu próprio script.

Substituir o cross_attention módulo com a versão otimizada:

# 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')

Estes são os nomes do arquivo de pesos compilados que usamos ao criar as compilações. Sinta-se à vontade para alterar os nomes dos arquivos, mas certifique-se de que os nomes dos arquivos de pesos correspondam ao que você especificou aqui.

Em seguida, precisamos carregá-los usando o Neuron SDK e defini-los nos pesos reais do modelo. Ao carregar os pesos otimizados da UNet, observe que também especificamos o número de núcleos do Neuron nos quais precisamos carregá-los. Aqui, carregamos em um único acelerador com dois núcleos:

# 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)

Executar a inferência com um prompt invoca o objeto pipe para gerar uma imagem.

Crie o ponto de extremidade do SageMaker

Usamos APIs Boto3 para criar um endpoint SageMaker. Conclua as seguintes etapas:

1. Crie o tarball apenas com a porção e o opcional model.py arquivos e carregue-os no Amazon S3.
2. Crie o modelo usando o contêiner de imagem e o tarball do modelo carregado anteriormente.
3. Crie a configuração do endpoint usando os seguintes parâmetros principais:
   1. Use um ml.inf2.xlarge instância.
   2. Conjunto ContainerStartupHealthCheckTimeoutInSeconds para 240 para garantir que a verificação de integridade seja iniciada após a implantação do modelo.
   3. Conjunto VolumeInGB para um valor maior para que possa ser usado para carregar pesos de modelo com tamanho de 32 GB.

Criar um modelo do SageMaker

Depois de criar o arquivo model.tar.gz e carregá-lo no Amazon S3, precisamos criar um modelo SageMaker. Usamos o contêiner LMI e o artefato do modelo da etapa anterior para criar o modelo SageMaker. SageMaker nos permite personalizar e injetar várias variáveis ​​de ambiente. Para este fluxo de trabalho, podemos deixar tudo como padrão. Veja o seguinte código:

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

Crie o objeto modelo, que essencialmente cria um contêiner de bloqueio que é carregado na instância e usado para inferência:

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},
)

Crie um endpoint SageMaker

Nesta demonstração, usamos uma instância ml.inf2.xlarge. Precisamos definir o VolumeSizeInGB parâmetros para fornecer o espaço em disco necessário para carregar o modelo e os pesos. Este parâmetro é aplicável a instâncias que suportam o Loja de blocos elásticos da Amazon (Amazon EBS) anexo de volume. Podemos deixar o tempo limite de download do modelo e a verificação de integridade de inicialização do contêiner com um valor mais alto, o que dará tempo adequado para o contêiner extrair os pesos do Amazon S3 e carregar nos aceleradores AWS Inferentia2. Para mais detalhes, consulte 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
        },
    ],
)

Por fim, criamos um endpoint SageMaker:

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

Chamar o endpoint do modelo

Este é um modelo generativo, então passamos o prompt que o modelo usa para gerar a imagem. A carga útil é do tipo JSON:

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

Comparando o modelo de difusão estável em Inf2

Executamos alguns testes para avaliar o modelo de difusão estável com o tipo de dados BF 16 em Inf2 e somos capazes de derivar números de latência que rivalizam ou excedem alguns dos outros aceleradores para difusão estável. Isso, juntamente com o custo mais baixo dos chips AWS Inferentia2, torna esta proposta extremamente valiosa.

Os números a seguir são do modelo Stable Diffusion implantado em uma instância inf2.xl. Para obter mais informações sobre custos, consulte Instâncias do Amazon EC2 Inf2.

Conclusão

Nesta postagem, nos aprofundamos na compilação, otimização e implantação do modelo Stable Diffusion 2.1 usando instâncias Inf2. Também demonstramos a implantação de modelos de difusão estável usando SageMaker. As instâncias Inf2 também oferecem ótimo desempenho de preço para Stable Diffusion 1.5. Para saber mais sobre por que as instâncias Inf2 são ótimas para IA generativa e grandes modelos de linguagem, consulte Instâncias Inf2 do Amazon EC2 para inferência de IA generativa de baixo custo e alto desempenho já estão disponíveis ao público em geral. Para detalhes de desempenho, consulte Desempenho Inf2. Confira exemplos adicionais no GitHub repo.

Agradecimentos especiais a Matthew Mcclain, Beni Hegedus, Kamran Khan, Shruti Koparkar e Qing Lan pela revisão e fornecimento de contribuições valiosas.

Sobre os autores

Vivek Gangasani é arquiteto sênior de soluções de aprendizado de máquina na Amazon Web Services. Ele trabalha com startups de aprendizado de máquina para criar e implantar aplicativos AI/ML na AWS. Atualmente, ele está focado em fornecer soluções para MLOps, inferência de ML e ML de baixo código. Ele trabalhou em projetos em diferentes domínios, incluindo processamento de linguagem natural e visão computacional.

KC Tung é Arquiteto de Soluções Sênior no AWS Annapurna Labs. Ele é especialista em treinamento e implantação de grandes modelos de deep learning em escala na nuvem. Ele tem um Ph.D. em biofísica molecular pela University of Texas Southwestern Medical Center, em Dallas. Ele falou em AWS Summits e AWS Reinvent. Hoje ele ajuda os clientes a treinar e implantar grandes modelos PyTorch e TensorFlow na nuvem AWS. É autor de dois livros: Conheça o TensorFlow Enterprise e Referência de bolso do TensorFlow 2.

Rupinder Grewal é um arquiteto de soluções especialista sênior em Ai/ML da AWS. Ele atualmente se concentra em servir de modelos e MLOps no SageMaker. Antes dessa função, ele trabalhou como engenheiro de aprendizado de máquina construindo e hospedando modelos. Fora do trabalho, ele gosta de jogar tênis e andar de bicicleta em trilhas nas montanhas.

• Conteúdo com tecnologia de SEO e distribuição de relações públicas. Seja amplificado hoje.
• PlatoData.Network Gerativa Vertical Ai. Capacite-se. Acesse aqui.
• PlatoAiStream. Inteligência Web3. Conhecimento Amplificado. Acesse aqui.
• PlatãoESG. Automotivo / EVs, Carbono Tecnologia Limpa, Energia, Ambiente, Solar, Gestão de resíduos. Acesse aqui.
• BlockOffsets. Modernizando a Propriedade de Compensação Ambiental. Acesse aqui.
• Fonte: https://aws.amazon.com/blogs/machine-learning/maximize-stable-diffusion-performance-and-lower-inference-costs-with-aws-inferentia2/