Repository navigation
Conversation
In a layerwise run of a model sharded over several GPUs, only the GPU holding the current layer worked on its expert updates while the others sat idle. The updates of a fused-MoE layer's experts are independent, so helpers whose update reads only the weight, its Hessian and the format's codec (the GGML formats) now run in one worker thread per GPU this process uses that has room for an update, each writing its result back to its weight. The first update runs in the caller's thread so lazy setup, such as loading a format's CUDA extension, happens once before the workers start. The results are bit-identical to running them in turn on one GPU. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
Contributor
|
No actionable comments were generated in the recent review.
 |
| Layer / File(s) | Summary |
|---|---|
Device-aware GPTQ updates modelopt/torch/quantization/utils/calib_utils.py, modelopt/torch/quantization/ggml/gptq.py |
GPTQHelper.update_weights accepts an optional target device and transfers its Hessian and working weight there. GGMLGPTQHelper supports updates on another device and moves packed weights to the selected weight’s device. |
GPU placement and expert scheduling modelopt/torch/quantization/utils/calib_utils.py |
Hessian placement uses available GPU memory and falls back to CPU under the stated conditions. update_expert_weights runs eligible updates on worker threads when enough GPUs and helpers qualify; other updates run on their weight devices. |
Calibration wiring and GPU test modelopt/torch/quantization/model_calib.py, tests/gpu/torch/quantization/test_gptq.py |
GPTQ calibration gathers expert handles and weights, then delegates the work to update_expert_weights. The two-GPU test checks the selected update device, expert weights, packed codes, and packed-weight placement against a one-GPU run. |
Priority:
Estimated code review effort: 3 (Moderate) | ~20 minutes
Change: Feature
Sequence Diagram(s)
sequenceDiagram
participant Calibration
participant update_expert_weights
participant ThreadPoolExecutor
participant GPTQHelper
participant CUDADevice
Calibration->>update_expert_weights: Submit collected expert work
update_expert_weights->>GPTQHelper: Run one eligible update before starting workers
update_expert_weights->>ThreadPoolExecutor: Queue remaining eligible updates
ThreadPoolExecutor->>GPTQHelper: Call update_weights with a target device
GPTQHelper->>CUDADevice: Transfer Hessian and working weight
ThreadPoolExecutor->>CUDADevice: Synchronize worker device
update_expert_weights->>CUDADevice: Synchronize source weight devices
Possibly related PRs
- NVIDIA/Model-Optimizer#2864: Changes the same fused-MoE GPTQ path and adds GPU-aware Hessian placement and expert-update handling.
- NVIDIA/Model-Optimizer#2945: Changes GPTQ Hessian placement and GPU headroom for layerwise calibration, which are also used in this change.
Suggested reviewers: sugunav14
Merge Risk: 
37c52
-
No merge-blocking issue was identified. The change is ready to merge after normal checks.
Pre-merge checks: 
-
Check name Status Explanation Description Check Check skipped - CodeRabbit’s high-level summary is enabled. Title check The title clearly and concisely describes the main change: parallel fused-MoE expert GPTQ updates across GPUs available to the process. Docstring Coverage Docstring coverage is 80.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 15 functions across 4 files. Linked Issues check Check skipped because no linked issues were found for this pull request. Out of Scope Changes check Check skipped because no linked issues were found for this pull request. Security Anti-Patterns No listed security anti-pattern was introduced. The pull request changes only four Python files, and added-line inspection found no unsafe torch.load/numpy.load, trust_remote_code=True, eval/exec, or … Full details: Security Anti-Patterns
-
Explanation
No listed security anti-pattern was introduced. The pull request changes only four Python files, and added-line inspection found no unsafe torch.load/numpy.load, trust_remote_code=True, eval/exec, or # nosec usage. No pyproject.toml or requirements.txt changes were present.
-
Finishing Touches
-
📝 Generate docstrings
-
- Commit to this branch
- Create a new PR
🧪 Generate unit tests (beta)
-
- Commit to this branch
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Comment @coderabbitai help to get the list of available commands.
Contributor
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #3002 +/- ##
==========================================
- Coverage 79.85% 79.12% -0.73%
==========================================
Files 661 661
Lines 74935 74973 +38
==========================================
- Hits 59836 59325 -511
- Misses 15099 15648 +549
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Type of change: new feature (performance)
In a layerwise GPTQ run of a model that
device_mapspreads over several GPUs, only the GPU holding the current layer did that layer's work. The fused-MoE expert updates, a large share of each MoE layer's time, ran one after another on that GPU while the others sat idle.The updates of one layer's experts are independent. For helpers whose update reads only the weight, its Hessian and the format's codec, the expert updates now run in one worker thread per GPU this process already uses that has room for an update. Each worker copies an expert's weight and Hessian to its GPU, runs the update there and writes the result back into the expert's weight. That covers the GGML formats (
GGMLGPTQHelper.UPDATES_ON_ANY_DEVICE); helpers that read quantizer state from the weight's device, such as NVFP4's, keep running in turn.The results are bit-identical to the sequential updates.
Independent of #2945; both touch
calib_utils.py, so whichever merges second needs a small rebase.Usage
No API change.
Testing
Real models on 4x GB300 (oci-jhb). Layerwise GPTQ on the routed experts at IQ2_XXS, 1,024 x 512 calibration tokens, KL to BF16 over wikitext-2 (100 x 512 tokens, second half scored). Measured on top of #2945, which sets the batch size and Hessian placement these runs used:
Qwen's KL is identical token by token. GLM's moves within run-to-run noise: its Hessian placement varies between runs with free memory (for example 222 vs 217 experts' Hessians on one GPU), which changes how they are summed.
test_expert_updates_on_another_gpu_match_one_gpu(2 GPUs): expert updates spread over two GPUs give bit-identical weights to running them in turn on one.tests/gpu/torch/quantization/{test_gptq,test_layerwise_calibrate}.py25 passed on 2x RTX PRO 6000 (CI's GPU runner has one GPU and skips the 2-GPU test).tests/unit/torch/quantization/{test_gptq,test_ggml_gptq,test_layerwise_calibrate}.py80 passed.Before your PR is "Ready for review"
CONTRIBUTING.md: N/A🤖 Generated with Claude Code
Summary by CodeRabbit