大规模 Embedding 训练优化实践

NVIDIA.

Large-ScaleEmbeddingTraininginTensorflow&Pytorch

HuiKang,ShijieLiu,NvidiaDevtchsoftwareengineer|2024/06/22

Agenda

·Introduction

●SparseOperationKitforTensorflow

·TorchRecOptimizationforPytorch

NVIDIA.

Introduction

Multi-stageRecommenderSystemPipeline

andScopeofThisPresentation

INFERENCEONLINE

REQUEST

TRAININGOFFLINE

RETRIEVAL

RetrieveTopKCandidates

EmbeddedInputItemorQuery

EmbeddingModelANNIndex

BuildAprox.NNIndex

TrainEmbedItems

EmbeddingModelfromCatalog

Trainingltem

DataData

FILTERING

FilterInvalidData

BloomFilter

BuildBloomFilters

Interaction

Data

AddFeaturestoCandidates

FeatureStore

BuildFeature

Store(item,user)

ItemandUserData

SCORING

Transform

Features

FeatureTfms

Feature

Transforms

TfmWorkflowData

ScoreTopKCandidates

RankingModel

TrainRankingModel

Training

Data

ORDERING

Ordering

BusinessLogic

OrderingPolicy

RESULTS

DefineOrderingPolicy

Business

Logic

OurFocUS

DeepLearningRecommendationModels(DLRM)

SimplifiedView

TopMLPs

Interaction

Bottom

MLPs

Numericalfeatures

EmbeddingTable0

Categoricalfeature0

EmbeddingTableN

CategoricalfeatureN

●BottomMLPs

oProcessdenseinputfeatures,oftenindependenttotheembeddinglookups

·Embeddingtables

oMapone-hot/multi-hotcategoricalfeaturesintothe

embeddingvectors

·Interaction

oLearneffectivefeaturecrosses

·TopMLPs

oOutputthefinalrecommendationresult,e.g.,clickprobability

NVIDIA.

GoalsofThisPresentation

•IntroduceSparseOperationKit,aTensorFlowplugintoacceleratelargescaleembeddingtraining

•IntroduceouroptimizationonTorchRec

•Sharetheachievedmulti-gputrainingperformance

SparseOperationKitForTensorflow

SparseOperationKit

·SparseOperationKit(SOK)isaPythonpackagewrappedGPUacceleratedoperationsdedicatedforsparsetraining/inferencecases.

·SOK,asapluginforTensorFlow,providesmodelparallelismfunctionalityfortheembeddingpart.

Machine0

Machine1

DP

OutputDenseValuesBatch0

OutputDenseValuesBatch1

SparseOperationKitEmbedding

OutputDenseValuesBatch0

OutputDenseValuesBatch1

DP

InputKeysBatch0

InputKeysBatch1

InputKeysBatch0

InputKeysBatch1

GPU1

GPU1

GPU0

GPU0

MP

NVIDIA.

Whatarethechallengesintheembeddingpartofmodelparallelism?

•Embeddingtablestorage

•Theembeddingtableisverylargeandrequiresasignificantamountofstoragespace.

•DuetothelimitedmemoryofGPUs,itisnotpossibletostorelargeembeddingtables.

•ManyuserschooseTensorFlow'sPSarchitecturetotrainlargeembeddingtables.

•However,whenusingthePSarchitecturefortraining,networkbandwidthoftenbecomesthebottleneckthatrestrictstrainingspeed.

•Theindicesoftheembeddingtablearegenerallyhashed,requiringtheembeddingtabletohavethecapabilitytostorehashindices.

•Theindexingoftf.Variablecannotfulfillthisfunctionality

•EmbeddinglookupOP

•Thelargenumberofembeddingtablesleadstoasignificantnumberoflookupprocesses,becomingabottleneckintraining.

•ThelookupprocessisoftenaccompaniedbymanysmallCUDAkernels,leadingtounderutilizationofGPUcomputationalresourcesandmemorybandwidth.

WhatcanSOKprovide?

•ModelParallelismforEmbeddingonmulti-gpus/multi-nodes

•Usinghashtableandhierarchicalmemory(GPUmemoryandhostmemory)tostoreembeddingtables.

•SOKusesHierarchicalKV(HKV)hashtabletostoreembeddingtables,whichcanfullyutilizebothGPUmemoryandhostmemoryforembeddingtablestorage.

•Atthesametime,asahashtable,HKVcanalsoaccepthashindicesforlookups.

