WIP

* Feat: init

* Feat: add validation + init from kwargs

* Fix: minor fixes

* Feat: more cleanup

* Minor refactor

* remove import

* adding support for pre-configured device mesh

* adding device mesh to fsdp2

* moving mesh dim defn to parralismconfig

* tests

* WIP device mesh/accelerator validation

* WIP more tests

* Test Driven Development (TDD)

* fixing build_device_mesh

* FSDP dim names

* adding example

* WIP

* fixing HSDP

* Feat: add back old options

* working example

* debugging

* adding parallelism config to partialstate

* Feat: revert ddp changes

* Revert DDP

* Feat: (untested) update mesh dims and some minor tweaks

* adding dp_cp dims

* updating comments

* WIP

* wip 2

* reverting

* storing state in accelerator rather than acceleratorstate

* Fix: minor tweaks

* wip example update

* Fixes for non-fsdp2 case

* Feat: ensure ddp/tp only works

* updating example

* updating example

* updating examples, fixing state

* fixed state

* comments

* fixing partial state check

* linting

* comments

* removing fn

* WIP: fix tp

* comments

* removing return

* reverting upcast

* add guards

* guards for empty self.parallelism_config

* use len on tuple to check if empty

* Feat: cleanup example

* Feat: some cleanup of example

* Feat: add trackio

* Fix: improve trackio

* Feat: TP works

* Feat: some fsdp2 improv

* Feat: working examples

* handle clipping for tensor parallel

* Implicit replicate

* Refactor: move to separate file + cleanup + basic comments

* Fix: add unadded files, fix circular import

* Feat: better readme

* Feat: add blog + ultrascale links

* Tmp: should_save_model now returns only true

* Fix: remove implicit_replication and style

* Fix: remove optional

* add guard on parallelism_config.tp_enabled

* fix import

* fixing empty parallelism_config

* fix import path for test patch

* fixing patch

---------

Co-authored-by: S1ro1 <matej.sirovatka@gmail.com>
Co-authored-by: Salman Mohammadi <“salman.mohammadi@outlook.com”>
Co-authored-by: Wing Lian <wing@axolotl.ai>

.devcontainer/devcontainer.json

@@ -0,0 +1,29 @@
+// File only needed for VSCode users to have proper Docker based interpreters
+{
+    "name": "accelerate_dev_environment",
+    "build": {
+        // ACTION NEEDED: comment/uncomment the relevant line depending on whether you are in a CPU/GPU environment
+         "dockerfile": "../docker/accelerate-cpu/Dockerfile"
+//        "dockerfile": "../docker/accelerate-gpu/Dockerfile"
+    },
+    "runArgs": [
+        // ACTION NEEDED: uncomment the next line if your local machine has GPUs available
+//        "--gpus", "all",
+        // Enable the docker container to access system resources
+        "--ipc", "host"
+    ],
+    "remoteEnv": {
+        "PYTHONPATH": "${containerEnv:PATH}:${containerWorkspaceFolder}"
+    },
+    "customizations": {
+        "vscode": {
+            "extensions": [
+                // Ensure we have IntelliSense in VSCode when running inside container
+                "ms-python.python"
+            ]
+        }
+    },
+    "workspaceFolder": "/workspaces/accelerate",
+    // Need git for VSCode to color code modifications. Only runs when building environment.
+    "onCreateCommand": "apt-get update && apt-get install -y git && pip install -e '.[dev]'"
+}

.github/ISSUE_TEMPLATE/bug-report.yml

@@ -0,0 +1,63 @@
+name: "\U0001F41B Bug Report"
+description: Submit a bug report to help us improve Accelerate
+body:
+  - type: markdown
+    attributes: 
+      value: | 
+        Thanks for taking the time to submit a bug report! 🐛 
+        If this is not a bug related to the Accelerate library directly, but instead a general question about your code or the library specifically please use the [forums](https://discuss.huggingface.co/c/accelerate/18).
+
+  - type: textarea
+    id: system-info
+    attributes:
+      label: System Info
+      description: Please share your accelerate configuration with us. You can run the command `accelerate env` and copy-paste its outputs below
+      render: Shell
+      placeholder: accelerate version, OS, python version, numpy version, torch version, and accelerate's configuration
+    validations:
+      required: true
+  
+  - type: checkboxes
+    id: information-scripts-examples
+    attributes:
+      label: Information
+      description: 'The problem arises when using:'
+      options:
+        - label: "The official example scripts"
+        - label: "My own modified scripts"
+  
+  - type: checkboxes
+    id: information-tasks
+    attributes:
+      label: Tasks
+      description: "The tasks I am working on are:"
+      options:
+        - label: "One of the scripts in the examples/ folder of Accelerate or an officially supported `no_trainer` script in the `examples` folder of the `transformers` repo (such as `run_no_trainer_glue.py`)"
+        - label: "My own task or dataset (give details below)"
+  
+  - type: textarea
+    id: reproduction
+    validations:
+      required: true
+    attributes:
+      label: Reproduction
+      description: |
+        Please provide a code sample that reproduces the problem you ran into. It can be a Colab link or just a code snippet.
+        If you have code snippets, error messages, stack traces please provide them here as well.
+        Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
+        Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.
+
+      placeholder: |
+        Steps to reproduce the behavior:
+          
+          1.
+          2.
+          3.
+
+  - type: textarea
+    id: expected-behavior
+    validations:
+      required: true
+    attributes:
+      label: Expected behavior
+      description: "A clear and concise description of what you would expect to happen."

.github/PULL_REQUEST_TEMPLATE.md

@@ -0,0 +1,47 @@
+# What does this PR do?
+
+<!--
+Congratulations! You've made it this far! You're not quite done yet though.
+
+Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.
+
+Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.
+
+Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
+-->
+
+<!-- Remove if not applicable -->
+
+Fixes # (issue)
+
+
+## Before submitting
+- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
+- [ ] Did you read the [contributor guideline](https://github.com/huggingface/accelerate/blob/main/CONTRIBUTING.md#submitting-a-pull-request-pr),
+      Pull Request section?
+- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
+      to it if that's the case.
+- [ ] Did you make sure to update the documentation with your changes? Here are the
+      [documentation guidelines](https://github.com/huggingface/accelerate/tree/main/docs), and
+      [here are tips on formatting docstrings](https://github.com/huggingface/accelerate/tree/main/docs#writing-documentation---specification).
+- [ ] Did you write any new necessary tests?
+
+
+## Who can review?
+
+Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
+members/contributors who may be interested in your PR.
+
+<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @
+
+ If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
+
+- Big modeling: @SunMarc
+- Fully-Sharded Data Parallism: @SunMarc @zach-huggingface
+- DeepSpeed: @SunMarc @zach-huggingface
+- Command Line Interface: @SunMarc @zach-huggingface
+- Documentation: @SunMarc @zach-huggingface
+- Core parts of the library: @BenjaminBossan @SunMarc @zach-huggingface 
+- Maintained examples: @SunMarc or  @zach-huggingface
+
+ -->

.github/deploy_doc.sh

@@ -1,38 +0,0 @@
-#!/bin/bash
-set -ex
-
-function deploy_doc(){
-	echo "Creating doc at commit $1 and pushing to folder $2"
-	git checkout $1
-	cd "$GITHUB_WORKSPACE"
-	pip install -U .
-	cd "$GITHUB_WORKSPACE/docs"
-	if [ ! -z "$2" ]
-	then
-		if [ "$2" == "main" ]; then
-		    echo "Pushing main"
-			make clean && make html && scp -r -oStrictHostKeyChecking=no _build/html/* $DOC_HOST:$DOC_PATH/$2/
-			cp -r _build/html/_static .
-		elif ssh -oStrictHostKeyChecking=no $DOC_HOST "[ -d $DOC_PATH/$2 ]"; then
-			echo "Directory" $2 "already exists"
-			scp -r -oStrictHostKeyChecking=no _static/* $DOC_HOST:$DOC_PATH/$2/_static/
-		else
-			echo "Pushing version" $2
-			make clean && make html
-			rm -rf _build/html/_static
-			cp -r _static _build/html
-			scp -r -oStrictHostKeyChecking=no _build/html $DOC_HOST:$DOC_PATH/$2
-		fi
-	else
-		echo "Pushing stable"
-		make clean && make html
-		rm -rf _build/html/_static
-		cp -r _static _build/html
-		scp -r -oStrictHostKeyChecking=no _build/html/* $DOC_HOST:$DOC_PATH
-	fi
-}
-
-
-# You can find the commit for each tag on https://github.com/huggingface/accelerate/tags
-deploy_doc "main" main
-deploy_doc "0fbbbc5" # v0.1.0 Latest stable release

.github/workflows/build-docker-images-release.yml

@@ -0,0 +1,104 @@
+name: Build Docker images (releases)
+
+on:
+  workflow_dispatch:
+  release:
+    types: [published]
+
+concurrency:
+  group: docker-image-builds
+  cancel-in-progress: false
+
+jobs:
+  get-version:
+    runs-on: ubuntu-latest
+    outputs:
+      version: ${{ steps.step1.outputs.version }}
+    steps:
+      - uses: actions/checkout@4
+      - id: step1
+        run: echo "version=$(python setup.py --version)" >> $GITHUB_OUTPUT
+
+  version-cpu:
+    name: "Latest Accelerate CPU [version]"
+    runs-on:
+      group: aws-general-8-plus
+    needs: get-version
+    steps:
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v2
+      - name: Login to DockerHub
+        uses: docker/login-action@v2
+        with:
+          username: ${{ secrets.DOCKERHUB_USERNAME }}
+          password: ${{ secrets.DOCKERHUB_PASSWORD }}
+
+      - name: Build and Push CPU
+        uses: docker/build-push-action@v4
+        with:
+          file: docker/accelerate-cpu/Dockerfile
+          push: true
+          tags: huggingface/accelerate:cpu-release-${{ needs.get-version.outputs.version }}
+
+  version-cuda:
+    name: "Latest Accelerate GPU [version]"
+    runs-on:
+      group: aws-g6-4xlarge-plus
+    needs: get-version
+    steps:
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v2
+      - name: Login to DockerHub
+        uses: docker/login-action@v2
+        with:
+          username: ${{ secrets.DOCKERHUB_USERNAME }}
+          password: ${{ secrets.DOCKERHUB_PASSWORD }}
+
+      - name: Build and Push GPU
+        uses: docker/build-push-action@v4
+        with:
+          file: docker/accelerate-gpu/Dockerfile
+          push: true
+          tags: huggingface/accelerate:gpu-release-${{needs.get-version.outputs.version}}
+
+  version-cuda-deepspeed:
+    name: "Latest Accelerate GPU DeepSpeed [version]"
+    runs-on:
+      group: aws-g6-4xlarge-plus
+    needs: get-version
+    steps:
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v2
+      - name: Login to DockerHub
+        uses: docker/login-action@v2
+        with:
+          username: ${{ secrets.DOCKERHUB_USERNAME }}
+          password: ${{ secrets.DOCKERHUB_PASSWORD }}
+
+      - name: Build and Push GPU
+        uses: docker/build-push-action@v4
+        with:
+          file: docker/accelerate-gpu-deepspeed/Dockerfile
+          push: true
+          tags: huggingface/accelerate:gpu-deepspeed-release-${{needs.get-version.outputs.version}}
+
+  version-cuda-fp8-transformerengine:
+    name: "Latest Accelerate GPU FP8 TransformerEngine [version]"
+    runs-on:
+      group: aws-g6-4xlarge-plus
+    needs: get-version
+    steps:
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v2
+      - name: Login to DockerHub
+        uses: docker/login-action@v2
+        with:
+          username: ${{ secrets.DOCKERHUB_USERNAME }}
+          password: ${{ secrets.DOCKERHUB_PASSWORD }}
+
+      - name: Build and Push GPU
+        uses: docker/build-push-action@v4
+        with:
+          file: docker/accelerate-gpu/Dockerfile
+          push: true
+          tags: huggingface/accelerate:gpu-fp8-transformerengine-release-${{needs.get-version.outputs.version}}

.github/workflows/build_and_run_tests.yml

@@ -0,0 +1,50 @@
+name: Trigger docker images and run tests
+
+on:
+  push:
+    branches:
+      - main
+  workflow_dispatch:
+
+env:
+  GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+
+jobs:
+  check-for-source:
+    runs-on: ubuntu-latest
+    name: Check if setup was changed
+    outputs:
+      changed: ${{ steps.was_changed.outputs.changed }}
+    steps:
+      - uses: actions/checkout@v4
+        with: 
+          fetch-depth: "2"
+      
+      - name: Get changed files
+        id: changed-files
+        uses: tj-actions/changed-files@3f54ebb830831fc121d3263c1857cfbdc310cdb9 #v42
+      
+      - name: Was setup changed 
+        id: was_changed
+        run: |
+          for file in ${{ steps.changed-files.outputs.all_changed_files }}; do
+            if [ `basename "${file}"` == "setup.py" ]; then
+              echo "changed=1" >> $GITHUB_OUTPUT
+            fi
+          done
+          
+  build-docker-containers:
+    needs: check-for-source
+    if: (github.event_name == 'push') && (needs.check-for-source.outputs.changed == '1')
+    uses: ./.github/workflows/build_docker_images.yml
+    secrets: inherit
+
+  run-merge-tests:
+    needs: build-docker-containers
+    if: always()
+    uses: ./.github/workflows/run_merge_tests.yml
+
+  run-integration-tests:
+    needs: build-docker-containers
+    if: always()
+    uses: ./.github/workflows/self_hosted_integration_tests.yml

.github/workflows/build_docker_images.yml

@@ -0,0 +1,116 @@
+name: Build Docker images (scheduled)
+
+on:
+  workflow_dispatch:
+  workflow_call:
+  schedule:
+    - cron: "0 1 * * *"
+
+concurrency:
+  group: docker-image-builds
+  cancel-in-progress: false
+
+jobs:
+  latest-cpu:
+    name: "Latest Accelerate CPU [dev]"
+    runs-on:
+      group: aws-general-8-plus
+    steps:
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v2
+      - name: Login to DockerHub
+        uses: docker/login-action@v2
+        with:
+          username: ${{ secrets.DOCKERHUB_USERNAME }}
+          password: ${{ secrets.DOCKERHUB_PASSWORD }}
+      - name: Get current date
+        id: date
+        run: |
+          echo "date=$(date '+%Y-%m-%d')" >> $GITHUB_ENV
+      - name: Build and Push CPU
+        uses: docker/build-push-action@v4
+        with:
+          file: docker/accelerate-cpu/Dockerfile
+          push: true
+          tags: |
+            huggingface/accelerate:cpu-nightly
+            huggingface/accelerate:cpu-nightly-${{ env.date }}
+
+  latest-cuda:
+    name: "Latest Accelerate GPU [dev]"
+    runs-on:
+      group: aws-g6-4xlarge-plus
+    steps:
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v2
+      - name: Login to DockerHub
+        uses: docker/login-action@v2
+        with:
+          username: ${{ secrets.DOCKERHUB_USERNAME }}
+          password: ${{ secrets.DOCKERHUB_PASSWORD }}
+      - name: Get current date
+        id: date
+        run: |
+          echo "date=$(date '+%Y-%m-%d')" >> $GITHUB_ENV
+      - name: Build and Push GPU
+        uses: docker/build-push-action@v4
+        with:
+          file: docker/accelerate-gpu/Dockerfile
+          push: true
+          tags: |
+            huggingface/accelerate:gpu-nightly
+            huggingface/accelerate:gpu-nightly-${{ env.date }}
+
+  latest-cuda-deepspeed:
+    name: "Latest Accelerate GPU DeepSpeed [dev]"
+    runs-on:
+      group: aws-g6-4xlarge-plus
+    steps:
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v2
+      - name: Login to DockerHub
+        uses: docker/login-action@v2
+        with:
+          username: ${{ secrets.DOCKERHUB_USERNAME }}
+          password: ${{ secrets.DOCKERHUB_PASSWORD }}
+      - name: Get current date
+        id: date
+        run: |
+          echo "date=$(date '+%Y-%m-%d')" >> $GITHUB_ENV
+      - name: Build and Push GPU
+        uses: docker/build-push-action@v4
+        with:
+          file: docker/accelerate-gpu-deepspeed/Dockerfile
+          push: true
+          tags: |
+            huggingface/accelerate:gpu-deepspeed-nightly
+            huggingface/accelerate:gpu-deepspeed-nightly-${{ env.date }}
+
+  latest-cuda-fp8-transformerengine:
+    name: "Latest Accelerate GPU FP8 TransformerEngine [dev]"
+    runs-on:
+      group: aws-g6-4xlarge-plus
+    steps:
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v2
+      - name: Login to DockerHub
+        uses: docker/login-action@v2
+        with:
+          username: ${{ secrets.DOCKERHUB_USERNAME }}
+          password: ${{ secrets.DOCKERHUB_PASSWORD }}
+      - name: Get current date
+        id: date
+        run: |
+          echo "date=$(date '+%Y-%m-%d')" >> $GITHUB_ENV
+          # Get the previous month
+          echo "base_year=$(date -d 'last month' '+%y')" >> $GITHUB_ENV
+          echo "base_month=$(date -d 'last month' '+%m')" >> $GITHUB_ENV
+      - name: Build and Push GPU
+        uses: docker/build-push-action@v4
+        with:
+          file: benchmarks/fp8/transformer_engine/Dockerfile
+          push: true
+          tags: huggingface/accelerate:gpu-fp8-transformerengine-nightly-${{ env.date }}
+          build-args: |
+            BASE_YEAR=${{ env.base_year }}
+            BASE_MONTH=${{ env.base_month }}

.github/workflows/build_documentation.yml

@@ -0,0 +1,18 @@
+name: Build documentation
+
+on:
+  push:
+    branches:
+      - main
+      - doc-builder*
+      - v*-release
+
+jobs:
+   build:
+    uses: huggingface/doc-builder/.github/workflows/build_main_documentation.yml@main
+    with:
+      commit_sha: ${{ github.sha }}
+      package: accelerate
+      custom_container: huggingface/transformers-doc-builder
+    secrets:
+      hf_token: ${{ secrets.HF_DOC_BUILD_PUSH }}

.github/workflows/build_pr_documentation.yml

@@ -0,0 +1,17 @@
+name: Build PR Documentation
+
+on:
+  pull_request:
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
+  cancel-in-progress: true
+
+jobs:
+  build:
+    uses: huggingface/doc-builder/.github/workflows/build_pr_documentation.yml@main
+    with:
+      commit_sha: ${{ github.event.pull_request.head.sha }}
+      pr_number: ${{ github.event.number }}
+      package: accelerate
+      custom_container: huggingface/transformers-doc-builder

.github/workflows/docs-deploy.yml

@@ -1,37 +0,0 @@
-name: Deploy Documentation
-
-on:
-  push:
-    branches:
-      - main
-
-jobs:
-  deploy:
-    runs-on: ubuntu-latest
-    steps:
-      - name: Checkout repository
-        uses: actions/checkout@v1
-        with:
-          fetch-depth: 0
-
-      - name: Install SSH Key
-        uses: shimataro/ssh-key-action@v2
-        with:
-          key: ${{ secrets.DOC_SSH_KEY }}
-          name: id_rsa
-          known_hosts: ${{ secrets.DOC_KNOWN_HOST }}
-
-      - name: Install Python
-        uses: actions/setup-python@v1
-        with:
-          python-version: 3.6
-
-      - name: Install Python dependencies
-        working-directory: ./
-        run: pip install -e .[docs]
-
-      - name: Deploy documentation
-        env:
-          DOC_HOST: ${{ secrets.DOC_HOST }}
-          DOC_PATH: ${{ secrets.DOC_PATH }}
-        run: ./.github/deploy_doc.sh

.github/workflows/fp8_runner.yml

@@ -0,0 +1,37 @@
+name: Test FP8 Runner
+
+on:
+  workflow_dispatch:
+
+env:
+  GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+jobs:
+  set-prev-day:
+    runs-on: ubuntu-latest
+    outputs:
+      prev-day: ${{ steps.set-prev-day.outputs.prev-day }}
+    steps:
+      - name: Set PREV_DAY
+        id: set-prev-day
+        run: |
+          PREV_DAY=$(date -d "yesterday" '+%Y-%m-%d')
+          echo "prev-day=$PREV_DAY" >> $GITHUB_OUTPUT
+  run-fp8-tests:
+    needs: set-prev-day
+    runs-on:
+      group: aws-g6e-12xlarge
+    container:
+      image: huggingface/accelerate:gpu-fp8-transformerengine-nightly-${{ needs.set-prev-day.outputs.prev-day }}
+      options: --gpus all --shm-size "16gb"
+    steps:
+      - uses: actions/checkout@v3
+      - name: Install the library
+        run: |
+            pip install -e .[test_prod,test_fp8]
+      - name: Show installed libraries
+        run: |
+          pip freeze
+      - name: Run TE FP8 tests
+        run: |
+          python -m pytest -s -v ./tests/test_fp8.py
+

.github/workflows/gaudi3_scheduled.yml

@@ -0,0 +1,87 @@
+name: Gaudi3 tests (scheduled)
+
+on:
+  workflow_dispatch:
+  schedule: # every day at 6 AM UTC
+    - cron: "0 6 * * *"
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
+  cancel-in-progress: true
+
+jobs:
+  run-gaudi3-tests:
+    runs-on:
+      group: itac-bm-emr-gaudi3-dell-2gaudi
+
+    container:
+      image: docker://vault.habana.ai/gaudi-docker/1.21.1/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest
+      options: --runtime=habana --shm-size=64G --cap-add=sys_nice --env HABANA_VISIBLE_DEVICES
+      env:
+        OMPI_MCA_btl_vader_single_copy_mechanism: none
+        PT_ENABLE_INT64_SUPPORT: 1
+        PT_HPU_LAZY_MODE: 0
+        RUN_SLOW: 1
+
+    steps:
+      - name: HL-SMI (1)
+        run: |
+          hl-smi
+          echo "HABANA_VISIBLE_DEVICES=${HABANA_VISIBLE_DEVICES}"
+          echo "HABANA_VISIBLE_MODULES=${HABANA_VISIBLE_MODULES}"
+
+      - name: Extract HPU visible modules
+        id: add-modules
+        run: |
+          export HABANA_VISIBLE_MODULES=$(hl-smi -Q module_id -f csv,noheader | tr '\n' ',' | sed 's/,$//')
+          echo "HABANA_VISIBLE_MODULES=${HABANA_VISIBLE_MODULES}" >> $GITHUB_ENV
+
+      - name: HL-SMI (2)
+        run: |
+          hl-smi
+          echo "HABANA_VISIBLE_DEVICES=${HABANA_VISIBLE_DEVICES}"
+          echo "HABANA_VISIBLE_MODULES=${HABANA_VISIBLE_MODULES}"
+
+      - name: Checkout to Accelerate
+        uses: actions/checkout@v4
+
+      - name: Install Accelerate with Transformers & DeepSpeed
+        run: |
+          pip install -e .[testing] \
+            git+https://github.com/HabanaAI/DeepSpeed.git@1.20.0 \
+            git+https://github.com/huggingface/transformers.git
+
+      - name: Run CLI tests
+        if: ${{ !cancelled() && (success() || failure()) }}
+        run: |
+          make test_cli
+
+      - name: Run Core tests
+        if: ${{ !cancelled() && (success() || failure()) }}
+        run: |
+          make test_core
+
+      - name: Run Big Modeling tests
+        if: ${{ !cancelled() && (success() || failure()) }}
+        run: |
+          make test_big_modeling
+
+      - name: Run DeepSpeed integration tests
+        if: ${{ !cancelled() && (success() || failure()) }}
+        run: |
+          make test_deepspeed
+
+      - name: Run FSDP integration tests
+        if: ${{ !cancelled() && (success() || failure()) }}
+        run: |
+          make test_fsdp
+
+      - name: Run TP integration tests
+        if: ${{ !cancelled() && (success() || failure()) }}
+        run: |
+          make test_tp
+
+      - name: Run Examples tests
+        if: ${{ !cancelled() && (success() || failure()) }}
+        run: |
+          make test_examples

.github/workflows/integration_tests.yml

@@ -0,0 +1,58 @@
+# CI for specifically ensuring integrations work fine (`transformers` mainly)
+# Useful tips:
+#  - New integrations to test should have its own job, and follow a strategy method where we check both
+#    the pypi and github versions.
+#  - When checking the latest release of the integration, use
+#    git checkout $(git describe --tags `git rev-list --tags --max-count=1`) to get the latest release.
+
+name: Integration Tests
+
+on:
+  pull_request:
+    paths:
+      - "src/**"
+      - "tests/**"
+      - ".github/**"
+      - "examples/**"
+      - "setup.py"
+    types: [opened, synchronize, reopened]
+
+env:
+  HF_HOME: ~/hf_cache
+
+jobs:
+  run-trainer-tests:
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+    steps:
+    - uses: actions/checkout@v4
+    - name: Set up python 3.9
+      uses: actions/setup-python@v5
+      with:
+        python-version: 3.9
+        cache: 'pip'
+        cache-dependency-path: 'setup.py'
+
+    - name: Install Accelerate from source
+      run: |
+        pip install --upgrade pip
+        pip install -e .
+    
+    - name: Clone and install transformers
+      run: |
+        cd ..
+        git clone https://github.com/huggingface/transformers
+        cd transformers
+        pip install .[torch,testing]
+
+    - name: Show installed libraries
+      run: |
+        pip freeze
+
+    - name: Run Trainer tests
+      env:
+        WANDB_DISABLED: true
+      run: |
+        cd ../transformers
+        pytest -sv tests/trainer

.github/workflows/nightly.yml

@@ -0,0 +1,233 @@
+name: Self-hosted runner with slow tests (scheduled)
+
+on:
+  workflow_dispatch:
+  schedule:
+    - cron: "0 2 * * *"
+
+env:
+  RUN_SLOW: "yes"
+  IS_GITHUB_CI: "1"
+  SLACK_API_TOKEN: ${{ secrets.SLACK_API_TOKEN }}
+
+
+jobs:
+  run_core_tests_single_gpu:
+    runs-on:
+      group: aws-g6-4xlarge-plus
+    env:
+      CUDA_VISIBLE_DEVICES: "0"
+      TEST_TYPE: "single_gpu"
+    container:
+      image: huggingface/accelerate:gpu-nightly
+      options: --gpus all --shm-size "16gb"
+    defaults:
+      run:
+        shell: bash
+    steps:
+      - name: Update clone & pip install
+        run: |
+          source activate accelerate
+          git clone https://github.com/huggingface/accelerate;
+          cd accelerate;
+          git checkout ${{ github.sha }};
+          pip install -e . --no-deps
+          pip install pytest-reportlog tabulate
+
+      - name: Show installed libraries
+        run: |
+          source activate accelerate;
+          pip freeze
+
+      - name: Run test on GPUs
+        working-directory: accelerate
+        run: |
+          source activate accelerate
+          make test
+
+      - name: Run examples on GPUs
+        working-directory: accelerate
+        if: always()
+        run: |
+          source activate accelerate
+          pip uninstall comet_ml -y
+          make test_examples
+
+      - name: Generate Report
+        working-directory: accelerate
+        if: always()
+        run: |
+          pip install slack_sdk tabulate
+          python utils/log_reports.py >> $GITHUB_STEP_SUMMARY
+
+  run_deepspeed_tests_single_gpu:
+    runs-on:
+      group: aws-g6-4xlarge-plus
+    env:
+      CUDA_VISIBLE_DEVICES: "0"
+      TEST_TYPE: "single_gpu_deepspeed"
+    container:
+      image: huggingface/accelerate:gpu-deepspeed-nightly
+      options: --gpus all --shm-size "16gb"
+    defaults:
+      run:
+        shell: bash
+    steps:
+      - name: Update clone & pip install
+        run: |
+          source activate accelerate
+          git clone https://github.com/huggingface/accelerate;
+          cd accelerate;
+          git checkout ${{ github.sha }};
+          pip install -e . --no-deps
+          pip install pytest-reportlog tabulate
+
+      - name: Show installed libraries
+        run: |
+          source activate accelerate;
+          pip freeze
+
+      - name: Run test on GPUs
+        working-directory: accelerate
+        run: |
+          source activate accelerate
+          make test_deepspeed
+
+      - name: Run Integration tests on GPUs
+        working-directory: accelerate
+        if: always()
+        run: |
+          source activate accelerate
+          make test_integrations
+
+      - name: Run examples on GPUs
+        working-directory: accelerate
+        if: always()
+        run: |
+          source activate accelerate
+          pip uninstall comet_ml -y
+          make test_examples
+
+      - name: Generate Report
+        working-directory: accelerate
+        if: always()
+        run: |
+          pip install slack_sdk tabulate
+          python utils/log_reports.py >> $GITHUB_STEP_SUMMARY
+
+  run_core_tests_multi_gpu:
+    runs-on:
+      group: aws-g6-12xlarge-plus
+    env:
+      CUDA_VISIBLE_DEVICES: "0,1"
+      TEST_TYPE: "multi_gpu"
+    container:
+      image: huggingface/accelerate:gpu-nightly
+      options: --gpus all --shm-size "16gb"
+    defaults:
+      run:
+        shell: bash
+    steps:
+      - name: Update clone
+        run: |
+          source activate accelerate
+          git clone https://github.com/huggingface/accelerate;
+          cd accelerate;
+          git checkout ${{ github.sha }};
+          pip install -e . --no-deps
+          pip install pytest-reportlog tabulate
+
+      - name: Show installed libraries
+        run: |
+          source activate accelerate;
+          pip freeze
+
+      - name: Run core and big modeling tests on GPUs
+        working-directory: accelerate
+        run: |
+          source activate accelerate
+          make test_core
+          make test_big_modeling
+          make test_cli
+
+      - name: Run Integration tests on GPUs
+        working-directory: accelerate
+        if: always()
+        run: |
+          source activate accelerate
+          make test_integrations
+
+      - name: Run examples on GPUs
+        working-directory: accelerate
+        if: always()
+        run: |
+          source activate accelerate
+          pip uninstall comet_ml -y
+          make test_examples
+
+      - name: Generate Report
+        working-directory: accelerate
+        if: always()
+        run: |
+          pip install slack_sdk tabulate
+          python utils/log_reports.py >> $GITHUB_STEP_SUMMARY
+
+  run_deepspeed_tests_multi_gpu:
+    runs-on:
+      group: aws-g6-12xlarge-plus
+    env:
+      CUDA_VISIBLE_DEVICES: "0,1"
+      TEST_TYPE: "multi_gpu_deepspeed"
+    container:
+      image: huggingface/accelerate:gpu-deepspeed-nightly
+      options: --gpus all --shm-size "16gb"
+    defaults:
+      run:
+        shell: bash
+    steps:
+      - name: Update clone
+        run: |
+          source activate accelerate
+          git clone https://github.com/huggingface/accelerate;
+          cd accelerate;
+          git checkout ${{ github.sha }};
+          pip install -e . --no-deps
+          pip install pytest-reportlog tabulate
+
+      - name: Show installed libraries
+        run: |
+          source activate accelerate;
+          pip freeze
+
+      - name: Run DeepSpeed tests
+        working-directory: accelerate
+        run: |
+          source activate accelerate
+          make test_deepspeed
+
+      - name: Run Integration tests on GPUs
+        working-directory: accelerate
+        if: always()
+        run: |
+          source activate accelerate
+          make test_integrations
+
+      - name: Run examples on GPUs
+        working-directory: accelerate
+        if: always()
+        run: |
+          source activate accelerate
+          pip uninstall comet_ml -y
+          make test_examples
+
+      - name: Generate Report
+        working-directory: accelerate
+        if: always()
+        run: |
+          pip install slack_sdk tabulate
+          python utils/log_reports.py >> $GITHUB_STEP_SUMMARY
+
+
+  run-integration-tests:
+    if: always()
+    uses: ./.github/workflows/self_hosted_integration_tests.yml

.github/workflows/pr_style_bot.yml

@@ -0,0 +1,19 @@
+# To run this bot, comment "@bot /style" on a PR
+name: Style Bot
+
+on:
+  issue_comment:
+    types: [created]
+
+permissions:
+  contents: write
+  pull-requests: write
+
+jobs:
+  style:
+    uses: huggingface/huggingface_hub/.github/workflows/style-bot-action.yml@main
+    with:
+      python_quality_dependencies: "[quality]"
+      style_command_type: "default"
+    secrets:
+      bot_token: ${{ secrets.GITHUB_TOKEN }}

.github/workflows/quality.yml

@@ -6,12 +6,19 @@ jobs:
   quality:
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@v2
-    - name: Set up Python 3.6
-      uses: actions/setup-python@v2
+    - uses: actions/checkout@v4
+    - name: Set up Python 3.9
+      uses: actions/setup-python@v5
       with:
-        python-version: 3.6
+        python-version: 3.9
+        cache: 'pip'
+        cache-dependency-path: 'setup.py'
     - name: Install Python dependencies
       run: pip install -e .[quality]
     - name: Run Quality check
-      run: make quality
+      run: make quality
+    - name: Check if failure
+      if: ${{ failure() }}
+      run: |
+        echo "Quality check failed. Please ensure the right dependency versions are installed with 'pip install -e .[quality]' and rerun 'make style; make quality;'" >> $GITHUB_STEP_SUMMARY
+

.github/workflows/run_merge_tests.yml

@@ -0,0 +1,188 @@
+name: Self-hosted runner tests (push to "main")
+
+on:
+  workflow_call:
+  workflow_dispatch:
+
+env:
+  TESTING_MOCKED_DATALOADERS: "1"
+  IS_GITHUB_CI: "1"
+
+jobs:
+  run_core_tests_single_gpu:
+    runs-on:
+      group: aws-g6-4xlarge-plus
+    env:
+      CUDA_VISIBLE_DEVICES: "0"
+    container:
+      image: huggingface/accelerate:gpu-nightly
+      options: --gpus all --shm-size "16gb"
+    defaults:
+      run:
+        shell: bash
+    steps:
+      - name: Install accelerate
+        run: |
+          source activate accelerate;
+          git clone https://github.com/huggingface/accelerate;
+          cd accelerate;
+          git checkout ${{ github.sha }};
+          pip install -e .[testing,test_trackers] -U;
+          pip install pytest-reportlog tabulate  ;
+
+      - name: Show installed libraries
+        run: |
+          source activate accelerate;
+          pip freeze
+
+      - name: Run CLI tests (use make cli)
+        working-directory: accelerate
+        run: |
+          source activate accelerate;
+          make test_cli
+
+      - name: Run test on GPUs
+        working-directory: accelerate
+        if: always()
+        run: |
+          source activate accelerate;
+          make test
+      - name: Run examples on GPUs
+        working-directory: accelerate
+        if: always()
+        run: |
+          source activate accelerate;
+          pip uninstall comet_ml -y;
+          make test_examples
+
+      - name: Generate Report
+        working-directory: accelerate
+        if: always()
+        run: |
+          pip install tabulate;
+          python utils/log_reports.py >> $GITHUB_STEP_SUMMARY
+
+  run_deepspeed_tests_single_gpu:
+    runs-on:
+      group: aws-g6-4xlarge-plus
+    env:
+      CUDA_VISIBLE_DEVICES: "0"
+    container:
+      image: huggingface/accelerate:gpu-deepspeed-nightly
+      options: --gpus all --shm-size "16gb"
+    defaults:
+      run:
+        shell: bash
+    steps:
+      - name: Install accelerate
+        run: |
+          source activate accelerate;
+          git clone https://github.com/huggingface/accelerate;
+          cd accelerate;
+          git checkout ${{ github.sha }};
+          pip install -e .[testing,test_trackers] -U;
+          pip install pytest-reportlog tabulate  ;
+
+      - name: Show installed libraries
+        run: |
+          source activate accelerate;
+          pip freeze
+
+      - name: Run test on GPUs
+        working-directory: accelerate
+        if: always()
+        run: |
+          source activate accelerate;
+          make test_deepspeed
+
+      - name: Generate Report
+        working-directory: accelerate
+        if: always()
+        run: |
+          pip install tabulate;
+          python utils/log_reports.py >> $GITHUB_STEP_SUMMARY
+
+  run_core_tests_multi_gpu:
+    runs-on:
+      group: aws-g6-12xlarge-plus
+    env:
+      CUDA_VISIBLE_DEVICES: 0,1
+    container:
+      image: huggingface/accelerate:gpu-nightly
+      options: --gpus all --shm-size "16gb"
+    defaults:
+      run:
+        shell: bash
+    steps:
+      - name: Update clone
+        run: |
+          source activate accelerate;
+          git clone https://github.com/huggingface/accelerate;
+          cd accelerate;
+          git checkout ${{ github.sha }};
+          pip install -e .[testing,test_trackers] -U;
+          pip install pytest-reportlog tabulate
+
+      - name: Show installed libraries
+        run: |
+          source activate accelerate;
+          pip freeze
+
+      - name: Run test on GPUs
+        working-directory: accelerate
+        run: |
+          source activate accelerate;
+          make test
+
+      - name: Run examples on GPUs
+        working-directory: accelerate
+        if: always()
+        run: |
+          source activate accelerate;
+          pip uninstall comet_ml -y;
+          make test_examples
+
+      - name: Generate Report
+        working-directory: accelerate
+        if: always()
+        run: |
+          source activate accelerate;
+          python utils/log_reports.py >> $GITHUB_STEP_SUMMARY
+
+  run_deepspeed_tests_multi_gpu:
+    runs-on:
+      group: aws-g6-12xlarge-plus
+    container:
+      image: huggingface/accelerate:gpu-deepspeed-nightly
+      options: --gpus all --shm-size "16gb"
+    defaults:
+      run:
+        shell: bash
+    steps:
+      - name: Install accelerate
+        run: |
+          source activate accelerate;
+          git clone https://github.com/huggingface/accelerate;
+          cd accelerate;
+          git checkout ${{ github.sha }};
+          pip install -e .[testing,test_trackers] -U;
+          pip install pytest-reportlog tabulate  ;
+
+      - name: Show installed libraries
+        run: |
+          source activate accelerate;
+          pip freeze
+
+      - name: Run test on GPUs
+        working-directory: accelerate
+        if: always()
+        run: |
+          source activate accelerate;
+          make test_deepspeed
+
+      - name: Generate Report
+        working-directory: accelerate
+        if: always()
+        run: |
+          pip install tabulate;
+          python utils/log_reports.py >> $GITHUB_STEP_SUMMARY

.github/workflows/self_hosted_integration_tests.yml

@@ -0,0 +1,127 @@
+# CI for specifically ensuring integrations work fine (`transformers` mainly) on GPUs
+# Useful tips:
+#  - `working-directory` should be set to the root of the repo, which is cloned on the actual CI runner.
+#    It follows the directory structure of `actions-runner/_work/{repo_name}/{repo_name}/{cloned_repo} on
+#    prem, but in Actions setting `working-directory` looks just in the `{repo_name}` level.
+#  - New integrations to test should have its own job, and follow a strategy method where we check both
+#    the pypi and github versions.
+#  - Workflow call lets this be called from `build_and_run_tests.yml`
+#  - When using a docker container, it's recommended to set `--shm-size`, we use 16gb.
+name: Integration Tests (push to "main")
+
+on:
+  workflow_call:
+  workflow_dispatch:
+
+env:
+  HF_HOME: ~/hf_cache
+
+defaults:
+  run:
+    shell: bash
+
+jobs:
+  run-trainer-tests:
+    container:
+      image: huggingface/accelerate:gpu-deepspeed-nightly
+      options: --gpus all --shm-size "16gb"
+    runs-on:
+      group: aws-g6-12xlarge-plus
+    strategy:
+      fail-fast: false
+      matrix:
+        cuda_visible_devices: [
+          "0",
+          "0,1"
+        ]
+    steps:
+      - name: Install transformers
+        run: |
+          source activate accelerate;
+          git clone https://github.com/huggingface/transformers --depth 1;
+          cd transformers;
+          pip install .[torch,deepspeed-testing];
+          cd ..;
+
+      - name: Install accelerate
+        run: |
+          source activate accelerate;
+          git clone https://github.com/huggingface/accelerate;
+          cd accelerate;
+          git checkout ${{ github.sha }} ;
+          pip install -e .[testing];
+          pip uninstall comet_ml wandb dvclive -y
+          cd ..;
+
+      - name: Show installed libraries
+        run: |
+          source activate accelerate;
+          pip freeze
+
+      - name: Run trainer tests
+        working-directory: transformers/
+        env:
+          CUDA_VISIBLE_DEVICES: ${{ matrix.cuda_visible_devices }}
+          WANDB_DISABLED: true
+        run: |
+          source activate accelerate;
+          pytest -sv tests/trainer
+
+      - name: Run deepspeed tests
+        working-directory: transformers/
+        env:
+          CUDA_VISIBLE_DEVICES: ${{ matrix.cuda_visible_devices }}
+          WANDB_DISABLED: true
+        if: always()
+        run: |
+          source activate accelerate;
+          pytest -sv tests/deepspeed
+
+      - name: Run transformers examples tests
+        working-directory: transformers/
+        env:
+          CUDA_VISIBLE_DEVICES: ${{ matrix.cuda_visible_devices }}
+          WANDB_DISABLED: true
+        run: |
+          source activate accelerate
+          pip install -r examples/pytorch/_tests_requirements.txt
+          pytest -sv examples/pytorch/test_accelerate_examples.py examples/pytorch/test_pytorch_examples.py
+
+  run-skorch-tests:
+    container:
+      image: huggingface/accelerate:gpu-nightly
+      options: --gpus all --shm-size "16gb"
+    runs-on:
+      group: aws-g6-12xlarge-plus
+    strategy:
+      fail-fast: false
+    steps:
+      - name: Install accelerate
+        run:
+          source activate accelerate;
+          git clone https://github.com/huggingface/accelerate;
+          cd accelerate;
+          git checkout ${{ github.sha }};
+          pip install -e .[testing];
+          cd ..
+
+      - name: Install skorch
+        run: |
+          source activate accelerate
+          git clone https://github.com/skorch-dev/skorch;
+          cd skorch;
+          git config --global --add safe.directory '*'
+          git checkout master && git pull
+          pip install .[testing]
+          pip install flaky
+
+      - name: Show installed libraries
+        run: |
+          source activate accelerate;
+          pip freeze
+
+      - name: Run skorch tests
+        working-directory: skorch/
+        run: |
+          source activate accelerate;
+          pytest -sv -k TestAccelerate

.github/workflows/stale.yml

@@ -0,0 +1,33 @@
+name: Stale Bot
+
+on:
+  schedule:
+    - cron: "0 15 * * *"
+  workflow_dispatch:
+
+jobs:
+  close_stale_issues:
+    name: Close Stale Issues
+    if: github.repository == 'huggingface/accelerate'
+    runs-on: ubuntu-latest
+    permissions:
+      issues: write
+      pull-requests: write
+    env:
+      GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+    steps:
+    - uses: actions/checkout@v4
+    
+    - name: Setup Python
+      uses: actions/setup-python@v5
+      with:
+        python-version: 3.9
+        cache: 'pip'
+        cache-dependency-path: 'setup.py'
+    
+    - name: Install requirements
+      run: |
+        pip install PyGithub
+    - name: Close stale issues
+      run: |
+        python utils/stale.py

.github/workflows/test.yml

@@ -1,17 +1,70 @@
 name: Run Tests
 
-on: [pull_request]
+on:
+  pull_request:
+    paths:
+      - "src/**"
+      - "tests/**"
+      - ".github/**"
+      - "examples/**"
+      - "setup.py"
+    types: [opened, synchronize, reopened]
+
+env:
+  HF_HOME: ~/hf_cache
+  TESTING_MOCKED_DATALOADERS: "1"
+  IS_GITHUB_CI: "1"
 
 jobs:
-  test:
+  run-tests:
     runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        pytorch-version: [
+          latest,
+          minimum,
+        ]
+        test-kind: [
+          test_prod,
+          test_core,
+          test_cli,
+          test_big_modeling,
+          test_deepspeed,
+          test_fsdp,
+          test_example_differences,
+          test_checkpoint_step,
+          test_checkpoint_epoch,
+          test_rest
+        ]
     steps:
-    - uses: actions/checkout@v2
-    - name: Set up Python 3.6
-      uses: actions/setup-python@v2
+    - uses: actions/checkout@v4
+    - name: Set up python 3.9
+      uses: actions/setup-python@v5
       with:
-        python-version: 3.6
-    - name: Install Python dependencies
-      run: pip install -e .[test]
+        python-version: 3.9
+        cache: 'pip'
+        cache-dependency-path: 'setup.py'
+    
+    - name: Install the library
+      run: |
+        if [[ ${{ matrix.test-kind }} = test_prod ]]; then pip install -e .[test_prod]; fi
+        if [[ ${{ matrix.test-kind }} != test_prod ]]; then pip install -e .[testing,test_trackers]; fi
+        if [[ ${{ matrix.test-kind }} = test_rest ]]; then pip uninstall comet_ml -y; fi
+        if [[ ${{ matrix.pytorch-version }} = minimum ]]; then pip install torchvision==0.18.1 torch==2.3.1; fi
+        pip install pytest-reportlog tabulate setuptools importlib_metadata
+
+    - name: Show installed libraries
+      run: |
+        pip freeze
+    
     - name: Run Tests
-      run: make test
+      env: 
+        PYTORCH_VERSION: ${{ matrix.pytorch-version }}
+      run: |
+        make ${{ matrix.test-kind }}
+
+    - name: Generate Report
+      if: always()
+      run: |
+        python utils/log_reports.py >> $GITHUB_STEP_SUMMARY

.github/workflows/test_imports.yml

@@ -0,0 +1,55 @@
+name: Run Import Tests
+
+on:
+  pull_request:
+    paths:
+      - "src/**"
+      - "tests/**"
+      - ".github/**"
+      - "examples/**"
+      - "setup.py"
+    types: [opened, synchronize, reopened]
+
+env:
+  HF_HOME: ~/hf_cache
+  TESTING_MOCKED_DATALOADERS: "1"
+  IS_GITHUB_CI: "1"
+
+jobs:
+  run-tests:
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        pytorch-version: [
+          latest,
+          minimum,
+        ]
+    steps:
+    - uses: actions/checkout@v4
+    - name: Set up python 3.9
+      uses: actions/setup-python@v5
+      with:
+        python-version: 3.9
+        cache: 'pip'
+        cache-dependency-path: 'setup.py'
+    
+    - name: Install the library
+      run: |
+        pip install -e .
+        pip install pytest-reportlog tabulate setuptools git+https://github.com/muellerzr/import-timer
+
+    - name: Show installed libraries
+      run: |
+        pip freeze
+    
+    - name: Run Import Tests
+      env: 
+        PYTORCH_VERSION: ${{ matrix.pytorch-version }}
+      run: |
+        pytest -sv tests/test_imports.py
+
+    - name: Generate Report
+      if: always()
+      run: |
+        python utils/log_reports.py >> $GITHUB_STEP_SUMMARY

.github/workflows/trufflehog.yml

@@ -0,0 +1,15 @@
+on:
+  push:
+
+name: Secret Leaks
+
+jobs:
+  trufflehog:
+    runs-on: ubuntu-latest
+    steps:
+    - name: Checkout code
+      uses: actions/checkout@v4
+      with:
+        fetch-depth: 0
+    - name: Secret Scanning
+      uses: trufflesecurity/trufflehog@main

.github/workflows/upload_pr_documentation.yml

@@ -0,0 +1,16 @@
+name: Upload PR Documentation
+
+on:
+  workflow_run:
+    workflows: ["Build PR Documentation"]
+    types:
+      - completed
+
+jobs:
+  build:
+    uses: huggingface/doc-builder/.github/workflows/upload_pr_documentation.yml@main
+    with:
+      package_name: accelerate
+    secrets:
+      hf_token: ${{ secrets.HF_DOC_BUILD_PUSH }}
+      comment_bot_token: ${{ secrets.COMMENT_BOT_TOKEN }}

.gitignore

@@ -130,3 +130,18 @@ dmypy.json
 
 # VSCode
 .vscode
+
+# IntelliJ
+.idea
+
+# Mac .DS_Store
+.DS_Store
+
+# More test things
+wandb
+
+# ruff
+.ruff_cache
+
+
+models

.pre-commit-config.yaml

@@ -0,0 +1,13 @@
+repos:
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.2.1
+    hooks:
+      - id: ruff
+        args:
+          - --fix
+      - id: ruff-format
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.5.0
+    hooks:
+      - id: check-merge-conflict
+      - id: check-yaml

CONTRIBUTING.md

@@ -0,0 +1,249 @@
+<!---
+Copyright 2022 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+-->
+
+# How to contribute to 🤗 Accelerate?
+
+Everyone is welcome to contribute, and we value everybody's contribution. Code
+is thus not the only way to help the community. Answering questions, helping
+others, reaching out and improving the documentations are immensely valuable to
+the community.
+
+It also helps us if you spread the word: reference the library from blog posts
+on the awesome projects it made possible, shout out on Twitter every time it has
+helped you, or simply star the repo to say "thank you".
+
+Whichever way you choose to contribute, please be mindful to respect our
+[code of conduct](https://github.com/huggingface/accelerate/blob/main/CODE_OF_CONDUCT.md).
+
+## You can contribute in so many ways!
+
+Some of the ways you can contribute to Accelerate:
+* Fixing outstanding issues with the existing code;
+* Contributing to the examples or to the documentation;
+* Submitting issues related to bugs or desired new features.
+
+## Submitting a new issue or feature request
+
+Do your best to follow these guidelines when submitting an issue or a feature
+request. It will make it easier for us to come back to you quickly and with good
+feedback.
+
+### Did you find a bug?
+
+The 🤗 Accelerate library is robust and reliable thanks to the users who notify us of
+the problems they encounter. So thank you for reporting an issue.
+
+First, we would really appreciate it if you could **make sure the bug was not
+already reported** (use the search bar on Github under Issues).
+
+Did not find it? :( So we can act quickly on it, please follow these steps:
+
+* Include your **OS type and version**, the versions of **Python** and **PyTorch**.
+* A short, self-contained, code snippet that allows us to reproduce the bug in
+  less than 30s;
+* Provide the with your Accelerate configuration (located by default in `~/.cache/huggingface/accelerate/default_config.yaml`)
+
+### Do you want a new feature?
+
+A good feature request addresses the following points:
+
+1. Motivation first:
+* Is it related to a problem/frustration with the library? If so, please explain
+  why. Providing a code snippet that demonstrates the problem is best.
+* Is it related to something you would need for a project? We'd love to hear
+  about it!
+* Is it something you worked on and think could benefit the community?
+  Awesome! Tell us what problem it solved for you.
+2. Write a *full paragraph* describing the feature;
+3. Provide a **code snippet** that demonstrates its future use;
+4. In case this is related to a paper, please attach a link;
+5. Attach any additional information (drawings, screenshots, etc.) you think may help.
+
+If your issue is well written we're already 80% of the way there by the time you
+post it.
+
+## Submitting a pull request (PR)
+
+Before writing code, we strongly advise you to search through the existing PRs or
+issues to make sure that nobody is already working on the same thing. If you are
+unsure, it is always a good idea to open an issue to get some feedback.
+
+You will need basic `git` proficiency to be able to contribute to
+🤗 Accelerate. `git` is not the easiest tool to use but it has the greatest
+manual. Type `git --help` in a shell and enjoy. If you prefer books, [Pro
+Git](https://git-scm.com/book/en/v2) is a very good reference.
+
+Follow these steps to start contributing:
+
+1. Fork the [repository](https://github.com/huggingface/accelerate) by
+   clicking on the 'Fork' button on the repository's page. This creates a copy of the code
+   under your GitHub user account.
+
+2. Clone your fork to your local disk, and add the base repository as a remote. The following command
+   assumes you have your public SSH key uploaded to GitHub. See the following guide for more
+   [information](https://docs.github.com/en/repositories/creating-and-managing-repositories/cloning-a-repository).
+
+   ```bash
+   $ git clone git@github.com:<your Github handle>/accelerate.git
+   $ cd accelerate
+   $ git remote add upstream https://github.com/huggingface/accelerate.git
+   ```
+
+3. Create a new branch to hold your development changes, and do this for every new PR you work on.
+
+   Start by synchronizing your `main` branch with the `upstream/main` branch (ore details in the [GitHub Docs](https://docs.github.com/en/github/collaborating-with-issues-and-pull-requests/syncing-a-fork)):
+
+   ```bash
+   $ git checkout main
+   $ git fetch upstream
+   $ git merge upstream/main
+   ```
+
+   Once your `main` branch is synchronized, create a new branch from it:
+
+   ```bash
+   $ git checkout -b a-descriptive-name-for-my-changes
+   ```
+
+   **Do not** work on the `main` branch.
+
+4. Set up a development environment by running the following command in a conda or a virtual environment you've created for working on this library:
+
+   ```bash
+   $ pip install -e ".[dev]"
+   ```
+   
+   This will install all testing and linting/code quality dependencies for the library (see `quality`, `test_dev`, 
+   `test_prod` targets in [`setup.py`](./setup.py)).
+
+   (If accelerate was already installed in the virtual environment, remove
+   it with `pip uninstall accelerate` before reinstalling it in editable
+   mode with the `-e` flag).
+
+   Alternatively, if you are using [Visual Studio Code](https://code.visualstudio.com/Download), the fastest way to get set up is by using
+   the provided Dev Container. Documentation on how to get started with dev containers is available [here](https://code.visualstudio.com/docs/remote/containers).
+
+5. Develop the features on your branch.
+
+   As you work on the features, you should make sure that the test suite
+   passes. You should run the tests impacted by your changes like this (see 
+   below an explanation regarding the environment variable):
+
+   ```bash
+   $ pytest tests/<TEST_TO_RUN>.py
+   ```
+   
+   > For the following commands leveraging the `make` utility, we recommend using the WSL system when running on
+   > Windows. More information [here](https://docs.microsoft.com/en-us/windows/wsl/about).
+
+   You can also run the full suite with the following command.
+
+   ```bash
+   $ make test
+   ```
+
+   `accelerate` relies on `ruff` to format its source code
+   consistently. After you make changes, apply automatic style corrections and code verifications
+   that can't be automated in one go with:
+
+   This target is also optimized to only work with files modified by the PR you're working on.
+
+   If you prefer to run the checks one after the other, the following command apply the
+   style corrections:
+
+   ```bash
+   $ make style
+   ```
+
+   `accelerate` also uses a few custom scripts to check for coding mistakes. Quality
+   control runs in CI, however you can also run the same checks with:
+
+   ```bash
+   $ make quality
+   ```
+
+   You can also set up [`pre-commit`](https://pre-commit.com/) to run these checks
+   automatically as Git commit hooks.
+
+   ```bash
+   $ pip install pre-commit
+   $ pre-commit install
+   ```
+
+   Once you're happy with your changes, add changed files using `git add` and
+   make a commit with `git commit` to record your changes locally:
+
+   ```bash
+   $ git add modified_file.py
+   $ git commit
+   ```
+
+   Please write [good commit messages](https://chris.beams.io/posts/git-commit/).
+
+   It is a good idea to sync your copy of the code with the original
+   repository regularly. This way you can quickly account for changes:
+
+   ```bash
+   $ git fetch upstream
+   $ git rebase upstream/main
+   ```
+
+   Push the changes to your account using:
+
+   ```bash
+   $ git push -u origin a-descriptive-name-for-my-changes
+   ```
+
+6. Once you are satisfied (**and the checklist below is happy too**), go to the
+   webpage of your fork on GitHub. Click on 'Pull request' to send your changes
+   to the project maintainers for review.
+
+7. It's ok if maintainers ask you for changes. It happens to core contributors
+   too! So everyone can see the changes in the Pull request, work in your local
+   branch and push the changes to your fork. They will automatically appear in
+   the pull request.
+
+
+### Checklist
+
+1. The title of your pull request should be a summary of its contribution;
+2. If your pull request addresses an issue, please mention the issue number in
+   the pull request description to make sure they are linked (and people
+   consulting the issue know you are working on it);
+3. To indicate a work in progress please prefix the title with `[WIP]`, or mark
+   the PR as a draft PR. These are useful to avoid duplicated work, and to differentiate
+   it from PRs ready to be merged;
+4. Make sure existing tests pass;
+5. Add high-coverage tests. No quality testing = no merge.
+
+See an example of a good PR here: https://github.com/huggingface/accelerate/pull/255
+
+### Tests
+
+An extensive test suite is included to test the library behavior and several examples. Library tests can be found in
+the [tests folder](https://github.com/huggingface/accelerate/tree/main/tests).
+
+We use `pytest` in order to run the tests. From the root of the
+repository, here's how to run tests with `pytest` for the library:
+
+```bash
+$ python -m pytest -sv ./tests
+```
+
+In fact, that's how `make test` is implemented (sans the `pip install` line)!
+
+You can specify a smaller set of tests in order to test only the feature
+you're working on.

Makefile

@@ -1,6 +1,6 @@
-.PHONY: quality style test docs
+.PHONY: quality style test docs utils
 
-check_dirs := tests src examples
+check_dirs := .
 
 # Check that source code meets quality standards
 
@@ -8,25 +8,94 @@ extra_quality_checks:
 	python utils/check_copies.py
 	python utils/check_dummies.py
 	python utils/check_repo.py
-	python utils/style_doc.py src/accelerate docs/source --max_len 119
+	doc-builder style src/accelerate docs/source --max_len 119
 
 # this target runs checks on all files
 quality:
-	black --check $(check_dirs)
-	isort --check-only $(check_dirs)
-	flake8 $(check_dirs)
-	python utils/style_doc.py src/accelerate docs/source --max_len 119 --check_only
+	ruff check $(check_dirs)
+	ruff format --check $(check_dirs)
+	doc-builder style src/accelerate docs/source --max_len 119 --check_only
 
 # Format source code automatically and check is there are any problems left that need manual fixing
 style:
-	black $(check_dirs)
-	isort $(check_dirs)
-	python utils/style_doc.py src/accelerate docs/source --max_len 119
+	ruff check $(check_dirs) --fix
+	ruff format $(check_dirs)
+	doc-builder style src/accelerate docs/source --max_len 119
 	
 # Run tests for the library
-test:
-	python -m pytest -n auto --dist=loadfile -s -v ./tests/
+test_core:
+	python -m pytest -s -v ./tests/ \
+	--ignore=./tests/test_big_modeling.py \
+	--ignore=./tests/test_modeling_utils.py \
+	--ignore=./tests/test_examples.py \
+	--ignore=./tests/test_cli.py \
+	--ignore=./tests/deepspeed \
+	--ignore=./tests/fsdp \
+	--ignore=./tests/tp \
+	$(if $(IS_GITHUB_CI),--report-log "$(PYTORCH_VERSION)_core.log",)
 
-# Check that docs can build
-docs:
-	cd docs && make html SPHINXOPTS="-W"
+test_cli:
+	python -m pytest -s -v ./tests/test_cli.py $(if $(IS_GITHUB_CI),--report-log "$(PYTORCH_VERSION)_cli.log",)
+
+test_big_modeling:
+	python -m pytest -s -v ./tests/test_big_modeling.py ./tests/test_modeling_utils.py $(if $(IS_GITHUB_CI),--report-log "$(PYTORCH_VERSION)_big_modeling.log",)
+
+test_deepspeed:
+	python -m pytest -s -v ./tests/deepspeed $(if $(IS_GITHUB_CI),--report-log "$(PYTORCH_VERSION)_deepspeed.log",)
+
+test_fsdp:
+	python -m pytest -s -v ./tests/fsdp $(if $(IS_GITHUB_CI),--report-log "$(PYTORCH_VERSION)_fsdp.log",)
+
+test_tp:
+	python -m pytest -s -v ./tests/tp $(if $(IS_GITHUB_CI),--report-log "$(PYTORCH_VERSION)_tp.log",)
+
+# Since the new version of pytest will *change* how things are collected, we need `deepspeed` to 
+# run after test_core and test_cli
+test:
+	$(MAKE) test_core
+	$(MAKE) test_cli
+	$(MAKE) test_big_modeling
+	$(MAKE) test_deepspeed
+	$(MAKE) test_fsdp
+	$(MAKE) test_tp
+
+test_examples:
+	python -m pytest -s -v ./tests/test_examples.py $(if $(IS_GITHUB_CI),--report-log "$(PYTORCH_VERSION)_examples.log",)
+
+# Broken down example tests for the CI runners
+test_integrations:
+	python -m pytest -s -v ./tests/deepspeed ./tests/fsdp ./tests/tp $(if $(IS_GITHUB_CI),--report-log "$(PYTORCH_VERSION)_integrations.log",)
+
+test_example_differences:
+	python -m pytest -s -v ./tests/test_examples.py::ExampleDifferenceTests $(if $(IS_GITHUB_CI),--report-log "$(PYTORCH_VERSION)_example_diff.log",)
+
+test_checkpoint_epoch:
+	python -m pytest -s -v ./tests/test_examples.py::FeatureExamplesTests -k "by_epoch" $(if $(IS_GITHUB_CI),--report-log "$(PYTORCH_VERSION)_checkpoint_epoch.log",)
+
+test_checkpoint_step:
+	python -m pytest -s -v ./tests/test_examples.py::FeatureExamplesTests -k "by_step" $(if $(IS_GITHUB_CI),--report-log "$(PYTORCH_VERSION)_checkpoint_step.log",)
+
+# Same as test but used to install only the base dependencies
+test_prod:
+	$(MAKE) test_core
+
+test_rest:
+	python -m pytest -s -v ./tests/test_examples.py::FeatureExamplesTests -k "not by_step and not by_epoch" $(if $(IS_GITHUB_CI),--report-log "$(PYTORCH_VERSION)_rest.log",)
+
+# For developers to prepare a release
+prepare_release:
+	rm -rf dist build
+	python setup.py bdist_wheel sdist
+
+# Make sure this is ran in a fresh venv of some form
+install_test_release:
+	pip uninstall accelerate -y
+	pip install -i https://testpypi.python.org/pypi --extra-index-url https://pypi.org/simple accelerate$(if $(version),==$(version),)
+
+# Run as `make target=testpypi upload_release`
+upload_release:
+	@if [ "$(target)" != "testpypi" ] && [ "$(target)" != "pypi" ]; then \
+		echo "Error: target must be either 'testpypi' or 'pypi'"; \
+		exit 1; \
+	fi
+	twine upload dist/* -r $(target)

README.md

@@ -16,34 +16,28 @@ limitations under the License.
 
 <p align="center">
     <br>
-    <img src="docs/source/imgs/accelerate_logo.png" width="400"/>
+    <img src="https://raw.githubusercontent.com/huggingface/accelerate/main/docs/source/imgs/accelerate_logo.png" width="400"/>
     <br>
 <p>
 
 <p align="center">
-    <!-- Uncomment when CircleCI is setup
-    <a href="https://circleci.com/gh/huggingface/accelerate">
-        <img alt="Build" src="https://img.shields.io/circleci/build/github/huggingface/transformers/master">
-    </a>
+    <!-- Uncomment when CircleCI is set up
+    <a href="https://circleci.com/gh/huggingface/accelerate"><img alt="Build" src="https://img.shields.io/circleci/build/github/huggingface/transformers/master"></a>
     -->
-    <a href="https://github.com/huggingface/accelerate/blob/master/LICENSE">
-        <img alt="License" src="https://img.shields.io/github/license/huggingface/accelerate.svg?color=blue">
-    </a>
-    <a href="https://huggingface.co/transformers/index.html">
-        <img alt="Documentation" src="https://img.shields.io/website/http/huggingface.co/transformers/index.html.svg?down_color=red&down_message=offline&up_message=online">
-    </a>
-    <a href="https://github.com/huggingface/accelerate/releases">
-        <img alt="GitHub release" src="https://img.shields.io/github/release/huggingface/accelerate.svg">
-    </a>
-    <a href="https://github.com/huggingface/accelerate/blob/master/CODE_OF_CONDUCT.md">
-        <img alt="Contributor Covenant" src="https://img.shields.io/badge/Contributor%20Covenant-v2.0%20adopted-ff69b4.svg">
-    </a>
+    <a href="https://github.com/huggingface/accelerate/blob/main/LICENSE"><img alt="License" src="https://img.shields.io/github/license/huggingface/accelerate.svg?color=blue"></a>
+    <a href="https://huggingface.co/docs/accelerate/index.html"><img alt="Documentation" src="https://img.shields.io/website/http/huggingface.co/docs/accelerate/index.html.svg?down_color=red&down_message=offline&up_message=online"></a>
+    <a href="https://github.com/huggingface/accelerate/releases"><img alt="GitHub release" src="https://img.shields.io/github/release/huggingface/accelerate.svg"></a>
+    <a href="https://github.com/huggingface/accelerate/blob/main/CODE_OF_CONDUCT.md"><img alt="Contributor Covenant" src="https://img.shields.io/badge/Contributor%20Covenant-v2.0%20adopted-ff69b4.svg"></a>
 </p>
 
 <h3 align="center">
 <p>Run your *raw* PyTorch training script on any kind of device
 </h3>
 
+<h3 align="center">
+    <a href="https://hf.co/course"><img src="https://raw.githubusercontent.com/huggingface/accelerate/main/docs/source/imgs/course_banner.png"></a>
+</h3>
+
 ## Easy to integrate
 
 🤗 Accelerate was created for PyTorch users who like to write the training loop of PyTorch models but are reluctant to write and maintain the boilerplate code needed to use multi-GPUs/TPU/fp16.
@@ -63,12 +57,12 @@ Here is an example:
 + device = accelerator.device
 
   model = torch.nn.Transformer().to(device)
-  optim = torch.optim.Adam(model.parameters())
+  optimizer = torch.optim.Adam(model.parameters())
 
   dataset = load_dataset('my_dataset')
   data = torch.utils.data.DataLoader(dataset, shuffle=True)
 
-+ model, optim, data = accelerator.prepare(model, optim, data)
++ model, optimizer, data = accelerator.prepare(model, optimizer, data)
 
   model.train()
   for epoch in range(10):
@@ -81,13 +75,13 @@ Here is an example:
           output = model(source)
           loss = F.cross_entropy(output, targets)
 
-+         accelerator.backward(loss)
 -         loss.backward()
++         accelerator.backward(loss)
 
           optimizer.step()
 ```
 
-As you can see in this example, by adding 5-lines to any standard PyTorch training script you can now run on any kind of single or distributed node setting (single CPU, single GPU, multi-GPUs and TPUs) as well as with or without mixed precision (fp16).
+As you can see in this example, by adding 5-lines to any standard PyTorch training script you can now run on any kind of single or distributed node setting (single CPU, single GPU, multi-GPUs and TPUs) as well as with or without mixed precision (fp8, fp16, bf16).
 
 In particular, the same code can then be run without modification on your local machine for debugging or your training environment.
 
@@ -99,17 +93,17 @@ In particular, the same code can then be run without modification on your local
   from datasets import load_dataset
 + from accelerate import Accelerator
 
-+ accelerator = Accelerator()
 - device = 'cpu'
++ accelerator = Accelerator()
 
-+ model = torch.nn.Transformer()
 - model = torch.nn.Transformer().to(device)
-  optim = torch.optim.Adam(model.parameters())
++ model = torch.nn.Transformer()
+  optimizer = torch.optim.Adam(model.parameters())
 
   dataset = load_dataset('my_dataset')
   data = torch.utils.data.DataLoader(dataset, shuffle=True)
 
-+ model, optim, data = accelerator.prepare(model, optim, data)
++ model, optimizer, data = accelerator.prepare(model, optimizer, data)
 
   model.train()
   for epoch in range(10):
@@ -122,15 +116,17 @@ In particular, the same code can then be run without modification on your local
           output = model(source)
           loss = F.cross_entropy(output, targets)
 
-+         accelerator.backward(loss)
 -         loss.backward()
++         accelerator.backward(loss)
 
           optimizer.step()
 ```
 
+Want to learn more? Check out the [documentation](https://huggingface.co/docs/accelerate) or have a look at our [examples](https://github.com/huggingface/accelerate/tree/main/examples).
+
 ## Launching script
 
-🤗 Accelerate also provides an optional CLI tool that allows you to quickly configure and test your training environment before launching the scripts. No need to remember how to use `torch.distributed.launch` or to write a specific launcher for TPU training!
+🤗 Accelerate also provides an optional CLI tool that allows you to quickly configure and test your training environment before launching the scripts. No need to remember how to use `torch.distributed.run` or to write a specific launcher for TPU training!
 On your machine(s) just run:
 
 ```bash
@@ -149,17 +145,98 @@ For instance, here is how you would run the GLUE example on the MRPC task (from
 accelerate launch examples/nlp_example.py
 ```
 
+This CLI tool is **optional**, and you can still use `python my_script.py` or `python -m torchrun my_script.py` at your convenience.
+
+You can also directly pass in the arguments you would to `torchrun` as arguments to `accelerate launch` if you wish to not run` accelerate config`.
+
+For example, here is how to launch on two GPUs:
+
+```bash
+accelerate launch --multi_gpu --num_processes 2 examples/nlp_example.py
+```
+
+To learn more, check the CLI documentation available [here](https://huggingface.co/docs/accelerate/package_reference/cli).
+
+Or view the configuration zoo [here](https://github.com/huggingface/accelerate/blob/main/examples/config_yaml_templates/)
+
+## Launching multi-CPU run using MPI
+
+🤗 Here is another way to launch multi-CPU run using MPI. You can learn how to install Open MPI on [this page](https://www.open-mpi.org/faq/?category=building#easy-build). You can use Intel MPI or MVAPICH as well.
+Once you have MPI setup on your cluster, just run:
+```bash
+accelerate config
+```
+Answer the questions that are asked, selecting to run using multi-CPU, and answer "yes" when asked if you want accelerate to launch mpirun.
+Then, use `accelerate launch` with your script like:
+```bash
+accelerate launch examples/nlp_example.py
+```
+Alternatively, you can use mpirun directly, without using the CLI like:
+```bash
+mpirun -np 2 python examples/nlp_example.py
+```
+
+## Launching training using DeepSpeed
+
+🤗 Accelerate supports training on single/multiple GPUs using DeepSpeed. To use it, you don't need to change anything in your training code; you can set everything using just `accelerate config`. However, if you desire to tweak your DeepSpeed related args from your Python script, we provide you the `DeepSpeedPlugin`.
+
+```python
+from accelerate import Accelerator, DeepSpeedPlugin
+
+# deepspeed needs to know your gradient accumulation steps beforehand, so don't forget to pass it
+# Remember you still need to do gradient accumulation by yourself, just like you would have done without deepspeed
+deepspeed_plugin = DeepSpeedPlugin(zero_stage=2, gradient_accumulation_steps=2)
+accelerator = Accelerator(mixed_precision='fp16', deepspeed_plugin=deepspeed_plugin)
+
+# How to save your 🤗 Transformer?
+accelerator.wait_for_everyone()
+unwrapped_model = accelerator.unwrap_model(model)
+unwrapped_model.save_pretrained(save_dir, save_function=accelerator.save, state_dict=accelerator.get_state_dict(model))
+```
+
+Note: DeepSpeed support is experimental for now. In case you get into some problem, please open an issue.
+
+## Launching your training from a notebook
+
+🤗 Accelerate also provides a `notebook_launcher` function you can use in a notebook to launch a distributed training. This is especially useful for Colab or Kaggle notebooks with a TPU backend. Just define your training loop in a `training_function` then in your last cell, add:
+
+```python
+from accelerate import notebook_launcher
+
+notebook_launcher(training_function)
+```
+
+An example can be found in [this notebook](https://github.com/huggingface/notebooks/blob/main/examples/accelerate_examples/simple_nlp_example.ipynb). [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/notebooks/blob/main/examples/accelerate_examples/simple_nlp_example.ipynb)
+
 ## Why should I use 🤗 Accelerate?
 
-You should use 🤗 Accelerate when you want to easily run your training scripts in a distributed environment without having to renounce full control over your training loop. This is not a high-level framework above PyTorch, just a thin wrapper so you don't have to learn a new library, In fact the whole API of 🤗 Accelerate is in one class, the `Accelerator` object.
+You should use 🤗 Accelerate when you want to easily run your training scripts in a distributed environment without having to renounce full control over your training loop. This is not a high-level framework above PyTorch, just a thin wrapper so you don't have to learn a new library. In fact, the whole API of 🤗 Accelerate is in one class, the `Accelerator` object.
 
 ## Why shouldn't I use 🤗 Accelerate?
 
 You shouldn't use 🤗 Accelerate if you don't want to write a training loop yourself. There are plenty of high-level libraries above PyTorch that will offer you that, 🤗 Accelerate is not one of them.
 
+## Frameworks using 🤗 Accelerate
+
+If you like the simplicity of 🤗 Accelerate but would prefer a higher-level abstraction around its capabilities, some frameworks and libraries that are built on top of 🤗 Accelerate are listed below:
+
+* [Amphion](https://github.com/open-mmlab/Amphion) is a toolkit for Audio, Music, and Speech Generation. Its purpose is to support reproducible research and help junior researchers and engineers get started in the field of audio, music, and speech generation research and development.
+* [Animus](https://github.com/Scitator/animus) is a minimalistic framework to run machine learning experiments. Animus highlights common "breakpoints" in ML experiments and provides a unified interface for them within [IExperiment](https://github.com/Scitator/animus/blob/main/animus/core.py#L76).
+* [Catalyst](https://github.com/catalyst-team/catalyst#getting-started) is a PyTorch framework for Deep Learning Research and Development. It focuses on reproducibility, rapid experimentation, and codebase reuse so you can create something new rather than write yet another train loop. Catalyst provides a [Runner](https://catalyst-team.github.io/catalyst/api/core.html#runner) to connect all parts of the experiment: hardware backend, data transformations, model training, and inference logic.
+* [fastai](https://github.com/fastai/fastai#installing) is a PyTorch framework for Deep Learning that simplifies training fast and accurate neural nets using modern best practices. fastai provides a [Learner](https://docs.fast.ai/learner.html#Learner) to handle the training, fine-tuning, and inference of deep learning algorithms.
+* [Finetuner](https://github.com/jina-ai/finetuner) is a service that enables models to create higher-quality embeddings for semantic search, visual similarity search, cross-modal text<->image search, recommendation systems, clustering, duplication detection, anomaly detection, or other uses.
+* [InvokeAI](https://github.com/invoke-ai/InvokeAI) is a creative engine for Stable Diffusion models, offering industry-leading WebUI, terminal usage support, and serves as the foundation for many commercial products.
+* [Kornia](https://kornia.readthedocs.io/en/latest/get-started/introduction.html) is a differentiable library that allows classical computer vision to be integrated into deep learning models. Kornia provides a [Trainer](https://kornia.readthedocs.io/en/latest/x.html#kornia.x.Trainer) with the specific purpose to train and fine-tune the supported deep learning algorithms within the library.
+* [Open Assistant](https://projects.laion.ai/Open-Assistant/) is a chat-based assistant that understands tasks, can interact with their party systems, and retrieve information dynamically to do so. 
+* [pytorch-accelerated](https://github.com/Chris-hughes10/pytorch-accelerated) is a lightweight training library, with a streamlined feature set centered around a general-purpose [Trainer](https://pytorch-accelerated.readthedocs.io/en/latest/trainer.html), that places a huge emphasis on simplicity and transparency; enabling users to understand exactly what is going on under the hood, but without having to write and maintain the boilerplate themselves!
+* [Stable Diffusion web UI](https://github.com/AUTOMATIC1111/stable-diffusion-webui) is an open-source browser-based easy-to-use interface based on the Gradio library for Stable Diffusion.
+* [torchkeras](https://github.com/lyhue1991/torchkeras) is a simple tool for training pytorch model just in a keras style, a dynamic and beautiful plot is provided in notebook to monitor your loss or metric.
+* [transformers](https://github.com/huggingface/transformers) as a tool for helping train state-of-the-art machine learning models in PyTorch, Tensorflow, and JAX. (Accelerate is the backend for the PyTorch side).
+
+
 ## Installation
 
-This repository is tested on Python 3.6+ and PyTorch 1.4.0+
+This repository is tested on Python 3.8+ and PyTorch 1.10.0+
 
 You should install 🤗 Accelerate in a [virtual environment](https://docs.python.org/3/library/venv.html). If you're unfamiliar with Python virtual environments, check out the [user guide](https://packaging.python.org/guides/installing-using-pip-and-virtual-environments/).
 
@@ -174,8 +251,27 @@ pip install accelerate
 ## Supported integrations
 
 - CPU only
+- multi-CPU on one node (machine)
+- multi-CPU on several nodes (machines)
 - single GPU
 - multi-GPU on one node (machine)
 - multi-GPU on several nodes (machines)
 - TPU
-- FP16 with native AMP (apex on the roadmap)
+- FP16/BFloat16 mixed precision
+- FP8 mixed precision with [Transformer Engine](https://github.com/NVIDIA/TransformerEngine) or [MS-AMP](https://github.com/Azure/MS-AMP/)
+- DeepSpeed support (Experimental)
+- PyTorch Fully Sharded Data Parallel (FSDP) support (Experimental)
+- Megatron-LM support (Experimental)
+
+## Citing 🤗 Accelerate
+
+If you use 🤗 Accelerate in your publication, please cite it by using the following BibTeX entry.
+
+```bibtex
+@Misc{accelerate,
+  title =        {Accelerate: Training and inference at scale made simple, efficient and adaptable.},
+  author =       {Sylvain Gugger and Lysandre Debut and Thomas Wolf and Philipp Schmid and Zachary Mueller and Sourab Mangrulkar and Marc Sun and Benjamin Bossan},
+  howpublished = {\url{https://github.com/huggingface/accelerate}},
+  year =         {2022}
+}
+```

benchmarks/README.md

@@ -0,0 +1,5 @@
+# Benchmarks
+
+The folders below contain suites to test various functionalities in Accelerate.
+
+See their relevant README.md's for more information.

benchmarks/big_model_inference/README.md

@@ -0,0 +1,46 @@
+# Big model inference benchmarks
+
+Running inference with Accelerate on big models.
+
+## Setup
+
+These benchmarks use the `transformers` library:
+
+```bash
+pip install transformers
+```
+
+To reproduce or test a new setup, run
+
+```py
+python big_model_inference.py model_name
+```
+
+This script supports `gpt-j-6b`, `gpt-neox`, `opt` (30B version) and `T0pp` out of the box, but you can specify any valid checkpoint for `model_name`.
+
+To force a different `torch_dtype` than the one in the config: `--torch_dtype xxx`.
+
+If you get an error linked to disk offload, you need to add the option `--disk-offload`
+
+## Results
+
+On a setup with two Titan RTXs (24GB of RAM) and 32GB of RAM, we get the following benchmarks (T0pp does not run in float16, which is why it's not included).
+
+| Model | Model load time | Generation time | dtype | GPU 0 use | GPU 1 use | CPU use | Disk offload |
+|:-----:|:---------------:|:---------------:|:-----:|:---------:|:---------:|:-------:|:------------:|
+| GPT-J-6B | 8.7s | 0.05s per token | float16 | 11.7GB | 0GB | 0GB | no |
+| GPT-J-6B | 12.4s | 0.06s per token | float32 | 21.9GB | 1.5GB | 0GB | no |
+| GPT-Neo-X-20B | 30.9s | 0.08s per token | float16 | 21.5GB | 18GB | 0GB | no |
+| GPT-Neo-X-20B | 78.2s | 10.72s per token | float32 | 20.3GB | 22.7 GB | 24.4GB | yes |
+| T0pp (11B) | 29.4s | 0.05s per token | float32 | 21.1GB | 21.3GB | 0GB | no |
+| OPT-30B | 34.5s | 2.37s per token | float16 | 20.7GB | 22.3GB | 14.1GB | no |
+| OPT-30B | 112.3s | 33.9s per token | float32 | 20.2GB | 21.2GB | 23.5GB | yes |
+
+Note on the results:
+- using two GPUs instead of one does not slow down generation
+- using CPU offload slows down a bit (see OPT-30b)
+- using disk offload slows down a lot (need to implement prefetching)
+
+You will also note that Accelerate does not use anymore GPU and CPU RAM than necessary:
+- peak GPU memory is exactly the size of the model put on a given GPU
+- peak CPU memory is either the size of the biggest checkpoint shard or the part of the model offloaded on CPU, whichever is bigger.

benchmarks/big_model_inference/big_model_inference.py

@@ -0,0 +1,143 @@
+# Copyright 2022 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import time
+
+import torch
+import transformers
+from measures_util import end_measure, log_measures, start_measure
+from transformers import AutoConfig, AutoModelForCausalLM, AutoModelForSeq2SeqLM, AutoTokenizer
+
+from accelerate.utils import compute_module_sizes
+
+
+DEFAULT_MODELS = {
+    "gpt-j-6b": {"is_causal": True, "model": "sgugger/sharded-gpt-j-6B", "tokenizer": "EleutherAI/gpt-j-6B"},
+    "gpt-neox": {"is_causal": True, "model": "EleutherAI/gpt-neox-20b"},
+    "opt": {"is_causal": True, "model": "facebook/opt-30b"},
+    "T0pp": {"is_causal": False, "model": "bigscience/T0pp", "model_revision": "sharded"},
+}
+
+PROMPTS = [
+    "Hello, my name is",
+    "Are unicorns real? Unicorns are",
+    "For the first time in several years,",
+    "My name is Julien and I am",
+    "The goal of life is",
+    "Whenever I'm sad, I like to",
+]
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Run and time generations on a big model using Accelerate.")
+    parser.add_argument("model_name", type=str, default=None, help="The name of the model to try.")
+    parser.add_argument(
+        "--tokenizer_name", type=str, default=None, help="The name of the tokenizer (if different from the model."
+    )
+    parser.add_argument("--is_causal", type=bool, default=None, help="Whether or not the model is causal.")
+    parser.add_argument(
+        "--model_revision", type=str, default=None, help="The revision to use for the model checkpoint."
+    )
+    parser.add_argument("--torch_dtype", type=str, default=None, help="The dtype for the model.")
+    parser.add_argument("--disk_offload", action="store_true")
+
+    args = parser.parse_args()
+
+    # Sanitize args
+    if args.model_name in DEFAULT_MODELS:
+        defaults = DEFAULT_MODELS[args.model_name]
+        args.model_name = defaults["model"]
+        if args.tokenizer_name is None:
+            args.tokenizer_name = defaults.get("tokenizer", args.model_name)
+        if args.is_causal is None:
+            args.is_causal = defaults["is_causal"]
+        if args.model_revision is None:
+            args.model_revision = defaults.get("model_revision", "main")
+
+    if args.is_causal is None:
+        raise ValueError("Could not infer the default for `--is_causal`, pass either True or False for it.")
+    if args.tokenizer_name is None:
+        args.tokenizer_name = args.model_name
+    if args.model_revision is None:
+        args.model_revision = "main"
+
+    return args
+
+
+def main():
+    transformers.utils.logging.set_verbosity_error()
+    args = parse_args()
+
+    if args.torch_dtype is None:
+        config = AutoConfig.from_pretrained(args.model_name)
+        torch_dtype = getattr(config, "torch_dtype", torch.float32)
+    else:
+        torch_dtype = getattr(torch, args.torch_dtype)
+    model_cls = AutoModelForCausalLM if args.is_causal else AutoModelForSeq2SeqLM
+    kwargs = {
+        "torch_dtype": torch_dtype,
+        "revision": args.model_revision,
+    }
+    if args.disk_offload:
+        kwargs["offload_folder"] = "tmp_offload"
+        kwargs["offload_state_dict"] = True
+
+    start_measures = start_measure()
+    model = model_cls.from_pretrained(args.model_name, device_map="auto", **kwargs)
+    end_measures = end_measure(start_measures)
+    log_measures(end_measures, "Model loading")
+
+    module_sizes = compute_module_sizes(model)
+    device_size = {v: 0 for v in model.hf_device_map.values()}
+    for module, device in model.hf_device_map.items():
+        device_size[device] += module_sizes[module]
+    message = "\n".join([f"- {device}: {size // 2**20}MiB" for device, size in device_size.items()])
+    print(f"\nTheoretical use:\n{message}")
+
+    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_name)
+
+    start_measures = start_measure()
+    generation_times = []
+    gen_tokens = []
+    texts_outs = []
+    for prompt in PROMPTS:
+        inputs = tokenizer(prompt, return_tensors="pt").to(0)
+        tokens = inputs["input_ids"][0].tolist()
+        before_generate = time.time()
+        outputs = model.generate(inputs["input_ids"])
+        after_generate = time.time()
+        outputs = outputs[0].tolist()
+        num_gen_tokens = len(outputs) if outputs[: len(tokens)] != tokens else len(outputs) - len(tokens)
+        generation_time = after_generate - before_generate
+
+        text_out = tokenizer.decode(outputs, skip_special_tokens=True)
+        texts_outs.append(text_out)
+        generation_times.append(generation_time)
+        gen_tokens.append(num_gen_tokens)
+        print(f"Prompt: {prompt}\nGeneration {text_out}\nIn {generation_time:.2f}s for {num_gen_tokens} tokens\n")
+
+    end_measures = end_measure(start_measures)
+    log_measures(end_measures, "Model generation")
+
+    generation_times_per_token = [gen / tok for gen, tok in zip(generation_times, gen_tokens)]
+    avg_gen = sum(generation_times_per_token) / len(generation_times)
+    print(f"Average time of generation per token: {avg_gen:.2f}s")
+    print(f"First generation (avg time per token): {generation_times_per_token[0]:.2f}s")
+    avg_gen = sum(generation_times_per_token[1:]) / (len(generation_times_per_token) - 1)
+    print(f"Average time of generation per token (excluding the first): {avg_gen:.2f}s")
+
+
+if __name__ == "__main__":
+    main()

benchmarks/big_model_inference/measures_util.py

@@ -0,0 +1,104 @@
+# Copyright 2023 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import gc
+import threading
+import time
+
+import psutil
+import torch
+
+from accelerate.test_utils.testing import get_backend
+
+
+torch_device_type, _, _ = get_backend()
+torch_accelerator_module = getattr(torch, torch_device_type, torch.cuda)
+
+
+class PeakCPUMemory:
+    def __init__(self):
+        self.process = psutil.Process()
+        self.peak_monitoring = False
+
+    def peak_monitor(self):
+        self.cpu_memory_peak = -1
+
+        while True:
+            self.cpu_memory_peak = max(self.process.memory_info().rss, self.cpu_memory_peak)
+
+            # can't sleep or will not catch the peak right (this comment is here on purpose)
+            if not self.peak_monitoring:
+                break
+
+    def start(self):
+        self.peak_monitoring = True
+        self.thread = threading.Thread(target=self.peak_monitor)
+        self.thread.daemon = True
+        self.thread.start()
+
+    def stop(self):
+        self.peak_monitoring = False
+        self.thread.join()
+        return self.cpu_memory_peak
+
+
+cpu_peak_tracker = PeakCPUMemory()
+
+
+def start_measure():
+    # Time
+    measures = {"time": time.time()}
+
+    gc.collect()
+    torch_accelerator_module.empty_cache()
+
+    # CPU mem
+    measures["cpu"] = psutil.Process().memory_info().rss
+    cpu_peak_tracker.start()
+
+    # GPU mem
+    for i in range(torch_accelerator_module.device_count()):
+        measures[str(i)] = torch_accelerator_module.memory_allocated(i)
+    torch_accelerator_module.reset_peak_memory_stats()
+
+    return measures
+
+
+def end_measure(start_measures):
+    # Time
+    measures = {"time": time.time() - start_measures["time"]}
+
+    gc.collect()
+    torch_accelerator_module.empty_cache()
+
+    # CPU mem
+    measures["cpu"] = (psutil.Process().memory_info().rss - start_measures["cpu"]) / 2**20
+    measures["cpu-peak"] = (cpu_peak_tracker.stop() - start_measures["cpu"]) / 2**20
+
+    # GPU mem
+    for i in range(torch_accelerator_module.device_count()):
+        measures[str(i)] = (torch_accelerator_module.memory_allocated(i) - start_measures[str(i)]) / 2**20
+        measures[f"{i}-peak"] = (torch_accelerator_module.max_memory_allocated(i) - start_measures[str(i)]) / 2**20
+
+    return measures
+
+
+def log_measures(measures, description):
+    print(f"{description}:")
+    print(f"- Time: {measures['time']:.2f}s")
+    for i in range(torch_accelerator_module.device_count()):
+        print(f"- {torch_device_type} {i} allocated: {measures[str(i)]:.2f}MiB")
+        peak = measures[f"{i}-peak"]
+        print(f"- {torch_device_type} {i} peak: {peak:.2f}MiB")
+    print(f"- CPU RAM allocated: {measures['cpu']:.2f}MiB")
+    print(f"- CPU RAM peak: {measures['cpu-peak']:.2f}MiB")

benchmarks/fp8/ms_amp/Dockerfile

@@ -0,0 +1,12 @@
+FROM ghcr.io/azure/msamp
+
+RUN pip install transformers evaluate datasets
+RUN git clone https://github.com/huggingface/accelerate
+
+RUN cd accelerate && \
+    pip install -e . && \
+    cd benchmarks/fp8
+
+CMD ["bash"]
+
+

benchmarks/fp8/ms_amp/ddp.py

@@ -0,0 +1,123 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+This script tests to ensure that `accelerate` performs at the same level as raw `MS-AMP`.
+
+This particular script verifies this for DDP training.
+"""
+
+import evaluate
+import msamp
+import torch
+from fp8_utils import evaluate_model, get_training_utilities
+from torch.nn.parallel import DistributedDataParallel as DDP
+
+from accelerate import Accelerator
+from accelerate.state import AcceleratorState
+from accelerate.utils import FP8RecipeKwargs, get_grad_scaler, set_seed
+
+
+MODEL_NAME = "bert-base-cased"
+METRIC = evaluate.load("glue", "mrpc")
+
+
+def train_baseline(opt_level="O2"):
+    set_seed(42)
+    scaler = get_grad_scaler()
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = get_training_utilities(MODEL_NAME)
+    accelerator = Accelerator()
+    device = accelerator.device
+
+    model, optimizer = msamp.initialize(model, optimizer, opt_level=opt_level)
+
+    model.to(device)
+
+    # Convert the model to DDP
+    device_ids, output_device = [accelerator.local_process_index], accelerator.local_process_index
+    model = DDP(model, device_ids=device_ids, output_device=output_device)
+
+    base_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)
+    model.train()
+
+    for i, batch in enumerate(train_dataloader):
+        with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+            outputs = model(**batch)
+            loss = outputs.loss
+        scaler.scale(loss).backward()
+        optimizer.step()
+        optimizer.zero_grad()
+        lr_scheduler.step()
+
+    trained_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)
+
+    assert trained_model_results["accuracy"] > base_model_results["accuracy"], (
+        f"Accuracy should be higher for the trained model: {trained_model_results['accuracy']} > {base_model_results['accuracy']}"
+    )
+    assert trained_model_results["f1"] > base_model_results["f1"], (
+        f"F1 score should be higher for the trained model: {trained_model_results['f1']} > {base_model_results['f1']}"
+    )
+
+    return base_model_results, trained_model_results
+
+
+def train_integration(opt_level="O2"):
+    kwargs_handlers = [FP8RecipeKwargs(backend="msamp", opt_level=opt_level)]
+    AcceleratorState()._reset_state(True)
+    accelerator = Accelerator(mixed_precision="fp8", kwargs_handlers=kwargs_handlers)
+    set_seed(42)
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = get_training_utilities(
+        MODEL_NAME, accelerator=accelerator
+    )
+
+    model, optimizer = accelerator.prepare(model, optimizer)
+    base_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)
+    model.train()
+    for i, batch in enumerate(train_dataloader):
+        with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+            outputs = model(**batch)
+            loss = outputs.loss
+        accelerator.backward(loss)
+        optimizer.step()
+        optimizer.zero_grad()
+        lr_scheduler.step()
+
+    trained_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)
+
+    assert trained_model_results["accuracy"] > base_model_results["accuracy"], (
+        f"Accuracy should be higher for the trained model: {trained_model_results['accuracy']} > {base_model_results['accuracy']}"
+    )
+    assert trained_model_results["f1"] > base_model_results["f1"], (
+        f"F1 score should be higher for the trained model: {trained_model_results['f1']} > {base_model_results['f1']}"
+    )
+
+    return base_model_results, trained_model_results
+
+
+if __name__ == "__main__":
+    for opt_level in ["O1", "O2"]:
+        baseline_not_trained, baseline_trained = train_baseline(opt_level)
+        accelerator_not_trained, accelerator_trained = train_integration(opt_level)
+        assert baseline_not_trained["accuracy"] == accelerator_not_trained["accuracy"], (
+            f"Accuracy not the same for untrained baseline and accelerator using opt_level={opt_level}: {baseline_not_trained['accuracy']} == {accelerator_not_trained['accuracy']}"
+        )
+        assert baseline_not_trained["f1"] == accelerator_not_trained["f1"], (
+            f"F1 not the same for untrained baseline and accelerator using opt_level={opt_level}: {baseline_not_trained['f1']} == {accelerator_not_trained['f1']}"
+        )
+        assert baseline_trained["accuracy"] == accelerator_trained["accuracy"], (
+            f"Accuracy not the same for trained baseline and accelerator using opt_level={opt_level}: {baseline_trained['accuracy']} == {accelerator_trained['accuracy']}"
+        )
+        assert baseline_trained["f1"] == accelerator_trained["f1"], (
+            f"F1 not the same for trained baseline and accelerator using opt_level={opt_level}: {baseline_trained['f1']} == {accelerator_trained['f1']}"
+        )

benchmarks/fp8/ms_amp/distrib_deepspeed.py

@@ -0,0 +1,161 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+This script tests to ensure that `accelerate` performs at the same level as raw `MS-AMP`.
+
+This particular script verifies this for DeepSpeed training.
+
+NOTE: MS-AMP does *not* support ZeRO-3.
+"""
+
+# import msamp.deepspeed as msamp_deepspeed
+import evaluate
+import torch
+from fp8_utils import evaluate_model, get_training_utilities
+from msamp import deepspeed as msamp_deepspeed
+
+from accelerate import Accelerator, DeepSpeedPlugin
+from accelerate.state import AcceleratorState
+from accelerate.utils import set_seed
+
+
+MODEL_NAME = "bert-base-cased"
+METRIC = evaluate.load("glue", "mrpc")
+
+
+def train_baseline(zero_stage: int = 1, opt_level: str = "O1"):
+    set_seed(42)
+    accelerator = Accelerator()
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = get_training_utilities(
+        MODEL_NAME, accelerator=accelerator
+    )
+
+    import numpy as np
+
+    config = {
+        "train_batch_size": 32,
+        "train_micro_batch_size_per_gpu": 16,
+        "gradient_accumulation_steps": 1,
+        "zero_optimization": {
+            "stage": zero_stage,
+            "offload_optimizer": {"device": "none", "nvme_path": None},
+            "offload_param": {"device": "none", "nvme_path": None},
+        },
+        "gradient_clipping": 1.0,
+        "steps_per_print": np.inf,
+        "bf16": {"enabled": True},
+        "fp16": {"enabled": False},
+        "zero_allow_untested_optimizer": True,
+        "msamp": {
+            "enabled": True,
+            "opt_level": opt_level,
+        },
+    }
+    (
+        model,
+        optimizer,
+        _,
+        _,
+    ) = msamp_deepspeed.initialize(
+        model=model,
+        optimizer=optimizer,
+        config_params=config,
+    )
+
+    base_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)
+    model.train()
+
+    for _ in range(2):
+        for batch in train_dataloader:
+            outputs = model(**batch)
+            loss = outputs.loss
+            model.backward(loss)
+            model.step()
+            for _ in range(accelerator.num_processes):
+                lr_scheduler.step()
+
+    trained_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)
+    model.destroy()
+    torch.cuda.empty_cache()
+    AcceleratorState()._reset_state(True)
+    assert trained_model_results["accuracy"] > base_model_results["accuracy"], (
+        f"Accuracy should be higher for the trained model: {trained_model_results['accuracy']} > {base_model_results['accuracy']}"
+    )
+    assert trained_model_results["f1"] > base_model_results["f1"], (
+        f"F1 score should be higher for the trained model: {trained_model_results['f1']} > {base_model_results['f1']}"
+    )
+
+    return base_model_results, trained_model_results
+
+
+def train_integration(zero_stage: int = 1, opt_level: str = "O1"):
+    set_seed(42)
+    deepspeed_plugin = DeepSpeedPlugin(
+        zero_stage=zero_stage,
+        enable_msamp=True,
+        msamp_opt_level=opt_level,
+    )
+    accelerator = Accelerator(mixed_precision="fp8", deepspeed_plugin=deepspeed_plugin)
+    accelerator.state.deepspeed_plugin.deepspeed_config["train_micro_batch_size_per_gpu"] = 16
+
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = get_training_utilities(
+        MODEL_NAME, accelerator=accelerator
+    )
+
+    model, optimizer, lr_scheduler = accelerator.prepare(model, optimizer, lr_scheduler)
+    base_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)
+    model.train()
+    for _ in range(2):
+        for batch in train_dataloader:
+            outputs = model(**batch)
+            loss = outputs.loss
+            accelerator.backward(loss)
+            optimizer.step()
+            lr_scheduler.step()
+            optimizer.zero_grad()
+
+    trained_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)
+    model.destroy()
+    torch.cuda.empty_cache()
+    assert trained_model_results["accuracy"] > base_model_results["accuracy"], (
+        f"Accuracy should be higher for the trained model: {trained_model_results['accuracy']} > {base_model_results['accuracy']}"
+    )
+    assert trained_model_results["f1"] > base_model_results["f1"], (
+        f"F1 score should be higher for the trained model: {trained_model_results['f1']} > {base_model_results['f1']}"
+    )
+
+    AcceleratorState()._reset_state(True)
+    return base_model_results, trained_model_results
+
+
+if __name__ == "__main__":
+    for zero_stage in [1, 2]:
+        for opt_level in ["O1", "O2", "O3"]:
+            baseline_not_trained, baseline_trained = train_baseline(zero_stage, opt_level)
+            accelerator_not_trained, accelerator_trained = train_integration(zero_stage, opt_level)
+            assert baseline_not_trained["accuracy"] == accelerator_not_trained["accuracy"], (
+                f"ZERO stage {zero_stage}, opt_level={opt_level}:\nAccuracy should be the same for the baseline and accelerator: {baseline_not_trained['accuracy']} == {accelerator_not_trained['accuracy']}"
+            )
+            assert baseline_not_trained["f1"] == accelerator_not_trained["f1"], (
+                f"ZERO stage {zero_stage}, opt_level={opt_level}:\nF1 score should be the same for the baseline and accelerator: {baseline_not_trained['f1']} == {accelerator_not_trained['f1']}"
+            )
+            assert baseline_trained["accuracy"] == accelerator_trained["accuracy"], (
+                f"ZERO stage {zero_stage}, opt_level={opt_level}:\nAccuracy should be the same for the baseline and accelerator: {baseline_trained['accuracy']} == {accelerator_trained['accuracy']}"
+            )
+            assert baseline_trained["f1"] == accelerator_trained["f1"], (
+                f"ZERO stage {zero_stage}, opt_level={opt_level}:\nF1 score should be the same for the baseline and accelerator: {baseline_trained['f1']} == {accelerator_trained['f1']}"
+            )
+
+    torch.distributed.destroy_process_group()

benchmarks/fp8/ms_amp/fp8_utils.py

@@ -0,0 +1,118 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import torch
+
+
+def get_dataloaders(model_name: str, batch_size: int = 16):
+    from datasets import load_dataset
+    from torch.utils.data import DataLoader
+    from transformers import AutoTokenizer
+
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    datasets = load_dataset("glue", "mrpc")
+
+    def tokenize_function(examples):
+        # max_length=None => use the model max length (it's actually the default)
+        outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None)
+        return outputs
+
+    # Apply the method we just defined to all the examples in all the splits of the dataset
+    # starting with the main process first:
+    tokenized_datasets = datasets.map(
+        tokenize_function,
+        batched=True,
+        remove_columns=["idx", "sentence1", "sentence2"],
+    )
+
+    # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
+    # transformers library
+    tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
+
+    def collate_fn(examples):
+        return tokenizer.pad(
+            examples,
+            padding="longest",
+            pad_to_multiple_of=16,  # Specific for FP8
+            return_tensors="pt",
+        )
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(
+        tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size, drop_last=True
+    )
+    eval_dataloader = DataLoader(
+        tokenized_datasets["validation"],
+        shuffle=False,
+        collate_fn=collate_fn,
+        batch_size=16,
+        drop_last=True,
+    )
+
+    return train_dataloader, eval_dataloader
+
+
+def get_training_utilities(model_name: str, batch_size: int = 16, accelerator=None):
+    """
+    Returns a tuple of:
+        - Model
+        - Optimizer
+        - Train dataloader (prepared)
+        - Eval dataloader (prepared)
+        - LR Scheduler
+    Suitable for training on the MRPC dataset
+    """
+    from torch.optim import AdamW
+    from transformers import AutoModelForSequenceClassification, get_linear_schedule_with_warmup
+
+    from accelerate import Accelerator
+
+    if accelerator is None:
+        accelerator = Accelerator()
+    model = AutoModelForSequenceClassification.from_pretrained(model_name)
+    train_dataloader, eval_dataloader = get_dataloaders(model_name, batch_size)
+    optimizer = AdamW(model.parameters(), lr=0.0001)
+    lr_scheduler = get_linear_schedule_with_warmup(
+        optimizer=optimizer,
+        num_warmup_steps=100,
+        num_training_steps=len(train_dataloader) * 2,
+    )
+    train_dataloader, eval_dataloader = accelerator.prepare(train_dataloader, eval_dataloader)
+    return model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+
+
+def get_named_parameters(model):
+    """
+    Same thing as `Accelerator.get_named_parameters` Returns a list of the named parameters of the model (extracted
+    from parallel)
+    """
+    from accelerate.utils import extract_model_from_parallel
+
+    model = extract_model_from_parallel(model)
+    return {n: p for n, p in model.named_parameters()}
+
+
+def evaluate_model(model, dataloader, metric, accelerator=None):
+    "Turns model to .eval(), runs dataloader, calculates metric, then turns eval back on"
+    model.eval()
+    for step, batch in enumerate(dataloader):
+        with torch.no_grad():
+            # W/ MS-AMP, we need to cast while evaluating
+            with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+                outputs = model(**batch)
+        predictions = outputs.logits.argmax(dim=-1)
+        references = batch["labels"]
+        if accelerator is not None and accelerator.num_processes > 1:
+            predictions, references = accelerator.gather_for_metrics((predictions, references))
+        metric.add_batch(predictions=predictions, references=references)
+    return metric.compute()

benchmarks/fp8/ms_amp/non_distributed.py

@@ -0,0 +1,118 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+This script tests to ensure that `accelerate` performs at the same level as raw `MS-AMP`.
+
+This particular script verifies this for single GPU training.
+"""
+
+import evaluate
+import msamp
+import torch
+from fp8_utils import evaluate_model, get_training_utilities
+
+from accelerate import Accelerator
+from accelerate.state import AcceleratorState
+from accelerate.utils import FP8RecipeKwargs, get_grad_scaler, set_seed
+
+
+MODEL_NAME = "bert-base-cased"
+METRIC = evaluate.load("glue", "mrpc")
+
+
+def train_baseline(opt_level="O2"):
+    set_seed(42)
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = get_training_utilities(MODEL_NAME)
+
+    model, optimizer = msamp.initialize(model, optimizer, opt_level=opt_level)
+    model.to("cuda")
+
+    base_model_results = evaluate_model(model, eval_dataloader, METRIC)
+    model.train()
+    scaler = get_grad_scaler()
+
+    for batch in train_dataloader:
+        batch = batch.to("cuda")
+        with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+            outputs = model(**batch)
+        loss = outputs.loss
+        loss = scaler.scale(loss)
+        loss.backward()
+        optimizer.step()
+        optimizer.zero_grad()
+        lr_scheduler.step()
+
+    trained_model_results = evaluate_model(model, eval_dataloader, METRIC)
+
+    assert trained_model_results["accuracy"] > base_model_results["accuracy"], (
+        f"Accuracy should be higher for the trained model: {trained_model_results['accuracy']} > {base_model_results['accuracy']}"
+    )
+    assert trained_model_results["f1"] > base_model_results["f1"], (
+        f"F1 score should be higher for the trained model: {trained_model_results['f1']} > {base_model_results['f1']}"
+    )
+
+    return base_model_results, trained_model_results
+
+
+def train_integration(opt_level="O2"):
+    kwargs_handlers = [FP8RecipeKwargs(backend="msamp", opt_level=opt_level)]
+    AcceleratorState()._reset_state(True)
+    accelerator = Accelerator(mixed_precision="fp8", kwargs_handlers=kwargs_handlers)
+    set_seed(42)
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = get_training_utilities(
+        MODEL_NAME, accelerator=accelerator
+    )
+
+    model, optimizer, lr_scheduler = accelerator.prepare(model, optimizer, lr_scheduler)
+    base_model_results = evaluate_model(model, eval_dataloader, METRIC)
+    model.train()
+
+    for batch in train_dataloader:
+        outputs = model(**batch)
+        loss = outputs.loss
+        accelerator.backward(loss)
+        optimizer.step()
+        optimizer.zero_grad()
+        lr_scheduler.step()
+
+    trained_model_results = evaluate_model(model, eval_dataloader, METRIC)
+
+    assert trained_model_results["accuracy"] > base_model_results["accuracy"], (
+        f"Accuracy should be higher for the trained model: {trained_model_results['accuracy']} > {base_model_results['accuracy']}"
+    )
+    assert trained_model_results["f1"] > base_model_results["f1"], (
+        f"F1 score should be higher for the trained model: {trained_model_results['f1']} > {base_model_results['f1']}"
+    )
+
+    return base_model_results, trained_model_results
+
+
+if __name__ == "__main__":
+    for opt_level in ["O1", "O2"]:
+        baseline_not_trained, baseline_trained = train_baseline(opt_level)
+        accelerator_not_trained, accelerator_trained = train_integration(opt_level)
+
+        assert baseline_not_trained["accuracy"] == accelerator_not_trained["accuracy"], (
+            f"Accuracy should be the same for the baseline and accelerator: {baseline_not_trained['accuracy']} == {accelerator_not_trained['accuracy']}"
+        )
+        assert baseline_not_trained["f1"] == accelerator_not_trained["f1"], (
+            f"F1 score should be the same for the baseline and accelerator: {baseline_not_trained['f1']} == {accelerator_not_trained['f1']}"
+        )
+        assert baseline_trained["accuracy"] == accelerator_trained["accuracy"], (
+            f"Accuracy should be the same for the baseline and accelerator: {baseline_trained['accuracy']} == {accelerator_trained['accuracy']}"
+        )
+        assert baseline_trained["f1"] == accelerator_trained["f1"], (
+            f"F1 score should be the same for the baseline and accelerator: {baseline_trained['f1']} == {accelerator_trained['f1']}"
+        )

benchmarks/fp8/torchao/Dockerfile

@@ -0,0 +1,12 @@
+FROM nvcr.io/nvidia/pytorch:24.07-py3
+
+RUN pip install transformers evaluate datasets
+RUN git clone https://github.com/huggingface/accelerate.git
+
+RUN cd accelerate && \
+    pip install -e . && \
+    cd benchmarks/fp8
+
+RUN /bin/bash
+
+

benchmarks/fp8/torchao/README.md

@@ -0,0 +1,32 @@
+# FP8 Benchmarks
+
+Comparing and running [torchao](https://github.com/pytorch/ao/tree/main/torchao/float8) FP8 with accelerate
+
+## Overview
+
+This repo provides scripts which compare native `torchao` model training against `accelerate`'s own integration. Each modeling type is segmented out via a script, supporting the following:
+
+* Single GPU training (`non_distributed.py`)
+* Multi-GPU training via DistributedDataParallelism (`ddp.py`)
+* Fully Sharded Data Parallelism (`fsdp.py`)
+* DeepSpeed ZeRO 1-3 (`deepspeed.py`)
+
+To run them, it's recommended to use a docker image (see the attached `Dockerfile`) and not install `torchao` manually.
+
+## Running:
+
+There are official Docker images located at `huggingface/accelerate:gpu-fp8-torchao-nightly` which can be used.
+
+You can run all scripts using the core `accelerate launch` command without any `accelerate config` being needed.
+
+For single GPU, run it via `python`:
+
+```bash
+python non_distributed.py
+```
+
+For the rest, run it via `accelerate launch`:
+
+```bash
+accelerate launch ddp.py # or distrib_deepspeed.py, ddp.py
+```

benchmarks/fp8/torchao/ddp.py

@@ -0,0 +1,158 @@
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+This script tests to ensure that `accelerate` performs at the same level as raw `torchao`.
+
+This particular script verifies this for DDP training.
+"""
+
+from functools import partial
+
+import evaluate
+import torch
+from fp8_utils import get_training_utilities
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torchao.float8 import convert_to_float8_training
+
+from accelerate import Accelerator
+from accelerate.state import AcceleratorState
+from accelerate.utils import AORecipeKwargs, set_seed
+
+
+MODEL_NAME = "bert-base-cased"
+METRIC = evaluate.load("glue", "mrpc")
+
+
+def evaluate_model(model, dataloader, metric, accelerator=None):
+    "Turns model to .eval(), runs dataloader, calculates metric, then turns eval back on"
+    model.eval()
+    for step, batch in enumerate(dataloader):
+        with torch.no_grad():
+            outputs = model(**batch)
+        predictions = outputs.logits.argmax(dim=-1)
+        references = batch["labels"]
+        if accelerator is not None and accelerator.num_processes > 1:
+            predictions, references = accelerator.gather_for_metrics((predictions, references))
+        metric.add_batch(predictions=predictions, references=references)
+    return metric.compute()
+
+
+def filter_linear_layers(module, fqn, first_layer_name=None, last_layer_name=None):
+    if isinstance(module, torch.nn.Linear):
+        if module.in_features % 16 != 0 or module.out_features % 16 != 0:
+            return False
+    # For stability reasons, we skip the first and last linear layers
+    # Otherwise can lead to the model not training or converging properly
+    if fqn in (first_layer_name, last_layer_name):
+        return False
+    return True
+
+
+def train_baseline():
+    set_seed(42)
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = get_training_utilities(MODEL_NAME)
+    first_linear = None
+    last_linear = None
+    for name, module in model.named_modules():
+        if isinstance(module, torch.nn.Linear):
+            if first_linear is None:
+                first_linear = name
+            last_linear = name
+    func = partial(filter_linear_layers, first_layer_name=first_linear, last_layer_name=last_linear)
+    accelerator = Accelerator()
+    device = accelerator.device
+    model.to(device)
+
+    convert_to_float8_training(model, module_filter_fn=func)
+
+    # Convert the model to DDP
+    device_ids, output_device = [accelerator.local_process_index], accelerator.local_process_index
+    model = DDP(model, device_ids=device_ids, output_device=output_device)
+
+    base_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)
+    model.train()
+
+    for batch in train_dataloader:
+        with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+            batch = batch.to(device)
+            outputs = model(**batch)
+            loss = outputs.loss
+            loss.backward()
+        optimizer.step()
+        optimizer.zero_grad()
+        lr_scheduler.step()
+
+    trained_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)
+
+    assert trained_model_results["accuracy"] > base_model_results["accuracy"], (
+        f"Accuracy should be higher for the trained model: {trained_model_results['accuracy']} > {base_model_results['accuracy']}"
+    )
+    assert trained_model_results["f1"] > base_model_results["f1"], (
+        f"F1 score should be higher for the trained model: {trained_model_results['f1']} > {base_model_results['f1']}"
+    )
+
+    return base_model_results, trained_model_results
+
+
+def train_integration():
+    AcceleratorState()._reset_state(True)
+    accelerator = Accelerator(mixed_precision="fp8", kwargs_handlers=[AORecipeKwargs()])
+    set_seed(42)
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = get_training_utilities(
+        MODEL_NAME, accelerator=accelerator
+    )
+
+    model, optimizer = accelerator.prepare(model, optimizer)
+    base_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)
+    model.train()
+
+    for batch in train_dataloader:
+        outputs = model(**batch)
+        loss = outputs.loss
+        accelerator.backward(loss)
+        optimizer.step()
+        optimizer.zero_grad()
+        lr_scheduler.step()
+
+    trained_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)
+
+    assert trained_model_results["accuracy"] > base_model_results["accuracy"], (
+        f"Accuracy should be higher for the trained model: {trained_model_results['accuracy']} > {base_model_results['accuracy']}"
+    )
+    assert trained_model_results["f1"] > base_model_results["f1"], (
+        f"F1 score should be higher for the trained model: {trained_model_results['f1']} > {base_model_results['f1']}"
+    )
+
+    return base_model_results, trained_model_results
+
+
+if __name__ == "__main__":
+    baseline_not_trained, baseline_trained = train_baseline()
+    accelerator_not_trained, accelerator_trained = train_integration()
+
+    assert baseline_not_trained["accuracy"] == accelerator_not_trained["accuracy"], (
+        f"Accuracy should be the same for the baseline and accelerator: {baseline_not_trained['accuracy']} == {accelerator_not_trained['accuracy']}"
+    )
+    assert baseline_not_trained["f1"] == accelerator_not_trained["f1"], (
+        f"F1 score should be the same for the baseline and accelerator: {baseline_not_trained['f1']} == {accelerator_not_trained['f1']}"
+    )
+    assert baseline_trained["accuracy"] == accelerator_trained["accuracy"], (
+        f"Accuracy should be the same for the baseline and accelerator: {baseline_trained['accuracy']} == {accelerator_trained['accuracy']}"
+    )
+    assert baseline_trained["f1"] == accelerator_trained["f1"], (
+        f"F1 score should be the same for the baseline and accelerator: {baseline_trained['f1']} == {accelerator_trained['f1']}"
+    )
+
+    torch.distributed.destroy_process_group()

benchmarks/fp8/torchao/distrib_deepspeed.py

@@ -0,0 +1,213 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+This script tests to ensure that `accelerate` performs at the same level as raw `torchao`.
+
+This particular script verifies this for deepspeed training.
+"""
+
+from functools import partial
+from unittest.mock import patch
+
+import deepspeed
+import evaluate
+import torch
+from fp8_utils import evaluate_model, get_training_utilities
+from torchao.float8 import convert_to_float8_training
+from transformers.integrations import HfDeepSpeedConfig
+
+from accelerate import Accelerator, DeepSpeedPlugin
+from accelerate.state import AcceleratorState
+from accelerate.utils import AORecipeKwargs, set_seed
+
+
+MODEL_NAME = "bert-base-cased"
+METRIC = evaluate.load("glue", "mrpc")
+
+
+def filter_linear_layers(module, fqn, first_layer_name=None, last_layer_name=None):
+    if isinstance(module, torch.nn.Linear):
+        if module.in_features % 16 != 0 or module.out_features % 16 != 0:
+            return False
+    # For stability reasons, we skip the first and last linear layers
+    # Otherwise can lead to the model not training or converging properly
+    if fqn in (first_layer_name, last_layer_name):
+        return False
+    return True
+
+
+def train_baseline(zero_stage: int = 1):
+    set_seed(42)
+    # This forces transformers to think Zero-3 Init should be used
+    with patch("transformers.integrations.deepspeed.is_deepspeed_zero3_enabled") as mock:
+        mock.return_value = zero_stage == 3
+
+    config = HfDeepSpeedConfig(
+        {
+            "train_micro_batch_size_per_gpu": 16,
+            "gradient_accumulation_steps": 1,
+            "zero_optimization": {"stage": zero_stage},
+        }
+    )
+    plugin = DeepSpeedPlugin(hf_ds_config=config)
+    accelerator = Accelerator(deepspeed_plugin=plugin)
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = get_training_utilities(
+        MODEL_NAME, accelerator=accelerator
+    )
+    first_linear = None
+    last_linear = None
+    for name, module in model.named_modules():
+        if isinstance(module, torch.nn.Linear):
+            if first_linear is None:
+                first_linear = name
+            last_linear = name
+    func = partial(filter_linear_layers, first_layer_name=first_linear, last_layer_name=last_linear)
+
+    convert_to_float8_training(model, module_filter_fn=func)
+
+    import numpy as np
+
+    config = {
+        "train_batch_size": 32,
+        "train_micro_batch_size_per_gpu": 16,
+        "gradient_accumulation_steps": 1,
+        "zero_optimization": {
+            "stage": zero_stage,
+            "offload_optimizer": {"device": "none", "nvme_path": None},
+            "offload_param": {"device": "none", "nvme_path": None},
+            "stage3_gather_16bit_weights_on_model_save": False,
+        },
+        "gradient_clipping": 1.0,
+        "steps_per_print": np.inf,
+        "bf16": {"enabled": True},
+        "fp16": {"enabled": False},
+        "zero_allow_untested_optimizer": True,
+    }
+
+    (
+        model,
+        optimizer,
+        _,
+        lr_scheduler,
+    ) = deepspeed.initialize(
+        model=model,
+        optimizer=optimizer,
+        lr_scheduler=lr_scheduler,
+        config_params=config,
+    )
+
+    base_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)
+    model.train()
+
+    model_outputs = []
+    data = []
+
+    for batch in train_dataloader:
+        outputs = model(**batch)
+        data.append(batch.to("cpu"))
+        model_outputs.append(outputs.logits.to("cpu"))
+        loss = outputs.loss
+        model.backward(loss)
+        model.step()
+        for _ in range(accelerator.num_processes):
+            lr_scheduler.step()
+
+    trained_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)
+    model.destroy()
+    assert trained_model_results["accuracy"] > base_model_results["accuracy"], (
+        f"Accuracy should be higher for the trained model: {trained_model_results['accuracy']} > {base_model_results['accuracy']}"
+    )
+    assert trained_model_results["f1"] > base_model_results["f1"], (
+        f"F1 score should be higher for the trained model: {trained_model_results['f1']} > {base_model_results['f1']}"
+    )
+
+    del config
+    return base_model_results, trained_model_results, model_outputs, data
+
+
+def train_integration(zero_stage: int = 1):
+    set_seed(42)
+    AcceleratorState()._reset_state(True)
+    config = HfDeepSpeedConfig(
+        {
+            "train_micro_batch_size_per_gpu": 16,
+            "gradient_accumulation_steps": 1,
+            "zero_optimization": {"stage": zero_stage},
+        }
+    )
+    deepspeed_plugin = DeepSpeedPlugin(
+        hf_ds_config=config,
+    )
+    # This forces transformers to think Zero-3 Init should be used
+    with patch("transformers.integrations.deepspeed.is_deepspeed_zero3_enabled") as mock:
+        mock.return_value = zero_stage == 3
+    accelerator = Accelerator(
+        mixed_precision="fp8", kwargs_handlers=[AORecipeKwargs()], deepspeed_plugin=deepspeed_plugin
+    )
+
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = get_training_utilities(
+        MODEL_NAME, accelerator=accelerator
+    )
+
+    model, optimizer, lr_scheduler, train_dataloader, eval_dataloader = accelerator.prepare(
+        model, optimizer, lr_scheduler, train_dataloader, eval_dataloader
+    )
+    base_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)
+    model.train()
+    model_outputs = []
+    data = []
+    for batch in train_dataloader:
+        outputs = model(**batch)
+        data.append(batch.to("cpu"))
+        model_outputs.append(outputs.logits.to("cpu"))
+        loss = outputs.loss
+        accelerator.backward(loss)
+        optimizer.step()
+        lr_scheduler.step()
+        optimizer.zero_grad()
+
+    trained_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)
+    model.destroy()
+    assert trained_model_results["accuracy"] > base_model_results["accuracy"], (
+        f"Accuracy should be higher for the trained model: {trained_model_results['accuracy']} > {base_model_results['accuracy']}"
+    )
+    assert trained_model_results["f1"] > base_model_results["f1"], (
+        f"F1 score should be higher for the trained model: {trained_model_results['f1']} > {base_model_results['f1']}"
+    )
+
+    del config
+    return base_model_results, trained_model_results, model_outputs, data
+
+
+if __name__ == "__main__":
+    for zero_stage in [1, 2, 3]:
+        baseline_not_trained, baseline_trained, baseline_outputs, baseline_data = train_baseline(zero_stage)
+        accelerator_not_trained, accelerator_trained, accelerator_outputs, accelerator_data = train_integration(
+            zero_stage
+        )
+        assert baseline_not_trained["accuracy"] == accelerator_not_trained["accuracy"], (
+            f"ZERO stage {zero_stage}: Accuracy should be the same for the baseline and accelerator: {baseline_not_trained['accuracy']} == {accelerator_not_trained['accuracy']}"
+        )
+        assert baseline_not_trained["f1"] == accelerator_not_trained["f1"], (
+            f"ZERO stage {zero_stage}: F1 score should be the same for the baseline and accelerator: {baseline_not_trained['f1']} == {accelerator_not_trained['f1']}"
+        )
+        assert baseline_trained["accuracy"] == accelerator_trained["accuracy"], (
+            f"ZERO stage {zero_stage}: Accuracy should be the same for the baseline and accelerator: {baseline_trained['accuracy']} == {accelerator_trained['accuracy']}"
+        )
+        assert baseline_trained["f1"] == accelerator_trained["f1"], (
+            f"ZERO stage {zero_stage}: F1 score should be the same for the baseline and accelerator: {baseline_trained['f1']} == {accelerator_trained['f1']}"
+        )
+        AcceleratorState()._reset_state(True)
+    torch.distributed.destroy_process_group()

benchmarks/fp8/torchao/fp8_utils.py

@@ -0,0 +1,116 @@
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import torch
+
+
+def get_dataloaders(model_name: str, batch_size: int = 16):
+    from datasets import load_dataset
+    from torch.utils.data import DataLoader
+    from transformers import AutoTokenizer
+
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    datasets = load_dataset("glue", "mrpc")
+
+    def tokenize_function(examples):
+        # max_length=None => use the model max length (it's actually the default)
+        outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None)
+        return outputs
+
+    # Apply the method we just defined to all the examples in all the splits of the dataset
+    # starting with the main process first:
+    tokenized_datasets = datasets.map(
+        tokenize_function,
+        batched=True,
+        remove_columns=["idx", "sentence1", "sentence2"],
+    )
+
+    # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
+    # transformers library
+    tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
+
+    def collate_fn(examples):
+        return tokenizer.pad(
+            examples,
+            padding="longest",
+            pad_to_multiple_of=16,  # Specific for FP8
+            return_tensors="pt",
+        )
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(
+        tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size, drop_last=True
+    )
+    eval_dataloader = DataLoader(
+        tokenized_datasets["validation"],
+        shuffle=False,
+        collate_fn=collate_fn,
+        batch_size=16,
+        drop_last=True,
+    )
+
+    return train_dataloader, eval_dataloader
+
+
+def get_training_utilities(model_name: str, batch_size: int = 16, accelerator=None, prepare=True):
+    """
+    Returns a tuple of:
+        - Model
+        - Optimizer
+        - Train dataloader (prepared)
+        - Eval dataloader (prepared)
+        - LR Scheduler
+    Suitable for training on the MRPC dataset
+    """
+    from torch.optim import AdamW
+    from transformers import AutoModelForSequenceClassification, get_linear_schedule_with_warmup
+
+    from accelerate import Accelerator
+
+    if accelerator is None:
+        accelerator = Accelerator()
+    model = AutoModelForSequenceClassification.from_pretrained(model_name)
+    train_dataloader, eval_dataloader = get_dataloaders(model_name, batch_size)
+    optimizer = AdamW(model.parameters(), lr=0.0001)
+    lr_scheduler = get_linear_schedule_with_warmup(
+        optimizer=optimizer,
+        num_warmup_steps=100,
+        num_training_steps=len(train_dataloader) * 2,
+    )
+    train_dataloader, eval_dataloader = accelerator.prepare(train_dataloader, eval_dataloader)
+    return model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+
+
+def get_named_parameters(model):
+    """
+    Same thing as `Accelerator.get_named_parameters` Returns a list of the named parameters of the model (extracted
+    from parallel)
+    """
+    from accelerate.utils import extract_model_from_parallel
+
+    model = extract_model_from_parallel(model)
+    return {n: p for n, p in model.named_parameters()}
+
+
+def evaluate_model(model, dataloader, metric, accelerator=None):
+    "Turns model to .eval(), runs dataloader, calculates metric, then turns eval back on"
+    model.eval()
+    for step, batch in enumerate(dataloader):
+        with torch.no_grad():
+            outputs = model(**batch)
+        predictions = outputs.logits.argmax(dim=-1)
+        references = batch["labels"]
+        if accelerator is not None and accelerator.num_processes > 1:
+            predictions, references = accelerator.gather_for_metrics((predictions, references))
+        metric.add_batch(predictions=predictions, references=references)
+    return metric.compute()

benchmarks/fp8/torchao/fsdp.py

@@ -0,0 +1,173 @@
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+This script tests to ensure that `accelerate` performs at the same level as raw `torchao`.
+
+This particular script verifies this for FSDP training.
+"""
+
+from functools import partial
+
+import evaluate
+import torch
+from fp8_utils import get_training_utilities
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+from torch.distributed.fsdp import MixedPrecision
+from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
+from torchao.float8 import convert_to_float8_training
+from transformers.models.bert import BertLayer
+
+from accelerate import Accelerator
+from accelerate import FullyShardedDataParallelPlugin as FSDPPlugin
+from accelerate.state import AcceleratorState
+from accelerate.utils import AORecipeKwargs, set_seed
+
+
+MODEL_NAME = "bert-base-cased"
+METRIC = evaluate.load("glue", "mrpc")
+
+FSDP_WRAP_POLICY = partial(transformer_auto_wrap_policy, transformer_layer_cls={BertLayer})
+
+
+def filter_linear_layers(module, fqn, first_layer_name=None, last_layer_name=None):
+    if isinstance(module, torch.nn.Linear):
+        if module.in_features % 16 != 0 or module.out_features % 16 != 0:
+            return False
+    # For stability reasons, we skip the first and last linear layers
+    # Otherwise can lead to the model not training or converging properly
+    if fqn in (first_layer_name, last_layer_name):
+        return False
+    return True
+
+
+def evaluate_model(model, dataloader, metric, accelerator=None):
+    "Turns model to .eval(), runs dataloader, calculates metric, then turns eval back on"
+    model.eval()
+    for step, batch in enumerate(dataloader):
+        with torch.no_grad():
+            outputs = model(**batch)
+        predictions = outputs.logits.argmax(dim=-1)
+        references = batch["labels"]
+        if accelerator is not None and accelerator.num_processes > 1:
+            predictions, references = accelerator.gather_for_metrics((predictions, references))
+        metric.add_batch(predictions=predictions, references=references)
+    return metric.compute()
+
+
+def train_baseline():
+    set_seed(42)
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = get_training_utilities(MODEL_NAME)
+    first_linear = None
+    last_linear = None
+    for name, module in model.named_modules():
+        if isinstance(module, torch.nn.Linear):
+            if first_linear is None:
+                first_linear = name
+            last_linear = name
+    func = partial(filter_linear_layers, first_layer_name=first_linear, last_layer_name=last_linear)
+    accelerator = Accelerator()
+    device = accelerator.device
+    model.to(device)
+
+    convert_to_float8_training(model, module_filter_fn=func)
+
+    # Convert the model to FSDP
+    model = FSDP(
+        model,
+        use_orig_params=True,
+        mixed_precision=MixedPrecision(param_dtype=torch.bfloat16, reduce_dtype=torch.float32),
+        auto_wrap_policy=FSDP_WRAP_POLICY,
+    )
+
+    base_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)
+    model.train()
+
+    for batch in train_dataloader:
+        with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+            batch = batch.to(device)
+            outputs = model(**batch)
+        loss = outputs.loss
+        loss.backward()
+        optimizer.step()
+        optimizer.zero_grad()
+        lr_scheduler.step()
+
+    trained_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)
+
+    assert trained_model_results["accuracy"] > base_model_results["accuracy"], (
+        f"Accuracy should be higher for the trained model: {trained_model_results['accuracy']} > {base_model_results['accuracy']}"
+    )
+    assert trained_model_results["f1"] > base_model_results["f1"], (
+        f"F1 score should be higher for the trained model: {trained_model_results['f1']} > {base_model_results['f1']}"
+    )
+
+    return base_model_results, trained_model_results
+
+
+def train_integration():
+    AcceleratorState()._reset_state(True)
+    fsdp_plugin = FSDPPlugin(
+        auto_wrap_policy=FSDP_WRAP_POLICY,
+        use_orig_params=True,
+        mixed_precision_policy=MixedPrecision(param_dtype=torch.bfloat16, reduce_dtype=torch.float32),
+    )
+    accelerator = Accelerator(mixed_precision="fp8", fsdp_plugin=fsdp_plugin, kwargs_handlers=[AORecipeKwargs()])
+    set_seed(42)
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = get_training_utilities(
+        MODEL_NAME, accelerator=accelerator
+    )
+
+    model, optimizer = accelerator.prepare(model, optimizer)
+    base_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)
+    model.train()
+
+    for batch in train_dataloader:
+        outputs = model(**batch)
+        loss = outputs.loss
+        accelerator.backward(loss)
+        optimizer.step()
+        optimizer.zero_grad()
+        lr_scheduler.step()
+
+    trained_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)
+
+    assert trained_model_results["accuracy"] > base_model_results["accuracy"], (
+        f"Accuracy should be higher for the trained model: {trained_model_results['accuracy']} > {base_model_results['accuracy']}"
+    )
+    assert trained_model_results["f1"] > base_model_results["f1"], (
+        f"F1 score should be higher for the trained model: {trained_model_results['f1']} > {base_model_results['f1']}"
+    )
+
+    return base_model_results, trained_model_results
+
+
+if __name__ == "__main__":
+    baseline_not_trained, baseline_trained = train_baseline()
+    accelerator_not_trained, accelerator_trained = train_integration()
+
+    assert baseline_not_trained["accuracy"] == accelerator_not_trained["accuracy"], (
+        f"Accuracy should be the same for the baseline and accelerator: {baseline_not_trained['accuracy']} == {accelerator_not_trained['accuracy']}"
+    )
+    assert baseline_not_trained["f1"] == accelerator_not_trained["f1"], (
+        f"F1 score should be the same for the baseline and accelerator: {baseline_not_trained['f1']} == {accelerator_not_trained['f1']}"
+    )
+    assert baseline_trained["accuracy"] == accelerator_trained["accuracy"], (
+        f"Accuracy should be the same for the baseline and accelerator: {baseline_trained['accuracy']} == {accelerator_trained['accuracy']}"
+    )
+    assert baseline_trained["f1"] == accelerator_trained["f1"], (
+        f"F1 score should be the same for the baseline and accelerator: {baseline_trained['f1']} == {accelerator_trained['f1']}"
+    )
+
+    torch.distributed.destroy_process_group()

benchmarks/fp8/torchao/non_distributed.py

@@ -0,0 +1,145 @@
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+This script tests to ensure that `accelerate` performs at the same level as raw `torchao`.
+
+This particular script verifies this for single GPU training.
+"""
+
+from functools import partial
+
+import evaluate
+import torch
+from fp8_utils import get_training_utilities
+from torchao.float8 import convert_to_float8_training
+
+from accelerate import Accelerator
+from accelerate.state import AcceleratorState
+from accelerate.utils import AORecipeKwargs, set_seed
+
+
+MODEL_NAME = "bert-base-cased"
+METRIC = evaluate.load("glue", "mrpc")
+
+
+def evaluate_model(model, dataloader, metric, accelerator=None):
+    "Turns model to .eval(), runs dataloader, calculates metric, then turns eval back on"
+    model.eval()
+    for step, batch in enumerate(dataloader):
+        with torch.no_grad():
+            outputs = model(**batch)
+        predictions = outputs.logits.argmax(dim=-1)
+        references = batch["labels"]
+        if accelerator is not None and accelerator.num_processes > 1:
+            predictions, references = accelerator.gather_for_metrics((predictions, references))
+        metric.add_batch(predictions=predictions, references=references)
+    return metric.compute()
+
+
+def filter_linear_layers(module, fqn, first_layer_name=None, last_layer_name=None):
+    if isinstance(module, torch.nn.Linear):
+        if module.in_features % 16 != 0 or module.out_features % 16 != 0:
+            return False
+    # For stability reasons, we skip the first and last linear layers
+    # Otherwise can lead to the model not training or converging properly
+    if fqn in (first_layer_name, last_layer_name):
+        return False
+    return True
+
+
+def train_baseline():
+    set_seed(42)
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = get_training_utilities(MODEL_NAME)
+    first_linear = None
+    last_linear = None
+    for name, module in model.named_modules():
+        if isinstance(module, torch.nn.Linear):
+            if first_linear is None:
+                first_linear = name
+            last_linear = name
+
+    func = partial(filter_linear_layers, first_layer_name=first_linear, last_layer_name=last_linear)
+    model.to("cuda")
+    convert_to_float8_training(model, module_filter_fn=func)
+    base_model_results = evaluate_model(model, eval_dataloader, METRIC)
+    model.train()
+
+    for batch in train_dataloader:
+        with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+            outputs = model(**batch)
+            loss = outputs.loss
+            loss.backward()
+        optimizer.step()
+        optimizer.zero_grad()
+        lr_scheduler.step()
+
+    trained_model_results = evaluate_model(model, eval_dataloader, METRIC)
+
+    assert trained_model_results["accuracy"] > base_model_results["accuracy"], (
+        f"Accuracy should be higher for the trained model: {trained_model_results['accuracy']} > {base_model_results['accuracy']}"
+    )
+    assert trained_model_results["f1"] > base_model_results["f1"], (
+        f"F1 score should be higher for the trained model: {trained_model_results['f1']} > {base_model_results['f1']}"
+    )
+
+    return base_model_results, trained_model_results
+
+
+def train_integration():
+    set_seed(42)
+    accelerator = Accelerator(mixed_precision="fp8", kwargs_handlers=[AORecipeKwargs()])
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = get_training_utilities(
+        MODEL_NAME, accelerator=accelerator
+    )
+    model = accelerator.prepare(model)
+    base_model_results = evaluate_model(model, eval_dataloader, METRIC)
+    model.train()
+
+    for batch in train_dataloader:
+        outputs = model(**batch)
+        loss = outputs.loss
+        loss.backward()
+        optimizer.step()
+        optimizer.zero_grad()
+        lr_scheduler.step()
+
+    trained_model_results = evaluate_model(model, eval_dataloader, METRIC)
+
+    assert trained_model_results["accuracy"] > base_model_results["accuracy"], (
+        f"Accuracy should be higher for the trained model: {trained_model_results['accuracy']} > {base_model_results['accuracy']}"
+    )
+    assert trained_model_results["f1"] > base_model_results["f1"], (
+        f"F1 score should be higher for the trained model: {trained_model_results['f1']} > {base_model_results['f1']}"
+    )
+
+    return base_model_results, trained_model_results
+
+
+if __name__ == "__main__":
+    baseline_not_trained, baseline_trained = train_baseline()
+    AcceleratorState._reset_state(True)
+    accelerator_not_trained, accelerator_trained = train_integration()
+    assert baseline_not_trained["accuracy"] == accelerator_not_trained["accuracy"], (
+        f"Accuracy should be the same for the baseline and accelerator: {baseline_not_trained['accuracy']} == {accelerator_not_trained['accuracy']}"
+    )
+    assert baseline_not_trained["f1"] == accelerator_not_trained["f1"], (
+        f"F1 score should be the same for the baseline and accelerator: {baseline_not_trained['f1']} == {accelerator_not_trained['f1']}"
+    )
+    assert baseline_trained["accuracy"] == accelerator_trained["accuracy"], (
+        f"Accuracy should be the same for the baseline and accelerator: {baseline_trained['accuracy']} == {accelerator_trained['accuracy']}"
+    )
+    assert baseline_trained["f1"] == accelerator_trained["f1"], (
+        f"F1 score should be the same for the baseline and accelerator: {baseline_trained['f1']} == {accelerator_trained['f1']}"
+    )

benchmarks/fp8/transformer_engine/Dockerfile

@@ -0,0 +1,15 @@
+ARG BASE_YEAR=25
+ARG BASE_MONTH=03
+
+FROM nvcr.io/nvidia/pytorch:${BASE_YEAR}.${BASE_MONTH}-py3
+
+RUN pip install transformers evaluate datasets
+RUN git clone https://github.com/huggingface/accelerate.git
+
+RUN cd accelerate && \
+    pip install -e .[deepspeed] && \
+    cd benchmarks/fp8
+
+RUN /bin/bash
+
+

benchmarks/fp8/transformer_engine/README.md

@@ -0,0 +1,32 @@
+# FP8 Benchmarks
+
+Comparing and running [TransformerEngine](https://github.com/NVIDIA/TransformerEngine) FP8 with accelerate
+
+## Overview
+
+This repo provides scripts which compare native TransformerEngine model training against `accelerate`'s own integration. Each modeling type is segmented out via a script, supporting the following:
+
+* Single GPU training (`non_distributed.py`)
+* Multi-GPU training via DistributedDataParallelism (`ddp.py`)
+* Fully Sharded Data Parallelism (`fsdp.py`)
+* DeepSpeed ZeRO 1-3 (`deepspeed.py`)
+
+To run them, it's recommended to use a docker image (see the attached `Dockerfile`) and not install `TransformerEngine` manually.
+
+## Running:
+
+There are official Docker images located at `huggingface/accelerate:gpu-fp8-transformerengine-nightly` which can be used.
+
+You can run all scripts using the core `accelerate launch` command without any `accelerate config` being needed.
+
+For single GPU, run it via `python`:
+
+```bash
+python non_distributed.py
+```
+
+For the rest, run it via `accelerate launch`:
+
+```bash
+accelerate launch ddp.py # or distrib_deepspeed.py, ddp.py
+```

benchmarks/fp8/transformer_engine/ddp.py

@@ -0,0 +1,144 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+This script tests to ensure that `accelerate` performs at the same level as raw `TransformersEngine`.
+
+This particular script verifies this for DDP training.
+"""
+
+import evaluate
+import torch
+import transformer_engine.common.recipe as te_recipe
+import transformer_engine.pytorch as te
+from fp8_utils import evaluate_model, get_named_parameters, get_training_utilities
+from torch.nn.parallel import DistributedDataParallel as DDP
+from transformer_engine.common.recipe import DelayedScaling
+
+from accelerate import Accelerator
+from accelerate.state import AcceleratorState
+from accelerate.utils import FP8RecipeKwargs, set_seed
+from accelerate.utils.transformer_engine import convert_model
+
+
+MODEL_NAME = "bert-base-cased"
+METRIC = evaluate.load("glue", "mrpc")
+
+
+def train_baseline():
+    set_seed(42)
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = get_training_utilities(MODEL_NAME)
+    accelerator = Accelerator()
+    device = accelerator.device
+    model.to(device)
+
+    # Convert the model to TE
+    old_named_params = get_named_parameters(model)
+
+    with torch.no_grad():
+        convert_model(model)
+
+    FP8_RECIPE_KWARGS = {"fp8_format": te_recipe.Format.HYBRID, "amax_history_len": 32, "amax_compute_algo": "max"}
+    fp8_recipe = DelayedScaling(**FP8_RECIPE_KWARGS)
+
+    new_named_params = get_named_parameters(model)
+
+    # Convert the model to DDP
+    device_ids, output_device = [accelerator.local_process_index], accelerator.local_process_index
+    model = DDP(model, device_ids=device_ids, output_device=output_device)
+
+    mapping = {p: new_named_params[n] for n, p in old_named_params.items()}
+    for param_group in optimizer.param_groups:
+        param_group["params"] = [mapping[p] for p in param_group["params"]]
+
+    base_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)
+    model.train()
+
+    for _ in range(2):
+        for batch in train_dataloader:
+            with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
+                with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+                    batch = batch.to(device)
+                    outputs = model(**batch)
+            loss = outputs.loss
+            loss.backward()
+            optimizer.step()
+            optimizer.zero_grad()
+            lr_scheduler.step()
+
+    trained_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)
+
+    assert trained_model_results["accuracy"] > base_model_results["accuracy"], (
+        f"Accuracy should be higher for the trained model: {trained_model_results['accuracy']} > {base_model_results['accuracy']}"
+    )
+    assert trained_model_results["f1"] > base_model_results["f1"], (
+        f"F1 score should be higher for the trained model: {trained_model_results['f1']} > {base_model_results['f1']}"
+    )
+
+    return base_model_results, trained_model_results
+
+
+def train_integration():
+    FP8_RECIPE_KWARGS = {"fp8_format": "HYBRID", "amax_history_len": 32, "amax_compute_algo": "max"}
+    kwargs_handlers = [FP8RecipeKwargs(backend="TE", **FP8_RECIPE_KWARGS)]
+    AcceleratorState()._reset_state(True)
+    accelerator = Accelerator(mixed_precision="fp8", kwargs_handlers=kwargs_handlers)
+    set_seed(42)
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = get_training_utilities(
+        MODEL_NAME, accelerator=accelerator
+    )
+
+    model, optimizer = accelerator.prepare(model, optimizer)
+    base_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)
+    model.train()
+
+    for _ in range(2):
+        for batch in train_dataloader:
+            outputs = model(**batch)
+            loss = outputs.loss
+            accelerator.backward(loss)
+            optimizer.step()
+            optimizer.zero_grad()
+            lr_scheduler.step()
+
+    trained_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)
+
+    assert trained_model_results["accuracy"] > base_model_results["accuracy"], (
+        f"Accuracy should be higher for the trained model: {trained_model_results['accuracy']} > {base_model_results['accuracy']}"
+    )
+    assert trained_model_results["f1"] > base_model_results["f1"], (
+        f"F1 score should be higher for the trained model: {trained_model_results['f1']} > {base_model_results['f1']}"
+    )
+
+    return base_model_results, trained_model_results
+
+
+if __name__ == "__main__":
+    baseline_not_trained, baseline_trained = train_baseline()
+    accelerator_not_trained, accelerator_trained = train_integration()
+
+    assert baseline_not_trained["accuracy"] == accelerator_not_trained["accuracy"], (
+        f"Accuracy should be the same for the baseline and accelerator: {baseline_not_trained['accuracy']} == {accelerator_not_trained['accuracy']}"
+    )
+    assert baseline_not_trained["f1"] == accelerator_not_trained["f1"], (
+        f"F1 score should be the same for the baseline and accelerator: {baseline_not_trained['f1']} == {accelerator_not_trained['f1']}"
+    )
+    assert baseline_trained["accuracy"] == accelerator_trained["accuracy"], (
+        f"Accuracy should be the same for the baseline and accelerator: {baseline_trained['accuracy']} == {accelerator_trained['accuracy']}"
+    )
+    assert baseline_trained["f1"] == accelerator_trained["f1"], (
+        f"F1 score should be the same for the baseline and accelerator: {baseline_trained['f1']} == {accelerator_trained['f1']}"
+    )
+
+    torch.distributed.destroy_process_group()

benchmarks/fp8/transformer_engine/distrib_deepspeed.py

@@ -0,0 +1,191 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+This script tests to ensure that `accelerate` performs at the same level as raw `TransformersEngine`.
+
+This particular script verifies this for DDP training.
+"""
+
+from unittest.mock import patch
+
+import deepspeed
+import evaluate
+import torch
+import transformer_engine.common.recipe as te_recipe
+import transformer_engine.pytorch as te
+from fp8_utils import evaluate_model, get_named_parameters, get_training_utilities
+from transformer_engine.common.recipe import DelayedScaling
+
+from accelerate import Accelerator, DeepSpeedPlugin
+from accelerate.state import AcceleratorState
+from accelerate.utils import FP8RecipeKwargs, set_seed
+from accelerate.utils.transformer_engine import convert_model
+
+
+MODEL_NAME = "bert-base-cased"
+METRIC = evaluate.load("glue", "mrpc")
+
+
+def train_baseline(zero_stage: int = 1):
+    # This forces transformers to think Zero-3 Init should be used
+    with patch("transformers.integrations.deepspeed.is_deepspeed_zero3_enabled") as mock:
+        mock.return_value = zero_stage == 3
+    set_seed(42)
+
+    accelerator = Accelerator()
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = get_training_utilities(
+        MODEL_NAME, accelerator=accelerator
+    )
+
+    # Convert the model to TE
+    old_named_params = get_named_parameters(model)
+
+    with torch.no_grad():
+        convert_model(model)
+    new_named_params = get_named_parameters(model)
+
+    mapping = {p: new_named_params[n] for n, p in old_named_params.items()}
+    for param_group in optimizer.param_groups:
+        param_group["params"] = [mapping[p] for p in param_group["params"]]
+
+    FP8_RECIPE_KWARGS = {"fp8_format": te_recipe.Format.HYBRID, "amax_history_len": 32, "amax_compute_algo": "max"}
+    fp8_recipe = DelayedScaling(**FP8_RECIPE_KWARGS)
+
+    import numpy as np
+
+    config = {
+        "train_batch_size": 16,
+        "train_micro_batch_size_per_gpu": 16,
+        "gradient_accumulation_steps": 1,
+        "zero_optimization": {
+            "stage": zero_stage,
+            "offload_optimizer": {"device": "none", "nvme_path": None},
+            "offload_param": {"device": "none", "nvme_path": None},
+            "stage3_gather_16bit_weights_on_model_save": False,
+        },
+        "gradient_clipping": 1.0,
+        "steps_per_print": np.inf,
+        "bf16": {"enabled": True},
+        "fp16": {"enabled": False},
+        "zero_allow_untested_optimizer": True,
+    }
+
+    (
+        model,
+        optimizer,
+        _,
+        _,
+    ) = deepspeed.initialize(
+        model=model,
+        optimizer=optimizer,
+        config_params=config,
+    )
+
+    base_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)
+    model.train()
+
+    model_outputs = []
+    data = []
+
+    for _ in range(2):
+        for batch in train_dataloader:
+            with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
+                outputs = model(**batch)
+                data.append(batch.to("cpu"))
+            model_outputs.append(outputs.logits.to("cpu"))
+            loss = outputs.loss
+            model.backward(loss)
+            model.step()
+            for _ in range(accelerator.num_processes):
+                lr_scheduler.step()
+
+    trained_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)
+    model.destroy()
+    assert trained_model_results["accuracy"] > base_model_results["accuracy"], (
+        f"Accuracy should be higher for the trained model: {trained_model_results['accuracy']} > {base_model_results['accuracy']}"
+    )
+    assert trained_model_results["f1"] > base_model_results["f1"], (
+        f"F1 score should be higher for the trained model: {trained_model_results['f1']} > {base_model_results['f1']}"
+    )
+
+    return base_model_results, trained_model_results, model_outputs, data
+
+
+def train_integration(zero_stage: int = 1):
+    set_seed(42)
+    FP8_RECIPE_KWARGS = {"fp8_format": "HYBRID", "amax_history_len": 32, "amax_compute_algo": "max"}
+    kwargs_handlers = [FP8RecipeKwargs(backend="TE", **FP8_RECIPE_KWARGS)]
+    AcceleratorState()._reset_state(True)
+    deepspeed_plugin = DeepSpeedPlugin(
+        zero_stage=zero_stage,
+        zero3_init_flag=zero_stage == 3,
+    )
+    accelerator = Accelerator(
+        mixed_precision="fp8", kwargs_handlers=kwargs_handlers, deepspeed_plugin=deepspeed_plugin
+    )
+    accelerator.state.deepspeed_plugin.deepspeed_config["train_micro_batch_size_per_gpu"] = 16
+
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = get_training_utilities(
+        MODEL_NAME, accelerator=accelerator
+    )
+
+    model, optimizer, lr_scheduler = accelerator.prepare(model, optimizer, lr_scheduler)
+    base_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)
+    model.train()
+    model_outputs = []
+    data = []
+    for _ in range(2):
+        for batch in train_dataloader:
+            outputs = model(**batch)
+            data.append(batch.to("cpu"))
+            model_outputs.append(outputs.logits.to("cpu"))
+            loss = outputs.loss
+            accelerator.backward(loss)
+            optimizer.step()
+            lr_scheduler.step()
+            optimizer.zero_grad()
+
+    trained_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)
+    model.destroy()
+    assert trained_model_results["accuracy"] > base_model_results["accuracy"], (
+        f"Accuracy should be higher for the trained model: {trained_model_results['accuracy']} > {base_model_results['accuracy']}"
+    )
+    assert trained_model_results["f1"] > base_model_results["f1"], (
+        f"F1 score should be higher for the trained model: {trained_model_results['f1']} > {base_model_results['f1']}"
+    )
+
+    return base_model_results, trained_model_results, model_outputs, data
+
+
+if __name__ == "__main__":
+    for zero_stage in [1, 2, 3]:
+        baseline_not_trained, baseline_trained, baseline_outputs, baseline_data = train_baseline(zero_stage)
+        accelerator_not_trained, accelerator_trained, accelerator_outputs, accelerator_data = train_integration(
+            zero_stage
+        )
+        assert baseline_not_trained["accuracy"] == accelerator_not_trained["accuracy"], (
+            f"ZERO stage {zero_stage}: Accuracy should be the same for the baseline and accelerator: {baseline_not_trained['accuracy']} == {accelerator_not_trained['accuracy']}"
+        )
+        assert baseline_not_trained["f1"] == accelerator_not_trained["f1"], (
+            f"ZERO stage {zero_stage}: F1 score should be the same for the baseline and accelerator: {baseline_not_trained['f1']} == {accelerator_not_trained['f1']}"
+        )
+        assert baseline_trained["accuracy"] == accelerator_trained["accuracy"], (
+            f"ZERO stage {zero_stage}: Accuracy should be the same for the baseline and accelerator: {baseline_trained['accuracy']} == {accelerator_trained['accuracy']}"
+        )
+        assert baseline_trained["f1"] == accelerator_trained["f1"], (
+            f"ZERO stage {zero_stage}: F1 score should be the same for the baseline and accelerator: {baseline_trained['f1']} == {accelerator_trained['f1']}"
+        )
+
+        torch.distributed.destroy_process_group()

benchmarks/fp8/transformer_engine/fp8_utils.py

@@ -0,0 +1,116 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import torch
+
+
+def get_dataloaders(model_name: str, batch_size: int = 16):
+    from datasets import load_dataset
+    from torch.utils.data import DataLoader
+    from transformers import AutoTokenizer
+
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    datasets = load_dataset("glue", "mrpc")
+
+    def tokenize_function(examples):
+        # max_length=None => use the model max length (it's actually the default)
+        outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None)
+        return outputs
+
+    # Apply the method we just defined to all the examples in all the splits of the dataset
+    # starting with the main process first:
+    tokenized_datasets = datasets.map(
+        tokenize_function,
+        batched=True,
+        remove_columns=["idx", "sentence1", "sentence2"],
+    )
+
+    # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
+    # transformers library
+    tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
+
+    def collate_fn(examples):
+        return tokenizer.pad(
+            examples,
+            padding="longest",
+            pad_to_multiple_of=16,  # Specific for FP8
+            return_tensors="pt",
+        )
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(
+        tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size, drop_last=True
+    )
+    eval_dataloader = DataLoader(
+        tokenized_datasets["validation"],
+        shuffle=False,
+        collate_fn=collate_fn,
+        batch_size=16,
+        drop_last=True,
+    )
+
+    return train_dataloader, eval_dataloader
+
+
+def get_training_utilities(model_name: str, batch_size: int = 16, accelerator=None):
+    """
+    Returns a tuple of:
+        - Model
+        - Optimizer
+        - Train dataloader (prepared)
+        - Eval dataloader (prepared)
+        - LR Scheduler
+    Suitable for training on the MRPC dataset
+    """
+    from torch.optim import AdamW
+    from transformers import AutoModelForSequenceClassification, get_linear_schedule_with_warmup
+
+    from accelerate import Accelerator
+
+    if accelerator is None:
+        accelerator = Accelerator()
+    model = AutoModelForSequenceClassification.from_pretrained(model_name)
+    train_dataloader, eval_dataloader = get_dataloaders(model_name, batch_size)
+    optimizer = AdamW(model.parameters(), lr=0.0001)
+    lr_scheduler = get_linear_schedule_with_warmup(
+        optimizer=optimizer,
+        num_warmup_steps=100,
+        num_training_steps=len(train_dataloader) * 2,
+    )
+    train_dataloader, eval_dataloader = accelerator.prepare(train_dataloader, eval_dataloader)
+    return model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+
+
+def get_named_parameters(model):
+    """
+    Same thing as `Accelerator.get_named_parameters` Returns a list of the named parameters of the model (extracted
+    from parallel)
+    """
+    from accelerate.utils import extract_model_from_parallel
+
+    model = extract_model_from_parallel(model)
+    return {n: p for n, p in model.named_parameters()}
+
+
+def evaluate_model(model, dataloader, metric, accelerator=None):
+    "Turns model to .eval(), runs dataloader, calculates metric, then turns eval back on"
+    model.eval()
+    for step, batch in enumerate(dataloader):
+        with torch.no_grad():
+            outputs = model(**batch)
+        predictions = outputs.logits.argmax(dim=-1)
+        references = batch["labels"]
+        if accelerator is not None and accelerator.num_processes > 1:
+            predictions, references = accelerator.gather_for_metrics((predictions, references))
+        metric.add_batch(predictions=predictions, references=references)
+    return metric.compute()

benchmarks/fp8/transformer_engine/fsdp.py

@@ -0,0 +1,161 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+This script tests to ensure that `accelerate` performs at the same level as raw `TransformersEngine`.
+
+This particular script verifies this for FSDP training.
+"""
+
+from functools import partial
+
+import evaluate
+import torch
+import transformer_engine.common.recipe as te_recipe
+import transformer_engine.pytorch as te
+from fp8_utils import evaluate_model, get_named_parameters, get_training_utilities
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+from torch.distributed.fsdp import MixedPrecision
+from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
+from transformer_engine.common.recipe import DelayedScaling
+from transformers.models.bert import BertLayer
+
+from accelerate import Accelerator
+from accelerate import FullyShardedDataParallelPlugin as FSDPPlugin
+from accelerate.state import AcceleratorState
+from accelerate.utils import FP8RecipeKwargs, set_seed
+from accelerate.utils.transformer_engine import convert_model
+
+
+MODEL_NAME = "bert-base-cased"
+METRIC = evaluate.load("glue", "mrpc")
+
+FSDP_WRAP_POLICY = partial(transformer_auto_wrap_policy, transformer_layer_cls={BertLayer})
+
+
+def train_baseline():
+    set_seed(42)
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = get_training_utilities(MODEL_NAME)
+    accelerator = Accelerator()
+    device = accelerator.device
+    model.to(device)
+
+    # Convert the model to TE
+    old_named_params = get_named_parameters(model)
+
+    with torch.no_grad():
+        convert_model(model)
+
+    FP8_RECIPE_KWARGS = {"fp8_format": te_recipe.Format.HYBRID, "amax_history_len": 32, "amax_compute_algo": "max"}
+    fp8_recipe = DelayedScaling(**FP8_RECIPE_KWARGS)
+
+    new_named_params = get_named_parameters(model)
+
+    # Convert the model to FSDP
+    model = FSDP(
+        model,
+        use_orig_params=True,
+        mixed_precision=MixedPrecision(param_dtype=torch.bfloat16, reduce_dtype=torch.float32),
+        auto_wrap_policy=FSDP_WRAP_POLICY,
+    )
+
+    mapping = {p: new_named_params[n] for n, p in old_named_params.items()}
+    for param_group in optimizer.param_groups:
+        param_group["params"] = [mapping[p] for p in param_group["params"]]
+
+    base_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)
+    model.train()
+
+    for _ in range(2):
+        for batch in train_dataloader:
+            with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
+                with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+                    batch = batch.to(device)
+                    outputs = model(**batch)
+            loss = outputs.loss
+            loss.backward()
+            optimizer.step()
+            optimizer.zero_grad()
+            lr_scheduler.step()
+
+    trained_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)
+
+    assert trained_model_results["accuracy"] > base_model_results["accuracy"], (
+        f"Accuracy should be higher for the trained model: {trained_model_results['accuracy']} > {base_model_results['accuracy']}"
+    )
+    assert trained_model_results["f1"] > base_model_results["f1"], (
+        f"F1 score should be higher for the trained model: {trained_model_results['f1']} > {base_model_results['f1']}"
+    )
+
+    return base_model_results, trained_model_results
+
+
+def train_integration():
+    FP8_RECIPE_KWARGS = {"fp8_format": "HYBRID", "amax_history_len": 32, "amax_compute_algo": "max"}
+    kwargs_handlers = [FP8RecipeKwargs(backend="TE", **FP8_RECIPE_KWARGS)]
+    AcceleratorState()._reset_state(True)
+    fsdp_plugin = FSDPPlugin(
+        auto_wrap_policy=FSDP_WRAP_POLICY,
+        use_orig_params=True,
+        mixed_precision_policy=MixedPrecision(param_dtype=torch.bfloat16, reduce_dtype=torch.float32),
+    )
+    accelerator = Accelerator(mixed_precision="fp8", fsdp_plugin=fsdp_plugin, kwargs_handlers=kwargs_handlers)
+    set_seed(42)
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = get_training_utilities(
+        MODEL_NAME, accelerator=accelerator
+    )
+
+    model, optimizer = accelerator.prepare(model, optimizer)
+    base_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)
+    model.train()
+
+    for _ in range(2):
+        for batch in train_dataloader:
+            outputs = model(**batch)
+            loss = outputs.loss
+            accelerator.backward(loss)
+            optimizer.step()
+            optimizer.zero_grad()
+            lr_scheduler.step()
+
+    trained_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)
+
+    assert trained_model_results["accuracy"] > base_model_results["accuracy"], (
+        f"Accuracy should be higher for the trained model: {trained_model_results['accuracy']} > {base_model_results['accuracy']}"
+    )
+    assert trained_model_results["f1"] > base_model_results["f1"], (
+        f"F1 score should be higher for the trained model: {trained_model_results['f1']} > {base_model_results['f1']}"
+    )
+
+    return base_model_results, trained_model_results
+
+
+if __name__ == "__main__":
+    baseline_not_trained, baseline_trained = train_baseline()
+    accelerator_not_trained, accelerator_trained = train_integration()
+
+    assert baseline_not_trained["accuracy"] == accelerator_not_trained["accuracy"], (
+        f"Accuracy should be the same for the baseline and accelerator: {baseline_not_trained['accuracy']} == {accelerator_not_trained['accuracy']}"
+    )
+    assert baseline_not_trained["f1"] == accelerator_not_trained["f1"], (
+        f"F1 score should be the same for the baseline and accelerator: {baseline_not_trained['f1']} == {accelerator_not_trained['f1']}"
+    )
+    assert baseline_trained["accuracy"] == accelerator_trained["accuracy"], (
+        f"Accuracy should be the same for the baseline and accelerator: {baseline_trained['accuracy']} == {accelerator_trained['accuracy']}"
+    )
+    assert baseline_trained["f1"] == accelerator_trained["f1"], (
+        f"F1 score should be the same for the baseline and accelerator: {baseline_trained['f1']} == {accelerator_trained['f1']}"
+    )
+
+    torch.distributed.destroy_process_group()

benchmarks/fp8/transformer_engine/non_distributed.py

@@ -0,0 +1,132 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+This script tests to ensure that `accelerate` performs at the same level as raw `TransformersEngine`.
+
+This particular script verifies this for single GPU training.
+"""
+
+import evaluate
+import torch
+import transformer_engine.common.recipe as te_recipe
+import transformer_engine.pytorch as te
+from fp8_utils import evaluate_model, get_named_parameters, get_training_utilities
+from transformer_engine.common.recipe import DelayedScaling
+
+from accelerate import Accelerator
+from accelerate.state import AcceleratorState
+from accelerate.utils import FP8RecipeKwargs, set_seed
+from accelerate.utils.transformer_engine import convert_model
+
+
+MODEL_NAME = "bert-base-cased"
+METRIC = evaluate.load("glue", "mrpc")
+
+
+def train_baseline():
+    set_seed(42)
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = get_training_utilities(MODEL_NAME)
+
+    # Convert the model to TE
+    old_named_params = get_named_parameters(model)
+
+    with torch.no_grad():
+        convert_model(model)
+
+    new_named_params = get_named_parameters(model)
+    mapping = {p: new_named_params[n] for n, p in old_named_params.items()}
+    for param_group in optimizer.param_groups:
+        param_group["params"] = [mapping[p] for p in param_group["params"]]
+
+    FP8_RECIPE_KWARGS = {"fp8_format": te_recipe.Format.HYBRID, "amax_history_len": 32, "amax_compute_algo": "max"}
+    fp8_recipe = DelayedScaling(**FP8_RECIPE_KWARGS)
+
+    model.to("cuda")
+    base_model_results = evaluate_model(model, eval_dataloader, METRIC)
+    model.train()
+
+    for batch in train_dataloader:
+        with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
+            with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+                batch = batch.to("cuda")
+                outputs = model(**batch)
+        loss = outputs.loss
+        loss.backward()
+        optimizer.step()
+        optimizer.zero_grad()
+        lr_scheduler.step()
+
+    trained_model_results = evaluate_model(model, eval_dataloader, METRIC)
+
+    assert trained_model_results["accuracy"] > base_model_results["accuracy"], (
+        f"Accuracy should be higher for the trained model: {trained_model_results['accuracy']} > {base_model_results['accuracy']}"
+    )
+    assert trained_model_results["f1"] > base_model_results["f1"], (
+        f"F1 score should be higher for the trained model: {trained_model_results['f1']} > {base_model_results['f1']}"
+    )
+
+    return base_model_results, trained_model_results
+
+
+def train_integration():
+    FP8_RECIPE_KWARGS = {"fp8_format": "HYBRID", "amax_history_len": 32, "amax_compute_algo": "max"}
+    kwargs_handlers = [FP8RecipeKwargs(backend="TE", **FP8_RECIPE_KWARGS)]
+    AcceleratorState()._reset_state(True)
+    accelerator = Accelerator(mixed_precision="fp8", kwargs_handlers=kwargs_handlers)
+    set_seed(42)
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = get_training_utilities(
+        MODEL_NAME, accelerator=accelerator
+    )
+
+    model, optimizer, lr_scheduler = accelerator.prepare(model, optimizer, lr_scheduler)
+    base_model_results = evaluate_model(model, eval_dataloader, METRIC)
+    model.train()
+
+    for batch in train_dataloader:
+        outputs = model(**batch)
+        loss = outputs.loss
+        accelerator.backward(loss)
+        optimizer.step()
+        optimizer.zero_grad()
+        lr_scheduler.step()
+
+    trained_model_results = evaluate_model(model, eval_dataloader, METRIC)
+
+    assert trained_model_results["accuracy"] > base_model_results["accuracy"], (
+        f"Accuracy should be higher for the trained model: {trained_model_results['accuracy']} > {base_model_results['accuracy']}"
+    )
+    assert trained_model_results["f1"] > base_model_results["f1"], (
+        f"F1 score should be higher for the trained model: {trained_model_results['f1']} > {base_model_results['f1']}"
+    )
+
+    return base_model_results, trained_model_results
+
+
+if __name__ == "__main__":
+    baseline_not_trained, baseline_trained = train_baseline()
+    accelerator_not_trained, accelerator_trained = train_integration()
+
+    assert baseline_not_trained["accuracy"] == accelerator_not_trained["accuracy"], (
+        f"Accuracy should be the same for the baseline and accelerator: {baseline_not_trained['accuracy']} == {accelerator_not_trained['accuracy']}"
+    )
+    assert baseline_not_trained["f1"] == accelerator_not_trained["f1"], (
+        f"F1 score should be the same for the baseline and accelerator: {baseline_not_trained['f1']} == {accelerator_not_trained['f1']}"
+    )
+    assert baseline_trained["accuracy"] == accelerator_trained["accuracy"], (
+        f"Accuracy should be the same for the baseline and accelerator: {baseline_trained['accuracy']} == {accelerator_trained['accuracy']}"
+    )
+    assert baseline_trained["f1"] == accelerator_trained["f1"], (
+        f"F1 score should be the same for the baseline and accelerator: {baseline_trained['f1']} == {accelerator_trained['f1']}"
+    )

benchmarks/fsdp2/README.md

@@ -0,0 +1,74 @@
+# FSDP2 Benchmarks
+
+This benchmark showcases `FSDP2` in 🤗 `accelerate` and compares it to `torch` baseline.
+
+## Overview
+
+This benchmark consists of two parts:
+- `main.py` is the main script that runs the benchmark
+- `visualize.py` is the script that visualizes the results (if `--output_dir` was specified for the previous command)
+
+## Motivation
+
+We want to showcase that 🤗 `accelerate`'s integration of `FSDP2` is on par raw PyTorch, and highlight a "broken" part in PyTorch that creating an optimizer before applying `FSDP2` **doesn't result in a working training loop**. (more on this later)
+This script showcases **matching memory usage and convergence between `accelerate` and `torch`'s baseline.**
+To deal with this breaking change (and maintain backward compatibility with FSDP1 in terms of an API), `accelerate` had to come up with a workaround since `accelerate` assumes that the user will nearly always create a model, optimizer, scheduler, etc beforehand and bring them themselves. This lead to an issue of a stark increase in memory as well as the model not even training if the user creates an optimizer beforehand. 
+To workaround this, we replace the parameters inside the optimizer with the newly created FSDP2 sharded ones. More about this can be found in this [blog post (TBD)](TODO)
+> [!WARNING]
+> This script is intended to fit on 2x 24GB GPUs, though on so few GPUs it's not possible to see the memory difference (discrepancies in grad allocation result in lower memory usage in the non-fixed case), only the difference in convergence. Below are attached results from 8x H100 GPUs where the difference is visible.
+> TLDR: more GPUs = bigger memory difference between fixed and non-fixed cases.
+
+## Results
+
+Here are the results from running the benchmark on 8x H100 GPUs:
+
+<p align="center">
+  <img src="imgs/allocated_memory.png" width="80%" alt="Allocated Memory Usage">
+</p>
+<p align="center">
+  <img src="imgs/reserved_memory.png" width="80%" alt="Reserved Memory Usage">
+</p>
+
+As you can see, the memory usage of `accelerate` and `torch_post_shard` (the **intended** way) are very similar, while `torch_pre_shard_not_fixed` uses significantly more memory. Our fix in `torch_pre_shard_fixed` brings the memory usage back in line with the **intended** approach.
+
+> [!WARNING]
+> Timing discrepancies are due to the benchmarks being ran in 1 script.
+
+
+## Running
+
+To run the benchmark, you can either use `accelerate launch` or `torchrun`:
+```bash
+accelerate launch main.py
+```
+```bash
+# For two GPUs
+torchrun --nproc_per_node 2 main.py
+```
+
+This supports multiple configurable options, you can learn about them by running:
+```bash
+python3 main.py --help
+```
+
+This script will run 4 different benchmarks:
+- `torch_optimizer_after_fsdp`: `torch` baseline where optimizer is created after applying `FSDP2`, this is the **intended** way to do it
+- `torch_optimizer_before_fsdp_not_fixed`: `torch` baseline where optimizer is created before applying `FSDP2` without fixing the optimizer parameters
+- `torch_optimizer_before_fsdp_fixed`: `torch` baseline where optimizer is created before applying `FSDP2` with our fix to the optimizer
+- `accelerate`: `accelerate`'s own integration of `FSDP2` where optimizer is created before applying `FSDP2`, but we apply our fix to the optimizer
+
+Memory results are saved in a folder specified by `--output_dir` argument.
+Optionally, you can specify `--save_memory_snapshot` to save the torch memory snapshot, which can then be viewed using [`torch memory viz`](https://pytorch.org/memory_viz)
+
+## Visualizing results
+
+To visualize the results, you can run:
+
+```bash
+python3 visualize.py --dir <path_to_output_dir>
+```
+
+This will then create two plots, showcasing allocated and reserved memory usage between all the different benchmarks discussed above.
+
+
+

benchmarks/fsdp2/main.py

@@ -0,0 +1,122 @@
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import functools
+from typing import Callable
+
+import torch
+
+from accelerate import Accelerator
+from utils import parse_args, prepare_accelerate, prepare_torch
+
+
+MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct"
+LEARNING_RATE = 3e-5
+
+CONFIG = {
+    "model_name": MODEL_NAME,
+    "learning_rate": LEARNING_RATE,
+}
+
+
+def train(
+    model: torch.nn.Module,
+    optimizer: torch.optim.Optimizer,
+    train_dataloader: torch.utils.data.DataLoader,
+    accelerator: Accelerator,
+) -> torch.Tensor:
+    losses = []
+    for batch in train_dataloader:
+        optimizer.zero_grad()
+        outputs = model(**batch, use_cache=False)
+
+        loss = outputs.loss
+        losses.append(loss.item())
+        accelerator.backward(loss)
+        optimizer.step()
+
+    return torch.tensor(losses)
+
+
+def evaluate(args, config: dict, init_fn: Callable, run_name: str) -> torch.Tensor:
+    model, optimizer, dataloader, accelerator, memory_tracker = init_fn(args, config)
+
+    loss = train(model, optimizer, dataloader, accelerator)
+
+    memory_tracker.stop()
+    msg = f"""Results for {run_name} (rank 0):
+Loss: {loss[-1].item()}
+Peak Allocated Memory: {float(memory_tracker.peak_allocated_memory):.2f} MB
+Peak Reserved Memory: {float(memory_tracker.peak_reserved_memory):.2f} MB
+{"-" * 34}"""
+    accelerator.print(msg)
+    return loss
+
+
+def main():
+    args = parse_args()
+    evaluations = [
+        functools.partial(
+            evaluate,
+            init_fn=functools.partial(prepare_torch, post_shard_optimizer=False, apply_optimizer_fix=True),
+            run_name="Optimizer Before FSDP (w/ fix)",
+        ),
+        functools.partial(
+            evaluate,
+            init_fn=functools.partial(prepare_torch, post_shard_optimizer=False, apply_optimizer_fix=False),
+            run_name="Optimizer Before FSDP (w/o fix)",
+        ),
+        functools.partial(
+            evaluate,
+            init_fn=functools.partial(prepare_torch, post_shard_optimizer=True),
+            run_name="Optimizer After FSDP",
+        ),
+        functools.partial(evaluate, init_fn=prepare_accelerate, run_name="Accelerate"),
+    ]
+    labels = [
+        "Optimizer Before FSDP (w/ fix)",
+        "Optimizer Before FSDP (w/o fix)",
+        "Optimizer After FSDP",
+        "Accelerate",
+    ]
+
+    results = {}
+    torch.use_deterministic_algorithms(True)
+
+    for evaluation, label in zip(evaluations, labels):
+        results[label] = evaluation(args, CONFIG)
+
+    torch.testing.assert_close(
+        results["Optimizer After FSDP"],
+        results["Optimizer Before FSDP (w/ fix)"],
+        msg="Optimizer After FSDP and Optimizer Before FSDP (w/ fix) should be the same",
+    )
+
+    torch.testing.assert_close(
+        results["Optimizer After FSDP"],
+        results["Accelerate"],
+        msg="Optimizer After FSDP and Accelerate should be the same",
+    )
+
+    torch.testing.assert_close(
+        results["Accelerate"],
+        results["Optimizer Before FSDP (w/ fix)"],
+        msg="Accelerate and Optimizer Before FSDP (w/ fix) should be the same",
+    )
+
+    torch.distributed.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()

benchmarks/fsdp2/measure_utils.py

@@ -0,0 +1,130 @@
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import gc
+import json
+import os
+import threading
+import time
+
+import psutil
+import torch
+
+from accelerate import PartialState
+
+
+class MemoryTracker:
+    def __init__(
+        self,
+        device: torch.device,
+        output_directory: str,
+        run_name: str,
+        save_memory_snapshot: bool,
+        log_interval: float = 0.01,
+    ):
+        """Class for tracking gpu and cpu memory usage of the process.
+
+        Args:
+            device (`torch.device`):
+                PyTorch device to monitor.
+            output_directory (`str`):
+                Directory to save the memory usage data to, will be created if it doesn't exist.
+            run_name (`str`):
+                Name of the run, will be used to name the output files.
+            save_memory_snapshot (`bool`):
+                Whether to also save `torch.cuda.memory._dump_snapshot` to the output directory.
+            log_interval (`float`, *optional*):
+                Interval in seconds between memory measurements. Defaults to 0.01.
+        """
+        self.log_interval = log_interval
+        self.save_memory_snapshot = save_memory_snapshot
+        self.output_directory = output_directory
+        self.run_name = run_name
+
+        self.timestamps = []
+        self.allocated_memory = []
+        self.reserved_memory = []
+        self.virtual_memory = []
+
+        self.start_time = None
+        self.running = False
+
+        self._thread = None
+        self._state = PartialState()
+        self._process = psutil.Process()
+        self._device = device
+        self.torch_accelerator_module = getattr(torch, device.type, torch.cuda)
+
+    def _monitor(self):
+        self.start_time = time.time()
+
+        while self.running:
+            allocated = self.torch_accelerator_module.memory_allocated(self._device) / (1024 * 1024)
+            reserved = self.torch_accelerator_module.memory_reserved(self._device) / (1024 * 1024)
+            virtual_memory = self._process.memory_info().rss / (1024 * 1024)
+
+            self.allocated_memory.append(allocated)
+            self.reserved_memory.append(reserved)
+            self.virtual_memory.append(virtual_memory)
+            self.timestamps.append(time.time() - self.start_time)
+
+            time.sleep(self.log_interval)
+
+    def start(self):
+        gc.collect()
+        self.torch_accelerator_module.empty_cache()
+
+        if self.output_directory:
+            os.makedirs(self.output_directory, exist_ok=True)
+
+        if self.save_memory_snapshot:
+            self.torch_accelerator_module.memory._record_memory_history()
+
+        self.running = True
+        self._thread = threading.Thread(target=self._monitor)
+        self._thread.daemon = True
+        self._thread.start()
+
+    def stop(self):
+        self.running = False
+        if self._thread:
+            self._thread.join()
+
+        if self.save_memory_snapshot and self._state.is_main_process and self.output_directory:
+            output_file = os.path.join(self.output_directory, f"{self.run_name}_memory_snapshot.pkl")
+            self.torch_accelerator_module.memory._dump_snapshot(output_file)
+
+        if self._state.is_main_process and self.output_directory:
+            path = os.path.join(self.output_directory, f"{self.run_name}_memory_usage.json")
+            with open(path, "w") as f:
+                json.dump(
+                    {
+                        "timestamps": self.timestamps,
+                        "allocated_memory": self.allocated_memory,
+                        "reserved_memory": self.reserved_memory,
+                        "virtual_memory": self.virtual_memory,
+                    },
+                    f,
+                )
+        if self.save_memory_snapshot:
+            self.torch_accelerator_module.memory._record_memory_history(False)
+        self.torch_accelerator_module.empty_cache()
+
+    @property
+    def peak_allocated_memory(self):
+        return max(self.allocated_memory)
+
+    @property
+    def peak_reserved_memory(self):
+        return max(self.reserved_memory)

benchmarks/fsdp2/utils.py

@@ -0,0 +1,290 @@
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+from types import MethodType
+from typing import Union
+
+import torch
+from datasets import load_dataset
+from measure_utils import MemoryTracker
+from torch.distributed.fsdp import MixedPrecisionPolicy, fully_shard
+from torch.optim import AdamW
+from torch.utils.data import DataLoader
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, DataCollatorForLanguageModeling
+from transformers.models.qwen2.modeling_qwen2 import Qwen2DecoderLayer
+
+from accelerate import Accelerator, FullyShardedDataParallelPlugin
+from accelerate.state import AcceleratorState, is_initialized
+from accelerate.utils import convert_outputs_to_fp32, set_seed
+
+
+SEED = 421
+
+
+def get_named_parameters(model: torch.nn.Module, drop_refs: bool = False) -> dict[str, Union[torch.Tensor, int]]:
+    """
+    This function returns a dictionary mapping the parameter names to their data pointers or
+    the original parameters if `drop_refs` is `False`.
+    It is used to get the original parameter names before `fully_shard` is applied.
+
+    We only return the data pointers, so we drop the references to the original parameters
+    and `fully_shard` will then trigger a new allocation for the sharded ones.
+
+    Args:
+        model (`torch.nn.Module`): Model instance to get the named parameters from
+        drop_refs (`bool`, *optional*, defaults to `False`): Whether to drop the references to the original parameters
+
+    Returns:
+        `dict[str, Union[torch.Tensor, int]]`: Dictionary mapping the parameter names to their data pointers or the original parameters if `drop_refs` is `False`
+    """
+    named_parameters = {}
+    for n, p in model.named_parameters():
+        # We only preserve the data pointers to have the unique 1:1 mapping between the original and the sharded parameters
+        named_parameters[n] = p.data_ptr() if drop_refs else p
+    return named_parameters
+
+
+def replace_optimizer_params(optimizer: torch.optim.Optimizer):
+    """
+    This function is called before using `fully_shard` on the model. It replaces the parameters of the optimizer with
+    empty tensors, so `fully_shard` can trigger a new allocation for the sharded ones. After this, we swap the parameters
+    `data_ptr` to the original one, so we can reuse that later to map the sharded parameters to the original ones.
+    This function modifies the optimizer in-place.
+
+    Args:
+        optimizer (torch.optim.Optimizer): Optimizer instance which contains the original model parameters
+    """
+
+    for param_group in optimizer.param_groups:
+        for i, p in enumerate(param_group["params"]):
+            # We drop a reference to the original param here, so that _move_states_to_device triggers a reallocation
+            # This is required or else the `fully_shard` -> `_move_states_to_device` uses the original memory address
+            # for the sharded parameters, and we get a weird/undefined behavior.
+            param_group["params"][i] = torch.empty_like(p)
+
+            # We save the original data_ptr, so we can swap back the parameters later
+            param_group["params"][i].data_ptr = p.data_ptr()
+
+
+def swap_back_optimizer_params(
+    model: torch.nn.Module, optimizer: torch.optim.Optimizer, old_named_parameter_pointers: dict[str, int]
+):
+    """
+    This function is the counterpart of `replace_optimizer_params`. It is called after `fully_shard` being applied to
+    the model. It swaps the parameters of the optimizer to their sharded counterparts.
+    It is done using the `data_ptr` mapping prepared in `replace_optimizer_params` and `get_named_parameters`.
+
+    Args:
+        model (`torch.nn.Module`): Model instance to get the new named parameters from
+        optimizer (`torch.optim.Optimizer`): Optimizer instance to swap the parameters of
+        old_named_parameter_pointers (`dict[str, int]`): Dictionary mapping the original parameter names: data_ptrs to the new ones
+    """
+    # We get the new named parameters after `fully_shard` being applied
+    # We don't drop the references as we need the sharded parameters now
+    new_named_parameters = get_named_parameters(model, drop_refs=False)
+
+    # We create a mapping from the original data_ptr to the new sharded param corresponding to it
+    mapping = {p: new_named_parameters[n] for n, p in old_named_parameter_pointers.items()}
+
+    for param_group in optimizer.param_groups:
+        # We swap the parameters of the optimizer to the new sharded ones
+        param_group["params"] = [mapping[p.data_ptr] for p in param_group["params"]]
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        help="Directory to save the benchmarking results.",
+    )
+    parser.add_argument(
+        "--save_memory_snapshot",
+        action="store_true",
+        default=False,
+        help="If True, `torch.cuda.memory._dump_snapshot` will be used to additionaly save the memory trace.",
+    )
+    ######################
+    # Training arguments #
+    ######################
+    parser.add_argument(
+        "--batch_size",
+        type=int,
+        default=2,
+        help="Batch size for the training loop.",
+    )
+    parser.add_argument(
+        "--block_size",
+        type=int,
+        default=128,
+        help="The maximum sequence length to use with the model.",
+    )
+    parser.add_argument(
+        "--dataset_fraction",
+        type=float,
+        default=1.0,
+        help="Fraction of the dataset to use.",
+    )
+    return parser.parse_args()
+
+
+def prepare_dataloader(tokenizer, args, accelerator: Accelerator) -> DataLoader:
+    dataset = load_dataset("tiny_shakespeare", split="train", trust_remote_code=True)
+
+    def tokenize_function(example):
+        return tokenizer(
+            example["text"],
+        )
+
+    dataset = dataset.map(
+        tokenize_function,
+        batched=True,
+        remove_columns=["text"],
+    )
+
+    block_size = min(tokenizer.model_max_length, args.block_size)
+
+    def group_texts(examples):
+        concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}
+        total_length = len(concatenated_examples[list(examples.keys())[0]])
+
+        total_length = (total_length // block_size) * block_size
+
+        result = {
+            k: [t[i : i + block_size] for i in range(0, total_length, block_size)]
+            for k, t in concatenated_examples.items()
+        }
+
+        result["labels"] = result["input_ids"].copy()
+        return result
+
+    dataset = dataset.map(group_texts, batched=True)
+    dataset = dataset.select(range(int(len(dataset) * args.dataset_fraction)))
+
+    def collate_fn(examples):
+        return DataCollatorForLanguageModeling(
+            tokenizer=tokenizer,
+            mlm=False,
+        )(examples)
+
+    dataloader = DataLoader(
+        dataset,
+        batch_size=args.batch_size,
+        collate_fn=collate_fn,
+    )
+    dataloader = accelerator.prepare(dataloader)
+    return dataloader
+
+
+def get_model(model_name: str):
+    # We reguire model to be loaded in fp32, otherwise benchmarks don't match as accelerate does upcasting of parameters to fp32
+    config = AutoConfig.from_pretrained(model_name, trust_remote_code=True, torch_dtype=torch.float32)
+    model = AutoModelForCausalLM.from_config(config)
+    return model
+
+
+def get_tokenizer(model_name: str):
+    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+    tokenizer.pad_token = tokenizer.eos_token
+    return tokenizer
+
+
+def prepare_torch(
+    args, config: dict, post_shard_optimizer: bool = False, apply_optimizer_fix: bool = False
+) -> tuple[torch.nn.Module, torch.optim.Optimizer, torch.utils.data.DataLoader, Accelerator]:
+    mp_policy = MixedPrecisionPolicy(
+        param_dtype=torch.bfloat16,
+        reduce_dtype=torch.bfloat16,
+        output_dtype=torch.bfloat16,
+    )
+
+    accelerator = Accelerator(mixed_precision="bf16")
+    set_seed(SEED)
+    is_fixed = "fixed" if apply_optimizer_fix else "not_fixed"
+    is_post_shard = "optimizer_after_fsdp" if post_shard_optimizer else "optimizer_before_fsdp"
+    run_name = f"torch_{is_post_shard}" if post_shard_optimizer else f"torch_{is_post_shard}_{is_fixed}"
+
+    tokenizer = get_tokenizer(config["model_name"])
+    train_dataloader = prepare_dataloader(tokenizer, args, accelerator)
+
+    memory_tracker = MemoryTracker(accelerator.device, args.output_dir, run_name, args.save_memory_snapshot)
+    memory_tracker.start()
+
+    model = get_model(config["model_name"])
+    optimizer = None
+
+    if not post_shard_optimizer:
+        optimizer = AdamW(model.parameters(), lr=config["learning_rate"])
+
+        if apply_optimizer_fix:
+            # We drop the references to the original parameters, so that `fully_shard` can trigger a new allocation
+            # Then we get the `module_name: data_ptr` mapping, so we can swap back the parameters later
+            old_named_parameters = get_named_parameters(model, drop_refs=True)
+
+            # We replace the parameters of the optimizer with empty tensors, so that `fully_shard` can trigger a new allocation
+            # We also change the `data_ptr` of the parameters to the original ones, so we can swap back the parameters later
+            replace_optimizer_params(optimizer)
+
+    for module in model.modules():
+        if isinstance(module, Qwen2DecoderLayer):
+            fully_shard(module, mp_policy=mp_policy)
+    fully_shard(model, mp_policy=mp_policy)
+
+    # We do this to imitate how accelerate forces outputs to be in fp32 via `convert_outputs_to_fp32`
+    autocast_context = torch.autocast(device_type=accelerator.state.device.type, dtype=torch.bfloat16)
+    model_forward_func = model.forward.__func__
+    new_forward = autocast_context(model_forward_func)
+    model.forward = MethodType(new_forward, model)
+    model.forward = MethodType(convert_outputs_to_fp32(model.forward.__func__), model)
+
+    if post_shard_optimizer:
+        optimizer = AdamW(model.parameters(), lr=config["learning_rate"])
+
+    if not post_shard_optimizer and apply_optimizer_fix:
+        # We swap back the parameters of the optimizer to the original ones
+        swap_back_optimizer_params(model, optimizer, old_named_parameters)
+
+    return model, optimizer, train_dataloader, accelerator, memory_tracker
+
+
+def prepare_accelerate(
+    args, config: dict
+) -> tuple[torch.nn.Module, torch.optim.Optimizer, torch.utils.data.DataLoader, Accelerator]:
+    if is_initialized():
+        AcceleratorState()._reset_state(True)
+
+    fsdp_plugin = FullyShardedDataParallelPlugin(
+        fsdp_version=2,
+        auto_wrap_policy="transformer_based_wrap",
+        transformer_cls_names_to_wrap=["Qwen2DecoderLayer"],
+    )
+    accelerator = Accelerator(
+        fsdp_plugin=fsdp_plugin,
+        mixed_precision="bf16",
+    )
+    set_seed(SEED)
+
+    tokenizer = get_tokenizer(config["model_name"])
+    train_dataloader = prepare_dataloader(tokenizer, args, accelerator)
+
+    memory_tracker = MemoryTracker(accelerator.device, args.output_dir, "accelerate", args.save_memory_snapshot)
+    memory_tracker.start()
+
+    model = get_model(config["model_name"])
+    optimizer = AdamW(model.parameters(), lr=config["learning_rate"])
+
+    model, optimizer = accelerator.prepare(model, optimizer)
+
+    return model, optimizer, train_dataloader, accelerator, memory_tracker

benchmarks/fsdp2/visualize.py

@@ -0,0 +1,114 @@
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import json
+
+import matplotlib.pyplot as plt
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--dir", type=str, help="Directory containing the memory usage data")
+    parser.add_argument(
+        "--memory_threshold",
+        type=int,
+        default=0,
+        help="Memory threshold to filter data that is below this value (only filters 1st `--filter_partition` of the points which should roughtly correspond to the model loading)",
+    )
+    parser.add_argument(
+        "--filter_partition",
+        type=float,
+        default=1 / 3,
+        help="Partition to drop data from that are below the memory threshold",
+    )
+    return parser.parse_args()
+
+
+def filter_data(data, memory_threshold, filter_partition, key):
+    timestamps = data["timestamps"]
+    memory = data[key]
+
+    mid_point = int(len(timestamps) * filter_partition)
+    filtered_times = []
+    filtered_memory = []
+    for i, (t, m) in enumerate(zip(timestamps, memory)):
+        if i < mid_point and m < memory_threshold:
+            continue
+        filtered_times.append(t)
+        filtered_memory.append(m)
+    return filtered_times, filtered_memory
+
+
+def compare_memory_usage(data, labels, memory_threshold, filter_partition):
+    plt.style.use("seaborn-v0_8")
+    colors = ["#2ecc71", "#e74c3c", "#3498db", "#f1c40f"]
+
+    fig1, ax1 = plt.subplots(figsize=(15, 5))
+    for data_item, label, color in zip(data, labels, colors):
+        timestamps, allocated = filter_data(data_item, memory_threshold, filter_partition, "allocated_memory")
+        ax1.plot(timestamps, allocated, label=label, color=color, linewidth=2)
+
+    ax1.set_xlabel("Time (s)", fontsize=12)
+    ax1.set_ylabel("Allocated Memory (GB)", fontsize=12)
+    ax1.set_title("Allocated Memory Usage Over Time", fontsize=14, pad=15)
+    ax1.grid(True, linestyle="--", alpha=0.7)
+    ax1.legend(frameon=True, fancybox=True, shadow=True, fontsize=10)
+    ax1.spines["top"].set_visible(False)
+    ax1.spines["right"].set_visible(False)
+    plt.tight_layout()
+
+    fig2, ax2 = plt.subplots(figsize=(15, 5))
+    for data_item, label, color in zip(data, labels, colors):
+        timestamps, reserved = filter_data(data_item, memory_threshold, filter_partition, "reserved_memory")
+        ax2.plot(timestamps, reserved, label=label, color=color, linewidth=2)
+
+    ax2.set_xlabel("Time (s)", fontsize=12)
+    ax2.set_ylabel("Reserved Memory (GB)", fontsize=12)
+    ax2.set_title("Reserved Memory Usage Over Time", fontsize=14, pad=15)
+    ax2.grid(True, linestyle="--", alpha=0.7)
+    ax2.legend(frameon=True, fancybox=True, shadow=True, fontsize=10)
+    ax2.spines["top"].set_visible(False)
+    ax2.spines["right"].set_visible(False)
+    plt.tight_layout()
+
+    return fig1, fig2
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    DIR = args.dir
+    with open(f"{DIR}/torch_optimizer_before_fsdp_not_fixed_memory_usage.json") as f:
+        optimizer_before_fsdp_not_fixed = json.load(f)
+
+    with open(f"{DIR}/torch_optimizer_after_fsdp_memory_usage.json") as f:
+        optimizer_after_fsdp = json.load(f)
+
+    with open(f"{DIR}/torch_optimizer_before_fsdp_fixed_memory_usage.json") as f:
+        optimizer_before_fsdp_fixed = json.load(f)
+
+    with open(f"{DIR}/accelerate_memory_usage.json") as f:
+        accelerate = json.load(f)
+
+    data = [optimizer_before_fsdp_not_fixed, optimizer_before_fsdp_fixed, optimizer_after_fsdp, accelerate]
+    labels = [
+        "Optimizer Before FSDP (w/o fix)",
+        "Optimizer Before FSDP (w/ fix)",
+        "Optimizer After FSDP",
+        "Accelerate",
+    ]
+
+    fig1, fig2 = compare_memory_usage(data, labels, args.memory_threshold, args.filter_partition)
+    fig1.savefig(f"{DIR}/allocated_memory.png")
+    fig2.savefig(f"{DIR}/reserved_memory.png")

benchmarks/torch.compile/README.md

@@ -0,0 +1,111 @@
+# Regional Compilation Benchmark
+
+This benchmark compares different compilation strategies using PyTorch's `torch.compile` and Accelerate's `compile_regions` utility, which is based on the recipe in [PyTorch documentation](https://pytorch.org/tutorials/recipes/regional_compilation.html).
+
+## Overview
+
+The benchmark evaluates three approaches:
+
+- **Baseline**: No compilation, standard PyTorch eager execution.
+- **Full compilation**: Using PyTorch's `torch.compile()` on the entire model.
+- **Regional compilation**: Using `accelerate.utils.compile_regions()` which targets specific blocks of the model to optimize compilation time.
+
+Each approach is tested with different batch sizes (1 and 4) and sequence lengths (128) on various LLaMA-based models ranging from 1B to 13B parameters. We purposefully run the forward pass outside of the `torch.no_grad()` context to simulate performance in a training environment, where gradients are needed.
+
+## Usage
+
+To run this benchmark:
+
+```bash
+python regional_compilation.py
+```
+
+The script will automatically download the model configurations, create models, and benchmark both compilation and inference times across different scenarios.
+
+## Requirements
+
+- Suitable GPU memory for the models being tested.
+- PyTorch with CUDA support.
+- Transformers library.
+- Accelerate library.
+
+## Results
+
+The benchmark results are summarized in the following figures:
+
+- Compilation time is how long it takes to run the first forward pass.
+- Speedup factor is the ratio of non-compiled baseline inference time to the fully/regionally compiled inference time.
+
+<p align="center">
+  <img src="imgs/compilation_time.png" width="80%" alt="Compilation Time">
+</p>
+<p align="center">
+  <img src="imgs/speedup_factor.png" width="80%" alt="Speedup Factor">
+</p>
+
+Full results are available in the tables below:
+
+```markdown
+[-------------------------------------------------- NousResearch/Llama-3.2-1B ---------------------------------------------------]
+                            |  Inference time (1x128)  |  Inference time (4x128)  |  Compile time (1x128)  |  Compile time (4x128)
+1 threads: -----------------------------------------------------------------------------------------------------------------------
+      Baseline              |           18.3           |           18.4           |                        |                      
+      Full compilation      |            6.3           |           10.0           |        10696.4         |        10248.0       
+      Regional compilation  |            9.7           |           10.0           |         1952.7         |         2903.9       
+
+Times are in milliseconds (ms).
+
+[---------------------------------------------- NousResearch/Hermes-3-Llama-3.2-3B ----------------------------------------------]
+                            |  Inference time (1x128)  |  Inference time (4x128)  |  Compile time (1x128)  |  Compile time (4x128)
+1 threads: -----------------------------------------------------------------------------------------------------------------------
+      Baseline              |           33.4           |           33.6           |                        |                      
+      Full compilation      |           11.2           |           23.9           |        17857.5         |        17736.5       
+      Regional compilation  |           17.3           |           23.7           |         2993.2         |         2478.8       
+
+Times are in milliseconds (ms).
+
+[---------------------------------------------- NousResearch/Hermes-3-Llama-3.1-8B ----------------------------------------------]
+                            |  Inference time (1x128)  |  Inference time (4x128)  |  Compile time (1x128)  |  Compile time (4x128)
+1 threads: -----------------------------------------------------------------------------------------------------------------------
+      Baseline              |           40.3           |           59.5           |                        |                      
+      Full compilation      |           18.9           |           54.4           |        20437.8         |        20152.3       
+      Regional compilation  |           19.7           |           54.0           |         2903.1         |         2438.0       
+
+Times are in milliseconds (ms).
+
+[--------------------------------------------- NousResearch/Nous-Hermes-Llama2-13b ----------------------------------------------]
+                            |  Inference time (1x128)  |  Inference time (4x128)  |  Compile time (1x128)  |  Compile time (4x128)
+1 threads: -----------------------------------------------------------------------------------------------------------------------
+      Baseline              |           45.5           |          100.4           |                        |                      
+      Full compilation      |           29.4           |           89.7           |        23099.4         |        22885.9       
+      Regional compilation  |           29.4           |           87.5           |         2945.5         |         2526.2       
+
+Times are in milliseconds (ms).
+```
+
+## Results Summary
+
+### Compilation Time
+
+Regional compilation provides significantly faster compilation times compared to full model compilation:
+
+- **Full compilation**: Takes ~10-23 seconds depending on model size.
+- **Regional compilation**: Takes only ~2-3 seconds across all model sizes.
+- **Speed improvement**: Regional compilation is **5-9x faster** to compile.
+
+### Inference Time
+
+Regional compilation delivers inference performance close to full compilation:
+
+- For batch size 1:
+  - For smaller models (1B-3B): Full compilation has a slight edge over regional compilation.
+  - For larger models (8B-13B): Regional compilation performs similarly to full compilation.
+- For batch size 4: Regional compilation performs similarly to full compilation across all models.
+
+## Key Takeaways
+
+1. **Comparable Performance**: Regional compilation delivers performance speedups similar to full compilation, especially for larger models.
+2. **Faster Compilation**: Regional compilation significantly reduces the time taken to compile models, making it a more efficient choice for deployment.
+3. **Batch Size Impact**: At batch size 4, full compilation and regional compilation perform nearly identically.
+4. **Model Size Impact**: Even with a small batch size, full compilation and regional compilation perform similarly for larger models (8B-13B).
+5. **Practical Application**: For real-world applications, regional compilation is a practical choice for optimizing training cold start times, especially when working with large models.

benchmarks/torch.compile/regional_compilation.py

@@ -0,0 +1,77 @@
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+from torch.utils.benchmark import Compare, Timer
+from transformers import AutoConfig, AutoModelForCausalLM
+
+from accelerate.test_utils.testing import get_backend
+from accelerate.utils import compile_regions
+
+
+torch.set_float32_matmul_precision("high")
+
+COMPILE_ITERS = 2
+INFERENCE_ITERS = 100
+
+BASELINE = "Baseline"
+COMPILE_TIME = "Compile time"
+INFRENCE_TIME = "Inference time"
+FULL_COMPILATION = "Full compilation"
+REGIONAL_COMPILATION = "Regional compilation"
+
+INFRENCE_STMT = "model(input_ids, use_cache=False)"
+COMPILE_STMT = f"torch._dynamo.reset(); torch._inductor.utils.clear_inductor_caches(); {INFRENCE_STMT}"
+
+torch_device_type, _, _ = get_backend()
+
+results = []
+for model_id in [
+    # non-gated llama models
+    "NousResearch/Llama-3.2-1B",
+    "NousResearch/Hermes-3-Llama-3.2-3B",
+    "NousResearch/Hermes-3-Llama-3.1-8B",
+    "NousResearch/Nous-Hermes-Llama2-13b",
+]:
+    with torch.device(torch_device_type):
+        config = AutoConfig.from_pretrained(model_id)
+        model = AutoModelForCausalLM.from_config(config).to(dtype=torch.float16).eval()
+
+    full_compilation_model = torch.compile(model)
+    regional_compilation_model = compile_regions(model)
+
+    for model, sub_label, description, stmt, iters in [
+        (model, BASELINE, INFRENCE_TIME, INFRENCE_STMT, INFERENCE_ITERS),
+        (full_compilation_model, FULL_COMPILATION, COMPILE_TIME, COMPILE_STMT, COMPILE_ITERS),
+        (full_compilation_model, FULL_COMPILATION, INFRENCE_TIME, INFRENCE_STMT, INFERENCE_ITERS),
+        (regional_compilation_model, REGIONAL_COMPILATION, COMPILE_TIME, COMPILE_STMT, COMPILE_ITERS),
+        (regional_compilation_model, REGIONAL_COMPILATION, INFRENCE_TIME, INFRENCE_STMT, INFERENCE_ITERS),
+    ]:
+        for batch_size, sequence_length in [(1, 128), (4, 128)]:
+            input_ids = torch.randint(
+                0, 1000, size=(batch_size, sequence_length), dtype=torch.int64, device=torch_device_type
+            )
+            results.append(
+                Timer(
+                    label=model_id,
+                    sub_label=sub_label,
+                    description=f"{description} ({batch_size}x{sequence_length})",
+                    globals={"model": model, "input_ids": input_ids},
+                    stmt=stmt,
+                ).timeit(number=iters)
+            )
+
+compare = Compare(results)
+compare.colorize()
+compare.print()

docker/README.md

@@ -0,0 +1,74 @@
+<!---
+Copyright 2024 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+-->
+
+# Official Hugging Face Accelerate Docker Images
+
+Accelerate publishes a variety of docker versions as part of our CI that users can also use. These are stable images that Accelerate can run off of which comes with a variety of different setup configurations, all of which are officially hosted on [Docker Hub](https://hub.docker.com/r/huggingface/accelerate).
+
+A breakdown of each are given below
+
+## Naming Conventions
+
+Accelerate docker images follow a tagging convention of:
+
+```bash
+huggingface/accelerate:{accelerator}-{nightly,release}
+```
+
+`accelerator` in this instance is one of many applical pre-configured backend supports:
+* `gpu`: Comes compiled off of the `nvidia/cuda` image and includes core parts like `bitsandbytes`. Runs off python 3.9.
+* `cpu`: Comes compiled off of `python:3.9-slim` and is designed for non-CUDA based workloads.
+* More to come soon
+* `gpu-deepspeed`: Comes compiled off of the `nvidia/cuda` image and includes core parts like `bitsandbytes` as well as the latest `deepspeed` version. Runs off python 3.10.
+* `gpu-fp8-transformerengine`: Comes compiled off of `nvcr.io/nvidia/pytorch` and is specifically for running the `benchmarks/fp8` scripts on devices which support FP8 operations using the `TransformerEngine` library (RTX 4090, H100, etc)
+
+## Nightlies vs Releases
+
+Each release a new build is pushed with a version number included in the name. For a GPU-supported image of version 0.28.0 for instance, it would look like the following:
+
+```bash
+huggingface/accelerate:gpu-release-0.28.0
+```
+
+Nightlies contain two different image tags. There is a general `nightly` tag which is built each night, and a `nightly-YYYY-MM-DD` which corresponds to a build from a particular date.
+
+For instance, here is an example nightly CPU image from 3/14/2024
+
+```bash
+huggingface/accelerate:cpu-nightly-2024-03-14
+```
+
+## Running the images
+
+Each image comes compiled with `conda` and an `accelerate` environment contains all of the installed dependencies. 
+
+To pull down the latest nightly run:
+
+```bash
+docker pull huggingface/accelerate:gpu-nightly
+```
+
+To then run it in interactive mode with GPU-memory available, run:
+
+```bash
+docker container run --gpus all -it huggingface/accelerate:gpu-nightly
+```
+
+## DEPRECATED IMAGES
+
+CPU and GPU docker images were hosted at `huggingface/accelerate-gpu` and `huggingface/accelerate-cpu`. These builds are now outdated and will not receive updates. 
+
+The builds at the corresponding `huggingface/accelerate:{gpu,cpu}` contain the same `Dockerfile`, so it's as simple as changing the docker image to the desired ones from above. We will not be deleting these images for posterity, but they will not be receiving updates going forward.

docker/accelerate-cpu/Dockerfile

@@ -0,0 +1,35 @@
+# Builds CPU-only Docker image of PyTorch
+# Uses multi-staged approach to reduce size
+# Stage 1
+FROM python:3.9-slim as compile-image
+
+ARG DEBIAN_FRONTEND=noninteractive
+
+RUN apt update
+RUN apt-get install -y --no-install-recommends \
+    build-essential \
+    git \
+    gcc
+
+# Setup virtual environment for Docker
+ENV VIRTUAL_ENV=/opt/venv
+RUN python3 -m venv ${VIRTUAL_ENV}
+# Make sure we use the virtualenv
+ENV PATH="${VIRTUAL_ENV}/bin:$PATH"
+WORKDIR /workspace
+# Install specific CPU torch wheel to save on space
+RUN python3 -m pip install --upgrade --no-cache-dir pip
+RUN python3 -m pip install --no-cache-dir \
+    jupyter \
+    git+https://github.com/huggingface/accelerate#egg=accelerate[testing,test_trackers] \
+    --extra-index-url https://download.pytorch.org/whl/cpu
+    
+# Stage 2
+FROM python:3.9-slim AS build-image
+COPY --from=compile-image /opt/venv /opt/venv
+RUN useradd -ms /bin/bash user
+USER user
+
+# Make sure we use the virtualenv
+ENV PATH="/opt/venv/bin:$PATH"
+CMD ["/bin/bash"]

docker/accelerate-gpu-deepspeed/Dockerfile

@@ -0,0 +1,46 @@
+# Builds GPU docker image of PyTorch specifically
+# Uses multi-staged approach to reduce size
+# Stage 1
+# Use base conda image to reduce time
+FROM continuumio/miniconda3:latest AS compile-image
+# Specify py version
+# Note: DeepSpeed beyond v0.12.6 requires py 3.10
+ENV PYTHON_VERSION=3.10
+# Install apt libs
+RUN apt-get update && \
+    apt-get install -y curl git wget && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists*
+
+# Create our conda env
+RUN conda create --name accelerate python=${PYTHON_VERSION} ipython jupyter pip
+# We don't install pytorch here yet since CUDA isn't available
+# instead we use the direct torch wheel
+ENV PATH /opt/conda/envs/accelerate/bin:$PATH
+# Activate our bash shell
+RUN chsh -s /bin/bash
+SHELL ["/bin/bash", "-c"]
+# Activate the conda env, install mpy4pi, and install torch + accelerate
+RUN source activate accelerate && conda install -c conda-forge mpi4py
+RUN source activate accelerate && \
+    python3 -m pip install --no-cache-dir \
+    git+https://github.com/huggingface/accelerate#egg=accelerate[testing,test_trackers,deepspeed] \
+    --extra-index-url https://download.pytorch.org/whl/cu126
+
+RUN python3 -m pip install --no-cache-dir bitsandbytes
+
+# Stage 2
+FROM nvidia/cuda:12.6.3-cudnn-devel-ubuntu22.04 AS build-image
+COPY --from=compile-image /opt/conda /opt/conda
+ENV PATH /opt/conda/bin:$PATH
+
+# Install apt libs
+RUN apt-get update && \
+    apt-get install -y curl git wget && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists*
+
+RUN echo "source activate accelerate" >> ~/.profile
+
+# Activate the virtualenv
+CMD ["/bin/bash"]

docker/accelerate-gpu/Dockerfile

@@ -0,0 +1,45 @@
+# Builds GPU docker image of PyTorch specifically
+# Uses multi-staged approach to reduce size
+# Stage 1
+# Use base conda image to reduce time
+FROM continuumio/miniconda3:latest AS compile-image
+# Specify py version
+ENV PYTHON_VERSION=3.9
+# Install apt libs
+RUN apt-get update && \
+    apt-get install -y curl git wget && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists*
+
+# Create our conda env
+RUN conda create --name accelerate python=${PYTHON_VERSION} ipython jupyter pip
+# We don't install pytorch here yet since CUDA isn't available
+# instead we use the direct torch wheel
+ENV PATH /opt/conda/envs/accelerate/bin:$PATH
+# Activate our bash shell
+RUN chsh -s /bin/bash
+SHELL ["/bin/bash", "-c"]
+# Activate the conda env, install mpy4pi, and install torch + accelerate
+RUN source activate accelerate && conda install -c conda-forge mpi4py
+RUN source activate accelerate && \
+    python3 -m pip install --no-cache-dir \
+    git+https://github.com/huggingface/accelerate#egg=accelerate[testing,test_trackers] \
+    --extra-index-url https://download.pytorch.org/whl/cu126
+
+RUN python3 -m pip install --no-cache-dir bitsandbytes
+
+# Stage 2
+FROM nvidia/cuda:12.6.3-cudnn-devel-ubuntu22.04  AS build-image
+COPY --from=compile-image /opt/conda /opt/conda
+ENV PATH /opt/conda/bin:$PATH
+
+# Install apt libs
+RUN apt-get update && \
+    apt-get install -y curl git wget && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists*
+
+RUN echo "source activate accelerate" >> ~/.profile
+
+# Activate the virtualenv
+CMD ["/bin/bash"]

docs/README.md

@@ -0,0 +1,267 @@
+<!---
+Copyright 2023 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+-->
+
+# Generating the documentation
+
+To generate the documentation, you first have to build it. Several packages are necessary to build the doc, 
+you can install them with the following command, at the root of the code repository:
+
+```bash
+pip install -e ".[docs]"
+```
+
+Then you need to install our special tool that builds the documentation:
+
+```bash
+pip install git+https://github.com/huggingface/doc-builder
+```
+
+---
+**NOTE**
+
+You only need to generate the documentation to inspect it locally (if you're planning changes and want to
+check how they look before committing for instance). You don't have to commit the built documentation.
+
+---
+
+## Building the documentation
+
+Once you have setup the `doc-builder` and additional packages, you can generate the documentation by 
+typing the following command:
+
+```bash
+doc-builder build accelerate docs/source/ --build_dir ~/tmp/test-build
+```
+
+You can adapt the `--build_dir` to set any temporary folder that you prefer. This command will create it and generate
+the MDX files that will be rendered as the documentation on the main website. You can inspect them in your favorite
+Markdown editor.
+
+## Previewing the documentation
+
+To preview the docs, first install the `watchdog` module with:
+
+```bash
+pip install watchdog
+```
+
+Then run the following command:
+
+```bash
+doc-builder preview {package_name} {path_to_docs}
+```
+
+For example:
+
+```bash
+doc-builder preview accelerate docs/source/
+```
+
+The docs will be viewable at [http://localhost:3000](http://localhost:3000). You can also preview the docs once you have opened a PR. You will see a bot add a comment to a link where the documentation with your changes lives.
+
+---
+**NOTE**
+
+The `preview` command only works with existing doc files. When you add a completely new file, you need to update `_toctree.yml` & restart `preview` command (`ctrl-c` to stop it & call `doc-builder preview ...` again).
+
+---
+
+## Adding a new element to the navigation bar
+
+Accepted files are Markdown (.md).
+
+Create a file with its extension and put it in the source directory. You can then link it to the toc-tree by putting
+the filename without the extension in the [`_toctree.yml`](https://github.com/huggingface/accelerate/blob/main/docs/source/_toctree.yml) file.
+
+## Renaming section headers and moving sections
+
+It helps to keep the old links working when renaming the section header and/or moving sections from one document to another. This is because the old links are likely to be used in Issues, Forums, and Social media and it'd make for a much more superior user experience if users reading those months later could still easily navigate to the originally intended information.
+
+Therefore, we simply keep a little map of moved sections at the end of the document where the original section was. The key is to preserve the original anchor.
+
+So if you renamed a section from: "Section A" to "Section B", then you can add at the end of the file:
+
+```
+Sections that were moved:
+
+[ <a href="#section-b">Section A</a><a id="section-a"></a> ]
+```
+and of course, if you moved it to another file, then:
+
+```
+Sections that were moved:
+
+[ <a href="../new-file#section-b">Section A</a><a id="section-a"></a> ]
+```
+
+Use the relative style to link to the new file so that the versioned docs continue to work.
+
+
+## Writing Documentation - Specification
+
+The `huggingface/accelerate` documentation follows the
+[Google documentation](https://sphinxcontrib-napoleon.readthedocs.io/en/latest/example_google.html) style for docstrings,
+although we can write them directly in Markdown.
+
+### Adding a new tutorial
+
+Adding a new tutorial or section is done in two steps:
+
+- Add a new file under `./source`. This file can either be ReStructuredText (.rst) or Markdown (.md).
+- Link that file in `./source/_toctree.yml` on the correct toc-tree.
+
+Make sure to put your new file under the proper section. It's unlikely to go in the first section (*Get Started*), so
+depending on the intended targets (beginners, more advanced users, or researchers) it should go in sections two, three, or
+four.
+
+### Writing source documentation
+
+Values that should be put in `code` should either be surrounded by backticks: \`like so\`. Note that argument names
+and objects like True, None, or any strings should usually be put in `code`.
+
+When mentioning a class, function, or method, it is recommended to use our syntax for internal links so that our tool
+adds a link to its documentation with this syntax: \[\`XXXClass\`\] or \[\`function\`\]. This requires the class or 
+function to be in the main package.
+
+If you want to create a link to some internal class or function, you need to
+provide its path. For instance: \[\`utils.gather\`\]. This will be converted into a link with
+`utils.gather` in the description. To get rid of the path and only keep the name of the object you are
+linking to in the description, add a ~: \[\`~utils.gather\`\] will generate a link with `gather` in the description.
+
+The same works for methods so you can either use \[\`XXXClass.method\`\] or \[~\`XXXClass.method\`\].
+
+#### Defining arguments in a method
+
+Arguments should be defined with the `Args:` (or `Arguments:` or `Parameters:`) prefix, followed by a line return and
+an indentation. The argument should be followed by its type, with its shape if it is a tensor, a colon, and its
+description:
+
+```
+    Args:
+        n_layers (`int`): The number of layers of the model.
+```
+
+If the description is too long to fit in one line (more than 119 characters in total), another indentation is necessary 
+before writing the description after the argument.
+
+Finally, to maintain uniformity if any *one* description is too long to fit on one line, the 
+rest of the parameters should follow suit and have an indention before their description.
+
+Here's an example showcasing everything so far:
+
+```
+    Args:
+        gradient_accumulation_steps (`int`, *optional*, default to 1):
+            The number of steps that should pass before gradients are accumulated. A number > 1 should be combined with `Accelerator.accumulate`.
+        cpu (`bool`, *optional*):
+            Whether or not to force the script to execute on CPU. Will ignore GPU available if set to `True` and force the execution on one process only.
+```
+
+For optional arguments or arguments with defaults we follow the following syntax: imagine we have a function with the
+following signature:
+
+```
+def my_function(x: str = None, a: float = 1):
+```
+
+then its documentation should look like this:
+
+```
+    Args:
+        x (`str`, *optional*):
+            This argument controls ... and has a description longer than 119 chars.
+        a (`float`, *optional*, defaults to 1):
+            This argument is used to ... and has a description longer than 119 chars.
+```
+
+Note that we always omit the "defaults to \`None\`" when None is the default for any argument. Also note that even
+if the first line describing your argument type and its default gets long, you can't break it on several lines. You can
+however write as many lines as you want in the indented description (see the example above with `input_ids`).
+
+#### Writing a multi-line code block
+
+Multi-line code blocks can be useful for displaying examples. They are done between two lines of three backticks as usual in Markdown:
+
+
+````
+```python
+# first line of code
+# second line
+# etc
+```
+````
+
+#### Writing a return block
+
+The return block should be introduced with the `Returns:` prefix, followed by a line return and an indentation.
+The first line should be the type of the return, followed by a line return. No need to indent further for the elements
+building the return.
+
+Here's an example of a single value return:
+
+```
+    Returns:
+        `List[int]`: A list of integers in the range [0, 1] --- 1 for a special token, 0 for a sequence token.
+```
+
+Here's an example of a tuple return, comprising several objects:
+
+```
+    Returns:
+        `tuple(torch.FloatTensor)` comprising various elements depending on the configuration ([`BertConfig`]) and inputs:
+        - ** loss** (*optional*, returned when `masked_lm_labels` is provided) `torch.FloatTensor` of shape `(1,)` --
+          Total loss is the sum of the masked language modeling loss and the next sequence prediction (classification) loss.
+        - **prediction_scores** (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`) --
+          Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
+```
+
+## Styling the docstring
+
+We have an automatic script running with the `make style` comment that will make sure that:
+- the docstrings fully take advantage of the line width
+- all code examples are formatted using black, like the code of the Transformers library
+
+This script may have some weird failures if you made a syntax mistake or if you uncover a bug. Therefore, it's
+recommended to commit your changes before running `make style`, so you can revert the changes done by that script
+easily.
+
+## Writing documentation examples
+
+The syntax for Example docstrings can look as follows:
+
+```
+    Example:
+
+    ```python
+    >>> import time
+    >>> from accelerate import Accelerator
+    >>> accelerator = Accelerator()
+    >>> if accelerator.is_main_process:
+    ...     time.sleep(2)
+    >>> else:
+    ...     print("I'm waiting for the main process to finish its sleep...")
+    >>> accelerator.wait_for_everyone()
+    >>> # Should print on every process at the same time
+    >>> print("Everyone is here")
+    ```
+```
+
+The docstring should give a minimal, clear example of how the respective function 
+is to be used in inference and also include the expected (ideally sensible)
+output.
+Often, readers will try out the example before even going through the function 
+or class definitions. Therefore, it is of utmost importance that the example 
+works as expected.

docs/source/_static/css/code-snippets.css

@@ -1,16 +0,0 @@
-
-.highlight .c1, .highlight .sd{
-    color: #999
-}
-
-.highlight .nn, .highlight .k, .highlight .s1, .highlight .nb, .highlight .bp, .highlight .kc {
-    color: #FB8D68;
-}
-
-.highlight .kn, .highlight .nv, .highlight .s2, .highlight .ow {
-    color: #6670FF;
-}
-
-.highlight .gp {
-    color: #FB8D68;
-}

docs/source/_static/css/huggingface.css

@@ -1,350 +0,0 @@
-/* Our DOM objects */
-
-/* Colab dropdown */
-
-table.center-aligned-table td {
-    text-align: center;
-}
-
-table.center-aligned-table th {
-    text-align: center;
-    vertical-align: middle;
-}
-
-.colab-dropdown {
-    position: relative;
-    display: inline-block;
-}
-  
-.colab-dropdown-content {
-    display: none;
-    position: absolute;
-    background-color: #f9f9f9;
-    min-width: 117px;
-    box-shadow: 0px 8px 16px 0px rgba(0,0,0,0.2);
-    z-index: 1;
-}
-  
-.colab-dropdown-content button {
-    color: #6670FF;
-    background-color: #f9f9f9;
-    font-size: 12px;
-    border: none;
-    min-width: 117px;
-    padding: 5px 5px;
-    text-decoration: none;
-    display: block;
-}
-  
-.colab-dropdown-content button:hover {background-color: #eee;}
-  
-.colab-dropdown:hover .colab-dropdown-content {display: block;}
-
-/* Version control */
-
-.version-button {
-    background-color: #6670FF;
-    color: white;
-    border: none;
-    padding: 5px;
-    font-size: 15px;
-    cursor: pointer;
-}
-
-.version-button:hover, .version-button:focus {
-    background-color: #A6B0FF;
-}
- 
-.version-dropdown {
-    display: none;
-    background-color: #6670FF;
-    min-width: 160px;
-    overflow: auto;
-    font-size: 15px;
-}
-  
-.version-dropdown a {
-    color: white;
-    padding: 3px 4px;
-    text-decoration: none;
-    display: block;
-}
-  
-.version-dropdown a:hover {
-    background-color: #A6B0FF;
-}
-  
-.version-show {
-    display: block;
-}
-
-/* Framework selector */
-
-.framework-selector {
-    display: flex;
-    flex-direction: row;
-    justify-content: flex-end;
-    margin-right: 30px;
-}
-
-.framework-selector > button {
-    background-color: white;
-    color: #6670FF;
-    border: 1px solid #6670FF;
-    padding: 5px;
-}
-
-.framework-selector > button.selected{
-    background-color: #6670FF;
-    color: white;
-    border: 1px solid #6670FF;
-    padding: 5px;
-}
-
-/* Copy button */
-
-a.copybtn {
-    margin: 3px;
-}
-
-/* The literal code blocks */
-.rst-content tt.literal, .rst-content tt.literal, .rst-content code.literal {
-    color: #6670FF;
-}
-
-/* To keep the logo centered */
-.wy-side-scroll {
-    width: auto;
-    font-size: 20px;
-}
-
-/* The div that holds the Hugging Face logo */
-.HuggingFaceDiv {
-    width: 100%
-}
-
-/* The research field on top of the toc tree */
-.wy-side-nav-search{
-    padding-top: 0;
-    background-color: #6670FF;
-}
-
-/* The toc tree */
-.wy-nav-side{
-    background-color: #6670FF;
-}
-
-/* The section headers in the toc tree */
-.wy-menu-vertical p.caption{
-    background-color: #4d59ff;
-    line-height: 40px;
-}
-
-/* The selected items in the toc tree */
-.wy-menu-vertical li.current{
-    background-color: #A6B0FF;
-}
-
-/* When a list item that does belong to the selected block from the toc tree is hovered */
-.wy-menu-vertical li.current a:hover{
-    background-color: #B6C0FF;
-}
-
-/* When a list item that does NOT belong to the selected block from the toc tree is hovered. */
-.wy-menu-vertical li a:hover{
-    background-color: #A7AFFB;
-}
-
-/* The text items on the toc tree */
-.wy-menu-vertical a {
-    color: #FFFFDD;
-    font-family: Calibre-Light, sans-serif;
-}
-.wy-menu-vertical header, .wy-menu-vertical p.caption{
-    color: white;
-    font-family: Calibre-Light, sans-serif;
-}
-
-/* The color inside the selected toc tree block */
-.wy-menu-vertical li.toctree-l2 a, .wy-menu-vertical li.toctree-l3 a, .wy-menu-vertical li.toctree-l4 a {
-    color: black;
-}
-
-/* Inside the depth-2 selected toc tree block */
-.wy-menu-vertical li.toctree-l2.current>a {
-    background-color: #B6C0FF
-}
-.wy-menu-vertical li.toctree-l2.current li.toctree-l3>a {
-    background-color: #C6D0FF
-}
-
-/* Inside the depth-3 selected toc tree block */
-.wy-menu-vertical li.toctree-l3.current li.toctree-l4>a{
-    background-color: #D6E0FF
-}
-
-/* Inside code snippets */
-.rst-content dl:not(.docutils) dt{
-    font-size: 15px;
-}
-
-/* Links */
-a {
-    color: #6670FF;
-}
-
-/* Content bars */
-.rst-content dl:not(.docutils) dt {
-    background-color: rgba(251, 141, 104, 0.1);
-    border-right: solid 2px #FB8D68;
-    border-left: solid 2px #FB8D68;
-    color: #FB8D68;
-    font-family: Calibre-Light, sans-serif;
-    border-top: none;
-    font-style: normal !important;
-}
-
-/* Expand button */
-.wy-menu-vertical li.toctree-l2 span.toctree-expand,
-.wy-menu-vertical li.on a span.toctree-expand, .wy-menu-vertical li.current>a span.toctree-expand,
-.wy-menu-vertical li.toctree-l3 span.toctree-expand{
-    color: black;
-}
-
-/* Max window size */
-.wy-nav-content{
-    max-width: 1200px;
-}
-
-/* Mobile header */
-.wy-nav-top{
-    background-color: #6670FF;
-}
-
-
-/* Source spans */
-.rst-content .viewcode-link, .rst-content .viewcode-back{
-    color: #6670FF;
-    font-size: 110%;
-    letter-spacing: 2px;
-    text-transform: uppercase;
-}
-
-/* It would be better for table to be visible without horizontal scrolling */
-.wy-table-responsive table td, .wy-table-responsive table th{
-    white-space: normal;
-}
-
-.footer {
-    margin-top: 20px;
-}
-
-.footer__Social {
-    display: flex;
-    flex-direction: row;
-}
-
-.footer__CustomImage {
-    margin: 2px 5px 0 0;
-}
-
-/* class and method names in doc */
-.rst-content dl:not(.docutils) tt.descname, .rst-content dl:not(.docutils) tt.descclassname, .rst-content dl:not(.docutils) tt.descname, .rst-content dl:not(.docutils) code.descname, .rst-content dl:not(.docutils) tt.descclassname, .rst-content dl:not(.docutils) code.descclassname{
-    font-family: Calibre, sans-serif;
-    font-size: 20px !important;
-}
-
-/* class name in doc*/
-.rst-content dl:not(.docutils) tt.descname, .rst-content dl:not(.docutils) tt.descname, .rst-content dl:not(.docutils) code.descname{
-    margin-right: 10px;
-    font-family: Calibre-Medium, sans-serif;
-}
-
-/* Method and class parameters */
-.sig-param{
-    line-height: 23px;
-}
-
-/* Class introduction "class" string at beginning */
-.rst-content dl:not(.docutils) .property{
-    font-size: 18px;
-    color: black;
-}
-
-
-/* FONTS */
-body{
-    font-family: Calibre, sans-serif;
-    font-size: 16px;
-}
-
-h1 {
-    font-family: Calibre-Thin, sans-serif;
-    font-size: 70px;
-}
-
-h2, .rst-content .toctree-wrapper p.caption, h3, h4, h5, h6, legend{
-    font-family: Calibre-Medium, sans-serif;
-}
-
-@font-face {
-    font-family: Calibre-Medium;
-    src: url(./Calibre-Medium.otf);
-    font-weight:400;
-}
-
-@font-face {
-    font-family: Calibre;
-    src: url(./Calibre-Regular.otf);
-    font-weight:400;
-}
-
-@font-face {
-    font-family: Calibre-Light;
-    src: url(./Calibre-Light.ttf);
-    font-weight:400;
-}
-
-@font-face {
-    font-family: Calibre-Thin;
-    src: url(./Calibre-Thin.otf);
-    font-weight:400;
-}
-
-
-/**
- * Nav Links to other parts of huggingface.co
- */
- div.menu {
-    position: absolute;
-    top: 0;
-    right: 0;
-    padding-top: 20px;
-    padding-right: 20px;
-    z-index: 1000;
-}
-div.menu a {
-    font-size: 14px;
-    letter-spacing: 0.3px;
-    text-transform: uppercase;
-    color: white;
-    -webkit-font-smoothing: antialiased;
-    background: linear-gradient(0deg, #6671ffb8, #9a66ffb8 50%);
-    padding: 10px 16px 6px 16px;
-    border-radius: 3px;
-    margin-left: 12px;
-    position: relative;
-}
-div.menu a:active {
-    top: 1px;
-}
-@media (min-width: 768px) and (max-width: 1750px) {
-    .wy-breadcrumbs {
-        margin-top: 32px;
-    }
-}
-@media (max-width: 768px) {
-    div.menu {
-        display: none;
-    }
-}

docs/source/_toctree.yml

@@ -0,0 +1,131 @@
+- sections:
+  - local: index
+    title: 🤗 Accelerate
+  - local: basic_tutorials/install
+    title: Installation
+  - local: quicktour
+    title: Quicktour
+  title: Getting started
+- sections:
+  - local: basic_tutorials/overview
+    title: Overview
+  - local: basic_tutorials/migration
+    title: Add Accelerate to your code
+  - local: basic_tutorials/execution
+    title: Execution process
+  - local: basic_tutorials/tpu
+    title: TPU training
+  - local: basic_tutorials/launch
+    title: Launching Accelerate scripts
+  - local: basic_tutorials/notebook
+    title: Launching distributed training from Jupyter Notebooks
+  title: Tutorials
+- sections:
+  - isExpanded: true
+    sections:
+    - local: usage_guides/explore
+      title: Start Here!
+    - local: usage_guides/model_size_estimator
+      title: Model memory estimator
+    - local: usage_guides/quantization
+      title: Model quantization
+    - local: usage_guides/tracking
+      title: Experiment trackers
+    - local: usage_guides/profiler
+      title: Profiler
+    - local: usage_guides/checkpoint
+      title: Checkpointing
+    - local: basic_tutorials/troubleshooting
+      title: Troubleshoot
+    - local: usage_guides/training_zoo
+      title: Example Zoo
+    title: Accelerate
+  - isExpanded: true
+    sections:
+    - local: usage_guides/gradient_accumulation
+      title: Gradient accumulation
+    - local: usage_guides/local_sgd
+      title: Local SGD
+    - local: usage_guides/low_precision_training
+      title: Low precision (FP8) training
+    - local: usage_guides/deepspeed
+      title: DeepSpeed
+    - local: usage_guides/deepspeed_multiple_model
+      title: Using multiple models with DeepSpeed
+    - local: usage_guides/ddp_comm_hook
+      title: DDP Communication Hooks
+    - local: usage_guides/fsdp
+      title: Fully Sharded Data Parallel
+    - local: usage_guides/megatron_lm
+      title: Megatron-LM
+    - local: usage_guides/sagemaker
+      title: Amazon SageMaker
+    - local: usage_guides/mps
+      title: Apple M1 GPUs
+    - local: usage_guides/intel_cpu
+      title: Intel CPU
+    - local: usage_guides/gaudi
+      title: Intel Gaudi
+    - local: usage_guides/compilation
+      title: Compilation
+    title: Training
+  - isExpanded: true
+    sections:
+    - local: usage_guides/big_modeling
+      title: Big Model Inference
+    - local: usage_guides/distributed_inference
+      title: Distributed inference
+    title: Inference
+  title: How to guides
+- sections:
+  - local: concept_guides/internal_mechanism
+    title: Accelerate's internal mechanism
+  - local: concept_guides/big_model_inference
+    title: Loading big models into memory
+  - local: concept_guides/performance
+    title: Comparing performance across distributed setups
+  - local: concept_guides/deferring_execution
+    title: Executing and deferring jobs
+  - local: concept_guides/gradient_synchronization
+    title: Gradient synchronization
+  - local: concept_guides/fsdp_and_deepspeed
+    title: FSDP vs DeepSpeed
+  - local: concept_guides/fsdp1_vs_fsdp2
+    title: FSDP1 vs FSDP2
+  - local: concept_guides/low_precision_training
+    title: Low precision training methods
+  - local: concept_guides/training_tpu
+    title: Training on TPUs
+  title: Concepts and fundamentals
+- sections:
+  - local: package_reference/accelerator
+    title: Accelerator
+  - local: package_reference/state
+    title: Stateful classes
+  - local: package_reference/cli
+    title: The Command Line
+  - local: package_reference/torch_wrappers
+    title: DataLoaders, Optimizers, Schedulers
+  - local: package_reference/tracking
+    title: Experiment trackers
+  - local: package_reference/launchers
+    title: Launchers
+  - local: package_reference/deepspeed
+    title: DeepSpeed utilities
+  - local: package_reference/logging
+    title: Logging
+  - local: package_reference/big_modeling
+    title: Working with large models
+  - local: package_reference/inference
+    title: Pipeline parallelism
+  - local: package_reference/kwargs
+    title: Kwargs handlers
+  - local: package_reference/fp8
+    title: FP8
+  - local: package_reference/utilities
+    title: Utility functions and classes
+  - local: package_reference/megatron_lm
+    title: Megatron-LM utilities
+  - local: package_reference/fsdp
+    title: Fully Sharded Data Parallel utilities
+  title: "Reference"

docs/source/accelerator.rst

@@ -1,43 +0,0 @@
-.. 
-    Copyright 2021 The HuggingFace Team. All rights reserved.
-
-    Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
-    the License. You may obtain a copy of the License at
-
-        http://www.apache.org/licenses/LICENSE-2.0
-
-    Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
-    an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
-    specific language governing permissions and limitations under the License.
-
-Accelerator
-=======================================================================================================================
-
-The :class:`~accelerate.Accelerator` is the main class provided by 🤗 Accelerate. It serves at the main entrypoint for
-the API. To quickly adapt your script to work on any kind of setup with 🤗 Accelerate juste:
-
-1. Initialize an :class:`~accelerate.Accelerator` object (that we will call :obj:`accelerator` in the rest of this
-   page) as early as possible in your script.
-2. Pass along your model(s), optimizer(s), dataloader(s) to the :meth:`~accelerate.Accelerator.prepare` method.
-3. (Optional but best practice) Remove all the :obj:`.cuda()` or :obj:`.to(device)` in your code and let the
-   :obj:`accelerator` handle device placement for you.
-4. Replace the :obj:`loss.backward()` in your code by :obj:`accelerator.backward(loss)`.
-5. (Optional, when using distributed evaluation) Gather your predictions and labelsbefore storing them or using them
-   for metric computation using :meth:`~accelerate.Accelerator.gather`.
-
-This is all what is needed in most cases. For more advanced case or a nicer experience here are the functions you
-should search for and replace by the corresponding methods of your :obj:`accelerator`:
-
-- :obj:`print` statements should be replaced by :meth:`~accelerate.Accelerator.print` to be only printed once per
-  process.
-- Use :meth:`~accelerate.Accelerator.is_local_main_process` for statements that should be executed once per server.
-- Use :meth:`~accelerate.Accelerator.is_main_process` for statements that should be executed once only.
-- Use :meth:`~accelerate.Accelerator.wait_for_everyone` to make sure all processes join that point before continuing
-  (useful before a model save for instance).
-- Use :meth:`~accelerate.Accelerator.unwrap_model` to unwrap your model before saving it.
-- Use :meth:`~accelerate.Accelerator.save` instead of :obj:`torch.save`.
-- Use :meth:`~accelerate.Accelerator.clip_grad_norm_` instead of :obj:`torch.nn.utils.clip_grad_norm_` and
-  :meth:`~accelerate.Accelerator.clip_grad_value_` instead of :obj:`torch.nn.utils.clip_grad_value_`.
-
-.. autoclass:: accelerate.Accelerator
-    :members:

docs/source/basic_tutorials/execution.md

@@ -0,0 +1,128 @@
+<!--Copyright 2024 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
+-->
+
+# Execution process
+
+When working with distributed training systems, it is important to manage how and when processes are executed across GPUs. Some processes are completed faster than others, and some processes shouldn't begin if others haven't finished yet. Accelerate provides tools for orchestrating when processes are executed to ensure everything remains synchronized across all devices.
+
+This tutorial will teach you how to execute a process on only one machine and how to delay execution until all processes have reached a certain point.
+
+## Execute on one process
+
+Certain code only needs to be run once on a given machine, such as printing a log statement or only displaying one progress bar on the local main process.
+
+<hfoptions id="local-execution">
+<hfoption id="statements">
+
+You should use `accelerator.is_local_main_process` to indicate code that should only be executed once.
+
+```py
+from tqdm.auto import tqdm
+
+progress_bar = tqdm(range(args.max_train_steps), disable=not accelerator.is_local_main_process)
+```
+
+You could also wrap a statement with `accelerator.is_local_main_process`.
+
+> [!TIP]
+> For standalone `print` statements that aren't wrapped in `accelerator.is_local_main_process`, replace `print` with Accelerate's [`~Accelerator.print`] method to only print once per process.
+
+```py
+if accelerator.is_local_main_process:
+    print("Accelerate is the best")
+```
+
+</hfoption>
+<hfoption id="function">
+
+For a function that should only be executed once, use [`~Accelerator.on_local_main_process`].
+
+```py
+@accelerator.on_local_main_process
+def do_my_thing():
+    "Something done once per server"
+    do_thing_once_per_server()
+```
+
+</hfoption>
+</hfoptions>
+
+You could also direct Accelerate to execute code once across *all processes* regardless of the number of machines. This is useful if you're uploading a final model to the Hub.
+
+<hfoptions id="main-execution">
+<hfoption id="statement">
+
+You should use `accelerator.is_main_process` to indicate code that should only be executed once across all processes.
+
+```py
+if accelerator.is_main_process:
+    repo.push_to_hub()
+```
+
+</hfoption>
+<hfoption id="function">
+
+For a function that should only be executed once across all processes, use [`~Accelerator.on_main_process`].
+
+```py
+@accelerator.on_main_process
+def do_my_thing():
+    "Something done once per server"
+    do_thing_once()
+```
+
+</hfoption>
+</hfoptions>
+
+## Execute on a specific process
+
+Accelerate can also help you execute functions that should only be executed on a specific process or a local process index.
+
+<hfoptions id="specific-execution">
+<hfoption id="specific process">
+
+Use the [`~Accelerator.on_process`] method and specify the process index to execute a function on.
+
+```py
+@accelerator.on_process(process_index=0)
+def do_my_thing():
+    "Something done on process index 0"
+    do_thing_on_index_zero()
+```
+
+</hfoption>
+<hfoption id="local process">
+
+Use the [`~Accelerator.on_local_process`] method and specify the local process index to execute a function on.
+
+```py
+@accelerator.on_local_process(local_process_idx=0)
+def do_my_thing():
+    "Something done on process index 0 on each server"
+    do_thing_on_index_zero_on_each_server()
+```
+
+</hfoption>
+</hfoptions>
+
+## Defer execution
+
+When you run your script on several GPUs at the same time, some code may be executed faster than others. You might need to wait for all processes to reach a certain point before executing the next set of instructions. For instance, you shouldn’t save a model before making sure every process is done with training.
+
+To do this, add [`~Accelerator.wait_for_everyone`] in your code. This blocks all processes that have finished first from continuing until all remaining processes have reached the same point (this has no effect if you're running on a single GPU or CPU).
+
+```py
+accelerator.wait_for_everyone()
+```

docs/source/basic_tutorials/install.md

@@ -0,0 +1,114 @@
+<!--Copyright 2022 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
+-->
+
+# Installation
+
+Before you start, you will need to setup your environment, install the appropriate packages, and configure Accelerate. Accelerate is tested on **Python 3.8+**.
+
+Accelerate is available on pypi and conda, as well as on GitHub. Details to install from each are below:
+
+## pip
+
+To install Accelerate from pypi, perform:
+
+```bash
+pip install accelerate
+```
+
+## conda
+
+Accelerate can also be installed with conda with:
+
+```bash
+conda install -c conda-forge accelerate
+```
+
+## Source
+
+New features are added every day that haven't been released yet. To try them out yourself, install
+from the GitHub repository:
+
+```bash
+pip install git+https://github.com/huggingface/accelerate
+```
+
+If you're working on contributing to the library or wish to play with the source code and see live 
+results as you run the code, an editable version can be installed from a locally-cloned version of the 
+repository:
+
+```bash
+git clone https://github.com/huggingface/accelerate
+cd accelerate
+pip install -e .
+```
+
+## Configuration
+
+After installing, you need to configure Accelerate for how the current system is setup for training. 
+To do so run the following and answer the questions prompted to you:
+
+```bash
+accelerate config
+```
+
+To write a barebones configuration that doesn't include options such as DeepSpeed configuration or running on TPUs, you can quickly run:
+
+```bash
+python -c "from accelerate.utils import write_basic_config; write_basic_config(mixed_precision='fp16')"
+```
+
+Accelerate will automatically utilize the maximum number of GPUs available and set the mixed precision mode.
+
+To check that your configuration looks fine, run:
+
+```bash
+accelerate env
+```
+
+An example output is shown below, which describes two GPUs on a single machine with no mixed precision being used:
+
+
+```bash
+- `Accelerate` version: 1.2.0.dev0
+- Platform: Linux-6.8.0-47-generic-x86_64-with-glibc2.35
+- `accelerate` bash location: /home/zach/miniconda3/envs/accelerate/bin/accelerate
+- Python version: 3.10.13
+- Numpy version: 1.26.4
+- PyTorch version (GPU?): 2.5.1+cu124 (True)
+- PyTorch XPU available: False
+- PyTorch NPU available: False
+- PyTorch MLU available: False
+- PyTorch MUSA available: False
+- System RAM: 187.91 GB
+- GPU type: NVIDIA GeForce RTX 4090
+- `Accelerate` default config:
+        - compute_environment: LOCAL_MACHINE
+        - distributed_type: MULTI_GPU
+        - mixed_precision: no
+        - use_cpu: False
+        - debug: False
+        - num_processes: 2
+        - machine_rank: 0
+        - num_machines: 1
+        - gpu_ids: all
+        - rdzv_backend: static
+        - same_network: True
+        - main_training_function: main
+        - enable_cpu_affinity: False
+        - downcast_bf16: no
+        - tpu_use_cluster: False
+        - tpu_use_sudo: False
+        - tpu_env: []
+```

docs/source/basic_tutorials/launch.md

@@ -0,0 +1,235 @@
+<!--Copyright 2022 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
+-->
+
+# Launching Accelerate scripts
+
+In the previous tutorial, you were introduced to how to modify your current training script to use Accelerate.
+The final version of that code is shown below:
+
+```python
+from accelerate import Accelerator
+
+accelerator = Accelerator()
+
+model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+    model, optimizer, training_dataloader, scheduler
+)
+
+for batch in training_dataloader:
+    optimizer.zero_grad()
+    inputs, targets = batch
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    accelerator.backward(loss)
+    optimizer.step()
+    scheduler.step()
+```
+
+But how do you run this code and have it utilize the special hardware available to it?
+
+First, you should rewrite the above code into a function, and make it callable as a script. For example:
+
+```diff
+  from accelerate import Accelerator
+  
++ def main():
+      accelerator = Accelerator()
+
+      model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+          model, optimizer, training_dataloader, scheduler
+      )
+
+      for batch in training_dataloader:
+          optimizer.zero_grad()
+          inputs, targets = batch
+          outputs = model(inputs)
+          loss = loss_function(outputs, targets)
+          accelerator.backward(loss)
+          optimizer.step()
+          scheduler.step()
+
++ if __name__ == "__main__":
++     main()
+```
+
+Next, you need to launch it with `accelerate launch`. 
+
+<Tip warning={true}>
+
+  It's recommended you run `accelerate config` before using `accelerate launch` to configure your environment to your liking. 
+  Otherwise Accelerate will use very basic defaults depending on your system setup.
+
+</Tip>
+
+
+## Using accelerate launch
+
+Accelerate has a special CLI command to help you launch your code in your system through `accelerate launch`.
+This command wraps around all of the different commands needed to launch your script on various platforms, without you having to remember what each of them is.
+
+<Tip>
+
+  If you are familiar with launching scripts in PyTorch yourself such as with `torchrun`, you can still do this. It is not required to use `accelerate launch`.
+
+</Tip>
+
+You can launch your script quickly by using:
+
+```bash
+accelerate launch {script_name.py} --arg1 --arg2 ...
+```
+
+Just put `accelerate launch` at the start of your command, and pass in additional arguments and parameters to your script afterward like normal!
+
+Since this runs the various torch spawn methods, all of the expected environment variables can be modified here as well.
+For example, here is how to use `accelerate launch` with a single GPU:
+
+```bash
+# for cuda device:
+CUDA_VISIBLE_DEVICES="0" accelerate launch {script_name.py} --arg1 --arg2 ...
+# for xpu device:
+ZE_AFFINITY_MASK="0" accelerate launch {script_name.py} --arg1 --arg2 ...
+```
+
+You can also use `accelerate launch` without performing `accelerate config` first, but you may need to manually pass in the right configuration parameters.
+In this case, Accelerate will make some hyperparameter decisions for you, e.g., if GPUs are available, it will use all of them by default without the mixed precision.
+Here is how you would use all GPUs and train with mixed precision disabled:
+
+```bash
+accelerate launch --multi_gpu {script_name.py} {--arg1} {--arg2} ...
+```
+
+Or by specifying a number of GPUs to use:
+
+```bash
+accelerate launch --num_processes=2 {script_name.py} {--arg1} {--arg2} ...
+```
+
+To get more specific you should pass in the needed parameters yourself. For instance, here is how you 
+would also launch that same script on two GPUs using mixed precision while avoiding all of the warnings: 
+
+```bash
+accelerate launch --multi_gpu --mixed_precision=fp16 --num_processes=2 {script_name.py} {--arg1} {--arg2} ...
+```
+
+For a complete list of parameters you can pass in, run:
+
+```bash
+accelerate launch -h
+```
+
+<Tip>
+
+  Even if you are not using Accelerate in your code, you can still use the launcher for starting your scripts!
+
+</Tip>
+
+For a visualization of this difference, that earlier `accelerate launch` on multi-gpu would look something like so with `torchrun`:
+
+```bash
+MIXED_PRECISION="fp16" torchrun --nproc_per_node=2 --nnodes=1 {script_name.py} {--arg1} {--arg2} ...
+```
+
+You can also launch your script utilizing the launch CLI as a python module itself, enabling the ability to pass in other python-specific
+launching behaviors. To do so, use `accelerate.commands.launch` instead of `accelerate launch`:
+
+```bash
+python -m accelerate.commands.launch --num_processes=2 {script_name.py} {--arg1} {--arg2}
+```
+
+If you want to execute the script with any other python flags, you can pass them in as well similar to `-m`, such as 
+the below example enabling unbuffered stdout and stderr:
+
+```bash
+python -u -m accelerate.commands.launch --num_processes=2 {script_name.py} {--arg1} {--arg2}
+```
+
+<Tip>
+
+  You can run your code on CPU as well! This is helpful for debugging and testing purposes on toy models and datasets. 
+
+```bash
+accelerate launch --cpu {script_name.py} {--arg1} {--arg2}
+```  
+
+</Tip>
+
+## Why you should always use `accelerate config`
+
+Why is it useful to the point you should **always** run `accelerate config`? 
+
+Remember that earlier call to `accelerate launch` as well as `torchrun`?
+Post configuration, to run that script with the needed parts you just need to use `accelerate launch` outright, without passing anything else in:
+
+```bash
+accelerate launch {script_name.py} {--arg1} {--arg2} ...
+```
+
+
+## Custom Configurations
+
+As briefly mentioned earlier, `accelerate launch` should be mostly used through combining set configurations 
+made with the `accelerate config` command. These configs are saved to a `default_config.yaml` file in your cache folder for Accelerate. 
+This cache folder is located at (with decreasing order of priority):
+
+- The content of your environment variable `HF_HOME` suffixed with `accelerate`.
+- If it does not exist, the content of your environment variable `XDG_CACHE_HOME` suffixed with
+  `huggingface/accelerate`.
+- If this does not exist either, the folder `~/.cache/huggingface/accelerate`.
+
+To have multiple configurations, the flag `--config_file` can be passed to the `accelerate launch` command paired 
+with the location of the custom yaml. 
+
+An example yaml may look something like the following for two GPUs on a single machine using `fp16` for mixed precision:
+```yaml
+compute_environment: LOCAL_MACHINE
+deepspeed_config: {}
+distributed_type: MULTI_GPU
+fsdp_config: {}
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+mixed_precision: fp16
+num_machines: 1
+num_processes: 2
+use_cpu: false
+```
+
+Launching a script from the location of that custom yaml file looks like the following:
+```bash
+accelerate launch --config_file {path/to/config/my_config_file.yaml} {script_name.py} {--arg1} {--arg2} ...
+```
+
+## Multi-node training
+Multi-node training with Accelerate is similar to [multi-node training with torchrun](https://pytorch.org/tutorials/intermediate/ddp_series_multinode.html). The simplest way to launch a multi-node training run is to do the following:
+
+- Copy your codebase and data to all nodes. (or place them on a shared filesystem)
+- Setup your python packages on all nodes.
+- Run `accelerate config` on the main single node first. After specifying the number of nodes, you will be asked to specify the rank of each node (this will be 0 for the main/master node), along with the IP address and port for the main process. This is required for the worker nodes to communicate with the main process. Afterwards, you can copy or send this config file across all of your nodes, changing the `machine_rank` to 1, 2,3, etc. to avoid having to run the command (or just follow their directions directly for launching with `torchrun` as well)
+
+Once you have done this, you can start your multi-node training run by running `accelerate launch` (or `torchrun`) on all nodes.
+
+<Tip>
+    It is required that the command be ran on all nodes for everything to start, not just running it from the main node. You can use something like SLURM or a different process executor to wrap around this requirement and call everything from a single command.
+</Tip>
+
+<Tip>
+
+ It is recommended to use the intranet IP of your main node over the public IP for better latency. This is the `192.168.x.x` or the `172.x.x.x` address you see when you run `hostname -I` on the main node.
+
+</Tip>
+
+To get a better idea about multi-node training, check out our example for [multi-node training with FSDP](https://huggingface.co/blog/ram-efficient-pytorch-fsdp).

docs/source/basic_tutorials/migration.md

@@ -0,0 +1,224 @@
+<!--Copyright 2022 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
+-->
+
+# Add Accelerate to your code
+
+Each distributed training framework has their own way of doing things which can require writing a lot of custom code to adapt it to your PyTorch training code and training environment. Accelerate offers a friendly way to interface with these distributed training frameworks without having to learn the specific details of each one. Accelerate takes care of those details for you, so you can focus on the training code and scale it to any distributed training environment.
+
+In this tutorial, you'll learn how to adapt your existing PyTorch code with Accelerate and get you on your way toward training on distributed systems with ease! You'll start with a basic PyTorch training loop (it assumes all the training objects like `model` and `optimizer` have been setup already) and progressively integrate Accelerate into it.
+
+```python
+device = "cuda"
+model.to(device)
+
+for batch in training_dataloader:
+    optimizer.zero_grad()
+    inputs, targets = batch
+    inputs = inputs.to(device)
+    targets = targets.to(device)
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    loss.backward()
+    optimizer.step()
+    scheduler.step()
+```
+
+## Accelerator
+
+The [`Accelerator`] is the main class for adapting your code to work with Accelerate. It knows about the distributed setup you're using such as the number of different processes and your hardware type. This class also provides access to many of the necessary methods for enabling your PyTorch code to work in any distributed training environment and for managing and executing processes across devices.
+
+That's why you should always start by importing and creating an [`Accelerator`] instance in your script.
+
+```python
+from accelerate import Accelerator
+
+accelerator = Accelerator()
+```
+
+The [`Accelerator`] also knows which device to move your PyTorch objects to, so it is recommended to let Accelerate handle this for you.
+
+```diff
+- device = "cuda"
++ device = accelerator.device
+  model.to(device)
+```
+
+## Prepare PyTorch objects
+
+Next, you need to prepare your PyTorch objects (model, optimizer, scheduler, etc.) for distributed training. The [`~Accelerator.prepare`] method takes care of placing your model in the appropriate container (like single GPU or multi-GPU) for your training setup, adapting the optimizer and scheduler to use Accelerate's [`~optimizer.AcceleratedOptimizer`] and [`~scheduler.AcceleratedScheduler`], and creating a new dataloader that can be sharded across processes.
+
+> [!TIP]
+> Accelerate only prepares objects that inherit from their respective PyTorch classes such as `torch.optim.Optimizer`.
+
+The PyTorch objects are returned in the same order they're sent.
+
+```py
+model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+    model, optimizer, training_dataloader, scheduler
+)
+```
+
+## Training loop
+
+Finally, remove the `to(device)` calls to the inputs and targets in the training loop because Accelerate's DataLoader classes automatically places them on the right device. You should also replace the usual `backward()` pass with Accelerate's [`~Accelerator.backward`] method which scales the gradients for you and uses the appropriate `backward()` method depending on your distributed setup (for example, DeepSpeed or Megatron).
+
+```diff
+-   inputs = inputs.to(device)
+-   targets = targets.to(device)
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+-   loss.backward()
++   accelerator.backward(loss)
+```
+
+Put everything together and your new Accelerate training loop should now look like this!
+
+```python
+from accelerate import Accelerator
+accelerator = Accelerator()
+
+device = accelerator.device
+model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+    model, optimizer, training_dataloader, scheduler
+)
+
+for batch in training_dataloader:
+    optimizer.zero_grad()
+    inputs, targets = batch
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    accelerator.backward(loss)
+    optimizer.step()
+    scheduler.step()
+```
+
+## Training features
+
+Accelerate offers additional features - like gradient accumulation, gradient clipping, mixed precision training and more - you can add to your script to improve your training run. Let's explore these three features.
+
+### Gradient accumulation
+
+Gradient accumulation enables you to train on larger batch sizes by accumulating the gradients over multiple batches before updating the weights. This can be useful for getting around memory limitations. To enable this feature in Accelerate, specify the `gradient_accumulation_steps` parameter in the [`Accelerator`] class and add the [`~Accelerator.accumulate`] context manager to your script.
+
+```diff
++ accelerator = Accelerator(gradient_accumulation_steps=2)
+  model, optimizer, training_dataloader = accelerator.prepare(model, optimizer, training_dataloader)
+
+  for input, label in training_dataloader:
++     with accelerator.accumulate(model):
+          predictions = model(input)
+          loss = loss_function(predictions, label)
+          accelerator.backward(loss)
+          optimizer.step()
+          scheduler.step()
+          optimizer.zero_grad()
+```
+
+### Gradient clipping
+
+Gradient clipping is a technique to prevent "exploding gradients", and Accelerate offers:
+
+* [`~Accelerator.clip_grad_value_`] to clip gradients to a minimum and maximum value
+* [`~Accelerator.clip_grad_norm_`] for normalizing gradients to a certain value
+
+### Mixed precision
+
+Mixed precision accelerates training by using a lower precision data type like fp16 (half-precision) to calculate the gradients. For the best performance with Accelerate, the loss should be computed inside your model (like in Transformers models) because computations outside of the model are computed in full precision.
+
+Set the mixed precision type to use in the [`Accelerator`], and then use the [`~Accelerator.autocast`] context manager to automatically cast the values to the specified data type.
+
+> [!WARNING]
+> Accelerate enables automatic mixed precision, so [`~Accelerator.autocast`] is only needed if there are other mixed precision operations besides those performed on loss by [`~Accelerator.backward`] which already handles the scaling.
+
+```diff
++ accelerator = Accelerator(mixed_precision="fp16")
++ with accelerator.autocast():
+      loss = complex_loss_function(outputs, target)
+```
+
+## Save and load
+
+Accelerate can also save and load a *model* once training is complete or you can also save the model and optimizer *state* which could be useful for resuming training.
+
+### Model
+
+Once all processes are complete, unwrap the model with the [`~Accelerator.unwrap_model`] method before saving it because the [`~Accelerator.prepare`] method wrapped your model into the proper interface for distributed training. If you don't unwrap the model, saving the model state dictionary also saves any potential extra layers from the larger model and you won't be able to load the weights back into your base model.
+
+You should use the [`~Accelerator.save_model`] method to unwrap and save the model state dictionary. This method can also save a model into sharded checkpoints or into the [safetensors](https://hf.co/docs/safetensors/index) format.
+
+<hfoptions id="save">
+<hfoption id="single checkpoint">
+
+```py
+accelerator.wait_for_everyone()
+accelerator.save_model(model, save_directory)
+```
+
+<Tip>
+
+For models from the [Transformers](https://hf.co/docs/transformers/index) library, save the model with the [`~transformers.PreTrainedModel.save_pretrained`] method so that it can be reloaded with the [`~transformers.PreTrainedModel.from_pretrained`] method.
+
+```py
+from transformers import AutoModel
+
+unwrapped_model = accelerator.unwrap_model(model)
+unwrapped_model.save_pretrained(
+    "path/to/my_model_directory",
+    is_main_process=accelerator.is_main_process,
+    save_function=accelerator.save,
+)
+
+model = AutoModel.from_pretrained("path/to/my_model_directory")
+```
+
+</Tip>
+
+To load your weights, use the [`~Accelerator.unwrap_model`] method to unwrap the model first before loading the weights. All model parameters are references to tensors, so this loads your weights inside `model`.
+
+```py
+unwrapped_model = accelerator.unwrap_model(model)
+path_to_checkpoint = os.path.join(save_directory,"pytorch_model.bin")
+unwrapped_model.load_state_dict(torch.load(path_to_checkpoint))
+```
+
+</hfoption>
+<hfoption id="sharded checkpoint">
+
+Set `safe_serialization=True` to save the model in the safetensor format.
+
+```py
+accelerator.wait_for_everyone()
+accelerator.save_model(model, save_directory, max_shard_size="1GB", safe_serialization=True)
+```
+
+To load a sharded checkpoint or a safetensor formatted checkpoint, use the [`~accelerate.load_checkpoint_in_model`] method. This method allows you to load a checkpoint onto a specific device.
+
+```py
+load_checkpoint_in_model(unwrapped_model, save_directory, device_map={"":device})
+```
+
+</hfoption>
+</hfoptions>
+
+### State
+
+During training, you may want to save the current state of the model, optimizer, random generators, and potentially learning rate schedulers so they can be restored in the *same script*. You should add the [`~Accelerator.save_state`] and [`~Accelerator.load_state`] methods to your script to save and load states.
+
+To further customize where and how states are saved through [`~Accelerator.save_state`], use the [`~utils.ProjectConfiguration`] class. For example, if `automatic_checkpoint_naming` is enabled, each saved checkpoint is stored at `Accelerator.project_dir/checkpoints/checkpoint_{checkpoint_number}`.
+
+Any other stateful items to be stored should be registered with the [`~Accelerator.register_for_checkpointing`] method so they can be saved and loaded. Every object passed to this method to be stored must have a `load_state_dict` and `state_dict` function.
+
+> [!TIP]
+> If you have [`torchdata>=0.8.0`](https://github.com/pytorch/data/tree/main) installed, you can additionally pass `use_stateful_dataloader=True` into your [`~utils.DataLoaderConfiguration`]. This extends Accelerate's DataLoader classes with a `load_state_dict` and `state_dict` function, and makes it so `Accelerator.save_state` and `Accelerator.load_state` also track how far into the training dataset it has read when persisting the model.

docs/source/basic_tutorials/notebook.md

@@ -0,0 +1,476 @@
+<!--Copyright 2022 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
+-->
+
+# Launching distributed training from Jupyter Notebooks
+
+This tutorial teaches you how to fine tune a computer vision model with 🤗 Accelerate from a Jupyter Notebook on a distributed system.
+You will also learn how to setup a few requirements needed for ensuring your environment is configured properly, your data has been prepared properly, and finally how to launch training.
+
+<Tip>
+
+    This tutorial is also available as a Jupyter Notebook [here](https://github.com/huggingface/notebooks/blob/main/examples/accelerate_examples/simple_cv_example.ipynb)
+
+</Tip>
+
+## Configuring the Environment
+
+Before any training can be performed, an Accelerate config file must exist in the system. Usually this can be done by running the following in a terminal and answering the prompts:
+
+```bash
+accelerate config
+```
+
+However, if general defaults are fine and you are *not* running on a TPU, Accelerate has a utility to quickly write your GPU configuration into a config file via [`utils.write_basic_config`].
+
+The following code will restart Jupyter after writing the configuration, as CUDA code was called to perform this. 
+
+<Tip warning={true}>
+
+    CUDA can't be initialized more than once on a multi-GPU system. It's fine to debug in the notebook and have calls to CUDA, but in order to finally train a full cleanup and restart will need to be performed.
+    
+</Tip>
+
+```python
+import os
+from accelerate.utils import write_basic_config
+
+write_basic_config()  # Write a config file
+os._exit(00)  # Restart the notebook
+```
+
+## Preparing the Dataset and Model
+
+Next you should prepare your dataset. As mentioned earlier, great care should be taken when preparing the `DataLoaders` and model to make sure that **nothing** is put on *any* GPU. 
+
+If you do, it is recommended to put that specific code into a function and call that from within the notebook launcher interface, which will be shown later. 
+
+Make sure the dataset is downloaded based on the directions [here](https://github.com/huggingface/accelerate/tree/main/examples#simple-vision-example)
+
+```python
+import os, re, torch, PIL
+import numpy as np
+
+from torch.optim.lr_scheduler import OneCycleLR
+from torch.utils.data import DataLoader, Dataset
+from torchvision.transforms import Compose, RandomResizedCrop, Resize, ToTensor
+
+from accelerate import Accelerator
+from accelerate.utils import set_seed
+from timm import create_model
+```
+
+First you need to create a function to extract the class name based on a filename:
+
+```python
+import os
+
+data_dir = "../../images"
+fnames = os.listdir(data_dir)
+fname = fnames[0]
+print(fname)
+```
+
+```python out
+beagle_32.jpg
+```
+
+In the case here, the label is `beagle`. Using regex you can extract the label from the filename:
+
+```python
+import re
+
+
+def extract_label(fname):
+    stem = fname.split(os.path.sep)[-1]
+    return re.search(r"^(.*)_\d+\.jpg$", stem).groups()[0]
+```
+
+```python
+extract_label(fname)
+```
+
+And you can see it properly returned the right name for our file:
+
+```python out
+"beagle"
+```
+
+Next a `Dataset` class should be made to handle grabbing the image and the label:
+
+```python
+class PetsDataset(Dataset):
+    def __init__(self, file_names, image_transform=None, label_to_id=None):
+        self.file_names = file_names
+        self.image_transform = image_transform
+        self.label_to_id = label_to_id
+
+    def __len__(self):
+        return len(self.file_names)
+
+    def __getitem__(self, idx):
+        fname = self.file_names[idx]
+        raw_image = PIL.Image.open(fname)
+        image = raw_image.convert("RGB")
+        if self.image_transform is not None:
+            image = self.image_transform(image)
+        label = extract_label(fname)
+        if self.label_to_id is not None:
+            label = self.label_to_id[label]
+        return {"image": image, "label": label}
+```
+
+Now to build the dataset. Outside the training function you can find and declare all the filenames and labels and use them as references inside the 
+launched function:
+
+```python
+fnames = [os.path.join("../../images", fname) for fname in fnames if fname.endswith(".jpg")]
+```
+
+Next gather all the labels:
+
+```python
+all_labels = [extract_label(fname) for fname in fnames]
+id_to_label = list(set(all_labels))
+id_to_label.sort()
+label_to_id = {lbl: i for i, lbl in enumerate(id_to_label)}
+```
+
+Next, you should make a `get_dataloaders` function that will return your built dataloaders for you. As mentioned earlier, if data is automatically 
+sent to the GPU or a TPU device when building your `DataLoaders`, they must be built using this method. 
+
+```python
+def get_dataloaders(batch_size: int = 64):
+    "Builds a set of dataloaders with a batch_size"
+    random_perm = np.random.permutation(len(fnames))
+    cut = int(0.8 * len(fnames))
+    train_split = random_perm[:cut]
+    eval_split = random_perm[cut:]
+
+    # For training a simple RandomResizedCrop will be used
+    train_tfm = Compose([RandomResizedCrop((224, 224), scale=(0.5, 1.0)), ToTensor()])
+    train_dataset = PetsDataset([fnames[i] for i in train_split], image_transform=train_tfm, label_to_id=label_to_id)
+
+    # For evaluation a deterministic Resize will be used
+    eval_tfm = Compose([Resize((224, 224)), ToTensor()])
+    eval_dataset = PetsDataset([fnames[i] for i in eval_split], image_transform=eval_tfm, label_to_id=label_to_id)
+
+    # Instantiate dataloaders
+    train_dataloader = DataLoader(train_dataset, shuffle=True, batch_size=batch_size, num_workers=4)
+    eval_dataloader = DataLoader(eval_dataset, shuffle=False, batch_size=batch_size * 2, num_workers=4)
+    return train_dataloader, eval_dataloader
+```
+
+Finally, you should import the scheduler to be used later:
+
+```python
+from torch.optim.lr_scheduler import CosineAnnealingLR
+```
+
+## Writing the Training Function
+
+Now you can build the training loop. [`notebook_launcher`] works by passing in a function to call that will be ran across the distributed system.
+
+Here is a basic training loop for the animal classification problem:
+
+<Tip>
+
+    The code has been split up to allow for explanations on each section. A full version that can be copy and pasted will be available at the end
+
+</Tip>
+
+
+```python
+def training_loop(mixed_precision="fp16", seed: int = 42, batch_size: int = 64):
+    set_seed(seed)
+    accelerator = Accelerator(mixed_precision=mixed_precision)
+```
+
+First you should set the seed and create an [`Accelerator`] object as early in the training loop as possible.
+
+<Tip warning={true}>
+
+    If training on the TPU, your training loop should take in the model as a parameter and it should be instantiated 
+    outside of the training loop function. See the [TPU best practices](../concept_guides/training_tpu) 
+    to learn why
+
+</Tip>
+
+Next you should build your dataloaders and create your model:
+
+```python
+    train_dataloader, eval_dataloader = get_dataloaders(batch_size)
+    model = create_model("resnet50d", pretrained=True, num_classes=len(label_to_id))
+```
+
+<Tip>
+
+    You build the model here so that the seed also controls the new weight initialization
+
+</Tip>
+
+As you are performing transfer learning in this example, the encoder of the model starts out frozen so the head of the model can be 
+trained only initially:
+
+```python
+    for param in model.parameters():
+        param.requires_grad = False
+    for param in model.get_classifier().parameters():
+        param.requires_grad = True
+```
+
+Normalizing the batches of images will make training a little faster:
+
+```python
+    mean = torch.tensor(model.default_cfg["mean"])[None, :, None, None]
+    std = torch.tensor(model.default_cfg["std"])[None, :, None, None]
+```
+
+To make these constants available on the active device, you should set it to the Accelerator's device:
+
+```python
+    mean = mean.to(accelerator.device)
+    std = std.to(accelerator.device)
+```
+
+Next instantiate the rest of the PyTorch classes used for training:
+
+```python
+    optimizer = torch.optim.Adam(params=model.parameters(), lr=3e-2 / 25)
+    lr_scheduler = OneCycleLR(optimizer=optimizer, max_lr=3e-2, epochs=5, steps_per_epoch=len(train_dataloader))
+```
+
+Before passing everything to [`~Accelerator.prepare`].
+
+<Tip>
+
+    There is no specific order to remember, you just need to unpack the objects in the same order you gave them to the prepare method.
+
+</Tip>
+
+```python
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+    )
+```
+
+Now train the model:
+
+```python
+    for epoch in range(5):
+        model.train()
+        for batch in train_dataloader:
+            inputs = (batch["image"] - mean) / std
+            outputs = model(inputs)
+            loss = torch.nn.functional.cross_entropy(outputs, batch["label"])
+            accelerator.backward(loss)
+            optimizer.step()
+            lr_scheduler.step()
+            optimizer.zero_grad()
+```
+
+The evaluation loop will look slightly different compared to the training loop. The number of elements passed as well as the overall 
+total accuracy of each batch will be added to two constants:
+
+```python
+        model.eval()
+        accurate = 0
+        num_elems = 0
+```
+
+Next you have the rest of your standard PyTorch loop:
+
+```python
+        for batch in eval_dataloader:
+            inputs = (batch["image"] - mean) / std
+            with torch.no_grad():
+                outputs = model(inputs)
+            predictions = outputs.argmax(dim=-1)
+```
+
+Before finally the last major difference. 
+
+When performing distributed evaluation, the predictions and labels need to be passed through 
+[`~Accelerator.gather`] so that all of the data is available on the current device and a properly calculated metric can be achieved:
+
+```python
+            accurate_preds = accelerator.gather(predictions) == accelerator.gather(batch["label"])
+            num_elems += accurate_preds.shape[0]
+            accurate += accurate_preds.long().sum()
+```
+
+Now you just need to calculate the actual metric for this problem, and you can print it on the main process using [`~Accelerator.print`]:
+
+```python
+        eval_metric = accurate.item() / num_elems
+        accelerator.print(f"epoch {epoch}: {100 * eval_metric:.2f}")
+```
+
+A full version of this training loop is available below:
+
+```python
+def training_loop(mixed_precision="fp16", seed: int = 42, batch_size: int = 64):
+    set_seed(seed)
+    # Initialize accelerator
+    accelerator = Accelerator(mixed_precision=mixed_precision)
+    # Build dataloaders
+    train_dataloader, eval_dataloader = get_dataloaders(batch_size)
+
+    # Instantiate the model (you build the model here so that the seed also controls new weight initializations)
+    model = create_model("resnet50d", pretrained=True, num_classes=len(label_to_id))
+
+    # Freeze the base model
+    for param in model.parameters():
+        param.requires_grad = False
+    for param in model.get_classifier().parameters():
+        param.requires_grad = True
+
+    # You can normalize the batches of images to be a bit faster
+    mean = torch.tensor(model.default_cfg["mean"])[None, :, None, None]
+    std = torch.tensor(model.default_cfg["std"])[None, :, None, None]
+
+    # To make these constants available on the active device, set it to the accelerator device
+    mean = mean.to(accelerator.device)
+    std = std.to(accelerator.device)
+
+    # Instantiate the optimizer
+    optimizer = torch.optim.Adam(params=model.parameters(), lr=3e-2 / 25)
+
+    # Instantiate the learning rate scheduler
+    lr_scheduler = OneCycleLR(optimizer=optimizer, max_lr=3e-2, epochs=5, steps_per_epoch=len(train_dataloader))
+
+    # Prepare everything
+    # There is no specific order to remember, you just need to unpack the objects in the same order you gave them to the
+    # prepare method.
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+    )
+
+    # Now you train the model
+    for epoch in range(5):
+        model.train()
+        for batch in train_dataloader:
+            inputs = (batch["image"] - mean) / std
+            outputs = model(inputs)
+            loss = torch.nn.functional.cross_entropy(outputs, batch["label"])
+            accelerator.backward(loss)
+            optimizer.step()
+            lr_scheduler.step()
+            optimizer.zero_grad()
+
+        model.eval()
+        accurate = 0
+        num_elems = 0
+        for batch in eval_dataloader:
+            inputs = (batch["image"] - mean) / std
+            with torch.no_grad():
+                outputs = model(inputs)
+            predictions = outputs.argmax(dim=-1)
+            accurate_preds = accelerator.gather(predictions) == accelerator.gather(batch["label"])
+            num_elems += accurate_preds.shape[0]
+            accurate += accurate_preds.long().sum()
+
+        eval_metric = accurate.item() / num_elems
+        # Use accelerator.print to print only on the main process.
+        accelerator.print(f"epoch {epoch}: {100 * eval_metric:.2f}")
+```
+
+## Using the notebook_launcher
+
+All that's left is to use the [`notebook_launcher`].
+
+You pass in the function, the arguments (as a tuple), and the number of processes to train on. (See the [documentation](../package_reference/launchers) for more information)
+
+```python
+from accelerate import notebook_launcher
+```
+
+```python
+args = ("fp16", 42, 64)
+notebook_launcher(training_loop, args, num_processes=2)
+```
+
+In the case of running on multiple nodes, you need to set up a Jupyter session at each node and run the launching cell at the same time.
+
+For an environment containing 2 nodes (computers) with 8 GPUs each and the main computer with an IP address of "172.31.43.8", it would look like so:
+
+```python
+notebook_launcher(training_loop, args, master_addr="172.31.43.8", node_rank=0, num_nodes=2, num_processes=8)
+```
+
+And in the second Jupyter session on the other machine:
+
+<Tip>
+
+    Notice how the `node_rank` has changed
+
+</Tip>
+
+```python
+notebook_launcher(training_loop, args, master_addr="172.31.43.8", node_rank=1, num_nodes=2, num_processes=8)
+```
+
+In the case of running on the TPU, it would look like so:
+
+```python
+model = create_model("resnet50d", pretrained=True, num_classes=len(label_to_id))
+
+args = (model, "fp16", 42, 64)
+notebook_launcher(training_loop, args, num_processes=8)
+```
+
+To launch the training process with elasticity, enabling fault tolerance, you can use the `elastic_launch` feature provided by PyTorch. This requires setting additional parameters such as `rdzv_backend` and `max_restarts`. Here is an example of how to use `notebook_launcher` with elastic capabilities:
+
+```python
+notebook_launcher(
+    training_loop,
+    args,
+    num_processes=2,
+    max_restarts=3
+)
+```
+
+As it's running it will print the progress as well as state how many devices you ran on. This tutorial was ran with two GPUs:
+
+```python out
+Launching training on 2 GPUs.
+epoch 0: 88.12
+epoch 1: 91.73
+epoch 2: 92.58
+epoch 3: 93.90
+epoch 4: 94.71
+```
+
+And that's it!
+
+Please note that [`notebook_launcher`] ignores the Accelerate config file, to launch based on the config use:
+
+```bash
+accelerate launch
+```
+
+## Debugging 
+
+A common issue when running the `notebook_launcher` is receiving a CUDA has already been initialized issue. This usually stems
+from an import or prior code in the notebook that makes a call to the PyTorch `torch.cuda` sublibrary. To help narrow down what went wrong,
+you can launch the `notebook_launcher` with `ACCELERATE_DEBUG_MODE=yes` in your environment and an additional check
+will be made when spawning that a regular process can be created and utilize CUDA without issue. (Your CUDA code can still be ran afterwards).
+
+## Conclusion
+
+This notebook showed how to perform distributed training from inside of a Jupyter Notebook. Some key notes to remember:
+
+- Make sure to save any code that use CUDA (or CUDA imports) for the function passed to [`notebook_launcher`]
+- Set the `num_processes` to be the number of devices used for training (such as number of GPUs, CPUs, TPUs, etc)
+- If using the TPU, declare your model outside the training loop function

docs/source/basic_tutorials/overview.md

@@ -0,0 +1,24 @@
+<!--Copyright 2022 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
+-->
+
+# Overview
+
+Welcome to the Accelerate tutorials! These introductory guides will help catch you up to speed on working with Accelerate.
+You'll learn how to modify your code to have it work with the API seamlessly, how to launch your script properly,
+and more!
+
+These tutorials assume some basic knowledge of Python and familiarity with the PyTorch framework.
+
+If you have any questions about Accelerate, feel free to join and ask the community on our [forum](https://discuss.huggingface.co/c/accelerate/18).

docs/source/basic_tutorials/tpu.md

@@ -0,0 +1,38 @@
+<!--Copyright 2024 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
+-->
+
+# TPU training
+
+A [TPU (Tensor Processing Unit)](https://cloud.google.com/tpu/docs/intro-to-tpu) is a type of hardware specifically designed for training models efficiently. Accelerate supports TPU training, but there are a few things you should be aware of, namely graph compilation. This tutorial briefly discusses compilation, and for more details, take a look at the [Training on TPUs with Accelerate](../concept_guides/training_tpu) guide.
+
+## Compilation
+
+A TPU creates a graph of all the operations in the training step such as the forward pass, backward pass and optimizer step. This is why the first training step always takes a while because building and compiling this graph takes time. But once compilation is complete, it is cached and all subsequent steps are much faster.
+
+The key is to avoid compiling your code again or else training is super slow. This means all your operations must be exactly the same:
+
+* all tensors in your batches must have the same length (for example, no dynamic padding for NLP tasks)
+* your code must be static (for example, no layers with for loops that have different lengths depending on the input such as a LSTM)
+
+## Weight tying
+
+A common language model design is to tie the weights of the embedding and softmax layers. However, moving the model to a TPU (either yourself or passing it to the [`~Accelerator.prepare`] method) breaks the weight tying and you'll need to retie the weights.
+
+To add special behavior (like weight tying) in your script for TPUs, set [`~Accelerator.distributed_type`] to `DistributedType.TPU` first. Then you can use the [`~transformers.PreTrainedModel.tie_weights`] method to tie the weights.
+
+```py
+if accelerator.distributed_type == DistributedType.TPU:
+    model.tie_weights()
+```

docs/source/basic_tutorials/troubleshooting.md

@@ -0,0 +1,211 @@
+<!--Copyright 2023 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
+-->
+
+# Troubleshoot
+
+This guide provides solutions to some issues you might encounter when using Accelerate. Not all errors are covered because Accelerate is an active library that is continuously evolving and there are many different use cases and distributed training setups. If the solutions described here don't help with your specific error, please take a look at the [Ask for help](#ask-for-help) section to learn where and how to get help.
+
+## Logging
+
+Logging can help you identify where an error is coming from. In a distributed setup with multiple processes, logging can be a challenge, but Accelerate provides the [`~accelerate.logging`] utility to ensure logs are synchronized.
+
+To troubleshoot an issue, use [`~accelerate.logging`] instead of the standard Python [`logging`](https://docs.python.org/3/library/logging.html#module-logging) module. Set the verbosity level (`INFO`, `DEBUG`, `WARNING`, `ERROR`, `CRITICAL`) with the `log_level` parameter, and then you can either:
+
+1. Export the `log_level` as the `ACCELERATE_LOG_LEVEL` environment variable.
+2. Pass the `log_level` directly to `get_logger`.
+
+For example, to set `log_level="INFO"`:
+
+```py
+from accelerate.logging import get_logger
+
+logger = get_logger(__name__, log_level="DEBUG")
+```
+
+By default, the log is called on main processes only. To call it on all processes, pass `main_process_only=False`.
+If a log should be called on all processes and in order, also pass `in_order=True`.
+
+```py
+from accelerate.logging import get_logger
+
+logger = get_logger(__name__, log_level="DEBUG")
+# log all processes
+logger.debug("thing_to_log", main_process_only=False)
+# log all processes in order
+logger.debug("thing_to_log", main_process_only=False, in_order=True)
+```
+
+## Hanging code and timeout errors
+
+There can be many reasons why your code is hanging. Let's take a look at how to solve some of the most common issues that can cause your code to hang.
+
+### Mismatched tensor shapes
+
+Mismatched tensor shapes is a common issue that can cause your code to hang for a significant amount of time on a distributed setup.
+
+When running scripts in a distributed setup, functions such as [`Accelerator.gather`] and [`Accelerator.reduce`] are necessary to grab tensors across devices to collectively perform operations on them. These (and other) functions rely on `torch.distributed` to perform a `gather` operation, which requires tensors to have the **exact same shape** across all processes. When the tensor shapes don't match, your code hangs and you'll eventually hit a timeout exception.
+
+You can use Accelerate's operational debug mode to immediately catch this issue. We recommend enabling this mode during the `accelerate config` setup, but you can also enable it from the CLI, as an environment variable, or by manually editing the `config.yaml` file.
+
+<hfoptions id="mismatch">
+<hfoption id="CLI">
+
+```bash
+accelerate launch --debug {my_script.py} --arg1 --arg2
+```
+
+</hfoption>
+<hfoption id="environment variable">
+
+If enabling debug mode as an environment variable, you don't need to call `accelerate launch`.
+
+```bash
+ACCELERATE_DEBUG_MODE="1" torchrun {my_script.py} --arg1 --arg2
+```
+
+</hfoption>
+<hfoption id="config.yaml">
+
+Add `debug: true` to your `config.yaml` file.
+
+```yaml
+compute_environment: LOCAL_MACHINE
+debug: true
+```
+
+</hfoption>
+</hfoptions>
+
+Once you enable debug mode, you should get a traceback that points to the tensor shape mismatch issue.
+
+```py
+Traceback (most recent call last):
+  File "/home/zach_mueller_huggingface_co/test.py", line 18, in <module>
+    main()
+  File "/home/zach_mueller_huggingface_co/test.py", line 15, in main
+    broadcast_tensor = broadcast(tensor)
+  File "/home/zach_mueller_huggingface_co/accelerate/src/accelerate/utils/operations.py", line 303, in wrapper
+accelerate.utils.operations.DistributedOperationException:
+
+Cannot apply desired operation due to shape mismatches. All shapes across devices must be valid.
+
+Operation: `accelerate.utils.operations.broadcast`
+Input shapes:
+  - Process 0: [1, 5]
+  - Process 1: [1, 2, 5]
+```
+
+### Early stopping
+
+For early stopping in distributed training, if each process has a specific stopping condition (e.g. validation loss), it may not be synchronized across all processes. As a result, a break can happen on process 0 but not on process 1 which will cause your code to hang indefinitely until a timeout occurs.
+
+If you have early stopping conditionals, use the `set_trigger` and `check_trigger` methods to make sure all the processes
+are ended correctly.
+
+```py
+# Assume `should_do_breakpoint` is a custom defined function that returns a conditional, 
+# and that conditional might be true only on process 1
+if should_do_breakpoint(loss):
+    accelerator.set_trigger()
+
+# Later in the training script when we need to check for the breakpoint
+if accelerator.check_trigger():
+    break
+```
+
+### Low kernel versions on Linux
+
+On Linux with kernel version < 5.5, hanging processes have been reported. To avoid this problem, upgrade your system to a later kernel version.
+
+### MPI
+
+If your distributed CPU training job using MPI is hanging, ensure that you have
+[passwordless SSH](https://www.open-mpi.org/faq/?category=rsh#ssh-keys) setup (using keys) between the nodes. This means
+that for all nodes in your hostfile, you should to be able to SSH from one node to another without being prompted for a password.
+
+Next, try to run the `mpirun` command as a sanity check. For example, the command below should print out the
+hostnames for each of the nodes.
+
+```bash
+mpirun -f hostfile -n {number of nodes} -ppn 1 hostname
+```
+
+## Out-of-Memory
+
+One of the most frustrating errors when it comes to running training scripts is hitting "Out-of-Memory" on devices like CUDA, XPU or CPU. The entire script needs to be restarted and any progress is lost.
+
+To address this problem, Accelerate provides the [`find_executable_batch_size`] utility that is heavily based on [toma](https://github.com/BlackHC/toma).
+This utility retries code that fails due to OOM (out-of-memory) conditions and automatically lowers batch sizes. For each OOM condition, the algorithm decreases the batch size by half and retries the code until it succeeds.
+
+To use [`find_executable_batch_size`], restructure your training function to include an inner function with `find_executable_batch_size` and build your dataloaders inside it. At a minimum, this only takes 4 new lines of code.
+
+<Tip warning={true}> 
+
+The inner function **must** take batch size as the first parameter, but we do not pass one to it when called. The wrapper will handle this for you. Any object (models, optimizers) that consumes device memory and is passed to the [`Accelerator`] also **must** be declared inside the inner function.
+
+</Tip>
+
+```diff
+def training_function(args):
+    accelerator = Accelerator()
+
++   @find_executable_batch_size(starting_batch_size=args.batch_size)
++   def inner_training_loop(batch_size):
++       nonlocal accelerator # Ensure they can be used in our context
++       accelerator.free_memory() # Free all lingering references
+        model = get_model()
+        model.to(accelerator.device)
+        optimizer = get_optimizer()
+        train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size)
+        lr_scheduler = get_scheduler(
+            optimizer, 
+            num_training_steps=len(train_dataloader)*num_epochs
+        )
+        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+            model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+        )
+        train(model, optimizer, train_dataloader, lr_scheduler)
+        validate(model, eval_dataloader)
++   inner_training_loop()
+```
+
+## Non-reproducible results between device setups
+
+If you changed the device setup and observe different model performance, it is likely you didn't update your script when moving from one setup to another. Even if you're using the same script with the same batch size, the results will still be different on a TPU, multi-GPU, and single GPU.
+
+For example, if you were training on a single GPU with a batch size of 16 and you move to a dual GPU setup, you need to change the batch size to 8 to have the same effective batch size. This is because when training with Accelerate, the batch size passed to the dataloader is the **batch size per GPU**.
+
+To make sure you can reproduce the results between the setups, make sure to use the same seed, adjust the batch size accordingly, and consider scaling the learning rate.
+
+For more details and a quick reference for batch sizes, check out the [Comparing performance between different device setups](../concept_guides/performance) guide.
+
+## Performance issues on different GPUs
+
+If your multi-GPU setup consists of different GPUs, you may encounter some performance issues:
+
+- There may be an imbalance in GPU memory between the GPUs. In this case, the GPU with the smaller memory will limit the batch size or the size of the model that can be loaded onto the GPUs.
+- If you are using GPUs with different performance profiles, the performance will be driven by the slowest GPU you are using because the other GPUs will have to wait for it to complete its workload.
+
+Vastly different GPUs within the same setup can lead to performance bottlenecks.
+
+## Ask for help
+
+If none of the solutions and advice here helped resolve your issue, you can always reach out to the community and Accelerate team for help.
+
+- Ask for help on the Hugging Face forums by posting your question in the [Accelerate category](https://discuss.huggingface.co/c/accelerate/18). Make sure to write a descriptive post with relevant context about your setup and reproducible code to maximize the likelihood that your problem is solved!
+
+- Post a question on [Discord](http://hf.co/join/discord), and let the team and the community help you.
+
+- Create an Issue on the Accelerate [GitHub repository](https://github.com/huggingface/accelerate/issues) if you think you've found a bug related to the library. Include context regarding the bug and details about your distributed setup to help us better figure out what's wrong and how we can fix it.

docs/source/concept_guides/big_model_inference.md

@@ -0,0 +1,341 @@
+<!--Copyright 2022 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
+-->
+
+# Loading big models into memory
+
+When loading a pre-trained model in PyTorch, the usual workflow looks like this:
+
+```py
+import torch
+
+my_model = ModelClass(...)
+state_dict = torch.load(checkpoint_file)
+my_model.load_state_dict(state_dict)
+```
+
+In plain English, those steps are:
+1. Create the model with randomly initialized weights
+2. Load the model weights (in a dictionary usually called a state dict) from the disk
+3. Load those weights inside the model
+
+While this works very well for regularly sized models, this workflow has some clear limitations when we deal with a huge model: in step 1, we load a full version of the model in RAM, and spend some time randomly initializing the weights (which will be discarded in step 3). In step 2, we load another full version of the model in RAM, with the pre-trained weights. If you're loading a model with 6 billion parameters, this means you will need 24GB of RAM for each copy of the model, so 48GB in total (half of it to load the model in FP16).
+
+<Tip warning={true}>
+
+This API is quite new and still in its experimental stage. While we strive to provide a stable API, it's possible some small parts of the public API will change in the future.
+
+</Tip>
+
+## How the Process Works: A Quick Overview
+
+<Youtube id="MWCSGj9jEAo" />
+
+## How the Process Works: Working with Code
+
+### Instantiating an empty model
+
+The first tool Accelerate introduces to help with big models is a context manager [`init_empty_weights`] that helps you initialize a model without using any RAM so that step 1 can be done on models of any size. Here is how it works:
+
+```py
+from accelerate import init_empty_weights
+
+with init_empty_weights():
+    my_model = ModelClass(...)
+```
+
+For instance:
+
+```py
+with init_empty_weights():
+    model = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)])
+```
+
+initializes an empty model with a bit more than 100B parameters. Behind the scenes, this relies on the meta device introduced in PyTorch 1.9. During the initialization under the context manager, each time a parameter is created, it is instantly moved to that device.
+
+<Tip warning={true}>
+
+    You can't move a model initialized like this on CPU or another device directly, since it doesn't have any data. It's also very likely that a forward pass with that empty model will fail, as not all operations are supported on the meta device.
+
+</Tip>
+
+### Sharded checkpoints
+
+It's possible your model is so big that even a single copy won't fit in RAM. That doesn't mean it can't be loaded: if you have one or several GPUs, this is more memory available to store your model. In this case, it's better if your checkpoint is split into several smaller files that we call checkpoint shards.
+
+Accelerate will handle sharded checkpoints as long as you follow the following format: your checkpoint should be in a folder, with several files containing the partial state dicts, and there should be an index in the JSON format that contains a dictionary mapping parameter names to the file containing their weights. You can easily shard your model with [`~Accelerator.save_model`]. For instance, we could have a folder containing:
+
+```bash
+first_state_dict.bin
+index.json
+second_state_dict.bin
+```
+
+with index.json being the following file:
+
+```
+{
+  "linear1.weight": "first_state_dict.bin",
+  "linear1.bias": "first_state_dict.bin",
+  "linear2.weight": "second_state_dict.bin",
+  "linear2.bias": "second_state_dict.bin"
+}
+```
+
+and `first_state_dict.bin` containing the weights for `"linear1.weight"` and `"linear1.bias"`, `second_state_dict.bin` the ones for `"linear2.weight"` and `"linear2.bias"`
+
+### Loading weights
+
+The second tool Accelerate introduces is a function [`load_checkpoint_and_dispatch`], that will allow you to load a checkpoint inside your empty model. This supports full checkpoints (a single file containing the whole state dict) as well as sharded checkpoints. It will also automatically dispatch those weights across the devices you have available (GPUs, CPU RAM), so if you are loading a sharded checkpoint, the maximum RAM usage will be the size of the biggest shard.
+
+If you want to use big model inference with Transformers models, check out this [documentation](https://huggingface.co/docs/transformers/main/en/main_classes/model#large-model-loading).
+
+Here is how we can use this to load the [GPT2-1.5B](https://huggingface.co/marcsun13/gpt2-xl-linear-sharded) model.
+
+Let's download the sharded version of this model.
+
+```bash
+pip install huggingface_hub
+```
+
+```py
+from huggingface_hub import snapshot_download
+checkpoint = "marcsun13/gpt2-xl-linear-sharded"
+weights_location = snapshot_download(repo_id=checkpoint)
+```
+
+In order to initialize the model, we will use the library minGPT. 
+
+```bash
+git clone https://github.com/karpathy/minGPT.git
+pip install minGPT/
+```
+
+```py
+from accelerate import init_empty_weights
+from mingpt.model import GPT
+
+model_config = GPT.get_default_config()
+model_config.model_type = 'gpt2-xl'
+model_config.vocab_size = 50257
+model_config.block_size = 1024
+
+with init_empty_weights():
+    model = GPT(model_config)
+```
+
+Then, load the checkpoint we just downloaded with:
+
+```py
+from accelerate import load_checkpoint_and_dispatch
+
+model = load_checkpoint_and_dispatch(
+    model, checkpoint=weights_location, device_map="auto", no_split_module_classes=['Block']
+)
+```
+
+By passing `device_map="auto"`, we tell Accelerate to determine automatically where to put each layer of the model depending on the available resources:
+- first, we use the maximum space available on the GPU(s)
+- if we still need space, we store the remaining weights on the CPU
+- if there is not enough RAM, we store the remaining weights on the hard drive as memory-mapped tensors
+
+
+#### `no_split_module_classes`
+
+This parameter will indicate that some of the modules with the name `"Block"` should not be split across different devices. You should set here all blocks that 
+include a residual connection of some kind.
+
+
+#### The `device_map`
+
+You can see the `device_map` that Accelerate picked by accessing the `hf_device_map` attribute of your model:
+
+```py
+model.hf_device_map
+```
+
+```python out
+{'transformer.wte': 0,
+ 'transformer.wpe': 0,
+ 'transformer.drop': 0,
+ 'transformer.h.0': 0,
+ ...
+ 'transformer.h.21': 0, 
+ 'transformer.h.22': 1, 
+ 'transformer.h.23': 1, 
+ 'transformer.h.24': 1,
+ ...
+ 'transformer.h.47': 1, 
+ 'transformer.ln_f': 1, 
+ 'lm_head': 1}
+ ```
+
+It's fully possible to create your own device map for the layers to use as well, specifying the GPU device to use (a number), `"cpu"`, or `"disk"` and pass this in:
+
+```python
+device_map = {
+    "transformer.wte": "cpu",
+    "transformer.wpe": 0,
+    "transformer.drop": "cpu",
+    "transformer.h.0": "disk"
+}
+
+model = load_checkpoint_and_dispatch(
+    model, checkpoint=weights_location, device_map=device_map
+)
+
+```
+
+### Run the model
+
+Now that we have done this, our model lies across several devices, and maybe the hard drive. But it can still be used as a regular PyTorch model:
+
+```py
+from mingpt.bpe import BPETokenizer
+tokenizer = BPETokenizer()
+inputs = tokenizer("Hello, my name is").to(0)
+
+outputs = model.generate(x1, max_new_tokens=10, do_sample=False)[0]
+tokenizer.decode(outputs.cpu().squeeze())
+```
+
+Behind the scenes, Accelerate added hooks to the model, so that:
+- at each layer, the inputs are put on the right device (so even if your model is spread across several GPUs, it works)
+- for the weights offloaded on the CPU, they are put on a GPU just before the forward pass and cleaned up just after
+- for the weights offloaded on the hard drive, they are loaded in RAM then put on a GPU just before the forward pass and cleaned up just after
+
+This way, your model can run for inference even if it doesn't fit on one of the GPUs or the CPU RAM!
+
+<Tip warning={true}>
+
+    This only supports the inference of your model, not training. Most of the computation happens behind `torch.no_grad()` context managers to avoid spending some GPU memory with intermediate activations.
+
+</Tip>
+
+### Designing a device map
+
+You can let Accelerate handle the device map computation by setting `device_map` to one of the supported options (`"auto"`, `"balanced"`, `"balanced_low_0"`, `"sequential"`) or create one yourself if you want more control over where each layer should go.
+
+<Tip>
+
+    You can derive all sizes of the model (and thus compute a `device_map`) on a model that is on the meta device.
+
+</Tip>
+
+All the options will produce the same result when you don't have enough GPU memory to accommodate the whole model (which is to fit everything that can on the GPU, then offload weights on the CPU or even on the disk if there is not enough RAM). 
+
+When you have more GPU memory available than the model size, here is the difference between each option:
+- `"auto"` and `"balanced"` evenly split the model on all available GPUs, making it possible for you to use a batch size greater than 1.
+- `"balanced_low_0"` evenly splits the model on all GPUs except the first one, and only puts on GPU 0 what does not fit on the others. This option is great when you need to use GPU 0 for some processing of the outputs, like when using the `generate` function for Transformers models
+- `"sequential"` will fit what it can on GPU 0, then move on GPU 1 and so forth (so won't use the last GPUs if it doesn't need to).
+
+<Tip>
+
+    The options `"auto"` and `"balanced"` produce the same results for now, but the behavior of `"auto"` might change in the future if we find a strategy that makes more sense, while `"balanced"` will stay stable.
+
+</Tip>
+
+First note that you can limit the memory used on each GPU by using the `max_memory` argument (available in [`infer_auto_device_map`] and in all functions using it). When setting `max_memory`, you should pass along a dictionary containing the GPU identifiers (for instance `0`, `1` etc.) and the `"cpu"` key for the maximum RAM you want to use for CPU offload. The values can either be an integer (in bytes) or a string representing a number with its unit, such as `"10GiB"` or `"10GB"`.
+
+Here is an example where we don't want to use more than 10GiB on each of the two GPUs and no more than 30GiB of CPU RAM for the model weights:
+
+```python
+from accelerate import infer_auto_device_map
+
+device_map = infer_auto_device_map(my_model, max_memory={0: "10GiB", 1: "10GiB", "cpu": "30GiB"})
+```
+
+<Tip warning={true}>
+
+    When a first allocation happens in PyTorch, it loads CUDA kernels which take about 1-2GB of memory depending on the GPU. Therefore you always have less usable memory than the actual size of the GPU. To see how much memory is actually used do `torch.ones(1).cuda()` and look at the memory usage.
+
+    Therefore when you create memory maps with `max_memory` make sure to adjust the available memory accordingly to avoid out-of-memory errors.
+
+</Tip>
+
+Additionally, if you do some additional operations with your outputs without placing them back on the CPU (for instance inside the `generate` method of Transformers) and if you placed your inputs on a GPU, that GPU will consume more memory than the others (Accelerate always place the output back to the device of the input). Therefore if you would like to optimize the maximum batch size and you have many GPUs, give the first GPU less memory. For example, with BLOOM-176B on 8x80 A100 setup, the close-to-ideal map is:
+
+```python
+max_memory = {0: "30GIB", 1: "46GIB", 2: "46GIB", 3: "46GIB", 4: "46GIB", 5: "46GIB", 6: "46GIB", 7: "46GIB"}
+```
+as you can see we gave the remaining 7 GPUs ~50% more memory than GPU 0.
+
+If you opt to fully design the `device_map` yourself, it should be a dictionary with keys being module names of your model and values being a valid device identifier (for instance an integer for the GPUs) or `"cpu"` for CPU offload, `"disk"` for disk offload. The keys need to cover the whole model, you can then define your device map as you wish: for instance, if your model has two blocks (let's say `block1` and `block2`) which each contain three linear layers (let's say `linear1`, `linear2` and `linear3`), a valid device map can be:
+
+```python
+device_map = {"block1": 0, "block2": 1}
+```
+
+another one that is valid could be:
+
+```python
+device_map = {"block1": 0, "block2.linear1": 0, "block2.linear2": 1, "block2.linear3": 1}
+```
+
+On the other hand, this one is not valid as it does not cover every parameter of the model:
+
+```python
+device_map = {"block1": 0, "block2.linear1": 1, "block2.linear2": 1}
+```
+
+<Tip>
+
+    To be the most efficient, make sure your device map puts the parameters on the GPUs in a sequential manner (e.g. don't put one of the first weights on GPU 0, then weights on GPU 1 and the last weight back to GPU 0) to avoid making many transfers of data between the GPUs.
+
+</Tip>
+
+## CPU offload only
+
+If you want to offload your model on CPU, you can use [`cpu_offload`]. As a result, all parameters of the model will be offloaded and only one copy of the state dict of the model will be kept. During the forward pass, parameters will be extracted from that state dict and put on the execution device and passed as they are needed, then offloaded again. 
+
+```python
+cpu_offload(model, execution_device)
+```
+
+You can also use [`cpu_offload_with_hook`]. This function will offloads a model on the CPU and puts it back to an execution device when executed. The difference with [`cpu_offload`] is that the model stays on the execution device after the forward and is only offloaded again when the `offload` method of the returned `hook` is called. Furthermore, [`cpu_offload_with_hook`] is more performant but less memory saving. It is useful for pipelines running a model in a loop:
+
+```python
+model_1, hook_1 = cpu_offload_with_hook(model_1, execution_device)
+model_2, hook_2 = cpu_offload_with_hook(model_2, execution_device, prev_module_hook=hook_1)
+model_3, hook_3 = cpu_offload_with_hook(model_3, execution_device, prev_module_hook=hook_2)
+
+hid_1 = model_1(input)
+for i in range(50):
+    # model1 is offloaded on the CPU at the first iteration, model 2 stays on the GPU for this whole loop.
+    hid_2 = model_2(hid_1)
+# model2 is offloaded to the CPU just before this forward.
+hid_3 = model_3(hid_3)
+
+# For model3, you need to manually call the hook offload method.
+hook_3.offload()
+```
+
+## Disk offload only
+
+To perform disk offload, you can use [`disk_offload`]. As a result, all parameters of the model will be offloaded as memory-mapped array in a given folder. During the forward pass, parameters will be accessed from that folder and put on the execution device passed as they are needed, then offloaded again.
+
+```python
+disk_offload(model, offload_dir, execution_device)
+```
+
+## Limits and further development
+
+We are aware of the current limitations in the API:
+
+- [`infer_auto_device_map`] (or `device_map="auto"` in [`load_checkpoint_and_dispatch`]) tries to maximize GPU and CPU RAM it sees available when you execute it. While PyTorch is very good at managing GPU RAM efficiently (and giving it back when not needed), it's not entirely true with Python and CPU RAM. Therefore, an automatically computed device map might be too intense on the CPU. Move a few modules to the disk device if you get crashes due to a lack of RAM.
+- [`infer_auto_device_map`] (or `device_map="auto"` in [`load_checkpoint_and_dispatch`]) attributes devices sequentially (to avoid moving things back and forth) so if your first layer is bigger than the size of the GPU you have, it will end up with everything on the CPU/Disk.
+- [`load_checkpoint_and_dispatch`] and [`load_checkpoint_in_model`] do not perform any check on the correctness of your state dict compared to your model at the moment (this will be fixed in a future version), so you may get some weird errors if trying to load a checkpoint with mismatched or missing keys.
+- The model parallelism used when your model is split on several GPUs is naive and not optimized, meaning that only one GPU works at a given time and the other sits idle.
+- When weights are offloaded on the CPU/hard drive, there is no pre-fetching (yet, we will work on this for future versions) which means the weights are put on the GPU when they are needed and not before.
+- Hard-drive offloading might be very slow if the hardware you run on does not have fast communication between disk and CPU (like NVMes).

docs/source/concept_guides/deferring_execution.md

@@ -0,0 +1,130 @@
+<!--Copyright 2022 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
+-->
+
+# Executing and deferring jobs
+
+When you run your usual script, instructions are executed in order. Using Accelerate to deploy your script on several
+GPUs at the same time introduces a complication: while each process executes all instructions in order, some may be
+faster than others.
+
+You might need to wait for all processes to have reached a certain point before executing a given instruction. For
+instance, you shouldn't save a model before being sure every process is done with training, and you wouldn't want to 
+continue training before all the model weights have been loaded in. To do this, just write the following line in your code:
+
+```
+accelerator.wait_for_everyone()
+```
+
+This instruction will block all the processes that arrive first until all the other processes have reached that
+point (if you run your script on just one GPU or CPU, this won't do anything).
+
+A few example cases of when to use this utility are listed below:
+
+<Tip>
+
+    Some of these are utilized with the [`~Accelerator.main_process_first`] context manager, which utilizes [`~Accelerator.wait_for_everyone`] to 
+    run a particular set of code on the main process beforehand before triggering and launching the other processes
+
+</Tip>
+
+## Downloading a Dataset 
+
+When downloading a dataset, you should download it first on the main process and then load the cached dataset afterward
+
+<Tip>
+
+    `load_dataset` will perform a lock under the hood to stop multiple downloads from happening at once, but if you are downloading something 
+    not using this library you should use this method.
+    
+</Tip>
+
+```python
+with accelerator.main_process_first():
+    datasets = load_dataset("glue", "mrpc")
+```
+
+Under the hood this is the same as calling: 
+
+```python
+# First do something on the main process
+if accelerator.is_main_process:
+    datasets = load_dataset("glue", "mrpc")
+else:
+    accelerator.wait_for_everyone()
+
+# And then send it to the rest of them
+if not accelerator.is_main_process:
+    datasets = load_dataset("glue", "mrpc")
+else:
+    accelerator.wait_for_everyone()
+```
+
+## Saving the `state_dict`
+
+When saving the `state_dict` of the model, since you would normally save one file on just the main process
+you should specify that:
+
+```python
+if accelerator.is_main_process:
+    model = accelerator.unwrap_model(model)
+    torch.save(model.state_dict(), "weights.pth")
+```
+
+## Loading in the `state_dict`
+
+When loading in the `state_dict` to a model, optimizer, or scheduler, you should wait 
+for all workers to have the weights loaded in before moving on to training
+
+```python
+with accelerator.main_process_first():
+    state = torch.load("weights.pth")
+    model.load_state_dict(state)
+```
+
+## Applying a multi-worker CPU operation 
+
+Applying a `map()` operation on multiple workers, such as tokenizing should be done on the 
+main process first, and then propagated to each one. 
+
+```python
+datasets = load_dataset("glue", "mrpc")
+
+with accelerator.main_process_first():
+    tokenized_datasets = datasets.map(
+        tokenize_function,
+        batched=True,
+        remove_columns=["idx", "sentence1", "sentence2"],
+    )
+```
+
+## Applying checks such as Early Stopping
+
+To have a check that works with a flag set by a particular process, the `set_trigger` and `check_trigger` API should be used. Useful examples
+for doing so can include situations such as using early stopping and monitoring the loss (as each loss slightly differs on each process).
+
+Call [`Accelerator.set_trigger`] when your condition has been met, and [`Accelerator.check_trigger`] when checking if that condition has been met in any process:
+
+```python
+for (x,y) in data_loader:
+    logits = model(x)
+    loss = loss_func(logits, y)
+    # Assume `should_do_early_stopping` is a custom defined function that returns a conditional
+    if should_do_early_stopping(loss):
+        accelerator.set_trigger()
+
+    # Later in the training script when we need to check for the breakpoint
+    if accelerator.check_trigger():
+        break
+```

docs/source/concept_guides/fsdp1_vs_fsdp2.md

@@ -0,0 +1,105 @@
+<!--Copyright 2025 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
+-->
+
+# FSDP1 vs FSDP2
+
+This guide explains the key differences between `FSDP1` and `FSDP2` and helps you migrate your existing code to use `FSDP2` with minimal changes.
+
+## How is FSDP2 better than FSDP1?
+
+First, we want to understand how `FSDP1` and `FSDP2` work internally to understand the differences between them. This also helps us understand the limitations of `FSDP1` and how `FSDP2` solves them.
+
+We'll be discussing a scenario where we have a single `Layer` that contains 3 `Linear` layers and is wrapped using `FSDP` to be sharded across 2 GPUs.
+
+<div align="center">
+  <img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/accelerate/layer.png" alt="Layer">
+</div>
+
+### FSDP1
+First, we have to understand the original `FSDP1` and the limitations it brings. It represents each `FSDP` module as a single `FlatParameter` which is a single 1D tensor that contains all of the module parameters, which then get sharded across ranks. I.e. if you wrap the `Layer` with `FSDP1`, you'd achieve something as such:
+
+<div align="center">
+  <img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/accelerate/fsdp1.png" alt="FSDP1">
+</div>
+
+You might notice a problem. The whole `Layer` gets flattened into a single `FlatParameter`, which then gets sharded across ranks. But if it's a single `FlatParameter` object, how do we store metadata? That is one of the limitations. Properly storing per-parameter metadata such as `dtype`, `requires_grad`, etc. is not possible without some ugly hacks.
+
+### FSDP2
+This is why `FSDP2` was introduced. It doesn't use `FlatParameter`, instead it uses `DTensor` which is short for "Distributed Tensor". Each `DTensor` basically represents a vanilla `torch.Tensor` that has been sharded across ranks. It contains metadata about the original `torch.Tensor` and how it's sharded, what is the [placement type](https://pytorch.org/docs/stable/distributed.tensor.html#module-torch.distributed.tensor.placement_types) and so on. This is why it's called `per-parameter sharding`. The following figure shows the difference:
+
+<div align="center">
+  <img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/accelerate/fsdp2.png" alt="FSDP2">
+</div>
+
+Each Parameter of the original `Layer` is sharded across the 0th dimension, and split between 2 GPUs. Now, each `Linear` layer is a separate `DTensor` and storing metadata per-parameter is possible and straightforward.
+
+
+> [!TIP] 
+> In the image above, the tensors were sharded across the 1st dimension for the sake of fitting the image on the screen, in reality, they are sharded across the 0th dimension as stated above
+
+## What does FSDP2 offer?
+
+`FSDP2` is a new and improved version of PyTorch's fully-sharded data parallel training API. Its main advantage is using `DTensor` to represent sharded parameters. Compared to `FSDP1`, it offers:
+- Simpler internal implementation, where each `Parameter` is a separate `DTensor`
+- Enables simple partial parameter freezing because of the above, which makes methods as [`LORA`](https://arxiv.org/abs/2106.09685) work out of the box
+- With `DTensor`, `FSDP2` supports mixing `fp8` and other parameter types in the same model out of the box
+- Faster and simpler checkpointing without extra communication across ranks using `SHARDED_STATE_DICT` and [`torch.distributed.checkpoint`](https://pytorch.org/docs/stable/distributed.checkpoint.html), this way, each rank only saves its own shard and corresponding metadata
+- For loading, it uses a `state_dict` of the sharded model to directly load the sharded parameters
+- Support for asynchronous checkpointing, where parameters are first copied to CPU memory, after this, main thread continues training while another thread stores the parameters on disk
+- Memory efficiency and deterministic memory usage, `FSDP2` doesn't use `recordStream` anymore and uses stream-to-stream synchronization (for more technical details see [this forum post](https://dev-discuss.pytorch.org/t/fsdp-cudacachingallocator-an-outsider-newb-perspective/1486) and [this issue](https://github.com/pytorch/pytorch/issues/114299))
+- In the future, optimizations of the communication patterns via `torch.compile` are planned, further improving the performance and memory efficiency
+
+
+## API Differences
+
+We have already discussed the internal differences, now let's discuss the differences, you, as a user, will need to know. 
+
+Here are the main changes in configuration options when using `FSDP2` through the `accelerate` CLI:
+
+Previous (`FSDP1`) | New (`FSDP2`) | What Changed
+-- | -- | --
+`--fsdp_sharding_strategy` | `--fsdp_reshard_after_forward` | replaces `--fsdp_sharding_strategy`, changed to `true` (previously `FULL_SHARD`) or `false` (previously `SHARD_GRAD_OP`)
+`--fsdp_backward_prefetch` | \*\***REMOVED**\*\* | `FSDP2` uses previous `BACKWARD_PRE` option by default, as only this allows communication and computation overlap
+`--fsdp_forward_prefetch` | \*\***NOT YET IMPLEMENTED**\*\* | How to implement this is under active discussion, for now it is not supported in `FSDP2`
+`--fsdp_sync_module_states` | \*\***REMOVED**\*\* | with `FSDP2`, this parameter becomes redundant
+`--fsdp_cpu_ram_efficient_loading` | `--fsdp_cpu_ram_efficient_loading` | if `true`, `FSDP2` will similarly load the model only on rank 0, and then parameters get synced to other ranks, this is the same behavior as `FSDP1`, however, setting `--fsdp_sync_module_states` isn't required anymore
+`--fsdp_state_dict_type` | `--fsdp_state_dict_type` | `LOCAL_STATE_DICT` becomes obsolete and with `FSDP2` `SHARDED_STATE_DICT` is the default option, which results in no extra communication and each rank saving its own shard, other possible option is `FULL_STATE_DICT` which results in extra communication and spike in memory usage but saves the full model from rank 0.
+`--fsdp_use_orig_params` | \*\***REMOVED**\*\* | `FSDP2` uses a `DTensor` class on the background, which means it *always* uses the original parameters by default
+\*\***NEW**\*\* | `--fsdp_version` | `1` is the default option, to not break existing code, set to `2` to use `FSDP2`
+
+For all other options that remain unchanged, see the [`FSDP` documentation](../usage_guides/fsdp.md).
+
+## How to Switch to FSDP2
+
+### If using Python code:
+Simply set `fsdp_version=2` when creating your plugin and replace options according to the table above.
+
+```python
+from accelerate import FullyShardedDataParallelPlugin, Accelerator
+
+fsdp_plugin = FullyShardedDataParallelPlugin(
+    fsdp_version=2
+    # other options...
+)
+accelerator = Accelerator(fsdp_plugin=fsdp_plugin)
+```
+
+### If using YAML config:
+Use our conversion tool:
+```bash
+accelerate to-fsdp2 --config_file config.yaml --output_file new_config.yaml
+```
+
+This will automatically convert all FSDP1 settings to their FSDP2 equivalents. Use `--overwrite` to update the existing file instead of creating a new one.

docs/source/concept_guides/fsdp_and_deepspeed.md

@@ -0,0 +1,192 @@
+<!--Copyright 2024 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
+-->
+
+# FSDP vs DeepSpeed
+
+Accelerate offers flexibilty of training frameworks, by integrating two extremely powerful tools for distributed training, namely [Pytorch FSDP](../usage_guides/fsdp) and [Microsoft DeepSpeed](../usage_guides/deepspeed). The aim of this tutorial is to draw parallels, as well as to outline potential differences, to empower the user to switch seamlessly between these two frameworks.
+
+<Tip>
+
+  To switch between the frameworks, we recommend launching code `accelerate launch` passing in the correct config file with `--config_file`, or passing in the respective arguments directly for [FSDP and DeepSpeed](../package_reference/cli#accelerate-launch) .
+
+  Example Accelerate configurations can be found here for [DeepSpeed](../usage_guides/deepspeed#accelerate-deepspeed-plugin) and [FSDP](../usage_guides/fsdp#how-it-works-out-of-the-box), or in the [example zoo under "Launch Configurations"](../usage_guides/explore)
+ 
+</Tip>
+
+<Tip warning={true}>
+
+This tutorial is for single-node, multi-GPU, scenarios only.
+
+</Tip>
+
+## Configuring Functionalities
+
+Model tensors are split into different GPUs in an attempt to scale up model sizes; this is termed *sharding* in FSDP, and *partitioning* in DeepSpeed. FSDP sharding and DeepSpeed ZeRO (partitioning) stages are configured by `--fsdp_sharding_strategy`, and `--zero_stage`, respectively.  In particular, FSDP `FULL_SHARD` maps to DeepSpeed ZeRO stage `3`; see this [comprehensive mapping between FSDP sharding and DeepSpeed ZeRO settings](../usage_guides/fsdp#mapping-between-fsdp-sharding-strategies-and-deepspeed-zero-stages). The below table summarizes and groups similar settings:
+
+Group | Framework | Configuration | Example | Restrictions (if any)
+--|--|--|--|--
+sharding / partitioning | FSDP<br>DeepSpeed | `--fsdp_sharding_strategy`<br>`--zero_stage` | `1` (`FULL_SHARD`) <br>`3` | 
+offload | FSDP<br>DeepSpeed | `--fsdp_offload_params`<br>`--offload_param_device`<br>`--offload_optimizer_device` | `true`<br>`cpu`<br>`cpu` | all or nothing <br><br> 
+model loading | FSDP<br>DeepSpeed | <span style="white-space:nowrap;">`--fsdp_cpu_ram_efficient_loading`</span><br>`--zero3_init_flag` | `true`<br>`true` | <br>only ZeRO 3
+efficient checkpointing | FSDP<br>DeepSpeed | `--fsdp_state_dict_type`<br>`--zero3_save_16bit_model` |  `SHARDED_STATE_DICT`<br>`true` |  <br>only ZeRO 3
+weights prefetching | FSDP<br><br>DeepSpeed | `--fsdp_forward_prefetch`<br>`--fsdp_backward_prefetch`<br>None | `true`<br>`BACKWARD_PRE` | <br><br>
+model | FSDP<br><br>DeepSpeed |  `--fsdp_auto_wrap_policy`<br><span style="white-space:nowrap;">`--fsdp_transformer_layer_cls_to_wrap`</span><br>None | `TRANSFORMER_BASED_WRAP`<br><Layer Class> |<br>Usually not needed <br>Transparent to user.
+parameters summoning | FSDP<br>DeepSpeed | `--fsdp_use_orig_params`<br>None | `true` | required for `torch.compile`<br>Transparent to user
+parameters syncing | FSDP<br>DeepSpeed | `--fsdp_sync_module_states`<br>None | `true` | 
+training | FSDP<br>DeepSpeed | None<br>`--gradient_accumulation_steps`<br>`--gradient_clipping` | <br>`auto`<br>`auto` | Transparent to user
+
+For detailed descriptions of the above, refer to [`Accelerate` launch documentation](../package_reference/cli#accelerate-launch).
+
+<Tip>
+
+    To access other DeepSpeed configurations, such as mixed precision settings, 
+    you need to pass in a `--deepspeed_config_file`, see the [documentation](../usage_guides/deepspeed#deepspeed-config-file).  
+
+    DeepSpeed can be also configured via [`DeepSpeedPlugin`], e.g., `DeepSpeedPlugin.zero_stage` is equivalent of `--zero_stage`, and `DeepSpeedPlugin.hf_ds_config` can be used to pass `--deepeed_config_file.`
+
+</Tip>
+
+<Tip>
+
+    FSDP can be also configured via [`FullyShardedDataParallelPlugin`], e.g., `FullyShardedDataParallelPlugin.sharding_strategy` is equivalent of `--fsdp_sharding_strategy`.
+    
+</Tip>
+
+### Checkpointing
+
+Do note that while FSDP can be configured via `--fsdp_state_dict_type` to save either full / sharded checkpoints.
+
+<Tip>
+
+    For DeepSpeed Zero3, one could pass a `--zero3_save_16bit_model true`, which conveniently consolidates the model to a single rank and saves; this is the FSDP equivalent of `fsdp_state_dict_type: FULL_STATE_DICT`. 
+
+</Tip>
+
+<Tip warning={true}>
+
+    For large models, consolidating the model to a single rank can be very slow.
+
+</Tip>
+
+<Tip>
+
+    For quicker checkpointing, for FSDP use `fsdp_state_dict_type: SHARDED_STATE_DICT`, and for DeepSpeed Zero3 [use the `zero_to_fp32.py` script to post-convert sharded checkpoints](https://www.deepspeed.ai/tutorials/zero/#extracting-weights).
+
+
+</Tip>
+
+### Offloading
+
+FSDP only allows *all-or-nothing* offload (i.e., either offload parameters, gradients, and optimizer, or keep them all in GPU), but DeepSpeed can offload parameters and optimizer differently. Furthermore, DeepSpeed also supports [offloading to NVME](https://www.deepspeed.ai/docs/config-json/#parameter-offloading).
+
+### Prefetching
+
+FSDP allows two prefetching configurations `--fsdp_forward_prefetch` and `--fsdp_backward_prefetch` to improve overlap of comms / computation at a cost of extra memory, see [FSDP documentation](https://pytorch.org/docs/stable/fsdp.html). 
+For DeepSpeed, the prefetching will be turned on when needed, and it turns on depending on certain hyper-params like `stage3_param_persistence_threshold`, `stage3_max_reuse_distance`, etc, [that can be configured for Zero3](https://www.deepspeed.ai/docs/config-json/#parameter-offloading); `accelerate` may set these hyper-params automatically if you don't set those explicitly in the deepspeed config file.
+
+<Tip>
+
+    For FSDP set `fsdp_backward_prefetch: BACKWARD_PRE` for improved throughputs if memory allows.
+
+</Tip>
+
+### Model Loading
+
+While FSDP require an explicit `--fsdp_cpu_ram_efficient_loading true` to activate efficient model loading, `transformers` will activate the similar feature whenever DeepSpeed Zero3 is used.
+
+<Tip>
+
+    For FSDP, whenever setting `--fsdp_cpu_ram_efficient_loading true`, `accelerate` will automatically set `sync_module_states` to true. 
+    For RAM efficient loading the weights will be loaded only in a single rank, and thus requires `sync_module_states` to broadcast weights to other ranks.
+
+</Tip>
+
+### Model
+
+FSDP requires an explicit `--fsdp_auto_wrap_policy` for the algorithm to decide how to schedule the all-gather and reduce-scatter operations. But for DeepSpeed this is transparent to the user.
+
+<Tip>
+
+    For FSDP, simply set `fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP`. With the latest [`transformers`] versions, we try our best to figure out the suitable `fsdp_transformer_layer_cls_to_wrap` for HF transformers models. However, if you get an error regarding it, please specify this.
+
+</Tip>
+
+### Parameters Summoning
+
+FSDP requires an explicit `--fsdp_use_orig_params` flag if using `torch.compile`, see [the pytorch documentation](https://pytorch.org/docs/stable/fsdp.html#module-torch.distributed.fsdp). For DeepSpeed this is transparent to the user.
+
+<Tip>
+
+    For FSDP, when using `torch.compile` please set `fsdp_use_orig_params: True`.
+
+</Tip>
+
+
+## Training
+
+Deepspeed requires explicit `--gradient_accumulation_steps` and `--gradient_clipping` flags. For FSDP this is transparent to the user.
+
+<Tip>
+
+    When using DeepSpeed, set `gradient_accumulation_steps: "auto"` and `gradient_clipping: "auto"` to automatically pick up values set in the [`Accelerator`] or [`TrainingArguments`] (if using `transformers`).
+
+</Tip>
+
+
+## On Differences in Data Precision Handling
+
+To discuss how data precision is handled in both FSDP and Deepspeed, it is instructive to first give an overview of how model parameters are handled in these frameworks. Before the model / optimizer parameters are distributed across GPUs, parameter preparation is involved to first "flatten" them to one-dimensional [`torch.Tensor`](https://pytorch.org/docs/stable/tensors.html#torch-tensor). The implementation of FSDP / DeepSpeed varies in the respect of the `dtype` in which these "flattened" parameters are stored, and there are ramifications with regards to how [`torch.Optimizer`](https://pytorch.org/docs/stable/optim.html#module-torch.optim) allocate their `dtype`s. The table below outlines the processes for both frameworks; the "Local" column indicates the process occurring at a per-gpu level, therefore any memory overheads by upcasting should be understood to be amortized by the number of gpus used.
+
+<Tip>
+
+    As a rule of thumb, for stable training with automatic mixed precision, all the trainable parameters have to be in `torch.float32`.
+
+</Tip>
+
+Process | Local | Framework | Details
+--|--|--|--
+Loading, i.e., [`AutoModel.from_pretrained(..., torch_dtype=torch_dtype)`] |  
+Preparation, i.e., creation of "flat params" | ✅ | FSDP<br>DeepSpeed | created in `torch_dtype`.<br> disregards `torch_dtype`, created in `float32`.
+Optimizer initialization | ✅ | FSDP<br>DeepSpeed  | creates parameters in `torch_dtype`<br> creates parameters in `float32`
+Training Step, i.e, forward, backward, reduction | | FSDP<br>DeepSpeed  | follows [`MixedPrecision`](https://pytorch.org/docs/stable/fsdp.html#torch.distributed.fsdp.MixedPrecision)<br> follows `deepspeed_config_file` mixed precision settings.
+Optimizer (Pre-Step) | ✅ | FSDP<br>DeepSpeed | upcasting (if any) to `torch_dtype`<br>upcasted to `float32`
+Optimizer (Actual Step) | ✅ | FSDP<br>DeepSpeed  | occurs in `torch_dtype` <br> occurs in `float32`.
+
+<Tip warning={true}>
+
+    Therefore when using DeepSpeed a small number of GPUs, be aware of potentially significant memory overheads due to the upcasting during preparation.
+
+</Tip>
+
+<Tip>
+
+    With FSDP, in the absence of mixed precision, it is possible to operate the [`torch.Optimizer`](https://pytorch.org/docs/stable/optim.html#module-torch.optim) in low precision `torch_dtype`, which may be helpful when using small number of GPUs. 
+
+</Tip>
+
+<Tip warning={true}>
+
+    With mixed precision, FSDP and DeepSpeed will upcast in the model preparation step (c.f. table above). But do note that FSDP will then save checkpoints in the upcasted precision; Deepspeed may still save low precision checkpoints if `--zero3_save_16bit_model` is specified.
+
+</Tip>
+
+
+To clarify the above table consider the concrete examples below; the optimizer pre- and actual step combined for brevity. With FSDP it is possible to operate in the two modes shown below, but DeepSpeed can only operate in one.
+
+Framework | Model Loading (`torch_dtype`) | Mixed Precision | Preparation (Local) | Training | Optimizer (Local)
+--|--|--|--|--|--
+FSDP | bf16 | default (none) | bf16 | bf16 | bf16
+FSDP | bf16 | bf16 | fp32 | bf16 | fp32
+DeepSpeed   | bf16 | bf16 | fp32 | bf16 | fp32

docs/source/concept_guides/gradient_synchronization.md

@@ -0,0 +1,184 @@
+<!--Copyright 2022 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
+-->
+
+# Gradient synchronization
+
+PyTorch's distributed module operates by communicating back and forth between all of the GPUs in your system.
+This communication takes time, and ensuring all processes know the states of each other happens at particular triggerpoints
+when using the `ddp` module. 
+
+These triggerpoints are added to the PyTorch model, specifically their `forward()` and `backward()` methods. 
+This happens when the model is wrapped with `DistributedDataParallel`:
+```python
+import torch.nn as nn
+from torch.nn.parallel import DistributedDataParallel
+
+model = nn.Linear(10, 10)
+ddp_model = DistributedDataParallel(model)
+```
+In Accelerate this conversion happens automatically when calling [`~Accelerator.prepare`] and passing in your model.
+
+```diff
++ from accelerate import Accelerator
++ accelerator = Accelerator()
+  import torch.nn as nn
+- from torch.nn.parallel import DistributedDataParallel
+
+  model = nn.Linear(10,10)
++ model = accelerator.prepare(model)
+```
+
+## The slowdown in gradient accumulation
+
+You now understand that PyTorch adds hooks to the `forward` and `backward` method of your PyTorch model when 
+training in a distributed setup. But how does this risk slowing down your code?
+
+In DDP (distributed data parallel), the specific order in which processes are performed and ran are expected
+at specific points and these must also occur at roughly the same time before moving on.
+
+The most direct example is when you update model parameters through
+`optimizer.step()`.
+Without gradient accumulation, all instances of the model need to have updated
+their gradients computed, collated, and updated before moving on to the next
+batch of data.
+When performing gradient accumulation, you accumulate `n` loss gradients and
+skip `optimizer.step()` until `n` batches have been reached. As all training
+processes only need to synchronize by the time `optimizer.step()` is called,
+without any modification to your training step, this needless inter-process
+communication can cause a significant slowdown.
+
+ How can you avoid this overhead?
+
+## Solving the slowdown problem
+
+Since you are skipping model parameter updates when training on these batches, their gradients do not need to be synchronized until the point where `optimizer.step()` is actually called. 
+PyTorch cannot automagically tell when you need to do this, but they do provide a tool to help through the [`no_sync`](https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html#torch.nn.parallel.DistributedDataParallel.no_sync) context manager
+that is added to your model after converting it to DDP.
+
+Under this context manager, PyTorch will skip synchronizing the gradients when
+`.backward()` is called, and the first call to `.backward()` outside this 
+context manager will trigger the synchronization. See an example below:
+```python
+ddp_model, dataloader, optimizer = accelerator.prepare(model, dataloader, optimizer)
+
+for index, batch in enumerate(dataloader):
+    inputs, targets = batch
+    # Trigger gradient synchronization on the last batch
+    if index != (len(dataloader) - 1):
+        with ddp_model.no_sync():
+            # Gradients only accumulate
+            outputs = ddp_model(inputs)
+            loss = loss_func(outputs)
+            accelerator.backward(loss)
+    else:
+        # Gradients finally sync
+        outputs = ddp_model(inputs)
+        loss = loss_func(outputs)
+        accelerator.backward(loss)
+        optimizer.step()
+```
+
+In Accelerate to make this an API that can be called no matter the training device (though it may not do anything if you are not in a distributed system!),
+`ddp_model.no_sync` gets replaced with [`~Accelerator.no_sync`] and operates the same way:
+
+```diff
+  ddp_model, dataloader, optimizer = accelerator.prepare(model, dataloader, optimizer)
+
+  for index, batch in enumerate(dataloader):
+      inputs, targets = batch
+      # Trigger gradient synchronization on the last batch
+      if index != (len(dataloader)-1):
+-         with ddp_model.no_sync():
++         with accelerator.no_sync(model):
+              # Gradients only accumulate
+              outputs = ddp_model(inputs)
+              loss = loss_func(outputs, targets)
+              accelerator.backward(loss)
+      else:
+          # Gradients finally sync
+          outputs = ddp_model(inputs)
+          loss = loss_func(outputs)
+          accelerator.backward(loss)
+          optimizer.step()
+          optimizer.zero_grad()
+```
+
+As you may expect, the [`~Accelerator.accumulate`] function wraps around this conditional check by keeping track of the current batch number, leaving you with the final
+gradient accumulation API:
+
+```python
+ddp_model, dataloader, optimizer = accelerator.prepare(model, dataloader, optimizer)
+
+for batch in dataloader:
+    with accelerator.accumulate(model):
+        optimizer.zero_grad()
+        inputs, targets = batch
+        outputs = model(inputs)
+        loss = loss_function(outputs, targets)
+        accelerator.backward(loss)
+        optimizer.step()
+        optimizer.zero_grad()
+```
+
+As a result, you should either use *`accelerator.accumulate` or `accelerator.no_sync`* when it comes to API choice. 
+
+## Just how much of a slowdown is there, and easy mistakes you can make
+
+To set up a realistic example, consider the following setup:
+
+* Two single-GPU T4 nodes and one node with two GPUs
+* Each GPU is a T4, and are hosted on GCP
+* The script used is a modification of the [NLP Example](https://github.com/muellerzr/timing_experiments/blob/main/baseline.py) script
+* Batch size per GPU is 16, and gradients are accumulated every 4 steps
+
+All scripts are available in [this repository](https://github.com/muellerzr/timing_experiments).
+
+If not careful about gradient synchronization and GPU communication, a *large* amount of time can be wasted 
+from when these GPUs communicate to each other during unnecessary periods.
+
+By how much?
+
+Reference:
+- Baseline: uses no synchronization practices discussed here
+- `no_sync` improperly: `no_sync` only around the `backward` call, not the `forward`
+- `no_sync`: using the `no_sync` pattern properly
+- `accumulate`: using [`~Accelerator.accumulate`] properly
+
+Below are the average seconds per batch iterating over 29 batches of data for each setup on both a single node and on the dual-node setup:
+
+|             | Baseline  | `no_sync` improperly | `no_sync` | `accumulate`| 
+| :---------: | :-------: | :------------------: | :-------: | :---------: |
+| Multi-Node  | 2±0.01s    | 2.13±0.08s | **0.91±0.11s** | **0.91±0.11s** |
+| Single Node | 0.50±0.01s | 0.50±0.01s | **0.41±0.015s** | **0.41±0.015s** |
+
+As you can see, if you are not careful about how you set up your gradient synchronization, you can get upwards of more than a 2x slowdown during training!
+
+If you are worried about making sure everything is done properly, we highly recommend utilizing the [`~Accelerator.accumulate`] function and passing in
+`gradient_accumulation_steps` or `gradient_accumulation_plugin` to the [`Accelerator`] object so Accelerate can handle this for you.
+
+### `no_sync` requires additional GPU memory when using FSDP
+
+Be aware that not syncing gradients can have adverse effects while performing FSDP training. As it has been warned in `torch`, the [`no_sync` context manager for FSDP](https://pytorch.org/docs/stable/fsdp.html#torch.distributed.fsdp.FullyShardedDataParallel.no_sync) will require additional memory.
+
+Therefore in memory intensive situations while using FSDP, we recommend to set `sync_each_batch` to `True` in the [`~utils.GradientAccumulationPlugin`] to disable `no_sync`.
+
+See the example below where we fine-tune Mixtral (47B parameters) on 8 A100-80GB GPUs. We see that even for a modest `gradient_accumulation_steps=2` we quickly go out-of-memory (OOM) if `no_sync` is enabled. Again, this is due to additional memory overheads due to FSDP's `no_sync`. However, if `no_sync` is disabled via `sync_each_batch=True`, then the memory consumption for `gradient_accumulation_steps=16` reverts to that of `gradient_accumulation_steps=1`.
+
+| Model           | `no_sync` (accum=1) | `no_sync` (accum=2) | `no_sync` disabled (accum=16)
+| :-------------: | :-----------------: | :-----------------: | :-----------------: 
+mixtral 8x7B      | 69G                 | OOM                 | 69G
+
+> [!WARNING] 
+> Disabling `no_sync` means there _will be slowdown_ due the extra data syncs, as explained by the earlier sections of this guide.

docs/source/concept_guides/internal_mechanism.md

@@ -0,0 +1,74 @@
+<!--Copyright 2021 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
+-->
+
+# Accelerate's internal mechanisms
+
+Internally, Accelerate works by first analyzing the environment in which the script is launched to determine which
+kind of distributed setup is used, how many different processes there are and which one the current script is in. All
+that information is stored in the [`~AcceleratorState`].
+
+This class is initialized the first time you instantiate an [`~Accelerator`] as well as performing any
+specific initialization your distributed setup needs. Its state is then uniquely shared through all instances of
+[`~state.AcceleratorState`]. (The same can also be done with the [`PartialState`], a more barebones version it inherits)
+
+Then, when calling [`~Accelerator.prepare`], the library:
+
+- wraps your model(s) in the container adapted for the distributed setup,
+- wraps your optimizer(s) in an [`~optimizer.AcceleratedOptimizer`],
+- wraps your scheduler(s) in an [`~scheduler.AcceleratedScheduler`]
+- creates a new version of your dataloader(s) in a [`~data_loader.DataLoaderShard`] or [`~data_loader.DataLoaderDispatcher`]
+
+While the model(s), optimizer(s), and scheduler(s) are just put in simple wrappers, the dataloader(s) are re-created. This is mostly
+because PyTorch does not let the user change the `batch_sampler` of a dataloader once it's been created and the
+library handles the sharding of your data between processes by changing that `batch_sampler` to yield every other
+`num_processes` batches (if enabled).
+
+The [`~data_loader.DataLoaderShard`] subclasses `DataLoader` to add the following functionality:
+
+- it synchronizes the appropriate random number generator of all processes at each new iteration, to ensure any
+  randomization (like shuffling) is done the exact same way across processes.
+- it puts the batches on the proper device before yielding them (unless you have opted out of
+  `device_placement=True`).
+  
+The [`~data_loader.DataLoaderDispatcher`] subclasses differs from the [`~data_loader.DataLoaderShard`] in that when iterating through the `DataLoader`, the data is all starting from process 0 and *then* split and sent off to each process rather than it happening at the dataset level.
+
+The random number generator synchronization will by default synchronize:
+
+- the `generator` attribute of a given sampler (like the PyTorch `RandomSampler`) for PyTorch >= 1.6
+- the main random number generator in PyTorch <=1.5.1
+
+You can choose which random number generator(s) to synchronize with the `rng_types` argument of the main
+[`Accelerator`]. In PyTorch >= 1.6, it is recommended to rely on a local `generator` to avoid
+setting the same seed in the main random number generator in all processes.
+
+<Tip warning={true}>
+
+    Synchronization of the main torch (or CUDA or XLA) random number generator will affect any other potential random
+    artifacts you could have in your dataset (like random data augmentation) in the sense that all processes will get
+    the same random numbers from the torch random modules (so will apply the same random data augmentation if it's
+    controlled by torch).
+
+</Tip>
+
+<Tip>
+
+    The randomization part of your custom sampler, batch sampler or iterable dataset should be done using a local
+    `torch.Generator` object (in PyTorch >= 1.6), see the traditional `RandomSampler`, as an example.
+
+</Tip>
+
+If you have [`torchdata>=0.8.0`](https://github.com/pytorch/data/tree/main) installed, and you have passed `use_stateful_dataloader=True` into your [`~utils.DataLoaderConfiguration`], these classes will directly inherit from `StatefulDataLoader` instead, and maintain a `state_dict`.
+
+For more details about the internals, see the [Internals page](../package_reference/torch_wrappers).

docs/source/concept_guides/low_precision_training.md

@@ -0,0 +1,74 @@
+<!--Copyright 2023 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
+-->
+
+# Low precision training methods
+
+The release of new kinds of hardware led to the emergence of new training paradigms that better utilize them. Currently, this is in the form of training
+in 8-bit precision using packages such as [TransformersEngine](https://github.com/NVIDIA/TransformerEngine) (TE) or [MS-AMP](https://github.com/Azure/MS-AMP/tree/main).
+
+For an introduction to the topics discussed today, we recommend reviewing the [low-precision usage guide](../usage_guides/low_precision_training) as this documentation will reference it regularly. 
+
+## A Quick Chart
+
+Below is a quick chart from the MS-AMP documentation showing the different bit-precisions for each solution during training:
+
+Optimization Level | Computation(GEMM) | Comm | Weight | Master Weight | Weight Gradient | Optimizer States
+-- | -- | -- | -- | -- | -- | --
+FP16 AMP | FP16 | FP32 | FP32 | N/A | FP32 | FP32+FP32
+Nvidia TE | FP8 | FP32 | FP32 | N/A | FP32 | FP32+FP32
+MS-AMP O1 | FP8 | FP8 | FP16 | N/A | FP8 | FP32+FP32
+MS-AMP O2 | FP8 | FP8 | FP16 | N/A | FP8 | FP8+FP16
+MS-AMP O3 | FP8 | FP8 | FP8 | FP16 | FP8 | FP8+FP16
+
+## `TransformersEngine`
+
+`TransformersEngine` is the first solution to trying to train in 8-bit floating point. It works by using drop-in replacement layers for certain ones in a model that utilizes their FP8-engine to reduce the number of bits (such as 32 to 8) without degrading the final accuracy of the model. 
+
+Specifically, Accelerate will find and replace the following layers with `TransformersEngine` versions:
+
+* `nn.LayerNorm` for `te.LayerNorm`
+* `nn.Linear` for `te.Linear`
+
+As a result we wind up with a model that has most of its layers in BF16, while some layers are in FP8 reducing some of the memory. 
+
+Anecdotally, we have noticed that performance gains don't really start showing when using `TransformerEngine` until a large majority of the layers
+in the model are made up of those two layers to replace. As a result, only larger models have shown performance improvements when the number of parameters is around and upwards of a few billion. 
+
+The `TransformerEngine` can receive many different arguments that customize how it performs FP8 calculations and what they do. A full list of the arguments is available below:
+
+* `margin`: The margin to use for the gradient scaling.
+* `interval`: The interval to use for how often the scaling factor is recomputed.
+* `fp8_format``: The format to use for the FP8 recipe. Must be one of `HYBRID` or `E4M3`. (Generally `HYBRID` for training, `E4M3` for evaluation)
+* `amax_history_len`: The length of the history to use for the scaling factor computation
+* `amax_compute_algo`: The algorithm to use for the scaling factor computation. Must be one of `max` or `most_recent`.
+* `override_linear_precision`: Whether or not to execute `fprop`, `dgrad`, and `wgrad` GEMMS in higher precision.
+
+You can customize each of these as part of [`utils.FP8RecipeKwargs`] to help optimize performance of your models.
+
+If we notice in the chart mentioned earlier, TE simply casts the computation layers into FP8, while everything else is in FP32. As a result this winds up utilizing the most memory but does so with the benefit of guaranteeing the least amount of loss in end accuracy during training. 
+
+## `MS-AMP`
+
+MS-AMP takes a different approach to `TransformersEngine` by providing three different optimization levels to convert more operations in FP8 or FP16.
+
+* The base optimization level (`O1`), passes communications of the weights (such as in DDP) in FP8, stores the weights of the model in FP16, and leaves the optimizer states in FP32. The main benefit of this optimization level is that we can reduce the communication bandwidth by essentially half. Additionally, more GPU memory is saved due to 1/2 of everything being cast in FP8, and the weights being cast to FP16. Notably, both the optimizer states remain in FP32.
+
+* The second optimization level (`O2`) improves upon this by also reducing the precision of the optimizer states. One is in FP8 while the other is in FP16. Generally it's been shown that this will only provide a net-gain of no degraded end accuracy, increased training speed, and reduced memory as now every state is either in FP16 or FP8. 
+
+* Finally, MS-AMP has a third optimization level (`O3`) which helps during DDP scenarios such as DeepSpeed. The weights of the model in memory are fully cast to FP8, and the master weights are now stored in FP16. This fully reduces memory by the highest factor as now not only is almost everything in FP8, only two states are left in FP16. Currently, only DeepSpeed versions up through 0.9.2 are supported, so this capability is not included in the Accelerate integration
+
+## Combining the two
+
+More experiments need to be performed but it's been noted that combining both MS-AMP and TransformersEngine can lead to the highest throughput by relying on NVIDIA's optimized FP8 operators and utilizing how MS-AMP reduces the memory overhead.

docs/source/concept_guides/performance.md

@@ -0,0 +1,103 @@
+<!--Copyright 2022 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
+-->
+
+# Comparing performance across distributed setups
+
+Evaluating and comparing the performance from different setups can be quite tricky if you don't know what to look for.
+For example, you cannot run the same script with the same batch size across TPU, multi-GPU, and single-GPU with Accelerate 
+and expect your results to line up. 
+
+But why?
+
+There are three reasons for this that this tutorial will cover: 
+
+1. **Setting the right seeds**
+2. **Observed Batch Sizes**
+3. **Learning Rates**
+
+## Setting the Seed 
+
+While this issue has not come up as much, make sure to use [`utils.set_seed`] to fully set the seed in all distributed cases so training will be reproducible:
+
+```python
+from accelerate.utils import set_seed
+
+set_seed(42)
+```
+
+Why is this important? Under the hood this will set **5** different seed settings:
+
+```python
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed) # or torch.xpu.manual_seed_all, etc
+    # ^^ safe to call this function even if cuda is not available
+    if is_torch_xla_available():
+        xm.set_rng_state(seed)
+```
+
+The random state, numpy's state, torch, torch's device state, and if TPUs are available torch_xla's cuda state.
+
+## Observed Batch Sizes 
+
+When training with Accelerate, the batch size passed to the dataloader is the **batch size per GPU**. What this entails is 
+a batch size of 64 on two GPUs is truly a batch size of 128. As a result, when testing on a single GPU this needs to be accounted for,
+as well as similarly for TPUs. 
+
+The below table can be used as a quick reference to try out different batch sizes:
+
+<Tip>
+
+In this example, there are two GPUs for "Multi-GPU" and a TPU pod with 8 workers
+
+</Tip>
+
+| Single GPU Batch Size | Multi-GPU Equivalent Batch Size | TPU Equivalent Batch Size |
+|-----------------------|---------------------------------|---------------------------|
+| 256                   | 128                             | 32                        |
+| 128                   | 64                              | 16                        |
+| 64                    | 32                              | 8                         |
+| 32                    | 16                              | 4                         |
+
+## Learning Rates 
+
+As noted in multiple sources[[1](https://aws.amazon.com/blogs/machine-learning/scalable-multi-node-deep-learning-training-using-gpus-in-the-aws-cloud/)][[2](https://docs.nvidia.com/clara/clara-train-sdk/pt/model.html#classification-models-multi-gpu-training)], the learning rate should be scaled *linearly* based on the number of devices present. The below 
+snippet shows doing so with Accelerate:
+
+<Tip>
+
+Since users can have their own learning rate schedulers defined, we leave this up to the user to decide if they wish to scale their 
+learning rate or not.
+ 
+</Tip>
+
+```python
+learning_rate = 1e-3
+accelerator = Accelerator()
+learning_rate *= accelerator.num_processes
+
+optimizer = AdamW(params=model.parameters(), lr=learning_rate)
+```
+
+You will also find that `accelerate` will step the learning rate based on the number of processes being trained on. This is because 
+of the observed batch size noted earlier. So in the case of 2 GPUs, the learning rate will be stepped twice as often as a single GPU
+to account for the batch size being twice as large (if no changes to the batch size on the single GPU instance are made).
+
+## Gradient Accumulation and Mixed Precision
+
+When using gradient accumulation and mixed precision, due to how gradient averaging works (accumulation) and the precision loss (mixed precision), 
+some degradation in performance is expected. This will be explicitly seen when comparing the batch-wise loss between different compute 
+setups. However, the overall loss, metric, and general performance at the end of training should be _roughly_ the same.

docs/source/concept_guides/training_tpu.md

@@ -0,0 +1,167 @@
+<!--Copyright 2022 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
+-->
+
+# Training on TPUs
+
+Training on TPUs can be slightly different from training on multi-gpu, even with Accelerate. This guide aims to show you 
+where you should be careful and why, as well as the best practices in general.
+
+## Training in a Notebook
+
+The main carepoint when training on TPUs comes from the [`notebook_launcher`]. As mentioned in the [notebook tutorial](../usage_guides/notebook), you need to 
+restructure your training code into a function that can get passed to the [`notebook_launcher`] function and be careful about not declaring any tensors on the GPU.
+
+While on a TPU that last part is not as important, a critical part to understand is that when you launch code from a notebook you do so through a process called **forking**. 
+When launching from the command-line, you perform **spawning**, where a python process is not currently running and you *spawn* a new process in. Since your Jupyter notebook is already 
+utilizing a python process, you need to *fork* a new process from it to launch your code. 
+
+Where this becomes important is in regard to declaring your model. On forked TPU processes, it is recommended that you instantiate your model *once* and pass this into your 
+training function. This is different than training on GPUs where you create `n` models that have their gradients synced and back-propagated at certain moments. Instead, one 
+model instance is shared between all the nodes and it is passed back and forth. This is important especially when training on low-resource TPUs such as those provided in Kaggle kernels or
+on Google Colaboratory. 
+
+Below is an example of a training function passed to the [`notebook_launcher`] if training on CPUs or GPUs:
+
+<Tip>
+
+    This code snippet is based off the one from the `simple_nlp_example` notebook found [here](https://github.com/huggingface/notebooks/blob/main/examples/accelerate_examples/simple_nlp_example.ipynb) with slight 
+    modifications for the sake of simplicity
+
+</Tip>
+
+```python
+def training_function():
+    # Initialize accelerator
+    accelerator = Accelerator()
+    model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=2)
+    train_dataloader, eval_dataloader = create_dataloaders(
+        train_batch_size=hyperparameters["train_batch_size"], eval_batch_size=hyperparameters["eval_batch_size"]
+    )
+
+    # Instantiate optimizer
+    optimizer = AdamW(params=model.parameters(), lr=hyperparameters["learning_rate"])
+
+    # Prepare everything
+    # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
+    # prepare method.
+    model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(
+        model, optimizer, train_dataloader, eval_dataloader
+    )
+
+    num_epochs = hyperparameters["num_epochs"]
+    # Now we train the model
+    for epoch in range(num_epochs):
+        model.train()
+        for step, batch in enumerate(train_dataloader):
+            outputs = model(**batch)
+            loss = outputs.loss
+            accelerator.backward(loss)
+
+            optimizer.step()
+            optimizer.zero_grad()
+```
+
+```python
+from accelerate import notebook_launcher
+
+notebook_launcher(training_function)
+```
+
+<Tip>
+
+    The `notebook_launcher` will default to 8 processes if Accelerate has been configured for a TPU
+
+</Tip>
+
+If you use this example and declare the model *inside* the training loop, then on a low-resource system you will potentially see an error 
+like:
+
+```
+ProcessExitedException: process 0 terminated with signal SIGSEGV
+```
+
+This error is *extremely* cryptic but the basic explanation is you ran out of system RAM. You can avoid this entirely by reconfiguring the training function to 
+accept a single `model` argument, and declare it in an outside cell:
+
+```python
+# In another Jupyter cell
+model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=2)
+```
+
+```diff
++ def training_function(model):
+      # Initialize accelerator
+      accelerator = Accelerator()
+-     model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=2)
+      train_dataloader, eval_dataloader = create_dataloaders(
+          train_batch_size=hyperparameters["train_batch_size"], eval_batch_size=hyperparameters["eval_batch_size"]
+      )
+  ...
+```
+
+And finally calling the training function with:
+
+```diff
+  from accelerate import notebook_launcher
+- notebook_launcher(training_function)
++ notebook_launcher(training_function, (model,))
+```
+
+<Tip>
+
+    The above workaround is only needed when launching a TPU instance from a Jupyter Notebook on a low-resource server such as Google Colaboratory or Kaggle. If 
+    using a script or launching on a much beefier server declaring the model beforehand is not needed.
+
+</Tip>
+
+## Mixed Precision and Global Variables 
+
+As mentioned in the [mixed precision tutorial](../usage_guides/mixed_precision), Accelerate supports fp16 and bf16, both of which can be used on TPUs.
+That being said, ideally `bf16` should be utilized as it is extremely efficient to use.
+
+There are two "layers" when using `bf16` and Accelerate on TPUs, at the base level and at the operation level. 
+
+At the base level, this is enabled when passing `mixed_precision="bf16"` to `Accelerator`, such as:
+```python
+accelerator = Accelerator(mixed_precision="bf16")
+```
+By default, this will cast `torch.float` and `torch.double` to `bfloat16` on TPUs. 
+The specific configuration being set is an environmental variable of `XLA_USE_BF16` is set to `1`.
+
+There is a further configuration you can perform which is setting the `XLA_DOWNCAST_BF16` environmental variable. If set to `1`, then 
+`torch.float` is `bfloat16` and `torch.double` is `float32`.
+
+This is performed in the `Accelerator` object when passing `downcast_bf16=True`:
+```python
+accelerator = Accelerator(mixed_precision="bf16", downcast_bf16=True)
+```
+
+Using downcasting instead of bf16 everywhere is good for when you are trying to calculate metrics, log values, and more where raw bf16 tensors would be unusable. 
+
+## Training Times on TPUs
+
+As you launch your script, you may notice that training seems exceptionally slow at first. This is because TPUs
+first run through a few batches of data to see how much memory to allocate before finally utilizing this configured 
+memory allocation extremely efficiently. 
+
+If you notice that your evaluation code to calculate the metrics of your model takes longer due to a larger batch size being used, 
+it is recommended to keep the batch size the same as the training data if it is too slow. Otherwise the memory will reallocate to this 
+new batch size after the first few iterations. 
+
+<Tip>
+
+    Just because the memory is allocated does not mean it will be used or that the batch size will increase when going back to your training dataloader.
+
+</Tip>

docs/source/conf.py

@@ -1,210 +0,0 @@
-# -*- coding: utf-8 -*-
-#
-# Configuration file for the Sphinx documentation builder.
-#
-# This file does only contain a selection of the most common options. For a
-# full list see the documentation:
-# http://www.sphinx-doc.org/en/master/config
-
-# -- Path setup --------------------------------------------------------------
-
-# If extensions (or modules to document with autodoc) are in another directory,
-# add these directories to sys.path here. If the directory is relative to the
-# documentation root, use os.path.abspath to make it absolute, like shown here.
-#
-import os
-import sys
-
-sys.path.insert(0, os.path.abspath("../../src"))
-
-
-# -- Project information -----------------------------------------------------
-
-project = "accelerate"
-copyright = "2020, The Hugging Face Team, Licenced under the Apache License, Version 2.0"
-author = "huggingface"
-
-# The short X.Y version
-version = "0.2.0"
-
-# -- General configuration ---------------------------------------------------
-
-# If your documentation needs a minimal Sphinx version, state it here.
-#
-# needs_sphinx = '1.0'
-
-# Add any Sphinx extension module names here, as strings. They can be
-# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
-# ones.
-extensions = [
-    "sphinx.ext.autodoc",
-    "sphinx.ext.extlinks",
-    "sphinx.ext.coverage",
-    "sphinx.ext.napoleon",
-    "recommonmark",
-    "sphinx.ext.viewcode",
-    "sphinx_markdown_tables",
-    "sphinx_copybutton",
-    "sphinxext.opengraph",
-]
-
-# Add any paths that contain templates here, relative to this directory.
-templates_path = ["_templates"]
-
-# The suffix(es) of source filenames.
-# You can specify multiple suffix as a list of string:
-#
-source_suffix = [".rst", ".md"]
-# source_suffix = '.rst'
-
-# The master toctree document.
-master_doc = "index"
-
-# The language for content autogenerated by Sphinx. Refer to documentation
-# for a list of supported languages.
-#
-# This is also used if you do content translation via gettext catalogs.
-# Usually you set "language" from the command line for these cases.
-language = None
-
-# List of patterns, relative to source directory, that match files and
-# directories to ignore when looking for source files.
-# This pattern also affects html_static_path and html_extra_path.
-exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"]
-
-# The name of the Pygments (syntax highlighting) style to use.
-pygments_style = None
-
-# Remove the prompt when copying examples
-copybutton_prompt_text = r">>> |\.\.\. "
-copybutton_prompt_is_regexp = True
-
-# -- Options for HTML output -------------------------------------------------
-
-# The theme to use for HTML and HTML Help pages.  See the documentation for
-# a list of builtin themes.
-#
-html_theme = "sphinx_rtd_theme"
-
-# Theme options are theme-specific and customize the look and feel of a theme
-# further.  For a list of options available for each theme, see the
-# documentation.
-#
-html_theme_options = {"analytics_id": "UA-83738774-2"}
-
-#  Configuration for OpenGraph and Twitter Card Tags.
-# These are responsible for creating nice shareable social images https://ahrefs.com/blog/open-graph-meta-tags/
-# https://ogp.me/#type_website
-ogp_image = "https://huggingface.co/front/thumbnails/docs/accelerate.png"
-ogp_description = "Run your raw PyTorch training script on any kind of device. 🤗 Accelerate provides an easy API to make your scripts run with mixed precision and on any kind of distributed setting (multi-GPUs, TPUs etc.)"
-ogp_description_length = 160
-
-ogp_custom_meta_tags = [
-    f'<meta name="twitter:image" content="{ogp_image}">',
-    f'<meta name="twitter:description" content="{ogp_description}">',
-]
-
-
-# Add any paths that contain custom static files (such as style sheets) here,
-# relative to this directory. They are copied after the builtin static files,
-# so a file named "default.css" will overwrite the builtin "default.css".
-html_static_path = ["_static"]
-
-# Custom sidebar templates, must be a dictionary that maps document names
-# to template names.
-#
-# The default sidebars (for documents that don't match any pattern) are
-# defined by theme itself.  Builtin themes are using these templates by
-# default: ``['localtoc.html', 'relations.html', 'sourcelink.html',
-# 'searchbox.html']``.
-#
-# html_sidebars = {}
-
-# This must be the name of an image file (path relative to the configuration
-# directory) that is the favicon of the docs. Modern browsers use this as
-# the icon for tabs, windows and bookmarks. It should be a Windows-style
-# icon file (.ico).
-html_favicon = "favicon.ico"
-
-
-# -- Options for HTMLHelp output ---------------------------------------------
-
-# Output file base name for HTML help builder.
-htmlhelp_basename = "acceleratedoc"
-
-
-# -- Options for LaTeX output ------------------------------------------------
-
-latex_elements = {
-    # The paper size ('letterpaper' or 'a4paper').
-    #
-    # 'papersize': 'letterpaper',
-    # The font size ('10pt', '11pt' or '12pt').
-    #
-    # 'pointsize': '10pt',
-    # Additional stuff for the LaTeX preamble.
-    #
-    # 'preamble': '',
-    # Latex figure (float) alignment
-    #
-    # 'figure_align': 'htbp',
-}
-
-# Grouping the document tree into LaTeX files. List of tuples
-# (source start file, target name, title,
-#  author, documentclass [howto, manual, or own class]).
-latex_documents = [
-    (master_doc, "accelerate.tex", "accelerate Documentation", "huggingface", "manual"),
-]
-
-
-# -- Options for manual page output ------------------------------------------
-
-# One entry per manual page. List of tuples
-# (source start file, name, description, authors, manual section).
-man_pages = [(master_doc, "accelerate", "accelerate Documentation", [author], 1)]
-
-
-# -- Options for Texinfo output ----------------------------------------------
-
-# Grouping the document tree into Texinfo files. List of tuples
-# (source start file, target name, title, author,
-#  dir menu entry, description, category)
-texinfo_documents = [
-    (
-        master_doc,
-        "accelerate",
-        "accelerate Documentation",
-        author,
-        "accelerate",
-        "One line description of project.",
-        "Miscellaneous",
-    ),
-]
-
-
-# -- Options for Epub output -------------------------------------------------
-
-# Bibliographic Dublin Core info.
-epub_title = project
-
-# The unique identifier of the text. This can be a ISBN number
-# or the project homepage.
-#
-# epub_identifier = ''
-
-# A unique identification for the text.
-#
-# epub_uid = ''
-
-# A list of files that should not be packed into the epub file.
-epub_exclude_files = ["search.html"]
-
-
-def setup(app):
-    app.add_css_file("css/huggingface.css")
-    app.add_css_file("css/code-snippets.css")
-    app.add_js_file("js/custom.js")
-
-
-# -- Extension configuration -------------------------------------------------

docs/source/index.md

@@ -0,0 +1,74 @@
+<!--Copyright 2022 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
+-->
+
+# Accelerate
+
+Accelerate is a library that enables the same PyTorch code to be run across any distributed configuration by adding just four lines of code! In short, training and inference at scale made simple, efficient and adaptable.
+
+```diff
++ from accelerate import Accelerator
++ accelerator = Accelerator()
+
++ model, optimizer, training_dataloader, scheduler = accelerator.prepare(
++     model, optimizer, training_dataloader, scheduler
++ )
+
+  for batch in training_dataloader:
+      optimizer.zero_grad()
+      inputs, targets = batch
+      inputs = inputs.to(device)
+      targets = targets.to(device)
+      outputs = model(inputs)
+      loss = loss_function(outputs, targets)
++     accelerator.backward(loss)
+      optimizer.step()
+      scheduler.step()
+```
+
+Built on `torch_xla` and `torch.distributed`, Accelerate takes care of the heavy lifting, so you don't have to write any custom code to adapt to these platforms.
+Convert existing codebases to utilize [DeepSpeed](usage_guides/deepspeed), perform [fully sharded data parallelism](usage_guides/fsdp), and have automatic support for mixed-precision training! 
+
+<Tip> 
+
+  To get a better idea of this process, make sure to check out the [Tutorials](basic_tutorials/overview)! 
+
+</Tip>
+
+
+This code can then be launched on any system through Accelerate's CLI interface:
+```bash
+accelerate launch {my_script.py}
+```
+
+<div class="mt-10">
+  <div class="w-full flex flex-col space-y-4 md:space-y-0 md:grid md:grid-cols-2 md:gap-y-4 md:gap-x-5">
+    <a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="./basic_tutorials/overview"
+      ><div class="w-full text-center bg-gradient-to-br from-blue-400 to-blue-500 rounded-lg py-1.5 font-semibold mb-5 text-white text-lg leading-relaxed">Tutorials</div>
+      <p class="text-gray-700">Learn the basics and become familiar with using Accelerate. Start here if you are using Accelerate for the first time!</p>
+    </a>
+    <a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="./usage_guides/explore"
+      ><div class="w-full text-center bg-gradient-to-br from-indigo-400 to-indigo-500 rounded-lg py-1.5 font-semibold mb-5 text-white text-lg leading-relaxed">How-to guides</div>
+      <p class="text-gray-700">Practical guides to help you achieve a specific goal. Take a look at these guides to learn how to use Accelerate to solve real-world problems.</p>
+    </a>
+    <a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="./concept_guides/gradient_synchronization"
+      ><div class="w-full text-center bg-gradient-to-br from-pink-400 to-pink-500 rounded-lg py-1.5 font-semibold mb-5 text-white text-lg leading-relaxed">Conceptual guides</div>
+      <p class="text-gray-700">High-level explanations for building a better understanding of important topics such as avoiding subtle nuances and pitfalls in distributed training and DeepSpeed.</p>
+   </a>
+    <a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="./package_reference/accelerator"
+      ><div class="w-full text-center bg-gradient-to-br from-purple-400 to-purple-500 rounded-lg py-1.5 font-semibold mb-5 text-white text-lg leading-relaxed">Reference</div>
+      <p class="text-gray-700">Technical descriptions of how Accelerate classes and methods work.</p>
+    </a>
+  </div>
+</div>

docs/source/index.rst

@@ -1,159 +0,0 @@
-.. 
-    Copyright 2021 The HuggingFace Team. All rights reserved.
-
-    Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
-    the License. You may obtain a copy of the License at
-
-        http://www.apache.org/licenses/LICENSE-2.0
-
-    Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
-    an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
-    specific language governing permissions and limitations under the License.
-
-Accelerate
-=======================================================================================================================
-
-Run your *raw* PyTorch training script on any kind of device
-
-Features
------------------------------------------------------------------------------------------------------------------------
-
-- 🤗 Accelerate provides an easy API to make your scripts run with mixed precision and on any kind of distributed
-  setting (multi-GPUs, TPUs etc.) while still letting you write your own training loop. The same code can then runs
-  seamlessly on your local machine for debugging or your training environment.
-
-- 🤗 Accelerate also provides a CLI tool that allows you to quickly configure and test your training environment then
-  launch the scripts.
-
-
-Easy to integrate
------------------------------------------------------------------------------------------------------------------------
-
-A traditional training loop in PyTorch looks like this:
-
-.. code-block:: python
-
-    my_model.to(device)
-
-    for batch in my_training_dataloader:
-        my_optimizer.zero_grad()
-        inputs, targets = batch
-        inputs = inputs.to(device)
-        targets = targets.to(device)
-        outputs = my_model(inputs)
-        loss = my_loss_function(outputs, targets)
-        loss.backward()
-        my_optimizer.step()
-
-Changing it to work with accelerate is really easy and only adds a few lines of code:
-
-.. code-block:: diff
-
-    + from accelerate import Accelerator
-
-    + accelerator = Accelerator()
-      # Use the device given by the `accelerator` object.
-    + device = accelerator.device
-      my_model.to(device)
-      # Pass every important object (model, optimizer, dataloader) to `accelerator.prepare`
-    + my_model, my_optimizer, my_training_dataloader = accelerate.prepare(
-    +     my_model, my_optimizer, my_training_dataloader
-    + )
-
-      for batch in my_training_dataloader:
-          my_optimizer.zero_grad()
-          inputs, targets = batch
-          inputs = inputs.to(device)
-          targets = targets.to(device)
-          outputs = my_model(inputs)
-          loss = my_loss_function(outputs, targets)
-          # Just a small change for the backward instruction
-    -     loss.backward()
-    +     accelerate.backward(loss)
-          my_optimizer.step()
-
-and with this, your script can now run in a distributed environment (multi-GPU, TPU).
-
-You can even simplify your script a bit by letting 🤗 Accelerate handle the device placement for you (which is safer,
-especially for TPU training):
-
-.. code-block:: diff
-
-    + from accelerate import Accelerator
-
-    + accelerator = Accelerator()
-    - my_model.to(device)
-      # Pass every important object (model, optimizer, dataloader) to `accelerator.prepare`
-    + my_model, my_optimizer, my_training_dataloader = accelerate.prepare(
-    +     my_model, my_optimizer, my_training_dataloader
-    + )
-
-      for batch in my_training_dataloader:
-          my_optimizer.zero_grad()
-          inputs, targets = batch
-    -     inputs = inputs.to(device)
-    -     targets = targets.to(device)
-          outputs = my_model(inputs)
-          loss = my_loss_function(outputs, targets)
-          # Just a small change for the backward instruction
-    -     loss.backward()
-    +     accelerate.backward(loss)
-          my_optimizer.step()
-
-
-Script launcher
------------------------------------------------------------------------------------------------------------------------
-
-No need to remember how to use ``torch.distributed.launch`` or to write a specific launcher for TPU training! 🤗
-Accelerate comes with a CLI tool that will make your life easier when launching distributed scripts.
-
-On your machine(s) just run:
-
-.. code-block:: bash
-
-    accelerate config
-
-and answer the questions asked. This will generate a config file that will be used automatically to properly set the
-default options when doing
-
-.. code-block:: bash
-
-    accelerate launch my_script.py --args_to_my_script
-
-For instance, here is how you would run the NLP example (from the root of the repo):
-
-.. code-block:: bash
-
-    accelerate launch examples/nlp_example.py
-
-
-Supported integrations
------------------------------------------------------------------------------------------------------------------------
-
-- CPU only
-- single GPU
-- multi-GPU on one node (machine)
-- multi-GPU on several nodes (machines)
-- TPU
-- FP16 with native AMP (apex on the roadmap)
-
-.. toctree::
-    :maxdepth: 2
-    :caption: Get started
-
-    quicktour
-    installation
-
-.. toctree::
-    :maxdepth: 2
-    :caption: Guides
-
-    sagemaker
-
-.. toctree::
-    :maxdepth: 2
-    :caption: API reference
-
-    accelerator
-    kwargs
-    internal