Repository: pytorch
Issue: Two independent future chains in PowerSGD cannot be kicked off asynchronously
Labels: oncall: distributed,triaged

## 🐛 Bug

PowerSGD communication hook handles the rank-1 tensors and high-rank tensors separately and hence creates two future chains that operate two non-overlapping subsets of the input gradient tensor. These two independent callbacks should be able to run asynchronously. However, currently they have to be kicked off sequentially.

## To Reproduce

Steps to reproduce the behavior:

1. Move `allreduce_contiguous_rank1_tensors_fut` ahead of "Step II" in torch/distributed/algorithms/ddp_comm_hooks/powerSGD_hook.py
2. Run a distributed training with PowerSGD turned on.

Facebook:
Run `scripts/wayi/mast_prof_gradient_compression.sh POWER_SGD 8` to launch a training job

The error message seems to be relevant to cache allocator.

```
[trainer0]:*** Aborted at 1608595708 (Unix time, try 'date -d @1608595708') ***
[trainer0]:*** Signal 15 (SIGTERM) (0x7fff5389aeb0) received by PID 1151 (pthread TID 0x7fb5b652b180) (linux TID 1151) (maybe from PID 1401532080, UID 32767) (code: -1241465184), stack trace: ***
[trainer0]:    @ 0000000019b5092e folly::symbolizer::(anonymous namespace)::signalHandler(int, siginfo_t*, void*)
[trainer0]:    @ 0000000000000000 (unknown)
[trainer0]:    @ 0000000018b23c77 at::Tensor::record_stream(c10::Stream) const
[trainer0]:    @ 000000000945f6e9 torch::autograd::THPVariable_record_stream(_object*, _object*, _object*)
[trainer0]:    @ 000000000011c9c3 _PyMethodDef_RawFastCallKeywords
[trainer0]:    @ 0000000000134751 _PyMethodDescr_FastCallKeywords
[trainer0]:    @ 0000000000174aee _PyEval_EvalFrameDefault
[trainer0]:    @ 000000000012c84d _PyFunction_FastCallKeywords
[trainer0]:    @ 000000000016f255 _PyEval_EvalFrameDefault
[trainer0]:    @ 000000000012c84d _PyFunction_FastCallKeywords
[trainer0]:    @ 000000000016f255 _PyEval_EvalFrameDefault
[trainer0]:    @ 000000000012c84d _PyFunction_FastCallKeywords
[trainer0]:    @ 000000000016f255 _PyEval_EvalFrameDefault
[trainer0]:    @ 000000000012c84d _PyFunction_FastCallKeywords
[trainer0]:    @ 000000000016f255 _PyEval_EvalFrameDefault
[trainer0]:    @ 000000000012c84d _PyFunction_FastCallKeywords
[trainer0]:    @ 0000000000174855 _PyEval_EvalFrameDefault
[trainer0]:    @ 000000000012cf49 _PyFunction_FastCallDict
[trainer0]:    @ 00000000000d8591 _PyObject_Call_Prepend
[trainer0]:    @ 0000000000125ef0 slot_tp_call
[trainer0]:    @ 0000000000134e2b _PyObject_FastCallKeywords
[trainer0]:    @ 0000000000174a90 _PyEval_EvalFrameDefault
[trainer0]:    @ 000000000012b602 _PyEval_EvalCodeWithName
[trainer0]:    @ 000000000012d109 _PyFunction_FastCallDict
[trainer0]:    @ 0000000000171644 _PyEval_EvalFrameDefault
[trainer0]:    @ 000000000012c84d _PyFunction_FastCallKeywords
[trainer0]:    @ 000000000016f255 _PyEval_EvalFrameDefault
[trainer0]:    @ 000000000012c84d _PyFunction_FastCallKeywords
[trainer0]:    @ 000000000016f255 _PyEval_EvalFrameDefault
[trainer0]:    @ 000000000012c84d _PyFunction_FastCallKeywords
[trainer0]:    @ 000000000016f07e _PyEval_EvalFrameDefault
[trainer0]:    @ 000000000012b602 _PyEval_EvalCodeWithName
[trainer0]:    @ 000000000012ca20 _PyFunction_FastCallKeywords
[trainer0]:    @ 0000000000170228 _PyEval_EvalFrameDefault
[trainer0]:    @ 000000000012b602 _PyEval_EvalCodeWithName
[trainer0]:    @ 000000000012c6db PyEval_EvalCodeEx
[trainer0]:    @ 000000000012c72e PyEval_EvalCode
[trainer0]:    @ 00000000002115f4 run_mod
[trainer0]:    @ 0000000000215974 PyRun_StringFlags
[trainer0]:    @ 00000000002178cf PyRun_SimpleStringFlags
[trainer0]:    @ 000000000021aa15 pymain_main
[trainer0]:    @ 000000000021b04f _Py_UnixMain
[trainer0]:    @ 00000000000211a5 __libc_start_main
[trainer0]:    @ 0000000000400689 _start
```


cc @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @agolynski @SciPioneer @H-Huang @mrzzd