•HKValsoprovidesanevictionfeature,whichcanusealgorithmslikeLRUandLFUtoevictunusedindices,preventingtheembeddingtablefromgrowingindefinitely.

•EfficientFuseLookupOperation

•SOKusesHugeCTR'sembeddinglookupasthebackendforlookups,whichcanfusethelookupsofdifferentembeddingtablestogether,reducingthelaunchoverheadofsmallkernels.

•SOK'sAPIssimilartoTensorFlowforembeddingtraining，easytomigrate.

SOKReleaseMilestone

V1.1.4(V2.0

Preview)

Supportsparse

lookupanddynamic

embeddingin

experimentalAPls

V2.0.0

Releasesparselook

upanddynamic

embeddingas

officalAPI

V1.0.0

FirstReleasewithdistributed

embeddinglookup

V2.1.0

Improvethe

stabilityand

documentationof

SOK

SOKArchitecture

·SOK'ssoftwaredependencies

FrameworkWrapper

SOKTensorFlow

EmbeddingStorageLibrary

HierarchicalDynamicEmbedding

KVStorage

EmbeddingOperatorsLibrary

HugeCTRHugeCTR

SparseLookupDenseLookup

models.

·HugeCTRisaGPU-acceleratedrecommenderframeworkdesignedfortrainingandinferenceoflargedeeplearning

·HugeCTRfusesalltheforwardlookupsintheforwardpass,reducingtheall-to-allsize.

·HugeCTRfusesallthebackwardlookupoperations'backwardoperatorsandprovidesaload-balancingWgradreductionkernel.

·Repolink:

/NVIDIA-Merlin/HugeCTR

HierarchicalKV

·HierarchicalKVisapartofNVIDIAMerlinandprovideshierarchicalkey-valuestoragetomeetRecSysrequirements

·ThekeycapabilityofHierarchicalKVistostorekey-value(feature-embedding)onhigh-bandwidthmemory(HBM)ofGPUsandinhostmemory.

HierarchicalKV

Bucketinfo

bucketO

key

score

digest

value

int64_tuint64_t

1byteDigests

64/128floats

Keys

Scores

Values

bucketn-1

Storageofbucketi(SOA)

EachKey-Valuepairhasadigest:

64bitshashedkey

bucketindex

HKVBuckets(AOS)

hashed

bucket

index

bucket

inner-offset

64bitskey

bucket

inner-offset

bucketi

unused

64bits

digest

hash

key

Digest

·For:accelerate

theprobingprocess.

1byte

·Howitwork:alldigestsinabucketformafilter.

·Repolink:

/NVIDIA-Merlin/HierarchicalKV

②nVIDIA.

HierarchicalKV

•HierarchicalKVisapartofNVIDIAMerlinandprovideshierarchicalkey-valuestoragetomeetRecSysrequirements.

•ThekeycapabilityofHierarchicalKVistostorekey-value(feature-embedding)onhigh-bandwidthmemory(HBM)of

API

GPUsandinhostmemory.

insert_or_assign

insert_and_evict

HKVhashtable

find_or_insert

assign

accum_or_assign

find

•HKVAPIlink:/NVIDIA-Merlin/HierarchicalKV#api-maturity-matrix

SOKAPls

·SOKsupportsthefollowingAPl:

·sok.lookup_sparse

·AlookupAPIthatisconsistentwiththefunctionalityandparametersofthetf.embedding_lookup_sparseAPI.

·sok.Variableandsok.DynamicVariable

■Inheritsfromtf.Variable,providingautomaticdistributionfunctionality.

■sok.DynamicVariableprovideshashtableandevictionfunctionality.

·sok.dump/sok.load

■Automaticallyparalleldump/loadembeddingdatatofilesystem.

·Sok.incremental_dump

■SOKprovidesanincrementaldumpfeaturethat,basedontheuser-inputtimethreshold,dumpstheindicesandvaluesupdatedafter

thetimethreshold.

FormoreSOKfeatures,pleasereadtheSOKdocument:

https://nvidia-merlin.github.io/HugeCTR/sparseoperationkit/master/index.html

nVIDIA.

SOKAPlsexample

Importsparseoperationkitassok

V1=sok.DynamicVariable(dimension=16,initializer=str(18)))

v2=sok.DynamicVariable(dimension=8,initializer=str(18)))

