v0.13.1: Release

* Fix num_processes is not defined

* Also reorganize questions

Co-authored-by: Sylvain Gugger <Sylvain.gugger@gmail.com>

.github/ISSUE_TEMPLATE/bug-report.yml

@@ -0,0 +1,58 @@
+name: "\U0001F41B Bug Report"
+description: Submit a bug report to help us improve Accelerate
+body:
+  - 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."
+      render: Shell

.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,64 @@
+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@v3
+      - id: step1
+        run: echo "::set-output name=version::$(python setup.py --version)"
+
+  version-cpu:
+    name: "Latest Accelerate CPU [version]"
+    runs-on: ubuntu-latest
+    needs: get-version
+    steps:
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v1
+      - name: Check out code
+        uses: actions/checkout@v2
+      - name: Login to DockerHub
+        uses: docker/login-action@v1
+        with:
+          username: ${{ secrets.DOCKERHUB_USERNAME }}
+          password: ${{ secrets.DOCKERHUB_PASSWORD }}
+
+      - name: Build and Push CPU
+        uses: docker/build-push-action@v2
+        with:
+          context: ./docker/accelerate-cpu
+          push: true
+          tags: huggingface/accelerate-cpu:${{needs.get-version.outputs.version}}
+
+  version-cuda:
+    name: "Latest Accelerate GPU [version]"
+    runs-on: ubuntu-latest
+    needs: get-version
+    steps:
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v1
+      - name: Check out code
+        uses: actions/checkout@v2
+      - name: Login to DockerHub
+        uses: docker/login-action@v1
+        with:
+          username: ${{ secrets.DOCKERHUB_USERNAME }}
+          password: ${{ secrets.DOCKERHUB_PASSWORD }}
+
+      - name: Build and Push GPU
+        uses: docker/build-push-action@v2
+        with:
+          context: ./docker/accelerate-gpu
+          push: true
+          tags: huggingface/accelerate-gpu:${{needs.get-version.outputs.version}}

.github/workflows/build_and_run_tests.yml

@@ -0,0 +1,45 @@
+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@v3
+        with: 
+          fetch-depth: "2"
+      
+      - name: Get changed files
+        id: changed-files
+        uses: tj-actions/changed-files@v22.2
+      
+      - 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 ::set-output name=changed::"1"
+            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

.github/workflows/build_docker_images.yml

@@ -0,0 +1,54 @@
+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: ubuntu-latest
+    steps:
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v1
+      - name: Check out code
+        uses: actions/checkout@v2
+      - name: Login to DockerHub
+        uses: docker/login-action@v1
+        with:
+          username: ${{ secrets.DOCKERHUB_USERNAME }}
+          password: ${{ secrets.DOCKERHUB_PASSWORD }}
+
+      - name: Build and Push CPU
+        uses: docker/build-push-action@v2
+        with:
+          context: ./docker/accelerate-cpu
+          push: true
+          tags: huggingface/accelerate-cpu
+
+  latest-cuda:
+    name: "Latest Accelerate GPU [dev]"
+    runs-on: ubuntu-latest
+    steps:
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v1
+      - name: Check out code
+        uses: actions/checkout@v2
+      - name: Login to DockerHub
+        uses: docker/login-action@v1
+        with:
+          username: ${{ secrets.DOCKERHUB_USERNAME }}
+          password: ${{ secrets.DOCKERHUB_PASSWORD }}
+
+      - name: Build and Push GPU
+        uses: docker/build-push-action@v2
+        with:
+          context: ./docker/accelerate-gpu
+          push: true
+          tags: huggingface/accelerate-gpu

.github/workflows/build_documentation.yml

@@ -0,0 +1,17 @@
+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
+    secrets:
+      token: ${{ secrets.HUGGINGFACE_PUSH }}

.github/workflows/build_pr_documentation.yml

@@ -0,0 +1,16 @@
+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

.github/workflows/delete_doc_comment.yml

@@ -0,0 +1,13 @@
+name: Delete dev documentation
+
+on:
+  pull_request:
+    types: [ closed ]
+
+
+jobs:
+  delete:
+    uses: huggingface/doc-builder/.github/workflows/delete_doc_comment.yml@main
+    with:
+      pr_number: ${{ github.event.number }}
+      package: accelerate

.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/nightly.yml

@@ -0,0 +1,88 @@
+name: Self-hosted runner with slow tests (scheduled)
+
+on:
+  workflow_dispatch:
+  schedule:
+    - cron: "0 2 * * *"
+
+env:
+  RUN_SLOW: "yes"
+  IS_GITHUB_CI: "1"
+
+jobs:
+  run_all_tests_single_gpu:
+    runs-on: [self-hosted, docker-gpu, multi-gpu]
+    env:
+      CUDA_VISIBLE_DEVICES: "0"
+    container:
+      image: huggingface/accelerate-gpu:latest
+      options: --gpus all --shm-size "16gb"
+    defaults:
+      run:
+        working-directory: accelerate/
+        shell: bash
+    steps:
+      - name: Update clone & pip install
+        run: |
+          source activate accelerate
+          git config --global --add safe.directory '*'
+          git fetch && git checkout ${{ github.sha }} 
+          pip install -e . --no-deps
+          pip install pytest-reportlog
+
+      - name: Run test on GPUs
+        run: |
+          source activate accelerate
+          make test
+      - name: Run examples on GPUs
+        run: |
+          source activate accelerate
+          pip uninstall comet_ml -y
+          make test_examples
+          
+      - name: Generate Report
+        if: always()
+        run: |
+          python utils/log_reports.py >> $GITHUB_STEP_SUMMARY
+
+  run_all_tests_multi_gpu:
+    runs-on: [self-hosted, docker-gpu, multi-gpu]
+    env:
+      CUDA_VISIBLE_DEVICES: "0,1"
+    container:
+      image: huggingface/accelerate-gpu:latest
+      options: --gpus all --shm-size "16gb"
+    defaults:
+      run:
+        working-directory: accelerate/
+        shell: bash
+    steps:
+      - name: Update clone
+        run: |
+          source activate accelerate
+          git config --global --add safe.directory '*'
+          git fetch && git checkout ${{ github.sha }}
+          pip install -e . --no-deps
+          pip install pytest-reportlog
+
+      - name: Run core and big modeling tests on GPUs
+        run: |
+          source activate accelerate
+          make test_big_modeling
+          make test_core
+
+      - name: Run Integration tests on GPUs
+        run: |
+          source activate accelerate
+          make test_integrations
+
+      - name: Run examples on GPUs
+        run: |
+          source activate accelerate
+          pip uninstall comet_ml -y
+          make test_examples
+
+      - name: Generate Report
+        if: always()
+        run: |
+          python utils/log_reports.py >> $GITHUB_STEP_SUMMARY

.github/workflows/quality.yml

@@ -7,10 +7,10 @@ jobs:
     runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@v2
-    - name: Set up Python 3.6
-      uses: actions/setup-python@v2
+    - name: Set up Python 3.7
+      uses: actions/setup-python@v3
       with:
-        python-version: 3.6
+        python-version: 3.7
     - name: Install Python dependencies
       run: pip install -e .[quality]
     - name: Run Quality check

.github/workflows/run_merge_tests.yml

@@ -0,0 +1,79 @@
+name: Self-hosted runner tests (push to "main")
+
+on:
+  workflow_call:
+  workflow_dispatch:
+
+env:
+  TESTING_MOCKED_DATALOADERS: "1"
+  IS_GITHUB_CI: "1"
+
+jobs:
+  run_all_tests_single_gpu:
+    runs-on: [self-hosted, docker-gpu, multi-gpu]
+    env:
+      CUDA_VISIBLE_DEVICES: "0"
+    container:
+      image: huggingface/accelerate-gpu:latest
+      options: --gpus all --shm-size "16gb"
+    defaults:
+      run:
+        working-directory: accelerate/
+        shell: bash
+    steps:
+      - name: Update clone & pip install
+        run: |
+          source activate accelerate
+          git config --global --add safe.directory '*'
+          git fetch && git checkout ${{ github.sha }}
+          pip install -e .[testing,test_trackers]
+          pip install pytest-reportlog
+
+      - name: Run test on GPUs
+        run: |
+          source activate accelerate
+          make test
+      - name: Run examples on GPUs
+        run: |
+          source activate accelerate
+          pip uninstall comet_ml -y
+          make test_examples
+
+      - name: Generate Report
+        if: always()
+        run: |
+          python utils/log_reports.py >> $GITHUB_STEP_SUMMARY
+
+  run_all_tests_multi_gpu:
+    runs-on: [self-hosted, docker-gpu, multi-gpu]
+    container:
+      image: huggingface/accelerate-gpu:latest
+      options: --gpus all --shm-size "16gb"
+    defaults:
+      run:
+        working-directory: accelerate/
+        shell: bash
+    steps:
+      - name: Update clone
+        run: |
+          source activate accelerate
+          git config --global --add safe.directory '*'
+          git fetch && git checkout ${{ github.sha }}
+          pip install -e .[testing,test_trackers]
+          pip install pytest-reportlog
+
+      - name: Run test on GPUs
+        run: |
+          source activate accelerate
+          make test
+
+      - name: Run examples on GPUs
+        run: |
+          source activate accelerate
+          pip uninstall comet_ml -y
+          make test_examples
+
+      - name: Generate Report
+        if: always()
+        run: |
+          python utils/log_reports.py >> $GITHUB_STEP_SUMMARY

.github/workflows/stale.yml

@@ -0,0 +1,28 @@
+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
+    env:
+      GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+    steps:
+    - uses: actions/checkout@v2
+    
+    - name: Setup Python
+      uses: actions/setup-python@v1
+      with:
+        python-version: 3.7
+    
+    - 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_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@v3
+    - name: Set up python 3.7
+      uses: actions/setup-python@v3
       with:
-        python-version: 3.6
-    - name: Install Python dependencies
-      run: pip install -e .[test]
+        python-version: 3.7
+    
+    - name: Activate python cache
+      uses: actions/cache@v3
+      with:
+        path: |
+          ${{ env.pythonLocation }}
+          ${{ env.HF_HOME }}
+        key: ${{ env.pythonLocation }}-${{ matrix.test-kind }}-${{ hashFiles('setup.py') }}
+    
+    - name: Install the library
+      run: |
+        pip install --upgrade pip
+        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 torch==1.6.0; fi
+        pip install pytest-reportlog
+    
     - name: Run Tests
-      run: make test
+      run: |
+        make ${{ matrix.test-kind }}
+
+    - name: Generate Report
+      if: always()
+      run: |
+        python utils/log_reports.py >> $GITHUB_STEP_SUMMARY

.gitignore

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

CONTRIBUTING.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.
+-->
+
+# 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 ".[quality]"
+   ```
+
+   (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.)
+
+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 `black` and `isort` 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 `flake8` and 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
+   ```
+
+   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
 
-check_dirs := tests src examples
+check_dirs := tests src examples benchmarks
 
 # Check that source code meets quality standards
 
@@ -25,8 +25,40 @@ style:
 	
 # Run tests for the library
 test:
-	python -m pytest -n auto --dist=loadfile -s -v ./tests/
+	python -m pytest -s -v ./tests/ --ignore=./tests/test_examples.py $(if $(IS_GITHUB_CI),--report-log 'all.log',)
 
-# Check that docs can build
-docs:
-	cd docs && make html SPHINXOPTS="-W"
+test_big_modeling:
+	python -m pytest -s -v ./tests/test_big_modeling.py $(if $(IS_GITHUB_CI),--report-log 'big_modeling.log',)
+
+test_core:
+	python -m pytest -s -v ./tests/ --ignore=./tests/test_examples.py --ignore=./tests/deepspeed --ignore=./tests/test_big_modeling.py \
+	--ignore=./tests/fsdp $(if $(IS_GITHUB_CI),--report-log 'core.log',)
+
+test_deepspeed:
+	python -m pytest -s -v ./tests/deepspeed $(if $(IS_GITHUB_CI),--report-log 'deepspeed.log',)
+
+test_fsdp:
+	python -m pytest -s -v ./tests/fsdp $(if $(IS_GITHUB_CI),--report-log 'fsdp.log',)
+
+test_examples:
+	python -m pytest -s -v ./tests/test_examples.py $(if $(IS_GITHUB_CI),--report-log 'examples.log',)
+
+# Broken down example tests for the CI runners
+test_integrations:
+	python -m pytest -s -v ./tests/deepspeed ./tests/fsdp $(if $(IS_GITHUB_CI),--report-log 'integrations.log',)
+
+test_example_differences:
+	python -m pytest -s -v ./tests/test_examples.py::ExampleDifferenceTests $(if $(IS_GITHUB_CI),--report-log '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 '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 '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 'rest.log',)

README.md

@@ -26,16 +26,16 @@ limitations under the License.
         <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">
+    <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/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 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/master/CODE_OF_CONDUCT.md">
+    <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>
@@ -44,6 +44,10 @@ limitations under the License.
 <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 +67,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,8 +85,8 @@ Here is an example:
           output = model(source)
           loss = F.cross_entropy(output, targets)
 
-+         accelerator.backward(loss)
 -         loss.backward()
++         accelerator.backward(loss)
 
           optimizer.step()
 ```
@@ -99,17 +103,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,12 +126,14 @@ 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 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!
@@ -149,6 +155,49 @@ 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 torch.distributed.launch my_script.py` at your convenance.
+
+## 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
+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 accelerator import Accelerator, DeepSpeedPlugin
+
+# deepspeed needs to know your gradient accumulation steps before hand, 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(fp16=True, 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.
@@ -157,6 +206,17 @@ You should use 🤗 Accelerate when you want to easily run your training scripts
 
 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 your training loop, some frameworks that are built on top of 🤗 Accelerate are listed below:
+
+* [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 train, 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.
+* [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.
+* [pytorch-accelerated](https://github.com/Chris-hughes10/pytorch-accelerated) is a lightweight training library, with a streamlined feature set centred 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!
+
+
 ## Installation
 
 This repository is tested on Python 3.6+ and PyTorch 1.4.0+
@@ -174,8 +234,25 @@ 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)
+- DeepSpeed support (Experimental)
+- PyTorch Fully Sharded Data Parallel (FSDP) 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, Lysandre Debut, Thomas Wolf, Philipp Schmid, Zachary Mueller, Sourab Mangrulkar},
+  howpublished = {\url{https://github.com/huggingface/accelerate}},
+  year =         {2022}
+}
+```

benchmarks/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 inference_acc.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.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 accelerate.utils import compute_module_sizes
+from measures_util import end_measure, log_measures, start_measure
+from transformers import AutoConfig, AutoModelForCausalLM, AutoModelForSeq2SeqLM, AutoTokenizer
+
+
+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/measures_util.py

@@ -0,0 +1,86 @@
+import gc
+import threading
+import time
+
+import torch
+
+import psutil
+
+
+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.cuda.empty_cache()
+
+    # CPU mem
+    measures["cpu"] = psutil.Process().memory_info().rss
+    cpu_peak_tracker.start()
+
+    # GPU mem
+    for i in range(torch.cuda.device_count()):
+        measures[str(i)] = torch.cuda.memory_allocated(i)
+    torch.cuda.reset_peak_memory_stats()
+
+    return measures
+
+
+def end_measure(start_measures):
+    # Time
+    measures = {"time": time.time() - start_measures["time"]}
+
+    gc.collect()
+    torch.cuda.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.cuda.device_count()):
+        measures[str(i)] = (torch.cuda.memory_allocated(i) - start_measures[str(i)]) / 2**20
+        measures[f"{i}-peak"] = (torch.cuda.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.cuda.device_count()):
+        print(f"- GPU {i} allocated: {measures[str(i)]:.2f}MiB")
+        peak = measures[f"{i}-peak"]
+        print(f"- GPU {i} peak: {peak:.2f}MiB")
+    print(f"- CPU RAM allocated: {measures['cpu']:.2f}MiB")
+    print(f"- CPU RAM peak: {measures['cpu-peak']:.2f}MiB")

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.7-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.7-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/Dockerfile

@@ -0,0 +1,42 @@
+# Builds GPU docker image of PyTorch
+# 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.7.3
+# 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 and install torch + accelerate
+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/cu113
+
+# Stage 2
+FROM nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.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/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,74 @@
+- 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: Migrating to 🤗 Accelerate
+  - local: basic_tutorials/launch
+    title: Launching distributed code
+  - local: basic_tutorials/notebook
+    title: Launching distributed training from Jupyter Notebooks
+  title: Tutorials
+- sections:
+  - local: usage_guides/gradient_accumulation
+    title: Performing gradient accumulation
+  - local: usage_guides/fsdp
+    title: Fully Sharded Data Parallelism
+  - local: usage_guides/checkpoint
+    title: Saving and loading training states
+  - local: usage_guides/deepspeed
+    title: How to use DeepSpeed
+  - local: usage_guides/tracking
+    title: Using experiment trackers
+  - local: usage_guides/big_modeling
+    title: How to use large models with small resources
+  - local: usage_guides/memory
+    title: How to avoid CUDA Out-of-Memory
+  - local: usage_guides/sagemaker
+    title: Using 🤗 Accelerate on SageMaker
+  - local: usage_guides/mps
+    title: How to use Apple Silicon M1 GPUs
+  - local: usage_guides/training_zoo
+    title: 🤗 Accelerate Example Zoo
+  title: How-To Guides
+- sections:
+  - local: concept_guides/performance
+    title: Comparing performance across distributed setups
+  - local: concept_guides/gradient_synchronization
+    title: Gradient synchronization
+  - local: concept_guides/deferring_execution
+    title: Executing and deferring jobs
+  - local: concept_guides/training_tpu
+    title: TPU best practices
+  title: Concepts and fundamentals
+- sections: 
+  - local: package_reference/accelerator
+    title: Main Accelerator class
+  - local: package_reference/state
+    title: Stateful configuration classes
+  - local: package_reference/cli
+    title: The Command Line
+  - local: package_reference/torch_wrappers
+    title: Torch wrapper classes
+  - local: package_reference/tracking
+    title: Experiment trackers
+  - local: package_reference/launchers
+    title: Distributed 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/kwargs
+    title: Kwargs handlers
+  - local: package_reference/utilities
+    title: Utility functions and classes
+  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/install.mdx

@@ -0,0 +1,99 @@
+<!--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.
+-->
+
+# Installation and Configuration
+
+Before you start, you will need to setup your environment, install the appropriate packages, and configure 🤗 Accelerate. 🤗 Accelerate is tested on **Python 3.7+**.
+
+## Installing 🤗 Accelerate
+
+🤗 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 .
+```
+
+## Configuring 🤗 Accelerate
+
+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: 0.11.0.dev0
+- Platform: Linux-5.10.0-15-cloud-amd64-x86_64-with-debian-11.3
+- Python version: 3.7.12
+- Numpy version: 1.19.5
+- PyTorch version (GPU?): 1.12.0+cu102 (True)
+- `Accelerate` default config:
+        - compute_environment: LOCAL_MACHINE
+        - distributed_type: MULTI_GPU
+        - mixed_precision: no
+        - use_cpu: False
+        - num_processes: 2
+        - machine_rank: 0
+        - num_machines: 1
+        - main_process_ip: None
+        - main_process_port: None
+        - main_training_function: main
+        - deepspeed_config: {}
+        - fsdp_config: {}
+```

docs/source/basic_tutorials/launch.mdx

@@ -0,0 +1,178 @@
+<!--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.
+-->
+
+# Launching your 🤗 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 are.
+
+<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 afterwards 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
+CUDA_VISIBLE_DEVICES="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} ...
+```
+
+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 --num_machines=1 {script_name.py} {--arg1} {--arg2} ...
+```
+
+## 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} ...
+```

docs/source/basic_tutorials/migration.mdx

@@ -0,0 +1,123 @@
+<!--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.
+-->
+
+# Migrating your code to 🤗 Accelerate
+
+This tutorial will detail how to easily convert existing PyTorch code to use 🤗 Accelerate!
+You'll see that by just changing a few lines of code, 🤗 Accelerate can perform its magic and get you on 
+your way towards running your code on distributed systems with ease!
+
+## The base training loop
+
+To begin, write out a very basic PyTorch training loop. 
+
+<Tip>
+
+    We are under the presumption that `training_dataloader`, `model`, `optimizer`, `scheduler`, and `loss_function` have been defined beforehand.
+
+</Tip>
+
+```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()
+```
+
+## Add in 🤗 Accelerate
+
+To start using 🤗 Accelerate, first import and create an [`Accelerator`] instance:
+```python
+from accelerate import Accelerator
+
+accelerator = Accelerator()
+```
+[`Accelerator`] is the main force behind utilizing all the possible options for distributed training!
+
+### Setting the right device
+
+The [`Accelerator`] class knows the right device to move any PyTorch object to at any time, so you should
+change the definition of `device` to come from [`Accelerator`]:
+
+```diff
+- device = 'cuda'
++ device = accelerator.device
+  model.to(device)
+```
+
+### Preparing your objects
+
+Next you need to pass all of the important objects related to training into [`~Accelerator.prepare`]. 🤗 Accelerate will
+make sure everything is setup in the current environment for you to start training:
+
+```
+model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+    model, optimizer, training_dataloader, scheduler
+)
+```
+These objects are returned in the same order they were sent in with. By default when using `device_placement=True`, all of the objects that can be sent to the right device will be.
+If you need to work with data that isn't passed to [~Accelerator.prepare] but should be on the active device, you should pass in the `device` you made earlier. 
+
+<Tip warning={true}>
+
+    Accelerate will only prepare objects that inherit from their respective PyTorch classes (such as `torch.optim.Optimizer`).
+
+</Tip>
+
+### Modifying the training loop
+
+Finally, three lines of code need to be changed in the training loop. 🤗 Accelerate's DataLoader classes will automatically handle the device placement by default,
+and [`~Accelerator.backward`] should be used for performing the backward pass:
+
+```diff
+-   inputs = inputs.to(device)
+-   targets = targets.to(device)
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+-   loss.backward()
++   accelerator.backward(loss)
+```
+
+With that, your training loop is now ready to use 🤗 Accelerate!
+
+## The finished code
+
+Below is the final version of the converted code: 
+
+```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()
+```
+

docs/source/basic_tutorials/notebook.mdx

@@ -0,0 +1,429 @@
+<!--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.
+-->
+
+# Launching Multi-Node Training from a Jupyter Environment
+
+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, a 🤗 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-node 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 at 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 explainations 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 initaliziations)
+    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 this constant available on the active device, set it to the accelerator device
+    mean = mean.to(accelerator.device)
+    std = std.to(accelerator.device)
+
+    # Intantiate 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 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)
+```
+
+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!
+
+## 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.mdx

@@ -0,0 +1,21 @@
+<!--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.
+-->
+
+# 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/concept_guides/deferring_execution.mdx

@@ -0,0 +1,107 @@
+<!--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.
+-->
+
+# Deferring Executions
+
+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 for 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 loading the cached dataset in afterwards
+
+<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"],
+    )
+```

docs/source/concept_guides/gradient_synchronization.mdx

@@ -0,0 +1,119 @@
+<!--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.
+-->
+
+# 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 all of the parameters in a model through `.backward()`. All instances of the model
+need to have updated their gradients, collated, and updated again before moving onto the next batch of data. But when performing 
+gradient accumulation, you accumulate `n` losses and skip `.backward()` until `n` batches have been reached. This 
+can cause a significant slowdown since all the processes need to communicate with them more times than needed. How 
+can you avoid this overhead?
+
+## Solving the slowdown problem
+
+Since you are skipping these batches, their gradients do not need to be synchronized until the point where `.backward()` 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 = accelerator.prepare(model, dataloader)
+
+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)
+```
+
+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 = accelerator.prepare(model, dataloader)
+
+  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)
+```
+
+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 = accelerator.prepare(model, dataloader)
+
+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)
+```
+
+As a result, you should either use *`accelerator.accumulate` or `accelerator.no_sync`* when it comes to API choice. 

docs/source/concept_guides/performance.mdx

@@ -0,0 +1,91 @@
+<!--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.
+-->
+
+# Comparing performance between different device 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's 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 reproducable:
+
+```python
+from accelerate 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)
+    # ^^ safe to call this function even if cuda is not available
+    if is_tpu_available():
+        xm.set_rng_state(seed)
+```
+
+The random state, numpy's state, torch, torch's cuda 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/tlt-mi_archive/clara-train-sdk-v2.0/nvmidl/appendix/training_with_multiple_gpus.html)], 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)
+```
+

docs/source/concept_guides/training_tpu.mdx

@@ -0,0 +1,164 @@
+<!--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.
+-->
+
+# Training on TPUs with 🤗 Accelerate
+
+Training on TPUs can be slightly different than 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 regards 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/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.mdx

@@ -0,0 +1,71 @@
+<!--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.
+-->
+
+# 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="/docs/accelerate/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="/docs/accelerate/usage_guides/gradient_accumulation"
+      ><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="/docs/accelerate/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="/docs/accelerate/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

docs/source/installation.md

