Skip to content

Add Float8Tensor#2463

Merged
jerryzh168 merged 1 commit into
mainfrom
jerryzh168/stack/9
Aug 6, 2025
Merged

Add Float8Tensor#2463
jerryzh168 merged 1 commit into
mainfrom
jerryzh168/stack/9

Conversation

@jerryzh168

@jerryzh168 jerryzh168 commented Jun 30, 2025 •
edited
Loading

Copy link
Copy Markdown
Contributor

Stacked PRs:

Add Float8Tensor

Summary:

  • Added Float8Tensor that's using fbgemm kernels and scaled_mm:
    • per row activation + per row weight linear calling torch._scaled_mm op (for compatibilty with SM 8.9)
    • per tensor activation + per tensor weight quant linear calling torch._scaled_mm op (for compatibilty with SM 8.9)
    • per row activation + per row weight bmm calling torch.ops.fbgemm.f8f8bf16_rowwise_batched kernel (only works for SM 9.0+) can use batched scaled mm from torch when it's supported: [RFC]: PyTorch Low-Precision GEMMs Public API pytorch#157950
  • dynamic quantization kwargs is added to the Float8Tensor directly
  • Added QuantizeTensorKwargs and QuantizeTensorToFloat8Kwargs to store key word args for Float8Tensor.to_float8
  • Updated Float8DynamicActivationFloat8WeightConfig and Float8WeightOnlyConfig to use Float8Tensor

Test Plan:
python test/dtypes/test_affine_quantized_float.py
python test/quantization/quantize_/workflows/float8/test_float8_tensor.py

Reviewers:

Subscribers:

Tasks:

Tags:

@pytorch-bot

pytorch-bot Bot commented Jun 30, 2025 •
edited
Loading

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/2463

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

❌ 1 New Failure

As of commit b0c2cf3 with merge base b757fb9 (image):

NEW FAILURE - The following job has failed:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

jerryzh168 added a commit that referenced this pull request Jun 30, 2025
@jerryzh168
Add Float8RowwiseTensor
5cae4d0 
Summary:
Splits out the float8 rowwise quantized path (both act and weight) of AQT to Float8RowwiseTensor

Next: could potentially incorporate the per tensor activation path there as well
Next: we can split the per tensor weight path to another Tensor as well, so we can deprecate AQT path for float8

Test Plan:
python test/dtypes/test_affine_quantized_float.py
python test/quantization/quantize_/test_float8_rowwise_tensor.py

Reviewers:

Subscribers:

Tasks:

Tags:

