Skip to content

Conversation

jerryzh168
Copy link
Contributor

@jerryzh168 jerryzh168 commented May 23, 2025

Summary:
we also plan to expose some other kernels like fp8xint4 and bf16xfp8, fp8xfp8 to compare with existing torchao kernels

Test Plan:
test/dtypes/test_fbgemm_int4_tensor.py

H100, with compile:

overall tokens/sec TTFT Peak Memory Model Size
baseline - 1 131.65 0.0220 16.24 GB 15.01 GB
baseline - 128 76.38 0.0544 26.92 GB 15.01 GB
int4wo - 1 207.69 0.0288 6.41 GB 3.99 GB
int4wo - 128 12.85 0.4223 16.01 GB 3.99 GB
fbgemm-int4 - 1 (no compile) 40.00 0.0508 29.03 GB 4.22 GB
fbgemm-int4 - 128 (no compile) 11.46 0.0846 28.96 GB 4.22 GB
export CHECKPOINT_PATH=../../../checkpoints # path to checkpoints folder
export MODEL_REPO=meta-llama/Meta-Llama-3.1-8B-Instruct
# default batch size 1
python generate.py --checkpoint_path $CHECKPOINT_PATH/$MODEL_REPO/model.pth --compile --write_result benchmark_results.txt
python generate.py --checkpoint_path $CHECKPOINT_PATH/$MODEL_REPO/model.pth --compile --quantization int4wo-128 --write_result benchmark_results.txt
python generate.py --checkpoint_path $CHECKPOINT_PATH/$MODEL_REPO/model.pth --quantization fbgemm-int4-128 --write_result benchmark_results.txt

python generate.py --checkpoint_path $CHECKPOINT_PATH/$MODEL_REPO/model.pth --compile --write_result benchmark_results.txt --batch_size 128
python generate.py --checkpoint_path $CHECKPOINT_PATH/$MODEL_REPO/model.pth --compile --quantization int4wo-128 --write_result benchmark_results.txt --batch_size 128
python generate.py --checkpoint_path $CHECKPOINT_PATH/$MODEL_REPO/model.pth --quantization fbgemm-int4-128 --write_result benchmark_results.txt --batch_size 128

Note: fbgemm-int4-128 does not work with compile yet since the fbgemm op does not have meta device implementation.

Reviewers:

Subscribers:

Tasks:

Tags:

Copy link

pytorch-bot bot commented May 23, 2025

🔗 Helpful Links

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

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

✅ No Failures

As of commit d2066dc with merge base b0cfeec (image):
💚 Looks good so far! There are no failures yet. 💚

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

@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 May 23, 2025
@samanamp
Copy link

Thank you! community really needs this.

@jerryzh168 jerryzh168 added the topic: new feature Use this tag if this PR adds a new feature label May 27, 2025
@jerryzh168 jerryzh168 requested a review from drisspg May 27, 2025 20:57
@jerryzh168 jerryzh168 requested a review from drisspg May 28, 2025 00:07
Copy link
Contributor

@drisspg drisspg left a comment

Choose a reason for hiding this comment

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

Okay everything looks pretty good but the API for the FBGEMM config feels gross imo I know its a thin wrapper around their op but I think we can do better than io string

Copy link
Contributor

@drisspg drisspg left a comment

Choose a reason for hiding this comment

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

Looks good, can you also add a serialization test entry, want to ensure we can seralize str eums

Summary:
we also plan to expose some other kernels like fp8xint4
and bf16xfp8, fp8xfp8 to compare with existing torchao kernels

Test Plan:
test/dtypes/test_fbgemm_int4_tensor.py

Reviewers:

Subscribers:

Tasks:

Tags:
@jerryzh168 jerryzh168 merged commit 0b33f12 into pytorch:main May 28, 2025
19 checks passed
liangel-02 pushed a commit that referenced this pull request Aug 25, 2025
* Add support for fbgemm int4 mm kernel

Summary:
we also plan to expose some other kernels like fp8xint4
and bf16xfp8, fp8xfp8 to compare with existing torchao kernels

Test Plan:
test/dtypes/test_fbgemm_int4_tensor.py

Reviewers:

Subscribers:

Tasks:

Tags:

* fix and test

* fix dtype

* use importlib

* add links to fbgemm code

* update io_dtype type

* renaming

* remove enum

* serializability update

* format

* fix tests

* skip fbgemm config tests for 2.5 and below
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

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

Development

Successfully merging this pull request may close these issues.

4 participants