Merge branch 'main' into gemma3

Author: krammnic

.github/workflows/export.yaml

@@ -30,7 +30,7 @@ jobs:
         python-version: ['3.9', '3.10', '3.11']
     steps:
       - name: Check out repo
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
       - name: Setup conda env
         uses: conda-incubator/setup-miniconda@v2
         with:

.github/workflows/gpu_test.yaml

@@ -24,7 +24,7 @@ defaults:
 jobs:
   gpu_test:
     if: github.repository_owner == 'pytorch'
-    runs-on: linux.8xlarge.nvidia.gpu
+    runs-on: linux.g5.12xlarge.nvidia.gpu
     strategy:
       matrix:
         python-version: ['3.9', '3.10', '3.11']
@@ -34,7 +34,7 @@ jobs:
           - torch-version: ${{ github.event_name == 'pull_request' && 'nightly' }}
     steps:
       - name: Check out repo
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
       - name: Setup conda env
         uses: conda-incubator/setup-miniconda@v2
         with:
@@ -44,17 +44,17 @@ jobs:
           python-version: ${{ matrix.python-version }}
       - name: Update pip
         run: python -m pip install --upgrade pip
-      - name: Install torch nightly
+      - name: Install nightly versions of PyTorch packages (if applicable)
         if: ${{ matrix.torch-version == 'nightly' }}
-        run: python -m pip install --pre torch==2.7.0.dev20250201 torchvision==0.22.0.dev20250201 torchao==0.9.0.dev20250201 --index-url https://download.pytorch.org/whl/nightly/cu126
-      - name: Install torch stable
+        run: python -m pip install --pre torch torchvision torchao --index-url https://download.pytorch.org/whl/nightly/cu126
+      - name: Install torch stable (if applicable)
         if: ${{ matrix.torch-version == 'stable' }}
         run: python -m pip install torch torchvision torchao
-      - name: Install remaining dependencies
-        run: |
-          python -m pip install -e ".[dev]"
-          python -m pip install lm-eval>=0.4.5
+      - name: Install recipe-specific dependencies
+        run: python -m pip install lm-eval==0.4.8
+      - name: Install the torchtune library with dev options
+        run: python -m pip install -e ".[dev]"
       - name: Run recipe and unit tests with coverage
         run: pytest tests --ignore tests/torchtune/modules/_export --with-integration --cov=. --cov-report=xml --durations=20 -vv
-      - name: Upload Coverage to Codecov
+      - name: Upload coverage to Codecov
         uses: codecov/codecov-action@v3

.github/workflows/lint.yaml

@@ -22,7 +22,7 @@ jobs:
         python-version: ['3.10']
     steps:
       - name: Check out repo
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
       - name: Setup python
         uses: actions/setup-python@v4
         with:
@@ -35,6 +35,6 @@ jobs:
           pre-commit install-hooks
       - name: Get changed files
         id: changed-files
-        uses: tj-actions/changed-files@v41.0.0
+        uses: tj-actions/changed-files@d6e91a2266cdb9d62096cebf1e8546899c6aa18f # v45.0.6
       - name: Lint modified files
         run: pre-commit run --files ${{ steps.changed-files.outputs.all_changed_files }}

.github/workflows/recipe_test.yaml

@@ -26,11 +26,9 @@ jobs:
         python-version: ['3.11']
         torch-version: ["stable", "nightly"]
       fail-fast: false
-    env:
-      PYTORCH_CUDA_ALLOC_CONF: expandable_segments:True
     steps:
       - name: Check out repo
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
       - name: Setup conda env
         uses: conda-incubator/setup-miniconda@v2
         with:
@@ -57,8 +55,8 @@ jobs:
         run: python -m pip install torch torchvision torchao
       - name: Install remaining dependencies
         run: |
+          python -m pip install lm-eval==0.4.8
           python -m pip install -e ".[dev]"
-          python -m pip install lm-eval>=0.4.5
       - name: Run regression tests with coverage
         run: pytest tests -m slow_integration_test --silence-s3-logs --cov=. --cov-report=xml --durations=20 -vv
       - name: Upload Coverage to Codecov

.github/workflows/regression_test.yaml

