Skip to content
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.

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

Triton #6798

Merged
merged 67 commits into from
Jun 7, 2024
Merged

Triton #6798

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
67 commits
Select commit Hold shift + click to select a range
6dccf0a
Update infra_triggers.tf
ManfeiBai Oct 4, 2023
9828123
Skeleton trition support
bhavya01 Mar 20, 2024
99bf48d
Merge branch 'master' into triton
bhavya01 Mar 20, 2024
b89e558
Fix bugs
bhavya01 Mar 21, 2024
64189bd
Fix custom call invocation
bhavya01 Mar 21, 2024
0c208ef
Refactor to include gpu custom call and create triton dir
bhavya01 Mar 22, 2024
b553ba7
Lint fixes
bhavya01 Mar 22, 2024
c5129e6
python lint fix
bhavya01 Mar 22, 2024
48e7127
Updated base image for CI
bhavya01 Mar 27, 2024
e04fc97
Update github workflow gcr image
bhavya01 Mar 28, 2024
37bf127
Merge branch 'master' into custom
bhavya01 Mar 28, 2024
6061895
Remove xrt build and test file
bhavya01 Mar 28, 2024
f59ddbf
Add temporary test to run triton kernel
bhavya01 Mar 28, 2024
158aed4
Fix tests
bhavya01 Mar 28, 2024
87b92c5
Update payload for xla gpu custom call
bhavya01 Mar 29, 2024
847ccc5
Update gpu runner
bhavya01 Mar 29, 2024
eca6d52
Merge branch 'master' into triton
bhavya01 Apr 4, 2024
2348ca3
Extract payload from triton kernel programatically
bhavya01 Apr 12, 2024
110c8c6
Merge branch 'master' into triton
bhavya01 Apr 12, 2024
a226150
Lint fixes
bhavya01 Apr 12, 2024
4c1f4f5
Only build triton files for GPU
bhavya01 Apr 12, 2024
431f822
build pytorch for ampere gpus
bhavya01 Apr 13, 2024
4bade16
c++ lint fix
bhavya01 Apr 13, 2024
1c5b47d
Python lint fix
bhavya01 Apr 13, 2024
3138a92
Fix torch cuda arch list
bhavya01 Apr 13, 2024
3f00cfd
Use a bigger machine for CI build
bhavya01 Apr 13, 2024
e729cfb
Add triton test to run_tests.sh
bhavya01 Apr 13, 2024
8e304c0
Update triton env variable
bhavya01 Apr 15, 2024
27bdc3a
Set up a separate CI for triton tests
bhavya01 Apr 15, 2024
9a3ef84
Fix github workflow to add _triton.yml
bhavya01 Apr 15, 2024
ade444d
Rebuild torch xla for triton tests
bhavya01 Apr 15, 2024
cb0bb85
Create a separate CI tab for triton tests
bhavya01 Apr 16, 2024
015b1ad
Separate build and test phase for triton
bhavya01 Apr 16, 2024
a18028a
Fix flags for docker run container
bhavya01 Apr 16, 2024
993ee92
Update triton.yml to output docker image
bhavya01 Apr 16, 2024
a87b782
Add a python binding to register custom calls and remove jax files
bhavya01 May 10, 2024
bf05d1b
Fix lint
bhavya01 May 10, 2024
4582fe8
Merge main
bhavya01 May 10, 2024
9680167
Merge master
bhavya01 May 10, 2024
a7b94c6
Merge master after updating
bhavya01 May 10, 2024
e14636a
Update CI to use cuda plugin
bhavya01 May 10, 2024
256d819
Install jaxlib while setting up triton tests
bhavya01 May 10, 2024
c616e64
Install triton package while running triton tests
bhavya01 May 10, 2024
60b8d18
Experimental: Build pytorch with cuda
bhavya01 May 13, 2024
2bde624
Revert build pytorch with CUDA
bhavya01 May 14, 2024
e6c4e0a
Merge branch 'master' into triton
bhavya01 May 14, 2024
14ee545
Remove ansible path for triton CI
bhavya01 May 14, 2024
25acb26
Style fixes
bhavya01 May 20, 2024
6b0ac18
[Experimental] test new CI
bhavya01 May 28, 2024
4d97150
[Experimental] Set XLA_CUDA=0 for cuda arch in ansible
bhavya01 May 28, 2024
e079049
[Experimental] Update CI to build pytorch cuda with ansible
bhavya01 May 29, 2024
d9c89b6
Update CI
bhavya01 May 30, 2024
7a6c809
Fix CI workflow file
bhavya01 May 30, 2024
6b1954d
Fix CI workflow
bhavya01 May 30, 2024
21797a6
Fix the wheels installed for tests requiring torch cuda
bhavya01 May 30, 2024
e6e89d3
Add compute_capability=8.6 for xla cuda plugin
bhavya01 May 31, 2024
ac45fe1
update TORCH_CUDA_ARCH_LIST
bhavya01 May 31, 2024
f828fbb
Experimental build torch and torch_xla cuda wheels
bhavya01 May 31, 2024
ac56c00
Merge branch 'master' into triton
bhavya01 May 31, 2024
c3b8653
Update build_and_test.yml
bhavya01 May 31, 2024
a1168c6
Update dlpack test to only use one device
bhavya01 May 31, 2024
39551a2
Remove compute capability 8.6 from cuda plugin
bhavya01 May 31, 2024
35e0869
Remove triton.sh
bhavya01 May 31, 2024
f95d898
Default empty torch_cuda_arch_list in ansible config
bhavya01 May 31, 2024
291104d
Merge branch 'master' into triton
bhavya01 Jun 5, 2024
f5c9b1a
Revert CI changes
bhavya01 Jun 6, 2024
5b23969
Revert CI changes pt2
bhavya01 Jun 6, 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
Prev Previous commit
Next Next commit
Install jaxlib while setting up triton tests
  • Loading branch information
