VLDB2022_HET：Scaling out Huge Embedding Model Training via Cache-enabled Distributed Framework_腾讯云数据库.pdf

tinge

9页

5次

2023-08-29

免费下载

HET: Scaling out Huge Embedding Model Training via

Cache-enabled Distributed Framework

Xupeng Miao

∗1

, Hailin Zhang

∗1

, Yining Shi

, Xiaonan Nie

, Zhi Yang

, Yangyu Tao

, Bin Cui

1,3

Department of Computer Science & Key Lab of High Condence Software Technologies (MOE), Peking University

Institute of Computational Social Science, Peking University (Qingdao),

Tencent Inc.

Center for Data Science, Peking University & National Engineering Laboratory for Big Data Analysis and Applications

{xupeng.miao, z.hl, shiyining, xiaonan.nie, yangzhi, bin.cui}@pku.edu.cn,

brucetao@tencent.com

ABSTRACT

Embedding models have been an eective learning paradigm for

high-dimensional data. However, one open issue of embedding

models is that their representations (latent factors) often result in

large parameter space. We observe that existing distributed train-

ing frameworks face a scalability issue of embedding models since

updating and retrieving the shared embedding parameters from

servers usually dominates the training cycle. In this paper, we pro-

pose HET, a new system framework that signicantly improves the

scalability of huge embedding model training. We embrace skewed

popularity distributions of embeddings as a performance oppor-

tunity and leverage it to address the communication bottleneck

with an embedding cache. To ensure consistency across the caches,

we incorporate a new consistency model into HET design, which

provides ne-grained consistency guarantees on a per-embedding

basis. Compared to previous work that only allows staleness for

read operations, HET also utilizes staleness for write operations.

Evaluations on six representative tasks show that HET achieves

up to 88% embedding communication reductions and up to 20

performance speedup over the state-of-the-art baselines.

PVLDB Reference Format:

Xupeng Miao, Hailin Zhang, Yining Shi, Xiaonan Nie, Zhi Yang, Yangyu

Tao, Bin Cui. HET: Scaling out Huge Embedding Model Training via

Cache-enabled Distributed Framework . PVLDB, 15(2): 312 - 320, 2022.

doi:10.14778/3489496.3489511

PVLDB Artifact Availability:

The source code of this research paper has been made publicly available at

https://github.com/PKU-DAIR/Hetu/.

1 INTRODUCTION

To train a model on high-dimensional data, such as words in a cor-

pus of text [

] or the user-item interaction data [

it is common to use an embedding model, which projects a sparse

high-dimensional feature space, into a continuous low-dimensional

embedding space. For example, in a language model, a training ex-

ample might be a sparse vector with non-zero entries corresponding

to the IDs of words in a vocabulary, and the distributed represen-

tation for each word will be a lower-dimensional vector. “Wide

∗

Equal contribution.

This work is licensed under the Creative Commons BY-NC-ND 4.0 International

License. Visit https://creativecommons.org/licenses/by-nc-nd/4.0/ to view a copy of

this license. For any use beyond those covered by this license, obtain permission by

emailing info@vldb.org. Copyright is held by the owner/author(s). Publication rights

licensed to the VLDB Endowment.

Proceedings of the VLDB Endowment, Vol. 15, No. 2 ISSN 2150-8097.

doi:10.14778/3489496.3489511

and deep learning” [

] creates distributed representations from

cross-product transformations on categorical features. Embedding

model is common at modern web companies (e.g., Facebook [

Google [

] and Tencent [

]), which have been recognized as

an eective learning paradigm to extract useful information for

downstream tasks such as recommendation.

As each feature needs to be represented by a set of embeddings

(i.e., latent vectors), many embedding models are at a giant scale and

are too large to copy to a worker on every use, or even to store in

RAM on a single host. For instance, the parameters of a real-world

document embedding model in Google [

] occupies several

terabytes, and the industrial click-through rate prediction model in

Baidu [

] has 10

input sparse features and also requires 10 Tb

parameters. For this reason, it is challenging to scale embedding

models up to large-scale use cases, in which millions or even billions

of parameters need to be learned.

Modern distributed ML systems (e.g., TensorFlow [

]) typically

adopt the parameter server [

] framework to scale out models.

The server usually maintains the globally shared parameters by

aggregating updates from the workers and updating the global

parameters. Workers communicate only with the server nodes,

updating and retrieving the shared parameters. Existing ML systems

usually support data parallelism where a worker usually contains a

replica of the ML model and is assigned an equal-sized partition of

the entire training data. Bulk Synchronous Parallel (BSP) [

] or

Asynchronous Parallel (ASP) [

] are usually adopted for updating

the model parameters during distributed training.

However, this setup faces a scalability issue for large embed-

ding models [

]. We observe that the greatest ineciency

comes from updating and retrieving the shared feature embedding

parameters through a limited bandwidth link. For example, using

TensorFlow with ASP, up to 86% of training time is spent on em-

bedding fetching and updating, which dominates the training cycle.

The major reason is that an embedding model often uses deep neu-

ral networks with low computational complexity, comparing with

the giant embedding data. Accordingly, the computation takes a

much shorter time than the reads and writes of remote embedding

data. Moreover, with the increasing gap between emerging pow-

erful accelerators and the slow growth of network bandwidth, the

embedding communication bottleneck would become even more

severe. To our knowledge, there is little prior work addressing the

scalability issue of embedding models in a distributed environment.

In this paper, we propose

HET

, a novel distributed system frame-

work to scale

uge

mbedding model

raining. Our key idea is

to exploit an ecient

embedding-cache-enabled

architecture,

arXiv:2112.07221v1 [cs.LG] 14 Dec 2021

Sparse

Feature

Dense Feature

Embedding Lookup

Prediction

Dense Part

Sparse Part

𝑒

𝑊

𝑥

𝑆 ×𝐷

Figure 1: Illustrate of embedding model architecture

which is mainly inspired by the critical characteristics for embed-

ding models: popularity skewness and staleness tolerant. Specically,

the popularity distribution of embeddings is often highly skewed,

typically following power-law distributions [

], implying a perfor-

mance opportunity: a small cache of hot embedding at each worker

can eectively save the network bandwidth while scaling training

throughput with the number of workers.

Replicating shared embedding data in multiple caches raises the

problem of consistency in the presence of writes. Fortunately, em-

bedding models are iterative convergent algorithms where some

staleness errors during training are acceptable and will not pre-

vent convergence. In other words, embedding models are robust

to a bounded amount of inconsistency (e.g., reading out-of-date

shared state). By relaxing the consistency guarantees properly, we

can exploit the opportunity of caches to gain signicant system

improvements. However, conventional relaxed consistency models

such as Stale Synchronous Parallel (SSP) [

] are not aware of the

presence of skew access and require that every single worker should

be able to hold an entire set of parameters. Moreover, they mainly

target straggler problems rather than communication overhead. For

example, SSP maintains an up-to-date global model through sent

out write updates to servers each clock. Considering the large scale

and communication cost of embedding models, the above issues

become a critical limitation for scaling out the training.

To address the above issues, we incorporate a new consistency

model into HET. Our consistency model diers from the traditional

ones in two aspects. First, we enable the

ne-grained

caching and

consistency that provides guarantees on a per-embedding basis.

Specically, for each cached embedding, we leverage an embedding-

specic Lamport clock to manage its ne-grained consistency ac-

tions (e.g., validation, synchronization) and provide the concept

of “per-embedding-clock-bounded” consistency — a worker can

see all updates of an embedding older than certain embedding-

specic clocks. We provide a ne-grained consistency model and

theoretically prove its convergence guarantees. Second, compared

to previous works that only allow staleness for read operations, we

further utilizes staleness for writes, allowing

stale-writes

based

on the timestamp deviation between the global and local clocks of

each embedding. This means that writing to an embedding residing

in the cache does not update the underlying global model until the

embedding is invalided or evicted from the cache. This feature is

critical in reducing communication overheads.

We summarize our contributions as follows: First, we reveal the

performance bottleneck and the opportunity for scaling huge em-

bedding models, and introduce a novel

system abstraction

with

embedding cache. Second, we employ a new cache

consistency

model

that provides (1) clock-bounded consistency at the ne-

grained of each embedding, and (2) allows staleness for both caches

read and write operations for minimizing communication overhead.

Finally, we build

HET system

, a new framework that implements

the proposed system abstraction and the consistency model, and

supports 10

parameters scale embedding model training, achiev-

ing 6

−

speedup and up to 88% embedding communication

reduction over the state-of-the-art baseline systems.

2 PRELIMINARY

2.1 Distributed Training

Distributed machine learning.

The target of machine learn-

ing is to nd a model

x ∈ R

𝑑

(

𝑑

is the total number of parameters

in the model) that minimizes the empirical risk:

min

𝐹 (x) :=

|𝜉 |



𝑖

𝑓 (x; 𝜉

𝑖

)

, (1)

where

𝑓 (·)

is the loss function,

𝜉

is the training dataset and

𝜉

𝑖

rep-

resents the

𝑖

-th data sample. Distributed ML systems have been

extensively studied in recent years to scale up ML for big data and

large models. Parameter Server (PS) is a trendy data parallelism ar-

chitecture for many existing systems (e.g., TensorFlow [

], PS2 [

]).

Another choice is All-Reduce and several recent systems (e.g., Py-

Torch [

], Horovod [

]) show superior performance over PS with

the help of NCCL [4], especially for dense models.

Parallel training paradigms.

Most of data parallelism studies

manage to keep consistent model performance as the standalone

mini-batch SGD. BSP assumes that all

𝑁

workers are fully synchro-

nized and performing the following update rule:

x(𝑡 + 1) = x (𝑡) −𝜂

𝑁



𝑖=1

𝐺

𝑖

(x (𝑡); 𝜉

𝑖

)

, (2)

where

𝜂

is the learning rate,

𝜉

𝑖

are randomly sampled from the

training set, and

𝐺

𝑖

(·)

denotes the gradient from the

𝑖

-th worker.

However, the frequent synchronization and straggler problem bring

signicant communication costs. ASP avoids such overheads by

allowing the workers to proceed without waiting for each other.

But the model degradation happens because of the stale gradients.

To balance the trade-o between training eciency and model

performance, SSP and several variants [

] have been proposed. It

has been proved that they could share the same convergence rate

with BSP when the staleness is upper bounded [

]. Unfortunately,

SSP requires to store the replication of the entire model inside every

single worker, which is impractical for giant models.

2.2 Embedding Models

Many types of embedding models (e.g., Wide & Deep [

], Deep

& Cross [

], DeepFM [

], xDeepFM [

] and Deep Interest Net-

work [

]) have been developed for high-dimension data, and have

achieved widespread success in recommender systems–a critical

service for internet companies. Figure 1 illustrates a common em-

bedding model architecture. In order to handle categorical data,

embedding tables map categorical features to dense representations

in an abstract space. In particular, each embedding lookup may be

interpreted as using a one-hot vector

𝑒

𝑖

(with the

𝑖

-th position being

1 while others are 0, where index

𝑖

corresponds to

𝑖

-th category)