Release: v0.23.0

* Remove checkpoints only on main process

shutil.rmtree might throw errors if called on multiply processes. Make a call only on main process

* Apply style

.devcontainer/devcontainer.json

@@ -15,10 +15,14 @@
     "remoteEnv": {
         "PYTHONPATH": "${containerEnv:PATH}:${containerWorkspaceFolder}"
     },
-    "extensions": [
-        // Ensure we have IntelliSense in VSCode when running inside container
-        "ms-python.python"
-    ],
+    "customizations": {
+        "vscode": {
+            "extensions": [
+                // Ensure we have IntelliSense in VSCode when running inside container
+                "ms-python.python"
+            ]
+        }
+    },
     "workspaceFolder": "/workspaces/accelerate",
     // Need git for VSCode to color code modifications. Only runs when building environment.
     "onCreateCommand": "apt-get update && apt-get install -y git && pip install -e '.[dev]'"

.github/ISSUE_TEMPLATE/bug-report.yml

@@ -1,6 +1,12 @@
 name: "\U0001F41B Bug Report"
 description: Submit a bug report to help us improve Accelerate
 body:
+  - type: markdown
+    attributes: 
+      value: | 
+        Thanks for taking the time to submit a bug report! 🐛 
+        If this is not a bug related to the Accelerate library directly, but instead a general question about your code or the library specifically please use the [forums](https://discuss.huggingface.co/c/accelerate/18).
+
   - type: textarea
     id: system-info
     attributes:
@@ -55,4 +61,3 @@ body:
     attributes:
       label: Expected behavior
       description: "A clear and concise description of what you would expect to happen."
-      render: Shell

.github/PULL_REQUEST_TEMPLATE.md

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

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

@@ -21,44 +21,40 @@ jobs:
 
   version-cpu:
     name: "Latest Accelerate CPU [version]"
-    runs-on: ubuntu-latest
+    runs-on: [self-hosted, docker-gpu, multi-gpu]
     needs: get-version
     steps:
       - name: Set up Docker Buildx
-        uses: docker/setup-buildx-action@v1
-      - name: Check out code
-        uses: actions/checkout@v2
+        uses: docker/setup-buildx-action@v2
       - name: Login to DockerHub
-        uses: docker/login-action@v1
+        uses: docker/login-action@v2
         with:
           username: ${{ secrets.DOCKERHUB_USERNAME }}
           password: ${{ secrets.DOCKERHUB_PASSWORD }}
 
       - name: Build and Push CPU
-        uses: docker/build-push-action@v2
+        uses: docker/build-push-action@v4
         with:
-          context: ./docker/accelerate-cpu
+          file: docker/accelerate-cpu/Dockerfile
           push: true
           tags: huggingface/accelerate-cpu:${{needs.get-version.outputs.version}}
 
   version-cuda:
     name: "Latest Accelerate GPU [version]"
-    runs-on: ubuntu-latest
+    runs-on: [self-hosted, docker-gpu, multi-gpu]
     needs: get-version
     steps:
       - name: Set up Docker Buildx
-        uses: docker/setup-buildx-action@v1
-      - name: Check out code
-        uses: actions/checkout@v2
+        uses: docker/setup-buildx-action@v2
       - name: Login to DockerHub
-        uses: docker/login-action@v1
+        uses: docker/login-action@v2
         with:
           username: ${{ secrets.DOCKERHUB_USERNAME }}
           password: ${{ secrets.DOCKERHUB_PASSWORD }}
 
       - name: Build and Push GPU
-        uses: docker/build-push-action@v2
+        uses: docker/build-push-action@v4
         with:
-          context: ./docker/accelerate-gpu
+          file: docker/accelerate-gpu/Dockerfile
           push: true
-          tags: huggingface/accelerate-gpu:${{needs.get-version.outputs.version}}
+          tags: huggingface/accelerate-gpu:${{needs.get-version.outputs.version}}

.github/workflows/build_and_run_tests.yml

@@ -42,4 +42,9 @@ jobs:
   run-merge-tests:
     needs: build-docker-containers
     if: always()
-    uses: ./.github/workflows/run_merge_tests.yml
+    uses: ./.github/workflows/run_merge_tests.yml
+
+  run-integration-tests:
+    needs: run-merge-tests
+    if: always()
+    uses: ./.github/workflows/self_hosted_integration_tests.yml

.github/workflows/build_docker_images.yml

@@ -11,44 +11,50 @@ concurrency:
   cancel-in-progress: false
 
 jobs:
+  clean-storage:
+    name: "Clean docker image storage"
+    runs-on: [self-hosted, docker-gpu, multi-gpu]
+    steps:
+      - name: Clean storage
+        run: |
+          docker image prune --all -f --filter "until=48h"
+          docker system prune --all -f --filter "until=48h"
+
   latest-cpu:
     name: "Latest Accelerate CPU [dev]"
-    runs-on: ubuntu-latest
+    runs-on: [self-hosted, docker-gpu, multi-gpu]
+    needs: clean-storage
     steps:
       - name: Set up Docker Buildx
-        uses: docker/setup-buildx-action@v1
-      - name: Check out code
-        uses: actions/checkout@v2
+        uses: docker/setup-buildx-action@v2
       - name: Login to DockerHub
-        uses: docker/login-action@v1
+        uses: docker/login-action@v2
         with:
           username: ${{ secrets.DOCKERHUB_USERNAME }}
           password: ${{ secrets.DOCKERHUB_PASSWORD }}
-
       - name: Build and Push CPU
-        uses: docker/build-push-action@v2
+        uses: docker/build-push-action@v4
         with:
-          context: ./docker/accelerate-cpu
+          file: docker/accelerate-cpu/Dockerfile          
           push: true
           tags: huggingface/accelerate-cpu
 
   latest-cuda:
     name: "Latest Accelerate GPU [dev]"
-    runs-on: ubuntu-latest
+    runs-on: [self-hosted, docker-gpu, multi-gpu]
+    needs: clean-storage
     steps:
       - name: Set up Docker Buildx
-        uses: docker/setup-buildx-action@v1
-      - name: Check out code
-        uses: actions/checkout@v2
+        uses: docker/setup-buildx-action@v2
       - name: Login to DockerHub
-        uses: docker/login-action@v1
+        uses: docker/login-action@v2
         with:
           username: ${{ secrets.DOCKERHUB_USERNAME }}
           password: ${{ secrets.DOCKERHUB_PASSWORD }}
 
       - name: Build and Push GPU
-        uses: docker/build-push-action@v2
+        uses: docker/build-push-action@v4
         with:
-          context: ./docker/accelerate-gpu
+          file: docker/accelerate-gpu/Dockerfile          
           push: true
-          tags: huggingface/accelerate-gpu
+          tags: huggingface/accelerate-gpu

.github/workflows/build_documentation.yml

@@ -15,3 +15,4 @@ jobs:
       package: accelerate
     secrets:
       token: ${{ secrets.HUGGINGFACE_PUSH }}
+      hf_token: ${{ secrets.HF_DOC_BUILD_PUSH }}

.github/workflows/delete_doc_comment.yml

@@ -1,13 +1,14 @@
-name: Delete dev documentation
+name: Delete doc comment
 
 on:
-  pull_request:
-    types: [ closed ]
+  workflow_run:
+    workflows: ["Delete doc comment trigger"]
+    types:
+      - completed
 
 
 jobs:
   delete:
     uses: huggingface/doc-builder/.github/workflows/delete_doc_comment.yml@main
-    with:
-      pr_number: ${{ github.event.number }}
-      package: accelerate
+    secrets:
+      comment_bot_token: ${{ secrets.COMMENT_BOT_TOKEN }}

.github/workflows/delete_doc_comment_trigger.yml

@@ -0,0 +1,12 @@
+name: Delete doc comment trigger
+
+on:
+  pull_request:
+    types: [ closed ]
+
+
+jobs:
+  delete:
+    uses: huggingface/doc-builder/.github/workflows/delete_doc_comment_trigger.yml@main
+    with:
+      pr_number: ${{ github.event.number }}

.github/workflows/integration_tests.yml

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

.github/workflows/nightly.yml

@@ -8,12 +8,15 @@ on:
 env:
   RUN_SLOW: "yes"
   IS_GITHUB_CI: "1"
+  SLACK_API_TOKEN: ${{ secrets.SLACK_API_TOKEN }}
+
 
 jobs:
   run_all_tests_single_gpu:
     runs-on: [self-hosted, docker-gpu, multi-gpu]
     env:
       CUDA_VISIBLE_DEVICES: "0"
+      TEST_TYPE: "single_gpu"
     container:
       image: huggingface/accelerate-gpu:latest
       options: --gpus all --shm-size "16gb"
@@ -28,13 +31,15 @@ jobs:
           git config --global --add safe.directory '*'
           git fetch && git checkout ${{ github.sha }} 
           pip install -e . --no-deps
-          pip install pytest-reportlog
+          pip install pytest-reportlog tabulate
 
       - name: Run test on GPUs
         run: |
           source activate accelerate
           make test
+          
       - name: Run examples on GPUs
+        if: always()
         run: |
           source activate accelerate
           pip uninstall comet_ml -y
@@ -43,12 +48,14 @@ jobs:
       - name: Generate Report
         if: always()
         run: |
+          pip install slack_sdk tabulate
           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"
+      TEST_TYPE: "multi_gpu"
     container:
       image: huggingface/accelerate-gpu:latest
       options: --gpus all --shm-size "16gb"
@@ -63,20 +70,23 @@ jobs:
           git config --global --add safe.directory '*'
           git fetch && git checkout ${{ github.sha }}
           pip install -e . --no-deps
-          pip install pytest-reportlog
+          pip install pytest-reportlog tabulate
 
       - name: Run core and big modeling tests on GPUs
         run: |
           source activate accelerate
-          make test_big_modeling
           make test_core
+          make test_big_modeling
+          make test_cli
 
       - name: Run Integration tests on GPUs
+        if: always()
         run: |
           source activate accelerate
           make test_integrations
 
       - name: Run examples on GPUs
+        if: always()
         run: |
           source activate accelerate
           pip uninstall comet_ml -y
@@ -85,4 +95,11 @@ jobs:
       - name: Generate Report
         if: always()
         run: |
-          python utils/log_reports.py >> $GITHUB_STEP_SUMMARY
+          pip install slack_sdk tabulate
+          python utils/log_reports.py >> $GITHUB_STEP_SUMMARY
+
+  
+  run-integration-tests:
+    needs: [run_all_tests_single_gpu, run_all_tests_multi_gpu]
+    if: always()
+    uses: ./.github/workflows/self_hosted_integration_tests.yml

.github/workflows/quality.yml

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

.github/workflows/run_merge_tests.yml

@@ -26,8 +26,8 @@ jobs:
           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
+          pip install -e .[testing,test_trackers] -U
+          pip install pytest-reportlog tabulate
 
       - name: Run CLI tests
         run: |
@@ -35,10 +35,12 @@ jobs:
           make test_cli
           
       - name: Run test on GPUs
+        if: always()
         run: |
           source activate accelerate
           make test
       - name: Run examples on GPUs
+        if: always()
         run: |
           source activate accelerate
           pip uninstall comet_ml -y
@@ -47,6 +49,7 @@ jobs:
       - name: Generate Report
         if: always()
         run: |
+          pip install tabulate
           python utils/log_reports.py >> $GITHUB_STEP_SUMMARY
 
   run_all_tests_multi_gpu:
@@ -64,13 +67,8 @@ jobs:
           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 CLI tests
-        run: |
-          source activate accelerate
-          make test_cli
+          pip install -e .[testing,test_trackers] -U
+          pip install pytest-reportlog tabulate
 
       - name: Run test on GPUs
         run: |
@@ -78,6 +76,7 @@ jobs:
           make test
 
       - name: Run examples on GPUs
+        if: always()
         run: |
           source activate accelerate
           pip uninstall comet_ml -y
@@ -86,4 +85,5 @@ jobs:
       - name: Generate Report
         if: always()
         run: |
+          pip install tabulate
           python utils/log_reports.py >> $GITHUB_STEP_SUMMARY

.github/workflows/self_hosted_integration_tests.yml

@@ -0,0 +1,126 @@
+# CI for specifically ensuring integrations work fine (`transformers` mainly) on GPUs
+# Useful tips:
+#  - `working-directory` should be set to the root of the repo, which is cloned on the actual CI runner.
+#    It follows the directory structure of `actions-runner/_work/{repo_name}/{repo_name}/{cloned_repo} on 
+#    prem, but in Actions setting `working-directory` looks just in the `{repo_name}` level.
+#  - New integrations to test should have its own job, and follow a strategy method where we check both
+#    the pypi and github versions.
+#  - Workflow call lets this be called from `build_and_run_tests.yml`
+#  - When using a docker container, it's recommended to set `--shm-size`, we use 16gb.
+name: Integration Tests (push to "main")
+
+on:
+  workflow_call:
+  workflow_dispatch:
+
+env:
+  HF_HOME: ~/hf_cache
+
+defaults:
+  run:
+    shell: bash
+
+jobs:
+  run-trainer-tests:
+    container:
+      image: huggingface/accelerate-gpu:latest
+      options: --gpus all --shm-size "16gb"
+    runs-on: [self-hosted, docker-gpu, multi-gpu]
+    strategy:
+      fail-fast: false
+      matrix:
+        transformers-version: [
+          pypi,
+          github
+        ]
+        cuda_visible_devices: [
+          "0", 
+          "0,1"
+        ]
+    steps:
+      - name: Update accelerate clone and pip install
+        working-directory: accelerate/
+        run: 
+          source activate accelerate;
+          git config --global --add safe.directory '*';
+          git checkout main && git fetch && git checkout ${{ github.sha }};
+          pip install -e .;
+
+      - name: Update transformers clone & pip install
+        working-directory: transformers/
+        run: |
+          source activate accelerate
+          git config --global --add safe.directory '*'
+          git checkout main && git pull
+          if [[ ${{ matrix.transformers-version }} = pypi ]]; then 
+            git checkout $(git describe --tags `git rev-list --tags --max-count=1`)
+          fi
+          pip install .[torch,deepspeed-testing]
+      
+      - name: Show installed libraries
+        run: |
+          source activate accelerate;
+          pip freeze
+
+      - name: Run trainer tests
+        working-directory: transformers/
+        env:
+          CUDA_VISIBLE_DEVICES: ${{ matrix.cuda_visible_devices }}
+          WANDB_DISABLED: true
+        run: |
+          source activate accelerate;
+          pytest -sv tests/trainer
+
+      - name: Run deepspeed tests
+        working-directory: transformers/
+        env:
+          CUDA_VISIBLE_DEVICES: ${{ matrix.cuda_visible_devices }}
+          WANDB_DISABLED: true
+        if: always()
+        run: |
+          source activate accelerate;
+          pytest -sv tests/deepspeed
+
+  run-skorch-tests:
+    container:
+      image: huggingface/accelerate-gpu:latest
+      options: --gpus all --shm-size "16gb"
+    runs-on: [self-hosted, docker-gpu, multi-gpu]
+    strategy:
+      fail-fast: false
+      matrix:
+        skorch-version: [
+          pypi,
+          github
+        ]
+    steps:
+      - name: Update accelerate clone and pip install
+        working-directory: accelerate/
+        run: 
+          source activate accelerate;
+          git config --global --add safe.directory '*';
+          git checkout main && git fetch && git checkout ${{ github.sha }};
+          pip install -e .;
+
+      - name: Update skorch clone & pip install
+        working-directory: skorch/
+        run: |
+          source activate accelerate
+          git config --global --add safe.directory '*'
+          git checkout master && git pull
+          if [[ ${{ matrix.skorch-version }} = pypi ]]; then 
+            git checkout $(git describe --tags `git rev-list --tags --max-count=1`)
+          fi
+          pip install .[testing]
+          pip install flaky
+
+      - name: Show installed libraries
+        run: |
+          source activate accelerate;
+          pip freeze
+
+      - name: Run skorch tests
+        working-directory: skorch/
+        run: |
+          source activate accelerate;
+          pytest -sv -k TestAccelerate

.github/workflows/stale.yml

@@ -18,7 +18,7 @@ jobs:
     - name: Setup Python
       uses: actions/setup-python@v1
       with:
-        python-version: 3.7
+        python-version: 3.8
     
     - name: Install requirements
       run: |

.github/workflows/test.yml

@@ -23,7 +23,7 @@ jobs:
       matrix:
         pytorch-version: [
           latest,
-          minimum
+          minimum,
         ]
         test-kind: [
           test_prod,
@@ -39,10 +39,10 @@ jobs:
         ]
     steps:
     - uses: actions/checkout@v3.1.0
-    - name: Set up python 3.7
+    - name: Set up python 3.8
       uses: actions/setup-python@v3
       with:
-        python-version: 3.7
+        python-version: 3.8
     
     - name: Activate python cache
       uses: actions/cache@v3
@@ -58,8 +58,8 @@ jobs:
         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
+        if [[ ${{ matrix.test-kind }} = minimum ]]; then pip install torch==1.10.0; fi
+        pip install pytest-reportlog tabulate
     
     - name: Run Tests
       env: 

.github/workflows/upload_pr_documentation.yml

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

.gitignore

@@ -138,4 +138,7 @@ dmypy.json
 .DS_Store
 
 # More test things
-wandb
+wandb
+
+# ruff
+.ruff_cache

CONTRIBUTING.md

@@ -152,7 +152,7 @@ Follow these steps to start contributing:
    $ make test
    ```
 
-   `accelerate` relies on `black` and `isort` to format its source code
+   `accelerate` relies on `black` and `ruff` to format its source code
    consistently. After you make changes, apply automatic style corrections and code verifications
    that can't be automated in one go with:
 
@@ -165,7 +165,7 @@ Follow these steps to start contributing:
    $ make style
    ```
 
-   `accelerate` also uses `flake8` and a few custom scripts to check for coding mistakes. Quality
+   `accelerate` also uses a few custom scripts to check for coding mistakes. Quality
    control runs in CI, however you can also run the same checks with:
 
    ```bash

Makefile

@@ -1,6 +1,6 @@
-.PHONY: quality style test docs
+.PHONY: quality style test docs utils
 
-check_dirs := tests src examples benchmarks
+check_dirs := tests src examples benchmarks utils
 
 # Check that source code meets quality standards
 
@@ -12,15 +12,14 @@ extra_quality_checks:
 
 # this target runs checks on all files
 quality:
-	black --check $(check_dirs)
-	isort --check-only $(check_dirs)
-	flake8 $(check_dirs)
+	black --required-version 23 --check $(check_dirs)
+	ruff $(check_dirs)
 	doc-builder style src/accelerate docs/source --max_len 119 --check_only
 
 # Format source code automatically and check is there are any problems left that need manual fixing
 style:
-	black $(check_dirs)
-	isort $(check_dirs)
+	black --required-version 23 $(check_dirs)
+	ruff $(check_dirs) --fix
 	doc-builder style src/accelerate docs/source --max_len 119
 	
 # Run tests for the library
@@ -28,7 +27,7 @@ test:
 	python -m pytest -s -v ./tests/ --ignore=./tests/test_examples.py $(if $(IS_GITHUB_CI),--report-log "$(PYTORCH_VERSION)_all.log",)
 
 test_big_modeling:
-	python -m pytest -s -v ./tests/test_big_modeling.py $(if $(IS_GITHUB_CI),--report-log "$(PYTORCH_VERSION)_big_modeling.log",)
+	python -m pytest -s -v ./tests/test_big_modeling.py ./tests/test_modeling_utils.py $(if $(IS_GITHUB_CI),--report-log "$(PYTORCH_VERSION)_big_modeling.log",)
 
 test_core:
 	python -m pytest -s -v ./tests/ --ignore=./tests/test_examples.py --ignore=./tests/deepspeed --ignore=./tests/test_big_modeling.py \

README.md

@@ -16,12 +16,12 @@ limitations under the License.
 
 <p align="center">
     <br>
-    <img src="docs/source/imgs/accelerate_logo.png" width="400"/>
+    <img src="https://raw.githubusercontent.com/huggingface/accelerate/main/docs/source/imgs/accelerate_logo.png" width="400"/>
     <br>
 <p>
 
 <p align="center">
-    <!-- Uncomment when CircleCI is setup
+    <!-- Uncomment when CircleCI is set up
     <a href="https://circleci.com/gh/huggingface/accelerate">
         <img alt="Build" src="https://img.shields.io/circleci/build/github/huggingface/transformers/master">
     </a>
@@ -91,7 +91,7 @@ Here is an example:
           optimizer.step()
 ```
 
-As you can see in this example, by adding 5-lines to any standard PyTorch training script you can now run on any kind of single or distributed node setting (single CPU, single GPU, multi-GPUs and TPUs) as well as with or without mixed precision (fp16).
+As you can see in this example, by adding 5-lines to any standard PyTorch training script you can now run on any kind of single or distributed node setting (single CPU, single GPU, multi-GPUs and TPUs) as well as with or without mixed precision (fp8, fp16, bf16).
 
 In particular, the same code can then be run without modification on your local machine for debugging or your training environment.
 
@@ -132,11 +132,11 @@ In particular, the same code can then be run without modification on your local
           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).
+Want to learn more? Check out the [documentation](https://huggingface.co/docs/accelerate) or have a look at our [examples](https://github.com/huggingface/accelerate/tree/main/examples).
 
 ## Launching script
 
-🤗 Accelerate also provides an optional CLI tool that allows you to quickly configure and test your training environment before launching the scripts. No need to remember how to use `torch.distributed.launch` or to write a specific launcher for TPU training!
+🤗 Accelerate also provides an optional CLI tool that allows you to quickly configure and test your training environment before launching the scripts. No need to remember how to use `torch.distributed.run` or to write a specific launcher for TPU training!
 On your machine(s) just run:
 
 ```bash
@@ -155,7 +155,17 @@ 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.
+This CLI tool is **optional**, and you can still use `python my_script.py` or `python -m torchrun my_script.py` at your convenience.
+
+You can also directly pass in the arguments you would to `torchrun` as arguments to `accelerate launch` if you wish to not run` accelerate config`.
+
+For example, here is how to launch on two GPUs:
+
+```bash
+accelerate launch --multi_gpu --num_processes 2 examples/nlp_example.py
+```
+
+To learn more, check the CLI documentation available [here](https://huggingface.co/docs/accelerate/package_reference/cli).
 
 ## Launching multi-CPU run using MPI
 
@@ -168,15 +178,15 @@ 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`.
+🤗 Accelerate supports training on single/multiple GPUs using DeepSpeed. To use it, you don't need to change anything in your training code; you can set everything using just `accelerate config`. However, if you desire to tweak your DeepSpeed related args from your Python script, we provide you the `DeepSpeedPlugin`.
 
 ```python
 from accelerate import Accelerator, DeepSpeedPlugin
 
-# deepspeed needs to know your gradient accumulation steps before hand, so don't forget to pass it
+# deepspeed needs to know your gradient accumulation steps beforehand, so don't forget to pass it
 # Remember you still need to do gradient accumulation by yourself, just like you would have done without deepspeed
 deepspeed_plugin = DeepSpeedPlugin(zero_stage=2, gradient_accumulation_steps=2)
-accelerator = Accelerator(fp16=True, deepspeed_plugin=deepspeed_plugin)
+accelerator = Accelerator(mixed_precision='fp16', deepspeed_plugin=deepspeed_plugin)
 
 # How to save your 🤗 Transformer?
 accelerator.wait_for_everyone()
@@ -200,7 +210,7 @@ An example can be found in [this notebook](https://github.com/huggingface/notebo
 
 ## Why should I use 🤗 Accelerate?
 
-You should use 🤗 Accelerate when you want to easily run your training scripts in a distributed environment without having to renounce full control over your training loop. This is not a high-level framework above PyTorch, just a thin wrapper so you don't have to learn a new library, In fact the whole API of 🤗 Accelerate is in one class, the `Accelerator` object.
+You should use 🤗 Accelerate when you want to easily run your training scripts in a distributed environment without having to renounce full control over your training loop. This is not a high-level framework above PyTorch, just a thin wrapper so you don't have to learn a new library. In fact, the whole API of 🤗 Accelerate is in one class, the `Accelerator` object.
 
 ## Why shouldn't I use 🤗 Accelerate?
 
@@ -208,18 +218,24 @@ You shouldn't use 🤗 Accelerate if you don't want to write a training loop you
 
 ## 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:
+If you like the simplicity of 🤗 Accelerate but would prefer a higher-level abstraction around its capabilities, some frameworks and libraries that are built on top of 🤗 Accelerate are listed below:
 
 * [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.
+* [Catalyst](https://github.com/catalyst-team/catalyst#getting-started) is a PyTorch framework for Deep Learning Research and Development. It focuses on reproducibility, rapid experimentation, and codebase reuse so you can create something new rather than write yet another train loop. Catalyst provides a [Runner](https://catalyst-team.github.io/catalyst/api/core.html#runner) to connect all parts of the experiment: hardware backend, data transformations, model training, and inference logic.
 * [fastai](https://github.com/fastai/fastai#installing) is a PyTorch framework for Deep Learning that simplifies training fast and accurate neural nets using modern best practices. fastai provides a [Learner](https://docs.fast.ai/learner.html#Learner) to handle the training, fine-tuning, and inference of deep learning algorithms.
+* [Finetuner](https://github.com/jina-ai/finetuner) is a service that enables models to create higher-quality embeddings for semantic search, visual similarity search, cross-modal text<->image search, recommendation systems, clustering, duplication detection, anomaly detection, or other uses.
+* [InvokeAI](https://github.com/invoke-ai/InvokeAI) is a creative engine for Stable Diffusion models, offering industry-leading WebUI, terminal usage support, and serves as the foundation for many commercial products.
 * [Kornia](https://kornia.readthedocs.io/en/latest/get-started/introduction.html) is a differentiable library that allows classical computer vision to be integrated into deep learning models. Kornia provides a [Trainer](https://kornia.readthedocs.io/en/latest/x.html#kornia.x.Trainer) with the specific purpose to train and fine-tune the supported deep learning algorithms within the library.
-* [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!
+* [Open Assistant](https://projects.laion.ai/Open-Assistant/) is a chat-based assistant that understands tasks, can interact with their party systems, and retrieve information dynamically to do so. 
+* [pytorch-accelerated](https://github.com/Chris-hughes10/pytorch-accelerated) is a lightweight training library, with a streamlined feature set centered around a general-purpose [Trainer](https://pytorch-accelerated.readthedocs.io/en/latest/trainer.html), that places a huge emphasis on simplicity and transparency; enabling users to understand exactly what is going on under the hood, but without having to write and maintain the boilerplate themselves!
+* [Stable Diffusion web UI](https://github.com/AUTOMATIC1111/stable-diffusion-webui) is an open-source browser-based easy-to-use interface based on the Gradio library for Stable Diffusion.
+* [torchkeras](https://github.com/lyhue1991/torchkeras) is a simple tool for training pytorch model just in a keras style, a dynamic and beautiful plot is provided in notebook to monitor your loss or metric.
+* [transformers](https://github.com/huggingface/transformers) as a tool for helping train state-of-the-art machine learning models in PyTorch, Tensorflow, and JAX. (Accelerate is the backend for the PyTorch side).
 
 
 ## Installation
 
-This repository is tested on Python 3.6+ and PyTorch 1.4.0+
+This repository is tested on Python 3.8+ and PyTorch 1.10.0+
 
 You should install 🤗 Accelerate in a [virtual environment](https://docs.python.org/3/library/venv.html). If you're unfamiliar with Python virtual environments, check out the [user guide](https://packaging.python.org/guides/installing-using-pip-and-virtual-environments/).
 
@@ -240,7 +256,8 @@ pip install accelerate
 - multi-GPU on one node (machine)
 - multi-GPU on several nodes (machines)
 - TPU
-- FP16 with native AMP (apex on the roadmap)
+- FP16/BFloat16 mixed precision
+- FP8 mixed precision with [Transformer Engine](https://github.com/NVIDIA/TransformerEngine)
 - DeepSpeed support (Experimental)
 - PyTorch Fully Sharded Data Parallel (FSDP) support (Experimental)
 - Megatron-LM support (Experimental)
@@ -252,7 +269,7 @@ If you use 🤗 Accelerate in your publication, please cite it by using the foll
 ```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},
