[PyTorch][torch.compile] Replace TensorProto with make_empty_traceable by kshitij12345 · Pull Request #13 · pggPL/TransformerEngine

kshitij12345 · 2026-07-02T10:33:19Z

Replace the 172-line TensorProto dataclass with a single function make_empty_traceable(quantizer, shape, dtype, device) that directly allocates traceable quantized tensors. The fake impls now return actual tensors (which become FakeTensors under register_fake) instead of intermediate descriptors.

The key insight: make_empty_traceable stashes _te_flat_names and _te_flat_ctx on the resulting tensor. Dynamo treats non-callable attributes on traceable wrapper subclasses as constant metadata, so forward_fn can read slot counts and reassembly info from these attributes without calling tensor_flatten (which would cause a graph break since it returns non-Tensor Python objects).

This eliminates:

TensorProto class and to_tensor_proto helper (tensor_proto.py deleted)
_proto_view (converted tensor fields to TensorProto before fake impls)
_tensor_field_names (identified fields for _proto_view)
_proto_slot_count / _proto_reassemble (operated on TensorProto objects)
TensorProto branch in _value_to_flat_tensors and _format_bwd_result

The fake impls in linear.py now use:

isinstance(inp, QuantizedTensorStorage) instead of inp.is_quantized
weight._quantizer instead of weight.quantizer (TensorProto field)
make_empty_traceable(...) instead of TensorProto(...)
Direct set_usage on quantizer instead of proto.update_usage()

Test Plan:

python -m pytest tests/pytorch/test_torch_compile.py -v

Authored with Claude.

Description

Please include a brief summary of the changes, relevant motivation and context.

Fixes # (issue)

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

Please list the changes introduced in this PR:

Change A
Change B

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

…ompile Give tensorless quantizers (MXFP8, FP8 blockwise, FP8 current-scaling, NVFP4) value-object semantics so torch.compile can treat them as baked-in constants: - Add opt-in value identity to the base Quantizer (_value_fields / _value_key / __eq__ / __hash__). Quantizers holding live tensors (delayed-scaling Float8Quantizer) and custom quantizers keep identity semantics. - New transformer_engine/pytorch/dynamo.py houses the torch.compile glue: __fx_repr__, value-key reconstruction and register_value_opaque_quantizer (gracefully a no-op on PyTorch builds without the opaque-object API). - Register the four tensorless quantizers as value opaque types. Also fix CustomRecipe state caching in TransformerEngineBaseModule: set_meta_tensor now rebuilds quantizers when the CustomRecipe instance changes (e.g. nested te.autocast regions) instead of reusing the first recipe's state, since every CustomRecipe shares the CustomRecipeState type but carries its own qfactory. Move the quantizer value-object tests into tests/pytorch/test_torch_compile.py and add that file to the L0 pytorch unittest QA suite. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

…globals Follow-up to the value-opaque quantizer support: - Remove the module-level _QUANTIZER_VALUE_REGISTRY (qualname -> class) and _quantizer_from_value_key. __fx_repr__ now captures the quantizer class directly in the FX globals and reconstructs via _rebuild_quantizer(cls, items), matching how PyTorch's own value opaque types (e.g. DTensor placements) reconstruct themselves. This removes global mutable state and the qualname collision risk. - Consolidate the quantizer value-object tests in test_torch_compile.py down to two functions and exercise reconstruction through the public __fx_repr__ path instead of internal helpers. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Replace the single dynamo.py module with a dynamo/ package so the torch.compile glue can grow with a clear responsibility split across the stacked branches. This branch owns the value-opaque quantizer layer. * dynamo/quantizer_opaque.py -- register_value_opaque_quantizer and helpers * dynamo/__init__.py -- re-exports the public API so callers keep importing from transformer_engine.pytorch.dynamo unchanged Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

A value-opaque quantizer must not carry live distributed state. Scan the quantizer attributes in __fx_repr__ and raise TypeError if any holds a torch.distributed.ProcessGroup (e.g. a non-None deprecated amax_reduction_group), so it cannot be silently baked into a torch.compile FX graph. Clarify the related comments accordingly. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

