[PyTorch][torch.compile] Add TensorProto mechanism by pggPL · Pull Request #8 · pggPL/TransformerEngine

pggPL · 2026-06-06T14:52:17Z

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 logical shape/dtype and, for quantized tensors, the value-opaque quantizer defining the layout.

The key property is that TensorProto.create_tensor() materializes a quantized tensor purely in Python (via Quantizer.alloc_tensors + the storage's __tensor_unflatten__), so it traces under torch.compile(fullgraph=True) with no graph break — unlike make_empty, which goes through the opaque C++ tex.create_empty_quantized_tensor. This is the foundation for writing torch.library custom-op fake implementations of quantized ops.

This builds on the value-opaque quantizer work (so a TensorProto is itself safe to treat as a compile-time constant).

Type of change

Documentation change (change only to the documentation, either a fix or a new content)
Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
Infra/Build change
Code refactoring

Changes

dynamo.py: Add TensorProto dataclass (shape, dtype, quantizer, requires_grad, device) with is_quantized, inner_names(), create_metadata() and create_tensor(), plus a to_tensor_proto() helper that builds a proto from a plain torch.Tensor or a QuantizedTensorStorage/QuantizedTensor.
quantized_tensor.py:
- Add the PyTorch wrapper-subclass flatten protocol (__tensor_flatten__ / __tensor_unflatten__) to QuantizedTensorStorage, driven by a per-class _FLATTEN_TENSOR_BUFFERS declaration of (attribute_name, constructor_kwarg) pairs.
- Add a _STORAGE_REGISTRY (populated via __init_subclass__) so __tensor_unflatten__ can resolve a concrete storage/wrapper class from its qualname inside an FX graph.
- Add pure-Python, traceable allocation hooks to Quantizer: alloc_tensors, create_metadata, and the opt-in overrides _describe_buffers, _storage_scalars, _resolve_storage_cls.
Quantizers: Implement the allocation hooks for Float8CurrentScalingQuantizer, MXFP8Quantizer and Float8BlockQuantizer.
Storage classes: Declare _FLATTEN_TENSOR_BUFFERS for Float8TensorStorage, MXFP8TensorStorage and Float8BlockwiseQTensorStorage.
ops/basic/basic_linear.py: Add allocation-free _functional_forward_fake / _functional_backward_fake that operate on TensorProto and return output/gradient protos, as a basis for custom-op fake impls (single-device only; TP/SP shape effects not yet modeled).
Tests: Add tests/pytorch/test_tensor_proto.py (CPU smoke tests for _describe_buffers/alloc_tensors/create_metadata, flatten round-trip, and to_tensor_proto) and torch.compile fullgraph tests in test_torch_compile.py.

Checklist:

I have read and followed the contributing guidelines
The functionality is complete
I have commented my code, particularly in hard-to-understand areas
I have made corresponding changes to the documentation
My changes generate no new warnings
I have added tests that prove my fix is effective or that my feature works
New and existing unit tests pass locally with my changes

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>

kshitij12345

Overall looks good

Would it be possible to reduce duplication between _linear_forward_impl_fake and _linear_forward_impl.

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>

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>

_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>

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>

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>

Signed-off-by: Tim Moon <tmoon@nvidia.com>

Bump collective stream min JAX version to 0.10.1 Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

* Add opt-in reduced precision output for cuDNN MXFP8 norm Signed-off-by: sraman-rgb <sraman@nvidia.com> * Warn when cuDNN MXFP8 norm changes intermediate dtype Signed-off-by: sraman-rgb <sraman@nvidia.com> * Scope cuDNN MXFP8 dtype override to FP8 output Signed-off-by: sraman-rgb <sraman@nvidia.com> --------- Signed-off-by: sraman-rgb <sraman@nvidia.com>

…CC 10.x (NVIDIA#3056) * Refine the support for D=256 on Blackwell server type GPUs Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Add deterministic tests for D=256 for sm10.x Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Test clean up Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Add PyT side tests for D=256 cuDNN fused attn on SM100-110 Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Refine the PyT D=256 tests Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix the filtering condition for bias type for D=256 on sm10x for cudnn fused attn Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Code clean up Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Switching test check logic around Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Add THD xfail marked tests for testing Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Fix the jax test bias condition fo skipping for D=256 Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * nit: Fix comment Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Move the D=256 changes so as to follow an order of increasing cuDNN versions Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add support for D=256 THD in TE Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Add TE JAX CP test support to verify D=256 for AG and Ring Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Add PyTorch D256 CP tests Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * nit: clean up comments Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Correct the atgs passed to test dpa Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Simplify THD fused attention layout check Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Rename PyTorch D256 fused attention tests Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Run JAX fused attention before softmax Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Relabel THD D256 cuDNN guard to 9.25 Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fold D256 PyTorch CP tests Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Rename D256 PyTorch attention tests Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Remove D256 PyTorch attention param ids Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Disable FA4 D256 SWA on SM100 Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Fix FA4 D256 packed QKV handling Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Use is_training in DPA test helper Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> --------- Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Signed-off-by: Fred Heinecke <fheinecke@nvidia.com> Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

remove nv-internal comments Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>

[JAX] Add NCCL EP drop-on-overflow policy and total_recv_tokens output Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

Addresses the latest review round: - Rename TensorProto -> TensorSpec, to_tensor_proto -> to_tensor_spec and dynamo/tensor_proto.py -> dynamo/tensor_spec.py. "Proto" collided with ONNX/protobuf and invented a new term for something PyTorch already has vocabulary for; "spec" matches DTensorSpec / tf.TensorSpec. It is not a tensor subclass and it is not fake-specific (create_tensor() in eager builds a real tensor), so FakeQuantizedTensor would not fit. - inner_names(): verify that inner_tensor_specs follows the storage's _INNER_TENSORS order instead of silently reordering. All four quantizers already emit that order, so the reorder was a no-op and the docstring rationale (NVFP4 grouping amax after each scale) was stale. A quantizer that breaks the contract now fails loudly instead of being papered over. - Use the real availability reasons (reason_for_no_nvfp4, reason_for_no_fp8_block_scaling) in _skip_if_dequantize_unsupported instead of hardcoded strings. - Speak of "inner tensors" consistently instead of "buffers", matching _INNER_TENSORS / inner_tensor_specs / create_inner_tensors. - Fold test_tensor_spec_create_tensor_{eager,fake} into one test parametrized on fake, and drop test_primitives_unflatten_compiles: its production-code coverage is a subset of test_tensor_spec_create_tensor_compiles, the only part unique to it being the test helper's meta-device stride computation. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

… rename The previous commit renamed "buffers" to "inner tensors"/"specs" with a word-boundary substitution, which also rewrote three comments in code this PR does not touch: the GPU-buffers and FP8-buffers notes in float8_tensor and the device-inference note in mxfp8_tensor. Restore their original wording. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Add NLTK pin <3.10.1 in encoder requirements Signed-off-by: Alex Y. Chan <alechan@nvidia.com>

Signed-off-by: Haijie Zhi <133995660+cupkk@users.noreply.github.com>

* Stage 1&2: Python containers + quantize/gemm dispatch/unwrap Signed-off-by: Evgeny <etsykunov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Enable quantized_model_init Signed-off-by: Evgeny <etsykunov@nvidia.com> * FSDP support Signed-off-by: Evgeny <etsykunov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Enable CPU offloading Signed-off-by: Evgeny <etsykunov@nvidia.com> * Activation recomputation Signed-off-by: Evgeny <etsykunov@nvidia.com> * TP/SP Signed-off-by: Evgeny <etsykunov@nvidia.com> * Resolve comments: hybrid uniform list, make_empty try, __repr__, etc. Signed-off-by: Evgeny <etsykunov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Respect usage Signed-off-by: Evgeny <etsykunov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Misc minor fixes: comments, tests, etc. Signed-off-by: Evgeny <etsykunov@nvidia.com> * Towards MLM integration Signed-off-by: Evgeny <etsykunov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Resolve comments: improve fsdp/tp/sp tests + amax reduction fix Signed-off-by: Evgeny <etsykunov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Enable FSDP2 hybrid protocol for Float8Block tensor Signed-off-by: Evgeny <etsykunov@nvidia.com> * Enable Identity (no-op) quantization Signed-off-by: Evgeny <etsykunov@nvidia.com> * Bug fixing Signed-off-by: Evgeny <etsykunov@nvidia.com> * More fixes Signed-off-by: Evgeny <etsykunov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Minor fixes Signed-off-by: Evgeny <etsykunov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Handle transpose-only Float8 in hybrid distopt Signed-off-by: Evgeny <etsykunov@nvidia.com> * Refactor hybrid tensors to own only parent quantizer Signed-off-by: Evgeny <etsykunov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Enable full hybrid nvfp4 recipe (col only nvfp4) + minor fixes Signed-off-by: Evgeny <etsykunov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * quantization_factory_base/zoo Signed-off-by: Evgeny <etsykunov@nvidia.com> * Enable columnwise_source and hybrid recipes Signed-off-by: Evgeny <etsykunov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Respect quantizer veto for save_original_inp Signed-off-by: Evgeny <etsykunov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * IMprove hybrid factory tests Signed-off-by: Evgeny <etsykunov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix attention factories and align with env-var setup Signed-off-by: Evgeny <etsykunov@nvidia.com> * Update quantization factory zoo Signed-off-by: Evgeny <etsykunov@nvidia.com> * Resolve comments Signed-off-by: Evgeny <etsykunov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add MHA boudary quantization tests Signed-off-by: root <root@prenyx0286.a51.clusters.nvidia.com> * Improve numerical correctness of the tests Signed-off-by: Evgeny <etsykunov@nvidia.com> * Resolve comments: fix rowwise_dequantized and improve quantizer validation in GroupedLinear, etc Signed-off-by: Evgeny <etsykunov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Finish comment resolving: materialize quantizer once, + minor fixes Signed-off-by: Evgeny <etsykunov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add automatic alignment contract for CustomRecipe Signed-off-by: Evgeny <etsykunov@nvidia.com> * Misc fixes: float8 fsdp, column-only float8, checks, etc. Signed-off-by: Evgeny <etsykunov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Hopper fixes Signed-off-by: Evgeny <etsykunov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * More Hopper fixes Signed-off-by: Evgeny <etsykunov@nvidia.com> * Guard hybrid partial-master distopt updates Signed-off-by: Evgeny <etsykunov@nvidia.com> * Blackwell CI fixes Signed-off-by: Evgeny <etsykunov@nvidia.com> * Docs minor update Signed-off-by: Evgeny <etsykunov@nvidia.com> * Improve fsdp2 and other tests Signed-off-by: Evgeny <etsykunov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix delayed-scaling hybrid FP8 FSDP guard Signed-off-by: Evgeny <etsykunov@nvidia.com> * Resolve comments Signed-off-by: Evgeny <etsykunov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * refactor custom recipe Signed-off-by: Evgeny <etsykunov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Minor changes Signed-off-by: Evgeny <etsykunov@nvidia.com> * Fix CI Signed-off-by: Evgeny <etsykunov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix CI Signed-off-by: Evgeny <etsykunov@nvidia.com> * Fix CI Signed-off-by: Evgeny <etsykunov@nvidia.com> --------- Signed-off-by: Evgeny <etsykunov@nvidia.com> Signed-off-by: Evgeny Tsykunov <etsykunov@nvidia.com> Signed-off-by: root <root@prenyx0286.a51.clusters.nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: root <root@prenyx0286.a51.clusters.nvidia.com>

Three conflicts, all "both sides added code in the same spot", resolved by keeping both: - quantized_tensor.py: main's FSDP2 buffer protocol next to this branch's subclass flatten protocol. - float8_tensor.py: main's is_requantization_safe next to the quantizer's storage_metadata / inner_tensor_specs. - storage/float8_tensor_storage.py: import line, both InnerTensor and _resolve_view_shape. The new storages main brings in (HybridQuantizedTensorStorage, IdentityTensorStorage) declare no InnerTensor fields and are deliberately not covered by the flatten protocol. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

@staticmethods

* Add TensorProto mechanism for data-free quantized tensor allocation 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> * [PyTorch] torch.compile: dedup cached FP8 weight from saved-for-backward 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> * [PyTorch] nvfp4: emit _describe_buffers in canonical flatten order 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> * Address review: error on undescribed buffers, gate nvfp4 test on HW support - 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> * [PyTorch] Workaround torch.compile staticmethod guard bug in NVFP4 _describe_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> * [PyTorch] Document TensorOrQuantized union (review feedback) 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> * [PyTorch] Add shape property to QuantizedTensorStorage 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> * [PyTorch] Test parity of Python alloc vs C++ make_empty 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> * [PyTorch] Simplify comment on captured requires_grad flags 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> * [PyTorch] Drop redundant None guards in wt_save alias dedup 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> * [PyTorch] Use torch._prims_common.make_contiguous_strides_for 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> * [PyTorch] Raise on update_usage of a non-quantized TensorProto 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> * [PyTorch] Collapse to_tensor_proto branches 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> * [PyTorch] Drop dead weight_fp8 fallback in fake backward 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> * [PyTorch] Drop redundant comment on saved_weight Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * [PyTorch] Add TensorProto.assemble for rebuilding from ready-made buffers 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> * Fix blockwise alloc parity test and drop dead CUDA skips 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> * Mirror quantize_weight's cache semantics in the Linear fake forward _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> * Honor the dequantized backward override in the Linear fake forward 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> * Include the bias in the Linear fake output's differentiability 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> * Move the Linear fake impls out to the custom-op branch _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> * Keep attributes attached to quantized parameters across _apply 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> * Fail loudly if a buffer vanishes from a parameter during _apply 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> * Raise when a parameter vanishes during _apply instead of skipping it 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> * Drop the vanished-buffer check from _apply 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> * Declare flat buffers on the field instead of in a parallel tuple 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> * Name the flat buffers after PyTorch's own term for them "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> * Carry the storage class itself through the flatten context __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> * Rename TensorProto to TensorSpec and tighten the inner-tensor contract Addresses the latest review round: - Rename TensorProto -> TensorSpec, to_tensor_proto -> to_tensor_spec and dynamo/tensor_proto.py -> dynamo/tensor_spec.py. "Proto" collided with ONNX/protobuf and invented a new term for something PyTorch already has vocabulary for; "spec" matches DTensorSpec / tf.TensorSpec. It is not a tensor subclass and it is not fake-specific (create_tensor() in eager builds a real tensor), so FakeQuantizedTensor would not fit. - inner_names(): verify that inner_tensor_specs follows the storage's _INNER_TENSORS order instead of silently reordering. All four quantizers already emit that order, so the reorder was a no-op and the docstring rationale (NVFP4 grouping amax after each scale) was stale. A quantizer that breaks the contract now fails loudly instead of being papered over. - Use the real availability reasons (reason_for_no_nvfp4, reason_for_no_fp8_block_scaling) in _skip_if_dequantize_unsupported instead of hardcoded strings. - Speak of "inner tensors" consistently instead of "buffers", matching _INNER_TENSORS / inner_tensor_specs / create_inner_tensors. - Fold test_tensor_spec_create_tensor_{eager,fake} into one test parametrized on fake, and drop test_primitives_unflatten_compiles: its production-code coverage is a subset of test_tensor_spec_create_tensor_compiles, the only part unique to it being the test helper's meta-device stride computation. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Restore three unrelated comments mangled by the buffers/inner-tensors rename The previous commit renamed "buffers" to "inner tensors"/"specs" with a word-boundary substitution, which also rewrote three comments in code this PR does not touch: the GPU-buffers and FP8-buffers notes in float8_tensor and the device-inference note in mxfp8_tensor. Restore their original wording. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> --------- Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

pggPL force-pushed the make_qunatizers_opaque branch from 33e9d73 to d341eeb Compare June 16, 2026 15:21

pggPL requested a review from cyanguwa as a code owner June 16, 2026 15:21

pggPL force-pushed the tensor_proto_mechanism branch from 2cccc30 to 2e252f9 Compare June 16, 2026 15:31

pggPL force-pushed the tensor_proto_mechanism branch from 2e252f9 to ba92f5b Compare June 16, 2026 16:05

pggPL force-pushed the tensor_proto_mechanism branch from ba92f5b to b1273ea Compare June 16, 2026 16:12

kshitij12345 reviewed Jun 22, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/dynamo/tensor_proto.py Outdated

Comment thread tests/pytorch/test_torch_compile.py Outdated

pggPL force-pushed the make_qunatizers_opaque branch from e4a879b to adc65f6 Compare June 29, 2026 07:14

pggPL force-pushed the tensor_proto_mechanism branch from 85355a6 to c1e40b2 Compare June 29, 2026 07:16

pggPL force-pushed the make_qunatizers_opaque branch from adc65f6 to c7bbc83 Compare June 29, 2026 07:33

pggPL force-pushed the tensor_proto_mechanism branch from c1e40b2 to e760487 Compare June 29, 2026 07:34

pggPL force-pushed the make_qunatizers_opaque branch from c7bbc83 to f592cbb Compare June 29, 2026 09:26

pggPL force-pushed the tensor_proto_mechanism branch from e760487 to 50d5c21 Compare June 29, 2026 09:26

pggPL force-pushed the make_qunatizers_opaque branch from f592cbb to 945f62d Compare June 29, 2026 09:34

pggPL force-pushed the tensor_proto_mechanism branch from 50d5c21 to da709e7 Compare June 29, 2026 09:35

pggPL force-pushed the tensor_proto_mechanism branch 3 times, most recently from 5131ebc to 77831be Compare June 29, 2026 10:24

pggPL force-pushed the tensor_proto_mechanism branch from 77831be to 29e5245 Compare June 29, 2026 12:47

pggPL force-pushed the tensor_proto_mechanism branch from 29e5245 to 99c1377 Compare June 29, 2026 13:10

pggPL added 14 commits July 28, 2026 14:21

Merge remote-tracking branch 'upstream/main' into tensor_proto_mechanism

a30ea45

Merge remote-tracking branch 'upstream/main' into tensor_proto_mechanism

0228c9c

pggPL changed the base branch from make_qunatizers_opaque to main July 31, 2026 12:40

timmoon10 and others added 14 commits July 31, 2026 05:48

Update list of authorized CI users (NVIDIA#3292)

60697de

Signed-off-by: Tim Moon <tmoon@nvidia.com>

[JAX] Bump collective stream min JAX version to 0.10.1 (NVIDIA#3295)

91d53a0

Bump collective stream min JAX version to 0.10.1 Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

[CI] Pin JAX image to 2026-07-21 (NVIDIA#3250)

a9e635a

Signed-off-by: Fred Heinecke <fheinecke@nvidia.com> Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

[Common] Remove nv-internal-* comments (NVIDIA#3280)

6b23ad3

remove nv-internal comments Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

Pin nltk to work around breaking change in tests (NVIDIA#3306)

cb7e2de

Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>

[JAX] EP Dispatch with overflow detection option (NVIDIA#3277)

72f4e23

[JAX] Add NCCL EP drop-on-overflow policy and total_recv_tokens output Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

[JAX] Add NLTK pin <3.10.1 in encoder requirements (NVIDIA#3314)

82a0f6e

Add NLTK pin <3.10.1 in encoder requirements Signed-off-by: Alex Y. Chan <alechan@nvidia.com>

Fix typo in contribution guidelines (NVIDIA#3308)

429c1bc

Signed-off-by: Haijie Zhi <133995660+cupkk@users.noreply.github.com>

Navigation Menu

[PyTorch][torch.compile] Add TensorProto mechanism - #8

[PyTorch][torch.compile] Add TensorProto mechanism #8
pggPL wants to merge 46 commits into
mainfrom
tensor_proto_mechanism

pggPL commented Jun 6, 2026

Uh oh!

kshitij12345 left a comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

12 participants

Conversation

pggPL commented Jun 6, 2026

Description

Type of change

Changes

Checklist:

Uh oh!

kshitij12345 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

12 participants