Hur Amazon Search minskade ML-slutledningskostnaderna med 85 % med AWS Inferentia

Amazons produktsökmotor indexerar miljarder produkter, betjänar hundratals miljoner kunder över hela världen och är en av de mest använda tjänsterna i världen. Amazon Search-teamet utvecklar maskininlärningsteknik (ML) som driver Amazon.com sökmotor och hjälper kunder att söka utan ansträngning. Att leverera en fantastisk kundupplevelse och verka i den enorma skala som krävs av Amazon.com sökmotor, det här teamet letar alltid efter sätt att bygga mer kostnadseffektiva system med realtidsfördröjning och genomströmningskrav. Teamet utforskar ständigt hårdvara och kompilatorer optimerade för djupinlärning för att påskynda modellträning och slutledning, samtidigt som de minskar driftskostnaderna över hela linjen.

I det här inlägget beskriver vi hur Amazon Search använder AWS slutledning, en högpresterande accelerator som byggts av AWS för att påskynda arbetsbelastningar för djupinlärning. Teamet kör ML-inferens med låg latens med transformatorbaserade NLP-modeller på AWS Inferentia-baserade Amazon Elastic Compute Cloud (Amazon EC2) Inf1-instanser, och sparar upp till 85 % i infrastrukturkostnader samtidigt som stark genomströmning och latensprestanda bibehålls.

Djup inlärning för förutsägelse av dubbletter och frågeavsikter

Att söka på Amazon Marketplace är ett multi-task, multi-modalt problem, som hanterar flera indata såsom ASINs (Amazon Standard Identification Number, ett 10-siffrigt alfanumeriskt nummer som unikt identifierar produkter), produktbilder, textbeskrivningar och frågor. För att skapa en skräddarsydd användarupplevelse används förutsägelser från många modeller för olika aspekter av sökning. Detta är en utmaning eftersom söksystemet har tusentals modeller med tiotusentals transaktioner per sekund (TPS) vid toppbelastning. Vi fokuserar på två komponenter i den upplevelsen:

• Kundupplevda dubbletter av förutsägelser – För att visa den mest relevanta listan över produkter som matchar en användares fråga är det viktigt att identifiera produkter som kunderna har svårt att skilja mellan
• Fråga avsiktsförutsägelse – För att anpassa söksidan och produktlayouten för att bättre passa vad kunden letar efter är det viktigt att förutsäga syftet och typen av användarens fråga (till exempel en mediarelaterad fråga, hjälpfråga och andra frågetyper)

Båda dessa förutsägelser görs med hjälp av transformatormodellarkitekturer, nämligen BERT-baserade modeller. Faktum är att båda delar samma BERT-baserade modell som bas, och var och en staplar ett klassificerings-/regressionshuvud ovanpå denna ryggrad.

Dubblettprediktion tar in olika textfunktioner för ett par utvärderade produkter som indata (som produkttyp, titel, beskrivning och så vidare) och beräknas periodiskt för stora datamängder. Denna modell tränas från början till slut på ett multi-task sätt. Sök efter Amazon SageMaker Processing används för att köra dessa batch-arbetsbelastningar regelbundet för att automatisera deras lansering och endast betala för den handläggningstid som används. För detta användningsfall för batch-arbetsbelastning var kravet på slutledningsgenomströmning 8,800 XNUMX totala TPS.

Avsiktsförutsägelse tar användarens textfråga som input och behövs i realtid för att dynamiskt betjäna daglig trafik och förbättra användarupplevelsen på Amazon Marketplace. Modellen är tränad på ett klassificeringsmål i flera klasser. Denna modell distribueras sedan på Amazon Elastic Container Service (Amazon ECS), som möjliggör snabb automatisk skalning och enkel distributionsdefinition och hantering. Eftersom detta är ett användningsfall i realtid krävde det att P99-fördröjningen var under 10 millisekunder för att säkerställa en härlig användarupplevelse.

AWS Inferentia och AWS Neuron SDK