@@ -1,96 +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.
--->
-
-# Installation
-
-🤗 Accelerate is tested on Python 3.6+, and PyTorch 1.6.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/). Create a virtual environment with the version of Python you're going
-to use and activate it.
-
-Now, if you want to use 🤗 Accelerate, you can install it with pip.
-
-## Installation with pip
-
-First you need to install PyTorch. Please refer to the
-[PyTorch installation page](https://pytorch.org/get-started/locally/#start-locally) regarding the specific install command for your platform.
-
-When PyTorch has been installed, 🤗 Accelerate can be installed using pip as follows:
-
-```bash
-pip install accelerate
-```
-
-Alternatively, for CPU-support only, you can install 🤗 Accelerate and PyTorch in one line with:
-
-```bash
-pip install accelerate[torch]
-```
-
-To check 🤗 Accelerate is properly installed, run the following command:
-
-```bash
-python -c "TODO write"
-```
-
-## Installing from source
-
-Here is how to quickly install `accelerate` from source:
-
-```bash
-pip install git+https://github.com/huggingface/accelerate
-```
-
-Note that this will install not the latest released version, but the bleeding edge `master` version, which you may want to use in case a bug has been fixed since the last official release and a new release hasn't  been yet rolled out.
-
-While we strive to keep `master` operational at all times, if you notice some issues, they usually get fixed within a few hours or a day and and you're more than welcome to help us detect any problems by opening an [Issue](https://github.com/huggingface/accelerate/issues) and this way, things will get fixed even sooner.
-
-Again, you can run:
-
-```bash
-python -c "TODO write"
-```
-
-to check 🤗 Accelerate is properly installed.
-
-## Editable install
-
-If you want to constantly use the bleeding edge `master` version of the source code, or if you want to contribute to the library and need to test the changes in the code you're making, you will need an editable install. This is done by cloning the repository and installing with the following commands:
-
-``` bash
-git clone https://github.com/huggingface/accelerate.git
-cd transformers
-pip install -e .
-```
-
-This command performs a magical link between the folder you cloned the repository to and your python library paths, and it'll look inside this folder in addition to the normal library-wide paths. So if normally your python packages get installed into:
-```
-~/anaconda3/envs/main/lib/python3.7/site-packages/
-```
-now this editable install will reside where you clone the folder to, e.g. `~/accelerate/` and python will search it too.
-
-Do note that you have to keep that `accelerate` folder around and not delete it to continue using the 🤗 Accelerate library.
-
-Now, let's get to the real benefit of this installation approach. Say, you saw some new feature has been just committed into `master`. If you have already performed all the steps above, to update your accelerate repo to include all the latest commits, all you need to do is to `cd` into that cloned repository folder and update the clone to the latest version:
-
-```
-cd ~/accelerate/
-git pull
-```
-
-There is nothing else to do. Your python environment will find the bleeding edge version of 🤗 Accelerate on the next run.
-

docs/source/internal.rst

@@ -1,85 +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.
-
-Internals
-=======================================================================================================================
-
-
-Optimizer
------------------------------------------------------------------------------------------------------------------------
-
-.. autoclass:: accelerate.optimizer.AcceleratedOptimizer
-
-
-DataLoader
------------------------------------------------------------------------------------------------------------------------
-
-The main work on your PyTorch :obj:`DataLoader` is done by the following function:
-
-.. autofunction:: accelerate.data_loader.prepare_data_loader
-
-
-BatchSamplerShard
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: accelerate.data_loader.DataLoaderShard
-    :members:
-
-
-BatchSamplerShard
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: accelerate.data_loader.BatchSamplerShard
-    :members:
-
-
-IterableDatasetShard
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: accelerate.data_loader.IterableDatasetShard
-    :members:
-
-
-Distributed Config
------------------------------------------------------------------------------------------------------------------------
-
-
-AcceleratorState
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: accelerate.state.AcceleratorState
-    :members:
-
-
-DistributedType
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: accelerate.state.DistributedType
-    :members:
-
-
-Utilities
------------------------------------------------------------------------------------------------------------------------
-
-.. autofunction:: accelerate.utils.extract_model_from_parallel
-
-.. autofunction:: accelerate.utils.gather
-
-.. autofunction:: accelerate.utils.send_to_device
-
-.. autofunction:: accelerate.utils.set_seed
-
-.. autofunction:: accelerate.utils.synchronize_rng_state
-
-.. autofunction:: accelerate.utils.synchronize_rng_states
-
-.. autofunction:: accelerate.utils.wait_for_everyone

docs/source/kwargs.rst

@@ -1,30 +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.
-
-
-Kwargs Handlers
-=======================================================================================================================
-
-The following objects can be passed to the main :class:`~accelerate.Accelerator` to customize how some PyTorch objects
-related to distributed training or mixed precision are created.
-
-
-DistributedDataParallelKwargs
------------------------------------------------------------------------------------------------------------------------
-
-.. autoclass:: accelerate.DistributedDataParallelKwargs
-
-
-GradScalerKwargs
------------------------------------------------------------------------------------------------------------------------
-
-.. autoclass:: accelerate.GradScalerKwargs

docs/source/package_reference/accelerator.mdx

@@ -0,0 +1,163 @@
+<!--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 [`Accelerator`] is the main class provided by 🤗 Accelerate. 
+It serves at the main entrypoint for the API. 
+
+## Quick adaptation of your code
+
+To quickly adapt your script to work on any kind of setup with 🤗 Accelerate just:
+
+1. Initialize an [`Accelerator`] object (that we will call `accelerator` throughout this page) as early as possible in your script.
+2. Pass your dataloader(s), model(s), optimizer(s), and scheduler(s) to the [`~Accelerator.prepare`] method.
+3. Remove all the `.cuda()` or `.to(device)` from your code and let the `accelerator` handle the device placement for you. 
+
+<Tip>
+
+    Step three is optional, but considered a best practice.
+
+</Tip>
+
+4. Replace `loss.backward()` in your code with `accelerator.backward(loss)`
+5. Gather your predictions and labels before storing them or using them for metric computation using [`~Accelerator.gather`]
+
+<Tip warning={true}>
+
+    Step five is mandatory when using distributed evaluation
+    
+</Tip>
+
+In most cases this is all that is needed. The next section lists a few more advanced use cases and nice features
+you should search for and replace by the corresponding methods of your `accelerator`:
+
+## Advanced recommendations
+
+### Printing
+
+`print` statements should be replaced by [`~Accelerator.print`] to be printed once per process
+
+```diff
+- print("My thing I want to print!")
++ accelerator.print("My thing I want to print!")
+```
+
+### Executing processes
+
+#### Once on a single server
+
+For statements that should be executed once per server, use [`~Accelerator.is_local_main_process`]:
+
+```python
+if accelerator.is_local_main_process:
+    do_thing_once_per_server()
+```
+
+A function can be wrapped using the [`~Accelerator.on_local_main_process`] function to achieve the same 
+behavior on a function's execution:
+
+```python
+@accelerator.on_local_main_process
+def do_my_thing():
+    "Something done once per server"
+    do_thing_once_per_server()
+```
+
+#### Only ever once across all servers
+
+For statements that should only ever be executed once, use [`~Accelerator.is_main_process`]:
+
+```python
+if accelerator.is_main_process:
+    do_thing_once()
+```
+
+A function can be wrapped using the [`~Accelerator.on_main_process`] function to achieve the same 
+behavior on a function's execution:
+
+```python
+@accelerator.on_main_process
+def do_my_thing():
+    "Something done once per server"
+    do_thing_once()
+```
+
+#### On specific processes
+
+If a function should be ran on a specific overall or local process index, there are similar decorators 
+to achieve this:
+
+```python
+@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()
+```
+
+```python
+@accelerator.on_process(process_index=0)
+def do_my_thing():
+    "Something done on process index 0"
+    do_thing_on_index_zero()
+```
+
+### Synchronicity control
+
+Use [`~Accelerator.wait_for_everyone`] to make sure all processes join that point before continuing. (Useful before a model save for instance)
+
+### Saving and loading
+
+Use [`~Accelerator.unwrap_model`] before saving to remove all special model wrappers added during the distributed process. 
+
+```python
+model = MyModel()
+model = accelerator.prepare(model)
+# Unwrap
+model = accelerator.unwrap_model(model)
+```
+
+Use [`~Accelerator.save`] instead of `torch.save`:
+
+```diff
+  state_dict = model.state_dict()
+- torch.save(state_dict, "my_state.pkl")
++ accelerator.save(state_dict, "my_state.pkl")
+```
+
+### Operations
+
+Use [`~Accelerator.clip_grad_norm_`] instead of ``torch.nn.utils.clip_grad_norm_`` and [`~Accelerator.clip_grad_value_`] instead of ``torch.nn.utils.clip_grad_value``
+
+### Gradient Accumulation
+
+To perform gradient accumulation use [`~Accelerator.accumulate`] and specify a gradient_accumulation_steps. 
+This will also automatically ensure the gradients are synced or unsynced when on 
+multi-device training, check if the step should actually be performed, and auto-scale the loss:
+
+```diff
+- accelerator = Accelerator()
++ accelerator = Accelerator(gradient_accumulation_steps=2)
+
+  for (input, label) in training_dataloader:
++     with accelerator.accumulate(model):
+          predictions = model(input)
+          loss = loss_function(predictions, labels)
+          accelerator.backward(loss)
+          optimizer.step()
+          scheduler.step()
+          optimizer.zero_grad()
+```
+
+## Overall API documentation:
+
+[[autodoc]] Accelerator

docs/source/package_reference/big_modeling.mdx

@@ -0,0 +1,41 @@
+<!--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.
+-->
+
+# Working with large models
+
+## Dispatching and Offloading Models
+
+[[autodoc]] big_modeling.init_empty_weights
+[[autodoc]] big_modeling.cpu_offload
+[[autodoc]] big_modeling.disk_offload
+[[autodoc]] big_modeling.dispatch_model
+[[autodoc]] big_modeling.load_checkpoint_and_dispatch
+
+## Model Hooks
+
+### Hook Classes
+
+[[autodoc]] hooks.ModelHook
+[[autodoc]] hooks.AlignDevicesHook
+[[autodoc]] hooks.SequentialHook
+
+### Adding Hooks
+
+[[autodoc]] hooks.add_hook_to_module
+[[autodoc]] hooks.attach_execution_device_hook
+[[autodoc]] hooks.attach_align_device_hook
+[[autodoc]] hooks.attach_align_device_hook_on_blocks
+
+### Removing Hooks
+
+[[autodoc]] hooks.remove_hook_from_module
+[[autodoc]] hooks.remove_hook_from_submodules

docs/source/package_reference/cli.mdx

@@ -0,0 +1,153 @@
+<!--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.
+-->
+
+# The Command Line 
+
+Below is a list of all the available commands 🤗 Accelerate with their parameters
+
+## accelerate config
+
+**Command**:
+
+`accelerate config` or `accelerate-config`
+
+Launches a series of prompts to create and save a `default_config.yml` configuration file for your training system. Should 
+always be ran first on your machine.
+
+**Usage**: 
+
+```bash
+accelerate config [arguments]
+```
+
+**Optional Arguments**:
+* `--config_file CONFIG_FILE` (`str`) -- The path to use to store the config file. Will default to a file named default_config.yaml in the cache location, which is the content
+                        of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have such an environment variable, your cache directory
+                        (`~/.cache` or the content of `XDG_CACHE_HOME`) suffixed with `huggingface`.
+* `-h`, `--help` (`bool`) -- Show a help message and exit
+
+## accelerate env
+
+**Command**:
+
+`accelerate env` or `accelerate-env`
+
+Lists the contents of the passed 🤗 Accelerate configuration file. Should always be used when opening an issue on the [GitHub repository](https://github.com/huggingface/accelerate).
+
+**Usage**:
+
+```bash
+accelerate env [arguments]
+```
+
+**Optional Arguments**:
+* `--config_file CONFIG_FILE` (`str`) -- The path to use to store the config file. Will default to a file named default_config.yaml in the cache location, which is the content
+                        of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have such an environment variable, your cache directory
+                        (`~/.cache` or the content of `XDG_CACHE_HOME`) suffixed with `huggingface`.
+* `-h`, `--help` (`bool`) -- Show a help message and exit
+
+## accelerate launch
+
+**Command**:
+
+`accelerate launch` or `accelerate-launch`
+
+Launches a specified script on a distributed system with the right parameters.
+
+**Usage**: 
+
+```bash
+accelerate launch [arguments] {training_script} --{training_script-argument-1} --{training_script-argument-2} ...
+```
+
+**Positional Arguments**:
+
+- `{training_script}` -- The full path to the script to be launched in parallel
+- `--{training_script-argument-1}` -- Arguments of the training script
+
+**Optional Arguments**:
+
+* `-h`, `--help` (`bool`) -- Show a help message and exit
+* `--config_file CONFIG_FILE` (`str`)-- The config file to use for the default values in the launching script.
+* `--cpu` (`bool`) -- Whether or not to force the training on the CPU.
+* `--mixed_precision {no,fp16,bf16}` (`str`) -- Whether or not to use mixed precision training. Choose between FP16 and BF16 (bfloat16) training. BF16 training is only supported on
+                    Nvidia Ampere GPUs and PyTorch 1.10 or later.
+* `--multi_gpu` (`bool`, defaults to `False`) -- Whether or not this should launch a distributed GPU training.
+* `-m`, `--module` (`bool`) -- Change each process to interpret the launch script as a Python module, executing with the same behavior as 'python -m'.
+* `--no_python` (`bool`) -- Skip prepending the training script with 'python' - just execute it directly. Useful when the script is not a Python script.
+
+The rest of these arguments are configured through `accelerate config` and are read in from the specified `--config_file` (or default configuration) for their 
+values. They can also be passed in manually.
+
+**Machine Configuration Arguments**:
+
+The following arguments are useful for customization of worker machines
+* `--machine_rank MACHINE_RANK` (`int`) -- The rank of the machine on which this script is launched.
+* `--num_machines NUM_MACHINES` (`int`) -- The total number of machines used in this training.
+* `--num_processes NUM_PROCESSES` (`int`) -- The total number of processes to be launched in parallel.
+* `--gpu_ids` (`str`) -- What GPUs (by id) should be used for training on this machine as a comma-seperated list
+* `--main_process_ip MAIN_PROCESS_IP` (`str`) -- The IP address of the machine of rank 0.
+* `--main_process_port MAIN_PROCESS_PORT` (`int`) -- The port to use to communicate with the machine of rank 0.
+* `--num_cpu_threads_per_process NUM_CPU_THREADS_PER_PROCESS` (`int`) -- The number of CPU threads per process. Can be tuned for optimal performance.
+
+
+**DeepSpeed Arguments**:
+
+The following arguments are only useful when `use_deepspeed` is passed: 
+* `--use_deepspeed` (`bool`) -- Whether to use deepspeed.
+* `--deepspeed_config_file DEEPSPEED_CONFIG_FILE` (`str`) -- DeepSpeed config file.
+* `--zero_stage ZERO_STAGE` (`str`) -- DeepSpeed's ZeRO optimization stage
+* `--offload_optimizer_device OFFLOAD_OPTIMIZER_DEVICE` (`str`) -- Decides where (none|cpu|nvme) to offload optimizer states
+* `--offload_param_device OFFLOAD_PARAM_DEVICE` (`str`) -- Decides where (none|cpu|nvme) to offload parameters
+* `--gradient_accumulation_steps GRADIENT_ACCUMULATION_STEPS` (`int`) -- Number of gradient_accumulation_steps used in your training script
+* `--gradient_clipping GRADIENT_CLIPPING` (`float`) -- gradient clipping value used in your training script
+The following arguments are related to using ZeRO Stage-3
+* `--zero3_init_flag ZERO3_INIT_FLAG` (`bool`) -- Decides Whether (true|false) to enable `deepspeed.zero.Init` for constructing massive models
+* `--zero3_save_16bit_model ZERO3_SAVE_16BIT_MODEL` (`bool`) -- Decides Whether (true|false) to save 16-bit model weights when using ZeRO Stage-3
+
+**Fully Sharded Data Parallelism Arguments**:
+
+The following arguments are only useful when `use_fdsp` is passed:
+* `--use_fsdp` (`bool`) -- Whether to use fsdp.
+* `--offload_params OFFLOAD_PARAMS` (`bool`) -- Decides Whether (true|false) to offload parameters and gradients to CPU.
+* `--min_num_params MIN_NUM_PARAMS` (`int`) -- FSDP's minimum number of parameters for Default Auto Wrapping.
+* `--sharding_strategy SHARDING_STRATEGY` (`str`) -- FSDP's Sharding Strategy.
+
+**TPU Arguments**:
+
+The following arguments are only useful when `tpu` is passed:
+* `--tpu` (`bool`) - Whether or not this should launch a TPU training.
+* `--main_training_function MAIN_TRAINING_FUNCTION` (`str`) -- The name of the main function to be executed in your script.
+
+**AWS SageMaker Arguments**:
+
+The following arguments are only useful when training in SageMaker
+* `--aws_access_key_id AWS_ACCESS_KEY_ID` (`str`) -- The AWS_ACCESS_KEY_ID used to launch the Amazon SageMaker training job
+* `--aws_secret_access_key AWS_SECRET_ACCESS_KEY` (`str`) -- The AWS_SECRET_ACCESS_KEY used to launch the Amazon SageMaker training job
+
+## accelerate test
+
+`accelerate test` or `accelerate-test`
+
+Runs `accelerate/test_utils/test_script.py` to verify that 🤗 Accelerate has been properly configured on your system and runs. 
+
+**Usage**: 
+
+```bash
+accelerate test [arguments]
+```
+
+**Optional Arguments**:
+* `--config_file CONFIG_FILE` (`str`) -- The path to use to store the config file. Will default to a file named default_config.yaml in the cache location, which is the content
+                        of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have such an environment variable, your cache directory
+                        (`~/.cache` or the content of `XDG_CACHE_HOME`) suffixed with `huggingface`.
+* `-h`, `--help` (`bool`) -- Show a help message and exit

docs/source/package_reference/deepspeed.mdx

@@ -0,0 +1,25 @@
+<!--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.
+-->
+
+# Utilities for DeepSpeed
+
+[[autodoc]] utils.DeepSpeedPlugin
+
+[[autodoc]] utils.DummyOptim
+
+[[autodoc]] utils.DummyScheduler
+
+[[autodoc]] utils.DeepSpeedEngineWrapper
+
+[[autodoc]] utils.DeepSpeedOptimizerWrapper
+
+[[autodoc]] utils.DeepSpeedSchedulerWrapper

docs/source/package_reference/kwargs.mdx

@@ -0,0 +1,29 @@
+<!--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.
+-->
+
+# Kwargs Handlers
+
+The following objects can be passed to the main [`Accelerator`] to customize how some PyTorch objects
+related to distributed training or mixed precision are created.
+
+
+## DistributedDataParallelKwargs
+
+[[autodoc]] DistributedDataParallelKwargs
+
+## GradScalerKwargs
+
+[[autodoc]] GradScalerKwargs
+
+## InitProcessGroupKwargs
+
+[[autodoc]] InitProcessGroupKwargs

docs/source/package_reference/launchers.mdx

@@ -0,0 +1,19 @@
+<!--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.
+-->
+
+# Launchers
+
+Functions for launching training on distributed processes.
+
+
+[[autodoc]] accelerate.notebook_launcher
+[[autodoc]] accelerate.debug_launcher

docs/source/package_reference/logging.mdx

@@ -0,0 +1,24 @@
+<!--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.
+-->
+
+# Logging with Accelerate
+
+Accelerate has its own logging utility to handle logging while in a distributed system.
+To utilize this replace cases of `logging` with `accelerate.logging`:
+```diff
+- import logging
++ from accelerate.logging import get_logger
+- logger = logging.getLogger(__name__)
++ logger = get_logger(__name__)
+```
+
+[[autodoc]] logging.get_logger

docs/source/package_reference/state.mdx

@@ -0,0 +1,23 @@
+<!--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.
+-->
+
+# Stateful Classes
+
+Below are variations of a [singleton class](https://en.wikipedia.org/wiki/Singleton_pattern) in the sense that all
+instances share the same state, which is initialized on the first instantiation.
+
+These classes are immutable and store information about certain configurations or 
+states.
+
+[[autodoc]] state.AcceleratorState
+
+[[autodoc]] state.GradientState

docs/source/package_reference/torch_wrappers.mdx

@@ -0,0 +1,33 @@
+<!--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.
+-->
+
+# Wrapper classes for torch Dataloaders, Optimizers, and Schedulers
+
+The internal classes Accelerate uses to prepare objects for distributed training
+when calling [`~Accelerator.prepare`].
+
+## Datasets and DataLoaders
+
+[[autodoc]] data_loader.prepare_data_loader
+
+[[autodoc]] data_loader.BatchSamplerShard
+[[autodoc]] data_loader.IterableDatasetShard
+[[autodoc]] data_loader.DataLoaderShard
+[[autodoc]] data_loader.DataLoaderDispatcher
+
+## Optimizers 
+
+[[autodoc]] optimizer.AcceleratedOptimizer
+
+## Schedulers 
+
+[[autodoc]] scheduler.AcceleratedScheduler

docs/source/package_reference/tracking.mdx

@@ -0,0 +1,26 @@
+<!--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.
+-->
+
+# Experiment Tracking
+
+## The Base Tracker Class
+
+[[autodoc]] tracking.GeneralTracker
+
+## Integrated Trackers
+
+[[autodoc]] tracking.TensorBoardTracker
+    - __init__
+[[autodoc]] tracking.WandBTracker
+    - __init__
+[[autodoc]] tracking.CometMLTracker
+    - __init__

docs/source/package_reference/utilities.mdx

@@ -0,0 +1,95 @@
+<!--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.
+-->
+
+# Helpful Utilities
+
+Below are a variety of utility functions that 🤗 Accelerate provides, broken down by use-case. 
+
+## Data Classes
+
+These are basic dataclasses used throughout 🤗 Accelerate and they can be passed in as parameters.
+
+[[autodoc]] utils.DistributedType
+
+[[autodoc]] utils.LoggerType
+
+[[autodoc]] utils.PrecisionType
+
+## Data Manipulation and Operations
+
+These include data operations that mimic the same `torch` ops but can be used on distributed processes.
+
+[[autodoc]] utils.broadcast
+
+[[autodoc]] utils.concatenate
+
+[[autodoc]] utils.gather
+
+[[autodoc]] utils.pad_across_processes
+
+[[autodoc]] utils.reduce
+
+[[autodoc]] utils.send_to_device
+
+## Environment Checks
+
+These functionalities check the state of the current working environment including information about the operating system itself, what it can support, and if particular dependencies are installed. 
+
+[[autodoc]] utils.is_bf16_available
+
+[[autodoc]] utils.is_torch_version
+
+[[autodoc]] utils.is_tpu_available
+
+## Environment Configuration
+
+[[autodoc]] utils.write_basic_config
+
+When setting up 🤗 Accelerate for the first time, rather than running `accelerate config` [~utils.write_basic_config] can be used as an alternative for quick configuration.
+
+## Memory
+
+[[autodoc]] utils.get_max_memory
+
+[[autodoc]] utils.find_executable_batch_size
+
+## Modeling
+
+These utilities relate to interacting with PyTorch models
+
+[[autodoc]] utils.extract_model_from_parallel
+
+[[autodoc]] utils.get_max_layer_size
+
+[[autodoc]] utils.offload_state_dict
+
+
+## Parallel
+
+These include general utilities that should be used when working in parallel.
+
+[[autodoc]] utils.extract_model_from_parallel
+
+[[autodoc]] utils.save
+
+[[autodoc]] utils.wait_for_everyone
+
+
+## Random
+
+These utilities relate to setting and synchronizing of all the random states.
+
+[[autodoc]] utils.set_seed
+
+[[autodoc]] utils.synchronize_rng_state
+
+[[autodoc]] utils.synchronize_rng_states

docs/source/quicktour.mdx

@@ -0,0 +1,501 @@
+<!--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.
+-->
+
+# Quick tour
+
+Let's have a look at the 🤗 Accelerate main features and traps to avoid.
+
+## Main use
+
+To use 🤗 Accelerate in your own script, you have to change four things:
+
+1. Import the [`Accelerator`] main class and instantiate one in an `accelerator` object:
+
+```python
+from accelerate import Accelerator
+
+accelerator = Accelerator()
+```
+
+This should happen as early as possible in your training script as it will initialize everything necessary for
+distributed training. You don't need to indicate the kind of environment you are in (just one machine with a GPU, one
+machines with several GPUs, several machines with multiple GPUs or a TPU), the library will detect this automatically.
+
+2. Remove the call `.to(device)` or `.cuda()` for your model and input data. The `accelerator` object
+will handle this for you and place all those objects on the right device for you. If you know what you're doing, you
+can leave those `.to(device)` calls but you should use the device provided by the `accelerator` object:
+`accelerator.device`.
+
+To fully deactivate the automatic device placement, pass along `device_placement=False` when initializing your
+[`Accelerator`].
+
+<Tip warning={true}>
+
+    If you place your objects manually on the proper device, be careful to create your optimizer after putting your
+    model on `accelerator.device` or your training will fail on TPU.
+
+</Tip>
+
+3. Pass all objects relevant to training (optimizer, model, training dataloader, learning rate scheduler) to the
+[`~Accelerator.prepare`] method. This will make sure everything is ready for training.
+
+```python
+model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
+    model, optimizer, train_dataloader, lr_scheduler
+)
+```
+
+In particular, your training dataloader will be sharded across all GPUs/TPU cores available so that each one sees a
+different portion of the training dataset. Also, the random states of all processes will be synchronized at the
+beginning of each iteration through your dataloader, to make sure the data is shuffled the same way (if you decided to
+use `shuffle=True` or any kind of random sampler).
+
+<Tip>
+
+    The actual batch size for your training will be the number of devices used multiplied by the batch size you set in
+    your script: for instance training on 4 GPUs with a batch size of 16 set when creating the training dataloader will
+    train at an actual batch size of 64.
+
+</Tip>
+
+Alternatively, you can use the option `split_batches=True` when creating initializing your
+[`Accelerator`], in which case the batch size will always stay the same, whether your run your
+script on 1, 2, 4 or 64 GPUs.
+
+You should execute this instruction as soon as all objects for training are created, before starting your actual
+training loop.
+
+<Tip warning={true}>
+
+    You should only pass the learning rate scheduler to [`~Accelerator.prepare`] when the scheduler needs to be stepped
+    at each optimizer step.
+
+</Tip>
+
+<Tip warning={true}>
+
+    Your training dataloader may change length when going through this method: if you run on X GPUs, it will have its
+    length divided by X (since your actual batch size will be multiplied by X), unless you set
+    `split_batches=True`.
+
+</Tip>
+
+Any instruction using your training dataloader length (for instance if you want to log the number of total training
+steps) should go after the call to [`~Accelerator.prepare`].
+
+You can perfectly send your dataloader to [`~Accelerator.prepare`] on its own, but it's best to send the
+model and optimizer to [`~Accelerator.prepare`] together.
+
+You may or may not want to send your validation dataloader to [`~Accelerator.prepare`], depending on
+whether you want to run distributed evaluation or not (see below).
+
+4. Replace the line `loss.backward()` by `accelerator.backward(loss)`.
+
+And you're all set! With all these changes, your script will run on your local machine as well as on multiple GPUs or a
+TPU! You can either use your favorite tool to launch the distributed training, or you can use the 🤗 Accelerate
+launcher.
+
+
+## Distributed evaluation
+
+You can perform regular evaluation in your training script, if you leave your validation dataloader out of the
+[`~Accelerator.prepare`] method. In this case, you will need to put the input data on the
+`accelerator.device` manually.
+
+To perform distributed evaluation, send along your validation dataloader to the [`~Accelerator.prepare`]
+method:
+
+```python
+validation_dataloader = accelerator.prepare(validation_dataloader)
+```
+
+As for your training dataloader, it will mean that (should you run your script on multiple devices) each device will
+only see part of the evaluation data. This means you will need to group your predictions together. This is very easy to
+do with the [`~Accelerator.gather_for_metrics`] method.
+
+```python
+for inputs, targets in validation_dataloader:
+    predictions = model(inputs)
+    # Gather all predictions and targets
+    all_predictions, all_targets = accelerator.gather_for_metrics((predictions, targets))
+    # Example of use with a *Datasets.Metric*
+    metric.add_batch(all_predictions, all_targets)
+```
+
+<Tip warning={true}>
+
+    Similar to the training dataloader, passing your validation dataloader through
+    [`~Accelerator.prepare`] may change it: if you run on X GPUs, it will have its length divided by X
+    (since your actual batch size will be multiplied by X), unless you set `split_batches=True`.
+
+</Tip>
+
+Any instruction using your training dataloader length (for instance if you need the number of total training steps
+to create a learning rate scheduler) should go after the call to [`~Accelerator.prepare`]. 
+
+Some data at the end of the dataset may be duplicated so the batch can be divided equally among all workers. As a result, metrics
+should be calculated through the [`~Accelerator.gather_for_metrics`] method to automatically remove the duplicated data while gathering.
+
+<Tip>
+
+    If for some reason you don't wish to have this automatically done, [`~Accelerator.gather`] can be used instead to gather 
+    the data across all processes and this can manually be done instead.
+
+</Tip>
+
+
+<Tip warning={true}>
+
+    The [`~Accelerator.gather`] and [`~Accelerator.gather_for_metrics`] methods require the tensors to be all the same size on each process. If
+    you have tensors of different sizes on each process (for instance when dynamically padding to the maximum length in
+    a batch), you should use the [`~Accelerator.pad_across_processes`] method to pad you tensor to the
+    biggest size across processes.
+
+</Tip>
+
+## Launching your distributed script
+
+You can use the regular commands to launch your distributed training (like `torch.distributed.launch` for
+PyTorch), they are fully compatible with 🤗 Accelerate. The only caveat here is that 🤗 Accelerate uses the environment
+to determine all useful information, so `torch.distributed.launch` should be used with the flag `--use_env`.
+
+🤗 Accelerate also provides a CLI tool that unifies all launchers, so you only have to remember one command. To use it,
+just run:
+
+```bash
+accelerate config
+```
+
+on your machine and reply to the questions asked. This will save a *default_config.yaml* file in your cache folder for
+🤗 Accelerate. That cache folder is (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*
+
+You can also specify with the flag `--config_file` the location of the file you want to save.
+
+Once this is done, you can test everything is going well on your setup by running:
+
+```bash
+accelerate test
+```
+
+This will launch a short script that will test the distributed environment. If it runs fine, you are ready for the next
+step!
+
+Note that if you specified a location for the config file in the previous step, you need to pass it here as well:
+
+```bash
+accelerate test --config_file path_to_config.yaml
+```
+
+Now that this is done, you can run your script with the following command:
+
+```bash
+accelerate launch path_to_script.py --args_for_the_script
+```
+
+If you stored the config file in a non-default location, you can indicate it to the launcher like his:
+
+```bash
+accelerate launch --config_file path_to_config.yaml path_to_script.py --args_for_the_script
+```
+
+You can also override any of the arguments determined by your config file. 
+To see the complete list of parameters that you can pass in, run `accelerate launch -h`. 
+
+Check out the [Launch tutorial](basic_tutorials/launch) for more information about launching your scripts. 
+
+
+## Launching training from a notebook
+
+In Accelerate 0.3.0, a new [`notebook_launcher`] has been introduced to help you launch your training
+function from a notebook. This launcher supports launching a training with TPUs on Colab or Kaggle, as well as training
+on several GPUs (if the machine on which you are running your notebook has them).
+
+Just define a function responsible for your whole training and/or evaluation in a cell of the notebook, then execute a
+cell with the following code:
+
+```python
+from accelerate import notebook_launcher
+
+notebook_launcher(training_function)
+```
+
+<Tip warning={true}>
+
+    Your [`Accelerator`] object should only be defined inside the training function. This is because the
+    initialization should be done inside the launcher only.
+
+</Tip>
+
+Check out the [Notebook Launcher tutorial](basic_tutorials/notebook) for more information about training on TPUs. 
+
+
+## Training on TPU
+
+If you want to launch your script on TPUs, there are a few caveats you should be aware of. Behind the scenes, the TPUs
+will create a graph of all the operations happening in your training step (forward pass, backward pass and optimizer
+step). This is why your first step of training will always be very long as building and compiling this graph for
+optimizations takes some time.
+
+The good news is that this compilation will be cached so the second step and all the following will be much faster. The
+bad news is that it only applies if all of your steps do exactly the same operations, which implies:
+
+- having all tensors of the same length in all your batches
+- having static code (i.e., not a for loop of length that could change from step to step)
+
+Having any of the things above change between two steps will trigger a new compilation which will, once again, take a
+lot of time. In practice, that means you must take special care to have all your tensors in your inputs of the same
+shape (so no dynamic padding for instance if you are in an NLP problem) and should not use layers with for loops that
+have different lengths depending on the inputs (such as an LSTM) or the training will be excruciatingly slow.
+
+To introduce special behavior in your script for TPUs you can check the `distributed_type` of your
+`accelerator`:
+
+```python docstyle-ignore
+from accelerate import DistributedType
+
+if accelerator.distributed_type == DistributedType.TPU:
+    # do something of static shape
+else:
+    # go crazy and be dynamic
+```
+
+The [NLP example](https://github.com/huggingface/accelerate/blob/main/examples/nlp_example.py) shows an example in a 
+situation with dynamic padding.
+
+One last thing to pay close attention to: if your model has tied weights (such as language models which tie the weights
+of the embedding matrix with the weights of the decoder), moving this model to the TPU (either yourself or after you
+passed your model to [`~Accelerator.prepare`]) will break the tying. You will need to retie the weights
+after. You can find an example of this in the [run_clm_no_trainer](https://github.com/huggingface/transformers/blob/master/examples/pytorch/language-modeling/run_clm.py) script in
+the Transformers repository.
+
+Check out the [TPU tutorial](concept_guides/training_tpu) for more information about training on TPUs. 
+
+
+## Other caveats
+
+We list here all smaller issues you could have in your script conversion and how to resolve them.
+
+### Execute a statement only on one processes
+
+Some of your instructions only need to run for one process on a given server: for instance a data download or a log
+statement. To do this, wrap the statement in a test like this:
+
+```python docstyle-ignore
+if accelerator.is_local_main_process:
+    # Is executed once per server
+```
+
+Another example is progress bars: to avoid having multiple progress bars in your output, you should only display one on
+the local main process:
+
+```python
+from tqdm.auto import tqdm
+
+progress_bar = tqdm(range(args.max_train_steps), disable=not accelerator.is_local_main_process)
+```
+
+The *local* means per machine: if you are running your training on two servers with several GPUs, the instruction will
+be executed once on each of those servers. If you need to execute something only once for all processes (and not per
+machine) for instance, uploading the final model to the 🤗 model hub, wrap it in a test like this:
+
+```python docstyle-ignore
+if accelerator.is_main_process:
+    # Is executed once only
+```
+
+For printing statements you only want executed once per machine, you can just replace the `print` function by
+`accelerator.print`.
+
+
+### Defer execution
+
+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. 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).
+
+
+### Saving/loading a model
+
+Saving the model you trained might need a bit of adjustment: first you should wait for all processes to reach that
+point in the script as shown above, and then, you should unwrap your model before saving it. This is because when going
+through the [`~Accelerator.prepare`] method, your model may have been placed inside a bigger model,
+which deals with the distributed training. This in turn means that saving your model state dictionary without taking
+any precaution will take that potential extra layer into account, and you will end up with weights you can't load back
+in your base model.
+
+This is why it's recommended to *unwrap* your model first. Here is an example:
+
+```
+accelerator.wait_for_everyone()
+unwrapped_model = accelerator.unwrap_model(model)
+accelerator.save(unwrapped_model.state_dict(), filename)
+```
+
+If your script contains logic to load a checkpoint, we also recommend you load your weights in the unwrapped model
+(this is only useful if you use the load function after making your model go through
+[`~Accelerator.prepare`]). Here is an example:
+
+```
+unwrapped_model = accelerator.unwrap_model(model)
+unwrapped_model.load_state_dict(torch.load(filename))
+```
+
+Note that since all the model parameters are references to tensors, this will load your weights inside `model`.
+
+## Saving/loading entire states
+
+When training your model, you may want to save the current state of the model, optimizer, random generators, and potentially LR schedulers to be restored in the _same script_.
+You can use [`~Accelerator.save_state`] and [`~Accelerator.load_state`] respectively to do so, just by simply passing in a save location. 
+If you have registered any other stateful items to be stored through [`~Accelerator.register_for_checkpointing`] they will also be saved and/or loaded.
+
+<Tip>
+
+    Every object passed to [`~Accelerator.register_for_checkpointing`] must have a `load_state_dict` and `state_dict` function to be stored
+
+</Tip>
+
+
+### Gradient clipping
+
+If you are using gradient clipping in your script, you should replace the calls to
+`torch.nn.utils.clip_grad_norm_` or `torch.nn.utils.clip_grad_value_` with [`~Accelerator.clip_grad_norm_`]
+and [`~Accelerator.clip_grad_value_`] respectively.
+
+
+### Mixed Precision training
+
+If you are running your training in Mixed Precision with 🤗 Accelerate, you will get the best result with your loss being
+computed inside your model (like in Transformer models for instance). Every computation outside of the model will be
+executed in full precision (which is generally what you want for loss computation, especially if it involves a
+softmax). However you might want to put your loss computation inside the *accelerator.autocast* context manager:
+
+```
+with accelerator.autocast():
+    loss = complex_loss_function(outputs, target):
+```
+
+Another caveat with Mixed Precision training is that the gradient will skip a few updates at the beginning and
+sometimes during training: because of the dynamic loss scaling strategy, there are points during training where the
+gradients have overflown, and the loss scaling factor is reduced to avoid this happening again at the next step.
+
+This means that you may update your learning rate scheduler when there was no update, which is fine in general, but may
+have an impact when you have very little training data, or if the first learning rate values of your scheduler are very
+important. In this case, you can skip the learning rate scheduler updates when the optimizer step was not done like
+this:
+
+```
+if not accelerator.optimizer_step_was_skipped:
+    lr_scheduler.step()
+```
+
+### Gradient Accumulation 
+
+To perform gradient accumulation use [`~Accelerator.accumulate`] and specify a `gradient_accumulation_steps`. 
+This will also automatically ensure the gradients are synced or unsynced when on multi-device training, check if the step should
+actually be performed, and auto-scale the loss:
+
+```python
+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()
+```
+
+### DeepSpeed
+
+DeepSpeed support is experimental, so the underlying API will evolve in the near future and may have some slight
+breaking changes. In particular, 🤗 Accelerate does not support DeepSpeed config you have written yourself yet, this
+will be added in a next version.
+
+<Tip warning={true}>
+
+    The [`notebook_launcher`] does not support the DeepSpeed integration yet.
+
+</Tip>
+
+## Internal mechanism
+
+Internally, the library 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`].
+
+Then, when calling [`~Accelerator.prepare`], the library:
+
+- wraps your model(s) in the container adapted for the distributed setup,
+- wraps your optimizer(s) in a [`~optimizer.AcceleratedOptimizer`],
+- creates a new version of your dataloader(s) in a [`~data_loader.DataLoaderShard`].
+
+While the model(s) and optimizer(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.
+
+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 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>
+
+For more details about the internals, see the [Internals page](package_reference/torch_wrappers).

docs/source/quicktour.rst

@@ -1,377 +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.
-
-Quick tour
-=======================================================================================================================
-
-Let's have a look at a look at 🤗 Accelerate main features and traps to avoid.
-
-Main use
------------------------------------------------------------------------------------------------------------------------
-
-To use 🤗 Accelerate in your own script, you have to change four things:
-
-1. Import the :class:`~accelerate.Accelerator` main class instantiate one in an :obj:`accelerator` object:
-
-.. code-block:: python
-
-    from accelerate import Accelerator
-
-    accelerator = Accelerator()
-
-This should happen as early as possible in your training script as it will initialize everything necessary for
-distributed training. You don't need to indicate the kind of environment you are in (just one machine with a GPU, one
-match with several GPUs, several machines with multiple GPUs or a TPU), the library will detect this automatically.
-
-2. Remove the call :obj:`.to(device)` or :obj:`.cuda()` for your model and input data. The :obj:`accelerator` object
-will handle this for you and place all those objects on the right device for you. If you know what you're doing, you
-can leave those :obj:`.to(device)` calls but you should use the device provided by the :obj:`accelerator` object:
-:obj:`accelerator.device`.
-
-To fully deactivate the automatic device placement, pass along :obj:`device_placement=False` when initializing your
-:class:`~accelerate.Accelerator`.
-
-.. Warning::
-
-    If you place your objects manually on the proper device, be careful to create your optimizer after putting your
-    model on :obj:`accelerator.device` or your training will fail on TPU.
-
-3. Pass all objects relevant to training (optimizer, model, training dataloader) to the
-:meth:`~accelerate.Accelerator.prepare` method. This will make sure everything is ready for training.
-
-.. code-block:: python
-
-    model, optimizer, train_dataloader = accelerator.prepare(model, optimizer, train_dataloader)
-
-In particular, your training dataloader will be sharded accross all GPUs/TPU cores available so that each one sees a
-different portion of the training dataset. Also, the random states of all processes will be synchronized at the
-beginning of each iteration through your dataloader, to make sure the data is shuffled the same way (if you decided to
-use :obj:`shuffle=True` or any kind of random sampler).
-
-.. Note::
-
-    The actual batch size for your training will be the number of devices used multiplied by the batch size you set in
-    your script: for instance training on 4 GPUs with a batch size of 16 set when creating the training dataloader will
-    train at an actual batch size of 64.
-
-    Alternatively, you can use the option :obj:`split_batches=True` when creating initializing your
-    :class:`~accelerate.Accelerator`, in which case the batch size will always stay the same, whether your run your
-    script on 1, 2, 4 or 64 GPUs.
-
-You should execute this instruction as soon as all objects for training are created, before starting your actual
-training loop.
-
-.. Warning::
-
-    Your training dataloader may change length when going through this method: if you run on X GPUs, it will have its
-    length divided by X (since your actual batch size will be multiplied by X), unless you set
-    :obj:`split_batches=True`.
-
-    Any instruction using your training dataloader length (for instance if you need the number of total training steps
-    to create a learning rate scheduler) should go after the call to :meth:`~accelerate.Accelerator.prepare`.
-
-You can perfectly send your dataloader to :meth:`~accelerate.Accelerator.prepare` on its own, but it's best to send the
-model and optimizer to :meth:`~accelerate.Accelerator.prepare` together.
-
-You may or may not want to send your validation dataloader to :meth:`~accelerate.Accelerator.prepare`, depending on
-whether you want to run distributed evaluation or not (see below).
-
-4. Replace the line :obj:`loss.backward()` by :obj:`accelerator.backward(loss)`.
-
-And you're all set! With all these changes, your script will run on your local machine as well as on multiple GPUs or a
-TPU! You can either use your favorite tool to launch the distributed training, or you can use the 🤗 Accelerate
-launcher.
-
-
-Distributed evaluation
------------------------------------------------------------------------------------------------------------------------
-
-You can perform regular evaluation in your training script, if you leave your validation dataloader out of the
-:meth:`~accelerate.Accelerator.prepare` method. In this case, you will need to put the input data on the
-:obj:`accelerator.device` manually.
-
-To perform distributed evaluation, send along your validation dataloader to the :meth:`~accelerate.Accelerator.prepare`
-method:
-
-.. code-block:: python
-
-    validation_dataloader = accelerator.prepare(validation_dataloader)
-
-Like for your training dataloader, it will mean that (should you run your script on multiple devices) each device will
-only see part of the evaluation data. This means you will need to group your predictions together. This is very easy to
-do with the :meth:`~accelerate.Accelerator.gather` method.
-
-.. code-block:: python
-
-    for inputs, targets in validation_dataloader:
-        predictions = model(inputs)
-        # Gather all predictions and targets
-        all_predictions = accelerator.gather(predictions)
-        all_targets = accelerator.gather(targets)
-        # Example of use with a `Datasets.Metric`
-        metric.add_batch(all_predictions, all_targets)
-
-
-.. Warning::
-
-    Like for the training dataloader, passing your validation dataloader through
-    :meth:`~accelerate.Accelerator.prepare` may change its: if you run on X GPUs, it will have its length divided by X
-    (since your actual batch size will be multiplied by X), unless you set :obj:`split_batches=True`.
-
-    Any instruction using your training dataloader length (for instance if you need the number of total training steps
-    to create a learning rate scheduler) should go after the call to :meth:`~accelerate.Accelerator.prepare`.
-
-.. Warning::
-
-    The :meth:`~accelerate.Accelerator.gather` method requires the tensors to be all the same size on each process. If
-    you have tensors of different sizes on each process (for instance when dynamically padding to the maximum length in
-    a batch), you should use the :meth:`~accelerate.Accelerator.pad_across_processes` method to pad you tensor to the
-    biggest size across processes.
-
-
-Launching your distributed script
------------------------------------------------------------------------------------------------------------------------
-
-You can use the regular commands to launch your distributed training (like :obj:`torch.distributed.launch` for
-PyTorch), they are fully compatible with 🤗 Accelerate. The only caveat here is that 🤗 Accelerate uses the environment
-to determine all useful information, so :obj:`torch.distributed.launch` should be used with the flag :obj:`--use_env`.
-
-🤗 Accelerate also provides a CLI tool that unifies all launcher, so you only have to remember one command. To use it,
-just run
-
-.. code-block:: bash
-
-    accelerate config
-
-on your machine and reply to the questions asked. This will save a `default_config.json` file in your cache folder for
-🤗 Accelerate. That cache folder is (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`
-
-You can also specify with the flag :obj:`--config_file` the location of the file you want to save.
-
-Once this is done, you can test everything is going well on your setup by running
-
-.. code-block:: bash
-
-    accelerate test
-
-
-This will launch a short script that will test the distributed environment. If it runs fine, you are ready for the next
-step!
-
-Note that if you specified a location for the config file in the previous step, you need to pass it here as well:
-
-.. code-block:: bash
-
-    accelerate test --config_file path_to_config.json
-
-
-Now that this is done, you can run your script with the following command:
-
-.. code-block:: bash
-
-    accelerate launch path_to_script.py --args_for_the_script
-
-
-If you stored the config file in a non-default location, you can indicate it to the launcher like his:
-
-.. code-block:: bash
-
-    accelerate launch --config_file path_to_config.json path_to_script.py --args_for_the_script
-
-You can also override any of the arguments determined by your config file, see TODO: insert ref here.
-
-
-Training on TPU
------------------------------------------------------------------------------------------------------------------------
-
-If you want to launch your script on TPUs, there are a few caveats you should be aware of. Behind the scenes, the TPUs
-will create a graph of all the operations happening im your training step (forward pass, backward pass and optimizer
-step). This is why your first step of training will always be very long as building and compiling this graph for
-optimizations takes some time.
-
-The good news is that this compilation will be cached so the second step and all the following will be much faster. The
-bas news is that it only applies if all of your steps do exactly the same operations, which implies:
-
-- having all tensors of the same length in all your lenghts
-- having static code (i.e., not a foor loop of length that could change from step to step)
-
-Having any of the things above change between two steps will trigger a new compilation which will, once again, take a
-lof of time. In practice, that means you must take special care to have all your tensors in your inputs of the same
-shape (so no dynamic padding for instance if you are in an NLP problem) and should not use layer with for loops that
-have different lengths depending on the inputs (such as an LSTM) or the training will be excruciatingly slow.
-
-To introduce special behavior in your script for TPUs you can check the :obj:`distributed_type` of your
-:obj:`accelerator`:
-
-.. code-block:: python
-
-    from accelerate import DistributedType
-
-    if accelerator.distributed_type == DistributedType.TPU:
-        # do something of static shape
-    else:
-        # go crazy and be dynamic
-
-The `NLP example <https://github.com/huggingface/accelerate/blob/main/examples/nlp_example.py>`__ shows an example in
-situation with dynamic padding.
-
-
-Other caveats
------------------------------------------------------------------------------------------------------------------------
-
-We list here all smaller issues you could have in your script conversion and how to resolve them.
-
-Execute a statement only on one processes
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-Some of your instructions only need to run for one process on a given server: for instance a data download or a log
-statement. To do this, wrap the statement in a test like this:
-
-.. code-block:: python
-
-    if accelerator.is_local_main_process:
-        # Is executed once per server
-
-Another example is progress bars: to avoid having multiple progress bars in your output, you should only display one on
-the local main process:
-
-.. code-block:: python
-
-    from tqdm.auto import tqdm
-
-    progress_bar = tqdm(range(args.max_train_steps), disable=not accelerator.is_local_main_process)
-
-The `local` means per machine: if you are running your training on two servers with several GPUs, the instruction will
-be executed once on each of those servers. If you need to execute something only once for all processes (and not per
-machine) for instance, uploading the final model to the 🤗 model hub, wrap it in a test like this:
-
-.. code-block:: python
-
-    if accelerator.is_main_process:
-        # Is executed once only
-
-For printing statements you only want executed once per machine, you can just replace the :obj:`print` function by
-:obj:`accelerator.print`.
-
-
-Defer execution
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-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. To do this, just write the
-following line in your code:
-
-.. code-block::
-
-    accelerator.wait_for_everyone()
-
-This instruction will block all the processes that arrive them first until all the other processes have reached that
-point (if you run your script on just one GPU or CPU, this wont' do anything).
-
-
-Saving/loading a model
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-Saving the model you trained might need a bit of adjustment: first you should wait for all processes to reach that
-point in the script as shown above, and then, you should unwrap your model before saving it. This is because when going
-through the :meth:`~accelerate.Accelerator.prepare` method, your model may have been placed inside a bigger model,
-which deals with the distributed training. This in turn means that saving your model state dictionary without taking
-any precaution will take that potential extra layer into account, and you will end up with weights you can't load back
-in your base model.
-
-This is why it's recommended to `unwrap` your model first. Here is an example:
-
-.. code-block::
-
-    accelerator.wait_for_everyone()
-    unwrapped_model = accelerator.unwrap_model(model)
-    accelerator.save(unwrapped_model.state_dict(), filename)
-
-If your script contains a logic to load checkpoint, we also recommend you load your weights in the unwrapped model
-(this is only useful if you use the load function after making your model go through
-:meth:`~accelerate.Accelerator.prepare`). Here is an example:
-
-.. code-block::
-
-    unwrapped_model = accelerator.unwrap_model(model)
-    unwrapped_model.load_state_dict(torch.load(filename))
-
-Note that since all the model parameters are references to tensors, this will load your weights inside :obj:`model`.
-
-Gradient clipping
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-If you are using gradient clipping in your script, you should replace the calls to
-:obj:`torch.nn.utils.clip_grad_norm_` or :obj:`torch.nn.utils.clip_grad_value_` with :obj:`accelerator.clip_grad_norm_`
-and :obj:`accelerator.clip_grad_value_` respectively.
-
-
-Internal mechanism
------------------------------------------------------------------------------------------------------------------------
-
-Internally, the library 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 :class:`~accelerate.state.AcceleratorState`.
-
-This class is initialized the first time you instantiate a :class:`~accelerate.Accelerator` as well as performing any
-specific initialization your distributed setup needs. Its state is then uniquely shared through all instances of
-:class:`~accelerate.state.AcceleratorState`.
-
-Then, when calling :meth:`~accelerate.Accelerator.prepare`, the library:
-
-- wraps your model(s) in the container adapted for the distributed setup,
-- wraps your optimizer(s) in a :class:`~accelerate.optimizer.AcceleratedOptimizer`,
-- creates a new version of your dataloader(s) in a :class:`~accelerate.data_loader.DataLoaderShard`.
-
-While the model(s) and optimizer(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 :obj:`batch_sampler` of a dataloader once it's been created and the
-library handles the sharding of your data between processes by changing that :obj:`batch_sampler` to yield every other
-:obj:`num_processes` batches.
-
-The :class:`~accelerate.data_loader.DataLoaderShard` subclasses :obj:`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
-  :obj:`device_placement=True`).
-
-The random number generator synchronization will by default synchronize:
-
-- the :obj:`generator` attribute of a given sampler (like the PyTorch :obj:`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 :obj:`rng_types` argument of the main
-:class:`~accelerate.Accelerator`. In PyTorch >= 1.6, it is recommended to rely on local :obj:`generator` to avoid
-setting the same seed in the main random number generator in all processes.
-
-.. Warning::
-
-    Synchronization 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 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).
-
-.. Note::
-
-    The randomization part of your custom sampler, batch sampler or iterable dataset should be done using a local
-    :obj:`torch.Generator` object (in PyTorch >= 1.6), see the traditional :obj:`RandomSampler`, as an example.
-
-See more details about the internal in the :doc:`Internals page <internal>`.

docs/source/sagemaker.rst

@@ -1,169 +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.
-
-Amazon SageMaker
-=======================================================================================================================
-
-Hugging Face and Amazon introduced new `Hugging Face Deep Learning Containers (DLCs)
-<https://github.com/aws/deep-learning-containers/blob/master/available_images.md#huggingface-training-containers>`_ to
-make it easier than ever to train Hugging Face Transformer models in `Amazon SageMaker
-<https://aws.amazon.com/sagemaker/>`_.
-
-To learn how to use the new 🤗 DLCs with the Amazon SageMaker to run your 🤗 Accelerate scripts and raw training loops.0
-
-
-
-Getting Started
------------------------------------------------------------------------------------------------------------------------
-
-Setup & Installation
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-
-Before you can run your 🤗 Accelerate scripts on Amazon SageMaker you need to sign up for an AWS account. If you do not
-have an AWS account yet learn more `here <https://docs.aws.amazon.com/sagemaker/latest/dg/gs-set-up.html>`__.
-
-After you have your AWS Account you need to install the ``sagemaker`` sdk for 🤗 Accelerate with.
-
-.. code-block:: 
-
-    pip install "accelerate[sagemaker]" --upgrade
-
-
-🤗 Accelerate currently uses the 🤗 DLCs, with ``transformers``, ``datasets`` and ``tokenizers`` pre-installed. 🤗
-Accelerate is not in the DLC yet (will soon be added!) so to use it within Amazon SageMaker you need to create a
-``requirements.txt`` in the same directory where your training script is located and add it as dependency.
-
-.. code-block:: 
-
-    accelerate
-
-You should also add any other dependencies you have to this ``requirements.txt``.
-
-
-Configure 🤗 Accelerate
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-You can configure the launch configuration for Amazon SageMaker the same as you do for non SageMaker training jobs with
-the 🤗 Accelerate CLI.
-
-.. code-block:: 
-
-    accelerate config
-    # In which compute environment are you running? ([0] This machine, [1] AWS (Amazon SageMaker)): 1
-
-
-🤗 Accelerate will go through a questionnaire about your Amazon SageMaker setup and create a config file you can edit.
-
-.. note::
-    🤗 Accelerate is not saving any of your credentials.
-
-
-Prepare a 🤗 Accelerate fine-tuning script
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-The training script is very similar to a training script you might run outside of SageMaker, but to save your model
-after training you need to specify either ``/opt/ml/model`` or use ``os.environ["SM_MODEL_DIR"]`` as your save
-directory. After training, artifacts in this directory are uploaded to S3.
-
-
-.. code-block:: diff
-
-    - torch.save('/opt/ml/model`)
-    + accelerator.save('/opt/ml/model')
-
-
-.. warning::
-    SageMaker doesn’t support argparse actions. If you want to use, for example, boolean hyperparameters, you need to
-    specify type as bool in your script and provide an explicit True or False value for this hyperparameter. `[REF]
-    <https://sagemaker.readthedocs.io/en/stable/frameworks/pytorch/using_pytorch.html#prepare-a-pytorch-training-script>`__.
-
-
-Launch Training
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-You can launch your training with 🤗 Accelerate CLI with
-
-.. code-block:: 
-
-    accelerate launch path_to_script.py --args_to_the_script
-
-
-This will launch your training script using your configuration. The only thing you have to do is provide all the
-arguments needed by your training script as named arguments.
-
-**Examples**
-
-.. note::
-    If you run one of the example scripts, don't forget to add ``accelerator.save('/opt/ml/model')`` to it.
-
-.. code-block:: 
-
-    accelerate launch ./examples/sagemaker_example.py 
-
-
-Outputs:
-
-.. code-block:: 
-
-    Configuring Amazon SageMaker environment
-    Converting Arguments to Hyperparameters
-    Creating Estimator
-    2021-04-08 11:56:50 Starting - Starting the training job...
-    2021-04-08 11:57:13 Starting - Launching requested ML instancesProfilerReport-1617883008: InProgress
-    .........
-    2021-04-08 11:58:54 Starting - Preparing the instances for training.........
-    2021-04-08 12:00:24 Downloading - Downloading input data
-    2021-04-08 12:00:24 Training - Downloading the training image..................
-    2021-04-08 12:03:39 Training - Training image download completed. Training in progress..
-    ........
-    epoch 0: {'accuracy': 0.7598039215686274, 'f1': 0.8178438661710037}
-    epoch 1: {'accuracy': 0.8357843137254902, 'f1': 0.882249560632689}
-    epoch 2: {'accuracy': 0.8406862745098039, 'f1': 0.8869565217391304}
-    ........
-    2021-04-08 12:05:40 Uploading - Uploading generated training model
-    2021-04-08 12:05:40 Completed - Training job completed
-    Training seconds: 331
-    Billable seconds: 331
-    You can find your model data at: s3://your-bucket/accelerate-sagemaker-1-2021-04-08-11-56-47-108/output/model.tar.gz
-
-
-
-Advanced Features
------------------------------------------------------------------------------------------------------------------------
-
-Distributed Training: Data Parallelism
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-*currently in development, will be supported soon.*
-
-Distributed Training: Model Parallelism
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-*currently in development, will be supported soon.*
-
-Python packages and dependencies
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-🤗 Accelerate currently uses the 🤗 DLCs, with ``transformers``, ``datasets`` and ``tokenizers`` pre-installed. If you
-want to use different/other Python packages you can do this by adding them to the ``requirements.txt``. These packages
-will be installed before your training script is started.
-
-Remote scripts: Use scripts located on Github
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-*undecided if feature is needed. Contact us if you would like this feature.*
-
-Use Spot Instances
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-*undecided if feature is needed. Contact us if you would like this feature.*

docs/source/usage_guides/big_modeling.mdx

@@ -0,0 +1,294 @@
+<!--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.
+-->
+
+# Handling big models
+
+When loading a pretrained 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 pretrained weights. If you're loading a model with 6 billions 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 on 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 in 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. 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.
+
+Here is how we can use this to load the [GPT-J-6B](https://huggingface.co/EleutherAI/gpt-j-6B) model. You clone the sharded version of this model with:
+
+```bash
+git clone https://huggingface.co/sgugger/sharded-gpt-j-6B
+cd sharded-gpt-j-6B
+git-lfs install
+git pull
+```
+
+then we can initialize the model with
+
+```py
+from accelerate import init_empty_weights
+from transformers import AutoConfig, AutoModelForCausalLM
+
+checkpoint = "EleutherAI/gpt-j-6B"
+config = AutoConfig.from_pretrained(checkpoint)
+
+with init_empty_weights():
+    model = AutoModelForCausalLM.from_config(config)
+```
+
+and load the checkpoint we just downloaded with:
+
+```py
+from accelerate import load_checkpoint_and_dispatch
+
+model = load_checkpoint_and_dispatch(
+    model, "sharded-gpt-j-6B", device_map="auto", no_split_module_classes=["GPTJBlock"]
+)
+```
+
+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=["GPTJBlock"]` indicates that the modules that are `GPTJBlock` should not be split on different devices. You should set here all blocks that include a residual connection of some kind.
+
+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.drop': 0,
+ 'transformer.h.0': 0,
+ 'transformer.h.1': 0,
+ 'transformer.h.2': 0,
+ 'transformer.h.3': 0,
+ 'transformer.h.4': 0,
+ 'transformer.h.5': 0,
+ 'transformer.h.6': 0,
+ 'transformer.h.7': 0,
+ 'transformer.h.8': 0,
+ 'transformer.h.9': 0,
+ 'transformer.h.10': 0,
+ 'transformer.h.11': 0,
+ 'transformer.h.12': 0,
+ 'transformer.h.13': 0,
+ 'transformer.h.14': 0,
+ 'transformer.h.15': 0,
+ 'transformer.h.16': 0,
+ 'transformer.h.17': 0,
+ 'transformer.h.18': 0,
+ 'transformer.h.19': 0,
+ 'transformer.h.20': 0,
+ 'transformer.h.21': 0,
+ 'transformer.h.22': 0,
+ 'transformer.h.23': 0,
+ 'transformer.h.24': 1,
+ 'transformer.h.25': 1,
+ 'transformer.h.26': 1,
+ 'transformer.h.27': 1,
+ 'transformer.ln_f': 1,
+ 'lm_head': 1}
+ ```
+
+You can also design your `device_map` yourself, if you prefer to explicitly decide where each layer should be. In this case, the command above becomes:
+
+```py
+model = load_checkpoint_and_dispatch(model, "sharded-gpt-j-6B", device_map=my_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 transformers import AutoTokenizer
+
+tokenizer = AutoTokenizer.from_pretrained(checkpoint)
+inputs = tokenizer("Hello, my name is", return_tensors="pt")
+inputs = inputs.to(0)
+output = model.generate(inputs["input_ids"])
+tokenizer.decode(output[0].tolist())
+```
+
+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, you 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 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 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 used 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 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 avaialble 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>
+
+## Limits and further development
+
+We are aware of the current limitations in the API:
+
+- While this could theoretically work on just one CPU with potential disk offload, you need at least one GPU to run this API. This will be fixed in further development.
+- [`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 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/usage_guides/checkpoint.mdx

@@ -0,0 +1,60 @@
+<!--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.
+-->
+
+# Checkpointing
+
+When training a PyTorch model with 🤗 Accelerate, you may often want to save and continue a state of training. Doing so requires
+saving and loading the model, optimizer, RNG generators, and the GradScaler. Inside 🤗 Accelerate are two convenience functions to achieve this quickly:
+- Use [`~Accelerator.save_state`] for saving everything mentioned above to a folder location
+- Use [`~Accelerator.load_state`] for loading everything stored from an earlier `save_state`
+
+It should be noted that the expectation is that those states come from the same training script, they should not be from two separate scripts.
+
+- By using [`~Accelerator.register_for_checkpointing`], you can register custom objects to be automatically stored or loaded from the two prior functions,
+so long as the object has a `state_dict` **and** a `load_state_dict` functionality. This could include objects such as a learning rate scheduler. 
+
+Below is a brief example using checkpointing to save and reload a state during training:
+
+```python
+from accelerate import Accelerator
+import torch
+
+accelerator = Accelerator()
+
+my_scheduler = torch.optim.lr_scheduler.StepLR(my_optimizer, step_size=1, gamma=0.99)
+my_model, my_optimizer, my_training_dataloader = accelerate.prepare(my_model, my_optimizer, my_training_dataloader)
+
+# Register the LR scheduler
+accelerate.register_for_checkpointing(my_scheduler)
+
+# Save the starting state
+accelerate.save_state("my/save/path")
+
+device = accelerator.device
+my_model.to(device)
+
+# Perform training
+for epoch in range(num_epochs):
+    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)
+        accelerator.backward(loss)
+        my_optimizer.step()
+    my_scheduler.step()
+
+# Restore previous state
+accelerate.load_state("my/save/path")
+```

docs/source/usage_guides/deepspeed.mdx

@@ -0,0 +1,494 @@
+<!--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.
+-->
+
+# DeepSpeed 
+
+[DeepSpeed](https://github.com/microsoft/DeepSpeed) implements everything described in the [ZeRO paper](https://arxiv.org/abs/1910.02054). Currently it provides full support for:
+
+1. Optimizer state partitioning (ZeRO stage 1)
+2. Gradient partitioning (ZeRO stage 2)
+3. Parameter partitioning (ZeRO stage 3)
+4. Custom mixed precision training handling
+5. A range of fast CUDA-extension-based optimizers
+6. ZeRO-Offload to CPU and Disk/NVMe
+
+ZeRO-Offload has its own dedicated paper: [ZeRO-Offload: Democratizing Billion-Scale Model Training](https://arxiv.org/abs/2101.06840). And NVMe-support is described in the paper [ZeRO-Infinity: Breaking the GPU
+Memory Wall for Extreme Scale Deep Learning](https://arxiv.org/abs/2104.07857).
+
+DeepSpeed ZeRO-2 is primarily used only for training, as its features are of no use to inference.
+
+DeepSpeed ZeRO-3 can be used for inference as well, since it allows huge models to be loaded on multiple GPUs, which
+won't be possible on a single GPU.
+
+🤗 Accelerate integrates [DeepSpeed](https://github.com/microsoft/DeepSpeed) via 2 options:
+
+1. Integration of the DeepSpeed features via `deepspeed config file` specification in `accelerate config` . You just supply your custom config file or use our template. Most of
+   this document is focused on this feature. This supports all the core features of DeepSpeed and gives user a lot of flexibility. 
+   User may have to change few lines of code depending on the config.
+2. Integration via `deepspeed_plugin`.This supports subset of the DeepSpeed features and uses default options for the rest of the configurations. 
+   User need not change any code and is good for those who are fine with most of the default settings of DeepSpeed.
+
+## What is integrated?
+
+Training:
+
+1. DeepSpeed ZeRO training supports the full ZeRO stages 1, 2 and 3 as well as CPU/Disk offload of optimizer states, gradients and parameters. 
+Below is a short description of Data Parallelism using ZeRO - Zero Redundancy Optimizer along with diagram from this [blog post](https://www.microsoft.com/en-us/research/blog/zero-deepspeed-new-system-optimizations-enable-training-models-with-over-100-billion-parameters/)
+![ZeRO Data Parallelism](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/parallelism-zero.png)
+
+(Source: [link](https://www.microsoft.com/en-us/research/blog/zero-deepspeed-new-system-optimizations-enable-training-models-with-over-100-billion-parameters/))
+
+ a. **Stage 1** : Shards optimizer states across data parallel workers/GPUs
+
+ b. **Stage 2** : Shards optimizer states + gradients across data parallel workers/GPUs
+
+ c. **Stage 3**: Shards optimizer states + gradients + model parameters across data parallel workers/GPUs
+
+ d. **Optimizer Offload**: Offloads the gradients + optimizer states to CPU/Disk building on top of ZERO Stage 2
+
+ e. **Param Offload**: Offloads the model parameters to CPU/Disk building on top of ZERO Stage 3
+
+<u>Note</u>: With respect to Disk Offload, the disk should be an NVME for decent speed but it technically work on any Disk
+
+Inference:
+
+1. DeepSpeed ZeRO Inference supports ZeRO stage 3 with ZeRO-Infinity. It uses the same ZeRO protocol as training, but
+   it doesn't use an optimizer and a lr scheduler and only stage 3 is relevant. For more details see:
+   [deepspeed-zero-inference](#deepspeed-zero-inference).
+
+
+## How it works?
+
+**Pre-Requisites**: Install DeepSpeed version >=0.6.5. Please refer to the [DeepSpeed Installation details](https://github.com/microsoft/DeepSpeed#installation)
+for more information.
+
+We will first look at easy to use integration via `accelerate config`. 
+Followed by more flexible and feature rich `deepspeed config file` integration. 
+
+### Accelerate DeepSpeed Plugin
+On your machine(s) just run:
+
+```bash
+accelerate config
+```
+
+and answer the questions asked. It will ask whether you want to use a config file for DeepSpeed to which you should answer no. Then answer the following questions to generate a basic DeepSpeed config.
+This will generate a config file that will be used automatically to properly set the
+default options when doing
+
+```bash
+accelerate launch my_script.py --args_to_my_script
+```
+
+For instance, here is how you would run the NLP example `examples/nlp_example.py` (from the root of the repo) with DeepSpeed Plugin:
+
+**ZeRO Stage-2 DeepSpeed Plugin Example**
+```bash
+compute_environment: LOCAL_MACHINE
+deepspeed_config:
+ gradient_accumulation_steps: 1
+ gradient_clipping: 1.0
+ offload_optimizer_device: none
+ offload_param_device: none
+ zero3_init_flag: true
+ zero_stage: 2
+distributed_type: DEEPSPEED
+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
+```
+
+```bash
+accelerate launch examples/nlp_example.py --mixed_precision fp16
+```
+
+**ZeRO Stage-3 with CPU Offload DeepSpeed Plugin Example**
+```bash
+compute_environment: LOCAL_MACHINE
+deepspeed_config:
+  gradient_accumulation_steps: 1
+  gradient_clipping: 1.0
+  offload_optimizer_device: cpu
+  offload_param_device: cpu
+  zero3_init_flag: true
+  zero3_save_16bit_model: true
+  zero_stage: 3
+distributed_type: DEEPSPEED
+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
+```
+
+```bash
+accelerate launch examples/nlp_example.py --mixed_precision fp16
+```
+
+Currently, `Accelerate` supports following config through the CLI:
+
+```bash
+`zero_stage`: [0] Disabled, [1] optimizer state partitioning, [2] optimizer+gradient state partitioning and [3] optimizer+gradient+parameter partitioning
+`gradient_accumulation_steps`: Number of training steps to accumulate gradients before averaging and applying them.
+`gradient_clipping`: Enable gradient clipping with value.
+`offload_optimizer_device`: [none] Disable optimizer offloading, [cpu] offload optimizer to CPU, [nvme] offload optimizer to NVMe SSD. Only applicable with ZeRO >= Stage-2.
+`offload_param_device`: [none] Disable parameter offloading, [cpu] offload parameters to CPU, [nvme] offload parameters to NVMe SSD. Only applicable with ZeRO Stage-3.
+`zero3_init_flag`: Decides whether to enable `deepspeed.zero.Init` for constructing massive models. Only applicable with ZeRO Stage-3.
+`zero3_save_16bit_model`: Decides whether to save 16-bit model weights when using ZeRO Stage-3.
+`mixed_precision`: `no` for FP32 training, `fp16` for FP16 mixed-precision training and `bf16` for BF16 mixed-precision training. 
+```
+To be able to tweak more options, you will need to use a DeepSpeed config file.
+
+### DeepSpeed Config File
+On your machine(s) just run:
+
+```bash
+accelerate config
+```
+
+and answer the questions asked. It will ask whether you want to use a config file for deepspeed to which you answer yes 
+and provide the path to the deepspeed config file. 
+This will generate a config file that will be used automatically to properly set the
+default options when doing
+
+```bash
+accelerate launch my_script.py --args_to_my_script
+```
+
+For instance, here is how you would run the NLP example `examples/by_feature/deepspeed_with_config_support.py` (from the root of the repo) with DeepSpeed Config File:
+
+**ZeRO Stage-2 DeepSpeed Config File Example**
+```bash
+compute_environment: LOCAL_MACHINE
+deepspeed_config:
+ deepspeed_config_file: /home/ubuntu/accelerate/examples/configs/deepspeed_config_templates/zero_stage2_config.json
+ zero3_init_flag: true
+distributed_type: DEEPSPEED
+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
+```
+
+with the contents of `zero_stage2_config.json` being:
+```json
+{
+    "fp16": {
+        "enabled": true,
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+    "optimizer": {
+        "type": "AdamW",
+        "params": {
+            "lr": "auto",
+            "weight_decay": "auto",
+            "torch_adam": true,
+            "adam_w_mode": true
+        }
+    },
+    "scheduler": {
+        "type": "WarmupDecayLR",
+        "params": {
+            "warmup_min_lr": "auto",
+            "warmup_max_lr": "auto",
+            "warmup_num_steps": "auto",
+            "total_num_steps": "auto"
+        }
+    },
+    "zero_optimization": {
+        "stage": 2,
+        "allgather_partitions": true,
+        "allgather_bucket_size": 2e8,
+        "overlap_comm": true,
+        "reduce_scatter": true,
+        "reduce_bucket_size": "auto",
+        "contiguous_gradients": true
+    },
+    "gradient_accumulation_steps": 1,
+    "gradient_clipping": "auto",
+    "steps_per_print": 2000,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}
+```
+
+```bash
+accelerate launch examples/by_feature/deepspeed_with_config_support.py \
+--config_name "gpt2-large" \
+--tokenizer_name "gpt2-large" \
+--dataset_name "wikitext" \
+--dataset_config_name "wikitext-2-raw-v1" \
+--block_size 128 \
+--output_dir "./clm/clm_deepspeed_stage2_accelerate" \
+--learning_rate 5e-4 \
+--per_device_train_batch_size 24 \
+--per_device_eval_batch_size 24 \
+--num_train_epochs 3 \
+--with_tracking \
+--report_to "wandb"\
+```
+
+**ZeRO Stage-3 with CPU offload DeepSpeed Config File Example**
+```bash
+compute_environment: LOCAL_MACHINE
+deepspeed_config:
+ deepspeed_config_file: /home/ubuntu/accelerate/examples/configs/deepspeed_config_templates/zero_stage3_offload_config.json
+ zero3_init_flag: true
+distributed_type: DEEPSPEED
+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
+```
+with the contents of `zero_stage3_offload_config.json` being:
+```json
+{
+    "fp16": {
+        "enabled": true,
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+    "optimizer": {
+        "type": "AdamW",
+        "params": {
+            "lr": "auto",
+            "weight_decay": "auto"
+        }
+    },
+    "scheduler": {
+        "type": "WarmupDecayLR",
+        "params": {
+            "warmup_min_lr": "auto",
+            "warmup_max_lr": "auto",
+            "warmup_num_steps": "auto",
+            "total_num_steps": "auto"
+        }
+    },
+    "zero_optimization": {
+        "stage": 3,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "offload_param": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "overlap_comm": true,
+        "contiguous_gradients": true,
+        "reduce_bucket_size": "auto",
+        "stage3_prefetch_bucket_size": "auto",
+        "stage3_param_persistence_threshold": "auto",
+        "sub_group_size": 1e9,
+        "stage3_max_live_parameters": 1e9,
+        "stage3_max_reuse_distance": 1e9,
+        "stage3_gather_16bit_weights_on_model_save": "auto"
+    },
+    "gradient_accumulation_steps": 1,
+    "gradient_clipping": "auto",
+    "steps_per_print": 2000,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}
+```
+
+```bash
+accelerate launch examples/by_feature/deepspeed_with_config_support.py \
+--config_name "gpt2-large" \
+--tokenizer_name "gpt2-large" \
+--dataset_name "wikitext" \
+--dataset_config_name "wikitext-2-raw-v1" \
+--block_size 128 \
+--output_dir "./clm/clm_deepspeed_stage3_offload_accelerate" \
+--learning_rate 5e-4 \
+--per_device_train_batch_size 32 \
+--per_device_eval_batch_size 32 \
+--num_train_epochs 3 \
+--with_tracking \
+--report_to "wandb"\
+```
+
+**Important code changes when using DeepSpeed Config File**
+
+1. DeepSpeed Optimizers and Schedulers. For more information on these, 
+see the [DeepSpeed Optimizers](https://deepspeed.readthedocs.io/en/latest/optimizers.html) and [DeepSpeed Schedulers](https://deepspeed.readthedocs.io/en/latest/schedulers.html) documentation.
+We will look at the changes needed in the code when using these.
+   
+   a. DS Optim + DS Scheduler: The case when both `optimizer` and `scheduler` keys present in the DeepSpeed config file.
+   In this situation, those will be used and user has to use `accelerate.utils.DummyOptim` and `accelerate.utils.DummyScheduler` to replace the PyTorch/Custom optimizers and schedulers in their code.
+   Below is the snippet from `examples/by_feature/deepspeed_with_config_support.py` showing this:
+   ```python
+    # Creates Dummy Optimizer if `optimizer` was spcified in the config file else creates Adam Optimizer
+    optimizer_cls = (
+        torch.optim.AdamW
+        if accelerator.state.deepspeed_plugin is None
+        or "optimizer" not in accelerator.state.deepspeed_plugin.deepspeed_config
+        else DummyOptim
+    )
+    optimizer = optimizer_cls(optimizer_grouped_parameters, lr=args.learning_rate)
+
+    # Creates Dummy Scheduler if `scheduler` was spcified in the config file else creates `args.lr_scheduler_type` Scheduler
+    if (
+        accelerator.state.deepspeed_plugin is None
+        or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config
+    ):
+        lr_scheduler = get_scheduler(
+            name=args.lr_scheduler_type,
+            optimizer=optimizer,
+            num_warmup_steps=args.num_warmup_steps,
+            num_training_steps=args.max_train_steps,
+        )
+    else:
+        lr_scheduler = DummyScheduler(
+            optimizer, total_num_steps=args.max_train_steps, warmup_num_steps=args.num_warmup_steps
+        )
+   ```
+   b. Custom Optim + Custom Scheduler: The case when both `optimizer` and `scheduler` keys are absent in the DeepSpeed config file.
+   In this situation, no code changes are needed from the user and this is the case when using integration via DeepSpeed Plugin.
+   In the above example we can see that the code remains unchanged if the `optimizer` and `scheduler` keys are absent in the DeepSpeed config file.
+
+   c. Custom Optim + DS Scheduler: The case when only `scheduler` key is present in the DeepSpeed config file. 
+   In this situation, user has to use `accelerate.utils.DummyScheduler` to replace the PyTorch/Custom scheduler in their code. 
+
+   d. DS Optim + Custom Scheduler: The case when only `optimizer` key is present in the DeepSpeed config file. 
+   This will result in an error because you can only use DS Scheduler when using DS Optim.
+
+2. Notice the `auto` values in the above example DeepSpeed config files. These are automatically handled by `prepare` method 
+based on model, dataloaders, dummy optimizer and dummy schedulers provided to `prepare` method. 
+Only the `auto` fields specified in above examples are handled by `prepare` method and the rest have to be explicitly specified by the user.
+
+## Saving and loading
+
+1. Saving and loading of models is unchanged for ZeRO Stage-1 and Stage-2.
+
+2. under ZeRO Stage-3, `state_dict` contains just the placeholders since the model weights are partitioned across multiple GPUs.
+ZeRO Stage-3 has 2 options:
+
+   a. Saving the entire 16bit model weights to directly load later on using `model.load_state_dict(torch.load(pytorch_model.bin))`.
+   For this, either set `zero_optimization.stage3_gather_16bit_weights_on_model_save` to True in DeepSpeed Config file or set
+   `zero3_save_16bit_model` to True in DeepSpeed Plugin. 
+   **Note that this option requires consolidation of the weights on one GPU it can be slow and memory demanding, so only use this feature when needed.**
+   Below is the snippet from `examples/by_feature/deepspeed_with_config_support.py` showing this:
+   ```python
+   unwrapped_model = accelerator.unwrap_model(model)
+
+   # New Code #
+   # Saves the whole/unpartitioned fp16 model when in ZeRO Stage-3 to the output directory if
+   # `stage3_gather_16bit_weights_on_model_save` is True in DeepSpeed Config file or
+   # `zero3_save_16bit_model` is True in DeepSpeed Plugin.
+   # For Zero Stages 1 and 2, models are saved as usual in the output directory.
+   # The model name saved is `pytorch_model.bin`
+   unwrapped_model.save_pretrained(
+       args.output_dir,
+       is_main_process=accelerator.is_main_process,
+       save_function=accelerator.save,
+       state_dict=accelerator.get_state_dict(model),
+   )
+   ```
+
+   b. To get 32bit weights, first save the model using `model.save_checkpoint()`.
+   Below is the snippet from `examples/by_feature/deepspeed_with_config_support.py` showing this:
+   ```python
+   success = model.save_checkpoint(PATH, ckpt_id, checkpoint_state_dict)
+   status_msg = "checkpointing: PATH={}, ckpt_id={}".format(PATH, ckpt_id)
+   if success:
+       logging.info(f"Success {status_msg}")
+   else:
+       logging.warning(f"Failure {status_msg}")
+   ``` 
+   This will create ZeRO model and optimizer partitions along with `zero_to_fp32.py` script in checkpoint directory.
+   You can use this script to do offline consolidation.  
+   It requires no configuration files or GPUs. Here is an example of its usage:  
+   ```bash
+   $ cd /path/to/checkpoint_dir
+   $ ./zero_to_fp32.py . pytorch_model.bin
+   Processing zero checkpoint at global_step1
+   Detected checkpoint of type zero stage 3, world_size: 2
+   Saving fp32 state dict to pytorch_model.bin (total_numel=60506624)
+   ```
+   To get 32bit model for saving/inference, you can perform:
+   ```python
+   from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint
+
+   unwrapped_model = accelerator.unwrap_model(model)
+   fp32_model = load_state_dict_from_zero_checkpoint(unwrapped_model, checkpoint_dir)
+   ```
+   If you are only interested in the `state_dict`, you can do the following:
+   ```python
+   from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
+
+   state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir)
+   ```
+   Note that all these functions require ~2x memory (general RAM) of the size of the final checkpoint.
+
+## ZeRO Inference
+DeepSpeed ZeRO Inference supports ZeRO stage 3 with ZeRO-Infinity. 
+It uses the same ZeRO protocol as training, but it doesn't use an optimizer and a lr scheduler and only stage 3 is relevant.
+With accelerate integration, you just need to prepare the model and dataloader as shown below:
+
+```python
+model, eval_dataloader = accelerator.prepare(model, eval_dataloader)
+```
+
+## Few caveats to be aware of 
+
+1. Current integration doesn’t support Pipeline Parallelism of DeepSpeed.
+2. Current integration doesn’t support `mpu`, limiting the tensor parallelism which is supported in Megatron-LM. 
+3. Current integration doesn’t support multiple models. 
+
+## DeepSpeed Resources
+
+The documentation for the internals related to deepspeed can be found [here](../package_reference/deepspeed).
+
+- [Project's github](https://github.com/microsoft/deepspeed)
+- [Usage docs](https://www.deepspeed.ai/getting-started/)
+- [API docs](https://deepspeed.readthedocs.io/en/latest/index.html)
+- [Blog posts](https://www.microsoft.com/en-us/research/search/?q=deepspeed)
+
+Papers:
+
+- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://arxiv.org/abs/1910.02054)
+- [ZeRO-Offload: Democratizing Billion-Scale Model Training](https://arxiv.org/abs/2101.06840)
+- [ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning](https://arxiv.org/abs/2104.07857)
+
+Finally, please, remember that, 🤗 `Accelerate` only integrates DeepSpeed, therefore if you
+have any problems or questions with regards to DeepSpeed usage, please, file an issue with [DeepSpeed GitHub](https://github.com/microsoft/DeepSpeed/issues).
+

docs/source/usage_guides/fsdp.mdx

@@ -0,0 +1,125 @@
+<!--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.
+-->
+
+# Fully Sharded Data Parallel
+
+To accelerate training huge models on larger batch sizes, we can use a fully sharded data parallel model.
+This type of data parallel paradigm enables fitting more data and larger models by sharding the optimizer states, gradients and parameters.
+To read more about it and the benefits, check out the [Fully Sharded Data Parallel blog](https://pytorch.org/blog/introducing-pytorch-fully-sharded-data-parallel-api/).
+We have integrated the latest PyTorch's Fully Sharded Data Parallel (FSDP) training feature.
+All you need to do is enable it through the config.
+
+## How it works out of the box
+
+On your machine(s) just run:
+
+```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
+
+```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) with FSDP enabled:
+
+```bash
+compute_environment: LOCAL_MACHINE
+deepspeed_config: {}
+distributed_type: FSDP
+downcast_bf16: 'no'
+fsdp_config:
+  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  fsdp_backward_prefetch_policy: BACKWARD_PRE
+  fsdp_offload_params: false
+  fsdp_sharding_strategy: 1
+  fsdp_state_dict_type: FULL_STATE_DICT
+  fsdp_transformer_layer_cls_to_wrap: GPT2Block
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+mixed_precision: 'no'
+num_machines: 1
+num_processes: 2
+use_cpu: false
+```
+
+```bash
+accelerate launch examples/nlp_example.py
+```
+
+Currently, `Accelerate` supports the following config through the CLI:
+
+```bash
+`Sharding Strategy`: [1] FULL_SHARD (shards optimizer states, gradients and parameters), [2] SHARD_GRAD_OP (shards optimizer states and gradients), [3] NO_SHARD
+`Offload Params`: Decides Whether to offload parameters and gradients to CPU
+`Auto Wrap Policy`: [1] TRANSFORMER_BASED_WRAP, [2] SIZE_BASED_WRAP, [3] NO_WRAP 
+`Transformer Layer Class to Wrap`: When using `TRANSFORMER_BASED_WRAP`, user specifies transformer layer class name (case-sensitive) to wrap ,e.g, `BertLayer`, `GPTJBlock`, `T5Block`...
+`Min Num Params`: minimum number of parameters when using `SIZE_BASED_WRAP`
+`Backward Prefetch`: [1] BACKWARD_PRE, [2] BACKWARD_POST, [3] NO_PREFETCH
+`State Dict Type`: [1] FULL_STATE_DICT, [2] LOCAL_STATE_DICT, [3] SHARDED_STATE_DICT  
+```
+
+## A few caveats to be aware of
+
+- PyTorch FSDP auto wraps sub-modules, flattens the parameters and shards the parameters in place.
+  Due to this, any optimizer created before model wrapping gets broken and occupies more memory.
+  Hence, it is highly recommended and efficient to prepare the model before creating the optimizer.
+  `Accelerate` will automatically wrap the model and create an optimizer for you in case of single model with a warning message.
+  > FSDP Warning: When using FSDP, it is efficient and recommended to call prepare for the model before creating the optimizer
+
+However, below is the recommended way to prepare model and optimizer while using FSDP:
+
+```diff
+  model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", return_dict=True)
++ model = accelerator.prepare(model)
+
+  optimizer = torch.optim.AdamW(params=model.parameters(), lr=lr)
+
+- model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+-        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+-    )
+
++ optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
++         optimizer, train_dataloader, eval_dataloader, lr_scheduler
++    )
+```
+
+- In case of a single model, if you have created the optimizer with multiple parameter groups and called prepare with them together,
+  then the parameter groups will be lost and the following warning is displayed:
+  > FSDP Warning: When using FSDP, several parameter groups will be conflated into
+  > a single one due to nested module wrapping and parameter flattening.
+  
+  This is because parameter groups created before wrapping will have no meaning post wrapping due to parameter flattening of nested FSDP modules into 1D arrays (which can consume many layers).
+  For instance, below are the named parameters of an FSDP model on GPU 0 (When using 2 GPUs. Around 55M (110M/2) params in 1D arrays as this will have the 1st shard of the parameters). 
+  Here, if one has applied no weight decay for [bias, LayerNorm.weight] the named parameters of an unwrapped BERT model, 
+  it can't be applied to the below FSDP wrapped model as there are no named parameters with either of those strings and 
+  the parameters of those layers are concatenated with parameters of various other layers.
+  ```
+  {
+    '_fsdp_wrapped_module.flat_param': torch.Size([494209]), 
+    '_fsdp_wrapped_module._fpw_module.bert.embeddings.word_embeddings._fsdp_wrapped_module.flat_param': torch.Size([11720448]), 
+    '_fsdp_wrapped_module._fpw_module.bert.encoder._fsdp_wrapped_module.flat_param': torch.Size([42527232])
+  }
+  ```
+
+
+- In case of multiple models, it is necessary to prepare the models before creating optimizers or else it will throw an error. 
+Then pass the optimizers to the prepare call in the same order as corresponding models else `accelerator.save_state()` and `accelerator.load_state()` will result in wrong/unexpected behaviour.
+- This feature is incompatible with `--predict_with_generate` in the `run_translation.py` script of 🤗 `Transformers` library.
+
+For more control, users can leverage the `FullyShardedDataParallelPlugin`. After creating an instance of this class, users can pass it to the Accelerator class instantiation.
+For more information on these options, please refer to the PyTorch [FullyShardedDataParallel](https://github.com/pytorch/pytorch/blob/0df2e863fbd5993a7b9e652910792bd21a516ff3/torch/distributed/fsdp/fully_sharded_data_parallel.py#L236) code.

docs/source/usage_guides/gradient_accumulation.mdx

@@ -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.
+-->
+
+# Performing gradient accumulation with 🤗 Accelerate
+
+Gradient accumulation is a technique where you can train on bigger batch sizes than 
+your machine would normally be able to fit into memory. This is done by accumulating gradients over
+several batches, and only stepping the optimizer after a certain number of batches have been performed.
+
+While technically standard gradient accumulation code would work fine in a distributed setup, it is not the most efficient
+method for doing so and you may experience considerable slowdowns!
+
+In this tutorial you will see how to quickly setup gradient accumulation and perform it with the utilities provided in 🤗 Accelerate,
+which can total to adding just one new line of code!
+
+This example will use a very simplistic PyTorch training loop that performs gradient accumulation every two batches:
+
+```python
+device = "cuda"
+model.to(device)
+
+gradient_accumulation_steps = 2
+
+for index, batch in enumerate(training_dataloader):
+    inputs, targets = batch
+    inputs = inputs.to(device)
+    targets = targets.to(device)
+    outputs = model(inputs)
+    loss = loss_function(outputs, targets)
+    loss = loss / gradient_accumulation_steps
+    loss.backward()
+    if (index + 1) % gradient_accumulation_steps == 0:
+        optimizer.step()
+        scheduler.step()
+        optimizer.zero_grad()
+```
+
+## Converting it to 🤗 Accelerate
+
+First the code shown earlier will be converted to utilize 🤗 Accelerate without the special gradient accumulation helper:
+
+```diff
++ from accelerate import Accelerator
++ accelerator = Accelerator()
+
++ model, optimizer, training_dataloader, scheduler = accelerator.prepare(
++     model, optimizer, training_dataloader, scheduler
++ )
+
+  for index, batch in enumerate(training_dataloader):
+      inputs, targets = batch
+-     inputs = inputs.to(device)
+-     targets = targets.to(device)
+      outputs = model(inputs)
+      loss = loss_function(outputs, targets)
+      loss = loss / gradient_accumulation_steps
++     accelerator.backward(loss)
+      if (index+1) % gradient_accumulation_steps == 0:
+          optimizer.step()
+          scheduler.step()
+          optimizer.zero_grad()
+```
+
+<Tip warning={true}>
+
+  In its current state, this code is not going to perform gradient accumulation efficiently due to a process called gradient synchronization. Read more about that in the [Concepts tutorial](concept_guides/gradient_synchronization)!
+
+</Tip>
+
+## Letting 🤗 Accelerate handle gradient accumulation
+
+All that is left now is to let 🤗 Accelerate handle the gradient accumulation for us. To do so you should pass in a `gradient_accumulation_steps` parameter to [`Accelerator`], dictating the number 
+of steps to perform before each call to `step()` and how to automatically adjust the loss during the call to [`~Accelerator.backward`]:
+
+```diff
+  from accelerate import Accelerator
+- accelerator = Accelerator()
++ accelerator = Accelerator(gradient_accumulation_steps=2)
+```
+
+From here you can use the [`~Accelerator.accumulate`] context manager from inside your training loop to automatically perform the gradient accumulation for you!
+You just wrap it around the entire training part of our code: 
+
+```diff
+- for index, batch in enumerate(training_dataloader):
++ for batch in training_dataloader:
++     with accelerator.accumulate(model):
+          inputs, targets = batch
+          outputs = model(inputs)
+```
+
+You can remove all the special checks for the step number and the loss adjustment:
+
+```diff
+- loss = loss / gradient_accumulation_steps
+  accelerator.backward(loss)
+- if (index+1) % gradient_accumulation_steps == 0:
+  optimizer.step()
+  scheduler.step()
+  optimizer.zero_grad()
+```
+
+As you can see the [`Accelerator`] is able to keep track of the batch number you are on and it will automatically know whether to step through the prepared optimizer and how to adjust the loss. 
+
+## The finished code
+
+Below is the finished implementation for performing gradient accumulation with 🤗 Accelerate
+
+```python
+for batch in training_dataloader:
+    with accelerator.accumulate(model):
+        inputs, targets = batch
+        outputs = model(inputs)
+        loss = loss_function(outputs, targets)
+        accelerator.backward(loss)
+        optimizer.step()
+        scheduler.step()
+        optimizer.zero_grad()
+```
+
+To learn more about what magic this wraps around, read the [Gradient Synchronization concept guide](/concept_guides/gradient_synchronization)

docs/source/usage_guides/memory.mdx

@@ -0,0 +1,55 @@
+<!--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.
+-->
+
+# Memory Utilities
+
+One of the most frustrating errors when it comes to running training scripts is hitting "CUDA Out-of-Memory", 
+as the entire script needs to be restarted, progress is lost, and typically a developer would want to simply
+start their script and let it run.
+
+`Accelerate` provides a utility heavily based on [toma](https://github.com/BlackHC/toma) to give this capability.
+
+## find_executable_batch_size
+
+This algorithm operates with exponential decay, decreasing the batch size in half after each failed run on some 
+training script. To use it, restructure your training function to include an inner function that includes this wrapper, 
+and build your dataloaders inside it. At a minimum, this could look like 4 new lines of code. 
+> Note: The inner function *must* take in the batch size as the first parameter, but we do not pass one to it when called. The wrapper handles this for us
+
+It should also be noted that anything which will consume CUDA memory and passed to the `accelerator` **must** be declared inside the inner function,
+such as models and optimizers.
+
+```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()
+```
+
+To find out more, check the documentation [here](../package_reference/utilities#accelerate.find_executable_batch_size).

docs/source/usage_guides/mps.mdx

@@ -0,0 +1,82 @@
+<!--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.
+-->
+
+# Accelerated PyTorch Training on Mac
+
+With PyTorch v1.12 release, developers and researchers can take advantage of Apple silicon GPUs for significantly faster model training. 
+This unlocks the ability to perform machine learning workflows like prototyping and fine-tuning locally, right on Mac.
+Apple's Metal Performance Shaders (MPS) as a backend for PyTorch enables this and can be used via the new `"mps"` device. 
+This will map computational graphs and primitives on the MPS Graph framework and tuned kernels provided by MPS.
+For more information please refer official documents [Introducing Accelerated PyTorch Training on Mac](https://pytorch.org/blog/introducing-accelerated-pytorch-training-on-mac/)
+and [MPS BACKEND](https://pytorch.org/docs/stable/notes/mps.html).
+
+### Benefits of Training and Inference using Apple Silicon Chips
+
+1. Enables users to train larger networks or batch sizes locally
+2. Reduces data retrieval latency and provides the GPU with direct access to the full memory store due to unified memory architecture. 
+Therefore, improving end-to-end performance.
+3. Reduces costs associated with cloud-based development or the need for additional local GPUs.
+
+**Pre-requisites**: To install torch with mps support, 
+please follow this nice medium article [GPU-Acceleration Comes to PyTorch on M1 Macs](https://medium.com/towards-data-science/gpu-acceleration-comes-to-pytorch-on-m1-macs-195c399efcc1).
+
+
+## How it works out of the box
+
+On your machine(s) just run:
+
+```bash
+accelerate config
+```
+
+and answer the questions asked, specifically choose `MPS` for the query:
+
+```
+ Which type of machine are you using?. 
+ ```
+
+This will generate a config file that will be used automatically to properly set 
+the default options when doing `accelerate launch`, such as the one shown below:
+
+```bash
+compute_environment: LOCAL_MACHINE
+deepspeed_config: {}
+distributed_type: MPS
+downcast_bf16: 'no'
+fsdp_config: {}
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+mixed_precision: 'no'
+num_machines: 1
+num_processes: 1
+use_cpu: false
+```
+
+After this configuration has been made, here is how you run the CV example 
+(from the root of the repo) with MPS enabled:
+
+```bash
+accelerate launch /examples/cv_example.py --data_dir images
+```
+
+## A few caveats to be aware of
+
+1. We strongly recommend to install PyTorch >= 1.13 (nightly version at the time of writing) on your MacOS machine. 
+It has major fixes related to model correctness and performance improvements for transformer based models.
+Please refer to https://github.com/pytorch/pytorch/issues/82707 for more details.
+2. Distributed setups `gloo` and `nccl` are not working with `mps` device. 
+This means that currently only single GPU of `mps` device type can be used.
+
+Finally, please, remember that, 🤗 `Accelerate` only integrates MPS backend, therefore if you
+have any problems or questions with regards to MPS backend usage, please, file an issue with [PyTorch GitHub](https://github.com/pytorch/pytorch/issues).

docs/source/usage_guides/sagemaker.mdx

@@ -0,0 +1,169 @@
+<!--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.
+-->
+
+# Amazon SageMaker
+
+Hugging Face and Amazon introduced new [Hugging Face Deep Learning Containers (DLCs)](https://github.com/aws/deep-learning-containers/blob/master/available_images.md#huggingface-training-containers) to
+make it easier than ever to train Hugging Face Transformer models in [Amazon SageMaker](https://aws.amazon.com/sagemaker/).
+
+## Getting Started
+
+### Setup & Installation
+
+
+Before you can run your 🤗 Accelerate scripts on Amazon SageMaker you need to sign up for an AWS account. If you do not
+have an AWS account yet learn more [here](https://docs.aws.amazon.com/sagemaker/latest/dg/gs-set-up.html).
+
+After you have your AWS Account you need to install the `sagemaker` sdk for 🤗 Accelerate with:
+
+```bash
+pip install "accelerate[sagemaker]" --upgrade
+```
+
+🤗 Accelerate currently uses the 🤗 DLCs, with `transformers`, `datasets` and `tokenizers` pre-installed. 🤗
+Accelerate is not in the DLC yet (will soon be added!) so to use it within Amazon SageMaker you need to create a
+`requirements.txt` in the same directory where your training script is located and add it as dependency:
+
+```
+accelerate
+```
+
+You should also add any other dependencies you have to this `requirements.txt`.
+
+
+### Configure 🤗 Accelerate
+
+You can configure the launch configuration for Amazon SageMaker the same as you do for non SageMaker training jobs with
+the 🤗 Accelerate CLI:
+
+```bash
+accelerate config
+# In which compute environment are you running? ([0] This machine, [1] AWS (Amazon SageMaker)): 1
+```
+
+🤗 Accelerate will go through a questionnaire about your Amazon SageMaker setup and create a config file you can edit.
+
+<Tip>
+
+    🤗 Accelerate is not saving any of your credentials.
+
+</Tip>
+
+### Prepare a 🤗 Accelerate fine-tuning script
+
+The training script is very similar to a training script you might run outside of SageMaker, but to save your model
+after training you need to specify either `/opt/ml/model` or use `os.environ["SM_MODEL_DIR"]` as your save
+directory. After training, artifacts in this directory are uploaded to S3:
+
+
+```diff
+- torch.save('/opt/ml/model`)
++ accelerator.save('/opt/ml/model')
+```
+
+<Tip warning={true}>
+
+    SageMaker doesn’t support argparse actions. If you want to use, for example, boolean hyperparameters, you need to
+    specify type as bool in your script and provide an explicit True or False value for this hyperparameter. [[REF]](https://sagemaker.readthedocs.io/en/stable/frameworks/pytorch/using_pytorch.html#prepare-a-pytorch-training-script).
+
+</Tip>
+
+### Launch Training
+
+You can launch your training with 🤗 Accelerate CLI with:
+
+```
+accelerate launch path_to_script.py --args_to_the_script
+```
+
+This will launch your training script using your configuration. The only thing you have to do is provide all the
+arguments needed by your training script as named arguments.
+
+**Examples**
+
+<Tip>
+
+    If you run one of the example scripts, don't forget to add `accelerator.save('/opt/ml/model')` to it.
+
+</Tip>
+
+```bash
+accelerate launch ./examples/sagemaker_example.py
+```
+
+Outputs:
+
+```
+Configuring Amazon SageMaker environment
+Converting Arguments to Hyperparameters
+Creating Estimator
+2021-04-08 11:56:50 Starting - Starting the training job...
+2021-04-08 11:57:13 Starting - Launching requested ML instancesProfilerReport-1617883008: InProgress
+.........
+2021-04-08 11:58:54 Starting - Preparing the instances for training.........
+2021-04-08 12:00:24 Downloading - Downloading input data
+2021-04-08 12:00:24 Training - Downloading the training image..................
+2021-04-08 12:03:39 Training - Training image download completed. Training in progress..
+........
+epoch 0: {'accuracy': 0.7598039215686274, 'f1': 0.8178438661710037}
+epoch 1: {'accuracy': 0.8357843137254902, 'f1': 0.882249560632689}
+epoch 2: {'accuracy': 0.8406862745098039, 'f1': 0.8869565217391304}
+........
+2021-04-08 12:05:40 Uploading - Uploading generated training model
+2021-04-08 12:05:40 Completed - Training job completed
+Training seconds: 331
+Billable seconds: 331
+You can find your model data at: s3://your-bucket/accelerate-sagemaker-1-2021-04-08-11-56-47-108/output/model.tar.gz
+```
+
+## Advanced Features
+
+### Distributed Training: Data Parallelism
+
+Set up the accelerate config by running `accelerate config` and answer the SageMaker questions and set it up.
+To use SageMaker DDP, select it when asked 
+`What is the distributed mode? ([0] No distributed training, [1] data parallelism):`.
+Example config below:
+```yaml
+base_job_name: accelerate-sagemaker-1
+compute_environment: AMAZON_SAGEMAKER
+distributed_type: DATA_PARALLEL
+ec2_instance_type: ml.p3.16xlarge
+iam_role_name: xxxxx
+image_uri: null
+mixed_precision: fp16
+num_machines: 1
+profile: xxxxx
+py_version: py38
+pytorch_version: 1.10.2
+region: us-east-1
+transformers_version: 4.17.0
+use_cpu: false
+```
+
+### Distributed Training: Model Parallelism
+
+*currently in development, will be supported soon.*
+
+### Python packages and dependencies
+
+🤗 Accelerate currently uses the 🤗 DLCs, with `transformers`, `datasets` and `tokenizers` pre-installed. If you
+want to use different/other Python packages you can do this by adding them to the `requirements.txt`. These packages
+will be installed before your training script is started.
+
+### Remote scripts: Use scripts located on Github
+
+*undecided if feature is needed. Contact us if you would like this feature.*
+
+### Use Spot Instances
+
+*undecided if feature is needed. Contact us if you would like this feature.*

docs/source/usage_guides/tracking.mdx

@@ -0,0 +1,200 @@
+<!--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.
+-->
+
+# Tracking
+
+There are a large number of experiment tracking API's available, however getting them all to work with in a multi-processing environment can oftentimes be complex.
+🤗 Accelerate provides a general tracking API that can be used to log useful items during your script through [`Accelerator.log`]
+
+## Integrated Trackers
+
+Currently `Accelerate` supports three trackers out-of-the-box:
+
+- TensorBoard
+- WandB
+- CometML
+
+To use any of them, pass in the selected type(s) to the `log_with` parameter in [`Accelerate`]:
+```python
+from accelerate import Accelerator
+from accelerate.utils import LoggerType
+
+accelerator = Accelerator(log_with="all")  # For all available trackers in the environment
+accelerator = Accelerator(log_with="wandb")
+accelerator = Accelerator(log_with=["wandb", LoggerType.TENSORBOARD])
+```
+
+At the start of your experiment [`Accelerator.init_trackers`] should be used to setup your project, and potentially add any experiment hyperparameters to be logged:
+```python
+hps = {"num_iterations": 5, "learning_rate": 1e-2}
+accelerator.init_trackers("my_project", config=hps)
+```
+
+When you are ready to log any data, [`Accelerator.log`] should be used.
+A `step` can also be passed in to correlate the data with a particular step in the training loop.
+```python
+accelerator.log({"train_loss": 1.12, "valid_loss": 0.8}, step=1)
+```
+
+Once you've finished training, make sure to run [`Accelerator.end_training`] so that all the trackers can run their finish functionalities if they have any.
+```python
+accelerator.end_training()
+```
+
+
+A full example is below:
+```python
+from accelerate import Accelerator
+
+accelerator = Accelerator(log_with="all")
+config = {
+    "num_iterations": 5,
+    "learning_rate": 1e-2,
+    "loss_function": str(my_loss_function),
+}
+
+accelerator.init_trackers("example_project", config=config)
+
+my_model, my_optimizer, my_training_dataloader = accelerate.prepare(my_model, my_optimizer, my_training_dataloader)
+device = accelerator.device
+my_model.to(device)
+
+for iteration in config["num_iterations"]:
+    for step, 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)
+        accelerator.backward(loss)
+        my_optimizer.step()
+        accelerator.log({"training_loss": loss}, step=step)
+accelerator.end_training()
+```
+
+
+## Implementing Custom Trackers
+
+To implement a new tracker to be used in `Accelerator`, a new one can be made through implementing the [`GeneralTracker`] class.
+Every tracker must implement three functions and have three properties:
+  - `__init__`: 
+    - Should store a `run_name` and initialize the tracker API of the integrated library. 
+    - If a tracker stores their data locally (such as TensorBoard), a `logging_dir` parameter can be added.
+  - `store_init_configuration`: 
+    - Should take in a `values` dictionary and store them as a one-time experiment configuration
+  - `log`: 
+    - Should take in a `values` dictionary and a `step`, and should log them to the run
+
+  - `name` (`str`):
+    - A unique string name for the tracker, such as `"wandb"` for the wandb tracker. 
+    - This will be used for interacting with this tracker specifically
+  - `requires_logging_directory` (`bool`):
+    - Whether a `logging_dir` is needed for this particular tracker and if it uses one.
+  - `tracker`: 
+    - This should be implemented as a `@property` function 
+    - Should return the internal tracking mechanism the library uses, such as the `run` object for `wandb`.
+
+A brief example can be seen below with an integration with Weights and Biases, containing only the relevant information:
+```python
+from accelerate.tracking import GeneralTracker
+from typing import Optional
+
+import wandb
+
+
+class MyCustomTracker(GeneralTracker):
+    name = "wandb"
+    requires_logging_directory = False
+
+    def __init__(self, run_name: str):
+        self.run_name = run_name
+        run = wandb.init(self.run_name)
+
+    @property
+    def tracker(self):
+        return self.run.run
+
+    def store_init_configuration(self, values: dict):
+        wandb.config(values)
+
+    def log(self, values: dict, step: Optional[int] = None):
+        wandb.log(values, step=step)
+```
+
+When you are ready to build your `Accelerator` object, pass in an **instance** of your tracker to [`Accelerator.log_with`] to have it automatically
+be used with the API:
+
+```python
+tracker = MyCustomTracker("some_run_name")
+accelerator = Accelerator(log_with=tracker)
+```
+
+These also can be mixed with existing trackers, including with `"all"`:
+
+```python
+tracker = MyCustomTracker("some_run_name")
+accelerator = Accelerator(log_with=[tracker, "all"])
+```
+
+## Accessing the internal tracker 
+
+If some custom interactions with a tracker might be wanted directly, you can quickly access one using the 
+[`Accelerator.get_tracker`] method. Just pass in the string corresponding to a tracker's `.name` attribute 
+and it will return that tracker on the main process.
+
+This example shows doing so with wandb:
+
+```python
+wandb_tracker = accelerator.get_tracker("wandb")
+```
+
+From there you can interact with `wandb`'s `run` object like normal:
+
+<Tip warning={true}>
+  Make sure to only interact with trackers on the main process!
+</Tip>
+
+
+```python
+if accelerator.is_main_process:
+    wandb_run.log_artifact(some_artifact_to_log)
+```
+
+## When a wrapper cannot work
+
+If a library has an API that does not follow a strict `.log` with an overall dictionary such as Neptune.AI, logging can be done manually under an `if accelerator.is_main_process` statement:
+```diff
+  from accelerate import Accelerator
++ import neptune.new as neptune
+
+  accelerator = Accelerator()
++ run = neptune.init(...)
+
+  my_model, my_optimizer, my_training_dataloader = accelerate.prepare(my_model, my_optimizer, my_training_dataloader)
+  device = accelerator.device
+  my_model.to(device)
+
+  for iteration in config["num_iterations"]:
+      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)
+          total_loss += loss
+          accelerator.backward(loss)
+          my_optimizer.step()
++         if accelerator.is_main_process:
++             run["logs/training/batch/loss"].log(loss)
+```

docs/source/usage_guides/training_zoo.mdx

@@ -0,0 +1,107 @@
+<!--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.
+-->
+
+# Example Zoo
+
+Below contains a non-exhuastive list of tutorials and scripts showcasing Accelerate
+
+## Official Accelerate Examples:
+
+### Basic Examples
+
+These examples showcase the base features of Accelerate and are a great starting point
+
+- [Barebones NLP example](https://github.com/huggingface/accelerate/blob/main/examples/nlp_example.py)
+- [Barebones computer vision example](https://github.com/huggingface/accelerate/blob/main/examples/cv_example.py)
+
+### Feature Specific Examples
+
+These examples showcase specific features that the Accelerate framework offers
+
+- [Automatic memory-aware gradient accumulation](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/automatic_gradient_accumulation.py)
+- [Checkpointing states](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/checkpointing.py)
+- [Cross validation](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/cross_validation.py)
+- [DeepSpeed](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/deepspeed_with_config_support.py)
+- [Fully Sharded Data Parallelism](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/fsdp_with_peak_mem_tracking.py)
+- [Gradient accumulation](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/gradient_accumulation.py)
+- [Memory-aware batch size finder](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/memory.py)
+- [Metric Computation](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/multi_process_metrics.py)
+- [Using Trackers](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/tracking.py)
+
+### Full Examples 
+
+These examples showcase every feature in Accelerate at once that was shown in "Feature Specific Examples"
+
+- [Complete NLP example](https://github.com/huggingface/accelerate/blob/main/examples/complete_nlp_example.py)
+- [Complete computer vision example](https://github.com/huggingface/accelerate/blob/main/examples/complete_cv_example.py)
+- [Causal language model fine-tuning example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/language-modeling/run_clm_no_trainer.py)
+- [Masked language model fine-tuning example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/language-modeling/run_mlm_no_trainer.py)
+- [Speech pretraining example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/speech-pretraining/run_wav2vec2_pretraining_no_trainer.py)
+- [Translation fine-tuning example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/translation/run_translation_no_trainer.py)
+- [Text classification fine-tuning example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/text-classification/run_glue_no_trainer.py)
+- [Semantic segmentation fine-tuning example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/semantic-segmentation/run_semantic_segmentation_no_trainer.py)
+- [Question answering fine-tuning example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/question-answering/run_qa_no_trainer.py)
+- [Beam search question answering fine-tuning example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/question-answering/run_qa_beam_search_no_trainer.py)
+- [Multiple choice question answering fine-tuning example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/multiple-choice/run_swag_no_trainer.py)
+- [Named entity recognition fine-tuning example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/token-classification/run_ner_no_trainer.py)
+- [Image classification fine-tuning example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/image-classification/run_image_classification_no_trainer.py)
+- [Summarization fine-tuning example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/summarization/run_summarization_no_trainer.py)
+
+## Integration Examples 
+
+These are tutorials from libraries that integrate with 🤗 Accelerate: 
+
+### Catalyst
+
+- [Distributed training tutorial with Catalyst](https://catalyst-team.github.io/catalyst/tutorials/ddp.html)
+
+### DALLE2-pytorch 
+
+- [Fine-tuning DALLE2](https://github.com/lucidrains/DALLE2-pytorch#usage)
+
+### 🤗 diffusers
+
+- [Performing textual inversion with diffusers](https://github.com/huggingface/diffusers/tree/main/examples/textual_inversion)
+- [Training DreamBooth with diffusers](https://github.com/huggingface/diffusers/tree/main/examples/dreambooth)
+
+### fastai 
+
+- [Distributed training from Jupyter Notebooks with fastai](https://docs.fast.ai/tutorial.distributed.html)
+- [Basic distributed training examples with fastai](https://docs.fast.ai/examples/distributed_app_examples.html)
+
+### GradsFlow
+
+- [Auto Image Classification with GradsFlow](https://docs.gradsflow.com/en/latest/examples/nbs/01-ImageClassification/)
+
+### imagen-pytorch 
+
+- [Fine-tuning Imagen](https://github.com/lucidrains/imagen-pytorch#usage)
+
+### Kornia
+
+- [Fine-tuning vision models with Kornia's Trainer](https://kornia.readthedocs.io/en/latest/get-started/training.html)
+
+### PyTorch Accelerated 
+
+- [Quickstart distributed training tutorial with PyTorch Accelerated](https://pytorch-accelerated.readthedocs.io/en/latest/quickstart.html)
+
+### PyTorch3D
+
+- [Perform Deep Learning with 3D data](https://pytorch3d.org/tutorials/)
+
+### Tez 
+
+- [Leaf disease detection with Tez and Accelerate](https://www.kaggle.com/code/abhishek/tez-faster-and-easier-training-for-leaf-detection/notebook)
+
+### trlx 
+
+- [How to implement a sentiment learning task with trlx](https://github.com/CarperAI/trlx#example-how-to-add-a-task)

examples/README.md

@@ -23,12 +23,12 @@ The [nlp_example.py](./nlp_example.py) script is a simple example to train a Ber
 Prior to running it you should install 🤗 Dataset and 🤗 Transformers:
 
 ```bash
-pip install datasets, transformers
+pip install datasets evaluate transformers
 ```
 
 The same script can be run in any of the following configurations:
 - single CPU or single GPU
-- multi GPUS (using PyTorch distributed mode)
+- multi GPUs (using PyTorch distributed mode)
 - (multi) TPUs
 - fp16 (mixed-precision) or fp32 (normal precision)
 
@@ -57,8 +57,8 @@ To run it in each of these various modes, use the following commands:
         ```
     * from any server with Accelerate launcher
         ```bash
-        accelerate launch --fb16 ./nlp_example.py
-- multi GPUS (using PyTorch distributed mode)
+        accelerate launch --fp16 ./nlp_example.py
+- multi GPUs (using PyTorch distributed mode)
     * With Accelerate config and launcher
         ```bash
         accelerate config  # This will create a config file on your server
@@ -103,14 +103,14 @@ The [cv_example.py](./cv_example.py) script is a simple example to fine-tune a R
 
 The same script can be run in any of the following configurations:
 - single CPU or single GPU
-- multi GPUS (using PyTorch distributed mode)
+- multi GPUs (using PyTorch distributed mode)
 - (multi) TPUs
 - fp16 (mixed-precision) or fp32 (normal precision)
 
 Prior to running it you should install timm and torchvision:
 
 ```bash
-pip install timm, torchvision
+pip install timm torchvision
 ```
 
 and you should download the data with the following commands:
@@ -136,7 +136,7 @@ To run it in each of these various modes, use the following commands:
         ```
 - single GPU:
     ```bash
-    python ./nlp_example.py  # from a server with a GPU
+    python ./cv_example.py  # from a server with a GPU
     ```
 - with fp16 (mixed-precision)
     * from any server by passing `fp16=True` to the `Accelerator`.
@@ -145,8 +145,8 @@ To run it in each of these various modes, use the following commands:
         ```
     * from any server with Accelerate launcher
         ```bash
-        accelerate launch --fb16 ./cv_example.py --data_dir path_to_data
-- multi GPUS (using PyTorch distributed mode)
+        accelerate launch --fp16 ./cv_example.py --data_dir path_to_data
+- multi GPUs (using PyTorch distributed mode)
     * With Accelerate config and launcher
         ```bash
         accelerate config  # This will create a config file on your server
@@ -183,3 +183,30 @@ To run it in each of these various modes, use the following commands:
         ```
     * In PyTorch:
         Add an `xmp.spawn` line in your script as you usually do.
+
+### Simple vision example (GANs)
+
+- [huggan project](https://github.com/huggingface/community-events/tree/main/huggan)
+
+### Using AWS SageMaker integration
+- [Examples showcasing AWS SageMaker integration of 🤗 Accelerate.](https://github.com/pacman100/accelerate-aws-sagemaker)
+    
+## Finer Examples
+
+While the first two scripts are extremely barebones when it comes to what you can do with accelerate, more advanced features are documented in two other locations.
+
+### `by_feature` examples
+
+These scripts are *individual* examples highlighting one particular feature or use-case within Accelerate. They all stem from the [nlp_example.py](./nlp_example.py) script, and any changes or modifications is denoted with a `# New Code #` comment.
+
+Read the README.md file located in the `by_feature` folder for more information.
+
+### `complete_*` examples
+
+These two scripts contain *every* single feature currently available in Accelerate in one place, as one giant script.
+
+New arguments that can be passed include:
+
+- `checkpointing_steps`, whether the various states should be saved at the end of every `n` steps, or `"epoch"` for each epoch. States are then saved to folders named `step_{n}` or `epoch_{n}`
+- `resume_from_checkpoint`, should be used if you want to resume training off of a previous call to the script and passed a `checkpointing_steps` to it.
+- `with_tracking`, should be used if you want to log the training run using all available experiment trackers in your environment. Currently supported trackers include TensorBoard, Weights and Biases, and CometML.

examples/by_feature/README.md

@@ -0,0 +1,80 @@
+# What are these scripts?
+
+All scripts in this folder originate from the `nlp_example.py` file, as it is a very simplistic NLP training example using Accelerate with zero extra features.
+
+From there, each further script adds in just **one** feature of Accelerate, showing how you can quickly modify your own scripts to implement these capabilities.
+
+A full example with all of these parts integrated together can be found in the `complete_nlp_example.py` script and `complete_cv_example.py` script.
+
+Adjustments to each script from the base `nlp_example.py` file can be found quickly by searching for "# New Code #"
+
+## Example Scripts by Feature and their Arguments
+
+### Base Example (`../nlp_example.py`)
+
+- Shows how to use `Accelerator` in an extremely simplistic PyTorch training loop
+- Arguments available:
+  - `mixed_precision`, whether to use mixed precision. ("no", "fp16", or "bf16")
+  - `cpu`, whether to train using only the CPU. (yes/no/1/0)
+
+All following scripts also accept these arguments in addition to their added ones.
+
+These arguments should be added at the end of any method for starting the python script (such as `python`, `accelerate launch`, `python -m torch.distributed.launch`), such as:
+
+```bash
+accelerate launch ../nlp_example.py --mixed_precision fp16 --cpu 0
+```
+
+### Checkpointing and Resuming Training (`checkpointing.py`)
+
+- Shows how to use `Accelerator.save_state` and `Accelerator.load_state` to save or continue training
+- **It is assumed you are continuing off the same training script**
+- Arguments available:
+  - `checkpointing_steps`, after how many steps the various states should be saved. ("epoch", 1, 2, ...)
+  - `output_dir`, where saved state folders should be saved to, default is current working directory
+  - `resume_from_checkpoint`, what checkpoint folder to resume from. ("epoch_0", "step_22", ...)
+
+These arguments should be added at the end of any method for starting the python script (such as `python`, `accelerate launch`, `python -m torch.distributed.launch`), such as:
+
+(Note, `resume_from_checkpoint` assumes that we've ran the script for one epoch with the `--checkpointing_steps epoch` flag)
+
+```bash
+accelerate launch ./checkpointing.py --checkpointing_steps epoch output_dir "checkpointing_tutorial" --resume_from_checkpoint "checkpointing_tutorial/epoch_0"
+```
+
+### Cross Validation (`cross_validation.py`)
+
+- Shows how to use `Accelerator.free_memory` and run cross validation efficiently with `datasets`.
+- Arguments available:
+  - `num_folds`, the number of folds the training dataset should be split into.
+
+These arguments should be added at the end of any method for starting the python script (such as `python`, `accelerate launch`, `python -m torch.distributed.launch`), such as:
+
+```bash
+accelerate launch ./cross_validation.py --num_folds 2
+```
+
+### Experiment Tracking (`tracking.py`)
+
+- Shows how to use `Accelerate.init_trackers` and `Accelerator.log`
+- Can be used with Weights and Biases, TensorBoard, or CometML.
+- Arguments available:
+  - `with_tracking`, whether to load in all available experiment trackers from the environment.
+
+These arguments should be added at the end of any method for starting the python script (such as `python`, `accelerate launch`, `python -m torch.distributed.launch`), such as:
+
+```bash
+accelerate launch ./tracking.py --with_tracking
+```
+
+### Gradient Accumulation (`gradient_accumulation.py`)
+
+- Shows how to use `Accelerator.no_sync` to prevent gradient averaging in a distributed setup.
+- Arguments available:
+  - `gradient_accumulation_steps`, the number of steps to perform before the gradients are accumulated and the optimizer and scheduler are stepped + zero_grad
+
+These arguments should be added at the end of any method for starting the python script (such as `python`, `accelerate launch`, `python -m torch.distributed.launch`), such as:
+
+```bash
+accelerate launch ./gradient_accumulation.py --gradient_accumulation_steps 5
+```

examples/by_feature/automatic_gradient_accumulation.py

@@ -0,0 +1,232 @@
+# 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 os
+
+import torch
+from torch.optim import AdamW
+from torch.utils.data import DataLoader
+
+# New Code #
+import evaluate
+from accelerate import Accelerator, DistributedType
+from accelerate.utils import find_executable_batch_size
+from datasets import load_dataset
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed
+
+
+########################################################################
+# This is a fully working simple example to use Accelerate,
+# specifically showcasing how to combine both the gradient accumulation
+# and automatic batch size finder utilities of Accelerate to perfrom
+# automatic gradient accumulation
+#
+# This example trains a Bert base model on GLUE MRPC
+# in any of the following settings (with the same script):
+#   - single CPU or single GPU
+#   - multi GPUS (using PyTorch distributed mode)
+#   - (multi) TPUs
+#   - fp16 (mixed-precision) or fp32 (normal precision)
+#
+# New additions from the base script can be found quickly by
+# looking for the # New Code # tags
+#
+# To run it in each of these various modes, follow the instructions
+# in the readme for examples:
+# https://github.com/huggingface/accelerate/tree/main/examples
+#
+########################################################################
+
+EVAL_BATCH_SIZE = 32
+
+
+def get_dataloaders(accelerator: Accelerator, batch_size: int = 16):
+    """
+    Creates a set of `DataLoader`s for the `glue` dataset,
+    using "bert-base-cased" as the tokenizer.
+
+    Args:
+        accelerator (`Accelerator`):
+            An `Accelerator` object
+        batch_size (`int`, *optional*):
+            The batch size for the train and validation DataLoaders.
+    """
+    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+    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:
+    with accelerator.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):
+        # On TPU it's best to pad everything to the same length or training will be very slow.
+        if accelerator.distributed_type == DistributedType.TPU:
+            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
+        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(
+        tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size
+    )
+    eval_dataloader = DataLoader(
+        tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
+    )
+
+    return train_dataloader, eval_dataloader
+
+
+# For testing only
+if os.environ.get("TESTING_MOCKED_DATALOADERS", None) == "1":
+    from accelerate.test_utils.training import mocked_dataloaders
+
+    get_dataloaders = mocked_dataloaders  # noqa: F811
+
+
+def training_function(config, args):
+    # For testing only
+    if os.environ.get("TESTING_MOCKED_DATALOADERS", None) == "1":
+        config["num_epochs"] = 2
+    # Initialize accelerator
+    accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)
+    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
+    lr = config["lr"]
+    num_epochs = int(config["num_epochs"])
+    seed = int(config["seed"])
+    observed_batch_size = int(config["batch_size"])
+
+    metric = evaluate.load("glue", "mrpc")
+
+    # New Code #
+    # We use the `find_executable_batch_size` decorator, passing in the desired observed batch size
+    # to train on. If a CUDA OOM error occurs, it will retry this loop cutting the batch size in
+    # half each time. From this, we can calculate the number of gradient accumulation steps needed
+    # and modify the Accelerator object as a result
+    @find_executable_batch_size(starting_batch_size=int(observed_batch_size))
+    def inner_training_loop(batch_size):
+        # Since we need to modify the outside accelerator object, we need to bring it
+        # to the local scope
+        nonlocal accelerator
+
+        # We can calculate the number of gradient accumulation steps based on the current
+        # batch size vs the starting batch size
+        num_gradient_accumulation_steps = observed_batch_size // batch_size
+
+        # And then set it in the Accelerator directly:
+        accelerator.gradient_accumulation_steps = num_gradient_accumulation_steps
+
+        # Next we need to free all of the stored model references in the Accelerator each time
+        accelerator.free_memory()
+
+        # And set the seed so our results are reproducable each reset
+        set_seed(seed)
+
+        # Instantiate the model (we build the model here so that the seed also control new weights initialization)
+        model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", return_dict=True)
+
+        # We could avoid this line since the accelerator is set with `device_placement=True` (default value).
+        # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
+        # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
+        model = model.to(accelerator.device)
+
+        # Instantiate optimizer
+        optimizer = AdamW(params=model.parameters(), lr=lr)
+        train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size)
+
+        # Instantiate scheduler
+        lr_scheduler = get_linear_schedule_with_warmup(
+            optimizer=optimizer,
+            num_warmup_steps=100,
+            num_training_steps=(len(train_dataloader) * num_epochs),
+        )
+
+        # 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, lr_scheduler = accelerator.prepare(
+            model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+        )
+
+        # Now we train the model
+        for epoch in range(num_epochs):
+            model.train()
+            for step, batch in enumerate(train_dataloader):
+                # And perform gradient accumulation
+                with accelerator.accumulate(model):
+                    # We could avoid this line since we set the accelerator with `device_placement=True`.
+                    batch.to(accelerator.device)
+                    outputs = model(**batch)
+                    loss = outputs.loss
+                    accelerator.backward(loss)
+                    optimizer.step()
+                    lr_scheduler.step()
+                    optimizer.zero_grad()
+
+            model.eval()
+            for step, batch in enumerate(eval_dataloader):
+                # We could avoid this line since we set the accelerator with `device_placement=True`.
+                batch.to(accelerator.device)
+                with torch.no_grad():
+                    outputs = model(**batch)
+                predictions = outputs.logits.argmax(dim=-1)
+                predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]))
+                metric.add_batch(
+                    predictions=predictions,
+                    references=references,
+                )
+
+            eval_metric = metric.compute()
+            # Use accelerator.print to print only on the main process.
+            accelerator.print(f"epoch {epoch}:", eval_metric)
+
+    # New Code #
+    # And call it at the end with no arguments
+    # Note: You could also refactor this outside of your training loop function
+    inner_training_loop()
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Simple example of training script.")
+    parser.add_argument(
+        "--mixed_precision",
+        type=str,
+        default="no",
+        choices=["no", "fp16", "bf16"],
+        help="Whether to use mixed precision. Choose"
+        "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
+        "and an Nvidia Ampere GPU.",
+    )
+    parser.add_argument("--cpu", action="store_true", help="If passed, will train on the CPU.")
+    args = parser.parse_args()
+    # New Code #
+    # We modify the starting batch size to be an observed batch size of 256, to guarentee an initial CUDA OOM
+    config = {"lr": 2e-5, "num_epochs": 3, "seed": 42, "batch_size": 256}
+    training_function(config, args)
+
+
+if __name__ == "__main__":
+    main()

