Parameter server

This is a summary of an e-mail between myself and @MLnick

### From Nick

Taking the simplest version of say logistic regression, the basic idea is to split up the parameter vector (beta) itself into chunks (so it could be a dask array potentially, or some other distributed dask datastructure). Training would then iterate over "mini-batches" using SGD (let's say each mini batch is a chunk of the X array). In each mini batch, the worker will "pull" the latest version of "beta" from the Parameter Server and compute the (local) gradient for the batch. The worker then sends this gradient to the PS, which then performs the update (i.e. update its part of "beta" using say the gradient update from the worker and the step size). The next iteration then proceeds in the same way. This can be sync or async (but typically is either fully async or "bounded stale" async).

The key is to do this effectively as direct communication from the worker doing the mini batch gradient computation, to the worker holding the parameters (the "parameter server"), without involving the master ("client" app) at all, and to only "pull" and "push" the part of beta required for local computation (due to sparsity this doesn't need to be the full beta in many cases). In situations where the data is very sparse (e.g. like the Criteo data) the communication is substantially reduced in this approach. And the model size can be scaled up significantly (e.g. for FMs the model size can be very large).

This is slightly different from the way say L-BFGS works currently (and the way I seem to understand ADMM works in dask-glm) - i.e. that more or less a set of local computations are performed on the distributed data on the workers, and the results collected back to the "master", where an update step is performed (using LBFGS or the averaging of ADMM, respectively). This is also the way Spark does things.

What I'm struggling with is quite how to achieve the PS approach in dask. It seems possible to do it in a few different ways, e.g. perhaps it's possible just using simple distributed dask arrays, or perhaps using "worker_client" and/or Channels. The issue I have is how to let each worker "pull" the latest view of "beta" in each iteration, and how to have each worker "push" its local gradient out to update the "beta" view, without the "master" being involved.

I'm looking into the async work in http://matthewrocklin.com/blog/work/2017/04/19/dask-glm-2 also to see if I can do something similar here.

### From me

First, there are two nodes that you might consider the "master", the scheduler and the client. This is somewhat of a deviation from Spark, where they are both in the same spot.

Second, what are your communication and computation requirements? A roundtrip from the client to scheduler to worker to scheduler to client takes around 10ms on a decent network. A worker-worker communication would be shorter, definitely, but may also involve more technology. We can do worker-to-worker direct, but I wanted to make sure that this was necessary.

Channels currently coordinate metadata through the scheduler. They work a bit like this:

1. Worker A subscribes to channel, tells scheduler
2. Worker B subscribes to channel, tells scheduler
3. Worker A creates some data and registers it on the channel, tells the scheduler
4. Scheduler tells all workers that are on this channel (A and B) that a new piece of data is on the channel
5. Worker B says great, I want this data, and asks the scheduler where it can get it
6. Scheduler tells Worker B that the data is on Worker A
7. Worker B gets data from Worker A

So there are a few network hops here, although each should be in the millisecond range (I think?).

We could also set up a proper parameter server structure with single hop communicatinos. Building these things isn't hard. As usual my goal is to extract from this experiment something slightly more general to see if we can hit a broader use case.

So I guess my questions become:

1. What are your communication requirements
2. How much data are you likely to shove through this
3. Are you likely to have multiple parameter servers? If so how would you anticipate sharding communication?

### From Nick

The PS idea is very simple at the high level. The "parameter server" can be thought of as a "distributed key-value store". There could be 1 or more PS nodes (the idea is precisely to allow scaling the size of model parameters across multiple nodes, such as in the case of factorization machines, neural networks etc).

A good reference paper is https://www.cs.cmu.edu/~muli/file/parameter_server_osdi14.pdf

So in theory, at the start of an iteration, a worker node asks the PS for only the parameters it needs to compute its update (in sparse data situations, this might only be a few % of the overall # features, per "partition" or "batch"). This can be thought of as a set of (key, value) pairs where the keys are vector indices and the values are vector values at the corresponding index, of the parameter vector. In practice, each PS node will hold a "slice" of the parameter vector (the paper uses a chord key layout for example), and will work with vectors rather than raw key-value pairs, for greater efficiency.

It seems like Channels might be a decent way to go about this. Yes, there is some network comm overhead but in practice for a large scale problem, the time to actually send the data (parameters and gradients say) would dominate the few ms of network hops. This cost could also be partly hidden through async operations.

The way I thought about it with Channels, which you touch on is:

1. Let's say we have 1x PS worker for simplicity, and some other "compute" workers. The "compute" workers will hold the chunks of data (X, y blocks). The PS will hold "beta".
2. PS creates an initial beta vector (random data, zeros, whatever). It could "publish" this vector (future?) on the Channel "params", saying "here is the latest version of beta".
3. Workers start iteration 1, and pull the new "beta" (future?) off the channel. Let's say Worker A perhaps needs 10% of the total vector - so it "pulls" beta[idx] from PS - where idx is the set of active feature indices it needs to compute it's gradient.
4. Worker A computes its partial gradient for the chunk. It needs to "push" this grad[idx] (or alternatively, a "sparse vector" version of grad) to the PS. It could push this as another future into the channel? Or perhaps another channel? But would the idea be that PS gets this future off the channel, knows that Worker A holds the data it needs, and does something like beta[idx] -= grad[idx] * step_size (simplified update), where it will know to pull grad[idx] from Worker A? And then "publishes" the new "beta future" on the "params" channel?
5. This all happens async - so effectively a "slow" worker may "miss" a few beta updates. Workers could always poll the head of the channel for the latest.
6. The PS could in this way implement some form of "bounded synchronous" updates.

To answer your specific questions:

1. As I mention above, of course we'd prefer to have lowest cost communication for the above scenario - but I would expect a few ms overhead from network hops to be marginal in terms of overall cost. I would tend to start with what is "built in" and see if it works well, before trying to build more custom stuff.
2. Quite a lot - that is the idea, to scale to large models. By large I would say typically 100s millions - billions of parameters in total. Each mini-batch would not typically communicate that entire parameter space, but it could still be a few million parameters per mini batch.
3. Yes - though even 1 PS can be useful in scaling. Sharding can range from simply splitting the param vector in contiguous chunks, to "key chord" layouts and other more involved architectures (mostly this is done for fault tolerance purposes).

### From me

So here is some code just to get things started off:

```python
def parameter_server():
beta = np.zeros(1000000)
with worker_client() as c:
betas = c.channel('betas', maxlen=1)
future_beta = c.scatter(beta)
betas.append(future_beta)

updates = c.channel('updates')
for update in updates:
beta = modify(beta, update)
future_beta = c.scatter(beta)
betas.append(future_beta)

def worker(idx, x):
with worker_client(separate_thread=False) as c:
betas = c.channel('betas', maxlen=1)
last_beta = betas.data[-1]
subset_beta = c.submit(operator.getitem, last_beta, idx).result()
params = subset_beta.result()

update = create_update(x, params)
updates = c.channel('updates')
updates.append(update)
updates.flush()
```

For what it's worth I expect this code to fail in some way. I think that channels will probably have to be slightly modified somehow. For example currently we're going to record all of the updates that have been sent to the updates channel. We need to have some way of stating that a reference is no longer needed. Channels need some mechanism to consume and destroy references to futures safely.