@@ -0,0 +1,71 @@
+name: RL tests
+
+on:
+  push:
+    paths:
+      - 'torchtune/recipes/dev/**grpo**'
+      - 'torchtune/recipes/configs/dev/**grpo**'
+      - 'torchtune/dev/rl/**'
+      - 'torchtune/dev/grpo/**'
+  pull_request:
+    paths:
+      - 'torchtune/recipes/dev/**grpo**'
+      - 'torchtune/recipes/configs/dev/**grpo**'
+      - 'torchtune/dev/rl/**'
+      - 'torchtune/dev/grpo/**'
+  schedule:
+    # Runs at midnight evvery day
+    - cron: '0 0 * * *'
+
+concurrency:
+  group: gpu-test-${{ github.workflow }}-${{ github.ref == 'refs/heads/main' && github.run_number || github.ref }}
+  cancel-in-progress: true
+
+permissions:
+  id-token: write
+  contents: read
+
+defaults:
+  run:
+    shell: bash -l -eo pipefail {0}
+
+jobs:
+  gpu_test:
+    if: github.repository_owner == 'pytorch'
+    runs-on: linux.g5.12xlarge.nvidia.gpu
+    strategy:
+      matrix:
+        python-version: ['3.9', '3.10', '3.11']
+        torch-version: ["stable", "nightly"]
+        # Do not run against nightlies on PR
+        exclude:
+          - torch-version: ${{ github.event_name == 'pull_request' && 'nightly' }}
+    steps:
+      - name: Check out repo
+        uses: actions/checkout@v4
+      - name: Setup conda env
+        uses: conda-incubator/setup-miniconda@v2
+        with:
+          auto-update-conda: true
+          miniconda-version: "latest"
+          activate-environment: test
+          python-version: ${{ matrix.python-version }}
+      - name: Update pip
+        run: python -m pip install --upgrade pip
+      - name: Install nightly versions of PyTorch packages (if applicable)
+        if: ${{ matrix.torch-version == 'nightly' }}
+        run: python -m pip install --pre torch torchvision torchao --index-url https://download.pytorch.org/whl/nightly/cu126
+      - name: Install torch stable (if applicable)
+        if: ${{ matrix.torch-version == 'stable' }}
+        run: python -m pip install torch torchvision torchao
+      - name: Install recipe-specific dependencies
+        run: python -m pip install lm-eval==0.4.8
+      - name: Install the torchtune library with dev options
+        run: python -m pip install -e ".[dev]"
+      - name: Install the torchtune libary with async_rl options
+        if: ${{ matrix.python-version != '3.9' }}
+        run:  python -m pip install -e ".[async_rl]"
+      - name: Run recipe and unit tests with coverage
+        run: pytest tests/torchtune/dev/rl tests/recipes/dev --run-rl-tests --with-integration --cov=. --cov-report=xml --durations=20 -vv
+      - name: Upload coverage to Codecov
+        uses: codecov/codecov-action@v3

.github/workflows/rl_test.yaml

@@ -22,7 +22,7 @@ jobs:
         python-version: ['3.9', '3.10', '3.11']
     steps:
       - name: Check out repo
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
       - name: Setup conda env
         uses: conda-incubator/setup-miniconda@v2
         with:

.github/workflows/unit_test.yaml

@@ -49,3 +49,9 @@ repos:
   hooks:
     - id: pydoclint
       args: [--config=pyproject.toml]
+
+- repo: https://github.com/astral-sh/ruff-pre-commit
+  rev: v0.11.10
+  hooks:
+    - id: ruff-check
+      args: [ --fix ]

.pre-commit-config.yaml

@@ -60,6 +60,9 @@ Whenever running tests in torchtune, favor using the command line flags as much
 
 Note that the above flags can be combined with other pytest flags, so e.g. `pytest tests -m integration_test -k 'test_loss'` will run only recipe tests matching the substring `test_loss`.
 
+> [!NOTE]
+> torchtune now contains a [prototype for an asynchronous implementation of GRPO](recipes/dev/async_grpo.md), complete with corresponding tests. Since this prototype brings with it additional dependencies, we do not run any of our async RL tests by default. To run an async RL test, you should append `--run-rl-tests` to your pytest command. E.g. `pytest --with-integration --run-rl-tests tests/recipes/dev/test_async_grpo_full_finetune_distributed.py`.
+
 > [!NOTE]
 > Expected reference values for many of our tests have been calculated on Intel-based CPUs, or CUDA-based GPUs. Precision differences when testing on other hardware may cause certain tests to fail due to calculated values falling outside the configured tolerance limits. These tests may be skipped based on detected hardware (e.g. for [MPS](https://github.com/pytorch/torchtune/blob/ca95345b732d41bab7261d208cd5c860a2f76a5a/tests/test_utils.py#L345)).