examples/by_feature/checkpointing.py

@@ -0,0 +1,303 @@
+# coding=utf-8
+# Copyright 2021 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 argparse
+import os
+
+import torch
+from torch.optim import AdamW
+from torch.utils.data import DataLoader
+
+import evaluate
+from accelerate import Accelerator, DistributedType
+from datasets import load_dataset
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed
+
+
+########################################################################
+# This is a fully working simple example to use Accelerate,
+# specifically showcasing the checkpointing capability,
+# and builds off the `nlp_example.py` script.
+#
+# This example trains a Bert base model on GLUE MRPC
+# in any of the following settings (with the same script):
+#   - single CPU or single GPU
+#   - multi GPUS (using PyTorch distributed mode)
+#   - (multi) TPUs
+#   - fp16 (mixed-precision) or fp32 (normal precision)
+#
+# To help focus on the differences in the code, building `DataLoaders`
+# was refactored into its own function.
+# New additions from the base script can be found quickly by
+# looking for the # New Code # tags
+#
+# To run it in each of these various modes, follow the instructions
+# in the readme for examples:
+# https://github.com/huggingface/accelerate/tree/main/examples
+#
+########################################################################
+
+MAX_GPU_BATCH_SIZE = 16
+EVAL_BATCH_SIZE = 32
+
+
+def get_dataloaders(accelerator: Accelerator, batch_size: int = 16):
+    """
+    Creates a set of `DataLoader`s for the `glue` dataset,
+    using "bert-base-cased" as the tokenizer.
+
+    Args:
+        accelerator (`Accelerator`):
+            An `Accelerator` object
+        batch_size (`int`, *optional*):
+            The batch size for the train and validation DataLoaders.
+    """
+    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+    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:
+    with accelerator.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):
+        # On TPU it's best to pad everything to the same length or training will be very slow.
+        if accelerator.distributed_type == DistributedType.TPU:
+            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
+        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(
+        tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size
+    )
+    eval_dataloader = DataLoader(
+        tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
+    )
+
+    return train_dataloader, eval_dataloader
+
+
+# For testing only
+if os.environ.get("TESTING_MOCKED_DATALOADERS", None) == "1":
+    from accelerate.test_utils.training import mocked_dataloaders
+
+    get_dataloaders = mocked_dataloaders  # noqa: F811
+
+
+def training_function(config, args):
+    # For testing only
+    if os.environ.get("TESTING_MOCKED_DATALOADERS", None) == "1":
+        config["num_epochs"] = 2
+    # Initialize accelerator
+    accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)
+    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
+    lr = config["lr"]
+    num_epochs = int(config["num_epochs"])
+    seed = int(config["seed"])
+    batch_size = int(config["batch_size"])
+
+    # New Code #
+    # Parse out whether we are saving every epoch or after a certain number of batches
+    if hasattr(args.checkpointing_steps, "isdigit"):
+        if args.checkpointing_steps == "epoch":
+            checkpointing_steps = args.checkpointing_steps
+        elif args.checkpointing_steps.isdigit():
+            checkpointing_steps = int(args.checkpointing_steps)
+        else:
+            raise ValueError(
+                f"Argument `checkpointing_steps` must be either a number or `epoch`. `{args.checkpointing_steps}` passed."
+            )
+    else:
+        checkpointing_steps = None
+
+    set_seed(seed)
+
+    train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size)
+    metric = evaluate.load("glue", "mrpc")
+
+    # If the batch size is too big we use gradient accumulation
+    gradient_accumulation_steps = 1
+    if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
+        gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
+        batch_size = MAX_GPU_BATCH_SIZE
+
+    # Instantiate the model (we build the model here so that the seed also control new weights initialization)
+    model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", return_dict=True)
+
+    # We could avoid this line since the accelerator is set with `device_placement=True` (default value).
+    # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
+    # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
+    model = model.to(accelerator.device)
+
+    # Instantiate optimizer
+    optimizer = AdamW(params=model.parameters(), lr=lr)
+
+    # Instantiate scheduler
+    lr_scheduler = get_linear_schedule_with_warmup(
+        optimizer=optimizer,
+        num_warmup_steps=100,
+        num_training_steps=(len(train_dataloader) * num_epochs) // gradient_accumulation_steps,
+    )
+
+    # 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, lr_scheduler = accelerator.prepare(
+        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+    )
+
+    # New Code #
+    # We need to keep track of how many total steps we have iterated over
+    overall_step = 0
+    # We also need to keep track of the stating epoch so files are named properly
+    starting_epoch = 0
+
+    # We need to load the checkpoint back in before training here with `load_state`
+    # The total number of epochs is adjusted based on where the state is being loaded from,
+    # as we assume continuation of the same training script
+    if args.resume_from_checkpoint:
+        if args.resume_from_checkpoint is not None or args.resume_from_checkpoint != "":
+            accelerator.print(f"Resumed from checkpoint: {args.resume_from_checkpoint}")
+            accelerator.load_state(args.resume_from_checkpoint)
+            path = os.path.basename(args.resume_from_checkpoint)
+        else:
+            # Get the most recent checkpoint
+            dirs = [f.name for f in os.scandir(os.getcwd()) if f.is_dir()]
+            dirs.sort(key=os.path.getctime)
+            path = dirs[-1]  # Sorts folders by date modified, most recent checkpoint is the last
+        # Extract `epoch_{i}` or `step_{i}`
+        training_difference = os.path.splitext(path)[0]
+
+        if "epoch" in training_difference:
+            starting_epoch = int(training_difference.replace("epoch_", "")) + 1
+            resume_step = None
+        else:
+            resume_step = int(training_difference.replace("step_", ""))
+            starting_epoch = resume_step // len(train_dataloader)
+            resume_step -= starting_epoch * len(train_dataloader)
+
+    # Now we train the model
+    for epoch in range(starting_epoch, num_epochs):
+        model.train()
+        for step, batch in enumerate(train_dataloader):
+            # New Code #
+            # We need to skip steps until we reach the resumed step during the first epoch
+            if args.resume_from_checkpoint and epoch == starting_epoch:
+                if resume_step is not None and step < resume_step:
+                    overall_step += 1
+                    continue
+            # We could avoid this line since we set the accelerator with `device_placement=True`.
+            batch.to(accelerator.device)
+            outputs = model(**batch)
+            loss = outputs.loss
+            loss = loss / gradient_accumulation_steps
+            accelerator.backward(loss)
+            if step % gradient_accumulation_steps == 0:
+                optimizer.step()
+                lr_scheduler.step()
+                optimizer.zero_grad()
+            # New Code #
+            overall_step += 1
+
+            # New Code #
+            # We save the model, optimizer, lr_scheduler, and seed states by calling `save_state`
+            # These are saved to folders named `step_{overall_step}`
+            # Will contain files: "pytorch_model.bin", "optimizer.bin", "scheduler.bin", and "random_states.pkl"
+            # If mixed precision was used, will also save a "scalar.bin" file
+            if isinstance(checkpointing_steps, int):
+                output_dir = f"step_{overall_step}"
+                if overall_step % checkpointing_steps == 0:
+                    if args.output_dir is not None:
+                        output_dir = os.path.join(args.output_dir, output_dir)
+                    accelerator.save_state(output_dir)
+
+        model.eval()
+        for step, batch in enumerate(eval_dataloader):
+            # We could avoid this line since we set the accelerator with `device_placement=True` (the default).
+            batch.to(accelerator.device)
+            with torch.no_grad():
+                outputs = model(**batch)
+            predictions = outputs.logits.argmax(dim=-1)
+            predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]))
+            metric.add_batch(
+                predictions=predictions,
+                references=references,
+            )
+
+        eval_metric = metric.compute()
+        # Use accelerator.print to print only on the main process.
+        accelerator.print(f"epoch {epoch}:", eval_metric)
+
+        # New Code #
+        # We save the model, optimizer, lr_scheduler, and seed states by calling `save_state`
+        # These are saved to folders named `epoch_{epoch}`
+        # Will contain files: "pytorch_model.bin", "optimizer.bin", "scheduler.bin", and "random_states.pkl"
+        # If mixed precision was used, will also save a "scalar.bin" file
+        if checkpointing_steps == "epoch":
+            output_dir = f"epoch_{epoch}"
+            if args.output_dir is not None:
+                output_dir = os.path.join(args.output_dir, output_dir)
+            accelerator.save_state(output_dir)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Simple example of training script.")
+    parser.add_argument(
+        "--mixed_precision",
+        type=str,
+        default="no",
+        choices=["no", "fp16", "bf16"],
+        help="Whether to use mixed precision. Choose"
+        "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
+        "and an Nvidia Ampere GPU.",
+    )
+    parser.add_argument("--cpu", action="store_true", help="If passed, will train on the CPU.")
+    parser.add_argument(
+        "--checkpointing_steps",
+        type=str,
+        default=None,
+        help="Whether the various states should be saved at the end of every n steps, or 'epoch' for each epoch.",
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default=".",
+        help="Optional save directory where all checkpoint folders will be stored. Default is the current working directory.",
+    )
+    parser.add_argument(
+        "--resume_from_checkpoint",
+        type=str,
+        default=None,
+        help="If the training should continue from a checkpoint folder.",
+    )
+    args = parser.parse_args()
+    config = {"lr": 2e-5, "num_epochs": 3, "seed": 42, "batch_size": 16}
+    training_function(config, args)
+
+
+if __name__ == "__main__":
+    main()