bhavya01 committed May 10, 2024
commit 256d8194be060d68ac37d69e4700a1eb240e3ebf
4 changes: 3 additions & 1 deletion .circleci/triton.sh
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,9 @@ fi
apply_patches

python -c "import fcntl; fcntl.fcntl(1, fcntl.F_SETFL, 0)"
pip install --upgrade "jax[cuda12]"
pip install -U --pre jaxlib -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html
pip install -U --pre jax-cuda12-pjrt jax-cuda12-plugin -f https://storage.googleapis.com/jax-releases/jax_cuda_plugin_nightly_releases.html
pip install -U --pre jax -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html

export PATH=$PATH:/usr/local/cuda-12.1/bin
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda-12.1/lib64
Expand Down
12 changes: 10 additions & 2 deletions .github/workflows/_triton.yml
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
name: xla-test
name: triton-test
on:
workflow_call:
inputs:
Expand Down Expand Up @@ -34,6 +34,12 @@ jobs:
container:
image: ${{ inputs.dev-image }}
options: "${{ inputs.install-cuda-plugin && '--gpus all' || '' }} --shm-size 16g"
timeout-minutes: ${{ inputs.timeout-minutes }}
env:
GCLOUD_SERVICE_KEY: ${{ secrets.gcloud-service-key }}
GOOGLE_APPLICATION_CREDENTIALS: /tmp/default_credentials.json
BAZEL_JOBS: 16
BAZEL_REMOTE_CACHE: 1
steps:
# See https://github.com/actions/checkout/issues/1014#issuecomment-1906802802
- name: Clean up workspace
Expand Down Expand Up @@ -102,7 +108,9 @@ jobs:
pip install fsspec
pip install rich
# Jax nightly is needed for Triton tests.
pip install --upgrade "jax[cuda12]"
pip install -U --pre jaxlib -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html
pip install -U --pre jax-cuda12-pjrt jax-cuda12-plugin -f https://storage.googleapis.com/jax-releases/jax_cuda_plugin_nightly_releases.html
pip install -U --pre jax -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html
- name: Run Tests
env:
PJRT_DEVICE: CUDA
Expand Down
2 changes: 1 addition & 1 deletion infra/tpu-pytorch/infra_triggers.tf
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ module "terraform_apply" {
source = "../terraform_modules/apply_terraform_trigger"

included_files = ["infra/**"]
branch = "main"
branch = "master"
config_directory = "infra/tpu-pytorch"

worker_pool_id = module.worker_pool.id
Expand Down
Loading