mrocklin
commented
6 years ago One general mechanism that might fit this need well is Pub/Sub. Under this construct any worker can join a topic as either a publisher or subscriber (or both). It learns from the scheduler who all of its peers are on that topic. When it publishes data on that topic it sends this data directly to all subscribers (or, in the future, to some subscribers that then forward it along to others). Under this model the parameter subscribes to updates and publishes parameters.
An implementation of a single parameter server might look like the following:
def parameter_server(...):
""" This runs on one special worker """
parameters = np.random.random(...)
updates_sub = Sub('updates') # subscribe to updates from workers
parameters_pub = Pub('parameters') # publish new parameters to workers
async def publish_parameters():
while not should_stop():
parameters_pub.put(parameters) # maybe send only to a few round-robin instead?
await gen.sleep(0.01) # maybe control this based on network traffic?
get_worker().loop.add_callback(publish_parameters)
for new_data in updates_sub:
update(parameters, new_data)
if should_stop():
return parameters
def worker(data: List[np.ndarray], # these correspond to chunks of data from a dask array
labels: List[np.ndarray]):
""" This runs on each worker and gets the arrays present on that worker """
data = np.concat(data)
labels = np.concat(labels)
updates_pub = Pub('updates')
parameters_sub = Sub('parameters')
async def update_parameters():
async for params in parameters_sub:
parameters = params
get_worker().loop.add_callback(update_parameters)
while not should_stop():
batch = data[np.random.randint(0, len(data))]
result = compute_update(parameters, batch)
updates_pub.put(result)
stsievert
/>
I'm restarting the discussion on parameter servers from https://github.com/dask/dask-glm/issues/57
When doing distributed training it is often useful to rapidly update and share parameters between distributed workers. The frequency with which these updates can be shared seems to strongly affect the rate of convergence. The common practice today seems to be to have dedicated nodes within the network that serve as parameter servers, accepting and aggregating updates from workers, and periodically publishing those back out to workers.
Dask's existing task scheduling model is not ideal for this. Every computation or communication checks in with the scheduler, which adds unnecessary network roundtrips and provides a central bottleneck for metadata tracking. To do parameter servers well we probably need to break out of this model and engage more peer-to-peer coordination, without frequent checking in with the scheduler.
Building this infrastructure is not hard, but if possible it would be good to build general machinery that can be used to solve both this problem, and potentially others. I would appreciate having a conversation with a few people to identify what programming and networking constructs might suffice for this.
cc @stsievert @fabianp @mlnick