examples/by_feature/cross_validation.py

@@ -0,0 +1,268 @@
+# coding=utf-8
+# Copyright 2022 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 argparse
+from typing import List
+
+import numpy as np
+import torch
+from torch.optim import AdamW
+from torch.utils.data import DataLoader
+
+import evaluate
+from accelerate import Accelerator, DistributedType
+from datasets import DatasetDict, load_dataset
+
+# New Code #
+# We'll be using StratifiedKFold for this example
+from sklearn.model_selection import StratifiedKFold
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed
+
+
+########################################################################
+# This is a fully working simple example to use Accelerate,
+# specifically showcasing how to perform Cross Validation,
+# and builds off the `nlp_example.py` script.
+#
+# This example trains a Bert base model on GLUE MRPC
+# in any of the following settings (with the same script):
+#   - single CPU or single GPU
+#   - multi GPUS (using PyTorch distributed mode)
+#   - (multi) TPUs
+#   - fp16 (mixed-precision) or fp32 (normal precision)
+#
+# To help focus on the differences in the code, building `DataLoaders`
+# was refactored into its own function.
+# New additions from the base script can be found quickly by
+# looking for the # New Code # tags
+#
+# To run it in each of these various modes, follow the instructions
+# in the readme for examples:
+# https://github.com/huggingface/accelerate/tree/main/examples
+#
+########################################################################
+
+
+MAX_GPU_BATCH_SIZE = 16
+EVAL_BATCH_SIZE = 32
+
+# New Code #
+# We need a different `get_dataloaders` function that will build dataloaders by index
+
+
+def get_fold_dataloaders(
+    accelerator: Accelerator, dataset: DatasetDict, train_idxs: List[int], valid_idxs: List[int], batch_size: int = 16
+):
+    """
+    Gets a set of train, valid, and test dataloaders for a particular fold
+
+    Args:
+        accelerator (`Accelerator`):
+            The main `Accelerator` object
+        train_idxs (list of `int`):
+            The split indices for the training dataset
+        valid_idxs (list of `int`):
+            The split indices for the validation dataset
+        batch_size (`int`):
+            The size of the minibatch. Default is 16
+    """
+    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+    datasets = DatasetDict(
+        {
+            "train": dataset["train"].select(train_idxs),
+            "validation": dataset["train"].select(valid_idxs),
+            "test": dataset["validation"],
+        }
+    )
+
+    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:
+    with accelerator.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):
+        # On TPU it's best to pad everything to the same length or training will be very slow.
+        if accelerator.distributed_type == DistributedType.TPU:
+            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
+        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(
+        tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size
+    )
+    eval_dataloader = DataLoader(
+        tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
+    )
+
+    test_dataloader = DataLoader(
+        tokenized_datasets["test"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
+    )
+
+    return train_dataloader, eval_dataloader, test_dataloader
+
+
+def training_function(config, args):
+    # New Code #
+    test_predictions = []
+    # Download the dataset
+    datasets = load_dataset("glue", "mrpc")
+    # Create our splits
+    kfold = StratifiedKFold(n_splits=int(args.num_folds))
+    # Initialize accelerator
+    accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)
+    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
+    lr = config["lr"]
+    num_epochs = int(config["num_epochs"])
+    seed = int(config["seed"])
+    batch_size = int(config["batch_size"])
+
+    metric = evaluate.load("glue", "mrpc")
+
+    # If the batch size is too big we use gradient accumulation
+    gradient_accumulation_steps = 1
+    if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
+        gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
+        batch_size = MAX_GPU_BATCH_SIZE
+
+    set_seed(seed)
+
+    # New Code #
+    # Create our folds:
+    folds = kfold.split(np.zeros(datasets["train"].num_rows), datasets["train"]["label"])
+    test_references = []
+    # Iterate over them
+    for i, (train_idxs, valid_idxs) in enumerate(folds):
+        train_dataloader, eval_dataloader, test_dataloader = get_fold_dataloaders(
+            accelerator,
+            datasets,
+            train_idxs,
+            valid_idxs,
+        )
+        # Instantiate the model (we build the model here so that the seed also control new weights initialization)
+        model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", return_dict=True)
+
+        # We could avoid this line since the accelerator is set with `device_placement=True` (default value).
+        # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
+        # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
+        model = model.to(accelerator.device)
+
+        # Instantiate optimizer
+        optimizer = AdamW(params=model.parameters(), lr=lr)
+
+        # Instantiate scheduler
+        lr_scheduler = get_linear_schedule_with_warmup(
+            optimizer=optimizer,
+            num_warmup_steps=100,
+            num_training_steps=(len(train_dataloader) * num_epochs) // gradient_accumulation_steps,
+        )
+
+        # 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, lr_scheduler = accelerator.prepare(
+            model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+        )
+
+        # Now we train the model
+        for epoch in range(num_epochs):
+            model.train()
+            for step, batch in enumerate(train_dataloader):
+                # We could avoid this line since we set the accelerator with `device_placement=True`.
+                batch.to(accelerator.device)
+                outputs = model(**batch)
+                loss = outputs.loss
+                loss = loss / gradient_accumulation_steps
+                accelerator.backward(loss)
+                if step % gradient_accumulation_steps == 0:
+                    optimizer.step()
+                    lr_scheduler.step()
+                    optimizer.zero_grad()
+
+            model.eval()
+            for step, batch in enumerate(eval_dataloader):
+                # We could avoid this line since we set the accelerator with `device_placement=True`.
+                batch.to(accelerator.device)
+                with torch.no_grad():
+                    outputs = model(**batch)
+                predictions = outputs.logits.argmax(dim=-1)
+                predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]))
+                metric.add_batch(
+                    predictions=predictions,
+                    references=references,
+                )
+
+            eval_metric = metric.compute()
+            # Use accelerator.print to print only on the main process.
+            accelerator.print(f"epoch {epoch}:", eval_metric)
+
+        # New Code #
+        # We also run predictions on the test set at the very end
+        fold_predictions = []
+        for step, batch in enumerate(test_dataloader):
+            # We could avoid this line since we set the accelerator with `device_placement=True`.
+            batch.to(accelerator.device)
+            with torch.no_grad():
+                outputs = model(**batch)
+            predictions = outputs.logits
+            predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]))
+            fold_predictions.append(predictions.cpu())
+            if i == 0:
+                # We need all of the test predictions
+                test_references.append(references.cpu())
+        # Use accelerator.print to print only on the main process.
+        test_predictions.append(torch.cat(fold_predictions, dim=0))
+        # We now need to release all our memory and get rid of the current model, optimizer, etc
+        accelerator.free_memory()
+    # New Code #
+    # Finally we check the accuracy of our folded results:
+    test_references = torch.cat(test_references, dim=0)
+    preds = torch.stack(test_predictions, dim=0).sum(dim=0).div(int(args.num_folds)).argmax(dim=-1)
+    test_metric = metric.compute(predictions=preds, references=test_references)
+    accelerator.print("Average test metrics from all folds:", test_metric)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Simple example of training script.")
+    parser.add_argument(
+        "--mixed_precision",
+        type=str,
+        default="no",
+        choices=["no", "fp16", "bf16"],
+        help="Whether to use mixed precision. Choose"
+        "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
+        "and an Nvidia Ampere GPU.",
+    )
+    parser.add_argument("--cpu", action="store_true", help="If passed, will train on the CPU.")
+    # New Code #
+    parser.add_argument("--num_folds", type=int, default=3, help="The number of splits to perform across the dataset")
+    args = parser.parse_args()
+    config = {"lr": 2e-5, "num_epochs": 3, "seed": 42, "batch_size": 16}
+    training_function(config, args)
+
+
+if __name__ == "__main__":
+    main()