indices1=tf.SparseTensor(

indices=[[0,0],[0,1],[1,0],[1,1],[1,2]],

values=[1,1,3,4,5],

denseshape=[2,3]

indices2=tf.SparseTensor(

indices=[[0,0],[1,0],[1,1]],

values=[1,2,3],

denseshape=[2,2]

emb1

emb2=tf.nn.

withtf.GradientTape()astape:equalsto

emb1,emb2=sok.lookup_sparse(

[v1,v2],

[indices1,indices2],

combiners=["sum","mean"]

loss=tf.reducesum(emb1)+tf.reducesum(emb2)

grads=tape.gradient(loss,[v1,v2])

createSOKvariables,similarbehaviorwithtf.Variable

=tf.nn.embedding_lookup_sparse(v1,indices1,combiner="sum")embeddinglookupsparse(v2,indices2,combiner="mean")

optimizer=tf.keras.optimizers.SGD(learningrate=1.0)optimizer.apply_gradients(zip(grads,[v1,v2]))

USECASE

·MomoacceleratestherecommendationsystemwithSOKwith5xspeedupe2e.

o

/blog/nvidia-merlin-helps

-momo-recommender-system-achieve-high-

performance-training-and-optimization/

·IntegratedintoAlibaba'sDeepRecasamulti-node

multi-GPUembeddingsynchronoustraining.

multi-GPUembeddingbackend,supportingmulti-node.

o

https://deeprec.readthedocs.io/en/latest/SOK.html

·InOppo'sexperiment:

oCPUPSe2etime760ms/GPU(original)PSe2e

time450ms->SOKe2etime300ms,e2e

performanceimprovementof60%/33%

oEmbeddingpartfrom250ms(GPUOriginal)->85ms(SoK),performanceimproveabout67%.

speedup

性能加速效果

10

2

0

SoK+通信优化+IO优化

SoK+通信优化+IO优化+XLA+其他

USECASE

·MomoacceleratestherecommendationsystemwithSOKwith5xspeedupe2e.

momo-recommender-system-achieve-high-

EndtoEndTrainingTime(ms/iter)760ms

800

700

600

500

450ms

400

300ms

300

250ms

200

100

0

85ms

WithCPUPSWithOPPO'soriginalGPUPSWithSoK

o

/blog/nvidia-merlin-helps

-

performance-training-and-optimization/

·IntegratedintoAlibaba'sDeepRecasamulti-node,

multi-GPUembeddingsynchronoustraining.

o

https://deeprec.readthedocs.io/en/latest/SOK.html

multi-GPUembeddingbackend,supportingmulti-node,

·InOppo'sexperiment:

oCPUPSe2etime760ms/GPU(original)PSe2e

time450ms->SOKe2etime300ms,e2e

performanceimprovementof60%/33%

oEmbeddingpartfrom250ms(GPUOriginal)->85ms(SoK),performanceimproveabout67%.

nVIDIA.

TorchRecOptimizationForPytorch

TorchRecIntroduction

Basic

•TorchRecisaPyTorchdomainlibraryfortrainingrecommendersystemsatscale,whichallowsuserstotrainmodelswithgiganticembeddingtablesshardedacrossmanyGPUs

•/pytorch/torchrec

•Ouroptimizationgoalis

oTargetedonMLPerfDLRM-DCNv2model,16nodescases

oDontbreakitsoriginalAPI

oDontbreakitsoriginalarchitecture.Preferhigh-level,lessmodificationchanges.

oSomeoptimizationsaremodel-architecturedependent.E.g.cudagraph.

TorchRecIntroduction

EmbeddingConfig

EmbeddingBagConfig

TrainPipelineSparseDist

Architecture

DistributedModelParallel

TorchRecAPI

EmbeddingShardingPlanner

ShardedEmbeddingBagCollectionShardedEmbeddingCollection

TorchRec

...

InternalModule

BaseSparseFeaturesDistBaseEmbeddingDistBaseEmbeddingLookup

SplitTableBatchedEmbeddingBagsCodegen

block_bucketize_sparse_features

FBGEMM

DenseTableBatchedEmbeddingBagsCodegen

•ThearchitectureofTorchReccanbespecifiedasthreelayers

oAPIlayer:aneasy-to-usewrapperforuserstoconfiguredifferenttypesofembeddings,shardingstrategies,andtoperformshardinganddooverlap

oTorchRecInternalModule:itconsistsofPythonlevelimplementation,whichbasicallyarenn.Module.

oFBGEMM:ahigh-performanceC++operator,usedforembeddinglookup,wgradreduction,andfeatureinputdistribution.

#1.ConstructDLRMmodule

ebconfigs=[

EmbeddingBagConfig(

name=tableidxtoname(featureidx),

embeddingdim=args.embeddingdim,

num_embeddings=args.numembeddingsperfeature[featureidx],

featurenames=[featureidxtoname(featureidx)],

forfeatureidxinrange(args.numembeddingtable)

trainmodel=DLRMDCNTrain(

embeddingbagcollection=EmbeddingBagcollection(

tables=ebconfigs,device=torch.device("meta")),

denseinfeatures=args.denseinfeatures,

densearchlayersizes=args.densearchlayersizes,

overarchlayersizes=args.overarchlayersizes,

dcnnum_layers=args.dcnnumlayers,

dcnlowrankdim=args.dcnlowrankdim,

batchsize=args.batchsize,

densedevice=device,

#2.embeddingsharding

model=DistributedModelParallel(module=trainmodel,

device=device,

plan=plan,

(

#3.Overlapcomm/compute/devicetransferduringtrainingthroughtrainpipeline

trainpipeline=TrainPipelineSparseDist(

model,

nonfusedoptimizer,

device,

)

trainiterator=iter(CriteoRecDataset(...))

foriinrange(int(args.numiterations)):

trainpipeline.progress(trainiterator)

TorchRecexample

AP

·ThisisanexampletoshowhowtouseTorchRecAP

·Thereare3steps:

oUseEmbeddingBagConfigandEmbeddingBagCollectiontoconstructembedding

oUseDistributedModelParalleltodosharding.ThedefaultbehaviorisdomodelparallelonembeddingandforotherpartdodataparallelorFSDP.Ithasinterfaceforusto

customizeifwewanttosupportothertypeofparallelism.

oUseTrainPipelineSparseDisttoconstructtrainingpipeline.

·MostofouroptimizationsonlyleadtochangesinsideTrainPipelineSparseDist.

nVIDIA.

MLPerfDLRM-DCNv2TrainingPerformance

Train4BillionSamplesin1Minute

MLPerfDLRM-DCNv2TrainingPerformanceon16DGX-H100Nodes

NativeTorchRec

OptimizedTorchRec

MLPerfRecord

0

7.66

2.3

3.4

2

6

iterationtime(ms)

*LowerisBetter

since

●DLRM-DCNv2:anewMLPerfrecommendationbenchmark

thev3.0submissionwiththefollowingupdates:

oCriteo4TBmulti-hotdataset

o3Crosslayers

oAdagradoptimizer

o26Bparameters

·Weachieved2.25×speedupafterouroptimizationon16DGXH100nodeswithbatch_size_per_gpu=1024(7.68ms->3.4msgputrainingtimeperiter)

·Asareference,MLPerfrecordtakes2.3msperiteration.

Mainstream

Commstream0

Commstream1

Inputdiststream

Inputdistcommstream

TorchRecOptimization

Traintimelinebreakdown

EmbfwdBMLPfwd|TMLPfwd+bckBMLPbckEmbbck+EmboptMLPopt

allreduceallreduceallreduce

Emba2a3Emba2a

Nextiterationprefetch

bucketizeD2H+syncD2H+sync

inputdista2a0inputdista2a1inputdista2a2

·Thereare5partsafterbreakdown,identifiedinpicwithMLPoperator

1.Inputfeaturedist:convertdata-parallelinputfeatureasmodel-parallel

2.Embforward:doingforwardlookupandemball-to-all.EachGPUwillgetdata-parallelembeddingvector

3.MLPforward&backward&allreduce:

BMLPisbottomMLPandTMLPistopMLP

4.Embbackward:backwardemball-to-alltoconvertdata-parallelgradientstomodel-parallel,anddoesgradientaccumulationandembeddingtableupdate

5.MLPoptimizer:MLPupdate

·Dependency:(excuteserially,while(①canbeprefetchedincurrentiteration.

Embeddingoperator

Communication

Host-devicesyncandcopy

TorchRecOptimization

Overview

•Optimizecpulaunchlatency

•CudagraphMLPandallreduce:MLPandallreducecanbegraphed.

•Multi-threadingkernellaunchforinputdistprefetch:sothatwheninputdistisdoingsync,forwardandbackwardcanstilllaunchkernels.

•PinnedD2Hcopy:makeD2HcopyinInputdistaspinnedhostmemorycopy

•TORCH_NCCL_AVOID_RECORD_STREAMS:envvariablethatcanreducepytorchallocatoractivitywhenusingNCCL

•Optimizeinputdistworkload

•Removeinputdista2a0andrelatedD2H+sync:firsta2aresultcanbereusedwhenbatch_size_per_gpuisnotchanged.

TorchRecOptimization

Embfwd

Emba2a

2

Embbck+Embopt

allreduce

Emba2a

4

CudagraphMLPandallreduce

MLPopt

Mainstream

Commstream0

Commstream1

5

BMLPfwdTMLPfwd+bckBMLPbck

allreduce

3

allreduce

Inputdiststream

Inputdistcommstream

bucketizeD2H+syncD2H+sync

inputdista2a0

inputdista2a1

inputdista2a2

1

Nextiterationprefetch

MLPoperator

EmbeddingoperatorCommunication

Host-devicesyncandcopy

•MLPandAllreducecanbegraphed,embeddingandinputdistcannotbegraphed

oBecauseMLPandallreducedoesnotrelyondynamicvalueandembeddingdoes.

TorchRecOptimization

Multi-threadingkernellaunchforinputdistprefetch

Embfwd

Emba2a

2

Embbck+EmboptMLPopt

Mainstream

Commstream0

Commstream1

Inputdiststream

5

BMLPfwdTMLPfwd+bckBMLPbck

allreduceallreduceallreduce

4

Nextiterationprefetch

3Emba2a

bucketize

D2H+syncD2H+sync

Inputdistcommstream1inputdista2a0inputdista2a1inputdista2a2

•PyTorchusessinglethreadtolaunchkernel.

•Thelaunchorderis:2->3->1->4->5

oWelaunch1inthemiddleinsteadofendbecausewecanmakecommunicationbetteroverlapwithcomputation.

•TheissueisD2H+syncininputfeaturedist,makingCPUwaitingfortheGPUandcannotlaunchkernel

MLPoperator

EmbeddingoperatorCommunication

Host-devicesyncandcopy

•Wecanuseaseperatethreadtolaunch1tomakekernellaunchin4and5notblockedbyD2H+sync

TorchRecOptimization

Embfwd

Emba2a

2

Embbck+Embopt

allreduce

Emba2a

4

PinnedD2Hcopy

MLPopt

Mainstream

Commstream0

Commstream1

5

BMLPfwdTMLPfwd+bckBMLPbck

allreduce

3

allreduce

Inputdiststream

Inputdistcommstream

D2H+sync

inputdista2a1inputdista2a2

inputdista2a0

bucketize

D2H+sync

1

EmbeddingoperatorCommunication

Host-devicesyncandcopy

Nextiterationprefetch

MLPoperator

•TorchRecusestensor.tolist()todoD2Hininputfeaturedist

•However,thehostdataisunpinnedmemoryandwillcausebadperf

•Soinsteadofusingdirectlyusingdevicetensor.tolist(),oWefirstcreateapinnedhosttensor

oThenhost_tensor.copy_(device_tensor).tolist()

TorchRecOptimization

TORCH_NCCL_AVOID_RECORD_STREAMS

#Scenario0,unsafe

a=torch.rand(...,

device="cuda")

s.wait_stream(

torch.cuda.current_stream()

)

withtorch.cuda.stream(s):

usea...

dela#sisstillusinga!

laterops...

#Scenario1,safebutinefficient

a=torch.rand(...,device="cuda")

s.wait_stream(

torch.cuda.current_stream()

)

withtorch.cuda.stream(s):

a.record_stream(s)

usea...

dela#willwaituntilworkinshas

beenfinished

laterops...

#Scenario2,safeandeffecient

a=torch.rand(...,device="cuda")

s.wait_stream(

torch.cuda.current_stream()

)

withtorch.cuda.stream(s):

usea...

torch.cuda.current_stream().wait_stream(s)#thissyncmakessurethata

willbedeallocateafterbeingusedonstreams.

dela

laterops...

•record_stream:Marksthetensorashavingbeenusedbythisstream.

oWhenthetensorisdeallocated,ensurethetensormemoryisnotreusedforanothertensoruntilallworkqueuedonstreamatthetimeofdeallocationiscomplete.

oThesesafetyguaranteescomeatsomeperformanceandpredictabilitycost(analogoustothetradeoffbetweenGCandmanualmemorymanagement),suchaspollingonevent

•Ifuserhasalreadyusestreamsynchronize,theydon’tneedrecord_stream.SeeScenario2.

•PyTorchNCCLwilldoesrecord_streambydefault,whichcanberemovedifwe