[PyTorch][torch.compile] Remove process group from quantizers#3104
Open
pggPL wants to merge 13 commits into
Open
[PyTorch][torch.compile] Remove process group from quantizers#3104pggPL wants to merge 13 commits into
pggPL wants to merge 13 commits into
Conversation
…ing on quantizer Remove the persistent amax reduction process group from quantizer objects (Float8CurrentScalingQuantizer, NVFP4Quantizer) and thread it as a call-time argument through quantize/quantize_impl/update_quantized down to the fused C++ kernels (tex.quantize, layernorm_fwd, rmsnorm_fwd). The C++ constructors no longer read with_amax_reduction/amax_reduction_group from Python; the group is attached transiently to the C++ quantizer per call. All call sites (linear, layernorm_linear, layernorm_mlp, ops/basic_linear, base.grad_output_preprocess, distributed gather, attention context-parallel, FSDP) now compute and forward the reduction group explicitly. Backward compatibility: with_amax_reduction/amax_reduction_group are still accepted (constructor or attribute) and honored as a fallback when no per-call group is given, with a DeprecationWarning, preserving 1:1 numerics for legacy callers. Also drop the per-quantizer RHT matrix tensor in favor of a property backed by the process-global cache. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
…uantizer Clean up explanatory comments and docstrings that narrated the migration of the amax reduction process group off the quantizer (e.g. "passed per call instead of being stored on the quantizer", "no longer stored on the Python quantizer"). No functional changes. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
… state Thread the amax reduction process group as a per-call argument to the C++ Quantizer::quantize / quantize_impl / quantize_with_amax methods and drop the with_amax_reduction / amax_reduction_group members from the C++ quantizer objects. The group is converted once via convert_amax_reduction_group and forwarded to the kernels, so the quantizer carries no distributed state. Also remove an orphaned comment in attention backends that described amax reduction group handling happening elsewhere. No public API change: these are internal transformer_engine::pytorch symbols; the pybind-exposed quantize already takes amax_reduction_group. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
The grouped NVFP4 quantize-with-amax path never receives a reduction group (the FC2 input is a local per-token activation, not an all-gathered shard), so allreduce_nvfp4_amax_tensors was always a no-op. Drop the helper and the amax-tensor list it consumed. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Drop dangling comments at the end of the per-module _customize_quantizers if/else blocks that only narrated that amax reduction is "supplied per quantize call" (and a "no amax reduction here" aside). They annotated no code and were migration artifacts; the if/else twins missed in the first cleanup pass. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
…s.partial Pass amax_reduction_group as an explicit argument to the _QuantizeFunc autograd function (and forward it to quantize_impl) rather than binding it via a functools.partial on every quantize() call. Avoids a per-call allocation on the hot path and is friendlier to torch.compile. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Drop the "(not on the quantizer, which must stay free of distributed state)" aside in fsdp_pre_all_gather, matching the earlier trim of the same phrase in the _stash_fsdp_amax_group docstring. The remaining comment still explains why the reduction group is stashed on the tensor. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
for more information, see https://pre-commit.ci
Collaborator
Author
|
/te-ci pytorch L1 |
Contributor
Greptile SummaryThis refactoring moves the amax reduction process group out of quantizer state and passes it per quantization call instead, eliminating mutable distributed state on quantizers and making
Confidence Score: 5/5Safe to merge. The refactoring is mechanically consistent across all 25 files: the per-call group threading correctly replaces stored-group mutations, deprecated fallbacks are preserved, and the All production quantization paths correctly thread transformer_engine/pytorch/custom_recipes/quantization_ref_current_scaling.py Important Files Changed
Reviews (4): Last reviewed commit: "Carry amax reduction group on QuantizedT..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Comment on lines
326
to
+329
| """Quantize tensor""" | ||
| return self.quantize(tensor) | ||
| if amax_reduction_group is None: | ||
| return self.quantize(tensor) | ||
| return self.quantize(tensor, amax_reduction_group=amax_reduction_group) |
Contributor
There was a problem hiding this comment.
None guard here is redundant: self.quantize(tensor) and self.quantize(tensor, amax_reduction_group=None) are identical because quantize defaults the argument to None. The branch just adds noise.
Suggested change
| """Quantize tensor""" | |
| return self.quantize(tensor) | |
| if amax_reduction_group is None: | |
| return self.quantize(tensor) | |
| return self.quantize(tensor, amax_reduction_group=amax_reduction_group) | |
| """Quantize tensor""" | |
| return self.quantize(tensor, amax_reduction_group=amax_reduction_group) |
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
Comment on lines
+200
to
+203
| @property | ||
| def rht_matrix(self) -> torch.Tensor: | ||
| """RHT matrix (fetched from the process-global cache, not stored per quantizer).""" | ||
| return get_rht_matrix(self._with_random_sign_mask, torch.cuda.current_device()) |
Contributor
There was a problem hiding this comment.
NVFP4Quantizer instances
rht_matrix is now a property that reads self._with_random_sign_mask, but _with_random_sign_mask is a new field that did not exist in pickled state produced before this change. When Python's default __setstate__ (i.e., self.__dict__.update(state)) loads an old pickle, _with_random_sign_mask is absent, so any access to the rht_matrix property raises AttributeError. A __setstate__ that infers _with_random_sign_mask from the old stored rht_matrix (or supplies a safe default) would preserve backward compatibility for serialized quantizers.
The amax-reduction refactor moved the reduction group from quantizer state to a per-call `amax_reduction_group` argument, but three call paths were not updated, so they silently skipped amax reduction (or raised): - DebugQuantizer.quantize()/update_quantized() did not accept `amax_reduction_group`, raising TypeError under sequence parallelism. - gather_along_first_dim() forwarded the group to the FP8 all-gather but not to _all_gather_nvfp4(), so NVFP4 row-parallel + sequence parallel quantized with a local amax (exact dgrad mismatch). - NVFP4Tensor had no FSDP2 amax-group stash/pop (unlike Float8Tensor), so FSDP2 weight shards re-quantized with inconsistent global scales. Mirror the existing Float8 handling for NVFP4 and forward the group through the NVFP4 all-gather and the debug quantizer. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
for more information, see https://pre-commit.ci
FSDP2 re-quantizes a sharded weight after the optimizer step through QuantizedTensor.__torch_dispatch__ -> quantize_ -> update_quantized (and via the Tensor.data setter, _set_data). On this branch the amax reduction group is no longer stored on the quantizer, but quantize_ did not thread a group through, so the weight shards re-quantized with a local per-rank amax instead of a global one. The all-gathered FP8/NVFP4 weight then dequantized to different values per shard, failing _check_fp8_fsdp2_allgather for Float8CurrentScaling and NVFP4BlockScaling with fp8_init=True. Add an explicit per-call `amax_reduction_group` argument to the `quantize_` hierarchy (QuantizedTensorStorage / QuantizedTensor / Float8/NVFP4/MXFP8/ Float8Blockwise), keeping it consistent with quantize()/update_quantized(). When no group is supplied, Float8Tensor/NVFP4Tensor fall back to the group stashed on the tensor by fsdp_pre_all_gather, so in-place optimizer re-quantization stays globally scaled without every caller having to thread the group. The FSDP stash is now kept (read, not popped) and refreshed each iteration, matching the previous behavior where the group lived on the quantizer. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Collaborator
Author
|
Blocked by FSDP bug, refactor in progress. I plan to store .amax_reduction_group in QuantizedTensor. |
…rough quantize_ Float8Tensor/NVFP4Tensor carry an optional ``amax_reduction_group`` attribute, set by FSDP2 in fsdp_pre_all_gather and kept (refreshed every iteration) rather than popped. ``update_quantized`` reads it from the destination tensor when no group is supplied per call, so in-place re-quantization via quantize_/_set_data (e.g. the FSDP2 optimizer step) stays globally scaled without threading a group through the quantize_ signature. The deprecated quantizer-stored group is still honored as a lowest-priority fallback. The current-scaling reference quantizer gains the same per-call + destination-tensor support. Supersedes the alternative of adding an explicit amax_reduction_group argument to the quantize_ hierarchy. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
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
1 participant
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.
Description
This makes adding torch.compile support much easier.
Move amax reduction process group handling out of quantizer state and pass it per quantization call instead. This avoids storing process groups inside quantizers while keeping deprecated stored-group fallback behavior for compatibility.
Fixes # (issue)
Type of change
Changes
Please list the changes introduced in this PR:
amax_reduction_groupthroughquantize/module call paths instead of storing it on quantizers.Checklist: