<commit_before># Embedding and Partitioned Variable in TF 2.0


| Status        | Proposed                                                                  |
| :------------ | :------------------------------------------------------------------------ |
| **Author** | <NAME> (<EMAIL>)                                        |
|**Sponsors**| <NAME> (<EMAIL>), <NAME> (<EMAIL>), <NAME>  (<EMAIL>) |
|**Updated** | 2019-1-16                                                                 |


Summary: this RFC describes the design and interfaces for embeddings and partitioned variables in TF 2.0 based on Distribution Strategy. We propose to support embeddings and sharded layers with DistributionStrategy and Keras' layer API. We also propose to shard embeddings and layers only in `ParameterServerStrategy` and mirror them in `MirroredStrategy` or `CollectiveAllReduceStrategy.` 


## Background


### Embeddings


An [embedding](https://www.tensorflow.org/guide/embedding) is a mapping from discrete objects, such as words, to vectors of real numbers. In TensorFlow implementation, we use a variable to hold all these vectors for a feature column. All vectors in an embedding variable have the same size but the sizes of embedding variables may vary. 

We don't always use the whole embedding variable in a training step, the result of embedding lookup is <code>[IndexedSlices](https://www.tensorflow.org/api_docs/python/tf/IndexedSlices)</code> which is essentially a dict mapping from indices to vectors. So is the gradient to an embedding variable.


#### Sharded Embedding


Sharded embedding is one use case of `PartitionedVariable`. People sometimes shard embedding variables across multiple parameter servers. There are many reasons why people shard embeddings:

1.  parallelize computations : we can do embedding lookups in parallel and at the same time keep communication between machines small; 
1.  balance parameter servers' load: if an embedding layer is large and placed on one ps, that ps can become a hotspot;
1.  shard storage: occasionally embedding variables can be too large to fit in the memory of one machine.
1.  reduce contentions in asynchronous training: updates from different workers could possibly end up on different partitions which reduces possibility of overriding each other's update and using locking, although this depends on the implementation of embedding update.


##### Partition Strategy: div or mod


There are two strategies to look up embedding vectors on a sharded embedding variable: "div" and "mod". See an explanation [here](https://www.tensorflow.org/api_docs/python/tf/nn/embedding_lookup). "Mod" strategy may be useful when references to embedding slices are not evenly distributed over their indices.  However, "mod" strategy is a poor approximation to the actual load balancing users want. For example, when vocabulary is sorted by frequency, with "mod" strategy the first parameter server always has larger load than the second one. Furthermore, its current checkpointing mechanism also prevents users from migrating to a cluster with different number of parameter servers.


#### Current Support in Distribution Strategy


Sharded embeddings and non-sharded embeddings are both supported by `ParameterServerStrategy` with TF 1.x API. See this [PR](https://github.com/tensorflow/tensorflow/pull/23254). Sharded embeddings are not supported by other distribution strategies.

Non-sharded embeddings are currently supported by `MirroredStrategy` which mirrors them on all replicas. A [naive implementation of allgather](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/distribute/cross_device_utils.py#L648) exists in TensorFlow codebase to collect their gradients onto each replica. See [the transformer model](https://github.com/tensorflow/models/tree/master/official/transformer) for the usage of embeddings with `MirroredStrategy`.

However, embeddings are not supported by multi-worker `CollectiveAllReduceStrategy` unless we implement multi-worker allgather primitive in collective ops or always cast their updates to dense tensors and use allreduce to aggregate gradients.

Note: our design proposal below assumes there is an efficient implementation of allgather but any detail of allgather is out of this doc's scope. We'll have allgather implementation in collective ops soon.


### Partitioned Layer


A layer of a network can sometimes be partitioned. We call it partitioned layer in this document. This is another use case of `PartitionedVariable`.

The main reason people partition a layer is to balance parameter servers' load.

Some layers can be huge. For example the logits layer can be extremely large when there are large number of classes. People may want to shard their computations accordingly to avoid concatenating them, e.g. the implementation of `tf.nn.nce_loss`.

Even for mid-sized layer, people want to partition them to have more even distribution of parameter servers' load because people may get bad luck when variables are assigned to parameter servers in a round-robin fashion. See this [issue](https://github.com/tensorflow/tensorflow/issues/24953).

Partitioned layer is only supported by lower level APIs with or without `ParameterServerStrategy` in 1.x.


### Current Usage of PartitionedVariable


With TF 1.x API users can create a `PartitionedVariable` via `tf.get_variable` API:

```python
partitioned_var = tf.get_variable(
    "embedding",
    shape=(10000, 128),
    dtype=tf.float32,
    partitioner=tf.min_max_variable_partitioner(
        max_partitions=num_ps_replicas, axis=0))
```

The `partitioned_var` here is an instance of `PartitionedVariable`.

Alternatively users can use `tf.variable_scope` to apply `partitioner` to all variables created under its scope:

```python
with tf.variable_scope(partitioner=tf.min_max_variable_partitioner(
        max_partitions=num_ps_replicas, axis=0)):
  tf.keras.layers.Embedding(...)
```

Both these cases will also work when they are under a `ParameterServerStrategy`'s scope. However `MirroredStrategy` and `CollectiveAllReduceStrategy` will [error out](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/distribute/distribute_lib.py#L872) if a `PartitionedVariable` is about to be created.

These APIs `tf.get_variable` and `tf.variable_scope` will be gone in TF 2.0.


### Possible Problems with The `PartitionedVariable` Class 


Although we have a class called `PartitionedVariable`,  we often concatenate all of their components to pretend it is a single variable.

For example, `tf.matmul(partitioned_var, inputs)` would concatenate the list of variables in the `partitioned_var` before the execution of `matmul` as if the `partitioned_var` is a single variable. With the concatenation, the `partitioned_var` would only serve for loadbalancing purposes.

To avoid issue, it requires people to write methods that respect the partitions, e.g. `tf.nn.sampled_softmax`. This has some flavor of model parallelism.


## Goals


1.  Support partitioned embedding and partitioned layer in TF 2.0 in parameter server architecture via Distribution Strategy's and Keras layer's API. 
1.  Better support for embeddings in mirrored and collective allreduce architecture. We will not shard them in this architecture in our pre-TF 2.0 design.
2.  We will only support "div" partition strategy but as a post-TF 2.0 work we will support re-balancing embeddings in the Keras' Embedding layer. That means we will not support "mod" partition strategy any more.
3.  We also would like the API to enable people to easily switch between different architectures without needing to modify their models.


## Non-goals


1.  We don't have plan to support any flavor of model parallelism beyond the current implementations of embedding lookup and loss functions that respect partitioning.
2.  Support for partitioned variable in `ParameterServerStrategy` in eager mode is not a near-term goal since parameter server support in eager runtime won't be ready in the near future.
3.  We will not support any flavor of online learning or liquid sharding, e.g. dynamically placing or resizing partitions.


## Design Proposal


### Overview


Since Distribution Strategy is a library that is designed to parallelize and scale up models. We think that support for embeddings and partitioned layers should exist in Distribution Strategys. We'll support them differently for different Distribution Strategy:

<table>
  <tr>
    <td></td>
    <td>Embeddings</td>
    <td>Partitioned Layer</td>
  </tr>
  <tr>
    <td>Mirrored/ CollectiveAllReduce Strategy</td>
    <td>
      Pre-TF 2.0:<br />
      w/ Keras' Embedding layer API: won't be sharded but will be mirrored on each replica. 
      <p><br />
      Post-TF 2.0:<br />
      w/ Distribution Strategy's API directly: allow sharding; being placed on host; or dense update
      <p>
      Heuristics in Embedding layer to auto-set performance improvement hints.
    </td>
    <td>
      Pre-TF 2.0:<br />
      w/ Keras' layer API: won't be sharded but will be mirrored on all replicas, just like a normal layer.
      <p><br />
      Post-TF 2.0:<br />
      w/ Distribution Strategy's API directly: allow being mirrored on each host
      <p>
      Heuristics in Embedding layer to auto-place larger embeddings on host.
    </td>
  </tr>
  <tr>
    <td>ParameterServer Strategy</td>
    <td>
      Pre-TF 2.0:<br />
      w/ Keras' Embedding layer API: probably sharded, according to <code>partitioner</code> passed to its constructor.
     <p><br />
     Post-TF 2.0:<br />
w/ Distribution Strategy's API directly: allow creating partitioned variable with different partitioners
    </td>
    <td>
      Pre-TF 2.0:<br />
      w/ Keras' layer API: probably sharded, according to <code>partitioner</code> passed to its constructor which applies globally to all layers.
      <p><br />
      Post-TF 2.0:<br />
      w/ Distribution Strategy's API directly: allow creating partitioned variable with different partitioners.
    </td>
  </tr>
</table>

In the `ParameterServerStrategy`, the goal of supporting partitioned variable is to match the existing behavior in 1.x. We will propose its APIs for TF 2.0 below. 

**In the mirrored and collective allreduce architecture, we will only mirror embeddings or logits layers, either on devices or on host.** 

In the future, we will probably support a new strategy that is a hybrid of these two architectures in the future.

Please refer to the discussion in [Alternatives Considered](#Alternatives-Considered) section for why we've made these decisions.


### Proposed Change for pre-TF 2.0


The following is planned to be done before TF 2.0's release.

<commit_msg>Use markdown language for the table instead of HTML.<commit_after># Embedding and Partitioned Variable in TF 2.0


| Status        | Proposed                                                                  |
| :------------ | :------------------------------------------------------------------------ |
| **Author** | <NAME> (<EMAIL>)                                        |
|**Sponsors**| <NAME> (<EMAIL>), <NAME> (<EMAIL>), <NAME>  (<EMAIL>) |
|**Updated** | 2019-1-16                                                                 |


Summary: this RFC describes the design and interfaces for embeddings and partitioned variables in TF 2.0 based on Distribution Strategy. We propose to support embeddings and sharded layers with DistributionStrategy and Keras' layer API. We also propose to shard embeddings and layers only in `ParameterServerStrategy` and mirror them in `MirroredStrategy` or `CollectiveAllReduceStrategy.` 


## Background


### Embeddings


An [embedding](https://www.tensorflow.org/guide/embedding) is a mapping from discrete objects, such as words, to vectors of real numbers. In TensorFlow implementation, we use a variable to hold all these vectors for a feature column. All vectors in an embedding variable have the same size but the sizes of embedding variables may vary. 

We don't always use the whole embedding variable in a training step, the result of embedding lookup is <code>[IndexedSlices](https://www.tensorflow.org/api_docs/python/tf/IndexedSlices)</code> which is essentially a dict mapping from indices to vectors. So is the gradient to an embedding variable.


#### Sharded Embedding


Sharded embedding is one use case of `PartitionedVariable`. People sometimes shard embedding variables across multiple parameter servers. There are many reasons why people shard embeddings:

1.  parallelize computations : we can do embedding lookups in parallel and at the same time keep communication between machines small; 
1.  balance parameter servers' load: if an embedding layer is large and placed on one ps, that ps can become a hotspot;
1.  shard storage: occasionally embedding variables can be too large to fit in the memory of one machine.
1.  reduce contentions in asynchronous training: updates from different workers could possibly end up on different partitions which reduces possibility of overriding each other's update and using locking, although this depends on the implementation of embedding update.


##### Partition Strategy: div or mod


There are two strategies to look up embedding vectors on a sharded embedding variable: "div" and "mod". See an explanation [here](https://www.tensorflow.org/api_docs/python/tf/nn/embedding_lookup). "Mod" strategy may be useful when references to embedding slices are not evenly distributed over their indices.  However, "mod" strategy is a poor approximation to the actual load balancing users want. For example, when vocabulary is sorted by frequency, with "mod" strategy the first parameter server always has larger load than the second one. Furthermore, its current checkpointing mechanism also prevents users from migrating to a cluster with different number of parameter servers.


#### Current Support in Distribution Strategy


Sharded embeddings and non-sharded embeddings are both supported by `ParameterServerStrategy` with TF 1.x API. See this [PR](https://github.com/tensorflow/tensorflow/pull/23254). Sharded embeddings are not supported by other distribution strategies.

Non-sharded embeddings are currently supported by `MirroredStrategy` which mirrors them on all replicas. A [naive implementation of allgather](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/distribute/cross_device_utils.py#L648) exists in TensorFlow codebase to collect their gradients onto each replica. See [the transformer model](https://github.com/tensorflow/models/tree/master/official/transformer) for the usage of embeddings with `MirroredStrategy`.

However, embeddings are not supported by multi-worker `CollectiveAllReduceStrategy` unless we implement multi-worker allgather primitive in collective ops or always cast their updates to dense tensors and use allreduce to aggregate gradients.

Note: our design proposal below assumes there is an efficient implementation of allgather but any detail of allgather is out of this doc's scope. We'll have allgather implementation in collective ops soon.


### Partitioned Layer


A layer of a network can sometimes be partitioned. We call it partitioned layer in this document. This is another use case of `PartitionedVariable`.

The main reason people partition a layer is to balance parameter servers' load.

Some layers can be huge. For example the logits layer can be extremely large when there are large number of classes. People may want to shard their computations accordingly to avoid concatenating them, e.g. the implementation of `tf.nn.nce_loss`.

Even for mid-sized layer, people want to partition them to have more even distribution of parameter servers' load because people may get bad luck when variables are assigned to parameter servers in a round-robin fashion. See this [issue](https://github.com/tensorflow/tensorflow/issues/24953).

Partitioned layer is only supported by lower level APIs with or without `ParameterServerStrategy` in 1.x.


### Current Usage of PartitionedVariable


With TF 1.x API users can create a `PartitionedVariable` via `tf.get_variable` API:

```python
partitioned_var = tf.get_variable(
    "embedding",
    shape=(10000, 128),
    dtype=tf.float32,
    partitioner=tf.min_max_variable_partitioner(
        max_partitions=num_ps_replicas, axis=0))
```

The `partitioned_var` here is an instance of `PartitionedVariable`.

Alternatively users can use `tf.variable_scope` to apply `partitioner` to all variables created under its scope:

```python
with tf.variable_scope(partitioner=tf.min_max_variable_partitioner(
        max_partitions=num_ps_replicas, axis=0)):
  tf.keras.layers.Embedding(...)
```

Both these cases will also work when they are under a `ParameterServerStrategy`'s scope. However `MirroredStrategy` and `CollectiveAllReduceStrategy` will [error out](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/distribute/distribute_lib.py#L872) if a `PartitionedVariable` is about to be created.

These APIs `tf.get_variable` and `tf.variable_scope` will be gone in TF 2.0.


### Possible Problems with The `PartitionedVariable` Class 


Although we have a class called `PartitionedVariable`,  we often concatenate all of their components to pretend it is a single variable.

For example, `tf.matmul(partitioned_var, inputs)` would concatenate the list of variables in the `partitioned_var` before the execution of `matmul` as if the `partitioned_var` is a single variable. With the concatenation, the `partitioned_var` would only serve for loadbalancing purposes.

To avoid issue, it requires people to write methods that respect the partitions, e.g. `tf.nn.sampled_softmax`. This has some flavor of model parallelism.


## Goals


1.  Support partitioned embedding and partitioned layer in TF 2.0 in parameter server architecture via Distribution Strategy's and Keras layer's API. 
1.  Better support for embeddings in mirrored and collective allreduce architecture. We will not shard them in this architecture in our pre-TF 2.0 design.
2.  We will only support "div" partition strategy but as a post-TF 2.0 work we will support re-balancing embeddings in the Keras' Embedding layer. That means we will not support "mod" partition strategy any more.
3.  We also would like the API to enable people to easily switch between different architectures without needing to modify their models.


## Non-goals


1.  We don't have plan to support any flavor of model parallelism beyond the current implementations of embedding lookup and loss functions that respect partitioning.
2.  Support for partitioned variable in `ParameterServerStrategy` in eager mode is not a near-term goal since parameter server support in eager runtime won't be ready in the near future.
3.  We will not support any flavor of online learning or liquid sharding, e.g. dynamically placing or resizing partitions.


## Design Proposal


### Overview


Since Distribution Strategy is a library that is designed to parallelize and scale up models. We think that support for embeddings and partitioned layers should exist in Distribution Strategys. We'll support them differently for different Distribution Strategy:

|          |Embeddings|Partitioned Layer|
|:-------- |:-------- |:----------------|
|Mirrored/ CollectiveAllReduce Strategy|Pre-TF 2.0: <br>w/ Keras' Embedding layer API: won't be sharded but will be mirrored on each replica. <br><br>Post-TF 2.0: <br>w/ Distribution Strategy's API directly: allow sharding; being placed on host; or dense update. <br><br>Heuristics in Embedding layer to auto-set performance improvement hints.|Pre-TF 2.0: <br>w/ Keras' layer API: won't be sharded but will be mirrored on all replicas, just like a normal layer. <br><br>Post-TF 2.0: w/ Distribution Strategy's API directly: allow being mirrored on each host. <br><br>Heuristics in Embedding layer to auto-place larger embeddings on host.|
|ParameterServer Strategy|Pre-TF 2.0:<br>w/ Keras' Embedding layer API: probably sharded, according to `partitioner` passed to its constructor.<br><br>Post-TF 2.0:<br>w/ Distribution Strategy's API directly: allow creating partitioned variable with different partitioners.|Pre-TF 2.0:<br>w/ Keras' layer API: probably sharded, according to `partitioner` passed to its constructor which applies globally to all layers.<br><br>Post-TF 2.0:<br>w/ Distribution Strategy's API directly: allow creating partitioned variable with different partitioners.|

In the `ParameterServerStrategy`, the goal of supporting partitioned variable is to match the existing behavior in 1.x. We will propose its APIs for TF 2.0 below. 

**In the mirrored and collective allreduce architecture, we will only mirror embeddings or logits layers, either on devices or on host.** 

In the future, we will probably support a new strategy that is a hybrid of these two architectures in the future.

Please refer to the discussion in [Alternatives Considered](#Alternatives-Considered) section for why we've made these decisions.


### Proposed Change for pre-TF 2.0


The following is planned to be done before TF 2.0's release.


{'fraction_non_alphanumeric': 0.07461276917264828, 'fraction_numerical': 0.007933509633547412, 'mean_word_length': 4.011121628017038, 'pattern_counts': {'":': 0, '<': 90, '<?xml version=': 0, '>': 90, 'https://': 16, 'lorem ipsum': 0, 'www.': 6, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4035151', 'n_tokens_mistral': 5533, 'n_tokens_neox': 5130, 'n_words': 2630}