EC2 Inf1-instanser drivs av AWS Inferentia, den första ML-acceleratorn som byggdes av AWS för att påskynda arbetsbelastningar för djupinlärning. Inf1-instanser levererar upp till 2.3 gånger högre genomströmning och upp till 70 % lägre kostnad per slutledning än jämförbara GPU-baserade EC2-instanser. Du kan fortsätta träna dina modeller med ditt val av ramverk (PyTorch, TensorFlow, MXNet) och sedan enkelt distribuera dem på AWS Inferentia för att dra nytta av de inbyggda prestandaoptimeringarna. Du kan distribuera ett brett utbud av modelltyper med hjälp av Inf1-instanser, från bildigenkänning, objektdetektering, naturlig språkbehandling (NLP) och moderna rekommendationsmodeller.

AWS Neuron är ett mjukvaruutvecklingskit (SDK) som består av en kompilator, runtime och profileringsverktyg som optimerar ML-inferensprestandan för EC2 Inf1-instanserna. Neuron är integrerat med populära ML-ramverk som TensorFlow och PyTorch. Därför kan du distribuera modeller för djupinlärning på AWS Inferentia med samma välbekanta API:er som ditt ramverk du väljer, och dra nytta av höjningen av prestanda och lägsta kostnad per slutledning i molnet.

Sedan lanseringen har Neuron SDK fortsatt att öka bredden av modeller som den stöder samtidigt som den fortsätter att förbättra prestandan och minska slutledningskostnaderna. Detta inkluderar NLP-modeller (BERT), bildklassificeringsmodeller (ResNet, VGG) och objektdetekteringsmodeller (OpenPose och SSD).

Distribuera på Inf1-instanser för låg latens, hög genomströmning och kostnadsbesparingar

Amazon Search-teamet ville spara kostnader samtidigt som de uppfyllde deras höga genomströmningskrav på dupliceringsförutsägelser och kravet på låg latens vid förutsägelse av frågeavsikt. De valde att distribuera på AWS Inferentia-baserade Inf1-instanser och uppfyllde inte bara de höga prestandakraven, utan sparade också upp till 85 % på slutledningskostnader.

Kundupplevda dubbletter av förutsägelser

Före användningen av Inf1, en dedikerad Amazon EMR klustret kördes med CPU-baserade instanser. Utan att förlita sig på hårdvaruacceleration var ett stort antal instanser nödvändiga för att möta det höga genomströmningskravet på 8,800 1.6 totala transaktioner per sekund. Teamet bytte till inf4xlarge instanser, var och en med 16 AWS Inferentia-acceleratorer och 4 NeuronCores (16 kärnor per AWS Inferentia-chip). De spårade den transformatorbaserade modellen för en enda NeuronCore och laddade ett läge per NeuronCore för att maximera genomströmningen. Genom att dra fördel av de 85 tillgängliga NeuronCores, minskade de slutledningskostnaderna med 2 % (baserat på den nuvarande offentliga Amazon ECXNUMX on-demand-prissättningen).

Fråga avsiktsförutsägelse

Med tanke på P99-latenskravet på 10 millisekunder eller mindre, laddade teamet modellen till alla tillgängliga NeuronCore på inf1.6xlarge-instanser. Du kan enkelt göra detta med PyTorch Neuron med hjälp av torch.neuron.DataParallel API. Med Inf1-utbyggnaden var modellfördröjningen 3 millisekunder, end-to-end latens cirka 10 millisekunder och maximal genomströmning vid toppbelastning nådde 16,000 XNUMX TPS.

Kom igång med exempelkompilering och distributionskod

Följande är några exempelkoder som hjälper dig att komma igång med Inf1-instanser och inse prestanda och kostnadsfördelar som Amazon Search-teamet. Vi visar hur man kompilerar och utför slutledning med en PyTorch-modell med hjälp av PyTorch Neuron.

Först sammanställs modellen med torch.neuron.trace():

m = torch.jit.load(f="./cpu_model.pt", map_location=torch.device('cpu'))
m.eval()
model_neuron = torch.neuron.trace(
    m,
    inputs,
    compiler_workdir="work_" + str(cores) + "_" + str(batch_size),
    compiler_args=[
        '--fp32-cast=all', '--neuroncore-pipeline-cores=' + str(cores)
    ])
model_neuron.save("m5_batch" + str(batch_size) + "_cores" + str(cores) +
                  "_with_extra_op_and_fp32cast.pt")

För hela listan över möjliga argument till trace metod, se PyTorch-Neuron spårar Python API. Som du kan se, kompilatorargument kan överföras till torch.neuron API direkt. Alla FP32-operatörer är castade till BF16 with --fp32-cast=all, ger högsta prestanda samtidigt som dynamiskt omfång bevaras. Fler gjutningsalternativ finns tillgängliga för att låta dig styra prestandan för att modellera precision. Modellerna som användes för båda användningsfallen kompilerades för en enda NeuronCore (nr rörledning).

