Enable PyPI publishing (#725)

Summary:
This PR enables building wheel files for pymomentum and publishing to PyPI ([`pymomentum-gpu`](https://pypi.org/project/pymomentum-gpu/) and [`pymomentum-cpu`](https://pypi.org/project/pymomentum-cpu/)).

*   New PyPI workflow with CPU/GPU wheel builds
*   Template-based config system for platform-specific PyTorch versions
*   CI cost optimizations (concurrency control, Python 3.13 skipping)
*   Ubuntu CUDA 12.9.0 installer workaround
*   Disk space optimization for GPU builds
*   Comprehensive documentation added

## Checklist:

- [x] Adheres to the [style guidelines](https://facebookresearch.github.io/momentum/docs/developer_guide/style_guide)
- [x] Codebase formatted by running `pixi run lint`

Pull Request resolved: #725

Test Plan:
https://fb.workplace.com/groups/1222607168126483/permalink/2638230143230838/

*   CI wheel builds (no publishing)
*   CI publish to PyPI on git tag
*   Local testing with pixi commands (e.g., `pixi run -e py313 build_pypi_wheel`)

Reviewed By: yutingye

Differential Revision: D85481027

Pulled By: jeongseok-meta

fbshipit-source-id: 76cfa74ba47e22bebe3e2ef62abc1600883fff99

Author: jeongseok-meta

.github/actions/setup-cuda/action.yml

@@ -0,0 +1,33 @@
+name: 'Setup CUDA Toolkit'
+description: 'Install CUDA toolkit and configure Conda environment variables for pixi'
+
+inputs:
+  cuda-version:
+    description: 'CUDA version to install (e.g., 12.9.0)'
+    required: true
+
+runs:
+  using: 'composite'
+  steps:
+    - name: Install CUDA Toolkit (Windows only)
+      # Skip Ubuntu due to CUDA 12.9.0 installer bug: boost::filesystem::copy_file error
+      # For wheel builds, we only need CONDA_OVERRIDE_CUDA (PyTorch includes CUDA runtime)
+      if: runner.os == 'Windows'
+      uses: Jimver/[email protected]
+      with:
+        cuda: ${{ inputs.cuda-version }}
+
+    - name: Check CUDA Version (Windows only)
+      if: runner.os == 'Windows'
+      shell: bash
+      run: |
+        nvcc --version
+
+    - name: Set Conda environment variables
+      shell: bash
+      run: |
+        CUDA_MAJOR=$(echo ${{ inputs.cuda-version }} | cut -d'.' -f1)
+        echo "CONDA_OVERRIDE_CUDA=$CUDA_MAJOR" >> $GITHUB_ENV
+        if [[ "$RUNNER_OS" == "Linux" ]]; then
+          echo "Skipping CUDA installation on Ubuntu (using mock CUDA $CUDA_MAJOR for pixi)"
+        fi

.github/workflows/ci_macos.yml

@@ -13,6 +13,10 @@ on:
       - "**/website/**"
   workflow_dispatch:
 
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
 jobs:
   momentum:
     # Enabled for minimal macOS platform coverage

.github/workflows/ci_ubuntu.yml

@@ -13,6 +13,10 @@ on:
       - "**/website/**"
   workflow_dispatch:
 
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
 jobs:
   momentum:
     name: cpp${{ matrix.simd == 'ON' && '-simd' || '' }}-${{ matrix.mode == '' && 'opt' || 'dev' }}-ubuntu
@@ -62,8 +66,8 @@ jobs:
       fail-fast: false
       matrix:
         include:
-          - pixi_env: "cpu"
-          - pixi_env: "gpu"
+          - pixi_env: "py312"
+          - pixi_env: "py312-cuda129"
             cuda_version: "12.9.0"
     env:
       FULL_CUDA_VERSION: ${{ matrix.cuda_version }}
@@ -82,23 +86,11 @@ jobs:
       - name: Checkout
         uses: actions/checkout@v5
 
-      - name: Install CUDA Toolkit
+      - name: Setup CUDA
         if: ${{ contains(matrix.pixi_env, 'cuda') || contains(matrix.pixi_env, 'gpu') }}
-        uses: Jimver/[email protected]
-        id: cuda-toolkit
+        uses: ./.github/actions/setup-cuda
         with:
-          # Available versions: https://github.com/Jimver/cuda-toolkit/blob/v0.2.24/src/links/linux-links.ts
-          cuda: ${{ matrix.cuda_version }}
-
-      - name: Check CUDA Version
-        if: ${{ contains(matrix.pixi_env, 'cuda') || contains(matrix.pixi_env, 'gpu') }}
-        run: |
-          nvcc --version
-
-      - name: Set Conda environment variables
-        if: ${{ contains(matrix.pixi_env, 'cuda') || contains(matrix.pixi_env, 'gpu') }}
-        run: |
-          echo "CONDA_OVERRIDE_CUDA=$(echo $FULL_CUDA_VERSION | cut -d'.' -f1)" >> $GITHUB_ENV
+          cuda-version: ${{ matrix.cuda_version }}
 
       - name: Set up pixi
         uses: prefix-dev/[email protected]

.github/workflows/ci_windows.yml

@@ -13,6 +13,10 @@ on:
       - "**/website/**"
   workflow_dispatch:
 
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
 jobs:
   momentum:
     if: false # Temporarily disabled to reduce CI cost
@@ -54,35 +58,21 @@ jobs:
       fail-fast: false
       matrix:
         include:
-          - pixi_env: "cpu"
+          - pixi_env: "py312"
           # GPU disabled to reduce CI cost
-          # - pixi_env: "gpu"
+          # - pixi_env: "py312-cuda129"
           #   cuda_version: "12.9.0"
     env:
       FULL_CUDA_VERSION: ${{ matrix.cuda_version }}
     steps:
       - name: Checkout
         uses: actions/checkout@v5
 
-      - name: Install CUDA Toolkit
+      - name: Setup CUDA
         if: ${{ contains(matrix.pixi_env, 'cuda') || contains(matrix.pixi_env, 'gpu') }}
-        uses: Jimver/[email protected]
-        id: cuda-toolkit
+        uses: ./.github/actions/setup-cuda
         with:
-          # Available versions: https://github.com/Jimver/cuda-toolkit/blob/v0.2.24/src/links/linux-links.ts
-          cuda: ${{ matrix.cuda_version }}
-
-      - name: Check CUDA Version
-        if: ${{ contains(matrix.pixi_env, 'cuda') || contains(matrix.pixi_env, 'gpu') }}
-        run: |
-          nvcc --version
-
-      - name: Set Conda environment variables
-        if: ${{ contains(matrix.pixi_env, 'cuda') || contains(matrix.pixi_env, 'gpu') }}
-        shell: pwsh
-        run: |
-          $majorVersion = $Env:FULL_CUDA_VERSION.Split('.')[0]
-          "CONDA_OVERRIDE_CUDA=$majorVersion" | Out-File -FilePath $Env:GITHUB_ENV -Encoding utf8 -Append
+          cuda-version: ${{ matrix.cuda_version }}
 
       - name: Set up pixi
         uses: prefix-dev/[email protected]