NVFP4Quantizer is registered as a value-opaque quantizer but was missing from the value-semantics / __fx_repr__ round-trip test. Add it to _VALUE_QUANTIZERS (skipped without CUDA, which it needs to construct). Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

…__/__hash__ The amax reduction group is excluded from the value key, so a value quantizer that stored one would compare/hash equal to a groupless one and let torch.compile reuse a graph that skips the reduction. __eq__/__hash__ now raise (mirroring __fx_repr__, which already rejects any process-group-bearing quantizer). The group should be passed per quantize call, not stored on the quantizer. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Add is_value_opaque_quantizer() + the _te_compile_value_opaque flag stamped at registration, so dynamo-traced code can detect registered quantizers (and fall back to eager for unregistered ones). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

…fp4 value key - Narrow register_opaque_type except to (RuntimeError, TypeError): the API is already imported above, so ImportError/AttributeError there only mask real errors. - Add test_quantizer_value_object_fullgraph exercising torch.compile(fullgraph=True) end-to-end to verify opaque-type registration took effect. - Restore missing NVFP4Quantizer._with_random_sign_mask assignment required by _value_fields()/_value_key(). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

…trip _rebuild_quantizer only restores value-key fields, so a reconstructed NVFP4Quantizer was missing the derived rht_matrix tensor (not hashable, so not in the value key) and failed at copy()/quantize time. Add a _rebuild_derived_state hook (called by _rebuild_quantizer) that NVFP4Quantizer uses to rebuild rht_matrix from _with_random_sign_mask (lru_cache -> cheap). Extend test_quantizer_value_object to also quantize with the original and the rebuilt quantizer and require bit-exact results (gated on HW support), so a field the kernel needs but the value key omits can no longer slip through. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Move the ProcessGroup guard out of the (overridable) __fx_repr__ into Quantizer._value_key -- the single point every value-materialization path (__eq__/__hash__/__fx_repr__) goes through -- so a custom __fx_repr__ can no longer bypass it. Generalizes the old amax-only check to any field holding a ProcessGroup. Add a test that a value quantizer carrying a live group raises. Addresses review on NVIDIA#3152. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

…assthrough Replace the trivial pass-through fullgraph test with one that drives each production quantizer through a minimal custom op (quantize + dequantize) under torch.compile(fullgraph=True) and compares to eager -- so the opaque-type registration is actually exercised inside the graph (a graph break would make fullgraph=True raise). Op registration sits right before the test. Also drop stale comments referencing the old __fx_repr__-side process-group guard. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

…paque flag - rht_matrix_random_sign_mask_t is a device-independent int derived from _with_random_sign_mask (the device only places a throwaway tensor); fix the misleading comment. - Explain why registration uses a class attribute, not a registry set: is_value_opaque_quantizer is traced inside the compile graph and dynamo can bake a getattr constant but cannot do 'type(q) in set' on the opaque class. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

is_opaque_value_type(cls) sat between the import guard and the register_opaque_type guard, so on a partial/experimental opaque-object build it could raise RuntimeError/TypeError and crash TE import. Move it inside the same except so the 'registration never crashes import' promise holds for both calls. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Squashed PR NVIDIA#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>

@staticmethods

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

…stants; fix SP memory leak; test suite hook-up Wrap CommOverlapCore pybind11 methods that return compile-time constants so torch.compile(fullgraph=True) can trace through them without graph breaks: - `is_fp8_ubuf()` → `ub_is_fp8()` / `get_ub_is_fp8()` in base.py; `_ub_is_fp8()` in gemm.py - `with_cublasmp()` → `ub_is_cublasmp()` in base.py All callers in linear.py, layernorm_linear.py, layernorm_mlp.py, base.py, gemm.py, userbuffers_backward_linear.py and userbuffers_forward_linear.py updated. Fix quantized grad_output not being freed early for column-parallel SP backward. Row-parallel SP already called clear_tensor_data(grad_output) to release the gathered tensor; column-parallel SP quantizes grad_output to Float8TensorStorage but never freed it before returning. Under torch.compile reduce-overhead this leaves 3 live pool tensors at recording end and triggers "Detected 3 tensor(s) in the cudagraph pool not tracked as outputs". Extend the existing clear_tensor_data guard to cover both parallel modes. Fix custom-recipe quantizer state being re-initialised on every forward call even when the recipe object has not changed. The existing early-exit for CustomRecipeState was missing an identity check on the recipe object, so any repeated call with the same recipe would bypass the early-return and rebuild quantizers unnecessarily. Add `if recipe_state.recipe is recipe: return` to restore the intended caching behaviour. Add test_torch_compile.py to L0_pytorch_unittest so the autocast and existing compile tests run in CI. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> (cherry picked from commit bfce3a7)

