Zmaksymalizuj wydajność stabilnej dyfuzji i obniż koszty wnioskowania dzięki AWS Inferentia2 | Usługi sieciowe Amazona

Modele generatywnej sztucznej inteligencji odnotowały w ostatnich miesiącach szybki rozwój ze względu na imponujące możliwości tworzenia realistycznego tekstu, obrazów, kodu i dźwięku. Wśród tych modeli modele Stable Diffusion wyróżniają się wyjątkową siłą w tworzeniu wysokiej jakości obrazów w oparciu o podpowiedzi tekstowe. Stable Diffusion może generować szeroką gamę obrazów wysokiej jakości, w tym realistyczne portrety, krajobrazy, a nawet dzieła sztuki abstrakcyjnej. Podobnie jak inne generatywne modele sztucznej inteligencji, modele Stable Diffusion wymagają wydajnych obliczeń, aby zapewnić wnioskowanie o niskim opóźnieniu.

W tym poście pokazujemy, jak uruchomić modele ze stabilną dyfuzją i osiągnąć wysoką wydajność przy najniższych kosztach Elastyczna chmura obliczeniowa Amazon (Amazon EC2) za pomocą Instancje Amazon EC2 Inf2 zasilany przez Inferencja AWS2. Przyjrzymy się architekturze modelu stabilnej dyfuzji i przejdziemy przez kolejne etapy kompilacji modelu stabilnej dyfuzji przy użyciu Neuron AWS i wdrożenie go w instancji Inf2. Omawiamy także optymalizacje, które pakiet Neuron SDK automatycznie wprowadza w celu poprawy wydajności. Na platformie AWS Inferentia2.1 można ekonomicznie uruchomić wersję Stable Diffusion 1.5 i 2. Na koniec pokazujemy, jak wdrożyć model Stable Diffusion w instancji Inf2 za pomocą Amazon Sage Maker.

Rozmiar modelu Stable Diffusion 2.1 w formacie zmiennoprzecinkowym 32 (FP32) wynosi 5 GB i 2.5 GB w bfoat16 (BF16). Pojedyncza instancja inf2.xlarge ma jeden akcelerator AWS Inferentia2 z 32 GB pamięci HBM. Model Stable Diffusion 2.1 może zmieścić się w pojedynczej instancji inf2.xlarge. Stable Diffusion to model zamiany tekstu na obraz, którego można używać do tworzenia obrazów o różnych stylach i treści, po prostu wprowadzając monit tekstowy jako dane wejściowe. Aby dowiedzieć się więcej o architekturze modelu Stable Diffusion, zobacz Twórz wysokiej jakości obrazy za pomocą modeli stabilnej dystrybucji i wdrażaj je w ekonomiczny sposób za pomocą Amazon SageMaker.

Jak pakiet Neuron SDK optymalizuje wydajność stabilnej dyfuzji

Zanim będziemy mogli wdrożyć model Stable Diffusion 2.1 w instancjach AWS Inferentia2, musimy skompilować komponenty modelu przy użyciu SDK neuronu. Zestaw Neuron SDK, który zawiera kompilator głębokiego uczenia się, środowisko wykonawcze i narzędzia, kompiluje i automatycznie optymalizuje modele głębokiego uczenia się, aby mogły wydajnie działać na instancjach Inf2 i wydobywać pełną wydajność akceleratora AWS Inferentia2. Mamy przykłady dostępne dla modelu Stable Diffusion 2.1 na stronie GitHub repo. W tym notatniku przedstawiono kompleksowy przykład kompilowania modelu Stable Diffusion, zapisywania skompilowanych modeli Neuron i ładowania ich do środowiska wykonawczego w celu wyciągnięcia wniosków.

Używamy pliki StableDiffusionPipeline z Przytulonej Twarzy diffusers bibliotekę do załadowania i skompilowania modelu. Następnie kompilujemy wszystkie elementy modelu dla Neurona przy użyciu torch_neuronx.trace() i zapisz zoptymalizowany model jako TorchScript. Procesy kompilacji mogą wymagać dużej ilości pamięci i wymagać znacznej ilości pamięci RAM. Aby to obejść, przed śledzeniem każdego modelu tworzymy plik deepcopy śledzonej części rurociągu. Następnie usuwamy obiekt potoku z pamięci za pomocą del pipe. Ta technika jest szczególnie przydatna podczas kompilacji na instancjach z małą ilością pamięci RAM.

