Metrics placement between CPU and GPU and syncs

[maximsch2]

Question/discussion on the boundary of metrics and lightning itself: by default, we suggest people to store metrics as a field on the LightningModule, which means the metric is moved to whatever device the model is being trained/executed on. This is great as it makes the use of the metric easy, but how should we approach it if we want to keep it on CPU. Motivation for this is trying to minimize GPU RAM usage to leave it for the model itself, and also to allow us to run more metrics and more expensive metrics. This is relatively non-trivial now: we need to move metric to CPU, setup separate distributed GLOO based group (to exchange metric states), and manually handle moves of data back and forth. Any thoughts on better ways to solve this?

[justusschock]

There were thoughts of using pytorch's TCP-Store for this (by wrapping that one, we probably could have one storage for all the states across metrics and it could be moved to cpu upon receiving and moved to GPU upon requesting it (maybe).

Thoughts @maximsch2 @SkafteNicki ?

[maximsch2]

The API seems pretty limiting. It can be fine for simpler metrics, but the fact that you need to increment counters one by one will likely be a bottleneck for bigger ones. precision/recall at threshold metrics will have 4*num_classes states, for metrics that need entire curve and use binning we'll need to multiply it further by num_thresholds

[justusschock]

@maximsch2 Yes, it can be limiting for some points, but tbh. I don't think it will.

When thinking of something like

 store = dist.TCPStore("127.0.0.1", 0, 1, True)
store.add("first_state", [3, 5]) # just some placeholders
store.add("first_state", [4,5])

This should be fine. And we could use it with lists like here or also with tensors/integers/floats.
The benefit would be, that by sending things there we can immediately move them to cpu and also there is no additional sync necessary between the processes in ddp (especially. no GPU sync) since all processes just need to gather the states from that store.

Where do you see limits?

[maximsch2]

I haven't used this before myself, but torch.distributed.Store.add docs seems to suggest that you can only do a single int.