Remove fp8_out from the LN API - #8
Merged
Merged
Conversation
Member
Author
|
/blossom-ci |
1 similar comment
Member
Author
|
/blossom-ci |
ptrendx
force-pushed
the
pr_rm_fp8_out
branch
from
October 10, 2022 22:45
b1bc0b4 to
b533ab7
Compare
ptrendx
marked this pull request as ready for review
October 10, 2022 22:47
Signed-off-by: Przemyslaw Tredak <ptredak@nvidia.com>
Signed-off-by: Przemyslaw Tredak <ptredak@nvidia.com>
ptrendx
force-pushed
the
pr_rm_fp8_out
branch
from
October 10, 2022 22:47
b533ab7 to
f122b0d
Compare
ptrendx
pushed a commit
that referenced
this pull request
Aug 5, 2026
* 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>
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
2 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.
fp8_out argument is not needed since the output Tensor already has the proper type set.