for more information, see https://pre-commit.ci (cherry picked from commit afe364b)

ToyLinear now overrides get_quantizer_roles so CustomRecipeState doesn't hit the no-roles warning, which graph-breaks under fullgraph=True. qfactory dispatches on role.tensor_type instead of a pre-baked string key. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> (cherry picked from commit 22f80e4)

Squashed PR NVIDIA#9 (linear_compile) onto the rebased base. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> (cherry picked from commit 84dbc6b)

…tom_op Replace the low-level torch.library.Library (_TE_LIB.define/.impl + functional register_fake/register_autograd/register_torch_dispatch with lib=) with the standard torch.library.custom_op API, passing the dynamically built schema explicitly via schema=. register_fake/register_autograd/register_torch_dispatch are now methods on the returned CustomOpDef. Drops the TOR901 Library usage and is robust to re-registration (get_library_allowing_overwrite). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> (cherry picked from commit d590560)

…e sentinel Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Replace the 172-line TensorProto dataclass with a single function make_empty_traceable(quantizer, shape, dtype, device) that directly allocates traceable quantized tensors. The fake impls now return actual tensors (which become FakeTensors under register_fake) instead of intermediate descriptors. The key insight: make_empty_traceable stashes _te_flat_names and _te_flat_ctx on the resulting tensor. Dynamo treats non-callable attributes on traceable wrapper subclasses as constant metadata, so forward_fn can read slot counts and reassembly info from these attributes without calling __tensor_flatten__ (which would cause a graph break since it returns non-Tensor Python objects). This eliminates: - TensorProto class and to_tensor_proto helper (tensor_proto.py deleted) - _proto_view (converted tensor fields to TensorProto before fake impls) - _tensor_field_names (identified fields for _proto_view) - _proto_slot_count / _proto_reassemble (operated on TensorProto objects) - TensorProto branch in _value_to_flat_tensors and _format_bwd_result The fake impls in linear.py now use: - isinstance(inp, QuantizedTensorStorage) instead of inp.is_quantized - weight._quantizer instead of weight.quantizer (TensorProto field) - make_empty_traceable(...) instead of TensorProto(...) - Direct set_usage on quantizer instead of proto.update_usage() Test Plan: ``` python -m pytest tests/pytorch/test_torch_compile.py -v -k 'not nvfp4' ``` Authored with Claude.

kshitij12345 · 2026-07-02T13:40:22Z

+    # TODO: understand why Dynamo does not recognize the quantizer retrieved via
+    # t._quantizer as the same value-opaque type it would if captured from a
+    # closure. If that is fixed upstream, the stashed attributes become
+    # unnecessary and we could compute slot counts directly from the quantizer.


PyTorch Issue: pytorch/pytorch#188796

kshitij12345 · 2026-07-02T15:01:06Z

        quantizer.optimize_for_gemm = self.optimize_for_gemm
-        quantizer.rht_matrix = self.rht_matrix
-        quantizer.rht_matrix_random_sign_mask_t = self.rht_matrix_random_sign_mask_t
+        if not torch.compiler.is_compiling():


TODO: Understand this better

This happens due to NVFPQuantizer being an OpaqueObject but then we try to attach a Tensor onto it.

…der is_compiling Under Dynamo tracing, rht_matrix is a FakeTensor attached to an opaque script object. Accessing it in copy() triggers SourcelessBuilder which cannot wrap FakeTensor, causing an InternalTorchDynamoError. The fake impl never runs real quantization, so rht_matrix is unnecessary during tracing. Guard the tensor field copies with torch.compiler.is_compiling() -- the matrix will be rebuilt lazily via _rebuild_derived_state if the quantizer is later used outside tracing. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