Dodatkowo wykonujemy również optymalizacje modeli Stable Diffusion. UNet posiada najbardziej intensywny obliczeniowo aspekt wnioskowania. Komponent UNet działa na tensorach wejściowych o rozmiarze wsadowym wynoszącym dwa, generując odpowiedni tensor wyjściowy również o rozmiarze wsadowym wynoszącym dwa, w celu utworzenia pojedynczego obrazu. Elementy wchodzące w skład tych partii są całkowicie od siebie niezależne. Możemy skorzystać z tego zachowania, aby uzyskać optymalne opóźnienie, uruchamiając jedną partię na każdym rdzeniu Neuron. Kompilujemy UNet dla jednej partii (używając tensorów wejściowych w jednej partii), a następnie używamy metody torch_neuronx.DataParallel API do ładowania tego pojedynczego modelu wsadowego na każdy rdzeń. Dane wyjściowe tego interfejsu API to płynny moduł składający się z dwóch partii: możemy przekazać do UNet dane wejściowe z dwóch partii i zwrócony zostanie wynik z dwóch partii, ale wewnętrznie dwa modele z pojedynczą partią działają na dwóch rdzeniach Neuron . Strategia ta optymalizuje wykorzystanie zasobów i zmniejsza opóźnienia.

Skompiluj i wdróż model stabilnej dyfuzji na instancji Inf2 EC2

Aby skompilować i wdrożyć model Stable Diffusion na instancji Inf2 EC2, zaloguj się do Konsola zarządzania AWS i utwórz instancję inf2.8xlarge. Należy pamiętać, że instancja inf2.8xlarge jest wymagana tylko do kompilacji modelu, ponieważ kompilacja wymaga większej pamięci hosta. Model Stable Diffusion może być hostowany w instancji inf2.xlarge. Najnowsze AMI z bibliotekami Neuron można znaleźć, korzystając z poniższych informacji Interfejs wiersza poleceń AWS (AWS CLI), polecenie:

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

Na potrzeby tego przykładu utworzyliśmy instancję EC2 przy użyciu Deep Learning AMI Neuron PyTorch 1.13 (Ubuntu 20.04). Następnie możesz utworzyć środowisko laboratoryjne JupyterLab, łącząc się z instancją i wykonując następujące kroki:

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

Znajduje się na nim notatnik zawierający wszystkie kroki kompilacji i hostowania modelu GitHub.

Przyjrzyjmy się etapom kompilacji jednego z bloków kodera tekstu. Inne bloki będące częścią potoku Stable Diffusion można skompilować w podobny sposób.

Pierwszym krokiem jest załadowanie wstępnie wytrenowanego modelu z Hugging Face. The StableDiffusionPipeline.from_pretrained metoda ładuje wstępnie wytrenowany model do naszego obiektu potoku, pipe. Następnie tworzymy deepcopy kodera tekstu z naszego potoku, skutecznie go klonując. The del pipe polecenie jest następnie używane do usunięcia oryginalnego obiektu potoku, zwalniając zajętą ​​przez niego pamięć. Tutaj kwantyzujemy model do wag 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

Ten krok polega na opakowaniu naszego kodera tekstu za pomocą pliku NeuronTextEncoder obwoluta. Dane wyjściowe skompilowanego modułu kodera tekstu będą miały postać dict. Konwertujemy go na list wpisz, używając tego opakowania:

text_encoder = NeuronTextEncoder(text_encoder)

Inicjujemy tensor PyTorch emb z pewnymi wartościami. The emb tensor jest używany jako przykładowe wejście dla torch_neuronx.trace funkcjonować. Ta funkcja śledzi nasz koder tekstu i kompiluje go do formatu zoptymalizowanego dla Neuron. Ścieżka katalogu skompilowanego modelu jest tworzona poprzez łączenie COMPILER_WORKDIR_ROOT z podkatalogiem 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'),
        )

Skompilowany koder tekstu jest zapisywany przy użyciu torch.jit.save. Jest on przechowywany pod nazwą pliku model.pt w folderze text_encoder katalog obszaru roboczego naszego kompilatora:

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

Połączenia notatnik obejmuje podobne kroki w celu kompilacji innych komponentów modelu: UNet, dekoder VAE i VAE post_quant_conv. Po skompilowaniu wszystkich modeli możesz załadować i uruchomić model, wykonując następujące kroki:

1. Zdefiniuj ścieżki dla skompilowanych modeli.
2. Załaduj wstępnie przeszkolony StableDiffusionPipeline model, którego konfiguracja została określona tak, aby korzystała z typu danych bfloat16.
3. Załaduj model UNet na dwa rdzenie Neuron za pomocą torch_neuronx.DataParallel API. Umożliwia to równoległe wnioskowanie o danych, co może znacznie przyspieszyć działanie modelu.
4. Załaduj pozostałe części modelu (text_encoder, decoder, post_quant_conv) na pojedynczy rdzeń Neuronu.

Następnie można uruchomić potok, podając tekst wejściowy jako monit. Poniżej znajdują się zdjęcia wygenerowane przez model dla podpowiedzi:

• Portret Renauda Sechana, piórem i tuszem, zawiłe rysunki liniowe, autor: Craig Mullins, Ruan Jia, Kentaro Miura, Greg Rutkowski, Loundraw

• Portret starego górnika z XIX w., piękne malowidło, z bardzo szczegółowym malowaniem twarzy autorstwa Grega Rutkowskiego

• Zamek w środku lasu

Host Stable Diffusion 2.1 na AWS Inferentia2 i SageMaker

Hostowanie modeli Stable Diffusion za pomocą SageMaker wymaga również kompilacji za pomocą pakietu Neuron SDK. Kompilację można ukończyć przed czasem lub w czasie wykonywania, korzystając z kontenerów LMI (ang. Large Model Inference). Kompilacja z wyprzedzeniem pozwala na szybsze ładowanie modelu i jest preferowaną opcją.

Kontenery LMI SageMaker umożliwiają wdrożenie modelu na dwa sposoby:

• Opcja bez kodu, w której po prostu udostępniamy plik serving.properties plik z wymaganymi konfiguracjami
• Przynieś swój własny skrypt wnioskowania

Przyglądamy się obu rozwiązaniom i przeglądamy konfiguracje oraz skrypt wnioskowania (model.py). W tym poście demonstrujemy wdrożenie przy użyciu wstępnie skompilowanego modelu przechowywanego w pliku Usługa Amazon Simple Storage Łyżka (Amazon S3). Możesz użyć tego wstępnie skompilowanego modelu do swoich wdrożeń.

Skonfiguruj model za pomocą dostarczonego skryptu

W tej sekcji pokażemy, jak skonfigurować kontener LMI do hostowania modeli Stable Diffusion. Notebook SD2.1 dostępny na GitHub. Pierwszym krokiem jest utworzenie pakietu konfiguracyjnego modelu zgodnie z następującą strukturą katalogów. Naszym celem jest użycie minimalnych konfiguracji modelu potrzebnych do hostowania modelu. Wymagana struktura katalogów jest następująca:

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

Następnie tworzymy plik serwowanie.właściwości plik z następującymi parametrami:

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

Parametry określają, co następuje:

• opcja.model_id – Kontenery LMI używają s5cmd do załadowania modelu z lokalizacji S3 i dlatego musimy określić lokalizację, w której znajdują się nasze skompilowane wagi.
• opcja.punkt wejścia – Aby skorzystać z wbudowanych handlerów, podajemy klasę transformatorów-neuronx. Jeśli masz niestandardowy skrypt wnioskowania, musisz go zamiast tego podać.
• opcja.dtyp – Określa ładowanie ciężarków o określonym rozmiarze. W tym poście używamy BF16, który dodatkowo zmniejsza nasze wymagania dotyczące pamięci w porównaniu z FP32 i zmniejsza z tego powodu nasze opóźnienia.
• opcja.tensor_równoległy_stopień – Ten parametr określa liczbę akceleratorów, które stosujemy w tym modelu. Akcelerator chipów AWS Inferentia2 posiada dwa rdzenie Neuron, dlatego podanie wartości 2 oznacza, że ​​używamy jednego akceleratora (dwa rdzenie). Oznacza to, że możemy teraz utworzyć wielu pracowników, aby zwiększyć przepustowość punktu końcowego.
• opcja.silnik – Jest to ustawione na Python, aby wskazać, że nie będziemy używać innych kompilatorów, takich jak DeepSpeed ​​lub Faster Transformer dla tego hostingu.

Przynieś własny scenariusz

Jeśli chcesz użyć własnego niestandardowego skryptu wnioskowania, musisz usunąć plik option.entryPoint od serving.properties. W takim przypadku kontener LMI będzie szukał pliku model.py plik w tej samej lokalizacji co plik serving.properties i użyj tego do uruchomienia wnioskowania.

Utwórz własny skrypt wnioskowania (model.py)

Tworzenie własnego skryptu wnioskowania przy użyciu kontenera LMI jest stosunkowo proste. Kontener wymaga Twojego model.py plik, aby mieć implementację następującej metody:

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

Przyjrzyjmy się niektórym krytycznym obszarom dołączony notatnik, który demonstruje funkcję „przynieś własny skrypt”.

Wymienić cross_attention moduł w wersji zoptymalizowanej:

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

Są to nazwy skompilowanych plików wag, których używaliśmy podczas tworzenia kompilacji. Możesz dowolnie zmieniać nazwy plików, ale upewnij się, że nazwy plików ciężarów odpowiadają tym, które tu określiłeś.

Następnie musimy je załadować za pomocą pakietu Neuron SDK i ustawić w rzeczywistych wagach modelu. Podczas ładowania zoptymalizowanych wag UNet należy pamiętać, że określamy również liczbę rdzeni Neuron, na które musimy je załadować. Tutaj ładujemy do pojedynczego akceleratora z dwoma rdzeniami:

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

Uruchomienie wnioskowania za pomocą zachęty wywołuje obiekt potoku w celu wygenerowania obrazu.

Utwórz punkt końcowy SageMaker

Do tworzenia punktu końcowego SageMaker używamy interfejsów API Boto3. Wykonaj następujące kroki:

1. Utwórz plik tar zawierający tylko porcję i opcjonalne model.py pliki i prześlij je do Amazon S3.
2. Utwórz model, korzystając z kontenera obrazów i przesłanego wcześniej archiwum modelu.
3. Utwórz konfigurację punktu końcowego, używając następujących kluczowych parametrów:
   1. Użyj ml.inf2.xlarge instancja.
   2. Zestaw ContainerStartupHealthCheckTimeoutInSeconds do 240, aby upewnić się, że kontrola kondycji rozpocznie się po wdrożeniu modelu.
   3. Zestaw VolumeInGB na większą wartość, aby można było jej używać do ładowania modeli o wadze 32 GB.

Stwórz model SageMakera

Po utworzeniu pliku model.tar.gz i przesłaniu go do Amazon S3 musimy utworzyć model SageMaker. Do utworzenia modelu SageMaker używamy kontenera LMI i artefaktu modelu z poprzedniego kroku. SageMaker pozwala nam dostosowywać i wstrzykiwać różne zmienne środowiskowe. W przypadku tego przepływu pracy możemy pozostawić wszystko jako domyślne. Zobacz następujący kod:

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

Utwórz obiekt modelu, który zasadniczo tworzy kontener blokujący, który jest ładowany do instancji i używany do wnioskowania:

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

Utwórz punkt końcowy SageMaker

W tym demo używamy instancji ml.inf2.xlarge. Musimy ustawić VolumeSizeInGB parametry, aby zapewnić niezbędną przestrzeń dyskową do załadowania modelu i wag. Ten parametr ma zastosowanie do instancji obsługujących Sklep Amazon Elastic Block (Amazon EBS) załącznik wolumenowy. Możemy pozostawić limit czasu pobierania modelu i kontrolę stanu uruchamiania kontenera na wyższą wartość, co zapewni kontenerowi odpowiednią ilość czasu na ściągnięcie ciężarów z Amazon S3 i załadowanie do akceleratorów AWS Inferentia2. Więcej szczegółów znajdziesz w UtwórzKonfiguracjęPunktuKońcowego.

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

Na koniec tworzymy punkt końcowy SageMaker:

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