examples/by_feature/deepspeed_with_config_support.py

@@ -0,0 +1,734 @@
+#!/usr/bin/env python
+# coding=utf-8
+# Copyright 2022 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.
+"""
+Fine-tuning the library models for causal language modeling (GPT, GPT-2, CTRL, ...)
+on a text file or a dataset without using HuggingFace Trainer.
+
+Here is the full list of checkpoints on the hub that can be fine-tuned by this script:
+https://huggingface.co/models?filter=text-generation
+"""
+# You can also adapt this script on your own causal language modeling task. Pointers for this are left as comments.
+
+import argparse
+import json
+import logging
+import math
+import os
+import random
+from itertools import chain
+from pathlib import Path
+
+import torch
+from torch.utils.data import DataLoader
+
+import datasets
+import transformers
+from accelerate import Accelerator, DistributedType
+from accelerate.logging import get_logger
+from accelerate.utils import DummyOptim, DummyScheduler, set_seed
+from datasets import load_dataset
+from huggingface_hub import Repository
+from tqdm.auto import tqdm
+from transformers import (
+    CONFIG_MAPPING,
+    MODEL_MAPPING,
+    AutoConfig,
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    SchedulerType,
+    default_data_collator,
+    get_scheduler,
+)
+from transformers.utils import get_full_repo_name
+from transformers.utils.versions import require_version
+
+
+logger = get_logger(__name__)
+
+require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/language-modeling/requirements.txt")
+
+MODEL_CONFIG_CLASSES = list(MODEL_MAPPING.keys())
+MODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Finetune a transformers model on a causal language modeling task")
+    parser.add_argument(
+        "--dataset_name",
+        type=str,
+        default=None,
+        help="The name of the dataset to use (via the datasets library).",
+    )
+    parser.add_argument(
+        "--dataset_config_name",
+        type=str,
+        default=None,
+        help="The configuration name of the dataset to use (via the datasets library).",
+    )
+    parser.add_argument(
+        "--train_file", type=str, default=None, help="A csv or a json file containing the training data."
+    )
+    parser.add_argument(
+        "--validation_file", type=str, default=None, help="A csv or a json file containing the validation data."
+    )
+    parser.add_argument(
+        "--validation_split_percentage",
+        default=5,
+        help="The percentage of the train set used as validation set in case there's no validation split",
+    )
+    parser.add_argument(
+        "--model_name_or_path",
+        type=str,
+        help="Path to pretrained model or model identifier from huggingface.co/models.",
+        required=False,
+    )
+    parser.add_argument(
+        "--config_name",
+        type=str,
+        default=None,
+        help="Pretrained config name or path if not the same as model_name",
+    )
+    parser.add_argument(
+        "--tokenizer_name",
+        type=str,
+        default=None,
+        help="Pretrained tokenizer name or path if not the same as model_name",
+    )
+    parser.add_argument(
+        "--use_slow_tokenizer",
+        action="store_true",
+        help="If passed, will use a slow tokenizer (not backed by the 🤗 Tokenizers library).",
+    )
+    parser.add_argument(
+        "--per_device_train_batch_size",
+        type=int,
+        default=8,
+        help="Batch size (per device) for the training dataloader.",
+    )
+    parser.add_argument(
+        "--per_device_eval_batch_size",
+        type=int,
+        default=8,
+        help="Batch size (per device) for the evaluation dataloader.",
+    )
+    parser.add_argument(
+        "--learning_rate",
+        type=float,
+        default=5e-5,
+        help="Initial learning rate (after the potential warmup period) to use.",
+    )
+    parser.add_argument("--weight_decay", type=float, default=0.0, help="Weight decay to use.")
+    parser.add_argument("--num_train_epochs", type=int, default=3, help="Total number of training epochs to perform.")
+    parser.add_argument(
+        "--max_train_steps",
+        type=int,
+        default=None,
+        help="Total number of training steps to perform. If provided, overrides num_train_epochs.",
+    )
+    parser.add_argument(
+        "--gradient_accumulation_steps",
+        type=int,
+        default=1,
+        help="Number of updates steps to accumulate before performing a backward/update pass.",
+    )
+    parser.add_argument(
+        "--lr_scheduler_type",
+        type=SchedulerType,
+        default="linear",
+        help="The scheduler type to use.",
+        choices=["linear", "cosine", "cosine_with_restarts", "polynomial", "constant", "constant_with_warmup"],
+    )
+    parser.add_argument(
+        "--num_warmup_steps", type=int, default=0, help="Number of steps for the warmup in the lr scheduler."
+    )
+    parser.add_argument("--output_dir", type=str, default=None, help="Where to store the final model.")
+    parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")
+    parser.add_argument(
+        "--model_type",
+        type=str,
+        default=None,
+        help="Model type to use if training from scratch.",
+        choices=MODEL_TYPES,
+    )
+    parser.add_argument(
+        "--block_size",
+        type=int,
+        default=None,
+        help=(
+            "Optional input sequence length after tokenization. The training dataset will be truncated in block of"
+            " this size for training. Default to the model max input length for single sentence inputs (take into"
+            " account special tokens)."
+        ),
+    )
+    parser.add_argument(
+        "--preprocessing_num_workers",
+        type=int,
+        default=None,
+        help="The number of processes to use for the preprocessing.",
+    )
+    parser.add_argument(
+        "--overwrite_cache", type=bool, default=False, help="Overwrite the cached training and evaluation sets"
+    )
+    parser.add_argument(
+        "--no_keep_linebreaks", action="store_true", help="Do not keep line breaks when using TXT files."
+    )
+    parser.add_argument("--push_to_hub", action="store_true", help="Whether or not to push the model to the Hub.")
+    parser.add_argument(
+        "--hub_model_id", type=str, help="The name of the repository to keep in sync with the local `output_dir`."
+    )
+    parser.add_argument("--hub_token", type=str, help="The token to use to push to the Model Hub.")
+    parser.add_argument(
+        "--checkpointing_steps",
+        type=str,
+        default=None,
+        help="Whether the various states should be saved at the end of every n steps, or 'epoch' for each epoch.",
+    )
+    parser.add_argument(
+        "--resume_from_checkpoint",
+        type=str,
+        default=None,
+        help="If the training should continue from a checkpoint folder.",
+    )
+    # New Code #
+    # Whether to load the best model at the end of training
+    parser.add_argument(
+        "--load_best_model",
+        action="store_true",
+        help="Whether to load the best model at the end of training",
+    )
+    parser.add_argument(
+        "--with_tracking",
+        action="store_true",
+        help="Whether to enable experiment trackers for logging.",
+    )
+    parser.add_argument(
+        "--report_to",
+        type=str,
+        default="all",
+        help=(
+            'The integration to report the results and logs to. Supported platforms are `"tensorboard"`,'
+            ' `"wandb"` and `"comet_ml"`. Use `"all"` (default) to report to all integrations.'
+            "Only applicable when `--with_tracking` is passed."
+        ),
+    )
+    args = parser.parse_args()
+
+    # Sanity checks
+    if args.dataset_name is None and args.train_file is None and args.validation_file is None:
+        raise ValueError("Need either a dataset name or a training/validation file.")
+    else:
+        if args.train_file is not None:
+            extension = args.train_file.split(".")[-1]
+            assert extension in ["csv", "json", "txt"], "`train_file` should be a csv, json or txt file."
+        if args.validation_file is not None:
+            extension = args.validation_file.split(".")[-1]
+            assert extension in ["csv", "json", "txt"], "`validation_file` should be a csv, json or txt file."
+
+    if args.push_to_hub:
+        assert args.output_dir is not None, "Need an `output_dir` to create a repo when `--push_to_hub` is passed."
+
+    return args
+
+
+# New Code #
+def checkpoint_model(checkpoint_folder, ckpt_id, model, epoch, last_global_step, **kwargs):
+    """Utility function for checkpointing model + optimizer dictionaries
+    The main purpose for this is to be able to resume training from that instant again
+    """
+    checkpoint_state_dict = {
+        "epoch": epoch,
+        "last_global_step": last_global_step,
+    }
+    # Add extra kwargs too
+    checkpoint_state_dict.update(kwargs)
+
+    success = model.save_checkpoint(checkpoint_folder, ckpt_id, checkpoint_state_dict)
+    status_msg = f"checkpointing: checkpoint_folder={checkpoint_folder}, ckpt_id={ckpt_id}"
+    if success:
+        logging.info(f"Success {status_msg}")
+    else:
+        logging.warning(f"Failure {status_msg}")
+    return
+
+
+# New Code #
+def load_training_checkpoint(model, load_dir, tag=None, **kwargs):
+    """Utility function for checkpointing model + optimizer dictionaries
+    The main purpose for this is to be able to resume training from that instant again
+    """
+    _, checkpoint_state_dict = model.load_checkpoint(load_dir, tag=tag, **kwargs)
+    epoch = checkpoint_state_dict["epoch"]
+    last_global_step = checkpoint_state_dict["last_global_step"]
+    del checkpoint_state_dict
+    return (epoch, last_global_step)
+
+
+# New Code #
+def evaluate(args, model, eval_dataloader, accelerator, eval_dataset):
+    model.eval()
+    losses = []
+    for step, batch in enumerate(eval_dataloader):
+        with torch.no_grad():
+            outputs = model(**batch)
+
+        loss = outputs.loss
+        losses.append(accelerator.gather(loss.repeat(args.per_device_eval_batch_size)))
+
+    losses = torch.cat(losses)
+    losses = losses[: len(eval_dataset)]
+    try:
+        eval_loss = torch.mean(losses)
+        perplexity = math.exp(eval_loss)
+    except OverflowError:
+        perplexity = float("inf")
+    return perplexity, eval_loss
+
+
+def main():
+    args = parse_args()
+
+    # Initialize the accelerator. We will let the accelerator handle device placement for us in this example.
+    # If we're using tracking, we also need to initialize it here and it will by default pick up all supported trackers
+    # in the environment
+    accelerator = (
+        Accelerator(log_with=args.report_to, logging_dir=args.output_dir) if args.with_tracking else Accelerator()
+    )
+    # Make one log on every process with the configuration for debugging.
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+        datefmt="%m/%d/%Y %H:%M:%S",
+        level=logging.INFO,
+    )
+    logger.info(accelerator.state, main_process_only=False)
+    if accelerator.is_local_main_process:
+        datasets.utils.logging.set_verbosity_warning()
+        transformers.utils.logging.set_verbosity_info()
+    else:
+        datasets.utils.logging.set_verbosity_error()
+        transformers.utils.logging.set_verbosity_error()
+
+    # If passed along, set the training seed now.
+    if args.seed is not None:
+        set_seed(args.seed)
+
+    # Handle the repository creation
+    if accelerator.is_main_process:
+        if args.push_to_hub:
+            if args.hub_model_id is None:
+                repo_name = get_full_repo_name(Path(args.output_dir).name, token=args.hub_token)
+            else:
+                repo_name = args.hub_model_id
+            repo = Repository(args.output_dir, clone_from=repo_name)
+
+            with open(os.path.join(args.output_dir, ".gitignore"), "w+") as gitignore:
+                if "step_*" not in gitignore:
+                    gitignore.write("step_*\n")
+                if "epoch_*" not in gitignore:
+                    gitignore.write("epoch_*\n")
+        elif args.output_dir is not None:
+            os.makedirs(args.output_dir, exist_ok=True)
+    accelerator.wait_for_everyone()
+
+    # Get the datasets: you can either provide your own CSV/JSON/TXT training and evaluation files (see below)
+    # or just provide the name of one of the public datasets available on the hub at https://huggingface.co/datasets/
+    # (the dataset will be downloaded automatically from the datasets Hub).
+    #
+    # For CSV/JSON files, this script will use the column called 'text' or the first column if no column called
+    # 'text' is found. You can easily tweak this behavior (see below).
+    #
+    # In distributed training, the load_dataset function guarantee that only one local process can concurrently
+    # download the dataset.
+    if args.dataset_name is not None:
+        # Downloading and loading a dataset from the hub.
+        raw_datasets = load_dataset(args.dataset_name, args.dataset_config_name)
+        if "validation" not in raw_datasets.keys():
+            raw_datasets["validation"] = load_dataset(
+                args.dataset_name,
+                args.dataset_config_name,
+                split=f"train[:{args.validation_split_percentage}%]",
+            )
+            raw_datasets["train"] = load_dataset(
+                args.dataset_name,
+                args.dataset_config_name,
+                split=f"train[{args.validation_split_percentage}%:]",
+            )
+    else:
+        data_files = {}
+        dataset_args = {}
+        if args.train_file is not None:
+            data_files["train"] = args.train_file
+        if args.validation_file is not None:
+            data_files["validation"] = args.validation_file
+        extension = args.train_file.split(".")[-1]
+        if extension == "txt":
+            extension = "text"
+            dataset_args["keep_linebreaks"] = not args.no_keep_linebreaks
+        raw_datasets = load_dataset(extension, data_files=data_files, **dataset_args)
+        # If no validation data is there, validation_split_percentage will be used to divide the dataset.
+        if "validation" not in raw_datasets.keys():
+            raw_datasets["validation"] = load_dataset(
+                extension,
+                data_files=data_files,
+                split=f"train[:{args.validation_split_percentage}%]",
+                **dataset_args,
+            )
+            raw_datasets["train"] = load_dataset(
+                extension,
+                data_files=data_files,
+                split=f"train[{args.validation_split_percentage}%:]",
+                **dataset_args,
+            )
+
+    # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at
+    # https://huggingface.co/docs/datasets/loading_datasets.html.
+
+    # Load pretrained model and tokenizer
+    #
+    # In distributed training, the .from_pretrained methods guarantee that only one local process can concurrently
+    # download model & vocab.
+    if args.config_name:
+        config = AutoConfig.from_pretrained(args.config_name)
+    elif args.model_name_or_path:
+        config = AutoConfig.from_pretrained(args.model_name_or_path)
+    else:
+        config = CONFIG_MAPPING[args.model_type]()
+        logger.warning("You are instantiating a new config instance from scratch.")
+
+    if args.tokenizer_name:
+        tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_name, use_fast=not args.use_slow_tokenizer)
+    elif args.model_name_or_path:
+        tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path, use_fast=not args.use_slow_tokenizer)
+    else:
+        raise ValueError(
+            "You are instantiating a new tokenizer from scratch. This is not supported by this script."
+            "You can do it from another script, save it, and load it from here, using --tokenizer_name."
+        )
+
+    if args.model_name_or_path:
+        model = AutoModelForCausalLM.from_pretrained(
+            args.model_name_or_path,
+            from_tf=bool(".ckpt" in args.model_name_or_path),
+            config=config,
+        )
+    else:
+        logger.info("Training new model from scratch")
+        model = AutoModelForCausalLM.from_config(config)
+
+    model.resize_token_embeddings(len(tokenizer))
+
+    # Preprocessing the datasets.
+    # First we tokenize all the texts.
+    column_names = raw_datasets["train"].column_names
+    text_column_name = "text" if "text" in column_names else column_names[0]
+
+    def tokenize_function(examples):
+        return tokenizer(examples[text_column_name])
+
+    with accelerator.main_process_first():
+        tokenized_datasets = raw_datasets.map(
+            tokenize_function,
+            batched=True,
+            num_proc=args.preprocessing_num_workers,
+            remove_columns=column_names,
+            load_from_cache_file=not args.overwrite_cache,
+            desc="Running tokenizer on dataset",
+        )
+
+    if args.block_size is None:
+        block_size = tokenizer.model_max_length
+        if block_size > 1024:
+            logger.warning(
+                f"The tokenizer picked seems to have a very large `model_max_length` ({tokenizer.model_max_length}). "
+                "Picking 1024 instead. You can change that default value by passing --block_size xxx."
+            )
+        block_size = 1024
+    else:
+        if args.block_size > tokenizer.model_max_length:
+            logger.warning(
+                f"The block_size passed ({args.block_size}) is larger than the maximum length for the model"
+                f"({tokenizer.model_max_length}). Using block_size={tokenizer.model_max_length}."
+            )
+        block_size = min(args.block_size, tokenizer.model_max_length)
+
+    # Main data processing function that will concatenate all texts from our dataset and generate chunks of block_size.
+    def group_texts(examples):
+        # Concatenate all texts.
+        concatenated_examples = {k: list(chain(*examples[k])) for k in examples.keys()}
+        total_length = len(concatenated_examples[list(examples.keys())[0]])
+        # We drop the small remainder, we could add padding if the model supported it instead of this drop, you can
+        # customize this part to your needs.
+        if total_length >= block_size:
+            total_length = (total_length // block_size) * block_size
+        # Split by chunks of max_len.
+        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
+
+    # Note that with `batched=True`, this map processes 1,000 texts together, so group_texts throws away a remainder
+    # for each of those groups of 1,000 texts. You can adjust that batch_size here but a higher value might be slower
+    # to preprocess.
+    #
+    # To speed up this part, we use multiprocessing. See the documentation of the map method for more information:
+    # https://huggingface.co/docs/datasets/package_reference/main_classes.html#datasets.Dataset.map
+
+    with accelerator.main_process_first():
+        lm_datasets = tokenized_datasets.map(
+            group_texts,
+            batched=True,
+            num_proc=args.preprocessing_num_workers,
+            load_from_cache_file=not args.overwrite_cache,
+            desc=f"Grouping texts in chunks of {block_size}",
+        )
+
+    train_dataset = lm_datasets["train"]
+    eval_dataset = lm_datasets["validation"]
+
+    # Log a few random samples from the training set:
+    for index in random.sample(range(len(train_dataset)), 3):
+        logger.info(f"Sample {index} of the training set: {train_dataset[index]}.")
+
+    # DataLoaders creation:
+    train_dataloader = DataLoader(
+        train_dataset, shuffle=True, collate_fn=default_data_collator, batch_size=args.per_device_train_batch_size
+    )
+    eval_dataloader = DataLoader(
+        eval_dataset, collate_fn=default_data_collator, batch_size=args.per_device_eval_batch_size
+    )
+
+    # Optimizer
+    # Split weights in two groups, one with weight decay and the other not.
+    no_decay = ["bias", "LayerNorm.weight"]
+    optimizer_grouped_parameters = [
+        {
+            "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
+            "weight_decay": args.weight_decay,
+        },
+        {
+            "params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)],
+            "weight_decay": 0.0,
+        },
+    ]
+    # New Code #
+    # Creates Dummy Optimizer if `optimizer` was specified in the config file else creates Adam Optimizer
+    optimizer_cls = (
+        torch.optim.AdamW
+        if accelerator.state.deepspeed_plugin is None
+        or "optimizer" not in accelerator.state.deepspeed_plugin.deepspeed_config
+        else DummyOptim
+    )
+    optimizer = optimizer_cls(optimizer_grouped_parameters, lr=args.learning_rate)
+
+    # On TPU, the tie weights in our model have been disconnected, so we need to restore the ties.
+    if accelerator.distributed_type == DistributedType.TPU:
+        model.tie_weights()
+
+    # Scheduler and math around the number of training steps.
+
+    # New Code
+    # Get gradient accumulation steps from deepspeed config if available
+    if accelerator.state.deepspeed_plugin is not None:
+        args.gradient_accumulation_steps = accelerator.state.deepspeed_plugin.deepspeed_config[
+            "gradient_accumulation_steps"
+        ]
+
+    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
+    if args.max_train_steps is None:
+        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
+    else:
+        args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
+
+    # New Code #
+    # Creates Dummy Scheduler if `scheduler` was specified in the config file else creates `args.lr_scheduler_type` Scheduler
+    if (
+        accelerator.state.deepspeed_plugin is None
+        or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config
+    ):
+        lr_scheduler = get_scheduler(
+            name=args.lr_scheduler_type,
+            optimizer=optimizer,
+            num_warmup_steps=args.num_warmup_steps,
+            num_training_steps=args.max_train_steps,
+        )
+    else:
+        lr_scheduler = DummyScheduler(
+            optimizer, total_num_steps=args.max_train_steps, warmup_num_steps=args.num_warmup_steps
+        )
+
+    # Prepare everything with our `accelerator`.
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+    )
+
+    # We need to recalculate our total training steps as the size of the training dataloader may have changed.
+    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
+    args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
+
+    # Figure out how many steps we should save the Accelerator states
+    if hasattr(args.checkpointing_steps, "isdigit"):
+        checkpointing_steps = args.checkpointing_steps
+        if args.checkpointing_steps.isdigit():
+            checkpointing_steps = int(args.checkpointing_steps)
+    else:
+        checkpointing_steps = None
+
+    # We need to initialize the trackers we use, and also store our configuration.
+    # The trackers initializes automatically on the main process.
+    if args.with_tracking:
+        experiment_config = vars(args)
+        # TensorBoard cannot log Enums, need the raw value
+        experiment_config["lr_scheduler_type"] = experiment_config["lr_scheduler_type"].value
+        accelerator.init_trackers("clm_no_trainer", experiment_config)
+
+    # Train!
+    total_batch_size = args.per_device_train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
+
+    logger.info("***** Running training *****")
+    logger.info(f"  Num examples = {len(train_dataset)}")
+    logger.info(f"  Num Epochs = {args.num_train_epochs}")
+    logger.info(f"  Instantaneous batch size per device = {args.per_device_train_batch_size}")
+    logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
+    logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
+    logger.info(f"  Total optimization steps = {args.max_train_steps}")
+    # Only show the progress bar once on each machine.
+    progress_bar = tqdm(range(args.max_train_steps), disable=not accelerator.is_local_main_process)
+    completed_steps = 0
+    starting_epoch = 0
+    best_metric = None
+    best_metric_checkpoint = None
+
+    # Potentially load in the weights and states from a previous save
+    if args.resume_from_checkpoint:
+        # New Code #
+        # Loads the DeepSpeed checkpoint from the specified path
+        _, last_global_step = load_training_checkpoint(
+            model,
+            args.resume_from_checkpoint,
+            **{"load_optimizer_states": True, "load_lr_scheduler_states": True},
+        )
+        accelerator.print(f"Resumed from checkpoint: {args.resume_from_checkpoint}")
+        resume_step = last_global_step
+        starting_epoch = resume_step // len(train_dataloader)
+        resume_step -= starting_epoch * len(train_dataloader)
+
+    for epoch in range(starting_epoch, args.num_train_epochs):
+        model.train()
+        if args.with_tracking:
+            total_loss = 0
+        for step, batch in enumerate(train_dataloader):
+            # We need to skip steps until we reach the resumed step
+            if args.resume_from_checkpoint and epoch == starting_epoch:
+                if resume_step is not None and step < resume_step:
+                    completed_steps += 1
+                    continue
+            outputs = model(**batch)
+            loss = outputs.loss
+            # We keep track of the loss at each epoch
+            if args.with_tracking:
+                total_loss += loss.detach().float()
+            loss = loss / args.gradient_accumulation_steps
+            accelerator.backward(loss)
+            if step % args.gradient_accumulation_steps == 0 or step == len(train_dataloader) - 1:
+                optimizer.step()
+                lr_scheduler.step()
+                optimizer.zero_grad()
+                progress_bar.update(1)
+                completed_steps += 1
+
+            if isinstance(checkpointing_steps, int):
+                if completed_steps % checkpointing_steps == 0:
+                    output_dir = f"step_{completed_steps }"
+                    if args.output_dir is not None:
+                        output_dir = os.path.join(args.output_dir, output_dir)
+                    accelerator.save_state(output_dir)
+            if completed_steps >= args.max_train_steps:
+                break
+
+        perplexity, eval_loss = evaluate(args, model, eval_dataloader, accelerator, eval_dataset)
+        logger.info(f"epoch {epoch}: perplexity: {perplexity} eval_loss: {eval_loss}")
+
+        if args.with_tracking:
+            accelerator.log(
+                {
+                    "perplexity": perplexity,
+                    "eval_loss": eval_loss,
+                    "train_loss": total_loss.item() / len(train_dataloader),
+                    "epoch": epoch,
+                    "step": completed_steps,
+                },
+                step=completed_steps,
+            )
+
+        # New Code #
+        # Save the DeepSpeed checkpoint to the specified path
+        checkpoint_model(args.output_dir, epoch, model, epoch, completed_steps)
+
+        # New Code #
+        # Tracks the best checkpoint and best metric
+        if best_metric is None or best_metric > perplexity:
+            best_metric = perplexity
+            best_metric_checkpoint = os.path.join(args.output_dir, str(epoch))
+            accelerator.print(f"New best metric: {best_metric} at epoch {epoch}")
+            accelerator.print(f"best_metric_checkpoint: {best_metric_checkpoint}")
+
+    # New Code #
+    # Loads the best checkpoint after the training is finished
+    if args.load_best_model:
+        _, last_global_step = load_training_checkpoint(
+            model,
+            "/".join(best_metric_checkpoint.split("/")[:-1]),
+            tag=best_metric_checkpoint.split("/")[-1],
+            **{"load_optimizer_states": True, "load_lr_scheduler_states": True},
+        )
+
+    # New Code #
+    # Evaluates using the best checkpoint
+    perplexity, eval_loss = evaluate(args, model, eval_dataloader, accelerator, eval_dataset)
+    logger.info(f"Best model metrics: perplexity: {perplexity} eval_loss: {eval_loss}")
+    if perplexity != best_metric:
+        raise AssertionError(
+            f"Best metric {best_metric} does not match the metric {perplexity} of the loaded best model."
+        )
+
+    if args.output_dir is not None:
+        accelerator.wait_for_everyone()
+        unwrapped_model = accelerator.unwrap_model(model)
+
+        # New Code #
+        # Saves the whole/unpartitioned fp16 model when in ZeRO Stage-3 to the output directory if
+        # `stage3_gather_16bit_weights_on_model_save` is True in DeepSpeed Config file or
+        # `zero3_save_16bit_model` is True in DeepSpeed Plugin.
+        # For Zero Stages 1 and 2, models are saved as usual in the output directory.
+        # The model name saved is `pytorch_model.bin`
+        unwrapped_model.save_pretrained(
+            args.output_dir,
+            is_main_process=accelerator.is_main_process,
+            save_function=accelerator.save,
+            state_dict=accelerator.get_state_dict(model),
+        )
+        if accelerator.is_main_process:
+            tokenizer.save_pretrained(args.output_dir)
+            if args.push_to_hub:
+                repo.push_to_hub(commit_message="End of training", auto_lfs_prune=True)
+
+        with open(os.path.join(args.output_dir, "all_results.json"), "w") as f:
+            json.dump({"perplexity": perplexity, "eval_loss": eval_loss.item()}, f)
+
+
+if __name__ == "__main__":
+    main()

