fix(dist): write back broadcast qparams to offload storage on non-src ranks - #3066
fix(dist): write back broadcast qparams to offload storage on non-src ranks#3066xesdiny wants to merge 2 commits into
Conversation
|
Important Review skippedAuto reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Team Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Merge Protections🔴 1 of 1 protections blocking · waiting on 👀 reviews
🔴 Require one maintainer reviewWaiting for any of
This rule is failing.All PRs must have at least one approving review from a maintainer before merging.
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request ensures that broadcasted quantization parameters are correctly written back to offload storage (CPUCache) on non-source ranks by calling update_offload_parameter after communication is complete. This resolves issues where the underlying CPU storage remains stale. Unit tests have been added to verify this behavior. The reviewer suggests guarding the broadcast logic with dist.is_initialized() to prevent potential crashes in non-distributed environments.
|
👋 Hi! Thank you for contributing to llm-compressor. Please add the ready label when the PR is ready for review. Note: This is required to complete the testing suite, please only add the label once the PR is code complete and local testing has been performed. |
|
The quality checks have failed. Please run |
|
The quality checks have failed. Please run |
999bd35 to
9506e3d
Compare
|
The quality checks have failed. Please run |
9506e3d to
3aca174
Compare
3aca174 to
7d63da1
Compare
|
@gemini-code-assist review |
There was a problem hiding this comment.
Code Review
This pull request introduces writeback support for CPU-offloaded modules during quantization parameter broadcasting by calling update_offload_parameter on non-source ranks, and adds corresponding unit tests. The review feedback points out a potential issue where as_broadcastable(param) might return a copy of the tensor rather than a view; if so, the broadcasted values would not be written back to the offload storage. It is recommended to save and write back the tensor returned by as_broadcastable(param) to ensure correctness.
7d63da1 to
7e5316a
Compare
|
@gemini-code-assist review |
There was a problem hiding this comment.
Code Review
This pull request updates broadcast_qparams_and_cleanup in src/llmcompressor/utils/dist.py to ensure that broadcasted quantization parameters are correctly written back to CPU-offloaded storage on non-source ranks using update_offload_parameter, and adds corresponding regression tests. The review feedback recommends appending the original param instead of broadcast_param to writeback_items to avoid potential shape or dtype mismatches, as the original parameter is already updated in-place by dist.broadcast.
7e5316a to
c179dc3
Compare
|
@gemini-code-assist review |
There was a problem hiding this comment.
Code Review
This pull request updates broadcast_qparams_and_cleanup in src/llmcompressor/utils/dist.py to ensure that broadcasted quantization parameters are correctly written back to offload storage (such as CPUCache) on non-source ranks using update_offload_parameter. It also adds a comprehensive suite of regression tests to verify this behavior. The review feedback suggests caching dist.is_initialized() outside the loop to improve efficiency and defensively copying data back to the original parameter tensor if as_broadcastable returns a copy or a view, ensuring correctness under all circumstances.
c179dc3 to
e78e258
Compare
|
@gemini-code-assist review |
There was a problem hiding this comment.
Code Review
This pull request updates the broadcast_qparams_and_cleanup function in src/llmcompressor/utils/dist.py to ensure that CPU-offloaded modules correctly persist broadcast results back to offload storage on non-source ranks using update_offload_parameter. It also introduces regression tests to validate this behavior. The review feedback suggests optimizing the writeback logic by checking data_ptr() to avoid redundant copy operations when tensors share the same underlying storage, and adding unit tests to cover the code path where an explicit copy is required.
e78e258 to
bca0a15
Compare
|
The quality checks have failed. Please run |
262f741 to
6066531
Compare
|
@gemini-code-assist review |
There was a problem hiding this comment.
Code Review
This pull request updates broadcast_qparams_and_cleanup to ensure that non-source ranks correctly write back broadcasted quantization parameters to offload storage (such as CPUCache) using update_offload_parameter, and adds comprehensive regression tests. The feedback suggests using reshape_as instead of view_as when copying the broadcasted parameter back to avoid potential RuntimeErrors if the tensor is non-contiguous.
… ranks broadcast_qparams_and_cleanup calls dist.broadcast on module attributes such as weight_scale, weight_zero_point, and weight_g_idx. When these tensors are stored in CPUCache (the common case with sequential pipeline auto_offload), getattr(module, name) returns a temporary onloaded GPU tensor — a new allocation each call. dist.broadcast modifies that temporary in-place, but the underlying CPU storage is never updated, leaving non-source ranks with stale or uninitialised quantization parameters. The quantization completes silently with no exception. Fix: collect (module, name, param, broadcast_param) tuples for non-src ranks before issuing async broadcasts. After _wait_for_comms, copy back if storage differs and call update_offload_parameter unconditionally to flush the broadcast result to CPUCache storage. Closes vllm-project#3119 Signed-off-by: xesdiny <xesdiny@gmail.com>
…cleanup writeback Three test files covering different levels of the bug: - test_broadcast_qparams_writeback.py: mock-based unit tests verifying the control-flow logic (update_offload_parameter called at the correct time for non-src ranks, not called for src rank, missing params handled, multi-module case). - test_cpucache_temporary_tensor.py (1 GPU, no dist): root-cause proof that CPUCache.__getattr__ returns a new temporary CUDA tensor on each call. Demonstrates that in-place modification without writeback is silently lost, and that update_offload_parameter correctly persists it. - test_broadcast_qparams_ddp_integration.py (2 GPU, NCCL): end-to-end integration test with real dist.broadcast and real CPUCache. Phase A confirms the upstream bug (non-src rank keeps stale weight_scale after broadcast). Phase B confirms the fix propagates correct values to all ranks. Modules are created before dist.init_process_group() to force CPUCache selection. Signed-off-by: xesdiny <xesdiny@gmail.com>
f6646c7 to
bfc9e1c
Compare
|
@gemini-code-assist review |
There was a problem hiding this comment.
Code Review
This pull request updates broadcast_qparams_and_cleanup in src/llmcompressor/utils/dist.py to explicitly write back broadcasted quantization parameters to CPU-offloaded storage on non-source ranks, preventing parameters from being lost when dist.broadcast modifies temporary onloaded tensors. It also adds comprehensive integration, regression, and root-cause tests. The review feedback suggests replacing torch.accelerator APIs (set_device_index and is_available) with their torch.cuda equivalents in the new tests to maintain backward compatibility with PyTorch versions older than 2.4.
| rank = int(os.environ["RANK"]) | ||
| local_rank = int(os.environ["LOCAL_RANK"]) | ||
| device = torch.device(f"cuda:{local_rank}") | ||
| torch.accelerator.set_device_index(local_rank) |
There was a problem hiding this comment.
Using torch.accelerator.set_device_index requires PyTorch 2.4 or newer. To maintain backward compatibility with older PyTorch versions (e.g., 2.1 to 2.3) which are commonly used in production environments, it is safer to use torch.cuda.set_device.
| torch.accelerator.set_device_index(local_rank) | |
| torch.cuda.set_device(local_rank) |
| This is the key property that makes dist.broadcast's in-place write | ||
| ineffective: it writes to a temporary that is immediately discarded. | ||
| """ | ||
| if not torch.accelerator.is_available(): |
There was a problem hiding this comment.
Using torch.accelerator.is_available requires PyTorch 2.4 or newer. To ensure backward compatibility with older PyTorch versions, use torch.cuda.is_available() instead.
| if not torch.accelerator.is_available(): | |
| if not torch.cuda.is_available(): |
|
|
||
| This is the bug: non-source ranks end up with stale weight_scale. | ||
| """ | ||
| if not torch.accelerator.is_available(): |
There was a problem hiding this comment.
Using torch.accelerator.is_available requires PyTorch 2.4 or newer. To ensure backward compatibility with older PyTorch versions, use torch.cuda.is_available() instead.
| if not torch.accelerator.is_available(): | |
| if not torch.cuda.is_available(): |
|
|
||
| After calling it, subsequent getattr returns the new value. | ||
| """ | ||
| if not torch.accelerator.is_available(): |
There was a problem hiding this comment.
Using torch.accelerator.is_available requires PyTorch 2.4 or newer. To ensure backward compatibility with older PyTorch versions, use torch.cuda.is_available() instead.
| if not torch.accelerator.is_available(): | |
| if not torch.cuda.is_available(): |
1 participant
Summary
broadcast_qparams_and_cleanupcallsdist.broadcaston module attributes such asweight_scale,weight_zero_point, andweight_g_idx. When these tensors are stored inCPUCache(the common case with sequential pipelineauto_offload),getattr(module, name)returns a temporary onloaded GPU tensor — a new allocation each call.dist.broadcastmodifies that temporary in-place, but the underlying CPU storage is never updated — leaving non-source ranks with stale or uninitialised quantization parameters.Trigger conditions: multi-GPU DDP GPTQ with sequential pipeline and
auto_offload=True, typically used to parallelize activation collection over large calibration sets. The quantization completes silently with no exception; the corruption only appears when inspecting saved checkpoints.This PR adds an explicit
update_offload_parameterwriteback for non-source ranks after_wait_for_comms, which correctly flushes the broadcast result back to CPUCache storage.Changes
src/llmcompressor/utils/dist.py: collect(module, name, param, broadcast_param)for non-src-rank entries before issuing async broadcasts; after comms complete, copy back and callupdate_offload_parameter(+40/-2 lines)tests/llmcompressor/utils/test_broadcast_qparams_writeback.py: unit tests covering non-src writeback, src no-writeback, missing params, and multi-module casestests/llmcompressor/utils/test_cpucache_temporary_tensor.py(new): 3 unit tests proving CPUCache returns a new temporary CUDA tensor pergetattrcall and that in-place modification is lost without writebacktests/llmcompressor/utils/test_broadcast_qparams_ddp_integration.py(new): 2-GPU integration test using realdist.broadcast+ real CPUCache, demonstrating the bug (Phase A) and fix (Phase B) in a single distributed sessionEmpirical validation
Controlled A/B on 8× L20,
Qwen3.5-122B-A10B, N=8192, seq=6144:Full details in issue #3119.
Fixes #3119