The C++ quantize kernel requires with_post_rht_amax=True when with_rht is enabled. The test factory was creating an NVFP4Quantizer with with_rht=True but with_post_rht_amax defaulting to False, causing 'Pre-RHT amax is not supported yet' at quantize time. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

kshitij12345 · 2026-07-14T13:58:14Z

cc: @shino16

shino16 · 2026-07-15T08:33:10Z

Interesting. @kshitij12345 Can I ask you for a bit more explanation on why you want to bring this change?

I think TensorProto will be useful as TE's FakeTensor, although the current functionality is limited. @pggPL What do you think?

kshitij12345 · 2026-07-15T09:10:35Z

Can I ask you for a bit more explanation on why you want to bring this change?

My main thought is that we should be able to rely on PyTorch (and FakeTensor's) to propagate the metadata without requiring a new abstraction of TensorProto. This in my opinion should help simplify readability as we don't need to learn/understand a new abstraction.

pggPL and others added 25 commits June 29, 2026 11:25

Reword opaque-flag comment: self-contained, no Linear reference

2c3c5df

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

598a07c

for more information, see https://pre-commit.ci (cherry picked from commit afe364b)

Add torch.compile custom-op path for Linear

9bd16fd

Squashed PR NVIDIA#9 (linear_compile) onto the rebased base. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> (cherry picked from commit 84dbc6b)

[PyTorch] custom_op: note pytorch/pytorch#187434 enables dropping Non…

554f5a8

…e sentinel Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

github-actions Bot added the community-contribution label Jul 2, 2026

kshitij12345 marked this pull request as draft July 2, 2026 10:35

kshitij12345 force-pushed the linear_compile_no_proto branch 2 times, most recently from bda5d6c to 8877da0 Compare July 2, 2026 10:42

kshitij12345 marked this pull request as ready for review July 2, 2026 10:53

kshitij12345 force-pushed the linear_compile_no_proto branch from 8877da0 to 3bd5ab5 Compare July 2, 2026 13:35

kshitij12345 force-pushed the linear_compile_no_proto branch from 3bd5ab5 to b24d259 Compare July 2, 2026 13:37

kshitij12345 commented Jul 2, 2026

View reviewed changes

kshitij12345 force-pushed the linear_compile_no_proto branch from f506cb1 to c33cd00 Compare July 2, 2026 15:17

pggPL force-pushed the linear_compile branch from 5052b1e to c6544d0 Compare August 5, 2026 15:56

Navigation Menu

[PyTorch][torch.compile] Replace TensorProto with make_empty_traceable - #13

[PyTorch][torch.compile] Replace TensorProto with make_empty_traceable#13
kshitij12345 wants to merge 28 commits into
pggPL:linear_compilefrom
kshitij12345:linear_compile_no_proto

kshitij12345 commented Jul 2, 2026 •

edited

Loading

Uh oh!

kshitij12345 Jul 2, 2026 •

edited

Loading

Uh oh!

kshitij12345 Jul 2, 2026

Uh oh!

kshitij12345 Jul 3, 2026 •

edited

Loading

Uh oh!

kshitij12345 commented Jul 14, 2026

Uh oh!

shino16 commented Jul 15, 2026 •

edited

Loading

Uh oh!

kshitij12345 commented Jul 15, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

kshitij12345 commented Jul 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Type of change

Changes

Checklist:

Uh oh!

kshitij12345 Jul 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kshitij12345 Jul 2, 2026

Choose a reason for hiding this comment

Uh oh!

kshitij12345 Jul 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kshitij12345 commented Jul 14, 2026

Uh oh!

shino16 commented Jul 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

kshitij12345 commented Jul 15, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

kshitij12345 commented Jul 2, 2026 •

edited

Loading

kshitij12345 Jul 2, 2026 •

edited

Loading

kshitij12345 Jul 3, 2026 •

edited

Loading

shino16 commented Jul 15, 2026 •

edited

Loading