examples/by_feature/fsdp_with_peak_mem_tracking.py

@@ -0,0 +1,381 @@
+# coding=utf-8
+# Copyright 2021 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 argparse
+import gc
+import os
+
+import torch
+from torch.utils.data import DataLoader
+
+import evaluate
+from accelerate import Accelerator, DistributedType
+from datasets import load_dataset
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed
+
+
+########################################################################
+# This is a fully working simple example to use Accelerate
+#
+# This example trains a Bert base model on GLUE MRPC
+# in any of the following settings (with the same script):
+#   - single CPU or single GPU
+#   - multi GPUS (using PyTorch distributed mode)
+#   - (multi) TPUs
+#   - fp16 (mixed-precision) or fp32 (normal precision)
+#   - FSDP
+#
+# This example also demonstrates the checkpointing and sharding capabilities
+#
+# To run it in each of these various modes, follow the instructions
+# in the readme for examples:
+# https://github.com/huggingface/accelerate/tree/main/examples
+#
+########################################################################
+
+
+MAX_GPU_BATCH_SIZE = 16
+EVAL_BATCH_SIZE = 32
+
+
+# New Code #
+# Converting Bytes to Megabytes
+def b2mb(x):
+    return int(x / 2**20)
+
+
+# New Code #
+# This context manager is used to track the peak memory usage of the process
+class TorchTracemalloc:
+    def __enter__(self):
+        gc.collect()
+        torch.cuda.empty_cache()
+        torch.cuda.reset_max_memory_allocated()  # reset the peak gauge to zero
+        self.begin = torch.cuda.memory_allocated()
+        return self
+
+    def __exit__(self, *exc):
+        gc.collect()
+        torch.cuda.empty_cache()
+        self.end = torch.cuda.memory_allocated()
+        self.peak = torch.cuda.max_memory_allocated()
+        self.used = b2mb(self.end - self.begin)
+        self.peaked = b2mb(self.peak - self.begin)
+        # print(f"delta used/peak {self.used:4d}/{self.peaked:4d}")
+
+
+# For testing only
+if os.environ.get("TESTING_MOCKED_DATALOADERS", None) == "1":
+    from accelerate.test_utils.training import mocked_dataloaders
+
+    get_dataloaders = mocked_dataloaders  # noqa: F811
+
+
+def training_function(config, args):
+    # For testing only
+    if os.environ.get("TESTING_MOCKED_DATALOADERS", None) == "1":
+        config["num_epochs"] = 2
+    # Initialize accelerator
+    if args.with_tracking:
+        accelerator = Accelerator(
+            cpu=args.cpu, mixed_precision=args.mixed_precision, log_with="wandb", logging_dir=args.logging_dir
+        )
+    else:
+        accelerator = Accelerator()
+    accelerator.print(accelerator.distributed_type)
+
+    if hasattr(args.checkpointing_steps, "isdigit"):
+        if args.checkpointing_steps == "epoch":
+            checkpointing_steps = args.checkpointing_steps
+        elif args.checkpointing_steps.isdigit():
+            checkpointing_steps = int(args.checkpointing_steps)
+        else:
+            raise ValueError(
+                f"Argument `checkpointing_steps` must be either a number or `epoch`. `{args.checkpointing_steps}` passed."
+            )
+    else:
+        checkpointing_steps = None
+    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
+    lr = config["lr"]
+    num_epochs = int(config["num_epochs"])
+    seed = int(config["seed"])
+    batch_size = int(config["batch_size"])
+
+    # We need to initialize the trackers we use, and also store our configuration
+    if args.with_tracking:
+        experiment_config = vars(args)
+        accelerator.init_trackers("fsdp_glue_no_trainer", experiment_config)
+
+    tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path)
+    datasets = load_dataset("glue", "mrpc")
+    metric = evaluate.load("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:
+    with accelerator.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")
+
+    # If the batch size is too big we use gradient accumulation
+    gradient_accumulation_steps = 1
+    if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
+        gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
+        batch_size = MAX_GPU_BATCH_SIZE
+
+    def collate_fn(examples):
+        # On TPU it's best to pad everything to the same length or training will be very slow.
+        if accelerator.distributed_type == DistributedType.TPU:
+            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
+        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(
+        tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size
+    )
+    eval_dataloader = DataLoader(
+        tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
+    )
+
+    set_seed(seed)
+
+    # Instantiate the model (we build the model here so that the seed also control new weights initialization)
+    model = AutoModelForSequenceClassification.from_pretrained(args.model_name_or_path, return_dict=True)
+    # New Code #
+    # For FSDP feature, it is highly recommended and efficient to prepare the model before creating optimizer
+    model = accelerator.prepare(model)
+    accelerator.print(model)
+
+    # Instantiate optimizer
+    # New Code #
+    # For FSDP feature, at present it doesn't support multiple parameter groups,
+    # so we need to create a single parameter group for the whole model
+    optimizer = torch.optim.AdamW(params=model.parameters(), lr=lr, weight_decay=2e-4)
+
+    # Instantiate scheduler
+    lr_scheduler = get_linear_schedule_with_warmup(
+        optimizer=optimizer,
+        num_warmup_steps=10,
+        num_training_steps=(len(train_dataloader) * num_epochs) // gradient_accumulation_steps,
+    )
+
+    # New Code #
+    # For FSDP feature, prepare everything except the model as we have already prepared the model
+    # before creating the optimizer
+    # 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.
+    optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+        optimizer, train_dataloader, eval_dataloader, lr_scheduler
+    )
+
+    overall_step = 0
+
+    # Potentially load in the weights and states from a previous save
+    if args.resume_from_checkpoint:
+        if args.resume_from_checkpoint is not None or args.resume_from_checkpoint != "":
+            accelerator.print(f"Resumed from checkpoint: {args.resume_from_checkpoint}")
+            accelerator.load_state(args.resume_from_checkpoint)
+            path = os.path.basename(args.resume_from_checkpoint)
+        else:
+            # Get the most recent checkpoint
+            dirs = [f.name for f in os.scandir(os.getcwd()) if f.is_dir()]
+            dirs.sort(key=os.path.getctime)
+            path = dirs[-1]  # Sorts folders by date modified, most recent checkpoint is the last
+        # Extract `epoch_{i}` or `step_{i}`
+        training_difference = os.path.splitext(path)[0]
+
+        if "epoch" in training_difference:
+            num_epochs -= int(training_difference.replace("epoch_", ""))
+            resume_step = None
+        else:
+            resume_step = int(training_difference.replace("step_", ""))
+            num_epochs -= resume_step // len(train_dataloader)
+            # If resuming by step, we also need to know exactly how far into the DataLoader we went
+            resume_step = (num_epochs * len(train_dataloader)) - resume_step
+
+    # Now we train the model
+    for epoch in range(num_epochs):
+        # New Code #
+        # context manager to track the peak memory usage during the training epoch
+        with TorchTracemalloc() as tracemalloc:
+            model.train()
+            if args.with_tracking:
+                total_loss = 0
+            for step, batch in enumerate(train_dataloader):
+                # We need to skip steps until we reach the resumed step
+                if args.resume_from_checkpoint and epoch == 0:
+                    if resume_step is not None and step < resume_step:
+                        pass
+                # We could avoid this line since we set the accelerator with `device_placement=True`.
+                batch.to(accelerator.device)
+                outputs = model(**batch)
+                loss = outputs.loss
+                loss = loss / gradient_accumulation_steps
+                # We keep track of the loss at each epoch
+                if args.with_tracking:
+                    total_loss += loss.detach().float()
+                accelerator.backward(loss)
+                if step % gradient_accumulation_steps == 0:
+                    optimizer.step()
+                    lr_scheduler.step()
+                    optimizer.zero_grad()
+                    # accelerator.print(lr_scheduler.get_lr())
+
+                overall_step += 1
+
+                if isinstance(checkpointing_steps, int):
+                    output_dir = f"step_{overall_step}"
+                    if overall_step % checkpointing_steps == 0:
+                        if args.output_dir is not None:
+                            output_dir = os.path.join(args.output_dir, output_dir)
+                        accelerator.save_state(output_dir)
+        # New Code #
+        # Printing the GPU memory usage details such as allocated memory, peak memory, and total memory usage
+        accelerator.print("Memory before entering the train : {}".format(b2mb(tracemalloc.begin)))
+        accelerator.print("Memory consumed at the end of the train (end-begin): {}".format(tracemalloc.used))
+        accelerator.print("Peak Memory consumed during the train (max-begin): {}".format(tracemalloc.peaked))
+        accelerator.print(
+            "Total Peak Memory consumed during the train (max): {}".format(
+                tracemalloc.peaked + b2mb(tracemalloc.begin)
+            )
+        )
+        # Logging the peak memory usage of the GPU to the tracker
+        if args.with_tracking:
+            accelerator.log(
+                {
+                    "train_total_peak_memory": tracemalloc.peaked + b2mb(tracemalloc.begin),
+                },
+                step=epoch,
+            )
+
+        # New Code #
+        # context manager to track the peak memory usage during the evaluation
+        with TorchTracemalloc() as tracemalloc:
+            model.eval()
+            for step, batch in enumerate(eval_dataloader):
+                # We could avoid this line since we set the accelerator with `device_placement=True`.
+                batch.to(accelerator.device)
+                with torch.no_grad():
+                    outputs = model(**batch)
+                predictions = outputs.logits.argmax(dim=-1)
+                predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]))
+                metric.add_batch(
+                    predictions=predictions,
+                    references=references,
+                )
+
+            eval_metric = metric.compute()
+            # Use accelerator.print to print only on the main process.
+            accelerator.print(f"epoch {epoch}:", eval_metric)
+            if args.with_tracking:
+                accelerator.log(
+                    {
+                        "accuracy": eval_metric["accuracy"],
+                        "f1": eval_metric["f1"],
+                        "train_loss": total_loss.item() / len(train_dataloader),
+                    },
+                    step=epoch,
+                )
+
+            if checkpointing_steps == "epoch":
+                output_dir = f"epoch_{epoch}"
+                if args.output_dir is not None:
+                    output_dir = os.path.join(args.output_dir, output_dir)
+                accelerator.save_state(output_dir)
+        # New Code #
+        # Printing the GPU memory usage details such as allocated memory, peak memory, and total memory usage
+        accelerator.print("Memory before entering the eval : {}".format(b2mb(tracemalloc.begin)))
+        accelerator.print("Memory consumed at the end of the eval (end-begin): {}".format(tracemalloc.used))
+        accelerator.print("Peak Memory consumed during the eval (max-begin): {}".format(tracemalloc.peaked))
+        accelerator.print(
+            "Total Peak Memory consumed during the eval (max): {}".format(tracemalloc.peaked + b2mb(tracemalloc.begin))
+        )
+        # Logging the peak memory usage of the GPU to the tracker
+        if args.with_tracking:
+            accelerator.log(
+                {
+                    "eval_total_peak_memory": tracemalloc.peaked + b2mb(tracemalloc.begin),
+                },
+                step=epoch,
+            )
+
+    if args.with_tracking:
+        accelerator.end_training()
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Simple example of training script.")
+    parser.add_argument(
+        "--mixed_precision",
+        type=str,
+        default="no",
+        choices=["no", "fp16", "bf16"],
+        help="Whether to use mixed precision. Choose"
+        "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
+        "and an Nvidia Ampere GPU.",
+    )
+    parser.add_argument("--cpu", action="store_true", help="If passed, will train on the CPU.")
+    parser.add_argument(
+        "--checkpointing_steps",
+        type=str,
+        default=None,
+        help="Whether the various states should be saved at the end of every n steps, or 'epoch' for each epoch.",
+    )
+    parser.add_argument(
+        "--resume_from_checkpoint",
+        type=str,
+        default=None,
+        help="If the training should continue from a checkpoint folder.",
+    )
+    parser.add_argument(
+        "--with_tracking",
+        action="store_true",
+        help="Whether to load in all available experiment trackers from the environment and use them for logging.",
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default=".",
+        help="Optional save directory where all checkpoint folders will be stored. Default is the current working directory.",
+    )
+    parser.add_argument(
+        "--logging_dir",
+        type=str,
+        default="logs",
+        help="Location on where to store experiment tracking logs`",
+    )
+    parser.add_argument(
+        "--model_name_or_path",
+        type=str,
+        help="Path to pretrained model or model identifier from huggingface.co/models.",
+        required=True,
+    )
+    args = parser.parse_args()
+    config = {"lr": 2e-5, "num_epochs": 3, "seed": 1, "batch_size": 16}
+    training_function(config, args)
+
+
+if __name__ == "__main__":
+    main()

examples/by_feature/gradient_accumulation.py

@@ -0,0 +1,215 @@
+# coding=utf-8
+# Copyright 2021 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 argparse
+import os
+
+import torch
+from torch.optim import AdamW
+from torch.utils.data import DataLoader
+
+import evaluate
+from accelerate import Accelerator, DistributedType
+from datasets import load_dataset
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed
+
+
+########################################################################
+# This is a fully working simple example to use Accelerate
+# and perform gradient accumulation
+#
+# This example trains a Bert base model on GLUE MRPC
+# in any of the following settings (with the same script):
+#   - single CPU or single GPU
+#   - multi GPUS (using PyTorch distributed mode)
+#   - (multi) TPUs
+#   - fp16 (mixed-precision) or fp32 (normal precision)
+#
+# To run it in each of these various modes, follow the instructions
+# in the readme for examples:
+# https://github.com/huggingface/accelerate/tree/main/examples
+#
+########################################################################
+
+
+MAX_GPU_BATCH_SIZE = 16
+EVAL_BATCH_SIZE = 32
+
+
+def get_dataloaders(accelerator: Accelerator, batch_size: int = 16):
+    """
+    Creates a set of `DataLoader`s for the `glue` dataset,
+    using "bert-base-cased" as the tokenizer.
+
+    Args:
+        accelerator (`Accelerator`):
+            An `Accelerator` object
+        batch_size (`int`, *optional*):
+            The batch size for the train and validation DataLoaders.
+    """
+    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+    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:
+    with accelerator.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):
+        # On TPU it's best to pad everything to the same length or training will be very slow.
+        if accelerator.distributed_type == DistributedType.TPU:
+            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
+        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(
+        tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size
+    )
+    eval_dataloader = DataLoader(
+        tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
+    )
+
+    return train_dataloader, eval_dataloader
+
+
+# For testing only
+if os.environ.get("TESTING_MOCKED_DATALOADERS", None) == "1":
+    from accelerate.test_utils.training import mocked_dataloaders
+
+    get_dataloaders = mocked_dataloaders  # noqa: F811
+
+
+def training_function(config, args):
+    # For testing only
+    if os.environ.get("TESTING_MOCKED_DATALOADERS", None) == "1":
+        config["num_epochs"] = 2
+    # New Code #
+    gradient_accumulation_steps = int(args.gradient_accumulation_steps)
+    # Initialize accelerator
+    accelerator = Accelerator(
+        cpu=args.cpu, mixed_precision=args.mixed_precision, gradient_accumulation_steps=gradient_accumulation_steps
+    )
+    if accelerator.distributed_type == DistributedType.TPU and gradient_accumulation_steps > 1:
+        raise NotImplementedError(
+            "Gradient accumulation on TPUs is currently not supported. Pass `gradient_accumulation_steps=1`"
+        )
+    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
+    lr = config["lr"]
+    num_epochs = int(config["num_epochs"])
+    seed = int(config["seed"])
+    batch_size = int(config["batch_size"])
+
+    metric = evaluate.load("glue", "mrpc")
+
+    set_seed(seed)
+    train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size)
+    # Instantiate the model (we build the model here so that the seed also control new weights initialization)
+    model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", return_dict=True)
+
+    # We could avoid this line since the accelerator is set with `device_placement=True` (default value).
+    # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
+    # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
+    model = model.to(accelerator.device)
+
+    # Instantiate optimizer
+    optimizer = AdamW(params=model.parameters(), lr=lr)
+
+    # Instantiate scheduler
+    lr_scheduler = get_linear_schedule_with_warmup(
+        optimizer=optimizer,
+        num_warmup_steps=100,
+        num_training_steps=(len(train_dataloader) * num_epochs),
+    )
+
+    # 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, lr_scheduler = accelerator.prepare(
+        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+    )
+
+    # Now we train the model
+    for epoch in range(num_epochs):
+        model.train()
+        for step, batch in enumerate(train_dataloader):
+            # We could avoid this line since we set the accelerator with `device_placement=True`.
+            batch.to(accelerator.device)
+            # New code #
+            # We use the new `accumulate` context manager to perform gradient accumulation
+            # We also currently do not support TPUs nor advise it as bugs were found on the XLA side when running our tests.
+            with accelerator.accumulate(model):
+                output = model(**batch)
+                loss = output.loss
+                accelerator.backward(loss)
+                optimizer.step()
+                lr_scheduler.step()
+                optimizer.zero_grad()
+
+        model.eval()
+        for step, batch in enumerate(eval_dataloader):
+            # We could avoid this line since we set the accelerator with `device_placement=True`.
+            batch.to(accelerator.device)
+            with torch.no_grad():
+                outputs = model(**batch)
+            predictions = outputs.logits.argmax(dim=-1)
+            predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]))
+            metric.add_batch(
+                predictions=predictions,
+                references=references,
+            )
+
+        eval_metric = metric.compute()
+        # Use accelerator.print to print only on the main process.
+        accelerator.print(f"epoch {epoch}:", eval_metric)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Simple example of training script.")
+    parser.add_argument(
+        "--mixed_precision",
+        type=str,
+        default="no",
+        choices=["no", "fp16", "bf16"],
+        help="Whether to use mixed precision. Choose"
+        "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
+        "and an Nvidia Ampere GPU.",
+    )
+    # New Code #
+    parser.add_argument(
+        "--gradient_accumulation_steps",
+        type=int,
+        default=1,
+        help="The number of minibatches to be ran before gradients are accumulated.",
+    )
+    parser.add_argument("--cpu", action="store_true", help="If passed, will train on the CPU.")
+    args = parser.parse_args()
+    config = {"lr": 2e-5, "num_epochs": 3, "seed": 42, "batch_size": 16}
+    training_function(config, args)
+
+
+if __name__ == "__main__":
+    main()

examples/by_feature/memory.py

@@ -0,0 +1,220 @@
+# 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 os
+
+import torch
+from torch.optim import AdamW
+from torch.utils.data import DataLoader
+
+# New Code #
+import evaluate
+from accelerate import Accelerator, DistributedType
+from accelerate.utils import find_executable_batch_size
+from datasets import load_dataset
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed
+
+
+########################################################################
+# This is a fully working simple example to use Accelerate,
+# specifically showcasing how to ensure out-of-memory errors never
+# interrupt training, and builds off the `nlp_example.py` script.
+#
+# This example trains a Bert base model on GLUE MRPC
+# in any of the following settings (with the same script):
+#   - single CPU or single GPU
+#   - multi GPUS (using PyTorch distributed mode)
+#   - (multi) TPUs
+#   - fp16 (mixed-precision) or fp32 (normal precision)
+#
+# New additions from the base script can be found quickly by
+# looking for the # New Code # tags
+#
+# To run it in each of these various modes, follow the instructions
+# in the readme for examples:
+# https://github.com/huggingface/accelerate/tree/main/examples
+#
+########################################################################
+
+
+MAX_GPU_BATCH_SIZE = 16
+EVAL_BATCH_SIZE = 32
+
+
+def get_dataloaders(accelerator: Accelerator, batch_size: int = 16):
+    """
+    Creates a set of `DataLoader`s for the `glue` dataset,
+    using "bert-base-cased" as the tokenizer.
+
+    Args:
+        accelerator (`Accelerator`):
+            An `Accelerator` object
+        batch_size (`int`, *optional*):
+            The batch size for the train and validation DataLoaders.
+    """
+    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+    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:
+    with accelerator.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):
+        # On TPU it's best to pad everything to the same length or training will be very slow.
+        if accelerator.distributed_type == DistributedType.TPU:
+            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
+        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(
+        tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size
+    )
+    eval_dataloader = DataLoader(
+        tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
+    )
+
+    return train_dataloader, eval_dataloader
+
+
+# For testing only
+if os.environ.get("TESTING_MOCKED_DATALOADERS", None) == "1":
+    from accelerate.test_utils.training import mocked_dataloaders
+
+    get_dataloaders = mocked_dataloaders  # noqa: F811
+
+
+def training_function(config, args):
+    # For testing only
+    if os.environ.get("TESTING_MOCKED_DATALOADERS", None) == "1":
+        config["num_epochs"] = 2
+    # Initialize accelerator
+    accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)
+    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
+    lr = config["lr"]
+    num_epochs = int(config["num_epochs"])
+    seed = int(config["seed"])
+    batch_size = int(config["batch_size"])
+
+    metric = evaluate.load("glue", "mrpc")
+
+    # New Code #
+    # We now can define an inner training loop function. It should take a batch size as the only parameter,
+    # and build the dataloaders in there.
+    # It also gets our decorator
+    @find_executable_batch_size(starting_batch_size=batch_size)
+    def inner_training_loop(batch_size):
+        # And now just move everything below under this function
+        # We need to bring in the Accelerator object from earlier
+        nonlocal accelerator
+        # And reset all of its attributes that could hold onto any memory:
+        accelerator.free_memory()
+
+        # Then we can declare the model, optimizer, and everything else:
+        set_seed(seed)
+
+        # Instantiate the model (we build the model here so that the seed also control new weights initialization)
+        model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", return_dict=True)
+
+        # We could avoid this line since the accelerator is set with `device_placement=True` (default value).
+        # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
+        # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
+        model = model.to(accelerator.device)
+
+        # Instantiate optimizer
+        optimizer = AdamW(params=model.parameters(), lr=lr)
+        train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size)
+
+        # Instantiate scheduler
+        lr_scheduler = get_linear_schedule_with_warmup(
+            optimizer=optimizer,
+            num_warmup_steps=100,
+            num_training_steps=(len(train_dataloader) * num_epochs),
+        )
+
+        # 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, lr_scheduler = accelerator.prepare(
+            model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+        )
+
+        # Now we train the model
+        for epoch in range(num_epochs):
+            model.train()
+            for step, batch in enumerate(train_dataloader):
+                # We could avoid this line since we set the accelerator with `device_placement=True`.
+                batch.to(accelerator.device)
+                outputs = model(**batch)
+                loss = outputs.loss
+                accelerator.backward(loss)
+                optimizer.step()
+                lr_scheduler.step()
+                optimizer.zero_grad()
+
+            model.eval()
+            for step, batch in enumerate(eval_dataloader):
+                # We could avoid this line since we set the accelerator with `device_placement=True`.
+                batch.to(accelerator.device)
+                with torch.no_grad():
+                    outputs = model(**batch)
+                predictions = outputs.logits.argmax(dim=-1)
+                predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]))
+                metric.add_batch(
+                    predictions=predictions,
+                    references=references,
+                )
+
+            eval_metric = metric.compute()
+            # Use accelerator.print to print only on the main process.
+            accelerator.print(f"epoch {epoch}:", eval_metric)
+
+    # New Code #
+    # And call it at the end with no arguments
+    # Note: You could also refactor this outside of your training loop function
+    inner_training_loop()
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Simple example of training script.")
+    parser.add_argument(
+        "--mixed_precision",
+        type=str,
+        default="no",
+        choices=["no", "fp16", "bf16"],
+        help="Whether to use mixed precision. Choose"
+        "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
+        "and an Nvidia Ampere GPU.",
+    )
+    parser.add_argument("--cpu", action="store_true", help="If passed, will train on the CPU.")
+    args = parser.parse_args()
+    config = {"lr": 2e-5, "num_epochs": 3, "seed": 42, "batch_size": 16}
+    training_function(config, args)
+
+
+if __name__ == "__main__":
+    main()

examples/by_feature/multi_process_metrics.py

