Amd add quickreduce bak - #8139
lihaoyang-amd wants to merge 3 commits into
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Summary of Changes
Hello @lihaoyang-amd, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!
This pull request introduces a new "Quick Allreduce" communication primitive specifically optimized for AMD ROCm MI300 series GPUs. This new all-reduce implementation leverages quantization (FP, INT8, INT6, INT4) to further accelerate collective communication, complementing the existing custom all-reduce. It integrates into the distributed communication framework, allowing for automatic fallback to other methods if conditions for quick all-reduce are not met.
Highlights
- New Quick Allreduce Implementation: Introduced a new, highly optimized all-reduce primitive (
QuickAllReduce) specifically designed for AMD ROCm MI300 series GPUs. This implementation supports various quantization levels (FP, INT8, INT6, INT4) for improved performance. - Quantization Support: The Quick Allreduce supports different quantization regimes, allowing for trade-offs between precision and communication speed, with configurable environment variables (
ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16,ROCM_QUICK_REDUCE_QUANTIZATION). - Dynamic Allreduce Selection: The distributed communication framework now intelligently selects the most efficient all-reduce method (Quick Allreduce, Custom Allreduce, PyMSCCLPP, or PyTorch NCCL) based on hardware capabilities, tensor size, and configuration, with automatic fallback.
- Centralized Device Utilities: Common utility functions for device introspection (like
is_full_nvlinkandis_weak_contiguous) have been refactored and moved to a dedicatedcustom_all_reduce_utils.pyfile for better organization. - Comprehensive Testing: A new test suite (
test_quick_allreduce.py) has been added to validate the functionality of the Quick Allreduce in both eager and graph modes across different world sizes and data types.
Using Gemini Code Assist
The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.
Invoking Gemini
You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| Feature | Command | Description |
|---|---|---|
| Code Review | /gemini review |
Performs a code review for the current pull request in its current state. |
| Pull Request Summary | /gemini summary |
Provides a summary of the current pull request in its current state. |
| Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in issue comments and review comments. |
| Help | /gemini help |
Displays a list of available commands. |
Customization
To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.
Limitations & Feedback
Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.
You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.
Footnotes
-
Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩
There was a problem hiding this comment.
Code Review
This pull request introduces a "quick all-reduce" implementation for AMD ROCm GPUs, which leverages quantization to accelerate the all-reduce operation. The changes are comprehensive, touching Python application logic, C++/HIP kernels, build files, and adding new tests. The overall structure and implementation are well-done. The refactoring of utility functions into custom_all_reduce_utils.py improves code organization. The new QuickAllReduce class is a significant addition and is integrated into the existing parallel state management. I've identified a few critical issues in the all-reduce dispatch logic and the C++ kernel implementation that need to be addressed. There are also some medium-severity issues related to environment variable handling in Python and a potential bug in the C++ buffer allocation. Once these issues are resolved, this will be a solid contribution to improving performance on ROCm platforms.
| return torch.ops.sglang.outplace_all_reduce( | ||
| input_, group_name=self.unique_name | ||
| outplace_all_reduce_method = "pymscclpp" | ||
| outplace_all_reduce_method = "ca" |
There was a problem hiding this comment.
This line unconditionally overwrites outplace_all_reduce_method to "ca". This negates the preceding logic that selects between qr, ca, and pymscclpp. As a result, quick_all_reduce will never be used. This seems to be a leftover from debugging and should be removed.
| torch.ops.sglang.outplace_all_reduce( | ||
| input_, | ||
| group_name=self.unique_name, | ||
| outplace_all_reduce_method=outplace_all_reduce_method, | ||
| ) |
There was a problem hiding this comment.
torch.ops.sglang.outplace_all_reduce is an out-of-place operation that returns a new tensor with the result. The current implementation calls the op but discards its return value. You need to return the result of this operation.
| torch.ops.sglang.outplace_all_reduce( | |
| input_, | |
| group_name=self.unique_name, | |
| outplace_all_reduce_method=outplace_all_reduce_method, | |
| ) | |
| return torch.ops.sglang.outplace_all_reduce( | |
| input_, | |
| group_name=self.unique_name, | |
| outplace_all_reduce_method=outplace_all_reduce_method, | |
| ) |
| } else { | ||
| int32_t int16_2 = (qw >> (i * 4)) & kMask000F; | ||
| int16_t low = static_cast<int16_t>(int16_2 & 0xFFFF); | ||
| int16_t high = static_cast<int16_t>((int16_2 >> 16) & 0xFFFF); | ||
| nv_bfloat16 bf_low = __float2bfloat16(static_cast<float>(low)); | ||
| nv_bfloat16 bf_high = __float2bfloat16(static_cast<float>(high)); | ||
| nv_bfloat162 bf2 = __halves2bfloat162(bf_low, bf_high); | ||
| int32_t packed_bf16 = *reinterpret_cast<int32_t*>(&bf2); | ||
| w[i] = packed_add<nv_bfloat16>(packed_bf16, kRangeMin); | ||
| } |
There was a problem hiding this comment.
The dequantization logic for bfloat16 in CodecQ4::recv appears to be incorrect. The send function adds kRangeBias to the quantized values. The recv function should perform the inverse operation, which is subtracting this bias. However, the current implementation for bfloat16 adds kRangeMin instead. This will lead to incorrect results for bfloat16 tensors.
A similar issue seems to exist in CodecQ6::recv (lines 322-331) and CodecQ8::recv (lines 467-475) for bfloat16 dequantization.
| } | ||
| // Allocate buffer size for worst case: F16 2-stage buffer. | ||
| uint32_t flags_buffer_size = 2 * world_size * kMaxNumBlocks * sizeof(uint32_t); | ||
| static int64_t data_buffer_size = 2 * this->kMaxProblemSize; |
There was a problem hiding this comment.
The data_buffer_size is declared as a static local variable inside DeviceComms::init. This means it will be initialized only once, using the kMaxProblemSize of the first DeviceComms instance created. If subsequent instances are created with a different max_problem_size, they will still use the buffer size calculated from the first instance, which could lead to buffer overflows or inefficient memory usage. The static keyword should be removed to make data_buffer_size a regular local variable.
int64_t data_buffer_size = 2 * this->kMaxProblemSize;| # On RocM, bfloat16 kernels are slower than fp16 | ||
| # due to slower match operations | ||
| # If environment variable is set to 1, we convert input to fp16 | ||
| self.use_fp16_kernels = os.environ.get("ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16", 1) |
There was a problem hiding this comment.
os.environ.get() returns a string. self.use_fp16_kernels is initialized with the string "1" by default, but it's used as a boolean/integer in should_quick_allreduce and passed to qr_all_reduce. While a non-empty string evaluates to True, it's safer and clearer to explicitly cast it to the expected type.
| self.use_fp16_kernels = os.environ.get("ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16", 1) | |
| self.use_fp16_kernels = int(os.environ.get("ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16", 1)) |
| qr_max_size = os.environ.get("ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB", None) | ||
| if qr_max_size is not None: | ||
| if qr_max_size < 1: | ||
| logger.info( | ||
| "You should not set a max_size smaller than 1MB, which can " | ||
| "lead to error or degradation to custom allreduce or rccl." | ||
| ) | ||
| qr_max_size = qr_max_size * MB |
There was a problem hiding this comment.
os.environ.get() returns a string, but qr_max_size is used in numerical operations on lines 210 and 215. This will cause a TypeError at runtime. You should convert the value from the environment variable to an integer before performing calculations.
qr_max_size = os.environ.get("ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB", None)
if qr_max_size is not None:
qr_max_size = int(qr_max_size) * MB| p.join() | ||
|
|
||
|
|
||
| class TestCustomAllReduce: |
There was a problem hiding this comment.
The test class is named TestCustomAllReduce, but it appears to be testing the new QuickAllReduce implementation. For clarity and consistency, consider renaming it to TestQuickAllReduce.
| class TestCustomAllReduce: | |
| class TestQuickAllReduce: |
|
Thanks @lihaoyang-amd. Closing this because it is still a draft and has not been updated in 255 days. Reopen it if the work is still relevant. Some directories moved recently, so an older branch may need retargeting: |
2 participants
test