Add DistributedHistory for multi-gpu training

[BenjaminBossan]

Description

(this is more or less copied from the docs)

When training a net in a distributed setting, e.g. when using torch.nn.parallel.DistributedDataParallel, directly or indirectly with the help of AccelerateMixin, the default history class should not be used. This is because each process will have its own history instance with no syncing happening between processes. Therefore, the information in the histories can diverge. When steering the training process through the histories, the resulting differences can cause trouble. When using early stopping, for instance, one process could receive the signal to stop but not the other.

DistributedHistory will take care of syncing the distributed batch information across processes, which will prevent the issue just described.

This class needs to be initialized with a distributed store provided by PyTorch (https://pytorch.org/docs/stable/distributed.html#distributed-key-value-store). I have only tested torch.distributed.TCPStore so far, but torch.distributed.FileStore should also work. The DistributedHistory also needs to be initialized with its rank and the world size (number of processes) so that it has all the required information to perform the syncing. When using accelerate, that information can be retrieved from the Accelerator instance.

Comments

1. Even though the batch information, which is split across processes, is synced, the epoch information, which is not split, is not synced. E.g. the recorded duration can be different between processes. It is not quite clear what the "correct" behavior should be here, it would probably depend on what is done based on this information.
2. To make it possible to use the new class, I had to change the net initialization code to not create a new history instance (except when it is None) -- instead the history is just cleared. Otherwise, calling fit would always overwrite the DistributedHistory with a normal History object.
3. Unfortunately, TCPStore cannot be pickled. Therefore, I set it to None when pickling. This is not tragic as long as users pickle the final model and only load it for inference. If they want to keep on training, they would need to set the net.history.store manually.
4. To store values in the kv store as strings, I json-serialize the values, which means that certain types are not supported (say, numpy arrays). I guess the values could be pickled, but that seems like overkill.
5. I had some bugs when testing this with certain Python versions with pytest ("Fatal Python error: Aborted"). This did not happen with other Python versions or when not using pytest. Fingers crossed that this works in CI. I also had to add some time.sleep calls in the multiprocessing tests to avoid "broken pipe" etc. Update: CI had segfaults for PyTorch 1.11, so I'm skipping tests with DistributedHistory for that PyTorch version. Maybe the tests are just flaky, but we don't want that.

[BenjaminBossan]

For the history of this feature, see this comment and associated thread. Ping @Raphaaal.

[githubnemo]

Looks clean and well thought out. LGTM

[githubnemo]

it would be beneficial to know why this is expected

[githubnemo]

As a single description of the parameter this is quite tautological. Maybe it is worthwhile to remark the properties and the possible origin of the rank? E.g. "Number differentiating the distributed training processes, e.g. as provided by accelerate.local_process_index."?

[BenjaminBossan]

@githubnemo Very good points, I have extended the docstrings and added more context to the comment in __getstate__. I also added a minor change there, namely also catching PicklingError. The PyTorch stores will raise TypeError, but maybe some users want to provide custom stores, so it's cautious to also catch PicklingError.

Oh and btw. coverage claims that these lines are not covered, which is not true (adding an error there will raise). I suspect that coverage is not correctly measured because of the use of multiprocessing.

[BenjaminBossan]

Failing CI seems to be unrelated.