[ { "path": ".devcontainer/devcontainer.json", "content": "// File only needed for VSCode users to have proper Docker based interpreters\n{\n \"name\": \"accelerate_dev_environment\",\n \"build\": {\n // ACTION NEEDED: comment/uncomment the relevant line depending on whether you are in a CPU/GPU environment\n \"dockerfile\": \"../docker/accelerate-cpu/Dockerfile\"\n// \"dockerfile\": \"../docker/accelerate-gpu/Dockerfile\"\n },\n \"runArgs\": [\n // ACTION NEEDED: uncomment the next line if your local machine has GPUs available\n// \"--gpus\", \"all\",\n // Enable the docker container to access system resources\n \"--ipc\", \"host\"\n ],\n \"remoteEnv\": {\n \"PYTHONPATH\": \"${containerEnv:PATH}:${containerWorkspaceFolder}\"\n },\n \"customizations\": {\n \"vscode\": {\n \"extensions\": [\n // Ensure we have IntelliSense in VSCode when running inside container\n \"ms-python.python\"\n ]\n }\n },\n \"workspaceFolder\": \"/workspaces/accelerate\",\n // Need git for VSCode to color code modifications. Only runs when building environment.\n \"onCreateCommand\": \"apt-get update && apt-get install -y git && pip install -e '.[dev]'\"\n}" }, { "path": ".github/ISSUE_TEMPLATE/bug-report.yml", "content": "name: \"\\U0001F41B Bug Report\"\ndescription: Submit a bug report to help us improve Accelerate\nbody:\n - type: markdown\n attributes: \n value: | \n Thanks for taking the time to submit a bug report! šŸ› \n 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).\n\n - type: textarea\n id: system-info\n attributes:\n label: System Info\n description: Please share your accelerate configuration with us. You can run the command `accelerate env` and copy-paste its outputs below\n render: Shell\n placeholder: accelerate version, OS, python version, numpy version, torch version, and accelerate's configuration\n validations:\n required: true\n \n - type: checkboxes\n id: information-scripts-examples\n attributes:\n label: Information\n description: 'The problem arises when using:'\n options:\n - label: \"The official example scripts\"\n - label: \"My own modified scripts\"\n \n - type: checkboxes\n id: information-tasks\n attributes:\n label: Tasks\n description: \"The tasks I am working on are:\"\n options:\n - label: \"One of the scripts in the examples/ folder of Accelerate or an officially supported `no_trainer` script in the `examples` folder of the `transformers` repo (such as `run_no_trainer_glue.py`)\"\n - label: \"My own task or dataset (give details below)\"\n \n - type: textarea\n id: reproduction\n validations:\n required: true\n attributes:\n label: Reproduction\n description: |\n Please provide a code sample that reproduces the problem you ran into. It can be a Colab link or just a code snippet.\n If you have code snippets, error messages, stack traces please provide them here as well.\n Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting\n Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.\n\n placeholder: |\n Steps to reproduce the behavior:\n \n 1.\n 2.\n 3.\n\n - type: textarea\n id: expected-behavior\n validations:\n required: true\n attributes:\n label: Expected behavior\n description: \"A clear and concise description of what you would expect to happen.\"\n" }, { "path": ".github/PULL_REQUEST_TEMPLATE.md", "content": "# What does this PR do?\n\n\n\n\n\nFixes # (issue)\n\n\n## Before submitting\n- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).\n- [ ] Did you read the [contributor guideline](https://github.com/huggingface/accelerate/blob/main/CONTRIBUTING.md#submitting-a-pull-request-pr),\n Pull Request section?\n- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link\n to it if that's the case.\n- [ ] Did you make sure to update the documentation with your changes? Here are the\n [documentation guidelines](https://github.com/huggingface/accelerate/tree/main/docs), and\n [here are tips on formatting docstrings](https://github.com/huggingface/accelerate/tree/main/docs#writing-documentation---specification).\n- [ ] Did you write any new necessary tests?\n\n\n## Who can review?\n\nAnyone in the community is free to review the PR once the tests have passed. Feel free to tag\nmembers/contributors who may be interested in your PR.\n\n" }, { "path": ".github/workflows/build-docker-images-release.yml", "content": "name: Build Docker images (releases)\n\non:\n workflow_dispatch:\n release:\n types: [published]\n\nconcurrency:\n group: docker-image-builds\n cancel-in-progress: false\n\njobs:\n get-version:\n runs-on: ubuntu-latest\n outputs:\n version: ${{ steps.step1.outputs.version }}\n steps:\n - uses: actions/checkout@v6\n - id: step1\n run: echo \"version=$(python setup.py --version)\" >> $GITHUB_OUTPUT\n\n version-cpu:\n name: \"Latest Accelerate CPU [version]\"\n runs-on:\n group: aws-general-8-plus\n needs: get-version\n steps:\n - name: Set up Docker Buildx\n uses: docker/setup-buildx-action@v3\n - name: Login to DockerHub\n uses: docker/login-action@v3\n with:\n username: ${{ secrets.DOCKERHUB_USERNAME }}\n password: ${{ secrets.DOCKERHUB_PASSWORD }}\n\n - name: Build and Push CPU\n uses: docker/build-push-action@v6\n with:\n file: docker/accelerate-cpu/Dockerfile\n push: true\n tags: huggingface/accelerate:cpu-release-${{ needs.get-version.outputs.version }}\n\n version-cuda:\n name: \"Latest Accelerate GPU [version]\"\n runs-on:\n group: aws-g6-4xlarge-plus\n needs: get-version\n steps:\n - name: Set up Docker Buildx\n uses: docker/setup-buildx-action@v3\n - name: Login to DockerHub\n uses: docker/login-action@v3\n with:\n username: ${{ secrets.DOCKERHUB_USERNAME }}\n password: ${{ secrets.DOCKERHUB_PASSWORD }}\n\n - name: Build and Push GPU\n uses: docker/build-push-action@v6\n with:\n file: docker/accelerate-gpu/Dockerfile\n push: true\n tags: huggingface/accelerate:gpu-release-${{needs.get-version.outputs.version}}\n\n version-cuda-deepspeed:\n name: \"Latest Accelerate GPU DeepSpeed [version]\"\n runs-on:\n group: aws-g6-4xlarge-plus\n needs: get-version\n steps:\n - name: Set up Docker Buildx\n uses: docker/setup-buildx-action@v3\n - name: Login to DockerHub\n uses: docker/login-action@v3\n with:\n username: ${{ secrets.DOCKERHUB_USERNAME }}\n password: ${{ secrets.DOCKERHUB_PASSWORD }}\n\n - name: Build and Push GPU\n uses: docker/build-push-action@v6\n with:\n file: docker/accelerate-gpu-deepspeed/Dockerfile\n push: true\n tags: huggingface/accelerate:gpu-deepspeed-release-${{needs.get-version.outputs.version}}\n\n version-cuda-fp8-transformerengine:\n name: \"Latest Accelerate GPU FP8 TransformerEngine [version]\"\n runs-on:\n group: aws-g6-4xlarge-plus\n needs: get-version\n steps:\n - name: Set up Docker Buildx\n uses: docker/setup-buildx-action@v3\n - name: Login to DockerHub\n uses: docker/login-action@v3\n with:\n username: ${{ secrets.DOCKERHUB_USERNAME }}\n password: ${{ secrets.DOCKERHUB_PASSWORD }}\n\n - name: Build and Push GPU\n uses: docker/build-push-action@v6\n with:\n file: docker/accelerate-gpu/Dockerfile\n push: true\n tags: huggingface/accelerate:gpu-fp8-transformerengine-release-${{needs.get-version.outputs.version}}" }, { "path": ".github/workflows/build_and_run_tests.yml", "content": "name: Trigger docker images and run tests\n\non:\n push:\n branches:\n - main\n workflow_dispatch:\n\nenv:\n GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}\n\njobs:\n check-for-source:\n runs-on: ubuntu-latest\n name: Check if setup was changed\n outputs:\n changed: ${{ steps.was_changed.outputs.changed }}\n steps:\n - uses: actions/checkout@v6\n with: \n fetch-depth: \"2\"\n \n - name: Get changed files\n id: changed-files\n uses: tj-actions/changed-files@3f54ebb830831fc121d3263c1857cfbdc310cdb9 #v42\n \n - name: Was setup changed \n id: was_changed\n run: |\n for file in ${{ steps.changed-files.outputs.all_changed_files }}; do\n if [ `basename \"${file}\"` == \"setup.py\" ]; then\n echo \"changed=1\" >> $GITHUB_OUTPUT\n fi\n done\n \n build-docker-containers:\n needs: check-for-source\n if: (github.event_name == 'push') && (needs.check-for-source.outputs.changed == '1')\n uses: ./.github/workflows/build_docker_images.yml\n secrets: inherit\n\n run-merge-tests:\n needs: build-docker-containers\n if: always()\n uses: ./.github/workflows/run_merge_tests.yml\n\n run-integration-tests:\n needs: build-docker-containers\n if: always()\n uses: ./.github/workflows/self_hosted_integration_tests.yml\n" }, { "path": ".github/workflows/build_docker_images.yml", "content": "name: Build Docker images (scheduled)\n\non:\n workflow_dispatch:\n workflow_call:\n schedule:\n - cron: \"0 1 * * *\"\n\nconcurrency:\n group: docker-image-builds\n cancel-in-progress: false\n\njobs:\n latest-cpu:\n name: \"Latest Accelerate CPU [dev]\"\n runs-on:\n group: aws-general-8-plus\n steps:\n - name: Set up Docker Buildx\n uses: docker/setup-buildx-action@v3\n - name: Login to DockerHub\n uses: docker/login-action@v3\n with:\n username: ${{ secrets.DOCKERHUB_USERNAME }}\n password: ${{ secrets.DOCKERHUB_PASSWORD }}\n - name: Get current date\n id: date\n run: |\n echo \"date=$(date '+%Y-%m-%d')\" >> $GITHUB_ENV\n - name: Build and Push CPU\n uses: docker/build-push-action@v6\n with:\n file: docker/accelerate-cpu/Dockerfile\n push: true\n tags: |\n huggingface/accelerate:cpu-nightly\n huggingface/accelerate:cpu-nightly-${{ env.date }}\n\n latest-cuda:\n name: \"Latest Accelerate GPU [dev]\"\n runs-on:\n group: aws-g6-4xlarge-plus\n steps:\n - name: Set up Docker Buildx\n uses: docker/setup-buildx-action@v3\n - name: Login to DockerHub\n uses: docker/login-action@v3\n with:\n username: ${{ secrets.DOCKERHUB_USERNAME }}\n password: ${{ secrets.DOCKERHUB_PASSWORD }}\n - name: Get current date\n id: date\n run: |\n echo \"date=$(date '+%Y-%m-%d')\" >> $GITHUB_ENV\n - name: Build and Push GPU\n uses: docker/build-push-action@v6\n with:\n file: docker/accelerate-gpu/Dockerfile\n push: true\n tags: |\n huggingface/accelerate:gpu-nightly\n huggingface/accelerate:gpu-nightly-${{ env.date }}\n\n latest-cuda-deepspeed:\n name: \"Latest Accelerate GPU DeepSpeed [dev]\"\n runs-on:\n group: aws-g6-4xlarge-plus\n steps:\n - name: Set up Docker Buildx\n uses: docker/setup-buildx-action@v3\n - name: Login to DockerHub\n uses: docker/login-action@v3\n with:\n username: ${{ secrets.DOCKERHUB_USERNAME }}\n password: ${{ secrets.DOCKERHUB_PASSWORD }}\n - name: Get current date\n id: date\n run: |\n echo \"date=$(date '+%Y-%m-%d')\" >> $GITHUB_ENV\n - name: Build and Push GPU\n uses: docker/build-push-action@v6\n with:\n file: docker/accelerate-gpu-deepspeed/Dockerfile\n push: true\n tags: |\n huggingface/accelerate:gpu-deepspeed-nightly\n huggingface/accelerate:gpu-deepspeed-nightly-${{ env.date }}\n\n latest-cuda-fp8-transformerengine:\n name: \"Latest Accelerate GPU FP8 TransformerEngine [dev]\"\n runs-on:\n group: aws-g6-4xlarge-plus\n steps:\n - name: Set up Docker Buildx\n uses: docker/setup-buildx-action@v3\n - name: Login to DockerHub\n uses: docker/login-action@v3\n with:\n username: ${{ secrets.DOCKERHUB_USERNAME }}\n password: ${{ secrets.DOCKERHUB_PASSWORD }}\n - name: Get current date\n id: date\n run: |\n echo \"date=$(date '+%Y-%m-%d')\" >> $GITHUB_ENV\n # Get the previous month\n echo \"base_year=$(date -d 'last month' '+%y')\" >> $GITHUB_ENV\n echo \"base_month=$(date -d 'last month' '+%m')\" >> $GITHUB_ENV\n - name: Build and Push GPU\n uses: docker/build-push-action@v6\n with:\n file: benchmarks/fp8/transformer_engine/Dockerfile\n push: true\n tags: huggingface/accelerate:gpu-fp8-transformerengine-nightly-${{ env.date }}\n build-args: |\n BASE_YEAR=${{ env.base_year }}\n BASE_MONTH=${{ env.base_month }}" }, { "path": ".github/workflows/build_documentation.yml", "content": "name: Build documentation\n\non:\n push:\n branches:\n - main\n - doc-builder*\n - v*-release\n\njobs:\n build:\n uses: huggingface/doc-builder/.github/workflows/build_main_documentation.yml@main\n with:\n commit_sha: ${{ github.sha }}\n package: accelerate\n custom_container: huggingface/transformers-doc-builder\n secrets:\n hf_token: ${{ secrets.HF_DOC_BUILD_PUSH }}\n" }, { "path": ".github/workflows/build_pr_documentation.yml", "content": "name: Build PR Documentation\n\non:\n pull_request:\n\nconcurrency:\n group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}\n cancel-in-progress: true\n\njobs:\n build:\n uses: huggingface/doc-builder/.github/workflows/build_pr_documentation.yml@main\n with:\n commit_sha: ${{ github.event.pull_request.head.sha }}\n pr_number: ${{ github.event.number }}\n package: accelerate\n custom_container: huggingface/transformers-doc-builder\n" }, { "path": ".github/workflows/fp8_runner.yml", "content": "name: Test FP8 Runner\n\non:\n workflow_dispatch:\n\nenv:\n GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}\njobs:\n set-prev-day:\n runs-on: ubuntu-latest\n outputs:\n prev-day: ${{ steps.set-prev-day.outputs.prev-day }}\n steps:\n - name: Set PREV_DAY\n id: set-prev-day\n run: |\n PREV_DAY=$(date -d \"yesterday\" '+%Y-%m-%d')\n echo \"prev-day=$PREV_DAY\" >> $GITHUB_OUTPUT\n run-fp8-tests:\n needs: set-prev-day\n runs-on:\n group: aws-g6e-12xlarge\n container:\n image: huggingface/accelerate:gpu-fp8-transformerengine-nightly-${{ needs.set-prev-day.outputs.prev-day }}\n options: --gpus all --shm-size \"16gb\"\n steps:\n - uses: actions/checkout@v6\n - name: Install the library\n run: |\n pip install -e .[test_prod,test_fp8]\n - name: Show installed libraries\n run: |\n pip freeze\n - name: Run TE FP8 tests\n run: |\n python -m pytest -s -v ./tests/test_fp8.py\n\n" }, { "path": ".github/workflows/gaudi3_scheduled.yml", "content": "name: Gaudi3 tests (scheduled)\n\non:\n workflow_dispatch:\n schedule: # every day at 6 AM UTC\n - cron: \"0 6 * * *\"\n\nconcurrency:\n group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}\n cancel-in-progress: true\n\njobs:\n run-gaudi3-tests:\n runs-on:\n group: itac-bm-emr-gaudi3-dell-2gaudi\n\n container:\n image: docker://vault.habana.ai/gaudi-docker/1.21.1/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest\n options: --runtime=habana --shm-size=64G --cap-add=sys_nice --env HABANA_VISIBLE_DEVICES\n env:\n OMPI_MCA_btl_vader_single_copy_mechanism: none\n PT_ENABLE_INT64_SUPPORT: 1\n PT_HPU_LAZY_MODE: 0\n RUN_SLOW: 1\n\n steps:\n - name: HL-SMI (1)\n run: |\n hl-smi\n echo \"HABANA_VISIBLE_DEVICES=${HABANA_VISIBLE_DEVICES}\"\n echo \"HABANA_VISIBLE_MODULES=${HABANA_VISIBLE_MODULES}\"\n\n - name: Extract HPU visible modules\n id: add-modules\n run: |\n export HABANA_VISIBLE_MODULES=$(hl-smi -Q module_id -f csv,noheader | tr '\\n' ',' | sed 's/,$//')\n echo \"HABANA_VISIBLE_MODULES=${HABANA_VISIBLE_MODULES}\" >> $GITHUB_ENV\n\n - name: HL-SMI (2)\n run: |\n hl-smi\n echo \"HABANA_VISIBLE_DEVICES=${HABANA_VISIBLE_DEVICES}\"\n echo \"HABANA_VISIBLE_MODULES=${HABANA_VISIBLE_MODULES}\"\n\n - name: Checkout to Accelerate\n uses: actions/checkout@v6\n\n - name: Install Accelerate with Transformers & DeepSpeed\n run: |\n pip install -e .[testing] \\\n git+https://github.com/HabanaAI/DeepSpeed.git@1.20.0 \\\n git+https://github.com/huggingface/transformers.git\n\n - name: Run CLI tests\n if: ${{ !cancelled() && (success() || failure()) }}\n run: |\n make test_cli\n\n - name: Run Core tests\n if: ${{ !cancelled() && (success() || failure()) }}\n run: |\n make test_core\n\n - name: Run Big Modeling tests\n if: ${{ !cancelled() && (success() || failure()) }}\n run: |\n make test_big_modeling\n\n - name: Run DeepSpeed integration tests\n if: ${{ !cancelled() && (success() || failure()) }}\n run: |\n make test_deepspeed\n\n - name: Run FSDP integration tests\n if: ${{ !cancelled() && (success() || failure()) }}\n run: |\n make test_fsdp\n\n - name: Run TP integration tests\n if: ${{ !cancelled() && (success() || failure()) }}\n run: |\n make test_tp\n\n - name: Run Examples tests\n if: ${{ !cancelled() && (success() || failure()) }}\n run: |\n make test_examples\n" }, { "path": ".github/workflows/integration_tests.yml", "content": "# CI for specifically ensuring integrations work fine (`transformers` mainly)\n# Useful tips:\n# - New integrations to test should have its own job, and follow a strategy method where we check both\n# the pypi and github versions.\n# - When checking the latest release of the integration, use\n# git checkout $(git describe --tags `git rev-list --tags --max-count=1`) to get the latest release.\n\nname: Integration Tests\n\non:\n pull_request:\n paths:\n - \"src/**\"\n - \"tests/**\"\n - \".github/**\"\n - \"examples/**\"\n - \"setup.py\"\n types: [opened, synchronize, reopened]\n\nenv:\n HF_HOME: ~/hf_cache\n\njobs:\n run-trainer-tests:\n runs-on: ubuntu-latest\n strategy:\n fail-fast: false\n steps:\n - uses: actions/checkout@v6\n - name: Set up python 3.10\n uses: actions/setup-python@v6\n with:\n python-version: '3.10'\n cache: 'pip'\n cache-dependency-path: 'setup.py'\n\n - name: Install Accelerate from source\n run: |\n pip install --upgrade pip\n pip install -e .\n \n - name: Clone and install transformers\n run: |\n cd ..\n git clone https://github.com/huggingface/transformers\n cd transformers\n pip install .[torch,testing]\n\n - name: Show installed libraries\n run: |\n pip freeze\n\n - name: Run Trainer tests\n env:\n WANDB_DISABLED: true\n run: |\n cd ../transformers\n pytest -sv tests/trainer\n" }, { "path": ".github/workflows/nightly.yml", "content": "name: Self-hosted runner with slow tests (scheduled)\n\non:\n workflow_dispatch:\n schedule:\n - cron: \"0 2 * * *\"\n\nenv:\n RUN_SLOW: \"yes\"\n IS_GITHUB_CI: \"1\"\n SLACK_API_TOKEN: ${{ secrets.SLACK_API_TOKEN }}\n\n\njobs:\n run_core_tests_single_gpu:\n runs-on:\n group: aws-g6-4xlarge-plus\n env:\n CUDA_VISIBLE_DEVICES: \"0\"\n TEST_TYPE: \"single_gpu\"\n container:\n image: huggingface/accelerate:gpu-nightly\n options: --gpus all --shm-size \"16gb\"\n defaults:\n run:\n shell: bash\n steps:\n - name: Update clone & pip install\n run: |\n source activate accelerate\n git clone https://github.com/huggingface/accelerate;\n cd accelerate;\n git checkout ${{ github.sha }};\n pip install -e . --no-deps\n pip install pytest-reportlog tabulate\n\n - name: Show installed libraries\n run: |\n source activate accelerate;\n pip freeze\n\n - name: Run test on GPUs\n working-directory: accelerate\n run: |\n source activate accelerate\n make test\n\n - name: Run examples on GPUs\n working-directory: accelerate\n if: always()\n run: |\n source activate accelerate\n pip uninstall comet_ml -y\n make test_examples\n\n - name: Generate Report\n working-directory: accelerate\n if: always()\n run: |\n pip install slack_sdk tabulate\n python utils/log_reports.py >> $GITHUB_STEP_SUMMARY\n\n run_deepspeed_tests_single_gpu:\n runs-on:\n group: aws-g6-4xlarge-plus\n env:\n CUDA_VISIBLE_DEVICES: \"0\"\n TEST_TYPE: \"single_gpu_deepspeed\"\n container:\n image: huggingface/accelerate:gpu-deepspeed-nightly\n options: --gpus all --shm-size \"16gb\"\n defaults:\n run:\n shell: bash\n steps:\n - name: Update clone & pip install\n run: |\n source activate accelerate\n git clone https://github.com/huggingface/accelerate;\n cd accelerate;\n git checkout ${{ github.sha }};\n pip install -e . --no-deps\n pip install pytest-reportlog tabulate\n\n - name: Show installed libraries\n run: |\n source activate accelerate;\n pip freeze\n\n - name: Run test on GPUs\n working-directory: accelerate\n run: |\n source activate accelerate\n make test_deepspeed\n\n - name: Run Integration tests on GPUs\n working-directory: accelerate\n if: always()\n run: |\n source activate accelerate\n make test_integrations\n\n - name: Run examples on GPUs\n working-directory: accelerate\n if: always()\n run: |\n source activate accelerate\n pip uninstall comet_ml -y\n make test_examples\n\n - name: Generate Report\n working-directory: accelerate\n if: always()\n run: |\n pip install slack_sdk tabulate\n python utils/log_reports.py >> $GITHUB_STEP_SUMMARY\n\n run_core_tests_multi_gpu:\n runs-on:\n group: aws-g6-12xlarge-plus\n env:\n CUDA_VISIBLE_DEVICES: \"0,1\"\n TEST_TYPE: \"multi_gpu\"\n container:\n image: huggingface/accelerate:gpu-nightly\n options: --gpus all --shm-size \"16gb\"\n defaults:\n run:\n shell: bash\n steps:\n - name: Update clone\n run: |\n source activate accelerate\n git clone https://github.com/huggingface/accelerate;\n cd accelerate;\n git checkout ${{ github.sha }};\n pip install -e . --no-deps\n pip install pytest-reportlog tabulate\n\n - name: Show installed libraries\n run: |\n source activate accelerate;\n pip freeze\n\n - name: Run core and big modeling tests on GPUs\n working-directory: accelerate\n run: |\n source activate accelerate\n make test_core\n make test_big_modeling\n make test_cli\n\n - name: Run Integration tests on GPUs\n working-directory: accelerate\n if: always()\n run: |\n source activate accelerate\n make test_integrations\n\n - name: Run examples on GPUs\n working-directory: accelerate\n if: always()\n run: |\n source activate accelerate\n pip uninstall comet_ml -y\n make test_examples\n\n - name: Generate Report\n working-directory: accelerate\n if: always()\n run: |\n pip install slack_sdk tabulate\n python utils/log_reports.py >> $GITHUB_STEP_SUMMARY\n\n run_deepspeed_tests_multi_gpu:\n runs-on:\n group: aws-g6-12xlarge-plus\n env:\n CUDA_VISIBLE_DEVICES: \"0,1\"\n TEST_TYPE: \"multi_gpu_deepspeed\"\n container:\n image: huggingface/accelerate:gpu-deepspeed-nightly\n options: --gpus all --shm-size \"16gb\"\n defaults:\n run:\n shell: bash\n steps:\n - name: Update clone\n run: |\n source activate accelerate\n git clone https://github.com/huggingface/accelerate;\n cd accelerate;\n git checkout ${{ github.sha }};\n pip install -e . --no-deps\n pip install pytest-reportlog tabulate\n\n - name: Show installed libraries\n run: |\n source activate accelerate;\n pip freeze\n\n - name: Run DeepSpeed tests\n working-directory: accelerate\n run: |\n source activate accelerate\n make test_deepspeed\n\n - name: Run Integration tests on GPUs\n working-directory: accelerate\n if: always()\n run: |\n source activate accelerate\n make test_integrations\n\n - name: Run examples on GPUs\n working-directory: accelerate\n if: always()\n run: |\n source activate accelerate\n pip uninstall comet_ml -y\n make test_examples\n\n - name: Generate Report\n working-directory: accelerate\n if: always()\n run: |\n pip install slack_sdk tabulate\n python utils/log_reports.py >> $GITHUB_STEP_SUMMARY\n\n\n run-integration-tests:\n if: always()\n uses: ./.github/workflows/self_hosted_integration_tests.yml\n" }, { "path": ".github/workflows/pr_style_bot.yml", "content": "# To run this bot, comment \"@bot /style\" on a PR\nname: Style Bot\n\non:\n issue_comment:\n types: [created]\n\npermissions:\n contents: write\n pull-requests: write\n\njobs:\n style:\n uses: huggingface/huggingface_hub/.github/workflows/style-bot-action.yml@main\n with:\n python_quality_dependencies: \"[quality]\"\n style_command_type: \"default\"\n secrets:\n bot_token: ${{ secrets.GITHUB_TOKEN }}" }, { "path": ".github/workflows/quality.yml", "content": "name: Quality Check\n\non: [pull_request]\n\njobs:\n quality:\n runs-on: ubuntu-latest\n steps:\n - uses: actions/checkout@v6\n - name: Set up Python 3.10\n uses: actions/setup-python@v6\n with:\n python-version: '3.10'\n cache: 'pip'\n cache-dependency-path: 'setup.py'\n - name: Install Python dependencies\n run: pip install -e .[quality]\n - name: Run Quality check\n run: make quality\n - name: Check if failure\n if: ${{ failure() }}\n run: |\n 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\n\n" }, { "path": ".github/workflows/run_merge_tests.yml", "content": "name: Self-hosted runner tests (push to \"main\")\n\non:\n workflow_call:\n workflow_dispatch:\n\nenv:\n TESTING_MOCKED_DATALOADERS: \"1\"\n IS_GITHUB_CI: \"1\"\n\njobs:\n run_core_tests_single_gpu:\n runs-on:\n group: aws-g6-4xlarge-plus\n env:\n CUDA_VISIBLE_DEVICES: \"0\"\n container:\n image: huggingface/accelerate:gpu-nightly\n options: --gpus all --shm-size \"16gb\"\n defaults:\n run:\n shell: bash\n steps:\n - name: Install accelerate\n run: |\n source activate accelerate;\n git clone https://github.com/huggingface/accelerate;\n cd accelerate;\n git checkout ${{ github.sha }};\n pip install -e .[testing,test_trackers] -U;\n pip install pytest-reportlog tabulate ;\n\n - name: Show installed libraries\n run: |\n source activate accelerate;\n pip freeze\n\n - name: Run CLI tests (use make cli)\n working-directory: accelerate\n run: |\n source activate accelerate;\n make test_cli\n\n - name: Run test on GPUs\n working-directory: accelerate\n if: always()\n run: |\n source activate accelerate;\n make test\n - name: Run examples on GPUs\n working-directory: accelerate\n if: always()\n run: |\n source activate accelerate;\n pip uninstall comet_ml -y;\n make test_examples\n\n - name: Generate Report\n working-directory: accelerate\n if: always()\n run: |\n pip install tabulate;\n python utils/log_reports.py >> $GITHUB_STEP_SUMMARY\n\n run_deepspeed_tests_single_gpu:\n runs-on:\n group: aws-g6-4xlarge-plus\n env:\n CUDA_VISIBLE_DEVICES: \"0\"\n container:\n image: huggingface/accelerate:gpu-deepspeed-nightly\n options: --gpus all --shm-size \"16gb\"\n defaults:\n run:\n shell: bash\n steps:\n - name: Install accelerate\n run: |\n source activate accelerate;\n git clone https://github.com/huggingface/accelerate;\n cd accelerate;\n git checkout ${{ github.sha }};\n pip install -e .[testing,test_trackers] -U;\n pip install pytest-reportlog tabulate ;\n\n - name: Show installed libraries\n run: |\n source activate accelerate;\n pip freeze\n\n - name: Run test on GPUs\n working-directory: accelerate\n if: always()\n run: |\n source activate accelerate;\n make test_deepspeed\n\n - name: Generate Report\n working-directory: accelerate\n if: always()\n run: |\n pip install tabulate;\n python utils/log_reports.py >> $GITHUB_STEP_SUMMARY\n\n run_core_tests_multi_gpu:\n runs-on:\n group: aws-g6-12xlarge-plus\n env:\n CUDA_VISIBLE_DEVICES: 0,1\n container:\n image: huggingface/accelerate:gpu-nightly\n options: --gpus all --shm-size \"16gb\"\n defaults:\n run:\n shell: bash\n steps:\n - name: Update clone\n run: |\n source activate accelerate;\n git clone https://github.com/huggingface/accelerate;\n cd accelerate;\n git checkout ${{ github.sha }};\n pip install -e .[testing,test_trackers] -U;\n pip install pytest-reportlog tabulate\n\n - name: Show installed libraries\n run: |\n source activate accelerate;\n pip freeze\n\n - name: Run test on GPUs\n working-directory: accelerate\n run: |\n source activate accelerate;\n make test\n\n - name: Run examples on GPUs\n working-directory: accelerate\n if: always()\n run: |\n source activate accelerate;\n pip uninstall comet_ml -y;\n make test_examples\n\n - name: Generate Report\n working-directory: accelerate\n if: always()\n run: |\n source activate accelerate;\n python utils/log_reports.py >> $GITHUB_STEP_SUMMARY\n\n run_deepspeed_tests_multi_gpu:\n runs-on:\n group: aws-g6-12xlarge-plus\n container:\n image: huggingface/accelerate:gpu-deepspeed-nightly\n options: --gpus all --shm-size \"16gb\"\n defaults:\n run:\n shell: bash\n steps:\n - name: Install accelerate\n run: |\n source activate accelerate;\n git clone https://github.com/huggingface/accelerate;\n cd accelerate;\n git checkout ${{ github.sha }};\n pip install -e .[testing,test_trackers] -U;\n pip install pytest-reportlog tabulate ;\n\n - name: Show installed libraries\n run: |\n source activate accelerate;\n pip freeze\n\n - name: Run test on GPUs\n working-directory: accelerate\n if: always()\n run: |\n source activate accelerate;\n make test_deepspeed\n\n - name: Generate Report\n working-directory: accelerate\n if: always()\n run: |\n pip install tabulate;\n python utils/log_reports.py >> $GITHUB_STEP_SUMMARY\n" }, { "path": ".github/workflows/self_hosted_integration_tests.yml", "content": "# CI for specifically ensuring integrations work fine (`transformers` mainly) on GPUs\n# Useful tips:\n# - `working-directory` should be set to the root of the repo, which is cloned on the actual CI runner.\n# It follows the directory structure of `actions-runner/_work/{repo_name}/{repo_name}/{cloned_repo} on\n# prem, but in Actions setting `working-directory` looks just in the `{repo_name}` level.\n# - New integrations to test should have its own job, and follow a strategy method where we check both\n# the pypi and github versions.\n# - Workflow call lets this be called from `build_and_run_tests.yml`\n# - When using a docker container, it's recommended to set `--shm-size`, we use 16gb.\nname: Integration Tests (push to \"main\")\n\non:\n workflow_call:\n workflow_dispatch:\n\nenv:\n HF_HOME: ~/hf_cache\n\ndefaults:\n run:\n shell: bash\n\njobs:\n run-trainer-tests:\n container:\n image: huggingface/accelerate:gpu-deepspeed-nightly\n options: --gpus all --shm-size \"16gb\"\n runs-on:\n group: aws-g6-12xlarge-plus\n strategy:\n fail-fast: false\n matrix:\n cuda_visible_devices: [\n \"0\",\n \"0,1\"\n ]\n steps:\n - name: Install transformers\n run: |\n source activate accelerate;\n git clone https://github.com/huggingface/transformers --depth 1;\n cd transformers;\n pip install .[torch,deepspeed-testing];\n cd ..;\n\n - name: Install accelerate\n run: |\n source activate accelerate;\n git clone https://github.com/huggingface/accelerate;\n cd accelerate;\n git checkout ${{ github.sha }} ;\n pip install -e .[testing];\n pip uninstall comet_ml wandb dvclive -y\n cd ..;\n\n - name: Show installed libraries\n run: |\n source activate accelerate;\n pip freeze\n\n - name: Run trainer tests\n working-directory: transformers/\n env:\n CUDA_VISIBLE_DEVICES: ${{ matrix.cuda_visible_devices }}\n WANDB_DISABLED: true\n run: |\n source activate accelerate;\n pytest -sv tests/trainer\n\n - name: Run deepspeed tests\n working-directory: transformers/\n env:\n CUDA_VISIBLE_DEVICES: ${{ matrix.cuda_visible_devices }}\n WANDB_DISABLED: true\n if: always()\n run: |\n source activate accelerate;\n pytest -sv tests/deepspeed\n\n - name: Run transformers examples tests\n working-directory: transformers/\n env:\n CUDA_VISIBLE_DEVICES: ${{ matrix.cuda_visible_devices }}\n WANDB_DISABLED: true\n run: |\n source activate accelerate\n pip install -r examples/pytorch/_tests_requirements.txt\n pytest -sv examples/pytorch/test_accelerate_examples.py examples/pytorch/test_pytorch_examples.py\n\n run-skorch-tests:\n container:\n image: huggingface/accelerate:gpu-nightly\n options: --gpus all --shm-size \"16gb\"\n runs-on:\n group: aws-g6-12xlarge-plus\n strategy:\n fail-fast: false\n steps:\n - name: Install accelerate\n run:\n source activate accelerate;\n git clone https://github.com/huggingface/accelerate;\n cd accelerate;\n git checkout ${{ github.sha }};\n pip install -e .[testing];\n cd ..\n\n - name: Install skorch\n run: |\n source activate accelerate\n git clone https://github.com/skorch-dev/skorch;\n cd skorch;\n git config --global --add safe.directory '*'\n git checkout master && git pull\n pip install .[test]\n pip install flaky\n\n - name: Show installed libraries\n run: |\n source activate accelerate;\n pip freeze\n\n - name: Run skorch tests\n working-directory: skorch/\n run: |\n source activate accelerate;\n pytest -sv -k TestAccelerate\n" }, { "path": ".github/workflows/stale.yml", "content": "name: Stale Bot\n\non:\n schedule:\n - cron: \"0 15 * * *\"\n workflow_dispatch:\n\njobs:\n close_stale_issues:\n name: Close Stale Issues\n if: github.repository == 'huggingface/accelerate'\n runs-on: ubuntu-latest\n permissions:\n issues: write\n pull-requests: write\n env:\n GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}\n steps:\n - uses: actions/checkout@v6\n \n - name: Setup Python\n uses: actions/setup-python@v6\n with:\n python-version: '3.10'\n cache: 'pip'\n cache-dependency-path: 'setup.py'\n \n - name: Install requirements\n run: |\n pip install PyGithub\n - name: Close stale issues\n run: |\n python utils/stale.py\n" }, { "path": ".github/workflows/test.yml", "content": "name: Run Tests\n\non:\n pull_request:\n paths:\n - \"src/**\"\n - \"tests/**\"\n - \".github/**\"\n - \"examples/**\"\n - \"setup.py\"\n types: [opened, synchronize, reopened]\n\nenv:\n HF_HOME: ~/hf_cache\n TESTING_MOCKED_DATALOADERS: \"1\"\n IS_GITHUB_CI: \"1\"\n\njobs:\n run-tests:\n runs-on:\n group: aws-general-8-plus\n strategy:\n fail-fast: false\n matrix:\n pytorch-version: [\n latest,\n minimum,\n ]\n test-kind: [\n test_prod,\n test_core,\n test_cli,\n test_big_modeling,\n test_deepspeed,\n test_fsdp,\n test_example_differences,\n test_checkpoint_step,\n test_checkpoint_epoch,\n test_rest\n ]\n steps:\n - uses: actions/checkout@v6\n - name: Set up python 3.10\n uses: actions/setup-python@v6\n with:\n python-version: '3.10'\n cache: 'pip'\n cache-dependency-path: 'setup.py'\n \n - name: Install the library\n run: |\n if [[ ${{ matrix.test-kind }} = test_prod ]]; then pip install -e .[test_prod]; fi\n if [[ ${{ matrix.test-kind }} != test_prod ]]; then pip install -e .[testing,test_trackers]; fi\n if [[ ${{ matrix.test-kind }} = test_rest ]]; then pip uninstall comet_ml -y; fi\n if [[ ${{ matrix.pytorch-version }} = minimum ]]; then pip install torchvision==0.19.0 torch==2.4.0; fi\n pip install pytest-reportlog tabulate setuptools importlib_metadata\n\n - name: Show installed libraries\n run: |\n pip freeze\n \n - name: Run Tests\n env: \n PYTORCH_VERSION: ${{ matrix.pytorch-version }}\n run: |\n make ${{ matrix.test-kind }}\n\n - name: Generate Report\n if: always()\n run: |\n python utils/log_reports.py >> $GITHUB_STEP_SUMMARY\n" }, { "path": ".github/workflows/test_imports.yml", "content": "name: Run Import Tests\n\non:\n pull_request:\n paths:\n - \"src/**\"\n - \"tests/**\"\n - \".github/**\"\n - \"examples/**\"\n - \"setup.py\"\n types: [opened, synchronize, reopened]\n\nenv:\n HF_HOME: ~/hf_cache\n TESTING_MOCKED_DATALOADERS: \"1\"\n IS_GITHUB_CI: \"1\"\n\njobs:\n run-tests:\n runs-on: ubuntu-latest\n strategy:\n fail-fast: false\n matrix:\n pytorch-version: [\n latest,\n minimum,\n ]\n steps:\n - uses: actions/checkout@v6\n - name: Set up python 3.10\n uses: actions/setup-python@v6\n with:\n python-version: '3.10'\n cache: 'pip'\n cache-dependency-path: 'setup.py'\n \n - name: Install the library\n run: |\n pip install -e .\n pip install pytest-reportlog tabulate setuptools git+https://github.com/muellerzr/import-timer\n\n - name: Show installed libraries\n run: |\n pip freeze\n \n - name: Run Import Tests\n env: \n PYTORCH_VERSION: ${{ matrix.pytorch-version }}\n run: |\n pytest -sv tests/test_imports.py\n\n - name: Generate Report\n if: always()\n run: |\n python utils/log_reports.py >> $GITHUB_STEP_SUMMARY\n" }, { "path": ".github/workflows/trufflehog.yml", "content": "on:\n push:\n\nname: Secret Leaks\n\njobs:\n trufflehog:\n runs-on: ubuntu-latest\n steps:\n - name: Checkout code\n uses: actions/checkout@v6\n with:\n fetch-depth: 0\n - name: Secret Scanning\n uses: trufflesecurity/trufflehog@main\n" }, { "path": ".github/workflows/upload_pr_documentation.yml", "content": "name: Upload PR Documentation\n\non:\n workflow_run:\n workflows: [\"Build PR Documentation\"]\n types:\n - completed\n\njobs:\n build:\n uses: huggingface/doc-builder/.github/workflows/upload_pr_documentation.yml@main\n with:\n package_name: accelerate\n secrets:\n hf_token: ${{ secrets.HF_DOC_BUILD_PUSH }}\n comment_bot_token: ${{ secrets.COMMENT_BOT_TOKEN }}" }, { "path": ".gitignore", "content": "# Byte-compiled / optimized / DLL files\n__pycache__/\n*.py[cod]\n*$py.class\n\n# C extensions\n*.so\n\n# Distribution / packaging\n.Python\nbuild/\ndevelop-eggs/\ndist/\ndownloads/\neggs/\n.eggs/\nlib/\nlib64/\nparts/\nsdist/\nvar/\nwheels/\npip-wheel-metadata/\nshare/python-wheels/\n*.egg-info/\n.installed.cfg\n*.egg\nMANIFEST\n\n# PyInstaller\n# Usually these files are written by a python script from a template\n# before PyInstaller builds the exe, so as to inject date/other infos into it.\n*.manifest\n*.spec\n\n# Installer logs\npip-log.txt\npip-delete-this-directory.txt\n\n# Unit test / coverage reports\nhtmlcov/\n.tox/\n.nox/\n.coverage\n.coverage.*\n.cache\nnosetests.xml\ncoverage.xml\n*.cover\n*.py,cover\n.hypothesis/\n.pytest_cache/\n\n# Translations\n*.mo\n*.pot\n\n# Django stuff:\n*.log\nlocal_settings.py\ndb.sqlite3\ndb.sqlite3-journal\n\n# Flask stuff:\ninstance/\n.webassets-cache\n\n# Scrapy stuff:\n.scrapy\n\n# Sphinx documentation\ndocs/_build/\n\n# PyBuilder\ntarget/\n\n# Jupyter Notebook\n.ipynb_checkpoints\n\n# IPython\nprofile_default/\nipython_config.py\n\n# pyenv\n.python-version\n\n# pipenv\n# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.\n# However, in case of collaboration, if having platform-specific dependencies or dependencies\n# having no cross-platform support, pipenv may install dependencies that don't work, or not\n# install all needed dependencies.\n#Pipfile.lock\n\n# PEP 582; used by e.g. github.com/David-OConnor/pyflow\n__pypackages__/\n\n# Celery stuff\ncelerybeat-schedule\ncelerybeat.pid\n\n# SageMath parsed files\n*.sage.py\n\n# Environments\n.env\n.venv\nenv/\nvenv/\nENV/\nenv.bak/\nvenv.bak/\n\n# Spyder project settings\n.spyderproject\n.spyproject\n\n# Rope project settings\n.ropeproject\n\n# mkdocs documentation\n/site\n\n# mypy\n.mypy_cache/\n.dmypy.json\ndmypy.json\n\n# Pyre type checker\n.pyre/\n\n# VSCode\n.vscode\n\n# IntelliJ\n.idea\n\n# Mac .DS_Store\n.DS_Store\n\n# More test things\nwandb\n\n# ruff\n.ruff_cache\n" }, { "path": ".pre-commit-config.yaml", "content": "repos:\n - repo: https://github.com/astral-sh/ruff-pre-commit\n rev: v0.2.1\n hooks:\n - id: ruff\n args:\n - --fix\n - id: ruff-format\n - repo: https://github.com/pre-commit/pre-commit-hooks\n rev: v4.5.0\n hooks:\n - id: check-merge-conflict\n - id: check-yaml\n" }, { "path": "CODE_OF_CONDUCT.md", "content": "\n# Contributor Covenant Code of Conduct\n\n## Our Pledge\n\nWe as members, contributors, and leaders pledge to make participation in our\ncommunity a harassment-free experience for everyone, regardless of age, body\nsize, visible or invisible disability, ethnicity, sex characteristics, gender\nidentity and expression, level of experience, education, socio-economic status,\nnationality, personal appearance, race, religion, or sexual identity\nand orientation.\n\nWe pledge to act and interact in ways that contribute to an open, welcoming,\ndiverse, inclusive, and healthy community.\n\n## Our Standards\n\nExamples of behavior that contributes to a positive environment for our\ncommunity include:\n\n* Demonstrating empathy and kindness toward other people\n* Being respectful of differing opinions, viewpoints, and experiences\n* Giving and gracefully accepting constructive feedback\n* Accepting responsibility and apologizing to those affected by our mistakes,\n and learning from the experience\n* Focusing on what is best not just for us as individuals, but for the\n overall community\n\nExamples of unacceptable behavior include:\n\n* The use of sexualized language or imagery, and sexual attention or\n advances of any kind\n* Trolling, insulting or derogatory comments, and personal or political attacks\n* Public or private harassment\n* Publishing others' private information, such as a physical or email\n address, without their explicit permission\n* Other conduct which could reasonably be considered inappropriate in a\n professional setting\n\n## Enforcement Responsibilities\n\nCommunity leaders are responsible for clarifying and enforcing our standards of\nacceptable behavior and will take appropriate and fair corrective action in\nresponse to any behavior that they deem inappropriate, threatening, offensive,\nor harmful.\n\nCommunity leaders have the right and responsibility to remove, edit, or reject\ncomments, commits, code, wiki edits, issues, and other contributions that are\nnot aligned to this Code of Conduct, and will communicate reasons for moderation\ndecisions when appropriate.\n\n## Scope\n\nThis Code of Conduct applies within all community spaces, and also applies when\nan individual is officially representing the community in public spaces.\nExamples of representing our community include using an official e-mail address,\nposting via an official social media account, or acting as an appointed\nrepresentative at an online or offline event.\n\n## Enforcement\n\nInstances of abusive, harassing, or otherwise unacceptable behavior may be\nreported to the community leaders responsible for enforcement at\nfeedback@huggingface.co.\nAll complaints will be reviewed and investigated promptly and fairly.\n\nAll community leaders are obligated to respect the privacy and security of the\nreporter of any incident.\n\n## Enforcement Guidelines\n\nCommunity leaders will follow these Community Impact Guidelines in determining\nthe consequences for any action they deem in violation of this Code of Conduct:\n\n### 1. Correction\n\n**Community Impact**: Use of inappropriate language or other behavior deemed\nunprofessional or unwelcome in the community.\n\n**Consequence**: A private, written warning from community leaders, providing\nclarity around the nature of the violation and an explanation of why the\nbehavior was inappropriate. A public apology may be requested.\n\n### 2. Warning\n\n**Community Impact**: A violation through a single incident or series\nof actions.\n\n**Consequence**: A warning with consequences for continued behavior. No\ninteraction with the people involved, including unsolicited interaction with\nthose enforcing the Code of Conduct, for a specified period of time. This\nincludes avoiding interactions in community spaces as well as external channels\nlike social media. Violating these terms may lead to a temporary or\npermanent ban.\n\n### 3. Temporary Ban\n\n**Community Impact**: A serious violation of community standards, including\nsustained inappropriate behavior.\n\n**Consequence**: A temporary ban from any sort of interaction or public\ncommunication with the community for a specified period of time. No public or\nprivate interaction with the people involved, including unsolicited interaction\nwith those enforcing the Code of Conduct, is allowed during this period.\nViolating these terms may lead to a permanent ban.\n\n### 4. Permanent Ban\n\n**Community Impact**: Demonstrating a pattern of violation of community\nstandards, including sustained inappropriate behavior, harassment of an\nindividual, or aggression toward or disparagement of classes of individuals.\n\n**Consequence**: A permanent ban from any sort of public interaction within\nthe community.\n\n## Attribution\n\nThis Code of Conduct is adapted from the [Contributor Covenant][homepage],\nversion 2.0, available at\nhttps://www.contributor-covenant.org/version/2/0/code_of_conduct.html.\n\nCommunity Impact Guidelines were inspired by [Mozilla's code of conduct\nenforcement ladder](https://github.com/mozilla/diversity).\n\n[homepage]: https://www.contributor-covenant.org\n\nFor answers to common questions about this code of conduct, see the FAQ at\nhttps://www.contributor-covenant.org/faq. Translations are available at\nhttps://www.contributor-covenant.org/translations.\n" }, { "path": "CONTRIBUTING.md", "content": "\n\n# How to contribute to šŸ¤— Accelerate?\n\nEveryone is welcome to contribute, and we value everybody's contribution. Code\nis thus not the only way to help the community. Answering questions, helping\nothers, reaching out and improving the documentations are immensely valuable to\nthe community.\n\nIt also helps us if you spread the word: reference the library from blog posts\non the awesome projects it made possible, shout out on Twitter every time it has\nhelped you, or simply star the repo to say \"thank you\".\n\nWhichever way you choose to contribute, please be mindful to respect our\n[code of conduct](https://github.com/huggingface/accelerate/blob/main/CODE_OF_CONDUCT.md).\n\n## You can contribute in so many ways!\n\nSome of the ways you can contribute to Accelerate:\n* Fixing outstanding issues with the existing code;\n* Contributing to the examples or to the documentation;\n* Submitting issues related to bugs or desired new features.\n\n## Submitting a new issue or feature request\n\nDo your best to follow these guidelines when submitting an issue or a feature\nrequest. It will make it easier for us to come back to you quickly and with good\nfeedback.\n\n### Did you find a bug?\n\nThe šŸ¤— Accelerate library is robust and reliable thanks to the users who notify us of\nthe problems they encounter. So thank you for reporting an issue.\n\nFirst, we would really appreciate it if you could **make sure the bug was not\nalready reported** (use the search bar on Github under Issues).\n\nDid not find it? :( So we can act quickly on it, please follow these steps:\n\n* Include your **OS type and version**, the versions of **Python** and **PyTorch**.\n* A short, self-contained, code snippet that allows us to reproduce the bug in\n less than 30s;\n* Provide the with your Accelerate configuration (located by default in `~/.cache/huggingface/accelerate/default_config.yaml`)\n\n### Do you want a new feature?\n\nA good feature request addresses the following points:\n\n1. Motivation first:\n* Is it related to a problem/frustration with the library? If so, please explain\n why. Providing a code snippet that demonstrates the problem is best.\n* Is it related to something you would need for a project? We'd love to hear\n about it!\n* Is it something you worked on and think could benefit the community?\n Awesome! Tell us what problem it solved for you.\n2. Write a *full paragraph* describing the feature;\n3. Provide a **code snippet** that demonstrates its future use;\n4. In case this is related to a paper, please attach a link;\n5. Attach any additional information (drawings, screenshots, etc.) you think may help.\n\nIf your issue is well written we're already 80% of the way there by the time you\npost it.\n\n## Submitting a pull request (PR)\n\nBefore writing code, we strongly advise you to search through the existing PRs or\nissues to make sure that nobody is already working on the same thing. If you are\nunsure, it is always a good idea to open an issue to get some feedback.\n\nYou will need basic `git` proficiency to be able to contribute to\nšŸ¤— Accelerate. `git` is not the easiest tool to use but it has the greatest\nmanual. Type `git --help` in a shell and enjoy. If you prefer books, [Pro\nGit](https://git-scm.com/book/en/v2) is a very good reference.\n\nFollow these steps to start contributing:\n\n1. Fork the [repository](https://github.com/huggingface/accelerate) by\n clicking on the 'Fork' button on the repository's page. This creates a copy of the code\n under your GitHub user account.\n\n2. Clone your fork to your local disk, and add the base repository as a remote. The following command\n assumes you have your public SSH key uploaded to GitHub. See the following guide for more\n [information](https://docs.github.com/en/repositories/creating-and-managing-repositories/cloning-a-repository).\n\n ```bash\n $ git clone git@github.com:/accelerate.git\n $ cd accelerate\n $ git remote add upstream https://github.com/huggingface/accelerate.git\n ```\n\n3. Create a new branch to hold your development changes, and do this for every new PR you work on.\n\n Start by synchronizing your `main` branch with the `upstream/main` branch (ore details in the [GitHub Docs](https://docs.github.com/en/github/collaborating-with-issues-and-pull-requests/syncing-a-fork)):\n\n ```bash\n $ git checkout main\n $ git fetch upstream\n $ git merge upstream/main\n ```\n\n Once your `main` branch is synchronized, create a new branch from it:\n\n ```bash\n $ git checkout -b a-descriptive-name-for-my-changes\n ```\n\n **Do not** work on the `main` branch.\n\n4. Set up a development environment by running the following command in a conda or a virtual environment you've created for working on this library:\n\n ```bash\n $ pip install -e \".[dev]\"\n ```\n \n This will install all testing and linting/code quality dependencies for the library (see `quality`, `test_dev`, \n `test_prod` targets in [`setup.py`](./setup.py)).\n\n (If accelerate was already installed in the virtual environment, remove\n it with `pip uninstall accelerate` before reinstalling it in editable\n mode with the `-e` flag).\n\n Alternatively, if you are using [Visual Studio Code](https://code.visualstudio.com/Download), the fastest way to get set up is by using\n the provided Dev Container. Documentation on how to get started with dev containers is available [here](https://code.visualstudio.com/docs/remote/containers).\n\n5. Develop the features on your branch.\n\n As you work on the features, you should make sure that the test suite\n passes. You should run the tests impacted by your changes like this (see \n below an explanation regarding the environment variable):\n\n ```bash\n $ pytest tests/.py\n ```\n \n > For the following commands leveraging the `make` utility, we recommend using the WSL system when running on\n > Windows. More information [here](https://docs.microsoft.com/en-us/windows/wsl/about).\n\n You can also run the full suite with the following command.\n\n ```bash\n $ make test\n ```\n\n `accelerate` relies on `ruff` to format its source code\n consistently. After you make changes, apply automatic style corrections and code verifications\n that can't be automated in one go with:\n\n This target is also optimized to only work with files modified by the PR you're working on.\n\n If you prefer to run the checks one after the other, the following command apply the\n style corrections:\n\n ```bash\n $ make style\n ```\n\n `accelerate` also uses a few custom scripts to check for coding mistakes. Quality\n control runs in CI, however you can also run the same checks with:\n\n ```bash\n $ make quality\n ```\n\n You can also set up [`pre-commit`](https://pre-commit.com/) to run these checks\n automatically as Git commit hooks.\n\n ```bash\n $ pip install pre-commit\n $ pre-commit install\n ```\n\n Once you're happy with your changes, add changed files using `git add` and\n make a commit with `git commit` to record your changes locally:\n\n ```bash\n $ git add modified_file.py\n $ git commit\n ```\n\n Please write [good commit messages](https://chris.beams.io/posts/git-commit/).\n\n It is a good idea to sync your copy of the code with the original\n repository regularly. This way you can quickly account for changes:\n\n ```bash\n $ git fetch upstream\n $ git rebase upstream/main\n ```\n\n Push the changes to your account using:\n\n ```bash\n $ git push -u origin a-descriptive-name-for-my-changes\n ```\n\n6. Once you are satisfied (**and the checklist below is happy too**), go to the\n webpage of your fork on GitHub. Click on 'Pull request' to send your changes\n to the project maintainers for review.\n\n7. It's ok if maintainers ask you for changes. It happens to core contributors\n too! So everyone can see the changes in the Pull request, work in your local\n branch and push the changes to your fork. They will automatically appear in\n the pull request.\n\n\n### Checklist\n\n1. The title of your pull request should be a summary of its contribution;\n2. If your pull request addresses an issue, please mention the issue number in\n the pull request description to make sure they are linked (and people\n consulting the issue know you are working on it);\n3. To indicate a work in progress please prefix the title with `[WIP]`, or mark\n the PR as a draft PR. These are useful to avoid duplicated work, and to differentiate\n it from PRs ready to be merged;\n4. Make sure existing tests pass;\n5. Add high-coverage tests. No quality testing = no merge.\n\nSee an example of a good PR here: https://github.com/huggingface/accelerate/pull/255\n\n### Tests\n\nAn extensive test suite is included to test the library behavior and several examples. Library tests can be found in\nthe [tests folder](https://github.com/huggingface/accelerate/tree/main/tests).\n\nWe use `pytest` in order to run the tests. From the root of the\nrepository, here's how to run tests with `pytest` for the library:\n\n```bash\n$ python -m pytest -sv ./tests\n```\n\nIn fact, that's how `make test` is implemented (sans the `pip install` line)!\n\nYou can specify a smaller set of tests in order to test only the feature\nyou're working on.\n" }, { "path": "LICENSE", "content": " Apache License\n Version 2.0, January 2004\n http://www.apache.org/licenses/\n\n TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION\n\n 1. Definitions.\n\n \"License\" shall mean the terms and conditions for use, reproduction,\n and distribution as defined by Sections 1 through 9 of this document.\n\n \"Licensor\" shall mean the copyright owner or entity authorized by\n the copyright owner that is granting the License.\n\n \"Legal Entity\" shall mean the union of the acting entity and all\n other entities that control, are controlled by, or are under common\n control with that entity. For the purposes of this definition,\n \"control\" means (i) the power, direct or indirect, to cause the\n direction or management of such entity, whether by contract or\n otherwise, or (ii) ownership of fifty percent (50%) or more of the\n outstanding shares, or (iii) beneficial ownership of such entity.\n\n \"You\" (or \"Your\") shall mean an individual or Legal Entity\n exercising permissions granted by this License.\n\n \"Source\" form shall mean the preferred form for making modifications,\n including but not limited to software source code, documentation\n source, and configuration files.\n\n \"Object\" form shall mean any form resulting from mechanical\n transformation or translation of a Source form, including but\n not limited to compiled object code, generated documentation,\n and conversions to other media types.\n\n \"Work\" shall mean the work of authorship, whether in Source or\n Object form, made available under the License, as indicated by a\n copyright notice that is included in or attached to the work\n (an example is provided in the Appendix below).\n\n \"Derivative Works\" shall mean any work, whether in Source or Object\n form, that is based on (or derived from) the Work and for which the\n editorial revisions, annotations, elaborations, or other modifications\n represent, as a whole, an original work of authorship. For the purposes\n of this License, Derivative Works shall not include works that remain\n separable from, or merely link (or bind by name) to the interfaces of,\n the Work and Derivative Works thereof.\n\n \"Contribution\" shall mean any work of authorship, including\n the original version of the Work and any modifications or additions\n to that Work or Derivative Works thereof, that is intentionally\n submitted to Licensor for inclusion in the Work by the copyright owner\n or by an individual or Legal Entity authorized to submit on behalf of\n the copyright owner. For the purposes of this definition, \"submitted\"\n means any form of electronic, verbal, or written communication sent\n to the Licensor or its representatives, including but not limited to\n communication on electronic mailing lists, source code control systems,\n and issue tracking systems that are managed by, or on behalf of, the\n Licensor for the purpose of discussing and improving the Work, but\n excluding communication that is conspicuously marked or otherwise\n designated in writing by the copyright owner as \"Not a Contribution.\"\n\n \"Contributor\" shall mean Licensor and any individual or Legal Entity\n on behalf of whom a Contribution has been received by Licensor and\n subsequently incorporated within the Work.\n\n 2. Grant of Copyright License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n copyright license to reproduce, prepare Derivative Works of,\n publicly display, publicly perform, sublicense, and distribute the\n Work and such Derivative Works in Source or Object form.\n\n 3. Grant of Patent License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n (except as stated in this section) patent license to make, have made,\n use, offer to sell, sell, import, and otherwise transfer the Work,\n where such license applies only to those patent claims licensable\n by such Contributor that are necessarily infringed by their\n Contribution(s) alone or by combination of their Contribution(s)\n with the Work to which such Contribution(s) was submitted. If You\n institute patent litigation against any entity (including a\n cross-claim or counterclaim in a lawsuit) alleging that the Work\n or a Contribution incorporated within the Work constitutes direct\n or contributory patent infringement, then any patent licenses\n granted to You under this License for that Work shall terminate\n as of the date such litigation is filed.\n\n 4. Redistribution. You may reproduce and distribute copies of the\n Work or Derivative Works thereof in any medium, with or without\n modifications, and in Source or Object form, provided that You\n meet the following conditions:\n\n (a) You must give any other recipients of the Work or\n Derivative Works a copy of this License; and\n\n (b) You must cause any modified files to carry prominent notices\n stating that You changed the files; and\n\n (c) You must retain, in the Source form of any Derivative Works\n that You distribute, all copyright, patent, trademark, and\n attribution notices from the Source form of the Work,\n excluding those notices that do not pertain to any part of\n the Derivative Works; and\n\n (d) If the Work includes a \"NOTICE\" text file as part of its\n distribution, then any Derivative Works that You distribute must\n include a readable copy of the attribution notices contained\n within such NOTICE file, excluding those notices that do not\n pertain to any part of the Derivative Works, in at least one\n of the following places: within a NOTICE text file distributed\n as part of the Derivative Works; within the Source form or\n documentation, if provided along with the Derivative Works; or,\n within a display generated by the Derivative Works, if and\n wherever such third-party notices normally appear. The contents\n of the NOTICE file are for informational purposes only and\n do not modify the License. You may add Your own attribution\n notices within Derivative Works that You distribute, alongside\n or as an addendum to the NOTICE text from the Work, provided\n that such additional attribution notices cannot be construed\n as modifying the License.\n\n You may add Your own copyright statement to Your modifications and\n may provide additional or different license terms and conditions\n for use, reproduction, or distribution of Your modifications, or\n for any such Derivative Works as a whole, provided Your use,\n reproduction, and distribution of the Work otherwise complies with\n the conditions stated in this License.\n\n 5. Submission of Contributions. Unless You explicitly state otherwise,\n any Contribution intentionally submitted for inclusion in the Work\n by You to the Licensor shall be under the terms and conditions of\n this License, without any additional terms or conditions.\n Notwithstanding the above, nothing herein shall supersede or modify\n the terms of any separate license agreement you may have executed\n with Licensor regarding such Contributions.\n\n 6. Trademarks. This License does not grant permission to use the trade\n names, trademarks, service marks, or product names of the Licensor,\n except as required for reasonable and customary use in describing the\n origin of the Work and reproducing the content of the NOTICE file.\n\n 7. Disclaimer of Warranty. Unless required by applicable law or\n agreed to in writing, Licensor provides the Work (and each\n Contributor provides its Contributions) on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or\n implied, including, without limitation, any warranties or conditions\n of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A\n PARTICULAR PURPOSE. You are solely responsible for determining the\n appropriateness of using or redistributing the Work and assume any\n risks associated with Your exercise of permissions under this License.\n\n 8. Limitation of Liability. In no event and under no legal theory,\n whether in tort (including negligence), contract, or otherwise,\n unless required by applicable law (such as deliberate and grossly\n negligent acts) or agreed to in writing, shall any Contributor be\n liable to You for damages, including any direct, indirect, special,\n incidental, or consequential damages of any character arising as a\n result of this License or out of the use or inability to use the\n Work (including but not limited to damages for loss of goodwill,\n work stoppage, computer failure or malfunction, or any and all\n other commercial damages or losses), even if such Contributor\n has been advised of the possibility of such damages.\n\n 9. Accepting Warranty or Additional Liability. While redistributing\n the Work or Derivative Works thereof, You may choose to offer,\n and charge a fee for, acceptance of support, warranty, indemnity,\n or other liability obligations and/or rights consistent with this\n License. However, in accepting such obligations, You may act only\n on Your own behalf and on Your sole responsibility, not on behalf\n of any other Contributor, and only if You agree to indemnify,\n defend, and hold each Contributor harmless for any liability\n incurred by, or claims asserted against, such Contributor by reason\n of your accepting any such warranty or additional liability.\n\n END OF TERMS AND CONDITIONS\n\n APPENDIX: How to apply the Apache License to your work.\n\n To apply the Apache License to your work, attach the following\n boilerplate notice, with the fields enclosed by brackets \"[]\"\n replaced with your own identifying information. (Don't include\n the brackets!) The text should be enclosed in the appropriate\n comment syntax for the file format. We also recommend that a\n file or class name and description of purpose be included on the\n same \"printed page\" as the copyright notice for easier\n identification within third-party archives.\n\n Copyright [yyyy] [name of copyright owner]\n\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n\n http://www.apache.org/licenses/LICENSE-2.0\n\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n" }, { "path": "Makefile", "content": ".PHONY: quality style test docs utils\n\ncheck_dirs := .\n\n# Check that source code meets quality standards\n\nextra_quality_checks:\n\tpython utils/check_copies.py\n\tpython utils/check_dummies.py\n\tpython utils/check_repo.py\n\n# this target runs checks on all files\nquality:\n\truff check $(check_dirs)\n\truff format --check $(check_dirs)\n\n# Format source code automatically and check is there are any problems left that need manual fixing\nstyle:\n\truff check $(check_dirs) --fix\n\truff format $(check_dirs)\n\t\n# Run tests for the library\ntest_core:\n\tpython -m pytest -s -v ./tests/ \\\n\t--ignore=./tests/test_big_modeling.py \\\n\t--ignore=./tests/test_modeling_utils.py \\\n\t--ignore=./tests/test_examples.py \\\n\t--ignore=./tests/test_cli.py \\\n\t--ignore=./tests/deepspeed \\\n\t--ignore=./tests/fsdp \\\n\t--ignore=./tests/tp \\\n\t$(if $(IS_GITHUB_CI),--report-log \"$(PYTORCH_VERSION)_core.log\",)\n\ntest_cli:\n\tpython -m pytest -s -v ./tests/test_cli.py $(if $(IS_GITHUB_CI),--report-log \"$(PYTORCH_VERSION)_cli.log\",)\n\ntest_big_modeling:\n\tpython -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\",)\n\ntest_deepspeed:\n\tpython -m pytest -s -v ./tests/deepspeed $(if $(IS_GITHUB_CI),--report-log \"$(PYTORCH_VERSION)_deepspeed.log\",)\n\ntest_fsdp:\n\tpython -m pytest -s -v ./tests/fsdp $(if $(IS_GITHUB_CI),--report-log \"$(PYTORCH_VERSION)_fsdp.log\",)\n\ntest_tp:\n\tpython -m pytest -s -v ./tests/tp $(if $(IS_GITHUB_CI),--report-log \"$(PYTORCH_VERSION)_tp.log\",)\n\n# Since the new version of pytest will *change* how things are collected, we need `deepspeed` to \n# run after test_core and test_cli\ntest:\n\t$(MAKE) test_core\n\t$(MAKE) test_cli\n\t$(MAKE) test_big_modeling\n\t$(MAKE) test_deepspeed\n\t$(MAKE) test_fsdp\n\t$(MAKE) test_tp\n\ntest_examples:\n\tpython -m pytest -s -v ./tests/test_examples.py $(if $(IS_GITHUB_CI),--report-log \"$(PYTORCH_VERSION)_examples.log\",)\n\n# Broken down example tests for the CI runners\ntest_integrations:\n\tpython -m pytest -s -v ./tests/fsdp ./tests/tp ./tests/deepspeed $(if $(IS_GITHUB_CI),--report-log \"$(PYTORCH_VERSION)_integrations.log\",)\n\ntest_example_differences:\n\tpython -m pytest -s -v ./tests/test_examples.py::ExampleDifferenceTests $(if $(IS_GITHUB_CI),--report-log \"$(PYTORCH_VERSION)_example_diff.log\",)\n\ntest_checkpoint_epoch:\n\tpython -m pytest -s -v ./tests/test_examples.py::FeatureExamplesTests -k \"by_epoch\" $(if $(IS_GITHUB_CI),--report-log \"$(PYTORCH_VERSION)_checkpoint_epoch.log\",)\n\ntest_checkpoint_step:\n\tpython -m pytest -s -v ./tests/test_examples.py::FeatureExamplesTests -k \"by_step\" $(if $(IS_GITHUB_CI),--report-log \"$(PYTORCH_VERSION)_checkpoint_step.log\",)\n\n# Same as test but used to install only the base dependencies\ntest_prod:\n\t$(MAKE) test_core\n\ntest_rest:\n\tpython -m pytest -s -v ./tests/test_examples.py::FeatureExamplesTests -k \"not by_step and not by_epoch\" $(if $(IS_GITHUB_CI),--report-log \"$(PYTORCH_VERSION)_rest.log\",)\n\n# For developers to prepare a release\nprepare_release:\n\trm -rf dist build\n\tpython setup.py bdist_wheel sdist\n\n# Make sure this is ran in a fresh venv of some form\ninstall_test_release:\n\tpip uninstall accelerate -y\n\tpip install -i https://testpypi.python.org/pypi --extra-index-url https://pypi.org/simple accelerate$(if $(version),==$(version),)\n\n# Run as `make target=testpypi upload_release`\nupload_release:\n\t@if [ \"$(target)\" != \"testpypi\" ] && [ \"$(target)\" != \"pypi\" ]; then \\\n\t\techo \"Error: target must be either 'testpypi' or 'pypi'\"; \\\n\t\texit 1; \\\n\tfi\n\ttwine upload dist/* -r $(target)" }, { "path": "README.md", "content": "\n\n

\n
\n \n
\n

\n\n

\n \n \"License\"\n \"Documentation\"\n \"GitHub\n \"Contributor\n

\n\n

\n

Run your *raw* PyTorch training script on any kind of device\n

\n\n

\n \n

\n\n## Easy to integrate\n\nšŸ¤— Accelerate was created for PyTorch users who like to write the training loop of PyTorch models but are reluctant to write and maintain the boilerplate code needed to use multi-GPUs/TPU/fp16.\n\nšŸ¤— Accelerate abstracts exactly and only the boilerplate code related to multi-GPUs/TPU/fp16 and leaves the rest of your code unchanged.\n\nHere is an example:\n\n```diff\n import torch\n import torch.nn.functional as F\n from datasets import load_dataset\n+ from accelerate import Accelerator\n\n+ accelerator = Accelerator()\n- device = 'cpu'\n+ device = accelerator.device\n\n model = torch.nn.Transformer().to(device)\n optimizer = torch.optim.Adam(model.parameters())\n\n dataset = load_dataset('my_dataset')\n data = torch.utils.data.DataLoader(dataset, shuffle=True)\n\n+ model, optimizer, data = accelerator.prepare(model, optimizer, data)\n\n model.train()\n for epoch in range(10):\n for source, targets in data:\n source = source.to(device)\n targets = targets.to(device)\n\n optimizer.zero_grad()\n\n output = model(source)\n loss = F.cross_entropy(output, targets)\n\n- loss.backward()\n+ accelerator.backward(loss)\n\n optimizer.step()\n```\n\nAs 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).\n\nIn particular, the same code can then be run without modification on your local machine for debugging or your training environment.\n\nšŸ¤— Accelerate even handles the device placement for you (which requires a few more changes to your code, but is safer in general), so you can even simplify your training loop further:\n\n```diff\n import torch\n import torch.nn.functional as F\n from datasets import load_dataset\n+ from accelerate import Accelerator\n\n- device = 'cpu'\n+ accelerator = Accelerator()\n\n- model = torch.nn.Transformer().to(device)\n+ model = torch.nn.Transformer()\n optimizer = torch.optim.Adam(model.parameters())\n\n dataset = load_dataset('my_dataset')\n data = torch.utils.data.DataLoader(dataset, shuffle=True)\n\n+ model, optimizer, data = accelerator.prepare(model, optimizer, data)\n\n model.train()\n for epoch in range(10):\n for source, targets in data:\n- source = source.to(device)\n- targets = targets.to(device)\n\n optimizer.zero_grad()\n\n output = model(source)\n loss = F.cross_entropy(output, targets)\n\n- loss.backward()\n+ accelerator.backward(loss)\n\n optimizer.step()\n```\n\nWant 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).\n\n## Launching script\n\nšŸ¤— 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!\nOn your machine(s) just run:\n\n```bash\naccelerate config\n```\n\nand answer the questions asked. This will generate a config file that will be used automatically to properly set the default options when doing\n\n```bash\naccelerate launch my_script.py --args_to_my_script\n``` \n\nFor instance, here is how you would run the GLUE example on the MRPC task (from the root of the repo):\n\n```bash\naccelerate launch examples/nlp_example.py\n```\n\nThis CLI tool is **optional**, and you can still use `python my_script.py` or `python -m torchrun my_script.py` at your convenience.\n\nYou can also directly pass in the arguments you would to `torchrun` as arguments to `accelerate launch` if you wish to not run` accelerate config`.\n\nFor example, here is how to launch on two GPUs:\n\n```bash\naccelerate launch --multi_gpu --num_processes 2 examples/nlp_example.py\n```\n\nTo learn more, check the CLI documentation available [here](https://huggingface.co/docs/accelerate/package_reference/cli).\n\nOr view the configuration zoo [here](https://github.com/huggingface/accelerate/blob/main/examples/config_yaml_templates/)\n\n## Launching multi-CPU run using MPI\n\nšŸ¤— Here is another way to launch multi-CPU run using MPI. You can learn how to install Open MPI on [this page](https://www.open-mpi.org/faq/?category=building#easy-build). You can use Intel MPI or MVAPICH as well.\nOnce you have MPI setup on your cluster, just run:\n```bash\naccelerate config\n```\nAnswer the questions that are asked, selecting to run using multi-CPU, and answer \"yes\" when asked if you want accelerate to launch mpirun.\nThen, use `accelerate launch` with your script like:\n```bash\naccelerate launch examples/nlp_example.py\n```\nAlternatively, you can use mpirun directly, without using the CLI like:\n```bash\nmpirun -np 2 python examples/nlp_example.py\n```\n\n## Launching training using DeepSpeed\n\nšŸ¤— 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`.\n\n```python\nfrom accelerate import Accelerator, DeepSpeedPlugin\n\n# deepspeed needs to know your gradient accumulation steps beforehand, so don't forget to pass it\n# Remember you still need to do gradient accumulation by yourself, just like you would have done without deepspeed\ndeepspeed_plugin = DeepSpeedPlugin(zero_stage=2, gradient_accumulation_steps=2)\naccelerator = Accelerator(mixed_precision='fp16', deepspeed_plugin=deepspeed_plugin)\n\n# How to save your šŸ¤— Transformer?\naccelerator.wait_for_everyone()\nunwrapped_model = accelerator.unwrap_model(model)\nunwrapped_model.save_pretrained(save_dir, save_function=accelerator.save, state_dict=accelerator.get_state_dict(model))\n```\n\nNote: DeepSpeed support is experimental for now. In case you get into some problem, please open an issue.\n\n## Launching your training from a notebook\n\nšŸ¤— Accelerate also provides a `notebook_launcher` function you can use in a notebook to launch a distributed training. This is especially useful for Colab or Kaggle notebooks with a TPU backend. Just define your training loop in a `training_function` then in your last cell, add:\n\n```python\nfrom accelerate import notebook_launcher\n\nnotebook_launcher(training_function)\n```\n\nAn example can be found in [this notebook](https://github.com/huggingface/notebooks/blob/main/examples/accelerate_examples/simple_nlp_example.ipynb). [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/notebooks/blob/main/examples/accelerate_examples/simple_nlp_example.ipynb)\n\n## Why should I use šŸ¤— Accelerate?\n\nYou 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.\n\n## Why shouldn't I use šŸ¤— Accelerate?\n\nYou shouldn't use šŸ¤— Accelerate if you don't want to write a training loop yourself. There are plenty of high-level libraries above PyTorch that will offer you that, šŸ¤— Accelerate is not one of them.\n\n## Frameworks using šŸ¤— Accelerate\n\nIf 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:\n\n* [Amphion](https://github.com/open-mmlab/Amphion) is a toolkit for Audio, Music, and Speech Generation. Its purpose is to support reproducible research and help junior researchers and engineers get started in the field of audio, music, and speech generation research and development.\n* [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).\n* [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.\n* [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.\n* [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.\n* [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.\n* [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.\n* [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. \n* [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!\n* [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.\n* [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.\n* [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).\n\n\n## Installation\n\nThis repository is tested on Python 3.8+ and PyTorch 1.10.0+\n\nYou 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/).\n\nFirst, create a virtual environment with the version of Python you're going to use and activate it.\n\nThen, you will need to install PyTorch: refer to the [official installation page](https://pytorch.org/get-started/locally/#start-locally) regarding the specific install command for your platform. Then šŸ¤— Accelerate can be installed using pip as follows:\n\n```bash\npip install accelerate\n```\n\n## Supported integrations\n\n- CPU only\n- multi-CPU on one node (machine)\n- multi-CPU on several nodes (machines)\n- single GPU\n- multi-GPU on one node (machine)\n- multi-GPU on several nodes (machines)\n- TPU\n- FP16/BFloat16 mixed precision\n- FP8 mixed precision with [Transformer Engine](https://github.com/NVIDIA/TransformerEngine) or [MS-AMP](https://github.com/Azure/MS-AMP/)\n- DeepSpeed support (Experimental)\n- PyTorch Fully Sharded Data Parallel (FSDP) support (Experimental)\n- Megatron-LM support (Experimental)\n\n## Citing šŸ¤— Accelerate\n\nIf you use šŸ¤— Accelerate in your publication, please cite it by using the following BibTeX entry.\n\n```bibtex\n@Misc{accelerate,\n title = {Accelerate: Training and inference at scale made simple, efficient and adaptable.},\n author = {Sylvain Gugger and Lysandre Debut and Thomas Wolf and Philipp Schmid and Zachary Mueller and Sourab Mangrulkar and Marc Sun and Benjamin Bossan},\n howpublished = {\\url{https://github.com/huggingface/accelerate}},\n year = {2022}\n}\n```\n" }, { "path": "benchmarks/README.md", "content": "# Benchmarks\n\nThe folders below contain suites to test various functionalities in Accelerate.\n\nSee their relevant README.md's for more information.\n" }, { "path": "benchmarks/big_model_inference/README.md", "content": "# Big model inference benchmarks\n\nRunning inference with Accelerate on big models.\n\n## Setup\n\nThese benchmarks use the `transformers` library:\n\n```bash\npip install transformers\n```\n\nTo reproduce or test a new setup, run\n\n```py\npython big_model_inference.py model_name\n```\n\nThis script supports `gpt-j-6b`, `gpt-neox`, `opt` (30B version) and `T0pp` out of the box, but you can specify any valid checkpoint for `model_name`.\n\nTo force a different `torch_dtype` than the one in the config: `--torch_dtype xxx`.\n\nIf you get an error linked to disk offload, you need to add the option `--disk-offload`\n\n## Results\n\nOn a setup with two Titan RTXs (24GB of RAM) and 32GB of RAM, we get the following benchmarks (T0pp does not run in float16, which is why it's not included).\n\n| Model | Model load time | Generation time | dtype | GPU 0 use | GPU 1 use | CPU use | Disk offload |\n|:-----:|:---------------:|:---------------:|:-----:|:---------:|:---------:|:-------:|:------------:|\n| GPT-J-6B | 8.7s | 0.05s per token | float16 | 11.7GB | 0GB | 0GB | no |\n| GPT-J-6B | 12.4s | 0.06s per token | float32 | 21.9GB | 1.5GB | 0GB | no |\n| GPT-Neo-X-20B | 30.9s | 0.08s per token | float16 | 21.5GB | 18GB | 0GB | no |\n| GPT-Neo-X-20B | 78.2s | 10.72s per token | float32 | 20.3GB | 22.7 GB | 24.4GB | yes |\n| T0pp (11B) | 29.4s | 0.05s per token | float32 | 21.1GB | 21.3GB | 0GB | no |\n| OPT-30B | 34.5s | 2.37s per token | float16 | 20.7GB | 22.3GB | 14.1GB | no |\n| OPT-30B | 112.3s | 33.9s per token | float32 | 20.2GB | 21.2GB | 23.5GB | yes |\n\nNote on the results:\n- using two GPUs instead of one does not slow down generation\n- using CPU offload slows down a bit (see OPT-30b)\n- using disk offload slows down a lot (need to implement prefetching)\n\nYou will also note that Accelerate does not use anymore GPU and CPU RAM than necessary:\n- peak GPU memory is exactly the size of the model put on a given GPU\n- peak CPU memory is either the size of the biggest checkpoint shard or the part of the model offloaded on CPU, whichever is bigger.\n" }, { "path": "benchmarks/big_model_inference/big_model_inference.py", "content": "# Copyright 2022 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport argparse\nimport time\n\nimport torch\nimport transformers\nfrom measures_util import end_measure, log_measures, start_measure\nfrom transformers import AutoConfig, AutoModelForCausalLM, AutoModelForSeq2SeqLM, AutoTokenizer\n\nfrom accelerate.utils import compute_module_sizes\n\n\nDEFAULT_MODELS = {\n \"gpt-j-6b\": {\"is_causal\": True, \"model\": \"sgugger/sharded-gpt-j-6B\", \"tokenizer\": \"EleutherAI/gpt-j-6B\"},\n \"gpt-neox\": {\"is_causal\": True, \"model\": \"EleutherAI/gpt-neox-20b\"},\n \"opt\": {\"is_causal\": True, \"model\": \"facebook/opt-30b\"},\n \"T0pp\": {\"is_causal\": False, \"model\": \"bigscience/T0pp\", \"model_revision\": \"sharded\"},\n}\n\nPROMPTS = [\n \"Hello, my name is\",\n \"Are unicorns real? Unicorns are\",\n \"For the first time in several years,\",\n \"My name is Julien and I am\",\n \"The goal of life is\",\n \"Whenever I'm sad, I like to\",\n]\n\n\ndef parse_args():\n parser = argparse.ArgumentParser(description=\"Run and time generations on a big model using Accelerate.\")\n parser.add_argument(\"model_name\", type=str, default=None, help=\"The name of the model to try.\")\n parser.add_argument(\n \"--tokenizer_name\", type=str, default=None, help=\"The name of the tokenizer (if different from the model.\"\n )\n parser.add_argument(\"--is_causal\", type=bool, default=None, help=\"Whether or not the model is causal.\")\n parser.add_argument(\n \"--model_revision\", type=str, default=None, help=\"The revision to use for the model checkpoint.\"\n )\n parser.add_argument(\"--torch_dtype\", type=str, default=None, help=\"The dtype for the model.\")\n parser.add_argument(\"--disk_offload\", action=\"store_true\")\n\n args = parser.parse_args()\n\n # Sanitize args\n if args.model_name in DEFAULT_MODELS:\n defaults = DEFAULT_MODELS[args.model_name]\n args.model_name = defaults[\"model\"]\n if args.tokenizer_name is None:\n args.tokenizer_name = defaults.get(\"tokenizer\", args.model_name)\n if args.is_causal is None:\n args.is_causal = defaults[\"is_causal\"]\n if args.model_revision is None:\n args.model_revision = defaults.get(\"model_revision\", \"main\")\n\n if args.is_causal is None:\n raise ValueError(\"Could not infer the default for `--is_causal`, pass either True or False for it.\")\n if args.tokenizer_name is None:\n args.tokenizer_name = args.model_name\n if args.model_revision is None:\n args.model_revision = \"main\"\n\n return args\n\n\ndef main():\n transformers.utils.logging.set_verbosity_error()\n args = parse_args()\n\n if args.torch_dtype is None:\n config = AutoConfig.from_pretrained(args.model_name)\n torch_dtype = getattr(config, \"torch_dtype\", torch.float32)\n else:\n torch_dtype = getattr(torch, args.torch_dtype)\n model_cls = AutoModelForCausalLM if args.is_causal else AutoModelForSeq2SeqLM\n kwargs = {\n \"torch_dtype\": torch_dtype,\n \"revision\": args.model_revision,\n }\n if args.disk_offload:\n kwargs[\"offload_folder\"] = \"tmp_offload\"\n kwargs[\"offload_state_dict\"] = True\n\n start_measures = start_measure()\n model = model_cls.from_pretrained(args.model_name, device_map=\"auto\", **kwargs)\n end_measures = end_measure(start_measures)\n log_measures(end_measures, \"Model loading\")\n\n module_sizes = compute_module_sizes(model)\n device_size = {v: 0 for v in model.hf_device_map.values()}\n for module, device in model.hf_device_map.items():\n device_size[device] += module_sizes[module]\n message = \"\\n\".join([f\"- {device}: {size // 2**20}MiB\" for device, size in device_size.items()])\n print(f\"\\nTheoretical use:\\n{message}\")\n\n tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_name)\n\n start_measures = start_measure()\n generation_times = []\n gen_tokens = []\n texts_outs = []\n for prompt in PROMPTS:\n inputs = tokenizer(prompt, return_tensors=\"pt\").to(0)\n tokens = inputs[\"input_ids\"][0].tolist()\n before_generate = time.time()\n outputs = model.generate(inputs[\"input_ids\"])\n after_generate = time.time()\n outputs = outputs[0].tolist()\n num_gen_tokens = len(outputs) if outputs[: len(tokens)] != tokens else len(outputs) - len(tokens)\n generation_time = after_generate - before_generate\n\n text_out = tokenizer.decode(outputs, skip_special_tokens=True)\n texts_outs.append(text_out)\n generation_times.append(generation_time)\n gen_tokens.append(num_gen_tokens)\n print(f\"Prompt: {prompt}\\nGeneration {text_out}\\nIn {generation_time:.2f}s for {num_gen_tokens} tokens\\n\")\n\n end_measures = end_measure(start_measures)\n log_measures(end_measures, \"Model generation\")\n\n generation_times_per_token = [gen / tok for gen, tok in zip(generation_times, gen_tokens)]\n avg_gen = sum(generation_times_per_token) / len(generation_times)\n print(f\"Average time of generation per token: {avg_gen:.2f}s\")\n print(f\"First generation (avg time per token): {generation_times_per_token[0]:.2f}s\")\n avg_gen = sum(generation_times_per_token[1:]) / (len(generation_times_per_token) - 1)\n print(f\"Average time of generation per token (excluding the first): {avg_gen:.2f}s\")\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "benchmarks/big_model_inference/measures_util.py", "content": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport gc\nimport threading\nimport time\n\nimport psutil\nimport torch\n\nfrom accelerate.test_utils.testing import get_backend\n\n\ntorch_device_type, _, _ = get_backend()\ntorch_accelerator_module = getattr(torch, torch_device_type, torch.cuda)\n\n\nclass PeakCPUMemory:\n def __init__(self):\n self.process = psutil.Process()\n self.peak_monitoring = False\n\n def peak_monitor(self):\n self.cpu_memory_peak = -1\n\n while True:\n self.cpu_memory_peak = max(self.process.memory_info().rss, self.cpu_memory_peak)\n\n # can't sleep or will not catch the peak right (this comment is here on purpose)\n if not self.peak_monitoring:\n break\n\n def start(self):\n self.peak_monitoring = True\n self.thread = threading.Thread(target=self.peak_monitor)\n self.thread.daemon = True\n self.thread.start()\n\n def stop(self):\n self.peak_monitoring = False\n self.thread.join()\n return self.cpu_memory_peak\n\n\ncpu_peak_tracker = PeakCPUMemory()\n\n\ndef start_measure():\n # Time\n measures = {\"time\": time.time()}\n\n gc.collect()\n torch_accelerator_module.empty_cache()\n\n # CPU mem\n measures[\"cpu\"] = psutil.Process().memory_info().rss\n cpu_peak_tracker.start()\n\n # GPU mem\n for i in range(torch_accelerator_module.device_count()):\n measures[str(i)] = torch_accelerator_module.memory_allocated(i)\n torch_accelerator_module.reset_peak_memory_stats()\n\n return measures\n\n\ndef end_measure(start_measures):\n # Time\n measures = {\"time\": time.time() - start_measures[\"time\"]}\n\n gc.collect()\n torch_accelerator_module.empty_cache()\n\n # CPU mem\n measures[\"cpu\"] = (psutil.Process().memory_info().rss - start_measures[\"cpu\"]) / 2**20\n measures[\"cpu-peak\"] = (cpu_peak_tracker.stop() - start_measures[\"cpu\"]) / 2**20\n\n # GPU mem\n for i in range(torch_accelerator_module.device_count()):\n measures[str(i)] = (torch_accelerator_module.memory_allocated(i) - start_measures[str(i)]) / 2**20\n measures[f\"{i}-peak\"] = (torch_accelerator_module.max_memory_allocated(i) - start_measures[str(i)]) / 2**20\n\n return measures\n\n\ndef log_measures(measures, description):\n print(f\"{description}:\")\n print(f\"- Time: {measures['time']:.2f}s\")\n for i in range(torch_accelerator_module.device_count()):\n print(f\"- {torch_device_type} {i} allocated: {measures[str(i)]:.2f}MiB\")\n peak = measures[f\"{i}-peak\"]\n print(f\"- {torch_device_type} {i} peak: {peak:.2f}MiB\")\n print(f\"- CPU RAM allocated: {measures['cpu']:.2f}MiB\")\n print(f\"- CPU RAM peak: {measures['cpu-peak']:.2f}MiB\")\n" }, { "path": "benchmarks/fp8/ms_amp/Dockerfile", "content": "FROM ghcr.io/azure/msamp\n\nRUN pip install transformers evaluate datasets\nRUN git clone https://github.com/huggingface/accelerate\n\nRUN cd accelerate && \\\n pip install -e . && \\\n cd benchmarks/fp8\n\nCMD [\"bash\"]\n\n\n" }, { "path": "benchmarks/fp8/ms_amp/ddp.py", "content": "# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nThis script tests to ensure that `accelerate` performs at the same level as raw `MS-AMP`.\n\nThis particular script verifies this for DDP training.\n\"\"\"\n\nimport evaluate\nimport msamp\nimport torch\nfrom fp8_utils import evaluate_model, get_training_utilities\nfrom torch.nn.parallel import DistributedDataParallel as DDP\n\nfrom accelerate import Accelerator\nfrom accelerate.state import AcceleratorState\nfrom accelerate.utils import FP8RecipeKwargs, get_grad_scaler, set_seed\n\n\nMODEL_NAME = \"bert-base-cased\"\nMETRIC = evaluate.load(\"glue\", \"mrpc\")\n\n\ndef train_baseline(opt_level=\"O2\"):\n set_seed(42)\n scaler = get_grad_scaler()\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = get_training_utilities(MODEL_NAME)\n accelerator = Accelerator()\n device = accelerator.device\n\n model, optimizer = msamp.initialize(model, optimizer, opt_level=opt_level)\n\n model.to(device)\n\n # Convert the model to DDP\n device_ids, output_device = [accelerator.local_process_index], accelerator.local_process_index\n model = DDP(model, device_ids=device_ids, output_device=output_device)\n\n base_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)\n model.train()\n\n for i, batch in enumerate(train_dataloader):\n with torch.autocast(device_type=\"cuda\", dtype=torch.bfloat16):\n outputs = model(**batch)\n loss = outputs.loss\n scaler.scale(loss).backward()\n optimizer.step()\n optimizer.zero_grad()\n lr_scheduler.step()\n\n trained_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)\n\n assert trained_model_results[\"accuracy\"] > base_model_results[\"accuracy\"], (\n f\"Accuracy should be higher for the trained model: {trained_model_results['accuracy']} > {base_model_results['accuracy']}\"\n )\n assert trained_model_results[\"f1\"] > base_model_results[\"f1\"], (\n f\"F1 score should be higher for the trained model: {trained_model_results['f1']} > {base_model_results['f1']}\"\n )\n\n return base_model_results, trained_model_results\n\n\ndef train_integration(opt_level=\"O2\"):\n kwargs_handlers = [FP8RecipeKwargs(backend=\"msamp\", opt_level=opt_level)]\n AcceleratorState()._reset_state(True)\n accelerator = Accelerator(mixed_precision=\"fp8\", kwargs_handlers=kwargs_handlers)\n set_seed(42)\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = get_training_utilities(\n MODEL_NAME, accelerator=accelerator\n )\n\n model, optimizer = accelerator.prepare(model, optimizer)\n base_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)\n model.train()\n for i, batch in enumerate(train_dataloader):\n with torch.autocast(device_type=\"cuda\", dtype=torch.bfloat16):\n outputs = model(**batch)\n loss = outputs.loss\n accelerator.backward(loss)\n optimizer.step()\n optimizer.zero_grad()\n lr_scheduler.step()\n\n trained_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)\n\n assert trained_model_results[\"accuracy\"] > base_model_results[\"accuracy\"], (\n f\"Accuracy should be higher for the trained model: {trained_model_results['accuracy']} > {base_model_results['accuracy']}\"\n )\n assert trained_model_results[\"f1\"] > base_model_results[\"f1\"], (\n f\"F1 score should be higher for the trained model: {trained_model_results['f1']} > {base_model_results['f1']}\"\n )\n\n return base_model_results, trained_model_results\n\n\nif __name__ == \"__main__\":\n for opt_level in [\"O1\", \"O2\"]:\n baseline_not_trained, baseline_trained = train_baseline(opt_level)\n accelerator_not_trained, accelerator_trained = train_integration(opt_level)\n assert baseline_not_trained[\"accuracy\"] == accelerator_not_trained[\"accuracy\"], (\n f\"Accuracy not the same for untrained baseline and accelerator using opt_level={opt_level}: {baseline_not_trained['accuracy']} == {accelerator_not_trained['accuracy']}\"\n )\n assert baseline_not_trained[\"f1\"] == accelerator_not_trained[\"f1\"], (\n f\"F1 not the same for untrained baseline and accelerator using opt_level={opt_level}: {baseline_not_trained['f1']} == {accelerator_not_trained['f1']}\"\n )\n assert baseline_trained[\"accuracy\"] == accelerator_trained[\"accuracy\"], (\n f\"Accuracy not the same for trained baseline and accelerator using opt_level={opt_level}: {baseline_trained['accuracy']} == {accelerator_trained['accuracy']}\"\n )\n assert baseline_trained[\"f1\"] == accelerator_trained[\"f1\"], (\n f\"F1 not the same for trained baseline and accelerator using opt_level={opt_level}: {baseline_trained['f1']} == {accelerator_trained['f1']}\"\n )\n" }, { "path": "benchmarks/fp8/ms_amp/distrib_deepspeed.py", "content": "# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nThis script tests to ensure that `accelerate` performs at the same level as raw `MS-AMP`.\n\nThis particular script verifies this for DeepSpeed training.\n\nNOTE: MS-AMP does *not* support ZeRO-3.\n\"\"\"\n\n# import msamp.deepspeed as msamp_deepspeed\nimport evaluate\nimport torch\nfrom fp8_utils import evaluate_model, get_training_utilities\nfrom msamp import deepspeed as msamp_deepspeed\n\nfrom accelerate import Accelerator, DeepSpeedPlugin\nfrom accelerate.state import AcceleratorState\nfrom accelerate.utils import set_seed\n\n\nMODEL_NAME = \"bert-base-cased\"\nMETRIC = evaluate.load(\"glue\", \"mrpc\")\n\n\ndef train_baseline(zero_stage: int = 1, opt_level: str = \"O1\"):\n set_seed(42)\n accelerator = Accelerator()\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = get_training_utilities(\n MODEL_NAME, accelerator=accelerator\n )\n\n import numpy as np\n\n config = {\n \"train_batch_size\": 32,\n \"train_micro_batch_size_per_gpu\": 16,\n \"gradient_accumulation_steps\": 1,\n \"zero_optimization\": {\n \"stage\": zero_stage,\n \"offload_optimizer\": {\"device\": \"none\", \"nvme_path\": None},\n \"offload_param\": {\"device\": \"none\", \"nvme_path\": None},\n },\n \"gradient_clipping\": 1.0,\n \"steps_per_print\": np.inf,\n \"bf16\": {\"enabled\": True},\n \"fp16\": {\"enabled\": False},\n \"zero_allow_untested_optimizer\": True,\n \"msamp\": {\n \"enabled\": True,\n \"opt_level\": opt_level,\n },\n }\n (\n model,\n optimizer,\n _,\n _,\n ) = msamp_deepspeed.initialize(\n model=model,\n optimizer=optimizer,\n config_params=config,\n )\n\n base_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)\n model.train()\n\n for _ in range(2):\n for batch in train_dataloader:\n outputs = model(**batch)\n loss = outputs.loss\n model.backward(loss)\n model.step()\n for _ in range(accelerator.num_processes):\n lr_scheduler.step()\n\n trained_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)\n model.destroy()\n torch.cuda.empty_cache()\n AcceleratorState()._reset_state(True)\n assert trained_model_results[\"accuracy\"] > base_model_results[\"accuracy\"], (\n f\"Accuracy should be higher for the trained model: {trained_model_results['accuracy']} > {base_model_results['accuracy']}\"\n )\n assert trained_model_results[\"f1\"] > base_model_results[\"f1\"], (\n f\"F1 score should be higher for the trained model: {trained_model_results['f1']} > {base_model_results['f1']}\"\n )\n\n return base_model_results, trained_model_results\n\n\ndef train_integration(zero_stage: int = 1, opt_level: str = \"O1\"):\n set_seed(42)\n deepspeed_plugin = DeepSpeedPlugin(\n zero_stage=zero_stage,\n enable_msamp=True,\n msamp_opt_level=opt_level,\n )\n accelerator = Accelerator(mixed_precision=\"fp8\", deepspeed_plugin=deepspeed_plugin)\n accelerator.state.deepspeed_plugin.deepspeed_config[\"train_micro_batch_size_per_gpu\"] = 16\n\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = get_training_utilities(\n MODEL_NAME, accelerator=accelerator\n )\n\n model, optimizer, lr_scheduler = accelerator.prepare(model, optimizer, lr_scheduler)\n base_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)\n model.train()\n for _ in range(2):\n for batch in train_dataloader:\n outputs = model(**batch)\n loss = outputs.loss\n accelerator.backward(loss)\n optimizer.step()\n lr_scheduler.step()\n optimizer.zero_grad()\n\n trained_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)\n model.destroy()\n torch.cuda.empty_cache()\n assert trained_model_results[\"accuracy\"] > base_model_results[\"accuracy\"], (\n f\"Accuracy should be higher for the trained model: {trained_model_results['accuracy']} > {base_model_results['accuracy']}\"\n )\n assert trained_model_results[\"f1\"] > base_model_results[\"f1\"], (\n f\"F1 score should be higher for the trained model: {trained_model_results['f1']} > {base_model_results['f1']}\"\n )\n\n AcceleratorState()._reset_state(True)\n return base_model_results, trained_model_results\n\n\nif __name__ == \"__main__\":\n for zero_stage in [1, 2]:\n for opt_level in [\"O1\", \"O2\", \"O3\"]:\n baseline_not_trained, baseline_trained = train_baseline(zero_stage, opt_level)\n accelerator_not_trained, accelerator_trained = train_integration(zero_stage, opt_level)\n assert baseline_not_trained[\"accuracy\"] == accelerator_not_trained[\"accuracy\"], (\n f\"ZERO stage {zero_stage}, opt_level={opt_level}:\\nAccuracy should be the same for the baseline and accelerator: {baseline_not_trained['accuracy']} == {accelerator_not_trained['accuracy']}\"\n )\n assert baseline_not_trained[\"f1\"] == accelerator_not_trained[\"f1\"], (\n f\"ZERO stage {zero_stage}, opt_level={opt_level}:\\nF1 score should be the same for the baseline and accelerator: {baseline_not_trained['f1']} == {accelerator_not_trained['f1']}\"\n )\n assert baseline_trained[\"accuracy\"] == accelerator_trained[\"accuracy\"], (\n f\"ZERO stage {zero_stage}, opt_level={opt_level}:\\nAccuracy should be the same for the baseline and accelerator: {baseline_trained['accuracy']} == {accelerator_trained['accuracy']}\"\n )\n assert baseline_trained[\"f1\"] == accelerator_trained[\"f1\"], (\n f\"ZERO stage {zero_stage}, opt_level={opt_level}:\\nF1 score should be the same for the baseline and accelerator: {baseline_trained['f1']} == {accelerator_trained['f1']}\"\n )\n\n torch.distributed.destroy_process_group()\n" }, { "path": "benchmarks/fp8/ms_amp/fp8_utils.py", "content": "# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport torch\n\n\ndef get_dataloaders(model_name: str, batch_size: int = 16):\n from datasets import load_dataset\n from torch.utils.data import DataLoader\n from transformers import AutoTokenizer\n\n tokenizer = AutoTokenizer.from_pretrained(model_name)\n datasets = load_dataset(\"glue\", \"mrpc\")\n\n def tokenize_function(examples):\n # max_length=None => use the model max length (it's actually the default)\n outputs = tokenizer(examples[\"sentence1\"], examples[\"sentence2\"], truncation=True, max_length=None)\n return outputs\n\n # Apply the method we just defined to all the examples in all the splits of the dataset\n # starting with the main process first:\n tokenized_datasets = datasets.map(\n tokenize_function,\n batched=True,\n remove_columns=[\"idx\", \"sentence1\", \"sentence2\"],\n )\n\n # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the\n # transformers library\n tokenized_datasets = tokenized_datasets.rename_column(\"label\", \"labels\")\n\n def collate_fn(examples):\n return tokenizer.pad(\n examples,\n padding=\"longest\",\n pad_to_multiple_of=16, # Specific for FP8\n return_tensors=\"pt\",\n )\n\n # Instantiate dataloaders.\n train_dataloader = DataLoader(\n tokenized_datasets[\"train\"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size, drop_last=True\n )\n eval_dataloader = DataLoader(\n tokenized_datasets[\"validation\"],\n shuffle=False,\n collate_fn=collate_fn,\n batch_size=16,\n drop_last=True,\n )\n\n return train_dataloader, eval_dataloader\n\n\ndef get_training_utilities(model_name: str, batch_size: int = 16, accelerator=None):\n \"\"\"\n Returns a tuple of:\n - Model\n - Optimizer\n - Train dataloader (prepared)\n - Eval dataloader (prepared)\n - LR Scheduler\n Suitable for training on the MRPC dataset\n \"\"\"\n from torch.optim import AdamW\n from transformers import AutoModelForSequenceClassification, get_linear_schedule_with_warmup\n\n from accelerate import Accelerator\n\n if accelerator is None:\n accelerator = Accelerator()\n model = AutoModelForSequenceClassification.from_pretrained(model_name)\n train_dataloader, eval_dataloader = get_dataloaders(model_name, batch_size)\n optimizer = AdamW(model.parameters(), lr=0.0001)\n lr_scheduler = get_linear_schedule_with_warmup(\n optimizer=optimizer,\n num_warmup_steps=100,\n num_training_steps=len(train_dataloader) * 2,\n )\n train_dataloader, eval_dataloader = accelerator.prepare(train_dataloader, eval_dataloader)\n return model, optimizer, train_dataloader, eval_dataloader, lr_scheduler\n\n\ndef get_named_parameters(model):\n \"\"\"\n Same thing as `Accelerator.get_named_parameters` Returns a list of the named parameters of the model (extracted\n from parallel)\n \"\"\"\n from accelerate.utils import extract_model_from_parallel\n\n model = extract_model_from_parallel(model)\n return {n: p for n, p in model.named_parameters()}\n\n\ndef evaluate_model(model, dataloader, metric, accelerator=None):\n \"Turns model to .eval(), runs dataloader, calculates metric, then turns eval back on\"\n model.eval()\n for step, batch in enumerate(dataloader):\n with torch.no_grad():\n # W/ MS-AMP, we need to cast while evaluating\n with torch.autocast(device_type=\"cuda\", dtype=torch.bfloat16):\n outputs = model(**batch)\n predictions = outputs.logits.argmax(dim=-1)\n references = batch[\"labels\"]\n if accelerator is not None and accelerator.num_processes > 1:\n predictions, references = accelerator.gather_for_metrics((predictions, references))\n metric.add_batch(predictions=predictions, references=references)\n return metric.compute()\n" }, { "path": "benchmarks/fp8/ms_amp/non_distributed.py", "content": "# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nThis script tests to ensure that `accelerate` performs at the same level as raw `MS-AMP`.\n\nThis particular script verifies this for single GPU training.\n\"\"\"\n\nimport evaluate\nimport msamp\nimport torch\nfrom fp8_utils import evaluate_model, get_training_utilities\n\nfrom accelerate import Accelerator\nfrom accelerate.state import AcceleratorState\nfrom accelerate.utils import FP8RecipeKwargs, get_grad_scaler, set_seed\n\n\nMODEL_NAME = \"bert-base-cased\"\nMETRIC = evaluate.load(\"glue\", \"mrpc\")\n\n\ndef train_baseline(opt_level=\"O2\"):\n set_seed(42)\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = get_training_utilities(MODEL_NAME)\n\n model, optimizer = msamp.initialize(model, optimizer, opt_level=opt_level)\n model.to(\"cuda\")\n\n base_model_results = evaluate_model(model, eval_dataloader, METRIC)\n model.train()\n scaler = get_grad_scaler()\n\n for batch in train_dataloader:\n batch = batch.to(\"cuda\")\n with torch.autocast(device_type=\"cuda\", dtype=torch.bfloat16):\n outputs = model(**batch)\n loss = outputs.loss\n loss = scaler.scale(loss)\n loss.backward()\n optimizer.step()\n optimizer.zero_grad()\n lr_scheduler.step()\n\n trained_model_results = evaluate_model(model, eval_dataloader, METRIC)\n\n assert trained_model_results[\"accuracy\"] > base_model_results[\"accuracy\"], (\n f\"Accuracy should be higher for the trained model: {trained_model_results['accuracy']} > {base_model_results['accuracy']}\"\n )\n assert trained_model_results[\"f1\"] > base_model_results[\"f1\"], (\n f\"F1 score should be higher for the trained model: {trained_model_results['f1']} > {base_model_results['f1']}\"\n )\n\n return base_model_results, trained_model_results\n\n\ndef train_integration(opt_level=\"O2\"):\n kwargs_handlers = [FP8RecipeKwargs(backend=\"msamp\", opt_level=opt_level)]\n AcceleratorState()._reset_state(True)\n accelerator = Accelerator(mixed_precision=\"fp8\", kwargs_handlers=kwargs_handlers)\n set_seed(42)\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = get_training_utilities(\n MODEL_NAME, accelerator=accelerator\n )\n\n model, optimizer, lr_scheduler = accelerator.prepare(model, optimizer, lr_scheduler)\n base_model_results = evaluate_model(model, eval_dataloader, METRIC)\n model.train()\n\n for batch in train_dataloader:\n outputs = model(**batch)\n loss = outputs.loss\n accelerator.backward(loss)\n optimizer.step()\n optimizer.zero_grad()\n lr_scheduler.step()\n\n trained_model_results = evaluate_model(model, eval_dataloader, METRIC)\n\n assert trained_model_results[\"accuracy\"] > base_model_results[\"accuracy\"], (\n f\"Accuracy should be higher for the trained model: {trained_model_results['accuracy']} > {base_model_results['accuracy']}\"\n )\n assert trained_model_results[\"f1\"] > base_model_results[\"f1\"], (\n f\"F1 score should be higher for the trained model: {trained_model_results['f1']} > {base_model_results['f1']}\"\n )\n\n return base_model_results, trained_model_results\n\n\nif __name__ == \"__main__\":\n for opt_level in [\"O1\", \"O2\"]:\n baseline_not_trained, baseline_trained = train_baseline(opt_level)\n accelerator_not_trained, accelerator_trained = train_integration(opt_level)\n\n assert baseline_not_trained[\"accuracy\"] == accelerator_not_trained[\"accuracy\"], (\n f\"Accuracy should be the same for the baseline and accelerator: {baseline_not_trained['accuracy']} == {accelerator_not_trained['accuracy']}\"\n )\n assert baseline_not_trained[\"f1\"] == accelerator_not_trained[\"f1\"], (\n f\"F1 score should be the same for the baseline and accelerator: {baseline_not_trained['f1']} == {accelerator_not_trained['f1']}\"\n )\n assert baseline_trained[\"accuracy\"] == accelerator_trained[\"accuracy\"], (\n f\"Accuracy should be the same for the baseline and accelerator: {baseline_trained['accuracy']} == {accelerator_trained['accuracy']}\"\n )\n assert baseline_trained[\"f1\"] == accelerator_trained[\"f1\"], (\n f\"F1 score should be the same for the baseline and accelerator: {baseline_trained['f1']} == {accelerator_trained['f1']}\"\n )\n" }, { "path": "benchmarks/fp8/torchao/Dockerfile", "content": "FROM nvcr.io/nvidia/pytorch:24.07-py3\n\nRUN pip install transformers evaluate datasets\nRUN git clone https://github.com/huggingface/accelerate.git\n\nRUN cd accelerate && \\\n pip install -e . && \\\n cd benchmarks/fp8\n\nRUN /bin/bash\n\n\n" }, { "path": "benchmarks/fp8/torchao/README.md", "content": "# FP8 Benchmarks\n\nComparing and running [torchao](https://github.com/pytorch/ao/tree/main/torchao/float8) FP8 with accelerate\n\n## Overview\n\nThis repo provides scripts which compare native `torchao` model training against `accelerate`'s own integration. Each modeling type is segmented out via a script, supporting the following:\n\n* Single GPU training (`non_distributed.py`)\n* Multi-GPU training via DistributedDataParallelism (`ddp.py`)\n* Fully Sharded Data Parallelism (`fsdp.py`)\n* DeepSpeed ZeRO 1-3 (`deepspeed.py`)\n\nTo run them, it's recommended to use a docker image (see the attached `Dockerfile`) and not install `torchao` manually.\n\n## Running:\n\nThere are official Docker images located at `huggingface/accelerate:gpu-fp8-torchao-nightly` which can be used.\n\nYou can run all scripts using the core `accelerate launch` command without any `accelerate config` being needed.\n\nFor single GPU, run it via `python`:\n\n```bash\npython non_distributed.py\n```\n\nFor the rest, run it via `accelerate launch`:\n\n```bash\naccelerate launch ddp.py # or distrib_deepspeed.py, ddp.py\n```" }, { "path": "benchmarks/fp8/torchao/ddp.py", "content": "# Copyright 2025 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nThis script tests to ensure that `accelerate` performs at the same level as raw `torchao`.\n\nThis particular script verifies this for DDP training.\n\"\"\"\n\nfrom functools import partial\n\nimport evaluate\nimport torch\nfrom fp8_utils import get_training_utilities\nfrom torch.nn.parallel import DistributedDataParallel as DDP\nfrom torchao.float8 import convert_to_float8_training\n\nfrom accelerate import Accelerator\nfrom accelerate.state import AcceleratorState\nfrom accelerate.utils import AORecipeKwargs, set_seed\n\n\nMODEL_NAME = \"bert-base-cased\"\nMETRIC = evaluate.load(\"glue\", \"mrpc\")\n\n\ndef evaluate_model(model, dataloader, metric, accelerator=None):\n \"Turns model to .eval(), runs dataloader, calculates metric, then turns eval back on\"\n model.eval()\n for step, batch in enumerate(dataloader):\n with torch.no_grad():\n outputs = model(**batch)\n predictions = outputs.logits.argmax(dim=-1)\n references = batch[\"labels\"]\n if accelerator is not None and accelerator.num_processes > 1:\n predictions, references = accelerator.gather_for_metrics((predictions, references))\n metric.add_batch(predictions=predictions, references=references)\n return metric.compute()\n\n\ndef filter_linear_layers(module, fqn, first_layer_name=None, last_layer_name=None):\n if isinstance(module, torch.nn.Linear):\n if module.in_features % 16 != 0 or module.out_features % 16 != 0:\n return False\n # For stability reasons, we skip the first and last linear layers\n # Otherwise can lead to the model not training or converging properly\n if fqn in (first_layer_name, last_layer_name):\n return False\n return True\n\n\ndef train_baseline():\n set_seed(42)\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = get_training_utilities(MODEL_NAME)\n first_linear = None\n last_linear = None\n for name, module in model.named_modules():\n if isinstance(module, torch.nn.Linear):\n if first_linear is None:\n first_linear = name\n last_linear = name\n func = partial(filter_linear_layers, first_layer_name=first_linear, last_layer_name=last_linear)\n accelerator = Accelerator()\n device = accelerator.device\n model.to(device)\n\n convert_to_float8_training(model, module_filter_fn=func)\n\n # Convert the model to DDP\n device_ids, output_device = [accelerator.local_process_index], accelerator.local_process_index\n model = DDP(model, device_ids=device_ids, output_device=output_device)\n\n base_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)\n model.train()\n\n for batch in train_dataloader:\n with torch.autocast(device_type=device.type, dtype=torch.bfloat16):\n batch = batch.to(device)\n outputs = model(**batch)\n loss = outputs.loss\n loss.backward()\n optimizer.step()\n optimizer.zero_grad()\n lr_scheduler.step()\n\n trained_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)\n\n assert trained_model_results[\"accuracy\"] > base_model_results[\"accuracy\"], (\n f\"Accuracy should be higher for the trained model: {trained_model_results['accuracy']} > {base_model_results['accuracy']}\"\n )\n assert trained_model_results[\"f1\"] > base_model_results[\"f1\"], (\n f\"F1 score should be higher for the trained model: {trained_model_results['f1']} > {base_model_results['f1']}\"\n )\n\n return base_model_results, trained_model_results\n\n\ndef train_integration():\n AcceleratorState()._reset_state(True)\n accelerator = Accelerator(mixed_precision=\"fp8\", kwargs_handlers=[AORecipeKwargs()])\n set_seed(42)\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = get_training_utilities(\n MODEL_NAME, accelerator=accelerator\n )\n\n model, optimizer = accelerator.prepare(model, optimizer)\n base_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)\n model.train()\n\n for batch in train_dataloader:\n outputs = model(**batch)\n loss = outputs.loss\n accelerator.backward(loss)\n optimizer.step()\n optimizer.zero_grad()\n lr_scheduler.step()\n\n trained_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)\n\n assert trained_model_results[\"accuracy\"] > base_model_results[\"accuracy\"], (\n f\"Accuracy should be higher for the trained model: {trained_model_results['accuracy']} > {base_model_results['accuracy']}\"\n )\n assert trained_model_results[\"f1\"] > base_model_results[\"f1\"], (\n f\"F1 score should be higher for the trained model: {trained_model_results['f1']} > {base_model_results['f1']}\"\n )\n\n return base_model_results, trained_model_results\n\n\nif __name__ == \"__main__\":\n baseline_not_trained, baseline_trained = train_baseline()\n accelerator_not_trained, accelerator_trained = train_integration()\n\n assert baseline_not_trained[\"accuracy\"] == accelerator_not_trained[\"accuracy\"], (\n f\"Accuracy should be the same for the baseline and accelerator: {baseline_not_trained['accuracy']} == {accelerator_not_trained['accuracy']}\"\n )\n assert baseline_not_trained[\"f1\"] == accelerator_not_trained[\"f1\"], (\n f\"F1 score should be the same for the baseline and accelerator: {baseline_not_trained['f1']} == {accelerator_not_trained['f1']}\"\n )\n assert baseline_trained[\"accuracy\"] == accelerator_trained[\"accuracy\"], (\n f\"Accuracy should be the same for the baseline and accelerator: {baseline_trained['accuracy']} == {accelerator_trained['accuracy']}\"\n )\n assert baseline_trained[\"f1\"] == accelerator_trained[\"f1\"], (\n f\"F1 score should be the same for the baseline and accelerator: {baseline_trained['f1']} == {accelerator_trained['f1']}\"\n )\n\n torch.distributed.destroy_process_group()\n" }, { "path": "benchmarks/fp8/torchao/distrib_deepspeed.py", "content": "# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nThis script tests to ensure that `accelerate` performs at the same level as raw `torchao`.\n\nThis particular script verifies this for deepspeed training.\n\"\"\"\n\nfrom functools import partial\nfrom unittest.mock import patch\n\nimport deepspeed\nimport evaluate\nimport torch\nfrom fp8_utils import evaluate_model, get_training_utilities\nfrom torchao.float8 import convert_to_float8_training\nfrom transformers.integrations import HfDeepSpeedConfig\n\nfrom accelerate import Accelerator, DeepSpeedPlugin\nfrom accelerate.state import AcceleratorState\nfrom accelerate.utils import AORecipeKwargs, set_seed\n\n\nMODEL_NAME = \"bert-base-cased\"\nMETRIC = evaluate.load(\"glue\", \"mrpc\")\n\n\ndef filter_linear_layers(module, fqn, first_layer_name=None, last_layer_name=None):\n if isinstance(module, torch.nn.Linear):\n if module.in_features % 16 != 0 or module.out_features % 16 != 0:\n return False\n # For stability reasons, we skip the first and last linear layers\n # Otherwise can lead to the model not training or converging properly\n if fqn in (first_layer_name, last_layer_name):\n return False\n return True\n\n\ndef train_baseline(zero_stage: int = 1):\n set_seed(42)\n # This forces transformers to think Zero-3 Init should be used\n with patch(\"transformers.integrations.deepspeed.is_deepspeed_zero3_enabled\") as mock:\n mock.return_value = zero_stage == 3\n\n config = HfDeepSpeedConfig(\n {\n \"train_micro_batch_size_per_gpu\": 16,\n \"gradient_accumulation_steps\": 1,\n \"zero_optimization\": {\"stage\": zero_stage},\n }\n )\n plugin = DeepSpeedPlugin(hf_ds_config=config)\n accelerator = Accelerator(deepspeed_plugin=plugin)\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = get_training_utilities(\n MODEL_NAME, accelerator=accelerator\n )\n first_linear = None\n last_linear = None\n for name, module in model.named_modules():\n if isinstance(module, torch.nn.Linear):\n if first_linear is None:\n first_linear = name\n last_linear = name\n func = partial(filter_linear_layers, first_layer_name=first_linear, last_layer_name=last_linear)\n\n convert_to_float8_training(model, module_filter_fn=func)\n\n import numpy as np\n\n config = {\n \"train_batch_size\": 32,\n \"train_micro_batch_size_per_gpu\": 16,\n \"gradient_accumulation_steps\": 1,\n \"zero_optimization\": {\n \"stage\": zero_stage,\n \"offload_optimizer\": {\"device\": \"none\", \"nvme_path\": None},\n \"offload_param\": {\"device\": \"none\", \"nvme_path\": None},\n \"stage3_gather_16bit_weights_on_model_save\": False,\n },\n \"gradient_clipping\": 1.0,\n \"steps_per_print\": np.inf,\n \"bf16\": {\"enabled\": True},\n \"fp16\": {\"enabled\": False},\n \"zero_allow_untested_optimizer\": True,\n }\n\n (\n model,\n optimizer,\n _,\n lr_scheduler,\n ) = deepspeed.initialize(\n model=model,\n optimizer=optimizer,\n lr_scheduler=lr_scheduler,\n config_params=config,\n )\n\n base_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)\n model.train()\n\n model_outputs = []\n data = []\n\n for batch in train_dataloader:\n outputs = model(**batch)\n data.append(batch.to(\"cpu\"))\n model_outputs.append(outputs.logits.to(\"cpu\"))\n loss = outputs.loss\n model.backward(loss)\n model.step()\n for _ in range(accelerator.num_processes):\n lr_scheduler.step()\n\n trained_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)\n model.destroy()\n assert trained_model_results[\"accuracy\"] > base_model_results[\"accuracy\"], (\n f\"Accuracy should be higher for the trained model: {trained_model_results['accuracy']} > {base_model_results['accuracy']}\"\n )\n assert trained_model_results[\"f1\"] > base_model_results[\"f1\"], (\n f\"F1 score should be higher for the trained model: {trained_model_results['f1']} > {base_model_results['f1']}\"\n )\n\n del config\n return base_model_results, trained_model_results, model_outputs, data\n\n\ndef train_integration(zero_stage: int = 1):\n set_seed(42)\n AcceleratorState()._reset_state(True)\n config = HfDeepSpeedConfig(\n {\n \"train_micro_batch_size_per_gpu\": 16,\n \"gradient_accumulation_steps\": 1,\n \"zero_optimization\": {\"stage\": zero_stage},\n }\n )\n deepspeed_plugin = DeepSpeedPlugin(\n hf_ds_config=config,\n )\n # This forces transformers to think Zero-3 Init should be used\n with patch(\"transformers.integrations.deepspeed.is_deepspeed_zero3_enabled\") as mock:\n mock.return_value = zero_stage == 3\n accelerator = Accelerator(\n mixed_precision=\"fp8\", kwargs_handlers=[AORecipeKwargs()], deepspeed_plugin=deepspeed_plugin\n )\n\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = get_training_utilities(\n MODEL_NAME, accelerator=accelerator\n )\n\n model, optimizer, lr_scheduler, train_dataloader, eval_dataloader = accelerator.prepare(\n model, optimizer, lr_scheduler, train_dataloader, eval_dataloader\n )\n base_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)\n model.train()\n model_outputs = []\n data = []\n for batch in train_dataloader:\n outputs = model(**batch)\n data.append(batch.to(\"cpu\"))\n model_outputs.append(outputs.logits.to(\"cpu\"))\n loss = outputs.loss\n accelerator.backward(loss)\n optimizer.step()\n lr_scheduler.step()\n optimizer.zero_grad()\n\n trained_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)\n model.destroy()\n assert trained_model_results[\"accuracy\"] > base_model_results[\"accuracy\"], (\n f\"Accuracy should be higher for the trained model: {trained_model_results['accuracy']} > {base_model_results['accuracy']}\"\n )\n assert trained_model_results[\"f1\"] > base_model_results[\"f1\"], (\n f\"F1 score should be higher for the trained model: {trained_model_results['f1']} > {base_model_results['f1']}\"\n )\n\n del config\n return base_model_results, trained_model_results, model_outputs, data\n\n\nif __name__ == \"__main__\":\n for zero_stage in [1, 2, 3]:\n baseline_not_trained, baseline_trained, baseline_outputs, baseline_data = train_baseline(zero_stage)\n accelerator_not_trained, accelerator_trained, accelerator_outputs, accelerator_data = train_integration(\n zero_stage\n )\n assert baseline_not_trained[\"accuracy\"] == accelerator_not_trained[\"accuracy\"], (\n f\"ZERO stage {zero_stage}: Accuracy should be the same for the baseline and accelerator: {baseline_not_trained['accuracy']} == {accelerator_not_trained['accuracy']}\"\n )\n assert baseline_not_trained[\"f1\"] == accelerator_not_trained[\"f1\"], (\n f\"ZERO stage {zero_stage}: F1 score should be the same for the baseline and accelerator: {baseline_not_trained['f1']} == {accelerator_not_trained['f1']}\"\n )\n assert baseline_trained[\"accuracy\"] == accelerator_trained[\"accuracy\"], (\n f\"ZERO stage {zero_stage}: Accuracy should be the same for the baseline and accelerator: {baseline_trained['accuracy']} == {accelerator_trained['accuracy']}\"\n )\n assert baseline_trained[\"f1\"] == accelerator_trained[\"f1\"], (\n f\"ZERO stage {zero_stage}: F1 score should be the same for the baseline and accelerator: {baseline_trained['f1']} == {accelerator_trained['f1']}\"\n )\n AcceleratorState()._reset_state(True)\n torch.distributed.destroy_process_group()\n" }, { "path": "benchmarks/fp8/torchao/fp8_utils.py", "content": "# Copyright 2025 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport torch\n\n\ndef get_dataloaders(model_name: str, batch_size: int = 16):\n from datasets import load_dataset\n from torch.utils.data import DataLoader\n from transformers import AutoTokenizer\n\n tokenizer = AutoTokenizer.from_pretrained(model_name)\n datasets = load_dataset(\"glue\", \"mrpc\")\n\n def tokenize_function(examples):\n # max_length=None => use the model max length (it's actually the default)\n outputs = tokenizer(examples[\"sentence1\"], examples[\"sentence2\"], truncation=True, max_length=None)\n return outputs\n\n # Apply the method we just defined to all the examples in all the splits of the dataset\n # starting with the main process first:\n tokenized_datasets = datasets.map(\n tokenize_function,\n batched=True,\n remove_columns=[\"idx\", \"sentence1\", \"sentence2\"],\n )\n\n # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the\n # transformers library\n tokenized_datasets = tokenized_datasets.rename_column(\"label\", \"labels\")\n\n def collate_fn(examples):\n return tokenizer.pad(\n examples,\n padding=\"longest\",\n pad_to_multiple_of=16, # Specific for FP8\n return_tensors=\"pt\",\n )\n\n # Instantiate dataloaders.\n train_dataloader = DataLoader(\n tokenized_datasets[\"train\"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size, drop_last=True\n )\n eval_dataloader = DataLoader(\n tokenized_datasets[\"validation\"],\n shuffle=False,\n collate_fn=collate_fn,\n batch_size=16,\n drop_last=True,\n )\n\n return train_dataloader, eval_dataloader\n\n\ndef get_training_utilities(model_name: str, batch_size: int = 16, accelerator=None, prepare=True):\n \"\"\"\n Returns a tuple of:\n - Model\n - Optimizer\n - Train dataloader (prepared)\n - Eval dataloader (prepared)\n - LR Scheduler\n Suitable for training on the MRPC dataset\n \"\"\"\n from torch.optim import AdamW\n from transformers import AutoModelForSequenceClassification, get_linear_schedule_with_warmup\n\n from accelerate import Accelerator\n\n if accelerator is None:\n accelerator = Accelerator()\n model = AutoModelForSequenceClassification.from_pretrained(model_name)\n train_dataloader, eval_dataloader = get_dataloaders(model_name, batch_size)\n optimizer = AdamW(model.parameters(), lr=0.0001)\n lr_scheduler = get_linear_schedule_with_warmup(\n optimizer=optimizer,\n num_warmup_steps=100,\n num_training_steps=len(train_dataloader) * 2,\n )\n train_dataloader, eval_dataloader = accelerator.prepare(train_dataloader, eval_dataloader)\n return model, optimizer, train_dataloader, eval_dataloader, lr_scheduler\n\n\ndef get_named_parameters(model):\n \"\"\"\n Same thing as `Accelerator.get_named_parameters` Returns a list of the named parameters of the model (extracted\n from parallel)\n \"\"\"\n from accelerate.utils import extract_model_from_parallel\n\n model = extract_model_from_parallel(model)\n return {n: p for n, p in model.named_parameters()}\n\n\ndef evaluate_model(model, dataloader, metric, accelerator=None):\n \"Turns model to .eval(), runs dataloader, calculates metric, then turns eval back on\"\n model.eval()\n for step, batch in enumerate(dataloader):\n with torch.no_grad():\n outputs = model(**batch)\n predictions = outputs.logits.argmax(dim=-1)\n references = batch[\"labels\"]\n if accelerator is not None and accelerator.num_processes > 1:\n predictions, references = accelerator.gather_for_metrics((predictions, references))\n metric.add_batch(predictions=predictions, references=references)\n return metric.compute()\n" }, { "path": "benchmarks/fp8/torchao/fsdp.py", "content": "# Copyright 2025 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nThis script tests to ensure that `accelerate` performs at the same level as raw `torchao`.\n\nThis particular script verifies this for FSDP training.\n\"\"\"\n\nfrom functools import partial\n\nimport evaluate\nimport torch\nfrom fp8_utils import get_training_utilities\nfrom torch.distributed.fsdp import FullyShardedDataParallel as FSDP\nfrom torch.distributed.fsdp import MixedPrecision\nfrom torch.distributed.fsdp.wrap import transformer_auto_wrap_policy\nfrom torchao.float8 import convert_to_float8_training\nfrom transformers.models.bert import BertLayer\n\nfrom accelerate import Accelerator\nfrom accelerate import FullyShardedDataParallelPlugin as FSDPPlugin\nfrom accelerate.state import AcceleratorState\nfrom accelerate.utils import AORecipeKwargs, set_seed\n\n\nMODEL_NAME = \"bert-base-cased\"\nMETRIC = evaluate.load(\"glue\", \"mrpc\")\n\nFSDP_WRAP_POLICY = partial(transformer_auto_wrap_policy, transformer_layer_cls={BertLayer})\n\n\ndef filter_linear_layers(module, fqn, first_layer_name=None, last_layer_name=None):\n if isinstance(module, torch.nn.Linear):\n if module.in_features % 16 != 0 or module.out_features % 16 != 0:\n return False\n # For stability reasons, we skip the first and last linear layers\n # Otherwise can lead to the model not training or converging properly\n if fqn in (first_layer_name, last_layer_name):\n return False\n return True\n\n\ndef evaluate_model(model, dataloader, metric, accelerator=None):\n \"Turns model to .eval(), runs dataloader, calculates metric, then turns eval back on\"\n model.eval()\n for step, batch in enumerate(dataloader):\n with torch.no_grad():\n outputs = model(**batch)\n predictions = outputs.logits.argmax(dim=-1)\n references = batch[\"labels\"]\n if accelerator is not None and accelerator.num_processes > 1:\n predictions, references = accelerator.gather_for_metrics((predictions, references))\n metric.add_batch(predictions=predictions, references=references)\n return metric.compute()\n\n\ndef train_baseline():\n set_seed(42)\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = get_training_utilities(MODEL_NAME)\n first_linear = None\n last_linear = None\n for name, module in model.named_modules():\n if isinstance(module, torch.nn.Linear):\n if first_linear is None:\n first_linear = name\n last_linear = name\n func = partial(filter_linear_layers, first_layer_name=first_linear, last_layer_name=last_linear)\n accelerator = Accelerator()\n device = accelerator.device\n model.to(device)\n\n convert_to_float8_training(model, module_filter_fn=func)\n\n # Convert the model to FSDP\n model = FSDP(\n model,\n use_orig_params=True,\n mixed_precision=MixedPrecision(param_dtype=torch.bfloat16, reduce_dtype=torch.float32),\n auto_wrap_policy=FSDP_WRAP_POLICY,\n )\n\n base_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)\n model.train()\n\n for batch in train_dataloader:\n with torch.autocast(device_type=device.type, dtype=torch.bfloat16):\n batch = batch.to(device)\n outputs = model(**batch)\n loss = outputs.loss\n loss.backward()\n optimizer.step()\n optimizer.zero_grad()\n lr_scheduler.step()\n\n trained_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)\n\n assert trained_model_results[\"accuracy\"] > base_model_results[\"accuracy\"], (\n f\"Accuracy should be higher for the trained model: {trained_model_results['accuracy']} > {base_model_results['accuracy']}\"\n )\n assert trained_model_results[\"f1\"] > base_model_results[\"f1\"], (\n f\"F1 score should be higher for the trained model: {trained_model_results['f1']} > {base_model_results['f1']}\"\n )\n\n return base_model_results, trained_model_results\n\n\ndef train_integration():\n AcceleratorState()._reset_state(True)\n fsdp_plugin = FSDPPlugin(\n auto_wrap_policy=FSDP_WRAP_POLICY,\n use_orig_params=True,\n mixed_precision_policy=MixedPrecision(param_dtype=torch.bfloat16, reduce_dtype=torch.float32),\n )\n accelerator = Accelerator(mixed_precision=\"fp8\", fsdp_plugin=fsdp_plugin, kwargs_handlers=[AORecipeKwargs()])\n set_seed(42)\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = get_training_utilities(\n MODEL_NAME, accelerator=accelerator\n )\n\n model, optimizer = accelerator.prepare(model, optimizer)\n base_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)\n model.train()\n\n for batch in train_dataloader:\n outputs = model(**batch)\n loss = outputs.loss\n accelerator.backward(loss)\n optimizer.step()\n optimizer.zero_grad()\n lr_scheduler.step()\n\n trained_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)\n\n assert trained_model_results[\"accuracy\"] > base_model_results[\"accuracy\"], (\n f\"Accuracy should be higher for the trained model: {trained_model_results['accuracy']} > {base_model_results['accuracy']}\"\n )\n assert trained_model_results[\"f1\"] > base_model_results[\"f1\"], (\n f\"F1 score should be higher for the trained model: {trained_model_results['f1']} > {base_model_results['f1']}\"\n )\n\n return base_model_results, trained_model_results\n\n\nif __name__ == \"__main__\":\n baseline_not_trained, baseline_trained = train_baseline()\n accelerator_not_trained, accelerator_trained = train_integration()\n\n assert baseline_not_trained[\"accuracy\"] == accelerator_not_trained[\"accuracy\"], (\n f\"Accuracy should be the same for the baseline and accelerator: {baseline_not_trained['accuracy']} == {accelerator_not_trained['accuracy']}\"\n )\n assert baseline_not_trained[\"f1\"] == accelerator_not_trained[\"f1\"], (\n f\"F1 score should be the same for the baseline and accelerator: {baseline_not_trained['f1']} == {accelerator_not_trained['f1']}\"\n )\n assert baseline_trained[\"accuracy\"] == accelerator_trained[\"accuracy\"], (\n f\"Accuracy should be the same for the baseline and accelerator: {baseline_trained['accuracy']} == {accelerator_trained['accuracy']}\"\n )\n assert baseline_trained[\"f1\"] == accelerator_trained[\"f1\"], (\n f\"F1 score should be the same for the baseline and accelerator: {baseline_trained['f1']} == {accelerator_trained['f1']}\"\n )\n\n torch.distributed.destroy_process_group()\n" }, { "path": "benchmarks/fp8/torchao/non_distributed.py", "content": "# Copyright 2025 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nThis script tests to ensure that `accelerate` performs at the same level as raw `torchao`.\n\nThis particular script verifies this for single GPU training.\n\"\"\"\n\nfrom functools import partial\n\nimport evaluate\nimport torch\nfrom fp8_utils import get_training_utilities\nfrom torchao.float8 import convert_to_float8_training\n\nfrom accelerate import Accelerator\nfrom accelerate.state import AcceleratorState\nfrom accelerate.utils import AORecipeKwargs, set_seed\n\n\nMODEL_NAME = \"bert-base-cased\"\nMETRIC = evaluate.load(\"glue\", \"mrpc\")\n\n\ndef evaluate_model(model, dataloader, metric, accelerator=None):\n \"Turns model to .eval(), runs dataloader, calculates metric, then turns eval back on\"\n model.eval()\n for step, batch in enumerate(dataloader):\n with torch.no_grad():\n outputs = model(**batch)\n predictions = outputs.logits.argmax(dim=-1)\n references = batch[\"labels\"]\n if accelerator is not None and accelerator.num_processes > 1:\n predictions, references = accelerator.gather_for_metrics((predictions, references))\n metric.add_batch(predictions=predictions, references=references)\n return metric.compute()\n\n\ndef filter_linear_layers(module, fqn, first_layer_name=None, last_layer_name=None):\n if isinstance(module, torch.nn.Linear):\n if module.in_features % 16 != 0 or module.out_features % 16 != 0:\n return False\n # For stability reasons, we skip the first and last linear layers\n # Otherwise can lead to the model not training or converging properly\n if fqn in (first_layer_name, last_layer_name):\n return False\n return True\n\n\ndef train_baseline():\n set_seed(42)\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = get_training_utilities(MODEL_NAME)\n first_linear = None\n last_linear = None\n for name, module in model.named_modules():\n if isinstance(module, torch.nn.Linear):\n if first_linear is None:\n first_linear = name\n last_linear = name\n\n func = partial(filter_linear_layers, first_layer_name=first_linear, last_layer_name=last_linear)\n accelerator = Accelerator()\n device = accelerator.device\n model.to(device)\n convert_to_float8_training(model, module_filter_fn=func)\n base_model_results = evaluate_model(model, eval_dataloader, METRIC)\n model.train()\n\n for batch in train_dataloader:\n with torch.autocast(device_type=device.type, dtype=torch.bfloat16):\n outputs = model(**batch)\n loss = outputs.loss\n loss.backward()\n optimizer.step()\n optimizer.zero_grad()\n lr_scheduler.step()\n\n trained_model_results = evaluate_model(model, eval_dataloader, METRIC)\n\n assert trained_model_results[\"accuracy\"] > base_model_results[\"accuracy\"], (\n f\"Accuracy should be higher for the trained model: {trained_model_results['accuracy']} > {base_model_results['accuracy']}\"\n )\n assert trained_model_results[\"f1\"] > base_model_results[\"f1\"], (\n f\"F1 score should be higher for the trained model: {trained_model_results['f1']} > {base_model_results['f1']}\"\n )\n\n return base_model_results, trained_model_results\n\n\ndef train_integration():\n set_seed(42)\n accelerator = Accelerator(mixed_precision=\"fp8\", kwargs_handlers=[AORecipeKwargs()])\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = get_training_utilities(\n MODEL_NAME, accelerator=accelerator\n )\n model = accelerator.prepare(model)\n base_model_results = evaluate_model(model, eval_dataloader, METRIC)\n model.train()\n\n for batch in train_dataloader:\n outputs = model(**batch)\n loss = outputs.loss\n loss.backward()\n optimizer.step()\n optimizer.zero_grad()\n lr_scheduler.step()\n\n trained_model_results = evaluate_model(model, eval_dataloader, METRIC)\n\n assert trained_model_results[\"accuracy\"] > base_model_results[\"accuracy\"], (\n f\"Accuracy should be higher for the trained model: {trained_model_results['accuracy']} > {base_model_results['accuracy']}\"\n )\n assert trained_model_results[\"f1\"] > base_model_results[\"f1\"], (\n f\"F1 score should be higher for the trained model: {trained_model_results['f1']} > {base_model_results['f1']}\"\n )\n\n return base_model_results, trained_model_results\n\n\nif __name__ == \"__main__\":\n baseline_not_trained, baseline_trained = train_baseline()\n AcceleratorState._reset_state(True)\n accelerator_not_trained, accelerator_trained = train_integration()\n assert baseline_not_trained[\"accuracy\"] == accelerator_not_trained[\"accuracy\"], (\n f\"Accuracy should be the same for the baseline and accelerator: {baseline_not_trained['accuracy']} == {accelerator_not_trained['accuracy']}\"\n )\n assert baseline_not_trained[\"f1\"] == accelerator_not_trained[\"f1\"], (\n f\"F1 score should be the same for the baseline and accelerator: {baseline_not_trained['f1']} == {accelerator_not_trained['f1']}\"\n )\n assert baseline_trained[\"accuracy\"] == accelerator_trained[\"accuracy\"], (\n f\"Accuracy should be the same for the baseline and accelerator: {baseline_trained['accuracy']} == {accelerator_trained['accuracy']}\"\n )\n assert baseline_trained[\"f1\"] == accelerator_trained[\"f1\"], (\n f\"F1 score should be the same for the baseline and accelerator: {baseline_trained['f1']} == {accelerator_trained['f1']}\"\n )\n" }, { "path": "benchmarks/fp8/transformer_engine/Dockerfile", "content": "ARG BASE_YEAR=25\nARG BASE_MONTH=03\n\nFROM nvcr.io/nvidia/pytorch:${BASE_YEAR}.${BASE_MONTH}-py3\n\nRUN pip install transformers evaluate datasets\nRUN git clone https://github.com/huggingface/accelerate.git\n\nRUN cd accelerate && \\\n pip install -e .[deepspeed] && \\\n cd benchmarks/fp8\n\nRUN /bin/bash\n\n\n" }, { "path": "benchmarks/fp8/transformer_engine/README.md", "content": "# FP8 Benchmarks\n\nComparing and running [TransformerEngine](https://github.com/NVIDIA/TransformerEngine) FP8 with accelerate\n\n## Overview\n\nThis repo provides scripts which compare native TransformerEngine model training against `accelerate`'s own integration. Each modeling type is segmented out via a script, supporting the following:\n\n* Single GPU training (`non_distributed.py`)\n* Multi-GPU training via DistributedDataParallelism (`ddp.py`)\n* Fully Sharded Data Parallelism (`fsdp.py`)\n* DeepSpeed ZeRO 1-3 (`deepspeed.py`)\n\nTo run them, it's recommended to use a docker image (see the attached `Dockerfile`) and not install `TransformerEngine` manually.\n\n## Running:\n\nThere are official Docker images located at `huggingface/accelerate:gpu-fp8-transformerengine-nightly` which can be used.\n\nYou can run all scripts using the core `accelerate launch` command without any `accelerate config` being needed.\n\nFor single GPU, run it via `python`:\n\n```bash\npython non_distributed.py\n```\n\nFor the rest, run it via `accelerate launch`:\n\n```bash\naccelerate launch ddp.py # or distrib_deepspeed.py, ddp.py\n```" }, { "path": "benchmarks/fp8/transformer_engine/ddp.py", "content": "# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nThis script tests to ensure that `accelerate` performs at the same level as raw `TransformersEngine`.\n\nThis particular script verifies this for DDP training.\n\"\"\"\n\nimport evaluate\nimport torch\nimport transformer_engine.common.recipe as te_recipe\nimport transformer_engine.pytorch as te\nfrom fp8_utils import evaluate_model, get_named_parameters, get_training_utilities\nfrom torch.nn.parallel import DistributedDataParallel as DDP\nfrom transformer_engine.common.recipe import DelayedScaling\n\nfrom accelerate import Accelerator\nfrom accelerate.state import AcceleratorState\nfrom accelerate.utils import FP8RecipeKwargs, set_seed\nfrom accelerate.utils.transformer_engine import convert_model\n\n\nMODEL_NAME = \"bert-base-cased\"\nMETRIC = evaluate.load(\"glue\", \"mrpc\")\n\n\ndef train_baseline():\n set_seed(42)\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = get_training_utilities(MODEL_NAME)\n accelerator = Accelerator()\n device = accelerator.device\n model.to(device)\n\n # Convert the model to TE\n old_named_params = get_named_parameters(model)\n\n with torch.no_grad():\n convert_model(model)\n\n FP8_RECIPE_KWARGS = {\"fp8_format\": te_recipe.Format.HYBRID, \"amax_history_len\": 32, \"amax_compute_algo\": \"max\"}\n fp8_recipe = DelayedScaling(**FP8_RECIPE_KWARGS)\n\n new_named_params = get_named_parameters(model)\n\n # Convert the model to DDP\n device_ids, output_device = [accelerator.local_process_index], accelerator.local_process_index\n model = DDP(model, device_ids=device_ids, output_device=output_device)\n\n mapping = {p: new_named_params[n] for n, p in old_named_params.items()}\n for param_group in optimizer.param_groups:\n param_group[\"params\"] = [mapping[p] for p in param_group[\"params\"]]\n\n base_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)\n model.train()\n\n for _ in range(2):\n for batch in train_dataloader:\n with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):\n with torch.autocast(device_type=\"cuda\", dtype=torch.bfloat16):\n batch = batch.to(device)\n outputs = model(**batch)\n loss = outputs.loss\n loss.backward()\n optimizer.step()\n optimizer.zero_grad()\n lr_scheduler.step()\n\n trained_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)\n\n assert trained_model_results[\"accuracy\"] > base_model_results[\"accuracy\"], (\n f\"Accuracy should be higher for the trained model: {trained_model_results['accuracy']} > {base_model_results['accuracy']}\"\n )\n assert trained_model_results[\"f1\"] > base_model_results[\"f1\"], (\n f\"F1 score should be higher for the trained model: {trained_model_results['f1']} > {base_model_results['f1']}\"\n )\n\n return base_model_results, trained_model_results\n\n\ndef train_integration():\n FP8_RECIPE_KWARGS = {\"fp8_format\": \"HYBRID\", \"amax_history_len\": 32, \"amax_compute_algo\": \"max\"}\n kwargs_handlers = [FP8RecipeKwargs(backend=\"TE\", **FP8_RECIPE_KWARGS)]\n AcceleratorState()._reset_state(True)\n accelerator = Accelerator(mixed_precision=\"fp8\", kwargs_handlers=kwargs_handlers)\n set_seed(42)\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = get_training_utilities(\n MODEL_NAME, accelerator=accelerator\n )\n\n model, optimizer = accelerator.prepare(model, optimizer)\n base_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)\n model.train()\n\n for _ in range(2):\n for batch in train_dataloader:\n outputs = model(**batch)\n loss = outputs.loss\n accelerator.backward(loss)\n optimizer.step()\n optimizer.zero_grad()\n lr_scheduler.step()\n\n trained_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)\n\n assert trained_model_results[\"accuracy\"] > base_model_results[\"accuracy\"], (\n f\"Accuracy should be higher for the trained model: {trained_model_results['accuracy']} > {base_model_results['accuracy']}\"\n )\n assert trained_model_results[\"f1\"] > base_model_results[\"f1\"], (\n f\"F1 score should be higher for the trained model: {trained_model_results['f1']} > {base_model_results['f1']}\"\n )\n\n return base_model_results, trained_model_results\n\n\nif __name__ == \"__main__\":\n baseline_not_trained, baseline_trained = train_baseline()\n accelerator_not_trained, accelerator_trained = train_integration()\n\n assert baseline_not_trained[\"accuracy\"] == accelerator_not_trained[\"accuracy\"], (\n f\"Accuracy should be the same for the baseline and accelerator: {baseline_not_trained['accuracy']} == {accelerator_not_trained['accuracy']}\"\n )\n assert baseline_not_trained[\"f1\"] == accelerator_not_trained[\"f1\"], (\n f\"F1 score should be the same for the baseline and accelerator: {baseline_not_trained['f1']} == {accelerator_not_trained['f1']}\"\n )\n assert baseline_trained[\"accuracy\"] == accelerator_trained[\"accuracy\"], (\n f\"Accuracy should be the same for the baseline and accelerator: {baseline_trained['accuracy']} == {accelerator_trained['accuracy']}\"\n )\n assert baseline_trained[\"f1\"] == accelerator_trained[\"f1\"], (\n f\"F1 score should be the same for the baseline and accelerator: {baseline_trained['f1']} == {accelerator_trained['f1']}\"\n )\n\n torch.distributed.destroy_process_group()\n" }, { "path": "benchmarks/fp8/transformer_engine/distrib_deepspeed.py", "content": "# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nThis script tests to ensure that `accelerate` performs at the same level as raw `TransformersEngine`.\n\nThis particular script verifies this for DDP training.\n\"\"\"\n\nfrom unittest.mock import patch\n\nimport deepspeed\nimport evaluate\nimport torch\nimport transformer_engine.common.recipe as te_recipe\nimport transformer_engine.pytorch as te\nfrom fp8_utils import evaluate_model, get_named_parameters, get_training_utilities\nfrom transformer_engine.common.recipe import DelayedScaling\n\nfrom accelerate import Accelerator, DeepSpeedPlugin\nfrom accelerate.state import AcceleratorState\nfrom accelerate.utils import FP8RecipeKwargs, set_seed\nfrom accelerate.utils.transformer_engine import convert_model\n\n\nMODEL_NAME = \"bert-base-cased\"\nMETRIC = evaluate.load(\"glue\", \"mrpc\")\n\n\ndef train_baseline(zero_stage: int = 1):\n # This forces transformers to think Zero-3 Init should be used\n with patch(\"transformers.integrations.deepspeed.is_deepspeed_zero3_enabled\") as mock:\n mock.return_value = zero_stage == 3\n set_seed(42)\n\n accelerator = Accelerator()\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = get_training_utilities(\n MODEL_NAME, accelerator=accelerator\n )\n\n # Convert the model to TE\n old_named_params = get_named_parameters(model)\n\n with torch.no_grad():\n convert_model(model)\n new_named_params = get_named_parameters(model)\n\n mapping = {p: new_named_params[n] for n, p in old_named_params.items()}\n for param_group in optimizer.param_groups:\n param_group[\"params\"] = [mapping[p] for p in param_group[\"params\"]]\n\n FP8_RECIPE_KWARGS = {\"fp8_format\": te_recipe.Format.HYBRID, \"amax_history_len\": 32, \"amax_compute_algo\": \"max\"}\n fp8_recipe = DelayedScaling(**FP8_RECIPE_KWARGS)\n\n import numpy as np\n\n config = {\n \"train_batch_size\": 16,\n \"train_micro_batch_size_per_gpu\": 16,\n \"gradient_accumulation_steps\": 1,\n \"zero_optimization\": {\n \"stage\": zero_stage,\n \"offload_optimizer\": {\"device\": \"none\", \"nvme_path\": None},\n \"offload_param\": {\"device\": \"none\", \"nvme_path\": None},\n \"stage3_gather_16bit_weights_on_model_save\": False,\n },\n \"gradient_clipping\": 1.0,\n \"steps_per_print\": np.inf,\n \"bf16\": {\"enabled\": True},\n \"fp16\": {\"enabled\": False},\n \"zero_allow_untested_optimizer\": True,\n }\n\n (\n model,\n optimizer,\n _,\n _,\n ) = deepspeed.initialize(\n model=model,\n optimizer=optimizer,\n config_params=config,\n )\n\n base_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)\n model.train()\n\n model_outputs = []\n data = []\n\n for _ in range(2):\n for batch in train_dataloader:\n with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):\n outputs = model(**batch)\n data.append(batch.to(\"cpu\"))\n model_outputs.append(outputs.logits.to(\"cpu\"))\n loss = outputs.loss\n model.backward(loss)\n model.step()\n for _ in range(accelerator.num_processes):\n lr_scheduler.step()\n\n trained_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)\n model.destroy()\n assert trained_model_results[\"accuracy\"] > base_model_results[\"accuracy\"], (\n f\"Accuracy should be higher for the trained model: {trained_model_results['accuracy']} > {base_model_results['accuracy']}\"\n )\n assert trained_model_results[\"f1\"] > base_model_results[\"f1\"], (\n f\"F1 score should be higher for the trained model: {trained_model_results['f1']} > {base_model_results['f1']}\"\n )\n\n return base_model_results, trained_model_results, model_outputs, data\n\n\ndef train_integration(zero_stage: int = 1):\n set_seed(42)\n FP8_RECIPE_KWARGS = {\"fp8_format\": \"HYBRID\", \"amax_history_len\": 32, \"amax_compute_algo\": \"max\"}\n kwargs_handlers = [FP8RecipeKwargs(backend=\"TE\", **FP8_RECIPE_KWARGS)]\n AcceleratorState()._reset_state(True)\n deepspeed_plugin = DeepSpeedPlugin(\n zero_stage=zero_stage,\n zero3_init_flag=zero_stage == 3,\n )\n accelerator = Accelerator(\n mixed_precision=\"fp8\", kwargs_handlers=kwargs_handlers, deepspeed_plugin=deepspeed_plugin\n )\n accelerator.state.deepspeed_plugin.deepspeed_config[\"train_micro_batch_size_per_gpu\"] = 16\n\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = get_training_utilities(\n MODEL_NAME, accelerator=accelerator\n )\n\n model, optimizer, lr_scheduler = accelerator.prepare(model, optimizer, lr_scheduler)\n base_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)\n model.train()\n model_outputs = []\n data = []\n for _ in range(2):\n for batch in train_dataloader:\n outputs = model(**batch)\n data.append(batch.to(\"cpu\"))\n model_outputs.append(outputs.logits.to(\"cpu\"))\n loss = outputs.loss\n accelerator.backward(loss)\n optimizer.step()\n lr_scheduler.step()\n optimizer.zero_grad()\n\n trained_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)\n model.destroy()\n assert trained_model_results[\"accuracy\"] > base_model_results[\"accuracy\"], (\n f\"Accuracy should be higher for the trained model: {trained_model_results['accuracy']} > {base_model_results['accuracy']}\"\n )\n assert trained_model_results[\"f1\"] > base_model_results[\"f1\"], (\n f\"F1 score should be higher for the trained model: {trained_model_results['f1']} > {base_model_results['f1']}\"\n )\n\n return base_model_results, trained_model_results, model_outputs, data\n\n\nif __name__ == \"__main__\":\n for zero_stage in [1, 2, 3]:\n baseline_not_trained, baseline_trained, baseline_outputs, baseline_data = train_baseline(zero_stage)\n accelerator_not_trained, accelerator_trained, accelerator_outputs, accelerator_data = train_integration(\n zero_stage\n )\n assert baseline_not_trained[\"accuracy\"] == accelerator_not_trained[\"accuracy\"], (\n f\"ZERO stage {zero_stage}: Accuracy should be the same for the baseline and accelerator: {baseline_not_trained['accuracy']} == {accelerator_not_trained['accuracy']}\"\n )\n assert baseline_not_trained[\"f1\"] == accelerator_not_trained[\"f1\"], (\n f\"ZERO stage {zero_stage}: F1 score should be the same for the baseline and accelerator: {baseline_not_trained['f1']} == {accelerator_not_trained['f1']}\"\n )\n assert baseline_trained[\"accuracy\"] == accelerator_trained[\"accuracy\"], (\n f\"ZERO stage {zero_stage}: Accuracy should be the same for the baseline and accelerator: {baseline_trained['accuracy']} == {accelerator_trained['accuracy']}\"\n )\n assert baseline_trained[\"f1\"] == accelerator_trained[\"f1\"], (\n f\"ZERO stage {zero_stage}: F1 score should be the same for the baseline and accelerator: {baseline_trained['f1']} == {accelerator_trained['f1']}\"\n )\n\n torch.distributed.destroy_process_group()\n" }, { "path": "benchmarks/fp8/transformer_engine/fp8_utils.py", "content": "# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport torch\n\n\ndef get_dataloaders(model_name: str, batch_size: int = 16):\n from datasets import load_dataset\n from torch.utils.data import DataLoader\n from transformers import AutoTokenizer\n\n tokenizer = AutoTokenizer.from_pretrained(model_name)\n datasets = load_dataset(\"glue\", \"mrpc\")\n\n def tokenize_function(examples):\n # max_length=None => use the model max length (it's actually the default)\n outputs = tokenizer(examples[\"sentence1\"], examples[\"sentence2\"], truncation=True, max_length=None)\n return outputs\n\n # Apply the method we just defined to all the examples in all the splits of the dataset\n # starting with the main process first:\n tokenized_datasets = datasets.map(\n tokenize_function,\n batched=True,\n remove_columns=[\"idx\", \"sentence1\", \"sentence2\"],\n )\n\n # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the\n # transformers library\n tokenized_datasets = tokenized_datasets.rename_column(\"label\", \"labels\")\n\n def collate_fn(examples):\n return tokenizer.pad(\n examples,\n padding=\"longest\",\n pad_to_multiple_of=16, # Specific for FP8\n return_tensors=\"pt\",\n )\n\n # Instantiate dataloaders.\n train_dataloader = DataLoader(\n tokenized_datasets[\"train\"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size, drop_last=True\n )\n eval_dataloader = DataLoader(\n tokenized_datasets[\"validation\"],\n shuffle=False,\n collate_fn=collate_fn,\n batch_size=16,\n drop_last=True,\n )\n\n return train_dataloader, eval_dataloader\n\n\ndef get_training_utilities(model_name: str, batch_size: int = 16, accelerator=None):\n \"\"\"\n Returns a tuple of:\n - Model\n - Optimizer\n - Train dataloader (prepared)\n - Eval dataloader (prepared)\n - LR Scheduler\n Suitable for training on the MRPC dataset\n \"\"\"\n from torch.optim import AdamW\n from transformers import AutoModelForSequenceClassification, get_linear_schedule_with_warmup\n\n from accelerate import Accelerator\n\n if accelerator is None:\n accelerator = Accelerator()\n model = AutoModelForSequenceClassification.from_pretrained(model_name)\n train_dataloader, eval_dataloader = get_dataloaders(model_name, batch_size)\n optimizer = AdamW(model.parameters(), lr=0.0001)\n lr_scheduler = get_linear_schedule_with_warmup(\n optimizer=optimizer,\n num_warmup_steps=100,\n num_training_steps=len(train_dataloader) * 2,\n )\n train_dataloader, eval_dataloader = accelerator.prepare(train_dataloader, eval_dataloader)\n return model, optimizer, train_dataloader, eval_dataloader, lr_scheduler\n\n\ndef get_named_parameters(model):\n \"\"\"\n Same thing as `Accelerator.get_named_parameters` Returns a list of the named parameters of the model (extracted\n from parallel)\n \"\"\"\n from accelerate.utils import extract_model_from_parallel\n\n model = extract_model_from_parallel(model)\n return {n: p for n, p in model.named_parameters()}\n\n\ndef evaluate_model(model, dataloader, metric, accelerator=None):\n \"Turns model to .eval(), runs dataloader, calculates metric, then turns eval back on\"\n model.eval()\n for step, batch in enumerate(dataloader):\n with torch.no_grad():\n outputs = model(**batch)\n predictions = outputs.logits.argmax(dim=-1)\n references = batch[\"labels\"]\n if accelerator is not None and accelerator.num_processes > 1:\n predictions, references = accelerator.gather_for_metrics((predictions, references))\n metric.add_batch(predictions=predictions, references=references)\n return metric.compute()\n" }, { "path": "benchmarks/fp8/transformer_engine/fsdp.py", "content": "# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nThis script tests to ensure that `accelerate` performs at the same level as raw `TransformersEngine`.\n\nThis particular script verifies this for FSDP training.\n\"\"\"\n\nfrom functools import partial\n\nimport evaluate\nimport torch\nimport transformer_engine.common.recipe as te_recipe\nimport transformer_engine.pytorch as te\nfrom fp8_utils import evaluate_model, get_named_parameters, get_training_utilities\nfrom torch.distributed.fsdp import FullyShardedDataParallel as FSDP\nfrom torch.distributed.fsdp import MixedPrecision\nfrom torch.distributed.fsdp.wrap import transformer_auto_wrap_policy\nfrom transformer_engine.common.recipe import DelayedScaling\nfrom transformers.models.bert import BertLayer\n\nfrom accelerate import Accelerator\nfrom accelerate import FullyShardedDataParallelPlugin as FSDPPlugin\nfrom accelerate.state import AcceleratorState\nfrom accelerate.utils import FP8RecipeKwargs, set_seed\nfrom accelerate.utils.transformer_engine import convert_model\n\n\nMODEL_NAME = \"bert-base-cased\"\nMETRIC = evaluate.load(\"glue\", \"mrpc\")\n\nFSDP_WRAP_POLICY = partial(transformer_auto_wrap_policy, transformer_layer_cls={BertLayer})\n\n\ndef train_baseline():\n set_seed(42)\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = get_training_utilities(MODEL_NAME)\n accelerator = Accelerator()\n device = accelerator.device\n model.to(device)\n\n # Convert the model to TE\n old_named_params = get_named_parameters(model)\n\n with torch.no_grad():\n convert_model(model)\n\n FP8_RECIPE_KWARGS = {\"fp8_format\": te_recipe.Format.HYBRID, \"amax_history_len\": 32, \"amax_compute_algo\": \"max\"}\n fp8_recipe = DelayedScaling(**FP8_RECIPE_KWARGS)\n\n new_named_params = get_named_parameters(model)\n\n # Convert the model to FSDP\n model = FSDP(\n model,\n use_orig_params=True,\n mixed_precision=MixedPrecision(param_dtype=torch.bfloat16, reduce_dtype=torch.float32),\n auto_wrap_policy=FSDP_WRAP_POLICY,\n )\n\n mapping = {p: new_named_params[n] for n, p in old_named_params.items()}\n for param_group in optimizer.param_groups:\n param_group[\"params\"] = [mapping[p] for p in param_group[\"params\"]]\n\n base_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)\n model.train()\n\n for _ in range(2):\n for batch in train_dataloader:\n with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):\n with torch.autocast(device_type=\"cuda\", dtype=torch.bfloat16):\n batch = batch.to(device)\n outputs = model(**batch)\n loss = outputs.loss\n loss.backward()\n optimizer.step()\n optimizer.zero_grad()\n lr_scheduler.step()\n\n trained_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)\n\n assert trained_model_results[\"accuracy\"] > base_model_results[\"accuracy\"], (\n f\"Accuracy should be higher for the trained model: {trained_model_results['accuracy']} > {base_model_results['accuracy']}\"\n )\n assert trained_model_results[\"f1\"] > base_model_results[\"f1\"], (\n f\"F1 score should be higher for the trained model: {trained_model_results['f1']} > {base_model_results['f1']}\"\n )\n\n return base_model_results, trained_model_results\n\n\ndef train_integration():\n FP8_RECIPE_KWARGS = {\"fp8_format\": \"HYBRID\", \"amax_history_len\": 32, \"amax_compute_algo\": \"max\"}\n kwargs_handlers = [FP8RecipeKwargs(backend=\"TE\", **FP8_RECIPE_KWARGS)]\n AcceleratorState()._reset_state(True)\n fsdp_plugin = FSDPPlugin(\n auto_wrap_policy=FSDP_WRAP_POLICY,\n use_orig_params=True,\n mixed_precision_policy=MixedPrecision(param_dtype=torch.bfloat16, reduce_dtype=torch.float32),\n )\n accelerator = Accelerator(mixed_precision=\"fp8\", fsdp_plugin=fsdp_plugin, kwargs_handlers=kwargs_handlers)\n set_seed(42)\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = get_training_utilities(\n MODEL_NAME, accelerator=accelerator\n )\n\n model, optimizer = accelerator.prepare(model, optimizer)\n base_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)\n model.train()\n\n for _ in range(2):\n for batch in train_dataloader:\n outputs = model(**batch)\n loss = outputs.loss\n accelerator.backward(loss)\n optimizer.step()\n optimizer.zero_grad()\n lr_scheduler.step()\n\n trained_model_results = evaluate_model(model, eval_dataloader, METRIC, accelerator=accelerator)\n\n assert trained_model_results[\"accuracy\"] > base_model_results[\"accuracy\"], (\n f\"Accuracy should be higher for the trained model: {trained_model_results['accuracy']} > {base_model_results['accuracy']}\"\n )\n assert trained_model_results[\"f1\"] > base_model_results[\"f1\"], (\n f\"F1 score should be higher for the trained model: {trained_model_results['f1']} > {base_model_results['f1']}\"\n )\n\n return base_model_results, trained_model_results\n\n\nif __name__ == \"__main__\":\n baseline_not_trained, baseline_trained = train_baseline()\n accelerator_not_trained, accelerator_trained = train_integration()\n\n assert baseline_not_trained[\"accuracy\"] == accelerator_not_trained[\"accuracy\"], (\n f\"Accuracy should be the same for the baseline and accelerator: {baseline_not_trained['accuracy']} == {accelerator_not_trained['accuracy']}\"\n )\n assert baseline_not_trained[\"f1\"] == accelerator_not_trained[\"f1\"], (\n f\"F1 score should be the same for the baseline and accelerator: {baseline_not_trained['f1']} == {accelerator_not_trained['f1']}\"\n )\n assert baseline_trained[\"accuracy\"] == accelerator_trained[\"accuracy\"], (\n f\"Accuracy should be the same for the baseline and accelerator: {baseline_trained['accuracy']} == {accelerator_trained['accuracy']}\"\n )\n assert baseline_trained[\"f1\"] == accelerator_trained[\"f1\"], (\n f\"F1 score should be the same for the baseline and accelerator: {baseline_trained['f1']} == {accelerator_trained['f1']}\"\n )\n\n torch.distributed.destroy_process_group()\n" }, { "path": "benchmarks/fp8/transformer_engine/non_distributed.py", "content": "# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nThis script tests to ensure that `accelerate` performs at the same level as raw `TransformersEngine`.\n\nThis particular script verifies this for single GPU training.\n\"\"\"\n\nimport evaluate\nimport torch\nimport transformer_engine.common.recipe as te_recipe\nimport transformer_engine.pytorch as te\nfrom fp8_utils import evaluate_model, get_named_parameters, get_training_utilities\nfrom transformer_engine.common.recipe import DelayedScaling\n\nfrom accelerate import Accelerator\nfrom accelerate.state import AcceleratorState\nfrom accelerate.utils import FP8RecipeKwargs, set_seed\nfrom accelerate.utils.transformer_engine import convert_model\n\n\nMODEL_NAME = \"bert-base-cased\"\nMETRIC = evaluate.load(\"glue\", \"mrpc\")\n\n\ndef train_baseline():\n set_seed(42)\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = get_training_utilities(MODEL_NAME)\n\n # Convert the model to TE\n old_named_params = get_named_parameters(model)\n\n with torch.no_grad():\n convert_model(model)\n\n new_named_params = get_named_parameters(model)\n mapping = {p: new_named_params[n] for n, p in old_named_params.items()}\n for param_group in optimizer.param_groups:\n param_group[\"params\"] = [mapping[p] for p in param_group[\"params\"]]\n\n FP8_RECIPE_KWARGS = {\"fp8_format\": te_recipe.Format.HYBRID, \"amax_history_len\": 32, \"amax_compute_algo\": \"max\"}\n fp8_recipe = DelayedScaling(**FP8_RECIPE_KWARGS)\n\n model.to(\"cuda\")\n base_model_results = evaluate_model(model, eval_dataloader, METRIC)\n model.train()\n\n for batch in train_dataloader:\n with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):\n with torch.autocast(device_type=\"cuda\", dtype=torch.bfloat16):\n batch = batch.to(\"cuda\")\n outputs = model(**batch)\n loss = outputs.loss\n loss.backward()\n optimizer.step()\n optimizer.zero_grad()\n lr_scheduler.step()\n\n trained_model_results = evaluate_model(model, eval_dataloader, METRIC)\n\n assert trained_model_results[\"accuracy\"] > base_model_results[\"accuracy\"], (\n f\"Accuracy should be higher for the trained model: {trained_model_results['accuracy']} > {base_model_results['accuracy']}\"\n )\n assert trained_model_results[\"f1\"] > base_model_results[\"f1\"], (\n f\"F1 score should be higher for the trained model: {trained_model_results['f1']} > {base_model_results['f1']}\"\n )\n\n return base_model_results, trained_model_results\n\n\ndef train_integration():\n FP8_RECIPE_KWARGS = {\"fp8_format\": \"HYBRID\", \"amax_history_len\": 32, \"amax_compute_algo\": \"max\"}\n kwargs_handlers = [FP8RecipeKwargs(backend=\"TE\", **FP8_RECIPE_KWARGS)]\n AcceleratorState()._reset_state(True)\n accelerator = Accelerator(mixed_precision=\"fp8\", kwargs_handlers=kwargs_handlers)\n set_seed(42)\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = get_training_utilities(\n MODEL_NAME, accelerator=accelerator\n )\n\n model, optimizer, lr_scheduler = accelerator.prepare(model, optimizer, lr_scheduler)\n base_model_results = evaluate_model(model, eval_dataloader, METRIC)\n model.train()\n\n for batch in train_dataloader:\n outputs = model(**batch)\n loss = outputs.loss\n accelerator.backward(loss)\n optimizer.step()\n optimizer.zero_grad()\n lr_scheduler.step()\n\n trained_model_results = evaluate_model(model, eval_dataloader, METRIC)\n\n assert trained_model_results[\"accuracy\"] > base_model_results[\"accuracy\"], (\n f\"Accuracy should be higher for the trained model: {trained_model_results['accuracy']} > {base_model_results['accuracy']}\"\n )\n assert trained_model_results[\"f1\"] > base_model_results[\"f1\"], (\n f\"F1 score should be higher for the trained model: {trained_model_results['f1']} > {base_model_results['f1']}\"\n )\n\n return base_model_results, trained_model_results\n\n\nif __name__ == \"__main__\":\n baseline_not_trained, baseline_trained = train_baseline()\n accelerator_not_trained, accelerator_trained = train_integration()\n\n assert baseline_not_trained[\"accuracy\"] == accelerator_not_trained[\"accuracy\"], (\n f\"Accuracy should be the same for the baseline and accelerator: {baseline_not_trained['accuracy']} == {accelerator_not_trained['accuracy']}\"\n )\n assert baseline_not_trained[\"f1\"] == accelerator_not_trained[\"f1\"], (\n f\"F1 score should be the same for the baseline and accelerator: {baseline_not_trained['f1']} == {accelerator_not_trained['f1']}\"\n )\n assert baseline_trained[\"accuracy\"] == accelerator_trained[\"accuracy\"], (\n f\"Accuracy should be the same for the baseline and accelerator: {baseline_trained['accuracy']} == {accelerator_trained['accuracy']}\"\n )\n assert baseline_trained[\"f1\"] == accelerator_trained[\"f1\"], (\n f\"F1 score should be the same for the baseline and accelerator: {baseline_trained['f1']} == {accelerator_trained['f1']}\"\n )\n" }, { "path": "benchmarks/fsdp2/README.md", "content": "# FSDP2 Benchmarks\n\nThis benchmark showcases `FSDP2` in šŸ¤— `accelerate` and compares it to `torch` baseline.\n\n## Overview\n\nThis benchmark consists of two parts:\n- `main.py` is the main script that runs the benchmark\n- `visualize.py` is the script that visualizes the results (if `--output_dir` was specified for the previous command)\n\n## Motivation\n\nWe want to showcase that šŸ¤— `accelerate`'s integration of `FSDP2` is on par raw PyTorch, and highlight a \"broken\" part in PyTorch that creating an optimizer before applying `FSDP2` **doesn't result in a working training loop**. (more on this later)\nThis script showcases **matching memory usage and convergence between `accelerate` and `torch`'s baseline.**\nTo deal with this breaking change (and maintain backward compatibility with FSDP1 in terms of an API), `accelerate` had to come up with a workaround since `accelerate` assumes that the user will nearly always create a model, optimizer, scheduler, etc beforehand and bring them themselves. This lead to an issue of a stark increase in memory as well as the model not even training if the user creates an optimizer beforehand. \nTo workaround this, we replace the parameters inside the optimizer with the newly created FSDP2 sharded ones. More about this can be found in this [blog post (TBD)](TODO)\n> [!WARNING]\n> This script is intended to fit on 2x 24GB GPUs, though on so few GPUs it's not possible to see the memory difference (discrepancies in grad allocation result in lower memory usage in the non-fixed case), only the difference in convergence. Below are attached results from 8x H100 GPUs where the difference is visible.\n> TLDR: more GPUs = bigger memory difference between fixed and non-fixed cases.\n\n## Results\n\nHere are the results from running the benchmark on 8x H100 GPUs:\n\n

\n \"Allocated\n

\n

\n \"Reserved\n

\n\nAs you can see, the memory usage of `accelerate` and `torch_post_shard` (the **intended** way) are very similar, while `torch_pre_shard_not_fixed` uses significantly more memory. Our fix in `torch_pre_shard_fixed` brings the memory usage back in line with the **intended** approach.\n\n> [!WARNING]\n> Timing discrepancies are due to the benchmarks being ran in 1 script.\n\n\n## Running\n\nTo run the benchmark, you can either use `accelerate launch` or `torchrun`:\n```bash\naccelerate launch main.py\n```\n```bash\n# For two GPUs\ntorchrun --nproc_per_node 2 main.py\n```\n\nThis supports multiple configurable options, you can learn about them by running:\n```bash\npython3 main.py --help\n```\n\nThis script will run 4 different benchmarks:\n- `torch_optimizer_after_fsdp`: `torch` baseline where optimizer is created after applying `FSDP2`, this is the **intended** way to do it\n- `torch_optimizer_before_fsdp_not_fixed`: `torch` baseline where optimizer is created before applying `FSDP2` without fixing the optimizer parameters\n- `torch_optimizer_before_fsdp_fixed`: `torch` baseline where optimizer is created before applying `FSDP2` with our fix to the optimizer\n- `accelerate`: `accelerate`'s own integration of `FSDP2` where optimizer is created before applying `FSDP2`, but we apply our fix to the optimizer\n\nMemory results are saved in a folder specified by `--output_dir` argument.\nOptionally, you can specify `--save_memory_snapshot` to save the torch memory snapshot, which can then be viewed using [`torch memory viz`](https://pytorch.org/memory_viz)\n\n## Visualizing results\n\nTo visualize the results, you can run:\n\n```bash\npython3 visualize.py --dir \n```\n\nThis will then create two plots, showcasing allocated and reserved memory usage between all the different benchmarks discussed above.\n\n\n\n" }, { "path": "benchmarks/fsdp2/main.py", "content": "# Copyright 2025 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport functools\nfrom typing import Callable\n\nimport torch\n\nfrom accelerate import Accelerator\nfrom utils import parse_args, prepare_accelerate, prepare_torch\n\n\nMODEL_NAME = \"Qwen/Qwen2.5-1.5B-Instruct\"\nLEARNING_RATE = 3e-5\n\nCONFIG = {\n \"model_name\": MODEL_NAME,\n \"learning_rate\": LEARNING_RATE,\n}\n\n\ndef train(\n model: torch.nn.Module,\n optimizer: torch.optim.Optimizer,\n train_dataloader: torch.utils.data.DataLoader,\n accelerator: Accelerator,\n) -> torch.Tensor:\n losses = []\n for batch in train_dataloader:\n optimizer.zero_grad()\n outputs = model(**batch, use_cache=False)\n\n loss = outputs.loss\n losses.append(loss.item())\n accelerator.backward(loss)\n optimizer.step()\n\n return torch.tensor(losses)\n\n\ndef evaluate(args, config: dict, init_fn: Callable, run_name: str) -> torch.Tensor:\n model, optimizer, dataloader, accelerator, memory_tracker = init_fn(args, config)\n\n loss = train(model, optimizer, dataloader, accelerator)\n\n memory_tracker.stop()\n msg = f\"\"\"Results for {run_name} (rank 0):\nLoss: {loss[-1].item()}\nPeak Allocated Memory: {float(memory_tracker.peak_allocated_memory):.2f} MB\nPeak Reserved Memory: {float(memory_tracker.peak_reserved_memory):.2f} MB\n{\"-\" * 34}\"\"\"\n accelerator.print(msg)\n return loss\n\n\ndef main():\n args = parse_args()\n evaluations = [\n functools.partial(\n evaluate,\n init_fn=functools.partial(prepare_torch, post_shard_optimizer=False, apply_optimizer_fix=True),\n run_name=\"Optimizer Before FSDP (w/ fix)\",\n ),\n functools.partial(\n evaluate,\n init_fn=functools.partial(prepare_torch, post_shard_optimizer=False, apply_optimizer_fix=False),\n run_name=\"Optimizer Before FSDP (w/o fix)\",\n ),\n functools.partial(\n evaluate,\n init_fn=functools.partial(prepare_torch, post_shard_optimizer=True),\n run_name=\"Optimizer After FSDP\",\n ),\n functools.partial(evaluate, init_fn=prepare_accelerate, run_name=\"Accelerate\"),\n ]\n labels = [\n \"Optimizer Before FSDP (w/ fix)\",\n \"Optimizer Before FSDP (w/o fix)\",\n \"Optimizer After FSDP\",\n \"Accelerate\",\n ]\n\n results = {}\n torch.use_deterministic_algorithms(True)\n\n for evaluation, label in zip(evaluations, labels):\n results[label] = evaluation(args, CONFIG)\n\n torch.testing.assert_close(\n results[\"Optimizer After FSDP\"],\n results[\"Optimizer Before FSDP (w/ fix)\"],\n msg=\"Optimizer After FSDP and Optimizer Before FSDP (w/ fix) should be the same\",\n )\n\n torch.testing.assert_close(\n results[\"Optimizer After FSDP\"],\n results[\"Accelerate\"],\n msg=\"Optimizer After FSDP and Accelerate should be the same\",\n )\n\n torch.testing.assert_close(\n results[\"Accelerate\"],\n results[\"Optimizer Before FSDP (w/ fix)\"],\n msg=\"Accelerate and Optimizer Before FSDP (w/ fix) should be the same\",\n )\n\n torch.distributed.destroy_process_group()\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "benchmarks/fsdp2/measure_utils.py", "content": "# Copyright 2025 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport gc\nimport json\nimport os\nimport threading\nimport time\n\nimport psutil\nimport torch\n\nfrom accelerate import PartialState\n\n\nclass MemoryTracker:\n def __init__(\n self,\n device: torch.device,\n output_directory: str,\n run_name: str,\n save_memory_snapshot: bool,\n log_interval: float = 0.01,\n ):\n \"\"\"Class for tracking gpu and cpu memory usage of the process.\n\n Args:\n device (`torch.device`):\n PyTorch device to monitor.\n output_directory (`str`):\n Directory to save the memory usage data to, will be created if it doesn't exist.\n run_name (`str`):\n Name of the run, will be used to name the output files.\n save_memory_snapshot (`bool`):\n Whether to also save `torch.cuda.memory._dump_snapshot` to the output directory.\n log_interval (`float`, *optional*):\n Interval in seconds between memory measurements. Defaults to 0.01.\n \"\"\"\n self.log_interval = log_interval\n self.save_memory_snapshot = save_memory_snapshot\n self.output_directory = output_directory\n self.run_name = run_name\n\n self.timestamps = []\n self.allocated_memory = []\n self.reserved_memory = []\n self.virtual_memory = []\n\n self.start_time = None\n self.running = False\n\n self._thread = None\n self._state = PartialState()\n self._process = psutil.Process()\n self._device = device\n self.torch_accelerator_module = getattr(torch, device.type, torch.cuda)\n\n def _monitor(self):\n self.start_time = time.time()\n\n while self.running:\n allocated = self.torch_accelerator_module.memory_allocated(self._device) / (1024 * 1024)\n reserved = self.torch_accelerator_module.memory_reserved(self._device) / (1024 * 1024)\n virtual_memory = self._process.memory_info().rss / (1024 * 1024)\n\n self.allocated_memory.append(allocated)\n self.reserved_memory.append(reserved)\n self.virtual_memory.append(virtual_memory)\n self.timestamps.append(time.time() - self.start_time)\n\n time.sleep(self.log_interval)\n\n def start(self):\n gc.collect()\n self.torch_accelerator_module.empty_cache()\n\n if self.output_directory:\n os.makedirs(self.output_directory, exist_ok=True)\n\n if self.save_memory_snapshot:\n self.torch_accelerator_module.memory._record_memory_history()\n\n self.running = True\n self._thread = threading.Thread(target=self._monitor)\n self._thread.daemon = True\n self._thread.start()\n\n def stop(self):\n self.running = False\n if self._thread:\n self._thread.join()\n\n if self.save_memory_snapshot and self._state.is_main_process and self.output_directory:\n output_file = os.path.join(self.output_directory, f\"{self.run_name}_memory_snapshot.pkl\")\n self.torch_accelerator_module.memory._dump_snapshot(output_file)\n\n if self._state.is_main_process and self.output_directory:\n path = os.path.join(self.output_directory, f\"{self.run_name}_memory_usage.json\")\n with open(path, \"w\") as f:\n json.dump(\n {\n \"timestamps\": self.timestamps,\n \"allocated_memory\": self.allocated_memory,\n \"reserved_memory\": self.reserved_memory,\n \"virtual_memory\": self.virtual_memory,\n },\n f,\n )\n if self.save_memory_snapshot:\n self.torch_accelerator_module.memory._record_memory_history(False)\n self.torch_accelerator_module.empty_cache()\n\n @property\n def peak_allocated_memory(self):\n return max(self.allocated_memory)\n\n @property\n def peak_reserved_memory(self):\n return max(self.reserved_memory)\n" }, { "path": "benchmarks/fsdp2/utils.py", "content": "# Copyright 2025 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport argparse\nfrom types import MethodType\nfrom typing import Union\n\nimport torch\nfrom datasets import load_dataset\nfrom measure_utils import MemoryTracker\nfrom torch.distributed.fsdp import MixedPrecisionPolicy, fully_shard\nfrom torch.optim import AdamW\nfrom torch.utils.data import DataLoader\nfrom transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, DataCollatorForLanguageModeling\nfrom transformers.models.qwen2.modeling_qwen2 import Qwen2DecoderLayer\n\nfrom accelerate import Accelerator, FullyShardedDataParallelPlugin\nfrom accelerate.state import AcceleratorState, is_initialized\nfrom accelerate.utils import convert_outputs_to_fp32, set_seed\n\n\nSEED = 421\n\n\ndef get_named_parameters(model: torch.nn.Module, drop_refs: bool = False) -> dict[str, Union[torch.Tensor, int]]:\n \"\"\"\n This function returns a dictionary mapping the parameter names to their data pointers or\n the original parameters if `drop_refs` is `False`.\n It is used to get the original parameter names before `fully_shard` is applied.\n\n We only return the data pointers, so we drop the references to the original parameters\n and `fully_shard` will then trigger a new allocation for the sharded ones.\n\n Args:\n model (`torch.nn.Module`): Model instance to get the named parameters from\n drop_refs (`bool`, *optional*, defaults to `False`): Whether to drop the references to the original parameters\n\n Returns:\n `dict[str, Union[torch.Tensor, int]]`: Dictionary mapping the parameter names to their data pointers or the original parameters if `drop_refs` is `False`\n \"\"\"\n named_parameters = {}\n for n, p in model.named_parameters():\n # We only preserve the data pointers to have the unique 1:1 mapping between the original and the sharded parameters\n named_parameters[n] = p.data_ptr() if drop_refs else p\n return named_parameters\n\n\ndef replace_optimizer_params(optimizer: torch.optim.Optimizer):\n \"\"\"\n This function is called before using `fully_shard` on the model. It replaces the parameters of the optimizer with\n empty tensors, so `fully_shard` can trigger a new allocation for the sharded ones. After this, we swap the parameters\n `data_ptr` to the original one, so we can reuse that later to map the sharded parameters to the original ones.\n This function modifies the optimizer in-place.\n\n Args:\n optimizer (torch.optim.Optimizer): Optimizer instance which contains the original model parameters\n \"\"\"\n\n for param_group in optimizer.param_groups:\n for i, p in enumerate(param_group[\"params\"]):\n # We drop a reference to the original param here, so that _move_states_to_device triggers a reallocation\n # This is required or else the `fully_shard` -> `_move_states_to_device` uses the original memory address\n # for the sharded parameters, and we get a weird/undefined behavior.\n param_group[\"params\"][i] = torch.empty_like(p)\n\n # We save the original data_ptr, so we can swap back the parameters later\n param_group[\"params\"][i].data_ptr = p.data_ptr()\n\n\ndef swap_back_optimizer_params(\n model: torch.nn.Module, optimizer: torch.optim.Optimizer, old_named_parameter_pointers: dict[str, int]\n):\n \"\"\"\n This function is the counterpart of `replace_optimizer_params`. It is called after `fully_shard` being applied to\n the model. It swaps the parameters of the optimizer to their sharded counterparts.\n It is done using the `data_ptr` mapping prepared in `replace_optimizer_params` and `get_named_parameters`.\n\n Args:\n model (`torch.nn.Module`): Model instance to get the new named parameters from\n optimizer (`torch.optim.Optimizer`): Optimizer instance to swap the parameters of\n old_named_parameter_pointers (`dict[str, int]`): Dictionary mapping the original parameter names: data_ptrs to the new ones\n \"\"\"\n # We get the new named parameters after `fully_shard` being applied\n # We don't drop the references as we need the sharded parameters now\n new_named_parameters = get_named_parameters(model, drop_refs=False)\n\n # We create a mapping from the original data_ptr to the new sharded param corresponding to it\n mapping = {p: new_named_parameters[n] for n, p in old_named_parameter_pointers.items()}\n\n for param_group in optimizer.param_groups:\n # We swap the parameters of the optimizer to the new sharded ones\n param_group[\"params\"] = [mapping[p.data_ptr] for p in param_group[\"params\"]]\n\n\ndef parse_args():\n parser = argparse.ArgumentParser()\n parser.add_argument(\n \"--output_dir\",\n type=str,\n help=\"Directory to save the benchmarking results.\",\n )\n parser.add_argument(\n \"--save_memory_snapshot\",\n action=\"store_true\",\n default=False,\n help=\"If True, `torch.cuda.memory._dump_snapshot` will be used to additionaly save the memory trace.\",\n )\n ######################\n # Training arguments #\n ######################\n parser.add_argument(\n \"--batch_size\",\n type=int,\n default=2,\n help=\"Batch size for the training loop.\",\n )\n parser.add_argument(\n \"--block_size\",\n type=int,\n default=128,\n help=\"The maximum sequence length to use with the model.\",\n )\n parser.add_argument(\n \"--dataset_fraction\",\n type=float,\n default=1.0,\n help=\"Fraction of the dataset to use.\",\n )\n return parser.parse_args()\n\n\ndef prepare_dataloader(tokenizer, args, accelerator: Accelerator) -> DataLoader:\n dataset = load_dataset(\"tiny_shakespeare\", split=\"train\", trust_remote_code=True)\n\n def tokenize_function(example):\n return tokenizer(\n example[\"text\"],\n )\n\n dataset = dataset.map(\n tokenize_function,\n batched=True,\n remove_columns=[\"text\"],\n )\n\n block_size = min(tokenizer.model_max_length, args.block_size)\n\n def group_texts(examples):\n concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}\n total_length = len(concatenated_examples[list(examples.keys())[0]])\n\n total_length = (total_length // block_size) * block_size\n\n result = {\n k: [t[i : i + block_size] for i in range(0, total_length, block_size)]\n for k, t in concatenated_examples.items()\n }\n\n result[\"labels\"] = result[\"input_ids\"].copy()\n return result\n\n dataset = dataset.map(group_texts, batched=True)\n dataset = dataset.select(range(int(len(dataset) * args.dataset_fraction)))\n\n def collate_fn(examples):\n return DataCollatorForLanguageModeling(\n tokenizer=tokenizer,\n mlm=False,\n )(examples)\n\n dataloader = DataLoader(\n dataset,\n batch_size=args.batch_size,\n collate_fn=collate_fn,\n )\n dataloader = accelerator.prepare(dataloader)\n return dataloader\n\n\ndef get_model(model_name: str):\n # We reguire model to be loaded in fp32, otherwise benchmarks don't match as accelerate does upcasting of parameters to fp32\n config = AutoConfig.from_pretrained(model_name, trust_remote_code=True, torch_dtype=torch.float32)\n model = AutoModelForCausalLM.from_config(config)\n return model\n\n\ndef get_tokenizer(model_name: str):\n tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)\n tokenizer.pad_token = tokenizer.eos_token\n return tokenizer\n\n\ndef prepare_torch(\n args, config: dict, post_shard_optimizer: bool = False, apply_optimizer_fix: bool = False\n) -> tuple[torch.nn.Module, torch.optim.Optimizer, torch.utils.data.DataLoader, Accelerator]:\n mp_policy = MixedPrecisionPolicy(\n param_dtype=torch.bfloat16,\n reduce_dtype=torch.bfloat16,\n output_dtype=torch.bfloat16,\n )\n\n accelerator = Accelerator(mixed_precision=\"bf16\")\n set_seed(SEED)\n is_fixed = \"fixed\" if apply_optimizer_fix else \"not_fixed\"\n is_post_shard = \"optimizer_after_fsdp\" if post_shard_optimizer else \"optimizer_before_fsdp\"\n run_name = f\"torch_{is_post_shard}\" if post_shard_optimizer else f\"torch_{is_post_shard}_{is_fixed}\"\n\n tokenizer = get_tokenizer(config[\"model_name\"])\n train_dataloader = prepare_dataloader(tokenizer, args, accelerator)\n\n memory_tracker = MemoryTracker(accelerator.device, args.output_dir, run_name, args.save_memory_snapshot)\n memory_tracker.start()\n\n model = get_model(config[\"model_name\"])\n optimizer = None\n\n if not post_shard_optimizer:\n optimizer = AdamW(model.parameters(), lr=config[\"learning_rate\"])\n\n if apply_optimizer_fix:\n # We drop the references to the original parameters, so that `fully_shard` can trigger a new allocation\n # Then we get the `module_name: data_ptr` mapping, so we can swap back the parameters later\n old_named_parameters = get_named_parameters(model, drop_refs=True)\n\n # We replace the parameters of the optimizer with empty tensors, so that `fully_shard` can trigger a new allocation\n # We also change the `data_ptr` of the parameters to the original ones, so we can swap back the parameters later\n replace_optimizer_params(optimizer)\n\n for module in model.modules():\n if isinstance(module, Qwen2DecoderLayer):\n fully_shard(module, mp_policy=mp_policy)\n fully_shard(model, mp_policy=mp_policy)\n\n # We do this to imitate how accelerate forces outputs to be in fp32 via `convert_outputs_to_fp32`\n autocast_context = torch.autocast(device_type=accelerator.state.device.type, dtype=torch.bfloat16)\n model_forward_func = model.forward.__func__\n new_forward = autocast_context(model_forward_func)\n model.forward = MethodType(new_forward, model)\n model.forward = MethodType(convert_outputs_to_fp32(model.forward.__func__), model)\n\n if post_shard_optimizer:\n optimizer = AdamW(model.parameters(), lr=config[\"learning_rate\"])\n\n if not post_shard_optimizer and apply_optimizer_fix:\n # We swap back the parameters of the optimizer to the original ones\n swap_back_optimizer_params(model, optimizer, old_named_parameters)\n\n return model, optimizer, train_dataloader, accelerator, memory_tracker\n\n\ndef prepare_accelerate(\n args, config: dict\n) -> tuple[torch.nn.Module, torch.optim.Optimizer, torch.utils.data.DataLoader, Accelerator]:\n if is_initialized():\n AcceleratorState()._reset_state(True)\n\n fsdp_plugin = FullyShardedDataParallelPlugin(\n fsdp_version=2,\n auto_wrap_policy=\"transformer_based_wrap\",\n transformer_cls_names_to_wrap=[\"Qwen2DecoderLayer\"],\n )\n accelerator = Accelerator(\n fsdp_plugin=fsdp_plugin,\n mixed_precision=\"bf16\",\n )\n set_seed(SEED)\n\n tokenizer = get_tokenizer(config[\"model_name\"])\n train_dataloader = prepare_dataloader(tokenizer, args, accelerator)\n\n memory_tracker = MemoryTracker(accelerator.device, args.output_dir, \"accelerate\", args.save_memory_snapshot)\n memory_tracker.start()\n\n model = get_model(config[\"model_name\"])\n optimizer = AdamW(model.parameters(), lr=config[\"learning_rate\"])\n\n model, optimizer = accelerator.prepare(model, optimizer)\n\n return model, optimizer, train_dataloader, accelerator, memory_tracker\n" }, { "path": "benchmarks/fsdp2/visualize.py", "content": "# Copyright 2025 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport argparse\nimport json\n\nimport matplotlib.pyplot as plt\n\n\ndef parse_args():\n parser = argparse.ArgumentParser()\n parser.add_argument(\"--dir\", type=str, help=\"Directory containing the memory usage data\")\n parser.add_argument(\n \"--memory_threshold\",\n type=int,\n default=0,\n help=\"Memory threshold to filter data that is below this value (only filters 1st `--filter_partition` of the points which should roughtly correspond to the model loading)\",\n )\n parser.add_argument(\n \"--filter_partition\",\n type=float,\n default=1 / 3,\n help=\"Partition to drop data from that are below the memory threshold\",\n )\n return parser.parse_args()\n\n\ndef filter_data(data, memory_threshold, filter_partition, key):\n timestamps = data[\"timestamps\"]\n memory = data[key]\n\n mid_point = int(len(timestamps) * filter_partition)\n filtered_times = []\n filtered_memory = []\n for i, (t, m) in enumerate(zip(timestamps, memory)):\n if i < mid_point and m < memory_threshold:\n continue\n filtered_times.append(t)\n filtered_memory.append(m)\n return filtered_times, filtered_memory\n\n\ndef compare_memory_usage(data, labels, memory_threshold, filter_partition):\n plt.style.use(\"seaborn-v0_8\")\n colors = [\"#2ecc71\", \"#e74c3c\", \"#3498db\", \"#f1c40f\"]\n\n fig1, ax1 = plt.subplots(figsize=(15, 5))\n for data_item, label, color in zip(data, labels, colors):\n timestamps, allocated = filter_data(data_item, memory_threshold, filter_partition, \"allocated_memory\")\n ax1.plot(timestamps, allocated, label=label, color=color, linewidth=2)\n\n ax1.set_xlabel(\"Time (s)\", fontsize=12)\n ax1.set_ylabel(\"Allocated Memory (GB)\", fontsize=12)\n ax1.set_title(\"Allocated Memory Usage Over Time\", fontsize=14, pad=15)\n ax1.grid(True, linestyle=\"--\", alpha=0.7)\n ax1.legend(frameon=True, fancybox=True, shadow=True, fontsize=10)\n ax1.spines[\"top\"].set_visible(False)\n ax1.spines[\"right\"].set_visible(False)\n plt.tight_layout()\n\n fig2, ax2 = plt.subplots(figsize=(15, 5))\n for data_item, label, color in zip(data, labels, colors):\n timestamps, reserved = filter_data(data_item, memory_threshold, filter_partition, \"reserved_memory\")\n ax2.plot(timestamps, reserved, label=label, color=color, linewidth=2)\n\n ax2.set_xlabel(\"Time (s)\", fontsize=12)\n ax2.set_ylabel(\"Reserved Memory (GB)\", fontsize=12)\n ax2.set_title(\"Reserved Memory Usage Over Time\", fontsize=14, pad=15)\n ax2.grid(True, linestyle=\"--\", alpha=0.7)\n ax2.legend(frameon=True, fancybox=True, shadow=True, fontsize=10)\n ax2.spines[\"top\"].set_visible(False)\n ax2.spines[\"right\"].set_visible(False)\n plt.tight_layout()\n\n return fig1, fig2\n\n\nif __name__ == \"__main__\":\n args = parse_args()\n DIR = args.dir\n with open(f\"{DIR}/torch_optimizer_before_fsdp_not_fixed_memory_usage.json\") as f:\n optimizer_before_fsdp_not_fixed = json.load(f)\n\n with open(f\"{DIR}/torch_optimizer_after_fsdp_memory_usage.json\") as f:\n optimizer_after_fsdp = json.load(f)\n\n with open(f\"{DIR}/torch_optimizer_before_fsdp_fixed_memory_usage.json\") as f:\n optimizer_before_fsdp_fixed = json.load(f)\n\n with open(f\"{DIR}/accelerate_memory_usage.json\") as f:\n accelerate = json.load(f)\n\n data = [optimizer_before_fsdp_not_fixed, optimizer_before_fsdp_fixed, optimizer_after_fsdp, accelerate]\n labels = [\n \"Optimizer Before FSDP (w/o fix)\",\n \"Optimizer Before FSDP (w/ fix)\",\n \"Optimizer After FSDP\",\n \"Accelerate\",\n ]\n\n fig1, fig2 = compare_memory_usage(data, labels, args.memory_threshold, args.filter_partition)\n fig1.savefig(f\"{DIR}/allocated_memory.png\")\n fig2.savefig(f\"{DIR}/reserved_memory.png\")\n" }, { "path": "benchmarks/torch.compile/README.md", "content": "# Regional Compilation Benchmark\n\nThis benchmark compares different compilation strategies using PyTorch's `torch.compile` and Accelerate's `compile_regions` utility, which is based on the recipe in [PyTorch documentation](https://pytorch.org/tutorials/recipes/regional_compilation.html).\n\n## Overview\n\nThe benchmark evaluates three approaches:\n\n- **Baseline**: No compilation, standard PyTorch eager execution.\n- **Full compilation**: Using PyTorch's `torch.compile()` on the entire model.\n- **Regional compilation**: Using `accelerate.utils.compile_regions()` which targets specific blocks of the model to optimize compilation time.\n\nEach approach is tested with different batch sizes (1 and 4) and sequence lengths (128) on various LLaMA-based models ranging from 1B to 13B parameters. We purposefully run the forward pass outside of the `torch.no_grad()` context to simulate performance in a training environment, where gradients are needed.\n\n## Usage\n\nTo run this benchmark:\n\n```bash\npython regional_compilation.py\n```\n\nThe script will automatically download the model configurations, create models, and benchmark both compilation and inference times across different scenarios.\n\n## Requirements\n\n- Suitable GPU memory for the models being tested.\n- PyTorch with CUDA support.\n- Transformers library.\n- Accelerate library.\n\n## Results\n\nThe benchmark results are summarized in the following figures:\n\n- Compilation time is how long it takes to run the first forward pass.\n- Speedup factor is the ratio of non-compiled baseline inference time to the fully/regionally compiled inference time.\n\n

\n \"Compilation\n

\n

\n \"Speedup\n

\n\nFull results are available in the tables below:\n\n```markdown\n[-------------------------------------------------- NousResearch/Llama-3.2-1B ---------------------------------------------------]\n | Inference time (1x128) | Inference time (4x128) | Compile time (1x128) | Compile time (4x128)\n1 threads: -----------------------------------------------------------------------------------------------------------------------\n Baseline | 18.3 | 18.4 | | \n Full compilation | 6.3 | 10.0 | 10696.4 | 10248.0 \n Regional compilation | 9.7 | 10.0 | 1952.7 | 2903.9 \n\nTimes are in milliseconds (ms).\n\n[---------------------------------------------- NousResearch/Hermes-3-Llama-3.2-3B ----------------------------------------------]\n | Inference time (1x128) | Inference time (4x128) | Compile time (1x128) | Compile time (4x128)\n1 threads: -----------------------------------------------------------------------------------------------------------------------\n Baseline | 33.4 | 33.6 | | \n Full compilation | 11.2 | 23.9 | 17857.5 | 17736.5 \n Regional compilation | 17.3 | 23.7 | 2993.2 | 2478.8 \n\nTimes are in milliseconds (ms).\n\n[---------------------------------------------- NousResearch/Hermes-3-Llama-3.1-8B ----------------------------------------------]\n | Inference time (1x128) | Inference time (4x128) | Compile time (1x128) | Compile time (4x128)\n1 threads: -----------------------------------------------------------------------------------------------------------------------\n Baseline | 40.3 | 59.5 | | \n Full compilation | 18.9 | 54.4 | 20437.8 | 20152.3 \n Regional compilation | 19.7 | 54.0 | 2903.1 | 2438.0 \n\nTimes are in milliseconds (ms).\n\n[--------------------------------------------- NousResearch/Nous-Hermes-Llama2-13b ----------------------------------------------]\n | Inference time (1x128) | Inference time (4x128) | Compile time (1x128) | Compile time (4x128)\n1 threads: -----------------------------------------------------------------------------------------------------------------------\n Baseline | 45.5 | 100.4 | | \n Full compilation | 29.4 | 89.7 | 23099.4 | 22885.9 \n Regional compilation | 29.4 | 87.5 | 2945.5 | 2526.2 \n\nTimes are in milliseconds (ms).\n```\n\n## Results Summary\n\n### Compilation Time\n\nRegional compilation provides significantly faster compilation times compared to full model compilation:\n\n- **Full compilation**: Takes ~10-23 seconds depending on model size.\n- **Regional compilation**: Takes only ~2-3 seconds across all model sizes.\n- **Speed improvement**: Regional compilation is **5-9x faster** to compile.\n\n### Inference Time\n\nRegional compilation delivers inference performance close to full compilation:\n\n- For batch size 1:\n - For smaller models (1B-3B): Full compilation has a slight edge over regional compilation.\n - For larger models (8B-13B): Regional compilation performs similarly to full compilation.\n- For batch size 4: Regional compilation performs similarly to full compilation across all models.\n\n## Key Takeaways\n\n1. **Comparable Performance**: Regional compilation delivers performance speedups similar to full compilation, especially for larger models.\n2. **Faster Compilation**: Regional compilation significantly reduces the time taken to compile models, making it a more efficient choice for deployment.\n3. **Batch Size Impact**: At batch size 4, full compilation and regional compilation perform nearly identically.\n4. **Model Size Impact**: Even with a small batch size, full compilation and regional compilation perform similarly for larger models (8B-13B).\n5. **Practical Application**: For real-world applications, regional compilation is a practical choice for optimizing training cold start times, especially when working with large models.\n" }, { "path": "benchmarks/torch.compile/regional_compilation.py", "content": "# Copyright 2025 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport torch\nfrom torch.utils.benchmark import Compare, Timer\nfrom transformers import AutoConfig, AutoModelForCausalLM\n\nfrom accelerate.test_utils.testing import get_backend\nfrom accelerate.utils import compile_regions\n\n\ntorch.set_float32_matmul_precision(\"high\")\n\nCOMPILE_ITERS = 2\nINFERENCE_ITERS = 100\n\nBASELINE = \"Baseline\"\nCOMPILE_TIME = \"Compile time\"\nINFRENCE_TIME = \"Inference time\"\nFULL_COMPILATION = \"Full compilation\"\nREGIONAL_COMPILATION = \"Regional compilation\"\n\nINFRENCE_STMT = \"model(input_ids, use_cache=False)\"\nCOMPILE_STMT = f\"torch._dynamo.reset(); torch._inductor.utils.clear_inductor_caches(); {INFRENCE_STMT}\"\n\ntorch_device_type, _, _ = get_backend()\n\nresults = []\nfor model_id in [\n # non-gated llama models\n \"NousResearch/Llama-3.2-1B\",\n \"NousResearch/Hermes-3-Llama-3.2-3B\",\n \"NousResearch/Hermes-3-Llama-3.1-8B\",\n \"NousResearch/Nous-Hermes-Llama2-13b\",\n]:\n with torch.device(torch_device_type):\n config = AutoConfig.from_pretrained(model_id)\n model = AutoModelForCausalLM.from_config(config).to(dtype=torch.float16).eval()\n\n full_compilation_model = torch.compile(model)\n regional_compilation_model = compile_regions(model)\n\n for model, sub_label, description, stmt, iters in [\n (model, BASELINE, INFRENCE_TIME, INFRENCE_STMT, INFERENCE_ITERS),\n (full_compilation_model, FULL_COMPILATION, COMPILE_TIME, COMPILE_STMT, COMPILE_ITERS),\n (full_compilation_model, FULL_COMPILATION, INFRENCE_TIME, INFRENCE_STMT, INFERENCE_ITERS),\n (regional_compilation_model, REGIONAL_COMPILATION, COMPILE_TIME, COMPILE_STMT, COMPILE_ITERS),\n (regional_compilation_model, REGIONAL_COMPILATION, INFRENCE_TIME, INFRENCE_STMT, INFERENCE_ITERS),\n ]:\n for batch_size, sequence_length in [(1, 128), (4, 128)]:\n input_ids = torch.randint(\n 0, 1000, size=(batch_size, sequence_length), dtype=torch.int64, device=torch_device_type\n )\n results.append(\n Timer(\n label=model_id,\n sub_label=sub_label,\n description=f\"{description} ({batch_size}x{sequence_length})\",\n globals={\"model\": model, \"input_ids\": input_ids},\n stmt=stmt,\n ).timeit(number=iters)\n )\n\ncompare = Compare(results)\ncompare.colorize()\ncompare.print()\n" }, { "path": "docker/README.md", "content": "\n\n# Official Hugging Face Accelerate Docker Images\n\nAccelerate publishes a variety of docker versions as part of our CI that users can also use. These are stable images that Accelerate can run off of which comes with a variety of different setup configurations, all of which are officially hosted on [Docker Hub](https://hub.docker.com/r/huggingface/accelerate).\n\nA breakdown of each are given below\n\n## Naming Conventions\n\nAccelerate docker images follow a tagging convention of:\n\n```bash\nhuggingface/accelerate:{accelerator}-{nightly,release}\n```\n\n`accelerator` in this instance is one of many applical pre-configured backend supports:\n* `gpu`: Comes compiled off of the `nvidia/cuda` image and includes core parts like `bitsandbytes`. Runs off python 3.9.\n* `cpu`: Comes compiled off of `python:3.9-slim` and is designed for non-CUDA based workloads.\n* More to come soon\n* `gpu-deepspeed`: Comes compiled off of the `nvidia/cuda` image and includes core parts like `bitsandbytes` as well as the latest `deepspeed` version. Runs off python 3.10.\n* `gpu-fp8-transformerengine`: Comes compiled off of `nvcr.io/nvidia/pytorch` and is specifically for running the `benchmarks/fp8` scripts on devices which support FP8 operations using the `TransformerEngine` library (RTX 4090, H100, etc)\n\n## Nightlies vs Releases\n\nEach release a new build is pushed with a version number included in the name. For a GPU-supported image of version 0.28.0 for instance, it would look like the following:\n\n```bash\nhuggingface/accelerate:gpu-release-0.28.0\n```\n\nNightlies contain two different image tags. There is a general `nightly` tag which is built each night, and a `nightly-YYYY-MM-DD` which corresponds to a build from a particular date.\n\nFor instance, here is an example nightly CPU image from 3/14/2024\n\n```bash\nhuggingface/accelerate:cpu-nightly-2024-03-14\n```\n\n## Running the images\n\nEach image comes compiled with `conda` and an `accelerate` environment contains all of the installed dependencies. \n\nTo pull down the latest nightly run:\n\n```bash\ndocker pull huggingface/accelerate:gpu-nightly\n```\n\nTo then run it in interactive mode with GPU-memory available, run:\n\n```bash\ndocker container run --gpus all -it huggingface/accelerate:gpu-nightly\n```\n\n## DEPRECATED IMAGES\n\nCPU and GPU docker images were hosted at `huggingface/accelerate-gpu` and `huggingface/accelerate-cpu`. These builds are now outdated and will not receive updates. \n\nThe builds at the corresponding `huggingface/accelerate:{gpu,cpu}` contain the same `Dockerfile`, so it's as simple as changing the docker image to the desired ones from above. We will not be deleting these images for posterity, but they will not be receiving updates going forward." }, { "path": "docker/accelerate-cpu/Dockerfile", "content": "# Builds CPU-only Docker image of PyTorch\n# Uses multi-staged approach to reduce size\n# Stage 1\nFROM python:3.10-slim as compile-image\n\nARG DEBIAN_FRONTEND=noninteractive\n\nRUN apt update\nRUN apt-get install -y --no-install-recommends \\\n build-essential \\\n git \\\n gcc\n\n# Setup virtual environment for Docker\nENV VIRTUAL_ENV=/opt/venv\nRUN python3 -m venv ${VIRTUAL_ENV}\n# Make sure we use the virtualenv\nENV PATH=\"${VIRTUAL_ENV}/bin:$PATH\"\nWORKDIR /workspace\n# Install specific CPU torch wheel to save on space\nRUN python3 -m pip install --upgrade --no-cache-dir pip\nRUN python3 -m pip install --no-cache-dir \\\n jupyter \\\n git+https://github.com/huggingface/accelerate#egg=accelerate[testing,test_trackers] \\\n --extra-index-url https://download.pytorch.org/whl/cpu\n \n# Stage 2\nFROM python:3.10-slim AS build-image\nCOPY --from=compile-image /opt/venv /opt/venv\nRUN useradd -ms /bin/bash user\nUSER user\n\n# Make sure we use the virtualenv\nENV PATH=\"/opt/venv/bin:$PATH\"\nCMD [\"/bin/bash\"]" }, { "path": "docker/accelerate-gpu/Dockerfile", "content": "# Builds GPU docker image of PyTorch specifically\n# Uses multi-staged approach to reduce size\n# Stage 1\n# Use base conda image to reduce time\nFROM continuumio/miniconda3:latest AS compile-image\n# Specify py version\nENV PYTHON_VERSION=3.10\n# Install apt libs\nRUN apt-get update && \\\n apt-get install -y curl git wget && \\\n apt-get clean && \\\n rm -rf /var/lib/apt/lists*\n\n# Create our conda env\nRUN conda create --name accelerate python=${PYTHON_VERSION} ipython jupyter pip\n# We don't install pytorch here yet since CUDA isn't available\n# instead we use the direct torch wheel\nENV PATH /opt/conda/envs/accelerate/bin:$PATH\n# Activate our bash shell\nRUN chsh -s /bin/bash\nSHELL [\"/bin/bash\", \"-c\"]\n# Activate the conda env, install mpy4pi, and install torch + accelerate\nRUN source activate accelerate && conda install -c conda-forge mpi4py\nRUN source activate accelerate && \\\n python3 -m pip install --no-cache-dir \\\n git+https://github.com/huggingface/accelerate#egg=accelerate[testing,test_trackers] \\\n --extra-index-url https://download.pytorch.org/whl/cu126\n\nRUN python3 -m pip install --no-cache-dir bitsandbytes\n\n# Stage 2\nFROM nvidia/cuda:12.6.3-cudnn-devel-ubuntu22.04 AS build-image\nCOPY --from=compile-image /opt/conda /opt/conda\nENV PATH /opt/conda/bin:$PATH\n\n# Install apt libs\nRUN apt-get update && \\\n apt-get install -y curl git wget && \\\n apt-get clean && \\\n rm -rf /var/lib/apt/lists*\n\nRUN echo \"source activate accelerate\" >> ~/.profile\n\n# Activate the virtualenv\nCMD [\"/bin/bash\"]" }, { "path": "docker/accelerate-gpu-deepspeed/Dockerfile", "content": "# Builds GPU docker image of PyTorch specifically\n# Uses multi-staged approach to reduce size\n# Stage 1\n# Use base conda image to reduce time\nFROM continuumio/miniconda3:latest AS compile-image\n# Specify py version\nENV PYTHON_VERSION=3.10\n# Install apt libs\nRUN apt-get update && \\\n apt-get install -y curl git wget && \\\n apt-get clean && \\\n rm -rf /var/lib/apt/lists*\n\n# Create our conda env\nRUN conda create --name accelerate python=${PYTHON_VERSION} ipython jupyter pip\n# We don't install pytorch here yet since CUDA isn't available\n# instead we use the direct torch wheel\nENV PATH /opt/conda/envs/accelerate/bin:$PATH\n# Activate our bash shell\nRUN chsh -s /bin/bash\nSHELL [\"/bin/bash\", \"-c\"]\n# Activate the conda env, install mpy4pi, and install torch + accelerate\nRUN source activate accelerate && conda install -c conda-forge mpi4py\nRUN source activate accelerate && \\\n python3 -m pip install --no-cache-dir \\\n git+https://github.com/huggingface/accelerate#egg=accelerate[testing,test_trackers,deepspeed] \\\n --extra-index-url https://download.pytorch.org/whl/cu126\n\nRUN python3 -m pip install --no-cache-dir bitsandbytes\n\n# Stage 2\nFROM nvidia/cuda:12.6.3-cudnn-devel-ubuntu22.04 AS build-image\nCOPY --from=compile-image /opt/conda /opt/conda\nENV PATH /opt/conda/bin:$PATH\n\n# Install apt libs\nRUN apt-get update && \\\n apt-get install -y curl git wget && \\\n apt-get clean && \\\n rm -rf /var/lib/apt/lists*\n\nRUN echo \"source activate accelerate\" >> ~/.profile\n\n# Activate the virtualenv\nCMD [\"/bin/bash\"]" }, { "path": "docs/Makefile", "content": "# Minimal makefile for Sphinx documentation\n#\n\n# You can set these variables from the command line.\nSPHINXOPTS =\nSPHINXBUILD = sphinx-build\nSOURCEDIR = source\nBUILDDIR = _build\n\n# Put it first so that \"make\" without argument is like \"make help\".\nhelp:\n\t@$(SPHINXBUILD) -M help \"$(SOURCEDIR)\" \"$(BUILDDIR)\" $(SPHINXOPTS) $(O)\n\n.PHONY: help Makefile\n\n# Catch-all target: route all unknown targets to Sphinx using the new\n# \"make mode\" option. $(O) is meant as a shortcut for $(SPHINXOPTS).\n%: Makefile\n\t@$(SPHINXBUILD) -M $@ \"$(SOURCEDIR)\" \"$(BUILDDIR)\" $(SPHINXOPTS) $(O)" }, { "path": "docs/README.md", "content": "\n\n# Generating the documentation\n\nTo generate the documentation, you first have to build it. Several packages are necessary to build the doc, \nyou can install them with the following command, at the root of the code repository:\n\n```bash\npip install -e \".[docs]\"\n```\n\nThen you need to install our special tool that builds the documentation:\n\n```bash\npip install git+https://github.com/huggingface/doc-builder\n```\n\n---\n**NOTE**\n\nYou only need to generate the documentation to inspect it locally (if you're planning changes and want to\ncheck how they look before committing for instance). You don't have to commit the built documentation.\n\n---\n\n## Building the documentation\n\nOnce you have setup the `doc-builder` and additional packages, you can generate the documentation by \ntyping the following command:\n\n```bash\ndoc-builder build accelerate docs/source/ --build_dir ~/tmp/test-build\n```\n\nYou can adapt the `--build_dir` to set any temporary folder that you prefer. This command will create it and generate\nthe MDX files that will be rendered as the documentation on the main website. You can inspect them in your favorite\nMarkdown editor.\n\n## Previewing the documentation\n\nTo preview the docs, first install the `watchdog` module with:\n\n```bash\npip install watchdog\n```\n\nThen run the following command:\n\n```bash\ndoc-builder preview {package_name} {path_to_docs}\n```\n\nFor example:\n\n```bash\ndoc-builder preview accelerate docs/source/\n```\n\nThe 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.\n\n---\n**NOTE**\n\nThe `preview` command only works with existing doc files. When you add a completely new file, you need to update `_toctree.yml` & restart `preview` command (`ctrl-c` to stop it & call `doc-builder preview ...` again).\n\n---\n\n## Adding a new element to the navigation bar\n\nAccepted files are Markdown (.md).\n\nCreate a file with its extension and put it in the source directory. You can then link it to the toc-tree by putting\nthe filename without the extension in the [`_toctree.yml`](https://github.com/huggingface/accelerate/blob/main/docs/source/_toctree.yml) file.\n\n## Renaming section headers and moving sections\n\nIt helps to keep the old links working when renaming the section header and/or moving sections from one document to another. This is because the old links are likely to be used in Issues, Forums, and Social media and it'd make for a much more superior user experience if users reading those months later could still easily navigate to the originally intended information.\n\nTherefore, we simply keep a little map of moved sections at the end of the document where the original section was. The key is to preserve the original anchor.\n\nSo if you renamed a section from: \"Section A\" to \"Section B\", then you can add at the end of the file:\n\n```\nSections that were moved:\n\n[ Section A ]\n```\nand of course, if you moved it to another file, then:\n\n```\nSections that were moved:\n\n[ Section A ]\n```\n\nUse the relative style to link to the new file so that the versioned docs continue to work.\n\n\n## Writing Documentation - Specification\n\nThe `huggingface/accelerate` documentation follows the\n[Google documentation](https://sphinxcontrib-napoleon.readthedocs.io/en/latest/example_google.html) style for docstrings,\nalthough we can write them directly in Markdown.\n\n### Adding a new tutorial\n\nAdding a new tutorial or section is done in two steps:\n\n- Add a new file under `./source`. This file can either be ReStructuredText (.rst) or Markdown (.md).\n- Link that file in `./source/_toctree.yml` on the correct toc-tree.\n\nMake sure to put your new file under the proper section. It's unlikely to go in the first section (*Get Started*), so\ndepending on the intended targets (beginners, more advanced users, or researchers) it should go in sections two, three, or\nfour.\n\n### Writing source documentation\n\nValues that should be put in `code` should either be surrounded by backticks: \\`like so\\`. Note that argument names\nand objects like True, None, or any strings should usually be put in `code`.\n\nWhen mentioning a class, function, or method, it is recommended to use our syntax for internal links so that our tool\nadds a link to its documentation with this syntax: \\[\\`XXXClass\\`\\] or \\[\\`function\\`\\]. This requires the class or \nfunction to be in the main package.\n\nIf you want to create a link to some internal class or function, you need to\nprovide its path. For instance: \\[\\`utils.gather\\`\\]. This will be converted into a link with\n`utils.gather` in the description. To get rid of the path and only keep the name of the object you are\nlinking to in the description, add a ~: \\[\\`~utils.gather\\`\\] will generate a link with `gather` in the description.\n\nThe same works for methods so you can either use \\[\\`XXXClass.method\\`\\] or \\[~\\`XXXClass.method\\`\\].\n\n#### Defining arguments in a method\n\nArguments should be defined with the `Args:` (or `Arguments:` or `Parameters:`) prefix, followed by a line return and\nan indentation. The argument should be followed by its type, with its shape if it is a tensor, a colon, and its\ndescription:\n\n```\n Args:\n n_layers (`int`): The number of layers of the model.\n```\n\nIf the description is too long to fit in one line (more than 119 characters in total), another indentation is necessary \nbefore writing the description after the argument.\n\nFinally, to maintain uniformity if any *one* description is too long to fit on one line, the \nrest of the parameters should follow suit and have an indention before their description.\n\nHere's an example showcasing everything so far:\n\n```\n Args:\n gradient_accumulation_steps (`int`, *optional*, default to 1):\n The number of steps that should pass before gradients are accumulated. A number > 1 should be combined with `Accelerator.accumulate`.\n cpu (`bool`, *optional*):\n Whether or not to force the script to execute on CPU. Will ignore GPU available if set to `True` and force the execution on one process only.\n```\n\nFor optional arguments or arguments with defaults we follow the following syntax: imagine we have a function with the\nfollowing signature:\n\n```\ndef my_function(x: str = None, a: float = 1):\n```\n\nthen its documentation should look like this:\n\n```\n Args:\n x (`str`, *optional*):\n This argument controls ... and has a description longer than 119 chars.\n a (`float`, *optional*, defaults to 1):\n This argument is used to ... and has a description longer than 119 chars.\n```\n\nNote that we always omit the \"defaults to \\`None\\`\" when None is the default for any argument. Also note that even\nif the first line describing your argument type and its default gets long, you can't break it on several lines. You can\nhowever write as many lines as you want in the indented description (see the example above with `input_ids`).\n\n#### Writing a multi-line code block\n\nMulti-line code blocks can be useful for displaying examples. They are done between two lines of three backticks as usual in Markdown:\n\n\n````\n```python\n# first line of code\n# second line\n# etc\n```\n````\n\n#### Writing a return block\n\nThe return block should be introduced with the `Returns:` prefix, followed by a line return and an indentation.\nThe first line should be the type of the return, followed by a line return. No need to indent further for the elements\nbuilding the return.\n\nHere's an example of a single value return:\n\n```\n Returns:\n `List[int]`: A list of integers in the range [0, 1] --- 1 for a special token, 0 for a sequence token.\n```\n\nHere's an example of a tuple return, comprising several objects:\n\n```\n Returns:\n `tuple(torch.FloatTensor)` comprising various elements depending on the configuration ([`BertConfig`]) and inputs:\n - ** loss** (*optional*, returned when `masked_lm_labels` is provided) `torch.FloatTensor` of shape `(1,)` --\n Total loss is the sum of the masked language modeling loss and the next sequence prediction (classification) loss.\n - **prediction_scores** (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`) --\n Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).\n```\n\n## Styling the docstring\n\nWe have an automatic script running with the `make style` comment that will make sure that:\n- the docstrings fully take advantage of the line width\n- all code examples are formatted using black, like the code of the Transformers library\n\nThis script may have some weird failures if you made a syntax mistake or if you uncover a bug. Therefore, it's\nrecommended to commit your changes before running `make style`, so you can revert the changes done by that script\neasily.\n\n## Writing documentation examples\n\nThe syntax for Example docstrings can look as follows:\n\n```\n Example:\n\n ```python\n >>> import time\n >>> from accelerate import Accelerator\n >>> accelerator = Accelerator()\n >>> if accelerator.is_main_process:\n ... time.sleep(2)\n >>> else:\n ... print(\"I'm waiting for the main process to finish its sleep...\")\n >>> accelerator.wait_for_everyone()\n >>> # Should print on every process at the same time\n >>> print(\"Everyone is here\")\n ```\n```\n\nThe docstring should give a minimal, clear example of how the respective function \nis to be used in inference and also include the expected (ideally sensible)\noutput.\nOften, readers will try out the example before even going through the function \nor class definitions. Therefore, it is of utmost importance that the example \nworks as expected." }, { "path": "docs/source/_toctree.yml", "content": "- sections:\n - local: index\n title: šŸ¤— Accelerate\n - local: basic_tutorials/install\n title: Installation\n - local: quicktour\n title: Quicktour\n title: Getting started\n- sections:\n - local: basic_tutorials/overview\n title: Overview\n - local: basic_tutorials/migration\n title: Add Accelerate to your code\n - local: basic_tutorials/execution\n title: Execution process\n - local: basic_tutorials/tpu\n title: TPU training\n - local: basic_tutorials/launch\n title: Launching Accelerate scripts\n - local: basic_tutorials/notebook\n title: Launching distributed training from Jupyter Notebooks\n title: Tutorials\n- sections:\n - isExpanded: true\n sections:\n - local: usage_guides/explore\n title: Start Here!\n - local: usage_guides/model_size_estimator\n title: Model memory estimator\n - local: usage_guides/quantization\n title: Model quantization\n - local: usage_guides/tracking\n title: Experiment trackers\n - local: usage_guides/profiler\n title: Profiler\n - local: usage_guides/checkpoint\n title: Checkpointing\n - local: basic_tutorials/troubleshooting\n title: Troubleshoot\n - local: usage_guides/training_zoo\n title: Example Zoo\n title: Accelerate\n - isExpanded: true\n sections:\n - local: usage_guides/gradient_accumulation\n title: Gradient accumulation\n - local: usage_guides/local_sgd\n title: Local SGD\n - local: usage_guides/low_precision_training\n title: Low precision (FP8) training\n - local: usage_guides/deepspeed\n title: DeepSpeed\n - local: usage_guides/deepspeed_multiple_model\n title: Using multiple models with DeepSpeed\n - local: usage_guides/ddp_comm_hook\n title: DDP Communication Hooks\n - local: usage_guides/fsdp\n title: Fully Sharded Data Parallel\n - local: usage_guides/megatron_lm\n title: Megatron-LM\n - local: usage_guides/sagemaker\n title: Amazon SageMaker\n - local: usage_guides/mps\n title: Apple M1 GPUs\n - local: usage_guides/intel_cpu\n title: Intel CPU\n - local: usage_guides/gaudi\n title: Intel Gaudi\n - local: usage_guides/compilation\n title: Compilation\n title: Training\n - isExpanded: true\n sections:\n - local: usage_guides/big_modeling\n title: Big Model Inference\n - local: usage_guides/distributed_inference\n title: Distributed inference\n title: Inference\n title: How to guides\n- sections:\n - local: concept_guides/internal_mechanism\n title: Accelerate's internal mechanism\n - local: concept_guides/big_model_inference\n title: Loading big models into memory\n - local: concept_guides/performance\n title: Comparing performance across distributed setups\n - local: concept_guides/deferring_execution\n title: Executing and deferring jobs\n - local: concept_guides/gradient_synchronization\n title: Gradient synchronization\n - local: concept_guides/fsdp_and_deepspeed\n title: FSDP vs DeepSpeed\n - local: concept_guides/fsdp1_vs_fsdp2\n title: FSDP1 vs FSDP2\n - local: concept_guides/context_parallelism\n title: Context parallelism\n - local: concept_guides/sequence_parallelism\n title: Sequence parallelism\n - local: concept_guides/low_precision_training\n title: Low precision training methods\n - local: concept_guides/training_tpu\n title: Training on TPUs\n title: Concepts and fundamentals\n- sections:\n - local: package_reference/accelerator\n title: Accelerator\n - local: package_reference/state\n title: Stateful classes\n - local: package_reference/cli\n title: The Command Line\n - local: package_reference/torch_wrappers\n title: DataLoaders, Optimizers, Schedulers\n - local: package_reference/tracking\n title: Experiment trackers\n - local: package_reference/launchers\n title: Launchers\n - local: package_reference/deepspeed\n title: DeepSpeed utilities\n - local: package_reference/logging\n title: Logging\n - local: package_reference/big_modeling\n title: Working with large models\n - local: package_reference/inference\n title: Pipeline parallelism\n - local: package_reference/kwargs\n title: Kwargs handlers\n - local: package_reference/fp8\n title: FP8\n - local: package_reference/utilities\n title: Utility functions and classes\n - local: package_reference/megatron_lm\n title: Megatron-LM utilities\n - local: package_reference/fsdp\n title: Fully Sharded Data Parallel utilities\n title: \"Reference\"\n" }, { "path": "docs/source/basic_tutorials/execution.md", "content": "\n\n# Execution process\n\nWhen working with distributed training systems, it is important to manage how and when processes are executed across GPUs. Some processes are completed faster than others, and some processes shouldn't begin if others haven't finished yet. Accelerate provides tools for orchestrating when processes are executed to ensure everything remains synchronized across all devices.\n\nThis tutorial will teach you how to execute a process on only one machine and how to delay execution until all processes have reached a certain point.\n\n## Execute on one process\n\nCertain code only needs to be run once on a given machine, such as printing a log statement or only displaying one progress bar on the local main process.\n\n\n\n\nYou should use `accelerator.is_local_main_process` to indicate code that should only be executed once.\n\n```py\nfrom tqdm.auto import tqdm\n\nprogress_bar = tqdm(range(args.max_train_steps), disable=not accelerator.is_local_main_process)\n```\n\nYou could also wrap a statement with `accelerator.is_local_main_process`.\n\n> [!TIP]\n> For standalone `print` statements that aren't wrapped in `accelerator.is_local_main_process`, replace `print` with Accelerate's [`~Accelerator.print`] method to only print once per process.\n\n```py\nif accelerator.is_local_main_process:\n print(\"Accelerate is the best\")\n```\n\n\n\n\nFor a function that should only be executed once, use [`~Accelerator.on_local_main_process`].\n\n```py\n@accelerator.on_local_main_process\ndef do_my_thing():\n \"Something done once per server\"\n do_thing_once_per_server()\n```\n\n\n\n\nYou could also direct Accelerate to execute code once across *all processes* regardless of the number of machines. This is useful if you're uploading a final model to the Hub.\n\n\n\n\nYou should use `accelerator.is_main_process` to indicate code that should only be executed once across all processes.\n\n```py\nif accelerator.is_main_process:\n repo.push_to_hub()\n```\n\n\n\n\nFor a function that should only be executed once across all processes, use [`~Accelerator.on_main_process`].\n\n```py\n@accelerator.on_main_process\ndef do_my_thing():\n \"Something done once per server\"\n do_thing_once()\n```\n\n\n\n\n## Execute on a specific process\n\nAccelerate can also help you execute functions that should only be executed on a specific process or a local process index.\n\n\n\n\nUse the [`~Accelerator.on_process`] method and specify the process index to execute a function on.\n\n```py\n@accelerator.on_process(process_index=0)\ndef do_my_thing():\n \"Something done on process index 0\"\n do_thing_on_index_zero()\n```\n\n\n\n\nUse the [`~Accelerator.on_local_process`] method and specify the local process index to execute a function on.\n\n```py\n@accelerator.on_local_process(local_process_idx=0)\ndef do_my_thing():\n \"Something done on process index 0 on each server\"\n do_thing_on_index_zero_on_each_server()\n```\n\n\n\n\n## Defer execution\n\nWhen you run your script on several GPUs at the same time, some code may be executed faster than others. You might need to wait for all processes to reach a certain point before executing the next set of instructions. For instance, you shouldnā€™t save a model before making sure every process is done with training.\n\nTo do this, add [`~Accelerator.wait_for_everyone`] in your code. This blocks all processes that have finished first from continuing until all remaining processes have reached the same point (this has no effect if you're running on a single GPU or CPU).\n\n```py\naccelerator.wait_for_everyone()\n```\n" }, { "path": "docs/source/basic_tutorials/install.md", "content": "\n\n# Installation\n\nBefore you start, you will need to setup your environment, install the appropriate packages, and configure Accelerate. Accelerate is tested on **Python 3.8+**.\n\nAccelerate is available on pypi and conda, as well as on GitHub. Details to install from each are below:\n\n## pip\n\nTo install Accelerate from pypi, perform:\n\n```bash\npip install accelerate\n```\n\n## conda\n\nAccelerate can also be installed with conda with:\n\n```bash\nconda install -c conda-forge accelerate\n```\n\n## Source\n\nNew features are added every day that haven't been released yet. To try them out yourself, install\nfrom the GitHub repository:\n\n```bash\npip install git+https://github.com/huggingface/accelerate\n```\n\nIf you're working on contributing to the library or wish to play with the source code and see live \nresults as you run the code, an editable version can be installed from a locally-cloned version of the \nrepository:\n\n```bash\ngit clone https://github.com/huggingface/accelerate\ncd accelerate\npip install -e .\n```\n\n## Configuration\n\nAfter installing, you need to configure Accelerate for how the current system is set up for training. \nTo do so run the following and answer the questions prompted to you:\n\n```bash\naccelerate config\n```\n\nTo write a barebones configuration that doesn't include options such as DeepSpeed configuration or running on TPUs, you can quickly run:\n\n```bash\npython -c \"from accelerate.utils import write_basic_config; write_basic_config(mixed_precision='fp16')\"\n```\n\nAccelerate will automatically utilize the maximum number of GPUs available and set the mixed precision mode.\n\nTo check that your configuration looks fine, run:\n\n```bash\naccelerate env\n```\n\nAn example output is shown below, which describes two GPUs on a single machine with no mixed precision being used:\n\n\n```bash\n- `Accelerate` version: 1.2.0.dev0\n- Platform: Linux-6.8.0-47-generic-x86_64-with-glibc2.35\n- `accelerate` bash location: /home/zach/miniconda3/envs/accelerate/bin/accelerate\n- Python version: 3.10.13\n- Numpy version: 1.26.4\n- PyTorch version (GPU?): 2.5.1+cu124 (True)\n- PyTorch XPU available: False\n- PyTorch NPU available: False\n- PyTorch MLU available: False\n- PyTorch MUSA available: False\n- System RAM: 187.91 GB\n- GPU type: NVIDIA GeForce RTX 4090\n- `Accelerate` default config:\n - compute_environment: LOCAL_MACHINE\n - distributed_type: MULTI_GPU\n - mixed_precision: no\n - use_cpu: False\n - debug: False\n - num_processes: 2\n - machine_rank: 0\n - num_machines: 1\n - gpu_ids: all\n - rdzv_backend: static\n - same_network: True\n - main_training_function: main\n - enable_cpu_affinity: False\n - downcast_bf16: no\n - tpu_use_cluster: False\n - tpu_use_sudo: False\n - tpu_env: []\n```\n" }, { "path": "docs/source/basic_tutorials/launch.md", "content": "\n\n# Launching Accelerate scripts\n\nIn the previous tutorial, you were introduced to how to modify your current training script to use Accelerate.\nThe final version of that code is shown below:\n\n```python\nfrom accelerate import Accelerator\n\naccelerator = Accelerator()\n\nmodel, optimizer, training_dataloader, scheduler = accelerator.prepare(\n model, optimizer, training_dataloader, scheduler\n)\n\nfor batch in training_dataloader:\n optimizer.zero_grad()\n inputs, targets = batch\n outputs = model(inputs)\n loss = loss_function(outputs, targets)\n accelerator.backward(loss)\n optimizer.step()\n scheduler.step()\n```\n\nBut how do you run this code and have it utilize the special hardware available to it?\n\nFirst, you should rewrite the above code into a function, and make it callable as a script. For example:\n\n```diff\n from accelerate import Accelerator\n \n+ def main():\n accelerator = Accelerator()\n\n model, optimizer, training_dataloader, scheduler = accelerator.prepare(\n model, optimizer, training_dataloader, scheduler\n )\n\n for batch in training_dataloader:\n optimizer.zero_grad()\n inputs, targets = batch\n outputs = model(inputs)\n loss = loss_function(outputs, targets)\n accelerator.backward(loss)\n optimizer.step()\n scheduler.step()\n\n+ if __name__ == \"__main__\":\n+ main()\n```\n\nNext, you need to launch it with `accelerate launch`. \n\n\n\n It's recommended you run `accelerate config` before using `accelerate launch` to configure your environment to your liking. \n Otherwise Accelerate will use very basic defaults depending on your system setup.\n\n\n\n\n## Using accelerate launch\n\nAccelerate has a special CLI command to help you launch your code in your system through `accelerate launch`.\nThis 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.\n\n\n\n If you are familiar with launching scripts in PyTorch yourself such as with `torchrun`, you can still do this. It is not required to use `accelerate launch`.\n\n\n\nYou can launch your script quickly by using:\n\n```bash\naccelerate launch {script_name.py} --arg1 --arg2 ...\n```\n\nJust put `accelerate launch` at the start of your command, and pass in additional arguments and parameters to your script afterward like normal!\n\nSince this runs the various torch spawn methods, all of the expected environment variables can be modified here as well.\nFor example, here is how to use `accelerate launch` with a single GPU:\n\n```bash\n# for cuda device:\nCUDA_VISIBLE_DEVICES=\"0\" accelerate launch {script_name.py} --arg1 --arg2 ...\n# for xpu device:\nZE_AFFINITY_MASK=\"0\" accelerate launch {script_name.py} --arg1 --arg2 ...\n```\n\nYou can also use `accelerate launch` without performing `accelerate config` first, but you may need to manually pass in the right configuration parameters.\nIn this case, Accelerate will make some hyperparameter decisions for you, e.g., if GPUs are available, it will use all of them by default without the mixed precision.\nHere is how you would use all GPUs and train with mixed precision disabled:\n\n```bash\naccelerate launch --multi_gpu {script_name.py} {--arg1} {--arg2} ...\n```\n\nOr by specifying a number of GPUs to use:\n\n```bash\naccelerate launch --num_processes=2 {script_name.py} {--arg1} {--arg2} ...\n```\n\nTo get more specific you should pass in the needed parameters yourself. For instance, here is how you \nwould also launch that same script on two GPUs using mixed precision while avoiding all of the warnings: \n\n```bash\naccelerate launch --multi_gpu --mixed_precision=fp16 --num_processes=2 {script_name.py} {--arg1} {--arg2} ...\n```\n\nFor a complete list of parameters you can pass in, run:\n\n```bash\naccelerate launch -h\n```\n\n\n\n Even if you are not using Accelerate in your code, you can still use the launcher for starting your scripts!\n\n\n\nFor a visualization of this difference, that earlier `accelerate launch` on multi-gpu would look something like so with `torchrun`:\n\n```bash\nMIXED_PRECISION=\"fp16\" torchrun --nproc_per_node=2 --nnodes=1 {script_name.py} {--arg1} {--arg2} ...\n```\n\nYou can also launch your script utilizing the launch CLI as a python module itself, enabling the ability to pass in other python-specific\nlaunching behaviors. To do so, use `accelerate.commands.launch` instead of `accelerate launch`:\n\n```bash\npython -m accelerate.commands.launch --num_processes=2 {script_name.py} {--arg1} {--arg2}\n```\n\nIf you want to execute the script with any other python flags, you can pass them in as well similar to `-m`, such as \nthe below example enabling unbuffered stdout and stderr:\n\n```bash\npython -u -m accelerate.commands.launch --num_processes=2 {script_name.py} {--arg1} {--arg2}\n```\n\n\n\n You can run your code on CPU as well! This is helpful for debugging and testing purposes on toy models and datasets. \n\n```bash\naccelerate launch --cpu {script_name.py} {--arg1} {--arg2}\n``` \n\n\n\n## Why you should always use `accelerate config`\n\nWhy is it useful to the point you should **always** run `accelerate config`? \n\nRemember that earlier call to `accelerate launch` as well as `torchrun`?\nPost configuration, to run that script with the needed parts you just need to use `accelerate launch` outright, without passing anything else in:\n\n```bash\naccelerate launch {script_name.py} {--arg1} {--arg2} ...\n```\n\n\n## Custom Configurations\n\nAs briefly mentioned earlier, `accelerate launch` should be mostly used through combining set configurations \nmade with the `accelerate config` command. These configs are saved to a `default_config.yaml` file in your cache folder for Accelerate. \nThis cache folder is located at (with decreasing order of priority):\n\n- The content of your environment variable `HF_HOME` suffixed with `accelerate`.\n- If it does not exist, the content of your environment variable `XDG_CACHE_HOME` suffixed with\n `huggingface/accelerate`.\n- If this does not exist either, the folder `~/.cache/huggingface/accelerate`.\n\nTo have multiple configurations, the flag `--config_file` can be passed to the `accelerate launch` command paired \nwith the location of the custom yaml. \n\nAn example yaml may look something like the following for two GPUs on a single machine using `fp16` for mixed precision:\n```yaml\ncompute_environment: LOCAL_MACHINE\ndeepspeed_config: {}\ndistributed_type: MULTI_GPU\nfsdp_config: {}\nmachine_rank: 0\nmain_process_ip: null\nmain_process_port: null\nmain_training_function: main\nmixed_precision: fp16\nnum_machines: 1\nnum_processes: 2\nuse_cpu: false\n```\n\nLaunching a script from the location of that custom yaml file looks like the following:\n```bash\naccelerate launch --config_file {path/to/config/my_config_file.yaml} {script_name.py} {--arg1} {--arg2} ...\n```\n\n## Multi-node training\nMulti-node training with Accelerate is similar to [multi-node training with torchrun](https://pytorch.org/tutorials/intermediate/ddp_series_multinode.html). The simplest way to launch a multi-node training run is to do the following:\n\n- Copy your codebase and data to all nodes. (or place them on a shared filesystem)\n- Setup your python packages on all nodes.\n- Run `accelerate config` on the main single node first. After specifying the number of nodes, you will be asked to specify the rank of each node (this will be 0 for the main/master node), along with the IP address and port for the main process. This is required for the worker nodes to communicate with the main process. Afterwards, you can copy or send this config file across all of your nodes, changing the `machine_rank` to 1, 2,3, etc. to avoid having to run the command (or just follow their directions directly for launching with `torchrun` as well)\n\nOnce you have done this, you can start your multi-node training run by running `accelerate launch` (or `torchrun`) on all nodes.\n\n\n It is required that the command be run on all nodes for everything to start, not just running it from the main node. You can use something like SLURM or a different process executor to wrap around this requirement and call everything from a single command.\n\n\n\n\n It is recommended to use the intranet IP of your main node over the public IP for better latency. This is the `192.168.x.x` or the `172.x.x.x` address you see when you run `hostname -I` on the main node.\n\n\n\nTo get a better idea about multi-node training, check out our example for [multi-node training with FSDP](https://huggingface.co/blog/ram-efficient-pytorch-fsdp).\n" }, { "path": "docs/source/basic_tutorials/migration.md", "content": "\n\n# Add Accelerate to your code\n\nEach distributed training framework has its own way of doing things which can require writing a lot of custom code to adapt it to your PyTorch training code and training environment. Accelerate offers a friendly way to interface with these distributed training frameworks without having to learn the specific details of each one. Accelerate takes care of those details for you, so you can focus on the training code and scale it to any distributed training environment.\n\nIn this tutorial, you'll learn how to adapt your existing PyTorch code with Accelerate and get you on your way toward training on distributed systems with ease! You'll start with a basic PyTorch training loop (it assumes all the training objects like `model` and `optimizer` have been set up already) and progressively integrate Accelerate into it.\n\n```python\ndevice = \"cuda\"\nmodel.to(device)\n\nfor batch in training_dataloader:\n optimizer.zero_grad()\n inputs, targets = batch\n inputs = inputs.to(device)\n targets = targets.to(device)\n outputs = model(inputs)\n loss = loss_function(outputs, targets)\n loss.backward()\n optimizer.step()\n scheduler.step()\n```\n\n## Accelerator\n\nThe [`Accelerator`] is the main class for adapting your code to work with Accelerate. It knows about the distributed setup you're using such as the number of different processes and your hardware type. This class also provides access to many of the necessary methods for enabling your PyTorch code to work in any distributed training environment and for managing and executing processes across devices.\n\nThat's why you should always start by importing and creating an [`Accelerator`] instance in your script.\n\n```python\nfrom accelerate import Accelerator\n\naccelerator = Accelerator()\n```\n\nThe [`Accelerator`] also knows which device to move your PyTorch objects to, so it is recommended to let Accelerate handle this for you.\n\n```diff\n- device = \"cuda\"\n+ device = accelerator.device\n model.to(device)\n```\n\n## Prepare PyTorch objects\n\nNext, you need to prepare your PyTorch objects (model, optimizer, scheduler, etc.) for distributed training. The [`~Accelerator.prepare`] method takes care of placing your model in the appropriate container (like single GPU or multi-GPU) for your training setup, adapting the optimizer and scheduler to use Accelerate's [`~optimizer.AcceleratedOptimizer`] and [`~scheduler.AcceleratedScheduler`], and creating a new dataloader that can be sharded across processes.\n\n> [!TIP]\n> Accelerate only prepares objects that inherit from their respective PyTorch classes such as `torch.optim.Optimizer`.\n\nThe PyTorch objects are returned in the same order they're sent.\n\n```py\nmodel, optimizer, training_dataloader, scheduler = accelerator.prepare(\n model, optimizer, training_dataloader, scheduler\n)\n```\n\n## Training loop\n\nFinally, remove the `to(device)` calls to the inputs and targets in the training loop because Accelerate's DataLoader classes automatically places them on the right device. You should also replace the usual `backward()` pass with Accelerate's [`~Accelerator.backward`] method which scales the gradients for you and uses the appropriate `backward()` method depending on your distributed setup (for example, DeepSpeed or Megatron).\n\n```diff\n- inputs = inputs.to(device)\n- targets = targets.to(device)\n outputs = model(inputs)\n loss = loss_function(outputs, targets)\n- loss.backward()\n+ accelerator.backward(loss)\n```\n\nPut everything together and your new Accelerate training loop should now look like this!\n\n```python\nfrom accelerate import Accelerator\naccelerator = Accelerator()\n\ndevice = accelerator.device\nmodel, optimizer, training_dataloader, scheduler = accelerator.prepare(\n model, optimizer, training_dataloader, scheduler\n)\n\nfor batch in training_dataloader:\n optimizer.zero_grad()\n inputs, targets = batch\n outputs = model(inputs)\n loss = loss_function(outputs, targets)\n accelerator.backward(loss)\n optimizer.step()\n scheduler.step()\n```\n\n## Training features\n\nAccelerate offers additional features - like gradient accumulation, gradient clipping, mixed precision training and more - you can add to your script to improve your training run. Let's explore these three features.\n\n### Gradient accumulation\n\nGradient accumulation enables you to train on larger batch sizes by accumulating the gradients over multiple batches before updating the weights. This can be useful for getting around memory limitations. To enable this feature in Accelerate, specify the `gradient_accumulation_steps` parameter in the [`Accelerator`] class and add the [`~Accelerator.accumulate`] context manager to your script.\n\n```diff\n+ accelerator = Accelerator(gradient_accumulation_steps=2)\n model, optimizer, training_dataloader = accelerator.prepare(model, optimizer, training_dataloader)\n\n for input, label in training_dataloader:\n+ with accelerator.accumulate(model):\n predictions = model(input)\n loss = loss_function(predictions, label)\n accelerator.backward(loss)\n optimizer.step()\n scheduler.step()\n optimizer.zero_grad()\n```\n\n### Gradient clipping\n\nGradient clipping is a technique to prevent \"exploding gradients\", and Accelerate offers:\n\n* [`~Accelerator.clip_grad_value_`] to clip gradients to a minimum and maximum value\n* [`~Accelerator.clip_grad_norm_`] for normalizing gradients to a certain value\n\n### Mixed precision\n\nMixed precision accelerates training by using a lower precision data type like fp16 (half-precision) to calculate the gradients. For the best performance with Accelerate, the loss should be computed inside your model (like in Transformers models) because computations outside of the model are computed in full precision.\n\nSet the mixed precision type to use in the [`Accelerator`], and then use the [`~Accelerator.autocast`] context manager to automatically cast the values to the specified data type.\n\n> [!WARNING]\n> Accelerate enables automatic mixed precision, so [`~Accelerator.autocast`] is only needed if there are other mixed precision operations besides those performed on loss by [`~Accelerator.backward`] which already handles the scaling.\n\n```diff\n+ accelerator = Accelerator(mixed_precision=\"fp16\")\n+ with accelerator.autocast():\n loss = complex_loss_function(outputs, target)\n```\n\n## Save and load\n\nAccelerate can also save and load a *model* once training is complete or you can also save the model and optimizer *state* which could be useful for resuming training.\n\n### Model\n\nOnce all processes are complete, unwrap the model with the [`~Accelerator.unwrap_model`] method before saving it because the [`~Accelerator.prepare`] method wrapped your model into the proper interface for distributed training. If you don't unwrap the model, saving the model state dictionary also saves any potential extra layers from the larger model and you won't be able to load the weights back into your base model.\n\nYou should use the [`~Accelerator.save_model`] method to unwrap and save the model state dictionary. This method can also save a model into sharded checkpoints or into the [safetensors](https://hf.co/docs/safetensors/index) format.\n\n\n\n\n```py\naccelerator.wait_for_everyone()\naccelerator.save_model(model, save_directory)\n```\n\n\n\nFor models from the [Transformers](https://hf.co/docs/transformers/index) library, save the model with the [`~transformers.PreTrainedModel.save_pretrained`] method so that it can be reloaded with the [`~transformers.PreTrainedModel.from_pretrained`] method.\n\n```py\nfrom transformers import AutoModel\n\nunwrapped_model = accelerator.unwrap_model(model)\nunwrapped_model.save_pretrained(\n \"path/to/my_model_directory\",\n is_main_process=accelerator.is_main_process,\n save_function=accelerator.save,\n)\n\nmodel = AutoModel.from_pretrained(\"path/to/my_model_directory\")\n```\n\n\n\nTo load your weights, use the [`~Accelerator.unwrap_model`] method to unwrap the model first before loading the weights. All model parameters are references to tensors, so this loads your weights inside `model`.\n\n```py\nunwrapped_model = accelerator.unwrap_model(model)\npath_to_checkpoint = os.path.join(save_directory,\"pytorch_model.bin\")\nunwrapped_model.load_state_dict(torch.load(path_to_checkpoint))\n```\n\n\n\n\nSet `safe_serialization=True` to save the model in the safetensor format.\n\n```py\naccelerator.wait_for_everyone()\naccelerator.save_model(model, save_directory, max_shard_size=\"1GB\", safe_serialization=True)\n```\n\nTo load a sharded checkpoint or a safetensor formatted checkpoint, use the [`~accelerate.load_checkpoint_in_model`] method. This method allows you to load a checkpoint onto a specific device.\n\n```py\nload_checkpoint_in_model(unwrapped_model, save_directory, device_map={\"\":device})\n```\n\n\n\n\n### State\n\nDuring training, you may want to save the current state of the model, optimizer, random generators, and potentially learning rate schedulers so they can be restored in the *same script*. You should add the [`~Accelerator.save_state`] and [`~Accelerator.load_state`] methods to your script to save and load states.\n\nTo further customize where and how states are saved through [`~Accelerator.save_state`], use the [`~utils.ProjectConfiguration`] class. For example, if `automatic_checkpoint_naming` is enabled, each saved checkpoint is stored at `Accelerator.project_dir/checkpoints/checkpoint_{checkpoint_number}`.\n\nAny other stateful items to be stored should be registered with the [`~Accelerator.register_for_checkpointing`] method so they can be saved and loaded. Every object passed to this method to be stored must have a `load_state_dict` and `state_dict` function.\n\n> [!TIP]\n> If you have [`torchdata>=0.8.0`](https://github.com/pytorch/data/tree/main) installed, you can additionally pass `use_stateful_dataloader=True` into your [`~utils.DataLoaderConfiguration`]. This extends Accelerate's DataLoader classes with a `load_state_dict` and `state_dict` function, and makes it so `Accelerator.save_state` and `Accelerator.load_state` also track how far into the training dataset it has read when persisting the model.\n" }, { "path": "docs/source/basic_tutorials/notebook.md", "content": "\n\n# Launching distributed training from Jupyter Notebooks\n\nThis tutorial teaches you how to fine-tune a computer vision model with šŸ¤— Accelerate from a Jupyter Notebook on a distributed system.\nYou will also learn how to set up a few requirements needed for ensuring your environment is configured properly, your data has been prepared properly, and finally how to launch training.\n\n\n\n This tutorial is also available as a Jupyter Notebook [here](https://github.com/huggingface/notebooks/blob/main/examples/accelerate_examples/simple_cv_example.ipynb)\n\n\n\n## Configuring the Environment\n\nBefore any training can be performed, an Accelerate config file must exist in the system. Usually this can be done by running the following in a terminal and answering the prompts:\n\n```bash\naccelerate config\n```\n\nHowever, if general defaults are fine and you are *not* running on a TPU, Accelerate has a utility to quickly write your device configuration into a config file via [`utils.write_basic_config`].\n\nThe following code will restart Jupyter after writing the configuration, as CUDA runtime or XPU runtime was called to perform this. \n\n\n\n CUDA and XPU can't be initialized more than once on a multi-device system. It's fine to debug in the notebook and have calls to CUDA/XPU, but in order to finally train a full cleanup and restart will need to be performed.\n \n\n\n```python\nimport os\nfrom accelerate.utils import write_basic_config\n\nwrite_basic_config() # Write a config file\nos._exit(00) # Restart the notebook\n```\n\n## Preparing the Dataset and Model\n\nNext you should prepare your dataset. As mentioned earlier, great care should be taken when preparing the `DataLoaders` and model to make sure that **nothing** is put on *any* GPU. \n\nIf you do, it is recommended to put that specific code into a function and call that from within the notebook launcher interface, which will be shown later. \n\nMake sure the dataset is downloaded based on the directions [here](https://github.com/huggingface/accelerate/tree/main/examples#simple-vision-example)\n\n```python\nimport os, re, torch, PIL\nimport numpy as np\n\nfrom torch.optim.lr_scheduler import OneCycleLR\nfrom torch.utils.data import DataLoader, Dataset\nfrom torchvision.transforms import Compose, RandomResizedCrop, Resize, ToTensor\n\nfrom accelerate import Accelerator\nfrom accelerate.utils import set_seed\nfrom timm import create_model\n```\n\nFirst you need to create a function to extract the class name based on a filename:\n\n```python\nimport os\n\ndata_dir = \"../../images\"\nfnames = os.listdir(data_dir)\nfname = fnames[0]\nprint(fname)\n```\n\n```python out\nbeagle_32.jpg\n```\n\nIn the case here, the label is `beagle`. Using regex you can extract the label from the filename:\n\n```python\nimport re\n\n\ndef extract_label(fname):\n stem = fname.split(os.path.sep)[-1]\n return re.search(r\"^(.*)_\\d+\\.jpg$\", stem).groups()[0]\n```\n\n```python\nextract_label(fname)\n```\n\nAnd you can see it properly returned the right name for our file:\n\n```python out\n\"beagle\"\n```\n\nNext a `Dataset` class should be made to handle grabbing the image and the label:\n\n```python\nclass PetsDataset(Dataset):\n def __init__(self, file_names, image_transform=None, label_to_id=None):\n self.file_names = file_names\n self.image_transform = image_transform\n self.label_to_id = label_to_id\n\n def __len__(self):\n return len(self.file_names)\n\n def __getitem__(self, idx):\n fname = self.file_names[idx]\n raw_image = PIL.Image.open(fname)\n image = raw_image.convert(\"RGB\")\n if self.image_transform is not None:\n image = self.image_transform(image)\n label = extract_label(fname)\n if self.label_to_id is not None:\n label = self.label_to_id[label]\n return {\"image\": image, \"label\": label}\n```\n\nNow to build the dataset. Outside the training function you can find and declare all the filenames and labels and use them as references inside the \nlaunched function:\n\n```python\nfnames = [os.path.join(\"../../images\", fname) for fname in fnames if fname.endswith(\".jpg\")]\n```\n\nNext gather all the labels:\n\n```python\nall_labels = [extract_label(fname) for fname in fnames]\nid_to_label = list(set(all_labels))\nid_to_label.sort()\nlabel_to_id = {lbl: i for i, lbl in enumerate(id_to_label)}\n```\n\nNext, you should make a `get_dataloaders` function that will return your built dataloaders for you. As mentioned earlier, if data is automatically \nsent to the GPU or a TPU device when building your `DataLoaders`, they must be built using this method. \n\n```python\ndef get_dataloaders(batch_size: int = 64):\n \"Builds a set of dataloaders with a batch_size\"\n random_perm = np.random.permutation(len(fnames))\n cut = int(0.8 * len(fnames))\n train_split = random_perm[:cut]\n eval_split = random_perm[cut:]\n\n # For training a simple RandomResizedCrop will be used\n train_tfm = Compose([RandomResizedCrop((224, 224), scale=(0.5, 1.0)), ToTensor()])\n train_dataset = PetsDataset([fnames[i] for i in train_split], image_transform=train_tfm, label_to_id=label_to_id)\n\n # For evaluation a deterministic Resize will be used\n eval_tfm = Compose([Resize((224, 224)), ToTensor()])\n eval_dataset = PetsDataset([fnames[i] for i in eval_split], image_transform=eval_tfm, label_to_id=label_to_id)\n\n # Instantiate dataloaders\n train_dataloader = DataLoader(train_dataset, shuffle=True, batch_size=batch_size, num_workers=4)\n eval_dataloader = DataLoader(eval_dataset, shuffle=False, batch_size=batch_size * 2, num_workers=4)\n return train_dataloader, eval_dataloader\n```\n\nFinally, you should import the scheduler to be used later:\n\n```python\nfrom torch.optim.lr_scheduler import CosineAnnealingLR\n```\n\n## Writing the Training Function\n\nNow you can build the training loop. [`notebook_launcher`] works by passing in a function to call that will be ran across the distributed system.\n\nHere is a basic training loop for the animal classification problem:\n\n\n\n The code has been split up to allow for explanations on each section. A full version that can be copy and pasted will be available at the end\n\n\n\n\n```python\ndef training_loop(mixed_precision=\"fp16\", seed: int = 42, batch_size: int = 64):\n set_seed(seed)\n accelerator = Accelerator(mixed_precision=mixed_precision)\n```\n\nFirst you should set the seed and create an [`Accelerator`] object as early in the training loop as possible.\n\n\n\n If training on the TPU, your training loop should take in the model as a parameter and it should be instantiated \n outside of the training loop function. See the [TPU best practices](../concept_guides/training_tpu) \n to learn why\n\n\n\nNext you should build your dataloaders and create your model:\n\n```python\n train_dataloader, eval_dataloader = get_dataloaders(batch_size)\n model = create_model(\"resnet50d\", pretrained=True, num_classes=len(label_to_id))\n```\n\n\n\n You build the model here so that the seed also controls the new weight initialization\n\n\n\nAs you are performing transfer learning in this example, the encoder of the model starts out frozen so the head of the model can be \ntrained only initially:\n\n```python\n for param in model.parameters():\n param.requires_grad = False\n for param in model.get_classifier().parameters():\n param.requires_grad = True\n```\n\nNormalizing the batches of images will make training a little faster:\n\n```python\n mean = torch.tensor(model.default_cfg[\"mean\"])[None, :, None, None]\n std = torch.tensor(model.default_cfg[\"std\"])[None, :, None, None]\n```\n\nTo make these constants available on the active device, you should set it to the Accelerator's device:\n\n```python\n mean = mean.to(accelerator.device)\n std = std.to(accelerator.device)\n```\n\nNext instantiate the rest of the PyTorch classes used for training:\n\n```python\n optimizer = torch.optim.Adam(params=model.parameters(), lr=3e-2 / 25)\n lr_scheduler = OneCycleLR(optimizer=optimizer, max_lr=3e-2, epochs=5, steps_per_epoch=len(train_dataloader))\n```\n\nBefore passing everything to [`~Accelerator.prepare`].\n\n\n\n There is no specific order to remember, you just need to unpack the objects in the same order you gave them to the prepare method.\n\n\n\n```python\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler\n )\n```\n\nNow train the model:\n\n```python\n for epoch in range(5):\n model.train()\n for batch in train_dataloader:\n inputs = (batch[\"image\"] - mean) / std\n outputs = model(inputs)\n loss = torch.nn.functional.cross_entropy(outputs, batch[\"label\"])\n accelerator.backward(loss)\n optimizer.step()\n lr_scheduler.step()\n optimizer.zero_grad()\n```\n\nThe evaluation loop will look slightly different compared to the training loop. The number of elements passed as well as the overall \ntotal accuracy of each batch will be added to two constants:\n\n```python\n model.eval()\n accurate = 0\n num_elems = 0\n```\n\nNext you have the rest of your standard PyTorch loop:\n\n```python\n for batch in eval_dataloader:\n inputs = (batch[\"image\"] - mean) / std\n with torch.no_grad():\n outputs = model(inputs)\n predictions = outputs.argmax(dim=-1)\n```\n\nBefore finally the last major difference. \n\nWhen performing distributed evaluation, the predictions and labels need to be passed through \n[`~Accelerator.gather`] so that all of the data is available on the current device and a properly calculated metric can be achieved:\n\n```python\n accurate_preds = accelerator.gather(predictions) == accelerator.gather(batch[\"label\"])\n num_elems += accurate_preds.shape[0]\n accurate += accurate_preds.long().sum()\n```\n\nNow you just need to calculate the actual metric for this problem, and you can print it on the main process using [`~Accelerator.print`]:\n\n```python\n eval_metric = accurate.item() / num_elems\n accelerator.print(f\"epoch {epoch}: {100 * eval_metric:.2f}\")\n```\n\nA full version of this training loop is available below:\n\n```python\ndef training_loop(mixed_precision=\"fp16\", seed: int = 42, batch_size: int = 64):\n set_seed(seed)\n # Initialize accelerator\n accelerator = Accelerator(mixed_precision=mixed_precision)\n # Build dataloaders\n train_dataloader, eval_dataloader = get_dataloaders(batch_size)\n\n # Instantiate the model (you build the model here so that the seed also controls new weight initializations)\n model = create_model(\"resnet50d\", pretrained=True, num_classes=len(label_to_id))\n\n # Freeze the base model\n for param in model.parameters():\n param.requires_grad = False\n for param in model.get_classifier().parameters():\n param.requires_grad = True\n\n # You can normalize the batches of images to be a bit faster\n mean = torch.tensor(model.default_cfg[\"mean\"])[None, :, None, None]\n std = torch.tensor(model.default_cfg[\"std\"])[None, :, None, None]\n\n # To make these constants available on the active device, set it to the accelerator device\n mean = mean.to(accelerator.device)\n std = std.to(accelerator.device)\n\n # Instantiate the optimizer\n optimizer = torch.optim.Adam(params=model.parameters(), lr=3e-2 / 25)\n\n # Instantiate the learning rate scheduler\n lr_scheduler = OneCycleLR(optimizer=optimizer, max_lr=3e-2, epochs=5, steps_per_epoch=len(train_dataloader))\n\n # Prepare everything\n # There is no specific order to remember, you just need to unpack the objects in the same order you gave them to the\n # prepare method.\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler\n )\n\n # Now you train the model\n for epoch in range(5):\n model.train()\n for batch in train_dataloader:\n inputs = (batch[\"image\"] - mean) / std\n outputs = model(inputs)\n loss = torch.nn.functional.cross_entropy(outputs, batch[\"label\"])\n accelerator.backward(loss)\n optimizer.step()\n lr_scheduler.step()\n optimizer.zero_grad()\n\n model.eval()\n accurate = 0\n num_elems = 0\n for batch in eval_dataloader:\n inputs = (batch[\"image\"] - mean) / std\n with torch.no_grad():\n outputs = model(inputs)\n predictions = outputs.argmax(dim=-1)\n accurate_preds = accelerator.gather(predictions) == accelerator.gather(batch[\"label\"])\n num_elems += accurate_preds.shape[0]\n accurate += accurate_preds.long().sum()\n\n eval_metric = accurate.item() / num_elems\n # Use accelerator.print to print only on the main process.\n accelerator.print(f\"epoch {epoch}: {100 * eval_metric:.2f}\")\n```\n\n## Using the notebook_launcher\n\nAll that's left is to use the [`notebook_launcher`].\n\nYou pass in the function, the arguments (as a tuple), and the number of processes to train on. (See the [documentation](../package_reference/launchers) for more information)\n\n```python\nfrom accelerate import notebook_launcher\n```\n\n```python\nargs = (\"fp16\", 42, 64)\nnotebook_launcher(training_loop, args, num_processes=2)\n```\n\nIn 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.\n\nFor 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:\n\n```python\nnotebook_launcher(training_loop, args, master_addr=\"172.31.43.8\", node_rank=0, num_nodes=2, num_processes=8)\n```\n\nAnd in the second Jupyter session on the other machine:\n\n\n\n Notice how the `node_rank` has changed\n\n\n\n```python\nnotebook_launcher(training_loop, args, master_addr=\"172.31.43.8\", node_rank=1, num_nodes=2, num_processes=8)\n```\n\nIn the case of running on the TPU, it would look like so:\n\n```python\nmodel = create_model(\"resnet50d\", pretrained=True, num_classes=len(label_to_id))\n\nargs = (model, \"fp16\", 42, 64)\nnotebook_launcher(training_loop, args, num_processes=8)\n```\n\nTo launch the training process with elasticity, enabling fault tolerance, you can use the `elastic_launch` feature provided by PyTorch. This requires setting additional parameters such as `rdzv_backend` and `max_restarts`. Here is an example of how to use `notebook_launcher` with elastic capabilities:\n\n```python\nnotebook_launcher(\n training_loop,\n args,\n num_processes=2,\n max_restarts=3\n)\n```\n\nAs it's running it will print the progress as well as state how many devices you ran on. This tutorial was ran with two GPUs:\n\n```python out\nLaunching training on 2 GPUs.\nepoch 0: 88.12\nepoch 1: 91.73\nepoch 2: 92.58\nepoch 3: 93.90\nepoch 4: 94.71\n```\n\nAnd that's it!\n\nPlease note that [`notebook_launcher`] ignores the Accelerate config file, to launch based on the config use:\n\n```bash\naccelerate launch\n```\n\n## Debugging \n\nA common issue when running the `notebook_launcher` is receiving a CUDA/XPU has already been initialized issue. This usually stems\nfrom an import or prior code in the notebook that makes a call to the PyTorch `torch.cuda` or `torch.xpu` sublibrary. To help narrow down what went wrong,\nyou can launch the `notebook_launcher` with `ACCELERATE_DEBUG_MODE=yes` in your environment and an additional check\nwill be made when spawning that a regular process can be created and utilize CUDA/XPU without issue. (Your CUDA/XPU code can still be ran afterwards).\n\n## Conclusion\n\nThis notebook showed how to perform distributed training from inside of a Jupyter Notebook. Some key notes to remember:\n\n- Make sure to save any code that use CUDA/XPU (or CUDA/XPU imports) for the function passed to [`notebook_launcher`]\n- Set the `num_processes` to be the number of devices used for training (such as number of GPUs, XPUs, CPUs, TPUs, etc)\n- If using the TPU, declare your model outside the training loop function\n" }, { "path": "docs/source/basic_tutorials/overview.md", "content": "\n\n# Overview\n\nWelcome to the Accelerate tutorials! These introductory guides will help catch you up to speed on working with Accelerate.\nYou'll learn how to modify your code to have it work with the API seamlessly, how to launch your script properly,\nand more!\n\nThese tutorials assume some basic knowledge of Python and familiarity with the PyTorch framework.\n\nIf you have any questions about Accelerate, feel free to join and ask the community on our [forum](https://discuss.huggingface.co/c/accelerate/18)." }, { "path": "docs/source/basic_tutorials/tpu.md", "content": "\n\n# TPU training\n\nA [TPU (Tensor Processing Unit)](https://cloud.google.com/tpu/docs/intro-to-tpu) is a type of hardware specifically designed for training models efficiently. Accelerate supports TPU training, but there are a few things you should be aware of, namely graph compilation. This tutorial briefly discusses compilation, and for more details, take a look at the [Training on TPUs with Accelerate](../concept_guides/training_tpu) guide.\n\n## Compilation\n\nA TPU creates a graph of all the operations in the training step such as the forward pass, backward pass and optimizer step. This is why the first training step always takes a while because building and compiling this graph takes time. But once compilation is complete, it is cached and all subsequent steps are much faster.\n\nThe key is to avoid compiling your code again or else training is super slow. This means all your operations must be exactly the same:\n\n* all tensors in your batches must have the same length (for example, no dynamic padding for NLP tasks)\n* your code must be static (for example, no layers with for loops that have different lengths depending on the input such as a LSTM)\n\n## Weight tying\n\nA common language model design is to tie the weights of the embedding and softmax layers. However, moving the model to a TPU (either yourself or passing it to the [`~Accelerator.prepare`] method) breaks the weight tying and you'll need to retie the weights.\n\nTo add special behavior (like weight tying) in your script for TPUs, set [`~Accelerator.distributed_type`] to `DistributedType.TPU` first. Then you can use the [`~transformers.PreTrainedModel.tie_weights`] method to tie the weights.\n\n```py\nif accelerator.distributed_type == DistributedType.TPU:\n model.tie_weights()\n```\n" }, { "path": "docs/source/basic_tutorials/troubleshooting.md", "content": "\n\n# Troubleshoot\n\nThis guide provides solutions to some issues you might encounter when using Accelerate. Not all errors are covered because Accelerate is an active library that is continuously evolving and there are many different use cases and distributed training setups. If the solutions described here don't help with your specific error, please take a look at the [Ask for help](#ask-for-help) section to learn where and how to get help.\n\n## Logging\n\nLogging can help you identify where an error is coming from. In a distributed setup with multiple processes, logging can be a challenge, but Accelerate provides the [`~accelerate.logging`] utility to ensure logs are synchronized.\n\nTo troubleshoot an issue, use [`~accelerate.logging`] instead of the standard Python [`logging`](https://docs.python.org/3/library/logging.html#module-logging) module. Set the verbosity level (`INFO`, `DEBUG`, `WARNING`, `ERROR`, `CRITICAL`) with the `log_level` parameter, and then you can either:\n\n1. Export the `log_level` as the `ACCELERATE_LOG_LEVEL` environment variable.\n2. Pass the `log_level` directly to `get_logger`.\n\nFor example, to set `log_level=\"INFO\"`:\n\n```py\nfrom accelerate.logging import get_logger\n\nlogger = get_logger(__name__, log_level=\"DEBUG\")\n```\n\nBy default, the log is called on main processes only. To call it on all processes, pass `main_process_only=False`.\nIf a log should be called on all processes and in order, also pass `in_order=True`.\n\n```py\nfrom accelerate.logging import get_logger\n\nlogger = get_logger(__name__, log_level=\"DEBUG\")\n# log all processes\nlogger.debug(\"thing_to_log\", main_process_only=False)\n# log all processes in order\nlogger.debug(\"thing_to_log\", main_process_only=False, in_order=True)\n```\n\n## Hanging code and timeout errors\n\nThere can be many reasons why your code is hanging. Let's take a look at how to solve some of the most common issues that can cause your code to hang.\n\n### Mismatched tensor shapes\n\nMismatched tensor shapes is a common issue that can cause your code to hang for a significant amount of time on a distributed setup.\n\nWhen running scripts in a distributed setup, functions such as [`Accelerator.gather`] and [`Accelerator.reduce`] are necessary to grab tensors across devices to collectively perform operations on them. These (and other) functions rely on `torch.distributed` to perform a `gather` operation, which requires tensors to have the **exact same shape** across all processes. When the tensor shapes don't match, your code hangs and you'll eventually hit a timeout exception.\n\nYou can use Accelerate's operational debug mode to immediately catch this issue. We recommend enabling this mode during the `accelerate config` setup, but you can also enable it from the CLI, as an environment variable, or by manually editing the `config.yaml` file.\n\n\n\n\n```bash\naccelerate launch --debug {my_script.py} --arg1 --arg2\n```\n\n\n\n\nIf enabling debug mode as an environment variable, you don't need to call `accelerate launch`.\n\n```bash\nACCELERATE_DEBUG_MODE=\"1\" torchrun {my_script.py} --arg1 --arg2\n```\n\n\n\n\nAdd `debug: true` to your `config.yaml` file.\n\n```yaml\ncompute_environment: LOCAL_MACHINE\ndebug: true\n```\n\n\n\n\nOnce you enable debug mode, you should get a traceback that points to the tensor shape mismatch issue.\n\n```py\nTraceback (most recent call last):\n File \"/home/zach_mueller_huggingface_co/test.py\", line 18, in \n main()\n File \"/home/zach_mueller_huggingface_co/test.py\", line 15, in main\n broadcast_tensor = broadcast(tensor)\n File \"/home/zach_mueller_huggingface_co/accelerate/src/accelerate/utils/operations.py\", line 303, in wrapper\naccelerate.utils.operations.DistributedOperationException:\n\nCannot apply desired operation due to shape mismatches. All shapes across devices must be valid.\n\nOperation: `accelerate.utils.operations.broadcast`\nInput shapes:\n - Process 0: [1, 5]\n - Process 1: [1, 2, 5]\n```\n\n### Early stopping\n\nFor early stopping in distributed training, if each process has a specific stopping condition (e.g. validation loss), it may not be synchronized across all processes. As a result, a break can happen on process 0 but not on process 1 which will cause your code to hang indefinitely until a timeout occurs.\n\nIf you have early stopping conditionals, use the `set_trigger` and `check_trigger` methods to make sure all the processes\nare ended correctly.\n\n```py\n# Assume `should_do_breakpoint` is a custom-defined function that returns a conditional, \n# and that conditional might be true only on process 1\nif should_do_breakpoint(loss):\n accelerator.set_trigger()\n\n# Later in the training script when we need to check for the breakpoint\nif accelerator.check_trigger():\n break\n```\n\n### Low kernel versions on Linux\n\nOn Linux with kernel version < 5.5, hanging processes have been reported. To avoid this problem, upgrade your system to a later kernel version.\n\n### MPI\n\nIf your distributed CPU training job using MPI is hanging, ensure that you have\n[passwordless SSH](https://www.open-mpi.org/faq/?category=rsh#ssh-keys) setup (using keys) between the nodes. This means\nthat for all nodes in your hostfile, you should be able to SSH from one node to another without being prompted for a password.\n\nNext, try to run the `mpirun` command as a sanity check. For example, the command below should print out the\nhostnames for each of the nodes.\n\n```bash\nmpirun -f hostfile -n {number of nodes} -ppn 1 hostname\n```\n\n## Out-of-Memory\n\nOne of the most frustrating errors when it comes to running training scripts is hitting \"Out-of-Memory\" on devices like CUDA, XPU or CPU. The entire script needs to be restarted and any progress is lost.\n\nTo address this problem, Accelerate provides the [`find_executable_batch_size`] utility that is heavily based on [toma](https://github.com/BlackHC/toma).\nThis utility retries code that fails due to OOM (out-of-memory) conditions and automatically lowers batch sizes. For each OOM condition, the algorithm decreases the batch size by half and retries the code until it succeeds.\n\nTo use [`find_executable_batch_size`], restructure your training function to include an inner function with `find_executable_batch_size` and build your dataloaders inside it. At a minimum, this only takes 4 new lines of code.\n\n \n\nThe inner function **must** take batch size as the first parameter, but we do not pass one to it when called. The wrapper will handle this for you. Any object (models, optimizers) that consumes device memory and is passed to the [`Accelerator`] also **must** be declared inside the inner function.\n\n\n\n```diff\ndef training_function(args):\n accelerator = Accelerator()\n\n+ @find_executable_batch_size(starting_batch_size=args.batch_size)\n+ def inner_training_loop(batch_size):\n+ nonlocal accelerator # Ensure they can be used in our context\n+ accelerator.free_memory() # Free all lingering references\n model = get_model()\n model.to(accelerator.device)\n optimizer = get_optimizer()\n train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size)\n lr_scheduler = get_scheduler(\n optimizer, \n num_training_steps=len(train_dataloader)*num_epochs\n )\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler\n )\n train(model, optimizer, train_dataloader, lr_scheduler)\n validate(model, eval_dataloader)\n+ inner_training_loop()\n```\n\n## Non-reproducible results between device setups\n\nIf you changed the device setup and observe different model performance, it is likely you didn't update your script when moving from one setup to another. Even if you're using the same script with the same batch size, the results will still be different on a TPU, multi-GPU, and single GPU.\n\nFor example, if you were training on a single GPU with a batch size of 16 and you move to a dual GPU setup, you need to change the batch size to 8 to have the same effective batch size. This is because when training with Accelerate, the batch size passed to the dataloader is the **batch size per GPU**.\n\nTo make sure you can reproduce the results between the setups, make sure to use the same seed, adjust the batch size accordingly, and consider scaling the learning rate.\n\nFor more details and a quick reference for batch sizes, check out the [Comparing performance between different device setups](../concept_guides/performance) guide.\n\n## Performance issues on different GPUs\n\nIf your multi-GPU setup consists of different GPUs, you may encounter some performance issues:\n\n- There may be an imbalance in GPU memory between the GPUs. In this case, the GPU with the smaller memory will limit the batch size or the size of the model that can be loaded onto the GPUs.\n- If you are using GPUs with different performance profiles, the performance will be driven by the slowest GPU you are using because the other GPUs will have to wait for it to complete its workload.\n\nVastly different GPUs within the same setup can lead to performance bottlenecks.\n\n## Ask for help\n\nIf none of the solutions and advice here helped resolve your issue, you can always reach out to the community and Accelerate team for help.\n\n- Ask for help on the Hugging Face forums by posting your question in the [Accelerate category](https://discuss.huggingface.co/c/accelerate/18). Make sure to write a descriptive post with relevant context about your setup and reproducible code to maximize the likelihood that your problem is solved!\n\n- Post a question on [Discord](http://hf.co/join/discord), and let the team and the community help you.\n\n- Create an Issue on the Accelerate [GitHub repository](https://github.com/huggingface/accelerate/issues) if you think you've found a bug related to the library. Include context regarding the bug and details about your distributed setup to help us better figure out what's wrong and how we can fix it.\n" }, { "path": "docs/source/concept_guides/big_model_inference.md", "content": "\n\n# Loading big models into memory\n\nWhen loading a pre-trained model in PyTorch, the usual workflow looks like this:\n\n```py\nimport torch\n\nmy_model = ModelClass(...)\nstate_dict = torch.load(checkpoint_file)\nmy_model.load_state_dict(state_dict)\n```\n\nIn plain English, those steps are:\n1. Create the model with randomly initialized weights\n2. Load the model weights (in a dictionary usually called a state dict) from the disk\n3. Load those weights inside the model\n\nWhile 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).\n\n\n\nThis 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.\n\n\n\n## How the Process Works: A Quick Overview\n\n\n\n## How the Process Works: Working with Code\n\n### Instantiating an empty model\n\nThe 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:\n\n```py\nfrom accelerate import init_empty_weights\n\nwith init_empty_weights():\n my_model = ModelClass(...)\n```\n\nFor instance:\n\n```py\nwith init_empty_weights():\n model = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)])\n```\n\ninitializes 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.\n\n\n\n You can't move a model initialized like this on CPU or another device directly, since it doesn't have any data. It's also very likely that a forward pass with that empty model will fail, as not all operations are supported on the meta device.\n\n\n\n### Sharded checkpoints\n\nIt'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.\n\nAccelerate 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:\n\n```bash\nfirst_state_dict.bin\nindex.json\nsecond_state_dict.bin\n```\n\nwith index.json being the following file:\n\n```\n{\n \"linear1.weight\": \"first_state_dict.bin\",\n \"linear1.bias\": \"first_state_dict.bin\",\n \"linear2.weight\": \"second_state_dict.bin\",\n \"linear2.bias\": \"second_state_dict.bin\"\n}\n```\n\nand `first_state_dict.bin` containing the weights for `\"linear1.weight\"` and `\"linear1.bias\"`, `second_state_dict.bin` the ones for `\"linear2.weight\"` and `\"linear2.bias\"`\n\n### Loading weights\n\nThe 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.\n\nIf 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).\n\nHere is how we can use this to load the [GPT2-1.5B](https://huggingface.co/marcsun13/gpt2-xl-linear-sharded) model.\n\nLet's download the sharded version of this model.\n\n```bash\npip install huggingface_hub\n```\n\n```py\nfrom huggingface_hub import snapshot_download\ncheckpoint = \"marcsun13/gpt2-xl-linear-sharded\"\nweights_location = snapshot_download(repo_id=checkpoint)\n```\n\nIn order to initialize the model, we will use the library minGPT. \n\n```bash\ngit clone https://github.com/karpathy/minGPT.git\npip install minGPT/\n```\n\n```py\nfrom accelerate import init_empty_weights\nfrom mingpt.model import GPT\n\nmodel_config = GPT.get_default_config()\nmodel_config.model_type = 'gpt2-xl'\nmodel_config.vocab_size = 50257\nmodel_config.block_size = 1024\n\nwith init_empty_weights():\n model = GPT(model_config)\n```\n\nThen, load the checkpoint we just downloaded with:\n\n```py\nfrom accelerate import load_checkpoint_and_dispatch\n\nmodel = load_checkpoint_and_dispatch(\n model, checkpoint=weights_location, device_map=\"auto\", no_split_module_classes=['Block']\n)\n```\n\nBy passing `device_map=\"auto\"`, we tell Accelerate to determine automatically where to put each layer of the model depending on the available resources:\n- first, we use the maximum space available on the GPU(s)\n- if we still need space, we store the remaining weights on the CPU\n- if there is not enough RAM, we store the remaining weights on the hard drive as memory-mapped tensors\n\n\n#### `no_split_module_classes`\n\nThis 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 \ninclude a residual connection of some kind.\n\n\n#### The `device_map`\n\nYou can see the `device_map` that Accelerate picked by accessing the `hf_device_map` attribute of your model:\n\n```py\nmodel.hf_device_map\n```\n\n```python out\n{'transformer.wte': 0,\n 'transformer.wpe': 0,\n 'transformer.drop': 0,\n 'transformer.h.0': 0,\n ...\n 'transformer.h.21': 0, \n 'transformer.h.22': 1, \n 'transformer.h.23': 1, \n 'transformer.h.24': 1,\n ...\n 'transformer.h.47': 1, \n 'transformer.ln_f': 1, \n 'lm_head': 1}\n ```\n\nIt'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:\n\n```python\ndevice_map = {\n \"transformer.wte\": \"cpu\",\n \"transformer.wpe\": 0,\n \"transformer.drop\": \"cpu\",\n \"transformer.h.0\": \"disk\"\n}\n\nmodel = load_checkpoint_and_dispatch(\n model, checkpoint=weights_location, device_map=device_map\n)\n\n```\n\n### Run the model\n\nNow 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:\n\n```py\nfrom mingpt.bpe import BPETokenizer\ntokenizer = BPETokenizer()\ninputs = tokenizer(\"Hello, my name is\").to(0)\n\noutputs = model.generate(x1, max_new_tokens=10, do_sample=False)[0]\ntokenizer.decode(outputs.cpu().squeeze())\n```\n\nBehind the scenes, Accelerate added hooks to the model, so that:\n- at each layer, the inputs are put on the right device (so even if your model is spread across several GPUs, it works)\n- for the weights offloaded on the CPU, they are put on a GPU just before the forward pass and cleaned up just after\n- 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\n\nThis way, your model can run for inference even if it doesn't fit on one of the GPUs or the CPU RAM!\n\n\n\n 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.\n\n\n\n### Designing a device map\n\nYou 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.\n\n\n\n You can derive all sizes of the model (and thus compute a `device_map`) on a model that is on the meta device.\n\n\n\nAll 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). \n\nWhen you have more GPU memory available than the model size, here is the difference between each option:\n- `\"auto\"` and `\"balanced\"` evenly split the model on all available GPUs, making it possible for you to use a batch size greater than 1.\n- `\"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\n- `\"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).\n\n\n\n The options `\"auto\"` and `\"balanced\"` produce the same results for now, but the behavior of `\"auto\"` might change in the future if we find a strategy that makes more sense, while `\"balanced\"` will stay stable.\n\n\n\nFirst 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\"`.\n\nHere 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:\n\n```python\nfrom accelerate import infer_auto_device_map\n\ndevice_map = infer_auto_device_map(my_model, max_memory={0: \"10GiB\", 1: \"10GiB\", \"cpu\": \"30GiB\"})\n```\n\n\n\n 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.\n\n Therefore when you create memory maps with `max_memory` make sure to adjust the available memory accordingly to avoid out-of-memory errors.\n\n\n\nAdditionally, 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:\n\n```python\nmax_memory = {0: \"30GIB\", 1: \"46GIB\", 2: \"46GIB\", 3: \"46GIB\", 4: \"46GIB\", 5: \"46GIB\", 6: \"46GIB\", 7: \"46GIB\"}\n```\nas you can see we gave the remaining 7 GPUs ~50% more memory than GPU 0.\n\nIf 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:\n\n```python\ndevice_map = {\"block1\": 0, \"block2\": 1}\n```\n\nanother one that is valid could be:\n\n```python\ndevice_map = {\"block1\": 0, \"block2.linear1\": 0, \"block2.linear2\": 1, \"block2.linear3\": 1}\n```\n\nOn the other hand, this one is not valid as it does not cover every parameter of the model:\n\n```python\ndevice_map = {\"block1\": 0, \"block2.linear1\": 1, \"block2.linear2\": 1}\n```\n\n\n\n To be the most efficient, make sure your device map puts the parameters on the GPUs in a sequential manner (e.g. don't put one of the first weights on GPU 0, then weights on GPU 1 and the last weight back to GPU 0) to avoid making many transfers of data between the GPUs.\n\n\n\n## CPU offload only\n\nIf you want to offload your model on CPU, you can use [`cpu_offload`]. As a result, all parameters of the model will be offloaded and only one copy of the state dict of the model will be kept. During the forward pass, parameters will be extracted from that state dict and put on the execution device and passed as they are needed, then offloaded again. \n\n```python\ncpu_offload(model, execution_device)\n```\n\nYou can also use [`cpu_offload_with_hook`]. This function will offloads a model on the CPU and puts it back to an execution device when executed. The difference with [`cpu_offload`] is that the model stays on the execution device after the forward and is only offloaded again when the `offload` method of the returned `hook` is called. Furthermore, [`cpu_offload_with_hook`] is more performant but less memory saving. It is useful for pipelines running a model in a loop:\n\n```python\nmodel_1, hook_1 = cpu_offload_with_hook(model_1, execution_device)\nmodel_2, hook_2 = cpu_offload_with_hook(model_2, execution_device, prev_module_hook=hook_1)\nmodel_3, hook_3 = cpu_offload_with_hook(model_3, execution_device, prev_module_hook=hook_2)\n\nhid_1 = model_1(input)\nfor i in range(50):\n # model1 is offloaded on the CPU at the first iteration, model 2 stays on the GPU for this whole loop.\n hid_2 = model_2(hid_1)\n# model2 is offloaded to the CPU just before this forward.\nhid_3 = model_3(hid_3)\n\n# For model3, you need to manually call the hook offload method.\nhook_3.offload()\n```\n\n## Disk offload only\n\nTo perform disk offload, you can use [`disk_offload`]. As a result, all parameters of the model will be offloaded as memory-mapped array in a given folder. During the forward pass, parameters will be accessed from that folder and put on the execution device passed as they are needed, then offloaded again.\n\n```python\ndisk_offload(model, offload_dir, execution_device)\n```\n\n## Limits and further development\n\nWe are aware of the current limitations in the API:\n\n- [`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.\n- [`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.\n- [`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.\n- 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.\n- 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.\n- Hard-drive offloading might be very slow if the hardware you run on does not have fast communication between disk and CPU (like NVMes).\n" }, { "path": "docs/source/concept_guides/context_parallelism.md", "content": "\n\n# Context Parallel in šŸ¤—`accelerate`\n\nThis guide will cover basics of using context parallelism in šŸ¤—`accelerate`, for the more curious readers, we will also cover some technicalities in the later sections.\n\nSee also the very related [Guide to Sequence Parallellism](./sequence_parallelism.md).\n\n## Why context parallelism?\n\nWith the advent of large language models, and recently reasoning models, the sequence length has been growing rapidly. This, combined with quadratic memory complexity of attention, has led to a need for more efficient ways to train models with long sequences.\nWith sequence length of 128k, the memory requirement of the attention matrix is `128k * 128k * 2 bytes * num_heads = ~32 GB * num_heads` for `bf16` precision, given vanilla attention implementation. Granted, with usage of `flash attention` or `SDPA` which do not materialize these attention weights, this decreases drastically, but the growth in memory requirements is still considerable.\n\nContext parallelism allows us to shard the inputs to the attention computation along the sequence dimension and compute the attention in parallel on multiple GPUs. With this, we can train models with long sequences, scaling potentially to 1M+ sequence length.\n\n## How to use context parallelism?\n\n```diff\nfrom accelerate.utils import ParallelismConfig, TorchContextParallelConfig\n\n+ cp_config = TorchContextParallelConfig(\n+ cp_comm_strategy=\"alltoall\", # no need to use cp_config at all, if you want to use the default \"allgather\"\n+ )\n\n+ parallelism_config = ParallelismConfig(\n+ cp_size=8,\n+ cp_handler=cp_config, # or just cp_size=8, if you want to use the default \"allgather\"\n+ )\n\naccelerator = Accelerator(\n ...,\n parallelism_config=parallelism_config,\n)\n```\n\nAs with any other feature in šŸ¤—`accelerate`, you can enable context parallelism also by passing the corresponding flags to `accelerate launch`.\nIn this case, it's no different:\n\n```bash\naccelerate launch --parallelism-config-cp-size 8 --parallelism-config-cp-comm-strategy [allgather|alltoall] ...\n```\n\n> [!Tip]\n> You can also set the `cp_size` and `cp_comm_strategy` in the `accelerate config` command, which will save them in your `accelerate` configuration file, so you don't have to pass them every time you launch your script.\n\n> [!Tip]\n> Context parallelism is compatible with other parallelism strategies, such as data parallelism, tensor parallelism and FSDP2.\n> You can simply combine them by setting your parallelism sizes to the desired values, e.g. `--parallelism-config-dp-size 8 --parallelism-config-tp-size 2 --parallelism-config-cp-size 8`. Or you can use the `ParallelismConfig` class to set them programmatically.\n\n> [!Warning]\n> Context parallelism is tightly coupled with `FSDP2`, which you can learn more about in the [FSDP2 introduction](fsdp1_vs_fsdp2.md). Meaning, context parallelism only works if you use `FullyShardedDataParallelPlugin` or `--use-fsdp` with version set to 2 to your\n> program. If no `FSDP2` is used, error will be raised.\n\n> [!Warning]\n> Context parallelism works only with [SDPA](https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html) and only with no mask or causal mask. We can't properly detect this for you, so it's your responsibility to ensure that you are using `SDPA` with no mask or causal mask. If you use any other attention implementation, it will raise an error.\n\nAfter enabling context parallelism with the methods mentioned above, you can then apply it to your training loop. We provide a thin wrapper around [`torch.distributed.tensor.experimental.context_parallel`](https://docs.pytorch.org/docs/stable/distributed.tensor.html#torch.distributed.tensor.experimental.context_parallel) that you can use in your training loop, that abstracts some of the complexity of using it (more on this later). To minimize the changes you have to do in your training loop, we provide a context manager that is a `noop` if context parallelism is not enabled, and applies the context parallelism if it is enabled. This way, you can use it in your training loop without changing any code based on your parallelism configuration.\nYou can use it as follows:\n\n```python\nfor batch in dataloader:\n with accelerator.maybe_context_parallel(\n buffers=[batch[\"input_ids\"], batch[\"attention_mask\"]],\n buffer_seq_dims=[1, 1],\n no_restore_buffers={batch[\"input_ids\"], batch[\"labels\"]},\n ):\n outputs = model(**batch)\n ...\n```\n\n> [!Warning]\n> This context manager has to be recreated with each training step, as shown in the example above. It's crucial to do so.\n\nThis can scale your context size to 1M+ sequence length potentially. Below, we showcase speed and memory usage of context parallelism for up-to 256k context size. We can see that when we double the context size and number of GPUs, we can achieve consistent memory usage, potentially enabling endless context length scaling.\n\n

\n \"context\n
\n Figure 1: Memory usage and speed of context parallelism for up-to 256k context size.\n

\n\n> [!Tip]\n> These examples were created with a script you can find [in the examples folder](https://github.com/huggingface/accelerate/blob/main/examples/fsdp2/nd_parallel.py). To run the example on 8 H100 GPUs (128k sequence length), you can use the following command:\n> ```bash\n> accelerate launch --use-fsdp --fsdp-activation-checkpointing=TRUE examples/fsdp2/nd_parallel.py --cp-size=8 --sequence-length=128000\n> ```\n\n\n## Accelerate's interface\n\nThe context manager takes a few arguments, that are used to configure the context parallelism.\n\n- `buffers`: This is a list of tensors that are to be sharded across the sequence dimension. These tensors are usually input ids, labels and attention mask.\n- `buffer_seq_dims`: This is a list of integers, that specify the sequence dimension of the buffers, in the order of the `buffers` list. If you pass `buffers=[input_ids, shift_labels]` with both having shape `[batch_size, sequence_length]`, you would pass `buffer_seq_dims=[1, 1]`.\n as the sequence dimension is the second dimension of the tensors. This is required for correct computation of the model outputs.\n- `no_restore_buffers`: The implementation of context parallelism modifies the buffers in-place, converting them to `torch.distributed.tensor.Dtensor`s. After the context manager exits, a communication kernel would need to be launched to restore the buffers to their original state (usually all-gather). This takes some time, so it is recommended to pass the same tensors as in the `buffers` argument, to avoid unnecessary communication, unless you are sure that you need to use the buffers after the context manager exits.\n\n\n> [!Warning]\n> Context parallelism is not compatible with `labels` that are a copy of `input_ids`, which models from šŸ¤— transformers can shift to enable causal language modeling themselves.\n> Imagine this case:\n> labels = [l1, l2, l3, l4, ... li]\n> if we apply context parallelism, each rank would end up with a part of labels, such as this:\n> labels_rank_0 = [l1, l2], labels_rank_1 = [l3, l4], ...\n> after transformers modelling code shifts the labels, it would end up with:\n> labels_rank_0 = [l2, PAD], labels_rank_1 = [l3, PAD], ...\n> where `PAD` is a padding token. This would result in incorrect loss computation, as the labels are not aligned with the inputs anymore.\n> Because of this, you need to manually shift the labels before passing them in the model\n\n\n## Configurable options\nAccelerate provides only a single option to configure context parallelism (except for `cp_size`)\n\n- `cp_comm_strategy`: The rotation method to use for the shards. We strongly recommend keeping this as `\"allgather\"`, as it's very likely it will outperform `\"alltoall\"` in most cases.\n\nContext parallel size is rather self-explanatory, it's the number of ranks across which the inputs are to be-sharded.\nContext parallel shard rotation defines how the shards of the inputs are rotated across ranks. We'll cover the 2 options in more detail in the next section.\n\nYou can see an end-to-end example in the [ND parallel example](https://github.com/huggingface/accelerate/blob/main/examples/fsdp2/nd_parallel.py) file, where you can train an 8B model with up-to 128k context length on a single 8xH100 node. Using multi-node training, you can scale this to 1M+ sequence length on multiple GPUs. You can also seamlessly combine it with other parallelism strategies to fit your needs.\n\n## Technical details\n\n> [!Tip]\n> This section is fairly technical, so if you don't need to learn the internals of context parallelism, you can skip it and start building šŸš€\n\nWe're going to be using word `shard` extensively in the following sections, so let's define it first. If we call tensor `sharded` across `Dth` dimension, across `N` ranks, we mean that this tensor is split into `N` parts, where each part of the tensor has shape `[..., D//N, ...]`.\n\n\n## So how does it work?\n\nContext parallelism works on sharding the `Q, K and V` matrices across the sequence dimension. Each rank has its assigned shard of `Q`, let's call it `Q_i`. This matrix stays only on this rank, during the whole computation. Similarly, each rank has its own shard of `K` and `V`, let's call them `K_i` and `V_i`. Then, each rank calculates attention with its own shard of `Q_i`, `K_i` and `V_i`, let's call it `attn_i`. During this computation, a communication kernel is launched to gather the `Ks` and `Vs` from all other ranks. What communication primitive is used, depends on the `context_parallel_shard_rotation` option.\nThis way, each rank gets to calculate local attention, first with `Q_i`, `K_i` and `V_i`, then with `K_j` and `V_j` from all other ranks. As each rank holds `Q, K and V` matrices that are sharded across the sequence dimension, the resulting matrices are smaller and can fit on a single GPU.\n\nWe can formalize this in the following pseudocode:\n```python\ncomm_kernel = {\"allgather\": allgather, \"alltoall\": alltoall}[context_parallel_shard_rotation]\nQi, Ki, Vi = shard(Q, K, V, seq_dim)\nattn[i] = attn(Qi, Ki, Vi)\nfor j in range(context_parallel_size):\n Kj, Vj = comm_kernel()\n attn[j] = attn(Qi, Kj, Vj) # [batch, num_heads, seq_len // context_parallel_size, head_dim]\n\nfinal_attn = combine(attn)\n```\n\n## all-to-all vs all-gather\n\n### all-gather\nSo what's the difference between all-to-all and all-gather? With all-gather, the communication is very simple. After (well, before, as it usually takes longer) we compute the local attention `attn_i` we launch an all-gather to gather all other `Ks` and `Vs` from all other ranks. As this communication is done, each rank has all the `Ks` and `Vs` from all other ranks, and can compute the attention with them sequentially.\nIn ideal scenario, all-gather finishes in the exact moment as the calculation of `attn_i` is done. However, this never happens in practice, so the ideal real overlap is achieved when the full `attn_i` is overlapped with a part of the communication, then to start the computation with `K_j` and `V_j`, we wait for the all-gather to finish.\n\n### all-to-all\nAll-to-all, or sometimes called `ring-rotation` utilizes a ring-like communication pattern. After concluding `attn_i` computation, an all-to-all is launched to send `K_i` and `V_i` to the neighbouring ranks. We then repeat this `context_parallel_size-1` times, so that each rank sees all the shards of `K` and `V` from all other ranks once. In ideal scenario, we prefetch shards `K_i+1` and `V_i+1` from the neighbouring rank and this communication is exactly overlapped with computation of our current `attn_i`. Again, realistically, this perfect overlap doesn't ever happen. Given the nature of this approach, if we don't achieve perfect overlap, the penalty is way larger than with all-gather.\n\n## How to choose the right rotation method?\nIn theory, all-to-all should be the better choice. Though in practice, it rarely is. Therefore, we default to all-gather, as it's more likely to achieve better performance. Extensive [benchmarks](https://discuss.pytorch.org/t/distributed-w-torchtitan-breaking-barriers-training-long-context-llms-with-1m-sequence-length-in-pytorch-using-context-parallel/215082) from the `torchtitan` team also show that all-to-all rarely outperforms all-gather. Though, we still provide both options, as you might find one to be better for your use case.\n\nYou can directly see this issue in the profiler output in the image below:\n

\n \"all-to-all\n
\n Figure 1: In red you can see the idle time, while we wait for the all-to-all kernel to finish. Highlighted in the first blue bar, you can see that it takes ~250us to finish, which is repeated N-1 times for each attention call, where N is the context parallel size.\n

\n\n\n## Why only FSDP2?\n\nWe only support context parallelism with `FSDP2`, as we create a joint mesh of `context_parallel_size` and `dp_shard_size` to\nutilize its full potential.\nHow it works is: we shard the model across the joint mesh of size `cp_size*dp_shard_size`, which maximizes the memory savings.\nThis is a \"free lunch\" of sorts, as `FSDP` communication is fully overlapped with the computation of attention, as shown in the images below.\n\n

\n \"why\n
\n Figure 2: In blue rectangles (Stream 23), you can see that the pre-fetch of `FSDP` shard is fully overlapped with the computation of attention (Stream 7), while in red rectangles (Stream 24), you can see that the all-gather kernel results in a bubble of idle time, in which our compute stream (7) is idle.\n

\n\nIn the figure above, you can also note the difference between all-to-all and all-gather. While in all-to-all (Figure 1), we launch a communication kernel N-1 times for each attention call, in all-gather (Figure 2), we launch a communication kernel only once. This results in a bigger bubble, but it only happens once per attention call, while in all-to-all, it happens N-1 times.\n\n## Data dispatching in joint mesh\n\nWe make sure to dispatch the same batch of data to the whole `cp` subgroup, so that the results are correct. (Meaning each rank in `cp` subgroup gets the same batch of data.) However, we also dispatch different batches to each rank of `dp_shard` group.\nImagine it like this:\n```\n# 8 GPUS, --dp_shard_size 4, --cp_size 2\n# mesh = [[0, 1], [2, 3], [4, 5], [6, 7]]\n# model is sharded across the whole mesh (each GPU holds 1/8 of the model)\n# GPUs 0,1 = batch 0\n# GPUs 2,3 = batch 1\n... and so on.\n```\n\n" }, { "path": "docs/source/concept_guides/deferring_execution.md", "content": "\n\n# Executing and deferring jobs\n\nWhen you run your usual script, instructions are executed in order. Using Accelerate to deploy your script on several\nGPUs at the same time introduces a complication: while each process executes all instructions in order, some may be\nfaster than others.\n\nYou might need to wait for all processes to have reached a certain point before executing a given instruction. For\ninstance, you shouldn't save a model before being sure every process is done with training, and you wouldn't want to \ncontinue training before all the model weights have been loaded in. To do this, just write the following line in your code:\n\n```\naccelerator.wait_for_everyone()\n```\n\nThis instruction will block all the processes that arrive first until all the other processes have reached that\npoint (if you run your script on just one GPU or CPU, this won't do anything).\n\nA few example cases of when to use this utility are listed below:\n\n\n\n Some of these are utilized with the [`~Accelerator.main_process_first`] context manager, which utilizes [`~Accelerator.wait_for_everyone`] to \n run a particular set of code on the main process beforehand before triggering and launching the other processes\n\n\n\n## Downloading a Dataset \n\nWhen downloading a dataset, you should download it first on the main process and then load the cached dataset afterward\n\n\n\n `load_dataset` will perform a lock under the hood to stop multiple downloads from happening at once, but if you are downloading something \n not using this library you should use this method.\n \n\n\n```python\nwith accelerator.main_process_first():\n datasets = load_dataset(\"glue\", \"mrpc\")\n```\n\nUnder the hood this is the same as calling: \n\n```python\n# First do something on the main process\nif accelerator.is_main_process:\n datasets = load_dataset(\"glue\", \"mrpc\")\nelse:\n accelerator.wait_for_everyone()\n\n# And then send it to the rest of them\nif not accelerator.is_main_process:\n datasets = load_dataset(\"glue\", \"mrpc\")\nelse:\n accelerator.wait_for_everyone()\n```\n\n## Saving the `state_dict`\n\nWhen saving the `state_dict` of the model, since you would normally save one file on just the main process\nyou should specify that:\n\n```python\nif accelerator.is_main_process:\n model = accelerator.unwrap_model(model)\n torch.save(model.state_dict(), \"weights.pth\")\n```\n\n## Loading in the `state_dict`\n\nWhen loading in the `state_dict` to a model, optimizer, or scheduler, you should wait \nfor all workers to have the weights loaded in before moving on to training\n\n```python\nwith accelerator.main_process_first():\n state = torch.load(\"weights.pth\")\n model.load_state_dict(state)\n```\n\n## Applying a multi-worker CPU operation \n\nApplying a `map()` operation on multiple workers, such as tokenizing should be done on the \nmain process first, and then propagated to each one. \n\n```python\ndatasets = load_dataset(\"glue\", \"mrpc\")\n\nwith accelerator.main_process_first():\n tokenized_datasets = datasets.map(\n tokenize_function,\n batched=True,\n remove_columns=[\"idx\", \"sentence1\", \"sentence2\"],\n )\n```\n\n## Applying checks such as Early Stopping\n\nTo 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\nfor doing so can include situations such as using early stopping and monitoring the loss (as each loss slightly differs on each process).\n\nCall [`Accelerator.set_trigger`] when your condition has been met, and [`Accelerator.check_trigger`] when checking if that condition has been met in any process:\n\n```python\nfor (x,y) in data_loader:\n logits = model(x)\n loss = loss_func(logits, y)\n # Assume `should_do_early_stopping` is a custom defined function that returns a conditional\n if should_do_early_stopping(loss):\n accelerator.set_trigger()\n\n # Later in the training script when we need to check for the breakpoint\n if accelerator.check_trigger():\n break\n```\n" }, { "path": "docs/source/concept_guides/fsdp1_vs_fsdp2.md", "content": "\n\n# FSDP1 vs FSDP2\n\nThis guide explains the key differences between `FSDP1` and `FSDP2` and helps you migrate your existing code to use `FSDP2` with minimal changes.\n\n## How is FSDP2 better than FSDP1?\n\nFirst, we want to understand how `FSDP1` and `FSDP2` work internally to understand the differences between them. This also helps us understand the limitations of `FSDP1` and how `FSDP2` solves them.\n\nWe'll be discussing a scenario where we have a single `Layer` that contains 3 `Linear` layers and is wrapped using `FSDP` to be sharded across 2 GPUs.\n\n
\n \"Layer\"\n
\n\n### FSDP1\nFirst, we have to understand the original `FSDP1` and the limitations it brings. It represents each `FSDP` module as a single `FlatParameter` which is a single 1D tensor that contains all of the module parameters, which then get sharded across ranks. I.e. if you wrap the `Layer` with `FSDP1`, you'd achieve something as such:\n\n
\n \"FSDP1\"\n
\n\nYou might notice a problem. The whole `Layer` gets flattened into a single `FlatParameter`, which then gets sharded across ranks. But if it's a single `FlatParameter` object, how do we store metadata? That is one of the limitations. Properly storing per-parameter metadata such as `dtype`, `requires_grad`, etc. is not possible without some ugly hacks.\n\n### FSDP2\nThis is why `FSDP2` was introduced. It doesn't use `FlatParameter`, instead it uses `DTensor` which is short for \"Distributed Tensor\". Each `DTensor` basically represents a vanilla `torch.Tensor` that has been sharded across ranks. It contains metadata about the original `torch.Tensor` and how it's sharded, what is the [placement type](https://pytorch.org/docs/stable/distributed.tensor.html#module-torch.distributed.tensor.placement_types) and so on. This is why it's called `per-parameter sharding`. The following figure shows the difference:\n\n
\n \"FSDP2\"\n
\n\nEach Parameter of the original `Layer` is sharded across the 0th dimension, and split between 2 GPUs. Now, each `Linear` layer is a separate `DTensor` and storing metadata per-parameter is possible and straightforward.\n\n\n> [!TIP] \n> In the image above, the tensors were sharded across the 1st dimension for the sake of fitting the image on the screen, in reality, they are sharded across the 0th dimension as stated above\n\n## What does FSDP2 offer?\n\n`FSDP2` is a new and improved version of PyTorch's fully-sharded data parallel training API. Its main advantage is using `DTensor` to represent sharded parameters. Compared to `FSDP1`, it offers:\n- Simpler internal implementation, where each `Parameter` is a separate `DTensor`\n- Enables simple partial parameter freezing because of the above, which makes methods as [`LORA`](https://huggingface.co/papers/2106.09685) work out of the box\n- With `DTensor`, `FSDP2` supports mixing `fp8` and other parameter types in the same model out of the box\n- Faster and simpler checkpointing without extra communication across ranks using `SHARDED_STATE_DICT` and [`torch.distributed.checkpoint`](https://pytorch.org/docs/stable/distributed.checkpoint.html), this way, each rank only saves its own shard and corresponding metadata\n- For loading, it uses a `state_dict` of the sharded model to directly load the sharded parameters\n- Support for asynchronous checkpointing, where parameters are first copied to CPU memory, after this, main thread continues training while another thread stores the parameters on disk\n- Memory efficiency and deterministic memory usage, `FSDP2` doesn't use `recordStream` anymore and uses stream-to-stream synchronization (for more technical details see [this forum post](https://dev-discuss.pytorch.org/t/fsdp-cudacachingallocator-an-outsider-newb-perspective/1486) and [this issue](https://github.com/pytorch/pytorch/issues/114299))\n- In the future, optimizations of the communication patterns via `torch.compile` are planned, further improving the performance and memory efficiency\n\n\n## API Differences\n\nWe have already discussed the internal differences, now let's discuss the differences, you, as a user, will need to know. \n\nHere are the main changes in configuration options when using `FSDP2` through the `accelerate` CLI:\n\nPrevious (`FSDP1`) | New (`FSDP2`) | What Changed\n-- | -- | --\n`--fsdp_sharding_strategy` | `--fsdp_reshard_after_forward` | replaces `--fsdp_sharding_strategy`, changed to `true` (previously `FULL_SHARD`) or `false` (previously `SHARD_GRAD_OP`)\n`--fsdp_backward_prefetch` | \\*\\***REMOVED**\\*\\* | `FSDP2` uses previous `BACKWARD_PRE` option by default, as only this allows communication and computation overlap\n`--fsdp_forward_prefetch` | \\*\\***NOT YET IMPLEMENTED**\\*\\* | How to implement this is under active discussion, for now it is not supported in `FSDP2`\n`--fsdp_sync_module_states` | \\*\\***REMOVED**\\*\\* | with `FSDP2`, this parameter becomes redundant\n`--fsdp_cpu_ram_efficient_loading` | `--fsdp_cpu_ram_efficient_loading` | if `true`, `FSDP2` will similarly load the model only on rank 0, and then parameters get synced to other ranks, this is the same behavior as `FSDP1`, however, setting `--fsdp_sync_module_states` isn't required anymore\n`--fsdp_state_dict_type` | `--fsdp_state_dict_type` | `LOCAL_STATE_DICT` becomes obsolete and with `FSDP2` `SHARDED_STATE_DICT` is the default option, which results in no extra communication and each rank saving its own shard, other possible option is `FULL_STATE_DICT` which results in extra communication and spike in memory usage but saves the full model from rank 0.\n`--fsdp_use_orig_params` | \\*\\***REMOVED**\\*\\* | `FSDP2` uses a `DTensor` class on the background, which means it *always* uses the original parameters by default\n\\*\\***NEW**\\*\\* | `--fsdp_version` | `1` is the default option, to not break existing code, set to `2` to use `FSDP2`\n\nFor all other options that remain unchanged, see the [`FSDP` documentation](../usage_guides/fsdp.md).\n\n## How to Switch to FSDP2\n\n### If using Python code:\nSimply set `fsdp_version=2` when creating your plugin and replace options according to the table above.\n\n```python\nfrom accelerate import FullyShardedDataParallelPlugin, Accelerator\n\nfsdp_plugin = FullyShardedDataParallelPlugin(\n fsdp_version=2\n # other options...\n)\naccelerator = Accelerator(fsdp_plugin=fsdp_plugin)\n```\n\n### If using YAML config:\nUse our conversion tool:\n```bash\naccelerate to-fsdp2 --config_file config.yaml --output_file new_config.yaml\n```\n\nThis will automatically convert all FSDP1 settings to their FSDP2 equivalents. Use `--overwrite` to update the existing file instead of creating a new one.\n" }, { "path": "docs/source/concept_guides/fsdp_and_deepspeed.md", "content": "\n\n# FSDP vs DeepSpeed\n\nAccelerate offers flexibility of training frameworks, by integrating two extremely powerful tools for distributed training, namely [Pytorch FSDP](../usage_guides/fsdp) and [Microsoft DeepSpeed](../usage_guides/deepspeed). The aim of this tutorial is to draw parallels, as well as to outline potential differences, to empower the user to switch seamlessly between these two frameworks.\n\n\n\n To switch between the frameworks, we recommend launching code `accelerate launch` passing in the correct config file with `--config_file`, or passing in the respective arguments directly for [FSDP and DeepSpeed](../package_reference/cli#accelerate-launch) .\n\n Example Accelerate configurations can be found here for [DeepSpeed](../usage_guides/deepspeed#accelerate-deepspeed-plugin) and [FSDP](../usage_guides/fsdp#how-it-works-out-of-the-box), or in the [example zoo under \"Launch Configurations\"](../usage_guides/explore)\n \n\n\n\n\nThis tutorial is for single-node, multi-GPU, scenarios only.\n\n\n\n## Configuring Functionalities\n\nModel tensors are split into different GPUs in an attempt to scale up model sizes; this is termed *sharding* in FSDP, and *partitioning* in DeepSpeed. FSDP sharding and DeepSpeed ZeRO (partitioning) stages are configured by `--fsdp_sharding_strategy`, and `--zero_stage`, respectively. In particular, FSDP `FULL_SHARD` maps to DeepSpeed ZeRO stage `3`; see this [comprehensive mapping between FSDP sharding and DeepSpeed ZeRO settings](../usage_guides/fsdp#mapping-between-fsdp-sharding-strategies-and-deepspeed-zero-stages). The below table summarizes and groups similar settings:\n\nGroup | Framework | Configuration | Example | Restrictions (if any)\n--|--|--|--|--\nsharding / partitioning | FSDP
DeepSpeed | `--fsdp_sharding_strategy`
`--zero_stage` | `1` (`FULL_SHARD`)
`3` | \noffload | FSDP
DeepSpeed | `--fsdp_offload_params`
`--offload_param_device`
`--offload_optimizer_device` | `true`
`cpu`
`cpu` | all or nothing

\nmodel loading | FSDP
DeepSpeed | `--fsdp_cpu_ram_efficient_loading`
`--zero3_init_flag` | `true`
`true` |
only ZeRO 3\nefficient checkpointing | FSDP
DeepSpeed | `--fsdp_state_dict_type`
`--zero3_save_16bit_model` | `SHARDED_STATE_DICT`
`true` |
only ZeRO 3\nweights prefetching | FSDP

DeepSpeed | `--fsdp_forward_prefetch`
`--fsdp_backward_prefetch`
None | `true`
`BACKWARD_PRE` |

\nmodel | FSDP

DeepSpeed | `--fsdp_auto_wrap_policy`
`--fsdp_transformer_layer_cls_to_wrap`
None | `TRANSFORMER_BASED_WRAP`
|
Usually not needed
Transparent to user.\nparameters summoning | FSDP
DeepSpeed | `--fsdp_use_orig_params`
None | `true` | required for `torch.compile`
Transparent to user\nparameters syncing | FSDP
DeepSpeed | `--fsdp_sync_module_states`
None | `true` | \ntraining | FSDP
DeepSpeed | None
`--gradient_accumulation_steps`
`--gradient_clipping` |
`auto`
`auto` | Transparent to user\n\nFor detailed descriptions of the above, refer to [`Accelerate` launch documentation](../package_reference/cli#accelerate-launch).\n\n\n\n To access other DeepSpeed configurations, such as mixed precision settings, \n you need to pass in a `--deepspeed_config_file`, see the [documentation](../usage_guides/deepspeed#deepspeed-config-file). \n\n DeepSpeed can be also configured via [`DeepSpeedPlugin`], e.g., `DeepSpeedPlugin.zero_stage` is equivalent of `--zero_stage`, and `DeepSpeedPlugin.hf_ds_config` can be used to pass `--deepeed_config_file.`\n\n\n\n\n\n FSDP can be also configured via [`FullyShardedDataParallelPlugin`], e.g., `FullyShardedDataParallelPlugin.sharding_strategy` is equivalent of `--fsdp_sharding_strategy`.\n \n\n\n### Checkpointing\n\nDo note that while FSDP can be configured via `--fsdp_state_dict_type` to save either full / sharded checkpoints.\n\n\n\n For DeepSpeed Zero3, one could pass a `--zero3_save_16bit_model true`, which conveniently consolidates the model to a single rank and saves; this is the FSDP equivalent of `fsdp_state_dict_type: FULL_STATE_DICT`. \n\n\n\n\n\n For large models, consolidating the model to a single rank can be very slow.\n\n\n\n\n\n For quicker checkpointing, for FSDP use `fsdp_state_dict_type: SHARDED_STATE_DICT`, and for DeepSpeed Zero3 [use the `zero_to_fp32.py` script to post-convert sharded checkpoints](https://www.deepspeed.ai/tutorials/zero/#extracting-weights).\n\n\n\n\n### Offloading\n\nFSDP only allows *all-or-nothing* offload (i.e., either offload parameters, gradients, and optimizer, or keep them all in GPU), but DeepSpeed can offload parameters and optimizer differently. Furthermore, DeepSpeed also supports [offloading to NVME](https://www.deepspeed.ai/docs/config-json/#parameter-offloading).\n\n### Prefetching\n\nFSDP allows two prefetching configurations `--fsdp_forward_prefetch` and `--fsdp_backward_prefetch` to improve overlap of comms / computation at a cost of extra memory, see [FSDP documentation](https://pytorch.org/docs/stable/fsdp.html). \nFor DeepSpeed, the prefetching will be turned on when needed, and it turns on depending on certain hyper-params like `stage3_param_persistence_threshold`, `stage3_max_reuse_distance`, etc, [that can be configured for Zero3](https://www.deepspeed.ai/docs/config-json/#parameter-offloading); `accelerate` may set these hyper-params automatically if you don't set those explicitly in the deepspeed config file.\n\n\n\n For FSDP set `fsdp_backward_prefetch: BACKWARD_PRE` for improved throughputs if memory allows.\n\n\n\n### Model Loading\n\nWhile FSDP require an explicit `--fsdp_cpu_ram_efficient_loading true` to activate efficient model loading, `transformers` will activate the similar feature whenever DeepSpeed Zero3 is used.\n\n\n\n For FSDP, whenever setting `--fsdp_cpu_ram_efficient_loading true`, `accelerate` will automatically set `sync_module_states` to true. \n For RAM efficient loading the weights will be loaded only in a single rank, and thus requires `sync_module_states` to broadcast weights to other ranks.\n\n\n\n### Model\n\nFSDP requires an explicit `--fsdp_auto_wrap_policy` for the algorithm to decide how to schedule the all-gather and reduce-scatter operations. But for DeepSpeed this is transparent to the user.\n\n\n\n For FSDP, simply set `fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP`. With the latest [`transformers`] versions, we try our best to figure out the suitable `fsdp_transformer_layer_cls_to_wrap` for HF transformers models. However, if you get an error regarding it, please specify this.\n\n\n\n### Parameters Summoning\n\nFSDP requires an explicit `--fsdp_use_orig_params` flag if using `torch.compile`, see [the pytorch documentation](https://pytorch.org/docs/stable/fsdp.html#module-torch.distributed.fsdp). For DeepSpeed this is transparent to the user.\n\n\n\n For FSDP, when using `torch.compile` please set `fsdp_use_orig_params: True`.\n\n\n\n\n## Training\n\nDeepspeed requires explicit `--gradient_accumulation_steps` and `--gradient_clipping` flags. For FSDP this is transparent to the user.\n\n\n\n When using DeepSpeed, set `gradient_accumulation_steps: \"auto\"` and `gradient_clipping: \"auto\"` to automatically pick up values set in the [`Accelerator`] or [`TrainingArguments`] (if using `transformers`).\n\n\n\n\n## On Differences in Data Precision Handling\n\nTo discuss how data precision is handled in both FSDP and Deepspeed, it is instructive to first give an overview of how model parameters are handled in these frameworks. Before the model / optimizer parameters are distributed across GPUs, parameter preparation is involved to first \"flatten\" them to one-dimensional [`torch.Tensor`](https://pytorch.org/docs/stable/tensors.html#torch-tensor). The implementation of FSDP / DeepSpeed varies in the respect of the `dtype` in which these \"flattened\" parameters are stored, and there are ramifications with regards to how [`torch.Optimizer`](https://pytorch.org/docs/stable/optim.html#module-torch.optim) allocate their `dtype`s. The table below outlines the processes for both frameworks; the \"Local\" column indicates the process occurring at a per-gpu level, therefore any memory overheads by upcasting should be understood to be amortized by the number of gpus used.\n\n\n\n As a rule of thumb, for stable training with automatic mixed precision, all the trainable parameters have to be in `torch.float32`.\n\n\n\nProcess | Local | Framework | Details\n--|--|--|--\nLoading, i.e., [`AutoModel.from_pretrained(..., torch_dtype=torch_dtype)`] | \nPreparation, i.e., creation of \"flat params\" | āœ… | FSDP
DeepSpeed | created in `torch_dtype`.
disregards `torch_dtype`, created in `float32`.\nOptimizer initialization | āœ… | FSDP
DeepSpeed | creates parameters in `torch_dtype`
creates parameters in `float32`\nTraining Step, i.e, forward, backward, reduction | | FSDP
DeepSpeed | follows [`MixedPrecision`](https://pytorch.org/docs/stable/fsdp.html#torch.distributed.fsdp.MixedPrecision)
follows `deepspeed_config_file` mixed precision settings.\nOptimizer (Pre-Step) | āœ… | FSDP
DeepSpeed | upcasting (if any) to `torch_dtype`
upcasted to `float32`\nOptimizer (Actual Step) | āœ… | FSDP
DeepSpeed | occurs in `torch_dtype`
occurs in `float32`.\n\n\n\n Therefore when using DeepSpeed a small number of GPUs, be aware of potentially significant memory overheads due to the upcasting during preparation.\n\n\n\n\n\n With FSDP, in the absence of mixed precision, it is possible to operate the [`torch.Optimizer`](https://pytorch.org/docs/stable/optim.html#module-torch.optim) in low precision `torch_dtype`, which may be helpful when using small number of GPUs. \n\n\n\n\n\n With mixed precision, FSDP and DeepSpeed will upcast in the model preparation step (c.f. table above). But do note that FSDP will then save checkpoints in the upcasted precision; Deepspeed may still save low precision checkpoints if `--zero3_save_16bit_model` is specified.\n\n\n\n\nTo clarify the above table consider the concrete examples below; the optimizer pre- and actual step combined for brevity. With FSDP it is possible to operate in the two modes shown below, but DeepSpeed can only operate in one.\n\nFramework | Model Loading (`torch_dtype`) | Mixed Precision | Preparation (Local) | Training | Optimizer (Local)\n--|--|--|--|--|--\nFSDP | bf16 | default (none) | bf16 | bf16 | bf16\nFSDP | bf16 | bf16 | fp32 | bf16 | fp32\nDeepSpeed | bf16 | bf16 | fp32 | bf16 | fp32\n" }, { "path": "docs/source/concept_guides/gradient_synchronization.md", "content": "\n\n# Gradient synchronization\n\nPyTorch's distributed module operates by communicating back and forth between all of the GPUs in your system.\nThis communication takes time, and ensuring all processes know the states of each other happens at particular triggerpoints\nwhen using the `ddp` module. \n\nThese triggerpoints are added to the PyTorch model, specifically their `forward()` and `backward()` methods. \nThis happens when the model is wrapped with `DistributedDataParallel`:\n```python\nimport torch.nn as nn\nfrom torch.nn.parallel import DistributedDataParallel\n\nmodel = nn.Linear(10, 10)\nddp_model = DistributedDataParallel(model)\n```\nIn Accelerate this conversion happens automatically when calling [`~Accelerator.prepare`] and passing in your model.\n\n```diff\n+ from accelerate import Accelerator\n+ accelerator = Accelerator()\n import torch.nn as nn\n- from torch.nn.parallel import DistributedDataParallel\n\n model = nn.Linear(10,10)\n+ model = accelerator.prepare(model)\n```\n\n## The slowdown in gradient accumulation\n\nYou now understand that PyTorch adds hooks to the `forward` and `backward` method of your PyTorch model when \ntraining in a distributed setup. But how does this risk slowing down your code?\n\nIn DDP (distributed data parallel), the specific order in which processes are performed and ran are expected\nat specific points and these must also occur at roughly the same time before moving on.\n\nThe most direct example is when you update model parameters through\n`optimizer.step()`.\nWithout gradient accumulation, all instances of the model need to have updated\ntheir gradients computed, collated, and updated before moving on to the next\nbatch of data.\nWhen performing gradient accumulation, you accumulate `n` loss gradients and\nskip `optimizer.step()` until `n` batches have been reached. As all training\nprocesses only need to synchronize by the time `optimizer.step()` is called,\nwithout any modification to your training step, this needless inter-process\ncommunication can cause a significant slowdown.\n\n How can you avoid this overhead?\n\n## Solving the slowdown problem\n\nSince 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. \nPyTorch 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\nthat is added to your model after converting it to DDP.\n\nUnder this context manager, PyTorch will skip synchronizing the gradients when\n`.backward()` is called, and the first call to `.backward()` outside this \ncontext manager will trigger the synchronization. See an example below:\n```python\nddp_model, dataloader, optimizer = accelerator.prepare(model, dataloader, optimizer)\n\nfor index, batch in enumerate(dataloader):\n inputs, targets = batch\n # Trigger gradient synchronization on the last batch\n if index != (len(dataloader) - 1):\n with ddp_model.no_sync():\n # Gradients only accumulate\n outputs = ddp_model(inputs)\n loss = loss_func(outputs)\n accelerator.backward(loss)\n else:\n # Gradients finally sync\n outputs = ddp_model(inputs)\n loss = loss_func(outputs)\n accelerator.backward(loss)\n optimizer.step()\n```\n\nIn 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!),\n`ddp_model.no_sync` gets replaced with [`~Accelerator.no_sync`] and operates the same way:\n\n```diff\n ddp_model, dataloader, optimizer = accelerator.prepare(model, dataloader, optimizer)\n\n for index, batch in enumerate(dataloader):\n inputs, targets = batch\n # Trigger gradient synchronization on the last batch\n if index != (len(dataloader)-1):\n- with ddp_model.no_sync():\n+ with accelerator.no_sync(model):\n # Gradients only accumulate\n outputs = ddp_model(inputs)\n loss = loss_func(outputs, targets)\n accelerator.backward(loss)\n else:\n # Gradients finally sync\n outputs = ddp_model(inputs)\n loss = loss_func(outputs)\n accelerator.backward(loss)\n optimizer.step()\n optimizer.zero_grad()\n```\n\nAs 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\ngradient accumulation API:\n\n```python\nddp_model, dataloader, optimizer = accelerator.prepare(model, dataloader, optimizer)\n\nfor batch in dataloader:\n with accelerator.accumulate(model):\n optimizer.zero_grad()\n inputs, targets = batch\n outputs = model(inputs)\n loss = loss_function(outputs, targets)\n accelerator.backward(loss)\n optimizer.step()\n optimizer.zero_grad()\n```\n\nAs a result, you should either use *`accelerator.accumulate` or `accelerator.no_sync`* when it comes to API choice. \n\n## Just how much of a slowdown is there, and easy mistakes you can make\n\nTo set up a realistic example, consider the following setup:\n\n* Two single-GPU T4 nodes and one node with two GPUs\n* Each GPU is a T4, and are hosted on GCP\n* The script used is a modification of the [NLP Example](https://github.com/muellerzr/timing_experiments/blob/main/baseline.py) script\n* Batch size per GPU is 16, and gradients are accumulated every 4 steps\n\nAll scripts are available in [this repository](https://github.com/muellerzr/timing_experiments).\n\nIf not careful about gradient synchronization and GPU communication, a *large* amount of time can be wasted \nfrom when these GPUs communicate to each other during unnecessary periods.\n\nBy how much?\n\nReference:\n- Baseline: uses no synchronization practices discussed here\n- `no_sync` improperly: `no_sync` only around the `backward` call, not the `forward`\n- `no_sync`: using the `no_sync` pattern properly\n- `accumulate`: using [`~Accelerator.accumulate`] properly\n\nBelow 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:\n\n| | Baseline | `no_sync` improperly | `no_sync` | `accumulate`| \n| :---------: | :-------: | :------------------: | :-------: | :---------: |\n| Multi-Node | 2Ā±0.01s | 2.13Ā±0.08s | **0.91Ā±0.11s** | **0.91Ā±0.11s** |\n| Single Node | 0.50Ā±0.01s | 0.50Ā±0.01s | **0.41Ā±0.015s** | **0.41Ā±0.015s** |\n\nAs 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!\n\nIf you are worried about making sure everything is done properly, we highly recommend utilizing the [`~Accelerator.accumulate`] function and passing in\n`gradient_accumulation_steps` or `gradient_accumulation_plugin` to the [`Accelerator`] object so Accelerate can handle this for you.\n\n### `no_sync` requires additional GPU memory when using FSDP\n\nBe aware that not syncing gradients can have adverse effects while performing FSDP training. As it has been warned in `torch`, the [`no_sync` context manager for FSDP](https://pytorch.org/docs/stable/fsdp.html#torch.distributed.fsdp.FullyShardedDataParallel.no_sync) will require additional memory.\n\nTherefore in memory intensive situations while using FSDP, we recommend to set `sync_each_batch` to `True` in the [`~utils.GradientAccumulationPlugin`] to disable `no_sync`.\n\nSee the example below where we fine-tune Mixtral (47B parameters) on 8 A100-80GB GPUs. We see that even for a modest `gradient_accumulation_steps=2` we quickly go out-of-memory (OOM) if `no_sync` is enabled. Again, this is due to additional memory overheads due to FSDP's `no_sync`. However, if `no_sync` is disabled via `sync_each_batch=True`, then the memory consumption for `gradient_accumulation_steps=16` reverts to that of `gradient_accumulation_steps=1`.\n\n| Model | `no_sync` (accum=1) | `no_sync` (accum=2) | `no_sync` disabled (accum=16)\n| :-------------: | :-----------------: | :-----------------: | :-----------------: \nmixtral 8x7B | 69G | OOM | 69G\n\n> [!WARNING] \n> Disabling `no_sync` means there _will be slowdown_ due the extra data syncs, as explained by the earlier sections of this guide." }, { "path": "docs/source/concept_guides/internal_mechanism.md", "content": "\n\n# Accelerate's internal mechanisms\n\nInternally, Accelerate works by first analyzing the environment in which the script is launched to determine which\nkind of distributed setup is used, how many different processes there are and which one the current script is in. All\nthat information is stored in the [`~AcceleratorState`].\n\nThis class is initialized the first time you instantiate an [`~Accelerator`] as well as performing any\nspecific initialization your distributed setup needs. Its state is then uniquely shared through all instances of\n[`~state.AcceleratorState`]. (The same can also be done with the [`PartialState`], a more barebones version it inherits)\n\nThen, when calling [`~Accelerator.prepare`], the library:\n\n- wraps your model(s) in the container adapted for the distributed setup,\n- wraps your optimizer(s) in an [`~optimizer.AcceleratedOptimizer`],\n- wraps your scheduler(s) in an [`~scheduler.AcceleratedScheduler`]\n- creates a new version of your dataloader(s) in a [`~data_loader.DataLoaderShard`] or [`~data_loader.DataLoaderDispatcher`]\n\nWhile the model(s), optimizer(s), and scheduler(s) are just put in simple wrappers, the dataloader(s) are re-created. This is mostly\nbecause PyTorch does not let the user change the `batch_sampler` of a dataloader once it's been created and the\nlibrary handles the sharding of your data between processes by changing that `batch_sampler` to yield every other\n`num_processes` batches (if enabled).\n\nThe [`~data_loader.DataLoaderShard`] subclasses `DataLoader` to add the following functionality:\n\n- it synchronizes the appropriate random number generator of all processes at each new iteration, to ensure any\n randomization (like shuffling) is done the exact same way across processes.\n- it puts the batches on the proper device before yielding them (unless you have opted out of\n `device_placement=True`).\n \nThe [`~data_loader.DataLoaderDispatcher`] subclasses differs from the [`~data_loader.DataLoaderShard`] in that when iterating through the `DataLoader`, the data is all starting from process 0 and *then* split and sent off to each process rather than it happening at the dataset level.\n\nThe random number generator synchronization will by default synchronize:\n\n- the `generator` attribute of a given sampler (like the PyTorch `RandomSampler`) for PyTorch >= 1.6\n- the main random number generator in PyTorch <=1.5.1\n\nYou can choose which random number generator(s) to synchronize with the `rng_types` argument of the main\n[`Accelerator`]. In PyTorch >= 1.6, it is recommended to rely on a local `generator` to avoid\nsetting the same seed in the main random number generator in all processes.\n\n\n\n Synchronization of the main torch (or CUDA or XLA) random number generator will affect any other potential random\n artifacts you could have in your dataset (like random data augmentation) in the sense that all processes will get\n the same random numbers from the torch random modules (so will apply the same random data augmentation if it's\n controlled by torch).\n\n\n\n\n\n The randomization part of your custom sampler, batch sampler or iterable dataset should be done using a local\n `torch.Generator` object (in PyTorch >= 1.6), see the traditional `RandomSampler`, as an example.\n\n\n\nIf you have [`torchdata>=0.8.0`](https://github.com/pytorch/data/tree/main) installed, and you have passed `use_stateful_dataloader=True` into your [`~utils.DataLoaderConfiguration`], these classes will directly inherit from `StatefulDataLoader` instead, and maintain a `state_dict`.\n\nFor more details about the internals, see the [Internals page](../package_reference/torch_wrappers).\n" }, { "path": "docs/source/concept_guides/low_precision_training.md", "content": "\n\n# Low precision training methods\n\nThe release of new kinds of hardware led to the emergence of new training paradigms that better utilize them. Currently, this is in the form of training\nin 8-bit precision using packages such as [TransformersEngine](https://github.com/NVIDIA/TransformerEngine) (TE), [torchao](https://github.com/pytorch/ao) (native PyTorch FP8), or the legacy [MS-AMP](https://github.com/Azure/MS-AMP/tree/main) (no longer maintained, see warning below).\n\nFor an introduction to the topics discussed today, we recommend reviewing the [low-precision usage guide](../usage_guides/low_precision_training) as this documentation will reference it regularly. \n\n## A Quick Chart\n\nBelow is a quick chart from the MS-AMP documentation showing the different bit-precisions for each solution during training:\n\nOptimization Level | Computation(GEMM) | Comm | Weight | Master Weight | Weight Gradient | Optimizer States\n-- | -- | -- | -- | -- | -- | --\nFP16 AMP | FP16 | FP32 | FP32 | N/A | FP32 | FP32+FP32\nNvidia TE | FP8 | FP32 | FP32 | N/A | FP32 | FP32+FP32\nMS-AMP O1 | FP8 | FP8 | FP16 | N/A | FP8 | FP32+FP32\nMS-AMP O2 | FP8 | FP8 | FP16 | N/A | FP8 | FP8+FP16\nMS-AMP O3 | FP8 | FP8 | FP8 | FP16 | FP8 | FP8+FP16\n\n## `TransformersEngine`\n\n`TransformersEngine` is the first solution to trying to train in 8-bit floating point. It works by using drop-in replacement layers for certain ones in a model that utilizes their FP8-engine to reduce the number of bits (such as 32 to 8) without degrading the final accuracy of the model. \n\nSpecifically, Accelerate will find and replace the following layers with `TransformersEngine` versions:\n\n* `nn.LayerNorm` for `te.LayerNorm`\n* `nn.Linear` for `te.Linear`\n\nAs a result we wind up with a model that has most of its layers in BF16, while some layers are in FP8 reducing some of the memory. \n\nAnecdotally, we have noticed that performance gains don't really start showing when using `TransformerEngine` until a large majority of the layers\nin the model are made up of those two layers to replace. As a result, only larger models have shown performance improvements when the number of parameters is around and upwards of a few billion. \n\nThe `TransformerEngine` can receive many different arguments that customize how it performs FP8 calculations and what they do. A full list of the arguments is available below:\n\n* `margin`: The margin to use for the gradient scaling.\n* `interval`: The interval to use for how often the scaling factor is recomputed.\n* `fp8_format``: The format to use for the FP8 recipe. Must be one of `HYBRID` or `E4M3`. (Generally `HYBRID` for training, `E4M3` for evaluation)\n* `amax_history_len`: The length of the history to use for the scaling factor computation\n* `amax_compute_algo`: The algorithm to use for the scaling factor computation. Must be one of `max` or `most_recent`.\n* `override_linear_precision`: Whether or not to execute `fprop`, `dgrad`, and `wgrad` GEMMS in higher precision.\n\nYou can customize each of these as part of [`utils.FP8RecipeKwargs`] to help optimize performance of your models.\n\nIf we notice in the chart mentioned earlier, TE simply casts the computation layers into FP8, while everything else is in FP32. As a result this winds up utilizing the most memory but does so with the benefit of guaranteeing the least amount of loss in end accuracy during training. \n\n## `MS-AMP`\n\n\n\n**āš ļø Deprecated / Unmaintained:** MS-AMP is no longer actively maintained by Microsoft. The repository has not seen updates since 2023 and has known compatibility issues with CUDA 12.x+, modern NCCL versions, and recent PyTorch releases (2.2+). **We strongly recommend using `TransformersEngine` or `torchao` instead.** See the [usage guide](../usage_guides/low_precision_training) for migration instructions.\n\n\n\nMS-AMP takes a different approach to `TransformersEngine` by providing three different optimization levels to convert more operations in FP8 or FP16.\n\n* The base optimization level (`O1`), passes communications of the weights (such as in DDP) in FP8, stores the weights of the model in FP16, and leaves the optimizer states in FP32. The main benefit of this optimization level is that we can reduce the communication bandwidth by essentially half. Additionally, more GPU memory is saved due to 1/2 of everything being cast in FP8, and the weights being cast to FP16. Notably, both the optimizer states remain in FP32.\n\n* The second optimization level (`O2`) improves upon this by also reducing the precision of the optimizer states. One is in FP8 while the other is in FP16. Generally it's been shown that this will only provide a net-gain of no degraded end accuracy, increased training speed, and reduced memory as now every state is either in FP16 or FP8. \n\n* Finally, MS-AMP has a third optimization level (`O3`) which helps during DDP scenarios such as DeepSpeed. The weights of the model in memory are fully cast to FP8, and the master weights are now stored in FP16. This fully reduces memory by the highest factor as now not only is almost everything in FP8, only two states are left in FP16. Currently, only DeepSpeed versions up through 0.9.2 are supported, so this capability is not included in the Accelerate integration\n\n## Combining the two\n\n\n\nSince MS-AMP is no longer maintained, this combination is not recommended for new projects.\n\n\n\nMore experiments need to be performed but it's been noted that combining both MS-AMP and TransformersEngine can lead to the highest throughput by relying on NVIDIA's optimized FP8 operators and utilizing how MS-AMP reduces the memory overhead.\n" }, { "path": "docs/source/concept_guides/performance.md", "content": "\n\n# Comparing performance across distributed setups\n\nEvaluating and comparing the performance from different setups can be quite tricky if you don't know what to look for.\nFor example, you cannot run the same script with the same batch size across TPU, multi-GPU, and single-GPU with Accelerate \nand expect your results to line up. \n\nBut why?\n\nThere are three reasons for this that this tutorial will cover: \n\n1. **Setting the right seeds**\n2. **Observed Batch Sizes**\n3. **Learning Rates**\n\n## Setting the Seed \n\nWhile 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:\n\n```python\nfrom accelerate.utils import set_seed\n\nset_seed(42)\n```\n\nWhy is this important? Under the hood this will set **5** different seed settings:\n\n```python\n random.seed(seed)\n np.random.seed(seed)\n torch.manual_seed(seed)\n torch.cuda.manual_seed_all(seed) # or torch.xpu.manual_seed_all, etc\n # ^^ safe to call this function even if cuda is not available\n if is_torch_xla_available():\n xm.set_rng_state(seed)\n```\n\nThe random state, numpy's state, torch, torch's device state, and if TPUs are available torch_xla's cuda state.\n\n## Observed Batch Sizes \n\nWhen training with Accelerate, the batch size passed to the dataloader is the **batch size per GPU**. What this entails is \na batch size of 64 on two GPUs is truly a batch size of 128. As a result, when testing on a single GPU this needs to be accounted for,\nas well as similarly for TPUs. \n\nThe below table can be used as a quick reference to try out different batch sizes:\n\n\n\nIn this example, there are two GPUs for \"Multi-GPU\" and a TPU pod with 8 workers\n\n\n\n| Single GPU Batch Size | Multi-GPU Equivalent Batch Size | TPU Equivalent Batch Size |\n|-----------------------|---------------------------------|---------------------------|\n| 256 | 128 | 32 |\n| 128 | 64 | 16 |\n| 64 | 32 | 8 |\n| 32 | 16 | 4 |\n\n## Learning Rates \n\nAs noted in multiple sources[[1](https://aws.amazon.com/blogs/machine-learning/scalable-multi-node-deep-learning-training-using-gpus-in-the-aws-cloud/)][[2](https://docs.nvidia.com/clara/clara-train-sdk/pt/model.html#classification-models-multi-gpu-training)], the learning rate should be scaled *linearly* based on the number of devices present. The below \nsnippet shows doing so with Accelerate:\n\n\n\nSince users can have their own learning rate schedulers defined, we leave this up to the user to decide if they wish to scale their \nlearning rate or not.\n \n\n\n```python\nlearning_rate = 1e-3\naccelerator = Accelerator()\nlearning_rate *= accelerator.num_processes\n\noptimizer = AdamW(params=model.parameters(), lr=learning_rate)\n```\n\nYou will also find that `accelerate` will step the learning rate based on the number of processes being trained on. This is because \nof 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\nto account for the batch size being twice as large (if no changes to the batch size on the single GPU instance are made).\n\n## Gradient Accumulation and Mixed Precision\n\nWhen using gradient accumulation and mixed precision, due to how gradient averaging works (accumulation) and the precision loss (mixed precision), \nsome degradation in performance is expected. This will be explicitly seen when comparing the batch-wise loss between different compute \nsetups. However, the overall loss, metric, and general performance at the end of training should be _roughly_ the same.\n" }, { "path": "docs/source/concept_guides/sequence_parallelism.md", "content": "\n\n# Sequence parallel in šŸ¤—`accelerate`\n\nThis guide will cover basics of using sequence parallelism in šŸ¤—`accelerate`.\n\nSee also the very related [Context Parallellism](./context_parallelism.md).\n\n## Why sequence parallelism?\n\nWith the advent of large language models, and recently reasoning models, the sequence length has been growing rapidly. This, combined with quadratic memory complexity of attention, has led to a need for more efficient ways to train models with long sequences.\nWith sequence length of 128k, the memory requirement of the attention matrix is `128k * 128k * 2 bytes * num_heads = ~32 GB * num_heads` for `bf16` precision, given vanilla attention implementation. Granted, with usage of `flash attention` or `SDPA` which do not materialize these attention weights, this decreases drastically, but the growth in memory requirements is still considerable.\n\nUlysses Sequence parallelism allows us to shard the inputs to the attention computation along the sequence dimension and compute the attention normally, but using only a slice of attention heads on each GPU. With this, we can train models with long sequences, with a few more tools, scaling to 15M+ sequence length. To see how to augment Ulysses SP with TiledMLP, Liger-Kernel, Activation checkpoint offload to cpu and a few other tricks pleae refer to the paper: [Arctic Long Sequence Training: Scalable And Efficient Training For Multi-Million Token Sequences](https://arxiv.org/abs/2506.13996).\n\n## How is Ulysses SP different from FSDP CP\n\nIn the document [Context Parallellism](./context_parallelism.md) you can learn about deploying another technology called Context Parallelism, which too slices on the sequence dimension but uses Ring Attention instead of slicing on the head dimension.\n\nThe following articles go into a very detailed explanation of the differences between the two technologies:\n- https://insujang.github.io/2024-01-11/tensor-parallelism-and-sequence-parallelism-detailed-analysis/\n- https://huggingface.co/blog/exploding-gradients/ulysses-ring-attention\n\nA quick summary adapting from one of the articles:\n- Ulysses SP has a relatively low communication overhead, but is limited by the number of Attention Heads and thus it has certain requirements for network topology (number of attention heads has has to be divisible by the number of participating gpus for a single replica). All-to-all communication is sensitive to latency and it requires Deepspeed.\n- FSDP CP Ring-Attention's P2P ring communication has no aforementioned divisibilty requirements, but has a higher communication volume.\n\nFinally it should be possible to combine SP + CP as explained in the paper [USP: A Unified Sequence Parallelism Approach for Long Context Generative AI](https://arxiv.org/abs/2405.07719) to support an even longer sequence length, albeit this is not yet integrated into šŸ¤—`accelerate`.\n\n\n## Supported sequence parallelism backends\n\nCurrently the only sequence parallelism backend is `deepspeed`, which comes from the modernized Ulysses SP which is part of the [Arctic Long Sequence Training technology](https://arxiv.org/abs/2506.13996). There is also a [tutorial](https://www.deepspeed.ai/tutorials/ulysses-alst-sequence-parallelism/) should you want to integrate it into your own code directly.\n\n## How to use sequence parallelism?\n\n```diff\nfrom accelerate.utils import ParallelismConfig, DeepSpeedSequenceParallelConfig\n\n+# Example: 4 GPUs with sp_size=4, dp_shard_size=1\n+# Ensure: dp_replicate_size Ɨ dp_shard_size Ɨ sp_size = 1 Ɨ 1 Ɨ 4 = 4 GPUs\nparallelism_config = ParallelismConfig(\n+ sp_backend=\"deepspeed\",\n+ sp_size=4,\n+ dp_shard_size=1, # Explicit: no data parallelism\n+ sp_handler=DeepSpeedSequenceParallelConfig(\n+ sp_seq_length_is_variable: true,\n+ sp_attn_implementation=\"sdpa\",\n+ ),\n+ )\n\naccelerator = Accelerator(\n ...,\n parallelism_config=parallelism_config,\n)\n```\n\nAs with any other feature in šŸ¤—`accelerate`, you can enable sequence parallelism also by passing the corresponding flags to `accelerate launch`. In this case, it's no different:\n\n```bash\naccelerate launch --parallelism-config-sp-size 8 ...\n```\n\n> [!Tip]\n> You can also set the `sp_size` and other configuration in the `accelerate config` command, which will save them in your `accelerate` configuration file, so you don't have to pass them every time you launch your script.\n\n> [!Tip]\n> sequence parallelism combines with data parallelism. It doesn't require additional GPUs.\n> So if you have 8 gpus you can do: `--parallelism-config-dp-shard-size 8 --parallelism-config-sp-size 8`. Or you can use the `ParallelismConfig` class to set them programmatically.\n>\n> **Important**: You must ensure `dp_replicate_size Ɨ dp_shard_size Ɨ sp_size = num_processes`. For example, with 8 GPUs and `sp_size=8`, you need `dp_shard_size=1` (since 1 Ɨ 1 Ɨ 8 = 8). With 4 GPUs and `sp_size=2`, you could use `dp_shard_size=2` (since 1 Ɨ 2 Ɨ 2 = 4) for 2D parallelism.\n\n\n## ALST/Ulysses SP backend configuration\n\nALST/UlyssesSP implements sequence parallelism using attention head parallelism, as explained in [this paper](https://arxiv.org/abs/2506.13996). For simplicity, we reuse the concept and setup of sequence parallelism, which, from the user's perspective, is the same: multiple GPUs are used to process a single batch.\n\nTo give a sense of what ALST made possible - it allowed us to train in bf16 with 500K tokens on a single H100 GPU, 3.7M on a single node, and 15M on Llama-8B using just four nodes. This feature of HF Accelerate enables only 1 of the 3 ALST components, so the achievable sequence length will be smaller. You'd want TiledMLP, Activation checkpoint offload to CPU, and a few other things enabled to get the full power of ALST. For details, please refer to [this tutorial](https://www.deepspeed.ai/tutorials/ulysses-alst-sequence-parallelism/).\n\nTo configure the `deepspeed` backend:\n\n```python\n# Example: 4 GPUs with sp_size=4, dp_shard_size=1\n# Ensure: dp_replicate_size Ɨ dp_shard_size Ɨ sp_size = 1 Ɨ 1 Ɨ 4 = 4 GPUs\nparallelism_config = ParallelismConfig(\n sp_backend=\"deepspeed\",\n sp_size=4,\n dp_shard_size=1, # Explicit: no data parallelism\n sp_handler=DeepSpeedSequenceParallelConfig(\n sp_seq_length=256,\n sp_seq_length_is_variable=True,\n sp_attn_implementation=\"sdpa\",\n ),\n)\naccelerator = Accelerator(\n ...,\n parallelism_config=parallelism_config,\n)\n```\n\n- `sp_backend`: set to `deepspeed` here\n- `sp_size` is the degree of the sequence parallelism - in the above example it's 4, therefore 4 gpus will be used to process a single batch (while doing DP=4 over the same gpus)\n- `sp_seq_length` and `sp_seq_length_is_variable` are used to deal with sequence lengths. If `sp_seq_length_is_variable=True` the backend will work with a sequence length that may change between batches, in which case `sp_seq_length` value can be set to anything divisible by the sequence parallel degree or not set at all. In this case on every `forward` the sequence variables will be derived from input. If `False` then `seq_length` needs to match the batch's sequence length dimension, which then will have to be padded to be always the same. The default is `True`.\n- `sp_attn_implementation` is one of `sdpa`, `flash_attention_2` or `flash_attention_3`. This sequence parallel implementation uses `position_ids` instead of `attention_mask` therefore, `eager` can't work here until it supports working with `position_ids`. Also, please note that `sdpa` doesn't handle multiple samples combined into one correctly; it will attend to the whole sample as one. If the samples aren't combined, `sdpa` will work correctly. Therefore, Flash Attention should be the ideal choice as it always works.\n\nInstead of setting these values in `DeepSpeedSequenceParallelConfig` object, you can also use the environment variables to accomplish the same - here they are correspondingly to the end of the list above.\n- `PARALLELISM_CONFIG_SP_BACKEND`\n- `PARALLELISM_CONFIG_SP_SEQ_LENGTH`\n- `PARALLELISM_CONFIG_SP_SEQ_LENGTH_IS_VARIABLE`\n- `PARALLELISM_CONFIG_SP_ATTN_IMPLEMENTATION`\n\nIf not passed in the code, `sp_size` can be set via `--parallelism_config_sp_size` CLI argument. Same for other arguments. You can also do the accelerate config file style config, e.g., for 2 GPUs:\n\n```yaml\ndistributed_type: DEEPSPEED\ndeepspeed_config:\n deepspeed_config_file: path/to/ds_config.json\nmachine_rank: 0\nnum_machines: 1\nnum_processes: 2\nparallelism_config:\n parallelism_config_dp_replicate_size: 1\n parallelism_config_dp_shard_size: 1 # Must satisfy: 1 Ɨ 1 Ɨ 2 = 2 num_processes\n parallelism_config_sp_size: 2\n parallelism_config_sp_backend: deepspeed\n parallelism_config_sp_seq_length_is_variable: true\n parallelism_config_sp_attn_implementation: sdpa\n\n```\n\nAs mentioned earlier Ulysses sequence parallelism is normally overlayed with data parallelism - same ranks are used for feeding unique data streams and also perform Ulysses Sequence Parallelism. But you could also create replicas like so:\n\n```python\n# Example: 4 GPUs with 2D parallelism (SP=2, DP=2)\n# Ensure: dp_replicate_size Ɨ dp_shard_size Ɨ sp_size = 2 Ɨ 1 Ɨ 2 = 4 GPUs\nparallelism_config = ParallelismConfig(\n dp_replicate_size=2,\n dp_shard_size=1, # Explicit: no sharding within replicas\n sp_size=2,\n sp_backend=\"deepspeed\",\n sp_handler=DeepSpeedSequenceParallelConfig(...),\n)\n```\nHere we use 4 gpus, with 2 sequence parallelism replicas. Deepspeed-ZeRO is what drives the data parallelism here.\n\nPlease note that a lot of magic is hidden inside [UlyssesSPDataLoaderAdapter](https://github.com/deepspeedai/DeepSpeed/blob/64c0052fa08438b4ecf4cae30af15091a92d2108/deepspeed/runtime/sequence_parallel/ulysses_sp.py#L442). It's used behind the scenes, wrapping your original DataLoader object, but you should be aware of it should you run into any problems. It also automatically injects the correct `shift_labels` into the batch dictionary, before the batch gets sharded across the participating ranks.\n\nNow the only remaining piece to start using ALST/UlyssesSP is to aggregate the loss across ranks using a differentiable `all_gather` to get the grads right. The following code does it, while also excluding any masked out with `-100` tokens, to get the correct average:\n\n```python\nsp_size = parallelism_config.sp_size if parallelism_config is not None else 1\nif sp_size > 1:\n sp_group = accelerator.torch_device_mesh[\"sp\"].get_group()\n sp_world_size = parallelism_config.sp_size\n\n# Normal training loop\nfor iter, batch in enumerate(dl):\n optimizer.zero_grad()\n\n batch = move_to_device(batch, model.device)\n\n # The model automatically receives shift_labels via **kwargs and uses it for loss computation.\n # Both standard transformers models and Liger-patched models handle this correctly.\n outputs = model(**batch)\n loss = outputs.loss\n shift_labels = batch[\"shift_labels\"]\n\n if sp_size > 1:\n # differentiable weighted per-shard-loss aggregation across ranks\n losses_per_rank = torch.distributed.nn.functional.all_gather(loss, group=sp_group)\n # special dealing with SFT that has prompt tokens that aren't used in loss computation\n good_tokens = (shift_labels != -100).view(-1).sum()\n good_tokens_per_rank = torch.distributed.nn.functional.all_gather(\n good_tokens, group=sp_group\n )\n # Skip ranks with zero valid tokens to avoid NaN contamination (NaN * 0 = NaN)\n total_loss = sum(\n losses_per_rank[rank] * good_tokens_per_rank[rank]\n for rank in range(sp_world_size)\n if good_tokens_per_rank[rank] > 0\n )\n total_good_tokens = sum(good_tokens_per_rank)\n loss = total_loss / max(total_good_tokens, 1)\n\n if rank == 0: accelerator.print(f\"{iter}: {loss=}\")\n accelerator.log(dict(train_loss=loss, step=iter))\n\n accelerator.backward(loss)\n optimizer.step()\n```\n\nNote that models automatically handle `shift_labels` when it's present in the batch. The model's forward pass receives `shift_labels` via `**kwargs` and passes it to the loss function, which correctly computes the loss for sequence parallelism. If you use [Liger Kernel](https://github.com/linkedin/Liger-Kernel), it also handles `shift_labels` seamlessly and computes loss in a very memory-efficient way. Liger is highly recommended for long sequence lengths, as it liberates GPU memory by using fused operations (e.g., fused logit-loss computation that never materializes the full logits tensor in memory).\n\nIf you want to see what HF Accelerate did behind the scenes please read [this full integration tutorial](https://www.deepspeed.ai/tutorials/ulysses-alst-sequence-parallelism/).\n\nFor an example of an Accelerate training loop with enabled ALST/UlyssesSP see [examples/alst_ulysses_sequence_parallelism](https://github.com/huggingface/accelerate/blob/main/examples/alst_ulysses_sequence_parallelism).\n\n[!Warning]\n> This API is quite new and still in its experimental stage. While we strive to provide a stable API, some small parts of the public API may change in the future.\n\nSince this is a Deepspeed backend the usual Deepspeed configuration applies, so you can combine sequence parallelism with optimizer states and/or weights offloading as well to liberate more gpu memory and enable an even longer sequence length. This technology has been tested to work with DeepSpeed ZeRO stage 2 and 3.\n\n" }, { "path": "docs/source/concept_guides/training_tpu.md", "content": "\n\n# Training on TPUs\n\nTraining on TPUs can be slightly different from training on multi-gpu, even with Accelerate. This guide aims to show you \nwhere you should be careful and why, as well as the best practices in general.\n\n## Training in a Notebook\n\nThe main carepoint when training on TPUs comes from the [`notebook_launcher`]. As mentioned in the [notebook tutorial](../usage_guides/notebook), you need to \nrestructure your training code into a function that can get passed to the [`notebook_launcher`] function and be careful about not declaring any tensors on the GPU.\n\nWhile on a TPU that last part is not as important, a critical part to understand is that when you launch code from a notebook you do so through a process called **forking**. \nWhen 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 \nutilizing a python process, you need to *fork* a new process from it to launch your code. \n\nWhere 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 \ntraining 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 \nmodel 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\non Google Colaboratory. \n\nBelow is an example of a training function passed to the [`notebook_launcher`] if training on CPUs or GPUs:\n\n\n\n This code snippet is based off the one from the `simple_nlp_example` notebook found [here](https://github.com/huggingface/notebooks/blob/main/examples/accelerate_examples/simple_nlp_example.ipynb) with slight \n modifications for the sake of simplicity\n\n\n\n```python\ndef training_function():\n # Initialize accelerator\n accelerator = Accelerator()\n model = AutoModelForSequenceClassification.from_pretrained(\"bert-base-cased\", num_labels=2)\n train_dataloader, eval_dataloader = create_dataloaders(\n train_batch_size=hyperparameters[\"train_batch_size\"], eval_batch_size=hyperparameters[\"eval_batch_size\"]\n )\n\n # Instantiate optimizer\n optimizer = AdamW(params=model.parameters(), lr=hyperparameters[\"learning_rate\"])\n\n # Prepare everything\n # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the\n # prepare method.\n model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(\n model, optimizer, train_dataloader, eval_dataloader\n )\n\n num_epochs = hyperparameters[\"num_epochs\"]\n # Now we train the model\n for epoch in range(num_epochs):\n model.train()\n for step, batch in enumerate(train_dataloader):\n outputs = model(**batch)\n loss = outputs.loss\n accelerator.backward(loss)\n\n optimizer.step()\n optimizer.zero_grad()\n```\n\n```python\nfrom accelerate import notebook_launcher\n\nnotebook_launcher(training_function)\n```\n\n\n\n The `notebook_launcher` will default to 8 processes if Accelerate has been configured for a TPU\n\n\n\nIf you use this example and declare the model *inside* the training loop, then on a low-resource system you will potentially see an error \nlike:\n\n```\nProcessExitedException: process 0 terminated with signal SIGSEGV\n```\n\nThis error is *extremely* cryptic but the basic explanation is you ran out of system RAM. You can avoid this entirely by reconfiguring the training function to \naccept a single `model` argument, and declare it in an outside cell:\n\n```python\n# In another Jupyter cell\nmodel = AutoModelForSequenceClassification.from_pretrained(\"bert-base-cased\", num_labels=2)\n```\n\n```diff\n+ def training_function(model):\n # Initialize accelerator\n accelerator = Accelerator()\n- model = AutoModelForSequenceClassification.from_pretrained(\"bert-base-cased\", num_labels=2)\n train_dataloader, eval_dataloader = create_dataloaders(\n train_batch_size=hyperparameters[\"train_batch_size\"], eval_batch_size=hyperparameters[\"eval_batch_size\"]\n )\n ...\n```\n\nAnd finally calling the training function with:\n\n```diff\n from accelerate import notebook_launcher\n- notebook_launcher(training_function)\n+ notebook_launcher(training_function, (model,))\n```\n\n\n\n The above workaround is only needed when launching a TPU instance from a Jupyter Notebook on a low-resource server such as Google Colaboratory or Kaggle. If \n using a script or launching on a much beefier server declaring the model beforehand is not needed.\n\n\n\n## Mixed Precision and Global Variables \n\nAs mentioned in the [mixed precision tutorial](../usage_guides/mixed_precision), Accelerate supports fp16 and bf16, both of which can be used on TPUs.\nThat being said, ideally `bf16` should be utilized as it is extremely efficient to use.\n\nThere are two \"layers\" when using `bf16` and Accelerate on TPUs, at the base level and at the operation level. \n\nAt the base level, this is enabled when passing `mixed_precision=\"bf16\"` to `Accelerator`, such as:\n```python\naccelerator = Accelerator(mixed_precision=\"bf16\")\n```\nBy default, this will cast `torch.float` and `torch.double` to `bfloat16` on TPUs. \nThe specific configuration being set is an environmental variable of `XLA_USE_BF16` is set to `1`.\n\nThere is a further configuration you can perform which is setting the `XLA_DOWNCAST_BF16` environmental variable. If set to `1`, then \n`torch.float` is `bfloat16` and `torch.double` is `float32`.\n\nThis is performed in the `Accelerator` object when passing `downcast_bf16=True`:\n```python\naccelerator = Accelerator(mixed_precision=\"bf16\", downcast_bf16=True)\n```\n\nUsing downcasting instead of bf16 everywhere is good for when you are trying to calculate metrics, log values, and more where raw bf16 tensors would be unusable. \n\n## Training Times on TPUs\n\nAs you launch your script, you may notice that training seems exceptionally slow at first. This is because TPUs\nfirst run through a few batches of data to see how much memory to allocate before finally utilizing this configured \nmemory allocation extremely efficiently. \n\nIf you notice that your evaluation code to calculate the metrics of your model takes longer due to a larger batch size being used, \nit is recommended to keep the batch size the same as the training data if it is too slow. Otherwise the memory will reallocate to this \nnew batch size after the first few iterations. \n\n\n\n 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.\n\n\n" }, { "path": "docs/source/index.md", "content": "\n\n# Accelerate\n\nAccelerate is a library that enables the same PyTorch code to be run across any distributed configuration by adding just four lines of code! In short, training and inference at scale made simple, efficient and adaptable.\n\n```diff\n+ from accelerate import Accelerator\n+ accelerator = Accelerator()\n\n+ model, optimizer, training_dataloader, scheduler = accelerator.prepare(\n+ model, optimizer, training_dataloader, scheduler\n+ )\n\n for batch in training_dataloader:\n optimizer.zero_grad()\n inputs, targets = batch\n inputs = inputs.to(device)\n targets = targets.to(device)\n outputs = model(inputs)\n loss = loss_function(outputs, targets)\n+ accelerator.backward(loss)\n optimizer.step()\n scheduler.step()\n```\n\nBuilt on `torch_xla` and `torch.distributed`, Accelerate takes care of the heavy lifting, so you don't have to write any custom code to adapt to these platforms.\nConvert existing codebases to utilize [DeepSpeed](usage_guides/deepspeed), perform [fully sharded data parallelism](usage_guides/fsdp), and have automatic support for mixed-precision training! \n\n \n\n To get a better idea of this process, make sure to check out the [Tutorials](basic_tutorials/overview)! \n\n\n\n\nThis code can then be launched on any system through Accelerate's CLI interface:\n```bash\naccelerate launch {my_script.py}\n```\n\n
\n
\n
Tutorials
\n

Learn the basics and become familiar with using Accelerate. Start here if you are using Accelerate for the first time!

\n \n
How-to guides
\n

Practical guides to help you achieve a specific goal. Take a look at these guides to learn how to use Accelerate to solve real-world problems.

\n \n
Conceptual guides
\n

High-level explanations for building a better understanding of important topics such as avoiding subtle nuances and pitfalls in distributed training and DeepSpeed.

\n \n
Reference
\n

Technical descriptions of how Accelerate classes and methods work.

\n \n
\n
\n" }, { "path": "docs/source/package_reference/accelerator.md", "content": "\n\n# Accelerator\n\nThe [`Accelerator`] is the main class for enabling distributed training on any type of training setup. Read the [Add Accelerator to your code](../basic_tutorials/migration) tutorial to learn more about how to add the [`Accelerator`] to your script.\n\n## Accelerator[[api]]\n\n[[autodoc]] Accelerator\n\n## Utilities\n\n[[autodoc]] accelerate.utils.gather_object\n" }, { "path": "docs/source/package_reference/big_modeling.md", "content": "\n\n# Working with large models\n\n## Dispatch and offload\n\n### init_empty_weights\n\n[[autodoc]] big_modeling.init_empty_weights\n\n### cpu_offload\n\n[[autodoc]] big_modeling.cpu_offload\n\n### cpu_offload_with_hook\n\n[[autodoc]] big_modeling.cpu_offload_with_hook\n\n### disk_offload\n\n[[autodoc]] big_modeling.disk_offload\n\n### dispatch_model\n\n[[autodoc]] big_modeling.dispatch_model\n\n### load_checkpoint_and_dispatch\n\n[[autodoc]] big_modeling.load_checkpoint_and_dispatch\n\n### load_checkpoint_in_model\n\n[[autodoc]] big_modeling.load_checkpoint_in_model\n\n### infer_auto_device_map\n\n[[autodoc]] utils.infer_auto_device_map\n\n## Hooks\n\n### ModelHook\n\n[[autodoc]] hooks.ModelHook\n\n### AlignDevicesHook\n\n[[autodoc]] hooks.AlignDevicesHook\n\n### SequentialHook\n\n[[autodoc]] hooks.SequentialHook\n\n### LayerwiseCastingHook\n\n[[autodoc]] hooks.LayerwiseCastingHook\n\n## Adding Hooks\n\n### add_hook_to_module\n\n[[autodoc]] hooks.add_hook_to_module\n\n### attach_execution_device_hook\n\n[[autodoc]] hooks.attach_execution_device_hook\n\n### attach_align_device_hook\n\n[[autodoc]] hooks.attach_align_device_hook\n\n### attach_align_device_hook_on_blocks\n\n[[autodoc]] hooks.attach_align_device_hook_on_blocks\n\n### attach_layerwise_casting_hooks\n\n[[autodoc]] big_modeling.attach_layerwise_casting_hooks\n\n## Removing Hooks\n\n### remove_hook_from_module\n\n[[autodoc]] hooks.remove_hook_from_module\n\n### remove_hook_from_submodules\n\n[[autodoc]] hooks.remove_hook_from_submodules\n\n## Utilities\n\n### has_offloaded_params\n\n[[autodoc]] utils.has_offloaded_params\n\n### align_module_device\n\n[[autodoc]] utils.align_module_device\n" }, { "path": "docs/source/package_reference/cli.md", "content": "\n\n# The Command Line \n\nBelow is a list of all the available commands šŸ¤— Accelerate with their parameters\n\n## accelerate config\n\n**Command**:\n\n`accelerate config` or `accelerate-config`\n\nLaunches a series of prompts to create and save a `default_config.yml` configuration file for your training system. Should \nalways be ran first on your machine.\n\n**Usage**: \n\n```bash\naccelerate config [arguments]\n```\n\n**Optional Arguments**:\n* `--config_file CONFIG_FILE` (`str`) -- The path to use to store the config file. Will default to a file named default_config.yaml in the cache location, which is the content\n of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have such an environment variable, your cache directory\n (`~/.cache` or the content of `XDG_CACHE_HOME`) suffixed with `huggingface`.\n* `-h`, `--help` (`bool`) -- Show a help message and exit\n\n## accelerate config default\n\n**Command**:\n\n`accelerate config default` or `accelerate-config default`\n\nCreate a default config file for Accelerate with only a few flags set.\n\n**Usage**: \n\n```bash\naccelerate config default [arguments]\n```\n\n**Optional Arguments**:\n* `--config_file CONFIG_FILE` (`str`) -- The path to use to store the config file. Will default to a file named default_config.yaml in the cache location, which is the content\n of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have such an environment variable, your cache directory\n (`~/.cache` or the content of `XDG_CACHE_HOME`) suffixed with `huggingface`.\n\n* `-h`, `--help` (`bool`) -- Show a help message and exit\n* `--mixed_precision {no,fp16,bf16}` (`str`) -- Whether or not to use mixed precision training. Choose between FP16 and BF16 (bfloat16) training. BF16 training is only supported on Nvidia Ampere GPUs and PyTorch 1.10 or later.\n\n## accelerate config update\n\n**Command**:\n\n`accelerate config update` or `accelerate-config update`\n\nUpdate an existing config file with the latest defaults while maintaining the old configuration.\n\n**Usage**: \n\n```bash\naccelerate config update [arguments]\n```\n\n**Optional Arguments**:\n* `--config_file CONFIG_FILE` (`str`) -- The path to the config file to update. Will default to a file named default_config.yaml in the cache location, which is the content\n of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have such an environment variable, your cache directory\n (`~/.cache` or the content of `XDG_CACHE_HOME`) suffixed with `huggingface`.\n\n* `-h`, `--help` (`bool`) -- Show a help message and exit\n\n\n## accelerate env\n\n**Command**:\n\n`accelerate env` or `accelerate-env` or `python -m accelerate.commands.env`\n\nLists 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).\n\n**Usage**:\n\n```bash\naccelerate env [arguments]\n```\n\n**Optional Arguments**:\n* `--config_file CONFIG_FILE` (`str`) -- The path to use to store the config file. Will default to a file named default_config.yaml in the cache location, which is the content\n of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have such an environment variable, your cache directory\n (`~/.cache` or the content of `XDG_CACHE_HOME`) suffixed with `huggingface`.\n* `-h`, `--help` (`bool`) -- Show a help message and exit\n\n## accelerate launch\n\n**Command**:\n\n`accelerate launch` or `accelerate-launch` or `python -m accelerate.commands.launch`\n\nLaunches a specified script on a distributed system with the right parameters.\n\n**Usage**: \n\n```bash\naccelerate launch [arguments] {training_script} --{training_script-argument-1} --{training_script-argument-2} ...\n```\n\n**Positional Arguments**:\n\n- `{training_script}` -- The full path to the script to be launched in parallel\n- `--{training_script-argument-1}` -- Arguments of the training script\n\n**Optional Arguments**:\n\n* `-h`, `--help` (`bool`) -- Show a help message and exit\n* `--config_file CONFIG_FILE` (`str`)-- The config file to use for the default values in the launching script.\n* `-m`, `--module` (`bool`) -- Change each process to interpret the launch script as a Python module, executing with the same behavior as 'python -m'.\n* `--no_python` (`bool`) -- Skip prepending the training script with 'python' - just execute it directly. Useful when the script is not a Python script.\n* `--debug` (`bool`) -- Whether to print out the torch.distributed stack trace when something fails.\n* `-q`, `--quiet` (`bool`) -- Silence subprocess errors from the launch stack trace to only show the relevant tracebacks. (Only applicable to DeepSpeed and single-process configurations).\n\n\nThe rest of these arguments are configured through `accelerate config` and are read in from the specified `--config_file` (or default configuration) for their \nvalues. They can also be passed in manually.\n\n**Hardware Selection Arguments**:\n\n* `--cpu` (`bool`) -- Whether or not to force the training on the CPU.\n* `--multi_gpu` (`bool`) -- Whether or not this should launch a distributed GPU training.\n* `--tpu` (`bool`) -- Whether or not this should launch a TPU training.\n\n**Resource Selection Arguments**:\n\nThe following arguments are useful for fine-tuning how available hardware should be used\n\n* `--mixed_precision {no,fp16,bf16,fp8}` (`str`) -- Whether or not to use mixed precision training. Choose between FP16 and BF16 (bfloat16) training. BF16 training is only supported on Nvidia Ampere GPUs and PyTorch 1.10 or later.\n* `--num_processes NUM_PROCESSES` (`int`) -- The total number of processes to be launched in parallel.\n* `--num_machines NUM_MACHINES` (`int`) -- The total number of machines used in this training.\n* `--num_cpu_threads_per_process NUM_CPU_THREADS_PER_PROCESS` (`int`) -- The number of CPU threads per process. Can be tuned for optimal performance.\n* `--enable_cpu_affinity` (`bool`) -- Whether or not CPU affinity and balancing should be enabled. Currently only supported on NVIDIA hardware.\n\n**Training Paradigm Arguments**:\n\nThe following arguments are useful for selecting which training paradigm to use.\n\n* `--use_deepspeed` (`bool`) -- Whether or not to use DeepSpeed for training.\n* `--use_fsdp` (`bool`) -- Whether or not to use FullyShardedDataParallel for training.\n* `--use_megatron_lm` (`bool`) -- Whether or not to use Megatron-LM for training.\n\n**Distributed GPU Arguments**:\n\nThe following arguments are only useful when `multi_gpu` is passed or multi-gpu training is configured through `accelerate config`: \n\n* `--gpu_ids` (`str`) -- What GPUs (by id) should be used for training on this machine as a comma-separated list\n* `--same_network` (`bool`) -- Whether all machines used for multinode training exist on the same local network.\n* `--machine_rank` (`int`) -- The rank of the machine on which this script is launched.\n* `--main_process_ip` (`str`) -- The IP address of the machine of rank 0.\n* `--main_process_port` (`int`) -- The port to use to communicate with the machine of rank 0.\n* `-t`, `--tee` (`str`) -- Tee std streams into a log file and also to console.\n* `--log_dir` (`str`) -- Base directory to use for log files when using torchrun/torch.distributed.run as launcher. Use with --tee to redirect std streams info log files.\n* `--role` (`str`) -- User-defined role for the workers.\n* `--rdzv_backend` (`str`) -- The rendezvous method to use, such as 'static' (the default) or 'c10d'\n* `--rdzv_conf` (`str`) -- Additional rendezvous configuration (=,=,...).\n* `--max_restarts` (`int`) -- Maximum number of worker group restarts before failing.\n* `--monitor_interval` (`int`) -- Interval, in seconds, to monitor the state of workers.\n\n**TPU Arguments**:\n\nThe following arguments are only useful when `tpu` is passed or TPU training is configured through `accelerate config`: \n\n* `--tpu_cluster` (`bool`) -- Whether to use a GCP TPU pod for training.\n* `--tpu_use_sudo` (`bool`) -- Whether to use `sudo` when running the TPU training script in each pod.\n* `--vm` (`str`) -- List of single Compute VM instance names. If not provided we assume usage of instance groups. For TPU pods.\n* `--env` (`str`) -- List of environment variables to set on the Compute VM instances. For TPU pods.\n* `--main_training_function` (`str`) -- The name of the main function to be executed in your script (only for TPU training).\n* `--downcast_bf16` (`bool`) -- Whether when using bf16 precision on TPUs if both float and double tensors are cast to bfloat16 or if double tensors remain as float32.\n\n**DeepSpeed Arguments**:\n\nThe following arguments are only useful when `use_deepspeed` is passed or `deepspeed` is configured through `accelerate config`: \n\n* `--deepspeed_config_file` (`str`) -- DeepSpeed config file.\n* `--zero_stage` (`int`) -- DeepSpeed's ZeRO optimization stage.\n* `--offload_optimizer_device` (`str`) -- Decides where (none|cpu|nvme) to offload optimizer states.\n* `--offload_param_device` (`str`) -- Decides where (none|cpu|nvme) to offload parameters.\n* `--offload_optimizer_nvme_path` (`str`) -- Decides Nvme Path to offload optimizer states.\n* `--gradient_accumulation_steps` (`int`) -- No of gradient_accumulation_steps used in your training script.\n* `--gradient_clipping` (`float`) -- Gradient clipping value used in your training script.\n* `--zero3_init_flag` (`str`) -- Decides Whether (true|false) to enable `deepspeed.zero.Init` for constructing massive models. Only applicable with DeepSpeed ZeRO Stage-3.\n* `--zero3_save_16bit_model` (`str`) -- Decides Whether (true|false) to save 16-bit model weights when using ZeRO Stage-3. Only applicable with DeepSpeed ZeRO Stage-3.\n* `--deepspeed_hostfile` (`str`) -- DeepSpeed hostfile for configuring multi-node compute resources.\n* `--deepspeed_exclusion_filter` (`str`) -- DeepSpeed exclusion filter string when using multi-node setup.\n* `--deepspeed_inclusion_filter` (`str`) -- DeepSpeed inclusion filter string when using multi-node setup.\n* `--deepspeed_multinode_launcher` (`str`) -- DeepSpeed multi-node launcher to use.\n* `--deepspeed_moe_layer_cls_names` (`str`) -- comma-separated list of transformer MoE layer class names (case-sensitive) to wrap, e.g, `MixtralSparseMoeBlock` `Qwen2MoeSparseMoeBlock`, `JetMoEAttention,JetMoEBlock`\n\n**Fully Sharded Data Parallelism Arguments**:\n\nThe following arguments are only useful when `use_fsdp` is passed or Fully Sharded Data Parallelism is configured through `accelerate config`:\n\n* `--fsdp_offload_params` (`str`) -- Decides Whether (true|false) to offload parameters and gradients to CPU.\n* `--fsdp_min_num_params` (`int`) -- FSDP's minimum number of parameters for Default Auto Wrapping.\n* `--fsdp_sharding_strategy` (`int`) -- FSDP's Sharding Strategy.\n* `--fsdp_auto_wrap_policy` (`str`) -- FSDP's auto wrap policy.\n* `--fsdp_transformer_layer_cls_to_wrap` (`str`) -- Transformer layer class name (case-sensitive) to wrap, e.g, `BertLayer`, `GPTJBlock`, `T5Block` ...\n* `--fsdp_backward_prefetch_policy` (`str`) -- FSDP's backward prefetch policy.\n* `--fsdp_state_dict_type` (`str`) -- FSDP's state dict type.\n* `--fsdp_forward_prefetch` (`str`) -- FSDP forward prefetch.\n* `--fsdp_use_orig_params` (`str`) -- If True, allows non-uniform `requires_grad` mixed in a FSDP unit.\n* `--fsdp_cpu_ram_efficient_loading` (`str`) -- If true, only the first process loads the pretrained model checkoint while all other processes have empty weights. When using this, `--fsdp_sync_module_states` needs to True.\n* `--fsdp_sync_module_states` (`str`) -- If true, each individually wrapped FSDP unit will broadcast module parameters from rank 0.\n* `--fsdp_activation_checkpointing` (`bool`) -- Decides Whether intermediate activations are freed during the forward pass, and a checkpoint is left as a placeholder\n\n**Megatron-LM Arguments**:\n\nThe following arguments are only useful when `use_megatron_lm` is passed or Megatron-LM is configured through `accelerate config`:\n\n* `--megatron_lm_tp_degree` (``) -- Megatron-LM's Tensor Parallelism (TP) degree.\n* `--megatron_lm_pp_degree` (``) -- Megatron-LM's Pipeline Parallelism (PP) degree.\n* `--megatron_lm_num_micro_batches` (``) -- Megatron-LM's number of micro batches when PP degree > 1.\n* `--megatron_lm_sequence_parallelism` (``) -- Decides Whether (true|false) to enable Sequence Parallelism when TP degree > 1.\n* `--megatron_lm_recompute_activations` (``) -- Decides Whether (true|false) to enable Selective Activation Recomputation.\n* `--megatron_lm_use_distributed_optimizer` (``) -- Decides Whether (true|false) to use distributed optimizer which shards optimizer state and gradients across Data Parallel (DP) ranks.\n* `--megatron_lm_gradient_clipping` (``) -- Megatron-LM's gradient clipping value based on global L2 Norm (0 to disable).\n\n**FP8 Arguments**:\n\n* `--fp8_backend` (`str`) -- Choose a backend to train with FP8 (`te` or `msamp`)\n* `--fp8_use_autocast_during_eval` (`bool`) -- Whether to use FP8 autocast during eval mode (useful only when `--fp8_backend=te` is passed). Generally better metrics are found when this is not passed.\n* `--fp8_margin` (`int`) -- The margin to use for the gradient scaling (useful only when `--fp8_backend=te` is passed).\n* `--fp8_interval` (`int`) -- The interval to use for how often the scaling factor is recomputed (useful only when `--fp8_backend=te` is passed).\n* `--fp8_format` (`str`) -- The format to use for the FP8 recipe (useful only when `--fp8_backend=te` is passed).\n* `--fp8_amax_history_len` (`int`) -- The length of the history to use for the scaling factor computation (useful only when `--fp8_backend=te` is passed).\n* `--fp8_amax_compute_algo` (`str`) -- The algorithm to use for the scaling factor computation. (useful only when `--fp8_backend=te` is passed).\n* `--fp8_override_linear_precision` (`Tuple[bool, bool, bool]`) -- Whether or not to execute `fprop`, `dgrad`, and `wgrad` GEMMS in higher precision.\n* `--fp8_opt_level` (`str`) -- What level of 8-bit collective communication should be used with MS-AMP (useful only when `--fp8_backend=msamp` is passed)\n\n**AWS SageMaker Arguments**:\n\nThe following arguments are only useful when training in SageMaker\n\n* `--aws_access_key_id AWS_ACCESS_KEY_ID` (`str`) -- The AWS_ACCESS_KEY_ID used to launch the Amazon SageMaker training job\n* `--aws_secret_access_key AWS_SECRET_ACCESS_KEY` (`str`) -- The AWS_SECRET_ACCESS_KEY used to launch the Amazon SageMaker training job\n\n## accelerate estimate-memory\n\n**Command**:\n\n`accelerate estimate-memory` or `accelerate-estimate-memory` or `python -m accelerate.commands.estimate`\n\nEstimates 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. \n\n\n\n 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.\n\n\n\n**Usage**: \n\n```bash\naccelerate estimate-memory {MODEL_NAME} --library_name {LIBRARY_NAME} --dtypes {dtype_1} {dtype_2} ...\n```\n\n**Required Arguments**:\n\n* `MODEL_NAME` (`str`)-- The model name on the Hugging Face Hub\n\n**Optional Arguments**:\n\n* `--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\n* `--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`\n* `--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.\n\n## accelerate tpu-config\n\n`accelerate tpu-config`\n\n**Usage**:\n\n```bash\naccelerate tpu-config [arguments]\n```\n\n**Optional Arguments**:\n* `-h`, `--help` (`bool`) -- Show a help message and exit\n\n**Config Arguments**:\n\nArguments that can be configured through `accelerate config`.\n\n* `--config_file` (`str`) -- Path to the config file to use for accelerate.\n* `--tpu_name` (`str`) -- The name of the TPU to use. If not specified, will use the TPU specified in the config file.\n* `--tpu_zone` (`str`) -- The zone of the TPU to use. If not specified, will use the zone specified in the config file.\n\n**TPU Arguments**:\n\nArguments for options ran inside the TPU.\n\n* `--command_file` (`str`) -- The path to the file containing the commands to run on the pod on startup.\n* `--command` (`str`) -- A command to run on the pod. Can be passed multiple times.\n* `--install_accelerate` (`bool`) -- Whether to install accelerate on the pod. Defaults to False.\n* `--accelerate_version` (`str`) -- The version of accelerate to install on the pod. If not specified, will use the latest pypi version. Specify 'dev' to install from GitHub.\n* `--debug` (`bool`) -- If set, will print the command that would be run instead of running it.\n\n## accelerate test\n\n`accelerate test` or `accelerate-test`\n\nRuns `accelerate/test_utils/test_script.py` to verify that šŸ¤— Accelerate has been properly configured on your system and runs. \n\n**Usage**: \n\n```bash\naccelerate test [arguments]\n```\n\n**Optional Arguments**:\n* `--config_file CONFIG_FILE` (`str`) -- The path to use to store the config file. Will default to a file named default_config.yaml in the cache location, which is the content\n of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have such an environment variable, your cache directory\n (`~/.cache` or the content of `XDG_CACHE_HOME`) suffixed with `huggingface`.\n* `-h`, `--help` (`bool`) -- Show a help message and exit\n" }, { "path": "docs/source/package_reference/deepspeed.md", "content": "\n\n# DeepSpeed utilities\n\n## DeepSpeedPlugin\n\n## get_active_deepspeed_plugin\n\n[[autodoc]] utils.get_active_deepspeed_plugin\n\n[[autodoc]] utils.DeepSpeedPlugin\n\n[[autodoc]] utils.deepspeed.DummyScheduler\n\n## DeepSpeedEnginerWrapper\n\n[[autodoc]] utils.deepspeed.DeepSpeedEngineWrapper\n\n## DeepSpeedOptimizerWrapper\n\n[[autodoc]] utils.deepspeed.DeepSpeedOptimizerWrapper\n\n## DeepSpeedSchedulerWrapper\n\n[[autodoc]] utils.deepspeed.DeepSpeedSchedulerWrapper\n\n## DummyOptim\n\n[[autodoc]] utils.deepspeed.DummyOptim\n\n## DummyScheduler" }, { "path": "docs/source/package_reference/fp8.md", "content": "\n\n# FP8\n\nBelow are functions and classes relative to the underlying FP8 implementation\n\n## FP8RecipeKwargs\n\n[[autodoc]] utils.FP8RecipeKwargs\n\n## convert_model\n\n[[autodoc]] utils.convert_model\n\n## has_transformer_engine_layers\n\n[[autodoc]] utils.has_transformer_engine_layers\n\n## contextual_fp8_autocast\n\n[[autodoc]] utils.contextual_fp8_autocast\n\n## apply_fp8_autowrap\n\n[[autodoc]] utils.apply_fp8_autowrap\n" }, { "path": "docs/source/package_reference/fsdp.md", "content": "\n\n# Fully Sharded Data Parallel utilities\n\n## enable_fsdp_ram_efficient_loading\n\n[[autodoc]] utils.enable_fsdp_ram_efficient_loading\n\n## disable_fsdp_ram_efficient_loading\n\n[[autodoc]] utils.disable_fsdp_ram_efficient_loading\n\n## merge_fsdp_weights\n\n[[autodoc]] utils.merge_fsdp_weights\n\n## FullyShardedDataParallelPlugin\n\n[[autodoc]] utils.FullyShardedDataParallelPlugin\n\n## fsdp2_load_full_state_dict\n\n[[autodoc]] utils.fsdp2_load_full_state_dict\n\n## fsdp2_switch_optimizer_parameters\n\n[[autodoc]] utils.fsdp2_switch_optimizer_parameters\n\n## fsdp2_prepare_model\n\n[[autodoc]] utils.fsdp2_prepare_model\n\n## fsdp2_prepare_auto_wrap_policy\n" }, { "path": "docs/source/package_reference/inference.md", "content": "\n\n# Pipeline parallelism\n\nAccelerate supports pipeline parallelism for large-scale training with the PyTorch [torch.distributed.pipelining](https://pytorch.org/docs/stable/distributed.pipelining.html) API.\n\n## prepare_pippy\n\n[[autodoc]] inference.prepare_pippy\n" }, { "path": "docs/source/package_reference/kwargs.md", "content": "\n\n# Kwargs handlers\n\nThe following objects can be passed to the main [`Accelerator`] to customize how some PyTorch objects\nrelated to distributed training or mixed precision are created.\n\n## AutocastKwargs\n\n[[autodoc]] AutocastKwargs\n\n## DistributedDataParallelKwargs\n\n[[autodoc]] DistributedDataParallelKwargs\n\n## FP8RecipeKwargs\n\n[[autodoc]] utils.FP8RecipeKwargs\n\n## ProfileKwargs\n\n[[autodoc]] utils.ProfileKwargs\n\n## GradScalerKwargs\n\n[[autodoc]] GradScalerKwargs\n\n## InitProcessGroupKwargs\n\n[[autodoc]] InitProcessGroupKwargs\n\n## KwargsHandler\n\n[[autodoc]] utils.KwargsHandler\n" }, { "path": "docs/source/package_reference/launchers.md", "content": "\n\n# Launchers\n\nFunctions for launching training on distributed processes.\n\n## notebook_launcher\n\n[[autodoc]] accelerate.notebook_launcher\n\n## debug_launcher\n\n[[autodoc]] accelerate.debug_launcher" }, { "path": "docs/source/package_reference/logging.md", "content": "\n\n# Logging\n\nRefer to the [Troubleshooting guide](../usage_guides/troubleshooting#logging) or to the example below to learn \nhow to use Accelerate's logger. \n\n[[autodoc]] logging.get_logger" }, { "path": "docs/source/package_reference/megatron_lm.md", "content": "\n\n# Megatron-LM utilities\n\n## MegatronLMPlugin\n\n[[autodoc]] utils.MegatronLMPlugin\n\n## MegatronLMDummyScheduler\n\n[[autodoc]] utils.MegatronLMDummyScheduler\n\n## MegatronLMDummyDataLoader\n\n[[autodoc]] utils.MegatronLMDummyDataLoader\n\n## AbstractTrainStep\n\n[[autodoc]] utils.AbstractTrainStep\n\n## GPTTrainStep\n\n[[autodoc]] utils.GPTTrainStep\n\n## BertTrainStep\n\n[[autodoc]] utils.BertTrainStep\n\n## T5TrainStep\n\n[[autodoc]] utils.T5TrainStep\n\n## avg_losses_across_data_parallel_group\n\n[[autodoc]] utils.avg_losses_across_data_parallel_group\n" }, { "path": "docs/source/package_reference/state.md", "content": "\n\n# Stateful Classes\n\nBelow are variations of a [singleton class](https://en.wikipedia.org/wiki/Singleton_pattern) in the sense that all\ninstances share the same state, which is initialized on the first instantiation.\n\nThese classes are immutable and store information about certain configurations or \nstates.\n\n## PartialState\n\n[[autodoc]] state.PartialState\n\n## AcceleratorState\n\n[[autodoc]] state.AcceleratorState\n\n## GradientState\n\n[[autodoc]] state.GradientState" }, { "path": "docs/source/package_reference/torch_wrappers.md", "content": "\n\n# DataLoaders, Optimizers, and Schedulers\n\nThe internal classes Accelerate uses to prepare objects for distributed training\nwhen calling [`~Accelerator.prepare`].\n\n## DataLoader utilities\n\n[[autodoc]] data_loader.prepare_data_loader\n[[autodoc]] data_loader.skip_first_batches\n\n## BatchSamplerShard\n\n[[autodoc]] data_loader.BatchSamplerShard\n\n## IterableDatasetShard\n\n[[autodoc]] data_loader.IterableDatasetShard\n\n## DataLoaderShard\n\n[[autodoc]] data_loader.DataLoaderShard\n\n## DataLoaderDispatcher\n\n[[autodoc]] data_loader.DataLoaderDispatcher\n\n## AcceleratedOptimizer\n\n[[autodoc]] optimizer.AcceleratedOptimizer\n\n## AcceleratedScheduler\n\n[[autodoc]] scheduler.AcceleratedScheduler" }, { "path": "docs/source/package_reference/tracking.md", "content": "\n\n# Experiment Trackers\n\n## GeneralTracker\n\n[[autodoc]] tracking.GeneralTracker\n\n## TensorBoardTracker\n\n[[autodoc]] tracking.TensorBoardTracker\n - __init__\n\n## WandBTracker\n\n[[autodoc]] tracking.WandBTracker\n - __init__\n\n## Trackio\n\n[[autodoc]] tracking.TrackioTracker\n - __init__\n\n## CometMLTracker\n\n[[autodoc]] tracking.CometMLTracker\n - __init__\n\n## AimTracker\n\n[[autodoc]] tracking.AimTracker\n - __init__\n\n## MLflowTracker\n\n[[autodoc]] tracking.MLflowTracker\n - __init__\n\n## ClearMLTracker\n\n[[autodoc]] tracking.ClearMLTracker\n - __init__\n\n## SwanLabTracker\n\n[[autodoc]] tracking.SwanLabTracker\n - __init__\n" }, { "path": "docs/source/package_reference/utilities.md", "content": "\n\n# Utility functions and classes\n\nBelow are a variety of utility functions that šŸ¤— Accelerate provides, broken down by use-case. \n\n## Constants\n\nConstants used throughout šŸ¤— Accelerate for reference\n\nThe following are constants used when utilizing [`Accelerator.save_state`]\n\n`utils.MODEL_NAME`: `\"pytorch_model\"`\n`utils.OPTIMIZER_NAME`: `\"optimizer\"`\n`utils.RNG_STATE_NAME`: `\"random_states\"`\n`utils.SCALER_NAME`: `\"scaler.pt`\n`utils.SCHEDULER_NAME`: `\"scheduler`\n\nThe following are constants used when utilizing [`Accelerator.save_model`]\n\n`utils.WEIGHTS_NAME`: `\"pytorch_model.bin\"`\n`utils.SAFE_WEIGHTS_NAME`: `\"model.safetensors\"`\n`utils.WEIGHTS_INDEX_NAME`: `\"pytorch_model.bin.index.json\"`\n`utils.SAFE_WEIGHTS_INDEX_NAME`: `\"model.safetensors.index.json\"`\n\n## Data Classes\n\nThese are basic dataclasses used throughout šŸ¤— Accelerate and they can be passed in as parameters.\n\n### Standalone\n\nThese are standalone dataclasses used for checks, such as the type of distributed system being used\n\n[[autodoc]] utils.ComputeEnvironment\n\n[[autodoc]] utils.DistributedType\n\n[[autodoc]] utils.DynamoBackend\n\n[[autodoc]] utils.LoggerType\n\n[[autodoc]] utils.PrecisionType\n\n[[autodoc]] utils.RNGType\n\n[[autodoc]] utils.SageMakerDistributedType\n\n### Kwargs\n\nThese are configurable arguments for specific interactions throughout the PyTorch ecosystem that Accelerate handles under the hood.\n\n[[autodoc]] utils.AutocastKwargs\n\n[[autodoc]] utils.DistributedDataParallelKwargs\n\n[[autodoc]] utils.FP8RecipeKwargs\n\n[[autodoc]] utils.GradScalerKwargs\n\n[[autodoc]] utils.InitProcessGroupKwargs\n\n[[autodoc]] utils.KwargsHandler\n\n## Plugins\n\nThese are plugins that can be passed to the [`Accelerator`] object. While they are defined elsewhere in the documentation, \nfor convenience all of them are available to see here:\n\n[[autodoc]] utils.DeepSpeedPlugin\n\n[[autodoc]] utils.FullyShardedDataParallelPlugin\n\n[[autodoc]] utils.GradientAccumulationPlugin\n\n[[autodoc]] utils.MegatronLMPlugin\n\n[[autodoc]] utils.TorchDynamoPlugin\n\n## Configurations\n\nThese are classes which can be configured and passed through to the appropriate integration\n\n[[autodoc]] utils.BnbQuantizationConfig\n\n[[autodoc]] utils.DataLoaderConfiguration\n\n[[autodoc]] utils.ProjectConfiguration\n\n## Environmental Variables\n\nThese are environmental variables that can be enabled for different use cases\n\n* `ACCELERATE_DEBUG_MODE` (`str`): Whether to run accelerate in debug mode. More info available [here](../usage_guides/debug.md).\n\n\n\n\n## Data Manipulation and Operations\n\nThese include data operations that mimic the same `torch` ops but can be used on distributed processes.\n\n[[autodoc]] utils.broadcast\n\n[[autodoc]] utils.broadcast_object_list\n\n[[autodoc]] utils.concatenate\n\n[[autodoc]] utils.convert_outputs_to_fp32\n\n[[autodoc]] utils.convert_to_fp32\n\n[[autodoc]] utils.gather\n\n[[autodoc]] utils.gather_object\n\n[[autodoc]] utils.get_grad_scaler\n\n[[autodoc]] utils.get_mixed_precision_context_manager\n\n[[autodoc]] utils.listify\n\n[[autodoc]] utils.pad_across_processes\n\n[[autodoc]] utils.recursively_apply\n\n[[autodoc]] utils.reduce\n\n[[autodoc]] utils.send_to_device\n\n[[autodoc]] utils.slice_tensors\n\n## Environment Checks\n\nThese functionalities check the state of the current working environment including information about the operating system itself, what it can support, and if particular dependencies are installed. \n\n[[autodoc]] utils.is_bf16_available\n\n[[autodoc]] utils.is_mps_available\n\n[[autodoc]] utils.is_npu_available\n\n[[autodoc]] utils.is_torch_version\n\n[[autodoc]] utils.is_torch_xla_available\n\n[[autodoc]] utils.is_xpu_available\n\n## Environment Manipulation\n\n[[autodoc]] utils.patch_environment\n\n[[autodoc]] utils.clear_environment\n\n[[autodoc]] utils.write_basic_config\n\nWhen setting up šŸ¤— Accelerate for the first time, rather than running `accelerate config` [~utils.write_basic_config] can be used as an alternative for quick configuration.\n\n[[autodoc]] utils.set_numa_affinity\n\n[[autodoc]] utils.environment.override_numa_affinity\n\n[[autodoc]] utils.purge_accelerate_environment\n\n## Memory\n\n[[autodoc]] utils.find_executable_batch_size\n\n## Modeling\n\nThese utilities relate to interacting with PyTorch models\n\n[[autodoc]] utils.calculate_maximum_sizes\n\n[[autodoc]] utils.compute_module_sizes\n\n[[autodoc]] utils.extract_model_from_parallel\n\n[[autodoc]] utils.get_balanced_memory\n\n[[autodoc]] utils.get_max_layer_size\n\n[[autodoc]] utils.infer_auto_device_map\n\n[[autodoc]] utils.load_checkpoint_in_model\n\n[[autodoc]] utils.load_offloaded_weights\n\n[[autodoc]] utils.load_state_dict\n\n[[autodoc]] utils.offload_state_dict\n\n[[autodoc]] utils.retie_parameters\n\n[[autodoc]] utils.set_module_tensor_to_device\n\n[[autodoc]] utils.get_module_children_bottom_up\n\n## Parallel\n\nThese include general utilities that should be used when working in parallel.\n\n[[autodoc]] utils.extract_model_from_parallel\n\n[[autodoc]] utils.save\n\n[[autodoc]] utils.load\n\n[[autodoc]] utils.wait_for_everyone\n\n\n## Random\n\nThese utilities relate to setting and synchronizing of all the random states.\n\n[[autodoc]] utils.set_seed\n\n[[autodoc]] utils.synchronize_rng_state\n\n[[autodoc]] utils.synchronize_rng_states\n\n\n## PyTorch XLA\n\nThese include utilities that are useful while using PyTorch with XLA.\n\n[[autodoc]] utils.install_xla\n\n## Loading model weights\n\nThese include utilities that are useful to load checkpoints.\n\n[[autodoc]] utils.load_checkpoint_in_model\n\n## Quantization\n\nThese include utilities that are useful to quantize model.\n\n[[autodoc]] utils.load_and_quantize_model\n" }, { "path": "docs/source/quicktour.md", "content": "\n\n# Quicktour\n\nThere are many ways to launch and run your code depending on your training environment ([torchrun](https://pytorch.org/docs/stable/elastic/run.html), [DeepSpeed](https://www.deepspeed.ai/), etc.) and available hardware. Accelerate offers a unified interface for launching and training on different distributed setups, allowing you to focus on your PyTorch training code instead of the intricacies of adapting your code to these different setups. This allows you to easily scale your PyTorch code for training and inference on distributed setups with hardware like GPUs and TPUs. Accelerate also provides Big Model Inference to make loading and running inference with really large models that usually don't fit in memory more accessible.\n\nThis quicktour introduces the three main features of Accelerate:\n\n* a unified command line launching interface for distributed training scripts\n* a training library for adapting PyTorch training code to run on different distributed setups\n* Big Model Inference\n\n## Unified launch interface\n\nAccelerate automatically selects the appropriate configuration values for any given distributed training framework (DeepSpeed, FSDP, etc.) through a unified configuration file generated from the [`accelerate config`](package_reference/cli#accelerate-config) command. You could also pass the configuration values explicitly to the command line which is helpful in certain situations like if you're using SLURM.\n\n\nBut in most cases, you should always run [`accelerate config`](package_reference/cli#accelerate-config) first to help Accelerate learn about your training setup.\n\n```bash\naccelerate config\n```\n\nThe [`accelerate config`](package_reference/cli#accelerate-config) command creates and saves a default_config.yaml file in Accelerate's cache folder. This file stores the configuration for your training environment, which helps Accelerate correctly launch your training script based on your machine.\n\nAfter you've configured your environment, you can test your setup with [`accelerate test`](package_reference/cli#accelerate-test), which launches a short script to test the distributed environment.\n\n```bash\naccelerate test\n```\n\n> [!TIP]\n> Add `--config_file` to the `accelerate test` or `accelerate launch` command to specify the location of the configuration file if it is saved in a non-default location like the cache.\n\nOnce your environment is set up, launch your training script with [`accelerate launch`](package_reference/cli#accelerate-launch)!\n\n```bash\naccelerate launch path_to_script.py --args_for_the_script\n```\n\nTo learn more, check out the [Launch distributed code](basic_tutorials/launch) tutorial for more information about launching your scripts.\n\nWe also have a [configuration zoo](https://github.com/huggingface/accelerate/blob/main/examples/config_yaml_templates) which showcases a number of premade **minimal** example configurations for a variety of setups you can run.\n\n## Adapt training code\n\nThe next main feature of Accelerate is the [`Accelerator`] class which adapts your PyTorch code to run on different distributed setups.\n\nYou only need to add a few lines of code to your training script to enable it to run on multiple GPUs or TPUs.\n\n```diff\n+ from accelerate import Accelerator\n+ accelerator = Accelerator()\n\n+ device = accelerator.device\n+ model, optimizer, training_dataloader, scheduler = accelerator.prepare(\n+ model, optimizer, training_dataloader, scheduler\n+ )\n\n for batch in training_dataloader:\n optimizer.zero_grad()\n inputs, targets = batch\n- inputs = inputs.to(device)\n- targets = targets.to(device)\n outputs = model(inputs)\n loss = loss_function(outputs, targets)\n+ accelerator.backward(loss)\n optimizer.step()\n scheduler.step()\n```\n\n1. Import and instantiate the [`Accelerator`] class at the beginning of your training script. The [`Accelerator`] class initializes everything necessary for distributed training, and it automatically detects your training environment (a single machine with a GPU, a machine with several GPUs, several machines with multiple GPUs or a TPU, etc.) based on how the code was launched.\n\n```python\nfrom accelerate import Accelerator\n\naccelerator = Accelerator()\n```\n\n2. Remove calls like `.cuda()` on your model and input data. The [`Accelerator`] class automatically places these objects on the appropriate device for you.\n\n> [!WARNING]\n> This step is *optional* but it is considered best practice to allow Accelerate to handle device placement. You could also deactivate automatic device placement by passing `device_placement=False` when initializing the [`Accelerator`]. If you want to explicitly place objects on a device with `.to(device)`, make sure you use `accelerator.device` instead. For example, if you create an optimizer before placing a model on `accelerator.device`, training fails on a TPU.\n\n> [!WARNING]\n> Accelerate does not use non-blocking transfers by default for its automatic device placement, which can result in potentially unwanted CUDA synchronizations. You can enable non-blocking transfers by passing a [`~utils.dataclasses.DataLoaderConfiguration`] with `non_blocking=True` set as the `dataloader_config` when initializing the [`Accelerator`]. As usual, non-blocking transfers will only work if the dataloader also has `pin_memory=True` set. Be wary that using non-blocking transfers from GPU to CPU may cause incorrect results if it results in CPU operations being performed on non-ready tensors.\n\n```py\ndevice = accelerator.device\n```\n\n3. Pass all relevant PyTorch objects for training (optimizer, model, dataloader(s), learning rate scheduler) to the [`~Accelerator.prepare`] method as soon as they're created. This method wraps the model in a container optimized for your distributed setup, uses Accelerates version of the optimizer and scheduler, and creates a sharded version of your dataloader for distribution across GPUs or TPUs.\n\n```python\nmodel, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(\n model, optimizer, train_dataloader, lr_scheduler\n)\n```\n\n4. Replace `loss.backward()` with [`~Accelerator.backward`] to use the correct `backward()` method for your training setup.\n\n```py\naccelerator.backward(loss)\n```\n\nRead [Accelerateā€™s internal mechanisms](concept_guides/internal_mechanism) guide to learn more details about how Accelerate adapts your code.\n\n### Distributed evaluation\n\nTo perform distributed evaluation, pass your validation dataloader to the [`~Accelerator.prepare`] method:\n\n```python\nvalidation_dataloader = accelerator.prepare(validation_dataloader)\n```\n\nEach device in your distributed setup only receives a part of the evaluation data, which means you should group your predictions together with the [`~Accelerator.gather_for_metrics`] method. This method requires all tensors to be the same size on each process, so if your tensors have different sizes on each process (for instance when dynamically padding to the maximum length in a batch), you should use the [`~Accelerator.pad_across_processes`] method to pad you tensor to the largest size across processes. Note that the tensors needs to be 1D and that we concatenate the tensors along the first dimension. \n\n```python\nfor inputs, targets in validation_dataloader:\n predictions = model(inputs)\n # Gather all predictions and targets\n all_predictions, all_targets = accelerator.gather_for_metrics((predictions, targets))\n # Example of use with a *Datasets.Metric*\n metric.add_batch(all_predictions, all_targets)\n```\n\nFor more complex cases (e.g. 2D tensors, don't want to concatenate tensors, dict of 3D tensors), you can pass `use_gather_object=True` in `gather_for_metrics`. This will return the list of objects after gathering. Note that using it with GPU tensors is not well supported and inefficient.\n\n> [!TIP]\n> Data at the end of a dataset may be duplicated so the batch can be equally divided among all workers. The [`~Accelerator.gather_for_metrics`] method automatically removes the duplicated data to calculate a more accurate metric.\n\n## Big Model Inference\n\nAccelerate's Big Model Inference has two main features, [`~accelerate.init_empty_weights`] and [`~accelerate.load_checkpoint_and_dispatch`], to load large models for inference that typically don't fit into memory.\n\n> [!TIP]\n> Take a look at the [Handling big models for inference](concept_guides/big_model_inference) guide for a better understanding of how Big Model Inference works under the hood.\n\n### Empty weights initialization\n\nThe [`~accelerate.init_empty_weights`] context manager initializes models of any size by creating a *model skeleton* and moving and placing parameters each time they're created to PyTorch's [**meta**](https://pytorch.org/docs/main/meta.html) device. This way, not all weights are immediately loaded and only a small part of the model is loaded into memory at a time.\n\nFor example, loading an empty [Mixtral-8x7B](https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1) model takes significantly less memory than fully loading the models and weights on the CPU.\n\n```py\nfrom accelerate import init_empty_weights\nfrom transformers import AutoConfig, AutoModelForCausalLM\n\nconfig = AutoConfig.from_pretrained(\"mistralai/Mixtral-8x7B-Instruct-v0.1\")\nwith init_empty_weights():\n model = AutoModelForCausalLM.from_config(config)\n```\n\n### Load and dispatch weights\n\nThe [`~accelerate.load_checkpoint_and_dispatch`] function loads full or sharded checkpoints into the empty model, and automatically distribute weights across all available devices.\n\nThe `device_map` parameter determines where to place each model layer, and specifying `\"auto\"` places them on the GPU first, then the CPU, and finally the hard drive as memory-mapped tensors if there's still not enough memory. Use the `no_split_module_classes` parameter to indicate which modules shouldn't be split across devices (typically those with a residual connection).\n\n```py\nfrom accelerate import load_checkpoint_and_dispatch\n\nmodel_checkpoint = \"your-local-model-folder\"\nmodel = load_checkpoint_and_dispatch(\n model, checkpoint=model_checkpoint, device_map=\"auto\", no_split_module_classes=['Block']\n)\n```\n\n## Next steps\n\nNow that you've been introduced to the main Accelerate features, your next steps could include:\n\n* Check out the [tutorials](basic_tutorials/overview) for a gentle walkthrough of Accelerate. This is especially useful if you're new to distributed training and the library.\n* Dive into the [guides](usage_guides/explore) to see how to use Accelerate for specific use-cases.\n* Deepen your conceptual understanding of how Accelerate works internally by reading the [concept guides](concept_guides/internal_mechanism).\n* Look up classes and commands in the [API reference](package_reference/accelerator) to see what parameters and options are available.\n" }, { "path": "docs/source/usage_guides/big_modeling.md", "content": "\n\n# Big Model Inference\n\nOne of the biggest advancements Accelerate provides is [Big Model Inference](../concept_guides/big_model_inference), which allows you to perform inference with models that don't fully fit on your graphics card.\n\nThis tutorial will show you how to use Big Model Inference in Accelerate and the Hugging Face ecosystem.\n\n## Accelerate\n\nA typical workflow for loading a PyTorch model is shown below. `ModelClass` is a model that exceeds the GPU memory of your device (mps or cuda or xpu).\n\n```py\nimport torch\n\nmy_model = ModelClass(...)\nstate_dict = torch.load(checkpoint_file)\nmy_model.load_state_dict(state_dict)\n```\n\nWith Big Model Inference, the first step is to init an empty skeleton of the model with the `init_empty_weights` context manager. This doesn't require any memory because `my_model` is \"parameterless\".\n\n```py\nfrom accelerate import init_empty_weights\nwith init_empty_weights():\n my_model = ModelClass(...)\n```\n\nNext, the weights are loaded into the model for inference.\n\nThe [`load_checkpoint_and_dispatch`] method loads a checkpoint inside your empty model and dispatches the weights for each layer across all available devices, starting with the fastest devices (GPU, MPS, XPU, NPU, MLU, SDAA, MUSA) first before moving to the slower ones (CPU and hard drive).\n\nSetting `device_map=\"auto\"` automatically fills all available space on the GPU(s) first, then the CPU, and finally, the hard drive (the absolute slowest option) if there is still not enough memory.\n\n> [!TIP]\n> Refer to the [Designing a device map](../concept_guides/big_model_inference#designing-a-device-map) guide for more details on how to design your own device map.\n\n```py\nfrom accelerate import load_checkpoint_and_dispatch\n\nmodel = load_checkpoint_and_dispatch(\n model, checkpoint=checkpoint_file, device_map=\"auto\"\n)\n```\n\nIf there are certain ā€œchunksā€ of layers that shouldnā€™t be split, pass them to `no_split_module_classes` (see [here](../concept_guides/big_model_inference#loading-weights) for more details).\n\nA models weights can also be sharded into multiple checkpoints to save memory, such as when the `state_dict` doesn't fit in memory (see [here](../concept_guides/big_model_inference#sharded-checkpoints) for more details).\n\nNow that the model is fully dispatched, you can perform inference.\n\n```py\ninput = torch.randn(2,3)\ndevice_type = next(iter(model.parameters())).device.type\ninput = input.to(device_type)\noutput = model(input)\n```\n\nEach time an input is passed through a layer, it is sent from the CPU to the GPU (or disk to CPU to GPU), the output is calculated, and the layer is removed from the GPU going back down the line. While this adds some overhead to inference, it enables you to run any size model on your system, as long as the largest layer fits on your GPU.\n\nMultiple GPUs, or \"model parallelism\", can be utilized but only one GPU will be active at any given moment. This forces the GPU to wait for the previous GPU to send it the output. You should launch your script normally with Python instead of other tools like torchrun and accelerate launch.\n\n> [!TIP]\n> You may also be interested in *pipeline parallelism* which utilizes all available GPUs at once, instead of only having one GPU active at a time. This approach is less flexible though. For more details, refer to the [Memory-efficient pipeline parallelism](./distributed_inference#memory-efficient-pipeline-parallelism-experimental) guide.\n\n\n\nTake a look at a full example of Big Model Inference below.\n\n```py\nimport torch\nfrom accelerate import init_empty_weights, load_checkpoint_and_dispatch\n\nwith init_empty_weights():\n model = MyModel(...)\n\nmodel = load_checkpoint_and_dispatch(\n model, checkpoint=checkpoint_file, device_map=\"auto\"\n)\n\ninput = torch.randn(2,3)\ndevice_type = next(iter(model.parameters())).device.type\ninput = input.to(device_type)\noutput = model(input)\n```\n\n## Hugging Face ecosystem\n\nOther libraries in the Hugging Face ecosystem, like Transformers or Diffusers, supports Big Model Inference in their [`~transformers.PreTrainedModel.from_pretrained`] constructors.\n\nYou just need to add `device_map=\"auto\"` in [`~transformers.PreTrainedModel.from_pretrained`] to enable Big Model Inference.\n\nFor example, load Big Sciences T0pp 11 billion parameter model with Big Model Inference.\n\n```py\nfrom transformers import AutoModelForSeq2SeqLM\n\nmodel = AutoModelForSeq2SeqLM.from_pretrained(\"bigscience/T0pp\", device_map=\"auto\")\n```\n\nAfter loading the model, the empty init and smart dispatch steps from before are executed 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 `torch_dtype` parameter to load a model in a lower precision.\n\n```py\nfrom transformers import AutoModelForSeq2SeqLM\n\nmodel = AutoModelForSeq2SeqLM.from_pretrained(\"bigscience/T0pp\", device_map=\"auto\", torch_dtype=torch.float16)\n```\n\n## Next steps\n\nFor a more detailed explanation of Big Model Inference, make sure to check out the [conceptual guide](../concept_guides/big_model_inference)!\n" }, { "path": "docs/source/usage_guides/checkpoint.md", "content": "\n\n# Checkpointing\n\nWhen training a PyTorch model with Accelerate, you may often want to save and continue a state of training. Doing so requires\nsaving and loading the model, optimizer, RNG generators, and the GradScaler. Inside Accelerate are two convenience functions to achieve this quickly:\n- Use [`~Accelerator.save_state`] for saving everything mentioned above to a folder location\n- Use [`~Accelerator.load_state`] for loading everything stored from an earlier `save_state`\n\nTo further customize where and how states are saved through [`~Accelerator.save_state`] the [`~utils.ProjectConfiguration`] class can be used. For example \nif `automatic_checkpoint_naming` is enabled each saved checkpoint will be located then at `Accelerator.project_dir/checkpoints/checkpoint_{checkpoint_number}`.\n\nIt should be noted that the expectation is that those states come from the same training script, they should not be from two separate scripts.\n\n- By using [`~Accelerator.register_for_checkpointing`], you can register custom objects to be automatically stored or loaded from the two prior functions,\nso long as the object has a `state_dict` **and** a `load_state_dict` functionality. This could include objects such as a learning rate scheduler. \n\n\nBelow is a brief example using checkpointing to save and reload a state during training:\n\n```python\nfrom accelerate import Accelerator\nimport torch\n\naccelerator = Accelerator(project_dir=\"my/save/path\")\n\nmy_scheduler = torch.optim.lr_scheduler.StepLR(my_optimizer, step_size=1, gamma=0.99)\nmy_model, my_optimizer, my_training_dataloader = accelerator.prepare(my_model, my_optimizer, my_training_dataloader)\n\n# Register the LR scheduler\naccelerator.register_for_checkpointing(my_scheduler)\n\n# Save the starting state\naccelerator.save_state()\n\ndevice = accelerator.device\nmy_model.to(device)\n\n# Perform training\nfor epoch in range(num_epochs):\n for batch in my_training_dataloader:\n my_optimizer.zero_grad()\n inputs, targets = batch\n inputs = inputs.to(device)\n targets = targets.to(device)\n outputs = my_model(inputs)\n loss = my_loss_function(outputs, targets)\n accelerator.backward(loss)\n my_optimizer.step()\n my_scheduler.step()\n\n# Restore the previous state\naccelerator.load_state(\"my/save/path/checkpointing/checkpoint_0\")\n```\n\n## Restoring the state of the DataLoader \n\nAfter resuming from a checkpoint, it may also be desirable to resume from a particular point in the active `DataLoader` if \nthe state was saved during the middle of an epoch. You can use [`~Accelerator.skip_first_batches`] to do so. \n\n```python\nfrom accelerate import Accelerator\n\naccelerator = Accelerator(project_dir=\"my/save/path\")\n\ntrain_dataloader = accelerator.prepare(train_dataloader)\naccelerator.load_state(\"my_state\")\n\n# Assume the checkpoint was saved 100 steps into the epoch\nskipped_dataloader = accelerator.skip_first_batches(train_dataloader, 100)\n\n# After the first iteration, go back to `train_dataloader`\n\n# First epoch\nfor batch in skipped_dataloader:\n # Do something\n pass\n\n# Second epoch\nfor batch in train_dataloader:\n # Do something\n pass\n```\n" }, { "path": "docs/source/usage_guides/compilation.md", "content": "# Compilation\n\n## Overview\n\nPytorch 2.0 introduced `torch.compile`, a powerful feature that makes PyTorch code run faster by JIT-compiling PyTorch code into optimized kernels. Key features of `torch.compile` include:\n\n- **Performance Improvement**: Significantly speeds up model execution by optimizing the computation graph.\n- **Ease of Use**: Requires minimal code changes to implement, making it highly accessible.\n- **Compatibility**: Works seamlessly with existing PyTorch code and models.\n\nWhen used with Accelerate, `torch.compile` integrates smoothly into distributed training workflows, allowing you to benefit from both distributed execution and compilation optimizations simultaneously.\n\nThe first execution of compiled code typically takes longer as it includes the compilation time, but subsequent runs are significantly faster. For optimal performance in different scenarios, `torch.compile` offers various modes like `\"default\"`, `\"reduce-overhead\"` (which uses CUDA graphs to further reduce overhead), and `\"max-autotune\"` (which performs extensive autotuning to find the best kernels for your model).\n\n## Using `torch.compile` with Accelerate\n\nAccelerate provides `TorchDynamoPlugin` for easy and seemless integration of `torch.compile` into your training scripts.\n\n```python\nfrom accelerate import Accelerator\nfrom accelerate.utils import TorchDynamoPlugin\n\n# Configure the compilation backend\ndynamo_plugin = TorchDynamoPlugin(\n backend=\"inductor\", # Options: \"inductor\", \"aot_eager\", \"aot_nvfuser\", etc.\n mode=\"default\", # Options: \"default\", \"reduce-overhead\", \"max-autotune\"\n fullgraph=True,\n dynamic=False\n)\n\n# Initialize accelerator with the plugin\naccelerator = Accelerator(dynamo_plugin=dynamo_plugin)\n# This will apply torch.compile to your model\nmodel = accelerator.prepare(model)\n```\n\nIt is compatible with all other features and plugins of Accelerate, including mixed precision, distributed training (DDP, FSDP, Deepspeed), etc.\n\n## Regional Compilation\n\nInstead of trying to compile the whole model, which usually has a big problem space for optimization. Regional compilation targets repeated blocks of the same class and compiles them sequentially to hit the compiler's cache. For example, in `GPT2LMHeadModel`, the repeated block/class is `GPT2Block`, and can be accessed as `model.transformer.h[0]`. The rest of the model (e.g model.lm_head) is compiled separately.\n\nThis allows us to speed up the compilation overhead / cold start of models like LLMs and Transformers in general.\nSee for more details.\n\n### How to Use Regional Compilation\n\nIt can be enabled by setting `use_regional_compilation=True` in the `TorchDynamoPlugin` configuration:\n\n```python\n# Configure the compilation backend\ndynamo_plugin = TorchDynamoPlugin(\n use_regional_compilation=True,\n ... # other parameters\n)\n# Initialize accelerator with the plugin\naccelerator = Accelerator(dynamo_plugin=dynamo_plugin)\n# This will apply compile_regions to your model\nmodel = accelerator.prepare(model)\n```\n\nYou could also use the `accelerate.utils.compile_regions` utility directly the same way you would use `torch.compile`.\n\n### Benefits of Regional Compilation\n\nWe have conducted extensive benchmarks comparing full compilation and regional compilation using the `torch.compile` feature in PyTorch. The full results are available in the [accelerate repository](https://github.com/huggingface/accelerate/tree/main/benchmarks/torch.compile/regional_compilation). The key findings from our benchmarks are:\n\n1. **Comparable Performance**: Regional compilation delivers performance speedups similar to full compilation, especially for larger models.\n2. **Faster Compilation**: Regional compilation significantly reduces the time taken to compile models, making it a more efficient choice for deployment.\n3. **Batch Size Impact**: The performance difference between compilation strategies diminishes with larger batch sizes, indicating that the overhead of compilation is less impactful in those scenarios.\n4. **Model Size Consideration**: The benefits of regional compilation are more pronounced in larger models, where the compilation time savings can be substantial.\n5. **Practical Application**: For real-world applications, regional compilation is a practical choice for optimizing training cold start times, especially when working with large models.\n\n## Conclusion\n\nBoth full and regional compilation can significantly speed up your models. Regional compilation offers a practical balance between compilation time and runtime performance, especially for training large models with substantial batch sizes.\n" }, { "path": "docs/source/usage_guides/ddp_comm_hook.md", "content": "\n\n# DDP Communication Hooks\n\nDistributed Data Parallel (DDP) communication hooks provide a generic interface to control how gradients are communicated across workers by overriding the vanilla allreduce in `DistributedDataParallel`. A few built-in communication hooks are provided, and users can easily apply any of these hooks to optimize communication.\n\n\n- **FP16 Compression Hook**: Compresses gradients by casting them to half-precision floating-point format (`torch.float16`), reducing communication overhead.\n- **BF16 Compression Hook**: Similar to FP16, but uses the Brain Floating Point format (`torch.bfloat16`), which can be more efficient on certain hardware.\n- **PowerSGD Hook**: An advanced gradient compression algorithm that provides high compression rates and can accelerate bandwidth-bound distributed training.\n\nIn this tutorial, you will see how to quickly set up DDP communication hooks and perform training with the utilities provided in Accelerate, which can be as simple as adding just one new line of code! This demonstrates how to use DDP communication hooks to optimize gradient communication in distributed training with the Accelerate library.\n\n## FP16 Compression Hook\n\n\n\n\n```python\nimport torch\nfrom torch.nn.parallel import DistributedDataParallel as DDP\nfrom torch.distributed.algorithms.ddp_comm_hooks import default_hooks\nfrom accelerate.test_utils.testing import get_backend\n\ndevice_type, _, _ = get_backend()\ndevice_id = getattr(torch, device_type, torch.cuda).current_device()\n\nclass MyModel(torch.nn.Module):\n def __init__(self):\n super().__init__()\n self.layer = torch.nn.Linear(10, 10)\n\n def forward(self, x):\n return self.layer(x)\n\nmodel = MyModel()\nmodel = DDP(model, device_ids=[device_id])\nmodel.register_comm_hook(state=None, hook=default_hooks.fp16_compress_hook)\n\n# Training loop\nfor data, targets in data_loader:\n outputs = model(data)\n loss = criterion(outputs, targets)\n loss.backward()\n optimizer.step()\n optimizer.zero_grad()\n```\n\n\n\n\n```python\nfrom accelerate import Accelerator, DDPCommunicationHookType, DistributedDataParallelKwargs\nimport torch\n\nclass MyModel(torch.nn.Module):\n def __init__(self):\n super().__init__()\n self.layer = torch.nn.Linear(10, 10)\n\n def forward(self, x):\n return self.layer(x)\n\n# DDP Communication Hook setup\nddp_kwargs = DistributedDataParallelKwargs(comm_hook=DDPCommunicationHookType.FP16)\naccelerator = Accelerator(kwargs_handlers=[ddp_kwargs])\n\nmodel = MyModel()\noptimizer = torch.optim.Adam(model.parameters())\ndata_loader = DataLoader(dataset, batch_size=16)\n\nmodel, optimizer, data_loader = accelerator.prepare(model, optimizer, data_loader)\n\n# Training loop\nfor data, targets in data_loader:\n outputs = model(data)\n loss = criterion(outputs, targets)\n accelerator.backward(loss)\n optimizer.step()\n optimizer.zero_grad()\n```\n\n\n\n\n### BF16 Compression Hook\n\n\n\nBF16 Compression Hook API is experimental, and it requires NCCL version later than 2.9.6.\n\n\n\n\n\n\n```python\nimport torch\nfrom torch.nn.parallel import DistributedDataParallel as DDP\nfrom torch.distributed.algorithms.ddp_comm_hooks import default_hooks\nfrom accelerate.test_utils.testing import get_backend\n\ndevice_type, _, _ = get_backend()\ndevice_id = getattr(torch, device_type, torch.cuda).current_device()\n\nclass MyModel(torch.nn.Module):\n def __init__(self):\n super().__init__()\n self.layer = torch.nn.Linear(10, 10)\n\n def forward(self, x):\n return self.layer(x)\n\nmodel = MyModel()\nmodel = DDP(model, device_ids=[device_id])\nmodel.register_comm_hook(state=None, hook=default_hooks.bf16_compress_hook)\n\n# Training loop\nfor data, targets in data_loader:\n outputs = model(data)\n loss = criterion(outputs, targets)\n loss.backward()\n optimizer.step()\n optimizer.zero_grad()\n```\n\n\n\n\n```python\nfrom accelerate import Accelerator, DDPCommunicationHookType, DistributedDataParallelKwargs\nimport torch\n\nclass MyModel(torch.nn.Module):\n def __init__(self):\n super().__init__()\n self.layer = torch.nn.Linear(10, 10)\n\n def forward(self, x):\n return self.layer(x)\n\n# DDP Communication Hook setup\nddp_kwargs = DistributedDataParallelKwargs(comm_hook=DDPCommunicationHookType.BF16)\naccelerator = Accelerator(kwargs_handlers=[ddp_kwargs])\n\nmodel = MyModel()\noptimizer = torch.optim.Adam(model.parameters())\ndata_loader = DataLoader(dataset, batch_size=16)\n\nmodel, optimizer, data_loader = accelerator.prepare(model, optimizer, data_loader)\n\n# Training loop\nfor data, targets in data_loader:\n outputs = model(data)\n loss = criterion(outputs, targets)\n accelerator.backward(loss)\n optimizer.step()\n optimizer.zero_grad()\n```\n\n\n\n\n### PowerSGD Hook\n\n\n\nPowerSGD typically requires extra memory of the same size as the modelā€™s gradients to enable error feedback, which can compensate for biased compressed communication and improve accuracy.\n\n\n\n\n\n\n```python\nimport torch\nfrom torch.nn.parallel import DistributedDataParallel as DDP\nfrom torch.distributed.algorithms.ddp_comm_hooks import powerSGD_hook\nfrom accelerate.test_utils.testing import get_backend\n\ndevice_type, _, _ = get_backend()\ndevice_id = getattr(torch, device_type, torch.cuda).current_device()\n\nclass MyModel(torch.nn.Module):\n def __init__(self):\n super().__init__()\n self.layer = torch.nn.Linear(10, 10)\n\n def forward(self, x):\n return self.layer(x)\n\nmodel = MyModel()\nmodel = DDP(model, device_ids=[device_id])\nstate = powerSGD_hook.PowerSGDState(process_group=None)\nmodel.register_comm_hook(state=state, hook=powerSGD_hook.powerSGD_hook)\n\n# Training loop\nfor data, targets in data_loader:\n outputs = model(data)\n loss = criterion(outputs, targets)\n loss.backward()\n optimizer.step()\n optimizer.zero_grad()\n```\n\n\n\n\n```python\nfrom accelerate import Accelerator, DDPCommunicationHookType, DistributedDataParallelKwargs\nimport torch\n\nclass MyModel(torch.nn.Module):\n def __init__(self):\n super().__init__()\n self.layer = torch.nn.Linear(10, 10)\n\n def forward(self, x):\n return self.layer(x)\n\n# DDP Communication Hook setup\nddp_kwargs = DistributedDataParallelKwargs(comm_hook=DDPCommunicationHookType.POWER_SGD)\naccelerator = Accelerator(kwargs_handlers=[ddp_kwargs])\n\nmodel = MyModel()\noptimizer = torch.optim.Adam(model.parameters())\ndata_loader = DataLoader(dataset, batch_size=16)\n\nmodel, optimizer, data_loader = accelerator.prepare(model, optimizer, data_loader)\n\n# Training loop\nfor data, targets in data_loader:\n outputs = model(data)\n loss = criterion(outputs, targets)\n accelerator.backward(loss)\n optimizer.step()\n optimizer.zero_grad()\n```\n\n\n\n\n## DDP Communication Hooks utilities\n\nThere are two additional utilities for supporting optional functionalities with the communication hooks.\n\n### comm_wrapper\n\n`comm_wrapper` is an option to wrap a communication hook with additional functionality. For example, it can be used to combine FP16 compression with other communication strategies. Currently supported wrappers are `no`, `fp16`, and `bf16`.\n\n```python\nfrom accelerate import Accelerator, DDPCommunicationHookType, DistributedDataParallelKwargs\nimport torch\n\nclass MyModel(torch.nn.Module):\n def __init__(self):\n super().__init__()\n self.layer = torch.nn.Linear(10, 10)\n\n def forward(self, x):\n return self.layer(x)\n\n# DDP Communication Hook setup\nddp_kwargs = DistributedDataParallelKwargs(\n comm_hook=DDPCommunicationHookType.POWER_SGD,\n comm_wrapper=DDPCommunicationHookType.FP16\n)\naccelerator = Accelerator(kwargs_handlers=[ddp_kwargs])\n\nmodel = MyModel()\noptimizer = torch.optim.Adam(model.parameters())\ndata_loader = DataLoader(dataset, batch_size=16)\n\nmodel, optimizer, data_loader = accelerator.prepare(model, optimizer, data_loader)\n\n# Training loop\nfor data, targets in data_loader:\n outputs = model(data)\n loss = criterion(outputs, targets)\n accelerator.backward(loss)\n optimizer.step()\n optimizer.zero_grad()\n```\n\n### comm_state_option\n\n`comm_state_option` allows you to pass additional state information required by certain communication hooks. This is particularly useful for stateful hooks like `PowerSGD`, which require maintaining hyperparameters and internal states across training steps. Below is an example showcasing the use of `comm_state_option` with the `PowerSGD` hook.\n\n```python\nfrom accelerate import Accelerator, DDPCommunicationHookType, DistributedDataParallelKwargs\nimport torch\n\nclass MyModel(torch.nn.Module):\n def __init__(self):\n super().__init__()\n self.layer = torch.nn.Linear(10, 10)\n\n def forward(self, x):\n return self.layer(x)\n\n# DDP Communication Hook setup\nddp_kwargs = DistributedDataParallelKwargs(\n comm_hook=DDPCommunicationHookType.POWER_SGD,\n comm_state_option={\"matrix_approximation_rank\": 2}\n)\naccelerator = Accelerator(kwargs_handlers=[ddp_kwargs])\n\nmodel = MyModel()\noptimizer = torch.optim.Adam(model.parameters())\ndata_loader = DataLoader(dataset, batch_size=16)\n\nmodel, optimizer, data_loader = accelerator.prepare(model, optimizer, data_loader)\n\n# Training loop\nfor data, targets in data_loader:\n outputs = model(data)\n loss = criterion(outputs, targets)\n accelerator.backward(loss)\n optimizer.step()\n optimizer.zero_grad()\n```\n\nFor more advanced usage and additional hooks, refer to the [PyTorch DDP Communication Hooks documentation](https://pytorch.org/docs/stable/ddp_comm_hooks.html).\n" }, { "path": "docs/source/usage_guides/deepspeed.md", "content": "\n\n# DeepSpeed\n\n[DeepSpeed](https://github.com/deepspeedai/DeepSpeed) implements everything described in the [ZeRO paper](https://huggingface.co/papers/1910.02054). Some of the salient optimizations are:\n\n1. Optimizer state partitioning (ZeRO stage 1)\n2. Gradient partitioning (ZeRO stage 2)\n3. Parameter partitioning (ZeRO stage 3)\n4. Custom mixed precision training handling\n5. A range of fast CUDA-extension-based optimizers\n6. ZeRO-Offload to CPU and Disk/NVMe\n7. Hierarchical partitioning of model parameters (ZeRO++)\n\nZeRO-Offload has its own dedicated paper: [ZeRO-Offload: Democratizing Billion-Scale Model Training](https://huggingface.co/papers/2101.06840). And NVMe-support is described in the paper [ZeRO-Infinity: Breaking the GPU\nMemory Wall for Extreme Scale Deep Learning](https://huggingface.co/papers/2104.07857).\n\nDeepSpeed ZeRO-2 is primarily used only for training, as its features are of no use to inference.\n\nDeepSpeed ZeRO-3 can be used for inference as well since it allows huge models to be loaded on multiple GPUs, which\nwon't be possible on a single GPU.\n\nAccelerate integrates [DeepSpeed](https://github.com/deepspeedai/DeepSpeed) via 2 options:\n\n1. 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\n this document is focused on this feature. This supports all the core features of DeepSpeed and gives user a lot of flexibility.\n User may have to change a few lines of code depending on the config.\n2. Integration via `deepspeed_plugin`.This supports subset of the DeepSpeed features and uses default options for the rest of the configurations.\n User need not change any code and is good for those who are fine with most of the default settings of DeepSpeed.\n\n## What is integrated?\n\nTraining:\n\n1. Accelerate integrates all features of DeepSpeed ZeRO. This includes all the ZeRO stages 1, 2 and 3 as well as ZeRO-Offload, ZeRO-Infinity (which can offload to disk/NVMe) and ZeRO++.\nBelow is a short description of Data Parallelism using ZeRO - Zero Redundancy Optimizer along with diagram from this [blog post](https://www.microsoft.com/en-us/research/blog/zero-deepspeed-new-system-optimizations-enable-training-models-with-over-100-billion-parameters/)\n![ZeRO Data Parallelism](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/parallelism-zero.png)\n\n(Source: [link](https://www.microsoft.com/en-us/research/blog/zero-deepspeed-new-system-optimizations-enable-training-models-with-over-100-billion-parameters/))\n\n a. **Stage 1** : Shards optimizer states across data parallel workers/GPUs\n\n b. **Stage 2** : Shards optimizer states + gradients across data parallel workers/GPUs\n\n c. **Stage 3**: Shards optimizer states + gradients + model parameters across data parallel workers/GPUs\n\n d. **Optimizer Offload**: Offloads the gradients + optimizer states to CPU/Disk building on top of ZERO Stage 2\n\n e. **Param Offload**: Offloads the model parameters to CPU/Disk building on top of ZERO Stage 3\n\n f. **Hierarchical Partitioning**: Enables efficient multi-node training with data-parallel training across nodes and ZeRO-3 sharding within a node, built on top of ZeRO Stage 3.\n\nNote: With respect to Disk Offload, the disk should be an NVME for decent speed but it technically works on any Disk\n\nInference:\n\n1. DeepSpeed ZeRO Inference supports ZeRO stage 3 with ZeRO-Infinity. It uses the same ZeRO protocol as training, but\n it doesn't use an optimizer and a lr scheduler and only stage 3 is relevant. For more details see:\n [deepspeed-zero-inference](#deepspeed-zero-inference).\n\n\n## How it works?\n\n**Pre-Requisites**: Install DeepSpeed version >=0.6.5. Please refer to the [DeepSpeed Installation details](https://github.com/deepspeedai/DeepSpeed#installation)\nfor more information.\n\nWe will first look at easy to use integration via `accelerate config`.\nFollowed by more flexible and feature rich `deepspeed config file` integration.\n\n### Accelerate DeepSpeed Plugin\nOn your machine(s) just run:\n\n```bash\naccelerate config\n```\n\nand answer the questions asked. It will ask whether you want to use a config file for DeepSpeed to which you should answer no. Then answer the following questions to generate a basic DeepSpeed config.\nThis will generate a config file that will be used automatically to properly set the\ndefault options when doing\n\n```bash\naccelerate launch my_script.py --args_to_my_script\n```\n\nFor instance, here is how you would run the NLP example `examples/nlp_example.py` (from the root of the repo) with DeepSpeed Plugin:\n\n**ZeRO Stage-2 DeepSpeed Plugin Example**\n```bash\ncompute_environment: LOCAL_MACHINE\ndeepspeed_config:\n gradient_accumulation_steps: 1\n gradient_clipping: 1.0\n offload_optimizer_device: none\n offload_param_device: none\n zero3_init_flag: true\n zero_stage: 2\ndistributed_type: DEEPSPEED\nfsdp_config: {}\nmachine_rank: 0\nmain_process_ip: null\nmain_process_port: null\nmain_training_function: main\nmixed_precision: fp16\nnum_machines: 1\nnum_processes: 2\nuse_cpu: false\n```\n\n```bash\naccelerate launch examples/nlp_example.py --mixed_precision fp16\n```\n\n**ZeRO Stage-3 with CPU Offload DeepSpeed Plugin Example**\n```bash\ncompute_environment: LOCAL_MACHINE\ndeepspeed_config:\n gradient_accumulation_steps: 1\n gradient_clipping: 1.0\n offload_optimizer_device: cpu\n offload_param_device: cpu\n zero3_init_flag: true\n zero3_save_16bit_model: true\n zero_stage: 3\ndistributed_type: DEEPSPEED\nfsdp_config: {}\nmachine_rank: 0\nmain_process_ip: null\nmain_process_port: null\nmain_training_function: main\nmixed_precision: fp16\nnum_machines: 1\nnum_processes: 2\nuse_cpu: false\n```\n\n```bash\naccelerate launch examples/nlp_example.py --mixed_precision fp16\n```\n\nCurrently, `Accelerate` supports following config through the CLI:\n\n```bash\n`zero_stage`: [0] Disabled, [1] optimizer state partitioning, [2] optimizer+gradient state partitioning and [3] optimizer+gradient+parameter partitioning\n`gradient_accumulation_steps`: Number of training steps to accumulate gradients before averaging and applying them.\n`gradient_clipping`: Enable gradient clipping with value.\n`offload_optimizer_device`: [none] Disable optimizer offloading, [cpu] offload optimizer to CPU, [nvme] offload optimizer to NVMe SSD. Only applicable with ZeRO >= Stage-2.\n`offload_optimizer_nvme_path`: Decides Nvme Path to offload optimizer states. If unspecified, will default to 'none'.\n`offload_param_device`: [none] Disable parameter offloading, [cpu] offload parameters to CPU, [nvme] offload parameters to NVMe SSD. Only applicable with ZeRO Stage-3.\n`offload_param_nvme_path`: Decides Nvme Path to offload parameters. If unspecified, will default to 'none'.\n`zero3_init_flag`: Decides whether to enable `deepspeed.zero.Init` for constructing massive models. Only applicable with ZeRO Stage-3.\n`zero3_save_16bit_model`: Decides whether to save 16-bit model weights when using ZeRO Stage-3.\n`mixed_precision`: `no` for FP32 training, `fp16` for FP16 mixed-precision training and `bf16` for BF16 mixed-precision training.\n`deepspeed_moe_layer_cls_names`: Comma-separated list of transformer Mixture-of-Experts (MoE) layer class names (case-sensitive) to wrap ,e.g, `MixtralSparseMoeBlock`, `Qwen2MoeSparseMoeBlock`, `JetMoEAttention,JetMoEBlock` ...\n`deepspeed_hostfile`: DeepSpeed hostfile for configuring multi-node compute resources.\n`deepspeed_exclusion_filter`: DeepSpeed exclusion filter string when using mutli-node setup.\n`deepspeed_inclusion_filter`: DeepSpeed inclusion filter string when using mutli-node setup.\n`deepspeed_multinode_launcher`: DeepSpeed multi-node launcher to use, e.g. `pdsh`, `standard`, `openmpi`, `mvapich`, `mpich`, `slurm`, `nossh` (requires DeepSpeed >= 0.14.5). If unspecified, will default to `pdsh`.\n`deepspeed_config_file`: path to the DeepSpeed config file in `json` format. See the next section for more details on this.\n```\nTo be able to tweak more options, you will need to use a DeepSpeed config file.\n\n### DeepSpeed Config File\nOn your machine(s) just run:\n\n```bash\naccelerate config\n```\n\nand answer the questions asked. It will ask whether you want to use a config file for deepspeed to which you answer yes\nand provide the path to the deepspeed config file.\nThis will generate a config file that will be used automatically to properly set the\ndefault options when doing\n\n```bash\naccelerate launch my_script.py --args_to_my_script\n```\n\nFor instance, here is how you would run the NLP example `examples/by_feature/deepspeed_with_config_support.py` (from the root of the repo) with DeepSpeed Config File:\n\n**ZeRO Stage-2 DeepSpeed Config File Example**\n```bash\ncompute_environment: LOCAL_MACHINE\ndeepspeed_config:\n deepspeed_config_file: /home/ubuntu/accelerate/examples/deepspeed_config_templates/zero_stage2_config.json\n zero3_init_flag: true\ndistributed_type: DEEPSPEED\nfsdp_config: {}\nmachine_rank: 0\nmain_process_ip: null\nmain_process_port: null\nmain_training_function: main\nmixed_precision: fp16\nnum_machines: 1\nnum_processes: 2\nuse_cpu: false\n```\n\nwith the contents of `zero_stage2_config.json` being:\n```json\n{\n \"fp16\": {\n \"enabled\": true,\n \"loss_scale\": 0,\n \"loss_scale_window\": 1000,\n \"initial_scale_power\": 16,\n \"hysteresis\": 2,\n \"min_loss_scale\": 1\n },\n \"optimizer\": {\n \"type\": \"AdamW\",\n \"params\": {\n \"lr\": \"auto\",\n \"weight_decay\": \"auto\",\n \"torch_adam\": true,\n \"adam_w_mode\": true\n }\n },\n \"scheduler\": {\n \"type\": \"WarmupDecayLR\",\n \"params\": {\n \"warmup_min_lr\": \"auto\",\n \"warmup_max_lr\": \"auto\",\n \"warmup_num_steps\": \"auto\",\n \"total_num_steps\": \"auto\"\n }\n },\n \"zero_optimization\": {\n \"stage\": 2,\n \"allgather_partitions\": true,\n \"allgather_bucket_size\": 2e8,\n \"overlap_comm\": true,\n \"reduce_scatter\": true,\n \"reduce_bucket_size\": \"auto\",\n \"contiguous_gradients\": true\n },\n \"gradient_accumulation_steps\": 1,\n \"gradient_clipping\": \"auto\",\n \"steps_per_print\": 2000,\n \"train_batch_size\": \"auto\",\n \"train_micro_batch_size_per_gpu\": \"auto\",\n \"wall_clock_breakdown\": false\n}\n```\n\n```bash\naccelerate launch examples/by_feature/deepspeed_with_config_support.py \\\n--config_name \"gpt2-large\" \\\n--tokenizer_name \"gpt2-large\" \\\n--dataset_name \"wikitext\" \\\n--dataset_config_name \"wikitext-2-raw-v1\" \\\n--block_size 128 \\\n--output_dir \"./clm/clm_deepspeed_stage2_accelerate\" \\\n--learning_rate 5e-4 \\\n--per_device_train_batch_size 24 \\\n--per_device_eval_batch_size 24 \\\n--num_train_epochs 3 \\\n--with_tracking \\\n--report_to \"wandb\"\\\n```\n\n**ZeRO Stage-3 with CPU offload DeepSpeed Config File Example**\n```bash\ncompute_environment: LOCAL_MACHINE\ndeepspeed_config:\n deepspeed_config_file: /home/ubuntu/accelerate/examples/deepspeed_config_templates/zero_stage3_offload_config.json\n zero3_init_flag: true\ndistributed_type: DEEPSPEED\nfsdp_config: {}\nmachine_rank: 0\nmain_process_ip: null\nmain_process_port: null\nmain_training_function: main\nmixed_precision: fp16\nnum_machines: 1\nnum_processes: 2\nuse_cpu: false\n```\nwith the contents of `zero_stage3_offload_config.json` being:\n```json\n{\n \"fp16\": {\n \"enabled\": true,\n \"loss_scale\": 0,\n \"loss_scale_window\": 1000,\n \"initial_scale_power\": 16,\n \"hysteresis\": 2,\n \"min_loss_scale\": 1\n },\n \"optimizer\": {\n \"type\": \"AdamW\",\n \"params\": {\n \"lr\": \"auto\",\n \"weight_decay\": \"auto\"\n }\n },\n \"scheduler\": {\n \"type\": \"WarmupDecayLR\",\n \"params\": {\n \"warmup_min_lr\": \"auto\",\n \"warmup_max_lr\": \"auto\",\n \"warmup_num_steps\": \"auto\",\n \"total_num_steps\": \"auto\"\n }\n },\n \"zero_optimization\": {\n \"stage\": 3,\n \"offload_optimizer\": {\n \"device\": \"cpu\",\n \"pin_memory\": true\n },\n \"offload_param\": {\n \"device\": \"cpu\",\n \"pin_memory\": true\n },\n \"overlap_comm\": true,\n \"contiguous_gradients\": true,\n \"reduce_bucket_size\": \"auto\",\n \"stage3_prefetch_bucket_size\": \"auto\",\n \"stage3_param_persistence_threshold\": \"auto\",\n \"sub_group_size\": 1e9,\n \"stage3_max_live_parameters\": 1e9,\n \"stage3_max_reuse_distance\": 1e9,\n \"stage3_gather_16bit_weights_on_model_save\": \"auto\"\n },\n \"gradient_accumulation_steps\": 1,\n \"gradient_clipping\": \"auto\",\n \"steps_per_print\": 2000,\n \"train_batch_size\": \"auto\",\n \"train_micro_batch_size_per_gpu\": \"auto\",\n \"wall_clock_breakdown\": false\n}\n```\n\n```bash\naccelerate launch examples/by_feature/deepspeed_with_config_support.py \\\n--config_name \"gpt2-large\" \\\n--tokenizer_name \"gpt2-large\" \\\n--dataset_name \"wikitext\" \\\n--dataset_config_name \"wikitext-2-raw-v1\" \\\n--block_size 128 \\\n--output_dir \"./clm/clm_deepspeed_stage3_offload_accelerate\" \\\n--learning_rate 5e-4 \\\n--per_device_train_batch_size 32 \\\n--per_device_eval_batch_size 32 \\\n--num_train_epochs 3 \\\n--with_tracking \\\n--report_to \"wandb\"\\\n```\n\n**ZeRO++ Config Example**\nYou can use the features of ZeRO++ by using the appropriate config parameters. Note that ZeRO++ is an extension for ZeRO Stage 3. Here is how the config file can be modified, from [DeepSpeed's ZeRO++ tutorial](https://www.deepspeed.ai/tutorials/zeropp/):\n\n```json\n{\n \"zero_optimization\": {\n \"stage\": 3,\n \"reduce_bucket_size\": \"auto\",\n\n \"zero_quantized_weights\": true,\n \"zero_hpz_partition_size\": 8,\n \"zero_quantized_gradients\": true,\n\n \"contiguous_gradients\": true,\n \"overlap_comm\": true\n }\n}\n```\n\nFor hierarchical partitioning, the partition size `zero_hpz_partition_size` should ideally be set to the number of GPUs per node. (For example, the above config file assumes 8 GPUs per node)\n\n**Important code changes when using DeepSpeed Config File**\n\n1. DeepSpeed Optimizers and Schedulers. For more information on these,\nsee the [DeepSpeed Optimizers](https://deepspeed.readthedocs.io/en/latest/optimizers.html) and [DeepSpeed Schedulers](https://deepspeed.readthedocs.io/en/latest/schedulers.html) documentation.\nWe will look at the changes needed in the code when using these.\n\n a. DS Optim + DS Scheduler: The case when both `optimizer` and `scheduler` keys are present in the DeepSpeed config file.\n 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.\n Below is the snippet from `examples/by_feature/deepspeed_with_config_support.py` showing this:\n ```python\n # Creates Dummy Optimizer if `optimizer` was specified in the config file else creates Adam Optimizer\n optimizer_cls = (\n torch.optim.AdamW\n if accelerator.state.deepspeed_plugin is None\n or \"optimizer\" not in accelerator.state.deepspeed_plugin.deepspeed_config\n else DummyOptim\n )\n optimizer = optimizer_cls(optimizer_grouped_parameters, lr=args.learning_rate)\n\n # Creates Dummy Scheduler if `scheduler` was specified in the config file else creates `args.lr_scheduler_type` Scheduler\n if (\n accelerator.state.deepspeed_plugin is None\n or \"scheduler\" not in accelerator.state.deepspeed_plugin.deepspeed_config\n ):\n lr_scheduler = get_scheduler(\n name=args.lr_scheduler_type,\n optimizer=optimizer,\n num_warmup_steps=args.num_warmup_steps,\n num_training_steps=args.max_train_steps,\n )\n else:\n lr_scheduler = DummyScheduler(\n optimizer, total_num_steps=args.max_train_steps, warmup_num_steps=args.num_warmup_steps\n )\n ```\n b. Custom Optim + Custom Scheduler: The case when both `optimizer` and `scheduler` keys are absent in the DeepSpeed config file.\n In this situation, no code changes are needed from the user and this is the case when using integration via DeepSpeed Plugin.\n 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.\n\n c. Custom Optim + DS Scheduler: The case when only `scheduler` key is present in the DeepSpeed config file.\n In this situation, the user has to use `accelerate.utils.DummyScheduler` to replace the PyTorch/Custom scheduler in their code.\n\n d. DS Optim + Custom Scheduler: The case when only `optimizer` key is present in the DeepSpeed config file.\n This will result in an error because you can only use DS Scheduler when using DS Optim.\n\n2. Notice the `auto` values in the above example DeepSpeed config files. These are automatically handled by `prepare` method\nbased on model, dataloaders, dummy optimizer and dummy schedulers provided to `prepare` method.\nOnly the `auto` fields specified in above examples are handled by `prepare` method and the rest have to be explicitly specified by the user.\n\nThe `auto` values are calculated as:\n\n- `reduce_bucket_size`: `hidden_size * hidden_size`\n- `stage3_prefetch_bucket_size`: `int(0.9 * hidden_size * hidden_size)`\n- `stage3_param_persistence_threshold`: `10 * hidden_size`\n\nFor the `auto` feature to work for these 3 config entries - Accelerate will use `model.config.hidden_size` or `max(model.config.hidden_sizes)` as `hidden_size`. If neither of these is available, the launching will fail and you will have to set these 3 config entries manually. Remember the first 2 config entries are the communication buffers - the larger they are the more efficient the comms will be, and the larger they are the more GPU memory they will consume, so it's a tunable performance trade-off.\n\n\n**Things to note when using DeepSpeed Config File**\n\nBelow is a sample script using `deepspeed_config_file` in different scenarios.\n\nCode `test.py`:\n\n```python\nfrom accelerate import Accelerator\nfrom accelerate.state import AcceleratorState\n\n\ndef main():\n accelerator = Accelerator()\n accelerator.print(f\"{AcceleratorState()}\")\n\n\nif __name__ == \"__main__\":\n main()\n```\n\n**Scenario 1**: Manually tampered accelerate config file having `deepspeed_config_file` along with other entries.\n\n1. Content of the `accelerate` config:\n\n```yaml\ncommand_file: null\ncommands: null\ncompute_environment: LOCAL_MACHINE\ndeepspeed_config:\n gradient_accumulation_steps: 1\n gradient_clipping: 1.0\n offload_optimizer_device: 'cpu'\n offload_param_device: 'cpu'\n zero3_init_flag: true\n zero3_save_16bit_model: true\n zero_stage: 3\n deepspeed_config_file: 'ds_config.json'\ndistributed_type: DEEPSPEED\ndowncast_bf16: 'no'\ndynamo_backend: 'NO'\nfsdp_config: {}\ngpu_ids: null\nmachine_rank: 0\nmain_process_ip: null\nmain_process_port: null\nmain_training_function: main\nmegatron_lm_config: {}\nnum_machines: 1\nnum_processes: 2\nrdzv_backend: static\nsame_network: true\ntpu_name: null\ntpu_zone: null\nuse_cpu: false\n```\n\n2. `ds_config.json`:\n\n```json\n{\n \"bf16\": {\n \"enabled\": true\n },\n \"zero_optimization\": {\n \"stage\": 3,\n \"stage3_gather_16bit_weights_on_model_save\": false,\n \"offload_optimizer\": {\n \"device\": \"none\"\n },\n \"offload_param\": {\n \"device\": \"none\"\n }\n },\n \"gradient_clipping\": 1.0,\n \"train_batch_size\": \"auto\",\n \"train_micro_batch_size_per_gpu\": \"auto\",\n \"gradient_accumulation_steps\": 10,\n \"steps_per_print\": 2000000\n}\n```\n\n3. Output of `accelerate launch test.py`:\n\n```bash\nValueError: When using `deepspeed_config_file`, the following accelerate config variables will be ignored:\n['gradient_accumulation_steps', 'gradient_clipping', 'zero_stage', 'offload_optimizer_device', 'offload_param_device',\n'zero3_save_16bit_model', 'mixed_precision'].\nPlease specify them appropriately in the DeepSpeed config file.\nIf you are using an accelerate config file, remove other config variables mentioned in the above specified list.\nThe easiest method is to create a new config following the questionnaire via `accelerate config`.\nIt will only ask for the necessary config variables when using `deepspeed_config_file`.\n```\n\n**Scenario 2**: Use the solution of the error to create new accelerate config and check that no ambiguity error is now thrown.\n\n1. Run `accelerate config`:\n\n```bash\n$ accelerate config\n-------------------------------------------------------------------------------------------------------------------------------\nIn which compute environment are you running?\nThis machine\n-------------------------------------------------------------------------------------------------------------------------------\nWhich type of machine are you using?\nmulti-GPU\nHow many different machines will you use (use more than 1 for multi-node training)? [1]:\nDo you wish to optimize your script with torch dynamo?[yes/NO]:\nDo you want to use DeepSpeed? [yes/NO]: yes\nDo you want to specify a json file to a DeepSpeed config? [yes/NO]: yes\nPlease enter the path to the json DeepSpeed config file: ds_config.json\nDo you want to enable `deepspeed.zero.Init` when using ZeRO Stage-3 for constructing massive models? [yes/NO]: yes\nHow many GPU(s) should be used for distributed training? [1]:4\naccelerate configuration saved at ds_config_sample.yaml\n```\n\n2. Content of the `accelerate` config:\n\n```yaml\ncompute_environment: LOCAL_MACHINE\ndeepspeed_config:\n deepspeed_config_file: ds_config.json\n zero3_init_flag: true\ndistributed_type: DEEPSPEED\ndowncast_bf16: 'no'\ndynamo_backend: 'NO'\nfsdp_config: {}\nmachine_rank: 0\nmain_training_function: main\nmegatron_lm_config: {}\nnum_machines: 1\nnum_processes: 4\nrdzv_backend: static\nsame_network: true\nuse_cpu: false\n```\n\n3. Output of `accelerate launch test.py`:\n\n```bash\nDistributed environment: DEEPSPEED Backend: nccl\nNum processes: 4\nProcess index: 0\nLocal process index: 0\nDevice: cuda:0\nMixed precision type: bf16\nds_config: {'bf16': {'enabled': True}, 'zero_optimization': {'stage': 3, 'stage3_gather_16bit_weights_on_model_save': False, 'offload_optimizer': {'device': 'none'}, 'offload_param': {'device': 'none'}}, 'gradient_clipping': 1.0, 'train_batch_size': 'auto', 'train_micro_batch_size_per_gpu': 'auto', 'gradient_accumulation_steps': 10, 'steps_per_print': inf, 'fp16': {'enabled': False}}\n```\n\n**Scenario 3**: Setting the `accelerate launch` command arguments related to DeepSpeed as `\"auto\"` in the DeepSpeed` configuration file and check that things work as expected.\n\n1. New `ds_config.json` with `\"auto\"` for the `accelerate launch` DeepSpeed command arguments:\n\n```json\n{\n \"bf16\": {\n \"enabled\": \"auto\"\n },\n \"zero_optimization\": {\n \"stage\": \"auto\",\n \"stage3_gather_16bit_weights_on_model_save\": \"auto\",\n \"offload_optimizer\": {\n \"device\": \"auto\"\n },\n \"offload_param\": {\n \"device\": \"auto\"\n }\n },\n \"gradient_clipping\": \"auto\",\n \"train_batch_size\": \"auto\",\n \"train_micro_batch_size_per_gpu\": \"auto\",\n \"gradient_accumulation_steps\": \"auto\",\n \"steps_per_print\": 2000000\n}\n```\n\n2. Output of `accelerate launch --mixed_precision=\"fp16\" --zero_stage=3 --gradient_accumulation_steps=5 --gradient_clipping=1.0 --offload_param_device=\"cpu\" --offload_optimizer_device=\"nvme\" --zero3_save_16bit_model=\"true\" test.py`:\n\n```bash\nDistributed environment: DEEPSPEED Backend: nccl\nNum processes: 4\nProcess index: 0\nLocal process index: 0\nDevice: cuda:0\nMixed precision type: fp16\nds_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}}\n```\n\n**Note**:\n1. Remaining `\"auto\"` values are handled in `accelerator.prepare()` call as explained in point 2 of\n`Important code changes when using DeepSpeed Config File`.\n2. 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.\n\n## Saving and loading\n\n1. Saving and loading of models is unchanged for ZeRO Stage-1 and Stage-2.\n\n2. under ZeRO Stage-3, `state_dict` contains just the placeholders since the model weights are partitioned across multiple GPUs.\nZeRO Stage-3 has 2 options:\n\n a. Saving the entire 16bit model weights to directly load later on using `model.load_state_dict(torch.load(pytorch_model.bin))`.\n For this, either set `zero_optimization.stage3_gather_16bit_weights_on_model_save` to True in DeepSpeed Config file or set\n `zero3_save_16bit_model` to True in DeepSpeed Plugin.\n **Note that this option requires consolidation of the weights on one GPU it can be slow and memory demanding, so only use this feature when needed.**\n Below is the snippet from `examples/by_feature/deepspeed_with_config_support.py` showing this:\n ```python\n unwrapped_model = accelerator.unwrap_model(model)\n\n # New Code #\n # Saves the whole/unpartitioned fp16 model when in ZeRO Stage-3 to the output directory if\n # `stage3_gather_16bit_weights_on_model_save` is True in DeepSpeed Config file or\n # `zero3_save_16bit_model` is True in DeepSpeed Plugin.\n # For Zero Stages 1 and 2, models are saved as usual in the output directory.\n # The model name saved is `pytorch_model.bin`\n unwrapped_model.save_pretrained(\n args.output_dir,\n is_main_process=accelerator.is_main_process,\n save_function=accelerator.save,\n state_dict=accelerator.get_state_dict(model),\n )\n ```\n\n b. To get 32bit weights, first save the model using `model.save_checkpoint()`.\n Below is the snippet from `examples/by_feature/deepspeed_with_config_support.py` showing this:\n ```python\n success = model.save_checkpoint(PATH, ckpt_id, checkpoint_state_dict)\n status_msg = f\"checkpointing: PATH={PATH}, ckpt_id={ckpt_id}\"\n if success:\n logging.info(f\"Success {status_msg}\")\n else:\n logging.warning(f\"Failure {status_msg}\")\n ```\n This will create ZeRO model and optimizer partitions along with `zero_to_fp32.py` script in checkpoint directory.\n You can use this script to do offline consolidation.\n It requires no configuration files or GPUs. Here is an example of its usage:\n ```bash\n $ cd /path/to/checkpoint_dir\n $ ./zero_to_fp32.py . pytorch_model.bin\n Processing zero checkpoint at global_step1\n Detected checkpoint of type zero stage 3, world_size: 2\n Saving fp32 state dict to pytorch_model.bin (total_numel=60506624)\n ```\n To get 32bit model for saving/inference, you can perform:\n ```python\n from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint\n\n unwrapped_model = accelerator.unwrap_model(model)\n fp32_model = load_state_dict_from_zero_checkpoint(unwrapped_model, checkpoint_dir)\n ```\n If you are only interested in the `state_dict`, you can do the following:\n ```python\n from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint\n\n state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir)\n ```\n Note that all these functions require ~2x memory (general RAM) of the size of the final checkpoint.\n\n## ZeRO Inference\nDeepSpeed ZeRO Inference supports ZeRO stage 3 with ZeRO-Infinity.\nIt uses the same ZeRO protocol as training, but it doesn't use an optimizer and a lr scheduler and only stage 3 is relevant.\nWith accelerate integration, you just need to prepare the model and dataloader as shown below:\n\n```python\nmodel, eval_dataloader = accelerator.prepare(model, eval_dataloader)\n```\n\n## Few caveats to be aware of\n\n1. Current integration doesnā€™t support Pipeline Parallelism of DeepSpeed.\n2. Current integration doesnā€™t support `mpu`, limiting the tensor parallelism which is supported in Megatron-LM.\n3. Current integration doesnā€™t support multiple models.\n\n## Multi-node DeepSpeed\nDeepSpeed supports multi-node inference and training over a variety of different launchers. You can specify a different launcher by setting the `deepspeed_multinode_launcher` config in the CLI or in the DeepSpeed config file.\n\nCurrently, accelerate supports passing configuration for the following DeepSpeed multi-node launchers: `pdsh` (default), `standard`, `openmpi`, `mvapich`, `mpich`, `slurm`, `nossh` (requires DeepSpeed >= 0.14.5).\n\nPlease read the [DeepSpeed documentation](https://www.deepspeed.ai/getting-started/#resource-configuration-multi-node) for more information on the different launchers. By default, DeepSpeed will attempt to use passwordless SSH from the main machine node to the other nodes to perform the launcher command. In this configuration, the accelerate launch command only needs to be run on the main node. If using the `nossh` launcher, you will need to run the accelerate launch command on every node using copied configuration. \n\n## DeepSpeed Resources\n\nThe documentation for the internals related to deepspeed can be found [here](../package_reference/deepspeed).\n\n- [Project's github](https://github.com/deepspeedai/DeepSpeed)\n- [Usage docs](https://www.deepspeed.ai/getting-started/)\n- [API docs](https://deepspeed.readthedocs.io/en/latest/index.html)\n- [Blog posts](https://www.microsoft.com/en-us/research/search/?q=deepspeed)\n\nPapers:\n\n- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://huggingface.co/papers/1910.02054)\n- [ZeRO-Offload: Democratizing Billion-Scale Model Training](https://huggingface.co/papers/2101.06840)\n- [ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning](https://huggingface.co/papers/2104.07857)\n- [ZeRO++: Extremely Efficient Collective Communication for Giant Model Training](https://huggingface.co/papers/2306.10209)\n\n\nFinally, please, remember that `Accelerate` only integrates DeepSpeed, therefore if you\nhave any problems or questions with regards to DeepSpeed usage, please, file an issue with [DeepSpeed GitHub](https://github.com/deepspeedai/DeepSpeed/issues).\n\n\n\n\n For those interested in the similarities and differences between FSDP and DeepSpeed, please check out the [concept guide here](../concept_guides/fsdp_and_deepspeed)!\n \n" }, { "path": "docs/source/usage_guides/deepspeed_multiple_model.md", "content": "\n\n# Using multiple models with DeepSpeed\n\n\n\n This guide assumes that you have read and understood the [DeepSpeed usage guide](./deepspeed.md).\n\n\n\nRunning multiple models with Accelerate and DeepSpeed is useful for:\n\n* Knowledge distillation\n* Post-training techniques like RLHF (see the [TRL](https://github.com/huggingface/trl) library for more examples)\n* Training multiple models at once\n\nCurrently, Accelerate has a **very experimental API** to help you use multiple models.\n\nThis tutorial will focus on two common use cases:\n\n1. Knowledge distillation, where a smaller student model is trained to mimic a larger, better-performing teacher. If the student model fits on a single GPU, we can use ZeRO-2 for training and ZeRO-3 to shard the teacher for inference. This is significantly faster than using ZeRO-3 for both models.\n2. Training multiple *disjoint* models at once.\n\n## Knowledge distillation\n\nKnowledge distillation is a good example of using multiple models, but only training one of them.\n\nNormally, you would use a single [`utils.DeepSpeedPlugin`] for both models. However, in this case, there are two separate configurations. Accelerate allows you to create and use multiple plugins **if and only if** they are in a `dict` so that you can reference and enable the proper plugin when needed.\n\n```python\nfrom accelerate.utils import DeepSpeedPlugin\n\nzero2_plugin = DeepSpeedPlugin(hf_ds_config=\"zero2_config.json\")\nzero3_plugin = DeepSpeedPlugin(hf_ds_config=\"zero3_config.json\")\n\ndeepspeed_plugins = {\"student\": zero2_plugin, \"teacher\": zero3_plugin}\n```\n\nThe `zero2_config.json` should be configured for full training (so specify `scheduler` and `optimizer` if you are not utilizing your own), while `zero3_config.json` should only be configured for the inference model, as shown in the example below.\n\n```json\n{\n \"bf16\": {\n \"enabled\": \"auto\"\n },\n \"zero_optimization\": {\n \"stage\": 3,\n \"overlap_comm\": true,\n \"reduce_bucket_size\": \"auto\",\n \"stage3_prefetch_bucket_size\": \"auto\",\n \"stage3_param_persistence_threshold\": \"auto\",\n \"stage3_max_live_parameters\": \"auto\",\n \"stage3_max_reuse_distance\": \"auto\",\n },\n \"train_micro_batch_size_per_gpu\": 1\n}\n```\n\nAn example `zero2_config.json` configuration is shown below.\n\n```json\n{\n \"bf16\": {\n \"enabled\": \"auto\"\n },\n \"optimizer\": {\n \"type\": \"AdamW\",\n \"params\": {\n \"lr\": \"auto\",\n \"weight_decay\": \"auto\",\n \"torch_adam\": true,\n \"adam_w_mode\": true\n }\n },\n \"scheduler\": {\n \"type\": \"WarmupLR\",\n \"params\": {\n \"warmup_min_lr\": \"auto\",\n \"warmup_max_lr\": \"auto\",\n \"warmup_num_steps\": \"auto\"\n }\n },\n \"zero_optimization\": {\n \"stage\": 2,\n \"offload_optimizer\": {\n \"device\": \"cpu\",\n \"pin_memory\": true\n },\n },\n \"gradient_accumulation_steps\": 1,\n \"gradient_clipping\": \"auto\",\n \"train_batch_size\": \"auto\",\n \"train_micro_batch_size_per_gpu\": \"auto\",\n}\n```\n\n\n\n DeepSpeed will raise an error if `train_micro_batch_size_per_gpu` isn't specified, even if this particular model isn't being trained.\n\n\n\nFrom here, create a single [`Accelerator`] and pass in both configurations.\n\n```python\nfrom accelerate import Accelerator\n\naccelerator = Accelerator(deepspeed_plugins=deepspeed_plugins)\n```\n\nNow let's see how to use them.\n\n### Student model\n\nBy default, Accelerate sets the first item in the `dict` as the default or enabled plugin (`\"student\"` plugin). Verify this by using the [`utils.deepspeed.get_active_deepspeed_plugin`] function to see which plugin is enabled.\n\n```python\nactive_plugin = get_active_deepspeed_plugin(accelerator.state)\nassert active_plugin is deepspeed_plugins[\"student\"]\n```\n\n[`AcceleratorState`] also keeps the active DeepSpeed plugin saved in `state.deepspeed_plugin`.\n```python\nassert active_plugin is accelerator.deepspeed_plugin\n```\n\nSince `student` is the currently active plugin, let's go ahead and prepare the model, optimizer, and scheduler.\n\n```python\nstudent_model, optimizer, scheduler = ...\nstudent_model, optimizer, scheduler, train_dataloader = accelerator.prepare(student_model, optimizer, scheduler, train_dataloader)\n```\n\nNow it's time to deal with the teacher model.\n\n### Teacher model\n\nFirst, you need to specify in [`Accelerator`] that the `zero3_config.json` configuration should be used.\n\n```python\naccelerator.state.select_deepspeed_plugin(\"teacher\")\n```\n\nThis disables the `\"student\"` plugin and enables the `\"teacher\"` plugin instead. The\nDeepSpeed stateful config inside of Transformers is updated, and it changes which plugin configuration gets called when using\n`deepspeed.initialize()`. This allows you to use the automatic `deepspeed.zero.Init` context manager integration Transformers provides.\n\n```python\nteacher_model = AutoModel.from_pretrained(...)\nteacher_model = accelerator.prepare(teacher_model)\n```\n\nOtherwise, you should manually initialize the model with `deepspeed.zero.Init`.\n```python\nwith deepspeed.zero.Init(accelerator.deepspeed_plugin.config):\n model = MyModel(...)\n```\n\n### Training\n\nFrom here, your training loop can be whatever you like, as long as `teacher_model` is never being trained on.\n\n```python\nteacher_model.eval()\nstudent_model.train()\nfor batch in train_dataloader:\n with torch.no_grad():\n output_teacher = teacher_model(**batch)\n output_student = student_model(**batch)\n # Combine the losses or modify it in some way\n loss = output_teacher.loss + output_student.loss\n accelerator.backward(loss)\n optimizer.step()\n scheduler.step()\n optimizer.zero_grad()\n```\n\n## Train multiple disjoint models\n\nTraining multiple models is a more complicated scenario.\nIn its current state, we assume each model is **completely disjointed** from the other during training.\n\nThis scenario still requires two [`utils.DeepSpeedPlugin`]'s to be made. However, you also need a second [`Accelerator`], since different `deepspeed` engines are being called at different times. A single [`Accelerator`] can only carry one instance at a time.\n\nSince the [`state.AcceleratorState`] is a stateful object though, it is already aware of both [`utils.DeepSpeedPlugin`]'s available. You can just instantiate a second [`Accelerator`] with no extra arguments.\n\n```python\nfirst_accelerator = Accelerator(deepspeed_plugins=deepspeed_plugins)\nsecond_accelerator = Accelerator()\n```\n\nYou can call either `first_accelerator.state.select_deepspeed_plugin()` to enable or disable\na particular plugin, and then call [`prepare`].\n\n```python\n# can be `accelerator_0`, `accelerator_1`, or by calling `AcceleratorState().select_deepspeed_plugin(...)`\nfirst_accelerator.state.select_deepspeed_plugin(\"first_model\")\nfirst_model = AutoModel.from_pretrained(...)\n# For this example, `get_training_items` is a nonexistent function that gets the setup we need for training\nfirst_optimizer, first_scheduler, train_dl, eval_dl = get_training_items(model1)\nfirst_model, first_optimizer, first_scheduler, train_dl, eval_dl = accelerator.prepare(\n first_model, first_optimizer, first_scheduler, train_dl, eval_dl\n)\n\nsecond_accelerator.state.select_deepspeed_plugin(\"second_model\")\nsecond_model = AutoModel.from_pretrained(...)\n# For this example, `get_training_items` is a nonexistent function that gets the setup we need for training\nsecond_optimizer, second_scheduler, _, _ = get_training_items(model2)\nsecond_model, second_optimizer, second_scheduler = accelerator.prepare(\n second_model, second_optimizer, second_scheduler\n)\n```\n\nAnd now you can train:\n\n```python\nfor batch in dl:\n outputs1 = first_model(**batch)\n first_accelerator.backward(outputs1.loss)\n first_optimizer.step()\n first_scheduler.step()\n first_optimizer.zero_grad()\n \n outputs2 = model2(**batch)\n second_accelerator.backward(outputs2.loss)\n second_optimizer.step()\n second_scheduler.step()\n second_optimizer.zero_grad()\n```\n\n## Resources\n\nTo see more examples, please check out the [related tests](https://github.com/huggingface/accelerate/blob/main/src/accelerate/test_utils/scripts/external_deps/test_ds_multiple_model.py) currently in [Accelerate].\n" }, { "path": "docs/source/usage_guides/distributed_inference.md", "content": "\n\n# Distributed inference\n\nDistributed inference can fall into three brackets:\n\n1. Loading an entire model onto each GPU and sending chunks of a batch through each GPU's model copy at a time\n2. Loading parts of a model onto each GPU and processing a single input at one time\n3. Loading parts of a model onto each GPU and using what is called scheduled Pipeline Parallelism to combine the two prior techniques. \n\nWe're going to go through the first and the last bracket, showcasing how to do each as they are more realistic scenarios.\n\n\n## Sending chunks of a batch automatically to each loaded model\n\nThis is the most memory-intensive solution, as it requires each GPU to keep a full copy of the model in memory at a given time. \n\nNormally 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. \n\nA basic pipeline using the `diffusers` library might look something like so:\n\n```python\nimport torch\nimport torch.distributed as dist\nfrom diffusers import DiffusionPipeline\n\npipe = DiffusionPipeline.from_pretrained(\"runwayml/stable-diffusion-v1-5\", torch_dtype=torch.float16)\n```\nFollowed then by performing inference based on the specific prompt:\n\n```python\ndef run_inference(rank, world_size):\n dist.init_process_group(\"nccl\", rank=rank, world_size=world_size)\n pipe.to(rank)\n\n if torch.distributed.get_rank() == 0:\n prompt = \"a dog\"\n elif torch.distributed.get_rank() == 1:\n prompt = \"a cat\"\n\n result = pipe(prompt).images[0]\n result.save(f\"result_{rank}.png\")\n```\nOne will notice how we have to check the rank to know what prompt to send, which can be a bit tedious.\n\nA user might then also think that with Accelerate, using the `Accelerator` to prepare a dataloader for such a task might also be \na simple way to manage this. (To learn more, check out the relevant section in the [Quick Tour](../quicktour#distributed-evaluation))\n\nCan it manage it? Yes. Does it add unneeded extra code however: also yes.\n\n\nWith Accelerate, we can simplify this process by using the [`Accelerator.split_between_processes`] context manager (which also exists in `PartialState` and `AcceleratorState`). \nThis 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\nto be padded) for you to use right away.\n\nLet's rewrite the above example using this context manager:\n\n```python\nimport torch\nfrom accelerate import PartialState # Can also be Accelerator or AcceleratorState\nfrom diffusers import DiffusionPipeline\n\npipe = DiffusionPipeline.from_pretrained(\"runwayml/stable-diffusion-v1-5\", torch_dtype=torch.float16)\ndistributed_state = PartialState()\npipe.to(distributed_state.device)\n\n# Assume two processes\nwith distributed_state.split_between_processes([\"a dog\", \"a cat\"]) as prompt:\n result = pipe(prompt).images[0]\n result.save(f\"result_{distributed_state.process_index}.png\")\n```\n\nAnd then to launch the code, we can use the Accelerate:\n\nIf you have generated a config file to be used using `accelerate config`:\n\n```bash\naccelerate launch distributed_inference.py\n```\n\nIf you have a specific config file you want to use:\n\n```bash\naccelerate launch --config_file my_config.json distributed_inference.py\n```\n\nOr if don't want to make any config files and launch on two GPUs:\n\n> 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.\n\n```bash\naccelerate launch --num_processes 2 distributed_inference.py\n```\n\nWe've now reduced the boilerplate code needed to split this data to a few lines of code quite easily.\n\nBut what if we have an odd distribution of prompts to GPUs? For example, what if we have 3 prompts, but only 2 GPUs? \n\nUnder the context manager, the first GPU would receive the first two prompts and the second GPU the third, ensuring that \nall prompts are split and no overhead is needed.\n\n*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)\nYou can pass in `apply_padding=True` to ensure that the lists of prompts are padded to the same length, with extra data being taken \nfrom the last sample. This way all GPUs will have the same number of prompts, and you can then gather the results.\n\n\n\nThis is only needed when trying to perform an action such as gathering the results, where the data on each device \nneeds to be the same length. Basic inference does not require this.\n\n\n\nFor instance:\n\n```python\nimport torch\nfrom accelerate import PartialState # Can also be Accelerator or AcceleratorState\nfrom diffusers import DiffusionPipeline\n\npipe = DiffusionPipeline.from_pretrained(\"runwayml/stable-diffusion-v1-5\", torch_dtype=torch.float16)\ndistributed_state = PartialState()\npipe.to(distributed_state.device)\n\n# Assume two processes\nwith distributed_state.split_between_processes([\"a dog\", \"a cat\", \"a chicken\"], apply_padding=True) as prompt:\n result = pipe(prompt).images\n```\n\nOn the first GPU, the prompts will be `[\"a dog\", \"a cat\"]`, and on the second GPU it will be `[\"a chicken\", \"a chicken\"]`.\nMake sure to drop the final sample, as it will be a duplicate of the previous one.\n\nYou can find more complex examples [here](https://github.com/huggingface/accelerate/tree/main/examples/inference/distributed) such as how to use it with LLMs.\n\n## Memory-efficient pipeline parallelism (experimental)\n\nThis next part will discuss using *pipeline parallelism*. This is an **experimental** API that utilizes [torch.distributed.pipelining](https://pytorch.org/docs/stable/distributed.pipelining.html#) as a native solution. \n\nThe general idea with pipeline parallelism is: say you have 4 GPUs and a model big enough it can be *split* on four GPUs using `device_map=\"auto\"`. With this method you can send in 4 inputs at a time (for example here, any amount works) and each model chunk will work on an input, then receive the next input once the prior chunk finished, making it *much* more efficient **and faster** than the method described earlier. Here's a visual taken from the PyTorch repository:\n\n![Pipeline parallelism example](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/accelerate/pipeline_parallel.png)\n\nTo illustrate how you can use this with Accelerate, we have created an [example zoo](https://github.com/huggingface/accelerate/tree/main/examples/inference) showcasing a number of different models and situations. In this tutorial, we'll show this method for GPT2 across two GPUs.\n\nBefore you proceed, please make sure you have the latest PyTorch version installed by running the following:\n\n```bash\npip install torch\n```\n\nStart by creating the model on the CPU:\n\n```{python}\nfrom transformers import GPT2ForSequenceClassification, GPT2Config\n\nconfig = GPT2Config()\nmodel = GPT2ForSequenceClassification(config)\nmodel.eval()\n```\n\nNext you'll need to create some example inputs to use. These help `torch.distributed.pipelining` trace the model.\n\n\n However you make this example will determine the relative batch size that will be used/passed\n through the model at a given time, so make sure to remember how many items there are!\n\n\n```{python}\ninput = torch.randint(\n low=0,\n high=config.vocab_size,\n size=(2, 1024), # bs x seq_len\n device=\"cpu\",\n dtype=torch.int64,\n requires_grad=False,\n)\n```\nNext we need to actually perform the tracing and get the model ready. To do so, use the [`inference.prepare_pippy`] function and it will fully wrap the model for pipeline parallelism automatically:\n\n```{python}\nfrom accelerate.inference import prepare_pippy\nexample_inputs = {\"input_ids\": input}\nmodel = prepare_pippy(model, example_args=(input,))\n```\n\n\n\n There are a variety of parameters you can pass through to `prepare_pippy`:\n \n * `split_points` lets you determine what layers to split the model at. By default we use wherever `device_map=\"auto\" declares, such as `fc` or `conv1`.\n\n * `num_chunks` determines how the batch will be split and sent to the model itself (so `num_chunks=1` with four split points/four GPUs will have a naive MP where a single input gets passed between the four layer split points)\n\n\n\nFrom here, all that's left is to actually perform the distributed inference!\n\n\n\nWhen passing inputs, we highly recommend to pass them in as a tuple of arguments. Using `kwargs` is supported, however, this approach is experimental.\n\n\n```{python}\nargs = some_more_arguments\nwith torch.no_grad():\n output = model(*args)\n```\n\nWhen finished all the data will be on the last process only:\n\n```{python}\nfrom accelerate import PartialState\nif PartialState().is_last_process:\n print(output)\n```\n\n\n\n If you pass in `gather_output=True` to [`inference.prepare_pippy`], the output will be sent\n across to all the GPUs afterwards without needing the `is_last_process` check. This is \n `False` by default as it incurs a communication call.\n \n\n\nAnd that's it! To explore more, please check out the inference examples in the [Accelerate repo](https://github.com/huggingface/accelerate/tree/main/examples/inference/pippy) and our [documentation](../package_reference/inference) as we work to improving this integration. \n" }, { "path": "docs/source/usage_guides/explore.md", "content": "\n\n# Start Here!\n\nPlease use the interactive tool below to help you get started with learning about a particular \nfeature of Accelerate and how to utilize it! It will provide you with a code diff, an explanation\ntowards what is going on, as well as provide you with some useful links to explore more within\nthe documentation!\n\nMost code examples start from the following python code before integrating Accelerate in some way:\n\n```python\nfor batch in dataloader:\n optimizer.zero_grad()\n inputs, targets = batch\n inputs = inputs.to(device)\n targets = targets.to(device)\n outputs = model(inputs)\n loss = loss_function(outputs, targets)\n loss.backward()\n optimizer.step()\n scheduler.step()\n```\n\n
\n\t\n
\n
\n \n
\n" }, { "path": "docs/source/usage_guides/fsdp.md", "content": "\n\n# Fully Sharded Data Parallel\n\nTo accelerate training huge models on larger batch sizes, we can use a fully sharded data parallel model.\nThis type of data parallel paradigm enables fitting more data and larger models by sharding the optimizer states, gradients and parameters.\nTo read more about it and the benefits, check out the [Fully Sharded Data Parallel blog](https://pytorch.org/blog/introducing-pytorch-fully-sharded-data-parallel-api/).\nWe have integrated the latest PyTorch's Fully Sharded Data Parallel (FSDP) training feature.\nAll you need to do is enable it through the config.\n\n## How it works out of the box\n\nOn your machine(s) just run:\n\n```bash\naccelerate config\n```\n\nand answer the questions asked. This will generate a config file that will be used automatically to properly set the\ndefault options when doing\n\n```bash\naccelerate launch my_script.py --args_to_my_script\n```\n\nFor instance, here is how you would run `examples/nlp_example.py` (from the root of the repo) with FSDP enabled:\n\n```bash\ncompute_environment: LOCAL_MACHINE\ndebug: false\ndistributed_type: FSDP\ndowncast_bf16: 'no'\nfsdp_config:\n fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP\n fsdp_backward_prefetch_policy: BACKWARD_PRE\n fsdp_forward_prefetch: false\n fsdp_cpu_ram_efficient_loading: true\n fsdp_offload_params: false\n fsdp_sharding_strategy: FULL_SHARD\n fsdp_state_dict_type: SHARDED_STATE_DICT\n fsdp_sync_module_states: true\n fsdp_transformer_layer_cls_to_wrap: BertLayer\n fsdp_use_orig_params: true\nmachine_rank: 0\nmain_training_function: main\nmixed_precision: bf16\nnum_machines: 1\nnum_processes: 2\nrdzv_backend: static\nsame_network: true\ntpu_env: []\ntpu_use_cluster: false\ntpu_use_sudo: false\nuse_cpu: false\n```\n\n```bash\naccelerate launch examples/nlp_example.py\n```\n\nCurrently, `Accelerate` supports the following config through the CLI:\n\n`fsdp_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). For more information, please refer the official [PyTorch docs](https://pytorch.org/docs/stable/fsdp.html#torch.distributed.fsdp.ShardingStrategy).\n\n`fsdp_offload_params` : Decides Whether to offload parameters and gradients to CPU\n\n`fsdp_auto_wrap_policy`: [1] TRANSFORMER_BASED_WRAP, [2] SIZE_BASED_WRAP, [3] NO_WRAP\n\n`fsdp_transformer_layer_cls_to_wrap`: Only applicable for Transformers. When using `fsdp_auto_wrap_policy=TRANSFORMER_BASED_WRAP`, a user may provide a 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 layers) should not end up in different FSDP wrapped units. Using this policy, wrapping happens for each block containing Multi-Head Attention followed by a 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. It will try to use `model._no_split_modules` when possible.\n\n`fsdp_min_num_params`: minimum number of parameters when using `fsdp_auto_wrap_policy=SIZE_BASED_WRAP`.\n\n`fsdp_backward_prefetch_policy`: [1] BACKWARD_PRE, [2] BACKWARD_POST, [3] NO_PREFETCH\n\n`fsdp_forward_prefetch`: if True, then FSDP explicitly prefetches the next upcoming all-gather while executing in the forward pass. Should only be used for static-graph models since the prefetching follows the first iterationā€™s execution order. i.e., if the sub-modules' order changes dynamically during the model's execution do not enable this feature.\n\n`fsdp_state_dict_type`: [1] FULL_STATE_DICT, [2] LOCAL_STATE_DICT, [3] SHARDED_STATE_DICT\n\n`fsdp_use_orig_params`: If True, allows non-uniform `requires_grad` during init, which means support for interspersed frozen and trainable parameters. This setting is useful in cases such as parameter-efficient fine-tuning as discussed in [this post](https://dev-discuss.pytorch.org/t/rethinking-pytorch-fully-sharded-data-parallel-fsdp-from-first-principles/1019). This option also allows one to have multiple optimizer param groups. This should be `True` when creating an optimizer before preparing/wrapping the model with FSDP.\n\n`fsdp_cpu_ram_efficient_loading`: Only applicable for Transformers models. If True, only the first process loads the pretrained model checkpoint while all other processes have empty weights. This should be set to False if you experience errors when loading the pretrained Transformers model via `from_pretrained` method. When this setting is True `fsdp_sync_module_states` also must to be True, otherwise all the processes except the main process would have random weights leading to unexpected behaviour during training. For this to work, make sure the distributed process group is initialized before calling Transformers `from_pretrained` method. When using Trainer API, the distributed process group is initialized when you create an instance of `TrainingArguments` class.\n\n`fsdp_sync_module_states`: If True, each individually wrapped FSDP unit will broadcast module parameters from rank 0.\n\n\nFor additional and more nuanced control, you can specify other FSDP parameters via `FullyShardedDataParallelPlugin`.\nWhen creating `FullyShardedDataParallelPlugin` object, pass it the parameters that weren't part of the accelerate config or if you want to override them.\nThe 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.\n\nBelow is an example:\n\n```py\nfrom accelerate import FullyShardedDataParallelPlugin\nfrom torch.distributed.fsdp.fully_sharded_data_parallel import FullOptimStateDictConfig, FullStateDictConfig\n\nfsdp_plugin = FullyShardedDataParallelPlugin(\n state_dict_config=FullStateDictConfig(offload_to_cpu=False, rank0_only=False),\n optim_state_dict_config=FullOptimStateDictConfig(offload_to_cpu=False, rank0_only=False),\n)\n\naccelerator = Accelerator(fsdp_plugin=fsdp_plugin)\n```\n\n## Saving and loading\n\nThe new recommended way of checkpointing when using FSDP models is to use `SHARDED_STATE_DICT` as `StateDictType` when setting up the accelerate config.\nBelow is the code snippet to save using `save_state` utility of accelerate.\n\n```py\naccelerator.save_state(\"ckpt\")\n```\n\nInspect the checkpoint folder to see model and optimizer as shards per process:\n```\nls ckpt\n# optimizer_0 pytorch_model_0 random_states_0.pkl random_states_1.pkl scheduler.bin\n\ncd ckpt\n\nls optimizer_0\n# __0_0.distcp __1_0.distcp\n\nls pytorch_model_0\n# __0_0.distcp __1_0.distcp\n```\n\nTo load them back for resuming the training, use the `load_state` utility of accelerate\n\n```py\naccelerator.load_state(\"ckpt\")\n```\n\nWhen using transformers `save_pretrained`, pass `state_dict=accelerator.get_state_dict(model)` to save the model state dict.\n Below is an example:\n\n```diff\n unwrapped_model.save_pretrained(\n args.output_dir,\n is_main_process=accelerator.is_main_process,\n save_function=accelerator.save,\n+ state_dict=accelerator.get_state_dict(model),\n)\n```\n\n### State Dict\n\n`accelerator.get_state_dict` will call the underlying `model.state_dict` implementation using `FullStateDictConfig(offload_to_cpu=True, rank0_only=True)` context manager to get the state dict only for rank 0 and it will be offloaded to CPU.\n\nYou 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).\n\nIf you choose to use `StateDictType.SHARDED_STATE_DICT`, the weights of the model during `Accelerator.save_state` will be split into `n` files for each sub-split on the model. To merge them back into\na single dictionary to load back into the model later after training you can use the `merge_weights` utility:\n\n```py\nfrom accelerate.utils import merge_fsdp_weights\n\n# Our weights are saved usually in a `pytorch_model_fsdp_{model_number}` folder\nmerge_fsdp_weights(\"pytorch_model_fsdp_0\", \"output_path\", safe_serialization=True)\n```\nThe final output will then either be saved to `model.safetensors` or `pytorch_model.bin` (if `safe_serialization=False` is passed). \n\nThis can also be called using the CLI:\n```bash\naccelerate merge-weights pytorch_model_fsdp_0/ output_path\n```\n\n\n## Mapping between FSDP sharding strategies and DeepSpeed ZeRO Stages\n* `FULL_SHARD` maps to the DeepSpeed `ZeRO Stage-3`. Shards optimizer states, gradients and parameters.\n* `SHARD_GRAD_OP` maps to the DeepSpeed `ZeRO Stage-2`. Shards optimizer states and gradients.\n* `NO_SHARD` maps to `ZeRO Stage-0`. No sharding wherein each GPU has full copy of model, optimizer states and gradients.\n* `HYBRID_SHARD` maps to `ZeRO++ Stage-3` wherein `zero_hpz_partition_size=`. Here, this will shard optimizer states, gradients and parameters within each node while each node has full copy.\n\n## A few caveats to be aware of\n\n- In case of multiple models, pass the optimizers to the prepare call in the same order as corresponding models else `accelerator.save_state()` and `accelerator.load_state()` will result in wrong/unexpected behaviour.\n- This feature is incompatible with `--predict_with_generate` in the `run_translation.py` script of `Transformers` library.\n\nFor more control, users can leverage the `FullyShardedDataParallelPlugin`. After creating an instance of this class, users can pass it to the Accelerator class instantiation.\nFor more information on these options, please refer to the PyTorch [FullyShardedDataParallel](https://github.com/pytorch/pytorch/blob/0df2e863fbd5993a7b9e652910792bd21a516ff3/torch/distributed/fsdp/fully_sharded_data_parallel.py#L236) code.\n\n\n\n\n For those interested in the similarities and differences between FSDP and DeepSpeed, please check out the [concept guide here](../concept_guides/fsdp_and_deepspeed)!\n \n" }, { "path": "docs/source/usage_guides/gaudi.md", "content": "\n\n# Intel Gaudi\n\nUsers can take advantage of Intel Gaudi AI accelerators for significantly faster and cost-effective model training and inference.\nThe Intel Gaudi AI accelerator family currently includes three product generations: [Intel Gaudi 1](https://habana.ai/products/gaudi/), [Intel Gaudi 2](https://habana.ai/products/gaudi2/), and [Intel Gaudi 3](https://habana.ai/products/gaudi3/). Each server is equipped with 8 devices, known as Habana Processing Units (HPUs), providing 128GB of memory on Gaudi 3, 96GB on Gaudi 2, and 32GB on the first-gen Gaudi. For more details on the underlying hardware architecture, check out the [Gaudi Architecture Overview](https://docs.habana.ai/en/latest/Gaudi_Overview/Gaudi_Architecture.html).\n\n## How it works out of the box\n\nIt is enabled by default if an Intel Gaudi device is detected.\nTo disable it, pass `--cpu` flag to `accelerate launch` command or answer the corresponding question when answering the `accelerate config` questionnaire.\n\nYou can directly run the following script to test it out on Intel Gaudi:\n\n```bash\naccelerate launch /examples/cv_example.py --data_dir images\n```\n\n## Limitations\n\nThe following features are not part of the Accelerate library and requires [Optimum for Intel Gaudi](https://huggingface.co/docs/optimum/main/en/habana/index):\n\n- `fast_ddp` which implements DDP by applying an all-reduce on gradients instead of the Torch DDP wrapper.\n- `minimize_memory` which is used for fp8 training and enables keeping fp8 weights in memory between the forward and backward passes, leading to a smaller memory footprint at the cost of additional fp8 casts.\n- `context_parallel_size` which is used for Context/Sequence Parallelism (CP/SP) and partitions the network inputs and activations along sequence dimension to reduce memory footprint and increase throughput.\n" }, { "path": "docs/source/usage_guides/gradient_accumulation.md", "content": "\n\n# Performing gradient accumulation with Accelerate\n\nGradient accumulation is a technique where you can train on bigger batch sizes than \nyour machine would normally be able to fit into memory. This is done by accumulating gradients over\nseveral batches, and only stepping the optimizer after a certain number of batches have been performed.\n\nWhile technically standard gradient accumulation code would work fine in a distributed setup, it is not the most efficient\nmethod for doing so and you may experience considerable slowdowns!\n\nIn this tutorial you will see how to quickly setup gradient accumulation and perform it with the utilities provided in Accelerate,\nwhich can total to adding just one new line of code!\n\nThis example will use a very simplistic PyTorch training loop that performs gradient accumulation every two batches:\n\n```python\ndevice = \"cuda\"\nmodel.to(device)\n\ngradient_accumulation_steps = 2\n\nfor index, batch in enumerate(training_dataloader):\n inputs, targets = batch\n inputs = inputs.to(device)\n targets = targets.to(device)\n outputs = model(inputs)\n loss = loss_function(outputs, targets)\n loss = loss / gradient_accumulation_steps\n loss.backward()\n if (index + 1) % gradient_accumulation_steps == 0:\n optimizer.step()\n scheduler.step()\n optimizer.zero_grad()\n```\n\n## Converting it to Accelerate\n\nFirst the code shown earlier will be converted to utilize Accelerate without the special gradient accumulation helper:\n\n```diff\n+ from accelerate import Accelerator\n+ accelerator = Accelerator()\n\n+ model, optimizer, training_dataloader, scheduler = accelerator.prepare(\n+ model, optimizer, training_dataloader, scheduler\n+ )\n\n for index, batch in enumerate(training_dataloader):\n inputs, targets = batch\n- inputs = inputs.to(device)\n- targets = targets.to(device)\n outputs = model(inputs)\n loss = loss_function(outputs, targets)\n loss = loss / gradient_accumulation_steps\n+ accelerator.backward(loss)\n if (index+1) % gradient_accumulation_steps == 0:\n optimizer.step()\n scheduler.step()\n optimizer.zero_grad()\n```\n\n\n\n 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)!\n\n\n\n## Letting Accelerate handle gradient accumulation\n\nAll that is left now is to let Accelerate handle the gradient accumulation for us. To do so you should pass in a `gradient_accumulation_steps` parameter to [`Accelerator`], dictating the number \nof steps to perform before each call to `step()` and how to automatically adjust the loss during the call to [`~Accelerator.backward`]:\n\n```diff\n from accelerate import Accelerator\n- accelerator = Accelerator()\n+ accelerator = Accelerator(gradient_accumulation_steps=2)\n```\n\nAlternatively, 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. \nRead more about that in the [GradientAccumulationPlugin](../package_reference/accelerator#accelerate.utils.GradientAccumulationPlugin) docs.\n\nFrom here you can use the [`~Accelerator.accumulate`] context manager from inside your training loop to automatically perform the gradient accumulation for you!\nYou just wrap it around the entire training part of our code: \n\n```diff\n- for index, batch in enumerate(training_dataloader):\n+ for batch in training_dataloader:\n+ with accelerator.accumulate(model):\n inputs, targets = batch\n outputs = model(inputs)\n```\n\nYou can remove all the special checks for the step number and the loss adjustment:\n\n```diff\n- loss = loss / gradient_accumulation_steps\n accelerator.backward(loss)\n- if (index+1) % gradient_accumulation_steps == 0:\n optimizer.step()\n scheduler.step()\n optimizer.zero_grad()\n```\n\nAs 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. \n\n\n\nTypically with gradient accumulation, you would need to adjust the number of steps to reflect the change in total batches you are \ntraining on. Accelerate automagically does this for you by default. Behind the scenes we instantiate a [`GradientAccumulationPlugin`] configured to do this.\n\n\n\n\n\nThe [`state.GradientState`] is sync'd with the active dataloader being iterated upon. As such it assumes naively that when we have reached the end of the dataloader everything will sync and a step will be performed. To disable this, set `sync_with_dataloader` to be `False` in the [`GradientAccumulationPlugin`]:\n\n```{python}\nfrom accelerate import Accelerator\nfrom accelerate.utils import GradientAccumulationPlugin\n\nplugin = GradientAccumulationPlugin(sync_with_dataloader=False)\naccelerator = Accelerator(..., gradient_accumulation_plugin=plugin)\n```\n\n\n\n## The finished code\n\nBelow is the finished implementation for performing gradient accumulation with Accelerate\n\n```python\nfrom accelerate import Accelerator\naccelerator = Accelerator(gradient_accumulation_steps=2)\nmodel, optimizer, training_dataloader, scheduler = accelerator.prepare(\n model, optimizer, training_dataloader, scheduler\n)\nfor batch in training_dataloader:\n with accelerator.accumulate(model):\n inputs, targets = batch\n outputs = model(inputs)\n loss = loss_function(outputs, targets)\n accelerator.backward(loss)\n optimizer.step()\n scheduler.step()\n optimizer.zero_grad()\n```\n\n\n\nIt's important that **only one forward/backward** should be done inside the context manager `with accelerator.accumulate(model)`.\n\n\n\n\nTo learn more about what magic this wraps around, read the [Gradient Synchronization concept guide](../concept_guides/gradient_synchronization)\n\n\n## Self-contained example\n\nHere is a self-contained example that you can run to see gradient accumulation in action with Accelerate:\n\n```python\nimport torch\nimport copy\nfrom accelerate import Accelerator\nfrom accelerate.utils import set_seed\nfrom torch.utils.data import TensorDataset, DataLoader\n\n# seed\nset_seed(0)\n\n# define toy inputs and labels\nx = torch.tensor([1., 2., 3., 4., 5., 6., 7., 8.])\ny = torch.tensor([2., 4., 6., 8., 10., 12., 14., 16.])\ngradient_accumulation_steps = 4\nper_device_batch_size = len(x) // gradient_accumulation_steps\n\n# define dataset and dataloader\ndataset = TensorDataset(x, y)\ndataloader = DataLoader(dataset, batch_size=per_device_batch_size)\n\n# define model, optimizer and loss function\nclass SimpleLinearModel(torch.nn.Module):\n def __init__(self):\n super(SimpleLinearModel, self).__init__()\n self.weight = torch.nn.Parameter(torch.zeros((1, 1)))\n\n def forward(self, inputs):\n return inputs @ self.weight\n\nmodel = SimpleLinearModel()\nmodel_clone = copy.deepcopy(model)\ncriterion = torch.nn.MSELoss()\nmodel_optimizer = torch.optim.SGD(model.parameters(), lr=0.02)\naccelerator = Accelerator(gradient_accumulation_steps=gradient_accumulation_steps)\nmodel, model_optimizer, dataloader = accelerator.prepare(model, model_optimizer, dataloader)\nmodel_clone_optimizer = torch.optim.SGD(model_clone.parameters(), lr=0.02)\nprint(f\"initial model weight is {model.weight.mean().item():.5f}\")\nprint(f\"initial model weight is {model_clone.weight.mean().item():.5f}\")\nfor i, (inputs, labels) in enumerate(dataloader):\n with accelerator.accumulate(model):\n inputs = inputs.view(-1, 1)\n print(i, inputs.flatten())\n labels = labels.view(-1, 1)\n outputs = model(inputs)\n loss = criterion(outputs, labels)\n accelerator.backward(loss)\n model_optimizer.step()\n model_optimizer.zero_grad()\nloss = criterion(x.view(-1, 1) @ model_clone.weight, y.view(-1, 1))\nmodel_clone_optimizer.zero_grad()\nloss.backward()\nmodel_clone_optimizer.step()\nprint(f\"w/ accumulation, the final model weight is {model.weight.mean().item():.5f}\")\nprint(f\"w/o accumulation, the final model weight is {model_clone.weight.mean().item():.5f}\")\n```\n```\ninitial model weight is 0.00000\ninitial model weight is 0.00000\n0 tensor([1., 2.])\n1 tensor([3., 4.])\n2 tensor([5., 6.])\n3 tensor([7., 8.])\nw/ accumulation, the final model weight is 2.04000\nw/o accumulation, the final model weight is 2.04000\n```\n\n## Gradient accumulation on training samples of variable size\n\nAs was pointed out in this [blog-post](https://huggingface.co/blog/gradient_accumulation), which points out a common error that occurs when performing gradient accumulation on training samples of variable size:\n\n> [...] for gradient accumulation across token-level tasks like causal LM training, the correct loss should be computed by the **total loss across all batches in a gradient accumulation step** divided by the **total number of all non padding tokens in those batches**. This is not the same as the average of the per-batch loss values. \n\nIn other words, some adjustments must be made on losses that operate on a token-level basis.\n\n### Skeleton code\n\n```python\nfrom accelerate import Accelerator\nimport math\nimport contextlib\n\ngradient_accumulation_steps = 2\naccelerator = Accelerator(gradient_accumulation_steps=gradient_accumulation_steps)\nmodel, optimizer, training_dataloader, scheduler = accelerator.prepare(\n model, optimizer, training_dataloader, scheduler\n)\n\ntraining_iterator = iter(training_dataloader)\nnum_samples_in_epoch = len(training_dataloader)\nremainder = num_samples_in_epoch % gradient_accumulation_steps\nremainder = remainder if remainder != 0 else gradient_accumulation_steps\ntotal_updates = math.ceil(num_samples_in_epoch / gradient_accumulation_steps)\n \n\ntotal_batched_samples = 0\nfor update_step in range(total_updates):\n # In order to correctly the total number of non-padded tokens on which we'll compute the cross-entropy loss\n # we need to pre-load the full local batch - i.e the next per_device_batch_size * accumulation_steps samples\n batch_samples = []\n num_batches_in_step = gradient_accumulation_steps if update_step != (total_updates - 1) else remainder\n for _ in range(num_batches_in_step):\n batch_samples += [next(training_iterator)]\n \n # get local num items in batch \n num_items_in_batch = sum([(batch[\"labels\"].ne(-100)).sum() for batch in batch_samples])\n # to compute it correctly in a multi-device DDP training, we need to gather the total number of items in the full batch.\n num_items_in_batch = accelerator.gather(num_items_in_batch).sum().item()\n \n for i, batch in enumerate(batch_samples):\n # if we perform gradient accumulation in a multi-devices set-up, we want to avoid unnecessary communications when accumulating\n # cf: https://muellerzr.github.io/blog/gradient_accumulation.html\n if (i < len(batch_samples) - 1 and accelerator.num_processes > 1):\n ctx = model.no_sync\n else:\n ctx = contextlib.nullcontext\n \n total_batched_samples += 1\n\n with ctx():\n inputs, targets = batch\n outputs = model(inputs)\n loss = loss_function(outputs, targets) # the loss function should sum over samples rather than averaging\n \n # We multiply by num_processes because the DDP calculates the average gradient across all devices whereas dividing by num_items_in_batch already takes into account all devices\n # Same reason for gradient_accumulation_steps, but this times it's Accelerate that calculate the average gradient across the accumulated steps\n loss = (loss * gradient_accumulation_steps * accelerator.num_processes) / num_items_in_batch\n \n accelerator.backward(loss)\n\n # Sync gradients and perform optimization steps once every gradient_accumulation_steps\n optimizer.step()\n scheduler.step()\n optimizer.zero_grad()\n```\n\n### Self-contained causal LM example\n\n```py\nimport torch\nimport copy\nfrom accelerate import Accelerator\nfrom accelerate.utils import set_seed\nfrom accelerate.logging import get_logger\nfrom torch.utils.data import Dataset, DataLoader\nimport math\nimport contexlib\n\n# seed\nset_seed(0)\nlogger = get_logger(__name__)\n\nclass MyDataset(Dataset):\n def __init__(self, num_samples):\n super().__init__()\n self.len = num_samples\n\n def __getitem__(self, index):\n input_ids = torch.arange(1, index+2, dtype=torch.float32)\n labels = torch.remainder(input_ids, 2)\n return {\"input_ids\": input_ids, \"labels\": labels}\n\n def __len__(self):\n return self.len\n \ndef collate_fn(features):\n input_ids = torch.nn.utils.rnn.pad_sequence([f[\"input_ids\"] for f in features], batch_first=True, padding_value=-100)\n labels = torch.nn.utils.rnn.pad_sequence([f[\"labels\"] for f in features], batch_first=True, padding_value=-100)\n return {\"input_ids\": input_ids[..., None], \"labels\": labels[..., None]}\n\n# define toy inputs and labels\ngradient_accumulation_steps = 2\nper_device_batch_size = 4\n\n# define accelerator\naccelerator = Accelerator(gradient_accumulation_steps=gradient_accumulation_steps)\n\n# define dataset and dataloader\n# for this toy example, we'll compute gradient descent over one single global batch\ndataset = MyDataset(per_device_batch_size*gradient_accumulation_steps*accelerator.num_processes)\ndataloader = DataLoader(dataset, batch_size=per_device_batch_size, collate_fn=collate_fn)\n\n# define model, model_optimizer and loss function\nmodel = torch.nn.Linear(1, 2, bias=False)\nmodel_clone = copy.deepcopy(model)\ncriterion = torch.nn.CrossEntropyLoss(reduction=\"sum\") # must sum over samples rather than averaging\nmodel_optimizer = torch.optim.SGD(model.parameters(), lr=0.08)\n\n\nlogger.warning(f\"initial model weight is {model.weight.detach().cpu().squeeze()}\")\nlogger.warning(f\"initial model clone weight is {model_clone.weight.detach().cpu().squeeze()}\")\n\n# prepare artifacts - accelerator handles device placement and dataloader splitting\nmodel, model_optimizer = accelerator.prepare(model, model_optimizer)\ndataloader = accelerator.prepare_data_loader(dataloader, device_placement=True)\ntraining_iterator = iter(dataloader)\n\nnum_samples_in_epoch = len(dataloader)\nremainder = num_samples_in_epoch % gradient_accumulation_steps\nremainder = remainder if remainder != 0 else gradient_accumulation_steps\ntotal_gradient_updates = math.ceil(num_samples_in_epoch / gradient_accumulation_steps)\n\ntotal_batched_samples = 0\nfor update_step in range(total_gradient_updates):\n # In order to correctly the total number of non-padded tokens on which we'll compute the cross-entropy loss\n # we need to pre-load the full local batch - i.e the next per_device_batch_size * accumulation_steps samples\n batch_samples = []\n num_batches_in_step = gradient_accumulation_steps if update_step != (total_gradient_updates - 1) else remainder\n for _ in range(num_batches_in_step):\n batch_samples += [next(training_iterator)]\n \n # get local num items in batch \n local_num_items_in_batch = sum([(batch[\"labels\"].ne(-100)).sum() for batch in batch_samples])\n logger.warning(f\"Step {update_step} - Device {accelerator.process_index} - num items in the local batch {local_num_items_in_batch}\", main_process_only=False)\n\n # to compute it correctly in a multi-device DDP training, we need to gather the total number of items in the full batch.\n num_items_in_batch = accelerator.gather(local_num_items_in_batch).sum().item()\n logger.warning(f\"Total num items {num_items_in_batch}\")\n\n for i, batch in enumerate(batch_samples):\n inputs, labels = batch[\"input_ids\"], batch[\"labels\"]\n total_batched_samples += 1\n # if we perform gradient accumulation in a multi-devices set-up, we want to avoid unnecessary communications when accumulating\n # cf: https://muellerzr.github.io/blog/gradient_accumulation.html\n if (i < len(batch_samples) - 1 and accelerator.num_processes > 1):\n ctx = model.no_sync\n else:\n ctx = contextlib.nullcontext\n with ctx():\n\n outputs = model(inputs)\n loss = criterion(outputs.view(-1, 2), labels.view(-1).to(torch.int64))\n \n # We multiply by num_processes because the DDP calculates the average gradient across all devices whereas dividing by num_items_in_batch already takes into account all devices\n # Same reason for gradient_accumulation_steps, but this times it's Accelerate that calculate the average gradient across the accumulated steps \n loss = (loss * gradient_accumulation_steps * accelerator.num_processes) / num_items_in_batch\n accelerator.backward(loss)\n model_optimizer.step()\n model_optimizer.zero_grad()\n \n\nlogger.warning(f\"Device {accelerator.process_index} - w/ accumulation, the final model weight is {accelerator.unwrap_model(model).weight.detach().cpu().squeeze()}\", main_process_only=False)\n\n# We know do the same operation but on a single device and without gradient accumulation\n\nif accelerator.is_main_process:\n # prepare one single entire batch\n dataloader = DataLoader(dataset, batch_size=len(dataset), collate_fn=collate_fn)\n full_batch_without_accum = next(iter(dataloader))\n total_inputs, total_labels = full_batch_without_accum[\"input_ids\"], full_batch_without_accum[\"labels\"]\n model_clone_optimizer = torch.optim.SGD(model_clone.parameters(), lr=0.08)\n \n # train the cloned model\n loss = torch.nn.CrossEntropyLoss(reduction=\"mean\")(model_clone(total_inputs).view(-1, 2), total_labels.view(-1).to(torch.int64))\n model_clone_optimizer.zero_grad()\n loss.backward()\n model_clone_optimizer.step()\n \n # We should have the same final weights.\n logger.warning(f\"w/o accumulation, the final model weight is {model_clone.weight.detach().cpu().squeeze()}\")\n\n```\n\nResults on a single device - gradient accumulation steps set to 1 and batch_size set to 8:\n```\ninitial model weight is tensor([-0.0075, 0.5364])\ninitial model clone weight is tensor([-0.0075, 0.5364])\nStep 0 - Device 0 - num items in the local batch 36\nTotal num items 36\nDevice 0 - w/ accumulation, the final model weight is tensor([0.0953, 0.4337])\nw/o accumulation, the final model weight is tensor([0.0953, 0.4337])\n```\n\nResults on a two devices set-up - gradient accumulation steps set to 2 and batch_size set to 4.\n```\ninitial model weight is tensor([-0.0075, 0.5364])\ninitial model clone weight is tensor([-0.0075, 0.5364])\nStep 0 - Device 0 - num items in the local batch 52\nStep 0 - Device 1 - num items in the local batch 84\nTotal num items 136\nDevice 1 - w/ accumulation, the final model weight is tensor([0.2117, 0.3172])\nDevice 0 - w/ accumulation, the final model weight is tensor([0.2117, 0.3172])\nw/o accumulation, the final model weight is tensor([0.2117, 0.3172])\n```\n\n### To go further:\n\nPlease find a complete example script on a real world training run in the examples folder at the path [`accelerate/examples/by_feature/gradient_accumulation_for_autoregressive_models.py`](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/gradient_accumulation_for_autoregressive_models.py).\n\nRunning it on several training configurations with constant global batch size equal to 32 gives the following graph:\n\n
\n\n
\n\nNote that the training losses are exactly the same up to training step 20. The small deviation after this training step occurs at the very end of the first epoch, because, by [default](https://huggingface.co/docs/accelerate/en/package_reference/torch_wrappers#accelerate.data_loader.prepare_data_loader.even_batches), the dataloader duplicates the samples at the beginning of the dataset when the total batch size doesn't exactly divide the dataset.\n" }, { "path": "docs/source/usage_guides/intel_cpu.md", "content": "\n\n# Training on Intel CPU\n\n## How It Works For Training optimization in CPU\n\nAccelerate has full support for Intel CPU, all you need to do is enabling it through the config.\n\n**Scenario 1**: Acceleration of No distributed CPU training\n\nRun accelerate config on your machine:\n\n```bash\n$ accelerate config\n-----------------------------------------------------------------------------------------------------------------------------------------------------------\nIn which compute environment are you running?\nThis machine\n-----------------------------------------------------------------------------------------------------------------------------------------------------------\nWhich type of machine are you using?\nNo distributed training\nDo you want to run your training on CPU only (even if a GPU / Apple Silicon device is available)? [yes/NO]:yes\nDo you wish to optimize your script with torch dynamo?[yes/NO]:NO\nDo you want to use DeepSpeed? [yes/NO]: NO\n-----------------------------------------------------------------------------------------------------------------------------------------------------------\nDo you wish to use FP16 or BF16 (mixed precision)?\nbf16\n```\nThis will generate a config file that will be used automatically to properly set the\ndefault options when doing\n\n```bash\naccelerate launch my_script.py --args_to_my_script\n```\n\nFor instance, here is how you would run the NLP example `examples/nlp_example.py` (from the root of the repo) with `default_config.yaml` which is generated by `accelerate config`\n\n```bash\ncompute_environment: LOCAL_MACHINE\ndistributed_type: 'NO'\ndowncast_bf16: 'no'\nmachine_rank: 0\nmain_training_function: main\nmixed_precision: bf16\nnum_machines: 1\nnum_processes: 1\nrdzv_backend: static\nsame_network: true\ntpu_env: []\ntpu_use_cluster: false\ntpu_use_sudo: false\nuse_cpu: true\n```\n```bash\naccelerate launch examples/nlp_example.py\n```\n\n> [!CAUTION]\n> `accelerator.prepare` can currently only handle simultaneously preparing multiple models (and no optimizer) OR a single model-optimizer pair for training. Other attempts (e.g., two model-optimizer pairs) will raise a verbose error. To work around this limitation, consider separately using `accelerator.prepare` for each model-optimizer pair.\n\n**Scenario 2**: Acceleration of distributed CPU training\nwe 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\n\nRun accelerate config on your machine(node0):\n\n```bash\n$ accelerate config\n-----------------------------------------------------------------------------------------------------------------------------------------------------------\nIn which compute environment are you running?\nThis machine\n-----------------------------------------------------------------------------------------------------------------------------------------------------------\nWhich type of machine are you using?\nmulti-CPU\nHow many different machines will you use (use more than 1 for multi-node training)? [1]: 4\n-----------------------------------------------------------------------------------------------------------------------------------------------------------\nWhat is the rank of this machine?\n0\nWhat is the IP address of the machine that will host the main process? 36.112.23.24\nWhat is the port you will use to communicate with the main process? 29500\nAre 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\nDo you want accelerate to launch mpirun? [yes/NO]: yes\nPlease enter the path to the hostfile to use with mpirun [~/hostfile]: ~/hostfile\nEnter the number of oneCCL worker threads [1]: 1\nDo you wish to optimize your script with torch dynamo?[yes/NO]:NO\nHow many processes should be used for distributed training? [1]:16\n-----------------------------------------------------------------------------------------------------------------------------------------------------------\nDo you wish to use FP16 or BF16 (mixed precision)?\nbf16\n```\nFor instance, here is how you would run the NLP example `examples/nlp_example.py` (from the root of the repo) for distributed CPU training.\n\n`default_config.yaml` which is generated by `accelerate config`\n```bash\ncompute_environment: LOCAL_MACHINE\ndistributed_type: MULTI_CPU\ndowncast_bf16: 'no'\nmachine_rank: 0\nmain_process_ip: 36.112.23.24\nmain_process_port: 29500\nmain_training_function: main\nmixed_precision: bf16\nmpirun_config:\n mpirun_hostfile: /home/user/hostfile\nnum_machines: 4\nnum_processes: 16\nrdzv_backend: static\nsame_network: true\ntpu_env: []\ntpu_use_cluster: false\ntpu_use_sudo: false\nuse_cpu: true\n```\n\nSet following env and using intel MPI to launch the training\n\nIn `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.\n\nIf you selected to let Accelerate launch `mpirun`, ensure that the location of your hostfile matches the path in the config.\n\n```bash\n$ cat hostfile\nxxx.xxx.xxx.xxx #node0 ip\nxxx.xxx.xxx.xxx #node1 ip\nxxx.xxx.xxx.xxx #node2 ip\nxxx.xxx.xxx.xxx #node3 ip\n```\n\n```bash\naccelerate launch examples/nlp_example.py\n```\n\nYou can also directly launch distributed training with `mpirun` command, you need to run the following command in node0 and **16DDP** will be enabled in node0,node1,node2,node3 with BF16 mixed precision. When using this method, the python script, python environment, and accelerate config file need to be available on all of the machines used for multi-CPU training.\n\n```bash\nexport MASTER_ADDR=xxx.xxx.xxx.xxx #node0 ip\nmpirun -f hostfile -n 16 -ppn 4 accelerate launch examples/nlp_example.py\n```\n" }, { "path": "docs/source/usage_guides/local_sgd.md", "content": "\n\n# Using Local SGD with Accelerate\n\nLocal 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.\nUnlike 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.\n\nIn 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.\n\nThis example will use a very simplistic PyTorch training loop that performs gradient accumulation every two batches:\n\n```python\ndevice = \"cuda\"\nmodel.to(device)\n\ngradient_accumulation_steps = 2\n\nfor index, batch in enumerate(training_dataloader):\n inputs, targets = batch\n inputs = inputs.to(device)\n targets = targets.to(device)\n outputs = model(inputs)\n loss = loss_function(outputs, targets)\n loss = loss / gradient_accumulation_steps\n loss.backward()\n if (index + 1) % gradient_accumulation_steps == 0:\n optimizer.step()\n scheduler.step()\n optimizer.zero_grad()\n```\n\n## Converting it to Accelerate\n\nFirst the code shown earlier will be converted to use Accelerate with neither a LocalSGD or a gradient accumulation helper:\n\n```diff\n+ from accelerate import Accelerator\n+ accelerator = Accelerator()\n\n+ model, optimizer, training_dataloader, scheduler = accelerator.prepare(\n+ model, optimizer, training_dataloader, scheduler\n+ )\n\n for index, batch in enumerate(training_dataloader):\n inputs, targets = batch\n- inputs = inputs.to(device)\n- targets = targets.to(device)\n outputs = model(inputs)\n loss = loss_function(outputs, targets)\n loss = loss / gradient_accumulation_steps\n+ accelerator.backward(loss)\n if (index+1) % gradient_accumulation_steps == 0:\n optimizer.step()\n scheduler.step()\n```\n\n## Letting Accelerate handle model synchronization \n\nAll 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\nachieved by adding one `with LocalSGD` statement and one call `local_sgd.step()` after every optimizer step:\n\n```diff\n+local_sgd_steps=8\n\n+with LocalSGD(accelerator=accelerator, model=model, local_sgd_steps=8, enabled=True) as local_sgd:\n for batch in training_dataloader:\n with accelerator.accumulate(model):\n inputs, targets = batch\n outputs = model(inputs)\n loss = loss_function(outputs, targets)\n accelerator.backward(loss)\n optimizer.step()\n scheduler.step()\n optimizer.zero_grad()\n+ local_sgd.step()\n```\n\nUnder the hood, the Local SGD code **disables** automatic gradient synchronization (but accumulation still works as expected!). Instead it averages model parameters every `local_sgd_steps` steps (as well as at the end of the training loop).\n\n## Limitations\n\nThe current implementation works only with basic multi-GPU (or multi-CPU) training without, e.g., [DeepSpeed.](https://github.com/deepspeedai/DeepSpeed).\n\n## References\n\n Although we are not aware of the true origins of this simple approach, the idea of local SGD is quite old and goes\n back to at least:\n\n Zhang, J., De Sa, C., Mitliagkas, I., & RĆ©, C. (2016). [Parallel SGD: When does averaging help?. arXiv preprint\n arXiv:1606.07365.](https://huggingface.co/papers/1606.07365)\n\n We credit the term Local SGD to the following paper (but there might be earlier references we are not aware of).\n\n Stich, Sebastian Urban. [\"Local SGD Converges Fast and Communicates Little.\" ICLR 2019-International Conference on\n Learning Representations. No. CONF. 2019.](https://huggingface.co/papers/1805.09767)\n" }, { "path": "docs/source/usage_guides/low_precision_training.md", "content": "\n\n# Low Precision Training Methods\n\nAccelerate provides integrations to train on lower precision methods using specified supported hardware through the `TransformersEngine`, `MS-AMP`, and `torchao` packages. This documentation will help guide you through what hardware is supported, how to configure your [`Accelerator`] to leverage the low precision methods, and what you can expect when training. \n\n## What training on FP8 means\n\nTo explore more of the nitty-gritty in training in FP8 with PyTorch and Accelerate, check out the [concept_guide](../concept_guides/low_precision_training) on why this can be difficult. But essentially rather than training in BF16, some (or all) aspects of training a model can be performed using 8 bits instead of 16. The challenge is doing so without degrading final performance. \n\nThis is only enabled on specific NVIDIA hardware, namely:\n\n* Anything after the 3000 series consumer graphics cards (such as the 4090)\n* Hopper-based GPU architectures (such as the `H100` and `H200`)\n\nWhat this will result in is some reduction in the memory used (as we've cut the needed memory in half for some parts of training) and an increase in throughput *should* be seen as well for larger models that can replace certain layers with FP8-enabled ones.\n\n## Configuring the Accelerator\n\nCurrently two actively maintained backends for FP8 are supported (`TransformersEngine` and `torchao`), each with different capabilities and configurations. A legacy `MS-AMP` backend also exists but is no longer recommended (see [below](#configuring-ms-amp) for details).\n\nTo use either, the same core API is used. Just pass `mixed_precision=\"fp8\"` to either the [`Accelerator`], during `accelerate config` when prompted about mixed precision, or as part of your `config.yaml` file in the `mixed_precision` key:\n\n```{python}\nfrom accelerate import Accelerator\naccelerator = Accelerator(mixed_precision=\"fp8\")\n```\n\nTo specify a backend (and customize other parts of the FP8 mixed precision setup), you can utilize one of the `RecipeKwargs` dataclasses such as [`utils.AORecipeKwargs`], [`utils.TERecipeKwargs`], or [`utils.MSAMPRecipeKwargs`]; you can also clarify it in your config `yaml`/during `accelerate launch`. We recommend using `TransformersEngine` or `torchao` for new projects:\n\n```{python}\nfrom accelerate import Accelerator\nfrom accelerate.utils import TERecipeKwargs, AORecipeKwargs\n# Use TransformersEngine\nkwargs = [TERecipeKwargs()]\n# Or to use torchao\n# kwargs = [AORecipeKwargs()]\naccelerator = Accelerator(mixed_precision=\"fp8\", kwarg_handlers=kwargs)\n```\n\n```{yaml}\nmixed_precision: fp8\nfp8_config:\n amax_compute_algo: max\n amax_history_len: 1024\n backend: TE\n fp8_format: HYBRID\n interval: 1\n margin: 0\n override_linear_precision: (false, false, false)\n use_autocast_during_eval: false\n```\n\n## Configuring MS-AMP\n\n\n\n**āš ļø Deprecated / Unmaintained:** MS-AMP is no longer actively maintained by Microsoft. The [MS-AMP repository](https://github.com/Azure/MS-AMP) has not received updates since 2023 and has known compatibility issues:\n\n- Requires CUDA 11.x (does not support CUDA 12.x+)\n- Requires older NCCL versions incompatible with recent PyTorch releases\n- Does not support recent PyTorch versions (2.2+)\n\n**We strongly recommend using [`TransformersEngine`](#configuring-transformersengine) or [`torchao`](#configuring-torchao) instead for all new and existing FP8 training workflows.** Both are actively maintained and support modern CUDA/PyTorch versions. Native PyTorch FP8 support via `torchao` is particularly promising as a vendor-neutral solution.\n\nThe MS-AMP backend is retained in Accelerate for legacy compatibility but may be removed in a future release.\n\n\n\n`MS-AMP` has a single configuration argument: the optimization level. \n\nCurrently two levels of optimization are supported in the Accelerate integration, `\"O1\"` and `\"O2\"` (using the letter 'o', not zero). \n\n* `\"O1\"` will cast the weight gradients and `all_reduce` communications to happen in 8-bit, while the rest are done in 16 bit. This reduces the general GPU memory usage and speeds up communication bandwidths.\n* `\"O2\"` will also cast first-order optimizer states into 8 bit, while the second order states are in FP16. (Currently just the `Adam` optimizer is supported). This tries its best to minimize final accuracy degradation and will save the highest potential memory.\n\nTo specify an optimization level, pass it to the `FP8KwargsHandler` by setting the `optimization_level` argument:\n\n```{python}\nfrom accelerate import Accelerator\nfrom accelerate.utils import FP8RecipeKwargs\nkwargs = [FP8RecipeKwargs(backend=\"msamp\", optimization_level=\"O2\")]\naccelerator = Accelerator(mixed_precision=\"fp8\", kwarg_handlers=kwargs)\n```\n\nOr during `accelerate launch` via `--fp8_backend=msamp --fp8_opt_level=O2`\n\nSimilarly this can be set in your `config.yaml`:\n\n```{yaml}\nmixed_precision: fp8\nfp8_config:\n backend: MSAMP\n opt_level: O2\n```\n\n## Configuring TransformersEngine\n\nTransformersEngine has many options for customizing how and what FP8 calculations are performed. A full list of supported arguments and what they mean are available in [NVIDIA's documentation](https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/api/common.html), however they are restated as part of [`FP8KwargsHandler`]'s docstring for your convenience. \n\nAccelerate tries to set sensible defaults, but exploring and tweaking the various parameters yourself can lead to better performance potentially.\n\nTo use it, specify `backend=\"te\"` and modify any of the arguments you want as part of your kwarg handler:\n\n```{python}\nfrom accelerate import Accelerator\nfrom accelerate.utils import FP8RecipeKwargs\nkwargs = [FP8RecipeKwargs(backend=\"te\", ...)]\naccelerator = Accelerator(mixed_precision=\"fp8\", kwarg_handlers=kwargs)\n```\n\nOr during `accelerate launch` via `--fp8_backend=te ...`. Use `accelerate launch --fp8_backend=te -h` to see relevent arguments.\n\nSimilarly this can be set in your `config.yaml`:\n\n```{yaml}\nmixed_precision: fp8\nfp8_config:\n amax_compute_algo: max\n amax_history_len: 1024\n backend: TE\n fp8_format: HYBRID\n interval: 1\n margin: 0\n override_linear_precision: (false, false, false)\n use_autocast_during_eval: false\n```\n\n## Configuring `torchao`\n\n`torchao` is a [PyTorch-driven](https://github.com/pytorch/ao/tree/main/torchao/float8) hackable FP8 backend, aiming to be more approchable than the prior two engines. One of the core differences with `ao` compared to the prior two is that for numerical stability, it's found to be generally better off keeping the first *and* last layers in the model at the regular precision (be it FP32 or BF16), and then the other layers quantized down to FP8. As a result, a config for `ao` looks a bit differently:\n\n> Note: this API is experimental and is subject to change\n\n```{python}\nfrom accelerate import Accelerator\nfrom accelerate.utils import AORecipeKwargs, TorchDynamoPlugin, FullyShardedDataParallelPlugin\nfrom torchao.float8 import Float8LinearConfig\n\nfsdp2_plugin = FullyShardedDataParallelPlugin(\n fsdp_version=2,\n cpu_ram_efficient_loading=False, # CPU RAM efficient loading CANNOT work with fp8 torchao\n fsdp_auto_wrap_policy=\"TRANSFORMER_BASED_WRAP\",\n)\ndynamo_plugin = TorchDynamoPlugin(\n backend=\"inductor\",\n use_regional_compilation=True,\n)\nfp8_config = Float8LinearConfig(\n enable_fsdp_float8_all_gather=True, # Use FP8 all_gather in FSDP2\n pad_inner_dim=True,\n)\nkwargs = [AORecipeKwargs(\n config=fp8_config\n)]\naccelerator = Accelerator(\n mixed_precision=\"fp8\",\n fsdp_plugin=fsdp2_plugin,\n dynamo_plugin=dynamo_plugin,\n kwarg_handlers=kwargs,\n)\n```\n\nOr during `accelerate launch` via `--fp8_backend=ao ...`. Use `accelerate launch --fp8_backend=ao -h` to see relevent arguments.\n\nSimilarly, this can be set in `config.yaml`:\n\n```{yaml}\nmixed_precision: fp8\nfsdp_config:\n fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP\n fsdp_cpu_ram_efficient_loading: false\n fsdp_version: 2\nfp8_config:\n backend: AO\n pad_inner_dim: true\n enable_fsdp_float8_all_gather: true\ndynamo_config:\n dynamo_backend: INDUCTOR\n dynamo_use_regional_compilation: true\n```\n\nTo learn more about the specific parameters to be used, please see the official `torchao` repo.\n\n\n## Example Zoo\n\nWe have examples showcasing training with FP8 both with accelerate and its underlying implementation available in the accelerate repo.\nCurrently we support scripts showcasing:\n\n* Single GPU\n* Distributed Data Parallelism (Multi-GPU)\n* Fully Sharded Data Parallelism\n* DeepSpeed ZeRO 1 through 3\n\nFind out more [here](https://github.com/huggingface/accelerate/tree/main/benchmarks/fp8)\n\n## Further Reading\n\nTo learn more about training in FP8 please check out the following resources:\n\n* [Our concept guide](../concept_guides/low_precision_training) detailing into more about TransformersEngine, torchao, and MS-AMP\n* [The `transformers-engine` documentation](https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/api/common.html)\n* [The `torchao` documentation](https://github.com/pytorch/ao/tree/main/torchao/float8)\n* [The `MS-AMP` documentation](https://azure.github.io/MS-AMP/docs/) (āš ļø no longer maintained)\n" }, { "path": "docs/source/usage_guides/megatron_lm.md", "content": "\n\n\n# Megatron-LM\n\n[Megatron-LM](https://github.com/NVIDIA/Megatron-LM) enables training large transformer language models at scale.\nIt provides efficient tensor, pipeline and sequence based model parallelism for pre-training transformer based\nLanguage Models such as [GPT](https://huggingface.co/papers/2005.14165) (Decoder Only), [BERT](https://huggingface.co/papers/1810.04805) (Encoder Only) and [T5](https://huggingface.co/papers/1910.10683) (Encoder-Decoder).\nFor detailed information and how things work behind the scene please refer to the github [repo](https://github.com/NVIDIA/Megatron-LM).\n\n## What is integrated?\n\nAccelerate integrates following feature of Megatron-LM to enable large scale pre-training/finetuning\nof BERT (Encoder), GPT (Decoder) or T5 models (Encoder and Decoder):\n\na. **Tensor Parallelism (TP)**: Reduces memory footprint without much additional communication on intra-node ranks.\nEach tensor is split into multiple chunks with each shard residing on separate GPU. At each step, the same mini-batch of data is processed\nindependently and in parallel by each shard followed by syncing across all GPUs (`all-reduce` operation). \nIn a simple transformer layer, this leads to 2 `all-reduces` in the forward path and 2 in the backward path.\nFor more details, please refer to the research paper [Megatron-LM: Training Multi-Billion Parameter Language Models Using\nModel Parallelism](https://huggingface.co/papers/1909.08053) and \nthis section of blogpost [The Technology Behind BLOOM Training](https://huggingface.co/blog/bloom-megatron-deepspeed#tensor-parallelism).\n\n\nb. **Pipeline Parallelism (PP)**: Reduces memory footprint and enables large scale training via inter-node parallelization. \nReduces the bubble of naive PP via PipeDream-Flush schedule/1F1B schedule and Interleaved 1F1B schedule. \nLayers are distributed uniformly across PP stages. For example, if a model has `24` layers and we have `4` GPUs for\npipeline parallelism, each GPU will have `6` layers (24/4). For more details on schedules to reduce the idle time of PP,\nplease refer to the research paper [Efficient Large-Scale Language Model Training on GPU Clusters\nUsing Megatron-LM](https://huggingface.co/papers/2104.04473) and \nthis section of blogpost [The Technology Behind BLOOM Training](https://huggingface.co/blog/bloom-megatron-deepspeed#pipeline-parallelism).\n\nc. **Sequence Parallelism (SP)**: Reduces memory footprint without any additional communication. Only applicable when using TP.\nIt reduces activation memory required as it prevents the same copies to be on the tensor parallel ranks \npost `all-reduce` by replacing them with `reduce-scatter` and `no-op` operation would be replaced by `all-gather`. \nAs `all-reduce = reduce-scatter + all-gather`, this saves a ton of activation memory at no added communication cost. \nTo put it simply, it shards the outputs of each transformer layer along sequence dimension, e.g., \nif the sequence length is `1024` and the TP size is `4`, each GPU will have `256` tokens (1024/4) for each sample. \nThis increases the batch size that can be supported for training. For more details, please refer to the research paper\n[Reducing Activation Recomputation in Large Transformer Models](https://huggingface.co/papers/2205.05198). \n\nd. **Data Parallelism (DP)** via Distributed Optimizer: Reduces the memory footprint by sharding optimizer states and gradients across DP ranks\n(versus the traditional method of replicating the optimizer state across data parallel ranks). \nFor example, when using Adam optimizer with mixed-precision training, each parameter accounts for 12 bytes of memory.\nThis gets distributed equally across the GPUs, i.e., each parameter would account for 3 bytes (12/4) if we have 4 GPUs.\nFor more details, please refer to the research paper [ZeRO: Memory Optimizations Toward Training Trillion\nParameter Models](https://huggingface.co/papers/1910.02054) and following section of blog \n[The Technology Behind BLOOM Training](https://huggingface.co/blog/bloom-megatron-deepspeed#zero-data-parallelism).\n\ne. **Expert Parallelism (EP)** Expert parallelism in Megatron-LM is used for Mixture-of-Experts (MoE) layers, where many ā€œexpertsā€ (small feed-forward networks) exist but only a few are activated for each token. Instead of putting all experts on every GPU, Megatron distributes different experts across different GPUsā€”this is expert parallelism. During training, tokens are routed to the GPUs that host their selected experts, computed there, and then sent back, reducing memory cost. It often combines with tensor/pipeline parallelism for large-scale models.\nf. **Full Activation Recomputation**: Reduces the memory footprint of activations significantly via smart activation checkpointing.\nIt doesn't store activations occupying large memory while being fast to recompute thereby achieving great tradeoff between memory and recomputation.\nFor example, for GPT-3, this leads to 70% reduction in required memory for activations at the expense of\nonly 2.7% FLOPs overhead for recomputation of activations. For more details, please refer to the research paper \n[Reducing Activation Recomputation in Large Transformer Models](https://huggingface.co/papers/2205.05198).\n\ng. **Fused Kernels**: Fused Softmax, Mixed Precision Fused Layer Norm and Fused gradient accumulation to weight gradient computation of linear layer.\nPyTorch JIT compiled Fused GeLU and Fused Bias+Dropout+Residual addition.\n\nh. **Support for Indexed datasets**: Efficient binary format of datasets for large scale training. Support for the `mmap`, `cached` index file and the `lazy` loader format.\n\ni. **Checkpoint reshaping and interoperability**: Utility for reshaping Megatron-LM checkpoints of variable \ntensor and pipeline parallel sizes to the beloved Transformers sharded checkpoints as it has great support with plethora of tools\nsuch as Accelerate Big Model Inference, Megatron-DeepSpeed Inference etc. \nSupport is also available for converting Transformers sharded checkpoints to Megatron-LM checkpoint of variable tensor and pipeline parallel sizes\nfor large scale training. \n\n\n## Pre-Requisites \n\nYou will need to install the latest pytorch, cuda, nccl, and NVIDIA [APEX](https://github.com/NVIDIA/apex#quick-start) releases and the nltk library.\nSee [documentation](https://github.com/NVIDIA/Megatron-LM#setup) for more details. \nAnother way to setup the environment is to pull an NVIDIA PyTorch Container that comes with all the required installations from NGC.\n\nBelow is a step-by-step method to set up the conda environment:\n\n1. Create a virtual environment\n```\nconda create --name ml\n```\n\n2. Assuming that the machine has CUDA 11.3 installed, installing the corresponding PyTorch GPU Version\n```\nconda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch\n```\n\n3. Install Nvidia APEX\n```\ngit clone https://github.com/NVIDIA/apex\ncd apex\npip install -v --disable-pip-version-check --no-cache-dir --global-option=\"--cpp_ext\" --global-option=\"--cuda_ext\" ./\ncd ..\n```\n\n4. Installing Megatron-LM\n\n```\ngit clone https://github.com/NVIDIA/Megatron-LM.git\ncd Megatron-LM\ngit checkout 9a1c0d05c992c8a241da384ab27dce2021bb56dd\nyou need to manually move gpt_builders.py to megatron/training and update\ninclude = [\n \"megatron.core\", \n \"megatron.core.*\",\n \"megatron.training\",\n \"megatron.training.*\",\n \"megatron.legacy\",\n \"megatron.legacy.*\",\n]\nin pyproject.toml file to unblock yourself from using Megatron\npip install --no-use-pep517 -e .\n```\n\n## Prepare Megaton-LM checkpoint\nIf you want to fine-tune a model, make sure you have a torch dist format checkpoint ready. If you only have access to the huggingface model, please consider converting it to a torch dist format checkpoint acceptable to Megatron. One examle can be using slime's script, take GLM models as an example:\n```\nsource /your/path/to/slime/scripts/models/glm4.5-355B-A32B.sh\nsrun torchrun --nproc-per-node 8 \\\n /your/path/to/slime/tools/convert_hf_to_torch_dist.py \\\n ${MODEL_ARGS[@]} \\\n --hf-checkpoint /your/path/to/huggingface/models/GLM4.5-355B-A32B \\\n --save /your/path/to/megatron/models/GLM4.5-355B-A32B_torch_dist\n\n```\nAfter the conversion, make sure: 1. under `/your/path/to/megatron/models/GLM4.5-355B-A32B_torch_dist`: change the `latest_checkpointed_iteration.txt`'s content from `release` to `0` and rename the directory `release` to `iter_0000000`; 2: in the config, make sure `megatron_lm_no_load_optim` to be true so that no optimizer states are needed.\n\n## Accelerate Megatron-LM Plugin\n\nImportant features are directly supported via the `accelerate config` command. \nAn example of the corresponding questions for using Megatron-LM features is shown below:\n\n```bash\n:~$ accelerate config --config_file \"megatron_gpt_config.yaml\"\nIn which compute environment are you running? ([0] This machine, [1] AWS (Amazon SageMaker)): 0\nWhich type of machine are you using? ([0] No distributed training, [1] multi-CPU, [2] multi-GPU, [3] TPU): 2\nHow many different machines will you use (use more than 1 for multi-node training)? [1]: \nDo you want to use DeepSpeed? [yes/NO]: \nDo you want to use FullyShardedDataParallel? [yes/NO]: \nDo you want to use Megatron-LM ? [yes/NO]: yes\nWhat is the Tensor Parallelism degree/size? [1]:2\nDo you want to enable Sequence Parallelism? [YES/no]: \nWhat is the Pipeline Parallelism degree/size? [1]:2\nWhat is the number of micro-batches? [1]:2\nDo you want to enable selective activation recomputation? [YES/no]: \nDo you want to use distributed optimizer which shards optimizer state and gradients across data parallel ranks? [YES/no]: \nWhat is the gradient clipping value based on global L2 Norm (0 to disable)? [1.0]: \nHow many GPU(s) should be used for distributed training? [1]:4\nDo you wish to use FP16 or BF16 (mixed precision)? [NO/fp16/bf16]: bf16\n```\n\nThe resulting config is shown below:\n\n```\n~$ cat megatron_gpt_config.yaml \ncompute_environment: LOCAL_MACHINE\ndeepspeed_config: {}\ndistributed_type: MEGATRON_LM\ndowncast_bf16: 'no'\nfsdp_config: {}\nmachine_rank: 0\nmain_process_ip: null\nmain_process_port: null\nmain_training_function: main\nmegatron_lm_config:\n megatron_lm_gradient_clipping: 1.0\n megatron_lm_num_micro_batches: 2\n megatron_lm_pp_degree: 2\n megatron_lm_recompute_activations: true\n megatron_lm_sequence_parallelism: true\n megatron_lm_tp_degree: 2\n megatron_lm_use_distributed_optimizer: true\nmixed_precision: bf16\nnum_machines: 1\nnum_processes: 4\nrdzv_backend: static\nsame_network: true\nuse_cpu: false\n```\n\nWe will take the example of GPT pre-training. The minimal changes required to the official `run_clm_no_trainer.py` \nto use Megatron-LM are as follows:\n\n1. As Megatron-LM uses its own implementation of Optimizer, the corresponding scheduler compatible with it needs to be used.\nAs such, support for only the Megatron-LM's scheduler is present. User will need to create `accelerate.utils.MegatronLMDummyScheduler`.\nExample is given below:\n\n```python\nfrom accelerate.utils import MegatronLMDummyScheduler\n\nif accelerator.distributed_type == DistributedType.MEGATRON_LM:\n lr_scheduler = MegatronLMDummyScheduler(\n optimizer=optimizer,\n total_num_steps=args.max_train_steps,\n warmup_num_steps=args.num_warmup_steps,\n )\nelse:\n lr_scheduler = get_scheduler(\n name=args.lr_scheduler_type,\n optimizer=optimizer,\n num_warmup_steps=args.num_warmup_steps * args.gradient_accumulation_steps,\n num_training_steps=args.max_train_steps * args.gradient_accumulation_steps,\n )\n```\n\n2. Getting the details of the total batch size now needs to be cognization of tensor and pipeline parallel sizes.\nExample of getting the effective total batch size is shown below:\n\n```python\nif accelerator.distributed_type == DistributedType.MEGATRON_LM:\n total_batch_size = accelerator.state.megatron_lm_plugin.global_batch_size\nelse:\n total_batch_size = args.per_device_train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps\n```\n\n3. When using Megatron-LM, the losses are already averaged across the data parallel group\n\n```python\nif accelerator.distributed_type == DistributedType.MEGATRON_LM:\n losses.append(loss)\nelse:\n losses.append(accelerator.gather_for_metrics(loss.repeat(args.per_device_eval_batch_size)))\n\nif accelerator.distributed_type == DistributedType.MEGATRON_LM:\n losses = torch.tensor(losses)\nelse:\n losses = torch.cat(losses)\n```\n\n4. For Megatron-LM, we need to save the model using `accelerator.save_state`\n\n```python\nif accelerator.distributed_type == DistributedType.MEGATRON_LM:\n accelerator.save_state(args.output_dir)\nelse:\n unwrapped_model = accelerator.unwrap_model(model)\n unwrapped_model.save_pretrained(\n args.output_dir, is_main_process=accelerator.is_main_process, save_function=accelerator.save\n )\n```\n\nThat's it! We are good to go šŸš€. Please find the example script in the examples folder at the path `accelerate/examples/by_feature/megatron_lm_gpt_pretraining.py`.\nLet's run it for `gpt-large` model architecture using 4 A100-80GB GPUs.\n\n```bash\naccelerate launch --config_file megatron_gpt_config.yaml \\\nexamples/by_feature/megatron_lm_gpt_pretraining.py \\\n--config_name \"gpt2-large\" \\\n--tokenizer_name \"gpt2-large\" \\\n--dataset_name wikitext \\\n--dataset_config_name wikitext-2-raw-v1 \\\n--block_size 1024 \\\n--learning_rate 5e-5 \\\n--per_device_train_batch_size 24 \\\n--per_device_eval_batch_size 24 \\\n--num_train_epochs 5 \\\n--with_tracking \\\n--report_to \"wandb\" \\\n--output_dir \"awesome_model\"\n```\n\nBelow are some important excerpts from the output logs:\n\n```bash\nLoading extension module fused_dense_cuda...\n>>> done with compiling and loading fused kernels. Compilation time: 3.569 seconds\n > padded vocab (size: 50257) with 175 dummy tokens (new size: 50432)\nBuilding gpt model in the pre-training mode.\nThe Megatron LM model weights are initialized at random in `accelerator.prepare`. Please use `accelerator.load_checkpoint` to load a pre-trained checkpoint matching the distributed setup.\nPreparing dataloader\nPreparing dataloader\nPreparing model\n > number of parameters on (tensor, pipeline) model parallel rank (1, 0): 210753280\n > number of parameters on (tensor, pipeline) model parallel rank (1, 1): 209445120\n > number of parameters on (tensor, pipeline) model parallel rank (0, 0): 210753280\n > number of parameters on (tensor, pipeline) model parallel rank (0, 1): 209445120\nPreparing optimizer\nPreparing scheduler\n> learning rate decay style: linear\n10/10/2022 22:57:22 - INFO - __main__ - ***** Running training *****\n10/10/2022 22:57:22 - INFO - __main__ - Num examples = 2318\n10/10/2022 22:57:22 - INFO - __main__ - Num Epochs = 5\n10/10/2022 22:57:22 - INFO - __main__ - Instantaneous batch size per device = 24\n10/10/2022 22:57:22 - INFO - __main__ - Total train batch size (w. parallel, distributed & accumulation) = 48\n10/10/2022 22:57:22 - INFO - __main__ - Gradient Accumulation steps = 1\n10/10/2022 22:57:22 - INFO - __main__ - Total optimization steps = 245\n 20%|ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā– | 49/245 [01:04<04:09, 1.27s/it]\n 10/10/2022 22:58:29 - INFO - __main__ - epoch 0: perplexity: 1222.1594275215962 eval_loss: 7.10837459564209\n 40%|ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–Š | 98/245 [02:10<03:07, 1.28s/it]\n 10/10/2022 22:59:35 - INFO - __main__ - epoch 1: perplexity: 894.5236583794557 eval_loss: 6.796291351318359\n 60%|ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–Œ | 147/245 [03:16<02:05, 1.28s/it]\n 10/10/2022 23:00:40 - INFO - __main__ - epoch 2: perplexity: 702.8458788508042 eval_loss: 6.555137634277344\n 80%|ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–Š | 196/245 [04:22<01:02, 1.28s/it]\n 10/10/2022 23:01:46 - INFO - __main__ - epoch 3: perplexity: 600.3220028695281 eval_loss: 6.39746618270874\n100%|ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ| 245/245 [05:27<00:00, 1.28s/it]\n```\n\nThere are a large number of other options/features that one can set using `accelerate.utils.MegatronLMPlugin`.\n\n## Advanced features to leverage writing custom train step and Megatron-LM Indexed Datasets\n\nFor leveraging more features, please go through below details.\n\n1. Below is an example of changes required to customize the Train Step while using Megatron-LM. \nYou will implement the `accelerate.utils.AbstractTrainStep` or inherit from their corresponding children \n`accelerate.utils.GPTTrainStep`, `accelerate.utils.BertTrainStep` or `accelerate.utils.T5TrainStep`.\n\n```python\nfrom accelerate.utils import MegatronLMDummyScheduler, GPTTrainStep, avg_losses_across_data_parallel_group\n\n\n# Custom loss function for the Megatron model\nclass GPTTrainStepWithCustomLoss(GPTTrainStep):\n def __init__(self, megatron_args, **kwargs):\n super().__init__(megatron_args)\n self.kwargs = kwargs\n\n def get_loss_func(self):\n def loss_func(inputs, loss_mask, output_tensor):\n batch_size, seq_length = output_tensor.shape\n losses = output_tensor.float()\n loss_mask = loss_mask.view(-1).float()\n loss = losses.view(-1) * loss_mask\n\n # Resize and average loss per sample\n loss_per_sample = loss.view(batch_size, seq_length).sum(axis=1)\n loss_mask_per_sample = loss_mask.view(batch_size, seq_length).sum(axis=1)\n loss_per_sample = loss_per_sample / loss_mask_per_sample\n\n # Calculate and scale weighting\n weights = torch.stack([(inputs == kt).float() for kt in self.kwargs[\"keytoken_ids\"]]).sum(axis=[0, 2])\n weights = 1.0 + self.kwargs[\"alpha\"] * weights\n # Calculate weighted average\n weighted_loss = (loss_per_sample * weights).mean()\n\n # Reduce loss across data parallel groups\n averaged_loss = avg_losses_across_data_parallel_group([weighted_loss])\n\n return weighted_loss, {\"lm loss\": averaged_loss[0]}\n\n return loss_func\n\n def get_forward_step_func(self):\n def forward_step(data_iterator, model):\n \"\"\"Forward step.\"\"\"\n # Get the batch.\n tokens, labels, loss_mask, attention_mask, position_ids = self.get_batch(data_iterator)\n output_tensor = model(tokens, position_ids, attention_mask, labels=labels)\n\n return output_tensor, partial(self.loss_func, tokens, loss_mask)\n\n return forward_step\n\n\ndef main():\n # Custom loss function for the Megatron model\n keytoken_ids = []\n keywords = [\"plt\", \"pd\", \"sk\", \"fit\", \"predict\", \" plt\", \" pd\", \" sk\", \" fit\", \" predict\"]\n for keyword in keywords:\n ids = tokenizer([keyword]).input_ids[0]\n if len(ids) == 1:\n keytoken_ids.append(ids[0])\n accelerator.print(f\"Keytoken ids: {keytoken_ids}\")\n accelerator.state.megatron_lm_plugin.custom_train_step_class = GPTTrainStepWithCustomLoss\n accelerator.state.megatron_lm_plugin.custom_train_step_kwargs = {\n \"keytoken_ids\": keytoken_ids,\n \"alpha\": 0.25,\n }\n```\n\n2. For using the Megatron-LM datasets, a few more changes are required. Dataloaders for these datasets\nare available only on rank 0 of each tensor parallel group. As such, there are rank where dataloader won't be\navailable and this requires tweaks to the training loop. Being able to do all this shows how\nflexible and extensible Accelerate is. The changes required are as follows.\n\na. For Megatron-LM indexed datasets, we need to use `MegatronLMDummyDataLoader` \nand pass the required dataset args to it such as `data_path`, `seq_length` etc. \nSee [here](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/arguments.py#L804) for the list of available args. \n \n```python\nfrom accelerate.utils import MegatronLMDummyDataLoader\n\nmegatron_dataloader_config = {\n \"data_path\": args.data_path,\n \"splits_string\": args.splits_string,\n \"seq_length\": args.block_size,\n \"micro_batch_size\": args.per_device_train_batch_size,\n}\nmegatron_dataloader = MegatronLMDummyDataLoader(**megatron_dataloader_config)\naccelerator.state.megatron_lm_plugin.megatron_dataset_flag = True\n```\n\nb. `megatron_dataloader` is repeated 3 times to get training, validation and test dataloaders\nas per the `args.splits_string` proportions\n \n```python\nmodel, optimizer, lr_scheduler, train_dataloader, eval_dataloader, _ = accelerator.prepare(\n model, optimizer, lr_scheduler, megatron_dataloader, megatron_dataloader, megatron_dataloader\n)\n```\n\nc. Changes to training and evaluation loops as dataloader is only available on tensor parallel ranks 0\nSo, we need to iterate only if the dataloader isn't `None` else provide empty dict\nAs such, we loop using `while` loop and break when `completed_steps` is equal to `args.max_train_steps`\nThis is similar to the Megatron-LM setup wherein user has to provide `max_train_steps` when using Megaton-LM indexed datasets.\nThis displays how flexible and extensible Accelerate is.\n\n```python\nwhile completed_steps < args.max_train_steps:\n model.train()\n batch = next(train_dataloader) if train_dataloader is not None else {}\n outputs = model(**batch)\n loss = outputs.loss\n ...\n\n if completed_steps % eval_interval == 0:\n eval_completed_steps = 0\n losses = []\n while eval_completed_steps < eval_iters:\n model.eval()\n with torch.no_grad():\n batch = next(eval_dataloader) if eval_dataloader is not None else {}\n outputs = model(**batch)\n```\n\n \n## Utility for Checkpoint reshaping and interoperability\n\n1. The scripts for these are present in Transformers library under respective models. \nCurrently, it is available for GPT model [checkpoint_reshaping_and_interoperability.py](https://github.com/huggingface/transformers/blob/main/src/transformers/models/megatron_gpt2/checkpoint_reshaping_and_interoperability.py)\n\n2. Below is an example of conversion of checkpoint from Megatron-LM to universal Transformers sharded checkpoint.\n```bash\npython checkpoint_reshaping_and_interoperability.py \\\n--convert_checkpoint_from_megatron_to_transformers \\\n--load_path \"gpt/iter_0005000\" \\\n--save_path \"gpt/trfs_checkpoint\" \\\n--max_shard_size \"200MB\" \\\n--tokenizer_name \"gpt2\" \\\n--print-checkpoint-structure\n```\n\n3. Conversion of checkpoint from transformers to megatron with `tp_size=2`, `pp_size=2` and `dp_size=2`.\n```bash\npython checkpoint_utils/megatgron_gpt2/checkpoint_reshaping_and_interoperability.py \\\n--load_path \"gpt/trfs_checkpoint\" \\\n--save_path \"gpt/megatron_lm_checkpoint\" \\\n--target_tensor_model_parallel_size 2 \\\n--target_pipeline_model_parallel_size 2 \\\n--target_data_parallel_size 2 \\\n--target_params_dtype \"bf16\" \\\n--make_vocab_size_divisible_by 128 \\\n--use_distributed_optimizer \\\n--print-checkpoint-structure\n```\n\n## Megatron-LM GPT models support returning logits and `megatron_generate` function for text generation\n\n1. Returning logits require setting `require_logits=True` in MegatronLMPlugin as shown below. \nThese would be available in the last stage of pipeline.\n```python\nmegatron_lm_plugin = MegatronLMPlugin(return_logits=True)\n```\n\n2. `megatron_generate` method for Megatron-LM GPT model: This will use Tensor and Pipeline Parallelism to complete \ngenerations for a batch of inputs when using greedy with/without top_k/top_p sampling and for individual prompt inputs when using beam search decoding. \nOnly a subset of features of transformers generate is supported. This will help in using large models via tensor and pipeline parallelism \nfor generation (already does key-value caching and uses fused kernels by default).\nThis requires data parallel size to be 1, sequence parallelism and activation checkpointing to be disabled.\nIt also requires specifying path to tokenizer's vocab file and merges file. \nBelow example shows how to configure and use `megatron_generate` method for Megatron-LM GPT model.\n```python\n# specifying tokenizer's vocab and merges file\nvocab_file = os.path.join(args.resume_from_checkpoint, \"vocab.json\")\nmerge_file = os.path.join(args.resume_from_checkpoint, \"merges.txt\")\nother_megatron_args = {\"vocab_file\": vocab_file, \"merge_file\": merge_file}\nmegatron_lm_plugin = MegatronLMPlugin(other_megatron_args=other_megatron_args)\n\n# inference using `megatron_generate` functionality\ntokenizer.pad_token = tokenizer.eos_token\nmax_new_tokens = 64\nbatch_texts = [\n \"Are you human?\",\n \"The purpose of life is\",\n \"The arsenal was constructed at the request of\",\n \"How are you doing these days?\",\n]\nbatch_encodings = tokenizer(batch_texts, return_tensors=\"pt\", padding=True)\n\n# top-p sampling\ngenerated_tokens = model.megatron_generate(\n batch_encodings[\"input_ids\"],\n batch_encodings[\"attention_mask\"],\n max_new_tokens=max_new_tokens,\n top_p=0.8,\n top_p_decay=0.5,\n temperature=0.9,\n)\ndecoded_preds = tokenizer.batch_decode(generated_tokens.cpu().numpy())\naccelerator.print(decoded_preds)\n\n# top-k sampling\ngenerated_tokens = model.megatron_generate(\n batch_encodings[\"input_ids\"],\n batch_encodings[\"attention_mask\"],\n max_new_tokens=max_new_tokens,\n top_k=50,\n temperature=0.9,\n)\ndecoded_preds = tokenizer.batch_decode(generated_tokens.cpu().numpy())\naccelerator.print(decoded_preds)\n\n# adding `bos` token at the start\ngenerated_tokens = model.megatron_generate(\n batch_encodings[\"input_ids\"], batch_encodings[\"attention_mask\"], max_new_tokens=max_new_tokens, add_BOS=True\n)\ndecoded_preds = tokenizer.batch_decode(generated_tokens.cpu().numpy())\naccelerator.print(decoded_preds)\n\n# beam search => only takes single prompt\nbatch_texts = [\"The purpose of life is\"]\nbatch_encodings = tokenizer(batch_texts, return_tensors=\"pt\", padding=True)\ngenerated_tokens = model.megatron_generate(\n batch_encodings[\"input_ids\"],\n batch_encodings[\"attention_mask\"],\n max_new_tokens=max_new_tokens,\n num_beams=20,\n length_penalty=1.5,\n)\ndecoded_preds = tokenizer.batch_decode(generated_tokens.cpu().numpy())\naccelerator.print(decoded_preds)\n```\n\n3. An end-to-end example of using `megatron_generate` method for Megatron-LM GPT model is available at\n[megatron_gpt2_generation.py](https://github.com/pacman100/accelerate-megatron-test/blob/main/src/inference/megatron_gpt2_generation.py) with \nconfig file [megatron_lm_gpt_generate_config.yaml](https://github.com/pacman100/accelerate-megatron-test/blob/main/src/Configs/megatron_lm_gpt_generate_config.yaml).\nThe bash script with accelerate launch command is available at [megatron_lm_gpt_generate.sh](https://github.com/pacman100/accelerate-megatron-test/blob/main/megatron_lm_gpt_generate.sh).\nThe output logs of the script are available at [megatron_lm_gpt_generate.log](https://github.com/pacman100/accelerate-megatron-test/blob/main/output_logs/megatron_lm_gpt_generate.log).\n\n## Support for ROPE and ALiBi Positional embeddings and Multi-Query Attention\n\n1. For ROPE/ALiBi attention, pass `position_embedding_type` with `(\"absolute\" | \"rotary\" | \"alibi\")` to `MegatronLMPlugin` as shown below.\n```python\nother_megatron_args = {\"position_embedding_type\": \"alibi\"}\nmegatron_lm_plugin = MegatronLMPlugin(other_megatron_args=other_megatron_args)\n```\n\n2. For Multi-Query Attention, pass `attention_head_type` with `(\"multihead\" | \"multiquery\")` to `MegatronLMPlugin` as shown below.\n```python\nother_megatron_args = {\"attention_head_type\": \"multiquery\"}\nmegatron_lm_plugin = MegatronLMPlugin(other_megatron_args=other_megatron_args)\n```\n\n## Caveats\n\n1. Supports Transformers GPT2, Megatron-BERT and T5 models.\nThis covers Decoder only, Encode only and Encoder-Decoder model classes.\n\n2. Only loss is returned from model forward pass as \nthere is quite complex interplay of pipeline, tensor and data parallelism behind the scenes.\nThe `model(**batch_data)` call return loss(es) averaged across the data parallel ranks.\nThis is fine for most cases wherein pre-training jobs are run using Megatron-LM features and\nyou can easily compute the `perplexity` using the loss. \nFor GPT model, returning logits in addition to loss(es) is supported. \nThese logits aren't gathered across data parallel ranks. Use `accelerator.utils.gather_across_data_parallel_groups`\nto gather logits across data parallel ranks. These logits along with labels can be used for computing various \nperformance metrics. \n\n3. The main process is the last rank as the losses/logits are available in the last stage of pipeline.\n`accelerator.is_main_process` and `accelerator.is_local_main_process` return `True` for last rank when using \nMegatron-LM integration.\n\n4. In `accelerator.prepare` call, a Megatron-LM model corresponding to a given Transformers model is created\nwith random weights. Please use `accelerator.load_state` to load the Megatron-LM checkpoint with matching TP, PP and DP partitions.\n\n5. Currently, checkpoint reshaping and interoperability support is only available for GPT. \nSoon it will be extended to BERT and T5.\n\n6. `gradient_accumulation_steps` needs to be 1. When using Megatron-LM, micro batches in pipeline parallelism \nsetting is synonymous with gradient accumulation. \n\n7. When using Megatron-LM, use `accelerator.save_state` and `accelerator.load_state` for saving and loading checkpoints.\n\n8. Below are the mapping from Megatron-LM model architectures to the equivalent transformers model architectures.\nOnly these transformers model architectures are supported.\n\na. Megatron-LM [BertModel](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/bert_model.py) : \ntransformers models with `megatron-bert` in config's model type, e.g., \n[MegatronBERT](https://huggingface.co/docs/transformers/model_doc/megatron-bert)\n \nb. Megatron-LM [GPTModel](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py) : \ntransformers models with `gpt2` in config's model type, e.g., \n[OpenAI GPT2](https://huggingface.co/docs/transformers/model_doc/gpt2)\n \nc. Megatron-LM [T5Model](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/t5_model.py) : \ntransformers models with `t5` in config's model type, e.g., \n[T5](https://huggingface.co/docs/transformers/model_doc/t5) and \n[MT5](https://huggingface.co/docs/transformers/model_doc/mt5)\n" }, { "path": "docs/source/usage_guides/model_size_estimator.md", "content": "\n\n# Model memory estimator\n\nOne 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 device (such as loading the model onto CUDA or XPU).\n\nTo help alleviate this, Accelerate has a CLI interface through `accelerate estimate-memory`. This tutorial will \nhelp walk you through using it, what to expect, and at the end link to the interactive demo hosted on the Hub which will \neven let you post those results directly on the model repo!\n\nCurrently we support searching for models that can be used in `timm` and `transformers`.\n\n\n\n This API will load the model into memory on the `meta` device, so we are not actually downloading \n and loading the full weights of the model into memory, nor do we need to. As a result it's \n perfectly fine to measure 8 billion parameter models (or more), without having to worry about \n if your CPU can handle it!\n\n\n\n## Gradio Demos\n\nBelow are a few gradio demos related to what was described above. The first is the official Hugging Face memory estimation space, utilizing Accelerate directly:\n\n
\n\t\n
\n
\n \n
\n\nA community member has taken the idea and expanded it further, allowing you to filter models directly and see if you can run a particular LLM given GPU constraints and LoRA configurations. To play with it, see [here](https://huggingface.co/spaces/Vokturz/can-it-run-llm) for more details.\n\n## The Command\n\nWhen using `accelerate estimate-memory`, you need to pass in the name of the model you want to use, potentially the framework\nthat model utilizing (if it can't be found automatically), and the data types you want the model to be loaded in with.\n\nFor example, here is how we can calculate the memory footprint for `bert-base-cased`:\n\n```bash\naccelerate estimate-memory bert-base-cased\n```\n\nThis will download the `config.json` for `bert-based-cased`, load the model on the `meta` device, and report back how much space\nit will use:\n\nMemory Usage for loading `bert-base-cased`:\n\n| dtype | Largest Layer | Total Size | Training using Adam |\n|---------|---------------|------------|---------------------|\n| float32 | 84.95 MB | 418.18 MB | 1.61 GB |\n| float16 | 42.47 MB | 206.59 MB | 826.36 MB |\n| int8 | 21.24 MB | 103.29 MB | 413.18 MB |\n| int4 | 10.62 MB | 51.65 MB | 206.59 MB |\n\nBy default it will return all the supported dtypes (`int4` through `float32`), but if you are interested in specific ones these can be filtered.\n\n### Specific libraries\n\nIf the source library cannot be determined automatically (like it could in the case of `bert-base-cased`), a library name can\nbe passed in. \n\n```bash\naccelerate estimate-memory HuggingFaceM4/idefics-80b-instruct --library_name transformers\n```\n\nMemory Usage for loading `HuggingFaceM4/idefics-80b-instruct`:\n\n| dtype | Largest Layer | Total Size | Training using Adam |\n|---------|---------------|------------|---------------------|\n| float32 | 3.02 GB | 297.12 GB | 1.16 TB |\n| float16 | 1.51 GB | 148.56 GB | 594.24 GB |\n| int8 | 772.52 MB | 74.28 GB | 297.12 GB |\n| int4 | 386.26 MB | 37.14 GB | 148.56 GB |\n\n\n```bash\naccelerate estimate-memory timm/resnet50.a1_in1k --library_name timm\n```\n\nMemory Usage for loading `timm/resnet50.a1_in1k`:\n\n| dtype | Largest Layer | Total Size | Training using Adam |\n|---------|---------------|------------|---------------------|\n| float32 | 9.0 MB | 97.7 MB | 390.78 MB |\n| float16 | 4.5 MB | 48.85 MB | 195.39 MB |\n| int8 | 2.25 MB | 24.42 MB | 97.7 MB |\n| int4 | 1.12 MB | 12.21 MB | 48.85 MB |\n\n### Specific dtypes\n\nAs mentioned earlier, while we return `int4` through `float32` by default, any dtype can be used from `float32`, `float16`, `int8`, and `int4`.\n\nTo do so, pass them in after specifying `--dtypes`:\n\n```bash\naccelerate estimate-memory bert-base-cased --dtypes float32 float16\n```\n\nMemory Usage for loading `bert-base-cased`:\n\n| dtype | Largest Layer | Total Size | Training using Adam |\n|---------|---------------|------------|---------------------|\n| float32 | 84.95 MB | 413.18 MB | 1.61 GB |\n| float16 | 42.47 MB | 206.59 MB | 826.36 MB |\n\n## Caveats with this calculator\n\nThis calculator will tell you how much memory is needed to purely load the model in, *not* to perform inference.\n\nThis 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`.\n\nWhen 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 \nthis calculator once done.\n" }, { "path": "docs/source/usage_guides/mps.md", "content": "\n\n# Accelerated PyTorch Training on Mac\n\nWith PyTorch v1.12 release, developers and researchers can take advantage of Apple silicon GPUs for significantly faster model training. \nThis unlocks the ability to perform machine learning workflows like prototyping and fine-tuning locally, right on Mac.\nApple's Metal Performance Shaders (MPS) as a backend for PyTorch enables this and can be used via the new `\"mps\"` device. \nThis will map computational graphs and primitives on the MPS Graph framework and tuned kernels provided by MPS.\nFor more information please refer official documents [Introducing Accelerated PyTorch Training on Mac](https://pytorch.org/blog/introducing-accelerated-pytorch-training-on-mac/)\nand [MPS BACKEND](https://pytorch.org/docs/stable/notes/mps.html).\n\n### Benefits of Training and Inference using Apple Silicon Chips\n\n1. Enables users to train larger networks or batch sizes locally\n2. Reduces data retrieval latency and provides the GPU with direct access to the full memory store due to unified memory architecture. \nTherefore, improving end-to-end performance.\n3. Reduces costs associated with cloud-based development or the need for additional local GPUs.\n\n**Pre-requisites**: To install torch with mps support, \nplease follow this nice medium article [GPU-Acceleration Comes to PyTorch on M1 Macs](https://medium.com/towards-data-science/gpu-acceleration-comes-to-pytorch-on-m1-macs-195c399efcc1).\n\n\n## How it works out of the box\nIt is enabled by default on MacOs machines with MPS enabled Apple Silicon GPUs.\nTo disable it, pass `--cpu` flag to `accelerate launch` command or answer the corresponding question when answering the `accelerate config` questionnaire.\n\nYou can directly run the following script to test it out on MPS enabled Apple Silicon machines:\n```bash\naccelerate launch /examples/cv_example.py --data_dir images\n```\n\n## A few caveats to be aware of\n\n1. Distributed setups `gloo` and `nccl` are not working with `mps` device. \nThis means that currently only single GPU of `mps` device type can be used.\n\nFinally, please, remember that, `Accelerate` only integrates MPS backend, therefore if you\nhave any problems or questions with regards to MPS backend usage, please, file an issue with [PyTorch GitHub](https://github.com/pytorch/pytorch/issues)." }, { "path": "docs/source/usage_guides/profiler.md", "content": "\n\n# Profiler\n\nProfiler is a tool that allows the collection of performance metrics during training and inference. Profilerā€™s context manager API can be used to better understand what model operators are the most expensive, examine their input shapes and stack traces, study device kernel activity, and visualize the execution trace. It provides insights into the performance of your model, allowing you to optimize and improve it.\n\nThis guide explains how to use PyTorch Profiler to measure the time and memory consumption of the modelā€™s operators and how to integrate this with Accelerate. We will cover various use cases and provide examples for each.\n\n## Using profiler to analyze execution time\n\nProfiler allows one to check which operators were called during the execution of a code range wrapped with a profiler context manager.\n\nLetā€™s see how we can use profiler to analyze the execution time:\n\n\n\n\n```python\nimport torch\nimport torchvision.models as models\nfrom torch.profiler import profile, record_function, ProfilerActivity\n\nmodel = models.resnet18()\ninputs = torch.randn(5, 3, 224, 224)\n\nwith profile(activities=[ProfilerActivity.CPU], record_shapes=True) as prof:\n model(inputs)\n\nprint(prof.key_averages().table(sort_by=\"cpu_time_total\", row_limit=10))\n```\n\n\n\n\n```python\nfrom accelerate import Accelerator, ProfileKwargs\nimport torch\nimport torchvision.models as models\n\nmodel = models.resnet18()\ninputs = torch.randn(5, 3, 224, 224)\n\nprofile_kwargs = ProfileKwargs(\n activities=[\"cpu\"],\n record_shapes=True\n)\n\naccelerator = Accelerator(cpu=True, kwargs_handlers=[profile_kwargs])\nmodel = accelerator.prepare(model)\n\nwith accelerator.profile() as prof:\n with torch.no_grad():\n model(inputs)\n\nprint(prof.key_averages().table(sort_by=\"cpu_time_total\", row_limit=10))\n```\n\n\n\n\nThe resulting table output (omitting some columns):\n\n```\n--------------------------------- ------------ ------------ ------------ ------------ \n Name Self CPU CPU total CPU time avg # of Calls \n--------------------------------- ------------ ------------ ------------ ------------ \n aten::conv2d 171.000us 52.260ms 2.613ms 20 \n aten::convolution 227.000us 52.089ms 2.604ms 20 \n aten::_convolution 270.000us 51.862ms 2.593ms 20 \n aten::mkldnn_convolution 51.273ms 51.592ms 2.580ms 20 \n aten::batch_norm 118.000us 7.059ms 352.950us 20 \n aten::_batch_norm_impl_index 315.000us 6.941ms 347.050us 20 \n aten::native_batch_norm 6.305ms 6.599ms 329.950us 20 \n aten::max_pool2d 40.000us 4.008ms 4.008ms 1 \n aten::max_pool2d_with_indices 3.968ms 3.968ms 3.968ms 1 \n aten::add_ 780.000us 780.000us 27.857us 28 \n--------------------------------- ------------ ------------ ------------ ------------ \nSelf CPU time total: 67.016ms\n```\n\nTo get a finer granularity of results and include operator input shapes, pass `group_by_input_shape=True` (note: this requires running the profiler with `record_shapes=True`):\n\n```python\nprint(prof.key_averages(group_by_input_shape=True).table(sort_by=\"cpu_time_total\", row_limit=10))\n```\n\n## Using profiler to analyze memory consumption\n\nProfiler can also show the amount of memory (used by the modelā€™s tensors) that was allocated (or released) during the execution of the modelā€™s operators. To enable memory profiling functionality pass `profile_memory=True`.\n\n\n\n\n```python\nmodel = models.resnet18()\ninputs = torch.randn(5, 3, 224, 224)\n\nwith profile(activities=[ProfilerActivity.CPU],\n profile_memory=True, record_shapes=True) as prof:\n model(inputs)\n\nprint(prof.key_averages().table(sort_by=\"self_cpu_memory_usage\", row_limit=10))\n```\n\n\n\n\n```python\nmodel = models.resnet18()\ninputs = torch.randn(5, 3, 224, 224)\n\nprofile_kwargs = ProfileKwargs(\n activities=[\"cpu\"],\n profile_memory=True,\n record_shapes=True\n)\n\naccelerator = Accelerator(cpu=True, kwargs_handlers=[profile_kwargs])\nmodel = accelerator.prepare(model)\n\nwith accelerator.profile() as prof:\n model(inputs)\n\nprint(prof.key_averages().table(sort_by=\"self_cpu_memory_usage\", row_limit=10))\n```\n\n\n\n\nThe resulting table output (omitting some columns):\n\n```\n--------------------------------- ------------ ------------ ------------ \n Name CPU Mem Self CPU Mem # of Calls \n--------------------------------- ------------ ------------ ------------ \n aten::empty 94.85 Mb 94.85 Mb 205 \n aten::max_pool2d_with_indices 11.48 Mb 11.48 Mb 1 \n aten::addmm 19.53 Kb 19.53 Kb 1 \n aten::mean 10.00 Kb 10.00 Kb 1 \n aten::empty_strided 492 b 492 b 5 \n aten::cat 240 b 240 b 6 \n aten::abs 480 b 240 b 4 \n aten::masked_select 120 b 112 b 1 \n aten::ne 61 b 53 b 3 \n aten::eq 30 b 30 b 1 \n--------------------------------- ------------ ------------ ------------ \nSelf CPU time total: 69.332ms\n```\n\n\n## Exporting chrome trace\n\nYou can examine the sequence of profiled operators and CUDA kernels in Chrome trace viewer (`chrome://tracing`):\n\n![profile_export](https://github.com/huggingface/accelerate/assets/100389977/5acb193f-6d11-4f7b-9873-c600c19e8172)\n\n\n\n\n```python\nmodel = models.resnet18().cuda()\ninputs = torch.randn(5, 3, 224, 224).cuda()\n\nwith profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:\n model(inputs)\n\nprof.export_chrome_trace(\"trace.json\")\n```\n\n\n\n\n```python\nmodel = models.resnet18()\ninputs = torch.randn(5, 3, 224, 224).cuda()\nprofile_kwargs = ProfileKwargs(\n activities=[\"cpu\", \"cuda\"],\n output_trace_dir=\"trace\"\n)\n\naccelerator = Accelerator(kwargs_handlers=[profile_kwargs])\nmodel = accelerator.prepare(model)\n\nwith accelerator.profile() as prof:\n model(inputs)\n\n# The trace will be saved to the specified directory\n```\nFor other hardware accelerators, e.g. XPU, you can change `cuda` to `xpu` in the above example code.\n\n\n\n\n## Using Profiler to Analyze Long-Running Jobs\n\nProfiler offers an additional API to handle long-running jobs (such as training loops). Tracing all of the execution can be slow and result in very large trace files. To avoid this, use optional arguments:\n\n- `schedule_option`: Scheduling options allow you to control when profiling is active. This is useful for long-running jobs to avoid collecting too much data. Available keys are `wait`, `warmup`, `active`, `repeat` and `skip_first`. The profiler will skip the first `skip_first` steps, then wait for `wait` steps, then do the warmup for the next `warmup` steps, then do the active recording for the next `active` steps and then repeat the cycle starting with `wait` steps. The optional number of cycles is specified with the `repeat` parameter, the zero value means that the cycles will continue until the profiling is finished.\n- `on_trace_ready`: specifies a function that takes a reference to the profiler as an input and is called by the profiler each time the new trace is ready.\n\nTo illustrate how the API works, consider the following example:\n\n\n\n\n```python\nfrom torch.profiler import schedule\n\nmy_schedule = schedule(\n skip_first=1,\n wait=5,\n warmup=1,\n active=3,\n repeat=2\n)\n\ndef trace_handler(p):\n output = p.key_averages().table(sort_by=\"self_cuda_time_total\", row_limit=10)\n print(output)\n p.export_chrome_trace(\"/tmp/trace_\" + str(p.step_num) + \".json\")\n\nwith profile(\n activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],\n schedule=my_schedule,\n on_trace_ready=trace_handler\n) as p:\n for idx in range(8):\n model(inputs)\n p.step()\n```\n\n\n\n\n```python\ndef trace_handler(p):\n output = p.key_averages().table(sort_by=\"self_cuda_time_total\", row_limit=10)\n print(output)\n p.export_chrome_trace(\"/tmp/trace_\" + str(p.step_num) + \".json\")\n\nprofile_kwargs = ProfileKwargs(\n activities=[\"cpu\", \"cuda\"],\n schedule_option={\"wait\": 5, \"warmup\": 1, \"active\": 3, \"repeat\": 2, \"skip_first\": 1},\n on_trace_ready=trace_handler\n)\n\naccelerator = Accelerator(kwargs_handlers=[profile_kwargs])\nmodel = accelerator.prepare(model)\n\nwith accelerator.profile() as prof:\n for idx in range(8):\n model(inputs)\n prof.step()\n```\n\n\n\n\n## FLOPS\n\nUse formula to estimate the FLOPs (floating point operations) of specific operators (matrix multiplication and 2D convolution).\n\nTo measure floating-point operations (FLOPS):\n\n\n\n\n```python\nwith profile(\n activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],\n with_flops=True\n) as prof:\n model(inputs)\n\nprint(prof.key_averages().table(sort_by=\"flops\", row_limit=10))\n```\n\n\n\n\n```python\nprofile_kwargs = ProfileKwargs(\n with_flops=True\n)\naccelerator = Accelerator(kwargs_handlers=[profile_kwargs])\n\nwith accelerator.profile() as prof:\n model(inputs)\n\nprint(prof.key_averages().table(sort_by=\"flops\", row_limit=10))\n```\n\n\n\n\nThe resulting table output (omitting some columns):\n\n```\n------------------------------------------------------- ------------ ------------ ------------ \n Name Self CPU Self CUDA Total FLOPs \n------------------------------------------------------- ------------ ------------ ------------ \n aten::conv2d 197.000us 0.000us 18135613440.000 \n aten::addmm 103.000us 17.000us 5120000.000 \n aten::mul 29.000us 2.000us 30.000 \n aten::convolution 409.000us 0.000us -- \n aten::_convolution 253.000us 0.000us -- \n aten::cudnn_convolution 5.465ms 2.970ms -- \n cudaEventRecord 138.000us 0.000us -- \n cudaStreamIsCapturing 43.000us 0.000us -- \n cudaStreamGetPriority 40.000us 0.000us -- \n cudaDeviceGetStreamPriorityRange 10.000us 0.000us -- \n------------------------------------------------------- ------------ ------------ ------------ \nSelf CPU time total: 21.938ms\nSelf CUDA time total: 4.165ms\n```\n\n\n\n## Conclusion and Further Information\n\nPyTorch Profiler is a powerful tool for analyzing the performance of your models. By integrating it with Accelerate, you can easily profile your models and gain insights into their performance, helping you to optimize and improve them.\n\nFor more detailed information, refer to the [PyTorch Profiler documentation](https://pytorch.org/docs/stable/profiler.html)." }, { "path": "docs/source/usage_guides/quantization.md", "content": "\n\n# Model quantization\n\n## `bitsandbytes` Integration\n\nAccelerate 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.\n\nIf you want to use Transformers models with `bitsandbytes`, you should follow this [documentation](https://huggingface.co/docs/transformers/main_classes/quantization). \n\nTo 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).\n\n### Pre-Requisites\nYou will need to install the following requirements:\n\n- Install `bitsandbytes` library\n```bash\npip install bitsandbytes\n```\nFor non-cuda devices, you can refer to the bitsandbytes installation guide [here](https://huggingface.co/docs/bitsandbytes/main/en/installation#multi-backend).\n\n- Install latest `accelerate` from source\n```bash\npip install git+https://github.com/huggingface/accelerate.git\n```\n- Install `minGPT` and `huggingface_hub` to run examples\n```bash\ngit clone https://github.com/karpathy/minGPT.git\npip install minGPT/\npip install huggingface_hub\n```\n\n### How it works\n\nFirst, we need to initialize our model. To save memory, we can initialize an empty model using the context manager [`init_empty_weights`]. \n\nLet's take the GPT2 model from minGPT library.\n```py\nfrom accelerate import init_empty_weights\nfrom mingpt.model import GPT\n\nmodel_config = GPT.get_default_config()\nmodel_config.model_type = 'gpt2-xl'\nmodel_config.vocab_size = 50257\nmodel_config.block_size = 1024\n\nwith init_empty_weights():\n empty_model = GPT(model_config)\n```\n\nThen, 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. \n\n```py\nfrom huggingface_hub import snapshot_download\nweights_location = snapshot_download(repo_id=\"marcsun13/gpt2-xl-linear-sharded\")\n```\n\nFinally, you need to set your quantization configuration with [`~utils.BnbQuantizationConfig`].\n\nHere's an example for 8-bit quantization:\n```py\nfrom accelerate.utils import BnbQuantizationConfig\nbnb_quantization_config = BnbQuantizationConfig(load_in_8bit=True, llm_int8_threshold = 6)\n```\n\nHere's an example for 4-bit quantization:\n```py\nfrom accelerate.utils import BnbQuantizationConfig\nbnb_quantization_config = BnbQuantizationConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type=\"nf4\")\n```\n\nTo quantize your empty model with the selected configuration, you need to use [`~utils.load_and_quantize_model`]. \n\n```py\nfrom accelerate.utils import load_and_quantize_model\nquantized_model = load_and_quantize_model(empty_model, weights_location=weights_location, bnb_quantization_config=bnb_quantization_config)\n```\n\n### Saving and loading 8-bit model\n\nYou can save your 8-bit model with accelerate using [`~Accelerator.save_model`]. \n\n```py\nfrom accelerate import Accelerator\naccelerate = Accelerator()\nnew_weights_location = \"path/to/save_directory\"\naccelerate.save_model(quantized_model, new_weights_location)\n\nquantized_model_from_saved = load_and_quantize_model(empty_model, weights_location=new_weights_location, bnb_quantization_config=bnb_quantization_config, device_map = \"auto\")\n```\n\nNote that 4-bit model serialization is currently not supported.\n\n### Offload modules to cpu and disk \n\nYou 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.\nThis uses big model inference under the hood. Check this [documentation](https://huggingface.co/docs/accelerate/usage_guides/big_modeling) for more details. \n\nFor 8-bit quantization, the selected modules will be converted to 8-bit precision. \n\nFor 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. \n\n 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 :\n\n```py\ndevice_map = {\n \"transformer.wte\": 0,\n \"transformer.wpe\": 0,\n \"transformer.drop\": 0,\n \"transformer.h\": \"cpu\",\n \"transformer.ln_f\": \"disk\",\n \"lm_head\": \"disk\",\n}\n```\n### Fine-tune a quantized model\n\nIt 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.\n\nCurrently, 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 fine-tune 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. \n\nNote 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.\n\n### Example demo - running GPT2 1.5b on a Google Colab\n\nCheck out the Google Colab [demo](https://colab.research.google.com/drive/1T1pOgewAWVpR9gKpaEWw4orOrzPFb3yM?usp=sharing) for running quantized models on a GPT2 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.\n" }, { "path": "docs/source/usage_guides/sagemaker.md", "content": "\n\n# Amazon SageMaker\n\nHugging Face and Amazon introduced new [Hugging Face Deep Learning Containers (DLCs)](https://github.com/aws/deep-learning-containers/blob/master/available_images.md#huggingface-training-containers) to\nmake it easier than ever to train Hugging Face Transformer models in [Amazon SageMaker](https://aws.amazon.com/sagemaker/).\n\n## Getting Started\n\n### Setup & Installation\n\n\nBefore you can run your Accelerate scripts on Amazon SageMaker you need to sign up for an AWS account. If you do not\nhave an AWS account yet learn more [here](https://docs.aws.amazon.com/sagemaker/latest/dg/gs-set-up.html).\n\nAfter you have your AWS Account you need to install the `sagemaker` sdk for Accelerate with:\n\n```bash\npip install \"accelerate[sagemaker]\" --upgrade\n```\n\nAccelerate currently uses the DLCs, with `transformers`, `datasets` and `tokenizers` pre-installed. Accelerate is not in the DLC yet (will soon be added!) so to use it within Amazon SageMaker you need to create a\n`requirements.txt` in the same directory where your training script is located and add it as dependency:\n\n```\naccelerate\n```\n\nYou should also add any other dependencies you have to this `requirements.txt`.\n\n\n### Configure Accelerate\n\nYou can configure the launch configuration for Amazon SageMaker the same as you do for non SageMaker training jobs with\nthe Accelerate CLI:\n\n```bash\naccelerate config\n# In which compute environment are you running? ([0] This machine, [1] AWS (Amazon SageMaker)): 1\n```\n\nAccelerate will go through a questionnaire about your Amazon SageMaker setup and create a config file you can edit.\n\n\n\n Accelerate is not saving any of your credentials.\n\n\n\n### Prepare a Accelerate fine-tuning script\n\nThe training script is very similar to a training script you might run outside of SageMaker, but to save your model\nafter training you need to specify either `/opt/ml/model` or use `os.environ[\"SM_MODEL_DIR\"]` as your save\ndirectory. After training, artifacts in this directory are uploaded to S3:\n\n\n```diff\n- torch.save('/opt/ml/model`)\n+ accelerator.save('/opt/ml/model')\n```\n\n\n\n SageMaker doesnā€™t support argparse actions. If you want to use, for example, boolean hyperparameters, you need to\n specify type as bool in your script and provide an explicit True or False value for this hyperparameter. [[REF]](https://sagemaker.readthedocs.io/en/stable/frameworks/pytorch/using_pytorch.html#prepare-a-pytorch-training-script).\n\n\n\n### Launch Training\n\nYou can launch your training with Accelerate CLI with:\n\n```\naccelerate launch path_to_script.py --args_to_the_script\n```\n\nThis will launch your training script using your configuration. The only thing you have to do is provide all the\narguments needed by your training script as named arguments.\n\n**Examples**\n\n\n\n If you run one of the example scripts, don't forget to add `accelerator.save('/opt/ml/model')` to it.\n\n\n\n```bash\naccelerate launch ./examples/sagemaker_example.py\n```\n\nOutputs:\n\n```\nConfiguring Amazon SageMaker environment\nConverting Arguments to Hyperparameters\nCreating Estimator\n2021-04-08 11:56:50 Starting - Starting the training job...\n2021-04-08 11:57:13 Starting - Launching requested ML instancesProfilerReport-1617883008: InProgress\n.........\n2021-04-08 11:58:54 Starting - Preparing the instances for training.........\n2021-04-08 12:00:24 Downloading - Downloading input data\n2021-04-08 12:00:24 Training - Downloading the training image..................\n2021-04-08 12:03:39 Training - Training image download completed. Training in progress..\n........\nepoch 0: {'accuracy': 0.7598039215686274, 'f1': 0.8178438661710037}\nepoch 1: {'accuracy': 0.8357843137254902, 'f1': 0.882249560632689}\nepoch 2: {'accuracy': 0.8406862745098039, 'f1': 0.8869565217391304}\n........\n2021-04-08 12:05:40 Uploading - Uploading generated training model\n2021-04-08 12:05:40 Completed - Training job completed\nTraining seconds: 331\nBillable seconds: 331\nYou can find your model data at: s3://your-bucket/accelerate-sagemaker-1-2021-04-08-11-56-47-108/output/model.tar.gz\n```\n\n## Advanced Features\n\n### Distributed Training: Data Parallelism\n\nSet up the accelerate config by running `accelerate config` and answer the SageMaker questions and set it up.\nTo use SageMaker DDP, select it when asked \n`What is the distributed mode? ([0] No distributed training, [1] data parallelism):`.\nExample config below:\n```yaml\nbase_job_name: accelerate-sagemaker-1\ncompute_environment: AMAZON_SAGEMAKER\ndistributed_type: DATA_PARALLEL\nec2_instance_type: ml.p3.16xlarge\niam_role_name: xxxxx\nimage_uri: null\nmixed_precision: fp16\nnum_machines: 1\nprofile: xxxxx\npy_version: py10\npytorch_version: 2.5.0\nregion: us-east-1\ntransformers_version: 4.17.0\nuse_cpu: false\n```\n\n### Distributed Training: Model Parallelism\n\n*currently in development, will be supported soon.*\n\n### Python packages and dependencies\n\nAccelerate currently uses the DLCs, with `transformers`, `datasets` and `tokenizers` pre-installed. If you\nwant to use different/other Python packages you can do this by adding them to the `requirements.txt`. These packages\nwill be installed before your training script is started.\n\n### Local Training: SageMaker Local mode\n\nThe local mode in the SageMaker SDK allows you to run your training script locally inside the HuggingFace DLC (Deep Learning container) \nor using your custom container image. This is useful for debugging and testing your training script inside the final container environment.\nLocal mode uses Docker compose (*Note: Docker Compose V2 is not supported yet*). The SDK will handle the authentication against ECR\nto pull the DLC to your local environment. You can emulate CPU (single and multi-instance) and GPU (single instance) SageMaker training jobs.\n\nTo use local mode, you need to set your `ec2_instance_type` to `local`.\n\n```yaml\nec2_instance_type: local\n```\n\n### Advanced configuration\n\nThe configuration allows you to override parameters for the [Estimator](https://sagemaker.readthedocs.io/en/stable/api/training/estimators.html).\nThese 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.\n\n```yaml\nadditional_args:\n # enable network isolation to restrict internet access for containers\n enable_network_isolation: True\n```\n\nYou can find all available configuration [here](https://sagemaker.readthedocs.io/en/stable/api/training/estimators.html).\n\n### Use Spot Instances\n\nYou can use Spot Instances e.g. using (see [Advanced configuration](#advanced-configuration)):\n```yaml\nadditional_args:\n use_spot_instances: True\n max_wait: 86400\n```\n\n*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.*\n\n### Remote scripts: Use scripts located on Github\n\n*undecided if feature is needed. Contact us if you would like this feature.*" }, { "path": "docs/source/usage_guides/tracking.md", "content": "\n\n# Experiment trackers\n\nThere are a large number of experiment tracking APIs available, however getting them all to work in a multi-processing environment can oftentimes be complex.\nAccelerate provides a general tracking API that can be used to log useful items during your script through [`Accelerator.log`]\n\n## Integrated Trackers\n\nCurrently `Accelerate` supports eight trackers out-of-the-box:\n\n- TensorBoard\n- WandB \n- Trackio\n- CometML\n- Aim\n- MLFlow\n- ClearML\n- DVCLive\n\nTo use any of them, pass in the selected type(s) to the `log_with` parameter in [`Accelerate`]:\n```python\nfrom accelerate import Accelerator\nfrom accelerate.utils import LoggerType\n\naccelerator = Accelerator(log_with=\"all\") # For all available trackers in the environment\naccelerator = Accelerator(log_with=\"wandb\")\naccelerator = Accelerator(log_with=[\"wandb\", LoggerType.TENSORBOARD])\n```\n\nAt the start of your experiment [`Accelerator.init_trackers`] should be used to setup your project, and potentially add any experiment hyperparameters to be logged:\n```python\nhps = {\"num_iterations\": 5, \"learning_rate\": 1e-2}\naccelerator.init_trackers(\"my_project\", config=hps)\n```\n\nWhen you are ready to log any data, [`Accelerator.log`] should be used.\nA `step` can also be passed in to correlate the data with a particular step in the training loop.\n```python\naccelerator.log({\"train_loss\": 1.12, \"valid_loss\": 0.8}, step=1)\n```\n\nOnce you've finished training, make sure to run [`Accelerator.end_training`] so that all the trackers can run their finish functionalities if they have any.\n```python\naccelerator.end_training()\n```\n\n\nA full example is below:\n```python\nfrom accelerate import Accelerator\n\naccelerator = Accelerator(log_with=\"all\")\nconfig = {\n \"num_iterations\": 5,\n \"learning_rate\": 1e-2,\n \"loss_function\": str(my_loss_function),\n}\n\naccelerator.init_trackers(\"example_project\", config=config)\n\nmy_model, my_optimizer, my_training_dataloader = accelerator.prepare(my_model, my_optimizer, my_training_dataloader)\ndevice = accelerator.device\nmy_model.to(device)\n\nfor iteration in range(config[\"num_iterations\"]):\n for step, batch in enumerate(my_training_dataloader):\n my_optimizer.zero_grad()\n inputs, targets = batch\n inputs = inputs.to(device)\n targets = targets.to(device)\n outputs = my_model(inputs)\n loss = my_loss_function(outputs, targets)\n accelerator.backward(loss)\n my_optimizer.step()\n accelerator.log({\"training_loss\": loss}, step=step)\naccelerator.end_training()\n```\n\nIf 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 \nwhen 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`: \n\n```python\naccelerator = Accelerator(log_with=\"tensorboard\", project_dir=\".\")\n\n# use with ProjectConfiguration\nconfig = ProjectConfiguration(project_dir=\".\", logging_dir=\"another/directory\")\naccelerator = Accelerator(log_with=\"tensorboard\", project_config=config)\n```\n\n## Implementing Custom Trackers\n\nTo implement a new tracker to be used in `Accelerator`, a new one can be made through implementing the [`GeneralTracker`] class.\nEvery tracker must implement three functions and have three properties:\n - `__init__`: \n - Should store a `run_name` and initialize the tracker API of the integrated library. \n - If a tracker stores their data locally (such as TensorBoard), a `logging_dir` parameter can be added.\n - `store_init_configuration`: \n - Should take in a `values` dictionary and store them as a one-time experiment configuration\n - `log`: \n - Should take in a `values` dictionary and a `step`, and should log them to the run\n\n - `name` (`str`):\n - A unique string name for the tracker, such as `\"wandb\"` for the wandb tracker. \n - This will be used for interacting with this tracker specifically\n - `requires_logging_directory` (`bool`):\n - Whether a `logging_dir` is needed for this particular tracker and if it uses one.\n - `tracker`: \n - This should be implemented as a `@property` function \n - Should return the internal tracking mechanism the library uses, such as the `run` object for `wandb`.\n\nEach method should also utilize the [`state.PartialState`] class if the logger should only be executed on the main process for instance.\n\nA brief example can be seen below with an integration with Weights and Biases, containing only the relevant information and logging just on \nthe main process:\n```python\nfrom accelerate.tracking import GeneralTracker, on_main_process\nfrom typing import Optional\n\nimport wandb\n\n\nclass MyCustomTracker(GeneralTracker):\n name = \"wandb\"\n requires_logging_directory = False\n\n @on_main_process\n def __init__(self, run_name: str):\n self.run_name = run_name\n run = wandb.init(self.run_name)\n\n @property\n def tracker(self):\n return self.run.run\n\n @on_main_process\n def store_init_configuration(self, values: dict):\n wandb.config(values)\n\n @on_main_process\n def log(self, values: dict, step: Optional[int] = None):\n wandb.log(values, step=step)\n```\n\nWhen you are ready to build your `Accelerator` object, pass in an **instance** of your tracker to [`Accelerator.log_with`] to have it automatically\nbe used with the API:\n\n```python\ntracker = MyCustomTracker(\"some_run_name\")\naccelerator = Accelerator(log_with=tracker)\n```\n\nThese also can be mixed with existing trackers, including with `\"all\"`:\n\n```python\ntracker = MyCustomTracker(\"some_run_name\")\naccelerator = Accelerator(log_with=[tracker, \"all\"])\n```\n\n## Accessing the internal tracker \n\nIf some custom interactions with a tracker might be wanted directly, you can quickly access one using the \n[`Accelerator.get_tracker`] method. Just pass in the string corresponding to a tracker's `.name` attribute \nand it will return that tracker on the main process.\n\nThis example shows doing so with wandb:\n\n```python\nwandb_tracker = accelerator.get_tracker(\"wandb\")\n```\n\nFrom there you can interact with `wandb`'s `run` object like normal:\n\n```python\nwandb_tracker.log_artifact(some_artifact_to_log)\n```\n\n\n Trackers built in Accelerate will automatically execute on the correct process, \n so if a tracker is only meant to be ran on the main process it will do so \n automatically.\n\n\nIf you want to truly remove Accelerate's wrapping entirely, you can \nachieve the same outcome with:\n\n```python\nwandb_tracker = accelerator.get_tracker(\"wandb\", unwrap=True)\nif accelerator.is_main_process:\n wandb_tracker.log_artifact(some_artifact_to_log)\n```\n\n\n## When a wrapper cannot work\n\nIf 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:\n```diff\n from accelerate import Accelerator\n+ import neptune\n\n accelerator = Accelerator()\n+ run = neptune.init_run(...)\n\n my_model, my_optimizer, my_training_dataloader = accelerate.prepare(my_model, my_optimizer, my_training_dataloader)\n device = accelerator.device\n my_model.to(device)\n\n for iteration in config[\"num_iterations\"]:\n for batch in my_training_dataloader:\n my_optimizer.zero_grad()\n inputs, targets = batch\n inputs = inputs.to(device)\n targets = targets.to(device)\n outputs = my_model(inputs)\n loss = my_loss_function(outputs, targets)\n total_loss += loss\n accelerator.backward(loss)\n my_optimizer.step()\n+ if accelerator.is_main_process:\n+ run[\"logs/training/batch/loss\"].log(loss)\n```\n" }, { "path": "docs/source/usage_guides/training_zoo.md", "content": "\n\n# Example Zoo\n\nBelow contains a non-exhaustive list of tutorials and scripts showcasing Accelerate.\n\n## Official Accelerate Examples:\n\n### Basic Examples\n\nThese examples showcase the base features of Accelerate and are a great starting point\n\n- [Barebones NLP example](https://github.com/huggingface/accelerate/blob/main/examples/nlp_example.py)\n- [Barebones distributed NLP example in a Jupyter Notebook](https://github.com/huggingface/notebooks/blob/main/examples/accelerate_examples/simple_nlp_example.ipynb)\n- [Barebones computer vision example](https://github.com/huggingface/accelerate/blob/main/examples/cv_example.py)\n- [Barebones distributed computer vision example in a Jupyter Notebook](https://github.com/huggingface/notebooks/blob/main/examples/accelerate_examples/simple_cv_example.ipynb)\n- [Using Accelerate in Kaggle](https://www.kaggle.com/code/muellerzr/multi-gpu-and-accelerate)\n\n### Feature Specific Examples\n\nThese examples showcase specific features that the Accelerate framework offers\n\n- [Automatic memory-aware gradient accumulation](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/automatic_gradient_accumulation.py)\n- [Checkpointing states](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/checkpointing.py)\n- [Cross validation](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/cross_validation.py)\n- [DeepSpeed](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/deepspeed_with_config_support.py)\n- [Fully Sharded Data Parallelism](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/fsdp_with_peak_mem_tracking.py)\n- [Gradient accumulation](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/gradient_accumulation.py)\n- [Memory-aware batch size finder](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/memory.py)\n- [Metric Computation](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/multi_process_metrics.py)\n- [Using Trackers](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/tracking.py)\n- [Using Megatron-LM](https://github.com/huggingface/accelerate/blob/main/examples/by_feature/megatron_lm_gpt_pretraining.py)\n\n### Full Examples \n\nThese examples showcase every feature in Accelerate at once that was shown in \"Feature Specific Examples\"\n\n- [Complete NLP example](https://github.com/huggingface/accelerate/blob/main/examples/complete_nlp_example.py)\n- [Complete computer vision example](https://github.com/huggingface/accelerate/blob/main/examples/complete_cv_example.py)\n- [Very complete and extensible vision example showcasing SLURM, hydra, and a very extensible usage of the framework](https://github.com/yuvalkirstain/PickScore)\n- [Causal language model fine-tuning example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/language-modeling/run_clm_no_trainer.py)\n- [Masked language model fine-tuning example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/language-modeling/run_mlm_no_trainer.py)\n- [Speech pretraining example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/speech-pretraining/run_wav2vec2_pretraining_no_trainer.py)\n- [Translation fine-tuning example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/translation/run_translation_no_trainer.py)\n- [Text classification fine-tuning example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/text-classification/run_glue_no_trainer.py)\n- [Semantic segmentation fine-tuning example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/semantic-segmentation/run_semantic_segmentation_no_trainer.py)\n- [Question answering fine-tuning example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/question-answering/run_qa_no_trainer.py)\n- [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)\n- [Multiple choice question answering fine-tuning example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/multiple-choice/run_swag_no_trainer.py)\n- [Named entity recognition fine-tuning example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/token-classification/run_ner_no_trainer.py)\n- [Image classification fine-tuning example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/image-classification/run_image_classification_no_trainer.py)\n- [Summarization fine-tuning example](https://github.com/huggingface/transformers/blob/main/examples/pytorch/summarization/run_summarization_no_trainer.py)\n- [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)\n- [Megatron-LM examples for various NLp tasks](https://github.com/pacman100/accelerate-megatron-test) \n\n## Integration Examples \n\nThese are tutorials from libraries that integrate with Accelerate: \n\n> Don't find your integration here? Make a PR to include it!\n\n### Amphion\n- [Training Text-to-Speech Models with Amphion](https://github.com/open-mmlab/Amphion/blob/main/egs/tts/README.md)\n- [Training Singing Voice Conversion Models with Amphion](https://github.com/open-mmlab/Amphion/blob/main/egs/svc/README.md)\n- [Training Vocoders with Amphion](https://github.com/open-mmlab/Amphion/blob/main/egs/vocoder/README.md)\n\n### Catalyst\n\n- [Distributed training tutorial with Catalyst](https://catalyst-team.github.io/catalyst/tutorials/ddp.html)\n\n### DALLE2-pytorch \n\n- [Fine-tuning DALLE2](https://github.com/lucidrains/DALLE2-pytorch#usage)\n\n### Diffusers\n\n- [Performing textual inversion with diffusers](https://github.com/huggingface/diffusers/tree/main/examples/textual_inversion)\n- [Training DreamBooth with diffusers](https://github.com/huggingface/diffusers/tree/main/examples/dreambooth)\n\n### fastai \n\n- [Distributed training from Jupyter Notebooks with fastai](https://docs.fast.ai/tutorial.distributed.html)\n- [Basic distributed training examples with fastai](https://docs.fast.ai/examples/distributed_app_examples.html)\n\n### GradsFlow\n\n- [Auto Image Classification with GradsFlow](https://docs.gradsflow.com/en/latest/examples/nbs/01-ImageClassification/)\n\n### imagen-pytorch \n\n- [Fine-tuning Imagen](https://github.com/lucidrains/imagen-pytorch#usage)\n\n### Kornia\n\n- [Fine-tuning vision models with Kornia's Trainer](https://kornia.readthedocs.io/en/latest/get-started/training.html)\n\n### PyTorch Accelerated \n\n- [Quickstart distributed training tutorial with PyTorch Accelerated](https://pytorch-accelerated.readthedocs.io/en/latest/quickstart.html)\n\n### PyTorch3D\n\n- [Perform Deep Learning with 3D data](https://pytorch3d.org/tutorials/)\n\n### Stable-Dreamfusion\n\n- [Training with Stable-Dreamfusion to convert text to a 3D model](https://colab.research.google.com/drive/1MXT3yfOFvO0ooKEfiUUvTKwUkrrlCHpF?usp=sharing)\n\n### Tez \n\n- [Leaf disease detection with Tez and Accelerate](https://www.kaggle.com/code/abhishek/tez-faster-and-easier-training-for-leaf-detection/notebook)\n\n### trlx \n\n- [How to implement a sentiment learning task with trlx](https://github.com/CarperAI/trlx#example-how-to-add-a-task)\n\n### Comfy-UI\n\n- [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)\n\n\n## In Science\n\nBelow contains a non-exhaustive list of papers utilizing Accelerate. \n\n> Don't find your paper here? Make a PR to include it!\n\n* 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://huggingface.co/papers/2305.01569).\n* 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://huggingface.co/papers/2305.04091).\n* 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://huggingface.co/papers/2205.08343).\n* 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://huggingface.co/papers/2303.06865).\n* Peter Melchior, Yan Liang, ChangHoon Hahn, Andy Goulding: ā€œAutoencoding Galaxy Spectra I: Architectureā€, 2022; [arXiv:2211.07890](http://huggingface.co/papers/2211.07890).\n* 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://huggingface.co/papers/2304.04746).\n* Ayaan Haque, Matthew Tancik, Alexei A. Efros, Aleksander Holynski, Angjoo Kanazawa: ā€œInstruct-NeRF2NeRF: Editing 3D Scenes with Instructionsā€, 2023; [arXiv:2303.12789](http://huggingface.co/papers/2303.12789).\n* Luke Melas-Kyriazi, Christian Rupprecht, Iro Laina, Andrea Vedaldi: ā€œRealFusion: 360Ā° Reconstruction of Any Object from a Single Imageā€, 2023; [arXiv:2302.10663](http://huggingface.co/papers/2302.10663).\n* Xiaoshi Wu, Keqiang Sun, Feng Zhu, Rui Zhao, Hongsheng Li: ā€œBetter Aligning Text-to-Image Models with Human Preferenceā€, 2023; [arXiv:2303.14420](http://huggingface.co/papers/2303.14420).\n* 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://huggingface.co/papers/2303.17580).\n* 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://huggingface.co/papers/2210.12261).\n* Sheng-Yen Chou, Pin-Yu Chen, Tsung-Yi Ho: ā€œHow to Backdoor Diffusion Models?ā€, 2022; [arXiv:2212.05400](http://huggingface.co/papers/2212.05400).\n* 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://huggingface.co/papers/2303.07937).\n* 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://huggingface.co/papers/2303.11306).\n* DĆ­dac SurĆ­s, Sachit Menon, Carl Vondrick: ā€œViperGPT: Visual Inference via Python Execution for Reasoningā€, 2023; [arXiv:2303.08128](http://huggingface.co/papers/2303.08128).\n* 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://huggingface.co/papers/2303.09535).\n* Sean Welleck, Jiacheng Liu, Ximing Lu, Hannaneh Hajishirzi, Yejin Choi: ā€œNaturalProver: Grounded Mathematical Proof Generation with Language Modelsā€, 2022; [arXiv:2205.12910](http://huggingface.co/papers/2205.12910).\n* Elad Richardson, Gal Metzer, Yuval Alaluf, Raja Giryes, Daniel Cohen-Or: ā€œTEXTure: Text-Guided Texturing of 3D Shapesā€, 2023; [arXiv:2302.01721](http://huggingface.co/papers/2302.01721).\n* 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://huggingface.co/papers/2303.04603).\n* Shun Shao, Yftah Ziser, Shay Cohen: ā€œErasure of Unaligned Attributes from Neural Representationsā€, 2023; [arXiv:2302.02997](http://huggingface.co/papers/2302.02997).\n* Seonghyeon Ye, Hyeonbin Hwang, Sohee Yang, Hyeongu Yun, Yireun Kim, Minjoon Seo: ā€œIn-Context Instruction Learningā€, 2023; [arXiv:2302.14691](http://huggingface.co/papers/2302.14691).\n* 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://huggingface.co/papers/2303.02506).\n* Haoyu Chen, Zhihua Wang, Yang Yang, Qilin Sun, Kede Ma: ā€œLearning a Deep Color Difference Metric for Photographic Imagesā€, 2023; [arXiv:2303.14964](http://huggingface.co/papers/2303.14964).\n* Van-Hoang Le, Hongyu Zhang: ā€œLog Parsing with Prompt-based Few-shot Learningā€, 2023; [arXiv:2302.07435](http://huggingface.co/papers/2302.07435).\n* 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://huggingface.co/papers/2302.07866).\n* Ruoyao Wang, Peter Jansen, Marc-Alexandre CĆ“tĆ©, Prithviraj Ammanabrolu: ā€œBehavior Cloned Transformers are Neurosymbolic Reasonersā€, 2022; [arXiv:2210.07382](http://huggingface.co/papers/2210.07382).\n* 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://huggingface.co/papers/2304.13148). DOI: [https://dx.doi.org/10.1145/3539618.3591882 10.1145/3539618.3591882].\n* 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://huggingface.co/papers/2301.13826).\n* Marcio Fonseca, Yftah Ziser, Shay B. Cohen: ā€œFactorizing Content and Budget Decisions in Abstractive Summarization of Long Documentsā€, 2022; [arXiv:2205.12486](http://huggingface.co/papers/2205.12486).\n* Elad Richardson, Gal Metzer, Yuval Alaluf, Raja Giryes, Daniel Cohen-Or: ā€œTEXTure: Text-Guided Texturing of 3D Shapesā€, 2023; [arXiv:2302.01721](http://huggingface.co/papers/2302.01721).\n* 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://huggingface.co/papers/2212.10020).\n* 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://huggingface.co/papers/2302.00083).\n* 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://huggingface.co/papers/2211.01452).\n* 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://huggingface.co/papers/2206.11309).\n* 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://huggingface.co/papers/2304.14241).\n* 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://huggingface.co/papers/2206.11871).\n* Zhiruo Wang, Shuyan Zhou, Daniel Fried, Graham Neubig: ā€œExecution-Based Evaluation for Open-Domain Code Generationā€, 2022; [arXiv:2212.10481](http://huggingface.co/papers/2212.10481).\n* 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://huggingface.co/papers/2212.04875).\n* Jun Hao Liew, Hanshu Yan, Daquan Zhou, Jiashi Feng: ā€œMagicMix: Semantic Mixing with Diffusion Modelsā€, 2022; [arXiv:2210.16056](http://huggingface.co/papers/2210.16056).\n* 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://huggingface.co/papers/2110.06274).\n" }, { "path": "examples/README.md", "content": "\n\n# In this folder we showcase various full examples using šŸ¤— Accelerate\n\n## Simple NLP example\n\nThe [nlp_example.py](./nlp_example.py) script is a simple example to train a Bert model on a classification task ([GLUE's MRPC](https://www.microsoft.com/en-us/download/details.aspx?id=52398)).\n\nPrior to running it you should install šŸ¤— Dataset and šŸ¤— Transformers:\n\n```bash\npip install datasets evaluate transformers\n```\n\nThe same script can be run in any of the following configurations:\n- single CPU or single GPU\n- multi CPUs\n- multi GPUs (using PyTorch distributed mode)\n- (multi) TPUs\n- fp16 (mixed-precision) or fp32 (normal precision)\n\nTo run it in each of these various modes, use the following commands:\n- single CPU:\n * from a server without GPU\n ```bash\n python ./nlp_example.py\n ```\n * from any server by passing `cpu=True` to the `Accelerator`.\n ```bash\n python ./nlp_example.py --cpu\n ```\n * from any server with Accelerate launcher\n ```bash\n accelerate launch --cpu ./nlp_example.py\n ```\n- single GPU:\n ```bash\n python ./nlp_example.py # from a server with a GPU\n ```\n- with fp16 (mixed-precision)\n * from any server by passing `mixed_precison=fp16` to the `Accelerator`.\n ```bash\n python ./nlp_example.py --mixed_precision fp16\n ```\n * from any server with Accelerate launcher\n ```bash\n accelerate launch --mixed_precision fp16 ./nlp_example.py\n- multi CPUs (requires Open MPI, Intel MPI, or MVAPICH)\n * With Accelerate config and launcher, execute the following from node 0:\n ```bash\n accelerate config # Select to have accelerate launch mpirun\n accelerate launch ./nlp_example.py # This will run the script on each server\n ```\n * With Intel MPI:\n ```bash\n export CCL_WORKER_COUNT=1\n export MASTER_ADDR=xxx.xxx.xxx.xxx #node0 ip\n mpirun -f hostfile -n 16 -ppn 4 python ./nlp_example.py\n ```\n- multi GPUs (using PyTorch distributed mode)\n * With Accelerate config and launcher\n ```bash\n accelerate config # This will create a config file on your server\n accelerate launch ./nlp_example.py # This will run the script on your server\n ```\n * With traditional PyTorch launcher (`python -m torch.distributed.run` can be used instead of `torchrun`)\n ```bash\n torchrun --nproc_per_node 2 ./nlp_example.py\n ```\n- multi GPUs, multi node (several machines, using PyTorch distributed mode)\n * With Accelerate config and launcher, on each machine:\n ```bash\n accelerate config # This will create a config file on each server\n accelerate launch ./nlp_example.py # This will run the script on each server\n ```\n * With PyTorch launcher only (`python -m torch.distributed.run` can be used instead of `torchrun`). Run this command on each node:\n ```bash\n torchrun \\ # python -m torch.distributed.run \n --nproc_per_node 2 \\\n --nnodes 2 \\\n --rdzv_id 2299 \\ # A unique job id \n --rdzv_backend c10d \\\n --rdzv_endpoint master_node_ip_address:29500 \\\n ./nlp_example.py\n ```\n- (multi) TPUs\n * With Accelerate config and launcher\n ```bash\n accelerate config # This will create a config file on your TPU server\n accelerate launch ./nlp_example.py # This will run the script on each server\n ```\n * In PyTorch:\n Add an `xmp.spawn` line in your script as you usually do.\n\n\n## Simple vision example\n\nThe [cv_example.py](./cv_example.py) script is a simple example to fine-tune a ResNet-50 on a classification task ([Oxford-IIT Pet Dataset](https://www.robots.ox.ac.uk/~vgg/data/pets/)).\n\nThe same script can be run in any of the following configurations:\n- single CPU or single GPU\n- multi CPUs\n- multi GPUs (using PyTorch distributed mode)\n- (multi) TPUs\n- fp16 (mixed-precision) or fp32 (normal precision)\n\nPrior to running it you should install timm and torchvision:\n\n```bash\npip install timm torchvision\n```\n\nand you should download the data with the following commands:\n\n```bash\nwget https://www.robots.ox.ac.uk/~vgg/data/pets/data/images.tar.gz\ntar -xzf images.tar.gz\n```\n\nTo run it in each of these various modes, use the following commands:\n- single CPU:\n * from a server without GPU\n ```bash\n python ./cv_example.py --data_dir path_to_data\n ```\n * from any server by passing `cpu=True` to the `Accelerator`.\n ```bash\n python ./cv_example.py --data_dir path_to_data --cpu\n ```\n * from any server with Accelerate launcher\n ```bash\n accelerate launch --cpu ./cv_example.py --data_dir path_to_data\n ```\n- single GPU:\n ```bash\n python ./cv_example.py # from a server with a GPU\n ```\n- with fp16 (mixed-precision)\n * from any server by passing `mixed_precison=fp16` to the `Accelerator`.\n ```bash\n python ./cv_example.py --data_dir path_to_data --mixed_precison fp16\n ```\n * from any server with Accelerate launcher\n ```bash\n accelerate launch --mixed_precison fp16 ./cv_example.py --data_dir path_to_data\n- multi CPUs (requires Open MPI, Intel MPI, or MVAPICH)\n * With Accelerate config and launcher, run the following from node 0:\n ```bash\n accelerate config --config_file config.yaml # Select to have accelerate launch mpirun\n accelerate launch ./cv_example.py --data_dir path_to_data # This will run the script on each server\n ```\n * With Intel MPI, execute mpirun from node 0:\n ```bash\n export CCL_WORKER_COUNT=1\n export MASTER_ADDR=xxx.xxx.xxx.xxx #node0 ip\n mpirun -f hostfile -n 16 -ppn 4 python ./cv_example.py --data_dir path_to_data\n ```\n- multi GPUs (using PyTorch distributed mode)\n * With Accelerate config and launcher\n ```bash\n accelerate config --config_file config.yaml # This will create a config file on your server to `config.yaml`\n accelerate launch --config_file config.yaml ./cv_example.py --data_dir path_to_data # This will run the script on your server\n ```\n * With traditional PyTorch launcher (`python -m torch.distributed.run` can be used instead of `torchrun`)\n ```bash\n torchrun --nproc_per_node 2 ./cv_example.py --data_dir path_to_data\n ```\n- multi GPUs, multi node (several machines, using PyTorch distributed mode)\n * With Accelerate config and launcher, on each machine:\n ```bash\n accelerate config --config_file config.yaml # This will create a config file on your server to `config.yaml`\n accelerate launch --config_file config.yaml ./cv_example.py --data_dir path_to_data # This will run the script on each server\n ```\n * With PyTorch launcher only (`python -m torch.distributed.run` can be used instead of `torchrun`). Run this command on each node:\n ```bash\n torchrun \\ # python -m torch.distributed.run\n --nproc_per_node 2 \\\n --nnodes 2 \\\n --rdzv_id 2299 \\ # A unique job id \n --rdzv_backend c10d \\\n --rdzv_endpoint master_node_ip_address:29500 \\\n ./cv_example.py --data_dir path_to_data\n ```\n- (multi) TPUs\n * With Accelerate config and launcher\n ```bash\n accelerate config --config_file config.yaml # This will create a config file on your server to `config.yaml`\n accelerate launch --config_file config.yaml ./cv_example.py --data_dir path_to_data # This will run the script on each server\n ```\n * In PyTorch:\n Add an `xmp.spawn` line in your script as you usually do.\n\n### Simple vision example (GANs)\n\n- [huggan project](https://github.com/huggingface/community-events/tree/main/huggan)\n\n\n### Using AWS SageMaker integration\n- [Examples showcasing AWS SageMaker integration of šŸ¤— Accelerate.](https://github.com/pacman100/accelerate-aws-sagemaker)\n\n## Configuration zoo\nIn [/config_yaml_templates](./config_yaml_templates/) we have a variety of *minimal* `config.yaml` templates and examples to help you learn\nhow to create your own configuration files depending on the scenario. \n\n## SLURM Scripts \nIn [/slurm/submit_multigpu.sh](./slurm/submit_multigpu.sh) and [/slurm/submit_multinode.sh](./slurm/submit_multinode.sh) we present two scripts for running the examples on a machine with [SLURM](https://slurm.schedmd.com/documentation.html) workload manager. \n\nIn [/slurm/submit_multigpu.sh](./slurm/submit_multigpu.sh) the only parameter in the launcher that needs to be modified is `--num_processes`, which determines the number of GPUs we will use. In this case, using the environment variable `$SLURM_GPUS`, we indicate that we want to utilize all the GPUs available on the node we have requested. \n\nIn [/slurm/submit_multinode.sh](./slurm/submit_multinode.sh) we must specify the number of nodes that will be part of the training (`--num_machines`), how many GPUs we will use in total (`--num_processes`), the [`backend`](https://pytorch.org/docs/stable/elastic/run.html#note-on-rendezvous-backend), `--main_process_ip` which will be the address the master node and the `--main_process_port`.\n\nIn [/slurm/submit_multicpu.sh](./slurm/submit_multicpu.sh) we must specify the number of nodes that will be part of the training (`--num_machines`), how many CPU processes we will use in total (`--num_processes`), the [`backend`](https://pytorch.org/docs/stable/elastic/run.html#note-on-rendezvous-backend), `--main_process_ip` which will be the address the master node and the `--main_process_port`. `mpirun_hostfile` specifies to run the job using MPIRun.\n\nIn both scripts, we run `activateEnvironment.sh` at the beginning. This script should contain the necessary instructions to initialize the environment for execution. Below, we show an example that loads the necessary libraries ([Environment modules](https://github.com/cea-hpc/modules)), activates the Python environment, and sets up various environment variables, most of them to run the scripts in offline mode in case we don't have internet connection from the cluster.\n\n```bash\n# activateEnvironment.sh \nmodule purge\nmodule load anaconda3/2020.02 cuda/10.2 cudnn/8.0.5 nccl/2.9.9 arrow/7.0.0 openmpi\nsource activate /home/nct01/nct01328/pytorch_antoni_local\n\nexport HF_HOME=/gpfs/projects/nct01/nct01328/\nexport HF_LOCAL_HOME=/gpfs/projects/nct01/nct01328/HF_LOCAL\nexport HF_DATASETS_OFFLINE=1\nexport TRANSFORMERS_OFFLINE=1\nexport PYTHONPATH=/home/nct01/nct01328/transformers-in-supercomputers:$PYTHONPATH \nexport GPUS_PER_NODE=4\n```\n\n## Simple Multi-GPU Hardware Launcher (using an external platform)\n\n[multigpu_remote_launcher.py](./multigpu_remote_launcher.py) is a minimal script that demonstrates launching accelerate\non multiple remote GPUs, and with automatic hardware environment and dependency setup for reproducibility. You can\neasily customize the training function used, training arguments, hyperparameters, and type of compute hardware, and then\nrun the script to automatically launch multi GPU training on remote hardware.\n\nThis script uses [Runhouse](https://github.com/run-house/runhouse) to launch on self-hosted hardware (e.g. in your own\ncloud account or on-premise cluster) but there are other options for running remotely as well. Runhouse can be installed\nwith `pip install runhouse`, and you can refer to\n[hardware setup](https://runhouse-docs.readthedocs-hosted.com/en/latest/api/python/cluster.html#hardware-setup)\nfor hardware setup instructions, or this\n[Colab tutorial](https://colab.research.google.com/drive/1qVwYyLTCPYPSdz9ZX7BZl9Qm0A3j7RJe) for a more in-depth walkthrough.\n\n## Simple fine-tuning script that works on TPU\n\n[finetune_lm_tpu.py](./finetune_lm_tpu.py) is a classical language modeling generation fine tuning script that has been\nadapted to run best on TPUs. It has been successfully run and tested on a TPU v5 litepod-8, and it shows how it is\npossible to perform a fine-tuning task on such hardware thanks to accelerate and FSDPv2, using transformers and Torch XLA.\n\n## Finer Examples\n\nWhile 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.\n\n### `by_feature` examples\n\nThese scripts are *individual* examples highlighting one particular feature or use-case within Accelerate. They all stem from the [nlp_example.py](./nlp_example.py) script, and any changes or modifications is denoted with a `# New Code #` comment.\n\nRead the README.md file located in the `by_feature` folder for more information.\n\n### `complete_*` examples\n\nThese two scripts contain *every* single feature currently available in Accelerate in one place, as one giant script.\n\nNew arguments that can be passed include:\n\n- `checkpointing_steps`, whether the various states should be saved at the end of every `n` steps, or `\"epoch\"` for each epoch. States are then saved to folders named `step_{n}` or `epoch_{n}`\n- `resume_from_checkpoint`, should be used if you want to resume training off of a previous call to the script and passed a `checkpointing_steps` to it.\n- `with_tracking`, should be used if you want to log the training run using all available experiment trackers in your environment. Currently supported trackers include TensorBoard, Weights and Biases, and CometML.\n" }, { "path": "examples/alst_ulysses_sequence_parallelism/README.md", "content": "# Deepspeed's ALST/Ulysses sequence parallelism\n\nThis is an example of the use of Ulysses Sequence Parallelism, which uses attention head parallelism and is part of the Arctic Long Sequence Training project at [ArcticTraining](https://github.com/snowflakedb/ArcticTraining). [This paper](https://arxiv.org/abs/2506.13996) goes into the details of this protocol.\n\nFor nuances of usage please refer to the main HF Accelerate tutorial on [Context Parallelism](https://huggingface.co/docs/accelerate/en/concept_guides/context_parallelism).\n\nYou need to use at least `2` gpus to enable ALST/Ulysses sequence parallelism.\n\nTo run the example with `4` gpus:\n\n```bash\nbash ./sp-alst.sh\n```\n\nChange `4` to the desired sequence parallelism degree in these 2 files:\n```\nsp-alst.accelerate-config.yml:num_processes: 4\nsp-alst.py: sp_size=4,\n```\n" }, { "path": "examples/alst_ulysses_sequence_parallelism/sp-alst.accelerate-config.yml", "content": "compute_environment: LOCAL_MACHINE\ndeepspeed_config:\n deepspeed_config_file: sp-alst.ds-config.json\n zero3_init_flag: false\ndistributed_type: DEEPSPEED\nmachine_rank: 0\nmain_training_function: main\nnum_machines: 1\nnum_processes: 4\nrdzv_backend: static\nsame_network: true\nuse_cpu: false" }, { "path": "examples/alst_ulysses_sequence_parallelism/sp-alst.ds-config.json", "content": "{\n \"bf16\": {\n \"enabled\": true\n },\n \"zero_optimization\": {\n \"stage\": 3\n },\n \"gradient_accumulation_steps\": 1,\n \"train_batch_size\": \"auto\",\n \"train_micro_batch_size_per_gpu\": \"auto\",\n \"seq_parallel_communication_data_type\": \"bf16\"\n}" }, { "path": "examples/alst_ulysses_sequence_parallelism/sp-alst.py", "content": "# Copyright 2025 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport torch\nfrom deepspeed.runtime.utils import move_to_device\nfrom transformers import AutoModelForCausalLM, AutoTokenizer\n\nfrom accelerate import Accelerator\nfrom accelerate.utils import ParallelismConfig, set_seed\nfrom accelerate.utils.dataclasses import DeepSpeedSequenceParallelConfig\n\n\nset_seed(42)\n\nmodel_name = \"TinyLlama/TinyLlama-1.1B-Chat-v1.0\"\n# to run the example faster switch to the random model\n# model_name = \"hf-internal-testing/tiny-random-LlamaForCausalLM\"\n\nmicro_batch_size = 1\n\nparallelism_config = ParallelismConfig(\n sp_backend=\"deepspeed\",\n sp_size=4,\n sp_handler=DeepSpeedSequenceParallelConfig(\n sp_seq_length=256,\n sp_seq_length_is_variable=True,\n sp_attn_implementation=\"sdpa\",\n ),\n)\n\naccelerator = Accelerator(\n parallelism_config=parallelism_config,\n # log_with=\"wandb\", # enable to log into wandb\n)\naccelerator.init_trackers(\n project_name=\"ulysses-accelerate\",\n config={},\n init_kwargs={\"wandb\": dict(entity=\"yak\", name=\"deepspeed\")},\n)\n\ntokenizer = AutoTokenizer.from_pretrained(model_name)\nmodel = AutoModelForCausalLM.from_pretrained(model_name)\n\n# 2 quick rough datasets to demonstrate the workings\nif 1: # real dataset\n from datasets import load_dataset\n\n ds = load_dataset(\"HuggingFaceH4/ultrachat_200k\", split=\"train_sft[:12]\")\n\n # this is a quick example, it should be made more efficient to be used in real application\n def convert(ex):\n texts = tokenizer.apply_chat_template(conversation=ex[\"messages\"], tokenize=False)\n tokenized_dict = tokenizer(texts, max_length=256, padding=True, truncation=True)\n return tokenized_dict\n\n ds = ds.map(convert, batched=False, remove_columns=[\"prompt\", \"prompt_id\", \"messages\"])\n\n def collate_fn(batch):\n input_ids = torch.tensor(batch[0][\"input_ids\"]).unsqueeze(0)\n attention_mask = torch.tensor(batch[0][\"attention_mask\"]).unsqueeze(0)\n position_ids = torch.arange(input_ids.shape[1]).unsqueeze(0)\n return dict(\n input_ids=input_ids,\n position_ids=position_ids,\n labels=input_ids,\n attention_mask=attention_mask,\n )\n\n dl = torch.utils.data.DataLoader(\n ds, batch_size=micro_batch_size, collate_fn=collate_fn, drop_last=True, shuffle=False\n )\n\nelse: # fake dataset\n samples = 16\n seqlen = 256\n input_ids = torch.arange(1, seqlen * samples + 1).view(-1, seqlen) + 100\n position_ids = torch.arange(seqlen * samples).view(-1, seqlen)\n\n ds = torch.utils.data.TensorDataset(input_ids, position_ids)\n\n def collate_fn(batch):\n input_ids, position_ids = batch[0]\n return dict(\n input_ids=input_ids.unsqueeze(0),\n position_ids=position_ids.unsqueeze(0),\n labels=input_ids.unsqueeze(0),\n )\n\n dl = torch.utils.data.DataLoader(ds, batch_size=micro_batch_size, collate_fn=collate_fn)\n\noptimizer = torch.optim.Adam(model.parameters(), lr=1e-5)\n\nrank = torch.distributed.get_rank()\n\nif rank == 0:\n print(f\"DL orig: {len(dl)} samples\")\n\nmodel, optimizer, dl = accelerator.prepare(model, optimizer, dl)\n\nif rank == 0:\n print(f\"DL w/ adapter: {len(dl)} samples\")\n\nsp_size = parallelism_config.sp_size if parallelism_config else 1\nif sp_size > 1:\n sp_group = accelerator.torch_device_mesh[\"sp\"].get_group()\n sp_world_size = parallelism_config.sp_size\n\nunwrapped_model = accelerator.unwrap_model(model)\n\n# Normal training loop\nfor iter, batch in enumerate(dl):\n optimizer.zero_grad()\n\n if rank == 0:\n print(f\"batch {iter}: seqlen: {len(batch['input_ids'][0])}\")\n batch = move_to_device(batch, model.device)\n\n # The model automatically receives shift_labels via **kwargs and uses it for loss computation.\n # Both standard transformer models and Liger-patched models handle this correctly.\n outputs = model(**batch)\n loss = outputs.loss\n\n if sp_size > 1:\n # differentiable weighted per-shard-loss aggregation across ranks\n losses_per_rank = torch.distributed.nn.functional.all_gather(loss, group=sp_group)\n # special dealing with SFT that has prompt tokens that aren't used in loss computation\n good_tokens = (batch[\"shift_labels\"] != -100).view(-1).sum()\n good_tokens_per_rank = torch.distributed.nn.functional.all_gather(good_tokens, group=sp_group)\n total_loss = sum(losses_per_rank[rank] * good_tokens_per_rank[rank] for rank in range(sp_world_size))\n total_good_tokens = sum(good_tokens_per_rank)\n loss = total_loss / max(total_good_tokens, 1)\n\n if rank == 0:\n accelerator.print(f\"{iter}: {loss=}\")\n accelerator.log(dict(train_loss=loss, step=iter))\n\n accelerator.backward(loss)\n optimizer.step()\n\naccelerator.end_training()\n" }, { "path": "examples/alst_ulysses_sequence_parallelism/sp-alst.sh", "content": "export MASTER_ADDR=localhost\nexport MASTER_PORT=9998\npython -u -m accelerate.commands.launch \\\n --rdzv_conf \"rdzv_backend=c10d,rdzv_endpoint=$MASTER_ADDR:$MASTER_PORT\" \\\n --main_process_ip $MASTER_ADDR \\\n --main_process_port $MASTER_PORT \\\n --config_file sp-alst.accelerate-config.yml \\\n sp-alst.py\n" }, { "path": "examples/by_feature/README.md", "content": "# What are these scripts?\n\nAll scripts in this folder originate from the `nlp_example.py` file, as it is a very simplistic NLP training example using Accelerate with zero extra features.\n\nFrom there, each further script adds in just **one** feature of Accelerate, showing how you can quickly modify your own scripts to implement these capabilities.\n\nA full example with all of these parts integrated together can be found in the `complete_nlp_example.py` script and `complete_cv_example.py` script.\n\nAdjustments to each script from the base `nlp_example.py` file can be found quickly by searching for \"# New Code #\"\n\n## Example Scripts by Feature and their Arguments\n\n### Base Example (`../nlp_example.py`)\n\n- Shows how to use `Accelerator` in an extremely simplistic PyTorch training loop\n- Arguments available:\n - `mixed_precision`, whether to use mixed precision. (\"no\", \"fp16\", or \"bf16\")\n - `cpu`, whether to train using only the CPU. (yes/no/1/0)\n\nAll following scripts also accept these arguments in addition to their added ones.\n\nThese 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:\n\n```bash\naccelerate launch ../nlp_example.py --mixed_precision fp16 --cpu 0\n```\n\n### Checkpointing and Resuming Training (`checkpointing.py`)\n\n- Shows how to use `Accelerator.save_state` and `Accelerator.load_state` to save or continue training\n- **It is assumed you are continuing off the same training script**\n- Arguments available:\n - `checkpointing_steps`, after how many steps the various states should be saved. (\"epoch\", 1, 2, ...)\n - `output_dir`, where saved state folders should be saved to, default is current working directory\n - `resume_from_checkpoint`, what checkpoint folder to resume from. (\"epoch_0\", \"step_22\", ...)\n\nThese 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:\n\n(Note, `resume_from_checkpoint` assumes that we've ran the script for one epoch with the `--checkpointing_steps epoch` flag)\n\n```bash\naccelerate launch ./checkpointing.py --checkpointing_steps epoch output_dir \"checkpointing_tutorial\" --resume_from_checkpoint \"checkpointing_tutorial/epoch_0\"\n```\n\n### Cross Validation (`cross_validation.py`)\n\n- Shows how to use `Accelerator.free_memory` and run cross validation efficiently with `datasets`.\n- Arguments available:\n - `num_folds`, the number of folds the training dataset should be split into.\n\nThese 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:\n\n```bash\naccelerate launch ./cross_validation.py --num_folds 2\n```\n\n### Experiment Tracking (`tracking.py`)\n\n- Shows how to use `Accelerate.init_trackers` and `Accelerator.log`\n- Can be used with Weights and Biases, TensorBoard, or CometML.\n- Arguments available:\n - `with_tracking`, whether to load in all available experiment trackers from the environment.\n\nThese 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:\n\n```bash\naccelerate launch ./tracking.py --with_tracking\n```\n\n### Gradient Accumulation (`gradient_accumulation.py`)\n\n- Shows how to use `Accelerator.no_sync` to prevent gradient averaging in a distributed setup.\n- Arguments available:\n - `gradient_accumulation_steps`, the number of steps to perform before the gradients are accumulated and the optimizer and scheduler are stepped + zero_grad\n\nThese 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:\n\n```bash\naccelerate launch ./gradient_accumulation.py --gradient_accumulation_steps 5\n```\n\n### LocalSGD (`local_sgd.py`)\n- 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.\n\nThese 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:\n\n```bash\naccelerate launch ./local_sgd.py --local_sgd_steps 4\n```\n\n### DDP Communication Hook (`ddp_comm_hook.py`)\n\n- Shows how to use DDP Communication Hooks to control and optimize gradient communication across workers in a DistributedDataParallel setup.\n- Arguments available:\n - `ddp_comm_hook`, the type of DDP communication hook to use. Choose between `no`, `fp16`, `bf16`, `power_sgd`, and `batched_power_sgd`.\n\nThese arguments should be added at the end of any method for starting the python script (such as `accelerate launch`, `python -m torch.distributed.run`), such as:\n\n```bash\naccelerate launch ./ddp_comm_hook.py --mixed_precision fp16 --ddp_comm_hook power_sgd\n```\n\n### Profiler (`profiler.py`)\n\n- Shows how to use the profiling capabilities of `Accelerate` to profile PyTorch models during training.\n- Uses the `ProfileKwargs` handler to customize profiling options, including activities, scheduling, and additional profiling options.\n- Can generate and save profiling traces in JSON format for visualization in Chrome's tracing tool.\n\nArguments available:\n- `--record_shapes`: If passed, records shapes for profiling.\n- `--profile_memory`: If passed, profiles memory usage.\n- `--with_stack`: If passed, profiles stack traces.\n- `--with_flops`: If passed, profiles floating point operations (FLOPS).\n- `--output_trace_dir`: If specified, saves the profiling trace to the given dir in JSON format.\n- `--cpu`: If passed, trains on the CPU instead of GPU.\n\nThese 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:\n\n```bash\naccelerate launch ./profiler.py --record_shapes --profile_memory --with_flops --output_trace_dir \"profiler\"\n```\n" }, { "path": "examples/by_feature/automatic_gradient_accumulation.py", "content": "# Copyright 2022 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport argparse\nimport os\n\n# New Code #\nimport evaluate\nimport torch\nfrom datasets import load_dataset\nfrom torch.optim import AdamW\nfrom torch.utils.data import DataLoader\nfrom transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed\n\nfrom accelerate import Accelerator\nfrom accelerate.utils import find_executable_batch_size\n\n\n########################################################################\n# This is a fully working simple example to use Accelerate,\n# specifically showcasing how to combine both the gradient accumulation\n# and automatic batch size finder utilities of Accelerate to perfrom\n# automatic gradient accumulation\n#\n# This example trains a Bert base model on GLUE MRPC\n# in any of the following settings (with the same script):\n# - single CPU or single GPU\n# - multi GPUS (using PyTorch distributed mode)\n# - (multi) TPUs\n# - fp16 (mixed-precision) or fp32 (normal precision)\n#\n# New additions from the base script can be found quickly by\n# looking for the # New Code # tags\n#\n# To run it in each of these various modes, follow the instructions\n# in the readme for examples:\n# https://github.com/huggingface/accelerate/tree/main/examples\n#\n########################################################################\n\nEVAL_BATCH_SIZE = 32\n\n\ndef get_dataloaders(accelerator: Accelerator, batch_size: int = 16):\n \"\"\"\n Creates a set of `DataLoader`s for the `glue` dataset,\n using \"bert-base-cased\" as the tokenizer.\n\n Args:\n accelerator (`Accelerator`):\n An `Accelerator` object\n batch_size (`int`, *optional*):\n The batch size for the train and validation DataLoaders.\n \"\"\"\n tokenizer = AutoTokenizer.from_pretrained(\"bert-base-cased\")\n datasets = load_dataset(\"glue\", \"mrpc\")\n\n def tokenize_function(examples):\n # max_length=None => use the model max length (it's actually the default)\n outputs = tokenizer(examples[\"sentence1\"], examples[\"sentence2\"], truncation=True, max_length=None)\n return outputs\n\n # Apply the method we just defined to all the examples in all the splits of the dataset\n # starting with the main process first:\n with accelerator.main_process_first():\n tokenized_datasets = datasets.map(\n tokenize_function,\n batched=True,\n remove_columns=[\"idx\", \"sentence1\", \"sentence2\"],\n )\n\n # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the\n # transformers library\n tokenized_datasets = tokenized_datasets.rename_column(\"label\", \"labels\")\n\n def collate_fn(examples):\n # When using mixed precision we want round multiples of 8/16\n if accelerator.mixed_precision == \"fp8\":\n pad_to_multiple_of = 16\n elif accelerator.mixed_precision != \"no\":\n pad_to_multiple_of = 8\n else:\n pad_to_multiple_of = None\n\n return tokenizer.pad(\n examples,\n padding=\"longest\",\n pad_to_multiple_of=pad_to_multiple_of,\n return_tensors=\"pt\",\n )\n\n # Instantiate dataloaders.\n train_dataloader = DataLoader(\n tokenized_datasets[\"train\"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size\n )\n eval_dataloader = DataLoader(\n tokenized_datasets[\"validation\"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE\n )\n\n return train_dataloader, eval_dataloader\n\n\n# For testing only\nif os.environ.get(\"TESTING_MOCKED_DATALOADERS\", None) == \"1\":\n from accelerate.test_utils.training import mocked_dataloaders\n\n get_dataloaders = mocked_dataloaders # noqa: F811\n\n\ndef training_function(config, args):\n # For testing only\n if os.environ.get(\"TESTING_MOCKED_DATALOADERS\", None) == \"1\":\n config[\"num_epochs\"] = 2\n # Initialize accelerator\n accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)\n # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs\n lr = config[\"lr\"]\n num_epochs = int(config[\"num_epochs\"])\n seed = int(config[\"seed\"])\n observed_batch_size = int(config[\"batch_size\"])\n\n metric = evaluate.load(\"glue\", \"mrpc\")\n\n # New Code #\n # We use the `find_executable_batch_size` decorator, passing in the desired observed batch size\n # to train on. If a device OOM error occurs, it will retry this loop cutting the batch size in\n # half each time. From this, we can calculate the number of gradient accumulation steps needed\n # and modify the Accelerator object as a result\n @find_executable_batch_size(starting_batch_size=int(observed_batch_size))\n def inner_training_loop(batch_size):\n # Since we need to modify the outside accelerator object, we need to bring it\n # to the local scope\n nonlocal accelerator\n\n # We can calculate the number of gradient accumulation steps based on the current\n # batch size vs the starting batch size\n num_gradient_accumulation_steps = observed_batch_size // batch_size\n\n # And then set it in the Accelerator directly:\n accelerator.gradient_accumulation_steps = num_gradient_accumulation_steps\n\n # Next we need to free all of the stored model references in the Accelerator each time\n accelerator.free_memory()\n\n # And set the seed so our results are reproducable each reset\n set_seed(seed)\n\n # Instantiate the model (we build the model here so that the seed also control new weights initialization)\n model = AutoModelForSequenceClassification.from_pretrained(\"bert-base-cased\", return_dict=True)\n\n # We could avoid this line since the accelerator is set with `device_placement=True` (default value).\n # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer\n # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).\n model = model.to(accelerator.device)\n\n # Instantiate optimizer\n optimizer = AdamW(params=model.parameters(), lr=lr)\n train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size)\n\n # Instantiate scheduler\n lr_scheduler = get_linear_schedule_with_warmup(\n optimizer=optimizer,\n num_warmup_steps=100,\n num_training_steps=(len(train_dataloader) * num_epochs),\n )\n\n # Prepare everything\n # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the\n # prepare method.\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler\n )\n\n # Now we train the model\n for epoch in range(num_epochs):\n model.train()\n for step, batch in enumerate(train_dataloader):\n # And perform gradient accumulation\n with accelerator.accumulate(model):\n # We could avoid this line since we set the accelerator with `device_placement=True`.\n batch.to(accelerator.device)\n outputs = model(**batch)\n loss = outputs.loss\n accelerator.backward(loss)\n optimizer.step()\n lr_scheduler.step()\n optimizer.zero_grad()\n\n model.eval()\n for step, batch in enumerate(eval_dataloader):\n # We could avoid this line since we set the accelerator with `device_placement=True`.\n batch.to(accelerator.device)\n with torch.no_grad():\n outputs = model(**batch)\n predictions = outputs.logits.argmax(dim=-1)\n predictions, references = accelerator.gather_for_metrics((predictions, batch[\"labels\"]))\n metric.add_batch(\n predictions=predictions,\n references=references,\n )\n\n eval_metric = metric.compute()\n # Use accelerator.print to print only on the main process.\n accelerator.print(f\"epoch {epoch}:\", eval_metric)\n\n # New Code #\n # And call it at the end with no arguments\n # Note: You could also refactor this outside of your training loop function\n inner_training_loop()\n accelerator.end_training()\n\n\ndef main():\n parser = argparse.ArgumentParser(description=\"Simple example of training script.\")\n parser.add_argument(\n \"--mixed_precision\",\n type=str,\n default=None,\n choices=[\"no\", \"fp16\", \"bf16\", \"fp8\"],\n help=\"Whether to use mixed precision. Choose\"\n \"between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10.\"\n \"and an Nvidia Ampere GPU.\",\n )\n parser.add_argument(\"--cpu\", action=\"store_true\", help=\"If passed, will train on the CPU.\")\n args = parser.parse_args()\n # New Code #\n # We modify the starting batch size to be an observed batch size of 256, to guarantee an initial device OOM\n config = {\"lr\": 2e-5, \"num_epochs\": 3, \"seed\": 42, \"batch_size\": 256}\n training_function(config, args)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "examples/by_feature/checkpointing.py", "content": "# Copyright 2021 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport argparse\nimport os\n\nimport evaluate\nimport torch\nfrom datasets import load_dataset\nfrom torch.optim import AdamW\nfrom torch.utils.data import DataLoader\nfrom transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup\n\nfrom accelerate import Accelerator, DataLoaderConfiguration, DistributedType\nfrom accelerate.utils import set_seed\n\n\n########################################################################\n# This is a fully working simple example to use Accelerate,\n# specifically showcasing the checkpointing capability,\n# and builds off the `nlp_example.py` script.\n#\n# This example trains a Bert base model on GLUE MRPC\n# in any of the following settings (with the same script):\n# - single CPU or single GPU\n# - multi GPUS (using PyTorch distributed mode)\n# - (multi) TPUs\n# - fp16 (mixed-precision) or fp32 (normal precision)\n#\n# To help focus on the differences in the code, building `DataLoaders`\n# was refactored into its own function.\n# New additions from the base script can be found quickly by\n# looking for the # New Code # tags\n#\n# To run it in each of these various modes, follow the instructions\n# in the readme for examples:\n# https://github.com/huggingface/accelerate/tree/main/examples\n#\n########################################################################\n\nMAX_GPU_BATCH_SIZE = 16\nEVAL_BATCH_SIZE = 32\n\n\ndef get_dataloaders(accelerator: Accelerator, batch_size: int = 16):\n \"\"\"\n Creates a set of `DataLoader`s for the `glue` dataset,\n using \"bert-base-cased\" as the tokenizer.\n\n Args:\n accelerator (`Accelerator`):\n An `Accelerator` object\n batch_size (`int`, *optional*):\n The batch size for the train and validation DataLoaders.\n \"\"\"\n tokenizer = AutoTokenizer.from_pretrained(\"bert-base-cased\")\n datasets = load_dataset(\"glue\", \"mrpc\")\n\n def tokenize_function(examples):\n # max_length=None => use the model max length (it's actually the default)\n outputs = tokenizer(examples[\"sentence1\"], examples[\"sentence2\"], truncation=True, max_length=None)\n return outputs\n\n # Apply the method we just defined to all the examples in all the splits of the dataset\n # starting with the main process first:\n with accelerator.main_process_first():\n tokenized_datasets = datasets.map(\n tokenize_function,\n batched=True,\n remove_columns=[\"idx\", \"sentence1\", \"sentence2\"],\n )\n\n # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the\n # transformers library\n tokenized_datasets = tokenized_datasets.rename_column(\"label\", \"labels\")\n\n def collate_fn(examples):\n # On TPU it's best to pad everything to the same length or training will be very slow.\n max_length = 128 if accelerator.distributed_type == DistributedType.XLA else None\n # When using mixed precision we want round multiples of 8/16\n if accelerator.mixed_precision == \"fp8\":\n pad_to_multiple_of = 16\n elif accelerator.mixed_precision != \"no\":\n pad_to_multiple_of = 8\n else:\n pad_to_multiple_of = None\n\n return tokenizer.pad(\n examples,\n padding=\"longest\",\n max_length=max_length,\n pad_to_multiple_of=pad_to_multiple_of,\n return_tensors=\"pt\",\n )\n\n # Instantiate dataloaders.\n train_dataloader = DataLoader(\n tokenized_datasets[\"train\"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size\n )\n eval_dataloader = DataLoader(\n tokenized_datasets[\"validation\"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE\n )\n\n return train_dataloader, eval_dataloader\n\n\n# For testing only\nif os.environ.get(\"TESTING_MOCKED_DATALOADERS\", None) == \"1\":\n from accelerate.test_utils.training import mocked_dataloaders\n\n get_dataloaders = mocked_dataloaders # noqa: F811\n\n\ndef training_function(config, args):\n # For testing only\n if os.environ.get(\"TESTING_MOCKED_DATALOADERS\", None) == \"1\":\n config[\"num_epochs\"] = 2\n # Initialize accelerator\n dataloader_config = DataLoaderConfiguration(use_stateful_dataloader=args.use_stateful_dataloader)\n accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision, dataloader_config=dataloader_config)\n # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs\n lr = config[\"lr\"]\n num_epochs = int(config[\"num_epochs\"])\n seed = int(config[\"seed\"])\n batch_size = int(config[\"batch_size\"])\n\n # New Code #\n # Parse out whether we are saving every epoch or after a certain number of batches\n if hasattr(args.checkpointing_steps, \"isdigit\"):\n if args.checkpointing_steps == \"epoch\":\n checkpointing_steps = args.checkpointing_steps\n elif args.checkpointing_steps.isdigit():\n checkpointing_steps = int(args.checkpointing_steps)\n else:\n raise ValueError(\n f\"Argument `checkpointing_steps` must be either a number or `epoch`. `{args.checkpointing_steps}` passed.\"\n )\n else:\n checkpointing_steps = None\n\n set_seed(seed)\n\n train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size)\n metric = evaluate.load(\"glue\", \"mrpc\")\n\n # If the batch size is too big we use gradient accumulation\n gradient_accumulation_steps = 1\n if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.XLA:\n gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE\n batch_size = MAX_GPU_BATCH_SIZE\n\n # Instantiate the model (we build the model here so that the seed also control new weights initialization)\n model = AutoModelForSequenceClassification.from_pretrained(\"bert-base-cased\", return_dict=True)\n\n # We could avoid this line since the accelerator is set with `device_placement=True` (default value).\n # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer\n # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).\n model = model.to(accelerator.device)\n\n # Instantiate optimizer\n optimizer = AdamW(params=model.parameters(), lr=lr)\n\n # Instantiate scheduler\n lr_scheduler = get_linear_schedule_with_warmup(\n optimizer=optimizer,\n num_warmup_steps=100,\n num_training_steps=(len(train_dataloader) * num_epochs) // gradient_accumulation_steps,\n )\n\n # Prepare everything\n # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the\n # prepare method.\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler\n )\n\n # New Code #\n # We need to keep track of how many total steps we have iterated over\n overall_step = 0\n # We also need to keep track of the stating epoch so files are named properly\n starting_epoch = 0\n\n # We need to load the checkpoint back in before training here with `load_state`\n # The total number of epochs is adjusted based on where the state is being loaded from,\n # as we assume continuation of the same training script\n if args.resume_from_checkpoint:\n if args.resume_from_checkpoint is not None or args.resume_from_checkpoint != \"\":\n accelerator.print(f\"Resumed from checkpoint: {args.resume_from_checkpoint}\")\n accelerator.load_state(args.resume_from_checkpoint)\n path = os.path.basename(args.resume_from_checkpoint)\n else:\n # Get the most recent checkpoint\n dirs = [f.name for f in os.scandir(os.getcwd()) if f.is_dir()]\n dirs.sort(key=os.path.getctime)\n path = dirs[-1] # Sorts folders by date modified, most recent checkpoint is the last\n # Extract `epoch_{i}` or `step_{i}`\n training_difference = os.path.splitext(path)[0]\n\n if \"epoch\" in training_difference:\n starting_epoch = int(training_difference.replace(\"epoch_\", \"\")) + 1\n resume_step = None\n else:\n resume_step = int(training_difference.replace(\"step_\", \"\"))\n starting_epoch = resume_step // len(train_dataloader)\n resume_step -= starting_epoch * len(train_dataloader)\n\n # Now we train the model\n for epoch in range(starting_epoch, num_epochs):\n model.train()\n # New Code #\n if args.resume_from_checkpoint and epoch == starting_epoch and resume_step is not None:\n # We need to skip steps until we reach the resumed step only if we are not using a stateful dataloader\n if not args.use_stateful_dataloader:\n active_dataloader = accelerator.skip_first_batches(train_dataloader, resume_step)\n else:\n active_dataloader = train_dataloader\n overall_step += resume_step\n else:\n # After the first iteration though, we need to go back to the original dataloader\n active_dataloader = train_dataloader\n for step, batch in enumerate(active_dataloader):\n # We could avoid this line since we set the accelerator with `device_placement=True`.\n batch.to(accelerator.device)\n outputs = model(**batch)\n loss = outputs.loss\n loss = loss / gradient_accumulation_steps\n accelerator.backward(loss)\n if step % gradient_accumulation_steps == 0:\n optimizer.step()\n lr_scheduler.step()\n optimizer.zero_grad()\n # New Code #\n overall_step += 1\n\n # New Code #\n # We save the model, optimizer, lr_scheduler, and seed states by calling `save_state`\n # These are saved to folders named `step_{overall_step}`\n # Will contain files: \"pytorch_model.bin\", \"optimizer.bin\", \"scheduler.bin\", and \"random_states.pkl\"\n # If mixed precision was used, will also save a \"scalar.bin\" file\n if isinstance(checkpointing_steps, int):\n output_dir = f\"step_{overall_step}\"\n if overall_step % checkpointing_steps == 0:\n if args.output_dir is not None:\n output_dir = os.path.join(args.output_dir, output_dir)\n accelerator.save_state(output_dir)\n model.eval()\n for step, batch in enumerate(eval_dataloader):\n # We could avoid this line since we set the accelerator with `device_placement=True` (the default).\n batch.to(accelerator.device)\n with torch.no_grad():\n outputs = model(**batch)\n predictions = outputs.logits.argmax(dim=-1)\n predictions, references = accelerator.gather_for_metrics((predictions, batch[\"labels\"]))\n metric.add_batch(\n predictions=predictions,\n references=references,\n )\n eval_metric = metric.compute()\n # Use accelerator.print to print only on the main process.\n accelerator.print(f\"epoch {epoch}:\", eval_metric)\n\n # New Code #\n # We save the model, optimizer, lr_scheduler, and seed states by calling `save_state`\n # These are saved to folders named `epoch_{epoch}`\n # Will contain files: \"pytorch_model.bin\", \"optimizer.bin\", \"scheduler.bin\", and \"random_states.pkl\"\n # If mixed precision was used, will also save a \"scalar.bin\" file\n if checkpointing_steps == \"epoch\":\n output_dir = f\"epoch_{epoch}\"\n if args.output_dir is not None:\n output_dir = os.path.join(args.output_dir, output_dir)\n accelerator.save_state(output_dir)\n accelerator.end_training()\n\n\ndef main():\n parser = argparse.ArgumentParser(description=\"Simple example of training script.\")\n parser.add_argument(\n \"--mixed_precision\",\n type=str,\n default=None,\n choices=[\"no\", \"fp16\", \"bf16\", \"fp8\"],\n help=\"Whether to use mixed precision. Choose\"\n \"between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10.\"\n \"and an Nvidia Ampere GPU.\",\n )\n parser.add_argument(\"--cpu\", action=\"store_true\", help=\"If passed, will train on the CPU.\")\n parser.add_argument(\n \"--checkpointing_steps\",\n type=str,\n default=None,\n help=\"Whether the various states should be saved at the end of every n steps, or 'epoch' for each epoch.\",\n )\n parser.add_argument(\n \"--output_dir\",\n type=str,\n default=\".\",\n help=\"Optional save directory where all checkpoint folders will be stored. Default is the current working directory.\",\n )\n parser.add_argument(\n \"--resume_from_checkpoint\",\n type=str,\n default=None,\n help=\"If the training should continue from a checkpoint folder.\",\n )\n parser.add_argument(\n \"--use_stateful_dataloader\",\n action=\"store_true\",\n help=\"If the dataloader should be a resumable stateful dataloader.\",\n )\n args = parser.parse_args()\n config = {\"lr\": 2e-5, \"num_epochs\": 3, \"seed\": 42, \"batch_size\": 16}\n training_function(config, args)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "examples/by_feature/cross_validation.py", "content": "# Copyright 2022 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport argparse\n\nimport evaluate\nimport numpy as np\nimport torch\nfrom datasets import DatasetDict, load_dataset\n\n# New Code #\n# We'll be using StratifiedKFold for this example\nfrom sklearn.model_selection import StratifiedKFold\nfrom torch.optim import AdamW\nfrom torch.utils.data import DataLoader\nfrom transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed\n\nfrom accelerate import Accelerator, DistributedType\n\n\n########################################################################\n# This is a fully working simple example to use Accelerate,\n# specifically showcasing how to perform Cross Validation,\n# and builds off the `nlp_example.py` script.\n#\n# This example trains a Bert base model on GLUE MRPC\n# in any of the following settings (with the same script):\n# - single CPU or single GPU\n# - multi GPUS (using PyTorch distributed mode)\n# - (multi) TPUs\n# - fp16 (mixed-precision) or fp32 (normal precision)\n#\n# To help focus on the differences in the code, building `DataLoaders`\n# was refactored into its own function.\n# New additions from the base script can be found quickly by\n# looking for the # New Code # tags\n#\n# To run it in each of these various modes, follow the instructions\n# in the readme for examples:\n# https://github.com/huggingface/accelerate/tree/main/examples\n#\n########################################################################\n\n\nMAX_GPU_BATCH_SIZE = 16\nEVAL_BATCH_SIZE = 32\n\n# New Code #\n# We need a different `get_dataloaders` function that will build dataloaders by index\n\n\ndef get_fold_dataloaders(\n accelerator: Accelerator, dataset: DatasetDict, train_idxs: list[int], valid_idxs: list[int], batch_size: int = 16\n):\n \"\"\"\n Gets a set of train, valid, and test dataloaders for a particular fold\n\n Args:\n accelerator (`Accelerator`):\n The main `Accelerator` object\n train_idxs (list of `int`):\n The split indices for the training dataset\n valid_idxs (list of `int`):\n The split indices for the validation dataset\n batch_size (`int`):\n The size of the minibatch. Default is 16\n \"\"\"\n tokenizer = AutoTokenizer.from_pretrained(\"bert-base-cased\")\n datasets = DatasetDict(\n {\n \"train\": dataset[\"train\"].select(train_idxs),\n \"validation\": dataset[\"train\"].select(valid_idxs),\n \"test\": dataset[\"validation\"],\n }\n )\n\n def tokenize_function(examples):\n # max_length=None => use the model max length (it's actually the default)\n outputs = tokenizer(examples[\"sentence1\"], examples[\"sentence2\"], truncation=True, max_length=None)\n return outputs\n\n # Apply the method we just defined to all the examples in all the splits of the dataset\n # starting with the main process first:\n with accelerator.main_process_first():\n tokenized_datasets = datasets.map(\n tokenize_function,\n batched=True,\n remove_columns=[\"idx\", \"sentence1\", \"sentence2\"],\n )\n\n # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the\n # transformers library\n tokenized_datasets = tokenized_datasets.rename_column(\"label\", \"labels\")\n\n def collate_fn(examples):\n # On TPU it's best to pad everything to the same length or training will be very slow.\n max_length = 128 if accelerator.distributed_type == DistributedType.XLA else None\n # When using mixed precision we want round multiples of 8/16\n if accelerator.mixed_precision == \"fp8\":\n pad_to_multiple_of = 16\n elif accelerator.mixed_precision != \"no\":\n pad_to_multiple_of = 8\n else:\n pad_to_multiple_of = None\n\n return tokenizer.pad(\n examples,\n padding=\"longest\",\n max_length=max_length,\n pad_to_multiple_of=pad_to_multiple_of,\n return_tensors=\"pt\",\n )\n\n # Instantiate dataloaders.\n train_dataloader = DataLoader(\n tokenized_datasets[\"train\"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size\n )\n eval_dataloader = DataLoader(\n tokenized_datasets[\"validation\"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE\n )\n\n test_dataloader = DataLoader(\n tokenized_datasets[\"test\"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE\n )\n\n return train_dataloader, eval_dataloader, test_dataloader\n\n\ndef training_function(config, args):\n # New Code #\n test_predictions = []\n # Download the dataset\n datasets = load_dataset(\"glue\", \"mrpc\")\n # Create our splits\n kfold = StratifiedKFold(n_splits=int(args.num_folds))\n # Initialize accelerator\n accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)\n # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs\n lr = config[\"lr\"]\n num_epochs = int(config[\"num_epochs\"])\n seed = int(config[\"seed\"])\n batch_size = int(config[\"batch_size\"])\n\n metric = evaluate.load(\"glue\", \"mrpc\")\n\n # If the batch size is too big we use gradient accumulation\n gradient_accumulation_steps = 1\n if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.XLA:\n gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE\n batch_size = MAX_GPU_BATCH_SIZE\n\n set_seed(seed)\n\n # New Code #\n # Create our folds:\n folds = kfold.split(np.zeros(datasets[\"train\"].num_rows), datasets[\"train\"][\"label\"])\n test_references = []\n # Iterate over them\n for i, (train_idxs, valid_idxs) in enumerate(folds):\n train_dataloader, eval_dataloader, test_dataloader = get_fold_dataloaders(\n accelerator,\n datasets,\n train_idxs,\n valid_idxs,\n )\n # Instantiate the model (we build the model here so that the seed also control new weights initialization)\n model = AutoModelForSequenceClassification.from_pretrained(\"bert-base-cased\", return_dict=True)\n\n # We could avoid this line since the accelerator is set with `device_placement=True` (default value).\n # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer\n # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).\n model = model.to(accelerator.device)\n\n # Instantiate optimizer\n optimizer = AdamW(params=model.parameters(), lr=lr)\n\n # Instantiate scheduler\n lr_scheduler = get_linear_schedule_with_warmup(\n optimizer=optimizer,\n num_warmup_steps=100,\n num_training_steps=(len(train_dataloader) * num_epochs) // gradient_accumulation_steps,\n )\n\n # Prepare everything\n # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the\n # prepare method.\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler\n )\n\n # Now we train the model\n for epoch in range(num_epochs):\n model.train()\n for step, batch in enumerate(train_dataloader):\n # We could avoid this line since we set the accelerator with `device_placement=True`.\n batch.to(accelerator.device)\n outputs = model(**batch)\n loss = outputs.loss\n loss = loss / gradient_accumulation_steps\n accelerator.backward(loss)\n if step % gradient_accumulation_steps == 0:\n optimizer.step()\n lr_scheduler.step()\n optimizer.zero_grad()\n\n model.eval()\n for step, batch in enumerate(eval_dataloader):\n # We could avoid this line since we set the accelerator with `device_placement=True`.\n batch.to(accelerator.device)\n with torch.no_grad():\n outputs = model(**batch)\n predictions = outputs.logits.argmax(dim=-1)\n predictions, references = accelerator.gather_for_metrics((predictions, batch[\"labels\"]))\n metric.add_batch(\n predictions=predictions,\n references=references,\n )\n\n eval_metric = metric.compute()\n # Use accelerator.print to print only on the main process.\n accelerator.print(f\"epoch {epoch}:\", eval_metric)\n\n # New Code #\n # We also run predictions on the test set at the very end\n fold_predictions = []\n for step, batch in enumerate(test_dataloader):\n # We could avoid this line since we set the accelerator with `device_placement=True`.\n batch.to(accelerator.device)\n with torch.no_grad():\n outputs = model(**batch)\n predictions = outputs.logits\n predictions, references = accelerator.gather_for_metrics((predictions, batch[\"labels\"]))\n fold_predictions.append(predictions.cpu())\n if i == 0:\n # We need all of the test predictions\n test_references.append(references.cpu())\n # Use accelerator.print to print only on the main process.\n test_predictions.append(torch.cat(fold_predictions, dim=0))\n # We now need to release all our memory and get rid of the current model, optimizer, etc\n model, optimizer = accelerator.free_memory(model, optimizer)\n # New Code #\n # Finally we check the accuracy of our folded results:\n test_references = torch.cat(test_references, dim=0)\n preds = torch.stack(test_predictions, dim=0).sum(dim=0).div(int(args.num_folds)).argmax(dim=-1)\n test_metric = metric.compute(predictions=preds, references=test_references)\n accelerator.print(\"Average test metrics from all folds:\", test_metric)\n accelerator.end_training()\n\n\ndef main():\n parser = argparse.ArgumentParser(description=\"Simple example of training script.\")\n parser.add_argument(\n \"--mixed_precision\",\n type=str,\n default=None,\n choices=[\"no\", \"fp16\", \"bf16\", \"fp8\"],\n help=\"Whether to use mixed precision. Choose\"\n \"between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10.\"\n \"and an Nvidia Ampere GPU.\",\n )\n parser.add_argument(\"--cpu\", action=\"store_true\", help=\"If passed, will train on the CPU.\")\n # New Code #\n parser.add_argument(\"--num_folds\", type=int, default=3, help=\"The number of splits to perform across the dataset\")\n args = parser.parse_args()\n config = {\"lr\": 2e-5, \"num_epochs\": 3, \"seed\": 42, \"batch_size\": 16}\n training_function(config, args)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "examples/by_feature/ddp_comm_hook.py", "content": "# Copyright 2021 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport argparse\nimport os\n\nimport evaluate\nimport torch\nfrom datasets import load_dataset\nfrom torch.optim import AdamW\nfrom torch.utils.data import DataLoader\nfrom transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed\n\nfrom accelerate import Accelerator, DistributedType\nfrom accelerate.utils import DDPCommunicationHookType, DistributedDataParallelKwargs\n\n\n########################################################################\n# This is a fully working simple example to use Accelerate\n# and perform ddp communication hook\n#\n# This example trains a Bert base model on GLUE MRPC\n# in any of the following settings (with the same script):\n# - single CPU or single GPU\n# - multi GPUS (using PyTorch distributed mode)\n# - (multi) TPUs\n# - fp16 (mixed-precision) or fp32 (normal precision)\n#\n# To run it in each of these various modes, follow the instructions\n# in the readme for examples:\n# https://github.com/huggingface/accelerate/tree/main/examples\n#\n########################################################################\n\n\nMAX_GPU_BATCH_SIZE = 16\nEVAL_BATCH_SIZE = 32\n\n\ndef get_dataloaders(accelerator: Accelerator, batch_size: int = 16):\n \"\"\"\n Creates a set of `DataLoader`s for the `glue` dataset,\n using \"bert-base-cased\" as the tokenizer.\n\n Args:\n accelerator (`Accelerator`):\n An `Accelerator` object\n batch_size (`int`, *optional*):\n The batch size for the train and validation DataLoaders.\n \"\"\"\n tokenizer = AutoTokenizer.from_pretrained(\"bert-base-cased\")\n datasets = load_dataset(\"glue\", \"mrpc\")\n\n def tokenize_function(examples):\n # max_length=None => use the model max length (it's actually the default)\n outputs = tokenizer(examples[\"sentence1\"], examples[\"sentence2\"], truncation=True, max_length=None)\n return outputs\n\n # Apply the method we just defined to all the examples in all the splits of the dataset\n # starting with the main process first:\n with accelerator.main_process_first():\n tokenized_datasets = datasets.map(\n tokenize_function,\n batched=True,\n remove_columns=[\"idx\", \"sentence1\", \"sentence2\"],\n )\n\n # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the\n # transformers library\n tokenized_datasets = tokenized_datasets.rename_column(\"label\", \"labels\")\n\n def collate_fn(examples):\n # On TPU it's best to pad everything to the same length or training will be very slow.\n max_length = 128 if accelerator.distributed_type == DistributedType.XLA else None\n # When using mixed precision we want round multiples of 8/16\n if accelerator.mixed_precision == \"fp8\":\n pad_to_multiple_of = 16\n elif accelerator.mixed_precision != \"no\":\n pad_to_multiple_of = 8\n else:\n pad_to_multiple_of = None\n\n return tokenizer.pad(\n examples,\n padding=\"longest\",\n max_length=max_length,\n pad_to_multiple_of=pad_to_multiple_of,\n return_tensors=\"pt\",\n )\n\n # Instantiate dataloaders.\n train_dataloader = DataLoader(\n tokenized_datasets[\"train\"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size\n )\n eval_dataloader = DataLoader(\n tokenized_datasets[\"validation\"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE\n )\n\n return train_dataloader, eval_dataloader\n\n\n# For testing only\nif os.environ.get(\"TESTING_MOCKED_DATALOADERS\", None) == \"1\":\n from accelerate.test_utils.training import mocked_dataloaders\n\n get_dataloaders = mocked_dataloaders # noqa: F811\n\n\ndef training_function(config, args):\n # For testing only\n if os.environ.get(\"TESTING_MOCKED_DATALOADERS\", None) == \"1\":\n config[\"num_epochs\"] = 2\n # New Code #\n ddp_comm_hook_type = DDPCommunicationHookType(args.ddp_comm_hook)\n ddp_comm_wrapper = DDPCommunicationHookType(args.ddp_comm_wrapper)\n ddp_kwargs = DistributedDataParallelKwargs(comm_hook=ddp_comm_hook_type, comm_wrapper=ddp_comm_wrapper)\n # Initialize accelerator\n accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision, kwargs_handlers=[ddp_kwargs])\n # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs\n lr = config[\"lr\"]\n num_epochs = int(config[\"num_epochs\"])\n seed = int(config[\"seed\"])\n batch_size = int(config[\"batch_size\"])\n\n metric = evaluate.load(\"glue\", \"mrpc\")\n\n set_seed(seed)\n train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size)\n # Instantiate the model (we build the model here so that the seed also control new weights initialization)\n model = AutoModelForSequenceClassification.from_pretrained(\"bert-base-cased\", return_dict=True)\n\n # We could avoid this line since the accelerator is set with `device_placement=True` (default value).\n # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer\n # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).\n model = model.to(accelerator.device)\n\n # Instantiate optimizer\n optimizer = AdamW(params=model.parameters(), lr=lr)\n\n # Instantiate scheduler\n lr_scheduler = get_linear_schedule_with_warmup(\n optimizer=optimizer,\n num_warmup_steps=100,\n num_training_steps=(len(train_dataloader) * num_epochs),\n )\n\n # Prepare everything\n # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the\n # prepare method.\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler\n )\n\n # Now we train the model\n for epoch in range(num_epochs):\n model.train()\n for step, batch in enumerate(train_dataloader):\n # We could avoid this line since we set the accelerator with `device_placement=True`.\n batch.to(accelerator.device)\n # We use the new `accumulate` context manager to perform gradient accumulation\n with accelerator.accumulate(model):\n output = model(**batch)\n loss = output.loss\n accelerator.backward(loss)\n optimizer.step()\n lr_scheduler.step()\n optimizer.zero_grad()\n\n model.eval()\n for step, batch in enumerate(eval_dataloader):\n # We could avoid this line since we set the accelerator with `device_placement=True`.\n batch.to(accelerator.device)\n with torch.no_grad():\n outputs = model(**batch)\n predictions = outputs.logits.argmax(dim=-1)\n predictions, references = accelerator.gather_for_metrics((predictions, batch[\"labels\"]))\n metric.add_batch(\n predictions=predictions,\n references=references,\n )\n\n eval_metric = metric.compute()\n # Use accelerator.print to print only on the main process.\n accelerator.print(f\"epoch {epoch}:\", eval_metric)\n accelerator.end_training()\n\n\ndef main():\n parser = argparse.ArgumentParser(description=\"Simple example of training script.\")\n parser.add_argument(\n \"--mixed_precision\",\n type=str,\n default=None,\n choices=[\"no\", \"fp16\", \"bf16\", \"fp8\"],\n help=\"Whether to use mixed precision. Choose\"\n \"between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10.\"\n \"and an Nvidia Ampere GPU.\",\n )\n # New Code #\n parser.add_argument(\n \"--ddp_comm_hook\",\n type=str,\n default=\"no\",\n choices=[\"no\", \"fp16\", \"bf16\", \"power_sgd\", \"batched_power_sgd\"],\n help=\"DDP Communication hook to use. Choose between `no`, `fp16`, `bf16`, `power_sgd`, and `batched_power_sgd`.\",\n )\n # New Code #\n parser.add_argument(\n \"--ddp_comm_wrapper\",\n type=str,\n default=\"no\",\n choices=[\"no\", \"fp16\", \"bf16\"],\n help=\"DDP Communication wrapper to use. Choose between `no`, `fp16`, and `bf16`.\",\n )\n parser.add_argument(\"--cpu\", action=\"store_true\", help=\"If passed, will train on the CPU.\")\n args = parser.parse_args()\n config = {\"lr\": 2e-5, \"num_epochs\": 3, \"seed\": 42, \"batch_size\": 16}\n training_function(config, args)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "examples/by_feature/deepspeed_with_config_support.py", "content": "#!/usr/bin/env python\n# Copyright 2022 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nFine-tuning the library models for causal language modeling (GPT, GPT-2, CTRL, ...)\non a text file or a dataset without using HuggingFace Trainer.\n\nHere is the full list of checkpoints on the hub that can be fine-tuned by this script:\nhttps://huggingface.co/models?filter=text-generation\n\"\"\"\n# You can also adapt this script on your own causal language modeling task. Pointers for this are left as comments.\n\nimport argparse\nimport json\nimport logging\nimport math\nimport os\nimport random\nfrom itertools import chain\nfrom pathlib import Path\n\nimport datasets\nimport torch\nimport transformers\nfrom datasets import load_dataset\nfrom huggingface_hub import HfApi\nfrom torch.utils.data import DataLoader\nfrom tqdm.auto import tqdm\nfrom transformers import (\n CONFIG_MAPPING,\n MODEL_MAPPING,\n AutoConfig,\n AutoModelForCausalLM,\n AutoTokenizer,\n SchedulerType,\n default_data_collator,\n get_scheduler,\n)\nfrom transformers.utils.versions import require_version\n\nfrom accelerate import Accelerator, DistributedType\nfrom accelerate.logging import get_logger\nfrom accelerate.utils import DummyOptim, DummyScheduler, set_seed\n\n\nlogger = get_logger(__name__)\n\nrequire_version(\"datasets>=1.8.0\", \"To fix: pip install -r examples/pytorch/language-modeling/requirements.txt\")\n\nMODEL_CONFIG_CLASSES = list(MODEL_MAPPING.keys())\nMODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)\n\n\ndef parse_args():\n parser = argparse.ArgumentParser(description=\"Finetune a transformers model on a causal language modeling task\")\n parser.add_argument(\n \"--dataset_name\",\n type=str,\n default=None,\n help=\"The name of the dataset to use (via the datasets library).\",\n )\n parser.add_argument(\n \"--dataset_config_name\",\n type=str,\n default=None,\n help=\"The configuration name of the dataset to use (via the datasets library).\",\n )\n parser.add_argument(\n \"--train_file\", type=str, default=None, help=\"A csv or a json file containing the training data.\"\n )\n parser.add_argument(\n \"--validation_file\", type=str, default=None, help=\"A csv or a json file containing the validation data.\"\n )\n parser.add_argument(\n \"--validation_split_percentage\",\n default=5,\n help=\"The percentage of the train set used as validation set in case there's no validation split\",\n )\n parser.add_argument(\n \"--model_name_or_path\",\n type=str,\n help=\"Path to pretrained model or model identifier from huggingface.co/models.\",\n required=False,\n )\n parser.add_argument(\n \"--config_name\",\n type=str,\n default=None,\n help=\"Pretrained config name or path if not the same as model_name\",\n )\n parser.add_argument(\n \"--tokenizer_name\",\n type=str,\n default=None,\n help=\"Pretrained tokenizer name or path if not the same as model_name\",\n )\n parser.add_argument(\n \"--use_slow_tokenizer\",\n action=\"store_true\",\n help=\"If passed, will use a slow tokenizer (not backed by the šŸ¤— Tokenizers library).\",\n )\n parser.add_argument(\n \"--per_device_train_batch_size\",\n type=int,\n default=8,\n help=\"Batch size (per device) for the training dataloader.\",\n )\n parser.add_argument(\n \"--per_device_eval_batch_size\",\n type=int,\n default=8,\n help=\"Batch size (per device) for the evaluation dataloader.\",\n )\n parser.add_argument(\n \"--learning_rate\",\n type=float,\n default=5e-5,\n help=\"Initial learning rate (after the potential warmup period) to use.\",\n )\n parser.add_argument(\"--weight_decay\", type=float, default=0.0, help=\"Weight decay to use.\")\n parser.add_argument(\"--num_train_epochs\", type=int, default=3, help=\"Total number of training epochs to perform.\")\n parser.add_argument(\n \"--max_train_steps\",\n type=int,\n default=None,\n help=\"Total number of training steps to perform. If provided, overrides num_train_epochs.\",\n )\n parser.add_argument(\n \"--gradient_accumulation_steps\",\n type=int,\n default=1,\n help=\"Number of updates steps to accumulate before performing a backward/update pass.\",\n )\n parser.add_argument(\n \"--lr_scheduler_type\",\n type=SchedulerType,\n default=\"linear\",\n help=\"The scheduler type to use.\",\n choices=[\"linear\", \"cosine\", \"cosine_with_restarts\", \"polynomial\", \"constant\", \"constant_with_warmup\"],\n )\n parser.add_argument(\n \"--num_warmup_steps\", type=int, default=0, help=\"Number of steps for the warmup in the lr scheduler.\"\n )\n parser.add_argument(\"--output_dir\", type=str, default=None, help=\"Where to store the final model.\")\n parser.add_argument(\"--seed\", type=int, default=None, help=\"A seed for reproducible training.\")\n parser.add_argument(\n \"--model_type\",\n type=str,\n default=None,\n help=\"Model type to use if training from scratch.\",\n choices=MODEL_TYPES,\n )\n parser.add_argument(\n \"--block_size\",\n type=int,\n default=None,\n help=(\n \"Optional input sequence length after tokenization. The training dataset will be truncated in block of\"\n \" this size for training. Default to the model max input length for single sentence inputs (take into\"\n \" account special tokens).\"\n ),\n )\n parser.add_argument(\n \"--preprocessing_num_workers\",\n type=int,\n default=None,\n help=\"The number of processes to use for the preprocessing.\",\n )\n parser.add_argument(\n \"--overwrite_cache\", type=bool, default=False, help=\"Overwrite the cached training and evaluation sets\"\n )\n parser.add_argument(\n \"--no_keep_linebreaks\", action=\"store_true\", help=\"Do not keep line breaks when using TXT files.\"\n )\n parser.add_argument(\"--push_to_hub\", action=\"store_true\", help=\"Whether or not to push the model to the Hub.\")\n parser.add_argument(\n \"--hub_model_id\", type=str, help=\"The name of the repository to keep in sync with the local `output_dir`.\"\n )\n parser.add_argument(\"--hub_token\", type=str, help=\"The token to use to push to the Model Hub.\")\n parser.add_argument(\n \"--checkpointing_steps\",\n type=str,\n default=None,\n help=\"Whether the various states should be saved at the end of every n steps, or 'epoch' for each epoch.\",\n )\n parser.add_argument(\n \"--resume_from_checkpoint\",\n type=str,\n default=None,\n help=\"If the training should continue from a checkpoint folder.\",\n )\n # New Code #\n # Whether to load the best model at the end of training\n parser.add_argument(\n \"--load_best_model\",\n action=\"store_true\",\n help=\"Whether to load the best model at the end of training\",\n )\n parser.add_argument(\n \"--with_tracking\",\n action=\"store_true\",\n help=\"Whether to enable experiment trackers for logging.\",\n )\n parser.add_argument(\n \"--report_to\",\n type=str,\n default=\"all\",\n help=(\n 'The integration to report the results and logs to. Supported platforms are `\"tensorboard\"`,'\n ' `\"wandb\"`, `\"comet_ml\"`, `\"dvclive\"`, and `\"swanlab\"`. Use `\"all\"` (default) to report to all integrations.'\n \"Only applicable when `--with_tracking` is passed.\"\n ),\n )\n args = parser.parse_args()\n\n # Sanity checks\n if args.dataset_name is None and args.train_file is None and args.validation_file is None:\n raise ValueError(\"Need either a dataset name or a training/validation file.\")\n else:\n if args.train_file is not None:\n extension = args.train_file.split(\".\")[-1]\n assert extension in [\"csv\", \"json\", \"txt\"], \"`train_file` should be a csv, json or txt file.\"\n if args.validation_file is not None:\n extension = args.validation_file.split(\".\")[-1]\n assert extension in [\"csv\", \"json\", \"txt\"], \"`validation_file` should be a csv, json or txt file.\"\n\n if args.push_to_hub:\n assert args.output_dir is not None, \"Need an `output_dir` to create a repo when `--push_to_hub` is passed.\"\n\n return args\n\n\n# New Code #\ndef evaluate(args, model, eval_dataloader, accelerator, eval_dataset):\n model.eval()\n losses = []\n for step, batch in enumerate(eval_dataloader):\n with torch.no_grad():\n outputs = model(**batch)\n\n loss = outputs.loss\n losses.append(accelerator.gather_for_metrics(loss.repeat(args.per_device_eval_batch_size)))\n\n losses = torch.cat(losses)\n try:\n eval_loss = torch.mean(losses)\n perplexity = math.exp(eval_loss)\n except OverflowError:\n perplexity = float(\"inf\")\n return perplexity, eval_loss\n\n\ndef main():\n args = parse_args()\n\n # Initialize the accelerator. We will let the accelerator handle device placement for us in this example.\n # If we're using tracking, we also need to initialize it here and it will by default pick up all supported trackers\n # in the environment\n\n # when using DeepSpeed, the `gradient_accumulation_steps` is properly set from the DeepSpeed plugin/config\n # or from `accelerate launch` via `--gradient_accumulation_steps` else\n # defaulting to the passed `args.gradient_accumulation_steps`\n accelerator = (\n Accelerator(\n log_with=args.report_to,\n project_dir=args.output_dir,\n gradient_accumulation_steps=args.gradient_accumulation_steps,\n )\n if args.with_tracking\n else Accelerator(gradient_accumulation_steps=args.gradient_accumulation_steps)\n )\n\n # Make one log on every process with the configuration for debugging.\n logging.basicConfig(\n format=\"%(asctime)s - %(levelname)s - %(name)s - %(message)s\",\n datefmt=\"%m/%d/%Y %H:%M:%S\",\n level=logging.INFO,\n )\n logger.info(accelerator.state, main_process_only=False)\n if accelerator.is_local_main_process:\n datasets.utils.logging.set_verbosity_warning()\n transformers.utils.logging.set_verbosity_info()\n else:\n datasets.utils.logging.set_verbosity_error()\n transformers.utils.logging.set_verbosity_error()\n\n # If passed along, set the training seed now.\n if args.seed is not None:\n set_seed(args.seed)\n\n # Handle the repository creation\n if accelerator.is_main_process:\n if args.push_to_hub:\n api = HfApi(token=args.hub_token)\n\n # Create repo (repo_name from args or inferred)\n repo_name = args.hub_model_id\n if repo_name is None:\n repo_name = Path(args.output_dir).absolute().name\n repo_id = api.create_repo(repo_name, exist_ok=True).repo_id\n\n with open(os.path.join(args.output_dir, \".gitignore\"), \"w+\") as gitignore:\n if \"step_*\" not in gitignore:\n gitignore.write(\"step_*\\n\")\n if \"epoch_*\" not in gitignore:\n gitignore.write(\"epoch_*\\n\")\n elif args.output_dir is not None:\n os.makedirs(args.output_dir, exist_ok=True)\n accelerator.wait_for_everyone()\n\n # Get the datasets: you can either provide your own CSV/JSON/TXT training and evaluation files (see below)\n # or just provide the name of one of the public datasets available on the hub at https://huggingface.co/datasets/\n # (the dataset will be downloaded automatically from the datasets Hub).\n #\n # For CSV/JSON files, this script will use the column called 'text' or the first column if no column called\n # 'text' is found. You can easily tweak this behavior (see below).\n #\n # In distributed training, the load_dataset function guarantee that only one local process can concurrently\n # download the dataset.\n if args.dataset_name is not None:\n # Downloading and loading a dataset from the hub.\n raw_datasets = load_dataset(args.dataset_name, args.dataset_config_name)\n if \"validation\" not in raw_datasets.keys():\n raw_datasets[\"validation\"] = load_dataset(\n args.dataset_name,\n args.dataset_config_name,\n split=f\"train[:{args.validation_split_percentage}%]\",\n )\n raw_datasets[\"train\"] = load_dataset(\n args.dataset_name,\n args.dataset_config_name,\n split=f\"train[{args.validation_split_percentage}%:]\",\n )\n else:\n data_files = {}\n dataset_args = {}\n if args.train_file is not None:\n data_files[\"train\"] = args.train_file\n if args.validation_file is not None:\n data_files[\"validation\"] = args.validation_file\n extension = args.train_file.split(\".\")[-1]\n if extension == \"txt\":\n extension = \"text\"\n dataset_args[\"keep_linebreaks\"] = not args.no_keep_linebreaks\n raw_datasets = load_dataset(extension, data_files=data_files, **dataset_args)\n # If no validation data is there, validation_split_percentage will be used to divide the dataset.\n if \"validation\" not in raw_datasets.keys():\n raw_datasets[\"validation\"] = load_dataset(\n extension,\n data_files=data_files,\n split=f\"train[:{args.validation_split_percentage}%]\",\n **dataset_args,\n )\n raw_datasets[\"train\"] = load_dataset(\n extension,\n data_files=data_files,\n split=f\"train[{args.validation_split_percentage}%:]\",\n **dataset_args,\n )\n\n # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at\n # https://huggingface.co/docs/datasets/loading_datasets.html.\n\n # Load pretrained model and tokenizer\n #\n # In distributed training, the .from_pretrained methods guarantee that only one local process can concurrently\n # download model & vocab.\n if args.config_name:\n config = AutoConfig.from_pretrained(args.config_name)\n elif args.model_name_or_path:\n config = AutoConfig.from_pretrained(args.model_name_or_path)\n else:\n config = CONFIG_MAPPING[args.model_type]()\n logger.warning(\"You are instantiating a new config instance from scratch.\")\n\n if args.tokenizer_name:\n tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_name, use_fast=not args.use_slow_tokenizer)\n elif args.model_name_or_path:\n tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path, use_fast=not args.use_slow_tokenizer)\n else:\n raise ValueError(\n \"You are instantiating a new tokenizer from scratch. This is not supported by this script.\"\n \"You can do it from another script, save it, and load it from here, using --tokenizer_name.\"\n )\n\n if args.model_name_or_path:\n model = AutoModelForCausalLM.from_pretrained(\n args.model_name_or_path,\n from_tf=bool(\".ckpt\" in args.model_name_or_path),\n config=config,\n )\n else:\n logger.info(\"Training new model from scratch\")\n model = AutoModelForCausalLM.from_config(config)\n\n model.resize_token_embeddings(len(tokenizer))\n\n # Preprocessing the datasets.\n # First we tokenize all the texts.\n column_names = raw_datasets[\"train\"].column_names\n text_column_name = \"text\" if \"text\" in column_names else column_names[0]\n\n def tokenize_function(examples):\n return tokenizer(examples[text_column_name])\n\n with accelerator.main_process_first():\n tokenized_datasets = raw_datasets.map(\n tokenize_function,\n batched=True,\n num_proc=args.preprocessing_num_workers,\n remove_columns=column_names,\n load_from_cache_file=not args.overwrite_cache,\n desc=\"Running tokenizer on dataset\",\n )\n\n if args.block_size is None:\n block_size = tokenizer.model_max_length\n if block_size > 1024:\n logger.warning(\n f\"The tokenizer picked seems to have a very large `model_max_length` ({tokenizer.model_max_length}). \"\n \"Picking 1024 instead. You can change that default value by passing --block_size xxx.\"\n )\n block_size = 1024\n else:\n if args.block_size > tokenizer.model_max_length:\n logger.warning(\n f\"The block_size passed ({args.block_size}) is larger than the maximum length for the model\"\n f\"({tokenizer.model_max_length}). Using block_size={tokenizer.model_max_length}.\"\n )\n block_size = min(args.block_size, tokenizer.model_max_length)\n\n # Main data processing function that will concatenate all texts from our dataset and generate chunks of block_size.\n def group_texts(examples):\n # Concatenate all texts.\n concatenated_examples = {k: list(chain(*examples[k])) for k in examples.keys()}\n total_length = len(concatenated_examples[list(examples.keys())[0]])\n # We drop the small remainder, we could add padding if the model supported it instead of this drop, you can\n # customize this part to your needs.\n if total_length >= block_size:\n total_length = (total_length // block_size) * block_size\n # Split by chunks of max_len.\n result = {\n k: [t[i : i + block_size] for i in range(0, total_length, block_size)]\n for k, t in concatenated_examples.items()\n }\n result[\"labels\"] = result[\"input_ids\"].copy()\n return result\n\n # Note that with `batched=True`, this map processes 1,000 texts together, so group_texts throws away a remainder\n # for each of those groups of 1,000 texts. You can adjust that batch_size here but a higher value might be slower\n # to preprocess.\n #\n # To speed up this part, we use multiprocessing. See the documentation of the map method for more information:\n # https://huggingface.co/docs/datasets/package_reference/main_classes.html#datasets.Dataset.map\n\n with accelerator.main_process_first():\n lm_datasets = tokenized_datasets.map(\n group_texts,\n batched=True,\n num_proc=args.preprocessing_num_workers,\n load_from_cache_file=not args.overwrite_cache,\n desc=f\"Grouping texts in chunks of {block_size}\",\n )\n\n train_dataset = lm_datasets[\"train\"]\n eval_dataset = lm_datasets[\"validation\"]\n\n # Log a few random samples from the training set:\n for index in random.sample(range(len(train_dataset)), 3):\n logger.info(f\"Sample {index} of the training set: {train_dataset[index]}.\")\n\n # DataLoaders creation:\n train_dataloader = DataLoader(\n train_dataset, shuffle=True, collate_fn=default_data_collator, batch_size=args.per_device_train_batch_size\n )\n eval_dataloader = DataLoader(\n eval_dataset, collate_fn=default_data_collator, batch_size=args.per_device_eval_batch_size\n )\n\n # Optimizer\n # Split weights in two groups, one with weight decay and the other not.\n no_decay = [\"bias\", \"LayerNorm.weight\"]\n optimizer_grouped_parameters = [\n {\n \"params\": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],\n \"weight_decay\": args.weight_decay,\n },\n {\n \"params\": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)],\n \"weight_decay\": 0.0,\n },\n ]\n # New Code #\n # Creates Dummy Optimizer if `optimizer` was specified in the config file else creates Adam Optimizer\n optimizer_cls = (\n torch.optim.AdamW\n if accelerator.state.deepspeed_plugin is None\n or \"optimizer\" not in accelerator.state.deepspeed_plugin.deepspeed_config\n else DummyOptim\n )\n optimizer = optimizer_cls(optimizer_grouped_parameters, lr=args.learning_rate)\n\n # On TPU, the tie weights in our model have been disconnected, so we need to restore the ties.\n if accelerator.distributed_type == DistributedType.XLA:\n model.tie_weights()\n\n # Scheduler and math around the number of training steps.\n num_update_steps_per_epoch = math.ceil(len(train_dataloader) / accelerator.gradient_accumulation_steps)\n overrode_max_train_steps = False\n if args.max_train_steps is None:\n args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch\n overrode_max_train_steps = True\n else:\n args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)\n\n # New Code #\n # Creates Dummy Scheduler if `scheduler` was specified in the config file else creates `args.lr_scheduler_type` Scheduler\n if (\n accelerator.state.deepspeed_plugin is None\n or \"scheduler\" not in accelerator.state.deepspeed_plugin.deepspeed_config\n ):\n lr_scheduler = get_scheduler(\n name=args.lr_scheduler_type,\n optimizer=optimizer,\n num_warmup_steps=args.num_warmup_steps,\n num_training_steps=args.max_train_steps,\n )\n else:\n lr_scheduler = DummyScheduler(\n optimizer, total_num_steps=args.max_train_steps, warmup_num_steps=args.num_warmup_steps\n )\n\n # Prepare everything with our `accelerator`.\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler\n )\n\n # We need to recalculate our total training steps as the size of the training dataloader may have changed.\n num_update_steps_per_epoch = math.ceil(len(train_dataloader) / accelerator.gradient_accumulation_steps)\n if overrode_max_train_steps:\n args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch\n # Afterwards we recalculate our number of training epochs\n args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)\n\n # Figure out how many steps we should save the Accelerator states\n checkpointing_steps = args.checkpointing_steps\n if checkpointing_steps is not None and checkpointing_steps.isdigit():\n checkpointing_steps = int(checkpointing_steps)\n\n # We need to initialize the trackers we use, and also store our configuration.\n # The trackers initializes automatically on the main process.\n if args.with_tracking:\n experiment_config = vars(args)\n # TensorBoard cannot log Enums, need the raw value\n experiment_config[\"lr_scheduler_type\"] = experiment_config[\"lr_scheduler_type\"].value\n accelerator.init_trackers(\"clm_no_trainer\", experiment_config)\n\n # Train!\n total_batch_size = (\n args.per_device_train_batch_size * accelerator.num_processes * accelerator.gradient_accumulation_steps\n )\n\n logger.info(\"***** Running training *****\")\n logger.info(f\" Num examples = {len(train_dataset)}\")\n logger.info(f\" Num Epochs = {args.num_train_epochs}\")\n logger.info(f\" Instantaneous batch size per device = {args.per_device_train_batch_size}\")\n logger.info(f\" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}\")\n logger.info(f\" Gradient Accumulation steps = {accelerator.gradient_accumulation_steps}\")\n logger.info(f\" Total optimization steps = {args.max_train_steps}\")\n # Only show the progress bar once on each machine.\n progress_bar = tqdm(range(args.max_train_steps), disable=not accelerator.is_local_main_process)\n completed_steps = 0\n starting_epoch = 0\n best_metric = None\n best_metric_checkpoint = None\n\n # Potentially load in the weights and states from a previous save\n if args.resume_from_checkpoint:\n accelerator.load_state(args.resume_from_checkpoint)\n accelerator.print(f\"Resumed from checkpoint: {args.resume_from_checkpoint}\")\n path = os.path.basename(args.resume_from_checkpoint)\n training_difference = os.path.splitext(path)[0]\n\n if \"epoch\" in training_difference:\n starting_epoch = int(training_difference.replace(\"epoch_\", \"\")) + 1\n resume_step = None\n completed_steps = starting_epoch * num_update_steps_per_epoch\n else:\n resume_step = int(training_difference.replace(\"step_\", \"\"))\n starting_epoch = resume_step // num_update_steps_per_epoch\n resume_step -= starting_epoch * num_update_steps_per_epoch\n completed_steps = resume_step\n\n # update progress bar if resumed from checkpoint\n progress_bar.update(completed_steps)\n\n for epoch in range(starting_epoch, args.num_train_epochs):\n model.train()\n if args.with_tracking:\n total_loss = 0\n\n # skip new `skip_first_batches` to skip the batches when resuming from ckpt\n if args.resume_from_checkpoint and epoch == starting_epoch and resume_step is not None:\n # We need to skip steps until we reach the resumed step\n active_dataloader = accelerator.skip_first_batches(train_dataloader, resume_step)\n else:\n # After the first iteration though, we need to go back to the original dataloader\n active_dataloader = train_dataloader\n for step, batch in enumerate(active_dataloader):\n # In particular, DeepSpeed handles `gradient_accumulation` via `DeepSpeedEngine`.\n # Below, we use `accelerator.accumulate` if the user\n # wants to switch to other approaches such as plain DDP, PyTorch FSDP ...\n # This avoids having to change any code as things are all handled across different distributed setups.\n with accelerator.accumulate(model):\n outputs = model(**batch)\n loss = outputs.loss\n accelerator.backward(loss)\n optimizer.step()\n lr_scheduler.step()\n optimizer.zero_grad()\n\n if accelerator.sync_gradients:\n progress_bar.update(1)\n completed_steps += 1\n\n # We keep track of the loss at each epoch\n if args.with_tracking:\n step_loss = accelerator.reduce(loss.detach().clone()).item()\n total_loss += step_loss\n\n if isinstance(checkpointing_steps, int):\n if completed_steps % checkpointing_steps == 0:\n output_dir = f\"step_{completed_steps}\"\n if args.output_dir is not None:\n output_dir = os.path.join(args.output_dir, output_dir)\n accelerator.save_state(output_dir)\n if completed_steps >= args.max_train_steps:\n break\n\n perplexity, eval_loss = evaluate(args, model, eval_dataloader, accelerator, eval_dataset)\n logger.info(f\"epoch {epoch}: perplexity: {perplexity} eval_loss: {eval_loss}\")\n\n if args.with_tracking:\n accelerator.log(\n {\n \"perplexity\": perplexity,\n \"eval_loss\": eval_loss,\n \"train_loss\": total_loss / len(train_dataloader),\n \"epoch\": epoch,\n \"step\": completed_steps,\n },\n step=completed_steps,\n )\n\n if isinstance(checkpointing_steps, str) and checkpointing_steps == \"epoch\":\n accelerator.save_state(os.path.join(args.output_dir, f\"epoch_{epoch}\"))\n\n # New Code #\n # Tracks the best checkpoint and best metric\n if best_metric is None or best_metric > perplexity:\n best_metric = perplexity\n best_metric_checkpoint = os.path.join(args.output_dir, \"best_checkpoint\")\n accelerator.save_state(best_metric_checkpoint)\n accelerator.print(f\"New best metric: {best_metric} at epoch {epoch}\")\n accelerator.print(f\"best_metric_checkpoint: {best_metric_checkpoint}\")\n\n # New Code #\n # Loads the best checkpoint after the training is finished\n if args.load_best_model:\n accelerator.load_state(best_metric_checkpoint)\n\n # New Code #\n # Evaluates using the best checkpoint\n perplexity, eval_loss = evaluate(args, model, eval_dataloader, accelerator, eval_dataset)\n logger.info(f\"Best model metrics: perplexity: {perplexity} eval_loss: {eval_loss}\")\n if perplexity != best_metric:\n raise AssertionError(\n f\"Best metric {best_metric} does not match the metric {perplexity} of the loaded best model.\"\n )\n\n if args.output_dir is not None:\n accelerator.wait_for_everyone()\n unwrapped_model = accelerator.unwrap_model(model)\n\n # New Code #\n # Saves the whole/unpartitioned fp16 model when in ZeRO Stage-3 to the output directory if\n # `stage3_gather_16bit_weights_on_model_save` is True in DeepSpeed Config file or\n # `zero3_save_16bit_model` is True in DeepSpeed Plugin.\n # For Zero Stages 1 and 2, models are saved as usual in the output directory.\n # The model name saved is `pytorch_model.bin`\n unwrapped_model.save_pretrained(\n args.output_dir,\n is_main_process=accelerator.is_main_process,\n save_function=accelerator.save,\n state_dict=accelerator.get_state_dict(model),\n )\n if accelerator.is_main_process:\n tokenizer.save_pretrained(args.output_dir)\n if args.push_to_hub:\n api.upload_folder(\n repo_id=repo_id,\n folder_path=args.output_dir,\n commit_message=\"End of training\",\n )\n\n with open(os.path.join(args.output_dir, \"all_results.json\"), \"w\") as f:\n json.dump({\"perplexity\": perplexity, \"eval_loss\": eval_loss.item()}, f)\n accelerator.end_training()\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "examples/by_feature/early_stopping.py", "content": "# Copyright 2021 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport argparse\n\nimport evaluate\nimport torch\nfrom datasets import load_dataset\nfrom torch.optim import AdamW\nfrom torch.utils.data import DataLoader\nfrom transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed\n\nfrom accelerate import Accelerator, DistributedType\n\n\n########################################################################\n# This is a fully working simple example to use Accelerate\n# specifically showcasing how to perform early stopping,\n# and builds off the `nlp_example.py` script\n#\n# This example trains a Bert base model on GLUE MRPC\n# in any of the following settings (with the same script):\n# - single CPU or single GPU\n# - multi GPUS (using PyTorch distributed mode)\n# - (multi) TPUs\n# - fp16 (mixed-precision) or fp32 (normal precision)\n#\n# To run it in each of these various modes, follow the instructions\n# in the readme for examples:\n# https://github.com/huggingface/accelerate/tree/main/examples\n#\n########################################################################\n\n\nMAX_GPU_BATCH_SIZE = 16\nEVAL_BATCH_SIZE = 32\n\n\ndef get_dataloaders(accelerator: Accelerator, batch_size: int = 16):\n \"\"\"\n Creates a set of `DataLoader`s for the `glue` dataset,\n using \"bert-base-cased\" as the tokenizer.\n\n Args:\n accelerator (`Accelerator`):\n An `Accelerator` object\n batch_size (`int`, *optional*):\n The batch size for the train and validation DataLoaders.\n \"\"\"\n tokenizer = AutoTokenizer.from_pretrained(\"bert-base-cased\")\n datasets = load_dataset(\"glue\", \"mrpc\")\n\n def tokenize_function(examples):\n # max_length=None => use the model max length (it's actually the default)\n outputs = tokenizer(examples[\"sentence1\"], examples[\"sentence2\"], truncation=True, max_length=None)\n return outputs\n\n # Apply the method we just defined to all the examples in all the splits of the dataset\n # starting with the main process first:\n with accelerator.main_process_first():\n tokenized_datasets = datasets.map(\n tokenize_function,\n batched=True,\n remove_columns=[\"idx\", \"sentence1\", \"sentence2\"],\n )\n\n # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the\n # transformers library\n tokenized_datasets = tokenized_datasets.rename_column(\"label\", \"labels\")\n\n def collate_fn(examples):\n # On TPU it's best to pad everything to the same length or training will be very slow.\n max_length = 128 if accelerator.distributed_type == DistributedType.XLA else None\n # When using mixed precision we want round multiples of 8/16\n if accelerator.mixed_precision == \"fp8\":\n pad_to_multiple_of = 16\n elif accelerator.mixed_precision != \"no\":\n pad_to_multiple_of = 8\n else:\n pad_to_multiple_of = None\n\n return tokenizer.pad(\n examples,\n padding=\"longest\",\n max_length=max_length,\n pad_to_multiple_of=pad_to_multiple_of,\n return_tensors=\"pt\",\n )\n\n # Instantiate dataloaders.\n train_dataloader = DataLoader(\n tokenized_datasets[\"train\"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size, drop_last=True\n )\n eval_dataloader = DataLoader(\n tokenized_datasets[\"validation\"],\n shuffle=False,\n collate_fn=collate_fn,\n batch_size=EVAL_BATCH_SIZE,\n drop_last=(accelerator.mixed_precision == \"fp8\"),\n )\n\n return train_dataloader, eval_dataloader\n\n\n# New code\nclass EarlyStoppingCallback:\n \"A callback class that helps with early stopping\"\n\n def __init__(self, min_delta=0, patience=5):\n self.min_delta = min_delta\n self.patience = patience\n self.counter = 0\n self.lowest_loss = float(\"inf\")\n\n def check_early_stopping(self, eval_loss):\n delta = self.lowest_loss - eval_loss\n if delta >= self.min_delta:\n self.lowest_loss = eval_loss\n self.counter = 0\n else:\n self.counter += 1\n if self.counter >= self.patience:\n return True\n return False\n\n\ncallback = EarlyStoppingCallback()\n\n\ndef training_function(config, args):\n # Initialize accelerator\n accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)\n # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs\n lr = config[\"lr\"]\n num_epochs = int(config[\"num_epochs\"])\n seed = int(config[\"seed\"])\n batch_size = int(config[\"batch_size\"])\n\n metric = evaluate.load(\"glue\", \"mrpc\")\n\n # If the batch size is too big we use gradient accumulation\n gradient_accumulation_steps = 1\n if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.XLA:\n gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE\n batch_size = MAX_GPU_BATCH_SIZE\n\n set_seed(seed)\n train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size)\n # Instantiate the model (we build the model here so that the seed also control new weights initialization)\n model = AutoModelForSequenceClassification.from_pretrained(\"bert-base-cased\", return_dict=True)\n\n # We could avoid this line since the accelerator is set with `device_placement=True` (default value).\n # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer\n # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).\n model = model.to(accelerator.device)\n # Instantiate optimizer\n optimizer = AdamW(params=model.parameters(), lr=lr)\n\n # Instantiate scheduler\n lr_scheduler = get_linear_schedule_with_warmup(\n optimizer=optimizer,\n num_warmup_steps=100,\n num_training_steps=(len(train_dataloader) * num_epochs) // gradient_accumulation_steps,\n )\n\n # Prepare everything\n # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the\n # prepare method.\n\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler\n )\n\n # Now we train the model\n for epoch in range(num_epochs):\n model.train()\n for step, batch in enumerate(train_dataloader):\n # We could avoid this line since we set the accelerator with `device_placement=True`.\n batch.to(accelerator.device)\n outputs = model(**batch)\n loss = outputs.loss\n loss = loss / gradient_accumulation_steps\n accelerator.backward(loss)\n if step % gradient_accumulation_steps == 0:\n optimizer.step()\n lr_scheduler.step()\n optimizer.zero_grad()\n\n # New code\n # Check if we should stop the training on any processes\n if callback.check_early_stopping(loss.item()):\n accelerator.set_trigger()\n\n # If so, we break the loop\n if accelerator.check_trigger():\n break\n\n model.eval()\n for step, batch in enumerate(eval_dataloader):\n # We could avoid this line since we set the accelerator with `device_placement=True`.\n batch.to(accelerator.device)\n with torch.no_grad():\n outputs = model(**batch)\n predictions = outputs.logits.argmax(dim=-1)\n predictions, references = accelerator.gather_for_metrics((predictions, batch[\"labels\"]))\n metric.add_batch(\n predictions=predictions,\n references=references,\n )\n\n eval_metric = metric.compute()\n\n # Use accelerator.print to print only on the main process.\n accelerator.print(f\"epoch {epoch}:\", eval_metric)\n accelerator.end_training()\n\n\ndef main():\n parser = argparse.ArgumentParser(description=\"Simple example of training script.\")\n parser.add_argument(\n \"--mixed_precision\",\n type=str,\n default=None,\n choices=[\"no\", \"fp16\", \"bf16\", \"fp8\"],\n help=\"Whether to use mixed precision. Choose\"\n \"between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10.\"\n \"and an Nvidia Ampere GPU.\",\n )\n parser.add_argument(\"--cpu\", action=\"store_true\", help=\"If passed, will train on the CPU.\")\n args = parser.parse_args()\n config = {\"lr\": 2e-5, \"num_epochs\": 3, \"seed\": 42, \"batch_size\": 16}\n training_function(config, args)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "examples/by_feature/fsdp_with_peak_mem_tracking.py", "content": "# Copyright 2021 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport argparse\nimport gc\nimport os\nimport threading\n\nimport evaluate\nimport psutil\nimport torch\nfrom datasets import load_dataset\nfrom torch.distributed.fsdp.fully_sharded_data_parallel import FullOptimStateDictConfig, FullStateDictConfig\nfrom torch.utils.data import DataLoader\nfrom transformers import (\n AutoModelForSequenceClassification,\n AutoTokenizer,\n get_linear_schedule_with_warmup,\n set_seed,\n)\n\nfrom accelerate import Accelerator, DistributedType, FullyShardedDataParallelPlugin\nfrom accelerate.utils import is_npu_available, is_xpu_available\n\n\n########################################################################\n# This is a fully working simple example to use Accelerate\n#\n# This example trains a Bert base model on GLUE MRPC\n# in any of the following settings (with the same script):\n# - single CPU or single GPU\n# - multi GPUS (using PyTorch distributed mode)\n# - (multi) TPUs\n# - fp16 (mixed-precision) or fp32 (normal precision)\n# - FSDP\n#\n# This example also demonstrates the checkpointing and sharding capabilities\n#\n# To run it in each of these various modes, follow the instructions\n# in the readme for examples:\n# https://github.com/huggingface/accelerate/tree/main/examples\n#\n########################################################################\n\n\nMAX_GPU_BATCH_SIZE = 16\nEVAL_BATCH_SIZE = 32\n\n\n# New Code #\n# Converting Bytes to Megabytes\ndef b2mb(x):\n return int(x / 2**20)\n\n\n# New Code #\n# This context manager is used to track the peak memory usage of the process\nclass TorchTracemalloc:\n def __enter__(self):\n gc.collect()\n if torch.cuda.is_available():\n torch.cuda.empty_cache()\n torch.cuda.reset_max_memory_allocated() # reset the peak gauge to zero\n self.begin = torch.cuda.memory_allocated()\n elif is_xpu_available():\n torch.xpu.empty_cache()\n torch.xpu.reset_max_memory_allocated() # reset the peak gauge to zero\n self.begin = torch.xpu.memory_allocated()\n elif is_npu_available():\n torch.npu.empty_cache()\n torch.npu.reset_max_memory_allocated() # reset the peak gauge to zero\n self.begin = torch.npu.memory_allocated()\n self.process = psutil.Process()\n\n self.cpu_begin = self.cpu_mem_used()\n self.peak_monitoring = True\n peak_monitor_thread = threading.Thread(target=self.peak_monitor_func)\n peak_monitor_thread.daemon = True\n peak_monitor_thread.start()\n return self\n\n def cpu_mem_used(self):\n \"\"\"get resident set size memory for the current process\"\"\"\n return self.process.memory_info().rss\n\n def peak_monitor_func(self):\n self.cpu_peak = -1\n\n while True:\n self.cpu_peak = max(self.cpu_mem_used(), self.cpu_peak)\n\n # can't sleep or will not catch the peak right (this comment is here on purpose)\n # time.sleep(0.001) # 1msec\n\n if not self.peak_monitoring:\n break\n\n def __exit__(self, *exc):\n self.peak_monitoring = False\n\n gc.collect()\n if torch.cuda.is_available():\n torch.cuda.empty_cache()\n self.end = torch.cuda.memory_allocated()\n self.peak = torch.cuda.max_memory_allocated()\n elif is_xpu_available():\n torch.xpu.empty_cache()\n self.end = torch.xpu.memory_allocated()\n self.peak = torch.xpu.max_memory_allocated()\n elif is_npu_available():\n torch.npu.empty_cache()\n self.end = torch.npu.memory_allocated()\n self.peak = torch.npu.max_memory_allocated()\n self.used = b2mb(self.end - self.begin)\n self.peaked = b2mb(self.peak - self.begin)\n\n self.cpu_end = self.cpu_mem_used()\n self.cpu_used = b2mb(self.cpu_end - self.cpu_begin)\n self.cpu_peaked = b2mb(self.cpu_peak - self.cpu_begin)\n # print(f\"delta used/peak {self.used:4d}/{self.peaked:4d}\")\n\n\n# For testing only\nif os.environ.get(\"TESTING_MOCKED_DATALOADERS\", None) == \"1\":\n from accelerate.test_utils.training import mocked_dataloaders\n\n get_dataloaders = mocked_dataloaders # noqa: F811\n\n\ndef training_function(config, args):\n # For testing only\n if os.environ.get(\"TESTING_MOCKED_DATALOADERS\", None) == \"1\":\n config[\"num_epochs\"] = 2\n\n # New Code #\n # Pass the advanced FSDP settings not part of the accelerate config by creating fsdp_plugin\n fsdp_plugin = FullyShardedDataParallelPlugin(\n state_dict_config=FullStateDictConfig(offload_to_cpu=False, rank0_only=False),\n optim_state_dict_config=FullOptimStateDictConfig(offload_to_cpu=False, rank0_only=False),\n )\n\n # Initialize accelerator\n if args.with_tracking:\n accelerator = Accelerator(\n cpu=args.cpu,\n mixed_precision=args.mixed_precision,\n log_with=\"wandb\",\n project_dir=args.logging_dir,\n fsdp_plugin=fsdp_plugin,\n )\n else:\n accelerator = Accelerator(fsdp_plugin=fsdp_plugin)\n accelerator.print(accelerator.distributed_type)\n\n if hasattr(args.checkpointing_steps, \"isdigit\"):\n if args.checkpointing_steps == \"epoch\":\n checkpointing_steps = args.checkpointing_steps\n elif args.checkpointing_steps.isdigit():\n checkpointing_steps = int(args.checkpointing_steps)\n else:\n raise ValueError(\n f\"Argument `checkpointing_steps` must be either a number or `epoch`. `{args.checkpointing_steps}` passed.\"\n )\n else:\n checkpointing_steps = None\n # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs\n lr = config[\"lr\"]\n num_epochs = int(config[\"num_epochs\"])\n seed = int(config[\"seed\"])\n batch_size = int(config[\"batch_size\"])\n\n # We need to initialize the trackers we use, and also store our configuration\n if args.with_tracking:\n experiment_config = vars(args)\n accelerator.init_trackers(\"fsdp_glue_no_trainer\", experiment_config)\n\n tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path)\n datasets = load_dataset(\"glue\", \"mrpc\")\n metric = evaluate.load(\"glue\", \"mrpc\")\n\n def tokenize_function(examples):\n # max_length=None => use the model max length (it's actually the default)\n outputs = tokenizer(examples[\"sentence1\"], examples[\"sentence2\"], truncation=True, max_length=None)\n return outputs\n\n # Apply the method we just defined to all the examples in all the splits of the dataset\n # starting with the main process first:\n with accelerator.main_process_first():\n tokenized_datasets = datasets.map(\n tokenize_function,\n batched=True,\n remove_columns=[\"idx\", \"sentence1\", \"sentence2\"],\n )\n\n # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the\n # transformers library\n tokenized_datasets = tokenized_datasets.rename_column(\"label\", \"labels\")\n\n # If the batch size is too big we use gradient accumulation\n gradient_accumulation_steps = 1\n if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.XLA:\n gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE\n batch_size = MAX_GPU_BATCH_SIZE\n\n def collate_fn(examples):\n # On TPU it's best to pad everything to the same length or training will be very slow.\n max_length = 128 if accelerator.distributed_type == DistributedType.XLA else None\n # When using mixed precision we want round multiples of 8/16\n if accelerator.mixed_precision == \"fp8\":\n pad_to_multiple_of = 16\n elif accelerator.mixed_precision != \"no\":\n pad_to_multiple_of = 8\n else:\n pad_to_multiple_of = None\n\n return tokenizer.pad(\n examples,\n padding=\"longest\",\n max_length=max_length,\n pad_to_multiple_of=pad_to_multiple_of,\n return_tensors=\"pt\",\n )\n\n # Instantiate dataloaders.\n train_dataloader = DataLoader(\n tokenized_datasets[\"train\"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size\n )\n eval_dataloader = DataLoader(\n tokenized_datasets[\"validation\"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE\n )\n\n set_seed(seed)\n\n # Instantiate the model (we build the model here so that the seed also control new weights initialization)\n model = AutoModelForSequenceClassification.from_pretrained(\n args.model_name_or_path, return_dict=True, low_cpu_mem_usage=True\n )\n\n no_decay = [\"bias\", \"LayerNorm.weight\"]\n optimizer_grouped_parameters = [\n {\n \"params\": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],\n \"weight_decay\": 0.003,\n },\n {\n \"params\": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)],\n \"weight_decay\": 0.0,\n },\n ]\n\n optimizer = torch.optim.AdamW(params=optimizer_grouped_parameters, lr=lr, weight_decay=2e-4)\n\n # Instantiate scheduler\n lr_scheduler = get_linear_schedule_with_warmup(\n optimizer=optimizer,\n num_warmup_steps=10,\n num_training_steps=(len(train_dataloader) * num_epochs) // gradient_accumulation_steps,\n )\n\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler\n )\n\n overall_step = 0\n\n # Potentially load in the weights and states from a previous save\n if args.resume_from_checkpoint:\n if args.resume_from_checkpoint is not None or args.resume_from_checkpoint != \"\":\n accelerator.print(f\"Resumed from checkpoint: {args.resume_from_checkpoint}\")\n accelerator.load_state(args.resume_from_checkpoint)\n path = os.path.basename(args.resume_from_checkpoint)\n else:\n # Get the most recent checkpoint\n dirs = [f.name for f in os.scandir(os.getcwd()) if f.is_dir()]\n dirs.sort(key=os.path.getctime)\n path = dirs[-1] # Sorts folders by date modified, most recent checkpoint is the last\n # Extract `epoch_{i}` or `step_{i}`\n training_difference = os.path.splitext(path)[0]\n\n if \"epoch\" in training_difference:\n num_epochs -= int(training_difference.replace(\"epoch_\", \"\"))\n resume_step = None\n else:\n resume_step = int(training_difference.replace(\"step_\", \"\"))\n num_epochs -= resume_step // len(train_dataloader)\n # If resuming by step, we also need to know exactly how far into the DataLoader we went\n resume_step = (num_epochs * len(train_dataloader)) - resume_step\n\n # Now we train the model\n for epoch in range(num_epochs):\n # New Code #\n # context manager to track the peak memory usage during the training epoch\n with TorchTracemalloc() as tracemalloc:\n model.train()\n if args.with_tracking:\n total_loss = 0\n for step, batch in enumerate(train_dataloader):\n # We need to skip steps until we reach the resumed step\n if args.resume_from_checkpoint and epoch == 0:\n if resume_step is not None and step < resume_step:\n pass\n # We could avoid this line since we set the accelerator with `device_placement=True`.\n batch.to(accelerator.device)\n outputs = model(**batch)\n loss = outputs.loss\n # We keep track of the loss at each epoch\n if args.with_tracking:\n total_loss += loss.detach().float()\n accelerator.backward(loss)\n if step % gradient_accumulation_steps == 0:\n optimizer.step()\n lr_scheduler.step()\n optimizer.zero_grad()\n # accelerator.print(lr_scheduler.get_lr())\n\n overall_step += 1\n\n if isinstance(checkpointing_steps, int):\n output_dir = f\"step_{overall_step}\"\n if overall_step % checkpointing_steps == 0:\n if args.output_dir is not None:\n output_dir = os.path.join(args.output_dir, output_dir)\n accelerator.save_state(output_dir)\n # New Code #\n # Printing the GPU memory usage details such as allocated memory, peak memory, and total memory usage\n accelerator.print(f\"Memory before entering the train : {b2mb(tracemalloc.begin)}\")\n accelerator.print(f\"Memory consumed at the end of the train (end-begin): {tracemalloc.used}\")\n accelerator.print(f\"Peak Memory consumed during the train (max-begin): {tracemalloc.peaked}\")\n accelerator.print(\n f\"Total Peak Memory consumed during the train (max): {tracemalloc.peaked + b2mb(tracemalloc.begin)}\"\n )\n # Logging the peak memory usage of the GPU to the tracker\n if args.with_tracking:\n accelerator.log(\n {\n \"train_total_peak_memory\": tracemalloc.peaked + b2mb(tracemalloc.begin),\n },\n step=epoch,\n )\n\n # New Code #\n # context manager to track the peak memory usage during the evaluation\n with TorchTracemalloc() as tracemalloc:\n model.eval()\n for step, batch in enumerate(eval_dataloader):\n # We could avoid this line since we set the accelerator with `device_placement=True`.\n batch.to(accelerator.device)\n with torch.no_grad():\n outputs = model(**batch)\n predictions = outputs.logits.argmax(dim=-1)\n predictions, references = accelerator.gather_for_metrics((predictions, batch[\"labels\"]))\n metric.add_batch(\n predictions=predictions,\n references=references,\n )\n\n eval_metric = metric.compute()\n # Use accelerator.print to print only on the main process.\n accelerator.print(f\"epoch {epoch}:\", eval_metric)\n if args.with_tracking:\n accelerator.log(\n {\n \"accuracy\": eval_metric[\"accuracy\"],\n \"f1\": eval_metric[\"f1\"],\n \"train_loss\": total_loss.item() / len(train_dataloader),\n },\n step=epoch,\n )\n\n if checkpointing_steps == \"epoch\":\n output_dir = f\"epoch_{epoch}\"\n if args.output_dir is not None:\n output_dir = os.path.join(args.output_dir, output_dir)\n accelerator.save_state(output_dir)\n # New Code #\n # Printing the GPU memory usage details such as allocated memory, peak memory, and total memory usage\n accelerator.print(f\"Memory before entering the eval : {b2mb(tracemalloc.begin)}\")\n accelerator.print(f\"Memory consumed at the end of the eval (end-begin): {tracemalloc.used}\")\n accelerator.print(f\"Peak Memory consumed during the eval (max-begin): {tracemalloc.peaked}\")\n accelerator.print(\n f\"Total Peak Memory consumed during the eval (max): {tracemalloc.peaked + b2mb(tracemalloc.begin)}\"\n )\n # Logging the peak memory usage of the GPU to the tracker\n if args.with_tracking:\n accelerator.log(\n {\n \"eval_total_peak_memory\": tracemalloc.peaked + b2mb(tracemalloc.begin),\n },\n step=epoch,\n )\n\n accelerator.end_training()\n\n\ndef main():\n parser = argparse.ArgumentParser(description=\"Simple example of training script.\")\n parser.add_argument(\n \"--mixed_precision\",\n type=str,\n default=None,\n choices=[\"no\", \"fp16\", \"bf16\", \"fp8\"],\n help=\"Whether to use mixed precision. Choose\"\n \"between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10.\"\n \"and an Nvidia Ampere GPU.\",\n )\n parser.add_argument(\"--cpu\", action=\"store_true\", help=\"If passed, will train on the CPU.\")\n parser.add_argument(\n \"--checkpointing_steps\",\n type=str,\n default=None,\n help=\"Whether the various states should be saved at the end of every n steps, or 'epoch' for each epoch.\",\n )\n parser.add_argument(\n \"--resume_from_checkpoint\",\n type=str,\n default=None,\n help=\"If the training should continue from a checkpoint folder.\",\n )\n parser.add_argument(\n \"--with_tracking\",\n action=\"store_true\",\n help=\"Whether to load in all available experiment trackers from the environment and use them for logging.\",\n )\n parser.add_argument(\n \"--output_dir\",\n type=str,\n default=\".\",\n help=\"Optional save directory where all checkpoint folders will be stored. Default is the current working directory.\",\n )\n parser.add_argument(\n \"--logging_dir\",\n type=str,\n default=\"logs\",\n help=\"Location on where to store experiment tracking logs`\",\n )\n parser.add_argument(\n \"--model_name_or_path\",\n type=str,\n help=\"Path to pretrained model or model identifier from huggingface.co/models.\",\n required=True,\n )\n args = parser.parse_args()\n config = {\"lr\": 2e-5, \"num_epochs\": 3, \"seed\": 42, \"batch_size\": 16}\n training_function(config, args)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "examples/by_feature/gradient_accumulation.py", "content": "# Copyright 2021 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport argparse\nimport os\n\nimport evaluate\nimport torch\nfrom datasets import load_dataset\nfrom torch.optim import AdamW\nfrom torch.utils.data import DataLoader\nfrom transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed\n\nfrom accelerate import Accelerator, DistributedType\n\n\n########################################################################\n# This is a fully working simple example to use Accelerate\n# and perform gradient accumulation\n#\n# This example trains a Bert base model on GLUE MRPC\n# in any of the following settings (with the same script):\n# - single CPU or single GPU\n# - multi GPUS (using PyTorch distributed mode)\n# - (multi) TPUs\n# - fp16 (mixed-precision) or fp32 (normal precision)\n#\n# To run it in each of these various modes, follow the instructions\n# in the readme for examples:\n# https://github.com/huggingface/accelerate/tree/main/examples\n#\n########################################################################\n\n\nMAX_GPU_BATCH_SIZE = 16\nEVAL_BATCH_SIZE = 32\n\n\ndef get_dataloaders(accelerator: Accelerator, batch_size: int = 16):\n \"\"\"\n Creates a set of `DataLoader`s for the `glue` dataset,\n using \"bert-base-cased\" as the tokenizer.\n\n Args:\n accelerator (`Accelerator`):\n An `Accelerator` object\n batch_size (`int`, *optional*):\n The batch size for the train and validation DataLoaders.\n \"\"\"\n tokenizer = AutoTokenizer.from_pretrained(\"bert-base-cased\")\n datasets = load_dataset(\"glue\", \"mrpc\")\n\n def tokenize_function(examples):\n # max_length=None => use the model max length (it's actually the default)\n outputs = tokenizer(examples[\"sentence1\"], examples[\"sentence2\"], truncation=True, max_length=None)\n return outputs\n\n # Apply the method we just defined to all the examples in all the splits of the dataset\n # starting with the main process first:\n with accelerator.main_process_first():\n tokenized_datasets = datasets.map(\n tokenize_function,\n batched=True,\n remove_columns=[\"idx\", \"sentence1\", \"sentence2\"],\n )\n\n # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the\n # transformers library\n tokenized_datasets = tokenized_datasets.rename_column(\"label\", \"labels\")\n\n def collate_fn(examples):\n # On TPU it's best to pad everything to the same length or training will be very slow.\n max_length = 128 if accelerator.distributed_type == DistributedType.XLA else None\n # When using mixed precision we want round multiples of 8/16\n if accelerator.mixed_precision == \"fp8\":\n pad_to_multiple_of = 16\n elif accelerator.mixed_precision != \"no\":\n pad_to_multiple_of = 8\n else:\n pad_to_multiple_of = None\n\n return tokenizer.pad(\n examples,\n padding=\"longest\",\n max_length=max_length,\n pad_to_multiple_of=pad_to_multiple_of,\n return_tensors=\"pt\",\n )\n\n # Instantiate dataloaders.\n train_dataloader = DataLoader(\n tokenized_datasets[\"train\"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size\n )\n eval_dataloader = DataLoader(\n tokenized_datasets[\"validation\"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE\n )\n\n return train_dataloader, eval_dataloader\n\n\n# For testing only\nif os.environ.get(\"TESTING_MOCKED_DATALOADERS\", None) == \"1\":\n from accelerate.test_utils.training import mocked_dataloaders\n\n get_dataloaders = mocked_dataloaders # noqa: F811\n\n\ndef training_function(config, args):\n # For testing only\n if os.environ.get(\"TESTING_MOCKED_DATALOADERS\", None) == \"1\":\n config[\"num_epochs\"] = 2\n # New Code #\n gradient_accumulation_steps = int(args.gradient_accumulation_steps)\n # Initialize accelerator\n accelerator = Accelerator(\n cpu=args.cpu, mixed_precision=args.mixed_precision, gradient_accumulation_steps=gradient_accumulation_steps\n )\n if accelerator.distributed_type == DistributedType.XLA and gradient_accumulation_steps > 1:\n raise NotImplementedError(\n \"Gradient accumulation on TPUs is currently not supported. Pass `gradient_accumulation_steps=1`\"\n )\n # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs\n lr = config[\"lr\"]\n num_epochs = int(config[\"num_epochs\"])\n seed = int(config[\"seed\"])\n batch_size = int(config[\"batch_size\"])\n\n metric = evaluate.load(\"glue\", \"mrpc\")\n\n set_seed(seed)\n train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size)\n # Instantiate the model (we build the model here so that the seed also control new weights initialization)\n model = AutoModelForSequenceClassification.from_pretrained(\"bert-base-cased\", return_dict=True)\n\n # We could avoid this line since the accelerator is set with `device_placement=True` (default value).\n # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer\n # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).\n model = model.to(accelerator.device)\n\n # Instantiate optimizer\n optimizer = AdamW(params=model.parameters(), lr=lr)\n\n # Instantiate scheduler\n lr_scheduler = get_linear_schedule_with_warmup(\n optimizer=optimizer,\n num_warmup_steps=100,\n num_training_steps=(len(train_dataloader) * num_epochs),\n )\n\n # Prepare everything\n # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the\n # prepare method.\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler\n )\n\n # Now we train the model\n for epoch in range(num_epochs):\n model.train()\n for step, batch in enumerate(train_dataloader):\n # We could avoid this line since we set the accelerator with `device_placement=True`.\n batch.to(accelerator.device)\n # New code #\n # We use the new `accumulate` context manager to perform gradient accumulation\n # We also currently do not support TPUs nor advise it as bugs were found on the XLA side when running our tests.\n with accelerator.accumulate(model):\n output = model(**batch)\n loss = output.loss\n accelerator.backward(loss)\n optimizer.step()\n lr_scheduler.step()\n optimizer.zero_grad()\n\n model.eval()\n for step, batch in enumerate(eval_dataloader):\n # We could avoid this line since we set the accelerator with `device_placement=True`.\n batch.to(accelerator.device)\n with torch.no_grad():\n outputs = model(**batch)\n predictions = outputs.logits.argmax(dim=-1)\n predictions, references = accelerator.gather_for_metrics((predictions, batch[\"labels\"]))\n metric.add_batch(\n predictions=predictions,\n references=references,\n )\n\n eval_metric = metric.compute()\n # Use accelerator.print to print only on the main process.\n accelerator.print(f\"epoch {epoch}:\", eval_metric)\n accelerator.end_training()\n\n\ndef main():\n parser = argparse.ArgumentParser(description=\"Simple example of training script.\")\n parser.add_argument(\n \"--mixed_precision\",\n type=str,\n default=None,\n choices=[\"no\", \"fp16\", \"bf16\", \"fp8\"],\n help=\"Whether to use mixed precision. Choose\"\n \"between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10.\"\n \"and an Nvidia Ampere GPU.\",\n )\n # New Code #\n parser.add_argument(\n \"--gradient_accumulation_steps\",\n type=int,\n default=1,\n help=\"The number of minibatches to be ran before gradients are accumulated.\",\n )\n parser.add_argument(\"--cpu\", action=\"store_true\", help=\"If passed, will train on the CPU.\")\n args = parser.parse_args()\n config = {\"lr\": 2e-5, \"num_epochs\": 3, \"seed\": 42, \"batch_size\": 16}\n training_function(config, args)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "examples/by_feature/gradient_accumulation_for_autoregressive_models.py", "content": "# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport argparse\nimport contextlib\nimport math\nimport os\n\nimport torch\nfrom datasets import load_dataset\nfrom torch.optim import AdamW\nfrom torch.utils.data import DataLoader\nfrom transformers import AutoModelForCausalLM, AutoTokenizer, get_constant_schedule, set_seed\n\nfrom accelerate import Accelerator, DistributedType\n\n\n########################################################################\n# This is a fully working simple example to use Accelerate\n# and perform gradient accumulation on samples of variable size\n#\n# This example trains a SmolLM base model on WikiText-2 v1\n# in any of the following settings (with the same script):\n# - single CPU or single GPU\n# - multi GPUS (using PyTorch distributed mode)\n# - (multi) TPUs\n# - fp16 (mixed-precision) or fp32 (normal precision)\n#\n# To run it in each of these various modes, follow the instructions\n# in the readme for examples:\n# https://github.com/huggingface/accelerate/tree/main/examples\n#\n########################################################################\n\n\nEVAL_BATCH_SIZE = 32\n\n\ndef get_dataloaders(accelerator: Accelerator, batch_size: int = 16, max_training_samples=500):\n \"\"\"\n Creates a set of `DataLoader`s for the `Salesforce/wikitext` dataset,\n using \"HuggingFaceTB/SmolLM-360M\" as the tokenizer.\n\n Args:\n accelerator (`Accelerator`):\n An `Accelerator` object\n batch_size (`int`, *optional*):\n The batch size for the train and validation DataLoaders.\n \"\"\"\n tokenizer = AutoTokenizer.from_pretrained(\"HuggingFaceTB/SmolLM-360M\")\n tokenizer.pad_token = tokenizer.eos_token\n with accelerator.local_main_process_first():\n datasets = load_dataset(\"Salesforce/wikitext\", \"wikitext-2-v1\")\n datasets[\"train\"] = datasets[\"train\"].select(range(max_training_samples))\n\n def tokenize_function(examples):\n # max_length=None => use the model max length (it's actually the default)\n outputs = tokenizer(examples[\"text\"], truncation=True, max_length=None, return_attention_mask=False)\n return outputs\n\n # Filter out empty texts\n with accelerator.main_process_first():\n datasets = datasets.filter(\n lambda x: len(x) > 0,\n input_columns=\"text\",\n )\n\n # Apply the method we just defined to all the examples in all the splits of the dataset\n # starting with the main process first:\n with accelerator.main_process_first():\n tokenized_datasets = datasets.map(\n tokenize_function,\n batched=True,\n remove_columns=[\"text\"],\n )\n\n # Filter out empty samples\n with accelerator.main_process_first():\n tokenized_datasets = tokenized_datasets.filter(\n lambda x: len(x) > 0,\n input_columns=\"input_ids\",\n )\n\n def collate_fn(examples):\n # On TPU it's best to pad everything to the same length or training will be very slow.\n max_length = (\n 128\n if accelerator.distributed_type == DistributedType.XLA\n else max([len(e[\"input_ids\"]) for e in examples])\n )\n # When using mixed precision we want round multiples of 8/16\n if accelerator.mixed_precision == \"fp8\":\n pad_to_multiple_of = 16\n elif accelerator.mixed_precision != \"no\":\n pad_to_multiple_of = 8\n else:\n pad_to_multiple_of = None\n\n batch = tokenizer.pad(\n examples,\n padding=\"max_length\",\n max_length=max_length + 1,\n pad_to_multiple_of=pad_to_multiple_of,\n return_tensors=\"pt\",\n )\n\n batch[\"labels\"] = batch[\"input_ids\"][:, 1:]\n batch[\"input_ids\"] = batch[\"input_ids\"][:, :-1]\n if \"attention_mask\" in batch:\n batch[\"attention_mask\"] = batch[\"attention_mask\"][:, :-1]\n\n batch[\"labels\"] = torch.where(batch[\"labels\"] == tokenizer.pad_token_id, -100, batch[\"labels\"])\n\n return batch\n\n # Instantiate dataloaders.\n train_dataloader = DataLoader(\n tokenized_datasets[\"train\"], shuffle=False, collate_fn=collate_fn, batch_size=batch_size\n )\n eval_dataloader = DataLoader(\n tokenized_datasets[\"validation\"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE\n )\n\n return train_dataloader, eval_dataloader\n\n\n# For testing only\nif os.environ.get(\"TESTING_MOCKED_DATALOADERS\", None) == \"1\":\n from accelerate.test_utils.training import mocked_dataloaders_for_autoregressive_models\n\n get_dataloaders = mocked_dataloaders_for_autoregressive_models # noqa: F811\n\n\ndef training_function(config, args):\n # For testing only\n if os.environ.get(\"TESTING_MOCKED_DATALOADERS\", None) == \"1\":\n config[\"num_epochs\"] = 2\n\n gradient_accumulation_steps = int(args.gradient_accumulation_steps)\n # Initialize accelerator\n if args.with_wandb_tracking:\n accelerator = Accelerator(\n cpu=args.cpu,\n mixed_precision=args.mixed_precision,\n gradient_accumulation_steps=gradient_accumulation_steps,\n log_with=\"wandb\",\n )\n else:\n accelerator = Accelerator(\n cpu=args.cpu, mixed_precision=args.mixed_precision, gradient_accumulation_steps=gradient_accumulation_steps\n )\n if accelerator.distributed_type == DistributedType.XLA and gradient_accumulation_steps > 1:\n raise NotImplementedError(\n \"Gradient accumulation on TPUs is currently not supported. Pass `gradient_accumulation_steps=1`\"\n )\n # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs\n lr = config[\"lr\"]\n num_epochs = int(config[\"num_epochs\"])\n seed = int(config[\"seed\"])\n batch_size = int(config[\"batch_size\"])\n max_grad_norm = config[\"max_grad_norm\"]\n\n # We need to initialize the trackers we use, and also store our configuration\n if args.with_wandb_tracking:\n run = os.path.split(__file__)[-1].split(\".\")[0]\n run_name = f\"{accelerator.num_processes}GPU-grad{gradient_accumulation_steps}-bs{batch_size}\"\n accelerator.init_trackers(\n run,\n config,\n init_kwargs={\"wandb\": {\"name\": run_name}},\n )\n\n set_seed(seed)\n train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size)\n # Instantiate the model (we build the model here so that the seed also control new weights initialization)\n model = AutoModelForCausalLM.from_pretrained(\"HuggingFaceTB/SmolLM-360M\")\n\n # We could avoid this line since the accelerator is set with `device_placement=True` (default value).\n # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer\n # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).\n model = model.to(accelerator.device)\n\n # Instantiate optimizer\n optimizer = AdamW(params=model.parameters(), lr=lr)\n\n # Instantiate scheduler\n lr_scheduler = get_constant_schedule(\n optimizer=optimizer,\n )\n\n # Prepare everything\n # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the\n # prepare method.\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler\n )\n\n num_samples_in_epoch = len(train_dataloader)\n remainder = num_samples_in_epoch % gradient_accumulation_steps\n remainder = remainder if remainder != 0 else gradient_accumulation_steps\n total_gradient_updates = math.ceil(num_samples_in_epoch / gradient_accumulation_steps)\n\n total_batched_samples = 0\n # Now we train the model\n for epoch in range(num_epochs):\n model.train()\n training_iterator = iter(train_dataloader)\n for update_step in range(total_gradient_updates):\n # In order to correctly the total number of non-padded tokens on which we'll compute the cross-entropy loss\n # we need to pre-load the full local batch - i.e the next per_device_batch_size * accumulation_steps samples\n batch_samples = []\n num_batches_in_step = (\n gradient_accumulation_steps if update_step != (total_gradient_updates - 1) else remainder\n )\n for _ in range(num_batches_in_step):\n batch_samples += [next(training_iterator)]\n # get local num items in batch\n local_num_items_in_batch = sum([(batch[\"labels\"].ne(-100)).sum() for batch in batch_samples])\n\n # to compute it correctly in a multi-device DDP training, we need to gather the total number of items in the full batch.\n num_items_in_batch = accelerator.gather(local_num_items_in_batch).sum().item()\n losses = []\n for i, batch in enumerate(batch_samples):\n # if we perform gradient accumulation in a multi-devices set-up, we want to avoid unecessary communications when accumulating\n # cf: https://muellerzr.github.io/blog/gradient_accumulation.html\n ctx = (\n model.no_sync\n if (i < len(batch_samples) - 1 and accelerator.num_processes > 1)\n else contextlib.nullcontext\n )\n with ctx():\n total_batched_samples += 1\n\n outputs = model(**batch, use_cache=False, num_items_in_batch=num_items_in_batch)\n loss = outputs.loss\n\n # We multiply by num_processes because the DDP calculates the average gradient across all devices whereas dividing by num_items_in_batch already takes into account all devices\n # Same reason for gradient_accumulation_steps, but this times it's Accelerate that calculate the average gradient across the accumulated steps\n # Because the loss is already divided by `num_items_in_batch` in the `transformers` code, we don't need to do it again\n loss = loss * gradient_accumulation_steps * accelerator.num_processes\n accelerator.backward(loss)\n losses.append(loss.detach())\n\n # Sync gradients and perform optimization steps once every gradient_accumulation_steps\n grad_norm = accelerator.clip_grad_norm_(model.parameters(), max_grad_norm)\n optimizer.step()\n lr_scheduler.step()\n optimizer.zero_grad()\n\n losses = accelerator.gather(sum(losses)).sum().item() / (\n accelerator.num_processes * gradient_accumulation_steps\n )\n\n grad_norm = grad_norm.detach().item() if isinstance(grad_norm, torch.Tensor) else grad_norm\n accelerator.print(\n f\"epoch {epoch} - update step {update_step}:: grad norm: {grad_norm} ::train loss: {losses}\"\n )\n if args.with_wandb_tracking:\n accelerator.log(\n {\n \"train/grad_norm\": grad_norm,\n \"train/epoch\": epoch,\n \"train/loss\": losses,\n },\n step=update_step + total_gradient_updates * epoch,\n )\n model.eval()\n losses = []\n for step, batch in enumerate(eval_dataloader):\n with torch.no_grad():\n outputs = model(**batch, use_cache=False)\n eval_loss = outputs.loss\n losses.append(accelerator.gather_for_metrics(loss.repeat(EVAL_BATCH_SIZE)))\n\n losses = torch.cat(losses)\n try:\n eval_loss = torch.mean(losses)\n perplexity = math.exp(eval_loss)\n except OverflowError:\n perplexity = float(\"inf\")\n\n # Use accelerator.print to print only on the main process.\n accelerator.print(f\"epoch {epoch}:: eval perplexity: {perplexity} eval_loss: {eval_loss}\")\n if args.with_wandb_tracking:\n accelerator.log(\n {\n \"eval/perplexity\": perplexity,\n \"eval/loss\": eval_loss,\n \"eval/epoch\": epoch,\n },\n step=update_step + total_gradient_updates * epoch,\n )\n accelerator.end_training()\n\n\ndef main():\n parser = argparse.ArgumentParser(description=\"Simple example of training script.\")\n parser.add_argument(\n \"--mixed_precision\",\n type=str,\n default=None,\n choices=[\"no\", \"fp16\", \"bf16\", \"fp8\"],\n help=\"Whether to use mixed precision. Choose\"\n \"between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10.\"\n \"and an Nvidia Ampere GPU.\",\n )\n\n parser.add_argument(\n \"--gradient_accumulation_steps\",\n type=int,\n default=1,\n help=\"The number of minibatches to be ran before gradients are accumulated.\",\n )\n parser.add_argument(\n \"--per_device_batch_size\",\n type=int,\n default=2,\n help=\"The size of each minibatch\",\n )\n\n parser.add_argument(\"--cpu\", action=\"store_true\", help=\"If passed, will train on the CPU.\")\n parser.add_argument(\n \"--with_wandb_tracking\",\n action=\"store_true\",\n help=\"Whether to load in wandb from the environment and use them for logging.\",\n )\n args = parser.parse_args()\n config = {\"lr\": 2e-5, \"num_epochs\": 3, \"seed\": 42, \"batch_size\": args.per_device_batch_size, \"max_grad_norm\": 1.0}\n training_function(config, args)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "examples/by_feature/local_sgd.py", "content": "# Copyright 2023 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport argparse\nimport os\n\nimport evaluate\nimport torch\nfrom datasets import load_dataset\nfrom torch.optim import AdamW\nfrom torch.utils.data import DataLoader\nfrom transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed\n\nfrom accelerate import Accelerator, DistributedType\nfrom accelerate.local_sgd import LocalSGD\n\n\n########################################################################\n# This is a fully working simple example to use Accelerate\n# with LocalSGD, which is a method to synchronize model\n# parameters every K batches. It is different, but complementary\n# to gradient accumulation.\n#\n# This example trains a Bert base model on GLUE MRPC\n# in any of the following settings (with the same script):\n# - single CPU or single GPU\n# - multi GPUS (using PyTorch distributed mode)\n# - (multi) TPUs\n# - fp16 (mixed-precision) or fp32 (normal precision)\n#\n# To run it in each of these various modes, follow the instructions\n# in the readme for examples:\n# https://github.com/huggingface/accelerate/tree/main/examples\n#\n########################################################################\n\n\nMAX_GPU_BATCH_SIZE = 16\nEVAL_BATCH_SIZE = 32\n\n\ndef get_dataloaders(accelerator: Accelerator, batch_size: int = 16):\n \"\"\"\n Creates a set of `DataLoader`s for the `glue` dataset,\n using \"bert-base-cased\" as the tokenizer.\n\n Args:\n accelerator (`Accelerator`):\n An `Accelerator` object\n batch_size (`int`, *optional*):\n The batch size for the train and validation DataLoaders.\n \"\"\"\n tokenizer = AutoTokenizer.from_pretrained(\"bert-base-cased\")\n datasets = load_dataset(\"glue\", \"mrpc\")\n\n def tokenize_function(examples):\n # max_length=None => use the model max length (it's actually the default)\n outputs = tokenizer(examples[\"sentence1\"], examples[\"sentence2\"], truncation=True, max_length=None)\n return outputs\n\n # Apply the method we just defined to all the examples in all the splits of the dataset\n # starting with the main process first:\n with accelerator.main_process_first():\n tokenized_datasets = datasets.map(\n tokenize_function,\n batched=True,\n remove_columns=[\"idx\", \"sentence1\", \"sentence2\"],\n )\n\n # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the\n # transformers library\n tokenized_datasets = tokenized_datasets.rename_column(\"label\", \"labels\")\n\n def collate_fn(examples):\n # On TPU it's best to pad everything to the same length or training will be very slow.\n max_length = 128 if accelerator.distributed_type == DistributedType.XLA else None\n # When using mixed precision we want round multiples of 8/16\n if accelerator.mixed_precision == \"fp8\":\n pad_to_multiple_of = 16\n elif accelerator.mixed_precision != \"no\":\n pad_to_multiple_of = 8\n else:\n pad_to_multiple_of = None\n\n return tokenizer.pad(\n examples,\n padding=\"longest\",\n max_length=max_length,\n pad_to_multiple_of=pad_to_multiple_of,\n return_tensors=\"pt\",\n )\n\n # Instantiate dataloaders.\n train_dataloader = DataLoader(\n tokenized_datasets[\"train\"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size\n )\n eval_dataloader = DataLoader(\n tokenized_datasets[\"validation\"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE\n )\n\n return train_dataloader, eval_dataloader\n\n\n# For testing only\nif os.environ.get(\"TESTING_MOCKED_DATALOADERS\", None) == \"1\":\n from accelerate.test_utils.training import mocked_dataloaders\n\n get_dataloaders = mocked_dataloaders # noqa: F811\n\n\ndef training_function(config, args):\n # For testing only\n if os.environ.get(\"TESTING_MOCKED_DATALOADERS\", None) == \"1\":\n config[\"num_epochs\"] = 2\n # New Code #\n gradient_accumulation_steps = int(args.gradient_accumulation_steps)\n local_sgd_steps = int(args.local_sgd_steps)\n # Initialize accelerator\n accelerator = Accelerator(\n cpu=args.cpu, mixed_precision=args.mixed_precision, gradient_accumulation_steps=gradient_accumulation_steps\n )\n # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs\n lr = config[\"lr\"]\n num_epochs = int(config[\"num_epochs\"])\n seed = int(config[\"seed\"])\n batch_size = int(config[\"batch_size\"])\n\n metric = evaluate.load(\"glue\", \"mrpc\")\n\n set_seed(seed)\n train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size)\n # Instantiate the model (we build the model here so that the seed also control new weights initialization)\n model = AutoModelForSequenceClassification.from_pretrained(\"bert-base-cased\", return_dict=True)\n\n # We could avoid this line since the accelerator is set with `device_placement=True` (default value).\n # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer\n # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).\n model = model.to(accelerator.device)\n\n # Instantiate optimizer\n optimizer = AdamW(params=model.parameters(), lr=lr)\n\n # Instantiate scheduler\n lr_scheduler = get_linear_schedule_with_warmup(\n optimizer=optimizer,\n num_warmup_steps=100,\n num_training_steps=(len(train_dataloader) * num_epochs),\n )\n\n # Prepare everything\n # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the\n # prepare method.\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler\n )\n\n # Now we train the model\n for epoch in range(num_epochs):\n model.train()\n with LocalSGD(\n accelerator=accelerator, model=model, local_sgd_steps=local_sgd_steps, enabled=local_sgd_steps is not None\n ) as local_sgd:\n for step, batch in enumerate(train_dataloader):\n # We could avoid this line since we set the accelerator with `device_placement=True`.\n batch.to(accelerator.device)\n # New code #\n # We use the new `accumulate` context manager to perform gradient accumulation\n # We also currently do not support TPUs nor advise it as bugs were found on the XLA side when running our tests.\n with accelerator.accumulate(model):\n output = model(**batch)\n loss = output.loss\n accelerator.backward(loss)\n optimizer.step()\n lr_scheduler.step()\n optimizer.zero_grad()\n # LocalSGD-specific line\n local_sgd.step()\n\n model.eval()\n for step, batch in enumerate(eval_dataloader):\n # We could avoid this line since we set the accelerator with `device_placement=True`.\n batch.to(accelerator.device)\n with torch.no_grad():\n outputs = model(**batch)\n predictions = outputs.logits.argmax(dim=-1)\n predictions, references = accelerator.gather_for_metrics((predictions, batch[\"labels\"]))\n metric.add_batch(\n predictions=predictions,\n references=references,\n )\n\n eval_metric = metric.compute()\n # Use accelerator.print to print only on the main process.\n accelerator.print(f\"epoch {epoch}:\", eval_metric)\n accelerator.end_training()\n\n\ndef main():\n parser = argparse.ArgumentParser(description=\"Simple example of training script.\")\n parser.add_argument(\n \"--mixed_precision\",\n type=str,\n default=None,\n choices=[\"no\", \"fp16\", \"bf16\", \"fp8\"],\n help=\"Whether to use mixed precision. Choose\"\n \"between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10.\"\n \"and an Nvidia Ampere GPU.\",\n )\n # New Code #\n parser.add_argument(\n \"--gradient_accumulation_steps\",\n type=int,\n default=1,\n help=\"The number of minibatches to be ran before gradients are accumulated.\",\n )\n parser.add_argument(\n \"--local_sgd_steps\", type=int, default=8, help=\"Number of local SGD steps or None to disable local SGD\"\n )\n parser.add_argument(\"--cpu\", action=\"store_true\", help=\"If passed, will train on the CPU.\")\n args = parser.parse_args()\n config = {\"lr\": 2e-5, \"num_epochs\": 3, \"seed\": 42, \"batch_size\": 16}\n training_function(config, args)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "examples/by_feature/megatron_lm_gpt_pretraining.py", "content": "#!/usr/bin/env python\n# Copyright 2021 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nFine-tuning the library models for causal language modeling (GPT, GPT-2, CTRL, ...)\non a text file or a dataset without using HuggingFace Trainer.\n\nHere is the full list of checkpoints on the hub that can be fine-tuned by this script:\nhttps://huggingface.co/models?filter=text-generation\n\"\"\"\n# You can also adapt this script on your own causal language modeling task. Pointers for this are left as comments.\n\nimport argparse\nimport json\nimport logging\nimport math\nimport os\nimport random\nfrom itertools import chain\nfrom pathlib import Path\n\nimport datasets\nimport torch\nimport transformers\nfrom datasets import load_dataset\nfrom huggingface_hub import HfApi\nfrom torch.utils.data import DataLoader\nfrom tqdm.auto import tqdm\nfrom transformers import (\n CONFIG_MAPPING,\n MODEL_MAPPING,\n AutoConfig,\n AutoModelForCausalLM,\n AutoTokenizer,\n SchedulerType,\n default_data_collator,\n get_scheduler,\n)\nfrom transformers.utils import check_min_version, send_example_telemetry\nfrom transformers.utils.versions import require_version\n\nfrom accelerate import Accelerator, DistributedType, init_empty_weights\nfrom accelerate.logging import get_logger\nfrom accelerate.utils import MegatronLMDummyScheduler, set_seed\n\n\n# Will error if the minimal version of Transformers is not installed. Remove at your own risks.\ncheck_min_version(\"4.23.0.dev0\")\n\nlogger = get_logger(__name__)\n\nrequire_version(\"datasets>=1.8.0\", \"To fix: pip install -r examples/pytorch/language-modeling/requirements.txt\")\n\nMODEL_CONFIG_CLASSES = list(MODEL_MAPPING.keys())\nMODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)\n\n\ndef parse_args():\n parser = argparse.ArgumentParser(description=\"Finetune a transformers model on a causal language modeling task\")\n parser.add_argument(\n \"--dataset_name\",\n type=str,\n default=None,\n help=\"The name of the dataset to use (via the datasets library).\",\n )\n parser.add_argument(\n \"--dataset_config_name\",\n type=str,\n default=None,\n help=\"The configuration name of the dataset to use (via the datasets library).\",\n )\n parser.add_argument(\n \"--train_file\", type=str, default=None, help=\"A csv or a json file containing the training data.\"\n )\n parser.add_argument(\n \"--validation_file\", type=str, default=None, help=\"A csv or a json file containing the validation data.\"\n )\n parser.add_argument(\n \"--validation_split_percentage\",\n default=5,\n help=\"The percentage of the train set used as validation set in case there's no validation split\",\n )\n parser.add_argument(\n \"--model_name_or_path\",\n type=str,\n help=\"Path to pretrained model or model identifier from huggingface.co/models.\",\n required=False,\n )\n parser.add_argument(\n \"--config_name\",\n type=str,\n default=None,\n help=\"Pretrained config name or path if not the same as model_name\",\n )\n parser.add_argument(\n \"--tokenizer_name\",\n type=str,\n default=None,\n help=\"Pretrained tokenizer name or path if not the same as model_name\",\n )\n parser.add_argument(\n \"--use_slow_tokenizer\",\n action=\"store_true\",\n help=\"If passed, will use a slow tokenizer (not backed by the šŸ¤— Tokenizers library).\",\n )\n parser.add_argument(\n \"--per_device_train_batch_size\",\n type=int,\n default=8,\n help=\"Batch size (per device) for the training dataloader.\",\n )\n parser.add_argument(\n \"--per_device_eval_batch_size\",\n type=int,\n default=8,\n help=\"Batch size (per device) for the evaluation dataloader.\",\n )\n parser.add_argument(\n \"--learning_rate\",\n type=float,\n default=5e-5,\n help=\"Initial learning rate (after the potential warmup period) to use.\",\n )\n parser.add_argument(\"--weight_decay\", type=float, default=0.0, help=\"Weight decay to use.\")\n parser.add_argument(\"--num_train_epochs\", type=int, default=3, help=\"Total number of training epochs to perform.\")\n parser.add_argument(\n \"--max_train_steps\",\n type=int,\n default=None,\n help=\"Total number of training steps to perform. If provided, overrides num_train_epochs.\",\n )\n parser.add_argument(\n \"--gradient_accumulation_steps\",\n type=int,\n default=1,\n help=\"Number of updates steps to accumulate before performing a backward/update pass.\",\n )\n parser.add_argument(\n \"--lr_scheduler_type\",\n type=SchedulerType,\n default=\"linear\",\n help=\"The scheduler type to use.\",\n choices=[\"linear\", \"cosine\", \"cosine_with_restarts\", \"polynomial\", \"constant\", \"constant_with_warmup\"],\n )\n parser.add_argument(\n \"--num_warmup_steps\", type=int, default=0, help=\"Number of steps for the warmup in the lr scheduler.\"\n )\n parser.add_argument(\"--output_dir\", type=str, default=None, help=\"Where to store the final model.\")\n parser.add_argument(\"--seed\", type=int, default=None, help=\"A seed for reproducible training.\")\n parser.add_argument(\n \"--model_type\",\n type=str,\n default=None,\n help=\"Model type to use if training from scratch.\",\n choices=MODEL_TYPES,\n )\n parser.add_argument(\n \"--block_size\",\n type=int,\n default=None,\n help=(\n \"Optional input sequence length after tokenization. The training dataset will be truncated in block of\"\n \" this size for training. Default to the model max input length for single sentence inputs (take into\"\n \" account special tokens).\"\n ),\n )\n parser.add_argument(\n \"--preprocessing_num_workers\",\n type=int,\n default=None,\n help=\"The number of processes to use for the preprocessing.\",\n )\n parser.add_argument(\n \"--overwrite_cache\", action=\"store_true\", help=\"Overwrite the cached training and evaluation sets\"\n )\n parser.add_argument(\n \"--no_keep_linebreaks\", action=\"store_true\", help=\"Do not keep line breaks when using TXT files.\"\n )\n parser.add_argument(\"--push_to_hub\", action=\"store_true\", help=\"Whether or not to push the model to the Hub.\")\n parser.add_argument(\n \"--hub_model_id\", type=str, help=\"The name of the repository to keep in sync with the local `output_dir`.\"\n )\n parser.add_argument(\"--hub_token\", type=str, help=\"The token to use to push to the Model Hub.\")\n parser.add_argument(\n \"--checkpointing_steps\",\n type=str,\n default=None,\n help=\"Whether the various states should be saved at the end of every n steps, or 'epoch' for each epoch.\",\n )\n parser.add_argument(\n \"--resume_from_checkpoint\",\n type=str,\n default=None,\n help=\"If the training should continue from a checkpoint folder.\",\n )\n parser.add_argument(\n \"--initial_megatron_lm_checkpoint\",\n type=str,\n default=None,\n help=\"If the training should start from a Megatron-LM checkpoint.\",\n )\n parser.add_argument(\n \"--with_tracking\",\n action=\"store_true\",\n help=\"Whether to enable experiment trackers for logging.\",\n )\n parser.add_argument(\n \"--report_to\",\n type=str,\n default=\"all\",\n help=(\n 'The integration to report the results and logs to. Supported platforms are `\"tensorboard\"`,'\n ' `\"wandb\"`, `\"comet_ml\"`, and `\"dvclive\"`, and `\"swanlab\"`. Use `\"all\"` (default) to report to all integrations.'\n \"Only applicable when `--with_tracking` is passed.\"\n ),\n )\n args = parser.parse_args()\n\n # Sanity checks\n if args.dataset_name is None and args.train_file is None and args.validation_file is None:\n raise ValueError(\"Need either a dataset name or a training/validation file.\")\n else:\n if args.train_file is not None:\n extension = args.train_file.split(\".\")[-1]\n assert extension in [\"csv\", \"json\", \"txt\"], \"`train_file` should be a csv, json or txt file.\"\n if args.validation_file is not None:\n extension = args.validation_file.split(\".\")[-1]\n assert extension in [\"csv\", \"json\", \"txt\"], \"`validation_file` should be a csv, json or txt file.\"\n\n if args.push_to_hub:\n assert args.output_dir is not None, \"Need an `output_dir` to create a repo when `--push_to_hub` is passed.\"\n\n return args\n\n\ndef main():\n args = parse_args()\n\n # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The\n # information sent is the one passed as arguments along with your Python/PyTorch versions.\n send_example_telemetry(\"run_clm_no_trainer\", args)\n\n # Initialize the accelerator. We will let the accelerator handle device placement for us in this example.\n # If we're using tracking, we also need to initialize it here and it will by default pick up all supported trackers\n # in the environment\n accelerator_log_kwargs = {}\n\n if args.with_tracking:\n accelerator_log_kwargs[\"log_with\"] = args.report_to\n accelerator_log_kwargs[\"project_dir\"] = args.output_dir\n\n accelerator = Accelerator(gradient_accumulation_steps=args.gradient_accumulation_steps, **accelerator_log_kwargs)\n\n # Make one log on every process with the configuration for debugging.\n logging.basicConfig(\n format=\"%(asctime)s - %(levelname)s - %(name)s - %(message)s\",\n datefmt=\"%m/%d/%Y %H:%M:%S\",\n level=logging.INFO,\n )\n logger.info(accelerator.state, main_process_only=False)\n if accelerator.is_local_main_process:\n datasets.utils.logging.set_verbosity_warning()\n transformers.utils.logging.set_verbosity_info()\n else:\n datasets.utils.logging.set_verbosity_error()\n transformers.utils.logging.set_verbosity_error()\n\n # If passed along, set the training seed now.\n if args.seed is not None:\n set_seed(args.seed)\n\n # Handle the repository creation\n if accelerator.is_main_process:\n if args.push_to_hub:\n api = HfApi(token=args.hub_token)\n\n # Create repo (repo_name from args or inferred)\n repo_name = args.hub_model_id\n if repo_name is None:\n repo_name = Path(args.output_dir).absolute().name\n repo_id = api.create_repo(repo_name, exist_ok=True).repo_id\n\n with open(os.path.join(args.output_dir, \".gitignore\"), \"w+\") as gitignore:\n if \"step_*\" not in gitignore:\n gitignore.write(\"step_*\\n\")\n if \"epoch_*\" not in gitignore:\n gitignore.write(\"epoch_*\\n\")\n elif args.output_dir is not None:\n os.makedirs(args.output_dir, exist_ok=True)\n accelerator.wait_for_everyone()\n\n # Get the datasets: you can either provide your own CSV/JSON/TXT training and evaluation files (see below)\n # or just provide the name of one of the public datasets available on the hub at https://huggingface.co/datasets/\n # (the dataset will be downloaded automatically from the datasets Hub).\n #\n # For CSV/JSON files, this script will use the column called 'text' or the first column if no column called\n # 'text' is found. You can easily tweak this behavior (see below).\n #\n # In distributed training, the load_dataset function guarantee that only one local process can concurrently\n # download the dataset.\n if args.dataset_name is not None:\n # Downloading and loading a dataset from the hub.\n raw_datasets = load_dataset(args.dataset_name, args.dataset_config_name)\n if \"validation\" not in raw_datasets.keys():\n raw_datasets[\"validation\"] = load_dataset(\n args.dataset_name,\n args.dataset_config_name,\n split=f\"train[:{args.validation_split_percentage}%]\",\n )\n raw_datasets[\"train\"] = load_dataset(\n args.dataset_name,\n args.dataset_config_name,\n split=f\"train[{args.validation_split_percentage}%:]\",\n )\n else:\n data_files = {}\n dataset_args = {}\n if args.train_file is not None:\n data_files[\"train\"] = args.train_file\n if args.validation_file is not None:\n data_files[\"validation\"] = args.validation_file\n extension = args.train_file.split(\".\")[-1]\n if extension == \"txt\":\n extension = \"text\"\n dataset_args[\"keep_linebreaks\"] = not args.no_keep_linebreaks\n raw_datasets = load_dataset(extension, data_files=data_files, **dataset_args)\n # If no validation data is there, validation_split_percentage will be used to divide the dataset.\n if \"validation\" not in raw_datasets.keys():\n raw_datasets[\"validation\"] = load_dataset(\n extension,\n data_files=data_files,\n split=f\"train[:{args.validation_split_percentage}%]\",\n **dataset_args,\n )\n raw_datasets[\"train\"] = load_dataset(\n extension,\n data_files=data_files,\n split=f\"train[{args.validation_split_percentage}%:]\",\n **dataset_args,\n )\n\n # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at\n # https://huggingface.co/docs/datasets/loading_datasets.html.\n\n # Load pretrained model and tokenizer\n #\n # In distributed training, the .from_pretrained methods guarantee that only one local process can concurrently\n # download model & vocab.\n if args.config_name:\n config = AutoConfig.from_pretrained(args.config_name)\n elif args.model_name_or_path:\n config = AutoConfig.from_pretrained(args.model_name_or_path)\n else:\n config = CONFIG_MAPPING[args.model_type]()\n logger.warning(\"You are instantiating a new config instance from scratch.\")\n\n if args.tokenizer_name:\n tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_name, use_fast=not args.use_slow_tokenizer)\n elif args.model_name_or_path:\n tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path, use_fast=not args.use_slow_tokenizer)\n else:\n raise ValueError(\n \"You are instantiating a new tokenizer from scratch. This is not supported by this script.\"\n \"You can do it from another script, save it, and load it from here, using --tokenizer_name.\"\n )\n\n if args.model_name_or_path:\n # if we are using Megatron-LM, we can use init_empty_weights to load the model without initializing the weights\n # since the weights are loaded later.\n if args.resume_from_checkpoint is not None or args.initial_megatron_lm_checkpoint is not None:\n assert config is not None, \"config should not be None for Megatron-LM\"\n with init_empty_weights():\n model = AutoModelForCausalLM.from_config(config)\n else:\n model = AutoModelForCausalLM.from_pretrained(\n args.model_name_or_path,\n from_tf=bool(\".ckpt\" in args.model_name_or_path),\n config=config,\n )\n else:\n logger.info(\"Training new model from scratch\")\n model = AutoModelForCausalLM.from_config(config)\n\n model.resize_token_embeddings(len(tokenizer))\n\n # Preprocessing the datasets.\n # First we tokenize all the texts.\n column_names = raw_datasets[\"train\"].column_names\n text_column_name = \"text\" if \"text\" in column_names else column_names[0]\n\n def tokenize_function(examples):\n return tokenizer(examples[text_column_name])\n\n with accelerator.main_process_first():\n tokenized_datasets = raw_datasets.map(\n tokenize_function,\n batched=True,\n num_proc=args.preprocessing_num_workers,\n remove_columns=column_names,\n load_from_cache_file=not args.overwrite_cache,\n desc=\"Running tokenizer on dataset\",\n )\n\n if args.block_size is None:\n block_size = tokenizer.model_max_length\n if block_size > 1024:\n logger.warning(\n f\"The tokenizer picked seems to have a very large `model_max_length` ({tokenizer.model_max_length}). \"\n \"Picking 1024 instead. You can change that default value by passing --block_size xxx.\"\n )\n block_size = 1024\n else:\n if args.block_size > tokenizer.model_max_length:\n logger.warning(\n f\"The block_size passed ({args.block_size}) is larger than the maximum length for the model\"\n f\"({tokenizer.model_max_length}). Using block_size={tokenizer.model_max_length}.\"\n )\n block_size = min(args.block_size, tokenizer.model_max_length)\n\n # Main data processing function that will concatenate all texts from our dataset and generate chunks of block_size.\n def group_texts(examples):\n # Concatenate all texts.\n concatenated_examples = {k: list(chain(*examples[k])) for k in examples.keys()}\n total_length = len(concatenated_examples[list(examples.keys())[0]])\n # We drop the small remainder, we could add padding if the model supported it instead of this drop, you can\n # customize this part to your needs.\n if total_length >= block_size:\n total_length = (total_length // block_size) * block_size\n # Split by chunks of max_len.\n result = {\n k: [t[i : i + block_size] for i in range(0, total_length, block_size)]\n for k, t in concatenated_examples.items()\n }\n result[\"labels\"] = result[\"input_ids\"].copy()\n return result\n\n # Note that with `batched=True`, this map processes 1,000 texts together, so group_texts throws away a remainder\n # for each of those groups of 1,000 texts. You can adjust that batch_size here but a higher value might be slower\n # to preprocess.\n #\n # To speed up this part, we use multiprocessing. See the documentation of the map method for more information:\n # https://huggingface.co/docs/datasets/package_reference/main_classes.html#datasets.Dataset.map\n\n with accelerator.main_process_first():\n lm_datasets = tokenized_datasets.map(\n group_texts,\n batched=True,\n num_proc=args.preprocessing_num_workers,\n load_from_cache_file=not args.overwrite_cache,\n desc=f\"Grouping texts in chunks of {block_size}\",\n )\n\n train_dataset = lm_datasets[\"train\"]\n eval_dataset = lm_datasets[\"validation\"]\n\n # Log a few random samples from the training set:\n for index in random.sample(range(len(train_dataset)), 3):\n logger.info(f\"Sample {index} of the training set: {train_dataset[index]}.\")\n\n # DataLoaders creation:\n train_dataloader = DataLoader(\n train_dataset, shuffle=True, collate_fn=default_data_collator, batch_size=args.per_device_train_batch_size\n )\n eval_dataloader = DataLoader(\n eval_dataset, collate_fn=default_data_collator, batch_size=args.per_device_eval_batch_size\n )\n\n # Optimizer\n # Split weights in two groups, one with weight decay and the other not.\n no_decay = [\"bias\", \"layer_norm.weight\"]\n optimizer_grouped_parameters = [\n {\n \"params\": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],\n \"weight_decay\": args.weight_decay,\n },\n {\n \"params\": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)],\n \"weight_decay\": 0.0,\n },\n ]\n optimizer = torch.optim.AdamW(optimizer_grouped_parameters, lr=args.learning_rate)\n\n # Scheduler and math around the number of training steps.\n overrode_max_train_steps = False\n num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)\n if args.max_train_steps is None:\n args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch\n overrode_max_train_steps = True\n\n # New Code\n # For Megatron-LM, we need to use `MegatronLMDummyScheduler` instead of regular schedulers\n if accelerator.distributed_type == DistributedType.MEGATRON_LM:\n lr_scheduler = MegatronLMDummyScheduler(\n optimizer=optimizer,\n total_num_steps=args.max_train_steps,\n warmup_num_steps=args.num_warmup_steps,\n )\n else:\n lr_scheduler = get_scheduler(\n name=args.lr_scheduler_type,\n optimizer=optimizer,\n num_warmup_steps=args.num_warmup_steps * args.gradient_accumulation_steps,\n num_training_steps=args.max_train_steps * args.gradient_accumulation_steps,\n )\n\n # Prepare everything with our `accelerator`.\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler\n )\n\n # On TPU, the tie weights in our model have been disconnected, so we need to restore the ties.\n if accelerator.distributed_type == DistributedType.XLA:\n model.tie_weights()\n\n # We need to recalculate our total training steps as the size of the training dataloader may have changed.\n num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)\n if overrode_max_train_steps:\n args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch\n # Afterwards we recalculate our number of training epochs\n args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)\n\n # Figure out how many steps we should save the Accelerator states\n checkpointing_steps = args.checkpointing_steps\n if checkpointing_steps is not None and checkpointing_steps.isdigit():\n checkpointing_steps = int(checkpointing_steps)\n\n # We need to initialize the trackers we use, and also store our configuration.\n # The trackers initializes automatically on the main process.\n if args.with_tracking:\n experiment_config = vars(args)\n # TensorBoard cannot log Enums, need the raw value\n experiment_config[\"lr_scheduler_type\"] = experiment_config[\"lr_scheduler_type\"].value\n accelerator.init_trackers(\"clm_no_trainer\", experiment_config)\n\n # Train!\n # New Code\n # For Megatron-LM, we need to get `global_batch_size` from megatron_lm_plugin\n # as it handles the specifics related to data parallelism, tensor model parallelism and pipeline parallelism\n if accelerator.distributed_type == DistributedType.MEGATRON_LM:\n total_batch_size = accelerator.state.megatron_lm_plugin.global_batch_size\n else:\n total_batch_size = (\n args.per_device_train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps\n )\n\n logger.info(\"***** Running training *****\")\n logger.info(f\" Num examples = {len(train_dataset)}\")\n logger.info(f\" Num Epochs = {args.num_train_epochs}\")\n logger.info(f\" Instantaneous batch size per device = {args.per_device_train_batch_size}\")\n logger.info(f\" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}\")\n logger.info(f\" Gradient Accumulation steps = {args.gradient_accumulation_steps}\")\n logger.info(f\" Total optimization steps = {args.max_train_steps}\")\n # Only show the progress bar once on each machine.\n progress_bar = tqdm(range(args.max_train_steps), disable=not accelerator.is_local_main_process)\n completed_steps = 0\n starting_epoch = 0\n\n # Potentially load in the weights and states from a previous save\n if args.resume_from_checkpoint:\n if args.resume_from_checkpoint is not None or args.resume_from_checkpoint != \"\":\n accelerator.print(f\"Resumed from checkpoint: {args.resume_from_checkpoint}\")\n accelerator.load_state(args.resume_from_checkpoint)\n path = os.path.basename(args.resume_from_checkpoint)\n else:\n # Get the most recent checkpoint\n dirs = [f.name for f in os.scandir(os.getcwd()) if f.is_dir()]\n dirs.sort(key=os.path.getctime)\n path = dirs[-1] # Sorts folders by date modified, most recent checkpoint is the last\n # Extract `epoch_{i}` or `step_{i}`\n training_difference = os.path.splitext(path)[0]\n\n if \"epoch\" in training_difference:\n starting_epoch = int(training_difference.replace(\"epoch_\", \"\")) + 1\n resume_step = None\n else:\n # need to multiply `gradient_accumulation_steps` to reflect real steps\n resume_step = int(training_difference.replace(\"step_\", \"\")) * args.gradient_accumulation_steps\n starting_epoch = resume_step // len(train_dataloader)\n resume_step -= starting_epoch * len(train_dataloader)\n\n if args.initial_megatron_lm_checkpoint:\n assert accelerator.distributed_type == DistributedType.MEGATRON_LM, (\n \"initial_megatron_lm_checkpoint should only be used with Megatron-LM\"\n )\n assert args.resume_from_checkpoint is None, (\n \"resume_from_checkpoint should not be provided when initial_megatron_lm_checkpoint is provided\"\n )\n accelerator.print(\n f\"Loading Megatron-LM checkpoint from the initial checkpoint (directly from the release directory converted using megatron bridge): {args.initial_megatron_lm_checkpoint}\"\n )\n checkpoint_dir = args.initial_megatron_lm_checkpoint\n latest_iter_file = os.path.join(checkpoint_dir, \"latest_checkpointed_iteration.txt\")\n assert os.path.isfile(latest_iter_file), f\"{latest_iter_file} does not exist in {checkpoint_dir}\"\n with open(latest_iter_file) as f:\n contents = f.read().strip()\n assert contents == \"0\", (\n f\"latest_checkpointed_iteration.txt in {checkpoint_dir} must contain only '0' (found '{contents}'), please mannually change it to '0' and rename the directory release to iter_0000000, also make sure megatron_lm_no_load_optim is set to true in the config file\"\n )\n # Also assert iter_0000000 directory exists\n iter0_dir = os.path.join(checkpoint_dir, \"iter_0000000\")\n assert os.path.isdir(iter0_dir), (\n f\"{iter0_dir} directory does not exist in {checkpoint_dir}, please rename the release directory to iter_0000000\"\n )\n accelerator.load_state(args.initial_megatron_lm_checkpoint)\n # update the progress_bar if load from checkpoint\n progress_bar.update(starting_epoch * num_update_steps_per_epoch)\n completed_steps = starting_epoch * num_update_steps_per_epoch\n\n for epoch in range(starting_epoch, args.num_train_epochs):\n model.train()\n if args.with_tracking:\n total_loss = 0\n for step, batch in enumerate(train_dataloader):\n # We need to skip steps until we reach the resumed step\n if args.resume_from_checkpoint and epoch == starting_epoch:\n if resume_step is not None and step < resume_step:\n if step % args.gradient_accumulation_steps == 0:\n progress_bar.update(1)\n completed_steps += 1\n continue\n\n with accelerator.accumulate(model):\n outputs = model(**batch)\n loss = outputs.loss\n # We keep track of the loss at each epoch\n if args.with_tracking:\n total_loss += loss.detach().float()\n accelerator.backward(loss)\n optimizer.step()\n lr_scheduler.step()\n optimizer.zero_grad()\n\n # Checks if the accelerator has performed an optimization step behind the scenes\n if accelerator.sync_gradients:\n progress_bar.update(1)\n completed_steps += 1\n\n if isinstance(checkpointing_steps, int):\n if completed_steps % checkpointing_steps == 0:\n output_dir = f\"step_{completed_steps}\"\n if args.output_dir is not None:\n output_dir = os.path.join(args.output_dir, output_dir)\n accelerator.save_state(output_dir)\n if completed_steps >= args.max_train_steps:\n break\n\n model.eval()\n losses = []\n for step, batch in enumerate(eval_dataloader):\n with torch.no_grad():\n outputs = model(**batch)\n\n loss = outputs.loss\n # New Code\n # For Megatron-LM, the losses are already averaged across the data parallel group\n if accelerator.distributed_type == DistributedType.MEGATRON_LM:\n losses.append(loss)\n else:\n losses.append(accelerator.gather_for_metrics(loss.repeat(args.per_device_eval_batch_size)))\n try:\n if accelerator.distributed_type == DistributedType.MEGATRON_LM:\n losses = torch.tensor(losses)\n else:\n losses = torch.cat(losses)\n eval_loss = torch.mean(losses)\n perplexity = math.exp(eval_loss)\n except OverflowError:\n perplexity = float(\"inf\")\n\n logger.info(f\"epoch {epoch}: perplexity: {perplexity} eval_loss: {eval_loss}\")\n\n if args.with_tracking:\n accelerator.log(\n {\n \"perplexity\": perplexity,\n \"eval_loss\": eval_loss,\n \"train_loss\": total_loss.item() / len(train_dataloader),\n \"epoch\": epoch,\n \"step\": completed_steps,\n },\n step=completed_steps,\n )\n\n if args.push_to_hub and epoch < args.num_train_epochs - 1:\n accelerator.wait_for_everyone()\n unwrapped_model = accelerator.unwrap_model(model)\n unwrapped_model.save_pretrained(\n args.output_dir, is_main_process=accelerator.is_main_process, save_function=accelerator.save\n )\n if accelerator.is_main_process:\n tokenizer.save_pretrained(args.output_dir)\n api.upload_folder(\n repo_id=repo_id,\n folder_path=args.output_dir,\n commit_message=f\"Training in progress epoch {epoch}\",\n run_as_future=True,\n )\n\n if args.checkpointing_steps == \"epoch\":\n output_dir = f\"epoch_{epoch}\"\n if args.output_dir is not None:\n output_dir = os.path.join(args.output_dir, output_dir)\n accelerator.save_state(output_dir)\n\n # this is causing some issue with Megatron-LM when using `wandb` at the end of the main function.\n # Everything works fine inspite of commenting this out. (wandb finishes/closes the run without error)\n # if args.with_tracking:\n # accelerator.end_training()\n\n if args.output_dir is not None:\n accelerator.wait_for_everyone()\n # New Code\n # For Megatron-LM, we need to save the model using `accelerator.save_state`\n if accelerator.distributed_type == DistributedType.MEGATRON_LM:\n accelerator.save_state(args.output_dir)\n else:\n unwrapped_model = accelerator.unwrap_model(model)\n unwrapped_model.save_pretrained(\n args.output_dir, is_main_process=accelerator.is_main_process, save_function=accelerator.save\n )\n if accelerator.is_main_process:\n tokenizer.save_pretrained(args.output_dir)\n if args.push_to_hub:\n api.upload_folder(\n repo_id=repo_id,\n folder_path=args.output_dir,\n commit_message=\"End of training\",\n )\n\n with open(os.path.join(args.output_dir, \"all_results.json\"), \"w\") as f:\n json.dump({\"perplexity\": perplexity}, f)\n accelerator.end_training()\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "examples/by_feature/memory.py", "content": "# Copyright 2022 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport argparse\nimport os\n\n# New Code #\nimport evaluate\nimport torch\nfrom datasets import load_dataset\nfrom torch.optim import AdamW\nfrom torch.utils.data import DataLoader\nfrom transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed\n\nfrom accelerate import Accelerator, DistributedType\nfrom accelerate.utils import find_executable_batch_size\n\n\n########################################################################\n# This is a fully working simple example to use Accelerate,\n# specifically showcasing how to ensure out-of-memory errors never\n# interrupt training, and builds off the `nlp_example.py` script.\n#\n# This example trains a Bert base model on GLUE MRPC\n# in any of the following settings (with the same script):\n# - single CPU or single GPU\n# - multi GPUS (using PyTorch distributed mode)\n# - (multi) TPUs\n# - fp16 (mixed-precision) or fp32 (normal precision)\n#\n# New additions from the base script can be found quickly by\n# looking for the # New Code # tags\n#\n# To run it in each of these various modes, follow the instructions\n# in the readme for examples:\n# https://github.com/huggingface/accelerate/tree/main/examples\n#\n########################################################################\n\n\nMAX_GPU_BATCH_SIZE = 16\nEVAL_BATCH_SIZE = 32\n\n\ndef get_dataloaders(accelerator: Accelerator, batch_size: int = 16):\n \"\"\"\n Creates a set of `DataLoader`s for the `glue` dataset,\n using \"bert-base-cased\" as the tokenizer.\n\n Args:\n accelerator (`Accelerator`):\n An `Accelerator` object\n batch_size (`int`, *optional*):\n The batch size for the train and validation DataLoaders.\n \"\"\"\n tokenizer = AutoTokenizer.from_pretrained(\"bert-base-cased\")\n datasets = load_dataset(\"glue\", \"mrpc\")\n\n def tokenize_function(examples):\n # max_length=None => use the model max length (it's actually the default)\n outputs = tokenizer(examples[\"sentence1\"], examples[\"sentence2\"], truncation=True, max_length=None)\n return outputs\n\n # Apply the method we just defined to all the examples in all the splits of the dataset\n # starting with the main process first:\n with accelerator.main_process_first():\n tokenized_datasets = datasets.map(\n tokenize_function,\n batched=True,\n remove_columns=[\"idx\", \"sentence1\", \"sentence2\"],\n )\n\n # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the\n # transformers library\n tokenized_datasets = tokenized_datasets.rename_column(\"label\", \"labels\")\n\n def collate_fn(examples):\n # On TPU it's best to pad everything to the same length or training will be very slow.\n max_length = 128 if accelerator.distributed_type == DistributedType.XLA else None\n # When using mixed precision we want round multiples of 8/16\n if accelerator.mixed_precision == \"fp8\":\n pad_to_multiple_of = 16\n elif accelerator.mixed_precision != \"no\":\n pad_to_multiple_of = 8\n else:\n pad_to_multiple_of = None\n\n return tokenizer.pad(\n examples,\n padding=\"longest\",\n max_length=max_length,\n pad_to_multiple_of=pad_to_multiple_of,\n return_tensors=\"pt\",\n )\n\n # Instantiate dataloaders.\n train_dataloader = DataLoader(\n tokenized_datasets[\"train\"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size\n )\n eval_dataloader = DataLoader(\n tokenized_datasets[\"validation\"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE\n )\n\n return train_dataloader, eval_dataloader\n\n\n# For testing only\nif os.environ.get(\"TESTING_MOCKED_DATALOADERS\", None) == \"1\":\n from accelerate.test_utils.training import mocked_dataloaders\n\n get_dataloaders = mocked_dataloaders # noqa: F811\n\n\ndef training_function(config, args):\n # For testing only\n if os.environ.get(\"TESTING_MOCKED_DATALOADERS\", None) == \"1\":\n config[\"num_epochs\"] = 2\n # Initialize accelerator\n accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)\n # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs\n lr = config[\"lr\"]\n num_epochs = int(config[\"num_epochs\"])\n seed = int(config[\"seed\"])\n batch_size = int(config[\"batch_size\"])\n\n metric = evaluate.load(\"glue\", \"mrpc\")\n\n # New Code #\n # We now can define an inner training loop function. It should take a batch size as the only parameter,\n # and build the dataloaders in there.\n # It also gets our decorator\n @find_executable_batch_size(starting_batch_size=batch_size)\n def inner_training_loop(batch_size):\n # And now just move everything below under this function\n # We need to bring in the Accelerator object from earlier\n nonlocal accelerator\n # And reset all of its attributes that could hold onto any memory:\n accelerator.free_memory()\n\n # Then we can declare the model, optimizer, and everything else:\n set_seed(seed)\n\n # Instantiate the model (we build the model here so that the seed also control new weights initialization)\n model = AutoModelForSequenceClassification.from_pretrained(\"bert-base-cased\", return_dict=True)\n\n # We could avoid this line since the accelerator is set with `device_placement=True` (default value).\n # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer\n # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).\n model = model.to(accelerator.device)\n\n # Instantiate optimizer\n optimizer = AdamW(params=model.parameters(), lr=lr)\n train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size)\n\n # Instantiate scheduler\n lr_scheduler = get_linear_schedule_with_warmup(\n optimizer=optimizer,\n num_warmup_steps=100,\n num_training_steps=(len(train_dataloader) * num_epochs),\n )\n\n # Prepare everything\n # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the\n # prepare method.\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler\n )\n\n # Now we train the model\n for epoch in range(num_epochs):\n model.train()\n for step, batch in enumerate(train_dataloader):\n # We could avoid this line since we set the accelerator with `device_placement=True`.\n batch.to(accelerator.device)\n outputs = model(**batch)\n loss = outputs.loss\n accelerator.backward(loss)\n optimizer.step()\n lr_scheduler.step()\n optimizer.zero_grad()\n\n model.eval()\n for step, batch in enumerate(eval_dataloader):\n # We could avoid this line since we set the accelerator with `device_placement=True`.\n batch.to(accelerator.device)\n with torch.no_grad():\n outputs = model(**batch)\n predictions = outputs.logits.argmax(dim=-1)\n predictions, references = accelerator.gather_for_metrics((predictions, batch[\"labels\"]))\n metric.add_batch(\n predictions=predictions,\n references=references,\n )\n\n eval_metric = metric.compute()\n # Use accelerator.print to print only on the main process.\n accelerator.print(f\"epoch {epoch}:\", eval_metric)\n\n # New Code #\n # And call it at the end with no arguments\n # Note: You could also refactor this outside of your training loop function\n inner_training_loop()\n accelerator.end_training()\n\n\ndef main():\n parser = argparse.ArgumentParser(description=\"Simple example of training script.\")\n parser.add_argument(\n \"--mixed_precision\",\n type=str,\n default=None,\n choices=[\"no\", \"fp16\", \"bf16\", \"fp8\"],\n help=\"Whether to use mixed precision. Choose\"\n \"between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10.\"\n \"and an Nvidia Ampere GPU.\",\n )\n parser.add_argument(\"--cpu\", action=\"store_true\", help=\"If passed, will train on the CPU.\")\n args = parser.parse_args()\n config = {\"lr\": 2e-5, \"num_epochs\": 3, \"seed\": 42, \"batch_size\": 16}\n training_function(config, args)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "examples/by_feature/multi_process_metrics.py", "content": "# Copyright 2022 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport argparse\nimport os\n\nimport evaluate\nimport torch\nfrom datasets import load_dataset\nfrom torch.optim import AdamW\nfrom torch.utils.data import DataLoader\nfrom transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed\n\nfrom accelerate import Accelerator, DistributedType\n\n\n########################################################################\n# This is a fully working simple example to use Accelerate,\n# specifically showcasing how to properly calculate the metrics on the\n# validation dataset when in a distributed system, and builds off the\n# `nlp_example.py` script.\n#\n# This example trains a Bert base model on GLUE MRPC\n# in any of the following settings (with the same script):\n# - single CPU or single GPU\n# - multi GPUS (using PyTorch distributed mode)\n# - (multi) TPUs\n# - fp16 (mixed-precision) or fp32 (normal precision)\n#\n# To help focus on the differences in the code, building `DataLoaders`\n# was refactored into its own function.\n# New additions from the base script can be found quickly by\n# looking for the # New Code # tags\n#\n# To run it in each of these various modes, follow the instructions\n# in the readme for examples:\n# https://github.com/huggingface/accelerate/tree/main/examples\n#\n########################################################################\n\n\nMAX_GPU_BATCH_SIZE = 16\nEVAL_BATCH_SIZE = 32\n\n\ndef get_dataloaders(accelerator: Accelerator, batch_size: int = 16):\n \"\"\"\n Creates a set of `DataLoader`s for the `glue` dataset,\n using \"bert-base-cased\" as the tokenizer.\n\n Args:\n accelerator (`Accelerator`):\n An `Accelerator` object\n batch_size (`int`, *optional*):\n The batch size for the train and validation DataLoaders.\n \"\"\"\n tokenizer = AutoTokenizer.from_pretrained(\"bert-base-cased\")\n datasets = load_dataset(\"glue\", \"mrpc\")\n\n def tokenize_function(examples):\n # max_length=None => use the model max length (it's actually the default)\n outputs = tokenizer(examples[\"sentence1\"], examples[\"sentence2\"], truncation=True, max_length=None)\n return outputs\n\n # Apply the method we just defined to all the examples in all the splits of the dataset\n # starting with the main process first:\n with accelerator.main_process_first():\n tokenized_datasets = datasets.map(\n tokenize_function,\n batched=True,\n remove_columns=[\"idx\", \"sentence1\", \"sentence2\"],\n )\n\n # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the\n # transformers library\n tokenized_datasets = tokenized_datasets.rename_column(\"label\", \"labels\")\n\n def collate_fn(examples):\n # On TPU it's best to pad everything to the same length or training will be very slow.\n max_length = 128 if accelerator.distributed_type == DistributedType.XLA else None\n # When using mixed precision we want round multiples of 8/16\n if accelerator.mixed_precision == \"fp8\":\n pad_to_multiple_of = 16\n elif accelerator.mixed_precision != \"no\":\n pad_to_multiple_of = 8\n else:\n pad_to_multiple_of = None\n\n return tokenizer.pad(\n examples,\n padding=\"longest\",\n max_length=max_length,\n pad_to_multiple_of=pad_to_multiple_of,\n return_tensors=\"pt\",\n )\n\n # Instantiate dataloaders.\n train_dataloader = DataLoader(\n tokenized_datasets[\"train\"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size\n )\n eval_dataloader = DataLoader(\n tokenized_datasets[\"validation\"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE\n )\n\n return train_dataloader, eval_dataloader\n\n\n# For testing only\nif os.environ.get(\"TESTING_MOCKED_DATALOADERS\", None) == \"1\":\n from accelerate.test_utils.training import mocked_dataloaders\n\n get_dataloaders = mocked_dataloaders # noqa: F811\n\n\ndef training_function(config, args):\n # For testing only\n if os.environ.get(\"TESTING_MOCKED_DATALOADERS\", None) == \"1\":\n config[\"num_epochs\"] = 2\n # Initialize accelerator\n accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)\n # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs\n lr = config[\"lr\"]\n num_epochs = int(config[\"num_epochs\"])\n seed = int(config[\"seed\"])\n batch_size = int(config[\"batch_size\"])\n\n metric = evaluate.load(\"glue\", \"mrpc\")\n\n # If the batch size is too big we use gradient accumulation\n gradient_accumulation_steps = 1\n if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.XLA:\n gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE\n batch_size = MAX_GPU_BATCH_SIZE\n\n set_seed(seed)\n train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size)\n # Instantiate the model (we build the model here so that the seed also control new weights initialization)\n model = AutoModelForSequenceClassification.from_pretrained(\"bert-base-cased\", return_dict=True)\n\n # We could avoid this line since the accelerator is set with `device_placement=True` (default value).\n # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer\n # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).\n model = model.to(accelerator.device)\n\n # Instantiate optimizer\n optimizer = AdamW(params=model.parameters(), lr=lr)\n\n # Instantiate scheduler\n lr_scheduler = get_linear_schedule_with_warmup(\n optimizer=optimizer,\n num_warmup_steps=100,\n num_training_steps=(len(train_dataloader) * num_epochs) // gradient_accumulation_steps,\n )\n\n # Prepare everything\n # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the\n # prepare method.\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler\n )\n\n # Now we train the model\n for epoch in range(num_epochs):\n model.train()\n for step, batch in enumerate(train_dataloader):\n # We could avoid this line since we set the accelerator with `device_placement=True`.\n batch.to(accelerator.device)\n outputs = model(**batch)\n loss = outputs.loss\n loss = loss / gradient_accumulation_steps\n accelerator.backward(loss)\n if step % gradient_accumulation_steps == 0:\n optimizer.step()\n lr_scheduler.step()\n optimizer.zero_grad()\n\n model.eval()\n samples_seen = 0\n for step, batch in enumerate(eval_dataloader):\n # We could avoid this line since we set the accelerator with `device_placement=True`.\n batch.to(accelerator.device)\n with torch.no_grad():\n outputs = model(**batch)\n predictions = outputs.logits.argmax(dim=-1)\n predictions, references = accelerator.gather((predictions, batch[\"labels\"]))\n # New Code #\n # First we check if it's a distributed system\n if accelerator.use_distributed:\n # Then see if we're on the last batch of our eval dataloader\n if step == len(eval_dataloader) - 1:\n # Last batch needs to be truncated on distributed systems as it contains additional samples\n predictions = predictions[: len(eval_dataloader.dataset) - samples_seen]\n references = references[: len(eval_dataloader.dataset) - samples_seen]\n else:\n # Otherwise we add the number of samples seen\n samples_seen += references.shape[0]\n # All of this can be avoided if you use `Accelerator.gather_for_metrics` instead of `Accelerator.gather`:\n # accelerator.gather_for_metrics((predictions, batch[\"labels\"]))\n metric.add_batch(\n predictions=predictions,\n references=references,\n )\n\n eval_metric = metric.compute()\n # Use accelerator.print to print only on the main process.\n accelerator.print(f\"epoch {epoch}:\", eval_metric)\n accelerator.end_training()\n\n\ndef main():\n parser = argparse.ArgumentParser(description=\"Simple example of training script.\")\n parser.add_argument(\n \"--mixed_precision\",\n type=str,\n default=None,\n choices=[\"no\", \"fp16\", \"bf16\", \"fp8\"],\n help=\"Whether to use mixed precision. Choose\"\n \"between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10.\"\n \"and an Nvidia Ampere GPU.\",\n )\n parser.add_argument(\"--cpu\", action=\"store_true\", help=\"If passed, will train on the CPU.\")\n args = parser.parse_args()\n config = {\"lr\": 2e-5, \"num_epochs\": 3, \"seed\": 42, \"batch_size\": 16}\n training_function(config, args)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "examples/by_feature/profiler.py", "content": "# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport argparse\nimport os\n\nimport evaluate\nimport torch\nfrom datasets import load_dataset\nfrom torch.optim import AdamW\nfrom torch.utils.data import DataLoader\nfrom transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed\n\nfrom accelerate import Accelerator, DistributedType\nfrom accelerate.utils import ProfileKwargs\n\n\n########################################################################\n# This is a fully working simple example to use Accelerate\n# and perform profiling\n#\n# This example trains a Bert base model on GLUE MRPC\n# in any of the following settings (with the same script):\n# - single CPU or single device (CUDA GPU, Intel XPU etc.)\n# - multi devices (using PyTorch distributed mode)\n# - (multi) TPUs\n# - fp16 (mixed-precision) or fp32 (normal precision)\n#\n# To run it in each of these various modes, follow the instructions\n# in the readme for examples:\n# https://github.com/huggingface/accelerate/tree/main/examples\n#\n########################################################################\n\n\nMAX_GPU_BATCH_SIZE = 16\nEVAL_BATCH_SIZE = 32\n\n\ndef get_dataloaders(accelerator: Accelerator, batch_size: int = 16):\n \"\"\"\n Creates a set of `DataLoader`s for the `glue` dataset,\n using \"bert-base-cased\" as the tokenizer.\n\n Args:\n accelerator (`Accelerator`):\n An `Accelerator` object\n batch_size (`int`, *optional*):\n The batch size for the train and validation DataLoaders.\n \"\"\"\n tokenizer = AutoTokenizer.from_pretrained(\"bert-base-cased\")\n datasets = load_dataset(\"glue\", \"mrpc\")\n\n def tokenize_function(examples):\n # max_length=None => use the model max length (it's actually the default)\n outputs = tokenizer(examples[\"sentence1\"], examples[\"sentence2\"], truncation=True, max_length=None)\n return outputs\n\n # Apply the method we just defined to all the examples in all the splits of the dataset\n # starting with the main process first:\n with accelerator.main_process_first():\n tokenized_datasets = datasets.map(\n tokenize_function,\n batched=True,\n remove_columns=[\"idx\", \"sentence1\", \"sentence2\"],\n )\n\n # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the\n # transformers library\n tokenized_datasets = tokenized_datasets.rename_column(\"label\", \"labels\")\n\n def collate_fn(examples):\n # On TPU it's best to pad everything to the same length or training will be very slow.\n max_length = 128 if accelerator.distributed_type == DistributedType.XLA else None\n # When using mixed precision we want round multiples of 8/16\n if accelerator.mixed_precision == \"fp8\":\n pad_to_multiple_of = 16\n elif accelerator.mixed_precision != \"no\":\n pad_to_multiple_of = 8\n else:\n pad_to_multiple_of = None\n\n return tokenizer.pad(\n examples,\n padding=\"longest\",\n max_length=max_length,\n pad_to_multiple_of=pad_to_multiple_of,\n return_tensors=\"pt\",\n )\n\n # Instantiate dataloaders.\n train_dataloader = DataLoader(\n tokenized_datasets[\"train\"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size\n )\n eval_dataloader = DataLoader(\n tokenized_datasets[\"validation\"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE\n )\n\n return train_dataloader, eval_dataloader\n\n\n# For testing only\nif os.environ.get(\"TESTING_MOCKED_DATALOADERS\", None) == \"1\":\n from accelerate.test_utils.training import mocked_dataloaders\n\n get_dataloaders = mocked_dataloaders # noqa: F811\n\n\ndef training_function(config, args):\n # For testing only\n if os.environ.get(\"TESTING_MOCKED_DATALOADERS\", None) == \"1\":\n config[\"num_epochs\"] = 2\n # New Code #\n profile_kwargs = ProfileKwargs(\n record_shapes=args.record_shapes,\n profile_memory=args.profile_memory,\n with_flops=args.with_flops,\n output_trace_dir=args.output_trace_dir,\n )\n # Initialize accelerator\n accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision, kwargs_handlers=[profile_kwargs])\n # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs\n lr = config[\"lr\"]\n num_epochs = int(config[\"num_epochs\"])\n seed = int(config[\"seed\"])\n batch_size = int(config[\"batch_size\"])\n\n metric = evaluate.load(\"glue\", \"mrpc\")\n\n set_seed(seed)\n train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size)\n # Instantiate the model (we build the model here so that the seed also control new weights initialization)\n model = AutoModelForSequenceClassification.from_pretrained(\"bert-base-cased\", return_dict=True)\n\n # We could avoid this line since the accelerator is set with `device_placement=True` (default value).\n # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer\n # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).\n model = model.to(accelerator.device)\n\n # Instantiate optimizer\n optimizer = AdamW(params=model.parameters(), lr=lr)\n\n # Instantiate scheduler\n lr_scheduler = get_linear_schedule_with_warmup(\n optimizer=optimizer,\n num_warmup_steps=100,\n num_training_steps=(len(train_dataloader) * num_epochs),\n )\n\n # Prepare everything\n # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the\n # prepare method.\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler\n )\n\n # Now we train the model\n for epoch in range(num_epochs):\n model.train()\n # New Code #\n with accelerator.profile() as prof:\n for step, batch in enumerate(train_dataloader):\n # We could avoid this line since we set the accelerator with `device_placement=True`.\n batch.to(accelerator.device)\n # We use the new `accumulate` context manager to perform gradient accumulation\n with accelerator.accumulate(model):\n output = model(**batch)\n loss = output.loss\n accelerator.backward(loss)\n optimizer.step()\n lr_scheduler.step()\n optimizer.zero_grad()\n # New Code #\n accelerator.print(\n prof.key_averages().table(\n sort_by=\"self_cpu_time_total\" if args.cpu else f\"self_{accelerator.device.type}_time_total\",\n row_limit=-1,\n )\n )\n\n model.eval()\n for step, batch in enumerate(eval_dataloader):\n # We could avoid this line since we set the accelerator with `device_placement=True`.\n batch.to(accelerator.device)\n with torch.no_grad():\n outputs = model(**batch)\n predictions = outputs.logits.argmax(dim=-1)\n predictions, references = accelerator.gather_for_metrics((predictions, batch[\"labels\"]))\n metric.add_batch(\n predictions=predictions,\n references=references,\n )\n\n eval_metric = metric.compute()\n # Use accelerator.print to print only on the main process.\n accelerator.print(f\"epoch {epoch}:\", eval_metric)\n accelerator.end_training()\n\n\ndef main():\n parser = argparse.ArgumentParser(description=\"Simple example of training script.\")\n parser.add_argument(\n \"--mixed_precision\",\n type=str,\n default=None,\n choices=[\"no\", \"fp16\", \"bf16\", \"fp8\"],\n help=\"Whether to use mixed precision. Choose\"\n \"between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10.\"\n \"and an Nvidia Ampere GPU or an Intel XPU.\",\n )\n # New Code #\n parser.add_argument(\n \"--record_shapes\",\n action=\"store_true\",\n default=False,\n help=\"If passed, will record shapes for profiling.\",\n )\n # New Code #\n parser.add_argument(\n \"--profile_memory\",\n action=\"store_true\",\n default=False,\n help=\"If passed, will profile memory.\",\n )\n # New Code #\n parser.add_argument(\n \"--with_flops\",\n action=\"store_true\",\n default=False,\n help=\"If passed, will profile flops.\",\n )\n # New Code #\n parser.add_argument(\n \"--output_trace_dir\",\n type=str,\n default=None,\n help=\"If passed, will save a json trace to the specified path.\",\n )\n parser.add_argument(\"--cpu\", action=\"store_true\", help=\"If passed, will train on the CPU.\")\n args = parser.parse_args()\n config = {\"lr\": 2e-5, \"num_epochs\": 3, \"seed\": 42, \"batch_size\": 16}\n training_function(config, args)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "examples/by_feature/schedule_free.py", "content": "# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport argparse\nimport os\n\nimport evaluate\nimport torch\nfrom datasets import load_dataset\nfrom torch.utils.data import DataLoader\nfrom transformers import AutoModelForSequenceClassification, AutoTokenizer, set_seed\n\nfrom accelerate import Accelerator, DistributedType\nfrom accelerate.utils import is_schedulefree_available\n\n\nif is_schedulefree_available():\n import schedulefree\nelse:\n raise ImportError(\n \"This example requires the `schedulefree` library. Please install it with `pip install schedulefree`\"\n )\n\n\n########################################################################\n# This is a fully working simple example to use Accelerate and Facebook's\n# scheduler-free optimizer: https://github.com/facebookresearch/schedule_free/\n#\n# This example trains a Bert base model on GLUE MRPC\n# in any of the following settings (with the same script):\n# - single CPU or single GPU\n# - multi GPUS (using PyTorch distributed mode)\n# - (multi) TPUs\n# - fp16 (mixed-precision) or fp32 (normal precision)\n#\n# To run it in each of these various modes, follow the instructions\n# in the readme for examples:\n# https://github.com/huggingface/accelerate/tree/main/examples\n#\n########################################################################\n\n\nMAX_GPU_BATCH_SIZE = 16\nEVAL_BATCH_SIZE = 32\n\n\ndef get_dataloaders(accelerator: Accelerator, batch_size: int = 16):\n \"\"\"\n Creates a set of `DataLoader`s for the `glue` dataset,\n using \"bert-base-cased\" as the tokenizer.\n\n Args:\n accelerator (`Accelerator`):\n An `Accelerator` object\n batch_size (`int`, *optional*):\n The batch size for the train and validation DataLoaders.\n \"\"\"\n tokenizer = AutoTokenizer.from_pretrained(\"bert-base-cased\")\n datasets = load_dataset(\"glue\", \"mrpc\")\n\n def tokenize_function(examples):\n # max_length=None => use the model max length (it's actually the default)\n outputs = tokenizer(examples[\"sentence1\"], examples[\"sentence2\"], truncation=True, max_length=None)\n return outputs\n\n # Apply the method we just defined to all the examples in all the splits of the dataset\n # starting with the main process first:\n with accelerator.main_process_first():\n tokenized_datasets = datasets.map(\n tokenize_function,\n batched=True,\n remove_columns=[\"idx\", \"sentence1\", \"sentence2\"],\n )\n\n # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the\n # transformers library\n tokenized_datasets = tokenized_datasets.rename_column(\"label\", \"labels\")\n\n def collate_fn(examples):\n # For Torchxla, it's best to pad everything to the same length or training will be very slow.\n max_length = 128 if accelerator.distributed_type == DistributedType.XLA else None\n # When using mixed precision we want round multiples of 8/16\n if accelerator.mixed_precision == \"fp8\":\n pad_to_multiple_of = 16\n elif accelerator.mixed_precision != \"no\":\n pad_to_multiple_of = 8\n else:\n pad_to_multiple_of = None\n\n return tokenizer.pad(\n examples,\n padding=\"longest\",\n max_length=max_length,\n pad_to_multiple_of=pad_to_multiple_of,\n return_tensors=\"pt\",\n )\n\n # Instantiate dataloaders.\n train_dataloader = DataLoader(\n tokenized_datasets[\"train\"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size, drop_last=True\n )\n eval_dataloader = DataLoader(\n tokenized_datasets[\"validation\"],\n shuffle=False,\n collate_fn=collate_fn,\n batch_size=EVAL_BATCH_SIZE,\n drop_last=(accelerator.mixed_precision == \"fp8\"),\n )\n\n return train_dataloader, eval_dataloader\n\n\n# For testing only\n\n\nif os.environ.get(\"TESTING_MOCKED_DATALOADERS\", None) == \"1\":\n from accelerate.test_utils.training import mocked_dataloaders\n\n get_dataloaders = mocked_dataloaders # noqa: F811\n\n\ndef training_function(config, args):\n # Initialize accelerator\n accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)\n # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs\n lr = config[\"lr\"]\n num_epochs = int(config[\"num_epochs\"])\n seed = int(config[\"seed\"])\n batch_size = int(config[\"batch_size\"])\n\n metric = evaluate.load(\"glue\", \"mrpc\")\n\n # If the batch size is too big we use gradient accumulation\n gradient_accumulation_steps = 1\n if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.XLA:\n gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE\n batch_size = MAX_GPU_BATCH_SIZE\n\n set_seed(seed)\n train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size)\n # Instantiate the model (we build the model here so that the seed also control new weights initialization)\n model = AutoModelForSequenceClassification.from_pretrained(\"bert-base-cased\", return_dict=True)\n\n # We could avoid this line since the accelerator is set with `device_placement=True` (default value).\n # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer\n # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).\n model = model.to(accelerator.device)\n # Instantiate optimizer with warmup steps\n optimizer = schedulefree.AdamWScheduleFree(\n model.parameters(),\n lr=lr,\n warmup_steps=100,\n )\n\n # Prepare everything\n # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the\n # prepare method.\n\n model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(\n model, optimizer, train_dataloader, eval_dataloader\n )\n\n # Now we train the model\n for epoch in range(num_epochs):\n model.train()\n optimizer.train()\n for step, batch in enumerate(train_dataloader):\n # We could avoid this line since we set the accelerator with `device_placement=True`.\n batch.to(accelerator.device)\n outputs = model(**batch)\n loss = outputs.loss\n loss = loss / gradient_accumulation_steps\n accelerator.backward(loss)\n if step % gradient_accumulation_steps == 0:\n optimizer.step()\n optimizer.zero_grad()\n\n model.eval()\n optimizer.eval()\n for step, batch in enumerate(eval_dataloader):\n # We could avoid this line since we set the accelerator with `device_placement=True`.\n batch.to(accelerator.device)\n with torch.no_grad():\n outputs = model(**batch)\n predictions = outputs.logits.argmax(dim=-1)\n predictions, references = accelerator.gather_for_metrics((predictions, batch[\"labels\"]))\n metric.add_batch(\n predictions=predictions,\n references=references,\n )\n\n eval_metric = metric.compute()\n # Use accelerator.print to print only on the main process.\n accelerator.print(f\"epoch {epoch}:\", eval_metric)\n accelerator.end_training()\n\n\ndef main():\n parser = argparse.ArgumentParser(description=\"Simple example of training script.\")\n parser.add_argument(\n \"--mixed_precision\",\n type=str,\n default=None,\n choices=[\"no\", \"fp16\", \"bf16\", \"fp8\"],\n help=\"Whether to use mixed precision. Choose\"\n \"between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10.\"\n \"and an Nvidia Ampere GPU.\",\n )\n parser.add_argument(\"--cpu\", action=\"store_true\", help=\"If passed, will train on the CPU.\")\n args = parser.parse_args()\n config = {\"lr\": 2e-5, \"num_epochs\": 3, \"seed\": 42, \"batch_size\": 16}\n training_function(config, args)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "examples/by_feature/tracking.py", "content": "# Copyright 2021 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport argparse\nimport os\n\nimport evaluate\nimport torch\nfrom datasets import load_dataset\nfrom torch.optim import AdamW\nfrom torch.utils.data import DataLoader\nfrom transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed\n\nfrom accelerate import Accelerator, DistributedType\n\n\n########################################################################\n# This is a fully working simple example to use Accelerate,\n# specifically showcasing the experiment tracking capability,\n# and builds off the `nlp_example.py` script.\n#\n# This example trains a Bert base model on GLUE MRPC\n# in any of the following settings (with the same script):\n# - single CPU or single GPU\n# - multi GPUS (using PyTorch distributed mode)\n# - (multi) TPUs\n# - fp16 (mixed-precision) or fp32 (normal precision)\n#\n# To help focus on the differences in the code, building `DataLoaders`\n# was refactored into its own function.\n# New additions from the base script can be found quickly by\n# looking for the # New Code # tags\n#\n# To run it in each of these various modes, follow the instructions\n# in the readme for examples:\n# https://github.com/huggingface/accelerate/tree/main/examples\n#\n########################################################################\n\nMAX_GPU_BATCH_SIZE = 16\nEVAL_BATCH_SIZE = 32\n\n\ndef get_dataloaders(accelerator: Accelerator, batch_size: int = 16):\n \"\"\"\n Creates a set of `DataLoader`s for the `glue` dataset,\n using \"bert-base-cased\" as the tokenizer.\n\n Args:\n accelerator (`Accelerator`):\n An `Accelerator` object\n batch_size (`int`, *optional*):\n The batch size for the train and validation DataLoaders.\n \"\"\"\n tokenizer = AutoTokenizer.from_pretrained(\"bert-base-cased\")\n datasets = load_dataset(\"glue\", \"mrpc\")\n\n def tokenize_function(examples):\n # max_length=None => use the model max length (it's actually the default)\n outputs = tokenizer(examples[\"sentence1\"], examples[\"sentence2\"], truncation=True, max_length=None)\n return outputs\n\n # Apply the method we just defined to all the examples in all the splits of the dataset\n # starting with the main process first:\n with accelerator.main_process_first():\n tokenized_datasets = datasets.map(\n tokenize_function,\n batched=True,\n remove_columns=[\"idx\", \"sentence1\", \"sentence2\"],\n )\n\n # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the\n # transformers library\n tokenized_datasets = tokenized_datasets.rename_column(\"label\", \"labels\")\n\n def collate_fn(examples):\n # On TPU it's best to pad everything to the same length or training will be very slow.\n max_length = 128 if accelerator.distributed_type == DistributedType.XLA else None\n # When using mixed precision we want round multiples of 8/16\n if accelerator.mixed_precision == \"fp8\":\n pad_to_multiple_of = 16\n elif accelerator.mixed_precision != \"no\":\n pad_to_multiple_of = 8\n else:\n pad_to_multiple_of = None\n\n return tokenizer.pad(\n examples,\n padding=\"longest\",\n max_length=max_length,\n pad_to_multiple_of=pad_to_multiple_of,\n return_tensors=\"pt\",\n )\n\n # Instantiate dataloaders.\n train_dataloader = DataLoader(\n tokenized_datasets[\"train\"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size\n )\n eval_dataloader = DataLoader(\n tokenized_datasets[\"validation\"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE\n )\n\n return train_dataloader, eval_dataloader\n\n\n# For testing only\nif os.environ.get(\"TESTING_MOCKED_DATALOADERS\", None) == \"1\":\n from accelerate.test_utils.training import mocked_dataloaders\n\n get_dataloaders = mocked_dataloaders # noqa: F811\n\n\ndef training_function(config, args):\n # For testing only\n if os.environ.get(\"TESTING_MOCKED_DATALOADERS\", None) == \"1\":\n config[\"num_epochs\"] = 2\n # Initialize Accelerator\n\n # New Code #\n # We pass in \"all\" to `log_with` to grab all available trackers in the environment\n # Note: If using a custom `Tracker` class, should be passed in here such as:\n # >>> log_with = [\"all\", MyCustomTrackerClassInstance()]\n if args.with_tracking:\n accelerator = Accelerator(\n cpu=args.cpu, mixed_precision=args.mixed_precision, log_with=\"all\", project_dir=args.project_dir\n )\n else:\n accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)\n # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs\n lr = config[\"lr\"]\n num_epochs = int(config[\"num_epochs\"])\n seed = int(config[\"seed\"])\n batch_size = int(config[\"batch_size\"])\n set_seed(seed)\n\n train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size)\n metric = evaluate.load(\"glue\", \"mrpc\")\n\n # If the batch size is too big we use gradient accumulation\n gradient_accumulation_steps = 1\n if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.XLA:\n gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE\n batch_size = MAX_GPU_BATCH_SIZE\n\n # Instantiate the model (we build the model here so that the seed also control new weights initialization)\n model = AutoModelForSequenceClassification.from_pretrained(\"bert-base-cased\", return_dict=True)\n\n # We could avoid this line since the accelerator is set with `device_placement=True` (default value).\n # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer\n # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).\n model = model.to(accelerator.device)\n\n # Instantiate optimizer\n optimizer = AdamW(params=model.parameters(), lr=lr)\n\n # Instantiate scheduler\n lr_scheduler = get_linear_schedule_with_warmup(\n optimizer=optimizer,\n num_warmup_steps=100,\n num_training_steps=(len(train_dataloader) * num_epochs) // gradient_accumulation_steps,\n )\n\n # Prepare everything\n # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the\n # prepare method.\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler\n )\n\n # New Code #\n # We need to initialize the trackers we use. Overall configurations can also be stored\n if args.with_tracking:\n run = os.path.split(__file__)[-1].split(\".\")[0]\n accelerator.init_trackers(run, config)\n\n # Now we train the model\n for epoch in range(num_epochs):\n model.train()\n # New Code #\n # For our tracking example, we will log the total loss of each epoch\n if args.with_tracking:\n total_loss = 0\n for step, batch in enumerate(train_dataloader):\n # We could avoid this line since we set the accelerator with `device_placement=True`.\n batch.to(accelerator.device)\n outputs = model(**batch)\n loss = outputs.loss\n # New Code #\n if args.with_tracking:\n total_loss += loss.detach().float()\n loss = loss / gradient_accumulation_steps\n accelerator.backward(loss)\n if step % gradient_accumulation_steps == 0:\n optimizer.step()\n lr_scheduler.step()\n optimizer.zero_grad()\n\n model.eval()\n for step, batch in enumerate(eval_dataloader):\n # We could avoid this line since we set the accelerator with `device_placement=True` (the default).\n batch.to(accelerator.device)\n with torch.no_grad():\n outputs = model(**batch)\n predictions = outputs.logits.argmax(dim=-1)\n predictions, references = accelerator.gather_for_metrics((predictions, batch[\"labels\"]))\n metric.add_batch(\n predictions=predictions,\n references=references,\n )\n\n eval_metric = metric.compute()\n # Use accelerator.print to print only on the main process.\n accelerator.print(f\"epoch {epoch}:\", eval_metric)\n\n # New Code #\n # To actually log, we call `Accelerator.log`\n # The values passed can be of `str`, `int`, `float` or `dict` of `str` to `float`/`int`\n if args.with_tracking:\n accelerator.log(\n {\n \"accuracy\": eval_metric[\"accuracy\"],\n \"f1\": eval_metric[\"f1\"],\n \"train_loss\": total_loss.item() / len(train_dataloader),\n \"epoch\": epoch,\n },\n step=epoch,\n )\n\n accelerator.end_training()\n\n\ndef main():\n parser = argparse.ArgumentParser(description=\"Simple example of training script.\")\n parser.add_argument(\n \"--mixed_precision\",\n type=str,\n default=None,\n choices=[\"no\", \"fp16\", \"bf16\", \"fp8\"],\n help=\"Whether to use mixed precision. Choose\"\n \"between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10.\"\n \"and an Nvidia Ampere GPU.\",\n )\n parser.add_argument(\"--cpu\", action=\"store_true\", help=\"If passed, will train on the CPU.\")\n parser.add_argument(\n \"--with_tracking\",\n action=\"store_true\",\n help=\"Whether to load in all available experiment trackers from the environment and use them for logging.\",\n )\n parser.add_argument(\n \"--project_dir\",\n type=str,\n default=\"logs\",\n help=\"Location on where to store experiment tracking logs` and relevent project information\",\n )\n args = parser.parse_args()\n config = {\"lr\": 2e-5, \"num_epochs\": 3, \"seed\": 42, \"batch_size\": 16}\n training_function(config, args)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "examples/complete_cv_example.py", "content": "# Copyright 2021 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport argparse\nimport os\nimport re\n\nimport numpy as np\nimport PIL\nimport torch\nfrom timm import create_model\nfrom torch.optim.lr_scheduler import OneCycleLR\nfrom torch.utils.data import DataLoader, Dataset\nfrom torchvision.transforms import Compose, RandomResizedCrop, Resize, ToTensor\n\nfrom accelerate import Accelerator, DataLoaderConfiguration\nfrom accelerate.utils import is_xpu_available\n\n\n########################################################################\n# This is a fully working simple example to use Accelerate\n#\n# This example trains a ResNet50 on the Oxford-IIT Pet Dataset\n# in any of the following settings (with the same script):\n# - single CPU or single GPU\n# - multi GPUS (using PyTorch distributed mode)\n# - (multi) TPUs\n# - fp16 (mixed-precision) or fp32 (normal precision)\n#\n# To run it in each of these various modes, follow the instructions\n# in the readme for examples:\n# https://github.com/huggingface/accelerate/tree/main/examples\n#\n########################################################################\n\n\n# Function to get the label from the filename\ndef extract_label(fname):\n stem = fname.split(os.path.sep)[-1]\n return re.search(r\"^(.*)_\\d+\\.jpg$\", stem).groups()[0]\n\n\nclass PetsDataset(Dataset):\n def __init__(self, file_names, image_transform=None, label_to_id=None):\n self.file_names = file_names\n self.image_transform = image_transform\n self.label_to_id = label_to_id\n\n def __len__(self):\n return len(self.file_names)\n\n def __getitem__(self, idx):\n fname = self.file_names[idx]\n raw_image = PIL.Image.open(fname)\n image = raw_image.convert(\"RGB\")\n if self.image_transform is not None:\n image = self.image_transform(image)\n label = extract_label(fname)\n if self.label_to_id is not None:\n label = self.label_to_id[label]\n return {\"image\": image, \"label\": label}\n\n\ndef training_function(config, args):\n # Initialize accelerator\n dataloader_config = DataLoaderConfiguration(use_stateful_dataloader=args.use_stateful_dataloader)\n if args.with_tracking:\n accelerator = Accelerator(\n cpu=args.cpu,\n mixed_precision=args.mixed_precision,\n log_with=\"all\",\n project_dir=args.project_dir,\n dataloader_config=dataloader_config,\n )\n else:\n accelerator = Accelerator(\n cpu=args.cpu, mixed_precision=args.mixed_precision, dataloader_config=dataloader_config\n )\n\n # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs\n lr = config[\"lr\"]\n num_epochs = int(config[\"num_epochs\"])\n seed = int(config[\"seed\"])\n batch_size = int(config[\"batch_size\"])\n image_size = config[\"image_size\"]\n if not isinstance(image_size, (list, tuple)):\n image_size = (image_size, image_size)\n\n # Parse out whether we are saving every epoch or after a certain number of batches\n if hasattr(args.checkpointing_steps, \"isdigit\"):\n if args.checkpointing_steps == \"epoch\":\n checkpointing_steps = args.checkpointing_steps\n elif args.checkpointing_steps.isdigit():\n checkpointing_steps = int(args.checkpointing_steps)\n else:\n raise ValueError(\n f\"Argument `checkpointing_steps` must be either a number or `epoch`. `{args.checkpointing_steps}` passed.\"\n )\n else:\n checkpointing_steps = None\n\n # We need to initialize the trackers we use, and also store our configuration\n if args.with_tracking:\n run = os.path.split(__file__)[-1].split(\".\")[0]\n accelerator.init_trackers(run, config)\n\n # Grab all the image filenames\n file_names = [os.path.join(args.data_dir, fname) for fname in os.listdir(args.data_dir) if fname.endswith(\".jpg\")]\n\n # Build the label correspondences\n all_labels = [extract_label(fname) for fname in file_names]\n id_to_label = list(set(all_labels))\n id_to_label.sort()\n label_to_id = {lbl: i for i, lbl in enumerate(id_to_label)}\n\n # Set the seed before splitting the data.\n np.random.seed(seed)\n torch.manual_seed(seed)\n if torch.cuda.is_available():\n torch.cuda.manual_seed_all(seed)\n elif is_xpu_available():\n torch.xpu.manual_seed_all(seed)\n\n # Split our filenames between train and validation\n random_perm = np.random.permutation(len(file_names))\n cut = int(0.8 * len(file_names))\n train_split = random_perm[:cut]\n eval_split = random_perm[cut:]\n\n # For training we use a simple RandomResizedCrop\n train_tfm = Compose([RandomResizedCrop(image_size, scale=(0.5, 1.0)), ToTensor()])\n train_dataset = PetsDataset(\n [file_names[i] for i in train_split], image_transform=train_tfm, label_to_id=label_to_id\n )\n\n # For evaluation, we use a deterministic Resize\n eval_tfm = Compose([Resize(image_size), ToTensor()])\n eval_dataset = PetsDataset([file_names[i] for i in eval_split], image_transform=eval_tfm, label_to_id=label_to_id)\n\n # Instantiate dataloaders.\n train_dataloader = DataLoader(train_dataset, shuffle=True, batch_size=batch_size, num_workers=4)\n eval_dataloader = DataLoader(eval_dataset, shuffle=False, batch_size=batch_size, num_workers=4)\n\n # Instantiate the model (we build the model here so that the seed also control new weights initialization)\n model = create_model(\"resnet50d\", pretrained=True, num_classes=len(label_to_id))\n\n # We could avoid this line since the accelerator is set with `device_placement=True` (default value).\n # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer\n # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).\n model = model.to(accelerator.device)\n\n # Freezing the base model\n for param in model.parameters():\n param.requires_grad = False\n for param in model.get_classifier().parameters():\n param.requires_grad = True\n\n # We normalize the batches of images to be a bit faster.\n mean = torch.tensor(model.default_cfg[\"mean\"])[None, :, None, None].to(accelerator.device)\n std = torch.tensor(model.default_cfg[\"std\"])[None, :, None, None].to(accelerator.device)\n\n # Instantiate optimizer\n optimizer = torch.optim.Adam(params=model.parameters(), lr=lr / 25)\n\n # Instantiate learning rate scheduler\n lr_scheduler = OneCycleLR(optimizer=optimizer, max_lr=lr, epochs=num_epochs, steps_per_epoch=len(train_dataloader))\n\n # Prepare everything\n # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the\n # prepare method.\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler\n )\n # We need to keep track of how many total steps we have iterated over\n overall_step = 0\n # We also need to keep track of the starting epoch so files are named properly\n starting_epoch = 0\n\n # Potentially load in the weights and states from a previous save\n if args.resume_from_checkpoint:\n if args.resume_from_checkpoint is not None or args.resume_from_checkpoint != \"\":\n accelerator.print(f\"Resumed from checkpoint: {args.resume_from_checkpoint}\")\n accelerator.load_state(args.resume_from_checkpoint)\n path = os.path.basename(args.resume_from_checkpoint)\n else:\n # Get the most recent checkpoint\n dirs = [f.name for f in os.scandir(os.getcwd()) if f.is_dir()]\n dirs.sort(key=os.path.getctime)\n path = dirs[-1] # Sorts folders by date modified, most recent checkpoint is the last\n # Extract `epoch_{i}` or `step_{i}`\n training_difference = os.path.splitext(path)[0]\n\n if \"epoch\" in training_difference:\n starting_epoch = int(training_difference.replace(\"epoch_\", \"\")) + 1\n resume_step = None\n else:\n resume_step = int(training_difference.replace(\"step_\", \"\"))\n starting_epoch = resume_step // len(train_dataloader)\n resume_step -= starting_epoch * len(train_dataloader)\n\n # Now we train the model\n for epoch in range(starting_epoch, num_epochs):\n model.train()\n if args.with_tracking:\n total_loss = 0\n if args.resume_from_checkpoint and epoch == starting_epoch and resume_step is not None:\n # We need to skip steps until we reach the resumed step\n active_dataloader = accelerator.skip_first_batches(train_dataloader, resume_step)\n overall_step += resume_step\n else:\n # After the first iteration though, we need to go back to the original dataloader\n active_dataloader = train_dataloader\n for batch in active_dataloader:\n # We could avoid this line since we set the accelerator with `device_placement=True`.\n batch = {k: v.to(accelerator.device) for k, v in batch.items()}\n inputs = (batch[\"image\"] - mean) / std\n outputs = model(inputs)\n loss = torch.nn.functional.cross_entropy(outputs, batch[\"label\"])\n # We keep track of the loss at each epoch\n if args.with_tracking:\n total_loss += loss.detach().float()\n accelerator.backward(loss)\n optimizer.step()\n lr_scheduler.step()\n optimizer.zero_grad()\n overall_step += 1\n if isinstance(checkpointing_steps, int):\n output_dir = f\"step_{overall_step}\"\n if overall_step % checkpointing_steps == 0:\n if args.output_dir is not None:\n output_dir = os.path.join(args.output_dir, output_dir)\n accelerator.save_state(output_dir)\n model.eval()\n accurate = 0\n num_elems = 0\n for step, batch in enumerate(eval_dataloader):\n # We could avoid this line since we set the accelerator with `device_placement=True`.\n batch = {k: v.to(accelerator.device) for k, v in batch.items()}\n inputs = (batch[\"image\"] - mean) / std\n with torch.no_grad():\n outputs = model(inputs)\n predictions = outputs.argmax(dim=-1)\n predictions, references = accelerator.gather_for_metrics((predictions, batch[\"label\"]))\n accurate_preds = predictions == references\n num_elems += accurate_preds.shape[0]\n accurate += accurate_preds.long().sum()\n\n eval_metric = accurate.item() / num_elems\n # Use accelerator.print to print only on the main process.\n accelerator.print(f\"epoch {epoch}: {100 * eval_metric:.2f}\")\n if args.with_tracking:\n accelerator.log(\n {\n \"accuracy\": 100 * eval_metric,\n \"train_loss\": total_loss.item() / len(train_dataloader),\n \"epoch\": epoch,\n },\n step=overall_step,\n )\n if checkpointing_steps == \"epoch\":\n output_dir = f\"epoch_{epoch}\"\n if args.output_dir is not None:\n output_dir = os.path.join(args.output_dir, output_dir)\n accelerator.save_state(output_dir)\n\n accelerator.end_training()\n\n\ndef main():\n parser = argparse.ArgumentParser(description=\"Simple example of training script.\")\n parser.add_argument(\"--data_dir\", required=True, help=\"The data folder on disk.\")\n parser.add_argument(\"--fp16\", action=\"store_true\", help=\"If passed, will use FP16 training.\")\n parser.add_argument(\n \"--mixed_precision\",\n type=str,\n default=None,\n choices=[\"no\", \"fp16\", \"bf16\", \"fp8\"],\n help=\"Whether to use mixed precision. Choose\"\n \"between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10.\"\n \"and an Nvidia Ampere GPU.\",\n )\n parser.add_argument(\"--cpu\", action=\"store_true\", help=\"If passed, will train on the CPU.\")\n parser.add_argument(\n \"--checkpointing_steps\",\n type=str,\n default=None,\n help=\"Whether the various states should be saved at the end of every n steps, or 'epoch' for each epoch.\",\n )\n parser.add_argument(\n \"--output_dir\",\n type=str,\n default=\".\",\n help=\"Optional save directory where all checkpoint folders will be stored. Default is the current working directory.\",\n )\n parser.add_argument(\n \"--resume_from_checkpoint\",\n type=str,\n default=None,\n help=\"If the training should continue from a checkpoint folder.\",\n )\n parser.add_argument(\n \"--use_stateful_dataloader\",\n action=\"store_true\",\n help=\"If the dataloader should be a resumable stateful dataloader.\",\n )\n parser.add_argument(\n \"--with_tracking\",\n action=\"store_true\",\n help=\"Whether to load in all available experiment trackers from the environment and use them for logging.\",\n )\n parser.add_argument(\n \"--project_dir\",\n type=str,\n default=\"logs\",\n help=\"Location on where to store experiment tracking logs` and relevent project information\",\n )\n args = parser.parse_args()\n config = {\"lr\": 3e-2, \"num_epochs\": 3, \"seed\": 42, \"batch_size\": 64, \"image_size\": 224}\n training_function(config, args)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "examples/complete_nlp_example.py", "content": "# Copyright 2021 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport argparse\nimport os\n\nimport evaluate\nimport torch\nfrom datasets import load_dataset\nfrom torch.optim import AdamW\nfrom torch.utils.data import DataLoader\nfrom transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed\n\nfrom accelerate import Accelerator, DataLoaderConfiguration, DistributedType\n\n\n########################################################################\n# This is a fully working simple example to use Accelerate\n#\n# This example trains a Bert base model on GLUE MRPC\n# in any of the following settings (with the same script):\n# - single CPU or single GPU\n# - multi GPUS (using PyTorch distributed mode)\n# - (multi) TPUs\n# - fp16 (mixed-precision) or fp32 (normal precision)\n#\n# This example also demonstrates the checkpointing and sharding capabilities\n#\n# To run it in each of these various modes, follow the instructions\n# in the readme for examples:\n# https://github.com/huggingface/accelerate/tree/main/examples\n#\n########################################################################\n\n\nMAX_GPU_BATCH_SIZE = 16\nEVAL_BATCH_SIZE = 32\n\n\ndef training_function(config, args):\n # Initialize accelerator\n dataloader_config = DataLoaderConfiguration(use_stateful_dataloader=args.use_stateful_dataloader)\n if args.with_tracking:\n accelerator = Accelerator(\n cpu=args.cpu,\n mixed_precision=args.mixed_precision,\n dataloader_config=dataloader_config,\n log_with=\"all\",\n project_dir=args.project_dir,\n )\n else:\n accelerator = Accelerator(\n cpu=args.cpu, mixed_precision=args.mixed_precision, dataloader_config=dataloader_config\n )\n\n if hasattr(args.checkpointing_steps, \"isdigit\"):\n if args.checkpointing_steps == \"epoch\":\n checkpointing_steps = args.checkpointing_steps\n elif args.checkpointing_steps.isdigit():\n checkpointing_steps = int(args.checkpointing_steps)\n else:\n raise ValueError(\n f\"Argument `checkpointing_steps` must be either a number or `epoch`. `{args.checkpointing_steps}` passed.\"\n )\n else:\n checkpointing_steps = None\n # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs\n lr = config[\"lr\"]\n num_epochs = int(config[\"num_epochs\"])\n seed = int(config[\"seed\"])\n batch_size = int(config[\"batch_size\"])\n\n # We need to initialize the trackers we use, and also store our configuration\n if args.with_tracking:\n run = os.path.split(__file__)[-1].split(\".\")[0]\n accelerator.init_trackers(run, config)\n\n tokenizer = AutoTokenizer.from_pretrained(\"bert-base-cased\")\n datasets = load_dataset(\"glue\", \"mrpc\")\n metric = evaluate.load(\"glue\", \"mrpc\")\n\n def tokenize_function(examples):\n # max_length=None => use the model max length (it's actually the default)\n outputs = tokenizer(examples[\"sentence1\"], examples[\"sentence2\"], truncation=True, max_length=None)\n return outputs\n\n # Apply the method we just defined to all the examples in all the splits of the dataset\n # starting with the main process first:\n with accelerator.main_process_first():\n tokenized_datasets = datasets.map(\n tokenize_function,\n batched=True,\n remove_columns=[\"idx\", \"sentence1\", \"sentence2\"],\n )\n\n # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the\n # transformers library\n tokenized_datasets = tokenized_datasets.rename_column(\"label\", \"labels\")\n\n # If the batch size is too big we use gradient accumulation\n gradient_accumulation_steps = 1\n if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.XLA:\n gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE\n batch_size = MAX_GPU_BATCH_SIZE\n\n def collate_fn(examples):\n # On TPU it's best to pad everything to the same length or training will be very slow.\n max_length = 128 if accelerator.distributed_type == DistributedType.XLA else None\n # When using mixed precision we want round multiples of 8/16\n if accelerator.mixed_precision == \"fp8\":\n pad_to_multiple_of = 16\n elif accelerator.mixed_precision != \"no\":\n pad_to_multiple_of = 8\n else:\n pad_to_multiple_of = None\n\n return tokenizer.pad(\n examples,\n padding=\"longest\",\n max_length=max_length,\n pad_to_multiple_of=pad_to_multiple_of,\n return_tensors=\"pt\",\n )\n\n # Instantiate dataloaders.\n train_dataloader = DataLoader(\n tokenized_datasets[\"train\"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size\n )\n eval_dataloader = DataLoader(\n tokenized_datasets[\"validation\"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE\n )\n\n set_seed(seed)\n\n # Instantiate the model (we build the model here so that the seed also control new weights initialization)\n model = AutoModelForSequenceClassification.from_pretrained(\"bert-base-cased\", return_dict=True)\n\n # We could avoid this line since the accelerator is set with `device_placement=True` (default value).\n # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer\n # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).\n model = model.to(accelerator.device)\n\n # Instantiate optimizer\n optimizer = AdamW(params=model.parameters(), lr=lr)\n\n # Instantiate scheduler\n lr_scheduler = get_linear_schedule_with_warmup(\n optimizer=optimizer,\n num_warmup_steps=100,\n num_training_steps=(len(train_dataloader) * num_epochs) // gradient_accumulation_steps,\n )\n\n # Prepare everything\n # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the\n # prepare method.\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler\n )\n\n # We need to keep track of how many total steps we have iterated over\n overall_step = 0\n # We also need to keep track of the stating epoch so files are named properly\n starting_epoch = 0\n\n # Potentially load in the weights and states from a previous save\n if args.resume_from_checkpoint:\n if args.resume_from_checkpoint is not None or args.resume_from_checkpoint != \"\":\n accelerator.print(f\"Resumed from checkpoint: {args.resume_from_checkpoint}\")\n accelerator.load_state(args.resume_from_checkpoint)\n path = os.path.basename(args.resume_from_checkpoint)\n else:\n # Get the most recent checkpoint\n dirs = [f.name for f in os.scandir(os.getcwd()) if f.is_dir()]\n dirs.sort(key=os.path.getctime)\n path = dirs[-1] # Sorts folders by date modified, most recent checkpoint is the last\n # Extract `epoch_{i}` or `step_{i}`\n training_difference = os.path.splitext(path)[0]\n\n if \"epoch\" in training_difference:\n starting_epoch = int(training_difference.replace(\"epoch_\", \"\")) + 1\n resume_step = None\n else:\n resume_step = int(training_difference.replace(\"step_\", \"\"))\n starting_epoch = resume_step // len(train_dataloader)\n resume_step -= starting_epoch * len(train_dataloader)\n\n # Now we train the model\n for epoch in range(starting_epoch, num_epochs):\n model.train()\n if args.with_tracking:\n total_loss = 0\n if args.resume_from_checkpoint and epoch == starting_epoch and resume_step is not None:\n # We need to skip steps until we reach the resumed step\n if not args.use_stateful_dataloader:\n active_dataloader = accelerator.skip_first_batches(train_dataloader, resume_step)\n else:\n active_dataloader = train_dataloader\n overall_step += resume_step\n else:\n # After the first iteration though, we need to go back to the original dataloader\n active_dataloader = train_dataloader\n for step, batch in enumerate(active_dataloader):\n # We could avoid this line since we set the accelerator with `device_placement=True`.\n batch.to(accelerator.device)\n outputs = model(**batch)\n loss = outputs.loss\n loss = loss / gradient_accumulation_steps\n # We keep track of the loss at each epoch\n if args.with_tracking:\n total_loss += loss.detach().float()\n accelerator.backward(loss)\n if step % gradient_accumulation_steps == 0:\n optimizer.step()\n lr_scheduler.step()\n optimizer.zero_grad()\n\n overall_step += 1\n\n if isinstance(checkpointing_steps, int):\n output_dir = f\"step_{overall_step}\"\n if overall_step % checkpointing_steps == 0:\n if args.output_dir is not None:\n output_dir = os.path.join(args.output_dir, output_dir)\n accelerator.save_state(output_dir)\n\n model.eval()\n for step, batch in enumerate(eval_dataloader):\n # We could avoid this line since we set the accelerator with `device_placement=True`.\n batch.to(accelerator.device)\n with torch.no_grad():\n outputs = model(**batch)\n predictions = outputs.logits.argmax(dim=-1)\n predictions, references = accelerator.gather_for_metrics((predictions, batch[\"labels\"]))\n metric.add_batch(\n predictions=predictions,\n references=references,\n )\n\n eval_metric = metric.compute()\n # Use accelerator.print to print only on the main process.\n accelerator.print(f\"epoch {epoch}:\", eval_metric)\n if args.with_tracking:\n accelerator.log(\n {\n \"accuracy\": eval_metric[\"accuracy\"],\n \"f1\": eval_metric[\"f1\"],\n \"train_loss\": total_loss.item() / len(train_dataloader),\n \"epoch\": epoch,\n },\n step=epoch,\n )\n\n if checkpointing_steps == \"epoch\":\n output_dir = f\"epoch_{epoch}\"\n if args.output_dir is not None:\n output_dir = os.path.join(args.output_dir, output_dir)\n accelerator.save_state(output_dir)\n\n accelerator.end_training()\n\n\ndef main():\n parser = argparse.ArgumentParser(description=\"Simple example of training script.\")\n parser.add_argument(\n \"--mixed_precision\",\n type=str,\n default=None,\n choices=[\"no\", \"fp16\", \"bf16\", \"fp8\"],\n help=\"Whether to use mixed precision. Choose\"\n \"between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10.\"\n \"and an Nvidia Ampere GPU.\",\n )\n parser.add_argument(\"--cpu\", action=\"store_true\", help=\"If passed, will train on the CPU.\")\n parser.add_argument(\n \"--checkpointing_steps\",\n type=str,\n default=None,\n help=\"Whether the various states should be saved at the end of every n steps, or 'epoch' for each epoch.\",\n )\n parser.add_argument(\n \"--resume_from_checkpoint\",\n type=str,\n default=None,\n help=\"If the training should continue from a checkpoint folder.\",\n )\n parser.add_argument(\n \"--use_stateful_dataloader\",\n action=\"store_true\",\n help=\"If the dataloader should be a resumable stateful dataloader.\",\n )\n parser.add_argument(\n \"--with_tracking\",\n action=\"store_true\",\n help=\"Whether to load in all available experiment trackers from the environment and use them for logging.\",\n )\n parser.add_argument(\n \"--output_dir\",\n type=str,\n default=\".\",\n help=\"Optional save directory where all checkpoint folders will be stored. Default is the current working directory.\",\n )\n parser.add_argument(\n \"--project_dir\",\n type=str,\n default=\"logs\",\n help=\"Location on where to store experiment tracking logs` and relevent project information\",\n )\n args = parser.parse_args()\n config = {\"lr\": 2e-5, \"num_epochs\": 3, \"seed\": 42, \"batch_size\": 16}\n training_function(config, args)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "examples/config_yaml_templates/README.md", "content": "# Config Zoo\n\nThis folder contains a variety of minimal configurations for `Accelerate` achieving certain goals. You can use these \ndirect config YAML's, or build off of them for your own YAML's.\n\nThese are highly annoted versions, aiming to teach you what each section does.\n\nEach config can be run via `accelerate launch --config_file {file} run_me.py`\n\n`run_me.py` will then print out how the current environment is setup (the contents of the `AcceleratorState`)" }, { "path": "examples/config_yaml_templates/deepspeed.yaml", "content": "# Similar to FSDP, we set the distributed type as DEEPSPEED\ndistributed_type: DEEPSPEED\n# With DeepSpeed, we utilize a deepspeed config file for the entire configuration\ndeepspeed_config:\n # Can also be any of the config json's in accelerate/examples/deepspeed_config_templates\n deepspeed_config_file: ../deepspeed_config_templates/zero_stage1_config.json\n # If using ZeRO-3 and wanting to load big models in, this should be set to `true` so \n # `transformers` uses the right `init` function\n zero3_init_flag: false # true \n\n# Finally we need to specify the number of accelerators to use\nnum_processes: 2\n# Optionally we can set the mixed precision now instead of in the deepspeed config file,\n# however this requires the `fp16` and `bf16` options to be set to `auto` in the deepspeed config file\n# mixed_precision: \"bf16\"\n" }, { "path": "examples/config_yaml_templates/fp8.yaml", "content": "# This config template simply setups up the TransformersEngine config (and a config for a single GPU),\n# this can interop with the other configs in this folder\ndistributed_type: \"NO\"\nmixed_precision: \"fp8\"\n# Then we specify the fp8 configuration:\nfp8_config:\n backend: TE # Can be TE | MS-AMP\n # The following are TE specific arguments.\n # See https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/api/common.html#common-api for more details\n amax_history_len: 1024\n fp8_format: E4M3\n interval: 1\n margin: 0\n override_linear_precision: [false, false, false]\n # Generally this should always be set to `false` to have the most realistic fp8 eval performance\n use_autocast_during_eval: false\n # If using MS-AMP, we ignore all of the prior and set a opt_level\n #opt_level: O1\n" }, { "path": "examples/config_yaml_templates/fsdp.yaml", "content": "# Since we are doing FSDP (even though it's multi-accelerator), we need to specify the distributed type as FSDP\ndistributed_type: FSDP\n# Can be one of \"no\", \"fp16\", or \"bf16\" (see `transformer_engine.yaml` for `fp8`, but it works for FSDP as well)\nmixed_precision: 'bf16'\n# Specify the number of accelerators to use\nnum_processes: 2\n# Then we can specify the FSDP config\nfsdp_config:\n fsdp_activation_checkpointing: false\n fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP\n fsdp_backward_prefetch: BACKWARD_PRE\n fsdp_cpu_ram_efficient_loading: true\n fsdp_forward_prefetch: false\n fsdp_offload_params: false\n fsdp_sharding_strategy: FULL_SHARD\n fsdp_state_dict_type: SHARDED_STATE_DICT\n fsdp_sync_module_states: true\n fsdp_use_orig_params: true\n" }, { "path": "examples/config_yaml_templates/multi_gpu.yaml", "content": "# Specify distributed_type as `MULTI_GPU` for DDP\ndistributed_type: \"MULTI_GPU\"\n# Can be one of \"no\", \"fp16\", or \"bf16\" (see `transformer_engine.yaml` for `fp8`)\nmixed_precision: \"bf16\"\n# Specify the number of GPUs to use\nnum_processes: 2" }, { "path": "examples/config_yaml_templates/multi_node.yaml", "content": "# This config template is for a multi-node setup. This assumes DDP, but can be interop'd with the other configs in this folder\n# Generally it's recommended to look at the SLURM config template for a more robust multi-node setup\ndistributed_type: MULTI_GPU\n# We need to specify the current machine's rank\nmachine_rank: 0\n# We then need to specify the IP address and port of the main process\nmain_process_ip: '1234'\nmain_process_port: 9999\n# We need to specify the number of machines\nnum_machines: 2\n# We need to specify the *total* number of processes\nnum_processes: 8\n# And then we need to specify how rdvz comms will be handled \nrdzv_backend: static # or c10d\n# If the compute nodes are on the same network (cloud will more than likely be false)\nsame_network: false\n" }, { "path": "examples/config_yaml_templates/multi_xpu.yaml", "content": "# Specify distributed_type as `MULTI_XPU` for DDP\ndistributed_type: \"MULTI_XPU\"\n# Can be one of \"no\", \"fp16\", or \"bf16\" (see `transformer_engine.yaml` for `fp8`)\nmixed_precision: \"bf16\"\n# Specify the number of XPUs to use\nnum_processes: 2\n" }, { "path": "examples/config_yaml_templates/run_me.py", "content": "# Copyright 2024 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nA base script which outputs the accelerate config for the given environment\n\"\"\"\n\nfrom accelerate import Accelerator\n\n\naccelerator = Accelerator()\n\naccelerator.print(f\"Accelerator state from the current environment:\\n{accelerator.state}\")\nif accelerator.fp8_recipe_handler is not None:\n accelerator.print(f\"FP8 config:\\n{accelerator.fp8_recipe_handler}\")\naccelerator.end_training()\n" }, { "path": "examples/config_yaml_templates/single_accelerator.yaml", "content": "# Since this is single GPU/XPU, we don't need distributed training\ndistributed_type: \"NO\"\n# Can be one of \"no\", \"fp16\", or \"bf16\" (see `transformer_engine.yaml` for `fp8`)\nmixed_precision: \"bf16\"\n" }, { "path": "examples/cv_example.py", "content": "# Copyright 2021 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport argparse\nimport os\nimport re\n\nimport numpy as np\nimport PIL\nimport torch\nfrom timm import create_model\nfrom torch.optim.lr_scheduler import OneCycleLR\nfrom torch.utils.data import DataLoader, Dataset\nfrom torchvision.transforms import Compose, RandomResizedCrop, Resize, ToTensor\n\nfrom accelerate import Accelerator\nfrom accelerate.utils import set_seed\n\n\n########################################################################\n# This is a fully working simple example to use Accelerate\n#\n# This example trains a ResNet50 on the Oxford-IIT Pet Dataset\n# in any of the following settings (with the same script):\n# - single CPU or single GPU\n# - multi GPUS (using PyTorch distributed mode)\n# - (multi) TPUs\n# - fp16 (mixed-precision) or fp32 (normal precision)\n#\n# To run it in each of these various modes, follow the instructions\n# in the readme for examples:\n# https://github.com/huggingface/accelerate/tree/main/examples\n#\n########################################################################\n\n\n# Function to get the label from the filename\ndef extract_label(fname):\n stem = fname.split(os.path.sep)[-1]\n return re.search(r\"^(.*)_\\d+\\.jpg$\", stem).groups()[0]\n\n\nclass PetsDataset(Dataset):\n def __init__(self, file_names, image_transform=None, label_to_id=None):\n self.file_names = file_names\n self.image_transform = image_transform\n self.label_to_id = label_to_id\n\n def __len__(self):\n return len(self.file_names)\n\n def __getitem__(self, idx):\n fname = self.file_names[idx]\n raw_image = PIL.Image.open(fname)\n image = raw_image.convert(\"RGB\")\n if self.image_transform is not None:\n image = self.image_transform(image)\n label = extract_label(fname)\n if self.label_to_id is not None:\n label = self.label_to_id[label]\n return {\"image\": image, \"label\": label}\n\n\ndef training_function(config, args):\n # Initialize accelerator\n accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)\n\n # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs\n lr = config[\"lr\"]\n num_epochs = int(config[\"num_epochs\"])\n seed = int(config[\"seed\"])\n batch_size = int(config[\"batch_size\"])\n image_size = config[\"image_size\"]\n if not isinstance(image_size, (list, tuple)):\n image_size = (image_size, image_size)\n\n # Grab all the image filenames\n file_names = [os.path.join(args.data_dir, fname) for fname in os.listdir(args.data_dir) if fname.endswith(\".jpg\")]\n\n # Build the label correspondences\n all_labels = [extract_label(fname) for fname in file_names]\n id_to_label = list(set(all_labels))\n id_to_label.sort()\n label_to_id = {lbl: i for i, lbl in enumerate(id_to_label)}\n\n # Set the seed before splitting the data.\n set_seed(seed)\n # Split our filenames between train and validation\n random_perm = np.random.permutation(len(file_names))\n cut = int(0.8 * len(file_names))\n train_split = random_perm[:cut]\n eval_split = random_perm[cut:]\n\n # For training we use a simple RandomResizedCrop\n train_tfm = Compose([RandomResizedCrop(image_size, scale=(0.5, 1.0)), ToTensor()])\n train_dataset = PetsDataset(\n [file_names[i] for i in train_split], image_transform=train_tfm, label_to_id=label_to_id\n )\n\n # For evaluation, we use a deterministic Resize\n eval_tfm = Compose([Resize(image_size), ToTensor()])\n eval_dataset = PetsDataset([file_names[i] for i in eval_split], image_transform=eval_tfm, label_to_id=label_to_id)\n\n # Instantiate dataloaders.\n train_dataloader = DataLoader(train_dataset, shuffle=True, batch_size=batch_size, num_workers=4)\n eval_dataloader = DataLoader(eval_dataset, shuffle=False, batch_size=batch_size, num_workers=4)\n\n # Instantiate the model (we build the model here so that the seed also control new weights initialization)\n model = create_model(\"resnet50d\", pretrained=True, num_classes=len(label_to_id))\n\n # We could avoid this line since the accelerator is set with `device_placement=True` (default value).\n # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer\n # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).\n model = model.to(accelerator.device)\n\n # Freezing the base model\n for param in model.parameters():\n param.requires_grad = False\n for param in model.get_classifier().parameters():\n param.requires_grad = True\n\n # We normalize the batches of images to be a bit faster.\n mean = torch.tensor(model.default_cfg[\"mean\"])[None, :, None, None].to(accelerator.device)\n std = torch.tensor(model.default_cfg[\"std\"])[None, :, None, None].to(accelerator.device)\n\n # Instantiate optimizer\n optimizer = torch.optim.Adam(params=model.parameters(), lr=lr / 25)\n\n # Instantiate learning rate scheduler\n lr_scheduler = OneCycleLR(optimizer=optimizer, max_lr=lr, epochs=num_epochs, steps_per_epoch=len(train_dataloader))\n\n # Prepare everything\n # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the\n # prepare method.\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(\n model, optimizer, train_dataloader, eval_dataloader, lr_scheduler\n )\n\n # Now we train the model\n for epoch in range(num_epochs):\n model.train()\n for step, batch in enumerate(train_dataloader):\n # We could avoid this line since we set the accelerator with `device_placement=True`.\n batch = {k: v.to(accelerator.device) for k, v in batch.items()}\n inputs = (batch[\"image\"] - mean) / std\n outputs = model(inputs)\n loss = torch.nn.functional.cross_entropy(outputs, batch[\"label\"])\n accelerator.backward(loss)\n optimizer.step()\n lr_scheduler.step()\n optimizer.zero_grad()\n\n model.eval()\n accurate = 0\n num_elems = 0\n for _, batch in enumerate(eval_dataloader):\n # We could avoid this line since we set the accelerator with `device_placement=True`.\n batch = {k: v.to(accelerator.device) for k, v in batch.items()}\n inputs = (batch[\"image\"] - mean) / std\n with torch.no_grad():\n outputs = model(inputs)\n predictions = outputs.argmax(dim=-1)\n predictions, references = accelerator.gather_for_metrics((predictions, batch[\"label\"]))\n accurate_preds = predictions == references\n num_elems += accurate_preds.shape[0]\n accurate += accurate_preds.long().sum()\n\n eval_metric = accurate.item() / num_elems\n # Use accelerator.print to print only on the main process.\n accelerator.print(f\"epoch {epoch}: {100 * eval_metric:.2f}\")\n accelerator.end_training()\n\n\ndef main():\n parser = argparse.ArgumentParser(description=\"Simple example of training script.\")\n parser.add_argument(\"--data_dir\", required=True, help=\"The data folder on disk.\")\n parser.add_argument(\n \"--mixed_precision\",\n type=str,\n default=None,\n choices=[\"no\", \"fp16\", \"bf16\", \"fp8\"],\n help=\"Whether to use mixed precision. Choose\"\n \"between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10.\"\n \"and an Nvidia Ampere GPU.\",\n )\n parser.add_argument(\"--cpu\", action=\"store_true\", help=\"If passed, will train on the CPU.\")\n args = parser.parse_args()\n config = {\"lr\": 3e-2, \"num_epochs\": 3, \"seed\": 42, \"batch_size\": 64, \"image_size\": 224}\n training_function(config, args)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "examples/deepspeed_config_templates/zero_stage1_config.json", "content": "{\n \"fp16\": {\n \"enabled\": true,\n \"loss_scale\": 0,\n \"loss_scale_window\": 1000,\n \"initial_scale_power\": 16,\n \"hysteresis\": 2,\n \"min_loss_scale\": 1\n },\n \"optimizer\": {\n \"type\": \"AdamW\",\n \"params\": {\n \"lr\": \"auto\",\n \"weight_decay\": \"auto\",\n \"torch_adam\": true,\n \"adam_w_mode\": true\n }\n },\n \"scheduler\": {\n \"type\": \"WarmupDecayLR\",\n \"params\": {\n \"warmup_min_lr\": \"auto\",\n \"warmup_max_lr\": \"auto\",\n \"warmup_num_steps\": \"auto\",\n \"total_num_steps\": \"auto\"\n }\n },\n \"zero_optimization\": {\n \"stage\": 1,\n \"allgather_partitions\": true,\n \"allgather_bucket_size\": 2e8,\n \"overlap_comm\": true,\n \"reduce_scatter\": true,\n \"reduce_bucket_size\": \"auto\",\n \"contiguous_gradients\": true\n },\n \"gradient_accumulation_steps\": 1,\n \"gradient_clipping\": \"auto\",\n \"steps_per_print\": 2000,\n \"train_batch_size\": \"auto\",\n \"train_micro_batch_size_per_gpu\": \"auto\",\n \"wall_clock_breakdown\": false\n}" }, { "path": "examples/deepspeed_config_templates/zero_stage2_config.json", "content": "{\n \"fp16\": {\n \"enabled\": true,\n \"loss_scale\": 0,\n \"loss_scale_window\": 1000,\n \"initial_scale_power\": 16,\n \"hysteresis\": 2,\n \"min_loss_scale\": 1\n },\n \"optimizer\": {\n \"type\": \"AdamW\",\n \"params\": {\n \"lr\": \"auto\",\n \"weight_decay\": \"auto\",\n \"torch_adam\": true,\n \"adam_w_mode\": true\n }\n },\n \"scheduler\": {\n \"type\": \"WarmupDecayLR\",\n \"params\": {\n \"warmup_min_lr\": \"auto\",\n \"warmup_max_lr\": \"auto\",\n \"warmup_num_steps\": \"auto\",\n \"total_num_steps\": \"auto\"\n }\n },\n \"zero_optimization\": {\n \"stage\": 2,\n \"allgather_partitions\": true,\n \"allgather_bucket_size\": 2e8,\n \"overlap_comm\": true,\n \"reduce_scatter\": true,\n \"reduce_bucket_size\": \"auto\",\n \"contiguous_gradients\": true\n },\n \"gradient_accumulation_steps\": 1,\n \"gradient_clipping\": \"auto\",\n \"steps_per_print\": 2000,\n \"train_batch_size\": \"auto\",\n \"train_micro_batch_size_per_gpu\": \"auto\",\n \"wall_clock_breakdown\": false\n}" }, { "path": "examples/deepspeed_config_templates/zero_stage2_offload_config.json", "content": "{\n \"fp16\": {\n \"enabled\": true,\n \"loss_scale\": 0,\n \"loss_scale_window\": 1000,\n \"initial_scale_power\": 16,\n \"hysteresis\": 2,\n \"min_loss_scale\": 1\n },\n \"optimizer\": {\n \"type\": \"AdamW\",\n \"params\": {\n \"lr\": \"auto\",\n \"weight_decay\": \"auto\",\n \"torch_adam\": true,\n \"adam_w_mode\": true\n }\n },\n \"scheduler\": {\n \"type\": \"WarmupDecayLR\",\n \"params\": {\n \"warmup_min_lr\": \"auto\",\n \"warmup_max_lr\": \"auto\",\n \"warmup_num_steps\": \"auto\",\n \"total_num_steps\": \"auto\"\n }\n },\n \"zero_optimization\": {\n \"stage\": 2,\n \"offload_optimizer\": {\n \"device\": \"cpu\",\n \"pin_memory\": true\n },\n \"allgather_partitions\": true,\n \"allgather_bucket_size\": 2e8,\n \"overlap_comm\": true,\n \"reduce_scatter\": true,\n \"reduce_bucket_size\": \"auto\",\n \"contiguous_gradients\": true\n },\n \"gradient_accumulation_steps\": 1,\n \"gradient_clipping\": \"auto\",\n \"steps_per_print\": 2000,\n \"train_batch_size\": \"auto\",\n \"train_micro_batch_size_per_gpu\": \"auto\",\n \"wall_clock_breakdown\": false\n}" }, { "path": "examples/deepspeed_config_templates/zero_stage3_config.json", "content": "{\n \"fp16\": {\n \"enabled\": true,\n \"loss_scale\": 0,\n \"loss_scale_window\": 1000,\n \"initial_scale_power\": 16,\n \"hysteresis\": 2,\n \"min_loss_scale\": 1\n },\n \"optimizer\": {\n \"type\": \"AdamW\",\n \"params\": {\n \"lr\": \"auto\",\n \"weight_decay\": \"auto\"\n }\n },\n \"scheduler\": {\n \"type\": \"WarmupDecayLR\",\n \"params\": {\n \"warmup_min_lr\": \"auto\",\n \"warmup_max_lr\": \"auto\",\n \"warmup_num_steps\": \"auto\",\n \"total_num_steps\": \"auto\"\n }\n },\n \"zero_optimization\": {\n \"stage\": 3,\n \"overlap_comm\": true,\n \"contiguous_gradients\": true,\n \"reduce_bucket_size\": \"auto\",\n \"stage3_prefetch_bucket_size\": \"auto\",\n \"stage3_param_persistence_threshold\": \"auto\",\n \"sub_group_size\": 1e9,\n \"stage3_max_live_parameters\": 1e9,\n \"stage3_max_reuse_distance\": 1e9,\n \"stage3_gather_16bit_weights_on_model_save\": \"auto\"\n },\n \"gradient_accumulation_steps\": 1,\n \"gradient_clipping\": \"auto\",\n \"steps_per_print\": 2000,\n \"train_batch_size\": \"auto\",\n \"train_micro_batch_size_per_gpu\": \"auto\",\n \"wall_clock_breakdown\": false\n}" }, { "path": "examples/deepspeed_config_templates/zero_stage3_offload_config.json", "content": "{\n \"fp16\": {\n \"enabled\": true,\n \"loss_scale\": 0,\n \"loss_scale_window\": 1000,\n \"initial_scale_power\": 16,\n \"hysteresis\": 2,\n \"min_loss_scale\": 1\n },\n \"optimizer\": {\n \"type\": \"AdamW\",\n \"params\": {\n \"lr\": \"auto\",\n \"weight_decay\": \"auto\"\n }\n },\n \"scheduler\": {\n \"type\": \"WarmupDecayLR\",\n \"params\": {\n \"warmup_min_lr\": \"auto\",\n \"warmup_max_lr\": \"auto\",\n \"warmup_num_steps\": \"auto\",\n \"total_num_steps\": \"auto\"\n }\n },\n \"zero_optimization\": {\n \"stage\": 3,\n \"offload_optimizer\": {\n \"device\": \"cpu\",\n \"pin_memory\": true\n },\n \"offload_param\": {\n \"device\": \"cpu\",\n \"pin_memory\": true\n },\n \"overlap_comm\": true,\n \"contiguous_gradients\": true,\n \"reduce_bucket_size\": \"auto\",\n \"stage3_prefetch_bucket_size\": \"auto\",\n \"stage3_param_persistence_threshold\": \"auto\",\n \"sub_group_size\": 1e9,\n \"stage3_max_live_parameters\": 1e9,\n \"stage3_max_reuse_distance\": 1e9,\n \"stage3_gather_16bit_weights_on_model_save\": \"auto\"\n },\n \"gradient_accumulation_steps\": 1,\n \"gradient_clipping\": \"auto\",\n \"steps_per_print\": 2000,\n \"train_batch_size\": \"auto\",\n \"train_micro_batch_size_per_gpu\": \"auto\",\n \"wall_clock_breakdown\": false\n}" }, { "path": "examples/finetune_lm_tpu.py", "content": "# Copyright 2025 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# Example of fine-tuning a model on a TPU using FSDPv2, TRL and PEFT.\n#\n# Run the script with:\n# python finetune_lm_tpu.py [--model_id MODEL_ID] [--dataset_id DATASET_ID]\n#\n# This script has been tested on a TPU v5 litepod-8.\n\nimport argparse\n\nimport torch\nimport torch_xla.runtime as xr\nfrom datasets import load_dataset\nfrom peft import LoraConfig\nfrom transformers import AutoModelForCausalLM, AutoTokenizer\nfrom trl import SFTConfig, SFTTrainer\n\n\n# FSDPv2 requires SPMD to be enabled.\nxr.use_spmd()\n\n\ndef format_dolly(example, tokenizer):\n \"\"\"Format Dolly dataset examples using the tokenizer's chat template.\"\"\"\n user_content = example[\"instruction\"]\n if len(example[\"context\"]) > 0:\n user_content += f\"\\n\\nContext: {example['context']}\"\n\n messages = [\n {\n \"role\": \"system\",\n \"content\": \"You are a helpful assistant\",\n },\n {\"role\": \"user\", \"content\": user_content},\n {\"role\": \"assistant\", \"content\": example[\"response\"]},\n ]\n\n return tokenizer.apply_chat_template(messages, tokenize=False)\n\n\ndef train(model_id, dataset):\n # Load model with low_cpu_mem_usage to avoid loading full model into CPU memory\n # FSDPv2 will handle sharding across TPUs\n model = AutoModelForCausalLM.from_pretrained(\n model_id,\n use_cache=False,\n torch_dtype=torch.bfloat16,\n low_cpu_mem_usage=True,\n device_map=None, # Let FSDP handle device placement\n )\n\n tokenizer = AutoTokenizer.from_pretrained(model_id)\n\n if tokenizer.pad_token is None:\n if model.config.model_type == \"llama\":\n # Vanilla Llama models have a finetune gith pad id token\n tokenizer.pad_token = \"<|finetune_right_pad_id|>\"\n elif tokenizer.eos_token is not None:\n tokenizer.pad_token = tokenizer.eos_token\n else:\n raise ValueError(f\"Cannot get or guess pad token for model {model_id}.\")\n\n if tokenizer.chat_template is None:\n # Set chat template for Llama 3.1 format\n tokenizer.chat_template = (\n \"{% for message in messages %}\"\n \"{% if message['role'] == 'system' %}\"\n \"<|begin_of_text|><|start_header_id|>system<|end_header_id|>\\n\\n{{ message['content'] }}<|eot_id|>\"\n \"{% elif message['role'] == 'user' %}\"\n \"<|start_header_id|>user<|end_header_id|>\\n\\n{{ message['content'] }}<|eot_id|>\"\n \"{% elif message['role'] == 'assistant' %}\"\n \"<|start_header_id|>assistant<|end_header_id|>\\n\\n{{ message['content'] }}<|eot_id|>\"\n \"{% endif %}\"\n \"{% endfor %}\"\n \"{% if add_generation_prompt %}\"\n \"<|start_header_id|>assistant<|end_header_id|>\\n\\n\"\n \"{% endif %}\"\n )\n\n # Try to guess the DecoderLayer class name, based on common model architectures\n transformer_layer_cls_to_wrap = model.model.layers[0].__class__.__name__\n\n # Get FSDP training arguments\n fsdp_training_args = {\n \"fsdp\": \"full_shard\",\n \"fsdp_config\": {\n \"transformer_layer_cls_to_wrap\": [transformer_layer_cls_to_wrap],\n \"xla\": True,\n \"xla_fsdp_v2\": True,\n \"xla_fsdp_grad_ckpt\": True,\n },\n }\n\n # Set up PEFT LoRA for fine-tuning.\n lora_config = LoraConfig(\n r=32,\n lora_alpha=128,\n lora_dropout=0.05,\n target_modules=[\"q_proj\", \"k_proj\"],\n task_type=\"CAUSAL_LM\",\n )\n\n sft_config = SFTConfig(\n gradient_checkpointing=False, # Required on TPU, not supported\n max_length=1024,\n per_device_train_batch_size=4,\n num_train_epochs=3,\n max_steps=-1,\n output_dir=\"./output\",\n optim=\"adafactor\",\n logging_steps=1,\n dataloader_drop_last=True, # Required for FSDPv2.\n dataset_text_field=\"text\",\n packing=True,\n **fsdp_training_args,\n )\n\n # Set up the trainer\n trainer = SFTTrainer(\n model=model,\n train_dataset=dataset,\n args=sft_config,\n peft_config=lora_config,\n processing_class=tokenizer,\n formatting_func=lambda example: format_dolly(example, tokenizer),\n )\n\n trainer.train()\n\n\n# =============================================================================\n# Main Function\n# =============================================================================\nif __name__ == \"__main__\":\n parser = argparse.ArgumentParser(description=\"Simple example of training script.\")\n\n parser.add_argument(\n \"--model_id\", \"-m\", type=str, default=\"meta-llama/Llama-3.2-1B\", help=\"Model id to use for training.\"\n )\n parser.add_argument(\n \"--dataset_id\",\n \"-d\",\n type=str,\n default=\"databricks/databricks-dolly-15k\",\n help=\"Dataset id to use for training.\",\n )\n\n args = parser.parse_args()\n\n # NOTE: this section can be adapted to load any dataset you want.\n dataset_id = args.dataset_id\n dolly_dataset = load_dataset(dataset_id, split=\"train\")\n\n train(\n model_id=args.model_id,\n dataset=dolly_dataset,\n )\n" }, { "path": "examples/inference/distributed/README.md", "content": "# Distributed inference examples\n\nThis folder contains a variety of tutorials for running distributed inference with the following strategy: \n\nLoad an entire model onto each GPU and sending chunks of a batch through each GPUā€™s model copy at a time\n\n## Installation\n\n```bash\npip install accelerate torch\n```\n\n## Running code\n\nYou can either use `torchrun` or the recommended way of `accelerate launch` (without needing to run `accelerate config`) on each script:\n\n```bash\naccelerate launch --num_processes {NUM_GPUS} phi2.py\n```\n\nOr:\n\n```bash\ntorchrun --nproc-per-node {NUM_GPUS} phi2.py\n```\n" }, { "path": "examples/inference/distributed/distributed_image_generation.py", "content": "# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nOriginally by jiwooya1000, put together together by sayakpaul.\nDocumentation: https://huggingface.co/docs/diffusers/main/en/training/distributed_inference\n\nRun:\n\naccelerate launch distributed_image_generation.py --batch_size 8\n\n# Enable memory optimizations for large models like SD3\naccelerate launch distributed_image_generation.py --batch_size 8 --low_mem\n\"\"\"\n\nimport os\nimport time\n\nimport fire\nimport torch\nfrom datasets import load_dataset\nfrom diffusers import DiffusionPipeline\nfrom tqdm import tqdm\n\nfrom accelerate import PartialState\nfrom accelerate.utils import gather_object\n\n\nSTART_TIME = time.strftime(\"%Y%m%d_%H%M%S\")\nDTYPE_MAP = {\"fp32\": torch.float32, \"fp16\": torch.float16, \"bf16\": torch.bfloat16}\n\n\ndef get_batches(items, batch_size):\n num_batches = (len(items) + batch_size - 1) // batch_size\n batches = []\n\n for i in range(num_batches):\n start_index = i * batch_size\n end_index = min((i + 1) * batch_size, len(items))\n batch = items[start_index:end_index]\n batches.append(batch)\n\n return batches\n\n\ndef main(\n ckpt_id: str = \"PixArt-alpha/PixArt-Sigma-XL-2-1024-MS\",\n save_dir: str = \"./evaluation/examples\",\n seed: int = 1,\n batch_size: int = 4,\n num_inference_steps: int = 20,\n guidance_scale: float = 4.5,\n dtype: str = \"fp16\",\n low_mem: bool = False,\n):\n pipeline = DiffusionPipeline.from_pretrained(ckpt_id, torch_dtype=DTYPE_MAP[dtype])\n\n save_dir = save_dir + f\"_{START_TIME}\"\n\n parti_prompts = load_dataset(\"nateraw/parti-prompts\", split=\"train\")\n data_loader = get_batches(items=parti_prompts[\"Prompt\"], batch_size=batch_size)\n\n distributed_state = PartialState()\n if low_mem:\n pipeline.enable_model_cpu_offload(gpu_id=distributed_state.device.index)\n else:\n pipeline = pipeline.to(distributed_state.device)\n\n if distributed_state.is_main_process:\n if not os.path.exists(save_dir):\n os.makedirs(save_dir)\n print(f\"Directory '{save_dir}' created successfully.\")\n else:\n print(f\"Directory '{save_dir}' already exists.\")\n\n count = 0\n for _, prompts_raw in tqdm(enumerate(data_loader), total=len(data_loader)):\n input_prompts = []\n\n with distributed_state.split_between_processes(prompts_raw) as prompts:\n generator = torch.manual_seed(seed)\n images = pipeline(\n prompts, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, generator=generator\n ).images\n input_prompts.extend(prompts)\n\n distributed_state.wait_for_everyone()\n\n images = gather_object(images)\n input_prompts = gather_object(input_prompts)\n\n if distributed_state.is_main_process:\n for image, prompt in zip(images, input_prompts):\n count += 1\n temp_dir = os.path.join(save_dir, f\"example_{count}\")\n\n os.makedirs(temp_dir)\n prompt = \"_\".join(prompt.split())\n image.save(f\"image_{prompt}.png\")\n\n if distributed_state.is_main_process:\n print(f\">>> Image Generation Finished. Saved in {save_dir}\")\n\n\nif __name__ == \"__main__\":\n fire.Fire(main)\n" }, { "path": "examples/inference/distributed/distributed_speech_generation.py", "content": "# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nimport json\nimport os\nimport pathlib\nimport queue\nfrom concurrent.futures import ThreadPoolExecutor\nfrom typing import Union\n\nimport fire\nimport scipy.io.wavfile\nimport torch\nfrom datasets import load_dataset\nfrom transformers import AutoTokenizer, VitsModel\n\nfrom accelerate import PartialState\nfrom accelerate.utils import tqdm\n\n\n\"\"\"\nRequirements: transformers accelerate fire scipy datasets\npip install transformers accelerate fire scipy datasets\nExample usage:\naccelerate launch distributed_speech_generation.py --output_path outputs --batch_size 8 --num_workers 2 --dataset_split train\n\"\"\"\n\n\"\"\"\nTo run the speech generation\nimport scipy.io.wavfile\nimport numpy as np\nfrom IPython.display import Audio\nsample_rate, audio_data = scipy.io.wavfile.read('path_to_you_wav_file.wav')\naudio_data = audio_data.astype(np.float32) / 32762.0\nAudio(audio_data, rate=sample_rate)\n\"\"\"\n\n\ndef load_pokemon_data(split: str, max_text_length: int):\n \"\"\"Load Pokemon descriptions from the dataset\"\"\"\n ds = load_dataset(\"svjack/pokemon-blip-captions-en-zh\", split=split)\n\n # Create dataset of dictionaries\n dataset = []\n for idx, text in enumerate(ds[\"en_text\"]):\n if len(text.strip()) > 0: # Skip empty descriptions\n dataset.append(\n {\n \"id\": f\"pokemon_{idx:06d}\",\n \"text\": text.strip()[:max_text_length], # Truncate long descriptions\n \"original_text\": text.strip(), # Keep original for metadata\n }\n )\n return dataset\n\n\nclass ExistsFilter:\n def __init__(self, output_dir: Union[pathlib.Path, str]):\n current_files = [f.split(\".wav\")[0] for f in os.listdir(output_dir) if f.endswith(\".wav\")]\n self.processed_files = set(current_files)\n print(f\"Existing audio files found: {len(self.processed_files)}.\")\n\n def __call__(self, x):\n return x[\"id\"] not in self.processed_files\n\n\ndef preprocess_fn(sample, tokenizer, max_text_length: int):\n inputs = tokenizer(sample[\"text\"], padding=False, truncation=True, max_length=max_text_length, return_tensors=\"pt\")\n\n return {\n \"input_ids\": inputs[\"input_ids\"][0].tolist(),\n \"attention_mask\": inputs[\"attention_mask\"][0].tolist(),\n \"id\": sample[\"id\"],\n \"text\": sample[\"text\"],\n \"original_text\": sample[\"original_text\"],\n }\n\n\ndef collate_fn(examples, tokenizer):\n \"\"\"Collate batch of examples with proper padding\"\"\"\n # Find max length in this batch\n max_length = max(len(example[\"input_ids\"]) for example in examples)\n\n # Pad sequences to max_length\n input_ids_list = []\n attention_mask_list = []\n\n for example in examples:\n # Get current lengths\n curr_len = len(example[\"input_ids\"])\n padding_length = max_length - curr_len\n\n # Pad sequences\n padded_input_ids = example[\"input_ids\"] + [tokenizer.pad_token_id] * padding_length\n padded_attention_mask = example[\"attention_mask\"] + [0] * padding_length\n\n input_ids_list.append(padded_input_ids)\n attention_mask_list.append(padded_attention_mask)\n\n # Convert to tensors\n input_ids = torch.tensor(input_ids_list, dtype=torch.long)\n attention_mask = torch.tensor(attention_mask_list, dtype=torch.long)\n\n ids = [example[\"id\"] for example in examples]\n texts = [example[\"text\"] for example in examples]\n original_texts = [example[\"original_text\"] for example in examples]\n\n return {\n \"input_ids\": input_ids,\n \"attention_mask\": attention_mask,\n \"ids\": ids,\n \"texts\": texts,\n \"original_texts\": original_texts,\n }\n\n\ndef create_dataloader(dataset, batch_size, distributed_state, tokenizer):\n \"\"\"Create dataloader with preprocessing\"\"\"\n processed_dataset = [preprocess_fn(item, tokenizer, max_text_length=200) for item in dataset]\n\n # Split dataset for distributed processing\n if distributed_state.num_processes > 1:\n chunk_size = len(processed_dataset) // distributed_state.num_processes\n start_idx = distributed_state.process_index * chunk_size\n end_idx = (\n start_idx + chunk_size\n if distributed_state.process_index < distributed_state.num_processes - 1\n else len(processed_dataset)\n )\n processed_dataset = processed_dataset[start_idx:end_idx]\n\n # Create batches\n batches = []\n for i in range(0, len(processed_dataset), batch_size):\n batch = processed_dataset[i : i + batch_size]\n batches.append(collate_fn(batch, tokenizer))\n return batches\n\n\ndef save_results(output_queue: queue.Queue, output_dir: pathlib.Path, sampling_rate: int):\n while True:\n try:\n item = output_queue.get(timeout=5)\n if item is None:\n break\n waveforms, ids, texts, original_texts = item\n\n # Save each audio file and its metadata\n for waveform, file_id, text, original_text in zip(waveforms, ids, texts, original_texts):\n # Save audio\n wav_path = output_dir / f\"{file_id}.wav\"\n scipy.io.wavfile.write(wav_path, rate=sampling_rate, data=waveform.cpu().float().numpy())\n\n # Save metadata with both truncated and original text\n metadata = {\n \"text_used\": text,\n \"original_text\": original_text,\n \"model\": \"facebook/mms-tts-eng\",\n \"sampling_rate\": sampling_rate,\n }\n metadata_path = output_dir / f\"{file_id}_metadata.json\"\n with metadata_path.open(\"w\") as f:\n json.dump(metadata, f, indent=4)\n\n except queue.Empty:\n continue\n\n\ndef main(\n output_path: str = \"speech_data\",\n batch_size: int = 8,\n num_workers: int = 2,\n dataset_split: str = \"train\",\n model_name: str = \"facebook/mms-tts-eng\",\n max_text_length: int = 200,\n):\n output_dir = pathlib.Path(output_path)\n output_dir.mkdir(parents=True, exist_ok=True)\n distributed_state = PartialState()\n\n # Load model and tokenizer\n model = VitsModel.from_pretrained(\n model_name,\n device_map=distributed_state.device,\n torch_dtype=torch.float32,\n )\n\n tokenizer = AutoTokenizer.from_pretrained(model_name)\n\n # Load and filter data\n dataset = load_pokemon_data(dataset_split, max_text_length)\n exist_filter = ExistsFilter(output_dir)\n dataset = [item for item in dataset if exist_filter(item)]\n\n distributed_state.print(f\"Processing {len(dataset)} Pokemon descriptions\")\n\n # Create dataloader\n batches = create_dataloader(dataset, batch_size, distributed_state, tokenizer)\n\n # Setup output queue and save thread\n output_queue = queue.Queue()\n save_thread = ThreadPoolExecutor(max_workers=num_workers)\n save_future = save_thread.submit(save_results, output_queue, output_dir, model.config.sampling_rate)\n\n try:\n for batch in tqdm(batches, desc=\"Generating Pokemon descriptions\"):\n with torch.no_grad():\n outputs = model(\n input_ids=batch[\"input_ids\"].to(distributed_state.device, dtype=torch.long),\n attention_mask=batch[\"attention_mask\"].to(distributed_state.device, dtype=torch.long),\n ).waveform\n\n output_queue.put((outputs, batch[\"ids\"], batch[\"texts\"], batch[\"original_texts\"]))\n finally:\n output_queue.put(None)\n save_thread.shutdown(wait=True)\n\n save_future.result()\n\n\nif __name__ == \"__main__\":\n fire.Fire(main)\n" }, { "path": "examples/inference/distributed/florence2.py", "content": "# Copyright 2024 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport json\nimport os\nimport pathlib\nimport queue\nfrom concurrent.futures import ThreadPoolExecutor\nfrom functools import partial\nfrom typing import Union\n\nimport fire\nimport torch\nimport webdataset as wds\nfrom huggingface_hub.utils import insecure_hashlib\nfrom PIL import Image\nfrom tqdm import tqdm\nfrom transformers import AutoModelForCausalLM, AutoProcessor\n\nfrom accelerate import PartialState\n\n\n\"\"\"\nAdditional requirements: flash_attn einops timm webdataset fire tqdm huggingface_hub\npip install flash_attn einops timm webdataset fire tqdm huggingface_hub\n\nExample:\n\naccelerate launch --num_processes=2 florence2.py --data_path \"https://huggingface.co/datasets/pixparse/cc3m-wds/resolve/main/cc3m-train-0000.tar\" --output_path outputs --batch_size 12 --num_workers 1 --prompt \"\"\n\"\"\"\n\n\ndef main(\n data_path: str,\n output_path: str,\n batch_size: int,\n num_workers: int,\n prompt: str = \"\",\n model_name: str = \"microsoft/Florence-2-large\",\n max_new_tokens: int = 1024,\n num_beams: int = 3,\n):\n output_dir = pathlib.Path(output_path)\n\n distributed_state = PartialState()\n\n if distributed_state.is_main_process:\n output_dir.mkdir(exist_ok=True)\n\n model = AutoModelForCausalLM.from_pretrained(\n model_name,\n device_map=distributed_state.device,\n torch_dtype=torch.float16,\n trust_remote_code=True,\n )\n\n processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True, clean_up_tokenization_spaces=True)\n\n class ExistsFilter:\n def __init__(self, output_dir: Union[pathlib.Path, str]):\n current_training_img_hashes = [f.split(\".jpg\")[0] for f in os.listdir(output_dir) if f.endswith(\".jpg\")]\n self.current_training_img_hashes = set(current_training_img_hashes)\n if distributed_state.is_main_process:\n print(f\"Existing images found: {len(self.current_training_img_hashes)}.\")\n\n def __call__(self, x):\n if len(self.current_training_img_hashes) > 0:\n if x[\"img_hash\"] in self.current_training_img_hashes:\n return False\n else:\n return True\n else:\n return True\n\n def preprocess_fn(sample, processor):\n image: Image.Image = sample[\"jpg\"].convert(\"RGB\")\n img_hash = insecure_hashlib.sha1(image.tobytes()).hexdigest()\n inputs = processor(\n text=prompt,\n images=image,\n return_tensors=\"pt\",\n )\n return {\n \"input_ids\": inputs[\"input_ids\"],\n \"pixel_values\": inputs[\"pixel_values\"],\n \"image\": image,\n \"img_hash\": img_hash,\n \"original_caption\": sample[\"txt\"],\n }\n\n def collate_fn(examples):\n input_ids = torch.cat([example[\"input_ids\"] for example in examples])\n pixel_values = torch.cat([example[\"pixel_values\"] for example in examples])\n images = [example[\"image\"] for example in examples]\n img_hashes = [example[\"img_hash\"] for example in examples]\n captions = [example[\"original_caption\"] for example in examples]\n return {\n \"input_ids\": input_ids,\n \"pixel_values\": pixel_values,\n \"images\": images,\n \"img_hashes\": img_hashes,\n \"original_captions\": captions,\n }\n\n exist_filter = ExistsFilter(output_dir)\n dataset = (\n wds.WebDataset(\n data_path,\n handler=wds.warn_and_continue,\n nodesplitter=None,\n shardshuffle=False,\n empty_check=False,\n )\n .decode(\"pil\", handler=wds.warn_and_continue)\n .map(partial(preprocess_fn, processor=processor), handler=wds.warn_and_continue)\n )\n if len(exist_filter.current_training_img_hashes) > 0:\n dataset = dataset.select(exist_filter)\n dataset = dataset.batched(\n batch_size,\n partial=False,\n collation_fn=collate_fn,\n )\n dataloader = wds.WebLoader(\n dataset,\n batch_size=None,\n num_workers=num_workers,\n pin_memory=True,\n persistent_workers=True,\n )\n\n def save_results(output_queue: queue.Queue, output_dir: pathlib.Path, processor):\n while True:\n try:\n item = output_queue.get(timeout=5)\n if item is None:\n break\n original_captions, predictions, images, img_hashes = item\n predicted_captions = processor.batch_decode(\n predictions,\n skip_special_tokens=False,\n )\n for caption, pred_caption, image, img_hash in zip(\n original_captions, predicted_captions, images, img_hashes\n ):\n processed_caption = processor.post_process_generation(\n pred_caption, task=prompt, image_size=(image.width, image.height)\n )[prompt]\n img_path = output_dir.joinpath(f\"{img_hash}.jpg\")\n image.save(img_path)\n\n caption_dict = {\"original\": caption, \"predicted\": processed_caption}\n with output_dir.joinpath(f\"{img_hash}_caption.json\").open(\"w\") as f:\n json.dump(caption_dict, f, indent=4)\n\n except queue.Empty:\n continue\n\n output_queue = queue.Queue()\n save_thread = ThreadPoolExecutor(max_workers=num_workers)\n save_future = save_thread.submit(save_results, output_queue, output_dir, processor)\n\n try:\n for _, batch_raw in tqdm(\n enumerate(dataloader),\n disable=not distributed_state.is_main_process,\n ):\n with distributed_state.split_between_processes(batch_raw) as batch:\n outputs = model.generate(\n input_ids=batch[\"input_ids\"].to(distributed_state.device),\n pixel_values=batch[\"pixel_values\"].to(distributed_state.device, model.dtype),\n max_new_tokens=max_new_tokens,\n num_beams=num_beams,\n )\n output_queue.put(\n (\n batch[\"original_captions\"],\n outputs,\n batch[\"images\"],\n batch[\"img_hashes\"],\n )\n )\n finally:\n output_queue.put(None)\n save_thread.shutdown(wait=True)\n\n save_future.result()\n\n\nif __name__ == \"__main__\":\n fire.Fire(main)\n" }, { "path": "examples/inference/distributed/llava_next_video.py", "content": "# Copyright 2025 The HuggingFace Inc. team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport json\nimport os\nimport pathlib\nimport queue\nimport time\nfrom concurrent.futures import ThreadPoolExecutor\n\nimport av\nimport fire\nimport numpy as np\nimport torch\nfrom huggingface_hub import snapshot_download\nfrom tqdm import tqdm\nfrom transformers import LlavaNextVideoForConditionalGeneration, LlavaNextVideoProcessor\n\nfrom accelerate import PartialState\n\n\nSTART_TIME = time.strftime(\"%Y%m%d_%H%M%S\")\nDTYPE_MAP = {\"fp32\": torch.float32, \"fp16\": torch.float16, \"bf16\": torch.bfloat16}\n\n\n\"\"\"\nExample:\n\naccelerate launch llava_next_video.py\n\"\"\"\n\n\ndef save_results(output_queue: queue.Queue, output_dir: pathlib.Path):\n count = 0\n while True:\n try:\n item = output_queue.get(timeout=5)\n if item is None:\n break\n prompt, video, generated_text = item\n example_file = f\"example_{count}\"\n temp_dir = os.path.join(output_dir, example_file)\n\n metadata = {\"prompt\": prompt, \"video\": video, \"generated_text\": generated_text}\n with open(temp_dir, \"w\") as f:\n json.dump(metadata, f, indent=4)\n count += 1\n\n except queue.Empty:\n continue\n\n\ndef get_batches(processed_videos, batch_size):\n num_batches = (len(processed_videos) + batch_size - 1) // batch_size\n batches = []\n\n for i in range(num_batches):\n start_index = i * batch_size\n end_index = min((i + 1) * batch_size, len(processed_videos))\n batch = processed_videos[start_index:end_index]\n batches.append(batch)\n\n return batches\n\n\ndef read_video_pyav(container, indices):\n \"\"\"\n Decode the video with PyAV decoder.\n Args:\n container (`av.container.input.InputContainer`): PyAV container.\n indices (`List[int]`): List of frame indices to decode.\n Returns:\n result (np.ndarray): np array of decoded frames of shape (num_frames, height, width, 3).\n \"\"\"\n frames = []\n container.seek(0)\n start_index = indices[0]\n end_index = indices[-1]\n for i, frame in enumerate(container.decode(video=0)):\n if i > end_index:\n break\n if i >= start_index and i in indices:\n frames.append(frame)\n return np.stack([x.to_ndarray(format=\"rgb24\") for x in frames])\n\n\ndef get_video_paths(video_dir):\n \"\"\"Get paths to all video files in the directory and its subdirectories.\"\"\"\n video_extensions = (\".mp4\", \".avi\", \".mov\", \".mkv\") # Add more extensions if needed\n video_paths = []\n\n for root, _, files in os.walk(video_dir):\n for file in files:\n if file.lower().endswith(video_extensions):\n video_paths.append(os.path.join(root, file))\n\n return video_paths\n\n\ndef process_videos(video_paths, processor, prompt, frames_per_video):\n \"\"\"Process a batch of videos and prepare them for the model.\"\"\"\n batch_inputs = []\n\n for video_path in video_paths:\n try:\n with av.open(video_path) as container:\n total_frames = container.streams.video[0].frames\n indices = np.arange(0, total_frames, total_frames / frames_per_video).astype(int)\n clip = read_video_pyav(container, indices)\n\n processed = processor(text=prompt, videos=clip, return_tensors=\"pt\")\n batch_inputs.append(\n {\n \"input_ids\": processed[\"input_ids\"],\n \"pixel_values_videos\": processed[\"pixel_values_videos\"],\n \"video\": video_path,\n }\n )\n\n except Exception as e:\n print(f\"Error processing video {video_path}: {str(e)}\")\n continue\n\n return batch_inputs\n\n\ndef main(\n model_name: str = \"llava-hf/LLaVA-NeXT-Video-7B-hf\",\n save_dir: str = \"./evaluation/examples\",\n prompt: str = \"USER: