Concepts Design Doc

An ElasticDL model consists of two kinds of parameters:

dense parameters, and

A dense parameter is a dense tensor with a name. An embedding table is a map from some ID to embedding vectors. An embedding table also has a name. Formalizing the concepts, we have

model = {dense parameter} + {embedding table}
dense parameter = tensor + name
embedding table = {id, tensor} + name

where the curly braces denote zero or more.

embedding table gradient

The content of dense gradient is the same as that of the dense parameter. The content of embedding table gradient is the same as that of the embedding table.

On the parameter server, we’d prefer to maintain each embedding table as a map from ID to embedding vectors. With such a data structure, it is efficient to allocate memory for new embedding vectors. On the contrary, we’d concatenate embedding vectors into protobuf messages for parameter pulling and gradient pushing. We cannot use concatenated embedding vectors as the in-memory data structure on the PS, because allocating new embedding vectors involves resize the space of concatenated embedding vectors.

Let’s make a short summary, following is all the core concepts of ElasticDL include:

model = {dense parameter} + {embedding table}
dense parameter = tensor + name

We introduce an proto message to represent the concatenated embedding vectors pulled from PS, and the concatenated embedding vectors of gradient waiting to be pushed to PS.

The definition of :

For in-memory part, we introduce an data structure.

RPC Service

Following is RPC services between PS and worker.