Skip to content

add cuda device specification and precision, and tests #519

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Apr 27, 2024
Merged

add cuda device specification and precision, and tests #519

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
dcff176
Add tests for quantize json, add cuda device specification and precis…
Apr 27, 2024
a641b39
merge and disable mobile.json on stories
Apr 27, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 24 additions & 0 deletions .github/workflows/pull.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -266,6 +266,12 @@ jobs:
bash .ci/scripts/validate.sh "./checkpoints/${REPO_NAME}/model.pth" "cuda" "aoti-float32"
echo "::endgroup::"

echo "::group::Run inference with quantize file"
if [ $(uname -s) != Darwin ]; then
python3 generate.py --quantize config/data/cuda.json --checkpoint "./checkpoints/${REPO_NAME}/model.pth"
fi
echo "::endgroup::"

test-gpu-aoti-float16:
uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
name: test-gpu-aoti-float16 (${{ matrix.platform }}, ${{ matrix.model_name }})
Expand Down Expand Up @@ -307,6 +313,13 @@ jobs:
bash .ci/scripts/validate.sh "./checkpoints/${REPO_NAME}/model.pth" "cuda" "aoti-float16"
echo "::endgroup::"

echo "::group::Run inference with quantize file"
if [ $(uname -s) == Darwin ]; then
python3 export.py --output-dso-path /tmp/model.so --quantize config/data/cuda.json --checkpoint "./checkpoints/${REPO_NAME}/model.pth"
python3 generate.py --dso-path /tmp/model.so --checkpoint "./checkpoints/${REPO_NAME}/model.pth"~
fi
echo "::endgroup::"

test-gpu-eval-sanity-check:
uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
name: test-gpu-eval-sanity-check (${{ matrix.platform }}, ${{ matrix.model_name }})
Expand Down Expand Up @@ -428,9 +441,20 @@ jobs:
export MODEL_PATH=checkpoints/stories15M/stories15M.pt
export MODEL_NAME=stories15M
export MODEL_DIR=/tmp

echo "******************************************"
echo "*** vanilla ***"
echo "******************************************"
python export.py --checkpoint-path ${MODEL_PATH} --output-pte-path ${MODEL_DIR}/${MODEL_NAME}.pte
python3 torchchat.py generate --checkpoint-path ${MODEL_PATH} --temperature 0 --pte-path ${MODEL_DIR}/${MODEL_NAME}.pte

echo "******************************************"
echo "*** --quantize config/data/mobile.json ***"
echo "******************************************"
# python export.py --quantize config/data/mobile.json --checkpoint-path ${MODEL_PATH} --output-pte-path ${MODEL_DIR}/${MODEL_NAME}.pte
# python3 torchchat.py generate --checkpoint-path ${MODEL_PATH} --temperature 0 --pte-path ${MODEL_DIR}/${MODEL_NAME}.pte


echo "******************************************"
echo "******* Emb: channel-wise quantized ******"
echo "******************************************"
Expand Down
2 changes: 2 additions & 0 deletions config/data/cuda.json
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
{
"executor": {"accelerator": "cuda"},
"precision": {"dtype": "bf16"},
"linear:int4": {"groupsize" : 256}
}
3 changes: 3 additions & 0 deletions quantize.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
import torch.nn.functional as F
from build.utils import (
find_multiple,
get_device_str,
get_precision,
name_to_dtype,
state_dict_device,
Expand Down Expand Up @@ -124,6 +125,8 @@ def quantized_model(self) -> nn.Module:

#########################################################################
### wrapper for setting device as a QuantHandler ###
### for onw select device for PyTorch eager and AOTI, in future ###
### also use this for selecting delegate when exporting with ET ###


class ExecutorHandler(QuantHandler):
Expand Down