+  author =       {Sylvain Gugger, Lysandre Debut, Thomas Wolf, Philipp Schmid, Zachary Mueller, Sourab Mangrulkar, Marc Sun, Benjamin Bossan},
   howpublished = {\url{https://github.com/huggingface/accelerate}},
   year =         {2022}
 }

benchmarks/big_model_inference.py

@@ -16,12 +16,12 @@ 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
 
+from accelerate.utils import compute_module_sizes
+
 
 DEFAULT_MODELS = {
     "gpt-j-6b": {"is_causal": True, "model": "sgugger/sharded-gpt-j-6B", "tokenizer": "EleutherAI/gpt-j-6B"},

benchmarks/measures_util.py

@@ -2,9 +2,8 @@ import gc
 import threading
 import time
 
-import torch
-
 import psutil
+import torch
 
 
 class PeakCPUMemory:

docker/accelerate-cpu/Dockerfile

@@ -1,7 +1,7 @@
 # Builds CPU-only Docker image of PyTorch
 # Uses multi-staged approach to reduce size
 # Stage 1
-FROM python:3.7-slim as compile-image
+FROM python:3.8-slim as compile-image
 
 ARG DEBIAN_FRONTEND=noninteractive
 
@@ -25,7 +25,7 @@ RUN python3 -m pip install --no-cache-dir \
     --extra-index-url https://download.pytorch.org/whl/cpu
     
 # Stage 2
-FROM python:3.7-slim AS build-image
+FROM python:3.8-slim AS build-image
 COPY --from=compile-image /opt/venv /opt/venv
 RUN useradd -ms /bin/bash user
 USER user

docker/accelerate-gpu/Dockerfile

@@ -4,7 +4,7 @@
 # Use base conda image to reduce time
 FROM continuumio/miniconda3:latest AS compile-image
 # Specify py version
-ENV PYTHON_VERSION=3.7.3
+ENV PYTHON_VERSION=3.8
 # Install apt libs
 RUN apt-get update && \
     apt-get install -y curl git wget && \
@@ -23,7 +23,9 @@ SHELL ["/bin/bash", "-c"]
 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
+    --extra-index-url https://download.pytorch.org/whl/cu117
+
+RUN python3 -m pip install --no-cache-dir bitsandbytes
 
 # Stage 2
 FROM nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04 AS build-image

docs/README.md

@@ -67,7 +67,7 @@ doc-builder preview {package_name} {path_to_docs}
 For example:
 
 ```bash
-doc-builder preview transformers docs/source/en/
+doc-builder preview accelerate docs/source/
 ```
 
 The docs will be viewable at [http://localhost:3000](http://localhost:3000). You can also preview the docs once you have opened a PR. You will see a bot add a comment to a link where the documentation with your changes lives.
@@ -81,7 +81,7 @@ The `preview` command only works with existing doc files. When you add a complet
 
 ## Adding a new element to the navigation bar
 
-Accepted files are Markdown (.md or .mdx).
+Accepted files are Markdown (.md).
 
 Create a file with its extension and put it in the source directory. You can then link it to the toc-tree by putting
 the filename without the extension in the [`_toctree.yml`](https://github.com/huggingface/accelerate/blob/main/docs/source/_toctree.yml) file.

docs/source/_toctree.yml

@@ -22,13 +22,23 @@
   - local: usage_guides/training_zoo
     title: Example Zoo
   - local: usage_guides/big_modeling
-    title: How perform inference on large models with small resources
+    title: How to perform inference on large models with small resources
+  - local: usage_guides/model_size_estimator
+    title: Knowing how big of a model you can fit into memory
+  - local: usage_guides/quantization
+    title: How to quantize model 
+  - local: usage_guides/distributed_inference
+    title: How to perform distributed inference with normal resources
   - local: usage_guides/gradient_accumulation
     title: Performing gradient accumulation
+  - local: usage_guides/local_sgd
+    title: Accelerating training with local SGD
   - local: usage_guides/checkpoint
     title: Saving and loading training states
   - local: usage_guides/tracking
     title: Using experiment trackers
+  - local: usage_guides/debug
+    title: Debugging timeout errors
   - local: usage_guides/memory
     title: How to avoid CUDA Out-of-Memory
   - local: usage_guides/mps
@@ -41,8 +51,12 @@
     title: How to use Megatron-LM
   - local: usage_guides/sagemaker
     title: How to use 🤗 Accelerate with SageMaker
+  - local: usage_guides/ipex
+    title: How to use 🤗 Accelerate with Intel® Extension for PyTorch for cpu
   title: How-To Guides
 - sections:
+  - local: concept_guides/big_model_inference
+    title: Loading big models into memory
   - local: concept_guides/performance
     title: Comparing performance across distributed setups
   - local: concept_guides/deferring_execution
@@ -77,4 +91,6 @@
     title: Utility functions and classes
   - local: package_reference/megatron_lm
     title: Megatron-LM Utilities
-  title: "Reference"
+  - local: package_reference/fsdp
+    title: Fully Sharded Data Parallelism Utilities
+  title: "Reference"

docs/source/basic_tutorials/install.md

@@ -8,11 +8,14 @@ http://www.apache.org/licenses/LICENSE-2.0
 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
 an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
 specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
 -->
 
 # Installation 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+**.
+Before you start, you will need to setup your environment, install the appropriate packages, and configure 🤗 Accelerate. 🤗 Accelerate is tested on **Python 3.8+**.
 
 ## Installing 🤗 Accelerate
 

docs/source/basic_tutorials/launch.md

@@ -8,6 +8,9 @@ http://www.apache.org/licenses/LICENSE-2.0
 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
 an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
 specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
 -->
 
 # Launching your 🤗 Accelerate scripts
@@ -36,7 +39,7 @@ for batch in training_dataloader:
 
 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:
+First, you should rewrite the above code into a function, and make it callable as a script. For example:
 
 ```diff
   from accelerate import Accelerator
@@ -61,7 +64,7 @@ First you should rewrite the above code into a function, and make it callable as
 +     main()
 ```
 
-Next you need to launch it with `accelerate launch`. 
+Next, you need to launch it with `accelerate launch`. 
 
 <Tip warning={true}>
 
@@ -74,7 +77,7 @@ Next you need to launch it with `accelerate launch`.
 ## 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.
+This command wraps around all of the different commands needed to launch your script on various platforms, without you having to remember what each of them is.
 
 <Tip>
 
@@ -88,7 +91,7 @@ You can launch your script quickly by using:
 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!
+Just put `accelerate launch` at the start of your command, and pass in additional arguments and parameters to your script afterward like normal!
 
 Since this runs the various torch spawn methods, all of the expected environment variables can be modified here as well.
 For example, here is how to use `accelerate launch` with a single GPU:
@@ -105,6 +108,12 @@ Here is how you would use all GPUs and train with mixed precision disabled:
 accelerate launch --multi_gpu {script_name.py} {--arg1} {--arg2} ...
 ```
 
+Or by specifying a number of GPUs to use:
+
+```bash
+accelerate launch --num_processes=2 {script_name.py} {--arg1} {--arg2} ...
+```
+
 To get more specific you should pass in the needed parameters yourself. For instance, here is how you 
 would also launch that same script on two GPUs using mixed precision while avoiding all of the warnings: 
 
@@ -130,6 +139,21 @@ For a visualization of this difference, that earlier `accelerate launch` on mult
 MIXED_PRECISION="fp16" torchrun --nproc_per_node=2 --num_machines=1 {script_name.py} {--arg1} {--arg2} ...
 ```
 
+You can also launch your script utilizing the launch CLI as a python module itself, enabling the ability to pass in other python-specific
+launching behaviors. To do so, use `accelerate.commands.launch` instead of `accelerate launch`:
+
+```bash
+python -m accelerate.commands.launch --num_processes=2 {script_name.py} {--arg1} {--arg2}
+```
+
+If you want to execute the script with any other python flags, you can pass them in as well similar to `-m`, such as 
+the below example enabling unbuffered stdout and stderr:
+
+```bash
+python -u -m accelerate.commands.launch --num_processes=2 {script_name.py} {--arg1} {--arg2}
+```
+
+
 ## Why you should always use `accelerate config`
 
 Why is it useful to the point you should **always** run `accelerate config`? 
@@ -175,4 +199,4 @@ 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.md

@@ -8,13 +8,16 @@ http://www.apache.org/licenses/LICENSE-2.0
 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
 an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
 specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
 -->
 
 # 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!
+your way toward running your code on distributed systems with ease!
 
 ## The base training loop
 
@@ -65,7 +68,7 @@ change the definition of `device` to come from [`Accelerator`]:
 
 ### Preparing your objects
 
-Next you need to pass all of the important objects related to training into [`~Accelerator.prepare`]. 🤗 Accelerate will
+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:
 
 ```
@@ -73,7 +76,7 @@ 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.
+These objects are returned in the same order they were sent in. 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}>
@@ -121,3 +124,6 @@ for batch in training_dataloader:
     scheduler.step()
 ```
 
+## More Resources
+
+To check out more ways on how to migrate to 🤗 Accelerate, check out our [interactive migration tutorial](https://huggingface.co/docs/accelerate/usage_guides/explore) which showcases other items that need to be watched for when using Accelerate and how to do so quickly.

docs/source/basic_tutorials/notebook.md

@@ -8,9 +8,12 @@ http://www.apache.org/licenses/LICENSE-2.0
 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
 an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
 specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
 -->
 
-# Launching Multi-Node Training from a Jupyter Environment
+# Launching Multi-GPU 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.
@@ -35,7 +38,7 @@ The following code will restart Jupyter after writing the configuration, as CUDA
 
 <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.
+    CUDA can't be initialized more than once on a multi-GPU system. It's fine to debug in the notebook and have calls to CUDA, but in order to finally train a full cleanup and restart will need to be performed.
     
 </Tip>
 
@@ -153,7 +156,7 @@ def get_dataloaders(batch_size: int = 64):
     random_perm = np.random.permutation(len(fnames))
     cut = int(0.8 * len(fnames))
     train_split = random_perm[:cut]
-    eval_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()])
@@ -337,7 +340,7 @@ def training_loop(mixed_precision="fp16", seed: int = 42, batch_size: int = 64):
     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
+    # To make these constants available on the active device, set it to the accelerator device
     mean = mean.to(accelerator.device)
     std = std.to(accelerator.device)
 
@@ -398,6 +401,26 @@ args = ("fp16", 42, 64)
 notebook_launcher(training_loop, args, num_processes=2)
 ```
 
+In the case of running on multiple nodes, you need to set up a Jupyter session at each node and run the launching cell at the same time.
+
+For an environment containing 2 nodes (computers) with 8 GPUs each and the main computer with an IP address of "172.31.43.8", it would look like so:
+
+```python
+notebook_launcher(training_loop, args, master_addr="172.31.43.8", node_rank=0, num_nodes=2, num_processes=8)
+```
+
+And in the second Jupyter session on the other machine:
+
+<Tip>
+
+    Notice how the `node_rank` has changed
+
+</Tip>
+
+```python
+notebook_launcher(training_loop, args, master_addr="172.31.43.8", node_rank=1, num_nodes=2, num_processes=8)
+```
+
 In the case of running on the TPU, it would look like so:
 
 ```python
@@ -420,10 +443,17 @@ epoch 4: 94.71
 
 And that's it!
 
+## Debugging 
+
+A common issue when running the `notebook_launcher` is receiving a CUDA has already been initialized issue. This usually stems
+from an import or prior code in the notebook that makes a call to the PyTorch `torch.cuda` sublibrary. To help narrow down what went wrong,
+you can launch the `notebook_launcher` with `ACCELERATE_DEBUG_MODE=yes` in your environment and an additional check
+will be made when spawning that a regular process can be created and utilize CUDA without issue. (Your CUDA code can still be ran afterwards).
+
 ## Conclusion
 
 This notebook showed how to perform distributed training from inside of a Jupyter Notebook. Some key notes to remember:
 
 - Make sure to save any code that use CUDA (or CUDA imports) for the function passed to [`notebook_launcher`]
 - Set the `num_processes` to be the number of devices used for training (such as number of GPUs, CPUs, TPUs, etc)
-- If using the TPU, declare your model outside the training loop function
+- If using the TPU, declare your model outside the training loop function

docs/source/basic_tutorials/overview.md

@@ -8,6 +8,9 @@ http://www.apache.org/licenses/LICENSE-2.0
 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
 an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
 specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
 -->
 
 # Overview

docs/source/concept_guides/big_model_inference.md

@@ -8,11 +8,14 @@ http://www.apache.org/licenses/LICENSE-2.0
 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
 an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
 specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
 -->
 
 # Handling big models for inference
 
-When loading a pretrained model in PyTorch, the usual workflow looks like this:
+When loading a pre-trained model in PyTorch, the usual workflow looks like this:
 
 ```py
 import torch
@@ -27,11 +30,11 @@ In plain English, those steps are:
 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).
+While this works very well for regularly sized models, this workflow has some clear limitations when we deal with a huge model: in step 1, we load a full version of the model in RAM, and spend some time randomly initializing the weights (which will be discarded in step 3). In step 2, we load another full version of the model in RAM, with the pre-trained weights. If you're loading a model with 6 billion parameters, this means you will need 24GB of RAM for each copy of the model, so 48GB in total (half of it to load the model in FP16).
 
 <Tip warning={true}>
 
-    This API is quite new and still in its experimental stage. While we strive to provide a stable API, it's possible some small parts of the public API will change in the future.
+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>
 
@@ -43,7 +46,7 @@ While this works very well for regularly sized models, this workflow has some cl
 
 ### 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:
+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
@@ -59,7 +62,7 @@ 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.
+initializes an empty model with a bit more than 100B parameters. Behind the scenes, this relies on the meta device introduced in PyTorch 1.9. During the initialization under the context manager, each time a parameter is created, it is instantly moved to that device.
 
 <Tip warning={true}>
 
@@ -69,9 +72,9 @@ initializes an empty model with a bit more than 100B parameters. Behind the scen
 
 ### 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.
+It's possible your model is so big that even a single copy won't fit in RAM. That doesn't mean it can't be loaded: if you have one or several GPUs, this is more memory available to store your model. In this case, it's better if your checkpoint is split into several smaller files that we call checkpoint shards.
 
-🤗 Accelerate will handle sharded checkpoints as long as you follow the following format: your checkpoint should be in a folder, with several files containing the partial state dicts, and there should be an index in the JSON format that contains a dictionary mapping parameter names to the file containing their weights. For instance we could have a folder containing:
+🤗 Accelerate will handle sharded checkpoints as long as you follow the following format: your checkpoint should be in a folder, with several files containing the partial state dicts, and there should be an index in the JSON format that contains a dictionary mapping parameter names to the file containing their weights. You can easily shard your model with [`~Accelerator.save_model`]. For instance, we could have a folder containing:
 
 ```bash
 first_state_dict.bin
@@ -96,44 +99,65 @@ and `first_state_dict.bin` containing the weights for `"linear1.weight"` and `"l
 
 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:
+If you want to use big model inference with 🤗 Transformers models, check out this [documentation](https://huggingface.co/docs/transformers/main/en/main_classes/model#large-model-loading).
+
+Here is how we can use this to load the [GPT2-1.5B](https://huggingface.co/marcsun13/gpt2-xl-linear-sharded) model.
+
+Let's download the sharded version of this model.
 
 ```bash
-git clone https://huggingface.co/sgugger/sharded-gpt-j-6B
-cd sharded-gpt-j-6B
-git-lfs install
-git pull
+pip install huggingface_hub
 ```
 
-then we can initialize the model with
+```py
+from huggingface_hub import snapshot_download
+checkpoint = "marcsun13/gpt2-xl-linear-sharded"
+weights_location = snapshot_download(repo_id=checkpoint)
+```
+
+In order to initialize the model, we will use the library minGPT. 
+
+```bash
+git clone https://github.com/karpathy/minGPT.git
+pip install minGPT/
+```
 
 ```py
 from accelerate import init_empty_weights
-from transformers import AutoConfig, AutoModelForCausalLM
+from mingpt.model import GPT
 
-checkpoint = "EleutherAI/gpt-j-6B"
-config = AutoConfig.from_pretrained(checkpoint)
+model_config = GPT.get_default_config()
+model_config.model_type = 'gpt2-xl'
+model_config.vocab_size = 50257
+model_config.block_size = 1024
 
 with init_empty_weights():
-    model = AutoModelForCausalLM.from_config(config)
+    model = GPT(model_config)
 ```
 
-and load the checkpoint we just downloaded with:
+Then, 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"]
+    model, checkpoint=weights_location, device_map="auto", no_split_module_classes=['Block']
 )
 ```
 
 By passing `device_map="auto"`, we tell 🤗 Accelerate to determine automatically where to put each layer of the model depending on the available resources:
-- first we use the maximum space available on the GPU(s)
+- 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.
+
+#### `no_split_module_classes`
+
+This parameter will indicate that some of the modules with the name `"Block"` should not be split across different devices. You should set here all blocks that 
+include a residutal connection of some kind.
+
+
+#### The `device_map`
 
 You can see the `device_map` that 🤗 Accelerate picked by accessing the `hf_device_map` attribute of your model:
 
@@ -143,43 +167,34 @@ model.hf_device_map
 
 ```python out
 {'transformer.wte': 0,
+ 'transformer.wpe': 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.21': 0, 
+ 'transformer.h.22': 1, 
+ 'transformer.h.23': 1, 
  'transformer.h.24': 1,
- 'transformer.h.25': 1,
- 'transformer.h.26': 1,
- 'transformer.h.27': 1,
- 'transformer.ln_f': 1,
+ ...
+ 'transformer.h.47': 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:
+It's fully possible to create your own device map for the layers to use as well, specifying the GPU device to use (a number), `"cpu"`, or `"disk"` and pass this in:
+
+```python
+device_map = {
+    "transformer.wte": "cpu",
+    "transformer.wpe": 0,
+    "transformer.drop": "cpu",
+    "transformer.h.0": "disk"
+}
+
+model = load_checkpoint_and_dispatch(
+    model, checkpoint=weights_location, device_map=device_map
+)
 
-```py
-model = load_checkpoint_and_dispatch(model, "sharded-gpt-j-6B", device_map=my_device_map)
 ```
 
 ### Run the model
@@ -187,31 +202,30 @@ model = load_checkpoint_and_dispatch(model, "sharded-gpt-j-6B", device_map=my_de
 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
+from mingpt.bpe import BPETokenizer
+tokenizer = BPETokenizer()
+inputs = tokenizer("Hello, my name is").to(0)
 
-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())
+outputs = model.generate(x1, max_new_tokens=10, do_sample=False)[0]
+tokenizer.decode(outputs.cpu().squeeze())
 ```
 
 Behind the scenes, 🤗 Accelerate added hooks to the model, so that:
 - at each layer, the inputs are put on the right device (so even if your model is spread across several GPUs, it works)
-- for the weights offloaded on the CPU, they are put on a GPU just before the forward pass, and cleaned up just after
-- for the weights offloaded on the hard drive, they are loaded in RAM then put on a GPU just before the forward pass, and cleaned up just after
+- 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!
+This way, your model can run for inference even if it doesn't fit on one of the GPUs or the CPU RAM!
 
 <Tip warning={true}>
 
-    This only supports 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.
+    This only supports the inference of your model, not training. Most of the computation happens behind `torch.no_grad()` context managers to avoid spending some GPU memory with intermediate activations.
 
 </Tip>
 
 ### Designing a device map
 
-You can let 🤗 Accelerate handle the device map computation by setting `device_map` to one of the supported options (`"auto"`, `"balanced"`, `"balanced_low_0"`, `"sequential"`) or create one yourself, if you want more control over where each layer should go.
+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>
 
@@ -221,7 +235,7 @@ You can let 🤗 Accelerate handle the device map computation by setting `device
 
 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:
+When you have more GPU memory available than the model size, here is the difference between each option:
 - `"auto"` and `"balanced"` evenly split the model on all available GPUs, making it possible for you to use a batch size greater than 1.
 - `"balanced_low_0"` evenly splits the model on all GPUs except the first one, and only puts on GPU 0 what does not fit on the others. This option is great when you need to use GPU 0 for some processing of the outputs, like when using the `generate` function for Transformers models
 - `"sequential"` will fit what it can on GPU 0, then move on GPU 1 and so forth (so won't use the last GPUs if it doesn't need to).
@@ -232,9 +246,9 @@ When you have more GPU memory available than the model size, here the difference
 
 </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"`.
+First note that you can limit the memory used on each GPU by using the `max_memory` argument (available in [`infer_auto_device_map`] and in all functions using it). When setting `max_memory`, you should pass along a dictionary containing the GPU identifiers (for instance `0`, `1` etc.) and the `"cpu"` key for the maximum RAM you want to use for CPU offload. The values can either be an integer (in bytes) or a string representing a number with its unit, such as `"10GiB"` or `"10GB"`.
 
-Here is an example where we don't want to use more than 10GiB on each of two GPUs and no more than 30GiB of CPU RAM for the model weights:
+Here is an example where we don't want to use more than 10GiB on each of the two GPUs and no more than 30GiB of CPU RAM for the model weights:
 
 ```python
 from accelerate import infer_auto_device_map
@@ -246,18 +260,18 @@ device_map = infer_auto_device_map(my_model, max_memory={0: "10GiB", 1: "10GiB",
 
     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.
+    Therefore when you create memory maps with `max_memory` make sure to adjust the available memory accordingly to avoid out-of-memory errors.
 
 </Tip>
 
-Additionally, if you do some additional operations with your outputs without placing them back on the CPU (for instance inside the `generate` method of Transformers) and if you placed your inputs on a GPU, that GPU will consume more memory than the others (Accelerate always place the output back to the device of the input). Therefore if you would like to optimize the maximum batch size and you have many GPUs, give the first GPU less memory. For example, with BLOOM-176B on 8x80 A100 setup the close to ideal map is:
+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:
+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}
@@ -286,9 +300,9 @@ device_map = {"block1": 0, "block2.linear1": 1, "block2.linear2": 1}
 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`]) tries to maximize GPU and CPU RAM it sees available when you execute it. While PyTorch is very good at managing GPU RAM efficiently (and giving it back when not needed), it's not entirely true with Python and CPU RAM. Therefore, an automatically computed device map might be too intense on the CPU. Move a few modules to the disk device if you get crashes due to a lack of RAM.
 - [`infer_auto_device_map`] (or `device_map="auto"` in [`load_checkpoint_and_dispatch`]) attributes devices sequentially (to avoid moving things back and forth) so if your first layer is bigger than the size of the GPU you have, it will end up with everything on the CPU/Disk.
 - [`load_checkpoint_and_dispatch`] and [`load_checkpoint_in_model`] do not perform any check on the correctness of your state dict compared to your model at the moment (this will be fixed in a future version), so you may get some weird errors if trying to load a checkpoint with mismatched or missing keys.
 - The model parallelism used when your model is split on several GPUs is naive and not optimized, meaning that only one GPU works at a given time and the other sits idle.
 - When weights are offloaded on the CPU/hard drive, there is no pre-fetching (yet, we will work on this for future versions) which means the weights are put on the GPU when they are needed and not before.
-- Hard-drive offloading might be very slow if the hardware you run on does not have fast communication between disk and CPU (like NVMes).
+- Hard-drive offloading might be very slow if the hardware you run on does not have fast communication between disk and CPU (like NVMes).

docs/source/concept_guides/deferring_execution.md

@@ -8,6 +8,9 @@ http://www.apache.org/licenses/LICENSE-2.0
 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
 an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
 specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
 -->
 
 # Deferring Executions
@@ -27,7 +30,7 @@ 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:
+A few example cases of when to use this utility are listed below:
 
 <Tip>
 
@@ -38,7 +41,7 @@ A few example cases for when to use this utility are listed below:
 
 ## Downloading a Dataset 
 
-When downloading a dataset, you should download it first on the main process and then loading the cached dataset in afterwards
+When downloading a dataset, you should download it first on the main process and then load the cached dataset afterward
 
 <Tip>
 
@@ -104,4 +107,24 @@ with accelerator.main_process_first():
         batched=True,
         remove_columns=["idx", "sentence1", "sentence2"],
     )
+```
+
+## Applying checks such as Early Stopping
+
+To have a check that works with a flag set by a particular process, the `set_trigger` and `check_trigger` API should be used. Useful examples
+for doing so can include situations such as using early stopping and monitoring the loss (as each loss slightly differs on each process).
+
+Call [`Accelerator.set_trigger`] when your condition has been met, and [`Accelerator.check_trigger`] when checking if that condition has been met in any process:
+
+```python
+for (x,y) in data_loader:
+    logits = model(x)
+    loss = loss_func(logits, y)
+    # Assume `should_do_early_stopping` is a custom defined function that returns a conditional
+    if should_do_early_stopping(loss):
+        accelerator.set_trigger()
+
+    # Later in the training script when we need to check for the breakpoint
+    if accelerator.check_trigger():
+        break
 ```

docs/source/concept_guides/gradient_synchronization.md

@@ -8,6 +8,9 @@ http://www.apache.org/licenses/LICENSE-2.0
 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
 an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
 specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
 -->
 
 # Gradient Synchronization
@@ -45,22 +48,30 @@ 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?
+The most direct example is when you update model parameters through
+`optimizer.step()`.
+Without gradient accumulation, all instances of the model need to have updated
+their gradients computed, collated, and updated before moving on to the next
+batch of data.
+When performing gradient accumulation, you accumulate `n` loss gradients and
+skip `optimizer.step()` until `n` batches have been reached. As all training
+processes only need to sychronize by the time `optimizer.step()` is called,
+without any modification to your training step, this neededless inter-process
+communication can cause a significant slowdown.
+
+ 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. 
+Since you are skipping model parameter updates when training on these batches, their gradients do not need to be synchronized until the point where `optimizer.step()` is actually called. 
 PyTorch cannot automagically tell when you need to do this, but they do provide a tool to help through the [`no_sync`](https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html#torch.nn.parallel.DistributedDataParallel.no_sync) context manager
 that is added to your model after converting it to DDP.
 
-Under this context manager, PyTorch will skip synchronizing the gradients when `.backward()` is called, and the first call to `.backward()` outside this 
+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)
+ddp_model, dataloader, optimizer = accelerator.prepare(model, dataloader, optimizer)
 
 for index, batch in enumerate(dataloader):
     inputs, targets = batch
@@ -76,13 +87,14 @@ for index, batch in enumerate(dataloader):
         outputs = ddp_model(inputs)
         loss = loss_func(outputs)
         accelerator.backward(loss)
+        optimizer.step()
 ```
 
 In 🤗 Accelerate to make this an API that can be called no matter the training device (though it may not do anything if you are not in a distributed system!),
 `ddp_model.no_sync` gets replaced with [`~Accelerator.no_sync`] and operates the same way:
 
 ```diff
-  ddp_model, dataloader = accelerator.prepare(model, dataloader)
+  ddp_model, dataloader, optimizer = accelerator.prepare(model, dataloader, optimizer)
 
   for index, batch in enumerate(dataloader):
       inputs, targets = batch
@@ -99,13 +111,15 @@ In 🤗 Accelerate to make this an API that can be called no matter the training
           outputs = ddp_model(inputs)
           loss = loss_func(outputs)
           accelerator.backward(loss)
+          optimizer.step()
+          optimizer.zero_grad()
 ```
 
 As you may expect, the [`~Accelerator.accumulate`] function wraps around this conditional check by keeping track of the current batch number, leaving you with the final
 gradient accumulation API:
 
 ```python
-ddp_model, dataloader = accelerator.prepare(model, dataloader)
+ddp_model, dataloader, optimizer = accelerator.prepare(model, dataloader, optimizer)
 
 for batch in dataloader:
     with accelerator.accumulate(model):
@@ -114,6 +128,42 @@ for batch in dataloader:
         outputs = model(inputs)
         loss = loss_function(outputs, targets)
         accelerator.backward(loss)
+        optimizer.step()
+        optimizer.zero_grad()
 ```
 