stack-info: PR: #2463, branch: jerryzh168/stack/9
@jerryzh168 jerryzh168 force-pushed the jerryzh168/stack/9 branch from da79207 to 5cae4d0 Compare June 30, 2025 23:01
This was referenced Jun 30, 2025
Merged
Merged
@facebook-github-bot facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jun 30, 2025
@jerryzh168 jerryzh168 mentioned this pull request Jun 30, 2025
Merged
@jerryzh168 jerryzh168 added the topic: new feature Use this tag if this PR adds a new feature label Jun 30, 2025
@jerryzh168 jerryzh168 changed the base branch from jerryzh168/stack/4 to main July 2, 2025 01:58
@jerryzh168 jerryzh168 force-pushed the jerryzh168/stack/9 branch from 5cae4d0 to 33ca58e Compare July 2, 2025 01:58
@jerryzh168 jerryzh168 changed the title Add Float8RowwiseTensor Add Float8Tensor Jul 2, 2025
@jerryzh168 jerryzh168 changed the base branch from main to jerryzh168/stack/4 July 2, 2025 01:58
@jerryzh168 jerryzh168 mentioned this pull request Jul 2, 2025
Merged
vkuzo
vkuzo reviewed Jul 2, 2025
View reviewed changes
Comment thread torchao/quantization/quantize_/float8_tensor.py Outdated
vkuzo
vkuzo reviewed Jul 2, 2025
View reviewed changes
Comment thread torchao/quantization/quant_api.py Outdated
vkuzo
vkuzo reviewed Jul 2, 2025
View reviewed changes
Comment thread torchao/quantization/quantize_/float8_tensor.py Outdated
@jerryzh168 jerryzh168 changed the base branch from jerryzh168/stack/4 to main July 2, 2025 20:35
@jerryzh168 jerryzh168 force-pushed the jerryzh168/stack/9 branch from 33ca58e to 897ec7e Compare July 2, 2025 20:36
@jerryzh168 jerryzh168 changed the base branch from main to jerryzh168/stack/4 July 2, 2025 20:36
@jerryzh168 jerryzh168 changed the base branch from jerryzh168/stack/4 to main July 2, 2025 21:42
@jerryzh168 jerryzh168 force-pushed the jerryzh168/stack/9 branch 2 times, most recently from 7897dcf to 99a1bb1 Compare July 2, 2025 21:42
@jerryzh168 jerryzh168 mentioned this pull request Jul 2, 2025
Merged
@jerryzh168 jerryzh168 changed the base branch from main to jerryzh168/stack/11 July 2, 2025 21:42
@jerryzh168 jerryzh168 changed the base branch from jerryzh168/stack/11 to main July 2, 2025 23:44
@jerryzh168 jerryzh168 force-pushed the jerryzh168/stack/9 branch from 99a1bb1 to 7e9f224 Compare July 2, 2025 23:44
@jerryzh168 jerryzh168 changed the base branch from main to jerryzh168/stack/11 July 2, 2025 23:44
@jerryzh168 jerryzh168 changed the base branch from jerryzh168/stack/11 to main July 3, 2025 00:09
@jerryzh168 jerryzh168 force-pushed the jerryzh168/stack/9 branch from 7e9f224 to 442bd6c Compare July 3, 2025 00:09
@jerryzh168 jerryzh168 changed the base branch from main to jerryzh168/stack/11 July 3, 2025 00:09
vkuzo
vkuzo reviewed Jul 25, 2025
View reviewed changes
Comment thread torchao/quantization/quantize_/common/general_kernel_preference.py Outdated
vkuzo
vkuzo reviewed Jul 25, 2025
View reviewed changes
Comment thread torchao/quantization/quantize_/workflows/_utils.py Outdated
from torchao.quantization.quantize_.common import QuantizeTensorKwargs