Wywołaj punkt końcowy modelu

Jest to model generatywny, dlatego przekazujemy zachętę, której model używa do wygenerowania obrazu. Ładunek jest typu JSON:

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

Test porównawczy modelu stabilnej dyfuzji na Inf2

Przeprowadziliśmy kilka testów w celu porównania modelu Stable Diffusion z typem danych BF 16 na Inf2 i jesteśmy w stanie wyprowadzić liczby opóźnień, które dorównują lub przekraczają niektóre inne akceleratory Stable Diffusion. To, w połączeniu z niższym kosztem chipów AWS Inferentia2, sprawia, że ​​jest to niezwykle cenna propozycja.

Poniższe liczby pochodzą z modelu Stable Diffusion wdrożonego w instancji inf2.xl. Więcej informacji na temat kosztów można znaleźć w artykule Instancje Amazon EC2 Inf2.

Wnioski

W tym poście zagłębiliśmy się w kompilację, optymalizację i wdrażanie modelu Stable Diffusion 2.1 przy użyciu instancji Inf2. Zademonstrowaliśmy także wdrożenie modeli Stable Diffusion przy użyciu SageMaker. Instancje Inf2 zapewniają także świetną wydajność cenową w przypadku wersji Stable Diffusion 1.5. Aby dowiedzieć się więcej o tym, dlaczego instancje Inf2 świetnie nadają się do generatywnej sztucznej inteligencji i dużych modeli językowych, zobacz Instancje Amazon EC2 Inf2 umożliwiające niedrogie i wydajne generatywne wnioskowanie AI są teraz ogólnie dostępne. Aby uzyskać szczegółowe informacje na temat wydajności, zobacz Wydajność Inf2. Sprawdź dodatkowe przykłady na stronie GitHub repo.

Specjalne podziękowania dla Matthew Mcclaina, Beni Hegedusa, Kamrana Khana, Shruti Koparkar i Qing Lan za recenzję i dostarczenie cennych informacji.

O autorach

Vivek Gangasani jest starszym architektem rozwiązań uczenia maszynowego w Amazon Web Services. Współpracuje ze startupami zajmującymi się uczeniem maszynowym przy tworzeniu i wdrażaniu aplikacji AI/ML w AWS. Obecnie koncentruje się na dostarczaniu rozwiązań dla MLOps, wnioskowania ML i ML low-code. Pracował nad projektami z różnych dziedzin, w tym przetwarzania języka naturalnego i widzenia komputerowego.

KC Tung jest starszym architektem rozwiązań w AWS Annapurna Labs. Specjalizuje się w szkoleniu i wdrażaniu modeli głębokiego uczenia na dużą skalę w chmurze. Posiada stopień doktora. w dziedzinie biofizyki molekularnej z University of Texas Southwestern Medical Center w Dallas. Przemawiał na AWS Summits i AWS Reinvent. Dziś pomaga klientom trenować i wdrażać duże modele PyTorch i TensorFlow w chmurze AWS. Jest autorem dwóch książek: Poznaj TensorFlow Enterprise i Odniesienie do kieszeni TensorFlow 2.

Rozbijacz Grewal jest Sr Ai/ML Specialist Solutions Architect z AWS. Obecnie skupia się na serwowaniu modeli i MLOpów na SageMakerze. Wcześniej pracował jako inżynier ds. uczenia maszynowego budując i hostując modele. Poza pracą lubi grać w tenisa i jeździć na rowerze po górskich szlakach.

• Dystrybucja treści i PR oparta na SEO. Uzyskaj wzmocnienie już dziś.
• PlatoData.Network Pionowe generatywne AI. Wzmocnij się. Dostęp tutaj.
• PlatoAiStream. Inteligencja Web3. Wiedza wzmocniona. Dostęp tutaj.
• PlatonESG. Motoryzacja / pojazdy elektryczne, Węgiel Czysta technologia, Energia, Środowisko, Słoneczny, Gospodarowanie odpadami. Dostęp tutaj.
• Przesunięcia bloków. Modernizacja własności offsetu środowiskowego. Dostęp tutaj.
• Źródło: https://aws.amazon.com/blogs/machine-learning/maximize-stable-diffusion-performance-and-lower-inference-costs-with-aws-inferentia2/