Use BFloat16 in distributed quantization when supported by NCCL #125113
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/125113
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 4694f99 with merge base 91a4740 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
pytorch-bot
bot
added
oncall: distributed
cyyever
force-pushed
the
quantization_bf16
branch
from
33ee903
to
4694f99
Compare
| @@ -69,15 +69,16 @@ at::Tensor _float_to_bfloat16_cuda(const at::Tensor& input) { | |||
|
|
|||
| auto output = at::empty( | |||
| {nrows, output_columns}, | |||
| input.options().dtype(at::kHalf)); // at::kHalf | |||
| #if HAS_NCCL_BF16_DATATYPE | |||
| input.options().dtype(at::kBFloat16)); | |||
There was a problem hiding this comment.
fyi i don't think you need to do this one in the preprocessor, you should be able to do it like:
input.options().dtype(HAS_NCCL_BF16_DATATYPE ? at::kBFloat16 : at::kHalf));
There was a problem hiding this comment.
HAS_NCCL_BF16_DATATYPE is a macro and I think it's better to format code like this so that it is easy to identify and remove the old branch in the future.
cpuhrsch
added
the
triaged
| reinterpret_cast<uint16_t*>(output.mutable_data_ptr<at::Half>()) | ||
| #endif | ||
| ); | ||
| C10_CUDA_KERNEL_LAUNCH_CHECK(); |
There was a problem hiding this comment.
what does the C10_CUDA_KERNEL_LAUNCH_CHECK function do? What's the purpose of uncommenting it?
|
@pytorchmergebot merge |
Merge startedYour change will be merged once all checks pass (ETA 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
…rch#125113) This PR enables BFloat16 in torch/csrc/distributed/c10d/quantization/quantization_gpu.cu . Pull Request resolved: pytorch#125113 Approved by: https://github.com/kwen2501
…rch#125113) This PR enables BFloat16 in torch/csrc/distributed/c10d/quantization/quantization_gpu.cu . Pull Request resolved: pytorch#125113 Approved by: https://github.com/kwen2501
This PR enables BFloat16 in torch/csrc/distributed/c10d/quantization/quantization_gpu.cu .
cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @penguinwu @fegin @XilunWu @wanchaol @fduwjj @wz337 @tianyu-l @wconstab @yf225 @chauhang @d4l3k