@@ -0,0 +1,225 @@
+# coding=utf-8
+# Copyright 2022 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 argparse
+import os
+
+import torch
+from torch.optim import AdamW
+from torch.utils.data import DataLoader
+
+import evaluate
+from accelerate import Accelerator, DistributedType
+from datasets import load_dataset
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed
+
+
+########################################################################
+# This is a fully working simple example to use Accelerate,
+# specifically showcasing how to properly calculate the metrics on the
+# validation dataset when in a distributed system, and builds off the
+# `nlp_example.py` script.
+#
+# This example trains a Bert base model on GLUE MRPC
+# in any of the following settings (with the same script):
+#   - single CPU or single GPU
+#   - multi GPUS (using PyTorch distributed mode)
+#   - (multi) TPUs
+#   - fp16 (mixed-precision) or fp32 (normal precision)
+#
+# To help focus on the differences in the code, building `DataLoaders`
+# was refactored into its own function.
+# New additions from the base script can be found quickly by
+# looking for the # New Code # tags
+#
+# To run it in each of these various modes, follow the instructions
+# in the readme for examples:
+# https://github.com/huggingface/accelerate/tree/main/examples
+#
+########################################################################
+
+
+MAX_GPU_BATCH_SIZE = 16
+EVAL_BATCH_SIZE = 32
+
+
+def get_dataloaders(accelerator: Accelerator, batch_size: int = 16):
+    """
+    Creates a set of `DataLoader`s for the `glue` dataset,
+    using "bert-base-cased" as the tokenizer.
+
+    Args:
+        accelerator (`Accelerator`):
+            An `Accelerator` object
+        batch_size (`int`, *optional*):
+            The batch size for the train and validation DataLoaders.
+    """
+    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+    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:
+    with accelerator.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):
+        # On TPU it's best to pad everything to the same length or training will be very slow.
+        if accelerator.distributed_type == DistributedType.TPU:
+            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
+        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(
+        tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size
+    )
+    eval_dataloader = DataLoader(
+        tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
+    )
+
+    return train_dataloader, eval_dataloader
+
+
+# For testing only
+if os.environ.get("TESTING_MOCKED_DATALOADERS", None) == "1":
+    from accelerate.test_utils.training import mocked_dataloaders
+
+    get_dataloaders = mocked_dataloaders  # noqa: F811
+
+
+def training_function(config, args):
+    # For testing only
+    if os.environ.get("TESTING_MOCKED_DATALOADERS", None) == "1":
+        config["num_epochs"] = 2
+    # Initialize accelerator
+    accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)
+    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
+    lr = config["lr"]
+    num_epochs = int(config["num_epochs"])
+    seed = int(config["seed"])
+    batch_size = int(config["batch_size"])
+
+    metric = evaluate.load("glue", "mrpc")
+
+    # If the batch size is too big we use gradient accumulation
+    gradient_accumulation_steps = 1
+    if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
+        gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
+        batch_size = MAX_GPU_BATCH_SIZE
+
+    set_seed(seed)
+    train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size)
+    # Instantiate the model (we build the model here so that the seed also control new weights initialization)
+    model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", return_dict=True)
+
+    # We could avoid this line since the accelerator is set with `device_placement=True` (default value).
+    # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
+    # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
+    model = model.to(accelerator.device)
+
+    # Instantiate optimizer
+    optimizer = AdamW(params=model.parameters(), lr=lr)
+
+    # Instantiate scheduler
+    lr_scheduler = get_linear_schedule_with_warmup(
+        optimizer=optimizer,
+        num_warmup_steps=100,
+        num_training_steps=(len(train_dataloader) * num_epochs) // gradient_accumulation_steps,
+    )
+
+    # 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, lr_scheduler = accelerator.prepare(
+        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+    )
+
+    # Now we train the model
+    for epoch in range(num_epochs):
+        model.train()
+        for step, batch in enumerate(train_dataloader):
+            # We could avoid this line since we set the accelerator with `device_placement=True`.
+            batch.to(accelerator.device)
+            outputs = model(**batch)
+            loss = outputs.loss
+            loss = loss / gradient_accumulation_steps
+            accelerator.backward(loss)
+            if step % gradient_accumulation_steps == 0:
+                optimizer.step()
+                lr_scheduler.step()
+                optimizer.zero_grad()
+
+        model.eval()
+        samples_seen = 0
+        for step, batch in enumerate(eval_dataloader):
+            # We could avoid this line since we set the accelerator with `device_placement=True`.
+            batch.to(accelerator.device)
+            with torch.no_grad():
+                outputs = model(**batch)
+            predictions = outputs.logits.argmax(dim=-1)
+            predictions, references = accelerator.gather((predictions, batch["labels"]))
+            # New Code #
+            # First we check if it's a distributed system
+            if accelerator.use_distributed:
+                # Then see if we're on the last batch of our eval dataloader
+                if step == len(eval_dataloader) - 1:
+                    # Last batch needs to be truncated on distributed systems as it contains additional samples
+                    predictions = predictions[: len(eval_dataloader.dataset) - samples_seen]
+                    references = references[: len(eval_dataloader.dataset) - samples_seen]
+                else:
+                    # Otherwise we add the number of samples seen
+                    samples_seen += references.shape[0]
+            # All of this can be avoided if you use `Accelerator.gather_for_metrics` instead of `Accelerator.gather`:
+            # accelerator.gather_for_metrics((predictions, batch["labels"]))
+            metric.add_batch(
+                predictions=predictions,
+                references=references,
+            )
+
+        eval_metric = metric.compute()
+        # Use accelerator.print to print only on the main process.
+        accelerator.print(f"epoch {epoch}:", eval_metric)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Simple example of training script.")
+    parser.add_argument(
+        "--mixed_precision",
+        type=str,
+        default="no",
+        choices=["no", "fp16", "bf16"],
+        help="Whether to use mixed precision. Choose"
+        "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
+        "and an Nvidia Ampere GPU.",
+    )
+    parser.add_argument("--cpu", action="store_true", help="If passed, will train on the CPU.")
+    args = parser.parse_args()
+    config = {"lr": 2e-5, "num_epochs": 3, "seed": 42, "batch_size": 16}
+    training_function(config, args)
+
+
+if __name__ == "__main__":
+    main()

examples/by_feature/tracking.py

@@ -0,0 +1,263 @@
+# coding=utf-8
+# Copyright 2021 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 argparse
+import os
+
+import torch
+from torch.optim import AdamW
+from torch.utils.data import DataLoader
+
+import evaluate
+from accelerate import Accelerator, DistributedType
+from datasets import load_dataset
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed
+
+
+########################################################################
+# This is a fully working simple example to use Accelerate,
+# specifically showcasing the experiment tracking capability,
+# and builds off the `nlp_example.py` script.
+#
+# This example trains a Bert base model on GLUE MRPC
+# in any of the following settings (with the same script):
+#   - single CPU or single GPU
+#   - multi GPUS (using PyTorch distributed mode)
+#   - (multi) TPUs
+#   - fp16 (mixed-precision) or fp32 (normal precision)
+#
+# To help focus on the differences in the code, building `DataLoaders`
+# was refactored into its own function.
+# New additions from the base script can be found quickly by
+# looking for the # New Code # tags
+#
+# To run it in each of these various modes, follow the instructions
+# in the readme for examples:
+# https://github.com/huggingface/accelerate/tree/main/examples
+#
+########################################################################
+
+MAX_GPU_BATCH_SIZE = 16
+EVAL_BATCH_SIZE = 32
+
+
+def get_dataloaders(accelerator: Accelerator, batch_size: int = 16):
+    """
+    Creates a set of `DataLoader`s for the `glue` dataset,
+    using "bert-base-cased" as the tokenizer.
+
+    Args:
+        accelerator (`Accelerator`):
+            An `Accelerator` object
+        batch_size (`int`, *optional*):
+            The batch size for the train and validation DataLoaders.
+    """
+    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+    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:
+    with accelerator.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):
+        # On TPU it's best to pad everything to the same length or training will be very slow.
+        if accelerator.distributed_type == DistributedType.TPU:
+            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
+        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(
+        tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size
+    )
+    eval_dataloader = DataLoader(
+        tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
+    )
+
+    return train_dataloader, eval_dataloader
+
+
+# For testing only
+if os.environ.get("TESTING_MOCKED_DATALOADERS", None) == "1":
+    from accelerate.test_utils.training import mocked_dataloaders
+
+    get_dataloaders = mocked_dataloaders  # noqa: F811
+
+
+def training_function(config, args):
+    # For testing only
+    if os.environ.get("TESTING_MOCKED_DATALOADERS", None) == "1":
+        config["num_epochs"] = 2
+    # Initialize Accelerator
+
+    # New Code #
+    # We pass in "all" to `log_with` to grab all available trackers in the environment
+    # Note: If using a custom `Tracker` class, should be passed in here such as:
+    # >>> log_with = ["all", MyCustomTrackerClassInstance()]
+    if args.with_tracking:
+        accelerator = Accelerator(
+            cpu=args.cpu, mixed_precision=args.mixed_precision, log_with="all", logging_dir=args.logging_dir
+        )
+    else:
+        accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)
+    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
+    lr = config["lr"]
+    num_epochs = int(config["num_epochs"])
+    seed = int(config["seed"])
+    batch_size = int(config["batch_size"])
+    set_seed(seed)
+
+    train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size)
+    metric = evaluate.load("glue", "mrpc")
+
+    # If the batch size is too big we use gradient accumulation
+    gradient_accumulation_steps = 1
+    if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
+        gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
+        batch_size = MAX_GPU_BATCH_SIZE
+
+    # Instantiate the model (we build the model here so that the seed also control new weights initialization)
+    model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", return_dict=True)
+
+    # We could avoid this line since the accelerator is set with `device_placement=True` (default value).
+    # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
+    # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
+    model = model.to(accelerator.device)
+
+    # Instantiate optimizer
+    optimizer = AdamW(params=model.parameters(), lr=lr)
+
+    # Instantiate scheduler
+    lr_scheduler = get_linear_schedule_with_warmup(
+        optimizer=optimizer,
+        num_warmup_steps=100,
+        num_training_steps=(len(train_dataloader) * num_epochs) // gradient_accumulation_steps,
+    )
+
+    # 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, lr_scheduler = accelerator.prepare(
+        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+    )
+
+    # New Code #
+    # We need to initialize the trackers we use. Overall configurations can also be stored
+    if args.with_tracking:
+        run = os.path.split(__file__)[-1].split(".")[0]
+        accelerator.init_trackers(run, config)
+
+    # Now we train the model
+    for epoch in range(num_epochs):
+        model.train()
+        # New Code #
+        # For our tracking example, we will log the total loss of each epoch
+        if args.with_tracking:
+            total_loss = 0
+        for step, batch in enumerate(train_dataloader):
+            # We could avoid this line since we set the accelerator with `device_placement=True`.
+            batch.to(accelerator.device)
+            outputs = model(**batch)
+            loss = outputs.loss
+            # New Code #
+            if args.with_tracking:
+                total_loss += loss.detach().float()
+            loss = loss / gradient_accumulation_steps
+            accelerator.backward(loss)
+            if step % gradient_accumulation_steps == 0:
+                optimizer.step()
+                lr_scheduler.step()
+                optimizer.zero_grad()
+
+        model.eval()
+        for step, batch in enumerate(eval_dataloader):
+            # We could avoid this line since we set the accelerator with `device_placement=True` (the default).
+            batch.to(accelerator.device)
+            with torch.no_grad():
+                outputs = model(**batch)
+            predictions = outputs.logits.argmax(dim=-1)
+            predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]))
+            metric.add_batch(
+                predictions=predictions,
+                references=references,
+            )
+
+        eval_metric = metric.compute()
+        # Use accelerator.print to print only on the main process.
+        accelerator.print(f"epoch {epoch}:", eval_metric)
+
+        # New Code #
+        # To actually log, we call `Accelerator.log`
+        # The values passed can be of `str`, `int`, `float` or `dict` of `str` to `float`/`int`
+        if args.with_tracking:
+            accelerator.log(
+                {
+                    "accuracy": eval_metric["accuracy"],
+                    "f1": eval_metric["f1"],
+                    "train_loss": total_loss.item() / len(train_dataloader),
+                    "epoch": epoch,
+                },
+                step=epoch,
+            )
+
+    # New Code #
+    # When a run is finished, you should call `accelerator.end_training()`
+    # to close all of the open trackers
+    if args.with_tracking:
+        accelerator.end_training()
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Simple example of training script.")
+    parser.add_argument(
+        "--mixed_precision",
+        type=str,
+        default="no",
+        choices=["no", "fp16", "bf16"],
+        help="Whether to use mixed precision. Choose"
+        "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
+        "and an Nvidia Ampere GPU.",
+    )
+    parser.add_argument("--cpu", action="store_true", help="If passed, will train on the CPU.")
+    parser.add_argument(
+        "--with_tracking",
+        action="store_true",
+        help="Whether to load in all available experiment trackers from the environment and use them for logging.",
+    )
+    parser.add_argument(
+        "--logging_dir",
+        type=str,
+        default="logs",
+        help="Location on where to store experiment tracking logs`",
+    )
+    args = parser.parse_args()
+    config = {"lr": 2e-5, "num_epochs": 3, "seed": 42, "batch_size": 16}
+    training_function(config, args)
+
+
+if __name__ == "__main__":
+    main()

examples/complete_cv_example.py

@@ -0,0 +1,315 @@
+# coding=utf-8
+# Copyright 2021 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 argparse
+import os
+import re
+
+import numpy as np
+import torch
+from torch.optim.lr_scheduler import OneCycleLR
+from torch.utils.data import DataLoader, Dataset
+
+import PIL
+from accelerate import Accelerator
+from timm import create_model
+from torchvision.transforms import Compose, RandomResizedCrop, Resize, ToTensor
+
+
+########################################################################
+# This is a fully working simple example to use Accelerate
+#
+# This example trains a ResNet50 on the Oxford-IIT Pet Dataset
+# in any of the following settings (with the same script):
+#   - single CPU or single GPU
+#   - multi GPUS (using PyTorch distributed mode)
+#   - (multi) TPUs
+#   - fp16 (mixed-precision) or fp32 (normal precision)
+#
+# To run it in each of these various modes, follow the instructions
+# in the readme for examples:
+# https://github.com/huggingface/accelerate/tree/main/examples
+#
+########################################################################
+
+
+# Function to get the label from the filename
+def extract_label(fname):
+    stem = fname.split(os.path.sep)[-1]
+    return re.search(r"^(.*)_\d+\.jpg$", stem).groups()[0]
+
+
+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}
+
+
+def training_function(config, args):
+    # Initialize accelerator
+    if args.with_tracking:
+        accelerator = Accelerator(
+            cpu=args.cpu, mixed_precision=args.mixed_precision, log_with="all", logging_dir=args.logging_dir
+        )
+    else:
+        accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)
+
+    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
+    lr = config["lr"]
+    num_epochs = int(config["num_epochs"])
+    seed = int(config["seed"])
+    batch_size = int(config["batch_size"])
+    image_size = config["image_size"]
+    if not isinstance(image_size, (list, tuple)):
+        image_size = (image_size, image_size)
+
+    # Parse out whether we are saving every epoch or after a certain number of batches
+    if hasattr(args.checkpointing_steps, "isdigit"):
+        if args.checkpointing_steps == "epoch":
+            checkpointing_steps = args.checkpointing_steps
+        elif args.checkpointing_steps.isdigit():
+            checkpointing_steps = int(args.checkpointing_steps)
+        else:
+            raise ValueError(
+                f"Argument `checkpointing_steps` must be either a number or `epoch`. `{args.checkpointing_steps}` passed."
+            )
+    else:
+        checkpointing_steps = None
+
+    # We need to initialize the trackers we use, and also store our configuration
+    if args.with_tracking:
+        run = os.path.split(__file__)[-1].split(".")[0]
+        accelerator.init_trackers(run, config)
+
+    # Grab all the image filenames
+    file_names = [os.path.join(args.data_dir, fname) for fname in os.listdir(args.data_dir) if fname.endswith(".jpg")]
+
+    # Build the label correspondences
+    all_labels = [extract_label(fname) for fname in file_names]
+    id_to_label = list(set(all_labels))
+    id_to_label.sort()
+    label_to_id = {lbl: i for i, lbl in enumerate(id_to_label)}
+
+    # Set the seed before splitting the data.
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+    # Split our filenames between train and validation
+    random_perm = np.random.permutation(len(file_names))
+    cut = int(0.8 * len(file_names))
+    train_split = random_perm[:cut]
+    eval_split = random_perm[cut:]
+
+    # For training we use a simple RandomResizedCrop
+    train_tfm = Compose([RandomResizedCrop(image_size, scale=(0.5, 1.0)), ToTensor()])
+    train_dataset = PetsDataset(
+        [file_names[i] for i in train_split], image_transform=train_tfm, label_to_id=label_to_id
+    )
+
+    # For evaluation, we use a deterministic Resize
+    eval_tfm = Compose([Resize(image_size), ToTensor()])
+    eval_dataset = PetsDataset([file_names[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, num_workers=4)
+
+    # Instantiate the model (we build the model here so that the seed also control new weights initialization)
+    model = create_model("resnet50d", pretrained=True, num_classes=len(label_to_id))
+
+    # We could avoid this line since the accelerator is set with `device_placement=True` (default value).
+    # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
+    # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
+    model = model.to(accelerator.device)
+
+    # Freezing the base model
+    for param in model.parameters():
+        param.requires_grad = False
+    for param in model.get_classifier().parameters():
+        param.requires_grad = True
+
+    # We normalize the batches of images to be a bit faster.
+    mean = torch.tensor(model.default_cfg["mean"])[None, :, None, None].to(accelerator.device)
+    std = torch.tensor(model.default_cfg["std"])[None, :, None, None].to(accelerator.device)
+
+    # Instantiate optimizer
+    optimizer = torch.optim.Adam(params=model.parameters(), lr=lr / 25)
+
+    # Instantiate learning rate scheduler
+    lr_scheduler = OneCycleLR(optimizer=optimizer, max_lr=lr, epochs=num_epochs, steps_per_epoch=len(train_dataloader))
+
+    # 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, lr_scheduler = accelerator.prepare(
+        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+    )
+    # We need to keep track of how many total steps we have iterated over
+    overall_step = 0
+    # We also need to keep track of the stating epoch so files are named properly
+    starting_epoch = 0
+
+    # Potentially load in the weights and states from a previous save
+    if args.resume_from_checkpoint:
+        if args.resume_from_checkpoint is not None or args.resume_from_checkpoint != "":
+            accelerator.print(f"Resumed from checkpoint: {args.resume_from_checkpoint}")
+            accelerator.load_state(args.resume_from_checkpoint)
+            path = os.path.basename(args.resume_from_checkpoint)
+        else:
+            # Get the most recent checkpoint
+            dirs = [f.name for f in os.scandir(os.getcwd()) if f.is_dir()]
+            dirs.sort(key=os.path.getctime)
+            path = dirs[-1]  # Sorts folders by date modified, most recent checkpoint is the last
+        # Extract `epoch_{i}` or `step_{i}`
+        training_difference = os.path.splitext(path)[0]
+
+        if "epoch" in training_difference:
+            starting_epoch = int(training_difference.replace("epoch_", "")) + 1
+            resume_step = None
+        else:
+            resume_step = int(training_difference.replace("step_", ""))
+            starting_epoch = resume_step // len(train_dataloader)
+            resume_step -= starting_epoch * len(train_dataloader)
+
+    # Now we train the model
+    for epoch in range(starting_epoch, num_epochs):
+        model.train()
+        if args.with_tracking:
+            total_loss = 0
+        for step, batch in enumerate(train_dataloader):
+            # We need to skip steps until we reach the resumed step
+            if args.resume_from_checkpoint and epoch == starting_epoch:
+                if resume_step is not None and step < resume_step:
+                    overall_step += 1
+                    continue
+            # We could avoid this line since we set the accelerator with `device_placement=True`.
+            batch = {k: v.to(accelerator.device) for k, v in batch.items()}
+            inputs = (batch["image"] - mean) / std
+            outputs = model(inputs)
+            loss = torch.nn.functional.cross_entropy(outputs, batch["label"])
+            # We keep track of the loss at each epoch
+            if args.with_tracking:
+                total_loss += loss.detach().float()
+            accelerator.backward(loss)
+            optimizer.step()
+            lr_scheduler.step()
+            optimizer.zero_grad()
+            overall_step += 1
+            if isinstance(checkpointing_steps, int):
+                output_dir = f"step_{overall_step}"
+                if overall_step % checkpointing_steps == 0:
+                    if args.output_dir is not None:
+                        output_dir = os.path.join(args.output_dir, output_dir)
+                    accelerator.save_state(output_dir)
+        model.eval()
+        accurate = 0
+        for step, batch in enumerate(eval_dataloader):
+            # We could avoid this line since we set the accelerator with `device_placement=True`.
+            batch = {k: v.to(accelerator.device) for k, v in batch.items()}
+            inputs = (batch["image"] - mean) / std
+            with torch.no_grad():
+                outputs = model(inputs)
+            predictions = outputs.argmax(dim=-1)
+            predictions, references = accelerator.gather_for_metrics((predictions, batch["label"]))
+            accurate_preds = predictions == references
+            accurate += accurate_preds.long().sum()
+
+        eval_metric = accurate.item() / accelerator.gradient_state.samples_seen
+        # Use accelerator.print to print only on the main process.
+        accelerator.print(f"epoch {epoch}: {100 * eval_metric:.2f}")
+        if args.with_tracking:
+            accelerator.log(
+                {
+                    "accuracy": 100 * eval_metric,
+                    "train_loss": total_loss.item() / len(train_dataloader),
+                    "epoch": epoch,
+                },
+                step=overall_step,
+            )
+        if checkpointing_steps == "epoch":
+            output_dir = f"epoch_{epoch}"
+            if args.output_dir is not None:
+                output_dir = os.path.join(args.output_dir, output_dir)
+            accelerator.save_state(output_dir)
+
+    if args.with_tracking:
+        accelerator.end_training()
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Simple example of training script.")
+    parser.add_argument("--data_dir", required=True, help="The data folder on disk.")
+    parser.add_argument("--fp16", action="store_true", help="If passed, will use FP16 training.")
+    parser.add_argument(
+        "--mixed_precision",
+        type=str,
+        default="no",
+        choices=["no", "fp16", "bf16"],
+        help="Whether to use mixed precision. Choose"
+        "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
+        "and an Nvidia Ampere GPU.",
+    )
+    parser.add_argument("--cpu", action="store_true", help="If passed, will train on the CPU.")
+    parser.add_argument(
+        "--checkpointing_steps",
+        type=str,
+        default=None,
+        help="Whether the various states should be saved at the end of every n steps, or 'epoch' for each epoch.",
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default=".",
+        help="Optional save directory where all checkpoint folders will be stored. Default is the current working directory.",
+    )
+    parser.add_argument(
+        "--resume_from_checkpoint",
+        type=str,
+        default=None,
+        help="If the training should continue from a checkpoint folder.",
+    )
+    parser.add_argument(
+        "--with_tracking",
+        action="store_true",
+        help="Whether to load in all available experiment trackers from the environment and use them for logging.",
+    )
+    parser.add_argument(
+        "--logging_dir",
+        type=str,
+        default="logs",
+        help="Location on where to store experiment tracking logs`",
+    )
+    args = parser.parse_args()
+    config = {"lr": 3e-2, "num_epochs": 3, "seed": 42, "batch_size": 64, "image_size": 224}
+    training_function(config, args)
+
+
+if __name__ == "__main__":
+    main()

examples/complete_nlp_example.py

@@ -0,0 +1,296 @@
+# coding=utf-8
+# Copyright 2021 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 argparse
+import os
+
+import torch
+from torch.optim import AdamW
+from torch.utils.data import DataLoader
+
+import evaluate
+from accelerate import Accelerator, DistributedType
+from datasets import load_dataset
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed
+
+
+########################################################################
+# This is a fully working simple example to use Accelerate
+#
+# This example trains a Bert base model on GLUE MRPC
+# in any of the following settings (with the same script):
+#   - single CPU or single GPU
+#   - multi GPUS (using PyTorch distributed mode)
+#   - (multi) TPUs
+#   - fp16 (mixed-precision) or fp32 (normal precision)
+#
+# This example also demonstrates the checkpointing and sharding capabilities
+#
+# To run it in each of these various modes, follow the instructions
+# in the readme for examples:
+# https://github.com/huggingface/accelerate/tree/main/examples
+#
+########################################################################
+
+
+MAX_GPU_BATCH_SIZE = 16
+EVAL_BATCH_SIZE = 32
+
+
+def training_function(config, args):
+    # Initialize accelerator
+    if args.with_tracking:
+        accelerator = Accelerator(
+            cpu=args.cpu, mixed_precision=args.mixed_precision, log_with="all", logging_dir=args.logging_dir
+        )
+    else:
+        accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)
+
+    if hasattr(args.checkpointing_steps, "isdigit"):
+        if args.checkpointing_steps == "epoch":
+            checkpointing_steps = args.checkpointing_steps
+        elif args.checkpointing_steps.isdigit():
+            checkpointing_steps = int(args.checkpointing_steps)
+        else:
+            raise ValueError(
+                f"Argument `checkpointing_steps` must be either a number or `epoch`. `{args.checkpointing_steps}` passed."
+            )
+    else:
+        checkpointing_steps = None
+    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
+    lr = config["lr"]
+    num_epochs = int(config["num_epochs"])
+    seed = int(config["seed"])
+    batch_size = int(config["batch_size"])
+
+    # We need to initialize the trackers we use, and also store our configuration
+    if args.with_tracking:
+        run = os.path.split(__file__)[-1].split(".")[0]
+        accelerator.init_trackers(run, config)
+
+    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+    datasets = load_dataset("glue", "mrpc")
+    metric = evaluate.load("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:
+    with accelerator.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")
+
+    # If the batch size is too big we use gradient accumulation
+    gradient_accumulation_steps = 1
+    if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
+        gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
+        batch_size = MAX_GPU_BATCH_SIZE
+
+    def collate_fn(examples):
+        # On TPU it's best to pad everything to the same length or training will be very slow.
+        if accelerator.distributed_type == DistributedType.TPU:
+            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
+        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(
+        tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size
+    )
+    eval_dataloader = DataLoader(
+        tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
+    )
+
+    set_seed(seed)
+
+    # Instantiate the model (we build the model here so that the seed also control new weights initialization)
+    model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", return_dict=True)
+
+    # We could avoid this line since the accelerator is set with `device_placement=True` (default value).
+    # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
+    # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
+    model = model.to(accelerator.device)
+
+    # Instantiate optimizer
+    optimizer = AdamW(params=model.parameters(), lr=lr)
+
+    # Instantiate scheduler
+    lr_scheduler = get_linear_schedule_with_warmup(
+        optimizer=optimizer,
+        num_warmup_steps=100,
+        num_training_steps=(len(train_dataloader) * num_epochs) // gradient_accumulation_steps,
+    )
+
+    # 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, lr_scheduler = accelerator.prepare(
+        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+    )
+
+    # We need to keep track of how many total steps we have iterated over
+    overall_step = 0
+    # We also need to keep track of the stating epoch so files are named properly
+    starting_epoch = 0
+
+    # Potentially load in the weights and states from a previous save
+    if args.resume_from_checkpoint:
+        if args.resume_from_checkpoint is not None or args.resume_from_checkpoint != "":
+            accelerator.print(f"Resumed from checkpoint: {args.resume_from_checkpoint}")
+            accelerator.load_state(args.resume_from_checkpoint)
+            path = os.path.basename(args.resume_from_checkpoint)
+        else:
+            # Get the most recent checkpoint
+            dirs = [f.name for f in os.scandir(os.getcwd()) if f.is_dir()]
+            dirs.sort(key=os.path.getctime)
+            path = dirs[-1]  # Sorts folders by date modified, most recent checkpoint is the last
+        # Extract `epoch_{i}` or `step_{i}`
+        training_difference = os.path.splitext(path)[0]
+
+        if "epoch" in training_difference:
+            starting_epoch = int(training_difference.replace("epoch_", "")) + 1
+            resume_step = None
+        else:
+            resume_step = int(training_difference.replace("step_", ""))
+            starting_epoch = resume_step // len(train_dataloader)
+            resume_step -= starting_epoch * len(train_dataloader)
+
+    # Now we train the model
+    for epoch in range(starting_epoch, num_epochs):
+        model.train()
+        if args.with_tracking:
+            total_loss = 0
+        for step, batch in enumerate(train_dataloader):
+            # We need to skip steps until we reach the resumed step
+            if args.resume_from_checkpoint and epoch == starting_epoch:
+                if resume_step is not None and step < resume_step:
+                    overall_step += 1
+                    continue
+            # We could avoid this line since we set the accelerator with `device_placement=True`.
+            batch.to(accelerator.device)
+            outputs = model(**batch)
+            loss = outputs.loss
+            loss = loss / gradient_accumulation_steps
+            # We keep track of the loss at each epoch
+            if args.with_tracking:
+                total_loss += loss.detach().float()
+            accelerator.backward(loss)
+            if step % gradient_accumulation_steps == 0:
+                optimizer.step()
+                lr_scheduler.step()
+                optimizer.zero_grad()
+
+            overall_step += 1
+
+            if isinstance(checkpointing_steps, int):
+                output_dir = f"step_{overall_step}"
+                if overall_step % checkpointing_steps == 0:
+                    if args.output_dir is not None:
+                        output_dir = os.path.join(args.output_dir, output_dir)
+                    accelerator.save_state(output_dir)
+
+        model.eval()
+        for step, batch in enumerate(eval_dataloader):
+            # We could avoid this line since we set the accelerator with `device_placement=True`.
+            batch.to(accelerator.device)
+            with torch.no_grad():
+                outputs = model(**batch)
+            predictions = outputs.logits.argmax(dim=-1)
+            predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]))
+            metric.add_batch(
+                predictions=predictions,
+                references=references,
+            )
+
+        eval_metric = metric.compute()
+        # Use accelerator.print to print only on the main process.
+        accelerator.print(f"epoch {epoch}:", eval_metric)
+        if args.with_tracking:
+            accelerator.log(
+                {
+                    "accuracy": eval_metric["accuracy"],
+                    "f1": eval_metric["f1"],
+                    "train_loss": total_loss.item() / len(train_dataloader),
+                    "epoch": epoch,
+                },
+                step=epoch,
+            )
+
+        if checkpointing_steps == "epoch":
+            output_dir = f"epoch_{epoch}"
+            if args.output_dir is not None:
+                output_dir = os.path.join(args.output_dir, output_dir)
+            accelerator.save_state(output_dir)
+
+    if args.with_tracking:
+        accelerator.end_training()
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Simple example of training script.")
+    parser.add_argument(
+        "--mixed_precision",
+        type=str,
+        default="no",
+        choices=["no", "fp16", "bf16"],
+        help="Whether to use mixed precision. Choose"
+        "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
+        "and an Nvidia Ampere GPU.",
+    )
+    parser.add_argument("--cpu", action="store_true", help="If passed, will train on the CPU.")
+    parser.add_argument(
+        "--checkpointing_steps",
+        type=str,
+        default=None,
+        help="Whether the various states should be saved at the end of every n steps, or 'epoch' for each epoch.",
+    )
+    parser.add_argument(
+        "--resume_from_checkpoint",
+        type=str,
+        default=None,
+        help="If the training should continue from a checkpoint folder.",
+    )
+    parser.add_argument(
+        "--with_tracking",
+        action="store_true",
+        help="Whether to load in all available experiment trackers from the environment and use them for logging.",
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default=".",
+        help="Optional save directory where all checkpoint folders will be stored. Default is the current working directory.",
+    )
+    parser.add_argument(
+        "--logging_dir",
+        type=str,
+        default="logs",
+        help="Location on where to store experiment tracking logs`",
+    )
+    args = parser.parse_args()
+    config = {"lr": 2e-5, "num_epochs": 3, "seed": 42, "batch_size": 16}
+    training_function(config, args)
+
+
+if __name__ == "__main__":
+    main()

examples/cv_example.py

@@ -73,7 +73,7 @@ class PetsDataset(Dataset):
 
 def training_function(config, args):
     # Initialize accelerator
-    accelerator = Accelerator(fp16=args.fp16, cpu=args.cpu)
+    accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)
 
     # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
     lr = config["lr"]
@@ -139,17 +139,16 @@ def training_function(config, args):
     # Instantiate optimizer
     optimizer = torch.optim.Adam(params=model.parameters(), lr=lr / 25)
 
+    # Instantiate learning rate scheduler
+    lr_scheduler = OneCycleLR(optimizer=optimizer, max_lr=lr, epochs=num_epochs, steps_per_epoch=len(train_dataloader))
+
     # 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
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
     )
 
-    # Instantiate learning rate scheduler after preparing the training dataloader as the prepare method
-    # may change its length.
-    lr_scheduler = OneCycleLR(optimizer=optimizer, max_lr=lr, epochs=num_epochs, steps_per_epoch=len(train_dataloader))
-
     # Now we train the model
     for epoch in range(num_epochs):
         model.train()
@@ -167,14 +166,15 @@ def training_function(config, args):
         model.eval()
         accurate = 0
         num_elems = 0
-        for step, batch in enumerate(eval_dataloader):
+        for _, batch in enumerate(eval_dataloader):
             # We could avoid this line since we set the accelerator with `device_placement=True`.
             batch = {k: v.to(accelerator.device) for k, v in batch.items()}
             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"])
+            predictions, references = accelerator.gather_for_metrics((predictions, batch["label"]))
+            accurate_preds = predictions == references
             num_elems += accurate_preds.shape[0]
             accurate += accurate_preds.long().sum()
 
@@ -186,7 +186,21 @@ def training_function(config, args):
 def main():
     parser = argparse.ArgumentParser(description="Simple example of training script.")
     parser.add_argument("--data_dir", required=True, help="The data folder on disk.")
-    parser.add_argument("--fp16", action="store_true", help="If passed, will use FP16 training.")
+    parser.add_argument(
+        "--mixed_precision",
+        type=str,
+        default="no",
+        choices=["no", "fp16", "bf16"],
+        help="Whether to use mixed precision. Choose"
+        "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
+        "and an Nvidia Ampere GPU.",
+    )
+    parser.add_argument(
+        "--checkpointing_steps",
+        type=str,
+        default=None,
+        help="Whether the various states should be saved at the end of every n steps, or 'epoch' for each epoch.",
+    )
     parser.add_argument("--cpu", action="store_true", help="If passed, will train on the CPU.")
     args = parser.parse_args()
     config = {"lr": 3e-2, "num_epochs": 3, "seed": 42, "batch_size": 64, "image_size": 224}

examples/deepspeed_config_templates/zero_stage1_config.json