-As a result, you should either use *`accelerator.accumulate` or `accelerator.no_sync`* when it comes to API choice. 
+As a result, you should either use *`accelerator.accumulate` or `accelerator.no_sync`* when it comes to API choice. 
+
+## Just how much of a slowdown is there, and easy mistakes you can make
+
+To set up a realistic example, consider the following setup:
+
+* Two single-GPU T4 nodes and one node with two GPUs
+* Each GPU is a T4, and are hosted on GCP
+* The script used is a modification of the [NLP Example](https://github.com/muellerzr/timing_experiments/blob/main/baseline.py) script
+* Batch size per GPU is 16, and gradients are accumulated every 4 steps
+
+All scripts are available in [this repository](https://github.com/muellerzr/timing_experiments).
+
+If not careful about gradient synchronization and GPU communication, a *large* amount of time can be wasted 
+from when these GPUs communicate to each other during unnecessary periods.
+
+By how much?
+
+Reference:
+- Baseline: uses no synchronization practices discussed here
+- `no_sync` improperly: `no_sync` only around the `backward` call, not the `forward`
+- `no_sync`: using the `no_sync` pattern properly
+- `accumulate`: using [`~Accelerator.accumulate`] properly
+
+Below are the average seconds per batch iterating over 29 batches of data for each setup on both a single node and on the dual-node setup:
+
+|             | Baseline  | `no_sync` improperly | `no_sync` | `accumulate`| 
+| :---------: | :-------: | :------------------: | :-------: | :---------: |
+| Multi-Node  | 2±0.01s    | 2.13±0.08s | **0.91±0.11s** | **0.91±0.11s** |
+| Single Node | 0.50±0.01s | 0.50±0.01s | **0.41±0.015s** | **0.41±0.015s** |
+
+As you can see, if you are not careful about how you set up your gradient synchronization, you can get upwards of more than a 2x slowdown during training!
+
+If you are worried about making sure everything is done properly, we highly recommend utilizing the [`~Accelerator.accumulate`] function and passing in
+`gradient_accumulation_steps` or `gradient_accumulation_plugin` to the [`Accelerator`] object so Accelerate can handle this for you.

docs/source/concept_guides/performance.md

@@ -8,6 +8,9 @@ http://www.apache.org/licenses/LICENSE-2.0
 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
 an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
 specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
 -->
 
 # Comparing performance between different device setups
@@ -18,7 +21,7 @@ and expect your results to line up.
 
 But why?
 
-There's three reasons for this that this tutorial will cover: 
+There are three reasons for this that this tutorial will cover: 
 
 1. **Setting the right seeds**
 2. **Observed Batch Sizes**
@@ -26,10 +29,10 @@ There's three reasons for this that this tutorial will cover:
 
 ## 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:
+While this issue has not come up as much, make sure to use [`utils.set_seed`] to fully set the seed in all distributed cases so training will be reproducible:
 
 ```python
-from accelerate import set_seed
+from accelerate.utils import set_seed
 
 set_seed(42)
 ```
@@ -58,7 +61,7 @@ The below table can be used as a quick reference to try out different batch size
 
 <Tip>
 
-In this example there are two GPUs for "Multi-GPU" and a TPU pod with 8 workers
+In this example, there are two GPUs for "Multi-GPU" and a TPU pod with 8 workers
 
 </Tip>
 
@@ -90,5 +93,11 @@ optimizer = AdamW(params=model.parameters(), lr=learning_rate)
 ```
 
 You will also find that `accelerate` will step the learning rate based on the number of processes being trained on. This is because 
-of the observed batch size noted earlier. So in a case of 2 GPUs, the learning rate will be stepped twice as often as a single GPU
-to account for the batch size being twice as large (if no changes to the batch size on the single GPU instance are made).
+of the observed batch size noted earlier. So in the case of 2 GPUs, the learning rate will be stepped twice as often as a single GPU
+to account for the batch size being twice as large (if no changes to the batch size on the single GPU instance are made).
+
+## Gradient Accumulation and Mixed Precision
+
+When using gradient accumulation and mixed precision, due to how gradient averaging works (accumulation) and the precision loss (mixed precision), 
+some degradation in performance is expected. This will be explicitly seen when comparing the batch-wise loss between different compute 
+setups. However, the overall loss, metric, and general performance at the end of training should be _roughly_ the same.

docs/source/concept_guides/training_tpu.md

@@ -8,11 +8,14 @@ http://www.apache.org/licenses/LICENSE-2.0
 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
 an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
 specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
 -->
 
 # Training on TPUs with 🤗 Accelerate
 
-Training on TPUs can be slightly different than training on multi-gpu, even with 🤗 Accelerate. This guide aims to show you 
+Training on TPUs can be slightly different from training on multi-gpu, even with 🤗 Accelerate. This guide aims to show you 
 where you should be careful and why, as well as the best practices in general.
 
 ## Training in a Notebook
@@ -24,8 +27,8 @@ While on a TPU that last part is not as important, a critical part to understand
 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 
+Where this becomes important is in regard to declaring your model. On forked TPU processes, it is recommended that you instantiate your model *once* and pass this into your 
+training function. This is different than training on GPUs where you create `n` models that have their gradients synced and back-propagated at certain moments. Instead, one 
 model instance is shared between all the nodes and it is passed back and forth. This is important especially when training on low-resource TPUs such as those provided in Kaggle kernels or
 on Google Colaboratory. 
 
@@ -134,7 +137,7 @@ At the base level, this is enabled when passing `mixed_precision="bf16"` to `Acc
 ```python
 accelerator = Accelerator(mixed_precision="bf16")
 ```
-By default this will cast `torch.float` and `torch.double` to `bfloat16` on TPUs. 
+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 
@@ -161,4 +164,4 @@ new batch size after the first few iterations.
 
     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>
+</Tip>

docs/source/index.md

@@ -8,6 +8,9 @@ http://www.apache.org/licenses/LICENSE-2.0
 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
 an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
 specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
 -->
 
 # Accelerate

docs/source/package_reference/accelerator.md

@@ -8,12 +8,15 @@ http://www.apache.org/licenses/LICENSE-2.0
 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
 an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
 specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
 -->
 
 # Accelerator
 
 The [`Accelerator`] is the main class provided by 🤗 Accelerate. 
-It serves at the main entrypoint for the API. 
+It serves at the main entry point for the API. 
 
 ## Quick adaptation of your code
 
@@ -45,7 +48,7 @@ you should search for and replace by the corresponding methods of your `accelera
 
 ### Printing
 
-`print` statements should be replaced by [`~Accelerator.print`] to be printed once per process
+`print` statements should be replaced by [`~Accelerator.print`] to be printed once per process:
 
 ```diff
 - print("My thing I want to print!")
@@ -113,25 +116,55 @@ def do_my_thing():
 
 ### Synchronicity control
 
-Use [`~Accelerator.wait_for_everyone`] to make sure all processes join that point before continuing. (Useful before a model save for instance)
+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`:
+Use [`~Accelerator.save_model`] instead of `torch.save` to save a model. It will remove all model wrappers added during the distributed process, get the state_dict of the model and save it. The state_dict will be in the same precision as the model being trained.
 
 ```diff
-  state_dict = model.state_dict()
 - torch.save(state_dict, "my_state.pkl")
-+ accelerator.save(state_dict, "my_state.pkl")
++ accelerator.save_model(model, save_directory)
+```
+
+[`~Accelerator.save_model`] can also save a model into sharded checkpoints or with safetensors format.
+Here is an example: 
+
+```python
+accelerator.save_model(model, save_directory, max_shard_size="1GB", safe_serialization=True)
+```
+
+#### 🤗 Transformers models
+
+If you are using models from the [🤗 Transformers](https://huggingface.co/docs/transformers/) library, you can use the `.save_pretrained()` method.
+
+```python
+from transformers import AutoModel
+
+model = AutoModel.from_pretrained("bert-base-cased")
+model = accelerator.prepare(model)
+
+# ...fine-tune with PyTorch...
+
+unwrapped_model = accelerator.unwrap_model(model)
+unwrapped_model.save_pretrained(
+    "path/to/my_model_directory",
+    is_main_process=accelerator.is_main_process,
+    save_function=accelerator.save,
+)
+```
+
+This will ensure your model stays compatible with other 🤗 Transformers functionality like the `.from_pretrained()` method.
+
+```python
+from transformers import AutoModel
+
+model = AutoModel.from_pretrained("path/to/my_model_directory")
 ```
 
 ### Operations
@@ -157,7 +190,22 @@ multi-device training, check if the step should actually be performed, and auto-
           scheduler.step()
           optimizer.zero_grad()
 ```
+#### GradientAccumulationPlugin
+[[autodoc]] utils.GradientAccumulationPlugin
+
+
+Instead of passing `gradient_accumulation_steps` you can instantiate a GradientAccumulationPlugin and pass it to the [`Accelerator`]'s `__init__`
+as `gradient_accumulation_plugin`. You can only pass either one of `gradient_accumulation_plugin` or `gradient_accumulation_steps` passing both will raise an error.
+```diff
+from accelerate.utils import GradientAccumulationPlugin
+
+gradient_accumulation_plugin = GradientAccumulationPlugin(num_steps=2)
+- accelerator = Accelerator()
++ accelerator = Accelerator(gradient_accumulation_plugin=gradient_accumulation_plugin)
+```
+
+In addition to the number of steps, this also lets you configure whether or not you adjust your learning rate scheduler to account for the change in steps due to accumulation.
 
 ## Overall API documentation:
 
-[[autodoc]] Accelerator
+[[autodoc]] Accelerator

docs/source/package_reference/big_modeling.md

@@ -8,6 +8,9 @@ http://www.apache.org/licenses/LICENSE-2.0
 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
 an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
 specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
 -->
 
 # Working with large models
@@ -19,6 +22,8 @@ specific language governing permissions and limitations under the License.
 [[autodoc]] big_modeling.disk_offload
 [[autodoc]] big_modeling.dispatch_model
 [[autodoc]] big_modeling.load_checkpoint_and_dispatch
+[[autodoc]] big_modeling.load_checkpoint_in_model
+[[autodoc]] utils.infer_auto_device_map
 
 ## Model Hooks
 

docs/source/package_reference/cli.md

@@ -8,6 +8,9 @@ http://www.apache.org/licenses/LICENSE-2.0
 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
 an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
 specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
 -->
 
 # The Command Line 
@@ -83,7 +86,7 @@ accelerate config update [arguments]
 
 **Command**:
 
-`accelerate env` or `accelerate-env`
+`accelerate env` or `accelerate-env` or `python -m accelerate.commands.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).
 
@@ -103,7 +106,7 @@ accelerate env [arguments]
 
 **Command**:
 
-`accelerate launch` or `accelerate-launch`
+`accelerate launch` or `accelerate-launch` or `python -m accelerate.commands.launch`
 
 Launches a specified script on a distributed system with the right parameters.
 
@@ -135,8 +138,8 @@ values. They can also be passed in manually.
 
 * `--cpu` (`bool`) -- Whether or not to force the training on the CPU.
 * `--multi_gpu` (`bool`) -- Whether or not this should launch a distributed GPU training.
-* `--mps` (`bool`) -- Whether or not this should use MPS-enabled GPU device on MacOS machines.
 * `--tpu` (`bool`) -- Whether or not this should launch a TPU training.
+* `--ipex` (`bool`) -- Whether or not this should launch an Intel Pytorch Extension (IPEX) training.
 
 **Resource Selection Arguments**:
 
@@ -154,6 +157,7 @@ The following arguments are useful for selecting which training paradigm to use.
 * `--use_deepspeed` (`bool`) -- Whether or not to use DeepSpeed for training.
 * `--use_fsdp` (`bool`) -- Whether or not to use FullyShardedDataParallel for training.
 * `--use_megatron_lm` (`bool`) -- Whether or not to use Megatron-LM for training.
+* `--use_xpu` (`bool`) -- Whether to use IPEX plugin to speed up training on XPU specifically.
 
 **Distributed GPU Arguments**:
 
@@ -164,6 +168,7 @@ The following arguments are only useful when `multi_gpu` is passed or multi-gpu
 * `--machine_rank MACHINE_RANK` (`int`) -- The rank of the machine on which this script is launched.
 * `--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.
+* `--rdzv_backend` (`str`) -- The rendezvous method to use, such as "static" or "c10d"
 * `--rdzv_conf` (`str`) -- Additional rendezvous configuration (<key1>=<value1>,<key2>=<value2>,...).
 * `--max_restarts` (`int`) -- Maximum number of worker group restarts before failing.
 * `--monitor_interval` (`float`) -- Interval, in seconds, to monitor the state of workers.
@@ -223,6 +228,36 @@ 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 estimate-memory
+
+**Command**:
+
+`accelerate estimate-memory` or `accelerate-estimate-memory` or `python -m accelerate.commands.estimate`
+
+Estimates the total vRAM a particular model hosted on the Hub needs to be loaded in with an estimate for training. Requires that `huggingface_hub` be installed. 
+
+<Tip>
+
+    When performing inference, typically add ≤20% to the result as overall allocation [as referenced here](https://blog.eleuther.ai/transformer-math/). We will have more extensive estimations in the future that will automatically be included in the calculation.
+
+</Tip>
+
+**Usage**: 
+
+```bash
+accelerate estimate-memory {MODEL_NAME} --library_name {LIBRARY_NAME} --dtypes {dtype_1} {dtype_2} ...
+```
+
+**Required Arguments**:
+
+* `MODEL_NAME` (`str`)-- The model name on the Hugging Face Hub
+
+**Optional Arguments**:
+
+* `--library_name {timm,transformers}` (`str`) -- The library the model has an integration with, such as `transformers`, needed only if this information is not stored on the Hub
+* `--dtypes {float32,float16,int8,int4}` (`[{float32,float16,int8,int4} ...]`) -- The dtypes to use for the model, must be one (or many) of `float32`, `float16`, `int8`, and `int4`
+* `--trust_remote_code` (`bool`) -- Whether or not to allow for custom models defined on the Hub in their own modeling files. This option should only be passed for repositories you trust and in which you have read the code, as it will execute code present on the Hub on your local machine.
+
 ## accelerate tpu-config
 
 `accelerate tpu-config`

docs/source/package_reference/deepspeed.md

@@ -8,6 +8,9 @@ http://www.apache.org/licenses/LICENSE-2.0
 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
 an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
 specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
 -->
 
 # Utilities for DeepSpeed

docs/source/package_reference/fsdp.md

@@ -0,0 +1,18 @@
+<!--Copyright 2023 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
+-->
+
+# Utilities for Fully Sharded Data Parallelism
+
+[[autodoc]] utils.FullyShardedDataParallelPlugin

docs/source/package_reference/kwargs.md

@@ -8,6 +8,9 @@ http://www.apache.org/licenses/LICENSE-2.0
 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
 an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
 specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
 -->
 
 # Kwargs Handlers
@@ -15,11 +18,18 @@ specific language governing permissions and limitations under the License.
 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.
 
+## AutocastKwargs
+
+[[autodoc]] AutocastKwargs
 
 ## DistributedDataParallelKwargs
 
 [[autodoc]] DistributedDataParallelKwargs
 
+## FP8RecipeKwargs
+
+[[autodoc]] utils.FP8RecipeKwargs
+
 ## GradScalerKwargs
 
 [[autodoc]] GradScalerKwargs

docs/source/package_reference/launchers.md

@@ -8,6 +8,9 @@ http://www.apache.org/licenses/LICENSE-2.0
 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
 an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
 specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
 -->
 
 # Launchers

docs/source/package_reference/logging.md

@@ -8,6 +8,9 @@ http://www.apache.org/licenses/LICENSE-2.0
 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
 an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
 specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
 -->
 
 # Logging with Accelerate

docs/source/package_reference/megatron_lm.md

@@ -8,6 +8,9 @@ http://www.apache.org/licenses/LICENSE-2.0
 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
 an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
 specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
 -->
 
 # Utilities for Megatron-LM

docs/source/package_reference/state.md

@@ -8,6 +8,9 @@ http://www.apache.org/licenses/LICENSE-2.0
 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
 an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
 specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
 -->
 
 # Stateful Classes
@@ -18,6 +21,8 @@ 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.PartialState
+
 [[autodoc]] state.AcceleratorState
 
 [[autodoc]] state.GradientState

docs/source/package_reference/torch_wrappers.md

@@ -8,6 +8,9 @@ http://www.apache.org/licenses/LICENSE-2.0
 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
 an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
 specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
 -->
 
 # Wrapper classes for torch Dataloaders, Optimizers, and Schedulers
@@ -18,6 +21,7 @@ when calling [`~Accelerator.prepare`].
 ## Datasets and DataLoaders
 
 [[autodoc]] data_loader.prepare_data_loader
+[[autodoc]] data_loader.skip_first_batches
 
 [[autodoc]] data_loader.BatchSamplerShard
 [[autodoc]] data_loader.IterableDatasetShard

docs/source/package_reference/tracking.md

@@ -8,6 +8,9 @@ http://www.apache.org/licenses/LICENSE-2.0
 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
 an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
 specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
 -->
 
 # Experiment Tracking
@@ -24,3 +27,7 @@ specific language governing permissions and limitations under the License.
     - __init__
 [[autodoc]] tracking.CometMLTracker
     - __init__
+[[autodoc]] tracking.AimTracker
+    - __init__
+[[autodoc]] tracking.MLflowTracker
+    - __init__

docs/source/package_reference/utilities.md

@@ -8,24 +8,64 @@ http://www.apache.org/licenses/LICENSE-2.0
 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
 an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
 specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
 -->
 
 # Helpful Utilities
 
 Below are a variety of utility functions that 🤗 Accelerate provides, broken down by use-case. 
 
+## Constants
+
+Constants used throughout 🤗 Accelerate for reference
+
+The following are constants used when utilizing [`Accelerator.save_state`]
+
+`utils.MODEL_NAME`: `"pytorch_model"`
+`utils.OPTIMIZER_NAME`: `"optimizer"`
+`utils.RNG_STATE_NAME`: `"random_states"`
+`utils.SCALER_NAME`: `"scaler.pt`
+`utils.SCHEDULER_NAME`: `"scheduler`
+
+The following are constants used when utilizing [`Accelerator.save_model`]
+
+`utils.WEIGHTS_NAME`: `"pytorch_model.bin"`
+`utils.SAFE_WEIGHTS_NAME`: `"model.safetensors"`
+`utils.WEIGHTS_INDEX_NAME`: `"pytorch_model.bin.index.json"`
+`utils.SAFE_WEIGHTS_INDEX_NAME`: `"model.safetensors.index.json"`
+
 ## Data Classes
 
 These are basic dataclasses used throughout 🤗 Accelerate and they can be passed in as parameters.
 
 [[autodoc]] utils.DistributedType
 
+[[autodoc]] utils.DynamoBackend
+
 [[autodoc]] utils.LoggerType
 
 [[autodoc]] utils.PrecisionType
 
 [[autodoc]] utils.ProjectConfiguration
 
+## Plugins
+
+These are plugins that can be passed to the [`Accelerator`] object. While they are defined elsewhere in the documentation, 
+for convience all of them are available to see here:
+
+[[autodoc]] utils.DeepSpeedPlugin
+
+[[autodoc]] utils.FullyShardedDataParallelPlugin
+
+[[autodoc]] utils.GradientAccumulationPlugin
+
+[[autodoc]] utils.MegatronLMPlugin
+
+[[autodoc]] utils.TorchDynamoPlugin
+
+
 ## Data Manipulation and Operations
 
 These include data operations that mimic the same `torch` ops but can be used on distributed processes.
@@ -48,11 +88,23 @@ These functionalities check the state of the current working environment includi
 
 [[autodoc]] utils.is_bf16_available
 
+[[autodoc]] utils.is_ipex_available
+
+[[autodoc]] utils.is_mps_available
+
+[[autodoc]] utils.is_npu_available
+
 [[autodoc]] utils.is_torch_version
 
 [[autodoc]] utils.is_tpu_available
 
-## Environment Configuration
+[[autodoc]] utils.is_xpu_available
+
+## Environment Manipulation
+
+[[autodoc]] utils.patch_environment
+
+[[autodoc]] utils.clear_environment
 
 [[autodoc]] utils.write_basic_config
 
@@ -102,3 +154,17 @@ These utilities relate to setting and synchronizing of all the random states.
 These include utilities that are useful while using PyTorch with XLA.
 
 [[autodoc]] utils.install_xla
+
+## Loading model weights
+
+These include utilities that are useful to load checkpoints.
+
+[[autodoc]] utils.load_checkpoint_in_model
+
+## Quantization
+
+These include utilities that are useful to quantize model.
+
+[[autodoc]] utils.load_and_quantize_model
+
+[[autodoc]] utils.BnbQuantizationConfig

docs/source/quicktour.md

@@ -8,6 +8,9 @@ http://www.apache.org/licenses/LICENSE-2.0
 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
 an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
 specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
 -->
 
 # Quick tour
@@ -67,9 +70,9 @@ use `shuffle=True` or any kind of random sampler).
 
 </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.
+Alternatively, you can use the option `split_batches=True` when creating and initializing your
+[`Accelerator`], in which case the batch size will always stay the same, whether you 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.
@@ -164,9 +167,8 @@ should be calculated through the [`~Accelerator.gather_for_metrics`] method to a
 
 ## 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`.
+You can use the regular commands to launch your distributed training (like `torch.distributed.run` for
+PyTorch), they are fully compatible with 🤗 Accelerate.
 
 🤗 Accelerate also provides a CLI tool that unifies all launchers, so you only have to remember one command. To use it,
 just run:
@@ -346,27 +348,40 @@ point in the script as shown above, and then, you should unwrap your model befor
 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:
+in your base model. The [`~Accelerator.save_model`] method will help you to achieve that. It will unwrap your model and save
+the model state dictionnary.
 
+Here is an example:
 ```
 accelerator.wait_for_everyone()
-unwrapped_model = accelerator.unwrap_model(model)
-accelerator.save(unwrapped_model.state_dict(), filename)
+accelerator.save_model(model, save_directory)
+```
+The [`~Accelerator.save_model`] method can also save a model into sharded checkpoints or with safetensors format.
+Here is an example: 
+
+```python
+accelerator.wait_for_everyone()
+accelerator.save_model(model, save_directory, max_shard_size="1GB", safe_serialization=True)
 ```
 
 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:
 
-```
+```python
 unwrapped_model = accelerator.unwrap_model(model)
-unwrapped_model.load_state_dict(torch.load(filename))
+path_to_checkpoint = os.path.join(save_directory,"pytorch_model.bin")
+unwrapped_model.load_state_dict(torch.load(path_to_checkpoint))
 ```
 
 Note that since all the model parameters are references to tensors, this will load your weights inside `model`.
 
+If you want to load a sharded checkpoint or a checkpoint with safetensors format into the model with a specific `device`, we recommend you to load it with [`~utils.load_checkpoint_in_model`] function. Here's an example:
+
+```python
+load_checkpoint_in_model(unwrapped_model, save_directory, device_map={"":device})
+```
+
 ## 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_.
@@ -396,7 +411,7 @@ and [`~Accelerator.clip_grad_value_`] respectively.
 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:
+softmax). However you might want to put your loss computation inside the [`~Accelerator.autocast`] context manager:
 
 ```
 with accelerator.autocast():

docs/source/usage_guides/big_modeling.md

@@ -0,0 +1,150 @@
+<!--Copyright 2022 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
+-->
+
+# Handling big models for inference
+
+One of the biggest advancements 🤗 Accelerate provides is the concept of [large model inference](../concept_guides/big_model_inference) wherein you can perform *inference* on models that cannot fully fit on your graphics card. 
+
+This tutorial will be broken down into two parts showcasing how to use both 🤗 Accelerate and 🤗 Transformers (a higher API-level) to make use of this idea.
+
+## Using 🤗 Accelerate
+
+For these tutorials, we'll assume a typical workflow for loading your model in such that:
+
+```py
+import torch
+
+my_model = ModelClass(...)
+state_dict = torch.load(checkpoint_file)
+my_model.load_state_dict(state_dict)
+```
+
+Note that here we assume that `ModelClass` is a model that takes up more video-card memory than what can fit on your device (be it `mps` or `cuda`).
+
+The first step is to init an empty skeleton of the model which won't take up any RAM using the [`init_empty_weights`] context manager:
+
+```py
+from accelerate import init_empty_weights
+with init_empty_weights():
+    my_model = ModelClass(...)
+```
+
+With this `my_model` currently is "parameterless", hence leaving the smaller footprint than what one would normally get loading this onto the CPU directly. 
+
+Next we need to load in the weights to our model so we can perform inference.
+
+For this we will use [`load_checkpoint_and_dispatch`], which as the name implies will load a checkpoint inside your empty model and dispatch the weights for each layer across all the devices you have available (GPU/MPS and CPU RAM). 
+
+To determine how this `dispatch` can be performed, generally specifying `device_map="auto"` will be good enough as 🤗 Accelerate
+will attempt to fill all the space in your GPU(s), then loading them to the CPU, and finally if there is not enough RAM it will be loaded to the disk (the absolute slowest option). 
+
+<Tip>
+
+For more details on desigining your own device map, see this section of the [concept guide](../concept_guide/big_model_inference#desigining-a-device-map)
+
+</Tip>
+
+See an example below:
+
+```py
+from accelerate import load_checkpoint_and_dispatch
+
+model = load_checkpoint_and_dispatch(
+    model, checkpoint=checkpoint_file, device_map="auto"
+)
+```
+
+<Tip>
+
+    If there are certain "chunks" of layers that shouldn't be split, you can pass them in as `no_split_module_classes`. Read more about it [here](../concept_guides/big_model_inference#loading-weights)
+
+</Tip>
+
+<Tip>
+
+    Also to save on memory (such as if the `state_dict` will not fit in RAM), a model's weights can be divided and split into multiple checkpoint files. Read more about it [here](../concept_guides/big_model_inference#sharded-checkpoints)
+
+</Tip>
+
+Now that the model is dispatched fully, you can perform inference as normal with the model:
+
+```py
+input = torch.randn(2,3)
+input = input.to("cuda")
+output = model(input)
+```
+
+What will happen now is each time the input gets passed through a layer, it will be sent from the CPU to the GPU (or disk to CPU to GPU), the output is calculated, and then the layer is pulled back off the GPU going back down the line. While this adds some overhead to the inference being performed, through this method it is possible to run **any size model** on your system, as long as the largest layer is capable of fitting on your GPU. 
+
+<Tip>
+
+    Multiple GPUs can be utilized, however this is considered "model parallism" and as a result only one GPU will be active at a given moment, waiting for the prior one to send it the output. You should launch your script normally with `python`
+    and not need `torchrun`, `accelerate launch`, etc.
+
+</Tip>
+
+For a visual representation of this, check out the animation below:
+
+<Youtube id="MWCSGj9jEAo" />
+
+### Complete Example
+
+Below is the full example showcasing what we performed above:
+
+```py
+import torch
+from accelerate import init_empty_weights, load_checkpoint_and_dispatch
+
+with init_empty_weights():
+    model = MyModel(...)
+
+model = load_checkpoint_and_dispatch(
+    model, checkpoint=checkpoint_file, device_map="auto"
+)
+
+input = torch.randn(2,3)
+input = input.to("cuda")
+output = model(input)
+```
+
+## Using 🤗 Transformers, 🤗 Diffusers, and other 🤗 Open Source Libraries
+
+Libraries that support 🤗 Accelerate big model inference include all of the earlier logic in their `from_pretrained` constructors. 
+
+These operate by specifying a string representing the model to download from the [🤗 Hub](https://hf.co/models) and then denoting `device_map="auto"` along with a few extra parameters. 
+
+As a brief example, we will look at using `transformers` and loading in Big Science's T0pp model. 
+
+```py
+from transformers import AutoModelForSeq2SeqLM
+
+model = AutoModelForSeq2SeqLM("bigscience/T0pp", device_map="auto")
+```
+
+After loading the model in, the initial steps from before to prepare a model have all been done and the model is fully
+ready to make use of all the resources in your machine. Through these constructors, you can also save *more* memory by
+specifying the precision the model is loaded into as well, through the `torch_dtype` parameter, such as:
+
+```py
+from transformers import AutoModelForSeq2SeqLM
+
+model = AutoModelForSeq2SeqLM("bigscience/T0pp", device_map="auto", torch_dtype=torch.float16)
+```
+
+To learn more about this, check out the 🤗 Transformers documentation available [here](https://huggingface.co/docs/transformers/main/en/main_classes/model#large-model-loading).
+
+## Where to go from here
+
+For a much more detailed look at big model inference, be sure to check out the [Conceptual Guide on it](../concept_guides/big_model_inference)

docs/source/usage_guides/checkpoint.md

@@ -8,6 +8,9 @@ http://www.apache.org/licenses/LICENSE-2.0
 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
 an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
 specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
 -->
 
 # Checkpointing
@@ -17,7 +20,7 @@ saving and loading the model, optimizer, RNG generators, and the GradScaler. Ins
 - 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`
 
-To further customize where and how states saved through [`~Accelerator.save_state`] the [`~utils.ProjectConfiguration`] class can be used. For example 
+To further customize where and how states are saved through [`~Accelerator.save_state`] the [`~utils.ProjectConfiguration`] class can be used. For example 
 if `automatic_checkpoint_naming` is enabled each saved checkpoint will be located then at `Accelerator.project_dir/checkpoints/checkpoint_{checkpoint_number}`.
 
 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.
@@ -25,6 +28,7 @@ It should be noted that the expectation is that those states come from the same
 - 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
@@ -58,6 +62,35 @@ for epoch in range(num_epochs):
         my_optimizer.step()
     my_scheduler.step()
 
-# Restore previous state
+# Restore the previous state
 accelerator.load_state("my/save/path/checkpointing/checkpoint_0")
 ```
+
+## Restoring the state of the DataLoader 
+
+After resuming from a checkpoint, it may also be desirable to resume from a particular point in the active `DataLoader` if 
+the state was saved during the middle of an epoch. You can use [`~Accelerator.skip_first_batches`] to do so. 
+
+```python
+from accelerate import Accelerator
+
+accelerator = Accelerator(project_dir="my/save/path")
+
+train_dataloader = accelerator.prepare(train_dataloader)
+accelerator.load_state("my_state")
+
+# Assume the checkpoint was saved 100 steps into the epoch
+skipped_dataloader = accelerator.skip_first_batches(train_dataloader, 100)
+
+# After the first iteration, go back to `train_dataloader`
+
+# First epoch
+for batch in skipped_dataloader:
+    # Do something
+    pass
+
+# Second epoch
+for batch in train_dataloader:
+    # Do something
+    pass
+```

docs/source/usage_guides/debug.md

@@ -0,0 +1,93 @@
+<!--Copyright 2022 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
+-->
+
+# Debugging Distributed Operations
+
+When running scripts in a distributed fashion, often functions such as [`Accelerator.gather`] and [`Accelerator.reduce`] (and others) are neccessary to grab tensors across devices and perform certain operations on them. However, if the tensors which are being grabbed are not the proper shapes then this will result in your code hanging forever. The only sign that exists of this truly happening is hitting a timeout exception from `torch.distributed`, but this can get quite costly as usually the timeout is 10 minutes.
+
+Accelerate now has a `debug` mode which adds a neglible amount of time to each operation, but allows it to verify that the inputs you are bringing in can *actually* perform the operation you want **without** hitting this timeout problem!
+
+## Visualizing the problem
+
+To have a tangible example of this issue, let's take the following setup (on 2 GPUs):
+
+```python
+from accelerate import PartialState
+
+state = PartialState()
+if state.process_index == 0:
+    tensor = torch.tensor([[0.0, 1, 2, 3, 4]]).to(state.device)
+else:
+    tensor = torch.tensor([[[0.0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]]).to(state.device)
+
+broadcast_tensor = broadcast(tensor)
+print(broadcast_tensor)
+```
+
+We've created a single tensor on each device, with two radically different shapes. With this setup if we want to perform an operation such as [`utils.broadcast`], we would forever hit a timeout because `torch.distributed` requires that these operations have the **exact same shape** across all processes for it to work.
+
+If you run this yourself, you will find that `broadcast_tensor` can be printed on the main process, but its results won't quite be right, and then it will just hang never printing it on any of the other processes:
+
+```
+>>> tensor([[0, 1, 2, 3, 4]], device='cuda:0')
+```
+
+## The solution
+
+By enabling Accelerate's operational debug mode, Accelerate will properly find and catch errors such as this and provide a very clear traceback immediatly: 
+
+```
+Traceback (most recent call last):
+  File "/home/zach_mueller_huggingface_co/test.py", line 18, in <module>
+    main()
+  File "/home/zach_mueller_huggingface_co/test.py", line 15, in main
+        main()broadcast_tensor = broadcast(tensor)
+  File "/home/zach_mueller_huggingface_co/accelerate/src/accelerate/utils/operations.py", line 303, in wrapper
+    broadcast_tensor = broadcast(tensor)
+accelerate.utils.operations.DistributedOperationException: Cannot apply desired operation due to shape mismatches. All shapes across devices must be valid.
+
+Operation: `accelerate.utils.operations.broadcast`
+Input shapes:
+  - Process 0: [1, 5]
+  - Process 1: [1, 2, 5]
+```
+
+This explains that the shapes across our devices were *not* the same, and that we should ensure that they match properly to be compatible. Typically this means that there is either an extra dimension, or certain dimensions are incompatible with the operation.
+
+To enable this please do one of the following:
+
+Enable it through the questionarre during `accelerate config` (recommended)
+
+From the CLI: 
+
+```
+accelerate launch --debug {my_script.py} --arg1 --arg2
+```
+
+As an environmental variable (which avoids the need for `accelerate launch`):
+
+```
+ACCELERATE_DEBUG_MODE="1" accelerate launch {my_script.py} --arg1 --arg2
+```
+
+Manually changing the `config.yaml` file:
+
+```diff
+ compute_environment: LOCAL_MACHINE
++debug: true
+```
+
+
+

docs/source/usage_guides/deepspeed.md

@@ -8,11 +8,14 @@ http://www.apache.org/licenses/LICENSE-2.0
 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
 an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
 specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
 -->
 
 # 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:
+[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)
@@ -26,14 +29,14 @@ 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
+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.
+   User may have to change a 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.
 
@@ -57,7 +60,7 @@ Below is a short description of Data Parallelism using ZeRO - Zero Redundancy Op
 
  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
+<u>Note</u>: With respect to Disk Offload, the disk should be an NVME for decent speed but it technically works on any Disk
 
 Inference:
 
@@ -352,8 +355,8 @@ accelerate launch examples/by_feature/deepspeed_with_config_support.py \
 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.
+   a. DS Optim + DS Scheduler: The case when both `optimizer` and `scheduler` keys are present in the DeepSpeed config file.
+   In this situation, those will be used and the 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
@@ -386,7 +389,7 @@ We will look at the changes needed in the code when using these.
    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. 
+   In this situation, the 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.
@@ -582,8 +585,10 @@ Mixed precision type: fp16
 ds_config: {'bf16': {'enabled': False}, 'zero_optimization': {'stage': 3, 'stage3_gather_16bit_weights_on_model_save': True, 'offload_optimizer': {'device': 'nvme'}, 'offload_param': {'device': 'cpu'}}, 'gradient_clipping': 1.0, 'train_batch_size': 'auto', 'train_micro_batch_size_per_gpu': 'auto', 'gradient_accumulation_steps': 5, 'steps_per_print': inf, 'fp16': {'enabled': True, 'auto_cast': True}}
 ```
 
-**Note**: Remaining `"auto"` values are handled in `accelerator.prepare()` call as explained in point 2 of 
+**Note**: 
+1. Remaining `"auto"` values are handled in `accelerator.prepare()` call as explained in point 2 of 
 `Important code changes when using DeepSpeed Config File`.
+2. Only when `gradient_accumulation_steps` is `auto`, the value passed while creating `Accelerator` object via `Accelerator(gradient_accumulation_steps=k)` will be used. When using DeepSpeed Plugin, the value from it will be used and it will overwrite the value passed while creating Accelerator object. 
 
 ## Saving and loading
 
@@ -679,6 +684,6 @@ Papers:
 - [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
+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/distributed_inference.md

@@ -0,0 +1,136 @@
+<!--Copyright 2023 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
+-->
+
+# Distributed Inference with 🤗 Accelerate
+
+Distributed inference is a common use case, especially with natural language processing (NLP) models. Users often want to
+send a number of different prompts, each to a different GPU, and then get the results back. This also has other cases
+outside of just NLP, however for this tutorial we will focus on just this idea of each GPU receiving a different prompt,
+and then returning the results.
+
+## The Problem
+
+Normally when doing this, users send the model to a specific device to load it from the CPU, and then move each prompt to a different device. 
+
+A basic pipeline using the `diffusers` library might look something like so:
+
+```python
+import torch
+import torch.distributed as dist
+from diffusers import DiffusionPipeline
+
+pipe = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
+```
+Followed then by performing inference based on the specific prompt:
+
+```python
+def run_inference(rank, world_size):
+    dist.init_process_group("nccl", rank=rank, world_size=world_size)
+    pipe.to(rank)
+
+    if torch.distributed.get_rank() == 0:
+        prompt = "a dog"
+    elif torch.distributed.get_rank() == 1:
+        prompt = "a cat"
+
+    result = pipe(prompt).images[0]
+    result.save(f"result_{rank}.png")
+```
+One will notice how we have to check the rank to know what prompt to send, which can be a bit tedious.
+
+A user might then also think that with 🤗 Accelerate, using the `Accelerator` to prepare a dataloader for such a task might also be 
+a simple way to manage this. (To learn more, check out the relvent section in the [Quick Tour](../quicktour#distributed-evaluation))
+
+Can it manage it? Yes. Does it add unneeded extra code however: also yes.
+
+## The Solution
+
+With 🤗 Accelerate, we can simplify this process by using the [`Accelerator.split_between_processes`] context manager (which also exists in `PartialState` and `AcceleratorState`). 
+This function will automatically split whatever data you pass to it (be it a prompt, a set of tensors, a dictionary of the prior data, etc.) across all the processes (with a potential
+to be padded) for you to use right away.
+
+Let's rewrite the above example using this context manager:
+
+```python
+from accelerate import PartialState  # Can also be Accelerator or AcceleratorState
+from diffusers import DiffusionPipeline
+
+pipe = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
+distributed_state = PartialState()
+pipe.to(distributed_state.device)
+
+# Assume two processes
+with distributed_state.split_between_processes(["a dog", "a cat"]) as prompt:
+    result = pipe(prompt).images[0]
+    result.save(f"result_{distributed_state.process_index}.png")
+```
+
+And then to launch the code, we can use the 🤗 Accelerate:
+
+If you have generated a config file to be used using `accelerate config`:
+
+```bash
+accelerate launch distributed_inference.py
+```
+
+If you have a specific config file you want to use:
+
+```bash
+accelerate launch --config_file my_config.json distributed_inference.py
+```
+
+Or if don't want to make any config files and launch on two GPUs:
+
+> Note: You will get some warnings about values being guessed based on your system. To remove these you can do `accelerate config default` or go through `accelerate config` to create a config file.
+
+```bash
+accelerate launch --num_processes 2 distributed_inference.py
+```
+
+We've now reduced the boilerplate code needed to split this data to a few lines of code quite easily.
+
+But what if we have an odd distribution of prompts to GPUs? For example, what if we have 3 prompts, but only 2 GPUs? 
+
+Under the context manager, the first GPU would receive the first two prompts and the second GPU the third, ensuring that 
+all prompts are split and no overhead is needed.
+
+*However*, what if we then wanted to do something with the results of *all the GPUs*? (Say gather them all and perform some kind of post processing)
+You can pass in `apply_padding=True` to ensure that the lists of prompts are padded to the same length, with extra data being taken 
+from the last sample. This way all GPUs will have the same number of prompts, and you can then gather the results.
+
+<Tip>
+
+This is only needed when trying to perform an action such as gathering the results, where the data on each device 
+needs to be the same length. Basic inference does not require this.
+
+</Tip>
+
+For instance:
+
+```python
+from accelerate import PartialState  # Can also be Accelerator or AcceleratorState
+from diffusers import DiffusionPipeline
+
+pipe = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
+distributed_state = PartialState()
+pipe.to(distributed_state.device)
+
+# Assume two processes
+with distributed_state.split_between_processes(["a dog", "a cat", "a chicken"], apply_padding=True) as prompt:
+    result = pipe(prompt).images
+```
+
+On the first GPU, the prompts will be `["a dog", "a cat"]`, and on the second GPU it will be `["a chicken", "a chicken"]`.
+Make sure to drop the final sample, as it will be a duplicate of the previous one.

docs/source/usage_guides/explore.md

@@ -8,6 +8,9 @@ http://www.apache.org/licenses/LICENSE-2.0
 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
 an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
 specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
 -->
 
 # Learning how to incorporate 🤗 Accelerate features quickly!
@@ -34,14 +37,14 @@ for batch in dataloader:
 
 <div class="block dark:hidden">
 	<iframe 
-        src="https://muellerzr-accelerate-examples.hf.space?__theme=light"
+        src="https://hf-accelerate-accelerate-examples.hf.space?__theme=light"
         width="850"
         height="1600"
     ></iframe>
 </div>
 <div class="hidden dark:block">
     <iframe 
-        src="https://muellerzr-accelerate-examples.hf.space?__theme=dark"
+        src="https://hf-accelerate-accelerate-examples.hf.space?__theme=dark"
         width="850"
         height="1600"
     ></iframe>

docs/source/usage_guides/fsdp.md

@@ -8,6 +8,9 @@ http://www.apache.org/licenses/LICENSE-2.0
 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
 an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
 specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
 -->
 
 # Fully Sharded Data Parallel
@@ -46,7 +49,7 @@ fsdp_config:
   fsdp_offload_params: false
   fsdp_sharding_strategy: 1
   fsdp_state_dict_type: FULL_STATE_DICT
-  fsdp_transformer_layer_cls_to_wrap: GPT2Block
+  fsdp_transformer_layer_cls_to_wrap: BertLayer
 machine_rank: 0
 main_process_ip: null
 main_process_port: null
@@ -64,15 +67,113 @@ 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
+`Sharding Strategy`: [1] FULL_SHARD (shards optimizer states, gradients and parameters), [2] SHARD_GRAD_OP (shards optimizer states and gradients), [3] NO_SHARD (DDP), [4] HYBRID_SHARD (shards optimizer states, gradients and parameters within each node while each node has full copy), [5] HYBRID_SHARD_ZERO2 (shards optimizer states and gradients within each node while each node has full copy)
+
 `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`...
+
+`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 comma-separated string of transformer layer class names (case-sensitive) to wrap ,e.g, 
+`BertLayer`, `GPTJBlock`, `T5Block`, `BertLayer,BertEmbeddings,BertSelfOutput`...
+This is important because submodules that share weights (e.g., embedding layer) should not end up in different FSDP wrapped units.
+Using this policy, wrapping happens for each block containing Multi-Head Attention followed by couple of MLP layers. 
+Remaining layers including the shared embeddings are conveniently wrapped in same outermost FSDP unit.
+Therefore, use this for transformer based models.
+You can use the `model._no_split_modules` for 🤗 Transformer models by answering `yes` to 
+`Do you want to use the model's `_no_split_modules` to wrap. Only applicable for 🤗 Transformers`. 
+It will try to use `model._no_split_modules` when available.  
+
 `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  
+
+`State Dict Type`: [1] FULL_STATE_DICT, [2] LOCAL_STATE_DICT, [3] SHARDED_STATE_DICT 
+
+`Forward Prefetch`: if True, then FSDP explicitly prefetches the next upcoming
+all-gather while executing in the forward pass. only use with Static graphs.
+
+`Use Orig Params`: If True, allows non-uniform `requires_grad` during init, which means support for interspersed frozen and trainable paramteres. 
+Useful in cases such as parameter-efficient fine-tuning. 
+Please refer this [blog](https://dev-discuss.pytorch.org/t/rethinking-pytorch-fully-sharded-data-parallel-fsdp-from-first-principles/1019)
+
+`Sync Module States`: If True, each individually wrapped FSDP unit will broadcast module parameters from rank 0
 ```
 
+For additional and more nuanced control, you can specify other FSDP parameters via `FullyShardedDataParallelPlugin`. 
+When creating `FullyShardedDataParallelPlugin` object, pass it the parameters that weren't part of the accelerate config or if you want to override them.
+The FSDP parameters will be picked based on the accelerate config file or launch command arguments and other parameters that you will pass directly through the `FullyShardedDataParallelPlugin` object will set/override that.
+
+Below is an example:
+
+```py
+from accelerate import FullyShardedDataParallelPlugin
+from torch.distributed.fsdp.fully_sharded_data_parallel import FullOptimStateDictConfig, FullStateDictConfig
+
+fsdp_plugin = FullyShardedDataParallelPlugin(
+    state_dict_config=FullStateDictConfig(offload_to_cpu=False, rank0_only=False),
+    optim_state_dict_config=FullOptimStateDictConfig(offload_to_cpu=False, rank0_only=False),
+)
+
+accelerator = Accelerator(fsdp_plugin=fsdp_plugin)
+```
+
+## Saving and loading
+
+The new recommended way of checkpointing when using FSDP models is to use `SHARDED_STATE_DICT` as `StateDictType` when setting up the accelerate config.
+Below is the code snippet to save using `save_state` utility of accelerate.
+
+```py
+accelerator.save_state("ckpt")
+```
+
+Inspect the ckeckpoint folder to see model and optimizer as shards per process:
+```
+ls ckpt 
+# optimizer_0  pytorch_model_0  random_states_0.pkl  random_states_1.pkl  scheduler.bin
+
+cd ckpt
+
+ls optimizer_0
+# __0_0.distcp  __1_0.distcp
+
+ls pytorch_model_0
+# __0_0.distcp  __1_0.distcp
+```
+
+To load them back for resuming the training, use the `load_state` utility of accelerate
+
+```py
+accelerator.load_state("ckpt")
+```
+
+When using transformers `save_pretrained`, pass `state_dict=accelerator.get_state_dict(model)` to save the model state dict. 
+  Below is an example:
+
+```diff
+  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),
+)
+```
+
+### State Dict
+
+`accelerator.get_state_dict` will call the underlying `model.state_dict` implementation.  With a model wrapped by FSDP, the default behavior of `state_dict` is to gather all of the state in the rank 0 device.  This can cause CUDA out of memory errors if the parameters don't fit on a single GPU.
+
+To avoid this, PyTorch provides a context manager that adjusts the behavior of `state_dict`.  To offload some of the state dict onto CPU, you can use the following code:
+
+```
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP, StateDictType, FullStateDictConfig
+
+full_state_dict_config = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
+with FSDP.state_dict_type(unwrapped_model, StateDictType.FULL_STATE_DICT, full_state_dict_config):
+    state = accelerator.get_state_dict(unwrapped_model)
+```
+
+You can then pass `state` into the `save_pretrained` method.  There are several modes for `StateDictType` and `FullStateDictConfig` that you can use to control the behavior of `state_dict`.  For more information, see the [PyTorch documentation](https://pytorch.org/docs/stable/fsdp.html).
+
 ## A few caveats to be aware of
 
 - PyTorch FSDP auto wraps sub-modules, flattens the parameters and shards the parameters in place.

docs/source/usage_guides/gradient_accumulation.md

@@ -8,6 +8,9 @@ http://www.apache.org/licenses/LICENSE-2.0
 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
 an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
 specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
 -->
 
 # Performing gradient accumulation with 🤗 Accelerate
@@ -72,7 +75,7 @@ First the code shown earlier will be converted to utilize 🤗 Accelerate withou
 
 <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)!
+  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>
 
@@ -87,6 +90,9 @@ of steps to perform before each call to `step()` and how to automatically adjust
 + accelerator = Accelerator(gradient_accumulation_steps=2)
 ```
 
+Alternatively, you can pass in a `gradient_accumulation_plugin` parameter to the [`Accelerator`] object's `__init__`, which will allow you to further customize the gradient accumulation behavior. 
+Read more about that in the [GradientAccumulationPlugin](../package_reference/accelerator#accelerate.utils.GradientAccumulationPlugin) docs.
+
 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: 
 
@@ -111,11 +117,21 @@ You can remove all the special checks for the step number and the loss adjustmen
 
 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. 
 
+<Tip>
+Typically with gradient accumulation, you would need to adjust the number of steps to reflect the change in total batches you are 
+training on. 🤗 Accelerate automagically does this for you by default. Behind the scenes we instantiate a GradientAccumulationPlugin configured to do this.
+</Tip>
+
 ## The finished code
 
 Below is the finished implementation for performing gradient accumulation with 🤗 Accelerate
 
 ```python
+from accelerate import Accelerator
+accelerator = Accelerator(gradient_accumulation_steps=2)
+model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+    model, optimizer, training_dataloader, scheduler
+)
 for batch in training_dataloader:
     with accelerator.accumulate(model):
         inputs, targets = batch
@@ -127,4 +143,74 @@ for batch in training_dataloader:
         optimizer.zero_grad()
 ```
 
-To learn more about what magic this wraps around, read the [Gradient Synchronization concept guide](/concept_guides/gradient_synchronization)
+<Tip warning={true}>
+
+It's important that **only one forward/backward** should be done inside the context manager `with accelerator.accumulate(model)`.
+
+</Tip>
+
+
+To learn more about what magic this wraps around, read the [Gradient Synchronization concept guide](../concept_guides/gradient_synchronization)
+
+
+## Self-contained example
+
+Here is a self-contained example that you can run to see gradient accumulation in action with 🤗 Accelerate:
+
+```python
+import torch
+import copy
+from accelerate import Accelerator
+from accelerate.utils import set_seed
+from torch.utils.data import TensorDataset, DataLoader
+
+# seed
+set_seed(0)
+
+# define toy inputs and labels
+x = torch.tensor([1., 2., 3., 4., 5., 6., 7., 8.])
+y = torch.tensor([2., 4., 6., 8., 10., 12., 14., 16.])
+gradient_accumulation_steps = 4
+batch_size = len(x) // gradient_accumulation_steps
+
+# define dataset and dataloader
+dataset = TensorDataset(x, y)
+dataloader = DataLoader(dataset, batch_size=batch_size)
+
+# define model, optimizer and loss function
+model = torch.zeros((1, 1), requires_grad=True)
+model_clone = copy.deepcopy(model)
+criterion = torch.nn.MSELoss()
+model_optimizer = torch.optim.SGD([model], lr=0.02)
+accelerator = Accelerator(gradient_accumulation_steps=gradient_accumulation_steps)
+model, model_optimizer, dataloader = accelerator.prepare(model, model_optimizer, dataloader)
+model_clone_optimizer = torch.optim.SGD([model_clone], lr=0.02)
+print(f"initial model weight is {model.mean().item():.5f}")
+print(f"initial model weight is {model_clone.mean().item():.5f}")
+for i, (inputs, labels) in enumerate(dataloader):
+    with accelerator.accumulate(model):
+        inputs = inputs.view(-1, 1)
+        print(i, inputs.flatten())
+        labels = labels.view(-1, 1)
+        outputs = inputs @ model
+        loss = criterion(outputs, labels)
+        accelerator.backward(loss)
+        model_optimizer.step()
+        model_optimizer.zero_grad()
+loss = criterion(x.view(-1, 1) @ model_clone, y.view(-1, 1))
+model_clone_optimizer.zero_grad()
+loss.backward()
+model_clone_optimizer.step()
+print(f"w/ accumulation, the final model weight is {model.mean().item():.5f}")
+print(f"w/o accumulation, the final model weight is {model_clone.mean().item():.5f}")
+```
+```
+initial model weight is 0.00000
+initial model weight is 0.00000
+0 tensor([1., 2.])
+1 tensor([3., 4.])
+2 tensor([5., 6.])
+3 tensor([7., 8.])
+w/ accumulation, the final model weight is 2.04000
+w/o accumulation, the final model weight is 2.04000
+```

docs/source/usage_guides/ipex.md

@@ -0,0 +1,174 @@
+<!--Copyright 2022 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
+-->
+
+# Intel® Extension for PyTorch
+
+[IPEX](https://github.com/intel/intel-extension-for-pytorch) is optimized for CPUs with AVX-512 or above, and functionally works for CPUs with only AVX2. So, it is expected to bring performance benefit for Intel CPU generations with AVX-512 or above while CPUs with only AVX2 (e.g., AMD CPUs or older Intel CPUs) might result in a better performance under IPEX, but not guaranteed. IPEX provides performance optimizations for CPU training with both Float32 and BFloat16. The usage of BFloat16 is the main focus of the following sections.
+
+Low precision data type BFloat16 has been natively supported on the 3rd Generation Xeon® Scalable Processors (aka Cooper Lake) with AVX512 instruction set and will be supported on the next generation of Intel® Xeon® Scalable Processors with Intel® Advanced Matrix Extensions (Intel® AMX) instruction set with further boosted performance. The Auto Mixed Precision for CPU backend has been enabled since PyTorch-1.10. At the same time, the support of Auto Mixed Precision with BFloat16 for CPU and BFloat16 optimization of operators has been massively enabled in Intel® Extension for PyTorch, and partially upstreamed to PyTorch master branch. Users can get better performance and user experience with IPEX Auto Mixed Precision.
+
+## IPEX installation:
+
+IPEX release is following PyTorch, to install via pip:
+
+| PyTorch Version   | IPEX version   |
+| :---------------: | :----------:   |
+| 2.0               |  2.0.0         |
+| 1.13              |  1.13.0        |
+| 1.12              |  1.12.300      |
+| 1.11              |  1.11.200      |
+| 1.10              |  1.10.100      |
+
+```
+pip install intel_extension_for_pytorch==<version_name> -f https://developer.intel.com/ipex-whl-stable-cpu
+```
+
+Check more approaches for [IPEX installation](https://intel.github.io/intel-extension-for-pytorch/cpu/latest/tutorials/installation.html).
+
+
+## How It Works For Training optimization in CPU
+
+🤗 Accelerate has integrated [IPEX](https://github.com/intel/intel-extension-for-pytorch), all you need to do is enabling it through the config.
+
+**Scenario 1**: Acceleration of No distributed CPU training
+
+Run <u>accelerate config</u> on your machine:
+
+```bash
+$ accelerate config
+-----------------------------------------------------------------------------------------------------------------------------------------------------------
+In which compute environment are you running?
+This machine
+-----------------------------------------------------------------------------------------------------------------------------------------------------------
+Which type of machine are you using?
+No distributed training
+Do you want to run your training on CPU only (even if a GPU / Apple Silicon device is available)? [yes/NO]:yes
+Do you want to use Intel PyTorch Extension (IPEX) to speed up training on CPU? [yes/NO]:yes
+Do you wish to optimize your script with torch dynamo?[yes/NO]:NO
+Do you want to use DeepSpeed? [yes/NO]: NO
+-----------------------------------------------------------------------------------------------------------------------------------------------------------
+Do you wish to use FP16 or BF16 (mixed precision)?
+bf16
+```
+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 IPEX enabled.
+default_config.yaml that is generated after `accelerate config`
+
+```bash
+compute_environment: LOCAL_MACHINE
+distributed_type: 'NO'
+downcast_bf16: 'no'
+ipex_config:
+  ipex: true
+machine_rank: 0
+main_training_function: main
+mixed_precision: bf16
+num_machines: 1
+num_processes: 1
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: true
+```
+```bash
+accelerate launch examples/nlp_example.py
+```
+
+**Scenario 2**: Acceleration of distributed CPU training
+we use Intel oneCCL for communication, combined with Intel® MPI library to deliver flexible, efficient, scalable cluster messaging on Intel® architecture. you could refer the [here](https://huggingface.co/docs/transformers/perf_train_cpu_many) for the installation guide
+
+Run <u>accelerate config</u> on your machine(node0):
+
+```bash
+$ accelerate config
+-----------------------------------------------------------------------------------------------------------------------------------------------------------
+In which compute environment are you running?
+This machine
+-----------------------------------------------------------------------------------------------------------------------------------------------------------
+Which type of machine are you using?
+multi-CPU
+How many different machines will you use (use more than 1 for multi-node training)? [1]: 4
+-----------------------------------------------------------------------------------------------------------------------------------------------------------
+What is the rank of this machine?
+0
+What is the IP address of the machine that will host the main process? 36.112.23.24
+What is the port you will use to communicate with the main process? 29500
+Are all the machines on the same local network? Answer `no` if nodes are on the cloud and/or on different network hosts [YES/no]: yes
+Do you want to use Intel PyTorch Extension (IPEX) to speed up training on CPU? [yes/NO]:yes
+Do you wish to optimize your script with torch dynamo?[yes/NO]:NO
+How many CPU(s) should be used for distributed training? [1]:16
+-----------------------------------------------------------------------------------------------------------------------------------------------------------
+Do you wish to use FP16 or BF16 (mixed precision)?
+bf16
+```
+For instance, here is how you would run the NLP example `examples/nlp_example.py` (from the root of the repo) with IPEX enabled for distributed CPU training.
+
+default_config.yaml that is generated after `accelerate config`
+```bash
+compute_environment: LOCAL_MACHINE
+distributed_type: MULTI_CPU
+downcast_bf16: 'no'
+ipex_config:
+  ipex: true
+machine_rank: 0
+main_process_ip: 36.112.23.24
+main_process_port: 29500
+main_training_function: main
+mixed_precision: bf16
+num_machines: 4
+num_processes: 16
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: true
+```
+
+Set following env and using intel MPI to launch the training
+
+In node0, you need to create a configuration file which contains the IP addresses of each node (for example hostfile) and pass that configuration file path as an argument.
+```bash
+$ cat hostfile
+xxx.xxx.xxx.xxx #node0 ip
+xxx.xxx.xxx.xxx #node1 ip
+xxx.xxx.xxx.xxx #node2 ip
+xxx.xxx.xxx.xxx #node3 ip
+```
+Now, run the following command in node0 and **16DDP** will be enabled in node0,node1,node2,node3 with BF16 mixed precision:
+```bash
+oneccl_bindings_for_pytorch_path=$(python -c "from oneccl_bindings_for_pytorch import cwd; print(cwd)")
+source $oneccl_bindings_for_pytorch_path/env/setvars.sh
+export CCL_WORKER_COUNT=1
+export MASTER_ADDR=xxx.xxx.xxx.xxx #node0 ip
+export CCL_ATL_TRANSPORT=ofi
+mpirun -f hostfile -n 16 -ppn 4 accelerate launch examples/nlp_example.py
+```
+
+## Related Resources
+
+- [Project's github](https://github.com/intel/intel-extension-for-pytorch)
+- [API docs](https://intel.github.io/intel-extension-for-pytorch/cpu/latest/tutorials/api_doc.html)
+- [Tuning guide](https://intel.github.io/intel-extension-for-pytorch/cpu/latest/tutorials/performance_tuning/tuning_guide.html)
+- [Blogs & Publications](https://intel.github.io/intel-extension-for-pytorch/cpu/latest/tutorials/blogs_publications.html)
+

docs/source/usage_guides/local_sgd.md

@@ -0,0 +1,108 @@
+<!--Copyright 2023 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
+-->
+
+# Using Local SGD with 🤗 Accelerate
+
+Local SGD is a technique for distributed training where gradients are not synchronized every step. Thus, each process updates its own version of the model weights and after a given number of steps these weights are synchronized by averaging across all processes. This improves communication efficiency and can lead to substantial training speed up especially when a computer lacks a faster interconnect such as NVLink.
+Unlike gradient accumulation (where improving communication efficiency requires increasing the effective batch size), Local SGD does not require changing a batch size or a learning rate / schedule. However, if necessary, Local SGD can be combined with gradient accumulation as well.
+
+In this tutorial you will see how to quickly setup  Local SGD 🤗 Accelerate. Compared to a standard Accelerate setup, this requires only two extra lines 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 use 🤗 Accelerate  with neither a LocalSGD or a 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()
+```
+
+## Letting 🤗 Accelerate handle model synchronization 
+
+All that is left now is to let 🤗 Accelerate handle model parameter synchronization **and** the gradient accumulation for us. For simplicity let us assume we need to synchronize every 8 steps. This is
+achieved by adding one `with LocalSGD` statement and one call `local_sgd.step()` after every optimizer step:
+
+```diff
++local_sgd_steps=8
+
++with LocalSGD(accelerator=accelerator, model=model, local_sgd_steps=8, enabled=True) as local_sgd:
+    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()
++           local_sgd.step()
+```
+
+Under the hood, the Local SGD code **disables** automatic gradient synchornization (but accumulation still works as expected!). Instead it averages model parameters every `local_sgd_steps` steps (as well as in the end of the training loop).
+
+## Limitations
+
+The current implementation works only with basic multi-GPU (or multi-CPU) training without, e.g., [DeepSpeed.](https://github.com/microsoft/DeepSpeed).
+
+## References
+
+    Although we are not aware of the true origins of this simple approach, the idea of local SGD is quite old and goes
+    back to at least:
+
+    Zhang, J., De Sa, C., Mitliagkas, I., & Ré, C. (2016). [Parallel SGD: When does averaging help?. arXiv preprint
+    arXiv:1606.07365.](https://arxiv.org/abs/1606.07365)
+
+    We credit the term Local SGD to the following paper (but there might be earlier references we are not aware of).
+
+    Stich, Sebastian Urban. ["Local SGD Converges Fast and Communicates Little." ICLR 2019-International Conference on
+    Learning Representations. No. CONF. 2019.](https://arxiv.org/abs/1805.09767)

docs/source/usage_guides/megatron_lm.md

@@ -8,6 +8,9 @@ http://www.apache.org/licenses/LICENSE-2.0
 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
 an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
 specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
 -->
 
 
@@ -115,7 +118,7 @@ An example of thr corresponding questions for using Megatron-LM features is show
 ```bash
 :~$ accelerate config --config_file "megatron_gpt_config.yaml"
 In which compute environment are you running? ([0] This machine, [1] AWS (Amazon SageMaker)): 0
-Which type of machine are you using? ([0] No distributed training, [1] multi-CPU, [2] multi-GPU, [3] TPU [4] MPS): 2
+Which type of machine are you using? ([0] No distributed training, [1] multi-CPU, [2] multi-GPU, [3] TPU): 2
 How many different machines will you use (use more than 1 for multi-node training)? [1]: 
 Do you want to use DeepSpeed? [yes/NO]: 
 Do you want to use FullyShardedDataParallel? [yes/NO]: 
@@ -290,6 +293,7 @@ You will implement the `accelerate.utils.AbstractTrainStep` or inherit from thei
 ```python
 from accelerate.utils import MegatronLMDummyScheduler, GPTTrainStep, avg_losses_across_data_parallel_group
 
+
 # Custom loss function for the Megatron model
 class GPTTrainStepWithCustomLoss(GPTTrainStep):
     def __init__(self, megatron_args, **kwargs):

docs/source/usage_guides/memory.md

@@ -8,6 +8,9 @@ http://www.apache.org/licenses/LICENSE-2.0
 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
 an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
 specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
 -->
 
 # Memory Utilities

docs/source/usage_guides/model_size_estimator.md

@@ -0,0 +1,121 @@
+<!--Copyright 2022 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
+-->
+
+# Understanding how big of a model can fit on your machine
+
+One very difficult aspect when exploring potential models to use on your machine is knowing just how big of a model will *fit* into memory with your current graphics card (such as loading the model onto CUDA).
+
+To help alleviate this, 🤗 Accelerate has a CLI interface through `accelerate estimate-memory`. This tutorial will 
+help walk you through using it, what to expect, and at the end link to the interactive demo hosted on the 🤗 Hub which will 
+even let you post those results directly on the model repo!
+
+Currently we support searching for models that can be used in `timm` and `transformers`.
+
+<Tip>
+
+    This API will load the model into memory on the `meta` device, so we are not actually downloading 
+    and loading the full weights of the model into memory, nor do we need to. As a result it's 
+    perfectly fine to measure 8 billion parameter models (or more), without having to worry about 
+    if your CPU can handle it!
+
+</Tip>
+
+## The Command
+
+When using `accelerate estimate-memory`, you need to pass in the name of the model you want to use, potentially the framework
+that model utilizing (if it can't be found automatically), and the data types you want the model to be loaded in with.
+
+For example, here is how we can calculate the memory footprint for `bert-base-cased`:
+
+```bash
+accelerate estimate-memory bert-base-cased
+```
+
+This will download the `config.json` for `bert-based-cased`, load the model on the `meta` device, and report back how much space
+it will use:
+
+Memory Usage for loading `bert-base-cased`:
+
+| dtype   | Largest Layer | Total Size | Training using Adam |
+|---------|---------------|------------|---------------------|
+| float32 | 84.95 MB      | 418.18 MB  | 1.61 GB             |
+| float16 | 42.47 MB      | 206.59 MB  | 826.36 MB           |
+| int8    | 21.24 MB      | 103.29 MB  | 413.18 MB           |
+| int4    | 10.62 MB      | 51.65 MB   | 206.59 MB           |
+
+By default it will return all the supported dtypes (`int4` through `float32`), but if you are interested in specific ones these can be filtered.
+
+### Specific libraries
+
+If the source library cannot be determined automatically (like it could in the case of `bert-base-cased`), a library name can
+be passed in. 
+
+```bash
+accelerate estimate-memory HuggingFaceM4/idefics-80b-instruct --library_name transformers
+```
+
+Memory Usage for loading `HuggingFaceM4/idefics-80b-instruct`:
+
+| dtype   | Largest Layer | Total Size | Training using Adam |
+|---------|---------------|------------|---------------------|
+| float32 | 3.02 GB       | 297.12 GB  | 1.16 TB             |
+| float16 | 1.51 GB       | 148.56 GB  | 594.24 GB           |
+| int8    | 772.52 MB     | 74.28 GB   | 297.12 GB           |
+| int4    | 386.26 MB     | 37.14 GB   | 148.56 GB           |
+
+
+```bash
+accelerate estimate-memory timm/resnet50.a1_in1k --library_name timm
+```
+
+Memory Usage for loading `timm/resnet50.a1_in1k`:
+
+| dtype   | Largest Layer | Total Size | Training using Adam |
+|---------|---------------|------------|---------------------|
+| float32 | 9.0 MB        | 97.7 MB    | 390.78 MB           |
+| float16 | 4.5 MB        | 48.85 MB   | 195.39 MB           |
+| int8    | 2.25 MB       | 24.42 MB   | 97.7 MB             |
+| int4    | 1.12 MB       | 12.21 MB   | 48.85 MB            |
+
+### Specific dtypes
+
+As mentioned earlier, while we return `int4` through `float32` by default, any dtype can be used from `float32`, `float16`, `int8`, and `int4`.
+
+To do so, pass them in after specifying `--dtypes`:
+
+```bash
+accelerate estimate-memory bert-base-cased --dtypes float32 float16
+```
+
+Memory Usage for loading `bert-base-cased`:
+
+| dtype   | Largest Layer | Total Size | Training using Adam |
+|---------|---------------|------------|---------------------|
+| float32 | 84.95 MB      | 413.18 MB  | 1.61 GB             |
+| float16 | 42.47 MB      | 206.59 MB  | 826.36 MB           |
+
+## Caveats with this calculator
+
+This calculator will tell you how much memory is needed to purely load the model in, *not* to perform inference.
+
+This calculation is accurate within a few % of the actual value, so it is a very good view of just how much memory it will take. For instance loading `bert-base-cased` actually takes `413.68 MB` when loaded on CUDA in full precision, and the calculator estimates `413.18 MB`.
+
+When performing inference you can expect to add up to an additional 20% as found by [EleutherAI](https://blog.eleuther.ai/transformer-math/). We'll be conducting research into finding a more accurate estimate to these values, and will update 
+this calculator once done.
+
+## Live Gradio Demo
+
+Lastly, we invite you to try the [live Gradio demo](https://huggingface.co/spaces/hf-accelerate/model-memory-usage) of this utility,
+which includes an option to post a discussion thread on a models repository with this data. Doing so will help provide access to these numbers in the community faster and help users know what you've learned!

docs/source/usage_guides/mps.md

@@ -8,6 +8,9 @@ http://www.apache.org/licenses/LICENSE-2.0
 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
 an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
 specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
 -->
 
 # Accelerated PyTorch Training on Mac
@@ -31,41 +34,10 @@ please follow this nice medium article [GPU-Acceleration Comes to PyTorch on M1
 
 
 ## How it works out of the box
+It is enabled by default on MacOs machines with MPS enabled Apple Silicon GPUs.
+To disable it, pass `--cpu` flag to `accelerate launch` command or answer the corresponding question when answering the `accelerate config` questionnaire.
 
-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:
-
+You can directly run the following script to test it out on MPS enabled Apple Silicon machines:
 ```bash
 accelerate launch /examples/cv_example.py --data_dir images
 ```

docs/source/usage_guides/quantization.md

@@ -0,0 +1,136 @@
+<!--Copyright 2023 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
+-->
+
+# Quantization
+
+## `bitsandbytes` Integration
+
+🤗 Accelerate brings `bitsandbytes` quantization to your model. You can now load any pytorch model in 8-bit or 4-bit with a few lines of code.
+
+If you want to use 🤗 Transformers models with `bitsandbytes`, you should follow this [documentation](https://huggingface.co/docs/transformers/main_classes/quantization). 
+
+To learn more about how the `bitsandbytes` quantization works, check out the blog posts on [8-bit quantization](https://huggingface.co/blog/hf-bitsandbytes-integration) and [4-bit quantization](https://huggingface.co/blog/4bit-transformers-bitsandbytes).
+
+### Pre-Requisites
+You will need to install the following requirements:
+
+- Install `bitsandbytes` library
+```bash
+pip install bitsandbytes
+```
+- Install latest `accelerate` from source
+```bash
+pip install git+https://github.com/huggingface/accelerate.git
+```
+- Install `minGPT` and `huggingface_hub` to run examples
+```bash
+git clone https://github.com/karpathy/minGPT.git
+pip install minGPT/
+pip install huggingface_hub
+```
+
+### How it works
+
+First, we need to initialize our model. To save memory, we can initialize an empty model using the context manager [`init_empty_weights`]. 
+
+Let's take the GPT2 model from minGPT library.
+```py
+from accelerate import init_empty_weights
+from mingpt.model import GPT
+
+model_config = GPT.get_default_config()
+model_config.model_type = 'gpt2-xl'
+model_config.vocab_size = 50257
+model_config.block_size = 1024
+
+with init_empty_weights():
+    empty_model = GPT(model_config)
+```
+
+Then, we need to get the path to the weights of your model. The path can be the state_dict file (e.g. "pytorch_model.bin") or a folder containing the sharded checkpoints. 
+
+```py
+from huggingface_hub import snapshot_download
+weights_location = snapshot_download(repo_id="marcsun13/gpt2-xl-linear-sharded")
+```
+
+Finally, you need to set your quantization configuration with [`~utils.BnbQuantizationConfig`].
+
+Here's an example for 8-bit quantization:
+```py
+from accelerate.utils import BnbQuantizationConfig
+bnb_quantization_config = BnbQuantizationConfig(load_in_8bit=True, llm_int8_threshold = 6)
+```
+
+Here's an example for 4-bit quantization:
+```py
+from accelerate.utils import BnbQuantizationConfig
+bnb_quantization_config = BnbQuantizationConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4")
+```
+
+To quantize your empty model with the selected configuration, you need to use [`~utils.load_and_quantize_model`]. 
+
+```py
+from accelerate.utils import load_and_quantize_model
+quantized_model = load_and_quantize_model(empty_model, weights_location=weights_location, bnb_quantization_config=bnb_quantization_config, device_map = "auto")
+```
+
+### Saving and loading 8-bit model
+
+You can save your 8-bit model with accelerate using [`~Accelerator.save_model`]. 
+
+```py
+from accelerate import Accelerator
+accelerate = Accelerator()
+new_weights_location = "path/to/save_directory"
+accelerate.save_model(quantized_model, new_weights_location)
+
+quantized_model_from_saved = load_and_quantize_model(empty_model, weights_location=new_weights_location, bnb_quantization_config=bnb_quantization_config, device_map = "auto")
+```
+
+Note that 4-bit model serialization is currently not supported.
+
+### Offload modules to cpu and disk 
+
+You can offload some modules to cpu/disk if you don't have enough space on the GPU to store the entire model on your GPUs.
+This uses big model inference under the hood. Check this [documentation](https://huggingface.co/docs/accelerate/usage_guides/big_modeling) for more details. 
+
+For 8-bit quantization, the selected modules will be converted to 8-bit precision. 
+
+For 4-bit quantization, the selected modules will be kept in `torch_dtype` that the user passed in `BnbQuantizationConfig`.  We will add support to convert these offloaded modules in 4-bit when 4-bit serialization will be possible. 
+
+ You just need to pass a custom `device_map` in order to offload modules on cpu/disk. The offload modules will be dispatched on the GPU when needed. Here's an example :
+
+```py
+device_map = {
+    "transformer.wte": 0,
+    "transformer.wpe": 0,
+    "transformer.drop": 0,
+    "transformer.h": "cpu",
+    "transformer.ln_f": "disk",
+    "lm_head": "disk",
+}
+```
+### Fine-tune a quantized model
+
+It is not possible to perform pure 8bit or 4bit training on these models. However, you can train these models by leveraging parameter efficient fine tuning methods (PEFT) and train for example adapters on top of them. Please have a look at [peft](https://github.com/huggingface/peft) library for more details.
+
+Currently, you can't add adapters on top of any quantized model. However, with the official support of adapters with 🤗 Transformers models, you can fine-tune quantized models. If you want to finetune a 🤗 Transformers model , follow this [documentation](https://huggingface.co/docs/transformers/main_classes/quantization) instead. Check out this [demo](https://colab.research.google.com/drive/1VoYNfYDKcKRQRor98Zbf2-9VQTtGJ24k?usp=sharing) on how to fine-tune a 4-bit 🤗 Transformers model. 
+
+Note that you don’t need to pass `device_map` when loading the model for training. It will automatically load your model on your GPU. Please note that `device_map=auto` should be used for inference only.
+
+### Example demo - running GPT2 1.5b on a Google Colab
+
+Check out the Google Colab [demo](https://colab.research.google.com/drive/1T1pOgewAWVpR9gKpaEWw4orOrzPFb3yM?usp=sharing) for running quantized models on a GTP2 model. The GPT2-1.5B model checkpoint is in FP32 which uses 6GB of memory. After quantization, it uses 1.6GB with 8-bit modules and 1.2GB with 4-bit modules.

docs/source/usage_guides/sagemaker.md

@@ -8,6 +8,9 @@ http://www.apache.org/licenses/LICENSE-2.0
 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
 an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
 specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
 -->
 
 # Amazon SageMaker
@@ -160,10 +163,43 @@ use_cpu: false
 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
+### Local Training: SageMaker Local mode
 
-*undecided if feature is needed. Contact us if you would like this feature.*
+The local mode in the SageMaker SDK allows you to run your training script locally inside the HuggingFace DLC (Deep Learning container) 
+or using your custom container image. This is useful for debugging and testing your training script inside the final container environment.
+Local mode uses Docker compose (*Note: Docker Compose V2 is not supported yet*). The SDK will handle the authentication against ECR
+to pull the DLC to your local environment. You can emulate CPU (single and multi-instance) and GPU (single instance) SageMaker training jobs.
+
+To use local mode, you need to set your `ec2_instance_type` to `local`.
+
+```yaml
+ec2_instance_type: local
+```
+
+### Advanced configuration
+
+The configuration allows you to override parameters for the [Estimator](https://sagemaker.readthedocs.io/en/stable/api/training/estimators.html).
+These settings have to be applied in the config file and are not part of `accelerate config`. You can control many additional aspects of the training job, e.g. use Spot instances, enable network isolation and many more.
+
+```yaml
+additional_args:
+  # enable network isolation to restrict internet access for containers
+  enable_network_isolation: True
+```
+
+You can find all available configuration [here](https://sagemaker.readthedocs.io/en/stable/api/training/estimators.html).
 
 ### Use Spot Instances
 
-*undecided if feature is needed. Contact us if you would like this feature.*
+You can use Spot Instances e.g. using (see [Advanced configuration](#advanced-configuration)):
+```yaml
+additional_args:
+  use_spot_instances: True
+  max_wait: 86400
+```
+
+*Note: Spot Instances are subject to be terminated and training to be continued from a checkpoint. This is not handled in 🤗 Accelerate out of the box. Contact us if you would like this feature.*
+
+### Remote scripts: Use scripts located on Github
+
+*undecided if feature is needed. Contact us if you would like this feature.*

docs/source/usage_guides/tracking.md

@@ -8,6 +8,9 @@ http://www.apache.org/licenses/LICENSE-2.0
 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
 an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
 specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
 -->
 
 # Tracking
@@ -83,11 +86,15 @@ for iteration in config["num_iterations"]:
 accelerator.end_training()
 ```
 
-If a tracker requires a directory to save data to such as `TensorBoard` then a `logging_dir` or `project_dir` can be passed in. `project_dir` is useful 
-if there are other further configurations such as those which can be combined with the [`~utils.ProjectConfiguration`] dataclass.
+If a tracker requires a directory to save data to, such as `TensorBoard`, then pass the directory path to `project_dir`. The `project_dir` parameter is useful 
+when there are other configurations to be combined with in the [`~utils.ProjectConfiguration`] data class. For example, you can save the TensorBoard data to `project_dir` and everything else can be logged in the `logging_dir` parameter of [`~utils.ProjectConfiguration`: 
 
 ```python
-accelerator = Accelerator(log_with="tensorboard", logging_dir=".")
+accelerator = Accelerator(log_with="tensorboard", project_dir=".")
+
+# use with ProjectConfiguration
+config = ProjectConfiguration(project_dir=".", logging_dir="another/directory")
+accelerator = Accelerator(log_with="tensorboard", project_config=config)
 ```
 
 ## Implementing Custom Trackers
@@ -111,9 +118,12 @@ Every tracker must implement three functions and have three properties:
     - 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:
+Each method should also utilize the [`state.PartialState`] class if the logger should only be executed on the main process for instance.
+
+A brief example can be seen below with an integration with Weights and Biases, containing only the relevant information and logging just on 
+the main process:
 ```python
-from accelerate.tracking import GeneralTracker
+from accelerate.tracking import GeneralTracker, on_main_process
 from typing import Optional
 
 import wandb
@@ -123,6 +133,7 @@ class MyCustomTracker(GeneralTracker):
     name = "wandb"
     requires_logging_directory = False
 
+    @on_main_process
     def __init__(self, run_name: str):
         self.run_name = run_name
         run = wandb.init(self.run_name)
@@ -131,9 +142,11 @@ class MyCustomTracker(GeneralTracker):
     def tracker(self):
         return self.run.run
 
+    @on_main_process
     def store_init_configuration(self, values: dict):
         wandb.config(values)
 
+    @on_main_process
     def log(self, values: dict, step: Optional[int] = None):
         wandb.log(values, step=step)
 ```
@@ -167,16 +180,26 @@ 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!
+```python
+wandb_run.log_artifact(some_artifact_to_log)
+```
+
+<Tip>
+  Trackers built in Accelerate will automatically execute on the correct process, 
+  so if a tracker is only meant to be ran on the main process it will do so 
+  automatically.
 </Tip>
 
+If you want to truly remove Accelerate's wrapping entirely, you can 
+achieve the same outcome with:
 
 ```python
-if accelerator.is_main_process:
-    wandb_run.log_artifact(some_artifact_to_log)
+wandb_tracker = accelerator.get_tracker("wandb", unwrap=True)
+with accelerator.on_main_process:
+    wandb_tracker.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:

docs/source/usage_guides/training_zoo.md

@@ -0,0 +1,175 @@
+<!--Copyright 2022 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
+-->
+
+# 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 distributed NLP example in a Jupyter Notebook](https://github.com/huggingface/notebooks/blob/main/examples/accelerate_examples/simple_nlp_example.ipynb)
+- [Barebones computer vision example](https://github.com/huggingface/accelerate/blob/main/examples/cv_example.py)
+- [Barebones distributed computer vision example in a Jupyter Notebook](https://github.com/huggingface/notebooks/blob/main/examples/accelerate_examples/simple_cv_example.ipynb)
+- [Using Accelerate in Kaggle](https://www.kaggle.com/code/muellerzr/multi-gpu-and-accelerate)
+
+### 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)
+- [Using Megatron-LM](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/megatron_lm_gpt_pretraining.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)
+- [Very complete and extensible vision example showcasing SLURM, hydra, and a very extensible usage of the framework](https://github.com/yuvalkirstain/PickScore)
+- [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)
+- [End-to-end examples on how to use AWS SageMaker integration of Accelerate](https://github.com/huggingface/notebooks/blob/main/sagemaker/22_accelerate_sagemaker_examples/README.md)
+- [Megatron-LM examples for various NLp tasks](https://github.com/pacman100/accelerate-megatron-test) 
+
+## Integration Examples 
+
+These are tutorials from libraries that integrate with 🤗 Accelerate: 
+
+> Don't find your integration here? Make a PR to include it!
+
+### 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/)
+
+### Stable-Dreamfusion
+
+- [Training with Stable-Dreamfusion to convert text to a 3D model](https://colab.research.google.com/drive/1MXT3yfOFvO0ooKEfiUUvTKwUkrrlCHpF?usp=sharing)
+
+### 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)
+
+### Comfy-UI
+
+- [Enabling using large Stable Diffusion Models in low-vram settings using Accelerate](https://github.com/comfyanonymous/ComfyUI/blob/master/comfy/model_management.py#L291-L296)
+
+
+## In Science
+
+Below contains a non-exhaustive list of papers utilizing 🤗 Accelerate. 
+
+> Don't find your paper here? Make a PR to include it!
+
+* Yuval Kirstain, Adam Polyak, Uriel Singer, Shahbuland Matiana, Joe Penna, Omer Levy: “Pick-a-Pic: An Open Dataset of User Preferences for Text-to-Image Generation”, 2023; [arXiv:2305.01569](http://arxiv.org/abs/2305.01569).
+* Lei Wang, Wanyu Xu, Yihuai Lan, Zhiqiang Hu, Yunshi Lan, Roy Ka-Wei Lee, Ee-Peng Lim: “Plan-and-Solve Prompting: Improving Zero-Shot Chain-of-Thought Reasoning by Large Language Models”, 2023; [arXiv:2305.04091](http://arxiv.org/abs/2305.04091).
+* Arthur Câmara, Claudia Hauff: “Moving Stuff Around: A study on efficiency of moving documents into memory for Neural IR models”, 2022; [arXiv:2205.08343](http://arxiv.org/abs/2205.08343).
+* Ying Sheng, Lianmin Zheng, Binhang Yuan, Zhuohan Li, Max Ryabinin, Daniel Y. Fu, Zhiqiang Xie, Beidi Chen, Clark Barrett, Joseph E. Gonzalez, Percy Liang, Christopher Ré, Ion Stoica, Ce Zhang: “High-throughput Generative Inference of Large Language Models with a Single GPU”, 2023; [arXiv:2303.06865](http://arxiv.org/abs/2303.06865).
+* Peter Melchior, Yan Liang, ChangHoon Hahn, Andy Goulding: “Autoencoding Galaxy Spectra I: Architecture”, 2022; [arXiv:2211.07890](http://arxiv.org/abs/2211.07890).
+* Jiaao Chen, Aston Zhang, Mu Li, Alex Smola, Diyi Yang: “A Cheaper and Better Diffusion Language Model with Soft-Masked Noise”, 2023; [arXiv:2304.04746](http://arxiv.org/abs/2304.04746).
+* Ayaan Haque, Matthew Tancik, Alexei A. Efros, Aleksander Holynski, Angjoo Kanazawa: “Instruct-NeRF2NeRF: Editing 3D Scenes with Instructions”, 2023; [arXiv:2303.12789](http://arxiv.org/abs/2303.12789).
+* Luke Melas-Kyriazi, Christian Rupprecht, Iro Laina, Andrea Vedaldi: “RealFusion: 360° Reconstruction of Any Object from a Single Image”, 2023; [arXiv:2302.10663](http://arxiv.org/abs/2302.10663).
+* Xiaoshi Wu, Keqiang Sun, Feng Zhu, Rui Zhao, Hongsheng Li: “Better Aligning Text-to-Image Models with Human Preference”, 2023; [arXiv:2303.14420](http://arxiv.org/abs/2303.14420).
+* Yongliang Shen, Kaitao Song, Xu Tan, Dongsheng Li, Weiming Lu, Yueting Zhuang: “HuggingGPT: Solving AI Tasks with ChatGPT and its Friends in HuggingFace”, 2023; [arXiv:2303.17580](http://arxiv.org/abs/2303.17580).
+* Yue Yang, Wenlin Yao, Hongming Zhang, Xiaoyang Wang, Dong Yu, Jianshu Chen: “Z-LaVI: Zero-Shot Language Solver Fueled by Visual Imagination”, 2022; [arXiv:2210.12261](http://arxiv.org/abs/2210.12261).
+* Sheng-Yen Chou, Pin-Yu Chen, Tsung-Yi Ho: “How to Backdoor Diffusion Models?”, 2022; [arXiv:2212.05400](http://arxiv.org/abs/2212.05400).
+* Junyoung Seo, Wooseok Jang, Min-Seop Kwak, Jaehoon Ko, Hyeonsu Kim, Junho Kim, Jin-Hwa Kim, Jiyoung Lee, Seungryong Kim: “Let 2D Diffusion Model Know 3D-Consistency for Robust Text-to-3D Generation”, 2023; [arXiv:2303.07937](http://arxiv.org/abs/2303.07937).
+* Or Patashnik, Daniel Garibi, Idan Azuri, Hadar Averbuch-Elor, Daniel Cohen-Or: “Localizing Object-level Shape Variations with Text-to-Image Diffusion Models”, 2023; [arXiv:2303.11306](http://arxiv.org/abs/2303.11306).
+* Dídac Surís, Sachit Menon, Carl Vondrick: “ViperGPT: Visual Inference via Python Execution for Reasoning”, 2023; [arXiv:2303.08128](http://arxiv.org/abs/2303.08128).
+* Chenyang Qi, Xiaodong Cun, Yong Zhang, Chenyang Lei, Xintao Wang, Ying Shan, Qifeng Chen: “FateZero: Fusing Attentions for Zero-shot Text-based Video Editing”, 2023; [arXiv:2303.09535](http://arxiv.org/abs/2303.09535).
+* Sean Welleck, Jiacheng Liu, Ximing Lu, Hannaneh Hajishirzi, Yejin Choi: “NaturalProver: Grounded Mathematical Proof Generation with Language Models”, 2022; [arXiv:2205.12910](http://arxiv.org/abs/2205.12910).
+* Elad Richardson, Gal Metzer, Yuval Alaluf, Raja Giryes, Daniel Cohen-Or: “TEXTure: Text-Guided Texturing of 3D Shapes”, 2023; [arXiv:2302.01721](http://arxiv.org/abs/2302.01721).
+* Puijin Cheng, Li Lin, Yijin Huang, Huaqing He, Wenhan Luo, Xiaoying Tang: “Learning Enhancement From Degradation: A Diffusion Model For Fundus Image Enhancement”, 2023; [arXiv:2303.04603](http://arxiv.org/abs/2303.04603).
+* Shun Shao, Yftah Ziser, Shay Cohen: “Erasure of Unaligned Attributes from Neural Representations”, 2023; [arXiv:2302.02997](http://arxiv.org/abs/2302.02997).
+* Seonghyeon Ye, Hyeonbin Hwang, Sohee Yang, Hyeongu Yun, Yireun Kim, Minjoon Seo: “In-Context Instruction Learning”, 2023; [arXiv:2302.14691](http://arxiv.org/abs/2302.14691).
+* Shikun Liu, Linxi Fan, Edward Johns, Zhiding Yu, Chaowei Xiao, Anima Anandkumar: “Prismer: A Vision-Language Model with An Ensemble of Experts”, 2023; [arXiv:2303.02506](http://arxiv.org/abs/2303.02506 ).
+* Haoyu Chen, Zhihua Wang, Yang Yang, Qilin Sun, Kede Ma: “Learning a Deep Color Difference Metric for Photographic Images”, 2023; [arXiv:2303.14964](http://arxiv.org/abs/2303.14964).
+* Van-Hoang Le, Hongyu Zhang: “Log Parsing with Prompt-based Few-shot Learning”, 2023; [arXiv:2302.07435](http://arxiv.org/abs/2302.07435).
+* Keito Kudo, Yoichi Aoki, Tatsuki Kuribayashi, Ana Brassard, Masashi Yoshikawa, Keisuke Sakaguchi, Kentaro Inui: “Do Deep Neural Networks Capture Compositionality in Arithmetic Reasoning?”, 2023; [arXiv:2302.07866](http://arxiv.org/abs/2302.07866).
+* Ruoyao Wang, Peter Jansen, Marc-Alexandre Côté, Prithviraj Ammanabrolu: “Behavior Cloned Transformers are Neurosymbolic Reasoners”, 2022; [arXiv:2210.07382](http://arxiv.org/abs/2210.07382).
+* Martin Wessel, Tomáš Horych, Terry Ruas, Akiko Aizawa, Bela Gipp, Timo Spinde: “Introducing MBIB -- the first Media Bias Identification Benchmark Task and Dataset Collection”, 2023; [arXiv:2304.13148](http://arxiv.org/abs/2304.13148 ). DOI: [https://dx.doi.org/10.1145/3539618.3591882 10.1145/3539618.3591882].
+* Hila Chefer, Yuval Alaluf, Yael Vinker, Lior Wolf, Daniel Cohen-Or: “Attend-and-Excite: Attention-Based Semantic Guidance for Text-to-Image Diffusion Models”, 2023; [arXiv:2301.13826](http://arxiv.org/abs/2301.13826).
+* Marcio Fonseca, Yftah Ziser, Shay B. Cohen: “Factorizing Content and Budget Decisions in Abstractive Summarization of Long Documents”, 2022; [arXiv:2205.12486](http://arxiv.org/abs/2205.12486).
+* Elad Richardson, Gal Metzer, Yuval Alaluf, Raja Giryes, Daniel Cohen-Or: “TEXTure: Text-Guided Texturing of 3D Shapes”, 2023; [arXiv:2302.01721](http://arxiv.org/abs/2302.01721).
+* Tianxing He, Jingyu Zhang, Tianle Wang, Sachin Kumar, Kyunghyun Cho, James Glass, Yulia Tsvetkov: “On the Blind Spots of Model-Based Evaluation Metrics for Text Generation”, 2022; [arXiv:2212.10020](http://arxiv.org/abs/2212.10020).
+* Ori Ram, Yoav Levine, Itay Dalmedigos, Dor Muhlgay, Amnon Shashua, Kevin Leyton-Brown, Yoav Shoham: “In-Context Retrieval-Augmented Language Models”, 2023; [arXiv:2302.00083](http://arxiv.org/abs/2302.00083).
+* Dacheng Li, Rulin Shao, Hongyi Wang, Han Guo, Eric P. Xing, Hao Zhang: “MPCFormer: fast, performant and private Transformer inference with MPC”, 2022; [arXiv:2211.01452](http://arxiv.org/abs/2211.01452).
+* Baolin Peng, Michel Galley, Pengcheng He, Chris Brockett, Lars Liden, Elnaz Nouri, Zhou Yu, Bill Dolan, Jianfeng Gao: “GODEL: Large-Scale Pre-Training for Goal-Directed Dialog”, 2022; [arXiv:2206.11309](http://arxiv.org/abs/2206.11309).
+* Egil Rønningstad, Erik Velldal, Lilja Øvrelid: “Entity-Level Sentiment Analysis (ELSA): An exploratory task survey”, 2023, Proceedings of the 29th International Conference on Computational Linguistics, 2022, pages 6773-6783; [arXiv:2304.14241](http://arxiv.org/abs/2304.14241).
+* Charlie Snell, Ilya Kostrikov, Yi Su, Mengjiao Yang, Sergey Levine: “Offline RL for Natural Language Generation with Implicit Language Q Learning”, 2022; [arXiv:2206.11871](http://arxiv.org/abs/2206.11871).
+* Zhiruo Wang, Shuyan Zhou, Daniel Fried, Graham Neubig: “Execution-Based Evaluation for Open-Domain Code Generation”, 2022; [arXiv:2212.10481](http://arxiv.org/abs/2212.10481).
+* Minh-Long Luu, Zeyi Huang, Eric P. Xing, Yong Jae Lee, Haohan Wang: “Expeditious Saliency-guided Mix-up through Random Gradient Thresholding”, 2022; [arXiv:2212.04875](http://arxiv.org/abs/2212.04875).
+* Jun Hao Liew, Hanshu Yan, Daquan Zhou, Jiashi Feng: “MagicMix: Semantic Mixing with Diffusion Models”, 2022; [arXiv:2210.16056](http://arxiv.org/abs/2210.16056).
+* Yaqing Wang, Subhabrata Mukherjee, Xiaodong Liu, Jing Gao, Ahmed Hassan Awadallah, Jianfeng Gao: “LiST: Lite Prompted Self-training Makes Parameter-Efficient Few-shot Learners”, 2021; [arXiv:2110.06274](http://arxiv.org/abs/2110.06274).

docs/source/usage_guides/training_zoo.mdx

@@ -1,117 +0,0 @@
-<!--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 distributed NLP example in a Jupyter Notebook](https://github.com/huggingface/notebooks/blob/main/examples/accelerate_examples/simple_nlp_example.ipynb)
-- [Barebones computer vision example](https://github.com/huggingface/accelerate/blob/main/examples/cv_example.py)
-- [Barebones distributed computer vision example in a Jupyter Notebook](https://github.com/huggingface/notebooks/blob/main/examples/accelerate_examples/simple_cv_example.ipynb)
-- [Using Accelerate in Kaggle](https://www.kaggle.com/code/muellerzr/multi-gpu-and-accelerate)
-
-### 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)
-- [Using Megatron-LM](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/megatron_lm_gpt_pretraining.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)
-- [End-to-end examples on how to use AWS SageMaker integration of Accelerate](https://github.com/huggingface/notebooks/blob/main/sagemaker/22_accelerate_sagemaker_examples/README.md)
-- [Megatron-LM examples for various NLp tasks](https://github.com/pacman100/accelerate-megatron-test) 
-
-## 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/)
-
-### Stable-Dreamfusion
-
-- [Training with Stable-Dreamfusion to convert text to a 3D model](https://colab.research.google.com/drive/1MXT3yfOFvO0ooKEfiUUvTKwUkrrlCHpF?usp=sharing)
-
-### 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

@@ -51,22 +51,22 @@ To run it in each of these various modes, use the following commands:
     python ./nlp_example.py  # from a server with a GPU
     ```
 - with fp16 (mixed-precision)
-    * from any server by passing `fp16=True` to the `Accelerator`.
+    * from any server by passing `mixed_precison=fp16` to the `Accelerator`.
         ```bash
-        python ./nlp_example.py --fp16
+        python ./nlp_example.py --mixed_precision fp16
         ```
     * from any server with Accelerate launcher
         ```bash
-        accelerate launch --fp16 ./nlp_example.py
+        accelerate launch --mixed_precision 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
         accelerate launch ./nlp_example.py  # This will run the script on your server
         ```
-    * With traditional PyTorch launcher
+    * With traditional PyTorch launcher (`torch.distributed.launch` can be used with older versions of PyTorch)
         ```bash
-        python -m torch.distributed.launch --nproc_per_node 2 --use_env ./nlp_example.py
+        python -m torchrun --nproc_per_node 2 --use_env ./nlp_example.py
         ```
 - multi GPUs, multi node (several machines, using PyTorch distributed mode)
     * With Accelerate config and launcher, on each machine:
@@ -74,14 +74,14 @@ To run it in each of these various modes, use the following commands:
         accelerate config  # This will create a config file on each server
         accelerate launch ./nlp_example.py  # This will run the script on each server
         ```
-    * With PyTorch launcher only
+    * With PyTorch launcher only (`torch.distributed.launch` can be used in older versions of PyTorch)
         ```bash
-        python -m torch.distributed.launch --nproc_per_node 2 \
+        python -m torchrun --nproc_per_node 2 \
             --use_env \
             --node_rank 0 \
             --master_addr master_node_ip_address \
             ./nlp_example.py  # On the first server
-        python -m torch.distributed.launch --nproc_per_node 2 \
+        python -m torchrun --nproc_per_node 2 \
             --use_env \
             --node_rank 1 \
             --master_addr master_node_ip_address \
@@ -139,22 +139,22 @@ To run it in each of these various modes, use the following commands:
     python ./cv_example.py  # from a server with a GPU
     ```
 - with fp16 (mixed-precision)
-    * from any server by passing `fp16=True` to the `Accelerator`.
+    * from any server by passing `mixed_precison=fp16` to the `Accelerator`.
         ```bash
-        python ./cv_example.py --data_dir path_to_data --fp16
+        python ./cv_example.py --data_dir path_to_data --mixed_precison fp16
         ```
     * from any server with Accelerate launcher
         ```bash
-        accelerate launch --fp16 ./cv_example.py --data_dir path_to_data
+        accelerate launch --mixed_precison 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
         accelerate launch ./cv_example.py --data_dir path_to_data  # This will run the script on your server
         ```
-    * With traditional PyTorch launcher
+    * With traditional PyTorch launcher (`torch.distributed.launch` can be used with older versions of PyTorch)
         ```bash
-        python -m torch.distributed.launch --nproc_per_node 2 --use_env ./cv_example.py --data_dir path_to_data
+        python -m torchrun --nproc_per_node 2 --use_env ./cv_example.py --data_dir path_to_data
         ```
 - multi GPUs, multi node (several machines, using PyTorch distributed mode)
     * With Accelerate config and launcher, on each machine:
@@ -162,14 +162,14 @@ To run it in each of these various modes, use the following commands:
         accelerate config  # This will create a config file on each server
         accelerate launch ./cv_example.py --data_dir path_to_data  # This will run the script on each server
         ```
-    * With PyTorch launcher only
+    * With PyTorch launcher only (`torch.distributed.launch` can be used with older versions of PyTorch)
         ```bash
-        python -m torch.distributed.launch --nproc_per_node 2 \
+        python -m torchrun --nproc_per_node 2 \
             --use_env \
             --node_rank 0 \
             --master_addr master_node_ip_address \
             ./cv_example.py --data_dir path_to_data  # On the first server
-        python -m torch.distributed.launch --nproc_per_node 2 \
+        python -m torchrun --nproc_per_node 2 \
             --use_env \
             --node_rank 1 \
             --master_addr master_node_ip_address \
@@ -190,7 +190,22 @@ To run it in each of these various modes, use the following commands:
 
 ### Using AWS SageMaker integration
 - [Examples showcasing AWS SageMaker integration of 🤗 Accelerate.](https://github.com/pacman100/accelerate-aws-sagemaker)
-    
+
+
+## Simple Multi-GPU Hardware Launcher
+
+[multigpu_remote_launcher.py](./multigpu_remote_launcher.py) is a minimal script that demonstrates launching accelerate
+on multiple remote GPUs, and with automatic hardware environment and dependency setup for reproducibility. You can
+easily customize the training function used, training arguments, hyperparameters, and type of compute hardware, and then
+run the script to automatically launch multi GPU training on remote hardware.
+
+This script uses [Runhouse](https://github.com/run-house/runhouse) to launch on self-hosted hardware (e.g. in your own
+cloud account or on-premise cluster) but there are other options for running remotely as well. Runhouse can be installed
+with `pip install runhouse`, and you can refer to
+[hardware setup](https://runhouse-docs.readthedocs-hosted.com/en/latest/api/python/cluster.html#hardware-setup)
+for hardware setup instructions, or this
+[Colab tutorial](https://colab.research.google.com/drive/1qVwYyLTCPYPSdz9ZX7BZl9Qm0A3j7RJe) for a more in-depth walkthrough.
+
 ## 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.

examples/by_feature/README.md

@@ -19,7 +19,7 @@ Adjustments to each script from the base `nlp_example.py` file can be found quic
 
 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:
+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.run`), such as:
 
 ```bash
 accelerate launch ../nlp_example.py --mixed_precision fp16 --cpu 0
@@ -34,7 +34,7 @@ accelerate launch ../nlp_example.py --mixed_precision fp16 --cpu 0
   - `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:
+These arguments should be added at the end of any method for starting the python script (such as `python`, `accelerate launch`, `python -m torchrun`), such as:
 
 (Note, `resume_from_checkpoint` assumes that we've ran the script for one epoch with the `--checkpointing_steps epoch` flag)
 
@@ -48,7 +48,7 @@ accelerate launch ./checkpointing.py --checkpointing_steps epoch output_dir "che
 - 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:
+These arguments should be added at the end of any method for starting the python script (such as `python`, `accelerate launch`, `python -m torchrun`), such as:
 
 ```bash
 accelerate launch ./cross_validation.py --num_folds 2
@@ -61,7 +61,7 @@ accelerate launch ./cross_validation.py --num_folds 2
 - 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:
+These arguments should be added at the end of any method for starting the python script (such as `python`, `accelerate launch`, `python -m torchrun`), such as:
 
 ```bash
 accelerate launch ./tracking.py --with_tracking
@@ -73,8 +73,19 @@ accelerate launch ./tracking.py --with_tracking
 - 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:
+These arguments should be added at the end of any method for starting the python script (such as `python`, `accelerate launch`, `python -m torchrun`), such as:
 
 ```bash
 accelerate launch ./gradient_accumulation.py --gradient_accumulation_steps 5
-```
+```
+
+### LocalSGD (`local_sgd.py`)
+- Shows how to use `Accelerator.no_sync` to prevent gradient averaging in a distributed setup. However, unlike gradient accumulation, this method does not change the effective batch size. Local SGD can be combined with gradient accumulation.
+
+These arguments should be added at the end of any method for starting the python script (such as `python`, `accelerate launch`, `python -m torchrun`), such as:
+
+```bash
+accelerate launch ./local_sgd.py --local_sgd_steps 4
+```
+
+

examples/by_feature/automatic_gradient_accumulation.py

@@ -14,17 +14,17 @@
 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
+import torch
 from datasets import load_dataset
+from torch.optim import AdamW
+from torch.utils.data import DataLoader
 from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed
 
+from accelerate import Accelerator
+from accelerate.utils import find_executable_batch_size
+
 
 ########################################################################
 # This is a fully working simple example to use Accelerate,
@@ -84,10 +84,20 @@ def get_dataloaders(accelerator: Accelerator, batch_size: int = 16):
     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")
+        # When using mixed precision we want round multiples of 8/16
+        if accelerator.mixed_precision == "fp8":
+            pad_to_multiple_of = 16
+        elif accelerator.mixed_precision != "no":
+            pad_to_multiple_of = 8
+        else:
+            pad_to_multiple_of = None
+
+        return tokenizer.pad(
+            examples,
+            padding="longest",
+            pad_to_multiple_of=pad_to_multiple_of,
+            return_tensors="pt",
+        )
 
     # Instantiate dataloaders.
     train_dataloader = DataLoader(
@@ -215,7 +225,7 @@ def main():
         "--mixed_precision",
         type=str,
         default=None,
-        choices=["no", "fp16", "bf16"],
+        choices=["no", "fp16", "bf16", "fp8"],
         help="Whether to use mixed precision. Choose"
         "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
         "and an Nvidia Ampere GPU.",

examples/by_feature/checkpointing.py

@@ -15,15 +15,15 @@
 import argparse
 import os
 
+import evaluate
 import torch
+from datasets import load_dataset
 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
 
+from accelerate import Accelerator, DistributedType
+
 
 ########################################################################
 # This is a fully working simple example to use Accelerate,
@@ -86,9 +86,22 @@ def get_dataloaders(accelerator: Accelerator, batch_size: int = 16):
 
     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")
+        max_length = 128 if accelerator.distributed_type == DistributedType.TPU else None
+        # When using mixed precision we want round multiples of 8/16
+        if accelerator.mixed_precision == "fp8":
+            pad_to_multiple_of = 16
+        elif accelerator.mixed_precision != "no":
+            pad_to_multiple_of = 8
+        else:
+            pad_to_multiple_of = None
+
+        return tokenizer.pad(
+            examples,
+            padding="longest",
+            max_length=max_length,
+            pad_to_multiple_of=pad_to_multiple_of,
+            return_tensors="pt",
+        )
 
     # Instantiate dataloaders.
     train_dataloader = DataLoader(
@@ -206,9 +219,12 @@ def training_function(config, args):
         # New Code #
         if args.resume_from_checkpoint and epoch == starting_epoch and resume_step is not None:
             # We need to skip steps until we reach the resumed step
-            train_dataloader = accelerator.skip_first_batches(train_dataloader, resume_step)
+            active_dataloader = accelerator.skip_first_batches(train_dataloader, resume_step)
             overall_step += resume_step
-        for step, batch in enumerate(train_dataloader):
+        else:
+            # After the first iteration though, we need to go back to the original dataloader
+            active_dataloader = train_dataloader
+        for step, batch in enumerate(active_dataloader):
             # We could avoid this line since we set the accelerator with `device_placement=True`.
             batch.to(accelerator.device)
             outputs = model(**batch)
@@ -269,7 +285,7 @@ def main():
         "--mixed_precision",
         type=str,
         default=None,
-        choices=["no", "fp16", "bf16"],
+        choices=["no", "fp16", "bf16", "fp8"],
         help="Whether to use mixed precision. Choose"
         "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
         "and an Nvidia Ampere GPU.",

examples/by_feature/cross_validation.py

@@ -15,20 +15,20 @@
 import argparse
 from typing import List
 
+import evaluate
 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 torch.optim import AdamW
+from torch.utils.data import DataLoader
 from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed
 
+from accelerate import Accelerator, DistributedType
+
 
 ########################################################################
 # This is a fully working simple example to use Accelerate,
@@ -106,9 +106,22 @@ def get_fold_dataloaders(
 
     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")
+        max_length = 128 if accelerator.distributed_type == DistributedType.TPU else None
+        # When using mixed precision we want round multiples of 8/16
+        if accelerator.mixed_precision == "fp8":
+            pad_to_multiple_of = 16
+        elif accelerator.mixed_precision != "no":
+            pad_to_multiple_of = 8
+        else:
+            pad_to_multiple_of = None
+
+        return tokenizer.pad(
+            examples,
+            padding="longest",
+            max_length=max_length,
+            pad_to_multiple_of=pad_to_multiple_of,
+            return_tensors="pt",
+        )
 
     # Instantiate dataloaders.
     train_dataloader = DataLoader(
@@ -251,7 +264,7 @@ def main():
         "--mixed_precision",
         type=str,
         default=None,
-        choices=["no", "fp16", "bf16"],
+        choices=["no", "fp16", "bf16", "fp8"],
         help="Whether to use mixed precision. Choose"
         "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
         "and an Nvidia Ampere GPU.",

examples/by_feature/deepspeed_with_config_support.py

@@ -31,16 +31,12 @@ import random
 from itertools import chain
 from pathlib import Path
 
-import torch
-from torch.utils.data import DataLoader
-
 import datasets
+import torch
 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 torch.utils.data import DataLoader
 from tqdm.auto import tqdm
 from transformers import (
     CONFIG_MAPPING,
@@ -55,6 +51,10 @@ from transformers import (
 from transformers.utils import get_full_repo_name
 from transformers.utils.versions import require_version
 
+from accelerate import Accelerator, DistributedType
+from accelerate.logging import get_logger
+from accelerate.utils import DummyOptim, DummyScheduler, set_seed
+
 
 logger = get_logger(__name__)
 
@@ -243,39 +243,6 @@ def parse_args():
     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()
@@ -302,9 +269,20 @@ def main():
     # 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
+
+    # when using DeepSpeed, the `gradient_accumulation_steps` is properly set from the DeepSpeed plugin/config
+    # or from `accelerate launch` via `--gradient_accumulation_steps`  else
+    # defaulting to the passed `args.gradient_accumulation_steps`
     accelerator = (
-        Accelerator(log_with=args.report_to, logging_dir=args.output_dir) if args.with_tracking else Accelerator()
+        Accelerator(
+            log_with=args.report_to,
+            project_dir=args.output_dir,
+            gradient_accumulation_steps=args.gradient_accumulation_steps,
+        )
+        if args.with_tracking
+        else Accelerator(gradient_accumulation_steps=args.gradient_accumulation_steps)
     )
+
     # Make one log on every process with the configuration for debugging.
     logging.basicConfig(
         format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
@@ -538,17 +516,11 @@ def main():
         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)
+    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / accelerator.gradient_accumulation_steps)
+    overrode_max_train_steps = False
     if args.max_train_steps is None:
         args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
+        overrode_max_train_steps = True
     else:
         args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
 
@@ -575,16 +547,16 @@ def main():
     )
 
     # 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
+    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / accelerator.gradient_accumulation_steps)
+    if overrode_max_train_steps:
+        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
+    # Afterwards we recalculate our number of training epochs
+    args.num_train_epochs = math.ceil(args.max_train_steps / 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
+    checkpointing_steps = args.checkpointing_steps
+    if checkpointing_steps is not None and checkpointing_steps.isdigit():
+        checkpointing_steps = int(checkpointing_steps)
 
     # We need to initialize the trackers we use, and also store our configuration.
     # The trackers initializes automatically on the main process.
@@ -595,14 +567,16 @@ def main():
         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
+    total_batch_size = (
+        args.per_device_train_batch_size * accelerator.num_processes * accelerator.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"  Gradient Accumulation steps = {accelerator.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)
@@ -613,45 +587,54 @@ def main():
 
     # 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.load_state(args.resume_from_checkpoint)
         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)
+        path = os.path.basename(args.resume_from_checkpoint)
+        training_difference = os.path.splitext(path)[0]
+
+        if "epoch" in training_difference:
+            starting_epoch = int(training_difference.replace("epoch_", "")) + 1
+            resume_step = None
+            completed_steps = starting_epoch * num_update_steps_per_epoch
+        else:
+            resume_step = int(training_difference.replace("step_", ""))
+            starting_epoch = resume_step // num_update_steps_per_epoch
+            resume_step -= starting_epoch * num_update_steps_per_epoch
+            completed_steps = resume_step
 
     for epoch in range(starting_epoch, args.num_train_epochs):
         model.train()
         if args.with_tracking:
             total_loss = 0
+
+        # skip new `skip_first_batches` to skip the batches when resuming from ckpt
+        if args.resume_from_checkpoint:
+            train_dataloader = accelerator.skip_first_batches(train_dataloader, num_batches=resume_step)
         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 + 1) % args.gradient_accumulation_steps == 0 or step == len(train_dataloader) - 1:
+            # In particular, DeepSpeed handles `gradient_accumulation` via `DeepSpeedEngine`.
+            # Below, we use `accelerator.accumulate` if the user
+            # wants to switch to other approaches such as plain DDP, PyTorch FSDP ...
+            # This avoids having to change any code as things are all handled across different distributed setups.
+            with accelerator.accumulate(model):
+                outputs = model(**batch)
+                loss = outputs.loss
+                accelerator.backward(loss)
                 optimizer.step()
                 lr_scheduler.step()
                 optimizer.zero_grad()
-                progress_bar.update(1)
-                completed_steps += 1
+
+                if accelerator.sync_gradients:
+                    progress_bar.update(1)
+                    completed_steps += 1
+
+            # We keep track of the loss at each epoch
+            if args.with_tracking:
+                step_loss = accelerator.reduce(loss.detach().clone()).item()
+                total_loss += step_loss
 
             if isinstance(checkpointing_steps, int):
                 if completed_steps % checkpointing_steps == 0:
-                    output_dir = f"step_{completed_steps }"
+                    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)
@@ -666,34 +649,29 @@ def main():
                 {
                     "perplexity": perplexity,
                     "eval_loss": eval_loss,
-                    "train_loss": total_loss.item() / len(train_dataloader),
+                    "train_loss": total_loss / 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)
+        if isinstance(checkpointing_steps, str) and checkpointing_steps == "epoch":
+            accelerator.save_state(os.path.join(args.output_dir, f"epoch_{epoch}"))
 
         # 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))
+            best_metric_checkpoint = os.path.join(args.output_dir, "best_checkpoint")
+            accelerator.save_state(best_metric_checkpoint)
             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},
-        )
+        accelerator.load_state(best_metric_checkpoint)
 
     # New Code #
     # Evaluates using the best checkpoint

examples/by_feature/early_stopping.py

@@ -0,0 +1,246 @@
+# 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 evaluate
+import torch
+from datasets import load_dataset
+from torch.optim import AdamW
+from torch.utils.data import DataLoader
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed
+
+from accelerate import Accelerator, DistributedType
+
+
+########################################################################
+# This is a fully working simple example to use Accelerate
+# specifically showcasing how to perform early stopping,
+# 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 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.
+        max_length = 128 if accelerator.distributed_type == DistributedType.TPU else None
+        # When using mixed precision we want round multiples of 8/16
+        if accelerator.mixed_precision == "fp8":
+            pad_to_multiple_of = 16
+        elif accelerator.mixed_precision != "no":
+            pad_to_multiple_of = 8
+        else:
+            pad_to_multiple_of = None
+
+        return tokenizer.pad(
+            examples,
+            padding="longest",
+            max_length=max_length,
+            pad_to_multiple_of=pad_to_multiple_of,
+            return_tensors="pt",
+        )
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(
+        tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size, drop_last=True
+    )
+    eval_dataloader = DataLoader(
+        tokenized_datasets["validation"],
+        shuffle=False,
+        collate_fn=collate_fn,
+        batch_size=EVAL_BATCH_SIZE,
+        drop_last=(accelerator.mixed_precision == "fp8"),
+    )
+
+    return train_dataloader, eval_dataloader
+
+
+# New code
+class EarlyStoppingCallback:
+    "A callback class that helps with early stopping"
+
+    def __init__(self, min_delta=0, patience=5):
+        self.min_delta = min_delta
+        self.patience = patience
+        self.counter = 0
+        self.lowest_loss = float("inf")
+
+    def check_early_stopping(self, eval_loss):
+        delta = self.lowest_loss - eval_loss
+        if delta >= self.min_delta:
+            self.lowest_loss = eval_loss
+            self.counter = 0
+        else:
+            self.counter += 1
+            if self.counter >= self.patience:
+                return True
+        return False
+
+
+callback = EarlyStoppingCallback()
+
+
+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)
+
+    # 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()
+
+            # New code
+            # Check if we should stop the training on any processes
+            if callback.check_early_stopping(loss.item()):
+                accelerator.set_trigger()
+
+            # If so, we break the loop
+            if accelerator.check_trigger():
+                break
+
+        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=None,
+        choices=["no", "fp16", "bf16", "fp8"],
+        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/fsdp_with_peak_mem_tracking.py

@@ -15,14 +15,23 @@
 import argparse
 import gc
 import os
-
-import torch
-from torch.utils.data import DataLoader
+import threading
 
 import evaluate
-from accelerate import Accelerator, DistributedType
+import psutil
+import torch
 from datasets import load_dataset
-from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed
+from torch.distributed.fsdp.fully_sharded_data_parallel import FullOptimStateDictConfig, FullStateDictConfig
+from torch.utils.data import DataLoader
+from transformers import (
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+    get_linear_schedule_with_warmup,
+    set_seed,
+)
+
+from accelerate import Accelerator, DistributedType, FullyShardedDataParallelPlugin
+from accelerate.utils import is_npu_available, is_xpu_available
 
 
 ########################################################################
@@ -60,18 +69,65 @@ def b2mb(x):
 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()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            torch.cuda.reset_max_memory_allocated()  # reset the peak gauge to zero
+            self.begin = torch.cuda.memory_allocated()
+        elif is_xpu_available():
+            torch.xpu.empty_cache()
+            torch.xpu.reset_max_memory_allocated()  # reset the peak gauge to zero
+            self.begin = torch.xpu.memory_allocated()
+        elif is_npu_available():
+            torch.npu.empty_cache()
+            torch.npu.reset_max_memory_allocated()  # reset the peak gauge to zero
+            self.begin = torch.npu.memory_allocated()
+        self.process = psutil.Process()
+
+        self.cpu_begin = self.cpu_mem_used()
+        self.peak_monitoring = True
+        peak_monitor_thread = threading.Thread(target=self.peak_monitor_func)
+        peak_monitor_thread.daemon = True
+        peak_monitor_thread.start()
         return self
 
+    def cpu_mem_used(self):
+        """get resident set size memory for the current process"""
+        return self.process.memory_info().rss
+
+    def peak_monitor_func(self):
+        self.cpu_peak = -1
+
+        while True:
+            self.cpu_peak = max(self.cpu_mem_used(), self.cpu_peak)
+
+            # can't sleep or will not catch the peak right (this comment is here on purpose)
+            # time.sleep(0.001) # 1msec
+
+            if not self.peak_monitoring:
+                break
+
     def __exit__(self, *exc):
+        self.peak_monitoring = False
+
         gc.collect()
-        torch.cuda.empty_cache()
-        self.end = torch.cuda.memory_allocated()
-        self.peak = torch.cuda.max_memory_allocated()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            self.end = torch.cuda.memory_allocated()
+            self.peak = torch.cuda.max_memory_allocated()
+        elif is_xpu_available():
+            torch.xpu.empty_cache()
+            self.end = torch.xpu.memory_allocated()
+            self.peak = torch.xpu.max_memory_allocated()
+        elif is_npu_available():
+            torch.npu.empty_cache()
+            self.end = torch.npu.memory_allocated()
+            self.peak = torch.npu.max_memory_allocated()
         self.used = b2mb(self.end - self.begin)
         self.peaked = b2mb(self.peak - self.begin)
+
+        self.cpu_end = self.cpu_mem_used()
+        self.cpu_used = b2mb(self.cpu_end - self.cpu_begin)
+        self.cpu_peaked = b2mb(self.cpu_peak - self.cpu_begin)
         # print(f"delta used/peak {self.used:4d}/{self.peaked:4d}")
 
 
@@ -86,13 +142,25 @@ def training_function(config, args):
     # For testing only
     if os.environ.get("TESTING_MOCKED_DATALOADERS", None) == "1":
         config["num_epochs"] = 2
+
+    # New Code #
+    # Pass the advanced FSDP settings not part of the accelerate config by creating fsdp_plugin
+    fsdp_plugin = FullyShardedDataParallelPlugin(
+        state_dict_config=FullStateDictConfig(offload_to_cpu=False, rank0_only=False),
+        optim_state_dict_config=FullOptimStateDictConfig(offload_to_cpu=False, rank0_only=False),
+    )
+
     # Initialize accelerator
     if args.with_tracking:
         accelerator = Accelerator(
-            cpu=args.cpu, mixed_precision=args.mixed_precision, log_with="wandb", logging_dir=args.logging_dir
+            cpu=args.cpu,
+            mixed_precision=args.mixed_precision,
+            log_with="wandb",
+            project_dir=args.logging_dir,
+            fsdp_plugin=fsdp_plugin,
         )
     else:
-        accelerator = Accelerator()
+        accelerator = Accelerator(fsdp_plugin=fsdp_plugin)
     accelerator.print(accelerator.distributed_type)
 
     if hasattr(args.checkpointing_steps, "isdigit"):
@@ -147,9 +215,22 @@ def training_function(config, args):
 
     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")
+        max_length = 128 if accelerator.distributed_type == DistributedType.TPU else None
+        # When using mixed precision we want round multiples of 8/16
+        if accelerator.mixed_precision == "fp8":
+            pad_to_multiple_of = 16
+        elif accelerator.mixed_precision != "no":
+            pad_to_multiple_of = 8
+        else:
+            pad_to_multiple_of = None
+
+        return tokenizer.pad(
+            examples,
+            padding="longest",
+            max_length=max_length,
+            pad_to_multiple_of=pad_to_multiple_of,
+            return_tensors="pt",
+        )
 
     # Instantiate dataloaders.
     train_dataloader = DataLoader(
@@ -162,7 +243,10 @@ def training_function(config, args):
     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)
+    model = AutoModelForSequenceClassification.from_pretrained(
+        args.model_name_or_path, return_dict=True, low_cpu_mem_usage=True
+    )
+
     # New Code #
     # For FSDP feature, it is highly recommended and efficient to prepare the model before creating optimizer
     model = accelerator.prepare(model)
@@ -232,7 +316,6 @@ def training_function(config, args):
                 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()
@@ -331,7 +414,7 @@ def main():
         "--mixed_precision",
         type=str,
         default=None,
-        choices=["no", "fp16", "bf16"],
+        choices=["no", "fp16", "bf16", "fp8"],
         help="Whether to use mixed precision. Choose"
         "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
         "and an Nvidia Ampere GPU.",
@@ -373,7 +456,7 @@ def main():
         required=True,
     )
     args = parser.parse_args()
-    config = {"lr": 2e-5, "num_epochs": 3, "seed": 1, "batch_size": 16}
+    config = {"lr": 2e-5, "num_epochs": 3, "seed": 42, "batch_size": 16}
     training_function(config, args)
 
 

examples/by_feature/gradient_accumulation.py

@@ -15,15 +15,15 @@
 import argparse
 import os
 
+import evaluate
 import torch
+from datasets import load_dataset
 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
 
+from accelerate import Accelerator, DistributedType
+
 
 ########################################################################
 # This is a fully working simple example to use Accelerate
@@ -81,9 +81,22 @@ def get_dataloaders(accelerator: Accelerator, batch_size: int = 16):
 
     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")
+        max_length = 128 if accelerator.distributed_type == DistributedType.TPU else None
+        # When using mixed precision we want round multiples of 8/16
+        if accelerator.mixed_precision == "fp8":
+            pad_to_multiple_of = 16
+        elif accelerator.mixed_precision != "no":
+            pad_to_multiple_of = 8
+        else:
+            pad_to_multiple_of = None
+
+        return tokenizer.pad(
+            examples,
+            padding="longest",
+            max_length=max_length,
+            pad_to_multiple_of=pad_to_multiple_of,
+            return_tensors="pt",
+        )
 
     # Instantiate dataloaders.
     train_dataloader = DataLoader(
@@ -193,7 +206,7 @@ def main():
         "--mixed_precision",
         type=str,
         default=None,
-        choices=["no", "fp16", "bf16"],
+        choices=["no", "fp16", "bf16", "fp8"],
         help="Whether to use mixed precision. Choose"
         "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
         "and an Nvidia Ampere GPU.",

examples/by_feature/local_sgd.py

@@ -0,0 +1,238 @@
+# coding=utf-8
+# Copyright 2023 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 evaluate
+import torch
+from datasets import load_dataset
+from torch.optim import AdamW
+from torch.utils.data import DataLoader
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed
+
+from accelerate import Accelerator, DistributedType
+from accelerate.local_sgd import LocalSGD
+
+
+########################################################################
+# This is a fully working simple example to use Accelerate
+# with LocalSGD, which is a method to synchronize model
+# parameters every K batches. It is different, but complementary
+# to 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.
+        max_length = 128 if accelerator.distributed_type == DistributedType.TPU else None
+        # When using mixed precision we want round multiples of 8/16
+        if accelerator.mixed_precision == "fp8":
+            pad_to_multiple_of = 16
+        elif accelerator.mixed_precision != "no":
+            pad_to_multiple_of = 8
+        else:
+            pad_to_multiple_of = None
+
+        return tokenizer.pad(
+            examples,
+            padding="longest",
+            max_length=max_length,
+            pad_to_multiple_of=pad_to_multiple_of,
+            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)
+    local_sgd_steps = int(args.local_sgd_steps)
+    # Initialize accelerator
+    accelerator = Accelerator(
+        cpu=args.cpu, mixed_precision=args.mixed_precision, gradient_accumulation_steps=gradient_accumulation_steps
+    )
+    if accelerator.distributed_type not in [DistributedType.NO, DistributedType.MULTI_CPU, DistributedType.MULTI_GPU]:
+        raise NotImplementedError("LocalSGD is supported only for CPUs and GPUs (no DeepSpeed or MegatronLM)")
+    # 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()
+        with LocalSGD(
+            accelerator=accelerator, model=model, local_sgd_steps=local_sgd_steps, enabled=local_sgd_steps is not None
+        ) as local_sgd:
+            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()
+                    # LocalSGD-specific line
+                    local_sgd.step()
+
+        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=None,
+        choices=["no", "fp16", "bf16", "fp8"],
+        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(
+        "--local_sgd_steps", type=int, default=8, help="Number of local SGD steps or None to disable local SGD"
+    )
+    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/megatron_lm_gpt_pretraining.py

@@ -31,16 +31,12 @@ import random
 from itertools import chain
 from pathlib import Path
 
-import torch
-from torch.utils.data import DataLoader
-
 import datasets
+import torch
 import transformers
-from accelerate import Accelerator, DistributedType
-from accelerate.logging import get_logger
-from accelerate.utils import MegatronLMDummyScheduler, set_seed
 from datasets import load_dataset
 from huggingface_hub import Repository
+from torch.utils.data import DataLoader
 from tqdm.auto import tqdm
 from transformers import (
     CONFIG_MAPPING,
@@ -55,6 +51,10 @@ from transformers import (
 from transformers.utils import check_min_version, get_full_repo_name, send_example_telemetry
 from transformers.utils.versions import require_version
 
+from accelerate import Accelerator, DistributedType
+from accelerate.logging import get_logger
+from accelerate.utils import MegatronLMDummyScheduler, set_seed
+
 
 # Will error if the minimal version of Transformers is not installed. Remove at your own risks.
 check_min_version("4.23.0.dev0")

examples/by_feature/memory.py

@@ -14,16 +14,16 @@
 import argparse
 import os
 
-import torch
-from torch.optim import AdamW
-from torch.utils.data import DataLoader
-
 # New Code #
 import evaluate
+import torch
+from datasets import load_dataset
+from torch.optim import AdamW
+from torch.utils.data import DataLoader
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed
+
 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
 
 
 ########################################################################
@@ -86,9 +86,22 @@ def get_dataloaders(accelerator: Accelerator, batch_size: int = 16):
 
     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")
+        max_length = 128 if accelerator.distributed_type == DistributedType.TPU else None
+        # When using mixed precision we want round multiples of 8/16
+        if accelerator.mixed_precision == "fp8":
+            pad_to_multiple_of = 16
+        elif accelerator.mixed_precision != "no":
+            pad_to_multiple_of = 8
+        else:
+            pad_to_multiple_of = None
+
+        return tokenizer.pad(
+            examples,
+            padding="longest",
+            max_length=max_length,
+            pad_to_multiple_of=pad_to_multiple_of,
+            return_tensors="pt",
+        )
 
     # Instantiate dataloaders.
     train_dataloader = DataLoader(
@@ -205,7 +218,7 @@ def main():
         "--mixed_precision",
         type=str,
         default=None,
-        choices=["no", "fp16", "bf16"],
+        choices=["no", "fp16", "bf16", "fp8"],
         help="Whether to use mixed precision. Choose"
         "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
         "and an Nvidia Ampere GPU.",

examples/by_feature/multi_process_metrics.py

@@ -15,15 +15,15 @@
 import argparse
 import os
 
+import evaluate
 import torch
+from datasets import load_dataset
 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
 
+from accelerate import Accelerator, DistributedType
+
 
 ########################################################################
 # This is a fully working simple example to use Accelerate,
@@ -88,9 +88,22 @@ def get_dataloaders(accelerator: Accelerator, batch_size: int = 16):
 
     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")
+        max_length = 128 if accelerator.distributed_type == DistributedType.TPU else None
+        # When using mixed precision we want round multiples of 8/16
+        if accelerator.mixed_precision == "fp8":
+            pad_to_multiple_of = 16
+        elif accelerator.mixed_precision != "no":
+            pad_to_multiple_of = 8
+        else:
+            pad_to_multiple_of = None
+
+        return tokenizer.pad(
+            examples,
+            padding="longest",
+            max_length=max_length,
+            pad_to_multiple_of=pad_to_multiple_of,
+            return_tensors="pt",
+        )
 
     # Instantiate dataloaders.
     train_dataloader = DataLoader(
@@ -210,7 +223,7 @@ def main():
         "--mixed_precision",
         type=str,
         default=None,
-        choices=["no", "fp16", "bf16"],
+        choices=["no", "fp16", "bf16", "fp8"],
         help="Whether to use mixed precision. Choose"
         "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
         "and an Nvidia Ampere GPU.",

examples/by_feature/tracking.py

@@ -15,15 +15,15 @@
 import argparse
 import os
 
+import evaluate
 import torch
+from datasets import load_dataset
 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
 
+from accelerate import Accelerator, DistributedType
+
 
 ########################################################################
 # This is a fully working simple example to use Accelerate,
@@ -86,9 +86,22 @@ def get_dataloaders(accelerator: Accelerator, batch_size: int = 16):
 
     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")
+        max_length = 128 if accelerator.distributed_type == DistributedType.TPU else None
+        # When using mixed precision we want round multiples of 8/16
+        if accelerator.mixed_precision == "fp8":
+            pad_to_multiple_of = 16
+        elif accelerator.mixed_precision != "no":
+            pad_to_multiple_of = 8
+        else:
+            pad_to_multiple_of = None
+
+        return tokenizer.pad(
+            examples,
+            padding="longest",
+            max_length=max_length,
+            pad_to_multiple_of=pad_to_multiple_of,
+            return_tensors="pt",
+        )
 
     # Instantiate dataloaders.
     train_dataloader = DataLoader(
@@ -120,7 +133,7 @@ def training_function(config, args):
     # >>> 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
+            cpu=args.cpu, mixed_precision=args.mixed_precision, log_with="all", project_dir=args.project_dir
         )
     else:
         accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)
@@ -237,7 +250,7 @@ def main():
         "--mixed_precision",
         type=str,
         default=None,
-        choices=["no", "fp16", "bf16"],
+        choices=["no", "fp16", "bf16", "fp8"],
         help="Whether to use mixed precision. Choose"
         "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
         "and an Nvidia Ampere GPU.",
@@ -249,10 +262,10 @@ def main():
         help="Whether to load in all available experiment trackers from the environment and use them for logging.",
     )
     parser.add_argument(
-        "--logging_dir",
+        "--project_dir",
         type=str,
         default="logs",
-        help="Location on where to store experiment tracking logs`",
+        help="Location on where to store experiment tracking logs` and relevent project information",
     )
     args = parser.parse_args()
     config = {"lr": 2e-5, "num_epochs": 3, "seed": 42, "batch_size": 16}

examples/complete_cv_example.py

@@ -17,15 +17,15 @@ import os
 import re
 
 import numpy as np
+import PIL
 import torch
+from timm import create_model
 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
 
+from accelerate import Accelerator
+
 
 ########################################################################
 # This is a fully working simple example to use Accelerate
@@ -75,7 +75,7 @@ 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
+            cpu=args.cpu, mixed_precision=args.mixed_precision, log_with="all", project_dir=args.project_dir
         )
     else:
         accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)
@@ -173,7 +173,7 @@ def training_function(config, args):
     )
     # 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
+    # We also need to keep track of the starting epoch so files are named properly
     starting_epoch = 0
 
     # Potentially load in the weights and states from a previous save
@@ -205,9 +205,12 @@ def training_function(config, args):
             total_loss = 0
         if args.resume_from_checkpoint and epoch == starting_epoch and resume_step is not None:
             # We need to skip steps until we reach the resumed step
-            train_dataloader = accelerator.skip_first_batches(train_dataloader, resume_step)
+            active_dataloader = accelerator.skip_first_batches(train_dataloader, resume_step)
             overall_step += resume_step
-        for batch in train_dataloader:
+        else:
+            # After the first iteration though, we need to go back to the original dataloader
+            active_dataloader = train_dataloader
+        for batch in active_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
@@ -272,7 +275,7 @@ def main():
         "--mixed_precision",
         type=str,
         default=None,
-        choices=["no", "fp16", "bf16"],
+        choices=["no", "fp16", "bf16", "fp8"],
         help="Whether to use mixed precision. Choose"
         "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
         "and an Nvidia Ampere GPU.",
@@ -302,10 +305,10 @@ def main():
         help="Whether to load in all available experiment trackers from the environment and use them for logging.",
     )
     parser.add_argument(
-        "--logging_dir",
+        "--project_dir",
         type=str,
         default="logs",
-        help="Location on where to store experiment tracking logs`",
+        help="Location on where to store experiment tracking logs` and relevent project information",
     )
     args = parser.parse_args()
     config = {"lr": 3e-2, "num_epochs": 3, "seed": 42, "batch_size": 64, "image_size": 224}

examples/complete_nlp_example.py

@@ -15,15 +15,15 @@
 import argparse
 import os
 
+import evaluate
 import torch
+from datasets import load_dataset
 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
 
+from accelerate import Accelerator, DistributedType
+
 
 ########################################################################
 # This is a fully working simple example to use Accelerate
@@ -52,7 +52,7 @@ 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
+            cpu=args.cpu, mixed_precision=args.mixed_precision, log_with="all", project_dir=args.project_dir
         )
     else:
         accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)
@@ -109,9 +109,22 @@ def training_function(config, args):
 
     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")
+        max_length = 128 if accelerator.distributed_type == DistributedType.TPU else None
+        # When using mixed precision we want round multiples of 8/16
+        if accelerator.mixed_precision == "fp8":
+            pad_to_multiple_of = 16
+        elif accelerator.mixed_precision != "no":
+            pad_to_multiple_of = 8
+        else:
+            pad_to_multiple_of = None
+
+        return tokenizer.pad(
+            examples,
+            padding="longest",
+            max_length=max_length,
+            pad_to_multiple_of=pad_to_multiple_of,
+            return_tensors="pt",
+        )
 
     # Instantiate dataloaders.
     train_dataloader = DataLoader(
@@ -182,9 +195,12 @@ def training_function(config, args):
             total_loss = 0
         if args.resume_from_checkpoint and epoch == starting_epoch and resume_step is not None:
             # We need to skip steps until we reach the resumed step
-            train_dataloader = accelerator.skip_first_batches(train_dataloader, resume_step)
+            active_dataloader = accelerator.skip_first_batches(train_dataloader, resume_step)
             overall_step += resume_step
-        for step, batch in enumerate(train_dataloader):
+        else:
+            # After the first iteration though, we need to go back to the original dataloader
+            active_dataloader = train_dataloader
+        for step, batch in enumerate(active_dataloader):
             # We could avoid this line since we set the accelerator with `device_placement=True`.
             batch.to(accelerator.device)
             outputs = model(**batch)
@@ -251,7 +267,7 @@ def main():
         "--mixed_precision",
         type=str,
         default=None,
-        choices=["no", "fp16", "bf16"],
+        choices=["no", "fp16", "bf16", "fp8"],
         help="Whether to use mixed precision. Choose"
         "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
         "and an Nvidia Ampere GPU.",
@@ -281,10 +297,10 @@ def main():
         help="Optional save directory where all checkpoint folders will be stored. Default is the current working directory.",
     )
     parser.add_argument(
-        "--logging_dir",
+        "--project_dir",
         type=str,
         default="logs",
-        help="Location on where to store experiment tracking logs`",
+        help="Location on where to store experiment tracking logs` and relevent project information",
     )
     args = parser.parse_args()
     config = {"lr": 2e-5, "num_epochs": 3, "seed": 42, "batch_size": 16}

examples/cv_example.py

@@ -17,15 +17,15 @@ import os
 import re
 
 import numpy as np
+import PIL
 import torch
+from timm import create_model
 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
 
+from accelerate import Accelerator
+
 
 ########################################################################
 # This is a fully working simple example to use Accelerate
@@ -190,7 +190,7 @@ def main():
         "--mixed_precision",
         type=str,
         default=None,
-        choices=["no", "fp16", "bf16"],
+        choices=["no", "fp16", "bf16", "fp8"],
         help="Whether to use mixed precision. Choose"
         "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
         "and an Nvidia Ampere GPU.",

examples/multigpu_remote_launcher.py

@@ -0,0 +1,55 @@
+import argparse
+
+import runhouse as rh
+import torch
+from nlp_example import training_function
+
+from accelerate.utils import PrepareForLaunch, patch_environment
+
+
+def launch_train(*args):
+    num_processes = torch.cuda.device_count()
+    print(f"Device count: {num_processes}")
+    with patch_environment(
+        world_size=num_processes, master_addr="127.0.01", master_port="29500", mixed_precision=args[1].mixed_precision
+    ):
+        launcher = PrepareForLaunch(training_function, distributed_type="MULTI_GPU")
+        torch.multiprocessing.start_processes(launcher, args=args, nprocs=num_processes, start_method="spawn")
+
+
+if __name__ == "__main__":
+    # Refer to https://runhouse-docs.readthedocs-hosted.com/en/main/rh_primitives/cluster.html#hardware-setup
+    # for cloud access setup instructions (if using on-demand hardware), and for API specifications.
+
+    # on-demand GPU
+    # gpu = rh.cluster(name='rh-cluster', instance_type='V100:1', provider='cheapest', use_spot=False)  # single GPU
+    gpu = rh.cluster(name="rh-cluster", instance_type="V100:4", provider="cheapest", use_spot=False)  # multi GPU
+    gpu.up_if_not()
+
+    # on-prem GPU
+    # gpu = rh.cluster(
+    #           ips=["ip_addr"], ssh_creds={ssh_user:"<username>", ssh_private_key:"<key_path>"}, name="rh-cluster"
+    #       )
+
+    # Set up remote function
+    reqs = [
+        "pip:./",
+        "transformers",
+        "datasets",
+        "evaluate",
+        "tqdm",
+        "scipy",
+        "scikit-learn",
+        "tensorboard",
+        "torch --upgrade --extra-index-url https://download.pytorch.org/whl/cu117",
+    ]
+    launch_train_gpu = rh.function(fn=launch_train, system=gpu, reqs=reqs, name="train_bert_glue")
+
+    # Define train args/config, run train function
+    train_args = argparse.Namespace(cpu=False, mixed_precision="fp16")
+    config = {"lr": 2e-5, "num_epochs": 3, "seed": 42, "batch_size": 16}
+    launch_train_gpu(config, train_args, stream_logs=True)
+
+    # Alternatively, we can just run as instructed in the README (but only because there's already a wrapper CLI):
+    # gpu.install_packages(reqs)
+    # gpu.run(['accelerate launch --multi_gpu accelerate/examples/nlp_example.py'])

examples/nlp_example.py

@@ -14,15 +14,15 @@
 # limitations under the License.
 import argparse
 
+import evaluate
 import torch
+from datasets import load_dataset
 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
 
+from accelerate import Accelerator, DistributedType
+
 
 ########################################################################
 # This is a fully working simple example to use Accelerate
@@ -79,16 +79,33 @@ def get_dataloaders(accelerator: Accelerator, batch_size: int = 16):
 
     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")
+        max_length = 128 if accelerator.distributed_type == DistributedType.TPU else None
+        # When using mixed precision we want round multiples of 8/16
+        if accelerator.mixed_precision == "fp8":
+            pad_to_multiple_of = 16
+        elif accelerator.mixed_precision != "no":
+            pad_to_multiple_of = 8
+        else:
+            pad_to_multiple_of = None
+
+        return tokenizer.pad(
+            examples,
+            padding="longest",
+            max_length=max_length,
+            pad_to_multiple_of=pad_to_multiple_of,
+            return_tensors="pt",
+        )
 
     # Instantiate dataloaders.
     train_dataloader = DataLoader(
-        tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size
+        tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size, drop_last=True
     )
     eval_dataloader = DataLoader(
-        tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
+        tokenized_datasets["validation"],
+        shuffle=False,
+        collate_fn=collate_fn,
+        batch_size=EVAL_BATCH_SIZE,
+        drop_last=(accelerator.mixed_precision == "fp8"),
     )
 
     return train_dataloader, eval_dataloader
@@ -120,7 +137,6 @@ def training_function(config, args):
     # 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)
 
@@ -134,6 +150,7 @@ def training_function(config, args):
     # 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
     )
@@ -177,7 +194,7 @@ def main():
         "--mixed_precision",
         type=str,
         default=None,
-        choices=["no", "fp16", "bf16"],
+        choices=["no", "fp16", "bf16", "fp8"],
         help="Whether to use mixed precision. Choose"
         "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
         "and an Nvidia Ampere GPU.",

pyproject.toml

@@ -1,3 +1,17 @@
 [tool.black]
 line-length = 119
-target-version = ['py36']
+target-version = ['py37']
+
+[tool.ruff]
+# Never enforce `E501` (line length violations).
+ignore = ["E501", "E741", "W605"]
+select = ["E", "F", "I", "W"]
+line-length = 119
+
+# Ignore import violations in all `__init__.py` files.
+[tool.ruff.per-file-ignores]
+"__init__.py" = ["E402", "F401", "F403", "F811"]
+
+[tool.ruff.isort]
+lines-after-imports = 2
+known-first-party = ["accelerate"]

setup.cfg

@@ -4,11 +4,6 @@ ensure_newline_before_comments = True
 force_grid_wrap = 0
 include_trailing_comma = True
 known_first_party = accelerate
-known_third_party =
-    numpy
-    torch
-    torch_xla
-
 line_length = 119
 lines_after_imports = 2
 multi_line_output = 3

setup.py

@@ -16,10 +16,10 @@ from setuptools import setup
 from setuptools import find_packages
 
 extras = {}
-extras["quality"] = ["black ~= 22.0", "isort >= 5.5.4", "flake8 >= 3.8.3", "hf-doc-builder >= 0.3.0"]
+extras["quality"] = ["black ~= 23.1", "ruff >= 0.0.241", "hf-doc-builder >= 0.3.0", "urllib3 < 2.0.0"]
 extras["docs"] = []
 extras["test_prod"] = ["pytest", "pytest-xdist", "pytest-subtests", "parameterized"]
-extras["test_dev"] = ["datasets", "evaluate", "transformers", "scipy", "scikit-learn", "deepspeed<0.7.0", "tqdm"]
+extras["test_dev"] = ["datasets", "evaluate", "transformers", "scipy", "scikit-learn", "deepspeed", "tqdm", "bitsandbytes", "timm"]
 extras["testing"] = extras["test_prod"] + extras["test_dev"]
 extras["rich"] = ["rich"]
 
@@ -32,7 +32,7 @@ extras["sagemaker"] = [
 
 setup(
     name="accelerate",
-    version="0.16.0.dev0",
+    version="0.23.0",
     description="Accelerate",
     long_description=open("README.md", "r", encoding="utf-8").read(),
     long_description_content_type="text/markdown",
@@ -47,11 +47,12 @@ setup(
         "console_scripts": [
             "accelerate=accelerate.commands.accelerate_cli:main",
             "accelerate-config=accelerate.commands.config:main",
+            "accelerate-estimate-memory=accelerate.commands.estimate:main",
             "accelerate-launch=accelerate.commands.launch:main",
         ]
     },
-    python_requires=">=3.7.0",
-    install_requires=["numpy>=1.17", "packaging>=20.0", "psutil", "pyyaml", "torch>=1.4.0"],
+    python_requires=">=3.8.0",
+    install_requires=["numpy>=1.17", "packaging>=20.0", "psutil", "pyyaml", "torch>=1.10.0", "huggingface_hub"],
     extras_require=extras,
     classifiers=[
         "Development Status :: 5 - Production/Stable",
@@ -61,27 +62,35 @@ setup(
         "License :: OSI Approved :: Apache Software License",
         "Operating System :: OS Independent",
         "Programming Language :: Python :: 3",
-        "Programming Language :: Python :: 3.7",
+        "Programming Language :: Python :: 3.8",
         "Topic :: Scientific/Engineering :: Artificial Intelligence",
     ],
 )
 
 # Release checklist
-# 1. Change the version in __init__.py and setup.py.
-# 2. Commit these changes with the message: "Release: VERSION"
-# 3. Add a tag in git to mark the release: "git tag VERSION -m 'Adds tag VERSION for pypi' "
-#    Push the tag to git: git push --tags origin main
-# 4. Run the following commands in the top-level directory:
+# 1. Checkout the release branch (for a patch the current release branch, for a new minor version, create one):
+#      git checkout -b vXX.xx-release
+#    The -b is only necessary for creation (so remove it when doing a patch)
+# 2. Change the version in __init__.py and setup.py to the proper value.
+# 3. Commit these changes with the message: "Release: v<VERSION>"
+# 4. Add a tag in git to mark the release:
+#      git tag v<VERSION> -m 'Adds tag v<VERSION> for pypi'
+#    Push the tag and release commit to git: git push --tags origin vXX.xx-release
+# 5. Run the following commands in the top-level directory:
+#      rm -rf dist
+#      rm -rf build
 #      python setup.py bdist_wheel
 #      python setup.py sdist
-# 5. Upload the package to the pypi test server first:
-#      twine upload dist/* -r pypitest
-#      twine upload dist/* -r pypitest --repository-url=https://test.pypi.org/legacy/
-# 6. Check that you can install it in a virtualenv by running:
+# 6. Upload the package to the pypi test server first:
+#      twine upload dist/* -r testpypi
+# 7. Check that you can install it in a virtualenv by running:
+#      pip install accelerate
+#      pip uninstall accelerate
 #      pip install -i https://testpypi.python.org/pypi accelerate
 #      accelerate env
 #      accelerate test
-# 7. Upload the final version to actual pypi:
+# 8. Upload the final version to actual pypi:
 #      twine upload dist/* -r pypi
-# 8. Add release notes to the tag in github once everything is looking hunky-dory.
-# 9. Update the version in __init__.py, setup.py to the new version "-dev" and push to master
+# 9. Add release notes to the tag in github once everything is looking hunky-dory.
+# 10. Go back to the main branch and update the version in __init__.py, setup.py to the new version ".dev" and push to
+#     main.

src/accelerate/__init__.py

@@ -1,12 +1,9 @@
-# flake8: noqa
-# There's no way to ignore "F401 '...' imported but unused" warnings in this
-# module, but to preserve other warnings. So, don't check this module at all.
-
-__version__ = "0.16.0.dev0"
+__version__ = "0.23.0"
 
 from .accelerator import Accelerator
 from .big_modeling import (
     cpu_offload,
+    cpu_offload_with_hook,
     disk_offload,
     dispatch_model,
     init_empty_weights,
@@ -15,7 +12,9 @@ from .big_modeling import (
 )
 from .data_loader import skip_first_batches
 from .launchers import debug_launcher, notebook_launcher
+from .state import PartialState
 from .utils import (
+    AutocastKwargs,
     DeepSpeedPlugin,
     DistributedDataParallelKwargs,
     DistributedType,

src/accelerate/big_modeling.py

@@ -12,14 +12,23 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import logging
 import os
 from contextlib import contextmanager
+from functools import wraps
 from typing import Dict, List, Optional, Union
 
 import torch
 import torch.nn as nn
 
-from .hooks import AlignDevicesHook, add_hook_to_module, attach_align_device_hook, attach_align_device_hook_on_blocks
+from .hooks import (
+    AlignDevicesHook,
+    CpuOffload,
+    UserCpuOffloadHook,
+    add_hook_to_module,
+    attach_align_device_hook,
+    attach_align_device_hook_on_blocks,
+)
 from .utils import (
     OffloadedWeightsLoader,
     check_device_map,
@@ -27,21 +36,25 @@ from .utils import (
     find_tied_parameters,
     get_balanced_memory,
     infer_auto_device_map,
+    is_torch_version,
     load_checkpoint_in_model,
     offload_state_dict,
+    parse_flag_from_env,
     retie_parameters,
 )
-from .utils.versions import is_torch_version
+
+
+logger = logging.getLogger(__name__)
 
 
 @contextmanager
-def init_empty_weights(include_buffers: bool = False):
+def init_empty_weights(include_buffers: bool = None):
     """
     A context manager under which models are initialized with all parameters on the meta device, therefore creating an
     empty model. Useful when just initializing the model would blow the available RAM.
 
     Args:
-        include_buffers (`bool`, *optional*, defaults to `False`):
+        include_buffers (`bool`, *optional*):
             Whether or not to also put all buffers on the meta device while initializing.
 
     Example:
@@ -62,21 +75,21 @@ def init_empty_weights(include_buffers: bool = False):
 
     </Tip>
     """
-    if not is_torch_version(">=", "1.9.0"):
-        raise NotImplementedError("Initializing empty weights to a meta device requires torch >= 1.9.0")
+    if include_buffers is None:
+        include_buffers = parse_flag_from_env("ACCELERATE_INIT_INCLUDE_BUFFERS", False)
     with init_on_device(torch.device("meta"), include_buffers=include_buffers) as f:
         yield f
 
 
 @contextmanager
-def init_on_device(device: torch.device, include_buffers: bool = False):
+def init_on_device(device: torch.device, include_buffers: bool = None):
     """
     A context manager under which models are initialized with all parameters on the specified device.
 
     Args:
         device (`torch.device`):
             Device to initialize all parameters on.
-        include_buffers (`bool`, *optional*, defaults to `False`):
+        include_buffers (`bool`, *optional*):
             Whether or not to also put all buffers on the meta device while initializing.
 
     Example:
@@ -89,6 +102,15 @@ def init_on_device(device: torch.device, include_buffers: bool = False):
         tst = nn.Liner(100, 100)  # on `cuda` device
     ```
     """
+    if include_buffers is None:
+        include_buffers = parse_flag_from_env("ACCELERATE_INIT_INCLUDE_BUFFERS", False)
+
+    # TODO(shingjan): remove the torch version check once older versions are deprecated
+    if is_torch_version(">=", "2.0") and include_buffers:
+        with device:
+            yield
+        return
+
     old_register_parameter = nn.Module.register_parameter
     if include_buffers:
         old_register_buffer = nn.Module.register_buffer
@@ -100,8 +122,8 @@ def init_on_device(device: torch.device, include_buffers: bool = False):
             kwargs = module._parameters[name].__dict__
             module._parameters[name] = param_cls(module._parameters[name].to(device), **kwargs)
 
-    def register_empty_buffer(module, name, buffer):
-        old_register_buffer(module, name, buffer)
+    def register_empty_buffer(module, name, buffer, persistent=True):
+        old_register_buffer(module, name, buffer, persistent=persistent)
         if buffer is not None:
             module._buffers[name] = module._buffers[name].to(device)
 
@@ -164,8 +186,6 @@ def cpu_offload(
             called directly during the forward, for instance if a `dense` linear layer is registered, but at forward,
             `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly.
     """
-    if not is_torch_version(">=", "1.9.0"):
-        raise NotImplementedError("CPU offloading requires torch >= 1.9.0")
     if execution_device is None:
         execution_device = next(iter(model.parameters())).device
     if state_dict is None:
@@ -184,6 +204,50 @@ def cpu_offload(
     return model
 
 
+def cpu_offload_with_hook(
+    model: torch.nn.Module,
+    execution_device: Optional[Union[int, str, torch.device]] = None,
+    prev_module_hook: Optional[UserCpuOffloadHook] = None,
+):
+    """
+    Offloads a model on the CPU and puts it back to an execution device when executed. The difference with
+    [`cpu_offload`] is that the model stays on the execution device after the forward and is only offloaded again when
+    the `offload` method of the returned `hook` is called. Useful for pipelines running a model in a loop.
+
+    Args:
+        model (`torch.nn.Module`):
+            The model to offload.
+        execution_device(`str`, `int` or `torch.device`, *optional*):
+            The device on which the model should be executed. Will default to the MPS device if it's available, then
+            GPU 0 if there is a GPU, and finally to the CPU.
+        prev_module_hook (`UserCpuOffloadHook`, *optional*):
+            The hook sent back by this function for a previous model in the pipeline you are running. If passed, its
+            offload method will be called just before the forward of the model to which this hook is attached.
+
+    Example:
+
+    ```py
+    model_1, hook_1 = cpu_offload_with_hook(model_1, cuda_device)
+    model_2, hook_2 = cpu_offload_with_hook(model_2, cuda_device, prev_module_hook=hook_1)
+    model_3, hook_3 = cpu_offload_with_hook(model_3, cuda_device, prev_module_hook=hook_2)
+
+    hid_1 = model_1(input)
+    for i in range(50):
+        # model1 is offloaded on the CPU at the first iteration, model 2 stays on the GPU for this whole loop.
+        hid_2 = model_2(hid_1)
+    # model2 is offloaded to the CPU just before this forward.
+    hid_3 = model_3(hid_3)
+
+    # For model3, you need to manually call the hook offload method.
+    hook_3.offload()
+    ```
+    """
+    hook = CpuOffload(execution_device=execution_device, prev_module_hook=prev_module_hook)
+    add_hook_to_module(model, hook, append=True)
+    user_hook = UserCpuOffloadHook(model, hook)
+    return model, user_hook
+
+
 def disk_offload(
     model: nn.Module,
     offload_dir: Union[str, os.PathLike],
@@ -211,8 +275,6 @@ def disk_offload(
             called directly during the forward, for instance if a `dense` linear layer is registered, but at forward,
             `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly.
     """
-    if not is_torch_version(">=", "1.9.0"):
-        raise NotImplementedError("Disk offloading requires torch >= 1.9.0")
     if not os.path.isdir(offload_dir) or not os.path.isfile(os.path.join(offload_dir, "index.json")):
         offload_state_dict(offload_dir, model.state_dict())
     if execution_device is None:
@@ -240,7 +302,9 @@ def dispatch_model(
     offload_dir: Optional[Union[str, os.PathLike]] = None,
     offload_index: Optional[Dict[str, str]] = None,
     offload_buffers: bool = False,
+    skip_keys: Optional[Union[str, List[str]]] = None,
     preload_module_classes: Optional[List[str]] = None,
+    force_hooks: bool = False,
 ):
     """
     Dispatches a model according to a given device map. Layers of the model might be spread across GPUs, offloaded on
@@ -264,67 +328,106 @@ def dispatch_model(
             to the index saved in `save_folder`.
         offload_buffers (`bool`, *optional*, defaults to `False`):
             Whether or not to offload the buffers with the model parameters.
+        skip_keys (`str` or `List[str]`, *optional*):
+            A list of keys to ignore when moving inputs or outputs between devices.
         preload_module_classes (`List[str]`, *optional*):
             A list of classes whose instances should load all their weights (even in the submodules) at the beginning
             of the forward. This should only be used for classes that have submodules which are registered but not
             called directly during the forward, for instance if a `dense` linear layer is registered, but at forward,
             `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly.
+        force_hooks (`bool`, *optional*, defaults to `False`):
+            Whether or not to force device hooks to be attached to the model even if all layers are dispatched to a
+            single device.
     """
-    if not is_torch_version(">=", "1.9.0"):
-        raise NotImplementedError("Model dispatching requires torch >= 1.9.0")
     # Error early if the device map is incomplete.
     check_device_map(model, device_map)
 
-    if main_device is None:
-        if set(device_map.values()) == {"cpu"} or set(device_map.values()) == {"cpu", "disk"}:
-            main_device = "cpu"
+    # for backward compatibility
+    is_bnb_quantized = (
+        getattr(model, "is_quantized", False) or getattr(model, "is_loaded_in_8bit", False)
+    ) and getattr(model, "quantization_method", "bitsandbytes") == "bitsandbytes"
+
+    # We attach hooks if the device_map has at least 2 different devices or if
+    # force_hooks is set to `True`. Otherwise, the model in already loaded
+    # in the unique device and the user can decide where to dispatch the model.
+    # If the model is quantized, we always force-dispatch the model
+    if (len(set(device_map.values())) > 1) or is_bnb_quantized or force_hooks:
+        if main_device is None:
+            if set(device_map.values()) == {"cpu"} or set(device_map.values()) == {"cpu", "disk"}:
+                main_device = "cpu"
+            else:
+                main_device = [d for d in device_map.values() if d not in ["cpu", "disk"]][0]
+
+        if main_device != "cpu":
+            cpu_modules = [name for name, device in device_map.items() if device == "cpu"]
+            if state_dict is None and len(cpu_modules) > 0:
+                state_dict = extract_submodules_state_dict(model.state_dict(), cpu_modules)
+
+        disk_modules = [name for name, device in device_map.items() if device == "disk"]
+        if offload_dir is None and offload_index is None and len(disk_modules) > 0:
+            raise ValueError(
+                "We need an `offload_dir` to dispatch this model according to this `device_map`, the following submodules "
+                f"need to be offloaded: {', '.join(disk_modules)}."
+            )
+        if (
+            len(disk_modules) > 0
+            and offload_index is None
+            and (not os.path.isdir(offload_dir) or not os.path.isfile(os.path.join(offload_dir, "index.json")))
+        ):
+            disk_state_dict = extract_submodules_state_dict(model.state_dict(), disk_modules)
+            offload_state_dict(offload_dir, disk_state_dict)
+
+        execution_device = {
+            name: main_device if device in ["cpu", "disk"] else device for name, device in device_map.items()
+        }
+        execution_device[""] = main_device
+        offloaded_devices = ["disk"] if main_device == "cpu" or main_device == "mps" else ["cpu", "disk"]
+        offload = {name: device in offloaded_devices for name, device in device_map.items()}
+        save_folder = offload_dir if len(disk_modules) > 0 else None
+        if state_dict is not None or save_folder is not None or offload_index is not None:
+            device = main_device if offload_index is not None else None
+            weights_map = OffloadedWeightsLoader(
+                state_dict=state_dict, save_folder=save_folder, index=offload_index, device=device
+            )
         else:
-            main_device = [d for d in device_map.values() if d not in ["cpu", "disk"]][0]
+            weights_map = None
 
-    if main_device != "cpu":
-        cpu_modules = [name for name, device in device_map.items() if device == "cpu"]
-        if state_dict is None and len(cpu_modules) > 0:
-            state_dict = extract_submodules_state_dict(model.state_dict(), cpu_modules)
-
-    disk_modules = [name for name, device in device_map.items() if device == "disk"]
-    if offload_dir is None and offload_index is None and len(disk_modules) > 0:
-        raise ValueError(
-            "We need an `offload_dir` to dispatch this model according to this `device_map`, the following submodules "
-            f"need to be offloaded: {', '.join(disk_modules)}."
+        tied_params = find_tied_parameters(model)
+        attach_align_device_hook_on_blocks(
+            model,
+            execution_device=execution_device,
+            offload=offload,
+            offload_buffers=offload_buffers,
+            weights_map=weights_map,
+            skip_keys=skip_keys,
+            preload_module_classes=preload_module_classes,
         )
-    if (
-        len(disk_modules) > 0
-        and offload_index is None
-        and (not os.path.isdir(offload_dir) or not os.path.isfile(os.path.join(offload_dir, "index.json")))
-    ):
-        disk_state_dict = extract_submodules_state_dict(model.state_dict(), disk_modules)
-        offload_state_dict(offload_dir, disk_state_dict)
+        # Attaching the hook may break tied weights, so we retie them
+        retie_parameters(model, tied_params)
+
+        # add warning to cuda and to method
+        def add_warning(fn, model):
+            @wraps(fn)
+            def wrapper(*args, **kwargs):
+                logger.warning("You shouldn't move a model when it is dispatched on multiple devices.")
+                for param in model.parameters():
+                    if param.device == torch.device("meta"):
+                        raise RuntimeError("You can't move a model that has some modules offloaded to cpu or disk.")
+                return fn(*args, **kwargs)
+
+            return wrapper
+
+        model.to = add_warning(model.to, model)
+        model.cuda = add_warning(model.cuda, model)
 
-    execution_device = {
-        name: main_device if device in ["cpu", "disk"] else device for name, device in device_map.items()
-    }
-    offloaded_devices = ["disk"] if main_device == "cpu" else ["cpu", "disk"]
-    offload = {name: device in offloaded_devices for name, device in device_map.items()}
-    save_folder = offload_dir if len(disk_modules) > 0 else None
-    if state_dict is not None or save_folder is not None or offload_index is not None:
-        device = main_device if offload_index is not None else None
-        weights_map = OffloadedWeightsLoader(
-            state_dict=state_dict, save_folder=save_folder, index=offload_index, device=device
-        )
     else:
-        weights_map = None
-
-    tied_params = find_tied_parameters(model)
-    attach_align_device_hook_on_blocks(
-        model,
-        execution_device=execution_device,
-        offload=offload,
-        offload_buffers=offload_buffers,
-        weights_map=weights_map,
-        preload_module_classes=preload_module_classes,
-    )
-    # Attaching the hook may break tied weights, so we retie them
-    retie_parameters(model, tied_params)
+        device = list(device_map.values())[0]
+        if device != "disk":
+            model.to(device)
+        else:
+            raise ValueError(
+                "You are trying to offload the whole model to the disk. Please use the `disk_offload` function instead."
+            )
     model.hf_device_map = device_map
     return model
 
@@ -339,7 +442,9 @@ def load_checkpoint_and_dispatch(
     offload_buffers: bool = False,
     dtype: Optional[Union[str, torch.dtype]] = None,
     offload_state_dict: Optional[bool] = None,
+    skip_keys: Optional[Union[str, List[str]]] = None,
     preload_module_classes: Optional[List[str]] = None,
+    force_hooks: bool = False,
 ):
     """
     Loads a (potentially sharded) checkpoint inside a model, potentially sending weights to a given device as they are
@@ -375,11 +480,16 @@ def load_checkpoint_and_dispatch(
             If `True`, will temporarily offload the CPU state dict on the hard drive to avoid getting out of CPU RAM if
             the weight of the CPU state dict + the biggest shard does not fit. Will default to `True` if the device map
             picked contains `"disk"` values.
+        skip_keys (`str` or `List[str]`, *optional*):
+            A list of keys to ignore when moving inputs or outputs between devices.
         preload_module_classes (`List[str]`, *optional*):
             A list of classes whose instances should load all their weights (even in the submodules) at the beginning
             of the forward. This should only be used for classes that have submodules which are registered but not
             called directly during the forward, for instance if a `dense` linear layer is registered, but at forward,
             `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly.
+        force_hooks (`bool`, *optional*, defaults to `False`):
+            Whether or not to force device hooks to be attached to the model even if all layers are dispatched to a
+            single device.
 
     Example:
 
@@ -403,26 +513,24 @@ def load_checkpoint_and_dispatch(
     ... )
     ```
     """
-    if not is_torch_version(">=", "1.9.0"):
-        raise NotImplementedError("Loading and dispatching requires torch >= 1.9.0")
     if isinstance(device_map, str) and device_map not in ["auto", "balanced", "balanced_low_0", "sequential"]:
         raise ValueError(
             "If passing a string for `device_map`, please choose 'auto', 'balanced', 'balanced_low_0' or "
             "'sequential'."
         )
-    if device_map != "sequential":
-        max_memory = get_balanced_memory(
-            model,
-            max_memory=max_memory,
-            no_split_module_classes=no_split_module_classes,
-            dtype=dtype,
-            low_zero=(device_map == "balanced_low_0"),
-        )
     if isinstance(device_map, str):
+        if device_map != "sequential":
+            max_memory = get_balanced_memory(
+                model,
+                max_memory=max_memory,
+                no_split_module_classes=no_split_module_classes,
+                dtype=dtype,
+                low_zero=(device_map == "balanced_low_0"),
+            )
         device_map = infer_auto_device_map(
             model, max_memory=max_memory, no_split_module_classes=no_split_module_classes, dtype=dtype
         )
-    if offload_state_dict is None and "disk" in device_map.values():
+    if offload_state_dict is None and device_map is not None and "disk" in device_map.values():
         offload_state_dict = True
     load_checkpoint_in_model(
         model,
@@ -440,5 +548,7 @@ def load_checkpoint_and_dispatch(
         device_map=device_map,
         offload_dir=offload_folder,
         offload_buffers=offload_buffers,
+        skip_keys=skip_keys,
         preload_module_classes=preload_module_classes,
+        force_hooks=force_hooks,
     )

src/accelerate/checkpointing.py

@@ -29,6 +29,7 @@ from .utils import (
     SCHEDULER_NAME,
     get_pretty_name,
     is_tpu_available,
+    is_xpu_available,
     save,
 )
 
@@ -37,6 +38,7 @@ if is_tpu_available(check_device=False):
     import torch_xla.core.xla_model as xm
 
 from .logging import get_logger
+from .state import PartialState
 
 
 logger = get_logger(__name__)
@@ -99,8 +101,10 @@ def save_accelerator_state(
     states["random_state"] = random.getstate()
     states["numpy_random_seed"] = np.random.get_state()
     states["torch_manual_seed"] = torch.get_rng_state()
-    states["torch_cuda_manual_seed"] = torch.cuda.get_rng_state_all()
-    # ^^ safe to call this function even if cuda is not available
+    if is_xpu_available():
+        states["torch_xpu_manual_seed"] = torch.xpu.get_rng_state_all()
+    else:
+        states["torch_cuda_manual_seed"] = torch.cuda.get_rng_state_all()
     if is_tpu_available():
         states["xm_seed"] = xm.get_rng_state()
     output_states_file = os.path.join(output_dir, states_name)
@@ -110,7 +114,14 @@ def save_accelerator_state(
 
 
 def load_accelerator_state(
-    input_dir, models, optimizers, schedulers, process_index, scaler=None, **load_model_func_kwargs
+    input_dir,
+    models,
+    optimizers,
+    schedulers,
+    process_index,
+    scaler=None,
+    map_location=None,
+    **load_model_func_kwargs,
 ):
     """
     Loads states of the models, optimizers, scaler, and RNG generators from a given directory.
@@ -128,21 +139,32 @@ def load_accelerator_state(
             The current process index in the Accelerator state
         scaler (`torch.cuda.amp.GradScaler`, *optional*):
             An optional *GradScaler* instance to load
+        map_location (`str`, *optional*):
+            What device to load the optimizer state onto. Should be one of either "cpu" or "on_device".
         load_model_func_kwargs (`dict`, *optional*):
             Additional arguments that can be passed to the model's `load_state_dict` method.
     """
+    if map_location not in [None, "cpu", "on_device"]:
+        raise TypeError(
+            "Unsupported optimizer map location passed, please choose one of `None`, `'cpu'`, or `'on_device'`"
+        )
+    if map_location is None:
+        map_location = "cpu"
+    elif map_location == "on_device":
+        map_location = PartialState().device
     # Model states
     for i, model in enumerate(models):
         weights_name = f"{MODEL_NAME}.bin" if i == 0 else f"{MODEL_NAME}_{i}.bin"
         input_model_file = os.path.join(input_dir, weights_name)
-        models[i].load_state_dict(torch.load(input_model_file, map_location="cpu"), **load_model_func_kwargs)
+        models[i].load_state_dict(torch.load(input_model_file, map_location=map_location), **load_model_func_kwargs)
     logger.info("All model weights loaded successfully")
 
     # Optimizer states
     for i, opt in enumerate(optimizers):
         optimizer_name = f"{OPTIMIZER_NAME}.bin" if i == 0 else f"{OPTIMIZER_NAME}_{i}.bin"
         input_optimizer_file = os.path.join(input_dir, optimizer_name)
-        optimizers[i].load_state_dict(torch.load(input_optimizer_file, map_location="cpu"))
+        optimizer_state = torch.load(input_optimizer_file, map_location=map_location)
+        optimizers[i].load_state_dict(optimizer_state)
     logger.info("All optimizer states loaded successfully")
 
     # Scheduler states
@@ -164,12 +186,14 @@ def load_accelerator_state(
         random.setstate(states["random_state"])
         np.random.set_state(states["numpy_random_seed"])
         torch.set_rng_state(states["torch_manual_seed"])
-        torch.cuda.set_rng_state_all(states["torch_cuda_manual_seed"])
-        # ^^ safe to call this function even if cuda is not available
+        if is_xpu_available():
+            torch.xpu.set_rng_state_all(states["torch_xpu_manual_seed"])
+        else:
+            torch.cuda.set_rng_state_all(states["torch_cuda_manual_seed"])
         if is_tpu_available():
             xm.set_rng_state(states["xm_seed"])
         logger.info("All random states loaded successfully")
-    except:
+    except Exception:
         logger.info("Could not load random states")
 
 
@@ -189,4 +213,4 @@ def load_custom_state(obj, path, index: int = 0):
     """
     load_location = f"{path}/custom_checkpoint_{index}.pkl"
     logger.info(f"Loading the state of {get_pretty_name(obj)} from {load_location}")
-    obj.load_state_dict(torch.load(load_location))
+    obj.load_state_dict(torch.load(load_location, map_location="cpu"))

src/accelerate/commands/accelerate_cli.py

@@ -18,17 +18,19 @@ from argparse import ArgumentParser
 
 from accelerate.commands.config import get_config_parser
 from accelerate.commands.env import env_command_parser
+from accelerate.commands.estimate import estimate_command_parser
 from accelerate.commands.launch import launch_command_parser
 from accelerate.commands.test import test_command_parser
 from accelerate.commands.tpu import tpu_command_parser
 
 
 def main():
-    parser = ArgumentParser("Accelerate CLI tool", usage="accelerate <command> [<args>]")
+    parser = ArgumentParser("Accelerate CLI tool", usage="accelerate <command> [<args>]", allow_abbrev=False)
     subparsers = parser.add_subparsers(help="accelerate command helpers")
 
     # Register commands
     get_config_parser(subparsers=subparsers)
+    estimate_command_parser(subparsers=subparsers)
     env_command_parser(subparsers=subparsers)
     launch_command_parser(subparsers=subparsers)
     tpu_command_parser(subparsers=subparsers)

src/accelerate/commands/config/__init__.py

@@ -23,7 +23,7 @@ from .update import update_command_parser
 
 
 def get_config_parser(subparsers=None):
-    parent_parser = argparse.ArgumentParser(add_help=False)
+    parent_parser = argparse.ArgumentParser(add_help=False, allow_abbrev=False)
     # The main config parser
     config_parser = config_command_parser(subparsers)
     # The subparser to add commands to

src/accelerate/commands/config/cluster.py

@@ -19,9 +19,11 @@ import os
 from ...utils import (
     ComputeEnvironment,
     DistributedType,
-    DynamoBackend,
     is_deepspeed_available,
+    is_mps_available,
+    is_npu_available,
     is_transformers_available,
+    is_xpu_available,
 )
 from ...utils.constants import (
     DEEPSPEED_MULTINODE_LAUNCHERS,
@@ -29,9 +31,11 @@ from ...utils.constants import (
     FSDP_BACKWARD_PREFETCH,
     FSDP_SHARDING_STRATEGY,
     FSDP_STATE_DICT_TYPE,
+    TORCH_DYNAMO_MODES,
 )
 from .config_args import ClusterConfig
 from .config_utils import (
+    DYNAMO_BACKENDS,
     _ask_field,
     _ask_options,
     _convert_distributed_mode,
@@ -44,7 +48,7 @@ from .config_utils import (
 def get_cluster_input():
     distributed_type = _ask_options(
         "Which type of machine are you using?",
-        ["No distributed training", "multi-CPU", "multi-GPU", "TPU", "MPS"],
+        ["No distributed training", "multi-CPU", "multi-XPU", "multi-GPU", "multi-NPU", "TPU"],
         _convert_distributed_mode,
     )
 
@@ -56,11 +60,14 @@ def get_cluster_input():
     main_process_port = None
     rdzv_backend = "static"
     same_network = True
-    tpu_name = None
-    tpu_zone = None
-    commands = None
-    command_file = None
-    if distributed_type in [DistributedType.MULTI_GPU, DistributedType.MULTI_CPU]:
+    debug = False
+
+    if distributed_type in [
+        DistributedType.MULTI_GPU,
+        DistributedType.MULTI_NPU,
+        DistributedType.MULTI_XPU,
+        DistributedType.MULTI_CPU,
+    ]:
         num_machines = _ask_field(
             "How many different machines will you use (use more than 1 for multi-node training)? [1]: ",
             int,
@@ -89,10 +96,16 @@ def get_cluster_input():
                 rdzv_backend = _ask_field(
                     "What rendezvous backend will you use? ('static', 'c10d', ...): ", default="static"
                 )
+        debug = _ask_field(
+            "Should distributed operations be checked while running for errors? This can avoid timeout issues but will be slower. [yes/NO]: ",
+            _convert_yes_no_to_bool,
+            default=False,
+            error_message="Please enter yes or no.",
+        )
 
     if distributed_type == DistributedType.NO:
         use_cpu = _ask_field(
-            "Do you want to run your training on CPU only (even if a GPU is available)? [yes/NO]:",
+            "Do you want to run your training on CPU only (even if a GPU / Apple Silicon / Ascend NPU device is available)? [yes/NO]:",
             _convert_yes_no_to_bool,
             default=False,
             error_message="Please enter yes or no.",
@@ -102,6 +115,27 @@ def get_cluster_input():
     else:
         use_cpu = False
 
+    ipex_config = {}
+    if use_cpu:
+        ipex_config["ipex"] = _ask_field(
+            "Do you want to use Intel PyTorch Extension (IPEX) to speed up training on CPU? [yes/NO]:",
+            _convert_yes_no_to_bool,
+            default=False,
+            error_message="Please enter yes or no.",
+        )
+    if (
+        not use_cpu
+        and is_xpu_available()
+        and distributed_type not in [DistributedType.MULTI_GPU, DistributedType.MULTI_NPU, DistributedType.TPU]
+    ):
+        ipex_config["use_xpu"] = _ask_field(
+            "Do you want to use XPU plugin to speed up training on XPU? [yes/NO]:",
+            _convert_yes_no_to_bool,
+            default=False,
+            error_message="Please enter yes or no.",
+        )
+
+    dynamo_config = {}
     use_dynamo = _ask_field(
         "Do you wish to optimize your script with torch dynamo?[yes/NO]:",
         _convert_yes_no_to_bool,
@@ -109,28 +143,43 @@ def get_cluster_input():
         error_message="Please enter yes or no.",
     )
     if use_dynamo:
-        dynamo_backend = _ask_options(
+        prefix = "dynamo_"
+        dynamo_config[prefix + "backend"] = _ask_options(
             "Which dynamo backend would you like to use?",
-            [
-                "eager",
-                "aot_eager",
-                "inductor",
-                "nvfuser",
-                "aot_nvfuser",
-                "aot_cudagraphs",
-                "ofi",
-                "fx2trt",
-                "onnxrt",
-                "ipex",
-            ],
+            [x.lower() for x in DYNAMO_BACKENDS],
             _convert_dynamo_backend,
             default=2,
         )
-    else:
-        dynamo_backend = DynamoBackend.NO
+        use_custom_options = _ask_field(
+            "Do you want to customize the defaults sent to torch.compile? [yes/NO]: ",
+            _convert_yes_no_to_bool,
+            default=False,
+            error_message="Please enter yes or no.",
+        )
 
+        if use_custom_options:
+            dynamo_config[prefix + "mode"] = _ask_options(
+                "Which mode do you want to use?",
+                TORCH_DYNAMO_MODES,
+                lambda x: TORCH_DYNAMO_MODES[int(x)],
+                default=0,
+            )
+            dynamo_config[prefix + "use_fullgraph"] = _ask_field(
+                "Do you want the fullgraph mode or it is ok to break model into several subgraphs? [yes/NO]: ",
+                _convert_yes_no_to_bool,
+                default=False,
+                error_message="Please enter yes or no.",
+            )
+            dynamo_config[prefix + "use_dynamic"] = _ask_field(
+                "Do you want to enable dynamic shape tracing? [yes/NO]: ",
+                _convert_yes_no_to_bool,
+                default=False,
+                error_message="Please enter yes or no.",
+            )
+
+    use_mps = not use_cpu and is_mps_available()
     deepspeed_config = {}
-    if distributed_type in [DistributedType.MULTI_GPU, DistributedType.NO]:
+    if distributed_type in [DistributedType.MULTI_GPU, DistributedType.NO] and not use_mps:
         use_deepspeed = _ask_field(
             "Do you want to use DeepSpeed? [yes/NO]: ",
             _convert_yes_no_to_bool,
@@ -172,6 +221,18 @@ def get_cluster_input():
                     deepspeed_config["offload_param_device"] = _ask_options(
                         "Where to offload parameters?", deepspeed_devices, lambda x: deepspeed_devices[int(x)]
                     )
+                    if deepspeed_config["offload_param_device"] == "nvme":
+                        deepspeed_config["offload_param_nvme_path"] = _ask_field(
+                            "Nvme Path to offload parameters?",
+                            str,
+                            default="/nvme",
+                        )
+                    if deepspeed_config["offload_optimizer_device"] == "nvme":
+                        deepspeed_config["offload_optimizer_nvme_path"] = _ask_field(
+                            "Nvme Path to offload optimizer states?",
+                            str,
+                            default="/nvme",
+                        )
                 deepspeed_config["gradient_accumulation_steps"] = _ask_field(
                     "How many gradient accumulation steps you're passing in your script? [1]: ",
                     int,
@@ -252,7 +313,7 @@ def get_cluster_input():
                         )
 
     fsdp_config = {}
-    if distributed_type in [DistributedType.MULTI_GPU]:
+    if distributed_type in [DistributedType.MULTI_GPU, DistributedType.MULTI_NPU, DistributedType.MULTI_XPU]:
         use_fsdp = _ask_field(
             "Do you want to use FullyShardedDataParallel? [yes/NO]: ",
             _convert_yes_no_to_bool,
@@ -282,15 +343,23 @@ def get_cluster_input():
                 lambda x: FSDP_AUTO_WRAP_POLICY[int(x)],
             )
             if fsdp_config["fsdp_auto_wrap_policy"] == FSDP_AUTO_WRAP_POLICY[0]:
-                fsdp_config["fsdp_transformer_layer_cls_to_wrap"] = _ask_field(
-                    "What is the transformer layer class name (case-sensitive) to wrap ,e.g, `BertLayer`, `GPTJBlock`, `T5Block` ...? : ",
-                    str,
+                use_no_split_modules = _ask_field(
+                    "Do you want to use the model's `_no_split_modules` to wrap. Only applicable for 🤗 Transformers [yes/NO]: ",
+                    _convert_yes_no_to_bool,
+                    default=False,
+                    error_message="Please enter yes or no.",
                 )
+                if not use_no_split_modules:
+                    fsdp_config["fsdp_transformer_layer_cls_to_wrap"] = _ask_field(
+                        "Specify the comma-separated list of transformer layer class names (case-sensitive) to wrap ,e.g, :"
+                        "`BertLayer`, `GPTJBlock`, `T5Block`, `BertLayer,BertEmbeddings,BertSelfOutput` ...? : ",
+                        str,
+                    )
             elif fsdp_config["fsdp_auto_wrap_policy"] == FSDP_AUTO_WRAP_POLICY[1]:
                 fsdp_config["fsdp_min_num_params"] = _ask_field(
                     "What should be your FSDP's minimum number of parameters for Default Auto Wrapping Policy? [1e8]: ",
                     int,
-                    default=1e8,
+                    default=100000000,
                 )
             fsdp_backward_prefetch_query = "What should be your FSDP's backward prefetch policy?"
             fsdp_config["fsdp_backward_prefetch_policy"] = _ask_options(
@@ -303,6 +372,25 @@ def get_cluster_input():
                 fsdp_state_dict_type_query,
                 FSDP_STATE_DICT_TYPE,
                 lambda x: FSDP_STATE_DICT_TYPE[int(x)],
+                default=2,
+            )
+            fsdp_config["fsdp_forward_prefetch"] = _ask_field(
+                "Do you want to enable FSDP's forward prefetch policy? [yes/NO]: ",
+                _convert_yes_no_to_bool,
+                default=False,
+                error_message="Please enter yes or no.",
+            )
+            fsdp_config["fsdp_use_orig_params"] = _ask_field(
+                "Do you want to enable FSDP's `use_orig_params` feature? [yes/NO]: ",
+                _convert_yes_no_to_bool,
+                default=False,
+                error_message="Please enter yes or no.",
+            )
+            fsdp_config["fsdp_sync_module_states"] = _ask_field(
+                "Do you want each individually wrapped FSDP unit to broadcast module parameters from rank 0 at the start? [YES/no]: ",
+                _convert_yes_no_to_bool,
+                default=True,
+                error_message="Please enter yes or no.",
             )
 
     megatron_lm_config = {}
@@ -365,72 +453,24 @@ def get_cluster_input():
                 float,
                 default=1.0,
             )
+    # TPU specific defaults
+    tpu_commands = None
+    tpu_command_file = None
+    tpu_downcast_bf16 = "no"
+    tpu_env = []
+    tpu_name = None
+    tpu_vm = None
+    tpu_zone = None
+    tpu_use_sudo = False
+    tpu_use_cluster = False
 
-    if distributed_type == DistributedType.TPU:
-        main_training_function = _ask_field(
-            "What is the name of the function in your script that should be launched in all parallel scripts? [main]: ",
-            default="main",
-        )
-        use_cluster = _ask_field(
-            "Are you using a TPU cluster? [yes/NO]: ",
-            _convert_yes_no_to_bool,
-            default=False,
-            error_message="Please enter yes or no.",
-        )
-        if use_cluster:
-            tpu_name = _ask_field(
-                "What is the name of your TPU cluster? ",
-                default=None,
-                error_message="Please enter the name of your TPU cluster.",
-            )
-            tpu_zone = _ask_field(
-                "What is the zone of your TPU cluster? ",
-                default=None,
-                error_message="Please enter the zone of your TPU cluster.",
-            )
-            run_commands = _ask_field(
-                "Do you have code you wish to run on startup in each pod? [yes/NO]: ",
-                _convert_yes_no_to_bool,
-                default=False,
-                error_message="Please enter yes or no.",
-            )
-            if run_commands:
-                use_command_file = _ask_field(
-                    "Is this code located in a bash script? [yes/NO]: ",
-                    _convert_yes_no_to_bool,
-                    default=False,
-                    error_message="Please enter yes or no.",
-                )
-                if use_command_file:
-                    command_file = _ask_field(
-                        "What is the path to your bash script? ",
-                        default=None,
-                        error_message="Please enter the path to your bash script.",
-                    )
-                    command_file = os.path.abspath(command_file)
-                else:
-                    print("Please enter each command seperately you wish to run on startup in each pod.")
-                    commands = []
-                    another_command = True
-                    while another_command:
-                        commands.append(
-                            _ask_field(
-                                "Please enter a single command to be ran ",
-                                default=None,
-                                error_message="Please enter the commands you wish to run on startup in each pod as a single string.",
-                            )
-                        )
-                        another_command = _ask_field(
-                            "Do you wish to add another command? [yes/NO]: ",
-                            _convert_yes_no_to_bool,
-                            default=False,
-                            error_message="Please enter yes or no.",
-                        )
-
-    else:
-        main_training_function = "main"
-
-    if distributed_type in [DistributedType.MULTI_CPU, DistributedType.MULTI_GPU, DistributedType.TPU]:
+    if distributed_type in [
+        DistributedType.MULTI_CPU,
+        DistributedType.MULTI_XPU,
+        DistributedType.MULTI_GPU,
+        DistributedType.MULTI_NPU,
+        DistributedType.TPU,
+    ]:
         machine_type = str(distributed_type).split(".")[1].replace("MULTI_", "")
         if machine_type == "TPU":
             machine_type += " cores"
@@ -452,32 +492,124 @@ def get_cluster_input():
     else:
         num_processes = 1
 
-    if distributed_type in [DistributedType.MULTI_GPU, DistributedType.NO] and not use_cpu:
+    if (distributed_type == DistributedType.MULTI_GPU) and (num_machines == 1) and (num_processes == 1):
+        raise ValueError(
+            f"Specified distributed type {distributed_type} but only using 1 GPU on a single machine. Please select `No distributed training` for the type of machine you are using."
+        )
+
+    if (
+        distributed_type
+        in [
+            DistributedType.MULTI_GPU,
+            DistributedType.MULTI_NPU,
+            DistributedType.MULTI_XPU,
+            DistributedType.NO,
+        ]
+        and not use_cpu
+        and not use_mps
+    ):
+        if is_npu_available():
+            machine_type = "NPU(s)"
+        else:
+            machine_type = "GPU(s)"
         gpu_ids = _ask_field(
-            "What GPU(s) (by id) should be used for training on this machine as a comma-seperated list? [all]:",
+            f"What {machine_type} (by id) should be used for training on this machine as a comma-seperated list? [all]:",
             default="all",
         )
 
-    if distributed_type != DistributedType.TPU:
+    if distributed_type == DistributedType.TPU:
+        mixed_precision = "no"
+        main_training_function = _ask_field(
+            "What is the name of the function in your script that should be launched in all parallel scripts? [main]: ",
+            default="main",
+        )
+        tpu_use_cluster = _ask_field(
+            "Are you using a TPU cluster? [yes/NO]: ",
+            _convert_yes_no_to_bool,
+            default=False,
+            error_message="Please enter yes or no.",
+        )
+        if tpu_use_cluster:
+            tpu_name = _ask_field(
+                "What is the name of your TPU cluster? ",
+                default=None,
+                error_message="Please enter the name of your TPU cluster.",
+            )
+            tpu_zone = _ask_field(
+                "What is the zone of your TPU cluster? ",
+                default=None,
+                error_message="Please enter the zone of your TPU cluster.",
+            )
+            tpu_use_sudo = _ask_field(
+                "To run a python script in a TPU pod, should `sudo` be used? [yes/NO]: ",
+                default=False,
+                error_message="Please enter yes or no.",
+            )
+            run_commands = _ask_field(
+                "Do you have code you wish to run on startup in each pod? [yes/NO]: ",
+                _convert_yes_no_to_bool,
+                default=False,
+                error_message="Please enter yes or no.",
+            )
+            if run_commands:
+                use_command_file = _ask_field(
+                    "Is this code located in a bash script? [yes/NO]: ",
+                    _convert_yes_no_to_bool,
+                    default=False,
+                    error_message="Please enter yes or no.",
+                )
+                if use_command_file:
+                    tpu_command_file = _ask_field(
+                        "What is the path to your bash script? ",
+                        default=None,
+                        error_message="Please enter the path to your bash script.",
+                    )
+                    tpu_command_file = os.path.abspath(tpu_command_file)
+                else:
+                    print("Please enter each command seperately you wish to run on startup in each pod.")
+                    tpu_commands = []
+                    another_command = True
+                    while another_command:
+                        tpu_commands.append(
+                            _ask_field(
+                                "Please enter a single command to be ran ",
+                                default=None,
+                                error_message="Please enter the commands you wish to run on startup in each pod as a single string.",
+                            )
+                        )
+                        another_command = _ask_field(
+                            "Do you wish to add another command? [yes/NO]: ",
+                            _convert_yes_no_to_bool,
+                            default=False,
+                            error_message="Please enter yes or no.",
+                        )
+            tpu_vm = _ask_field(
+                "If not using an instance group, what are the names of the Compute VM instances to be used, seperated by a comma: ",
+                default="",
+            ).split(",")
+            tpu_env = _ask_field(
+                "What environment variables do you wish to set in each pod, seperated by a comma: ",
+                default="",
+            ).split(",")
+
+    else:
+        main_training_function = "main"
         if distributed_type == DistributedType.DEEPSPEED and use_deepspeed_config:
             mixed_precision = None
         else:
             mixed_precision = _ask_options(
                 "Do you wish to use FP16 or BF16 (mixed precision)?",
-                ["no", "fp16", "bf16"],
+                ["no", "fp16", "bf16", "fp8"],
                 _convert_mixed_precision,
             )
-    else:
-        mixed_precision = "no"
 
     if use_dynamo and mixed_precision == "no" and not use_cpu:
         print(
             "Torch dynamo used without mixed precision requires TF32 to be efficient. Accelerate will enable it by default when launching your scripts."
         )
 
-    downcast_bf16 = "no"
     if distributed_type == DistributedType.TPU and mixed_precision == "bf16":
-        downcast_bf16 = _ask_field(
+        tpu_downcast_bf16 = _ask_field(
             "Should `torch.float` be cast as `bfloat16` and `torch.double` remain `float32` on TPUs?", default="no"
         )
 
@@ -487,7 +619,7 @@ def get_cluster_input():
         num_processes=num_processes,
         gpu_ids=gpu_ids,
         mixed_precision=mixed_precision,
-        downcast_bf16=downcast_bf16,
+        downcast_bf16=tpu_downcast_bf16,
         machine_rank=machine_rank,
         num_machines=num_machines,
         main_process_ip=main_process_ip,
@@ -496,12 +628,18 @@ def get_cluster_input():
         deepspeed_config=deepspeed_config,
         fsdp_config=fsdp_config,
         megatron_lm_config=megatron_lm_config,
+        ipex_config=ipex_config,
         use_cpu=use_cpu,
         rdzv_backend=rdzv_backend,
         same_network=same_network,
+        commands=tpu_commands,
+        command_file=tpu_command_file,
+        tpu_env=tpu_env,
         tpu_name=tpu_name,
+        tpu_vm=tpu_vm,
         tpu_zone=tpu_zone,
-        commands=commands,
-        command_file=command_file,
-        dynamo_backend=dynamo_backend,
+        tpu_use_sudo=tpu_use_sudo,
+        tpu_use_cluster=tpu_use_cluster,
+        dynamo_config=dynamo_config,
+        debug=debug,
     )