def _choose_quant_func_and_quantize_tensor(

@vkuzo vkuzo Jul 25, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: colocate with QuantizeTensorKwargs and move out from utils, since this is more of a developer facing function which is central to the design?

@jerryzh168 jerryzh168 Jul 25, 2025 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

hmmm, I have thought about move this to common/ but then common is going to import from the workflow specific code, this seems weird, like:

from torchao.quantization.quantize_.workflows import (
        Float8Tensor,
        QuantizeTensorToFloat8Kwargs,
    )

any ideas to resolve that? or do you feel it's fine for functions in common to import from workflow?

huydhn
huydhn reviewed Jul 25, 2025
View reviewed changes
Comment thread .github/workflows/lib_integration_test.yml Outdated
git submodule update --init --recursive
python use_existing_torch.py
pip install -r requirements/build.txt
pip install --no-build-isolation -e .

@huydhn huydhn Jul 25, 2025 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@jerryzh168 This setup, unfortunately, destroys our H100 fleet because:

  1. It builds vLLM from scratch without any cache and that could take hours. So, it's a very inefficient use of H100 runners
  2. The build somehow hang Docker daemon, maybe memory related. At least, you want to set MAX_JOBS here to match what vLLM does https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile#L198-L202

As it's not trivial to build vLLM from source locally, I think we should explore the option to use pre-build image from vLLM where you can pull them from public.ecr.aws/q9t5s3a7/vllm-ci-postmerge-repo:<COMMIT_HASH>. Here is an example https://github.com/pytorch/pytorch-integration-testing/blob/main/.github/workflows/vllm-benchmark.yml#L149

Another way is to run this test on PyTorch CI where @yangw-dev are working on a CI job to test vLLM main v.s. PyTorch (we also need to build vLLM there, but with caching properly setup). We could also run pytest test/integration from AO there I think

@jerryzh168 jerryzh168 Jul 25, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I see thanks, will try out the suggestions in #2601

@huydhn huydhn mentioned this pull request Jul 26, 2025
Closed
vkuzo
vkuzo reviewed Jul 28, 2025
View reviewed changes
Comment thread torchao/quantization/quantize_/common/general_kernel_preference.py Outdated
vkuzo
vkuzo reviewed Jul 28, 2025
View reviewed changes
Comment thread torchao/quantization/quantize_/common/general_kernel_preference.py Outdated
vkuzo
vkuzo reviewed Jul 28, 2025
View reviewed changes
Comment thread torchao/quantization/quantize_/common/general_kernel_preference.py Outdated
vkuzo
vkuzo reviewed Jul 29, 2025
View reviewed changes
Comment thread torchao/quantization/quantize_/workflows/float8/float8_tensor.py Outdated
vkuzo
vkuzo reviewed Jul 29, 2025
View reviewed changes
Comment thread torchao/quantization/quantize_/workflows/float8/float8_tensor.py Outdated
vkuzo
vkuzo reviewed Jul 29, 2025
View reviewed changes
Comment thread torchao/quantization/quantize_/workflows/float8/float8_tensor.py Outdated
vkuzo
vkuzo reviewed Jul 29, 2025
View reviewed changes
Comment thread torchao/quantization/quantize_/workflows/float8/float8_tensor.py Outdated
vkuzo
vkuzo reviewed Jul 29, 2025
View reviewed changes
Comment thread torchao/quantization/quantize_/workflows/float8/float8_tensor.py Outdated
vkuzo
vkuzo reviewed Jul 29, 2025
View reviewed changes
Comment thread torchao/quantization/quant_api.py Outdated
vkuzo
vkuzo reviewed Jul 29, 2025
View reviewed changes
Comment thread test/dtypes/test_affine_quantized_float.py Outdated
vkuzo
vkuzo reviewed Jul 29, 2025
View reviewed changes
Comment thread test/float8/test_base.py Outdated
vkuzo
vkuzo reviewed Jul 29, 2025
View reviewed changes
Comment thread torchao/quantization/quantize_/common/general_kernel_preference.py Outdated
vkuzo
vkuzo reviewed Jul 31, 2025
View reviewed changes
Comment thread test/quantization/quantize_/workflows/float8/test_float8_tensor.py Outdated
vkuzo
vkuzo reviewed Jul 31, 2025
View reviewed changes
Comment thread torchao/quantization/quant_api.py
vkuzo
vkuzo reviewed Jul 31, 2025
View reviewed changes
Comment thread torchao/quantization/quantize_/common/kernel_preference.py Outdated
vkuzo
vkuzo reviewed Jul 31, 2025
View reviewed changes
Comment thread torchao/quantization/quantize_/workflows/float8/float8_tensor.py
vkuzo
vkuzo reviewed Jul 31, 2025
View reviewed changes
Comment thread test/quantization/quantize_/workflows/float8/test_float8_tensor.py
vkuzo
vkuzo reviewed Jul 31, 2025
View reviewed changes
Comment thread torchao/quantization/quantize_/workflows/float8/float8_tensor.py
This was referenced Jul 31, 2025
Open
Merged
Merged
Merged
@jerryzh168
Add Float8Tensor
b0c2cf3 
Summary:
* Added Float8Tensor that's using fbgemm kernels and scaled_mm:
    * per row activation + per row weight linear calling torch._scaled_mm op (for compatibilty with SM 8.9)
    * per tensor activation + per tensor weight quant linear calling torch._scaled_mm op (for compatibilty with SM 8.9)
    * per row activation + per row weight bmm calling torch.ops.fbgemm.f8f8bf16_rowwise_batched kernel (only works for SM 9.0+) can use batched scaled mm from torch when it's supported: pytorch/pytorch#157950
* dynamic quantization kwargs is added to the Float8Tensor directly
* Added QuantizeTensorKwargs and QuantizeTensorToFloat8Kwargs to store key word args for Float8Tensor.to_float8
* Updated Float8DynamicActivationFloat8WeightConfig and Float8WeightOnlyConfig to use Float8Tensor

Test Plan:
python test/dtypes/test_affine_quantized_float.py
python test/quantization/quantize_/workflows/float8/test_float8_tensor.py

Reviewers:

Subscribers:

Tasks:

Tags:

stack-info: PR: #2463, branch: jerryzh168/stack/9
@clee2000

clee2000 commented Aug 6, 2025

Copy link
Copy Markdown
Contributor

/easycla

@jerryzh168 jerryzh168 mentioned this pull request Aug 8, 2025
Merged
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@huydhn huydhn huydhn left review comments

@vkuzo vkuzo vkuzo approved these changes

@drisspg drisspg Awaiting requested review from drisspg

Assignees

No one assigned

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. topic: new feature Use this tag if this PR adds a new feature

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

6 participants

@jerryzh168 @clee2000 @huydhn @vkuzo @drisspg @facebook-github-bot