Vi laddar sedan modellen på Inferentia med torch.jit.load, och använd den för förutsägelse. De Neuron körtid laddar automatiskt modellen till NeuronCores.

cm_cpd_preprocessing_jit = torch.jit.load(f=CM_CPD_PROC,
                                          map_location=torch.device('cpu'))
cm_cpd_preprocessing_jit.eval()
m5_model = torch.jit.load(f=CM_CPD_M5)
m5_model.eval()

input = get_input()
with torch.no_grad():
    batch_cm_cpd = cm_cpd_preprocessing_jit(input)
    input_ids, attention_mask, position_ids, valid_length, token_type_ids = (
        batch_cm_cpd['input_ids'].type(torch.IntTensor),
        batch_cm_cpd['attention_mask'].type(torch.HalfTensor),
        batch_cm_cpd['position_ids'].type(torch.IntTensor),
        batch_cm_cpd['valid_length'].type(torch.IntTensor),
        batch_cm_cpd['token_type_ids'].type(torch.IntTensor))
    model_res = m5_model(input_ids, attention_mask, position_ids, valid_length,
                         token_type_ids)

Slutsats

Amazon Search-teamet kunde minska sina slutledningskostnader med 85 % genom att använda AWS Inferentia-baserade Inf1-instanser, under tung trafik och krävande prestandakrav. AWS Inferentia och Neuron SDK gav teamet flexibiliteten att optimera driftsättningsprocessen separat från träning och lägga fram en ytlig inlärningskurva via väl avrundade verktyg och välbekanta ram-API:er.

Du kan låsa upp prestanda och kostnadsfördelar genom att komma igång med exempelkoden i det här inlägget. Kolla också in slutet till slutet självstudiekurser att köra ML-modeller på Inferentia med PyTorch och TensorFlow.

Om författarna

João Moura är en AI/ML Specialist Solutions Architect på Amazon Web Services. Han är mestadels fokuserad på NLP-användningsfall och att hjälpa kunder att optimera utbildning och implementering av modeller för djupinlärning. Han är också en aktiv förespråkare för ML-specialiserad hårdvara och ML-lösningar med låg kod.

Weiqi Zhang är Software Engineering Manager på Search M5, där han arbetar med att produktisera storskaliga modeller för Amazons maskininlärningsapplikationer. Hans intressen inkluderar informationssökning och maskininlärningsinfrastruktur.

Jason Carlson är en mjukvaruingenjör för att utveckla pipelines för maskininlärning för att hjälpa till att minska antalet stulna sökvisningar på grund av kundupplevda dubbletter. Han arbetar mestadels med Apache Spark, AWS och PyTorch för att hjälpa till att distribuera och mata/bearbeta data för ML-modeller. På fritiden gillar han att läsa och springa.

Shaohui Xi är en SDE på Search Query Understanding Infra-teamet. Han leder arbetet med att bygga storskaliga djupinlärningstjänster online med låg latens och hög tillgänglighet. Utanför jobbet tycker han om att åka skidor och att utforska god mat.

Zhuoqi Zhang är en mjukvaruutvecklingsingenjör på Search Query Understanding Infra-teamet. Han arbetar med att bygga ramverk för modellbetjäning för att förbättra latens och genomströmning för djupinlärningstjänster online. Utanför jobbet gillar han att spela basket, åka snowboard och köra bil.

Haowei Sun är en mjukvaruingenjör i Search Query Understanding Infra-teamet. Hon arbetar med att designa API:er och infrastruktur som stöder djupinlärningstjänster online. Hennes intressen inkluderar service API design, infrastruktur setup och underhåll. Utanför jobbet tycker hon om att springa, vandra och resa.

Jaspreet Singh är en tillämpad forskare i M5-teamet, där han arbetar med storskaliga grundmodeller för att förbättra kundens shoppingupplevelse. Hans forskningsintressen inkluderar lärande i flera uppgifter, informationssökning och representationsinlärning.

Shruti Koparkar är Senior Product Marketing Manager på AWS. Hon hjälper kunder att utforska, utvärdera och använda EC2 accelererad datorinfrastruktur för deras maskininlärningsbehov.