Skip to content

[Feature] Make benchmarked losses compatible with torch.compile #2405

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 13 commits into from
Sep 17, 2024
Merged

[Feature] Make benchmarked losses compatible with torch.compile #2405

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
f6865af
Update
vmoens Aug 30, 2024
2548a64
Update
vmoens Sep 2, 2024
e60bc17
Update
vmoens Sep 2, 2024
ef7376b
Update
vmoens Sep 2, 2024
09d2372
Update
vmoens Sep 4, 2024
c35f722
Update
vmoens Sep 12, 2024
f9aa095
Update
vmoens Sep 12, 2024
98e3914
Update
vmoens Sep 13, 2024
ba42d13
Update
vmoens Sep 13, 2024
56c93cb
Update
vmoens Sep 13, 2024
771217c
Update
vmoens Sep 16, 2024
0e30b4b
Update
vmoens Sep 17, 2024
e4441a0
Update
vmoens Sep 17, 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
6 changes: 3 additions & 3 deletions .github/unittest/linux/scripts/run_all.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -127,13 +127,13 @@ if [[ "$TORCH_VERSION" == "nightly" ]]; then
if [ "${CU_VERSION:-}" == cpu ] ; then
pip3 install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cpu -U
else
pip3 install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/$CU_VERSION
pip3 install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/$CU_VERSION -U
fi
elif [[ "$TORCH_VERSION" == "stable" ]]; then
if [ "${CU_VERSION:-}" == cpu ] ; then
pip3 install torch torchvision --index-url https://download.pytorch.org/whl/cpu
pip3 install torch torchvision --index-url https://download.pytorch.org/whl/cpu -U
else
pip3 install torch torchvision --index-url https://download.pytorch.org/whl/$CU_VERSION
pip3 install torch torchvision --index-url https://download.pytorch.org/whl/$CU_VERSION -U
fi
else
printf "Failed to install pytorch"
Expand Down
8 changes: 4 additions & 4 deletions .github/unittest/linux_examples/scripts/run_all.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -150,15 +150,15 @@ git submodule sync && git submodule update --init --recursive
printf "Installing PyTorch with %s\n" "${CU_VERSION}"
if [[ "$TORCH_VERSION" == "nightly" ]]; then
if [ "${CU_VERSION:-}" == cpu ] ; then
pip3 install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cpu -U
pip3 install --pre torch torchvision numpy==1.26.4 --index-url https://download.pytorch.org/whl/nightly/cpu -U
else
pip3 install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/$CU_VERSION
pip3 install --pre torch torchvision numpy==1.26.4 --index-url https://download.pytorch.org/whl/nightly/$CU_VERSION
fi
elif [[ "$TORCH_VERSION" == "stable" ]]; then
if [ "${CU_VERSION:-}" == cpu ] ; then
pip3 install torch torchvision --index-url https://download.pytorch.org/whl/cpu
pip3 install torch torchvision numpy==1.26.4 --index-url https://download.pytorch.org/whl/cpu
else
pip3 install torch torchvision --index-url https://download.pytorch.org/whl/$CU_VERSION
pip3 install torch torchvision numpy==1.26.4 --index-url https://download.pytorch.org/whl/$CU_VERSION
fi
else
printf "Failed to install pytorch"
Expand Down
2 changes: 1 addition & 1 deletion .github/unittest/linux_libs/scripts_brax/install.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ if [[ "$TORCH_VERSION" == "nightly" ]]; then
fi
elif [[ "$TORCH_VERSION" == "stable" ]]; then
if [ "${CU_VERSION:-}" == cpu ] ; then
pip3 install torch --index-url https://download.pytorch.org/whl/cpu
pip3 install torch --index-url https://download.pytorch.org/whl/cpu -U
else
pip3 install torch --index-url https://download.pytorch.org/whl/cu121
fi
Expand Down
4 changes: 2 additions & 2 deletions .github/unittest/linux_libs/scripts_openx/install.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,9 +37,9 @@ if [[ "$TORCH_VERSION" == "nightly" ]]; then
fi
elif [[ "$TORCH_VERSION" == "stable" ]]; then
if [ "${CU_VERSION:-}" == cpu ] ; then
pip3 install torch torchvision --index-url https://download.pytorch.org/whl/cpu
pip3 install torch torchvision --index-url https://download.pytorch.org/whl/cpu -U
else
pip3 install torch torchvision --index-url https://download.pytorch.org/whl/cu121
pip3 install torch torchvision --index-url https://download.pytorch.org/whl/cu121 -U
fi
else
printf "Failed to install pytorch"
Expand Down
8 changes: 4 additions & 4 deletions .github/unittest/linux_libs/scripts_rlhf/install.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,15 +31,15 @@ git submodule sync && git submodule update --init --recursive
printf "Installing PyTorch with cu121"
if [[ "$TORCH_VERSION" == "nightly" ]]; then
if [ "${CU_VERSION:-}" == cpu ] ; then
pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cpu -U
pip3 install --pre torch numpy==1.26.4 --index-url https://download.pytorch.org/whl/nightly/cpu -U
else
pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu121 -U
pip3 install --pre torch numpy==1.26.4 --index-url https://download.pytorch.org/whl/nightly/cu121 -U
fi
elif [[ "$TORCH_VERSION" == "stable" ]]; then
if [ "${CU_VERSION:-}" == cpu ] ; then
pip3 install torch --index-url https://download.pytorch.org/whl/cpu
pip3 install torch numpy==1.26.4 --index-url https://download.pytorch.org/whl/cpu
else
pip3 install torch --index-url https://download.pytorch.org/whl/cu121
pip3 install torch numpy==1.26.4 --index-url https://download.pytorch.org/whl/cu121
fi
else
printf "Failed to install pytorch"
Expand Down
2 changes: 1 addition & 1 deletion .github/unittest/linux_libs/scripts_vd4rl/install.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,7 @@ if [[ "$TORCH_VERSION" == "nightly" ]]; then
fi
elif [[ "$TORCH_VERSION" == "stable" ]]; then
if [ "${CU_VERSION:-}" == cpu ] ; then
pip3 install torch torchvision --index-url https://download.pytorch.org/whl/cpu
pip3 install torch torchvision --index-url https://download.pytorch.org/whl/cpu -U
else
pip3 install torch torchvision --index-url https://download.pytorch.org/whl/cu121
fi
Expand Down
2 changes: 1 addition & 1 deletion .github/unittest/linux_olddeps/scripts_gym_0_13/install.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@ printf "Installing PyTorch with %s\n" "${CU_VERSION}"
if [ "${CU_VERSION:-}" == cpu ] ; then
conda install pytorch==2.0 torchvision==0.15 cpuonly -c pytorch -y
else
conda install pytorch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 pytorch-cuda=11.8 numpy==1.26 numpy-base==1.26 -c pytorch -c nvidia -y
conda install pytorch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 pytorch-cuda=11.8 numpy==1.26 -c pytorch -c nvidia -y
fi

# Solving circular import: https://stackoverflow.com/questions/75501048/how-to-fix-attributeerror-partially-initialized-module-charset-normalizer-has
Expand Down
2 changes: 1 addition & 1 deletion .github/unittest/linux_optdeps/scripts/install.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ version="$(python -c "print('.'.join(\"${CUDA_VERSION}\".split('.')[:2]))")"
git submodule sync && git submodule update --init --recursive

printf "Installing PyTorch with %s\n" "${CU_VERSION}"
pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/$CU_VERSION
pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/$CU_VERSION -U

# install tensordict
if [[ "$RELEASE" == 0 ]]; then
Expand Down
2 changes: 2 additions & 0 deletions .github/workflows/benchmarks.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@ jobs:
- name: Run benchmarks
run: |
cd benchmarks/
export TORCHDYNAMO_INLINE_INBUILT_NN_MODULES=1
python -m pytest --benchmark-json output.json
- name: Store benchmark results
if: ${{ github.ref == 'refs/heads/main' || github.event_name == 'workflow_dispatch' }}
Expand Down Expand Up @@ -107,6 +108,7 @@ jobs:
- name: Run benchmarks
run: |
cd benchmarks/
export TORCHDYNAMO_INLINE_INBUILT_NN_MODULES=1
python3 -m pytest --benchmark-json output.json
- name: Store benchmark results
uses: benchmark-action/github-action-benchmark@v1
Expand Down
2 changes: 2 additions & 0 deletions .github/workflows/benchmarks_pr.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@ jobs:
- name: Run benchmarks
run: |
cd benchmarks/
export TORCHDYNAMO_INLINE_INBUILT_NN_MODULES=1
RUN_BENCHMARK="pytest --rank 0 --benchmark-json "
git checkout ${{ github.event.pull_request.base.sha }}
$RUN_BENCHMARK ${{ env.BASELINE_JSON }}
Expand Down Expand Up @@ -125,6 +126,7 @@ jobs:
- name: Run benchmarks
run: |
cd benchmarks/
export TORCHDYNAMO_INLINE_INBUILT_NN_MODULES=1
RUN_BENCHMARK="pytest --rank 0 --benchmark-json "
git checkout ${{ github.event.pull_request.base.sha }}
$RUN_BENCHMARK ${{ env.BASELINE_JSON }}
Expand Down
Loading
Loading