Support different AMP & buffer configurations in one experiment, fix minor bugs

[justheuristic]

New features:

• CollaborativeOptimizer can now combine fp16=True and reuse_grad_buffers=True with a special scaler
• CollaborativeOptimizer peers with reuse_grad_buffers=True and reuse_grad_buffers=False can now co-exist
• CollaborativeOptimizer peers with and without AMP can now co-exist

The new behavior of CollaborativeOptimizer with fp16 is:

• grad_scaler=None: regular fp32 behavior
• reuse_grad_buffers=False with GradScaler: works as usual, independently un-scales each tensor before accumulation, does not affect internal optimizer
• reuse_grad_buffers=True with GradScaler: when calling scaler.step(opt), it will raise error and complain that it requires HivemindGradScaler
• reuse_grad_buffers=False with HivemindGradScaler: applies unscale/update only around global optimizer step

[justheuristic]

This previously caused a bug where reuse=True peers would not be scaled by scale_by. As a result, if there was a mix of reuse=True and reuse=False peers, reuse=True would have larger gradients and dominate the reuse=False peers.

[justheuristic]

This wold actually an error with reuse_grad_buffers=True, but it worked because noone asked for more than len(grad_buffers) elements

[codecov]

Codecov Report

Merging #389 (773a3bb) into master (4a9bc92) will decrease coverage by 0.74%.
The diff coverage is 23.65%.

@@            Coverage Diff             @@
##           master     #389      +/-   ##
==========================================
- Coverage   84.11%   83.37%   -0.75%     
==========================================
  Files          70       71       +1     
  Lines        6423     6497      +74     
==========================================
+ Hits         5403     5417      +14     
- Misses       1020     1080      +60     

Impacted Files	Coverage Δ
hivemind/optim/collaborative.py	23.77% <7.89%> (-1.32%)	⬇️
hivemind/optim/grad_scaler.py	33.33% <33.33%> (ø)
hivemind/optim/__init__.py	100.00% <100.00%> (ø)
hivemind/dht/node.py	91.44% <0.00%> (-1.19%)	⬇️
hivemind/averaging/matchmaking.py	83.75% <0.00%> (-0.32%)	⬇️

[borzunov]

Why does it ignore the allow_fp16 value always setting it to True?

[justheuristic]

that's the same trick that fairscale uses to allow training without master fp32 weights
https://github.com/facebookresearch/fairscale/blob/main/fairscale/optim/grad_scaler.py
(got referred there by @TimDettmers)

Added a quick comment explaining that

[borzunov]

nit (suggestion): Since the if branch ends with return, we can remove else: and one indent level for the remaining code. This may make it look simpler.

[justheuristic]

I respectfully disagree here: it does reduce indentation, but:

• the code in the else clause is quite simple as it is
• crucially, last time the code was wrong because i've accidentally removed return and didn't think about the ramifications

If you do insist, please state that explicitly, i'll remove the else clause anyway

[borzunov]

I would say that else: after return is obviously redundant and seems like a bug, so I insist on choosing one of the two options:

1. return without else: (preferred since it implies less indentation, more clarity)
2. else: without return