DDP validation: All gather for flattened 1D tensors taking long time to complete

[microsoft]

DDP validation: All gather for flattened 1D tensors taking long time to complete

[krishansubudhi]

I am working on the fix will raise a PR soon.

[krishansubudhi]

From more debugging it seems like a bug in gather_tensors_on_cpu

def gather_tensors_on_cpu(self, x: torch.tensor):
        n_samples = len(x)
        self._set_gather_frequency(n_samples)
        gathered = []
        n_chunks = n_samples // self.gather_frequency + 1
        print(n_chunks, n_samples, self.gather_frequency) # Debug code introduced

Output

1510 3018 2

This bug is because of self._set_gather_frequency(n_samples).
In case of multiple outputs, if dimension 0 of first output was 2, then gather_frequency will be set as 2 for rest of the outputs. Class variable assignment needs to be avoided here.

[krishansubudhi]

Also if n_chunks is different in 2 processes, all-gather gets stuck as the process with higher number of chunks keeps waiting. For this, either make num_chunks = 1 or gather the num_chunks tensor first and take the maximum.

[jsleep]

Thanks for debugging this @krishansubudhi - any chance you'd be able to put these changes in or should we try to resource for next sprint?