From df694f34016bec026283e5222b77f8c4c0341e86 Mon Sep 17 00:00:00 2001 From: Will Cromar Date: Tue, 9 Apr 2024 22:45:23 +0000 Subject: [PATCH 1/4] Simplify developer setup steps --- experimental/torch_xla2/README.md | 62 ++++++++----------- experimental/torch_xla2/dev-requirements.txt | 15 +++-- experimental/torch_xla2/pyproject.toml | 31 +++++----- experimental/torch_xla2/test-requirements.txt | 3 + experimental/torch_xla2/test_requirements.txt | 5 -- 5 files changed, 51 insertions(+), 65 deletions(-) create mode 100644 experimental/torch_xla2/test-requirements.txt delete mode 100644 experimental/torch_xla2/test_requirements.txt diff --git a/experimental/torch_xla2/README.md b/experimental/torch_xla2/README.md index adf3a4de83c..fba08f40498 100644 --- a/experimental/torch_xla2/README.md +++ b/experimental/torch_xla2/README.md @@ -4,7 +4,8 @@ Currently this is only source-installable. Requires Python version >= 3.10. -### NOTE: +### NOTE: + Please don't install torch-xla from instructions in https://github.com/pytorch/xla/blob/master/CONTRIBUTING.md . In particular, the following are not needed: @@ -18,57 +19,44 @@ TorchXLA2 and torch-xla have different installation instructions, please follow the instructions below from scratch (fresh venv / conda environment.) -### 1. Install dependencies +### 1. Installing `torch_xla2` -#### 1.0 (optional) Make a virtualenv / conda env, and activate it. +#### 1.0 (recommended) Make a virtualenv / conda env -```bash -conda create --name python=3.10 -conda activate -``` -Or, -```bash -python -m venv create my_venv -source my_venv/bin/activate -``` +If you are using VSCode, then [you can create a new environment from +UI](https://code.visualstudio.com/docs/python/environments). Select the +`dev-requirements.txt` when asked to install project dependencies. -#### 1.1 Install torch CPU, even if your device has GPU or TPU: +Otherwise create a new environment from the command line. ```bash -pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu -``` - -Or, follow official instructions in [pytorch.org](https://pytorch.org/get-started/locally/) to install for your OS. - -#### 1.2 Install Jax for either GPU or TPU - -If you are using Google Cloud TPU, then -```bash -pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html -``` +# Option 1: venv +python -m venv create my_venv +source my_venv/bin/activate -If you are using a machine with NVidia GPU: +# Option 2: conda +conda create --name python=3.10 +conda activate -```bash -pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html +# Either way, install the dev requirements. +pip install -r dev-requirements.txt ``` -If you are using a CPU-only machine: -```bash -pip install --upgrade "jax[cpu]" -``` +Note: `dev-requirements.txt` will install the CPU-only version of PyTorch. -Or, follow the official instructions in https://jax.readthedocs.io/en/latest/installation.html to install for your OS or Device. +#### 1.1 Install this package -#### 1.3 Install this package +Install `torch_xla2` from source for your platform: ```bash -pip install -e . +pip install -e .[cpu] +pip install -e .[cuda] +pip install -e .[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html ``` -#### 1.4 (optional) verify installation by running tests +#### 1.2 (optional) verify installation by running tests ```bash -pip install -r test_requirements.txt +pip install -r test-requirements.txt pytest test -``` \ No newline at end of file +``` diff --git a/experimental/torch_xla2/dev-requirements.txt b/experimental/torch_xla2/dev-requirements.txt index 4a32310fbda..bdab196f0d7 100644 --- a/experimental/torch_xla2/dev-requirements.txt +++ b/experimental/torch_xla2/dev-requirements.txt @@ -1,9 +1,8 @@ -absl-py==2.0.0 -flatbuffers==23.5.26 -jax==0.4.23 -jaxlib==0.4.23 -pytest -tensorflow +-f https://download.pytorch.org/whl/torch torch==2.2.1+cpu -immutabledict -sentencepiece \ No newline at end of file +torchvision +torchaudio +torchtext +ruff~=0.3.5 +pytest +pytest-xdist diff --git a/experimental/torch_xla2/pyproject.toml b/experimental/torch_xla2/pyproject.toml index d24c7a32ede..fc5a251726d 100644 --- a/experimental/torch_xla2/pyproject.toml +++ b/experimental/torch_xla2/pyproject.toml @@ -2,29 +2,30 @@ requires = ["hatchling"] build-backend = "hatchling.build" - [project] version = "0.0.1" name = "torch_xla2" dependencies = [ "absl-py", - "flatbuffers", + "immutabledict", + "jax>=0.4.24", "pytest", - "tensorflow", - - # Note: Exclude these because otherwise on pip install . - # pip will install libs from pypi which is the GPU version - # of these libs. - # We most likely need CPU version of torch and TPU version of - # jax. So it's best for users to install them by hand - # See more at README.md - # "jax>=0.4.24", - # "jaxlib>=0.4.24", - # "torch", + "tensorflow-cpu", + # Developers should use `dev-requirements.txt` first + "torch>=2.2.1", ] +requires-python = ">=3.10" +license = {file = "LICENSE"} + +[project.optional-dependencies] +cpu = ["jax[cpu]"] +# Add libtpu index `-f https://storage.googleapis.com/libtpu-releases/index.html` +tpu = ["jax[tpu]"] +cuda = ["jax[cuda12]"] [tool.pytest.ini_options] addopts="-n auto" -requires-python = ">=3.10" -license = {file = "LICENSE"} +[tool.ruff] +line-length = 80 +indent-width = 2 diff --git a/experimental/torch_xla2/test-requirements.txt b/experimental/torch_xla2/test-requirements.txt new file mode 100644 index 00000000000..854dc21a573 --- /dev/null +++ b/experimental/torch_xla2/test-requirements.txt @@ -0,0 +1,3 @@ +-r dev-requirements.txt +sentencepiece +expecttest diff --git a/experimental/torch_xla2/test_requirements.txt b/experimental/torch_xla2/test_requirements.txt deleted file mode 100644 index c8596327236..00000000000 --- a/experimental/torch_xla2/test_requirements.txt +++ /dev/null @@ -1,5 +0,0 @@ -pytest -immutabledict -sentencepiece -pytest-xdist -expecttest \ No newline at end of file From 6ce70ae1f4799afc9616ff43b5204112467f22b1 Mon Sep 17 00:00:00 2001 From: Will Cromar Date: Tue, 9 Apr 2024 22:55:50 +0000 Subject: [PATCH 2/4] adjust requirements --- experimental/torch_xla2/dev-requirements.txt | 5 ----- experimental/torch_xla2/pyproject.toml | 2 +- experimental/torch_xla2/test-requirements.txt | 2 ++ 3 files changed, 3 insertions(+), 6 deletions(-) diff --git a/experimental/torch_xla2/dev-requirements.txt b/experimental/torch_xla2/dev-requirements.txt index bdab196f0d7..004a1d71ad7 100644 --- a/experimental/torch_xla2/dev-requirements.txt +++ b/experimental/torch_xla2/dev-requirements.txt @@ -1,8 +1,3 @@ -f https://download.pytorch.org/whl/torch torch==2.2.1+cpu -torchvision -torchaudio -torchtext ruff~=0.3.5 -pytest -pytest-xdist diff --git a/experimental/torch_xla2/pyproject.toml b/experimental/torch_xla2/pyproject.toml index fc5a251726d..0c2101dbcb9 100644 --- a/experimental/torch_xla2/pyproject.toml +++ b/experimental/torch_xla2/pyproject.toml @@ -11,7 +11,7 @@ dependencies = [ "jax>=0.4.24", "pytest", "tensorflow-cpu", - # Developers should use `dev-requirements.txt` first + # Developers should install `dev-requirements.txt` first "torch>=2.2.1", ] requires-python = ">=3.10" diff --git a/experimental/torch_xla2/test-requirements.txt b/experimental/torch_xla2/test-requirements.txt index 854dc21a573..1deead455a1 100644 --- a/experimental/torch_xla2/test-requirements.txt +++ b/experimental/torch_xla2/test-requirements.txt @@ -1,3 +1,5 @@ -r dev-requirements.txt +pytest +pytest-xdist sentencepiece expecttest From de25dd42706510ddb7de4494e6fb7879b0473b13 Mon Sep 17 00:00:00 2001 From: Will Cromar Date: Tue, 9 Apr 2024 22:59:06 +0000 Subject: [PATCH 3/4] update github workflow --- .github/workflows/torch_xla2.yml | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/.github/workflows/torch_xla2.yml b/.github/workflows/torch_xla2.yml index 7c5a88bf430..72732286157 100644 --- a/.github/workflows/torch_xla2.yml +++ b/.github/workflows/torch_xla2.yml @@ -34,9 +34,7 @@ jobs: shell: bash working-directory: experimental/torch_xla2 run: | - pip install pytest absl-py jax[cpu] flatbuffers tensorflow - pip install torch --index-url https://download.pytorch.org/whl/cpu - pip install -r test_requirements.txt + pip install -r test-requirements.txt pip install -e . - name: Run tests working-directory: experimental/torch_xla2 From e7ca0a2182ffafbbca082ca0b2e1bed69632ce6a Mon Sep 17 00:00:00 2001 From: Will Cromar Date: Tue, 9 Apr 2024 23:13:56 +0000 Subject: [PATCH 4/4] install cpu jaxlib --- .github/workflows/torch_xla2.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/torch_xla2.yml b/.github/workflows/torch_xla2.yml index 72732286157..441addad422 100644 --- a/.github/workflows/torch_xla2.yml +++ b/.github/workflows/torch_xla2.yml @@ -35,7 +35,7 @@ jobs: working-directory: experimental/torch_xla2 run: | pip install -r test-requirements.txt - pip install -e . + pip install -e .[cpu] - name: Run tests working-directory: experimental/torch_xla2 shell: bash