[PyTorch][torch.compile] Add TensorProto mechanism - #3153
Open
pggPL wants to merge 32 commits into
Open
Conversation
Contributor
Greptile SummaryThis PR introduces
Confidence Score: 5/5Safe to merge. The new TensorProto mechanism, flatten/unflatten protocol, and _apply attribute-restore fix are well-contained, well-tested, and do not alter existing quantization or training behaviour. No functional regressions were found. The TensorProto/flatten/unflatten protocol is validated by end-to-end tests covering CPU compilation, FakeTensorMode, structural and functional parity with the C++ allocator, and flatten round-trips across all four quantizer formats. The only findings are annotation style issues and an uncached hardware-capability query, neither of which affects correctness. Files Needing Attention: No files require special attention; the annotation style in mxfp8_tensor_storage.py and float8_blockwise_tensor_storage.py can be tidied at leisure. Important Files Changed
Reviews (20): Last reviewed commit: "Carry the storage class itself through t..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
pggPL
force-pushed
the
tensor_proto_mechanism
branch
8 times, most recently
from
June 29, 2026 13:46
9e78a6c to
50c11cd
Compare
Squashed PR #8 (tensor_proto_mechanism) onto the rebased base. Adds TensorProto (pure-Python, torch.compile-traceable quantized-tensor allocation via Quantizer.alloc_tensors + storage __tensor_flatten__/__tensor_unflatten__), Linear fake fwd/bwd impls for the custom-op path, and tests. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
The cached FP8 weight is the same tensor returned as new_weight_workspace (cache miss) or passed in as weight_workspace (cache hit). A custom op may not return a tensor that aliases an input or another return, so mark those slots and reconstruct wt_save in _linear_setup_ctx instead of saving it twice. Mirrored in the fake impl so the saved-slot layout matches. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
NVFP4Quantizer._describe_buffers grouped each amax right after its scale (per-usage), diverging from NVFP4TensorStorage._FLATTEN_TENSOR_BUFFERS (amax buffers last). The order is functionally irrelevant (buffers are consumed by name in alloc_tensors and reordered in TensorProto.inner_names), but aligning it makes describe/flatten agree and fixes test_to_tensor_proto_quantized[nvfp4]. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
…upport - TensorProto.inner_names now raises if the quantizer describes buffer(s) absent from the storage's _FLATTEN_TENSOR_BUFFERS, instead of silently appending them. - Gate the nvfp4 proto-quantizer param on nvfp4_available so it skips on hardware without NVFP4 support rather than failing. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
…escribe_buffers Access NVFP4Quantizer @staticmethods (convert_shape_for_fp4, get_columnwise_shape) via the class instead of the instance. Under torch.compile, instance access of a @staticmethod on a value-opaque object crashes Dynamo guard generation with "'function' object has no attribute '__func__'" (pytorch/pytorch#182741). Temporary workaround until the PyTorch-side fix lands. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
pggPL
force-pushed
the
tensor_proto_mechanism
branch
from
July 7, 2026 10:04
ff48e52 to
e36cf6d
Compare
ptrendx
reviewed
Jul 7, 2026
7 tasks
The union is intentional: fields may carry bare QuantizedTensorStorage objects (internal-quantizer optimization), and the annotation is introspected in the follow-up custom-op PR to build the op schema with flatten/unflatten slots. Also note the size()/.shape asymmetry and how TensorProto handles it. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Make .shape valid on bare storages (derived from size()), so Tensor, QuantizedTensor, bare storage and TensorProto all expose the same attribute. Wrapper subclasses defer to the native TensorBase.shape. Simplifies the shape fallback in to_tensor_proto. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Address review: build the same quantized tensor via make_empty (C++, tex.create_empty_quantized_tensor) and via the Python primitives (_describe_buffers + create_metadata + alloc_tensors + __tensor_unflatten__) and check structural parity (class, buffer set, per-buffer shape/dtype/device, flatten context) and functional parity (the real quantize kernel writes bit-identical results into both, dequantize matches), across quantizer families x rowwise/columnwise x wrapper/internal. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Address review: the change stands on its own as a correctness fix; drop the detailed (and imprecise) fake-impl/cudagraph justification. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Address review: wt_save is known non-None past the first branch, so 'X is not None and wt_save is X' reduces to 'wt_save is X'. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Address review: replace the local _contiguous_stride helper with the torch one (stable at this path since v1.13); it also matches the ATen contiguous-stride convention for zero-size dims and handles SymInts. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Address review: a silent no-op diverges from the real object's behavior (plain torch.Tensor has no update_usage), which is exactly the class of fake/real mismatches the proto is meant to avoid. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Address review: after the QuantizedTensorStorage.shape property the storage and plain-tensor paths differed only in getattr fallbacks (dtype/_dtype, _quantizer), which work uniformly for all input kinds; drop the isinstance branch and the local import it needed. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Address review: the saved_weight slot is unconditionally aliased to the weight parameter in forward, so it is never None in backward. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
…to tensor_proto_mechanism
shino16
approved these changes
Jul 14, 2026
Comment on lines
+89
to
+90
| # machine-read when the args bag becomes a torch.compile custom op (follow-up | ||
| # PR): such fields get flatten/unflatten slots so a bare storage can cross the |
There was a problem hiding this comment.
Note: We will update this comment as part of the custom op PR.
…fers Factor the quantized-reassembly tail of create_tensor into assemble(), which rebuilds a tensor from already-materialized inner buffers (in inner_names() order). create_tensor becomes assemble(create_inner_tensors()). This is the reusable primitive the torch.compile custom-op boundary uses to rebuild an op's quantized outputs from its flat Tensor[] payload. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Collaborator
Author
|
/te-ci pytorch |
13 tasks
ptrendx
reviewed
Jul 28, 2026
Comment on lines
+842
to
+855
| _PROTO_QUANTIZERS = [ | ||
| pytest.param(_current_scaling, (4, 8), id="fp8_current_scaling"), | ||
| pytest.param(_mxfp8, (64, 128), id="mxfp8"), | ||
| pytest.param(_blockwise, (128, 256), id="fp8_blockwise"), | ||
| pytest.param( | ||
| _nvfp4, | ||
| (64, 128), | ||
| id="nvfp4", | ||
| marks=pytest.mark.skipif( | ||
| not nvfp4_available, | ||
| reason="NVFP4 is not available", | ||
| ), | ||
| ), | ||
| ] |
Member
There was a problem hiding this comment.
Why is nvfp4 skipped when not available but not MXFP8?
test_python_alloc_matches_cpp_make_empty compared buffers the quantize kernel never writes: the scale-inv padding is allocated uninitialized by both paths, so the bit-exact comparison saw random bytes and failed on H100/B200 for fp8_blockwise. Zero every buffer before quantizing, so the comparison covers kernel output only. Also drop the param-level skips on the nvfp4 entries of _PROTO_QUANTIZERS and _VALUE_QUANTIZERS. is_fp8_available() and friends run at import time and go through torch.cuda.current_device(), so this module cannot be collected without CUDA at all and skipif(not torch.cuda.is_available()) never fires; the same goes for the torch.cuda.is_available() halves of the _hw_available() guards. Gating nvfp4 on nvfp4_available was also inconsistent with MXFP8 and blockwise, which are gated at runtime and only in the tests that run a kernel -- the allocation primitives themselves are pure Python and describe the layout on any HW. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
_linear_forward_impl_fake diverged from quantize_weight on the weight workspace in three ways: - it produced a new workspace only when update_ws was true, but the real cache-miss path returns (out, out) whenever cache=True, regardless of update_workspace; a first call with is_first_microbatch=False therefore lost the workspace and the "new_workspace" saved-weight alias; - it treated any non-None cached workspace as a hit, while the real path runs _is_weight_workspace_valid() first and falls through to a miss when the cached buffer layout no longer matches the quantizer's usage; - it kept quantizer.internal, so the descriptor resolved to a bare storage class, while the real path quantizes persistent workspaces with internal=False and caches wrapper tensors. On a cache hit the weightmat is now the workspace descriptor itself, and on a miss with cache_weight it is the same proto object returned as the new workspace, matching quantize_weight's aliasing. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Eager forward forces save_original_input=False for backward_override="dequantized", but the fake only handled "high_precision". With save_original_input=True and that override, the fake aliased the original input into saved-tensor slot 0 while eager saved a quantized input with rowwise-only usage, so the saved payload layout and the compiled backward setup disagreed. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
The output proto's requires_grad considered only the input and the weight, so a frozen input and weight with a trainable bias described the output as non-differentiable while eager _Linear.apply produces a differentiable one. bias_requires_grad is already False when there is no bias. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Collaborator
Author
|
/te-ci pytorch |
_linear_forward_impl_fake / _linear_backward_impl_fake, and the eager-side changes that existed only to support them (reading the requires_grad flags off LinearFwdArgs, the new_workspace/weight_workspace alias dedup and the _linear_setup_ctx signature carrying (out, new_weight_workspace)), have no caller in this PR: nothing registers them as a custom op's fake, so nothing exercises them here. They belong with the custom-op registration that consumes them. This PR is left as the TensorProto mechanism proper -- the proto, the storage flatten protocol and the pure-Python quantizer allocation hooks -- which the new tests do cover. linear.py returns to its upstream state. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Collaborator
Author
|
/te-ci pytorch |
Quantized tensors now implement the wrapper-subclass flatten protocol, so nn.Module._apply moves them with torch.utils.swap_tensors instead of the `param.data = ...` path. The swap exchanges the parameter's entire __dict__: that is how the inner buffers reach the surviving object, but it also carries off everything attached to the parameter from the outside. TE relies on several such attributes: _high_precision_init_val and its two accessors (quantized_model_init(preserve_high_precision_init_val=True)), plus main_grad, grad_added_to_main_grad and overwrite_main_grad, which Megatron-Core attaches. They survived before only because `param.data = ...` is a no-op for a wrapper subclass -- the outer tensor is a zero-storage shell and the assignment never touched __dict__, so device moves silently did nothing at all. Snapshot the parameters' __dict__ before delegating to nn.Module._apply and restore the entries the swap dropped, rebinding bound accessors to the surviving parameter. Entries still present afterwards are the tensor's own state, where the post-swap value is the correct one. Covers the two test_sanity grouped-linear high-precision-init tests that broke on B200, and adds a direct test over .cuda() / .cpu() / .half(). Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Collaborator
Author
|
/te-ci pytorch L1 |
pggPL
added a commit
to pggPL/TransformerEngine
that referenced
this pull request
Jul 29, 2026
Carried over from the TensorProto PR (NVIDIA#3153), where these impls used to live without a caller; they belong here, with the custom-op registration that consumes them. - Weight workspace: quantize_weight returns a fresh workspace on every cache miss with cache=True, not only when update_workspace is set; it discards a cached workspace that fails _is_weight_workspace_valid; and it quantizes persistent workspaces with internal=False so the cache holds wrapper tensors. The fake did none of the three. - backward_override="dequantized" forces save_original_input=False in the eager forward; the fake only handled "high_precision", so it aliased the original input where eager saves a rowwise-only quantized one. - The output's requires_grad ignored the bias, describing the output of a bias-only-trainable Linear as non-differentiable. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
The restore loop keys off "present after the swap": what survived is the tensor's own state, what did not is an externally attached annotation. That holds only as long as every declared buffer really is present afterwards. If one were not, the loop would quietly put the pre-move value back and splice a buffer from the old device (or from before a dtype conversion) into the moved parameter -- silently wrong numerics rather than a crash. Raise instead when a name from _FLATTEN_TENSOR_BUFFERS is about to be restored. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
nn.Module._apply only assigns to self._parameters, never removes entries, so a missing parameter after it returns means something unexpected happened. Skipping it silently dropped every attribute attached to that parameter -- the failure this override exists to prevent. Match the buffer check and fail loudly. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Reverts 6e61d36. The check guarded a case that cannot arise today: the storages always set every declared buffer attribute, to None when unused, so the key is present whatever the usage flags say. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Each storage class listed its tensor buffers twice: once as a field
annotation, once as an (attribute, constructor kwarg) pair in
_FLATTEN_TENSOR_BUFFERS, in a different order and further down the file.
Adding a buffer meant remembering both.
Mark the field instead -- _scale_inv: Annotated[torch.Tensor,
Buffer("fp8_scale_inv")] -- and collect the declarations in
__init_subclass__, which already runs there for the storage registry.
_FLATTEN_TENSOR_BUFFERS survives as the derived attribute, so every consumer
is untouched, and the collected values are identical to the hand-written
tuples for all nine registered classes.
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
"Buffer" collides with nn.Module's buffers, which are a different thing, and _FLATTEN_TENSOR_BUFFERS named a consumer (__tensor_flatten__) rather than the thing itself -- the list has four of them. PyTorch calls exactly this concept "inner tensors", which TensorProto.inner_names() already follows. Also drop the underscore from the two hooks every quantizer has to implement. They were the only members of the extension contract marked private, which is why the tests needed seven protected-access waivers to call them; the members nobody overrides (alloc_tensors, create_metadata) were public already. Buffer -> InnerTensor _FLATTEN_TENSOR_BUFFERS -> _INNER_TENSORS _describe_buffers -> inner_tensor_specs _storage_metadata -> storage_metadata Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
__tensor_flatten__ put the class qualname in the context and __tensor_unflatten__ looked it up in a module-level registry, populated from __init_subclass__. The indirection bought nothing: dynamo bakes the class object into the graph as a constant just as happily, which is what the custom-op branch already relies on. Store type(self) directly and drop _STORAGE_REGISTRY. __init_subclass__ stays for collecting the InnerTensor field annotations. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Collaborator
Author
|
/te-ci pytorch L1 |
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
3 participants
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 PR introduces
TensorProto— a data-free prototype of a tensor (or quantized tensor) that captures everything needed to reason about and rebuild a tensor without holding any storage: its logicalshape/dtypeand, for quantized tensors, the value-opaquequantizerdefining the layout.The key property is that
TensorProto.create_tensor()materializes a quantized tensor purely in Python (viaQuantizer.alloc_tensors+ the storage's__tensor_unflatten__), so it traces undertorch.compile(fullgraph=True)with no graph break — unlikemake_empty, which goes through the opaque C++tex.create_empty_quantized_tensor. This is the foundation for writingtorch.librarycustom-op fake implementations of quantized ops.This builds on the value-opaque quantizer work (so a
TensorProtois itself safe to treat as a compile-time constant).Type of change
Changes
dynamo.py: AddTensorProtodataclass (shape,dtype,quantizer,requires_grad,device) withis_quantized,inner_names(),create_metadata()andcreate_tensor(), plus ato_tensor_proto()helper that builds a proto from a plaintorch.Tensoror aQuantizedTensorStorage/QuantizedTensor.quantized_tensor.py:__tensor_flatten__/__tensor_unflatten__) toQuantizedTensorStorage, driven by a per-class_FLATTEN_TENSOR_BUFFERSdeclaration of(attribute_name, constructor_kwarg)pairs._STORAGE_REGISTRY(populated via__init_subclass__) so__tensor_unflatten__can resolve a concrete storage/wrapper class from its qualname inside an FX graph.Quantizer:alloc_tensors,create_metadata, and the opt-in overrides_describe_buffers,_storage_scalars,_resolve_storage_cls.Float8CurrentScalingQuantizer,MXFP8QuantizerandFloat8BlockQuantizer._FLATTEN_TENSOR_BUFFERSforFloat8TensorStorage,MXFP8TensorStorageandFloat8BlockwiseQTensorStorage.ops/basic/basic_linear.py: Add allocation-free_functional_forward_fake/_functional_backward_fakethat operate onTensorProtoand return output/gradient protos, as a basis for custom-op fake impls (single-device only; TP/SP shape effects not yet modeled).tests/pytorch/test_tensor_proto.py(CPU smoke tests for_describe_buffers/alloc_tensors/create_metadata, flatten round-trip, andto_tensor_proto) andtorch.compilefullgraph tests intest_torch_compile.py.Checklist: