Skip to content

Commit

Permalink
[Feature] Make benchmarked losses compatible with torch.compile
Browse files Browse the repository at this point in the history 
ghstack-source-id: 825ded593dcffcecf626a705d8b7c0c5e0839719
Pull Request resolved: #2405
  • Loading branch information
@vmoens
vmoens committed Sep 17, 2024
1 parent 605b4aa commit 2332909
Show file tree
Hide file tree
Showing 26 changed files with 940 additions and 277 deletions.
6 changes: 3 additions & 3 deletions .github/unittest/linux/scripts/run_all.sh
View file
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
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
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
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
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
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
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
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
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
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

1 comment on commit 2332909

@github-actions
Copy link

Choose a reason for hiding this comment

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

⚠️ Performance Alert ⚠️

Possible performance regression was detected for benchmark 'CPU Benchmark Results'.
Benchmark result of this commit is worse than the previous benchmark result exceeding threshold 2.

Benchmark suite Current: 2332909 Previous: 605b4aa Ratio
benchmarks/test_replaybuffer_benchmark.py::test_rb_populate[TensorDictReplayBuffer-ListStorage-SamplerWithoutReplacement-400] 30.183476406713112 iter/sec (stddev: 0.20221241229003356) 157.5775708449948 iter/sec (stddev: 0.016796232595089552) 5.22

This comment was automatically generated by workflow using github-action-benchmark.

CC: @vmoens

Please sign in to comment.