Add Github workflow to build and publish wheel to PyTorch Index nightly (#1392)

Author: joecummings

.github/scripts/update_version.sh

@@ -0,0 +1,11 @@
+version_file="assets/version.txt"
+init_file="torchtitan/__init__.py"
+if [[ -n "$BUILD_VERSION" ]]; then
+    # Update the version in version.txt
+    echo "$BUILD_VERSION" > "$version_file"
+    # Create a variable named __version__ at the end of __init__.py
+    echo "__version__ = \"$BUILD_VERSION\"" >> "$init_file"
+else
+    echo "Error: BUILD_VERSION environment variable is not set or empty."
+    exit 1
+fi

.github/workflows/build_whl_and_publish.yaml

@@ -0,0 +1,40 @@
+name: Build nightly wheels and publish to PyTorch Index
+
+on:
+  push:
+    branches:
+      - nightly
+  workflow_dispatch:
+
+permissions:
+  id-token: write
+  contents: read
+
+jobs:
+  generate-matrix:
+    if: github.repository_owner == 'pytorch'
+    uses: pytorch/test-infra/.github/workflows/generate_binary_build_matrix.yml@main
+    with:
+      package-type: wheel
+      os: linux
+      test-infra-repository: pytorch/test-infra
+      test-infra-ref: main
+      with-cuda: enable
+      with-rocm: enable
+      python-versions: '["3.10", "3.11", "3.12"]'
+  build:
+    needs: generate-matrix
+    name: ${{ matrix.repository }}
+    uses: pytorch/test-infra/.github/workflows/build_wheels_linux.yml@main
+    strategy:
+      fail-fast: false
+    with:
+      repository: pytorch/torchtitan
+      ref: ""
+      test-infra-repository: pytorch/test-infra
+      test-infra-ref: main
+      package-name: torchtitan
+      build-matrix: ${{ needs.generate-matrix.outputs.matrix }}
+      pre-script: .github/scripts/update_version.sh
+      trigger-event: ${{ github.event_name }}
+      build-platform: 'python-build-package'

README.md

@@ -89,9 +89,7 @@ You may want to see how the model is defined or how parallelism techniques are a
 
 ## Installation
 
-### Nightly
-
-Coming soon.
+> [Install PyTorch](https://pytorch.org/get-started/locally/) before proceeding.
 
 ### Stable
 
@@ -103,14 +101,24 @@ Or via conda:
 ```sh
 conda install conda-forge::torchtitan
 ```
-### Sources
-   
+
+### Nightly
+
+> This method requires the nightly build of PyTorch.
+
+```sh
+pip install --pre torchtitan --index-url https://download.pytorch.org/whl/nightly/cu126
+```
+You can replace `cu126` with another version of cuda (e.g. `cu128`) or an AMD GPU (e.g. `rocm6.3`).
+
+### From source
+
+> This method requires the nightly build of PyTorch or PyTorch built from source.
+
 ```bash
 git clone https://github.com/pytorch/torchtitan
 cd torchtitan
 pip install -r requirements.txt
-pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu126 --force-reinstall
-[For AMD GPU] pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/rocm6.3 --force-reinstall
 ```
 
 ### Downloading a tokenizer