Add CollaborativeOptimizer

[leshanbog]

This PR introduces necessary objects for collaborative training so that client code for participation is clean and concise (like default train loop).

• WeightedAverager - DecentralizedAverager that averages trainable params or gradients with peer-wise weights
• CollaborativeOptimizer - performs model updates after collaboratively accumulating a target (large) batch size across peers

[justheuristic]

[example usage]

import time, socket
import torch, torch.nn as nn
import hivemind

with socket.socket() as sock:
    coordinator_exists = sock.connect_ex(("127.0.0.1", 1337)) == 0
if not coordinator_exists:
    dht = hivemind.DHT(listen_on='127.0.0.1:1337', start=True),
    
model = nn.Sequential(nn.Linear(32, 16), nn.ReLU(), nn.Linear(16, 32))

opt = hivemind.CollaborativeOptimizer(
    opt=torch.optim.Adam(model.parameters()),
    dht=hivemind.DHT(initial_peers=['127.0.0.1:1337'], start=True),
    prefix='test_exp', target_group_size=2,
    target_batch_size=32, batch_size_per_step=1, verbose=True,
    start=True
)


while True:
    x = torch.randn(10, 32)
    time.sleep(1)
    loss = torch.mean((x - model(x)) ** 2)
    loss.backward()
    opt.step()

[mryab]

The rationale for all logger.debug replacements is that we don't want to pollute the output of training scripts unless asked by the user. An alternative solution is to control the verbosity by an argument

[nevec]

Maybe we should accept averager class in constructor or give a possibility to override averager in subclass?
I don't see an easy way to change averaging logic (e.g. average opt stats) without copypasting the whole class. Correct me if I'm wrong

[justheuristic]

Good point, that might prove useful in subsequent research for secure/robust averagers
/* Added averager_cls parameter that allows user to override the averager class */

[justheuristic]

Conversation log with @nevec : we have instead opted for a _make_averager method for more versatility