@@ -0,0 +1,43 @@
+{
+    "fp16": {
+        "enabled": true,
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+    "optimizer": {
+        "type": "AdamW",
+        "params": {
+            "lr": "auto",
+            "weight_decay": "auto",
+            "torch_adam": true,
+            "adam_w_mode": true
+        }
+    },
+    "scheduler": {
+        "type": "WarmupDecayLR",
+        "params": {
+            "warmup_min_lr": "auto",
+            "warmup_max_lr": "auto",
+            "warmup_num_steps": "auto",
+            "total_num_steps": "auto"
+        }
+    },
+    "zero_optimization": {
+        "stage": 1,
+        "allgather_partitions": true,
+        "allgather_bucket_size": 2e8,
+        "overlap_comm": true,
+        "reduce_scatter": true,
+        "reduce_bucket_size": "auto",
+        "contiguous_gradients": true
+    },
+    "gradient_accumulation_steps": 1,
+    "gradient_clipping": "auto",
+    "steps_per_print": 2000,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}

examples/deepspeed_config_templates/zero_stage2_config.json

@@ -0,0 +1,43 @@
+{
+    "fp16": {
+        "enabled": true,
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+    "optimizer": {
+        "type": "AdamW",
+        "params": {
+            "lr": "auto",
+            "weight_decay": "auto",
+            "torch_adam": true,
+            "adam_w_mode": true
+        }
+    },
+    "scheduler": {
+        "type": "WarmupDecayLR",
+        "params": {
+            "warmup_min_lr": "auto",
+            "warmup_max_lr": "auto",
+            "warmup_num_steps": "auto",
+            "total_num_steps": "auto"
+        }
+    },
+    "zero_optimization": {
+        "stage": 2,
+        "allgather_partitions": true,
+        "allgather_bucket_size": 2e8,
+        "overlap_comm": true,
+        "reduce_scatter": true,
+        "reduce_bucket_size": "auto",
+        "contiguous_gradients": true
+    },
+    "gradient_accumulation_steps": 1,
+    "gradient_clipping": "auto",
+    "steps_per_print": 2000,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}

examples/deepspeed_config_templates/zero_stage2_offload_config.json

@@ -0,0 +1,47 @@
+{
+    "fp16": {
+        "enabled": true,
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+    "optimizer": {
+        "type": "AdamW",
+        "params": {
+            "lr": "auto",
+            "weight_decay": "auto",
+            "torch_adam": true,
+            "adam_w_mode": true
+        }
+    },
+    "scheduler": {
+        "type": "WarmupDecayLR",
+        "params": {
+            "warmup_min_lr": "auto",
+            "warmup_max_lr": "auto",
+            "warmup_num_steps": "auto",
+            "total_num_steps": "auto"
+        }
+    },
+    "zero_optimization": {
+        "stage": 2,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "allgather_partitions": true,
+        "allgather_bucket_size": 2e8,
+        "overlap_comm": true,
+        "reduce_scatter": true,
+        "reduce_bucket_size": "auto",
+        "contiguous_gradients": true
+    },
+    "gradient_accumulation_steps": 1,
+    "gradient_clipping": "auto",
+    "steps_per_print": 2000,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}

examples/deepspeed_config_templates/zero_stage3_config.json

@@ -0,0 +1,44 @@
+{
+    "fp16": {
+        "enabled": true,
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+    "optimizer": {
+        "type": "AdamW",
+        "params": {
+            "lr": "auto",
+            "weight_decay": "auto"
+        }
+    },
+    "scheduler": {
+        "type": "WarmupDecayLR",
+        "params": {
+            "warmup_min_lr": "auto",
+            "warmup_max_lr": "auto",
+            "warmup_num_steps": "auto",
+            "total_num_steps": "auto"
+        }
+    },
+    "zero_optimization": {
+        "stage": 3,
+        "overlap_comm": true,
+        "contiguous_gradients": true,
+        "reduce_bucket_size": "auto",
+        "stage3_prefetch_bucket_size": "auto",
+        "stage3_param_persistence_threshold": "auto",
+        "sub_group_size": 1e9,
+        "stage3_max_live_parameters": 1e9,
+        "stage3_max_reuse_distance": 1e9,
+        "stage3_gather_16bit_weights_on_model_save": "auto"
+    },
+    "gradient_accumulation_steps": 1,
+    "gradient_clipping": "auto",
+    "steps_per_print": 2000,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}

examples/deepspeed_config_templates/zero_stage3_offload_config.json

@@ -0,0 +1,52 @@
+{
+    "fp16": {
+        "enabled": true,
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+    "optimizer": {
+        "type": "AdamW",
+        "params": {
+            "lr": "auto",
+            "weight_decay": "auto"
+        }
+    },
+    "scheduler": {
+        "type": "WarmupDecayLR",
+        "params": {
+            "warmup_min_lr": "auto",
+            "warmup_max_lr": "auto",
+            "warmup_num_steps": "auto",
+            "total_num_steps": "auto"
+        }
+    },
+    "zero_optimization": {
+        "stage": 3,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "offload_param": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "overlap_comm": true,
+        "contiguous_gradients": true,
+        "reduce_bucket_size": "auto",
+        "stage3_prefetch_bucket_size": "auto",
+        "stage3_param_persistence_threshold": "auto",
+        "sub_group_size": 1e9,
+        "stage3_max_live_parameters": 1e9,
+        "stage3_max_reuse_distance": 1e9,
+        "stage3_gather_16bit_weights_on_model_save": "auto"
+    },
+    "gradient_accumulation_steps": 1,
+    "gradient_clipping": "auto",
+    "steps_per_print": 2000,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}

examples/nlp_example.py

@@ -15,17 +15,13 @@
 import argparse
 
 import torch
+from torch.optim import AdamW
 from torch.utils.data import DataLoader
 
+import evaluate
 from accelerate import Accelerator, DistributedType
-from datasets import load_dataset, load_metric
-from transformers import (
-    AdamW,
-    AutoModelForSequenceClassification,
-    AutoTokenizer,
-    get_linear_schedule_with_warmup,
-    set_seed,
-)
+from datasets import load_dataset
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed
 
 
 ########################################################################
@@ -49,20 +45,19 @@ MAX_GPU_BATCH_SIZE = 16
 EVAL_BATCH_SIZE = 32
 
 
-def training_function(config, args):
-    # Initialize accelerator
-    accelerator = Accelerator(fp16=args.fp16, cpu=args.cpu)
-
-    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
-    lr = config["lr"]
-    num_epochs = int(config["num_epochs"])
-    correct_bias = config["correct_bias"]
-    seed = int(config["seed"])
-    batch_size = int(config["batch_size"])
+def get_dataloaders(accelerator: Accelerator, batch_size: int = 16):
+    """
+    Creates a set of `DataLoader`s for the `glue` dataset,
+    using "bert-base-cased" as the tokenizer.
 
+    Args:
+        accelerator (`Accelerator`):
+            An `Accelerator` object
+        batch_size (`int`, *optional*):
+            The batch size for the train and validation DataLoaders.
+    """
     tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
     datasets = load_dataset("glue", "mrpc")
-    metric = load_metric("glue", "mrpc")
 
     def tokenize_function(examples):
         # max_length=None => use the model max length (it's actually the default)
@@ -70,21 +65,17 @@ def training_function(config, args):
         return outputs
 
     # Apply the method we just defined to all the examples in all the splits of the dataset
-    tokenized_datasets = datasets.map(
-        tokenize_function,
-        batched=True,
-        remove_columns=["idx", "sentence1", "sentence2"],
-    )
+    # starting with the main process first:
+    with accelerator.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.rename_column_("label", "labels")
-
-    # If the batch size is too big we use gradient accumulation
-    gradient_accumulation_steps = 1
-    if batch_size > MAX_GPU_BATCH_SIZE:
-        gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
-        batch_size = MAX_GPU_BATCH_SIZE
+    tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
 
     def collate_fn(examples):
         # On TPU it's best to pad everything to the same length or training will be very slow.
@@ -100,8 +91,28 @@ def training_function(config, args):
         tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
     )
 
-    set_seed(seed)
+    return train_dataloader, eval_dataloader
 
+
+def training_function(config, args):
+    # Initialize accelerator
+    accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)
+    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
+    lr = config["lr"]
+    num_epochs = int(config["num_epochs"])
+    seed = int(config["seed"])
+    batch_size = int(config["batch_size"])
+
+    metric = evaluate.load("glue", "mrpc")
+
+    # If the batch size is too big we use gradient accumulation
+    gradient_accumulation_steps = 1
+    if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
+        gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
+        batch_size = MAX_GPU_BATCH_SIZE
+
+    set_seed(seed)
+    train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size)
     # Instantiate the model (we build the model here so that the seed also control new weights initialization)
     model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", return_dict=True)
 
@@ -111,21 +122,20 @@ def training_function(config, args):
     model = model.to(accelerator.device)
 
     # Instantiate optimizer
-    optimizer = AdamW(params=model.parameters(), lr=lr, correct_bias=correct_bias)
+    optimizer = AdamW(params=model.parameters(), lr=lr)
+
+    # Instantiate scheduler
+    lr_scheduler = get_linear_schedule_with_warmup(
+        optimizer=optimizer,
+        num_warmup_steps=100,
+        num_training_steps=(len(train_dataloader) * num_epochs) // gradient_accumulation_steps,
+    )
 
     # 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
-    )
-
-    # Instantiate learning rate scheduler after preparing the training dataloader as the prepare method
-    # may change its length.
-    lr_scheduler = get_linear_schedule_with_warmup(
-        optimizer=optimizer,
-        num_warmup_steps=100,
-        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
     )
 
     # Now we train the model
@@ -150,9 +160,10 @@ def training_function(config, args):
             with torch.no_grad():
                 outputs = model(**batch)
             predictions = outputs.logits.argmax(dim=-1)
+            predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]))
             metric.add_batch(
-                predictions=accelerator.gather(predictions),
-                references=accelerator.gather(batch["labels"]),
+                predictions=predictions,
+                references=references,
             )
 
         eval_metric = metric.compute()
@@ -162,10 +173,18 @@ def training_function(config, args):
 
 def main():
     parser = argparse.ArgumentParser(description="Simple example of training script.")
-    parser.add_argument("--fp16", type=bool, default=False, help="If passed, will use FP16 training.")
-    parser.add_argument("--cpu", type=bool, default=False, help="If passed, will train on the CPU.")
+    parser.add_argument(
+        "--mixed_precision",
+        type=str,
+        default="no",
+        choices=["no", "fp16", "bf16"],
+        help="Whether to use mixed precision. Choose"
+        "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
+        "and an Nvidia Ampere GPU.",
+    )
+    parser.add_argument("--cpu", action="store_true", help="If passed, will train on the CPU.")
     args = parser.parse_args()
-    config = {"lr": 2e-5, "num_epochs": 3, "correct_bias": True, "seed": 42, "batch_size": 16}
+    config = {"lr": 2e-5, "num_epochs": 3, "seed": 42, "batch_size": 16}
     training_function(config, args)
 
 

examples/requirements.txt

@@ -1 +1,3 @@
-accelerate # used to be installed in Amazon SageMaker environment
+accelerate # used to be installed in Amazon SageMaker environment
+evaluate
+datasets==2.3.2

manim_animations/big_model_inference/stage_1.py

@@ -0,0 +1,108 @@
+# 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.
+
+from manim import *
+
+
+class Stage1(Scene):
+    def construct(self):
+        mem = Rectangle(height=0.5,width=0.5)
+        fill = Rectangle(height=0.46,width=0.46).set_stroke(width=0)
+
+        cpu_left_col_base = [mem.copy() for i in range(6)]
+        cpu_right_col_base = [mem.copy() for i in range(6)]
+        cpu_left_col = VGroup(*cpu_left_col_base).arrange(UP, buff=0)
+        cpu_right_col = VGroup(*cpu_right_col_base).arrange(UP, buff=0)
+        cpu_rects = VGroup(cpu_left_col,cpu_right_col).arrange(RIGHT, buff=0)
+        cpu_text = Text("CPU", font_size=24)
+        cpu = Group(cpu_rects,cpu_text).arrange(DOWN, buff=0.5, aligned_edge=DOWN)
+        cpu.move_to([-2.5,-.5,0])
+        self.add(cpu)
+
+        gpu_base = [mem.copy() for i in range(1)]
+        gpu_rect = VGroup(*gpu_base).arrange(UP,buff=0)
+        gpu_text = Text("GPU", font_size=24)
+        gpu = Group(gpu_rect,gpu_text).arrange(DOWN, buff=0.5, aligned_edge=DOWN)
+        gpu.align_to(cpu, DOWN)
+        gpu.set_x(gpu.get_x() - 1)
+        
+        self.add(gpu)
+
+        model_base = [mem.copy() for i in range(6)]
+        model_rect = VGroup(*model_base).arrange(RIGHT,buff=0)
+
+        model_text = Text("Model", font_size=24)
+        model = Group(model_rect,model_text).arrange(DOWN, buff=0.5, aligned_edge=DOWN)
+        model.move_to([3, -1., 0])
+        
+        self.play(
+            Create(cpu_left_col, run_time=1),
+            Create(cpu_right_col, run_time=1),
+            Create(gpu_rect, run_time=1),
+        )
+
+        step_1 = MarkupText(
+            f"First, an empty model skeleton is loaded\ninto <span fgcolor='{YELLOW}'>memory</span> without using much RAM.", 
+            font_size=24
+        )
+
+        key = Square(side_length=2.2)
+        key.move_to([-5, 2, 0])
+
+        key_text = MarkupText(
+            f"<b>Key:</b>\n\n<span fgcolor='{YELLOW}'>●</span> Empty Model",
+            font_size=18,
+        )
+
+        key_text.move_to([-5, 2.4, 0])
+
+
+        step_1.move_to([2, 2, 0])
+        self.play(
+            Write(step_1, run_time=2.5),
+            Write(key_text),
+            Write(key)
+        )
+
+        self.add(model)
+        
+
+        cpu_targs = []
+        first_animations = []
+        second_animations = []
+        for i,rect in enumerate(model_base):
+
+            cpu_target = Rectangle(height=0.46,width=0.46).set_stroke(width=0.).set_fill(YELLOW, opacity=0.7)
+            cpu_target.move_to(rect)
+            cpu_target.generate_target()
+            cpu_target.target.height = 0.46/4
+            cpu_target.target.width = 0.46/3
+            
+            if i == 0:
+                cpu_target.target.next_to(cpu_left_col_base[0].get_corner(DOWN+LEFT), buff=0.02, direction=UP)
+                cpu_target.target.set_x(cpu_target.target.get_x()+0.1)
+            elif i == 3:
+                cpu_target.target.next_to(cpu_targs[0].target, direction=UP, buff=0.)
+            else:
+                cpu_target.target.next_to(cpu_targs[i-1].target, direction=RIGHT, buff=0.)
+            cpu_targs.append(cpu_target)
+
+            first_animations.append(rect.animate(run_time=0.5).set_stroke(YELLOW))
+            second_animations.append(MoveToTarget(cpu_target, run_time=1.5))
+
+        self.play(*first_animations)
+        self.play(*second_animations)
+                 
+
+        self.wait()

manim_animations/big_model_inference/stage_2.py

@@ -0,0 +1,126 @@
+# 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.
+
+from manim import *
+
+class Stage2(Scene):
+    def construct(self):
+        mem = Rectangle(height=0.5,width=0.5)
+        fill = Rectangle(height=0.46,width=0.46).set_stroke(width=0)
+
+        cpu_left_col_base = [mem.copy() for i in range(6)]
+        cpu_right_col_base = [mem.copy() for i in range(6)]
+        cpu_left_col = VGroup(*cpu_left_col_base).arrange(UP, buff=0)
+        cpu_right_col = VGroup(*cpu_right_col_base).arrange(UP, buff=0)
+        cpu_rects = VGroup(cpu_left_col,cpu_right_col).arrange(RIGHT, buff=0)
+        cpu_text = Text("CPU", font_size=24)
+        cpu = Group(cpu_rects,cpu_text).arrange(DOWN, buff=0.5, aligned_edge=DOWN)
+        cpu.move_to([-2.5,-.5,0])
+        self.add(cpu)
+
+        gpu_base = [mem.copy() for i in range(4)]
+        gpu_rect = VGroup(*gpu_base).arrange(UP,buff=0)
+        gpu_text = Text("GPU", font_size=24)
+        gpu = Group(gpu_rect,gpu_text).arrange(DOWN, buff=0.5, aligned_edge=DOWN)
+        gpu.move_to([-1,-1,0])
+        self.add(gpu)
+
+        model_base = [mem.copy() for i in range(6)]
+        model_rect = VGroup(*model_base).arrange(RIGHT,buff=0)
+
+        model_text = Text("Model", font_size=24)
+        model = Group(model_rect,model_text).arrange(DOWN, buff=0.5, aligned_edge=DOWN)
+        model.move_to([3, -1., 0])
+        self.add(model)
+        
+        cpu_targs = []
+        for i,rect in enumerate(model_base):
+            rect.set_stroke(YELLOW)
+            # target = fill.copy().set_fill(YELLOW, opacity=0.7)
+            # target.move_to(rect)
+            # self.add(target)
+
+            cpu_target = Rectangle(height=0.46/4,width=0.46/3).set_stroke(width=0.).set_fill(YELLOW, opacity=0.7)
+            
+            if i == 0:
+                cpu_target.next_to(cpu_left_col_base[0].get_corner(DOWN+LEFT), buff=0.02, direction=UP)
+                cpu_target.set_x(cpu_target.get_x()+0.1)
+            elif i == 3:
+                cpu_target.next_to(cpu_targs[0], direction=UP, buff=0.)
+            else:
+                cpu_target.next_to(cpu_targs[i-1], direction=RIGHT, buff=0.)
+            self.add(cpu_target)
+            cpu_targs.append(cpu_target)
+
+              
+
+        checkpoint_base = [mem.copy() for i in range(6)]
+        checkpoint_rect = VGroup(*checkpoint_base).arrange(RIGHT,buff=0)
+
+        checkpoint_text = Text("Loaded Checkpoint", font_size=24)
+        checkpoint = Group(checkpoint_rect,checkpoint_text).arrange(DOWN, aligned_edge=DOWN, buff=0.4)
+        checkpoint.move_to([3, .5, 0])
+            
+        key = Square(side_length=2.2)
+        key.move_to([-5, 2, 0])
+
+        key_text = MarkupText(
+            f"<b>Key:</b>\n\n<span fgcolor='{YELLOW}'>●</span> Empty Model",
+            font_size=18,
+        )
+
+        key_text.move_to([-5, 2.4, 0])
+
+        self.add(key_text, key)
+
+        blue_text = MarkupText(
+            f"<span fgcolor='{BLUE}'>●</span> Checkpoint",
+            font_size=18,
+        )
+
+        blue_text.next_to(key_text, DOWN*2.4, aligned_edge=key_text.get_left())
+
+        step_2 = MarkupText(
+            f'Next, a <i><span fgcolor="{BLUE}">second</span></i> model is loaded into memory,\nwith the weights of a <span fgcolor="{BLUE}">single shard</span>.', 
+            font_size=24
+        )
+        step_2.move_to([2, 2, 0])
+        self.play(
+            Write(step_2),
+            Write(blue_text)
+        )
+
+        self.play(
+            Write(checkpoint_text, run_time=1),
+            Create(checkpoint_rect, run_time=1)
+        )
+
+        first_animations = []
+        second_animations = []
+        for i,rect in enumerate(checkpoint_base):
+            target = fill.copy().set_fill(BLUE, opacity=0.7)
+            target.move_to(rect)
+            first_animations.append(GrowFromCenter(target, run_time=1))
+
+            cpu_target = target.copy()
+            cpu_target.generate_target()
+            if i < 5:
+                cpu_target.target.move_to(cpu_left_col_base[i+1])
+            else:
+                cpu_target.target.move_to(cpu_right_col_base[i-5])
+            second_animations.append(MoveToTarget(cpu_target, run_time=1.5))
+            
+        self.play(*first_animations)
+        self.play(*second_animations)
+        self.wait()

manim_animations/big_model_inference/stage_3.py

@@ -0,0 +1,158 @@
+# 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.
+
+from manim import *
+
+class Stage3(Scene):
+    def construct(self):
+        mem = Rectangle(height=0.5,width=0.5)
+        meta_mem = Rectangle(height=0.25,width=0.25)
+        fill = Rectangle(height=0.46,width=0.46).set_stroke(width=0)
+
+        cpu_left_col_base = [mem.copy() for i in range(6)]
+        cpu_right_col_base = [mem.copy() for i in range(6)]
+        cpu_left_col = VGroup(*cpu_left_col_base).arrange(UP, buff=0)
+        cpu_right_col = VGroup(*cpu_right_col_base).arrange(UP, buff=0)
+        cpu_rects = VGroup(cpu_left_col,cpu_right_col).arrange(RIGHT, buff=0)
+        cpu_text = Text("CPU", font_size=24)
+        cpu = Group(cpu_rects,cpu_text).arrange(DOWN, buff=0.5, aligned_edge=DOWN)
+        cpu.move_to([-2.5,-.5,0])
+        self.add(cpu)
+
+        gpu_base = [mem.copy() for i in range(4)]
+        gpu_rect = VGroup(*gpu_base).arrange(UP,buff=0)
+        gpu_text = Text("GPU", font_size=24)
+        gpu = Group(gpu_rect,gpu_text).arrange(DOWN, buff=0.5, aligned_edge=DOWN)
+        gpu.move_to([-1,-1,0])
+        self.add(gpu)
+
+        model_base = [mem.copy() for i in range(6)]
+        model_rect = VGroup(*model_base).arrange(RIGHT,buff=0)
+
+        model_text = Text("Model", font_size=24)
+        model = Group(model_rect,model_text).arrange(DOWN, buff=0.5, aligned_edge=DOWN)
+        model.move_to([3, -1., 0])
+        self.add(model)
+
+        model_arr = []
+        model_cpu_arr = []
+        model_meta_arr = []
+        
+        for i,rect in enumerate(model_base):
+            rect.set_stroke(YELLOW)
+
+            cpu_target = Rectangle(height=0.46/4,width=0.46/3).set_stroke(width=0.).set_fill(YELLOW, opacity=0.7)
+            
+            if i == 0:
+                cpu_target.next_to(cpu_left_col_base[0].get_corner(DOWN+LEFT), buff=0.02, direction=UP)
+                cpu_target.set_x(cpu_target.get_x()+0.1)
+            elif i == 3:
+                cpu_target.next_to(model_cpu_arr[0], direction=UP, buff=0.)
+            else:
+                cpu_target.next_to(model_cpu_arr[i-1], direction=RIGHT, buff=0.)
+            self.add(cpu_target)
+            model_cpu_arr.append(cpu_target)
+
+        self.add(*model_arr, *model_cpu_arr, *model_meta_arr)
+
+        checkpoint_base = [mem.copy() for i in range(6)]
+        checkpoint_rect = VGroup(*checkpoint_base).arrange(RIGHT,buff=0)
+
+        checkpoint_text = Text("Loaded Checkpoint", font_size=24)
+        checkpoint = Group(checkpoint_rect,checkpoint_text).arrange(DOWN, buff=0.5, aligned_edge=DOWN)
+        checkpoint.move_to([3, .5, 0])
+            
+        self.add(checkpoint)
+
+        ckpt_arr = []
+        ckpt_cpu_arr = []
+
+        for i,rect in enumerate(checkpoint_base):
+            target = fill.copy().set_fill(BLUE, opacity=0.7)
+            target.move_to(rect)
+            ckpt_arr.append(target)
+
+            cpu_target = target.copy()
+            if i < 5:
+                cpu_target.move_to(cpu_left_col_base[i+1])
+            else:
+                cpu_target.move_to(cpu_right_col_base[i-5])
+            ckpt_cpu_arr.append(cpu_target)
+        self.add(*ckpt_arr, *ckpt_cpu_arr)
+
+        key = Square(side_length=2.2)
+        key.move_to([-5, 2, 0])
+
+        key_text = MarkupText(
+            f"<b>Key:</b>\n\n<span fgcolor='{YELLOW}'>●</span> Empty Model",
+            font_size=18,
+        )
+
+        key_text.move_to([-5, 2.4, 0])
+
+        self.add(key_text, key)
+
+        blue_text = MarkupText(
+            f"<span fgcolor='{BLUE}'>●</span> Checkpoint",
+            font_size=18,
+        )
+
+        blue_text.next_to(key_text, DOWN*2.4, aligned_edge=key_text.get_left())
+        self.add(blue_text)
+
+        step_3 = MarkupText(
+            f'Based on the passed in configuration, weights are stored in\na variety of np.memmaps on disk or to a particular device.', 
+            font_size=24
+        )
+        step_3.move_to([2, 2, 0])
+
+        disk_left_col_base = [meta_mem.copy() for i in range(6)]
+        disk_right_col_base = [meta_mem.copy() for i in range(6)]
+        disk_left_col = VGroup(*disk_left_col_base).arrange(UP, buff=0)
+        disk_right_col = VGroup(*disk_right_col_base).arrange(UP, buff=0)
+        disk_rects = VGroup(disk_left_col,disk_right_col).arrange(RIGHT, buff=0)
+        disk_text = Text("Disk", font_size=24)
+        disk = Group(disk_rects,disk_text).arrange(DOWN, buff=0.5, aligned_edge=DOWN)
+        disk.move_to([-4.,-1.25,0])
+        self.play(
+            Write(step_3, run_time=3),
+            Write(disk_text, run_time=1),
+            Create(disk_rects, run_time=1)
+        )
+
+        animations = []
+        for i,rect in enumerate(ckpt_cpu_arr):
+            target = rect.copy()
+            target.generate_target()
+            target.target.move_to(disk_left_col_base[i]).scale(0.5)
+            animations.append(MoveToTarget(target, run_time=1.5))
+        self.play(*animations)
+
+        self.play(FadeOut(step_3))
+
+        step_4 = MarkupText(
+            f'Then, the checkpoint is removed from memory\nthrough garbage collection.', 
+            font_size=24
+        )
+        step_4.move_to([2, 2, 0])
+
+        self.play(
+            Write(step_4, run_time=3)
+        )
+
+        self.play(
+            FadeOut(checkpoint_rect, checkpoint_text, *ckpt_arr, *ckpt_cpu_arr),
+        )
+
+        self.wait()      

manim_animations/big_model_inference/stage_4.py

@@ -0,0 +1,156 @@
+# 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.
+
+from manim import *
+
+class Stage4(Scene):
+    def construct(self):
+        mem = Rectangle(height=0.5,width=0.5)
+        fill = Rectangle(height=0.46,width=0.46).set_stroke(width=0)
+        meta_mem = Rectangle(height=0.25,width=0.25)
+
+        cpu_left_col_base = [mem.copy() for i in range(6)]
+        cpu_right_col_base = [mem.copy() for i in range(6)]
+        cpu_left_col = VGroup(*cpu_left_col_base).arrange(UP, buff=0)
+        cpu_right_col = VGroup(*cpu_right_col_base).arrange(UP, buff=0)
+        cpu_rects = VGroup(cpu_left_col,cpu_right_col).arrange(RIGHT, buff=0)
+        cpu_text = Text("CPU", font_size=24)
+        cpu = Group(cpu_rects,cpu_text).arrange(DOWN, buff=0.5, aligned_edge=DOWN)
+        cpu.move_to([-2.5,-.5,0])
+        self.add(cpu)
+
+        gpu_base = [mem.copy() for i in range(4)]
+        gpu_rect = VGroup(*gpu_base).arrange(UP,buff=0)
+        gpu_text = Text("GPU", font_size=24)
+        gpu = Group(gpu_rect,gpu_text).arrange(DOWN, buff=0.5, aligned_edge=DOWN)
+        gpu.move_to([-1,-1,0])
+        self.add(gpu)
+
+        model_base = [mem.copy() for i in range(6)]
+        model_rect = VGroup(*model_base).arrange(RIGHT,buff=0)
+
+        model_text = Text("Model", font_size=24)
+        model = Group(model_rect,model_text).arrange(DOWN, buff=0.5, aligned_edge=DOWN)
+        model.move_to([3, -1., 0])
+        self.add(model)
+
+        model_cpu_arr = []
+        model_meta_arr = []
+        
+        for i,rect in enumerate(model_base):
+            rect.set_stroke(YELLOW)
+
+            cpu_target = Rectangle(height=0.46/4,width=0.46/3).set_stroke(width=0.).set_fill(YELLOW, opacity=0.7)
+            
+            if i == 0:
+                cpu_target.next_to(cpu_left_col_base[0].get_corner(DOWN+LEFT), buff=0.02, direction=UP)
+                cpu_target.set_x(cpu_target.get_x()+0.1)
+            elif i == 3:
+                cpu_target.next_to(model_cpu_arr[0], direction=UP, buff=0.)
+            else:
+                cpu_target.next_to(model_cpu_arr[i-1], direction=RIGHT, buff=0.)
+            self.add(cpu_target)
+            model_cpu_arr.append(cpu_target)
+
+        self.add(*model_cpu_arr, *model_meta_arr)
+
+        disk_left_col_base = [meta_mem.copy() for i in range(6)]
+        disk_right_col_base = [meta_mem.copy() for i in range(6)]
+        disk_left_col = VGroup(*disk_left_col_base).arrange(UP, buff=0)
+        disk_right_col = VGroup(*disk_right_col_base).arrange(UP, buff=0)
+        disk_rects = VGroup(disk_left_col,disk_right_col).arrange(RIGHT, buff=0)
+        disk_text = Text("Disk", font_size=24)
+        disk = Group(disk_rects,disk_text).arrange(DOWN, buff=0.5, aligned_edge=DOWN)
+        disk.move_to([-4.,-1.25,0])
+        self.add(disk_text, disk_rects)
+
+        cpu_disk_arr = []
+
+        for i in range(6):
+            target = fill.copy().set_fill(BLUE, opacity=0.8)
+            target.move_to(disk_left_col_base[i]).scale(0.5)
+            cpu_disk_arr.append(target)
+
+        self.add(*cpu_disk_arr)
+
+        key = Square(side_length=2.2)
+        key.move_to([-5, 2, 0])
+
+        key_text = MarkupText(
+            f"<b>Key:</b>\n\n<span fgcolor='{YELLOW}'>●</span> Empty Model",
+            font_size=18,
+        )
+
+        key_text.move_to([-5, 2.4, 0])
+
+        self.add(key_text, key)
+
+        blue_text = MarkupText(
+            f"<span fgcolor='{BLUE}'>●</span> Checkpoint",
+            font_size=18,
+        )
+
+        blue_text.next_to(key_text, DOWN*2.4, aligned_edge=key_text.get_left())
+        self.add(blue_text)
+
+        step_5 = MarkupText(
+            f'The offloaded weights are all sent to the CPU.', 
+            font_size=24
+        )
+        step_5.move_to([2, 2, 0])
+
+        self.play(Write(step_5, run_time=3))
+
+        for i in range(6):
+            rect = cpu_disk_arr[i]
+            cp2 = rect.copy().set_fill(BLUE, opacity=0.8).scale(2.0)
+            cp2.generate_target()
+            cp2.target.move_to(model_base[i])
+
+            if i == 0:
+                rect.set_fill(BLUE, opacity=0.8)
+                rect.generate_target()
+                rect.target.move_to(cpu_left_col_base[0]).scale(2.0)
+                
+                self.remove(*model_meta_arr, 
+                    *model_cpu_arr,
+                )
+
+            else:
+                rect.generate_target()
+                rect.target.move_to(cpu_left_col_base[i]).scale(2.0)
+            self.play(
+                MoveToTarget(rect),
+                MoveToTarget(cp2),
+                model_base[i].animate.set_stroke(WHITE)
+            )
+        self.play(FadeOut(step_5))
+
+        step_5 = MarkupText(
+            f'Finally, hooks are added to each weight in the model\nto transfer the weights from CPU to GPU\n\t\tand back when needed.', 
+            font_size=24
+        )
+        step_5.move_to([2, 2, 0])
+
+        self.play(Write(step_5, run_time=3))
+
+        arrows = []
+        animations = []
+        for i in range(6):
+            a = Arrow(start=UP, end=DOWN, color=RED, buff=.5)
+            a.next_to(model_base[i].get_left(), UP, buff=0.2)
+            arrows.append(a)
+            animations.append(Write(a))
+        self.play(*animations)
+        self.wait()  

manim_animations/big_model_inference/stage_5.py

@@ -0,0 +1,221 @@
+# 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.
+
+from manim import *
+
+class Stage5(Scene):
+    def construct(self):
+        mem = Rectangle(height=0.5,width=0.5)
+        fill = Rectangle(height=0.46,width=0.46).set_stroke(width=0)
+
+        meta_mem = Rectangle(height=0.25,width=0.25)
+
+        cpu_left_col_base = [mem.copy() for i in range(6)]
+        cpu_right_col_base = [mem.copy() for i in range(6)]
+        cpu_left_col = VGroup(*cpu_left_col_base).arrange(UP, buff=0)
+        cpu_right_col = VGroup(*cpu_right_col_base).arrange(UP, buff=0)
+        cpu_rects = VGroup(cpu_left_col,cpu_right_col).arrange(RIGHT, buff=0)
+        cpu_text = Text("CPU", font_size=24)
+        cpu = Group(cpu_rects,cpu_text).arrange(DOWN, buff=0.5, aligned_edge=DOWN)
+        cpu.move_to([-2.5,-.5,0])
+        self.add(cpu)
+
+        gpu_base = [mem.copy() for i in range(4)]
+        gpu_rect = VGroup(*gpu_base).arrange(UP,buff=0)
+        gpu_text = Text("GPU", font_size=24)
+        gpu = Group(gpu_rect,gpu_text).arrange(DOWN, buff=0.5, aligned_edge=DOWN)
+        gpu.move_to([-1,-1,0])
+        self.add(gpu)
+
+        model_base = [mem.copy() for i in range(6)]
+        model_rect = VGroup(*model_base).arrange(RIGHT,buff=0)
+
+        model_text = Text("Model", font_size=24)
+        model = Group(model_rect,model_text).arrange(DOWN, buff=0.5, aligned_edge=DOWN)
+        model.move_to([3, -1., 0])
+        self.add(model)
+
+        model_arr = []
+        model_cpu_arr = []
+        
+        for i,rect in enumerate(model_base):
+            target = fill.copy().set_fill(BLUE, opacity=0.8)
+            target.move_to(rect)
+            model_arr.append(target)
+
+            cpu_target = Rectangle(height=0.46,width=0.46).set_stroke(width=0.).set_fill(BLUE, opacity=0.8)
+            cpu_target.move_to(cpu_left_col_base[i])
+            model_cpu_arr.append(cpu_target)
+
+        self.add(*model_arr, *model_cpu_arr)
+
+        disk_left_col_base = [meta_mem.copy() for i in range(6)]
+        disk_right_col_base = [meta_mem.copy() for i in range(6)]
+        disk_left_col = VGroup(*disk_left_col_base).arrange(UP, buff=0)
+        disk_right_col = VGroup(*disk_right_col_base).arrange(UP, buff=0)
+        disk_rects = VGroup(disk_left_col,disk_right_col).arrange(RIGHT, buff=0)
+        disk_text = Text("Disk", font_size=24)
+        disk = Group(disk_rects,disk_text).arrange(DOWN, buff=0.5, aligned_edge=DOWN)
+        disk.move_to([-4,-1.25,0])
+        self.add(disk_text, disk_rects)
+
+        key = Square(side_length=2.2)
+        key.move_to([-5, 2, 0])
+
+        key_text = MarkupText(
+            f"<b>Key:</b>\n\n<span fgcolor='{YELLOW}'>●</span> Empty Model",
+            font_size=18,
+        )
+
+        key_text.move_to([-5, 2.4, 0])
+
+        self.add(key_text, key)
+
+        blue_text = MarkupText(
+            f"<span fgcolor='{BLUE}'>●</span> Checkpoint",
+            font_size=18,
+        )
+
+        blue_text.next_to(key_text, DOWN*2.4, aligned_edge=key_text.get_left())
+        self.add(blue_text)
+
+        step_6 = MarkupText(
+            f'Now watch as an input is passed through the model\nand how the memory is utilized and handled.', 
+            font_size=24
+        )
+        step_6.move_to([2, 2, 0])
+
+        self.play(Write(step_6))
+
+        input = Square(0.3)
+        input.set_fill(RED, opacity=1.)
+        input.set_stroke(width=0.)
+        input.next_to(model_base[0], LEFT, buff=.5)
+
+        self.play(Write(input))
+
+        input.generate_target()
+        input.target.next_to(model_arr[0], direction=LEFT, buff=0.02)
+        self.play(MoveToTarget(input))
+
+        self.play(FadeOut(step_6))
+
+
+        a = Arrow(start=UP, end=DOWN, color=RED, buff=.5)
+        a.next_to(model_arr[0].get_left(), UP, buff=0.2)
+
+        model_cpu_arr[0].generate_target()
+        model_cpu_arr[0].target.move_to(gpu_rect[0])
+
+        step_7 = MarkupText(
+            f'As the input reaches a layer, the hook triggers\nand weights are moved from the CPU\nto the GPU and back.', 
+            font_size=24
+        )
+        step_7.move_to([2, 2, 0])
+
+        self.play(Write(step_7, run_time=3))
+
+        circ_kwargs = {"run_time":1, "fade_in":True, "fade_out":True, "buff":0.02}
+
+        self.play(
+            Write(a), 
+            Circumscribe(model_arr[0], color=ORANGE, **circ_kwargs),
+            Circumscribe(model_cpu_arr[0], color=ORANGE, **circ_kwargs),
+            Circumscribe(gpu_rect[0], color=ORANGE, **circ_kwargs),
+        )
+        self.play(
+            MoveToTarget(model_cpu_arr[0])
+        )
+
+        a_c = a.copy()
+        for i in range(6):
+            a_c.next_to(model_arr[i].get_right()+0.02, UP, buff=0.2)
+
+            input.generate_target()
+            input.target.move_to(model_arr[i].get_right()+0.02)
+
+            grp = AnimationGroup(
+                FadeOut(a, run_time=.5), 
+                MoveToTarget(input, run_time=.5), 
+                FadeIn(a_c, run_time=.5),
+                lag_ratio=0.2
+            )
+
+            self.play(grp)
+
+
+            model_cpu_arr[i].generate_target()
+            model_cpu_arr[i].target.move_to(cpu_left_col_base[i])
+
+
+            if i < 5:
+                model_cpu_arr[i+1].generate_target()
+                model_cpu_arr[i+1].target.move_to(gpu_rect[0])
+                if i >= 1:
+                    circ_kwargs["run_time"] = .7
+
+                self.play(
+                    Circumscribe(model_arr[i], **circ_kwargs),
+                    Circumscribe(cpu_left_col_base[i], **circ_kwargs),
+                    Circumscribe(cpu_left_col_base[i+1], color=ORANGE, **circ_kwargs),                    
+                    Circumscribe(gpu_rect[0], color=ORANGE, **circ_kwargs),
+                    Circumscribe(model_arr[i+1], color=ORANGE, **circ_kwargs),
+                )
+                if i < 1:
+                    self.play(
+                        MoveToTarget(model_cpu_arr[i]), 
+                        MoveToTarget(model_cpu_arr[i+1]),
+                    )
+                else:
+                    self.play(
+                        MoveToTarget(model_cpu_arr[i], run_time=.7), 
+                        MoveToTarget(model_cpu_arr[i+1], run_time=.7),
+                    )
+            else:
+                model_cpu_arr[i].generate_target()
+                model_cpu_arr[i].target.move_to(cpu_left_col_base[-1])
+                input.generate_target()
+                input.target.next_to(model_arr[-1].get_right(), RIGHT+0.02, buff=0.2)
+
+                self.play(
+                    Circumscribe(model_arr[-1], color=ORANGE, **circ_kwargs),
+                    Circumscribe(cpu_left_col_base[-1], color=ORANGE, **circ_kwargs),
+                    Circumscribe(gpu_rect[0], color=ORANGE, **circ_kwargs),
+                )
+
+                self.play(
+                    MoveToTarget(model_cpu_arr[i])
+                )
+
+            a = a_c
+            a_c = a_c.copy()
+
+        input.generate_target()
+        input.target.next_to(model_base[-1], RIGHT+0.02, buff=.5)
+        self.play(
+            FadeOut(step_7),
+            FadeOut(a, run_time=.5), 
+        )
+
+        step_8 = MarkupText(
+            f'Inference on a model too large for GPU memory\nis successfully completed.', font_size=24
+        )
+        step_8.move_to([2, 2, 0])
+
+        self.play(
+            Write(step_8, run_time=3),
+            MoveToTarget(input)
+        )
+
+        self.wait()