2 NCCL all_reduces per step for large models on multiNode training

[puririshi98]

🐛 Bug

When using large models with DDP, we witness two NCCL all reduces per iteration and have no ability to increase the first bucket size. This is due to reducer.hpp KDefaultFirstBucketBytes being a const expr int that cannot be modified through the python api, yet it is exposed as an attribute of torch.distributed.
The value is defined here:

pytorch/torch/csrc/distributed/c10d/reducer.hpp

Line 26 in 43fb39c

constexpr int kDefaultFirstBucketBytes = int(1024 * 1024);

The bucket size is set here:
https://github.com/pytorch/pytorch/blob/master/torch/csrc/distributed/c10d/reducer.cpp#L1649

There is a Python API exposure, however since the value is a constexpr it is defined at compile time which would make changing it through Python at Runtime problematic. Changing the value through the Python does not currently work.
Python API:

pytorch/torch/csrc/distributed/c10d/init.cpp

Line 1623 in 43fb39c

module.attr("_DEFAULT_FIRST_BUCKET_BYTES") = ::c10d::kDefaultFirstBucketBytes;

The user should be able to increase this KDefaultFirstBucketBytes value through the python API.

Reproduction Steps (Should Work in any pytorch container from 21.05 and onwards):

1. tar -xvzf ddp_overlap.tgz
2. cp -r ddp_overlap /opt/
3. cd /opt/ddp_overlap
4. bash run.sh

Repro Scripts:
ddp_overlap.zip

Example of issue:

This issue persists if the KDefaultFirstBucketBytes value is edited through the Python binding showing that it does not work. In contrast, changing the hardcoded value of the constexpr defined in reducer.hpp manually and recompiling pytorch does remove the 2 all_reduce calls issue. This value should be properly defined in the backend of PyTorch such that it can be edited by a user in the PyTorch Python API without having to recompile for situations such as this with very large models.

cc @ezyang @gchanan @zou3519 @bdhirsh @jbschlosser @anjali411 @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @agolynski @SciPioneer @H-Huang @mrzzd @cbalioglu @gcramer23

[puririshi98]

@rohan-varma

[rohan-varma]

@puririshi98 Thanks for raising this issue! Can you confirm that if you edit dist._DEFAULT_FIRST_BUCKET_BYTES here: https://github.com/pytorch/pytorch/blob/master/torch/nn/parallel/distributed.py#L590-L594 the buckets are created with the right sizes at least for the first iteration? This changed python value should directly be input to the C++ compute_bucket_assignments so that should work IIUC.

Although, there is a larger issue where it looks like we always use the default first bucket size when rebuilding buckets, so the user-config value won't have any effect after we rebuild buckets the second iteration: https://github.com/pytorch/pytorch/blob/master/torch/csrc/distributed/c10d/reducer.cpp#L1649 cc @zhaojuanmao

[puririshi98]

@rohan-varma you are correct, adding that to the suggested ditributed.py init code did fix the issue for the first iteration only.

[zhaojuanmao]

first bucket size was not exposed for users to tune. So this is not a bug, but a good feature request. In fact, in our future auto tuning project, we would like to support tuning first bucket size and last bucket size.

[rohan-varma]

first bucket can now be tuned via dist._DEFAULT_FIRST_BUCKET_BYTES as of #62748